diff --git a/src/features/background-agent/manager.test.ts b/src/features/background-agent/manager.test.ts index 3cc7569d6..0ae1c2661 100644 --- a/src/features/background-agent/manager.test.ts +++ b/src/features/background-agent/manager.test.ts @@ -170,6 +170,7 @@ function createBackgroundManager(): BackgroundManager { const client = { session: { prompt: async () => ({}), + abort: async () => ({}), }, } return new BackgroundManager({ client, directory: tmpdir() } as unknown as PluginInput) @@ -1053,6 +1054,7 @@ describe("BackgroundManager.resume model persistence", () => { promptCalls.push(args) return {} }, + abort: async () => ({}), }, } manager = new BackgroundManager({ client, directory: tmpdir() } as unknown as PluginInput) @@ -1926,3 +1928,162 @@ describe("BackgroundManager.checkAndInterruptStaleTasks", () => { }) }) +describe("BackgroundManager.shutdown session abort", () => { + test("should call session.abort for all running tasks during shutdown", () => { + // #given + const abortedSessionIDs: string[] = [] + const client = { + session: { + prompt: async () => ({}), + abort: async (args: { path: { id: string } }) => { + abortedSessionIDs.push(args.path.id) + return {} + }, + }, + } + const manager = new BackgroundManager({ client, directory: tmpdir() } as unknown as PluginInput) + + const task1: BackgroundTask = { + id: "task-1", + sessionID: "session-1", + parentSessionID: "parent-1", + parentMessageID: "msg-1", + description: "Running task 1", + prompt: "Test", + agent: "test-agent", + status: "running", + startedAt: new Date(), + } + const task2: BackgroundTask = { + id: "task-2", + sessionID: "session-2", + parentSessionID: "parent-2", + parentMessageID: "msg-2", + description: "Running task 2", + prompt: "Test", + agent: "test-agent", + status: "running", + startedAt: new Date(), + } + + getTaskMap(manager).set(task1.id, task1) + getTaskMap(manager).set(task2.id, task2) + + // #when + manager.shutdown() + + // #then + expect(abortedSessionIDs).toContain("session-1") + expect(abortedSessionIDs).toContain("session-2") + expect(abortedSessionIDs).toHaveLength(2) + }) + + test("should not call session.abort for completed or cancelled tasks", () => { + // #given + const abortedSessionIDs: string[] = [] + const client = { + session: { + prompt: async () => ({}), + abort: async (args: { path: { id: string } }) => { + abortedSessionIDs.push(args.path.id) + return {} + }, + }, + } + const manager = new BackgroundManager({ client, directory: tmpdir() } as unknown as PluginInput) + + const completedTask: BackgroundTask = { + id: "task-completed", + sessionID: "session-completed", + parentSessionID: "parent-1", + parentMessageID: "msg-1", + description: "Completed task", + prompt: "Test", + agent: "test-agent", + status: "completed", + startedAt: new Date(), + completedAt: new Date(), + } + const cancelledTask: BackgroundTask = { + id: "task-cancelled", + sessionID: "session-cancelled", + parentSessionID: "parent-2", + parentMessageID: "msg-2", + description: "Cancelled task", + prompt: "Test", + agent: "test-agent", + status: "cancelled", + startedAt: new Date(), + completedAt: new Date(), + } + const pendingTask: BackgroundTask = { + id: "task-pending", + parentSessionID: "parent-3", + parentMessageID: "msg-3", + description: "Pending task", + prompt: "Test", + agent: "test-agent", + status: "pending", + queuedAt: new Date(), + } + + getTaskMap(manager).set(completedTask.id, completedTask) + getTaskMap(manager).set(cancelledTask.id, cancelledTask) + getTaskMap(manager).set(pendingTask.id, pendingTask) + + // #when + manager.shutdown() + + // #then + expect(abortedSessionIDs).toHaveLength(0) + }) + + test("should call onShutdown callback during shutdown", () => { + // #given + let shutdownCalled = false + const client = { + session: { + prompt: async () => ({}), + abort: async () => ({}), + }, + } + const manager = new BackgroundManager( + { client, directory: tmpdir() } as unknown as PluginInput, + undefined, + { + onShutdown: () => { + shutdownCalled = true + }, + } + ) + + // #when + manager.shutdown() + + // #then + expect(shutdownCalled).toBe(true) + }) + + test("should not throw when onShutdown callback throws", () => { + // #given + const client = { + session: { + prompt: async () => ({}), + abort: async () => ({}), + }, + } + const manager = new BackgroundManager( + { client, directory: tmpdir() } as unknown as PluginInput, + undefined, + { + onShutdown: () => { + throw new Error("cleanup failed") + }, + } + ) + + // #when / #then + expect(() => manager.shutdown()).not.toThrow() + }) +}) + diff --git a/src/features/background-agent/manager.ts b/src/features/background-agent/manager.ts index 4545050f1..ab564eba9 100644 --- a/src/features/background-agent/manager.ts +++ b/src/features/background-agent/manager.ts @@ -79,6 +79,7 @@ export class BackgroundManager { private config?: BackgroundTaskConfig private tmuxEnabled: boolean private onSubagentSessionCreated?: OnSubagentSessionCreated + private onShutdown?: () => void private queuesByKey: Map = new Map() private processingKeys: Set = new Set() @@ -89,6 +90,7 @@ export class BackgroundManager { options?: { tmuxConfig?: TmuxConfig onSubagentSessionCreated?: OnSubagentSessionCreated + onShutdown?: () => void } ) { this.tasks = new Map() @@ -100,6 +102,7 @@ export class BackgroundManager { this.config = config this.tmuxEnabled = options?.tmuxConfig?.enabled ?? false this.onSubagentSessionCreated = options?.onSubagentSessionCreated + this.onShutdown = options?.onShutdown this.registerProcessCleanup() } @@ -1346,7 +1349,25 @@ Use \`background_output(task_id="${task.id}")\` to retrieve this result when rea log("[background-agent] Shutting down BackgroundManager") this.stopPolling() - // Release concurrency for all running tasks first + // Abort all running sessions to prevent zombie processes (#1240) + for (const task of this.tasks.values()) { + if (task.status === "running" && task.sessionID) { + this.client.session.abort({ + path: { id: task.sessionID }, + }).catch(() => {}) + } + } + + // Notify shutdown listeners (e.g., tmux cleanup) + if (this.onShutdown) { + try { + this.onShutdown() + } catch (error) { + log("[background-agent] Error in onShutdown callback:", error) + } + } + + // Release concurrency for all running tasks for (const task of this.tasks.values()) { if (task.concurrencyKey) { this.concurrencyManager.release(task.concurrencyKey) diff --git a/src/hooks/interactive-bash-session/index.ts b/src/hooks/interactive-bash-session/index.ts index 324793d47..307441629 100644 --- a/src/hooks/interactive-bash-session/index.ts +++ b/src/hooks/interactive-bash-session/index.ts @@ -6,6 +6,7 @@ import { } from "./storage"; import { OMO_SESSION_PREFIX, buildSessionReminderMessage } from "./constants"; import type { InteractiveBashSessionState } from "./types"; +import { subagentSessions } from "../../features/claude-code-session-state"; interface ToolExecuteInput { tool: string; @@ -146,7 +147,7 @@ function findSubcommand(tokens: string[]): string { return "" } -export function createInteractiveBashSessionHook(_ctx: PluginInput) { +export function createInteractiveBashSessionHook(ctx: PluginInput) { const sessionStates = new Map(); function getOrCreateState(sessionID: string): InteractiveBashSessionState { @@ -178,6 +179,10 @@ export function createInteractiveBashSessionHook(_ctx: PluginInput) { await proc.exited; } catch {} } + + for (const sessionId of subagentSessions) { + ctx.client.session.abort({ path: { id: sessionId } }).catch(() => {}) + } } const toolExecuteAfter = async ( diff --git a/src/index.ts b/src/index.ts index f8c9b4f36..240844d96 100644 --- a/src/index.ts +++ b/src/index.ts @@ -264,6 +264,11 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { }); log("[index] onSubagentSessionCreated callback completed"); }, + onShutdown: () => { + tmuxSessionManager.cleanup().catch((error) => { + log("[index] tmux cleanup error during shutdown:", error) + }) + }, }); const atlasHook = isHookEnabled("atlas")