diff --git a/apps/sim/app/api/auth/oauth/token/route.test.ts b/apps/sim/app/api/auth/oauth/token/route.test.ts index b1c07e96b68..821fe2bfa66 100644 --- a/apps/sim/app/api/auth/oauth/token/route.test.ts +++ b/apps/sim/app/api/auth/oauth/token/route.test.ts @@ -24,6 +24,7 @@ vi.mock('@/lib/oauth/credential-service', () => ({ vi.mock('@/lib/auth/credential-access', () => ({ authorizeCredentialUse: mockAuthorizeCredentialUse, + authorizeCredentialUseForAuth: mockAuthorizeCredentialUse, })) import { TokenServiceAccountValidationError } from '@/lib/credentials/token-service-accounts/errors' diff --git a/apps/sim/app/api/auth/oauth/token/route.ts b/apps/sim/app/api/auth/oauth/token/route.ts index cc66068135a..c3e1744dc1f 100644 --- a/apps/sim/app/api/auth/oauth/token/route.ts +++ b/apps/sim/app/api/auth/oauth/token/route.ts @@ -11,17 +11,9 @@ import { authorizeCredentialUse } from '@/lib/auth/credential-access' import { AuthType, checkSessionOrInternalAuth } from '@/lib/auth/hybrid' import { generateRequestId } from '@/lib/core/utils/request' import { withRouteHandler } from '@/lib/core/utils/with-route-handler' -import { TokenServiceAccountValidationError } from '@/lib/credentials/token-service-accounts/errors' -import { - getCredential, - getOAuthToken, - refreshTokenIfNeeded, - resolveOAuthAccountId, - resolveServiceAccountToken, -} from '@/lib/oauth/credential-service' -import { extractSalesforceInstanceUrl, isSalesforceOAuthProviderId } from '@/lib/oauth/salesforce' +import { getCredential, getOAuthToken } from '@/lib/oauth/credential-service' +import { completeOAuthCredentialToken, resolveCredentialToken } from '@/lib/oauth/token-resolution' import { captureServerEvent } from '@/lib/posthog/server' -import { extractZohoDeskBaseFromScope } from '@/tools/zoho_desk/host-allowlist' export const dynamic = 'force-dynamic' @@ -123,194 +115,25 @@ export const POST = withRouteHandler(async (request: NextRequest) => { } } - if (!credentialId) { - return NextResponse.json({ error: 'Credential ID is required' }, { status: 400 }) - } - - const resolved = await resolveOAuthAccountId(credentialId) - if (resolved?.credentialType === 'service_account' && resolved.credentialId) { - const authz = await authorizeCredentialUse(request, { - credentialId, - workflowId: workflowId ?? undefined, - requireWorkflowIdForInternal: false, - callerUserId, - }) - if (!authz.ok) { - return NextResponse.json({ error: authz.error || 'Unauthorized' }, { status: 403 }) - } - - const saActorId = authz.requesterUserId - const saWorkspaceId = resolved.workspaceId ?? authz.workspaceId ?? null - const emitServiceAccountAccess = () => { - if (!saActorId) return - recordAudit({ - workspaceId: saWorkspaceId, - actorId: saActorId, - action: AuditAction.CREDENTIAL_ACCESSED, - resourceType: AuditResourceType.CREDENTIAL, - resourceId: resolved.credentialId ?? credentialId, - description: `Accessed service account credential for provider ${resolved.providerId ?? 'unknown'}`, - metadata: { - provider: resolved.providerId, - credentialType: 'service_account', - }, - request, - }) - captureServerEvent( - saActorId, - 'credential_used', - { - credential_type: 'service_account', - provider_id: resolved.providerId ?? 'unknown', - ...(saWorkspaceId ? { workspace_id: saWorkspaceId } : {}), - }, - saWorkspaceId ? { groups: { workspace: saWorkspaceId } } : undefined - ) - } - - try { - const result = await resolveServiceAccountToken( - resolved.credentialId, - resolved.providerId, - scopes ?? [], - impersonateEmail - ) - emitServiceAccountAccess() - return NextResponse.json( - { - accessToken: result.accessToken, - cloudId: result.cloudId, - domain: result.domain, - instanceUrl: result.instanceUrl, - apiDomain: result.apiDomain, - authStyle: result.authStyle, - }, - { status: 200 } - ) - } catch (error) { - logger.error(`[${requestId}] Service account token error:`, error) - if (error instanceof TokenServiceAccountValidationError) { - // Classified provider outages are infra failures, not bad credentials. - if (error.code === 'provider_unavailable') { - return NextResponse.json( - { error: 'Credential provider is temporarily unavailable' }, - { status: 502 } - ) - } - // A stored host that no longer resolves is a configuration failure — - // surface the code so runtime consumers can say "check the host" - // instead of a generic auth error. - if (error.code === 'site_not_found') { - return NextResponse.json( - { - code: error.code, - error: 'Credential host not found — reconnect the credential with a valid host', - }, - { status: 400 } - ) - } - // A revoked/rotated-away or misconfigured stored secret — surface the - // code so runtime consumers can prompt to reconnect the credential - // rather than showing a generic auth failure. - if (error.code === 'invalid_credentials') { - return NextResponse.json( - { - code: error.code, - error: 'Credential rejected by the provider — reconnect the credential', - }, - { status: 401 } - ) - } - } - return NextResponse.json({ error: 'Failed to get service account token' }, { status: 401 }) - } - } - - const authz = await authorizeCredentialUse(request, { + const auth = await checkSessionOrInternalAuth(request, { requireWorkflowId: false }) + const result = await resolveCredentialToken(auth, { + requestId, credentialId, workflowId: workflowId ?? undefined, - requireWorkflowIdForInternal: false, + scopes, + impersonateEmail, callerUserId, + auditRequest: request, }) - if (!authz.ok || !authz.credentialOwnerUserId) { - return NextResponse.json({ error: authz.error || 'Unauthorized' }, { status: 403 }) - } - - const resolvedCredentialId = authz.resolvedCredentialId || credentialId - const credential = await getCredential( - requestId, - resolvedCredentialId, - authz.credentialOwnerUserId - ) - - if (!credential) { - return NextResponse.json({ error: 'Credential not found' }, { status: 404 }) - } - - const oauthActorId = authz.requesterUserId - const oauthWorkspaceId = authz.workspaceId ?? null - - try { - const { accessToken } = await refreshTokenIfNeeded( - requestId, - credential, - resolvedCredentialId - ) - - if (oauthActorId) { - recordAudit({ - workspaceId: oauthWorkspaceId, - actorId: oauthActorId, - action: AuditAction.CREDENTIAL_ACCESSED, - resourceType: AuditResourceType.CREDENTIAL, - resourceId: resolvedCredentialId, - description: `Accessed OAuth credential for provider ${credential.providerId}`, - metadata: { - provider: credential.providerId, - credentialType: 'oauth', - }, - request, - }) - captureServerEvent( - oauthActorId, - 'credential_used', - { - credential_type: 'oauth', - provider_id: credential.providerId, - ...(oauthWorkspaceId ? { workspace_id: oauthWorkspaceId } : {}), - }, - oauthWorkspaceId ? { groups: { workspace: oauthWorkspaceId } } : undefined - ) - } - - const instanceUrl = isSalesforceOAuthProviderId(credential.providerId) - ? extractSalesforceInstanceUrl(credential.scope) - : undefined - - // Zoho Desk persists its data-center-specific REST base URL in the scope - // string (derived from the token response api_domain) so callers never - // assume a host. Surface it as apiDomain for tool param injection. - let apiDomain: string | undefined - if (credential.providerId === 'zoho-desk' && credential.scope) { - // Use the shared extractor, not a local regex: it also enforces https + - // the Zoho apex allowlist. This value is injected into EVERY tool call, - // so an unvalidated host here would receive the OAuth token. - apiDomain = extractZohoDeskBaseFromScope(credential.scope) - } + if (!result.ok) { return NextResponse.json( - { - accessToken, - idToken: credential.idToken || undefined, - ...(instanceUrl && { instanceUrl }), - ...(apiDomain && { apiDomain }), - }, - { status: 200 } + { ...(result.code ? { code: result.code } : {}), error: result.error }, + { status: result.status } ) - } catch (error) { - logger.error(`[${requestId}] Failed to refresh access token:`, error) - return NextResponse.json({ error: 'Failed to refresh access token' }, { status: 401 }) } + + return NextResponse.json(result.token, { status: 200 }) } catch (error) { logger.error(`[${requestId}] Error getting access token`, error) return NextResponse.json({ error: 'Internal server error' }, { status: 500 }) @@ -366,70 +189,20 @@ export const GET = withRouteHandler(async (request: NextRequest) => { return NextResponse.json({ error: 'No access token available' }, { status: 400 }) } - const actorId = authz.requesterUserId - const workspaceId = authz.workspaceId ?? null - - try { - const { accessToken } = await refreshTokenIfNeeded( - requestId, - credential, - resolvedCredentialId - ) - - if (actorId) { - recordAudit({ - workspaceId, - actorId, - action: AuditAction.CREDENTIAL_ACCESSED, - resourceType: AuditResourceType.CREDENTIAL, - resourceId: resolvedCredentialId, - description: `Accessed OAuth credential for provider ${credential.providerId}`, - metadata: { - provider: credential.providerId, - credentialType: 'oauth', - }, - request, - }) - captureServerEvent( - actorId, - 'credential_used', - { - credential_type: 'oauth', - provider_id: credential.providerId, - ...(workspaceId ? { workspace_id: workspaceId } : {}), - }, - workspaceId ? { groups: { workspace: workspaceId } } : undefined - ) - } - - const instanceUrl = isSalesforceOAuthProviderId(credential.providerId) - ? extractSalesforceInstanceUrl(credential.scope) - : undefined - - // Zoho Desk persists its data-center-specific REST base URL in the scope - // string (derived from the token response api_domain) so callers never - // assume a host. Surface it as apiDomain for tool param injection. - let apiDomain: string | undefined - if (credential.providerId === 'zoho-desk' && credential.scope) { - // Use the shared extractor, not a local regex: it also enforces https + - // the Zoho apex allowlist. This value is injected into EVERY tool call, - // so an unvalidated host here would receive the OAuth token. - apiDomain = extractZohoDeskBaseFromScope(credential.scope) - } + const result = await completeOAuthCredentialToken({ + requestId, + credential, + resolvedCredentialId, + actorId: authz.requesterUserId, + workspaceId: authz.workspaceId ?? null, + auditRequest: request, + }) - return NextResponse.json( - { - accessToken, - idToken: credential.idToken || undefined, - ...(instanceUrl && { instanceUrl }), - ...(apiDomain && { apiDomain }), - }, - { status: 200 } - ) - } catch (error) { - logger.error(`[${requestId}] Failed to refresh access token:`, error) - return NextResponse.json({ error: 'Failed to refresh access token' }, { status: 401 }) + if (!result.ok) { + return NextResponse.json({ error: result.error }, { status: result.status }) } + + return NextResponse.json(result.token, { status: 200 }) } catch (error) { logger.error(`[${requestId}] Error fetching access token`, error) return NextResponse.json({ error: 'Internal server error' }, { status: 500 }) diff --git a/apps/sim/app/api/copilot/checkpoints/revert/route.test.ts b/apps/sim/app/api/copilot/checkpoints/revert/route.test.ts index 785718e9ba2..621256b8ac3 100644 --- a/apps/sim/app/api/copilot/checkpoints/revert/route.test.ts +++ b/apps/sim/app/api/copilot/checkpoints/revert/route.test.ts @@ -17,12 +17,23 @@ import { import { NextRequest } from 'next/server' import { afterAll, afterEach, beforeEach, describe, expect, it, vi } from 'vitest' -const { mockGetAccessibleCopilotChat } = vi.hoisted(() => ({ +const { + mockGetAccessibleCopilotChat, + mockParseWorkflowStateForPersistence, + mockSaveWorkflowNormalizedState, +} = vi.hoisted(() => ({ mockGetAccessibleCopilotChat: vi.fn(), + mockParseWorkflowStateForPersistence: vi.fn(), + mockSaveWorkflowNormalizedState: vi.fn(), })) vi.mock('@/lib/workflows/utils', () => workflowsUtilsMock) +vi.mock('@/lib/workflows/persistence/save-normalized-state', () => ({ + parseWorkflowStateForPersistence: mockParseWorkflowStateForPersistence, + saveWorkflowNormalizedState: mockSaveWorkflowNormalizedState, +})) + vi.mock('@/lib/copilot/chat/lifecycle', () => ({ getAccessibleCopilotChat: mockGetAccessibleCopilotChat, getAccessibleCopilotChatAuth: mockGetAccessibleCopilotChat, @@ -38,13 +49,21 @@ describe('Copilot Checkpoints Revert API Route', () => { authMockFns.mockGetSession.mockResolvedValue(null) + /** Authorization is the route's workflow read, so an allowed result always carries one. */ workflowAuthzMockFns.mockAuthorizeWorkflowByWorkspacePermission.mockResolvedValue({ allowed: true, status: 200, + workflow: { id: 'b2c3d4e5-f6a7-4b89-a0d1-e2f3a4b5c6d7', workspaceId: 'ws-123' }, }) mockGetAccessibleCopilotChat.mockResolvedValue({ id: 'chat-123', userId: 'user-123' }) + mockParseWorkflowStateForPersistence.mockImplementation((value: unknown) => ({ + success: true, + data: value, + })) + mockSaveWorkflowNormalizedState.mockResolvedValue({ success: true, warnings: [] }) + global.fetch = vi.fn() vi.spyOn(Date, 'now').mockReturnValue(1640995200000) @@ -184,7 +203,12 @@ describe('Copilot Checkpoints Revert API Route', () => { } queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint]) - queueTableRows(schemaMock.workflow, []) + /** Authorization performs the workflow read, so a missing workflow surfaces through it. */ + workflowAuthzMockFns.mockAuthorizeWorkflowByWorkspacePermission.mockResolvedValueOnce({ + allowed: false, + status: 404, + workflow: null, + }) const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', { method: 'POST', @@ -197,6 +221,7 @@ describe('Copilot Checkpoints Revert API Route', () => { expect(response.status).toBe(404) const responseData = await response.json() expect(responseData.error).toBe('Workflow not found') + expect(mockSaveWorkflowNormalizedState).not.toHaveBeenCalled() }) it('should return 401 when workflow belongs to different user', async () => { @@ -220,6 +245,7 @@ describe('Copilot Checkpoints Revert API Route', () => { workflowAuthzMockFns.mockAuthorizeWorkflowByWorkspacePermission.mockResolvedValueOnce({ allowed: false, status: 403, + workflow: { id: 'b2c3d4e5-f6a7-4b89-a0d1-e2f3a4b5c6d7', workspaceId: 'ws-123' }, }) const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', { @@ -297,24 +323,19 @@ describe('Copilot Checkpoints Revert API Route', () => { }, }) - // Verify fetch was called with correct parameters - expect(global.fetch).toHaveBeenCalledWith( - 'http://localhost:3000/api/workflows/c3d4e5f6-a7b8-4c09-a1e2-f3a4b5c6d7e8/state', - { - method: 'PUT', - headers: { - 'Content-Type': 'application/json', - Cookie: 'session=test-session', - }, - body: JSON.stringify({ + expect(mockSaveWorkflowNormalizedState).toHaveBeenCalledWith( + expect.objectContaining({ + workflowId: 'c3d4e5f6-a7b8-4c09-a1e2-f3a4b5c6d7e8', + userId: 'user-123', + state: { blocks: { block1: { type: 'start' } }, edges: [{ from: 'block1', to: 'block2' }], loops: {}, parallels: {}, isDeployed: true, lastSaved: 1640995200000, - }), - } + }, + }) ) }) @@ -452,7 +473,7 @@ describe('Copilot Checkpoints Revert API Route', () => { }) }) - it('should return 500 when state API call fails', async () => { + it('should return 500 when the state write fails', async () => { setAuthenticated() const mockCheckpoint = { @@ -470,9 +491,10 @@ describe('Copilot Checkpoints Revert API Route', () => { queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint]) queueTableRows(schemaMock.workflow, [mockWorkflow]) - ;(global.fetch as any).mockResolvedValue({ - ok: false, - text: () => Promise.resolve('State validation failed'), + mockSaveWorkflowNormalizedState.mockResolvedValueOnce({ + success: false, + status: 500, + error: 'Failed to save workflow state', }) const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', { @@ -488,6 +510,36 @@ describe('Copilot Checkpoints Revert API Route', () => { expect(responseData.error).toBe('Failed to revert workflow to checkpoint') }) + it('should return 500 when the checkpoint state fails validation', async () => { + setAuthenticated() + + const mockCheckpoint = { + id: 'checkpoint-123', + workflowId: 'a7b8c9d0-e1f2-4a34-b5c6-d7e8f9a0b1c2', + userId: 'user-123', + workflowState: { blocks: {}, edges: [] }, + } + + queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint]) + queueTableRows(schemaMock.workflow, [{ id: mockCheckpoint.workflowId, userId: 'user-123' }]) + + mockParseWorkflowStateForPersistence.mockReturnValueOnce({ + success: false, + error: { issues: [{ message: 'blocks: invalid' }] }, + }) + + const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ checkpointId: 'checkpoint-123' }), + }) + + const response = await POST(req) + + expect(response.status).toBe(500) + expect(mockSaveWorkflowNormalizedState).not.toHaveBeenCalled() + }) + it('should handle database errors during checkpoint lookup', async () => { setAuthenticated() @@ -519,8 +571,9 @@ describe('Copilot Checkpoints Revert API Route', () => { } dbChainMockFns.where.mockReturnValueOnce(Promise.resolve([mockCheckpoint])) - dbChainMockFns.where.mockReturnValueOnce( - Promise.reject(new Error('Database error during workflow lookup')) + /** Authorization performs the workflow read, so a failed lookup surfaces through it. */ + workflowAuthzMockFns.mockAuthorizeWorkflowByWorkspacePermission.mockRejectedValueOnce( + new Error('Database error during workflow lookup') ) const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', { @@ -536,7 +589,7 @@ describe('Copilot Checkpoints Revert API Route', () => { expect(responseData.error).toBe('Failed to revert to checkpoint') }) - it('should handle fetch network errors', async () => { + it('should handle unexpected errors from the state write', async () => { setAuthenticated() const mockCheckpoint = { @@ -554,7 +607,7 @@ describe('Copilot Checkpoints Revert API Route', () => { queueTableRows(schemaMock.workflowCheckpoints, [mockCheckpoint]) queueTableRows(schemaMock.workflow, [mockWorkflow]) - ;(global.fetch as any).mockRejectedValue(new Error('Network error')) + mockSaveWorkflowNormalizedState.mockRejectedValueOnce(new Error('Network error')) const req = new NextRequest('http://localhost:3000/api/copilot/checkpoints/revert', { method: 'POST', @@ -587,7 +640,7 @@ describe('Copilot Checkpoints Revert API Route', () => { expect(responseData.error).toBe('Failed to revert to checkpoint') }) - it('should forward cookies to state API call', async () => { + it('should apply the state in-process instead of re-authenticating over HTTP', async () => { setAuthenticated() const mockCheckpoint = { @@ -623,17 +676,15 @@ describe('Copilot Checkpoints Revert API Route', () => { await POST(req) - expect(global.fetch).toHaveBeenCalledWith( - 'http://localhost:3000/api/workflows/d0e1f2a3-b4c5-4d67-a8f9-a0b1c2d3e4f5/state', - { - method: 'PUT', - headers: { - 'Content-Type': 'application/json', - Cookie: 'session=test-session; auth=token123', - }, - body: expect.any(String), - } + expect(mockSaveWorkflowNormalizedState).toHaveBeenCalledWith( + expect.objectContaining({ + workflowId: 'd0e1f2a3-b4c5-4d67-a8f9-a0b1c2d3e4f5', + userId: 'user-123', + }) ) + for (const call of (global.fetch as any).mock.calls) { + expect(String(call[0])).not.toContain('/state') + } }) it('should handle missing cookies gracefully', async () => { @@ -673,16 +724,11 @@ describe('Copilot Checkpoints Revert API Route', () => { const response = await POST(req) expect(response.status).toBe(200) - expect(global.fetch).toHaveBeenCalledWith( - 'http://localhost:3000/api/workflows/e1f2a3b4-c5d6-4e78-a9a0-b1c2d3e4f5a6/state', - { - method: 'PUT', - headers: { - 'Content-Type': 'application/json', - Cookie: '', // Empty string when no cookies - }, - body: expect.any(String), - } + expect(mockSaveWorkflowNormalizedState).toHaveBeenCalledWith( + expect.objectContaining({ + workflowId: 'e1f2a3b4-c5d6-4e78-a9a0-b1c2d3e4f5a6', + userId: 'user-123', + }) ) }) diff --git a/apps/sim/app/api/copilot/checkpoints/revert/route.ts b/apps/sim/app/api/copilot/checkpoints/revert/route.ts index f784dc48d84..1543372f773 100644 --- a/apps/sim/app/api/copilot/checkpoints/revert/route.ts +++ b/apps/sim/app/api/copilot/checkpoints/revert/route.ts @@ -1,5 +1,5 @@ import { db } from '@sim/db' -import { workflowCheckpoints, workflow as workflowTable } from '@sim/db/schema' +import { workflowCheckpoints } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { authorizeWorkflowByWorkspacePermission } from '@sim/platform-authz/workflow' import { and, eq } from 'drizzle-orm' @@ -15,8 +15,11 @@ import { createRequestTracker, createUnauthorizedResponse, } from '@/lib/copilot/request/http' -import { getInternalApiBaseUrl } from '@/lib/core/utils/urls' import { withRouteHandler } from '@/lib/core/utils/with-route-handler' +import { + parseWorkflowStateForPersistence, + saveWorkflowNormalizedState, +} from '@/lib/workflows/persistence/save-normalized-state' import { isUuidV4 } from '@/executor/constants' const logger = createLogger('CheckpointRevertAPI') @@ -62,21 +65,15 @@ export const POST = withRouteHandler(async (request: NextRequest) => { return createNotFoundResponse('Checkpoint not found or access denied') } - const workflowData = await db - .select() - .from(workflowTable) - .where(eq(workflowTable.id, checkpoint.workflowId)) - .then((rows) => rows[0]) - - if (!workflowData) { - return createNotFoundResponse('Workflow not found') - } - + /** Authorization already loads the workflow, so its absence is the not-found signal. */ const authorization = await authorizeWorkflowByWorkspacePermission({ workflowId: checkpoint.workflowId, userId, action: 'write', }) + if (!authorization.workflow) { + return createNotFoundResponse('Workflow not found') + } if (!authorization.allowed) { return createUnauthorizedResponse() } @@ -121,28 +118,40 @@ export const POST = withRouteHandler(async (request: NextRequest) => { return NextResponse.json({ error: 'Invalid workflow ID format' }, { status: 400 }) } - const stateResponse = await fetch( - `${getInternalApiBaseUrl()}/api/workflows/${checkpoint.workflowId}/state`, - { - method: 'PUT', - headers: { - 'Content-Type': 'application/json', - Cookie: request.headers.get('Cookie') || '', - }, - body: JSON.stringify(cleanedState), - } - ) + /** + * The checkpoint blob is persisted JSONB, so it goes through the same + * schema the PUT state contract applies before it is written back — the + * validation the removed HTTP hop used to provide. + */ + const parsedState = parseWorkflowStateForPersistence(cleanedState) + if (!parsedState.success) { + logger.error( + `[${tracker.requestId}] Checkpoint state failed validation`, + parsedState.error.issues + ) + return NextResponse.json( + { error: 'Failed to revert workflow to checkpoint' }, + { status: 500 } + ) + } + + const saveResult = await saveWorkflowNormalizedState({ + requestId: tracker.requestId, + workflowId: checkpoint.workflowId, + userId, + state: parsedState.data, + /** Already resolved above; re-deriving it would repeat 2-3 sequential reads. */ + authorization, + }) - if (!stateResponse.ok) { - const errorData = await stateResponse.text() - logger.error(`[${tracker.requestId}] Failed to apply checkpoint state: ${errorData}`) + if (!saveResult.success) { + logger.error(`[${tracker.requestId}] Failed to apply checkpoint state: ${saveResult.error}`) return NextResponse.json( { error: 'Failed to revert workflow to checkpoint' }, { status: 500 } ) } - const result = await stateResponse.json() logger.info( `[${tracker.requestId}] Successfully reverted workflow ${checkpoint.workflowId} to checkpoint ${checkpointId}` ) diff --git a/apps/sim/app/api/workflows/[id]/state/route.ts b/apps/sim/app/api/workflows/[id]/state/route.ts index 209cfd226e8..27ef1e72c2d 100644 --- a/apps/sim/app/api/workflows/[id]/state/route.ts +++ b/apps/sim/app/api/workflows/[id]/state/route.ts @@ -1,28 +1,17 @@ import { db } from '@sim/db' import { workflow } from '@sim/db/schema' import { createLogger } from '@sim/logger' -import { - assertWorkflowMutable, - authorizeWorkflowByWorkspacePermission, - WorkflowLockedError, -} from '@sim/platform-authz/workflow' +import { authorizeWorkflowByWorkspacePermission } from '@sim/platform-authz/workflow' import { toError } from '@sim/utils/errors' import { eq, sql } from 'drizzle-orm' import { type NextRequest, NextResponse } from 'next/server' import { putWorkflowNormalizedStateContract } from '@/lib/api/contracts/workflows' import { parseRequest } from '@/lib/api/server' import { checkSessionOrInternalAuth } from '@/lib/auth/hybrid' -import { env } from '@/lib/core/config/env' import { generateRequestId } from '@/lib/core/utils/request' -import { getSocketServerUrl } from '@/lib/core/utils/urls' import { withRouteHandler } from '@/lib/core/utils/with-route-handler' -import { extractAndPersistCustomTools } from '@/lib/workflows/persistence/custom-tools-persistence' -import { prepareWorkflowStateForPersistence } from '@/lib/workflows/persistence/prepare-state' -import { - loadWorkflowFromNormalizedTables, - saveWorkflowToNormalizedTables, -} from '@/lib/workflows/persistence/utils' -import type { BlockState, WorkflowState } from '@/stores/workflows/workflow/types' +import { saveWorkflowNormalizedState } from '@/lib/workflows/persistence/save-normalized-state' +import { loadWorkflowFromNormalizedTables } from '@/lib/workflows/persistence/utils' const logger = createLogger('WorkflowStateAPI') @@ -117,164 +106,29 @@ export const PUT = withRouteHandler( const parsed = await parseRequest(putWorkflowNormalizedStateContract, request, context) if (!parsed.success) return parsed.response - const state = parsed.data.body - const authorization = await authorizeWorkflowByWorkspacePermission({ + const result = await saveWorkflowNormalizedState({ + requestId, workflowId, userId, - action: 'write', + state: parsed.data.body, }) - const workflowData = authorization.workflow - - if (!workflowData) { - logger.warn(`[${requestId}] Workflow ${workflowId} not found for state update`) - return NextResponse.json({ error: 'Workflow not found' }, { status: 404 }) - } - const canUpdate = authorization.allowed - - if (!canUpdate) { - logger.warn( - `[${requestId}] User ${userId} denied permission to update workflow state ${workflowId}` - ) + if (!result.success) { return NextResponse.json( - { error: authorization.message || 'Access denied' }, - { status: authorization.status || 403 } - ) - } - - await assertWorkflowMutable(workflowId) - - // Note: prior versions cross-checked that each variable's `workflowId` - // equalled the path param. The write contract does not carry `workflowId` - // per variable (the path param is the source of truth), so the check - // is unreachable and was removed. - - const { state: preparedState, warnings: preparationWarnings } = - prepareWorkflowStateForPersistence({ - blocks: state.blocks as Record, - edges: state.edges as WorkflowState['edges'], - }) - - const workflowState = { - ...preparedState, - lastSaved: state.lastSaved || Date.now(), - isDeployed: state.isDeployed || false, - deployedAt: state.deployedAt, - } - - const saveResult = await db.transaction(async (tx) => { - await tx - .select({ id: workflow.id }) - .from(workflow) - .where(eq(workflow.id, workflowId)) - .limit(1) - .for('update') - - const result = await saveWorkflowToNormalizedTables( - workflowId, - workflowState as WorkflowState, - tx - ) - - if (!result.success) return result - - // Update workflow's lastSynced timestamp and variables if provided - const updateData: { - lastSynced: Date - updatedAt: Date - variables?: typeof state.variables - } = { - lastSynced: new Date(), - updatedAt: new Date(), - } - - // If variables are provided in the state, update them in the workflow record - if (state.variables !== undefined) { - updateData.variables = state.variables - } - - await tx.update(workflow).set(updateData).where(eq(workflow.id, workflowId)) - - return result - }) - - if (!saveResult.success) { - logger.error( - `[${requestId}] Failed to save workflow ${workflowId} state:`, - saveResult.error - ) - return NextResponse.json( - { error: 'Failed to save workflow state', details: saveResult.error }, - { status: 500 } + { + error: result.error, + ...(result.details !== undefined ? { details: result.details } : {}), + }, + { status: result.status } ) } - // Extract and persist custom tools to database - try { - const workspaceId = workflowData.workspaceId - if (workspaceId) { - const { saved, errors } = await extractAndPersistCustomTools( - workflowState, - workspaceId, - userId - ) - - if (saved > 0) { - logger.info(`[${requestId}] Persisted ${saved} custom tool(s) to database`, { - workflowId, - }) - } - - if (errors.length > 0) { - logger.warn(`[${requestId}] Some custom tools failed to persist`, { - errors, - workflowId, - }) - } - } else { - logger.warn( - `[${requestId}] Workflow has no workspaceId, skipping custom tools persistence`, - { - workflowId, - } - ) - } - } catch (error) { - logger.error(`[${requestId}] Failed to persist custom tools`, { error, workflowId }) - } - const elapsed = Date.now() - startTime logger.info(`[${requestId}] Successfully saved workflow ${workflowId} state in ${elapsed}ms`) - try { - const notifyResponse = await fetch(`${getSocketServerUrl()}/api/workflow-updated`, { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - 'x-api-key': env.INTERNAL_API_SECRET, - }, - body: JSON.stringify({ workflowId }), - }) - - if (!notifyResponse.ok) { - logger.warn( - `[${requestId}] Failed to notify Socket.IO server about workflow ${workflowId} update` - ) - } - } catch (notificationError) { - logger.warn( - `[${requestId}] Error notifying Socket.IO server about workflow ${workflowId} update`, - notificationError - ) - } - - return NextResponse.json({ success: true, warnings: preparationWarnings }, { status: 200 }) + return NextResponse.json({ success: true, warnings: result.warnings }, { status: 200 }) } catch (error: any) { - if (error instanceof WorkflowLockedError) { - return NextResponse.json({ error: error.message }, { status: error.status }) - } - const elapsed = Date.now() - startTime logger.error( `[${requestId}] Error saving workflow ${workflowId} state after ${elapsed}ms`, diff --git a/apps/sim/app/invite/[id]/invite.tsx b/apps/sim/app/invite/[id]/invite.tsx index 232dfbdb884..ffc75e40866 100644 --- a/apps/sim/app/invite/[id]/invite.tsx +++ b/apps/sim/app/invite/[id]/invite.tsx @@ -15,7 +15,7 @@ import { InviteLayout, InviteStatusCard } from '@/app/invite/components' import { useInvitationDetails } from '@/hooks/queries/invitations' import { organizationKeys } from '@/hooks/queries/organization' import { refreshSessionQuery } from '@/hooks/queries/session' -import { subscriptionKeys } from '@/hooks/queries/subscription' +import { subscriptionKeys } from '@/hooks/queries/utils/subscription-keys' import { workspaceKeys } from '@/hooks/queries/workspace' const logger = createLogger('InviteById') diff --git a/apps/sim/app/workspace/[workspaceId]/files/prefetch.ts b/apps/sim/app/workspace/[workspaceId]/files/prefetch.ts index 6ed88de422d..54d0276fcef 100644 --- a/apps/sim/app/workspace/[workspaceId]/files/prefetch.ts +++ b/apps/sim/app/workspace/[workspaceId]/files/prefetch.ts @@ -1,4 +1,5 @@ import type { QueryClient } from '@tanstack/react-query' +import { listWorkspaceFileFoldersContract } from '@/lib/api/contracts/workspace-file-folders' import { listWorkspaceFileFolders } from '@/lib/uploads/contexts/workspace/workspace-file-folder-manager' import { getWorkspaceHostContextForViewer } from '@/lib/workspaces/host-context' import { prefetchResourceListChrome } from '@/app/workspace/[workspaceId]/lib/prefetch-resource-list-chrome' @@ -15,10 +16,13 @@ import { * (scope `active`), so the browser paints populated on first render. * * The FILE LIST itself is deliberately not here: the sidebar reads it on every workspace route, so - * it is prefetched by `prefetchWorkspaceSidebar` in the layout — the only boundary that renders + * it is seeded by `prefetchWorkspaceSidebar` in the layout — the only boundary that renders * before the sidebar registers the query. Prefetching it again here would re-read it per request * and still not reach the server render (`HydrationBoundary` defers an already-seen query to an - * effect, which SSR never runs). See the note on that entry. + * effect, which SSR never runs). See the note on that entry. The layout declines to seed a + * workspace whose file list exceeds its payload budget; recovering those here would mean + * mirroring that budget check inversely, since an unconditional prefetch would re-read and + * duplicate the entry for every workspace under the budget. * * Folders and the chrome reads all go through the data layer, shaped to their route contracts so a * hydrated entry matches a client fetch. @@ -40,7 +44,17 @@ export async function prefetchFilesBrowser( await Promise.all([ queryClient.prefetchQuery({ queryKey: workspaceFileFolderKeys.list(workspaceId, 'active'), - queryFn: () => listWorkspaceFileFolders(workspaceId, { scope: 'active' }), + /** + * Parsed through the route's own response schema rather than seeded raw. The + * manager's record type and `workspaceFileFolderSchema` are two independent + * declarations that happen to agree today; without this parse, adding a column + * to one silently seeds a shape a client fetch would have stripped — the exact + * divergence that put ISO strings under `workspaceFilesKeys.list`. + */ + queryFn: async () => { + const folders = await listWorkspaceFileFolders(workspaceId, { scope: 'active' }) + return listWorkspaceFileFoldersContract.response.schema.shape.folders.parse(folders) + }, staleTime: WORKSPACE_FILE_FOLDERS_STALE_TIME, }), prefetchResourceListChrome(queryClient, workspaceId, 'file', userId), diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-tool-event.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-tool-event.ts index e00e851626d..296f48bc559 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-tool-event.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-tool-event.ts @@ -27,7 +27,7 @@ import { } from '@/app/workspace/[workspaceId]/home/hooks/stream/turn-model' import { deploymentKeys } from '@/hooks/queries/deployments' import { folderKeys } from '@/hooks/queries/utils/folder-keys' -import { workflowKeys } from '@/hooks/queries/workflows' +import { invalidateWorkflowLists } from '@/hooks/queries/utils/invalidate-workflow-lists' type ToolEvent = Extract @@ -58,7 +58,7 @@ function runToolResultSideEffects(ctx: StreamLoopContext, node: ToolNode): void if (deployedWorkflowId && typeof out?.isDeployed === 'boolean') { deps.queryClient.invalidateQueries({ queryKey: deploymentKeys.info(deployedWorkflowId) }) deps.queryClient.invalidateQueries({ queryKey: deploymentKeys.versions(deployedWorkflowId) }) - deps.queryClient.invalidateQueries({ queryKey: workflowKeys.list(deps.workspaceId) }) + void invalidateWorkflowLists(deps.queryClient, deps.workspaceId) } } @@ -66,7 +66,9 @@ function runToolResultSideEffects(ctx: StreamLoopContext, node: ToolNode): void deps.queryClient.invalidateQueries({ queryKey: folderKeys.list(deps.workspaceId) }) } if (WORKFLOW_MUTATION_TOOL_NAMES.has(name) && isSuccess) { - deps.queryClient.invalidateQueries({ queryKey: workflowKeys.list(deps.workspaceId) }) + // `rm` archives, so the archived list moves too — and the shared helper also + // refreshes the workflow selector lists that `@`-mentions and pickers read. + void invalidateWorkflowLists(deps.queryClient, deps.workspaceId, ['active', 'archived']) } const extractedResources = diff --git a/apps/sim/app/workspace/[workspaceId]/knowledge/hooks/use-knowledge-upload.ts b/apps/sim/app/workspace/[workspaceId]/knowledge/hooks/use-knowledge-upload.ts index c9cd91416ef..5923e9e03f6 100644 --- a/apps/sim/app/workspace/[workspaceId]/knowledge/hooks/use-knowledge-upload.ts +++ b/apps/sim/app/workspace/[workspaceId]/knowledge/hooks/use-knowledge-upload.ts @@ -209,7 +209,11 @@ export function useKnowledgeUpload(options: UseKnowledgeUploadOptions = {}) { setUploadProgress((prev) => ({ ...prev, stage: 'processing' })) logger.info(`Successfully started processing ${uploadedDocuments.length} documents`) - await queryClient.invalidateQueries({ queryKey: knowledgeKeys.detail(knowledgeBaseId) }) + await Promise.all([ + queryClient.invalidateQueries({ queryKey: knowledgeKeys.detail(knowledgeBaseId) }), + /** The knowledge-base list rows carry `docCount`, so an upload changes them too. */ + queryClient.invalidateQueries({ queryKey: knowledgeKeys.lists() }), + ]) return uploadedDocuments } catch (err) { diff --git a/apps/sim/app/workspace/[workspaceId]/lib/prefetch.test.ts b/apps/sim/app/workspace/[workspaceId]/lib/prefetch.test.ts index dee79520936..c27a77959b0 100644 --- a/apps/sim/app/workspace/[workspaceId]/lib/prefetch.test.ts +++ b/apps/sim/app/workspace/[workspaceId]/lib/prefetch.test.ts @@ -97,7 +97,10 @@ vi.mock('@sim/emcn', () => ({ import { prefetchFilesBrowser } from '@/app/workspace/[workspaceId]/files/prefetch' import { prefetchKnowledgeBases } from '@/app/workspace/[workspaceId]/knowledge/prefetch' -import { prefetchWorkspaceSidebar } from '@/app/workspace/[workspaceId]/prefetch' +import { + prefetchWorkspaceSidebar, + WORKSPACE_FILE_SEED_MAX, +} from '@/app/workspace/[workspaceId]/prefetch' import { prefetchTables } from '@/app/workspace/[workspaceId]/tables/prefetch' import { folderKeys } from '@/hooks/queries/utils/folder-keys' import { knowledgeKeys } from '@/hooks/queries/utils/knowledge-keys' @@ -290,17 +293,67 @@ describe('workspace list prefetches', () => { }) }) describe('prefetchFilesBrowser', () => { + /** + * The sibling `workspaceFilesKeys.list` once held ISO strings from one producer and + * `Date`s from another because a seed skipped the contract parse. This key is fed by + * a manager whose record type and the contract schema are independent declarations, + * so the parse — and this assertion — are what stop that recurring here. + */ + it('seeds the shape a client fetch caches, not the raw manager row', async () => { + mockListWorkspaceFileFolders.mockResolvedValue([ + { + id: 'folder-1', + workspaceId: WORKSPACE_ID, + userId: USER_ID, + name: 'Docs', + parentId: null, + path: '/Docs', + sortOrder: 0, + deletedAt: null, + createdAt: '2026-01-01T00:00:00.000Z', + updatedAt: '2026-01-02T00:00:00.000Z', + serverOnlyColumn: 'should-be-stripped', + }, + ]) + const client = makeClient() + + await prefetchFilesBrowser(client, WORKSPACE_ID, USER_ID) + + const [cached] = client.getQueryData( + workspaceFileFolderKeys.list(WORKSPACE_ID, 'active') + ) as Array> + expect(cached.createdAt).toBeInstanceOf(Date) + expect(cached.updatedAt).toBeInstanceOf(Date) + expect(cached).not.toHaveProperty('serverOnlyColumn') + }) + it('primes the folder key the client hook reads', async () => { - const folders = [{ id: 'folder-1' }] + const folders = [ + { + id: 'folder-1', + workspaceId: WORKSPACE_ID, + userId: USER_ID, + name: 'Docs', + parentId: null, + path: '/Docs', + sortOrder: 0, + deletedAt: null, + createdAt: '2026-01-01T00:00:00.000Z', + updatedAt: '2026-01-02T00:00:00.000Z', + }, + ] mockListWorkspaceFileFolders.mockResolvedValue(folders) const client = makeClient() await prefetchFilesBrowser(client, WORKSPACE_ID, USER_ID) expect(mockListWorkspaceFileFolders).toHaveBeenCalledWith(WORKSPACE_ID, { scope: 'active' }) - expect(client.getQueryData(workspaceFileFolderKeys.list(WORKSPACE_ID, 'active'))).toEqual( - folders - ) + /** Shape parity is asserted by the sibling test; this one pins the key and the args. */ + expect( + client.getQueryData(workspaceFileFolderKeys.list(WORKSPACE_ID, 'active')) as Array<{ + id: string + }> + ).toHaveLength(folders.length) }) /** @@ -460,6 +513,52 @@ describe('workspace list prefetches', () => { expect(client.getQueryData(workspaceKeys.list('active'))).toBeUndefined() }) + /** + * The file list is seeded on every workspace route, so it is the one entry whose size + * scales with a workspace's content on routes that never read it. The budget is passed + * down rather than applied here, so the read can stop before the share join. + */ + it('seeds the file list, bounded by the document payload budget', async () => { + const files = [{ id: 'file-1', name: 'a.txt' }] + mockListWorkspaceFilesWithShares.mockResolvedValue(files) + const client = makeClient() + + await prefetchWorkspaceSidebar(client, WORKSPACE_ID, USER_ID, HOST_CONTEXT, null) + + expect(mockListWorkspaceFilesWithShares).toHaveBeenCalledWith(WORKSPACE_ID, 'active', { + maxRows: WORKSPACE_FILE_SEED_MAX, + /** A failed read must reach the catch, not degrade to a cached empty list. */ + throwOnError: true, + }) + expect(client.getQueryData(workspaceFilesKeys.list(WORKSPACE_ID, 'active'))).toEqual(files) + }) + + /** + * The load-bearing half of the budget: a workspace over it seeds NOTHING rather than the + * prefix that was read. The sidebar search filters this list client-side and the Files + * browser renders it as the workspace's files, so a truncated seed would silently hide + * files — the client fetch must reach the route for the complete list instead. + */ + it('seeds nothing when the workspace exceeds the budget', async () => { + mockListWorkspaceFilesWithShares.mockResolvedValue(null) + const client = makeClient() + + await prefetchWorkspaceSidebar(client, WORKSPACE_ID, USER_ID, HOST_CONTEXT, null) + + expect(client.getQueryData(workspaceFilesKeys.list(WORKSPACE_ID, 'active'))).toBeUndefined() + }) + + /** A failed file read is an optimization loss, not a render failure. */ + it('does not throw when the file read rejects, and seeds no files', async () => { + mockListWorkspaceFilesWithShares.mockRejectedValue(new Error('500')) + const client = makeClient() + + await expect( + prefetchWorkspaceSidebar(client, WORKSPACE_ID, USER_ID, HOST_CONTEXT, null) + ).resolves.toBeUndefined() + expect(client.getQueryData(workspaceFilesKeys.list(WORKSPACE_ID, 'active'))).toBeUndefined() + }) + /** Guards the mismatch check that keeps one workspace's data out of another's cache. */ it('seeds nothing when the host context is for a different workspace', async () => { mockListWorkspacesForViewer.mockResolvedValue(LIST_PAYLOAD) diff --git a/apps/sim/app/workspace/[workspaceId]/prefetch.ts b/apps/sim/app/workspace/[workspaceId]/prefetch.ts index c24198fb51f..e19e970b581 100644 --- a/apps/sim/app/workspace/[workspaceId]/prefetch.ts +++ b/apps/sim/app/workspace/[workspaceId]/prefetch.ts @@ -25,10 +25,7 @@ import { workflowKeys } from '@/hooks/queries/utils/workflow-keys' import { mapWorkflow, WORKFLOW_LIST_STALE_TIME } from '@/hooks/queries/utils/workflow-list-query' import { normalizeWorkspacesResponse } from '@/hooks/queries/utils/workspace-list-query' import { WORKSPACE_PERMISSIONS_STALE_TIME, workspaceKeys } from '@/hooks/queries/workspace' -import { - WORKSPACE_FILES_LIST_STALE_TIME, - workspaceFilesKeys, -} from '@/hooks/queries/workspace-files' +import { workspaceFilesKeys } from '@/hooks/queries/workspace-files' import { WORKSPACE_HOST_CONTEXT_STALE_TIME, workspaceHostKeys, @@ -95,6 +92,46 @@ async function seedWorkspaceList( } } +/** + * How many files the layout is willing to inline into the document. Seeded on EVERY + * workspace route, so at ~500 bytes of JSON per file this budgets the entry at ~150 KB. + * + * A workspace above the budget seeds NOTHING rather than a prefix: the sidebar filters + * this list client-side, so a truncated seed would silently hide files. + */ +export const WORKSPACE_FILE_SEED_MAX = 300 + +/** + * Seeds the workspace's file list, which sidebar chrome registers on EVERY workspace + * route. It must be seeded HERE, not by the Files pages: `HydrationBoundary` defers a + * query the cache has already seen to a `useEffect`, which SSR never runs. + * + * Seeded rather than prefetched so it can decline to create an entry at all above + * {@link WORKSPACE_FILE_SEED_MAX} — `prefetchQuery` always creates one, and a partial + * entry would be read as the whole list. Parsed through the route's response contract. + */ +async function seedWorkspaceFiles(queryClient: QueryClient, workspaceId: string): Promise { + try { + const files = await listWorkspaceFilesWithShares(workspaceId, 'active', { + maxRows: WORKSPACE_FILE_SEED_MAX, + /** + * A failed read must reach the catch below, not degrade to an empty list: seeding + * `[]` would cache "this workspace has no files" as authoritative for the entry's + * lifetime, which is worse than seeding nothing and letting the client fetch. + */ + throwOnError: true, + }) + if (!files) return + queryClient.setQueryData(workspaceFilesKeys.list(workspaceId, 'active'), files) + } catch (error) { + /** Optimization only: the client fetch reaches the route instead. Logged so drift between + * this read and the contract's response schema doesn't degrade silently into a waterfall. */ + logger.warn('Workspace file list seed failed; client will fetch', { + error: getErrorMessage(error), + }) + } +} + /** * Prefetches the sidebar's workflow, chat, folder, workspace-permissions, * workspace, and viewer-profile reads for a workspace and stores them under the @@ -153,22 +190,7 @@ export async function prefetchWorkspaceSidebar( ] : []), prefetchResourceFolders(queryClient, workspaceId, 'workflow', userId), - /** - * The sidebar reads the workspace's files for its search modal, on EVERY workspace route — so this - * query is registered by sidebar chrome before any page renders. That ordering is why it has to be - * prefetched HERE and not only by the Files pages: `HydrationBoundary` hydrates a query the cache - * has already seen from a `useEffect`, which never runs during SSR, so a page-level boundary can - * only ever hand this entry to the client. Seeding it with the layout's own boundary — the first - * one to render — is what lets the server paint the Files browser and the open file's header - * populated instead of shipping a spinner and resolving it a beat later on the client. - * - * Same key + shape as {@link prefetchFilesBrowser}, so whichever runs is a no-op for the other. - */ - queryClient.prefetchQuery({ - queryKey: workspaceFilesKeys.list(workspaceId, 'active'), - queryFn: () => listWorkspaceFilesWithShares(workspaceId, 'active'), - staleTime: WORKSPACE_FILES_LIST_STALE_TIME, - }), + seedWorkspaceFiles(queryClient, workspaceId), queryClient.prefetchQuery({ queryKey: workspaceKeys.permissions(workspaceId), queryFn: () => diff --git a/apps/sim/app/workspace/[workspaceId]/upgrade/hooks/use-upgrade-state.ts b/apps/sim/app/workspace/[workspaceId]/upgrade/hooks/use-upgrade-state.ts index 1792360d7d6..15152b4a57b 100644 --- a/apps/sim/app/workspace/[workspaceId]/upgrade/hooks/use-upgrade-state.ts +++ b/apps/sim/app/workspace/[workspaceId]/upgrade/hooks/use-upgrade-state.ts @@ -9,6 +9,8 @@ import type { WorkspaceHostContext } from '@/lib/api/contracts/workspaces' import { useSubscriptionUpgrade } from '@/lib/billing/client/upgrade' import { CREDIT_TIERS } from '@/lib/billing/constants' import { getPlanTierCredits, isEnterprise, isFree, isPro, isTeam } from '@/lib/billing/plan-helpers' +import { invalidateWorkspaceUsage } from '@/hooks/queries/utils/invalidate-usage' +import { subscriptionKeys } from '@/hooks/queries/utils/subscription-keys' import { workspaceHostKeys } from '@/hooks/queries/workspace-host' const PRO_TIER = CREDIT_TIERS[0] @@ -89,8 +91,22 @@ export function useUpgradeState({ } }, [ownerBilling.billingInterval, subscription.isPaid]) - const refreshHostContext = useCallback( - () => queryClient.invalidateQueries({ queryKey: workspaceHostKeys.detail(workspaceId) }), + /** + * A non-redirect plan switch settles server-side immediately, so every read that + * describes the plan has to be refetched — the host context the page renders from, + * the subscription/usage reads the billing surfaces share, the proration invoice the + * switch just produced, and the workspace credit availability that drives the credits + * chip and the run gate. + */ + const refreshBillingState = useCallback( + () => + Promise.all([ + queryClient.invalidateQueries({ queryKey: workspaceHostKeys.detail(workspaceId) }), + queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }), + queryClient.invalidateQueries({ queryKey: subscriptionKeys.usage() }), + queryClient.invalidateQueries({ queryKey: subscriptionKeys.invoicesAll() }), + invalidateWorkspaceUsage(queryClient), + ]), [queryClient, workspaceId] ) @@ -123,9 +139,9 @@ export function useUpgradeState({ await requestJson(billingSwitchPlanContract, { body: { targetPlanName: subscription.plan, interval, workspaceId }, }) - await refreshHostContext() + await refreshBillingState() }, - [isLegacyPlan, refreshHostContext, subscription.plan, workspaceId] + [isLegacyPlan, refreshBillingState, subscription.plan, workspaceId] ) const currentCredits = getPlanTierCredits(subscription.plan) @@ -154,11 +170,11 @@ export function useUpgradeState({ workspaceId, }, }) - await refreshHostContext() + await refreshBillingState() } catch (e) { toast.error(getErrorMessage(e, 'Failed to upgrade')) } - }, [subscription.isTeam, isAnnual, refreshHostContext, workspaceId]) + }, [subscription.isTeam, isAnnual, refreshBillingState, workspaceId]) const onUpgradeToOtherTier = useCallback(async () => { const onMax = @@ -170,11 +186,11 @@ export function useUpgradeState({ await requestJson(billingSwitchPlanContract, { body: { targetPlanName, workspaceId }, }) - await refreshHostContext() + await refreshBillingState() } catch (e) { toast.error(getErrorMessage(e, 'Failed to switch plan')) } - }, [subscription.plan, subscription.isTeam, refreshHostContext, workspaceId]) + }, [subscription.plan, subscription.isTeam, refreshBillingState, workspaceId]) return { isLoading: false, diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts index 0e0a8b795fc..cb6573d6db0 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts +++ b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts @@ -10,7 +10,7 @@ import { wandGenerateStreamContract } from '@/lib/api/contracts' import { readSSEStream } from '@/lib/core/utils/sse' import { shouldStripCodeFences, stripCodeFences } from '@/lib/wand/strip-code-fences' import type { GenerationType } from '@/blocks/types' -import { subscriptionKeys } from '@/hooks/queries/subscription' +import { scheduleUsageRefresh } from '@/hooks/queries/utils/invalidate-usage' import { useSettingsNavigation } from '@/hooks/use-settings-navigation' import { useWorkflowRegistry } from '@/stores/workflows/registry/store' @@ -300,9 +300,7 @@ export function useWand({ strippedFences: generatedContent !== accumulatedContent, }) - setTimeout(() => { - queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }) - }, 1000) + scheduleUsageRefresh(queryClient) } catch (error: any) { if (error.name === 'AbortError') { logger.debug('Wand generation cancelled') diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.test.tsx b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.test.tsx index c9b2fd0a6a4..15c7171064c 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.test.tsx +++ b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.test.tsx @@ -191,10 +191,6 @@ vi.mock('@/executor/utils/start-block', () => ({ coerceValue: (_type: string, value: unknown) => value, })) -vi.mock('@/hooks/queries/subscription', () => ({ - subscriptionKeys: { users: () => ['subscription', 'users'] }, -})) - vi.mock('@/hooks/queries/utils/workflow-cache', () => ({ getWorkflows: () => [], })) diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts index 37d7fb5d5cb..97fb6a3f507 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts +++ b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts @@ -61,7 +61,7 @@ import type { SerializableExecutionState } from '@/executor/execution/types' import type { BlockLog, BlockState, ExecutionResult, StreamingExecution } from '@/executor/types' import { hasExecutionResult } from '@/executor/utils/errors' import { coerceValue } from '@/executor/utils/start-block' -import { subscriptionKeys } from '@/hooks/queries/subscription' +import { scheduleUsageRefresh } from '@/hooks/queries/utils/invalidate-usage' import { getWorkflows } from '@/hooks/queries/utils/workflow-cache' import { isExecutionStreamHttpError, @@ -922,9 +922,7 @@ export function useWorkflowExecution() { } // Invalidate subscription queries to update usage - setTimeout(() => { - queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }) - }, 1000) + scheduleUsageRefresh(queryClient) safeEnqueue(encodeSSE({ event: 'final', data: result })) // Note: Logs are already persisted server-side via execution-core.ts @@ -1458,9 +1456,7 @@ export function useWorkflowExecution() { setIsExecuting(activeWorkflowId, false) setActiveBlocks(activeWorkflowId, new Set()) } - setTimeout(() => { - queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }) - }, 1000) + scheduleUsageRefresh(queryClient) } }, diff --git a/apps/sim/blocks/blocks/credential.ts b/apps/sim/blocks/blocks/credential.ts index 2ea6ccaec85..d1e7baab8f4 100644 --- a/apps/sim/blocks/blocks/credential.ts +++ b/apps/sim/blocks/blocks/credential.ts @@ -3,7 +3,10 @@ import { getServiceConfigByProviderId } from '@/lib/oauth/utils' import { getQueryClient } from '@/app/_shell/providers/get-query-client' import type { BlockConfig } from '@/blocks/types' import { workspaceCredentialKeys } from '@/hooks/queries/utils/credential-keys' -import { fetchWorkspaceCredentialList } from '@/hooks/queries/utils/fetch-workspace-credentials' +import { + fetchWorkspaceCredentialList, + WORKSPACE_CREDENTIAL_LIST_STALE_TIME, +} from '@/hooks/queries/utils/fetch-workspace-credentials' import { useWorkflowRegistry } from '@/stores/workflows/registry/store' interface CredentialBlockOutput { @@ -72,7 +75,7 @@ export const CredentialBlock: BlockConfig = { const credentials = await getQueryClient().fetchQuery({ queryKey: workspaceCredentialKeys.list(workspaceId), queryFn: () => fetchWorkspaceCredentialList(workspaceId), - staleTime: 60 * 1000, + staleTime: WORKSPACE_CREDENTIAL_LIST_STALE_TIME, }) const seen = new Set() diff --git a/apps/sim/executor/handlers/evaluator/evaluator-handler.test.ts b/apps/sim/executor/handlers/evaluator/evaluator-handler.test.ts index 756b98798e2..6df5974aefa 100644 --- a/apps/sim/executor/handlers/evaluator/evaluator-handler.test.ts +++ b/apps/sim/executor/handlers/evaluator/evaluator-handler.test.ts @@ -4,10 +4,15 @@ import { createLogger } from '@sim/logger' import { authOAuthUtilsMock, authOAuthUtilsMockFns } from '@sim/testing' import { beforeEach, describe, expect, it, type Mock, vi } from 'vitest' -const { mockResolveAutoModel } = vi.hoisted(() => ({ +const { mockResolveAutoModel, mockCheckWorkspaceAccess } = vi.hoisted(() => ({ + mockCheckWorkspaceAccess: vi.fn(), mockResolveAutoModel: vi.fn(), })) +vi.mock('@/lib/workspaces/permissions/utils', () => ({ + checkWorkspaceAccess: mockCheckWorkspaceAccess, +})) + vi.mock('@/lib/oauth/credential-service', () => authOAuthUtilsMock) vi.mock('@/lib/credentials/access', () => ({ @@ -34,24 +39,30 @@ vi.mock('@/lib/model-router/resolve', () => ({ SIM_AUTO_SYSTEM_PREAMBLE: 'Sim auto system preamble', })) -import { - PRIVATE_MODEL_INPUT_PROVENANCE_HEADER, - PRIVATE_MODEL_INPUT_STATE_HEADER, - PROJECTED_MODEL_INPUT_PATHS_V1, -} from '@/lib/execution/model-input-provenance' -import { - RESOLVED_SECRET_PROVENANCE_FIELD, - RESOLVED_SECRET_PROVENANCE_METADATA_V1, -} from '@/lib/execution/private-tool-metadata' import { BlockType } from '@/executor/constants' import { EvaluatorBlockHandler } from '@/executor/handlers/evaluator/evaluator-handler' import type { ExecutionContext } from '@/executor/types' import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' +import { executeProviderRequest } from '@/providers' +import type { ProviderRequest } from '@/providers/types' import { getProviderFromModel } from '@/providers/utils' import type { SerializedBlock } from '@/serializer/types' const mockGetProviderFromModel = getProviderFromModel as Mock -const mockFetch = vi.fn() +const mockExecuteProviderRequest = executeProviderRequest as Mock + +/** The provider request the handler built, keyed the way the old wire body was. */ +function providerRequestBody(index = 0): ProviderRequest & { provider: string } { + const [provider, request] = mockExecuteProviderRequest.mock.calls[index] as [ + string, + ProviderRequest, + ] + return { provider, ...request } +} + +function providerRuntimeRegistry(index = 0): ResolvedSecretTraceRegistry | undefined { + return mockExecuteProviderRequest.mock.calls[index][2]?.resolvedSecretTraceRegistry +} const mockLogger = vi.mocked(createLogger).mock.results[ vi.mocked(createLogger).mock.calls.findIndex(([name]) => name === 'EvaluatorBlockHandler') @@ -97,8 +108,7 @@ describe('EvaluatorBlockHandler', () => { // Reset mocks using vi vi.clearAllMocks() - // unstubGlobals removes any module-scope fetch stub before each test, so re-stub here - vi.stubGlobal('fetch', mockFetch) + mockCheckWorkspaceAccess.mockResolvedValue({ hasAccess: true }) // Default mock implementations authOAuthUtilsMockFns.mockResolveOAuthAccountId.mockResolvedValue({ @@ -117,19 +127,12 @@ describe('EvaluatorBlockHandler', () => { billableRoutingCost: 0.002, }) - // Set up fetch mock to return a successful response - mockFetch.mockImplementation(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ score1: 5, score2: 8 }), - model: 'mock-model', - tokens: { input: 50, output: 10, total: 60 }, - cost: 0.002, - timing: { total: 200 }, - }), - }) + mockExecuteProviderRequest.mockResolvedValue({ + content: JSON.stringify({ score1: 5, score2: 8 }), + model: 'mock-model', + tokens: { input: 50, output: 10, total: 60 }, + cost: 0.002, + timing: { total: 200 }, }) }) @@ -142,6 +145,37 @@ describe('EvaluatorBlockHandler', () => { expect(handler.canHandle(nonEvalBlock)).toBe(false) }) + /** + * The admission checks the removed `/api/providers` hop owned. Mirrors the router's + * coverage — both handlers reach the provider through the same shared entry point. + */ + const admissionInputs = { + content: 'Evaluate this.', + metrics: [{ name: 'score1', description: 'First score', range: { min: 0, max: 10 } }], + model: 'gpt-4o', + apiKey: 'test-api-key', + } + + it('refuses to reach the provider without an execution subject', async () => { + mockContext.userId = undefined + + await expect(handler.execute(mockContext, mockBlock, admissionInputs)).rejects.toThrow( + 'Unauthorized' + ) + expect(mockExecuteProviderRequest).not.toHaveBeenCalled() + }) + + it('refuses to reach the provider when the subject lost workspace access', async () => { + mockContext.workspaceId = 'test-workspace' + mockCheckWorkspaceAccess.mockResolvedValue({ hasAccess: false }) + + await expect(handler.execute(mockContext, mockBlock, admissionInputs)).rejects.toThrow( + 'Forbidden' + ) + expect(mockCheckWorkspaceAccess).toHaveBeenCalledWith('test-workspace', 'test-user') + expect(mockExecuteProviderRequest).not.toHaveBeenCalled() + }) + it('should execute evaluator block correctly with basic inputs', async () => { const inputs = { content: 'This is the content to evaluate.', @@ -165,17 +199,9 @@ describe('EvaluatorBlockHandler', () => { const result = await handler.execute(mockContext, mockBlock, inputs) expect(mockGetProviderFromModel).toHaveBeenCalledWith('gpt-4o') - expect(mockFetch).toHaveBeenCalledWith( - expect.any(String), - expect.objectContaining({ - method: 'POST', - headers: expect.any(Object), - body: expect.any(String), - }) - ) + expect(mockExecuteProviderRequest).toHaveBeenCalledTimes(1) - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ provider: 'openai', model: 'gpt-4o', @@ -257,15 +283,8 @@ describe('EvaluatorBlockHandler', () => { apiKey: credentialSecret, }) - const request = mockFetch.mock.calls[0][1] - const requestBody = JSON.parse(request.body) - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBe( - RESOLVED_SECRET_PROVENANCE_METADATA_V1 - ) - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBe( - PROJECTED_MODEL_INPUT_PATHS_V1 - ) - expect(requestBody[RESOLVED_SECRET_PROVENANCE_FIELD]).toEqual({ + const requestBody = providerRequestBody() + expect(providerRuntimeRegistry()?.exportProvenance()).toEqual({ version: 1, complete: true, entries: [ @@ -336,15 +355,11 @@ describe('EvaluatorBlockHandler', () => { registry.recordResolvedInputProjection(secret.path, secret.plaintext, secret.projected) } mockContext.resolvedSecretTraceRegistry = registry - mockFetch.mockResolvedValueOnce({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ [projectedMetric.name.toLowerCase()]: 7 }), - model: 'mock-model', - tokens: {}, - cost: 0, - }), + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ [projectedMetric.name.toLowerCase()]: 7 }), + model: 'mock-model', + tokens: {}, + cost: 0, }) const result = await handler.execute(mockContext, mockBlock, { @@ -354,7 +369,7 @@ describe('EvaluatorBlockHandler', () => { apiKey: 'test-api-key', }) - const requestBody = JSON.parse(mockFetch.mock.calls[0][1].body) + const requestBody = providerRequestBody() const serializedRequest = JSON.stringify(requestBody) for (const secret of secrets) { expect(serializedRequest).not.toContain(secret.plaintext) @@ -365,8 +380,9 @@ describe('EvaluatorBlockHandler', () => { [projectedMetric.name.toLowerCase()]: { type: 'number' }, }) expect( - requestBody[RESOLVED_SECRET_PROVENANCE_FIELD].entries - .map((entry: { name: string }) => entry.name) + providerRuntimeRegistry() + ?.exportProvenance() + .entries.map((entry: { name: string }) => entry.name) .sort() ).toEqual(secrets.map((secret) => secret.name).sort()) expect(result).toMatchObject({ [rawMetric.name.toLowerCase()]: 7 }) @@ -380,11 +396,7 @@ describe('EvaluatorBlockHandler', () => { apiKey: 'test-api-key', }) - const request = mockFetch.mock.calls[0][1] - const requestBody = JSON.parse(request.body) - expect(Object.hasOwn(requestBody, RESOLVED_SECRET_PROVENANCE_FIELD)).toBe(false) - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBeNull() - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBeNull() + expect(providerRuntimeRegistry()).toBeUndefined() }) it('resolves sim-auto before executing evaluator and preserves its public identity', async () => { @@ -400,15 +412,11 @@ describe('EvaluatorBlockHandler', () => { model: 'sim-auto', } - mockFetch.mockResolvedValueOnce({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ quality: 5 }), - model: 'fireworks/glm-5.2', - tokens: { input: 80, output: 10, total: 90 }, - cost: { input: 0.001, output: 0.0005, total: 0.0015 }, - }), + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ quality: 5 }), + model: 'fireworks/glm-5.2', + tokens: { input: 80, output: 10, total: 90 }, + cost: { input: 0.001, output: 0.0005, total: 0.0015 }, }) const result = await handler.execute(mockContext, mockBlock, inputs) @@ -427,7 +435,7 @@ describe('EvaluatorBlockHandler', () => { }) expect(mockGetProviderFromModel).toHaveBeenCalledWith('fireworks/glm-5.2') - const requestBody = JSON.parse(mockFetch.mock.calls[0][1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ provider: 'openai', model: 'fireworks/glm-5.2', @@ -448,19 +456,13 @@ describe('EvaluatorBlockHandler', () => { it('bills the cost the provider proxy decided rather than recomputing it', async () => { // The proxy already resolved key provenance and the margin; recomputing // here would re-charge a BYOK caller the proxy correctly zeroed. - mockFetch.mockImplementation(() => - Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ score1: 5, score2: 8 }), - model: 'mock-model', - tokens: { input: 50, output: 10, total: 60 }, - cost: { input: 0.001, output: 0.0005, total: 0.0015 }, - timing: { total: 200 }, - }), - }) - ) + mockExecuteProviderRequest.mockResolvedValue({ + content: JSON.stringify({ score1: 5, score2: 8 }), + model: 'mock-model', + tokens: { input: 50, output: 10, total: 60 }, + cost: { input: 0.001, output: 0.0005, total: 0.0015 }, + timing: { total: 200 }, + }) const result = await handler.execute(mockContext, mockBlock, { content: 'This is the content to evaluate.', @@ -487,24 +489,17 @@ describe('EvaluatorBlockHandler', () => { apiKey: 'test-api-key', } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ clarity: 4 }), - model: 'm', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ clarity: 4 }), + model: 'm', + tokens: {}, + cost: 0, + timing: {}, }) await handler.execute(mockContext, mockBlock, inputs) - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ systemPrompt: expect.stringContaining(JSON.stringify(contentObj, null, 2)), }) @@ -524,24 +519,17 @@ describe('EvaluatorBlockHandler', () => { apiKey: 'test-api-key', } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ completeness: 1 }), - model: 'm', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ completeness: 1 }), + model: 'm', + tokens: {}, + cost: 0, + timing: {}, }) await handler.execute(mockContext, mockBlock, inputs) - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ systemPrompt: expect.stringContaining(JSON.stringify(contentObj, null, 2)), }) @@ -560,18 +548,12 @@ describe('EvaluatorBlockHandler', () => { apiKey: 'test-api-key', } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: '```json\n{ "quality": 9 }\n```', - model: 'm', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: '```json\n{ "quality": 9 }\n```', + model: 'm', + tokens: {}, + cost: 0, + timing: {}, }) const result = await handler.execute(mockContext, mockBlock, inputs) @@ -586,18 +568,12 @@ describe('EvaluatorBlockHandler', () => { apiKey: 'test-api-key', } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: 'Sorry, I cannot provide a score.', - model: 'm', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: 'Sorry, I cannot provide a score.', + model: 'm', + tokens: {}, + cost: 0, + timing: {}, }) const result = await handler.execute(mockContext, mockBlock, inputs) @@ -615,18 +591,12 @@ describe('EvaluatorBlockHandler', () => { apiKey: 'test-api-key', } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: '{ "accuracy": 1, "fluency": invalid }', - model: 'm', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: '{ "accuracy": 1, "fluency": invalid }', + model: 'm', + tokens: {}, + cost: 0, + timing: {}, }) const result = await handler.execute(mockContext, mockBlock, inputs) @@ -647,18 +617,12 @@ describe('EvaluatorBlockHandler', () => { apiKey: 'test-api-key', } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ camelcasescore: 7 }), - model: 'm', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ camelcasescore: 7 }), + model: 'm', + tokens: {}, + cost: 0, + timing: {}, }) const result = await handler.execute(mockContext, mockBlock, inputs) @@ -684,18 +648,12 @@ describe('EvaluatorBlockHandler', () => { apiKey: 'test-api-key', } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ presentScore: 4 }), - model: 'm', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ presentScore: 4 }), + model: 'm', + tokens: {}, + cost: 0, + timing: {}, }) const result = await handler.execute(mockContext, mockBlock, inputs) @@ -708,13 +666,7 @@ describe('EvaluatorBlockHandler', () => { const inputs = { content: 'Test error handling.', apiKey: 'test-api-key' } // Override fetch mock to return an error - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: false, - status: 500, - json: () => Promise.resolve({ error: 'Server error' }), - }) - }) + mockExecuteProviderRequest.mockRejectedValueOnce(new Error('Server error')) await expect(handler.execute(mockContext, mockBlock, inputs)).rejects.toThrow('Server error') }) @@ -730,11 +682,7 @@ describe('EvaluatorBlockHandler', () => { ]) registry.recordResolved('CONTENT_SECRET', 'resolved-evaluator-secret') mockContext.resolvedSecretTraceRegistry = registry - mockFetch.mockResolvedValueOnce({ - ok: false, - status: 500, - json: () => Promise.resolve({ error: providerError }), - }) + mockExecuteProviderRequest.mockRejectedValueOnce(new Error(providerError)) await expect( handler.execute(mockContext, mockBlock, { @@ -767,24 +715,17 @@ describe('EvaluatorBlockHandler', () => { mockGetProviderFromModel.mockReturnValue('azure-openai') - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ quality: 8 }), - model: 'gpt-4o', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ quality: 8 }), + model: 'gpt-4o', + tokens: {}, + cost: 0, + timing: {}, }) await handler.execute(mockContext, mockBlock, inputs) - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ provider: 'azure-openai', @@ -824,24 +765,17 @@ describe('EvaluatorBlockHandler', () => { ;(mockDb.db.query as any).account = { findFirst: vi.fn() } vi.spyOn(mockDb.db.query.account, 'findFirst').mockResolvedValue(mockAccount as any) - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ quality: 9 }), - model: 'gemini-2.0-flash-exp', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ quality: 9 }), + model: 'gemini-2.0-flash-exp', + tokens: {}, + cost: 0, + timing: {}, }) await handler.execute(mockContext, mockBlock, inputs) - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ provider: 'vertex', @@ -860,24 +794,17 @@ describe('EvaluatorBlockHandler', () => { // No model provided - should use default } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ score: 7 }), - model: 'claude-sonnet-5', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ score: 7 }), + model: 'claude-sonnet-5', + tokens: {}, + cost: 0, + timing: {}, }) await handler.execute(mockContext, mockBlock, inputs) - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody.model).toBe('claude-sonnet-5') }) diff --git a/apps/sim/executor/handlers/evaluator/evaluator-handler.ts b/apps/sim/executor/handlers/evaluator/evaluator-handler.ts index 6a094f64082..f162ae25c38 100644 --- a/apps/sim/executor/handlers/evaluator/evaluator-handler.ts +++ b/apps/sim/executor/handlers/evaluator/evaluator-handler.ts @@ -1,10 +1,5 @@ import { createLogger } from '@sim/logger' -import { - addModelInputProvenanceToRequest, - createModelInputProvenanceRequestMetadata, - markModelInputProjected, - projectResolvedModelInput, -} from '@/lib/execution/model-input-provenance' +import { projectResolvedModelInput } from '@/lib/execution/model-input-provenance' import { type AutoRoutingResult, addAutoRoutingCost, @@ -15,8 +10,8 @@ import type { BlockOutput } from '@/blocks/types' import { validateModelProvider } from '@/ee/access-control/utils/permission-check' import { BlockType, DEFAULTS, EVALUATOR } from '@/executor/constants' import type { BlockHandler, ExecutionContext } from '@/executor/types' -import { buildAPIUrl, buildAuthHeaders, extractAPIErrorMessage } from '@/executor/utils/http' import { isJSONString, parseJSON, stringifyJSON } from '@/executor/utils/json' +import { executeBlockProviderRequest } from '@/executor/utils/provider-request' import { projectResolvedSecretDiagnosticError } from '@/executor/utils/resolved-secret-content-projection' import { refuseResolvedSecretProjection } from '@/executor/utils/resolved-secret-projection-refusal' import type { @@ -192,8 +187,6 @@ export class EvaluatorBlockHandler implements BlockHandler { } try { - const url = buildAPIUrl('/api/providers', ctx.userId ? { userId: ctx.userId } : {}) - const providerRequest: ProviderRequest = { model, systemPrompt: systemPromptObj.systemPrompt, @@ -219,30 +212,13 @@ export class EvaluatorBlockHandler implements BlockHandler { workspaceId: ctx.workspaceId, } - const headers = new Headers(await buildAuthHeaders(ctx.userId)) - const modelInputMetadata = createModelInputProvenanceRequestMetadata( - modelInputProjection.registry, - modelInputPaths - ) - const requestBody = addModelInputProvenanceToRequest( - { provider: providerId, ...providerRequest }, - headers, - modelInputMetadata - ) - if (modelInputMetadata) markModelInputProjected(headers) - const response = await fetch(url.toString(), { - method: 'POST', - headers, - body: stringifyJSON(requestBody), + const result = await executeBlockProviderRequest({ + ctx, + providerId, + request: providerRequest, + resolvedSecretTraceRegistry: modelInputProjection.registry, }) - if (!response.ok) { - const errorMessage = await extractAPIErrorMessage(response) - throw new Error(errorMessage) - } - - const result = await response.json() - const parsedContent = this.extractJSONFromResponse( result.content, ctx.resolvedSecretTraceRegistry @@ -250,9 +226,8 @@ export class EvaluatorBlockHandler implements BlockHandler { const metricScores = this.extractMetricScores(parsedContent, metrics, projectedMetrics) - const inputTokens = result.tokens?.input || result.tokens?.prompt || DEFAULTS.TOKENS.PROMPT - const outputTokens = - result.tokens?.output || result.tokens?.completion || DEFAULTS.TOKENS.COMPLETION + const inputTokens = result.tokens?.input || DEFAULTS.TOKENS.PROMPT + const outputTokens = result.tokens?.output || DEFAULTS.TOKENS.COMPLETION const cost = addAutoRoutingCost( resolveProxiedModelCost(result.cost), diff --git a/apps/sim/executor/handlers/router/router-handler.test.ts b/apps/sim/executor/handlers/router/router-handler.test.ts index a7fe9e45141..8c7f7b0e945 100644 --- a/apps/sim/executor/handlers/router/router-handler.test.ts +++ b/apps/sim/executor/handlers/router/router-handler.test.ts @@ -9,8 +9,13 @@ import { } from '@sim/testing' import { beforeEach, describe, expect, it, type Mock, vi } from 'vitest' -const { mockResolveAutoModel } = vi.hoisted(() => ({ +const { mockResolveAutoModel, mockCheckWorkspaceAccess } = vi.hoisted(() => ({ mockResolveAutoModel: vi.fn(), + mockCheckWorkspaceAccess: vi.fn(), +})) + +vi.mock('@/lib/workspaces/permissions/utils', () => ({ + checkWorkspaceAccess: mockCheckWorkspaceAccess, })) vi.mock('@/lib/oauth/credential-service', () => authOAuthUtilsMock) @@ -40,27 +45,30 @@ vi.mock('@/lib/model-router/resolve', () => ({ SIM_AUTO_SYSTEM_PREAMBLE: 'Sim auto system preamble', })) -import { - PRIVATE_MODEL_INPUT_PROVENANCE_HEADER, - PRIVATE_MODEL_INPUT_STATE_HEADER, - PROJECTED_MODEL_INPUT_PATHS_V1, -} from '@/lib/execution/model-input-provenance' -import { - RESOLVED_SECRET_PROVENANCE_FIELD, - RESOLVED_SECRET_PROVENANCE_METADATA_V1, -} from '@/lib/execution/private-tool-metadata' import { generateRouterPrompt, generateRouterV2Prompt } from '@/blocks/blocks/router' import { BlockType } from '@/executor/constants' import { RouterBlockHandler } from '@/executor/handlers/router/router-handler' import type { ExecutionContext } from '@/executor/types' import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' +import { executeProviderRequest } from '@/providers' import { getProviderFromModel } from '@/providers/utils' import type { SerializedBlock, SerializedWorkflow } from '@/serializer/types' const mockGenerateRouterPrompt = generateRouterPrompt as Mock const mockGenerateRouterV2Prompt = generateRouterV2Prompt as Mock const mockGetProviderFromModel = getProviderFromModel as Mock -const mockFetch = vi.fn() +const mockExecuteProviderRequest = executeProviderRequest as Mock + +/** The provider request the handler built, keyed the way the old wire body was. */ +function providerRequestBody(index = 0): Record { + const [provider, request] = mockExecuteProviderRequest.mock.calls[index] + return { provider, ...request } +} + +function providerRuntimeRegistry(index = 0): ResolvedSecretTraceRegistry | undefined { + return mockExecuteProviderRequest.mock.calls[index][2]?.resolvedSecretTraceRegistry +} + const mockLogger = vi.mocked(createLogger).mock.results[ vi.mocked(createLogger).mock.calls.findIndex(([name]) => name === 'RouterBlockHandler') @@ -138,8 +146,7 @@ describe('RouterBlockHandler', () => { vi.clearAllMocks() encryptionMockFns.mockDecryptSecret.mockResolvedValue({ decrypted: 'test-decrypted' }) - // unstubGlobals removes any module-scope fetch stub before each test, so re-stub here - vi.stubGlobal('fetch', mockFetch) + mockCheckWorkspaceAccess.mockResolvedValue({ hasAccess: true }) authOAuthUtilsMockFns.mockResolveOAuthAccountId.mockResolvedValue({ accountId: 'test-vertex-credential-id', @@ -158,18 +165,12 @@ describe('RouterBlockHandler', () => { billableRoutingCost: 0.002, }) - mockFetch.mockImplementation(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: 'target-block-1', - model: 'mock-model', - tokens: { input: 100, output: 5, total: 105 }, - cost: 0.003, - timing: { total: 300 }, - }), - }) + mockExecuteProviderRequest.mockResolvedValue({ + content: 'target-block-1', + model: 'mock-model', + tokens: { input: 100, output: 5, total: 105 }, + cost: 0.003, + timing: { total: 300 }, }) }) @@ -219,17 +220,9 @@ describe('RouterBlockHandler', () => { expect(mockGenerateRouterPrompt).toHaveBeenCalledWith(inputs.prompt, expectedTargetBlocks) expect(mockGetProviderFromModel).toHaveBeenCalledWith('gpt-4o') - expect(mockFetch).toHaveBeenCalledWith( - expect.any(String), - expect.objectContaining({ - method: 'POST', - headers: expect.any(Object), - body: expect.any(String), - }) - ) + expect(mockExecuteProviderRequest).toHaveBeenCalledTimes(1) - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ provider: 'openai', model: 'gpt-4o', @@ -282,15 +275,8 @@ describe('RouterBlockHandler', () => { apiKey: credentialSecret, }) - const request = mockFetch.mock.calls[0][1] - const requestBody = JSON.parse(request.body) - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBe( - RESOLVED_SECRET_PROVENANCE_METADATA_V1 - ) - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBe( - PROJECTED_MODEL_INPUT_PATHS_V1 - ) - expect(requestBody[RESOLVED_SECRET_PROVENANCE_FIELD]).toEqual({ + const requestBody = providerRequestBody() + expect(providerRuntimeRegistry()?.exportProvenance()).toEqual({ version: 1, complete: true, entries: [ @@ -353,8 +339,7 @@ describe('RouterBlockHandler', () => { expect(rawState).toEqual({ result: stateSecret, ordinary: 'Box remains raw state' }) expect(mockTargetBlock1.config.params).toEqual({ p: 'a' }) - const requestBody = JSON.parse(mockFetch.mock.calls[0][1].body) - expect(requestBody[RESOLVED_SECRET_PROVENANCE_FIELD]).toEqual({ + expect(providerRuntimeRegistry()?.exportProvenance()).toEqual({ version: 1, complete: true, entries: [], @@ -404,29 +389,19 @@ describe('RouterBlockHandler', () => { apiKey: 'test-api-key', }) - const request = mockFetch.mock.calls[0][1] - const requestBody = JSON.parse(request.body) - expect(Object.hasOwn(requestBody, RESOLVED_SECRET_PROVENANCE_FIELD)).toBe(false) - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBeNull() - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBeNull() + expect(providerRuntimeRegistry()).toBeUndefined() }) it('bills the cost the provider proxy decided rather than recomputing it', async () => { // The proxy already resolved key provenance and the margin; recomputing // here would re-charge a BYOK caller the proxy correctly zeroed. - mockFetch.mockImplementation(() => - Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: 'target-block-1', - model: 'mock-model', - tokens: { input: 100, output: 5, total: 105 }, - cost: { input: 0.004, output: 0.002, total: 0.006 }, - timing: { total: 300 }, - }), - }) - ) + mockExecuteProviderRequest.mockResolvedValue({ + content: 'target-block-1', + model: 'mock-model', + tokens: { input: 100, output: 5, total: 105 }, + cost: { input: 0.004, output: 0.002, total: 0.006 }, + timing: { total: 300 }, + }) const result = await handler.execute(mockContext, mockBlock, { prompt: 'Choose the best option.', @@ -439,6 +414,26 @@ describe('RouterBlockHandler', () => { }) }) + it('refuses to reach the provider without an execution subject', async () => { + mockContext.userId = undefined + + await expect( + handler.execute(mockContext, mockBlock, { prompt: 'Choose the best option.' }) + ).rejects.toThrow('Unauthorized') + expect(mockExecuteProviderRequest).not.toHaveBeenCalled() + }) + + it('refuses to reach the provider when the subject lost workspace access', async () => { + mockContext.workspaceId = 'test-workspace' + mockCheckWorkspaceAccess.mockResolvedValue({ hasAccess: false }) + + await expect( + handler.execute(mockContext, mockBlock, { prompt: 'Choose the best option.' }) + ).rejects.toThrow('Forbidden') + expect(mockCheckWorkspaceAccess).toHaveBeenCalledWith('test-workspace', 'test-user') + expect(mockExecuteProviderRequest).not.toHaveBeenCalled() + }) + it('should throw error if target block is missing', async () => { const inputs = { prompt: 'Test' } mockContext.workflow!.blocks = [mockBlock, mockTargetBlock2] @@ -446,24 +441,18 @@ describe('RouterBlockHandler', () => { await expect(handler.execute(mockContext, mockBlock, inputs)).rejects.toThrow( 'Target block target-block-1 not found' ) - expect(mockFetch).not.toHaveBeenCalled() + expect(mockExecuteProviderRequest).not.toHaveBeenCalled() }) it('should throw error if LLM response is not a valid target block ID', async () => { const inputs = { prompt: 'Test', apiKey: 'test-api-key' } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: 'invalid-block-id', - model: 'mock-model', - tokens: {}, - cost: 0, - timing: {}, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: 'invalid-block-id', + model: 'mock-model', + tokens: {}, + cost: 0, + timing: {}, }) await expect(handler.execute(mockContext, mockBlock, inputs)).rejects.toThrow( @@ -475,16 +464,12 @@ describe('RouterBlockHandler', () => { const plaintext = 'router-provider-plaintext-secret' const content = `${plaintext} __var_API_KEY __sim_runtime` - mockFetch.mockResolvedValueOnce({ - ok: true, - json: () => - Promise.resolve({ - content, - model: 'mock-model', - tokens: {}, - cost: 0, - timing: {}, - }), + mockExecuteProviderRequest.mockResolvedValueOnce({ + content, + model: 'mock-model', + tokens: {}, + cost: 0, + timing: {}, }) await expect(handler.execute(mockContext, mockBlock, { prompt: 'Test' })).rejects.toThrow( @@ -512,8 +497,7 @@ describe('RouterBlockHandler', () => { expect(mockGetProviderFromModel).toHaveBeenCalledWith('claude-sonnet-5') - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ model: 'claude-sonnet-5', temperature: 0.1, @@ -523,13 +507,7 @@ describe('RouterBlockHandler', () => { it('should handle server error responses', async () => { const inputs = { prompt: 'Test error handling.', apiKey: 'test-api-key' } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: false, - status: 500, - json: () => Promise.resolve({ error: 'Server error' }), - }) - }) + mockExecuteProviderRequest.mockRejectedValueOnce(new Error('Server error')) await expect(handler.execute(mockContext, mockBlock, inputs)).rejects.toThrow('Server error') }) @@ -537,11 +515,7 @@ describe('RouterBlockHandler', () => { it('does not log sensitive provider errors while preserving the thrown error', async () => { const providerError = 'provider-plaintext-secret __var_API_KEY __sim_runtime' - mockFetch.mockResolvedValueOnce({ - ok: false, - status: 500, - json: () => Promise.resolve({ error: providerError }), - }) + mockExecuteProviderRequest.mockRejectedValueOnce(new Error(providerError)) await expect(handler.execute(mockContext, mockBlock, { prompt: 'Test' })).rejects.toThrow( providerError @@ -569,8 +543,7 @@ describe('RouterBlockHandler', () => { await handler.execute(mockContext, mockBlock, inputs) - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ provider: 'azure-openai', @@ -604,8 +577,7 @@ describe('RouterBlockHandler', () => { await handler.execute(mockContext, mockBlock, inputs) - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ provider: 'vertex', @@ -688,8 +660,7 @@ describe('RouterBlockHandler V2', () => { vi.clearAllMocks() - // unstubGlobals removes any module-scope fetch stub before each test, so re-stub here - vi.stubGlobal('fetch', mockFetch) + mockCheckWorkspaceAccess.mockResolvedValue({ hasAccess: true }) authOAuthUtilsMockFns.mockResolveOAuthAccountId.mockResolvedValue({ accountId: 'test-vertex-credential-id', @@ -732,19 +703,13 @@ describe('RouterBlockHandler V2', () => { ]), } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ - route: 'route-support', - reasoning: 'The user mentioned a billing issue which is a customer support matter.', - }), - model: 'gpt-4o', - tokens: { input: 150, output: 25, total: 175 }, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ + route: 'route-support', + reasoning: 'The user mentioned a billing issue which is a customer support matter.', + }), + model: 'gpt-4o', + tokens: { input: 150, output: 25, total: 175 }, }) const result = await handler.execute(mockContext, mockRouterV2Block, inputs) @@ -781,14 +746,10 @@ describe('RouterBlockHandler V2', () => { registry.recordResolvedInputProjection(['context'], contextSecret, '{{CONTEXT_SECRET}}') registry.recordResolved('API_KEY', credentialSecret) mockContext.resolvedSecretTraceRegistry = registry - mockFetch.mockResolvedValueOnce({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ route: 'route-support', reasoning: 'Matched support.' }), - model: 'gpt-4o', - tokens: { input: 10, output: 5, total: 15 }, - }), + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ route: 'route-support', reasoning: 'Matched support.' }), + model: 'gpt-4o', + tokens: { input: 10, output: 5, total: 15 }, }) await handler.execute(mockContext, mockRouterV2Block, { @@ -798,15 +759,8 @@ describe('RouterBlockHandler V2', () => { routes: [{ id: 'route-support', title: 'Support', value: 'Support requests' }], }) - const request = mockFetch.mock.calls[0][1] - const requestBody = JSON.parse(request.body) - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBe( - RESOLVED_SECRET_PROVENANCE_METADATA_V1 - ) - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBe( - PROJECTED_MODEL_INPUT_PATHS_V1 - ) - expect(requestBody[RESOLVED_SECRET_PROVENANCE_FIELD]).toEqual({ + const requestBody = providerRequestBody() + expect(providerRuntimeRegistry()?.exportProvenance()).toEqual({ version: 1, complete: true, entries: [ @@ -821,14 +775,10 @@ describe('RouterBlockHandler V2', () => { }) it('keeps the router V2 request shape when no provenance registry exists', async () => { - mockFetch.mockResolvedValueOnce({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ route: 'route-support', reasoning: 'Matched support.' }), - model: 'gpt-4o', - tokens: { input: 10, output: 5, total: 15 }, - }), + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ route: 'route-support', reasoning: 'Matched support.' }), + model: 'gpt-4o', + tokens: { input: 10, output: 5, total: 15 }, }) await handler.execute(mockContext, mockRouterV2Block, { @@ -838,11 +788,7 @@ describe('RouterBlockHandler V2', () => { routes: [{ id: 'route-support', title: 'Support', value: 'Support requests' }], }) - const request = mockFetch.mock.calls[0][1] - const requestBody = JSON.parse(request.body) - expect(Object.hasOwn(requestBody, RESOLVED_SECRET_PROVENANCE_FIELD)).toBe(false) - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_PROVENANCE_HEADER)).toBeNull() - expect((request.headers as Headers).get(PRIVATE_MODEL_INPUT_STATE_HEADER)).toBeNull() + expect(providerRuntimeRegistry()).toBeUndefined() }) it('resolves sim-auto before executing router V2 and preserves its public identity', async () => { @@ -859,18 +805,14 @@ describe('RouterBlockHandler V2', () => { ], } - mockFetch.mockResolvedValueOnce({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ - route: 'route-sales', - reasoning: 'This is a new request.', - }), - model: 'fireworks/glm-5.2', - tokens: { input: 100, output: 20, total: 120 }, - cost: { input: 0.001, output: 0.0005, total: 0.0015 }, - }), + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ + route: 'route-sales', + reasoning: 'This is a new request.', + }), + model: 'fireworks/glm-5.2', + tokens: { input: 100, output: 20, total: 120 }, + cost: { input: 0.001, output: 0.0005, total: 0.0015 }, }) const result = await handler.execute(mockContext, mockRouterV2Block, inputs) @@ -889,7 +831,7 @@ describe('RouterBlockHandler V2', () => { }) expect(mockGetProviderFromModel).toHaveBeenCalledWith('fireworks/glm-5.2') - const requestBody = JSON.parse(mockFetch.mock.calls[0][1].body) + const requestBody = providerRequestBody() expect(requestBody).toMatchObject({ provider: 'openai', model: 'fireworks/glm-5.2', @@ -915,25 +857,18 @@ describe('RouterBlockHandler V2', () => { routes: JSON.stringify([{ id: 'route-1', title: 'Route 1', value: 'Description 1' }]), } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ - route: 'route-1', - reasoning: 'Test reasoning', - }), - model: 'gpt-4o', - tokens: { input: 100, output: 20, total: 120 }, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ + route: 'route-1', + reasoning: 'Test reasoning', + }), + model: 'gpt-4o', + tokens: { input: 100, output: 20, total: 120 }, }) await handler.execute(mockContext, mockRouterV2Block, inputs) - const fetchCallArgs = mockFetch.mock.calls[0] - const requestBody = JSON.parse(fetchCallArgs[1].body) + const requestBody = providerRequestBody() expect(requestBody.responseFormat).toEqual({ name: 'router_response', @@ -964,19 +899,13 @@ describe('RouterBlockHandler V2', () => { routes: JSON.stringify([{ id: 'route-1', title: 'Route 1', value: 'Specific topic' }]), } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ - route: 'NO_MATCH', - reasoning: 'The query does not relate to any available route.', - }), - model: 'gpt-4o', - tokens: { input: 100, output: 20, total: 120 }, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ + route: 'NO_MATCH', + reasoning: 'The query does not relate to any available route.', + }), + model: 'gpt-4o', + tokens: { input: 100, output: 20, total: 120 }, }) await expect(handler.execute(mockContext, mockRouterV2Block, inputs)).rejects.toThrow( @@ -992,19 +921,13 @@ describe('RouterBlockHandler V2', () => { routes: JSON.stringify([{ id: 'route-1', title: 'Route 1', value: 'Description' }]), } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ - route: 'invalid-route', - reasoning: 'Some reasoning', - }), - model: 'gpt-4o', - tokens: { input: 100, output: 20, total: 120 }, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ + route: 'invalid-route', + reasoning: 'Some reasoning', + }), + model: 'gpt-4o', + tokens: { input: 100, output: 20, total: 120 }, }) await expect(handler.execute(mockContext, mockRouterV2Block, inputs)).rejects.toThrow( @@ -1020,19 +943,13 @@ describe('RouterBlockHandler V2', () => { routes: [{ id: 'route-1', title: 'Route 1', value: 'Description' }], } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: JSON.stringify({ - route: 'route-1', - reasoning: 'Matched route 1', - }), - model: 'gpt-4o', - tokens: { input: 100, output: 20, total: 120 }, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: JSON.stringify({ + route: 'route-1', + reasoning: 'Matched route 1', + }), + model: 'gpt-4o', + tokens: { input: 100, output: 20, total: 120 }, }) const result = await handler.execute(mockContext, mockRouterV2Block, inputs) @@ -1084,16 +1001,10 @@ describe('RouterBlockHandler V2', () => { routes: JSON.stringify([{ id: 'route-1', title: 'Route 1', value: 'Description' }]), } - mockFetch.mockImplementationOnce(() => { - return Promise.resolve({ - ok: true, - json: () => - Promise.resolve({ - content: 'route-1', - model: 'gpt-4o', - tokens: { input: 100, output: 5, total: 105 }, - }), - }) + mockExecuteProviderRequest.mockResolvedValueOnce({ + content: 'route-1', + model: 'gpt-4o', + tokens: { input: 100, output: 5, total: 105 }, }) const result = await handler.execute(mockContext, mockRouterV2Block, inputs) @@ -1111,14 +1022,10 @@ describe('RouterBlockHandler V2', () => { routes: [{ id: 'route-1', title: 'Route 1', value: 'Description' }], } - mockFetch.mockResolvedValueOnce({ - ok: true, - json: () => - Promise.resolve({ - content, - model: 'gpt-4o', - tokens: { input: 100, output: 5, total: 105 }, - }), + mockExecuteProviderRequest.mockResolvedValueOnce({ + content, + model: 'gpt-4o', + tokens: { input: 100, output: 5, total: 105 }, }) await expect(handler.execute(mockContext, mockRouterV2Block, inputs)).rejects.toThrow(content) diff --git a/apps/sim/executor/handlers/router/router-handler.ts b/apps/sim/executor/handlers/router/router-handler.ts index 11e5445889a..3292e2c57c2 100644 --- a/apps/sim/executor/handlers/router/router-handler.ts +++ b/apps/sim/executor/handlers/router/router-handler.ts @@ -1,11 +1,5 @@ import { createLogger } from '@sim/logger' -import { getInternalApiBaseUrl } from '@/lib/core/utils/urls' -import { - addModelInputProvenanceToRequest, - createModelInputProvenanceRequestMetadata, - markModelInputProjected, - projectResolvedModelInput, -} from '@/lib/execution/model-input-provenance' +import { projectResolvedModelInput } from '@/lib/execution/model-input-provenance' import { type AutoRoutingResult, addAutoRoutingCost, @@ -23,7 +17,7 @@ import { ROUTER, } from '@/executor/constants' import type { BlockHandler, ExecutionContext } from '@/executor/types' -import { buildAuthHeaders } from '@/executor/utils/http' +import { executeBlockProviderRequest } from '@/executor/utils/provider-request' import { refuseResolvedSecretProjection } from '@/executor/utils/resolved-secret-projection-refusal' import type { ResolvedSecretInputPath } from '@/executor/utils/resolved-secret-trace-registry' import { resolveVertexCredential } from '@/executor/utils/vertex-credential' @@ -103,9 +97,6 @@ export class RouterBlockHandler implements BlockHandler { } try { - const url = new URL('/api/providers', getInternalApiBaseUrl()) - if (ctx.userId) url.searchParams.set('userId', ctx.userId) - const messages = [{ role: 'user', content: routerConfig.prompt }] const systemPrompt = generateRouterPrompt(routerConfig.prompt, targetBlocks) const resolved = await this.resolveModel( @@ -147,36 +138,13 @@ export class RouterBlockHandler implements BlockHandler { workspaceId: ctx.workspaceId, } - const headers = new Headers(await buildAuthHeaders(ctx.userId)) - const modelInputMetadata = createModelInputProvenanceRequestMetadata( - modelInputProjection.registry, - promptModelInputPaths - ) - const requestBody = addModelInputProvenanceToRequest( - { provider: providerId, ...providerRequest }, - headers, - modelInputMetadata - ) - if (modelInputMetadata) markModelInputProjected(headers) - const response = await fetch(url.toString(), { - method: 'POST', - headers, - body: JSON.stringify(requestBody), + const result = await executeBlockProviderRequest({ + ctx, + providerId, + request: providerRequest, + resolvedSecretTraceRegistry: modelInputProjection.registry, }) - if (!response.ok) { - let errorMessage = `Provider API request failed with status ${response.status}` - try { - const errorData = await response.json() - if (errorData.error) { - errorMessage = errorData.error - } - } catch (_e) {} - throw new Error(errorMessage) - } - - const result = await response.json() - const chosenBlockId = result.content.trim().toLowerCase() const chosenBlock = targetBlocks?.find((b) => b.id === chosenBlockId) @@ -291,9 +259,6 @@ export class RouterBlockHandler implements BlockHandler { } try { - const url = new URL('/api/providers', getInternalApiBaseUrl()) - if (ctx.userId) url.searchParams.set('userId', ctx.userId) - const messages = [{ role: 'user', content: routerConfig.context }] const systemPrompt = generateRouterV2Prompt(routerConfig.context, modelRoutes) const resolved = await this.resolveModel( @@ -354,36 +319,13 @@ export class RouterBlockHandler implements BlockHandler { }, } - const headers = new Headers(await buildAuthHeaders(ctx.userId)) - const modelInputMetadata = createModelInputProvenanceRequestMetadata( - modelInputProjection.registry, - modelInputPaths - ) - const requestBody = addModelInputProvenanceToRequest( - { provider: providerId, ...providerRequest }, - headers, - modelInputMetadata - ) - if (modelInputMetadata) markModelInputProjected(headers) - const response = await fetch(url.toString(), { - method: 'POST', - headers, - body: JSON.stringify(requestBody), + const result = await executeBlockProviderRequest({ + ctx, + providerId, + request: providerRequest, + resolvedSecretTraceRegistry: modelInputProjection.registry, }) - if (!response.ok) { - let errorMessage = `Provider API request failed with status ${response.status}` - try { - const errorData = await response.json() - if (errorData.error) { - errorMessage = errorData.error - } - } catch (_e) {} - throw new Error(errorMessage) - } - - const result = await response.json() - let chosenRouteId: string let reasoning = '' diff --git a/apps/sim/executor/utils/provider-request.ts b/apps/sim/executor/utils/provider-request.ts new file mode 100644 index 00000000000..622d5f77a35 --- /dev/null +++ b/apps/sim/executor/utils/provider-request.ts @@ -0,0 +1,70 @@ +import { createLogger } from '@sim/logger' +import { checkWorkspaceAccess } from '@/lib/workspaces/permissions/utils' +import type { ExecutionContext } from '@/executor/types' +import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' +import { executeProviderRequest } from '@/providers' +import type { ProviderRequest, ProviderResponse } from '@/providers/types' + +const logger = createLogger('ExecutorProviderRequest') + +interface ExecuteBlockProviderRequestInput { + ctx: ExecutionContext + providerId: string + request: ProviderRequest + /** Supplied in place of the provenance envelope the HTTP boundary serialized and re-imported. */ + resolvedSecretTraceRegistry: ResolvedSecretTraceRegistry | undefined +} + +/** + * Runs one non-streaming provider request for a block handler in-process, replacing the + * executor's `POST /api/providers` round trip. The route's two admission checks are + * reproduced so the outcome is unchanged: an internal token with no user is rejected (the + * executor mints it from `ctx.userId`), and an execution subject who has left the billed + * workspace is rejected. The route's remaining work is already done by the caller or lives + * inside `executeProviderRequest`. + */ +export async function executeBlockProviderRequest({ + ctx, + providerId, + request, + resolvedSecretTraceRegistry, +}: ExecuteBlockProviderRequestInput): Promise { + if (!ctx.userId) { + throw new Error('Unauthorized') + } + + if (request.workspaceId) { + const workspaceAccess = await checkWorkspaceAccess(request.workspaceId, ctx.userId) + if (!workspaceAccess.hasAccess) { + throw new Error('Forbidden') + } + } + + /** + * No `executionContext`: it is only inherited by model-emitted tool calls, and the route + * this replaces never carried one. The whole context is omitted when there is no registry + * rather than passed carrying `undefined` — `executeProviderTool` reads that as missing + * provenance and fails the call closed with no error text. + */ + const response = await executeProviderRequest( + providerId, + { ...request, userId: ctx.userId }, + resolvedSecretTraceRegistry ? { resolvedSecretTraceRegistry } : undefined + ) + + if ( + response instanceof ReadableStream || + (typeof response === 'object' && response !== null && 'stream' in response) + ) { + logger.error('Provider returned a stream for a non-streaming block request', { providerId }) + throw new Error('Provider returned a streaming response for a non-streaming request') + } + + logger.info('Provider request completed', { + providerId, + model: request.model, + workflowId: ctx.workflowId, + }) + + return response +} diff --git a/apps/sim/hooks/queries/invitations.ts b/apps/sim/hooks/queries/invitations.ts index 394ff0b9113..3f93f6638ab 100644 --- a/apps/sim/hooks/queries/invitations.ts +++ b/apps/sim/hooks/queries/invitations.ts @@ -22,8 +22,8 @@ import { import { updateWorkspacePermissionsContract } from '@/lib/api/contracts/workspaces' import { organizationKeys } from '@/hooks/queries/organization' import { refreshSessionQuery } from '@/hooks/queries/session' -import { subscriptionKeys } from '@/hooks/queries/subscription' import { workspaceCredentialKeys } from '@/hooks/queries/utils/credential-keys' +import { subscriptionKeys } from '@/hooks/queries/utils/subscription-keys' import { workspaceKeys } from '@/hooks/queries/workspace' export const invitationKeys = { diff --git a/apps/sim/hooks/queries/kb/knowledge.test.ts b/apps/sim/hooks/queries/kb/knowledge.test.ts new file mode 100644 index 00000000000..4c01ae0fcf5 --- /dev/null +++ b/apps/sim/hooks/queries/kb/knowledge.test.ts @@ -0,0 +1,75 @@ +/** + * @vitest-environment node + */ + +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const mocks = vi.hoisted(() => ({ + requestJson: vi.fn(), + useMutation: vi.fn(), + invalidateQueries: vi.fn(), +})) + +vi.mock('@tanstack/react-query', () => ({ + keepPreviousData: Symbol('keepPreviousData'), + useInfiniteQuery: vi.fn(), + useMutation: mocks.useMutation, + useQuery: vi.fn(), + useQueryClient: vi.fn(() => ({ invalidateQueries: mocks.invalidateQueries })), +})) + +vi.mock('@sim/emcn', () => ({ + toast: { error: vi.fn(), success: vi.fn() }, +})) + +vi.mock('@/lib/api/client/request', () => ({ + requestJson: mocks.requestJson, +})) + +import { useBulkDocumentOperation, useDeleteDocument } from '@/hooks/queries/kb/knowledge' +import { knowledgeKeys } from '@/hooks/queries/utils/knowledge-keys' + +interface CapturedMutation { + onSettled: (data: unknown, error: unknown, variables: Record) => void +} + +function captureMutation(build: () => unknown): CapturedMutation { + let captured: CapturedMutation | undefined + mocks.useMutation.mockImplementation((options: CapturedMutation) => { + captured = options + return {} + }) + build() + if (!captured) throw new Error('useMutation was not called') + return captured +} + +describe('knowledge document mutations', () => { + beforeEach(() => { + vi.clearAllMocks() + }) + + it('invalidates the knowledge-base lists when a document is deleted', () => { + const mutation = captureMutation(() => useDeleteDocument()) + + mutation.onSettled(undefined, undefined, { knowledgeBaseId: 'kb-1', documentId: 'doc-1' }) + + expect(mocks.invalidateQueries).toHaveBeenCalledWith({ queryKey: knowledgeKeys.lists() }) + }) + + it('invalidates the knowledge-base lists on a bulk delete', () => { + const mutation = captureMutation(() => useBulkDocumentOperation()) + + mutation.onSettled(undefined, undefined, { knowledgeBaseId: 'kb-1', operation: 'delete' }) + + expect(mocks.invalidateQueries).toHaveBeenCalledWith({ queryKey: knowledgeKeys.lists() }) + }) + + it('leaves the knowledge-base lists alone on a bulk enable', () => { + const mutation = captureMutation(() => useBulkDocumentOperation()) + + mutation.onSettled(undefined, undefined, { knowledgeBaseId: 'kb-1', operation: 'enable' }) + + expect(mocks.invalidateQueries).not.toHaveBeenCalledWith({ queryKey: knowledgeKeys.lists() }) + }) +}) diff --git a/apps/sim/hooks/queries/kb/knowledge.ts b/apps/sim/hooks/queries/kb/knowledge.ts index 0dd4606b44f..f46457282fc 100644 --- a/apps/sim/hooks/queries/kb/knowledge.ts +++ b/apps/sim/hooks/queries/kb/knowledge.ts @@ -539,6 +539,10 @@ export function useDeleteDocument() { queryClient.invalidateQueries({ queryKey: knowledgeKeys.detail(knowledgeBaseId), }) + /** The knowledge-base list rows carry `docCount`, so removing a document changes them too. */ + queryClient.invalidateQueries({ + queryKey: knowledgeKeys.lists(), + }) }, }) } @@ -573,10 +577,16 @@ export function useBulkDocumentOperation() { return useMutation({ mutationFn: bulkDocumentOperation, - onSettled: (_data, _error, { knowledgeBaseId }) => { + onSettled: (_data, _error, { knowledgeBaseId, operation }) => { queryClient.invalidateQueries({ queryKey: knowledgeKeys.detail(knowledgeBaseId), }) + /** Only a bulk delete changes the `docCount` the knowledge-base list rows render. */ + if (operation === 'delete') { + queryClient.invalidateQueries({ + queryKey: knowledgeKeys.lists(), + }) + } }, }) } diff --git a/apps/sim/hooks/queries/oauth/oauth-connections.ts b/apps/sim/hooks/queries/oauth/oauth-connections.ts index 6a08f723d1f..25338567464 100644 --- a/apps/sim/hooks/queries/oauth/oauth-connections.ts +++ b/apps/sim/hooks/queries/oauth/oauth-connections.ts @@ -56,6 +56,18 @@ function defineServices(): ServiceInfo[] { return servicesList } +/** + * Resolves the service catalog merged with the caller's connections. + * + * A failed request resolves with the bare catalog rather than rejecting, so + * consumers keep correct service names and ids when the merge data is + * unavailable. The cost is that `isConnected`/`accounts` then report *unknown* + * as *disconnected*, which the result cannot distinguish. Read connection + * state from the workspace credentials query (`useWorkspaceCredentials`), which + * surfaces its own errors; a consumer that must branch on `isConnected` here + * needs this fallback removed first, or it will tell a connected user they are + * not. + */ async function fetchOAuthConnections(signal?: AbortSignal): Promise { try { const serviceDefinitions = defineServices() diff --git a/apps/sim/hooks/queries/organization.ts b/apps/sim/hooks/queries/organization.ts index 0845c99f71d..a0ce1295633 100644 --- a/apps/sim/hooks/queries/organization.ts +++ b/apps/sim/hooks/queries/organization.ts @@ -41,8 +41,8 @@ import { import { client } from '@/lib/auth/auth-client' import { isEnterprise, isPaid, isTeam } from '@/lib/billing/plan-helpers' import { hasPaidSubscriptionStatus } from '@/lib/billing/subscriptions/utils' -import { subscriptionKeys } from '@/hooks/queries/subscription' import { workspaceCredentialKeys } from '@/hooks/queries/utils/credential-keys' +import { subscriptionKeys } from '@/hooks/queries/utils/subscription-keys' import { workspaceKeys } from '@/hooks/queries/workspace' const logger = createLogger('OrganizationQueries') diff --git a/apps/sim/hooks/queries/schedules.ts b/apps/sim/hooks/queries/schedules.ts index a2db0ad31e9..a22fe07d4fe 100644 --- a/apps/sim/hooks/queries/schedules.ts +++ b/apps/sim/hooks/queries/schedules.ts @@ -27,7 +27,13 @@ export const scheduleKeys = { details: () => [...scheduleKeys.all, 'detail'] as const, schedule: (workflowId: string, blockId: string) => [...scheduleKeys.details(), workflowId, blockId] as const, - byId: (scheduleId: string) => [...scheduleKeys.details(), scheduleId] as const, + /** + * By-id reads sit under their own segment rather than directly under `details()`: + * a bare `[...details(), scheduleId]` is a prefix of `schedule(scheduleId, blockId)`, + * so the two addressings of the same schedule would alias in the cache. + */ + byIds: () => [...scheduleKeys.details(), 'by-id'] as const, + byId: (scheduleId: string) => [...scheduleKeys.byIds(), scheduleId] as const, } export type ScheduleData = WorkflowScheduleRow @@ -196,7 +202,7 @@ export function useReactivateSchedule() { body: { action: 'reactivate' }, }) - return { workflowId, blockId, workspaceId } + return { scheduleId, workflowId, blockId, workspaceId } }, onSuccess: ({ workflowId, blockId }) => { logger.info('Schedule reactivated', { workflowId, blockId }) @@ -206,9 +212,10 @@ export function useReactivateSchedule() { }, onSettled: async (data) => { if (!data) return - const { workflowId, blockId, workspaceId } = data + const { scheduleId, workflowId, blockId, workspaceId } = data await Promise.all([ queryClient.invalidateQueries({ queryKey: scheduleKeys.schedule(workflowId, blockId) }), + queryClient.invalidateQueries({ queryKey: scheduleKeys.byId(scheduleId) }), workspaceId ? queryClient.invalidateQueries({ queryKey: scheduleKeys.list(workspaceId) }) : Promise.resolve(), @@ -242,6 +249,8 @@ export function useRedeployWorkflowSchedule() { const { workflowId, blockId } = data await Promise.all([ queryClient.invalidateQueries({ queryKey: scheduleKeys.schedule(workflowId, blockId) }), + /** A redeploy recreates the schedule; the id-keyed reads are a separate subtree. */ + queryClient.invalidateQueries({ queryKey: scheduleKeys.byIds() }), queryClient.invalidateQueries({ queryKey: deploymentKeys.info(workflowId) }), queryClient.invalidateQueries({ queryKey: deploymentKeys.versions(workflowId) }), ]) diff --git a/apps/sim/hooks/queries/subscription.ts b/apps/sim/hooks/queries/subscription.ts index bafee85ad5b..5809f6b2eba 100644 --- a/apps/sim/hooks/queries/subscription.ts +++ b/apps/sim/hooks/queries/subscription.ts @@ -13,6 +13,8 @@ import { updateUsageLimitContract, } from '@/lib/api/contracts/subscription' import { organizationKeys } from '@/hooks/queries/organization' +import { invalidateWorkspaceUsage } from '@/hooks/queries/utils/invalidate-usage' +import { subscriptionKeys } from '@/hooks/queries/utils/subscription-keys' import { workspaceKeys } from '@/hooks/queries/workspace' export type { SubscriptionApiResponse } @@ -21,19 +23,6 @@ export const SUBSCRIPTION_DATA_STALE_TIME = 5 * 60 * 1000 export const USAGE_LIMIT_STALE_TIME = 30 * 1000 export const INVOICES_STALE_TIME = 5 * 60 * 1000 -/** - * Query key factories for subscription-related queries - */ -export const subscriptionKeys = { - all: ['subscription'] as const, - users: () => [...subscriptionKeys.all, 'user'] as const, - user: (includeOrg?: boolean) => [...subscriptionKeys.users(), { includeOrg }] as const, - usage: () => [...subscriptionKeys.all, 'usage'] as const, - invoicesAll: () => [...subscriptionKeys.all, 'invoices'] as const, - invoices: (context: 'user' | 'organization' = 'user', organizationId?: string) => - [...subscriptionKeys.invoicesAll(), context, organizationId ?? ''] as const, -} - /** * Fetch user subscription data * @param includeOrg - Whether to include organization role data @@ -265,6 +254,7 @@ export function useUpdateUsageLimit() { return Promise.all([ queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }), queryClient.invalidateQueries({ queryKey: subscriptionKeys.usage() }), + invalidateWorkspaceUsage(queryClient), ]) }, }) @@ -291,6 +281,7 @@ export function useUpgradeSubscription() { queryClient.invalidateQueries({ queryKey: subscriptionKeys.usage() }), queryClient.invalidateQueries({ queryKey: subscriptionKeys.invoicesAll() }), queryClient.invalidateQueries({ queryKey: workspaceKeys.lists() }), + invalidateWorkspaceUsage(queryClient), ...(variables.orgId ? [ queryClient.invalidateQueries({ @@ -328,6 +319,7 @@ export function usePurchaseCredits() { return Promise.all([ queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }), queryClient.invalidateQueries({ queryKey: subscriptionKeys.usage() }), + invalidateWorkspaceUsage(queryClient), ...(variables.orgId ? [ queryClient.invalidateQueries({ diff --git a/apps/sim/hooks/queries/tables.test.ts b/apps/sim/hooks/queries/tables.test.ts index 876fe2a981e..b90836494f4 100644 --- a/apps/sim/hooks/queries/tables.test.ts +++ b/apps/sim/hooks/queries/tables.test.ts @@ -69,6 +69,13 @@ import { tableKeys } from '@/hooks/queries/utils/table-keys' const TABLE_ID = 'tbl-1' const WORKSPACE_ID = 'ws-1' +/** + * Where a paged row list actually lives. Seeding at the bare `rowsRoot` prefix would + * exercise a key no hook writes, and would keep matching a cache walk that has been + * narrowed away from the `find` sibling hanging off the same parent. + */ +const ROWS_KEY = tableKeys.infiniteRows(TABLE_ID, tableRowsParamsKey({ pageSize: 1000 })) + function setCache(key: readonly unknown[], value: unknown) { cacheStore.set(JSON.stringify(key), value) } @@ -96,7 +103,7 @@ describe('useDeleteColumn optimistic update', () => { columnWidths: { name: 200, age: 100 }, }, }) - setCache(tableKeys.rowsRoot(TABLE_ID), { + setCache(ROWS_KEY, { rows: [ { id: 'r1', data: { name: 'a', age: 1 } }, { id: 'r2', data: { name: 'b', age: 2 } }, @@ -114,9 +121,7 @@ describe('useDeleteColumn optimistic update', () => { expect(detail?.schema.columns.map((c) => c.name)).toEqual(['name']) expect(detail?.metadata.columnWidths).toEqual({ name: 200 }) - const rows = getCache<{ rows: Array<{ data: Record }> }>( - tableKeys.rowsRoot(TABLE_ID) - ) + const rows = getCache<{ rows: Array<{ data: Record }> }>(ROWS_KEY) expect(rows?.rows.every((r) => !('age' in r.data))).toBe(true) expect(rows?.rows[0]?.data).toEqual({ name: 'a' }) @@ -124,6 +129,33 @@ describe('useDeleteColumn optimistic update', () => { expect(ctx?.rowSnapshots?.length).toBeGreaterThan(0) }) + /** + * The `find` cache hangs off the same `rowsRoot` parent as the paged rows but holds + * `{matches, truncated}` — no `pages`, no `rows`. A cache walk starting at the shared + * parent reaches it and throws inside `onMutate`, rejecting the mutation before it ever + * reaches the server: search a table, dismiss the search, then edit a cell. + */ + it('survives a cached search result hanging off the shared rows prefix', async () => { + setCache(tableKeys.detail(TABLE_ID), { + id: TABLE_ID, + schema: { columns: [{ name: 'age', type: 'number' }] }, + }) + setCache(ROWS_KEY, { + rows: [{ id: 'r1', data: { age: 1 } }], + totalCount: 1, + }) + setCache(tableKeys.find(TABLE_ID, 'q'), { matches: [{ rowId: 'r1', column: 'age' }] }) + + const hook = useDeleteColumn({ workspaceId: WORKSPACE_ID, tableId: TABLE_ID }) + + await expect(hook.onMutate?.('age')).resolves.toBeDefined() + + const rows = getCache<{ rows: Array<{ data: Record }> }>(ROWS_KEY) + expect(rows?.rows[0]?.data).toEqual({}) + /** The find entry is match coordinates, not row values — it must be left untouched. */ + expect(getCache<{ matches: unknown[] }>(tableKeys.find(TABLE_ID, 'q'))?.matches).toHaveLength(1) + }) + it('rolls back schema and rows on error using snapshots', async () => { const originalDetail = { id: TABLE_ID, @@ -135,7 +167,7 @@ describe('useDeleteColumn optimistic update', () => { totalCount: 1, } setCache(tableKeys.detail(TABLE_ID), originalDetail) - setCache(tableKeys.rowsRoot(TABLE_ID), originalRows) + setCache(ROWS_KEY, originalRows) const hook = useDeleteColumn({ workspaceId: WORKSPACE_ID, tableId: TABLE_ID }) const ctx = await hook.onMutate?.('age') @@ -145,7 +177,7 @@ describe('useDeleteColumn optimistic update', () => { hook.onError?.(new Error('boom'), 'age', ctx) expect(getCache(tableKeys.detail(TABLE_ID))).toEqual(originalDetail) - expect(getCache(tableKeys.rowsRoot(TABLE_ID))).toEqual(originalRows) + expect(getCache(ROWS_KEY)).toEqual(originalRows) }) it('invalidates schema, rows, and lists in onSettled', () => { @@ -194,7 +226,7 @@ describe('useUpdateColumn optimistic update', () => { id: TABLE_ID, schema: { columns: [{ name: 'age', type: 'number' }] }, }) - setCache(tableKeys.rowsRoot(TABLE_ID), { + setCache(ROWS_KEY, { rows: [ { id: 'r1', data: { age: 30 } }, { id: 'r2', data: { age: 40 } }, @@ -207,9 +239,7 @@ describe('useUpdateColumn optimistic update', () => { // Row data is id-keyed; a rename never moves it. The stored key (`age`) // becomes the column's stamped id, so cells stay reachable via getColumnId. - const rows = getCache<{ rows: Array<{ data: Record }> }>( - tableKeys.rowsRoot(TABLE_ID) - ) + const rows = getCache<{ rows: Array<{ data: Record }> }>(ROWS_KEY) expect(rows?.rows[0]?.data).toEqual({ age: 30 }) expect(rows?.rows[1]?.data).toEqual({ age: 40 }) @@ -265,7 +295,7 @@ describe('useDeleteColumn case-insensitive row cleanup', () => { id: TABLE_ID, schema: { columns: [{ name: 'Age', type: 'number' }] }, }) - setCache(tableKeys.rowsRoot(TABLE_ID), { + setCache(ROWS_KEY, { rows: [{ id: 'r1', data: { Age: 30, name: 'a' } }], totalCount: 1, }) @@ -273,9 +303,7 @@ describe('useDeleteColumn case-insensitive row cleanup', () => { const hook = useDeleteColumn({ workspaceId: WORKSPACE_ID, tableId: TABLE_ID }) await hook.onMutate?.('age') - const rows = getCache<{ rows: Array<{ data: Record }> }>( - tableKeys.rowsRoot(TABLE_ID) - ) + const rows = getCache<{ rows: Array<{ data: Record }> }>(ROWS_KEY) expect(rows?.rows[0]?.data).toEqual({ name: 'a' }) }) }) diff --git a/apps/sim/hooks/queries/tables.ts b/apps/sim/hooks/queries/tables.ts index 6d0848b4caf..bbc3718fd2f 100644 --- a/apps/sim/hooks/queries/tables.ts +++ b/apps/sim/hooks/queries/tables.ts @@ -832,6 +832,11 @@ function withOptimisticAutoFireExec(groups: WorkflowGroup[], row: TableRow): Tab /** * Apply a row-level transformation to all cached infinite row queries for this * table. Used for cell edits where positions don't change. + * + * Walks {@link tableKeys.infiniteRowsRoot} rather than `rowsRoot`: the latter is a + * shared parent, and handing this updater a `find` entry — a flat + * {@link TableFindResult}, not pages — throws on `old.pages` inside `onMutate`, so + * the whole cell edit would reject before reaching the server. */ function patchCachedRows( queryClient: ReturnType, @@ -839,7 +844,7 @@ function patchCachedRows( patchRow: (row: TableRow) => TableRow ) { queryClient.setQueriesData>( - { queryKey: tableKeys.rowsRoot(tableId), exact: false }, + { queryKey: tableKeys.infiniteRowsRoot(tableId), exact: false }, (old) => { if (!old) return old return { @@ -2187,7 +2192,7 @@ export async function snapshotAndMutateRows( ): Promise { const scope = options?.onlyKey ? ({ queryKey: options.onlyKey, exact: true } as const) - : ({ queryKey: tableKeys.rowsRoot(tableId) } as const) + : ({ queryKey: tableKeys.infiniteRowsRoot(tableId) } as const) if (options?.cancelInFlight !== false) { await queryClient.cancelQueries(scope) } diff --git a/apps/sim/hooks/queries/utils/invalidate-usage.ts b/apps/sim/hooks/queries/utils/invalidate-usage.ts new file mode 100644 index 00000000000..2c08f009d72 --- /dev/null +++ b/apps/sim/hooks/queries/utils/invalidate-usage.ts @@ -0,0 +1,32 @@ +import type { QueryClient } from '@tanstack/react-query' +import { subscriptionKeys } from '@/hooks/queries/utils/subscription-keys' +import { workspaceUsageKeys } from '@/hooks/queries/utils/workspace-usage-keys' + +/** + * Usage is written asynchronously as a run settles, so a refetch fired on completion + * races the write and re-reads the old balance. + */ +const USAGE_SETTLE_DELAY_MS = 1000 + +/** + * Invalidates the workspace credit/usage reads after anything that moves the balance. + * Both families are keyed per workspace but derive from one billing account, so the + * family prefixes are what must refetch. + */ +export function invalidateWorkspaceUsage(queryClient: QueryClient) { + return Promise.all([ + queryClient.invalidateQueries({ queryKey: workspaceUsageKeys.creditAvailabilities() }), + queryClient.invalidateQueries({ queryKey: workspaceUsageKeys.gates() }), + ]) +} + +/** + * Refreshes the billing reads a run touches, after {@link USAGE_SETTLE_DELAY_MS}. + * Shared by workflow execution and wand generation. + */ +export function scheduleUsageRefresh(queryClient: QueryClient) { + setTimeout(() => { + void queryClient.invalidateQueries({ queryKey: subscriptionKeys.users() }) + void invalidateWorkspaceUsage(queryClient) + }, USAGE_SETTLE_DELAY_MS) +} diff --git a/apps/sim/hooks/queries/utils/optimistic-mutation.test.ts b/apps/sim/hooks/queries/utils/optimistic-mutation.test.ts new file mode 100644 index 00000000000..936b1bab173 --- /dev/null +++ b/apps/sim/hooks/queries/utils/optimistic-mutation.test.ts @@ -0,0 +1,24 @@ +/** + * @vitest-environment node + */ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { generateTempId } from '@/hooks/queries/utils/optimistic-mutation' + +describe('generateTempId', () => { + afterEach(() => { + vi.useRealTimers() + }) + + it('is unique for ids created within the same millisecond', () => { + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z')) + + const ids = new Set(Array.from({ length: 100 }, () => generateTempId('temp-folder'))) + + expect(ids.size).toBe(100) + }) + + it('keeps the prefix so callers can identify optimistic rows', () => { + expect(generateTempId('temp-folder').startsWith('temp-folder-')).toBe(true) + }) +}) diff --git a/apps/sim/hooks/queries/utils/optimistic-mutation.ts b/apps/sim/hooks/queries/utils/optimistic-mutation.ts index b734ba941eb..2af55b543af 100644 --- a/apps/sim/hooks/queries/utils/optimistic-mutation.ts +++ b/apps/sim/hooks/queries/utils/optimistic-mutation.ts @@ -1,4 +1,5 @@ import { createLogger } from '@sim/logger' +import { generateId } from '@sim/utils/id' import type { QueryClient } from '@tanstack/react-query' const logger = createLogger('OptimisticMutation') @@ -72,6 +73,12 @@ export function createOptimisticMutationHandlers( } } +/** + * Placeholder id for an optimistic row, held only until the server response + * replaces it. Uses `generateId()` rather than a timestamp so two rows created + * in the same millisecond cannot collide — a collision would make + * `replaceOptimisticEntry` overwrite both entries with one server row. + */ export function generateTempId(prefix: string): string { - return `${prefix}-${Date.now()}` + return `${prefix}-${generateId()}` } diff --git a/apps/sim/hooks/queries/utils/subscription-keys.ts b/apps/sim/hooks/queries/utils/subscription-keys.ts new file mode 100644 index 00000000000..f2a879400a4 --- /dev/null +++ b/apps/sim/hooks/queries/utils/subscription-keys.ts @@ -0,0 +1,13 @@ +/** + * React Query key factory for subscription and billing reads. Standalone so the shared + * billing invalidations can use it without closing an import cycle through the hook module. + */ +export const subscriptionKeys = { + all: ['subscription'] as const, + users: () => [...subscriptionKeys.all, 'user'] as const, + user: (includeOrg?: boolean) => [...subscriptionKeys.users(), { includeOrg }] as const, + usage: () => [...subscriptionKeys.all, 'usage'] as const, + invoicesAll: () => [...subscriptionKeys.all, 'invoices'] as const, + invoices: (context: 'user' | 'organization' = 'user', organizationId?: string) => + [...subscriptionKeys.invoicesAll(), context, organizationId ?? ''] as const, +} diff --git a/apps/sim/hooks/queries/utils/table-keys.ts b/apps/sim/hooks/queries/utils/table-keys.ts index 0721bbfb308..853103b5122 100644 --- a/apps/sim/hooks/queries/utils/table-keys.ts +++ b/apps/sim/hooks/queries/utils/table-keys.ts @@ -25,8 +25,14 @@ export const tableKeys = { exportJobs: (workspaceId?: string) => [...tableKeys.all, 'export-jobs', workspaceId ?? ''] as const, rowsRoot: (tableId: string) => [...tableKeys.detail(tableId), 'rows'] as const, + /** + * Prefix covering only the paged row lists. `rowsRoot` is a shared parent — `find` + * hangs off it holding a different shape — so anything walking the cache for row + * pages must start here. + */ + infiniteRowsRoot: (tableId: string) => [...tableKeys.rowsRoot(tableId), 'infinite'] as const, infiniteRows: (tableId: string, paramsKey: string) => - [...tableKeys.rowsRoot(tableId), 'infinite', paramsKey] as const, + [...tableKeys.infiniteRowsRoot(tableId), paramsKey] as const, rowWrites: (tableId: string) => [...tableKeys.rowsRoot(tableId), 'write'] as const, find: (tableId: string, paramsKey: string) => [...tableKeys.rowsRoot(tableId), 'find', paramsKey] as const, diff --git a/apps/sim/hooks/queries/utils/workspace-usage-keys.ts b/apps/sim/hooks/queries/utils/workspace-usage-keys.ts new file mode 100644 index 00000000000..8bb1a0f1b72 --- /dev/null +++ b/apps/sim/hooks/queries/utils/workspace-usage-keys.ts @@ -0,0 +1,12 @@ +/** + * React Query key factory for the per-workspace credit and usage-gate reads. Standalone + * for the same import-cycle reason as {@link file://./subscription-keys.ts}. + */ +export const workspaceUsageKeys = { + all: ['workspace-usage'] as const, + creditAvailabilities: () => [...workspaceUsageKeys.all, 'credit-availability'] as const, + creditAvailability: (workspaceId: string) => + [...workspaceUsageKeys.creditAvailabilities(), workspaceId] as const, + gates: () => [...workspaceUsageKeys.all, 'gate'] as const, + gate: (workspaceId: string) => [...workspaceUsageKeys.gates(), workspaceId] as const, +} diff --git a/apps/sim/hooks/queries/workspace-files.ts b/apps/sim/hooks/queries/workspace-files.ts index a236a623f2f..69c775f2dde 100644 --- a/apps/sim/hooks/queries/workspace-files.ts +++ b/apps/sim/hooks/queries/workspace-files.ts @@ -521,6 +521,13 @@ export function useCloudStorageConfigured(enabled = true) { enabled, retry: false, staleTime: CLOUD_STORAGE_CONFIGURED_STALE_TIME, + /** + * Escapes the global `retryOnMount: false`: with an infinite `staleTime` and + * `retry: false`, one transient error leaves this query errored for the tab's lifetime, + * and the upload path reads "unknown" as "not configured" — disabling cloud uploads + * until a full reload. The key is global, so navigation cannot recover it. + */ + retryOnMount: true, }) } diff --git a/apps/sim/hooks/queries/workspace-usage.test.ts b/apps/sim/hooks/queries/workspace-usage.test.ts index f468d2b07e2..011b091ae94 100644 --- a/apps/sim/hooks/queries/workspace-usage.test.ts +++ b/apps/sim/hooks/queries/workspace-usage.test.ts @@ -15,12 +15,13 @@ import { getWorkspaceCreditAvailabilityContract, getWorkspaceUsageGateContract, } from '@/lib/api/contracts/workspaces' +import { invalidateWorkspaceUsage } from '@/hooks/queries/utils/invalidate-usage' +import { workspaceUsageKeys } from '@/hooks/queries/utils/workspace-usage-keys' import { fetchWorkspaceCreditAvailability, fetchWorkspaceUsageGate, WORKSPACE_CREDIT_AVAILABILITY_STALE_TIME, WORKSPACE_USAGE_GATE_STALE_TIME, - workspaceUsageKeys, } from '@/hooks/queries/workspace-usage' describe('workspace usage gate query', () => { @@ -65,4 +66,18 @@ describe('workspace usage gate query', () => { signal, }) }) + + it('invalidates both usage families so a spend or top-up refetches them', async () => { + const invalidateQueries = vi.fn().mockResolvedValue(undefined) + const queryClient = { invalidateQueries } as unknown as Parameters< + typeof invalidateWorkspaceUsage + >[0] + + await invalidateWorkspaceUsage(queryClient) + + expect(invalidateQueries).toHaveBeenCalledWith({ + queryKey: workspaceUsageKeys.creditAvailabilities(), + }) + expect(invalidateQueries).toHaveBeenCalledWith({ queryKey: workspaceUsageKeys.gates() }) + }) }) diff --git a/apps/sim/hooks/queries/workspace-usage.ts b/apps/sim/hooks/queries/workspace-usage.ts index df3ad7e0d08..f9f46152fee 100644 --- a/apps/sim/hooks/queries/workspace-usage.ts +++ b/apps/sim/hooks/queries/workspace-usage.ts @@ -6,15 +6,7 @@ import { type WorkspaceCreditAvailability, type WorkspaceUsageGate, } from '@/lib/api/contracts/workspaces' - -export const workspaceUsageKeys = { - all: ['workspace-usage'] as const, - creditAvailabilities: () => [...workspaceUsageKeys.all, 'credit-availability'] as const, - creditAvailability: (workspaceId: string) => - [...workspaceUsageKeys.creditAvailabilities(), workspaceId] as const, - gates: () => [...workspaceUsageKeys.all, 'gate'] as const, - gate: (workspaceId: string) => [...workspaceUsageKeys.gates(), workspaceId] as const, -} +import { workspaceUsageKeys } from '@/hooks/queries/utils/workspace-usage-keys' export const WORKSPACE_CREDIT_AVAILABILITY_STALE_TIME = 30 * 1000 export const WORKSPACE_USAGE_GATE_STALE_TIME = 30 * 1000 diff --git a/apps/sim/hooks/selectors/providers/cloudwatch/selectors.test.ts b/apps/sim/hooks/selectors/providers/cloudwatch/selectors.test.ts new file mode 100644 index 00000000000..66f2ce150ac --- /dev/null +++ b/apps/sim/hooks/selectors/providers/cloudwatch/selectors.test.ts @@ -0,0 +1,30 @@ +/** + * @vitest-environment node + */ +import { describe, expect, it } from 'vitest' +import { cloudwatchSelectors } from '@/hooks/selectors/providers/cloudwatch/selectors' +import type { SelectorQueryArgs } from '@/hooks/selectors/types' + +const AWS_CONTEXT: SelectorQueryArgs['context'] = { + awsAccessKeyId: 'AKIA', + awsSecretAccessKey: 'secret', + awsRegion: 'us-east-1', + logGroupName: '/aws/lambda/fn', +} + +describe('cloudwatch selector query keys', () => { + it.each([['cloudwatch.logGroups' as const], ['cloudwatch.logStreams' as const]])( + '%s scopes its key by search, which fetchList forwards as `prefix`', + (key) => { + const definition = cloudwatchSelectors[key] + const base: SelectorQueryArgs = { key, context: AWS_CONTEXT } + + const noSearch = definition.getQueryKey(base) + const withSearch = definition.getQueryKey({ ...base, search: 'api' }) + const otherSearch = definition.getQueryKey({ ...base, search: 'worker' }) + + expect(withSearch).not.toEqual(noSearch) + expect(withSearch).not.toEqual(otherSearch) + } + ) +}) diff --git a/apps/sim/hooks/selectors/providers/cloudwatch/selectors.ts b/apps/sim/hooks/selectors/providers/cloudwatch/selectors.ts index 065db6680d0..6d7b7480183 100644 --- a/apps/sim/hooks/selectors/providers/cloudwatch/selectors.ts +++ b/apps/sim/hooks/selectors/providers/cloudwatch/selectors.ts @@ -20,11 +20,12 @@ export const cloudwatchSelectors = { key: 'cloudwatch.logGroups', contracts: [selectorContracts.cloudwatchLogGroupsSelectorContract], staleTime: SELECTOR_STALE, - getQueryKey: ({ context }: SelectorQueryArgs) => [ + getQueryKey: ({ context, search }: SelectorQueryArgs) => [ 'selectors', 'cloudwatch.logGroups', context.awsAccessKeyId ?? 'none', context.awsRegion ?? 'none', + search ?? '', ], enabled: ({ context }) => Boolean(context.awsAccessKeyId && context.awsSecretAccessKey && context.awsRegion), @@ -51,12 +52,13 @@ export const cloudwatchSelectors = { key: 'cloudwatch.logStreams', contracts: [selectorContracts.cloudwatchLogStreamsSelectorContract], staleTime: SELECTOR_STALE, - getQueryKey: ({ context }: SelectorQueryArgs) => [ + getQueryKey: ({ context, search }: SelectorQueryArgs) => [ 'selectors', 'cloudwatch.logStreams', context.awsAccessKeyId ?? 'none', context.awsRegion ?? 'none', context.logGroupName ?? 'none', + search ?? '', ], enabled: ({ context }) => Boolean( diff --git a/apps/sim/hooks/selectors/use-selector-query.ts b/apps/sim/hooks/selectors/use-selector-query.ts index ea95d7e879d..213afcdab91 100644 --- a/apps/sim/hooks/selectors/use-selector-query.ts +++ b/apps/sim/hooks/selectors/use-selector-query.ts @@ -70,7 +70,12 @@ export function useSelectorOptions( context: args.context, search: args.search, } - const isEnabled = args.enabled ?? (definition.enabled ? definition.enabled(queryArgs) : true) + /** + * `definition.enabled` mirrors the preconditions the definition's own fetchers assert, so + * it is a hard precondition for the list, not a default a caller may replace. A caller's + * `enabled` only narrows — widening would run a fetch guaranteed to reject and cache it. + */ + const isEnabled = args.enabled !== false && (definition.enabled?.(queryArgs) ?? true) const supportsPagination = Boolean(definition.fetchPage) const flatQuery = useQuery({ diff --git a/apps/sim/lib/api/contracts/oauth-connections.ts b/apps/sim/lib/api/contracts/oauth-connections.ts index 7effa93f300..c9c4951efd2 100644 --- a/apps/sim/lib/api/contracts/oauth-connections.ts +++ b/apps/sim/lib/api/contracts/oauth-connections.ts @@ -62,6 +62,9 @@ const trelloCallbackQuerySchema = z }) .passthrough() +/** Google domain-wide-delegation subject. Also applied by in-process credential callers. */ +export const impersonateEmailSchema = z.string().email() + export const oauthTokenRequestBodySchema = z .object({ credentialId: z.string().min(1).optional(), @@ -69,7 +72,7 @@ export const oauthTokenRequestBodySchema = z providerId: z.string().min(1).optional(), workflowId: z.string().min(1).nullish(), scopes: z.array(z.string()).optional(), - impersonateEmail: z.string().email().optional(), + impersonateEmail: impersonateEmailSchema.optional(), }) .refine( (data) => data.credentialId || (data.credentialAccountUserId && data.providerId), @@ -99,6 +102,9 @@ const oauthTokenResponseSchema = z.object({ authStyle: z.enum(['x-api-token']).optional(), }) +/** Token material a resolved credential yields, on the wire and in-process alike. */ +export type OAuthTokenResponse = z.output + export const oauthTokenGetContract = defineRouteContract({ method: 'GET', path: '/api/auth/oauth/token', diff --git a/apps/sim/lib/auth/credential-access.test.ts b/apps/sim/lib/auth/credential-access.test.ts index 4a30c270b64..3f0f1022edb 100644 --- a/apps/sim/lib/auth/credential-access.test.ts +++ b/apps/sim/lib/auth/credential-access.test.ts @@ -23,7 +23,7 @@ vi.mock('@/lib/workspaces/permissions/utils', () => ({ resolveWorkspaceAccess: mockResolveWorkspaceAccess, })) -import { authorizeCredentialUse } from '@/lib/auth/credential-access' +import { authorizeCredentialUse, authorizeCredentialUseForAuth } from '@/lib/auth/credential-access' afterAll(resetDbChainMock) @@ -235,4 +235,41 @@ describe('authorizeCredentialUse', () => { expect(result.error).toBe('Credential not found') }) }) + + /** + * The in-process tool executor synthesizes the AuthResult an internal JWT + * would have produced instead of minting one and POSTing to ourselves, so the + * subject-less case must still fail closed here. + */ + describe('authorizeCredentialUseForAuth', () => { + it('fails closed when the authenticated caller carries no user id', async () => { + const result = await authorizeCredentialUseForAuth( + { success: true, authType: 'internal_jwt' }, + { credentialId: ACCOUNT_ID } + ) + + expect(result.ok).toBe(false) + expect(result.error).toBe('Authentication required') + }) + + it('fails closed when authentication did not succeed', async () => { + const result = await authorizeCredentialUseForAuth( + { success: false, error: 'Unauthorized' }, + { credentialId: ACCOUNT_ID } + ) + + expect(result.ok).toBe(false) + expect(result.error).toBe('Unauthorized') + }) + + it('rejects an asserted caller that does not match the internal token subject', async () => { + const result = await authorizeCredentialUseForAuth( + { success: true, userId: OWNER, authType: 'internal_jwt' }, + { credentialId: ACCOUNT_ID, callerUserId: 'someone-else' } + ) + + expect(result.ok).toBe(false) + expect(result.error).toBe('Caller user does not match internal token subject') + }) + }) }) diff --git a/apps/sim/lib/auth/credential-access.ts b/apps/sim/lib/auth/credential-access.ts index 125511b67a8..719cd30f8ca 100644 --- a/apps/sim/lib/auth/credential-access.ts +++ b/apps/sim/lib/auth/credential-access.ts @@ -2,7 +2,7 @@ import { db } from '@sim/db' import { account, credential, workflow as workflowTable } from '@sim/db/schema' import { and, asc, eq } from 'drizzle-orm' import type { NextRequest } from 'next/server' -import { AuthType, checkSessionOrInternalAuth } from '@/lib/auth/hybrid' +import { type AuthResult, AuthType, checkSessionOrInternalAuth } from '@/lib/auth/hybrid' import { type CredentialActorContext, canUseCredential, @@ -59,11 +59,27 @@ export async function authorizeCredentialUse( callerUserId?: string } ): Promise { - const { credentialId, workflowId, requireWorkflowIdForInternal = true, callerUserId } = params - const auth = await checkSessionOrInternalAuth(request, { - requireWorkflowId: requireWorkflowIdForInternal, + requireWorkflowId: params.requireWorkflowIdForInternal ?? true, }) + return authorizeCredentialUseForAuth(auth, params) +} + +/** + * Credential authorization for an already-authenticated caller. + * {@link authorizeCredentialUse} is the HTTP wrapper; in-process callers build the same + * {@link AuthResult} directly, so both paths run one identical rule. + */ +export async function authorizeCredentialUseForAuth( + auth: AuthResult, + params: { + credentialId: string + workflowId?: string + callerUserId?: string + } +): Promise { + const { credentialId, workflowId, callerUserId } = params + if (!auth.success || !auth.userId) { return { ok: false, error: auth.error || 'Authentication required' } } diff --git a/apps/sim/lib/billing/core/subscription.ts b/apps/sim/lib/billing/core/subscription.ts index 44345553e4d..d697c9fd7fa 100644 --- a/apps/sim/lib/billing/core/subscription.ts +++ b/apps/sim/lib/billing/core/subscription.ts @@ -1,3 +1,4 @@ +import { cache } from 'react' import { db } from '@sim/db' import { member, organization, subscription, user } from '@sim/db/schema' import { createLogger } from '@sim/logger' @@ -429,11 +430,7 @@ export async function isEnterpriseOrgAdminOrOwner(userId: string): Promise { +async function resolveOrganizationEnterprisePlan(organizationId: string): Promise { try { if (!isBillingEnabled) { return true @@ -456,6 +453,15 @@ export async function isOrganizationOnEnterprisePlan(organizationId: string): Pr } } +/** + * Check if an organization has an enterprise plan + * Used for Access Control (Permission Groups) feature gating + * + * Request-memoized: a settings render gates several sections on the same + * organization's plan, and it cannot change mid-render. + */ +export const isOrganizationOnEnterprisePlan = cache(resolveOrganizationEnterprisePlan) + /** * Entitlement for a single org-scoped enterprise feature. * @@ -612,10 +618,15 @@ async function hasWorkspaceTierAccess( * Whether the workspace's payer is on a usable Max-or-Enterprise subscription. * Shared by the inbox (Sim Mailer), live sync, and custom sandboxes, which all * sit on the same entitlement tier. + * + * Request-memoized: these features are gated side by side on one settings render, + * each otherwise repeating the identical workspace and subscription reads. The + * per-feature deployment and env short-circuits live in the wrappers and still run + * per call. */ -async function hasMaxTierWorkspaceAccess(workspaceId: string): Promise { - return hasWorkspaceTierAccess(workspaceId, isMaxTier) -} +const hasMaxTierWorkspaceAccess = cache( + (workspaceId: string): Promise => hasWorkspaceTierAccess(workspaceId, isMaxTier) +) /** * Check whether a workspace is entitled to the inbox (Sim Mailer) feature. diff --git a/apps/sim/lib/oauth/token-resolution.test.ts b/apps/sim/lib/oauth/token-resolution.test.ts new file mode 100644 index 00000000000..27d239091ae --- /dev/null +++ b/apps/sim/lib/oauth/token-resolution.test.ts @@ -0,0 +1,214 @@ +/** + * @vitest-environment node + */ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const { + mockAuthorizeCredentialUseForAuth, + mockGetCredential, + mockRecordAudit, + mockRefreshTokenIfNeeded, + mockResolveOAuthAccountId, + mockResolveServiceAccountToken, +} = vi.hoisted(() => ({ + mockAuthorizeCredentialUseForAuth: vi.fn(), + mockGetCredential: vi.fn(), + mockRecordAudit: vi.fn(), + mockRefreshTokenIfNeeded: vi.fn(), + mockResolveOAuthAccountId: vi.fn(), + mockResolveServiceAccountToken: vi.fn(), +})) + +vi.mock('@sim/audit', () => ({ + AuditAction: { CREDENTIAL_ACCESSED: 'credential.accessed' }, + AuditResourceType: { CREDENTIAL: 'credential' }, + recordAudit: mockRecordAudit, +})) + +vi.mock('@/lib/auth/credential-access', () => ({ + authorizeCredentialUseForAuth: mockAuthorizeCredentialUseForAuth, +})) + +vi.mock('@/lib/oauth/credential-service', () => ({ + getCredential: mockGetCredential, + refreshTokenIfNeeded: mockRefreshTokenIfNeeded, + resolveOAuthAccountId: mockResolveOAuthAccountId, + resolveServiceAccountToken: mockResolveServiceAccountToken, +})) + +vi.mock('@/lib/posthog/server', () => ({ + captureServerEvent: vi.fn(), +})) + +import { TokenServiceAccountValidationError } from '@/lib/credentials/token-service-accounts/errors' +import { resolveCredentialToken } from '@/lib/oauth/token-resolution' + +const INTERNAL_AUTH = { success: true, userId: 'user-1', authType: 'internal_jwt' } as const + +describe('resolveCredentialToken', () => { + beforeEach(() => { + vi.clearAllMocks() + mockResolveOAuthAccountId.mockResolvedValue(null) + }) + + it('fails closed when the credential is not authorized', async () => { + mockAuthorizeCredentialUseForAuth.mockResolvedValue({ + ok: false, + error: 'You do not have access to this credential.', + }) + + const result = await resolveCredentialToken(INTERNAL_AUTH, { + requestId: 'req-1', + credentialId: 'cred-1', + }) + + expect(result).toEqual({ + ok: false, + status: 403, + error: 'You do not have access to this credential.', + }) + expect(mockGetCredential).not.toHaveBeenCalled() + expect(mockRefreshTokenIfNeeded).not.toHaveBeenCalled() + expect(mockRecordAudit).not.toHaveBeenCalled() + }) + + it('fails closed when the caller carries no user id', async () => { + mockAuthorizeCredentialUseForAuth.mockResolvedValue({ + ok: false, + error: 'Authentication required', + }) + + const result = await resolveCredentialToken( + { success: true, authType: 'internal_jwt' }, + { requestId: 'req-1', credentialId: 'cred-1' } + ) + + expect(result).toEqual({ ok: false, status: 403, error: 'Authentication required' }) + }) + + it('refreshes the token, records the access trail, and returns the payload', async () => { + mockAuthorizeCredentialUseForAuth.mockResolvedValue({ + ok: true, + requesterUserId: 'user-1', + credentialOwnerUserId: 'owner-1', + workspaceId: 'ws-1', + resolvedCredentialId: 'account-1', + }) + mockGetCredential.mockResolvedValue({ + providerId: 'google', + idToken: 'id-token', + scope: 'https://www.googleapis.com/auth/gmail.send', + }) + mockRefreshTokenIfNeeded.mockResolvedValue({ accessToken: 'fresh', refreshed: true }) + + const result = await resolveCredentialToken(INTERNAL_AUTH, { + requestId: 'req-1', + credentialId: 'cred-1', + workflowId: 'wf-1', + }) + + expect(result).toEqual({ ok: true, token: { accessToken: 'fresh', idToken: 'id-token' } }) + expect(mockGetCredential).toHaveBeenCalledWith('req-1', 'account-1', 'owner-1') + expect(mockRefreshTokenIfNeeded).toHaveBeenCalled() + expect(mockRecordAudit).toHaveBeenCalledWith( + expect.objectContaining({ + actorId: 'user-1', + workspaceId: 'ws-1', + resourceId: 'account-1', + action: 'credential.accessed', + }) + ) + }) + + it('returns 404 when the authorized credential is missing', async () => { + mockAuthorizeCredentialUseForAuth.mockResolvedValue({ + ok: true, + requesterUserId: 'user-1', + credentialOwnerUserId: 'owner-1', + }) + mockGetCredential.mockResolvedValue(undefined) + + const result = await resolveCredentialToken(INTERNAL_AUTH, { + requestId: 'req-1', + credentialId: 'cred-1', + }) + + expect(result).toEqual({ ok: false, status: 404, error: 'Credential not found' }) + }) + + it('reports a failed refresh as 401 without recording access', async () => { + mockAuthorizeCredentialUseForAuth.mockResolvedValue({ + ok: true, + requesterUserId: 'user-1', + credentialOwnerUserId: 'owner-1', + }) + mockGetCredential.mockResolvedValue({ providerId: 'google' }) + mockRefreshTokenIfNeeded.mockRejectedValue(new Error('refresh token revoked')) + + const result = await resolveCredentialToken(INTERNAL_AUTH, { + requestId: 'req-1', + credentialId: 'cred-1', + }) + + expect(result).toEqual({ ok: false, status: 401, error: 'Failed to refresh access token' }) + expect(mockRecordAudit).not.toHaveBeenCalled() + }) + + it('authorizes service-account credentials before minting a token', async () => { + mockResolveOAuthAccountId.mockResolvedValue({ + credentialType: 'service_account', + credentialId: 'sa-1', + providerId: 'google', + workspaceId: 'ws-1', + accountId: '', + usedCredentialTable: true, + }) + mockAuthorizeCredentialUseForAuth.mockResolvedValue({ ok: false, error: 'Unauthorized' }) + + const result = await resolveCredentialToken(INTERNAL_AUTH, { + requestId: 'req-1', + credentialId: 'cred-1', + }) + + expect(result).toEqual({ ok: false, status: 403, error: 'Unauthorized' }) + expect(mockResolveServiceAccountToken).not.toHaveBeenCalled() + }) + + it('surfaces the classified service-account failure code', async () => { + mockResolveOAuthAccountId.mockResolvedValue({ + credentialType: 'service_account', + credentialId: 'sa-1', + providerId: 'atlassian', + workspaceId: 'ws-1', + accountId: '', + usedCredentialTable: true, + }) + mockAuthorizeCredentialUseForAuth.mockResolvedValue({ ok: true, requesterUserId: 'user-1' }) + mockResolveServiceAccountToken.mockRejectedValue( + new TokenServiceAccountValidationError('invalid_credentials', 401) + ) + + const result = await resolveCredentialToken(INTERNAL_AUTH, { + requestId: 'req-1', + credentialId: 'cred-1', + }) + + expect(result).toEqual({ + ok: false, + status: 401, + code: 'invalid_credentials', + error: 'Credential rejected by the provider — reconnect the credential', + }) + }) + + it('rejects a malformed impersonation subject before touching the credential', async () => { + const result = await resolveCredentialToken(INTERNAL_AUTH, { + requestId: 'req-1', + credentialId: 'cred-1', + impersonateEmail: 'not-an-email', + }) + + expect(result.ok).toBe(false) + expect(mockAuthorizeCredentialUseForAuth).not.toHaveBeenCalled() + }) +}) diff --git a/apps/sim/lib/oauth/token-resolution.ts b/apps/sim/lib/oauth/token-resolution.ts new file mode 100644 index 00000000000..4ee18823f8b --- /dev/null +++ b/apps/sim/lib/oauth/token-resolution.ts @@ -0,0 +1,297 @@ +import { AuditAction, AuditResourceType, recordAudit } from '@sim/audit' +import { createLogger } from '@sim/logger' +import { + impersonateEmailSchema, + type OAuthTokenResponse, +} from '@/lib/api/contracts/oauth-connections' +import { authorizeCredentialUseForAuth } from '@/lib/auth/credential-access' +import type { AuthResult } from '@/lib/auth/hybrid' +import { TokenServiceAccountValidationError } from '@/lib/credentials/token-service-accounts/errors' +import { + getCredential, + refreshTokenIfNeeded, + resolveOAuthAccountId, + resolveServiceAccountToken, +} from '@/lib/oauth/credential-service' +import { extractSalesforceInstanceUrl, isSalesforceOAuthProviderId } from '@/lib/oauth/salesforce' +import { captureServerEvent } from '@/lib/posthog/server' +import { extractZohoDeskBaseFromScope } from '@/tools/zoho_desk/host-allowlist' + +const logger = createLogger('OAuthTokenResolution') + +/** + * Duck type of the inbound request, used only so audit rows record IP and user agent. + * In-process callers omit it. + */ +export interface CredentialAuditRequest { + headers: { get(name: string): string | null } +} + +/** Token material a resolved credential yields; taken from the contract so it cannot drift. */ +export type CredentialTokenPayload = OAuthTokenResponse + +export interface ResolveCredentialTokenInput { + /** Correlation id used by the credential service's own logging. */ + requestId: string + credentialId?: string + workflowId?: string + /** Canonical provider scopes, used only by service-account token minting. */ + scopes?: string[] + /** Google domain-wide-delegation subject for service-account credentials. */ + impersonateEmail?: string + /** + * Asserted acting user. When the caller authenticated with an internal JWT it + * must equal the token subject, so a forged assertion cannot widen access. + */ + callerUserId?: string + auditRequest?: CredentialAuditRequest +} + +export type ResolveCredentialTokenResult = + | { ok: true; token: CredentialTokenPayload } + | { ok: false; status: number; error: string; code?: string } + +/** + * Emits the semantic "credential used" trail for one resolved credential. + * Both the audit row and the analytics event are fire-and-forget. + */ +function recordCredentialAccess(params: { + actorId: string + workspaceId: string | null + resourceId: string + providerId: string | null | undefined + credentialType: 'oauth' | 'service_account' + auditRequest?: CredentialAuditRequest +}): void { + const { actorId, workspaceId, resourceId, providerId, credentialType } = params + recordAudit({ + workspaceId, + actorId, + action: AuditAction.CREDENTIAL_ACCESSED, + resourceType: AuditResourceType.CREDENTIAL, + resourceId, + description: `Accessed ${credentialType === 'oauth' ? 'OAuth' : 'service account'} credential for provider ${providerId ?? 'unknown'}`, + metadata: { + provider: providerId, + credentialType, + }, + request: params.auditRequest, + }) + captureServerEvent( + actorId, + 'credential_used', + { + credential_type: credentialType, + provider_id: providerId ?? 'unknown', + ...(workspaceId ? { workspace_id: workspaceId } : {}), + }, + workspaceId ? { groups: { workspace: workspaceId } } : undefined + ) +} + +/** + * Projects a stored OAuth credential plus its access token into the wire payload. + * Provider hosts come out of the scope string through shared allowlisted helpers, never a + * local regex — these values are injected into tool calls that carry the token. + */ +function buildOAuthTokenPayload( + credential: { providerId: string; scope?: string | null; idToken?: string | null }, + accessToken: string +): CredentialTokenPayload { + const instanceUrl = isSalesforceOAuthProviderId(credential.providerId) + ? extractSalesforceInstanceUrl(credential.scope ?? undefined) + : undefined + + let apiDomain: string | undefined + if (credential.providerId === 'zoho-desk' && credential.scope) { + apiDomain = extractZohoDeskBaseFromScope(credential.scope) + } + + return { + accessToken, + idToken: credential.idToken || undefined, + ...(instanceUrl && { instanceUrl }), + ...(apiDomain && { apiDomain }), + } +} + +/** + * Refreshes an authorized OAuth credential, records its access trail, and + * projects the token payload. Shared by every surface that has already + * authorized the credential and loaded it. + */ +export async function completeOAuthCredentialToken(params: { + requestId: string + credential: { providerId: string; scope?: string | null; idToken?: string | null } + resolvedCredentialId: string + actorId?: string + workspaceId: string | null + auditRequest?: CredentialAuditRequest +}): Promise { + const { requestId, credential, resolvedCredentialId, actorId, workspaceId, auditRequest } = params + try { + const { accessToken } = await refreshTokenIfNeeded(requestId, credential, resolvedCredentialId) + + if (actorId) { + recordCredentialAccess({ + actorId, + workspaceId, + resourceId: resolvedCredentialId, + providerId: credential.providerId, + credentialType: 'oauth', + auditRequest, + }) + } + + return { ok: true, token: buildOAuthTokenPayload(credential, accessToken) } + } catch (error) { + logger.error(`[${requestId}] Failed to refresh access token:`, error) + return { ok: false, status: 401, error: 'Failed to refresh access token' } + } +} + +/** + * Authorized application operation behind `POST /api/auth/oauth/token`. Every surface that + * needs a credential token — the route and the in-process tool executor — goes through + * here, so authorization, refresh, and audit cannot drift between them. + * + * @param auth Result of authenticating the caller (session or internal JWT). + */ +export async function resolveCredentialToken( + auth: AuthResult, + input: ResolveCredentialTokenInput +): Promise { + const { + requestId, + credentialId, + workflowId, + scopes, + impersonateEmail, + callerUserId, + auditRequest, + } = input + + try { + if (!credentialId) { + return { ok: false, status: 400, error: 'Credential ID is required' } + } + if ( + impersonateEmail !== undefined && + !impersonateEmailSchema.safeParse(impersonateEmail).success + ) { + return { ok: false, status: 400, error: 'impersonateEmail must be a valid email address' } + } + + /** + * Both branches below authorize with the same arguments, and neither read depends + * on the other, so they resolve together — this runs per credentialed tool call. + */ + const [resolved, authz] = await Promise.all([ + resolveOAuthAccountId(credentialId), + authorizeCredentialUseForAuth(auth, { credentialId, workflowId, callerUserId }), + ]) + + if (resolved?.credentialType === 'service_account' && resolved.credentialId) { + if (!authz.ok) { + return { ok: false, status: 403, error: authz.error || 'Unauthorized' } + } + + const saActorId = authz.requesterUserId + const saWorkspaceId = resolved.workspaceId ?? authz.workspaceId ?? null + + try { + const result = await resolveServiceAccountToken( + resolved.credentialId, + resolved.providerId, + scopes ?? [], + impersonateEmail + ) + + if (saActorId) { + recordCredentialAccess({ + actorId: saActorId, + workspaceId: saWorkspaceId, + resourceId: resolved.credentialId, + providerId: resolved.providerId, + credentialType: 'service_account', + auditRequest, + }) + } + + return { + ok: true, + token: { + accessToken: result.accessToken, + cloudId: result.cloudId, + domain: result.domain, + instanceUrl: result.instanceUrl, + apiDomain: result.apiDomain, + authStyle: result.authStyle, + }, + } + } catch (error) { + logger.error(`[${requestId}] Service account token error:`, error) + if (error instanceof TokenServiceAccountValidationError) { + // Classified provider outages are infra failures, not bad credentials. + if (error.code === 'provider_unavailable') { + return { + ok: false, + status: 502, + error: 'Credential provider is temporarily unavailable', + } + } + // A stored host that no longer resolves is a configuration failure — + // surface the code so runtime consumers can say "check the host" + // instead of a generic auth error. + if (error.code === 'site_not_found') { + return { + ok: false, + status: 400, + code: error.code, + error: 'Credential host not found — reconnect the credential with a valid host', + } + } + // A revoked/rotated-away or misconfigured stored secret — surface the + // code so runtime consumers can prompt to reconnect the credential + // rather than showing a generic auth failure. + if (error.code === 'invalid_credentials') { + return { + ok: false, + status: 401, + code: error.code, + error: 'Credential rejected by the provider — reconnect the credential', + } + } + } + return { ok: false, status: 401, error: 'Failed to get service account token' } + } + } + + if (!authz.ok || !authz.credentialOwnerUserId) { + return { ok: false, status: 403, error: authz.error || 'Unauthorized' } + } + + const resolvedCredentialId = authz.resolvedCredentialId || credentialId + const credential = await getCredential( + requestId, + resolvedCredentialId, + authz.credentialOwnerUserId + ) + + if (!credential) { + return { ok: false, status: 404, error: 'Credential not found' } + } + + return completeOAuthCredentialToken({ + requestId, + credential, + resolvedCredentialId, + actorId: authz.requesterUserId, + workspaceId: authz.workspaceId ?? null, + auditRequest, + }) + } catch (error) { + logger.error(`[${requestId}] Error getting access token`, error) + return { ok: false, status: 500, error: 'Internal server error' } + } +} diff --git a/apps/sim/lib/permissions/super-user.ts b/apps/sim/lib/permissions/super-user.ts index 597ca135e4c..0e17b8b3afa 100644 --- a/apps/sim/lib/permissions/super-user.ts +++ b/apps/sim/lib/permissions/super-user.ts @@ -1,3 +1,4 @@ +import { cache } from 'react' import { db, dbReplica } from '@sim/db' import { settings, user } from '@sim/db/schema' import { eq } from 'drizzle-orm' @@ -41,8 +42,11 @@ export async function verifyEffectiveSuperUser(userId: string): Promise<{ * served from the replica: this gates features, not security-critical auth, so it * tolerates the replica's bounded staleness (admin role rarely changes). Falls back * to the primary when no replica is configured. + * + * Request-memoized: an account-settings render checks the same viewer in both the + * layout and the page. */ -export async function isPlatformAdmin(userId: string): Promise { +export const isPlatformAdmin = cache(async (userId: string): Promise => { const [row] = await dbReplica .select({ role: user.role }) .from(user) @@ -50,4 +54,4 @@ export async function isPlatformAdmin(userId: string): Promise { .limit(1) return row?.role === 'admin' -} +}) diff --git a/apps/sim/lib/uploads/contexts/workspace/workspace-file-manager.ts b/apps/sim/lib/uploads/contexts/workspace/workspace-file-manager.ts index cc83cc5e98d..fad6d9750a9 100644 --- a/apps/sim/lib/uploads/contexts/workspace/workspace-file-manager.ts +++ b/apps/sim/lib/uploads/contexts/workspace/workspace-file-manager.ts @@ -161,6 +161,12 @@ interface ListWorkspaceFilesOptions { hydrateFolderPaths?: boolean /** Propagate storage errors when an incomplete list would be unsafe. */ throwOnError?: boolean + /** + * Row cap for callers that only need to know whether the workspace fits a budget. + * The result is a prefix of the full list, so a caller that reads it as "the + * workspace's files" must not set this. + */ + limit?: number } /** @@ -223,7 +229,9 @@ interface WorkspaceFileMetadataInsert { size: number } -function workspaceFileSize(file: typeof workspaceFiles.$inferSelect): number { +function workspaceFileSize( + file: Pick +): number { return file.sizeBytes ?? file.size } @@ -1007,7 +1015,7 @@ export async function fileExistsInWorkspace( } function mapWorkspaceFileRecord( - file: typeof workspaceFiles.$inferSelect, + file: WorkspaceFileListRow, workspaceId: string, folderPaths: Map ): WorkspaceFileRecord { @@ -1144,9 +1152,37 @@ function workspaceFileScopeCondition(workspaceId: string, scope: WorkspaceFileSc : and(...base, isNull(workspaceFiles.deletedAt)) } +/** + * The columns {@link mapWorkspaceFileRecord} reads. These list reads are workspace-wide, + * so `select()` would ship five unprojected columns for every row of the scan. + */ +const workspaceFileListColumns = { + id: workspaceFiles.id, + key: workspaceFiles.key, + userId: workspaceFiles.userId, + workspaceId: workspaceFiles.workspaceId, + folderId: workspaceFiles.folderId, + originalName: workspaceFiles.originalName, + contentType: workspaceFiles.contentType, + size: workspaceFiles.size, + sizeBytes: workspaceFiles.sizeBytes, + width: workspaceFiles.width, + height: workspaceFiles.height, + deletedAt: workspaceFiles.deletedAt, + uploadedAt: workspaceFiles.uploadedAt, + updatedAt: workspaceFiles.updatedAt, + contentUpdatedAt: workspaceFiles.contentUpdatedAt, +} as const + +/** A row carrying exactly the columns {@link mapWorkspaceFileRecord} needs; a full row satisfies it. */ +type WorkspaceFileListRow = Pick< + typeof workspaceFiles.$inferSelect, + keyof typeof workspaceFileListColumns +> + /** Resolves `folderPath` for a page of rows, reading the folder tree only if any row needs it. */ async function hydrateWorkspaceFilePaths( - files: (typeof workspaceFiles.$inferSelect)[], + files: WorkspaceFileListRow[], workspaceId: string, options?: { folders?: WorkspaceFileFolderRecord[]; hydrateFolderPaths?: boolean } ): Promise { @@ -1167,12 +1203,13 @@ export async function listWorkspaceFiles( options?: ListWorkspaceFilesOptions ): Promise { try { - const { scope = 'active' } = options ?? {} - const files = await db - .select() + const { scope = 'active', limit } = options ?? {} + const query = db + .select(workspaceFileListColumns) .from(workspaceFiles) .where(workspaceFileScopeCondition(workspaceId, scope)) .orderBy(workspaceFiles.uploadedAt) + const files = await (limit === undefined ? query : query.limit(limit)) return hydrateWorkspaceFilePaths(files, workspaceId, options) } catch (error) { @@ -1265,7 +1302,7 @@ export async function queryWorkspaceFiles( ] const rows = await db - .select() + .select(workspaceFileListColumns) .from(workspaceFiles) .where(and(...conditions)) .orderBy(...listOrderBy(keysetColumns(keys), sortOrder)) diff --git a/apps/sim/lib/uploads/contexts/workspace/workspace-file-query.test.ts b/apps/sim/lib/uploads/contexts/workspace/workspace-file-query.test.ts index f4c81ef8bc3..a0fa4782194 100644 --- a/apps/sim/lib/uploads/contexts/workspace/workspace-file-query.test.ts +++ b/apps/sim/lib/uploads/contexts/workspace/workspace-file-query.test.ts @@ -46,7 +46,10 @@ vi.mock('@/lib/uploads/contexts/workspace/workspace-file-folder-manager', () => resolveWorkspaceFileFolderTarget: vi.fn(async () => null), })) -import { queryWorkspaceFiles } from '@/lib/uploads/contexts/workspace/workspace-file-manager' +import { + listWorkspaceFiles, + queryWorkspaceFiles, +} from '@/lib/uploads/contexts/workspace/workspace-file-manager' const WS = 'workspace-1' @@ -205,3 +208,52 @@ describe('queryWorkspaceFiles', () => { ).rejects.toMatchObject({ code: 'validation' }) }) }) + +/** + * `listWorkspaceFiles` materializes a whole scope, so what it reads per row is + * multiplied by the size of the workspace. These assertions pin the two ways that + * stays bounded: the projection, and the optional row cap its one budgeted caller + * (the workspace layout's server seed) uses to detect an oversized workspace. + */ +describe('listWorkspaceFiles', () => { + beforeEach(() => { + vi.clearAllMocks() + resetDbChainMock() + }) + + const lastProjection = () => + Object.keys((dbChainMockFns.select.mock.calls.at(-1)?.[0] ?? {}) as Record) + + it('projects only the columns the record mapper reads', async () => { + queueTableRows(schemaMock.workspaceFiles, [buildRow()]) + + await listWorkspaceFiles(WS) + + const projected = lastProjection() + expect(projected).toEqual( + expect.arrayContaining(['id', 'key', 'originalName', 'sizeBytes', 'contentUpdatedAt']) + ) + /** Columns no reader projects; `select()` would ship them for every row of the scan. */ + expect(projected).not.toContain('context') + expect(projected).not.toContain('chatId') + expect(projected).not.toContain('messageId') + expect(projected).not.toContain('displayName') + expect(projected).not.toContain('secretProvenanceVersion') + }) + + it('reads the whole scope when no cap is given', async () => { + queueTableRows(schemaMock.workspaceFiles, [buildRow()]) + + await listWorkspaceFiles(WS) + + expect(dbChainMockFns.limit).not.toHaveBeenCalled() + }) + + it('caps the rows read when the caller only needs to fit a budget', async () => { + queueTableRows(schemaMock.workspaceFiles, [buildRow()]) + + await listWorkspaceFiles(WS, { limit: 2 }) + + expect(dbChainMockFns.limit).toHaveBeenCalledWith(2) + }) +}) diff --git a/apps/sim/lib/workflows/persistence/save-normalized-state.test.ts b/apps/sim/lib/workflows/persistence/save-normalized-state.test.ts new file mode 100644 index 00000000000..9240f827fc9 --- /dev/null +++ b/apps/sim/lib/workflows/persistence/save-normalized-state.test.ts @@ -0,0 +1,83 @@ +/** + * @vitest-environment node + */ +import { describe, expect, it } from 'vitest' +import { parseWorkflowStateForPersistence } from '@/lib/workflows/persistence/save-normalized-state' + +/** + * A checkpoint blob as the revert route builds it: JSONB-derived blocks and edges, plus a + * real `Date` for `deployedAt`. + */ +function checkpointState(overrides?: Record) { + return { + blocks: { + 'block-1': { + id: 'block-1', + type: 'starter', + name: 'Start', + position: { x: 0, y: 0 }, + subBlocks: {}, + outputs: {}, + enabled: true, + }, + }, + edges: [], + loops: {}, + parallels: {}, + isDeployed: false, + lastSaved: 1_754_000_000_000, + ...overrides, + } +} + +describe('parseWorkflowStateForPersistence', () => { + /** + * The revert used to reach this schema by POSTing the blob over HTTP, so every value + * arrived JSON-serialized. In-process the blob keeps its runtime types. Both forms must + * parse identically, or a checkpoint that reverted before would start failing. + */ + it('accepts a Date for deployedAt exactly as it accepted the serialized string', () => { + const deployedAt = new Date('2026-01-02T03:04:05.678Z') + + const fromDate = parseWorkflowStateForPersistence(checkpointState({ deployedAt })) + const overTheWire = parseWorkflowStateForPersistence( + JSON.parse(JSON.stringify(checkpointState({ deployedAt }))) + ) + + expect(fromDate.success).toBe(true) + expect(overTheWire.success).toBe(true) + expect(fromDate.data?.deployedAt).toEqual(deployedAt) + expect(overTheWire.data?.deployedAt).toEqual(fromDate.data?.deployedAt) + }) + + it('round-trips a JSONB-shaped blob without dropping blocks or edges', () => { + const state = checkpointState() + + const parsed = parseWorkflowStateForPersistence(state) + + expect(parsed.success).toBe(true) + expect(Object.keys(parsed.data?.blocks ?? {})).toEqual(['block-1']) + expect(parsed.data?.lastSaved).toBe(1_754_000_000_000) + }) + + it('accepts a null deployedAt, which the revert passes for a never-deployed checkpoint', () => { + const parsed = parseWorkflowStateForPersistence(checkpointState({ deployedAt: null })) + + expect(parsed.success).toBe(true) + expect(parsed.data?.deployedAt).toBeNull() + }) + + /** The validation the removed HTTP hop used to provide: a malformed blob must not be written. */ + it('rejects a blob whose blocks are malformed', () => { + const parsed = parseWorkflowStateForPersistence({ + blocks: { 'block-1': { id: 'block-1' } }, + edges: [], + }) + + expect(parsed.success).toBe(false) + }) + + it('rejects a blob missing blocks entirely', () => { + expect(parseWorkflowStateForPersistence({ edges: [] }).success).toBe(false) + }) +}) diff --git a/apps/sim/lib/workflows/persistence/save-normalized-state.ts b/apps/sim/lib/workflows/persistence/save-normalized-state.ts new file mode 100644 index 00000000000..6ecfb9cb5c0 --- /dev/null +++ b/apps/sim/lib/workflows/persistence/save-normalized-state.ts @@ -0,0 +1,176 @@ +import { db } from '@sim/db' +import { workflow } from '@sim/db/schema' +import { createLogger } from '@sim/logger' +import { + assertWorkflowMutable, + authorizeWorkflowByWorkspacePermission, + WorkflowLockedError, + type WorkflowWorkspaceAuthorizationResult, +} from '@sim/platform-authz/workflow' +import { eq } from 'drizzle-orm' +import type { z } from 'zod' +import { + type WorkflowStateContractOutput, + workflowStateSchema, +} from '@/lib/api/contracts/workflows' +import { notifyWorkflowUpdated } from '@/lib/realtime/notify' +import { extractAndPersistCustomTools } from '@/lib/workflows/persistence/custom-tools-persistence' +import { prepareWorkflowStateForPersistence } from '@/lib/workflows/persistence/prepare-state' +import { saveWorkflowToNormalizedTables } from '@/lib/workflows/persistence/utils' +import type { BlockState, WorkflowState } from '@/stores/workflows/workflow/types' + +const logger = createLogger('WorkflowStatePersistence') + +export type SaveWorkflowNormalizedStateResult = + | { success: true; warnings: string[] } + | { success: false; status: number; error: string; details?: string } + +/** + * Validates an untrusted state blob against the same schema `PUT /api/workflows/[id]/state` + * applies, so in-process callers get the coercion and rejection the HTTP hop gave them. + */ +export function parseWorkflowStateForPersistence( + value: unknown +): z.ZodSafeParseResult { + return workflowStateSchema.safeParse(value) +} + +/** + * Writes a complete workflow state to the normalized tables: write authorization, + * the lock check, block/edge preparation, the row-locked save transaction, + * `lastSynced`/variables, custom-tool extraction, and the socket notification. + * Every surface that replaces a workflow's state calls this, so no step can be + * skipped by going through a different door. + * + * Every refusal, the lock included, comes back as a failure result — callers need + * one branch. + * + * `authorization` lets a caller that already resolved the same decision hand it in + * rather than pay for it twice; it must be the `write` decision for this workflow + * and user. + */ +export async function saveWorkflowNormalizedState(params: { + requestId: string + workflowId: string + userId: string + state: WorkflowStateContractOutput + authorization?: WorkflowWorkspaceAuthorizationResult +}): Promise { + const { requestId, workflowId, userId, state } = params + + const authorization = + params.authorization ?? + (await authorizeWorkflowByWorkspacePermission({ workflowId, userId, action: 'write' })) + const workflowData = authorization.workflow + + if (!workflowData) { + logger.warn(`[${requestId}] Workflow ${workflowId} not found for state update`) + return { success: false, status: 404, error: 'Workflow not found' } + } + + if (!authorization.allowed) { + logger.warn( + `[${requestId}] User ${userId} denied permission to update workflow state ${workflowId}` + ) + return { + success: false, + status: authorization.status || 403, + error: authorization.message || 'Access denied', + } + } + + try { + await assertWorkflowMutable(workflowId) + } catch (error) { + if (error instanceof WorkflowLockedError) { + return { success: false, status: error.status, error: error.message } + } + throw error + } + + const { state: preparedState, warnings: preparationWarnings } = + prepareWorkflowStateForPersistence({ + blocks: state.blocks as Record, + edges: state.edges as WorkflowState['edges'], + }) + + const workflowState = { + ...preparedState, + lastSaved: state.lastSaved || Date.now(), + isDeployed: state.isDeployed || false, + deployedAt: state.deployedAt, + } + + const saveResult = await db.transaction(async (tx) => { + await tx + .select({ id: workflow.id }) + .from(workflow) + .where(eq(workflow.id, workflowId)) + .limit(1) + .for('update') + + const result = await saveWorkflowToNormalizedTables( + workflowId, + workflowState as WorkflowState, + tx + ) + + if (!result.success) return result + + const updateData: { + lastSynced: Date + updatedAt: Date + variables?: typeof state.variables + } = { + lastSynced: new Date(), + updatedAt: new Date(), + } + + if (state.variables !== undefined) { + updateData.variables = state.variables + } + + await tx.update(workflow).set(updateData).where(eq(workflow.id, workflowId)) + + return result + }) + + if (!saveResult.success) { + logger.error(`[${requestId}] Failed to save workflow ${workflowId} state:`, saveResult.error) + return { + success: false, + status: 500, + error: 'Failed to save workflow state', + details: saveResult.error, + } + } + + try { + const workspaceId = workflowData.workspaceId + if (workspaceId) { + const { saved, errors } = await extractAndPersistCustomTools( + workflowState, + workspaceId, + userId + ) + + if (saved > 0) { + logger.info(`[${requestId}] Persisted ${saved} custom tool(s) to database`, { workflowId }) + } + + if (errors.length > 0) { + logger.warn(`[${requestId}] Some custom tools failed to persist`, { errors, workflowId }) + } + } else { + logger.warn(`[${requestId}] Workflow has no workspaceId, skipping custom tools persistence`, { + workflowId, + }) + } + } catch (error) { + logger.error(`[${requestId}] Failed to persist custom tools`, { error, workflowId }) + } + + await notifyWorkflowUpdated(workflowId) + + return { success: true, warnings: preparationWarnings } +} diff --git a/apps/sim/lib/workspace-files/queries.test.ts b/apps/sim/lib/workspace-files/queries.test.ts index e3212079206..0f69aa9ce40 100644 --- a/apps/sim/lib/workspace-files/queries.test.ts +++ b/apps/sim/lib/workspace-files/queries.test.ts @@ -53,6 +53,62 @@ describe('listWorkspaceFilesWithShares', () => { expect(file.uploadedAt).toEqual(new Date('2026-01-01T00:00:00.000Z')) }) + /** + * `maxRows` exists for a caller that will only use the list if the whole workspace fits + * its payload budget, so overflow must be reported as `null` — a prefix returned here + * would be presented as the workspace's complete file list. The share read still runs + * concurrently and is discarded: the under-budget workspaces are the common case, and + * serializing the two reads to save this one would tax every normal request. + */ + it('returns null when the workspace exceeds maxRows', async () => { + mockListWorkspaceFiles.mockResolvedValue([STORED_FILE, STORED_FILE, STORED_FILE]) + + const result = await listWorkspaceFilesWithShares('ws-1', 'active', { maxRows: 2 }) + + expect(result).toBeNull() + expect(mockListWorkspaceFiles).toHaveBeenCalledWith('ws-1', { scope: 'active', limit: 3 }) + }) + + it('returns the list when it fits maxRows', async () => { + mockListWorkspaceFiles.mockResolvedValue([STORED_FILE]) + + const result = await listWorkspaceFilesWithShares('ws-1', 'active', { maxRows: 2 }) + + expect(result).toHaveLength(1) + expect(mockGetWorkspaceShares).toHaveBeenCalledWith('file', 'ws-1') + }) + + /** The boundary the `>` comparison turns on: exactly maxRows must still be the list. */ + it('returns the list when it sits exactly on maxRows', async () => { + mockListWorkspaceFiles.mockResolvedValue([STORED_FILE, STORED_FILE]) + + const result = await listWorkspaceFilesWithShares('ws-1', 'active', { maxRows: 2 }) + + expect(result).toHaveLength(2) + }) + + /** + * The file read swallows errors and returns `[]` by default. A caller seeding a cache + * must not receive that: an empty list would be cached as "this workspace has no files". + */ + it('propagates a failed read instead of degrading to an empty list', async () => { + await listWorkspaceFilesWithShares('ws-1', 'active', { throwOnError: true }) + + expect(mockListWorkspaceFiles).toHaveBeenCalledWith( + 'ws-1', + expect.objectContaining({ throwOnError: true }) + ) + }) + + it('does not ask the file read to throw unless the caller opts in', async () => { + await listWorkspaceFilesWithShares('ws-1', 'active') + + expect(mockListWorkspaceFiles).toHaveBeenCalledWith( + 'ws-1', + expect.not.objectContaining({ throwOnError: true }) + ) + }) + it('joins each file public share onto its row', async () => { const share = { id: 'share-1', diff --git a/apps/sim/lib/workspace-files/queries.ts b/apps/sim/lib/workspace-files/queries.ts index 9a5f0d6b8ca..95c856903eb 100644 --- a/apps/sim/lib/workspace-files/queries.ts +++ b/apps/sim/lib/workspace-files/queries.ts @@ -6,21 +6,40 @@ import { } from '@/lib/uploads/contexts/workspace/workspace-file-manager' /** - * Lists a workspace's files with each file's public share joined on — shared by - * `GET /api/workspaces/[id]/files` and the Files browser's server prefetch so both cache - * one shape. + * Lists a workspace's files with each file's public share joined on, parsed through the + * `GET /api/workspaces/[id]/files` response contract so the workspace layout's server seed + * caches exactly the shape that route returns. * * Parsing through the route contract's response schema strips the server-only fields * `requestJson` strips on the client (`contentUpdatedAt`), so a prefetched entry is identical * to a client fetch rather than carrying a field that vanishes on the next refetch. * * Callers authorize the viewer against `workspaceId` first. + * + * `maxRows` bounds the result for a caller that only uses the list if the whole workspace + * fits a payload budget: on overflow it returns `null` rather than the prefix, so a + * truncated read is never presented as the workspace's files. The two reads still run + * concurrently, since under-budget workspaces are the common case. + * + * `throwOnError` propagates a failed file read instead of degrading to an empty list, + * which a cache seed would store as authoritative. */ -export async function listWorkspaceFilesWithShares(workspaceId: string, scope: WorkspaceFileScope) { +export async function listWorkspaceFilesWithShares( + workspaceId: string, + scope: WorkspaceFileScope, + options?: { maxRows?: number; throwOnError?: boolean } +) { + const maxRows = options?.maxRows const [files, shares] = await Promise.all([ - listWorkspaceFiles(workspaceId, { scope }), + listWorkspaceFiles(workspaceId, { + scope, + limit: maxRows === undefined ? undefined : maxRows + 1, + throwOnError: options?.throwOnError, + }), getWorkspaceShares('file', workspaceId), ]) + if (maxRows !== undefined && files.length > maxRows) return null + const withShares = files.map((file) => ({ ...file, share: shares.get(file.id) ?? null })) return listWorkspaceFilesContract.response.schema.shape.files.parse(withShares) } diff --git a/apps/sim/lib/workspaces/permissions/utils.test.ts b/apps/sim/lib/workspaces/permissions/utils.test.ts index 64d1aa2011d..9d132ff53bf 100644 --- a/apps/sim/lib/workspaces/permissions/utils.test.ts +++ b/apps/sim/lib/workspaces/permissions/utils.test.ts @@ -745,6 +745,33 @@ describe('Permission Utils', () => { expect(result).toEqual({ id: 'workspace123', ownerId: null }) }) + + /** + * Archived visibility is applied in JS, not SQL, so the read can be shared by the + * gates that disagree about it. That makes these two the boundary worth pinning: + * if the filter ever stops matching the old `archived_at IS NULL` predicate, archived + * workspaces silently become visible to callers that asked not to see them. + */ + it.concurrent('should hide an archived workspace by default', async () => { + const chain = createMockChain([ + { id: 'workspace123', ownerId: 'owner456', archivedAt: new Date('2026-01-01') }, + ]) + mockDb.select.mockReturnValue(chain) + + const result = await getWorkspaceWithOwner('workspace123') + + expect(result).toBeNull() + }) + + it.concurrent('should return an archived workspace when asked to include them', async () => { + const archivedAt = new Date('2026-01-01') + const chain = createMockChain([{ id: 'workspace123', ownerId: 'owner456', archivedAt }]) + mockDb.select.mockReturnValue(chain) + + const result = await getWorkspaceWithOwner('workspace123', { includeArchived: true }) + + expect(result).toEqual({ id: 'workspace123', ownerId: 'owner456', archivedAt }) + }) }) describe('workspaceExists', () => { diff --git a/apps/sim/lib/workspaces/permissions/utils.ts b/apps/sim/lib/workspaces/permissions/utils.ts index 16bb3251d5a..9634157d18d 100644 --- a/apps/sim/lib/workspaces/permissions/utils.ts +++ b/apps/sim/lib/workspaces/permissions/utils.ts @@ -1,3 +1,4 @@ +import { cache } from 'react' import { db } from '@sim/db' import { member, permissions, user, type WorkspaceMode, workspace } from '@sim/db/schema' import { @@ -77,17 +78,12 @@ export async function getWorkspaceById( return exists ? { id: workspaceId } : null } -/** - * Get a workspace with owner info by ID - * - * @param workspaceId - The workspace ID to look up - * @returns The workspace with owner info if found, null otherwise - */ -export async function getWorkspaceWithOwner( +async function selectWorkspaceWithOwner( workspaceId: string, - options?: { includeArchived?: boolean; executor?: DbOrTx; forUpdate?: boolean } + includeArchived: boolean, + executor: DbOrTx, + forUpdate: boolean ): Promise { - const { includeArchived = false, executor = db, forUpdate = false } = options ?? {} const query = executor .select({ id: workspace.id, @@ -110,6 +106,46 @@ export async function getWorkspaceWithOwner( return ws || null } +/** + * Request-memoized plain workspace read, keyed by id alone. One render pass resolves the + * same row through several independent gates. + * + * Keyed on the id and NOT on archived visibility: the gates disagree about archived + * workspaces, so memoizing that argument would give each answer its own entry and dedupe + * nothing. + * + * The returned row is SHARED by every consumer in the render. Treat it as immutable — an + * in-place edit poisons every gate that reads it for the rest of the pass. + */ +const readWorkspaceWithOwner = cache( + (workspaceId: string): Promise => + selectWorkspaceWithOwner(workspaceId, true, db, false) +) + +/** + * Get a workspace with owner info by ID + * + * Transaction-scoped (`executor`) and lock-acquiring (`forUpdate`) reads + * deliberately bypass {@link readWorkspaceWithOwner}: a row read inside one + * caller's transaction, or under a row lock only that caller holds, must never + * be handed to a later caller that took neither. + * + * @param workspaceId - The workspace ID to look up + * @returns The workspace with owner info if found, null otherwise + */ +export async function getWorkspaceWithOwner( + workspaceId: string, + options?: { includeArchived?: boolean; executor?: DbOrTx; forUpdate?: boolean } +): Promise { + const { includeArchived = false, executor, forUpdate = false } = options ?? {} + if (executor || forUpdate) { + return selectWorkspaceWithOwner(workspaceId, includeArchived, executor ?? db, forUpdate) + } + const ws = await readWorkspaceWithOwner(workspaceId) + if (!ws) return null + return includeArchived || !ws.archivedAt ? ws : null +} + /** * Resolve the effective workspace permission for a user under the governance * inheritance model: the owners/admins of the organization that owns the @@ -131,18 +167,7 @@ export async function getEffectiveWorkspacePermission( return resolveEffectiveWorkspacePermission(userId, ws.id, ws.organizationId, executor) } -/** - * Check workspace access for a user - * - * Verifies the workspace exists and the user has access to it. - * Returns access level (read/write) based on ownership, explicit permissions, - * and organization-admin inheritance. - * - * @param workspaceId - The workspace ID to check - * @param userId - The user ID to check access for - * @returns WorkspaceAccess object with exists, hasAccess, canWrite, and workspace data - */ -export async function checkWorkspaceAccess( +async function resolveWorkspaceAccessForUser( workspaceId: string, userId: string ): Promise { @@ -167,6 +192,25 @@ export async function checkWorkspaceAccess( return { exists: true, hasAccess, canWrite, canAdmin, workspace: ws, permission } } +/** + * Check workspace access for a user + * + * Verifies the workspace exists and the user has access to it. + * Returns access level (read/write) based on ownership, explicit permissions, + * and organization-admin inheritance. + * + * Request-memoized: a Server Component render pass authorizes the same + * (workspace, viewer) pair through several gates, and the answer cannot change + * mid-render. Outside a render React evaluates the resolver normally, so route + * handlers and background work re-read exactly as before. Takes no executor, so + * no transaction-scoped or locked read can ever be memoized here. + * + * @param workspaceId - The workspace ID to check + * @param userId - The user ID to check access for + * @returns WorkspaceAccess object with exists, hasAccess, canWrite, and workspace data + */ +export const checkWorkspaceAccess = cache(resolveWorkspaceAccessForUser) + /** * Returns `provided` when it was resolved for this exact workspace, otherwise * resolves fresh. The id match is what keeps a caller from authorizing against diff --git a/apps/sim/tools/index.ts b/apps/sim/tools/index.ts index 575b6da5c35..7a708b9db0b 100644 --- a/apps/sim/tools/index.ts +++ b/apps/sim/tools/index.ts @@ -56,6 +56,7 @@ import { } from '@/lib/execution/private-tool-metadata' import { parseMcpToolId } from '@/lib/mcp/utils' import { hostedKeyMetrics } from '@/lib/monitoring/metrics' +import type { CredentialTokenPayload } from '@/lib/oauth/token-resolution' import { resolveWorkspaceFileReference } from '@/lib/uploads/contexts/workspace/workspace-file-manager' import { markWorkspaceFileSecretProvenanceUnknown } from '@/lib/uploads/contexts/workspace/workspace-file-secret-provenance' import { assertPermissionsAllowed } from '@/ee/access-control/utils/permission-check' @@ -1702,8 +1703,6 @@ async function executeToolImplementation( `[${requestId}] Tool ${toolId} needs access token for credential: ${contextParams.credential}` ) try { - const baseUrl = getInternalApiBaseUrl() - const workflowId = contextParams._context?.workflowId const userId = contextParams._context?.userId @@ -1724,51 +1723,86 @@ async function executeToolImplementation( } } - logger.info(`[${requestId}] Fetching access token from ${baseUrl}/api/auth/oauth/token`) + /** + * The acting user asserted alongside an internal token. Only sent when the + * run enforces credential access, matching the `userId` query param the HTTP + * surface accepted — it never widens access, it only pins the assertion to + * the token subject. + */ + const callerUserId = + userId && contextParams._context?.enforceCredentialAccess ? userId : undefined - const tokenUrlObj = new URL('/api/auth/oauth/token', baseUrl) - if (workflowId) { - tokenUrlObj.searchParams.set('workflowId', workflowId) - } - if (userId && contextParams._context?.enforceCredentialAccess) { - tokenUrlObj.searchParams.set('userId', userId) - } + let data: CredentialTokenPayload - // Always send Content-Type; add internal auth on server-side runs - const tokenHeaders: Record = { 'Content-Type': 'application/json' } if (typeof window === 'undefined') { - try { - const internalToken = await generateInternalToken(userId) - tokenHeaders.Authorization = `Bearer ${internalToken}` - } catch (_e) { - // Swallow token generation errors; the request will fail and be reported upstream + // Server-side runs resolve the credential through the same application + // operation the route calls, rather than minting an internal JWT and + // POSTing to ourselves through the load balancer. The synthesized + // `AuthResult` is exactly what verifying that self-issued token would + // have produced, so authorization, refresh, and audit are unchanged — + // including failing closed when the run carries no user id. + const { resolveCredentialToken } = await import('@/lib/oauth/token-resolution') + const result = await resolveCredentialToken( + { success: true, authType: 'internal_jwt', userId }, + { + requestId, + credentialId: contextParams.credential as string, + workflowId, + scopes: tokenPayload.scopes, + impersonateEmail: tokenPayload.impersonateEmail, + callerUserId, + } + ) + + if (!result.ok) { + logger.error(`[${requestId}] Token fetch failed for ${toolId}:`, { + status: result.status, + error: result.error, + }) + const toolLabel = tool?.name || toolId + throw new Error(`Failed to obtain credential for ${toolLabel}: ${result.error}`) } - } - const response = await fetch(tokenUrlObj.toString(), { - method: 'POST', - headers: tokenHeaders, - body: JSON.stringify(tokenPayload), - }) + data = result.token + } else { + const baseUrl = getInternalApiBaseUrl() + logger.info(`[${requestId}] Fetching access token from ${baseUrl}/api/auth/oauth/token`) + + const tokenUrlObj = new URL('/api/auth/oauth/token', baseUrl) + if (workflowId) { + tokenUrlObj.searchParams.set('workflowId', workflowId) + } + if (callerUserId) { + tokenUrlObj.searchParams.set('userId', callerUserId) + } - if (!response.ok) { - const errorText = await response.text() - logger.error(`[${requestId}] Token fetch failed for ${toolId}:`, { - status: response.status, - error: errorText, + // boundary-raw-fetch: browser-side tool runs authenticate with the session cookie against the same-origin token route + const response = await fetch(tokenUrlObj.toString(), { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(tokenPayload), }) - let parsedError = errorText - try { - const parsed = JSON.parse(errorText) - if (parsed.error) parsedError = parsed.error - } catch { - // Use raw text + + if (!response.ok) { + const errorText = await response.text() + logger.error(`[${requestId}] Token fetch failed for ${toolId}:`, { + status: response.status, + error: errorText, + }) + let parsedError = errorText + try { + const parsed = JSON.parse(errorText) + if (parsed.error) parsedError = parsed.error + } catch { + // Use raw text + } + const toolLabel = tool?.name || toolId + throw new Error(`Failed to obtain credential for ${toolLabel}: ${parsedError}`) } - const toolLabel = tool?.name || toolId - throw new Error(`Failed to obtain credential for ${toolLabel}: ${parsedError}`) + + data = (await response.json()) as CredentialTokenPayload } - const data = await response.json() contextParams.accessToken = data.accessToken if (data.idToken) { contextParams.idToken = data.idToken