From 1b199674fa6bf3534aac51b31fd6ade7cb1b1edf 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 e143dc37f0..6149df9fba 100644
--- a/src/core/checkpoints/index.ts
+++ b/src/core/checkpoints/index.ts
@@ -370,24 +370,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(task: Task, force = false, files?: vscode.Uri[], suppressMessage = false) {
-	// 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 = `${task.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 (task.ongoingCheckpointSaves.has(saveKey)) {
 		const provider = task.providerRef.deref()
 		provider?.log(`[checkpointSave] duplicate checkpoint save detected for ${saveKey}, using existing operation`)
-		return ongoingCheckpointSaves.get(saveKey)
+		return task.ongoingCheckpointSaves.get(saveKey)
 	}
 	const service = await getInitializedCheckpointService(task)
 
@@ -432,10 +429,10 @@ export async function checkpointSave(task: Task, force = false, files?: vscode.U
 		})
 		.finally(() => {
 			// Clean up the tracking once completed
-			ongoingCheckpointSaves.delete(saveKey)
+			task.ongoingCheckpointSaves.delete(saveKey)
 		})
 
-	ongoingCheckpointSaves.set(saveKey, savePromise)
+	task.ongoingCheckpointSaves.set(saveKey, savePromise)
 	return savePromise
 }
 
diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts
index c5be865731..5e59f45938 100644
--- a/src/core/task/Task.ts
+++ b/src/core/task/Task.ts
@@ -63,6 +63,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"
@@ -268,6 +269,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
 	enableBridge: boolean
@@ -1525,6 +1527,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)