/** * Contract: while a tool call's arguments stream (`partialJson` still open), * the pending tool preview is paced by ToolArgsRevealController — frames carry * growing prefixes of the raw stream re-parsed into display args — and once * the JSON closes the final parsed arguments render as-is (snap), mirroring * how assistant text snaps at message_end. */ import { afterEach, beforeAll, describe, expect, it, vi } from "bun:test"; import type { AssistantMessage } from "@oh-my-pi/pi-ai"; import { kStreamingPartialJson } from "@oh-my-pi/pi-ai/utils/block-symbols"; import { resetSettingsForTest, Settings, settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import { ToolExecutionComponent } from "@oh-my-pi/pi-coding-agent/modes/components/tool-execution"; import { EventController } from "@oh-my-pi/pi-coding-agent/modes/controllers/event-controller"; import { STREAMING_REVEAL_FRAME_MS } from "@oh-my-pi/pi-coding-agent/modes/controllers/streaming-reveal"; 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 type { AgentSessionEvent } from "@oh-my-pi/pi-coding-agent/session/agent-session"; beforeAll(async () => { await initTheme(); }); function makeStreamingMessage(content: AssistantMessage["content"]): AssistantMessage { return { role: "assistant", content, api: "anthropic-messages", provider: "anthropic", model: "claude-sonnet-4-5", stopReason: "stop", usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, timestamp: Date.now(), }; } function createFixture(streamingMessage: AssistantMessage) { const pendingTools = new Map(); const ctx = { isInitialized: true, init: vi.fn(async () => {}), ui: { requestRender: vi.fn(), requestComponentRender: vi.fn() }, settings, statusLine: { invalidate: vi.fn() }, updateEditorTopBorder: vi.fn(), streamingComponent: { updateContent: vi.fn(), markTranscriptBlockFinalized: vi.fn() }, streamingMessage, transcriptMessageComponents: new WeakMap(), pendingTools, noteDisplayableThinkingContent: vi.fn(() => false), chatContainer: { addChild: vi.fn() }, toolOutputExpanded: false, session: { getToolByName: () => undefined }, viewSession: { getToolByName: () => undefined }, sessionManager: { getCwd: () => process.cwd() }, } as unknown as InteractiveModeContext; return { controller: new EventController(ctx), pendingTools }; } async function dispatch(controller: EventController, message: AssistantMessage) { const event = { type: "message_update", message, assistantMessageEvent: undefined as never, } as Extract; await controller.handleEvent(event); } async function dispatchToolStart( controller: EventController, payload: { toolCallId: string; toolName: string; args: Record }, ) { await controller.handleEvent({ type: "tool_execution_start", toolCallId: payload.toolCallId, toolName: payload.toolName, args: payload.args, } as Extract); } describe("EventController paces streamed tool args", () => { afterEach(() => { vi.useRealTimers(); resetSettingsForTest(); vi.restoreAllMocks(); }); it("reveals the initial slice immediately, then paces growth across message_updates", async () => { await Settings.init({ inMemory: true, cwd: process.cwd() }); vi.useFakeTimers(); const updateArgsSpy = vi.spyOn(ToolExecutionComponent.prototype, "updateArgs"); const content = "x".repeat(400); const target = `{"path":"/tmp/a.ts","content":"${content}"}`; // Seed includes the complete `path` field (closing quote at byte 20) plus // the opening of `content`, so the rendered preview must show the real // path on the very first dispatch. const seed = target.slice(0, 35); // First message_update: only a small slice has arrived. The reveal // MUST surface it as-is (no empty initial frame). const seedStreaming = makeStreamingMessage([ { type: "toolCall", id: "tc-1", name: "write", arguments: {}, [kStreamingPartialJson]: seed }, ]); const { controller, pendingTools } = createFixture(seedStreaming); await dispatch(controller, seedStreaming); expect(pendingTools.size).toBe(1); // Component constructor consumes the initial render args directly; no // updateArgs has been invoked yet, but the seeded prefix is already on // the pending preview. const componentRender = pendingTools.get("tc-1")!.render(80).join("\n"); expect(Bun.stripANSI(componentRender)).toContain("/tmp/a.ts"); // Second message_update: the rest of the payload arrives. The controller // paces the new backlog through reveal ticks. const fullStreaming = makeStreamingMessage([ { type: "toolCall", id: "tc-1", name: "write", arguments: {}, [kStreamingPartialJson]: target }, ]); await dispatch(controller, fullStreaming); for (let i = 0; i < 3; i++) { vi.advanceTimersByTime(STREAMING_REVEAL_FRAME_MS); } const pacedFrames = updateArgsSpy.mock.calls.map(call => call[0] as Record); expect(pacedFrames.length).toBeGreaterThan(0); let previousLength = seed.length; for (const frame of pacedFrames) { const prefix = frame.__partialJson; if (typeof prefix !== "string") throw new Error("Expected __partialJson string on paced frame"); expect(target.startsWith(prefix)).toBe(true); expect(prefix.length).toBeLessThan(target.length); expect(prefix.length).toBeGreaterThanOrEqual(previousLength); previousLength = prefix.length; } // The JSON closed: providers drop `partialJson` and deliver final args. const finalArgs = { path: "/tmp/a.ts", content }; await dispatch( controller, makeStreamingMessage([{ type: "toolCall", id: "tc-1", name: "write", arguments: finalArgs }]), ); expect(updateArgsSpy.mock.calls.at(-1)?.[0]).toBe(finalArgs); // The reveal entry is gone: no further paced frames tick in. const calls = updateArgsSpy.mock.calls.length; vi.advanceTimersByTime(STREAMING_REVEAL_FRAME_MS * 5); expect(updateArgsSpy.mock.calls.length).toBe(calls); }); it("streams the full target through unpaced when smoothing is disabled", async () => { await Settings.init({ inMemory: true, cwd: process.cwd() }); settings.set("display.smoothStreaming", false); vi.useFakeTimers(); const updateArgsSpy = vi.spyOn(ToolExecutionComponent.prototype, "updateArgs"); const target = `{"path":"/tmp/a.ts","content":"abc"}`; const streaming = makeStreamingMessage([ { type: "toolCall", id: "tc-1", name: "write", arguments: { path: "/tmp/a.ts" }, [kStreamingPartialJson]: target, }, ]); const { controller } = createFixture(streaming); await dispatch(controller, streaming); await dispatch(controller, streaming); const frame = updateArgsSpy.mock.calls.at(-1)?.[0] as Record; expect(frame.__partialJson).toBe(target); const calls = updateArgsSpy.mock.calls.length; vi.advanceTimersByTime(STREAMING_REVEAL_FRAME_MS * 5); expect(updateArgsSpy.mock.calls.length).toBe(calls); }); it("reconciles validated full args on tool_execution_start when the closing args update never lands", async () => { await Settings.init({ inMemory: true, cwd: process.cwd() }); vi.useFakeTimers(); const content = "y".repeat(50); const target = `{"path":"/tmp/exec.ts","content":"${content}"}`; const streaming = makeStreamingMessage([ { type: "toolCall", id: "tc-1", name: "write", arguments: {}, [kStreamingPartialJson]: target }, ]); const { controller, pendingTools } = createFixture(streaming); // Args still streaming, but the reveal now seeds the preview with the // full available partialJson on the very first message_update — so the // path is already visible before the tool starts executing. await dispatch(controller, streaming); const component = pendingTools.get("tc-1"); if (!component) throw new Error("expected a pending write component"); expect(Bun.stripANSI(component.render(80).join("\n"))).toContain("/tmp/exec.ts"); // The closing full-args message_update never arrives (throttled `arguments` // with smoothing off, an owned-dialect projector, or a superseded turn that // still runs the call). The tool executes anyway: tool_execution_start is the // one event every path emits with the validated args, so it must reconcile. await dispatchToolStart(controller, { toolCallId: "tc-1", toolName: "write", args: { path: "/tmp/exec.ts", content }, }); expect(Bun.stripANSI(component.render(80).join("\n"))).toContain("/tmp/exec.ts"); // The reveal entry was cancelled: a late tick cannot re-truncate the body // back to a streaming prefix. vi.advanceTimersByTime(STREAMING_REVEAL_FRAME_MS * 5); expect(Bun.stripANSI(component.render(80).join("\n"))).toContain("/tmp/exec.ts"); }); });