Recover agent context from torn run-state.json backups

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

@@ -1,4 +1,12 @@
-import { describe, test, expect, afterAll, beforeEach, afterEach, mock } from 'bun:test'
+import {
+  describe,
+  test,
+  expect,
+  afterAll,
+  beforeEach,
+  afterEach,
+  mock,
+} from 'bun:test'
 import * as fs from 'fs'
 import * as path from 'path'
 import * as os from 'os'
@@ -917,3 +925,174 @@ 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). mtimes are pinned because
+    // back-to-back writes can land inside one mtime quantum on some
+    // filesystems, which would make "newest" nondeterministic.
+    const oldTemp = runStatePath + '.999.oldest.tmp'
+    const newTemp = runStatePath + '.1234.newest.tmp'
+    fs.writeFileSync(oldTemp, '{"half":')
+    fs.writeFileSync(newTemp, JSON.stringify(runStateWithSession('from-tmp')))
+    const now = new Date()
+    fs.utimesSync(
+      oldTemp,
+      new Date(now.getTime() - 10_000),
+      new Date(now.getTime() - 10_000),
+    )
+    fs.utimesSync(newTemp, now, now)
+
+    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('prefers a newer complete checkpoint temp over the .bak', () => {
+    // Both fallbacks are intact generations; the newest one lost the least
+    // agent context, so recency — not a fixed .bak-first order — decides.
+    writePrimary('{ torn')
+    fs.writeFileSync(messagesPath, validMessages)
+    fs.writeFileSync(bakPath, JSON.stringify(runStateWithSession('from-bak')))
+    const tempPath = runStatePath + '.1234.checkpoint.tmp'
+    fs.writeFileSync(tempPath, JSON.stringify(runStateWithSession('from-tmp')))
+    const now = new Date()
+    fs.utimesSync(
+      bakPath,
+      new Date(now.getTime() - 20_000),
+      new Date(now.getTime() - 20_000),
+    )
+    fs.utimesSync(tempPath, now, now)
+
+    const loaded = loadMostRecentChatState()
+    expect(loaded).not.toBeNull()
+    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)
+  })
+})