@cortexkit/opencode-antigravity-auth 2.0.0 → 2.1.0

package/src/plugin/command-data.ts

@@ -40,10 +40,14 @@
  *   wrong account.
  */
 
+import { createHash } from 'node:crypto'
+
 import {
   buildSidebarMachineStateFromAccounts,
+  normalizeLegacyCachedQuota,
   type SidebarAccountRedactionInput,
   setSidebarMachineState,
+  toCapturedTier,
 } from '../sidebar-state'
 
 /**
@@ -70,19 +74,21 @@ type CommandDataAccountMetadata = {
     Record<CommandDataQuotaGroup, CommandDataQuotaGroupSummary>
   >
   cachedQuotaUpdatedAt?: number
+  cachedQuotaAccountId?: string
   accountIneligible?: boolean
 }
 
-type CommandDataQuotaGroup =
-  | 'claude'
-  | 'gemini-pro'
-  | 'gemini-flash'
-  | 'gpt-oss'
+type CommandDataQuotaGroup = 'gemini' | 'non-gemini'
 
 type CommandDataQuotaGroupSummary = {
   remainingFraction?: number
   resetTime?: string
   modelCount: number
+  windows?: Array<{
+    window: 'weekly' | '5h'
+    remainingFraction: number
+    resetTime: string
+  }>
 }
 
 type CommandDataAccountQuotaResult = {
@@ -105,7 +111,7 @@ type CommandDataAccountQuotaResult = {
   updatedAccount?: CommandDataAccountMetadata
 }
 
-type CommandDataAccountStorage = {
+export interface CommandDataAccountStorage {
   version: 4
   activeIndex: number
   activeIndexByFamily?: { claude?: number; gemini?: number }
@@ -133,12 +139,27 @@ export interface CommandAccountRow {
   enabled: boolean
   /** `true` when this row matches the harness-active account. */
   current: boolean
+  /**
+   * Absolute timestamp (ms) until which this account is rate-limited.
+   * Carried on the row so the sidebar refresher can match cooldowns
+   * without relying on unstable numeric indices.
+   */
+  coolingDownUntil?: number
+  /** Health score in [0, 100], forwarded from the live tracker. */
+  healthScore?: number
   quota: Array<{
-    key: 'claude' | 'gemini-pro' | 'gemini-flash' | 'gpt-oss'
+    key: 'gemini' | 'non-gemini'
     label: string
     remainingPercent: number | null
     resetAt?: number
+    windows?: Array<{
+      window: 'weekly' | '5h'
+      remainingPercent: number
+      resetAt?: number
+    }>
   }>
+  /** Captured plan tier. Absent when unknown; never defaulted. */
+  tier?: { id: string; capturedAt: number }
 }
 
 /**
@@ -149,17 +170,13 @@ const QUOTA_GROUP_LABELS: Record<
   CommandAccountRow['quota'][number]['key'],
   string
 > = {
-  claude: 'Claude',
-  'gemini-pro': 'Gemini Pro',
-  'gemini-flash': 'Gemini Flash',
-  'gpt-oss': 'GPT-OSS',
+  gemini: 'Gemini',
+  'non-gemini': 'Non-Gemini',
 }
 
 const SUPPORTED_QUOTA_KEYS = [
-  'claude',
-  'gemini-pro',
-  'gemini-flash',
-  'gpt-oss',
+  'gemini',
+  'non-gemini',
 ] as const satisfies readonly CommandAccountRow['quota'][number]['key'][]
 
 interface LiveAccountSnapshot {
@@ -172,11 +189,36 @@ interface LiveAccountSnapshot {
     Record<CommandDataQuotaGroup, CommandDataQuotaGroupSummary>
   >
   cachedQuotaUpdatedAt?: number
+  cachedQuotaAccountId?: string
   accountIneligible?: boolean
+  coolingDownUntil?: number
+  healthScore?: number
+  /** Captured plan tier from the most recent loadCodeAssist response. */
+  capturedTierId?: string
+  /** Captured paid-tier ID from the most recent loadCodeAssist response. */
+  capturedPaidTierId?: string
+  /** Epoch ms when capturedTierId was last recorded. */
+  capturedTierAt?: number
 }
 
 function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow {
-  const cached = entry.cachedQuota
+  // Stamp mismatch: the cached quota was captured for a different account
+  // (the refresh token changed, or an index shift placed another account's
+  // snapshot at this position). Drop the stale cache rather than rendering
+  // the wrong account's quota percentages.
+  const rawCached =
+    entry.cachedQuotaAccountId &&
+    entry.cachedQuotaAccountId !== quotaAccountIdentity(entry.refreshToken)
+      ? undefined
+      : entry.cachedQuota
+  // Normalize legacy pool keys (gemini-pro, gemini-flash, claude, gpt-oss)
+  // to canonical keys before the SUPPORTED_QUOTA_KEYS loop. Without this
+  // a legacy snapshot produces empty quota rows, and writeSidebar would
+  // permanently re-emit legacy keys on every dialog open until the next
+  // real quota refresh rewrites the on-disk snapshot.
+  const cached = normalizeLegacyCachedQuota(
+    rawCached as Parameters<typeof normalizeLegacyCachedQuota>[0],
+  )
   const quota: CommandAccountRow['quota'] = []
   for (const key of SUPPORTED_QUOTA_KEYS) {
     const cachedEntry = cached?.[key]
@@ -194,11 +236,27 @@ function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow {
       const parsed = Date.parse(cachedEntry.resetTime)
       if (Number.isFinite(parsed)) resetAt = parsed
     }
+    // Carry per-window breakdown when the cache was populated from the
+    // windowed RUQS path — the sidebar writer projects through here.
+    const windows = cachedEntry.windows?.length
+      ? cachedEntry.windows.map((w) => ({
+          window: w.window,
+          remainingPercent: Math.round(w.remainingFraction * 100),
+          resetAt:
+            typeof w.resetTime === 'string' && w.resetTime.length > 0
+              ? (() => {
+                  const parsed = Date.parse(w.resetTime)
+                  return Number.isFinite(parsed) ? parsed : undefined
+                })()
+              : undefined,
+        }))
+      : undefined
     quota.push({
       key,
       label: QUOTA_GROUP_LABELS[key],
       remainingPercent,
       resetAt,
+      windows,
     })
   }
   const label = `Account ${entry.index + 1}`
@@ -208,10 +266,41 @@ function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow {
     label,
     enabled: entry.enabled,
     current: entry.active,
+    coolingDownUntil: entry.coolingDownUntil,
+    healthScore: entry.healthScore,
     quota,
+    tier: toCapturedTier(entry),
   }
 }
 
+/**
+ * Opaque identity derived from a refresh token. Antigravity refresh tokens
+ * are stable (they do not rotate), so this hash is a durable, prunable
+ * identity to detect a stale cached quota after an account-index shift.
+ */
+function quotaAccountIdentity(refreshToken: string): string {
+  return createHash('sha256').update(refreshToken).digest('hex').slice(0, 16)
+}
+
+export function projectCommandAccountRows(
+  storage: CommandDataAccountStorage | null | undefined,
+): CommandAccountRow[] {
+  if (!storage) return []
+  const activeIndex = storage.activeIndexByFamily?.claude ?? storage.activeIndex
+  return storage.accounts.map((entry, index) =>
+    toCommandAccountRow({
+      index,
+      refreshToken: entry.refreshToken,
+      label: entry.label,
+      enabled: entry.enabled !== false,
+      active: index === activeIndex,
+      cachedQuota: entry.cachedQuota,
+      cachedQuotaUpdatedAt: entry.cachedQuotaUpdatedAt,
+      cachedQuotaAccountId: entry.cachedQuotaAccountId,
+    }),
+  )
+}
+
 /**
  * Live AccountManager view the data service needs.
  *
@@ -234,8 +323,6 @@ function toCommandAccountRow(entry: LiveAccountSnapshot): CommandAccountRow {
  *   dialog-triggered mutation lands on disk before the dialog's response
  *   returns. Without it, the dialog could toast "Account removed" while
  *   the file still carries the old pool.
- * - `activeIndex()` returns the position the harness considers active;
- *   the service uses it to mark the matching row as `current: true`.
  */
 export interface CommandDataAccountManagerView {
   getAccounts(): LiveAccountSnapshot[]
@@ -245,10 +332,17 @@ export interface CommandDataAccountManagerView {
     groups: Partial<
       Record<CommandDataQuotaGroup, CommandDataQuotaGroupSummary>
     >,
+    expectedRefreshToken?: string,
   ): void
   requestSaveToDisk(): void
   flushSaveToDisk(): Promise<void>
-  activeIndex(): number
+  /**
+   * Per-family live cursor. Used by the remove path so the persisted
+   * `activeIndexByFamily` follows each model's currently-active account
+   * independently — collapsing both families to a single index would
+   * silently unelect one family on every restart.
+   */
+  getActiveIndexByFamily(): { claude: number; gemini: number }
   /** Enable/disable the account at `index`. Returns true when it changed. */
   setAccountEnabled(index: number, enabled: boolean): boolean
   /** Pin `index` as the active account for every family the dialog cares about. */
@@ -328,8 +422,20 @@ export interface CommandDataService {
    * refresh token, bump `cachedQuotaUpdatedAt`, and push a label-only
    * sidebar snapshot. Returns the freshly persisted rows so the
    * dialog can re-render in place.
+   *
+   * Bypass flag is intentional: explicit user-triggered refreshes
+   * ("Refresh" button, direct slash-command with `refresh` argument)
+   * must always fetch, even if the quota manager has backed off.
    */
   refreshQuota(): Promise<CommandAccountRow[]>
+  /**
+   * Non-forced quota check that RESPECTS the quota manager's per-account
+   * backoff. Used by the dialog OPEN path — which is a VIEW, not a user
+   * request to fetch. Accounts already fresh (or in backoff) are skipped;
+   * the result is discarded (the sidebar poller carries the freshness
+   * guarantee for idle accounts).
+   */
+  refreshQuotaRespectingBackoff(): Promise<void>
   /**
    * Pin `index` as the active account for every family the dialog
    * tracks (claude + gemini). Mutates the live AccountManager AND the
@@ -381,39 +487,70 @@ export function createCommandDataService(
   const projectRows = (): CommandAccountRow[] =>
     accountManagerView.getAccounts().map(toCommandAccountRow)
 
-  const writeSidebar = (rows: CommandAccountRow[]): void => {
-    const accounts: SidebarAccountRedactionInput[] = rows.map((row) => {
-      const claude = row.quota.find((q) => q.key === 'claude')
-      const geminiPro = row.quota.find((q) => q.key === 'gemini-pro')
-      const geminiFlash = row.quota.find((q) => q.key === 'gemini-flash')
-      const gptOss = row.quota.find((q) => q.key === 'gpt-oss')
-      const toFraction = (
-        q: { remainingPercent: number | null } | undefined,
-      ): { remainingFraction?: number; resetTime?: string } | undefined => {
-        if (!q || q.remainingPercent == null) return undefined
-        return { remainingFraction: q.remainingPercent / 100 }
-      }
+  // Map a CommandAccountRow back to the SidebarAccountRedactionInput shape.
+  // Extracted to prevent the "sixth dropped field" pattern: every field in
+  // CommandAccountRow that must reach the sidebar lives here once, not as an
+  // inline literal that a future author has to remember to update in two places.
+  const toRedactionInput = (
+    row: CommandAccountRow,
+  ): SidebarAccountRedactionInput => {
+    const gemini = row.quota.find((q) => q.key === 'gemini')
+    const nonGemini = row.quota.find((q) => q.key === 'non-gemini')
+    // Map a CommandAccountRow quota entry to cachedQuota pool shape
+    // so projectQuotaPoolForSidebar (the canonical projection seam)
+    // carries the per-window breakdown into the sidebar state.
+    const toPool = (
+      q:
+        | {
+            remainingPercent: number | null
+            resetAt?: number
+            windows?: CommandAccountRow['quota'][number]['windows']
+          }
+        | undefined,
+    ) => {
+      if (!q || q.remainingPercent == null) return undefined
       return {
-        index: row.index,
-        label: row.label,
-        enabled: row.enabled,
-        current: row.current,
-        cachedQuota: {
-          claude: toFraction(claude),
-          'gemini-pro': toFraction(geminiPro),
-          'gemini-flash': toFraction(geminiFlash),
-          'gpt-oss': toFraction(gptOss),
-        },
+        remainingFraction: q.remainingPercent / 100,
+        resetTime:
+          typeof q.resetAt === 'number' && Number.isFinite(q.resetAt)
+            ? new Date(q.resetAt).toISOString()
+            : undefined,
+        windows: q.windows?.map((w) => ({
+          window: w.window,
+          remainingFraction: w.remainingPercent / 100,
+          resetTime:
+            typeof w.resetAt === 'number' && Number.isFinite(w.resetAt)
+              ? new Date(w.resetAt).toISOString()
+              : '',
+        })),
       }
-    })
-    // Fire-and-forget — the sidebar writer is fenced by its own queue,
+    }
+    return {
+      index: row.index,
+      label: row.label,
+      enabled: row.enabled,
+      current: row.current,
+      healthScore: row.healthScore,
+      // Tier passes the redaction boundary unchanged (plan metadata, not PII).
+      // The stamp check was already done by toCommandAccountRow.
+      tier: row.tier,
+      cachedQuota: {
+        gemini: toPool(gemini),
+        'non-gemini': toPool(nonGemini),
+      },
+    }
+  }
+
+  const writeSidebar = (rows: CommandAccountRow[]): void => {
+    const accounts = rows.map(toRedactionInput)
+    // Fire-and-forget -- the sidebar writer is fenced by its own queue,
     // so a transient lock contention cannot block the dialog response.
     void setSidebarMachineState(
       buildSidebarMachineStateFromAccounts(accounts, { checkedAt: now() }),
       { stateFile: sidebarStateFile },
     ).catch(() => {
-      // Lock contention is the only realistic failure mode; the next
-      // periodic writer will catch up.
+      // Sidebar writes are best-effort; the next command or quota refresh
+      // will publish the current snapshot.
     })
   }
 
@@ -466,7 +603,11 @@ export function createCommandDataService(
       for (const update of updates) {
         const liveIndex = liveIndexByRefreshToken.get(update.refreshToken)
         if (liveIndex === undefined || !update.groups) continue
-        accountManagerView.updateQuotaCache(liveIndex, update.groups)
+        accountManagerView.updateQuotaCache(
+          liveIndex,
+          update.groups,
+          update.refreshToken,
+        )
         liveQuotaChanged = true
       }
       if (liveQuotaChanged) accountManagerView.requestSaveToDisk()
@@ -486,6 +627,10 @@ export function createCommandDataService(
               return {
                 ...entry,
                 cachedQuota: update.groups,
+                // Stamp the persisted quota with an opaque identity derived
+                // from the refresh token so a later projection can detect
+                // a stale snapshot after an account-index shift.
+                cachedQuotaAccountId: quotaAccountIdentity(entry.refreshToken),
                 cachedQuotaUpdatedAt: refreshedAt,
               }
             }
@@ -510,6 +655,77 @@ export function createCommandDataService(
       return rows
     },
 
+    async refreshQuotaRespectingBackoff() {
+      const accountsForQuota = accountManagerView.getAccountsForQuotaCheck()
+      if (accountsForQuota.length === 0) return
+
+      // Non-forced: the quota manager's per-account backoff and dedup
+      // apply. Accounts that are fresh or in backoff are skipped. Unlike
+      // the earlier fire-and-forget path, successful results ARE folded
+      // into the live cache so the sidebar snapshot reflects fresh data
+      // immediately — the dialog renders ONLY from the sidebar file, and
+      // the quota manager's push path does not update AccountManager
+      // entries for this caller.
+      let results: Awaited<ReturnType<typeof quotaManager.refreshAccounts>>
+      try {
+        results = await quotaManager.refreshAccounts(accountsForQuota, {
+          indexFor: (account) => accountsForQuota.indexOf(account),
+          force: false,
+        })
+      } catch {
+        // Non-critical: dialog open must never fail because a background
+        // refresh check encountered backoff or a transient error.
+        return
+      }
+
+      // Build updates akin to the forced refresh path — resolve live
+      // indexes after the network call (a concurrent add/remove may have
+      // shifted positions), fold successful results into the live cache,
+      // then push a sidebar snapshot.
+      const updates: Array<{
+        refreshToken: string
+        groups?: Partial<
+          Record<CommandDataQuotaGroup, CommandDataQuotaGroupSummary>
+        >
+      }> = []
+      for (const result of results) {
+        const refreshToken =
+          result.updatedAccount?.refreshToken ??
+          accountsForQuota[result.index]?.refreshToken
+        if (!refreshToken) continue
+        const groups =
+          result.status === 'ok' && result.quota?.groups
+            ? result.quota.groups
+            : undefined
+        updates.push({ refreshToken, groups })
+      }
+
+      if (updates.length === 0) return
+
+      const liveIndexByRefreshToken = new Map<string, number>()
+      for (const entry of accountManagerView.getAccounts()) {
+        liveIndexByRefreshToken.set(entry.refreshToken, entry.index)
+      }
+      let liveQuotaChanged = false
+      for (const update of updates) {
+        const liveIndex = liveIndexByRefreshToken.get(update.refreshToken)
+        if (liveIndex === undefined || !update.groups) continue
+        accountManagerView.updateQuotaCache(
+          liveIndex,
+          update.groups,
+          update.refreshToken,
+        )
+        liveQuotaChanged = true
+      }
+      if (liveQuotaChanged) accountManagerView.requestSaveToDisk()
+
+      // Push a fresh sidebar snapshot — the live view now carries the
+      // just-refreshed quota plus any skipped (backoff/fresh) accounts
+      // whose cached percentages are unchanged.
+      const rows = projectRows()
+      writeSidebar(rows)
+    },
+
     async setCurrentAccount(index) {
       return mutateLiveAndStorage({ action: 'setCurrent', index })
     },
@@ -549,9 +765,32 @@ export function createCommandDataService(
       )
     }
 
+    // Capture the live view's current-account identity per family BEFORE
+    // the storage mutation. The remove action must persist the index that
+    // the same account will occupy AFTER the removal — unconditionally
+    // resetting to 0 made a non-current removal promote whichever account
+    // shifted into slot 0 to "active" on the next restart. Each model's
+    // cursor is tracked independently, so collapse-tokens does not work
+    // here: a Claude-active account and a Gemini-active account can
+    // point at different rows.
+    const liveCurrentTokens: { claude?: string; gemini?: string } = {}
+    if (action === 'remove') {
+      const liveAccounts = accountManagerView.getAccounts()
+      const liveIndexes = accountManagerView.getActiveIndexByFamily()
+      liveCurrentTokens.claude =
+        liveAccounts[liveIndexes.claude]?.refreshToken ?? undefined
+      liveCurrentTokens.gemini =
+        liveAccounts[liveIndexes.gemini]?.refreshToken ?? undefined
+    }
+
     let foundInStorage = false
     let desiredEnabled: boolean | undefined
     let previousEnabled: boolean | undefined
+    let nextActiveIndex = 0
+    let nextActiveIndexByFamily: { claude?: number; gemini?: number } = {
+      claude: 0,
+      gemini: 0,
+    }
     await storage.mutate((current) => {
       const tokenIdx = current.accounts.findIndex(
         (entry) => entry.refreshToken === refreshToken,
@@ -587,13 +826,54 @@ export function createCommandDataService(
         }
       }
 
+      // remove: build the post-removal account list, then resolve the
+      // current-account's NEW index in that list. If the removed
+      // account was the live current, the captured token falls out of
+      // the list — the live AccountManager leaves the numeric current
+      // cursor at the SAME index (now occupied by whichever account
+      // shifted in), so the persisted index must follow that numeric
+      // slot rather than resetting to 0.
+      const nextAccounts = current.accounts.filter(
+        (account) => account.refreshToken !== refreshToken,
+      )
+      const resolveNextIndex = (
+        liveToken: string | undefined,
+        legacyIndex: number,
+      ): number => {
+        if (!liveToken)
+          return Math.max(0, Math.min(legacyIndex, nextAccounts.length - 1))
+        const found = nextAccounts.findIndex(
+          (account) => account.refreshToken === liveToken,
+        )
+        if (found !== -1) return found
+        // The captured current token was the one removed. Mirror the
+        // core's in-memory semantics: keep the cursor at the removed
+        // slot when it still exists in the new array; when the removed
+        // slot was the last position and no longer exists, persist 0
+        // (the core's buildStorageSnapshot clamps negative sentinels
+        // to 0, and auth-doctor treats a negative activeIndex as
+        // corruption — cf. auth-doctor.ts:149-156).
+        //
+        // Use tokenIdx (the storage-snapshot position) rather than the
+        // outer-scope `index` (live-read position): a concurrent mutation
+        // between the live read and the lock can diverge the two, and the
+        // storage-level cursor must match the storage-level array.
+        return tokenIdx < nextAccounts.length ? tokenIdx : 0
+      }
+      const legacyClaude =
+        current.activeIndexByFamily?.claude ?? current.activeIndex
+      const legacyGemini =
+        current.activeIndexByFamily?.gemini ?? current.activeIndex
+      nextActiveIndex = resolveNextIndex(liveCurrentTokens.claude, legacyClaude)
+      nextActiveIndexByFamily = {
+        claude: resolveNextIndex(liveCurrentTokens.claude, legacyClaude),
+        gemini: resolveNextIndex(liveCurrentTokens.gemini, legacyGemini),
+      }
       return {
         ...current,
-        accounts: current.accounts.filter(
-          (account) => account.refreshToken !== refreshToken,
-        ),
-        activeIndex: 0,
-        activeIndexByFamily: { claude: 0, gemini: 0 },
+        accounts: nextAccounts,
+        activeIndex: nextActiveIndex,
+        activeIndexByFamily: nextActiveIndexByFamily,
       }
     })