Skip to content

Commit d0694e5

Browse files
committed
fix(background-agent): prevent memory leaks by cleaning notifications in finally block and add TTL-based task pruning
- Move clearNotificationsForTask() to finally block to ensure cleanup even on success - Add TASK_TTL_MS (30 min) constant for stale task detection - Implement pruneStaleTasksAndNotifications() to remove expired tasks and notifications - Add comprehensive tests for pruning functionality (fresh tasks, stale tasks, notifications) - Prevents indefinite Map growth when tasks complete without errors 🤖 GENERATED WITH ASSISTANCE OF [OhMyOpenCode](https://github.com/code-yeongyu/oh-my-opencode)
1 parent 4a9bdc8 commit d0694e5

File tree

2 files changed

+223
-1
lines changed

2 files changed

+223
-1
lines changed

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

Lines changed: 184 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
11
import { describe, test, expect, beforeEach } from "bun:test"
22
import type { BackgroundTask } from "./types"
33

4+
const TASK_TTL_MS = 30 * 60 * 1000
5+
46
class MockBackgroundManager {
57
private tasks: Map<string, BackgroundTask> = new Map()
8+
private notifications: Map<string, BackgroundTask[]> = new Map()
69

710
addTask(task: BackgroundTask): void {
811
this.tasks.set(task.id, task)
@@ -34,6 +37,74 @@ class MockBackgroundManager {
3437

3538
return result
3639
}
40+
41+
markForNotification(task: BackgroundTask): void {
42+
const queue = this.notifications.get(task.parentSessionID) ?? []
43+
queue.push(task)
44+
this.notifications.set(task.parentSessionID, queue)
45+
}
46+
47+
getPendingNotifications(sessionID: string): BackgroundTask[] {
48+
return this.notifications.get(sessionID) ?? []
49+
}
50+
51+
private clearNotificationsForTask(taskId: string): void {
52+
for (const [sessionID, tasks] of this.notifications.entries()) {
53+
const filtered = tasks.filter((t) => t.id !== taskId)
54+
if (filtered.length === 0) {
55+
this.notifications.delete(sessionID)
56+
} else {
57+
this.notifications.set(sessionID, filtered)
58+
}
59+
}
60+
}
61+
62+
pruneStaleTasksAndNotifications(): { prunedTasks: string[]; prunedNotifications: number } {
63+
const now = Date.now()
64+
const prunedTasks: string[] = []
65+
let prunedNotifications = 0
66+
67+
for (const [taskId, task] of this.tasks.entries()) {
68+
const age = now - task.startedAt.getTime()
69+
if (age > TASK_TTL_MS) {
70+
prunedTasks.push(taskId)
71+
this.clearNotificationsForTask(taskId)
72+
this.tasks.delete(taskId)
73+
}
74+
}
75+
76+
for (const [sessionID, notifications] of this.notifications.entries()) {
77+
if (notifications.length === 0) {
78+
this.notifications.delete(sessionID)
79+
continue
80+
}
81+
const validNotifications = notifications.filter((task) => {
82+
const age = now - task.startedAt.getTime()
83+
return age <= TASK_TTL_MS
84+
})
85+
const removed = notifications.length - validNotifications.length
86+
prunedNotifications += removed
87+
if (validNotifications.length === 0) {
88+
this.notifications.delete(sessionID)
89+
} else if (validNotifications.length !== notifications.length) {
90+
this.notifications.set(sessionID, validNotifications)
91+
}
92+
}
93+
94+
return { prunedTasks, prunedNotifications }
95+
}
96+
97+
getTaskCount(): number {
98+
return this.tasks.size
99+
}
100+
101+
getNotificationCount(): number {
102+
let count = 0
103+
for (const notifications of this.notifications.values()) {
104+
count += notifications.length
105+
}
106+
return count
107+
}
37108
}
38109

39110
function createMockTask(overrides: Partial<BackgroundTask> & { id: string; sessionID: string; parentSessionID: string }): BackgroundTask {
@@ -230,3 +301,116 @@ describe("BackgroundManager.getAllDescendantTasks", () => {
230301
expect(result[0].id).toBe("task-b")
231302
})
232303
})
304+
305+
describe("BackgroundManager.pruneStaleTasksAndNotifications", () => {
306+
let manager: MockBackgroundManager
307+
308+
beforeEach(() => {
309+
// #given
310+
manager = new MockBackgroundManager()
311+
})
312+
313+
test("should not prune fresh tasks", () => {
314+
// #given
315+
const task = createMockTask({
316+
id: "task-fresh",
317+
sessionID: "session-fresh",
318+
parentSessionID: "session-parent",
319+
startedAt: new Date(),
320+
})
321+
manager.addTask(task)
322+
323+
// #when
324+
const result = manager.pruneStaleTasksAndNotifications()
325+
326+
// #then
327+
expect(result.prunedTasks).toHaveLength(0)
328+
expect(manager.getTaskCount()).toBe(1)
329+
})
330+
331+
test("should prune tasks older than 30 minutes", () => {
332+
// #given
333+
const staleDate = new Date(Date.now() - 31 * 60 * 1000)
334+
const task = createMockTask({
335+
id: "task-stale",
336+
sessionID: "session-stale",
337+
parentSessionID: "session-parent",
338+
startedAt: staleDate,
339+
})
340+
manager.addTask(task)
341+
342+
// #when
343+
const result = manager.pruneStaleTasksAndNotifications()
344+
345+
// #then
346+
expect(result.prunedTasks).toContain("task-stale")
347+
expect(manager.getTaskCount()).toBe(0)
348+
})
349+
350+
test("should prune stale notifications", () => {
351+
// #given
352+
const staleDate = new Date(Date.now() - 31 * 60 * 1000)
353+
const task = createMockTask({
354+
id: "task-stale",
355+
sessionID: "session-stale",
356+
parentSessionID: "session-parent",
357+
startedAt: staleDate,
358+
})
359+
manager.markForNotification(task)
360+
361+
// #when
362+
const result = manager.pruneStaleTasksAndNotifications()
363+
364+
// #then
365+
expect(result.prunedNotifications).toBe(1)
366+
expect(manager.getNotificationCount()).toBe(0)
367+
})
368+
369+
test("should clean up notifications when task is pruned", () => {
370+
// #given
371+
const staleDate = new Date(Date.now() - 31 * 60 * 1000)
372+
const task = createMockTask({
373+
id: "task-stale",
374+
sessionID: "session-stale",
375+
parentSessionID: "session-parent",
376+
startedAt: staleDate,
377+
})
378+
manager.addTask(task)
379+
manager.markForNotification(task)
380+
381+
// #when
382+
manager.pruneStaleTasksAndNotifications()
383+
384+
// #then
385+
expect(manager.getTaskCount()).toBe(0)
386+
expect(manager.getNotificationCount()).toBe(0)
387+
})
388+
389+
test("should keep fresh tasks while pruning stale ones", () => {
390+
// #given
391+
const staleDate = new Date(Date.now() - 31 * 60 * 1000)
392+
const staleTask = createMockTask({
393+
id: "task-stale",
394+
sessionID: "session-stale",
395+
parentSessionID: "session-parent",
396+
startedAt: staleDate,
397+
})
398+
const freshTask = createMockTask({
399+
id: "task-fresh",
400+
sessionID: "session-fresh",
401+
parentSessionID: "session-parent",
402+
startedAt: new Date(),
403+
})
404+
manager.addTask(staleTask)
405+
manager.addTask(freshTask)
406+
407+
// #when
408+
const result = manager.pruneStaleTasksAndNotifications()
409+
410+
// #then
411+
expect(result.prunedTasks).toHaveLength(1)
412+
expect(result.prunedTasks).toContain("task-stale")
413+
expect(manager.getTaskCount()).toBe(1)
414+
expect(manager.getTask("task-fresh")).toBeDefined()
415+
})
416+
})

src/features/background-agent/manager.ts

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ import {
1212
} from "../hook-message-injector"
1313
import { subagentSessions } from "../claude-code-session-state"
1414

15+
const TASK_TTL_MS = 30 * 60 * 1000
16+
1517
type OpencodeClient = PluginInput["client"]
1618

1719
interface MessagePartInfo {
@@ -345,11 +347,12 @@ export class BackgroundManager {
345347
},
346348
query: { directory: this.directory },
347349
})
348-
this.clearNotificationsForTask(taskId)
349350
log("[background-agent] Successfully sent prompt to parent session:", { parentSessionID: task.parentSessionID })
350351
} catch (error) {
351352
log("[background-agent] prompt failed:", String(error))
352353
} finally {
354+
// Always clean up both maps to prevent memory leaks
355+
this.clearNotificationsForTask(taskId)
353356
this.tasks.delete(taskId)
354357
log("[background-agent] Removed completed task from memory:", taskId)
355358
}
@@ -377,7 +380,42 @@ export class BackgroundManager {
377380
return false
378381
}
379382

383+
private pruneStaleTasksAndNotifications(): void {
384+
const now = Date.now()
385+
386+
for (const [taskId, task] of this.tasks.entries()) {
387+
const age = now - task.startedAt.getTime()
388+
if (age > TASK_TTL_MS) {
389+
log("[background-agent] Pruning stale task:", { taskId, age: Math.round(age / 1000) + "s" })
390+
task.status = "error"
391+
task.error = "Task timed out after 30 minutes"
392+
task.completedAt = new Date()
393+
this.clearNotificationsForTask(taskId)
394+
this.tasks.delete(taskId)
395+
subagentSessions.delete(task.sessionID)
396+
}
397+
}
398+
399+
for (const [sessionID, notifications] of this.notifications.entries()) {
400+
if (notifications.length === 0) {
401+
this.notifications.delete(sessionID)
402+
continue
403+
}
404+
const validNotifications = notifications.filter((task) => {
405+
const age = now - task.startedAt.getTime()
406+
return age <= TASK_TTL_MS
407+
})
408+
if (validNotifications.length === 0) {
409+
this.notifications.delete(sessionID)
410+
} else if (validNotifications.length !== notifications.length) {
411+
this.notifications.set(sessionID, validNotifications)
412+
}
413+
}
414+
}
415+
380416
private async pollRunningTasks(): Promise<void> {
417+
this.pruneStaleTasksAndNotifications()
418+
381419
const statusResult = await this.client.session.status()
382420
const allStatuses = (statusResult.data ?? {}) as Record<string, { type: string }>
383421

0 commit comments

Comments
 (0)