From 2dbe3196abd874fdf9f34606855371d2a1b02dfd Mon Sep 17 00:00:00 2001 From: li-xiu-qi Date: Wed, 21 Oct 2026 23:05:18 +0800 Subject: [PATCH] fix(cli): render thinking blocks above the answer when the model streams them last Provider streaming assembles message.content in arrival order, so turns that flush the reasoning stream after the answer text rendered the thinking block below the answer while other turns rendered it above. Normalize the render path: when no thinking block precedes the first answer block, hoist thinking ahead of the answer in both the pi and step skins. Alternating runs are left untouched. --- .../ui/view/transcript/assistant-message.ts | 49 ++++++++++-- .../src/ui/view/transcript/step-message.ts | 4 +- apps/cli/test/assistant-message.test.ts | 74 ++++++++++++++++- apps/cli/test/step-message.test.ts | 80 +++++++++++++++++++ 4 files changed, 199 insertions(+), 8 deletions(-) diff --git a/apps/cli/src/ui/view/transcript/assistant-message.ts b/apps/cli/src/ui/view/transcript/assistant-message.ts index 9525a482c..333ef824f 100644 --- a/apps/cli/src/ui/view/transcript/assistant-message.ts +++ b/apps/cli/src/ui/view/transcript/assistant-message.ts @@ -21,6 +21,43 @@ type ZoneCache = { width: number; native: string[]; lines: string[] }; /** A `Markdown` instance kept alive across deltas of the same message. */ type MarkdownSlot = { kind: "text" | "thinking"; markdown: Markdown }; +/** + * Normalize a provider-assembled content array so that thinking blocks always + * render before the answer. Most models emit the reasoning stream first, but + * some turns flush it only after the answer text, which would otherwise paint + * the thinking block below the answer in those turns and above it in every + * other turn of the same conversation. + * + * Only the degenerate text/thinking shape is rewritten: every visible answer + * block precedes every visible thinking block, with no tool-call boundary to + * preserve. In that case the thinking suffix is hoisted ahead of the answer + * (stable within each group). Interleaved answers and thinking stay in their + * original order, so a mid-turn reasoning block never moves across an answer + * or tool call that it followed. + */ +export function withThinkingBeforeAnswer(content: AssistantMessage["content"]): AssistantMessage["content"] { + let firstText = -1; + let firstThinking = -1; + let lastText = -1; + for (let i = 0; i < content.length; i++) { + const block = content[i]; + if (block === undefined) continue; + if (block.type === "toolCall") return content; + if (block.type === "text" && block.text.trim().length > 0) { + if (firstText === -1) firstText = i; + lastText = i; + } + if (firstThinking === -1 && block.type === "thinking" && block.thinking.trim().length > 0) firstThinking = i; + } + if (firstText === -1 || firstThinking === -1 || firstThinking < firstText || lastText > firstThinking) { + return content; + } + return [ + ...content.filter((block) => block.type === "thinking"), + ...content.filter((block) => block.type !== "thinking"), + ]; +} + /** * Component that renders a complete assistant message */ @@ -150,17 +187,19 @@ export class AssistantMessageComponent extends Container { this.contentContainer.addChild(new Spacer(1)); } + const orderedContent = withThinkingBeforeAnswer(message.content); + // Render content in order - for (let i = 0; i < message.content.length; i++) { - const content = message.content[i]; + for (let i = 0; i < orderedContent.length; i++) { + const content = orderedContent[i]; if (content.type === "text" && content.text.trim()) { // Assistant text messages with no background - trim the text // Set paddingY=0 to avoid extra spacing before tool executions this.contentContainer.addChild(this.textMarkdown(content.text.trim())); } else if (content.type === "thinking") { const thinkingBlocks: string[] = []; - for (; i < message.content.length; i++) { - const thinkingContent = message.content[i]; + for (; i < orderedContent.length; i++) { + const thinkingContent = orderedContent[i]; if (thinkingContent.type !== "thinking") { break; } @@ -177,7 +216,7 @@ export class AssistantMessageComponent extends Container { // Add spacing only when another visible assistant content block follows. // This avoids a superfluous blank line before separately-rendered tool execution blocks. - const hasVisibleContentAfter = message.content + const hasVisibleContentAfter = orderedContent .slice(i + 1) .some((c) => (c.type === "text" && c.text.trim()) || (c.type === "thinking" && c.thinking.trim())); diff --git a/apps/cli/src/ui/view/transcript/step-message.ts b/apps/cli/src/ui/view/transcript/step-message.ts index dbba2ff95..1176b9932 100644 --- a/apps/cli/src/ui/view/transcript/step-message.ts +++ b/apps/cli/src/ui/view/transcript/step-message.ts @@ -3,7 +3,7 @@ import { getMarkdownTheme, theme } from "@step-harness/coding-agent"; import { type MarkdownTheme, stripTerminalSequences, truncateToWidth, visibleWidth } from "@step-harness/pi-tui"; import type { AssistantMessage } from "@step-harness/providers"; import { estimateTokens, formatElapsedTime, formatTokens } from "../chrome/footer.ts"; -import { AssistantMessageComponent } from "./assistant-message.ts"; +import { AssistantMessageComponent, withThinkingBeforeAnswer } from "./assistant-message.ts"; import { UserMessageComponent } from "./user-message.ts"; /** OSC 133 prompt-zone markers emitted by pi's message components. */ @@ -107,7 +107,7 @@ function partitionAssistantContent(message: AssistantMessage): AssistantContentR const runs: AssistantContentRun[] = []; let previousKind: AssistantContentRun["kind"] | undefined; - for (const content of message.content) { + for (const content of withThinkingBeforeAnswer(message.content)) { if (content.type !== "thinking" && content.type !== "text") { previousKind = undefined; continue; diff --git a/apps/cli/test/assistant-message.test.ts b/apps/cli/test/assistant-message.test.ts index 161fe360f..1accd0d5d 100644 --- a/apps/cli/test/assistant-message.test.ts +++ b/apps/cli/test/assistant-message.test.ts @@ -1,6 +1,9 @@ import type { AssistantMessage } from "@step-harness/providers"; import { describe, expect, test } from "vitest"; -import { AssistantMessageComponent } from "../src/ui/view/transcript/assistant-message.ts"; +import { + AssistantMessageComponent, + withThinkingBeforeAnswer, +} from "../src/ui/view/transcript/assistant-message.ts"; import { UserMessageComponent } from "../src/ui/view/transcript/user-message.ts"; import { initTheme } from "../../../packages/coding-agent/src/theme/theme.ts"; import { stripAnsi } from "../../../packages/coding-agent/src/utils/ansi.ts"; @@ -91,6 +94,75 @@ describe("AssistantMessageComponent", () => { expect(rendered).toContain("answer"); }); + test("hoists thinking ahead of the answer when the model streams it last", () => { + initTheme("dark"); + + const component = new AssistantMessageComponent( + createAssistantMessage([ + { type: "text", text: "answer text" }, + { type: "thinking", thinking: "late reasoning" }, + ]), + ); + const lines = stripAnsi(component.render(80).join("\n")).split("\n"); + const thinkingLine = lines.findIndex((line) => line.includes("late reasoning")); + const answerLine = lines.findIndex((line) => line.includes("answer text")); + + expect(thinkingLine).toBeGreaterThanOrEqual(0); + expect(answerLine).toBeGreaterThanOrEqual(0); + expect(thinkingLine).toBeLessThan(answerLine); + }); + + test("preserves answer segments when thinking arrives between them", () => { + const content: AssistantMessage["content"] = [ + { type: "text", text: "first answer" }, + { type: "thinking", thinking: "middle reasoning" }, + { type: "text", text: "second answer" }, + ]; + expect(withThinkingBeforeAnswer(content)).toBe(content); + + initTheme("dark"); + const component = new AssistantMessageComponent(createAssistantMessage(content)); + const lines = stripAnsi(component.render(80).join("\n")).split("\n"); + const firstAnswer = lines.findIndex((line) => line.includes("first answer")); + const thinking = lines.findIndex((line) => line.includes("middle reasoning")); + const secondAnswer = lines.findIndex((line) => line.includes("second answer")); + expect(firstAnswer).toBeLessThan(thinking); + expect(thinking).toBeLessThan(secondAnswer); + }); + + test("does not move late thinking across a tool-call boundary", () => { + const content: AssistantMessage["content"] = [ + { type: "text", text: "answer before tool" }, + { type: "toolCall", id: "tool-1", name: "read", arguments: { path: "file.txt" } }, + { type: "thinking", thinking: "reasoning after tool" }, + ]; + expect(withThinkingBeforeAnswer(content)).toBe(content); + }); + + test("keeps segmented reasoning in place when thinking already leads", () => { + initTheme("dark"); + + const component = new AssistantMessageComponent( + createAssistantMessage([ + { type: "thinking", thinking: "first reasoning" }, + { type: "text", text: "first answer" }, + { type: "thinking", thinking: "follow-up reasoning" }, + { type: "text", text: "final answer" }, + ]), + ); + const lines = stripAnsi(component.render(80).join("\n")).split("\n"); + + expect(lines.findIndex((line) => line.includes("first reasoning"))).toBeLessThan( + lines.findIndex((line) => line.includes("first answer")), + ); + expect(lines.findIndex((line) => line.includes("follow-up reasoning"))).toBeGreaterThan( + lines.findIndex((line) => line.includes("first answer")), + ); + expect(lines.findIndex((line) => line.includes("final answer"))).toBeGreaterThan( + lines.findIndex((line) => line.includes("follow-up reasoning")), + ); + }); + test("uses configured output padding for text and thinking", () => { initTheme("dark"); diff --git a/apps/cli/test/step-message.test.ts b/apps/cli/test/step-message.test.ts index 683ebb254..a24a76f02 100644 --- a/apps/cli/test/step-message.test.ts +++ b/apps/cli/test/step-message.test.ts @@ -176,6 +176,86 @@ describe("Step message presentation", () => { expect(finalAnswer).toBeGreaterThan(followUpThinking); }); + test("preserves answer segments around a middle thinking block", () => { + initTheme("step-blue"); + const component = new StepAssistantMessageComponent( + assistantMessage([ + { type: "text", text: "first answer" }, + { type: "thinking", thinking: "middle reasoning" }, + { type: "text", text: "second answer" }, + ]), + false, + getMarkdownTheme(), + ); + const lines = plain(component.render(80)); + const firstAnswer = lines.findIndex((line) => line.includes("first answer")); + const thinking = lines.findIndex((line) => line.includes("middle reasoning")); + const secondAnswer = lines.findIndex((line) => line.includes("second answer")); + expect(firstAnswer).toBeLessThan(thinking); + expect(thinking).toBeLessThan(secondAnswer); + }); + + test("hoists thinking ahead of the answer when the model streams it last", () => { + initTheme("step-blue"); + const component = new StepAssistantMessageComponent( + assistantMessage([ + { type: "text", text: "the answer" }, + { type: "thinking", thinking: "late reasoning" }, + ]), + false, + getMarkdownTheme(), + ); + const lines = plain(component.render(80)); + const thinking = lines.findIndex((line) => line.includes("late reasoning")); + const answer = lines.findIndex((line) => line.includes("the answer")); + expect(thinking).toBeGreaterThanOrEqual(0); + expect(answer).toBeGreaterThanOrEqual(0); + expect(thinking).toBeLessThan(answer); + }); + + test("keeps streamed thinking placement consistent across content updates", () => { + initTheme("step-blue"); + const component = new StepAssistantMessageComponent( + assistantMessage([{ type: "text", text: "first answer" }]), + false, + getMarkdownTheme(), + ); + expect(plain(component.render(80)).join("\n")).toContain("• first answer"); + + component.updateContent( + assistantMessage([ + { type: "text", text: "first answer" }, + { type: "thinking", thinking: "late reasoning" }, + ]), + true, + ); + let lines = plain(component.render(80)); + const hoistedThinking = lines.findIndex((line) => line.includes("late reasoning")); + const hoistedAnswer = lines.findIndex((line) => line.includes("first answer")); + expect(hoistedThinking).toBeGreaterThanOrEqual(0); + expect(hoistedAnswer).toBeGreaterThanOrEqual(0); + expect(hoistedThinking).toBeLessThan(hoistedAnswer); + + component.updateContent( + assistantMessage([ + { type: "text", text: "first answer" }, + { type: "thinking", thinking: "late reasoning" }, + { type: "text", text: "second answer" }, + ]), + true, + ); + lines = plain(component.render(80)); + const firstAnswer = lines.findIndex((line) => line.includes("first answer")); + const thinking = lines.findIndex((line) => line.includes("late reasoning")); + const secondAnswer = lines.findIndex((line) => line.includes("second answer")); + expect(firstAnswer).toBeGreaterThanOrEqual(0); + expect(thinking).toBeGreaterThanOrEqual(0); + expect(secondAnswer).toBeGreaterThanOrEqual(0); + expect(firstAnswer).toBeLessThan(thinking); + expect(thinking).toBeLessThan(secondAnswer); + expect(lines.filter((line) => line.includes("late reasoning"))).toHaveLength(1); + }); + test("updates output padding without dropping the first character", () => { initTheme("step-blue"); const assistant = new StepAssistantMessageComponent(