diff --git a/cli/src/hooks/use-send-message.ts b/cli/src/hooks/use-send-message.ts index 698859a44e..de9d5ba263 100644 --- a/cli/src/hooks/use-send-message.ts +++ b/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) } diff --git a/cli/src/utils/__tests__/run-state-storage.test.ts b/cli/src/utils/__tests__/run-state-storage.test.ts index 5fa887cf41..a4d532058d 100644 --- a/cli/src/utils/__tests__/run-state-storage.test.ts +++ b/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) + }) +}) diff --git a/cli/src/utils/run-state-storage.ts b/cli/src/utils/run-state-storage.ts index 698b503fbd..254f08806d 100644 --- a/cli/src/utils/run-state-storage.ts +++ b/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,120 @@ 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...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 every fallback candidate — the .bak and each leftover + * checkpoint .tmp — ordered newest-first by mtime (path comparison breaks + * ties on coarse-mtime filesystems), 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 + } + + // Candidates in recency order, newest first. Each .tmp is a complete write + // that never got renamed; the .bak is the previous completed generation. + // mtimeMs alone can be coarse (FAT, some CI filesystems quantize it), so + // ties fall back to path comparison, and the recency probe is best-effort — + // an unreadable candidate is skipped by the loop below, not fatal. + const bakPath = runStatePath + RUN_STATE_BACKUP_SUFFIX + const candidates = [bakPath] + try { + candidates.push( + ...fs + .readdirSync(chatDir) + .filter( + (file) => + file.startsWith(RUN_STATE_FILENAME + '.') && file.endsWith('.tmp'), + ) + .map((file) => path.join(chatDir, file)), + ) + } catch { + // Directory vanished mid-load; the loop below just finds nothing extra. + } + const candidatesWithMtime = candidates.map((filePath) => { + let mtimeMs = 0 + try { + mtimeMs = fs.statSync(filePath).mtimeMs + } catch { + // Candidate vanished between readdir and stat: try it last. + } + return { filePath, mtimeMs } + }) + candidatesWithMtime.sort( + (a, b) => b.mtimeMs - a.mtimeMs || b.filePath.localeCompare(a.filePath), + ) + for (const { filePath: candidatePath } of candidatesWithMtime) { + const recovered = tryReadRunState(candidatePath) + if (recovered !== undefined) { + bestEffortLog( + 'warn', + { runStatePath, recoveredFrom: candidatePath }, + 'run-state.json was unreadable; restored agent context from a backup generation', + ) + // Self-heal: make the recovered generation the primary again. + try { + writeFileAtomic(runStatePath, JSON.stringify(recovered)) + } catch { + // Best-effort; the candidate file still exists for the next load. + } + return { runState: recovered, fromPrimary: false } + } + } + + bestEffortLog( + 'warn', + { + runStatePath, + primaryError: + primaryError instanceof Error + ? primaryError.message + : String(primaryError), + candidatesTried: candidatesWithMtime.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 +460,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 +634,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 +692,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 +717,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) {