Files
oh-my-openagent/src/features/background-agent/manager.test.ts
YeonGyu-Kim d0694e5aa4 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)
2026-01-02 22:30:55 +09:00

417 lines
11 KiB
TypeScript

import { describe, test, expect, beforeEach } from "bun:test"
import type { BackgroundTask } from "./types"
const TASK_TTL_MS = 30 * 60 * 1000
class MockBackgroundManager {
private tasks: Map<string, BackgroundTask> = new Map()
private notifications: Map<string, BackgroundTask[]> = new Map()
addTask(task: BackgroundTask): void {
this.tasks.set(task.id, task)
}
getTask(id: string): BackgroundTask | undefined {
return this.tasks.get(id)
}
getTasksByParentSession(sessionID: string): BackgroundTask[] {
const result: BackgroundTask[] = []
for (const task of this.tasks.values()) {
if (task.parentSessionID === sessionID) {
result.push(task)
}
}
return result
}
getAllDescendantTasks(sessionID: string): BackgroundTask[] {
const result: BackgroundTask[] = []
const directChildren = this.getTasksByParentSession(sessionID)
for (const child of directChildren) {
result.push(child)
const descendants = this.getAllDescendantTasks(child.sessionID)
result.push(...descendants)
}
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 {
return {
parentMessageID: "mock-message-id",
description: "test task",
prompt: "test prompt",
agent: "test-agent",
status: "running",
startedAt: new Date(),
...overrides,
}
}
describe("BackgroundManager.getAllDescendantTasks", () => {
let manager: MockBackgroundManager
beforeEach(() => {
// #given
manager = new MockBackgroundManager()
})
test("should return empty array when no tasks exist", () => {
// #given - empty manager
// #when
const result = manager.getAllDescendantTasks("session-a")
// #then
expect(result).toEqual([])
})
test("should return direct children only when no nested tasks", () => {
// #given
const taskB = createMockTask({
id: "task-b",
sessionID: "session-b",
parentSessionID: "session-a",
})
manager.addTask(taskB)
// #when
const result = manager.getAllDescendantTasks("session-a")
// #then
expect(result).toHaveLength(1)
expect(result[0].id).toBe("task-b")
})
test("should return all nested descendants (2 levels deep)", () => {
// #given
// Session A -> Task B -> Task C
const taskB = createMockTask({
id: "task-b",
sessionID: "session-b",
parentSessionID: "session-a",
})
const taskC = createMockTask({
id: "task-c",
sessionID: "session-c",
parentSessionID: "session-b",
})
manager.addTask(taskB)
manager.addTask(taskC)
// #when
const result = manager.getAllDescendantTasks("session-a")
// #then
expect(result).toHaveLength(2)
expect(result.map(t => t.id)).toContain("task-b")
expect(result.map(t => t.id)).toContain("task-c")
})
test("should return all nested descendants (3 levels deep)", () => {
// #given
// Session A -> Task B -> Task C -> Task D
const taskB = createMockTask({
id: "task-b",
sessionID: "session-b",
parentSessionID: "session-a",
})
const taskC = createMockTask({
id: "task-c",
sessionID: "session-c",
parentSessionID: "session-b",
})
const taskD = createMockTask({
id: "task-d",
sessionID: "session-d",
parentSessionID: "session-c",
})
manager.addTask(taskB)
manager.addTask(taskC)
manager.addTask(taskD)
// #when
const result = manager.getAllDescendantTasks("session-a")
// #then
expect(result).toHaveLength(3)
expect(result.map(t => t.id)).toContain("task-b")
expect(result.map(t => t.id)).toContain("task-c")
expect(result.map(t => t.id)).toContain("task-d")
})
test("should handle multiple branches (tree structure)", () => {
// #given
// Session A -> Task B1 -> Task C1
// -> Task B2 -> Task C2
const taskB1 = createMockTask({
id: "task-b1",
sessionID: "session-b1",
parentSessionID: "session-a",
})
const taskB2 = createMockTask({
id: "task-b2",
sessionID: "session-b2",
parentSessionID: "session-a",
})
const taskC1 = createMockTask({
id: "task-c1",
sessionID: "session-c1",
parentSessionID: "session-b1",
})
const taskC2 = createMockTask({
id: "task-c2",
sessionID: "session-c2",
parentSessionID: "session-b2",
})
manager.addTask(taskB1)
manager.addTask(taskB2)
manager.addTask(taskC1)
manager.addTask(taskC2)
// #when
const result = manager.getAllDescendantTasks("session-a")
// #then
expect(result).toHaveLength(4)
expect(result.map(t => t.id)).toContain("task-b1")
expect(result.map(t => t.id)).toContain("task-b2")
expect(result.map(t => t.id)).toContain("task-c1")
expect(result.map(t => t.id)).toContain("task-c2")
})
test("should not include tasks from unrelated sessions", () => {
// #given
// Session A -> Task B
// Session X -> Task Y (unrelated)
const taskB = createMockTask({
id: "task-b",
sessionID: "session-b",
parentSessionID: "session-a",
})
const taskY = createMockTask({
id: "task-y",
sessionID: "session-y",
parentSessionID: "session-x",
})
manager.addTask(taskB)
manager.addTask(taskY)
// #when
const result = manager.getAllDescendantTasks("session-a")
// #then
expect(result).toHaveLength(1)
expect(result[0].id).toBe("task-b")
expect(result.map(t => t.id)).not.toContain("task-y")
})
test("getTasksByParentSession should only return direct children (not recursive)", () => {
// #given
// Session A -> Task B -> Task C
const taskB = createMockTask({
id: "task-b",
sessionID: "session-b",
parentSessionID: "session-a",
})
const taskC = createMockTask({
id: "task-c",
sessionID: "session-c",
parentSessionID: "session-b",
})
manager.addTask(taskB)
manager.addTask(taskC)
// #when
const result = manager.getTasksByParentSession("session-a")
// #then
expect(result).toHaveLength(1)
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()
})
})