diff --git a/src/features/background-agent/manager.test.ts b/src/features/background-agent/manager.test.ts index d93246447..e83054d5a 100644 --- a/src/features/background-agent/manager.test.ts +++ b/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> }).completionTimers } +function getRootDescendantCounts(manager: BackgroundManager): Map { + return (manager as unknown as { rootDescendantCounts: Map }).rootDescendantCounts +} + function getQueuesByKey( manager: BackgroundManager ): Map> { @@ -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 { diff --git a/src/features/background-agent/manager.ts b/src/features/background-agent/manager.ts index aa5979cf4..743d7769a 100644 --- a/src/features/background-agent/manager.ts +++ b/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 } diff --git a/src/features/tmux-subagent/manager.test.ts b/src/features/tmux-subagent/manager.test.ts index d82a17083..419e4b42d 100644 --- a/src/features/tmux-subagent/manager.test.ts +++ b/src/features/tmux-subagent/manager.test.ts @@ -168,6 +168,10 @@ function createTmuxConfig(overrides?: Partial): TmuxConfig { } } +function getTrackedSessions(manager: object): Map { + return Reflect.get(manager, 'sessions') as Map +} + 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 + + // 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 + const retryPendingCloses = Reflect.get(manager, 'retryPendingCloses') as () => Promise + + // 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) diff --git a/src/features/tmux-subagent/manager.ts b/src/features/tmux-subagent/manager.ts index bd5f8bf5d..2b688d998 100644 --- a/src/features/tmux-subagent/manager.ts +++ b/src/features/tmux-subagent/manager.ts @@ -283,9 +283,11 @@ export class TmuxSessionManager { } } - private async tryCloseTrackedSession(tracked: TrackedSession): Promise { - const state = await this.queryWindowStateSafely() - if (!state) return false + private async closeTrackedSessionPane(args: { + tracked: TrackedSession + state: WindowState + }): Promise { + 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 { + const { tracked, state, isolatedPaneAlreadyClosed } = args + this.removeTrackedSession(tracked.sessionId) + await this.cleanupIsolatedContainerAfterSessionDeletion( + tracked, + isolatedPaneAlreadyClosed, + state, + ) + } + + private async closeTrackedSession(tracked: TrackedSession): Promise { + 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 { 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.cleanupIsolatedContainerAfterSessionDeletion(tracked, false, state) + await this.finalizeTrackedSessionClose({ + tracked, + state, + isolatedPaneAlreadyClosed: false, + }) return } @@ -854,12 +889,11 @@ export class TmuxSessionManager { return } - this.removeTrackedSession(event.sessionID) - await this.cleanupIsolatedContainerAfterSessionDeletion( + await this.finalizeTrackedSessionClose({ 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 {