395 lines
15 KiB
TypeScript
395 lines
15 KiB
TypeScript
/**
|
|
* Pi status bar (custom footer)
|
|
*
|
|
* Replaces pi's default footer via `ctx.ui.setFooter()`. It mirrors the layout
|
|
* and *meaning* of the Claude Code status line (`~/.config/claude/statusline.py`),
|
|
* but swaps the emoji for Nerd Font glyphs — the same Material Design set already
|
|
* used across tmux (waybar cpu/mem/net glyphs) and nvim — so it reads natively in
|
|
* kitty (CaskaydiaCove Nerd Font Mono). Colors come from the active pi theme
|
|
* (catppuccin-mocha), not hard-coded ANSI.
|
|
*
|
|
* Segments render left → right and truncate at the terminal edge, so the most
|
|
* useful info stays visible on a narrow pane. Each icon echoes the Claude emoji
|
|
* it replaces:
|
|
*
|
|
* dir cwd, ~-collapsed (Claude 📁) U+F024B nf-md-folder
|
|
* git current branch (only inside a repo) (Claude 🌿) U+F062C nf-md-source_branch
|
|
* model model id + · thinking level (Claude 🤖) U+F06A9 nf-md-robot
|
|
* context % of context window used (+ tokens) (Claude 📝) U+F021A nf-md-text_box
|
|
* t/s last response's decode throughput U+F04C5 nf-md-speedometer
|
|
* (llama.cpp's own server-side number + dim ` pp N` prompt-processing
|
|
* speed when the final SSE chunk carries `timings`; otherwise the
|
|
* client-side clock -- see "server-side timings" below)
|
|
* ↑ ↓ tokens session prompt / output tokens (Claude 🔢) plain arrows
|
|
* (prompt = input + cache read + cache write: everything the server
|
|
* processed, matching the Claude bar — llama.cpp reports a cached
|
|
* prefix as cacheRead, so bare usage.input is a few dozen tokens)
|
|
*
|
|
* git, context, t/s and tokens only appear once there is data for them (i.e. after
|
|
* the first response, and git only inside a git repo). Extension statuses (e.g.
|
|
* plan-mode's "⏸ plan") are preserved at the far left. The Claude bar's system row
|
|
* (RAM/CPU/temp/disk) is intentionally omitted — that data isn't in the footer API,
|
|
* and the tmux bar below pi already shows cpu/mem/net.
|
|
*
|
|
* Two extras beyond the footer line:
|
|
* - a themed pulse working-indicator (the spinner shown while pi streams), and
|
|
* - a context-budget warning widget above the editor once the window passes
|
|
* 80% full, nudging a /compact before things get truncated.
|
|
*
|
|
* Restore the built-in footer (and reset the indicator/widget) with `/statusbar`.
|
|
*/
|
|
|
|
import type { AssistantMessage } from "@earendil-works/pi-ai";
|
|
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
|
|
import { truncateToWidth } from "@earendil-works/pi-tui";
|
|
|
|
// Nerd Font (Material Design) glyphs — verified to render in the user's font.
|
|
const ICON = {
|
|
dir: "\u{F024B}", // nf-md-folder
|
|
git: "\u{F062C}", // nf-md-source_branch
|
|
model: "\u{F06A9}", // nf-md-robot
|
|
context: "\u{F021A}", // nf-md-text_box
|
|
tps: "\u{F04C5}", // nf-md-speedometer
|
|
warn: "\u{F002A}", // nf-md-alert_outline (same glyph nvim uses for warnings)
|
|
up: "↑", // ↑
|
|
down: "↓", // ↓
|
|
};
|
|
|
|
// Widget key for the context-budget warning shown above the editor.
|
|
const CTX_WIDGET = "statusbar-context-warning";
|
|
// Show the warning once the context window is this full (%).
|
|
const CTX_WARN_AT = 80;
|
|
|
|
// Current thinking level, tracked via thinking_level_select (footer render has no
|
|
// direct access to it). Defaults to the settings.json default of "off".
|
|
let thinkingLevel = "off";
|
|
|
|
// Decode throughput (tokens/sec) of the most recent assistant response. Preferred
|
|
// source: the server's own timings (below). Fallback: a client-side clock from the
|
|
// first streamed token to message end, so prompt-eval time is excluded.
|
|
let lastTps = 0;
|
|
let lastPp = 0; // prompt-processing t/s, server-side only (0 = unknown)
|
|
let genFirstTokenAt = 0;
|
|
let sawToken = false;
|
|
|
|
// --- server-side timings (llama.cpp) ----------------------------------------
|
|
// llama.cpp puts a `timings` object on the final SSE chunk (the one carrying `usage`):
|
|
// prompt_n / prompt_ms / prompt_per_second, predicted_n / predicted_ms /
|
|
// predicted_per_second, cache_n, draft_n / draft_n_accepted. pi-ai's openai-completions
|
|
// provider reads only `usage` and drops the rest, and the extension API has no HTTP
|
|
// hook, so the one way to see it is to wrap globalThis.fetch and tee the SSE body.
|
|
// That is sanctioned: pi's http-dispatcher installs undici's fetch at startup and
|
|
// deliberately keeps a later override ("If a caller replaced fetch after module load,
|
|
// preserve that deliberate override"), and pi-ai builds `new OpenAI({ fetch: undefined })`
|
|
// per request, which falls back to globalThis.fetch -- so a wrapper installed at
|
|
// session_start sees every chat completion. Only POST */chat/completions responses
|
|
// with a text/event-stream body are teed; other providers' streams carry no `timings`
|
|
// and are passed through untouched apart from the tee. The tap stays installed for the
|
|
// life of the process (`/statusbar` off just stops displaying it; `/reload` rebinds the
|
|
// sink instead of wrapping twice).
|
|
interface ServerTimings {
|
|
tps: number; // predicted_per_second
|
|
pp: number; // prompt_per_second
|
|
predictedN: number;
|
|
promptN: number;
|
|
cacheN: number;
|
|
}
|
|
// Latest timings seen on the wire, consumed by message_end. Paired with the message by
|
|
// predicted_n === usage.output, so a compaction/summary call's timings never get
|
|
// attributed to the next visible response.
|
|
let pendingServerTimings: ServerTimings | undefined;
|
|
|
|
const FETCH_TAP_KEY = "__piStatusbarFetchTap";
|
|
|
|
/** Read one branch of a teed SSE body and report every `timings` object it carries. */
|
|
async function scanSse(stream: ReadableStream<Uint8Array>, onTimings: (t: ServerTimings) => void): Promise<void> {
|
|
const reader = stream.getReader();
|
|
const dec = new TextDecoder();
|
|
let buf = "";
|
|
try {
|
|
for (;;) {
|
|
const { value, done } = await reader.read();
|
|
if (done) break;
|
|
buf += dec.decode(value, { stream: true });
|
|
let nl: number;
|
|
while ((nl = buf.indexOf("\n")) >= 0) {
|
|
const line = buf.slice(0, nl).trim();
|
|
buf = buf.slice(nl + 1);
|
|
if (!line.startsWith("data:")) continue;
|
|
const payload = line.slice(5).trim();
|
|
if (!payload.includes('"timings"')) continue; // cheap pre-filter: parse only the final chunk
|
|
try {
|
|
const t = JSON.parse(payload)?.timings;
|
|
if (t && typeof t.predicted_per_second === "number") {
|
|
onTimings({
|
|
tps: t.predicted_per_second,
|
|
pp: typeof t.prompt_per_second === "number" ? t.prompt_per_second : 0,
|
|
predictedN: t.predicted_n ?? 0,
|
|
promptN: t.prompt_n ?? 0,
|
|
cacheN: t.cache_n ?? 0,
|
|
});
|
|
}
|
|
} catch {
|
|
// not JSON: ignore
|
|
}
|
|
}
|
|
}
|
|
} catch {
|
|
// aborted / cancelled stream: the primary branch reports that, we just stop
|
|
}
|
|
}
|
|
|
|
/** Wrap globalThis.fetch once; on /reload only rebind the sink. */
|
|
function installFetchTap(): void {
|
|
const g = globalThis as unknown as Record<string, unknown>;
|
|
const sink = (t: ServerTimings) => {
|
|
pendingServerTimings = t;
|
|
};
|
|
const existing = g[FETCH_TAP_KEY] as { sink: (t: ServerTimings) => void } | undefined;
|
|
if (existing) {
|
|
existing.sink = sink;
|
|
return;
|
|
}
|
|
const tap = { sink };
|
|
g[FETCH_TAP_KEY] = tap;
|
|
const orig = globalThis.fetch;
|
|
globalThis.fetch = (async (input: Parameters<typeof fetch>[0], init?: Parameters<typeof fetch>[1]) => {
|
|
const res = await orig(input, init);
|
|
let url = "";
|
|
try {
|
|
url = typeof input === "string" ? input : input instanceof URL ? input.href : ((input as Request)?.url ?? "");
|
|
} catch {
|
|
return res;
|
|
}
|
|
if (!url.endsWith("/chat/completions") || !res.body) return res;
|
|
if (!(res.headers.get("content-type") ?? "").includes("text/event-stream")) return res;
|
|
const [ours, theirs] = res.body.tee();
|
|
void scanSse(ours, (t) => tap.sink(t));
|
|
return new Response(theirs, res);
|
|
}) as typeof fetch;
|
|
}
|
|
|
|
// Set from inside the footer factory so events can trigger a re-render.
|
|
let requestRender: (() => void) | undefined;
|
|
|
|
function collapseHome(p: string): string {
|
|
const home = process.env.HOME || "";
|
|
if (p === home) return "~";
|
|
if (home && p.startsWith(`${home}/`)) return `~${p.slice(home.length)}`;
|
|
return p;
|
|
}
|
|
|
|
/** 1234 -> 1.2k, 843000 -> 843k, 1234567 -> 1.2M (same boundaries as the Claude bar's
|
|
* fmt_tokens: 0.9995 / 9.95 so a value never prints as "10.0k" or "1000k"). */
|
|
function fmtTokens(n: number): string {
|
|
for (const [unit, div] of [
|
|
["M", 1e6],
|
|
["k", 1e3],
|
|
] as const) {
|
|
const v = n / div;
|
|
if (v < 0.9995) continue;
|
|
return v >= 9.95 ? `${v.toFixed(0)}${unit}` : `${v.toFixed(1)}${unit}`;
|
|
}
|
|
return `${Math.round(n)}`;
|
|
}
|
|
|
|
/** Cumulative session tokens, plus the last request's prompt size (~ current
|
|
* context-window fill).
|
|
*
|
|
* Prompt = input + cacheRead + cacheWrite. pi-ai splits a cached prefix out of
|
|
* `usage.input` (llama.cpp too: the second call of a session reports e.g. input 20,
|
|
* cacheRead 2417), so summing bare `input` undercounted ↑ and -- worse -- made the
|
|
* context % and the ≥80% warning read near zero on every cached call. */
|
|
function tokenStats(ctx: ExtensionContext): {
|
|
input: number;
|
|
output: number;
|
|
lastInput: number;
|
|
} {
|
|
let input = 0;
|
|
let output = 0;
|
|
let lastInput = 0;
|
|
for (const entry of ctx.sessionManager.getBranch()) {
|
|
if (entry.type === "message" && entry.message.role === "assistant") {
|
|
const usage = (entry.message as AssistantMessage).usage;
|
|
if (!usage) continue; // aborted stream: no usage — a throwing render kills the TUI
|
|
const prompt = (usage.input ?? 0) + (usage.cacheRead ?? 0) + (usage.cacheWrite ?? 0);
|
|
input += prompt;
|
|
output += usage.output ?? 0;
|
|
lastInput = prompt;
|
|
}
|
|
}
|
|
return { input, output, lastInput };
|
|
}
|
|
|
|
/** Show/hide the above-editor warning when the context window is nearly full. */
|
|
function updateContextWidget(ctx: ExtensionContext): void {
|
|
if (ctx.mode !== "tui") return;
|
|
const contextWindow = ctx.model?.contextWindow;
|
|
const { lastInput } = tokenStats(ctx);
|
|
if (!contextWindow || lastInput <= 0) {
|
|
ctx.ui.setWidget(CTX_WIDGET, undefined);
|
|
return;
|
|
}
|
|
const pct = (lastInput / contextWindow) * 100;
|
|
if (pct < CTX_WARN_AT) {
|
|
ctx.ui.setWidget(CTX_WIDGET, undefined);
|
|
return;
|
|
}
|
|
const th = ctx.ui.theme;
|
|
const col = pct >= 90 ? "error" : "warning";
|
|
const line =
|
|
th.fg(col, `${ICON.warn} context ${pct.toFixed(0)}%`) +
|
|
th.fg("dim", ` (${fmtTokens(lastInput)}/${fmtTokens(contextWindow)}) — /compact soon`);
|
|
ctx.ui.setWidget(CTX_WIDGET, [line]);
|
|
}
|
|
|
|
/** A gentle catppuccin "breathing" pulse for the streaming working-indicator. */
|
|
function pulseIndicator(ctx: ExtensionContext) {
|
|
const th = ctx.ui.theme;
|
|
return {
|
|
frames: [th.fg("dim", "·"), th.fg("muted", "•"), th.fg("accent", "●"), th.fg("muted", "•")],
|
|
intervalMs: 120,
|
|
};
|
|
}
|
|
|
|
export default function (pi: ExtensionAPI) {
|
|
let enabled = true;
|
|
|
|
pi.on("thinking_level_select", async (event) => {
|
|
thinkingLevel = event.level;
|
|
requestRender?.();
|
|
});
|
|
|
|
pi.on("model_select", async () => {
|
|
requestRender?.();
|
|
});
|
|
|
|
// --- t/s timing: bracket each assistant response and clock its generation ---
|
|
pi.on("message_start", async (event) => {
|
|
if (event.message?.role !== "assistant") return;
|
|
sawToken = false;
|
|
genFirstTokenAt = 0;
|
|
});
|
|
|
|
pi.on("message_update", async (event) => {
|
|
if (event.message?.role !== "assistant") return;
|
|
if (!sawToken) {
|
|
sawToken = true;
|
|
genFirstTokenAt = Date.now();
|
|
}
|
|
});
|
|
|
|
pi.on("message_end", async (event, ctx) => {
|
|
if (event.message?.role !== "assistant") return;
|
|
const output = (event.message as AssistantMessage).usage?.output ?? 0;
|
|
// The tee reader and pi-ai consume the same final chunk; pi-ai's path to this
|
|
// event is longer, but give the reader one macrotask in case it is still behind.
|
|
if (!pendingServerTimings) await new Promise((r) => setTimeout(r, 20));
|
|
const srv = pendingServerTimings;
|
|
pendingServerTimings = undefined;
|
|
if (srv && srv.predictedN > 0 && (output === 0 || srv.predictedN === output)) {
|
|
lastTps = srv.tps;
|
|
lastPp = srv.pp;
|
|
} else {
|
|
const secs = genFirstTokenAt ? (Date.now() - genFirstTokenAt) / 1000 : 0;
|
|
if (secs > 0 && output > 0) lastTps = output / secs;
|
|
lastPp = 0;
|
|
}
|
|
requestRender?.();
|
|
if (enabled) updateContextWidget(ctx);
|
|
});
|
|
|
|
const install = (ctx: ExtensionContext) => {
|
|
if (ctx.mode !== "tui") return;
|
|
ctx.ui.setWorkingIndicator(pulseIndicator(ctx));
|
|
updateContextWidget(ctx);
|
|
ctx.ui.setFooter((tui, theme, footerData) => {
|
|
const sep = theme.fg("dim", " │ ");
|
|
requestRender = () => tui.requestRender();
|
|
const unsub = footerData.onBranchChange(() => tui.requestRender());
|
|
return {
|
|
dispose: unsub,
|
|
invalidate() {},
|
|
render(width: number): string[] {
|
|
const parts: string[] = [];
|
|
|
|
// Preserve extension statuses (plan-mode "⏸ plan", etc.) up front.
|
|
for (const status of footerData.getExtensionStatuses().values()) {
|
|
if (status) parts.push(status);
|
|
}
|
|
|
|
// directory
|
|
parts.push(theme.fg("accent", ICON.dir) + " " + theme.fg("text", collapseHome(ctx.cwd)));
|
|
|
|
// git branch (only inside a repo)
|
|
const branch = footerData.getGitBranch();
|
|
if (branch) parts.push(theme.fg("accent", ICON.git) + " " + theme.fg("success", branch));
|
|
|
|
// model (+ · thinking level)
|
|
if (ctx.model?.id) {
|
|
let model = theme.fg("accent", ICON.model) + " " + theme.fg("text", ctx.model.id);
|
|
if (thinkingLevel && thinkingLevel !== "off") {
|
|
model += theme.fg("dim", ` · ${thinkingLevel}`);
|
|
}
|
|
parts.push(model);
|
|
}
|
|
|
|
// context window usage (% when the window is known, else raw tokens)
|
|
const { input, output, lastInput } = tokenStats(ctx);
|
|
if (lastInput > 0) {
|
|
const contextWindow = ctx.model?.contextWindow;
|
|
let seg = theme.fg("accent", ICON.context) + " ";
|
|
if (contextWindow) {
|
|
const pct = Math.min(100, (lastInput / contextWindow) * 100);
|
|
const col = pct < 50 ? "success" : pct < 80 ? "warning" : "error";
|
|
seg += theme.fg(col, `${pct.toFixed(0)}%`) + theme.fg("dim", ` (${fmtTokens(lastInput)})`);
|
|
} else {
|
|
seg += theme.fg("text", fmtTokens(lastInput));
|
|
}
|
|
parts.push(seg);
|
|
}
|
|
|
|
// tokens/sec of the last response (+ dim prompt-processing t/s when the
|
|
// server reported it)
|
|
if (lastTps > 0) {
|
|
let seg = theme.fg("accent", ICON.tps) + " " + theme.fg("text", lastTps.toFixed(1)) + theme.fg("dim", " t/s");
|
|
if (lastPp > 0) seg += theme.fg("dim", ` pp ${lastPp.toFixed(0)}`);
|
|
parts.push(seg);
|
|
}
|
|
|
|
// ↑ ↓ cumulative session prompt / output tokens (Claude's 🔢 segment)
|
|
if (input || output) {
|
|
parts.push(theme.fg("dim", `${ICON.up}${fmtTokens(input)} ${ICON.down}${fmtTokens(output)}`));
|
|
}
|
|
|
|
return [truncateToWidth(parts.join(sep), width)];
|
|
},
|
|
};
|
|
});
|
|
};
|
|
|
|
pi.on("session_start", async (_event, ctx) => {
|
|
// Installed here, not at load time: pi's http-dispatcher only swaps in undici's
|
|
// fetch (idle timeout, proxy) while globalThis.fetch is still the original, so the
|
|
// tap must not get there first. By session_start, startup is done.
|
|
installFetchTap();
|
|
if (enabled) install(ctx);
|
|
});
|
|
|
|
pi.registerCommand("statusbar", {
|
|
description: "Toggle the custom pi status bar footer",
|
|
handler: async (_args, ctx) => {
|
|
enabled = !enabled;
|
|
if (enabled) {
|
|
install(ctx);
|
|
ctx.ui.notify("Custom status bar enabled", "info");
|
|
} else {
|
|
ctx.ui.setFooter(undefined);
|
|
ctx.ui.setWorkingIndicator(); // restore pi's default spinner
|
|
ctx.ui.setWidget(CTX_WIDGET, undefined);
|
|
ctx.ui.notify("Default footer restored", "info");
|
|
}
|
|
},
|
|
});
|
|
}
|