diff --git a/src/features/background-agent/task-poller.test.ts b/src/features/background-agent/task-poller.test.ts index 6da3ffb0e..6716ec694 100644 --- a/src/features/background-agent/task-poller.test.ts +++ b/src/features/background-agent/task-poller.test.ts @@ -420,6 +420,21 @@ describe("checkAndInterruptStaleTasks", () => { }) describe("pruneStaleTasksAndNotifications", () => { + function createTerminalTask(overrides: Partial = {}): 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() @@ -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() - 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([[task.id, task]]) + const notifications = new Map([[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) }) }) diff --git a/src/features/background-agent/task-poller.ts b/src/features/background-agent/task-poller.ts index bc4f049a4..a62d3b442 100644 --- a/src/features/background-agent/task-poller.ts +++ b/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([ "completed", "error", @@ -27,9 +29,28 @@ export function pruneStaleTasksAndNotifications(args: { }): void { const { tasks, notifications, onTaskPruned } = args const now = Date.now() + const tasksWithPendingNotifications = new Set() + + 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()