pi-extensions/shared/token-stats.ts

export interface PiTokenStats {
	requestStartMs: number;
	firstOutputMs?: number;
	responseEndMs?: number;
	inputTokens?: number;
	outputTokens?: number;
	thinkingTokens?: number;
}

export interface PiTokenRateSummary {
	processingTokensPerSecond: number | null;
	generationTokensPerSecond: number | null;
}

function readNumber(value: unknown): number | null {
	if (typeof value !== "number" || !Number.isFinite(value)) {
		return null;
	}
	return value;
}

export function captureFirstOutput(stats: PiTokenStats, nowMs: number): void {
	if (stats.firstOutputMs === undefined && Number.isFinite(nowMs)) {
		stats.firstOutputMs = nowMs;
	}
}

export function finalizePiTokenStats(
	stats: PiTokenStats,
	usage: { input: number; output: number; thinking?: number },
	nowMs: number,
): PiTokenStats {
	return {
		requestStartMs: stats.requestStartMs,
		firstOutputMs: stats.firstOutputMs,
		responseEndMs: nowMs,
		inputTokens: usage.input,
		outputTokens: usage.output,
		thinkingTokens: usage.thinking,
	};
}

export function readPiTokenStats(message: unknown): PiTokenStats | null {
	if (!message || typeof message !== "object") {
		return null;
	}
	const candidate = (message as { piTokenStats?: unknown }).piTokenStats;
	if (!candidate || typeof candidate !== "object") {
		return null;
	}
	const stats = candidate as Record<string, unknown>;
	const requestStartMs = readNumber(stats.requestStartMs);
	if (requestStartMs === null) {
		return null;
	}
	const firstOutputMs = readNumber(stats.firstOutputMs) ?? undefined;
	const responseEndMs = readNumber(stats.responseEndMs) ?? undefined;
	const inputTokens = readNumber(stats.inputTokens) ?? undefined;
	const outputTokens = readNumber(stats.outputTokens) ?? undefined;
	const thinkingTokens = readNumber(stats.thinkingTokens) ?? undefined;
	return {
		requestStartMs,
		firstOutputMs,
		responseEndMs,
		inputTokens,
		outputTokens,
		thinkingTokens,
	};
}

export function normalizePiEntries(entries: unknown): readonly unknown[] {
	if (Array.isArray(entries)) {
		return entries;
	}
	if (entries && typeof entries === "object") {
		const candidate = entries as {
			entries?: unknown;
			[Symbol.iterator]?: unknown;
		};
		if (Array.isArray(candidate.entries)) {
			return candidate.entries;
		}
		if (typeof candidate[Symbol.iterator] === "function") {
			try {
				return Array.from(candidate as Iterable<unknown>);
			} catch {
				return [];
			}
		}
	}
	return [];
}

export function findLatestPiTokenStats(entries: unknown): PiTokenStats | null {
	const normalizedEntries = normalizePiEntries(entries);
	for (let index = normalizedEntries.length - 1; index >= 0; index--) {
		const entry = normalizedEntries[index];
		if (!entry || typeof entry !== "object") {
			continue;
		}
		const candidate = entry as {
			type?: string;
			message?: {
				role?: string;
			};
		};
		if (candidate.type !== "message") {
			continue;
		}
		if (candidate.message?.role !== "assistant") {
			continue;
		}
		const stats = readPiTokenStats(candidate.message);
		if (stats) {
			return stats;
		}
		break;
	}
	return null;
}

function computeRate(tokens: number | undefined, startMs: number | undefined, endMs: number | undefined): number | null {
	if (
		tokens === undefined
		|| startMs === undefined
		|| endMs === undefined
		|| tokens <= 0
		|| endMs <= startMs
	) {
		return null;
	}
	return tokens / ((endMs - startMs) / 1000);
}

export function summarizePiTokenStats(stats: PiTokenStats): PiTokenRateSummary {
	const processingEndMs = stats.firstOutputMs ?? stats.responseEndMs;
	const generatedTokens = (stats.outputTokens ?? 0) + (stats.thinkingTokens ?? 0);
	return {
		processingTokensPerSecond: computeRate(stats.inputTokens, stats.requestStartMs, processingEndMs),
		generationTokensPerSecond: computeRate(generatedTokens, stats.firstOutputMs, stats.responseEndMs),
	};
}

export function formatPiTokenRateStatus(summary: PiTokenRateSummary): string | null {
	const parts: string[] = [];
	if (summary.processingTokensPerSecond !== null) {
		parts.push(`P ${summary.processingTokensPerSecond.toFixed(1)}tok/s`);
	}
	if (summary.generationTokensPerSecond !== null) {
		parts.push(`G ${summary.generationTokensPerSecond.toFixed(1)}tok/s`);
	}
	return parts.length > 0 ? parts.join("  ") : null;
}