Merge pull request #4100 from jeremymcs/claude/cleanup-mcp-stream-output-9uCeK

Improve MCP tool rendering with name parsing and compact args

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<string, unknown>,
+	result?: {
+		content: Array<{ type: string; text?: string }>;
+		isError: boolean;
+		details?: Record<string, unknown>;
+	},
+): 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__<server>__ 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"/);
+	});
 });

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__<server>__<tool>`) 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<string, unknown>);
+	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)
 }
@@ -990,19 +1044,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__<server>__<tool>`;
+			// 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")})`;
+					}
+				}
 			}
 		}
 

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__<server>__<tool>`) 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__<server>__` 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<string, unknown>,
+): 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") {

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__<server>__<tool> 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__<server>__ 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" });
+	});
+});