Skip to content

Commit 5c49a92

Browse files
committed
Recover agent context from torn run-state.json backups
When run-state.json is unreadable, loadMostRecentChatState now falls back to the previous generation (.bak rotated aside by each synchronous save) and to complete checkpoint temps left by SIGKILLed writes, ordered newest-first with a deterministic tie-break, self-heals the primary from whichever recovered, and clears the .bak on chat deletion. The async atomic write retries EPERM/EBUSY/EACCES renames briefly for Windows AV/indexer locks. The resume flow prepends an error-variant notice when nothing could be recovered, so the context loss is surfaced instead of reading as a broken assistant. Companion to #1166; details and repro in #1168. Recovery half of #1169, split per review so the fsync durability fix can land on its own.
1 parent a49b46d commit 5c49a92

3 files changed

Lines changed: 344 additions & 16 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: 180 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,12 @@
1-
import { describe, test, expect, afterAll, beforeEach, afterEach, mock } from 'bun:test'
1+
import {
2+
describe,
3+
test,
4+
expect,
5+
afterAll,
6+
beforeEach,
7+
afterEach,
8+
mock,
9+
} from 'bun:test'
210
import * as fs from 'fs'
311
import * as path from 'path'
412
import * as os from 'os'
@@ -917,3 +925,174 @@ describe('poisoned payload persistence', () => {
917925
expect(block.outputRaw.self).toBe('[Circular]')
918926
})
919927
})
928+
929+
describe('run state recovery', () => {
930+
// Point persistence at a temp dir via the explicit test override.
931+
const chatDir = path.join(TEST_ROOT, 'codebuff-test-recovery')
932+
933+
const runStateWithSession = (marker: string): RunState =>
934+
({
935+
sessionState: {
936+
mainAgentState: {
937+
messageHistory: [{ role: 'user', content: marker }],
938+
},
939+
},
940+
output: { type: 'lastMessage', value: marker },
941+
traceSessionId: 'trace-1',
942+
}) as unknown as RunState
943+
944+
const runStatePath = path.join(chatDir, 'run-state.json')
945+
const bakPath = runStatePath + '.bak'
946+
const messagesPath = path.join(chatDir, 'chat-messages.json')
947+
948+
const writePrimary = (contents: string) =>
949+
fs.writeFileSync(runStatePath, contents)
950+
const validMessages = JSON.stringify([
951+
{
952+
id: 'msg-1',
953+
variant: 'user',
954+
content: 'the prompt',
955+
timestamp: new Date().toISOString(),
956+
},
957+
] as ChatMessage[])
958+
959+
beforeEach(() => {
960+
fs.rmSync(chatDir, { recursive: true, force: true })
961+
fs.mkdirSync(chatDir, { recursive: true })
962+
setChatDirOverrideForTesting(chatDir)
963+
})
964+
965+
afterEach(() => {
966+
setChatDirOverrideForTesting(undefined)
967+
})
968+
969+
test('recovers agent context from the .bak when the primary is torn', () => {
970+
writePrimary('{"sessionState": {"main"') // torn: power loss after rename
971+
fs.writeFileSync(bakPath, JSON.stringify(runStateWithSession('from-bak')))
972+
fs.writeFileSync(messagesPath, validMessages)
973+
974+
const loaded = loadMostRecentChatState()
975+
expect(loaded).not.toBeNull()
976+
// Agent context survived — the model is NOT amnesiac next turn.
977+
expect((loaded!.runState as any).sessionState).toBeDefined()
978+
expect(loaded!.runStateRestored).toBe(true)
979+
// Self-healed: the primary is the recovered generation again.
980+
expect(JSON.parse(fs.readFileSync(runStatePath, 'utf8')).output.value).toBe(
981+
'from-bak',
982+
)
983+
})
984+
985+
test('recovers from the newest complete checkpoint temp when bak is absent', () => {
986+
writePrimary('{ torn')
987+
fs.writeFileSync(messagesPath, validMessages)
988+
// Two temps: an older torn one and a newer complete one (SIGKILL between
989+
// write and rename leaves the latter behind). mtimes are pinned because
990+
// back-to-back writes can land inside one mtime quantum on some
991+
// filesystems, which would make "newest" nondeterministic.
992+
const oldTemp = runStatePath + '.999.oldest.tmp'
993+
const newTemp = runStatePath + '.1234.newest.tmp'
994+
fs.writeFileSync(oldTemp, '{"half":')
995+
fs.writeFileSync(newTemp, JSON.stringify(runStateWithSession('from-tmp')))
996+
const now = new Date()
997+
fs.utimesSync(
998+
oldTemp,
999+
new Date(now.getTime() - 10_000),
1000+
new Date(now.getTime() - 10_000),
1001+
)
1002+
fs.utimesSync(newTemp, now, now)
1003+
1004+
const loaded = loadMostRecentChatState()
1005+
expect(loaded).not.toBeNull()
1006+
expect((loaded!.runState as any).sessionState).toBeDefined()
1007+
expect(loaded!.runStateRestored).toBe(true)
1008+
// Self-healed into the primary.
1009+
expect(JSON.parse(fs.readFileSync(runStatePath, 'utf8')).output.value).toBe(
1010+
'from-tmp',
1011+
)
1012+
})
1013+
1014+
test('prefers a newer complete checkpoint temp over the .bak', () => {
1015+
// Both fallbacks are intact generations; the newest one lost the least
1016+
// agent context, so recency — not a fixed .bak-first order — decides.
1017+
writePrimary('{ torn')
1018+
fs.writeFileSync(messagesPath, validMessages)
1019+
fs.writeFileSync(bakPath, JSON.stringify(runStateWithSession('from-bak')))
1020+
const tempPath = runStatePath + '.1234.checkpoint.tmp'
1021+
fs.writeFileSync(tempPath, JSON.stringify(runStateWithSession('from-tmp')))
1022+
const now = new Date()
1023+
fs.utimesSync(
1024+
bakPath,
1025+
new Date(now.getTime() - 20_000),
1026+
new Date(now.getTime() - 20_000),
1027+
)
1028+
fs.utimesSync(tempPath, now, now)
1029+
1030+
const loaded = loadMostRecentChatState()
1031+
expect(loaded).not.toBeNull()
1032+
expect(JSON.parse(fs.readFileSync(runStatePath, 'utf8')).output.value).toBe(
1033+
'from-tmp',
1034+
)
1035+
})
1036+
1037+
test('flags a healthy primary as fully restored', () => {
1038+
writePrimary(JSON.stringify(runStateWithSession('healthy')))
1039+
fs.writeFileSync(messagesPath, validMessages)
1040+
1041+
const loaded = loadMostRecentChatState()
1042+
expect(loaded!.runStateRestored).toBe(true)
1043+
expect((loaded!.runState as any).sessionState).toBeDefined()
1044+
})
1045+
1046+
test('falls back to a context-less placeholder with the loss flagged when nothing recovers', () => {
1047+
writePrimary('{ torn')
1048+
fs.writeFileSync(messagesPath, validMessages)
1049+
1050+
const loaded = loadMostRecentChatState()
1051+
expect(loaded).not.toBeNull()
1052+
// The amnesia carrier: no sessionState — the SDK will start a fresh
1053+
// session next turn. runStateRestored=false is what makes the UI say so
1054+
// instead of the model silently forgetting every earlier turn.
1055+
expect((loaded!.runState as any).sessionState).toBeUndefined()
1056+
expect(loaded!.runStateRestored).toBe(false)
1057+
// The transcript still survives.
1058+
expect(loaded!.messages.length).toBe(1)
1059+
})
1060+
1061+
test('saveChatState rotates the previous primary into .bak', () => {
1062+
saveChatState(runStateWithSession('generation-1'), [
1063+
{
1064+
id: 'msg-1',
1065+
variant: 'user',
1066+
content: 'first',
1067+
timestamp: new Date().toISOString(),
1068+
},
1069+
])
1070+
expect(fs.existsSync(bakPath)).toBe(false)
1071+
1072+
saveChatState(runStateWithSession('generation-2'), [
1073+
{
1074+
id: 'msg-2',
1075+
variant: 'user',
1076+
content: 'second',
1077+
timestamp: new Date().toISOString(),
1078+
},
1079+
])
1080+
1081+
expect(JSON.parse(fs.readFileSync(runStatePath, 'utf8')).output.value).toBe(
1082+
'generation-2',
1083+
)
1084+
expect(JSON.parse(fs.readFileSync(bakPath, 'utf8')).output.value).toBe(
1085+
'generation-1',
1086+
)
1087+
})
1088+
1089+
test('clearChatState removes the backup too', () => {
1090+
writePrimary(JSON.stringify(runStateWithSession('x')))
1091+
fs.writeFileSync(bakPath, JSON.stringify(runStateWithSession('bak')))
1092+
fs.writeFileSync(messagesPath, validMessages)
1093+
1094+
clearChatState()
1095+
expect(fs.existsSync(runStatePath)).toBe(false)
1096+
expect(fs.existsSync(bakPath)).toBe(false)
1097+
})
1098+
})

0 commit comments

Comments
 (0)