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(