Merge pull request #3115 from code-yeongyu/fix/prepublish-tmux-lifecycle

fix(tmux): clean cancel bookkeeping and unify isolated cleanup

src/features/background-agent/manager.test.ts

@@ -15,6 +15,7 @@ afterAll(() => { mock.restore() })
 import { getSessionPromptParams, clearSessionPromptParams } from "../../shared/session-prompt-params-state"
 import { tmpdir } from "node:os"
 import type { PluginInput } from "@opencode-ai/plugin"
+import { _resetForTesting as resetClaudeCodeSessionState, subagentSessions } from "../claude-code-session-state"
 import type { BackgroundTask, ResumeInput } from "./types"
 import { MIN_IDLE_TIME_MS } from "./constants"
 import { BackgroundManager } from "./manager"
@@ -213,6 +214,10 @@ function getCompletionTimers(manager: BackgroundManager): Map<string, ReturnType
   return (manager as unknown as { completionTimers: Map<string, ReturnType<typeof setTimeout>> }).completionTimers
 }
 
+function getRootDescendantCounts(manager: BackgroundManager): Map<string, number> {
+  return (manager as unknown as { rootDescendantCounts: Map<string, number> }).rootDescendantCounts
+}
+
 function getQueuesByKey(
   manager: BackgroundManager
 ): Map<string, Array<{ task: BackgroundTask; input: import("./types").LaunchInput }>> {
@@ -2585,6 +2590,7 @@ describe("BackgroundManager - Non-blocking Queue Integration", () => {
 
     test("should keep task cancelled when cancelled during tmux callback before running state is assigned", async () => {
       // given
+      resetClaudeCodeSessionState()
       const originalTmuxEnvironment = process.env.TMUX
       process.env.TMUX = "test-session"
 
@@ -2674,7 +2680,10 @@ describe("BackgroundManager - Non-blocking Queue Integration", () => {
         expect(promptAsyncSessionIDs).not.toContain(createdSessionID)
         expect(abortCalls).toEqual([createdSessionID])
         expect(getConcurrencyManager(manager).getCount("test-agent")).toBe(0)
+        expect(getRootDescendantCounts(manager).has("parent-session")).toBe(false)
+        expect(subagentSessions.has(createdSessionID)).toBe(false)
       } finally {
+        resetClaudeCodeSessionState()
         if (originalTmuxEnvironment === undefined) {
           delete process.env.TMUX
         } else {

src/features/background-agent/manager.ts

@@ -493,6 +493,10 @@ export class BackgroundManager {
 
     if (this.tasks.get(task.id)?.status === "cancelled") {
       await this.abortSessionWithLogging(sessionID, "cancelled during tmux setup")
+      subagentSessions.delete(sessionID)
+      if (task.rootSessionID) {
+        this.unregisterRootDescendant(task.rootSessionID)
+      }
       this.concurrencyManager.release(concurrencyKey)
       return
     }

src/features/tmux-subagent/manager.test.ts

@@ -168,6 +168,10 @@ function createTmuxConfig(overrides?: Partial<TmuxConfig>): TmuxConfig {
   }
 }
 
+function getTrackedSessions(manager: object): Map<string, { paneId: string; closePending: boolean; closeRetryCount: number }> {
+  return Reflect.get(manager, 'sessions') as Map<string, { paneId: string; closePending: boolean; closeRetryCount: number }>
+}
+
 describe('TmuxSessionManager', () => {
   beforeEach(() => {
     mockQueryWindowState.mockClear()
@@ -1532,6 +1536,280 @@ describe('TmuxSessionManager', () => {
   })
 
   describe('cleanup', () => {
+    test('#given session isolation with two tracked panes #when polling closes both sessions #then it reassigns the anchor and cleans up the isolated container', async () => {
+      // given
+      mockIsInsideTmux.mockReturnValue(true)
+      mockQueryWindowState.mockImplementation(async (paneId: string) => {
+        if (paneId === '%isolated-session-ses_first') {
+          return createWindowState({
+            mainPane: {
+              paneId: '%isolated-session-ses_first',
+              width: 110,
+              height: 44,
+              left: 0,
+              top: 0,
+              title: 'isolated',
+              isActive: true,
+            },
+            agentPanes: [
+              {
+                paneId: '%mock',
+                width: 40,
+                height: 44,
+                left: 110,
+                top: 0,
+                title: 'omo-subagent-Second Task',
+                isActive: false,
+              },
+            ],
+          })
+        }
+
+        if (paneId === '%mock') {
+          return createWindowState({
+            mainPane: {
+              paneId: '%isolated-session-ses_first',
+              width: 110,
+              height: 44,
+              left: 0,
+              top: 0,
+              title: 'isolated',
+              isActive: true,
+            },
+            agentPanes: [
+              {
+                paneId: '%mock',
+                width: 40,
+                height: 44,
+                left: 110,
+                top: 0,
+                title: 'omo-subagent-Second Task',
+                isActive: false,
+              },
+            ],
+          })
+        }
+
+        return createWindowState()
+      })
+
+      const { TmuxSessionManager } = await import('./manager')
+      const manager = new TmuxSessionManager(createMockContext(), createTmuxConfig({
+        enabled: true,
+        isolation: 'session',
+      }), mockTmuxDeps)
+
+      await manager.onSessionCreated(createSessionCreatedEvent('ses_first', 'ses_parent', 'First Task'))
+      await manager.onSessionCreated(createSessionCreatedEvent('ses_second', 'ses_parent', 'Second Task'))
+      mockExecuteAction.mockClear()
+
+      const closeSessionById = Reflect.get(manager, 'closeSessionById') as (sessionId: string) => Promise<void>
+
+      // when
+      await closeSessionById.call(manager, 'ses_first')
+
+      // then
+      expect(mockExecuteAction.mock.calls[0]?.[0]).toEqual({
+        type: 'close',
+        paneId: '%isolated-session-ses_first',
+        sessionId: 'ses_first',
+      })
+      expect(Reflect.get(manager, 'isolatedContainerPaneId')).toBe('%isolated-session-ses_first')
+      expect(Reflect.get(manager, 'isolatedWindowPaneId')).toBe('%mock')
+
+      // when
+      await closeSessionById.call(manager, 'ses_second')
+
+      // then
+      expect(mockExecuteAction).toHaveBeenCalledTimes(3)
+      expect(mockExecuteAction.mock.calls[1]?.[0]).toEqual({
+        type: 'close',
+        paneId: '%mock',
+        sessionId: 'ses_second',
+      })
+      expect(mockExecuteAction.mock.calls[2]?.[0]).toEqual({
+        type: 'close',
+        paneId: '%isolated-session-ses_first',
+        sessionId: 'ses_second',
+      })
+      expect(Reflect.get(manager, 'isolatedContainerPaneId')).toBeUndefined()
+      expect(Reflect.get(manager, 'isolatedWindowPaneId')).toBeUndefined()
+    })
+
+    test('#given session isolation with two tracked panes #when process shutdown cleanup runs #then it closes panes and the isolated container through the shared close path', async () => {
+      // given
+      mockIsInsideTmux.mockReturnValue(true)
+      mockQueryWindowState.mockImplementation(async (paneId: string) => {
+        if (paneId === '%isolated-session-ses_first') {
+          return createWindowState({
+            mainPane: {
+              paneId: '%isolated-session-ses_first',
+              width: 110,
+              height: 44,
+              left: 0,
+              top: 0,
+              title: 'isolated',
+              isActive: true,
+            },
+            agentPanes: [
+              {
+                paneId: '%mock',
+                width: 40,
+                height: 44,
+                left: 110,
+                top: 0,
+                title: 'omo-subagent-Second Task',
+                isActive: false,
+              },
+            ],
+          })
+        }
+
+        if (paneId === '%mock') {
+          return createWindowState({
+            mainPane: {
+              paneId: '%isolated-session-ses_first',
+              width: 110,
+              height: 44,
+              left: 0,
+              top: 0,
+              title: 'isolated',
+              isActive: true,
+            },
+            agentPanes: [
+              {
+                paneId: '%mock',
+                width: 40,
+                height: 44,
+                left: 110,
+                top: 0,
+                title: 'omo-subagent-Second Task',
+                isActive: false,
+              },
+            ],
+          })
+        }
+
+        return createWindowState()
+      })
+
+      const { TmuxSessionManager } = await import('./manager')
+      const manager = new TmuxSessionManager(createMockContext(), createTmuxConfig({
+        enabled: true,
+        isolation: 'session',
+      }), mockTmuxDeps)
+
+      await manager.onSessionCreated(createSessionCreatedEvent('ses_first', 'ses_parent', 'First Task'))
+      await manager.onSessionCreated(createSessionCreatedEvent('ses_second', 'ses_parent', 'Second Task'))
+      mockExecuteAction.mockClear()
+
+      // when
+      await manager.cleanup()
+
+      // then
+      expect(mockExecuteAction).toHaveBeenCalledTimes(3)
+      expect(mockExecuteAction.mock.calls[0]?.[0]).toEqual({
+        type: 'close',
+        paneId: '%isolated-session-ses_first',
+        sessionId: 'ses_first',
+      })
+      expect(mockExecuteAction.mock.calls[1]?.[0]).toEqual({
+        type: 'close',
+        paneId: '%mock',
+        sessionId: 'ses_second',
+      })
+      expect(mockExecuteAction.mock.calls[2]?.[0]).toEqual({
+        type: 'close',
+        paneId: '%isolated-session-ses_first',
+        sessionId: 'ses_second',
+      })
+      expect(Reflect.get(manager, 'isolatedContainerPaneId')).toBeUndefined()
+      expect(Reflect.get(manager, 'isolatedWindowPaneId')).toBeUndefined()
+    })
+
+    test('#given an isolated anchor close that fails once #when retryPendingCloses succeeds on retry #then it reassigns the isolated anchor through the shared cleanup path', async () => {
+      // given
+      mockIsInsideTmux.mockReturnValue(true)
+      mockQueryWindowState.mockImplementation(async (paneId: string) => {
+        if (paneId === '%isolated-session-ses_first') {
+          return createWindowState({
+            mainPane: {
+              paneId: '%isolated-session-ses_first',
+              width: 110,
+              height: 44,
+              left: 0,
+              top: 0,
+              title: 'isolated',
+              isActive: true,
+            },
+            agentPanes: [
+              {
+                paneId: '%mock',
+                width: 40,
+                height: 44,
+                left: 110,
+                top: 0,
+                title: 'omo-subagent-Second Task',
+                isActive: false,
+              },
+            ],
+          })
+        }
+
+        return createWindowState()
+      })
+
+      let closeAttemptCount = 0
+      mockExecuteAction.mockImplementation(async (action: PaneAction) => {
+        if (action.type === 'close' && action.sessionId === 'ses_first') {
+          closeAttemptCount += 1
+          if (closeAttemptCount === 1) {
+            return { success: false }
+          }
+        }
+
+        return { success: true }
+      })
+
+      const { TmuxSessionManager } = await import('./manager')
+      const manager = new TmuxSessionManager(createMockContext(), createTmuxConfig({
+        enabled: true,
+        isolation: 'session',
+      }), mockTmuxDeps)
+
+      await manager.onSessionCreated(createSessionCreatedEvent('ses_first', 'ses_parent', 'First Task'))
+      await manager.onSessionCreated(createSessionCreatedEvent('ses_second', 'ses_parent', 'Second Task'))
+      mockExecuteAction.mockClear()
+
+      const closeSessionById = Reflect.get(manager, 'closeSessionById') as (sessionId: string) => Promise<void>
+      const retryPendingCloses = Reflect.get(manager, 'retryPendingCloses') as () => Promise<void>
+
+      // when
+      await closeSessionById.call(manager, 'ses_first')
+
+      // then
+      expect(getTrackedSessions(manager).get('ses_first')?.closePending).toBe(true)
+      expect(Reflect.get(manager, 'isolatedWindowPaneId')).toBe('%isolated-session-ses_first')
+
+      // when
+      await retryPendingCloses.call(manager)
+
+      // then
+      expect(getTrackedSessions(manager).has('ses_first')).toBe(false)
+      expect(Reflect.get(manager, 'isolatedContainerPaneId')).toBe('%isolated-session-ses_first')
+      expect(Reflect.get(manager, 'isolatedWindowPaneId')).toBe('%mock')
+      expect(mockExecuteAction.mock.calls[0]?.[0]).toEqual({
+        type: 'close',
+        paneId: '%isolated-session-ses_first',
+        sessionId: 'ses_first',
+      })
+      expect(mockExecuteAction.mock.calls[1]?.[0]).toEqual({
+        type: 'close',
+        paneId: '%isolated-session-ses_first',
+        sessionId: 'ses_first',
+      })
+    })
+
     test('closes all tracked panes', async () => {
       // given
       mockIsInsideTmux.mockReturnValue(true)

src/features/tmux-subagent/manager.ts

@@ -283,9 +283,11 @@ export class TmuxSessionManager {
     }
   }
 
-  private async tryCloseTrackedSession(tracked: TrackedSession): Promise<boolean> {
+  private async closeTrackedSessionPane(args: {
-    const state = await this.queryWindowStateSafely()
+    tracked: TrackedSession
-    if (!state) return false
+    state: WindowState
+  }): Promise<boolean> {
+    const { tracked, state } = args
 
     try {
       const result = await executeAction(
@@ -309,6 +311,37 @@ export class TmuxSessionManager {
     }
   }
 
+  private async finalizeTrackedSessionClose(args: {
+    tracked: TrackedSession
+    state: WindowState
+    isolatedPaneAlreadyClosed: boolean
+  }): Promise<void> {
+    const { tracked, state, isolatedPaneAlreadyClosed } = args
+    this.removeTrackedSession(tracked.sessionId)
+    await this.cleanupIsolatedContainerAfterSessionDeletion(
+      tracked,
+      isolatedPaneAlreadyClosed,
+      state,
+    )
+  }
+
+  private async closeTrackedSession(tracked: TrackedSession): Promise<boolean> {
+    const state = await this.queryWindowStateSafely()
+    if (!state) return false
+
+    const closed = await this.closeTrackedSessionPane({ tracked, state })
+    if (!closed) {
+      return false
+    }
+
+    await this.finalizeTrackedSessionClose({
+      tracked,
+      state,
+      isolatedPaneAlreadyClosed: true,
+    })
+    return true
+  }
+
   private async retryPendingCloses(): Promise<void> {
     const pendingSessions = Array.from(this.sessions.values()).filter(
       (tracked) => tracked.closePending,
@@ -327,14 +360,13 @@ export class TmuxSessionManager {
         continue
       }
 
-      const closed = await this.tryCloseTrackedSession(tracked)
+      const closed = await this.closeTrackedSession(tracked)
       if (closed) {
         log("[tmux-session-manager] retried close succeeded", {
           sessionId: tracked.sessionId,
           paneId: tracked.paneId,
           closeRetryCount: tracked.closeRetryCount,
         })
-        this.removeTrackedSession(tracked.sessionId)
         continue
       }
 
@@ -825,8 +857,11 @@ export class TmuxSessionManager {
 
     const closeAction = decideCloseAction(state, event.sessionID, this.getSessionMappings())
     if (!closeAction) {
-      this.removeTrackedSession(event.sessionID)
+      await this.finalizeTrackedSessionClose({
-      await this.cleanupIsolatedContainerAfterSessionDeletion(tracked, false, state)
+        tracked,
+        state,
+        isolatedPaneAlreadyClosed: false,
+      })
       return
     }
 
@@ -854,12 +889,11 @@ export class TmuxSessionManager {
       return
     }
 
-    this.removeTrackedSession(event.sessionID)
+    await this.finalizeTrackedSessionClose({
-    await this.cleanupIsolatedContainerAfterSessionDeletion(
       tracked,
-      isolatedPaneAlreadyClosed,
       state,
-    )
+      isolatedPaneAlreadyClosed,
+    })
   }
 
 
@@ -882,13 +916,11 @@ export class TmuxSessionManager {
       paneId: tracked.paneId,
     })
 
-    const closed = await this.tryCloseTrackedSession(tracked)
+    const closed = await this.closeTrackedSession(tracked)
     if (!closed) {
       this.markSessionClosePending(sessionId)
       return
     }
-
-    this.removeTrackedSession(sessionId)
   }
 
   createEventHandler(): (input: { event: { type: string; properties?: unknown } }) => Promise<void> {