Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 41 additions & 0 deletions apps/sim/lib/auth/internal-delegation.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)
})
})
69 changes: 48 additions & 21 deletions apps/sim/lib/auth/internal-delegation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand All @@ -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') {
Expand All @@ -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()
}
Expand Down
Loading