Refactor extension structure
This commit is contained in:
Tobias Addicks
@@ -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: `⚙ <quote> · <elapsed>`, 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 `<adj>-<noun> · <NNN>` 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: `⚙ <quote> · <elapsed>`, 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 `<adj>-<noun> · <NNN>` 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
|
||||
|
||||
@@ -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");
|
||||
|
||||
+12
-37
@@ -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);
|
||||
});
|
||||
}
|
||||
|
||||
@@ -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<string, string>);
|
||||
}
|
||||
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(" ");
|
||||
|
||||
@@ -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 `/<name>-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);
|
||||
|
||||
@@ -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",
|
||||
);
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
@@ -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<T>(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");
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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<string, unknown> {
|
||||
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<string, unknown> | null {
|
||||
if (!existsSync(PI_SETTINGS_PATH)) {
|
||||
return null;
|
||||
}
|
||||
try {
|
||||
return JSON.parse(readFileSync(PI_SETTINGS_PATH, "utf8")) as Record<string, unknown>;
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
export function readExtensionSettings(extensionName: string): Record<string, unknown> | 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<T extends string>(
|
||||
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;
|
||||
}
|
||||
@@ -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<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;
|
||||
}
|
||||
@@ -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), []);
|
||||
});
|
||||
Reference in New Issue
Block a user