Skip to content
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 @@ -2,6 +2,7 @@ import { describe, it, expect, vi, beforeEach, afterEach, Mock } from "vitest"
import { Task } from "../../task/Task"
import { ClineProvider } from "../../webview/ClineProvider"
import { checkpointSave, checkpointRestore, checkpointDiff, getCheckpointService } from "../index"
import { MessageManager } from "../../message-manager"
import * as vscode from "vscode"

// Mock vscode
Expand Down Expand Up @@ -102,6 +103,7 @@ describe("Checkpoint functionality", () => {
overwriteApiConversationHistory: vi.fn(),
combineMessages: vi.fn().mockReturnValue([]),
}
mockTask.messageManager = new MessageManager(mockTask)

// Update the mock to return our mockCheckpointService
const checkpointsModule = await import("../../../services/checkpoints")
Expand Down
14 changes: 7 additions & 7 deletions src/core/checkpoints/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -258,20 +258,20 @@ 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))

// Calculate metrics from messages that will be deleted (must be done before rewind)
const deletedMessages = task.clineMessages.slice(index + 1)

const { totalTokensIn, totalTokensOut, totalCacheWrites, totalCacheReads, totalCost } = getApiMetrics(
task.combineMessages(deletedMessages),
)

// For delete operations, exclude the checkpoint message itself
// For edit operations, include the checkpoint message (to be edited)
const endIndex = operation === "edit" ? index + 1 : index
await task.overwriteClineMessages(task.clineMessages.slice(0, endIndex))
// Use MessageManager to properly handle context-management events
// This ensures orphaned Summary messages and truncation markers are cleaned up
await task.messageManager.rewindToTimestamp(ts, {
includeTargetMessage: operation === "edit",
})
Comment on lines +270 to +272
Copy link
Collaborator

Choose a reason for hiding this comment

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

Not a blocker, but I continue to think that it would be nice if we didn't need to special case edit in the way that we have been doing

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Tomorrow's problem?


// TODO: Verify that this is working as expected.
// Report the deleted API request metrics
await task.say(
"api_req_deleted",
JSON.stringify({
Expand Down
Loading
Loading