Luigit
repositories / pi-ext

pi-ext

bugabingas pi extensions

owned by admin

extensions/afk/__serial__/context.test.ts

Raw
import { fileURLToPath } from "node:url";
import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai";
import { estimateTokens } from "@earendil-works/pi-coding-agent";
import { afterEach, expect, it } from "vitest";
import {
	createPiSession,
	type PiSession,
} from "../../../test/fixture/create-pi-session.ts";

let session: PiSession | undefined;
afterEach(async () => {
	await session?.dispose();
	session = undefined;
});

it("measures same-history inactive and activated context without live API calls", async () => {
	const startup: number[] = [];
	const continued: number[] = [];
	for (const mode of ["stock", "inactive", "active"] as const) {
		session = await createPiSession({
			extensionPaths:
				mode === "stock"
					? []
					: [fileURLToPath(new URL("../index.ts", import.meta.url))],
			responses: [
				(context) => {
					startup.push(
						context.messages.reduce(
							(sum, message) => sum + estimateTokens(message),
							0,
						),
					);
					return fauxAssistantMessage("Started.");
				},
				(context) => {
					continued.push(
						context.messages.reduce(
							(sum, message) => sum + estimateTokens(message),
							0,
						),
					);
					return mode === "active"
						? fauxAssistantMessage(fauxToolCall("afk_complete", {}), {
								stopReason: "toolUse",
							})
						: fauxAssistantMessage("Done.");
				},
				fauxAssistantMessage("Done."),
			],
		});
		await session.session.prompt("Finish the current task.");
		if (mode === "active") {
			const command = session.extensions[0]?.commands.get("afk");
			if (!command) throw new Error("Missing /afk");
			await command.handler(
				"",
				session.session.extensionRunner.createContext() as never,
			);
			await session.session.waitForIdle();
		} else {
			await session.session.sendCustomMessage(
				{
					customType: "continue",
					content: "Continue requested work.",
					display: false,
				},
				{ triggerTurn: true },
			);
		}
		await session.dispose();
		session = undefined;
	}
	expect(startup[1]).toBe(startup[0]);
	expect(startup[2]).toBe(startup[0]);
	expect(continued[1]).toBe(continued[0]);
	expect(continued[2]).toBeGreaterThan(continued[0]!);
	console.info(
		`AFK context heuristic: inactive startup delta=${startup[1]! - startup[0]!}; activated same-history request delta=${continued[2]! - continued[0]!}. Provider-token delta not measured.`,
	);
});