From 4ca2514809bfa6b24d19c91f0630775b937827ca Mon Sep 17 00:00:00 2001 From: Ben Brandt Date: Thu, 23 Apr 2026 16:04:54 +0200 Subject: [PATCH] Track MCP startup per server status updates Replace the legacy startup-complete event flow with typed `mcpServer/startupStatus/updated` handling and wait for requested servers to reach terminal states. Also stop restoring persisted MCP servers on load when the request omits them, and update tests to cover the new behavior. --- src/CodexAcpClient.ts | 21 +- src/CodexAcpServer.ts | 73 ++++--- src/CodexAppServerClient.ts | 180 +++++++++++------- .../CodexACPAgent/CodexAcpClient.test.ts | 52 ++--- .../CodexACPAgent/load-session.test.ts | 22 +-- 5 files changed, 195 insertions(+), 153 deletions(-) diff --git a/src/CodexAcpClient.ts b/src/CodexAcpClient.ts index 923ecbcf..9566dc51 100644 --- a/src/CodexAcpClient.ts +++ b/src/CodexAcpClient.ts @@ -2,12 +2,16 @@ import {isCodexAuthRequest} from "./CodexAuthMethod"; import type {EmbeddedResourceResource} from "@agentclientprotocol/sdk"; import * as acp from "@agentclientprotocol/sdk"; import {type McpServer, RequestError} from "@agentclientprotocol/sdk"; -import type {ApprovalHandler, CodexAppServerClient, ElicitationHandler} from "./CodexAppServerClient"; +import type { + ApprovalHandler, + CodexAppServerClient, + ElicitationHandler, + McpStartupResult, +} from "./CodexAppServerClient"; import open from "open"; import type {Disposable} from "vscode-jsonrpc"; import type { ClientInfo, - McpStartupCompleteEvent, ReasoningEffort, ServerNotification } from "./app-server"; @@ -274,17 +278,12 @@ export class CodexAcpClient { }; } - async awaitMcpStartup(mcpStartupVersion: number): Promise> { - const startup = await this.codexClient.awaitMcpStartup(mcpStartupVersion); - return startup.ready; - } - - async awaitMcpStartupResult(mcpStartupVersion: number): Promise { - return await this.codexClient.awaitMcpStartup(mcpStartupVersion); + async awaitMcpServerStartup(serverNames: Array, afterVersion: number): Promise { + return await this.codexClient.awaitMcpServerStartup(serverNames, afterVersion); } - getMcpStartupCompleteVersion(): number { - return this.codexClient.getMcpStartupCompleteVersion(); + getMcpServerStartupVersion(): number { + return this.codexClient.getMcpServerStartupVersion(); } private createSessionConfig(projectPath: string, mcpServers: Array): JsonObject { diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts index f7425ff6..a8932fb0 100644 --- a/src/CodexAcpServer.ts +++ b/src/CodexAcpServer.ts @@ -10,8 +10,9 @@ import {CodexApprovalHandler} from "./CodexApprovalHandler"; import {CodexElicitationHandler} from "./CodexElicitationHandler"; import {CodexAuthMethods, type CodexAuthRequest} from "./CodexAuthMethod"; import {CodexAcpClient, type SessionMetadata, type SessionMetadataWithThread} from "./CodexAcpClient"; +import type {McpStartupResult} from "./CodexAppServerClient"; import {ACPSessionConnection, type UpdateSessionEvent} from "./ACPSessionConnection"; -import type {McpStartupCompleteEvent, InputModality, ReasoningEffort} from "./app-server"; +import type {InputModality, ReasoningEffort} from "./app-server"; import type { Account, CollabAgentToolCallStatus, @@ -55,6 +56,7 @@ export interface SessionState { interface PendingMcpStartupSession { requestedServers: Set; + afterVersion: number; } export class CodexAcpServer implements acp.Agent { @@ -143,7 +145,10 @@ export class CodexAcpServer implements acp.Agent { async getOrCreateSession(request: acp.NewSessionRequest | acp.ResumeSessionRequest): Promise<[SessionId, SessionModelState, SessionModeState]> { await this.checkAuthorization(); - const mcpStartupVersion = this.codexAcpClient.getMcpStartupCompleteVersion(); + const requestedMcpServers = request.mcpServers ?? []; + const mcpServerStartupVersion = requestedMcpServers.length > 0 + ? this.codexAcpClient.getMcpServerStartupVersion() + : null; let sessionMetadata: SessionMetadata; if ("sessionId" in request) { @@ -156,7 +161,7 @@ export class CodexAcpServer implements acp.Agent { const accountResponse = await this.runWithProcessCheck(() => this.codexAcpClient.getAccount()); const {sessionId, currentModelId, models} = sessionMetadata; - const sessionMcpServers = await this.resolveSessionMcpServers(request.mcpServers ?? [], mcpStartupVersion, "sessionId" in request); + const sessionMcpServers = this.resolveSessionMcpServers(requestedMcpServers, "sessionId" in request); const currentModel = this.findCurrentModel(models, currentModelId); const sessionState: SessionState = { sessionId: sessionId, @@ -175,12 +180,12 @@ export class CodexAcpServer implements acp.Agent { } this.sessions.set(sessionId, sessionState); - const requestedMcpServers = request.mcpServers ?? []; - if (requestedMcpServers.length > 0) { + if (requestedMcpServers.length > 0 && mcpServerStartupVersion !== null) { this.pendingMcpStartupSessions.set(sessionId, { requestedServers: new Set(requestedMcpServers.map(server => server.name)), + afterVersion: mcpServerStartupVersion, }); - this.publishMcpStartupStatusAsync(sessionId, mcpStartupVersion); + this.publishMcpStartupStatusAsync(sessionId); } this.publishAvailableCommandsAsync(sessionId); @@ -355,7 +360,10 @@ export class CodexAcpServer implements acp.Agent { thread: Thread; }> { await this.checkAuthorization(); - const mcpStartupVersion = this.codexAcpClient.getMcpStartupCompleteVersion(); + const requestedMcpServers = request.mcpServers ?? []; + const mcpServerStartupVersion = requestedMcpServers.length > 0 + ? this.codexAcpClient.getMcpServerStartupVersion() + : null; logger.log(`Load existing session: ${request.sessionId}...`); const sessionMetadata: SessionMetadataWithThread = await this.runWithProcessCheck(() => @@ -364,7 +372,7 @@ export class CodexAcpServer implements acp.Agent { const accountResponse = await this.runWithProcessCheck(() => this.codexAcpClient.getAccount()); const {sessionId, currentModelId, models, thread} = sessionMetadata; - const sessionMcpServers = await this.resolveSessionMcpServers(request.mcpServers ?? [], mcpStartupVersion, true); + const sessionMcpServers = this.resolveSessionMcpServers(requestedMcpServers, true); const currentModel = this.findCurrentModel(models, currentModelId); const sessionState: SessionState = { sessionId: sessionId, @@ -383,12 +391,12 @@ export class CodexAcpServer implements acp.Agent { }; this.sessions.set(sessionId, sessionState); - const requestedMcpServers = request.mcpServers ?? []; - if (requestedMcpServers.length > 0) { + if (requestedMcpServers.length > 0 && mcpServerStartupVersion !== null) { this.pendingMcpStartupSessions.set(sessionId, { requestedServers: new Set(requestedMcpServers.map(server => server.name)), + afterVersion: mcpServerStartupVersion, }); - this.publishMcpStartupStatusAsync(sessionId, mcpStartupVersion); + this.publishMcpStartupStatusAsync(sessionId); } await this.availableCommands.publish(sessionId); @@ -633,11 +641,10 @@ export class CodexAcpServer implements acp.Agent { return sessionState; } - private async resolveSessionMcpServers( + private resolveSessionMcpServers( mcpServers: Array, - mcpStartupVersion: number, recoverFromStartup: boolean, - ): Promise> { + ): Array { // Explicit MCP servers from the request are the primary source of truth for the session. const requestedServerNames = getRequestedMcpServerNames(mcpServers); if (requestedServerNames.length > 0) { @@ -647,27 +654,31 @@ export class CodexAcpServer implements acp.Agent { if (!recoverFromStartup) { return []; } - // loadSession/resumeSession may omit mcpServers; in that case recover the ready names - // from the startup event associated with this thread start/resume checkpoint. - logger.log("Recovering MCP servers from startup state..."); - return await this.runWithProcessCheck(() => this.codexAcpClient.awaitMcpStartup(mcpStartupVersion)); + // Without a thread-scoped startup completion event, loadSession/resumeSession can no longer + // recover omitted session MCP server names. Treat the session set as unknown unless ACP + // explicitly provided mcpServers in the request. + logger.log("Skipping MCP server recovery for load/resume without explicit mcpServers"); + return []; } - private publishMcpStartupStatusAsync(sessionId: string, mcpStartupVersion: number): void { - void this.doPublishMcpStartupStatus(sessionId, mcpStartupVersion); + private publishMcpStartupStatusAsync(sessionId: string): void { + void this.doPublishMcpStartupStatus(sessionId); } - private async doPublishMcpStartupStatus(sessionId: string, mcpStartupVersion: number): Promise { + private async doPublishMcpStartupStatus(sessionId: string): Promise { + const pendingStartup = this.pendingMcpStartupSessions.get(sessionId); + if (!pendingStartup) { + return; + } + try { - const mcpStartup = await this.runWithProcessCheck(() => this.codexAcpClient.awaitMcpStartupResult(mcpStartupVersion)); - 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); + const mcpStartup = await this.runWithProcessCheck(() => + this.codexAcpClient.awaitMcpServerStartup( + Array.from(pendingStartup.requestedServers), + pendingStartup.afterVersion, + ) + ); + await this.publishMcpStartupStatus(sessionId, mcpStartup, pendingStartup.requestedServers); } catch (err) { logger.error(`Failed to publish MCP startup status for session ${sessionId}`, err); } finally { @@ -677,7 +688,7 @@ export class CodexAcpServer implements acp.Agent { private async publishMcpStartupStatus( sessionId: string, - mcpStartup: McpStartupCompleteEvent, + mcpStartup: McpStartupResult, requestedServers?: Set ): Promise { const filteredStartup = requestedServers diff --git a/src/CodexAppServerClient.ts b/src/CodexAppServerClient.ts index 376332af..afd333ff 100644 --- a/src/CodexAppServerClient.ts +++ b/src/CodexAppServerClient.ts @@ -3,22 +3,37 @@ import type { ClientRequest, InitializeParams, InitializeResponse, - McpStartupCompleteEvent, ServerNotification } from "./app-server"; import type { AccountLoginCompletedNotification, AccountUpdatedNotification, + ConfigReadParams, + ConfigReadResponse, GetAccountParams, - GetAccountResponse, LoginAccountParams, LoginAccountResponse, LogoutAccountResponse, ModelListParams, + GetAccountResponse, + ListMcpServerStatusParams, + ListMcpServerStatusResponse, + LoginAccountParams, + LoginAccountResponse, + LogoutAccountResponse, + McpServerElicitationRequestParams, + McpServerElicitationRequestResponse, + McpServerStartupState, + McpServerStatusUpdatedNotification, + ModelListParams, ModelListResponse, - ThreadStartParams, - ThreadStartResponse, + SkillsListParams, + SkillsListResponse, ThreadLoadedListParams, ThreadLoadedListResponse, ThreadListParams, ThreadListResponse, ThreadReadParams, ThreadReadResponse, + ThreadResumeParams, + ThreadResumeResponse, + ThreadStartParams, + ThreadStartResponse, TurnCompletedNotification, TurnInterruptParams, TurnInterruptResponse, @@ -28,14 +43,6 @@ import type { CommandExecutionRequestApprovalResponse, FileChangeRequestApprovalParams, FileChangeRequestApprovalResponse, - McpServerElicitationRequestParams, - McpServerElicitationRequestResponse, - ThreadResumeParams, - ThreadResumeResponse, - SkillsListParams, - SkillsListResponse, - ListMcpServerStatusParams, - ListMcpServerStatusResponse, ConfigReadParams, ConfigReadResponse, } from "./app-server/v2"; export interface ApprovalHandler { @@ -47,6 +54,17 @@ export interface ElicitationHandler { handleElicitation(params: McpServerElicitationRequestParams): Promise; } +export type McpStartupFailure = { + server: string; + error: string; +}; + +export type McpStartupResult = { + ready: Array; + failed: Array; + cancelled: Array; +}; + const CommandExecutionApprovalRequest = new RequestType< CommandExecutionRequestApprovalParams, CommandExecutionRequestApprovalResponse, @@ -73,23 +91,23 @@ export class CodexAppServerClient { readonly connection: MessageConnection; private approvalHandlers = new Map(); private elicitationHandlers = new Map(); - private mcpStartupCompleteVersion = 0; - private lastMcpStartupComplete: McpStartupCompleteEvent | null = null; - private readonly mcpStartupCompleteResolvers: Array> = []; + private mcpServerStartupVersion = 0; + private readonly mcpServerStartupStates = new Map(); + private readonly mcpServerStartupResolvers: Array = []; constructor(connection: MessageConnection) { this.connection = connection; this.connection.onUnhandledNotification((data) => { - if (isMcpStartupCompleteNotification(data)) { - this.mcpStartupCompleteVersion += 1; - this.lastMcpStartupComplete = data.params.msg; - this.resolveSignal(data.params.msg, this.mcpStartupCompleteVersion, this.mcpStartupCompleteResolvers); - for (const callback of this.codexEventHandlers) { - callback({ eventType: "notification", ...data }); - } - return; - } const serverNotification = data as ServerNotification; + if (isMcpServerStatusUpdatedNotification(serverNotification)) { + this.mcpServerStartupVersion += 1; + this.mcpServerStartupStates.set(serverNotification.params.name, { + status: serverNotification.params.status, + error: serverNotification.params.error, + version: this.mcpServerStartupVersion, + }); + this.resolveMcpServerStartupResolvers(); + } this.notify(serverNotification); for (const callback of this.codexEventHandlers) { callback({ eventType: "notification", ...serverNotification }); @@ -177,17 +195,28 @@ export class CodexAppServerClient { return await this.sendRequest({ method: "config/read", params: params }); } - getMcpStartupCompleteVersion(): number { - return this.mcpStartupCompleteVersion; + getMcpServerStartupVersion(): number { + return this.mcpServerStartupVersion; } - async awaitMcpStartup(afterVersion: number): Promise { - return await this.awaitSignal( - this.lastMcpStartupComplete, - this.mcpStartupCompleteVersion, - afterVersion, - this.mcpStartupCompleteResolvers - ); + async awaitMcpServerStartup(serverNames: Array, afterVersion: number): Promise { + const uniqueServerNames = Array.from(new Set(serverNames.map(serverName => serverName.trim()).filter(serverName => serverName.length > 0))); + if (uniqueServerNames.length === 0) { + return { ready: [], failed: [], cancelled: [] }; + } + + const result = this.tryBuildMcpStartupResult(uniqueServerNames, afterVersion); + if (result !== null) { + return result; + } + + return await new Promise((resolve) => { + this.mcpServerStartupResolvers.push({ + serverNames: uniqueServerNames, + afterVersion, + resolve, + }); + }); } async accountRead(params: GetAccountParams): Promise { @@ -231,34 +260,49 @@ export class CodexAppServerClient { } } - private resolveSignal( - event: T, - version: number, - resolvers: Array> - ): void { - const pendingResolvers: Array> = []; - for (const resolver of resolvers) { - if (resolver.afterVersion < version) { - resolver.resolve(event); + private resolveMcpServerStartupResolvers(): void { + const pendingResolvers: Array = []; + for (const resolver of this.mcpServerStartupResolvers) { + const result = this.tryBuildMcpStartupResult(resolver.serverNames, resolver.afterVersion); + if (result !== null) { + resolver.resolve(result); } else { pendingResolvers.push(resolver); } } - resolvers.splice(0, resolvers.length, ...pendingResolvers); + this.mcpServerStartupResolvers.splice(0, this.mcpServerStartupResolvers.length, ...pendingResolvers); } - private async awaitSignal( - lastEvent: T | null, - currentVersion: number, - afterVersion: number, - resolvers: Array> - ): Promise { - if (lastEvent !== null && currentVersion > afterVersion) { - return lastEvent; + private tryBuildMcpStartupResult(serverNames: Array, afterVersion: number): McpStartupResult | null { + const ready: Array = []; + const failed: Array = []; + const cancelled: Array = []; + + for (const serverName of serverNames) { + const state = this.mcpServerStartupStates.get(serverName); + if (!state || state.version <= afterVersion) { + return null; + } + + switch (state.status) { + case "starting": + return null; + case "ready": + ready.push(serverName); + break; + case "failed": + failed.push({ + server: serverName, + error: state.error ?? "unknown MCP startup error", + }); + break; + case "cancelled": + cancelled.push(serverName); + break; + } } - return await new Promise((resolve) => { - resolvers.push({afterVersion, resolve}); - }); + + return { ready, failed, cancelled }; } private async sendRequest(request: CodexRequest): Promise { @@ -282,7 +326,7 @@ export class CodexAppServerClient { export type CodexConnectionEvent = | ({ eventType: "request" } & CodexRequest) | ({ eventType: "response" } & unknown) - | ({ eventType: "notification" } & (ServerNotification | McpStartupCompleteNotification)); + | ({ eventType: "notification" } & ServerNotification); type CodexRequest = DistributiveOmit @@ -290,21 +334,21 @@ type DistributiveOmit = T extends any ? Omit : never; -type SignalResolver = { - afterVersion: number; - resolve: (event: T) => void; +type McpServerStartupSnapshot = { + status: McpServerStartupState; + error: string | null; + version: number; }; -type McpStartupCompleteNotification = { - method: "codex/event/mcp_startup_complete", - params: { - msg: McpStartupCompleteEvent & { type: "mcp_startup_complete" } - } +type McpServerStartupResolver = { + serverNames: Array; + afterVersion: number; + resolve: (result: McpStartupResult) => void; }; -function isMcpStartupCompleteNotification(value: unknown): value is McpStartupCompleteNotification { - return typeof value === "object" - && value !== null - && "method" in value - && value.method === "codex/event/mcp_startup_complete"; +function isMcpServerStatusUpdatedNotification(notification: ServerNotification): notification is { + method: "mcpServer/startupStatus/updated"; + params: McpServerStatusUpdatedNotification; +} { + return notification.method === "mcpServer/startupStatus/updated"; } diff --git a/src/__tests__/CodexACPAgent/CodexAcpClient.test.ts b/src/__tests__/CodexACPAgent/CodexAcpClient.test.ts index 38b4ce01..e3622324 100644 --- a/src/__tests__/CodexACPAgent/CodexAcpClient.test.ts +++ b/src/__tests__/CodexACPAgent/CodexAcpClient.test.ts @@ -191,28 +191,40 @@ describe('ACP server test', { timeout: 40_000 }, () => { expect(listSkillsSpy.mock.invocationCallOrder[0]!).toBeLessThan(threadStartSpy.mock.invocationCallOrder[0]!); }); - it('waits for mcp startup completion and returns ready servers', async () => { + it('waits for typed mcp startup status updates and returns terminal states', async () => { const mockFixture = createCodexMockTestFixture(); const codexAcpClient = mockFixture.getCodexAcpClient(); const codexAppServerClient = mockFixture.getCodexAppServerClient(); - const startupPromise = codexAcpClient.awaitMcpStartup(codexAppServerClient.getMcpStartupCompleteVersion()); + const startupPromise = codexAcpClient.awaitMcpServerStartup( + ["alpha", "beta"], + codexAppServerClient.getMcpServerStartupVersion() + ); mockFixture.sendServerNotification({ - method: "codex/event/mcp_startup_complete", - params: { - msg: { - type: "mcp_startup_complete", - ready: ["alpha", "beta"], - failed: [], - cancelled: [], - } - } + method: "mcpServer/startupStatus/updated", + params: { name: "alpha", status: "starting", error: null } + }); + mockFixture.sendServerNotification({ + method: "mcpServer/startupStatus/updated", + params: { name: "beta", status: "starting", error: null } + }); + mockFixture.sendServerNotification({ + method: "mcpServer/startupStatus/updated", + params: { name: "alpha", status: "ready", error: null } + }); + mockFixture.sendServerNotification({ + method: "mcpServer/startupStatus/updated", + params: { name: "beta", status: "ready", error: null } }); - const mcpServers = await startupPromise; + const startup = await startupPromise; - expect(mcpServers).toEqual(["alpha", "beta"]); + expect(startup).toEqual({ + ready: ["alpha", "beta"], + failed: [], + cancelled: [], + }); }); it('forwards failed MCP startup as failed tool call updates after new session', async () => { @@ -253,18 +265,8 @@ describe('ACP server test', { timeout: 40_000 }, () => { }); mockFixture.sendServerNotification({ - method: "codex/event/mcp_startup_complete", - params: { - msg: { - type: "mcp_startup_complete", - ready: [], - failed: [{ - server: "broken-mcp", - error: "boom", - }], - cancelled: [], - } - } + method: "mcpServer/startupStatus/updated", + params: { name: "broken-mcp", status: "failed", error: "boom" } }); await vi.waitFor(() => { diff --git a/src/__tests__/CodexACPAgent/load-session.test.ts b/src/__tests__/CodexACPAgent/load-session.test.ts index c19b955c..cb401746 100644 --- a/src/__tests__/CodexACPAgent/load-session.test.ts +++ b/src/__tests__/CodexACPAgent/load-session.test.ts @@ -15,7 +15,6 @@ describe("CodexACPAgent - loadSession", () => { account: null, requiresOpenaiAuth: false, }); - codexAcpClient.awaitMcpStartup = vi.fn().mockResolvedValue([]); codexAcpClient.listSkills = vi.fn().mockResolvedValue({ data: [] }); const model: Model = { @@ -169,7 +168,7 @@ describe("CodexACPAgent - loadSession", () => { ); }); - it("should recover session mcp servers during loadSession when request omits them", async () => { + it("should not recover session mcp servers during loadSession when request omits them", async () => { const fixture = createCodexMockTestFixture(); const codexAcpAgent = fixture.getCodexAcpAgent(); const codexAcpClient = fixture.getCodexAcpClient(); @@ -182,8 +181,6 @@ describe("CodexACPAgent - loadSession", () => { }); codexAcpClient.listSkills = vi.fn().mockResolvedValue({ data: [] }); - const awaitMcpStartupSpy = vi.spyOn(codexAcpClient, "awaitMcpStartup").mockResolvedValue(["persisted-mcp"]); - const model: Model = { id: "gpt-5.2", model: "gpt-5.2", @@ -240,8 +237,7 @@ describe("CodexACPAgent - loadSession", () => { mcpServers: [], }); - expect(awaitMcpStartupSpy).toHaveBeenCalledWith(0); - expect(codexAcpAgent.getSessionState("session-1").sessionMcpServers).toEqual(["persisted-mcp"]); + expect(codexAcpAgent.getSessionState("session-1").sessionMcpServers).toEqual([]); }); it("publishes MCP startup failure for explicitly requested servers during loadSession", async () => { @@ -324,18 +320,8 @@ describe("CodexACPAgent - loadSession", () => { }); fixture.sendServerNotification({ - method: "codex/event/mcp_startup_complete", - params: { - msg: { - type: "mcp_startup_complete", - ready: [], - failed: [{ - server: "broken-mcp", - error: "boom", - }], - cancelled: [], - } - } + method: "mcpServer/startupStatus/updated", + params: { name: "broken-mcp", status: "failed", error: "boom" } }); await loadPromise;