fix(setup): keep comments around REA's table in Grok Build TOML (#6261)

The shared TOML server-table editor kept comment lines immediately only
for the Codex dialect. For Grok Build it queued them with blank lines,
then discarded the queue when an REA table header followed or while it
skipped REA's table, so a comment introducing the stale REA table and a
trailing file note after it were lost on setup and uninstall.

Every comment line is now kept in place for both TOML dialects; only
blank lines that separated a removed REA table go with it. The Codex
wrapper's terminal line-break restoration moves into the shared editor,
so removing a final REA table keeps the file's last line break for
Grok Build too.

Fixes #6207

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

src/application/ClientConfigurationDocument.ts

@@ -978,10 +978,11 @@ const withoutTrailingLineFeeds = (text: string): string => {
 };
 
 /**
- * Replace REA's TOML server tables and assignments. Comments that introduce a
- * following table stay. Codex also retains comments around REA's own table.
- * Text inside strings is not treated as a header. Only Grok Build reads the
- * root `disabled_mcp_servers` list that registration also edits.
+ * Replace REA's TOML server tables and assignments. Every comment line stays
+ * where it was, including comments before, inside, and after REA's own table;
+ * only blank lines that separated a removed table go with it. Text inside
+ * strings is not treated as a header. Only Grok Build reads the root
+ * `disabled_mcp_servers` list that registration also edits.
  */
 const upsertTomlServerSection = (
   originalText: string,
@@ -1005,6 +1006,8 @@ const upsertTomlServerSection = (
     cursor += line.length + 1;
   }
   const kept: string[] = [];
+  // Blank lines held until the next statement shows whether they belong to a
+  // removed REA table. Comments are user text and are never held.
   const pending: string[] = [];
   let skipping = false;
   let inlineServerTable = false;
@@ -1128,7 +1131,7 @@ const upsertTomlServerSection = (
       continue;
     }
     if (isBlankOrComment(line)) {
-      if (format === "toml" && line.trimStart().startsWith("#")) {
+      if (line.trimStart().startsWith("#")) {
         flushPending();
         kept.push(physical);
       } else pending.push(physical);
@@ -1168,6 +1171,9 @@ const upsertTomlServerSection = (
         ? `${fragment}${nl}`
         : `${trimmed}${nl}${nl}${fragment}${nl}`;
   }
+  // Removing a final table also drops the file's last line break.
+  if (originalText.endsWith("\n") && !body.endsWith("\n"))
+    body += originalText.endsWith("\r\n") ? "\r\n" : "\n";
   return bom ? `\uFEFF${body}` : body;
 };
 
@@ -1270,17 +1276,12 @@ const serializeCodexConfiguration = (
     isEmptyClientConfigurationText(originalText)
   )
     return stringifyToml(document);
-  const upserted = upsertTomlServerSection(
+  const edited = upsertTomlServerSection(
     originalText,
     PRODUCT_IDENTITY.mcpServerKey,
     tomlServerEntry(document, PRODUCT_IDENTITY.mcpServerKey),
     "toml",
   );
-  // 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);
 };
 

tests/acceptance/setup/tomlCommentPreservationCli.test.ts

@@ -0,0 +1,162 @@
+import { mkdir, readFile, writeFile } from "node:fs/promises";
+import { dirname, join } from "node:path";
+import { parse as parseToml } from "smol-toml";
+import { expect } from "vitest";
+
+import { createTestTempDirectory } from "../../fixtures/temporaryDirectory.js";
+import { cliTest } from "../../support/cli/cliFixture.js";
+
+const clients = [
+  { name: "codex", path: ".codex/config.toml" },
+  { name: "grok_build", path: ".grok/config.toml" },
+] as const;
+
+/** Parsed values other than REA's registration. */
+const unrelatedValues = (text: string) => {
+  const document = parseToml(text);
+  const servers = document.mcp_servers;
+  if (typeof servers === "object" && servers !== null) {
+    Reflect.deleteProperty(servers, "rea");
+    if (Object.keys(servers).length === 0) delete document.mcp_servers;
+  }
+  return document;
+};
+
+const scenarios = [
+  {
+    name: "a stale REA table between other tables",
+    original: [
+      "# root explanation",
+      'model = "x" # inline model note',
+      "",
+      "# about rea",
+      "[mcp_servers.rea]",
+      "# inside the rea table",
+      'command = "old"',
+      "",
+      "# between rea and its environment",
+      "[mcp_servers.rea.env]",
+      'OLD = "value"',
+      "# after the rea environment",
+      "",
+      "# about the other server",
+      "[mcp_servers.other]",
+      'command = "other" # other inline note',
+      "",
+      "# trailing note about the file",
+      "",
+    ].join("\n"),
+    comments: [
+      "# root explanation",
+      'model = "x" # inline model note',
+      "# about rea",
+      "# inside the rea table",
+      "# between rea and its environment",
+      "# after the rea environment",
+      "# about the other server",
+      'command = "other" # other inline note',
+      "# trailing note about the file",
+    ],
+  },
+  {
+    name: "a leading REA table and a final note without a line break",
+    original: [
+      "# about rea",
+      "[mcp_servers.rea]",
+      'command = "old"',
+      "# trailing note without a line break",
+    ].join("\n"),
+    comments: ["# about rea", "# trailing note without a line break"],
+  },
+] as const;
+
+for (const scenario of scenarios)
+  cliTest(
+    `TOML client setup and uninstall keep comments around ${scenario.name}`,
+    async ({ cli }) => {
+      const home = await createTestTempDirectory("rea-toml-comments-");
+      const paths = clients.map(({ path }) => join(home, path));
+      for (const path of paths) {
+        await mkdir(dirname(path), { recursive: true });
+        await writeFile(path, scenario.original);
+      }
+      const run = (arguments_: readonly string[]) =>
+        cli.run({
+          arguments: arguments_,
+          cwd: home,
+          environment: {
+            HOME: home,
+            USERPROFILE: home,
+            CODEX_HOME: join(home, ".codex"),
+            GROK_HOME: join(home, ".grok"),
+            XDG_CONFIG_HOME: join(home, ".config"),
+            TMPDIR: home,
+          },
+        });
+      const setup = [
+        "setup",
+        ...clients.flatMap(({ name }) => ["--client", name]),
+        "--skill=false",
+        "--json",
+      ];
+
+      const preview = await run([...setup, "--dry-run"]);
+      expect(preview.exitCode, preview.stdout).toBe(0);
+      expect(preview.json).toMatchObject({ status: "planned" });
+      for (const path of paths)
+        expect(await readFile(path, "utf8")).toBe(scenario.original);
+
+      const applied = await run([...setup, "--yes"]);
+      expect(applied.exitCode, applied.stdout).toBe(0);
+      expect(applied.json).toMatchObject({
+        status: "ready",
+        clients: {
+          codex: { status: "configured" },
+          grok_build: { status: "configured" },
+        },
+      });
+      const configured: string[] = [];
+      for (const path of paths) {
+        const text = await readFile(path, "utf8");
+        configured.push(text);
+        for (const comment of scenario.comments)
+          expect(text, path).toContain(comment);
+        expect(text, path).not.toContain('command = "old"');
+        expect(unrelatedValues(text), path).toEqual(
+          unrelatedValues(scenario.original),
+        );
+        expect(parseToml(text), path).toHaveProperty("mcp_servers.rea.command");
+        expect(await readFile(`${path}.rea.backup`, "utf8"), path).toBe(
+          scenario.original,
+        );
+      }
+
+      const repeated = await run([...setup, "--yes"]);
+      expect(repeated.exitCode, repeated.stdout).toBe(0);
+      expect(repeated.json).toMatchObject({
+        status: "ready",
+        appliedActions: [],
+      });
+      for (const [index, path] of paths.entries())
+        expect(await readFile(path, "utf8"), path).toBe(configured[index]);
+
+      const removed = await run(["uninstall", "--json"]);
+      expect(removed.exitCode, removed.stdout).toBe(0);
+      expect(removed.json).toMatchObject({ status: "complete" });
+      const remaining: string[] = [];
+      for (const path of paths) {
+        const text = await readFile(path, "utf8");
+        remaining.push(text);
+        for (const comment of scenario.comments)
+          expect(text, path).toContain(comment);
+        expect(parseToml(text), path).not.toHaveProperty("mcp_servers.rea");
+        expect(unrelatedValues(text), path).toEqual(
+          unrelatedValues(scenario.original),
+        );
+        // Setup ended the file with REA's table; removing it keeps that break.
+        expect(text.endsWith("\n"), path).toBe(true);
+      }
+      // Both TOML clients share one editor, so their text matches.
+      expect(remaining[1]).toBe(remaining[0]);
+    },
+  );

tests/boundary/filesystem/clientLifecycle.test.ts

@@ -262,10 +262,12 @@ describe("Grok Build server table edits", () => {
       'command = "blender"',
       'args = ["--mcp"]',
       "",
+      "# about rea",
       "[mcp_servers.rea]",
       'command = "old"',
       'args = ["mcp"]',
       "",
+      "# rea environment",
       "[mcp_servers.rea.env]",
       'OLD = "1"',
       "",
@@ -276,6 +278,8 @@ describe("Grok Build server table edits", () => {
       "[privacy]",
       "share = false",
       "",
+      "# trailing note",
+      "",
     ].join("\r\n");
     await writeFile(configPath, original);
     const environment = {
@@ -292,7 +296,14 @@ describe("Grok Build server table edits", () => {
     ).toMatchObject({ status: "configured" });
     const configured = await readFile(configPath, "utf8");
     expect(configured).toContain("\r\n");
-    expect(configured).toContain("# keep this comment");
+    expect(configured.replaceAll("\r\n", "")).not.toContain("\n");
+    for (const comment of [
+      "# keep this comment",
+      "# about rea",
+      "# rea environment",
+      "# trailing note",
+    ])
+      expect(configured).toContain(`${comment}\r\n`);
     expect(configured).toContain('disabled_mcp_servers = ["xapi"]');
     expect(configured).toContain('default_model = "grok-4"');
     expect(configured).toContain("[mcp_servers.blender]");
@@ -337,7 +348,13 @@ describe("Grok Build server table edits", () => {
       "complete",
     );
     const removed = await readFile(configPath, "utf8");
-    expect(removed).toContain("# keep this comment");
+    for (const comment of [
+      "# keep this comment",
+      "# about rea",
+      "# rea environment",
+      "# trailing note",
+    ])
+      expect(removed).toContain(`${comment}\r\n`);
     expect(removed).toContain("[mcp_servers.blender]");
     expect(removed).toContain("[[marketplace.sources]]");
     expect(removed).toContain('url = "https://example.test/market"');