Skip to content
Merged
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 {
Comment thread
NikolaiSviridov marked this conversation as resolved.
Outdated
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
74 changes: 72 additions & 2 deletions src/CodexAcpServer.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,17 @@ 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,
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 +51,10 @@ export interface SessionState {
sessionMcpServers?: Array<string>;
}

interface PendingMcpStartupSession {
requestedServers: Set<string>;
}

export class CodexAcpServer implements acp.Agent {
private readonly codexAcpClient: CodexAcpClient;
private readonly connection: acp.AgentSideConnection;
Expand All @@ -51,6 +63,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 +72,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 Down Expand Up @@ -159,6 +173,14 @@ 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)),
});
this.publishMcpStartupStatusAsync(sessionId, mcpStartupVersion);
}

this.publishAvailableCommandsAsync(sessionId);
const sessionModelState: SessionModelState = this.createModelState(models, currentModelId);
const sessionModeState: SessionModeState = sessionState.agentMode.toSessionModeState();
Expand Down Expand Up @@ -359,6 +381,14 @@ 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)),
});
this.publishMcpStartupStatusAsync(sessionId, mcpStartupVersion);
Comment thread
NikolaiSviridov marked this conversation as resolved.
}

await this.availableCommands.publish(sessionId);
const sessionModelState: SessionModelState = this.createModelState(models, currentModelId);
const sessionModeState: SessionModeState = sessionState.agentMode.toSessionModeState();
Expand Down Expand Up @@ -621,6 +651,46 @@ 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);
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);
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,
requestedServers?: Set<string>
): Promise<void> {
const filteredStartup = requestedServers
? {
ready: mcpStartup.ready.filter(server => requestedServers.has(server)),
failed: mcpStartup.failed.filter(server => requestedServers.has(server.server)),
cancelled: mcpStartup.cancelled.filter(server => requestedServers.has(server)),
}
: mcpStartup;

for (const update of CodexEventHandler.createMcpStartupUpdates(filteredStartup)) {
await this.connection.sessionUpdate({
sessionId,
update,
});
}
}

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") {
Comment thread
NikolaiSviridov marked this conversation as resolved.
Outdated
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
64 changes: 61 additions & 3 deletions src/CodexEventHandler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,11 +20,14 @@ import type {
ThreadTokenUsageUpdatedNotification,
TurnPlanUpdatedNotification
} from "./app-server/v2";
import type { McpStartupCompleteEvent } from "./app-server";
import {toTokenCount} from "./TokenCount";
import {
createCommandExecutionUpdate,
createDynamicToolCallUpdate,
createFileChangeUpdate,
createMcpRawInput,
createMcpRawOutput,
createFuzzyFileSearchComplete,
createFuzzyFileSearchStartOrUpdate,
createMcpToolCallUpdate,
Expand Down Expand Up @@ -101,12 +104,13 @@ export class CodexEventHandler {
case "turn/diff/updated":
case "item/commandExecution/terminalInteraction":
case "item/fileChange/outputDelta":
case "item/mcpToolCall/progress":
case "serverRequest/resolved":
case "account/updated":
case "fs/changed":
case "mcpServer/startupStatus/updated":
return null;
case "item/mcpToolCall/progress":
return this.createMcpToolProgressEvent(notification.params);
case "account/rateLimits/updated":
this.handleRateLimitsUpdated(notification.params);
return null;
Expand Down Expand Up @@ -210,13 +214,20 @@ export class CodexEventHandler {

private async completeItemEvent(event: ItemCompletedNotification): Promise<UpdateSessionEvent | null> {
switch (event.item.type) {
case "mcpToolCall":
case "fileChange":
case "dynamicToolCall":
return {
sessionUpdate: "tool_call_update",
toolCallId: event.item.id,
status: event.item.status === "completed" ? "completed" : "failed"
status: event.item.status === "completed" ? "completed" : "failed",
}
case "mcpToolCall":
return {
sessionUpdate: "tool_call_update",
toolCallId: event.item.id,
status: event.item.status === "completed" ? "completed" : "failed",
rawInput: createMcpRawInput(event.item.server, event.item.tool, event.item.arguments),
rawOutput: createMcpRawOutput(event.item.result, event.item.error),
}
case "commandExecution":
return this.completeCommandExecutionEvent(event.item);
Expand Down Expand Up @@ -258,6 +269,53 @@ export class CodexEventHandler {
}
}

private createMcpToolProgressEvent(event: { itemId: string, message: string }): UpdateSessionEvent {
const logDelta = event.message.trim();
return {
sessionUpdate: "tool_call_update",
toolCallId: event.itemId,
_meta: {
mcp_output_delta: {
data: logDelta,
}
}
};
}

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