Inject agent-control MCP via codex wrapper

packages/server/src/server/agent/providers/codex-agent.test.ts

@@ -1,8 +1,13 @@
 import { describe, expect, test } from "vitest";
+import { createServer as createHTTPServer } from "http";
+import { randomUUID } from "node:crypto";
 import { mkdtempSync, rmSync, promises as fs } from "node:fs";
 import type { Dirent } from "node:fs";
 import os from "node:os";
 import path from "node:path";
+import express from "express";
+import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
+import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
 
 import { CodexAgentClient, isSyntheticRolloutUserMessage } from "./codex-agent.js";
 import {
@@ -23,6 +28,9 @@ import type {
   AgentPersistenceHandle,
   AgentTimelineItem,
 } from "../agent-sdk-types.js";
+import { AgentManager } from "../agent-manager.js";
+import { AgentRegistry } from "../agent-registry.js";
+import { createAgentMcpServer } from "../mcp-server.js";
 
 function tmpCwd(): string {
   return mkdtempSync(path.join(os.tmpdir(), "codex-agent-e2e-"));
@@ -51,6 +59,130 @@ function log(message: string): void {
 
 type ToolCallItem = Extract<AgentTimelineItem, { type: "tool_call" }>;
 
+type AgentMcpServerHandle = {
+  url: string;
+  close: () => Promise<void>;
+};
+
+async function startAgentMcpServer(): Promise<AgentMcpServerHandle> {
+  const app = express();
+  app.use(express.json());
+  const httpServer = createHTTPServer(app);
+
+  const registryDir = await fs.mkdtemp(path.join(os.tmpdir(), "agent-mcp-registry-"));
+  const registryPath = path.join(registryDir, "agents.json");
+  const agentRegistry = new AgentRegistry(registryPath);
+  const agentManager = new AgentManager({
+    clients: {},
+    registry: agentRegistry,
+  });
+
+  let allowedHosts: string[] | undefined;
+  const agentMcpTransports = new Map<string, StreamableHTTPServerTransport>();
+
+  const createAgentMcpTransport = async (callerAgentId?: string) => {
+    const agentMcpServer = await createAgentMcpServer({
+      agentManager,
+      agentRegistry,
+      callerAgentId,
+    });
+
+    const transport = new StreamableHTTPServerTransport({
+      sessionIdGenerator: () => randomUUID(),
+      onsessioninitialized: (sessionId) => {
+        agentMcpTransports.set(sessionId, transport);
+        log(`Agent MCP session initialized: ${sessionId}`);
+      },
+      onsessionclosed: (sessionId) => {
+        agentMcpTransports.delete(sessionId);
+        log(`Agent MCP session closed: ${sessionId}`);
+      },
+      enableDnsRebindingProtection: true,
+      ...(allowedHosts ? { allowedHosts } : {}),
+    });
+
+    transport.onclose = () => {
+      if (transport.sessionId) {
+        agentMcpTransports.delete(transport.sessionId);
+      }
+    };
+    transport.onerror = (error) => {
+      console.error("[Agent MCP] Transport error:", error);
+    };
+
+    await agentMcpServer.connect(transport);
+    return transport;
+  };
+
+  const handleAgentMcpRequest: express.RequestHandler = async (req, res) => {
+    try {
+      const sessionId = req.header("mcp-session-id");
+      let transport = sessionId ? agentMcpTransports.get(sessionId) : undefined;
+
+      if (!transport) {
+        if (req.method !== "POST") {
+          res.status(400).json({
+            jsonrpc: "2.0",
+            error: { code: -32000, message: "Missing or invalid MCP session" },
+            id: null,
+          });
+          return;
+        }
+        if (!isInitializeRequest(req.body)) {
+          res.status(400).json({
+            jsonrpc: "2.0",
+            error: { code: -32000, message: "Initialization request expected" },
+            id: null,
+          });
+          return;
+        }
+        log("Agent MCP initialize request received");
+        const callerAgentIdRaw = req.query.callerAgentId;
+        const callerAgentId =
+          typeof callerAgentIdRaw === "string"
+            ? callerAgentIdRaw
+            : Array.isArray(callerAgentIdRaw)
+              ? callerAgentIdRaw[0]
+              : undefined;
+        transport = await createAgentMcpTransport(callerAgentId);
+      }
+
+      await transport.handleRequest(req as any, res as any, req.body);
+    } catch (error) {
+      console.error("[Agent MCP] Failed to handle request:", error);
+      if (!res.headersSent) {
+        res.status(500).json({
+          jsonrpc: "2.0",
+          error: { code: -32603, message: "Internal MCP server error" },
+          id: null,
+        });
+      }
+    }
+  };
+
+  app.post("/mcp/agents", handleAgentMcpRequest);
+  app.get("/mcp/agents", handleAgentMcpRequest);
+  app.delete("/mcp/agents", handleAgentMcpRequest);
+
+  const port = await new Promise<number>((resolve) => {
+    httpServer.listen(0, () => {
+      const address = httpServer.address();
+      resolve(typeof address === "object" && address ? address.port : 0);
+    });
+  });
+
+  allowedHosts = [`127.0.0.1:${port}`, `localhost:${port}`];
+  const url = `http://127.0.0.1:${port}/mcp/agents`;
+
+  return {
+    url,
+    close: async () => {
+      await new Promise<void>((resolve) => httpServer.close(() => resolve()));
+      rmSync(registryDir, { recursive: true, force: true });
+    },
+  };
+}
+
 function commandTextFromInput(input: unknown): string | null {
   if (!input || typeof input !== "object") {
     return null;
@@ -386,7 +518,7 @@ describe("CodexAgentClient (SDK integration)", () => {
       const cwd = tmpCwd();
       const restoreSessionDir = useTempCodexSessionDir();
       const client = new CodexAgentClient();
-      const config: AgentSessionConfig = { provider: "codex", cwd };
+      const config: AgentSessionConfig = { provider: "codex", cwd, modeId: "full-access" };
       let session: Awaited<ReturnType<typeof client.createSession>> | null = null;
       try {
         session = await client.createSession(config);
@@ -699,6 +831,74 @@ describe("CodexAgentClient (SDK integration)", () => {
     },
     180_000
   );
+
+  test(
+    "sees agent-control MCP tools (list_agents) via codex",
+    async () => {
+      const cwd = tmpCwd();
+      const restoreSessionDir = useTempCodexSessionDir();
+      const mcpServer = await startAgentMcpServer();
+      const client = new CodexAgentClient();
+      const config: AgentSessionConfig = {
+        provider: "codex",
+        cwd,
+        modeId: "full-access",
+        extra: { codex: { agentControlMcpUrl: mcpServer.url } },
+      };
+      let session: Awaited<ReturnType<typeof client.createSession>> | null = null;
+      try {
+        session = await client.createSession(config);
+        log(`Agent MCP URL: ${mcpServer.url}`);
+
+        const prompt = [
+          "Use the MCP tool agent_control.list_agents and report the agent IDs you receive.",
+          "If the tool is unavailable, say exactly: MCP tool list_agents unavailable.",
+          "Then stop.",
+        ].join("\n");
+
+        const toolCalls: ToolCallItem[] = [];
+        let finalText = "";
+        const eventLog: string[] = [];
+
+        for await (const event of session.stream(prompt)) {
+          eventLog.push(event.type);
+          if (event.type === "timeline" && event.provider === "codex") {
+            if (event.item.type === "tool_call") {
+              toolCalls.push(event.item);
+            } else if (event.item.type === "assistant_message") {
+              finalText += `${event.item.text}\n`;
+            }
+          }
+
+          if (event.type === "turn_failed") {
+            console.info("[CodexAgentTest] Turn failed:", event.error);
+          }
+
+          if (event.type === "turn_completed" || event.type === "turn_failed") {
+            break;
+          }
+        }
+
+        const listAgentsCall = toolCalls.find((call) => call.tool === "list_agents");
+        if (!listAgentsCall) {
+          console.info(
+            "[CodexAgentTest] MCP tool calls:",
+            toolCalls.map((call) => `${call.server}.${call.tool}`)
+          );
+          console.info("[CodexAgentTest] Final text:", finalText.trim());
+          console.info("[CodexAgentTest] Event types:", eventLog);
+        }
+
+        expect(listAgentsCall).toBeDefined();
+      } finally {
+        await session?.close();
+        await mcpServer.close();
+        rmSync(cwd, { recursive: true, force: true });
+        restoreSessionDir();
+      }
+    },
+    240_000
+  );
 });
 
 describe("isSyntheticRolloutUserMessage", () => {

packages/server/src/server/agent/providers/codex-agent.ts

@@ -42,6 +42,20 @@ import type {
 
 type CodexAgentConfig = AgentSessionConfig & { provider: "codex" };
 
+type CodexExtraConfig = {
+  agentControlMcpUrl?: string;
+};
+
+type CodexWrapperInfo = {
+  path: string;
+  dir: string;
+};
+
+type CodexSessionBootstrap = {
+  codex: Codex;
+  wrapper?: CodexWrapperInfo;
+};
+
 const CODEX_CAPABILITIES: AgentCapabilityFlags = {
   supportsStreaming: true,
   supportsSessionPersistence: true,
@@ -252,10 +266,77 @@ function detectSystemCodexPath(): string | undefined {
   return undefined;
 }
 
+function resolveAgentControlMcpUrl(config: CodexAgentConfig): string | null {
+  const extras = config.extra?.codex as CodexExtraConfig | undefined;
+  if (!extras || typeof extras !== "object") {
+    return null;
+  }
+  const baseUrl = extras.agentControlMcpUrl;
+  if (!baseUrl || typeof baseUrl !== "string") {
+    return null;
+  }
+  return config.parentAgentId
+    ? appendCallerAgentId(baseUrl, config.parentAgentId)
+    : baseUrl;
+}
+
+async function createCodexWrapperScript(
+  realCodexPath: string,
+  mcpUrl: string
+): Promise<CodexWrapperInfo> {
+  const templatePath = await resolveCodexWrapperTemplatePath();
+  const template = await fs.readFile(templatePath, "utf8");
+  const dir = await fs.mkdtemp(path.join(os.tmpdir(), "codex-wrapper-"));
+  const wrapperPath = path.join(dir, "codex-wrapper");
+  const content = template
+    .replace("__REAL_CODEX__", escapeBashValue(realCodexPath))
+    .replace("__MCP_URL__", escapeBashValue(mcpUrl));
+  await fs.writeFile(wrapperPath, content, { mode: 0o700 });
+  await fs.chmod(wrapperPath, 0o700);
+  return { path: wrapperPath, dir };
+}
+
+async function resolveCodexWrapperTemplatePath(): Promise<string> {
+  let dir = process.cwd();
+  for (let i = 0; i < 8; i += 1) {
+    const candidate = path.join(dir, "scripts", "codex-mcp-wrapper-template.sh");
+    if (await fileExists(candidate)) {
+      return candidate;
+    }
+    const parent = path.dirname(dir);
+    if (parent === dir) {
+      break;
+    }
+    dir = parent;
+  }
+  throw new Error(
+    "Unable to locate scripts/codex-mcp-wrapper-template.sh for Codex MCP wrapper"
+  );
+}
+
+function appendCallerAgentId(url: string, agentId: string): string {
+  try {
+    const parsed = new URL(url);
+    parsed.searchParams.set("callerAgentId", agentId);
+    return parsed.toString();
+  } catch {
+    const separator = url.includes("?") ? "&" : "?";
+    return `${url}${separator}callerAgentId=${encodeURIComponent(agentId)}`;
+  }
+}
+
+function escapeBashValue(value: string): string {
+  return value
+    .replace(/\\/g, "\\\\")
+    .replace(/"/g, '\\"')
+    .replace(/\$/g, "\\$");
+}
+
 export class CodexAgentClient implements AgentClient {
   readonly provider = "codex" as const;
   readonly capabilities = CODEX_CAPABILITIES;
-  private readonly codex: Codex;
+  private readonly baseOptions: CodexOptions;
+  private readonly baseCodexPath: string | null;
 
   constructor(options?: CodexOptions) {
     const codexOptions = { ...options };
@@ -273,12 +354,19 @@ export class CodexAgentClient implements AgentClient {
       }
     }
 
-    this.codex = new Codex(codexOptions);
+    this.baseOptions = codexOptions;
+    this.baseCodexPath = codexOptions.codexPathOverride ?? null;
   }
 
   async createSession(config: AgentSessionConfig): Promise<AgentSession> {
     const codexConfig = this.assertConfig(config);
-    return CodexAgentSession.create(this.codex, codexConfig);
+    const bootstrap = await this.createCodexForSession(codexConfig);
+    return CodexAgentSession.create(
+      bootstrap.codex,
+      codexConfig,
+      undefined,
+      bootstrap.wrapper
+    );
   }
 
   async resumeSession(
@@ -294,7 +382,13 @@ export class CodexAgentClient implements AgentClient {
     }
     const mergedConfig = { ...merged, provider: "codex" } as AgentSessionConfig;
     const codexConfig = this.assertConfig(mergedConfig);
-    return CodexAgentSession.create(this.codex, codexConfig, handle);
+    const bootstrap = await this.createCodexForSession(codexConfig);
+    return CodexAgentSession.create(
+      bootstrap.codex,
+      codexConfig,
+      handle,
+      bootstrap.wrapper
+    );
   }
 
   async listPersistedAgents(
@@ -334,15 +428,43 @@ export class CodexAgentClient implements AgentClient {
     }
     return config as CodexAgentConfig;
   }
+
+  private async createCodexForSession(
+    config: CodexAgentConfig
+  ): Promise<CodexSessionBootstrap> {
+    const mcpUrl = resolveAgentControlMcpUrl(config);
+    if (!mcpUrl || !this.baseCodexPath) {
+      if (mcpUrl && !this.baseCodexPath) {
+        console.warn("[Codex] MCP wrapper skipped (no codex binary path available)");
+      }
+      return { codex: new Codex(this.baseOptions) };
+    }
+
+    const wrapper = await createCodexWrapperScript(this.baseCodexPath, mcpUrl);
+    const env = {
+      ...process.env,
+      ...(this.baseOptions.env ?? {}),
+    };
+
+    return {
+      codex: new Codex({
+        ...this.baseOptions,
+        codexPathOverride: wrapper.path,
+        env,
+      }),
+      wrapper,
+    };
+  }
 }
 
 class CodexAgentSession implements AgentSession {
   static async create(
     codex: Codex,
     config: CodexAgentConfig,
-    handle?: AgentPersistenceHandle
+    handle?: AgentPersistenceHandle,
+    wrapper?: CodexWrapperInfo
   ): Promise<CodexAgentSession> {
-    const session = new CodexAgentSession(codex, config, handle);
+    const session = new CodexAgentSession(codex, config, handle, wrapper);
     if (handle) {
       await session.loadReplayHistory(handle);
     }
@@ -363,6 +485,7 @@ class CodexAgentSession implements AgentSession {
   private availableModes: AgentMode[] = CODEX_MODES;
   private readonly codexSessionDir: string | null;
   private rolloutPath: string | null;
+  private readonly wrapper: CodexWrapperInfo | null;
   private historyEvents: AgentStreamEvent[] = [];
   private pendingPermissions = new Map<string, AgentPermissionRequest>();
   private cancelCurrentTurn: (() => void) | null = null;
@@ -371,10 +494,12 @@ class CodexAgentSession implements AgentSession {
   constructor(
     codex: Codex,
     config: CodexAgentConfig,
-    handle?: AgentPersistenceHandle
+    handle?: AgentPersistenceHandle,
+    wrapper?: CodexWrapperInfo
   ) {
     this.codex = codex;
     this.config = { ...config };
+    this.wrapper = wrapper ?? null;
 
     // Validate mode if provided
     if (config.modeId && !VALID_CODEX_MODES.has(config.modeId)) {
@@ -634,6 +759,14 @@ class CodexAgentSession implements AgentSession {
   async close(): Promise<void> {
     this.thread = null;
     this.pendingPermissions.clear();
+    if (this.wrapper) {
+      try {
+        await fs.rm(this.wrapper.path, { force: true });
+        await fs.rm(this.wrapper.dir, { recursive: true, force: true });
+      } catch {
+        // ignore cleanup failures
+      }
+    }
   }
 
   // eslint-disable-next-line @typescript-eslint/no-unused-vars
@@ -697,6 +830,14 @@ class CodexAgentSession implements AgentSession {
       Object.assign(options, extra);
     }
 
+    // Codex CLI currently disables MCP tool availability when approval_policy is set.
+    // When injecting MCP servers, omit approvalPolicy/webSearchEnabled/networkAccessEnabled to ensure MCP tools are discoverable.
+    if (resolveAgentControlMcpUrl(this.config)) {
+      delete (options as Partial<ThreadOptions>).approvalPolicy;
+      delete (options as Partial<ThreadOptions>).webSearchEnabled;
+      delete (options as Partial<ThreadOptions>).networkAccessEnabled;
+    }
+
     return options;
   }