Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
12 kB · 322 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323import type { AssistantMessage } from "@earendil-works/pi-ai";import { setKeybindings } from "@earendil-works/pi-tui";import stripAnsi from "strip-ansi";import { describe, expect, test } from "vitest";import { KeybindingsManager } from "../src/core/keybindings.js";import { AssistantMessageComponent, thinkingRecap } from "../src/modes/interactive/components/assistant-message.js";import { initTheme, theme } from "../src/modes/interactive/theme/theme.js";
const OSC133_ZONE_START = "\x1b]133;A\x07";const OSC133_ZONE_END = "\x1b]133;B\x07";const OSC133_ZONE_FINAL = "\x1b]133;C\x07";
function createAssistantMessage(content: AssistantMessage["content"]): AssistantMessage { return { role: "assistant", content, api: "openai-responses", provider: "openai", model: "gpt-4o-mini", usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "stop", timestamp: Date.now(), };}
describe("AssistantMessageComponent", () => { test("adds OSC 133 zone markers to assistant messages without tool calls", () => { initTheme("dark");
const component = new AssistantMessageComponent(createAssistantMessage([{ type: "text", text: "hello" }])); const lines = component.render(40);
expect(lines).not.toHaveLength(0); expect(lines[0]).toContain(OSC133_ZONE_START); expect(lines[lines.length - 1].startsWith(OSC133_ZONE_END + OSC133_ZONE_FINAL)).toBe(true); });
test("ignores null content blocks from malformed provider responses", () => { initTheme("dark");
const malformedContent = [null, { type: "text", text: "hello" }] as unknown as AssistantMessage["content"]; const component = new AssistantMessageComponent(createAssistantMessage(malformedContent)); expect(stripAnsi(component.render(40).join("\n"))).toContain("hello");
const updatedContent = [null, { type: "text", text: "hello again" }] as unknown as AssistantMessage["content"]; component.updateContent(createAssistantMessage(updatedContent)); expect(stripAnsi(component.render(40).join("\n"))).toContain("hello again"); });
test("does not add OSC 133 zone markers when assistant message contains tool calls", () => { initTheme("dark");
const component = new AssistantMessageComponent( createAssistantMessage([ { type: "text", text: "calling tool" }, { type: "toolCall", id: "tool-1", name: "ipython", arguments: { code: "open('file.txt').read()" } }, ]), ); const rendered = component.render(60).join("\n");
expect(rendered.includes(OSC133_ZONE_START)).toBe(false); expect(rendered.includes(OSC133_ZONE_END)).toBe(false); expect(rendered.includes(OSC133_ZONE_FINAL)).toBe(false); });
test("renders an abort status for messages with tool calls", () => { initTheme("dark");
const message = { ...createAssistantMessage([ { type: "toolCall" as const, id: "tool-1", name: "ipython", arguments: { code: "while True: pass" } }, ]), stopReason: "aborted" as const, errorMessage: "Operation aborted", }; const rendered = stripAnsi(new AssistantMessageComponent(message).render(80).join("\n"));
expect(rendered).toContain("Operation aborted"); });
test("honors initial expansion for multiline assistant errors", () => { initTheme("dark");
const message = { ...createAssistantMessage([]), stopReason: "error" as const, errorMessage: [ "Provider request failed", "Traceback (most recent call last):", ' File "/tmp/internal.py", line 12, in run', "RuntimeError: backend crashed", ].join("\n"), }; const component = new AssistantMessageComponent(message, false, undefined, "Thinking...", { expanded: true }); const rendered = stripAnsi(component.render(100).join("\n"));
expect(rendered).toContain("/tmp/internal.py"); expect(rendered).not.toContain("Ctrl+O to expand"); });
test("renders auth recovery guidance inline for simple provider errors", () => { initTheme("dark");
const message = { ...createAssistantMessage([]), stopReason: "error" as const, errorMessage: "401 status code (no body)\n\nRun /login to update credentials.", }; const component = new AssistantMessageComponent(message); const raw = component.render(120).join("\n"); const rendered = stripAnsi(raw);
expect(rendered).toContain("Error: 401 status code (no body) · Run /login to update credentials."); expect(rendered).not.toContain("Ctrl+O to expand"); expect(raw).toContain(theme.getFgAnsi("error")); });
test("renders collapsed multiline assistant errors as errors", () => { initTheme("dark");
const message = { ...createAssistantMessage([]), stopReason: "error" as const, errorMessage: [ "Provider request failed", "Traceback (most recent call last):", ' File "/tmp/internal.py", line 12, in run', "RuntimeError: backend crashed", ].join("\n"), }; const component = new AssistantMessageComponent(message); const raw = component.render(100).join("\n"); const rendered = stripAnsi(raw);
expect(rendered).toContain("Error: Provider request failed"); expect(rendered).toContain("to expand"); expect(raw).toContain(theme.getFgAnsi("error")); });});
describe("AssistantMessageComponent streaming identity", () => { // updateContent applied incrementally (as during daemon streaming) must render // byte-identical to a fresh component built from the final message. Guards the // reconcile-instead-of-rebuild optimization. function expectIdentity(component: AssistantMessageComponent, message: AssistantMessage, width = 90) { const incremental = component.render(width); const fresh = new AssistantMessageComponent(message).render(width); expect(incremental).toEqual(fresh); }
test("growing text block renders identically", () => { initTheme("dark"); const corpus = "## Heading\n\nSome paragraph text that wraps.\n\n- one\n- two\n\n```js\nconst a = 1;\n```\n\nEnd."; const component = new AssistantMessageComponent(); let text = ""; for (let offset = 0; offset < corpus.length; offset += 5) { text += corpus.slice(offset, offset + 5); const message = createAssistantMessage([{ type: "text", text }]); component.updateContent(message); expectIdentity(component, message); } });
test("thinking then text then tool call renders identically", () => { initTheme("dark"); const component = new AssistantMessageComponent(); const steps: AssistantMessage["content"][] = [ [{ type: "thinking", thinking: "Let me think" }], [{ type: "thinking", thinking: "Let me think about this more carefully." }], [ { type: "thinking", thinking: "Let me think about this more carefully." }, { type: "text", text: "Here is" }, ], [ { type: "thinking", thinking: "Let me think about this more carefully." }, { type: "text", text: "Here is the answer with **bold** text." }, ], [ { type: "thinking", thinking: "Let me think about this more carefully." }, { type: "text", text: "Here is the answer with **bold** text." }, { type: "toolCall", id: "t1", name: "bash", arguments: { command: "ls" } }, ], ]; for (const content of steps) { const message = createAssistantMessage(content); component.updateContent(message); expectIdentity(component, message); } });
test("aborted and error stop reasons render identically after streaming", () => { initTheme("dark"); for (const final of [ { stopReason: "aborted" as const, errorMessage: undefined }, { stopReason: "error" as const, errorMessage: "Provider exploded" }, ]) { const component = new AssistantMessageComponent(); component.updateContent(createAssistantMessage([{ type: "text", text: "Partial out" }])); component.render(90); const message = { ...createAssistantMessage([{ type: "text", text: "Partial output" }]), ...final, }; component.updateContent(message); const incremental = component.render(90); const fresh = new AssistantMessageComponent(message).render(90); expect(incremental).toEqual(fresh); } });
test("collapsed thinking shows a bold label, recap, and bracketed hint", () => { initTheme("dark"); setKeybindings(new KeybindingsManager());
const thinking = [ "**Weighing options**", "", "Some detail about the options.", "", "**Deciding the approach**", "", "More detail.", ].join("\n"); const message = createAssistantMessage([ { type: "thinking", thinking }, { type: "text", text: "Answer." }, ]); const rendered = stripAnsi(new AssistantMessageComponent(message, true).render(120).join("\n"));
expect(rendered).toContain("Thinking... · Deciding the approach (Ctrl+T to expand)"); expect(rendered).not.toContain("Some detail");
const expanded = stripAnsi(new AssistantMessageComponent(message, false).render(120).join("\n")); expect(expanded).toContain("Thinking... (Ctrl+T to collapse)"); expect(expanded).toContain("Some detail about the options.");
// A whitespace-only trace falls back to the label instead of an empty recap. expect(thinkingRecap(" \n\t\n", "Thinking...")).toBe("Thinking..."); });
test("recap text with delimiters cannot mask structural changes", () => { initTheme("dark"); setKeybindings(new KeybindingsManager());
// Unescaped, the first recap "X|1:text:1" makes this signature identical // to the next structure's (recap "X" plus a real text block), so the // rebuild that renders the new text block would be skipped. const component = new AssistantMessageComponent(undefined, true); component.updateContent(createAssistantMessage([{ type: "thinking", thinking: "X|1:text:1" }])); component.render(120);
component.updateContent( createAssistantMessage([ { type: "thinking", thinking: "X" }, { type: "text", text: "Visible answer." }, ]), ); const rendered = stripAnsi(component.render(120).join("\n"));
expect(rendered).toContain("Visible answer."); });
test("collapsed thinking row truncates instead of wrapping on narrow widths", () => { initTheme("dark"); setKeybindings(new KeybindingsManager());
const thinking = `**${"A deliberately verbose reasoning summary header that keeps going ".repeat(3).trim()}**`; const message = createAssistantMessage([{ type: "thinking", thinking }]); const lines = new AssistantMessageComponent(message, true) .render(60) .map((line) => stripAnsi(line)) .filter((line) => line.trim().length > 0);
expect(lines).toHaveLength(1); expect(lines[0]).toContain("Thinking..."); expect(lines[0]).toContain("to expand"); });
test("setHideThinkingBlock and setExpanded mid-stream render identically", () => { initTheme("dark"); const component = new AssistantMessageComponent(); const content: AssistantMessage["content"] = [ { type: "thinking", thinking: "Deep thoughts here." }, { type: "text", text: "Visible answer." }, ]; const message = createAssistantMessage(content); component.updateContent(message); component.render(90);
component.setHideThinkingBlock(true); const hidden = component.render(90); const freshHidden = new AssistantMessageComponent(message, true).render(90); expect(hidden).toEqual(freshHidden);
component.setHideThinkingBlock(false); const shown = component.render(90); const freshShown = new AssistantMessageComponent(message, false).render(90); expect(shown).toEqual(freshShown); });
test("width change mid-stream renders identically", () => { initTheme("dark"); const component = new AssistantMessageComponent(); const corpus = "A paragraph long enough to wrap at narrow widths with **style** and `code`."; let text = ""; const widths = [40, 90, 60]; for (let offset = 0, i = 0; offset < corpus.length; offset += 8, i++) { text += corpus.slice(offset, offset + 8); const message = createAssistantMessage([{ type: "text", text }]); component.updateContent(message); expectIdentity(component, message, widths[i % widths.length]); } });});