pi-ext

import { describe, expect, it } from "vitest";
import { __test } from "../index.ts";

const { TpsRun, deltaLength, finalText, sumAssistantOutput } = __test;

function assistant(output: number) {
	return { role: "assistant", usage: { output } };
}

describe("tps calculation", () => {
	it("uses first streamed delta as baseline, not numerator", () => {
		let now = 0;
		const run = new TpsRun(() => now);

		run.startAgent();
		now = 1_000;
		run.startAssistant();
		now = 2_000;
		run.addDelta(400);
		now = 3_000;
		run.addDelta(4);
		run.endAssistant();
		now = 11_000;

		expect(run.finish([assistant(101)])).toBe("1.0/s");
	});

	it("sums multiple streamed spans and excludes tool gaps", () => {
		let now = 0;
		const run = new TpsRun(() => now);

		run.startAgent();
		run.startAssistant();
		run.addDelta(4);
		now = 1_000;
		run.addDelta(4);
		run.endAssistant();
		now = 5_000;
		run.startAssistant();
		run.addDelta(4);
		now = 6_000;
		run.addDelta(4);
		run.endAssistant();

		expect(
			run.finish([assistant(50), { role: "toolResult" }, assistant(50)]),
		).toBe("25.0/s");
	});

	it("closes an open streamed span at finish", () => {
		let now = 0;
		const run = new TpsRun(() => now);

		run.startAgent();
		run.startAssistant();
		now = 1_000;
		run.addDelta(4);
		now = 3_000;
		run.addDelta(4);

		expect(run.finish([assistant(40)])).toBe("10.0/s");
	});

	it("falls back to assistant span when only one delta streams", () => {
		let now = 0;
		const run = new TpsRun(() => now);

		run.startAgent();
		run.startAssistant();
		now = 1_000;
		run.addDelta(400);
		now = 3_000;
		run.endAssistant();

		expect(run.finish([assistant(40)])).toBe("13.3/s");
	});

	it("falls back to assistant span when no deltas stream", () => {
		let now = 0;
		const run = new TpsRun(() => now);

		run.startAgent();
		run.startAssistant();
		now = 2_000;
		run.endAssistant();

		expect(run.finish([assistant(40)])).toBe("20.0/s");
	});

	it("returns idle text without output tokens or measured duration", () => {
		expect(finalText(0, 1_000)).toBe("—/s");
		expect(finalText(10, 0)).toBe("—/s");
	});

	it("sums only assistant usage output", () => {
		expect(
			sumAssistantOutput([
				assistant(3),
				{ role: "user", usage: { output: 100 } },
				assistant(Number.NaN),
				assistant(2),
			]),
		).toBe(5);
	});

	it("estimates live deltas only from string delta events", () => {
		expect(
			deltaLength({
				assistantMessageEvent: { type: "text_delta", delta: "abcd" },
			}),
		).toBe(4);
		expect(
			deltaLength({
				assistantMessageEvent: { type: "thinking_delta", delta: "ab" },
			}),
		).toBe(2);
		expect(deltaLength({ assistantMessageEvent: { type: "text_delta" } })).toBe(
			0,
		);
		expect(
			deltaLength({
				assistantMessageEvent: { type: "text" as string, delta: "abcd" },
			}),
		).toBe(0);
	});
});