@selesai/code 0.5.20 → 0.5.22

package/dist/extensions/pi-powerline-footer/tps.ts

@@ -1,39 +1,98 @@
 /**
  * TPS Tracker — tokens-per-second tracking merged into the powerline footer.
  *
- * Writes live and final TPS stats to `ctx.ui.setStatus("tps", ...)`, which the
- * powerline `extension_statuses` segment surfaces in the status bar (and the
- * `setExtensionStatus` repaint hook turns into a re-render automatically).
- *
- * Tracks:
- *  - Main model output tokens vs. streaming wall time (excludes first-token
- *    latency and tool execution gaps).
- *  - pi-subagents child output tokens via tool_execution_* events, aggregated
- *    into a combined "main + subagents" TPS at agent_end.
+ * Live status remains intentionally approximate. Final main-model TPS uses
+ * provider output usage plus pi-tps' reliability, stall, and tool-call gates.
+ * Subagent UI/aggregation remains separate.
  */
 
 import type { ExtensionAPI } from "@selesai/code";
 
+const STALL_THRESHOLD_MS = 500;
+const MAX_PLAUSIBLE_TPS = 10_000;
+const MIN_STREAM_MS = 1;
+const MIN_STREAM_UPDATES = 5;
+const MIN_INTER_CHUNK_MS = 1;
+const MIN_GENERATION_MS = 200;
+const ACTIVE_TIME_THRESHOLD_MS = 200;
+const STALL_REDUCTION_DENOM = 2;
+const STALL_DOMINANCE_RATIO = 0.85;
+
+export interface TpsTiming {
+	messageStartMs: number;
+	lastUpdateMs: number;
+	firstTokenMs: number | null;
+	updateCount: number;
+	firstStreamUpdateMs: number | null;
+	lastStreamUpdateMs: number;
+	stallMs: number;
+	generationMs: number;
+}
+
+export interface ReliableTps {
+	tps: number;
+	effectiveMs: number;
+	isPrimary: boolean;
+}
+
+/** Port of pi-tps' primary/fallback reliability and volume gates. */
+export function calculateReliableTps(outputTokens: number, timing: TpsTiming): ReliableTps | null {
+	if (outputTokens <= 0 || timing.firstTokenMs === null) return null;
+
+	const streamMs = timing.updateCount > 0 && timing.firstStreamUpdateMs !== null
+		? timing.lastStreamUpdateMs - timing.firstStreamUpdateMs
+		: null;
+	const avgInterChunkGap = streamMs !== null && timing.updateCount > 1
+		? streamMs / (timing.updateCount - 1)
+		: 0;
+
+	let tps: number | null = null;
+	let effectiveMs = 0;
+	let isPrimary = false;
+	if (
+		streamMs !== null &&
+		streamMs >= MIN_STREAM_MS &&
+		timing.updateCount >= MIN_STREAM_UPDATES &&
+		avgInterChunkGap >= MIN_INTER_CHUNK_MS &&
+		timing.stallMs < streamMs &&
+		streamMs - timing.stallMs >= MIN_GENERATION_MS &&
+		timing.stallMs < streamMs - timing.stallMs
+	) {
+		effectiveMs = streamMs - timing.stallMs;
+		tps = Math.round((outputTokens / (effectiveMs / 1000)) * 10) / 10;
+		isPrimary = true;
+	} else if (timing.updateCount >= 2 && timing.generationMs >= MIN_GENERATION_MS) {
+		const stallsDominate =
+			timing.generationMs - timing.stallMs < ACTIVE_TIME_THRESHOLD_MS ||
+			timing.stallMs > timing.generationMs * STALL_DOMINANCE_RATIO;
+		effectiveMs = stallsDominate
+			? Math.max(timing.generationMs - timing.stallMs / STALL_REDUCTION_DENOM, MIN_GENERATION_MS)
+			: Math.max(timing.generationMs - timing.stallMs, MIN_GENERATION_MS);
+		tps = Math.round((outputTokens / (effectiveMs / 1000)) * 10) / 10;
+	}
+
+	if (tps === null || tps > MAX_PLAUSIBLE_TPS) return null;
+	return { tps, effectiveMs, isPrimary };
+}
+
 export function setupTpsTracker(pi: ExtensionAPI): void {
-	/** Timestamp when the current assistant message event started. Used as a fallback. */
 	let messageStart: number | null = null;
-	/** Timestamp of the first streamed output delta for the current assistant message. */
 	let streamStart: number | null = null;
-	/** Estimated streamed output tokens for live display before providers report final usage. */
 	let estimatedStreamedTokens = 0;
-	/** Cumulative official output tokens across all assistant messages in this agent run. */
-	let totalOutputTokens = 0;
-	/** Cumulative time (ms) spent actually streaming output deltas (excludes tool execution and first-token latency). */
-	let totalStreamMs = 0;
-	/** Subagent tool-call start times, used to include pi-subagents child output in aggregate TPS. */
+	let mainTiming: TpsTiming | null = null;
+	const mainMeasurements: Array<{
+		outputTokens: number;
+		timing: TpsTiming;
+		modelKey: string | null;
+		isToolCall: boolean;
+	}> = [];
+	let pendingToolMeasurement: (typeof mainMeasurements)[number] | null = null;
+	const tpsCaps = new Map<string, number>();
+
 	const subagentStarts = new Map<string, number>();
-	/** Latest cumulative foreground subagent tokens seen during tool_execution_update. */
 	const subagentLiveTokens = new Map<string, number>();
-	/** Cumulative subagent output tokens in this agent run. */
 	let totalSubagentTokens = 0;
-	/** Cumulative subagent wall time (ms) in this agent run. */
 	let totalSubagentMs = 0;
-	/** Completed subagent call summaries in this agent run. */
 	const subagentSummaries: Array<{ tokens: number; ms: number; tps: number }> = [];
 
 	function asRecord(value: unknown): Record<string, unknown> | null {
@@ -61,25 +120,8 @@ export function setupTpsTracker(pi: ExtensionAPI): void {
 		return numberValue(usage?.output ?? usage?.outputTokens ?? usage?.completionTokens ?? usage?.completion_tokens);
 	}
 
-	function generatedTokensFromMessage(message: { usage?: unknown; content?: unknown }): number {
-		const usageTokens = generatedTokensFromUsage(asRecord(message.usage));
-		if (usageTokens > 0) return usageTokens;
-
-		const content = Array.isArray(message.content) ? message.content : [];
-		const chars = content.reduce((sum, block) => {
-			const item = asRecord(block);
-			if (!item) return sum;
-			if (item.type === "text") return sum + String(item.text ?? "").length;
-			if (item.type === "thinking") return sum + String(item.thinking ?? "").length;
-			if (item.type === "toolCall") return sum + String(item.name ?? "").length + JSON.stringify(item.arguments ?? {}).length;
-			return sum;
-		}, 0);
-		return Math.ceil(chars / 4);
-	}
-
 	function outputTokensFromResult(value: unknown): number {
-		const results = getResults(value);
-		return results.reduce((sum, result) => sum + generatedTokensFromUsage(asRecord(result.usage)), 0);
+		return getResults(value).reduce((sum, result) => sum + generatedTokensFromUsage(asRecord(result.usage)), 0);
 	}
 
 	function liveTokensFromResult(value: unknown): number {
@@ -117,8 +159,9 @@ export function setupTpsTracker(pi: ExtensionAPI): void {
 	}
 
 	pi.on("agent_start", async (_event, ctx) => {
-		totalOutputTokens = 0;
-		totalStreamMs = 0;
+		mainMeasurements.length = 0;
+		pendingToolMeasurement = null;
+		mainTiming = null;
 		totalSubagentTokens = 0;
 		totalSubagentMs = 0;
 		subagentSummaries.length = 0;
@@ -127,77 +170,99 @@ export function setupTpsTracker(pi: ExtensionAPI): void {
 		estimatedStreamedTokens = 0;
 		subagentStarts.clear();
 		subagentLiveTokens.clear();
-		const theme = ctx.ui.theme;
-		ctx.ui.setStatus("tps", theme.fg("success", "generating..."));
+		ctx.ui.setStatus("tps", ctx.ui.theme.fg("success", "generating..."));
 	});
 
 	pi.on("message_start", async (event) => {
 		if (event.message.role !== "assistant") return;
-		messageStart = Date.now();
+		const now = performance.now();
+		messageStart = now;
 		streamStart = null;
 		estimatedStreamedTokens = 0;
+		mainTiming = {
+			messageStartMs: now,
+			lastUpdateMs: now,
+			firstTokenMs: null,
+			updateCount: 0,
+			firstStreamUpdateMs: null,
+			lastStreamUpdateMs: 0,
+			stallMs: 0,
+			generationMs: 0,
+		};
+		pendingToolMeasurement = null;
 	});
 
 	pi.on("message_update", async (event, ctx) => {
 		if (event.message.role !== "assistant") return;
 
+		const now = performance.now();
+		if (mainTiming) {
+			if (mainTiming.firstTokenMs === null) {
+				mainTiming.firstTokenMs = now;
+				mainTiming.lastUpdateMs = now;
+			} else {
+				mainTiming.updateCount++;
+				mainTiming.firstStreamUpdateMs ??= now;
+				mainTiming.lastStreamUpdateMs = now;
+				const gap = now - mainTiming.lastUpdateMs;
+				if (gap >= STALL_THRESHOLD_MS) mainTiming.stallMs += gap;
+				mainTiming.lastUpdateMs = now;
+			}
+		}
+
 		const streamEvent = event.assistantMessageEvent;
 		const isOutputDelta =
 			streamEvent.type === "text_delta" ||
 			streamEvent.type === "thinking_delta" ||
 			streamEvent.type === "toolcall_delta";
-
 		if (!isOutputDelta) return;
 
-		const now = Date.now();
 		streamStart ??= now;
 		estimatedStreamedTokens += Math.max(0, streamEvent.delta.length / 4);
-
 		const elapsed = (now - streamStart) / 1000;
 		const officialTokens = generatedTokensFromUsage(asRecord(event.message.usage));
 		const currentTokens = officialTokens > 0 ? officialTokens : estimatedStreamedTokens;
-
 		if (elapsed > 0 && currentTokens > 0) {
-			const tps = Math.round(currentTokens / elapsed);
-			const tokenLabel = officialTokens > 0
-				? `${officialTokens} tok`
-				: `~${Math.round(estimatedStreamedTokens)} tok`;
-			const theme = ctx.ui.theme;
-			ctx.ui.setStatus(
-				"tps",
-				`${theme.fg("accent", `${tps} tok/s`)}`,
-			);
+			ctx.ui.setStatus("tps", ctx.ui.theme.fg("accent", `${Math.round(currentTokens / elapsed)} tok/s`));
 		}
 	});
 
 	pi.on("message_end", async (event) => {
 		if (event.message.role !== "assistant") return;
 
-		const messageTokens = generatedTokensFromMessage(event.message);
-		const timingStart = streamStart ?? messageStart;
-		if (!timingStart || messageTokens <= 0) {
-			messageStart = null;
-			streamStart = null;
-			estimatedStreamedTokens = 0;
-			return;
+		const now = performance.now();
+		if (mainTiming) {
+			mainTiming.generationMs = now - mainTiming.messageStartMs;
+			const message = asRecord(event.message);
+			const provider = typeof message?.provider === "string" ? message.provider : null;
+			const model = typeof message?.model === "string" ? message.model : null;
+			const outputTokens = generatedTokensFromUsage(asRecord(message?.usage));
+			if (outputTokens > 0) {
+				const measurement = {
+					outputTokens,
+					timing: mainTiming,
+					modelKey: provider && model ? `${provider}:${model}` : null,
+					isToolCall: false,
+				};
+				mainMeasurements.push(measurement);
+				pendingToolMeasurement = measurement;
+			}
 		}
 
-		totalOutputTokens += messageTokens;
-		totalStreamMs += Math.max(0, Date.now() - timingStart);
-
+		mainTiming = null;
 		messageStart = null;
 		streamStart = null;
 		estimatedStreamedTokens = 0;
 	});
 
 	pi.on("tool_execution_start", async (event) => {
+		if (pendingToolMeasurement) pendingToolMeasurement.isToolCall = true;
 		if (event.toolName !== "subagent") return;
 		subagentStarts.set(event.toolCallId, Date.now());
 	});
 
 	pi.on("tool_execution_update", async (event, ctx) => {
 		if (event.toolName !== "subagent") return;
-
 		const start = subagentStarts.get(event.toolCallId);
 		if (!start) return;
 
@@ -209,17 +274,13 @@ export function setupTpsTracker(pi: ExtensionAPI): void {
 
 		const tps = Math.round(tokens / elapsed);
 		const mode = subagentModeLabel(event.partialResult);
-		const agentSummary = items
-			.filter((item) => item.tokens > 0 || item.status === "running")
-			.map((item) => `${item.name} ${Math.round(item.tokens / elapsed)} t/s`)
-			.join(" ");
-		const theme = ctx.ui.theme;
-		ctx.ui.setStatus("tps", `${theme.fg("accent", `main+${mode} ${tps} t/s`)} ${theme.fg("success", `(${agentSummary || `${tps} t/s`})`)}`);
+		const agentSummary = items.filter((item) => item.tokens > 0 || item.status === "running")
+			.map((item) => `${item.name} ${Math.round(item.tokens / elapsed)} t/s`).join(" ");
+		ctx.ui.setStatus("tps", `${ctx.ui.theme.fg("accent", `main+${mode} ${tps} t/s`)} ${ctx.ui.theme.fg("success", `(${agentSummary || `${tps} t/s`})`)}`);
 	});
 
 	pi.on("tool_execution_end", async (event, ctx) => {
 		if (event.toolName !== "subagent") return;
-
 		const start = subagentStarts.get(event.toolCallId);
 		subagentStarts.delete(event.toolCallId);
 		const liveTokens = subagentLiveTokens.get(event.toolCallId) ?? liveTokensFromResult(event.result);
@@ -228,59 +289,57 @@ export function setupTpsTracker(pi: ExtensionAPI): void {
 
 		const outputTokens = outputTokensFromResult(event.result);
 		const displayTokens = outputTokens || liveTokens;
-		const mode = subagentModeLabel(event.result);
 		const elapsedMs = Math.max(0, Date.now() - start);
 		if (displayTokens <= 0 || elapsedMs <= 0) return;
 		const elapsedSeconds = elapsedMs / 1000;
 		const items = subagentProgressItems(event.result);
-		const agentSummary = items
-			.filter((item) => item.tokens > 0)
-			.map((item) => `${item.name} ${Math.round(item.tokens / elapsedSeconds)} t/s`)
-			.join(" ");
+		const agentSummary = items.filter((item) => item.tokens > 0)
+			.map((item) => `${item.name} ${Math.round(item.tokens / elapsedSeconds)} t/s`).join(" ");
+		const tps = Math.round(displayTokens / elapsedSeconds);
+		const mode = subagentModeLabel(event.result);
 
-		totalOutputTokens += displayTokens;
-		totalStreamMs += elapsedMs;
 		totalSubagentTokens += displayTokens;
 		totalSubagentMs += elapsedMs;
-
-		const tps = Math.round(displayTokens / elapsedSeconds);
 		subagentSummaries.push({ tokens: displayTokens, ms: elapsedMs, tps });
-		const theme = ctx.ui.theme;
-		ctx.ui.notify(`${theme.fg("success", "✓")} ${theme.fg("accent", `${tps} t/s`)}`, "info");
-		ctx.ui.setStatus(
-			"tps",
-			`${theme.fg("accent", `main+${mode} ${tps} t/s`)} ${theme.fg("success", `(${agentSummary || `${tps} t/s`})`)}`,
-		);
+		ctx.ui.notify(`${ctx.ui.theme.fg("success", "✓")} ${ctx.ui.theme.fg("accent", `${tps} t/s`)}`, "info");
+		ctx.ui.setStatus("tps", `${ctx.ui.theme.fg("accent", `main+${mode} ${tps} t/s`)} ${ctx.ui.theme.fg("success", `(${agentSummary || `${tps} t/s`})`)}`);
 	});
 
 	pi.on("agent_end", async (_event, ctx) => {
-		const elapsed = totalStreamMs / 1000;
-		const tps = totalOutputTokens > 0 && elapsed > 0 ? Math.round(totalOutputTokens / elapsed) : 0;
+		let measuredMainTokens = 0;
+		let measuredMainMs = 0;
+		for (const measurement of mainMeasurements) {
+			if (!measurement.modelKey) continue;
+			const reliable = calculateReliableTps(measurement.outputTokens, measurement.timing);
+			if (!reliable) continue;
+
+			let tps = reliable.tps;
+			if (measurement.isToolCall) {
+				const cap = tpsCaps.get(measurement.modelKey);
+				if (cap === undefined) continue;
+				tps = Math.min(tps, cap);
+			} else if (reliable.isPrimary) {
+				const cap = tpsCaps.get(measurement.modelKey);
+				if (cap === undefined || tps > cap) tpsCaps.set(measurement.modelKey, tps);
+			}
+			measuredMainTokens += tps * (reliable.effectiveMs / 1000);
+			measuredMainMs += reliable.effectiveMs;
+		}
 
-		const theme = ctx.ui.theme;
-		const icon = theme.fg("success", "✓");
-		const tpsLabel = tps > 0
-			? theme.fg("accent", `${tps} t/s`)
-			: theme.fg("success", "N/A");
-		const mainTokens = totalOutputTokens - totalSubagentTokens;
-		const mainMs = totalStreamMs - totalSubagentMs;
-		const mainTps = mainTokens > 0 && mainMs > 0 ? Math.round(mainTokens / (mainMs / 1000)) : 0;
+		const totalTokens = measuredMainTokens + totalSubagentTokens;
+		const totalMs = measuredMainMs + totalSubagentMs;
+		const tps = totalTokens > 0 && totalMs > 0 ? Math.round(totalTokens / (totalMs / 1000)) : 0;
+		const mainTps = measuredMainTokens > 0 && measuredMainMs > 0
+			? Math.round(measuredMainTokens / (measuredMainMs / 1000)) : 0;
 		const subagentTps = totalSubagentTokens > 0 && totalSubagentMs > 0
-			? Math.round(totalSubagentTokens / (totalSubagentMs / 1000))
-			: 0;
+			? Math.round(totalSubagentTokens / (totalSubagentMs / 1000)) : 0;
+		const theme = ctx.ui.theme;
 		const mainDetail = mainTps > 0 ? theme.fg("success", ` | main ${mainTps} t/s`) : "";
 		const subagentDetail = subagentTps > 0
-			? theme.fg("success", ` | subagents ${subagentTps} t/s (${subagentSummaries.length} call${subagentSummaries.length === 1 ? "" : "s"})`)
-			: "";
-
-		ctx.ui.notify(`${icon} total ${tpsLabel}${mainDetail}${subagentDetail}`, "info");
+			? theme.fg("success", ` | subagents ${subagentTps} t/s (${subagentSummaries.length} call${subagentSummaries.length === 1 ? "" : "s"})`) : "";
+		ctx.ui.notify(`${theme.fg("success", "✓")} total ${tps > 0 ? theme.fg("accent", `${tps} t/s`) : theme.fg("success", "N/A")}${mainDetail}${subagentDetail}`, "info");
 		const finalMainStatus = mainTps > 0 ? `main ${mainTps} t/s` : "main N/A";
-		const finalSubagentStatus = totalSubagentTokens > 0 && totalSubagentMs > 0
-			? `subagents ${Math.round(totalSubagentTokens / (totalSubagentMs / 1000))} t/s`
-			: "";
-		ctx.ui.setStatus(
-			"tps",
-			`${theme.fg("accent", `done ${tps} t/s`)} ${theme.fg("success", `(${[finalMainStatus, finalSubagentStatus].filter(Boolean).join(" ")})`)}`,
-		);
+		const finalSubagentStatus = subagentTps > 0 ? `subagents ${subagentTps} t/s` : "";
+		ctx.ui.setStatus("tps", `${theme.fg("accent", `done ${tps} t/s`)} ${theme.fg("success", `(${[finalMainStatus, finalSubagentStatus].filter(Boolean).join(" ")})`)}`);
 	});
 }