Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/core/checkpoints/__tests__/checkpoint.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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),
}
Expand Down
27 changes: 25 additions & 2 deletions src/core/checkpoints/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
}
Comment on lines +230 to +238
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The API history snapshot is saved after awaiting the checkpoint result, but if checkpointResult is undefined (when saveCheckpoint returns undefined for no changes), the snapshot save will be skipped. This could lead to inconsistent state where some checkpoints have snapshots and others don't, even when the API history should be preserved. Consider saving the snapshot before returning, or ensure it's saved even when no file changes occurred but API history exists.

Fix it with Roo Code or mention @roomote and request a fix.


return checkpointResult
}

export type CheckpointRestoreOptions = {
Expand Down Expand Up @@ -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)

Expand Down
9 changes: 9 additions & 0 deletions src/core/task/Task.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1079,6 +1079,11 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
}

public async condenseContext(): Promise<void> {
// 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
Expand Down Expand Up @@ -3011,6 +3016,10 @@ export class Task extends EventEmitter<TaskEvents> 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) {
Expand Down
35 changes: 35 additions & 0 deletions src/services/checkpoints/ShadowCheckpointService.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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<void> {
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<ApiMessage[] | null> {
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`)
Expand Down
196 changes: 196 additions & 0 deletions src/services/checkpoints/__tests__/api-snapshots.spec.ts
Original file line number Diff line number Diff line change
@@ -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)
})
})
})
Loading