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 41a1fd8f9..df5e14b00 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 @@ -320,6 +320,105 @@ 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 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, + 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-shrink-1", + name: "glob", + mcpServer: "filesystem", + arguments: { pattern: "**/*" }, + }; + + await handleAgentEvent(host, { type: "message_start", message: makeAssistant([]) } as any); + + // Sub-turn 1: generate longer provisional text content. + await handleAgentEvent( + host, + { + type: "message_update", + message: makeAssistant([{ type: "text", text: "Old provisional preface." }, { type: "text", text: "More old text." }]), + assistantMessageEvent: { + type: "text_delta", + contentIndex: 1, + delta: "More old text.", + partial: makeAssistant([{ type: "text", text: "Old provisional preface." }, { type: "text", text: "More old text." }]), + }, + } as any, + ); + assert.equal(host.chatContainer.children.length, 1, "first sub-turn text run should render"); + + // Sub-turn 2 starts (content shrink): old component is orphaned by design. + await handleAgentEvent( + host, + { + type: "message_update", + message: makeAssistant([{ type: "text", text: "New provisional text before tool." }]), + assistantMessageEvent: { + type: "text_delta", + contentIndex: 0, + delta: "New provisional text before tool.", + partial: makeAssistant([{ type: "text", text: "New provisional text before tool." }]), + }, + } as any, + ); + 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. + await handleAgentEvent( + host, + { + type: "message_update", + message: makeAssistant([{ type: "text", text: "New provisional text before tool." }, mcpTool]), + assistantMessageEvent: { + type: "toolcall_end", + contentIndex: 1, + toolCall: { + ...mcpTool, + externalResult: { + content: [{ type: "text", text: "glob output" }], + details: {}, + isError: false, + }, + }, + partial: makeAssistant([{ type: "text", text: "New provisional text before tool." }, mcpTool]), + }, + } 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"); + + const finalContent = [mcpTool, { type: "text", text: "Final visible question?" }]; + await handleAgentEvent( + host, + { + type: "message_update", + message: makeAssistant(finalContent), + assistantMessageEvent: { + type: "text_delta", + contentIndex: 1, + delta: "Final visible question?", + partial: makeAssistant(finalContent), + }, + } as any, + ); + assert.equal(host.chatContainer.children.length, 2); + assert.equal(host.chatContainer.children[0]?.constructor?.name, "ToolExecutionComponent"); + assert.equal(host.chatContainer.children[1]?.constructor?.name, "AssistantMessageComponent"); + + await handleAgentEvent(host, { type: "message_end", message: makeAssistant(finalContent) } as any); +}); + test("chat-controller pins latest assistant text above editor when tool calls are present", async () => { (globalThis as any)[Symbol.for("@gsd/pi-coding-agent:theme")] = { fg: (_key: string, text: string) => text, 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 556c87c54..9c606363f 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 @@ -20,6 +20,10 @@ type RenderedSegment = | { kind: "tool"; contentIndex: number; component: ToolExecutionComponent }; let renderedSegments: RenderedSegment[] = []; +// When providers reuse one assistant lifecycle across internal sub-turns, +// a content[] shrink resets renderedSegments. Keep the displaced segments so +// claude-code MCP pruning can remove stale provisional text later. +let orphanedSegments: RenderedSegment[] = []; function hasVisibleAssistantContent(message: { content: Array }): boolean { return message.content.some( @@ -93,6 +97,7 @@ export async function handleAgentEvent(host: InteractiveModeStateHost & { lastPinnedText = ""; hasToolsInTurn = false; renderedSegments = []; + orphanedSegments = []; if (pinnedBorder) pinnedBorder.stopSpinner(); pinnedBorder = undefined; pinnedTextComponent = undefined; @@ -113,6 +118,7 @@ export async function handleAgentEvent(host: InteractiveModeStateHost & { lastPinnedText = ""; hasToolsInTurn = false; renderedSegments = []; + orphanedSegments = []; lastContentLength = 0; if (pinnedBorder) pinnedBorder.stopSpinner(); pinnedBorder = undefined; @@ -226,6 +232,7 @@ export async function handleAgentEvent(host: InteractiveModeStateHost & { // content (#4144 regression). Prior sub-turn children stay in // chatContainer as frozen history; new segments append after them. if (contentBlocks.length < lastContentLength) { + orphanedSegments = [...renderedSegments]; renderedSegments = []; lastPinnedText = ""; lastProcessedContentIndex = 0; @@ -346,6 +353,20 @@ export async function handleAgentEvent(host: InteractiveModeStateHost & { // superseded by post-tool output. Prune stale text-run segments so // the final assistant output remains below tool output. if (shouldDropPreToolText && firstToolIdx >= 0) { + if (orphanedSegments.length > 0) { + const remainingOrphans: RenderedSegment[] = []; + for (const orphan of orphanedSegments) { + if (orphan.kind === "text-run") { + host.chatContainer.removeChild(orphan.component); + if (host.streamingComponent === orphan.component) { + host.streamingComponent = undefined; + } + continue; + } + remainingOrphans.push(orphan); + } + orphanedSegments = remainingOrphans; + } const desiredTextStarts = new Set( desired .filter((seg): seg is Extract => seg.kind === "text-run") @@ -536,6 +557,7 @@ export async function handleAgentEvent(host: InteractiveModeStateHost & { host.streamingComponent = undefined; host.streamingMessage = undefined; renderedSegments = []; + orphanedSegments = []; lastContentLength = 0; // Clear pinned output once the message is finalized in the chat // container — prevents duplicate display when the agent continues @@ -599,6 +621,7 @@ export async function handleAgentEvent(host: InteractiveModeStateHost & { host.streamingComponent = undefined; host.streamingMessage = undefined; renderedSegments = []; + orphanedSegments = []; lastContentLength = 0; host.pendingTools.clear(); // Pinned output is only useful while work is actively streaming.