@pi-vault/pi-status 0.1.0 → 0.2.1

package/src/index.ts

@@ -2,20 +2,13 @@ import type {
   ExtensionAPI,
   ExtensionContext,
 } from "@earendil-works/pi-coding-agent";
-import {
-  USAGE_CORE_READY_EVENT,
-  USAGE_CORE_REQUEST_EVENT,
-  USAGE_CORE_UPDATE_CURRENT_EVENT,
-} from "@pi-vault/pi-usage/events";
-import type { UsageCoreState } from "@pi-vault/pi-usage/types";
-import {
-  loadConfig,
-  saveConfigToSettings,
-  type PiStatusConfig,
-} from "./config.ts";
-import { buildFooterLine } from "./render.ts";
-import { createStatuslineEditor } from "./ui/statusline-editor.ts";
-import { fromPiTheme, noTheme, type StatuslineMenuTheme } from "./ui/statusline-theme.ts";
+import { loadConfig, saveConfigToSettings } from "./core/config.ts";
+import { buildSnapshot } from "./core/snapshot.ts";
+import { createUsageRuntime } from "./core/usage-runtime.ts";
+import type { PiStatusConfig } from "./shared/types.ts";
+import { createStatusLineEditor } from "./tui/editor.ts";
+import { buildFooterLine } from "./tui/render.ts";
+import { fromPiTheme, noTheme, type StatusLineTheme } from "./tui/theme.ts";
 
 type FooterComponent = {
   render: (width: number) => string[];
@@ -35,85 +28,55 @@ type FooterFactory = (
   footerData: FooterDataLike,
 ) => FooterComponent;
 
-// Used only to suppress the live footer while a custom UI (e.g. /statusline)
-// is open. Renders no lines and is a no-op otherwise.
+type RuntimeState = {
+  config: PiStatusConfig;
+  ctx: ExtensionContext | undefined;
+  requestRender: (() => void) | undefined;
+  gitBranch: string | null;
+  extensionStatuses: Map<string, string>;
+};
+
+function createRuntimeState(): RuntimeState {
+  return {
+    config: loadConfig().config,
+    ctx: undefined,
+    requestRender: undefined,
+    gitBranch: null,
+    extensionStatuses: new Map(),
+  };
+}
+
 const EMPTY_FOOTER_FACTORY: FooterFactory = () => ({
-  render(_width: number): string[] {
+  render(): string[] {
     return [];
   },
   invalidate(): void {},
   dispose(): void {},
 });
 
-// `ctx.ui.custom(...)` is expected to hand us a live Pi theme, but the
-// runtime contract is loose and the theme may be missing in some test or
-// boot contexts. We only adapt it when both `fg` and `bold` are callable so
-// the editor always gets a well-formed `StatuslineMenuTheme` and never has
-// to defend against partial themes at render time.
 function isLiveTheme(value: unknown): boolean {
   if (!value || typeof value !== "object") return false;
   const candidate = value as { fg?: unknown; bold?: unknown };
-  return typeof candidate.fg === "function" && typeof candidate.bold === "function";
-}
-
-function aggregateBranchTotals(ctx: ExtensionContext): {
-  input: number;
-  output: number;
-  totalTokens: number;
-} {
-  const totals = { input: 0, output: 0, totalTokens: 0 };
-  const branch = ctx.sessionManager.getBranch() as unknown[];
-
-  for (const entry of branch ?? []) {
-    if (!entry || typeof entry !== "object") continue;
-    const type = (entry as { type?: unknown }).type;
-    if (type !== "message") continue;
-    const message = (
-      entry as {
-        message?: {
-          role?: unknown;
-          usage?: { input?: number; output?: number; totalTokens?: number };
-        };
-      }
-    ).message;
-    if (message?.role !== "assistant") continue;
-    const usage = message.usage;
-    if (!usage) continue;
-    if (typeof usage.input === "number") totals.input += usage.input;
-    if (typeof usage.output === "number") totals.output += usage.output;
-    if (typeof usage.totalTokens === "number")
-      totals.totalTokens += usage.totalTokens;
-  }
-
-  return totals;
+  return (
+    typeof candidate.fg === "function" && typeof candidate.bold === "function"
+  );
 }
 
 export default function createExtension(pi: ExtensionAPI): void {
-  let runtimeConfig: PiStatusConfig = loadConfig().config;
-  let currentCtx: ExtensionContext | undefined;
-  let requestRender: (() => void) | undefined;
-  let usageState: UsageCoreState | undefined;
-  let lastGitBranch: string | null = null;
-  let lastExtensionStatuses = new Map<string, string>();
+  const state = createRuntimeState();
 
-  function refreshRuntimeConfig(cwd?: string): void {
-    runtimeConfig = loadConfig(cwd ? { cwd } : undefined).config;
-  }
+  const usageRuntime = createUsageRuntime(pi);
 
-  function acceptUsageState(payload: unknown): void {
-    if (!payload || typeof payload !== "object") return;
-    const maybe = payload as { state?: unknown };
-    const next =
-      maybe.state && typeof maybe.state === "object" ? maybe.state : payload;
-    usageState = next as UsageCoreState;
-    requestRender?.();
+  function refreshRuntimeConfig(cwd?: string): void {
+    state.config = loadConfig(cwd ? { cwd } : undefined).config;
   }
 
   function installFooter(ctx: ExtensionContext): void {
     if (!ctx.hasUI) return;
 
     const factory: FooterFactory = (tui, theme, footerData) => {
-      requestRender = () => tui.requestRender?.();
+      state.requestRender = () => tui.requestRender?.();
+      usageRuntime.setOnChange(state.requestRender);
       const unsubscribe = footerData.onBranchChange?.(() =>
         tui.requestRender?.(),
       );
@@ -121,42 +84,41 @@ export default function createExtension(pi: ExtensionAPI): void {
       return {
         dispose() {
           unsubscribe?.();
-          if (requestRender === tui.requestRender) {
-            requestRender = undefined;
-          }
+          if (state.requestRender === tui.requestRender)
+            state.requestRender = undefined;
+          usageRuntime.setOnChange(state.requestRender);
         },
         invalidate() {
-          requestRender?.();
+          state.requestRender?.();
         },
         render(width: number) {
-          const activeCtx = currentCtx ?? ctx;
-          lastGitBranch = footerData.getGitBranch();
-          lastExtensionStatuses = new Map(
+          const activeCtx = state.ctx ?? ctx;
+          state.gitBranch = footerData.getGitBranch();
+          state.extensionStatuses = new Map(
             footerData.getExtensionStatuses().entries(),
           );
+          const snapshot = buildSnapshot({
+            model: activeCtx.model,
+            cwd: activeCtx.cwd,
+            thinkingLevel: String(pi.getThinkingLevel()),
+            gitBranch: state.gitBranch,
+            isIdle: activeCtx.isIdle(),
+            hasPendingMessages: activeCtx.hasPendingMessages(),
+            contextUsage: activeCtx.getContextUsage(),
+            branch: activeCtx.sessionManager.getBranch() as unknown[],
+            sessionId: activeCtx.sessionManager.getSessionId(),
+            usageState: usageRuntime.getState(),
+            extensionStatuses: state.extensionStatuses,
+          });
           const line = buildFooterLine(
             {
-              model: activeCtx.model,
-              cwd: activeCtx.cwd,
-              thinkingLevel: String(pi.getThinkingLevel()),
-              gitBranch: lastGitBranch,
-              runState: !activeCtx.isIdle()
-                ? "busy"
-                : activeCtx.hasPendingMessages()
-                  ? "queued"
-                  : "idle",
-              contextUsage: activeCtx.getContextUsage(),
-              branchTotals: aggregateBranchTotals(activeCtx),
-              sessionId: activeCtx.sessionManager.getSessionId(),
-              usageState,
-              extensionStatuses: lastExtensionStatuses,
-              filter: runtimeConfig.filter,
-              segments: runtimeConfig.segments,
+              ...snapshot,
+              filter: state.config.filter,
+              segments: state.config.segments,
             },
             theme,
             width,
           );
-
           return [line];
         },
       };
@@ -166,37 +128,15 @@ export default function createExtension(pi: ExtensionAPI): void {
   }
 
   function installEmptyFooter(ctx: ExtensionContext): void {
-    if (!ctx.hasUI) return;
-    ctx.ui.setFooter(EMPTY_FOOTER_FACTORY as never);
+    if (ctx.hasUI) ctx.ui.setFooter(EMPTY_FOOTER_FACTORY as never);
   }
 
   function refresh(ctx: ExtensionContext): void {
-    currentCtx = ctx;
+    state.ctx = ctx;
     refreshRuntimeConfig(ctx.cwd);
-    requestRender?.();
+    state.requestRender?.();
   }
 
-  const unsubscribeUsageReady = pi.events.on(
-    USAGE_CORE_READY_EVENT,
-    (payload: unknown) => {
-      acceptUsageState(payload);
-    },
-  );
-
-  const unsubscribeUsageUpdate = pi.events.on(
-    USAGE_CORE_UPDATE_CURRENT_EVENT,
-    (payload: unknown) => {
-      acceptUsageState(payload);
-    },
-  );
-
-  pi.events.emit(USAGE_CORE_REQUEST_EVENT, {
-    type: "current",
-    reply(payload: unknown) {
-      acceptUsageState(payload);
-    },
-  });
-
   pi.registerCommand("statusline", {
     description: "Configure statusline segments and extension-status filters",
     handler: async (_args, ctx) => {
@@ -205,7 +145,7 @@ export default function createExtension(pi: ExtensionAPI): void {
         return;
       }
 
-      const discovered = [...lastExtensionStatuses.keys()].sort((a, b) =>
+      const discovered = [...state.extensionStatuses.keys()].sort((a, b) =>
         a.localeCompare(b),
       );
 
@@ -214,32 +154,31 @@ export default function createExtension(pi: ExtensionAPI): void {
         installEmptyFooter(ctx);
         result = await ctx.ui.custom<PiStatusConfig | null>(
           (tui, theme, _keys, done) => {
-            const activeCtx = currentCtx ?? ctx;
-            const menuTheme: StatuslineMenuTheme = isLiveTheme(theme)
+            const activeCtx = state.ctx ?? ctx;
+            const menuTheme: StatusLineTheme = isLiveTheme(theme)
               ? fromPiTheme(theme)
               : noTheme;
-            return createStatuslineEditor({
-              config: runtimeConfig,
+            const snapshot = buildSnapshot({
+              model: activeCtx.model,
+              cwd: activeCtx.cwd,
+              thinkingLevel: String(pi.getThinkingLevel()),
+              gitBranch: state.gitBranch,
+              isIdle: activeCtx.isIdle(),
+              hasPendingMessages: activeCtx.hasPendingMessages(),
+              contextUsage: activeCtx.getContextUsage(),
+              branch: activeCtx.sessionManager.getBranch() as unknown[],
+              sessionId: activeCtx.sessionManager.getSessionId(),
+              usageState: usageRuntime.getState(),
+              extensionStatuses: state.extensionStatuses,
+            });
+            return createStatusLineEditor({
+              config: state.config,
               discoveredStatuses: discovered,
-              previewInput: {
-                model: activeCtx.model,
-                cwd: activeCtx.cwd,
-                thinkingLevel: String(pi.getThinkingLevel()),
-                gitBranch: lastGitBranch,
-                runState: !activeCtx.isIdle()
-                  ? "busy"
-                  : activeCtx.hasPendingMessages()
-                    ? "queued"
-                    : "idle",
-                contextUsage: activeCtx.getContextUsage(),
-                branchTotals: aggregateBranchTotals(activeCtx),
-                sessionId: activeCtx.sessionManager.getSessionId(),
-                usageState,
-                extensionStatuses: lastExtensionStatuses,
-              },
+              previewInput: snapshot,
               theme: menuTheme,
               done,
               requestRender: () => tui.requestRender?.(),
+              usageAvailable: usageRuntime.getAvailable(),
             });
           },
         );
@@ -251,8 +190,8 @@ export default function createExtension(pi: ExtensionAPI): void {
 
       try {
         saveConfigToSettings(result, { cwd: ctx.cwd });
-        runtimeConfig = result;
-        requestRender?.();
+        state.config = result;
+        state.requestRender?.();
       } catch (error) {
         const message =
           error instanceof Error
@@ -264,14 +203,15 @@ export default function createExtension(pi: ExtensionAPI): void {
   });
 
   pi.on("session_start", (_event, ctx) => {
+    usageRuntime.requestCurrent();
     refreshRuntimeConfig(ctx.cwd);
-    currentCtx = ctx;
+    state.ctx = ctx;
     installFooter(ctx);
   });
 
   pi.on("session_tree", (_event, ctx) => {
     refreshRuntimeConfig(ctx.cwd);
-    currentCtx = ctx;
+    state.ctx = ctx;
     installFooter(ctx);
   });
 
@@ -284,10 +224,9 @@ export default function createExtension(pi: ExtensionAPI): void {
   });
 
   pi.on("session_shutdown", (_event, ctx) => {
-    currentCtx = undefined;
-    requestRender = undefined;
-    unsubscribeUsageReady();
-    unsubscribeUsageUpdate();
+    state.ctx = undefined;
+    state.requestRender = undefined;
+    usageRuntime.setOnChange(undefined);
     if (ctx.hasUI) ctx.ui.setFooter(undefined);
   });
 }

package/src/shared/types.ts

@@ -0,0 +1,63 @@
+export type StatusLineSegmentId =
+  | "model"
+  | "model-with-reasoning"
+  | "project-name"
+  | "current-dir"
+  | "git-branch"
+  | "run-state"
+  | "context-remaining"
+  | "context-used"
+  | "context-window-size"
+  | "used-tokens"
+  | "total-input-tokens"
+  | "total-output-tokens"
+  | "session-id"
+  | "five-hour-limit"
+  | "weekly-limit"
+  | "extension-statuses";
+
+export type StatusFilter =
+  | { mode: "all"; hidden: string[] }
+  | { mode: "only"; shown: string[] };
+
+export type PiStatusConfig = {
+  segments: StatusLineSegmentId[];
+  filter: StatusFilter;
+};
+
+export const KNOWN_SEGMENTS: readonly StatusLineSegmentId[] = [
+  "model",
+  "model-with-reasoning",
+  "project-name",
+  "current-dir",
+  "git-branch",
+  "run-state",
+  "context-remaining",
+  "context-used",
+  "context-window-size",
+  "used-tokens",
+  "total-input-tokens",
+  "total-output-tokens",
+  "session-id",
+  "five-hour-limit",
+  "weekly-limit",
+  "extension-statuses",
+] as const;
+
+export const DEFAULT_SEGMENTS: readonly StatusLineSegmentId[] = [
+  "model-with-reasoning",
+  "current-dir",
+] as const;
+
+export const USAGE_SEGMENTS = new Set<StatusLineSegmentId>([
+  "five-hour-limit",
+  "weekly-limit",
+]);
+
+export function isKnownSegment(value: string): value is StatusLineSegmentId {
+  return (KNOWN_SEGMENTS as readonly string[]).includes(value);
+}
+
+export function isUsageSegment(id: StatusLineSegmentId): boolean {
+  return USAGE_SEGMENTS.has(id);
+}