fix(codex): preserve nested MCP env tables when updating config

src/utils/code-tools/codex-toml-updater.ts

@@ -9,11 +9,12 @@
  * - MCP modifications should NOT affect API configurations
  */
 
-import type { CodexMcpService, CodexProvider } from './codex'
+import type { CodexConfigData, CodexMcpService, CodexProvider } from './codex'
 import { CODEX_CONFIG_FILE, CODEX_DIR } from '../../constants'
 import { ensureDir, exists, readFile, writeFile } from '../fs-operations'
 import { normalizeTomlPath } from '../platform'
 import { editToml, parseToml } from '../toml-edit'
+import { parseCodexConfig } from './codex'
 
 /**
  * Update top-level TOML fields using regex-based replacement
@@ -194,7 +195,7 @@
   }
 
   let content = readFile(CODEX_CONFIG_FILE) || ''
   const basePath = `mcp_servers.${serviceId}`
 
   // Check if this is an existing service with 'url' field (SSE protocol)
   // If so, we should NOT add command/args fields
@@ -203,25 +204,10 @@
 
   if (existingService?.url && !existingService?.command) {
     // This is an SSE-type service, only update non-conflicting fields
-    if (service.env && Object.keys(service.env).length > 0) {
-      content = editToml(content, `${basePath}.env`, service.env)
-    }
-    if (service.startup_timeout_sec) {
-      content = editToml(content, `${basePath}.startup_timeout_sec`, service.startup_timeout_sec)
-    }
+    content = upsertMcpSection(content, serviceId, service, { preserveCommandAndArgs: true })
   }
   else {
-    // This is a stdio-type service or new service, update all fields
-    const normalizedCommand = normalizeTomlPath(service.command)
-    content = editToml(content, `${basePath}.command`, normalizedCommand)
-    content = editToml(content, `${basePath}.args`, service.args || [])
-
-    if (service.env && Object.keys(service.env).length > 0) {
-      content = editToml(content, `${basePath}.env`, service.env)
-    }
-    if (service.startup_timeout_sec) {
-      content = editToml(content, `${basePath}.startup_timeout_sec`, service.startup_timeout_sec)
-    }
+    content = upsertMcpSection(content, serviceId, service)
   }
 
   writeFile(CODEX_CONFIG_FILE, content)
@@ -242,10 +228,7 @@
   const content = readFile(CODEX_CONFIG_FILE) || ''
 
   // Use regex to remove the entire section
-  const sectionRegex = new RegExp(
-    `\\n?\\[mcp_servers\\.${escapeRegex(serviceId)}\\][\\s\\S]*?(?=\\n\\[|$)`,
-    'g',
-  )
+  const sectionRegex = createMcpSectionRegex(serviceId)
 
   const updatedContent = content.replace(sectionRegex, '')
   writeFile(CODEX_CONFIG_FILE, updatedContent)
@@ -290,3 +273,178 @@
 function escapeRegex(str: string): string {
   return str.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')
 }
+
+function createMcpSectionRegex(serviceId: string): RegExp {
+  const escapedServiceId = escapeRegex(serviceId)
+  return new RegExp(
+    `(?:^|\\n)[ \\t]*\\[mcp_servers\\.${escapedServiceId}(?:\\.[^\\]]+)?\\][\\s\\S]*?(?=(?:\\n[ \\t]*\\[(?!mcp_servers\\.${escapedServiceId}(?:\\.|\\]))|$))`,
+    'g',
+  )
+}
+
+function upsertMcpSection(
+  content: string,
+  serviceId: string,
+  service: CodexMcpService,
+  options: { preserveCommandAndArgs?: boolean } = {},
+): string {
+  const existingConfig = content ? parseCodexConfig(content) : null
+  const existingService = existingConfig?.mcpServices.find(candidate => candidate.id === serviceId)
+  const mergedService = mergeMcpService(serviceId, existingService, service, options)
+  const section = renderMcpSection(serviceId, mergedService, options)
+  const sectionRegex = createMcpSectionRegex(serviceId)
+
+  if (sectionRegex.test(content)) {
+    const normalized = content.replace(sectionRegex, `\n${section.trimEnd()}\n`)
+    return normalized.replace(/^\n/, '')
+  }
+
+  const separator = content.trimEnd().length > 0 ? '\n\n' : ''
+  return `${content.trimEnd()}${separator}${section}\n`
+}
+
+function mergeMcpService(
+  serviceId: string,
+  existingService: CodexMcpService | undefined,
+  nextService: CodexMcpService,
+  options: { preserveCommandAndArgs?: boolean } = {},
+): CodexMcpService {
+  const mergedExtraFields = {
+    ...(existingService?.extraFields || {}),
+    ...(nextService.extraFields || {}),
+  }
+
+  if (options.preserveCommandAndArgs && existingService?.command) {
+    return {
+      ...nextService,
+      id: serviceId,
+      command: existingService.command,
+      args: existingService.args || [],
+      env: nextService.env ?? existingService.env,
+      startup_timeout_sec: nextService.startup_timeout_sec ?? existingService.startup_timeout_sec,
+      extraFields: Object.keys(mergedExtraFields).length > 0 ? mergedExtraFields : undefined,
+    }
+  }
+
+  return {
+    ...existingService,
+    ...nextService,
+    id: serviceId,
+    env: nextService.env ?? existingService?.env,
+    startup_timeout_sec: nextService.startup_timeout_sec ?? existingService?.startup_timeout_sec,
+    extraFields: Object.keys(mergedExtraFields).length > 0 ? mergedExtraFields : undefined,
+  }
+}
+
+function renderMcpSection(
+  serviceId: string,
+  service: CodexMcpService,
+  options: { preserveCommandAndArgs?: boolean } = {},
+): string {
+  const lines = [`[mcp_servers.${serviceId}]`]
+
+  if (!options.preserveCommandAndArgs) {
+    const normalizedCommand = normalizeTomlStringValue(service.command)
+    lines.push(`command = "${escapeTomlString(normalizedCommand)}"`)
+    lines.push(`args = ${formatTomlArray(service.args || [])}`)
+  }
+
+  if (service.env && Object.keys(service.env).length > 0) {
+    lines.push(`env = ${formatInlineTable(service.env)}`)
+  }
+
+  if (service.startup_timeout_sec) {
+    lines.push(`startup_timeout_sec = ${service.startup_timeout_sec}`)
+  }
+
+  if (service.extraFields) {
+    for (const [key, value] of Object.entries(service.extraFields)) {
+      const formatted = formatTomlField(key, value)
+      if (formatted) {
+        lines.push(formatted)
+      }
+    }
+  }
+
+  return lines.join('\n')
+}
+
+function formatTomlField(key: string, value: unknown): string {
+  if (value === null || value === undefined) {
+    return ''
+  }
+
+  if (typeof value === 'string') {
+    return `${key} = "${escapeTomlString(normalizeTomlStringValue(value))}"`
+  }
+
+  if (typeof value === 'number' || typeof value === 'boolean') {
+    return `${key} = ${value}`
+  }
+
+  if (Array.isArray(value)) {
+    return `${key} = ${formatTomlArray(value)}`
+  }
+
+  if (typeof value === 'object') {
+    return `${key} = ${formatInlineTable(value as Record<string, unknown>)}`
+  }
+
+  return ''
+}
+
+function formatInlineTable(obj: Record<string, unknown>): string {
+  const entries = Object.entries(obj)
+    .filter(([, value]) => value !== null && value !== undefined)
+    .map(([key, value]) => `${key} = ${formatInlineTableValue(value)}`)
+    .join(', ')
+  return `{${entries}}`
+}
+
+function formatInlineTableValue(value: unknown): string {
+  if (typeof value === 'string') {
+    return `'${normalizeTomlStringValue(value).replace(/'/g, "''")}'`
+  }
+
+  if (typeof value === 'number' || typeof value === 'boolean') {
+    return String(value)
+  }
+
+  if (Array.isArray(value)) {
+    return formatTomlArray(value)
+  }
+
+  if (value && typeof value === 'object') {
+    return formatInlineTable(value as Record<string, unknown>)
+  }
+
+  return `''`
+}
+
+function formatTomlArray(values: unknown[]): string {
+  const items = values.map((value) => {
+    if (typeof value === 'string') {
+      return `"${escapeTomlString(normalizeTomlStringValue(value))}"`
+    }
+
+    if (typeof value === 'number' || typeof value === 'boolean') {
+      return String(value)
+    }
+
+    if (value && typeof value === 'object') {
+      return formatInlineTable(value as Record<string, unknown>)
+    }
+
+    return '""'
+  })
+
+  return `[${items.join(', ')}]`
+}
+
+function escapeTomlString(value: string): string {
+  return value.replace(/\\/g, '\\\\').replace(/"/g, '\\"')
+}
+
+function normalizeTomlStringValue(value: string): string {
+  return value.includes('\\') ? normalizeTomlPath(value) : value
+}

tests/unit/utils/code-tools/codex-platform.test.ts

@@ -1,3 +1,4 @@
+import type { CodexConfigData } from '../../../../src/utils/code-tools/codex'
 import { describe, expect, it, vi } from 'vitest'
 
 const mockSelectMcpServices = vi.fn()
@@ -74,18 +75,27 @@ const codexModule = await import('../../../../src/utils/code-tools/codex')
 const { configureCodexMcp } = codexModule
 const { writeFile } = await import('../../../../src/utils/fs-operations')
 
+function createMockCodexConfig(overrides: Partial<CodexConfigData> = {}): CodexConfigData {
+  return {
+    model: null,
+    modelProvider: null,
+    providers: [],
+    mcpServices: [],
+    managed: false,
+    otherConfig: [],
+    modelProviderCommented: undefined,
+    ...overrides,
+  }
+}
+
 describe('applyCodexPlatformCommand integration', () => {
   it('should rewrite npx commands using platform-specific MCP command', async () => {
     mockSelectMcpServices.mockResolvedValue(['SERVICE'])
     mockGetMcpServices.mockResolvedValue([
       { id: 'SERVICE', name: 'Service', description: 'desc' },
     ])
 
-    vi.spyOn(codexModule, 'readCodexConfig').mockReturnValue({
-      providers: [],
-      mcpServices: [],
-      managed: false,
-    } as any)
+    vi.spyOn(codexModule, 'readCodexConfig').mockReturnValue(createMockCodexConfig())
     vi.spyOn(codexModule, 'backupCodexComplete').mockReturnValue(null)
 
     await configureCodexMcp()
@@ -108,11 +118,7 @@ describe('applyCodexPlatformCommand integration', () => {
       { id: 'serena', name: 'Serena', description: 'Serena MCP service' },
     ])
 
-    vi.spyOn(codexModule, 'readCodexConfig').mockReturnValue({
-      providers: [],
-      mcpServices: [],
-      managed: false,
-    } as any)
+    vi.spyOn(codexModule, 'readCodexConfig').mockReturnValue(createMockCodexConfig())
     vi.spyOn(codexModule, 'backupCodexComplete').mockReturnValue(null)
 
     await configureCodexMcp()
@@ -128,4 +134,111 @@ describe('applyCodexPlatformCommand integration', () => {
     expect(lastConfigContent).toContain('[mcp_servers.serena]')
     expect(lastConfigContent).toContain('command = "cmd"')
   })
+
+  it('should preserve existing node_repl env tables when adding MCP services', async () => {
+    mockSelectMcpServices.mockResolvedValue(['SERVICE'])
+    mockGetMcpServices.mockResolvedValue([
+      { id: 'SERVICE', name: 'Service', description: 'desc' },
+    ])
+
+    const { readFile } = await import('../../../../src/utils/fs-operations')
+    vi.mocked(readFile).mockReturnValue(`model_provider = "jjj"
+model = "gpt-5.2"
+
+[mcp_servers.node_repl]
+args = []
+command = 'C:/Users/yukaidi/AppData/Local/OpenAI/Codex/bin/node_repl.exe'
+startup_timeout_sec = 120
+
+[mcp_servers.node_repl.env]
+NODE_REPL_NATIVE_PIPE_CONNECT_TIMEOUT_MS = "1000"
+NODE_REPL_NODE_MODULE_DIRS = ""
+NODE_REPL_NODE_PATH = 'C:/Users/yukaidi/AppData/Local/OpenAI/Codex/bin/node.exe'
+CODEX_HOME = 'C:/Users/yukaidi/.codex'
+`)
+
+    vi.spyOn(codexModule, 'readCodexConfig').mockReturnValue(createMockCodexConfig())
+    vi.spyOn(codexModule, 'backupCodexComplete').mockReturnValue(null)
+
+    await configureCodexMcp()
+
+    const writeFileMock = vi.mocked(writeFile)
+    const configCalls = writeFileMock.mock.calls.filter(call => call[0].includes('config.toml'))
+    expect(configCalls.length).toBeGreaterThan(0)
+    const lastConfigContent = configCalls[configCalls.length - 1][1] as string
+
+    expect(lastConfigContent).toContain('[mcp_servers.node_repl.env]')
+    expect(lastConfigContent).not.toContain('{ NODE_REPL_NATIVE_PIPE_CONNECT_TIMEOUT_MS')
+    expect(lastConfigContent).not.toContain('[mcp_servers.node_repl]\nenv = {')
+    expect(lastConfigContent).toContain('[mcp_servers.service]')
+  })
+
+  it('should preserve existing SSE MCP url and extra fields during updates', async () => {
+    mockSelectMcpServices.mockResolvedValue([])
+    mockGetMcpServices.mockResolvedValue([])
+
+    const { readFile } = await import('../../../../src/utils/fs-operations')
+    vi.mocked(readFile).mockReturnValue(`
+[mcp_servers.remote-docs]
+url = "https://example.com/sse"
+startup_timeout_sec = 15
+retries = 3
+`)
+
+    vi.spyOn(codexModule, 'readCodexConfig').mockReturnValue(createMockCodexConfig({
+      mcpServices: [{
+        id: 'remote-docs',
+        command: 'remote-docs',
+        args: [],
+        startup_timeout_sec: 15,
+        extraFields: {
+          url: 'https://example.com/sse',
+          retries: 3,
+        },
+      }],
+    }))
+    vi.spyOn(codexModule, 'backupCodexComplete').mockReturnValue(null)
+
+    await configureCodexMcp()
+
+    const writeFileMock = vi.mocked(writeFile)
+    const configCalls = writeFileMock.mock.calls.filter(call => call[0].includes('config.toml'))
+    expect(configCalls.length).toBeGreaterThan(0)
+    const lastConfigContent = configCalls[configCalls.length - 1][1] as string
+
+    expect(lastConfigContent).toContain('[mcp_servers.remote-docs]')
+    expect(lastConfigContent).toContain('url = "https://example.com/sse"')
+    expect(lastConfigContent).toContain('retries = 3')
+  })
+
+  it('should stop at indented non-MCP section headers when replacing MCP sections', async () => {
+    mockSelectMcpServices.mockResolvedValue(['SERVICE'])
+    mockGetMcpServices.mockResolvedValue([
+      { id: 'SERVICE', name: 'Service', description: 'desc' },
+    ])
+
+    const { readFile } = await import('../../../../src/utils/fs-operations')
+    vi.mocked(readFile).mockReturnValue(`
+  [mcp_servers.service]
+  command = "old-command"
+  args = []
+
+  [custom_section]
+  key = "must-stay"
+`)
+
+    vi.spyOn(codexModule, 'readCodexConfig').mockReturnValue(createMockCodexConfig())
+    vi.spyOn(codexModule, 'backupCodexComplete').mockReturnValue(null)
+
+    await configureCodexMcp()
+
+    const writeFileMock = vi.mocked(writeFile)
+    const configCalls = writeFileMock.mock.calls.filter(call => call[0].includes('config.toml'))
+    expect(configCalls.length).toBeGreaterThan(0)
+    const lastConfigContent = configCalls[configCalls.length - 1][1] as string
+
+    expect(lastConfigContent).toContain('[mcp_servers.service]')
+    expect(lastConfigContent).toContain('[custom_section]')
+    expect(lastConfigContent).toContain('key = "must-stay"')
+  })
 })