|
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' |
2 | 10 | import * as fs from 'fs' |
3 | 11 | import * as path from 'path' |
4 | 12 | import * as os from 'os' |
@@ -917,3 +925,174 @@ describe('poisoned payload persistence', () => { |
917 | 925 | expect(block.outputRaw.self).toBe('[Circular]') |
918 | 926 | }) |
919 | 927 | }) |
| 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