import type { ContextUsage, ExtensionAPI, ExtensionContext, ThemeColor, } from "@earendil-works/pi-coding-agent"; import { closeDebug, dbg } from "./src/debug.ts"; import { offerFooterSegment, removeFooterSegment, } from "./src/pi-ext-footer-segment.ts"; import { parseNumberSetting, resolveSetting, type SettingDeclaration, type SettingSource, } from "./src/pi-ext-settings.ts"; const STATUS_ID = "context-size"; const FOOTER_SEGMENT = { id: "context-size", icon: "", zone: "llm", order: 1, } as const; type Limits = { warningPercent: number; errorPercent: number }; const DEFAULT_LIMITS: Limits = { warningPercent: 70, errorPercent: 85 }; function parsePercent(raw: unknown, source: SettingSource): number | undefined { const value = parseNumberSetting(raw, source); if (value !== undefined && (value <= 0 || value > 100)) throw new Error("must be greater than 0 and at most 100"); return value; } const WARNING_SETTING: SettingDeclaration = { key: "context-size.warningPercent", env: "PI_CONTEXT_SIZE_WARNING_PERCENT", parse: parsePercent, default: DEFAULT_LIMITS.warningPercent, }; const ERROR_SETTING: SettingDeclaration = { key: "context-size.errorPercent", env: "PI_CONTEXT_SIZE_ERROR_PERCENT", parse: parsePercent, default: DEFAULT_LIMITS.errorPercent, }; const BAR = { width: 10, empty: "─", full: "━", } as const; const UPDATE_EVENTS = [ "session_start", "agent_start", "turn_end", "agent_end", "model_select", "session_compact", ] as const; type SegmentView = { text: string; color: ThemeColor; }; function clamp(value: number, min: number, max: number): number { return Math.min(max, Math.max(min, value)); } function toneForPercent(percent: number | null, limits: Limits): ThemeColor { if (percent === null) return "muted"; if (percent >= limits.errorPercent) return "error"; if (percent >= limits.warningPercent) return "warning"; return "muted"; } function barTone(textTone: ThemeColor): ThemeColor { return textTone === "muted" ? "dim" : textTone; } function renderThinBar(percent: number | null): string { const filled = percent === null ? 0 : clamp(Math.round((percent / 100) * BAR.width), 0, BAR.width); return BAR.full.repeat(filled) + BAR.empty.repeat(BAR.width - filled); } /** Human-readable token count: 940, 1.2k, 262k, 1.2M, 12M. */ function formatTokens(tokens: number): string { if (tokens < 1_000) return `${Math.round(tokens)}`; if (tokens < 10_000) return `${(tokens / 1_000).toFixed(1)}k`; if (tokens < 1_000_000) return `${Math.round(tokens / 1_000)}k`; if (tokens < 10_000_000) return `${(tokens / 1_000_000).toFixed(1)}M`; return `${Math.round(tokens / 1_000_000)}M`; } /** * TUI view: the bar alone. Its length encodes the percentage and its colour * encodes the pressure, so percent and token counts would only repeat it. */ function buildSegment( ctx: ExtensionContext, usage: ContextUsage, limits: Limits, ): SegmentView { const percent = usage.percent ?? null; const textTone = toneForPercent(percent, limits); return { text: ctx.ui.theme.fg(barTone(textTone), renderThinBar(percent)), color: textTone, }; } /** * Non-TUI view: no bar and no glyphs, because RPC clients render the status * line as plain text in an unknown font. */ function buildPlainLabel(usage: ContextUsage): string { const percent = usage.percent ?? null; if (percent === null || usage.tokens === null) return "context unknown"; const left = Math.max(0, usage.contextWindow - usage.tokens); return `context ${Math.round(percent)}% used, ${formatTokens(left)} left`; } export default function contextSizeExtension(pi: ExtensionAPI) { let limits = DEFAULT_LIMITS; function resolveLimits(ctx: ExtensionContext): Limits { const warning = resolveSetting(pi, ctx, WARNING_SETTING); const error = resolveSetting(pi, ctx, ERROR_SETTING); let problem: string; if (!warning.ok || !error.ok) problem = [warning, error] .flatMap((result) => (result.ok ? [] : [result.error])) .join("; "); else if (warning.value < error.value) return { warningPercent: warning.value, errorPercent: error.value }; else problem = `${WARNING_SETTING.key} (${warning.value}) must be below ${ERROR_SETTING.key} (${error.value})`; if (ctx.hasUI) ctx.ui.notify( `${problem}; using defaults ${DEFAULT_LIMITS.warningPercent}% and ${DEFAULT_LIMITS.errorPercent}%`, "warning", ); return DEFAULT_LIMITS; } function clearSegment() { removeFooterSegment(pi, FOOTER_SEGMENT.id); } function clearPublished(ctx: ExtensionContext) { clearSegment(); ctx.ui.setStatus(STATUS_ID, undefined); } function updateSegment(ctx: ExtensionContext) { try { if (!ctx.hasUI) { dbg?.("update", { outcome: "no-ui" }); return; } const usage = ctx.getContextUsage(); if (!usage) { clearPublished(ctx); dbg?.("update", { outcome: "unknown" }); return; } if (ctx.mode !== "tui") { clearSegment(); ctx.ui.setStatus(STATUS_ID, buildPlainLabel(usage)); dbg?.("update", { outcome: "plain" }); return; } const view = buildSegment(ctx, usage, limits); const accepted = offerFooterSegment(pi, { ...FOOTER_SEGMENT, text: view.text, color: view.color, }); if (!accepted) clearSegment(); ctx.ui.setStatus(STATUS_ID, accepted ? undefined : view.text); dbg?.("update", { outcome: accepted ? "footer" : "fallback" }); } catch { // Ignore late events during reload/switch/shutdown; Pi marks old ctx/pi stale. dbg?.("update", { outcome: "error" }); } } pi.on("session_start", (_event, ctx) => { dbg?.("session.start"); limits = resolveLimits(ctx); }); for (const eventName of UPDATE_EVENTS) { pi.on(eventName as any, async (_event: any, ctx: any) => updateSegment(ctx), ); } pi.on("session_shutdown", async (_event, ctx) => { if (ctx.hasUI) ctx.ui.setStatus(STATUS_ID, undefined); dbg?.("session.shutdown"); closeDebug(); }); // Footer owns its own teardown. Emitting through `pi.events` during session_shutdown // can trip Pi's stale-context guard after reload/switch. } export const __test = { buildPlainLabel, formatTokens, renderThinBar };