diff --git a/src/features/background-agent/manager.ts b/src/features/background-agent/manager.ts index f418580c2..cb563d0aa 100644 --- a/src/features/background-agent/manager.ts +++ b/src/features/background-agent/manager.ts @@ -138,6 +138,7 @@ export class BackgroundManager { private tasks: Map + private tasksByParentSession: Map> private notifications: Map private pendingNotifications: Map private pendingByParent: Map> // Track pending tasks per parent for batching @@ -177,6 +178,7 @@ export class BackgroundManager { } ) { this.tasks = new Map() + this.tasksByParentSession = new Map() this.notifications = new Map() this.pendingNotifications = new Map() this.pendingByParent = new Map() @@ -280,6 +282,50 @@ export class BackgroundManager { this.unregisterRootDescendant(task.rootSessionID) } + private addTask(task: BackgroundTask): void { + this.tasks.set(task.id, task) + if (!task.parentSessionID) { + return + } + + const taskIDs = this.tasksByParentSession.get(task.parentSessionID) ?? new Set() + taskIDs.add(task.id) + this.tasksByParentSession.set(task.parentSessionID, taskIDs) + } + + private removeTask(task: BackgroundTask): void { + this.tasks.delete(task.id) + this.removeTaskFromParentIndex(task.id, task.parentSessionID) + } + + private updateTaskParent(task: BackgroundTask, parentSessionID: string): void { + if (task.parentSessionID === parentSessionID) { + return + } + + this.removeTaskFromParentIndex(task.id, task.parentSessionID) + task.parentSessionID = parentSessionID + const taskIDs = this.tasksByParentSession.get(parentSessionID) ?? new Set() + taskIDs.add(task.id) + this.tasksByParentSession.set(parentSessionID, taskIDs) + } + + private removeTaskFromParentIndex(taskID: string, parentSessionID: string | undefined): void { + if (!parentSessionID) { + return + } + + const taskIDs = this.tasksByParentSession.get(parentSessionID) + if (!taskIDs) { + return + } + + taskIDs.delete(taskID) + if (taskIDs.size === 0) { + this.tasksByParentSession.delete(parentSessionID) + } + } + async launch(input: LaunchInput): Promise { log("[background-agent] launch() called with:", { agent: input.agent, @@ -325,7 +371,7 @@ export class BackgroundManager { category: input.category, } - this.tasks.set(task.id, task) + this.addTask(task) this.taskHistory.record(input.parentSessionID, { id: task.id, agent: input.agent, description: input.description, status: "pending", category: input.category }) // Track for batched notifications immediately (pending state) @@ -636,13 +682,25 @@ export class BackgroundManager { } getTasksByParentSession(sessionID: string): BackgroundTask[] { - const result: BackgroundTask[] = [] - for (const task of this.tasks.values()) { - if (task.parentSessionID === sessionID) { - result.push(task) + const taskIDs = this.tasksByParentSession.get(sessionID) + if (!taskIDs) { + const result: BackgroundTask[] = [] + for (const task of this.tasks.values()) { + if (task.parentSessionID === sessionID) { + result.push(task) + } + } + return result + } + + const tasks: BackgroundTask[] = [] + for (const taskID of taskIDs) { + const task = this.tasks.get(taskID) + if (task) { + tasks.push(task) } } - return result + return tasks } getAllDescendantTasks(sessionID: string): BackgroundTask[] { @@ -696,7 +754,7 @@ export class BackgroundManager { const parentChanged = input.parentSessionID !== existingTask.parentSessionID if (parentChanged) { this.cleanupPendingByParent(existingTask) // Clean from OLD parent - existingTask.parentSessionID = input.parentSessionID + this.updateTaskParent(existingTask, input.parentSessionID) } if (input.parentAgent !== undefined) { existingTask.parentAgent = input.parentAgent @@ -751,7 +809,7 @@ export class BackgroundManager { concurrencyGroup, } - this.tasks.set(task.id, task) + this.addTask(task) subagentSessions.add(input.sessionID) this.startPolling() this.taskHistory.record(input.parentSessionID, { id: task.id, sessionID: input.sessionID, agent: input.agent || "task", description: input.description, status: "running", startedAt: task.startedAt }) @@ -801,7 +859,7 @@ export class BackgroundManager { existingTask.status = "running" existingTask.completedAt = undefined existingTask.error = undefined - existingTask.parentSessionID = input.parentSessionID + this.updateTaskParent(existingTask, input.parentSessionID) existingTask.parentMessageID = input.parentMessageID existingTask.parentModel = input.parentModel existingTask.parentAgent = input.parentAgent @@ -1493,7 +1551,7 @@ export class BackgroundManager { } this.clearNotificationsForTask(taskId) - this.tasks.delete(taskId) + this.removeTask(task) this.clearTaskHistoryWhenParentTasksGone(task.parentSessionID) if (task.sessionID) { subagentSessions.delete(task.sessionID) @@ -2161,6 +2219,7 @@ export class BackgroundManager { this.concurrencyManager.clear() this.tasks.clear() + this.tasksByParentSession.clear() this.notifications.clear() this.pendingNotifications.clear() this.pendingByParent.clear()