From 2d1081f1cc4237844a70507c97ddef4cb9e378bb Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 13 Apr 2026 05:46:35 +0000 Subject: [PATCH] fix: clean up MCP tool rendering in Claude Code CLI stream MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Strip the `mcp____` prefix from tool_use blocks emitted by the Claude Agent SDK so registered GSD extension renderers (gsd_plan_milestone, gsd_task_complete, etc.) match instead of falling through to the generic JSON-dump fallback. The original server name is preserved on the toolCall block under `mcpServer` for downstream rendering. Tighten the generic ToolExecutionComponent fallback for any remaining prefixed names (third-party MCP servers): show a muted `server·tool` title, render primitive args as compact `key=value` pairs, and truncate output to 10 lines when collapsed. --- .../__tests__/tool-execution.test.ts | 72 ++++++++++++++ .../interactive/components/tool-execution.ts | 94 ++++++++++++++++--- .../claude-code-cli/partial-builder.ts | 57 ++++++++--- .../tests/partial-builder.test.ts | 93 +++++++++++++++++- 4 files changed, 291 insertions(+), 25 deletions(-) diff --git a/packages/pi-coding-agent/src/modes/interactive/components/__tests__/tool-execution.test.ts b/packages/pi-coding-agent/src/modes/interactive/components/__tests__/tool-execution.test.ts index 9b3123fa5..43fd3f7a5 100644 --- a/packages/pi-coding-agent/src/modes/interactive/components/__tests__/tool-execution.test.ts +++ b/packages/pi-coding-agent/src/modes/interactive/components/__tests__/tool-execution.test.ts @@ -27,6 +27,26 @@ function renderTool( return stripAnsi(component.render(120).join("\n")); } +function renderToolCollapsed( + toolName: string, + args: Record, + result?: { + content: Array<{ type: string; text?: string }>; + isError: boolean; + details?: Record; + }, +): string { + const component = new ToolExecutionComponent( + toolName, + args, + {}, + undefined, + { requestRender() {} } as any, + ); + if (result) component.updateResult(result); + return stripAnsi(component.render(120).join("\n")); +} + describe("ToolExecutionComponent", () => { test("renders capitalized Claude Code Bash tool names with bash output instead of generic args JSON", () => { const rendered = renderTool( @@ -51,4 +71,56 @@ describe("ToolExecutionComponent", () => { assert.match(rendered, /hello/); assert.match(rendered, /world/); }); + + test("generic fallback strips mcp____ prefix and shows server·tool title", () => { + const rendered = renderTool( + "mcp__context7__resolve_library_id", + { name: "react" }, + { content: [{ type: "text", text: "react@18.3.1" }], isError: false }, + ); + + assert.match(rendered, /context7\u00b7resolve_library_id/); + assert.doesNotMatch(rendered, /mcp__/); + assert.match(rendered, /name="react"/); + assert.match(rendered, /react@18\.3\.1/); + }); + + test("generic fallback renders compact key=value args for primitive args", () => { + const rendered = renderTool( + "some_unknown_tool", + { count: 3, enabled: true, label: "hello" }, + ); + + assert.match(rendered, /some_unknown_tool/); + assert.match(rendered, /count=3/); + assert.match(rendered, /enabled=true/); + assert.match(rendered, /label="hello"/); + assert.doesNotMatch(rendered, /^\{$/m); + }); + + test("generic fallback truncates long output when collapsed", () => { + const longOutput = Array.from({ length: 25 }, (_, i) => `line ${i + 1}`).join("\n"); + const rendered = renderToolCollapsed( + "mcp__demo__do_thing", + { ok: true }, + { content: [{ type: "text", text: longOutput }], isError: false }, + ); + + assert.match(rendered, /line 1\b/); + assert.match(rendered, /line 10\b/); + assert.doesNotMatch(rendered, /line 20\b/); + assert.match(rendered, /\(15 more lines/); + }); + + test("generic fallback falls back to truncated JSON for complex args", () => { + const rendered = renderTool( + "mcp__demo__nested", + { payload: { nested: { deeply: ["a", "b", "c"] } }, name: "x" }, + ); + + assert.match(rendered, /demo\u00b7nested/); + // Multi-line JSON dump for the complex payload + assert.match(rendered, /"payload"/); + assert.match(rendered, /"nested"/); + }); }); diff --git a/packages/pi-coding-agent/src/modes/interactive/components/tool-execution.ts b/packages/pi-coding-agent/src/modes/interactive/components/tool-execution.ts index 4f7bcb641..a4bf41837 100644 --- a/packages/pi-coding-agent/src/modes/interactive/components/tool-execution.ts +++ b/packages/pi-coding-agent/src/modes/interactive/components/tool-execution.ts @@ -51,6 +51,60 @@ function str(value: unknown): string | null { return null; // Invalid type } +/** + * Split a Claude Code MCP tool name (`mcp____`) into its parts. + * Returns null for non-prefixed names. Duplicated from the claude-code-cli + * extension (parseMcpToolName) so this package doesn't have to import across + * the resources/extensions boundary. + */ +function parseMcpToolName(name: string): { server: string; tool: string } | null { + if (!name.startsWith("mcp__")) return null; + const rest = name.slice("mcp__".length); + const delim = rest.indexOf("__"); + if (delim <= 0 || delim === rest.length - 2) return null; + return { server: rest.slice(0, delim), tool: rest.slice(delim + 2) }; +} + +const COMPACT_ARG_VALUE_LIMIT = 60; +const GENERIC_OUTPUT_PREVIEW_LINES = 10; +const GENERIC_ARGS_JSON_PREVIEW_LINES = 10; + +/** + * Format tool args for the generic-renderer fallback. Produces a one-line + * `k=v, k=v` summary when every value is a primitive that fits inline; falls + * back to a truncated JSON dump for structurally complex args. + */ +function formatCompactArgs(args: unknown, expanded: boolean): string { + if (args == null) return ""; + if (typeof args !== "object") return String(args); + + const entries = Object.entries(args as Record); + if (entries.length === 0) return ""; + + const allPrimitive = entries.every(([, value]) => { + const t = typeof value; + if (t === "number" || t === "boolean") return true; + if (t === "string") return (value as string).length <= COMPACT_ARG_VALUE_LIMIT; + return value == null; + }); + + if (allPrimitive) { + return entries + .map(([key, value]) => { + if (typeof value === "string") return `${key}=${JSON.stringify(value)}`; + if (value == null) return `${key}=null`; + return `${key}=${String(value)}`; + }) + .join(", "); + } + + // Complex args: show truncated JSON. + const lines = JSON.stringify(args, null, 2).split("\n"); + const maxLines = expanded ? lines.length : GENERIC_ARGS_JSON_PREVIEW_LINES; + if (lines.length <= maxLines) return lines.join("\n"); + return lines.slice(0, maxLines).join("\n") + "\n..."; +} + export interface ToolExecutionOptions { showImages?: boolean; // default: true (only used if terminal supports images) } @@ -943,19 +997,37 @@ export class ToolExecutionComponent extends Container { } } } else { - // Generic tool (shouldn't reach here for custom tools) - text = theme.fg("toolTitle", theme.bold(this.toolName)); + // Generic tool / MCP tool without a registered renderer. + // MCP tool names from Claude Code arrive as `mcp____`; + // render the server prefix in muted style so the tool name reads + // cleanly. GSD-registered MCP tools have already had their prefix + // stripped upstream in partial-builder.ts and won't reach this branch. + const parsed = parseMcpToolName(this.toolName); + const displayName = parsed ? parsed.tool : this.toolName; + const serverPrefix = parsed ? theme.fg("muted", `${parsed.server}\u00b7`) : ""; + text = serverPrefix + theme.fg("toolTitle", theme.bold(displayName)); - const contentLines = JSON.stringify(this.args, null, 2).split("\n"); - const maxContentLines = 20; - const truncatedContent = contentLines.slice(0, maxContentLines); - if (contentLines.length > maxContentLines) { - truncatedContent.push("..."); + const argsText = formatCompactArgs(this.args, this.expanded); + if (argsText) { + if (argsText.includes("\n")) { + text += `\n\n${theme.fg("toolOutput", argsText)}`; + } else { + text += " " + theme.fg("toolOutput", argsText); + } } - text += `\n\n${truncatedContent.join("\n")}`; - const output = this.getTextOutput(); - if (output) { - text += `\n${output}`; + + if (this.result) { + const output = this.getTextOutput().trim(); + if (output) { + const lines = output.split("\n"); + const maxLines = this.expanded ? lines.length : GENERIC_OUTPUT_PREVIEW_LINES; + const displayLines = lines.slice(0, maxLines); + const remaining = lines.length - maxLines; + text += `\n\n${displayLines.map((line: string) => theme.fg("toolOutput", line)).join("\n")}`; + if (remaining > 0) { + text += `${theme.fg("muted", `\n... (${remaining} more lines,`)} ${keyHint("expandTools", "to expand")})`; + } + } } } diff --git a/src/resources/extensions/claude-code-cli/partial-builder.ts b/src/resources/extensions/claude-code-cli/partial-builder.ts index c1d011e14..0f52bc220 100644 --- a/src/resources/extensions/claude-code-cli/partial-builder.ts +++ b/src/resources/extensions/claude-code-cli/partial-builder.ts @@ -19,6 +19,49 @@ import type { import { hasXmlParameterTags, repairToolJson } from "@gsd/pi-ai"; import type { BetaContentBlock, BetaRawMessageStreamEvent, NonNullableUsage } from "./sdk-types.js"; +// --------------------------------------------------------------------------- +// MCP tool name parsing +// --------------------------------------------------------------------------- + +/** + * Split a Claude Code MCP tool name (`mcp____`) into its parts. + * Returns null for non-prefixed names so callers can fall through unchanged. + * + * Server names may contain hyphens (`gsd-workflow`); the SDK uses the literal + * `__` delimiter between the server name and the tool name. + */ +export function parseMcpToolName(name: string): { server: string; tool: string } | null { + if (!name.startsWith("mcp__")) return null; + const rest = name.slice("mcp__".length); + const delim = rest.indexOf("__"); + if (delim <= 0 || delim === rest.length - 2) return null; + return { server: rest.slice(0, delim), tool: rest.slice(delim + 2) }; +} + +/** + * Build a GSD ToolCall block from a Claude Code SDK tool_use block, stripping + * the `mcp____` prefix from the name so registered extension renderers + * (which use the unprefixed canonical names) can match. The original server + * name is preserved on the block for diagnostics and rendering. + */ +function toolCallFromBlock( + id: string, + rawName: string, + input: Record, +): ToolCall { + const parsed = parseMcpToolName(rawName); + const toolCall: ToolCall = { + type: "toolCall", + id, + name: parsed ? parsed.tool : rawName, + arguments: input, + }; + if (parsed) { + (toolCall as ToolCall & { mcpServer?: string }).mcpServer = parsed.server; + } + return toolCall; +} + // --------------------------------------------------------------------------- // Content-block mapping helpers // --------------------------------------------------------------------------- @@ -41,12 +84,7 @@ export function mapContentBlock( } satisfies ThinkingContent; case "tool_use": - return { - type: "toolCall", - id: block.id, - name: block.name, - arguments: block.input, - } satisfies ToolCall; + return toolCallFromBlock(block.id, block.name, block.input); case "server_tool_use": return { @@ -183,12 +221,7 @@ export class PartialMessageBuilder { } if (block.type === "tool_use") { this.toolJsonAccum.set(streamIndex, ""); - this.partial.content.push({ - type: "toolCall", - id: block.id, - name: block.name, - arguments: {}, - }); + this.partial.content.push(toolCallFromBlock(block.id, block.name, {})); return { type: "toolcall_start", contentIndex, partial: this.partial }; } if (block.type === "server_tool_use") { diff --git a/src/resources/extensions/claude-code-cli/tests/partial-builder.test.ts b/src/resources/extensions/claude-code-cli/tests/partial-builder.test.ts index 01c853a14..cff2a6830 100644 --- a/src/resources/extensions/claude-code-cli/tests/partial-builder.test.ts +++ b/src/resources/extensions/claude-code-cli/tests/partial-builder.test.ts @@ -1,7 +1,7 @@ import { describe, test } from "node:test"; import assert from "node:assert/strict"; -import { PartialMessageBuilder } from "../partial-builder.ts"; -import type { BetaRawMessageStreamEvent } from "../sdk-types.ts"; +import { mapContentBlock, parseMcpToolName, PartialMessageBuilder } from "../partial-builder.ts"; +import type { BetaContentBlock, BetaRawMessageStreamEvent } from "../sdk-types.ts"; describe("PartialMessageBuilder — malformed tool arguments (#2574)", () => { /** @@ -148,3 +148,92 @@ describe("PartialMessageBuilder — malformed tool arguments (#2574)", () => { } }); }); + +describe("parseMcpToolName", () => { + test("splits mcp____ into parts", () => { + assert.deepEqual( + parseMcpToolName("mcp__gsd-workflow__gsd_plan_milestone"), + { server: "gsd-workflow", tool: "gsd_plan_milestone" }, + ); + }); + + test("preserves server names containing hyphens", () => { + assert.deepEqual( + parseMcpToolName("mcp__my-cool-server__do_thing"), + { server: "my-cool-server", tool: "do_thing" }, + ); + }); + + test("preserves tool names containing underscores", () => { + assert.deepEqual( + parseMcpToolName("mcp__srv__a_b_c_d"), + { server: "srv", tool: "a_b_c_d" }, + ); + }); + + test("returns null for non-prefixed names", () => { + assert.equal(parseMcpToolName("Bash"), null); + assert.equal(parseMcpToolName("gsd_plan_milestone"), null); + }); + + test("returns null for malformed prefixes", () => { + assert.equal(parseMcpToolName("mcp__"), null); + assert.equal(parseMcpToolName("mcp__server"), null); + assert.equal(parseMcpToolName("mcp__server__"), null); + assert.equal(parseMcpToolName("mcp____tool"), null); + }); +}); + +describe("PartialMessageBuilder — MCP tool name normalization", () => { + test("strips mcp____ prefix on content_block_start", () => { + const builder = new PartialMessageBuilder("claude-sonnet-4-20250514"); + const event = builder.handleEvent({ + type: "content_block_start", + index: 0, + content_block: { + type: "tool_use", + id: "tool_1", + name: "mcp__gsd-workflow__gsd_plan_milestone", + input: {}, + }, + } as BetaRawMessageStreamEvent); + + assert.ok(event, "event should not be null"); + assert.equal(event!.type, "toolcall_start"); + if (event!.type === "toolcall_start") { + const toolCall = (event.partial.content[event.contentIndex] as any); + assert.equal(toolCall.name, "gsd_plan_milestone"); + assert.equal(toolCall.mcpServer, "gsd-workflow"); + } + }); + + test("leaves non-MCP tool names untouched", () => { + const builder = new PartialMessageBuilder("claude-sonnet-4-20250514"); + const event = builder.handleEvent({ + type: "content_block_start", + index: 0, + content_block: { type: "tool_use", id: "tool_1", name: "Bash", input: {} }, + } as BetaRawMessageStreamEvent); + + assert.ok(event); + if (event!.type === "toolcall_start") { + const toolCall = (event.partial.content[event.contentIndex] as any); + assert.equal(toolCall.name, "Bash"); + assert.equal(toolCall.mcpServer, undefined); + } + }); + + test("mapContentBlock strips MCP prefix on full tool_use blocks", () => { + const block: BetaContentBlock = { + type: "tool_use", + id: "tool_2", + name: "mcp__gsd-workflow__gsd_task_complete", + input: { taskId: "T001" }, + }; + const mapped = mapContentBlock(block) as any; + assert.equal(mapped.type, "toolCall"); + assert.equal(mapped.name, "gsd_task_complete"); + assert.equal(mapped.mcpServer, "gsd-workflow"); + assert.deepEqual(mapped.arguments, { taskId: "T001" }); + }); +});