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;