diff --git a/.changeset/sandbox-takeover-keeps-prefix.md b/.changeset/sandbox-takeover-keeps-prefix.md new file mode 100644 index 0000000000..c8f6595c26 --- /dev/null +++ b/.changeset/sandbox-takeover-keeps-prefix.md @@ -0,0 +1,8 @@ +--- +'@tanstack/ai-sandbox': patch +'@tanstack/ai-claude-code': patch +'@tanstack/ai-codex': patch +'@tanstack/ai-grok-build': patch +--- + +Keep text and tool calls the first host already delivered in the saved transcript after a sandbox takeover. The event log still stores that prefix once. diff --git a/packages/ai-claude-code/tests/attach.test.ts b/packages/ai-claude-code/tests/attach.test.ts index ff763677ae..f138780d52 100644 --- a/packages/ai-claude-code/tests/attach.test.ts +++ b/packages/ai-claude-code/tests/attach.test.ts @@ -18,6 +18,7 @@ import { localProcessSandbox } from '@tanstack/ai-sandbox-local-process' import { chunkFingerprint, createRunScopedIdGen, + isStoredReplayChunk, exitSentinelLine, journalPaths, } from '@tanstack/ai-sandbox' @@ -365,13 +366,16 @@ describe('claude-code durable-run wiring (attach path)', () => { }), ) - // Every chunk in `alreadyDelivered` is suppressed; only the tail this - // durability adapter had NOT yet seen comes through. Compared via - // `chunkFingerprint` (key-order independent, excludes the wall-clock - // `timestamp`) since the two translations run at different instants. - expect(chunks.map(chunkFingerprint)).toEqual( + // Stored chunks are yielded marked so chat() can record them. The tail + // this durability adapter had NOT yet seen is the unmarked part. + // Compared via `chunkFingerprint` (key-order independent, excludes the + // wall-clock `timestamp`) since the two translations run at different + // instants. + const delivered = chunks.filter((chunk) => !isStoredReplayChunk(chunk)) + expect(delivered.map(chunkFingerprint)).toEqual( notYetDelivered.map(chunkFingerprint), ) + expect(chunks.length - delivered.length).toBe(alreadyDelivered.length) await sbx.destroy() }) diff --git a/packages/ai-codex/tests/attach.test.ts b/packages/ai-codex/tests/attach.test.ts index 0847e26881..bccf6d1b0d 100644 --- a/packages/ai-codex/tests/attach.test.ts +++ b/packages/ai-codex/tests/attach.test.ts @@ -27,6 +27,7 @@ import { localProcessSandbox } from '@tanstack/ai-sandbox-local-process' import { DurableRunIdRequiredError, SandboxCapability, + isStoredReplayChunk, journalPaths, journaledCommand, provideSandboxDurability, @@ -304,11 +305,13 @@ describe('codex attach wiring', () => { }), ) - // The RUN_STARTED chunk was already "delivered" per the stored log, so - // alignment must suppress it — if `alignedIfAttaching` were bypassed, the - // freshly-translated RUN_STARTED would come through unsuppressed and this - // would fail. - expect(chunks.some((c) => c.type === 'RUN_STARTED')).toBe(false) + // The RUN_STARTED chunk was already in the stored log. It is yielded + // marked, for the transcript, and is not part of what gets appended. + // If `alignedIfAttaching` were bypassed, a fresh unmarked RUN_STARTED + // would come through and this would fail. + const delivered = chunks.filter((chunk) => !isStoredReplayChunk(chunk)) + expect(delivered.some((c) => c.type === 'RUN_STARTED')).toBe(false) + expect(chunks.some(isStoredReplayChunk)).toBe(true) // Everything after the aligned prefix must still arrive. expect(textOf(chunks)).toContain('pong') expect(chunks.some((c) => c.type === 'RUN_FINISHED')).toBe(true) diff --git a/packages/ai-grok-build/tests/attach.test.ts b/packages/ai-grok-build/tests/attach.test.ts index 63da39b306..bd5cc727cf 100644 --- a/packages/ai-grok-build/tests/attach.test.ts +++ b/packages/ai-grok-build/tests/attach.test.ts @@ -32,6 +32,7 @@ import { SandboxCapability, SandboxDurabilityCapability, chunkFingerprint, + isStoredReplayChunk, exitSentinelLine, journalPaths, } from '@tanstack/ai-sandbox' @@ -362,7 +363,12 @@ describe( durability: durabilityWith(fakeLog(replayed), true), }), ) - expect(nothingNew).toEqual([]) + // Stored chunks are yielded marked so chat() can record them. What would + // be appended to the log is the unmarked tail, and here that tail is empty. + expect(nothingNew.filter((chunk) => !isStoredReplayChunk(chunk))).toEqual( + [], + ) + expect(nothingNew.every(isStoredReplayChunk)).toBe(true) // Seeded with all but the last: exactly the undelivered tail comes through, // and it is fingerprint-identical to the final translated chunk. @@ -374,7 +380,9 @@ describe( durability: durabilityWith(fakeLog(replayed.slice(0, -1)), true), }), ) - expect(fingerprints(tail)).toEqual(fingerprints(replayed.slice(-1))) + expect( + fingerprints(tail.filter((chunk) => !isStoredReplayChunk(chunk))), + ).toEqual(fingerprints(replayed.slice(-1))) await sbx.destroy() }) diff --git a/packages/ai-sandbox/src/align.ts b/packages/ai-sandbox/src/align.ts index 0049889923..a5d2d9c4ab 100644 --- a/packages/ai-sandbox/src/align.ts +++ b/packages/ai-sandbox/src/align.ts @@ -195,6 +195,39 @@ export interface AlignToStoredLogOptions { * {@link DEFAULT_MAX_OUT_OF_BAND_SKIP}. Ignored when `isOutOfBand` is absent. */ maxOutOfBandSkip?: number + /** + * Yield a stored chunk instead of dropping it. The copy is marked with + * {@link isStoredReplayChunk} so `withSandbox` can record it and then keep + * it off the public stream, and so `pipeToRunLog` does not append it again. + */ + yieldStored?: boolean +} + +/** Field set on a replayed chunk the event log already holds. Not an AG-UI key. */ +const STORED_REPLAY = 'tanstackStoredReplay' + +/** + * True when this chunk is the replay of one the event log already holds. + * + * The chat engine must see it (text and tool calls), and the live stream must + * not. `withSandbox` drops it after recording. `pipeToRunLog` does not append it. + */ +export function isStoredReplayChunk(chunk: StreamChunk): boolean { + if ( + typeof chunk !== 'object' || + chunk === null || + !(STORED_REPLAY in chunk) + ) { + return false + } + return (chunk as Record)[STORED_REPLAY] === true +} + +function asStoredReplay(chunk: StreamChunk): StreamChunk { + // The flag is not an AG-UI field, so the spread is not a StreamChunk until + // it is checked on the way out. + const copy: Record = { ...chunk, [STORED_REPLAY]: true } + return copy as unknown as StreamChunk } /** @@ -254,6 +287,9 @@ export async function* alignToStoredLog( if (expected === actual) { cursor += 1 suppressed += 1 + // The engine still has to fold this into the transcript. The mark is + // how the public stream and the event log skip it. + if (options.yieldStored) yield asStoredReplay(chunk) break } // Mismatch. Only a stored chunk the replay provably cannot reproduce may diff --git a/packages/ai-sandbox/src/durability.ts b/packages/ai-sandbox/src/durability.ts index f540999ece..61a8d8d717 100644 --- a/packages/ai-sandbox/src/durability.ts +++ b/packages/ai-sandbox/src/durability.ts @@ -294,6 +294,11 @@ export function journalOptionsFor( * cannot reproduce; without it a bridged-tool run could not be taken over at * all. Wrap the merge RESULT, never the pre-merge translator, or the comparison * is against a stream the log never contained. + * + * A matching stored chunk is yielded marked ({@link isStoredReplayChunk}), not + * dropped. `chat()` records that prefix into the transcript. `withSandbox` + * then drops the mark from the public stream, and `pipeToRunLog` does not + * append it, so the event log still holds the prefix once. */ export function alignedIfAttaching( chunks: AsyncIterable, @@ -304,6 +309,7 @@ export function alignedIfAttaching( return alignToStoredLog(chunks, { durability: durability.adapter, isOutOfBand: isBridgeCustomChunk, + yieldStored: true, ...(logger === undefined ? {} : { logger }), }) } diff --git a/packages/ai-sandbox/src/index.ts b/packages/ai-sandbox/src/index.ts index d89c7b65b6..60191efb8b 100644 --- a/packages/ai-sandbox/src/index.ts +++ b/packages/ai-sandbox/src/index.ts @@ -310,6 +310,7 @@ export { export { alignToStoredLog, isBridgeCustomChunk, + isStoredReplayChunk, JournalReplayDivergedError, JournalReplayThreadIdMismatchError, DEFAULT_MAX_OUT_OF_BAND_SKIP, diff --git a/packages/ai-sandbox/src/middleware.ts b/packages/ai-sandbox/src/middleware.ts index caf0791bae..044fb6cafe 100644 --- a/packages/ai-sandbox/src/middleware.ts +++ b/packages/ai-sandbox/src/middleware.ts @@ -44,6 +44,7 @@ import { computeWorkspaceHash } from './key' import { buildFileHookEvent, resolveFileEvents } from './file-diff' import { ProjectionCapability, provideWorkspaceProjection } from './projection' import { resolveAllSecrets, resolveSecret } from './secrets' +import { isStoredReplayChunk } from './align' import { createToolHistoryRecorder, stripObservedToolCalls, @@ -1027,11 +1028,12 @@ export function withSandbox( runState.get(ctx)?.toolHistory.reconcile(ctx) }, - // Record the harness's own tool calls as transcript messages. Observe only: - // returning nothing passes every chunk through untouched. + // Record the harness's own tool calls as transcript messages. A stored + // replay chunk is recorded, then dropped: the event log already has it. async onChunk(ctx, chunk) { const state = runState.get(ctx) state?.toolHistory.observe(chunk, ctx) + if (isStoredReplayChunk(chunk)) return null if ( state && chunk.type === 'RUN_FINISHED' && @@ -1041,6 +1043,7 @@ export function withSandbox( await drainWatcher(state, 'pause') await stopSnapshotLease(state, { closePortable: true }) } + return undefined }, async onFinish(ctx) { diff --git a/packages/ai-sandbox/src/run.ts b/packages/ai-sandbox/src/run.ts index 15fd3d4d25..3124268f71 100644 --- a/packages/ai-sandbox/src/run.ts +++ b/packages/ai-sandbox/src/run.ts @@ -22,6 +22,7 @@ * invisible, each guard reports through the optional {@link RunDeps.logger}. */ import { EventType } from '@tanstack/ai' +import { isStoredReplayChunk } from './align' import { toRunErrorPayload } from '@tanstack/ai/adapter-internals' import type { InternalLogger } from '@tanstack/ai/adapter-internals' import type { @@ -306,6 +307,8 @@ export async function pipeToRunLog( for await (const chunk of stream) { if (signal?.aborted) return finish(ctx, 'aborted') + // Already in the log. Appending it again is the doubled live tail. + if (isStoredReplayChunk(chunk)) continue await durability.append([chunk]) if (isRunErrorChunk(chunk)) { return finish( diff --git a/packages/ai-sandbox/tests/durability.test.ts b/packages/ai-sandbox/tests/durability.test.ts index 488c69201b..f649ff940e 100644 --- a/packages/ai-sandbox/tests/durability.test.ts +++ b/packages/ai-sandbox/tests/durability.test.ts @@ -8,6 +8,7 @@ import { resolveDurableRunId, resolveSandboxDurability, } from '../src/durability' +import { isStoredReplayChunk } from '../src/align' import { captureLogger, collectChunks, fakeLog, fromChunkValues } from './fakes' import type { SandboxDurabilityOptions, @@ -317,21 +318,27 @@ describe('alignedIfAttaching', () => { expect(await collectChunks(wrapped)).toEqual([text('m1', 'a')]) }) - it('suppresses the stored prefix when attaching', async () => { + it('marks the stored prefix when attaching, and still yields the rest', async () => { const adapter = fakeLog([text('m1', 'a')]) expect(await adapter.snapshot()).toHaveLength(1) const durability = resolvedDurability({ runs: new InMemoryRunStore(), durability: { adapter, attach: true }, }) - expect( - await collectChunks( - alignedIfAttaching( - fromChunkValues([text('m1', 'a'), text('m1', 'b')]), - durability, - ), + const chunks = await collectChunks( + alignedIfAttaching( + fromChunkValues([text('m1', 'a'), text('m1', 'b')]), + durability, ), - ).toEqual([text('m1', 'b')]) + ) + // The engine needs the stored chunk. The live stream does not: the mark + // is what `withSandbox` and `pipeToRunLog` drop. + expect(chunks.filter(isStoredReplayChunk)).toEqual([ + { ...text('m1', 'a'), tanstackStoredReplay: true }, + ]) + expect(chunks.filter((chunk) => !isStoredReplayChunk(chunk))).toEqual([ + text('m1', 'b'), + ]) }) it('tolerates a bridged CUSTOM chunk in the stored log when attaching', async () => { @@ -343,14 +350,16 @@ describe('alignedIfAttaching', () => { runs: new InMemoryRunStore(), durability: { adapter, attach: true }, }) - expect( - await collectChunks( - alignedIfAttaching( - fromChunkValues([text('m1', 'a'), text('m1', 'b')]), - durability, - ), + const chunks = await collectChunks( + alignedIfAttaching( + fromChunkValues([text('m1', 'a'), text('m1', 'b')]), + durability, ), - ).toEqual([text('m1', 'b')]) + ) + expect(chunks.filter((chunk) => !isStoredReplayChunk(chunk))).toEqual([ + text('m1', 'b'), + ]) + expect(chunks.filter(isStoredReplayChunk)).toHaveLength(1) }) it('forwards the logger so the alignment summary is observable', async () => { diff --git a/packages/ai-sandbox/tests/takeover-transcript.test.ts b/packages/ai-sandbox/tests/takeover-transcript.test.ts new file mode 100644 index 0000000000..5cec46ec4f --- /dev/null +++ b/packages/ai-sandbox/tests/takeover-transcript.test.ts @@ -0,0 +1,328 @@ +/** + * A successor host must save the same transcript as a run that never restarted. + * + * Alignment used to drop the stored prefix inside the adapter, before `chat()` + * and `withSandbox` could record it. The live log was already correct. The + * saved thread was not: the first host's text and tool calls were gone, and a + * tool result that arrived after the cut was left with no call. + */ +import { describe, expect, it } from 'vitest' +import { EventType, InMemoryRunStore, chat } from '@tanstack/ai' +import { memoryPersistence, withPersistence } from '@tanstack/ai-persistence' +import { defineSandbox } from '../src/sandbox' +import { withSandbox } from '../src/middleware' +import { alignedIfAttaching, getSandboxDurability } from '../src/durability' +import { pipeToRunLog } from '../src/run' +import { InMemorySandboxInstanceStore } from '../src/instance-store' +import { FULL_CAPS, fakeLog, fromChunkValues, makeFakeHandle } from './fakes' +import type { + AnyTextAdapter, + ModelMessage, + StreamChunk, + StreamDurability, +} from '@tanstack/ai' +import type { SandboxProvider } from '../src/contracts' + +function fakeProvider(): SandboxProvider { + return { + name: 'fake', + capabilities: () => FULL_CAPS, + create: () => Promise.resolve(makeFakeHandle('sbx', 'fake', FULL_CAPS)), + resume: () => Promise.resolve(makeFakeHandle('sbx', 'fake', FULL_CAPS)), + destroy: () => Promise.resolve(), + } +} + +function stamp( + runId: string, + threadId: string, + chunk: StreamChunk, +): StreamChunk { + return { ...chunk, runId, threadId, timestamp: 1 } +} + +/** One harness reply: opening text, a bash call, its result, then the close. */ +function jobChunks(runId: string, threadId: string): Array { + const at = (chunk: StreamChunk) => stamp(runId, threadId, chunk) + return [ + at({ type: EventType.RUN_STARTED, runId, threadId }), + at({ + type: EventType.TEXT_MESSAGE_START, + messageId: 'open', + role: 'assistant', + }), + at({ + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'open', + delta: 'Starting the job.', + }), + at({ type: EventType.TEXT_MESSAGE_END, messageId: 'open' }), + at({ + type: EventType.TOOL_CALL_START, + toolCallId: 'bash-1', + toolCallName: 'bash', + }), + at({ + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'bash-1', + delta: '{"cmd":"sleep"}', + }), + at({ type: EventType.TOOL_CALL_END, toolCallId: 'bash-1' }), + at({ + type: EventType.TOOL_CALL_RESULT, + messageId: 'bash-1-result', + toolCallId: 'bash-1', + content: 'done', + }), + at({ + type: EventType.TEXT_MESSAGE_START, + messageId: 'close', + role: 'assistant', + }), + at({ + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'close', + delta: 'Finished.', + }), + at({ type: EventType.TEXT_MESSAGE_END, messageId: 'close' }), + at({ + type: EventType.RUN_FINISHED, + runId, + threadId, + finishReason: 'stop', + }), + ] +} + +/** The bash call is cut between two argument deltas. */ +function splitChunks(runId: string, threadId: string): Array { + const at = (chunk: StreamChunk) => stamp(runId, threadId, chunk) + return [ + at({ type: EventType.RUN_STARTED, runId, threadId }), + at({ + type: EventType.TOOL_CALL_START, + toolCallId: 'bash-1', + toolCallName: 'bash', + }), + at({ + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'bash-1', + delta: '{"cmd":', + }), + at({ + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'bash-1', + delta: '"sleep"}', + }), + at({ type: EventType.TOOL_CALL_END, toolCallId: 'bash-1' }), + at({ + type: EventType.TOOL_CALL_RESULT, + messageId: 'bash-1-result', + toolCallId: 'bash-1', + content: 'done', + }), + at({ + type: EventType.TEXT_MESSAGE_START, + messageId: 'close', + role: 'assistant', + }), + at({ + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'close', + delta: 'Finished.', + }), + at({ type: EventType.TEXT_MESSAGE_END, messageId: 'close' }), + at({ + type: EventType.RUN_FINISHED, + runId, + threadId, + finishReason: 'stop', + }), + ] +} + +function textDeltas(chunks: ReadonlyArray): Array { + return chunks.flatMap((chunk) => + chunk.type === EventType.TEXT_MESSAGE_CONTENT && + typeof chunk.delta === 'string' + ? [chunk.delta] + : [], + ) +} + +/** Role, text, and tool cards. Ids and timestamps differ between two runs. */ +function outline(messages: ReadonlyArray) { + return messages.map((message) => ({ + role: message.role, + content: message.content ?? null, + ...(message.role === 'tool' ? { toolCallId: message.toolCallId } : {}), + ...(message.toolCalls + ? { + toolCalls: message.toolCalls.map((call) => ({ + id: call.id, + name: call.function.name, + arguments: call.function.arguments, + })), + } + : {}), + })) +} + +function harness(chunks: Array): AnyTextAdapter { + return { + kind: 'text', + name: 'fixed', + model: 'test-model', + '~types': {}, + chatStream: (options: { + runId?: string + threadId?: string + capabilities?: Parameters[0] + }) => { + const durability = options.capabilities + ? getSandboxDurability(options.capabilities, { optional: true }) + : undefined + return alignedIfAttaching(fromChunkValues(chunks), durability) + }, + structuredOutput: () => Promise.resolve({ data: {}, rawText: '{}' }), + } as unknown as AnyTextAdapter +} + +async function drive(input: { + chunks: Array + /** Chunks the first host already appended. Empty for a run that did not restart. */ + stored: number + attach: boolean + threadId: string + runId: string +}): Promise<{ + saved: ReadonlyArray + live: Array + log: Array +}> { + const log = fakeLog(input.chunks.slice(0, input.stored)) + const runs = new InMemoryRunStore() + const persistence = memoryPersistence() + const sandbox = withSandbox( + defineSandbox({ + id: 's', + provider: fakeProvider(), + lifecycle: { snapshot: 'none' }, + fileEvents: false, + }), + { + instances: new InMemorySandboxInstanceStore(), + runs, + durability: { adapter: log, attach: input.attach }, + }, + ) + const live: Array = [] + async function* tee( + source: AsyncIterable, + ): AsyncIterable { + for await (const chunk of source) { + live.push(chunk) + yield chunk + } + } + await pipeToRunLog( + tee( + chat({ + adapter: harness(input.chunks), + messages: [{ role: 'user', content: 'go' }], + threadId: input.threadId, + runId: input.runId, + middleware: [withPersistence(persistence), sandbox], + }), + ), + { + runs, + runId: input.runId, + threadId: input.threadId, + durability: () => log as StreamDurability, + }, + ) + return { + saved: await persistence.stores.messages.loadThread(input.threadId), + live, + log: (await log.snapshot()).map((entry) => entry.chunk), + } +} + +describe('sandbox takeover transcript', () => { + it('keeps the first host text and the bash call in the saved thread', async () => { + const freshChunks = jobChunks('run-fresh', 'thread-fresh') + const fresh = await drive({ + chunks: freshChunks, + stored: 0, + attach: false, + threadId: 'thread-fresh', + runId: 'run-fresh', + }) + const takeoverChunks = jobChunks('run-take', 'thread-take') + // Through TOOL_CALL_END: the opening text and the complete bash call. + const takeover = await drive({ + chunks: takeoverChunks, + stored: 7, + attach: true, + threadId: 'thread-take', + runId: 'run-take', + }) + + expect(outline(takeover.saved)).toEqual(outline(fresh.saved)) + expect(JSON.stringify(outline(takeover.saved))).toContain( + 'Starting the job.', + ) + expect(outline(takeover.saved)).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + toolCalls: [ + { + id: 'bash-1', + name: 'bash', + arguments: '{"cmd":"sleep"}', + }, + ], + }), + ]), + ) + // The client tails the log. The public chat stream is only the rest. + expect(textDeltas(takeover.live)).toEqual(['Finished.']) + expect(textDeltas(takeover.log)).toEqual(['Starting the job.', 'Finished.']) + expect(takeover.log).toHaveLength(takeover.live.length + 7) + }) + + it('joins tool arguments that were cut across the takeover', async () => { + const fresh = await drive({ + chunks: splitChunks('run-fresh-split', 'thread-fresh-split'), + stored: 0, + attach: false, + threadId: 'thread-fresh-split', + runId: 'run-fresh-split', + }) + // RUN_STARTED, TOOL_CALL_START, and the first ARGS delta are already stored. + const takeover = await drive({ + chunks: splitChunks('run-take-split', 'thread-take-split'), + stored: 3, + attach: true, + threadId: 'thread-take-split', + runId: 'run-take-split', + }) + + expect(outline(takeover.saved)).toEqual(outline(fresh.saved)) + expect(outline(takeover.saved)).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + toolCalls: [ + { + id: 'bash-1', + name: 'bash', + arguments: '{"cmd":"sleep"}', + }, + ], + }), + ]), + ) + expect(textDeltas(takeover.live)).toEqual(['Finished.']) + }) +}) diff --git a/testing/e2e/tests/durable-takeover.spec.ts b/testing/e2e/tests/durable-takeover.spec.ts index 08c02feccc..b20d7dc863 100644 --- a/testing/e2e/tests/durable-takeover.spec.ts +++ b/testing/e2e/tests/durable-takeover.spec.ts @@ -1057,11 +1057,9 @@ test.describe('durable runs — the reaper', () => { // conversation never lands at all. This is that transcript, and the // still-producing run's empty one (asserted below) is the control. // - // It holds the REMAINDER, not the whole message: alignment suppresses the - // delivered prefix inside the adapter, so `withPersistence`'s accumulator - // never sees '1','2' — the delivery log is the complete record, the saved - // message is not. That is a property of every takeover, not of the reaper, - // and it is pinned here because it is invisible in the log assertion above. + // The saved message is the whole answer, including deltas the first host + // already put in the log. The log assertion above is the other half: those + // chunks are not appended again. expect( after.messages.map((message) => `${message.role}:${message.text}`), ).toEqual(['user:go', 'assistant:123456'])