Merge pull request #1178 from morluto/fix/codex-toml-preserve-comments

fix(setup): keep Codex TOML comments when adding or removing REA

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<string, unknown> =>
+    parseClientConfiguration(text, "toml").document;
+  const serialize = (document: Record<string, unknown>, 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);
+  });
+});

src/application/ClientConfigurationDocument.ts

@@ -399,7 +399,7 @@ const looksLikeTomlHeader = (line: string): boolean =>
   line.trimStart().startsWith("[");
 
 /** `mcp_servers.<server>` 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<string, unknown>,
   serverKey: string,
 ): unknown => {
@@ -1033,15 +1038,71 @@ const serializeGrokConfiguration = (
   document: Record<string, unknown>,
   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<string, unknown>,
+): Record<string, unknown> => {
+  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<string, unknown>,
+  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<string, unknown>,
+): 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<string, unknown>,
   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

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-");