diff --git a/src/main/rate-limits/antigravity-log-discovery.ts b/src/main/rate-limits/antigravity-log-discovery.ts new file mode 100644 index 000000000000..69639106e34f --- /dev/null +++ b/src/main/rate-limits/antigravity-log-discovery.ts @@ -0,0 +1,157 @@ +import type { Dir } from 'node:fs' +import { open, opendir, type FileHandle } from 'node:fs/promises' + +const CLI_LOG_LIMIT = 12 +const LOG_SECTION_LIMIT_BYTES = 128 * 1024 +const LOG_READ_LIMIT_BYTES = LOG_SECTION_LIMIT_BYTES * 2 + +/** Races filesystem work because Node's promise APIs do not accept AbortSignal. */ +function awaitWithAbort(operation: Promise, signal: AbortSignal): Promise { + signal.throwIfAborted() + return new Promise((resolve, reject) => { + const cleanup = (): void => signal.removeEventListener('abort', onAbort) + const onAbort = (): void => { + cleanup() + reject(signal.reason) + } + signal.addEventListener('abort', onAbort, { once: true }) + if (signal.aborted) { + onAbort() + return + } + operation.then( + (value) => { + cleanup() + resolve(value) + }, + (error: unknown) => { + cleanup() + reject(error) + } + ) + }) +} + +/** Disposes a resource that finishes opening after cancellation wins the race. */ +async function acquireWithAbort( + operation: Promise, + signal: AbortSignal, + dispose: (resource: T) => Promise +): Promise { + try { + return await awaitWithAbort(operation, signal) + } catch (error) { + if (signal.aborted) { + void operation.then(dispose).catch(() => undefined) + } + throw error + } +} + +/** Closes resources without letting a stuck close extend the fetch deadline. */ +async function closeWithAbort(operation: Promise, signal: AbortSignal): Promise { + if (signal.aborted) { + void operation.catch(() => undefined) + return + } + await awaitWithAbort(operation, signal) +} + +/** Keeps only the newest bounded set while streaming an arbitrarily large directory. */ +function retainNewestLogName(logNames: string[], candidate: string): void { + const insertionIndex = logNames.findIndex((name) => candidate.localeCompare(name) > 0) + if (insertionIndex === -1) { + if (logNames.length < CLI_LOG_LIMIT) { + logNames.push(candidate) + } + return + } + logNames.splice(insertionIndex, 0, candidate) + if (logNames.length > CLI_LOG_LIMIT) { + logNames.pop() + } +} + +/** Streams CLI log names so directory size cannot dictate peak memory use. */ +export async function findNewestAntigravityCliLogNames( + logDirectory: string, + signal: AbortSignal +): Promise { + const openDirectory = opendir(logDirectory) + const directory = await acquireWithAbort(openDirectory, signal, (resource: Dir) => + resource.close() + ) + const logNames: string[] = [] + try { + while (true) { + const entry = await awaitWithAbort(directory.read(), signal) + if (!entry) { + break + } + if (entry.isFile() && /^cli-.*\.log$/i.test(entry.name)) { + retainNewestLogName(logNames, entry.name) + } + } + return logNames + } finally { + await closeWithAbort(directory.close(), signal) + } +} + +/** Fills a bounded section despite short filesystem reads. */ +async function readLogSection( + handle: FileHandle, + position: number, + byteLength: number, + signal: AbortSignal +): Promise { + const buffer = Buffer.allocUnsafe(byteLength) + let totalBytesRead = 0 + while (totalBytesRead < byteLength) { + const { bytesRead } = await awaitWithAbort( + handle.read(buffer, totalBytesRead, byteLength - totalBytesRead, position + totalBytesRead), + signal + ) + if (bytesRead === 0) { + break + } + totalBytesRead += bytesRead + } + // Why: cancellation can land after the final read resolves but before its caller parses the log. + signal.throwIfAborted() + return buffer.subarray(0, totalBytesRead) +} + +/** Reads bounded head and tail sections because listener announcements occur at startup. */ +export async function readAntigravityLogExcerpt( + filePath: string, + signal: AbortSignal +): Promise { + signal.throwIfAborted() + const openFile = open(filePath, 'r') + const handle = await acquireWithAbort(openFile, signal, (resource: FileHandle) => + resource.close() + ) + try { + const stats = await awaitWithAbort(handle.stat(), signal) + if (!stats.isFile()) { + throw new Error('Antigravity log target is not a file') + } + if (stats.size === 0) { + return '' + } + if (stats.size <= LOG_READ_LIMIT_BYTES) { + return (await readLogSection(handle, 0, stats.size, signal)).toString('utf8') + } + const head = await readLogSection(handle, 0, LOG_SECTION_LIMIT_BYTES, signal) + const tail = await readLogSection( + handle, + stats.size - LOG_SECTION_LIMIT_BYTES, + LOG_SECTION_LIMIT_BYTES, + signal + ) + return `${head.toString('utf8')}\n${tail.toString('utf8')}` + } finally { + await closeWithAbort(handle.close(), signal) + } +} diff --git a/src/main/rate-limits/antigravity-loopback-client.test.ts b/src/main/rate-limits/antigravity-loopback-client.test.ts new file mode 100644 index 000000000000..283441931f41 --- /dev/null +++ b/src/main/rate-limits/antigravity-loopback-client.test.ts @@ -0,0 +1,115 @@ +import { createServer } from 'node:http' +import { join } from 'node:path' +import { describe, expect, it } from 'vitest' +import { + AntigravityLoopbackResponseError, + fetchAntigravityQuotaEndpoint, + getAntigravityCliLogDirectory, + getAntigravityLanguageServerLogPath, + parseAntigravityAppConfig, + parseAntigravityCliServerPorts, + parseAntigravityLanguageServerPort +} from './antigravity-loopback-client' + +describe('Antigravity loopback client', () => { + it('uses cross-platform CLI and desktop log paths', () => { + expect(getAntigravityCliLogDirectory('/home/lee')).toBe( + join('/home/lee', '.gemini', 'antigravity-cli', 'log') + ) + expect(getAntigravityLanguageServerLogPath('darwin', '/Users/lee', '/app-data')).toBe( + join('/Users/lee', 'Library', 'Logs', 'Antigravity', 'language_server.log') + ) + for (const platform of ['linux', 'win32'] as const) { + expect(getAntigravityLanguageServerLogPath(platform, '/home/lee', '/home/lee/.config')).toBe( + join('/home/lee/.config', 'Antigravity', 'logs', 'language_server.log') + ) + } + }) + + it('uses the newest listener announcement after a server restart', () => { + const log = [ + 'Language server listening on random port at 40100 for HTTPS (gRPC)', + 'Language server listening on fixed port at 40200 for HTTPS (gRPC)', + 'Language server listening on random port at 40201 for HTTP' + ].join('\n') + + expect(parseAntigravityCliServerPorts(log)).toEqual({ http: 40201, https: 40200 }) + expect(parseAntigravityLanguageServerPort(log)).toBe(40200) + }) + + it('accepts only Antigravity app configuration with a CSRF token', () => { + expect( + parseAntigravityAppConfig( + '' + ) + ).toEqual({ productName: 'antigravity', csrfToken: 'token' }) + expect( + parseAntigravityAppConfig( + '' + ) + ).toBeNull() + }) + + it('accepts formatted app configuration followed by other script statements', () => { + expect( + parseAntigravityAppConfig(``) + ).toEqual({ productName: 'antigravity', csrfToken: 'formatted-token' }) + }) + + it('does not parse app configuration across a closing script boundary', () => { + expect( + parseAntigravityAppConfig( + 'after"};' + ) + ).toBeNull() + }) + + it('rejects a loopback response that exceeds the byte limit', async () => { + const server = createServer((_request, response) => { + response.writeHead(200, { 'content-type': 'application/json' }) + response.end(Buffer.alloc(1024 * 1024 + 1, 'x')) + }) + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)) + const address = server.address() + if (!address || typeof address === 'string') { + throw new Error('Expected a TCP listener') + } + + try { + await expect( + fetchAntigravityQuotaEndpoint('http:', address.port, new AbortController().signal) + ).rejects.toBeInstanceOf(AntigravityLoopbackResponseError) + } finally { + await new Promise((resolve) => server.close(() => resolve())) + } + }) + + it('rejects a loopback response whose body is truncated', async () => { + const server = createServer((_request, response) => { + response.writeHead(200, { 'content-length': '100', 'content-type': 'application/json' }) + response.flushHeaders() + response.write('{}') + setImmediate(() => response.destroy()) + }) + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)) + const address = server.address() + if (!address || typeof address === 'string') { + throw new Error('Expected a TCP listener') + } + + try { + await expect( + fetchAntigravityQuotaEndpoint('http:', address.port, new AbortController().signal) + ).rejects.toBeInstanceOf(AntigravityLoopbackResponseError) + } finally { + await new Promise((resolve) => server.close(() => resolve())) + } + }) +}) diff --git a/src/main/rate-limits/antigravity-loopback-client.ts b/src/main/rate-limits/antigravity-loopback-client.ts new file mode 100644 index 000000000000..7f763eba291c --- /dev/null +++ b/src/main/rate-limits/antigravity-loopback-client.ts @@ -0,0 +1,315 @@ +import { request as httpRequest } from 'node:http' +import { request as httpsRequest } from 'node:https' +import { join } from 'node:path' +import type { ProviderRateLimits } from '../../shared/rate-limit-types' +import { parseAntigravityQuotaSummary } from './antigravity-quota-parser' + +const LANGUAGE_SERVER_LOG_NAME = 'language_server.log' +const REQUEST_TIMEOUT_MS = 1_250 +const MAX_RESPONSE_BYTES = 1024 * 1024 +const QUOTA_SUMMARY_PATH = '/exa.language_server_pb.LanguageServerService/RetrieveUserQuotaSummary' + +export type AntigravityLoopbackProtocol = 'http:' | 'https:' + +type AntigravityAppConfig = { + csrfToken: string + productName: string +} + +export type AntigravityServerPorts = { + http: number | null + https: number | null +} + +/** Distinguishes a complete service reply from a partial response that may be stale. */ +export class AntigravityLoopbackResponseError extends Error { + constructor( + message: string, + readonly responseCompleted = true + ) { + super(message) + this.name = 'AntigravityLoopbackResponseError' + } +} + +/** Rejects untrusted log-derived values outside the TCP port range. */ +function validPort(value: string | undefined): number | null { + const port = Number(value) + return Number.isInteger(port) && port > 0 && port <= 65_535 ? port : null +} + +/** Resolves the per-user directory containing Antigravity CLI logs. */ +export function getAntigravityCliLogDirectory(homePath: string): string { + return join(homePath, '.gemini', 'antigravity-cli', 'log') +} + +/** Returns the newest HTTP and HTTPS listener announcements in one log. */ +export function parseAntigravityCliServerPorts(log: string): AntigravityServerPorts { + const httpsMatches = [ + ...log.matchAll(/language server listening on (?:random|fixed)(?: port)? at (\d+) for HTTPS/gi) + ] + const httpMatches = [ + ...log.matchAll( + /language server listening on (?:random|fixed)(?: port)? at (\d+) for HTTP(?!S)/gi + ) + ] + return { + http: validPort(httpMatches.at(-1)?.[1]), + https: validPort(httpsMatches.at(-1)?.[1]) + } +} + +/** Resolves the Antigravity desktop language-server log for the host OS. */ +export function getAntigravityLanguageServerLogPath( + platform: NodeJS.Platform, + homePath: string, + appDataPath: string +): string { + return platform === 'darwin' + ? join(homePath, 'Library', 'Logs', 'Antigravity', LANGUAGE_SERVER_LOG_NAME) + : join(appDataPath, 'Antigravity', 'logs', LANGUAGE_SERVER_LOG_NAME) +} + +/** Returns the newest desktop HTTPS listener announcement. */ +export function parseAntigravityLanguageServerPort(log: string): number | null { + return parseAntigravityCliServerPorts(log).https +} + +/** Extracts the assigned JSON object without depending on surrounding script formatting. */ +function extractAppConfigJson(html: string): string | null { + const assignment = /window\.__APP_CONFIG__\s*=\s*/.exec(html) + if (!assignment) { + return null + } + const start = assignment.index + assignment[0].length + if (html[start] !== '{') { + return null + } + const scriptEndPattern = /<\/script\s*>/gi + scriptEndPattern.lastIndex = start + const end = scriptEndPattern.exec(html)?.index ?? html.length + let depth = 0 + let escaped = false + let inString = false + + for (let index = start; index < end; index += 1) { + const character = html[index] + if (inString) { + if (escaped) { + escaped = false + } else if (character === '\\') { + escaped = true + } else if (character === '"') { + inString = false + } + continue + } + if (character === '"') { + inString = true + } else if (character === '{') { + depth += 1 + } else if (character === '}') { + depth -= 1 + if (depth === 0) { + return html.slice(start, index + 1) + } + if (depth < 0) { + return null + } + } + } + return null +} + +/** Accepts CSRF configuration only from a page identifying Antigravity. */ +export function parseAntigravityAppConfig(html: string): AntigravityAppConfig | null { + const configJson = extractAppConfigJson(html) + if (!configJson) { + return null + } + try { + const config = JSON.parse(configJson) as Record + return config.productName === 'antigravity' && + typeof config.csrfToken === 'string' && + config.csrfToken.length > 0 + ? { productName: config.productName, csrfToken: config.csrfToken } + : null + } catch { + return null + } +} + +/** Sends a bounded request to a fixed loopback host. */ +export function requestAntigravityLoopbackPage( + protocol: AntigravityLoopbackProtocol, + port: number, + path: string, + signal: AbortSignal, + options?: { body?: string; csrfToken?: string } +): Promise { + return new Promise((resolve, reject) => { + let responseStarted = false + let responseCompleted = false + let requestDeadline: ReturnType | null = null + const clearRequestDeadline = (): void => { + if (requestDeadline) { + clearTimeout(requestDeadline) + requestDeadline = null + } + } + const body = options?.body + const headers: Record = { Connection: 'close' } + if (body !== undefined) { + headers['Content-Type'] = 'application/json' + headers['Content-Length'] = Buffer.byteLength(body) + headers['Connect-Protocol-Version'] = '1' + } + if (options?.csrfToken) { + headers['x-codeium-csrf-token'] = options.csrfToken + } + + // Why: older AGY servers use self-signed HTTPS. The exception is fixed to + // loopback so it cannot weaken certificate checks for a network request. + const request = protocol === 'https:' ? httpsRequest : httpRequest + const req = request( + { + protocol, + hostname: '127.0.0.1', + port, + path, + method: body === undefined ? 'GET' : 'POST', + headers, + rejectUnauthorized: protocol === 'https:' ? false : undefined, + signal, + timeout: REQUEST_TIMEOUT_MS + }, + (response) => { + responseStarted = true + const chunks: Buffer[] = [] + let byteLength = 0 + response.on('data', (chunk: Buffer) => { + byteLength += chunk.length + if (byteLength > MAX_RESPONSE_BYTES) { + req.destroy( + new AntigravityLoopbackResponseError( + 'Antigravity quota response exceeded size limit', + false + ) + ) + return + } + chunks.push(chunk) + }) + response.on('end', () => { + responseCompleted = true + clearRequestDeadline() + if (response.statusCode !== 200) { + reject( + new AntigravityLoopbackResponseError( + `Antigravity quota request failed (${response.statusCode ?? 'unknown'})` + ) + ) + return + } + resolve(Buffer.concat(chunks).toString('utf8')) + }) + response.on('error', (error) => { + clearRequestDeadline() + reject( + new AntigravityLoopbackResponseError( + `Antigravity quota response ended unexpectedly: ${error.message}`, + false + ) + ) + }) + } + ) + // Why: ClientRequest.timeout resets on socket activity, so a trickling + // loopback service also needs a fixed wall-clock deadline. + requestDeadline = setTimeout( + () => + req.destroy( + responseStarted + ? new AntigravityLoopbackResponseError('Antigravity quota response timed out', false) + : new Error('Antigravity quota request timed out') + ), + REQUEST_TIMEOUT_MS + ) + requestDeadline.unref() + req.on('timeout', () => + req.destroy( + responseStarted + ? new AntigravityLoopbackResponseError('Antigravity quota response timed out', false) + : new Error('Antigravity quota request timed out') + ) + ) + req.on('error', (error) => { + clearRequestDeadline() + if ( + responseStarted && + !signal.aborted && + !(error instanceof AntigravityLoopbackResponseError) + ) { + reject( + new AntigravityLoopbackResponseError( + `Antigravity quota response failed: ${error.message}`, + false + ) + ) + return + } + reject(error) + }) + req.on('close', () => { + clearRequestDeadline() + if (!responseCompleted) { + // Why: some supported Node runtimes can close a request without + // forwarding a separate response error. + reject( + new AntigravityLoopbackResponseError( + 'Antigravity quota request closed before the response completed', + false + ) + ) + } + }) + req.end(body) + }) +} + +/** Marks loopback data so downstream telemetry cannot mistake it for OAuth usage. */ +function withLiveSessionMetadata(limits: ProviderRateLimits): ProviderRateLimits { + return { + ...limits, + usageMetadata: { + source: 'live-session', + attemptedSources: ['live-session'], + credentialSource: 'agy-local-service', + authProvenance: 'antigravity' + } + } +} + +/** Fetches and validates one Antigravity quota endpoint response. */ +export async function fetchAntigravityQuotaEndpoint( + protocol: AntigravityLoopbackProtocol, + port: number, + signal: AbortSignal, + csrfToken?: string +): Promise { + const response = await requestAntigravityLoopbackPage( + protocol, + port, + QUOTA_SUMMARY_PATH, + signal, + { body: '{}', csrfToken } + ) + let data: unknown + try { + data = JSON.parse(response) as unknown + } catch { + throw new AntigravityLoopbackResponseError('Antigravity quota response was not valid JSON') + } + const parsed = parseAntigravityQuotaSummary(data) + return parsed ? withLiveSessionMetadata(parsed) : null +} diff --git a/src/main/rate-limits/antigravity-loopback-request-close.test.ts b/src/main/rate-limits/antigravity-loopback-request-close.test.ts new file mode 100644 index 000000000000..69580a09691f --- /dev/null +++ b/src/main/rate-limits/antigravity-loopback-request-close.test.ts @@ -0,0 +1,118 @@ +import { EventEmitter } from 'node:events' +import { afterEach, describe, expect, it, vi } from 'vitest' + +const { httpRequestMock } = vi.hoisted(() => ({ httpRequestMock: vi.fn() })) + +vi.mock('node:http', () => ({ request: httpRequestMock })) + +import { + AntigravityLoopbackResponseError, + requestAntigravityLoopbackPage +} from './antigravity-loopback-client' + +describe('Antigravity loopback request closure', () => { + afterEach(() => { + vi.useRealTimers() + }) + + it('rejects when the request closes without a response event', async () => { + const request = new EventEmitter() as EventEmitter & { + destroy: (error?: Error) => void + end: () => void + } + request.destroy = (error?: Error) => { + if (error) { + request.emit('error', error) + } + request.emit('close') + } + request.end = () => queueMicrotask(() => request.emit('close')) + httpRequestMock.mockReturnValueOnce(request) + + await expect( + requestAntigravityLoopbackPage( + 'http:', + 40_200, + '/exa.language_server_pb.LanguageServerService/RetrieveUserQuotaSummary', + new AbortController().signal + ) + ).rejects.toThrow( + new AntigravityLoopbackResponseError( + 'Antigravity quota request closed before the response completed', + false + ) + ) + }) + + it('enforces a wall-clock deadline after a response starts', async () => { + vi.useFakeTimers() + const response = Object.assign(new EventEmitter(), { statusCode: 200 }) + const request = new EventEmitter() as EventEmitter & { + destroy: (error?: Error) => void + end: () => void + } + const destroy = vi.fn((error?: Error) => { + if (error) { + request.emit('error', error) + } + request.emit('close') + }) + request.destroy = destroy + request.end = () => undefined + httpRequestMock.mockImplementationOnce( + (_options: unknown, onResponse: (value: EventEmitter & { statusCode: number }) => void) => { + queueMicrotask(() => { + onResponse(response) + response.emit('data', Buffer.from('{')) + }) + return request + } + ) + + const result = requestAntigravityLoopbackPage( + 'http:', + 40_200, + '/exa.language_server_pb.LanguageServerService/RetrieveUserQuotaSummary', + new AbortController().signal + ) + const rejection = expect(result).rejects.toMatchObject({ + name: 'AntigravityLoopbackResponseError', + message: 'Antigravity quota response timed out', + responseCompleted: false + }) + await vi.advanceTimersByTimeAsync(1_249) + expect(destroy).not.toHaveBeenCalled() + await vi.advanceTimersByTimeAsync(1) + + await rejection + expect(destroy).toHaveBeenCalledOnce() + }) + + it('clears the wall-clock deadline after an early close', async () => { + vi.useFakeTimers() + const request = new EventEmitter() as EventEmitter & { + destroy: (error?: Error) => void + end: () => void + } + const destroy = vi.fn() + request.destroy = destroy + request.end = () => queueMicrotask(() => request.emit('close')) + httpRequestMock.mockReturnValueOnce(request) + + const result = requestAntigravityLoopbackPage( + 'http:', + 40_200, + '/exa.language_server_pb.LanguageServerService/RetrieveUserQuotaSummary', + new AbortController().signal + ) + await expect(result).rejects.toThrow( + new AntigravityLoopbackResponseError( + 'Antigravity quota request closed before the response completed', + false + ) + ) + await vi.advanceTimersByTimeAsync(1_250) + + expect(destroy).not.toHaveBeenCalled() + }) +}) diff --git a/src/main/rate-limits/antigravity-quota-parser.test.ts b/src/main/rate-limits/antigravity-quota-parser.test.ts new file mode 100644 index 000000000000..75e8a9ca590c --- /dev/null +++ b/src/main/rate-limits/antigravity-quota-parser.test.ts @@ -0,0 +1,97 @@ +import { describe, expect, it } from 'vitest' +import { parseAntigravityQuotaSummary } from './antigravity-quota-parser' + +const quotaSummary = { + response: { + groups: [ + { + displayName: 'Gemini Models', + buckets: [ + { + bucketId: 'gemini-weekly', + remainingFraction: 0.916, + resetTime: '2026-07-16T21:59:04Z' + }, + { + bucketId: 'gemini-5h', + remainingFraction: 1, + resetTime: '2026-07-14T16:30:11Z' + } + ] + }, + { + displayName: 'Claude and GPT models', + buckets: [ + { + bucketId: '3p-weekly', + remainingFraction: 0.988, + resetTime: '2026-07-21T11:28:50Z' + }, + { + bucketId: '3p-5h', + remainingFraction: 0.964, + resetTime: '2026-07-14T16:28:50Z' + } + ] + } + ] + } +} + +describe('parseAntigravityQuotaSummary', () => { + it('keeps both quota families while deriving compatible summary windows', () => { + const result = parseAntigravityQuotaSummary(quotaSummary, 1234) + + expect(result).toMatchObject({ + provider: 'antigravity', + session: { usedPercent: 4, windowMinutes: 300 }, + weekly: { usedPercent: 8, windowMinutes: 10_080 }, + updatedAt: 1234, + status: 'ok' + }) + expect(result?.buckets?.map(({ name, usedPercent }) => ({ name, usedPercent }))).toEqual([ + { name: 'Gemini 5h', usedPercent: 0 }, + { name: 'Gemini wk', usedPercent: 8 }, + { name: 'Claude/GPT 5h', usedPercent: 4 }, + { name: 'Claude/GPT wk', usedPercent: 1 } + ]) + }) + + it('accepts the root-level groups shape returned by some AGY versions', () => { + expect(parseAntigravityQuotaSummary(quotaSummary.response)?.buckets).toHaveLength(4) + }) + + it('keeps recognized buckets when reset metadata is absent', () => { + const withoutReset = structuredClone(quotaSummary) + delete (withoutReset.response.groups[0].buckets[0] as { resetTime?: string }).resetTime + + expect(parseAntigravityQuotaSummary(withoutReset)?.buckets?.[1]).toMatchObject({ + name: 'Gemini wk', + usedPercent: 8, + resetsAt: null + }) + }) + + it('keeps the most constrained value when a response duplicates a bucket identity', () => { + const duplicated = structuredClone(quotaSummary) + duplicated.response.groups[1].buckets.push({ + bucketId: 'gemini-weekly', + remainingFraction: 0.25, + resetTime: '2026-07-16T21:59:04Z' + }) + + expect(parseAntigravityQuotaSummary(duplicated)?.buckets?.[1]?.usedPercent).toBe(75) + }) + + it('rejects malformed and unknown-only responses', () => { + expect(parseAntigravityQuotaSummary({ response: { groups: [] } })).toBeNull() + expect( + parseAntigravityQuotaSummary({ + response: { + groups: [{ buckets: [{ bucketId: 'future-window', remainingFraction: 0.5 }] }] + } + }) + ).toBeNull() + expect(parseAntigravityQuotaSummary(null)).toBeNull() + }) +}) diff --git a/src/main/rate-limits/antigravity-quota-parser.ts b/src/main/rate-limits/antigravity-quota-parser.ts new file mode 100644 index 000000000000..797df8c499fa --- /dev/null +++ b/src/main/rate-limits/antigravity-quota-parser.ts @@ -0,0 +1,143 @@ +import type { + ProviderRateLimits, + RateLimitBucket, + RateLimitWindow +} from '../../shared/rate-limit-types' + +type AntigravityBucketSpec = { + name: string + windowMinutes: number +} + +const ANTIGRAVITY_BUCKET_SPECS: Record = { + 'gemini-5h': { name: 'Gemini 5h', windowMinutes: 300 }, + 'gemini-weekly': { name: 'Gemini wk', windowMinutes: 10_080 }, + '3p-5h': { name: 'Claude/GPT 5h', windowMinutes: 300 }, + '3p-weekly': { name: 'Claude/GPT wk', windowMinutes: 10_080 } +} + +const ANTIGRAVITY_BUCKET_ORDER = ['gemini-5h', 'gemini-weekly', '3p-5h', '3p-weekly'] + +type ParsedQuotaBucket = { + bucketId: string + remainingFraction: number + resetTime: string | null +} + +/** Accepts the response envelopes emitted by current and earlier AGY runtimes. */ +function quotaGroups(value: unknown): unknown[] | null { + if (!value || typeof value !== 'object') { + return null + } + const root = value as { groups?: unknown; response?: { groups?: unknown } } + const groups = root.response?.groups ?? root.groups + return Array.isArray(groups) ? groups : null +} + +/** Rejects malformed quota entries before they can affect UI percentages. */ +function parseQuotaBucket(value: unknown): ParsedQuotaBucket | null { + if (!value || typeof value !== 'object') { + return null + } + const bucket = value as Record + if ( + typeof bucket.bucketId !== 'string' || + typeof bucket.remainingFraction !== 'number' || + !Number.isFinite(bucket.remainingFraction) || + (bucket.resetTime !== undefined && + bucket.resetTime !== null && + typeof bucket.resetTime !== 'string') + ) { + return null + } + return { + bucketId: bucket.bucketId, + remainingFraction: bucket.remainingFraction, + resetTime: typeof bucket.resetTime === 'string' ? bucket.resetTime : null + } +} + +/** Keeps only known identities and the safest value for duplicate entries. */ +function collectQuotaBuckets(groups: unknown[]): Map { + const buckets = new Map() + for (const value of groups) { + if (!value || typeof value !== 'object') { + continue + } + const rawBuckets = (value as { buckets?: unknown }).buckets + if (!Array.isArray(rawBuckets)) { + continue + } + for (const rawBucket of rawBuckets) { + const bucket = parseQuotaBucket(rawBucket) + if (bucket && Object.hasOwn(ANTIGRAVITY_BUCKET_SPECS, bucket.bucketId)) { + const current = buckets.get(bucket.bucketId) + // Why: duplicate identities are ambiguous; keeping the lower remaining + // value avoids quietly understating usage from a malformed response. + if (!current || bucket.remainingFraction < current.remainingFraction) { + buckets.set(bucket.bucketId, bucket) + } + } + } + } + return buckets +} + +/** Normalizes AGY fractions and reset timestamps into Orca's shared model. */ +function toRateLimitBucket(bucket: ParsedQuotaBucket): RateLimitBucket { + const spec = ANTIGRAVITY_BUCKET_SPECS[bucket.bucketId]! + const remainingFraction = Math.min(1, Math.max(0, bucket.remainingFraction)) + const resetsAt = bucket.resetTime ? Date.parse(bucket.resetTime) : Number.NaN + return { + name: spec.name, + usedPercent: Math.round((1 - remainingFraction) * 100), + windowMinutes: spec.windowMinutes, + resetsAt: Number.isFinite(resetsAt) ? resetsAt : null, + resetDescription: null + } +} + +/** Derives compatibility summaries without discarding the four detailed buckets. */ +function mostConstrainedWindow( + buckets: RateLimitBucket[], + windowMinutes: number +): RateLimitWindow | null { + const matching = buckets.filter((bucket) => bucket.windowMinutes === windowMinutes) + if (matching.length === 0) { + return null + } + const mostConstrained = matching.reduce((current, candidate) => + candidate.usedPercent > current.usedPercent ? candidate : current + ) + const { name: _name, ...window } = mostConstrained + return window +} + +/** Maps AGY's two quota families into the current Detailed/Compact bucket model. */ +export function parseAntigravityQuotaSummary( + value: unknown, + updatedAt = Date.now() +): ProviderRateLimits | null { + const groups = quotaGroups(value) + if (!groups) { + return null + } + const rawBuckets = collectQuotaBuckets(groups) + const buckets = ANTIGRAVITY_BUCKET_ORDER.map((id) => rawBuckets.get(id)) + .filter((bucket): bucket is ParsedQuotaBucket => bucket !== undefined) + .map(toRateLimitBucket) + if (buckets.length === 0) { + return null + } + return { + provider: 'antigravity', + // Why: summary windows keep non-roster consumers useful; Detailed renders + // the four named buckets and Compact independently selects their tightest. + session: mostConstrainedWindow(buckets, 300), + weekly: mostConstrainedWindow(buckets, 10_080), + buckets, + updatedAt, + error: null, + status: 'ok' + } +} diff --git a/src/main/rate-limits/antigravity-usage-fetcher.test.ts b/src/main/rate-limits/antigravity-usage-fetcher.test.ts new file mode 100644 index 000000000000..6363dd53f903 --- /dev/null +++ b/src/main/rate-limits/antigravity-usage-fetcher.test.ts @@ -0,0 +1,363 @@ +import { createServer } from 'node:http' +import { createServer as createHttpsServer } from 'node:https' +import { mkdir, mkdtemp, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { describe, expect, it, vi } from 'vitest' + +vi.mock('electron', () => ({ app: { getPath: vi.fn() } })) + +import { fetchAntigravityRateLimits } from './antigravity-usage-fetcher' +import { getAntigravityCliLogDirectory } from './antigravity-loopback-client' +import { + LOCAL_HTTPS_TEST_CERTIFICATE, + LOCAL_HTTPS_TEST_PRIVATE_KEY +} from '../browser/browser-local-https-test-certificate' + +const quotaSummary = { + response: { + groups: [ + { + buckets: [ + { bucketId: 'gemini-weekly', remainingFraction: 0.92 }, + { bucketId: 'gemini-5h', remainingFraction: 1 }, + { bucketId: '3p-weekly', remainingFraction: 0.99 }, + { bucketId: '3p-5h', remainingFraction: 0.96 } + ] + } + ] + } +} + +describe('Antigravity language-server discovery', () => { + it.each([ + { + announcementPosition: 'beginning', + makeLog: (announcement: string) => `${announcement}\n${'long log data\n'.repeat(25_000)}` + }, + { + announcementPosition: 'end', + makeLog: (announcement: string) => `${'long log data\n'.repeat(25_000)}${announcement}` + } + ])( + 'falls past a newer stale log when the live listener announcement is at the $announcementPosition of a long log', + async ({ makeLog }) => { + const homePath = await mkdtemp(join(tmpdir(), 'orca-antigravity-cli-')) + const logDirectory = getAntigravityCliLogDirectory(homePath) + let requestBody = '' + const server = createServer((request, response) => { + request.on('data', (chunk: Buffer) => { + requestBody += chunk.toString('utf8') + }) + request.on('end', () => { + response.writeHead(200, { 'content-type': 'application/json' }) + response.end(JSON.stringify(quotaSummary)) + }) + }) + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)) + const address = server.address() + if (!address || typeof address === 'string') { + throw new Error('Expected a TCP listener') + } + + try { + await mkdir(logDirectory, { recursive: true }) + await writeFile( + join(logDirectory, 'cli-20260714_123131.log'), + 'Language server listening on random port at 1 for HTTP' + ) + await writeFile( + join(logDirectory, 'cli-20260714_103225.log'), + makeLog(`Language server listening on random port at ${address.port} for HTTP`) + ) + + const result = await fetchAntigravityRateLimits({ + homePath, + appDataPath: join(homePath, 'app-data') + }) + + expect(requestBody).toBe('{}') + expect(result).toMatchObject({ + provider: 'antigravity', + session: { usedPercent: 4 }, + weekly: { usedPercent: 8 }, + status: 'ok', + usageMetadata: { + source: 'live-session', + credentialSource: 'agy-local-service', + authProvenance: 'antigravity' + } + }) + expect(result.buckets).toHaveLength(4) + } finally { + await new Promise((resolve, reject) => { + server.close((error) => (error ? reject(error) : resolve())) + }) + await rm(homePath, { recursive: true, force: true }) + } + } + ) + + it.each([ + { responseName: 'not-signed-in error', statusCode: 500, body: 'not signed in' }, + { responseName: 'empty quota', statusCode: 200, body: '{"response":{"groups":[]}}' } + ])( + 'does not fall back to an older account after the newest runtime answers with $responseName', + async ({ statusCode, body }) => { + const homePath = await mkdtemp(join(tmpdir(), 'orca-antigravity-account-')) + const logDirectory = getAntigravityCliLogDirectory(homePath) + let olderRequests = 0 + const newestServer = createServer((_request, response) => { + response.writeHead(statusCode, { 'content-type': 'application/json' }) + response.end(body) + }) + const olderServer = createServer((_request, response) => { + olderRequests += 1 + response.writeHead(200, { 'content-type': 'application/json' }) + response.end(JSON.stringify(quotaSummary)) + }) + await new Promise((resolve) => newestServer.listen(0, '127.0.0.1', resolve)) + await new Promise((resolve) => olderServer.listen(0, '127.0.0.1', resolve)) + const newestAddress = newestServer.address() + const olderAddress = olderServer.address() + if ( + !newestAddress || + typeof newestAddress === 'string' || + !olderAddress || + typeof olderAddress === 'string' + ) { + throw new Error('Expected TCP listeners') + } + + try { + await mkdir(logDirectory, { recursive: true }) + await writeFile( + join(logDirectory, 'cli-20260714_123131.log'), + `Language server listening on random port at ${newestAddress.port} for HTTP` + ) + await writeFile( + join(logDirectory, 'cli-20260714_103225.log'), + `Language server listening on random port at ${olderAddress.port} for HTTP` + ) + + const result = await fetchAntigravityRateLimits({ + homePath, + appDataPath: join(homePath, 'app-data') + }) + + expect(result).toMatchObject({ + provider: 'antigravity', + status: 'error', + usageMetadata: { failureKind: 'usage-unavailable' } + }) + expect(olderRequests).toBe(0) + } finally { + await Promise.all([ + new Promise((resolve) => newestServer.close(() => resolve())), + new Promise((resolve) => olderServer.close(() => resolve())) + ]) + await rm(homePath, { recursive: true, force: true }) + } + } + ) + + it('uses desktop CSRF configuration after an unauthenticated quota response', async () => { + const homePath = await mkdtemp(join(tmpdir(), 'orca-antigravity-desktop-')) + const appDataPath = join(homePath, 'app-data') + const logDirectory = join(appDataPath, 'Antigravity', 'logs') + const requestHeaders: (string | undefined)[] = [] + const server = createHttpsServer( + { key: LOCAL_HTTPS_TEST_PRIVATE_KEY, cert: LOCAL_HTTPS_TEST_CERTIFICATE }, + (request, response) => { + if (request.url === '/') { + response.writeHead(200) + response.end( + '' + ) + return + } + requestHeaders.push(request.headers['x-codeium-csrf-token'] as string | undefined) + if (request.headers['x-codeium-csrf-token'] !== 'desktop-token') { + response.writeHead(403) + response.end() + return + } + response.writeHead(200, { 'content-type': 'application/json' }) + response.end(JSON.stringify(quotaSummary)) + } + ) + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)) + const address = server.address() + if (!address || typeof address === 'string') { + throw new Error('Expected a TCP listener') + } + + try { + await mkdir(logDirectory, { recursive: true }) + await writeFile( + join(logDirectory, 'language_server.log'), + `Language server listening on random port at ${address.port} for HTTPS` + ) + + const result = await fetchAntigravityRateLimits({ + homePath, + appDataPath, + platform: 'linux' + }) + + expect(result.status).toBe('ok') + expect(requestHeaders).toEqual([undefined, 'desktop-token']) + } finally { + await new Promise((resolve) => server.close(() => resolve())) + await rm(homePath, { recursive: true, force: true }) + } + }) + + it('reports unavailable when no local runtime can be discovered', async () => { + const homePath = await mkdtemp(join(tmpdir(), 'orca-antigravity-missing-')) + try { + await expect( + fetchAntigravityRateLimits({ + homePath, + appDataPath: join(homePath, 'app-data') + }) + ).resolves.toMatchObject({ + provider: 'antigravity', + status: 'unavailable', + usageMetadata: { failureKind: 'cli-unavailable' } + }) + } finally { + await rm(homePath, { recursive: true, force: true }) + } + }) + + it('reports unavailable when only a stale listener announcement remains', async () => { + const homePath = await mkdtemp(join(tmpdir(), 'orca-antigravity-stale-listener-')) + const logDirectory = getAntigravityCliLogDirectory(homePath) + const server = createServer() + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)) + const address = server.address() + if (!address || typeof address === 'string') { + throw new Error('Expected a TCP listener') + } + await new Promise((resolve) => server.close(() => resolve())) + + try { + await mkdir(logDirectory, { recursive: true }) + await writeFile( + join(logDirectory, 'cli-20260714_123131.log'), + `Language server listening on random port at ${address.port} for HTTP` + ) + + await expect( + fetchAntigravityRateLimits({ + homePath, + appDataPath: join(homePath, 'app-data') + }) + ).resolves.toMatchObject({ + provider: 'antigravity', + status: 'unavailable', + usageMetadata: { failureKind: 'cli-unavailable' } + }) + } finally { + await rm(homePath, { recursive: true, force: true }) + } + }) + + it('honors service cancellation before starting discovery', async () => { + const controller = new AbortController() + controller.abort() + + await expect(fetchAntigravityRateLimits({ signal: controller.signal })).rejects.toThrow( + /aborted/i + ) + }) + + it('reports the total discovery deadline as a timeout', async () => { + const homePath = await mkdtemp(join(tmpdir(), 'orca-antigravity-timeout-')) + const logDirectory = getAntigravityCliLogDirectory(homePath) + let markStarted: (() => void) | undefined + const started = new Promise((resolve) => { + markStarted = resolve + }) + // Why: enough distinct silent endpoints let the outer discovery deadline + // win while each request remains independently wall-clock bounded. + const servers = Array.from({ length: 6 }, () => createServer(() => markStarted?.())) + await Promise.all( + servers.map( + (server) => new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)) + ) + ) + const ports = servers.map((server) => { + const address = server.address() + if (!address || typeof address === 'string') { + throw new Error('Expected a TCP listener') + } + return address.port + }) + + try { + await mkdir(logDirectory, { recursive: true }) + await Promise.all( + ports.map((port, index) => + writeFile( + join(logDirectory, `cli-20260714_12313${index}.log`), + `Language server listening on random port at ${port} for HTTP` + ) + ) + ) + const resultPromise = fetchAntigravityRateLimits({ + homePath, + appDataPath: join(homePath, 'app-data') + }) + await started + + await expect(resultPromise).resolves.toMatchObject({ + status: 'error', + error: 'Antigravity usage lookup timed out', + usageMetadata: { failureKind: 'usage-unavailable' } + }) + } finally { + await Promise.all( + servers.map((server) => new Promise((resolve) => server.close(() => resolve()))) + ) + await rm(homePath, { recursive: true, force: true }) + } + }, 10_000) + + it('aborts an in-flight loopback request', async () => { + const homePath = await mkdtemp(join(tmpdir(), 'orca-antigravity-abort-')) + const logDirectory = getAntigravityCliLogDirectory(homePath) + let markStarted: (() => void) | undefined + const started = new Promise((resolve) => { + markStarted = resolve + }) + const server = createServer(() => markStarted?.()) + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)) + const address = server.address() + if (!address || typeof address === 'string') { + throw new Error('Expected a TCP listener') + } + const controller = new AbortController() + + try { + await mkdir(logDirectory, { recursive: true }) + await writeFile( + join(logDirectory, 'cli-20260714_123131.log'), + `Language server listening on random port at ${address.port} for HTTP` + ) + const fetchResult = fetchAntigravityRateLimits({ + signal: controller.signal, + homePath, + appDataPath: join(homePath, 'app-data') + }) + await started + controller.abort() + + await expect(fetchResult).rejects.toThrow(/aborted/i) + } finally { + await new Promise((resolve) => server.close(() => resolve())) + await rm(homePath, { recursive: true, force: true }) + } + }) +}) diff --git a/src/main/rate-limits/antigravity-usage-fetcher.ts b/src/main/rate-limits/antigravity-usage-fetcher.ts new file mode 100644 index 000000000000..12fc71335cda --- /dev/null +++ b/src/main/rate-limits/antigravity-usage-fetcher.ts @@ -0,0 +1,235 @@ +import { homedir } from 'node:os' +import { join } from 'node:path' +import { app } from 'electron' +import type { ProviderRateLimits } from '../../shared/rate-limit-types' +import { + findNewestAntigravityCliLogNames, + readAntigravityLogExcerpt +} from './antigravity-log-discovery' +import { + fetchAntigravityQuotaEndpoint, + getAntigravityCliLogDirectory, + getAntigravityLanguageServerLogPath, + parseAntigravityAppConfig, + parseAntigravityCliServerPorts, + parseAntigravityLanguageServerPort, + requestAntigravityLoopbackPage, + AntigravityLoopbackResponseError, + type AntigravityLoopbackProtocol +} from './antigravity-loopback-client' + +const ENDPOINT_ATTEMPT_LIMIT = 8 +const FETCH_TIMEOUT_MS = 6_000 +const FETCH_TIMEOUT_MESSAGE = 'Antigravity usage lookup timed out' + +type FetchAttempt = { + discovered: boolean + answered: boolean + limits: ProviderRateLimits | null +} + +export type AntigravityUsageFetchOptions = { + signal?: AbortSignal + homePath?: string + appDataPath?: string + platform?: NodeJS.Platform +} + +/** Tries newest CLI runtimes first so older processes cannot replace the active account. */ +async function fetchFromCliLogs(homePath: string, signal: AbortSignal): Promise { + const logDirectory = getAntigravityCliLogDirectory(homePath) + let logNames: string[] + try { + logNames = await findNewestAntigravityCliLogNames(logDirectory, signal) + } catch { + signal.throwIfAborted() + return { discovered: false, answered: false, limits: null } + } + const attemptedEndpoints = new Set() + let discovered = false + + for (const logName of logNames) { + signal.throwIfAborted() + let log: string + try { + log = await readAntigravityLogExcerpt(join(logDirectory, logName), signal) + } catch { + signal.throwIfAborted() + // Why: AGY rotates logs between directory listing and reading; one + // vanished candidate must not hide another live runtime. + continue + } + const ports = parseAntigravityCliServerPorts(log) + const endpoints: { protocol: AntigravityLoopbackProtocol; port: number | null }[] = [ + { protocol: 'http:', port: ports.http }, + { protocol: 'https:', port: ports.https } + ] + let answered = false + for (const endpoint of endpoints) { + if (!endpoint.port) { + continue + } + discovered = true + const key = `${endpoint.protocol}//127.0.0.1:${endpoint.port}` + if (attemptedEndpoints.has(key) || attemptedEndpoints.size >= ENDPOINT_ATTEMPT_LIMIT) { + continue + } + attemptedEndpoints.add(key) + try { + const limits = await fetchAntigravityQuotaEndpoint(endpoint.protocol, endpoint.port, signal) + answered = true + if (limits) { + return { discovered: true, answered: true, limits } + } + } catch (error) { + answered ||= error instanceof AntigravityLoopbackResponseError && error.responseCompleted + signal.throwIfAborted() + // A newer one-shot AGY command can leave a stale log above a live session. + } + } + if (answered) { + // Why: a responding newest runtime owns the current account even when it + // is warming up; older runtimes can retain a signed-out account in memory. + return { discovered: true, answered: true, limits: null } + } + } + return { discovered, answered: false, limits: null } +} + +/** Falls back to the desktop runtime only when no newer CLI runtime answered. */ +async function fetchFromDesktopApp( + platform: NodeJS.Platform, + homePath: string, + appDataPath: string, + signal: AbortSignal +): Promise { + let port: number | null + try { + const logPath = getAntigravityLanguageServerLogPath(platform, homePath, appDataPath) + port = parseAntigravityLanguageServerPort(await readAntigravityLogExcerpt(logPath, signal)) + } catch { + signal.throwIfAborted() + return { discovered: false, answered: false, limits: null } + } + if (!port) { + return { discovered: false, answered: false, limits: null } + } + + let answered = false + try { + const limits = await fetchAntigravityQuotaEndpoint('https:', port, signal) + answered = true + if (limits) { + return { discovered: true, answered: true, limits } + } + } catch (error) { + answered ||= error instanceof AntigravityLoopbackResponseError && error.responseCompleted + signal.throwIfAborted() + } + + try { + const html = await requestAntigravityLoopbackPage('https:', port, '/', signal) + answered = true + const config = parseAntigravityAppConfig(html) + const limits = config + ? await fetchAntigravityQuotaEndpoint('https:', port, signal, config.csrfToken) + : null + return { discovered: true, answered: true, limits } + } catch (error) { + answered ||= error instanceof AntigravityLoopbackResponseError && error.responseCompleted + signal.throwIfAborted() + return { discovered: true, answered, limits: null } + } +} + +/** Clears stale quota windows while retaining a machine-readable failure reason. */ +function emptyResult( + status: 'error' | 'unavailable', + error: string, + failureKind: 'usage-unavailable' | 'cli-unavailable' +): ProviderRateLimits { + return { + provider: 'antigravity', + session: null, + weekly: null, + updatedAt: Date.now(), + error, + status, + usageMetadata: { + source: 'live-session', + attemptedSources: ['live-session'], + failureKind, + credentialSource: 'agy-local-service', + authProvenance: 'antigravity' + } + } +} + +/** Reads Antigravity quotas from the newest responsive host-local runtime. */ +export async function fetchAntigravityRateLimits( + options: AntigravityUsageFetchOptions = {} +): Promise { + const controller = new AbortController() + let timedOut = false + const timeout = setTimeout(() => { + timedOut = true + controller.abort(new Error(FETCH_TIMEOUT_MESSAGE)) + }, FETCH_TIMEOUT_MS) + const onAbort = (): void => controller.abort(options.signal?.reason) + options.signal?.addEventListener('abort', onAbort, { once: true }) + if (options.signal?.aborted) { + onAbort() + } + + try { + controller.signal.throwIfAborted() + const homePath = options.homePath ?? homedir() + const cliAttempt = await fetchFromCliLogs(homePath, controller.signal) + if (cliAttempt.limits) { + return cliAttempt.limits + } + if (cliAttempt.answered) { + return emptyResult( + 'error', + 'Antigravity model quota summary is unavailable', + 'usage-unavailable' + ) + } + const appDataPath = options.appDataPath ?? app.getPath('appData') + const desktopAttempt = await fetchFromDesktopApp( + options.platform ?? process.platform, + homePath, + appDataPath, + controller.signal + ) + if (desktopAttempt.limits) { + return desktopAttempt.limits + } + controller.signal.throwIfAborted() + // Why: listener announcements outlive their processes; only a runtime that + // completed a response should keep a stopped desktop/CLI surfaced as an error. + return desktopAttempt.answered + ? emptyResult('error', 'Antigravity model quota summary is unavailable', 'usage-unavailable') + : emptyResult( + 'unavailable', + 'Antigravity local usage service is not running', + 'cli-unavailable' + ) + } catch (error) { + if (options.signal?.aborted) { + throw error + } + return emptyResult( + 'error', + timedOut + ? FETCH_TIMEOUT_MESSAGE + : error instanceof Error + ? error.message + : 'Unknown Antigravity usage error', + 'usage-unavailable' + ) + } finally { + clearTimeout(timeout) + options.signal?.removeEventListener('abort', onAbort) + } +} diff --git a/src/main/rate-limits/antigravity-usage-log-rotation.test.ts b/src/main/rate-limits/antigravity-usage-log-rotation.test.ts new file mode 100644 index 000000000000..1f99c6e97330 --- /dev/null +++ b/src/main/rate-limits/antigravity-usage-log-rotation.test.ts @@ -0,0 +1,210 @@ +import type { Dirent } from 'node:fs' +import type * as FsPromises from 'node:fs/promises' +import { afterEach, describe, expect, it, vi } from 'vitest' +import type { ProviderRateLimits } from '../../shared/rate-limit-types' +import type * as AntigravityLoopbackClient from './antigravity-loopback-client' + +const { fetchQuotaEndpointMock, openMock, opendirMock } = vi.hoisted(() => ({ + fetchQuotaEndpointMock: vi.fn(), + openMock: vi.fn(), + opendirMock: vi.fn() +})) + +vi.mock('electron', () => ({ app: { getPath: vi.fn() } })) + +vi.mock('node:fs/promises', async (importOriginal) => ({ + ...(await importOriginal()), + open: openMock, + opendir: opendirMock +})) + +vi.mock('./antigravity-loopback-client', async (importOriginal) => ({ + ...(await importOriginal()), + fetchAntigravityQuotaEndpoint: fetchQuotaEndpointMock +})) + +import { fetchAntigravityRateLimits } from './antigravity-usage-fetcher' +import { AntigravityLoopbackResponseError } from './antigravity-loopback-client' +import { readAntigravityLogExcerpt } from './antigravity-log-discovery' + +function logEntry(name: string): Dirent { + return { name, isFile: () => true } as Dirent +} + +function readableLog(contents: string): unknown { + const bytes = Buffer.from(contents) + return { + stat: vi.fn(async () => ({ isFile: () => true, size: bytes.length })), + read: vi.fn(async (buffer: Buffer, offset: number, length: number, position: number) => { + const bytesRead = bytes.copy(buffer, offset, position, position + length) + return { buffer, bytesRead } + }), + close: vi.fn(async () => undefined) + } +} + +function readableDirectory(entries: Dirent[]): unknown { + let index = 0 + return { + read: vi.fn(async () => entries[index++] ?? null), + close: vi.fn(async () => undefined) + } +} + +function generatedLogDirectory(entryCount: number): unknown { + let index = 0 + return { + read: vi.fn(async () => { + if (index >= entryCount) { + return null + } + const name = `cli-20260815_${String(index).padStart(6, '0')}.log` + index += 1 + return logEntry(name) + }), + close: vi.fn(async () => undefined) + } +} + +describe('Antigravity rotating-log discovery', () => { + afterEach(() => { + vi.clearAllMocks() + vi.useRealTimers() + }) + + it('continues to an older candidate when the newest selected log disappears', async () => { + const limits: ProviderRateLimits = { + provider: 'antigravity', + session: { usedPercent: 25, windowMinutes: 300, resetsAt: null, resetDescription: null }, + weekly: null, + updatedAt: 1, + error: null, + status: 'ok' + } + opendirMock.mockResolvedValue( + readableDirectory([logEntry('cli-20260815_120000.log'), logEntry('cli-20260815_110000.log')]) + ) + openMock + .mockRejectedValueOnce(Object.assign(new Error('rotated away'), { code: 'ENOENT' })) + .mockResolvedValueOnce( + readableLog('Language server listening on random port at 40200 for HTTP') + ) + fetchQuotaEndpointMock.mockResolvedValue(limits) + + await expect( + fetchAntigravityRateLimits({ homePath: '/home/test', appDataPath: '/app-data' }) + ).resolves.toBe(limits) + expect(openMock).toHaveBeenCalledTimes(2) + expect(fetchQuotaEndpointMock).toHaveBeenCalledWith('http:', 40_200, expect.any(AbortSignal)) + }) + + it('falls back after a newer runtime starts but does not complete its response', async () => { + const limits: ProviderRateLimits = { + provider: 'antigravity', + session: { usedPercent: 25, windowMinutes: 300, resetsAt: null, resetDescription: null }, + weekly: null, + updatedAt: 1, + error: null, + status: 'ok' + } + opendirMock.mockResolvedValue( + readableDirectory([logEntry('cli-20260815_120000.log'), logEntry('cli-20260815_110000.log')]) + ) + openMock + .mockResolvedValueOnce( + readableLog('Language server listening on random port at 40201 for HTTP') + ) + .mockResolvedValueOnce( + readableLog('Language server listening on random port at 40200 for HTTP') + ) + fetchQuotaEndpointMock + .mockRejectedValueOnce( + new AntigravityLoopbackResponseError('Antigravity quota response timed out', false) + ) + .mockResolvedValueOnce(limits) + + await expect( + fetchAntigravityRateLimits({ homePath: '/home/test', appDataPath: '/app-data' }) + ).resolves.toBe(limits) + expect(fetchQuotaEndpointMock).toHaveBeenNthCalledWith( + 1, + 'http:', + 40_201, + expect.any(AbortSignal) + ) + expect(fetchQuotaEndpointMock).toHaveBeenNthCalledWith( + 2, + 'http:', + 40_200, + expect.any(AbortSignal) + ) + }) + + it('retains only the twelve newest candidates from a large directory stream', async () => { + opendirMock.mockResolvedValue(generatedLogDirectory(10_000)) + openMock.mockImplementation(async (filePath: string) => { + if (filePath.endsWith('language_server.log')) { + throw Object.assign(new Error('missing'), { code: 'ENOENT' }) + } + return readableLog('no listener in this log') + }) + + await fetchAntigravityRateLimits({ homePath: '/home/test', appDataPath: '/app-data' }) + + const cliPaths = openMock.mock.calls + .map(([filePath]) => filePath as string) + .filter((filePath) => filePath.includes('cli-')) + expect(cliPaths).toHaveLength(12) + expect(cliPaths[0]).toMatch(/cli-20260815_009999\.log$/) + expect(cliPaths.at(-1)).toMatch(/cli-20260815_009988\.log$/) + }) + + it('applies the total discovery deadline to a stalled directory open', async () => { + vi.useFakeTimers() + const close = vi.fn(async () => undefined) + const directory = { read: vi.fn(async () => null), close } + let finishOpen: ((value: unknown) => void) | undefined + opendirMock.mockReturnValue( + new Promise((resolve) => { + finishOpen = resolve + }) + ) + + const result = fetchAntigravityRateLimits({ + homePath: '/home/test', + appDataPath: '/app-data' + }) + await vi.advanceTimersByTimeAsync(6_000) + + await expect(result).resolves.toMatchObject({ + status: 'error', + error: 'Antigravity usage lookup timed out', + usageMetadata: { failureKind: 'usage-unavailable' } + }) + + finishOpen?.(directory) + await Promise.resolve() + await Promise.resolve() + expect(close).toHaveBeenCalledOnce() + }) + + it('observes cancellation after the final log read before returning the excerpt', async () => { + const controller = new AbortController() + const contents = 'Language server listening on random port at 40200 for HTTP' + openMock.mockResolvedValue(readableLog(contents)) + const nativeThrowIfAborted = controller.signal.throwIfAborted.bind(controller.signal) + let abortChecks = 0 + vi.spyOn(controller.signal, 'throwIfAborted').mockImplementation(() => { + abortChecks += 1 + if (abortChecks === 5) { + controller.abort() + } + nativeThrowIfAborted() + }) + + await expect( + readAntigravityLogExcerpt('/tmp/antigravity.log', controller.signal) + ).rejects.toThrow(/aborted/i) + expect(abortChecks).toBe(5) + }) +}) diff --git a/src/main/rate-limits/antigravity-usage-mirror.test.ts b/src/main/rate-limits/antigravity-usage-mirror.test.ts deleted file mode 100644 index 84c92e1ad22f..000000000000 --- a/src/main/rate-limits/antigravity-usage-mirror.test.ts +++ /dev/null @@ -1,70 +0,0 @@ -import { describe, expect, it } from 'vitest' -import type { ProviderRateLimits, ProviderRateLimitStatus } from '../../shared/rate-limit-types' -import { deriveAntigravityRateLimits } from './antigravity-usage-mirror' - -function geminiSnapshot( - status: ProviderRateLimitStatus, - error: string | null, - usedPercent: number | null = null -): ProviderRateLimits { - return { - provider: 'gemini', - session: - usedPercent === null - ? null - : { usedPercent, windowMinutes: 300, resetsAt: null, resetDescription: null }, - weekly: null, - updatedAt: 1_700_000_000_000, - error, - status - } -} - -describe('deriveAntigravityRateLimits', () => { - it('mirrors a successful Gemini read as shared Code Assist quota', () => { - const antigravity = deriveAntigravityRateLimits(geminiSnapshot('ok', null, 42)) - - expect(antigravity.provider).toBe('antigravity') - expect(antigravity.status).toBe('ok') - expect(antigravity.session?.usedPercent).toBe(42) - expect(antigravity.error).toBeNull() - }) - - it('reports unavailable without quoting the Gemini failure', () => { - const antigravity = deriveAntigravityRateLimits( - geminiSnapshot('error', 'Gemini project ID not found') - ) - - expect(antigravity.provider).toBe('antigravity') - expect(antigravity.status).toBe('unavailable') - expect(antigravity.error).not.toContain('Gemini project ID not found') - expect(antigravity.error).toContain('Antigravity usage is not available') - expect(antigravity.session).toBeNull() - expect(antigravity.weekly).toBeNull() - }) - - it('does not blame a missing sign-in when the quota read itself failed', () => { - const antigravity = deriveAntigravityRateLimits(geminiSnapshot('error', 'Token refresh failed')) - - // Why: the reported symptom is a connected sign-in whose Code Assist read failed. - expect(antigravity.error).toContain('could not be read right now') - expect(antigravity.error).not.toContain('sign-in is connected') - }) - - it('keeps the Gemini timestamp so activation freshness checks are not forced to refetch', () => { - const antigravity = deriveAntigravityRateLimits(geminiSnapshot('error', 'Token refresh failed')) - - expect(antigravity.updatedAt).toBe(1_700_000_000_000) - }) - - it('points at the missing sign-in when the Gemini opt-in is off', () => { - const antigravity = deriveAntigravityRateLimits( - geminiSnapshot('unavailable', 'Gemini CLI OAuth is disabled in settings') - ) - - expect(antigravity.status).toBe('unavailable') - expect(antigravity.error).not.toContain('Gemini CLI OAuth is disabled in settings') - expect(antigravity.error).toContain('Antigravity usage is not available') - expect(antigravity.error).toContain('Gemini CLI sign-in is connected') - }) -}) diff --git a/src/main/rate-limits/antigravity-usage-mirror.ts b/src/main/rate-limits/antigravity-usage-mirror.ts deleted file mode 100644 index 3cea6decc111..000000000000 --- a/src/main/rate-limits/antigravity-usage-mirror.ts +++ /dev/null @@ -1,29 +0,0 @@ -import type { ProviderRateLimits } 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. -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.' - -export function deriveAntigravityRateLimits(gemini: ProviderRateLimits): ProviderRateLimits { - if (gemini.status === 'ok') { - return { ...gemini, provider: 'antigravity' } - } - return { - provider: 'antigravity', - session: null, - weekly: null, - // Why: reuse the Gemini timestamp so activation freshness checks don't force a refetch every cycle. - updatedAt: gemini.updatedAt, - error: - gemini.status === 'unavailable' - ? ANTIGRAVITY_NO_SIGN_IN_REASON - : ANTIGRAVITY_QUOTA_UNREADABLE_REASON, - status: 'unavailable' - } -} 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..f7a7ceed4239 100644 --- a/src/main/rate-limits/rate-limit-service-test-harness.ts +++ b/src/main/rate-limits/rate-limit-service-test-harness.ts @@ -4,6 +4,7 @@ import type { ProviderRateLimits } from '../../shared/rate-limit-types' import type { RateLimitService } from './service' import { fetchCodexRateLimits } from './codex-fetcher' import { fetchGeminiRateLimits } from './gemini-usage-fetcher' +import { fetchAntigravityRateLimits } from './antigravity-usage-fetcher' import { fetchKimiRateLimits } from './kimi-fetcher' import { fetchMiniMaxRateLimits } from './minimax-fetcher' import { fetchGrokRateLimits } from './grok-fetcher' @@ -85,6 +86,7 @@ export function unavailableProvider( export function mockFreshBackgroundProviderFetches(): void { vi.mocked(fetchCodexRateLimits).mockImplementation(async () => okProvider('codex', 24)) vi.mocked(fetchGeminiRateLimits).mockImplementation(async () => okProvider('gemini', 0)) + vi.mocked(fetchAntigravityRateLimits).mockImplementation(async () => okProvider('antigravity', 0)) vi.mocked(fetchOpenCodeGoRateLimits).mockImplementation(async () => okProvider('opencode-go', 0)) vi.mocked(fetchKimiRateLimits).mockImplementation(async () => okProvider('kimi', 0)) vi.mocked(fetchMiniMaxRateLimits).mockImplementation(async () => okProvider('minimax', 0)) @@ -95,6 +97,7 @@ export function mockFreshBackgroundProviderFetches(): void { export function resetRateLimitProviderMocks(): void { vi.clearAllMocks() vi.mocked(fetchGeminiRateLimits).mockResolvedValue(okProvider('gemini', 0, Date.now())) + vi.mocked(fetchAntigravityRateLimits).mockResolvedValue(okProvider('antigravity', 0, Date.now())) vi.mocked(fetchOpenCodeGoRateLimits).mockResolvedValue(okProvider('opencode-go', 0, Date.now())) vi.mocked(fetchKimiRateLimits).mockResolvedValue(okProvider('kimi', 0, Date.now())) vi.mocked(fetchMiniMaxRateLimits).mockResolvedValue(okProvider('minimax', 0, Date.now())) 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..790ccc01e883 100644 --- a/src/main/rate-limits/service-account-target-selection.test.ts +++ b/src/main/rate-limits/service-account-target-selection.test.ts @@ -24,6 +24,10 @@ vi.mock('./gemini-usage-fetcher', () => ({ fetchGeminiRateLimits: vi.fn() })) +vi.mock('./antigravity-usage-fetcher', () => ({ + fetchAntigravityRateLimits: vi.fn() +})) + vi.mock('./kimi-fetcher', () => ({ fetchKimiRateLimits: 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..d2b22ced7552 100644 --- a/src/main/rate-limits/service-antigravity-usage.test.ts +++ b/src/main/rate-limits/service-antigravity-usage.test.ts @@ -3,10 +3,12 @@ import { RateLimitService } from './service' import { fetchClaudeRateLimits } from './claude-fetcher' import { fetchCodexRateLimits } from './codex-fetcher' import { fetchGeminiRateLimits } from './gemini-usage-fetcher' +import { fetchAntigravityRateLimits } from './antigravity-usage-fetcher' import { errorProvider, okProvider, - resetRateLimitProviderMocks + resetRateLimitProviderMocks, + unavailableProvider } from './rate-limit-service-test-harness' vi.mock('./claude-fetcher', () => ({ @@ -23,6 +25,10 @@ vi.mock('./gemini-usage-fetcher', () => ({ fetchGeminiRateLimits: vi.fn() })) +vi.mock('./antigravity-usage-fetcher', () => ({ + fetchAntigravityRateLimits: vi.fn() +})) + vi.mock('./kimi-fetcher', () => ({ fetchKimiRateLimits: vi.fn() })) @@ -58,42 +64,48 @@ describe('RateLimitService Antigravity usage', () => { vi.mocked(fetchGeminiRateLimits).mockResolvedValue( errorProvider('gemini', 'Gemini project ID not found') ) + vi.mocked(fetchAntigravityRateLimits).mockResolvedValue(okProvider('antigravity', 42)) const service = new RateLimitService() await service.refresh() const state = service.getState() - expect(state.antigravity?.status).toBe('unavailable') - expect(state.antigravity?.error).not.toContain('Gemini project ID not found') - expect(state.antigravity?.session).toBeNull() + expect(state.antigravity?.status).toBe('ok') + expect(state.antigravity?.session?.usedPercent).toBe(42) // Why: the real Gemini failure must still surface under its own provider. expect(state.gemini?.status).toBe('error') expect(state.gemini?.error).toBe('Gemini project ID not found') }) - it('keeps mirroring a successful Gemini read under the Antigravity provider', async () => { + it('does not mirror a successful Gemini read when the local service is unavailable', async () => { vi.mocked(fetchGeminiRateLimits).mockResolvedValue(okProvider('gemini', 42, Date.now())) + vi.mocked(fetchAntigravityRateLimits).mockResolvedValue( + unavailableProvider('antigravity', 'Antigravity local usage service is not running') + ) const service = new RateLimitService() await service.refresh() const state = service.getState() - expect(state.antigravity?.status).toBe('ok') + expect(state.antigravity?.status).toBe('unavailable') expect(state.antigravity?.provider).toBe('antigravity') - expect(state.antigravity?.session?.usedPercent).toBe(42) + expect(state.antigravity?.session).toBeNull() + expect(state.gemini?.session?.usedPercent).toBe(42) }) it('never leaves a cached Antigravity snapshot in the error retry lane', async () => { - vi.mocked(fetchGeminiRateLimits).mockResolvedValueOnce(okProvider('gemini', 42, Date.now())) + vi.mocked(fetchAntigravityRateLimits).mockResolvedValueOnce( + okProvider('antigravity', 42, Date.now()) + ) const service = new RateLimitService() await service.refresh() - vi.mocked(fetchGeminiRateLimits).mockResolvedValue( - errorProvider('gemini', 'Token refresh failed') + vi.mocked(fetchAntigravityRateLimits).mockResolvedValue( + unavailableProvider('antigravity', 'Antigravity local usage service is not running') ) await service.refresh() - // Why: stale-retention would otherwise show Gemini numbers as "Refresh failed" Antigravity usage. + // Why: a stopped local runtime must clear its old quota instead of leaving a stale snapshot. expect(service.getState().antigravity?.status).toBe('unavailable') expect(service.getState().antigravity?.session).toBeNull() }) 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..15b64b1e843a 100644 --- a/src/main/rate-limits/service-inactive-account-previews.test.ts +++ b/src/main/rate-limits/service-inactive-account-previews.test.ts @@ -31,6 +31,10 @@ vi.mock('./gemini-usage-fetcher', () => ({ fetchGeminiRateLimits: vi.fn() })) +vi.mock('./antigravity-usage-fetcher', () => ({ + fetchAntigravityRateLimits: vi.fn() +})) + vi.mock('./kimi-fetcher', () => ({ fetchKimiRateLimits: 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..949bc740be42 100644 --- a/src/main/rate-limits/service-live-claude-usage.test.ts +++ b/src/main/rate-limits/service-live-claude-usage.test.ts @@ -28,6 +28,10 @@ vi.mock('./gemini-usage-fetcher', () => ({ fetchGeminiRateLimits: vi.fn() })) +vi.mock('./antigravity-usage-fetcher', () => ({ + fetchAntigravityRateLimits: vi.fn() +})) + vi.mock('./kimi-fetcher', () => ({ fetchKimiRateLimits: 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 737ca0ff03d1..0376674a0f1d 100644 --- a/src/main/rate-limits/service-minimax-usage.test.ts +++ b/src/main/rate-limits/service-minimax-usage.test.ts @@ -25,6 +25,10 @@ vi.mock('./gemini-usage-fetcher', () => ({ fetchGeminiRateLimits: vi.fn() })) +vi.mock('./antigravity-usage-fetcher', () => ({ + fetchAntigravityRateLimits: vi.fn() +})) + vi.mock('./kimi-fetcher', () => ({ fetchKimiRateLimits: 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..6beff3460d12 100644 --- a/src/main/rate-limits/service-refresh-orchestration.test.ts +++ b/src/main/rate-limits/service-refresh-orchestration.test.ts @@ -4,6 +4,7 @@ import { RateLimitService } from './service' import { fetchClaudeRateLimits } from './claude-fetcher' import { fetchCodexRateLimits } from './codex-fetcher' import { fetchGeminiRateLimits } from './gemini-usage-fetcher' +import { fetchAntigravityRateLimits } from './antigravity-usage-fetcher' import { fetchKimiRateLimits } from './kimi-fetcher' import { fetchMiniMaxRateLimits } from './minimax-fetcher' import { fetchGrokRateLimits } from './grok-fetcher' @@ -15,7 +16,8 @@ import { flushMicrotasks, mockFreshBackgroundProviderFetches, okProvider, - resetRateLimitProviderMocks + resetRateLimitProviderMocks, + unavailableProvider } from './rate-limit-service-test-harness' vi.mock('./claude-fetcher', () => ({ @@ -32,6 +34,10 @@ vi.mock('./gemini-usage-fetcher', () => ({ fetchGeminiRateLimits: vi.fn() })) +vi.mock('./antigravity-usage-fetcher', () => ({ + fetchAntigravityRateLimits: vi.fn() +})) + vi.mock('./kimi-fetcher', () => ({ fetchKimiRateLimits: vi.fn() })) @@ -361,7 +367,7 @@ describe('RateLimitService', () => { expect(fetchGrokRateLimits).toHaveBeenCalledTimes(1) }) - it('fetches Gemini and OpenCode Go alongside Claude and Codex', async () => { + it('fetches Gemini, Antigravity, and OpenCode Go alongside Claude and Codex', async () => { const service = new RateLimitService() service.setOpenCodeGoConfigResolver(() => ({ sessionCookie: 'session=abc123', @@ -377,6 +383,9 @@ describe('RateLimitService', () => { vi.mocked(fetchClaudeRateLimits).mockResolvedValueOnce(okProvider('claude', 10, Date.now())) vi.mocked(fetchCodexRateLimits).mockResolvedValueOnce(okProvider('codex', 20, Date.now())) vi.mocked(fetchGeminiRateLimits).mockResolvedValueOnce(okProvider('gemini', 30, Date.now())) + vi.mocked(fetchAntigravityRateLimits).mockResolvedValueOnce( + okProvider('antigravity', 35, Date.now()) + ) vi.mocked(fetchOpenCodeGoRateLimits).mockResolvedValueOnce( okProvider('opencode-go', 40, Date.now()) ) @@ -395,6 +404,10 @@ describe('RateLimitService', () => { expect(fetchCodexRateLimits).toHaveBeenCalledTimes(1) expect(fetchGeminiRateLimits).toHaveBeenCalledTimes(1) expect(fetchGeminiRateLimits).toHaveBeenCalledWith(true) + expect(fetchAntigravityRateLimits).toHaveBeenCalledTimes(1) + expect(fetchAntigravityRateLimits).toHaveBeenCalledWith({ + signal: expect.any(AbortSignal) + }) expect(fetchOpenCodeGoRateLimits).toHaveBeenCalledTimes(1) expect(fetchOpenCodeGoRateLimits).toHaveBeenCalledWith( 'session=abc123', @@ -413,10 +426,29 @@ describe('RateLimitService', () => { expect(state.codex?.session?.usedPercent).toBe(20) expect(state.gemini?.status).toBe('ok') expect(state.gemini?.session?.usedPercent).toBe(30) + expect(state.antigravity?.status).toBe('ok') + expect(state.antigravity?.session?.usedPercent).toBe(35) expect(state.opencodeGo?.status).toBe('ok') expect(state.opencodeGo?.session?.usedPercent).toBe(40) }) + it('keeps Antigravity usage when Gemini usage is unavailable', async () => { + const service = new RateLimitService() + vi.mocked(fetchClaudeRateLimits).mockResolvedValueOnce(okProvider('claude', 10, Date.now())) + vi.mocked(fetchCodexRateLimits).mockResolvedValueOnce(okProvider('codex', 20, Date.now())) + vi.mocked(fetchGeminiRateLimits).mockResolvedValueOnce(unavailableProvider('gemini')) + vi.mocked(fetchAntigravityRateLimits).mockResolvedValueOnce( + okProvider('antigravity', 35, Date.now()) + ) + + await service.refresh() + + const state = service.getState() + expect(state.gemini?.status).toBe('unavailable') + expect(state.antigravity?.status).toBe('ok') + expect(state.antigravity?.session?.usedPercent).toBe(35) + }) + it('passes the resolved Kimi home into each fetch cycle', async () => { const service = new RateLimitService() const home = { diff --git a/src/main/rate-limits/service-window-activation.test.ts b/src/main/rate-limits/service-window-activation.test.ts index a4a1dc14f767..1c87a7707374 100644 --- a/src/main/rate-limits/service-window-activation.test.ts +++ b/src/main/rate-limits/service-window-activation.test.ts @@ -33,6 +33,10 @@ vi.mock('./gemini-usage-fetcher', () => ({ fetchGeminiRateLimits: vi.fn() })) +vi.mock('./antigravity-usage-fetcher', () => ({ + fetchAntigravityRateLimits: vi.fn() +})) + vi.mock('./kimi-fetcher', () => ({ fetchKimiRateLimits: vi.fn() })) diff --git a/src/main/rate-limits/service.ts b/src/main/rate-limits/service.ts index d6f291d6b38e..366d09ccc0b1 100644 --- a/src/main/rate-limits/service.ts +++ b/src/main/rate-limits/service.ts @@ -20,7 +20,7 @@ import { type NormalizedClaudeAccountSelectionTarget } from '../claude-accounts/runtime-selection' import { fetchGeminiRateLimits } from './gemini-usage-fetcher' -import { deriveAntigravityRateLimits } from './antigravity-usage-mirror' +import { fetchAntigravityRateLimits } from './antigravity-usage-fetcher' import { fetchKimiRateLimits } from './kimi-fetcher' import type { KimiHomeResolution } from '../kimi/kimi-runtime-home' import { fetchGrokRateLimits } from './grok-fetcher' @@ -1689,40 +1689,50 @@ export class RateLimitService { 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 - }) - ]) + const [ + claudeResult, + codexResult, + geminiResult, + antigravityResult, + 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), + // Why: Antigravity availability belongs to the desktop host running its + // loopback service; an SSH worktree must not redirect this discovery. + fetchAntigravityRateLimits({ signal }), + fetchOpenCodeGoRateLimits( + cookie, + workspaceIdOverride || undefined, + this.networkProxySettingsResolver?.() + ), + this.fetchKimiWithResolvedHome(), + miniMaxConfigResult.error + ? Promise.resolve(this.getMiniMaxCredentialError(miniMaxConfigResult.error)) + : fetchMiniMaxRateLimits({ + cookie: miniMaxCookie, + groupId: miniMaxGroupId, + models: miniMaxModels + }) + ]) if (signal.aborted) { return @@ -1767,8 +1777,20 @@ export class RateLimitService { 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) + const antigravity = + antigravityResult.status === 'fulfilled' + ? antigravityResult.value + : ({ + provider: 'antigravity', + session: null, + weekly: null, + updatedAt: Date.now(), + error: + antigravityResult.reason instanceof Error + ? antigravityResult.reason.message + : 'Unknown error', + status: 'error' + } satisfies ProviderRateLimits) const opencodeGo = opencodeGoResult.status === 'fulfilled' diff --git a/src/renderer/src/components/status-bar/StatusBar.tsx b/src/renderer/src/components/status-bar/StatusBar.tsx index 436da40e59c3..da7611637c0f 100644 --- a/src/renderer/src/components/status-bar/StatusBar.tsx +++ b/src/renderer/src/components/status-bar/StatusBar.tsx @@ -1179,7 +1179,9 @@ function VerboseProviderUsage({ p: ProviderRateLimits display: UsagePercentageDisplay }): React.JSX.Element { - if (p.buckets && p.buckets.length > 0) { + // Why: Antigravity buckets power the roster's Detailed/Compact modes, while + // this narrow footer stays readable with its derived 5h and weekly summary. + if (p.provider !== 'antigravity' && p.buckets && p.buckets.length > 0) { const visibleBuckets = p.buckets.filter((bucket) => STATUS_BAR_BUCKET_NAMES.has(bucket.name)) return ( <> @@ -2109,11 +2111,10 @@ function StatusBarInner({ floatingTerminalOpen }: StatusBarProps): React.JSX.Ele 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 may be provided by the desktop runtime even when `agy` is not on PATH. const antigravityUsageConfigured = statusBarItems.includes('antigravity') && - isStatusBarItemAvailable('antigravity', detectedAgentIds) + isStatusBarItemAvailable('antigravity', detectedAgentIds, antigravity) // Why: thread non-GlobalSettings durability flags so bars stay visible across reloads and snapshot refreshes. const usageSettings = { ...settings, @@ -2147,7 +2148,7 @@ function StatusBarInner({ floatingTerminalOpen }: StatusBarProps): React.JSX.Ele const showAntigravity = visibleAntigravity !== null && statusBarItems.includes('antigravity') && - isStatusBarItemAvailable('antigravity', detectedAgentIds) + isStatusBarItemAvailable('antigravity', detectedAgentIds, antigravity) // 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 = @@ -2487,7 +2488,7 @@ function StatusBarInner({ floatingTerminalOpen }: StatusBarProps): React.JSX.Ele {translate('auto.components.status.bar.StatusBar.c1df0d67ec', 'Gemini Usage')} )} - {isStatusBarItemAvailable('antigravity', detectedAgentIds) && ( + {isStatusBarItemAvailable('antigravity', detectedAgentIds, antigravity) && ( { diff --git a/src/renderer/src/components/status-bar/UsageRosterPanel.test.tsx b/src/renderer/src/components/status-bar/UsageRosterPanel.test.tsx index de4f5a0e06f8..18db2b391518 100644 --- a/src/renderer/src/components/status-bar/UsageRosterPanel.test.tsx +++ b/src/renderer/src/components/status-bar/UsageRosterPanel.test.tsx @@ -40,6 +40,57 @@ const signedOutCodex: ProviderRateLimits = { status: 'error' } +function antigravityLimits(tightestResetsAt: number | null = null): ProviderRateLimits { + return { + provider: 'antigravity', + session: { + usedPercent: 40, + windowMinutes: 300, + resetsAt: null, + resetDescription: null + }, + weekly: { + usedPercent: 90, + windowMinutes: 10_080, + resetsAt: null, + resetDescription: null + }, + buckets: [ + { + name: 'Gemini 5h', + usedPercent: 20, + windowMinutes: 300, + resetsAt: null, + resetDescription: null + }, + { + name: 'Gemini wk', + usedPercent: 60, + windowMinutes: 10_080, + resetsAt: null, + resetDescription: null + }, + { + name: 'Claude/GPT 5h', + usedPercent: 40, + windowMinutes: 300, + resetsAt: null, + resetDescription: null + }, + { + name: 'Claude/GPT wk', + usedPercent: 90, + windowMinutes: 10_080, + resetsAt: tightestResetsAt, + resetDescription: null + } + ], + updatedAt: 0, + error: null, + status: 'ok' + } +} + describe('UsageRow', () => { beforeEach(() => { mocks.useResetCountdownClock.mockClear() @@ -247,6 +298,49 @@ describe('UsageRow', () => { expect(markup).toContain('25%') expect(markup).toContain('60%') }) + + it('renders all four Antigravity buckets in Detailed mode', () => { + const markup = renderToStaticMarkup( + + ) + + expect(markup.match(/data-usage-window=/g)).toHaveLength(4) + expect(markup.match(/data-usage-bar/g)).toHaveLength(4) + expect(markup).toContain('Gemini 5h') + expect(markup).toContain('Gemini wk') + expect(markup).toContain('Claude/GPT 5h') + expect(markup).toContain('Claude/GPT wk') + }) + + it('renders only the tightest Antigravity bucket in Compact mode', () => { + const dateNow = vi.spyOn(Date, 'now').mockReturnValue(mocks.now) + try { + const markup = renderToStaticMarkup( + + ) + + expect(markup.match(/data-usage-window=/g)).toHaveLength(1) + expect(markup).toContain('43m') + expect(markup).toContain('90%') + expect(markup).not.toContain('>Claude/GPT wk<') + expect(markup).not.toContain('>Gemini wk<') + } finally { + dateNow.mockRestore() + } + }) }) describe('UsageRosterPanel density picker', () => { diff --git a/src/renderer/src/components/status-bar/UsageRosterPanel.tsx b/src/renderer/src/components/status-bar/UsageRosterPanel.tsx index feb066fadca4..5c51f188eb55 100644 --- a/src/renderer/src/components/status-bar/UsageRosterPanel.tsx +++ b/src/renderer/src/components/status-bar/UsageRosterPanel.tsx @@ -41,7 +41,10 @@ function shortLabel( section: UsageSection, useRemainingDuration = false ): string { - if (p.buckets?.some((b) => b.name === section.label)) { + const isNamedBucket = p.buckets?.some((b) => b.name === section.label) + // Why: Antigravity bucket names distinguish Detailed rows, but Compact + // should match other providers by showing the selected window's reset time. + if (isNamedBucket && !(p.provider === 'antigravity' && useRemainingDuration)) { return section.label } // fableWeekly shares the 7d window with weekly; label it distinctly so the two diff --git a/src/renderer/src/components/status-bar/provider-segment-monthly-window.test.tsx b/src/renderer/src/components/status-bar/provider-segment-monthly-window.test.tsx index 3bab0e1bfcb3..bd95b681c6db 100644 --- a/src/renderer/src/components/status-bar/provider-segment-monthly-window.test.tsx +++ b/src/renderer/src/components/status-bar/provider-segment-monthly-window.test.tsx @@ -43,6 +43,23 @@ function grokMonthlyLimits(status: ProviderRateLimits['status']): ProviderRateLi } } +function antigravityLimits(tightestResetsAt: number | null = null): ProviderRateLimits { + return { + provider: 'antigravity', + session: windowOf(40, 300), + weekly: windowOf(80, 10_080), + buckets: [ + { ...windowOf(20, 300), name: 'Gemini 5h' }, + { ...windowOf(80, 10_080, tightestResetsAt), name: 'Gemini wk' }, + { ...windowOf(40, 300), name: 'Claude/GPT 5h' }, + { ...windowOf(10, 10_080), name: 'Claude/GPT wk' } + ], + updatedAt: Date.now(), + error: null, + status: 'ok' + } +} + describe('ProviderSegment monthly window', () => { it('renders a monthly-only snapshot in the chip instead of a bare icon', async () => { const { ProviderSegment } = await import('./StatusBar') @@ -182,6 +199,42 @@ describe('ProviderSegment monthly window', () => { expect(markup).toContain('30% used Fable') expect(markup).not.toContain('40% used') }) + + it('keeps Antigravity verbose output to its derived summary windows', async () => { + const { ProviderSegment } = await import('./StatusBar') + + const markup = renderToStaticMarkup( + + ) + + expect(markup).toContain('40% used 5h') + expect(markup).toContain('80% used wk') + expect(markup).not.toContain('Gemini wk') + expect(markup).not.toContain('Claude/GPT') + }) + + it("selects Antigravity's tightest named bucket in compact mode", async () => { + const { ProviderSegment } = await import('./StatusBar') + const now = 1_700_000_000_000 + const dateNow = vi.spyOn(Date, 'now').mockReturnValue(now) + + try { + const markup = renderToStaticMarkup( + + ) + + expect(markup).toContain('80% used 43m') + expect(markup).not.toContain('Gemini wk') + expect(markup).not.toContain('Claude/GPT') + } finally { + dateNow.mockRestore() + } + }) }) describe('undefined provider window safety (crash d2c1da69 / bb74236c)', () => { 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..842ff2ebd256 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,6 +1,30 @@ import { describe, expect, it } from 'vitest' +import type { ProviderRateLimits } from '../../../../shared/rate-limit-types' import { isStatusBarItemAvailable } from './status-bar-agent-gating' +function antigravitySnapshot( + status: ProviderRateLimits['status'], + failureKind?: 'usage-unavailable' | 'cli-unavailable' +): ProviderRateLimits { + return { + provider: 'antigravity', + session: null, + weekly: null, + updatedAt: 1, + error: status === 'error' ? 'Quota unavailable' : null, + status, + usageMetadata: failureKind + ? { + source: 'live-session', + attemptedSources: ['live-session'], + failureKind, + credentialSource: 'agy-local-service', + authProvenance: 'antigravity' + } + : undefined + } +} + describe('isStatusBarItemAvailable', () => { it('shows non-CLI items regardless of detection', () => { // Why: ssh, resource-usage, and opencode-go aren't CLIs on PATH, so @@ -26,7 +50,6 @@ describe('isStatusBarItemAvailable', () => { expect(isStatusBarItemAvailable('claude', [])).toBe(false) expect(isStatusBarItemAvailable('codex', ['claude'])).toBe(false) expect(isStatusBarItemAvailable('gemini', ['claude', 'codex'])).toBe(false) - expect(isStatusBarItemAvailable('antigravity', ['claude', 'codex'])).toBe(false) expect(isStatusBarItemAvailable('grok', ['claude', 'kimi'])).toBe(false) }) @@ -34,7 +57,40 @@ describe('isStatusBarItemAvailable', () => { expect(isStatusBarItemAvailable('claude', ['claude'])).toBe(true) expect(isStatusBarItemAvailable('codex', ['codex', 'claude'])).toBe(true) expect(isStatusBarItemAvailable('gemini', ['gemini'])).toBe(true) - expect(isStatusBarItemAvailable('antigravity', ['antigravity'])).toBe(true) expect(isStatusBarItemAvailable('grok', ['grok'])).toBe(true) }) + + it('shows Antigravity for a desktop-only runtime snapshot', () => { + expect( + isStatusBarItemAvailable('antigravity', ['claude', 'codex'], antigravitySnapshot('ok')) + ).toBe(true) + expect( + isStatusBarItemAvailable( + 'antigravity', + ['claude', 'codex'], + antigravitySnapshot('error', 'usage-unavailable') + ) + ).toBe(true) + }) + + it('hides Antigravity after runtime discovery confirms it is unavailable', () => { + expect( + isStatusBarItemAvailable( + 'antigravity', + ['claude', 'codex'], + antigravitySnapshot('unavailable', 'cli-unavailable') + ) + ).toBe(false) + expect( + isStatusBarItemAvailable( + 'antigravity', + ['claude', 'codex'], + antigravitySnapshot('fetching', 'cli-unavailable') + ) + ).toBe(false) + }) + + it('shows Antigravity when its CLI is detected before a runtime snapshot arrives', () => { + expect(isStatusBarItemAvailable('antigravity', ['antigravity'])).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..ba85a2dbccbe 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,4 +1,5 @@ import type { TuiAgent } from '../../../../shared/tui-agent' +import type { ProviderRateLimits } from '../../../../shared/rate-limit-types' import type { StatusBarItem } from '../../../../shared/ui-chrome-types' // Why: CLI-backed usage bars are surface noise when the underlying @@ -12,14 +13,25 @@ const CLI_GATED_ITEMS: ReadonlySet = new Set([ 'codex', 'gemini', 'kimi', - 'antigravity', 'grok' ]) export function isStatusBarItemAvailable( id: StatusBarItem, - detectedAgentIds: TuiAgent[] | null + detectedAgentIds: TuiAgent[] | null, + usageSnapshot?: ProviderRateLimits | null ): boolean { + if (id === 'antigravity') { + if (detectedAgentIds === null || detectedAgentIds.includes('antigravity')) { + return true + } + // Why: Antigravity Desktop exposes the same host-local quota service without installing `agy` on PATH. + return ( + usageSnapshot != null && + usageSnapshot.status !== 'unavailable' && + usageSnapshot.usageMetadata?.failureKind !== 'cli-unavailable' + ) + } if (!CLI_GATED_ITEMS.has(id)) { return true } 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 e833bce7fa4c..3858dac7410c 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 @@ -121,11 +121,7 @@ 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 - ) + expect(hasUsageProviderSettings(usageSettings({ antigravityUsageConfigured: true }))).toBe(true) expect(hasUsageProviderSettings(usageSettings({ minimaxCookieConfigured: true }))).toBe(true) expect(hasUsageProviderSettings(usageSettings({ grokAuthConfigured: true }))).toBe(true) }) @@ -160,21 +156,13 @@ describe('hasUsageProviderSettingsForProvider', () => { expect(hasUsageProviderSettingsForProvider('grok', usageSettings())).toBe(false) }) - it('requires both a checked Antigravity item and Gemini OAuth as the durable Antigravity signal', () => { - 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. + it('uses the checked Antigravity item as its durable configuration signal', () => { expect( hasUsageProviderSettingsForProvider( 'antigravity', usageSettings({ antigravityUsageConfigured: true }) ) - ).toBe(false) + ).toBe(true) expect( hasUsageProviderSettingsForProvider( 'antigravity', @@ -324,11 +312,11 @@ describe('getVisibleUsageProvider', () => { ).toBe(null) }) - it('keeps Antigravity visible while the snapshot is pending when checked and Gemini OAuth is on', () => { + it('keeps Antigravity visible while the snapshot is pending when checked', () => { const visible = getVisibleUsageProvider( 'antigravity', null, - usageSettings({ antigravityUsageConfigured: true, geminiCliOAuthEnabled: true }) + usageSettings({ antigravityUsageConfigured: true }) ) expect(visible).toMatchObject({ provider: 'antigravity', @@ -338,26 +326,24 @@ 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('keeps a checked Antigravity item independent of Gemini OAuth', () => { expect( getVisibleUsageProvider( 'antigravity', null, usageSettings({ antigravityUsageConfigured: true }) ) - ).toBe(null) + ).toMatchObject({ provider: 'antigravity', status: 'fetching' }) expect( getVisibleUsageProvider( 'antigravity', provider('unavailable', { provider: 'antigravity', - error: 'Gemini CLI OAuth is disabled in settings' + error: 'Antigravity local usage service is not running' }), usageSettings({ antigravityUsageConfigured: true }) ) - ).toBe(null) + ).toMatchObject({ provider: 'antigravity', status: 'unavailable' }) }) }) @@ -494,14 +480,12 @@ describe('isUsageEmptyState', () => { grok: provider('unavailable', { provider: 'grok' }), minimax: provider('unavailable', { provider: 'minimax' }) }, - usageSettings({ antigravityUsageConfigured: true, geminiCliOAuthEnabled: true }) + usageSettings({ antigravityUsageConfigured: true }) ) ).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('does not show the setup CTA for checked Antigravity when Gemini OAuth is off', () => { expect( isUsageEmptyState( { @@ -516,6 +500,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 f1afd97f5f4b..69e52885961e 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,8 @@ 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 quota comes from its local runtime, so the checked item is + // the durable signal while the runtime may be between sessions. antigravityUsageConfigured: boolean // Why: MiniMax/Grok sign-in live on disk, not in settings; main sets these each poll. minimaxCookieConfigured: boolean @@ -74,8 +71,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?.grokAuthConfigured === true ) @@ -101,10 +97,7 @@ 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 + return settings.antigravityUsageConfigured === true } if (providerId === 'minimax') { return settings.minimaxCookieConfigured === true diff --git a/src/renderer/src/components/status-bar/tooltip.test.ts b/src/renderer/src/components/status-bar/tooltip.test.ts index 2dc23cf100a5..cedd9076bbe1 100644 --- a/src/renderer/src/components/status-bar/tooltip.test.ts +++ b/src/renderer/src/components/status-bar/tooltip.test.ts @@ -326,6 +326,54 @@ describe('getWindowSections', () => { ]) }) + it('returns only Antigravity buckets because its summaries duplicate them', () => { + const p: ProviderRateLimits = { + provider: 'antigravity', + session: { usedPercent: 40, windowMinutes: 300, resetsAt: null, resetDescription: null }, + weekly: { usedPercent: 80, windowMinutes: 10_080, resetsAt: null, resetDescription: null }, + buckets: [ + { + name: 'Gemini 5h', + usedPercent: 20, + windowMinutes: 300, + resetsAt: null, + resetDescription: null + }, + { + name: 'Gemini wk', + usedPercent: 80, + windowMinutes: 10_080, + resetsAt: null, + resetDescription: null + }, + { + name: 'Claude/GPT 5h', + usedPercent: 40, + windowMinutes: 300, + resetsAt: null, + resetDescription: null + }, + { + name: 'Claude/GPT wk', + usedPercent: 10, + windowMinutes: 10_080, + resetsAt: null, + resetDescription: null + } + ], + updatedAt: Date.now(), + error: null, + status: 'ok' + } + + expect(getWindowSections(p)).toEqual( + p.buckets?.map((bucket) => ({ + label: bucket.name, + window: bucket + })) + ) + }) + it('returns session and weekly when buckets are absent', () => { const p: ProviderRateLimits = { provider: 'claude', diff --git a/src/renderer/src/components/status-bar/tooltip.tsx b/src/renderer/src/components/status-bar/tooltip.tsx index db585ab64f7f..805cc328f09f 100644 --- a/src/renderer/src/components/status-bar/tooltip.tsx +++ b/src/renderer/src/components/status-bar/tooltip.tsx @@ -143,6 +143,11 @@ export function getWindowSections( ): { label: string; window: RateLimitWindow | null }[] { if (p.buckets?.length) { const bucketSections = p.buckets.map((b) => ({ label: b.name, window: b as RateLimitWindow })) + // Why: Antigravity's summaries are derived from these same four buckets; + // appending Weekly would duplicate one family in Detailed mode. + if (p.provider === 'antigravity') { + return bucketSections + } return [ ...bucketSections, { diff --git a/src/renderer/src/components/status-bar/use-available-status-bar-toggles.test.ts b/src/renderer/src/components/status-bar/use-available-status-bar-toggles.test.ts new file mode 100644 index 000000000000..e024401bfa5d --- /dev/null +++ b/src/renderer/src/components/status-bar/use-available-status-bar-toggles.test.ts @@ -0,0 +1,52 @@ +import { describe, expect, it, vi } from 'vitest' +import type { ProviderRateLimits } from '../../../../shared/rate-limit-types' +import type { TuiAgent } from '../../../../shared/tui-agent' +import type { StatusBarItem } from '../../../../shared/ui-chrome-types' + +const store = vi.hoisted(() => ({ + state: { + detectedAgentIds: [] as TuiAgent[], + rateLimits: { antigravity: null as ProviderRateLimits | null } + } +})) + +vi.mock('../../store', () => ({ + useAppStore: (selector: (state: typeof store.state) => T): T => selector(store.state) +})) + +import { useAvailableStatusBarToggles } from './use-available-status-bar-toggles' + +const toggles: { id: StatusBarItem }[] = [{ id: 'antigravity' }, { id: 'claude' }, { id: 'ssh' }] + +function antigravitySnapshot(status: ProviderRateLimits['status']): ProviderRateLimits { + return { + provider: 'antigravity', + session: null, + weekly: null, + updatedAt: 1, + error: null, + status + } +} + +describe('useAvailableStatusBarToggles', () => { + it('keeps the Antigravity toggle for a desktop-only runtime snapshot', () => { + store.state.detectedAgentIds = [] + store.state.rateLimits.antigravity = antigravitySnapshot('ok') + + expect(useAvailableStatusBarToggles(toggles).map((toggle) => toggle.id)).toEqual([ + 'antigravity', + 'ssh' + ]) + }) + + it('hides the Antigravity toggle when neither runtime is available', () => { + store.state.detectedAgentIds = [] + store.state.rateLimits.antigravity = { + ...antigravitySnapshot('unavailable'), + usageMetadata: { failureKind: 'cli-unavailable' } + } + + expect(useAvailableStatusBarToggles(toggles).map((toggle) => toggle.id)).toEqual(['ssh']) + }) +}) diff --git a/src/renderer/src/components/status-bar/use-available-status-bar-toggles.ts b/src/renderer/src/components/status-bar/use-available-status-bar-toggles.ts index cda54a7b8bc7..3ed83186b25a 100644 --- a/src/renderer/src/components/status-bar/use-available-status-bar-toggles.ts +++ b/src/renderer/src/components/status-bar/use-available-status-bar-toggles.ts @@ -8,5 +8,12 @@ export function useAvailableStatusBarToggles( toggles: readonly T[] ): T[] { const detectedAgentIds = useAppStore((s) => s.detectedAgentIds) - return toggles.filter((t) => isStatusBarItemAvailable(t.id, detectedAgentIds)) + const antigravity = useAppStore((s) => s.rateLimits.antigravity) + return toggles.filter((toggle) => + isStatusBarItemAvailable( + toggle.id, + detectedAgentIds, + toggle.id === 'antigravity' ? antigravity : null + ) + ) } diff --git a/src/shared/rate-limit-types.ts b/src/shared/rate-limit-types.ts index 83210fba2ccb..cbbc2bb86693 100644 --- a/src/shared/rate-limit-types.ts +++ b/src/shared/rate-limit-types.ts @@ -63,7 +63,7 @@ export type ProviderRateLimits = { fableWeekly?: RateLimitWindow | null /** 30-day monthly window (OpenCode Go, Grok unified billing), null if not available. */ monthly?: RateLimitWindow | null - /** Named per-model buckets (Gemini only). */ + /** Named provider quota buckets (Gemini and Antigravity). */ buckets?: RateLimitBucket[] /** Available earned Codex rate-limit reset credits, if reported. */ rateLimitResetCredits?: {