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..fd15dc4f5d99 --- /dev/null +++ b/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 + 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) + }) +}) 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..1a3b759414a6 --- /dev/null +++ b/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 { + 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 { + 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 +} 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/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/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..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 { @@ -32,7 +32,8 @@ export abstract class RateLimitServiceFullCycleApplication extends RateLimitServ geminiResult, opencodeGoResult, kimiResult, - miniMaxResult + miniMaxResult, + antigravityResult ], grokResultPromise } = prepared @@ -79,8 +80,13 @@ 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, 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/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 } 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..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 @@ -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,46 @@ describe('isStatusBarItemAvailable', () => { expect(isStatusBarItemAvailable('grok', ['grok'])).toBe(true) }) }) + +describe('isAntigravityStatusBarAvailable', () => { + 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([], 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 no host holds a sign-in and the CLI is absent', () => { + expect(isAntigravityStatusBarAvailable([], null)).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', () => { + 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..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 @@ -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,17 @@ 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 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 +): boolean { + 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 41a933a850bc..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) @@ -162,21 +160,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 +355,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 +508,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 +525,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..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 @@ -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 @@ -75,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 @@ -103,10 +100,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,