fix(ai-sandbox): keep the first host's transcript on takeover

.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.

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()
   })

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)

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()
     })

packages/ai-sandbox/src/align.ts

@@ -195,6 +195,39 @@ export interface AlignToStoredLogOptions<TOffset extends string = string> {
    * {@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<string, unknown>)[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<string, unknown> = { ...chunk, [STORED_REPLAY]: true }
+  return copy as unknown as StreamChunk
 }
 
 /**
@@ -254,6 +287,9 @@ export async function* alignToStoredLog<TOffset extends string = string>(
       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

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<StreamChunk>,
@@ -304,6 +309,7 @@ export function alignedIfAttaching(
   return alignToStoredLog(chunks, {
     durability: durability.adapter,
     isOutOfBand: isBridgeCustomChunk,
+    yieldStored: true,
     ...(logger === undefined ? {} : { logger }),
   })
 }

packages/ai-sandbox/src/index.ts

@@ -310,6 +310,7 @@ export {
 export {
   alignToStoredLog,
   isBridgeCustomChunk,
+  isStoredReplayChunk,
   JournalReplayDivergedError,
   JournalReplayThreadIdMismatchError,
   DEFAULT_MAX_OUT_OF_BAND_SKIP,

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<TOffset extends string = string>(
       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<TOffset extends string = string>(
         await drainWatcher(state, 'pause')
         await stopSnapshotLease(state, { closePortable: true })
       }
+      return undefined
     },
 
     async onFinish(ctx) {

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<TOffset extends string = string>(
 
     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(

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 () => {