Skip to content

Commit e9d48bc

Browse files
committed
fix(cli): recover agent context when run-state.json is torn
A resumed chat could start amnesiac — the transcript intact, the model with no memory of earlier turns — through a chain with no single point of failure: - writeFileAtomic renamed the temp over the target without fsync, so a power loss could land the rename while the data blocks were never written: the 'atomic' file contained garbage. fsync the temp before renaming (both sync and async paths). - loadMostRecentChatState tried only the primary and fell back to a RunState placeholder with no sessionState. The SDK starts a fresh session when previousRun.sessionState is absent, so the next turn silently lost every earlier turn. Now the load tries the rotated .bak (the previous complete generation) and then the newest complete checkpoint temp (a SIGKILL between write and rename leaves one behind), self-heals the primary from whichever recovered, and only gives up when all three are unreadable. - The loss used to be invisible to the user: the transcript rendered normally and the model just 'forgot'. loadMostRecentChatState now reports whether agent context survived, and the resume flow prepends an error-variant notice when it did not. Also: retry the async rename briefly on EPERM/EBUSY/EACCES — on Windows the just-closed handle can still be held by AV/indexer scans, which the new fsync widens the window for. Refs #1166 (persistence-side companion to the compaction wipes)
1 parent f3e9786 commit e9d48bc

4 files changed

Lines changed: 380 additions & 25 deletions

File tree

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

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -187,7 +187,21 @@ export const useSendMessage = ({
187187
if (loadedState) {
188188
previousRunStateRef.current = loadedState.runState
189189
setRunState(loadedState.runState)
190-
setMessages(sanitizeRestoredMessages(loadedState.messages))
190+
const restoredMessages = sanitizeRestoredMessages(loadedState.messages)
191+
if (loadedState.runStateRestored) {
192+
setMessages(restoredMessages)
193+
} else {
194+
// The agent's context was lost (torn run-state.json, nothing
195+
// recoverable) while the transcript survived. Surface it: without
196+
// this the model just answers as if the earlier turns never
197+
// happened, which reads as the assistant being broken.
198+
setMessages([
199+
createErrorChatMessage(
200+
'The saved agent context could not be restored, so the assistant starts this chat without memory of earlier turns. The transcript below is intact.',
201+
),
202+
...restoredMessages,
203+
])
204+
}
191205
if (loadedState.chatId) {
192206
setCurrentChatId(loadedState.chatId)
193207
}

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

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -917,3 +917,146 @@ describe('poisoned payload persistence', () => {
917917
expect(block.outputRaw.self).toBe('[Circular]')
918918
})
919919
})
920+
921+
describe('run state recovery', () => {
922+
// Point persistence at a temp dir via the explicit test override.
923+
const chatDir = path.join(TEST_ROOT, 'codebuff-test-recovery')
924+
925+
const runStateWithSession = (marker: string): RunState =>
926+
({
927+
sessionState: {
928+
mainAgentState: {
929+
messageHistory: [{ role: 'user', content: marker }],
930+
},
931+
},
932+
output: { type: 'lastMessage', value: marker },
933+
traceSessionId: 'trace-1',
934+
}) as unknown as RunState
935+
936+
const runStatePath = path.join(chatDir, 'run-state.json')
937+
const bakPath = runStatePath + '.bak'
938+
const messagesPath = path.join(chatDir, 'chat-messages.json')
939+
940+
const writePrimary = (contents: string) =>
941+
fs.writeFileSync(runStatePath, contents)
942+
const validMessages = JSON.stringify([
943+
{
944+
id: 'msg-1',
945+
variant: 'user',
946+
content: 'the prompt',
947+
timestamp: new Date().toISOString(),
948+
},
949+
] as ChatMessage[])
950+
951+
beforeEach(() => {
952+
fs.rmSync(chatDir, { recursive: true, force: true })
953+
fs.mkdirSync(chatDir, { recursive: true })
954+
setChatDirOverrideForTesting(chatDir)
955+
})
956+
957+
afterEach(() => {
958+
setChatDirOverrideForTesting(undefined)
959+
})
960+
961+
test('recovers agent context from the .bak when the primary is torn', () => {
962+
writePrimary('{"sessionState": {"main"') // torn: power loss after rename
963+
fs.writeFileSync(bakPath, JSON.stringify(runStateWithSession('from-bak')))
964+
fs.writeFileSync(messagesPath, validMessages)
965+
966+
const loaded = loadMostRecentChatState()
967+
expect(loaded).not.toBeNull()
968+
// Agent context survived — the model is NOT amnesiac next turn.
969+
expect((loaded!.runState as any).sessionState).toBeDefined()
970+
expect(loaded!.runStateRestored).toBe(true)
971+
// Self-healed: the primary is the recovered generation again.
972+
expect(
973+
JSON.parse(fs.readFileSync(runStatePath, 'utf8')).output.value,
974+
).toBe('from-bak')
975+
})
976+
977+
test('recovers from the newest complete checkpoint temp when bak is absent', () => {
978+
writePrimary('{ torn')
979+
fs.writeFileSync(messagesPath, validMessages)
980+
// Two temps: an older torn one and a newer complete one (SIGKILL between
981+
// write and rename leaves the latter behind).
982+
fs.writeFileSync(
983+
runStatePath + '.999.oldest.tmp',
984+
'{"half":',
985+
)
986+
fs.writeFileSync(
987+
runStatePath + '.1234.newest.tmp',
988+
JSON.stringify(runStateWithSession('from-tmp')),
989+
)
990+
991+
const loaded = loadMostRecentChatState()
992+
expect(loaded).not.toBeNull()
993+
expect((loaded!.runState as any).sessionState).toBeDefined()
994+
expect(loaded!.runStateRestored).toBe(true)
995+
// Self-healed into the primary.
996+
expect(
997+
JSON.parse(fs.readFileSync(runStatePath, 'utf8')).output.value,
998+
).toBe('from-tmp')
999+
})
1000+
1001+
test('flags a healthy primary as fully restored', () => {
1002+
writePrimary(JSON.stringify(runStateWithSession('healthy')))
1003+
fs.writeFileSync(messagesPath, validMessages)
1004+
1005+
const loaded = loadMostRecentChatState()
1006+
expect(loaded!.runStateRestored).toBe(true)
1007+
expect((loaded!.runState as any).sessionState).toBeDefined()
1008+
})
1009+
1010+
test('falls back to a context-less placeholder with the loss flagged when nothing recovers', () => {
1011+
writePrimary('{ torn')
1012+
fs.writeFileSync(messagesPath, validMessages)
1013+
1014+
const loaded = loadMostRecentChatState()
1015+
expect(loaded).not.toBeNull()
1016+
// The amnesia carrier: no sessionState — the SDK will start a fresh
1017+
// session next turn. runStateRestored=false is what makes the UI say so
1018+
// instead of the model silently forgetting every earlier turn.
1019+
expect((loaded!.runState as any).sessionState).toBeUndefined()
1020+
expect(loaded!.runStateRestored).toBe(false)
1021+
// The transcript still survives.
1022+
expect(loaded!.messages.length).toBe(1)
1023+
})
1024+
1025+
test('saveChatState rotates the previous primary into .bak', () => {
1026+
saveChatState(runStateWithSession('generation-1'), [
1027+
{
1028+
id: 'msg-1',
1029+
variant: 'user',
1030+
content: 'first',
1031+
timestamp: new Date().toISOString(),
1032+
},
1033+
])
1034+
expect(fs.existsSync(bakPath)).toBe(false)
1035+
1036+
saveChatState(runStateWithSession('generation-2'), [
1037+
{
1038+
id: 'msg-2',
1039+
variant: 'user',
1040+
content: 'second',
1041+
timestamp: new Date().toISOString(),
1042+
},
1043+
])
1044+
1045+
expect(JSON.parse(fs.readFileSync(runStatePath, 'utf8')).output.value).toBe(
1046+
'generation-2',
1047+
)
1048+
expect(JSON.parse(fs.readFileSync(bakPath, 'utf8')).output.value).toBe(
1049+
'generation-1',
1050+
)
1051+
})
1052+
1053+
test('clearChatState removes the backup too', () => {
1054+
writePrimary(JSON.stringify(runStateWithSession('x')))
1055+
fs.writeFileSync(bakPath, JSON.stringify(runStateWithSession('bak')))
1056+
fs.writeFileSync(messagesPath, validMessages)
1057+
1058+
clearChatState()
1059+
expect(fs.existsSync(runStatePath)).toBe(false)
1060+
expect(fs.existsSync(bakPath)).toBe(false)
1061+
})
1062+
})

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

Lines changed: 150 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,11 @@ type SavedChatState = {
2525
runState: RunState
2626
messages: ChatMessage[]
2727
chatId?: string
28+
/** False only when run-state.json was unreadable AND nothing could be
29+
* recovered from the backup or checkpoint temps, so the restored RunState
30+
* is a placeholder without agent context — the model starts amnesiac next
31+
* turn. The UI surfaces this so the loss is not silent. */
32+
runStateRestored: boolean
2833
}
2934

3035
type LiveChatState = {
@@ -244,6 +249,121 @@ type SerializedChatState = {
244249
messagesJson?: string
245250
}
246251

252+
/** The previous generation of run-state.json, rotated aside by saveChatState.
253+
* loadMostRecentChatState recovers from it when the primary is torn. */
254+
const RUN_STATE_BACKUP_SUFFIX = '.bak'
255+
256+
function tryReadRunState(filePath: string): RunState | undefined {
257+
try {
258+
return JSON.parse(fs.readFileSync(filePath, 'utf8')) as RunState
259+
} catch {
260+
return undefined
261+
}
262+
}
263+
264+
/**
265+
* Read run-state.json with best-effort recovery when it is torn.
266+
*
267+
* The atomic write makes the rename indivisible, but without an fsync a power
268+
* loss can still land the rename while the file's data blocks were never
269+
* written — leaving a truncated or empty primary (and even with the fsync,
270+
* external corruption exists). Two older complete generations can be lying in
271+
* the chat directory when that happens:
272+
*
273+
* - `run-state.json.bak` — the state the primary replaced at the last
274+
* synchronous save (saveChatState rotates the previous file aside).
275+
* - `run-state.json.<pid>.<uuid>.tmp` — a write that was killed in the
276+
* window between its write and its rename (writeFileAtomic unlinks its own
277+
* temp on error, but a SIGKILL cannot).
278+
*
279+
* Try the primary, then the backup, then the newest readable temp, and
280+
* self-heal the primary from whichever recovered. `fromPrimary` is false when
281+
* recovery had to fall back, so the caller can warn that some agent context
282+
* is missing instead of losing it silently.
283+
*/
284+
function readRunStateWithRecovery(chatDir: string): {
285+
runState: RunState
286+
fromPrimary: boolean
287+
} | null {
288+
const runStatePath = path.join(chatDir, RUN_STATE_FILENAME)
289+
let primaryError: unknown
290+
try {
291+
return {
292+
runState: JSON.parse(fs.readFileSync(runStatePath, 'utf8')) as RunState,
293+
fromPrimary: true,
294+
}
295+
} catch (error) {
296+
primaryError = error
297+
}
298+
299+
const bakPath = runStatePath + RUN_STATE_BACKUP_SUFFIX
300+
const fromBackup = tryReadRunState(bakPath)
301+
if (fromBackup !== undefined) {
302+
bestEffortLog(
303+
'warn',
304+
{ runStatePath, bakPath },
305+
'run-state.json was unreadable; restored agent context from the previous save',
306+
)
307+
// Self-heal: make the recovered generation the primary again.
308+
try {
309+
writeFileAtomic(
310+
runStatePath,
311+
JSON.stringify(fromBackup),
312+
)
313+
} catch {
314+
// Best-effort; the backup file still exists for the next load.
315+
}
316+
return { runState: fromBackup, fromPrimary: false }
317+
}
318+
319+
// Newest temp first — each is a complete write that never got renamed.
320+
const tmpPaths = fs
321+
.readdirSync(chatDir)
322+
.filter(
323+
(file) =>
324+
file.startsWith(RUN_STATE_FILENAME + '.') && file.endsWith('.tmp'),
325+
)
326+
.map((file) => path.join(chatDir, file))
327+
.sort((a, b) => {
328+
try {
329+
return fs.statSync(b).mtimeMs - fs.statSync(a).mtimeMs
330+
} catch {
331+
return 0
332+
}
333+
})
334+
for (const tmpPath of tmpPaths) {
335+
const recovered = tryReadRunState(tmpPath)
336+
if (recovered !== undefined) {
337+
bestEffortLog(
338+
'warn',
339+
{ runStatePath, tmpPath },
340+
'run-state.json was unreadable; restored agent context from an unfinished checkpoint write',
341+
)
342+
try {
343+
writeFileAtomic(runStatePath, JSON.stringify(recovered))
344+
} catch {
345+
// Best-effort; the temp file still exists for the next load.
346+
}
347+
return { runState: recovered, fromPrimary: false }
348+
}
349+
}
350+
351+
bestEffortLog(
352+
'warn',
353+
{
354+
runStatePath,
355+
primaryError:
356+
primaryError instanceof Error
357+
? primaryError.message
358+
: String(primaryError),
359+
backupTried: bakPath,
360+
tempsTried: tmpPaths.length,
361+
},
362+
'Could not read run state; restoring transcript without agent context',
363+
)
364+
return null
365+
}
366+
247367
/**
248368
* Serialize the two chat-state files independently, so a poisoned run state
249369
* cannot block persisting the transcript (and vice versa). Cyclic and
@@ -341,10 +461,18 @@ export function saveChatState(
341461
// since (e.g. the chat deleted from /history mid-run).
342462
fs.mkdirSync(chatDir, { recursive: true })
343463
if (serialized.runStateJson) {
344-
writeFileAtomic(
345-
path.join(chatDir, RUN_STATE_FILENAME),
346-
serialized.runStateJson,
347-
)
464+
const runStatePath = path.join(chatDir, RUN_STATE_FILENAME)
465+
// Rotate the previous generation aside before overwriting: it is the
466+
// newest complete state readRunStateWithRecovery can fall back to if
467+
// this write's rename lands but its data is later found torn.
468+
try {
469+
if (fs.existsSync(runStatePath)) {
470+
fs.renameSync(runStatePath, runStatePath + RUN_STATE_BACKUP_SUFFIX)
471+
}
472+
} catch {
473+
// Rotation is best-effort; the overwrite below still proceeds.
474+
}
475+
writeFileAtomic(runStatePath, serialized.runStateJson)
348476
}
349477
if (serialized.messagesJson) {
350478
writeFileAtomic(
@@ -507,14 +635,16 @@ export function loadMostRecentChatState(
507635
// must not lose the transcript, and vice versa. Restore whatever is
508636
// readable and fall back for the rest.
509637
let runState: RunState | null = null
510-
try {
511-
runState = JSON.parse(fs.readFileSync(runStatePath, 'utf8')) as RunState
512-
} catch (error) {
638+
let runStateRestored = false
639+
const recovered = readRunStateWithRecovery(chatDir)
640+
if (recovered) {
641+
// Agent context is present — whether straight from the primary or
642+
// recovered from a backup/temp (both carry a real sessionState).
643+
runState = recovered.runState
644+
runStateRestored = true
645+
} else {
513646
logger.warn(
514-
{
515-
runStatePath,
516-
error: error instanceof Error ? error.message : String(error),
517-
},
647+
{ runStatePath },
518648
'Could not read run state; restoring transcript without agent context',
519649
)
520650
}
@@ -563,7 +693,12 @@ export function loadMostRecentChatState(
563693
'Loaded chat state from chat directory',
564694
)
565695

566-
return { runState, messages, chatId: resolvedChatId }
696+
return {
697+
runState,
698+
messages,
699+
chatId: resolvedChatId,
700+
runStateRestored,
701+
}
567702
} catch (error) {
568703
logger.error(
569704
{
@@ -583,13 +718,14 @@ export function clearChatState(): void {
583718
const runStatePath = getRunStatePath()
584719
const messagesPath = getChatMessagesPath()
585720
const metaPath = path.join(resolveCurrentChatDir(), CHAT_META_FILENAME)
721+
const backupPath = runStatePath + RUN_STATE_BACKUP_SUFFIX
586722

587-
for (const filePath of [runStatePath, messagesPath, metaPath]) {
723+
for (const filePath of [runStatePath, backupPath, messagesPath, metaPath]) {
588724
fs.rmSync(filePath, { force: true })
589725
}
590726

591727
logger.debug(
592-
{ runStatePath, messagesPath, metaPath },
728+
{ runStatePath, backupPath, messagesPath, metaPath },
593729
'Cleared chat state files',
594730
)
595731
} catch (error) {

0 commit comments

Comments
 (0)