From 94c46f44fb4d83eac3742d5a2d1cf30d1bed3f7b Mon Sep 17 00:00:00 2001
From: Shawn <5414767+playcations@users.noreply.github.com>
Date: Mon, 1 Sep 2025 11:33:44 -0400
Subject: [PATCH] possible checkpoint memory leak

- Fix checkpoint memory leak by making
  ongoingCheckpointSaves task-scoped
- Moved ongoingCheckpointSaves Map from module-level to
  Task class property
- Add cleanup in Task.dispose() method to prevent memory
   leaks
- Update checkpoint functions to use task-scoped Map
- Fix test mock to include ongoingCheckpointSaves
  property
---
 src/core/checkpoints/__tests__/index.spec.ts |  1 +
 src/core/checkpoints/index.ts                | 15 ++++++---------
 src/core/task/Task.ts                        |  9 +++++++++
 3 files changed, 16 insertions(+), 9 deletions(-)

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<string, Promise<void | CheckpointResult | undefined>>()
-
 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<TaskEvents> implements TaskLike {
 	enableCheckpoints: boolean
 	checkpointService?: RepoPerTaskCheckpointService
 	checkpointServiceInitializing = false
+	ongoingCheckpointSaves = new Map<string, Promise<void | CheckpointResult | undefined>>()
 
 	// Task Bridge
 	enableTaskBridge: boolean
@@ -1397,6 +1399,13 @@ export class Task extends EventEmitter<TaskEvents> 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)