fix(codex): write mcp config to worktree

Author: mistercrunch

packages/core/package.json

@@ -161,6 +161,7 @@
     "@feathersjs/typebox": "^5.0.37",
     "@google/gemini-cli-core": "^0.9.0",
     "@google/genai": "^1.26.0",
+    "@iarna/toml": "^3.0.0",
     "@libsql/client": "^0.15.15",
     "@openai/codex-sdk": "^0.47.0",
     "@opencode-ai/sdk": "^1.0.30",
@@ -176,6 +177,7 @@
   "devDependencies": {
     "@types/bcryptjs": "^2.4.6",
     "@types/js-yaml": "^4.0.9",
+    "@types/iarna__toml": "^2.0.2",
     "@types/node": "^22.10.2",
     "@vitest/coverage-v8": "^4.0.3",
     "drizzle-kit": "^0.31.5",

packages/core/src/tools/codex/prompt-service.ts

@@ -12,6 +12,7 @@
 
 import * as fs from 'node:fs/promises';
 import * as path from 'node:path';
+import { type JsonMap, parse as parseToml, stringify as stringifyToml } from '@iarna/toml';
 import { Codex, type Thread, type ThreadItem } from '@openai/codex-sdk';
 import { getCredential, resolveApiKey, resolveUserEnvironment } from '../../config';
 import type { Database } from '../../db/client';
@@ -24,6 +25,16 @@ import type { TokenUsage } from '../../utils/pricing';
 import { DEFAULT_CODEX_MODEL } from './models';
 import { extractCodexTokenUsage } from './usage';
 
+type CodexTomlServerConfig = Record<string, unknown>;
+
+interface CodexTomlConfig {
+  approval_policy?: string;
+  sandbox_workspace_write?: Record<string, unknown>;
+  mcp_servers?: Record<string, CodexTomlServerConfig>;
+  agor_managed_servers?: string[];
+  [key: string]: unknown;
+}
+
 export interface CodexPromptResult {
   /** Complete assistant response from Codex */
   messages: Array<{
@@ -165,7 +176,7 @@ export class CodexPromptService {
   }
 
   /**
-   * Generate ~/.codex/config.toml with approval_policy, network_access, and MCP servers
+   * Generate Codex configuration files with approval_policy, network_access, and MCP servers
    *
    * NOTE: approval_policy, network_access, and MCP servers must be configured via config.toml
    * (not available in ThreadOptions). We minimize file writes by tracking a hash
@@ -174,12 +185,14 @@ export class CodexPromptService {
    * @param approvalPolicy - Codex approval policy (untrusted, on-request, on-failure, never)
    * @param networkAccess - Whether to allow outbound network access in workspace-write mode
    * @param sessionId - Session ID for fetching MCP servers
+   * @param worktreePath - Optional worktree path for repo-local .codex/config.toml
    * @returns Number of MCP servers configured
    */
   private async ensureCodexConfig(
     approvalPolicy: 'untrusted' | 'on-request' | 'on-failure' | 'never',
     networkAccess: boolean,
-    sessionId: SessionID
+    sessionId: SessionID,
+    worktreePath?: string
   ): Promise<number> {
     // Fetch MCP servers for this session (if repository is available)
     console.log(`🔍 [Codex MCP] Fetching MCP servers for session ${sessionId.substring(0, 8)}...`);
@@ -208,83 +221,170 @@ export class CodexPromptService {
       }
     }
 
-    // Generate MCP servers TOML blocks (stdio only)
-    let mcpServersToml = '';
+    // Build normalized server configs for cautious upsert
+    const managedServerConfigs: Record<string, Record<string, unknown>> = {};
+    const managedServerNames = new Set<string>();
     for (const server of stdioServers) {
-      // Normalize server name to lowercase for TOML convention
       const serverName = server.name.toLowerCase().replace(/[^a-z0-9_-]/g, '_');
+      managedServerNames.add(serverName);
       console.log(`   📝 [Codex MCP] Configuring server: ${server.name} -> ${serverName}`);
 
-      mcpServersToml += `\n[mcp_servers.${serverName}]\n`;
+      const serverConfig: Record<string, unknown> = {};
       if (server.command) {
-        mcpServersToml += `command = "${server.command}"\n`;
+        serverConfig.command = server.command;
         console.log(`      command: ${server.command}`);
       }
       if (server.args && server.args.length > 0) {
-        const argsJson = JSON.stringify(server.args);
-        mcpServersToml += `args = ${argsJson}\n`;
-        console.log(`      args: ${argsJson}`);
+        serverConfig.args = server.args;
+        console.log(`      args: ${JSON.stringify(server.args)}`);
       }
-
-      // Add environment variables if present
       if (server.env && Object.keys(server.env).length > 0) {
-        mcpServersToml += `\n[mcp_servers.${serverName}.env]\n`;
-        const envCount = Object.keys(server.env).length;
-        console.log(`      env vars: ${envCount} variable(s)`);
-        for (const [key, value] of Object.entries(server.env)) {
-          mcpServersToml += `${key} = "${value}"\n`;
-        }
+        serverConfig.env = server.env;
+        console.log(`      env vars: ${Object.keys(server.env).length} variable(s)`);
       }
-    }
-
-    // Generate network access TOML section (only for workspace-write sandbox)
-    const networkAccessToml = networkAccess
-      ? `\n[sandbox_workspace_write]\nnetwork_access = true\n`
-      : '';
 
-    // Generate complete config content
-    const configContent = `# Codex configuration
-# Generated by Agor - ${new Date().toISOString()}
-
-# Approval policy controls when Codex asks before running commands
-# Options: "untrusted", "on-request", "on-failure", "never"
-approval_policy = "${approvalPolicy}"
-${networkAccessToml}${mcpServersToml}`;
+      managedServerConfigs[serverName] = serverConfig;
+    }
 
     // Create hash to detect changes (include network access in hash)
-    const configHash = `${approvalPolicy}:${networkAccess}:${JSON.stringify(stdioServers.map(s => s.mcp_server_id))}`;
+    const configHash = `${approvalPolicy}:${networkAccess}:${worktreePath ?? ''}:${JSON.stringify(
+      stdioServers.map(s => s.mcp_server_id)
+    )}`;
 
     // Skip if config hasn't changed (avoid unnecessary file I/O)
     if (this.lastMCPServersHash === configHash) {
       console.log(`✅ [Codex MCP] Config unchanged, skipping write`);
       return stdioServers.length;
     }
 
+    const configTargets: Array<{ label: 'global' | 'worktree'; path: string }> = [];
+
     const homeDir = process.env.HOME || process.env.USERPROFILE;
-    if (!homeDir) {
-      console.warn('⚠️  [Codex MCP] Could not determine home directory, skipping Codex config');
-      return 0;
+    if (homeDir) {
+      configTargets.push({
+        label: 'global',
+        path: path.join(homeDir, '.codex', 'config.toml'),
+      });
+    } else {
+      console.warn(
+        '⚠️  [Codex MCP] Could not determine home directory, skipping global Codex config'
+      );
     }
 
-    const codexConfigDir = path.join(homeDir, '.codex');
-    const configPath = path.join(codexConfigDir, 'config.toml');
+    if (worktreePath) {
+      const repoConfigPath = path.join(worktreePath, '.codex', 'config.toml');
+      configTargets.push({
+        label: 'worktree',
+        path: repoConfigPath,
+      });
+    }
 
-    console.log(`📁 [Codex MCP] Writing config to: ${configPath}`);
-    console.log(`📄 [Codex MCP] Config content:\n${configContent}`);
+    if (configTargets.length === 0) {
+      console.warn('⚠️  [Codex MCP] No writable locations found for Codex config');
+      return stdioServers.length;
+    }
 
-    await fs.mkdir(codexConfigDir, { recursive: true });
-    await fs.writeFile(configPath, configContent, 'utf-8');
+    for (const target of configTargets) {
+      let config: CodexTomlConfig = {};
+      let previouslyManaged = new Set<string>();
+
+      try {
+        const existing = await fs.readFile(target.path, 'utf-8');
+        config = parseToml(existing) as CodexTomlConfig;
+        const managedList = Array.isArray(config.agor_managed_servers)
+          ? config.agor_managed_servers.map((value: unknown) => String(value)).filter(Boolean)
+          : [];
+        previouslyManaged = new Set<string>(managedList);
+      } catch (error) {
+        if ((error as NodeJS.ErrnoException)?.code !== 'ENOENT') {
+          console.warn(
+            `⚠️  [Codex MCP] Failed to read existing ${target.label} config (${target.path}):`,
+            error
+          );
+        }
+        config = {};
+        previouslyManaged = new Set<string>();
+      }
 
-    this.lastMCPServersHash = configHash;
-    console.log(
-      `✅ [Codex] Updated config.toml with approval_policy = "${approvalPolicy}", network_access = ${networkAccess}`
-    );
+      // Upsert approval policy
+      config.approval_policy = approvalPolicy;
+
+      // Upsert sandbox workspace settings
+      if (networkAccess) {
+        const sandbox =
+          typeof config.sandbox_workspace_write === 'object' &&
+          config.sandbox_workspace_write !== null
+            ? { ...config.sandbox_workspace_write }
+            : {};
+        sandbox.network_access = true;
+        config.sandbox_workspace_write = sandbox;
+      } else if (
+        typeof config.sandbox_workspace_write === 'object' &&
+        config.sandbox_workspace_write !== null
+      ) {
+        const sandbox = { ...config.sandbox_workspace_write };
+        delete sandbox.network_access;
+        if (Object.keys(sandbox).length === 0) {
+          delete config.sandbox_workspace_write;
+        } else {
+          config.sandbox_workspace_write = sandbox;
+        }
+      }
 
-    // Reinitialize Codex SDK to pick up the new config
+      // Prepare MCP servers table
+      const mcpServersTable: Record<string, CodexTomlServerConfig> =
+        typeof config.mcp_servers === 'object' && config.mcp_servers !== null
+          ? { ...config.mcp_servers }
+          : {};
+
+      // Remove previously managed entries no longer present
+      for (const name of previouslyManaged) {
+        if (!managedServerNames.has(name)) {
+          delete mcpServersTable[name];
+        }
+      }
+
+      // Upsert current managed servers
+      for (const [name, serverConfig] of Object.entries(managedServerConfigs)) {
+        mcpServersTable[name] = serverConfig;
+      }
+
+      if (Object.keys(mcpServersTable).length > 0) {
+        config.mcp_servers = mcpServersTable;
+      } else {
+        delete config.mcp_servers;
+      }
+
+      if (managedServerNames.size > 0) {
+        config.agor_managed_servers = Array.from(managedServerNames).sort();
+      } else {
+        delete config.agor_managed_servers;
+      }
+
+      const header = `# Codex configuration\n# Generated by Agor - ${new Date().toISOString()}\n\n`;
+      const body = stringifyToml(config as JsonMap);
+
+      try {
+        await fs.mkdir(path.dirname(target.path), { recursive: true });
+        await fs.writeFile(target.path, header + body, 'utf-8');
+        console.log(
+          `✅ [Codex MCP] Updated ${target.label} config with ${managedServerNames.size} managed MCP server(s)`
+        );
+      } catch (error) {
+        console.error(
+          `❌ [Codex MCP] Failed to write ${target.label} config (${target.path}):`,
+          error
+        );
+      }
+    }
+
+    this.lastMCPServersHash = configHash;
     this.reinitializeCodex();
     if (stdioServers.length > 0) {
       console.log(
-        `✅ [Codex MCP] Configured ${stdioServers.length} STDIO MCP server(s): ${stdioServers.map(s => s.name).join(', ')}`
+        `✅ [Codex MCP] Configured ${stdioServers.length} STDIO MCP server(s): ${stdioServers
+          .map(s => s.name)
+          .join(', ')}`
       );
     }
 
@@ -423,8 +523,23 @@ ${networkAccessToml}${mcpServersToml}`;
       `   Using Codex permissions: sandboxMode=${sandboxMode}, approvalPolicy=${approvalPolicy}, networkAccess=${networkAccess}`
     );
 
+    // Fetch worktree to determine working directory (needed before writing repo-level config)
+    const worktree = this.worktreesRepo
+      ? await this.worktreesRepo.findById(session.worktree_id)
+      : null;
+    if (!worktree) {
+      throw new Error(`Worktree ${session.worktree_id} not found for session ${sessionId}`);
+    }
+
+    const worktreePath = worktree.path;
+
     // Set approval_policy, network_access, and MCP servers in config.toml (required because they're not available in ThreadOptions)
-    const mcpServerCount = await this.ensureCodexConfig(approvalPolicy, networkAccess, sessionId);
+    const mcpServerCount = await this.ensureCodexConfig(
+      approvalPolicy,
+      networkAccess,
+      sessionId,
+      worktreePath
+    );
 
     const totalMcpServers = this.sessionMCPServerRepo
       ? (await this.sessionMCPServerRepo.listServers(sessionId, true)).length
@@ -439,19 +554,11 @@ ${networkAccessToml}${mcpServersToml}`;
       );
     }
 
-    // Fetch worktree to get working directory
-    const worktree = this.worktreesRepo
-      ? await this.worktreesRepo.findById(session.worktree_id)
-      : null;
-    if (!worktree) {
-      throw new Error(`Worktree ${session.worktree_id} not found for session ${sessionId}`);
-    }
-
-    console.log(`   Working directory: ${worktree.path}`);
+    console.log(`   Working directory: ${worktreePath}`);
 
     // Build thread options with sandbox mode and worktree working directory
     const threadOptions = {
-      workingDirectory: worktree.path,
+      workingDirectory: worktreePath,
       skipGitRepoCheck: false,
       sandboxMode,
     };