fix(cli): recover agent context when run-state.json is torn

cli/src/hooks/use-send-message.ts

@@ -187,7 +187,21 @@ export const useSendMessage = ({
       if (loadedState) {
         previousRunStateRef.current = loadedState.runState
         setRunState(loadedState.runState)
-        setMessages(sanitizeRestoredMessages(loadedState.messages))
+        const restoredMessages = sanitizeRestoredMessages(loadedState.messages)
+        if (loadedState.runStateRestored) {
+          setMessages(restoredMessages)
+        } else {
+          // The agent's context was lost (torn run-state.json, nothing
+          // recoverable) while the transcript survived. Surface it: without
+          // this the model just answers as if the earlier turns never
+          // happened, which reads as the assistant being broken.
+          setMessages([
+            createErrorChatMessage(
+              'The saved agent context could not be restored, so the assistant starts this chat without memory of earlier turns. The transcript below is intact.',
+            ),
+            ...restoredMessages,
+          ])
+        }
         if (loadedState.chatId) {
           setCurrentChatId(loadedState.chatId)
         }

cli/src/utils/__tests__/run-state-storage.test.ts

@@ -917,3 +917,146 @@ describe('poisoned payload persistence', () => {
     expect(block.outputRaw.self).toBe('[Circular]')
   })
 })
+
+describe('run state recovery', () => {
+  // Point persistence at a temp dir via the explicit test override.
+  const chatDir = path.join(TEST_ROOT, 'codebuff-test-recovery')
+
+  const runStateWithSession = (marker: string): RunState =>
+    ({
+      sessionState: {
+        mainAgentState: {
+          messageHistory: [{ role: 'user', content: marker }],
+        },
+      },
+      output: { type: 'lastMessage', value: marker },
+      traceSessionId: 'trace-1',
+    }) as unknown as RunState
+
+  const runStatePath = path.join(chatDir, 'run-state.json')
+  const bakPath = runStatePath + '.bak'
+  const messagesPath = path.join(chatDir, 'chat-messages.json')
+
+  const writePrimary = (contents: string) =>
+    fs.writeFileSync(runStatePath, contents)
+  const validMessages = JSON.stringify([
+    {
+      id: 'msg-1',
+      variant: 'user',
+      content: 'the prompt',
+      timestamp: new Date().toISOString(),
+    },
+  ] as ChatMessage[])
+
+  beforeEach(() => {
+    fs.rmSync(chatDir, { recursive: true, force: true })
+    fs.mkdirSync(chatDir, { recursive: true })
+    setChatDirOverrideForTesting(chatDir)
+  })
+
+  afterEach(() => {
+    setChatDirOverrideForTesting(undefined)
+  })
+
+  test('recovers agent context from the .bak when the primary is torn', () => {
+    writePrimary('{"sessionState": {"main"') // torn: power loss after rename
+    fs.writeFileSync(bakPath, JSON.stringify(runStateWithSession('from-bak')))
+    fs.writeFileSync(messagesPath, validMessages)
+
+    const loaded = loadMostRecentChatState()
+    expect(loaded).not.toBeNull()
+    // Agent context survived — the model is NOT amnesiac next turn.
+    expect((loaded!.runState as any).sessionState).toBeDefined()
+    expect(loaded!.runStateRestored).toBe(true)
+    // Self-healed: the primary is the recovered generation again.
+    expect(
+      JSON.parse(fs.readFileSync(runStatePath, 'utf8')).output.value,
+    ).toBe('from-bak')
+  })
+
+  test('recovers from the newest complete checkpoint temp when bak is absent', () => {
+    writePrimary('{ torn')
+    fs.writeFileSync(messagesPath, validMessages)
+    // Two temps: an older torn one and a newer complete one (SIGKILL between
+    // write and rename leaves the latter behind).
+    fs.writeFileSync(
+      runStatePath + '.999.oldest.tmp',
+      '{"half":',
+    )
+    fs.writeFileSync(
+      runStatePath + '.1234.newest.tmp',
+      JSON.stringify(runStateWithSession('from-tmp')),
+    )
+
+    const loaded = loadMostRecentChatState()
+    expect(loaded).not.toBeNull()
+    expect((loaded!.runState as any).sessionState).toBeDefined()
+    expect(loaded!.runStateRestored).toBe(true)
+    // Self-healed into the primary.
+    expect(
+      JSON.parse(fs.readFileSync(runStatePath, 'utf8')).output.value,
+    ).toBe('from-tmp')
+  })
+
+  test('flags a healthy primary as fully restored', () => {
+    writePrimary(JSON.stringify(runStateWithSession('healthy')))
+    fs.writeFileSync(messagesPath, validMessages)
+
+    const loaded = loadMostRecentChatState()
+    expect(loaded!.runStateRestored).toBe(true)
+    expect((loaded!.runState as any).sessionState).toBeDefined()
+  })
+
+  test('falls back to a context-less placeholder with the loss flagged when nothing recovers', () => {
+    writePrimary('{ torn')
+    fs.writeFileSync(messagesPath, validMessages)
+
+    const loaded = loadMostRecentChatState()
+    expect(loaded).not.toBeNull()
+    // The amnesia carrier: no sessionState — the SDK will start a fresh
+    // session next turn. runStateRestored=false is what makes the UI say so
+    // instead of the model silently forgetting every earlier turn.
+    expect((loaded!.runState as any).sessionState).toBeUndefined()
+    expect(loaded!.runStateRestored).toBe(false)
+    // The transcript still survives.
+    expect(loaded!.messages.length).toBe(1)
+  })
+
+  test('saveChatState rotates the previous primary into .bak', () => {
+    saveChatState(runStateWithSession('generation-1'), [
+      {
+        id: 'msg-1',
+        variant: 'user',
+        content: 'first',
+        timestamp: new Date().toISOString(),
+      },
+    ])
+    expect(fs.existsSync(bakPath)).toBe(false)
+
+    saveChatState(runStateWithSession('generation-2'), [
+      {
+        id: 'msg-2',
+        variant: 'user',
+        content: 'second',
+        timestamp: new Date().toISOString(),
+      },
+    ])
+
+    expect(JSON.parse(fs.readFileSync(runStatePath, 'utf8')).output.value).toBe(
+      'generation-2',
+    )
+    expect(JSON.parse(fs.readFileSync(bakPath, 'utf8')).output.value).toBe(
+      'generation-1',
+    )
+  })
+
+  test('clearChatState removes the backup too', () => {
+    writePrimary(JSON.stringify(runStateWithSession('x')))
+    fs.writeFileSync(bakPath, JSON.stringify(runStateWithSession('bak')))
+    fs.writeFileSync(messagesPath, validMessages)
+
+    clearChatState()
+    expect(fs.existsSync(runStatePath)).toBe(false)
+    expect(fs.existsSync(bakPath)).toBe(false)
+  })
+})

cli/src/utils/run-state-storage.ts

@@ -25,6 +25,11 @@ type SavedChatState = {
   runState: RunState
   messages: ChatMessage[]
   chatId?: string
+  /** False only when run-state.json was unreadable AND nothing could be
+   *  recovered from the backup or checkpoint temps, so the restored RunState
+   *  is a placeholder without agent context — the model starts amnesiac next
+   *  turn. The UI surfaces this so the loss is not silent. */
+  runStateRestored: boolean
 }
 
 type LiveChatState = {
@@ -244,6 +249,121 @@ type SerializedChatState = {
   messagesJson?: string
 }
 
+/** The previous generation of run-state.json, rotated aside by saveChatState.
+ *  loadMostRecentChatState recovers from it when the primary is torn. */
+const RUN_STATE_BACKUP_SUFFIX = '.bak'
+
+function tryReadRunState(filePath: string): RunState | undefined {
+  try {
+    return JSON.parse(fs.readFileSync(filePath, 'utf8')) as RunState
+  } catch {
+    return undefined
+  }
+}
+
+/**
+ * Read run-state.json with best-effort recovery when it is torn.
+ *
+ * The atomic write makes the rename indivisible, but without an fsync a power
+ * loss can still land the rename while the file's data blocks were never
+ * written — leaving a truncated or empty primary (and even with the fsync,
+ * external corruption exists). Two older complete generations can be lying in
+ * the chat directory when that happens:
+ *
+ *  - `run-state.json.bak` — the state the primary replaced at the last
+ *    synchronous save (saveChatState rotates the previous file aside).
+ *  - `run-state.json.<pid>.<uuid>.tmp` — a write that was killed in the
+ *    window between its write and its rename (writeFileAtomic unlinks its own
+ *    temp on error, but a SIGKILL cannot).
+ *
+ * Try the primary, then the backup, then the newest readable temp, and
+ * self-heal the primary from whichever recovered. `fromPrimary` is false when
+ * recovery had to fall back, so the caller can warn that some agent context
+ * is missing instead of losing it silently.
+ */
+function readRunStateWithRecovery(chatDir: string): {
+  runState: RunState
+  fromPrimary: boolean
+} | null {
+  const runStatePath = path.join(chatDir, RUN_STATE_FILENAME)
+  let primaryError: unknown
+  try {
+    return {
+      runState: JSON.parse(fs.readFileSync(runStatePath, 'utf8')) as RunState,
+      fromPrimary: true,
+    }
+  } catch (error) {
+    primaryError = error
+  }
+
+  const bakPath = runStatePath + RUN_STATE_BACKUP_SUFFIX
+  const fromBackup = tryReadRunState(bakPath)
+  if (fromBackup !== undefined) {
+    bestEffortLog(
+      'warn',
+      { runStatePath, bakPath },
+      'run-state.json was unreadable; restored agent context from the previous save',
+    )
+    // Self-heal: make the recovered generation the primary again.
+    try {
+      writeFileAtomic(
+        runStatePath,
+        JSON.stringify(fromBackup),
+      )
+    } catch {
+      // Best-effort; the backup file still exists for the next load.
+    }
+    return { runState: fromBackup, fromPrimary: false }
+  }
+
+  // Newest temp first — each is a complete write that never got renamed.
+  const tmpPaths = fs
+    .readdirSync(chatDir)
+    .filter(
+      (file) =>
+        file.startsWith(RUN_STATE_FILENAME + '.') && file.endsWith('.tmp'),
+    )
+    .map((file) => path.join(chatDir, file))
+    .sort((a, b) => {
+      try {
+        return fs.statSync(b).mtimeMs - fs.statSync(a).mtimeMs
+      } catch {
+        return 0
+      }
+    })
+  for (const tmpPath of tmpPaths) {
+    const recovered = tryReadRunState(tmpPath)
+    if (recovered !== undefined) {
+      bestEffortLog(
+        'warn',
+        { runStatePath, tmpPath },
+        'run-state.json was unreadable; restored agent context from an unfinished checkpoint write',
+      )
+      try {
+        writeFileAtomic(runStatePath, JSON.stringify(recovered))
+      } catch {
+        // Best-effort; the temp file still exists for the next load.
+      }
+      return { runState: recovered, fromPrimary: false }
+    }
+  }
+
+  bestEffortLog(
+    'warn',
+    {
+      runStatePath,
+      primaryError:
+        primaryError instanceof Error
+          ? primaryError.message
+          : String(primaryError),
+      backupTried: bakPath,
+      tempsTried: tmpPaths.length,
+    },
+    'Could not read run state; restoring transcript without agent context',
+  )
+  return null
+}
+
 /**
  * Serialize the two chat-state files independently, so a poisoned run state
  * cannot block persisting the transcript (and vice versa). Cyclic and
@@ -341,10 +461,18 @@ export function saveChatState(
     // since (e.g. the chat deleted from /history mid-run).
     fs.mkdirSync(chatDir, { recursive: true })
     if (serialized.runStateJson) {
-      writeFileAtomic(
-        path.join(chatDir, RUN_STATE_FILENAME),
-        serialized.runStateJson,
-      )
+      const runStatePath = path.join(chatDir, RUN_STATE_FILENAME)
+      // Rotate the previous generation aside before overwriting: it is the
+      // newest complete state readRunStateWithRecovery can fall back to if
+      // this write's rename lands but its data is later found torn.
+      try {
+        if (fs.existsSync(runStatePath)) {
+          fs.renameSync(runStatePath, runStatePath + RUN_STATE_BACKUP_SUFFIX)
+        }
+      } catch {
+        // Rotation is best-effort; the overwrite below still proceeds.
+      }
+      writeFileAtomic(runStatePath, serialized.runStateJson)
     }
     if (serialized.messagesJson) {
       writeFileAtomic(
@@ -507,14 +635,16 @@ export function loadMostRecentChatState(
     // must not lose the transcript, and vice versa. Restore whatever is
     // readable and fall back for the rest.
     let runState: RunState | null = null
-    try {
-      runState = JSON.parse(fs.readFileSync(runStatePath, 'utf8')) as RunState
-    } catch (error) {
+    let runStateRestored = false
+    const recovered = readRunStateWithRecovery(chatDir)
+    if (recovered) {
+      // Agent context is present — whether straight from the primary or
+      // recovered from a backup/temp (both carry a real sessionState).
+      runState = recovered.runState
+      runStateRestored = true
+    } else {
       logger.warn(
-        {
-          runStatePath,
-          error: error instanceof Error ? error.message : String(error),
-        },
+        { runStatePath },
         'Could not read run state; restoring transcript without agent context',
       )
     }
@@ -563,7 +693,12 @@ export function loadMostRecentChatState(
       'Loaded chat state from chat directory',
     )
 
-    return { runState, messages, chatId: resolvedChatId }
+    return {
+      runState,
+      messages,
+      chatId: resolvedChatId,
+      runStateRestored,
+    }
   } catch (error) {
     logger.error(
       {
@@ -583,13 +718,14 @@ export function clearChatState(): void {
     const runStatePath = getRunStatePath()
     const messagesPath = getChatMessagesPath()
     const metaPath = path.join(resolveCurrentChatDir(), CHAT_META_FILENAME)
+    const backupPath = runStatePath + RUN_STATE_BACKUP_SUFFIX
 
-    for (const filePath of [runStatePath, messagesPath, metaPath]) {
+    for (const filePath of [runStatePath, backupPath, messagesPath, metaPath]) {
       fs.rmSync(filePath, { force: true })
     }
 
     logger.debug(
-      { runStatePath, messagesPath, metaPath },
+      { runStatePath, backupPath, messagesPath, metaPath },
       'Cleared chat state files',
     )
   } catch (error) {