Skip to content

Commit 1b19967

Browse files
playcationshannesrudolph
authored andcommitted
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
1 parent 8715471 commit 1b19967

File tree

3 files changed

+16
-9
lines changed

3 files changed

+16
-9
lines changed

src/core/checkpoints/__tests__/index.spec.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@ const createMockTask = (options: { taskId: string; hasExistingCheckpoints: boole
7979
enableCheckpoints: options.enableCheckpoints ?? true,
8080
checkpointService: null as any,
8181
checkpointServiceInitializing: false,
82+
ongoingCheckpointSaves: new Map(),
8283
clineMessages: options.hasExistingCheckpoints
8384
? [{ say: "checkpoint_saved", ts: Date.now(), text: "existing-checkpoint-hash" }]
8485
: [],

src/core/checkpoints/index.ts

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -370,24 +370,21 @@ export async function getInitializedCheckpointService(
370370
}
371371
}
372372

373-
// Track ongoing checkpoint saves per task to prevent duplicates
374-
const ongoingCheckpointSaves = new Map<string, Promise<void | CheckpointResult | undefined>>()
375-
376373
export async function checkpointSave(task: Task, force = false, files?: vscode.Uri[], suppressMessage = false) {
377-
// Create a unique key for this checkpoint save operation
374+
// Create a unique key for this checkpoint save operation (task-scoped, no need for taskId in key)
378375
const filesKey = files
379376
? files
380377
.map((f) => f.fsPath)
381378
.sort()
382379
.join("|")
383380
: "all"
384-
const saveKey = `${task.taskId}-${force}-${filesKey}`
381+
const saveKey = `${force}-${filesKey}`
385382

386383
// If there's already an ongoing checkpoint save for this exact operation, return the existing promise
387-
if (ongoingCheckpointSaves.has(saveKey)) {
384+
if (task.ongoingCheckpointSaves.has(saveKey)) {
388385
const provider = task.providerRef.deref()
389386
provider?.log(`[checkpointSave] duplicate checkpoint save detected for ${saveKey}, using existing operation`)
390-
return ongoingCheckpointSaves.get(saveKey)
387+
return task.ongoingCheckpointSaves.get(saveKey)
391388
}
392389
const service = await getInitializedCheckpointService(task)
393390

@@ -432,10 +429,10 @@ export async function checkpointSave(task: Task, force = false, files?: vscode.U
432429
})
433430
.finally(() => {
434431
// Clean up the tracking once completed
435-
ongoingCheckpointSaves.delete(saveKey)
432+
task.ongoingCheckpointSaves.delete(saveKey)
436433
})
437434

438-
ongoingCheckpointSaves.set(saveKey, savePromise)
435+
task.ongoingCheckpointSaves.set(saveKey, savePromise)
439436
return savePromise
440437
}
441438

src/core/task/Task.ts

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ import { BrowserSession } from "../../services/browser/BrowserSession"
6363
import { McpHub } from "../../services/mcp/McpHub"
6464
import { McpServerManager } from "../../services/mcp/McpServerManager"
6565
import { RepoPerTaskCheckpointService } from "../../services/checkpoints"
66+
import { CheckpointResult } from "../../services/checkpoints/types"
6667

6768
// integrations
6869
import { DiffViewProvider } from "../../integrations/editor/DiffViewProvider"
@@ -268,6 +269,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
268269
enableCheckpoints: boolean
269270
checkpointService?: RepoPerTaskCheckpointService
270271
checkpointServiceInitializing = false
272+
ongoingCheckpointSaves = new Map<string, Promise<void | CheckpointResult | undefined>>()
271273

272274
// Task Bridge
273275
enableBridge: boolean
@@ -1525,6 +1527,13 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
15251527
console.error("Error removing event listeners:", error)
15261528
}
15271529

1530+
// Clean up ongoing checkpoint saves to prevent memory leaks
1531+
try {
1532+
this.ongoingCheckpointSaves.clear()
1533+
} catch (error) {
1534+
console.error("Error clearing ongoing checkpoint saves:", error)
1535+
}
1536+
15281537
// Stop waiting for child task completion.
15291538
if (this.pauseInterval) {
15301539
clearInterval(this.pauseInterval)

0 commit comments

Comments
 (0)