fix(background-agent): clean global subagentSessions and SessionCategoryRegistry on dispose

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<BackgroundTask> & { 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<string, BackgroundTask>([
+      [
+        "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)
+  })
+})

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<string>()
 
     // 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()

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();

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<ReturnType<typeof createEventHandler>>[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)
+  })
+})