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.
15 kB · 381 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382import type { AssistantMessage, Usage } from "@earendil-works/pi-ai";import { Container, type MarkdownTheme, type TUI } from "@earendil-works/pi-tui";import stripAnsi from "strip-ansi";import { beforeAll, describe, expect, test, vi } from "vitest";import type { AgentConnectionSessionEvent } from "../src/modes/agent-connection/index.js";import { AgentActivityTracker } from "../src/modes/interactive/agent-activity.js";import type { AssistantMessageComponent } from "../src/modes/interactive/components/assistant-message.js";import type { FileChangeSummary } from "../src/modes/interactive/components/edit-summary.js";import { createMermaidMarkdownTransform } from "../src/modes/interactive/components/mermaid.js";import type { ToolExecutionComponent } from "../src/modes/interactive/components/tool-execution.js";import { InteractiveMode } from "../src/modes/interactive/interactive-mode.js";import { getMarkdownTheme, initTheme } from "../src/modes/interactive/theme/theme.js";
const EMPTY_USAGE: Usage = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0, },};
type HandleEventThis = { isInitialized: boolean; settingsManager: { getShowTerminalProgress(): boolean }; connectionState: { isStreaming: boolean }; toolOutputExpanded: boolean; footer: { invalidate(): void }; ui: TUI; chatContainer: Container; recapContainer: Container; sessionRecap: string | undefined; hideThinkingBlock: boolean; hiddenThinkingLabel: string; streamingComponent: AssistantMessageComponent | undefined; streamingMessage: AssistantMessage | undefined; pendingTools: Map<string, ToolExecutionComponent>; agentRunFileChanges: Map<string, FileChangeSummary>; updateConnectionStateFromEvent(event: AgentConnectionSessionEvent): void; getMarkdownThemeWithSettings(): MarkdownTheme; getOrCreatePendingToolComponent(): Promise<ToolExecutionComponent | undefined>; getRetryAttempt(): number; getCurrentCwd(): string; stopWorkingLoader(): void; resetPendingToolState(): void; checkShutdownRequested(): Promise<void>; applyOptimisticContextUsage(): void; refreshConnectionContextUsage(): Promise<void>; clearShortcutGuide(): void; addMessageToChat(): void;};
type HandleEvent = (this: HandleEventThis, event: AgentConnectionSessionEvent) => Promise<void>;type GetUserInput = (this: { agentsViewRequest?: "agents_view" | "scoped_agents_view"; onInputCallback?: (text: string | undefined) => void;}) => Promise<string | undefined>;type HandleSubagentSummaryChatAction = ( this: { keybindings: { matches(data: string, action: string): boolean }; editor: { handleInput(data: string): void }; focusEditor(): void; toggleToolOutputExpansion(): void; toggleThinkingBlockVisibility(): void; }, data: string,) => void;
function createFakeInteractiveModeThis(): HandleEventThis { const fakeThis = { isInitialized: true, settingsManager: { getShowTerminalProgress: () => false }, connectionState: { isStreaming: false }, toolOutputExpanded: false, footer: { invalidate: vi.fn() }, activityTracker: new AgentActivityTracker(), ui: { requestRender: vi.fn() } as unknown as TUI, chatContainer: new Container(), recapContainer: new Container(), sessionRecap: "Updated files", hideThinkingBlock: false, hiddenThinkingLabel: "Thinking...", streamingComponent: undefined, streamingMessage: undefined, pendingMessagesContainer: new Container(), pendingBashComponents: [], pendingTools: new Map<string, ToolExecutionComponent>(), agentRunFileChanges: new Map<string, FileChangeSummary>(), updateConnectionStateFromEvent: vi.fn(), getMarkdownThemeWithSettings: () => getMarkdownTheme(), getOrCreatePendingToolComponent: vi.fn(async () => undefined), getRetryAttempt: () => 0, getCurrentCwd: () => "/tmp", stopWorkingLoader: vi.fn(), resetPendingToolState: vi.fn(), checkShutdownRequested: vi.fn(async () => {}), applyOptimisticContextUsage: vi.fn(), refreshConnectionContextUsage: vi.fn(async () => {}), clearShortcutGuide: vi.fn(), addMessageToChat: vi.fn(), }; Object.setPrototypeOf(fakeThis, InteractiveMode.prototype); return fakeThis;}
function createAssistantMessage(text: string): AssistantMessage { return { role: "assistant", content: [{ type: "text", text }], api: "test-api", provider: "test-provider", model: "test-model", usage: EMPTY_USAGE, stopReason: "stop", timestamp: Date.now(), };}
function renderChat(container: Container): string { return stripAnsi(container.render(120).join("\n"));}
describe("InteractiveMode streaming events", () => { beforeAll(() => { initTheme("dark"); });
test("renders assistant updates when attaching after message_start", async () => { const fakeThis = createFakeInteractiveModeThis(); const handleEvent = (InteractiveMode.prototype as unknown as { handleEvent: HandleEvent }).handleEvent;
await handleEvent.call(fakeThis, { type: "message_update", message: createAssistantMessage("partial response"), assistantMessageEvent: { type: "text_delta", contentIndex: 0, delta: "partial response", partial: createAssistantMessage("partial response"), }, });
expect(renderChat(fakeThis.chatContainer)).toContain("partial response");
await handleEvent.call(fakeThis, { type: "message_end", message: createAssistantMessage("final response"), });
expect(renderChat(fakeThis.chatContainer)).toContain("final response"); expect(fakeThis.streamingComponent).toBeUndefined(); expect(fakeThis.streamingMessage).toBeUndefined(); });
test("renders assistant end events when attaching after all updates", async () => { const fakeThis = createFakeInteractiveModeThis(); const handleEvent = (InteractiveMode.prototype as unknown as { handleEvent: HandleEvent }).handleEvent;
await handleEvent.call(fakeThis, { type: "message_end", message: createAssistantMessage("final response"), });
expect(renderChat(fakeThis.chatContainer)).toContain("final response"); expect(fakeThis.streamingComponent).toBeUndefined(); expect(fakeThis.streamingMessage).toBeUndefined(); });
test("does not block later compaction events on the agent-end stats refresh", async () => { const fakeThis = createFakeInteractiveModeThis(); let resolveRefresh!: () => void; fakeThis.refreshConnectionContextUsage = vi.fn( () => new Promise<void>((resolve) => { resolveRefresh = resolve; }), ); const handleEvent = (InteractiveMode.prototype as unknown as { handleEvent: HandleEvent }).handleEvent;
await expect(handleEvent.call(fakeThis, { type: "agent_end", messages: [] })).resolves.toBeUndefined(); expect(fakeThis.refreshConnectionContextUsage).toHaveBeenCalledOnce(); resolveRefresh(); });
test("keeps attached partial assistant text when agent_end arrives without message_end", async () => { const fakeThis = createFakeInteractiveModeThis(); const handleEvent = (InteractiveMode.prototype as unknown as { handleEvent: HandleEvent }).handleEvent;
await handleEvent.call(fakeThis, { type: "message_update", message: createAssistantMessage("partial response"), assistantMessageEvent: { type: "text_delta", contentIndex: 0, delta: "partial response", partial: createAssistantMessage("partial response"), }, }); await handleEvent.call(fakeThis, { type: "agent_end", messages: [] });
expect(renderChat(fakeThis.chatContainer)).toContain("partial response"); expect(fakeThis.streamingComponent).toBeUndefined(); expect(fakeThis.streamingMessage).toBeUndefined(); });
test("defers mermaid rendering to message_end in final mode", async () => { const fakeThis = createFakeInteractiveModeThis() as HandleEventThis & { mermaidMarkdownTransform?: ReturnType<typeof createMermaidMarkdownTransform>; }; fakeThis.mermaidMarkdownTransform = createMermaidMarkdownTransform({ getMode: () => "final" }); const handleEvent = (InteractiveMode.prototype as unknown as { handleEvent: HandleEvent }).handleEvent; const mermaidText = "```mermaid\nflowchart LR\n A[Start] --> B[Done]\n```";
await handleEvent.call(fakeThis, { type: "message_update", message: createAssistantMessage(mermaidText), assistantMessageEvent: { type: "text_delta", contentIndex: 0, delta: mermaidText, partial: createAssistantMessage(mermaidText), }, });
expect(renderChat(fakeThis.chatContainer)).not.toContain("───▶");
await handleEvent.call(fakeThis, { type: "message_end", message: createAssistantMessage(mermaidText), });
expect(renderChat(fakeThis.chatContainer)).toContain("───▶"); });
test("renders one agent-run edit total only when files changed", async () => { const fakeThis = createFakeInteractiveModeThis(); const handleEvent = (InteractiveMode.prototype as unknown as { handleEvent: HandleEvent }).handleEvent; const message = createAssistantMessage(""); message.content = [{ type: "toolCall", id: "edit-1", name: "edit", arguments: { path: "a.ts" } }];
await handleEvent.call(fakeThis, { type: "turn_end", message, toolResults: [ { role: "toolResult", toolCallId: "edit-1", toolName: "edit", content: [], details: { diff: "-1 old\n+1 new" }, isError: false, timestamp: 0, }, ], }); await handleEvent.call(fakeThis, { type: "agent_end", messages: [] }); const recap = renderChat(fakeThis.recapContainer); expect(recap).toContain("Recap: Updated files"); expect(recap).toContain("1 file changed | +1 -1"); expect(recap.indexOf("1 file changed")).toBeLessThan(recap.indexOf("Recap:")); expect(renderChat(fakeThis.chatContainer)).not.toContain("file changed");
const unchanged = createFakeInteractiveModeThis(); await handleEvent.call(unchanged, { type: "agent_end", messages: [] }); expect(renderChat(unchanged.recapContainer)).not.toContain("file changed"); });
test("keeps edit totals across automatic retries", async () => { const fakeThis = createFakeInteractiveModeThis(); fakeThis.agentRunFileChanges.set("/tmp/a.ts", { path: "a.ts", added: 1, removed: 1 }); fakeThis.getRetryAttempt = () => 1; const handleEvent = (InteractiveMode.prototype as unknown as { handleEvent: HandleEvent }).handleEvent;
await handleEvent.call(fakeThis, { type: "agent_start" });
expect([...fakeThis.agentRunFileChanges.values()]).toEqual([{ path: "a.ts", added: 1, removed: 1 }]); });
test("keeps edit totals when compaction restarts the agent", async () => { const fakeThis = createFakeInteractiveModeThis(); fakeThis.agentRunFileChanges.set("/tmp/a.ts", { path: "a.ts", added: 1, removed: 1 }); const handleEvent = (InteractiveMode.prototype as unknown as { handleEvent: HandleEvent }).handleEvent;
await handleEvent.call(fakeThis, { type: "agent_start" });
expect([...fakeThis.agentRunFileChanges.values()]).toEqual([{ path: "a.ts", added: 1, removed: 1 }]); });
test("clears edit totals when a new user prompt starts", async () => { const fakeThis = createFakeInteractiveModeThis(); fakeThis.agentRunFileChanges.set("/tmp/a.ts", { path: "a.ts", added: 1, removed: 1 }); const handleEvent = (InteractiveMode.prototype as unknown as { handleEvent: HandleEvent }).handleEvent; await handleEvent.call(fakeThis, { type: "agent_end", messages: [] }); expect(renderChat(fakeThis.recapContainer)).toContain("1 file changed");
await handleEvent.call(fakeThis, { type: "message_start", message: { role: "user", content: "next task", timestamp: Date.now() }, });
expect(fakeThis.agentRunFileChanges.size).toBe(0); expect(renderChat(fakeThis.recapContainer)).not.toContain("file changed"); expect(renderChat(fakeThis.recapContainer)).toContain("Recap: Updated files"); });
test("resolves input immediately after return to agents view was requested", async () => { const getUserInput = (InteractiveMode.prototype as unknown as { getUserInput: GetUserInput }).getUserInput;
await expect(getUserInput.call({ agentsViewRequest: "agents_view" })).resolves.toBeUndefined(); });
test("forwards typed keys from focused subagent summary back to the editor", () => { const handleSubagentSummaryChatAction = ( InteractiveMode.prototype as unknown as { handleSubagentSummaryChatAction: HandleSubagentSummaryChatAction } ).handleSubagentSummaryChatAction; const fakeThis = { keybindings: { matches: vi.fn(() => false) }, editor: { handleInput: vi.fn() }, focusEditor: vi.fn(), toggleToolOutputExpansion: vi.fn(), toggleThinkingBlockVisibility: vi.fn(), };
handleSubagentSummaryChatAction.call(fakeThis, "x");
expect(fakeThis.focusEditor).toHaveBeenCalledOnce(); expect(fakeThis.editor.handleInput).toHaveBeenCalledWith("x"); expect(fakeThis.toggleToolOutputExpansion).not.toHaveBeenCalled(); expect(fakeThis.toggleThinkingBlockVisibility).not.toHaveBeenCalled(); });
test("keeps focused subagent summary shortcuts in the chat surface", () => { const handleSubagentSummaryChatAction = ( InteractiveMode.prototype as unknown as { handleSubagentSummaryChatAction: HandleSubagentSummaryChatAction } ).handleSubagentSummaryChatAction; const fakeThis = { keybindings: { matches: vi.fn((_data: string, action: string) => action === "app.tools.expand") }, editor: { handleInput: vi.fn() }, focusEditor: vi.fn(), toggleToolOutputExpansion: vi.fn(), toggleThinkingBlockVisibility: vi.fn(), };
handleSubagentSummaryChatAction.call(fakeThis, "\x0f");
expect(fakeThis.toggleToolOutputExpansion).toHaveBeenCalledOnce(); expect(fakeThis.focusEditor).not.toHaveBeenCalled(); expect(fakeThis.editor.handleInput).not.toHaveBeenCalled(); });
test("does not pulse renders for background-only subagent work", () => { vi.useFakeTimers(); try { const requestRender = vi.fn(); const mode = Object.create(InteractiveMode.prototype) as InteractiveMode & Record<string, unknown>; Object.assign(mode, { connectionState: { isStreaming: false }, subagentSnapshots: new Map([["worker", { id: "worker", status: "running" }]]), pulseTimer: undefined, ui: { requestRender }, }); const updatePulse = Reflect.get(InteractiveMode.prototype, "updateWorkingPulse") as ( this: typeof mode, ) => void;
updatePulse.call(mode); vi.advanceTimersByTime(1000);
expect(requestRender).not.toHaveBeenCalled(); expect(Reflect.get(mode, "pulseTimer")).toBeUndefined(); } finally { vi.useRealTimers(); } });});