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(