Merge pull request #2470 from code-yeongyu/fix/terminal-task-retention-ttl

fix(background-agent): add TTL for terminal task retention

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

@@ -420,6 +420,21 @@ describe("checkAndInterruptStaleTasks", () => {
 })
 
 describe("pruneStaleTasksAndNotifications", () => {
+  function createTerminalTask(overrides: Partial<BackgroundTask> = {}): BackgroundTask {
+    return {
+      id: "terminal-task",
+      parentSessionID: "parent",
+      parentMessageID: "msg",
+      description: "terminal",
+      prompt: "terminal",
+      agent: "explore",
+      status: "completed",
+      startedAt: new Date(Date.now() - 40 * 60 * 1000),
+      completedAt: new Date(Date.now() - 31 * 60 * 1000),
+      ...overrides,
+    }
+  }
+
   it("should prune tasks that exceeded TTL", () => {
     //#given
     const tasks = new Map<string, BackgroundTask>()
@@ -449,24 +464,18 @@ describe("pruneStaleTasksAndNotifications", () => {
     expect(pruned).toContain("old-task")
   })
 
-  it("should skip terminal tasks even when they exceeded TTL", () => {
+  it("should prune terminal tasks when completion time exceeds terminal 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, {
+      tasks.set(status, createTerminalTask({
         id: status,
-        parentSessionID: "parent",
-        parentMessageID: "msg",
         description: status,
         prompt: status,
-        agent: "explore",
         status,
-        startedAt: oldStartedAt,
-        completedAt: new Date(),
-      })
+      }))
     }
 
     const pruned: string[] = []
@@ -480,6 +489,26 @@ describe("pruneStaleTasksAndNotifications", () => {
 
     //#then
     expect(pruned).toEqual([])
-    expect(Array.from(tasks.keys())).toEqual(terminalStatuses)
+    expect(Array.from(tasks.keys())).toEqual([])
+  })
+
+  it("should keep terminal tasks with pending notifications until notification cleanup", () => {
+    //#given
+    const task = createTerminalTask()
+    const tasks = new Map<string, BackgroundTask>([[task.id, task]])
+    const notifications = new Map<string, BackgroundTask[]>([[task.parentSessionID, [task]]])
+    const pruned: string[] = []
+
+    //#when
+    pruneStaleTasksAndNotifications({
+      tasks,
+      notifications,
+      onTaskPruned: (taskId) => pruned.push(taskId),
+    })
+
+    //#then
+    expect(pruned).toEqual([])
+    expect(tasks.has(task.id)).toBe(true)
+    expect(notifications.has(task.parentSessionID)).toBe(false)
   })
 })

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

@@ -13,6 +13,8 @@ import {
 } from "./constants"
 import { removeTaskToastTracking } from "./remove-task-toast-tracking"
 
+const TERMINAL_TASK_TTL_MS = 30 * 60 * 1000
+
 const TERMINAL_TASK_STATUSES = new Set<BackgroundTask["status"]>([
   "completed",
   "error",
@@ -27,9 +29,28 @@ export function pruneStaleTasksAndNotifications(args: {
 }): void {
   const { tasks, notifications, onTaskPruned } = args
   const now = Date.now()
+  const tasksWithPendingNotifications = new Set<string>()
+
+  for (const queued of notifications.values()) {
+    for (const task of queued) {
+      tasksWithPendingNotifications.add(task.id)
+    }
+  }
 
   for (const [taskId, task] of tasks.entries()) {
-    if (TERMINAL_TASK_STATUSES.has(task.status)) continue
+    if (TERMINAL_TASK_STATUSES.has(task.status)) {
+      if (tasksWithPendingNotifications.has(taskId)) continue
+
+      const completedAt = task.completedAt?.getTime()
+      if (!completedAt) continue
+
+      const age = now - completedAt
+      if (age <= TERMINAL_TASK_TTL_MS) continue
+
+      removeTaskToastTracking(taskId)
+      tasks.delete(taskId)
+      continue
+    }
 
     const timestamp = task.status === "pending"
       ? task.queuedAt?.getTime()