Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions src/CodexAcpClient.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import open from "open";
import type {Disposable} from "vscode-jsonrpc";
import type {
ClientInfo,
McpStartupCompleteEvent,
ReasoningEffort,
ServerNotification
} from "./app-server";
Expand All @@ -21,6 +22,7 @@ import type {
GetAccountResponse,
ListMcpServerStatusParams,
ListMcpServerStatusResponse,
McpServerStatusUpdatedNotification,
Model,
SkillsListParams,
SkillsListResponse,
Expand Down Expand Up @@ -275,6 +277,22 @@ export class CodexAcpClient {
return startup.ready;
}

async awaitMcpStartupResult(mcpStartupVersion: number): Promise<McpStartupCompleteEvent> {
return await this.codexClient.awaitMcpStartup(mcpStartupVersion);
}

onMcpServerStatusUpdated(handler: (event: McpServerStatusUpdatedNotification) => void): void {
this.codexClient.onMcpServerStatusUpdated(handler);
}

getMcpServerStatusVersion(): number {
return this.codexClient.getMcpServerStatusVersion();
}

getMcpServerStatusUpdates(afterVersion: number): Array<McpServerStatusUpdatedNotification> {
return this.codexClient.getMcpServerStatusUpdates(afterVersion);
}

getMcpStartupCompleteVersion(): number {
return this.codexClient.getMcpStartupCompleteVersion();
}
Expand Down
110 changes: 108 additions & 2 deletions src/CodexAcpServer.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,18 @@ import {CodexApprovalHandler} from "./CodexApprovalHandler";
import {CodexAuthMethods, type CodexAuthRequest} from "./CodexAuthMethod";
import {CodexAcpClient, type SessionMetadata, type SessionMetadataWithThread} from "./CodexAcpClient";
import {ACPSessionConnection, type UpdateSessionEvent} from "./ACPSessionConnection";
import type {Account, CollabAgentToolCallStatus, Model, Thread, ThreadItem, UserInput, ReasoningEffortOption} from "./app-server/v2";
import type {McpStartupCompleteEvent, InputModality, ReasoningEffort} from "./app-server";
import type {
Account,
CollabAgentToolCallStatus,
McpServerStatusUpdatedNotification,
Model,
Thread,
ThreadItem,
UserInput,
ReasoningEffortOption
} from "./app-server/v2";
import type {RateLimitsMap} from "./RateLimitsMap";
import type {InputModality, ReasoningEffort} from "./app-server";
import {ModelId} from "./ModelId";
import {AgentMode} from "./AgentMode";
import type {TokenCount} from "./TokenCount";
Expand Down Expand Up @@ -43,6 +52,12 @@ export interface SessionState {
sessionMcpServers?: Array<string>;
}

interface PendingMcpStartupSession {
requestedServers: Set<string>;
terminalServers: Set<string>;
publishedFailures: Set<string>;
}

export class CodexAcpServer implements acp.Agent {
private readonly codexAcpClient: CodexAcpClient;
private readonly connection: acp.AgentSideConnection;
Expand All @@ -51,6 +66,7 @@ export class CodexAcpServer implements acp.Agent {
private readonly availableCommands: CodexCommands;

private readonly sessions: Map<string, SessionState>;
private readonly pendingMcpStartupSessions: Map<string, PendingMcpStartupSession>;

constructor(
connection: acp.AgentSideConnection,
Expand All @@ -59,6 +75,7 @@ export class CodexAcpServer implements acp.Agent {
getExitCode?: () => number | null,
) {
this.sessions = new Map();
this.pendingMcpStartupSessions = new Map();
this.connection = connection;
this.codexAcpClient = codexAcpClient;
this.defaultAuthRequest = defaultAuthRequest ?? null;
Expand All @@ -68,6 +85,9 @@ export class CodexAcpServer implements acp.Agent {
codexAcpClient,
(operation) => this.runWithProcessCheck(operation)
);
this.codexAcpClient.onMcpServerStatusUpdated((event) => {
void this.handleMcpServerStatusUpdated(event);
});
}

async initialize(
Expand Down Expand Up @@ -128,6 +148,7 @@ 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 mcpStatusVersion = this.codexAcpClient.getMcpServerStatusVersion();

let sessionMetadata: SessionMetadata;
if ("sessionId" in request) {
Expand Down Expand Up @@ -159,6 +180,17 @@ export class CodexAcpServer implements acp.Agent {
}
this.sessions.set(sessionId, sessionState);

const requestedMcpServers = request.mcpServers ?? [];
if (requestedMcpServers.length > 0) {
this.pendingMcpStartupSessions.set(sessionId, {
requestedServers: new Set(requestedMcpServers.map(server => server.name)),
terminalServers: new Set(),
publishedFailures: new Set(),
});
await this.replayBufferedMcpStartupStatus(sessionId, mcpStatusVersion);
this.publishMcpStartupStatusAsync(sessionId, mcpStartupVersion);
}

this.publishAvailableCommandsAsync(sessionId);
const sessionModelState: SessionModelState = this.createModelState(models, currentModelId);
const sessionModeState: SessionModeState = sessionState.agentMode.toSessionModeState();
Expand Down Expand Up @@ -621,6 +653,80 @@ export class CodexAcpServer implements acp.Agent {
return await this.runWithProcessCheck(() => this.codexAcpClient.awaitMcpStartup(mcpStartupVersion));
}

private publishMcpStartupStatusAsync(sessionId: string, mcpStartupVersion: number): void {
void (async () => {
try {
const mcpStartup = await this.runWithProcessCheck(() => this.codexAcpClient.awaitMcpStartupResult(mcpStartupVersion));
const sessionState = this.sessions.get(sessionId);
if (sessionState) {
sessionState.sessionMcpServers = mcpStartup.ready;
}
await this.publishMcpStartupStatus(sessionId, mcpStartup);
this.pendingMcpStartupSessions.delete(sessionId);
} catch (err) {
logger.error(`Failed to publish MCP startup status for session ${sessionId}`, err);
}
})();
}

private async publishMcpStartupStatus(sessionId: string, mcpStartup: McpStartupCompleteEvent): Promise<void> {
for (const update of CodexEventHandler.createMcpStartupUpdates(mcpStartup)) {
await this.connection.sessionUpdate({
sessionId,
update,
});
}
}

private async handleMcpServerStatusUpdated(event: McpServerStatusUpdatedNotification): Promise<void> {
for (const [sessionId, pending] of this.pendingMcpStartupSessions) {
if (!pending.requestedServers.has(event.name)) {
continue;
}

if (event.status === "ready") {
pending.terminalServers.add(event.name);
const sessionState = this.sessions.get(sessionId);
if (sessionState && !(sessionState.sessionMcpServers ?? []).includes(event.name)) {
sessionState.sessionMcpServers = [...(sessionState.sessionMcpServers ?? []), event.name];
}
} else if (event.status === "failed" || event.status === "cancelled") {
pending.terminalServers.add(event.name);
if (!pending.publishedFailures.has(event.name)) {
pending.publishedFailures.add(event.name);
const updates = CodexEventHandler.createMcpStartupUpdates({
ready: [],
failed: event.status === "failed" ? [{
server: event.name,
error: event.error ?? "Unknown MCP startup error",
}] : [],
cancelled: event.status === "cancelled" ? [event.name] : [],
});
for (const update of updates) {
await this.connection.sessionUpdate({
sessionId,
update,
});
}
}
}

if (pending.terminalServers.size >= pending.requestedServers.size) {
this.pendingMcpStartupSessions.delete(sessionId);
}
}
}

private async replayBufferedMcpStartupStatus(sessionId: string, afterVersion: number): Promise<void> {
const updates = this.codexAcpClient.getMcpServerStatusUpdates(afterVersion);
for (const update of updates) {
await this.handleMcpServerStatusUpdated(update);
if (!this.pendingMcpStartupSessions.has(sessionId)) {
return;
}
}
}

async prompt(params: acp.PromptRequest): Promise<acp.PromptResponse> {
logger.log("Prompt received", {
sessionId: params.sessionId,
Expand Down
36 changes: 35 additions & 1 deletion src/CodexAppServerClient.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import type {
GetAccountParams,
GetAccountResponse, LoginAccountParams, LoginAccountResponse, LogoutAccountResponse, ModelListParams,
ModelListResponse,
McpServerStatusUpdatedNotification,
ThreadStartParams,
ThreadStartResponse,
ThreadLoadedListParams,
Expand Down Expand Up @@ -63,6 +64,12 @@ export class CodexAppServerClient {
private mcpStartupCompleteVersion = 0;
private lastMcpStartupComplete: McpStartupCompleteEvent | null = null;
private readonly mcpStartupCompleteResolvers: Array<SignalResolver<McpStartupCompleteEvent>> = [];
private mcpServerStatusVersion = 0;
private readonly mcpServerStatusUpdatedHandlers: Array<(event: McpServerStatusUpdatedNotification) => void> = [];
private readonly mcpServerStatusHistory: Array<{
version: number;
event: McpServerStatusUpdatedNotification;
}> = [];

constructor(connection: MessageConnection) {
this.connection = connection;
Expand All @@ -76,8 +83,10 @@ export class CodexAppServerClient {
}
return;
}

const serverNotification = data as ServerNotification;
if (serverNotification.method === "mcpServer/startupStatus/updated") {
this.recordMcpServerStatusUpdated(serverNotification.params);
}
this.notify(serverNotification);
for (const callback of this.codexEventHandlers) {
callback({ eventType: "notification", ...serverNotification });
Expand Down Expand Up @@ -166,6 +175,20 @@ export class CodexAppServerClient {
);
}

onMcpServerStatusUpdated(handler: (event: McpServerStatusUpdatedNotification) => void): void {
this.mcpServerStatusUpdatedHandlers.push(handler);
}

getMcpServerStatusVersion(): number {
return this.mcpServerStatusVersion;
}

getMcpServerStatusUpdates(afterVersion: number): Array<McpServerStatusUpdatedNotification> {
return this.mcpServerStatusHistory
.filter(entry => entry.version > afterVersion)
.map(entry => entry.event);
}

async accountRead(params: GetAccountParams): Promise<GetAccountResponse> {
return await this.sendRequest({ method: "account/read", params: params });
}
Expand Down Expand Up @@ -237,6 +260,17 @@ export class CodexAppServerClient {
});
}

private recordMcpServerStatusUpdated(event: McpServerStatusUpdatedNotification): void {
this.mcpServerStatusVersion += 1;
this.mcpServerStatusHistory.push({
version: this.mcpServerStatusVersion,
event,
});
for (const handler of this.mcpServerStatusUpdatedHandlers) {
handler(event);
}
}

private async sendRequest<R>(request: CodexRequest): Promise<R> {
for (const callback of this.codexEventHandlers) {
callback({ eventType: "request", ...request});
Expand Down
35 changes: 35 additions & 0 deletions src/CodexEventHandler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ import type {
ThreadTokenUsageUpdatedNotification,
TurnPlanUpdatedNotification
} from "./app-server/v2";
import type { McpStartupCompleteEvent } from "./app-server";
import {toTokenCount} from "./TokenCount";
import {
createCommandExecutionUpdate,
Expand Down Expand Up @@ -258,6 +259,40 @@ export class CodexEventHandler {
}
}

static createMcpStartupUpdates(event: McpStartupCompleteEvent): UpdateSessionEvent[] {
const failedUpdates = event.failed.map((server: McpStartupCompleteEvent["failed"][number]) => this.createMcpStartupToolCallUpdate(
server.server,
`[codex-acp forwarded startup error] MCP server \`${server.server}\` failed to start: ${server.error}`
));
const cancelledUpdates = event.cancelled.map((server: McpStartupCompleteEvent["cancelled"][number]) => this.createMcpStartupToolCallUpdate(
server,
`[codex-acp forwarded startup error] MCP server \`${server}\` startup was cancelled.`
));

return [...failedUpdates, ...cancelledUpdates];
}

private static createMcpStartupToolCallUpdate(serverName: string, message: string): UpdateSessionEvent {
return {
sessionUpdate: "tool_call",
toolCallId: this.getMcpStartupToolCallId(serverName),
kind: "other",
title: `mcp__${serverName}__startup`,
status: "failed",
content: [{
type: "content",
content: {
type: "text",
text: message,
},
}],
};
}

private static getMcpStartupToolCallId(serverName: string): string {
return `mcp_startup.${encodeURIComponent(serverName)}`;
}

private completeCommandExecutionEvent(item: ThreadItem & { "type": "commandExecution" }): UpdateSessionEvent {
return {
sessionUpdate: "tool_call_update",
Expand Down
Loading
Loading