omnirush 0.3.3 → 0.4.2

package/assets/extensions/omnirush/collector.ts

@@ -42,10 +42,18 @@ export default function (pi: any) {
           return ok ? refresher.auth?.accessToken ?? null : null;
         }
       : undefined,
-    log: (level, message) => {
+    log: (level, message, attributes) => {
       if (level === "warn") {
         try {
-          process.stderr.write(`omnirush collector: ${message}\n`);
+          // Attributes carry the diagnosis (HTTP status, error message,
+          // path, byte counts) — without them a remote failure is
+          // undebuggable ("collection operation failed" and nothing
+          // else). Attribute values are telemetry-safe (ids, sizes,
+          // error strings; never tokens).
+          const detail = attributes && Object.keys(attributes).length > 0
+            ? ` ${JSON.stringify(attributes)}`
+            : "";
+          process.stderr.write(`omnirush collector: ${message}${detail}\n`);
         } catch {
           /* stderr may be closed late in shutdown */
         }

package/assets/extensions/omnirush/commands.ts

@@ -1,48 +1,173 @@
 /**
- * Omnirush commands: /grant (today's token grant + usage) and /verify
- * (prove the served model is genuine — full implementation lands with the
- * gateway's model-attestation endpoint), plus the footer status line that
- * keeps grant/model info visible while the agent runs.
+ * Omnirush commands:
+ *   /grant   today's token grant + usage (headers when a request ran,
+ *            /device/me otherwise)
+ *   /usage   full usage from the manager (grant date, token counters,
+ *            devices)
+ *   /status  one-screen session summary (identity, model/effort, grant,
+ *            collector, origin, version)
+ *   /verify  prove the served model is genuine — full implementation
+ *            lands with the gateway's model-attestation endpoint (M2)
+ * plus the footer status line that keeps grant/model info visible while
+ * the agent runs.
+ *
+ * Handler contract (pi): handler(args: string, ctx: ExtensionCommandContext)
+ * — the FIRST parameter is the argument string, ctx comes SECOND. Getting
+ * this wrong silently no-ops the command (ctx.ui is undefined on a string),
+ * which is exactly how /grant and /verify broke the first time.
  */
 
+import { readSessionLedger, SESSION_LEDGER_FILE } from "./collector-lib";
+import { deviceMe, omniDir, resolveOrigin } from "./auth";
+import { sharedRefresher } from "./refresh";
+import { formatCollectorLine, formatStatusLines } from "./status-lib";
+import {
+  formatDeviceLines,
+  formatGrantLines,
+  formatGrantStatusLine,
+  formatTokenCount,
+  formatUsageDetails,
+  formatUsageSummary,
+  gatewayUsage,
+  grantFromUsage,
+} from "./usage";
+
 import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
-import { deviceMe, resolveOrigin } from "./auth";
-import { formatGrantLines, formatGrantStatusLine, gatewayUsage } from "./usage";
 
-/** Fetch the signed-in identity + server-side usage, best effort. */
-async function fetchIdentity(): Promise<any> {
-  const token = (process.env.OMNIRUSH_TOKEN || "").trim();
-  if (!token) return null;
+const CLI_VERSION = (process.env.OMNIRUSH_VERSION || "").trim() || "unknown";
+
+interface MeResult {
+  me: any;
+  error?: string;
+}
+
+/**
+ * Fetch the signed-in identity + server-side usage with the broker
+ * pattern: 401 triggers one single-flight refresh (both tokens rotate,
+ * persisted 0600) and a single retry with the rotated token. Never
+ * throws; errors come back as `error` strings safe to display.
+ */
+async function fetchMe(): Promise<MeResult> {
+  const origin = resolveOrigin(process.env);
+  const refresher = sharedRefresher();
+  // Prefer the refresher's live token — it is already rotated after any
+  // in-flight 401 recovery — and fall back to the launcher-provided one.
+  const token = refresher.auth?.accessToken || (process.env.OMNIRUSH_TOKEN || "").trim();
+  if (!token) return { me: null, error: "not signed in — run `omnirush login` first" };
   try {
-    return await deviceMe(resolveOrigin(process.env), {
-      accessToken: token,
-      fetchImpl: globalThis.fetch,
-    });
-  } catch {
-    return null;
+    let me = await deviceMe(origin, { accessToken: token, fetchImpl: globalThis.fetch });
+    if (me !== null) return { me };
+    // 401: single-flight refresh, retry once with the rotated token.
+    const refreshed = await refresher.refresh(token).catch(() => false);
+    const next = refresher.auth?.accessToken;
+    if (refreshed && next && next !== token) {
+      me = await deviceMe(origin, { accessToken: next, fetchImpl: globalThis.fetch });
+      if (me !== null) return { me };
+    }
+    return { me: null, error: "credentials rejected — run `omnirush login` again" };
+  } catch (error) {
+    return { me: null, error: `manager unreachable (${error?.message ?? error})` };
   }
 }
 
+function identityOf(me: any): string | null {
+  if (!me || typeof me !== "object") return null;
+  for (const key of ["email", "device_name", "user_id", "id"] as const) {
+    const value = me[key];
+    if (typeof value === "string" && value) return value;
+  }
+  return null;
+}
+
+/** The grant summary string shared by /grant and /status, or null. */
+function grantSummary(me: any): { text: string; date?: string } | null {
+  const usage = gatewayUsage();
+  const fromHeaders = grantFromUsage(usage);
+  const meUsage = me?.usage ?? null;
+  const limit = fromHeaders?.limitTokens ?? Number(meUsage?.token_limit);
+  const remaining = fromHeaders?.remainingTokens ?? Number(meUsage?.remaining_tokens);
+  if (!Number.isFinite(limit) || !Number.isFinite(remaining)) return null;
+  const used = Number(meUsage?.used_tokens);
+  const usedPart = Number.isFinite(used) ? ` (used ${formatTokenCount(used)})` : "";
+  return {
+    text: `${formatTokenCount(remaining)} of ${formatTokenCount(limit)} tokens remaining${usedPart}`,
+    date: typeof meUsage?.grant_date === "string" ? meUsage.grant_date : undefined,
+  };
+}
+
+function notify(ctx: any, text: string, level: "info" | "error" = "info"): void {
+  ctx?.ui?.notify?.(text, level);
+}
+
 export default function (pi: ExtensionAPI) {
   pi.registerCommand("grant", {
     description: "Show today's Omnirush token grant and usage",
-    handler: async (ctx: any) => {
+    handler: async (_args: string, ctx: any) => {
       const usage = gatewayUsage();
-      const me = await fetchIdentity();
+      const { me } = await fetchMe();
       const sessionId = ctx?.sessionManager?.getSessionId?.() ?? undefined;
       const lines = formatGrantLines({ usage, me, sessionId });
-      ctx.ui?.notify?.(
-        ["Omnirush usage:", ...lines.map((line) => `  ${line}`)].join("\n"),
-      );
+      notify(ctx, ["Omnirush usage:", ...lines.map((line) => `  ${line}`)].join("\n"));
+    },
+  });
+
+  pi.registerCommand("usage", {
+    description: "Show detailed Omnirush usage from the manager (counters + devices)",
+    handler: async (_args: string, ctx: any) => {
+      const { me, error } = await fetchMe();
+      if (!me) {
+        notify(ctx, `Omnirush usage unavailable: ${error}`, "error");
+        return;
+      }
+      const lines = [
+        "Omnirush usage (from the manager):",
+        ...formatUsageDetails({ me }).map((line) => `  ${line}`),
+      ];
+      const summary = formatUsageSummary(me);
+      if (summary) lines.push(`  summary: ${summary}`);
+      lines.push("", ...formatDeviceLines(me));
+      notify(ctx, lines.join("\n"));
+    },
+  });
+
+  pi.registerCommand("status", {
+    description: "Show the Omnirush session status (identity, model, grant, collector)",
+    handler: async (_args: string, ctx: any) => {
+      const { me } = await fetchMe();
+      const usage = gatewayUsage();
+      const grant = grantSummary(me);
+      const sessionId = String(ctx?.sessionManager?.getSessionId?.() ?? "");
+      let record = null;
+      try {
+        const ledger = await readSessionLedger(`${omniDir()}/${SESSION_LEDGER_FILE}`);
+        record = ledger.sessions[sessionId] ?? null;
+      } catch {
+        /* collector state is best effort */
+      }
+      const model = ctx?.model;
+      const statusLines = formatStatusLines({
+        signedIn: Boolean(me),
+        identity: identityOf(me),
+        model: model ? `${model.provider ?? "omnirush"}/${model.id}` : null,
+        thinkingLevel: typeof ctx?.thinkingLevel === "string" ? ctx.thinkingLevel : null,
+        effort: usage?.["x-omnirush-effort"] ?? null,
+        grant: grant?.text ?? null,
+        grantDate: grant?.date ?? null,
+        collector: { sessionId, record },
+        origin: resolveOrigin(process.env),
+        version: CLI_VERSION,
+      });
+      notify(ctx, ["Omnirush status:", ...statusLines.map((line) => `  ${line}`)].join("\n"));
     },
   });
 
   pi.registerCommand("verify", {
     description: "Verify the model serving your requests is genuine GPT-6 Astra",
-    handler: async (ctx: any) => {
+    handler: async (_args: string, ctx: any) => {
       // M2: POST /omnirush/verify → signed attestation compared against
       // observed model behavior. v0 placeholder.
-      ctx.ui?.notify?.(
+      notify(
+        ctx,
         "verify: model attestation arrives with the gateway lane (M2).",
       );
     },

package/assets/extensions/omnirush/index.ts

@@ -3,10 +3,13 @@
 // pi discovers either extensions/*.ts files or extensions/<dir>/index.ts —
 // this index makes the whole omnirush extension directory load as one
 // extension (without it, none of these files were ever registered).
+// EVERY module that registers commands/tools/handlers must be wired here:
+// a module left out of this list is silently dead in sessions.
 
 import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
 import collector from "./collector";
 import commands from "./commands";
+import mcp from "./mcp";
 import recorder from "./recorder";
 import sota from "./sota";
 
@@ -15,4 +18,5 @@ export default function (pi: ExtensionAPI) {
   recorder(pi as any);
   sota(pi as any);
   collector(pi as any);
+  mcp(pi as any);
 }