diff --git a/packages/pi-coding-agent/src/core/chat-controller-ordering.test.ts b/packages/pi-coding-agent/src/core/chat-controller-ordering.test.ts index df5e14b00..a22a84043 100644 --- a/packages/pi-coding-agent/src/core/chat-controller-ordering.test.ts +++ b/packages/pi-coding-agent/src/core/chat-controller-ordering.test.ts @@ -234,7 +234,7 @@ test("chat-controller renders serverToolUse before trailing text matching conten assert.equal(host.chatContainer.children[1]?.constructor?.name, "AssistantMessageComponent"); }); -test("chat-controller drops provisional pre-tool text for claude-code MCP turns", async () => { +test("chat-controller keeps pre-tool prose visible until post-tool prose arrives, then prunes it", async () => { (globalThis as any)[Symbol.for("@gsd/pi-coding-agent:theme")] = { fg: (_key: string, text: string) => text, bg: (_key: string, text: string) => text, @@ -273,7 +273,7 @@ test("chat-controller drops provisional pre-tool text for claude-code MCP turns" assert.equal(host.chatContainer.children.length, 1); assert.equal(host.chatContainer.children[0]?.constructor?.name, "AssistantMessageComponent"); - // MCP tool appears; provisional text should be removed from the chat stack. + // MCP tool appears; provisional text should remain visible until post-tool prose exists. await handleAgentEvent( host, { @@ -294,11 +294,16 @@ test("chat-controller drops provisional pre-tool text for claude-code MCP turns" }, } as any, ); - assert.equal(host.chatContainer.children.length, 1, "provisional pre-tool text should be pruned"); - assert.equal(host.chatContainer.children[0]?.constructor?.name, "ToolExecutionComponent"); + assert.equal(host.chatContainer.children.length, 2, "pre-tool prose should remain during tool-only window"); + assert.equal(host.chatContainer.children[0]?.constructor?.name, "AssistantMessageComponent"); + assert.equal(host.chatContainer.children[1]?.constructor?.name, "ToolExecutionComponent"); - // Final assistant output should render below the tool. - const finalContent = [mcpTool, { type: "text", text: "Which missing feature matters most to you?" }]; + // Post-tool prose arrives: pre-tool prose should now be pruned. + const finalContent = [ + { type: "text", text: "Let me inspect the workspace first." }, + mcpTool, + { type: "text", text: "Which missing feature matters most to you?" }, + ]; await handleAgentEvent( host, { @@ -306,7 +311,7 @@ test("chat-controller drops provisional pre-tool text for claude-code MCP turns" message: makeAssistant(finalContent), assistantMessageEvent: { type: "text_delta", - contentIndex: 1, + contentIndex: 2, delta: "Which missing feature matters most to you?", partial: makeAssistant(finalContent), }, @@ -320,6 +325,73 @@ test("chat-controller drops provisional pre-tool text for claude-code MCP turns" await handleAgentEvent(host, { type: "message_end", message: makeAssistant(finalContent) } as any); }); +test("chat-controller keeps pre-tool thinking visible for claude-code MCP turns without post-tool prose", async () => { + (globalThis as any)[Symbol.for("@gsd/pi-coding-agent:theme")] = { + fg: (_key: string, text: string) => text, + bg: (_key: string, text: string) => text, + bold: (text: string) => text, + italic: (text: string) => text, + truncate: (text: string) => text, + }; + + const host = createHost(); + host.getMarkdownThemeWithSettings = () => ({}); + + const mcpTool = { + type: "toolCall", + id: "mcp-tool-thinking-1", + name: "read", + mcpServer: "filesystem", + arguments: { filePath: "/tmp/demo.txt" }, + }; + + await handleAgentEvent(host, { type: "message_start", message: makeAssistant([]) } as any); + + const thinkingOnly = [{ type: "thinking", thinking: "I should inspect the workspace." }]; + await handleAgentEvent( + host, + { + type: "message_update", + message: makeAssistant(thinkingOnly), + assistantMessageEvent: { + type: "thinking_delta", + contentIndex: 0, + delta: "I should inspect the workspace.", + partial: makeAssistant(thinkingOnly), + }, + } as any, + ); + assert.equal(host.chatContainer.children.length, 1); + assert.equal(host.chatContainer.children[0]?.constructor?.name, "AssistantMessageComponent"); + + await handleAgentEvent( + host, + { + type: "message_update", + message: makeAssistant([thinkingOnly[0], mcpTool]), + assistantMessageEvent: { + type: "toolcall_end", + contentIndex: 1, + toolCall: { + ...mcpTool, + externalResult: { + content: [{ type: "text", text: "file preview" }], + details: {}, + isError: false, + }, + }, + partial: makeAssistant([thinkingOnly[0], mcpTool]), + }, + } as any, + ); + + assert.equal(host.chatContainer.children.length, 2, "thinking should remain visible while only tool output is present"); + assert.equal(host.chatContainer.children[0]?.constructor?.name, "AssistantMessageComponent"); + assert.equal(host.chatContainer.children[1]?.constructor?.name, "ToolExecutionComponent"); + + await handleAgentEvent(host, { type: "message_end", message: makeAssistant([thinkingOnly[0], mcpTool]) } as any); +}); + test("chat-controller prunes orphaned provisional text after claude-code sub-turn shrink when MCP tools appear", async () => { (globalThis as any)[Symbol.for("@gsd/pi-coding-agent:theme")] = { fg: (_key: string, text: string) => text, @@ -374,7 +446,7 @@ test("chat-controller prunes orphaned provisional text after claude-code sub-tur ); assert.equal(host.chatContainer.children.length, 2, "shrink keeps prior text until MCP tool context appears"); - // MCP tool appears in sub-turn 2: both old orphaned text and current pre-tool text should be pruned. + // MCP tool appears in sub-turn 2: tool-only windows keep provisional prose visible. await handleAgentEvent( host, { @@ -395,8 +467,10 @@ test("chat-controller prunes orphaned provisional text after claude-code sub-tur }, } as any, ); - assert.equal(host.chatContainer.children.length, 1, "stale text runs should be removed once MCP tool is present"); - assert.equal(host.chatContainer.children[0]?.constructor?.name, "ToolExecutionComponent"); + assert.equal(host.chatContainer.children.length, 3, "stale text runs are deferred until post-tool prose arrives"); + assert.equal(host.chatContainer.children[0]?.constructor?.name, "AssistantMessageComponent"); + assert.equal(host.chatContainer.children[1]?.constructor?.name, "AssistantMessageComponent"); + assert.equal(host.chatContainer.children[2]?.constructor?.name, "ToolExecutionComponent"); const finalContent = [mcpTool, { type: "text", text: "Final visible question?" }]; await handleAgentEvent( diff --git a/packages/pi-coding-agent/src/modes/interactive/components/assistant-message.ts b/packages/pi-coding-agent/src/modes/interactive/components/assistant-message.ts index 996354512..1b32fa6c7 100644 --- a/packages/pi-coding-agent/src/modes/interactive/components/assistant-message.ts +++ b/packages/pi-coding-agent/src/modes/interactive/components/assistant-message.ts @@ -89,6 +89,10 @@ export class AssistantMessageComponent extends Container { ); const hasTextContent = message.content.some((c) => c.type === "text" && c.text.trim().length > 0); const hasToolContent = message.content.some((c) => c.type === "toolCall" || c.type === "serverToolUse"); + // Claude Code often emits long reasoning blocks ahead of user-visible text/tool + // output in the same lifecycle. Keep chat output visible without requiring a + // manual thinking toggle every turn. + const shouldCapThinking = hasTextContent || hasToolContent || message.provider === "claude-code"; if (hasVisibleContent) { this.contentContainer.addChild(new Spacer(1)); @@ -122,7 +126,7 @@ export class AssistantMessageComponent extends Container { }); // Keep visible chat output readable when thinking traces are long. // Tool-bearing turns can stream text in a later assistant message. - if (hasTextContent || hasToolContent) { + if (shouldCapThinking) { thinkingMarkdown.maxLines = 8; } this.contentContainer.addChild(thinkingMarkdown); diff --git a/packages/pi-coding-agent/src/modes/interactive/controllers/chat-controller.ts b/packages/pi-coding-agent/src/modes/interactive/controllers/chat-controller.ts index 9c606363f..a69fceb1a 100644 --- a/packages/pi-coding-agent/src/modes/interactive/controllers/chat-controller.ts +++ b/packages/pi-coding-agent/src/modes/interactive/controllers/chat-controller.ts @@ -16,7 +16,13 @@ let lastContentLength = 0; // --- Segment walker state (per streaming assistant turn) --- type RenderedSegment = - | { kind: "text-run"; startIndex: number; endIndex: number; component: AssistantMessageComponent } + | { + kind: "text-run"; + startIndex: number; + endIndex: number; + contentType: "text" | "thinking"; + component: AssistantMessageComponent; + } | { kind: "tool"; contentIndex: number; component: ToolExecutionComponent }; let renderedSegments: RenderedSegment[] = []; @@ -319,44 +325,75 @@ export async function handleAgentEvent(host: InteractiveModeStateHost & { } return false; }); - const shouldDropPreToolText = isClaudeCodeProvider && hasMcpToolBlock; const firstToolIdx = blocks.findIndex((b: any) => b.type === "toolCall" || b.type === "serverToolUse"); + const hasPostToolText = firstToolIdx >= 0 + && blocks.some( + (b: any, idx: number) => ( + idx > firstToolIdx + && b?.type === "text" + && typeof b?.text === "string" + && b.text.trim().length > 0 + ), + ); + // Only prune provisional pre-tool prose after post-tool prose exists, + // so MCP tool-only windows do not blank the assistant content. + const shouldDropPreToolProse = isClaudeCodeProvider && hasMcpToolBlock && hasPostToolText; type DesiredSegment = - | { kind: "text-run"; startIndex: number; endIndex: number } + | { kind: "text-run"; startIndex: number; endIndex: number; contentType: "text" | "thinking" } | { kind: "tool"; contentIndex: number; toolId: string }; const desired: DesiredSegment[] = []; let runStart = -1; + let runEnd = -1; + let runType: "text" | "thinking" | undefined; + const closeRun = () => { + if (runStart !== -1 && runType) { + desired.push({ kind: "text-run", startIndex: runStart, endIndex: runEnd, contentType: runType }); + runStart = -1; + runEnd = -1; + runType = undefined; + } + }; for (let i = 0; i < blocks.length; i++) { const b = blocks[i]; - const isText = b.type === "text" || b.type === "thinking"; + const blockType = b.type === "text" || b.type === "thinking" ? b.type : undefined; + const isTextLike = blockType === "text" || blockType === "thinking"; const isTool = b.type === "toolCall" || b.type === "serverToolUse"; - if (isText) { - if (shouldDropPreToolText && firstToolIdx >= 0 && i < firstToolIdx) { - continue; + // For Claude Code MCP turns, prune only pre-tool prose, never thinking. + const shouldSkipProse = shouldDropPreToolProse && firstToolIdx >= 0 && i < firstToolIdx && blockType === "text"; + if (shouldSkipProse) { + closeRun(); + continue; + } + if (isTextLike) { + if (runStart === -1) { + runStart = i; + runEnd = i; + runType = blockType; + } else if (runType !== blockType) { + closeRun(); + runStart = i; + runEnd = i; + runType = blockType; + } else { + runEnd = i; } - if (runStart === -1) runStart = i; } else { - if (runStart !== -1) { - desired.push({ kind: "text-run", startIndex: runStart, endIndex: i - 1 }); - runStart = -1; - } + closeRun(); if (isTool) { desired.push({ kind: "tool", contentIndex: i, toolId: b.id }); } } } - if (runStart !== -1) { - desired.push({ kind: "text-run", startIndex: runStart, endIndex: blocks.length - 1 }); - } + closeRun(); // Claude Code MCP can emit provisional pre-tool prose that gets // superseded by post-tool output. Prune stale text-run segments so // the final assistant output remains below tool output. - if (shouldDropPreToolText && firstToolIdx >= 0) { + if (shouldDropPreToolProse && firstToolIdx >= 0) { if (orphanedSegments.length > 0) { const remainingOrphans: RenderedSegment[] = []; for (const orphan of orphanedSegments) { - if (orphan.kind === "text-run") { + if (orphan.kind === "text-run" && orphan.contentType === "text") { host.chatContainer.removeChild(orphan.component); if (host.streamingComponent === orphan.component) { host.streamingComponent = undefined; @@ -367,10 +404,10 @@ export async function handleAgentEvent(host: InteractiveModeStateHost & { } orphanedSegments = remainingOrphans; } - const desiredTextStarts = new Set( + const desiredTextKeys = new Set( desired .filter((seg): seg is Extract => seg.kind === "text-run") - .map((seg) => seg.startIndex), + .map((seg) => `${seg.contentType}:${seg.startIndex}`), ); const desiredToolIndices = new Set( desired @@ -379,7 +416,11 @@ export async function handleAgentEvent(host: InteractiveModeStateHost & { ); const nextRendered: RenderedSegment[] = []; for (const seg of renderedSegments) { - if (seg.kind === "text-run" && !desiredTextStarts.has(seg.startIndex)) { + if ( + seg.kind === "text-run" + && seg.contentType === "text" + && !desiredTextKeys.has(`${seg.contentType}:${seg.startIndex}`) + ) { host.chatContainer.removeChild(seg.component); if (host.streamingComponent === seg.component) { host.streamingComponent = undefined; @@ -411,7 +452,7 @@ export async function handleAgentEvent(host: InteractiveModeStateHost & { } else { // text-run segment const existing = renderedSegments.find( - (s) => s.kind === "text-run" && s.startIndex === seg.startIndex, + (s) => s.kind === "text-run" && s.startIndex === seg.startIndex && s.contentType === seg.contentType, ); if (!existing) { const comp = new AssistantMessageComponent( @@ -422,7 +463,13 @@ export async function handleAgentEvent(host: InteractiveModeStateHost & { { startIndex: seg.startIndex, endIndex: seg.endIndex }, ); host.chatContainer.addChild(comp); - renderedSegments.push({ kind: "text-run", startIndex: seg.startIndex, endIndex: seg.endIndex, component: comp }); + renderedSegments.push({ + kind: "text-run", + startIndex: seg.startIndex, + endIndex: seg.endIndex, + contentType: seg.contentType, + component: comp, + }); host.streamingComponent = comp; } } @@ -433,7 +480,9 @@ export async function handleAgentEvent(host: InteractiveModeStateHost & { for (const seg of renderedSegments) { if (seg.kind === "text-run") { // Find corresponding desired segment to get current endIndex - const d = desired.find((ds) => ds.kind === "text-run" && ds.startIndex === seg.startIndex); + const d = desired.find( + (ds) => ds.kind === "text-run" && ds.startIndex === seg.startIndex && ds.contentType === seg.contentType, + ); if (d && d.kind === "text-run" && d.endIndex !== seg.endIndex) { seg.endIndex = d.endIndex; seg.component.setRange({ startIndex: seg.startIndex, endIndex: seg.endIndex }); diff --git a/src/tests/assistant-message-thinking-visibility.test.ts b/src/tests/assistant-message-thinking-visibility.test.ts index 2821cf64a..0891f1223 100644 --- a/src/tests/assistant-message-thinking-visibility.test.ts +++ b/src/tests/assistant-message-thinking-visibility.test.ts @@ -34,7 +34,13 @@ test("assistant-message caps thinking block height when text content is present" assert.match( src, - /if \(hasTextContent \|\| hasToolContent\)\s*\{\s*thinkingMarkdown\.maxLines = 8;\s*\}/s, - "assistant-message should cap visible thinking lines when assistant text exists or tool blocks are present", + /const shouldCapThinking = hasTextContent \|\| hasToolContent \|\| message\.provider === "claude-code";/, + "assistant-message should derive a cap policy that also covers claude-code long reasoning traces", + ); + + assert.match( + src, + /if \(shouldCapThinking\)\s*\{\s*thinkingMarkdown\.maxLines = 8;\s*\}/s, + "assistant-message should cap visible thinking lines when the cap policy is active", ); });