From ccc7a2a271d5b2facd3af8fc4dc08f7e6eafd7a6 Mon Sep 17 00:00:00 2001
From: Tanner Linsley <tannerlinsley@gmail.com>
Date: Fri, 11 Sep 2026 13:24:30 -0600
Subject: [PATCH 1/2] Preserve validated Builder progress and distinguish
 incomplete runs

---
 .../builder/BuilderAssistant.client.tsx       |  91 +++++++++++++-
 src/utils/builder-ai-checkpoints.client.ts    |  14 ++-
 src/utils/builder-ai-run-outcome.ts           |  41 +++++++
 src/utils/builder-ai-stream.client.ts         |  42 ++++---
 src/utils/builder-ai.ts                       |  40 +++---
 tests/builder-ai-checkpoints.test.ts          |  33 +++++
 tests/builder-ai-run-outcome.test.ts          |  80 ++++++++++++
 tests/builder-ai-stream.test.ts               | 114 ++++++++++++++++++
 8 files changed, 415 insertions(+), 40 deletions(-)
 create mode 100644 src/utils/builder-ai-run-outcome.ts
 create mode 100644 tests/builder-ai-run-outcome.test.ts

diff --git a/src/components/builder/BuilderAssistant.client.tsx b/src/components/builder/BuilderAssistant.client.tsx
index ede068be2..dfa7c1b53 100644
--- a/src/components/builder/BuilderAssistant.client.tsx
+++ b/src/components/builder/BuilderAssistant.client.tsx
@@ -166,11 +166,13 @@ type TranscriptRow =
       message: string
     }
   | { id: 'error'; kind: 'error'; message: string }
-  | { id: 'checkpoint'; kind: 'checkpoint' }
+  | { id: 'checkpoint'; kind: 'checkpoint'; validated: boolean }
 
 type RollbackCheckpoint = {
   id: string
+  kind?: 'validated'
   execution: BuilderAiExecution
+  expectedExecution?: BuilderAiExecution
   persisted?: BuilderAiCheckpoint
 }
 
@@ -544,7 +546,11 @@ export const BuilderAssistant = React.forwardRef<
     )
     if (error) rows.push({ id: 'error', kind: 'error', message: error })
     if (rollbackCheckpoint) {
-      rows.push({ id: 'checkpoint', kind: 'checkpoint' })
+      rows.push({
+        id: 'checkpoint',
+        kind: 'checkpoint',
+        validated: rollbackCheckpoint.kind === 'validated',
+      })
     }
     return rows
   }, [
@@ -617,6 +623,7 @@ export const BuilderAssistant = React.forwardRef<
         return {
           checkpoint: {
             id: snapshot.checkpoint.id,
+            kind: snapshot.checkpoint.kind,
             execution: snapshot.execution,
             persisted: snapshot.checkpoint,
           },
@@ -1304,8 +1311,19 @@ export const BuilderAssistant = React.forwardRef<
   async function restoreRollbackCheckpoint() {
     const checkpoint = rollbackCheckpointRef.current
     if (!checkpoint) return
+    const generation = checkpointHydrationGenerationRef.current
+    const expectedCurrentExecution = serializeBuilderAiExecution(getExecution())
 
     try {
+      if (
+        checkpoint.expectedExecution &&
+        serializeBuilderAiExecution(checkpoint.expectedExecution) !==
+          serializeBuilderAiExecution(getExecution())
+      ) {
+        await dismissRollbackCheckpoint()
+        setError('The project changed after this checkpoint was created.')
+        return
+      }
       if (
         checkpoint.persisted &&
         !(await builderAiCheckpointMatchesExecution(
@@ -1317,6 +1335,18 @@ export const BuilderAssistant = React.forwardRef<
         setError('The project changed after this checkpoint was created.')
         return
       }
+      if (
+        generation !== checkpointHydrationGenerationRef.current ||
+        rollbackCheckpointRef.current !== checkpoint
+      )
+        return
+      if (
+        serializeBuilderAiExecution(getExecution()) !== expectedCurrentExecution
+      ) {
+        await dismissRollbackCheckpoint()
+        setError('The project changed while the checkpoint was being checked.')
+        return
+      }
       await onRestore(cloneBuilderAiExecution(checkpoint.execution), 'manual')
       await dismissRollbackCheckpoint()
       setError('')
@@ -1942,6 +1972,7 @@ export const BuilderAssistant = React.forwardRef<
     let checkpoint: RollbackCheckpoint | undefined
     let didStageExecution = false
     let lastStagedExecution: BuilderAiExecution | undefined
+    let validatedCheckpoint: RollbackCheckpoint | undefined
     let preserveStoppedExecution = false
     const activityId = pendingPrompt?.runId ?? crypto.randomUUID()
     let currentActivity: BuilderAiActivity | undefined
@@ -2074,6 +2105,7 @@ export const BuilderAssistant = React.forwardRef<
           setError(withTerminalSyncFailure(message, syncFailure))
           return 'error'
         }
+        await discardValidatedCheckpoint()
         stopActivity()
         const syncFailure = await finishDurableRun({ status: 'cancelled' })
         if (syncFailure) {
@@ -2109,6 +2141,7 @@ export const BuilderAssistant = React.forwardRef<
       }
 
       if (response.changedFiles.length === 0 && !response.runtimeChanged) {
+        await discardValidatedCheckpoint()
         completeActivity()
         const assistantMessage = createTranscriptMessage(
           'assistant',
@@ -2159,6 +2192,7 @@ export const BuilderAssistant = React.forwardRef<
         return 'error'
       }
       completeActivity()
+      await discardValidatedCheckpoint()
       const assistantMessage = createTranscriptMessage(
         'assistant',
         response.message,
@@ -2222,6 +2256,7 @@ export const BuilderAssistant = React.forwardRef<
           setError(withTerminalSyncFailure(message, syncFailure))
           return 'error'
         }
+        await discardValidatedCheckpoint()
         stopActivity()
         const syncFailure = await finishDurableRun({ status: 'cancelled' })
         if (syncFailure) {
@@ -2233,7 +2268,17 @@ export const BuilderAssistant = React.forwardRef<
       } else {
         const message = formatError(cause)
         if (preserveStoppedExecution) await discardCheckpoint()
-        else await rollbackStagedExecution()
+        else {
+          await rollbackStagedExecution()
+          if (
+            validatedCheckpoint &&
+            serializeBuilderAiExecution(getExecution()) ===
+              serializeBuilderAiExecution(checkpointExecution)
+          ) {
+            rollbackCheckpointRef.current = validatedCheckpoint
+            setRollbackCheckpoint(validatedCheckpoint)
+          }
+        }
         failActivity(message)
         const syncFailure = await finishDurableRun({
           status: 'failed',
@@ -2454,6 +2499,7 @@ export const BuilderAssistant = React.forwardRef<
         timestamp: Date.now(),
       })
       const runResult = await onApply(state.execution, abortController.signal)
+      abortController.signal.throwIfAborted()
       updateActivity({
         type: 'item-completed',
         runId: activityId,
@@ -2471,6 +2517,31 @@ export const BuilderAssistant = React.forwardRef<
       })
       const completion = validateBuilderAiCompletion(environmentSnapshot)
       if (completion.status === 'complete') {
+        if (
+          serializeBuilderAiExecution(state.execution) !==
+          serializeBuilderAiExecution(checkpointExecution)
+        ) {
+          validatedCheckpoint = {
+            id: `${activityId}:validated`,
+            kind: 'validated',
+            execution: cloneBuilderAiExecution(state.execution),
+            expectedExecution: checkpointExecution,
+          }
+          if (storageScope) {
+            validatedCheckpoint.persisted = await createBuilderAiCheckpoint(
+              storageScope,
+              validatedCheckpoint.id,
+              validatedCheckpoint.execution,
+              { kind: 'validated', expectedExecution: checkpointExecution },
+            )
+          }
+          if (abortController.signal.aborted) {
+            await discardValidatedCheckpoint()
+            abortController.signal.throwIfAborted()
+          }
+        } else {
+          await discardValidatedCheckpoint()
+        }
         updateActivity({
           type: 'item-completed',
           runId: activityId,
@@ -2582,6 +2653,13 @@ export const BuilderAssistant = React.forwardRef<
       setRollbackCheckpoint(checkpoint)
     }
 
+    async function discardValidatedCheckpoint() {
+      if (!validatedCheckpoint) return
+      if (storageScope)
+        await removeBuilderAiCheckpoint(storageScope, validatedCheckpoint.id)
+      validatedCheckpoint = undefined
+    }
+
     async function rollbackStagedExecution() {
       if (!checkpoint || !didStageExecution) return
       if (
@@ -3352,6 +3430,7 @@ function TranscriptRowView({
           />
         ) : row.kind === 'checkpoint' ? (
           <CheckpointNotice
+            validated={row.validated}
             onDismiss={onDismissCheckpoint}
             onRestore={onRestoreCheckpoint}
           />
@@ -3367,16 +3446,20 @@ function TranscriptRowView({
 }
 
 function CheckpointNotice({
+  validated,
   onDismiss,
   onRestore,
 }: {
+  validated: boolean
   onDismiss: () => void
   onRestore: () => void
 }) {
   return (
     <div className="flex items-center justify-between gap-3 rounded-lg border border-border-default bg-background-elevated px-3 py-2">
       <p className="text-xs/5 text-text-secondary">
-        The pre-run builder is available as a checkpoint.
+        {validated
+          ? 'Validated changes are available to restore.'
+          : 'The pre-run builder is available as a checkpoint.'}
       </p>
       <div className="flex shrink-0 items-center gap-1">
         <Button type="button" variant="secondary" size="xs" onClick={onRestore}>
diff --git a/src/utils/builder-ai-checkpoints.client.ts b/src/utils/builder-ai-checkpoints.client.ts
index 0776fb4c0..89843af05 100644
--- a/src/utils/builder-ai-checkpoints.client.ts
+++ b/src/utils/builder-ai-checkpoints.client.ts
@@ -15,6 +15,7 @@ let databasePromise: Promise<IDBDatabase> | undefined
 
 export type BuilderAiCheckpoint = {
   id: string
+  kind?: 'validated'
   createdAt: number
   lastAccessedAt: number
   sizeBytes: number
@@ -42,15 +43,22 @@ export async function createBuilderAiCheckpoint(
   scope: string,
   checkpointId: string,
   execution: BuilderAiExecution,
+  options?: { kind: 'validated'; expectedExecution: BuilderAiExecution },
 ) {
   if (!scope || !checkpointId) return undefined
 
   let serialized: string
   let executionId: string
+  let expectedExecutionId: string
 
   try {
     serialized = serializeCanonicalExecution(execution)
     executionId = await createExecutionId(serialized)
+    expectedExecutionId = options
+      ? await createExecutionId(
+          serializeCanonicalExecution(options.expectedExecution),
+        )
+      : executionId
     await writeExecution(executionId, serialized)
   } catch {
     return undefined
@@ -65,11 +73,12 @@ export async function createBuilderAiCheckpoint(
   const updated: StoredCheckpoint = {
     scope,
     id: checkpointId,
+    ...(options ? { kind: options.kind } : {}),
     createdAt: existing?.createdAt ?? now,
     lastAccessedAt: now,
     sizeBytes: getUtf8ByteLength(serialized),
     executionId,
-    expectedExecutionId: executionId,
+    expectedExecutionId,
   }
   const candidates = [
     updated,
@@ -368,6 +377,7 @@ function serializeCanonicalExecution(execution: BuilderAiExecution) {
 function toCheckpoint(checkpoint: StoredCheckpoint): BuilderAiCheckpoint {
   return {
     id: checkpoint.id,
+    ...(checkpoint.kind ? { kind: checkpoint.kind } : {}),
     createdAt: checkpoint.createdAt,
     lastAccessedAt: checkpoint.lastAccessedAt,
     sizeBytes: checkpoint.sizeBytes,
@@ -381,6 +391,7 @@ function isStoredCheckpoint(value: unknown): value is StoredCheckpoint {
     isRecord(value) &&
     typeof value.scope === 'string' &&
     typeof value.id === 'string' &&
+    (value.kind === undefined || value.kind === 'validated') &&
     typeof value.createdAt === 'number' &&
     Number.isFinite(value.createdAt) &&
     typeof value.lastAccessedAt === 'number' &&
@@ -395,6 +406,7 @@ function isStoredCheckpoint(value: unknown): value is StoredCheckpoint {
     hasOnlyKeys(value, [
       'scope',
       'id',
+      'kind',
       'createdAt',
       'lastAccessedAt',
       'sizeBytes',
diff --git a/src/utils/builder-ai-run-outcome.ts b/src/utils/builder-ai-run-outcome.ts
new file mode 100644
index 000000000..131c9f8e7
--- /dev/null
+++ b/src/utils/builder-ai-run-outcome.ts
@@ -0,0 +1,41 @@
+import type { RunFinishedEvent } from '@tanstack/ai'
+
+type BuilderAiRunOutcome =
+  | { status: 'continue' | 'interrupt' | 'complete' }
+  | { status: 'failed'; message: string }
+
+// A provider turn ending is not necessarily the Builder task completing.
+// Keep the server and browser on the same interpretation of the protocol.
+export function getBuilderAiRunOutcome(
+  event: RunFinishedEvent,
+): BuilderAiRunOutcome {
+  const reason = event.metadata?.tanstack?.finishReason ?? event.finishReason
+  if (event.outcome?.type === 'interrupt') return { status: 'interrupt' }
+  if (event.outcome && event.outcome.type !== 'success') {
+    return {
+      status: 'failed',
+      message: 'The provider stopped the run without completing it.',
+    }
+  }
+  if (reason === 'tool_calls') return { status: 'continue' }
+  if (reason === 'length') {
+    return {
+      status: 'failed',
+      message: 'The model reached its response limit before finishing.',
+    }
+  }
+  if (reason === 'content_filter') {
+    return {
+      status: 'failed',
+      message:
+        'The provider stopped the response because of its content policy.',
+    }
+  }
+  if (reason === 'stop' || (!reason && event.outcome?.type === 'success')) {
+    return { status: 'complete' }
+  }
+  return {
+    status: 'failed',
+    message: 'The provider ended the response without confirming completion.',
+  }
+}
diff --git a/src/utils/builder-ai-stream.client.ts b/src/utils/builder-ai-stream.client.ts
index 8afb8c60b..fe53345dc 100644
--- a/src/utils/builder-ai-stream.client.ts
+++ b/src/utils/builder-ai-stream.client.ts
@@ -1,4 +1,5 @@
 import type { StreamChunk } from '@tanstack/ai'
+import { getBuilderAiRunOutcome } from './builder-ai-run-outcome'
 import {
   ChatClient,
   clientTools,
@@ -249,9 +250,9 @@ export async function runBuilderAiStream({
           chunk.outcome?.type === 'interrupt' &&
           !receivedValidationState
         ) {
-          rejectValidationState?.(
-            new Error('Builder AI returned no validation state'),
-          )
+          const error = new Error('Builder AI returned no validation state')
+          rejectValidationState?.(error)
+          throw error
         }
         if (chunk.type === 'RUN_ERROR') {
           rejectValidationState?.(new Error(chunk.message))
@@ -261,7 +262,10 @@ export async function runBuilderAiStream({
     const updateClientForwardedProps = (
       nextForwardedProps: Record<string, unknown>,
     ) => client.updateOptions({ forwardedProps: nextForwardedProps })
-    const abort = () => client.stop()
+    const abort = () => {
+      rejectValidationState?.(new DOMException('Aborted', 'AbortError'))
+      client.stop()
+    }
     signal.addEventListener('abort', abort, { once: true })
 
     try {
@@ -299,10 +303,16 @@ export async function runBuilderAiStream({
     }
 
     function createValidationStatePromise() {
-      return new Promise<BuilderAiValidationState>((resolve, reject) => {
-        resolveValidationState = resolve
-        rejectValidationState = reject
-      })
+      const promise = new Promise<BuilderAiValidationState>(
+        (resolve, reject) => {
+          resolveValidationState = resolve
+          rejectValidationState = reject
+        },
+      )
+      // A run can fail before its validation tool starts awaiting this state.
+      // Mark it handled without changing the rejection seen by that tool.
+      void promise.catch(() => undefined)
+      return promise
     }
   }
 
@@ -358,19 +368,21 @@ export async function runBuilderAiStream({
     }
     if (chunk.type === 'RUN_ERROR') throw new Error(chunk.message)
     if (chunk.type !== 'RUN_FINISHED') return
-    if (
-      (chunk.metadata?.tanstack?.finishReason ?? chunk.finishReason) ===
-      'tool_calls'
-    ) {
+    const outcome = getBuilderAiRunOutcome(chunk)
+    if (outcome.status === 'failed') throw new Error(outcome.message)
+    if (outcome.status === 'interrupt') {
       if (pendingResult) {
         throw new Error('Builder AI returned a partial execution result')
       }
-      return
-    }
-    if (chunk.outcome?.type === 'interrupt') {
       if (allowInterrupt) return
       throw new Error('Builder AI was interrupted')
     }
+    if (outcome.status === 'continue') {
+      if (pendingResult) {
+        throw new Error('Builder AI returned a partial execution result')
+      }
+      return
+    }
     if (!pendingResult) {
       throw new Error('Builder AI returned no execution result')
     }
diff --git a/src/utils/builder-ai.ts b/src/utils/builder-ai.ts
index 49203ceff..7337f1411 100644
--- a/src/utils/builder-ai.ts
+++ b/src/utils/builder-ai.ts
@@ -1,4 +1,5 @@
 import { EventType, type StreamChunk } from '@tanstack/ai'
+import { getBuilderAiRunOutcome } from './builder-ai-run-outcome'
 import { isExampleRuntime } from './example-project'
 import {
   parseExampleWorkspace,
@@ -204,29 +205,28 @@ export async function* streamBuilderAiResponse(
               }
             : {}),
         }
-        continue
+        return
       }
 
-      if (
-        chunk.type === 'RUN_FINISHED' &&
-        (chunk.metadata?.tanstack?.finishReason ?? chunk.finishReason) !==
-          'tool_calls' &&
-        chunk.outcome?.type !== 'interrupt'
-      ) {
-        yield {
-          type: 'CUSTOM',
-          name: 'builder.project.execution',
-          value: createResponse(message.trim() || 'Builder changes are ready.'),
+      if (chunk.type === 'RUN_FINISHED') {
+        const outcome = getBuilderAiRunOutcome(chunk)
+        if (outcome.status === 'failed') {
+          sawTerminal = true
+          yield { type: EventType.RUN_ERROR, message: outcome.message }
+          return
+        }
+        if (outcome.status === 'complete') {
+          yield {
+            type: 'CUSTOM',
+            name: 'builder.project.execution',
+            value: createResponse(
+              message.trim() || 'Builder changes are ready.',
+            ),
+          }
+        }
+        if (outcome.status !== 'continue') {
+          sawTerminal = true
         }
-      }
-
-      if (
-        chunk.type === 'RUN_FINISHED' &&
-        ((chunk.metadata?.tanstack?.finishReason ?? chunk.finishReason) !==
-          'tool_calls' ||
-          chunk.outcome?.type === 'interrupt')
-      ) {
-        sawTerminal = true
       }
 
       yield chunk
diff --git a/tests/builder-ai-checkpoints.test.ts b/tests/builder-ai-checkpoints.test.ts
index fe5283107..910be8871 100644
--- a/tests/builder-ai-checkpoints.test.ts
+++ b/tests/builder-ai-checkpoints.test.ts
@@ -327,6 +327,39 @@ test('builder AI checkpoints deduplicate, roll back, and stay LRU bounded', asyn
       fallbackExecution,
     )
 
+    const validated = createExecution('export default "validated candidate"')
+    const recovery = await createBuilderAiCheckpoint(
+      'builder-recovery',
+      'run:validated',
+      validated,
+      { kind: 'validated', expectedExecution: fallbackExecution },
+    )
+    assert.ok(recovery)
+    assert.equal(recovery.kind, 'validated')
+    assert.equal(
+      await builderAiCheckpointMatchesExecution(recovery, fallbackExecution),
+      true,
+    )
+    assert.equal(
+      await builderAiCheckpointMatchesExecution(recovery, validated),
+      false,
+    )
+    assert.equal(
+      await builderAiCheckpointMatchesExecution(
+        recovery,
+        createExecution('manual edit'),
+      ),
+      false,
+    )
+    const recovered = await loadLatestBuilderAiCheckpoint('builder-recovery')
+    assert.equal(recovered?.checkpoint.kind, 'validated')
+    assert.deepEqual(recovered?.execution, validated)
+    assert.equal(
+      await loadLatestBuilderAiCheckpoint('another-owner'),
+      undefined,
+    )
+    await removeBuilderAiCheckpoint('builder-recovery', 'run:validated')
+
     for (let index = 0; index < 51; index += 1) {
       now += 1
       await createBuilderAiCheckpoint(
diff --git a/tests/builder-ai-run-outcome.test.ts b/tests/builder-ai-run-outcome.test.ts
new file mode 100644
index 000000000..0ca93a594
--- /dev/null
+++ b/tests/builder-ai-run-outcome.test.ts
@@ -0,0 +1,80 @@
+import assert from 'node:assert/strict'
+import test from 'node:test'
+import {
+  EventType,
+  type RunFinishedEvent,
+  type StreamChunk,
+} from '@tanstack/ai'
+import { getBuilderAiRunOutcome } from '../src/utils/builder-ai-run-outcome'
+import { streamBuilderAiResponse } from '../src/utils/builder-ai'
+
+function finish(fields: Partial<RunFinishedEvent>): RunFinishedEvent {
+  return {
+    type: EventType.RUN_FINISHED,
+    threadId: 'thread',
+    runId: 'run',
+    ...fields,
+  }
+}
+
+test('Builder distinguishes successful completion, tool continuation, and incomplete responses', () => {
+  assert.equal(
+    getBuilderAiRunOutcome(finish({ finishReason: 'stop' })).status,
+    'complete',
+  )
+  assert.equal(
+    getBuilderAiRunOutcome(finish({ outcome: { type: 'success' } })).status,
+    'complete',
+  )
+  assert.equal(
+    getBuilderAiRunOutcome(finish({ finishReason: 'tool_calls' })).status,
+    'continue',
+  )
+  assert.equal(getBuilderAiRunOutcome(finish({})).status, 'failed')
+  for (const event of [
+    finish({ finishReason: 'length', outcome: { type: 'success' } }),
+    finish({ finishReason: 'content_filter' }),
+    finish({
+      finishReason: 'stop',
+      metadata: { tanstack: { finishReason: 'length' } },
+    }),
+  ])
+    assert.equal(getBuilderAiRunOutcome(event).status, 'failed')
+})
+
+test('the server never creates a successful project result for an incomplete model response', async () => {
+  for (const event of [
+    finish({ finishReason: 'length' }),
+    finish({ finishReason: 'content_filter' }),
+    finish({}),
+  ]) {
+    async function* source(): AsyncGenerator<StreamChunk> {
+      yield event
+    }
+    const chunks: StreamChunk[] = []
+    for await (const chunk of streamBuilderAiResponse(
+      source(),
+      'secret',
+      () => {
+        throw new Error('An incomplete run must not produce a project result')
+      },
+    ))
+      chunks.push(chunk)
+    assert.equal(chunks.length, 1)
+    assert.equal(chunks[0]?.type, EventType.RUN_ERROR)
+  }
+})
+
+test('a provider error cannot be followed by a successful Builder result', async () => {
+  async function* source(): AsyncGenerator<StreamChunk> {
+    yield { type: EventType.RUN_ERROR, message: 'Connection lost' }
+    yield finish({ finishReason: 'stop' })
+  }
+  const chunks: StreamChunk[] = []
+  for await (const chunk of streamBuilderAiResponse(source(), 'secret', () => {
+    throw new Error('The run already failed')
+  }))
+    chunks.push(chunk)
+  assert.equal(chunks.length, 1)
+  assert.equal(chunks[0]?.type, EventType.RUN_ERROR)
+})
diff --git a/tests/builder-ai-stream.test.ts b/tests/builder-ai-stream.test.ts
index aa7bfdc14..4fc17f644 100644
--- a/tests/builder-ai-stream.test.ts
+++ b/tests/builder-ai-stream.test.ts
@@ -101,6 +101,46 @@ test('builder stream rejects execution without a final run event', async () => {
   })
 })
 
+test('builder stream does not accept a token-limited response as success', async () => {
+  await withStream(
+    [
+      runStarted,
+      executionEvent,
+      {
+        ...finalEvent,
+        metadata: { tanstack: { finishReason: 'length' } },
+      },
+    ],
+    async () => {
+      await assert.rejects(runStream(), /response limit/)
+    },
+  )
+})
+
+test('missing validation state fails without running a tool or an unhandled rejection', async () => {
+  const chunks = nativeValidationStream({
+    execution: response.execution,
+    index: 0,
+    runId: 'run-1',
+  }).filter((chunk) => chunk.type !== 'STATE_SNAPSHOT')
+  await withStream(chunks, async () => {
+    await assert.rejects(
+      runBuilderAiStream({
+        endpoint: 'http://builder.test/assist',
+        forwardedProps: { execution: response.execution },
+        messages: [{ role: 'user', content: 'Update the builder.' }],
+        signal: new AbortController().signal,
+        threadId: 'thread-1',
+        activityId: 'activity-1',
+        onValidate: async () => {
+          throw new Error('Validation must not run without state')
+        },
+      }),
+      /no validation state/,
+    )
+  })
+})
+
 test('builder stream rejects events between execution and the final run event', async () => {
   await withStream(
     [
@@ -386,6 +426,80 @@ test('builder stream validates and repairs more than twice inside one client-too
   })
 })
 
+test('a lost continuation does not rerun successful browser validation', async () => {
+  const requests: Array<Record<string, unknown>> = []
+  let validations = 0
+  await withStreams(
+    [
+      (body) =>
+        nativeValidationStream({
+          execution: response.execution,
+          index: 0,
+          runId: readRequestRunId(body),
+        }),
+      () => {
+        throw new TypeError('Network disconnected')
+      },
+    ],
+    requests,
+    [],
+    async () => {
+      await assert.rejects(
+        runBuilderAiStream({
+          endpoint: 'http://builder.test/assist',
+          forwardedProps: { execution: response.execution },
+          messages: [{ role: 'user', content: 'Update the builder.' }],
+          signal: new AbortController().signal,
+          threadId: 'thread-1',
+          activityId: 'activity-1',
+          onValidate: async () => {
+            validations++
+            return { result: { status: 'complete' } }
+          },
+        }),
+        /Stream response body read failed/,
+      )
+    },
+  )
+  assert.equal(validations, 1)
+  assert.equal(requests.length, 2)
+})
+
+test('cancelling during validation never submits a continuation', async () => {
+  const requests: Array<Record<string, unknown>> = []
+  const controller = new AbortController()
+  await withStreams(
+    [
+      (body) =>
+        nativeValidationStream({
+          execution: response.execution,
+          index: 0,
+          runId: readRequestRunId(body),
+        }),
+    ],
+    requests,
+    [],
+    async () => {
+      await assert.rejects(
+        runBuilderAiStream({
+          endpoint: 'http://builder.test/assist',
+          forwardedProps: { execution: response.execution },
+          messages: [{ role: 'user', content: 'Update the builder.' }],
+          signal: controller.signal,
+          threadId: 'thread-1',
+          activityId: 'activity-1',
+          onValidate: async () => {
+            controller.abort()
+            return { result: { status: 'complete' } }
+          },
+        }),
+        { name: 'AbortError' },
+      )
+    },
+  )
+  assert.equal(requests.length, 1)
+})
+
 test('builder stream rejects a changed execution that was not validated', async () => {
   const baselineExecution = {
     runtime: null,

From 89ad0ec02d0ab63ee30a242f64624f2414f642e0 Mon Sep 17 00:00:00 2001
From: Tanner Linsley <tannerlinsley@gmail.com>
Date: Fri, 11 Sep 2026 13:51:42 -0600
Subject: [PATCH 2/2] Harden Builder checkpoint recovery and storage migration

---
 .../builder/BuilderAssistant.client.tsx       |  25 ++--
 src/utils/builder-ai-checkpoints.client.ts    | 123 +++++++++++++++---
 tests/builder-ai-checkpoints.test.ts          |  81 +++++++++++-
 3 files changed, 193 insertions(+), 36 deletions(-)

diff --git a/src/components/builder/BuilderAssistant.client.tsx b/src/components/builder/BuilderAssistant.client.tsx
index dfa7c1b53..12feff5ff 100644
--- a/src/components/builder/BuilderAssistant.client.tsx
+++ b/src/components/builder/BuilderAssistant.client.tsx
@@ -605,7 +605,10 @@ export const BuilderAssistant = React.forwardRef<
       return
     }
 
-    const hydration = loadLatestBuilderAiCheckpoint(storageScope)
+    const hydration = loadLatestBuilderAiCheckpoint(
+      storageScope,
+      getExecutionRef.current(),
+    )
       .then(async (snapshot) => {
         if (!snapshot) return undefined
         const currentExecution = getExecutionRef.current()
@@ -1324,22 +1327,22 @@ export const BuilderAssistant = React.forwardRef<
         setError('The project changed after this checkpoint was created.')
         return
       }
-      if (
-        checkpoint.persisted &&
-        !(await builderAiCheckpointMatchesExecution(
+      const matchesExpected =
+        !checkpoint.persisted ||
+        (await builderAiCheckpointMatchesExecution(
           checkpoint.persisted,
           getExecution(),
         ))
-      ) {
-        await dismissRollbackCheckpoint()
-        setError('The project changed after this checkpoint was created.')
-        return
-      }
       if (
         generation !== checkpointHydrationGenerationRef.current ||
         rollbackCheckpointRef.current !== checkpoint
       )
         return
+      if (!matchesExpected) {
+        await dismissRollbackCheckpoint()
+        setError('The project changed after this checkpoint was created.')
+        return
+      }
       if (
         serializeBuilderAiExecution(getExecution()) !== expectedCurrentExecution
       ) {
@@ -2141,7 +2144,6 @@ export const BuilderAssistant = React.forwardRef<
       }
 
       if (response.changedFiles.length === 0 && !response.runtimeChanged) {
-        await discardValidatedCheckpoint()
         completeActivity()
         const assistantMessage = createTranscriptMessage(
           'assistant',
@@ -2161,6 +2163,7 @@ export const BuilderAssistant = React.forwardRef<
           )
           return 'error'
         }
+        await discardValidatedCheckpoint()
         await discardCheckpoint()
         return 'success'
       }
@@ -2192,7 +2195,6 @@ export const BuilderAssistant = React.forwardRef<
         return 'error'
       }
       completeActivity()
-      await discardValidatedCheckpoint()
       const assistantMessage = createTranscriptMessage(
         'assistant',
         response.message,
@@ -2213,6 +2215,7 @@ export const BuilderAssistant = React.forwardRef<
         )
         return 'error'
       }
+      await discardValidatedCheckpoint()
       await discardCheckpoint()
       return 'success'
     } catch (cause) {
diff --git a/src/utils/builder-ai-checkpoints.client.ts b/src/utils/builder-ai-checkpoints.client.ts
index 89843af05..ca35ea1ae 100644
--- a/src/utils/builder-ai-checkpoints.client.ts
+++ b/src/utils/builder-ai-checkpoints.client.ts
@@ -4,8 +4,10 @@ import {
   type BuilderAiExecution,
 } from './builder-ai'
 
-const metadataStorageKey = 'tanstack-builder-ai:checkpoints:v1'
-const databaseName = 'tanstack-builder-ai-checkpoints'
+const metadataStorageKey = 'tanstack-builder-ai:checkpoints:v2'
+const databaseName = 'tanstack-builder-ai-checkpoints-v2'
+const legacyMetadataStorageKey = 'tanstack-builder-ai:checkpoints:v1'
+const legacyDatabaseName = 'tanstack-builder-ai-checkpoints'
 const objectStoreName = 'executions'
 const maxCheckpointCount = 50
 const maxStoredBytes = 50 * 1024 * 1024
@@ -129,6 +131,7 @@ export async function updateBuilderAiCheckpointExpectedExecution(
 
   let expectedExecutionId: string
   try {
+    await getCheckpointDatabase()
     expectedExecutionId = await createExecutionId(
       serializeCanonicalExecution(execution),
     )
@@ -180,6 +183,11 @@ export async function loadBuilderAiCheckpoint(
   scope: string,
   checkpointId: string,
 ) {
+  try {
+    await getCheckpointDatabase()
+  } catch {
+    return undefined
+  }
   const checkpoints = readCheckpointIndex()
   const checkpoint = checkpoints.find(
     (candidate) => candidate.scope === scope && candidate.id === checkpointId,
@@ -198,12 +206,30 @@ export async function loadBuilderAiCheckpoint(
 
 export async function loadLatestBuilderAiCheckpoint(
   scope: string,
+  currentExecution?: BuilderAiExecution,
 ): Promise<BuilderAiCheckpointSnapshot | undefined> {
+  try {
+    await getCheckpointDatabase()
+  } catch {
+    return undefined
+  }
   const checkpoints = listBuilderAiCheckpoints(scope)
 
   for (const checkpoint of checkpoints) {
     const execution = await loadBuilderAiCheckpoint(scope, checkpoint.id)
-    if (execution) return { checkpoint, execution }
+    if (!execution) continue
+    if (
+      currentExecution &&
+      (!(await builderAiCheckpointMatchesExecution(
+        checkpoint,
+        currentExecution,
+      )) ||
+        serializeBuilderAiExecution(execution) ===
+          serializeBuilderAiExecution(currentExecution))
+    ) {
+      continue
+    }
+    return { checkpoint, execution }
   }
 
   return undefined
@@ -213,6 +239,11 @@ export async function removeBuilderAiCheckpoint(
   scope: string,
   checkpointId: string,
 ) {
+  try {
+    await getCheckpointDatabase()
+  } catch {
+    return
+  }
   const checkpoints = readCheckpointIndex()
   const checkpoint = checkpoints.find(
     (candidate) => candidate.scope === scope && candidate.id === checkpointId,
@@ -293,22 +324,22 @@ async function removeUnreferencedExecutions(
   )
 }
 
-function readCheckpointIndex() {
+function readCheckpointIndex(storageKey = metadataStorageKey, version = 2) {
   const storage = getLocalStorage()
   if (!storage) return []
 
   try {
-    const serialized = storage.getItem(metadataStorageKey)
+    const serialized = storage.getItem(storageKey)
     if (!serialized) return []
 
     const value: unknown = JSON.parse(serialized)
     if (
       !isRecord(value) ||
-      value.version !== 1 ||
+      value.version !== version ||
       !Array.isArray(value.checkpoints) ||
       !value.checkpoints.every(isStoredCheckpoint)
     ) {
-      storage.removeItem(metadataStorageKey)
+      storage.removeItem(storageKey)
       return []
     }
 
@@ -323,7 +354,7 @@ function readCheckpointIndex() {
     return Array.from(deduplicated.values())
   } catch {
     try {
-      storage.removeItem(metadataStorageKey)
+      storage.removeItem(storageKey)
     } catch {
       // Browser storage is best-effort and must not interrupt builder editing.
     }
@@ -338,7 +369,7 @@ function writeCheckpointIndex(checkpoints: ReadonlyArray<StoredCheckpoint>) {
   try {
     storage.setItem(
       metadataStorageKey,
-      JSON.stringify({ version: 1, checkpoints }),
+      JSON.stringify({ version: 2, checkpoints }),
     )
     return true
   } catch {
@@ -447,6 +478,14 @@ async function runExecutionRequest<TResult>(
   createRequest: (store: IDBObjectStore) => IDBRequest<TResult>,
 ) {
   const database = await getCheckpointDatabase()
+  return runDatabaseRequest(database, mode, createRequest)
+}
+
+function runDatabaseRequest<TResult>(
+  database: IDBDatabase,
+  mode: IDBTransactionMode,
+  createRequest: (store: IDBObjectStore) => IDBRequest<TResult>,
+) {
   return new Promise<TResult>((resolve, reject) => {
     let result: TResult
 
@@ -477,7 +516,56 @@ async function runExecutionRequest<TResult>(
 function getCheckpointDatabase() {
   if (databasePromise) return databasePromise
 
-  databasePromise = new Promise<IDBDatabase>((resolve, reject) => {
+  databasePromise = openCheckpointDatabase(databaseName)
+    .then(async (database) => {
+      try {
+        await migrateLegacyCheckpoints(database)
+        return database
+      } catch (error) {
+        database.close()
+        throw error
+      }
+    })
+    .catch((error: unknown) => {
+      databasePromise = undefined
+      throw error
+    })
+  return databasePromise
+}
+
+async function migrateLegacyCheckpoints(database: IDBDatabase) {
+  const storage = getLocalStorage()
+  if (!storage || storage.getItem(metadataStorageKey) !== null) return
+  const checkpoints = readCheckpointIndex(legacyMetadataStorageKey, 1)
+  if (checkpoints.length) {
+    const legacyDatabase = await openCheckpointDatabase(legacyDatabaseName)
+    try {
+      for (const executionId of new Set(
+        checkpoints.map((item) => item.executionId),
+      )) {
+        const value = await runDatabaseRequest(
+          legacyDatabase,
+          'readonly',
+          (store) => store.get(executionId),
+        )
+        if (typeof value === 'string') {
+          await runDatabaseRequest(database, 'readwrite', (store) =>
+            store.put(value, executionId),
+          )
+        }
+      }
+    } finally {
+      legacyDatabase.close()
+    }
+  }
+  // Another tab may have finished migration and added checkpoints while we copied.
+  if (storage.getItem(metadataStorageKey) !== null) return
+  if (!writeCheckpointIndex(checkpoints))
+    throw new Error('Could not migrate builder checkpoints')
+}
+
+function openCheckpointDatabase(name: string) {
+  return new Promise<IDBDatabase>((resolve, reject) => {
     const factory = globalThis.indexedDB
     if (!factory) {
       reject(new Error('indexedDB is not available in this environment.'))
@@ -485,7 +573,7 @@ function getCheckpointDatabase() {
     }
 
     let openFailed = false
-    const request = factory.open(databaseName, 1)
+    const request = factory.open(name, 1)
     request.onupgradeneeded = () => {
       if (!request.result.objectStoreNames.contains(objectStoreName)) {
         request.result.createObjectStore(objectStoreName)
@@ -493,13 +581,11 @@ function getCheckpointDatabase() {
     }
     request.onerror = () => {
       openFailed = true
-      reject(request.error ?? new Error(`Failed to open ${databaseName}.`))
+      reject(request.error ?? new Error(`Failed to open ${name}.`))
     }
     request.onblocked = () => {
       openFailed = true
-      reject(
-        new Error(`Opening IndexedDB database "${databaseName}" was blocked.`),
-      )
+      reject(new Error(`Opening IndexedDB database "${name}" was blocked.`))
     }
     request.onsuccess = () => {
       const database = request.result
@@ -509,14 +595,9 @@ function getCheckpointDatabase() {
       }
       database.onversionchange = () => {
         database.close()
-        databasePromise = undefined
+        if (name === databaseName) databasePromise = undefined
       }
       resolve(database)
     }
-  }).catch((error: unknown) => {
-    databasePromise = undefined
-    throw error
   })
-
-  return databasePromise
 }
diff --git a/tests/builder-ai-checkpoints.test.ts b/tests/builder-ai-checkpoints.test.ts
index 910be8871..ba5fd04d3 100644
--- a/tests/builder-ai-checkpoints.test.ts
+++ b/tests/builder-ai-checkpoints.test.ts
@@ -9,7 +9,11 @@ import {
   removeBuilderAiCheckpoint,
   updateBuilderAiCheckpointExpectedExecution,
 } from '../src/utils/builder-ai-checkpoints.client'
-import type { BuilderAiExecution } from '../src/utils/builder-ai'
+import {
+  parseBuilderAiExecution,
+  serializeBuilderAiExecution,
+  type BuilderAiExecution,
+} from '../src/utils/builder-ai'
 
 class MemoryStorage implements Storage {
   readonly values = new Map<string, string>()
@@ -116,9 +120,16 @@ class FakeOpenRequest {
 
 class FakeIndexedDb {
   readonly values = new Map<string, unknown>()
-
-  open() {
-    const request = new FakeOpenRequest(new FakeDatabase(this.values))
+  readonly legacyValues = new Map<string, unknown>()
+
+  open(name: string) {
+    const request = new FakeOpenRequest(
+      new FakeDatabase(
+        name === 'tanstack-builder-ai-checkpoints'
+          ? this.legacyValues
+          : this.values,
+      ),
+    )
     queueMicrotask(() => {
       request.onupgradeneeded?.()
       request.onsuccess?.()
@@ -172,6 +183,47 @@ test('builder AI checkpoints deduplicate, roll back, and stay LRU bounded', asyn
   Date.now = () => now
 
   try {
+    const legacyExecution = parseBuilderAiExecution(
+      createExecution('export default "legacy"'),
+    )
+    const legacySerialized = serializeBuilderAiExecution(legacyExecution)
+    const legacyBytes = new TextEncoder().encode(legacySerialized)
+    const legacyId = Array.from(
+      new Uint8Array(await crypto.subtle.digest('SHA-256', legacyBytes)),
+      (byte) => byte.toString(16).padStart(2, '0'),
+    ).join('')
+    const legacyIndex = JSON.stringify({
+      version: 1,
+      checkpoints: [
+        {
+          scope: 'legacy-project',
+          id: 'legacy-run',
+          createdAt: now,
+          lastAccessedAt: now,
+          sizeBytes: legacyBytes.byteLength,
+          executionId: legacyId,
+          expectedExecutionId: legacyId,
+        },
+      ],
+    })
+    localStorage.setItem('tanstack-builder-ai:checkpoints:v1', legacyIndex)
+    indexedDb.legacyValues.set(legacyId, legacySerialized)
+    assert.deepEqual(
+      (await loadLatestBuilderAiCheckpoint('legacy-project'))?.execution,
+      legacyExecution,
+    )
+    assert.equal(
+      localStorage.getItem('tanstack-builder-ai:checkpoints:v1'),
+      legacyIndex,
+    )
+    // An older tab can delete its own copy without removing the migrated checkpoint.
+    indexedDb.legacyValues.clear()
+    assert.deepEqual(
+      await loadBuilderAiCheckpoint('legacy-project', 'legacy-run'),
+      legacyExecution,
+    )
+    await removeBuilderAiCheckpoint('legacy-project', 'legacy-run')
+    localStorage.removeItem('tanstack-builder-ai:checkpoints:v1')
     const firstExecution = createExecution('export default "first"', {
       '/other.ts': 'export const other = true',
       '/index.tsx': 'export default "first"',
@@ -354,6 +406,27 @@ test('builder AI checkpoints deduplicate, roll back, and stay LRU bounded', asyn
     const recovered = await loadLatestBuilderAiCheckpoint('builder-recovery')
     assert.equal(recovered?.checkpoint.kind, 'validated')
     assert.deepEqual(recovered?.execution, validated)
+    now += 1
+    await createBuilderAiCheckpoint(
+      'builder-recovery',
+      'newer-unrelated',
+      firstExecution,
+    )
+    assert.equal(
+      (
+        await loadLatestBuilderAiCheckpoint(
+          'builder-recovery',
+          fallbackExecution,
+        )
+      )?.checkpoint.id,
+      'run:validated',
+    )
+    assert.equal(
+      await loadLatestBuilderAiCheckpoint('builder-recovery', validated),
+      undefined,
+    )
+    assert.equal(listBuilderAiCheckpoints('builder-recovery').length, 2)
+    await removeBuilderAiCheckpoint('builder-recovery', 'newer-unrelated')
     assert.equal(
       await loadLatestBuilderAiCheckpoint('another-owner'),
       undefined,