Add session close support and cleanup handling

benbrandt · benbrandt · commit 7b6100fdabbd · 2026-04-23T16:24:23.000+02:00
Advertise ACP close capability and implement
`unstable_closeSession` to interrupt active turns, unsubscribe
threads, remove per-session listeners, and forget session state.
diff --git a/src/CodexAcpClient.ts b/src/CodexAcpClient.ts
@@ -370,6 +370,28 @@ export class CodexAcpClient {
         this.codexClient.onElicitationRequest(sessionId, elicitationHandler);
     }
 
+    unsubscribeFromSessionEvents(sessionId: string): void {
+        this.codexClient.removeServerNotification(sessionId);
+        this.codexClient.removeApprovalRequest(sessionId);
+        this.codexClient.removeElicitationRequest(sessionId);
+    }
+
+    async closeSession(sessionId: string, currentTurnId: string | null): Promise<void> {
+        if (currentTurnId) {
+            try {
+                await this.codexClient.turnInterrupt({
+                    threadId: sessionId,
+                    turnId: currentTurnId,
+                });
+            } catch (err) {
+                logger.error(`Failed to interrupt turn while closing session ${sessionId}`, err);
+            }
+        }
+
+        await this.codexClient.threadUnsubscribe({threadId: sessionId});
+        this.unsubscribeFromSessionEvents(sessionId);
+    }
+
     async sendPrompt(
         request: acp.PromptRequest,
         agentMode: AgentMode,
diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts
@@ -101,6 +101,7 @@ export class CodexAcpServer implements acp.Agent {
                     image: true
                 },
                 sessionCapabilities: {
+                    close: {},
                     resume: { },
                     list: { }
                 },
@@ -238,6 +239,19 @@ export class CodexAcpServer implements acp.Agent {
         return await this.runWithProcessCheck(() => this.codexAcpClient.listSessions(params));
     }
 
+    async unstable_closeSession(params: acp.CloseSessionRequest): Promise<acp.CloseSessionResponse> {
+        logger.log("Closing session...", {sessionId: params.sessionId});
+        const sessionState = this.getSessionState(params.sessionId);
+
+        await this.runWithProcessCheck(() =>
+            this.codexAcpClient.closeSession(params.sessionId, sessionState.currentTurnId)
+        );
+        this.forgetSession(params.sessionId);
+
+        logger.log("Session closed", {sessionId: params.sessionId});
+        return {};
+    }
+
     async newSession(
         params: acp.NewSessionRequest,
     ): Promise<acp.NewSessionResponse> {
@@ -641,6 +655,11 @@ export class CodexAcpServer implements acp.Agent {
         return sessionState;
     }
 
+    private forgetSession(sessionId: string): void {
+        this.sessions.delete(sessionId);
+        this.pendingMcpStartupSessions.delete(sessionId);
+    }
+
     private resolveSessionMcpServers(
         mcpServers: Array<acp.McpServer>,
         recoverFromStartup: boolean,
@@ -678,7 +697,16 @@ export class CodexAcpServer implements acp.Agent {
                     pendingStartup.afterVersion,
                 )
             );
-            await this.publishMcpStartupStatus(sessionId, mcpStartup, pendingStartup.requestedServers);
+            const sessionState = this.sessions.get(sessionId);
+            const pendingStartup = this.pendingMcpStartupSessions.get(sessionId);
+            if (sessionState && pendingStartup) {
+                sessionState.sessionMcpServers = mcpStartup.ready.filter(serverName =>
+                    pendingStartup.requestedServers.has(serverName)
+                );
+await this.publishMcpStartupStatus(sessionId, mcpStartup, pendingStartup.requestedServers);
+            } else {
+                logger.log("Skipping MCP startup status for closed session", {sessionId});
+            }
         } catch (err) {
             logger.error(`Failed to publish MCP startup status for session ${sessionId}`, err);
         } finally {
diff --git a/src/CodexAppServerClient.ts b/src/CodexAppServerClient.ts
@@ -34,6 +34,8 @@ import type {
     ThreadResumeResponse,
     ThreadStartParams,
     ThreadStartResponse,
+    ThreadUnsubscribeParams,
+    ThreadUnsubscribeResponse,
     TurnCompletedNotification,
     TurnInterruptParams,
     TurnInterruptResponse,
@@ -143,10 +145,18 @@ export class CodexAppServerClient {
         this.approvalHandlers.set(threadId, handler);
     }
 
+    removeApprovalRequest(threadId: string): void {
+        this.approvalHandlers.delete(threadId);
+    }
+
     onElicitationRequest(threadId: string, handler: ElicitationHandler): void {
         this.elicitationHandlers.set(threadId, handler);
     }
 
+    removeElicitationRequest(threadId: string): void {
+        this.elicitationHandlers.delete(threadId);
+    }
+
     async initialize(params: InitializeParams): Promise<InitializeResponse> {
         return await this.sendRequest({ method: "initialize", params: params });
     }
@@ -179,6 +189,10 @@ export class CodexAppServerClient {
         return await this.sendRequest({ method: "thread/read", params: params });
     }
 
+    async threadUnsubscribe(params: ThreadUnsubscribeParams): Promise<ThreadUnsubscribeResponse> {
+        return await this.sendRequest({ method: "thread/unsubscribe", params });
+    }
+
     async listMcpServerStatus(params: ListMcpServerStatusParams): Promise<ListMcpServerStatusResponse> {
         return await this.sendRequest({ method: "mcpServerStatus/list", params });
     }
@@ -248,6 +262,10 @@ export class CodexAppServerClient {
         this.notificationHandlers.set(sessionId, callback);
     }
 
+    removeServerNotification(sessionId: string): void {
+        this.notificationHandlers.delete(sessionId);
+    }
+
     private codexEventHandlers: Array<(event: CodexConnectionEvent) => void> = [];
     onClientTransportEvent(callback: (event: CodexConnectionEvent) => void){
         this.codexEventHandlers.push(callback);
diff --git a/src/__tests__/CodexACPAgent/CodexAcpClient.test.ts b/src/__tests__/CodexACPAgent/CodexAcpClient.test.ts
@@ -915,4 +915,117 @@ describe('ACP server test', { timeout: 40_000 }, () => {
             }
         });
     });
+
+    it('closes an active session by interrupting the turn, unsubscribing the thread, and removing listeners', async () => {
+        const mockFixture = createCodexMockTestFixture();
+        const codexAcpClient = mockFixture.getCodexAcpClient();
+        const codexAppServerClient = mockFixture.getCodexAppServerClient();
+
+        await codexAcpClient.subscribeToSessionEvents(
+            "session-close",
+            vi.fn(),
+            {
+                handleCommandExecution: vi.fn(),
+                handleFileChange: vi.fn(),
+            },
+            {
+                handleElicitation: vi.fn(),
+            }
+        );
+
+        // @ts-expect-error verifying test-only access to private maps
+        expect(codexAppServerClient.notificationHandlers.has("session-close")).toBe(true);
+        // @ts-expect-error verifying test-only access to private maps
+        expect(codexAppServerClient.approvalHandlers.has("session-close")).toBe(true);
+        // @ts-expect-error verifying test-only access to private maps
+        expect(codexAppServerClient.elicitationHandlers.has("session-close")).toBe(true);
+
+        await codexAcpClient.closeSession("session-close", "turn-close");
+
+        const closeRequests = mockFixture.getCodexConnectionEvents([])
+            .filter((event) => event.eventType === "request");
+        expect(closeRequests).toEqual([
+            {
+                eventType: "request",
+                method: "turn/interrupt",
+                params: {threadId: "session-close", turnId: "turn-close"},
+            },
+            {
+                eventType: "request",
+                method: "thread/unsubscribe",
+                params: {threadId: "session-close"},
+            },
+        ]);
+
+        // @ts-expect-error verifying test-only access to private maps
+        expect(codexAppServerClient.notificationHandlers.has("session-close")).toBe(false);
+        // @ts-expect-error verifying test-only access to private maps
+        expect(codexAppServerClient.approvalHandlers.has("session-close")).toBe(false);
+        // @ts-expect-error verifying test-only access to private maps
+        expect(codexAppServerClient.elicitationHandlers.has("session-close")).toBe(false);
+    });
+
+    it('removes session bookkeeping after unstable_closeSession', async () => {
+        const mockFixture = createCodexMockTestFixture();
+        const codexAcpAgent = mockFixture.getCodexAcpAgent();
+        const codexAcpClient = mockFixture.getCodexAcpClient();
+        const sessionState = createTestSessionState({
+            sessionId: "session-close",
+            currentTurnId: "turn-close",
+        });
+        const closeSpy = vi.spyOn(codexAcpClient, "closeSession").mockResolvedValue();
+
+        // @ts-expect-error seeding private session store for focused close-session test
+        codexAcpAgent.sessions.set(sessionState.sessionId, sessionState);
+        // @ts-expect-error seeding private session store for focused close-session test
+        codexAcpAgent.pendingMcpStartupSessions.set(sessionState.sessionId, {
+            requestedServers: new Set(["alpha"]),
+        });
+
+        await expect(
+            codexAcpAgent.unstable_closeSession({sessionId: sessionState.sessionId})
+        ).resolves.toEqual({});
+        expect(closeSpy).toHaveBeenCalledWith(sessionState.sessionId, sessionState.currentTurnId);
+        expect(() => codexAcpAgent.getSessionState(sessionState.sessionId)).toThrow(
+            `Session ${sessionState.sessionId} not found`
+        );
+        // @ts-expect-error verifying test-only access to private map
+        expect(codexAcpAgent.pendingMcpStartupSessions.has(sessionState.sessionId)).toBe(false);
+    });
+
+    it('skips late MCP startup updates after a session is closed', async () => {
+        const mockFixture = createCodexMockTestFixture();
+        const codexAcpAgent = mockFixture.getCodexAcpAgent();
+        const codexAcpClient = mockFixture.getCodexAcpClient();
+        const sessionState = createTestSessionState({
+            sessionId: "session-close",
+        });
+
+        let resolveStartup!: (event: any) => void;
+        const startupPromise = new Promise((resolve) => {
+            resolveStartup = resolve;
+        });
+        vi.spyOn(codexAcpClient, "awaitMcpStartupResult").mockReturnValue(startupPromise as Promise<any>);
+
+        // @ts-expect-error seeding private session store for focused close-session test
+        codexAcpAgent.sessions.set(sessionState.sessionId, sessionState);
+        // @ts-expect-error seeding private session store for focused close-session test
+        codexAcpAgent.pendingMcpStartupSessions.set(sessionState.sessionId, {
+            requestedServers: new Set(["alpha"]),
+        });
+
+        // @ts-expect-error exercising private helper to verify close-session race handling
+        const publishPromise = codexAcpAgent.doPublishMcpStartupStatus(sessionState.sessionId, 1);
+        // @ts-expect-error exercising private helper to simulate close-session cleanup
+        codexAcpAgent.forgetSession(sessionState.sessionId);
+
+        resolveStartup({
+            ready: ["alpha"],
+            failed: [],
+            cancelled: [],
+        });
+        await publishPromise;
+
+        expect(mockFixture.getAcpConnectionEvents([])).toEqual([]);
+    });
 });
diff --git a/src/__tests__/CodexACPAgent/initialize.test.ts b/src/__tests__/CodexACPAgent/initialize.test.ts
@@ -37,6 +37,7 @@ describe('CodexACPAgent - initialize', () => {
                     image: true
                 },
                 sessionCapabilities: {
+                    close: {},
                     resume: {},
                     list: {},
                 },
