From 615762b0ac19fcf0ddcac2f51107a000c167921f Mon Sep 17 00:00:00 2001
From: Evan Reichard <evan@reichard.io>
Date: Sun, 3 May 2026 11:58:13 -0400
Subject: [PATCH] chore: use auth storage

---
 index.ts | 102 +++++++++++++++++++++++++++++++++----------------------
 1 file changed, 62 insertions(+), 40 deletions(-)

diff --git a/index.ts b/index.ts
index afb9d98..8e7e093 100644
--- a/index.ts
+++ b/index.ts
@@ -1,7 +1,4 @@
-import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
-import { readFileSync } from "node:fs";
-import { homedir } from "node:os";
-import { join } from "node:path";
+import { AuthStorage, type AuthCredential, type ExtensionAPI } from "@mariozechner/pi-coding-agent";
 import { statusbarConfig } from "./config";
 import { contextModule, costModule, directoryModule, modelModule, thinkingModule } from "./modules/basic";
 import { usageModule } from "./modules/usage";
@@ -12,7 +9,6 @@ import { zaiUsageProvider } from "./usage/zai";
 import type { ModuleContext, ModuleSpec, RenderedModule, StatusbarState } from "./types";
 import type { Provider, UsageCredential, UsageProvider, UsageReport } from "./usage";
 
-const AUTH_PATH = join(homedir(), ".pi", "agent", "auth.json");
 const REFRESH_MS = 60_000;
 const ANTHROPIC_REFRESH_MS = 15 * 60_000;
 
@@ -32,36 +28,49 @@ function activeProvider(ctx: any): Provider | undefined {
   return undefined;
 }
 
-function readPiCredential(provider: Provider): UsageCredential | undefined {
-  try {
-    const auth = JSON.parse(readFileSync(AUTH_PATH, "utf8"));
-    if (!isRecord(auth) || !isRecord(auth[provider])) return undefined;
+function credentialString(raw: unknown, key: string): string | undefined {
+  return isRecord(raw) && typeof raw[key] === "string" ? raw[key] : undefined;
+}
 
-    const raw = auth[provider];
-    if (raw.type === "oauth") {
-      return {
-        type: "oauth",
-        accessToken: typeof raw.access === "string" ? raw.access : undefined,
-        refreshToken: typeof raw.refresh === "string" ? raw.refresh : undefined,
-        expiresAt: typeof raw.expires === "number" ? raw.expires : undefined,
-        accountId: typeof raw.accountId === "string" ? raw.accountId : undefined,
-        email: typeof raw.email === "string" ? raw.email : undefined,
-        metadata: raw,
-      };
-    }
+function buildUsageCredential(raw: AuthCredential | undefined, apiKey: string): UsageCredential {
+  if (raw?.type === "oauth") {
+    return {
+      type: "oauth",
+      accessToken: apiKey,
+      refreshToken: raw.refresh,
+      expiresAt: raw.expires,
+      accountId: credentialString(raw, "accountId"),
+      email: credentialString(raw, "email"),
+      metadata: raw,
+    };
+  }
 
-    if (raw.type === "api_key") {
-      return {
-        type: "api_key",
-        apiKey: typeof raw.key === "string" ? raw.key : undefined,
-        accountId: typeof raw.accountId === "string" ? raw.accountId : undefined,
-        email: typeof raw.email === "string" ? raw.email : undefined,
-        metadata: raw,
-      };
-    }
-  } catch {}
+  return {
+    type: "api_key",
+    apiKey,
+    accountId: credentialString(raw, "accountId"),
+    email: credentialString(raw, "email"),
+    metadata: raw,
+  };
+}
 
-  return undefined;
+async function readPiCredential(authStorage: AuthStorage, provider: Provider): Promise<UsageCredential | undefined> {
+  authStorage.reload();
+  const apiKey = await authStorage.getApiKey(provider);
+  if (!apiKey) return undefined;
+  return buildUsageCredential(authStorage.get(provider), apiKey);
+}
+
+async function forceRefreshPiCredential(authStorage: AuthStorage, provider: Provider): Promise<UsageCredential> {
+  authStorage.reload();
+  const raw = authStorage.get(provider);
+  const oauthProvider = authStorage.getOAuthProviders().find(candidate => candidate.id === provider);
+  if (raw?.type !== "oauth" || !oauthProvider) throw new Error("login expired");
+
+  // Refresh Provider OAuth Token
+  const refreshed = await oauthProvider.refreshToken(raw);
+  authStorage.set(provider, { type: "oauth", ...refreshed });
+  return buildUsageCredential(authStorage.get(provider), oauthProvider.getApiKey(refreshed));
 }
 
 function renderModule(moduleCtx: ModuleContext, spec: ModuleSpec): RenderedModule {
@@ -105,6 +114,7 @@ export default function piStatusbarExtension(pi: ExtensionAPI) {
   let latestCtx: any;
   let requestRender: (() => void) | undefined;
   const statusbarState: StatusbarState = {};
+  const authStorage = AuthStorage.create();
 
   function updateThinkingLevel() {
     try {
@@ -154,7 +164,7 @@ export default function piStatusbarExtension(pi: ExtensionAPI) {
       return;
     }
 
-    const credential = readPiCredential(provider);
+    const credential = await readPiCredential(authStorage, provider);
     if (!credential) {
       statusbarState.report = { provider, fetchedAt: Date.now(), limits: [] };
       statusbarState.error = "not logged in";
@@ -163,25 +173,37 @@ export default function piStatusbarExtension(pi: ExtensionAPI) {
     }
 
     inFlight?.abort();
-    inFlight = new AbortController();
+    const controller = new AbortController();
+    inFlight = controller;
     statusbarState.report = undefined;
     statusbarState.error = undefined;
     rerender(ctx);
 
     try {
-      const report = await usageProviders[provider].fetchUsage({
-        provider,
-        credential,
-        baseUrl: ctx.model?.baseUrl,
-        signal: inFlight.signal,
-      }, {
+      let activeCredential = credential;
+      const usageCtx = {
         fetch: globalThis.fetch.bind(globalThis),
         logger: {
           debug: () => undefined,
           warn: () => undefined,
         },
+      };
+      const fetchParams = () => ({
+        provider,
+        credential: activeCredential,
+        baseUrl: ctx.model?.baseUrl,
+        signal: controller.signal,
       });
 
+      let report: UsageReport | null;
+      try {
+        report = await usageProviders[provider].fetchUsage(fetchParams(), usageCtx);
+      } catch (error) {
+        if (!(error instanceof Error) || error.message !== "unauthorized") throw error;
+        activeCredential = await forceRefreshPiCredential(authStorage, provider);
+        report = await usageProviders[provider].fetchUsage(fetchParams(), usageCtx);
+      }
+
       statusbarState.report = report ?? { provider, fetchedAt: Date.now(), limits: [] };
       statusbarState.error = report ? undefined : "unavailable";
     } catch (error) {