From 9b0e4af548b21c63befa5d40c68673ffad74c8d5 Mon Sep 17 00:00:00 2001 From: abruption Date: Mon, 7 Sep 2026 13:59:28 +0900 Subject: [PATCH 1/3] feat(usage): read Antigravity quota from its own credential, including SSH hosts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Antigravity usage was a relabelled copy of the Gemini Code Assist read, so it showed the wrong pools and went unavailable whenever no Gemini CLI sign-in existed. This adds a real read: the Antigravity CLI's own OAuth credential, the Antigravity `ideType`, and `retrieveUserQuotaSummary`. The credential is looked up on this machine first and then on connected SSH targets, so a dev box where `agy` runs but no Antigravity process is reachable from this machine now reports its quota. The Gemini mirror stays as the fallback whenever no direct read succeeds. Two details the endpoint requires that are easy to miss: - cloudcode-pa selects the product from the User-Agent, not from the request body. Without "antigravity" in it, `loadCodeAssist` still answers 200 but omits `cloudaicompanionProject`, and the quota call then fails 403 — which reads as a bad sign-in. Electron's default Chromium User-Agent resolves to a different product with no quota grant. - the response is grouped, and two groups routinely report identical percentages and reset times, so buckets stay keyed by group rather than being deduplicated. A remote host that drops mid-read is reported as unchecked rather than as a missing sign-in; a target that was never connected is not treated as evidence either way. --- .../antigravity-cloud-code-api.test.ts | 138 ++++++++++++++++ .../rate-limits/antigravity-cloud-code-api.ts | 148 ++++++++++++++++++ .../antigravity-credential-hosts.test.ts | 136 ++++++++++++++++ .../antigravity-credential-hosts.ts | 123 +++++++++++++++ .../antigravity-oauth-credential.test.ts | 67 ++++++++ .../antigravity-oauth-credential.ts | 56 +++++++ .../antigravity-quota-buckets.test.ts | 127 +++++++++++++++ .../rate-limits/antigravity-quota-buckets.ts | 58 +++++++ .../antigravity-quota-fetch.test.ts | 124 +++++++++++++++ .../rate-limits/antigravity-quota-fetch.ts | 105 +++++++++++++ .../rate-limit-service-test-harness.ts | 6 + .../service-account-target-selection.test.ts | 4 + .../service-antigravity-usage.test.ts | 57 ++++++- .../service-inactive-account-previews.test.ts | 4 + .../service-live-claude-usage.test.ts | 4 + .../rate-limits/service-minimax-usage.test.ts | 4 + .../service-refresh-orchestration.test.ts | 4 + .../service-window-activation.test.ts | 4 + .../service/service-full-cycle-application.ts | 11 +- .../service/service-full-cycle-preparation.ts | 85 +++++----- 20 files changed, 1224 insertions(+), 41 deletions(-) create mode 100644 src/main/rate-limits/antigravity-cloud-code-api.test.ts create mode 100644 src/main/rate-limits/antigravity-cloud-code-api.ts create mode 100644 src/main/rate-limits/antigravity-credential-hosts.test.ts create mode 100644 src/main/rate-limits/antigravity-credential-hosts.ts create mode 100644 src/main/rate-limits/antigravity-oauth-credential.test.ts create mode 100644 src/main/rate-limits/antigravity-oauth-credential.ts create mode 100644 src/main/rate-limits/antigravity-quota-buckets.test.ts create mode 100644 src/main/rate-limits/antigravity-quota-buckets.ts create mode 100644 src/main/rate-limits/antigravity-quota-fetch.test.ts create mode 100644 src/main/rate-limits/antigravity-quota-fetch.ts diff --git a/src/main/rate-limits/antigravity-cloud-code-api.test.ts b/src/main/rate-limits/antigravity-cloud-code-api.test.ts new file mode 100644 index 000000000000..6970da20784f --- /dev/null +++ b/src/main/rate-limits/antigravity-cloud-code-api.test.ts @@ -0,0 +1,138 @@ +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 + 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('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) + }) +}) diff --git a/src/main/rate-limits/antigravity-cloud-code-api.ts b/src/main/rate-limits/antigravity-cloud-code-api.ts new file mode 100644 index 000000000000..936263bf5ab7 --- /dev/null +++ b/src/main/rate-limits/antigravity-cloud-code-api.ts @@ -0,0 +1,148 @@ +/** + * 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 } +} + +async function fetchFromHost( + host: string, + accessToken: string, + signal: AbortSignal +): Promise { + 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 { + const timeout = AbortSignal.timeout(API_TIMEOUT_MS) + const signal = callerSignal ? AbortSignal.any([callerSignal, timeout]) : timeout + let last: AntigravityQuotaFetch = { status: 'error', message: 'No Cloud Code host answered' } + for (const host of API_HOSTS) { + last = await fetchFromHost(host, accessToken, signal) + // 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 +} diff --git a/src/main/rate-limits/antigravity-credential-hosts.test.ts b/src/main/rate-limits/antigravity-credential-hosts.test.ts new file mode 100644 index 000000000000..4e9e4d18dfd7 --- /dev/null +++ b/src/main/rate-limits/antigravity-credential-hosts.test.ts @@ -0,0 +1,136 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const { readFileMock, listTargetsMock, resolveRouteMock, resolveHomeMock, unverifiableMock } = + vi.hoisted(() => ({ + readFileMock: vi.fn(), + listTargetsMock: vi.fn(), + resolveRouteMock: vi.fn(), + resolveHomeMock: vi.fn(), + unverifiableMock: vi.fn() + })) + +vi.mock('node:fs/promises', () => ({ readFile: readFileMock })) +vi.mock('node:os', () => ({ homedir: () => '/home/local' })) +vi.mock('../ssh/ssh-target-registry', () => ({ listRegisteredSshTargets: listTargetsMock })) +vi.mock('../providers/execution-host-provider-dispatch', () => ({ + resolveFilesystemRouteForHost: resolveRouteMock +})) +vi.mock('../ipc/repos/remote-home-path', () => ({ resolveRemoteHomePath: resolveHomeMock })) +vi.mock('../ssh/ssh-channel-multiplexer', () => ({ + isSshRequestOutcomeUnverifiable: unverifiableMock +})) + +import { findAntigravityCredential } from './antigravity-credential-hosts' + +const NOW = 1_700_000_000_000 +const CREDENTIAL = (expiry: string) => + JSON.stringify({ token: { access_token: 'access-1', expiry } }) +const FRESH = CREDENTIAL(new Date(NOW + 60_000).toISOString()) +const STALE = CREDENTIAL(new Date(NOW - 60_000).toISOString()) + +function sshTarget(id: string, label: string) { + return { id, label, host: `${id}.example`, port: 22, username: 'ubuntu' } +} + +function reachableRoute(readFile: ReturnType) { + return { kind: 'ssh', hostId: 'ssh:box', connectionId: 'box', provider: { readFile } } +} + +describe('findAntigravityCredential', () => { + beforeEach(() => { + readFileMock.mockReset().mockRejectedValue(new Error('ENOENT')) + listTargetsMock.mockReset().mockReturnValue([]) + resolveRouteMock.mockReset() + resolveHomeMock.mockReset().mockResolvedValue('/home/ubuntu') + unverifiableMock.mockReset().mockReturnValue(false) + }) + + it('reads the Antigravity CLI path on this machine first', async () => { + readFileMock.mockResolvedValue(FRESH) + listTargetsMock.mockReturnValue([sshTarget('box', 'Dev box')]) + const result = await findAntigravityCredential(NOW) + expect(result).toMatchObject({ status: 'ok', found: { hostId: 'local' } }) + expect(readFileMock).toHaveBeenCalledWith( + '/home/local/.gemini/antigravity-cli/antigravity-oauth-token', + 'utf8' + ) + expect(resolveRouteMock).not.toHaveBeenCalled() + }) + + it('falls through to an SSH host when this machine has no sign-in', async () => { + const remoteRead = vi.fn().mockResolvedValue({ content: FRESH, isBinary: false }) + listTargetsMock.mockReturnValue([sshTarget('box', 'Dev box')]) + resolveRouteMock.mockReturnValue(reachableRoute(remoteRead)) + const result = await findAntigravityCredential(NOW) + expect(result).toMatchObject({ + status: 'ok', + found: { hostId: 'ssh:box', hostLabel: 'Dev box' } + }) + expect(remoteRead).toHaveBeenCalledWith( + '/home/ubuntu/.gemini/antigravity-cli/antigravity-oauth-token' + ) + }) + + it('skips targets Orca owns as runtime plumbing', async () => { + listTargetsMock.mockReturnValue([ + { ...sshTarget('box', 'Owned'), owner: { type: 'on-demand-runtime', runtimeId: 'r1' } }, + sshTarget('runtime-ssh-1', 'Nested') + ]) + expect(await findAntigravityCredential(NOW)).toEqual({ status: 'missing' }) + expect(resolveRouteMock).not.toHaveBeenCalled() + }) + + it('reports an expired sign-in against the host that holds it', async () => { + readFileMock.mockResolvedValue(STALE) + expect(await findAntigravityCredential(NOW)).toEqual({ + status: 'expired', + hostLabel: 'this machine' + }) + }) + + it('prefers a usable remote sign-in over an expired local one', async () => { + readFileMock.mockResolvedValue(STALE) + listTargetsMock.mockReturnValue([sshTarget('box', 'Dev box')]) + resolveRouteMock.mockReturnValue( + reachableRoute(vi.fn().mockResolvedValue({ content: FRESH, isBinary: false })) + ) + expect(await findAntigravityCredential(NOW)).toMatchObject({ + status: 'ok', + found: { hostId: 'ssh:box' } + }) + }) + + it('does not raise an error for a target the user has not connected', async () => { + listTargetsMock.mockReturnValue([sshTarget('box', 'Dev box')]) + resolveRouteMock.mockReturnValue({ + kind: 'ssh', + hostId: 'ssh:box', + connectionId: 'box', + provider: null + }) + expect(await findAntigravityCredential(NOW)).toEqual({ status: 'missing' }) + }) + + it('keeps looking when resolving a remote home throws', async () => { + listTargetsMock.mockReturnValue([sshTarget('box', 'Dev box')]) + resolveRouteMock.mockReturnValue(reachableRoute(vi.fn())) + resolveHomeMock.mockRejectedValue(new Error('relay gone')) + expect(await findAntigravityCredential(NOW)).toEqual({ status: 'missing' }) + }) + + it('treats a lost link as unverifiable but a real read failure as absent', async () => { + listTargetsMock.mockReturnValue([sshTarget('box', 'Dev box')]) + resolveRouteMock.mockReturnValue(reachableRoute(vi.fn().mockRejectedValue(new Error('boom')))) + + unverifiableMock.mockReturnValue(true) + expect(await findAntigravityCredential(NOW)).toMatchObject({ status: 'unverifiable' }) + + unverifiableMock.mockReturnValue(false) + expect(await findAntigravityCredential(NOW)).toEqual({ status: 'missing' }) + }) + + it('reports missing when a host holds an unparseable blob', async () => { + readFileMock.mockResolvedValue('not json') + expect(await findAntigravityCredential(NOW)).toEqual({ status: 'missing' }) + }) +}) diff --git a/src/main/rate-limits/antigravity-credential-hosts.ts b/src/main/rate-limits/antigravity-credential-hosts.ts new file mode 100644 index 000000000000..ac8e16051f1c --- /dev/null +++ b/src/main/rate-limits/antigravity-credential-hosts.ts @@ -0,0 +1,123 @@ +/** + * Locates the Antigravity CLI credential on the execution hosts Orca knows about. + * + * `agy` writes this file itself; reading it is what lets Orca report Antigravity quota + * for a host with no `agy` process reachable from this machine — an SSH dev box, most + * commonly. Only the one well-known `agy` path is ever read. + */ + +import { readFile } from 'node:fs/promises' +import { homedir } from 'node:os' +import { join, posix } from 'node:path' +import { + isRuntimeOwnedSshTargetId, + LOCAL_EXECUTION_HOST_ID, + toSshExecutionHostId, + type ExecutionHostId +} from '../../shared/execution-host' +import { resolveRemoteHomePath } from '../ipc/repos/remote-home-path' +import { resolveFilesystemRouteForHost } from '../providers/execution-host-provider-dispatch' +import { isSshRequestOutcomeUnverifiable } from '../ssh/ssh-channel-multiplexer' +import { listRegisteredSshTargets } from '../ssh/ssh-target-registry' +import { + isAntigravityCredentialExpired, + parseAntigravityOAuthCredential, + type AntigravityOAuthCredential +} from './antigravity-oauth-credential' + +const CREDENTIAL_PATH_SEGMENTS = ['.gemini', 'antigravity-cli', 'antigravity-oauth-token'] as const + +export type AntigravityHostCredential = { + credential: AntigravityOAuthCredential + hostId: ExecutionHostId + hostLabel: string +} + +export type AntigravityCredentialLookup = + | { status: 'ok'; found: AntigravityHostCredential } + | { status: 'missing' } + | { status: 'expired'; hostLabel: string } + /** Contact with a remote host was lost; absence was never established. */ + | { status: 'unverifiable'; hostLabel: string; message: string } + +type HostReadOutcome = + | { kind: 'ok'; credential: AntigravityOAuthCredential } + | { kind: 'absent' } + | { kind: 'unverifiable'; message: string } + +async function readLocalCredential(): Promise { + try { + const raw = await readFile(join(homedir(), ...CREDENTIAL_PATH_SEGMENTS), 'utf8') + const credential = parseAntigravityOAuthCredential(raw) + return credential ? { kind: 'ok', credential } : { kind: 'absent' } + } catch { + return { kind: 'absent' } + } +} + +async function readSshCredential(targetId: string): Promise { + const route = resolveFilesystemRouteForHost(toSshExecutionHostId(targetId)) + // Why: a target the user has not connected was never contacted, so it is not evidence + // either way — turning every idle target into an error would drown the real answer. + if (route.kind !== 'ssh' || !route.provider) { + return { kind: 'absent' } + } + const provider = route.provider + try { + const home = await resolveRemoteHomePath(route.connectionId, '~') + const result = await provider.readFile(posix.join(home, ...CREDENTIAL_PATH_SEGMENTS)) + const credential = parseAntigravityOAuthCredential(result.content) + return credential ? { kind: 'ok', credential } : { kind: 'absent' } + } catch (error) { + // Why: a relayed error loses errno, so "no sign-in here" and "link dropped" are + // only separable by the transport verdict (docs/reference/ssh-execution-boundary.md). + return isSshRequestOutcomeUnverifiable(error) + ? { kind: 'unverifiable', message: 'Remote host did not answer.' } + : { kind: 'absent' } + } +} + +type HostCandidate = { + hostId: ExecutionHostId + label: string + read: () => Promise +} + +function candidateHosts(): HostCandidate[] { + const remote = listRegisteredSshTargets() + .filter((target) => !target.owner && !isRuntimeOwnedSshTargetId(target.id)) + .map((target) => ({ + hostId: toSshExecutionHostId(target.id), + label: target.label || target.host, + read: () => readSshCredential(target.id) + })) + return [ + { hostId: LOCAL_EXECUTION_HOST_ID, label: 'this machine', read: readLocalCredential }, + ...remote + ] +} + +/** + * First host holding a usable credential wins, local before remote. An expired or + * unverifiable host is remembered only as the answer to report when no host succeeds. + */ +export async function findAntigravityCredential( + now: number = Date.now() +): Promise { + let fallback: AntigravityCredentialLookup | null = null + for (const host of candidateHosts()) { + const outcome = await host.read() + if (outcome.kind === 'ok' && !isAntigravityCredentialExpired(outcome.credential, now)) { + return { + status: 'ok', + found: { credential: outcome.credential, hostId: host.hostId, hostLabel: host.label } + } + } + if (outcome.kind === 'ok') { + fallback ??= { status: 'expired', hostLabel: host.label } + } else if (outcome.kind === 'unverifiable') { + fallback ??= { status: 'unverifiable', hostLabel: host.label, message: outcome.message } + } + } + return fallback ?? { status: 'missing' } +} diff --git a/src/main/rate-limits/antigravity-oauth-credential.test.ts b/src/main/rate-limits/antigravity-oauth-credential.test.ts new file mode 100644 index 000000000000..973909ad1b9c --- /dev/null +++ b/src/main/rate-limits/antigravity-oauth-credential.test.ts @@ -0,0 +1,67 @@ +import { describe, expect, it } from 'vitest' +import { + isAntigravityCredentialExpired, + parseAntigravityOAuthCredential +} from './antigravity-oauth-credential' + +const NESTED = { + auth_method: 'consumer', + token: { + access_token: 'access-1', + token_type: 'Bearer', + refresh_token: 'refresh-1', + expiry: '2026-09-07T12:44:03.30722+09:00' + } +} + +describe('parseAntigravityOAuthCredential', () => { + it('reads the nested token the Antigravity CLI writes', () => { + expect(parseAntigravityOAuthCredential(JSON.stringify(NESTED))).toEqual({ + accessToken: 'access-1', + expiresAt: new Date('2026-09-07T12:44:03.30722+09:00').getTime() + }) + }) + + it('reads the same payload behind the go-keyring base64 prefix', () => { + const encoded = `go-keyring-base64:${Buffer.from(JSON.stringify(NESTED)).toString('base64')}` + expect(parseAntigravityOAuthCredential(encoded)?.accessToken).toBe('access-1') + }) + + it('reads a flat payload that has no token wrapper', () => { + const flat = JSON.stringify({ access_token: 'access-2', expiry: '2026-01-01T00:00:00Z' }) + expect(parseAntigravityOAuthCredential(flat)?.accessToken).toBe('access-2') + }) + + it('keeps a usable token whose expiry is unparseable', () => { + const odd = JSON.stringify({ token: { access_token: 'access-3', expiry: 'not-a-date' } }) + expect(parseAntigravityOAuthCredential(odd)).toEqual({ + accessToken: 'access-3', + expiresAt: null + }) + }) + + it('rejects blobs with no access token', () => { + expect(parseAntigravityOAuthCredential('not json')).toBeNull() + expect(parseAntigravityOAuthCredential('{}')).toBeNull() + expect( + parseAntigravityOAuthCredential(JSON.stringify({ token: { access_token: '' } })) + ).toBeNull() + }) +}) + +describe('isAntigravityCredentialExpired', () => { + const at = (expiresAt: number | null) => ({ accessToken: 'a', expiresAt }) + + it('treats an elapsed expiry as expired', () => { + expect(isAntigravityCredentialExpired(at(1_000), 1_001)).toBe(true) + expect(isAntigravityCredentialExpired(at(1_000), 1_000)).toBe(true) + }) + + it('treats a future expiry as usable', () => { + expect(isAntigravityCredentialExpired(at(2_000), 1_000)).toBe(false) + }) + + it('lets the API decide when the expiry is unknown', () => { + expect(isAntigravityCredentialExpired(at(null), 1_000)).toBe(false) + }) +}) diff --git a/src/main/rate-limits/antigravity-oauth-credential.ts b/src/main/rate-limits/antigravity-oauth-credential.ts new file mode 100644 index 000000000000..5564f2d0301c --- /dev/null +++ b/src/main/rate-limits/antigravity-oauth-credential.ts @@ -0,0 +1,56 @@ +/** + * Parses the OAuth credential the Antigravity CLI (`agy`) persists for itself. + * + * Two storage encodings exist for the same payload: plain JSON, and the + * zalando/go-keyring convention of `go-keyring-base64:` followed by base64(JSON) + * that `agy` uses when it writes through an OS keyring. + */ + +const GO_KEYRING_BASE64_PREFIX = 'go-keyring-base64:' + +export type AntigravityOAuthCredential = { + accessToken: string + /** Unix ms; NaN-safe. `null` when the stored value has no parseable expiry. */ + expiresAt: number | null +} + +function decodeStoredBlob(raw: string): unknown { + const trimmed = raw.trim() + const payload = trimmed.startsWith(GO_KEYRING_BASE64_PREFIX) + ? Buffer.from(trimmed.slice(GO_KEYRING_BASE64_PREFIX.length), 'base64').toString('utf8') + : trimmed + try { + return JSON.parse(payload) as unknown + } catch { + return null + } +} + +export function parseAntigravityOAuthCredential(raw: string): AntigravityOAuthCredential | null { + const parsed = decodeStoredBlob(raw) + if (!parsed || typeof parsed !== 'object') { + return null + } + // Why: `agy` nests the OAuth grant under `token`, but older writes stored it flat. + const nested = (parsed as { token?: unknown }).token + const token = (nested && typeof nested === 'object' ? nested : parsed) as { + access_token?: unknown + expiry?: unknown + } + if (typeof token.access_token !== 'string' || token.access_token.length === 0) { + return null + } + const expiryMs = typeof token.expiry === 'string' ? new Date(token.expiry).getTime() : Number.NaN + return { + accessToken: token.access_token, + expiresAt: Number.isNaN(expiryMs) ? null : expiryMs + } +} + +/** An unknown expiry is treated as usable; the API call is the authority on a dead token. */ +export function isAntigravityCredentialExpired( + credential: AntigravityOAuthCredential, + now: number = Date.now() +): boolean { + return credential.expiresAt !== null && credential.expiresAt <= now +} diff --git a/src/main/rate-limits/antigravity-quota-buckets.test.ts b/src/main/rate-limits/antigravity-quota-buckets.test.ts new file mode 100644 index 000000000000..d6ec9cafe58c --- /dev/null +++ b/src/main/rate-limits/antigravity-quota-buckets.test.ts @@ -0,0 +1,127 @@ +import { describe, expect, it } from 'vitest' +import type { AntigravityQuotaGroup } from './antigravity-cloud-code-api' +import { + buildAntigravityBuckets, + mostConstrainedWindow, + toRateLimitBuckets +} from './antigravity-quota-buckets' + +const WEEKLY_RESET = '2026-09-11T05:00:13Z' +const FIVE_HOUR_RESET = '2026-09-07T08:00:13Z' + +function group(displayName: string, remaining: { weekly: number; fiveHour: number }) { + return { + displayName, + buckets: [ + { + displayName: 'Weekly Limit Remaining', + window: 'weekly', + resetTime: WEEKLY_RESET, + remainingFraction: remaining.weekly + }, + { + displayName: 'Five Hour Limit Remaining', + window: '5h', + resetTime: FIVE_HOUR_RESET, + remainingFraction: remaining.fiveHour + } + ] + } satisfies AntigravityQuotaGroup +} + +describe('buildAntigravityBuckets', () => { + it('maps the reported window rather than assuming one duration', () => { + const [weekly, fiveHour] = buildAntigravityBuckets([ + group('Gemini Models', { weekly: 0.79675514, fiveHour: 0.94602 }) + ]) + expect(weekly).toMatchObject({ + name: 'Gemini Models (Weekly)', + windowMinutes: 10080, + usedPercent: 20, + resetsAt: new Date(WEEKLY_RESET).getTime() + }) + expect(fiveHour).toMatchObject({ + name: 'Gemini Models (5h)', + windowMinutes: 300, + usedPercent: 5 + }) + }) + + it('keeps groups apart when their numbers are identical', () => { + const buckets = buildAntigravityBuckets([ + group('Gemini Models', { weekly: 1, fiveHour: 1 }), + group('Claude and GPT models', { weekly: 1, fiveHour: 1 }) + ]) + expect(buckets.map((b) => b.name)).toEqual([ + 'Gemini Models (Weekly)', + 'Gemini Models (5h)', + 'Claude and GPT models (Weekly)', + 'Claude and GPT models (5h)' + ]) + }) + + it('drops windows it cannot express instead of guessing a duration', () => { + const buckets = buildAntigravityBuckets([ + { + displayName: 'Gemini Models', + buckets: [ + { + displayName: 'Monthly', + window: 'monthly', + resetTime: WEEKLY_RESET, + remainingFraction: 0.5 + } + ] + } + ]) + expect(buckets).toEqual([]) + }) + + it('clamps a reset time it cannot parse to null', () => { + const [bucket] = buildAntigravityBuckets([ + { + displayName: 'Gemini Models', + buckets: [ + { + displayName: 'Weekly Limit Remaining', + window: 'weekly', + resetTime: 'nope', + remainingFraction: 0.5 + } + ] + } + ]) + expect(bucket?.resetsAt).toBeNull() + }) +}) + +describe('mostConstrainedWindow', () => { + const buckets = buildAntigravityBuckets([ + group('Gemini Models', { weekly: 0.2, fiveHour: 0.9 }), + group('Claude and GPT models', { weekly: 0.95, fiveHour: 0.1 }) + ]) + + it('reports the fullest group for each window', () => { + expect(mostConstrainedWindow(buckets, 'weekly')).toMatchObject({ + usedPercent: 80, + windowMinutes: 10080 + }) + expect(mostConstrainedWindow(buckets, '5h')).toMatchObject({ + usedPercent: 90, + windowMinutes: 300 + }) + }) + + it('returns null when the window is absent', () => { + expect(mostConstrainedWindow(buckets, 'monthly')).toBeNull() + }) +}) + +describe('toRateLimitBuckets', () => { + it('drops the internal window discriminator', () => { + const [bucket] = toRateLimitBuckets( + buildAntigravityBuckets([group('Gemini Models', { weekly: 1, fiveHour: 1 })]) + ) + expect(bucket && 'window' in bucket).toBe(false) + }) +}) diff --git a/src/main/rate-limits/antigravity-quota-buckets.ts b/src/main/rate-limits/antigravity-quota-buckets.ts new file mode 100644 index 000000000000..98037f15975c --- /dev/null +++ b/src/main/rate-limits/antigravity-quota-buckets.ts @@ -0,0 +1,58 @@ +import type { RateLimitBucket, RateLimitWindow } from '../../shared/rate-limit-types' +import type { AntigravityQuotaGroup } from './antigravity-cloud-code-api' + +const WINDOW_MINUTES: Record = { '5h': 300, weekly: 10080 } +const WINDOW_SUFFIX: Record = { '5h': '5h', weekly: 'Weekly' } + +export const ANTIGRAVITY_SESSION_WINDOW = '5h' +export const ANTIGRAVITY_WEEKLY_WINDOW = 'weekly' + +export type AntigravityBucket = RateLimitBucket & { window: string } + +function toUsedPercent(remainingFraction: number): number { + return Math.min(100, Math.max(0, Math.round((1 - remainingFraction) * 100))) +} + +/** + * Buckets stay keyed by group because two groups routinely report identical + * percentages and reset times, and collapsing them would hide a whole model family. + */ +export function buildAntigravityBuckets(groups: AntigravityQuotaGroup[]): AntigravityBucket[] { + return groups.flatMap((group) => + group.buckets.flatMap((bucket) => { + const windowMinutes = WINDOW_MINUTES[bucket.window] + if (windowMinutes === undefined) { + return [] + } + const resetsAt = new Date(bucket.resetTime).getTime() + return [ + { + name: `${group.displayName} (${WINDOW_SUFFIX[bucket.window]})`, + usedPercent: toUsedPercent(bucket.remainingFraction), + windowMinutes, + resetsAt: Number.isNaN(resetsAt) ? null : resetsAt, + resetDescription: null, + window: bucket.window + } + ] + }) + ) +} + +/** The window a user is actually constrained by is the fullest one in that period. */ +export function mostConstrainedWindow( + buckets: AntigravityBucket[], + window: string +): RateLimitWindow | null { + const inWindow = buckets.filter((bucket) => bucket.window === window) + if (inWindow.length === 0) { + return null + } + const worst = inWindow.reduce((a, b) => (b.usedPercent > a.usedPercent ? b : a)) + const { name: _name, window: _window, ...rest } = worst + return rest +} + +export function toRateLimitBuckets(buckets: AntigravityBucket[]): RateLimitBucket[] { + return buckets.map(({ window: _window, ...rest }) => rest) +} diff --git a/src/main/rate-limits/antigravity-quota-fetch.test.ts b/src/main/rate-limits/antigravity-quota-fetch.test.ts new file mode 100644 index 000000000000..378187bccf56 --- /dev/null +++ b/src/main/rate-limits/antigravity-quota-fetch.test.ts @@ -0,0 +1,124 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const { findCredentialMock, fetchQuotaMock } = vi.hoisted(() => ({ + findCredentialMock: vi.fn(), + fetchQuotaMock: vi.fn() +})) + +vi.mock('./antigravity-credential-hosts', () => ({ + findAntigravityCredential: findCredentialMock +})) +vi.mock('./antigravity-cloud-code-api', () => ({ fetchAntigravityQuota: fetchQuotaMock })) + +import { fetchAntigravityRateLimits } from './antigravity-quota-fetch' + +const FOUND = { + status: 'ok', + found: { + credential: { accessToken: 'access-1', expiresAt: null }, + hostId: 'ssh:box', + hostLabel: 'Dev box' + } +} + +const GROUPS = [ + { + displayName: 'Gemini Models', + buckets: [ + { + displayName: 'Weekly Limit Remaining', + window: 'weekly', + resetTime: '2026-09-11T05:00:13Z', + remainingFraction: 0.8 + }, + { + displayName: 'Five Hour Limit Remaining', + window: '5h', + resetTime: '2026-09-07T08:00:13Z', + remainingFraction: 0.5 + } + ] + } +] + +describe('fetchAntigravityRateLimits', () => { + beforeEach(() => { + findCredentialMock.mockReset().mockResolvedValue(FOUND) + fetchQuotaMock.mockReset().mockResolvedValue({ status: 'ok', groups: GROUPS }) + }) + + it('publishes both windows and names the host the sign-in came from', async () => { + const result = await fetchAntigravityRateLimits() + expect(result).toMatchObject({ + provider: 'antigravity', + status: 'ok', + error: null, + session: { usedPercent: 50, windowMinutes: 300 }, + weekly: { usedPercent: 20, windowMinutes: 10080 }, + usageMetadata: { source: 'oauth', credentialSource: 'Dev box' } + }) + expect(result.buckets).toHaveLength(2) + }) + + it('uses the stored access token as-is', async () => { + await fetchAntigravityRateLimits() + expect(fetchQuotaMock).toHaveBeenCalledWith('access-1', undefined) + }) + + it('asks the user to sign in when no host holds a credential', async () => { + findCredentialMock.mockResolvedValue({ status: 'missing' }) + const result = await fetchAntigravityRateLimits() + expect(result).toMatchObject({ + status: 'unavailable', + usageMetadata: { failureKind: 'missing-credentials' } + }) + expect(fetchQuotaMock).not.toHaveBeenCalled() + }) + + it('names the host whose sign-in expired', async () => { + findCredentialMock.mockResolvedValue({ status: 'expired', hostLabel: 'Dev box' }) + const result = await fetchAntigravityRateLimits() + expect(result.status).toBe('unavailable') + expect(result.error).toContain('Dev box') + expect(result.usageMetadata?.failureKind).toBe('stale-token') + }) + + it('does not report a lost remote host as a missing sign-in', async () => { + findCredentialMock.mockResolvedValue({ + status: 'unverifiable', + hostLabel: 'Dev box', + message: 'Remote host did not answer.' + }) + const result = await fetchAntigravityRateLimits() + expect(result.status).toBe('error') + expect(result.error).not.toMatch(/sign in to/i) + expect(result.usageMetadata?.failureKind).toBe('network') + }) + + it('separates a rejected token from an account without a grant', async () => { + fetchQuotaMock.mockResolvedValue({ status: 'unauthorized' }) + expect((await fetchAntigravityRateLimits()).usageMetadata?.failureKind).toBe('stale-token') + + fetchQuotaMock.mockResolvedValue({ status: 'not-entitled' }) + const notEntitled = await fetchAntigravityRateLimits() + expect(notEntitled.status).toBe('unavailable') + expect(notEntitled.usageMetadata?.failureKind).toBe('usage-unavailable') + }) + + it('surfaces a transport error without claiming quota is empty', async () => { + fetchQuotaMock.mockResolvedValue({ status: 'error', message: 'Antigravity quota fetch failed' }) + const result = await fetchAntigravityRateLimits() + expect(result).toMatchObject({ status: 'error', usageMetadata: { failureKind: 'server' } }) + }) + + it('reports a response with no usable window as a parse failure', async () => { + fetchQuotaMock.mockResolvedValue({ status: 'ok', groups: [] }) + expect((await fetchAntigravityRateLimits()).usageMetadata?.failureKind).toBe('parse') + }) + + it('forwards the caller abort signal', async () => { + const controller = new AbortController() + await fetchAntigravityRateLimits(controller.signal) + expect(fetchQuotaMock).toHaveBeenCalledWith('access-1', controller.signal) + }) +}) diff --git a/src/main/rate-limits/antigravity-quota-fetch.ts b/src/main/rate-limits/antigravity-quota-fetch.ts new file mode 100644 index 000000000000..ff571efd32b8 --- /dev/null +++ b/src/main/rate-limits/antigravity-quota-fetch.ts @@ -0,0 +1,105 @@ +import type { + ProviderRateLimits, + UsageRateLimitFailureKind, + UsageRateLimitMetadata +} from '../../shared/rate-limit-types' +import { fetchAntigravityQuota } from './antigravity-cloud-code-api' +import { findAntigravityCredential } from './antigravity-credential-hosts' +import { + ANTIGRAVITY_SESSION_WINDOW, + ANTIGRAVITY_WEEKLY_WINDOW, + buildAntigravityBuckets, + mostConstrainedWindow, + toRateLimitBuckets +} from './antigravity-quota-buckets' + +function unusable( + status: 'error' | 'unavailable', + error: string, + failureKind: UsageRateLimitFailureKind, + metadata: UsageRateLimitMetadata = {} +): ProviderRateLimits { + return { + provider: 'antigravity', + session: null, + weekly: null, + updatedAt: Date.now(), + error, + status, + usageMetadata: { source: 'oauth', failureKind, ...metadata } + } +} + +/** + * Reads Antigravity's own quota from the Antigravity credential, on whichever + * execution host holds one. No `agy` process has to be running or reachable. + */ +export async function fetchAntigravityRateLimits( + signal?: AbortSignal +): Promise { + const lookup = await findAntigravityCredential() + if (lookup.status === 'missing') { + return unusable( + 'unavailable', + 'Antigravity usage is not available. Sign in to the Antigravity CLI on this machine or on a connected SSH host.', + 'missing-credentials' + ) + } + if (lookup.status === 'expired') { + return unusable( + 'unavailable', + `Antigravity sign-in on ${lookup.hostLabel} has expired. Run the Antigravity CLI there to refresh it.`, + 'stale-token', + { credentialSource: lookup.hostLabel } + ) + } + if (lookup.status === 'unverifiable') { + // Why: contact was lost before absence could be established, so this must not be + // reported as "no sign-in" (docs/reference/ssh-execution-boundary.md). + return unusable( + 'error', + `Antigravity usage could not be checked on ${lookup.hostLabel}. ${lookup.message}`, + 'network', + { credentialSource: lookup.hostLabel } + ) + } + + const { credential, hostLabel } = lookup.found + const quota = await fetchAntigravityQuota(credential.accessToken, signal) + if (quota.status === 'unauthorized') { + return unusable( + 'error', + `Antigravity rejected the sign-in stored on ${hostLabel}.`, + 'stale-token', + { credentialSource: hostLabel } + ) + } + if (quota.status === 'not-entitled') { + return unusable( + 'unavailable', + 'This Google account has no Antigravity quota grant.', + 'usage-unavailable', + { credentialSource: hostLabel } + ) + } + if (quota.status === 'error') { + return unusable('error', quota.message, 'server', { credentialSource: hostLabel }) + } + + const buckets = buildAntigravityBuckets(quota.groups) + if (buckets.length === 0) { + return unusable('error', 'Antigravity returned no recognisable quota windows.', 'parse', { + credentialSource: hostLabel + }) + } + return { + provider: 'antigravity', + session: mostConstrainedWindow(buckets, ANTIGRAVITY_SESSION_WINDOW), + weekly: mostConstrainedWindow(buckets, ANTIGRAVITY_WEEKLY_WINDOW), + buckets: toRateLimitBuckets(buckets), + updatedAt: Date.now(), + error: null, + status: 'ok', + usageMetadata: { source: 'oauth', credentialSource: hostLabel } + } +} diff --git a/src/main/rate-limits/rate-limit-service-test-harness.ts b/src/main/rate-limits/rate-limit-service-test-harness.ts index 00d3db160029..50cf075ce9f5 100644 --- a/src/main/rate-limits/rate-limit-service-test-harness.ts +++ b/src/main/rate-limits/rate-limit-service-test-harness.ts @@ -2,6 +2,7 @@ import { EventEmitter } from 'node:events' import { vi, type Mock } from 'vitest' import type { ProviderRateLimits } from '../../shared/rate-limit-types' import type { RateLimitService } from './service' +import { fetchAntigravityRateLimits } from './antigravity-quota-fetch' import { fetchCodexRateLimits } from './codex-fetcher' import { fetchGeminiRateLimits } from './gemini-usage-fetcher' import { fetchKimiRateLimits } from './kimi-fetcher' @@ -89,6 +90,10 @@ export function mockFreshBackgroundProviderFetches(): void { vi.mocked(fetchKimiRateLimits).mockImplementation(async () => okProvider('kimi', 0)) vi.mocked(fetchMiniMaxRateLimits).mockImplementation(async () => okProvider('minimax', 0)) vi.mocked(fetchGrokRateLimits).mockImplementation(async () => unavailableProvider('grok')) + // Why: the mirror is the documented fallback, so an un-stubbed Antigravity read must leave it in charge. + vi.mocked(fetchAntigravityRateLimits).mockImplementation(async () => + unavailableProvider('antigravity') + ) } /** Shared `beforeEach` body: healthy stubs for every provider the service polls. */ @@ -106,6 +111,7 @@ export function resetRateLimitProviderMocks(): void { error: null, status: 'unavailable' }) + vi.mocked(fetchAntigravityRateLimits).mockResolvedValue(unavailableProvider('antigravity')) vi.mocked(hasMiniMaxSessionCookie).mockReturnValue(false) vi.mocked(readGrokAuthSession).mockReturnValue({ status: 'missing' }) } diff --git a/src/main/rate-limits/service-account-target-selection.test.ts b/src/main/rate-limits/service-account-target-selection.test.ts index 86008a4bd846..2a100e670acc 100644 --- a/src/main/rate-limits/service-account-target-selection.test.ts +++ b/src/main/rate-limits/service-account-target-selection.test.ts @@ -20,6 +20,10 @@ vi.mock('./codex-fetcher', () => ({ fetchCodexRateLimits: vi.fn() })) +vi.mock('./antigravity-quota-fetch', () => ({ + fetchAntigravityRateLimits: vi.fn() +})) + vi.mock('./gemini-usage-fetcher', () => ({ fetchGeminiRateLimits: vi.fn() })) diff --git a/src/main/rate-limits/service-antigravity-usage.test.ts b/src/main/rate-limits/service-antigravity-usage.test.ts index b617d9099443..e5ae2b2df694 100644 --- a/src/main/rate-limits/service-antigravity-usage.test.ts +++ b/src/main/rate-limits/service-antigravity-usage.test.ts @@ -2,11 +2,13 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' import { RateLimitService } from './service' import { fetchClaudeRateLimits } from './claude-fetcher' import { fetchCodexRateLimits } from './codex-fetcher' +import { fetchAntigravityRateLimits } from './antigravity-quota-fetch' import { fetchGeminiRateLimits } from './gemini-usage-fetcher' import { errorProvider, okProvider, - resetRateLimitProviderMocks + resetRateLimitProviderMocks, + unavailableProvider } from './rate-limit-service-test-harness' vi.mock('./claude-fetcher', () => ({ @@ -19,6 +21,10 @@ vi.mock('./codex-fetcher', () => ({ fetchCodexRateLimits: vi.fn() })) +vi.mock('./antigravity-quota-fetch', () => ({ + fetchAntigravityRateLimits: vi.fn() +})) + vi.mock('./gemini-usage-fetcher', () => ({ fetchGeminiRateLimits: vi.fn() })) @@ -97,4 +103,53 @@ describe('RateLimitService Antigravity usage', () => { expect(service.getState().antigravity?.status).toBe('unavailable') expect(service.getState().antigravity?.session).toBeNull() }) + + it("publishes Antigravity's own quota instead of the Gemini mirror when a read succeeds", async () => { + vi.mocked(fetchGeminiRateLimits).mockResolvedValue(okProvider('gemini', 42, Date.now())) + vi.mocked(fetchAntigravityRateLimits).mockResolvedValue({ + ...okProvider('antigravity', 90, Date.now()), + buckets: [ + { + name: 'Claude and GPT models (5h)', + usedPercent: 90, + windowMinutes: 300, + resetsAt: null, + resetDescription: null + } + ] + }) + const service = new RateLimitService() + + await service.refresh() + + const state = service.getState() + expect(state.antigravity?.status).toBe('ok') + expect(state.antigravity?.session?.usedPercent).toBe(90) + expect(state.antigravity?.buckets?.[0]?.name).toBe('Claude and GPT models (5h)') + // Why: the mirror must not overwrite a real read, and Gemini keeps its own numbers. + expect(state.gemini?.session?.usedPercent).toBe(42) + }) + + it('falls back to the mirror when no host holds an Antigravity sign-in', async () => { + vi.mocked(fetchGeminiRateLimits).mockResolvedValue(okProvider('gemini', 42, Date.now())) + vi.mocked(fetchAntigravityRateLimits).mockResolvedValue(unavailableProvider('antigravity')) + const service = new RateLimitService() + + await service.refresh() + + expect(service.getState().antigravity?.session?.usedPercent).toBe(42) + }) + + it('does not let a failed Antigravity read blank out a usable mirror', async () => { + vi.mocked(fetchGeminiRateLimits).mockResolvedValue(okProvider('gemini', 42, Date.now())) + vi.mocked(fetchAntigravityRateLimits).mockResolvedValue( + errorProvider('antigravity', 'Remote host did not answer.') + ) + const service = new RateLimitService() + + await service.refresh() + + expect(service.getState().antigravity?.status).toBe('ok') + expect(service.getState().antigravity?.session?.usedPercent).toBe(42) + }) }) diff --git a/src/main/rate-limits/service-inactive-account-previews.test.ts b/src/main/rate-limits/service-inactive-account-previews.test.ts index e734b52172a4..2c898d96efca 100644 --- a/src/main/rate-limits/service-inactive-account-previews.test.ts +++ b/src/main/rate-limits/service-inactive-account-previews.test.ts @@ -27,6 +27,10 @@ vi.mock('./codex-fetcher', () => ({ fetchCodexRateLimits: vi.fn() })) +vi.mock('./antigravity-quota-fetch', () => ({ + fetchAntigravityRateLimits: vi.fn() +})) + vi.mock('./gemini-usage-fetcher', () => ({ fetchGeminiRateLimits: vi.fn() })) diff --git a/src/main/rate-limits/service-live-claude-usage.test.ts b/src/main/rate-limits/service-live-claude-usage.test.ts index 59c1300532da..978267e47e7f 100644 --- a/src/main/rate-limits/service-live-claude-usage.test.ts +++ b/src/main/rate-limits/service-live-claude-usage.test.ts @@ -24,6 +24,10 @@ vi.mock('./codex-fetcher', () => ({ fetchCodexRateLimits: vi.fn() })) +vi.mock('./antigravity-quota-fetch', () => ({ + fetchAntigravityRateLimits: vi.fn() +})) + vi.mock('./gemini-usage-fetcher', () => ({ fetchGeminiRateLimits: vi.fn() })) diff --git a/src/main/rate-limits/service-minimax-usage.test.ts b/src/main/rate-limits/service-minimax-usage.test.ts index 7c3db21c63e7..aaf20a72411c 100644 --- a/src/main/rate-limits/service-minimax-usage.test.ts +++ b/src/main/rate-limits/service-minimax-usage.test.ts @@ -21,6 +21,10 @@ vi.mock('./codex-fetcher', () => ({ fetchCodexRateLimits: vi.fn() })) +vi.mock('./antigravity-quota-fetch', () => ({ + fetchAntigravityRateLimits: vi.fn() +})) + vi.mock('./gemini-usage-fetcher', () => ({ fetchGeminiRateLimits: vi.fn() })) diff --git a/src/main/rate-limits/service-refresh-orchestration.test.ts b/src/main/rate-limits/service-refresh-orchestration.test.ts index 7ac7f22f1649..dfd40d538d7a 100644 --- a/src/main/rate-limits/service-refresh-orchestration.test.ts +++ b/src/main/rate-limits/service-refresh-orchestration.test.ts @@ -28,6 +28,10 @@ vi.mock('./codex-fetcher', () => ({ fetchCodexRateLimits: vi.fn() })) +vi.mock('./antigravity-quota-fetch', () => ({ + fetchAntigravityRateLimits: vi.fn() +})) + vi.mock('./gemini-usage-fetcher', () => ({ fetchGeminiRateLimits: vi.fn() })) diff --git a/src/main/rate-limits/service-window-activation.test.ts b/src/main/rate-limits/service-window-activation.test.ts index a4a1dc14f767..839cfa495bf0 100644 --- a/src/main/rate-limits/service-window-activation.test.ts +++ b/src/main/rate-limits/service-window-activation.test.ts @@ -29,6 +29,10 @@ vi.mock('./codex-fetcher', () => ({ fetchCodexRateLimits: vi.fn() })) +vi.mock('./antigravity-quota-fetch', () => ({ + fetchAntigravityRateLimits: vi.fn() +})) + vi.mock('./gemini-usage-fetcher', () => ({ fetchGeminiRateLimits: vi.fn() })) diff --git a/src/main/rate-limits/service/service-full-cycle-application.ts b/src/main/rate-limits/service/service-full-cycle-application.ts index 7104face38fa..9b4bc5e8c643 100644 --- a/src/main/rate-limits/service/service-full-cycle-application.ts +++ b/src/main/rate-limits/service/service-full-cycle-application.ts @@ -32,7 +32,8 @@ export abstract class RateLimitServiceFullCycleApplication extends RateLimitServ geminiResult, opencodeGoResult, kimiResult, - miniMaxResult + miniMaxResult, + antigravityResult ], grokResultPromise } = prepared @@ -79,8 +80,12 @@ export abstract class RateLimitServiceFullCycleApplication extends RateLimitServ status: 'error' } satisfies ProviderRateLimits) - // Why: Antigravity can only borrow a *successful* Gemini read; a Gemini failure is not an Antigravity failure. - const antigravity = deriveAntigravityRateLimits(gemini) + // Why: the mirror stays as the fallback for accounts whose quota only exists as the shared + // Google Code Assist pool; a direct Antigravity read supersedes it whenever one succeeds. + const antigravity = + antigravityResult.status === 'fulfilled' && antigravityResult.value.status === 'ok' + ? antigravityResult.value + : deriveAntigravityRateLimits(gemini) const opencodeGo = opencodeGoResult.status === 'fulfilled' diff --git a/src/main/rate-limits/service/service-full-cycle-preparation.ts b/src/main/rate-limits/service/service-full-cycle-preparation.ts index c5bf533bc865..63012feb97a0 100644 --- a/src/main/rate-limits/service/service-full-cycle-preparation.ts +++ b/src/main/rate-limits/service/service-full-cycle-preparation.ts @@ -1,3 +1,4 @@ +import { fetchAntigravityRateLimits } from '../antigravity-quota-fetch' import { fetchClaudeRateLimits } from '../claude-fetcher' import { fetchCodexRateLimits } from '../codex-fetcher' import { fetchGeminiRateLimits } from '../gemini-usage-fetcher' @@ -36,6 +37,7 @@ export type FetchAllCyclePrepared = { PromiseSettledResult, PromiseSettledResult, PromiseSettledResult, + PromiseSettledResult, PromiseSettledResult ] grokResultPromise: Promise< @@ -138,42 +140,50 @@ export abstract class RateLimitServiceFullCyclePreparation extends RateLimitServ const claudeFetchGated = !options?.force && this.shouldSkipAutomatedClaudeFetch(previousState.claude) - const [claudeResult, codexResult, geminiResult, opencodeGoResult, kimiResult, miniMaxResult] = - await Promise.allSettled([ - claudeFetchGated - ? Promise.resolve(previousState.claude as ProviderRateLimits) - : fetchClaudeRateLimits({ - authPreparation: claudeAuthPreparation, - allowPtyFallback: this.shouldAllowClaudePtyFallback(claudeAuthPreparation), - allowUsagePanelSupplement: this.shouldAllowClaudeUsagePanelSupplement(), - networkProxySettings: this.networkProxySettingsResolver?.(), - signal - }), - codexFetchGated - ? Promise.resolve(previousState.codex as ProviderRateLimits) - : (missingWslCodexHome ?? - fetchCodexRateLimits({ - codexHomePath, - allowPtyFallback: this.shouldAllowCodexPtyFallback(), - signal - })), - fetchGeminiRateLimits(geminiCliOAuthEnabled), - fetchOpenCodeGoRateLimits( - cookie, - workspaceIdOverride || undefined, - this.networkProxySettingsResolver?.() - ), - this.fetchKimiWithResolvedHome(), - miniMaxConfigResult.error - ? Promise.resolve(this.getMiniMaxCredentialError(miniMaxConfigResult.error)) - : fetchMiniMaxRateLimits({ - cookie: miniMaxCookie, - groupId: miniMaxGroupId, - models: miniMaxModels, - endpointMode: miniMaxEndpoint, - apiKey: miniMaxApiKey - }) - ]) + const [ + claudeResult, + codexResult, + geminiResult, + opencodeGoResult, + kimiResult, + miniMaxResult, + antigravityResult + ] = await Promise.allSettled([ + claudeFetchGated + ? Promise.resolve(previousState.claude as ProviderRateLimits) + : fetchClaudeRateLimits({ + authPreparation: claudeAuthPreparation, + allowPtyFallback: this.shouldAllowClaudePtyFallback(claudeAuthPreparation), + allowUsagePanelSupplement: this.shouldAllowClaudeUsagePanelSupplement(), + networkProxySettings: this.networkProxySettingsResolver?.(), + signal + }), + codexFetchGated + ? Promise.resolve(previousState.codex as ProviderRateLimits) + : (missingWslCodexHome ?? + fetchCodexRateLimits({ + codexHomePath, + allowPtyFallback: this.shouldAllowCodexPtyFallback(), + signal + })), + fetchGeminiRateLimits(geminiCliOAuthEnabled), + fetchOpenCodeGoRateLimits( + cookie, + workspaceIdOverride || undefined, + this.networkProxySettingsResolver?.() + ), + this.fetchKimiWithResolvedHome(), + miniMaxConfigResult.error + ? Promise.resolve(this.getMiniMaxCredentialError(miniMaxConfigResult.error)) + : fetchMiniMaxRateLimits({ + cookie: miniMaxCookie, + groupId: miniMaxGroupId, + models: miniMaxModels, + endpointMode: miniMaxEndpoint, + apiKey: miniMaxApiKey + }), + fetchAntigravityRateLimits(signal) + ]) if (signal.aborted) { return null @@ -200,7 +210,8 @@ export abstract class RateLimitServiceFullCyclePreparation extends RateLimitServ geminiResult, opencodeGoResult, kimiResult, - miniMaxResult + miniMaxResult, + antigravityResult ], grokResultPromise } From 8e1ac1632aedb37ec3781efd414d9fbd10f72c6e Mon Sep 17 00:00:00 2001 From: abruption Date: Mon, 7 Sep 2026 14:39:45 +0900 Subject: [PATCH 2/3] fix(usage): give each Cloud Code host its own attempt, and stop hiding Antigravity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-ups. Host fallback never worked for the failure it existed for. `fetchFromHost` did not catch a rejected request, so a network-level failure on the daily host left `fetchAntigravityQuota` before the stable host was tried, and both hosts shared one `AbortSignal.timeout`, so a slow first host handed the fallback a signal that had already fired. Each host now gets its own deadline and its own catch, and a caller abort stops the loop instead of burning a second attempt. The direct read's verdicts were unreachable. Every non-`ok` result was replaced by the Gemini mirror, so "the sign-in on has expired" and "the remote host did not answer" could never be shown — including the `unverifiable` distinction this path exists to preserve. Those two verdicts now survive the fallback; no sign-in anywhere, no Antigravity grant, and transient endpoint failures still degrade to the mirror, which is the only case where it has something to say. The status bar hid Antigravity in exactly the case this targets. The slot was gated on local PATH detection, so a client with no `agy` of its own never rendered the bar even when the remote read succeeded, and the slot additionally required the Gemini CLI opt-in because the snapshot used to mirror that fetch. A snapshot carrying quota now counts as availability on its own, and the Gemini opt-in is no longer a precondition. --- .../antigravity-cloud-code-api.test.ts | 40 ++++++++++ .../rate-limits/antigravity-cloud-code-api.ts | 21 +++++- .../antigravity-usage-mirror.test.ts | 73 ++++++++++++++++++- .../rate-limits/antigravity-usage-mirror.ts | 34 +++++++-- .../service/service-full-cycle-application.ts | 13 ++-- .../status-bar/StatusBarVisibilityMenu.tsx | 3 +- .../status-bar-agent-gating.test.ts | 28 ++++++- .../status-bar/status-bar-agent-gating.ts | 11 +++ .../status-bar-provider-visibility.test.ts | 35 ++++----- .../status-bar-provider-visibility.ts | 15 ++-- .../status-bar/use-status-bar-controller.ts | 16 ++-- 11 files changed, 234 insertions(+), 55 deletions(-) diff --git a/src/main/rate-limits/antigravity-cloud-code-api.test.ts b/src/main/rate-limits/antigravity-cloud-code-api.test.ts index 6970da20784f..fd15dc4f5d99 100644 --- a/src/main/rate-limits/antigravity-cloud-code-api.test.ts +++ b/src/main/rate-limits/antigravity-cloud-code-api.test.ts @@ -130,6 +130,46 @@ describe('fetchAntigravityQuota', () => { 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) diff --git a/src/main/rate-limits/antigravity-cloud-code-api.ts b/src/main/rate-limits/antigravity-cloud-code-api.ts index 936263bf5ab7..1a3b759414a6 100644 --- a/src/main/rate-limits/antigravity-cloud-code-api.ts +++ b/src/main/rate-limits/antigravity-cloud-code-api.ts @@ -91,6 +91,11 @@ async function postInternal( 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, @@ -134,11 +139,21 @@ export async function fetchAntigravityQuota( accessToken: string, callerSignal?: AbortSignal ): Promise { - const timeout = AbortSignal.timeout(API_TIMEOUT_MS) - const signal = callerSignal ? AbortSignal.any([callerSignal, timeout]) : timeout let last: AntigravityQuotaFetch = { status: 'error', message: 'No Cloud Code host answered' } for (const host of API_HOSTS) { - last = await fetchFromHost(host, accessToken, signal) + 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 diff --git a/src/main/rate-limits/antigravity-usage-mirror.test.ts b/src/main/rate-limits/antigravity-usage-mirror.test.ts index 84c92e1ad22f..c39b6017646e 100644 --- a/src/main/rate-limits/antigravity-usage-mirror.test.ts +++ b/src/main/rate-limits/antigravity-usage-mirror.test.ts @@ -1,6 +1,13 @@ import { describe, expect, it } from 'vitest' -import type { ProviderRateLimits, ProviderRateLimitStatus } from '../../shared/rate-limit-types' -import { deriveAntigravityRateLimits } from './antigravity-usage-mirror' +import type { + ProviderRateLimits, + ProviderRateLimitStatus, + UsageRateLimitFailureKind +} from '../../shared/rate-limit-types' +import { + deriveAntigravityRateLimits, + resolveAntigravityRateLimits +} from './antigravity-usage-mirror' function geminiSnapshot( status: ProviderRateLimitStatus, @@ -68,3 +75,65 @@ describe('deriveAntigravityRateLimits', () => { expect(antigravity.error).toContain('Gemini CLI sign-in is connected') }) }) + +describe('resolveAntigravityRateLimits', () => { + const gemini: ProviderRateLimits = { + provider: 'gemini', + session: { usedPercent: 42, windowMinutes: 300, resetsAt: null, resetDescription: null }, + weekly: null, + updatedAt: 1, + error: null, + status: 'ok' + } + const direct = ( + status: ProviderRateLimits['status'], + failureKind?: UsageRateLimitFailureKind + ): ProviderRateLimits => ({ + provider: 'antigravity', + session: null, + weekly: null, + updatedAt: 2, + error: status === 'ok' ? null : 'nope', + status, + usageMetadata: failureKind ? { source: 'oauth', failureKind } : { source: 'oauth' } + }) + + it('publishes a successful direct read over the mirror', () => { + const ok = { ...direct('ok'), session: gemini.session } + expect(resolveAntigravityRateLimits(ok, gemini)).toBe(ok) + }) + + it('keeps an expired sign-in visible instead of showing Gemini numbers', () => { + const expired = direct('unavailable', 'stale-token') + expect(resolveAntigravityRateLimits(expired, gemini)).toBe(expired) + }) + + it('keeps an unreachable host visible so absence is never implied', () => { + const unreachable = direct('error', 'network') + expect(resolveAntigravityRateLimits(unreachable, gemini)).toBe(unreachable) + }) + + it('falls back to the mirror when no Antigravity sign-in exists anywhere', () => { + const result = resolveAntigravityRateLimits( + direct('unavailable', 'missing-credentials'), + gemini + ) + expect(result.provider).toBe('antigravity') + expect(result.session?.usedPercent).toBe(42) + }) + + it('falls back to the mirror for an account with no Antigravity grant', () => { + const result = resolveAntigravityRateLimits(direct('unavailable', 'usage-unavailable'), gemini) + expect(result.session?.usedPercent).toBe(42) + }) + + it('falls back to the mirror for a transient endpoint failure', () => { + expect( + resolveAntigravityRateLimits(direct('error', 'server'), gemini).session?.usedPercent + ).toBe(42) + }) + + it('falls back to the mirror when the direct read threw', () => { + expect(resolveAntigravityRateLimits(null, gemini).session?.usedPercent).toBe(42) + }) +}) diff --git a/src/main/rate-limits/antigravity-usage-mirror.ts b/src/main/rate-limits/antigravity-usage-mirror.ts index 3cea6decc111..8b6e5ba8635b 100644 --- a/src/main/rate-limits/antigravity-usage-mirror.ts +++ b/src/main/rate-limits/antigravity-usage-mirror.ts @@ -1,15 +1,39 @@ -import type { ProviderRateLimits } from '../../shared/rate-limit-types' +import type { ProviderRateLimits, UsageRateLimitFailureKind } from '../../shared/rate-limit-types' -// Why: the Antigravity CLI keeps its token in the OS keyring, not in the files the Gemini -// fetcher reads, so Orca never actually queries Antigravity. Only a *successful* Gemini read -// describes shared Google Code Assist quota; republishing a Gemini failure under the -// Antigravity provider id surfaced "Refresh failed" for a request that was never attempted. +// Why: this is the fallback for when the direct Antigravity read has no answer of its own — no +// sign-in on any host, or an account whose quota exists only as the shared Google Code Assist +// pool. Only a *successful* Gemini read describes that pool; republishing a Gemini failure under +// the Antigravity provider id surfaced "Refresh failed" for a request that was never attempted. const ANTIGRAVITY_NO_SIGN_IN_REASON = 'Antigravity usage is not available. Orca can only show shared Google Code Assist quota while a Gemini CLI sign-in is connected.' // Why: a Gemini `error` means the sign-in exists and the quota read failed, so blaming a missing sign-in would misdirect the user. const ANTIGRAVITY_QUOTA_UNREADABLE_REASON = 'Antigravity usage is not available. Orca reads it from the shared Google Code Assist quota, which could not be read right now.' +// Why: an expired sign-in and a host we could not reach are real answers about Antigravity, and +// each carries its own remedy; replacing them with Gemini's numbers hides the failure and the fix. +// Every other verdict — no sign-in anywhere, no Antigravity grant, a transient endpoint failure — +// leaves the shared-pool mirror as the best answer still available. +const DIRECT_VERDICTS_WORTH_SURFACING: ReadonlySet = new Set([ + 'stale-token', + 'network' +]) + +/** Picks between a direct Antigravity read and the shared Code Assist mirror. */ +export function resolveAntigravityRateLimits( + direct: ProviderRateLimits | null, + gemini: ProviderRateLimits +): ProviderRateLimits { + if (direct?.status === 'ok') { + return direct + } + const failureKind = direct?.usageMetadata?.failureKind + if (failureKind && DIRECT_VERDICTS_WORTH_SURFACING.has(failureKind)) { + return direct as ProviderRateLimits + } + return deriveAntigravityRateLimits(gemini) +} + export function deriveAntigravityRateLimits(gemini: ProviderRateLimits): ProviderRateLimits { if (gemini.status === 'ok') { return { ...gemini, provider: 'antigravity' } diff --git a/src/main/rate-limits/service/service-full-cycle-application.ts b/src/main/rate-limits/service/service-full-cycle-application.ts index 9b4bc5e8c643..095a784576cc 100644 --- a/src/main/rate-limits/service/service-full-cycle-application.ts +++ b/src/main/rate-limits/service/service-full-cycle-application.ts @@ -1,5 +1,5 @@ import { RateLimitServiceFullCyclePreparation } from './service-full-cycle-preparation' -import { deriveAntigravityRateLimits } from '../antigravity-usage-mirror' +import { resolveAntigravityRateLimits } from '../antigravity-usage-mirror' import type { ProviderRateLimits } from './service-types' export abstract class RateLimitServiceFullCycleApplication extends RateLimitServiceFullCyclePreparation { @@ -81,11 +81,12 @@ export abstract class RateLimitServiceFullCycleApplication extends RateLimitServ } satisfies ProviderRateLimits) // Why: the mirror stays as the fallback for accounts whose quota only exists as the shared - // Google Code Assist pool; a direct Antigravity read supersedes it whenever one succeeds. - const antigravity = - antigravityResult.status === 'fulfilled' && antigravityResult.value.status === 'ok' - ? antigravityResult.value - : deriveAntigravityRateLimits(gemini) + // Google Code Assist pool, but a direct read that reached a verdict about the user's + // Antigravity sign-in has to survive it. + const antigravity = resolveAntigravityRateLimits( + antigravityResult.status === 'fulfilled' ? antigravityResult.value : null, + gemini + ) const opencodeGo = opencodeGoResult.status === 'fulfilled' diff --git a/src/renderer/src/components/status-bar/StatusBarVisibilityMenu.tsx b/src/renderer/src/components/status-bar/StatusBarVisibilityMenu.tsx index 7b47b6bb6d29..ac2892811184 100644 --- a/src/renderer/src/components/status-bar/StatusBarVisibilityMenu.tsx +++ b/src/renderer/src/components/status-bar/StatusBarVisibilityMenu.tsx @@ -18,6 +18,7 @@ export function StatusBarVisibilityMenu({ controller: StatusBarController }): React.JSX.Element { const { + antigravitySlotAvailable, detectedAgentIds, menuOpen, menuPoint, @@ -74,7 +75,7 @@ export function StatusBarVisibilityMenu({ {translate('auto.components.status.bar.StatusBar.c1df0d67ec', 'Gemini Usage')} )} - {isStatusBarItemAvailable('antigravity', detectedAgentIds) && ( + {antigravitySlotAvailable && ( { diff --git a/src/renderer/src/components/status-bar/status-bar-agent-gating.test.ts b/src/renderer/src/components/status-bar/status-bar-agent-gating.test.ts index cd98315df505..6c5f714a5b3a 100644 --- a/src/renderer/src/components/status-bar/status-bar-agent-gating.test.ts +++ b/src/renderer/src/components/status-bar/status-bar-agent-gating.test.ts @@ -1,5 +1,8 @@ import { describe, expect, it } from 'vitest' -import { isStatusBarItemAvailable } from './status-bar-agent-gating' +import { + isAntigravityStatusBarAvailable, + isStatusBarItemAvailable +} from './status-bar-agent-gating' describe('isStatusBarItemAvailable', () => { it('shows non-CLI items regardless of detection', () => { @@ -38,3 +41,26 @@ describe('isStatusBarItemAvailable', () => { expect(isStatusBarItemAvailable('grok', ['grok'])).toBe(true) }) }) + +describe('isAntigravityStatusBarAvailable', () => { + const okSnapshot = { status: 'ok' } as const + + it('shows the slot for a sign-in that lives only on a remote host', () => { + // Why: `agy` may exist only on the SSH host, so PATH detection alone would hide a bar + // that a successful read has already proved useful. + expect(isAntigravityStatusBarAvailable([], okSnapshot)).toBe(true) + }) + + it('shows the slot when the CLI is on PATH even with no snapshot yet', () => { + expect(isAntigravityStatusBarAvailable(['antigravity'], null)).toBe(true) + }) + + it('hides the slot when neither the CLI nor a reading is present', () => { + expect(isAntigravityStatusBarAvailable([], null)).toBe(false) + expect(isAntigravityStatusBarAvailable([], { status: 'unavailable' })).toBe(false) + }) + + it('keeps the pre-detection default so the bar does not flicker on cold start', () => { + expect(isAntigravityStatusBarAvailable(null, null)).toBe(true) + }) +}) diff --git a/src/renderer/src/components/status-bar/status-bar-agent-gating.ts b/src/renderer/src/components/status-bar/status-bar-agent-gating.ts index 145abef3a900..bba36ee8dc54 100644 --- a/src/renderer/src/components/status-bar/status-bar-agent-gating.ts +++ b/src/renderer/src/components/status-bar/status-bar-agent-gating.ts @@ -1,3 +1,4 @@ +import type { ProviderRateLimits } from '../../../../shared/rate-limit-types' import type { TuiAgent } from '../../../../shared/tui-agent' import type { StatusBarItem } from '../../../../shared/ui-chrome-types' @@ -28,3 +29,13 @@ export function isStatusBarItemAvailable( } return detectedAgentIds.includes(id as TuiAgent) } + +// Why: Antigravity is read from a credential, so it can be signed in only on a remote execution +// host with no `agy` on this machine's PATH. A snapshot that actually carries quota is its own +// proof the slot is useful; PATH detection alone would hide the bar in exactly that case. +export function isAntigravityStatusBarAvailable( + detectedAgentIds: TuiAgent[] | null, + antigravity: Pick | null | undefined +): boolean { + return isStatusBarItemAvailable('antigravity', detectedAgentIds) || antigravity?.status === 'ok' +} diff --git a/src/renderer/src/components/status-bar/status-bar-provider-visibility.test.ts b/src/renderer/src/components/status-bar/status-bar-provider-visibility.test.ts index 41a933a850bc..faaf96771a1d 100644 --- a/src/renderer/src/components/status-bar/status-bar-provider-visibility.test.ts +++ b/src/renderer/src/components/status-bar/status-bar-provider-visibility.test.ts @@ -162,21 +162,20 @@ describe('hasUsageProviderSettingsForProvider', () => { expect(hasUsageProviderSettingsForProvider('grok', usageSettings())).toBe(false) }) - it('requires both a checked Antigravity item and Gemini OAuth as the durable Antigravity signal', () => { + it('takes a checked Antigravity item as the durable Antigravity signal on its own', () => { expect( hasUsageProviderSettingsForProvider( 'antigravity', usageSettings({ antigravityUsageConfigured: true, geminiCliOAuthEnabled: true }) ) ).toBe(true) - // Why: the snapshot mirrors the Gemini fetch — without the OAuth opt-in it - // is permanently unavailable, so the checked item alone is not durable. + // Why: Antigravity is read from its own credential, so the Gemini opt-in is not a precondition. expect( hasUsageProviderSettingsForProvider( 'antigravity', usageSettings({ antigravityUsageConfigured: true }) ) - ).toBe(false) + ).toBe(true) expect( hasUsageProviderSettingsForProvider( 'antigravity', @@ -358,26 +357,18 @@ describe('getVisibleUsageProvider', () => { }) }) - it('hides Antigravity while Gemini OAuth is off even when its status item is checked', () => { - // Why: without the OAuth opt-in the mirrored snapshot is permanently - // 'unavailable'; the default-on item must not pin a dead bar. + it('shows Antigravity with the Gemini opt-in off once its status item is checked', () => { + // Why: a direct Antigravity read does not go through the Gemini fetch, so gating the slot on + // that opt-in would hide a bar that can fill from the host holding the sign-in. expect( getVisibleUsageProvider( 'antigravity', null, usageSettings({ antigravityUsageConfigured: true }) ) - ).toBe(null) - expect( - getVisibleUsageProvider( - 'antigravity', - provider('unavailable', { - provider: 'antigravity', - error: 'Gemini CLI OAuth is disabled in settings' - }), - usageSettings({ antigravityUsageConfigured: true }) - ) - ).toBe(null) + ).not.toBe(null) + // Why: the slot itself is still earned by the checked item, not by the default alone. + expect(getVisibleUsageProvider('antigravity', null, usageSettings({}))).toBe(null) }) }) @@ -519,9 +510,9 @@ describe('isUsageEmptyState', () => { ).toBe(false) }) - it('still shows the setup CTA when Antigravity is checked but Gemini OAuth is off', () => { - // Why: the default-on Antigravity item is not configured usage on its own; - // it must not hide the teaching CTA from users who set nothing up. + it('drops the setup CTA once Antigravity is checked, with or without the Gemini opt-in', () => { + // Why: the checked item is only set after the slot is available — `agy` on PATH, or a + // snapshot already carrying quota from a remote host — so it is real setup, not a default. expect( isUsageEmptyState( { @@ -536,6 +527,6 @@ describe('isUsageEmptyState', () => { }, usageSettings({ antigravityUsageConfigured: true }) ) - ).toBe(true) + ).toBe(false) }) }) diff --git a/src/renderer/src/components/status-bar/status-bar-provider-visibility.ts b/src/renderer/src/components/status-bar/status-bar-provider-visibility.ts index 19258592f1fd..9f0bd6dd0954 100644 --- a/src/renderer/src/components/status-bar/status-bar-provider-visibility.ts +++ b/src/renderer/src/components/status-bar/status-bar-provider-visibility.ts @@ -8,11 +8,9 @@ export type UsageProviderSettings = Pick< | 'opencodeSessionCookie' | 'geminiCliOAuthEnabled' > & { - // Why: Antigravity has no separate persisted usage credential in Orca. The - // checked status-bar item is the durable user signal; StatusBar only sets - // this after PATH detection says the agent is available. Durability further - // requires geminiCliOAuthEnabled — the snapshot mirrors the Gemini fetch, - // which never yields data while that opt-in is off. + // Why: Antigravity has no separate persisted usage credential in Orca. The checked status-bar + // item is the durable user signal; StatusBar only sets this once the slot is available — the + // agent is on PATH, or a snapshot already carries quota from a remote host. antigravityUsageConfigured: boolean // Why: MiniMax/Grok sign-in live on disk, not in settings; main sets these each poll. minimaxCookieConfigured: boolean @@ -103,10 +101,9 @@ export function hasUsageProviderSettingsForProvider( return Boolean(settings.opencodeSessionCookie?.trim()) } if (providerId === 'antigravity') { - // Why: the Antigravity snapshot mirrors the Gemini fetch, which stays - // 'unavailable' until the user opts into Gemini CLI OAuth. Without that - // gate the default-on checked item would pin a permanently dead bar. - return settings.antigravityUsageConfigured === true && settings.geminiCliOAuthEnabled === true + // Why: Antigravity is read from its own credential now, so it no longer depends on the Gemini + // CLI opt-in. The slot gate upstream already keeps the default-on item from pinning a dead bar. + return settings.antigravityUsageConfigured === true } if (providerId === 'minimax') { return settings.minimaxCookieConfigured === true || settings.minimaxApiKeyConfigured === true diff --git a/src/renderer/src/components/status-bar/use-status-bar-controller.ts b/src/renderer/src/components/status-bar/use-status-bar-controller.ts index 33db1976b0da..69280bafec19 100644 --- a/src/renderer/src/components/status-bar/use-status-bar-controller.ts +++ b/src/renderer/src/components/status-bar/use-status-bar-controller.ts @@ -5,7 +5,10 @@ import { selectFloatingWorkspaceHasUnread } from '../../store/selectors' import type { ProviderRateLimits } from '../../../../shared/rate-limit-types' import { normalizeUsagePercentageDisplay } from '../../../../shared/usage-percentage-display' import { normalizeStatusBarUsageMode } from '../../../../shared/status-bar-usage-mode' -import { isStatusBarItemAvailable } from './status-bar-agent-gating' +import { + isAntigravityStatusBarAvailable, + isStatusBarItemAvailable +} from './status-bar-agent-gating' import { getVisibleUsageProvider, isUsageEmptyState } from './status-bar-provider-visibility' import { getUsageProviderAccountsSectionId } from './usage-provider-settings-target' import { CLOSE_ALL_CONTEXT_MENUS_EVENT, useStatusBarMenuFocusHandoff } from './ProviderDetailsMenu' @@ -102,11 +105,11 @@ export function useStatusBarController(floatingTerminalOpen: boolean) { const { claude, codex, gemini, opencodeGo, kimi, antigravity, minimax, grok } = rateLimits // Why: a bar is earned by a live snapshot or durable Settings setup; detection-gating hides per-CLI bars when the agent isn't on PATH. - // Why: Antigravity has no persisted credential, so a checked status item + detected CLI is the durable "show its slot" signal. - // Why: Antigravity visibility also requires geminiCliOAuthEnabled because its usage snapshot mirrors the Gemini fetch. + // Why: Antigravity has no persisted credential, so a checked status item plus either a detected + // CLI or a snapshot that carries quota is the durable "show its slot" signal. + const antigravitySlotAvailable = isAntigravityStatusBarAvailable(detectedAgentIds, antigravity) const antigravityUsageConfigured = - statusBarItems.includes('antigravity') && - isStatusBarItemAvailable('antigravity', detectedAgentIds) + statusBarItems.includes('antigravity') && antigravitySlotAvailable // Why: thread non-GlobalSettings durability flags so bars stay visible across reloads and snapshot refreshes. const usageSettings = { ...settings, @@ -141,7 +144,7 @@ export function useStatusBarController(floatingTerminalOpen: boolean) { const showAntigravity = visibleAntigravity !== null && statusBarItems.includes('antigravity') && - isStatusBarItemAvailable('antigravity', detectedAgentIds) + antigravitySlotAvailable // Why: MiniMax is cookie-auth, not a CLI on PATH, so detection-gating doesn't apply. const showMiniMax = visibleMiniMax !== null && statusBarItems.includes('minimax') const showGrok = @@ -236,6 +239,7 @@ export function useStatusBarController(floatingTerminalOpen: boolean) { anyVisible, compact, containerRefCallback, + antigravitySlotAvailable, detectedAgentIds, floatingTerminalActionLabel, floatingTerminalShortcut, From 904e51c6dfa31c0ad6cbc4453a804bb540b2929a Mon Sep 17 00:00:00 2001 From: abruption Date: Mon, 7 Sep 2026 14:53:14 +0900 Subject: [PATCH 3/3] fix(usage): keep the Antigravity bar once a remote sign-in expires or its host drops MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two fixes from the last round cancelled each other out. The direct read's `stale-token` and `network` verdicts are surfaced on purpose, but the status-bar slot was gated on `status === 'ok'`, so on a client with no local `agy` the bar disappeared the moment either verdict landed — burying the message that explains what to do about it, in exactly the case this targets. Availability is now a named credential source rather than a healthy status: it marks that a sign-in was found on some host, which is what makes the slot worth showing whatever state that sign-in is in. No sign-in anywhere still hides the bar, and the Gemini mirror names no source, so mirror-only users keep the local-CLI gate they had before. Also aligns `hasUsageProviderSettings` with the per-provider rule it had drifted from: a checked Antigravity item no longer needs the Gemini opt-in to count as configured usage, so the setup CTA does not contradict a live Antigravity bar. --- .../status-bar-agent-gating.test.ts | 28 ++++++++++++++++--- .../status-bar/status-bar-agent-gating.ts | 12 +++++--- .../status-bar-provider-visibility.test.ts | 8 ++---- .../status-bar-provider-visibility.ts | 3 +- 4 files changed, 36 insertions(+), 15 deletions(-) diff --git a/src/renderer/src/components/status-bar/status-bar-agent-gating.test.ts b/src/renderer/src/components/status-bar/status-bar-agent-gating.test.ts index 6c5f714a5b3a..f9e2f8e58727 100644 --- a/src/renderer/src/components/status-bar/status-bar-agent-gating.test.ts +++ b/src/renderer/src/components/status-bar/status-bar-agent-gating.test.ts @@ -43,21 +43,41 @@ describe('isStatusBarItemAvailable', () => { }) describe('isAntigravityStatusBarAvailable', () => { - const okSnapshot = { status: 'ok' } as const + const fromHost = { usageMetadata: { source: 'oauth', credentialSource: 'Dev box' } } as const it('shows the slot for a sign-in that lives only on a remote host', () => { // Why: `agy` may exist only on the SSH host, so PATH detection alone would hide a bar // that a successful read has already proved useful. - expect(isAntigravityStatusBarAvailable([], okSnapshot)).toBe(true) + expect(isAntigravityStatusBarAvailable([], fromHost)).toBe(true) + }) + + it('keeps the slot once that remote sign-in expires or its host drops', () => { + // Why: those verdicts are surfaced deliberately; hiding the bar would bury the message. + expect( + isAntigravityStatusBarAvailable([], { + usageMetadata: { source: 'oauth', failureKind: 'stale-token', credentialSource: 'Dev box' } + }) + ).toBe(true) + expect( + isAntigravityStatusBarAvailable([], { + usageMetadata: { source: 'oauth', failureKind: 'network', credentialSource: 'Dev box' } + }) + ).toBe(true) }) it('shows the slot when the CLI is on PATH even with no snapshot yet', () => { expect(isAntigravityStatusBarAvailable(['antigravity'], null)).toBe(true) }) - it('hides the slot when neither the CLI nor a reading is present', () => { + it('hides the slot when no host holds a sign-in and the CLI is absent', () => { expect(isAntigravityStatusBarAvailable([], null)).toBe(false) - expect(isAntigravityStatusBarAvailable([], { status: 'unavailable' })).toBe(false) + expect( + isAntigravityStatusBarAvailable([], { + usageMetadata: { source: 'oauth', failureKind: 'missing-credentials' } + }) + ).toBe(false) + // Why: the Gemini mirror names no credential source, so it must not widen the slot. + expect(isAntigravityStatusBarAvailable([], {})).toBe(false) }) it('keeps the pre-detection default so the bar does not flicker on cold start', () => { diff --git a/src/renderer/src/components/status-bar/status-bar-agent-gating.ts b/src/renderer/src/components/status-bar/status-bar-agent-gating.ts index bba36ee8dc54..41d84aad0ba7 100644 --- a/src/renderer/src/components/status-bar/status-bar-agent-gating.ts +++ b/src/renderer/src/components/status-bar/status-bar-agent-gating.ts @@ -31,11 +31,15 @@ export function isStatusBarItemAvailable( } // Why: Antigravity is read from a credential, so it can be signed in only on a remote execution -// host with no `agy` on this machine's PATH. A snapshot that actually carries quota is its own -// proof the slot is useful; PATH detection alone would hide the bar in exactly that case. +// host with no `agy` on this machine's PATH. A named credential source is proof we found a +// sign-in somewhere, which is what makes the slot useful — keying on `ok` instead would hide the +// bar the moment that sign-in expires or its host drops, burying the very message that says so. export function isAntigravityStatusBarAvailable( detectedAgentIds: TuiAgent[] | null, - antigravity: Pick | null | undefined + antigravity: Pick | null | undefined ): boolean { - return isStatusBarItemAvailable('antigravity', detectedAgentIds) || antigravity?.status === 'ok' + return ( + isStatusBarItemAvailable('antigravity', detectedAgentIds) || + Boolean(antigravity?.usageMetadata?.credentialSource) + ) } diff --git a/src/renderer/src/components/status-bar/status-bar-provider-visibility.test.ts b/src/renderer/src/components/status-bar/status-bar-provider-visibility.test.ts index faaf96771a1d..b0801ec5ae06 100644 --- a/src/renderer/src/components/status-bar/status-bar-provider-visibility.test.ts +++ b/src/renderer/src/components/status-bar/status-bar-provider-visibility.test.ts @@ -122,11 +122,9 @@ describe('hasUsageProviderSettings', () => { expect( hasUsageProviderSettings(usageSettings({ opencodeSessionCookie: ' session=abc ' })) ).toBe(true) - // Why: antigravity durability requires the Gemini OAuth opt-in; the - // checked item alone must not suppress the usage setup CTA. - expect(hasUsageProviderSettings(usageSettings({ antigravityUsageConfigured: true }))).toBe( - false - ) + // Why: the checked item is durable on its own now that the read no longer borrows the + // Gemini fetch — this must agree with hasUsageProviderSettingsForProvider('antigravity'). + expect(hasUsageProviderSettings(usageSettings({ antigravityUsageConfigured: true }))).toBe(true) expect(hasUsageProviderSettings(usageSettings({ minimaxCookieConfigured: true }))).toBe(true) expect(hasUsageProviderSettings(usageSettings({ minimaxApiKeyConfigured: true }))).toBe(true) expect(hasUsageProviderSettings(usageSettings({ grokAuthConfigured: true }))).toBe(true) diff --git a/src/renderer/src/components/status-bar/status-bar-provider-visibility.ts b/src/renderer/src/components/status-bar/status-bar-provider-visibility.ts index 9f0bd6dd0954..bfc8942bf0d0 100644 --- a/src/renderer/src/components/status-bar/status-bar-provider-visibility.ts +++ b/src/renderer/src/components/status-bar/status-bar-provider-visibility.ts @@ -73,8 +73,7 @@ export function hasUsageProviderSettings( (settings?.claudeManagedAccounts?.length ?? 0) > 0 || settings?.geminiCliOAuthEnabled === true || Boolean(settings?.opencodeSessionCookie?.trim()) || - // Antigravity's durable signal requires geminiCliOAuthEnabled, so it is - // already covered by the gemini term above. + settings?.antigravityUsageConfigured === true || settings?.minimaxCookieConfigured === true || settings?.minimaxApiKeyConfigured === true || settings?.grokAuthConfigured === true