diff --git a/apps/sim/lib/auth/internal-delegation.test.ts b/apps/sim/lib/auth/internal-delegation.test.ts index 2450d21bdb6..e4bac81615e 100644 --- a/apps/sim/lib/auth/internal-delegation.test.ts +++ b/apps/sim/lib/auth/internal-delegation.test.ts @@ -11,6 +11,8 @@ const { mockResolveWorkflow, mockResolveRun, mockResolveExecution, mockResolveDe mockResolveDeploymentVersion: vi.fn(), })) +vi.mock('@sim/utils/helpers', () => ({ sleep: vi.fn().mockResolvedValue(undefined) })) + vi.mock('@/lib/workflows/application/context', () => ({ resolveActiveWorkflowApplicationContext: mockResolveWorkflow, resolveActiveWorkflowRunApplicationContext: mockResolveRun, @@ -323,5 +325,44 @@ describe('bindInternalExecutorDelegation', () => { await expect( bindInternalExecutorDelegation(claims, { audience: 'sim:workspace-files' }) ).rejects.toBe(infrastructureError) + expect(mockResolveWorkflow).toHaveBeenCalledTimes(1) + }) + + it('retries a canonical load that failed on a reset database connection', async () => { + const connectionReset = Object.assign(new Error('Failed query'), { + cause: Object.assign(new Error('read ECONNRESET'), { code: 'ECONNRESET' }), + }) + mockResolveRun.mockRejectedValueOnce(connectionReset) + + const principal = await bindInternalExecutorDelegation( + { ...claims, executionId: 'execution-1' }, + { audience: 'sim:function-execute' } + ) + + expect(principal.workspaceId).toBe('workspace-1') + expect(mockResolveRun).toHaveBeenCalledTimes(2) + }) + + it('retries a current-workflow load that failed on a reset database connection', async () => { + const connectionReset = Object.assign(new Error('Failed query'), { + cause: Object.assign(new Error('read ECONNRESET'), { code: 'ECONNRESET' }), + }) + mockResolveDeploymentVersion.mockRejectedValueOnce(connectionReset) + + const principal = await bindInternalExecutorDelegation( + { + ...claims, + executionId: 'execution-1', + currentWorkflow: { + workflowId: 'child-workflow', + mode: 'deployment', + deploymentVersionId: 'deployment-version-1', + }, + }, + { audience: 'sim:credential-groups' } + ) + + expect(principal.workspaceId).toBe('workspace-1') + expect(mockResolveDeploymentVersion).toHaveBeenCalledTimes(2) }) }) diff --git a/apps/sim/lib/auth/internal-delegation.ts b/apps/sim/lib/auth/internal-delegation.ts index 502ab3859b8..57d8338a9b2 100644 --- a/apps/sim/lib/auth/internal-delegation.ts +++ b/apps/sim/lib/auth/internal-delegation.ts @@ -4,6 +4,7 @@ import type { } from '@sim/auth/principal' import type { VerifiedInternalDelegation } from '@/lib/auth/internal' import { asOrchestrationError } from '@/lib/core/orchestration/types' +import { withDatabaseReadRetry } from '@/lib/db/read-retry' import { type ActiveWorkflowApplicationContext, resolveActiveWorkflowApplicationContext, @@ -25,7 +26,15 @@ export class InvalidInternalDelegationBindingError extends Error { } } -/** Binds signed executor claims to the workflow's canonical active workspace. */ +const CANONICAL_LOAD_RETRY = { label: 'internal delegation canonical load' } as const + +/** + * Binds signed executor claims to the workflow's canonical active workspace. + * + * Every tool call a workflow run delegates (Function, MCP, files, knowledge) binds here first, and + * nothing above it retries. The canonical loads are independent reads, so a transient connection + * failure such as a reset pooled socket is retried instead of failing the block. + */ export async function bindInternalExecutorDelegation( claims: VerifiedInternalDelegation, options: BindInternalExecutorDelegationOptions @@ -43,19 +52,32 @@ export async function bindInternalExecutorDelegation( try { if (claims.currentWorkflow) { if (!claims.executionId) throw new InvalidInternalDelegationBindingError() - const executionContext = await resolveActiveWorkflowExecutionApplicationContext({ - runId: claims.executionId, - assertedWorkflowId: claims.workflowId, - }) + const runId = claims.executionId + const executionContext = await withDatabaseReadRetry( + () => + resolveActiveWorkflowExecutionApplicationContext({ + runId, + assertedWorkflowId: claims.workflowId, + }), + CANONICAL_LOAD_RETRY + ) context = executionContext rootDeploymentVersionId = executionContext.deploymentVersionId } else if (claims.executionId) { - context = await resolveActiveWorkflowRunApplicationContext({ - runId: claims.executionId, - assertedWorkflowId: claims.workflowId, - }) + const runId = claims.executionId + context = await withDatabaseReadRetry( + () => + resolveActiveWorkflowRunApplicationContext({ + runId, + assertedWorkflowId: claims.workflowId, + }), + CANONICAL_LOAD_RETRY + ) } else { - context = await resolveActiveWorkflowApplicationContext({ workflowId: claims.workflowId }) + context = await withDatabaseReadRetry( + () => resolveActiveWorkflowApplicationContext({ workflowId: claims.workflowId }), + CANONICAL_LOAD_RETRY + ) } } catch (error) { if (asOrchestrationError(error)?.code === 'not_found') { @@ -74,18 +96,23 @@ export async function bindInternalExecutorDelegation( throw new InvalidInternalDelegationBindingError() } } else { + const currentWorkflow = claims.currentWorkflow + const workspaceId = context.workspaceId try { - const currentContext = - claims.currentWorkflow.mode === 'deployment' - ? await resolveActiveWorkflowDeploymentVersionApplicationContext({ - workflowId: claims.currentWorkflow.workflowId, - deploymentVersionId: claims.currentWorkflow.deploymentVersionId, - assertedWorkspaceId: context.workspaceId, - }) - : await resolveActiveWorkflowApplicationContext({ - workflowId: claims.currentWorkflow.workflowId, - assertedWorkspaceId: context.workspaceId, - }) + const currentContext = await withDatabaseReadRetry( + () => + currentWorkflow.mode === 'deployment' + ? resolveActiveWorkflowDeploymentVersionApplicationContext({ + workflowId: currentWorkflow.workflowId, + deploymentVersionId: currentWorkflow.deploymentVersionId, + assertedWorkspaceId: workspaceId, + }) + : resolveActiveWorkflowApplicationContext({ + workflowId: currentWorkflow.workflowId, + assertedWorkspaceId: workspaceId, + }), + CANONICAL_LOAD_RETRY + ) if (currentContext.workspaceId !== context.workspaceId) { throw new InvalidInternalDelegationBindingError() }