@monotykamary/pi-better-grok 0.3.0 → 0.3.2

package/index.ts

@@ -7,7 +7,12 @@
  * cli-chat-proxy.grok.com identity-first billing surface using pi's native
  * xAI OAuth credential.
  */
-import { type ExtensionAPI, type ExtensionContext } from "@earendil-works/pi-coding-agent";
+import {
+  type ExtensionAPI,
+  type ExtensionContext,
+  type Theme,
+  type ThemeColor,
+} from "@earendil-works/pi-coding-agent";
 import { Container, SettingsList, type SettingsListTheme } from "@earendil-works/pi-tui";
 import { CONFIG_BASENAME, STATUS_KEY } from "./src/identity.ts";
 import {
@@ -56,6 +61,8 @@ import {
   parseUsageSnapshot,
   requestGrokUsage,
   UsageError,
+  type UsageSegment,
+  type UsageSeverity,
   type UsageSnapshot,
 } from "./src/usage.ts";
 import {
@@ -85,6 +92,46 @@ const requireSettingsListTheme = (): SettingsListTheme => {
   return loadedSettingsListTheme;
 };
 
+const USAGE_SEVERITY_COLORS: Record<UsageSeverity, ThemeColor> = {
+  ok: "success",
+  warning: "warning",
+  critical: "error",
+  muted: "dim",
+};
+
+/**
+ * Keep a coloured segment on one line. Unlike sanitizeStatusText this leaves the
+ * spacing shared by neighbouring segments ("5h: " + "90%") intact.
+ */
+function flattenSegment(text: string): string {
+  return text.replace(/[\r\n\t]+/g, " ");
+}
+
+/**
+ * Colour a usage line segment by segment: the percentage carries the severity
+ * colour while the label, countdown and banked resets stay dim.
+ */
+function colorizeUsageSegments(segments: UsageSegment[], theme: Theme): string {
+  return segments
+    .map((segment) =>
+      theme.fg(USAGE_SEVERITY_COLORS[segment.severity], flattenSegment(segment.text)),
+    )
+    .join("");
+}
+
+/** Pair the fast-mode segment with the usage segments for the status widget. */
+function statusWidgetParts(
+  fast: string | undefined,
+  usage: UsageSegment[] | undefined,
+): UsageSegment[] | undefined {
+  const parts: UsageSegment[] = fast ? [{ text: fast, severity: "muted" }] : [];
+  if (usage?.length) {
+    if (parts.length > 0) parts.push({ text: " · ", severity: "muted" });
+    parts.push(...usage);
+  }
+  return parts.length > 0 ? parts : undefined;
+}
+
 class DynamicBorder {
   readonly #color: (text: string) => string;
 
@@ -548,8 +595,10 @@ export default function betterGrok(pi: ExtensionAPI): void {
           parts.push(contextText);
 
           const cfg = config(ctx);
-          const usageStatusLine = usageController.statusLine(ctx, cfg, usingSubscription);
-          const usageLine = usageStatusLine ? theme.fg("dim", usageStatusLine) : undefined;
+          const usageStatusSegments = usageController.statusSegments(ctx, cfg, usingSubscription);
+          const usageLine = usageStatusSegments
+            ? colorizeUsageSegments(usageStatusSegments, theme)
+            : undefined;
 
           let statsLeft = parts.join(" ");
           let statsLeftWidth = visibleWidth(statsLeft);
@@ -630,22 +679,22 @@ export default function betterGrok(pi: ExtensionAPI): void {
     statusInstalled = text !== undefined;
   }
 
-  function setStatusWidget(ctx: ExtensionContext, text: string | undefined): void {
-    if (!text && !statusWidgetInstalled) return;
+  function setStatusWidget(ctx: ExtensionContext, parts: UsageSegment[] | undefined): void {
+    if (!parts && !statusWidgetInstalled) return;
     ctx.ui.setWidget(
       STATUS_KEY,
-      text
+      parts
         ? (_tui, theme) => ({
             invalidate() {},
             render(width: number): string[] {
-              const line = theme.fg("dim", sanitizeStatusText(text));
+              const line = colorizeUsageSegments(parts, theme);
               return [truncateToWidth(line, width, theme.fg("dim", "..."))];
             },
           })
         : undefined,
       { placement: "belowEditor" },
     );
-    statusWidgetInstalled = text !== undefined;
+    statusWidgetInstalled = parts !== undefined;
   }
 
   function updateFooter(ctx: ExtensionContext): void {
@@ -678,8 +727,8 @@ export default function betterGrok(pi: ExtensionAPI): void {
     }
 
     const fast = fastController.statusSegment(ctx, cfg);
-    const usage = usageController.statusLine(ctx, cfg);
-    setStatusWidget(ctx, [fast, usage].filter(Boolean).join(" · ") || undefined);
+    const usage = usageController.statusSegments(ctx, cfg);
+    setStatusWidget(ctx, statusWidgetParts(fast, usage));
   }
 
   pi.on("session_start", (_event, ctx) => {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@monotykamary/pi-better-grok",
-  "version": "0.3.0",
+  "version": "0.3.2",
   "description": "Improve Grok/xAI in pi with fast mode, subscription usage stats, banked reset redemption, multiprovider pools, footer polish, and settings — mirroring pi-better-openai.",
   "keywords": [
     "footer",

package/src/usage-controller.ts

@@ -3,7 +3,13 @@ import type { ResolvedConfig } from "./config.ts";
 import { isXaiProvider, readPiStoredOAuthToken } from "./grok-auth.ts";
 import { sanitizeDiagnosticError } from "./format.ts";
 import { currentModelKey } from "./fast-controller.ts";
-import { formatUsageDetail, formatUsageSnapshot, type UsageSnapshot } from "./usage.ts";
+import {
+  type UsageSegment,
+  type UsageSnapshot,
+  formatUsageDetail,
+  formatUsageSnapshot,
+  usageSegments,
+} from "./usage.ts";
 
 export function isGrokSubscriptionModel(
   ctx: Pick<ExtensionContext, "model" | "modelRegistry">,
@@ -77,22 +83,32 @@ export class UsageController {
     return this.usageSnapshot;
   }
 
-  statusLine(
+  statusSegments(
     ctx: ExtensionContext,
     cfg = this.getConfig(ctx),
     isUsingOAuth?: boolean,
-  ): string | undefined {
+  ): UsageSegment[] | undefined {
     return this.usageSnapshot &&
       !this.usageError &&
       cfg.usage.enabled &&
       isGrokSubscriptionModel(ctx, cfg, isUsingOAuth)
-      ? formatUsageSnapshot(this.usageSnapshot, {
+      ? usageSegments(this.usageSnapshot, {
           ...cfg.usage,
           bankedResets: this.getBankedResets?.() ?? null,
         })
       : undefined;
   }
 
+  statusLine(
+    ctx: ExtensionContext,
+    cfg = this.getConfig(ctx),
+    isUsingOAuth?: boolean,
+  ): string | undefined {
+    return this.statusSegments(ctx, cfg, isUsingOAuth)
+      ?.map((segment) => segment.text)
+      .join("");
+  }
+
   formatStatus(ctx: ExtensionContext): string {
     const cfg = this.getConfig(ctx);
     if (!cfg.usage.enabled) return "Usage display is disabled.";

package/src/usage.ts

@@ -241,24 +241,59 @@ export function formatBankedResetsSuffix(count: number | null | undefined): stri
   return `${count} banked reset${count === 1 ? "" : "s"}`;
 }
 
-export function formatUsageSnapshot(
+export type UsageSeverity = "ok" | "warning" | "critical" | "muted";
+
+export type UsageSegment = {
+  text: string;
+  severity: UsageSeverity;
+};
+
+/** Remaining budget at or below these thresholds turns the percentage amber/red. */
+const WARNING_LEFT_PERCENT = 30;
+const CRITICAL_LEFT_PERCENT = 10;
+
+export function severityForLeftPercent(percent: number | null): UsageSeverity {
+  if (percent === null) return "muted";
+  if (percent <= CRITICAL_LEFT_PERCENT) return "critical";
+  if (percent <= WARNING_LEFT_PERCENT) return "warning";
+  return "ok";
+}
+
+export function usageSegments(
   snapshot: UsageSnapshot,
   options: UsageStatusOptions,
   now = Date.now(),
-): string {
+): UsageSegment[] {
   const used = snapshot.creditUsagePercent;
   const left = used === null ? null : clampPercent(100 - used);
-  const parts = [`Usage: ${formatPercent(left)} left`];
+  const segments: UsageSegment[] = [
+    { text: "Usage: ", severity: "muted" },
+    { text: formatPercent(left), severity: severityForLeftPercent(left) },
+    { text: " left", severity: "muted" },
+  ];
   if (options.showResetTimes) {
     const seconds = periodSecondsLeft(snapshot, now);
     const countdown = formatResetCountdown(seconds);
     const clock = formatResetClock(seconds, now);
-    if (countdown && clock) parts.push(`↺ ${countdown} - ${clock}`);
+    if (countdown && clock) {
+      segments.push({ text: ` · ↺ ${countdown} - ${clock}`, severity: "muted" });
+    }
   }
   const banked =
     options.showBankedResets === false ? null : formatBankedResetsSuffix(options.bankedResets);
-  if (banked) parts.push(banked);
-  return parts.join(" · ");
+  if (banked) segments.push({ text: ` · ${banked}`, severity: "muted" });
+  return segments;
+}
+
+/** Flat text form of {@link usageSegments}, used for status lines and notifications. */
+export function formatUsageSnapshot(
+  snapshot: UsageSnapshot,
+  options: UsageStatusOptions,
+  now = Date.now(),
+): string {
+  return usageSegments(snapshot, options, now)
+    .map((segment) => segment.text)
+    .join("");
 }
 
 export function formatUsageDetail(snapshot: UsageSnapshot): string {