diff --git a/src/application/ClientConfigurationDocument.test.ts b/src/application/ClientConfigurationDocument.test.ts index f706cea6f4..83906e7ee4 100644 --- a/src/application/ClientConfigurationDocument.test.ts +++ b/src/application/ClientConfigurationDocument.test.ts @@ -1,8 +1,11 @@ import { describe, expect, it } from "vitest"; +import { parse as parseToml } from "smol-toml"; + import { clientConfigurationServersKey, parseClientConfiguration, + serializeClientConfiguration, } from "./ClientConfigurationDocument.js"; describe("JSON client configuration BOM handling", () => { @@ -40,3 +43,55 @@ describe("TOML client configuration BOM handling", () => { }); }); }); + +describe("Codex TOML serialization", () => { + const unrelated = [ + "# Keep this explanation.", + 'notify = ["a", "b"]', + "literal = 'C:\\demo\\path'", + "", + ].join("\n"); + const rea = { command: "npx", args: ["-y", "rea-agents@6.1.0", "mcp"] }; + const registered = `${unrelated}\n[mcp_servers.rea]\ncommand = "npx"\nargs = ["-y", "rea-agents@6.1.0", "mcp"]\n`; + const documentOf = (text: string): Record => + parseClientConfiguration(text, "toml").document; + const serialize = (document: Record, original: string) => + serializeClientConfiguration(document, "toml", original, [ + ["mcp_servers", "rea"], + ]); + + it("adds REA's table without rewriting comments or unrelated values", () => { + const document = { ...documentOf(unrelated), mcp_servers: { rea } }; + const written = serialize(document, unrelated); + + expect(written.startsWith(unrelated)).toBe(true); + expect(parseToml(written)).toEqual(document); + }); + + it("removes REA's table and leaves the rest of the file byte for byte", () => { + const document = { ...documentOf(registered), mcp_servers: {} }; + + expect(serialize(document, registered)).toBe(unrelated); + }); + + it("leaves a root disabled_mcp_servers list, which Codex does not read", () => { + const original = `# Codex ignores this list.\ndisabled_mcp_servers = ["rea"]\n`; + const document = { ...documentOf(original), mcp_servers: { rea } }; + const written = serialize(document, original); + + expect(written.startsWith(original)).toBe(true); + expect(parseToml(written)).toEqual(document); + }); + + it("rewrites the document when an edit would not hold the requested values", () => { + const document = { + ...documentOf(unrelated), + notify: ["changed"], + mcp_servers: { rea }, + }; + const written = serialize(document, unrelated); + + expect(written).not.toContain("# Keep this explanation."); + expect(parseToml(written)).toEqual(document); + }); +}); diff --git a/src/application/ClientConfigurationDocument.ts b/src/application/ClientConfigurationDocument.ts index 2b8a83a225..991dcd0b4e 100644 --- a/src/application/ClientConfigurationDocument.ts +++ b/src/application/ClientConfigurationDocument.ts @@ -399,7 +399,7 @@ const looksLikeTomlHeader = (line: string): boolean => line.trimStart().startsWith("["); /** `mcp_servers.` and tables nested under it, compared by key segment. */ -const isGrokServerTable = ( +const isTomlServerTable = ( path: readonly string[], serverKey: string, ): boolean => path[0] === "mcp_servers" && path[1] === serverKey; @@ -612,7 +612,7 @@ const hasSeparatorComma = (trivia: string): boolean => * Replace or remove `serverKey` inside one `{ ... }` table. A closed inline * table cannot later grow a `[mcp_servers.rea]` header. */ -const editGrokInlineTable = ( +const editInlineServerTable = ( table: string, serverKey: string, entry: unknown, @@ -831,13 +831,16 @@ const disabledServerListLines = ( }; /** - * Replace REA's Grok server tables and assignments. Comments that introduce a - * following table stay. Text inside strings is not treated as a header. + * Replace REA's TOML server tables and assignments. Comments that introduce a + * following table stay. Text inside strings is not treated as a header. Only + * Grok Build reads the root `disabled_mcp_servers` list that registration + * also edits. */ -const upsertGrokServerSection = ( +const upsertTomlServerSection = ( originalText: string, serverKey: string, entry: unknown, + editDisabledList: boolean, ): string => { // smol-toml accepts a leading BOM. Leave it in place or the first header is missed. const bom = originalText.startsWith("\uFEFF"); @@ -887,7 +890,7 @@ const upsertGrokServerSection = ( const header = parseTomlTableHeader(line); if (header !== undefined) { tablePath = header.path; - if (isGrokServerTable(header.path, serverKey)) { + if (isTomlServerTable(header.path, serverKey)) { discardPending(); skipping = true; lineIndex += 1; @@ -914,22 +917,24 @@ const upsertGrokServerSection = ( mode = "none"; continue; } - const disabledLines = disabledServerListLines( - disabledServerListEdit({ - text, - lineStart, - valueOffset, - valueEnd, - statementEnd, - tablePath, - segments: assignment.segments, - serverKey, - entry, - }), - lines, - lineIndex, - endLine, - ); + const disabledLines = editDisabledList + ? disabledServerListLines( + disabledServerListEdit({ + text, + lineStart, + valueOffset, + valueEnd, + statementEnd, + tablePath, + segments: assignment.segments, + serverKey, + entry, + }), + lines, + lineIndex, + endLine, + ) + : undefined; if (disabledLines !== undefined) { flushPending(); kept.push(...disabledLines); @@ -943,10 +948,10 @@ const upsertGrokServerSection = ( absolute.length === 1 && absolute[0] === "mcp_servers" && text[open] === "{"; - if (closesServerTable || isGrokServerTable(absolute, serverKey)) { + if (closesServerTable || isTomlServerTable(absolute, serverKey)) { if (closesServerTable) { flushPending(); - const edited = editGrokInlineTable( + const edited = editInlineServerTable( text.slice(open, valueEnd), serverKey, entry, @@ -1010,7 +1015,7 @@ const upsertGrokServerSection = ( return bom ? `\uFEFF${body}` : body; }; -const grokServerEntry = ( +const tomlServerEntry = ( document: Record, serverKey: string, ): unknown => { @@ -1033,15 +1038,71 @@ const serializeGrokConfiguration = ( document: Record, originalText: string | undefined, ): string => { - const entry = grokServerEntry(document, PRODUCT_IDENTITY.mcpServerKey); + const entry = tomlServerEntry(document, PRODUCT_IDENTITY.mcpServerKey); if (originalText === undefined) return stringifyToml(document); - return upsertGrokServerSection( + return upsertTomlServerSection( originalText, PRODUCT_IDENTITY.mcpServerKey, entry, + true, ); }; +/** Removing the last server leaves an empty table, which equals an absent one. */ +const withoutEmptyServerTable = ( + document: Record, +): Record => { + const servers = objectSchema.safeParse(document.mcp_servers); + if (!servers.success || Object.keys(servers.data).length > 0) return document; + const { mcp_servers: empty, ...rest } = document; + void empty; + return rest; +}; + +/** + * Keep an existing Codex document intact aside from the REA server tables. + * An edit that would change any other parsed value falls back to a full + * rewrite, which keeps every value but not comments or spelling. + */ +const serializeCodexConfiguration = ( + document: Record, + originalText: string | undefined, +): string => { + if ( + originalText === undefined || + isEmptyClientConfigurationText(originalText) + ) + return stringifyToml(document); + const upserted = upsertTomlServerSection( + originalText, + PRODUCT_IDENTITY.mcpServerKey, + tomlServerEntry(document, PRODUCT_IDENTITY.mcpServerKey), + false, + ); + // Removing a final table also drops the file's last line break. + const edited = + originalText.endsWith("\n") && !upserted.endsWith("\n") + ? `${upserted}${originalText.endsWith("\r\n") ? "\r\n" : "\n"}` + : upserted; + return tomlTextHolds(edited, document) ? edited : stringifyToml(document); +}; + +const tomlTextHolds = ( + text: string, + document: Record, +): boolean => { + try { + return clientConfigurationValuesEqual( + withoutEmptyServerTable(objectSchema.parse(parseToml(text))), + withoutEmptyServerTable(document), + ); + } catch (cause: unknown) { + // An unparsable edit is rejected in favor of the full rewrite. + void cause; + return false; + } +}; + /** Comments in the trivia between two properties. `undefined` when there are none. */ const commentsBetweenProperties = (gap: string): string | undefined => { const pieces: string[] = []; @@ -1146,7 +1207,7 @@ const deleteJsonPropertyPreservingComments = ( return `${text.slice(0, previous.offset + previous.length)}${suffix}${remainder}`; }; -/** Serialize a validated client document, retaining JSONC comments. */ +/** Serialize a validated client document, retaining JSONC and TOML comments. */ export const serializeClientConfiguration = ( document: Record, format: ClientConfigurationFormat | undefined, @@ -1155,7 +1216,8 @@ export const serializeClientConfiguration = ( ): string => { if (format === undefined || format === "unsupported") throw new TypeError("client does not have a supported MCP config format"); - if (format === "toml") return stringifyToml(document); + if (format === "toml") + return serializeCodexConfiguration(document, originalText); if (format === "grok") return serializeGrokConfiguration(document, originalText); // An empty original holds no comments or settings to preserve. Writing a diff --git a/tests/boundary/filesystem/clientLifecycle.test.ts b/tests/boundary/filesystem/clientLifecycle.test.ts index bf277a24ff..6c314eec45 100644 --- a/tests/boundary/filesystem/clientLifecycle.test.ts +++ b/tests/boundary/filesystem/clientLifecycle.test.ts @@ -1310,6 +1310,35 @@ describe("Grok Bot registration", () => { }); }); +describe("Codex configuration comments", () => { + it("keeps Codex comments and value spelling through setup and uninstall", async () => { + const home = await createTestTempDirectory("rea-toml-comments-"); + roots.push(home); + const configPath = join(home, ".codex/config.toml"); + const original = + '# Keep this explanation.\nnotify = ["a", "b"]\nliteral = \'C:\\demo\\path\'\n'; + await mkdir(dirname(configPath), { recursive: true }); + await writeFile(configPath, original); + + expect( + await configureClientConfiguration( + { name: "codex", configPath, format: "toml" }, + {}, + ["/isolated prefix/bin/rea", "mcp"], + ), + ).toMatchObject({ status: "configured" }); + const configured = await readFile(configPath, "utf8"); + expect(configured.startsWith(original)).toBe(true); + expect(configured).toContain("[mcp_servers.rea]"); + + const removed = await runUninstall(false, systemUninstallHost(home)); + expect(removed.items).toContainEqual( + expect.objectContaining({ name: "codex", status: "removed" }), + ); + expect(await readFile(configPath, "utf8")).toBe(original); + }); +}); + describe("client configuration filesystem removal", () => { it("uninstalls only owned entries and refuses purge symlinks", async () => { const home = await createTestTempDirectory("rea-uninstall-");