Luigit
repositories / pi-ext

pi-ext

bugabingas pi extensions

owned by admin

extensions/ultra/__tests__/action-summarizer.test.ts

Raw
import { describe, expect, it, vi } from "vitest";
import { __test, createActionSummarizer } from "../action-summarizer.ts";

function context() {
	const model = { provider: "fast", id: "tiny" };
	const find = vi.fn(() => model);
	const complete = vi.fn(async () => response("Routing through registry"));
	return {
		model,
		find,
		complete,
		ctx: { modelRegistry: { find, complete } },
	};
}

function fakeScheduler() {
	const timers: Array<() => void> = [];
	return {
		scheduler: {
			setTimeout(callback: () => void) {
				timers.push(callback);
				return callback;
			},
			clearTimeout(handle: () => void) {
				const index = timers.indexOf(handle);
				if (index >= 0) timers.splice(index, 1);
			},
		},
		flush() {
			for (const timer of timers.splice(0)) timer();
		},
		pending() {
			return timers.length;
		},
	};
}

async function settle() {
	await new Promise((resolve) => setTimeout(resolve, 0));
}

const input = (agentId: string, toolName = "read") => ({
	agentId,
	toolName,
	args: { path: "extensions/ultra/settings.ts" },
	taskSummary: "Inspect Ultra settings",
});

const response = (text: string) => ({
	stopReason: "stop",
	content: [{ type: "text", text }],
});

describe("action summarizer", () => {
	it("is disabled when the configured model is null", () => {
		const { ctx, find } = context();
		expect(createActionSummarizer(ctx as never, null)).toBeUndefined();
		expect(find).not.toHaveBeenCalled();
	});

	it("uses the configured cheap model with bounded low-latency options", async () => {
		const { ctx, model, find } = context();
		const scheduler = fakeScheduler();
		const generate = vi.fn(async () => response('"Reading Ultra settings."\n'));
		const summarize = createActionSummarizer(ctx as never, "fast/tiny", {
			generate: generate as never,
			scheduler: scheduler.scheduler,
		});
		const pending = summarize?.({
			...input("a"),
			args: { path: "extensions/ultra/settings.ts", noise: "x".repeat(3000) },
		});

		expect(generate).not.toHaveBeenCalled();
		scheduler.flush();
		await expect(pending).resolves.toBe("Reading Ultra settings");
		expect(find).toHaveBeenCalledWith("fast", "tiny");
		expect(generate).toHaveBeenCalledTimes(1);
		expect(generate.mock.calls[0][0]).toBe(model);
		expect(generate.mock.calls[0][1].systemPrompt).toContain(
			"Start with an -ing verb",
		);
		expect(
			generate.mock.calls[0][1].messages[0].content[0].text.length,
		).toBeLessThan(1700);
		expect(generate.mock.calls[0][2]).toMatchObject({
			maxTokens: 40,
			maxRetries: 0,
			cacheRetention: "none",
			timeoutMs: 2000,
		});
	});

	it("routes default generation through the configured model registry", async () => {
		const { ctx, complete } = context();
		const scheduler = fakeScheduler();
		const summarize = createActionSummarizer(ctx as never, "fast/tiny", {
			scheduler: scheduler.scheduler,
		});
		const pending = summarize?.(input("a"));

		scheduler.flush();

		await expect(pending).resolves.toBe("Routing through registry");
		expect(complete).toHaveBeenCalledOnce();
	});

	it("coalesces rapid actions per agent to the latest call", async () => {
		const { ctx } = context();
		const scheduler = fakeScheduler();
		const generate = vi.fn(async () => response("Finding current references"));
		const summarize = createActionSummarizer(ctx as never, "fast/tiny", {
			generate: generate as never,
			scheduler: scheduler.scheduler,
		});

		const stale = summarize?.(input("a", "read"));
		const latest = summarize?.(input("a", "grep"));
		scheduler.flush();

		await expect(stale).resolves.toBeUndefined();
		await expect(latest).resolves.toBe("Finding current references");
		expect(generate).toHaveBeenCalledTimes(1);
		expect(generate.mock.calls[0][1].messages[0].content[0].text).toContain(
			"Tool: grep",
		);
	});

	it("debounces agents independently", () => {
		const { ctx } = context();
		const scheduler = fakeScheduler();
		const summarize = createActionSummarizer(ctx as never, "fast/tiny", {
			generate: vi.fn() as never,
			scheduler: scheduler.scheduler,
		});

		void summarize?.(input("a"));
		void summarize?.(input("b"));
		void summarize?.(input("a", "grep"));

		expect(scheduler.pending()).toBe(2);
		summarize?.dispose();
	});

	it("flushes a fast agent's latest action when it finishes", async () => {
		const { ctx } = context();
		const scheduler = fakeScheduler();
		const generate = vi.fn(async () => response("Reading Ultra settings"));
		const summarize = createActionSummarizer(ctx as never, "fast/tiny", {
			generate: generate as never,
			scheduler: scheduler.scheduler,
		});
		const pending = summarize?.(input("a"));

		summarize?.finish("a");

		await expect(pending).resolves.toBe("Reading Ultra settings");
		expect(scheduler.pending()).toBe(0);
		expect(generate).toHaveBeenCalledTimes(1);
	});

	it("allows only one model request in flight across agents", async () => {
		const { ctx } = context();
		const scheduler = fakeScheduler();
		const releases: Array<(value: unknown) => void> = [];
		const generate = vi.fn(
			() => new Promise((resolve) => releases.push(resolve)),
		);
		const summarize = createActionSummarizer(ctx as never, "fast/tiny", {
			generate: generate as never,
			scheduler: scheduler.scheduler,
		});

		const first = summarize?.(input("a"));
		const second = summarize?.(input("b"));
		scheduler.flush();
		await settle();
		expect(generate).toHaveBeenCalledTimes(1);

		releases.shift()?.(response("Reading first action"));
		await settle();
		expect(generate).toHaveBeenCalledTimes(2);
		releases.shift()?.(response("Reading second action"));
		await expect(first).resolves.toBe("Reading first action");
		await expect(second).resolves.toBe("Reading second action");
	});

	it("cancels queued work when its agent is aborted", async () => {
		const { ctx } = context();
		const scheduler = fakeScheduler();
		const generate = vi.fn(async () => response("Should not run"));
		const summarize = createActionSummarizer(ctx as never, "fast/tiny", {
			generate: generate as never,
			scheduler: scheduler.scheduler,
		});
		const pending = summarize?.(input("a"));

		summarize?.cancel?.("a");
		scheduler.flush();
		await expect(pending).resolves.toBeUndefined();
		expect(generate).not.toHaveBeenCalled();
	});

	it("cancels every pending action when disposed", async () => {
		const { ctx } = context();
		const scheduler = fakeScheduler();
		const generate = vi.fn(async () => response("Should not run"));
		const summarize = createActionSummarizer(ctx as never, "fast/tiny", {
			generate: generate as never,
			scheduler: scheduler.scheduler,
		});
		const a = summarize?.(input("a"));
		const b = summarize?.(input("b"));

		summarize?.dispose();
		scheduler.flush();

		await expect(a).resolves.toBeUndefined();
		await expect(b).resolves.toBeUndefined();
		expect(generate).not.toHaveBeenCalled();
	});

	it("fails silently when generation rejects", async () => {
		const { ctx } = context();
		ctx.modelRegistry.complete = vi.fn(async () => {
			throw new Error("Provider is not configured");
		});
		const scheduler = fakeScheduler();
		const summarize = createActionSummarizer(ctx as never, "fast/tiny", {
			scheduler: scheduler.scheduler,
		});
		const pending = summarize?.(input("a"));
		scheduler.flush();
		await expect(pending).resolves.toBeUndefined();
	});

	it("rejects malformed model refs", () => {
		const { ctx } = context();
		expect(createActionSummarizer(ctx as never, "tiny")).toBeUndefined();
		expect(__test.parseModelRef("fast/tiny")).toEqual({
			provider: "fast",
			id: "tiny",
		});
	});
});