fix(usage): read Antigravity quota from its own credential, including SSH hosts

src/main/rate-limits/antigravity-cloud-code-api.test.ts

@@ -0,0 +1,178 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const { netFetchMock } = vi.hoisted(() => ({ netFetchMock: vi.fn() }))
+vi.mock('electron', () => ({ net: { fetch: netFetchMock } }))
+
+import { fetchAntigravityQuota } from './antigravity-cloud-code-api'
+
+const cancel = vi.fn()
+
+function makeResponse(status: number, body: unknown): Response {
+  return {
+    ok: status >= 200 && status < 300,
+    status,
+    json: () => Promise.resolve(body),
+    body: { cancel }
+  } as unknown as Response
+}
+
+const QUOTA_BODY = {
+  groups: [
+    {
+      displayName: 'Gemini Models',
+      buckets: [
+        {
+          bucketId: 'gemini-weekly',
+          displayName: 'Weekly Limit Remaining',
+          window: 'weekly',
+          resetTime: '2026-09-11T05:00:13Z',
+          remainingFraction: 0.79675514
+        }
+      ]
+    }
+  ]
+}
+
+function routeOk(url: string): Response {
+  if (url.includes('loadCodeAssist')) {
+    return makeResponse(200, {
+      cloudaicompanionProject: 'proj-1',
+      currentTier: { id: 'free-tier' }
+    })
+  }
+  if (url.includes('retrieveUserQuotaSummary')) {
+    return makeResponse(200, QUOTA_BODY)
+  }
+  // Why: an unrouted URL must fail loudly rather than pass the test silently.
+  return makeResponse(500, {})
+}
+
+function lastInit(match: string): RequestInit {
+  const call = netFetchMock.mock.calls.findLast(([url]) => String(url).includes(match))
+  return (call?.[1] ?? {}) as RequestInit
+}
+
+describe('fetchAntigravityQuota', () => {
+  beforeEach(() => {
+    cancel.mockReset()
+    netFetchMock.mockReset()
+    netFetchMock.mockImplementation((url: string) => Promise.resolve(routeOk(url)))
+  })
+
+  it('identifies itself as Antigravity, which is what unlocks the quota grant', async () => {
+    await fetchAntigravityQuota('token-1')
+    for (const method of ['loadCodeAssist', 'retrieveUserQuotaSummary']) {
+      const headers = lastInit(method).headers as Record<string, string>
+      expect(headers['User-Agent']).toMatch(/antigravity/i)
+    }
+  })
+
+  it('asks for the Antigravity project rather than the Gemini CLI one', async () => {
+    await fetchAntigravityQuota('token-1')
+    expect(JSON.parse(String(lastInit('loadCodeAssist').body))).toEqual({
+      metadata: { ideType: 'ANTIGRAVITY' }
+    })
+  })
+
+  it('returns the grouped windows on success', async () => {
+    const result = await fetchAntigravityQuota('token-1')
+    expect(result).toEqual({ status: 'ok', groups: QUOTA_BODY.groups })
+  })
+
+  it('passes the resolved project to the quota call', async () => {
+    await fetchAntigravityQuota('token-1')
+    expect(JSON.parse(String(lastInit('retrieveUserQuotaSummary').body))).toEqual({
+      project: 'proj-1'
+    })
+  })
+
+  it('reports a missing project as an entitlement gap, not an error', async () => {
+    netFetchMock.mockImplementation((url: string) =>
+      Promise.resolve(
+        url.includes('loadCodeAssist') ? makeResponse(200, { allowedTiers: [] }) : routeOk(url)
+      )
+    )
+    expect(await fetchAntigravityQuota('token-1')).toEqual({ status: 'not-entitled' })
+  })
+
+  it('reports a 403 from the quota call as an entitlement gap', async () => {
+    netFetchMock.mockImplementation((url: string) =>
+      Promise.resolve(
+        url.includes('retrieveUserQuotaSummary') ? makeResponse(403, {}) : routeOk(url)
+      )
+    )
+    expect(await fetchAntigravityQuota('token-1')).toEqual({ status: 'not-entitled' })
+  })
+
+  it('reports a rejected token as unauthorized', async () => {
+    netFetchMock.mockImplementation(() => Promise.resolve(makeResponse(401, {})))
+    expect(await fetchAntigravityQuota('token-1')).toEqual({ status: 'unauthorized' })
+  })
+
+  it('falls back to the stable host when the daily host errors', async () => {
+    netFetchMock.mockImplementation((url: string) =>
+      Promise.resolve(url.includes('daily-') ? makeResponse(500, {}) : routeOk(url))
+    )
+    const result = await fetchAntigravityQuota('token-1')
+    expect(result.status).toBe('ok')
+    expect(netFetchMock.mock.calls.some(([url]) => String(url).includes('daily-'))).toBe(true)
+  })
+
+  it('does not retry the second host on an account-level verdict', async () => {
+    netFetchMock.mockImplementation(() => Promise.resolve(makeResponse(401, {})))
+    await fetchAntigravityQuota('token-1')
+    expect(netFetchMock).toHaveBeenCalledTimes(1)
+  })
+
+  it('cancels response bodies it will not read', async () => {
+    netFetchMock.mockImplementation(() => Promise.resolve(makeResponse(500, {})))
+    await fetchAntigravityQuota('token-1')
+    expect(cancel).toHaveBeenCalled()
+  })
+
+  it('still tries the stable host when the daily request rejects outright', async () => {
+    netFetchMock.mockImplementation((url: string) =>
+      String(url).includes('daily-')
+        ? Promise.reject(new Error('socket hang up'))
+        : Promise.resolve(routeOk(String(url)))
+    )
+    const result = await fetchAntigravityQuota('token-1')
+    expect(result.status).toBe('ok')
+  })
+
+  it('gives each host its own deadline instead of one shared budget', async () => {
+    const signals: (AbortSignal | undefined)[] = []
+    netFetchMock.mockImplementation((url: string, init: RequestInit) => {
+      signals.push(init?.signal ?? undefined)
+      return String(url).includes('daily-')
+        ? Promise.reject(new Error('timed out'))
+        : Promise.resolve(routeOk(String(url)))
+    })
+    await fetchAntigravityQuota('token-1')
+    // Why: a shared deadline would hand the fallback the same, already-spent signal.
+    expect(signals[0]).not.toBe(signals.at(-1))
+  })
+
+  it('reports the failing host when every request rejects', async () => {
+    netFetchMock.mockImplementation(() => Promise.reject(new Error('ENOTFOUND')))
+    const result = await fetchAntigravityQuota('token-1')
+    expect(result).toMatchObject({ status: 'error' })
+    expect(result.status === 'error' && result.message).toContain('cloudcode-pa.googleapis.com')
+  })
+
+  it('stops before the second host when the caller aborts', async () => {
+    const controller = new AbortController()
+    netFetchMock.mockImplementation(() => {
+      controller.abort()
+      return Promise.reject(new Error('aborted'))
+    })
+    await fetchAntigravityQuota('token-1', controller.signal)
+    expect(netFetchMock).toHaveBeenCalledTimes(1)
+  })
+
+  it('sends the caller abort signal alongside its own timeout', async () => {
+    const controller = new AbortController()
+    await fetchAntigravityQuota('token-1', controller.signal)
+    expect(lastInit('loadCodeAssist').signal).toBeInstanceOf(AbortSignal)
+  })
+})

src/main/rate-limits/antigravity-cloud-code-api.ts

@@ -0,0 +1,163 @@
+/**
+ * Two-step Cloud Code call that returns Antigravity's own quota pools.
+ *
+ * `loadCodeAssist` resolves the project the quota is billed against, and
+ * `retrieveUserQuotaSummary` returns the grouped windows Antigravity shows in its
+ * own usage panel. Both are undocumented `v1internal` endpoints.
+ */
+
+import { net } from 'electron'
+import { cancelUnreadResponseBody } from '../lib/unread-response-body'
+
+const API_TIMEOUT_MS = 10_000
+
+// Why: cloudcode-pa selects the product from the User-Agent, not from the request body. With a
+// UA that does not contain "antigravity" the call still answers 200 but omits
+// cloudaicompanionProject, and the quota step then fails 403 — indistinguishable from a bad
+// sign-in. Electron's default Chromium UA resolves to a different product with no quota grant.
+const USER_AGENT = 'orca-antigravity-usage/1.0'
+
+// Why: Antigravity ships against the "daily" host; stable builds answer on the plain one.
+const API_HOSTS = ['daily-cloudcode-pa.googleapis.com', 'cloudcode-pa.googleapis.com'] as const
+
+export type AntigravityQuotaBucket = {
+  displayName: string
+  window: string
+  resetTime: string
+  remainingFraction: number
+  description?: string
+}
+
+export type AntigravityQuotaGroup = {
+  displayName: string
+  buckets: AntigravityQuotaBucket[]
+}
+
+export type AntigravityQuotaFetch =
+  | { status: 'ok'; groups: AntigravityQuotaGroup[] }
+  /** The credential authenticated but carries no Antigravity entitlement. */
+  | { status: 'not-entitled' }
+  | { status: 'unauthorized' }
+  | { status: 'error'; message: string }
+
+function isBucket(value: unknown): value is AntigravityQuotaBucket {
+  const b = value as AntigravityQuotaBucket
+  return (
+    typeof value === 'object' &&
+    value !== null &&
+    typeof b.displayName === 'string' &&
+    typeof b.window === 'string' &&
+    typeof b.resetTime === 'string' &&
+    typeof b.remainingFraction === 'number' &&
+    Number.isFinite(b.remainingFraction)
+  )
+}
+
+function parseGroups(data: unknown): AntigravityQuotaGroup[] {
+  const raw = (data as { groups?: unknown })?.groups
+  if (!Array.isArray(raw)) {
+    return []
+  }
+  return raw.flatMap((entry) => {
+    const group = entry as { displayName?: unknown; buckets?: unknown }
+    if (typeof group.displayName !== 'string' || !Array.isArray(group.buckets)) {
+      return []
+    }
+    return [{ displayName: group.displayName, buckets: group.buckets.filter(isBucket) }]
+  })
+}
+
+async function postInternal(
+  host: string,
+  method: string,
+  accessToken: string,
+  body: unknown,
+  signal: AbortSignal
+): Promise<{ ok: true; data: unknown } | { ok: false; status: number }> {
+  const res = await net.fetch(`https://${host}/v1internal:${method}`, {
+    method: 'POST',
+    headers: {
+      'Content-Type': 'application/json',
+      Authorization: `Bearer ${accessToken}`,
+      'User-Agent': USER_AGENT
+    },
+    body: JSON.stringify(body),
+    signal
+  })
+  if (!res.ok) {
+    await cancelUnreadResponseBody(res)
+    return { ok: false, status: res.status }
+  }
+  return { ok: true, data: (await res.json()) as unknown }
+}
+
+function hostFailureMessage(host: string, error: unknown): string {
+  const reason = error instanceof Error ? error.message : 'request failed'
+  return `Antigravity quota request to ${host} failed: ${reason}`
+}
+
+async function fetchFromHost(
+  host: string,
+  accessToken: string,
+  signal: AbortSignal
+): Promise<AntigravityQuotaFetch> {
+  const loaded = await postInternal(
+    host,
+    'loadCodeAssist',
+    accessToken,
+    { metadata: { ideType: 'ANTIGRAVITY' } },
+    signal
+  )
+  if (!loaded.ok) {
+    return loaded.status === 401
+      ? { status: 'unauthorized' }
+      : { status: 'error', message: `Antigravity project lookup failed (HTTP ${loaded.status})` }
+  }
+  const project = (loaded.data as { cloudaicompanionProject?: unknown }).cloudaicompanionProject
+  if (typeof project !== 'string' || project.length === 0) {
+    return { status: 'not-entitled' }
+  }
+  const quota = await postInternal(
+    host,
+    'retrieveUserQuotaSummary',
+    accessToken,
+    { project },
+    signal
+  )
+  if (!quota.ok) {
+    if (quota.status === 401) {
+      return { status: 'unauthorized' }
+    }
+    return quota.status === 403
+      ? { status: 'not-entitled' }
+      : { status: 'error', message: `Antigravity quota fetch failed (HTTP ${quota.status})` }
+  }
+  return { status: 'ok', groups: parseGroups(quota.data) }
+}
+
+export async function fetchAntigravityQuota(
+  accessToken: string,
+  callerSignal?: AbortSignal
+): Promise<AntigravityQuotaFetch> {
+  let last: AntigravityQuotaFetch = { status: 'error', message: 'No Cloud Code host answered' }
+  for (const host of API_HOSTS) {
+    if (callerSignal?.aborted) {
+      return last
+    }
+    // Why: one deadline shared across hosts lets a slow first host spend the budget the
+    // fallback needs, so the second attempt starts already aborted.
+    const timeout = AbortSignal.timeout(API_TIMEOUT_MS)
+    const signal = callerSignal ? AbortSignal.any([callerSignal, timeout]) : timeout
+    try {
+      last = await fetchFromHost(host, accessToken, signal)
+    } catch (error) {
+      // Why: a rejected request is this host's verdict, not the account's — the next host still gets a turn.
+      last = { status: 'error', message: hostFailureMessage(host, error) }
+    }
+    // Why: an auth or entitlement verdict is the account's, not the host's — retrying decides nothing.
+    if (last.status !== 'error') {
+      return last
+    }
+  }
+  return last
+}