391 lines
14 KiB
TypeScript
391 lines
14 KiB
TypeScript
/**
|
|
* Contract: renderInitialMessages renders the collapsed live DISPLAY TRANSCRIPT,
|
|
* not the LLM context. The transcript comes from
|
|
* `session.buildTranscriptSessionContext({ collapseCompactedHistory: true })`;
|
|
* `sessionManager.buildSessionContext()` — the LLM-context builder — must not be
|
|
* consulted for display.
|
|
*
|
|
* Also guards the cold-launch terminal cleanup: `omp` / `omp -c` leave the
|
|
* previous run's transcript in native scrollback because the TUI's initial
|
|
* paint preserves it, so the cold-launch render must request a
|
|
* scrollback-clearing repaint (`clearTerminalHistory`).
|
|
*/
|
|
|
|
import { afterEach, beforeAll, describe, expect, it, type Mock, vi } from "bun:test";
|
|
import type { AgentMessage } from "@oh-my-pi/pi-agent-core";
|
|
import type { AssistantMessage, ImageContent, Usage } from "@oh-my-pi/pi-ai";
|
|
import { kStreamingPartialJson } from "@oh-my-pi/pi-ai/utils/block-symbols";
|
|
import { resetSettingsForTest, Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
|
|
import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme";
|
|
import type { InteractiveModeContext } from "@oh-my-pi/pi-coding-agent/modes/types";
|
|
import { UiHelpers } from "@oh-my-pi/pi-coding-agent/modes/utils/ui-helpers";
|
|
import type { SessionContext } from "@oh-my-pi/pi-coding-agent/session/session-context";
|
|
import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
|
|
import { type Component, Container, Image, ImageProtocol, setTerminalImageProtocol, TERMINAL } from "@oh-my-pi/pi-tui";
|
|
import { TempDir } from "@oh-my-pi/pi-utils";
|
|
|
|
beforeAll(() => {
|
|
initTheme();
|
|
});
|
|
|
|
const originalImageProtocol = TERMINAL.imageProtocol;
|
|
|
|
afterEach(() => {
|
|
resetSettingsForTest();
|
|
setTerminalImageProtocol(originalImageProtocol);
|
|
vi.restoreAllMocks();
|
|
});
|
|
|
|
function makeEmptyContext(): SessionContext {
|
|
return {
|
|
messages: [],
|
|
thinkingLevel: "off",
|
|
serviceTier: undefined,
|
|
models: {},
|
|
injectedTtsrRules: [],
|
|
selectedMCPToolNames: [],
|
|
hasPersistedMCPToolSelection: false,
|
|
mode: "none",
|
|
};
|
|
}
|
|
|
|
/** Build a minimal InteractiveModeContext mock, returning spies for assertions. */
|
|
function makeCtx(): {
|
|
ctx: InteractiveModeContext;
|
|
transcriptSpy: Mock<(options?: { collapseCompactedHistory?: boolean }) => SessionContext>;
|
|
llmContextSpy: Mock<() => SessionContext>;
|
|
renderSessionContextSpy: Mock<(...args: unknown[]) => void>;
|
|
} {
|
|
const transcriptSpy = vi.fn(() => makeEmptyContext());
|
|
const llmContextSpy = vi.fn(() => makeEmptyContext());
|
|
const renderSessionContextSpy = vi.fn();
|
|
|
|
const ctx = {
|
|
chatContainer: { clear: vi.fn(), addChild: vi.fn() },
|
|
pendingMessagesContainer: { clear: vi.fn(), disposeChildren: vi.fn() },
|
|
pendingBashComponents: [],
|
|
pendingPythonComponents: [],
|
|
session: { buildTranscriptSessionContext: transcriptSpy },
|
|
viewSession: {
|
|
buildTranscriptSessionContext: transcriptSpy,
|
|
sessionManager: {
|
|
buildSessionContext: llmContextSpy,
|
|
getEntries: vi.fn(() => []),
|
|
getCwd: vi.fn(() => "/tmp"),
|
|
},
|
|
},
|
|
sessionManager: {
|
|
buildSessionContext: llmContextSpy,
|
|
getEntries: vi.fn(() => []),
|
|
getCwd: vi.fn(() => "/tmp"),
|
|
},
|
|
renderSessionContext: renderSessionContextSpy,
|
|
showStatus: vi.fn(),
|
|
ui: { requestRender: vi.fn() },
|
|
resetTranscript: () => ctx.chatContainer.clear(),
|
|
} as unknown as InteractiveModeContext;
|
|
|
|
return { ctx, transcriptSpy, llmContextSpy, renderSessionContextSpy };
|
|
}
|
|
|
|
const emptyUsage: Usage = {
|
|
input: 0,
|
|
output: 0,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
totalTokens: 0,
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
|
};
|
|
|
|
const pngImage: ImageContent = {
|
|
type: "image",
|
|
data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg==",
|
|
mimeType: "image/png",
|
|
};
|
|
|
|
function assistantToolCall(id: string, name: string, args: Record<string, unknown>): AssistantMessage {
|
|
return {
|
|
role: "assistant",
|
|
content: [{ type: "toolCall", id, name, arguments: args }],
|
|
api: "anthropic-messages",
|
|
provider: "anthropic",
|
|
model: "claude-sonnet",
|
|
usage: emptyUsage,
|
|
stopReason: "toolUse",
|
|
timestamp: 1,
|
|
};
|
|
}
|
|
|
|
function transcriptWith(messages: AgentMessage[]): SessionContext {
|
|
return { ...makeEmptyContext(), messages };
|
|
}
|
|
|
|
function countImageComponents(component: Component): number {
|
|
const own = component instanceof Image ? 1 : 0;
|
|
const children = (component as { children?: unknown }).children;
|
|
if (!Array.isArray(children)) return own;
|
|
return own + children.reduce((count, child) => count + countImageComponents(child as Component), 0);
|
|
}
|
|
|
|
function hasImageComponent(component: Component): boolean {
|
|
return countImageComponents(component) > 0;
|
|
}
|
|
|
|
function makeRenderCtx(transcript: SessionContext): { ctx: InteractiveModeContext; chatContainer: Container } {
|
|
const chatContainer = new Container();
|
|
let helpers: UiHelpers;
|
|
const ctx = {
|
|
chatContainer,
|
|
pendingMessagesContainer: new Container(),
|
|
pendingBashComponents: [],
|
|
pendingPythonComponents: [],
|
|
pendingTools: new Map(),
|
|
statusLine: { invalidate: vi.fn() },
|
|
updateEditorBorderColor: vi.fn(),
|
|
updateEditorTopBorder: vi.fn(),
|
|
ui: { requestRender: vi.fn(), imageBudget: undefined },
|
|
resetTranscript: () => chatContainer.clear(),
|
|
settings: { get: () => false },
|
|
toolOutputExpanded: false,
|
|
hideThinkingBlock: false,
|
|
focusedAgentId: undefined,
|
|
editor: { addToHistory: vi.fn() },
|
|
viewSession: {
|
|
buildTranscriptSessionContext: () => transcript,
|
|
getToolByName: () => undefined,
|
|
extensionRunner: undefined,
|
|
sessionManager: {
|
|
getEntries: vi.fn(() => []),
|
|
getCwd: vi.fn(() => "/tmp"),
|
|
},
|
|
},
|
|
sessionManager: {
|
|
getEntries: vi.fn(() => []),
|
|
getCwd: vi.fn(() => "/tmp"),
|
|
putBlobSync: vi.fn(() => ({
|
|
hash: "hash",
|
|
path: "/tmp/hash",
|
|
displayPath: "/tmp/hash.png",
|
|
ref: "blob:sha256:hash",
|
|
})),
|
|
},
|
|
addMessageToChat: (message: AgentMessage, options?: { populateHistory?: boolean }) =>
|
|
helpers.addMessageToChat(message, options),
|
|
renderSessionContext: (
|
|
context: SessionContext,
|
|
options?: { updateFooter?: boolean; populateHistory?: boolean },
|
|
) => helpers.renderSessionContext(context, options),
|
|
showStatus: vi.fn(),
|
|
} as unknown as InteractiveModeContext;
|
|
helpers = new UiHelpers(ctx);
|
|
return { ctx, chatContainer };
|
|
}
|
|
|
|
describe("UiHelpers.renderInitialMessages — transcript source", () => {
|
|
it("renders the collapsed live display transcript, never the LLM context", () => {
|
|
const { ctx, transcriptSpy, llmContextSpy, renderSessionContextSpy } = makeCtx();
|
|
const transcript = makeEmptyContext();
|
|
transcriptSpy.mockReturnValue(transcript);
|
|
|
|
new UiHelpers(ctx).renderInitialMessages();
|
|
|
|
expect(transcriptSpy).toHaveBeenCalledWith({ collapseCompactedHistory: true });
|
|
expect(llmContextSpy).not.toHaveBeenCalled();
|
|
expect(renderSessionContextSpy).toHaveBeenCalledWith(transcript, {
|
|
updateFooter: true,
|
|
populateHistory: true,
|
|
});
|
|
});
|
|
});
|
|
|
|
describe("UiHelpers.renderInitialMessages — clearTerminalHistory", () => {
|
|
it("requests a scrollback-clearing repaint when clearTerminalHistory is set", () => {
|
|
const { ctx } = makeCtx();
|
|
new UiHelpers(ctx).renderInitialMessages({ clearTerminalHistory: true });
|
|
expect(ctx.ui.requestRender).toHaveBeenCalledWith(true, { clearScrollback: true });
|
|
});
|
|
|
|
it("never clears scrollback when clearTerminalHistory is unset", () => {
|
|
const { ctx } = makeCtx();
|
|
new UiHelpers(ctx).renderInitialMessages();
|
|
const clearedCall = (ctx.ui.requestRender as Mock<(...a: unknown[]) => void>).mock.calls.find(
|
|
([force, opts]) => force === true && (opts as { clearScrollback?: boolean } | undefined)?.clearScrollback,
|
|
);
|
|
expect(clearedCall).toBeUndefined();
|
|
});
|
|
});
|
|
|
|
describe("UiHelpers.renderInitialMessages — image replay", () => {
|
|
it("restores read tool image blocks onto the rebuilt assistant transcript", async () => {
|
|
await Settings.init({ inMemory: true, overrides: { "terminal.showImages": true } });
|
|
setTerminalImageProtocol(ImageProtocol.Sixel);
|
|
const transcript = transcriptWith([
|
|
assistantToolCall("read-image", "read", { path: "sample.png" }),
|
|
{
|
|
role: "toolResult",
|
|
toolCallId: "read-image",
|
|
toolName: "read",
|
|
content: [{ type: "text", text: "Read image: sample.png" }, pngImage],
|
|
isError: false,
|
|
timestamp: 2,
|
|
},
|
|
]);
|
|
const { ctx, chatContainer } = makeRenderCtx(transcript);
|
|
|
|
new UiHelpers(ctx).renderInitialMessages();
|
|
|
|
expect(hasImageComponent(chatContainer)).toBe(true);
|
|
expect(Bun.stripANSI(chatContainer.render(100).join("\n"))).toContain("Read sample.png");
|
|
});
|
|
|
|
it("restores eval display image blocks onto rebuilt tool output", async () => {
|
|
await Settings.init({ inMemory: true, overrides: { "terminal.showImages": true } });
|
|
setTerminalImageProtocol(ImageProtocol.Sixel);
|
|
const transcript = transcriptWith([
|
|
assistantToolCall("eval-image", "eval", { language: "py", code: "display(image)" }),
|
|
{
|
|
role: "toolResult",
|
|
toolCallId: "eval-image",
|
|
toolName: "eval",
|
|
content: [{ type: "text", text: "(displayed 1 image; no text output)" }, pngImage],
|
|
details: {
|
|
language: "python",
|
|
cells: [{ index: 0, code: "display(image)", output: "display image 1: 1x1", status: "complete" }],
|
|
},
|
|
isError: false,
|
|
timestamp: 2,
|
|
},
|
|
]);
|
|
|
|
const { ctx, chatContainer } = makeRenderCtx(transcript);
|
|
|
|
new UiHelpers(ctx).renderInitialMessages();
|
|
|
|
expect(hasImageComponent(chatContainer)).toBe(true);
|
|
expect(Bun.stripANSI(chatContainer.render(100).join("\n"))).toContain("display image 1: 1x1");
|
|
});
|
|
|
|
it("replays reopened session image blocks through the cold-start rebuild path", async () => {
|
|
await Settings.init({ inMemory: true, overrides: { "terminal.showImages": true } });
|
|
setTerminalImageProtocol(ImageProtocol.Sixel);
|
|
using tempDir = TempDir.createSync("@pi-render-initial-image-replay-");
|
|
const session = SessionManager.create(tempDir.path(), tempDir.path());
|
|
session.appendMessage(assistantToolCall("read-reopened", "read", { path: "reopened.png" }));
|
|
session.appendMessage({
|
|
role: "toolResult",
|
|
toolCallId: "read-reopened",
|
|
toolName: "read",
|
|
content: [{ type: "text", text: "Read image: reopened.png" }, pngImage],
|
|
isError: false,
|
|
timestamp: 2,
|
|
});
|
|
session.appendMessage(assistantToolCall("eval-reopened", "eval", { language: "py", code: "display(image)" }));
|
|
session.appendMessage({
|
|
role: "toolResult",
|
|
toolCallId: "eval-reopened",
|
|
toolName: "eval",
|
|
content: [{ type: "text", text: "(displayed 1 image; no text output)" }, pngImage],
|
|
details: {
|
|
language: "python",
|
|
cells: [{ index: 0, code: "display(image)", output: "display image 1: 1x1", status: "complete" }],
|
|
},
|
|
isError: false,
|
|
timestamp: 4,
|
|
});
|
|
await session.flush();
|
|
const sessionFile = session.getSessionFile();
|
|
if (!sessionFile) throw new Error("Expected persisted session file");
|
|
const reloaded = await SessionManager.open(sessionFile);
|
|
const transcript = reloaded.buildSessionContext({ transcript: true });
|
|
const { ctx, chatContainer } = makeRenderCtx(transcript);
|
|
|
|
new UiHelpers(ctx).renderInitialMessages({ clearTerminalHistory: true });
|
|
|
|
expect(countImageComponents(chatContainer)).toBe(2);
|
|
expect(Bun.stripANSI(chatContainer.render(100).join("\n"))).toContain("Read reopened.png");
|
|
expect(ctx.ui.requestRender).toHaveBeenCalledWith(true, { clearScrollback: true });
|
|
});
|
|
});
|
|
|
|
describe("UiHelpers.renderSessionContext — error-stop tool calls", () => {
|
|
it("keeps the synthetic assistant error result instead of replaying a later tool result", async () => {
|
|
await Settings.init({ inMemory: true });
|
|
const transcript = transcriptWith([
|
|
{
|
|
role: "assistant",
|
|
content: [
|
|
{
|
|
type: "toolCall",
|
|
id: "error-tool",
|
|
name: "eval",
|
|
arguments: { language: "py", code: "raise RuntimeError('boom')" },
|
|
},
|
|
],
|
|
api: "anthropic-messages",
|
|
provider: "anthropic",
|
|
model: "claude-sonnet",
|
|
usage: emptyUsage,
|
|
stopReason: "error",
|
|
errorMessage: "synthetic assistant stop error",
|
|
timestamp: 1,
|
|
},
|
|
{
|
|
role: "toolResult",
|
|
toolCallId: "error-tool",
|
|
toolName: "eval",
|
|
content: [{ type: "text", text: "late tool result must not replace the assistant stop error" }],
|
|
isError: false,
|
|
timestamp: 2,
|
|
},
|
|
]);
|
|
const { ctx, chatContainer } = makeRenderCtx(transcript);
|
|
|
|
new UiHelpers(ctx).renderInitialMessages();
|
|
|
|
const rendered = Bun.stripANSI(chatContainer.render(120).join("\n"));
|
|
expect(rendered).toContain("synthetic assistant stop error");
|
|
expect(rendered).not.toContain("late tool result must not replace the assistant stop error");
|
|
});
|
|
});
|
|
|
|
describe("UiHelpers.renderSessionContext — mid-stream tool call rebuild", () => {
|
|
it("decodes streamed write content from partialJson, not the provider's stale parsed arguments", async () => {
|
|
// A transcript rebuild (theme change, settings edit, focus replay) can land
|
|
// while a write's args still stream. The provider re-parses `arguments`
|
|
// only every STREAMING_JSON_PARSE_MIN_GROWTH bytes, so the parsed snapshot
|
|
// lags the raw buffer. The rebuilt preview must decode from the buffer —
|
|
// exactly like the live reveal path — or the write body freezes at the
|
|
// last throttled parse until more bytes arrive.
|
|
await Settings.init({ inMemory: true });
|
|
const staleContent = "line one of the streamed write";
|
|
const grownBuffer = `{"path":"/tmp/mid.ts","content":"${staleContent}\\nGROWN_TAIL_SENTINEL`;
|
|
const transcript = transcriptWith([
|
|
{
|
|
role: "assistant",
|
|
content: [
|
|
{
|
|
type: "toolCall",
|
|
id: "write-mid",
|
|
name: "write",
|
|
// Provider-parsed snapshot from BEFORE the buffer grew.
|
|
arguments: { path: "/tmp/mid.ts", content: staleContent },
|
|
[kStreamingPartialJson]: grownBuffer,
|
|
},
|
|
],
|
|
api: "anthropic-messages",
|
|
provider: "anthropic",
|
|
model: "claude-sonnet",
|
|
usage: emptyUsage,
|
|
stopReason: "toolUse",
|
|
timestamp: 1,
|
|
},
|
|
]);
|
|
const { ctx, chatContainer } = makeRenderCtx(transcript);
|
|
|
|
new UiHelpers(ctx).renderInitialMessages();
|
|
|
|
const rendered = Bun.stripANSI(chatContainer.render(120).join("\n"));
|
|
expect(rendered).toContain("GROWN_TAIL_SENTINEL");
|
|
});
|
|
});
|