From a1a645b175bce64d90f77188ae894c2a2d000df7 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 | 43 ++++++++++++++++--- .../src/ui/view/transcript/step-message.ts | 4 +- apps/cli/test/assistant-message.test.ts | 42 ++++++++++++++++++ apps/cli/test/step-message.test.ts | 18 ++++++++ 4 files changed, 100 insertions(+), 7 deletions(-) diff --git a/apps/cli/src/ui/view/transcript/assistant-message.ts b/apps/cli/src/ui/view/transcript/assistant-message.ts index 0cf39aaf..4fec0874 100644 --- a/apps/cli/src/ui/view/transcript/assistant-message.ts +++ b/apps/cli/src/ui/view/transcript/assistant-message.ts @@ -18,6 +18,37 @@ const OSC133_ZONE_FINAL = "\x1b]133;C\x07"; /** Cached zone-marked output, keyed by the child lines it was derived from. */ type ZoneCache = { width: number; native: string[]; lines: string[] }; +/** + * 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 shape is rewritten: when no thinking block precedes the + * first answer block, every thinking block is hoisted ahead of the answer + * (stable within each group). Alternating runs (thinking, answer, thinking, + * answer) already lead with thinking and are returned untouched, preserving + * segmented reasoning where each answer follows its own thinking. + */ +export function withThinkingBeforeAnswer(content: AssistantMessage["content"]): AssistantMessage["content"] { + let firstText = -1; + let firstThinking = -1; + for (let i = 0; i < content.length; i++) { + const block = content[i]; + if (block === undefined) continue; + if (firstText === -1 && block.type === "text" && block.text.trim().length > 0) firstText = i; + if (firstThinking === -1 && block.type === "thinking" && block.thinking.trim().length > 0) firstThinking = i; + } + if (firstText === -1 || firstThinking < firstText) { + return content; + } + return [ + ...content.filter((block) => block.type === "thinking"), + ...content.filter((block) => block.type !== "thinking"), + ]; +} + /** * Component that renders a complete assistant message */ @@ -130,9 +161,11 @@ 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 @@ -143,8 +176,8 @@ export class AssistantMessageComponent extends Container { ); } 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; } @@ -161,7 +194,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 dbba2ff9..1176b993 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 161fe360..3e19a666 100644 --- a/apps/cli/test/assistant-message.test.ts +++ b/apps/cli/test/assistant-message.test.ts @@ -91,6 +91,48 @@ 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("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 683ebb25..cf3ddec7 100644 --- a/apps/cli/test/step-message.test.ts +++ b/apps/cli/test/step-message.test.ts @@ -176,6 +176,24 @@ describe("Step message presentation", () => { expect(finalAnswer).toBeGreaterThan(followUpThinking); }); + 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("updates output padding without dropping the first character", () => { initTheme("step-blue"); const assistant = new StepAssistantMessageComponent(