fix(background-agent): preserve terminal tasks until notification cleanup

src/features/background-agent/manager.test.ts

@@ -224,6 +224,12 @@ function stubNotifyParentSession(manager: BackgroundManager): void {
   ;(manager as unknown as { notifyParentSession: () => Promise<void> }).notifyParentSession = async () => {}
 }
 
+async function flushBackgroundNotifications(): Promise<void> {
+  for (let i = 0; i < 6; i++) {
+    await Promise.resolve()
+  }
+}
+
 function createToastRemoveTaskTracker(): { removeTaskCalls: string[]; resetToastManager: () => void } {
   _resetTaskToastManagerForTesting()
   const toastManager = initTaskToastManager({
@@ -1306,11 +1312,20 @@ describe("BackgroundManager.tryCompleteTask", () => {
     expect(abortedSessionIDs).toEqual(["session-1"])
   })
 
-  test("should clean pendingByParent even when notifyParentSession throws", async () => {
+  test("should clean pendingByParent even when promptAsync notification fails", async () => {
     // given
-    ;(manager as unknown as { notifyParentSession: () => Promise<void> }).notifyParentSession = async () => {
-      throw new Error("notify failed")
+    const client = {
+      session: {
+        prompt: async () => ({}),
+        promptAsync: async () => {
+          throw new Error("notify failed")
+        },
+        abort: async () => ({}),
+        messages: async () => ({ data: [] }),
+      },
     }
+    manager.shutdown()
+    manager = new BackgroundManager({ client, directory: tmpdir() } as unknown as PluginInput)
 
     const task: BackgroundTask = {
       id: "task-pending-cleanup",
@@ -1424,7 +1439,7 @@ describe("BackgroundManager.tryCompleteTask", () => {
     // then
     expect(rejectedCount).toBe(0)
     expect(promptBodies.length).toBe(2)
-    expect(promptBodies.some((b) => b.noReply === false)).toBe(true)
+    expect(promptBodies.filter((body) => body.noReply === false)).toHaveLength(1)
   })
 })
 
@@ -1932,7 +1947,6 @@ describe("BackgroundManager - Non-blocking Queue Integration", () => {
     test("should cancel running task and release concurrency", async () => {
       // given
       const manager = createBackgroundManager()
-      stubNotifyParentSession(manager)
 
       const concurrencyManager = getConcurrencyManager(manager)
       const concurrencyKey = "test-provider/test-model"
@@ -2890,7 +2904,7 @@ describe("BackgroundManager.shutdown session abort", () => {
 })
 
 describe("BackgroundManager.handleEvent - session.deleted cascade", () => {
-  test("should cancel descendant tasks when parent session is deleted", () => {
+  test("should cancel descendant tasks and keep them until delayed cleanup", async () => {
     // given
     const manager = createBackgroundManager()
     const parentSessionID = "session-parent"
@@ -2937,21 +2951,26 @@ describe("BackgroundManager.handleEvent - session.deleted cascade", () => {
       properties: { info: { id: parentSessionID } },
     })
 
+    await flushBackgroundNotifications()
+
     // then
-    expect(taskMap.has(childTask.id)).toBe(false)
-    expect(taskMap.has(siblingTask.id)).toBe(false)
-    expect(taskMap.has(grandchildTask.id)).toBe(false)
+    expect(taskMap.has(childTask.id)).toBe(true)
+    expect(taskMap.has(siblingTask.id)).toBe(true)
+    expect(taskMap.has(grandchildTask.id)).toBe(true)
     expect(taskMap.has(unrelatedTask.id)).toBe(true)
     expect(childTask.status).toBe("cancelled")
     expect(siblingTask.status).toBe("cancelled")
     expect(grandchildTask.status).toBe("cancelled")
     expect(pendingByParent.get(parentSessionID)).toBeUndefined()
     expect(pendingByParent.get("session-child")).toBeUndefined()
+    expect(getCompletionTimers(manager).has(childTask.id)).toBe(true)
+    expect(getCompletionTimers(manager).has(siblingTask.id)).toBe(true)
+    expect(getCompletionTimers(manager).has(grandchildTask.id)).toBe(true)
 
     manager.shutdown()
   })
 
-  test("should remove tasks from toast manager when session is deleted", () => {
+  test("should remove cancelled tasks from toast manager while preserving delayed cleanup", async () => {
     //#given
     const { removeTaskCalls, resetToastManager } = createToastRemoveTaskTracker()
     const manager = createBackgroundManager()
@@ -2980,9 +2999,13 @@ describe("BackgroundManager.handleEvent - session.deleted cascade", () => {
       properties: { info: { id: parentSessionID } },
     })
 
+    await flushBackgroundNotifications()
+
     //#then
     expect(removeTaskCalls).toContain(childTask.id)
     expect(removeTaskCalls).toContain(grandchildTask.id)
+    expect(getCompletionTimers(manager).has(childTask.id)).toBe(true)
+    expect(getCompletionTimers(manager).has(grandchildTask.id)).toBe(true)
 
     manager.shutdown()
     resetToastManager()
@@ -3045,7 +3068,7 @@ describe("BackgroundManager.handleEvent - session.error", () => {
     return task
   }
 
-  test("sets task to error, releases concurrency, and cleans up", async () => {
+  test("sets task to error, releases concurrency, and keeps it until delayed cleanup", async () => {
     //#given
     const manager = createBackgroundManager()
     const concurrencyManager = getConcurrencyManager(manager)
@@ -3078,18 +3101,21 @@ describe("BackgroundManager.handleEvent - session.error", () => {
       },
     })
 
+    await flushBackgroundNotifications()
+
     //#then
     expect(task.status).toBe("error")
     expect(task.error).toBe("Model not found: kimi-for-coding/k2p5.")
     expect(task.completedAt).toBeInstanceOf(Date)
     expect(concurrencyManager.getCount(concurrencyKey)).toBe(0)
-    expect(getTaskMap(manager).has(task.id)).toBe(false)
+    expect(getTaskMap(manager).has(task.id)).toBe(true)
     expect(getPendingByParent(manager).get(task.parentSessionID)).toBeUndefined()
+    expect(getCompletionTimers(manager).has(task.id)).toBe(true)
 
     manager.shutdown()
   })
 
-  test("removes errored task from toast manager", () => {
+  test("should remove errored task from toast manager while preserving delayed cleanup", async () => {
     //#given
     const { removeTaskCalls, resetToastManager } = createToastRemoveTaskTracker()
     const manager = createBackgroundManager()
@@ -3111,8 +3137,11 @@ describe("BackgroundManager.handleEvent - session.error", () => {
       },
     })
 
+    await flushBackgroundNotifications()
+
     //#then
     expect(removeTaskCalls).toContain(task.id)
+    expect(getCompletionTimers(manager).has(task.id)).toBe(true)
 
     manager.shutdown()
     resetToastManager()
@@ -3518,7 +3547,7 @@ describe("BackgroundManager.completionTimers - Memory Leak Fix", () => {
     expect(completionTimers.size).toBe(0)
   })
 
-  test("should cancel timer when task is deleted via session.deleted", () => {
+  test("should preserve cleanup timer when terminal task session is deleted", () => {
     // given
     const manager = createBackgroundManager()
     const task: BackgroundTask = {
@@ -3547,7 +3576,7 @@ describe("BackgroundManager.completionTimers - Memory Leak Fix", () => {
     })
 
     // then
-    expect(completionTimers.has(task.id)).toBe(false)
+    expect(completionTimers.has(task.id)).toBe(true)
 
     manager.shutdown()
   })

src/features/background-agent/manager.ts

@@ -390,7 +390,6 @@ export class BackgroundManager {
         }).catch(() => {})
 
         this.markForNotification(existingTask)
-        this.cleanupPendingByParent(existingTask)
         this.enqueueNotificationForParent(existingTask.parentSessionID, () => this.notifyParentSession(existingTask)).catch(err => {
           log("[background-agent] Failed to notify on error:", err)
         })
@@ -661,7 +660,6 @@ export class BackgroundManager {
       }
 
       this.markForNotification(existingTask)
-      this.cleanupPendingByParent(existingTask)
       this.enqueueNotificationForParent(existingTask.parentSessionID, () => this.notifyParentSession(existingTask)).catch(err => {
         log("[background-agent] Failed to notify on resume error:", err)
       })
@@ -804,16 +802,14 @@ export class BackgroundManager {
         this.idleDeferralTimers.delete(task.id)
       }
 
-      this.cleanupPendingByParent(task)
-      this.tasks.delete(task.id)
-      this.clearNotificationsForTask(task.id)
-      const toastManager = getTaskToastManager()
-      if (toastManager) {
-        toastManager.removeTask(task.id)
-      }
       if (task.sessionID) {
-        subagentSessions.delete(task.sessionID)
+        SessionCategoryRegistry.remove(task.sessionID)
       }
+
+      this.markForNotification(task)
+      this.enqueueNotificationForParent(task.parentSessionID, () => this.notifyParentSession(task)).catch(err => {
+        log("[background-agent] Error in notifyParentSession for errored task:", { taskId: task.id, error: err })
+      })
     }
 
     if (event.type === "session.deleted") {
@@ -834,47 +830,30 @@ export class BackgroundManager {
 
       if (tasksToCancel.size === 0) return
 
+      const deletedSessionIDs = new Set<string>([sessionID])
+      for (const task of tasksToCancel.values()) {
+        if (task.sessionID) {
+          deletedSessionIDs.add(task.sessionID)
+        }
+      }
+
       for (const task of tasksToCancel.values()) {
         if (task.status === "running" || task.status === "pending") {
           void this.cancelTask(task.id, {
             source: "session.deleted",
             reason: "Session deleted",
-            skipNotification: true,
+          }).then(() => {
+            if (deletedSessionIDs.has(task.parentSessionID)) {
+              this.pendingNotifications.delete(task.parentSessionID)
+            }
           }).catch(err => {
+            if (deletedSessionIDs.has(task.parentSessionID)) {
+              this.pendingNotifications.delete(task.parentSessionID)
+            }
             log("[background-agent] Failed to cancel task on session.deleted:", { taskId: task.id, error: err })
           })
         }
-
-        const existingTimer = this.completionTimers.get(task.id)
-        if (existingTimer) {
-          clearTimeout(existingTimer)
-          this.completionTimers.delete(task.id)
-        }
-
-        const idleTimer = this.idleDeferralTimers.get(task.id)
-        if (idleTimer) {
-          clearTimeout(idleTimer)
-          this.idleDeferralTimers.delete(task.id)
-        }
-
-        this.cleanupPendingByParent(task)
-        this.tasks.delete(task.id)
-        this.clearNotificationsForTask(task.id)
-        const toastManager = getTaskToastManager()
-        if (toastManager) {
-          toastManager.removeTask(task.id)
-        }
-        if (task.sessionID) {
-          subagentSessions.delete(task.sessionID)
-        }
-      }
-
-      for (const task of tasksToCancel.values()) {
-        if (task.parentSessionID) {
-          this.pendingNotifications.delete(task.parentSessionID)
-        }
       }
-
       SessionCategoryRegistry.remove(sessionID)
     }
 
@@ -1094,8 +1073,6 @@ export class BackgroundManager {
       this.idleDeferralTimers.delete(task.id)
     }
 
-    this.cleanupPendingByParent(task)
-
     if (abortSession && task.sessionID) {
       this.client.session.abort({
         path: { id: task.sessionID },
@@ -1202,9 +1179,6 @@ export class BackgroundManager {
 
     this.markForNotification(task)
 
-    // Ensure pending tracking is cleaned up even if notification fails
-    this.cleanupPendingByParent(task)
-
     const idleTimer = this.idleDeferralTimers.get(task.id)
     if (idleTimer) {
       clearTimeout(idleTimer)
@@ -1260,15 +1234,24 @@ export class BackgroundManager {
         this.pendingByParent.delete(task.parentSessionID)
       }
     } else {
-      allComplete = true
+      remainingCount = Array.from(this.tasks.values())
+        .filter(t => t.parentSessionID === task.parentSessionID && t.id !== task.id && (t.status === "running" || t.status === "pending"))
+        .length
+      allComplete = remainingCount === 0
     }
 
     const completedTasks = allComplete
       ? Array.from(this.tasks.values())
         .filter(t => t.parentSessionID === task.parentSessionID && t.status !== "running" && t.status !== "pending")
       : []
 
-    const statusText = task.status === "completed" ? "COMPLETED" : task.status === "interrupt" ? "INTERRUPTED" : "CANCELLED"
+    const statusText = task.status === "completed"
+      ? "COMPLETED"
+      : task.status === "interrupt"
+        ? "INTERRUPTED"
+        : task.status === "error"
+          ? "ERROR"
+          : "CANCELLED"
     const errorInfo = task.error ? `\n**Error:** ${task.error}` : ""
 
     let notification: string

src/features/background-agent/task-poller.test.ts

@@ -422,4 +422,38 @@ describe("pruneStaleTasksAndNotifications", () => {
     //#then
     expect(pruned).toContain("old-task")
   })
+
+  it("should skip terminal tasks even when they exceeded TTL", () => {
+    //#given
+    const tasks = new Map<string, BackgroundTask>()
+    const oldStartedAt = new Date(Date.now() - 31 * 60 * 1000)
+    const terminalStatuses: BackgroundTask["status"][] = ["completed", "error", "cancelled", "interrupt"]
+
+    for (const status of terminalStatuses) {
+      tasks.set(status, {
+        id: status,
+        parentSessionID: "parent",
+        parentMessageID: "msg",
+        description: status,
+        prompt: status,
+        agent: "explore",
+        status,
+        startedAt: oldStartedAt,
+        completedAt: new Date(),
+      })
+    }
+
+    const pruned: string[] = []
+
+    //#when
+    pruneStaleTasksAndNotifications({
+      tasks,
+      notifications: new Map<string, BackgroundTask[]>(),
+      onTaskPruned: (taskId) => pruned.push(taskId),
+    })
+
+    //#then
+    expect(pruned).toEqual([])
+    expect(Array.from(tasks.keys())).toEqual(terminalStatuses)
+  })
 })

src/features/background-agent/task-poller.ts

@@ -12,6 +12,13 @@ import {
   TASK_TTL_MS,
 } from "./constants"
 
+const TERMINAL_TASK_STATUSES = new Set<BackgroundTask["status"]>([
+  "completed",
+  "error",
+  "cancelled",
+  "interrupt",
+])
+
 export function pruneStaleTasksAndNotifications(args: {
   tasks: Map<string, BackgroundTask>
   notifications: Map<string, BackgroundTask[]>
@@ -21,6 +28,8 @@ export function pruneStaleTasksAndNotifications(args: {
   const now = Date.now()
 
   for (const [taskId, task] of tasks.entries()) {
+    if (TERMINAL_TASK_STATUSES.has(task.status)) continue
+
     const timestamp = task.status === "pending"
       ? task.queuedAt?.getTime()
       : task.startedAt?.getTime()