From 4f3d486d8f909f65f2f411aa801fd4e8e4466e29 Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Sat, 24 Oct 2026 01:48:35 +0800 Subject: [PATCH] fix(setup): retain TOML comments from replaced registration code --- ...ConfigurationDocument.tomlComments.test.ts | 102 ++++++++++++++++++ .../ClientConfigurationDocument.ts | 53 +++++++-- 2 files changed, 147 insertions(+), 8 deletions(-) create mode 100644 src/application/ClientConfigurationDocument.tomlComments.test.ts diff --git a/src/application/ClientConfigurationDocument.tomlComments.test.ts b/src/application/ClientConfigurationDocument.tomlComments.test.ts new file mode 100644 index 0000000000..9f415643e6 --- /dev/null +++ b/src/application/ClientConfigurationDocument.tomlComments.test.ts @@ -0,0 +1,102 @@ +import { expect, it } from "vitest"; + +import { + parseClientConfiguration, + serializeClientConfiguration, + withClientServers, +} from "./ClientConfigurationDocument.js"; + +const original = [ + 'root = "retained # root string"', + "[mcp_servers.rea] # server header note", + 'command = "old # command token" # command note', + "args = [", + ' "argument # value token", # array note', + "]", + 'description = """', + "# multiline value token", + '""" # multiline note', + "literal = 'literal # value token' # literal note", + "[mcp_servers.rea.tools.binary_session] # tool header note", + 'approval_mode = "approve" # caller approval', + "[mcp_servers.other]", + 'command = "other" # peer note', + "", +].join("\n"); + +for (const format of ["toml", "grok"] as const) + for (const lineBreak of ["\n", "\r\n"] as const) + it.each(["replace", "remove"] as const)( + `retains real comments through %s in ${format} with ${JSON.stringify(lineBreak)}`, + (operation) => { + const source = original.replaceAll("\n", lineBreak); + const parsed = parseClientConfiguration(source, format); + const { rea: removed, ...peers } = parsed.servers; + expect(removed).toBeDefined(); + const servers = + operation === "replace" + ? { ...peers, rea: { command: "new", args: [] } } + : peers; + const result = serializeClientConfiguration( + withClientServers(parsed, servers), + format, + source, + ); + expect(result.ok).toBe(true); + if (!result.ok) throw new Error("TOML edit was refused"); + for (const comment of [ + "# server header note", + "# command note", + "# array note", + "# multiline note", + "# literal note", + "# tool header note", + "# caller approval", + "# peer note", + ]) + expect(result.value.split(comment)).toHaveLength(2); + for (const stringContent of [ + "# command token", + "# value token", + "# multiline value token", + ]) + expect(result.value).not.toContain(stringContent); + expect(parseClientConfiguration(result.value, format).document).toEqual( + { + root: "retained # root string", + mcp_servers: servers, + }, + ); + if (lineBreak === "\r\n") + expect(result.value.replaceAll("\r\n", "")).not.toContain("\n"); + }, + ); + +it("retains comments from dotted REA assignments and nested inline values", () => { + for (const source of [ + 'mcp_servers.rea.command = "old # string token" # dotted note\n', + 'mcp_servers = { rea = { args = ["old # string token", # inline value note\n] }, other = { command = "other" } }\n', + ]) { + const parsed = parseClientConfiguration(source, "toml"); + const { rea: removed, ...servers } = parsed.servers; + expect(removed).toBeDefined(); + const result = serializeClientConfiguration( + withClientServers(parsed, servers), + "toml", + source, + ); + expect(result.ok).toBe(true); + if (!result.ok) throw new Error("TOML removal was refused"); + expect(result.value).toContain( + source.includes("# dotted note") + ? "# dotted note" + : "# inline value note", + ); + expect(result.value).not.toContain("# string token"); + if (source.includes("# inline value note")) + expect(result.value.split("\n")[0]).not.toContain("# inline value note"); + expect(parseClientConfiguration(result.value, "toml").servers).toEqual( + servers, + ); + } +}); diff --git a/src/application/ClientConfigurationDocument.ts b/src/application/ClientConfigurationDocument.ts index ad90e4d55a..f66f167911 100644 --- a/src/application/ClientConfigurationDocument.ts +++ b/src/application/ClientConfigurationDocument.ts @@ -771,11 +771,14 @@ const editInlineServerTable = ( table: string, serverKey: string, entry: unknown, -): string | undefined => { +): + | { readonly text: string; readonly comments: readonly string[] } + | undefined => { if (!table.startsWith("{")) return undefined; const end = endOfContainer(table, 0); if (table[end - 1] !== "}") return undefined; const pairs: InlinePair[] = []; + const comments: string[] = []; let index = 1; let trailing = ""; while (index < end - 1) { @@ -790,10 +793,14 @@ const editInlineServerTable = ( if (table[equals] !== "=") return undefined; const valueEnd = endOfTomlValue(table, equals + 1); if (valueEnd > end - 1) return undefined; + const pairText = table.slice(trivia.next, valueEnd); + const drop = key.segments[0] === serverKey; + if (drop) + for (const comment of tomlCommentLines(pairText)) comments.push(comment); pairs.push({ leading: trivia.text, - text: table.slice(trivia.next, valueEnd), - drop: key.segments[0] === serverKey, + text: pairText, + drop, }); index = valueEnd; } @@ -802,8 +809,8 @@ const editInlineServerTable = ( ? undefined : `${tomlInlineKey(serverKey)} = ${tomlInlineValue(entry)}`; if (insertion === undefined && pairs.every((pair) => !pair.drop)) - return table; - return rebuildInlineTable(pairs, trailing, insertion); + return { text: table, comments }; + return { text: rebuildInlineTable(pairs, trailing, insertion), comments }; }; const rebuildInlineTable = ( @@ -992,10 +999,32 @@ const withoutTrailingLineFeeds = (text: string): string => { return text.slice(0, end); }; +/** Retain comment trivia from removed code without promoting string contents. */ +const tomlCommentLines = (text: string): readonly string[] => { + const comments: string[] = []; + let index = 0; + while (index < text.length) { + if (text.startsWith('"""', index) || text.startsWith("'''", index)) { + const basic = text[index] === '"'; + index = + closeMultiline(text, index + 3, basic ? '"""' : "'''", basic) ?? + text.length; + } else if (text[index] === '"' || text[index] === "'") { + index = parseQuotedKey(text, index)?.next ?? text.length; + } else if (text[index] === "#") { + const end = text.indexOf("\n", index); + comments.push(text.slice(index, end < 0 ? text.length : end)); + index = end < 0 ? text.length : end + 1; + } else index += 1; + } + return comments; +}; + /** * 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 + * inline comments on replaced code become standalone comment lines. + * 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. */ @@ -1055,6 +1084,7 @@ const upsertTomlServerSection = ( tablePath = header.path; if (isTomlServerTable(header.path, serverKey)) { discardPending(); + appendLines(tomlCommentLines(physical)); skipping = true; lineIndex += 1; continue; @@ -1083,6 +1113,7 @@ const upsertTomlServerSection = ( endLine += 1; if (skipping) { discardPending(); + appendLines(tomlCommentLines(text.slice(lineStart, statementEnd))); lineIndex = endLine + 1; mode = "none"; continue; @@ -1132,10 +1163,16 @@ const upsertTomlServerSection = ( const replacement = edited === undefined ? text.slice(lineStart, statementEnd) - : `${prefix}${edited}${suffix}`; + : `${prefix}${edited.text}${suffix}`; appendLines(replacement.split("\n")); + // TOML 1.0 inline tables cannot gain standalone comments between + // their fields. Retain removed value comments after the assignment. + if (edited !== undefined) appendLines(edited.comments); inlineServerTable = true; - } else discardPending(); + } else { + discardPending(); + appendLines(tomlCommentLines(text.slice(lineStart, statementEnd))); + } } else { flushPending(); for (let index = lineIndex; index <= endLine; index += 1)