From 5b4dedbc8ab7c866e5bcb64eee9e5ef7103c79a7 Mon Sep 17 00:00:00 2001 From: Roo Code Date: Mon, 17 Nov 2025 16:14:02 +0000 Subject: [PATCH] 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 --- .../checkpoints/__tests__/checkpoint.test.ts | 2 + src/core/checkpoints/index.ts | 27 ++- src/core/task/Task.ts | 9 + .../checkpoints/ShadowCheckpointService.ts | 35 ++++ .../__tests__/api-snapshots.spec.ts | 196 ++++++++++++++++++ 5 files changed, 267 insertions(+), 2 deletions(-) create mode 100644 src/services/checkpoints/__tests__/api-snapshots.spec.ts diff --git a/src/core/checkpoints/__tests__/checkpoint.test.ts b/src/core/checkpoints/__tests__/checkpoint.test.ts index e9e565853f..4ed946b7fe 100644 --- a/src/core/checkpoints/__tests__/checkpoint.test.ts +++ b/src/core/checkpoints/__tests__/checkpoint.test.ts @@ -70,6 +70,8 @@ describe("Checkpoint functionality", () => { saveCheckpoint: vi.fn().mockResolvedValue({ commit: "test-commit-hash" }), restoreCheckpoint: vi.fn().mockResolvedValue(undefined), getDiff: vi.fn().mockResolvedValue([]), + saveApiHistorySnapshot: vi.fn().mockResolvedValue(undefined), + restoreApiHistorySnapshot: vi.fn().mockResolvedValue(null), // Return null to test fallback behavior on: vi.fn(), initShadowGit: vi.fn().mockResolvedValue(undefined), } diff --git a/src/core/checkpoints/index.ts b/src/core/checkpoints/index.ts index 3efdb466e6..f8f20f6108 100644 --- a/src/core/checkpoints/index.ts +++ b/src/core/checkpoints/index.ts @@ -219,12 +219,25 @@ export async function checkpointSave(task: Task, force = false, suppressMessage TelemetryService.instance.captureCheckpointCreated(task.taskId) // Start the checkpoint process in the background. - return service + const checkpointResult = await service .saveCheckpoint(`Task: ${task.taskId}, Time: ${Date.now()}`, { allowEmpty: force, suppressMessage }) .catch((err) => { console.error("[Task#checkpointSave] caught unexpected error, disabling checkpoints", err) task.enableCheckpoints = false + return undefined }) + + // Save API conversation history snapshot along with checkpoint + if (checkpointResult && checkpointResult.commit) { + try { + await service.saveApiHistorySnapshot(checkpointResult.commit, task.apiConversationHistory) + } catch (err) { + console.error("[Task#checkpointSave] Failed to save API history snapshot:", err) + // Don't disable checkpoints for this error, continue with file checkpoint only + } + } + + return checkpointResult } export type CheckpointRestoreOptions = { @@ -258,7 +271,17 @@ export async function checkpointRestore( await provider?.postMessageToWebview({ type: "currentCheckpointUpdated", text: commitHash }) if (mode === "restore") { - await task.overwriteApiConversationHistory(task.apiConversationHistory.filter((m) => !m.ts || m.ts < ts)) + // Try to restore the API conversation history snapshot + const restoredApiHistory = await service.restoreApiHistorySnapshot(commitHash) + if (restoredApiHistory) { + // Use the restored snapshot instead of filtering by timestamp + await task.overwriteApiConversationHistory(restoredApiHistory) + } else { + // Fallback to the old behavior if no snapshot exists + await task.overwriteApiConversationHistory( + task.apiConversationHistory.filter((m) => !m.ts || m.ts < ts), + ) + } const deletedMessages = task.clineMessages.slice(index + 1) diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index dbb9529c2e..d0d5a28424 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -1079,6 +1079,11 @@ export class Task extends EventEmitter implements TaskLike { } public async condenseContext(): Promise { + // Create a checkpoint before context compression to preserve the full conversation history + if (this.enableCheckpoints) { + await this.checkpointSave(true, true) // force=true to ensure checkpoint even if no file changes, suppressMessage=true to avoid cluttering chat + } + const systemPrompt = await this.getSystemPrompt() // Get condensing configuration @@ -3011,6 +3016,10 @@ export class Task extends EventEmitter implements TaskLike { currentProfileId, }) if (truncateResult.messages !== this.apiConversationHistory) { + // Create a checkpoint before overwriting with compressed history + if (truncateResult.summary && this.enableCheckpoints) { + await this.checkpointSave(true, true) // force=true, suppressMessage=true + } await this.overwriteApiConversationHistory(truncateResult.messages) } if (truncateResult.error) { diff --git a/src/services/checkpoints/ShadowCheckpointService.ts b/src/services/checkpoints/ShadowCheckpointService.ts index fee08b2fa4..b24cb81473 100644 --- a/src/services/checkpoints/ShadowCheckpointService.ts +++ b/src/services/checkpoints/ShadowCheckpointService.ts @@ -11,6 +11,8 @@ import * as vscode from "vscode" import { fileExistsAtPath } from "../../utils/fs" import { executeRipgrep } from "../../services/search/file-search" import { t } from "../../i18n" +import { safeWriteJson } from "../../utils/safeWriteJson" +import { type ApiMessage } from "../../core/task-persistence" import { CheckpointDiff, CheckpointResult, CheckpointEventMap } from "./types" import { getExcludePatterns } from "./excludes" @@ -333,6 +335,39 @@ export abstract class ShadowCheckpointService extends EventEmitter { } } + /** + * Save API conversation history snapshot for a checkpoint + * @param commitHash The checkpoint commit hash to associate with this snapshot + * @param apiHistory The API conversation history to save + */ + public async saveApiHistorySnapshot(commitHash: string, apiHistory: ApiMessage[]): Promise { + const snapshotPath = path.join(this.checkpointsDir, "api_snapshots", `${commitHash}.json`) + await safeWriteJson(snapshotPath, apiHistory) + } + + /** + * Restore API conversation history snapshot for a checkpoint + * @param commitHash The checkpoint commit hash to restore from + * @returns The restored API conversation history, or null if not found + */ + public async restoreApiHistorySnapshot(commitHash: string): Promise { + const snapshotPath = path.join(this.checkpointsDir, "api_snapshots", `${commitHash}.json`) + + if (await fileExistsAtPath(snapshotPath)) { + try { + const content = await fs.readFile(snapshotPath, "utf8") + return JSON.parse(content) as ApiMessage[] + } catch (error) { + this.log( + `[${this.constructor.name}#restoreApiHistorySnapshot] Failed to restore snapshot: ${error.message}`, + ) + return null + } + } + + return null + } + public async restoreCheckpoint(commitHash: string) { try { this.log(`[${this.constructor.name}#restoreCheckpoint] starting checkpoint restore`) diff --git a/src/services/checkpoints/__tests__/api-snapshots.spec.ts b/src/services/checkpoints/__tests__/api-snapshots.spec.ts new file mode 100644 index 0000000000..d9b96b79fb --- /dev/null +++ b/src/services/checkpoints/__tests__/api-snapshots.spec.ts @@ -0,0 +1,196 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest" +import * as fs from "fs/promises" +import * as path from "path" +import { ShadowCheckpointService } from "../ShadowCheckpointService" +import { fileExistsAtPath } from "../../../utils/fs" +import { safeWriteJson } from "../../../utils/safeWriteJson" +import type { ApiMessage } from "../../../core/task-persistence" + +vi.mock("fs/promises") +vi.mock("../../../utils/fs") +vi.mock("../../../utils/safeWriteJson") + +describe("API History Snapshots", () => { + let service: ShadowCheckpointService + const mockTaskId = "test-task-123" + const mockCheckpointsDir = "/mock/checkpoints" + const mockWorkspaceDir = "/mock/workspace" + const mockLog = vi.fn() + + beforeEach(() => { + vi.clearAllMocks() + // Create a concrete subclass for testing + class TestCheckpointService extends ShadowCheckpointService { + constructor() { + super(mockTaskId, mockCheckpointsDir, mockWorkspaceDir, mockLog) + } + } + service = new TestCheckpointService() + }) + + describe("saveApiHistorySnapshot", () => { + it("should save API history snapshot to the correct path", async () => { + const commitHash = "abc123" + const apiHistory: ApiMessage[] = [ + { + role: "user", + content: "Test message 1", + ts: 1000, + }, + { + role: "assistant", + content: "Test response 1", + ts: 2000, + }, + ] + + await service.saveApiHistorySnapshot(commitHash, apiHistory) + + const expectedPath = path.join(mockCheckpointsDir, "api_snapshots", `${commitHash}.json`) + expect(safeWriteJson).toHaveBeenCalledWith(expectedPath, apiHistory) + }) + + it("should handle empty API history", async () => { + const commitHash = "def456" + const apiHistory: ApiMessage[] = [] + + await service.saveApiHistorySnapshot(commitHash, apiHistory) + + const expectedPath = path.join(mockCheckpointsDir, "api_snapshots", `${commitHash}.json`) + expect(safeWriteJson).toHaveBeenCalledWith(expectedPath, apiHistory) + }) + + it("should handle API history with context compression markers", async () => { + const commitHash = "ghi789" + const apiHistory: ApiMessage[] = [ + { + role: "user", + content: "Initial message", + ts: 1000, + }, + { + role: "assistant", + content: "Response", + ts: 2000, + }, + { + role: "user", + content: "Compressed context summary", + ts: 3000, + isSummary: true, + }, + { + role: "assistant", + content: "After compression", + ts: 4000, + }, + ] + + await service.saveApiHistorySnapshot(commitHash, apiHistory) + + const expectedPath = path.join(mockCheckpointsDir, "api_snapshots", `${commitHash}.json`) + expect(safeWriteJson).toHaveBeenCalledWith(expectedPath, apiHistory) + }) + }) + + describe("restoreApiHistorySnapshot", () => { + it("should restore API history snapshot when file exists", async () => { + const commitHash = "abc123" + const expectedHistory: ApiMessage[] = [ + { + role: "user", + content: "Test message 1", + ts: 1000, + }, + { + role: "assistant", + content: "Test response 1", + ts: 2000, + }, + ] + + const snapshotPath = path.join(mockCheckpointsDir, "api_snapshots", `${commitHash}.json`) + vi.mocked(fileExistsAtPath).mockResolvedValue(true) + vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify(expectedHistory)) + + const result = await service.restoreApiHistorySnapshot(commitHash) + + expect(result).toEqual(expectedHistory) + expect(fileExistsAtPath).toHaveBeenCalledWith(snapshotPath) + expect(fs.readFile).toHaveBeenCalledWith(snapshotPath, "utf8") + }) + + it("should return null when snapshot file does not exist", async () => { + const commitHash = "nonexistent" + const snapshotPath = path.join(mockCheckpointsDir, "api_snapshots", `${commitHash}.json`) + vi.mocked(fileExistsAtPath).mockResolvedValue(false) + + const result = await service.restoreApiHistorySnapshot(commitHash) + + expect(result).toBeNull() + expect(fileExistsAtPath).toHaveBeenCalledWith(snapshotPath) + expect(fs.readFile).not.toHaveBeenCalled() + }) + + it("should handle invalid JSON gracefully", async () => { + const commitHash = "invalid" + const snapshotPath = path.join(mockCheckpointsDir, "api_snapshots", `${commitHash}.json`) + vi.mocked(fileExistsAtPath).mockResolvedValue(true) + vi.mocked(fs.readFile).mockResolvedValue("{ invalid json") + + const result = await service.restoreApiHistorySnapshot(commitHash) + + expect(result).toBeNull() + expect(mockLog).toHaveBeenCalledWith(expect.stringContaining("Failed to restore snapshot")) + }) + + it("should handle file read errors gracefully", async () => { + const commitHash = "error" + const snapshotPath = path.join(mockCheckpointsDir, "api_snapshots", `${commitHash}.json`) + vi.mocked(fileExistsAtPath).mockResolvedValue(true) + vi.mocked(fs.readFile).mockRejectedValue(new Error("File read error")) + + const result = await service.restoreApiHistorySnapshot(commitHash) + + expect(result).toBeNull() + expect(mockLog).toHaveBeenCalledWith(expect.stringContaining("Failed to restore snapshot")) + }) + + it("should restore API history with preserved context compression state", async () => { + const commitHash = "with-compression" + const expectedHistory: ApiMessage[] = [ + { + role: "user", + content: "Original message 1", + ts: 1000, + }, + { + role: "assistant", + content: "Original response 1", + ts: 2000, + }, + { + role: "user", + content: "Original message 2", + ts: 3000, + }, + { + role: "assistant", + content: "Original response 2", + ts: 4000, + }, + // No compression markers - full history preserved + ] + + const snapshotPath = path.join(mockCheckpointsDir, "api_snapshots", `${commitHash}.json`) + vi.mocked(fileExistsAtPath).mockResolvedValue(true) + vi.mocked(fs.readFile).mockResolvedValue(JSON.stringify(expectedHistory)) + + const result = await service.restoreApiHistorySnapshot(commitHash) + + expect(result).toEqual(expectedHistory) + expect(result).toHaveLength(4) + expect(result?.every((msg) => !msg.isSummary)).toBe(true) + }) + }) +})