diff --git a/src/utils/code-tools/codex-toml-updater.ts b/src/utils/code-tools/codex-toml-updater.ts index a1923446..61f1573c 100644 --- a/src/utils/code-tools/codex-toml-updater.ts +++ b/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 @@ -203,25 +204,10 @@ export function upsertCodexMcpService(serviceId: string, service: CodexMcpServic 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 @@ export function deleteCodexMcpService(serviceId: string): void { 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 @@ export function batchUpdateCodexMcpServices( 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)}` + } + + return '' +} + +function formatInlineTable(obj: Record): 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) + } + + 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) + } + + 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 +} diff --git a/tests/unit/utils/code-tools/codex-platform.test.ts b/tests/unit/utils/code-tools/codex-platform.test.ts index e5ea43e4..28843efa 100644 --- a/tests/unit/utils/code-tools/codex-platform.test.ts +++ b/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,6 +75,19 @@ 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 { + 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']) @@ -81,11 +95,7 @@ describe('applyCodexPlatformCommand integration', () => { { 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"') + }) })