Skip to content

Commit 9efb061

Browse files
committed
fix(provenance): preserve completions and isolate tool identities
1 parent 676f719 commit 9efb061

33 files changed

Lines changed: 1822 additions & 103 deletions

apps/sim/app/api/chat/[identifier]/route.test.ts

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -534,6 +534,38 @@ describe('Chat Identifier API Route', () => {
534534
)
535535
}, 10000)
536536

537+
it('projects the internal completion envelope to the public streaming callback', async () => {
538+
const response = await POST(createMockNextRequest('POST', { input: 'Hello' }), {
539+
params: Promise.resolve({ identifier: 'test-chat' }),
540+
})
541+
expect(response.status).toBe(200)
542+
const onBlockComplete = vi.fn()
543+
await vi.mocked(createStreamingResponse).mock.calls[0][0].executeFn({
544+
onStream: vi.fn(),
545+
onBlockComplete,
546+
abortSignal: new AbortController().signal,
547+
})
548+
await vi.mocked(executeWorkflow).mock.calls[0][4]?.onBlockComplete?.('block-1', {
549+
output: { value: 'public output' },
550+
outputBlockId: 'child:block-1',
551+
resolvedSecretTraceProvenance: {
552+
version: 1,
553+
complete: true,
554+
entries: [{ encryptedValue: 'private-ciphertext' }],
555+
scope: { userId: 'user-1', workspaceId: 'workspace-1' },
556+
},
557+
executionTime: 1,
558+
executionOrder: 0,
559+
startedAt: '2026-01-01T00:00:00Z',
560+
endedAt: '2026-01-01T00:00:01Z',
561+
})
562+
expect(onBlockComplete).toHaveBeenCalledExactlyOnceWith(
563+
'block-1',
564+
{ value: 'public output' },
565+
'child:block-1'
566+
)
567+
})
568+
537569
it('executes with the email proven by the chat authentication gate', async () => {
538570
mockValidateChatAuth.mockResolvedValueOnce({
539571
authorized: true,

apps/sim/app/api/chat/[identifier]/route.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -414,7 +414,8 @@ export const POST = withRouteHandler(
414414
isSecureMode: true,
415415
workflowTriggerType: 'chat',
416416
onStream,
417-
onBlockComplete,
417+
onBlockComplete: (blockId, data) =>
418+
onBlockComplete(blockId, data.output, data.outputBlockId),
418419
skipLoggingComplete: true,
419420
abortSignal,
420421
executionMode: 'stream',

apps/sim/app/api/workflows/[id]/execute/route.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1718,7 +1718,8 @@ async function handleExecutePost(
17181718
isSecureMode: false,
17191719
workflowTriggerType: triggerType === 'chat' ? 'chat' : 'api',
17201720
onStream,
1721-
onBlockComplete,
1721+
onBlockComplete: (blockId, data) =>
1722+
onBlockComplete(blockId, data.output, data.outputBlockId),
17221723
skipLoggingComplete: true,
17231724
includeFileBase64,
17241725
base64MaxBytes,

apps/sim/background/resume-execution.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import {
1414
getTimeoutErrorMessage,
1515
} from '@/lib/core/execution-limits'
1616
import { withCascadeLock } from '@/lib/table/cascade-lock'
17+
import type { WorkflowCellProgressWriter } from '@/lib/table/cell-write'
1718
import { isExecCancelled } from '@/lib/table/deps'
1819
import type { RowExecutionMetadata } from '@/lib/table/types'
1920
import { classifyWorkflowCellTerminalResult } from '@/lib/table/workflow-cell-result'
@@ -247,7 +248,7 @@ function throwIfResumeAttemptTimedOut(
247248
}
248249

249250
type CellWriters = {
250-
cellOnBlockComplete: (blockId: string, output: unknown) => Promise<void>
251+
cellOnBlockComplete: WorkflowCellProgressWriter['onBlockComplete']
251252
writeCellTerminal: (
252253
status: 'completed' | 'error' | 'cancelled' | 'paused',
253254
error: string | null
Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
/** A durable tool identity must never be reused by a different run. */
2+
export class AsyncToolCallOwnershipError extends Error {
3+
constructor() {
4+
super('Async tool call belongs to another run')
5+
this.name = 'AsyncToolCallOwnershipError'
6+
}
7+
}

apps/sim/lib/copilot/async-runs/repository.test.ts

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -406,4 +406,28 @@ describe('async tool repository single-row semantics', () => {
406406
expect(dbChainMockFns.values).not.toHaveBeenCalled()
407407
}
408408
)
409+
410+
it('refuses a provider-ID collision with an existing row in another run', async () => {
411+
dbChainMockFns.limit.mockResolvedValueOnce([
412+
{ runId: 'old-run', toolCallId: 'provider-shared-call', toolName: 'browser_close_tab' },
413+
])
414+
await expect(
415+
upsertAsyncToolCall({
416+
runId: 'current-run',
417+
toolCallId: 'provider-shared-call',
418+
toolName: 'glob',
419+
})
420+
).rejects.toThrow('Async tool call belongs to another run')
421+
expect(dbChainMockFns.values).not.toHaveBeenCalled()
422+
})
423+
424+
it('refuses a foreign row that wins the insert race', async () => {
425+
dbChainMockFns.limit
426+
.mockResolvedValueOnce([])
427+
.mockResolvedValueOnce([{ runId: 'other-run', toolCallId: 'call-race' }])
428+
dbChainMockFns.returning.mockResolvedValueOnce([])
429+
await expect(
430+
upsertAsyncToolCall({ runId: 'current-run', toolCallId: 'call-race', toolName: 'glob' })
431+
).rejects.toThrow('Async tool call belongs to another run')
432+
})
409433
})

apps/sim/lib/copilot/async-runs/repository.ts

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import { createLogger } from '@sim/logger'
1111
import { filterUndefined } from '@sim/utils/object'
1212
import { sanitizeValueForJsonb } from '@sim/utils/string'
1313
import { and, desc, eq, inArray, isNull, or, sql } from 'drizzle-orm'
14+
import { AsyncToolCallOwnershipError } from '@/lib/copilot/async-runs/errors'
1415
import { TraceAttr } from '@/lib/copilot/generated/trace-attributes-v1'
1516
import { TraceSpan } from '@/lib/copilot/generated/trace-spans-v1'
1617
import { markSpanForError } from '@/lib/copilot/request/otel'
@@ -219,7 +220,12 @@ export async function upsertAsyncToolCall(input: {
219220
},
220221
async () => {
221222
const existing = await getAsyncToolCall(input.toolCallId)
222-
if (existing) return existing
223+
if (existing) {
224+
if (input.runId && existing.runId !== input.runId) {
225+
throw new AsyncToolCallOwnershipError()
226+
}
227+
return existing
228+
}
223229

224230
const incomingStatus = input.status ?? 'pending'
225231
const effectiveRunId = input.runId ?? null
@@ -250,7 +256,11 @@ export async function upsertAsyncToolCall(input: {
250256
.onConflictDoNothing()
251257
.returning()
252258

253-
return row ?? getAsyncToolCall(input.toolCallId)
259+
const persisted = row ?? (await getAsyncToolCall(input.toolCallId))
260+
if (persisted && persisted.runId !== effectiveRunId) {
261+
throw new AsyncToolCallOwnershipError()
262+
}
263+
return persisted
254264
}
255265
)
256266
}

0 commit comments

Comments
 (0)