bermudis-pi-goodies 0.27.6 → 0.27.7

package/README.md

@@ -25,7 +25,7 @@ extensions. One entry point, thirteen independent features.
 First install, pinned to a published version:
 
 ```bash
-pi install npm:bermudis-pi-goodies@0.27.6
+pi install npm:bermudis-pi-goodies@0.27.7
 ```
 
 Picking up a new release on an existing install (`pi install` is not the
@@ -279,15 +279,23 @@ catalog no more than every four hours. Balance and quota requests run in the bac
 never delay session startup, model selection, or post-run input readiness.
 The footer also reads OpenRouter remaining credits for `openrouter`, z.ai GLM
 Coding Plan token quota for `zai` (Global) and `zai-coding-cn` (BigModel China),
-OpenAI Codex's ChatGPT subscription quota when using `openai-codex` OAuth,
-and CommandCode credit balance plus 5-hour/weekly usage windows for
+OpenAI Codex's ChatGPT subscription quota when using `openai-codex` OAuth, and
+CommandCode credit balance plus 5-hour/weekly usage windows for
 `commandcode` and `commandcode-anthropic` (same account, same key);
 it skips platform API-key auth because that has no ChatGPT subscription quota.
+For `openai` OAuth (Sign in with ChatGPT), OpenAI rejects the new token on its
+usage endpoints (`401 rejected_by_access_enforcement`), so the footer reads the
+plan windows through a companion `openai-codex` login on the same ChatGPT
+account. The active app's `client_id` must appear exactly once in
+`GET /wham/usage/chatpass/apps` queried with the Codex token; otherwise no
+balance is shown rather than another account's quota. Run `/login openai-codex`
+once with the same account and keep `openai` as the active model.
 OpenRouter uses `GET /api/v1/credits`; z.ai uses
-`GET /api/monitor/usage/quota/limit`; Codex uses `GET /wham/usage`;
-CommandCode uses `GET /alpha/billing/credits` (a private endpoint its CLI
+`GET /api/monitor/usage/quota/limit`; Codex and OpenAI use `GET /wham/usage`
+(plus `/chatpass/apps` for OpenAI); CommandCode uses
+`GET /alpha/billing/credits` (a private endpoint its CLI
 also uses; `COMMANDCODE_API_URL` can override the base URL). z.ai, Codex, and
-Codex show both each quota's window length and its `nextResetTime`/`reset_at`
+CommandCode show both each quota's window length and its `nextResetTime`/`reset_at`
 countdown (when supplied), and `CODEX_API_URL` or `CHATGPT_BASE_URL` can
 override the Codex base URL.
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "bermudis-pi-goodies",
-  "version": "0.27.6",
+  "version": "0.27.7",
   "repository": {
     "type": "git",
     "url": "git+https://github.com/bermudi/agent-extensions.git",

package/provider-balance.ts

@@ -51,6 +51,14 @@ const CODEX_API_BASE = (
   "https://chatgpt.com/backend-api"
 ).replace(/\/+$/, "");
 const CODEX_USAGE_ENDPOINT = `${CODEX_API_BASE}/wham/usage`;
+/**
+ * App registrations for the ChatGPT account behind a Codex credential. The
+ * new `openai` (Sign in with ChatGPT) token cannot read usage itself — the
+ * backend rejects it with 401 rejected_by_access_enforcement — so usage is
+ * read with a companion `openai-codex` login. Matching the new login's
+ * client_id here proves both logins belong to the same ChatGPT account.
+ */
+const CODEX_CHATPASS_APPS_ENDPOINT = `${CODEX_USAGE_ENDPOINT}/chatpass/apps`;
 const CODEX_AUTH_CLAIM = "https://api.openai.com/auth";
 const COMMANDCODE_API_BASE = (
   process.env.COMMANDCODE_API_URL || "https://api.commandcode.ai"
@@ -66,16 +74,20 @@ const IDLE_REFRESH_INTERVAL_MS = 60_000;
 const IDLE_REFRESH_JITTER_MS = 15_000;
 /** Minimum time between turn_end/delegate:usage-triggered balance refreshes. See throttledRefresh. */
 const TURN_REFRESH_MIN_INTERVAL_MS = 30_000;
+/** How often the identity watchdog compares the live credentials against the
+ *  balance shown in the footer. Pi's built-in /login and /logout commands are
+ *  dispatched by the TUI before extension input events run and fire no
+ *  extension-visible event when a model is already selected, so the only
+ *  supported way to notice a credential swap is re-reading the model registry.
+ *  Each tick costs one or two local credential-store reads (network only when
+ *  an OAuth token is inside its refresh window, which is shared and locked).
+ *  5s keeps the post-login footer lag well under the old 60–75s worst case
+ *  (the next idle poll) without meaningful cost. */
+export const IDENTITY_WATCHDOG_INTERVAL_MS = 5_000;
 /** On session start, adopt a sibling/previous session's cache entry younger
  *  than this instead of refetching — rapid session switches then cost nothing. */
 const SESSION_START_CACHE_ADOPT_MS = 15_000;
-/** How long the post-/login poller waits for the new credential to appear,
- *  at 1s per attempt. Kilo device logins take 30–60s of human time (browser
- *  round-trip), so a ~9s window (the old 10-attempt cap) usually expired
- *  before the user finished and the footer kept the old account's balance
- *  until the next idle poll. 90s covers the human with margin; the poll is a
- *  local keychain/auth-store read, not a network request. */
-export const AUTH_TRANSITION_MAX_ATTEMPTS = 90;
+
 /** After a failed provider fetch, idle polls machine-wide back off for this
  *  long via a shared marker in the balance cache. */
 export const FAILURE_BACKOFF_MS = 60_000;
@@ -110,9 +122,27 @@ const BALANCE_CACHE_MAX_ENTRIES = 256;
 const BALANCE_CACHE_MAX_BYTES = 1_000_000;
 const BALANCE_CACHE_ACCOUNT_DIR_PATTERN = /^[a-f0-9]{64}$/;
 
+export interface BalanceFetchOptions {
+  /**
+   * Credential for `companionProvider`, resolved by refreshBalance from the
+   * same model registry. Never persisted; used only for this request.
+   */
+  companionToken?: string;
+}
+
 export interface BalanceAdapter {
-  fetch(token: string, signal: AbortSignal): Promise<Balance>;
+  fetch(
+    token: string,
+    signal: AbortSignal,
+    options?: BalanceFetchOptions,
+  ): Promise<Balance>;
   requiresOAuth?: boolean;
+  /**
+   * Provider whose login is required to read this balance. `openai` needs an
+   * `openai-codex` companion on the same ChatGPT account because OpenAI
+   * rejects Sign in with ChatGPT tokens on the usage endpoints.
+   */
+  companionProvider?: string;
 }
 
 interface CodexQuotaWindow {
@@ -175,22 +205,53 @@ function decodeBase64Url(value: string): string {
   return new TextDecoder().decode(bytes);
 }
 
-export function parseCodexAccountId(token: string): string | null {
+function decodeJwtPayload(token: string): Record<string, unknown> | null {
   const parts = token.split(".");
   if (parts.length !== 3 || !parts[1]) return null;
 
   try {
-    const payload = asRecord(JSON.parse(decodeBase64Url(parts[1])));
-    const auth = payload ? asRecord(payload[CODEX_AUTH_CLAIM]) : null;
-    const accountId = auth?.chatgpt_account_id;
-    return typeof accountId === "string" && accountId.trim()
-      ? accountId.trim()
-      : null;
+    return asRecord(JSON.parse(decodeBase64Url(parts[1])));
   } catch {
     return null;
   }
 }
 
+export function parseCodexAccountId(token: string): string | null {
+  const payload = decodeJwtPayload(token);
+  const auth = payload ? asRecord(payload[CODEX_AUTH_CLAIM]) : null;
+  const accountId = auth?.chatgpt_account_id;
+  return typeof accountId === "string" && accountId.trim()
+    ? accountId.trim()
+    : null;
+}
+
+/**
+ * Identity of the application registration behind an `openai` OAuth login.
+ * The legacy Codex account lists this ID under /wham/usage/chatpass/apps,
+ * which is how the companion credential is proven to be the same account.
+ */
+export function parseOpenAIClientId(token: string): string | null {
+  const payload = decodeJwtPayload(token);
+  const clientId = payload?.client_id;
+  return typeof clientId === "string" && clientId.trim()
+    ? clientId.trim()
+    : null;
+}
+
+/** Exactly one registration must match: never fall back to another app. */
+export function openAIAppIsRegistered(
+  value: unknown,
+  clientId: string,
+): boolean {
+  const payload = asRecord(value);
+  const items = Array.isArray(payload?.items) ? payload.items : [];
+  let matches = 0;
+  for (const candidate of items) {
+    if (stringProperty(candidate, "id") === clientId) matches++;
+  }
+  return matches === 1;
+}
+
 function parseCodexQuotaWindow(value: unknown): CodexQuotaWindow | null {
   const usedPercent = numericProperty(value, "used_percent");
   const windowSeconds = numericProperty(value, "limit_window_seconds");
@@ -586,23 +647,26 @@ async function fetchZaiCodingCnQuota(
   return fetchZaiQuotaAt(ZAI_CODING_CN_QUOTA_ENDPOINT, token, signal);
 }
 
-async function fetchCodexQuota(
+function codexUsageHeaders(
   token: string,
+  accountId: string,
+): Record<string, string> {
+  return {
+    Authorization: `Bearer ${token}`,
+    Accept: "application/json",
+    "User-Agent": "pi",
+    originator: "pi",
+    "chatgpt-account-id": accountId,
+  };
+}
+
+async function fetchCodexQuotaAt(
+  token: string,
+  accountId: string,
   signal: AbortSignal,
 ): Promise<Balance> {
-  const accountId = parseCodexAccountId(token);
-  if (!accountId) {
-    throw new Error("Codex access token did not contain an account ID");
-  }
-
   const response = await fetch(CODEX_USAGE_ENDPOINT, {
-    headers: {
-      Authorization: `Bearer ${token}`,
-      Accept: "application/json",
-      "User-Agent": "pi",
-      originator: "pi",
-      "chatgpt-account-id": accountId,
-    },
+    headers: codexUsageHeaders(token, accountId),
     signal: timeoutSignal(BALANCE_FETCH_TIMEOUT_MS, signal),
   });
   if (!response.ok) {
@@ -616,6 +680,61 @@ async function fetchCodexQuota(
   return codexQuotaToBalance(quota);
 }
 
+async function fetchCodexQuota(
+  token: string,
+  signal: AbortSignal,
+): Promise<Balance> {
+  const accountId = parseCodexAccountId(token);
+  if (!accountId) {
+    throw new Error("Codex access token did not contain an account ID");
+  }
+  return fetchCodexQuotaAt(token, accountId, signal);
+}
+
+/**
+ * `openai` (Sign in with ChatGPT) tokens are rejected by ChatGPT's usage
+ * backend. Read the plan windows with the companion `openai-codex` login
+ * instead, but only after /chatpass/apps confirms that the active app
+ * registration belongs to that Codex account — otherwise the footer would
+ * show a different account's quota.
+ */
+export async function fetchOpenAIQuota(
+  token: string,
+  signal: AbortSignal,
+  options?: BalanceFetchOptions,
+): Promise<Balance> {
+  const companionToken = options?.companionToken;
+  if (!companionToken) {
+    throw new Error("OpenAI quota requires an openai-codex companion login");
+  }
+
+  const clientId = parseOpenAIClientId(token);
+  if (!clientId) {
+    throw new Error("OpenAI access token did not contain an application ID");
+  }
+  const accountId = parseCodexAccountId(companionToken);
+  if (!accountId) {
+    throw new Error("OpenAI Codex access token did not contain an account ID");
+  }
+
+  const registrationResponse = await fetch(CODEX_CHATPASS_APPS_ENDPOINT, {
+    headers: codexUsageHeaders(companionToken, accountId),
+    signal: timeoutSignal(BALANCE_FETCH_TIMEOUT_MS, signal),
+  });
+  if (!registrationResponse.ok) {
+    throw new Error(
+      `OpenAI app registration request failed: ${registrationResponse.status}`,
+    );
+  }
+  if (!openAIAppIsRegistered(await registrationResponse.json(), clientId)) {
+    throw new Error(
+      "OpenAI app registration was not found in the companion Codex account; sign in to openai and openai-codex with the same ChatGPT account",
+    );
+  }
+
+  return fetchCodexQuotaAt(companionToken, accountId, signal);
+}
+
 // --- CommandCode -----------------------------------------------------------
 // CommandCode documents only chat/messages/models endpoints, but its CLI
 // drives `/usage` from private alpha endpoints on the same host with the
@@ -742,6 +861,11 @@ const BALANCE_ADAPTERS: Readonly<Record<string, BalanceAdapter>> = {
   zai: { fetch: fetchZaiQuota },
   "zai-coding-cn": { fetch: fetchZaiCodingCnQuota },
   "openai-codex": { fetch: fetchCodexQuota, requiresOAuth: true },
+  openai: {
+    fetch: fetchOpenAIQuota,
+    requiresOAuth: true,
+    companionProvider: "openai-codex",
+  },
   // Both CommandCode transports share one account and one key.
   commandcode: { fetch: fetchCommandcodeBalance },
   "commandcode-anthropic": { fetch: fetchCommandcodeBalance },
@@ -758,17 +882,29 @@ interface BalanceCacheEntry {
  * Keep accounts isolated without persisting the credential itself. Sessions
  * using the same credential get the same cache key and can share a reading.
  */
-export function balanceCacheKey(provider: string, token: string): string {
+export function balanceCacheKey(
+  provider: string,
+  token: string,
+  companionToken?: string,
+): string {
   // Codex access tokens rotate, but their account ID is stable. Other
   // providers expose no account identifier here, so the token is the best
   // available identity. Include the endpoint because custom backends can use
   // overlapping account IDs while reporting unrelated balances.
+  //
+  // `openai` balances belong to the companion Codex account but are scoped to
+  // the active app registration, so key on both identities: a different Codex
+  // account or a different `openai` login must not share a reading.
   const identity =
-    provider === "openai-codex" ? (parseCodexAccountId(token) ?? token) : token;
+    provider === "openai-codex"
+      ? (parseCodexAccountId(token) ?? token)
+      : provider === "openai" && companionToken
+        ? `${parseCodexAccountId(companionToken) ?? companionToken}:${parseOpenAIClientId(token) ?? token}`
+        : token;
   const endpoint =
     provider === "kilo"
       ? KILO_BALANCE_ENDPOINT
-      : provider === "openai-codex"
+      : provider === "openai-codex" || provider === "openai"
         ? CODEX_USAGE_ENDPOINT
         : provider === "openrouter"
           ? OPENROUTER_CREDITS_ENDPOINT
@@ -1408,13 +1544,21 @@ export default function providerBalance(
   /** Provider and account the currently displayed balance belongs to. */
   let displayedProvider: string | undefined;
   let displayedCacheKey: string | undefined;
+  /**
+   * Cache key of the companion credential backing the displayed balance
+   * (`openai` -> `openai-codex`). Lets `/login openai-codex` wait for the new
+   * companion before refreshing the active `openai` row.
+   */
+  let displayedCompanionCacheKey: string | undefined;
+  /** Providers already told (once per session) that a companion login is missing. */
+  const companionWarningsLogged = new Set<string>();
   let refreshGeneration = 0;
   let refreshInFlight = false;
   let refreshController: AbortController | undefined;
   /** Stamp shared by every throttled refresh trigger (turn_end, delegate:usage). */
   let lastThrottledRefreshAt: number | undefined;
   let idleRefreshTimer: ReturnType<typeof setTimeout> | undefined;
-  let authTransitionTimer: ReturnType<typeof setTimeout> | undefined;
+  let identityWatchdogTimer: ReturnType<typeof setTimeout> | undefined;
   let requestRender: (() => void) | undefined;
   /** True while our footer component is the one mounted in the TUI. */
   let footerInstalled = false;
@@ -1430,68 +1574,104 @@ export default function providerBalance(
     requestRender?.();
   }
 
-  /** Login/logout is delivered as input before the command changes auth. */
-  function invalidateAuthTransition(provider: string): void {
-    if (provider !== displayedProvider) return;
-    refreshGeneration++;
-    refreshInFlight = false;
-    refreshController?.abort();
-    refreshController = undefined;
-    displayedCacheKey = undefined;
-    clearBalance();
+  /**
+   * Missing companion is an expected state, not a fetch failure: no balance is
+   * shown and no failure marker is written. Record it once per session so the
+   * reason survives outside the footer (which simply stays blank).
+   */
+  function noteCompanionUnavailable(
+    providerId: string,
+    companionProvider: string,
+  ): void {
+    if (companionWarningsLogged.has(providerId)) return;
+    companionWarningsLogged.add(providerId);
+    logGoodiesEvent({
+      type: "provider_balance_unavailable",
+      provider: providerId,
+      companion: companionProvider,
+      message: `${providerId} ChatGPT usage needs a companion ${companionProvider} login on the same account; run /login ${companionProvider} with that account.`,
+    });
   }
 
-  function scheduleAuthTransitionCheck(
-    ctx: ExtensionContext,
-    provider: string,
-    previousCacheKey: string | undefined,
-  ): void {
-    if (authTransitionTimer !== undefined) clearTimer(authTransitionTimer);
-    let attempts = 0;
-    // Deliberately not gated on refreshGeneration: any unrelated refresh
-    // (turn_end, agent_settled, an idle poll) landing while the user is still
-    // completing /login would otherwise kill the poller, and that refresh
-    // resolved the OLD credential before it started, so nothing would pick up
-    // the new one. Only a session switch (activeContext) ends the poll early.
-    const check = async (): Promise<void> => {
-      attempts++;
-      if (activeContext !== ctx) return;
-      let token: string | undefined;
-      try {
-        token = await ctx.modelRegistry.getApiKeyForProvider(provider);
-      } catch (error) {
-        // pi's getApiKeyForProvider swallows keychain errors and resolves to
-        // undefined, so this catch only fires if that contract changes (the
-        // test double rejects to exercise it). Log the failure — a silent
-        // catch here would hide exactly the breakage that stalls the
-        // post-/login footer — while keeping the retry-until-cap flow.
-        reportFailure(
-          "provider_balance_warning",
-          `[provider-balance] keychain read for "${provider}" failed while waiting for the /login credential swap (attempt ${attempts}/${AUTH_TRANSITION_MAX_ATTEMPTS}): ${describeError(error)}`,
-        );
-        if (attempts < AUTH_TRANSITION_MAX_ATTEMPTS) {
-          authTransitionTimer = setTimer(() => void check(), 1_000);
-        } else {
-          authTransitionTimer = undefined;
-        }
-        return;
+  /**
+   * Periodically compare the LIVE credentials against the account the shown
+   * balance was fetched for, and refresh on any mismatch.
+   *
+   * This replaced an `input`-event hook for /login|/logout: pi's built-in
+   * command dispatcher (interactive-mode setupEditorSubmitHandler) handles
+   * both and returns before extension input handlers run, and the login
+   * completion path emits no extension-visible event when a model is already
+   * selected. The hook could therefore never fire from the TUI. Re-reading
+   * the model registry is a supported API, catches every credential-change
+   * path (TUI /login, /logout, another session's login on the same machine,
+   * a rotated token), and needs no assumptions about command routing.
+   */
+  function scheduleIdentityWatchdog(): void {
+    if (identityWatchdogTimer !== undefined) clearTimer(identityWatchdogTimer);
+    identityWatchdogTimer = setTimer(async () => {
+      identityWatchdogTimer = undefined;
+      await checkDisplayedIdentity();
+      if (activeContext !== undefined) scheduleIdentityWatchdog();
+    }, IDENTITY_WATCHDOG_INTERVAL_MS);
+  }
+
+  async function checkDisplayedIdentity(): Promise<void> {
+    const ctx = activeContext;
+    if (ctx === undefined || ctx.mode !== "tui") return;
+    const provider = displayedProvider;
+    if (provider === undefined) return;
+    const adapter = adapters[provider];
+    if (!adapter) return;
+    // An in-flight refresh re-resolves identity after its fetch (and its
+    // mismatch branch triggers a follow-up refresh), so it already covers any
+    // swap this tick could have seen.
+    if (refreshInFlight) return;
+    let token: string | undefined;
+    let companionToken: string | undefined;
+    try {
+      token = await ctx.modelRegistry.getApiKeyForProvider(provider);
+      if (token && adapter.companionProvider) {
+        companionToken = (
+          await ctx.modelRegistry.getProviderAuth(adapter.companionProvider)
+        )?.auth.apiKey;
       }
-      if (activeContext !== ctx) return;
-      const currentCacheKey = token
-        ? balanceCacheKey(provider, token)
-        : undefined;
-      if (currentCacheKey === previousCacheKey) {
-        if (attempts < AUTH_TRANSITION_MAX_ATTEMPTS) {
-          authTransitionTimer = setTimer(() => void check(), 1_000);
-        } else {
-          authTransitionTimer = undefined;
-        }
-        return;
+    } catch (error) {
+      // Boundary: pi's getApiKeyForProvider swallows keychain errors and
+      // resolves to undefined, so a rejection here means that contract broke
+      // (the test double rejects to exercise it). This runs in a timer
+      // callback — rethrowing would escape as an unhandled rejection — so the
+      // failure is reported to the structured log and the watchdog keeps
+      // ticking; a silent catch would strand the footer on a stale account.
+      reportFailure(
+        "provider_balance_warning",
+        `[provider-balance] identity watchdog read for "${provider}" failed: ${describeError(error)}`,
+      );
+      return;
+    }
+    const model = ctx.model;
+    const oauthUsable =
+      !adapter.requiresOAuth ||
+      (model !== undefined &&
+        model.provider === provider &&
+        ctx.modelRegistry.isUsingOAuth(model));
+    let expectedCacheKey: string | undefined;
+    let expectedCompanionCacheKey: string | undefined;
+    if (token !== undefined && oauthUsable) {
+      expectedCacheKey = balanceCacheKey(provider, token, companionToken);
+      if (adapter.companionProvider && companionToken !== undefined) {
+        expectedCompanionCacheKey = balanceCacheKey(
+          adapter.companionProvider,
+          companionToken,
+        );
       }
-      authTransitionTimer = undefined;
-      void refreshForModel(ctx, ctx.model);
-    };
-    authTransitionTimer = setTimer(() => void check(), 250);
+    }
+    if (
+      expectedCacheKey === displayedCacheKey &&
+      expectedCompanionCacheKey === displayedCompanionCacheKey
+    ) {
+      return;
+    }
+    void refreshForModel(ctx, model);
   }
 
   /** Refresh the balance for whatever model is active. Provider-agnostic: the
@@ -1529,6 +1709,7 @@ export default function providerBalance(
     if (provider !== displayedProvider) {
       displayedProvider = provider;
       displayedCacheKey = undefined;
+      displayedCompanionCacheKey = undefined;
       clearBalance();
     }
 
@@ -1547,6 +1728,7 @@ export default function providerBalance(
           !ctx.modelRegistry.isUsingOAuth(model))
       ) {
         displayedCacheKey = undefined;
+        displayedCompanionCacheKey = undefined;
         clearBalance();
         return;
       }
@@ -1564,6 +1746,7 @@ export default function providerBalance(
       } catch (error) {
         if (generation === refreshGeneration) {
           displayedCacheKey = undefined;
+          displayedCompanionCacheKey = undefined;
           clearBalance();
         }
         throw error;
@@ -1571,11 +1754,43 @@ export default function providerBalance(
       if (generation !== refreshGeneration) return;
       if (!token) {
         displayedCacheKey = undefined;
+        displayedCompanionCacheKey = undefined;
         identityPending = false;
         clearBalance();
         return;
       }
-      const tokenCacheKey = balanceCacheKey(providerId, token);
+      // `openai` usage is read with the legacy `openai-codex` login. Resolve
+      // it through the model registry (not auth.json) so its token is
+      // refreshed and scoped exactly like a normal request.
+      let companionToken: string | undefined;
+      if (adapter.companionProvider) {
+        try {
+          companionToken = (
+            await ctx.modelRegistry.getProviderAuth(adapter.companionProvider)
+          )?.auth.apiKey;
+        } catch (error) {
+          if (generation === refreshGeneration) {
+            displayedCacheKey = undefined;
+            displayedCompanionCacheKey = undefined;
+            clearBalance();
+          }
+          throw error;
+        }
+        if (generation !== refreshGeneration) return;
+        if (!companionToken) {
+          displayedCacheKey = undefined;
+          displayedCompanionCacheKey = undefined;
+          identityPending = false;
+          clearBalance();
+          noteCompanionUnavailable(providerId, adapter.companionProvider);
+          return;
+        }
+        displayedCompanionCacheKey = balanceCacheKey(
+          adapter.companionProvider,
+          companionToken,
+        );
+      }
+      const tokenCacheKey = balanceCacheKey(providerId, token, companionToken);
       cacheKey = tokenCacheKey;
       if (displayedCacheKey !== tokenCacheKey) clearBalance();
       displayedCacheKey = tokenCacheKey;
@@ -1609,7 +1824,11 @@ export default function providerBalance(
         }
       }
 
-      const nextBalance = await adapter.fetch(token, controller.signal);
+      const nextBalance = await adapter.fetch(
+        token,
+        controller.signal,
+        companionToken ? { companionToken } : undefined,
+      );
       const observedAtMs = now();
       if (generation !== refreshGeneration) return;
 
@@ -1621,6 +1840,7 @@ export default function providerBalance(
         (!model || !ctx.modelRegistry.isUsingOAuth(model))
       ) {
         displayedCacheKey = undefined;
+        displayedCompanionCacheKey = undefined;
         clearBalance();
         return;
       }
@@ -1630,17 +1850,31 @@ export default function providerBalance(
       // flicker proportional to keychain latency. The re-check only needs to
       // clear the display if it FAILS (handled below); a passing check just
       // commits the fresh balance on top of the old one.
+      // Re-read every credential the cache key covers before accepting the
+      // fetch, and read the ACTIVE token LAST — twice. Any credential change
+      // must land in one of the awaits below to go undetected: before the
+      // companion read (mismatched composite key), or between the two active
+      // reads (double read disagrees). After the second active read resolves,
+      // the remaining checks and the commit are synchronous, so there is no
+      // window left in which a swap could slip through and paint the previous
+      // account's quota. Reading the companion after the active token (as this
+      // block once did) reopened exactly such a window after the double read.
       let currentToken: string | undefined;
       let confirmedToken: string | undefined;
+      let companionTokenAfter: string | undefined;
       try {
+        if (adapter.companionProvider) {
+          companionTokenAfter = (
+            await ctx.modelRegistry.getProviderAuth(adapter.companionProvider)
+          )?.auth.apiKey;
+        }
         currentToken = await ctx.modelRegistry.getApiKeyForProvider(providerId);
-        // A credential can change while the first final lookup is pending.
-        // Resolve it once more before accepting the provider response.
         confirmedToken =
           await ctx.modelRegistry.getApiKeyForProvider(providerId);
       } catch (error) {
         if (generation === refreshGeneration) {
           displayedCacheKey = undefined;
+          displayedCompanionCacheKey = undefined;
           clearBalance();
         }
         throw error;
@@ -1653,11 +1887,14 @@ export default function providerBalance(
           (!currentModel || !ctx.modelRegistry.isUsingOAuth(currentModel))) ||
         !currentToken ||
         !confirmedToken ||
-        balanceCacheKey(providerId, currentToken) !== cacheKey ||
-        balanceCacheKey(providerId, confirmedToken) !== cacheKey
+        balanceCacheKey(providerId, currentToken, companionTokenAfter) !==
+          cacheKey ||
+        balanceCacheKey(providerId, confirmedToken, companionTokenAfter) !==
+          cacheKey
       ) {
         if (generation === refreshGeneration) {
           displayedCacheKey = undefined;
+          displayedCompanionCacheKey = undefined;
           clearBalance();
           if (currentToken || confirmedToken) {
             void refreshForModel(ctx, currentModel);
@@ -1667,6 +1904,12 @@ export default function providerBalance(
       }
 
       identityPending = false;
+      if (adapter.companionProvider && companionTokenAfter) {
+        displayedCompanionCacheKey = balanceCacheKey(
+          adapter.companionProvider,
+          companionTokenAfter,
+        );
+      }
       // Recovery detected: clear the marker so sibling idle polls resume.
       clearBalanceFailureMarker(cacheKey, cacheDir);
       const persisted = writeCachedBalance(
@@ -1827,9 +2070,13 @@ export default function providerBalance(
   pi.on("session_start", (_event, ctx) => {
     activeContext = ctx;
     activeThinkingLevel = restoredThinkingLevel(ctx);
+    companionWarningsLogged.clear();
     installFooter(ctx);
     cleanupBalanceCache(cacheDir, now());
-    if (ctx.mode === "tui") scheduleIdleRefresh();
+    if (ctx.mode === "tui") {
+      scheduleIdleRefresh();
+      scheduleIdentityWatchdog();
+    }
     // Footer data is supplemental. Never hold up session readiness on network.
     // Adopt a seconds-old sibling/previous reading instead of refetching, so
     // rapid session switches and staggered process starts don't stack
@@ -1839,15 +2086,12 @@ export default function providerBalance(
     });
   });
 
-  pi.on("input", (event, ctx) => {
-    const match = /^\/(login|logout)(?:\s+(\S+))?$/.exec(event.text.trim());
-    if (!match) return;
-    const provider = match[2] ?? ctx.model?.provider;
-    if (!provider || provider !== ctx.model?.provider) return;
-    const previousCacheKey = displayedCacheKey;
-    invalidateAuthTransition(provider);
-    scheduleAuthTransitionCheck(ctx, provider, previousCacheKey);
-  });
+  // NOTE: there is deliberately no `input` handler here for /login or
+  // /logout. Pi's TUI dispatches those built-in commands directly in the
+  // editor submit handler and returns before extension input events are
+  // emitted, so such a hook can never fire. Credential changes (either login
+  // of the active provider or of the openai-codex companion) are picked up by
+  // the identity watchdog instead.
 
   pi.on("model_select", (event, ctx) => {
     activeContext = ctx;
@@ -1936,8 +2180,8 @@ export default function providerBalance(
     lastThrottledRefreshAt = undefined;
     if (idleRefreshTimer !== undefined) clearTimer(idleRefreshTimer);
     idleRefreshTimer = undefined;
-    if (authTransitionTimer !== undefined) clearTimer(authTransitionTimer);
-    authTransitionTimer = undefined;
+    if (identityWatchdogTimer !== undefined) clearTimer(identityWatchdogTimer);
+    identityWatchdogTimer = undefined;
     activeContext = undefined;
     activeThinkingLevel = "off";
     requestRender = undefined;
@@ -1947,5 +2191,6 @@ export default function providerBalance(
     balanceFetchedAt = undefined;
     displayedProvider = undefined;
     displayedCacheKey = undefined;
+    displayedCompanionCacheKey = undefined;
   });
 }