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)
This commit is contained in:
@@ -1,8 +1,11 @@
|
|||||||
import { describe, test, expect, beforeEach } from "bun:test"
|
import { describe, test, expect, beforeEach } from "bun:test"
|
||||||
import type { BackgroundTask } from "./types"
|
import type { BackgroundTask } from "./types"
|
||||||
|
|
||||||
|
const TASK_TTL_MS = 30 * 60 * 1000
|
||||||
|
|
||||||
class MockBackgroundManager {
|
class MockBackgroundManager {
|
||||||
private tasks: Map<string, BackgroundTask> = new Map()
|
private tasks: Map<string, BackgroundTask> = new Map()
|
||||||
|
private notifications: Map<string, BackgroundTask[]> = new Map()
|
||||||
|
|
||||||
addTask(task: BackgroundTask): void {
|
addTask(task: BackgroundTask): void {
|
||||||
this.tasks.set(task.id, task)
|
this.tasks.set(task.id, task)
|
||||||
@@ -34,6 +37,74 @@ class MockBackgroundManager {
|
|||||||
|
|
||||||
return result
|
return result
|
||||||
}
|
}
|
||||||
|
|
||||||
|
markForNotification(task: BackgroundTask): void {
|
||||||
|
const queue = this.notifications.get(task.parentSessionID) ?? []
|
||||||
|
queue.push(task)
|
||||||
|
this.notifications.set(task.parentSessionID, queue)
|
||||||
|
}
|
||||||
|
|
||||||
|
getPendingNotifications(sessionID: string): BackgroundTask[] {
|
||||||
|
return this.notifications.get(sessionID) ?? []
|
||||||
|
}
|
||||||
|
|
||||||
|
private clearNotificationsForTask(taskId: string): void {
|
||||||
|
for (const [sessionID, tasks] of this.notifications.entries()) {
|
||||||
|
const filtered = tasks.filter((t) => t.id !== taskId)
|
||||||
|
if (filtered.length === 0) {
|
||||||
|
this.notifications.delete(sessionID)
|
||||||
|
} else {
|
||||||
|
this.notifications.set(sessionID, filtered)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pruneStaleTasksAndNotifications(): { prunedTasks: string[]; prunedNotifications: number } {
|
||||||
|
const now = Date.now()
|
||||||
|
const prunedTasks: string[] = []
|
||||||
|
let prunedNotifications = 0
|
||||||
|
|
||||||
|
for (const [taskId, task] of this.tasks.entries()) {
|
||||||
|
const age = now - task.startedAt.getTime()
|
||||||
|
if (age > TASK_TTL_MS) {
|
||||||
|
prunedTasks.push(taskId)
|
||||||
|
this.clearNotificationsForTask(taskId)
|
||||||
|
this.tasks.delete(taskId)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
for (const [sessionID, notifications] of this.notifications.entries()) {
|
||||||
|
if (notifications.length === 0) {
|
||||||
|
this.notifications.delete(sessionID)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
const validNotifications = notifications.filter((task) => {
|
||||||
|
const age = now - task.startedAt.getTime()
|
||||||
|
return age <= TASK_TTL_MS
|
||||||
|
})
|
||||||
|
const removed = notifications.length - validNotifications.length
|
||||||
|
prunedNotifications += removed
|
||||||
|
if (validNotifications.length === 0) {
|
||||||
|
this.notifications.delete(sessionID)
|
||||||
|
} else if (validNotifications.length !== notifications.length) {
|
||||||
|
this.notifications.set(sessionID, validNotifications)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return { prunedTasks, prunedNotifications }
|
||||||
|
}
|
||||||
|
|
||||||
|
getTaskCount(): number {
|
||||||
|
return this.tasks.size
|
||||||
|
}
|
||||||
|
|
||||||
|
getNotificationCount(): number {
|
||||||
|
let count = 0
|
||||||
|
for (const notifications of this.notifications.values()) {
|
||||||
|
count += notifications.length
|
||||||
|
}
|
||||||
|
return count
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
function createMockTask(overrides: Partial<BackgroundTask> & { id: string; sessionID: string; parentSessionID: string }): BackgroundTask {
|
function createMockTask(overrides: Partial<BackgroundTask> & { id: string; sessionID: string; parentSessionID: string }): BackgroundTask {
|
||||||
@@ -230,3 +301,116 @@ describe("BackgroundManager.getAllDescendantTasks", () => {
|
|||||||
expect(result[0].id).toBe("task-b")
|
expect(result[0].id).toBe("task-b")
|
||||||
})
|
})
|
||||||
})
|
})
|
||||||
|
|
||||||
|
describe("BackgroundManager.pruneStaleTasksAndNotifications", () => {
|
||||||
|
let manager: MockBackgroundManager
|
||||||
|
|
||||||
|
beforeEach(() => {
|
||||||
|
// #given
|
||||||
|
manager = new MockBackgroundManager()
|
||||||
|
})
|
||||||
|
|
||||||
|
test("should not prune fresh tasks", () => {
|
||||||
|
// #given
|
||||||
|
const task = createMockTask({
|
||||||
|
id: "task-fresh",
|
||||||
|
sessionID: "session-fresh",
|
||||||
|
parentSessionID: "session-parent",
|
||||||
|
startedAt: new Date(),
|
||||||
|
})
|
||||||
|
manager.addTask(task)
|
||||||
|
|
||||||
|
// #when
|
||||||
|
const result = manager.pruneStaleTasksAndNotifications()
|
||||||
|
|
||||||
|
// #then
|
||||||
|
expect(result.prunedTasks).toHaveLength(0)
|
||||||
|
expect(manager.getTaskCount()).toBe(1)
|
||||||
|
})
|
||||||
|
|
||||||
|
test("should prune tasks older than 30 minutes", () => {
|
||||||
|
// #given
|
||||||
|
const staleDate = new Date(Date.now() - 31 * 60 * 1000)
|
||||||
|
const task = createMockTask({
|
||||||
|
id: "task-stale",
|
||||||
|
sessionID: "session-stale",
|
||||||
|
parentSessionID: "session-parent",
|
||||||
|
startedAt: staleDate,
|
||||||
|
})
|
||||||
|
manager.addTask(task)
|
||||||
|
|
||||||
|
// #when
|
||||||
|
const result = manager.pruneStaleTasksAndNotifications()
|
||||||
|
|
||||||
|
// #then
|
||||||
|
expect(result.prunedTasks).toContain("task-stale")
|
||||||
|
expect(manager.getTaskCount()).toBe(0)
|
||||||
|
})
|
||||||
|
|
||||||
|
test("should prune stale notifications", () => {
|
||||||
|
// #given
|
||||||
|
const staleDate = new Date(Date.now() - 31 * 60 * 1000)
|
||||||
|
const task = createMockTask({
|
||||||
|
id: "task-stale",
|
||||||
|
sessionID: "session-stale",
|
||||||
|
parentSessionID: "session-parent",
|
||||||
|
startedAt: staleDate,
|
||||||
|
})
|
||||||
|
manager.markForNotification(task)
|
||||||
|
|
||||||
|
// #when
|
||||||
|
const result = manager.pruneStaleTasksAndNotifications()
|
||||||
|
|
||||||
|
// #then
|
||||||
|
expect(result.prunedNotifications).toBe(1)
|
||||||
|
expect(manager.getNotificationCount()).toBe(0)
|
||||||
|
})
|
||||||
|
|
||||||
|
test("should clean up notifications when task is pruned", () => {
|
||||||
|
// #given
|
||||||
|
const staleDate = new Date(Date.now() - 31 * 60 * 1000)
|
||||||
|
const task = createMockTask({
|
||||||
|
id: "task-stale",
|
||||||
|
sessionID: "session-stale",
|
||||||
|
parentSessionID: "session-parent",
|
||||||
|
startedAt: staleDate,
|
||||||
|
})
|
||||||
|
manager.addTask(task)
|
||||||
|
manager.markForNotification(task)
|
||||||
|
|
||||||
|
// #when
|
||||||
|
manager.pruneStaleTasksAndNotifications()
|
||||||
|
|
||||||
|
// #then
|
||||||
|
expect(manager.getTaskCount()).toBe(0)
|
||||||
|
expect(manager.getNotificationCount()).toBe(0)
|
||||||
|
})
|
||||||
|
|
||||||
|
test("should keep fresh tasks while pruning stale ones", () => {
|
||||||
|
// #given
|
||||||
|
const staleDate = new Date(Date.now() - 31 * 60 * 1000)
|
||||||
|
const staleTask = createMockTask({
|
||||||
|
id: "task-stale",
|
||||||
|
sessionID: "session-stale",
|
||||||
|
parentSessionID: "session-parent",
|
||||||
|
startedAt: staleDate,
|
||||||
|
})
|
||||||
|
const freshTask = createMockTask({
|
||||||
|
id: "task-fresh",
|
||||||
|
sessionID: "session-fresh",
|
||||||
|
parentSessionID: "session-parent",
|
||||||
|
startedAt: new Date(),
|
||||||
|
})
|
||||||
|
manager.addTask(staleTask)
|
||||||
|
manager.addTask(freshTask)
|
||||||
|
|
||||||
|
// #when
|
||||||
|
const result = manager.pruneStaleTasksAndNotifications()
|
||||||
|
|
||||||
|
// #then
|
||||||
|
expect(result.prunedTasks).toHaveLength(1)
|
||||||
|
expect(result.prunedTasks).toContain("task-stale")
|
||||||
|
expect(manager.getTaskCount()).toBe(1)
|
||||||
|
expect(manager.getTask("task-fresh")).toBeDefined()
|
||||||
|
})
|
||||||
|
})
|
||||||
|
|||||||
@@ -12,6 +12,8 @@ import {
|
|||||||
} from "../hook-message-injector"
|
} from "../hook-message-injector"
|
||||||
import { subagentSessions } from "../claude-code-session-state"
|
import { subagentSessions } from "../claude-code-session-state"
|
||||||
|
|
||||||
|
const TASK_TTL_MS = 30 * 60 * 1000
|
||||||
|
|
||||||
type OpencodeClient = PluginInput["client"]
|
type OpencodeClient = PluginInput["client"]
|
||||||
|
|
||||||
interface MessagePartInfo {
|
interface MessagePartInfo {
|
||||||
@@ -345,11 +347,12 @@ export class BackgroundManager {
|
|||||||
},
|
},
|
||||||
query: { directory: this.directory },
|
query: { directory: this.directory },
|
||||||
})
|
})
|
||||||
this.clearNotificationsForTask(taskId)
|
|
||||||
log("[background-agent] Successfully sent prompt to parent session:", { parentSessionID: task.parentSessionID })
|
log("[background-agent] Successfully sent prompt to parent session:", { parentSessionID: task.parentSessionID })
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
log("[background-agent] prompt failed:", String(error))
|
log("[background-agent] prompt failed:", String(error))
|
||||||
} finally {
|
} finally {
|
||||||
|
// Always clean up both maps to prevent memory leaks
|
||||||
|
this.clearNotificationsForTask(taskId)
|
||||||
this.tasks.delete(taskId)
|
this.tasks.delete(taskId)
|
||||||
log("[background-agent] Removed completed task from memory:", taskId)
|
log("[background-agent] Removed completed task from memory:", taskId)
|
||||||
}
|
}
|
||||||
@@ -377,7 +380,42 @@ export class BackgroundManager {
|
|||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
|
||||||
|
private pruneStaleTasksAndNotifications(): void {
|
||||||
|
const now = Date.now()
|
||||||
|
|
||||||
|
for (const [taskId, task] of this.tasks.entries()) {
|
||||||
|
const age = now - task.startedAt.getTime()
|
||||||
|
if (age > TASK_TTL_MS) {
|
||||||
|
log("[background-agent] Pruning stale task:", { taskId, age: Math.round(age / 1000) + "s" })
|
||||||
|
task.status = "error"
|
||||||
|
task.error = "Task timed out after 30 minutes"
|
||||||
|
task.completedAt = new Date()
|
||||||
|
this.clearNotificationsForTask(taskId)
|
||||||
|
this.tasks.delete(taskId)
|
||||||
|
subagentSessions.delete(task.sessionID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
for (const [sessionID, notifications] of this.notifications.entries()) {
|
||||||
|
if (notifications.length === 0) {
|
||||||
|
this.notifications.delete(sessionID)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
const validNotifications = notifications.filter((task) => {
|
||||||
|
const age = now - task.startedAt.getTime()
|
||||||
|
return age <= TASK_TTL_MS
|
||||||
|
})
|
||||||
|
if (validNotifications.length === 0) {
|
||||||
|
this.notifications.delete(sessionID)
|
||||||
|
} else if (validNotifications.length !== notifications.length) {
|
||||||
|
this.notifications.set(sessionID, validNotifications)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
private async pollRunningTasks(): Promise<void> {
|
private async pollRunningTasks(): Promise<void> {
|
||||||
|
this.pruneStaleTasksAndNotifications()
|
||||||
|
|
||||||
const statusResult = await this.client.session.status()
|
const statusResult = await this.client.session.status()
|
||||||
const allStatuses = (statusResult.data ?? {}) as Record<string, { type: string }>
|
const allStatuses = (statusResult.data ?? {}) as Record<string, { type: string }>
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user