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
This commit is contained in:
Shawn 2025-09-01 11:33:44 -04:00 committed by Hannes Rudolph
parent 8715471bb5
commit 1b199674fa
3 changed files with 16 additions and 9 deletions

View file

@ -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" }]
: [],

View file

@ -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
}

View file

@ -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)