Skip to content

Commit 5b4dedb

Browse files
committed
fix: preserve full conversation history when restoring checkpoints before context compression
- Save API conversation history snapshots along with file checkpoints - Restore full API history when reverting to checkpoints made before compression - Automatically create checkpoint before context compression operations - Add comprehensive tests for API history snapshot functionality Fixes #9308
1 parent 3e0bd0e commit 5b4dedb

File tree

5 files changed

+267
-2
lines changed

5 files changed

+267
-2
lines changed

src/core/checkpoints/__tests__/checkpoint.test.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,8 @@ describe("Checkpoint functionality", () => {
7070
saveCheckpoint: vi.fn().mockResolvedValue({ commit: "test-commit-hash" }),
7171
restoreCheckpoint: vi.fn().mockResolvedValue(undefined),
7272
getDiff: vi.fn().mockResolvedValue([]),
73+
saveApiHistorySnapshot: vi.fn().mockResolvedValue(undefined),
74+
restoreApiHistorySnapshot: vi.fn().mockResolvedValue(null), // Return null to test fallback behavior
7375
on: vi.fn(),
7476
initShadowGit: vi.fn().mockResolvedValue(undefined),
7577
}

src/core/checkpoints/index.ts

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -219,12 +219,25 @@ export async function checkpointSave(task: Task, force = false, suppressMessage
219219
TelemetryService.instance.captureCheckpointCreated(task.taskId)
220220

221221
// Start the checkpoint process in the background.
222-
return service
222+
const checkpointResult = await service
223223
.saveCheckpoint(`Task: ${task.taskId}, Time: ${Date.now()}`, { allowEmpty: force, suppressMessage })
224224
.catch((err) => {
225225
console.error("[Task#checkpointSave] caught unexpected error, disabling checkpoints", err)
226226
task.enableCheckpoints = false
227+
return undefined
227228
})
229+
230+
// Save API conversation history snapshot along with checkpoint
231+
if (checkpointResult && checkpointResult.commit) {
232+
try {
233+
await service.saveApiHistorySnapshot(checkpointResult.commit, task.apiConversationHistory)
234+
} catch (err) {
235+
console.error("[Task#checkpointSave] Failed to save API history snapshot:", err)
236+
// Don't disable checkpoints for this error, continue with file checkpoint only
237+
}
238+
}
239+
240+
return checkpointResult
228241
}
229242

230243
export type CheckpointRestoreOptions = {
@@ -258,7 +271,17 @@ export async function checkpointRestore(
258271
await provider?.postMessageToWebview({ type: "currentCheckpointUpdated", text: commitHash })
259272

260273
if (mode === "restore") {
261-
await task.overwriteApiConversationHistory(task.apiConversationHistory.filter((m) => !m.ts || m.ts < ts))
274+
// Try to restore the API conversation history snapshot
275+
const restoredApiHistory = await service.restoreApiHistorySnapshot(commitHash)
276+
if (restoredApiHistory) {
277+
// Use the restored snapshot instead of filtering by timestamp
278+
await task.overwriteApiConversationHistory(restoredApiHistory)
279+
} else {
280+
// Fallback to the old behavior if no snapshot exists
281+
await task.overwriteApiConversationHistory(
282+
task.apiConversationHistory.filter((m) => !m.ts || m.ts < ts),
283+
)
284+
}
262285

263286
const deletedMessages = task.clineMessages.slice(index + 1)
264287

src/core/task/Task.ts

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1079,6 +1079,11 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
10791079
}
10801080

10811081
public async condenseContext(): Promise<void> {
1082+
// Create a checkpoint before context compression to preserve the full conversation history
1083+
if (this.enableCheckpoints) {
1084+
await this.checkpointSave(true, true) // force=true to ensure checkpoint even if no file changes, suppressMessage=true to avoid cluttering chat
1085+
}
1086+
10821087
const systemPrompt = await this.getSystemPrompt()
10831088

10841089
// Get condensing configuration
@@ -3011,6 +3016,10 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
30113016
currentProfileId,
30123017
})
30133018
if (truncateResult.messages !== this.apiConversationHistory) {
3019+
// Create a checkpoint before overwriting with compressed history
3020+
if (truncateResult.summary && this.enableCheckpoints) {
3021+
await this.checkpointSave(true, true) // force=true, suppressMessage=true
3022+
}
30143023
await this.overwriteApiConversationHistory(truncateResult.messages)
30153024
}
30163025
if (truncateResult.error) {

src/services/checkpoints/ShadowCheckpointService.ts

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ import * as vscode from "vscode"
1111
import { fileExistsAtPath } from "../../utils/fs"
1212
import { executeRipgrep } from "../../services/search/file-search"
1313
import { t } from "../../i18n"
14+
import { safeWriteJson } from "../../utils/safeWriteJson"
15+
import { type ApiMessage } from "../../core/task-persistence"
1416

1517
import { CheckpointDiff, CheckpointResult, CheckpointEventMap } from "./types"
1618
import { getExcludePatterns } from "./excludes"
@@ -333,6 +335,39 @@ export abstract class ShadowCheckpointService extends EventEmitter {
333335
}
334336
}
335337

338+
/**
339+
* Save API conversation history snapshot for a checkpoint
340+
* @param commitHash The checkpoint commit hash to associate with this snapshot
341+
* @param apiHistory The API conversation history to save
342+
*/
343+
public async saveApiHistorySnapshot(commitHash: string, apiHistory: ApiMessage[]): Promise<void> {
344+
const snapshotPath = path.join(this.checkpointsDir, "api_snapshots", `${commitHash}.json`)
345+
await safeWriteJson(snapshotPath, apiHistory)
346+
}
347+
348+
/**
349+
* Restore API conversation history snapshot for a checkpoint
350+
* @param commitHash The checkpoint commit hash to restore from
351+
* @returns The restored API conversation history, or null if not found
352+
*/
353+
public async restoreApiHistorySnapshot(commitHash: string): Promise<ApiMessage[] | null> {
354+
const snapshotPath = path.join(this.checkpointsDir, "api_snapshots", `${commitHash}.json`)
355+
356+
if (await fileExistsAtPath(snapshotPath)) {
357+
try {
358+
const content = await fs.readFile(snapshotPath, "utf8")
359+
return JSON.parse(content) as ApiMessage[]
360+
} catch (error) {
361+
this.log(
362+
`[${this.constructor.name}#restoreApiHistorySnapshot] Failed to restore snapshot: ${error.message}`,
363+
)
364+
return null
365+
}
366+
}
367+
368+
return null
369+
}
370+
336371
public async restoreCheckpoint(commitHash: string) {
337372
try {
338373
this.log(`[${this.constructor.name}#restoreCheckpoint] starting checkpoint restore`)
Lines changed: 196 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,196 @@
1+
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"
2+
import * as fs from "fs/promises"
3+
import * as path from "path"
4+
import { ShadowCheckpointService } from "../ShadowCheckpointService"
5+
import { fileExistsAtPath } from "../../../utils/fs"
6+
import { safeWriteJson } from "../../../utils/safeWriteJson"
7+
import type { ApiMessage } from "../../../core/task-persistence"
8+
9+
vi.mock("fs/promises")
10+
vi.mock("../../../utils/fs")
11+
vi.mock("../../../utils/safeWriteJson")
12+
13+
describe("API History Snapshots", () => {
14+
let service: ShadowCheckpointService
15+
const mockTaskId = "test-task-123"
16+
const mockCheckpointsDir = "/mock/checkpoints"
17+
const mockWorkspaceDir = "/mock/workspace"
18+
const mockLog = vi.fn()
19+
20+
beforeEach(() => {
21+
vi.clearAllMocks()
22+
// Create a concrete subclass for testing
23+
class TestCheckpointService extends ShadowCheckpointService {
24+
constructor() {
25+
super(mockTaskId, mockCheckpointsDir, mockWorkspaceDir, mockLog)
26+
}
27+
}
28+
service = new TestCheckpointService()
29+
})
30+
31+
describe("saveApiHistorySnapshot", () => {
32+
it("should save API history snapshot to the correct path", async () => {
33+
const commitHash = "abc123"
34+
const apiHistory: ApiMessage[] = [
35+
{
36+
role: "user",
37+
content: "Test message 1",
38+
ts: 1000,
39+
},
40+
{
41+
role: "assistant",
42+
content: "Test response 1",
43+
ts: 2000,
44+
},
45+
]
46+
47+
await service.saveApiHistorySnapshot(commitHash, apiHistory)
48+
49+
const expectedPath = path.join(mockCheckpointsDir, "api_snapshots", `${commitHash}.json`)
50+
expect(safeWriteJson).toHaveBeenCalledWith(expectedPath, apiHistory)
51+
})
52+
53+
it("should handle empty API history", async () => {
54+
const commitHash = "def456"
55+
const apiHistory: ApiMessage[] = []
56+
57+
await service.saveApiHistorySnapshot(commitHash, apiHistory)
58+
59+
const expectedPath = path.join(mockCheckpointsDir, "api_snapshots", `${commitHash}.json`)
60+
expect(safeWriteJson).toHaveBeenCalledWith(expectedPath, apiHistory)
61+
})
62+
63+
it("should handle API history with context compression markers", async () => {
64+
const commitHash = "ghi789"
65+
const apiHistory: ApiMessage[] = [
66+
{
67+
role: "user",
68+
content: "Initial message",
69+
ts: 1000,
70+
},
71+
{
72+
role: "assistant",
73+
content: "Response",
74+
ts: 2000,
75+
},
76+
{
77+
role: "user",
78+
content: "Compressed context summary",
79+
ts: 3000,
80+
isSummary: true,
81+
},
82+
{
83+
role: "assistant",
84+
content: "After compression",
85+
ts: 4000,
86+
},
87+
]
88+
89+
await service.saveApiHistorySnapshot(commitHash, apiHistory)
90+
91+
const expectedPath = path.join(mockCheckpointsDir, "api_snapshots", `${commitHash}.json`)
92+
expect(safeWriteJson).toHaveBeenCalledWith(expectedPath, apiHistory)
93+
})
94+
})
95+
96+
describe("restoreApiHistorySnapshot", () => {
97+
it("should restore API history snapshot when file exists", async () => {
98+
const commitHash = "abc123"
99+
const expectedHistory: ApiMessage[] = [
100+
{
101+
role: "user",
102+
content: "Test message 1",
103+
ts: 1000,
104+
},
105+
{
106+
role: "assistant",
107+
content: "Test response 1",
108+
ts: 2000,
109+
},
110+
]
111+
112+
const snapshotPath = path.join(mockCheckpointsDir, "api_snapshots", `${commitHash}.json`)
113+
vi.mocked(fileExistsAtPath).mockResolvedValue(true)
114+
vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify(expectedHistory))
115+
116+
const result = await service.restoreApiHistorySnapshot(commitHash)
117+
118+
expect(result).toEqual(expectedHistory)
119+
expect(fileExistsAtPath).toHaveBeenCalledWith(snapshotPath)
120+
expect(fs.readFile).toHaveBeenCalledWith(snapshotPath, "utf8")
121+
})
122+
123+
it("should return null when snapshot file does not exist", async () => {
124+
const commitHash = "nonexistent"
125+
const snapshotPath = path.join(mockCheckpointsDir, "api_snapshots", `${commitHash}.json`)
126+
vi.mocked(fileExistsAtPath).mockResolvedValue(false)
127+
128+
const result = await service.restoreApiHistorySnapshot(commitHash)
129+
130+
expect(result).toBeNull()
131+
expect(fileExistsAtPath).toHaveBeenCalledWith(snapshotPath)
132+
expect(fs.readFile).not.toHaveBeenCalled()
133+
})
134+
135+
it("should handle invalid JSON gracefully", async () => {
136+
const commitHash = "invalid"
137+
const snapshotPath = path.join(mockCheckpointsDir, "api_snapshots", `${commitHash}.json`)
138+
vi.mocked(fileExistsAtPath).mockResolvedValue(true)
139+
vi.mocked(fs.readFile).mockResolvedValue("{ invalid json")
140+
141+
const result = await service.restoreApiHistorySnapshot(commitHash)
142+
143+
expect(result).toBeNull()
144+
expect(mockLog).toHaveBeenCalledWith(expect.stringContaining("Failed to restore snapshot"))
145+
})
146+
147+
it("should handle file read errors gracefully", async () => {
148+
const commitHash = "error"
149+
const snapshotPath = path.join(mockCheckpointsDir, "api_snapshots", `${commitHash}.json`)
150+
vi.mocked(fileExistsAtPath).mockResolvedValue(true)
151+
vi.mocked(fs.readFile).mockRejectedValue(new Error("File read error"))
152+
153+
const result = await service.restoreApiHistorySnapshot(commitHash)
154+
155+
expect(result).toBeNull()
156+
expect(mockLog).toHaveBeenCalledWith(expect.stringContaining("Failed to restore snapshot"))
157+
})
158+
159+
it("should restore API history with preserved context compression state", async () => {
160+
const commitHash = "with-compression"
161+
const expectedHistory: ApiMessage[] = [
162+
{
163+
role: "user",
164+
content: "Original message 1",
165+
ts: 1000,
166+
},
167+
{
168+
role: "assistant",
169+
content: "Original response 1",
170+
ts: 2000,
171+
},
172+
{
173+
role: "user",
174+
content: "Original message 2",
175+
ts: 3000,
176+
},
177+
{
178+
role: "assistant",
179+
content: "Original response 2",
180+
ts: 4000,
181+
},
182+
// No compression markers - full history preserved
183+
]
184+
185+
const snapshotPath = path.join(mockCheckpointsDir, "api_snapshots", `${commitHash}.json`)
186+
vi.mocked(fileExistsAtPath).mockResolvedValue(true)
187+
vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify(expectedHistory))
188+
189+
const result = await service.restoreApiHistorySnapshot(commitHash)
190+
191+
expect(result).toEqual(expectedHistory)
192+
expect(result).toHaveLength(4)
193+
expect(result?.every((msg) => !msg.isSummary)).toBe(true)
194+
})
195+
})
196+
})

0 commit comments

Comments
 (0)