pi-ext

import type {
	ExtensionAPI,
	ExtensionContext,
	ThemeColor,
} from "@earendil-works/pi-coding-agent";
import { closeDebug, dbg } from "./src/debug.ts";
import {
	offerFooterSegment,
	removeFooterSegment,
} from "./src/pi-ext-footer-segment.ts";

const STATUS_ID = "tps";
const SEGMENT_ID = "tps";
const ICON = "󰓅"; // nf-md-speedometer U+F04C5
const IDLE_TEXT = "—/s";
const LIVE_UPDATE_MS = 250;
const CHARS_PER_TOKEN = 4;

type Usage = {
	output?: number;
};

type AssistantLike = {
	role?: unknown;
	usage?: Usage;
};

type Clock = () => number;

function isAssistantMessage(message: unknown): message is AssistantLike {
	return (
		!!message &&
		typeof message === "object" &&
		(message as AssistantLike).role === "assistant"
	);
}

function add(value: number | undefined): number {
	return typeof value === "number" && Number.isFinite(value) ? value : 0;
}

function sumAssistantOutput(messages: readonly unknown[]): number {
	let output = 0;
	for (const message of messages) {
		if (!isAssistantMessage(message)) continue;
		output += add(message.usage?.output);
	}
	return output;
}

function finalText(outputTokens: number, assistantElapsedMs: number): string {
	if (outputTokens <= 0 || assistantElapsedMs <= 0) return IDLE_TEXT;
	return `${(outputTokens / (assistantElapsedMs / 1000)).toFixed(1)}/s`;
}

function liveText(chars: number, elapsedMs: number): string {
	const seconds = elapsedMs / 1000;
	if (seconds <= 0) return "0.0/s";
	return `${(chars / CHARS_PER_TOKEN / seconds).toFixed(1)}/s`;
}

function deltaLength(event: {
	assistantMessageEvent: { type: string; delta?: string };
}): number {
	const update = event.assistantMessageEvent;
	return update.type.endsWith("_delta") && typeof update.delta === "string"
		? update.delta.length
		: 0;
}

class TpsRun {
	private assistantStartMs: number | null = null;
	private decodeStartMs: number | null = null;
	private assistantElapsedMs = 0;
	private measuredElapsedMs = 0;
	private streamedChars = 0;
	private measuredChars = 0;
	private spanDeltaCount = 0;
	private lastLiveUpdateMs = 0;
	outcome: "measured" | "fallback" | "idle" = "idle";

	constructor(private readonly now: Clock = Date.now) {}

	startAgent(): void {
		this.assistantStartMs = null;
		this.decodeStartMs = null;
		this.assistantElapsedMs = 0;
		this.measuredElapsedMs = 0;
		this.streamedChars = 0;
		this.measuredChars = 0;
		this.spanDeltaCount = 0;
		this.lastLiveUpdateMs = 0;
	}

	startAssistant(): void {
		if (this.assistantStartMs !== null) return;
		this.assistantStartMs = this.now();
		this.decodeStartMs = null;
		this.spanDeltaCount = 0;
	}

	endAssistant(): void {
		const now = this.now();
		if (this.assistantStartMs !== null)
			this.assistantElapsedMs += Math.max(0, now - this.assistantStartMs);
		if (this.decodeStartMs !== null && this.spanDeltaCount > 1)
			this.measuredElapsedMs += Math.max(0, now - this.decodeStartMs);
		this.assistantStartMs = null;
		this.decodeStartMs = null;
		this.spanDeltaCount = 0;
	}

	addDelta(chars: number): void {
		const safeChars = Math.max(0, chars);
		if (safeChars === 0) return;
		if (this.assistantStartMs === null) this.startAssistant();
		this.streamedChars += safeChars;
		this.spanDeltaCount++;
		if (this.spanDeltaCount === 1) {
			this.decodeStartMs = this.now();
			return;
		}
		this.measuredChars += safeChars;
	}

	shouldPublishLive(): boolean {
		const now = this.now();
		if (now - this.lastLiveUpdateMs < LIVE_UPDATE_MS) return false;
		this.lastLiveUpdateMs = now;
		return true;
	}

	liveText(): string {
		return liveText(this.measuredChars, this.currentMeasuredElapsedMs());
	}

	finish(messages: readonly unknown[]): string {
		this.endAssistant();
		const outputTokens = sumAssistantOutput(messages);
		if (
			this.measuredChars > 0 &&
			this.streamedChars > 0 &&
			this.measuredElapsedMs > 0
		) {
			this.outcome = outputTokens > 0 ? "measured" : "idle";
			return finalText(
				outputTokens * (this.measuredChars / this.streamedChars),
				this.measuredElapsedMs,
			);
		}
		this.outcome =
			outputTokens > 0 && this.assistantElapsedMs > 0 ? "fallback" : "idle";
		return finalText(outputTokens, this.assistantElapsedMs);
	}

	private currentMeasuredElapsedMs(): number {
		if (this.decodeStartMs !== null && this.spanDeltaCount > 1)
			return (
				this.measuredElapsedMs + Math.max(0, this.now() - this.decodeStartMs)
			);
		if (this.streamedChars === 0 && this.assistantStartMs !== null)
			return (
				this.assistantElapsedMs +
				Math.max(0, this.now() - this.assistantStartMs)
			);
		return this.measuredElapsedMs;
	}
}

function clearSegment(pi: ExtensionAPI): void {
	removeFooterSegment(pi, SEGMENT_ID);
}

function publishPresentation(
	pi: ExtensionAPI,
	ctx: ExtensionContext,
	text: string | undefined,
	color: ThemeColor,
): void {
	if (!ctx.hasUI) return;

	try {
		let accepted = false;
		if (ctx.mode === "tui") {
			accepted =
				text === undefined
					? removeFooterSegment(pi, SEGMENT_ID)
					: offerFooterSegment(pi, {
							id: SEGMENT_ID,
							icon: ICON,
							text,
							variants: [text.replace(/\.\d\//, "/")],
							color,
							zone: "llm",
							order: 4,
						});
		}

		if (!accepted) clearSegment(pi);
		ctx.ui.setStatus(
			STATUS_ID,
			accepted || text === undefined ? undefined : `${ICON} ${text}`,
		);
	} catch {
		// Ignore stale ctx during reload/switch/shutdown.
	}
}

export default function tpsExtension(pi: ExtensionAPI) {
	let run: TpsRun | null = null;

	function publish(
		ctx: ExtensionContext,
		text: string | undefined,
		color: ThemeColor = "accent",
	): void {
		publishPresentation(pi, ctx, text, color);
	}

	pi.on("session_start", async (_event, ctx) => {
		dbg?.("session.start");
		publish(ctx, IDLE_TEXT, "dim");
	});

	pi.on("agent_start", async (_event, ctx) => {
		run = new TpsRun();
		run.startAgent();
		publish(ctx, "0.0/s", "dim");
	});

	pi.on("message_start", async (event, _ctx) => {
		if (!run || !isAssistantMessage(event.message)) return;
		run.startAssistant();
	});

	pi.on("message_update", async (event, ctx) => {
		if (!run || !isAssistantMessage(event.message)) return;
		run.startAssistant();
		run.addDelta(deltaLength(event));
		if (!run.shouldPublishLive()) return;
		publish(ctx, run.liveText(), "accent");
	});

	pi.on("message_end", async (event, _ctx) => {
		if (!run || !isAssistantMessage(event.message)) return;
		run.endAssistant();
	});

	pi.on("agent_end", async (event, ctx) => {
		if (!run) return;
		const text = run.finish(event.messages);
		dbg?.("rate.finish", { outcome: run.outcome });
		run = null;
		publish(ctx, text, text === IDLE_TEXT ? "dim" : "success");
	});

	pi.on("session_shutdown", async (_event, ctx) => {
		dbg?.("session.shutdown");
		closeDebug();
		run = null;
		publish(ctx, undefined);
	});
}

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