From 01564df5be7c906f54fe19350baba99f5d5476b9 Mon Sep 17 00:00:00 2001 From: Tobias Addicks Date: Sun, 17 May 2026 22:44:26 +0200 Subject: [PATCH] Refactor extension structure --- README.md | 8 +- ai-server/stream.ts | 25 ++++++ dark-mechanicus/banner.ts | 49 +++-------- dark-mechanicus/footer.ts | 32 ++++++- dark-mechanicus/index.ts | 10 +-- dark-mechanicus/indicator.ts | 45 ++++------ dark-mechanicus/status-line.ts | 23 +++-- scripts/install-client.sh | 22 ++++- shared/pi-settings.ts | 59 +++++++++++++ shared/token-stats.ts | 153 +++++++++++++++++++++++++++++++++ tests/token-stats.test.ts | 111 ++++++++++++++++++++++++ 11 files changed, 445 insertions(+), 92 deletions(-) create mode 100644 shared/pi-settings.ts create mode 100644 shared/token-stats.ts create mode 100644 tests/token-stats.test.ts diff --git a/README.md b/README.md index 42ec469..5f71510 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,8 @@ Warhammer 40k "Dark Mechanicum" flavoring on top of pi's interactive TUI. | Extension | What it does | |---|---| | [`ai-server/`](ai-server/) | Remote llama.cpp provider over mTLS. Dynamic model discovery. Admin slash commands (load / unload / ctx / preset / restart / refresh). Custom SSE stream implementation with tool calls, reasoning, cache token reporting. See [ai-server/README.md](ai-server/README.md) for the full setup. | -| [`dark-mechanicus/`](dark-mechanicus/) | TUI customization bundle for the dark-mechanicus theme — loaded as one extension via `index.ts`. Includes: `indicator.ts` (working indicator: `⚙ · `, pulsing cog, 45-quote pool), `banner.ts` (cog-and-skull header art), `footer.ts` (custom context+compaction+tok/s footer), `status-line.ts` (third footer line with rotating flavor text), `session-names.ts` (auto `- · ` session names + tab title), `thinking-label.ts` (`Cogitating...` for folded thinking blocks), `markdown-body-color.ts` (forces lavender body text). | +| [`token-stats.ts`](token-stats.ts) + [`shared/token-stats.ts`](shared/token-stats.ts) | Footer owner for context-window + token-rate display. Tracks prefill and generation speed, including reasoning/thinking tokens, and reads `tokenStats.enabled` from `~/.pi/agent/settings.json`. | +| [`dark-mechanicus/`](dark-mechanicus/) | TUI customization bundle for the dark-mechanicus theme — loaded as one extension via `index.ts`. Includes: `indicator.ts` (working indicator: `⚙ · `, pulsing cog, 45-quote pool), `banner.ts` (cog-and-skull header art), `status-line.ts` (third footer line with rotating flavor text), `session-names.ts` (auto `- · ` session names + tab title), `thinking-label.ts` (`Cogitating...` for folded thinking blocks), `markdown-body-color.ts` (forces lavender body text). Display toggles now come from `darkMechanicus` settings instead of slash commands. | | [`local-llama.ts`](local-llama.ts) | Tiny provider registration for a local llama-server (e.g. `127.0.0.1:8088`). | ## Theme @@ -36,7 +37,7 @@ Activate with `/settings` inside pi (or set `"theme": "dark-mechanicus"` in ## Tests -Twenty-eight tests total, no external dependencies. Runs with Node 22+'s +Thirty-two tests total, no external dependencies. Runs with Node 22+'s built-in test runner + `--experimental-strip-types`: ```bash @@ -48,6 +49,7 @@ node --experimental-strip-types --test tests/*.test.ts | `tests/messages.test.ts` | 13 unit tests over `ai-server/messages.ts` — pi Context → OpenAI payload conversion (system prompts, user/assistant/tool-result roles, tool calls, image-only messages). | | `tests/router-utils.test.ts` | 9 unit tests over `ai-server/router-utils.ts` — `extractCtxSize`, `isShardArtefact` (the filter that hides GGUF multi-shard phantoms from the model picker). | | `tests/integration.test.ts` | 6 live-endpoint tests: `/health`, `/models`, model-entry shape, mTLS enforcement, publicly-trusted cert (Let's Encrypt contract), chat completion usage shape including `prompt_tokens_details.cached_tokens`. Auto-skip if the server is unreachable. | +| `tests/token-stats.test.ts` | 4 unit tests over `shared/token-stats.ts` — timing metadata parsing and rate calculation, including thinking-token-aware generation speed. | Stream-parsing edge cases (SSE framing, tool-call splits across chunks, reasoning deltas, abort mid-stream) remain deferred — they need a mock @@ -80,11 +82,11 @@ pi-extensions/ │ ├── admin.ts router HTTP client + SSH helpers │ ├── router-utils.ts pure helpers (test-friendly) │ └── README.md full mTLS + systemd + Caddy setup notes +├── token-stats.ts footer owner for context + token rate display ├── dark-mechanicus/ theme TUI bundle (one extension, multi-file) │ ├── index.ts entry — sequences each module's registrar │ ├── indicator.ts working indicator (cog + quote + timer) │ ├── banner.ts TUI header art -│ ├── footer.ts custom footer layout │ ├── status-line.ts rotating flavor status line │ ├── session-names.ts auto-name generator + tab title │ ├── thinking-label.ts "Cogitating..." for folded thinking blocks diff --git a/ai-server/stream.ts b/ai-server/stream.ts index 68a9f39..04e1d4e 100644 --- a/ai-server/stream.ts +++ b/ai-server/stream.ts @@ -12,6 +12,7 @@ import { createAssistantMessageEventStream, parseStreamingJson, } from "@mariozechner/pi-ai"; +import { captureFirstOutput, finalizePiTokenStats, type PiTokenStats } from "../shared/token-stats.js"; import { AI_SERVER_CHAT_PATH, AI_SERVER_URL, @@ -56,6 +57,9 @@ export function streamAiServer( const stream = createAssistantMessageEventStream(); (async () => { + const tokenStats: PiTokenStats = { + requestStartMs: Date.now(), + }; const output: AssistantMessage = { role: "assistant", content: [], @@ -76,6 +80,7 @@ export function streamAiServer( let currentBlock: CurrentBlock | null = null; let streamEnded = false; + let thinkingTokens: number | undefined; const endWithError = (reason: "error" | "aborted", message: string) => { if (streamEnded) return; @@ -249,6 +254,14 @@ export function streamAiServer( if (data.id && !output.responseId) output.responseId = data.id; if (data.usage) { + const reportedThinkingTokens = + data.usage.completion_tokens_details?.reasoning_tokens; + if ( + typeof reportedThinkingTokens === "number" + && Number.isFinite(reportedThinkingTokens) + ) { + thinkingTokens = reportedThinkingTokens; + } output.usage.input = data.usage.prompt_tokens ?? output.usage.input; output.usage.output = @@ -278,6 +291,7 @@ export function streamAiServer( const reasoning: string | undefined = delta.reasoning_content ?? delta.reasoning; if (reasoning) { + captureFirstOutput(tokenStats, Date.now()); if (!currentBlock || currentBlock.kind !== "thinking") { finishCurrentBlock(); currentBlock = { kind: "thinking", thinking: "" }; @@ -301,6 +315,7 @@ export function streamAiServer( // ── Text ── if (typeof delta.content === "string" && delta.content.length > 0) { + captureFirstOutput(tokenStats, Date.now()); if (!currentBlock || currentBlock.kind !== "text") { finishCurrentBlock(); currentBlock = { kind: "text", text: "" }; @@ -325,6 +340,7 @@ export function streamAiServer( // ── Tool calls ── if (Array.isArray(delta.tool_calls)) { for (const tc of delta.tool_calls) { + captureFirstOutput(tokenStats, Date.now()); const tcId: string | undefined = tc.id; const tcName: string | undefined = tc.function?.name; const tcArgs: string | undefined = tc.function?.arguments; @@ -402,6 +418,15 @@ export function streamAiServer( finishCurrentBlock(); calculateCost(model, output.usage); + (output as AssistantMessage & { piTokenStats?: PiTokenStats }).piTokenStats = finalizePiTokenStats( + tokenStats, + { + input: output.usage.input, + output: output.usage.output, + thinking: thinkingTokens, + }, + Date.now(), + ); if (options?.signal?.aborted) { endWithError("aborted", "Request aborted"); diff --git a/dark-mechanicus/banner.ts b/dark-mechanicus/banner.ts index 0d4edae..661857d 100644 --- a/dark-mechanicus/banner.ts +++ b/dark-mechanicus/banner.ts @@ -1,4 +1,5 @@ import type { ExtensionAPI, Theme } from "@mariozechner/pi-coding-agent"; +import { readExtensionEnumSetting } from "../shared/pi-settings.js"; // ─── Compact cog-and-skull art (17 lines × ~38 cols) ───────────────────── @@ -187,12 +188,17 @@ function renderFull(theme: Theme, width: number): string[] { // ─── Extension entry ───────────────────────────────────────────────────── type Mode = "compact" | "full" | "off"; -let currentMode: Mode = "compact"; +const EXTENSION_NAME = "dark-mechanicus"; +const BANNER_MODES: readonly Mode[] = ["compact", "full", "off"]; + +function readBannerMode(): Mode { + return readExtensionEnumSetting(EXTENSION_NAME, "bannerMode", BANNER_MODES, "full"); +} export default function (pi: ExtensionAPI) { - const apply = (ctx: any, mode: Mode) => { + const apply = (ctx: any) => { if (!ctx.hasUI) return; - currentMode = mode; + const mode = readBannerMode(); if (mode === "off") { ctx.ui.setHeader(undefined); return; @@ -207,41 +213,10 @@ export default function (pi: ExtensionAPI) { }; pi.on("session_start", async (_event, ctx) => { - apply(ctx, "full"); + apply(ctx); }); - pi.registerCommand("mechanicus-banner-compact", { - description: "Compact text-only mechanicus banner (default)", - handler: async (_args, ctx) => { - apply(ctx, "compact"); - ctx.ui.notify("Compact banner active.", "info"); - }, - }); - - pi.registerCommand("mechanicus-banner-full", { - description: "Full banner with the large cog-and-skull art (44 lines tall)", - handler: async (_args, ctx) => { - apply(ctx, "full"); - ctx.ui.notify( - `Full banner active (${ART_HEIGHT} lines). /mechanicus-banner-compact to revert.`, - "info", - ); - }, - }); - - pi.registerCommand("mechanicus-banner-off", { - description: "Restore pi's default header", - handler: async (_args, ctx) => { - apply(ctx, "off"); - ctx.ui.notify("Default header restored.", "info"); - }, - }); - - pi.registerCommand("mechanicus-banner-on", { - description: "Re-install the current mechanicus banner mode", - handler: async (_args, ctx) => { - apply(ctx, currentMode === "off" ? "compact" : currentMode); - ctx.ui.notify(`Banner active (${currentMode === "off" ? "compact" : currentMode}).`, "info"); - }, + pi.on("turn_end", async (_event, ctx) => { + apply(ctx); }); } diff --git a/dark-mechanicus/footer.ts b/dark-mechanicus/footer.ts index e73bb90..9d07869 100644 --- a/dark-mechanicus/footer.ts +++ b/dark-mechanicus/footer.ts @@ -23,6 +23,7 @@ import type { Theme, } from "@mariozechner/pi-coding-agent"; import { formatTokens } from "../shared/format.js"; +import { normalizePiEntries } from "../shared/token-stats.js"; export default function (pi: ExtensionAPI) { let cachedCtx: ExtensionContext | null = null; @@ -54,6 +55,29 @@ export default function (pi: ExtensionAPI) { width: number, ): string[] => { const dim = (s: string) => theme.fg("dim", s); + const readExtensionStatuses = (): Array<[string, string]> => { + const statuses = footerData?.getExtensionStatuses?.(); + if (!statuses) { + return []; + } + if (Array.isArray(statuses)) { + return statuses as Array<[string, string]>; + } + if (typeof statuses === "object") { + const candidate = statuses as { + [Symbol.iterator]?: unknown; + }; + if (typeof candidate[Symbol.iterator] === "function") { + try { + return Array.from(candidate as Iterable<[string, string]>); + } catch { + return []; + } + } + return Object.entries(statuses as Record); + } + return []; + }; // ── Walk session entries for token totals + compactions ─ let totalIn = 0; @@ -62,7 +86,7 @@ export default function (pi: ExtensionAPI) { let totalCost = 0; let compactions = 0; let lastAssistantOutput = 0; - const entries = c.sessionManager.getEntries(); + const entries = normalizePiEntries(c.sessionManager.getEntries()); for (const entry of entries) { if (entry.type === "compaction") compactions++; if ( @@ -165,9 +189,9 @@ export default function (pi: ExtensionAPI) { const lines = [pwdLine, statsLine]; // ── Line 3+: extension statuses (alphabetical) ────────── - const extStatuses = footerData.getExtensionStatuses(); - if (extStatuses.size > 0) { - const sorted = Array.from(extStatuses.entries()).sort(([a], [b]) => + const extStatuses = readExtensionStatuses(); + if (extStatuses.length > 0) { + const sorted = extStatuses.sort(([a], [b]) => a.localeCompare(b), ); const statusLine = sorted.map(([, text]) => text).join(" "); diff --git a/dark-mechanicus/index.ts b/dark-mechanicus/index.ts index e6e67e7..bb94f36 100644 --- a/dark-mechanicus/index.ts +++ b/dark-mechanicus/index.ts @@ -1,18 +1,15 @@ /** * dark-mechanicus — bundle of TUI customization extensions for the * "dark-mechanicus" theme. Loaded as a single pi extension because the - * pieces are designed to compose (banner + status line + footer + indicator + * pieces are designed to compose (banner + status line + indicator * + thinking label all share visual language with themes/dark-mechanicus.json). * - * Each module is self-contained — registers its own commands and event - * listeners. This index only sequences their initialization. Disabling any - * single module is done via its own `/-off` command, not by editing - * this file. + * Display toggles are configured via ~/.pi/agent/settings.json under + * extensions["dark-mechanicus"]. This index only sequences initialization. */ import type { ExtensionAPI } from "@mariozechner/pi-coding-agent"; import banner from "./banner.js"; -import footer from "./footer.js"; import indicator from "./indicator.js"; import markdownBodyColor from "./markdown-body-color.js"; import sessionNames from "./session-names.js"; @@ -21,7 +18,6 @@ import thinkingLabel from "./thinking-label.js"; export default function (pi: ExtensionAPI) { banner(pi); - footer(pi); indicator(pi); markdownBodyColor(pi); sessionNames(pi); diff --git a/dark-mechanicus/indicator.ts b/dark-mechanicus/indicator.ts index b1a602d..69aad91 100644 --- a/dark-mechanicus/indicator.ts +++ b/dark-mechanicus/indicator.ts @@ -14,14 +14,14 @@ * Fisher-Yates shuffle on every turn_start so the order differs per * assistant response. * - * Commands: - * /dark-mechanicus-indicator-off Restore pi's default spinner - * /dark-mechanicus-indicator-on Re-enable the mechanicum indicator + * Controlled via ~/.pi/agent/settings.json under + * extensions["dark-mechanicus"].indicatorEnabled. */ import type { ExtensionAPI } from "@mariozechner/pi-coding-agent"; import { safely } from "../shared/ctx.js"; import { formatElapsed } from "../shared/format.js"; +import { readExtensionBooleanSetting } from "../shared/pi-settings.js"; // ─── Quote pools (15 per category) ─────────────────────────────────────── @@ -118,10 +118,13 @@ function buildFrames(quote: string, elapsedMs: number): string[] { // ─── Extension ─────────────────────────────────────────────────────────── -type Mode = "on" | "off"; +const EXTENSION_NAME = "dark-mechanicus"; + +function readIndicatorEnabled(): boolean { + return readExtensionBooleanSetting(EXTENSION_NAME, "indicatorEnabled", true); +} export default function (pi: ExtensionAPI) { - let mode: Mode = "on"; let shuffled = shuffle(ALL_QUOTES); let startedAt: number | null = null; let tickerHandle: NodeJS.Timeout | null = null; @@ -146,7 +149,7 @@ export default function (pi: ExtensionAPI) { }; const applyFrame = (ctx: any) => { - if (!ctx.hasUI || mode === "off" || startedAt === null) return; + if (!ctx.hasUI || !readIndicatorEnabled() || startedAt === null) return; // Guard against stale-ctx: the setInterval ticker can fire between // session_shutdown being queued and our turn_end/session_shutdown // handlers running. Any ctx.ui.* access on a stale runner throws. @@ -181,12 +184,18 @@ export default function (pi: ExtensionAPI) { // ── Seed at session start so pi's default spinner never appears ── pi.on("session_start", async (_event, ctx) => { + if (!readIndicatorEnabled()) return; suppressWorkingMessage(ctx); }); // ── Turn lifecycle ─────────────────────────────────────────────── pi.on("turn_start", async (_event, ctx) => { - if (mode === "off") return; + if (!readIndicatorEnabled()) { + stopTicker(); + startedAt = null; + safely(() => ctx.ui.setWorkingIndicator(undefined)); + return; + } suppressWorkingMessage(ctx); startedAt = Date.now(); shuffled = shuffle(ALL_QUOTES); @@ -204,26 +213,4 @@ export default function (pi: ExtensionAPI) { stopTicker(); startedAt = null; }); - - // ── Commands ───────────────────────────────────────────────────── - pi.registerCommand("dark-mechanicus-indicator-off", { - description: "Restore pi's default spinner and skip cogitation summaries", - handler: async (_args, ctx) => { - mode = "off"; - stopTicker(); - ctx.ui.setWorkingIndicator(undefined); - ctx.ui.notify("Default spinner restored.", "info"); - }, - }); - - pi.registerCommand("dark-mechanicus-indicator-on", { - description: "Re-enable the pulsing cog indicator + cogitation summary", - handler: async (_args, ctx) => { - mode = "on"; - ctx.ui.notify( - "Dark mechanicum indicator re-armed. Takes effect on the next turn.", - "info", - ); - }, - }); } diff --git a/dark-mechanicus/status-line.ts b/dark-mechanicus/status-line.ts index 2a52bcc..faf1a64 100644 --- a/dark-mechanicus/status-line.ts +++ b/dark-mechanicus/status-line.ts @@ -12,6 +12,7 @@ */ import type { ExtensionAPI } from "@mariozechner/pi-coding-agent"; +import { readExtensionBooleanSetting } from "../shared/pi-settings.js"; const STATUSES = [ "HERETEK FORGE · ACTIVE", @@ -35,9 +36,19 @@ function pick(arr: readonly T[]): T { return arr[Math.floor(Math.random() * arr.length)]!; } +const EXTENSION_NAME = "dark-mechanicus"; + +function readStatusLineEnabled(): boolean { + return readExtensionBooleanSetting(EXTENSION_NAME, "statusLineEnabled", true); +} + export default function (pi: ExtensionAPI) { const applyStatus = (ctx: any, body?: string) => { if (!ctx.hasUI) return; + if (!readStatusLineEnabled()) { + ctx.ui.setStatus("mechanicus", undefined); + return; + } const text = body ?? pick(STATUSES); const theme = ctx.ui.theme; const line = `${theme.fg("warning", "⚙")} ${theme.fg("muted", text)}`; @@ -57,16 +68,12 @@ export default function (pi: ExtensionAPI) { pi.registerCommand("mechanicus-status-cycle", { description: "Force a new mechanicum status line immediately", handler: async (_args, ctx) => { + if (!readStatusLineEnabled()) { + ctx.ui.notify("mechanicus status line is disabled in settings.", "info"); + return; + } applyStatus(ctx); ctx.ui.notify("Status re-transmuted.", "info"); }, }); - - pi.registerCommand("mechanicus-status-off", { - description: "Remove the mechanicum status line from the footer", - handler: async (_args, ctx) => { - ctx.ui.setStatus("mechanicus", undefined); - ctx.ui.notify("Status line cleared.", "info"); - }, - }); } diff --git a/scripts/install-client.sh b/scripts/install-client.sh index 32ce1bc..35113de 100755 --- a/scripts/install-client.sh +++ b/scripts/install-client.sh @@ -15,7 +15,8 @@ # certs mTLS client cert + root CA from the Caddy host # ssh SSH key-auth to the ai-server host (admin commands) # ai-server the pi extension (mTLS llama.cpp router provider) -# dark-mechanicus theme TUI bundle (banner, footer, indicator, etc.) +# dark-mechanicus theme TUI bundle (banner, indicator, status line, etc.) +# token-stats footer owner for context + token rate display # themes JSON theme palette files (themes/*.json) # local-llama tiny stub provider for a local llama-server # ai-complete shell CLI for direct llama-server access (installed to @@ -63,7 +64,7 @@ SKIP_VERIFY=0 LEGACY_NO_CERTS=0 LEGACY_NO_SSH=0 -ALL_COMPONENTS=(certs ssh ai-server dark-mechanicus themes local-llama ai-complete shared) +ALL_COMPONENTS=(certs ssh ai-server dark-mechanicus token-stats themes local-llama ai-complete shared) usage() { sed -n '2,/^$/p' "$0" | sed 's/^#\{0,1\} \{0,1\}//'; exit 0; } @@ -110,9 +111,11 @@ if [[ $LEGACY_NO_SSH -eq 1 ]]; then COMPONENTS=$(echo "$COMPONENTS" | tr ',' '\n' | grep -vx 'ssh' | paste -sd, -) fi -# Auto-add 'shared' if dark-mechanicus is selected (it imports from ../shared/) -if [[ ",$COMPONENTS," == *",dark-mechanicus,"* && ",$COMPONENTS," != *",shared,"* ]]; then +# Auto-add 'shared' if dark-mechanicus or token-stats is selected (they import from ../shared/ or ./shared/) +if [[ ",$COMPONENTS," == *",dark-mechanicus,"* || ",$COMPONENTS," == *",token-stats,"* ]]; then + if [[ ",$COMPONENTS," != *",shared,"* ]]; then COMPONENTS="$COMPONENTS,shared" + fi fi # Normalize and dedupe @@ -182,6 +185,17 @@ if want dark-mechanicus; then else echo " (dark-mechanicus/ not in repo, skipping)" fi + + if want token-stats; then + if [[ -f token-stats.ts ]]; then + echo; echo "==> [token-stats] syncing footer extension" + if [[ -f "$PI_DIR/extensions/token-stats.ts" && $UPDATE_ONLY -eq 1 ]]; then + echo " keep existing token-stats.ts (--update-only)" + else + rsync "${RSYNC_FLAGS[@]}" token-stats.ts "$PI_DIR/extensions/" + fi + fi + fi fi if want shared; then diff --git a/shared/pi-settings.ts b/shared/pi-settings.ts new file mode 100644 index 0000000..eadab9d --- /dev/null +++ b/shared/pi-settings.ts @@ -0,0 +1,59 @@ +import { existsSync, readFileSync } from "node:fs"; +import { join } from "node:path"; + +const HOME = process.env.HOME ?? process.env.USERPROFILE ?? ""; +const PI_SETTINGS_PATH = join(HOME, ".pi", "agent", "settings.json"); + +function isPlainObject(value: unknown): value is Record { + return !!value && typeof value === "object" && !Array.isArray(value); +} + +function toCamelCase(value: string): string { + return value.replace(/-([a-z])/g, (_match, letter: string) => letter.toUpperCase()); +} + +function readPiSettings(): Record | null { + if (!existsSync(PI_SETTINGS_PATH)) { + return null; + } + try { + return JSON.parse(readFileSync(PI_SETTINGS_PATH, "utf8")) as Record; + } catch { + return null; + } +} + +export function readExtensionSettings(extensionName: string): Record | null { + const settings = readPiSettings(); + if (!settings) { + return null; + } + const config = settings[extensionName] ?? settings[toCamelCase(extensionName)]; + return isPlainObject(config) ? config : null; +} + +export function readExtensionBooleanSetting( + extensionName: string, + key: string, + defaultValue: boolean, +): boolean { + const config = readExtensionSettings(extensionName); + if (!config || typeof config[key] !== "boolean") { + return defaultValue; + } + return config[key] as boolean; +} + +export function readExtensionEnumSetting( + extensionName: string, + key: string, + allowedValues: readonly T[], + defaultValue: T, +): T { + const config = readExtensionSettings(extensionName); + const value = config?.[key]; + if (typeof value !== "string") { + return defaultValue; + } + return (allowedValues as readonly string[]).includes(value) ? (value as T) : defaultValue; +} diff --git a/shared/token-stats.ts b/shared/token-stats.ts new file mode 100644 index 0000000..09b721d --- /dev/null +++ b/shared/token-stats.ts @@ -0,0 +1,153 @@ +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; + 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); + } 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; +} diff --git a/tests/token-stats.test.ts b/tests/token-stats.test.ts new file mode 100644 index 0000000..30d4a62 --- /dev/null +++ b/tests/token-stats.test.ts @@ -0,0 +1,111 @@ +import assert from "node:assert/strict"; +import { test } from "node:test"; +import { + findLatestPiTokenStats, + formatPiTokenRateStatus, + normalizePiEntries, + readPiTokenStats, + summarizePiTokenStats, +} from "../shared/token-stats.ts"; + +test("summarizePiTokenStats reports separate processing and generation rates", () => { + const summary = summarizePiTokenStats({ + requestStartMs: 1000, + firstOutputMs: 3000, + responseEndMs: 7000, + inputTokens: 800, + outputTokens: 200, + thinkingTokens: 300, + }); + assert.equal(summary.processingTokensPerSecond, 400); + assert.equal(summary.generationTokensPerSecond, 125); + assert.equal(formatPiTokenRateStatus(summary), "P 400.0tok/s G 125.0tok/s"); +}); + +test("summarizePiTokenStats falls back to response end when no output token arrived", () => { + const summary = summarizePiTokenStats({ + requestStartMs: 1000, + responseEndMs: 5000, + inputTokens: 200, + outputTokens: 0, + }); + assert.equal(summary.processingTokensPerSecond, 50); + assert.equal(summary.generationTokensPerSecond, null); + assert.equal(formatPiTokenRateStatus(summary), "P 50.0tok/s"); +}); + +test("readPiTokenStats rejects missing timing metadata", () => { + assert.equal(readPiTokenStats({ role: "assistant" }), null); + assert.deepEqual( + readPiTokenStats({ + piTokenStats: { + requestStartMs: 10, + firstOutputMs: 20, + responseEndMs: 30, + inputTokens: 40, + outputTokens: 50, + thinkingTokens: 60, + }, + }), + { + requestStartMs: 10, + firstOutputMs: 20, + responseEndMs: 30, + inputTokens: 40, + outputTokens: 50, + thinkingTokens: 60, + }, + ); +}); + +test("findLatestPiTokenStats ignores older assistant messages after the latest assistant turn", () => { + assert.deepEqual( + findLatestPiTokenStats([ + { + type: "message", + message: { + role: "assistant", + piTokenStats: { + requestStartMs: 1, + responseEndMs: 2, + outputTokens: 3, + }, + }, + }, + { + type: "message", + message: { + role: "user", + }, + }, + { + type: "message", + message: { + role: "assistant", + piTokenStats: { + requestStartMs: 10, + firstOutputMs: 20, + responseEndMs: 30, + inputTokens: 40, + outputTokens: 50, + thinkingTokens: 60, + }, + }, + }, + ]), + { + requestStartMs: 10, + firstOutputMs: 20, + responseEndMs: 30, + inputTokens: 40, + outputTokens: 50, + thinkingTokens: 60, + }, + ); +}); + +test("findLatestPiTokenStats treats non-array entries as empty", () => { + assert.equal(findLatestPiTokenStats(undefined), null); + assert.equal(findLatestPiTokenStats({}), null); + assert.deepEqual(normalizePiEntries(undefined), []); +});