diff --git a/src/main/ipc/dashboard-payload-validation.ts b/src/main/ipc/dashboard-payload-validation.ts index e8f8560aa518..59248ba1a35c 100644 --- a/src/main/ipc/dashboard-payload-validation.ts +++ b/src/main/ipc/dashboard-payload-validation.ts @@ -2,7 +2,8 @@ import { DASHBOARD_MAX_LABEL_LENGTH, type DashboardRevealAgentArgs, type DashboardSleepWorkspaceArgs, - type DashboardSnapshot + type DashboardSnapshot, + type DashboardStopAgentArgs } from '../../shared/dashboard-snapshot' import { BoundedMap } from '../../shared/bounded-map' import { normalizeExecutionHostId } from '../../shared/execution-host' @@ -83,6 +84,20 @@ export function isDashboardRevealAgentArgs(value: unknown): value is DashboardRe ) } +export function isDashboardStopAgentArgs(value: unknown): value is DashboardStopAgentArgs { + if (!value || typeof value !== 'object' || Array.isArray(value)) { + return false + } + const args = value as Record + return ( + isBoundedString(args.paneKey, MAX_ID_LENGTH) && + isBoundedString(args.worktreeId, MAX_ID_LENGTH) && + isBoundedString(args.tabId, MAX_ID_LENGTH) && + (args.leafId === null || isBoundedString(args.leafId, MAX_ID_LENGTH)) && + (args.ptyId === null || isBoundedString(args.ptyId, MAX_ID_LENGTH)) + ) +} + export function isDashboardPaneKey(value: unknown): value is string { return isBoundedString(value, MAX_ID_LENGTH) } diff --git a/src/main/ipc/dashboard-popout.test.ts b/src/main/ipc/dashboard-popout.test.ts index ce10b3556fce..28eb2e259722 100644 --- a/src/main/ipc/dashboard-popout.test.ts +++ b/src/main/ipc/dashboard-popout.test.ts @@ -274,6 +274,41 @@ describe('registerDashboardPopoutHandlers', () => { expect(sendToTrustedMock).toHaveBeenCalledWith('ui:sleepDashboardWorkspace', args) }) + it('relays valid stop requests from only the popout, without raising the main window', () => { + const args = { + paneKey: 'tab1:leaf1', + worktreeId: 'w1', + tabId: 't1', + leafId: 'l1', + ptyId: 'pty1' + } + + handlers.get('dashboardPopout:stopAgent')!({ sender: untrustedSender } as never, args) + handlers.get('dashboardPopout:stopAgent')!({ sender: popoutSender } as never, { + ...args, + paneKey: '' + }) + expect(sendToTrustedMock).not.toHaveBeenCalled() + + handlers.get('dashboardPopout:stopAgent')!({ sender: popoutSender } as never, args) + expect(sendToTrustedMock).toHaveBeenCalledWith('ui:stopDashboardAgent', args) + // Stopping must not pull the user off the board. + expect(safelyRevealMock).not.toHaveBeenCalled() + expect(appMock.focus).not.toHaveBeenCalled() + }) + + it('accepts a stop request for a retained card with no live pane', () => { + const args = { + paneKey: 'tab1:leaf1', + worktreeId: 'w1', + tabId: 't1', + leafId: null, + ptyId: null + } + + handlers.get('dashboardPopout:stopAgent')!({ sender: popoutSender } as never, args) + expect(sendToTrustedMock).toHaveBeenCalledWith('ui:stopDashboardAgent', args) + }) it('reveals an agent in only the trusted main window', () => { const main = makeWindow(mainSender) getTrustedWindowMock.mockReturnValue(main) diff --git a/src/main/ipc/dashboard-popout.ts b/src/main/ipc/dashboard-popout.ts index 5b4781ec761a..20134d516220 100644 --- a/src/main/ipc/dashboard-popout.ts +++ b/src/main/ipc/dashboard-popout.ts @@ -16,7 +16,8 @@ import { isDashboardPaneKey, isDashboardRevealAgentArgs, isDashboardSleepWorkspaceArgs, - isDashboardSpawnAgentArgs + isDashboardSpawnAgentArgs, + isDashboardStopAgentArgs } from './dashboard-payload-validation' // The most recent snapshot the main renderer published, replayed to the popout @@ -40,6 +41,7 @@ export function registerDashboardPopoutHandlers( ipcMain.removeHandler('dashboardPopout:ackAgent') ipcMain.removeHandler('dashboardPopout:spawnAgent') ipcMain.removeHandler('dashboardPopout:sleepWorkspace') + ipcMain.removeHandler('dashboardPopout:stopAgent') onDashboardPopoutOpenChanged((open) => { if (!open) { @@ -127,6 +129,20 @@ export function registerDashboardPopoutHandlers( sendToTrustedUIRenderer('ui:ackDashboardAgent', (args as { paneKey: string }).paneKey) }) + // Stop-from-board: the main renderer owns the terminal/pty teardown, so relay + // the request to it. Deliberately does NOT raise the main window — the whole + // point is clearing an agent without leaving the board. + ipcMain.handle('dashboardPopout:stopAgent', (event, args: unknown): void => { + if ( + !isDashboardPopoutRenderer(event.sender) || + !isDashboardEnabled(store) || + !isDashboardStopAgentArgs(args) + ) { + return + } + sendToTrustedUIRenderer('ui:stopDashboardAgent', args) + }) + // Click-to-focus: raise the main window and route it to the agent's pane. ipcMain.handle('dashboardPopout:revealAgent', (event, args: unknown): void => { if ( diff --git a/src/main/rate-limits/antigravity-auth-types.ts b/src/main/rate-limits/antigravity-auth-types.ts new file mode 100644 index 000000000000..e76c6857037f --- /dev/null +++ b/src/main/rate-limits/antigravity-auth-types.ts @@ -0,0 +1,53 @@ +import type { UsageRateLimitFailureKind } from '../../shared/rate-limit-types' + +export type AntigravityCredentialSource = 'official-keychain' | 'official-token-file' + +export type AntigravityAccessToken = { + accessToken: string + credentialSource: AntigravityCredentialSource + sourceKey: string +} + +export class AntigravityAuthError extends Error { + readonly failureKind: UsageRateLimitFailureKind + readonly status: number | null + + constructor( + message: string, + failureKind: UsageRateLimitFailureKind, + status: number | null = null + ) { + super(message) + this.name = 'AntigravityAuthError' + this.failureKind = failureKind + this.status = status + } +} + +export type TokenEnvelope = { + token: { + access_token: string + refresh_token?: string + token_type?: string + expiry?: string | number + } + [key: string]: unknown +} + +export type ParsedCredentials = { + source: AntigravityCredentialSource + sourceKey: string + tokenPath: string | null + envelope: TokenEnvelope + accessToken: string + refreshToken: string + expiresAtMs: number | null + baseHomeDir: string +} + +export type RefreshedCredentials = { + accessToken: string + refreshToken: string + expiresAtMs: number + idToken?: string +} diff --git a/src/main/rate-limits/antigravity-auth.test.ts b/src/main/rate-limits/antigravity-auth.test.ts new file mode 100644 index 000000000000..a052c5190253 --- /dev/null +++ b/src/main/rate-limits/antigravity-auth.test.ts @@ -0,0 +1,336 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' +import { basename, join } from 'node:path' + +const { + readKeyringMock, + writeKeyringMock, + readFileMock, + writeFileMock, + renameMock, + rmMock, + netFetchMock +} = vi.hoisted(() => ({ + readKeyringMock: vi.fn(), + writeKeyringMock: vi.fn(), + readFileMock: vi.fn(), + writeFileMock: vi.fn(), + renameMock: vi.fn(), + rmMock: vi.fn(), + netFetchMock: vi.fn() +})) + +vi.mock('./antigravity-keychain', () => ({ + readAntigravityKeyring: readKeyringMock, + writeAntigravityKeyring: writeKeyringMock +})) + +vi.mock('node:fs/promises', () => ({ + readFile: readFileMock, + writeFile: writeFileMock, + rename: renameMock, + rm: rmMock +})) + +vi.mock('electron', () => ({ net: { fetch: netFetchMock } })) + +import { + getAntigravityAccessToken, + invalidateAntigravityAccessToken, + readAntigravityCredentials, + type AntigravityCredentialSource +} from './antigravity-auth' + +function credentialRecord(overrides: Record = {}): Record { + return { + auth_method: 'consumer', + token: { + access_token: 'access-token', + refresh_token: 'refresh-token', + expiry: '2026-08-01T12:00:00.000Z', + token_type: 'Bearer' + }, + ...overrides + } +} + +function keyringValue(record = credentialRecord()): string { + return `go-keyring-base64:${Buffer.from(JSON.stringify(record), 'utf8').toString('base64')}` +} + +function response(body: unknown, status = 200): Response { + return { + ok: status >= 200 && status < 300, + status, + json: async () => body + } as Response +} + +let homeCounter = 0 +function testHome(): string { + homeCounter += 1 + return `/tmp/antigravity-auth-test-${homeCounter}` +} + +describe('Antigravity credential authority', () => { + beforeEach(() => { + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-08-01T11:00:00.000Z')) + readKeyringMock.mockReset() + writeKeyringMock.mockReset() + readFileMock.mockReset() + writeFileMock.mockReset() + renameMock.mockReset() + rmMock.mockReset() + netFetchMock.mockReset() + readKeyringMock.mockResolvedValue({ status: 'missing' }) + readFileMock.mockRejectedValue({ code: 'ENOENT' }) + writeKeyringMock.mockResolvedValue(undefined) + writeFileMock.mockResolvedValue(undefined) + renameMock.mockResolvedValue(undefined) + rmMock.mockResolvedValue(undefined) + }) + + it('reads the official keyring before the token file', async () => { + const home = testHome() + readKeyringMock.mockResolvedValue({ status: 'found', value: keyringValue() }) + + const result = await readAntigravityCredentials(home) + + expect(result.source satisfies AntigravityCredentialSource).toBe('official-keychain') + expect(result.accessToken).toBe('access-token') + expect(readFileMock).not.toHaveBeenCalled() + }) + + it('falls back only to the official token file when the keyring is missing', async () => { + const home = testHome() + const raw = JSON.stringify(credentialRecord()) + readFileMock.mockResolvedValue(raw) + const controller = new AbortController() + + const result = await readAntigravityCredentials(home, controller.signal) + + expect(result.source).toBe('official-token-file') + expect(readFileMock).toHaveBeenCalledWith( + join(home, '.gemini', 'antigravity-cli', 'antigravity-oauth-token'), + { encoding: 'utf8', signal: controller.signal } + ) + }) + + it('does not treat an unrelated auth file as an Antigravity credential', async () => { + const home = testHome() + readFileMock.mockImplementation(async (filePath: string) => { + if (basename(filePath) === 'antigravity-oauth-token') { + throw { code: 'ENOENT' } + } + return JSON.stringify(credentialRecord()) + }) + + await expect(readAntigravityCredentials(home)).rejects.toMatchObject({ + failureKind: 'missing-credentials' + }) + }) + + it('treats numeric Unix-second expiry values as epoch milliseconds', async () => { + const home = testHome() + const expiresAtSeconds = Math.floor(Date.now() / 1000) + 3_600 + readKeyringMock.mockResolvedValue({ + status: 'found', + value: keyringValue({ + token: { access_token: 'fresh-access-token', expiry: expiresAtSeconds } + }) + }) + + await expect(getAntigravityAccessToken({ baseHomeDir: home })).resolves.toMatchObject({ + accessToken: 'fresh-access-token' + }) + expect(netFetchMock).not.toHaveBeenCalled() + }) + + it('shares one refresh request across concurrent callers', async () => { + const home = testHome() + readKeyringMock.mockResolvedValue({ + status: 'found', + value: keyringValue({ + token: { + access_token: 'expired-access-token', + refresh_token: 'refresh-token', + expiry: '2026-08-01T09:00:00.000Z' + } + }) + }) + let resolveRefresh!: (value: Response) => void + netFetchMock.mockReturnValue( + new Promise((resolve) => { + resolveRefresh = resolve + }) + ) + + const first = getAntigravityAccessToken({ baseHomeDir: home }) + const second = getAntigravityAccessToken({ baseHomeDir: home }) + await vi.advanceTimersByTimeAsync(0) + + expect(netFetchMock).toHaveBeenCalledTimes(1) + resolveRefresh(response({ access_token: 'refreshed-access-token', expires_in: 3600 })) + + await expect(first).resolves.toMatchObject({ accessToken: 'refreshed-access-token' }) + await expect(second).resolves.toMatchObject({ accessToken: 'refreshed-access-token' }) + }) + + it('lets a remaining caller keep a shared refresh alive after another aborts', async () => { + const home = testHome() + readKeyringMock.mockResolvedValue({ + status: 'found', + value: keyringValue({ + token: { + access_token: 'expired-access-token', + refresh_token: 'refresh-token', + expiry: '2026-08-01T09:00:00.000Z' + } + }) + }) + let resolveRefresh!: (value: Response) => void + netFetchMock.mockReturnValue( + new Promise((resolve) => { + resolveRefresh = resolve + }) + ) + const firstController = new AbortController() + const first = getAntigravityAccessToken({ + baseHomeDir: home, + signal: firstController.signal + }) + const second = getAntigravityAccessToken({ baseHomeDir: home }) + await vi.advanceTimersByTimeAsync(0) + + firstController.abort() + await expect(first).rejects.toMatchObject({ name: 'AbortError' }) + const refreshCall = netFetchMock.mock.calls[0] + expect(refreshCall?.[1]).toEqual(expect.objectContaining({ signal: expect.any(AbortSignal) })) + const refreshOptions = refreshCall?.[1] as RequestInit | undefined + expect(refreshOptions?.signal?.aborted).toBe(false) + + resolveRefresh(response({ access_token: 'refreshed-access-token', expires_in: 3600 })) + await expect(second).resolves.toMatchObject({ accessToken: 'refreshed-access-token' }) + }) + + it('persists a rotated refresh token back to the official keyring', async () => { + const home = testHome() + readKeyringMock.mockResolvedValue({ + status: 'found', + value: keyringValue({ + token: { + access_token: 'expired-access-token', + refresh_token: 'old-refresh-token', + expiry: '2026-08-01T09:00:00.000Z' + } + }) + }) + netFetchMock.mockResolvedValue( + response({ + access_token: 'refreshed-access-token', + refresh_token: 'new-refresh-token', + expires_in: 3600 + }) + ) + + await expect(getAntigravityAccessToken({ baseHomeDir: home })).resolves.toMatchObject({ + credentialSource: 'official-keychain' + }) + expect(writeKeyringMock).toHaveBeenCalledTimes(1) + const written = String(writeKeyringMock.mock.calls[0]?.[0]) + expect(written.startsWith('go-keyring-base64:')).toBe(true) + const decoded = JSON.parse( + Buffer.from(written.slice('go-keyring-base64:'.length), 'base64').toString('utf8') + ) as { token: { refresh_token: string } } + expect(decoded.token.refresh_token).toBe('new-refresh-token') + }) + + it('continues credential persistence after the only consumer aborts', async () => { + const home = testHome() + readKeyringMock.mockResolvedValue({ + status: 'found', + value: keyringValue({ + token: { + access_token: 'expired-access-token', + refresh_token: 'old-refresh-token', + expiry: '2026-08-01T09:00:00.000Z' + } + }) + }) + let resolveWrite!: () => void + writeKeyringMock.mockReturnValue( + new Promise((resolve) => { + resolveWrite = resolve + }) + ) + netFetchMock.mockResolvedValue( + response({ + access_token: 'refreshed-access-token', + refresh_token: 'new-refresh-token', + expires_in: 3600 + }) + ) + const controller = new AbortController() + const pending = getAntigravityAccessToken({ baseHomeDir: home, signal: controller.signal }) + await vi.advanceTimersByTimeAsync(0) + expect(writeKeyringMock).toHaveBeenCalledTimes(1) + + controller.abort() + resolveWrite() + + await expect(pending).rejects.toMatchObject({ name: 'AbortError' }) + expect(writeKeyringMock.mock.calls[0]?.[1]).toBeUndefined() + }) + + it('removes a temporary token file when the atomic rename fails', async () => { + const home = testHome() + readKeyringMock.mockResolvedValue({ status: 'missing' }) + readFileMock.mockResolvedValue( + JSON.stringify( + credentialRecord({ + token: { + access_token: 'expired-access-token', + refresh_token: 'old-refresh-token', + expiry: '2026-08-01T09:00:00.000Z' + } + }) + ) + ) + renameMock.mockRejectedValue(new Error('rename failed')) + netFetchMock.mockResolvedValue( + response({ + access_token: 'refreshed-access-token', + refresh_token: 'new-refresh-token', + expires_in: 3600 + }) + ) + + await expect(getAntigravityAccessToken({ baseHomeDir: home })).rejects.toMatchObject({ + failureKind: 'keychain-unavailable' + }) + expect(rmMock).toHaveBeenCalledWith( + `${join(home, '.gemini', 'antigravity-cli', 'antigravity-oauth-token')}.${process.pid}.tmp`, + { force: true } + ) + }) + + it('re-reads the official authority after an access token is invalidated', async () => { + const home = testHome() + readKeyringMock.mockResolvedValue({ + status: 'found', + value: keyringValue({ token: { access_token: 'first-access-token' } }) + }) + const first = await getAntigravityAccessToken({ baseHomeDir: home }) + + readKeyringMock.mockResolvedValue({ + status: 'found', + value: keyringValue({ token: { access_token: 'new-access-token' } }) + }) + invalidateAntigravityAccessToken(first.sourceKey) + + await expect(getAntigravityAccessToken({ baseHomeDir: home })).resolves.toMatchObject({ + accessToken: 'new-access-token' + }) + expect(readKeyringMock).toHaveBeenCalledTimes(2) + }) +}) diff --git a/src/main/rate-limits/antigravity-auth.ts b/src/main/rate-limits/antigravity-auth.ts new file mode 100644 index 000000000000..87243170643e --- /dev/null +++ b/src/main/rate-limits/antigravity-auth.ts @@ -0,0 +1,202 @@ +import { readFile } from 'node:fs/promises' +import { homedir } from 'node:os' +import { join } from 'node:path' +import { readAntigravityKeyring } from './antigravity-keychain' +import { + AntigravityAuthError, + type AntigravityAccessToken, + type AntigravityCredentialSource, + type ParsedCredentials, + type TokenEnvelope +} from './antigravity-auth-types' +import { refreshAntigravitySingleFlight } from './antigravity-token-refresh' + +export { ANTIGRAVITY_CLIENT_ID, ANTIGRAVITY_CLIENT_SECRET } from './antigravity-token-refresh' +export { AntigravityAuthError } from './antigravity-auth-types' +export type { + AntigravityAccessToken, + AntigravityCredentialSource, + ParsedCredentials +} from './antigravity-auth-types' + +const ANTIGRAVITY_TOKEN_PATH_PARTS = [ + '.gemini', + 'antigravity-cli', + 'antigravity-oauth-token' +] as const +const TOKEN_SKEW_MS = 60_000 +const KEYRING_VALUE_PREFIX = 'go-keyring-base64:' + +let cachedCredentials: ParsedCredentials | null = null + +export async function getAntigravityAccessToken( + options: { + baseHomeDir?: string + forceRefresh?: boolean + signal?: AbortSignal + } = {} +): Promise { + const baseHomeDir = options.baseHomeDir ?? homedir() + throwIfAborted(options.signal) + let credentials = + cachedCredentials?.baseHomeDir === baseHomeDir + ? cachedCredentials + : await readAntigravityCredentials(baseHomeDir, options.signal) + + if (!options.forceRefresh && isAccessTokenFresh(credentials)) { + return toAccessToken(credentials) + } + if (!credentials.refreshToken) { + throw new AntigravityAuthError('Antigravity sign-in expired', 'stale-token') + } + + const refreshed = await refreshAntigravitySingleFlight(credentials, options.signal) + credentials = { + ...credentials, + accessToken: refreshed.accessToken, + refreshToken: refreshed.refreshToken, + expiresAtMs: refreshed.expiresAtMs + } + cachedCredentials = credentials + return toAccessToken(credentials) +} + +export function invalidateAntigravityAccessToken(sourceKey: string): void { + if (cachedCredentials?.sourceKey === sourceKey) { + // Why: a 401 can mean the CLI was re-authenticated elsewhere; force the + // next request to re-read the official keyring/token file authority. + cachedCredentials = null + } +} + +export async function readAntigravityCredentials( + baseHomeDir = homedir(), + signal?: AbortSignal +): Promise { + throwIfAborted(signal) + const keyring = await readAntigravityKeyring(signal) + if (keyring.status === 'found') { + const parsed = parseCredentialValue(keyring.value, 'official-keychain', baseHomeDir) + cachedCredentials = parsed + return parsed + } + + const tokenPath = join(baseHomeDir, ...ANTIGRAVITY_TOKEN_PATH_PARTS) + try { + const raw = await readFile(tokenPath, { encoding: 'utf8', signal }) + const parsed = parseCredentialValue(raw, 'official-token-file', baseHomeDir) + cachedCredentials = parsed + return parsed + } catch (error) { + if (error instanceof AntigravityAuthError) { + throw error + } + if (!isFileNotFoundError(error)) { + throw new AntigravityAuthError('Unable to read Antigravity credentials', 'parse') + } + if (keyring.status === 'unavailable') { + throw new AntigravityAuthError( + 'Antigravity system keyring is unavailable', + 'keychain-unavailable' + ) + } + throw new AntigravityAuthError('Antigravity credentials not found', 'missing-credentials') + } +} + +function toAccessToken(credentials: ParsedCredentials): AntigravityAccessToken { + if (!credentials.accessToken) { + throw new AntigravityAuthError('Antigravity access token is missing', 'stale-token') + } + return { + accessToken: credentials.accessToken, + credentialSource: credentials.source, + sourceKey: credentials.sourceKey + } +} + +function parseCredentialValue( + raw: string, + source: AntigravityCredentialSource, + baseHomeDir: string +): ParsedCredentials { + let decoded = raw.trim() + if (decoded.startsWith(KEYRING_VALUE_PREFIX)) { + decoded = Buffer.from(decoded.slice(KEYRING_VALUE_PREFIX.length), 'base64').toString('utf8') + } + + let envelope: unknown + try { + envelope = JSON.parse(decoded) + } catch { + throw new AntigravityAuthError('Antigravity credential record is invalid', 'parse') + } + if (!isTokenEnvelope(envelope)) { + throw new AntigravityAuthError('Antigravity credential record is invalid', 'parse') + } + return { + source, + sourceKey: `${source}:${baseHomeDir}`, + tokenPath: + source === 'official-token-file' ? join(baseHomeDir, ...ANTIGRAVITY_TOKEN_PATH_PARTS) : null, + envelope, + accessToken: envelope.token.access_token, + refreshToken: + typeof envelope.token.refresh_token === 'string' ? envelope.token.refresh_token : '', + expiresAtMs: parseExpiry(envelope.token.expiry), + baseHomeDir + } +} + +function isTokenEnvelope(value: unknown): value is TokenEnvelope { + if (typeof value !== 'object' || value === null || !('token' in value)) { + return false + } + const token = (value as { token?: unknown }).token + return ( + typeof token === 'object' && + token !== null && + 'access_token' in token && + typeof (token as { access_token?: unknown }).access_token === 'string' && + (token as { access_token: string }).access_token.length > 0 + ) +} + +function parseExpiry(value: string | number | undefined): number | null { + if (typeof value === 'number') { + // Why: the official agy CLI writes ISO-8601 strings, but older credential + // records may contain Unix seconds while other OAuth stores use millis. + if (!Number.isFinite(value)) { + return null + } + return value > 1_000_000_000_000 ? value : value * 1000 + } + if (typeof value === 'string') { + const parsed = Date.parse(value) + return Number.isFinite(parsed) ? parsed : null + } + return null +} + +function isAccessTokenFresh(credentials: ParsedCredentials): boolean { + return credentials.expiresAtMs === null || credentials.expiresAtMs - Date.now() > TOKEN_SKEW_MS +} + +function isFileNotFoundError(error: unknown): boolean { + return Boolean( + error && + typeof error === 'object' && + 'code' in error && + (error as { code?: unknown }).code === 'ENOENT' + ) +} + +function throwIfAborted(signal?: AbortSignal): void { + if (signal?.aborted) { + throw createAbortError() + } +} + +function createAbortError(): Error { + return Object.assign(new Error('Antigravity request aborted'), { name: 'AbortError' }) +} diff --git a/src/main/rate-limits/antigravity-keychain.test.ts b/src/main/rate-limits/antigravity-keychain.test.ts new file mode 100644 index 000000000000..c006f063a494 --- /dev/null +++ b/src/main/rate-limits/antigravity-keychain.test.ts @@ -0,0 +1,132 @@ +import { execFile } from 'node:child_process' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +vi.mock('node:child_process', () => ({ + execFile: vi.fn() +})) + +import { readAntigravityKeyring, writeAntigravityKeyring } from './antigravity-keychain' + +const execFileMock = vi.mocked(execFile) +const originalPlatform = Object.getOwnPropertyDescriptor(process, 'platform') + +function setPlatform(platform: NodeJS.Platform): void { + Object.defineProperty(process, 'platform', { configurable: true, value: platform }) +} + +function invokeCallback(callback: unknown, error: Error | null, stdout = '', stderr = ''): void { + const execCallback = callback as (error: Error | null, stdout: string, stderr: string) => void + execCallback(error, stdout, stderr) +} + +describe('Antigravity system keyring adapter', () => { + beforeEach(() => { + setPlatform('darwin') + execFileMock.mockReset() + }) + + afterEach(() => { + vi.useRealTimers() + if (originalPlatform) { + Object.defineProperty(process, 'platform', originalPlatform) + } + }) + + it('reads the official macOS service/account pair', async () => { + execFileMock.mockImplementationOnce((_file, _args, _options, callback) => { + invokeCallback(callback, null, 'go-keyring-base64:record\n') + return null as never + }) + + await expect(readAntigravityKeyring()).resolves.toEqual({ + status: 'found', + value: 'go-keyring-base64:record' + }) + expect(execFileMock.mock.calls[0]?.[1]).toEqual([ + 'find-generic-password', + '-s', + 'gemini', + '-a', + 'antigravity', + '-w' + ]) + }) + + it('writes rotated credentials back to the official macOS item', async () => { + execFileMock.mockImplementationOnce((_file, _args, _options, callback) => { + invokeCallback(callback, null) + return null as never + }) + + await writeAntigravityKeyring('rotated-record') + + expect(execFileMock.mock.calls[0]?.[1]).toEqual([ + 'add-generic-password', + '-U', + '-s', + 'gemini', + '-a', + 'antigravity', + '-w', + 'rotated-record' + ]) + }) + + it('treats an absent macOS item as missing rather than a keyring failure', async () => { + const notFound = Object.assign(new Error('not found'), { code: 44 }) + execFileMock.mockImplementationOnce((_file, _args, _options, callback) => { + invokeCallback(callback, notFound, '', 'could not be found') + return null as never + }) + + await expect(readAntigravityKeyring()).resolves.toEqual({ status: 'missing' }) + }) + + it('kills a pending keyring command when the fetch cycle aborts', async () => { + const killMock = vi.fn() + execFileMock.mockImplementationOnce(() => ({ kill: killMock }) as never) + const controller = new AbortController() + const read = readAntigravityKeyring(controller.signal) + + controller.abort() + + await expect(read).rejects.toMatchObject({ name: 'AbortError' }) + expect(killMock).toHaveBeenCalled() + }) + + it('uses PowerShell credential-manager access on Windows', async () => { + setPlatform('win32') + execFileMock.mockImplementationOnce((_file, _args, _options, callback) => { + invokeCallback(callback, null, 'record') + return null as never + }) + + await expect(readAntigravityKeyring()).resolves.toEqual({ status: 'found', value: 'record' }) + expect(execFileMock.mock.calls[0]?.[0]).toBe('powershell.exe') + expect(execFileMock.mock.calls[0]?.[1]).toEqual( + expect.arrayContaining(['-EncodedCommand', expect.any(String)]) + ) + }) + + it('classifies an absent Linux Secret Service item as missing', async () => { + setPlatform('linux') + const notFound = Object.assign(new Error('secret not found'), { code: 1 }) + execFileMock.mockImplementationOnce((_file, _args, _options, callback) => { + invokeCallback(callback, notFound, '', 'secret not found') + return null as never + }) + + await expect(readAntigravityKeyring()).resolves.toEqual({ status: 'missing' }) + }) + + it('classifies a Linux Secret Service failure as unavailable', async () => { + setPlatform('linux') + const failure = Object.assign(new Error('dbus unavailable'), { code: 1 }) + execFileMock.mockImplementationOnce((_file, _args, _options, callback) => { + invokeCallback(callback, failure, '', 'dbus unavailable') + return null as never + }) + + await expect(readAntigravityKeyring()).resolves.toEqual({ status: 'unavailable' }) + }) +}) diff --git a/src/main/rate-limits/antigravity-keychain.ts b/src/main/rate-limits/antigravity-keychain.ts new file mode 100644 index 000000000000..58edd2d840b5 --- /dev/null +++ b/src/main/rate-limits/antigravity-keychain.ts @@ -0,0 +1,311 @@ +import { execFile } from 'node:child_process' +import { encodePowerShellCommand } from '../../shared/powershell-command-encoding' + +const KEYCHAIN_SERVICE = 'gemini' +const KEYCHAIN_ACCOUNT = 'antigravity' +const KEYCHAIN_COMMAND_TIMEOUT_MS = 3_000 +const KEYCHAIN_MAX_BUFFER = 128 * 1024 + +export type AntigravityKeyringReadResult = + | { status: 'found'; value: string } + | { status: 'missing' } + | { status: 'unavailable' } + +type CommandOptions = { + input?: string + signal?: AbortSignal +} + +type CommandResult = { + stdout: string + stderr: string +} + +// Why: keep the official Go keyring authority separate from Orca safeStorage. +export async function readAntigravityKeyring( + signal?: AbortSignal +): Promise { + if (process.platform === 'darwin') { + return readMacKeychain(signal) + } + if (process.platform === 'linux') { + return readLinuxSecretService(signal) + } + if (process.platform === 'win32') { + return readWindowsCredentialManager(signal) + } + return { status: 'unavailable' } +} + +export async function writeAntigravityKeyring(value: string, signal?: AbortSignal): Promise { + if (process.platform === 'darwin') { + await runCommand( + 'security', + ['add-generic-password', '-U', '-s', KEYCHAIN_SERVICE, '-a', KEYCHAIN_ACCOUNT, '-w', value], + { signal } + ) + return + } + if (process.platform === 'linux') { + await runCommand( + 'secret-tool', + [ + 'store', + '--label=Antigravity OAuth token', + 'service', + KEYCHAIN_SERVICE, + 'account', + KEYCHAIN_ACCOUNT + ], + { input: value, signal } + ) + return + } + if (process.platform === 'win32') { + await runCommand( + 'powershell.exe', + [ + '-NoLogo', + '-NoProfile', + '-NonInteractive', + '-EncodedCommand', + encodePowerShellCommand(WINDOWS_CREDENTIAL_WRITE_SCRIPT) + ], + { input: value, signal } + ) + return + } + throw new Error('System keyring is unavailable on this platform') +} + +async function readMacKeychain(signal?: AbortSignal): Promise { + try { + const { stdout } = await runCommand( + 'security', + ['find-generic-password', '-s', KEYCHAIN_SERVICE, '-a', KEYCHAIN_ACCOUNT, '-w'], + { signal } + ) + return stdout.trim() ? { status: 'found', value: stdout.trim() } : { status: 'missing' } + } catch (error) { + if (isAbortError(error)) { + throw error + } + if (isNotFoundError(error)) { + return { status: 'missing' } + } + return { status: 'unavailable' } + } +} + +async function readLinuxSecretService(signal?: AbortSignal): Promise { + try { + const { stdout } = await runCommand( + 'secret-tool', + ['lookup', 'service', KEYCHAIN_SERVICE, 'account', KEYCHAIN_ACCOUNT], + { signal } + ) + return stdout.trim() ? { status: 'found', value: stdout.trim() } : { status: 'missing' } + } catch (error) { + if (isAbortError(error)) { + throw error + } + if (isCommandMissingError(error) || !isNotFoundError(error)) { + return { status: 'unavailable' } + } + return { status: 'missing' } + } +} + +async function readWindowsCredentialManager( + signal?: AbortSignal +): Promise { + try { + const { stdout } = await runCommand( + 'powershell.exe', + [ + '-NoLogo', + '-NoProfile', + '-NonInteractive', + '-EncodedCommand', + encodePowerShellCommand(WINDOWS_CREDENTIAL_READ_SCRIPT) + ], + { signal } + ) + return stdout.trim() ? { status: 'found', value: stdout.trim() } : { status: 'missing' } + } catch (error) { + if (isAbortError(error)) { + throw error + } + return isNotFoundError(error) ? { status: 'missing' } : { status: 'unavailable' } + } +} + +function runCommand( + command: string, + args: string[], + options: CommandOptions = {} +): Promise { + return new Promise((resolve, reject) => { + let settled = false + let child: ReturnType | undefined + const timer = setTimeout(() => { + if (settled) { + return + } + settled = true + child?.kill() + reject(Object.assign(new Error('System keyring command timed out'), { code: 'ETIMEDOUT' })) + }, KEYCHAIN_COMMAND_TIMEOUT_MS) + + const finish = (callback: () => void): void => { + if (settled) { + return + } + settled = true + clearTimeout(timer) + options.signal?.removeEventListener('abort', abort) + callback() + } + + const abort = (): void => { + child?.kill() + finish(() => + reject(Object.assign(new Error('System keyring command aborted'), { name: 'AbortError' })) + ) + } + + if (options.signal?.aborted) { + abort() + return + } + options.signal?.addEventListener('abort', abort, { once: true }) + + try { + child = execFile( + command, + args, + { + encoding: 'utf8', + maxBuffer: KEYCHAIN_MAX_BUFFER, + timeout: KEYCHAIN_COMMAND_TIMEOUT_MS, + windowsHide: true + }, + (error, stdout, stderr) => { + if (error) { + finish(() => + reject( + Object.assign(error, { + stdout: String(stdout), + stderr: String(stderr) + }) + ) + ) + return + } + finish(() => resolve({ stdout: String(stdout), stderr: String(stderr) })) + } + ) + if (options.input !== undefined && child.stdin) { + // Why: ignore EPIPE from a child killed before stdin is read. + child.stdin.once('error', () => undefined) + child.stdin.end(options.input) + } + } catch (error) { + finish(() => reject(error)) + } + }) +} + +function isCommandMissingError(error: unknown): boolean { + return Boolean( + error && + typeof error === 'object' && + 'code' in error && + (error as { code?: unknown }).code === 'ENOENT' + ) +} + +function isNotFoundError(error: unknown): boolean { + const code = + error && typeof error === 'object' && 'code' in error + ? (error as { code?: unknown }).code + : undefined + const errObj = + error && typeof error === 'object' ? (error as { stderr?: unknown; message?: unknown }) : null + const stderrStr = typeof errObj?.stderr === 'string' ? errObj.stderr : '' + const messageStr = typeof errObj?.message === 'string' ? errObj.message : '' + const message = errObj ? `${stderrStr} ${messageStr}` : String(error) + return code === 44 || /could not be found|not found/i.test(message) +} + +function isAbortError(error: unknown): boolean { + return Boolean( + error && typeof error === 'object' && 'name' in error && error.name === 'AbortError' + ) +} + +const WINDOWS_CREDENTIAL_LIBRARY = ` +Add-Type -TypeDefinition @' +using System; +using System.Runtime.InteropServices; +using System.Text; +public static class OrcaAntigravityCredential { + [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)] + public struct CREDENTIAL { + public uint Flags, Type; + public string TargetName, Comment; + public System.Runtime.InteropServices.ComTypes.FILETIME LastWritten; + public uint CredentialBlobSize; + public IntPtr CredentialBlob; + public uint Persist, AttributeCount; + public IntPtr Attributes; + public string TargetAlias, UserName; + } + [DllImport("advapi32.dll", CharSet = CharSet.Unicode, SetLastError = true)] + public static extern bool CredRead(string target, uint type, uint flags, out IntPtr credential); + [DllImport("advapi32.dll", SetLastError = true)] + public static extern void CredFree(IntPtr credential); + [DllImport("advapi32.dll", CharSet = CharSet.Unicode, SetLastError = true)] + public static extern bool CredWrite(ref CREDENTIAL credential, uint flags); + public static string Read(string target) { + IntPtr pointer; + if (!CredRead(target, 1, 0, out pointer)) return null; + try { + CREDENTIAL cred = (CREDENTIAL)Marshal.PtrToStructure(pointer, typeof(CREDENTIAL)); + if (cred.CredentialBlob == IntPtr.Zero || cred.CredentialBlobSize == 0) return null; + byte[] bytes = new byte[cred.CredentialBlobSize]; + Marshal.Copy(cred.CredentialBlob, bytes, 0, (int)cred.CredentialBlobSize); + return Encoding.UTF8.GetString(bytes); + } finally { CredFree(pointer); } + } + public static bool Write(string target, string value) { + byte[] bytes = Encoding.UTF8.GetBytes(value); + IntPtr blob = Marshal.AllocHGlobal(bytes.Length); + try { + Marshal.Copy(bytes, 0, blob, bytes.Length); + CREDENTIAL cred = new CREDENTIAL { + Type = 1, TargetName = target, CredentialBlobSize = (uint)bytes.Length, + CredentialBlob = blob, Persist = 2, UserName = target + }; + return CredWrite(ref cred, 0); + } finally { Marshal.FreeHGlobal(blob); } + } +} +'@ +` + +const WINDOWS_CREDENTIAL_READ_SCRIPT = `${WINDOWS_CREDENTIAL_LIBRARY} +$target = 'gemini:antigravity' +$value = [OrcaAntigravityCredential]::Read($target) +if ($null -eq $value) { $value = [OrcaAntigravityCredential]::Read('gemini/antigravity') } +if ($null -ne $value) { + [Console]::Out.Write($value) +} +` + +const WINDOWS_CREDENTIAL_WRITE_SCRIPT = `${WINDOWS_CREDENTIAL_LIBRARY} +$inputValue = [Console]::In.ReadToEnd() +if ($inputValue.Length -eq 0 -or -not [OrcaAntigravityCredential]::Write('gemini:antigravity', $inputValue)) { + exit 1 +} +` 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..104c8a72306b --- /dev/null +++ b/src/main/rate-limits/antigravity-quota-parser.ts @@ -0,0 +1,78 @@ +import type { RateLimitBucket } from '../../shared/rate-limit-types' + +const QUOTA_BUCKETS = { + 'gemini-5h': { name: 'Antigravity 5h', windowMinutes: 300 }, + 'gemini-weekly': { name: 'Antigravity weekly', windowMinutes: 10_080 }, + '3p-5h': { name: '3-party 5h', windowMinutes: 300 }, + '3p-weekly': { name: '3-party weekly', windowMinutes: 10_080 } +} as const + +type QuotaBucketId = keyof typeof QUOTA_BUCKETS + +type RawQuotaBucket = { + bucketId?: unknown + remainingFraction?: unknown + resetTime?: unknown +} + +export type AntigravityQuotaBucket = { + id: QuotaBucketId + bucket: RateLimitBucket +} + +export function parseAntigravityQuotaBuckets(data: unknown): AntigravityQuotaBucket[] { + const byId = new Map() + for (const raw of collectRawQuotaBuckets(data)) { + const normalized = normalizeQuotaBucket(raw) + if (normalized && !byId.has(normalized.id)) { + byId.set(normalized.id, normalized) + } + } + return (Object.keys(QUOTA_BUCKETS) as QuotaBucketId[]) + .map((id) => byId.get(id)) + .filter((bucket): bucket is AntigravityQuotaBucket => bucket !== undefined) +} + +function collectRawQuotaBuckets(data: unknown): RawQuotaBucket[] { + if (typeof data !== 'object' || data === null) { + return [] + } + const value = data as { groups?: unknown; buckets?: unknown } + const grouped = Array.isArray(value.groups) + ? value.groups.flatMap((group) => { + if (typeof group !== 'object' || group === null || !('buckets' in group)) { + return [] + } + const buckets = (group as { buckets?: unknown }).buckets + return Array.isArray(buckets) ? buckets : [] + }) + : [] + const flat = Array.isArray(value.buckets) ? value.buckets : [] + return [...grouped, ...flat].filter(isRawQuotaBucket) +} + +function isRawQuotaBucket(value: unknown): value is RawQuotaBucket { + return typeof value === 'object' && value !== null +} + +function normalizeQuotaBucket(raw: RawQuotaBucket): AntigravityQuotaBucket | null { + if (typeof raw.bucketId !== 'string' || !(raw.bucketId in QUOTA_BUCKETS)) { + return null + } + const id = raw.bucketId as QuotaBucketId + if (typeof raw.remainingFraction !== 'number' || !Number.isFinite(raw.remainingFraction)) { + return null + } + const resetTime = typeof raw.resetTime === 'string' ? Date.parse(raw.resetTime) : Number.NaN + const definition = QUOTA_BUCKETS[id] + return { + id, + bucket: { + name: definition.name, + usedPercent: Math.round((1 - Math.min(1, Math.max(0, raw.remainingFraction))) * 100), + windowMinutes: definition.windowMinutes, + resetsAt: Number.isFinite(resetTime) ? resetTime : null, + resetDescription: null + } + } +} diff --git a/src/main/rate-limits/antigravity-token-refresh.ts b/src/main/rate-limits/antigravity-token-refresh.ts new file mode 100644 index 000000000000..ac01485ec6ac --- /dev/null +++ b/src/main/rate-limits/antigravity-token-refresh.ts @@ -0,0 +1,212 @@ +import { rename, rm, writeFile } from 'node:fs/promises' +import { net } from 'electron' +import { writeAntigravityKeyring } from './antigravity-keychain' +import { + AntigravityAuthError, + type ParsedCredentials, + type RefreshedCredentials, + type TokenEnvelope +} from './antigravity-auth-types' + +const GOOGLE_TOKEN_URL = 'https://oauth2.googleapis.com/token' +const API_TIMEOUT_MS = 10_000 +const KEYRING_VALUE_PREFIX = 'go-keyring-base64:' + +// Why: these public native-app credentials are embedded in the official agy binary. +export const ANTIGRAVITY_CLIENT_ID = + '1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com' +export const ANTIGRAVITY_CLIENT_SECRET = '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' + +type RefreshFlight = { + sourceKey: string + controller: AbortController + promise: Promise + consumers: number +} + +const refreshFlights = new Map() + +export function refreshAntigravitySingleFlight( + credentials: ParsedCredentials, + signal?: AbortSignal +): Promise { + const existing = refreshFlights.get(credentials.sourceKey) + if (existing) { + return waitForRefresh(existing, signal) + } + + const controller = new AbortController() + let flight!: RefreshFlight + const promise = refreshCredential(credentials, controller.signal) + .then(async (refreshed) => { + if (refreshed.refreshToken !== credentials.refreshToken) { + // Why: once OAuth has returned a rotated credential, persistence must + // finish even if the last consumer aborts its wait for the result. + await persistRotatedCredentials(credentials, refreshed) + } + return refreshed + }) + .finally(() => { + if (refreshFlights.get(credentials.sourceKey) === flight) { + refreshFlights.delete(credentials.sourceKey) + } + }) + flight = { sourceKey: credentials.sourceKey, controller, promise, consumers: 0 } + refreshFlights.set(credentials.sourceKey, flight) + return waitForRefresh(flight, signal) +} + +function waitForRefresh( + flight: RefreshFlight, + signal?: AbortSignal +): Promise { + flight.consumers += 1 + return new Promise((resolve, reject) => { + let settled = false + const release = (): void => { + flight.consumers = Math.max(0, flight.consumers - 1) + if (flight.consumers === 0) { + if (refreshFlights.get(flight.sourceKey) === flight) { + refreshFlights.delete(flight.sourceKey) + } + flight.controller.abort() + } + } + const finish = (callback: () => void): void => { + if (settled) { + return + } + settled = true + signal?.removeEventListener('abort', onAbort) + release() + callback() + } + const onAbort = (): void => { + finish(() => reject(createAbortError())) + } + + signal?.addEventListener('abort', onAbort, { once: true }) + flight.promise.then( + (value) => finish(() => resolve(value)), + (error: unknown) => finish(() => reject(error)) + ) + if (signal?.aborted) { + onAbort() + } + }) +} + +async function refreshCredential( + credentials: ParsedCredentials, + signal: AbortSignal +): Promise { + const requestSignal = AbortSignal.any([signal, AbortSignal.timeout(API_TIMEOUT_MS)]) + let response: Response + try { + response = await net.fetch(GOOGLE_TOKEN_URL, { + method: 'POST', + headers: { 'Content-Type': 'application/x-www-form-urlencoded' }, + body: new URLSearchParams({ + client_id: ANTIGRAVITY_CLIENT_ID, + client_secret: ANTIGRAVITY_CLIENT_SECRET, + refresh_token: credentials.refreshToken, + grant_type: 'refresh_token' + }).toString(), + signal: requestSignal + }) + } catch (error) { + if (isAbortError(error)) { + throw error + } + throw new AntigravityAuthError('Antigravity token refresh failed', 'network') + } + + const body = (await response.json().catch(() => ({}))) as { + access_token?: unknown + refresh_token?: unknown + expires_in?: unknown + id_token?: unknown + } + if (!response.ok) { + const failureKind = + response.status === 400 || response.status === 401 ? 'stale-token' : 'server' + throw new AntigravityAuthError('Antigravity token refresh failed', failureKind, response.status) + } + if ( + typeof body.access_token !== 'string' || + typeof body.expires_in !== 'number' || + !Number.isFinite(body.expires_in) || + body.expires_in <= 0 + ) { + throw new AntigravityAuthError('Antigravity token refresh response is invalid', 'parse') + } + return { + accessToken: body.access_token, + refreshToken: + typeof body.refresh_token === 'string' && body.refresh_token.length > 0 + ? body.refresh_token + : credentials.refreshToken, + expiresAtMs: Date.now() + body.expires_in * 1000, + ...(typeof body.id_token === 'string' ? { idToken: body.id_token } : {}) + } +} + +async function persistRotatedCredentials( + credentials: ParsedCredentials, + refreshed: RefreshedCredentials +): Promise { + const envelope = buildUpdatedEnvelope(credentials.envelope, refreshed) + try { + if (credentials.source === 'official-keychain') { + await writeAntigravityKeyring( + `${KEYRING_VALUE_PREFIX}${Buffer.from(JSON.stringify(envelope), 'utf8').toString('base64')}` + ) + return + } + if (!credentials.tokenPath) { + throw new Error('Antigravity token file path is missing') + } + const temporaryPath = `${credentials.tokenPath}.${process.pid}.tmp` + try { + await writeFile(temporaryPath, JSON.stringify(envelope, null, 2), { + encoding: 'utf8', + mode: 0o600 + }) + await rename(temporaryPath, credentials.tokenPath) + } catch (error) { + await rm(temporaryPath, { force: true }).catch(() => undefined) + throw error + } + } catch { + throw new AntigravityAuthError( + 'Antigravity rotated credentials could not be saved', + 'keychain-unavailable' + ) + } +} + +function buildUpdatedEnvelope( + envelope: TokenEnvelope, + refreshed: RefreshedCredentials +): TokenEnvelope { + return { + ...envelope, + ...(refreshed.idToken ? { id_token: refreshed.idToken } : {}), + token: { + ...envelope.token, + access_token: refreshed.accessToken, + refresh_token: refreshed.refreshToken, + expiry: new Date(refreshed.expiresAtMs).toISOString() + } + } +} + +function createAbortError(): Error { + return Object.assign(new Error('Antigravity request aborted'), { name: 'AbortError' }) +} + +function isAbortError(error: unknown): boolean { + return Boolean( + error && typeof error === 'object' && 'name' in error && error.name === 'AbortError' + ) +} 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..d6597c8b050a --- /dev/null +++ b/src/main/rate-limits/antigravity-usage-fetcher.test.ts @@ -0,0 +1,280 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +const { getTokenMock, invalidateTokenMock, netFetchMock, MockAntigravityAuthError } = vi.hoisted( + () => { + class HoistedMockAntigravityAuthError extends Error { + readonly failureKind: string + readonly status: number | null + + constructor(message: string, failureKind: string, status: number | null = null) { + super(message) + this.name = 'AntigravityAuthError' + this.failureKind = failureKind + this.status = status + } + } + return { + getTokenMock: vi.fn(), + invalidateTokenMock: vi.fn(), + netFetchMock: vi.fn(), + MockAntigravityAuthError: HoistedMockAntigravityAuthError + } + } +) + +vi.mock('./antigravity-auth', () => ({ + AntigravityAuthError: MockAntigravityAuthError, + getAntigravityAccessToken: getTokenMock, + invalidateAntigravityAccessToken: invalidateTokenMock +})) + +vi.mock('electron', () => ({ net: { fetch: netFetchMock } })) + +import { + fetchAntigravityRateLimits, + parseAntigravityQuotaBuckets +} from './antigravity-usage-fetcher' + +function response(body: unknown, status = 200, headers: Record = {}): Response { + return { + ok: status >= 200 && status < 300, + status, + headers: new Headers(headers), + json: async () => body + } as Response +} + +function token(accessToken = 'access-token') { + return { + accessToken, + credentialSource: 'official-keychain' as const, + sourceKey: 'official-keychain:/tmp/test' + } +} + +const quotaGroups = [ + { + buckets: [ + { + bucketId: 'gemini-weekly', + remainingFraction: 0.75, + resetTime: '2026-08-04T17:38:47Z' + }, + { + bucketId: 'gemini-5h', + remainingFraction: 0.9, + resetTime: '2026-08-01T18:30:17Z' + } + ] + }, + { + buckets: [ + { bucketId: '3p-weekly', remainingFraction: 0.5, resetTime: '2026-08-06T13:00:53Z' }, + { bucketId: '3p-5h', remainingFraction: 1, resetTime: '2026-08-01T18:37:49Z' } + ] + } +] + +describe('fetchAntigravityRateLimits', () => { + beforeEach(() => { + vi.useFakeTimers() + vi.setSystemTime(new Date('2026-08-01T11:00:00.000Z')) + getTokenMock.mockReset() + invalidateTokenMock.mockReset() + netFetchMock.mockReset() + getTokenMock.mockResolvedValue(token()) + netFetchMock.mockImplementation((url: string) => { + if (url.includes('loadCodeAssist')) { + return Promise.resolve(response({ cloudaicompanionProject: 'project-123' })) + } + if (url.includes('retrieveUserQuotaSummary')) { + return Promise.resolve(response({ groups: quotaGroups })) + } + return Promise.resolve(response({}, 404)) + }) + }) + + it('uses Antigravity metadata and maps all current quota identities', async () => { + const signal = new AbortController().signal + const result = await fetchAntigravityRateLimits({ signal }) + + expect(result.status).toBe('ok') + expect(result.provider).toBe('antigravity') + expect(result.session?.usedPercent).toBe(10) + expect(result.weekly?.usedPercent).toBe(25) + expect(result.buckets?.map((bucket) => bucket.name)).toEqual([ + 'Antigravity 5h', + 'Antigravity weekly', + '3-party 5h', + '3-party weekly' + ]) + expect(getTokenMock).toHaveBeenCalledWith(expect.objectContaining({ signal })) + + const loadCall = netFetchMock.mock.calls.find((call) => + String(call[0]).includes('loadCodeAssist') + ) + expect(loadCall).toBeDefined() + const loadBody = JSON.parse((loadCall![1] as RequestInit).body as string) as { + metadata: { platform: string; ideType: string; pluginType: string } + } + expect(loadBody.metadata).toEqual({ + ideType: 'ANTIGRAVITY', + platform: 'PLATFORM_UNSPECIFIED', + pluginType: 'GEMINI' + }) + expect((loadCall![1] as RequestInit).headers).toEqual( + expect.objectContaining({ + Authorization: 'Bearer access-token', + 'Client-Metadata': JSON.stringify(loadBody.metadata) + }) + ) + }) + + it('deduplicates by bucket identity and ignores unknown or invalid buckets', () => { + const parsed = parseAntigravityQuotaBuckets({ + groups: [ + { + buckets: [ + { bucketId: 'gemini-5h', remainingFraction: 0.8, resetTime: '2026-08-01T12:00:00Z' }, + { bucketId: 'gemini-5h', remainingFraction: 0.2, resetTime: '2026-08-01T12:00:00Z' }, + { bucketId: 'unknown-window', remainingFraction: 0.1 }, + { bucketId: '3p-5h', remainingFraction: Number.NaN } + ] + } + ] + }) + + expect(parsed).toHaveLength(1) + expect(parsed[0]?.id).toBe('gemini-5h') + expect(parsed[0]?.bucket.usedPercent).toBe(20) + }) + + it('does not turn an empty successful summary into an ok result or a fallback storm', async () => { + netFetchMock.mockImplementation((url: string) => { + if (url.includes('loadCodeAssist')) { + return Promise.resolve(response({ cloudaicompanionProject: 'project-123' })) + } + if (url.includes('retrieveUserQuotaSummary')) { + return Promise.resolve(response({ groups: [] })) + } + return Promise.resolve(response({ buckets: quotaGroups[0].buckets })) + }) + + const result = await fetchAntigravityRateLimits() + + expect(result.status).toBe('error') + expect(result.usageMetadata?.failureKind).toBe('usage-unavailable') + expect(netFetchMock).toHaveBeenCalledTimes(2) + expect( + netFetchMock.mock.calls.some((call) => + String(call[0]).endsWith('v1internal:retrieveUserQuota') + ) + ).toBe(false) + }) + + it('uses the legacy endpoint only when the summary endpoint is unsupported', async () => { + netFetchMock.mockImplementation((url: string) => { + if (url.includes('loadCodeAssist')) { + return Promise.resolve(response({ cloudaicompanionProject: 'project-123' })) + } + if (url.includes('retrieveUserQuotaSummary')) { + return Promise.resolve(response({}, 404)) + } + return Promise.resolve(response({ buckets: [quotaGroups[0].buckets[1]] })) + }) + + const result = await fetchAntigravityRateLimits() + + expect(result.status).toBe('ok') + expect(result.session?.usedPercent).toBe(10) + expect(netFetchMock).toHaveBeenCalledTimes(3) + }) + + it('reports an empty legacy response with the summary endpoint failure context', async () => { + netFetchMock.mockImplementation((url: string) => { + if (url.includes('loadCodeAssist')) { + return Promise.resolve(response({ cloudaicompanionProject: 'project-123' })) + } + if (url.includes('retrieveUserQuotaSummary')) { + return Promise.resolve(response({}, 404)) + } + return Promise.resolve(response({ groups: [] })) + }) + + const result = await fetchAntigravityRateLimits() + + expect(result.status).toBe('error') + expect(result.usageMetadata?.failureKind).toBe('usage-unavailable') + expect(netFetchMock).toHaveBeenCalledTimes(3) + }) + + it('refreshes once after a quota API 401 or 403', async () => { + let loadCalls = 0 + netFetchMock.mockImplementation((url: string) => { + if (url.includes('loadCodeAssist')) { + loadCalls += 1 + return Promise.resolve( + loadCalls === 1 + ? response({ error: 'Unauthenticated' }, 401) + : response({ cloudaicompanionProject: 'project-123' }) + ) + } + return Promise.resolve(response({ groups: quotaGroups })) + }) + getTokenMock + .mockResolvedValueOnce(token('old-access-token')) + .mockResolvedValueOnce(token('new-access-token')) + + const result = await fetchAntigravityRateLimits() + + expect(result.status).toBe('ok') + expect(invalidateTokenMock).toHaveBeenCalledWith('official-keychain:/tmp/test') + expect(getTokenMock).toHaveBeenNthCalledWith(2, { forceRefresh: true }) + expect( + netFetchMock.mock.calls.some( + (call) => + (call[1] as RequestInit).headers && JSON.stringify(call[1]).includes('new-access-token') + ) + ).toBe(true) + + netFetchMock.mockReset() + let forbiddenLoadCalls = 0 + netFetchMock.mockImplementation((url: string) => { + if (url.includes('loadCodeAssist')) { + forbiddenLoadCalls += 1 + return Promise.resolve( + forbiddenLoadCalls === 1 + ? response({ error: 'Forbidden' }, 403) + : response({ cloudaicompanionProject: 'project-123' }) + ) + } + return Promise.resolve(response({ groups: quotaGroups })) + }) + getTokenMock.mockReset() + getTokenMock + .mockResolvedValueOnce(token('old-forbidden-token')) + .mockResolvedValueOnce(token('new-forbidden-token')) + const forbidden = await fetchAntigravityRateLimits() + expect(forbidden.status).toBe('ok') + expect(getTokenMock).toHaveBeenNthCalledWith(2, { forceRefresh: true }) + expect( + netFetchMock.mock.calls.some( + (call) => + (call[1] as RequestInit).headers && + JSON.stringify(call[1]).includes('new-forbidden-token') + ) + ).toBe(true) + }) + + it('reports missing credentials as unavailable', async () => { + getTokenMock.mockRejectedValue( + new MockAntigravityAuthError('Antigravity credentials not found', 'missing-credentials') + ) + + const result = await fetchAntigravityRateLimits() + + expect(result.status).toBe('unavailable') + expect(result.usageMetadata?.failureKind).toBe('missing-credentials') + expect(netFetchMock).not.toHaveBeenCalled() + }) +}) 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..6d5c1ae32ac5 --- /dev/null +++ b/src/main/rate-limits/antigravity-usage-fetcher.ts @@ -0,0 +1,320 @@ +import { net } from 'electron' +import type { + ProviderRateLimits, + RateLimitBucket, + RateLimitWindow, + UsageRateLimitFailureKind, + UsageRateLimitMetadata +} from '../../shared/rate-limit-types' +import { + AntigravityAuthError, + getAntigravityAccessToken, + invalidateAntigravityAccessToken, + type AntigravityAccessToken +} from './antigravity-auth' +import { + parseAntigravityQuotaBuckets, + type AntigravityQuotaBucket +} from './antigravity-quota-parser' + +export { parseAntigravityQuotaBuckets } from './antigravity-quota-parser' +export type { AntigravityQuotaBucket } from './antigravity-quota-parser' + +const API_TIMEOUT_MS = 10_000 +const API_BASE_URL = 'https://cloudcode-pa.googleapis.com' +const LOAD_CODE_ASSIST_URL = `${API_BASE_URL}/v1internal:loadCodeAssist` +const QUOTA_SUMMARY_URL = `${API_BASE_URL}/v1internal:retrieveUserQuotaSummary` +const LEGACY_QUOTA_URL = `${API_BASE_URL}/v1internal:retrieveUserQuota` +const ANTIGRAVITY_PLATFORM = 'PLATFORM_UNSPECIFIED' +const API_CLIENT_METADATA = { + ideType: 'ANTIGRAVITY', + platform: ANTIGRAVITY_PLATFORM, + pluginType: 'GEMINI' +} as const + +type AntigravityApiErrorOptions = { + status?: number | null + failureKind?: UsageRateLimitFailureKind + retryAtMs?: number +} + +class AntigravityApiError extends Error { + readonly status: number | null + readonly failureKind: UsageRateLimitFailureKind + readonly retryAtMs: number | undefined + + constructor(message: string, options: AntigravityApiErrorOptions = {}) { + super(message) + this.name = 'AntigravityApiError' + this.status = options.status ?? null + this.failureKind = options.failureKind ?? 'unknown' + this.retryAtMs = options.retryAtMs + } +} + +export type AntigravityFetcherOptions = { + baseHomeDir?: string + signal?: AbortSignal +} + +export async function fetchAntigravityRateLimits( + options: AntigravityFetcherOptions = {} +): Promise { + const { signal, baseHomeDir } = options + try { + const token = await getAntigravityAccessToken({ baseHomeDir, signal }) + return await fetchWithToken(token, { baseHomeDir, signal }) + } catch (error) { + if (isAbortError(error)) { + throw error + } + return errorResult(error, null) + } +} + +async function fetchWithToken( + token: AntigravityAccessToken, + options: AntigravityFetcherOptions +): Promise { + let currentToken = token + for (let attempt = 0; attempt < 2; attempt += 1) { + try { + const projectId = await loadProjectId(currentToken.accessToken, options.signal) + const quota = await loadQuota(projectId, currentToken.accessToken, options.signal) + return quotaResult(quota, currentToken.credentialSource) + } catch (error) { + if (isAbortError(error)) { + throw error + } + if (attempt === 0 && isUnauthorizedApiError(error)) { + invalidateAntigravityAccessToken(currentToken.sourceKey) + try { + currentToken = await getAntigravityAccessToken({ + baseHomeDir: options.baseHomeDir, + forceRefresh: true, + signal: options.signal + }) + } catch (refreshError) { + return errorResult(refreshError, currentToken.credentialSource) + } + continue + } + return errorResult(error, currentToken.credentialSource) + } + } + return errorResult( + new AntigravityApiError('Antigravity quota request failed'), + currentToken.credentialSource + ) +} + +async function loadProjectId(accessToken: string, signal?: AbortSignal): Promise { + const data = await postJson( + LOAD_CODE_ASSIST_URL, + { + metadata: API_CLIENT_METADATA + }, + accessToken, + signal + ) + if ( + typeof data !== 'object' || + data === null || + !('cloudaicompanionProject' in data) || + typeof (data as { cloudaicompanionProject?: unknown }).cloudaicompanionProject !== 'string' || + !(data as { cloudaicompanionProject: string }).cloudaicompanionProject.trim() + ) { + throw new AntigravityApiError('Antigravity project ID is missing from the API response', { + failureKind: 'parse' + }) + } + return (data as { cloudaicompanionProject: string }).cloudaicompanionProject +} + +async function loadQuota( + projectId: string, + accessToken: string, + signal?: AbortSignal +): Promise { + let summaryError: AntigravityApiError | null = null + try { + const summary = await postJson(QUOTA_SUMMARY_URL, { project: projectId }, accessToken, signal) + return parseAntigravityQuotaBuckets(summary) + } catch (error) { + if (isAbortError(error)) { + throw error + } + if (!(error instanceof AntigravityApiError) || !isUnsupportedEndpoint(error.status)) { + throw error + } + summaryError = error + } + + // Why: only an explicitly unsupported summary endpoint may use the legacy + // endpoint. Auth, network, rate-limit, and empty-success responses must not + // fan out into a second request or silently change quota identity. + const legacy = await postJson(LEGACY_QUOTA_URL, { project: projectId }, accessToken, signal) + const legacyBuckets = parseAntigravityQuotaBuckets(legacy) + if (legacyBuckets.length === 0) { + throw new AntigravityApiError( + 'Antigravity quota endpoint returned no recognized quota buckets', + { status: summaryError?.status ?? null, failureKind: 'usage-unavailable' } + ) + } + return legacyBuckets +} + +function quotaResult( + buckets: AntigravityQuotaBucket[], + credentialSource: string +): ProviderRateLimits { + if (buckets.length === 0) { + throw new AntigravityApiError( + 'Antigravity quota response contained no recognized quota buckets', + { failureKind: 'usage-unavailable' } + ) + } + const session = buckets.find((entry) => entry.id === 'gemini-5h')?.bucket ?? null + const weekly = buckets.find((entry) => entry.id === 'gemini-weekly')?.bucket ?? null + return { + provider: 'antigravity', + session: session ? withoutName(session) : null, + weekly: weekly ? withoutName(weekly) : null, + buckets: buckets.map((entry) => entry.bucket), + updatedAt: Date.now(), + error: null, + status: 'ok', + usageMetadata: { + source: 'oauth', + credentialSource + } + } +} + +function withoutName(bucket: RateLimitBucket): RateLimitWindow { + const { name: _name, ...window } = bucket + return window +} + +async function postJson( + url: string, + body: unknown, + accessToken: string, + signal?: AbortSignal +): Promise { + const requestSignal = signal + ? AbortSignal.any([signal, AbortSignal.timeout(API_TIMEOUT_MS)]) + : AbortSignal.timeout(API_TIMEOUT_MS) + let response: Response + try { + response = await net.fetch(url, { + method: 'POST', + headers: antigravityHeaders(accessToken), + body: JSON.stringify(body), + signal: requestSignal + }) + } catch (error) { + if (isAbortError(error)) { + throw error + } + throw new AntigravityApiError('Antigravity API request failed', { failureKind: 'network' }) + } + if (!response.ok) { + throw new AntigravityApiError(`Antigravity API request failed (HTTP ${response.status})`, { + status: response.status, + failureKind: failureKindForStatus(response.status), + retryAtMs: getRetryAtMs(response) + }) + } + try { + return await response.json() + } catch { + throw new AntigravityApiError('Antigravity API response is invalid', { failureKind: 'parse' }) + } +} + +function antigravityHeaders(accessToken: string): Record { + return { + Authorization: `Bearer ${accessToken}`, + 'Content-Type': 'application/json', + 'User-Agent': `antigravity/1.1.9 ${process.platform}/${process.arch}`, + 'X-Goog-Api-Client': 'google-cloud-sdk vscode_cloudshelleditor/0.1', + 'Client-Metadata': JSON.stringify(API_CLIENT_METADATA) + } +} + +function failureKindForStatus(status: number): UsageRateLimitFailureKind { + if (status === 401) { + return 'stale-token' + } + if (status === 403) { + return 'missing-scope' + } + if (status === 429) { + return 'rate-limited' + } + if (status >= 500) { + return 'server' + } + return 'unknown' +} + +function getRetryAtMs(response: Response): number | undefined { + const retryAfter = response.headers?.get('retry-after') + if (!retryAfter) { + return undefined + } + const seconds = Number(retryAfter) + if (Number.isFinite(seconds)) { + return Date.now() + Math.max(0, seconds) * 1000 + } + const date = Date.parse(retryAfter) + return Number.isFinite(date) ? date : undefined +} + +function isUnsupportedEndpoint(status: number | null): boolean { + return status === 404 || status === 405 || status === 501 +} + +function isUnauthorizedApiError(error: unknown): error is AntigravityApiError { + return error instanceof AntigravityApiError && (error.status === 401 || error.status === 403) +} + +function errorResult(error: unknown, credentialSource: string | null): ProviderRateLimits { + const metadata: UsageRateLimitMetadata = { + source: 'oauth', + ...(credentialSource ? { credentialSource } : {}) + } + if (error instanceof AntigravityAuthError) { + metadata.failureKind = error.failureKind + } else if (error instanceof AntigravityApiError) { + metadata.failureKind = error.failureKind + if (error.retryAtMs !== undefined) { + metadata.retryAtMs = error.retryAtMs + } + } else { + metadata.failureKind = 'unknown' + } + const status = + error instanceof AntigravityAuthError || error instanceof AntigravityApiError + ? error.failureKind === 'missing-credentials' + ? 'unavailable' + : 'error' + : 'error' + return { + provider: 'antigravity', + session: null, + weekly: null, + buckets: [], + updatedAt: Date.now(), + error: error instanceof Error ? error.message : 'Antigravity quota request failed', + status, + usageMetadata: metadata + } +} + +function isAbortError(error: unknown): boolean { + return Boolean( + error && typeof error === 'object' && 'name' in error && error.name === 'AbortError' + ) +} 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..08b9d0a468f5 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,9 @@ 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 () => + unavailableProvider('antigravity') + ) 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 +99,7 @@ export function mockFreshBackgroundProviderFetches(): void { export function resetRateLimitProviderMocks(): void { vi.clearAllMocks() vi.mocked(fetchGeminiRateLimits).mockResolvedValue(okProvider('gemini', 0, Date.now())) + vi.mocked(fetchAntigravityRateLimits).mockResolvedValue(unavailableProvider('antigravity')) 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 a7a32a24b7e3..f8334f4d5525 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-inactive-account-previews.test.ts b/src/main/rate-limits/service-inactive-account-previews.test.ts index 53338098074d..3c23c877d723 100644 --- a/src/main/rate-limits/service-inactive-account-previews.test.ts +++ b/src/main/rate-limits/service-inactive-account-previews.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-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..800becc9c8c0 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' @@ -32,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() })) @@ -377,6 +382,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 +403,8 @@ 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,6 +423,8 @@ 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) }) 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 a8db8ec6db6c..05987059b9ad 100644 --- a/src/main/rate-limits/service.ts +++ b/src/main/rate-limits/service.ts @@ -20,6 +20,7 @@ import { type NormalizedClaudeAccountSelectionTarget } from '../claude-accounts/runtime-selection' import { fetchGeminiRateLimits } from './gemini-usage-fetcher' +import { fetchAntigravityRateLimits } from './antigravity-usage-fetcher' import { fetchKimiRateLimits } from './kimi-fetcher' import type { KimiHomeResolution } from '../kimi/kimi-runtime-home' import { fetchGrokRateLimits } from './grok-fetcher' @@ -1655,38 +1656,46 @@ 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 - }), - missingWslCodexHome ?? - fetchCodexRateLimits({ - codexHomePath, - allowPtyFallback: this.shouldAllowCodexPtyFallback(), + 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 }), - 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 - }) - ]) + missingWslCodexHome ?? + fetchCodexRateLimits({ + codexHomePath, + allowPtyFallback: this.shouldAllowCodexPtyFallback(), + signal + }), + fetchGeminiRateLimits(geminiCliOAuthEnabled), + 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 @@ -1731,11 +1740,20 @@ export class RateLimitService { status: 'error' } satisfies ProviderRateLimits) - // Why: Antigravity shares Gemini credentials today; mirror the Gemini snapshot so its status-bar UI gets a real lifecycle instead of null. - const antigravity: ProviderRateLimits = { - ...gemini, - provider: 'antigravity' - } + 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/preload/api/dashboard-api.ts b/src/preload/api/dashboard-api.ts index aa8d6dabd8ad..115b651457d5 100644 --- a/src/preload/api/dashboard-api.ts +++ b/src/preload/api/dashboard-api.ts @@ -2,7 +2,8 @@ import type { DashboardRevealAgentArgs, DashboardSleepWorkspaceArgs, DashboardSnapshot, - DashboardSpawnAgentArgs + DashboardSpawnAgentArgs, + DashboardStopAgentArgs } from '../../shared/dashboard-snapshot' import type { TerminalPreviewConnectResult, @@ -17,6 +18,7 @@ export type DashboardApi = { onSnapshotRequested: (callback: () => void) => () => void onRevealAgent: (callback: (args: DashboardRevealAgentArgs) => void) => () => void onAckAgent: (callback: (paneKey: string) => void) => () => void + onStopAgent: (callback: (args: DashboardStopAgentArgs) => void) => () => void onSpawnAgent: (callback: (args: DashboardSpawnAgentArgs) => void) => () => void onSleepWorkspace: (callback: (args: DashboardSleepWorkspaceArgs) => void) => () => void requestSnapshot: () => Promise @@ -24,6 +26,7 @@ export type DashboardApi = { onViewRequested: (callback: (view: 'board' | 'map') => void) => () => void revealAgent: (args: DashboardRevealAgentArgs) => Promise ackAgent: (paneKey: string) => Promise + stopAgent: (args: DashboardStopAgentArgs) => Promise spawnAgent: (args: DashboardSpawnAgentArgs) => Promise sleepWorkspace: (args: DashboardSleepWorkspaceArgs) => Promise } diff --git a/src/preload/index.ts b/src/preload/index.ts index 9be513fa9809..f2abd5ee3fb7 100644 --- a/src/preload/index.ts +++ b/src/preload/index.ts @@ -9,7 +9,8 @@ import type { DashboardRevealAgentArgs, DashboardSleepWorkspaceArgs, DashboardSnapshot, - DashboardSpawnAgentArgs + DashboardSpawnAgentArgs, + DashboardStopAgentArgs } from '../shared/dashboard-snapshot' import type { TerminalPreviewConnectResult, @@ -2391,6 +2392,12 @@ const api = { ipcRenderer.on('ui:ackDashboardAgent', listener) return () => ipcRenderer.removeListener('ui:ackDashboardAgent', listener) }, + onStopAgent: (callback: (args: DashboardStopAgentArgs) => void): (() => void) => { + const listener = (_event: Electron.IpcRendererEvent, args: DashboardStopAgentArgs): void => + callback(args) + ipcRenderer.on('ui:stopDashboardAgent', listener) + return () => ipcRenderer.removeListener('ui:stopDashboardAgent', listener) + }, onSpawnAgent: (callback: (args: DashboardSpawnAgentArgs) => void): (() => void) => { const listener = (_event: Electron.IpcRendererEvent, args: DashboardSpawnAgentArgs): void => callback(args) @@ -2424,6 +2431,8 @@ const api = { ipcRenderer.invoke('dashboardPopout:revealAgent', args), ackAgent: (paneKey: string): Promise => ipcRenderer.invoke('dashboardPopout:ackAgent', { paneKey }), + stopAgent: (args: DashboardStopAgentArgs): Promise => + ipcRenderer.invoke('dashboardPopout:stopAgent', args), spawnAgent: (args: DashboardSpawnAgentArgs): Promise => ipcRenderer.invoke('dashboardPopout:spawnAgent', args), sleepWorkspace: (args: DashboardSleepWorkspaceArgs): Promise => diff --git a/src/renderer/src/components/dashboard-popout/AgentKanbanBoard.test.tsx b/src/renderer/src/components/dashboard-popout/AgentKanbanBoard.test.tsx index fc3389d0dbcf..2985ea8356d3 100644 --- a/src/renderer/src/components/dashboard-popout/AgentKanbanBoard.test.tsx +++ b/src/renderer/src/components/dashboard-popout/AgentKanbanBoard.test.tsx @@ -24,23 +24,28 @@ vi.mock('./AgentKanbanCard', () => ({ card, repoIcon, now, - onOpenTerminal + onOpenTerminal, + onStop }: { card: DashboardCard repoIcon?: RepoIcon | null now: number onOpenTerminal: (card: DashboardCard) => void + onStop: (card: DashboardCard) => void }) => ( -
onOpenTerminal(card)} - > - {card.worktreeName} -
+ <> +
onOpenTerminal(card)} + > + {card.worktreeName} +
+ + ) }, (previous, next) => previous.onOpenTerminal === next.onOpenTerminal && + previous.onStop === next.onStop && sameCard(previous.card, next.card) && sameRepoIcon(previous.repoIcon, next.repoIcon) && (displayTimestamp(previous.card) <= 0 || diff --git a/src/renderer/src/components/dashboard-popout/AgentKanbanCardStopControl.tsx b/src/renderer/src/components/dashboard-popout/AgentKanbanCardStopControl.tsx new file mode 100644 index 000000000000..8e5dfa3911de --- /dev/null +++ b/src/renderer/src/components/dashboard-popout/AgentKanbanCardStopControl.tsx @@ -0,0 +1,111 @@ +import { useCallback, useEffect } from 'react' +import { X } from 'lucide-react' +import { cn } from '@/lib/utils' +import { translate } from '@/i18n/i18n' +import type { DashboardCard } from '../../../../shared/dashboard-snapshot' + +/** How long the inline "Stop agent?" confirmation stays armed before reverting. */ +const CONFIRM_TIMEOUT_MS = 3_000 + +/** Idle and done agents confirm inline; anything still running routes through + * the board's modal, mirroring Cmd+W in the main window (idle panes close + * outright, running ones prompt). */ +export function stopNeedsDialog(card: DashboardCard): boolean { + return card.dotState !== 'idle' && card.dotState !== 'done' +} + +type AgentKanbanCardStopControlProps = { + card: DashboardCard + /** Armed state lives on the card so moving the pointer anywhere off the card + * — not merely off this small control — cancels a pending confirmation. */ + confirming: boolean + onArm: () => void + onDisarm: () => void + /** Fires once the stop is committed — the board decides whether that means + * stopping now or opening the confirmation dialog. */ + onStop: (card: DashboardCard) => void +} + +/** + * Stop affordance for one card, overlaid on the top-right corner where the + * state dot sits (the dot fades out on hover so they never collide). It is a + * SIBLING of the card's button — nesting a button inside one is invalid HTML — + * so clicking it can never also open the terminal. + */ +export function AgentKanbanCardStopControl({ + card, + confirming, + onArm, + onDisarm, + onStop +}: AgentKanbanCardStopControlProps): React.JSX.Element { + // Why: an armed confirmation the user walked away from must not stay armed + // for the next visit to this card. + useEffect(() => { + if (!confirming) { + return + } + const timer = setTimeout(onDisarm, CONFIRM_TIMEOUT_MS) + return () => clearTimeout(timer) + }, [confirming, onDisarm]) + + const handleStopClick = useCallback(() => { + if (confirming || stopNeedsDialog(card)) { + onDisarm() + onStop(card) + return + } + onArm() + }, [card, confirming, onArm, onDisarm, onStop]) + + // Matches the card's p-2.5 padding so the control lands exactly on the dot. + const anchor = 'absolute right-2.5 top-2.5 z-10 flex items-center' + + if (confirming) { + return ( + + + + ) + } + + return ( + + + + ) +} diff --git a/src/renderer/src/components/dashboard-popout/AgentKanbanColumn.tsx b/src/renderer/src/components/dashboard-popout/AgentKanbanColumn.tsx new file mode 100644 index 000000000000..5dbc752e8468 --- /dev/null +++ b/src/renderer/src/components/dashboard-popout/AgentKanbanColumn.tsx @@ -0,0 +1,66 @@ +import type { DashboardBucket, DashboardCard } from '../../../../shared/dashboard-snapshot' +import type { RepoIcon } from '../../../../shared/repo-icon' +import { translate } from '@/i18n/i18n' +import { AgentKanbanCard } from './AgentKanbanCard' + +function bucketLabel(bucket: DashboardBucket): string { + switch (bucket) { + case 'attention': + return translate('dashboardPopout.bucket.attention', 'Needs You') + case 'working': + return translate('dashboardPopout.bucket.working', 'Working') + case 'done': + return translate('dashboardPopout.bucket.done', 'Done') + case 'idle': + return translate('dashboardPopout.bucket.idle', 'Idle') + } +} + +export function AgentKanbanColumn({ + bucket, + cards, + repoIconsByRepoId, + now, + onOpenTerminal, + onStop +}: { + bucket: DashboardBucket + cards: DashboardCard[] + repoIconsByRepoId: Record | undefined + now: number + onOpenTerminal: (card: DashboardCard) => void + onStop: (card: DashboardCard) => void +}): React.JSX.Element { + return ( + // Why: attention no longer tints the whole column — the cards inside carry + // their own state color, so a column border would double-signal it. +
+
+ + {bucketLabel(bucket)} + + + {cards.length} + +
+
+ {cards.length === 0 ? ( +

+ {translate('dashboardPopout.bucket.empty', 'None')} +

+ ) : ( + cards.map((card) => ( + + )) + )} +
+
+ ) +} diff --git a/src/renderer/src/components/dashboard-popout/AgentStopConfirmDialog.tsx b/src/renderer/src/components/dashboard-popout/AgentStopConfirmDialog.tsx new file mode 100644 index 000000000000..9928b0d81d62 --- /dev/null +++ b/src/renderer/src/components/dashboard-popout/AgentStopConfirmDialog.tsx @@ -0,0 +1,69 @@ +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle +} from '@/components/ui/dialog' +import { Button } from '@/components/ui/button' +import { translate } from '@/i18n/i18n' +import type { DashboardCard } from '../../../../shared/dashboard-snapshot' + +type AgentStopConfirmDialogProps = { + /** The agent awaiting confirmation; null renders the dialog closed. */ + card: DashboardCard | null + onCancel: () => void + onConfirm: () => void +} + +/** + * Confirms stopping an agent that is still running. Reuses CloseTerminalDialog's + * copy keys verbatim so the board and Cmd+W say the same thing about the same + * action. The "don't ask again" opt-out is deliberately omitted: it belongs to + * the terminal's own close flow, not to a board that stops agents by the card. + */ +export function AgentStopConfirmDialog({ + card, + onCancel, + onConfirm +}: AgentStopConfirmDialogProps): React.JSX.Element { + return ( + { + if (!isOpen) { + onCancel() + } + }} + > + + + + {translate( + 'auto.components.terminal.pane.CloseTerminalDialog.stop_agent_title', + 'Stop this agent?' + )} + + + {translate( + 'auto.components.terminal.pane.CloseTerminalDialog.stop_agent_description', + "Closing this terminal will stop the agent's current work." + )} + + + + + + + + + ) +} diff --git a/src/renderer/src/components/dashboard/useDashboardPopoutBridge.ts b/src/renderer/src/components/dashboard/useDashboardPopoutBridge.ts index 3d58b9100c83..207e218618da 100644 --- a/src/renderer/src/components/dashboard/useDashboardPopoutBridge.ts +++ b/src/renderer/src/components/dashboard/useDashboardPopoutBridge.ts @@ -3,6 +3,7 @@ import { useAppStore, type AppState } from '@/store' import { activateTabAndFocusPane } from '@/lib/activate-tab-and-focus-pane' import { runSleepWorktree } from '../sidebar/sleep-worktree-flow' import type { RepoIcon } from '../../../../shared/repo-icon' +import { stopDashboardAgent } from '@/lib/stop-dashboard-agent' import { buildDashboardSnapshot, type DashboardSnapshotState } from './build-dashboard-snapshot' import { launchDashboardAgent } from './launch-dashboard-agent' @@ -149,6 +150,17 @@ export function useDashboardPopoutBridge(enabled: boolean): void { }) }, [enabled]) + // Stopping an agent from the board runs its terminal teardown here, where the + // store lives. Same ?. shield as the ack listener above. + useEffect(() => { + if (!enabled) { + return + } + return window.api.dashboard.onStopAgent?.((args) => { + stopDashboardAgent(args) + }) + }, [enabled]) + useEffect(() => { if (!enabled) { return diff --git a/src/renderer/src/components/status-bar/StatusBar.tsx b/src/renderer/src/components/status-bar/StatusBar.tsx index 0243c5f8deec..e5e4690428d0 100644 --- a/src/renderer/src/components/status-bar/StatusBar.tsx +++ b/src/renderer/src/components/status-bar/StatusBar.tsx @@ -2070,8 +2070,7 @@ 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 reads its own CLI keyring credentials; the checked item + detected CLI is its durable signal. const antigravityUsageConfigured = statusBarItems.includes('antigravity') && isStatusBarItemAvailable('antigravity', detectedAgentIds) diff --git a/src/renderer/src/components/status-bar/UsageRosterPanel.tsx b/src/renderer/src/components/status-bar/UsageRosterPanel.tsx index feb066fadca4..77db86acb82d 100644 --- a/src/renderer/src/components/status-bar/UsageRosterPanel.tsx +++ b/src/renderer/src/components/status-bar/UsageRosterPanel.tsx @@ -35,7 +35,7 @@ function providerMaxUsed(sections: UsageSection[]): number { : 0 } -// Buckets (Gemini Flash/Pro) keep their model name; windows use their duration. +// Buckets keep their provider-specific name; windows use their duration. function shortLabel( p: ProviderRateLimits, section: UsageSection, 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..9863da004425 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, detected Antigravity item as its durable 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 the CLI is detected', () => { 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 visible while credentials are unavailable', () => { 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 credentials not found' }), usageSettings({ antigravityUsageConfigured: true }) ) - ).toBe(null) + ).toMatchObject({ provider: 'antigravity', status: 'unavailable' }) }) }) @@ -499,9 +485,7 @@ describe('isUsageEmptyState', () => { ).toBe(false) }) - it('still shows the setup CTA when Antigravity is checked but Gemini OAuth is off', () => { - // Why: the default-on Antigravity item is not configured usage on its own; - // it must not hide the teaching CTA from users who set nothing up. + it('does not show the setup CTA when the detected Antigravity CLI is checked', () => { 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..7a778d462f87 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 credentials live in the CLI keyring; StatusBar sets this + // after PATH detection, so its checked item and detected CLI are durable. antigravityUsageConfigured: boolean // Why: MiniMax/Grok sign-in live on disk, not in settings; main sets these each poll. minimaxCookieConfigured: boolean @@ -73,9 +70,8 @@ export function hasUsageProviderSettings( (settings?.codexManagedAccounts?.length ?? 0) > 0 || (settings?.claudeManagedAccounts?.length ?? 0) > 0 || settings?.geminiCliOAuthEnabled === true || + settings?.antigravityUsageConfigured === true || Boolean(settings?.opencodeSessionCookie?.trim()) || - // Antigravity's durable signal requires geminiCliOAuthEnabled, so it is - // already covered by the gemini term above. 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..4f906b11658e 100644 --- a/src/renderer/src/components/status-bar/tooltip.test.ts +++ b/src/renderer/src/components/status-bar/tooltip.test.ts @@ -326,6 +326,26 @@ describe('getWindowSections', () => { ]) }) + it('does not repeat a weekly bucket in the generic weekly section', () => { + const weekly = { + usedPercent: 25, + windowMinutes: 10080, + resetsAt: 1234, + resetDescription: null + } + const p: ProviderRateLimits = { + provider: 'antigravity', + session: null, + weekly, + buckets: [{ name: 'Antigravity weekly', ...weekly }], + updatedAt: Date.now(), + error: null, + status: 'ok' + } + + expect(getWindowSections(p)).toEqual([{ label: 'Antigravity weekly', window: p.buckets![0] }]) + }) + 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..67ce33b2c351 100644 --- a/src/renderer/src/components/status-bar/tooltip.tsx +++ b/src/renderer/src/components/status-bar/tooltip.tsx @@ -143,12 +143,24 @@ 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 })) + const weeklyIsAlreadyBucket = p.weekly + ? bucketSections.some( + (section) => + section.window.windowMinutes === p.weekly?.windowMinutes && + section.window.usedPercent === p.weekly?.usedPercent && + section.window.resetsAt === p.weekly?.resetsAt + ) + : false return [ ...bucketSections, - { - label: translate('auto.components.status.bar.tooltip.252c096536', 'Weekly'), - window: p.weekly - } + ...(weeklyIsAlreadyBucket + ? [] + : [ + { + label: translate('auto.components.status.bar.tooltip.252c096536', 'Weekly'), + window: p.weekly + } + ]) ] } const sections: { label: string; window: RateLimitWindow | null }[] = [ diff --git a/src/renderer/src/components/status-bar/usage-provider-settings-target.test.ts b/src/renderer/src/components/status-bar/usage-provider-settings-target.test.ts index cbcfcda04929..2f20c8a41f6b 100644 --- a/src/renderer/src/components/status-bar/usage-provider-settings-target.test.ts +++ b/src/renderer/src/components/status-bar/usage-provider-settings-target.test.ts @@ -6,7 +6,7 @@ describe('getUsageProviderAccountsSectionId', () => { expect(getUsageProviderAccountsSectionId('claude')).toBe('accounts-claude') expect(getUsageProviderAccountsSectionId('codex')).toBe('accounts-codex') expect(getUsageProviderAccountsSectionId('gemini')).toBe('accounts-gemini') - expect(getUsageProviderAccountsSectionId('antigravity')).toBe('accounts-gemini') + expect(getUsageProviderAccountsSectionId('antigravity')).toBeNull() expect(getUsageProviderAccountsSectionId('opencode-go')).toBe('accounts-opencode-go') expect(getUsageProviderAccountsSectionId('minimax')).toBe('accounts-minimax') expect(getUsageProviderAccountsSectionId('grok')).toBe('accounts-grok') diff --git a/src/renderer/src/components/status-bar/usage-provider-settings-target.ts b/src/renderer/src/components/status-bar/usage-provider-settings-target.ts index 59ea71042490..daedaff19bdd 100644 --- a/src/renderer/src/components/status-bar/usage-provider-settings-target.ts +++ b/src/renderer/src/components/status-bar/usage-provider-settings-target.ts @@ -9,9 +9,10 @@ export function getUsageProviderAccountsSectionId( case 'codex': return 'accounts-codex' case 'gemini': - case 'antigravity': - // Why: Antigravity usage currently shares Gemini's OAuth configuration. return 'accounts-gemini' + case 'antigravity': + // Why: agy owns its browser/keyring sign-in; Orca has no sign-in section. + return null case 'opencode-go': return 'accounts-opencode-go' case 'minimax': diff --git a/src/renderer/src/i18n/locales/en.json b/src/renderer/src/i18n/locales/en.json index a349b2c68120..1946f0b1181f 100644 --- a/src/renderer/src/i18n/locales/en.json +++ b/src/renderer/src/i18n/locales/en.json @@ -15672,6 +15672,11 @@ "settingsTooltip": "Board settings", "card": { "you": "You", + "stop": { + "action": "Stop agent", + "confirm": "Stop agent?", + "tooltip": "Stop agent — kills the process and closes its terminal" + }, "time": { "justNow": "just now", "minutes": "{{count}}m", diff --git a/src/renderer/src/i18n/locales/es.json b/src/renderer/src/i18n/locales/es.json index b08ae94d09ed..1acee8ee7e28 100644 --- a/src/renderer/src/i18n/locales/es.json +++ b/src/renderer/src/i18n/locales/es.json @@ -14526,6 +14526,11 @@ "total": "{{count}} total", "card": { "you": "You", + "stop": { + "action": "Detener agente", + "confirm": "¿Detener agente?", + "tooltip": "Detener agente: termina el proceso y cierra su terminal" + }, "time": { "justNow": "ahora mismo", "minutes": "{{count}} min", diff --git a/src/renderer/src/i18n/locales/ja.json b/src/renderer/src/i18n/locales/ja.json index 30facdfa6e13..edf01ed4cbac 100644 --- a/src/renderer/src/i18n/locales/ja.json +++ b/src/renderer/src/i18n/locales/ja.json @@ -14526,6 +14526,11 @@ "total": "合計 {{count}} 件", "card": { "you": "あなた", + "stop": { + "action": "エージェントを停止", + "confirm": "エージェントを停止?", + "tooltip": "エージェントを停止 — プロセスを終了してターミナルを閉じます" + }, "time": { "justNow": "たった今", "minutes": "{{count}}分", diff --git a/src/renderer/src/i18n/locales/ko.json b/src/renderer/src/i18n/locales/ko.json index b1a2d8bc63f8..dc15e1d36c15 100644 --- a/src/renderer/src/i18n/locales/ko.json +++ b/src/renderer/src/i18n/locales/ko.json @@ -14537,6 +14537,11 @@ "total": "{{count}} total", "card": { "you": "You", + "stop": { + "action": "에이전트 중지", + "confirm": "에이전트를 중지할까요?", + "tooltip": "에이전트 중지 — 프로세스를 종료하고 터미널을 닫습니다" + }, "time": { "justNow": "방금", "minutes": "{{count}}분", diff --git a/src/renderer/src/i18n/locales/zh.json b/src/renderer/src/i18n/locales/zh.json index 921ca072c1c7..e6318ad07eb8 100644 --- a/src/renderer/src/i18n/locales/zh.json +++ b/src/renderer/src/i18n/locales/zh.json @@ -14546,6 +14546,11 @@ "total": "共 {{count}} 个", "card": { "you": "你", + "stop": { + "action": "停止 Agent", + "confirm": "停止 Agent?", + "tooltip": "停止 Agent — 结束进程并关闭其终端" + }, "time": { "justNow": "刚刚", "minutes": "{{count}} 分钟", diff --git a/src/renderer/src/lib/stop-dashboard-agent.test.ts b/src/renderer/src/lib/stop-dashboard-agent.test.ts new file mode 100644 index 000000000000..6b5d0809abc4 --- /dev/null +++ b/src/renderer/src/lib/stop-dashboard-agent.test.ts @@ -0,0 +1,122 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' +import { stopDashboardAgent } from './stop-dashboard-agent' + +const dropAgentStatus = vi.fn() +const dismissRetainedAgent = vi.fn() +const closeTerminalTab = vi.fn() +const closeWebRuntimeTerminal = vi.fn((_ptyId: string | null | undefined) => false) +const ptyKill = vi.fn() + +let terminalLayoutsByTabId: Record }> = {} + +vi.mock('@/store', () => ({ + useAppStore: { + getState: () => ({ + dropAgentStatus, + dismissRetainedAgent, + terminalLayoutsByTabId + }) + } +})) + +vi.mock('@/components/terminal/terminal-tab-actions', () => ({ + closeTerminalTab: (...args: unknown[]) => closeTerminalTab(...args) +})) + +vi.mock('@/runtime/web-runtime-session', () => ({ + closeWebRuntimeTerminal: (ptyId: string | null | undefined) => closeWebRuntimeTerminal(ptyId) +})) + +function args( + overrides: Partial[0]> = {} +): Parameters[0] { + return { + paneKey: 'tab-1:leaf-1', + worktreeId: 'worktree-1', + tabId: 'tab-1', + leafId: 'leaf-1', + ptyId: 'pty-1', + ...overrides + } +} + +describe('stopDashboardAgent', () => { + beforeEach(() => { + vi.clearAllMocks() + closeWebRuntimeTerminal.mockReturnValue(false) + terminalLayoutsByTabId = { + 'tab-1': { ptyIdsByLeafId: { 'leaf-1': 'pty-1' } } + } + vi.stubGlobal('window', { api: { pty: { kill: ptyKill } } }) + }) + + it('drops the row before tearing the pane down so retention cannot resurrect the card', () => { + const order: string[] = [] + dropAgentStatus.mockImplementation(() => order.push('drop')) + dismissRetainedAgent.mockImplementation(() => order.push('dismiss')) + closeTerminalTab.mockImplementation(() => order.push('close')) + + stopDashboardAgent(args()) + + expect(order).toEqual(['drop', 'dismiss', 'close']) + expect(dropAgentStatus).toHaveBeenCalledWith('tab-1:leaf-1') + expect(dismissRetainedAgent).toHaveBeenCalledWith('tab-1:leaf-1') + }) + + it('closes the whole tab when the agent is its only pane', () => { + stopDashboardAgent(args()) + + expect(closeTerminalTab).toHaveBeenCalledWith( + 'tab-1', + expect.objectContaining({ + reason: 'user', + rejectPinned: true + }) + ) + expect(ptyKill).not.toHaveBeenCalled() + }) + + it('kills only the agent pane when the tab is split', () => { + terminalLayoutsByTabId = { + 'tab-1': { ptyIdsByLeafId: { 'leaf-1': 'pty-1', 'leaf-2': 'pty-2' } } + } + + stopDashboardAgent(args()) + + expect(ptyKill).toHaveBeenCalledWith('pty-1') + expect(closeTerminalTab).not.toHaveBeenCalled() + }) + + it('routes a host-owned runtime pane through the runtime close instead of pty.kill', () => { + terminalLayoutsByTabId = { + 'tab-1': { + ptyIdsByLeafId: { 'leaf-1': 'remote:env/1', 'leaf-2': 'pty-2' } + } + } + closeWebRuntimeTerminal.mockReturnValue(true) + + stopDashboardAgent(args({ ptyId: 'remote:env/1' })) + + expect(closeWebRuntimeTerminal).toHaveBeenCalledWith('remote:env/1') + expect(ptyKill).not.toHaveBeenCalled() + }) + + it('stops the agent anyway when a pinned tab refuses to close', () => { + closeTerminalTab.mockImplementation((_tabId: string, options: { onCancel?: () => void }) => + options.onCancel?.() + ) + + stopDashboardAgent(args()) + + expect(ptyKill).toHaveBeenCalledWith('pty-1') + }) + + it('only drops the row for a retained card whose pane is already gone', () => { + stopDashboardAgent(args({ ptyId: null })) + + expect(dropAgentStatus).toHaveBeenCalledWith('tab-1:leaf-1') + expect(dismissRetainedAgent).toHaveBeenCalledWith('tab-1:leaf-1') + expect(closeTerminalTab).not.toHaveBeenCalled() + expect(ptyKill).not.toHaveBeenCalled() + }) +}) diff --git a/src/renderer/src/lib/stop-dashboard-agent.ts b/src/renderer/src/lib/stop-dashboard-agent.ts new file mode 100644 index 000000000000..313c19ee1248 --- /dev/null +++ b/src/renderer/src/lib/stop-dashboard-agent.ts @@ -0,0 +1,56 @@ +import { useAppStore } from '@/store' +import { closeTerminalTab } from '@/components/terminal/terminal-tab-actions' +import { closeWebRuntimeTerminal } from '@/runtime/web-runtime-session' +import type { DashboardStopAgentArgs } from '../../../shared/dashboard-snapshot' + +/** Kills the PTY behind one agent pane, matching TerminalPane's own split-pane + * teardown order: host-owned runtime terminals first, local/SSH otherwise. */ +function killAgentPane(ptyId: string): void { + if (closeWebRuntimeTerminal(ptyId)) { + return + } + void window.api.pty.kill(ptyId) +} + +function paneCountForTab(tabId: string): number { + const layout = useAppStore.getState().terminalLayoutsByTabId[tabId] + return Object.keys(layout?.ptyIdsByLeafId ?? {}).length +} + +/** + * Stops one agent from the pop-out board: kills its process and takes its row + * off the board. Runs in the MAIN renderer — the pop-out has no store of its + * own — and never activates the worktree, so the board stays put. + */ +export function stopDashboardAgent(args: DashboardStopAgentArgs): void { + const state = useAppStore.getState() + // Why: drop the live row BEFORE the pane dies so dropAgentStatus plants its + // retention suppressor while a live entry still exists; otherwise the + // retention sync re-adds the agent as a retained "done" row and the card + // never actually leaves the board. + state.dropAgentStatus(args.paneKey) + state.dismissRetainedAgent(args.paneKey) + + if (!args.ptyId) { + // Retained card whose pane is already gone — the row was the only thing left. + return + } + + // A split tab hosts other panes the user didn't ask to close, so kill just + // this agent's PTY and leave the tab standing. + if (paneCountForTab(args.tabId) > 1) { + killAgentPane(args.ptyId) + return + } + + const ptyId = args.ptyId + closeTerminalTab(args.tabId, { + reason: 'user', + // Why: the pinned-close confirmation would open a modal in the main window, + // which is behind the board and may be unattended. Stop the agent anyway + // and leave the pinned tab itself alone — silently doing nothing would read + // as a broken button. + rejectPinned: true, + onCancel: () => killAgentPane(ptyId) + }) +} diff --git a/src/shared/dashboard-snapshot.ts b/src/shared/dashboard-snapshot.ts index 16ed5420aa17..9fde6bd488f5 100644 --- a/src/shared/dashboard-snapshot.ts +++ b/src/shared/dashboard-snapshot.ts @@ -217,7 +217,6 @@ export type DashboardRevealAgentArgs = { tabId: string leafId: string | null } - export type DashboardSpawnAgentArgs = { worktreeId: string agent: TuiAgent @@ -228,3 +227,14 @@ export type DashboardSpawnAgentArgs = { export type DashboardSleepWorkspaceArgs = { worktreeId: string } + +/** Payload for stopping an agent from the board: kills its process and drops + * its row. ptyId is null for a retained card whose pane is already gone — + * there is nothing left to kill, so the stop degrades to dropping the row. */ +export type DashboardStopAgentArgs = { + paneKey: string + worktreeId: string + tabId: string + leafId: string | null + ptyId: string | null +}