Luigit
repositories / pi-ext

pi-ext

bugabingas pi extensions

owned by admin

extensions/ultra/__tests__/progress.test.ts

Raw
import type { AgentSessionEvent } from "@earendil-works/pi-coding-agent";
import { describe, expect, it, vi } from "vitest";
import {
	type AgentControls,
	initialReducerState,
	mapSessionEvent,
	reduceProgress,
	type UltraProgressEvent,
} from "../progress.ts";

// AgentSessionEvent is a large discriminated union whose message-bearing
// variants carry full AgentMessage / AssistantMessage objects. The mapping
// only ever reads the validated field paths, so the fixtures build minimal
// shapes and cast — exercising exactly the paths the real session emits.
const ev = (e: unknown): AgentSessionEvent => e as AgentSessionEvent;

const AGENT = "review#0";
const PHASE = "review";

function map(e: unknown): UltraProgressEvent | null {
	return mapSessionEvent(ev(e), AGENT, PHASE);
}

describe("mapSessionEvent", () => {
	it("maps a scripted session sequence to tagged ultra events", () => {
		// tool_execution_start -> message_update(text_delta) -> tool_execution_end -> agent_end
		const start = map({ type: "agent_start" });
		const toolStart = map({
			type: "tool_execution_start",
			toolCallId: "t1",
			toolName: "read",
			args: { file: "a.ts" },
		});
		const delta = map({
			type: "message_update",
			message: {},
			assistantMessageEvent: {
				type: "text_delta",
				contentIndex: 0,
				delta: "hello ",
				partial: {},
			},
		});
		const toolEnd = map({
			type: "tool_execution_end",
			toolCallId: "t1",
			toolName: "read",
			// content and details deliberately differ: the mapping must read
			// `result.details` (the validated structured-output path), never content.
			result: { content: "CONTENT_TEXT", details: { verdict: "DETAILS_VAL" } },
			isError: false,
		});
		const end = map({ type: "agent_end", messages: [], willRetry: false });

		// every mapped event is tagged with the agent id + phase it was given
		for (const e of [start, toolStart, delta, toolEnd, end]) {
			expect(e).not.toBeNull();
			expect(e?.agentId).toBe(AGENT);
			expect(e?.phase).toBe(PHASE);
		}

		expect(start).toMatchObject({ kind: "start" });
		expect(toolStart).toMatchObject({ kind: "action", toolName: "read" });
		expect(toolStart?.text).toContain("a.ts");
		expect(delta).toMatchObject({ kind: "delta", text: "hello " });
		expect(toolEnd).toMatchObject({ kind: "action", toolName: "read" });
		expect(end).toMatchObject({ kind: "end", status: "done" });

		// the discriminating assertion: details is read, content is not
		expect(toolEnd?.text).toContain("DETAILS_VAL");
		expect(toolEnd?.text).not.toContain("CONTENT_TEXT");
	});

	it("maps thinking deltas to delta events", () => {
		const e = map({
			type: "message_update",
			message: {},
			assistantMessageEvent: {
				type: "thinking_delta",
				contentIndex: 0,
				delta: "pondering",
				partial: {},
			},
		});
		expect(e).toMatchObject({ kind: "delta", text: "pondering" });
	});

	it("flags a tool error on the action event", () => {
		const e = map({
			type: "tool_execution_end",
			toolCallId: "t9",
			toolName: "bash",
			result: { content: "boom", details: undefined },
			isError: true,
		});
		expect(e).toMatchObject({
			kind: "action",
			toolName: "bash",
			status: "error",
		});
	});

	it("renders known and arbitrary JSON tool calls tersely", () => {
		let state = reduceProgress(
			initialReducerState(),
			map({
				type: "tool_execution_start",
				toolCallId: "read-1",
				toolName: "read",
				args: { path: "src/settings.ts", offset: 4, limit: 20 },
			}) as UltraProgressEvent,
		);
		expect(state.rows[0]?.action).toBe("Reading src/settings.ts");

		state = reduceProgress(
			state,
			map({
				type: "tool_execution_start",
				toolCallId: "custom-1",
				toolName: "custom_tool",
				args: {
					action: "sync",
					items: [1, 2, 3],
					options: { force: true },
				},
			}) as UltraProgressEvent,
		);
		expect(state.rows[0]?.action).toBe(
			"Using custom_tool · Action: sync · Items: 3 items · Options: 1 field",
		);
	});

	it("keeps complete tool results while rendering their JSON tersely", () => {
		const result = {
			content: [{ type: "text", text: "Verbose body that stays expandable" }],
			details: {
				ok: true,
				files: ["a.ts", "b.ts"],
				metadata: { source: "index" },
			},
		};
		let state = reduceProgress(
			initialReducerState(),
			map({
				type: "tool_execution_start",
				toolCallId: "custom-1",
				toolName: "custom_tool",
				args: { action: "index" },
			}) as UltraProgressEvent,
		);
		state = reduceProgress(
			state,
			map({
				type: "tool_execution_end",
				toolCallId: "custom-1",
				toolName: "custom_tool",
				result,
				isError: false,
			}) as UltraProgressEvent,
		);

		const item = state.transcriptItems[AGENT]?.[0];
		expect(item).toMatchObject({
			kind: "tool",
			result,
			fallback: "Result · Ok: true · Files: 2 items · Metadata: 1 field",
		});
	});

	it.each([
		[null, "Result · null"],
		[true, "Result · true"],
		[42, "Result · 42"],
		[[1, 2, 3], "Result · 3 items"],
		[
			{ content: [{ type: "text", text: "Indexed 3 files\nverbose" }] },
			"Result · Indexed 3 files",
		],
	])("renders JSON result %j correctly", (result, expected) => {
		const state = reduceProgress(
			initialReducerState(),
			map({
				type: "tool_execution_end",
				toolCallId: "custom-1",
				toolName: "custom_tool",
				result,
				isError: false,
			}) as UltraProgressEvent,
		);
		expect(state.transcriptItems[AGENT]?.[0]).toMatchObject({
			fallback: expected,
			result,
		});
	});

	it("summarizes JSON text results without changing their payload", () => {
		const result = {
			content: [{ type: "text", text: '{"status":"ready","files":["a.ts"]}' }],
		};
		const event = map({
			type: "tool_execution_end",
			toolCallId: "json",
			toolName: "custom",
			result,
			isError: false,
		});
		expect(event).toMatchObject({
			text: "Status: ready · Files: 1 item",
			result,
		});
	});

	it("maps tool execution updates as partial results", () => {
		const partialResult = {
			content: [{ type: "text", text: "first chunk" }],
		};
		expect(
			map({
				type: "tool_execution_update",
				toolCallId: "t1",
				toolName: "read",
				args: { path: "a.ts" },
				partialResult,
			}),
		).toMatchObject({
			kind: "action",
			toolCallId: "t1",
			toolName: "read",
			result: partialResult,
			isPartial: true,
		});
	});

	it("returns null for events that do not map", () => {
		expect(map({ type: "turn_start" })).toBeNull();
		expect(map({ type: "message_start", message: {} })).toBeNull();
		expect(
			map({ type: "queue_update", steering: [], followUp: [] }),
		).toBeNull();
		// a message_update whose inner event is not a text/thinking delta
		expect(
			map({
				type: "message_update",
				message: {},
				assistantMessageEvent: {
					type: "text_start",
					contentIndex: 0,
					partial: {},
				},
			}),
		).toBeNull();
	});

	it("never attaches controls (the runner owns the live session)", () => {
		const e = map({ type: "agent_start" });
		expect(e?.controls).toBeUndefined();
	});
});

// --- reducer -------------------------------------------------------------

const pe = (
	e: Partial<UltraProgressEvent> &
		Pick<UltraProgressEvent, "agentId" | "phase" | "kind">,
): UltraProgressEvent => e;

function fold(events: UltraProgressEvent[]) {
	return events.reduce(reduceProgress, initialReducerState());
}

describe("reduceProgress", () => {
	it("starts empty", () => {
		const s = initialReducerState();
		expect(s.rows).toEqual([]);
		expect(s.transcripts).toEqual({});
		expect(s.phases).toEqual([]);
		expect(s.phase).toBeUndefined();
		expect(s).toMatchObject({
			total: 0,
			queued: 0,
			running: 0,
			done: 0,
			dropped: 0,
			later: 0,
		});
	});

	it("seeds queued agents and pending phases from a workflow plan", () => {
		const s = reduceProgress(initialReducerState(), {
			kind: "plan",
			startedAt: 1000,
			phases: [
				{
					phase: "research",
					status: "resolved",
					agentIds: ["research#0"],
				},
				{ phase: "verify", status: "agents-pending", agentIds: [] },
				{
					phase: "publish",
					status: "condition-pending",
					agentIds: [],
				},
			],
		} as unknown as UltraProgressEvent);

		expect(s.phases.map((phase) => phase.phase)).toEqual([
			"research",
			"verify",
			"publish",
		]);
		expect(s.rows).toEqual([
			expect.objectContaining({
				agentId: "research#0",
				status: "queued",
			}),
		]);
		expect(s.phase).toBeUndefined();
		expect(s).toMatchObject({
			total: 1,
			queued: 1,
			running: 0,
			done: 0,
			dropped: 0,
			later: 2,
		});
	});

	it("stores the resolved summary and prompt without starting the queued agent", () => {
		let state = reduceProgress(initialReducerState(), {
			kind: "plan",
			startedAt: 1000,
			phases: [{ phase: "review", status: "resolved", agentIds: ["review#0"] }],
		} as UltraProgressEvent);
		state = reduceProgress(
			state,
			pe({
				agentId: "review#0",
				phase: "review",
				kind: "prompt",
				summary: "Review security",
				prompt: "Inspect every authentication path.",
			}),
		);

		expect(state.rows[0]).toMatchObject({
			status: "queued",
			summary: "Review security",
			prompt: "Inspect every authentication path.",
		});
		expect(state.transcriptItems["review#0"]).toEqual([
			{ kind: "user", text: "Inspect every authentication path." },
		]);
	});

	it("starts wall time with the plan and freezes it when every phase settles", () => {
		vi.useFakeTimers();
		try {
			vi.setSystemTime(1000);
			const planned = reduceProgress(initialReducerState(), {
				kind: "plan",
				startedAt: 1000,
				phases: [{ phase: "one", status: "resolved", agentIds: ["one#0"] }],
			} as unknown as UltraProgressEvent);
			expect(planned.startedAt).toBe(1000);
			expect(planned.endedAt).toBeUndefined();

			vi.setSystemTime(61_000);
			const done = reduceProgress(planned, {
				kind: "phase",
				phase: "one",
				status: "done",
				agentIds: ["one#0"],
			} as unknown as UltraProgressEvent);
			expect(done.endedAt).toBe(61_000);

			vi.setSystemTime(121_000);
			const repeated = reduceProgress(done, {
				kind: "phase",
				phase: "one",
				status: "done",
				agentIds: ["one#0"],
			} as unknown as UltraProgressEvent);
			expect(repeated.endedAt).toBe(61_000);
		} finally {
			vi.useRealTimers();
		}
	});

	it("resolves pending agents before start and transitions queued to running", () => {
		const planned = reduceProgress(initialReducerState(), {
			kind: "plan",
			startedAt: 1000,
			phases: [{ phase: "verify", status: "agents-pending", agentIds: [] }],
		} as unknown as UltraProgressEvent);
		const resolved = reduceProgress(planned, {
			kind: "phase",
			phase: "verify",
			status: "resolved",
			agentIds: ["verify#0", "verify#1"],
		} as unknown as UltraProgressEvent);

		expect(resolved.rows.map((row) => row.status)).toEqual([
			"queued",
			"queued",
		]);
		expect(resolved).toMatchObject({ queued: 2, running: 0, later: 0 });
		const running = reduceProgress(
			resolved,
			pe({ agentId: "verify#0", phase: "verify", kind: "start" }),
		);
		expect(running.rows.map((row) => row.status)).toEqual([
			"running",
			"queued",
		]);
		expect(running).toMatchObject({ queued: 1, running: 1, later: 0 });
	});

	it("removes provisional agent rows when a phase is skipped", () => {
		const planned = reduceProgress(initialReducerState(), {
			kind: "plan",
			startedAt: 1000,
			phases: [
				{ phase: "publish", status: "resolved", agentIds: ["publish#0"] },
			],
		} as unknown as UltraProgressEvent);
		const skipped = reduceProgress(planned, {
			kind: "phase",
			phase: "publish",
			status: "skipped",
			agentIds: [],
		} as unknown as UltraProgressEvent);

		expect(skipped.rows).toEqual([]);
		expect(skipped.phases[0]).toMatchObject({
			phase: "publish",
			status: "skipped",
		});
		expect(skipped).toMatchObject({ total: 0, queued: 0, later: 0 });
	});

	it("formats tool call rows from structured args and does not replace them with raw results", () => {
		const s = fold([
			pe({
				agentId: "a",
				phase: "p1",
				kind: "action",
				toolName: "bash",
				args: { command: "git status --short\nsecond line" },
				text: JSON.stringify({ command: "git status --short" }),
			}),
			pe({
				agentId: "a",
				phase: "p1",
				kind: "action",
				toolName: "bash",
				result: { content: [{ type: "text", text: "RAW RESULT" }] },
				text: "RAW RESULT",
			}),
		]);

		expect(s.rows[0].action).toBe("Running git status --short");
		expect(s.rows[0].action).not.toContain("RAW RESULT");
	});

	it("prefers an LLM action summary over raw tool args", () => {
		const s = fold([
			pe({
				agentId: "a",
				phase: "p1",
				kind: "action",
				toolName: "read",
				args: { path: "extensions/ultra/progress.ts" },
				actionSummary: "Inspecting workflow progress state",
			}),
		]);

		expect(s.rows[0].action).toBe("Inspecting workflow progress state");
		expect(s.rows[0].action).not.toContain("progress.ts");
	});

	it("keeps blocked outcomes visible instead of summarizing them as zero findings", () => {
		const s = fold([
			pe({
				agentId: "report#0",
				phase: "report",
				kind: "action",
				toolName: "structured_output",
				args: { status: "blocked", summary: "Missing evidence", findings: [] },
			}),
		]);
		expect(s.rows[0].action).toBe(
			"structured output · blocked · Missing evidence",
		);
	});

	it("summarizes structured_output rows instead of dumping JSON", () => {
		const s = fold([
			pe({
				agentId: "a",
				phase: "review",
				kind: "action",
				toolName: "structured_output",
				args: {
					findings: [
						{ title: "A very long finding title", file: "a.ts" },
						{ title: "Another finding", file: "b.ts" },
					],
				},
				text: JSON.stringify({
					findings: [{ title: "A very long finding title" }],
				}),
			}),
		]);

		expect(s.rows[0].action).toBe("structured output · 2 findings");
		expect(s.rows[0].action).not.toContain('{"findings"');
	});

	it("folds a scripted stream into one upserted row + transcript", () => {
		const s = fold([
			pe({ agentId: "a", phase: "p1", kind: "start" }),
			pe({
				agentId: "a",
				phase: "p1",
				kind: "action",
				toolName: "read",
				text: "a.ts",
			}),
			pe({ agentId: "a", phase: "p1", kind: "delta", text: "hello " }),
			pe({ agentId: "a", phase: "p1", kind: "delta", text: "world" }),
			pe({ agentId: "a", phase: "p1", kind: "end", status: "done" }),
		]);

		expect(s.rows).toHaveLength(1);
		expect(s.rows[0]).toMatchObject({
			agentId: "a",
			phase: "p1",
			status: "done",
		});
		expect(s.rows[0].action).toContain("read");
		// transcript accumulates the streamed deltas verbatim + the tool action
		expect(s.transcripts.a).toContain("hello world");
		expect(s.transcripts.a).toContain("read");
		expect(s).toMatchObject({ total: 1, done: 1, dropped: 0 });
		expect(s.phases.map((phase) => phase.phase)).toEqual(["p1"]);
		expect(s.phase).toBe("p1");
	});

	it("upserts idempotently — a later event updates the row, never appends", () => {
		const s = fold([
			pe({ agentId: "a", phase: "p1", kind: "start" }),
			pe({ agentId: "b", phase: "p1", kind: "start" }),
			pe({ agentId: "a", phase: "p1", kind: "end", status: "done" }),
		]);
		expect(s.rows.map((r) => r.agentId)).toEqual(["a", "b"]); // order preserved
		expect(s.rows[0]).toMatchObject({ agentId: "a", status: "done" }); // updated in place
		expect(s.rows[1]).toMatchObject({ agentId: "b", status: "running" });
		expect(s.total).toBe(2);
	});

	it("accumulates transcripts per agent and isolates agents", () => {
		const s = fold([
			pe({ agentId: "a", phase: "p1", kind: "delta", text: "secretA" }),
			pe({ agentId: "b", phase: "p1", kind: "delta", text: "secretB" }),
			pe({ agentId: "a", phase: "p1", kind: "delta", text: "-more" }),
		]);
		expect(s.transcripts.a).toBe("secretA-more");
		expect(s.transcripts.b).toBe("secretB");
		expect(s.transcripts.a).not.toContain("secretB");
		expect(s.transcripts.b).not.toContain("secretA");
	});

	it("updates done/dropped tallies on terminal events", () => {
		const s = fold([
			pe({ agentId: "a", phase: "p1", kind: "end", status: "done" }),
			pe({ agentId: "b", phase: "p1", kind: "end", status: "dropped" }),
			pe({ agentId: "c", phase: "p1", kind: "end", status: "failed" }),
		]);
		expect(s).toMatchObject({ total: 3, done: 1, dropped: 2 });
	});

	it("retains typed terminal failure diagnostics on dropped rows", () => {
		const failure = {
			code: "transport" as const,
			message: "rate limited",
			retryable: true,
			attempts: 2,
			model: "openai/gpt",
			transcriptTail: "bad response",
		};
		const s = fold([
			pe({
				agentId: "a",
				phase: "p1",
				kind: "end",
				status: "dropped",
				failure,
			}),
		]);
		expect(s.rows[0].failure).toEqual(failure);
	});

	it("tracks the phase counter as distinct phases arrive in order", () => {
		const s = fold([
			pe({ agentId: "a", phase: "review", kind: "start" }),
			pe({ agentId: "a", phase: "review", kind: "end", status: "done" }),
			pe({ agentId: "v0", phase: "verify", kind: "start" }),
		]);
		expect(s.phases.map((phase) => phase.phase)).toEqual(["review", "verify"]);
		expect(s.phase).toBe("verify");
	});

	it("retains model, timing, and usage metadata across later events", () => {
		const usage = {
			input: 100,
			output: 20,
			total: 120,
			cacheRead: 10,
			cacheWrite: 0,
			cost: 0.1,
		};
		const s = fold([
			pe({
				agentId: "a",
				phase: "p1",
				kind: "start",
				tier: "small",
				model: "openai/gpt-5-mini",
				thinkingLevel: "low",
				startedAt: 1000,
			}),
			pe({ agentId: "a", phase: "p1", kind: "delta", text: "x" }),
			pe({ agentId: "a", phase: "p1", kind: "end", status: "done", usage }),
		]);

		expect(s.rows[0]).toMatchObject({
			tier: "small",
			model: "openai/gpt-5-mini",
			thinkingLevel: "low",
			startedAt: 1000,
			usage,
		});
	});

	it("binds per-agent controls from the event onto the row", () => {
		const controls: AgentControls = { steer: vi.fn(), abort: vi.fn() };
		const s = fold([
			pe({ agentId: "a", phase: "p1", kind: "start", controls }),
			pe({ agentId: "a", phase: "p1", kind: "delta", text: "x" }), // later event w/o controls
		]);
		expect(s.rows[0].controls).toBe(controls); // retained across later events
	});

	it("returns a new state object and preserves identity of unchanged rows", () => {
		const s0 = fold([
			pe({ agentId: "a", phase: "p1", kind: "start" }),
			pe({ agentId: "b", phase: "p1", kind: "start" }),
		]);
		const s1 = reduceProgress(
			s0,
			pe({ agentId: "a", phase: "p1", kind: "end", status: "done" }),
		);
		expect(s1).not.toBe(s0); // pure: new state
		expect(s1.rows[1]).toBe(s0.rows[1]); // unchanged row keeps reference (ui memoization)
		expect(s1.rows[0]).not.toBe(s0.rows[0]); // changed row is fresh
	});

	it("merges tool start and end into one transcript item by tool call id", () => {
		const result = {
			content: [{ type: "text", text: "file contents" }],
			details: { lines: 1 },
		};
		const s = fold([
			pe({
				agentId: "a",
				phase: "p1",
				kind: "action",
				toolCallId: "tool-1",
				toolName: "read",
				args: { path: "file.ts" },
			}),
			pe({
				agentId: "a",
				phase: "p1",
				kind: "action",
				toolCallId: "tool-1",
				toolName: "read",
				result: { content: [{ type: "text", text: "partial" }] },
				isPartial: true,
			}),
			pe({
				agentId: "a",
				phase: "p1",
				kind: "action",
				toolCallId: "tool-1",
				toolName: "read",
				result,
				isError: false,
				isPartial: false,
			}),
		]);

		expect(s.transcriptItems.a).toHaveLength(1);
		expect(s.transcriptItems.a[0]).toMatchObject({
			kind: "tool",
			toolCallId: "tool-1",
			toolName: "read",
			args: { path: "file.ts" },
			result,
			isError: false,
			isPartial: false,
		});
	});

	it("end-to-end: mapped session events fold into the expected board state", () => {
		const raw = [
			{ type: "agent_start" },
			{
				type: "tool_execution_start",
				toolCallId: "t1",
				toolName: "grep",
				args: { pattern: "x" },
			},
			{
				type: "message_update",
				message: {},
				assistantMessageEvent: {
					type: "text_delta",
					contentIndex: 0,
					delta: "done.",
					partial: {},
				},
			},
			{ type: "agent_end", messages: [], willRetry: false },
		];
		const mapped = raw
			.map((r) => mapSessionEvent(ev(r), "g0", "review"))
			.filter((e): e is UltraProgressEvent => e !== null);
		const s = mapped.reduce(reduceProgress, initialReducerState());

		expect(s.rows).toHaveLength(1);
		expect(s.rows[0]).toMatchObject({
			agentId: "g0",
			phase: "review",
			status: "done",
		});
		expect(s.rows[0].action).toBe("Searching x");
		expect(s.transcripts.g0).toContain("done.");
		expect(s).toMatchObject({ total: 1, done: 1 });
	});
});