diff --git a/src/core/checkpoints/__tests__/index.spec.ts b/src/core/checkpoints/__tests__/index.spec.ts index 97090b7296..30dcd45312 100644 --- a/src/core/checkpoints/__tests__/index.spec.ts +++ b/src/core/checkpoints/__tests__/index.spec.ts @@ -79,6 +79,7 @@ const createMockTask = (options: { taskId: string; hasExistingCheckpoints: boole enableCheckpoints: options.enableCheckpoints ?? true, checkpointService: null as any, checkpointServiceInitializing: false, + ongoingCheckpointSaves: new Map(), clineMessages: options.hasExistingCheckpoints ? [{ say: "checkpoint_saved", ts: Date.now(), text: "existing-checkpoint-hash" }] : [], diff --git a/src/core/checkpoints/index.ts b/src/core/checkpoints/index.ts index 1f443c06c1..9f58e9490e 100644 --- a/src/core/checkpoints/index.ts +++ b/src/core/checkpoints/index.ts @@ -360,24 +360,21 @@ export async function getInitializedCheckpointService( } } -// Track ongoing checkpoint saves per task to prevent duplicates -const ongoingCheckpointSaves = new Map>() - export async function checkpointSave(cline: Task, force = false, files?: vscode.Uri[]) { - // Create a unique key for this checkpoint save operation + // Create a unique key for this checkpoint save operation (task-scoped, no need for taskId in key) const filesKey = files ? files .map((f) => f.fsPath) .sort() .join("|") : "all" - const saveKey = `${cline.taskId}-${force}-${filesKey}` + const saveKey = `${force}-${filesKey}` // If there's already an ongoing checkpoint save for this exact operation, return the existing promise - if (ongoingCheckpointSaves.has(saveKey)) { + if (cline.ongoingCheckpointSaves.has(saveKey)) { const provider = cline.providerRef.deref() provider?.log(`[checkpointSave] duplicate checkpoint save detected for ${saveKey}, using existing operation`) - return ongoingCheckpointSaves.get(saveKey) + return cline.ongoingCheckpointSaves.get(saveKey) } const service = await getInitializedCheckpointService(cline) @@ -422,10 +419,10 @@ export async function checkpointSave(cline: Task, force = false, files?: vscode. }) .finally(() => { // Clean up the tracking once completed - ongoingCheckpointSaves.delete(saveKey) + cline.ongoingCheckpointSaves.delete(saveKey) }) - ongoingCheckpointSaves.set(saveKey, savePromise) + cline.ongoingCheckpointSaves.set(saveKey, savePromise) return savePromise } diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index 104cb87206..04b0a5aca4 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -60,6 +60,7 @@ import { BrowserSession } from "../../services/browser/BrowserSession" import { McpHub } from "../../services/mcp/McpHub" import { McpServerManager } from "../../services/mcp/McpServerManager" import { RepoPerTaskCheckpointService } from "../../services/checkpoints" +import { CheckpointResult } from "../../services/checkpoints/types" // integrations import { DiffViewProvider } from "../../integrations/editor/DiffViewProvider" @@ -252,6 +253,7 @@ export class Task extends EventEmitter implements TaskLike { enableCheckpoints: boolean checkpointService?: RepoPerTaskCheckpointService checkpointServiceInitializing = false + ongoingCheckpointSaves = new Map>() // Task Bridge enableTaskBridge: boolean @@ -1397,6 +1399,13 @@ export class Task extends EventEmitter implements TaskLike { console.error("Error removing event listeners:", error) } + // Clean up ongoing checkpoint saves to prevent memory leaks + try { + this.ongoingCheckpointSaves.clear() + } catch (error) { + console.error("Error clearing ongoing checkpoint saves:", error) + } + // Stop waiting for child task completion. if (this.pauseInterval) { clearInterval(this.pauseInterval)