diff --git a/src/features/background-agent/manager-shutdown-global-cleanup.test.ts b/src/features/background-agent/manager-shutdown-global-cleanup.test.ts new file mode 100644 index 000000000..d238b2dc0 --- /dev/null +++ b/src/features/background-agent/manager-shutdown-global-cleanup.test.ts @@ -0,0 +1,97 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test" +import { tmpdir } from "node:os" + +import { _resetForTesting, subagentSessions } from "../claude-code-session-state" +import { SessionCategoryRegistry } from "../../shared/session-category-registry" +import { BackgroundManager } from "./manager" +import type { BackgroundTask } from "./types" + +function createTask(overrides: Partial & { id: string; sessionID: string }): BackgroundTask { + return { + parentSessionID: "parent-session", + parentMessageID: "parent-message", + description: "test task", + prompt: "test prompt", + agent: "explore", + status: "running", + startedAt: new Date(), + ...overrides, + } +} + +function createBackgroundManager(): BackgroundManager { + return new BackgroundManager({ + client: { + session: { + abort: async () => ({}), + prompt: async () => ({}), + promptAsync: async () => ({}), + }, + } as never, + project: {} as never, + directory: tmpdir(), + worktree: tmpdir(), + serverUrl: new URL("https://example.com"), + $: {} as never, + } as never) +} + +describe("BackgroundManager shutdown global cleanup", () => { + beforeEach(() => { + // given + _resetForTesting() + SessionCategoryRegistry.clear() + }) + + afterEach(() => { + // given + _resetForTesting() + SessionCategoryRegistry.clear() + }) + + test("removes tracked session IDs from subagentSessions and SessionCategoryRegistry on shutdown", async () => { + // given + const runningSessionID = "ses-running-shutdown-cleanup" + const completedSessionID = "ses-completed-shutdown-cleanup" + const unrelatedSessionID = "ses-unrelated-shutdown-cleanup" + const manager = createBackgroundManager() + const tasks = new Map([ + [ + "task-running-shutdown-cleanup", + createTask({ + id: "task-running-shutdown-cleanup", + sessionID: runningSessionID, + }), + ], + [ + "task-completed-shutdown-cleanup", + createTask({ + id: "task-completed-shutdown-cleanup", + sessionID: completedSessionID, + status: "completed", + completedAt: new Date(), + }), + ], + ]) + + Object.assign(manager, { tasks }) + + subagentSessions.add(runningSessionID) + subagentSessions.add(completedSessionID) + subagentSessions.add(unrelatedSessionID) + SessionCategoryRegistry.register(runningSessionID, "quick") + SessionCategoryRegistry.register(completedSessionID, "deep") + SessionCategoryRegistry.register(unrelatedSessionID, "test") + + // when + await manager.shutdown() + + // then + expect(subagentSessions.has(runningSessionID)).toBe(false) + expect(subagentSessions.has(completedSessionID)).toBe(false) + expect(subagentSessions.has(unrelatedSessionID)).toBe(true) + expect(SessionCategoryRegistry.has(runningSessionID)).toBe(false) + expect(SessionCategoryRegistry.has(completedSessionID)).toBe(false) + expect(SessionCategoryRegistry.has(unrelatedSessionID)).toBe(true) + }) +}) diff --git a/src/features/background-agent/manager.ts b/src/features/background-agent/manager.ts index d133f2d77..130bf6fde 100644 --- a/src/features/background-agent/manager.ts +++ b/src/features/background-agent/manager.ts @@ -1707,9 +1707,14 @@ Use \`background_output(task_id="${task.id}")\` to retrieve this result when rea this.shutdownTriggered = true log("[background-agent] Shutting down BackgroundManager") this.stopPolling() + const trackedSessionIDs = new Set() // Abort all running sessions to prevent zombie processes (#1240) for (const task of this.tasks.values()) { + if (task.sessionID) { + trackedSessionIDs.add(task.sessionID) + } + if (task.status === "running" && task.sessionID) { this.client.session.abort({ path: { id: task.sessionID }, @@ -1744,6 +1749,11 @@ Use \`background_output(task_id="${task.id}")\` to retrieve this result when rea } this.idleDeferralTimers.clear() + for (const sessionID of trackedSessionIDs) { + subagentSessions.delete(sessionID) + SessionCategoryRegistry.remove(sessionID) + } + this.concurrencyManager.clear() this.tasks.clear() this.notifications.clear() diff --git a/src/plugin/event.ts b/src/plugin/event.ts index 0246c45c4..d18ec3691 100644 --- a/src/plugin/event.ts +++ b/src/plugin/event.ts @@ -319,6 +319,7 @@ export function createEventHandler(args: { } if (sessionInfo?.id) { + const wasSyncSubagentSession = syncSubagentSessions.has(sessionInfo.id); clearSessionAgent(sessionInfo.id); lastHandledModelErrorMessageID.delete(sessionInfo.id); lastHandledRetryStatusKey.delete(sessionInfo.id); @@ -329,6 +330,9 @@ export function createEventHandler(args: { firstMessageVariantGate.clear(sessionInfo.id); clearSessionModel(sessionInfo.id); syncSubagentSessions.delete(sessionInfo.id); + if (wasSyncSubagentSession) { + subagentSessions.delete(sessionInfo.id); + } deleteSessionTools(sessionInfo.id); await managers.skillMcpManager.disconnectSession(sessionInfo.id); await lspManager.cleanupTempDirectoryClients(); diff --git a/src/tools/call-omo-agent/reused-sync-session-delete-cleanup.test.ts b/src/tools/call-omo-agent/reused-sync-session-delete-cleanup.test.ts new file mode 100644 index 000000000..cfaf9f497 --- /dev/null +++ b/src/tools/call-omo-agent/reused-sync-session-delete-cleanup.test.ts @@ -0,0 +1,87 @@ +import { afterEach, describe, expect, it } from "bun:test" + +import { + _resetForTesting, + subagentSessions, + syncSubagentSessions, +} from "../../features/claude-code-session-state" +import { createEventHandler } from "../../plugin/event" + +function createMinimalEventHandler() { + return createEventHandler({ + ctx: {} as never, + pluginConfig: {} as never, + firstMessageVariantGate: { + markSessionCreated: () => {}, + clear: () => {}, + }, + managers: { + tmuxSessionManager: { + onSessionCreated: async () => {}, + onSessionDeleted: async () => {}, + }, + skillMcpManager: { + disconnectSession: async () => {}, + }, + } as never, + hooks: { + autoUpdateChecker: { event: async () => {} }, + claudeCodeHooks: { event: async () => {} }, + backgroundNotificationHook: { event: async () => {} }, + sessionNotification: async () => {}, + todoContinuationEnforcer: { handler: async () => {} }, + unstableAgentBabysitter: { event: async () => {} }, + contextWindowMonitor: { event: async () => {} }, + directoryAgentsInjector: { event: async () => {} }, + directoryReadmeInjector: { event: async () => {} }, + rulesInjector: { event: async () => {} }, + thinkMode: { event: async () => {} }, + anthropicContextWindowLimitRecovery: { event: async () => {} }, + runtimeFallback: undefined, + modelFallback: undefined, + agentUsageReminder: { event: async () => {} }, + categorySkillReminder: { event: async () => {} }, + interactiveBashSession: { event: async () => {} }, + ralphLoop: { event: async () => {} }, + stopContinuationGuard: { event: async () => {}, isStopped: () => false }, + compactionTodoPreserver: { event: async () => {} }, + writeExistingFileGuard: { event: async () => {} }, + atlasHook: { handler: async () => {} }, + } as never, + }) +} + +describe("reused sync session delete cleanup", () => { + afterEach(() => { + _resetForTesting() + }) + + it("removes reused sync sessions from subagentSessions when session.deleted fires", async () => { + // given + const syncSessionID = "ses-reused-sync-delete-cleanup" + const unrelatedSubagentSessionID = "ses-unrelated-subagent-delete-cleanup" + const eventHandler = createMinimalEventHandler() + const input = { + event: { + type: "session.deleted", + properties: { + info: { + id: syncSessionID, + }, + }, + }, + } as Parameters>[0] + + subagentSessions.add(syncSessionID) + syncSubagentSessions.add(syncSessionID) + subagentSessions.add(unrelatedSubagentSessionID) + + // when + await eventHandler(input) + + // then + expect(syncSubagentSessions.has(syncSessionID)).toBe(false) + expect(subagentSessions.has(syncSessionID)).toBe(false) + expect(subagentSessions.has(unrelatedSubagentSessionID)).toBe(true) + }) +})