diff --git a/src/core/checkpoints/__tests__/checkpointAfterDelete.test.ts b/src/core/checkpoints/__tests__/checkpointAfterDelete.test.ts index 8984f92d5c..2e628d06e7 100644 --- a/src/core/checkpoints/__tests__/checkpointAfterDelete.test.ts +++ b/src/core/checkpoints/__tests__/checkpointAfterDelete.test.ts @@ -8,9 +8,15 @@ import * as vscode from "vscode" vi.mock("vscode", () => ({ window: { showErrorMessage: vi.fn(), + createTextEditorDecorationType: vi.fn(() => ({})), + showInformationMessage: vi.fn(), }, Uri: { file: vi.fn((path: string) => ({ fsPath: path })), + parse: vi.fn((uri: string) => ({ with: vi.fn(() => ({})) })), + }, + commands: { + executeCommand: vi.fn(), }, })) @@ -35,7 +41,7 @@ describe("Checkpoint after message deletion", () => { beforeEach(() => { // Create mock checkpoint service mockCheckpointService = { - isInitialized: false, + isInitialized: true, // Set to true by default for most tests saveCheckpoint: vi.fn().mockResolvedValue({ commit: "test-commit-hash" }), on: vi.fn(), initShadowGit: vi.fn().mockResolvedValue(undefined), @@ -55,7 +61,7 @@ describe("Checkpoint after message deletion", () => { mockTask = { taskId: "test-task-id", enableCheckpoints: true, - checkpointService: undefined, + checkpointService: mockCheckpointService, // Set the service directly for most tests checkpointServiceInitializing: false, providerRef: { deref: () => mockProvider, @@ -77,6 +83,10 @@ describe("Checkpoint after message deletion", () => { }) it("should wait for checkpoint service initialization before saving", async () => { + // Set up task with uninitialized service + mockCheckpointService.isInitialized = false + mockTask.checkpointService = mockCheckpointService + // Simulate service initialization after a delay setTimeout(() => { mockCheckpointService.isInitialized = true @@ -85,15 +95,9 @@ describe("Checkpoint after message deletion", () => { // Call checkpointSave const savePromise = checkpointSave(mockTask, true) - // Initially, service should not be initialized - expect(mockCheckpointService.isInitialized).toBe(false) - // Wait for the save to complete const result = await savePromise - // Service should now be initialized - expect(mockCheckpointService.isInitialized).toBe(true) - // saveCheckpoint should have been called expect(mockCheckpointService.saveCheckpoint).toHaveBeenCalledWith( expect.stringContaining("Task: test-task-id"), @@ -130,6 +134,7 @@ describe("Checkpoint after message deletion", () => { it("should preserve checkpoint data through message deletion flow", async () => { // Initialize service mockCheckpointService.isInitialized = true + mockTask.checkpointService = mockCheckpointService // Simulate saving checkpoint before user message const checkpointResult = await checkpointSave(mockTask, true) @@ -150,15 +155,9 @@ describe("Checkpoint after message deletion", () => { // Simulate message deletion and reinitialization mockTask.clineMessages = [] - mockTask.checkpointService = undefined + mockTask.checkpointService = mockCheckpointService // Keep service available mockTask.checkpointServiceInitializing = false - // Re-initialize checkpoint service - setTimeout(() => { - mockCheckpointService.isInitialized = true - mockTask.checkpointService = mockCheckpointService - }, 50) - // Save checkpoint again after deletion const newCheckpointResult = await checkpointSave(mockTask, true) diff --git a/src/core/checkpoints/__tests__/utils.test.ts b/src/core/checkpoints/__tests__/utils.test.ts new file mode 100644 index 0000000000..eb4fccb96c --- /dev/null +++ b/src/core/checkpoints/__tests__/utils.test.ts @@ -0,0 +1,107 @@ +import { describe, it, expect } from "vitest" +import { isValidCheckpoint, hasValidCheckpoint, extractCheckpoint, type ValidCheckpoint } from "../utils" + +describe("checkpoint utils", () => { + describe("isValidCheckpoint", () => { + it("should return true for valid checkpoint", () => { + const checkpoint: ValidCheckpoint = { hash: "abc123" } + expect(isValidCheckpoint(checkpoint)).toBe(true) + }) + + it("should return false for null or undefined", () => { + expect(isValidCheckpoint(null)).toBe(false) + expect(isValidCheckpoint(undefined)).toBe(false) + }) + + it("should return false for non-object types", () => { + expect(isValidCheckpoint("string")).toBe(false) + expect(isValidCheckpoint(123)).toBe(false) + expect(isValidCheckpoint(true)).toBe(false) + expect(isValidCheckpoint([])).toBe(false) + }) + + it("should return false for objects without hash property", () => { + expect(isValidCheckpoint({})).toBe(false) + expect(isValidCheckpoint({ other: "property" })).toBe(false) + }) + + it("should return false for objects with non-string hash", () => { + expect(isValidCheckpoint({ hash: 123 })).toBe(false) + expect(isValidCheckpoint({ hash: null })).toBe(false) + expect(isValidCheckpoint({ hash: undefined })).toBe(false) + expect(isValidCheckpoint({ hash: {} })).toBe(false) + expect(isValidCheckpoint({ hash: [] })).toBe(false) + }) + + it("should return false for empty hash string", () => { + expect(isValidCheckpoint({ hash: "" })).toBe(false) + }) + + it("should return true for valid hash strings", () => { + expect(isValidCheckpoint({ hash: "a" })).toBe(true) + expect(isValidCheckpoint({ hash: "abc123def456" })).toBe(true) + expect(isValidCheckpoint({ hash: "commit-hash-with-dashes" })).toBe(true) + }) + }) + + describe("hasValidCheckpoint", () => { + it("should return true for message with valid checkpoint", () => { + const message = { checkpoint: { hash: "abc123" } } + expect(hasValidCheckpoint(message)).toBe(true) + }) + + it("should return false for null or undefined message", () => { + expect(hasValidCheckpoint(null)).toBe(false) + expect(hasValidCheckpoint(undefined)).toBe(false) + }) + + it("should return false for non-object message", () => { + expect(hasValidCheckpoint("string")).toBe(false) + expect(hasValidCheckpoint(123)).toBe(false) + }) + + it("should return false for message without checkpoint property", () => { + expect(hasValidCheckpoint({})).toBe(false) + expect(hasValidCheckpoint({ text: "message" })).toBe(false) + }) + + it("should return false for message with invalid checkpoint", () => { + expect(hasValidCheckpoint({ checkpoint: null })).toBe(false) + expect(hasValidCheckpoint({ checkpoint: "invalid" })).toBe(false) + expect(hasValidCheckpoint({ checkpoint: {} })).toBe(false) + expect(hasValidCheckpoint({ checkpoint: { hash: "" } })).toBe(false) + expect(hasValidCheckpoint({ checkpoint: { hash: 123 } })).toBe(false) + }) + + it("should work as type guard", () => { + const message: unknown = { checkpoint: { hash: "abc123" }, other: "data" } + if (hasValidCheckpoint(message)) { + // TypeScript should know message has checkpoint property + expect(message.checkpoint.hash).toBe("abc123") + } + }) + }) + + describe("extractCheckpoint", () => { + it("should extract valid checkpoint from message", () => { + const message = { checkpoint: { hash: "abc123" } } + const result = extractCheckpoint(message) + expect(result).toEqual({ hash: "abc123" }) + }) + + it("should return undefined for message without valid checkpoint", () => { + expect(extractCheckpoint({})).toBeUndefined() + expect(extractCheckpoint({ checkpoint: null })).toBeUndefined() + expect(extractCheckpoint({ checkpoint: { hash: "" } })).toBeUndefined() + expect(extractCheckpoint(null)).toBeUndefined() + expect(extractCheckpoint(undefined)).toBeUndefined() + }) + + it("should return the same checkpoint object reference", () => { + const checkpoint = { hash: "abc123" } + const message = { checkpoint } + const result = extractCheckpoint(message) + expect(result).toBe(checkpoint) + }) + }) +}) diff --git a/src/core/checkpoints/index.ts b/src/core/checkpoints/index.ts index cf0cb61f5e..6ece497bd7 100644 --- a/src/core/checkpoints/index.ts +++ b/src/core/checkpoints/index.ts @@ -14,6 +14,9 @@ import { DIFF_VIEW_URI_SCHEME } from "../../integrations/editor/DiffViewProvider import { CheckpointServiceOptions, RepoPerTaskCheckpointService } from "../../services/checkpoints" +// Map to store pending checkpoint operations by taskId to prevent race conditions +const pendingCheckpointOperations = new Map>() + export function getCheckpointService(cline: Task) { if (!cline.enableCheckpoints) { return undefined @@ -150,20 +153,48 @@ async function getInitializedCheckpointService( } export async function checkpointSave(cline: Task, force = false) { - // Use getInitializedCheckpointService to wait for initialization - const service = await getInitializedCheckpointService(cline) + const taskId = cline.taskId - if (!service) { - return + // Check if there's already a pending checkpoint operation for this task + const existingOperation = pendingCheckpointOperations.get(taskId) + if (existingOperation) { + // Return the existing Promise to prevent duplicate operations + return existingOperation } - TelemetryService.instance.captureCheckpointCreated(cline.taskId) + // Create a new checkpoint operation Promise + const checkpointOperation = (async () => { + try { + // Use getInitializedCheckpointService to wait for initialization + const service = await getInitializedCheckpointService(cline) - // Start the checkpoint process in the background. - return service.saveCheckpoint(`Task: ${cline.taskId}, Time: ${Date.now()}`, { allowEmpty: force }).catch((err) => { - console.error("[Task#checkpointSave] caught unexpected error, disabling checkpoints", err) - cline.enableCheckpoints = false - }) + if (!service) { + return + } + + TelemetryService.instance.captureCheckpointCreated(cline.taskId) + + // Start the checkpoint process in the background. + return await service.saveCheckpoint(`Task: ${cline.taskId}, Time: ${Date.now()}`, { allowEmpty: force }) + } catch (err) { + console.error("[Task#checkpointSave] caught unexpected error, disabling checkpoints", err) + cline.enableCheckpoints = false + } + })() + + // Store the operation in the Map + pendingCheckpointOperations.set(taskId, checkpointOperation) + + // Clean up the Map entry after the operation completes (success or failure) + checkpointOperation + .finally(() => { + pendingCheckpointOperations.delete(taskId) + }) + .catch(() => { + // Error already handled above, this catch prevents unhandled rejection + }) + + return checkpointOperation } export type CheckpointRestoreOptions = { diff --git a/src/core/checkpoints/utils.ts b/src/core/checkpoints/utils.ts new file mode 100644 index 0000000000..0b197d221d --- /dev/null +++ b/src/core/checkpoints/utils.ts @@ -0,0 +1,51 @@ +/** + * Checkpoint-related utilities and type definitions + */ + +/** + * Represents a valid checkpoint with required properties + */ +export interface ValidCheckpoint { + hash: string +} + +/** + * Type guard to check if an object is a valid checkpoint + * @param checkpoint - The object to check + * @returns True if the checkpoint is valid, false otherwise + */ +export function isValidCheckpoint(checkpoint: unknown): checkpoint is ValidCheckpoint { + return ( + checkpoint !== null && + checkpoint !== undefined && + typeof checkpoint === "object" && + "hash" in checkpoint && + typeof (checkpoint as any).hash === "string" && + (checkpoint as any).hash.length > 0 // Ensure hash is not empty + ) +} + +/** + * Validates if a message has a valid checkpoint for restoration + * @param message - The message object to check + * @returns True if the message contains a valid checkpoint, false otherwise + */ +export function hasValidCheckpoint(message: unknown): message is { checkpoint: ValidCheckpoint } { + if (!message || typeof message !== "object" || !("checkpoint" in message)) { + return false + } + + return isValidCheckpoint((message as any).checkpoint) +} + +/** + * Extracts a valid checkpoint from a message if it exists + * @param message - The message object to extract from + * @returns The valid checkpoint or undefined + */ +export function extractCheckpoint(message: unknown): ValidCheckpoint | undefined { + if (hasValidCheckpoint(message)) { + return message.checkpoint + } + return undefined +} diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index 16d88e4834..fea4a9f067 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -352,14 +352,7 @@ export class Task extends EventEmitter { // Cline Messages private async getSavedClineMessages(): Promise { - const messages = await readTaskMessages({ taskId: this.taskId, globalStoragePath: this.globalStoragePath }) - console.log("[Task#getSavedClineMessages] Loaded messages from disk:", messages.length) - const messagesWithCheckpoints = messages.filter((m) => m.checkpoint) - console.log("[Task#getSavedClineMessages] Messages with checkpoints:", messagesWithCheckpoints.length) - if (messagesWithCheckpoints.length > 0) { - console.log("[Task#getSavedClineMessages] Sample checkpoint:", messagesWithCheckpoints[0].checkpoint) - } - return messages + return await readTaskMessages({ taskId: this.taskId, globalStoragePath: this.globalStoragePath }) } private async addToClineMessages(message: ClineMessage) { @@ -534,7 +527,14 @@ export class Task extends EventEmitter { await this.addToClineMessages({ ts: askTs, type: "ask", ask: type, text, isProtected }) } - await pWaitFor(() => this.askResponse !== undefined || this.lastMessageTs !== askTs, { interval: 100 }) + await pWaitFor(() => this.askResponse !== undefined || this.lastMessageTs !== askTs || this.abort, { + interval: 100, + }) + + if (this.abort) { + // Task was aborted, return a default response + return { response: "messageResponse", text: undefined, images: undefined } + } if (this.lastMessageTs !== askTs) { // Could happen if we send multiple asks in a row i.e. with @@ -734,24 +734,6 @@ export class Task extends EventEmitter { const feedbackCheckpoint = checkpoint || this.pendingUserMessageCheckpoint this.pendingUserMessageCheckpoint = undefined // Clear it after use - console.log("[Task#say] Adding user_feedback message with checkpoint:", feedbackCheckpoint) - console.log( - "[Task#say] Full message object:", - JSON.stringify( - { - ts: sayTs, - type: "say", - say: type, - text, - images, - checkpoint: feedbackCheckpoint, - contextCondense, - }, - null, - 2, - ), - ) - await this.addToClineMessages({ ts: sayTs, type: "say", @@ -842,18 +824,6 @@ export class Task extends EventEmitter { private async resumeTaskFromHistory() { const modifiedClineMessages = await this.getSavedClineMessages() - // Debug: Check if any messages have checkpoints - const messagesWithCheckpoints = modifiedClineMessages.filter((m) => m.checkpoint) - console.log("[Task#resumeTaskFromHistory] Total messages loaded:", modifiedClineMessages.length) - console.log("[Task#resumeTaskFromHistory] Messages with checkpoints:", messagesWithCheckpoints.length) - messagesWithCheckpoints.forEach((msg, idx) => { - console.log(`[Task#resumeTaskFromHistory] Message ${idx} with checkpoint:`, { - ts: msg.ts, - say: msg.say, - checkpoint: msg.checkpoint, - }) - }) - // Remove any resume messages that may have been added before const lastRelevantMessageIndex = findLastIndex( modifiedClineMessages, @@ -1161,6 +1131,13 @@ export class Task extends EventEmitter { this.abandoned = true } + // Resolve any pending ask operations to prevent "Current ask promise was ignored" errors + if (this.askResponse === undefined) { + this.askResponse = "messageResponse" + this.askResponseText = undefined + this.askResponseImages = undefined + } + this.abort = true this.emit("taskAborted") diff --git a/src/core/task/__tests__/Task.spec.ts b/src/core/task/__tests__/Task.spec.ts index 693f72d1c7..9d342d324c 100644 --- a/src/core/task/__tests__/Task.spec.ts +++ b/src/core/task/__tests__/Task.spec.ts @@ -1334,5 +1334,89 @@ describe("Cline", () => { expect(task.diffStrategy).toBeUndefined() }) }) + + describe("Ask Operation Abort Handling", () => { + let mockProvider: any + let mockApiConfig: any + + beforeEach(() => { + vi.clearAllMocks() + + mockApiConfig = { + apiProvider: "anthropic", + apiKey: "test-key", + } + + mockProvider = { + context: { + globalStorageUri: { fsPath: "/test/storage" }, + }, + getState: vi.fn().mockResolvedValue({}), + postMessageToWebview: vi.fn(), + postStateToWebview: vi.fn().mockResolvedValue(undefined), + } + }) + + it("should handle pending ask operations gracefully when task is aborted", async () => { + const [cline, task] = Task.create({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "test task", + }) + + // Start an ask operation but don't respond to it + const askPromise = cline.ask("tool", "Test question") + + // Abort the task while ask is pending + await cline.abortTask() + + // The ask should resolve with a default response instead of throwing + const result = await askPromise + expect(result.response).toBe("messageResponse") + expect(result.text).toBeUndefined() + expect(result.images).toBeUndefined() + + // Ensure the task was properly aborted + expect(cline.abort).toBe(true) + }) + + it("should not throw 'Current ask promise was ignored' error when task is aborted", async () => { + const [cline, task] = Task.create({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "test task", + }) + + // Start multiple ask operations + const askPromise1 = cline.ask("tool", "Question 1") + const askPromise2 = cline.ask("tool", "Question 2") + + // Abort the task + await cline.abortTask() + + // Both asks should resolve without throwing + await expect(askPromise1).resolves.toBeTruthy() + await expect(askPromise2).resolves.toBeTruthy() + }) + + it("should resolve pending ask with messageResponse when abortTask is called", async () => { + const [cline, task] = Task.create({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "test task", + }) + + // Spy on the ask response properties + // Start an ask operation + const askPromise = cline.ask("tool", "Test question") + + // Abort the task + await cline.abortTask() + + // The ask response should have been set to messageResponse + const result = await askPromise + expect(result.response).toBe("messageResponse") + }) + }) }) }) diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index 550d258d84..b07ac498f9 100644 --- a/src/core/webview/ClineProvider.ts +++ b/src/core/webview/ClineProvider.ts @@ -79,6 +79,17 @@ export type ClineProviderEvents = { clineCreated: [cline: Task] } +interface PendingEditOperation { + messageTs: number + editedContent: string + images?: string[] + messageIndex: number + apiConversationHistoryIndex: number + originalCheckpoint: { hash: string } + timeoutId: NodeJS.Timeout + createdAt: number +} + class OrganizationAllowListViolationError extends Error { constructor(message: string) { super(message) @@ -107,6 +118,8 @@ export class ClineProvider protected mcpHub?: McpHub // Change from private to protected private marketplaceManager: MarketplaceManager private mdmService?: MdmService + private pendingOperations: Map = new Map() + private static readonly PENDING_OPERATION_TIMEOUT_MS = 30000 // 30 seconds public isViewLaunched = false public settingsImportedAt?: number @@ -240,6 +253,73 @@ export class ClineProvider await this.removeClineFromStack() } + // Pending Edit Operations Management + + /** + * Sets a pending edit operation with automatic timeout cleanup + */ + public setPendingEditOperation( + operationId: string, + editData: { + messageTs: number + editedContent: string + images?: string[] + messageIndex: number + apiConversationHistoryIndex: number + originalCheckpoint: { hash: string } + }, + ): void { + // Clear any existing operation with the same ID + this.clearPendingEditOperation(operationId) + + // Create timeout for automatic cleanup + const timeoutId = setTimeout(() => { + this.clearPendingEditOperation(operationId) + this.log(`[setPendingEditOperation] Automatically cleared stale pending operation: ${operationId}`) + }, ClineProvider.PENDING_OPERATION_TIMEOUT_MS) + + // Store the operation + this.pendingOperations.set(operationId, { + ...editData, + timeoutId, + createdAt: Date.now(), + }) + + this.log(`[setPendingEditOperation] Set pending operation: ${operationId}`) + } + + /** + * Gets a pending edit operation by ID + */ + private getPendingEditOperation(operationId: string): PendingEditOperation | undefined { + return this.pendingOperations.get(operationId) + } + + /** + * Clears a specific pending edit operation + */ + private clearPendingEditOperation(operationId: string): boolean { + const operation = this.pendingOperations.get(operationId) + if (operation) { + clearTimeout(operation.timeoutId) + this.pendingOperations.delete(operationId) + this.log(`[clearPendingEditOperation] Cleared pending operation: ${operationId}`) + return true + } + return false + } + + /** + * Clears all pending edit operations + */ + private clearAllPendingEditOperations(): void { + for (const [operationId, operation] of this.pendingOperations) { + clearTimeout(operation.timeoutId) + } + this.pendingOperations.clear() + this.log(`[clearAllPendingEditOperations] Cleared all pending operations`) + } + /* VSCode extensions use the disposable pattern to clean up resources when the sidebar/editor tab is closed by the user or system. This applies to event listening, commands, interacting with the UI, etc. - https://vscode-docs.readthedocs.io/en/stable/extensions/patterns-and-principles/ @@ -259,6 +339,10 @@ export class ClineProvider await this.removeClineFromStack() this.log("Cleared task") + // Clear all pending edit operations to prevent memory leaks + this.clearAllPendingEditOperations() + this.log("Cleared pending operations") + if (this.view && "dispose" in this.view) { this.view.dispose() this.log("Disposed webview") @@ -607,9 +691,10 @@ export class ClineProvider ) // Check if there's a pending edit after checkpoint restoration - if ((this as any).pendingEditAfterRestore) { - const pendingEdit = (this as any).pendingEditAfterRestore - ;(this as any).pendingEditAfterRestore = undefined // Clear the pending edit + const operationId = `task-${cline.taskId}` + const pendingEdit = this.getPendingEditOperation(operationId) + if (pendingEdit) { + this.clearPendingEditOperation(operationId) // Clear the pending edit this.log(`[initClineWithHistoryItem] Processing pending edit after checkpoint restoration`) @@ -1502,13 +1587,6 @@ export class ClineProvider "[ClineProvider#getStateToPostToWebview] Messages with checkpoints:", messagesWithCheckpoints.length, ) - if (messagesWithCheckpoints.length > 0) { - console.log("[ClineProvider#getStateToPostToWebview] Sample message with checkpoint:", { - ts: messagesWithCheckpoints[0].ts, - say: messagesWithCheckpoints[0].say, - checkpoint: messagesWithCheckpoints[0].checkpoint, - }) - } return messages })(), taskHistory: (taskHistory || []) diff --git a/src/core/webview/__tests__/checkpointRestoreHandler.test.ts b/src/core/webview/__tests__/checkpointRestoreHandler.test.ts index d79349ed63..468e5188d5 100644 --- a/src/core/webview/__tests__/checkpointRestoreHandler.test.ts +++ b/src/core/webview/__tests__/checkpointRestoreHandler.test.ts @@ -1,5 +1,6 @@ import { describe, it, expect, vi, beforeEach } from "vitest" -import { handleCheckpointRestoreOperation, hasValidCheckpoint } from "../checkpointRestoreHandler" +import { handleCheckpointRestoreOperation } from "../checkpointRestoreHandler" +import { hasValidCheckpoint } from "../../checkpoints/utils" import { saveTaskMessages } from "../../task-persistence" import * as vscode from "vscode" @@ -27,13 +28,14 @@ describe("checkpointRestoreHandler", () => { }, }, getTaskWithId: vi.fn().mockResolvedValue({ - historyItem: { id: "test-task", messages: [] }, + historyItem: { id: "task123", messages: [] }, }), initClineWithHistoryItem: vi.fn(), + setPendingEditOperation: vi.fn(), } mockCline = { - taskId: "test-task", + taskId: "task123", clineMessages: [ { ts: 1, text: "Message 1" }, { ts: 2, text: "Message 2", checkpoint: { hash: "abc123" } }, @@ -59,6 +61,10 @@ describe("checkpointRestoreHandler", () => { expect(hasValidCheckpoint({ checkpoint: {} })).toBe(false) expect(hasValidCheckpoint({ checkpoint: { hash: 123 } })).toBe(false) }) + + it("should return false for empty hash", () => { + expect(hasValidCheckpoint({ checkpoint: { hash: "" } })).toBe(false) + }) }) describe("handleCheckpointRestoreOperation", () => { @@ -84,14 +90,14 @@ describe("checkpointRestoreHandler", () => { // Should save messages after restoration expect(saveTaskMessages).toHaveBeenCalledWith({ messages: mockCline.clineMessages, - taskId: "test-task", + taskId: "task123", globalStoragePath: "/test/global/storage", }) // Should reinitialize the task - expect(mockProvider.getTaskWithId).toHaveBeenCalledWith("test-task") + expect(mockProvider.getTaskWithId).toHaveBeenCalledWith("task123") expect(mockProvider.initClineWithHistoryItem).toHaveBeenCalledWith({ - id: "test-task", + id: "task123", messages: [], }) }) @@ -115,8 +121,8 @@ describe("checkpointRestoreHandler", () => { editData, }) - // Should set pendingEditAfterRestore on provider - expect(mockProvider.pendingEditAfterRestore).toEqual({ + // Should call setPendingEditOperation on provider + expect(mockProvider.setPendingEditOperation).toHaveBeenCalledWith("task-task123", { messageTs: 2, editedContent: "Edited content", images: ["image1.png"], diff --git a/src/core/webview/checkpointRestoreHandler.ts b/src/core/webview/checkpointRestoreHandler.ts index a9a6c83f66..d25bf34ea6 100644 --- a/src/core/webview/checkpointRestoreHandler.ts +++ b/src/core/webview/checkpointRestoreHandler.ts @@ -4,13 +4,14 @@ import { saveTaskMessages } from "../task-persistence" import * as vscode from "vscode" import pWaitFor from "p-wait-for" import { t } from "../../i18n" +import { ValidCheckpoint } from "../checkpoints/utils" export interface CheckpointRestoreConfig { provider: ClineProvider currentCline: Task messageTs: number messageIndex: number - checkpoint: { hash: string } + checkpoint: ValidCheckpoint operation: "delete" | "edit" editData?: { editedContent: string @@ -29,14 +30,15 @@ export async function handleCheckpointRestoreOperation(config: CheckpointRestore try { // For edit operations, set up pending edit data before restoration if (operation === "edit" && editData) { - ;(provider as any).pendingEditAfterRestore = { + const operationId = `task-${currentCline.taskId}` + provider.setPendingEditOperation(operationId, { messageTs, editedContent: editData.editedContent, images: editData.images, messageIndex: config.messageIndex, apiConversationHistoryIndex: editData.apiConversationHistoryIndex, originalCheckpoint: checkpoint, - } + }) } // Perform the checkpoint restoration @@ -73,19 +75,6 @@ export async function handleCheckpointRestoreOperation(config: CheckpointRestore } } -/** - * Validates if a message has a valid checkpoint for restoration - */ -export function hasValidCheckpoint(message: any): boolean { - return ( - (message?.checkpoint && - typeof message.checkpoint === "object" && - "hash" in message.checkpoint && - typeof message.checkpoint.hash === "string") || - false - ) -} - /** * Common checkpoint restore validation and initialization utility. * This can be used by any checkpoint restore flow that needs to wait for initialization. diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index 0258fe1657..c767c7d6a3 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -19,11 +19,8 @@ import { type ApiMessage } from "../task-persistence/apiMessages" import { saveTaskMessages } from "../task-persistence" import { ClineProvider } from "./ClineProvider" -import { - handleCheckpointRestoreOperation, - hasValidCheckpoint, - waitForClineInitialization, -} from "./checkpointRestoreHandler" +import { handleCheckpointRestoreOperation } from "./checkpointRestoreHandler" +import { ValidCheckpoint, hasValidCheckpoint } from "../checkpoints/utils" import { changeLanguage, t } from "../../i18n" import { Package } from "../../shared/package" import { RouterName, toRouterName, ModelRecord } from "../../shared/api" @@ -124,11 +121,8 @@ export const webviewMessageHandler = async ( ) const { messageIndex } = findMessageIndices(messageTs, currentCline) - console.log("[webviewMessageHandler] Checking for checkpoint at messageIndex:", messageIndex) if (messageIndex !== -1) { const targetMessage = currentCline.clineMessages[messageIndex] - console.log("[webviewMessageHandler] Target message:", JSON.stringify(targetMessage, null, 2)) - console.log("[webviewMessageHandler] Target message checkpoint:", targetMessage?.checkpoint) hasCheckpoint = !!( targetMessage?.checkpoint && typeof targetMessage.checkpoint === "object" && @@ -152,61 +146,68 @@ export const webviewMessageHandler = async ( * Handles confirmed message deletion from webview dialog */ const handleDeleteMessageConfirm = async (messageTs: number, restoreCheckpoint?: boolean): Promise => { - // Only proceed if we have a current cline - if (provider.getCurrentCline()) { - const currentCline = provider.getCurrentCline()! - const { messageIndex, apiConversationHistoryIndex } = findMessageIndices(messageTs, currentCline) + const currentCline = provider.getCurrentCline() + if (!currentCline) { + console.error("[handleDeleteMessageConfirm] No current cline available") + return + } - if (messageIndex !== -1) { - try { - const targetMessage = currentCline.clineMessages[messageIndex] + const { messageIndex, apiConversationHistoryIndex } = findMessageIndices(messageTs, currentCline) - // If checkpoint restoration is requested, restore to the checkpoint first - if (restoreCheckpoint && hasValidCheckpoint(targetMessage)) { - await handleCheckpointRestoreOperation({ - provider, - currentCline, - messageTs: targetMessage.ts!, - messageIndex, - checkpoint: targetMessage.checkpoint as { hash: string }, - operation: "delete", - }) - } else { - // For non-checkpoint deletes, preserve checkpoint associations for remaining messages - // Store checkpoints from messages that will be preserved - const preservedCheckpoints = new Map() - for (let i = 0; i < messageIndex; i++) { - const msg = currentCline.clineMessages[i] - if (msg?.checkpoint && msg.ts) { - preservedCheckpoints.set(msg.ts, msg.checkpoint) - } - } + if (messageIndex === -1) { + const errorMessage = `Message with timestamp ${messageTs} not found` + console.error("[handleDeleteMessageConfirm]", errorMessage) + await vscode.window.showErrorMessage(errorMessage) + return + } - // Delete this message and all subsequent messages - await removeMessagesThisAndSubsequent(currentCline, messageIndex, apiConversationHistoryIndex) + try { + const targetMessage = currentCline.clineMessages[messageIndex] - // Restore checkpoint associations for preserved messages - for (const [ts, checkpoint] of preservedCheckpoints) { - const msgIndex = currentCline.clineMessages.findIndex((msg) => msg.ts === ts) - if (msgIndex !== -1) { - currentCline.clineMessages[msgIndex].checkpoint = checkpoint - } - } - - // Save the updated messages with restored checkpoints - await saveTaskMessages({ - messages: currentCline.clineMessages, - taskId: currentCline.taskId, - globalStoragePath: provider.contextProxy.globalStorageUri.fsPath, - }) + // If checkpoint restoration is requested, restore to the checkpoint first + if (restoreCheckpoint && hasValidCheckpoint(targetMessage)) { + await handleCheckpointRestoreOperation({ + provider, + currentCline, + messageTs: targetMessage.ts!, + messageIndex, + checkpoint: targetMessage.checkpoint as ValidCheckpoint, + operation: "delete", + }) + } else { + // For non-checkpoint deletes, preserve checkpoint associations for remaining messages + // Store checkpoints from messages that will be preserved + const preservedCheckpoints = new Map() + for (let i = 0; i < messageIndex; i++) { + const msg = currentCline.clineMessages[i] + if (msg?.checkpoint && msg.ts) { + preservedCheckpoints.set(msg.ts, msg.checkpoint) } - } catch (error) { - console.error("Error in delete message:", error) - vscode.window.showErrorMessage( - `Error deleting message: ${error instanceof Error ? error.message : String(error)}`, - ) } + + // Delete this message and all subsequent messages + await removeMessagesThisAndSubsequent(currentCline, messageIndex, apiConversationHistoryIndex) + + // Restore checkpoint associations for preserved messages + for (const [ts, checkpoint] of preservedCheckpoints) { + const msgIndex = currentCline.clineMessages.findIndex((msg) => msg.ts === ts) + if (msgIndex !== -1) { + currentCline.clineMessages[msgIndex].checkpoint = checkpoint + } + } + + // Save the updated messages with restored checkpoints + await saveTaskMessages({ + messages: currentCline.clineMessages, + taskId: currentCline.taskId, + globalStoragePath: provider.contextProxy.globalStorageUri.fsPath, + }) } + } catch (error) { + console.error("Error in delete message:", error) + vscode.window.showErrorMessage( + `Error deleting message: ${error instanceof Error ? error.message : String(error)}`, + ) } } @@ -232,11 +233,8 @@ export const webviewMessageHandler = async ( }) const { messageIndex } = findMessageIndices(messageTs, currentCline) - console.log("[webviewMessageHandler] Edit - Checking for checkpoint at messageIndex:", messageIndex) if (messageIndex !== -1) { const targetMessage = currentCline.clineMessages[messageIndex] - console.log("[webviewMessageHandler] Edit - Target message:", JSON.stringify(targetMessage, null, 2)) - console.log("[webviewMessageHandler] Edit - Target message checkpoint:", targetMessage?.checkpoint) hasCheckpoint = !!( targetMessage?.checkpoint && typeof targetMessage.checkpoint === "object" && @@ -279,91 +277,97 @@ export const webviewMessageHandler = async ( restoreCheckpoint?: boolean, images?: string[], ): Promise => { - // Only proceed if we have a current cline - if (provider.getCurrentCline()) { - const currentCline = provider.getCurrentCline()! + const currentCline = provider.getCurrentCline() + if (!currentCline) { + console.error("[handleEditMessageConfirm] No current cline available") + return + } - // Use findMessageIndices to find messages based on timestamp - const { messageIndex, apiConversationHistoryIndex } = findMessageIndices(messageTs, currentCline) + // Use findMessageIndices to find messages based on timestamp + const { messageIndex, apiConversationHistoryIndex } = findMessageIndices(messageTs, currentCline) - if (messageIndex !== -1) { - try { - const targetMessage = currentCline.clineMessages[messageIndex] + if (messageIndex === -1) { + const errorMessage = `Message with timestamp ${messageTs} not found` + console.error("[handleEditMessageConfirm]", errorMessage) + await vscode.window.showErrorMessage(errorMessage) + return + } - // Preserve the original checkpoint data for the edited message - const originalCheckpoint = targetMessage?.checkpoint + try { + const targetMessage = currentCline.clineMessages[messageIndex] - // If checkpoint restoration is requested, restore to the checkpoint first - if (restoreCheckpoint && hasValidCheckpoint(targetMessage)) { - await handleCheckpointRestoreOperation({ - provider, - currentCline, - messageTs: targetMessage.ts!, - messageIndex, - checkpoint: targetMessage.checkpoint as { hash: string }, - operation: "edit", - editData: { - editedContent, - images, - apiConversationHistoryIndex, - }, - }) - // The task will be cancelled and reinitialized by checkpointRestore - // The pending edit will be processed in the reinitialized task - return - } + // Preserve the original checkpoint data for the edited message + const originalCheckpoint = targetMessage?.checkpoint - // For non-checkpoint edits, preserve checkpoint associations for remaining messages - // Store checkpoints from messages that will be preserved - const preservedCheckpoints = new Map() - for (let i = 0; i < messageIndex; i++) { - const msg = currentCline.clineMessages[i] - if (msg?.checkpoint && msg.ts) { - preservedCheckpoints.set(msg.ts, msg.checkpoint) - } - } - - // Edit this message and delete subsequent - await removeMessagesThisAndSubsequent(currentCline, messageIndex, apiConversationHistoryIndex) - - // Restore checkpoint associations for preserved messages - for (const [ts, checkpoint] of preservedCheckpoints) { - const msgIndex = currentCline.clineMessages.findIndex((msg) => msg.ts === ts) - if (msgIndex !== -1) { - currentCline.clineMessages[msgIndex].checkpoint = checkpoint - } - } - - // Save the updated messages with restored checkpoints - await saveTaskMessages({ - messages: currentCline.clineMessages, - taskId: currentCline.taskId, - globalStoragePath: provider.contextProxy.globalStorageUri.fsPath, - }) - - // Process the edited message as a regular user message - // Preserve the original checkpoint for the new message - if (originalCheckpoint) { - // Store the checkpoint to be attached to the new message - currentCline.pendingUserMessageCheckpoint = originalCheckpoint - } - - webviewMessageHandler(provider, { - type: "askResponse", - askResponse: "messageResponse", - text: editedContent, + // If checkpoint restoration is requested, restore to the checkpoint first + if (restoreCheckpoint && hasValidCheckpoint(targetMessage)) { + await handleCheckpointRestoreOperation({ + provider, + currentCline, + messageTs: targetMessage.ts!, + messageIndex, + checkpoint: targetMessage.checkpoint as ValidCheckpoint, + operation: "edit", + editData: { + editedContent, images, - }) + apiConversationHistoryIndex, + }, + }) + // The task will be cancelled and reinitialized by checkpointRestore + // The pending edit will be processed in the reinitialized task + return + } - // Don't initialize with history item for edit operations - // The webviewMessageHandler will handle the conversation state - } catch (error) { - console.error("Error in edit message:", error) - vscode.window.showErrorMessage( - `Error editing message: ${error instanceof Error ? error.message : String(error)}`, - ) + // For non-checkpoint edits, preserve checkpoint associations for remaining messages + // Store checkpoints from messages that will be preserved + const preservedCheckpoints = new Map() + for (let i = 0; i < messageIndex; i++) { + const msg = currentCline.clineMessages[i] + if (msg?.checkpoint && msg.ts) { + preservedCheckpoints.set(msg.ts, msg.checkpoint) } } + + // Edit this message and delete subsequent + await removeMessagesThisAndSubsequent(currentCline, messageIndex, apiConversationHistoryIndex) + + // Restore checkpoint associations for preserved messages + for (const [ts, checkpoint] of preservedCheckpoints) { + const msgIndex = currentCline.clineMessages.findIndex((msg) => msg.ts === ts) + if (msgIndex !== -1) { + currentCline.clineMessages[msgIndex].checkpoint = checkpoint + } + } + + // Save the updated messages with restored checkpoints + await saveTaskMessages({ + messages: currentCline.clineMessages, + taskId: currentCline.taskId, + globalStoragePath: provider.contextProxy.globalStorageUri.fsPath, + }) + + // Process the edited message as a regular user message + // Preserve the original checkpoint for the new message + if (originalCheckpoint) { + // Store the checkpoint to be attached to the new message + currentCline.pendingUserMessageCheckpoint = originalCheckpoint + } + + webviewMessageHandler(provider, { + type: "askResponse", + askResponse: "messageResponse", + text: editedContent, + images, + }) + + // Don't initialize with history item for edit operations + // The webviewMessageHandler will handle the conversation state + } catch (error) { + console.error("Error in edit message:", error) + vscode.window.showErrorMessage( + `Error editing message: ${error instanceof Error ? error.message : String(error)}`, + ) } } @@ -523,10 +527,8 @@ export const webviewMessageHandler = async ( // Save checkpoint BEFORE processing the user message if checkpoints are enabled const currentCline = provider.getCurrentCline() if (currentCline && currentCline.enableCheckpoints && message.askResponse === "messageResponse") { - console.log("[webviewMessageHandler] Saving checkpoint before user message processing") try { const checkpointResult = await currentCline.checkpointSave(true) // Force checkpoint save - console.log("[webviewMessageHandler] Checkpoint result:", checkpointResult) if (checkpointResult?.commit) { // Store checkpoint data temporarily to be used when creating the user_feedback message currentCline.pendingUserMessageCheckpoint = { @@ -534,12 +536,6 @@ export const webviewMessageHandler = async ( timestamp: Date.now(), type: "user_message", } - console.log( - "[webviewMessageHandler] Set pendingUserMessageCheckpoint:", - currentCline.pendingUserMessageCheckpoint, - ) - } else { - console.log("[webviewMessageHandler] No commit in checkpoint result") } } catch (error) { console.error("[webviewMessageHandler] Failed to save checkpoint before user message:", error) diff --git a/webview-ui/src/components/chat/ChatView.tsx b/webview-ui/src/components/chat/ChatView.tsx index 1de8c096d9..7ee175d6c1 100644 --- a/webview-ui/src/components/chat/ChatView.tsx +++ b/webview-ui/src/components/chat/ChatView.tsx @@ -783,18 +783,24 @@ const ChatViewComponent: React.ForwardRefRenderFunction { + // Pre-compute checkpoint hashes that have associated user messages for O(1) lookup + const userMessageCheckpointHashes = new Set() + modifiedMessages.forEach((msg) => { + if ( + msg.say === "user_feedback" && + msg.checkpoint && + (msg.checkpoint as any).type === "user_message" && + (msg.checkpoint as any).hash + ) { + userMessageCheckpointHashes.add((msg.checkpoint as any).hash) + } + }) + const newVisibleMessages = modifiedMessages.filter((message) => { // Filter out checkpoint_saved messages that are associated with user messages if (message.say === "checkpoint_saved" && message.text) { - // Check if there's a user_feedback message with a checkpoint that has this hash - const hasAssociatedUserMessage = modifiedMessages.some( - (msg) => - msg.say === "user_feedback" && - msg.checkpoint && - (msg.checkpoint as any).type === "user_message" && - (msg.checkpoint as any).hash === message.text, - ) - if (hasAssociatedUserMessage) { + // Use O(1) Set lookup instead of O(n) array search + if (userMessageCheckpointHashes.has(message.text)) { return false } }