From 44ca8171d2beb820839c9264f42f474933e7eee0 Mon Sep 17 00:00:00 2001 From: Brent Duarte Date: Wed, 9 Sep 2026 16:27:03 -0700 Subject: [PATCH 1/7] feat(core): share account pool integration primitives with Pi --- packages/core/src/account-manager.test.ts | 106 +++++++++++ packages/core/src/account-manager.ts | 122 +++++++++--- packages/core/src/index.ts | 2 + packages/core/src/persist-account-pool.ts | 171 +++++++++++++++++ packages/core/src/rate-limit-response.ts | 157 ++++++++++++++++ .../opencode/src/plugin/fetch-interceptor.ts | 164 +--------------- .../src/plugin/persist-account-pool.ts | 176 +----------------- 7 files changed, 539 insertions(+), 359 deletions(-) create mode 100644 packages/core/src/persist-account-pool.ts create mode 100644 packages/core/src/rate-limit-response.ts diff --git a/packages/core/src/account-manager.test.ts b/packages/core/src/account-manager.test.ts index 3146ed90..8b179d6d 100644 --- a/packages/core/src/account-manager.test.ts +++ b/packages/core/src/account-manager.test.ts @@ -42,6 +42,112 @@ const stored: AccountStorageV4 = { } describe('core AccountManager', () => { + it.each([ + 'sticky', + 'hybrid', + 'round-robin', + ] as const)('applies PID offset to %s before selection, only once', (strategy) => { + const memory = createStore(stored) + const manager = new AccountManager(undefined, structuredClone(stored), { + store: memory.store, + pid: 1, + now: () => 1_000, + }) + expect( + manager.getCurrentOrNextForFamily( + 'gemini', + 'gemini-3.8-flash', + strategy, + 'antigravity', + true, + )?.index, + ).toBe(1) + expect( + manager.getCurrentOrNextForFamily( + 'gemini', + 'gemini-3.8-flash', + strategy, + 'antigravity', + true, + )?.index, + ).toBe(strategy === 'round-robin' ? 0 : 1) + }) + + it('reconciles durable metadata while preserving routing, credentials and per-token failure counters', () => { + const memory = createStore(stored) + const manager = new AccountManager( + { + type: 'oauth', + refresh: 'r1|p1', + access: 'access', + expires: 9999999999999, + }, + structuredClone(stored), + { store: memory.store }, + ) + const first = manager.getCurrentOrNextForFamily( + 'gemini', + null, + 'round-robin', + )! + first.consecutiveFailures = 2 + const next = structuredClone(stored) + next.accounts[0]!.enabled = false + next.accounts[1]!.rateLimitResetTimes = { claude: Date.now() + 60_000 } + manager.reconcileStorage(next) + expect(manager.getAccounts()[0]).toMatchObject({ + access: 'access', + consecutiveFailures: 2, + enabled: false, + }) + expect( + manager.getCurrentOrNextForFamily('gemini', null, 'round-robin')?.index, + ).toBe(1) + expect( + manager.getAccounts()[1]?.rateLimitResetTimes.claude, + ).toBeGreaterThan(Date.now()) + expect(memory.mergedSaves()).toBe(0) + next.accounts[0]!.refreshToken = 'rotated' + manager.reconcileStorage(next) + expect(manager.getAccounts()[0]?.access).toBeUndefined() + expect(manager.getAccounts()[0]?.consecutiveFailures).toBeUndefined() + }) + + it('reconciles removals without resurrecting accounts or transferring session pins', () => { + const memory = createStore(stored) + const manager = new AccountManager(undefined, structuredClone(stored), { + store: memory.store, + }) + const identity = { id: 'session' } + manager.getCurrentOrNextForFamily( + 'gemini', + null, + 'sticky', + 'antigravity', + false, + 100, + 60_000, + identity, + ) + manager.reconcileStorage({ + version: 4, + activeIndex: 0, + accounts: [stored.accounts[1]!], + }) + expect(manager.getTotalAccountCount()).toBe(1) + expect( + manager.getCurrentOrNextForFamily( + 'gemini', + null, + 'sticky', + 'antigravity', + false, + 100, + 60_000, + identity, + )?.parts.refreshToken, + ).toBe('r2') + }) it('constructs from stored and fallback auth', () => { const memory = createStore(stored) const manager = new AccountManager( diff --git a/packages/core/src/account-manager.ts b/packages/core/src/account-manager.ts index 9be86b47..190ee534 100644 --- a/packages/core/src/account-manager.ts +++ b/packages/core/src/account-manager.ts @@ -54,6 +54,8 @@ export interface AccountManagerOptions { now?: () => number random?: () => number pid?: number + /** Hosts doing field-level persistence can own fingerprint writes themselves. */ + persistFingerprintUpdates?: boolean onDiagnostic?: (message: string, fields?: Record) => void } @@ -481,7 +483,10 @@ export class AccountManager { } // Persist updated fingerprint versions to disk - if (fingerprintVersionChanged) { + if ( + fingerprintVersionChanged && + options.persistFingerprintUpdates !== false + ) { this.requestSaveToDisk() } @@ -542,6 +547,58 @@ export class AccountManager { return this.getEnabledAccounts().length } + /** Reload durable metadata without resetting this process's routing state. + * Access tokens and failure counters are transient and survive only an exact + * refresh-token match. Removed accounts are never resurrected. + */ + reconcileStorage(stored: AccountStorageV4): void { + const previous = this.accounts + const byToken = new Map( + previous.map((account) => [account.parts.refreshToken, account]), + ) + const fresh = new AccountManager(undefined, stored, { + store: this.store, + now: this.now, + random: this.random, + pid: this.pid, + persistFingerprintUpdates: false, + }) + this.accounts = fresh.accounts.map((account) => { + const old = byToken.get(account.parts.refreshToken) + if (!old) return account + return { + ...account, + access: old.access, + expires: old.expires, + fingerprint: + stored.accounts[account.index]?.fingerprint ?? + old.fingerprint ?? + account.fingerprint, + touchedForQuota: old.touchedForQuota, + consecutiveFailures: old.consecutiveFailures, + lastFailureTime: old.lastFailureTime, + } + }) + const remap = (index: number): number => { + const token = previous[index]?.parts.refreshToken + return this.accounts.findIndex( + (account) => account.parts.refreshToken === token, + ) + } + for (const family of ['claude', 'gemini'] as const) { + this.currentAccountIndexByFamily[family] = remap( + this.currentAccountIndexByFamily[family], + ) + } + for (const state of this.requestSessionStates.values()) { + for (const family of ['claude', 'gemini'] as const) { + state.currentAccountIndexByFamily[family] = remap( + state.currentAccountIndexByFamily[family], + ) + } + } + } + getTotalAccountCount(): number { return this.accounts.length } @@ -811,6 +868,40 @@ export class AccountManager { } } + // PID-based offset for multi-session distribution (opt-in) + // Different sessions (PIDs) will prefer different starting accounts + const offsetApplied = identity + ? this.getRequestSessionState(identity).offsetAppliedByFamily + : this.sessionOffsetApplied + if ( + pidOffsetEnabled && + !offsetApplied[family] && + this.accounts.length > 1 + ) { + const pidOffset = this.pid % this.accounts.length + const activeIndex = this.getActiveIndex(family, identity) + const baseIndex = + activeIndex >= 0 ? activeIndex : this.getCursor(family, identity) + const newIndex = (baseIndex + pidOffset) % this.accounts.length + + this.onDiagnostic?.('Applying PID account offset', { + pid: this.pid, + offset: pidOffset, + family, + fromIndex: baseIndex, + toIndex: newIndex, + }) + + this.setActiveIndex(family, newIndex, identity) + if (strategy === 'round-robin') { + const cursors = identity + ? this.getRequestSessionState(identity).cursorByFamily + : this.cursorByFamily + cursors[family] = newIndex + } + offsetApplied[family] = true + } + if (strategy === 'round-robin') { const next = this.getNextForFamily( family, @@ -888,35 +979,6 @@ export class AccountManager { } } - // Fallback: sticky selection (used when hybrid finds no candidates) - // PID-based offset for multi-session distribution (opt-in) - // Different sessions (PIDs) will prefer different starting accounts - const offsetApplied = identity - ? this.getRequestSessionState(identity).offsetAppliedByFamily - : this.sessionOffsetApplied - if ( - pidOffsetEnabled && - !offsetApplied[family] && - this.accounts.length > 1 - ) { - const pidOffset = this.pid % this.accounts.length - const activeIndex = this.getActiveIndex(family, identity) - const baseIndex = - activeIndex >= 0 ? activeIndex : this.getCursor(family, identity) - const newIndex = (baseIndex + pidOffset) % this.accounts.length - - this.onDiagnostic?.('Applying PID account offset', { - pid: this.pid, - offset: pidOffset, - family, - fromIndex: baseIndex, - toIndex: newIndex, - }) - - this.setActiveIndex(family, newIndex, identity) - offsetApplied[family] = true - } - const current = this.getCurrentAccountForFamily(family, identity) if (current && !excludeIndexes?.has(current.index)) { clearExpiredRateLimits(current, this.now) diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 27e35f57..b1b37b72 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -22,9 +22,11 @@ export * from './fingerprint.ts' export * from './logger.ts' export * from './model-registry.ts' export * from './model-types.ts' +export * from './persist-account-pool.ts' export * from './project.ts' export * from './quota-manager.ts' export * from './quota-types.ts' +export * from './rate-limit-response.ts' export * from './rotation.ts' export * from './transform/index.ts' export * from './version.ts' diff --git a/packages/core/src/persist-account-pool.ts b/packages/core/src/persist-account-pool.ts new file mode 100644 index 00000000..4d1d897b --- /dev/null +++ b/packages/core/src/persist-account-pool.ts @@ -0,0 +1,171 @@ +/** + * Account pool persistence for OAuth flows. + * + * Merges a batch of successful OAuth token-exchange results into the + * persisted pool. All reads + writes happen inside the core + * `mutateAccountStorage` callback so the mutator sees the freshest + * state read while the lock is held — without it, a concurrent add + * would race the read-modify-write and silently disappear. + * + * Two upsert keys are honored, in priority order: + * 1. email — survives refresh-token rotation for the same Google account + * 2. refresh token — handles the no-email case and out-of-band rotations + * + * Destructive (`replaceAll: true`) writes start from an empty v4 inside + * the same locked callback so a stale merge cannot resurrect a removed + * account. + */ + +import { mutateAccountStorage } from './account-storage.ts' +import type { AccountMetadataV3, AccountStorageV4 } from './account-types.ts' +import type { AntigravityTokenExchangeResult } from './antigravity/oauth.ts' +import { parseRefreshParts } from './auth.ts' + +type TokenSuccess = Extract + +function clampInt(value: number, min: number, max: number): number { + if (!Number.isFinite(value)) { + return min + } + return Math.min(max, Math.max(min, Math.floor(value))) +} + +function applyUpserts( + current: AccountStorageV4, + results: TokenSuccess[], + replaceAll: boolean, +): AccountStorageV4 | undefined { + const now = Date.now() + + // For fresh logins, start from empty inside the locked callback so + // a stale merge cannot resurrect a removed account. + const accounts: AccountMetadataV3[] = replaceAll ? [] : [...current.accounts] + + const indexByRefreshToken = new Map() + const indexByEmail = new Map() + for (let i = 0; i < accounts.length; i++) { + const acc = accounts[i] + if (!acc) continue + if (acc.refreshToken) { + indexByRefreshToken.set(acc.refreshToken, i) + } + if (acc.email) { + indexByEmail.set(acc.email, i) + } + } + + for (const result of results) { + const parts = parseRefreshParts(result.refresh) + if (!parts.refreshToken) { + continue + } + + // Email match wins over token match — handles refresh-token rotation + // for the same Google account. + const existingByEmail = result.email + ? indexByEmail.get(result.email) + : undefined + const existingByToken = indexByRefreshToken.get(parts.refreshToken) + const existingIndex = existingByEmail ?? existingByToken + + if (existingIndex === undefined) { + const newIndex = accounts.length + indexByRefreshToken.set(parts.refreshToken, newIndex) + if (result.email) { + indexByEmail.set(result.email, newIndex) + } + accounts.push({ + email: result.email, + label: result.label, + refreshToken: parts.refreshToken, + projectId: parts.projectId, + managedProjectId: parts.managedProjectId, + addedAt: now, + lastUsed: now, + enabled: true, + }) + continue + } + + const existing = accounts[existingIndex] + if (!existing) continue + + const oldToken = existing.refreshToken + accounts[existingIndex] = { + ...existing, + email: result.email ?? existing.email, + label: result.label ?? existing.label, + refreshToken: parts.refreshToken, + projectId: parts.projectId ?? existing.projectId, + managedProjectId: parts.managedProjectId ?? existing.managedProjectId, + lastUsed: now, + } + + if (oldToken !== parts.refreshToken) { + indexByRefreshToken.delete(oldToken) + indexByRefreshToken.set(parts.refreshToken, existingIndex) + } + } + + if (accounts.length === 0) { + return undefined + } + + const activeIndex = replaceAll + ? 0 + : typeof current.activeIndex === 'number' && + Number.isFinite(current.activeIndex) + ? current.activeIndex + : 0 + + const clamped = clampInt(activeIndex, 0, accounts.length - 1) + return { + version: 4, + accounts, + activeIndex: clamped, + activeIndexByFamily: { + claude: clamped, + gemini: clamped, + }, + } +} + +/** + * Merge a batch of successful OAuth results into the persisted pool. + * + * - `replaceAll: true` — start from empty (fresh login) + * - `replaceAll: false` — preserve existing accounts, upsert by email + * then refresh token, bump `lastUsed` + * + * Both branches run their mutator INSIDE the locked callback. The + * `replaceAll` branch seeds the mutator from an empty v4 rather than + * reading the disk state, but the file lock is still required so the + * write is atomic against concurrent writers — a deleted-account merge + * would resurrect a stale account if we wrote without the lock. + */ +export async function persistAccountPoolAtPath( + path: string, + results: TokenSuccess[], + replaceAll: boolean = false, +): Promise { + if (results.length === 0) { + return + } + + const emptyV4 = (): AccountStorageV4 => ({ + version: 4, + accounts: [], + activeIndex: 0, + }) + + if (replaceAll) { + await mutateAccountStorage(path, () => + applyUpserts(emptyV4(), results, true), + ) + return + } + + await mutateAccountStorage(path, (current) => + applyUpserts(current, results, false), + ) +} diff --git a/packages/core/src/rate-limit-response.ts b/packages/core/src/rate-limit-response.ts new file mode 100644 index 00000000..3a7f6a8d --- /dev/null +++ b/packages/core/src/rate-limit-response.ts @@ -0,0 +1,157 @@ +/** Reads `retry-after-ms` / `retry-after` headers, in that order. */ +export function retryAfterMsFromResponse( + response: Response, + defaultRetryMs: number = 60_000, +): number { + const retryAfterMsHeader = response.headers.get('retry-after-ms') + if (retryAfterMsHeader) { + const parsed = Number.parseInt(retryAfterMsHeader, 10) + if (!Number.isNaN(parsed) && parsed > 0) { + return parsed + } + } + + const retryAfterHeader = response.headers.get('retry-after') + if (retryAfterHeader) { + const parsed = Number.parseInt(retryAfterHeader, 10) + if (!Number.isNaN(parsed) && parsed > 0) { + return parsed * 1000 + } + } + + return defaultRetryMs +} + +export interface RateLimitBodyInfo { + retryDelayMs: number | null + message?: string + quotaResetTime?: string + reason?: string +} + +export function extractRateLimitBodyInfo(body: unknown): RateLimitBodyInfo { + if (!body || typeof body !== 'object') return { retryDelayMs: null } + + const error = (body as { error?: unknown }).error + const message = + error && typeof error === 'object' + ? (error as { message?: string }).message + : undefined + + const details = + error && typeof error === 'object' + ? (error as { details?: unknown[] }).details + : undefined + + let reason: string | undefined + if (Array.isArray(details)) { + for (const detail of details) { + if (!detail || typeof detail !== 'object') continue + const type = (detail as { '@type'?: string })['@type'] + if (typeof type === 'string' && type.includes('google.rpc.ErrorInfo')) { + const detailReason = (detail as { reason?: string }).reason + if (typeof detailReason === 'string') { + reason = detailReason + break + } + } + } + + for (const detail of details) { + if (!detail || typeof detail !== 'object') continue + const type = (detail as { '@type'?: string })['@type'] + if (typeof type === 'string' && type.includes('google.rpc.RetryInfo')) { + const retryDelay = (detail as { retryDelay?: string }).retryDelay + if (typeof retryDelay === 'string') { + const retryDelayMs = parseDurationToMs(retryDelay) + if (retryDelayMs !== null) { + return { retryDelayMs, message, reason } + } + } + } + } + + for (const detail of details) { + if (!detail || typeof detail !== 'object') continue + const metadata = (detail as { metadata?: Record }) + .metadata + if (metadata && typeof metadata === 'object') { + const quotaResetDelay = metadata.quotaResetDelay + const quotaResetTime = metadata.quotaResetTimeStamp + if (typeof quotaResetDelay === 'string') { + const quotaResetDelayMs = parseDurationToMs(quotaResetDelay) + if (quotaResetDelayMs !== null) { + return { + retryDelayMs: quotaResetDelayMs, + message, + quotaResetTime, + reason, + } + } + } + } + } + } + + if (message) { + const afterMatch = message.match(/reset after\s+([0-9hms.]+)/i) + const rawDuration = afterMatch?.[1] + if (rawDuration) { + const parsed = parseDurationToMs(rawDuration) + if (parsed !== null) { + return { retryDelayMs: parsed, message, reason } + } + } + } + + return { retryDelayMs: null, message, reason } +} + +function parseDurationToMs(duration: string): number | null { + const simpleMatch = duration.match(/^(\d+(?:\.\d+)?)(ms|s|m|h)?$/i) + if (simpleMatch) { + const value = parseFloat(simpleMatch[1]!) + const unit = (simpleMatch[2] || 's').toLowerCase() + switch (unit) { + case 'h': + return value * 3600 * 1000 + case 'm': + return value * 60 * 1000 + case 's': + return value * 1000 + case 'ms': + return value + default: + return value * 1000 + } + } + + const compoundRegex = /(\d+(?:\.\d+)?)(h|m(?!s)|s|ms)/gi + let totalMs = 0 + let matchFound = false + let match: RegExpExecArray | null = null + + while (true) { + match = compoundRegex.exec(duration) + if (match === null) break + matchFound = true + const value = parseFloat(match[1]!) + const unit = match[2]?.toLowerCase() + switch (unit) { + case 'h': + totalMs += value * 3600 * 1000 + break + case 'm': + totalMs += value * 60 * 1000 + break + case 's': + totalMs += value * 1000 + break + case 'ms': + totalMs += value + break + } + } + + return matchFound ? totalMs : null +} diff --git a/packages/opencode/src/plugin/fetch-interceptor.ts b/packages/opencode/src/plugin/fetch-interceptor.ts index 18d2a9bb..5f59092c 100644 --- a/packages/opencode/src/plugin/fetch-interceptor.ts +++ b/packages/opencode/src/plugin/fetch-interceptor.ts @@ -1,4 +1,8 @@ -import { fetchWithAgyCliTransport } from '@cortexkit/antigravity-auth-core' +import { + extractRateLimitBodyInfo, + fetchWithAgyCliTransport, + retryAfterMsFromResponse, +} from '@cortexkit/antigravity-auth-core' import { ANTIGRAVITY_ENDPOINT_FALLBACKS } from '../constants' import { type SidebarRoutingEntry, @@ -123,30 +127,6 @@ function terminalFetchError( return lastError ?? new Error(fallbackMessage) } -/** Reads `retry-after-ms` / `retry-after` headers, in that order. */ -function retryAfterMsFromResponse( - response: Response, - defaultRetryMs: number = 60_000, -): number { - const retryAfterMsHeader = response.headers.get('retry-after-ms') - if (retryAfterMsHeader) { - const parsed = Number.parseInt(retryAfterMsHeader, 10) - if (!Number.isNaN(parsed) && parsed > 0) { - return parsed - } - } - - const retryAfterHeader = response.headers.get('retry-after') - if (retryAfterHeader) { - const parsed = Number.parseInt(retryAfterHeader, 10) - if (!Number.isNaN(parsed) && parsed > 0) { - return parsed * 1000 - } - } - - return defaultRetryMs -} - /** Formats a millisecond duration the way the legacy toast pipeline did. */ function formatWaitTime(ms: number): string { if (ms < 1000) return `${ms}ms` @@ -2147,137 +2127,3 @@ export function createFetchInterceptor( return { fetch, dispose } } - -interface RateLimitBodyInfo { - retryDelayMs: number | null - message?: string - quotaResetTime?: string - reason?: string -} - -function extractRateLimitBodyInfo(body: unknown): RateLimitBodyInfo { - if (!body || typeof body !== 'object') return { retryDelayMs: null } - - const error = (body as { error?: unknown }).error - const message = - error && typeof error === 'object' - ? (error as { message?: string }).message - : undefined - - const details = - error && typeof error === 'object' - ? (error as { details?: unknown[] }).details - : undefined - - let reason: string | undefined - if (Array.isArray(details)) { - for (const detail of details) { - if (!detail || typeof detail !== 'object') continue - const type = (detail as { '@type'?: string })['@type'] - if (typeof type === 'string' && type.includes('google.rpc.ErrorInfo')) { - const detailReason = (detail as { reason?: string }).reason - if (typeof detailReason === 'string') { - reason = detailReason - break - } - } - } - - for (const detail of details) { - if (!detail || typeof detail !== 'object') continue - const type = (detail as { '@type'?: string })['@type'] - if (typeof type === 'string' && type.includes('google.rpc.RetryInfo')) { - const retryDelay = (detail as { retryDelay?: string }).retryDelay - if (typeof retryDelay === 'string') { - const retryDelayMs = parseDurationToMs(retryDelay) - if (retryDelayMs !== null) { - return { retryDelayMs, message, reason } - } - } - } - } - - for (const detail of details) { - if (!detail || typeof detail !== 'object') continue - const metadata = (detail as { metadata?: Record }) - .metadata - if (metadata && typeof metadata === 'object') { - const quotaResetDelay = metadata.quotaResetDelay - const quotaResetTime = metadata.quotaResetTimeStamp - if (typeof quotaResetDelay === 'string') { - const quotaResetDelayMs = parseDurationToMs(quotaResetDelay) - if (quotaResetDelayMs !== null) { - return { - retryDelayMs: quotaResetDelayMs, - message, - quotaResetTime, - reason, - } - } - } - } - } - } - - if (message) { - const afterMatch = message.match(/reset after\s+([0-9hms.]+)/i) - const rawDuration = afterMatch?.[1] - if (rawDuration) { - const parsed = parseDurationToMs(rawDuration) - if (parsed !== null) { - return { retryDelayMs: parsed, message, reason } - } - } - } - - return { retryDelayMs: null, message, reason } -} - -function parseDurationToMs(duration: string): number | null { - const simpleMatch = duration.match(/^(\d+(?:\.\d+)?)(ms|s|m|h)?$/i) - if (simpleMatch) { - const value = parseFloat(simpleMatch[1]!) - const unit = (simpleMatch[2] || 's').toLowerCase() - switch (unit) { - case 'h': - return value * 3600 * 1000 - case 'm': - return value * 60 * 1000 - case 's': - return value * 1000 - case 'ms': - return value - default: - return value * 1000 - } - } - - const compoundRegex = /(\d+(?:\.\d+)?)(h|m(?!s)|s|ms)/gi - let totalMs = 0 - let matchFound = false - let match: RegExpExecArray | null = null - - while (true) { - match = compoundRegex.exec(duration) - if (match === null) break - matchFound = true - const value = parseFloat(match[1]!) - const unit = match[2]?.toLowerCase() - switch (unit) { - case 'h': - totalMs += value * 3600 * 1000 - break - case 'm': - totalMs += value * 60 * 1000 - break - case 's': - totalMs += value * 1000 - break - case 'ms': - totalMs += value - break - } - } - - return matchFound ? totalMs : null -} diff --git a/packages/opencode/src/plugin/persist-account-pool.ts b/packages/opencode/src/plugin/persist-account-pool.ts index 76e9140d..d2676289 100644 --- a/packages/opencode/src/plugin/persist-account-pool.ts +++ b/packages/opencode/src/plugin/persist-account-pool.ts @@ -1,176 +1,12 @@ -/** - * Account pool persistence for OAuth flows. - * - * Merges a batch of successful OAuth token-exchange results into the - * persisted pool. All reads + writes happen inside the core - * `mutateAccountStorage` callback so the mutator sees the freshest - * state read while the lock is held — without it, a concurrent add - * would race the read-modify-write and silently disappear. - * - * Two upsert keys are honored, in priority order: - * 1. email — survives refresh-token rotation for the same Google account - * 2. refresh token — handles the no-email case and out-of-band rotations - * - * Destructive (`replaceAll: true`) writes start from an empty v4 inside - * the same locked callback so a stale merge cannot resurrect a removed - * account. - */ - -import type { - AccountMetadataV3, - AccountStorageV4, +import { + type AntigravityTokenExchangeResult, + persistAccountPoolAtPath, } from '@cortexkit/antigravity-auth-core' -import { mutateAccountStorage } from '@cortexkit/antigravity-auth-core' - -import type { AntigravityTokenExchangeResult } from '../antigravity/oauth' -import { parseRefreshParts } from './auth' import { getStoragePath } from './storage' -type TokenSuccess = Extract - -function clampInt(value: number, min: number, max: number): number { - if (!Number.isFinite(value)) { - return min - } - return Math.min(max, Math.max(min, Math.floor(value))) -} - -function applyUpserts( - current: AccountStorageV4, - results: TokenSuccess[], - replaceAll: boolean, -): AccountStorageV4 | undefined { - const now = Date.now() - - // For fresh logins, start from empty inside the locked callback so - // a stale merge cannot resurrect a removed account. - const accounts: AccountMetadataV3[] = replaceAll ? [] : [...current.accounts] - - const indexByRefreshToken = new Map() - const indexByEmail = new Map() - for (let i = 0; i < accounts.length; i++) { - const acc = accounts[i] - if (!acc) continue - if (acc.refreshToken) { - indexByRefreshToken.set(acc.refreshToken, i) - } - if (acc.email) { - indexByEmail.set(acc.email, i) - } - } - - for (const result of results) { - const parts = parseRefreshParts(result.refresh) - if (!parts.refreshToken) { - continue - } - - // Email match wins over token match — handles refresh-token rotation - // for the same Google account. - const existingByEmail = result.email - ? indexByEmail.get(result.email) - : undefined - const existingByToken = indexByRefreshToken.get(parts.refreshToken) - const existingIndex = existingByEmail ?? existingByToken - - if (existingIndex === undefined) { - const newIndex = accounts.length - indexByRefreshToken.set(parts.refreshToken, newIndex) - if (result.email) { - indexByEmail.set(result.email, newIndex) - } - accounts.push({ - email: result.email, - label: result.label, - refreshToken: parts.refreshToken, - projectId: parts.projectId, - managedProjectId: parts.managedProjectId, - addedAt: now, - lastUsed: now, - enabled: true, - }) - continue - } - - const existing = accounts[existingIndex] - if (!existing) continue - - const oldToken = existing.refreshToken - accounts[existingIndex] = { - ...existing, - email: result.email ?? existing.email, - label: result.label ?? existing.label, - refreshToken: parts.refreshToken, - projectId: parts.projectId ?? existing.projectId, - managedProjectId: parts.managedProjectId ?? existing.managedProjectId, - lastUsed: now, - } - - if (oldToken !== parts.refreshToken) { - indexByRefreshToken.delete(oldToken) - indexByRefreshToken.set(parts.refreshToken, existingIndex) - } - } - - if (accounts.length === 0) { - return undefined - } - - const activeIndex = replaceAll - ? 0 - : typeof current.activeIndex === 'number' && - Number.isFinite(current.activeIndex) - ? current.activeIndex - : 0 - - const clamped = clampInt(activeIndex, 0, accounts.length - 1) - return { - version: 4, - accounts, - activeIndex: clamped, - activeIndexByFamily: { - claude: clamped, - gemini: clamped, - }, - } -} - -/** - * Merge a batch of successful OAuth results into the persisted pool. - * - * - `replaceAll: true` — start from empty (fresh login) - * - `replaceAll: false` — preserve existing accounts, upsert by email - * then refresh token, bump `lastUsed` - * - * Both branches run their mutator INSIDE the locked callback. The - * `replaceAll` branch seeds the mutator from an empty v4 rather than - * reading the disk state, but the file lock is still required so the - * write is atomic against concurrent writers — a deleted-account merge - * would resurrect a stale account if we wrote without the lock. - */ export async function persistAccountPool( - results: TokenSuccess[], - replaceAll: boolean = false, + results: Extract[], + replaceAll = false, ): Promise { - if (results.length === 0) { - return - } - - const path = getStoragePath() - const emptyV4 = (): AccountStorageV4 => ({ - version: 4, - accounts: [], - activeIndex: 0, - }) - - if (replaceAll) { - await mutateAccountStorage(path, () => - applyUpserts(emptyV4(), results, true), - ) - return - } - - await mutateAccountStorage(path, (current) => - applyUpserts(current, results, false), - ) + await persistAccountPoolAtPath(getStoragePath(), results, replaceAll) } From 82a05e98bb3991ff1b28ea9e196e39847f6886d3 Mon Sep 17 00:00:00 2001 From: Brent Duarte Date: Wed, 9 Sep 2026 16:27:13 -0700 Subject: [PATCH 2/7] feat(pi): integrate durable multi-account routing and operator controls --- ARCHITECTURE.md | 22 +- README.md | 2 +- packages/pi/README.md | 144 +++++++- packages/pi/src/commands.ts | 117 +++++++ packages/pi/src/credential-cache.ts | 2 +- packages/pi/src/index.test.ts | 6 +- packages/pi/src/index.ts | 57 ++-- packages/pi/src/provider.test.ts | 267 +++++++++++++++ packages/pi/src/runtime.test.ts | 509 ++++++++++++++++++++++++++++ packages/pi/src/runtime.ts | 505 +++++++++++++++++++++++++++ packages/pi/src/settings.ts | 62 ++++ packages/pi/src/stream.ts | 62 +++- 12 files changed, 1707 insertions(+), 48 deletions(-) create mode 100644 packages/pi/src/commands.ts create mode 100644 packages/pi/src/provider.test.ts create mode 100644 packages/pi/src/runtime.test.ts create mode 100644 packages/pi/src/runtime.ts create mode 100644 packages/pi/src/settings.ts diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 38bf7ffc..70b582ac 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -127,7 +127,7 @@ graph TB Metadata --> Project ``` -The dependency graph is acyclic. No layer reaches into `account-manager` from `auth` — the auth module intentionally treats the refresh token as opaque so it can be reused by the Pi extension which has no concept of an account pool. +The dependency graph is acyclic. No layer reaches into `account-manager` from `auth` — the auth module treats the refresh token as opaque; both host adapters compose it with the shared account pool. ## OpenCode server plugin @@ -272,10 +272,28 @@ The notification queue is `packages/opencode/src/rpc/notifications.ts:1-59`. `pu - `pi.registerProvider(ANTIGRAVITY_PROVIDER_ID, { name, baseUrl, api, models, oauth, streamSimple })` (line 93-110). - `models` is `getPublicModelDefinitions()` filtered to drop image-output (Pi's `AssistantMessage` protocol has no image output type) and re-mapped onto Pi's `Model` shape (`packages/pi/src/index.ts:78-91`). - `oauth.login` invokes `authorizeAntigravity` from core, asks the host for the callback URL/code via `callbacks.onPrompt`, and calls `exchangeAntigravity` (line 22-58). -- `oauth.refreshToken` reads the packed `refreshToken|projectId|managedProjectId` triple, calls `refreshAntigravityToken` for the bare refresh, and re-packs the project segments (line 60-75). +- `oauth.login` persists successful exchanges via core's path-parameterized `persistAccountPoolAtPath`, also used by OpenCode, before returning the host credential. Email/token upserts preserve other accounts and disabled flags. +- `oauth.refreshToken` delegates to `PiAccountRuntime.refreshHost`: refreshes an enabled pool credential under core's fenced storage lock and can recover through a peer when the host's last credential fails. - `oauth.getApiKey` bridges the packed refresh into the stream by stashing it in `credential-cache.ts` so the stream can rejoin project context after the access token is stripped (line 102-107). - `streamSimple` is `streamCortexKitAntigravity` from `packages/pi/src/stream.ts`. It preserves same-model thinking/text/tool signatures, keeps native function-call IDs, emits same-target function responses with the AGY CLI model role, and tracks `last_execution_id` plus CLI-compatible step indexes per Pi session. +`packages/pi/src/runtime.ts` wires the existing `AccountManager`, quota manager, +quota probes and transport into Pi. The manager reconciles durable state before +selection while retaining process-local cursors and token-matched transient state. +Pi persists field-level updates through `mutateAccountStorage`; it never saves a +stale whole-pool snapshot over a peer's enable/disable or OAuth changes. The +default v4 pool is `$PI_AGENT_DIR/antigravity-accounts.json`, overridden by +`PI_ANTIGRAVITY_AUTH_FILE`. Legacy host auth is imported only into an empty pool. + +The two Pi settings (hybrid strategy and PID offset enabled by default) live at +`.config.json`. Core applies PID offset before all three strategy +branches. Quota uses current gemini/non-gemini groups, 30-minute on-demand refresh +and a 60-minute stale TTL; core's single-account and stale-cache exceptions remain. +Pre-stream rate-limit failover shares OpenCode's extracted response parsing and +core cooldown classification. Partial streams and ambiguous transport failures +are not replayed. Minimal `/agy-*` commands expose redacted pool controls; see +the [Pi README](packages/pi/README.md) for the operator contract and smoke test. + The package's `package.json` (`packages/pi/package.json:34-58`) declares `pi.extensions: ['./dist/index.js']` and pulls Pi's three peer dependencies from `@earendil-works/`. ### Pi's package-name dependency on the host diff --git a/README.md b/README.md index 296d2e8e..3adfcfc9 100644 --- a/README.md +++ b/README.md @@ -238,7 +238,7 @@ pi /login google-antigravity # OAuth flow ``` -The Pi extension registers the `google-antigravity` provider automatically through the package's `pi.extensions` field; no further wiring is needed. Both packages delegate to `@cortexkit/antigravity-auth-core`, so the request transport and model transforms are shared. The account-pool layer is not: Pi holds a single credential locally (no `antigravity-accounts.json`, no rotation, no killswitch, no operator settings). Multi-account rotation, quota routing, and the killswitch are OpenCode-only. +The Pi extension registers the `google-antigravity` provider automatically through the package's `pi.extensions` field; no further wiring is needed. Both packages delegate transport, transforms, multi-account storage, selection and quota routing to `@cortexkit/antigravity-auth-core`. Repeat `/login google-antigravity` to add accounts; use `/agy-accounts`, `/agy-quota refresh`, `/agy-strategy`, `/agy-enable` and `/agy-disable` to manage the pool. Pi defaults to hybrid selection with PID offset enabled and safely imports existing single-account host credentials. See the [Pi README](packages/pi/README.md) for storage paths, configuration, concurrency semantics and a manual smoke test. OpenCode's killswitch and Gemini CLI header fallback remain host-specific. ## Configuration reference diff --git a/packages/pi/README.md b/packages/pi/README.md index 402fdacb..9d6f4d10 100644 --- a/packages/pi/README.md +++ b/packages/pi/README.md @@ -23,6 +23,44 @@ pi install npm:@cortexkit/pi-antigravity-auth A browser URL is shown. Complete the Google OAuth flow and paste the resulting callback URL (or authorization code) back into the prompt. +Repeat `/login google-antigravity` (or `/agy-add`) and choose a different Google +account to add agy2, agy3, and so on. Login upserts by email, then refresh token; +re-authenticating an existing account updates its credentials without replacing +the pool or changing an operator's disabled flag. `/agy-add` invokes the same +registered OAuth flow through Pi's auth storage. + +## Account controls + +| Command | Effect | +| --- | --- | +| `/agy-accounts` | List agy1-based indexes, partially redacted email, enabled state, process-local health, cooldown, last selected account, and cached quota. | +| `/agy-add` | Start OAuth to add or re-authenticate an account. | +| `/agy-quota` | Display cached quota for every account, including stale/unknown state. | +| `/agy-quota refresh` | Refresh enabled accounts through core's quota manager; report failures while retaining cached values. | +| `/agy-strategy` | Show strategy and PID-offset setting. | +| `/agy-strategy sticky` | Keep using an eligible account until it becomes unavailable. | +| `/agy-strategy hybrid` | Use core health, token-bucket, freshness/LRU scoring and stickiness (default). | +| `/agy-strategy round-robin` | Rotate eligible accounts on each request, including tool continuations. | +| `/agy-disable agy2` | Disable an account without deleting its credentials (`2` also works). | +| `/agy-enable agy2` | Re-enable an account. Existing upstream eligibility/verification blocks must be resolved first. | + +Pi uses the existing core AccountManager and v4 storage. Current quota groups +are `gemini` (Flash/Pro) and `non-gemini` (Claude/GPT-OSS). With multiple enabled +accounts, core skips accounts at 80% usage while their quota cache is fresh. +Quota is refreshed on demand before requests when missing or 30 minutes old; +the cache expires after 60 minutes. Failed refreshes use core's bounded backoff. +Core's single-enabled-account exception and stale/unknown-cache fail-open +behavior are preserved. Disabled accounts are never selected. + +HTTP 429, 500, 503 and 529 responses before streaming use core classification, +RetryInfo/Retry-After parsing and cooldowns, then try another eligible account. +Each credential is attempted at most once per request, with one extra selection +to tolerate a concurrent token rotation. All-unavailable pools return an error +without waiting/spinning. Non-retryable HTTP errors, ambiguous transport failures, +and streams that have already emitted content are not replayed. Pi retains its +existing Antigravity transport; Gemini CLI header fallback, OpenCode killswitch +controls, and image-output routes are not enabled by this integration. + ## Models The extension registers the Antigravity model catalog under the @@ -49,13 +87,113 @@ Select a model with `/model` or `pi -m google-antigravity/antigravity-gemini-3.8 | `PI_AGENT_DIR` | Override the pi agent directory (default `~/.pi/agent`). | | `PI_ANTIGRAVITY_AUTH_FILE` | Override the account storage file path. | +The pool defaults to `~/.pi/agent/antigravity-accounts.json`, or +`$PI_AGENT_DIR/antigravity-accounts.json`. A small JSON settings file lives at +`.config.json` (so the auth-file override also relocates settings): + +```json +{ + "account_selection_strategy": "hybrid", + "pid_offset_enabled": true +} +``` + +Missing settings use these defaults. `/agy-strategy` persists only the strategy; +edit the JSON file to configure PID offset. Invalid settings fail closed. Pi +does not load OpenCode's host configuration. Restart Pi after changing PID offset +to observe a new initial assignment. + +PID offset seeds core's selection once per model family using PID modulo pool +size, for all three strategies. It spreads independent workers' starting points; +it is not an exclusive reservation, and different PIDs can have the same offset. +Routing cursors, health and token-bucket scores are process-local, as in core; +credentials, quota and cooldowns survive restarts. `/agy-accounts` marks the last +account dispatched by the current process, not a global cross-process selection. + +### Migration, concurrency and security + +Existing Pi `auth.json` remains host-managed. At session start (or first use of a +host credential), the extension imports that single credential into a missing or +empty v4 pool. It never overwrites a nonempty pool with stale host auth and never +deletes the old credential. A malformed/unsupported pool is left intact and core +attempts a protected `.corrupt-*` backup; repair the original file before retrying. +Core also handles its existing v1–v3 account-file migrations. +Legacy credentials without an email can only be matched by refresh token. If +Google issues a different token on re-login, that unidentifiable legacy entry +may remain separately; disable it after confirming the new account works. + +All pool updates use core's renewable fenced lock and atomic `0600` writes. +Credential refresh holds that lock so two processes cannot overwrite a rotated +refresh token with stale data. Access tokens are cached only in process memory +by this extension (Pi still maintains its own host credential). Long refreshes +can cause another process to exhaust core's bounded lock-wait budget; retry the +request if storage is busy. + +Each selection reloads durable state while retaining local routing state. +Cooldown writes patch the current token's record and retain the longest deadline; +they do not overwrite enable flags, concurrent additions, or newer credentials. +An already dispatched request may finish after an account is disabled; the next +selection sees the new flag. Quota in-flight deduplication is per process, so +concurrent workers can still perform duplicate quota probes. + +Both `auth.json` and the account pool contain sensitive tokens. Keep them private, +including backups; never attach them to bug reports. Account/quota commands omit +tokens, project IDs, fingerprints and arbitrary account labels. The unofficial +authentication warning above applies to every account added to the pool. + +## Manual smoke test + +Build this checkout with `bun install --frozen-lockfile && bun run build`. +To load the local build without installing a release, from the repository root: + +```bash +pi -e ./packages/pi/dist/index.js +``` + +Use an otherwise unconfigured Pi extension list to avoid loading the released +and local versions together. In Pi: + +```text +/login google-antigravity +/login google-antigravity +/login google-antigravity +/agy-accounts +/agy-quota refresh +/agy-strategy round-robin +``` + +Choose a different Google account for each login. For the exact sequence below, +set `pid_offset_enabled` to `false` in `.config.json` and restart Pi. +Keep `account_selection_strategy` set to `round-robin`. With all three accounts +eligible, send four simple prompts with no tool calls, running `/agy-accounts` +after each: the selected markers should be agy1, agy2, agy3, agy1. With PID offset +enabled the same cycle can start at another account. Tool calls are additional +requests and also advance round-robin. + +```text +/agy-strategy hybrid +/agy-disable agy1 +/agy-accounts +``` + +Send another prompt; agy1 must be bypassed. Re-enable with `/agy-enable agy1`. +Refresh quota and, if an account is naturally exhausted or cooling down, confirm +that hybrid selects an eligible peer. Do not deliberately exhaust a live account; +the automated tests inject exhausted quota and HTTP rate limits deterministically. + +For parallel verification, set PID offset back to `true`, start three terminals +with the same `PI_AGENT_DIR`/`PI_ANTIGRAVITY_AUTH_FILE` and the local extension, +then inspect `/agy-accounts` after requests. Disable an account in one terminal; +subsequent requests in the others must bypass it. PID collisions are possible; +this test checks shared durable state and offset behavior, not exclusive leases. + ## Notes This package shares its transport, OAuth, fingerprint, and request-transform logic with the OpenCode plugin via -[`@cortexkit/antigravity-auth-core`](../core). The current pi release targets a -single authenticated account; multi-account rotation and quota gating are -provided by the OpenCode plugin and are planned for pi in a later release. +[`@cortexkit/antigravity-auth-core`](../core), including multi-account storage, +selection and quota routing. Pi's existing models and streaming event protocol +remain unchanged. ## License diff --git a/packages/pi/src/commands.ts b/packages/pi/src/commands.ts new file mode 100644 index 00000000..f2d46d98 --- /dev/null +++ b/packages/pi/src/commands.ts @@ -0,0 +1,117 @@ +import type { ExtensionAPI } from '@earendil-works/pi-coding-agent' +import type { PiAccountRuntime } from './runtime.ts' +import { isStrategy, readSettings, writeStrategy } from './settings.ts' + +export function registerAccountCommands( + pi: ExtensionAPI, + runtime: PiAccountRuntime, +): void { + const register = ( + name: string, + description: string, + handler: Parameters[1]['handler'], + ) => { + pi.registerCommand(name, { + description, + handler: async (args, context) => { + try { + await handler(args, context) + } catch { + context.ui.notify( + 'Antigravity command failed. Check the account/ settings file, account number, or re-authenticate. Existing credentials were retained.', + 'error', + ) + } + }, + }) + } + register( + 'agy-accounts', + 'List Antigravity accounts (no tokens)', + async (_args, ctx) => { + ctx.ui.notify(await runtime.describe(), 'info') + }, + ) + register( + 'agy-quota', + 'Show cached quota; /agy-quota refresh fetches current quota', + async (args, ctx) => { + if (args.trim() === 'refresh') { + const failures = await runtime.refreshQuota(true) + if (failures) + ctx.ui.notify( + `Quota refresh failed for ${failures} account(s); cached values retained.`, + 'warning', + ) + } else if (args.trim()) throw new Error('Use /agy-quota [refresh]') + ctx.ui.notify(await runtime.describe(), 'info') + }, + ) + register( + 'agy-strategy', + 'Show or set sticky, hybrid, round-robin', + async (args, ctx) => { + const value = args.trim() + if (value) { + if (!isStrategy(value)) { + ctx.ui.notify( + 'Usage: /agy-strategy [sticky|hybrid|round-robin]', + 'warning', + ) + return + } + await writeStrategy(runtime.settingsPath, value) + } + const config = await readSettings(runtime.settingsPath) + ctx.ui.notify( + `Strategy: ${config.account_selection_strategy}; PID offset: ${config.pid_offset_enabled}`, + 'info', + ) + }, + ) + for (const enabled of [true, false]) { + register( + enabled ? 'agy-enable' : 'agy-disable', + 'Toggle an account: agy1 or 1', + async (args, ctx) => { + const match = args.trim().match(/^(?:agy)?([1-9]\d*)$/) + if (!match) { + ctx.ui.notify( + `Usage: /agy-${enabled ? 'enable' : 'disable'} agy1`, + 'warning', + ) + return + } + await runtime.setEnabled(Number(match[1]) - 1, enabled) + ctx.ui.notify(await runtime.describe(), 'info') + }, + ) + } + register( + 'agy-add', + 'Add an account using the provider OAuth login', + async (_args, ctx) => { + if (!ctx.hasUI) throw new Error('OAuth requires interactive Pi') + await ctx.modelRegistry.authStorage.login('google-antigravity', { + onAuth: ({ url }) => + ctx.ui.notify(`Open this URL in your browser:\n${url}`, 'info'), + onPrompt: async ({ message }) => { + const value = await ctx.ui.input(message) + if (!value) throw new Error('Login cancelled') + return value + }, + onDeviceCode: () => { + throw new Error('Unexpected device-code flow') + }, + onSelect: async ({ message, options }) => { + const label = await ctx.ui.select( + message, + options.map((option) => option.label), + ) + return options.find((option) => option.label === label)?.id + }, + }) + ctx.ui.notify(await runtime.describe(), 'info') + }, + ) +} diff --git a/packages/pi/src/credential-cache.ts b/packages/pi/src/credential-cache.ts index 186ea627..a12f17e3 100644 --- a/packages/pi/src/credential-cache.ts +++ b/packages/pi/src/credential-cache.ts @@ -10,7 +10,7 @@ */ const packedRefreshByAccessToken = new Map() -// Single-account extension: keep the map tiny. +// Only recent host credentials use this bridge; the runtime owns pool tokens. const MAX_ENTRIES = 4 export function rememberPackedRefresh( diff --git a/packages/pi/src/index.test.ts b/packages/pi/src/index.test.ts index fcf0d991..b0cdfa5e 100644 --- a/packages/pi/src/index.test.ts +++ b/packages/pi/src/index.test.ts @@ -5,7 +5,11 @@ import cortexKitPiAntigravityAuth from './index.ts' describe('Pi Antigravity model catalog', () => { it('exposes the live GPT-OSS route but not unsupported image-output chat models', () => { const registerProvider = mock() - cortexKitPiAntigravityAuth({ registerProvider } as never) + cortexKitPiAntigravityAuth({ + registerProvider, + registerCommand: mock(), + on: mock(), + } as never) expect(registerProvider).toHaveBeenCalledTimes(1) const [, config] = registerProvider.mock.calls[0] as [ diff --git a/packages/pi/src/index.ts b/packages/pi/src/index.ts index f142a163..8e1de654 100644 --- a/packages/pi/src/index.ts +++ b/packages/pi/src/index.ts @@ -2,15 +2,15 @@ import { authorizeAntigravity, exchangeAntigravity, getPublicModelDefinitions, - refreshAntigravityToken, } from '@cortexkit/antigravity-auth-core' import type { OAuthCredentials, OAuthLoginCallbacks, } from '@earendil-works/pi-ai' import type { ExtensionAPI } from '@earendil-works/pi-coding-agent' - +import { registerAccountCommands } from './commands.ts' import { rememberPackedRefresh } from './credential-cache.ts' +import { PiAccountRuntime } from './runtime.ts' import { streamCortexKitAntigravity } from './stream.ts' const ANTIGRAVITY_PROVIDER_ID = 'google-antigravity' @@ -21,8 +21,10 @@ function textImageInput(): Array<'text' | 'image'> { async function loginAntigravity( callbacks: OAuthLoginCallbacks, + runtime: PiAccountRuntime, ): Promise { const auth = await authorizeAntigravity() + callbacks.signal?.throwIfAborted() callbacks.onAuth({ url: auth.url }) const code = await callbacks.onPrompt({ message: 'Paste the Antigravity OAuth callback URL or code:', @@ -40,41 +42,48 @@ async function loginAntigravity( const codeParam = url.searchParams.get('code') const stateParam = url.searchParams.get('state') if (codeParam) rawCode = codeParam + if (stateParam && stateParam !== authState) + throw new Error('OAuth state mismatch') if (stateParam) state = stateParam - } catch { + } catch (error) { + if (!(error instanceof TypeError)) throw error // Not a URL — treat the input as a bare authorization code. } + callbacks.signal?.throwIfAborted() const result = await exchangeAntigravity(rawCode, state) if (result.type !== 'success') { throw new Error(`Antigravity OAuth exchange failed: ${result.error}`) } + callbacks.signal?.throwIfAborted() + await runtime.login(result) + return { refresh: result.refresh, access: result.access, expires: result.expires, - } -} - -async function refreshAntigravityCredentials( - credentials: OAuthCredentials, -): Promise { - // Stored refresh is `refreshToken|projectId|managedProjectId`. - const refreshToken = credentials.refresh.split('|')[0] ?? credentials.refresh - const refreshed = await refreshAntigravityToken(refreshToken) - // Preserve the project segments packed into the stored refresh string. - const projectSegments = credentials.refresh.includes('|') - ? credentials.refresh.slice(credentials.refresh.indexOf('|')) - : '' - return { - refresh: `${refreshed.refresh}${projectSegments}`, - access: refreshed.access, - expires: refreshed.expires, + email: result.email, } } export default function cortexKitPiAntigravityAuth(pi: ExtensionAPI): void { + const runtime = new PiAccountRuntime() + registerAccountCommands(pi, runtime) + pi.on('session_start', async (_event, context) => { + const auth = context.modelRegistry.authStorage.get(ANTIGRAVITY_PROVIDER_ID) + if (auth?.type === 'oauth') { + try { + await runtime.migrate(auth) + } catch { + context.ui.notify( + 'Antigravity account migration failed; existing auth and pool were retained. Repair the account file before retrying.', + 'error', + ) + } + } + }) + pi.on('session_shutdown', async () => runtime.dispose()) const models = Object.values(getPublicModelDefinitions()) // Pi's AssistantMessage protocol has no image-output content type. Keep // generation-only image routes out of the chat model catalog rather than @@ -97,15 +106,17 @@ export default function cortexKitPiAntigravityAuth(pi: ExtensionAPI): void { models, oauth: { name: 'Google Antigravity (CortexKit)', - login: loginAntigravity, - refreshToken: refreshAntigravityCredentials, + login: (callbacks) => loginAntigravity(callbacks, runtime), + refreshToken: (credentials) => runtime.refreshHost(credentials), getApiKey: (credentials) => { // Bridge the packed refresh (refreshToken|projectId|managedProjectId) // to the stream, which otherwise only receives the bare access token. rememberPackedRefresh(credentials.access, credentials.refresh) + runtime.remember(credentials) return credentials.access }, }, - streamSimple: streamCortexKitAntigravity, + streamSimple: (model, context, options) => + streamCortexKitAntigravity(model, context, options, runtime), }) } diff --git a/packages/pi/src/provider.test.ts b/packages/pi/src/provider.test.ts new file mode 100644 index 00000000..5fb9986f --- /dev/null +++ b/packages/pi/src/provider.test.ts @@ -0,0 +1,267 @@ +import { afterEach, beforeEach, describe, expect, it, mock } from 'bun:test' +import { mkdtemp, rm } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import * as core from '@cortexkit/antigravity-auth-core' +import type { + Api, + Model, + OAuthCredentials, + OAuthLoginCallbacks, +} from '@earendil-works/pi-ai' +import type { + ExtensionAPI, + ExtensionCommandContext, +} from '@earendil-works/pi-coding-agent' + +const exchange = mock() +const transport = mock() +mock.module('@cortexkit/antigravity-auth-core', () => ({ + ...core, + authorizeAntigravity: async () => ({ + url: 'https://accounts.google.com/auth?state=expected-state', + verifier: 'verifier', + projectId: '', + }), + exchangeAntigravity: exchange, + fetchWithAgyCliTransport: transport, +})) +const { default: register } = await import('./index.ts') + +type Provider = Parameters[1] +type Command = Parameters[1] +let provider: Provider +const commands = new Map() +const hooks = new Map Promise>() +let directory: string +let previousPath: string | undefined +let previousFetch: typeof fetch +let sequence: number +let hostAuth: OAuthCredentials | undefined +const notify = mock() +const callbacks = { + onAuth: mock(), + onPrompt: async () => 'code', + onDeviceCode: mock(), + onSelect: async () => undefined, +} satisfies OAuthLoginCallbacks + +beforeEach(async () => { + directory = await mkdtemp(join(tmpdir(), 'pi-provider-')) + previousPath = process.env.PI_ANTIGRAVITY_AUTH_FILE + process.env.PI_ANTIGRAVITY_AUTH_FILE = join(directory, 'accounts.json') + previousFetch = globalThis.fetch + globalThis.fetch = mock(async () => + Response.json({ groups: [] }), + ) as unknown as typeof fetch + core.initHealthTracker({}) + core.initTokenTracker({}) + sequence = 0 + hostAuth = undefined + exchange.mockReset() + transport.mockReset() + notify.mockClear() + commands.clear() + hooks.clear() + exchange.mockImplementation(async () => { + const i = ++sequence + return { + type: 'success', + email: `person${i}@example.com`, + refresh: `refresh-${i}|project|managed`, + access: `access-${i}`, + expires: Date.now() + 3600_000, + projectId: 'project', + } + }) + register({ + registerProvider: (id: string, config: Provider) => { + expect(id).toBe('google-antigravity') + provider = config + }, + registerCommand: (name: string, config: Command) => + commands.set(name, config), + on: (name: string, callback: (...args: unknown[]) => Promise) => + hooks.set(name, callback), + } as unknown as ExtensionAPI) +}) + +afterEach(async () => { + await hooks.get('session_shutdown')?.() + globalThis.fetch = previousFetch + if (previousPath === undefined) delete process.env.PI_ANTIGRAVITY_AUTH_FILE + else process.env.PI_ANTIGRAVITY_AUTH_FILE = previousPath + await rm(directory, { recursive: true, force: true }) +}) + +function context() { + return { + hasUI: true, + ui: { notify, input: async () => 'code', select: async () => undefined }, + modelRegistry: { + authStorage: { + get: () => (hostAuth ? { type: 'oauth', ...hostAuth } : undefined), + login: async (_id: string, cb: OAuthLoginCallbacks) => { + hostAuth = await provider.oauth!.login(cb) + }, + }, + }, + } as unknown as ExtensionCommandContext +} + +describe('Pi provider multi-account integration', () => { + it('repeated /login and /agy-add share the OAuth flow and retain all accounts', async () => { + hostAuth = await provider.oauth!.login(callbacks) + hostAuth = await provider.oauth!.login(callbacks) + await commands.get('agy-add')!.handler('', context()) + const accounts = await core.loadAccountStorage( + process.env.PI_ANTIGRAVITY_AUTH_FILE!, + ) + expect(accounts?.accounts).toHaveLength(3) + expect(exchange).toHaveBeenCalledTimes(3) + for (const name of [ + 'agy-accounts', + 'agy-add', + 'agy-quota', + 'agy-strategy', + 'agy-enable', + 'agy-disable', + ]) + expect(commands.has(name)).toBe(true) + }) + + it('migrates host auth at session start before a second login replaces the host credential', async () => { + hostAuth = { + refresh: 'legacy|project|managed', + access: 'legacy-access', + expires: Date.now() + 3600_000, + } + await hooks.get('session_start')?.({}, context()) + await provider.oauth!.login(callbacks) + expect( + (await core.loadAccountStorage(process.env.PI_ANTIGRAVITY_AUTH_FILE!)) + ?.accounts, + ).toHaveLength(2) + }) + + it('rejects a mismatched OAuth state and cancellation before exchange/persistence', async () => { + await expect( + provider.oauth!.login({ + ...callbacks, + onPrompt: async () => 'http://localhost/?code=code&state=wrong', + }), + ).rejects.toThrow('state mismatch') + await expect( + provider.oauth!.login({ ...callbacks, signal: AbortSignal.abort() }), + ).rejects.toThrow() + expect(exchange).not.toHaveBeenCalled() + expect( + await core.loadAccountStorage(process.env.PI_ANTIGRAVITY_AUTH_FILE!), + ).toBeNull() + }) + + it('operator commands persist strategy/enable state and never display tokens', async () => { + await provider.oauth!.login(callbacks) + await provider.oauth!.login(callbacks) + const ctx = context() + await commands.get('agy-strategy')!.handler('round-robin', ctx) + await commands.get('agy-disable')!.handler('agy1', ctx) + expect( + (await core.loadAccountStorage(process.env.PI_ANTIGRAVITY_AUTH_FILE!)) + ?.accounts[0]?.enabled, + ).toBe(false) + await commands.get('agy-enable')!.handler('1', ctx) + await commands.get('agy-quota')!.handler('refresh', ctx) + await commands.get('agy-accounts')!.handler('', ctx) + const output = JSON.stringify(notify.mock.calls) + expect(output).toContain('round-robin') + expect(output).toContain('agy2') + expect(output).not.toContain('access-') + expect(output).not.toContain('refresh-') + expect( + (await core.loadAccountStorage(process.env.PI_ANTIGRAVITY_AUTH_FILE!)) + ?.accounts[0]?.enabled, + ).toBe(true) + }) + + it('registered stream routes agy1, agy2, agy3, agy1 and keeps request/session transforms', async () => { + for (let i = 0; i < 3; i++) + hostAuth = await provider.oauth!.login(callbacks) + const authFile = process.env.PI_ANTIGRAVITY_AUTH_FILE! + await core.writeJsonAtomic(`${authFile}.config.json`, { + account_selection_strategy: 'round-robin', + pid_offset_enabled: false, + }) + transport.mockImplementation( + async () => + new Response( + 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"hello"}]},"finishReason":"STOP"}]}}\n\n', + ), + ) + const model = { + ...provider.models![0], + api: 'google-generative-ai', + provider: 'google-antigravity', + } as Model + const apiKey = provider.oauth!.getApiKey(hostAuth!) + for (let i = 0; i < 4; i++) { + const result = provider.streamSimple!( + model, + { messages: [{ role: 'user', content: 'hello', timestamp: 1 }] }, + { apiKey, sessionId: 'session' }, + ) + expect((await result.result()).stopReason).toBe('stop') + } + const headers = transport.mock.calls.map( + (call) => (call[1] as RequestInit).headers as Record, + ) + expect(headers.map((header) => header.Authorization)).toEqual([ + 'Bearer access-1', + 'Bearer access-2', + 'Bearer access-3', + 'Bearer access-1', + ]) + const bodies = transport.mock.calls.map((call) => + JSON.parse((call[1] as RequestInit).body as string), + ) + expect(bodies[0].project).toBe('managed') + expect(bodies[0].request.sessionId).toBe(bodies[1].request.sessionId) + expect(headers.every((header) => !!header['User-Agent'])).toBe(true) + }) + + it('fails over before streaming, but does not replay a partial stream', async () => { + for (let i = 0; i < 3; i++) + hostAuth = await provider.oauth!.login(callbacks) + await core.writeJsonAtomic( + `${process.env.PI_ANTIGRAVITY_AUTH_FILE!}.config.json`, + { account_selection_strategy: 'sticky', pid_offset_enabled: false }, + ) + transport.mockImplementationOnce( + async () => new Response('', { status: 429 }), + ) + transport.mockImplementationOnce( + async () => + new Response( + 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"partial"}]}}]}}\n\n', + ), + ) + const model = { + ...provider.models![0], + api: 'google-generative-ai', + provider: 'google-antigravity', + } as Model + const stream = provider.streamSimple!( + model, + { messages: [] }, + { apiKey: provider.oauth!.getApiKey(hostAuth!) }, + ) + const result = await stream.result() + expect(result.stopReason).toBe('error') + expect(result.content).toContainEqual({ type: 'text', text: 'partial' }) + expect(transport).toHaveBeenCalledTimes(2) + const bodies = transport.mock.calls.map((call) => + JSON.parse((call[1] as RequestInit).body as string), + ) + expect(bodies[0].requestId).toBe(bodies[1].requestId) + }) +}) diff --git a/packages/pi/src/runtime.test.ts b/packages/pi/src/runtime.test.ts new file mode 100644 index 00000000..4d286b68 --- /dev/null +++ b/packages/pi/src/runtime.test.ts @@ -0,0 +1,509 @@ +import { afterEach, beforeEach, describe, expect, it, mock } from 'bun:test' +import { mkdtemp, readFile, rm, stat, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { + type AntigravityRefreshResult, + getHealthTracker, + initHealthTracker, + initTokenTracker, + loadAccountStorage, + mutateAccountStorage, +} from '@cortexkit/antigravity-auth-core' +import { PiAccountRuntime } from './runtime.ts' +import { readSettings, writeStrategy } from './settings.ts' + +const model = 'gemini-3.8-flash' +let directory: string +let path: string +const runtimes: PiAccountRuntime[] = [] +let previousFetch: typeof fetch +const quotaFetch = mock() +const refresh = mock( + async (token: string): Promise => ({ + refresh: token, + access: `access-${token}`, + expires: Date.now() + 3600_000, + }), +) + +function login(index: number, token = `secret-refresh-${index}`) { + return { + type: 'success' as const, + email: `user${index}@example.com`, + refresh: `${token}|project-${index}|managed-${index}`, + access: `secret-access-${index}`, + expires: Date.now() + 3600_000, + projectId: `project-${index}`, + } +} + +function runtime(pid = 0) { + const result = new PiAccountRuntime({ path, pid, refreshToken: refresh }) + runtimes.push(result) + return result +} + +async function pool(count = 3) { + const result = runtime() + for (let i = 1; i <= count; i++) await result.login(login(i)) + await result.refreshQuota() + return result +} + +async function dispatch(result: PiAccountRuntime) { + const request = await result.dispatch(model, async () => new Response('ok')) + result.complete(request.account, true) + return request.account.index +} + +beforeEach(async () => { + directory = await mkdtemp(join(tmpdir(), 'pi-pool-')) + path = join(directory, 'accounts.json') + initHealthTracker({}) + initTokenTracker({}) + refresh.mockClear() + previousFetch = globalThis.fetch + quotaFetch.mockReset() + quotaFetch.mockImplementation(async () => + Response.json({ + groups: [ + { + displayName: 'Gemini', + buckets: [ + { + bucketId: 'gemini-test', + displayName: 'Gemini', + window: '5h', + remainingFraction: 0.9, + resetTime: '2099-01-01T00:00:00Z', + }, + ], + }, + { + displayName: 'Other', + buckets: [ + { + bucketId: '3p-test', + displayName: 'Other', + window: 'weekly', + remainingFraction: 0.8, + resetTime: '2099-01-01T00:00:00Z', + }, + ], + }, + ], + }), + ) + globalThis.fetch = quotaFetch as unknown as typeof fetch +}) + +afterEach(async () => { + for (const result of runtimes.splice(0)) await result.dispose() + globalThis.fetch = previousFetch + await rm(directory, { recursive: true, force: true }) +}) + +describe('Pi shared account runtime', () => { + it('persists first, second and third logins using v4 and secure permissions', async () => { + const result = runtime() + for (let i = 1; i <= 3; i++) { + await result.login(login(i)) + expect((await loadAccountStorage(path))?.accounts).toHaveLength(i) + } + expect((await stat(path)).mode & 0o777).toBe(0o600) + expect(await readFile(path, 'utf8')).not.toContain('secret-access') + }) + + it('deduplicates by email and token, updates rotated credentials, preserves disabled state', async () => { + const result = await pool() + await result.setEnabled(0, false) + await result.login(login(1, 'rotated-secret')) + await result.login(login(2)) + const stored = await loadAccountStorage(path) + expect(stored?.accounts).toHaveLength(3) + expect(stored?.accounts[0]).toMatchObject({ + refreshToken: 'rotated-secret', + enabled: false, + }) + expect(await dispatch(result)).not.toBe(0) + }) + + it('reloads across sessions and refreshes the selected stored credential', async () => { + await pool() + const next = runtime() + expect(await dispatch(next)).toBe(0) + expect(refresh).toHaveBeenCalledWith('secret-refresh-1') + expect(await next.describe()).toContain('agy3') + }) + + it.each([ + 'sticky', + 'hybrid', + ] as const)('%s keeps a healthy current account', async (strategy) => { + const result = await pool() + await writeStrategy(result.settingsPath, strategy) + expect(await dispatch(result)).toBe(0) + expect(await dispatch(result)).toBe(0) + }) + + it('round-robin rotates per request while preserving the cursor through reloads', async () => { + const result = await pool() + await writeStrategy(result.settingsPath, 'round-robin') + const selected = [] + for (let i = 0; i < 4; i++) selected.push(await dispatch(result)) + expect(selected).toEqual([0, 1, 2, 0]) + }) + + it.each([ + 'sticky', + 'hybrid', + 'round-robin', + ] as const)('%s applies PID offset once for independent workers', async (strategy) => { + const result = await pool() + await writeStrategy(result.settingsPath, strategy) + expect(await dispatch(runtime(3))).toBe(0) + expect(await dispatch(runtime(4))).toBe(1) + expect(await dispatch(runtime(5))).toBe(2) + }) + + it('hybrid uses core health and token-bucket scoring', async () => { + const result = await pool() + getHealthTracker().recordFailure(0) + getHealthTracker().recordFailure(0) + expect(await dispatch(result)).toBe(1) + }) + + it.each([ + 429, 503, 529, 500, + ])('records HTTP %s cooldown and fails over once', async (status) => { + const result = await pool() + const send = mock(async () => new Response('ok')) + send.mockImplementationOnce(async () => + Response.json( + { + error: { + details: [ + { + '@type': 'type.googleapis.com/google.rpc.ErrorInfo', + reason: 'QUOTA_EXHAUSTED', + }, + { + '@type': 'type.googleapis.com/google.rpc.RetryInfo', + retryDelay: '120s', + }, + ], + }, + }, + { status }, + ), + ) + const selected = await result.dispatch(model, send) + expect(selected.account.index).toBe(1) + expect(send).toHaveBeenCalledTimes(2) + const stored = await loadAccountStorage(path) + expect( + stored?.accounts[0]?.rateLimitResetTimes?.[`gemini-antigravity:${model}`], + ).toBeGreaterThan(Date.now() + 110_000) + expect(getHealthTracker().getScore(0)).toBe(60) + expect(await dispatch(runtime())).not.toBe(0) + }) + + it('terminates when all accounts become rate-limited, without exposing provider secrets', async () => { + const result = await pool() + const send = mock( + async () => new Response('secret-refresh-1', { status: 429 }), + ) + await expect(result.dispatch(model, send)).rejects.toThrow('HTTP 429') + expect(send).toHaveBeenCalledTimes(3) + await expect(result.dispatch(model, send)).rejects.toThrow( + 'All Antigravity', + ) + expect(send).toHaveBeenCalledTimes(3) + expect(await result.describe()).not.toContain('secret-refresh') + }) + + it('does not retry non-retryable responses or ambiguous transport failures', async () => { + const result = await pool() + const send = mock(async () => new Response('bad request', { status: 400 })) + expect((await result.dispatch(model, send)).response.status).toBe(400) + expect(send).toHaveBeenCalledTimes(1) + const broken = mock(async (): Promise => { + throw new Error('transport failure') + }) + await expect(result.dispatch(model, broken)).rejects.toThrow( + 'transport failure', + ) + expect(broken).toHaveBeenCalledTimes(1) + }) + + it.each([ + 'gemini-3.8-flash', + 'gemini-3.1-pro', + 'claude-sonnet-4-6-thinking', + 'gpt-oss-120b-medium', + ])('bypasses exhausted quota for %s', async (requested) => { + const result = await pool() + await mutateAccountStorage(path, (current) => { + current.accounts[0]!.cachedQuota = { + gemini: { remainingFraction: 0, modelCount: 1 }, + 'non-gemini': { remainingFraction: 0, modelCount: 1 }, + } + return current + }) + expect( + (await result.dispatch(requested, async () => new Response('ok'))).account + .index, + ).toBe(1) + }) + + it('retains core single-account quota exception', async () => { + const result = await pool(1) + await mutateAccountStorage(path, (current) => { + current.accounts[0]!.cachedQuota!.gemini!.remainingFraction = 0 + return current + }) + expect(await dispatch(result)).toBe(0) + }) + + it('fails open on stale quota when refresh is unavailable and retains the stale cache', async () => { + const result = await pool() + await mutateAccountStorage(path, (current) => { + current.accounts[0]!.cachedQuota!.gemini!.remainingFraction = 0 + current.accounts[0]!.cachedQuotaUpdatedAt = Date.now() - 3 * 3600_000 + return current + }) + quotaFetch.mockImplementation(async () => { + throw new Error('network unavailable') + }) + expect(await dispatch(result)).toBe(0) + expect( + (await loadAccountStorage(path))?.accounts[0]?.cachedQuota?.gemini + ?.remainingFraction, + ).toBe(0) + expect(await result.describe()).toContain('stale') + }) + + it('excludes disabled accounts, including changes made by another running instance', async () => { + const first = await pool() + const second = runtime() + expect(await dispatch(first)).toBe(0) + await second.setEnabled(0, false) + expect(await dispatch(first)).toBe(1) + await second.setEnabled(1, false) + await second.setEnabled(2, false) + await expect(dispatch(first)).rejects.toThrow('All Antigravity') + await second.setEnabled(2, true) + expect(await dispatch(first)).toBe(2) + }) + + it('an in-flight cooldown write cannot undo another process disable or add', async () => { + const first = await pool() + const second = runtime() + let started!: () => void + let release!: () => void + const ready = new Promise((resolve) => { + started = resolve + }) + const gate = new Promise((resolve) => { + release = resolve + }) + let calls = 0 + const pending = first.dispatch(model, async () => { + if (++calls > 1) return new Response('ok') + started() + await gate + return new Response('', { status: 429 }) + }) + await ready + await second.setEnabled(0, false) + await second.login(login(4)) + release() + await pending + const stored = await loadAccountStorage(path) + expect(stored?.accounts).toHaveLength(4) + expect(stored?.accounts[0]?.enabled).toBe(false) + expect(stored?.accounts[0]?.rateLimitResetTimes).toBeDefined() + }) + + it('concurrent logins preserve every account through fenced writes', async () => { + await Promise.all([1, 2, 3].map((i) => runtime(i).login(login(i)))) + expect((await loadAccountStorage(path))?.accounts).toHaveLength(3) + }) + + it('independent OS processes reload the same pool and apply their actual PID offset', async () => { + await pool() + const script = ` + import { PiAccountRuntime } from ${JSON.stringify(new URL('./runtime.ts', import.meta.url).pathname)}; + const runtime = new PiAccountRuntime({ path: process.argv[1], refreshToken: async token => ({ refresh: token, access: 'test-access', expires: Date.now() + 3600000 }) }); + const result = await runtime.dispatch('gemini-3.8-flash', async () => new Response('ok')); + console.log(JSON.stringify({ pid: process.pid, index: result.account.index })); + await runtime.dispose(); + ` + const children = Array.from({ length: 3 }, () => + Bun.spawn([process.execPath, '--eval', script, path], { + stdout: 'pipe', + stderr: 'pipe', + }), + ) + const results = await Promise.all( + children.map(async (child) => { + const [code, output, errors] = await Promise.all([ + child.exited, + new Response(child.stdout).text(), + new Response(child.stderr).text(), + ]) + expect(errors).toBe('') + expect(code).toBe(0) + return JSON.parse(output) as { pid: number; index: number } + }), + ) + for (const result of results) expect(result.index).toBe(result.pid % 3) + expect((await loadAccountStorage(path))?.accounts).toHaveLength(3) + }) + + it('serializes simultaneous token rotations and never restores the old token', async () => { + await pool(1) + const calls: string[] = [] + const rotatingRefresh = async (token: string) => { + calls.push(token) + await new Promise((resolve) => setTimeout(resolve, 20)) + return { + refresh: `${token}-rotated`, + access: 'new-access', + expires: Date.now() + 3600_000, + } + } + const workers = [1, 2].map( + () => + new PiAccountRuntime({ path, pid: 0, refreshToken: rotatingRefresh }), + ) + runtimes.push(...workers) + await Promise.all(workers.map(dispatch)) + expect(calls).toEqual(['secret-refresh-1', 'secret-refresh-1-rotated']) + const stored = await loadAccountStorage(path) + expect(stored?.accounts).toHaveLength(1) + expect(stored?.accounts[0]?.refreshToken).toBe( + 'secret-refresh-1-rotated-rotated', + ) + }) + + it('simultaneous cooldown writes retain the longest deadline for the same quota key', async () => { + await pool(1) + const workers = [runtime(), runtime()] + let waiting = 0 + let release!: () => void + const gate = new Promise((resolve) => { + release = resolve + }) + const results = await Promise.allSettled( + workers.map((worker, index) => + worker.dispatch(model, async () => { + if (++waiting === 2) release() + await gate + return new Response('', { + status: 429, + headers: { 'retry-after': index === 0 ? '120' : '60' }, + }) + }), + ), + ) + expect(waiting).toBe(2) + for (const result of results) { + expect(result.status).toBe('rejected') + if (result.status === 'rejected') + expect(String(result.reason)).toContain('HTTP 429') + } + expect( + (await loadAccountStorage(path))?.accounts[0]?.rateLimitResetTimes?.[ + `gemini-antigravity:${model}` + ], + ).toBeGreaterThan(Date.now() + 110_000) + }) + + it('a stale refresh cannot overwrite a token rotated by a peer', async () => { + const first = await pool() + const second = runtime() + await first.login(login(1, 'new-refresh')) + const result = await second.dispatch(model, async (auth) => { + expect(auth.refresh).toStartWith('new-refresh|') + return new Response('ok') + }) + expect(result.account.index).toBe(0) + expect((await loadAccountStorage(path))?.accounts).toHaveLength(3) + }) + + it('fails over a selected credential refresh failure and persists its cooldown', async () => { + await pool() + const broken = new PiAccountRuntime({ + path, + pid: 0, + refreshToken: async (token) => { + if (token === 'secret-refresh-1') + throw new Error('invalid_grant secret-access-1') + return refresh(token) + }, + }) + runtimes.push(broken) + expect(await dispatch(broken)).toBe(1) + expect((await loadAccountStorage(path))?.accounts[0]?.cooldownReason).toBe( + 'auth-failure', + ) + expect(await broken.describe()).not.toContain('secret-access') + }) + + it('the Pi host refresh can recover through another pool member', async () => { + const result = await pool() + await result.setEnabled(0, false) + const next = await result.refreshHost({ ...login(1), expires: 0 }) + expect(next.refresh).toStartWith('secret-refresh-2|') + }) + + it('migrates a legacy Pi credential once without overwriting or resurrecting accounts', async () => { + const result = runtime() + await result.migrate({ ...login(1) }) + await result.login(login(2)) + await result.login(login(1, 'new-refresh')) + await runtime().migrate({ ...login(1) }) + const stored = await loadAccountStorage(path) + expect(stored?.accounts).toHaveLength(2) + expect(stored?.accounts[0]?.refreshToken).toBe('new-refresh') + }) + + it.each([ + '{broken secret-refresh-1', + '{"version":99,"accounts":[]}', + '{"version":4,"accounts":"invalid"}', + ])('fails closed on malformed/future storage', async (text) => { + await writeFile(path, text) + const result = runtime() + await expect(result.login(login(1))).rejects.toThrow() + await expect(result.migrate({ ...login(2) })).rejects.toThrow() + expect(await readFile(path, 'utf8')).toBe(text) + }) + + it('aborted requests do not dispatch or consume routing attempts', async () => { + const result = await pool() + const send = mock(async () => new Response('ok')) + await expect( + result.dispatch(model, send, AbortSignal.abort()), + ).rejects.toThrow() + expect(send).not.toHaveBeenCalled() + }) + + it('operator output contains neither credentials nor arbitrary account labels', async () => { + const result = await pool() + await mutateAccountStorage(path, (current) => { + current.accounts[0]!.label = 'secret-access-1' + return current + }) + const output = await result.describe() + expect(output).toContain('u***@example.com') + expect(output).not.toContain('secret-') + expect(output).not.toContain('user1@') + expect(await readSettings(result.settingsPath)).toEqual({ + account_selection_strategy: 'hybrid', + pid_offset_enabled: true, + }) + }) +}) diff --git a/packages/pi/src/runtime.ts b/packages/pi/src/runtime.ts new file mode 100644 index 00000000..b2dcb225 --- /dev/null +++ b/packages/pi/src/runtime.ts @@ -0,0 +1,505 @@ +import { + AccountManager, + type AccountMetadataV3, + type AccountModelFamily, + ANTIGRAVITY_ENDPOINT_FALLBACKS, + type AntigravityTokenExchangeResult, + accessTokenExpired, + aggregateQuota, + aggregateQuotaSummary, + computeSoftQuotaCacheTtlMs, + createQuotaManager, + defaultAccountStorageStore, + defaultKeyOf, + ensureProjectContext, + extractRateLimitBodyInfo, + type FetchQuotaSummaryOptions, + fetchAvailableModels, + fetchQuotaSummary, + fetchWithActiveTimeout, + formatRefreshParts, + getHealthTracker, + getTokenTracker, + loadAccountStorage, + type ManagedAccount, + mutateAccountStorage, + type OAuthAuthDetails, + parseRateLimitReason, + parseRefreshParts, + persistAccountPoolAtPath, + refreshAntigravityToken, + resolveQuotaGroup, + retryAfterMsFromResponse, +} from '@cortexkit/antigravity-auth-core' +import type { OAuthCredentials } from '@earendil-works/pi-ai' +import { getPiAntigravityAuthFile } from './paths.ts' +import { readSettings } from './settings.ts' + +const QUOTA_REFRESH_MS = 30 * 60_000 +const QUOTA_TTL_MS = computeSoftQuotaCacheTtlMs('auto', 30) +type LoginResult = Extract + +class AccountUnavailable extends Error {} +class CredentialRefreshFailed extends Error {} + +export interface PiRuntimeOptions { + path?: string + pid?: number + refreshToken?: typeof refreshAntigravityToken +} + +/** Host wiring only: selection/scoring, quota aggregation, cooldowns and all + * durable writes remain owned by core. No provider request runs under a lock. + * Token refresh does, so a rotated token cannot be overwritten by a stale peer. + */ +export class PiAccountRuntime { + readonly path: string + readonly settingsPath: string + private manager?: AccountManager + private readonly credentials = new Map() + private readonly refreshToken: typeof refreshAntigravityToken + private readonly pid: number + private selected?: string + private legacy?: OAuthCredentials + private readonly quota + + constructor(options: PiRuntimeOptions = {}) { + this.path = options.path ?? getPiAntigravityAuthFile() + this.settingsPath = `${this.path}.config.json` + this.pid = options.pid ?? process.pid + this.refreshToken = options.refreshToken ?? refreshAntigravityToken + this.quota = createQuotaManager({ + keyOf: defaultKeyOf, + fetchAccountQuota: async (account, signal) => { + try { + const auth = await this.credentialFor(account.refreshToken) + signal.throwIfAborted() + const context = await ensureProjectContext(auth) + signal.throwIfAborted() + const parts = parseRefreshParts(context.auth.refresh) + const fetchVia: NonNullable = ( + url, + init, + extra, + ) => + fetchWithActiveTimeout( + url, + { + ...init, + signal: extra.signal + ? AbortSignal.any([signal, extra.signal]) + : signal, + }, + { timeoutMs: extra.timeoutMs }, + ) + const common = { + accessToken: context.auth.access ?? '', + projectId: context.effectiveProjectId, + managedProjectId: parts.managedProjectId, + endpoints: ANTIGRAVITY_ENDPOINT_FALLBACKS, + fetchVia, + } + let quota: ReturnType + try { + quota = aggregateQuotaSummary( + (await fetchQuotaSummary(common)).summary, + ) + } catch (error) { + if (signal.aborted) throw error + quota = aggregateQuota((await fetchAvailableModels(common)).models) + } + signal.throwIfAborted() + await this.patch(parts.refreshToken, (current) => { + current.projectId = parts.projectId ?? current.projectId + current.managedProjectId = + parts.managedProjectId ?? current.managedProjectId + current.cachedQuota = quota.groups + current.cachedQuotaUpdatedAt = Date.now() + }) + return { index: 0, status: 'ok', quota } + } catch { + // Provider bodies can contain credentials. Never relay them to quota + // manager diagnostics or operator output. + return { + index: 0, + status: 'error', + error: 'Quota refresh failed; cached quota retained', + } + } + }, + }) + } + + remember(credentials: OAuthCredentials): void { + this.legacy = credentials + this.credentials.set(parseRefreshParts(credentials.refresh).refreshToken, { + type: 'oauth', + ...credentials, + }) + } + + async login(result: LoginResult): Promise { + if (this.legacy) await this.migrate(this.legacy) + await persistAccountPoolAtPath(this.path, [result]) + this.remember({ + refresh: result.refresh, + access: result.access, + expires: result.expires, + email: result.email, + }) + } + + /** Import the host credential only into a missing/empty pool. An existing + * pool is authoritative: a stale auth.json must not re-add rotated accounts. + */ + async migrate(credentials: OAuthCredentials): Promise { + this.remember(credentials) + const parts = parseRefreshParts(credentials.refresh) + if (!parts.refreshToken) return + const stored = await loadAccountStorage(this.path) + if (stored?.accounts.length) return + await mutateAccountStorage(this.path, (current) => { + if (current.accounts.length) return current + current.accounts.push({ + ...parts, + email: + typeof credentials.email === 'string' ? credentials.email : undefined, + addedAt: Date.now(), + lastUsed: 0, + enabled: true, + }) + return current + }) + } + + private async reload(): Promise { + const stored = (await loadAccountStorage(this.path)) ?? { + version: 4 as const, + accounts: [], + activeIndex: 0, + } + const tokens = new Set( + stored.accounts.map((account) => account.refreshToken), + ) + for (const token of this.credentials.keys()) { + if (!tokens.has(token)) this.credentials.delete(token) + } + if (this.manager) this.manager.reconcileStorage(stored) + else + this.manager = new AccountManager(undefined, stored, { + store: defaultAccountStorageStore, + storagePath: this.path, + pid: this.pid, + persistFingerprintUpdates: false, + }) + return this.manager + } + + private async patch( + token: string, + update: (account: AccountMetadataV3) => void, + ): Promise { + await mutateAccountStorage(this.path, (current) => { + const account = current.accounts.find( + (entry) => entry.refreshToken === token, + ) + if (account) update(account) + return current + }) + } + + private async credentialFor(token: string): Promise { + let auth: OAuthAuthDetails | undefined + await mutateAccountStorage(this.path, async (current) => { + const account = current.accounts.find( + (entry) => entry.refreshToken === token, + ) + if ( + !account || + account.enabled === false || + account.accountIneligible || + account.verificationRequired + ) { + throw new AccountUnavailable('Account changed or is disabled') + } + const cached = this.credentials.get(token) + if (cached && !accessTokenExpired(cached)) { + auth = { ...cached, refresh: formatRefreshParts(account) } + return current + } + let refreshed: Awaited> + try { + refreshed = await this.refreshToken(token) + if ( + !refreshed.access || + !refreshed.refresh || + !Number.isFinite(refreshed.expires) + ) { + throw new Error('Invalid refresh result') + } + } catch { + throw new CredentialRefreshFailed( + 'Antigravity token refresh failed; re-authenticate the account', + ) + } + account.refreshToken = refreshed.refresh + auth = { + type: 'oauth', + access: refreshed.access, + expires: refreshed.expires, + refresh: formatRefreshParts({ + ...account, + refreshToken: refreshed.refresh, + }), + } + return current + }) + if (!auth) throw new AccountUnavailable('Account is unavailable') + this.credentials.set(parseRefreshParts(auth.refresh).refreshToken, auth) + return auth + } + + /** Pi refreshes its host credential before stream dispatch. A failed last + * login must not prevent a healthy pool member from reaching the runtime. + */ + async refreshHost(credentials: OAuthCredentials): Promise { + await this.migrate(credentials) + const manager = await this.reload() + for (const account of manager.getEnabledAccounts()) { + try { + const auth = await this.credentialFor(account.parts.refreshToken) + return { + refresh: auth.refresh, + access: auth.access ?? '', + expires: auth.expires ?? 0, + } + } catch (error) { + if ( + !( + error instanceof CredentialRefreshFailed || + error instanceof AccountUnavailable + ) + ) + throw error + if (error instanceof CredentialRefreshFailed) + await this.authFailure(account) + } + } + throw new Error( + 'No usable Antigravity credentials; /login google-antigravity to re-authenticate', + ) + } + + private async authFailure(account: ManagedAccount): Promise { + getHealthTracker().recordFailure(account.index) + this.manager?.markAccountCoolingDown(account, 60_000, 'auth-failure') + await this.patch(account.parts.refreshToken, (current) => { + current.coolingDownUntil = Math.max( + current.coolingDownUntil ?? 0, + account.coolingDownUntil ?? 0, + ) + current.cooldownReason = 'auth-failure' + }) + } + + async refreshQuota(force = false): Promise { + const accounts = (await loadAccountStorage(this.path))?.accounts ?? [] + const results = await this.quota.refreshAccounts( + accounts.filter( + (account) => + force || + account.cachedQuotaUpdatedAt == null || + Date.now() - account.cachedQuotaUpdatedAt >= QUOTA_REFRESH_MS, + ), + { force, indexFor: (account) => accounts.indexOf(account) }, + ) + return results.filter((result) => result.status === 'error').length + } + + async dispatch( + model: string, + send: ( + auth: OAuthAuthDetails, + account: ManagedAccount, + ) => Promise, + signal?: AbortSignal, + ): Promise<{ response: Response; account: ManagedAccount }> { + signal?.throwIfAborted() + if (this.legacy) await this.migrate(this.legacy) + await this.refreshQuota() + const config = await readSettings(this.settingsPath) + const attempted = new Set() + // One extra selection tolerates a peer rotating a token between reload and + // credential acquisition. The attempted-token set still bounds sends. + const budget = (await this.reload()).getTotalAccountCount() + 1 + let lastError = + 'All Antigravity accounts are disabled, cooling down, or over cached quota' + const family: AccountModelFamily = + resolveQuotaGroup('gemini', model) === 'non-gemini' ? 'claude' : 'gemini' + for (let attempt = 0; attempt < budget; attempt++) { + signal?.throwIfAborted() + const manager = await this.reload() + const excluded = new Set( + manager + .getAccounts() + .filter( + (a) => + attempted.has(a.parts.refreshToken) || + a.verificationRequired || + a.accountIneligible, + ) + .map((a) => a.index), + ) + const account = manager.getCurrentOrNextForFamily( + family, + model, + config.account_selection_strategy, + 'antigravity', + config.pid_offset_enabled, + 80, + QUOTA_TTL_MS, + undefined, + excluded, + ) + if (!account) break + attempted.add(account.parts.refreshToken) + let auth: OAuthAuthDetails + try { + auth = await this.credentialFor(account.parts.refreshToken) + } catch (error) { + if ( + !( + error instanceof CredentialRefreshFailed || + error instanceof AccountUnavailable + ) + ) + throw error + lastError = error.message + if (error instanceof CredentialRefreshFailed) + await this.authFailure(account) + continue + } + manager.updateFromAuth(account, auth) + const token = account.parts.refreshToken + attempted.add(token) + signal?.throwIfAborted() + manager.markAccountUsed(account.index) + await this.patch(token, (current) => { + current.lastUsed = Math.max(current.lastUsed, account.lastUsed) + current.fingerprint ??= account.fingerprint + }) + const consumed = getTokenTracker().consume(account.index) + let response: Response + try { + response = await send(auth, account) + } catch (error) { + if (consumed) getTokenTracker().refund(account.index) + if (!signal?.aborted) getHealthTracker().recordFailure(account.index) + // Unknown transport failures may occur after upstream accepted work. + // Surface them without replaying a potentially billable request. + throw error + } + if (response.ok) { + this.selected = token + return { response, account } + } + if (consumed) getTokenTracker().refund(account.index) + if (![429, 503, 529, 500].includes(response.status)) + return { response, account } + let body: unknown + try { + body = await response.json() + } catch { + body = undefined + } + const info = extractRateLimitBodyInfo(body) + manager.markRateLimitedWithReason( + account, + family, + 'antigravity', + model, + parseRateLimitReason(info.reason, info.message, response.status), + info.retryDelayMs ?? retryAfterMsFromResponse(response), + ) + getHealthTracker().recordRateLimit(account.index) + await this.patch(token, (current) => { + for (const [key, until] of Object.entries( + account.rateLimitResetTimes, + )) { + current.rateLimitResetTimes ??= {} + current.rateLimitResetTimes[key] = Math.max( + current.rateLimitResetTimes[key] ?? 0, + until ?? 0, + ) + } + }) + lastError = `Antigravity HTTP ${response.status}; no eligible account remains (cooldown recorded)` + } + throw new Error(lastError) + } + + complete(account: ManagedAccount, success: boolean): void { + const current = this.manager + ?.getAccounts() + .find((entry) => entry.parts.refreshToken === account.parts.refreshToken) + if (!current) return + if (success) { + this.manager?.markRequestSuccess(current) + getHealthTracker().recordSuccess(current.index) + } else getHealthTracker().recordFailure(current.index) + } + + async setEnabled(index: number, enabled: boolean): Promise { + await mutateAccountStorage(this.path, (current) => { + const account = current.accounts[index] + if (!account) throw new Error('Unknown account; use /agy-accounts') + if ( + enabled && + (account.accountIneligible || account.verificationRequired) + ) { + throw new Error( + 'Resolve upstream account eligibility or verification before enabling', + ) + } + account.enabled = enabled + return current + }) + } + + async describe(): Promise { + const manager = await this.reload() + return ( + manager + .getAccounts() + .map((account) => { + const email = account.email?.match(/^(.)([^@]*)@(.+)$/) + const label = email + ? `${email[1]}***@${email[3]}` + : '(email unavailable)' + const until = Math.max( + account.coolingDownUntil ?? 0, + ...Object.values(account.rateLimitResetTimes).map((n) => n ?? 0), + ) + const quota = + Object.entries(account.cachedQuota ?? {}) + .map( + ([group, value]) => + `${group}=${value.remainingFraction == null ? '?' : `${Math.round(value.remainingFraction * 100)}%`} remaining`, + ) + .join(', ') || 'quota unknown' + const age = + account.cachedQuotaUpdatedAt == null + ? '' + : ` (${Date.now() - account.cachedQuotaUpdatedAt > QUOTA_TTL_MS ? 'stale' : 'cached'})` + return `agy${account.index + 1} ${label} ${account.enabled ? 'enabled' : 'disabled'} health=${getHealthTracker().getScore(account.index)}${this.selected === account.parts.refreshToken ? ' selected' : ''} ${until > Date.now() ? `cooldown=${Math.ceil((until - Date.now()) / 1000)}s` : 'ready'} ${quota}${age}` + }) + .join('\n') || 'No Antigravity accounts. Use /login google-antigravity.' + ) + } + + async dispose(): Promise { + await this.quota.dispose() + await this.manager?.dispose() + this.credentials.clear() + } +} diff --git a/packages/pi/src/settings.ts b/packages/pi/src/settings.ts new file mode 100644 index 00000000..61b16f15 --- /dev/null +++ b/packages/pi/src/settings.ts @@ -0,0 +1,62 @@ +import { mkdir, readFile } from 'node:fs/promises' +import { dirname } from 'node:path' +import { + type AccountSelectionStrategy, + acquireFencedFileLock, + writeJsonAtomic, +} from '@cortexkit/antigravity-auth-core' + +export interface PiRoutingSettings { + account_selection_strategy: AccountSelectionStrategy + pid_offset_enabled: boolean +} + +export function isStrategy(value: unknown): value is AccountSelectionStrategy { + return value === 'sticky' || value === 'hybrid' || value === 'round-robin' +} + +export async function readSettings(path: string): Promise { + let value: unknown + try { + value = JSON.parse(await readFile(path, 'utf8')) + } catch (error) { + if ((error as NodeJS.ErrnoException).code === 'ENOENT') { + return { account_selection_strategy: 'hybrid', pid_offset_enabled: true } + } + throw new Error( + 'Cannot read Pi Antigravity routing settings; repair the settings file', + ) + } + if (!value || typeof value !== 'object') + throw new Error('Invalid Pi Antigravity settings') + const config = value as Record + const strategy = config.account_selection_strategy ?? 'hybrid' + const offset = config.pid_offset_enabled ?? true + if (!isStrategy(strategy) || typeof offset !== 'boolean') { + throw new Error('Invalid Pi Antigravity routing strategy or PID offset') + } + return { account_selection_strategy: strategy, pid_offset_enabled: offset } +} + +export async function writeStrategy( + path: string, + strategy: AccountSelectionStrategy, +): Promise { + await mkdir(dirname(path), { recursive: true }) + const lock = await acquireFencedFileLock({ + path, + name: 'pi-routing', + ttlMs: 10_000, + renew: true, + }) + if (!lock) + throw new Error('Pi Antigravity settings are busy; retry the command') + try { + const config = await readSettings(path) + config.account_selection_strategy = strategy + await lock.assertOwned() + await writeJsonAtomic(path, config) + } finally { + await lock.release() + } +} diff --git a/packages/pi/src/stream.ts b/packages/pi/src/stream.ts index e4d221ba..82f8c49a 100644 --- a/packages/pi/src/stream.ts +++ b/packages/pi/src/stream.ts @@ -4,8 +4,11 @@ import { ANTIGRAVITY_ENDPOINT, buildAgyAgentRequestMetadata, buildAntigravityHarnessUserAgent, + buildFingerprintHeaders, ensureProjectContext, fetchWithAgyCliTransport, + type ManagedAccount, + type OAuthAuthDetails, orderAgyRequestPayloadInPlace, resolveModelForHeaderStyle, } from '@cortexkit/antigravity-auth-core' @@ -27,6 +30,7 @@ import { import { buildGeminiRequest } from './convert.ts' import { getPackedRefresh } from './credential-cache.ts' +import type { PiAccountRuntime } from './runtime.ts' const STREAM_ACTION = 'streamGenerateContent' const FALLBACK_SESSION_KEY = '__default__' @@ -339,7 +343,9 @@ async function sendAntigravityRequest(options: { context: Context streamOptions?: SimpleStreamOptions accessToken: string - sessionKey: string + auth?: OAuthAuthDetails + account?: ManagedAccount + requestScope: AgyRequestScope signal?: AbortSignal }): Promise { const resolved = resolvePiAntigravityModel( @@ -353,12 +359,14 @@ async function sendAntigravityRequest(options: { // With it, ensureProjectContext returns the cached managedProjectId directly // instead of re-running loadCodeAssist every turn. const packedRefresh = getPackedRefresh(options.accessToken) ?? '' - const projectContext = await ensureProjectContext({ - type: 'oauth', - refresh: packedRefresh, - access: options.accessToken, - expires: Date.now() + 60_000, - }) + const projectContext = await ensureProjectContext( + options.auth ?? { + type: 'oauth', + refresh: packedRefresh, + access: options.accessToken, + expires: Date.now() + 60_000, + }, + ) const request = buildGeminiRequest(options.context, { provider: options.model.provider, @@ -387,11 +395,10 @@ async function sendAntigravityRequest(options: { request.generationConfig = generationConfig } - const requestScope = requestSessions.beginRequest(options.sessionKey) const requestId = finalizePiAntigravityRequest( request, wireModel, - requestScope, + options.requestScope, ) const envelope = { @@ -413,6 +420,9 @@ async function sendAntigravityRequest(options: { Authorization: `Bearer ${options.accessToken}`, 'Content-Type': 'application/json', 'User-Agent': buildAntigravityHarnessUserAgent(), + ...(options.account?.fingerprint + ? buildFingerprintHeaders(options.account.fingerprint) + : {}), 'Accept-Encoding': 'gzip', }, body: JSON.stringify(envelope), @@ -425,6 +435,7 @@ export function streamCortexKitAntigravity( model: Model, context: Context, options?: SimpleStreamOptions, + runtime?: PiAccountRuntime, ): AssistantMessageEventStream { const stream = createAssistantMessageEventStream() @@ -434,6 +445,7 @@ export function streamCortexKitAntigravity( let response: Response | undefined let requestAbort: AbortController | undefined let chunkIterator: AsyncIterator | undefined + let selectedAccount: ManagedAccount | undefined try { const accessToken = options?.apiKey ?? '' @@ -441,18 +453,31 @@ export function streamCortexKitAntigravity( throw new Error('Missing Antigravity OAuth access token') const sessionKey = getRequestSessionKey(context, options) + const requestScope = requestSessions.beginRequest(sessionKey) requestAbort = new AbortController() const requestSignal = options?.signal ? AbortSignal.any([options.signal, requestAbort.signal]) : requestAbort.signal - response = await sendAntigravityRequest({ - model, - context, - streamOptions: options, - accessToken, - sessionKey, - signal: requestSignal, - }) + const send = (auth?: OAuthAuthDetails, account?: ManagedAccount) => + sendAntigravityRequest({ + model, + context, + streamOptions: options, + accessToken: auth?.access ?? accessToken, + auth, + account, + requestScope, + signal: requestSignal, + }) + if (runtime) { + const result = await runtime.dispatch( + resolvePiAntigravityModel(model, options?.reasoning).actualModel, + send, + requestSignal, + ) + response = result.response + selectedAccount = result.account + } else response = await send() if (!response.ok) { throw new Error( @@ -661,6 +686,7 @@ export function streamCortexKitAntigravity( ) } + if (selectedAccount) runtime?.complete(selectedAccount, true) stream.push({ type: 'done', reason: output.stopReason as 'stop' | 'length' | 'toolUse', @@ -671,6 +697,8 @@ export function streamCortexKitAntigravity( } stream.end() } catch (error) { + if (selectedAccount && !options?.signal?.aborted) + runtime?.complete(selectedAccount, false) requestAbort?.abort() await chunkIterator?.return?.(undefined).catch(() => {}) await response?.body?.cancel().catch(() => {}) From 21ac5f87e05da75f412256bf83974d38c9946bf4 Mon Sep 17 00:00:00 2001 From: Brent Duarte Date: Wed, 9 Sep 2026 16:50:26 -0700 Subject: [PATCH 3/7] fix: preserve account identity across Pi retries and reconciliation --- packages/core/src/account-manager.test.ts | 125 +++++++++++++++++++++- packages/core/src/account-manager.ts | 57 ++++++++-- packages/core/src/rotation.ts | 20 ++++ packages/pi/src/runtime.test.ts | 112 +++++++++++++++++++ packages/pi/src/runtime.ts | 114 ++++++++++++++------ 5 files changed, 383 insertions(+), 45 deletions(-) diff --git a/packages/core/src/account-manager.test.ts b/packages/core/src/account-manager.test.ts index 8b179d6d..ba8c71f2 100644 --- a/packages/core/src/account-manager.test.ts +++ b/packages/core/src/account-manager.test.ts @@ -1,7 +1,8 @@ -import { describe, expect, it } from 'bun:test' +import { afterEach, describe, expect, it } from 'bun:test' import { AccountManager } from './account-manager.ts' import type { AccountStorageStore } from './account-storage.ts' import type { AccountStorageV4 } from './account-types.ts' +import { initHealthTracker, initTokenTracker } from './rotation.ts' function createStore(initial: AccountStorageV4 | null = null) { let state = initial @@ -42,6 +43,128 @@ const stored: AccountStorageV4 = { } describe('core AccountManager', () => { + afterEach(() => { + initHealthTracker({}) + initTokenTracker({}) + }) + + it('clears transient penalties when a no-email account is replaced', () => { + const health = initHealthTracker({}) + const tokens = initTokenTracker({}) + const manager = new AccountManager(undefined, structuredClone(stored), { + store: createStore(stored).store, + persistFingerprintUpdates: false, + }) + const initialScore = health.getScore(0) + const initialTokens = tokens.getTokens(0) + health.recordFailure(0) + health.recordFailure(0) + tokens.consume(0, 3) + manager.recordSessionUsage(0) + const next = structuredClone(stored) + next.accounts[0]!.refreshToken = 'unknown-replacement' + manager.reconcileStorage(next) + expect(health.getScore(0)).toBe(initialScore) + expect(tokens.getTokens(0)).toBe(initialTokens) + expect(manager.wasUsedInSession(0)).toBe(false) + }) + + it.each([ + { order: [1, 2], next: 'b' }, + { order: [0, 2], next: 'c' }, + { order: [0, 1], next: 'b' }, + { order: [2, 0, 1], next: 'b' }, + ])('keeps the next surviving RR account across $order', ({ order, next }) => { + for (const identity of [undefined, { id: 'request' }]) { + const pool: AccountStorageV4 = { + version: 4, + activeIndex: 0, + accounts: ['a', 'b', 'c'].map((refreshToken) => ({ + refreshToken, + email: `${refreshToken}@example.com`, + addedAt: 1, + lastUsed: 0, + })), + } + const manager = new AccountManager(undefined, structuredClone(pool), { + store: createStore(pool).store, + persistFingerprintUpdates: false, + }) + const select = () => + manager.getNextForFamily( + 'gemini', + null, + 'antigravity', + 100, + 60_000, + identity, + )?.parts.refreshToken + expect(select()).toBe('a') + manager.reconcileStorage({ + ...pool, + accounts: order.map((index) => pool.accounts[index]!), + }) + expect(select()).toBe(next) + } + }) + + it('remaps usage, health and balances by unique identity through removal, reorder and token rotation', () => { + const health = initHealthTracker({ recoveryRatePerHour: 0 }) + const tokens = initTokenTracker({ regenerationRatePerMinute: 0 }) + const pool: AccountStorageV4 = { + version: 4, + activeIndex: 0, + accounts: ['a', 'b', 'c'].map((refreshToken) => ({ + refreshToken, + email: `${refreshToken}@example.com`, + addedAt: 1, + lastUsed: 0, + })), + } + const manager = new AccountManager(undefined, structuredClone(pool), { + store: createStore(pool).store, + persistFingerprintUpdates: false, + }) + const session = { id: 'request' } + manager.recordSessionUsage(0) + manager.recordSessionUsage(1, session) + health.recordFailure(0) + health.recordFailure(0) + health.recordSuccess(1) + health.recordRateLimit(2) + tokens.consume(0, 3) + tokens.consume(1, 1) + tokens.consume(2, 2) + const expected = [1, 2].map((index) => ({ + score: health.getScore(index), + tokens: tokens.getTokens(index), + })) + manager.reconcileStorage({ + ...pool, + accounts: [pool.accounts[1]!, pool.accounts[2]!], + }) + expect(health.getScore(0)).toBe(expected[0]!.score) + expect(tokens.getTokens(0)).toBe(expected[0]!.tokens) + manager.reconcileStorage({ + ...pool, + accounts: [ + pool.accounts[2]!, + { + ...pool.accounts[1]!, + email: ' B@EXAMPLE.COM ', + refreshToken: 'rotated', + }, + ], + }) + for (const [index, original] of [1, 0].entries()) { + expect(health.getScore(index)).toBe(expected[original]!.score) + expect(tokens.getTokens(index)).toBe(expected[original]!.tokens) + } + expect(manager.wasUsedInSession(0)).toBe(false) + expect(manager.wasUsedInSession(1, session)).toBe(true) + expect(manager.wasUsedInSession(0, session)).toBe(false) + }) + it.each([ 'sticky', 'hybrid', diff --git a/packages/core/src/account-manager.ts b/packages/core/src/account-manager.ts index 190ee534..2c160138 100644 --- a/packages/core/src/account-manager.ts +++ b/packages/core/src/account-manager.ts @@ -548,8 +548,8 @@ export class AccountManager { } /** Reload durable metadata without resetting this process's routing state. - * Access tokens and failure counters are transient and survive only an exact - * refresh-token match. Removed accounts are never resurrected. + * Access tokens survive only an exact credential match. Routing and scoring + * follow unique normalized email, falling back to token for legacy accounts. */ reconcileStorage(stored: AccountStorageV4): void { const previous = this.accounts @@ -563,13 +563,32 @@ export class AccountManager { pid: this.pid, persistFingerprintUpdates: false, }) + const keyOf = (account: ManagedAccount): string => { + const email = account.email?.trim().toLowerCase() + return email ? `email:${email}` : `token:${account.parts.refreshToken}` + } + const remap = (index: number): number => { + const old = previous[index] + if (!old) return -1 + const key = keyOf(old) + if (previous.filter((account) => keyOf(account) === key).length !== 1) + return -1 + const matches = fresh.accounts.filter((account) => keyOf(account) === key) + return matches.length === 1 ? matches[0]!.index : -1 + } + const indexMap = new Map( + previous.map((account) => [account.index, remap(account.index)]), + ) this.accounts = fresh.accounts.map((account) => { - const old = byToken.get(account.parts.refreshToken) + const old = previous.find( + (entry) => indexMap.get(entry.index) === account.index, + ) if (!old) return account + const credential = byToken.get(account.parts.refreshToken) return { ...account, - access: old.access, - expires: old.expires, + access: credential?.access, + expires: credential?.expires, fingerprint: stored.accounts[account.index]?.fingerprint ?? old.fingerprint ?? @@ -579,19 +598,37 @@ export class AccountManager { lastFailureTime: old.lastFailureTime, } }) - const remap = (index: number): number => { - const token = previous[index]?.parts.refreshToken - return this.accounts.findIndex( - (account) => account.parts.refreshToken === token, - ) + const remapCursor = (cursor: number): number => { + for (let offset = 0; offset < previous.length; offset++) { + const next = remap((cursor + offset) % previous.length) + if (next >= 0) return next + } + return 0 + } + const remapUsed = (used: Set): Set => + new Set([...used].map(remap).filter((index) => index >= 0)) + const changed = + previous.length !== this.accounts.length || + previous.some((account) => remap(account.index) !== account.index) + if (changed) { + getHealthTracker().remapAccounts(indexMap) + getTokenTracker().remapAccounts(indexMap) + this.sessionUsedAccounts = remapUsed(this.sessionUsedAccounts) } for (const family of ['claude', 'gemini'] as const) { + if (changed) + this.cursorByFamily[family] = remapCursor(this.cursorByFamily[family]) this.currentAccountIndexByFamily[family] = remap( this.currentAccountIndexByFamily[family], ) } for (const state of this.requestSessionStates.values()) { + if (changed) state.usedAccounts = remapUsed(state.usedAccounts) for (const family of ['claude', 'gemini'] as const) { + if (changed) + state.cursorByFamily[family] = remapCursor( + state.cursorByFamily[family], + ) state.currentAccountIndexByFamily[family] = remap( state.currentAccountIndexByFamily[family], ) diff --git a/packages/core/src/rotation.ts b/packages/core/src/rotation.ts index 082ce452..22b963cd 100644 --- a/packages/core/src/rotation.ts +++ b/packages/core/src/rotation.ts @@ -257,6 +257,16 @@ export class HealthScoreTracker { this.scores.delete(accountIndex) } + /** Move transient state with surviving accounts, dropping unknown identities. */ + remapAccounts(indexMap: ReadonlyMap): void { + const previous = new Map(this.scores) + this.scores.clear() + for (const [oldIndex, newIndex] of indexMap) { + const state = previous.get(oldIndex) + if (newIndex >= 0 && state) this.scores.set(newIndex, state) + } + } + /** * Get all scores for debugging/logging. */ @@ -546,6 +556,16 @@ export class TokenBucketTracker { getMaxTokens(): number { return this.config.maxTokens } + + /** Move balances with surviving accounts, dropping unknown identities. */ + remapAccounts(indexMap: ReadonlyMap): void { + const previous = new Map(this.buckets) + this.buckets.clear() + for (const [oldIndex, newIndex] of indexMap) { + const state = previous.get(oldIndex) + if (newIndex >= 0 && state) this.buckets.set(newIndex, state) + } + } } // ============================================================================ diff --git a/packages/pi/src/runtime.test.ts b/packages/pi/src/runtime.test.ts index 4d286b68..380d64b9 100644 --- a/packages/pi/src/runtime.test.ts +++ b/packages/pi/src/runtime.test.ts @@ -5,6 +5,7 @@ import { join } from 'node:path' import { type AntigravityRefreshResult, getHealthTracker, + getTokenTracker, initHealthTracker, initTokenTracker, loadAccountStorage, @@ -105,6 +106,117 @@ afterEach(async () => { }) describe('Pi shared account runtime', () => { + it('attributes in-flight health and refunds after a local reconciliation reorders accounts', async () => { + const result = await pool(2) + initTokenTracker({ regenerationRatePerMinute: 0 }) + const balance = getTokenTracker().getTokens(0) + const score = getHealthTracker().getScore(0) + let calls = 0 + const response = await result.dispatch(model, async () => { + if (++calls > 1) return new Response('ok') + await mutateAccountStorage(path, (current) => { + current.accounts[0]!.refreshToken = 'rotated' + current.accounts.reverse() + return current + }) + await result.describe() + return new Response('', { status: 429 }) + }) + expect(response.account.email).toBe('user2@example.com') + expect(getHealthTracker().getScore(0)).toBe(score) + expect(getHealthTracker().getScore(1)).toBeLessThan(score) + expect(getTokenTracker().getTokens(0)).toBe(balance - 1) + expect(getTokenTracker().getTokens(1)).toBe(balance) + }) + + it.each([ + 429, 500, 503, 529, + ])('attributes in-flight HTTP %s to the account after credential rotation', async (status) => { + const result = await pool(2) + const sends: string[] = [] + await result.dispatch(model, async (_auth, account) => { + sends.push(account.parts.refreshToken) + if (sends.length > 1) return new Response('ok') + await mutateAccountStorage(path, (current) => { + current.accounts[0]!.refreshToken = 'rotated' + current.accounts[0]!.email = ' USER1@EXAMPLE.COM ' + return current + }) + return new Response('', { status }) + }) + expect(sends).toEqual(['secret-refresh-1', 'secret-refresh-2']) + const account = (await loadAccountStorage(path))!.accounts[0]! + expect(account.refreshToken).toBe('rotated') + expect( + account.rateLimitResetTimes?.[`gemini-antigravity:${model}`], + ).toBeGreaterThan(Date.now()) + }) + + it('retains attempted identity even when a peer clears the persisted cooldown', async () => { + const result = await pool(2) + const sends: string[] = [] + await expect( + result.dispatch(model, async (_auth, account) => { + sends.push(account.parts.refreshToken) + await mutateAccountStorage(path, (current) => { + current.accounts[0]!.refreshToken = 'rotated' + current.accounts[0]!.rateLimitResetTimes = {} + return current + }) + return new Response('', { status: 429 }) + }), + ).rejects.toThrow('HTTP 429') + expect(sends).toEqual(['secret-refresh-1', 'secret-refresh-2']) + }) + + it('stops failover when a dispatched no-email credential disappears', async () => { + const result = await pool(2) + await mutateAccountStorage(path, (current) => { + delete current.accounts[0]!.email + return current + }) + const send = mock(async () => { + await mutateAccountStorage(path, (current) => { + current.accounts[0]!.refreshToken = 'unknown-replacement' + return current + }) + return new Response('', { status: 429 }) + }) + await expect(result.dispatch(model, send)).rejects.toThrow( + 'cooldown not recorded', + ) + expect(send).toHaveBeenCalledTimes(1) + expect( + (await loadAccountStorage(path))!.accounts[0]!.rateLimitResetTimes, + ).toBeUndefined() + }) + + it.each([ + true, + false, + ])('reports whether in-flight quota can be attributed with email=%s', async (withEmail) => { + const result = await pool(1) + const before = 1 + await mutateAccountStorage(path, (current) => { + current.accounts[0]!.cachedQuotaUpdatedAt = before + if (!withEmail) delete current.accounts[0]!.email + return current + }) + const response = quotaFetch.getMockImplementation()! + quotaFetch.mockImplementation(async () => { + await mutateAccountStorage(path, (current) => { + current.accounts[0]!.refreshToken = 'quota-rotated' + return current + }) + return response() + }) + expect(await result.refreshQuota(true)).toBe(withEmail ? 0 : 1) + const account = (await loadAccountStorage(path))!.accounts[0]! + expect(account.refreshToken).toBe('quota-rotated') + if (withEmail) expect(account.cachedQuotaUpdatedAt).toBeGreaterThan(before) + else expect(account.cachedQuotaUpdatedAt).toBe(before) + }) + it('persists first, second and third logins using v4 and secure permissions', async () => { const result = runtime() for (let i = 1; i <= 3; i++) { diff --git a/packages/pi/src/runtime.ts b/packages/pi/src/runtime.ts index b2dcb225..81a609f0 100644 --- a/packages/pi/src/runtime.ts +++ b/packages/pi/src/runtime.ts @@ -42,6 +42,15 @@ type LoginResult = Extract class AccountUnavailable extends Error {} class CredentialRefreshFailed extends Error {} +function attemptKey(account: ManagedAccount): string { + return defaultKeyOf({ + ...account.parts, + email: account.email?.trim().toLowerCase(), + addedAt: account.addedAt, + lastUsed: account.lastUsed, + }) +} + export interface PiRuntimeOptions { path?: string pid?: number @@ -109,13 +118,17 @@ export class PiAccountRuntime { quota = aggregateQuota((await fetchAvailableModels(common)).models) } signal.throwIfAborted() - await this.patch(parts.refreshToken, (current) => { - current.projectId = parts.projectId ?? current.projectId - current.managedProjectId = - parts.managedProjectId ?? current.managedProjectId - current.cachedQuota = quota.groups - current.cachedQuotaUpdatedAt = Date.now() - }) + const applied = await this.patch( + { ...account, refreshToken: parts.refreshToken }, + (current) => { + current.projectId = parts.projectId ?? current.projectId + current.managedProjectId = + parts.managedProjectId ?? current.managedProjectId + current.cachedQuota = quota.groups + current.cachedQuotaUpdatedAt = Date.now() + }, + ) + if (!applied) throw new AccountUnavailable('Quota account changed') return { index: 0, status: 'ok', quota } } catch { // Provider bodies can contain credentials. Never relay them to quota @@ -196,16 +209,24 @@ export class PiAccountRuntime { } private async patch( - token: string, + target: { email?: string; refreshToken: string }, update: (account: AccountMetadataV3) => void, - ): Promise { + ): Promise { + let applied = false await mutateAccountStorage(this.path, (current) => { - const account = current.accounts.find( - (entry) => entry.refreshToken === token, + const email = target.email?.trim().toLowerCase() + const matches = current.accounts.filter((entry) => + email + ? entry.email?.trim().toLowerCase() === email + : entry.refreshToken === target.refreshToken, ) - if (account) update(account) + if (matches.length === 1 && matches[0]) { + update(matches[0]) + applied = true + } return current }) + return applied } private async credentialFor(token: string): Promise { @@ -291,15 +312,21 @@ export class PiAccountRuntime { } private async authFailure(account: ManagedAccount): Promise { - getHealthTracker().recordFailure(account.index) + const current = this.currentAccount(account) + if (current) getHealthTracker().recordFailure(current.index) this.manager?.markAccountCoolingDown(account, 60_000, 'auth-failure') - await this.patch(account.parts.refreshToken, (current) => { - current.coolingDownUntil = Math.max( - current.coolingDownUntil ?? 0, - account.coolingDownUntil ?? 0, - ) - current.cooldownReason = 'auth-failure' - }) + const applied = await this.patch( + { ...account.parts, email: account.email }, + (current) => { + current.coolingDownUntil = Math.max( + current.coolingDownUntil ?? 0, + account.coolingDownUntil ?? 0, + ) + current.cooldownReason = 'auth-failure' + }, + ) + if (!applied) + throw new AccountUnavailable('Account changed; cooldown not recorded') } async refreshQuota(force = false): Promise { @@ -330,7 +357,7 @@ export class PiAccountRuntime { const config = await readSettings(this.settingsPath) const attempted = new Set() // One extra selection tolerates a peer rotating a token between reload and - // credential acquisition. The attempted-token set still bounds sends. + // credential acquisition. Stable attempted identities still bound sends. const budget = (await this.reload()).getTotalAccountCount() + 1 let lastError = 'All Antigravity accounts are disabled, cooling down, or over cached quota' @@ -344,7 +371,7 @@ export class PiAccountRuntime { .getAccounts() .filter( (a) => - attempted.has(a.parts.refreshToken) || + attempted.has(attemptKey(a)) || a.verificationRequired || a.accountIneligible, ) @@ -362,7 +389,6 @@ export class PiAccountRuntime { excluded, ) if (!account) break - attempted.add(account.parts.refreshToken) let auth: OAuthAuthDetails try { auth = await this.credentialFor(account.parts.refreshToken) @@ -375,26 +401,34 @@ export class PiAccountRuntime { ) throw error lastError = error.message - if (error instanceof CredentialRefreshFailed) + if (error instanceof CredentialRefreshFailed) { + attempted.add(attemptKey(account)) await this.authFailure(account) + } continue } manager.updateFromAuth(account, auth) const token = account.parts.refreshToken - attempted.add(token) + attempted.add(attemptKey(account)) + const target = { ...account.parts, email: account.email } signal?.throwIfAborted() manager.markAccountUsed(account.index) - await this.patch(token, (current) => { + const located = await this.patch(target, (current) => { current.lastUsed = Math.max(current.lastUsed, account.lastUsed) current.fingerprint ??= account.fingerprint }) + if (!located) + throw new AccountUnavailable('Account changed before dispatch') const consumed = getTokenTracker().consume(account.index) let response: Response try { response = await send(auth, account) } catch (error) { - if (consumed) getTokenTracker().refund(account.index) - if (!signal?.aborted) getHealthTracker().recordFailure(account.index) + const current = this.currentAccount(account) + if (current) { + if (consumed) getTokenTracker().refund(current.index) + if (!signal?.aborted) getHealthTracker().recordFailure(current.index) + } // Unknown transport failures may occur after upstream accepted work. // Surface them without replaying a potentially billable request. throw error @@ -403,7 +437,8 @@ export class PiAccountRuntime { this.selected = token return { response, account } } - if (consumed) getTokenTracker().refund(account.index) + const current = this.currentAccount(account) + if (consumed && current) getTokenTracker().refund(current.index) if (![429, 503, 529, 500].includes(response.status)) return { response, account } let body: unknown @@ -421,8 +456,9 @@ export class PiAccountRuntime { parseRateLimitReason(info.reason, info.message, response.status), info.retryDelayMs ?? retryAfterMsFromResponse(response), ) - getHealthTracker().recordRateLimit(account.index) - await this.patch(token, (current) => { + const rateLimited = this.currentAccount(account) + if (rateLimited) getHealthTracker().recordRateLimit(rateLimited.index) + const applied = await this.patch(target, (current) => { for (const [key, until] of Object.entries( account.rateLimitResetTimes, )) { @@ -433,15 +469,25 @@ export class PiAccountRuntime { ) } }) + if (!applied) + throw new AccountUnavailable( + `Antigravity HTTP ${response.status}; account changed; cooldown not recorded`, + ) lastError = `Antigravity HTTP ${response.status}; no eligible account remains (cooldown recorded)` } throw new Error(lastError) } + private currentAccount(account: ManagedAccount): ManagedAccount | undefined { + const matches = + this.manager + ?.getAccounts() + .filter((entry) => attemptKey(entry) === attemptKey(account)) ?? [] + return matches.length === 1 ? matches[0] : undefined + } + complete(account: ManagedAccount, success: boolean): void { - const current = this.manager - ?.getAccounts() - .find((entry) => entry.parts.refreshToken === account.parts.refreshToken) + const current = this.currentAccount(account) if (!current) return if (success) { this.manager?.markRequestSuccess(current) From 372e2c7fc92fdcfb90d02f162f5bce85139b99e4 Mon Sep 17 00:00:00 2001 From: Brent Duarte Date: Wed, 9 Sep 2026 17:06:39 -0700 Subject: [PATCH 4/7] fix(core): make shared tracker reconciliation idempotent --- packages/core/src/account-manager.test.ts | 221 ++++++++++++++++++++++ packages/core/src/account-manager.ts | 10 +- packages/core/src/rotation.ts | 47 +++++ 3 files changed, 276 insertions(+), 2 deletions(-) diff --git a/packages/core/src/account-manager.test.ts b/packages/core/src/account-manager.test.ts index ba8c71f2..b34e5336 100644 --- a/packages/core/src/account-manager.test.ts +++ b/packages/core/src/account-manager.test.ts @@ -48,6 +48,227 @@ describe('core AccountManager', () => { initTokenTracker({}) }) + function sharedManagers(emails = true) { + const health = initHealthTracker({ recoveryRatePerHour: 0 }) + const tokens = initTokenTracker({ regenerationRatePerMinute: 0 }) + const pool: AccountStorageV4 = { + version: 4, + activeIndex: 0, + accounts: ['a', 'b', 'c'].map((refreshToken) => ({ + refreshToken, + email: emails ? `${refreshToken}@example.com` : undefined, + addedAt: 1, + lastUsed: 0, + })), + } + const memory = createStore(pool) + const managers = Array.from( + { length: 3 }, + () => + new AccountManager(undefined, structuredClone(pool), { + store: memory.store, + persistFingerprintUpdates: false, + now: () => 10_000, + }), + ) + const values = () => + [0, 1, 2].map((index) => [ + health.getScore(index), + tokens.getTokens(index), + ]) + return { health, tokens, pool, managers, values } + } + + it.each([ + true, + false, + ])('keeps shared tracker ownership across three stale managers (emails: %s)', (emails) => { + const { health, tokens, pool, managers, values } = sharedManagers(emails) + health.recordFailure(0) + health.recordFailure(0) + tokens.consume(0, 7) + const next = { + ...pool, + accounts: [pool.accounts[1]!, pool.accounts[2]!, pool.accounts[0]!], + } + for (const manager of [...managers, ...managers]) { + manager.reconcileStorage(next) + expect(values()).toEqual([ + [70, 50], + [70, 50], + [30, 43], + ]) + expect( + manager.getCurrentOrNextForFamily( + 'gemini', + null, + 'hybrid', + 'antigravity', + false, + )?.parts.refreshToken, + ).toBe('b') + } + }) + + it('keeps depleted accounts out of hybrid selection after stale reconciliation', () => { + const { tokens, pool, managers, values } = sharedManagers() + tokens.consume(0, 50) + const next = { + ...pool, + accounts: [pool.accounts[1]!, pool.accounts[2]!, pool.accounts[0]!], + } + for (const manager of managers) { + manager.reconcileStorage(next) + expect(values()).toEqual([ + [70, 50], + [70, 50], + [70, 0], + ]) + expect( + manager.getCurrentOrNextForFamily( + 'gemini', + null, + 'hybrid', + 'antigravity', + false, + )?.parts.refreshToken, + ).toBe('b') + } + }) + + it('uses shared ownership when stale managers skip an intermediate layout', () => { + const { health, tokens, pool, managers, values } = sharedManagers() + health.recordFailure(0) + tokens.consume(0, 7) + managers[0]!.reconcileStorage({ + ...pool, + accounts: [pool.accounts[1]!, pool.accounts[2]!, pool.accounts[0]!], + }) + const next = { + ...pool, + accounts: [pool.accounts[2]!, pool.accounts[0]!, pool.accounts[1]!], + } + for (const manager of managers) { + manager.reconcileStorage(next) + expect(values()).toEqual([ + [70, 50], + [50, 43], + [70, 50], + ]) + } + }) + + it.each([ + 'health', + 'tokens', + ] as const)('reinitializing %s does not lose the other shared tracker layout', (reset) => { + const { health, tokens, pool, managers } = sharedManagers() + health.recordFailure(0) + tokens.consume(0, 7) + const next = { + ...pool, + accounts: [pool.accounts[1]!, pool.accounts[2]!, pool.accounts[0]!], + } + managers[0]!.reconcileStorage(next) + if (reset === 'health') initHealthTracker({ recoveryRatePerHour: 0 }) + else initTokenTracker({ regenerationRatePerMinute: 0 }) + for (const manager of managers.slice(1)) { + manager.reconcileStorage(next) + if (reset === 'health') { + expect([0, 1, 2].map((index) => tokens.getTokens(index))).toEqual([ + 50, 50, 43, + ]) + } else { + expect([0, 1, 2].map((index) => health.getScore(index))).toEqual([ + 70, 70, 50, + ]) + } + } + }) + + it('preserves every survivor through multiple shared tracker layout changes', () => { + const { health, tokens, pool, managers, values } = sharedManagers() + health.recordFailure(0) + health.recordSuccess(1) + health.recordRateLimit(2) + tokens.consume(0, 7) + tokens.consume(1, 11) + tokens.consume(2, 19) + const original = values() + for (const order of [ + [1, 2, 0], + [2, 0, 1], + ]) { + const next = { + ...pool, + accounts: order.map((index) => ({ + ...pool.accounts[index]!, + email: ` ${pool.accounts[index]!.email!.toUpperCase()} `, + refreshToken: `rotated-${index}`, + })), + } + for (const manager of managers) { + manager.reconcileStorage(next) + expect(values()).toEqual(order.map((index) => original[index]!)) + } + } + }) + + it('drops removed state across stale managers and does not reuse it for additions', () => { + const { health, tokens, pool, managers, values } = sharedManagers() + health.recordFailure(0) + health.recordFailure(0) + tokens.consume(0, 7) + const removed = { ...pool, accounts: pool.accounts.slice(1) } + for (const manager of managers) { + manager.reconcileStorage(removed) + expect(values()).toEqual([ + [70, 50], + [70, 50], + [70, 50], + ]) + expect(health.getSnapshot().size).toBe(0) + } + health.recordRateLimit(0) + tokens.consume(0, 11) + const added = { + ...pool, + accounts: [ + { ...pool.accounts[0]!, email: 'd@example.com', refreshToken: 'd' }, + ...removed.accounts, + ], + } + for (const manager of managers) { + manager.reconcileStorage(added) + expect(values()).toEqual([ + [70, 50], + [60, 39], + [70, 50], + ]) + } + }) + + it.each([ + true, + false, + ])('clears ambiguous shared identities across stale managers (emails: %s)', (emails) => { + const { health, tokens, pool, managers, values } = sharedManagers(emails) + health.recordFailure(0) + tokens.consume(0, 7) + const ambiguous = { + ...pool, + accounts: [pool.accounts[0]!, pool.accounts[0]!, pool.accounts[2]!], + } + for (const manager of managers) { + manager.reconcileStorage(ambiguous) + expect(values()).toEqual([ + [70, 50], + [70, 50], + [70, 50], + ]) + } + }) + it('clears transient penalties when a no-email account is replaced', () => { const health = initHealthTracker({}) const tokens = initTokenTracker({}) diff --git a/packages/core/src/account-manager.ts b/packages/core/src/account-manager.ts index 2c160138..fd22098f 100644 --- a/packages/core/src/account-manager.ts +++ b/packages/core/src/account-manager.ts @@ -29,6 +29,7 @@ import { type AccountWithMetrics, getHealthTracker, getTokenTracker, + reconcileAccountTrackers, selectHybridAccount, } from './rotation.ts' @@ -610,9 +611,14 @@ export class AccountManager { const changed = previous.length !== this.accounts.length || previous.some((account) => remap(account.index) !== account.index) + // Retain opaque identities in process-global state, never raw credentials. + const trackerIdentity = (account: ManagedAccount): string => + createHash('sha256').update(keyOf(account)).digest('hex') + reconcileAccountTrackers( + previous.map(trackerIdentity), + this.accounts.map(trackerIdentity), + ) if (changed) { - getHealthTracker().remapAccounts(indexMap) - getTokenTracker().remapAccounts(indexMap) this.sessionUsedAccounts = remapUsed(this.sessionUsedAccounts) } for (const family of ['claude', 'gemini'] as const) { diff --git a/packages/core/src/rotation.ts b/packages/core/src/rotation.ts index 22b963cd..58681f9c 100644 --- a/packages/core/src/rotation.ts +++ b/packages/core/src/rotation.ts @@ -572,6 +572,53 @@ export class TokenBucketTracker { // SINGLETON TRACKERS // ============================================================================ +// Ownership belongs to each shared tracker instance, never to a manager's +// potentially stale index layout. Weak keys also make tracker reinitialization +// discard the corresponding layout without resetting the other tracker. +const trackerLayouts = new WeakMap< + HealthScoreTracker | TokenBucketTracker, + readonly (string | null)[] +>() + +/** Reconcile both global trackers from their own last known identity layout. */ +export function reconcileAccountTrackers( + previousIdentities: readonly string[], + nextIdentities: readonly string[], +): void { + const uniqueLayout = (identities: readonly string[]) => { + const counts = new Map() + for (const identity of identities) { + counts.set(identity, (counts.get(identity) ?? 0) + 1) + } + return identities.map((identity) => + counts.get(identity) === 1 ? identity : null, + ) + } + const next = uniqueLayout(nextIdentities) + for (const tracker of [getHealthTracker(), getTokenTracker()]) { + const previous = + trackerLayouts.get(tracker) ?? uniqueLayout(previousIdentities) + // A stale manager targeting an already-applied layout is a no-op. Null + // identities must still be cleared, even if the positional layout matches. + if ( + previous.length !== next.length || + previous.some( + (identity, index) => identity === null || identity !== next[index], + ) + ) { + const indexMap = new Map() + for (const [index, identity] of previous.entries()) { + if (identity !== null) { + const nextIndex = next.indexOf(identity) + if (nextIndex >= 0) indexMap.set(index, nextIndex) + } + } + tracker.remapAccounts(indexMap) + } + trackerLayouts.set(tracker, next) + } +} + let globalTokenTracker: TokenBucketTracker | null = null export function getTokenTracker(): TokenBucketTracker { From 83530f4a2aec5b0d8911c79c7eade22ce01bf07b Mon Sep 17 00:00:00 2001 From: Brent Duarte Date: Wed, 9 Sep 2026 22:14:55 -0700 Subject: [PATCH 5/7] fix(pi): support runtime 0.85 credential APIs --- bun.lock | 126 ++++++++++++++++++------- packages/core/src/antigravity/oauth.ts | 2 + packages/pi/package.json | 12 +-- packages/pi/src/commands.ts | 19 +++- packages/pi/src/extension-load.test.ts | 111 ++++++++++++++++++++++ packages/pi/src/index.ts | 28 +++--- packages/pi/src/provider.test.ts | 38 ++++++-- packages/pi/src/runtime.test.ts | 37 +++++++- packages/pi/src/runtime.ts | 26 ++++- packages/pi/src/stream.test.ts | 29 ++++++ packages/pi/src/stream.ts | 15 ++- 11 files changed, 372 insertions(+), 71 deletions(-) create mode 100644 packages/pi/src/extension-load.test.ts diff --git a/bun.lock b/bun.lock index 0c6d2291..bee8458a 100644 --- a/bun.lock +++ b/bun.lock @@ -19,7 +19,7 @@ }, "packages/core": { "name": "@cortexkit/antigravity-auth-core", - "version": "2.1.0", + "version": "2.2.1", "dependencies": { "xdg-basedir": "^5.1.0", "zod": "^4.0.0", @@ -40,7 +40,7 @@ }, "packages/opencode": { "name": "@cortexkit/opencode-antigravity-auth", - "version": "2.1.0", + "version": "2.2.1", "bin": { "antigravity-auth": "./dist/cli.js", }, @@ -74,20 +74,20 @@ }, "packages/pi": { "name": "@cortexkit/pi-antigravity-auth", - "version": "2.1.0", + "version": "2.2.1", "dependencies": { - "@cortexkit/antigravity-auth-core": "2.0.0", + "@cortexkit/antigravity-auth-core": "2.2.1", }, "devDependencies": { - "@earendil-works/pi-ai": "^0.79.1", - "@earendil-works/pi-coding-agent": "^0.79.1", - "@earendil-works/pi-tui": "^0.79.1", + "@earendil-works/pi-ai": "0.85.1", + "@earendil-works/pi-coding-agent": "0.85.1", + "@earendil-works/pi-tui": "0.85.1", "esbuild": "^0.27.7", }, "peerDependencies": { - "@earendil-works/pi-ai": "*", - "@earendil-works/pi-coding-agent": "*", - "@earendil-works/pi-tui": "*", + "@earendil-works/pi-ai": "^0.85.1", + "@earendil-works/pi-coding-agent": "^0.85.1", + "@earendil-works/pi-tui": "^0.85.1", }, }, }, @@ -96,7 +96,7 @@ "@ampproject/remapping": ["@ampproject/remapping@2.3.0", "", { "dependencies": { "@jridgewell/gen-mapping": "^0.3.5", "@jridgewell/trace-mapping": "^0.3.24" } }, "sha512-30iZtAPgz+LTIYoeivqYo853f02jBYSd5uGnGpkFV0M3xOt9aN73erkgYAmZU43x4VfqcnLxW9Kpg3R5LC4YYw=="], - "@anthropic-ai/sdk": ["@anthropic-ai/sdk@0.91.1", "", { "dependencies": { "json-schema-to-ts": "^3.1.1" }, "peerDependencies": { "zod": "^3.25.0 || ^4.0.0" }, "optionalPeers": ["zod"], "bin": { "anthropic-ai-sdk": "bin/cli" } }, "sha512-LAmu761tSN9r66ixvmciswUj/ZC+1Q4iAfpedTfSVLeswRwnY3n2Nb6Tsk+cLPP28aLOPWeMgIuTuCcMC6W/iw=="], + "@anthropic-ai/sdk": ["@anthropic-ai/sdk@0.123.0", "", { "dependencies": { "json-schema-to-ts": "^3.1.1", "standardwebhooks": "^1.0.0" }, "peerDependencies": { "zod": "^3.25.0 || ^4.0.0" }, "optionalPeers": ["zod"], "bin": { "anthropic-ai-sdk": "bin/cli" } }, "sha512-Y9oX9mPNGZClHQOFqrWRk43Srcu/UHuPq3rfxxOq7JgW0gi+lJA2MAOK4Ul3k/+AUrwRWFJvd0tK3oC0Pw25dw=="], "@aws-crypto/sha256-browser": ["@aws-crypto/sha256-browser@5.2.0", "", { "dependencies": { "@aws-crypto/sha256-js": "^5.2.0", "@aws-crypto/supports-web-crypto": "^5.2.0", "@aws-crypto/util": "^5.2.0", "@aws-sdk/types": "^3.222.0", "@aws-sdk/util-locate-window": "^3.0.0", "@smithy/util-utf8": "^2.0.0", "tslib": "^2.6.2" } }, "sha512-AXfN/lGotSQwu6HNcEsIASo7kWXZ5HYWvfOmSNKDsEqC4OashTp8alTmaz+F7TC2L083SFv5RdB+qU3Vs1kZqw=="], @@ -230,13 +230,17 @@ "@cortexkit/pi-antigravity-auth": ["@cortexkit/pi-antigravity-auth@workspace:packages/pi"], - "@earendil-works/pi-agent-core": ["@earendil-works/pi-agent-core@0.79.10", "", { "dependencies": { "@earendil-works/pi-ai": "^0.79.10", "ignore": "7.0.5", "typebox": "1.1.38", "yaml": "2.9.0" } }, "sha512-XKxgdjhcPuyjrthCOFSgfzT3xZ1uBrJ1IMVDxci1to6hIN6BIg9J5iY8q0pGXK1DLgATLP23da+1UyZLwA360Q=="], + "@earendil-works/chord": ["@earendil-works/chord@0.85.1", "", { "dependencies": { "esbuild": "0.28.1" } }, "sha512-VDlkEC3dhCzQ5fcyH1OhG19dq+6jCn+rqc/iXFivwDYGR5anwo2RCiXij9PpHhqNR5GuhhE+Er69Zi1Sn4eY6w=="], - "@earendil-works/pi-ai": ["@earendil-works/pi-ai@0.79.10", "", { "dependencies": { "@anthropic-ai/sdk": "0.91.1", "@aws-sdk/client-bedrock-runtime": "3.1048.0", "@google/genai": "1.52.0", "@mistralai/mistralai": "2.2.6", "@opentelemetry/api": "1.9.0", "@smithy/node-http-handler": "4.7.3", "http-proxy-agent": "7.0.2", "https-proxy-agent": "7.0.6", "openai": "6.26.0", "partial-json": "0.1.7", "typebox": "1.1.38" }, "bin": { "pi-ai": "dist/cli.js" } }, "sha512-9jR23tOl0BIUdQMn70Gr72xYBpM7Xgl9Lyv7gAnU1USfkNRuYG/f/edLl+n/Dp/RafDW3JI4DF7y/GhgkORuew=="], + "@earendil-works/pi-agent-core": ["@earendil-works/pi-agent-core@0.85.1", "", { "dependencies": { "@earendil-works/chord": "^0.85.1", "@earendil-works/pi-ai": "^0.85.1", "@earendil-works/pi-telemetry": "^0.85.1", "diff": "8.0.4", "ignore": "7.0.5", "typebox": "1.3.7", "yaml": "2.9.0" } }, "sha512-hIXIP3eAWueAYiAl8aMvWCvvZ8Q5gT3Dip5bE5uJyIGh4+YlWRjtMLI4BaeoXoSs93zndjue61u1B/vhefLnuA=="], - "@earendil-works/pi-coding-agent": ["@earendil-works/pi-coding-agent@0.79.10", "", { "dependencies": { "@earendil-works/pi-agent-core": "^0.79.10", "@earendil-works/pi-ai": "^0.79.10", "@earendil-works/pi-tui": "^0.79.10", "@silvia-odwyer/photon-node": "0.3.4", "chalk": "5.6.2", "cross-spawn": "7.0.6", "diff": "8.0.4", "glob": "13.0.6", "highlight.js": "10.7.3", "hosted-git-info": "9.0.3", "ignore": "7.0.5", "jiti": "2.7.0", "minimatch": "10.2.5", "proper-lockfile": "4.1.2", "semver": "7.8.0", "typebox": "1.1.38", "undici": "8.5.0", "yaml": "2.9.0" }, "optionalDependencies": { "@mariozechner/clipboard": "0.3.9" }, "bin": { "pi": "dist/cli.js" } }, "sha512-YxaRhmgyDTvLDdGVbe7YzTHV80oL5mX5odg6EhGHz3w5Wu1Ix8DCw7bhtiOBLGQNFRcknia0zPmVWIj30XP1EA=="], + "@earendil-works/pi-ai": ["@earendil-works/pi-ai@0.85.1", "", { "dependencies": { "@anthropic-ai/sdk": "0.123.0", "@aws-sdk/client-bedrock-runtime": "3.1048.0", "@earendil-works/pi-telemetry": "^0.85.1", "@google/genai": "1.52.0", "@smithy/node-http-handler": "4.7.3", "http-proxy-agent": "7.0.2", "https-proxy-agent": "7.0.6", "openai": "6.40.0", "partial-json": "0.1.7", "typebox": "1.3.7" }, "bin": { "pi-ai": "dist/cli.js" } }, "sha512-+VgVIJDkDO2efYJKEEqvPTH4zmnIaXdAppGbO+vKFA9qy5PdhFiAenuFAkU+oiCSfOC4dMHDyrjdQeL4ZoC5CQ=="], - "@earendil-works/pi-tui": ["@earendil-works/pi-tui@0.79.10", "", { "dependencies": { "get-east-asian-width": "1.6.0", "marked": "18.0.5" } }, "sha512-FUVOjDn1DVwM1uHD5MNYboXQrXjIDbSt+BQ3py7nQWCY62tKfxgiM1OBMxTcwRWLfSdZHUPpV0hm1loIdUJnPw=="], + "@earendil-works/pi-coding-agent": ["@earendil-works/pi-coding-agent@0.85.1", "", { "dependencies": { "@earendil-works/chord": "^0.85.1", "@earendil-works/pi-agent-core": "^0.85.1", "@earendil-works/pi-ai": "^0.85.1", "@earendil-works/pi-tui": "^0.85.1", "@silvia-odwyer/photon-node": "0.3.4", "chalk": "5.6.2", "cross-spawn": "7.0.6", "diff": "8.0.4", "grok-mermaid": "0.2.2", "highlight.js": "10.7.3", "hosted-git-info": "9.0.3", "ignore": "7.0.5", "jiti": "2.7.0", "minimatch": "10.2.5", "proper-lockfile": "4.1.2", "semver": "7.8.0", "typebox": "1.3.7", "undici": "8.9.0", "yaml": "2.9.0" }, "optionalDependencies": { "@mariozechner/clipboard": "0.3.9" }, "bin": { "pi": "dist/bundle/cli.js" } }, "sha512-FGRN+OHbWaefBPGaTggAdLjrIHW+s2PzLyglz/5dfLzb9of7uuXMXYC0fJIeZTw+shS32o2cuQ9jF7YSDuL/oQ=="], + + "@earendil-works/pi-telemetry": ["@earendil-works/pi-telemetry@0.85.1", "", {}, "sha512-Bg/YN6kA7Swja/NQxka8xFdecb4E/auIEGF2G5A25EaQXhRnPj300/7/KpgsDDMYUzHTDAv4RyUxaQPJKW81Rw=="], + + "@earendil-works/pi-tui": ["@earendil-works/pi-tui@0.85.1", "", { "dependencies": { "get-east-asian-width": "1.6.0", "marked": "18.0.5" } }, "sha512-OIzw9efInmO4WOBnD4TxcTdBjmzvYJpzslkgoUro946nEGoYWg5rwv1p4fDt3/JvMx9QybryUCUwlm7j8Dreig=="], "@esbuild/aix-ppc64": ["@esbuild/aix-ppc64@0.27.7", "", { "os": "aix", "cpu": "ppc64" }, "sha512-EKX3Qwmhz1eMdEJokhALr0YiD0lhQNwDqkPYyPhiSwKrh7/4KRjQc04sZ8db+5DVVnZ1LmbNDI1uAMPEUBnQPg=="], @@ -322,8 +326,6 @@ "@mariozechner/clipboard-win32-x64-msvc": ["@mariozechner/clipboard-win32-x64-msvc@0.3.9", "", { "os": "win32", "cpu": "x64" }, "sha512-ihQC3EufqEY81vhXBgVBtK4prL+wc62zJsSvxrgz7K1hsdt6OObz6v9p3Rn1OG3GJksTTKMJF0u/guMISHPhSA=="], - "@mistralai/mistralai": ["@mistralai/mistralai@2.2.6", "", { "dependencies": { "@opentelemetry/semantic-conventions": "^1.40.0", "ws": "^8.18.0", "zod": "^3.25.0 || ^4.0.0", "zod-to-json-schema": "^3.25.0" }, "peerDependencies": { "@opentelemetry/api": "^1.9.0" }, "optionalPeers": ["@opentelemetry/api"] }, "sha512-W8pX7zHxjJvMIpw8JMxeJEleapXX0Q9NPszdNzqkM3MIEoIGPObdodujj+WHteXEvGfaP/AMwlNyRfEzSY6dQQ=="], - "@msgpackr-extract/msgpackr-extract-darwin-arm64": ["@msgpackr-extract/msgpackr-extract-darwin-arm64@3.0.4", "", { "os": "darwin", "cpu": "arm64" }, "sha512-LCkGo6JDfaBhgST7UpPWgNgLINpcpabaHfyz5OBx75nUYxBsaEPxjnyNjWpeb/xBup/682QnBfRBy2/LvPutZQ=="], "@msgpackr-extract/msgpackr-extract-darwin-x64": ["@msgpackr-extract/msgpackr-extract-darwin-x64@3.0.4", "", { "os": "darwin", "cpu": "x64" }, "sha512-zExlW9zUJKZH/tOtVMttwjKa4Xm/3KcNjnE3dPN92uCktwavMxpgCA3MoJK/DOnTWsQgo224OaST27/mPNAf+w=="], @@ -340,10 +342,6 @@ "@opencode-ai/sdk": ["@opencode-ai/sdk@1.17.13", "", { "dependencies": { "cross-spawn": "7.0.6" } }, "sha512-VItOGjMzRQx3zypwmeFLNhCiIx32kxS7FqzIJvVZLfyNGCifs3rfGC9qzNKWcxQo4SjNvAw++v4gWWU6Inv+JQ=="], - "@opentelemetry/api": ["@opentelemetry/api@1.9.0", "", {}, "sha512-3giAOQvZiH5F9bMlMiv8+GSPMeqg0dbaeo58/0SlA9sxSqZhnUtxzX9/2FzyhS9sWQf5S0GJE0AKBrFqjpeYcg=="], - - "@opentelemetry/semantic-conventions": ["@opentelemetry/semantic-conventions@1.43.0", "", {}, "sha512-eSYWTm620tTk45EKSedaUL8MFYI8hW164hIXsgIHyxu3VobUB3fFCu5t0hQby6OoWRPsG1KkKUG2M5UadiLiVg=="], - "@opentui/core": ["@opentui/core@0.4.5", "", { "dependencies": { "bun-ffi-structs": "0.2.4", "diff": "9.0.0", "marked": "17.0.1", "string-width": "7.2.0", "strip-ansi": "7.1.2" }, "optionalDependencies": { "@opentui/core-darwin-arm64": "0.4.5", "@opentui/core-darwin-x64": "0.4.5", "@opentui/core-linux-arm64": "0.4.5", "@opentui/core-linux-arm64-musl": "0.4.5", "@opentui/core-linux-x64": "0.4.5", "@opentui/core-linux-x64-musl": "0.4.5", "@opentui/core-win32-arm64": "0.4.5", "@opentui/core-win32-x64": "0.4.5" }, "peerDependencies": { "web-tree-sitter": "0.25.10" } }, "sha512-JsgRTPkA6e+Vxmumxai6SElOSlRQkbzNKHlCfemlArRiLhfC1IZ9RXJo2QH4xSu+uBOWAM90uss73/pPlkdEig=="], "@opentui/core-darwin-arm64": ["@opentui/core-darwin-arm64@0.4.5", "", { "os": "darwin", "cpu": "arm64" }, "sha512-8KUG0oRidnR+oW1RSZJ72/PhZLl+qRRMk5U/mieF4c0SJ5V3tYACpBZAKzQfHNd1f7QzD8FHZct1lPpQgtmkWg=="], @@ -404,6 +402,8 @@ "@smithy/util-utf8": ["@smithy/util-utf8@2.3.0", "", { "dependencies": { "@smithy/util-buffer-from": "^2.2.0", "tslib": "^2.6.2" } }, "sha512-R8Rdn8Hy72KKcebgLiv8jQcQkXoLMOGGv5uI1/k0l+snqkOzQ1R0ChUBCxWMlBsFMekWjq0wRudIweFs7sKT5A=="], + "@stablelib/base64": ["@stablelib/base64@1.0.1", "", {}, "sha512-1bnPQqSxSuc3Ii6MhBysoWCg58j97aUjuCSZrGSmDxNqtytIi0k8utUenAwTZN4V5mXXYGsVUI9zeBqy+jBOSQ=="], + "@standard-schema/spec": ["@standard-schema/spec@1.1.0", "", {}, "sha512-l2aFy5jALhniG5HgqrD6jXLi/rUWrKvqN/qJx6yoJsgKhblVd+iqqU4RCXavm/jPityDo5TCvKMnpjKnOriy0w=="], "@tsconfig/bun": ["@tsconfig/bun@1.0.10", "", {}, "sha512-5AV5YknQjNyoYzZ/8NG0dawqew/wH+x7ANiCfCIn29qo0cdbd1EryvFD1k5NSZWLBMOI/fGqMIaxi58GPIP9Cg=="], @@ -482,6 +482,8 @@ "fast-check": ["fast-check@4.9.0", "", { "dependencies": { "pure-rand": "^8.0.0" } }, "sha512-7ms6T7SybUev/PQITciI0yLM2pOSFy5zpG8Ty7tQofcVaQUvrMXp6CBwqF6fThLCLOrfBtuHAtwq6Yu4XPCllg=="], + "fast-sha256": ["fast-sha256@1.3.0", "", {}, "sha512-n11RGP/lrWEFI/bWdygLxhI+pVeo1ZYIVwvvPkW7azl/rOy+F3HYRZ2K5zeE9mmkhQppyv9sQFx0JM9UabnpPQ=="], + "fetch-blob": ["fetch-blob@3.2.0", "", { "dependencies": { "node-domexception": "^1.0.0", "web-streams-polyfill": "^3.0.3" } }, "sha512-7yAQpD2UMJzLi1Dqv7qFYnPbaPx7ZfFK6PiIxQ4PfkGPyNyl2Ugx+a/umUonmKqjhM4DnfbMvdX6otXq83soQQ=="], "find-babel-config": ["find-babel-config@2.1.2", "", { "dependencies": { "json5": "^2.2.3" } }, "sha512-ZfZp1rQyp4gyuxqt1ZqjFGVeVBvmpURMqdIWXbPRfB97Bf6BzdK/xSIbylEINzQ0kB5tlDQfn9HkNXXWsqTqLg=="], @@ -504,7 +506,7 @@ "get-east-asian-width": ["get-east-asian-width@1.6.0", "", {}, "sha512-QRbvDIbx6YklUe6RxeTeleMR0yv3cYH6PsPZHcnVn7xv7zO1BHN8r0XETu8n6Ye3Q+ahtSarc3WgtNWmehIBfA=="], - "glob": ["glob@13.0.6", "", { "dependencies": { "minimatch": "^10.2.2", "minipass": "^7.1.3", "path-scurry": "^2.0.2" } }, "sha512-Wjlyrolmm8uDpm/ogGyXZXb1Z+Ca2B8NbJwqBVg0axK9GbBeoS7yGV6vjXnYdGm6X53iehEuxxbyiKp8QmN4Vw=="], + "glob": ["glob@9.3.5", "", { "dependencies": { "fs.realpath": "^1.0.0", "minimatch": "^8.0.2", "minipass": "^4.2.4", "path-scurry": "^1.6.1" } }, "sha512-e1LleDykUz2Iu+MTYdkSsuWX8lvAjAcs0Xef0lNIu0S2wOAzuTxCJtcd9S3cijlwYF18EsU3rzb8jPVobxDh9Q=="], "google-auth-library": ["google-auth-library@10.9.0", "", { "dependencies": { "base64-js": "^1.3.0", "ecdsa-sig-formatter": "^1.0.11", "gaxios": "^7.1.4", "gcp-metadata": "8.1.2", "google-logging-utils": "1.1.3", "jws": "^4.0.0" } }, "sha512-xtvUqvINPhTaBm7nXqlYPcrMHJPm1lCNdSovxnKKhTm+4JsvQ+KGVYJViLoH9Yxu8w+T0Qv5HubzYT9BLrppJg=="], @@ -512,6 +514,8 @@ "graceful-fs": ["graceful-fs@4.2.11", "", {}, "sha512-RbJ5/jmFcNNCcDV5o9eTnBLJ/HszWV0P73bc+Ff4nS/rJj+YaS6IGyiOL0VoBYX+l1Wrl3k63h/KrH+nhJ0XvQ=="], + "grok-mermaid": ["grok-mermaid@0.2.2", "", {}, "sha512-XcJEP5dDC8liHBh52mlLjU18fNvu1ckFsu0QpIG3+APZ270fsj9wxpiA6cOURmbUEuoMVgjbC2+UYgTdCqqgzA=="], + "hasown": ["hasown@2.0.4", "", { "dependencies": { "function-bind": "^1.1.2" } }, "sha512-T2UbfbBEF32wiepXIsMlTW9+dDYC6wMh/t/vYA4tuOMKqWz/n3vr1NFSxQiyP+zk2mXsoMA/i/7qV6LKut1t1A=="], "highlight.js": ["highlight.js@10.7.3", "", {}, "sha512-tzcUFauisWKNHaRkN4Wjl/ZA07gENAjFl3J/c480dprkGTg5EQstgaNFqBfUqCq54kZRIEcreTsAgF/m2quD7A=="], @@ -586,7 +590,7 @@ "minimatch": ["minimatch@10.2.5", "", { "dependencies": { "brace-expansion": "^5.0.5" } }, "sha512-MULkVLfKGYDFYejP07QOurDLLQpcjk7Fw+7jXS2R2czRQzR56yHRveU5NDJEOviH+hETZKSkIk5c+T23GjFUMg=="], - "minipass": ["minipass@7.1.3", "", {}, "sha512-tEBHqDnIoM/1rXME1zgka9g6Q2lcoCkxHLuc7ODJ5BxbP5d4c2Z5cGgtXAku59200Cx7diuHTOYfSBD8n6mm8A=="], + "minipass": ["minipass@4.2.8", "", {}, "sha512-fNzuVyifolSLFL4NzpF+wEF4qrgqaaKX0haXPQEdQ7NKAN+WecoKMHV09YcuL/DHxrUsYQOK3MiuDf7Ip2OXfQ=="], "ms": ["ms@2.1.3", "", {}, "sha512-6FlzubTLZG3J2a/NVCAleEhjzq5oxgHyaCU9yYXvcLsvoVaHJq/s5xXI6/XXP6tz7R9xAOtHnSO/tXtF3WRTlA=="], @@ -604,7 +608,7 @@ "node-releases": ["node-releases@2.0.51", "", {}, "sha512-wRNIrw4DmVLKQlbgOMdkMx27Wrpzes2hh5Jtbi2bjPd+4wJstWIqP5A+lscnqbm0xxmT5Bpg8Lec5ItEBwx6BQ=="], - "openai": ["openai@6.26.0", "", { "peerDependencies": { "ws": "^8.18.0", "zod": "^3.25 || ^4.0" }, "optionalPeers": ["ws", "zod"], "bin": { "openai": "bin/cli" } }, "sha512-zd23dbWTjiJ6sSAX6s0HrCZi41JwTA1bQVs0wLQPZ2/5o2gxOJA5wh7yOAUgwYybfhDXyhwlpeQf7Mlgx8EOCA=="], + "openai": ["openai@6.40.0", "", { "peerDependencies": { "ws": "^8.18.0", "zod": "^3.25 || ^4.0" }, "optionalPeers": ["ws", "zod"] }, "sha512-MWtTjd/gQt4jpbji61NTgFWJLoY/PdRJ6wG9/ZDRMYNMlBKrCrSlkLI+KgHP1vR1qT6LKSAyAqIxno6lcK9JiA=="], "p-limit": ["p-limit@2.3.0", "", { "dependencies": { "p-try": "^2.0.0" } }, "sha512-//88mFWSJx8lxCzwdAABTJL2MyWB12+eIY7MDL2SqLmAkeKU9qxRvWuSyTjm3FUmpBEMuFfckAIqEaVGUDxb6w=="], @@ -624,7 +628,7 @@ "path-parse": ["path-parse@1.0.7", "", {}, "sha512-LDJzPVEEEPR+y48z93A0Ed0yXb8pAByGWo/k5YYdYgpY2/2EsOsksJrq7lOHxryrVOn1ejG6oAp8ahvOIQD8sw=="], - "path-scurry": ["path-scurry@2.0.2", "", { "dependencies": { "lru-cache": "^11.0.0", "minipass": "^7.1.2" } }, "sha512-3O/iVVsJAPsOnpwWIeD+d6z/7PmqApyQePUtCndjatj/9I5LylHvt5qluFaBT3I5h3r1ejfR056c+FCv+NnNXg=="], + "path-scurry": ["path-scurry@1.11.1", "", { "dependencies": { "lru-cache": "^10.2.0", "minipass": "^5.0.0 || ^6.0.2 || ^7.0.0" } }, "sha512-Xa4Nw17FS9ApQFJ9umLiJS4orGjm7ZzwUrwamcGQuHSzDyth9boKDaycYdDcZDuqYATXw4HFXgaqWTctW/v1HA=="], "picocolors": ["picocolors@1.1.1", "", {}, "sha512-xceH2snhtb5M9liqDsmEw56le376mTZkEX/jEb/RxNFyegNul7eNslCXP9FDj/Lcu0X8KEyMceP2ntpaHrDEVA=="], @@ -660,6 +664,8 @@ "solid-js": ["solid-js@1.9.12", "", { "dependencies": { "csstype": "^3.1.0", "seroval": "~1.5.0", "seroval-plugins": "~1.5.0" } }, "sha512-QzKaSJq2/iDrWR1As6MHZQ8fQkdOBf8GReYb7L5iKwMGceg7HxDcaOHk0at66tNgn9U2U7dXo8ZZpLIAmGMzgw=="], + "standardwebhooks": ["standardwebhooks@1.1.1", "", { "dependencies": { "@stablelib/base64": "^1.0.0", "fast-sha256": "^1.3.0" } }, "sha512-bCbX9ZEyFkWPsRz7Bl3NuQUJohmwGSev/yhr7vhaGPlc4AfIrspIRa6cPTBuI1ItmrTDJ4d/S2hCsfe4+vQGnQ=="], + "string-width": ["string-width@7.2.0", "", { "dependencies": { "emoji-regex": "^10.3.0", "get-east-asian-width": "^1.0.0", "strip-ansi": "^7.1.0" } }, "sha512-tsaTIkKW9b4N+AEj+SVA+WhJzV7/zMhcSu78mLKWSk7cXMOSHsBKFWUs0fWwq8QyK3MgJBQRX6Gbi4kYbdvGkQ=="], "strip-ansi": ["strip-ansi@7.1.2", "", { "dependencies": { "ansi-regex": "^6.0.1" } }, "sha512-gmBGslpoQJtgnMAvOVqGZpEz9dyoKTCzy2nfz/n8aIFhN/jCE/rCmcxabB6jOOHV+0WNnylOxaxBQPSvcWklhA=="], @@ -672,11 +678,11 @@ "tslib": ["tslib@2.8.1", "", {}, "sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w=="], - "typebox": ["typebox@1.1.38", "", {}, "sha512-pZ0aQPmMmXoUvSbeuWf/Hzsc+avNw/Zd6VeE8CFgkVGWyuHPJvqeJJDeJqLve+K70LvjYIoleGcoJHPT17cWoA=="], + "typebox": ["typebox@1.3.7", "", {}, "sha512-meKuifc33Pccx0O6PdIzYMq3Og8zvP4TIi/a+Bw3AEMZMxOD0+RHGQvpglEe6Zdy3wZ8nqn/j95h8LUZLk/6Hg=="], "typescript": ["typescript@6.0.3", "", { "bin": { "tsc": "bin/tsc", "tsserver": "bin/tsserver" } }, "sha512-y2TvuxSZPDyQakkFRPZHKFm+KKVqIisdg9/CZwm9ftvKXLP8NRWj38/ODjNbr43SsoXqNuAisEf1GdCxqWcdBw=="], - "undici": ["undici@8.5.0", "", {}, "sha512-xamtWoB1EshgjpmlXd7GGm2VfdDtw1+rD8uhry8pSNW3If6S8E0m2T2+orSKeZXEn/aPJMviCpDBA65WJt8zhg=="], + "undici": ["undici@8.9.0", "", {}, "sha512-aWZpUj7XoGonMClx4gdDRfgBjqeA+F473aDmROQQbM9n6PRfK/u1q/a0X4wMTgcHfT8H6fpbt98PFuDUwFg2YA=="], "undici-types": ["undici-types@7.16.0", "", {}, "sha512-Zz+aZWSj8LE6zoxD+xrjh4VfkIG8Ya6LvYkZqtUQGJPZjYl53ypCaUwWqo7eI0x66KBGeRo+mlBEkMSeSZ38Nw=="], @@ -716,6 +722,10 @@ "@babel/helper-create-class-features-plugin/semver": ["semver@6.3.1", "", { "bin": { "semver": "bin/semver.js" } }, "sha512-BR7VvDCVHO+q2xBEWskxS6DJE1qRnb7DxzUrogb71CWoSficBxYsiAGd+Kl0mmq/MprG9yArRkyrQxTO6XjMzA=="], + "@earendil-works/chord/esbuild": ["esbuild@0.28.1", "", { "optionalDependencies": { "@esbuild/aix-ppc64": "0.28.1", "@esbuild/android-arm": "0.28.1", "@esbuild/android-arm64": "0.28.1", "@esbuild/android-x64": "0.28.1", "@esbuild/darwin-arm64": "0.28.1", "@esbuild/darwin-x64": "0.28.1", "@esbuild/freebsd-arm64": "0.28.1", "@esbuild/freebsd-x64": "0.28.1", "@esbuild/linux-arm": "0.28.1", "@esbuild/linux-arm64": "0.28.1", "@esbuild/linux-ia32": "0.28.1", "@esbuild/linux-loong64": "0.28.1", "@esbuild/linux-mips64el": "0.28.1", "@esbuild/linux-ppc64": "0.28.1", "@esbuild/linux-riscv64": "0.28.1", "@esbuild/linux-s390x": "0.28.1", "@esbuild/linux-x64": "0.28.1", "@esbuild/netbsd-arm64": "0.28.1", "@esbuild/netbsd-x64": "0.28.1", "@esbuild/openbsd-arm64": "0.28.1", "@esbuild/openbsd-x64": "0.28.1", "@esbuild/openharmony-arm64": "0.28.1", "@esbuild/sunos-x64": "0.28.1", "@esbuild/win32-arm64": "0.28.1", "@esbuild/win32-ia32": "0.28.1", "@esbuild/win32-x64": "0.28.1" }, "bin": { "esbuild": "bin/esbuild" } }, "sha512-HrJrvZv5ayxBzPfwphOoNzkzOIIlifzk0KJrGK2c8R4+LKpMtpYLQeUdjnwjWv/LZlkH2laZk+4w78pi99D4Vw=="], + + "@earendil-works/pi-agent-core/diff": ["diff@8.0.4", "", {}, "sha512-DPi0FmjiSU5EvQV0++GFDOJ9ASQUVFh5kD+OzOnYdi7n3Wpm9hWWGfB/O2blfHcMVTL5WkQXSnRiK9makhrcnw=="], + "@earendil-works/pi-coding-agent/diff": ["diff@8.0.4", "", {}, "sha512-DPi0FmjiSU5EvQV0++GFDOJ9ASQUVFh5kD+OzOnYdi7n3Wpm9hWWGfB/O2blfHcMVTL5WkQXSnRiK9makhrcnw=="], "@earendil-works/pi-tui/marked": ["marked@18.0.5", "", { "bin": { "marked": "bin/marked.js" } }, "sha512-S6GcvALHg6K4ohtu4E7x0a1AqhAjp6cV8KhLSyN9qVapnzJkusVBxZRcIU9AeYsbe6P1hKDusSbEOzGyyuce6w=="], @@ -724,24 +734,70 @@ "babel-plugin-jsx-dom-expressions/@babel/helper-module-imports": ["@babel/helper-module-imports@7.18.6", "", { "dependencies": { "@babel/types": "^7.18.6" } }, "sha512-0NFvs3VkuSYbFi1x2Vd6tKrywq+z/cLeYC/RJNFrIX/30Bf5aiGYbtvGXolEktzJH8o5E5KJ3tT+nkxuuZFVlA=="], - "babel-plugin-module-resolver/glob": ["glob@9.3.5", "", { "dependencies": { "fs.realpath": "^1.0.0", "minimatch": "^8.0.2", "minipass": "^4.2.4", "path-scurry": "^1.6.1" } }, "sha512-e1LleDykUz2Iu+MTYdkSsuWX8lvAjAcs0Xef0lNIu0S2wOAzuTxCJtcd9S3cijlwYF18EsU3rzb8jPVobxDh9Q=="], + "glob/minimatch": ["minimatch@8.0.7", "", { "dependencies": { "brace-expansion": "^2.0.1" } }, "sha512-V+1uQNdzybxa14e/p00HZnQNNcTjnRJjDxg2V8wtkjFctq4M7hXFws4oekyTP0Jebeq7QYtpFyOeBAjc88zvYg=="], "p-retry/retry": ["retry@0.13.1", "", {}, "sha512-XQBQ3I8W1Cge0Seh+6gjj03LbmRFWuoszgK9ooCpwYIrhhoO80pfq4cUkU5DkknwfOfFteRwlZ56PYOGYyFWdg=="], "parse5/entities": ["entities@6.0.1", "", {}, "sha512-aN97NXWF6AWBTahfVOIrB/NShkzi5H7F9r1s9mD3cDj4Ko5f2qhhVoYMibXF7GlLveb/D2ioWay8lxI97Ven3g=="], - "babel-plugin-module-resolver/glob/minimatch": ["minimatch@8.0.7", "", { "dependencies": { "brace-expansion": "^2.0.1" } }, "sha512-V+1uQNdzybxa14e/p00HZnQNNcTjnRJjDxg2V8wtkjFctq4M7hXFws4oekyTP0Jebeq7QYtpFyOeBAjc88zvYg=="], + "path-scurry/lru-cache": ["lru-cache@10.4.3", "", {}, "sha512-JNAzZcXrCt42VGLuYz0zfAzDfAvJWW6AfYlDBQyDV5DClI2m5sAmK+OIO7s59XfsRsWHp02jAJrRadPRGTt6SQ=="], + + "path-scurry/minipass": ["minipass@7.1.3", "", {}, "sha512-tEBHqDnIoM/1rXME1zgka9g6Q2lcoCkxHLuc7ODJ5BxbP5d4c2Z5cGgtXAku59200Cx7diuHTOYfSBD8n6mm8A=="], + + "@earendil-works/chord/esbuild/@esbuild/aix-ppc64": ["@esbuild/aix-ppc64@0.28.1", "", { "os": "aix", "cpu": "ppc64" }, "sha512-Svl7tq8k/08+p6CXPpRjQ1fKX+1odH/BQbb48fV6fj3CWHhsoIOoY87w1oHXm0qEpkIK3ZfVgp0hed3XBXzXMQ=="], + + "@earendil-works/chord/esbuild/@esbuild/android-arm": ["@esbuild/android-arm@0.28.1", "", { "os": "android", "cpu": "arm" }, "sha512-0k2F129Xdio1TdJfzJ8sy1Q47vUD2NnwdhiAf7drUN1EBTfPf4hsFCtmMgu/6m8JSzsBrlmVjudMBQqOfG8usQ=="], + + "@earendil-works/chord/esbuild/@esbuild/android-arm64": ["@esbuild/android-arm64@0.28.1", "", { "os": "android", "cpu": "arm64" }, "sha512-34EGEbCIAgosYz6goLcopX6Mo7NyGv9tfwEM2/7Ce2VcVRk568iSvniGWcUXIy7wEDR1wzolcxcriFVrWYcwBg=="], + + "@earendil-works/chord/esbuild/@esbuild/android-x64": ["@esbuild/android-x64@0.28.1", "", { "os": "android", "cpu": "x64" }, "sha512-dbwY7ltSMDWsRatcRpCnES4F+im88OCUgGZjy52shC7GqHRE/cYlxNbB4Z4UpJswpcc4Qxd2oE/ufM0p61IKng=="], + + "@earendil-works/chord/esbuild/@esbuild/darwin-arm64": ["@esbuild/darwin-arm64@0.28.1", "", { "os": "darwin", "cpu": "arm64" }, "sha512-TZbWkQY7kvTAXbXUT7uVACR5cMHsDiSz9z7ZKAX/RTq/WJEk3QyRr0wZpNhBDX+/0CtdqUIJlOiodQcta6tY3Q=="], + + "@earendil-works/chord/esbuild/@esbuild/darwin-x64": ["@esbuild/darwin-x64@0.28.1", "", { "os": "darwin", "cpu": "x64" }, "sha512-zfdzgK9ACBNZLI/CyHTOx81SyNbM6YXn7rxSgX97VjyiPl9W1i4Ka4fgKECEoFCKGpvBj5qArWIGgQjOwkgskQ=="], + + "@earendil-works/chord/esbuild/@esbuild/freebsd-arm64": ["@esbuild/freebsd-arm64@0.28.1", "", { "os": "freebsd", "cpu": "arm64" }, "sha512-wG2EA8ENdEI0qhkSZMjfqrdY+ziCYCPMmtZjjIwOmXFjmyzEHn+UUxk5of+SYsjtfs3VpnlC7QLzSI5hY/rOAw=="], + + "@earendil-works/chord/esbuild/@esbuild/freebsd-x64": ["@esbuild/freebsd-x64@0.28.1", "", { "os": "freebsd", "cpu": "x64" }, "sha512-i7dZ9vQgnvSCzi/rYCXNgtF/U+eKZNJBzu3eTQbRgHnM7tNSizLOkRFAl3qzVc/Op/u5YkHHa4pf/3DOYHthLQ=="], + + "@earendil-works/chord/esbuild/@esbuild/linux-arm": ["@esbuild/linux-arm@0.28.1", "", { "os": "linux", "cpu": "arm" }, "sha512-qVXBOHQS+d5Y722GwJzJUtOLlX7km3CraOaGormF1pDtPd2C/l1SHRPgjLunLGe51Sh5YYWKMFDyV4SxgMQYTQ=="], + + "@earendil-works/chord/esbuild/@esbuild/linux-arm64": ["@esbuild/linux-arm64@0.28.1", "", { "os": "linux", "cpu": "arm64" }, "sha512-yHs+0uc8+nvEAfAfxrWQKK5peSNzBc4PegcMO0EJ2hT71uA7vB8Ihg2e77R2P7SG5uYjPbHlLLmve4LLLRCf0g=="], + + "@earendil-works/chord/esbuild/@esbuild/linux-ia32": ["@esbuild/linux-ia32@0.28.1", "", { "os": "linux", "cpu": "ia32" }, "sha512-d1z4ZuP0ajrfz/FhGT4vv278rX8KnPPJx8i5+AtK7TYbx9Le9F1hyzurZpkEyjkGa9dUGhQow4C1NmeGvqxN2w=="], + + "@earendil-works/chord/esbuild/@esbuild/linux-loong64": ["@esbuild/linux-loong64@0.28.1", "", { "os": "linux", "cpu": "none" }, "sha512-M5sRjUVZrkm1OAPR3dlOYzNmN+loZKGVi1VUQGrwuqLcbR6qeAz+famMhjASeH3YVKvZz+zT1jlh/keC3Rj/lg=="], + + "@earendil-works/chord/esbuild/@esbuild/linux-mips64el": ["@esbuild/linux-mips64el@0.28.1", "", { "os": "linux", "cpu": "none" }, "sha512-mRObBZeHh2OxcBFPWE/FjylkRgZdYuiTR3vaTozquCGOH14iP9oN4x4Ge81CoIDYQrXmIxpFumJBu5MtZpnQJQ=="], + + "@earendil-works/chord/esbuild/@esbuild/linux-ppc64": ["@esbuild/linux-ppc64@0.28.1", "", { "os": "linux", "cpu": "ppc64" }, "sha512-slScBsMAb3GFDcdrCgLwZtPYRoH2H/youv10QiZyRjmsP48fznoveWytSgCI/R0ZcUgpc0ZhIUEx6LHts8yrfQ=="], + + "@earendil-works/chord/esbuild/@esbuild/linux-riscv64": ["@esbuild/linux-riscv64@0.28.1", "", { "os": "linux", "cpu": "none" }, "sha512-kw0owk1o0GFETUJyW0jc0G4Yzs0BHZn0JDZ8JRT088vjJYX777BAs1fDGxAC+q831qOs2DTC96mNsG2opdfyyQ=="], + + "@earendil-works/chord/esbuild/@esbuild/linux-s390x": ["@esbuild/linux-s390x@0.28.1", "", { "os": "linux", "cpu": "s390x" }, "sha512-/lAIjX8aYFRByhh6L5rYtPEDRqa9de/4V/juOXcta5frjvzXO4/sqEtyytse0g3zZFuWu5cDN0MkLz2qRDD2Ag=="], + + "@earendil-works/chord/esbuild/@esbuild/linux-x64": ["@esbuild/linux-x64@0.28.1", "", { "os": "linux", "cpu": "x64" }, "sha512-u/anNYF2mmVOEDwLtnQ1wOr3EZ9sTNGLWrsYGYwHWzGA3Si84IOkHXlbWTD1NB+9/1lcnweYKO54uhxZydNzfA=="], + + "@earendil-works/chord/esbuild/@esbuild/netbsd-arm64": ["@esbuild/netbsd-arm64@0.28.1", "", { "os": "none", "cpu": "arm64" }, "sha512-oks0DYbLwWMmaakTsCb+zL4E+aHRVLom9IJZOAthMQEPiQmydXHkziYEsGYRx0uNV/IjEKGAV941JzH02pflqw=="], + + "@earendil-works/chord/esbuild/@esbuild/netbsd-x64": ["@esbuild/netbsd-x64@0.28.1", "", { "os": "none", "cpu": "x64" }, "sha512-aeL6lAnN89Hz43Mlh1G8ARasbuoYvSITDEx0tHh5b7jJnHcssqgjy9Yx430GDpmCa6OyrKoS0aNRjKundRizGg=="], + + "@earendil-works/chord/esbuild/@esbuild/openbsd-arm64": ["@esbuild/openbsd-arm64@0.28.1", "", { "os": "openbsd", "cpu": "arm64" }, "sha512-MEFJe5C3R8pwXdZ5Y21oo6m7ePiS0d9pWucn99O/wvyJZChoIQKrQDxKrGeW8F5+T0okTHesAmDeiHDTIq0V/Q=="], + + "@earendil-works/chord/esbuild/@esbuild/openbsd-x64": ["@esbuild/openbsd-x64@0.28.1", "", { "os": "openbsd", "cpu": "x64" }, "sha512-i/ZLIOafE0Z8cI/XANJAixoJL/uRAoS2xOA3rb0xN+KK0K177cMAsQYkzHtBrtMXAKuAc7HGgcWiZ/sRC1Nxgw=="], + + "@earendil-works/chord/esbuild/@esbuild/openharmony-arm64": ["@esbuild/openharmony-arm64@0.28.1", "", { "os": "none", "cpu": "arm64" }, "sha512-ge+Z7EXFNt2BO1oAMsVpiQ8EwndV9i1xXerAeTIK7AtPs3bKFXQM7nlRxDSIUIMeueR1CNXxqztLzdNeReKBJg=="], - "babel-plugin-module-resolver/glob/minipass": ["minipass@4.2.8", "", {}, "sha512-fNzuVyifolSLFL4NzpF+wEF4qrgqaaKX0haXPQEdQ7NKAN+WecoKMHV09YcuL/DHxrUsYQOK3MiuDf7Ip2OXfQ=="], + "@earendil-works/chord/esbuild/@esbuild/sunos-x64": ["@esbuild/sunos-x64@0.28.1", "", { "os": "sunos", "cpu": "x64" }, "sha512-BEjgtECkL3vY+SaSQ6nzVfiALUeFxpawyp8Jmf5PtYhf1Ug40N1h/hxlhts+f1FvSvarEigdxS3BlSMI2PJLcQ=="], - "babel-plugin-module-resolver/glob/path-scurry": ["path-scurry@1.11.1", "", { "dependencies": { "lru-cache": "^10.2.0", "minipass": "^5.0.0 || ^6.0.2 || ^7.0.0" } }, "sha512-Xa4Nw17FS9ApQFJ9umLiJS4orGjm7ZzwUrwamcGQuHSzDyth9boKDaycYdDcZDuqYATXw4HFXgaqWTctW/v1HA=="], + "@earendil-works/chord/esbuild/@esbuild/win32-arm64": ["@esbuild/win32-arm64@0.28.1", "", { "os": "win32", "cpu": "arm64" }, "sha512-lCv9eK/H6ZJWbE7bh2nw54CZ9M2nupBxJcTsdk/QQnWkdSjKGuxmmH8/GWrlT1eMmZfn4dGcCjRte397WqfQXA=="], - "babel-plugin-module-resolver/glob/minimatch/brace-expansion": ["brace-expansion@2.1.2", "", { "dependencies": { "balanced-match": "^1.0.0" } }, "sha512-w5JZcKgdhDOgOwm8H+KgbosopHMuGcl6qbulwjtz3SM7I7P3yW1eAjzMPLrIE+NQ9vjgANKHWeMHnrT0OXW1oA=="], + "@earendil-works/chord/esbuild/@esbuild/win32-ia32": ["@esbuild/win32-ia32@0.28.1", "", { "os": "win32", "cpu": "ia32" }, "sha512-zvb/mB2bSCoJOpoCBgYKKpX6YM6mJBlBUVUtVj41DlZJVEB6/0CKlRYxP5wWl1C1ILiCoAU5wZZ4q1P3qeS6Eg=="], - "babel-plugin-module-resolver/glob/path-scurry/lru-cache": ["lru-cache@10.4.3", "", {}, "sha512-JNAzZcXrCt42VGLuYz0zfAzDfAvJWW6AfYlDBQyDV5DClI2m5sAmK+OIO7s59XfsRsWHp02jAJrRadPRGTt6SQ=="], + "@earendil-works/chord/esbuild/@esbuild/win32-x64": ["@esbuild/win32-x64@0.28.1", "", { "os": "win32", "cpu": "x64" }, "sha512-bm4Mowrv+GXMlpWX++EcXw/iLyd1o3+bJkC2DkWXYVvgZCqD/bSj9ctZeAMC3cIxgjRVR2Dufaiu4YPxr5gW1A=="], - "babel-plugin-module-resolver/glob/path-scurry/minipass": ["minipass@7.1.3", "", {}, "sha512-tEBHqDnIoM/1rXME1zgka9g6Q2lcoCkxHLuc7ODJ5BxbP5d4c2Z5cGgtXAku59200Cx7diuHTOYfSBD8n6mm8A=="], + "glob/minimatch/brace-expansion": ["brace-expansion@2.1.2", "", { "dependencies": { "balanced-match": "^1.0.0" } }, "sha512-w5JZcKgdhDOgOwm8H+KgbosopHMuGcl6qbulwjtz3SM7I7P3yW1eAjzMPLrIE+NQ9vjgANKHWeMHnrT0OXW1oA=="], - "babel-plugin-module-resolver/glob/minimatch/brace-expansion/balanced-match": ["balanced-match@1.0.2", "", {}, "sha512-3oSeUO0TMV67hN1AmbXsK4yaqU7tjiHlbxRDZOpH0KW9+CeX4bRAaX0Anxt0tx2MrpRpWwQaPwIlISEJhYU5Pw=="], + "glob/minimatch/brace-expansion/balanced-match": ["balanced-match@1.0.2", "", {}, "sha512-3oSeUO0TMV67hN1AmbXsK4yaqU7tjiHlbxRDZOpH0KW9+CeX4bRAaX0Anxt0tx2MrpRpWwQaPwIlISEJhYU5Pw=="], } } diff --git a/packages/core/src/antigravity/oauth.ts b/packages/core/src/antigravity/oauth.ts index 4995fd14..f514b3ac 100644 --- a/packages/core/src/antigravity/oauth.ts +++ b/packages/core/src/antigravity/oauth.ts @@ -78,6 +78,7 @@ export interface AntigravityRefreshResult { */ export async function refreshAntigravityToken( refreshToken: string, + signal?: AbortSignal, ): Promise { const startTime = Date.now() const response = await fetchWithActiveTimeout( @@ -91,6 +92,7 @@ export async function refreshAntigravityToken( client_id: ANTIGRAVITY_CLIENT_ID, client_secret: ANTIGRAVITY_CLIENT_SECRET, }), + signal, }, ) diff --git a/packages/pi/package.json b/packages/pi/package.json index 7099c0c8..8bd47963 100644 --- a/packages/pi/package.json +++ b/packages/pi/package.json @@ -46,14 +46,14 @@ "@cortexkit/antigravity-auth-core": "2.2.1" }, "peerDependencies": { - "@earendil-works/pi-ai": "*", - "@earendil-works/pi-coding-agent": "*", - "@earendil-works/pi-tui": "*" + "@earendil-works/pi-ai": "^0.85.1", + "@earendil-works/pi-coding-agent": "^0.85.1", + "@earendil-works/pi-tui": "^0.85.1" }, "devDependencies": { - "@earendil-works/pi-ai": "^0.79.1", - "@earendil-works/pi-coding-agent": "^0.79.1", - "@earendil-works/pi-tui": "^0.79.1", + "@earendil-works/pi-ai": "0.85.1", + "@earendil-works/pi-coding-agent": "0.85.1", + "@earendil-works/pi-tui": "0.85.1", "esbuild": "^0.27.7" } } diff --git a/packages/pi/src/commands.ts b/packages/pi/src/commands.ts index f2d46d98..2981a33a 100644 --- a/packages/pi/src/commands.ts +++ b/packages/pi/src/commands.ts @@ -1,10 +1,18 @@ -import type { ExtensionAPI } from '@earendil-works/pi-coding-agent' +import type { + OAuthCredentials, + OAuthLoginCallbacks, +} from '@earendil-works/pi-ai' +import { + type ExtensionAPI, + readStoredCredential, +} from '@earendil-works/pi-coding-agent' import type { PiAccountRuntime } from './runtime.ts' import { isStrategy, readSettings, writeStrategy } from './settings.ts' export function registerAccountCommands( pi: ExtensionAPI, runtime: PiAccountRuntime, + login: (callbacks: OAuthLoginCallbacks) => Promise, ): void { const register = ( name: string, @@ -92,7 +100,14 @@ export function registerAccountCommands( 'Add an account using the provider OAuth login', async (_args, ctx) => { if (!ctx.hasUI) throw new Error('OAuth requires interactive Pi') - await ctx.modelRegistry.authStorage.login('google-antigravity', { + if (readStoredCredential('google-antigravity')?.type !== 'oauth') { + ctx.ui.notify( + 'Authenticate the provider first with /login google-antigravity, then use /agy-add for additional accounts.', + 'warning', + ) + return + } + await login({ onAuth: ({ url }) => ctx.ui.notify(`Open this URL in your browser:\n${url}`, 'info'), onPrompt: async ({ message }) => { diff --git a/packages/pi/src/extension-load.test.ts b/packages/pi/src/extension-load.test.ts new file mode 100644 index 00000000..289a2852 --- /dev/null +++ b/packages/pi/src/extension-load.test.ts @@ -0,0 +1,111 @@ +import { afterEach, beforeEach, describe, expect, it } from 'bun:test' +import { mkdir, mkdtemp, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { loadAccountStorage } from '@cortexkit/antigravity-auth-core' +import { + discoverAndLoadExtensions, + ExtensionRunner, + ModelRegistry, + ModelRuntime, + SessionManager, +} from '@earendil-works/pi-coding-agent' + +let directory: string +let previousAccountPath: string | undefined +let previousAgentDir: string | undefined + +beforeEach(async () => { + directory = await mkdtemp(join(tmpdir(), 'pi-extension-load-')) + previousAccountPath = process.env.PI_ANTIGRAVITY_AUTH_FILE + previousAgentDir = process.env.PI_CODING_AGENT_DIR + process.env.PI_ANTIGRAVITY_AUTH_FILE = join(directory, 'accounts.json') + process.env.PI_CODING_AGENT_DIR = join(directory, 'agent') +}) + +afterEach(async () => { + if (previousAccountPath === undefined) + delete process.env.PI_ANTIGRAVITY_AUTH_FILE + else process.env.PI_ANTIGRAVITY_AUTH_FILE = previousAccountPath + if (previousAgentDir === undefined) delete process.env.PI_CODING_AGENT_DIR + else process.env.PI_CODING_AGENT_DIR = previousAgentDir + await rm(directory, { recursive: true, force: true }) +}) + +async function writeHostCredential(credential: object): Promise { + const agentDir = process.env.PI_CODING_AGENT_DIR! + await mkdir(agentDir, { recursive: true }) + await writeFile( + join(agentDir, 'auth.json'), + JSON.stringify({ 'google-antigravity': credential }), + ) +} + +describe('Pi 0.85 extension loading', () => { + it('loads through the real extension runner and migrates only stored OAuth credentials at session_start', async () => { + const loaded = await discoverAndLoadExtensions( + [join(import.meta.dir, 'index.ts')], + directory, + process.env.PI_CODING_AGENT_DIR, + ) + expect(loaded.errors).toEqual([]) + expect(loaded.extensions).toHaveLength(1) + expect(loaded.extensions[0]?.commands.has('agy-accounts')).toBe(true) + expect( + loaded.runtime.pendingProviderRegistrations.map(({ name }) => name), + ).toEqual(['google-antigravity']) + + const modelRuntime = await ModelRuntime.create({ + authPath: join(process.env.PI_CODING_AGENT_DIR!, 'auth.json'), + modelsPath: null, + refreshOnCreate: false, + }) + const runner = new ExtensionRunner( + loaded.extensions, + loaded.runtime, + directory, + SessionManager.inMemory(directory), + new ModelRegistry(modelRuntime), + ) + const errors: string[] = [] + runner.onError(({ error }) => errors.push(error)) + + await runner.emit({ type: 'session_start', reason: 'startup' }) + expect( + await loadAccountStorage(process.env.PI_ANTIGRAVITY_AUTH_FILE!), + ).toBeNull() + + await writeFile( + join(process.env.PI_CODING_AGENT_DIR!, 'auth.json'), + '{malformed', + ) + await runner.emit({ type: 'session_start', reason: 'reload' }) + expect( + await loadAccountStorage(process.env.PI_ANTIGRAVITY_AUTH_FILE!), + ).toBeNull() + + await writeHostCredential({ type: 'api_key', key: 'ignored-api-key' }) + await runner.emit({ type: 'session_start', reason: 'reload' }) + expect( + await loadAccountStorage(process.env.PI_ANTIGRAVITY_AUTH_FILE!), + ).toBeNull() + + await writeHostCredential({ + type: 'oauth', + refresh: 'stored-refresh|project|managed', + access: 'stored-access', + expires: Date.now() + 3_600_000, + email: 'stored@example.com', + }) + await runner.emit({ type: 'session_start', reason: 'reload' }) + + const storage = await loadAccountStorage( + process.env.PI_ANTIGRAVITY_AUTH_FILE!, + ) + expect(storage?.accounts).toHaveLength(1) + expect(storage?.accounts[0]?.email).toBe('stored@example.com') + expect(errors).toEqual([]) + + await runner.emit({ type: 'session_shutdown', reason: 'quit' }) + }) +}) diff --git a/packages/pi/src/index.ts b/packages/pi/src/index.ts index 8e1de654..bf083494 100644 --- a/packages/pi/src/index.ts +++ b/packages/pi/src/index.ts @@ -7,7 +7,10 @@ import type { OAuthCredentials, OAuthLoginCallbacks, } from '@earendil-works/pi-ai' -import type { ExtensionAPI } from '@earendil-works/pi-coding-agent' +import { + type ExtensionAPI, + readStoredCredential, +} from '@earendil-works/pi-coding-agent' import { registerAccountCommands } from './commands.ts' import { rememberPackedRefresh } from './credential-cache.ts' import { PiAccountRuntime } from './runtime.ts' @@ -69,18 +72,20 @@ async function loginAntigravity( export default function cortexKitPiAntigravityAuth(pi: ExtensionAPI): void { const runtime = new PiAccountRuntime() - registerAccountCommands(pi, runtime) + registerAccountCommands(pi, runtime, (callbacks) => + loginAntigravity(callbacks, runtime), + ) pi.on('session_start', async (_event, context) => { - const auth = context.modelRegistry.authStorage.get(ANTIGRAVITY_PROVIDER_ID) - if (auth?.type === 'oauth') { - try { + try { + const auth = readStoredCredential(ANTIGRAVITY_PROVIDER_ID) + if (auth?.type === 'oauth') { await runtime.migrate(auth) - } catch { - context.ui.notify( - 'Antigravity account migration failed; existing auth and pool were retained. Repair the account file before retrying.', - 'error', - ) } + } catch { + context.ui.notify( + 'Antigravity account migration failed; existing auth and pool were retained. Repair the account file before retrying.', + 'error', + ) } }) pi.on('session_shutdown', async () => runtime.dispose()) @@ -107,7 +112,8 @@ export default function cortexKitPiAntigravityAuth(pi: ExtensionAPI): void { oauth: { name: 'Google Antigravity (CortexKit)', login: (callbacks) => loginAntigravity(callbacks, runtime), - refreshToken: (credentials) => runtime.refreshHost(credentials), + refreshToken: (credentials, signal) => + runtime.refreshHost(credentials, signal), getApiKey: (credentials) => { // Bridge the packed refresh (refreshToken|projectId|managedProjectId) // to the stream, which otherwise only receives the bare access token. diff --git a/packages/pi/src/provider.test.ts b/packages/pi/src/provider.test.ts index 5fb9986f..b219c7f9 100644 --- a/packages/pi/src/provider.test.ts +++ b/packages/pi/src/provider.test.ts @@ -1,5 +1,5 @@ import { afterEach, beforeEach, describe, expect, it, mock } from 'bun:test' -import { mkdtemp, rm } from 'node:fs/promises' +import { mkdir, mkdtemp, rm, writeFile } from 'node:fs/promises' import { tmpdir } from 'node:os' import { join } from 'node:path' import * as core from '@cortexkit/antigravity-auth-core' @@ -35,6 +35,7 @@ const commands = new Map() const hooks = new Map Promise>() let directory: string let previousPath: string | undefined +let previousAgentDir: string | undefined let previousFetch: typeof fetch let sequence: number let hostAuth: OAuthCredentials | undefined @@ -50,6 +51,8 @@ beforeEach(async () => { directory = await mkdtemp(join(tmpdir(), 'pi-provider-')) previousPath = process.env.PI_ANTIGRAVITY_AUTH_FILE process.env.PI_ANTIGRAVITY_AUTH_FILE = join(directory, 'accounts.json') + previousAgentDir = process.env.PI_CODING_AGENT_DIR + process.env.PI_CODING_AGENT_DIR = join(directory, 'agent') previousFetch = globalThis.fetch globalThis.fetch = mock(async () => Response.json({ groups: [] }), @@ -91,6 +94,8 @@ afterEach(async () => { globalThis.fetch = previousFetch if (previousPath === undefined) delete process.env.PI_ANTIGRAVITY_AUTH_FILE else process.env.PI_ANTIGRAVITY_AUTH_FILE = previousPath + if (previousAgentDir === undefined) delete process.env.PI_CODING_AGENT_DIR + else process.env.PI_CODING_AGENT_DIR = previousAgentDir await rm(directory, { recursive: true, force: true }) }) @@ -98,20 +103,34 @@ function context() { return { hasUI: true, ui: { notify, input: async () => 'code', select: async () => undefined }, - modelRegistry: { - authStorage: { - get: () => (hostAuth ? { type: 'oauth', ...hostAuth } : undefined), - login: async (_id: string, cb: OAuthLoginCallbacks) => { - hostAuth = await provider.oauth!.login(cb) - }, - }, - }, + modelRegistry: {}, } as unknown as ExtensionCommandContext } +async function persistHostAuth(): Promise { + const agentDir = process.env.PI_CODING_AGENT_DIR! + await mkdir(agentDir, { recursive: true }) + await writeFile( + join(agentDir, 'auth.json'), + JSON.stringify({ + 'google-antigravity': { type: 'oauth', ...hostAuth }, + }), + ) +} + describe('Pi provider multi-account integration', () => { + it('directs first-time /agy-add users through the public provider login flow', async () => { + await commands.get('agy-add')!.handler('', context()) + expect(exchange).not.toHaveBeenCalled() + expect(notify).toHaveBeenCalledWith( + 'Authenticate the provider first with /login google-antigravity, then use /agy-add for additional accounts.', + 'warning', + ) + }) + it('repeated /login and /agy-add share the OAuth flow and retain all accounts', async () => { hostAuth = await provider.oauth!.login(callbacks) + await persistHostAuth() hostAuth = await provider.oauth!.login(callbacks) await commands.get('agy-add')!.handler('', context()) const accounts = await core.loadAccountStorage( @@ -136,6 +155,7 @@ describe('Pi provider multi-account integration', () => { access: 'legacy-access', expires: Date.now() + 3600_000, } + await persistHostAuth() await hooks.get('session_start')?.({}, context()) await provider.oauth!.login(callbacks) expect( diff --git a/packages/pi/src/runtime.test.ts b/packages/pi/src/runtime.test.ts index 380d64b9..566b30ad 100644 --- a/packages/pi/src/runtime.test.ts +++ b/packages/pi/src/runtime.test.ts @@ -245,7 +245,7 @@ describe('Pi shared account runtime', () => { await pool() const next = runtime() expect(await dispatch(next)).toBe(0) - expect(refresh).toHaveBeenCalledWith('secret-refresh-1') + expect(refresh.mock.calls[0]?.[0]).toBe('secret-refresh-1') expect(await next.describe()).toContain('agy3') }) @@ -571,6 +571,41 @@ describe('Pi shared account runtime', () => { expect(next.refresh).toStartWith('secret-refresh-2|') }) + it('forwards and honors the Pi OAuth refresh abort signal', async () => { + const seeded = runtime() + await seeded.login(login(1)) + const controller = new AbortController() + let startRefresh: (() => void) | undefined + const refreshStarted = new Promise((resolve) => { + startRefresh = resolve + }) + const refreshing = new PiAccountRuntime({ + path, + pid: 0, + refreshToken: async (_token, signal) => { + startRefresh?.() + return await new Promise((_resolve, reject) => { + signal?.addEventListener('abort', () => reject(signal.reason), { + once: true, + }) + }) + }, + }) + runtimes.push(refreshing) + + const result = refreshing.refreshHost( + { ...login(1), expires: 0 }, + controller.signal, + ) + await refreshStarted + controller.abort(new Error('cancelled by Pi')) + + await expect(result).rejects.toThrow('cancelled by Pi') + expect( + (await loadAccountStorage(path))?.accounts[0]?.cooldownReason, + ).toBeUndefined() + }) + it('migrates a legacy Pi credential once without overwriting or resurrecting accounts', async () => { const result = runtime() await result.migrate({ ...login(1) }) diff --git a/packages/pi/src/runtime.ts b/packages/pi/src/runtime.ts index 81a609f0..a0513e4e 100644 --- a/packages/pi/src/runtime.ts +++ b/packages/pi/src/runtime.ts @@ -229,9 +229,14 @@ export class PiAccountRuntime { return applied } - private async credentialFor(token: string): Promise { + private async credentialFor( + token: string, + signal?: AbortSignal, + ): Promise { + signal?.throwIfAborted() let auth: OAuthAuthDetails | undefined await mutateAccountStorage(this.path, async (current) => { + signal?.throwIfAborted() const account = current.accounts.find( (entry) => entry.refreshToken === token, ) @@ -250,7 +255,8 @@ export class PiAccountRuntime { } let refreshed: Awaited> try { - refreshed = await this.refreshToken(token) + refreshed = await this.refreshToken(token, signal) + signal?.throwIfAborted() if ( !refreshed.access || !refreshed.refresh || @@ -258,7 +264,8 @@ export class PiAccountRuntime { ) { throw new Error('Invalid refresh result') } - } catch { + } catch (error) { + if (signal?.aborted) throw error throw new CredentialRefreshFailed( 'Antigravity token refresh failed; re-authenticate the account', ) @@ -283,12 +290,21 @@ export class PiAccountRuntime { /** Pi refreshes its host credential before stream dispatch. A failed last * login must not prevent a healthy pool member from reaching the runtime. */ - async refreshHost(credentials: OAuthCredentials): Promise { + async refreshHost( + credentials: OAuthCredentials, + signal?: AbortSignal, + ): Promise { + signal?.throwIfAborted() await this.migrate(credentials) + signal?.throwIfAborted() const manager = await this.reload() for (const account of manager.getEnabledAccounts()) { + signal?.throwIfAborted() try { - const auth = await this.credentialFor(account.parts.refreshToken) + const auth = await this.credentialFor( + account.parts.refreshToken, + signal, + ) return { refresh: auth.refresh, access: auth.access ?? '', diff --git a/packages/pi/src/stream.test.ts b/packages/pi/src/stream.test.ts index 5e061e3f..be9fe563 100644 --- a/packages/pi/src/stream.test.ts +++ b/packages/pi/src/stream.test.ts @@ -403,6 +403,35 @@ describe('convertGeminiToolCallPart', () => { }) describe('streamCortexKitAntigravity', () => { + it('runs Pi payload and response hooks around the provider request', async () => { + const calls: string[] = [] + const response = sseResponse([ + 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"answer"}]},"finishReason":"STOP"}]}}\n\n', + ]) + fetchWithAgyCliTransportMock.mockImplementationOnce(async () => response) + + const eventStream = streamCortexKitAntigravity(fakeModel(), userContext(), { + apiKey: 'test-token', + sessionId: 'pi-hooks', + onPayload: (payload) => { + calls.push('payload') + return { ...(payload as object), requestType: 'hooked' } + }, + onResponse: (received) => { + calls.push('response') + expect(received.status).toBe(response.status) + expect(response.body?.locked).toBe(false) + }, + }) + + expect((await eventStream.result()).stopReason).toBe('stop') + expect(calls).toEqual(['payload', 'response']) + expect( + JSON.parse(String(fetchWithAgyCliTransportMock.mock.calls[0]?.[1]?.body)) + .requestType, + ).toBe('hooked') + }) + it('surfaces an embedded SSE error instead of returning an empty success', async () => { const { events, result } = await runStream( fakeModel(), diff --git a/packages/pi/src/stream.ts b/packages/pi/src/stream.ts index 82f8c49a..46fef023 100644 --- a/packages/pi/src/stream.ts +++ b/packages/pi/src/stream.ts @@ -409,10 +409,13 @@ async function sendAntigravityRequest(options: { userAgent: 'antigravity', requestType: 'agent', } + const payload = + (await options.streamOptions?.onPayload?.(envelope, options.model)) ?? + envelope const url = `${ANTIGRAVITY_ENDPOINT}/v1internal:${STREAM_ACTION}?alt=sse` - return fetchWithAgyCliTransport( + const response = await fetchWithAgyCliTransport( url, { method: 'POST', @@ -425,10 +428,18 @@ async function sendAntigravityRequest(options: { : {}), 'Accept-Encoding': 'gzip', }, - body: JSON.stringify(envelope), + body: JSON.stringify(payload), }, { signal: options.signal ?? options.streamOptions?.signal ?? null }, ) + await options.streamOptions?.onResponse?.( + { + status: response.status, + headers: Object.fromEntries(response.headers), + }, + options.model, + ) + return response } export function streamCortexKitAntigravity( From 1e5759b6ee048f759b6ec8723d2cc71c83256eb4 Mon Sep 17 00:00:00 2001 From: Brent Duarte Date: Wed, 9 Sep 2026 22:56:04 -0700 Subject: [PATCH 6/7] fix(pi): reconcile canonical OAuth account identity --- packages/core/src/account-manager.ts | 33 ++- packages/core/src/account-storage.test.ts | 14 ++ packages/core/src/account-storage.ts | 16 +- packages/core/src/account-types.ts | 2 + packages/core/src/antigravity/oauth.test.ts | 3 +- packages/core/src/antigravity/oauth.ts | 58 +++-- packages/core/src/persist-account-pool.ts | 217 +++++++++++++++++- packages/core/src/quota-manager.ts | 8 +- packages/core/src/rotation.ts | 23 +- .../src/plugin/persist-account-pool.test.ts | 39 ++++ packages/pi/src/commands.ts | 10 +- packages/pi/src/index.ts | 8 +- packages/pi/src/provider.test.ts | 117 ++++++++++ packages/pi/src/runtime.test.ts | 213 ++++++++++++++++- packages/pi/src/runtime.ts | 184 ++++++++++++++- 15 files changed, 894 insertions(+), 51 deletions(-) diff --git a/packages/core/src/account-manager.ts b/packages/core/src/account-manager.ts index fd22098f..26316d6b 100644 --- a/packages/core/src/account-manager.ts +++ b/packages/core/src/account-manager.ts @@ -76,6 +76,7 @@ export type QuotaKey = BaseQuotaKey | `${BaseQuotaKey}:${string}` export interface ManagedAccount { index: number email?: string + accountId?: string label?: string addedAt: number lastUsed: number @@ -411,6 +412,7 @@ export class AccountManager { return { index, email: acc.email, + accountId: acc.accountId, label: acc.label, addedAt: clampNonNegativeInt(acc.addedAt, baseNow), lastUsed: clampNonNegativeInt(acc.lastUsed, 0), @@ -550,7 +552,8 @@ export class AccountManager { /** Reload durable metadata without resetting this process's routing state. * Access tokens survive only an exact credential match. Routing and scoring - * follow unique normalized email, falling back to token for legacy accounts. + * follow unique normalized email, then Google identity, falling back to the + * token only for legacy accounts without either stable identity. */ reconcileStorage(stored: AccountStorageV4): void { const previous = this.accounts @@ -566,7 +569,9 @@ export class AccountManager { }) const keyOf = (account: ManagedAccount): string => { const email = account.email?.trim().toLowerCase() - return email ? `email:${email}` : `token:${account.parts.refreshToken}` + if (email) return `email:${email}` + if (account.accountId) return `account:${account.accountId}` + return `token:${account.parts.refreshToken}` } const remap = (index: number): number => { const old = previous[index] @@ -575,11 +580,30 @@ export class AccountManager { if (previous.filter((account) => keyOf(account) === key).length !== 1) return -1 const matches = fresh.accounts.filter((account) => keyOf(account) === key) - return matches.length === 1 ? matches[0]!.index : -1 + if (matches.length === 1) return matches[0]!.index + const tokenMatches = fresh.accounts.filter( + (account) => account.parts.refreshToken === old.parts.refreshToken, + ) + return tokenMatches.length === 1 ? tokenMatches[0]!.index : -1 } const indexMap = new Map( previous.map((account) => [account.index, remap(account.index)]), ) + const trackerIndexMap = new Map(indexMap) + for (const freshAccount of fresh.accounts) { + const previousMatches = previous.filter( + (account) => indexMap.get(account.index) === freshAccount.index, + ) + if (previousMatches.length < 2) continue + const exactToken = previousMatches.find( + (account) => + account.parts.refreshToken === freshAccount.parts.refreshToken, + ) + const survivor = exactToken ?? previousMatches[0] + for (const account of previousMatches) { + if (account !== survivor) trackerIndexMap.set(account.index, -1) + } + } this.accounts = fresh.accounts.map((account) => { const old = previous.find( (entry) => indexMap.get(entry.index) === account.index, @@ -617,6 +641,7 @@ export class AccountManager { reconcileAccountTrackers( previous.map(trackerIdentity), this.accounts.map(trackerIdentity), + trackerIndexMap, ) if (changed) { this.sessionUsedAccounts = remapUsed(this.sessionUsedAccounts) @@ -1784,6 +1809,7 @@ export class AccountManager { version: 4, accounts: this.accounts.map((a) => ({ email: a.email, + accountId: a.accountId, label: a.label, refreshToken: a.parts.refreshToken, projectId: a.parts.projectId ?? a.projectId, @@ -2255,6 +2281,7 @@ export class AccountManager { getAccountsForQuotaCheck(): AccountMetadataV3[] { return this.accounts.map((a) => ({ email: a.email, + accountId: a.accountId, refreshToken: a.parts.refreshToken, projectId: a.parts.projectId ?? a.projectId, managedProjectId: a.parts.managedProjectId ?? a.managedProjectId, diff --git a/packages/core/src/account-storage.test.ts b/packages/core/src/account-storage.test.ts index b1cd03c6..248e3d4d 100644 --- a/packages/core/src/account-storage.test.ts +++ b/packages/core/src/account-storage.test.ts @@ -128,6 +128,20 @@ describe('deduplicateAccountsByEmail', () => { expect(deduplicateAccountsByEmail(accounts)).toHaveLength(1) expect(deduplicateAccountsByEmail(accounts)[0]?.refreshToken).toBe('new') }) + + it('normalizes email before deduplicating', () => { + const accounts: AccountMetadataV3[] = [ + { + email: ' A@Example.com ', + refreshToken: 'old', + addedAt: 1, + lastUsed: 1, + }, + { email: 'a@example.com', refreshToken: 'new', addedAt: 2, lastUsed: 9 }, + ] + expect(deduplicateAccountsByEmail(accounts)).toHaveLength(1) + expect(deduplicateAccountsByEmail(accounts)[0]?.refreshToken).toBe('new') + }) }) describe('mergeAccountStorage', () => { diff --git a/packages/core/src/account-storage.ts b/packages/core/src/account-storage.ts index 61968737..9645c2fe 100644 --- a/packages/core/src/account-storage.ts +++ b/packages/core/src/account-storage.ts @@ -211,15 +211,20 @@ export function deduplicateAccountsByEmail< continue } - const existingIndex = emailToNewestIndex.get(acc.email) + const email = acc.email.trim().toLowerCase() + if (!email) { + indicesToKeep.add(i) + continue + } + const existingIndex = emailToNewestIndex.get(email) if (existingIndex === undefined) { - emailToNewestIndex.set(acc.email, i) + emailToNewestIndex.set(email, i) continue } const existing = accounts[existingIndex] if (!existing) { - emailToNewestIndex.set(acc.email, i) + emailToNewestIndex.set(email, i) continue } @@ -233,7 +238,7 @@ export function deduplicateAccountsByEmail< (currLastUsed === existLastUsed && currAddedAt > existAddedAt) if (isNewer) { - emailToNewestIndex.set(acc.email, i) + emailToNewestIndex.set(email, i) } } @@ -561,6 +566,9 @@ function validateV4AccountRecord( if (typeof acc.lastUsed !== 'number' || !Number.isFinite(acc.lastUsed)) { return `accounts[${index}].lastUsed is missing or not a finite number` } + if (acc.accountId !== undefined && typeof acc.accountId !== 'string') { + return `accounts[${index}].accountId is not a string` + } return null } diff --git a/packages/core/src/account-types.ts b/packages/core/src/account-types.ts index 47ae0220..ebdecaa5 100644 --- a/packages/core/src/account-types.ts +++ b/packages/core/src/account-types.ts @@ -36,6 +36,8 @@ export type CooldownReason = export interface AccountMetadataV3 { email?: string + /** Stable Google userinfo identity. Never shown in operator output. */ + accountId?: string refreshToken: string projectId?: string managedProjectId?: string diff --git a/packages/core/src/antigravity/oauth.test.ts b/packages/core/src/antigravity/oauth.test.ts index 496884fb..8771f8a2 100644 --- a/packages/core/src/antigravity/oauth.test.ts +++ b/packages/core/src/antigravity/oauth.test.ts @@ -47,7 +47,7 @@ function userInfoBody( email = 'user@example.com', name = 'Alice Example', ): string { - return JSON.stringify({ email, name }) + return JSON.stringify({ id: 'google-account-1', email, name }) } describe('Antigravity OAuth', () => { @@ -147,6 +147,7 @@ describe('Antigravity OAuth', () => { expect(result.refresh).toBe('refresh-1|project-1') expect(result.access).toBe('access-1') expect(result.email).toBe('alice@example.com') + expect(result.accountId).toBe('google-account-1') expect(result.label).toBe('Alice Example') expect(result.projectId).toBe('project-1') diff --git a/packages/core/src/antigravity/oauth.ts b/packages/core/src/antigravity/oauth.ts index f514b3ac..a0f3f45e 100644 --- a/packages/core/src/antigravity/oauth.ts +++ b/packages/core/src/antigravity/oauth.ts @@ -52,6 +52,7 @@ interface AntigravityTokenExchangeSuccess { access: string expires: number email?: string + accountId?: string label?: string projectId: string } @@ -123,10 +124,50 @@ interface AntigravityTokenResponse { } interface AntigravityUserInfo { + id?: string email?: string name?: string } +export interface AntigravityAccountIdentity { + accountId?: string + email?: string + label?: string +} + +/** Resolve the stable Google identity carried by an OAuth access token. */ +export async function fetchAntigravityAccountIdentity( + accessToken: string, + signal?: AbortSignal, +): Promise { + const response = await fetchWithActiveTimeout( + 'https://www.googleapis.com/oauth2/v1/userinfo?alt=json', + { + headers: { + Authorization: `Bearer ${accessToken}`, + 'User-Agent': GEMINI_CLI_HEADERS['User-Agent'], + }, + signal, + }, + ) + if (!response.ok) return {} + const userInfo = (await response.json()) as AntigravityUserInfo + const email = + typeof userInfo.email === 'string' + ? userInfo.email.trim().toLowerCase() + : undefined + const accountId = + typeof userInfo.id === 'string' ? userInfo.id.trim() : undefined + return { + email: email || undefined, + accountId: accountId || undefined, + label: + typeof userInfo.name === 'string' + ? userInfo.name.trim() || undefined + : undefined, + } +} + /** * Encode an object into a URL-safe base64 string. */ @@ -292,20 +333,10 @@ export async function exchangeAntigravity( const tokenPayload = (await tokenResponse.json()) as AntigravityTokenResponse - const userInfoResponse = await fetchWithActiveTimeout( - 'https://www.googleapis.com/oauth2/v1/userinfo?alt=json', - { - headers: { - Authorization: `Bearer ${tokenPayload.access_token}`, - 'User-Agent': GEMINI_CLI_HEADERS['User-Agent'], - }, - }, + const userInfo = await fetchAntigravityAccountIdentity( + tokenPayload.access_token, ) - const userInfo = userInfoResponse.ok - ? ((await userInfoResponse.json()) as AntigravityUserInfo) - : {} - const refreshToken = tokenPayload.refresh_token if (!refreshToken) { return { type: 'failed', error: 'Missing refresh token in response' } @@ -324,7 +355,8 @@ export async function exchangeAntigravity( access: tokenPayload.access_token, expires: calculateTokenExpiry(startTime, tokenPayload.expires_in), email: userInfo.email, - label: userInfo.name?.trim() || undefined, + accountId: userInfo.accountId, + label: userInfo.label, projectId: effectiveProjectId || '', } } catch (error) { diff --git a/packages/core/src/persist-account-pool.ts b/packages/core/src/persist-account-pool.ts index 4d1d897b..25fd86d8 100644 --- a/packages/core/src/persist-account-pool.ts +++ b/packages/core/src/persist-account-pool.ts @@ -7,9 +7,10 @@ * state read while the lock is held — without it, a concurrent add * would race the read-modify-write and silently disappear. * - * Two upsert keys are honored, in priority order: + * Three upsert keys are honored, in priority order: * 1. email — survives refresh-token rotation for the same Google account - * 2. refresh token — handles the no-email case and out-of-band rotations + * 2. Google account ID — enriches a canonical login whose email was absent + * 3. refresh token — handles the no-email case and out-of-band rotations * * Destructive (`replaceAll: true`) writes start from an empty v4 inside * the same locked callback so a stale merge cannot resurrect a removed @@ -23,6 +24,21 @@ import { parseRefreshParts } from './auth.ts' type TokenSuccess = Extract +export interface ResolvedAccountIdentity { + refreshToken: string + email?: string + accountId?: string +} + +export class AccountIdentityAmbiguityError extends Error { + override readonly name = 'AccountIdentityAmbiguityError' +} + +function normalizeEmail(email: string | undefined): string | undefined { + const normalized = email?.trim().toLowerCase() + return normalized || undefined +} + function clampInt(value: number, min: number, max: number): number { if (!Number.isFinite(value)) { return min @@ -43,6 +59,7 @@ function applyUpserts( const indexByRefreshToken = new Map() const indexByEmail = new Map() + const indexByAccountId = new Map() for (let i = 0; i < accounts.length; i++) { const acc = accounts[i] if (!acc) continue @@ -50,7 +67,11 @@ function applyUpserts( indexByRefreshToken.set(acc.refreshToken, i) } if (acc.email) { - indexByEmail.set(acc.email, i) + const email = normalizeEmail(acc.email) + if (email) indexByEmail.set(email, i) + } + if (acc.accountId) { + indexByAccountId.set(acc.accountId, i) } } @@ -62,20 +83,36 @@ function applyUpserts( // Email match wins over token match — handles refresh-token rotation // for the same Google account. - const existingByEmail = result.email - ? indexByEmail.get(result.email) + const email = normalizeEmail(result.email) + const existingByEmail = email ? indexByEmail.get(email) : undefined + const existingByAccountId = result.accountId + ? indexByAccountId.get(result.accountId) : undefined const existingByToken = indexByRefreshToken.get(parts.refreshToken) - const existingIndex = existingByEmail ?? existingByToken + if ( + existingByEmail !== undefined && + existingByAccountId !== undefined && + existingByEmail !== existingByAccountId + ) { + throw new AccountIdentityAmbiguityError( + 'OAuth email and Google account identity resolve to different stored accounts', + ) + } + const existingIndex = + existingByEmail ?? existingByAccountId ?? existingByToken if (existingIndex === undefined) { const newIndex = accounts.length indexByRefreshToken.set(parts.refreshToken, newIndex) - if (result.email) { - indexByEmail.set(result.email, newIndex) + if (email) { + indexByEmail.set(email, newIndex) + } + if (result.accountId) { + indexByAccountId.set(result.accountId, newIndex) } accounts.push({ - email: result.email, + email, + accountId: result.accountId, label: result.label, refreshToken: parts.refreshToken, projectId: parts.projectId, @@ -89,11 +126,21 @@ function applyUpserts( const existing = accounts[existingIndex] if (!existing) continue + if ( + existing.accountId && + result.accountId && + existing.accountId !== result.accountId + ) { + throw new AccountIdentityAmbiguityError( + 'OAuth identity conflicts with the matched stored account', + ) + } const oldToken = existing.refreshToken accounts[existingIndex] = { ...existing, - email: result.email ?? existing.email, + email: email ?? existing.email, + accountId: result.accountId ?? existing.accountId, label: result.label ?? existing.label, refreshToken: parts.refreshToken, projectId: parts.projectId ?? existing.projectId, @@ -105,6 +152,8 @@ function applyUpserts( indexByRefreshToken.delete(oldToken) indexByRefreshToken.set(parts.refreshToken, existingIndex) } + if (email) indexByEmail.set(email, existingIndex) + if (result.accountId) indexByAccountId.set(result.accountId, existingIndex) } if (accounts.length === 0) { @@ -130,6 +179,154 @@ function applyUpserts( } } +function maxRateLimits( + first: AccountMetadataV3['rateLimitResetTimes'], + second: AccountMetadataV3['rateLimitResetTimes'], +): AccountMetadataV3['rateLimitResetTimes'] { + const merged = { ...first } + for (const [key, value] of Object.entries(second ?? {})) { + merged[key] = Math.max(merged[key] ?? 0, value ?? 0) + } + return Object.keys(merged).length ? merged : undefined +} + +function mergeProvenDuplicate( + survivor: AccountMetadataV3, + duplicate: AccountMetadataV3, +): AccountMetadataV3 { + const duplicateQuotaIsNewer = + (duplicate.cachedQuotaUpdatedAt ?? 0) > (survivor.cachedQuotaUpdatedAt ?? 0) + const coolingDownUntil = Math.max( + survivor.coolingDownUntil ?? 0, + duplicate.coolingDownUntil ?? 0, + ) + return { + ...duplicate, + ...survivor, + email: survivor.email ?? duplicate.email, + accountId: survivor.accountId ?? duplicate.accountId, + label: survivor.label ?? duplicate.label, + projectId: survivor.projectId ?? duplicate.projectId, + managedProjectId: survivor.managedProjectId ?? duplicate.managedProjectId, + addedAt: Math.min(survivor.addedAt, duplicate.addedAt), + lastUsed: Math.max(survivor.lastUsed, duplicate.lastUsed), + enabled: + survivor.accountIneligible || duplicate.accountIneligible + ? false + : survivor.enabled, + rateLimitResetTimes: maxRateLimits( + survivor.rateLimitResetTimes, + duplicate.rateLimitResetTimes, + ), + coolingDownUntil: coolingDownUntil || undefined, + cooldownReason: + (duplicate.coolingDownUntil ?? 0) > (survivor.coolingDownUntil ?? 0) + ? duplicate.cooldownReason + : survivor.cooldownReason, + cachedQuota: duplicateQuotaIsNewer + ? duplicate.cachedQuota + : survivor.cachedQuota, + cachedPerModelQuota: duplicateQuotaIsNewer + ? duplicate.cachedPerModelQuota + : survivor.cachedPerModelQuota, + cachedQuotaUpdatedAt: duplicateQuotaIsNewer + ? duplicate.cachedQuotaUpdatedAt + : survivor.cachedQuotaUpdatedAt, + verificationRequired: + survivor.verificationRequired || duplicate.verificationRequired, + accountIneligible: + survivor.accountIneligible || duplicate.accountIneligible, + } +} + +function remapRemovedIndex( + index: number | undefined, + removed: number, + survivor: number, +): number | undefined { + if (index === undefined) return undefined + if (index === removed) return survivor + return index > removed ? index - 1 : index +} + +/** + * Enrich an exact token-only account with identity resolved from Google and + * collapse only duplicates proven by that identity. The exact token is the + * fence: if a peer has replaced it, reconciliation fails without guessing. + */ +export async function reconcileAccountIdentityAtPath( + path: string, + identity: ResolvedAccountIdentity, +): Promise { + const email = normalizeEmail(identity.email) + const accountId = identity.accountId?.trim() || undefined + if (!email && !accountId) { + throw new AccountIdentityAmbiguityError( + 'Google account identity was unavailable; no accounts were merged', + ) + } + await mutateAccountStorage(path, (current) => { + let survivorIndex = current.accounts.findIndex( + (account) => account.refreshToken === identity.refreshToken, + ) + if (survivorIndex < 0) { + throw new AccountIdentityAmbiguityError( + 'The token-only account changed before identity reconciliation', + ) + } + const survivor = current.accounts[survivorIndex]! + if (survivor.accountId && accountId && survivor.accountId !== accountId) { + throw new AccountIdentityAmbiguityError( + 'The stored Google account identity conflicts with the resolved identity', + ) + } + survivor.email = email ?? survivor.email + survivor.accountId = accountId ?? survivor.accountId + + for (let index = current.accounts.length - 1; index >= 0; index--) { + if (index === survivorIndex) continue + const candidate = current.accounts[index]! + const candidateEmail = normalizeEmail(candidate.email) + const emailMatch = !!email && candidateEmail === email + const accountIdMatch = !!accountId && candidate.accountId === accountId + if (!emailMatch && !accountIdMatch) continue + if ( + emailMatch && + candidate.accountId && + accountId && + candidate.accountId !== accountId + ) { + throw new AccountIdentityAmbiguityError( + 'Matching email records have conflicting Google account identities', + ) + } + current.accounts[survivorIndex] = mergeProvenDuplicate( + current.accounts[survivorIndex]!, + candidate, + ) + const survivorAfterRemoval = + index < survivorIndex ? survivorIndex - 1 : survivorIndex + current.accounts.splice(index, 1) + current.activeIndex = + remapRemovedIndex(current.activeIndex, index, survivorAfterRemoval) ?? 0 + if (current.activeIndexByFamily) { + current.activeIndexByFamily.claude = remapRemovedIndex( + current.activeIndexByFamily.claude, + index, + survivorAfterRemoval, + ) + current.activeIndexByFamily.gemini = remapRemovedIndex( + current.activeIndexByFamily.gemini, + index, + survivorAfterRemoval, + ) + } + survivorIndex = survivorAfterRemoval + } + return current + }) +} + /** * Merge a batch of successful OAuth results into the persisted pool. * diff --git a/packages/core/src/quota-manager.ts b/packages/core/src/quota-manager.ts index 5d7f8642..061dfb71 100644 --- a/packages/core/src/quota-manager.ts +++ b/packages/core/src/quota-manager.ts @@ -111,11 +111,13 @@ interface AccountState { } /** - * Default keyOf — prefers email, falls back to refresh-token hash so the - * same identity is keyed even when emails are missing. + * Default keyOf — prefers normalized email, then stable Google identity, + * falling back to a refresh-token hash only when identity is unavailable. */ export function defaultKeyOf(account: AccountMetadataV3): string { - if (account.email) return `e:${account.email.toLowerCase()}` + const email = account.email?.trim().toLowerCase() + if (email) return `e:${email}` + if (account.accountId) return `a:${account.accountId}` const token = account.refreshToken || '' return `t:${createHash('sha256').update(token).digest('hex').slice(0, 16)}` } diff --git a/packages/core/src/rotation.ts b/packages/core/src/rotation.ts index 58681f9c..529af5e3 100644 --- a/packages/core/src/rotation.ts +++ b/packages/core/src/rotation.ts @@ -584,6 +584,7 @@ const trackerLayouts = new WeakMap< export function reconcileAccountTrackers( previousIdentities: readonly string[], nextIdentities: readonly string[], + authoritativeIndexMap?: ReadonlyMap, ): void { const uniqueLayout = (identities: readonly string[]) => { const counts = new Map() @@ -606,14 +607,24 @@ export function reconcileAccountTrackers( (identity, index) => identity === null || identity !== next[index], ) ) { - const indexMap = new Map() - for (const [index, identity] of previous.entries()) { - if (identity !== null) { - const nextIndex = next.indexOf(identity) - if (nextIndex >= 0) indexMap.set(index, nextIndex) + const expectedPrevious = uniqueLayout(previousIdentities) + const ownsExpectedLayout = + previous.length === expectedPrevious.length && + previous.every( + (identity, index) => identity === expectedPrevious[index], + ) + if (authoritativeIndexMap && ownsExpectedLayout) { + tracker.remapAccounts(authoritativeIndexMap) + } else { + const indexMap = new Map() + for (const [index, identity] of previous.entries()) { + if (identity !== null) { + const nextIndex = next.indexOf(identity) + if (nextIndex >= 0) indexMap.set(index, nextIndex) + } } + tracker.remapAccounts(indexMap) } - tracker.remapAccounts(indexMap) } trackerLayouts.set(tracker, next) } diff --git a/packages/opencode/src/plugin/persist-account-pool.test.ts b/packages/opencode/src/plugin/persist-account-pool.test.ts index 3c4a812d..45abfd93 100644 --- a/packages/opencode/src/plugin/persist-account-pool.test.ts +++ b/packages/opencode/src/plugin/persist-account-pool.test.ts @@ -430,6 +430,45 @@ describe('persistAccountPool behavior (lock-held persistence)', () => { expect(matches).toHaveLength(1) }) + it('enriches a token-only account by stable Google account identity after token rotation', async () => { + await storageModule.saveAccountsReplace( + createMockStorage([ + createMockAccount({ + email: undefined, + accountId: 'google-account-a', + refreshToken: 'canonical-token', + enabled: false, + }), + ]), + ) + + await persistAccountPool( + [ + { + type: 'success', + refresh: 'rotated-token|project-a', + access: 'new-access', + expires: Date.now() + 3_600_000, + email: ' A@Example.com ', + accountId: 'google-account-a', + projectId: 'project-a', + }, + ], + false, + ) + + expect(await storageModule.loadAccounts()).toMatchObject({ + accounts: [ + { + email: 'a@example.com', + accountId: 'google-account-a', + refreshToken: 'rotated-token', + enabled: false, + }, + ], + }) + }) + it('preserves activeIndex when adding new accounts (replaceAll=false)', async () => { await storageModule.saveAccountsReplace( createMockStorage( diff --git a/packages/pi/src/commands.ts b/packages/pi/src/commands.ts index 2981a33a..a0b99465 100644 --- a/packages/pi/src/commands.ts +++ b/packages/pi/src/commands.ts @@ -1,3 +1,4 @@ +import { AccountIdentityAmbiguityError } from '@cortexkit/antigravity-auth-core' import type { OAuthCredentials, OAuthLoginCallbacks, @@ -24,7 +25,14 @@ export function registerAccountCommands( handler: async (args, context) => { try { await handler(args, context) - } catch { + } catch (error) { + if (error instanceof AccountIdentityAmbiguityError) { + context.ui.notify( + 'Antigravity account identity is ambiguous; no accounts were merged. Re-authenticate or repair the token-only entries before retrying.', + 'error', + ) + return + } context.ui.notify( 'Antigravity command failed. Check the account/ settings file, account number, or re-authenticate. Existing credentials were retained.', 'error', diff --git a/packages/pi/src/index.ts b/packages/pi/src/index.ts index bf083494..fb62d47d 100644 --- a/packages/pi/src/index.ts +++ b/packages/pi/src/index.ts @@ -1,4 +1,5 @@ import { + AccountIdentityAmbiguityError, authorizeAntigravity, exchangeAntigravity, getPublicModelDefinitions, @@ -67,6 +68,7 @@ async function loginAntigravity( access: result.access, expires: result.expires, email: result.email, + accountId: result.accountId, } } @@ -81,9 +83,11 @@ export default function cortexKitPiAntigravityAuth(pi: ExtensionAPI): void { if (auth?.type === 'oauth') { await runtime.migrate(auth) } - } catch { + } catch (error) { context.ui.notify( - 'Antigravity account migration failed; existing auth and pool were retained. Repair the account file before retrying.', + error instanceof AccountIdentityAmbiguityError + ? 'Antigravity account migration found ambiguous token-only identity; no accounts were merged. Re-authenticate or repair the account file before retrying.' + : 'Antigravity account migration failed; existing auth and pool were retained. Repair the account file before retrying.', 'error', ) } diff --git a/packages/pi/src/provider.test.ts b/packages/pi/src/provider.test.ts index b219c7f9..d4a399b4 100644 --- a/packages/pi/src/provider.test.ts +++ b/packages/pi/src/provider.test.ts @@ -71,6 +71,7 @@ beforeEach(async () => { return { type: 'success', email: `person${i}@example.com`, + accountId: `google-person-${i}`, refresh: `refresh-${i}|project|managed`, access: `access-${i}`, expires: Date.now() + 3600_000, @@ -128,6 +129,54 @@ describe('Pi provider multi-account integration', () => { ) }) + it('canonical login persists and returns the stable OAuth identity', async () => { + hostAuth = await provider.oauth!.login(callbacks) + expect(hostAuth).toMatchObject({ + email: 'person1@example.com', + accountId: 'google-person-1', + }) + const accounts = await core.loadAccountStorage( + process.env.PI_ANTIGRAVITY_AUTH_FILE!, + ) + expect(accounts?.accounts).toEqual([ + expect.objectContaining({ + email: 'person1@example.com', + accountId: 'google-person-1', + }), + ]) + }) + + it('canonical login and agy-add re-authenticate the same stable account in place', async () => { + exchange.mockImplementation(async () => { + const i = ++sequence + return { + type: 'success', + email: 'same@example.com', + accountId: 'google-same-account', + refresh: `same-refresh-${i}|project|managed`, + access: `same-access-${i}`, + expires: Date.now() + 3_600_000, + projectId: 'project', + } + }) + hostAuth = await provider.oauth!.login(callbacks) + await persistHostAuth() + await commands.get('agy-disable')!.handler('agy1', context()) + await commands.get('agy-add')!.handler('', context()) + + expect( + (await core.loadAccountStorage(process.env.PI_ANTIGRAVITY_AUTH_FILE!)) + ?.accounts, + ).toEqual([ + expect.objectContaining({ + email: 'same@example.com', + accountId: 'google-same-account', + refreshToken: 'same-refresh-2', + enabled: false, + }), + ]) + }) + it('repeated /login and /agy-add share the OAuth flow and retain all accounts', async () => { hostAuth = await provider.oauth!.login(callbacks) await persistHostAuth() @@ -154,6 +203,7 @@ describe('Pi provider multi-account integration', () => { refresh: 'legacy|project|managed', access: 'legacy-access', expires: Date.now() + 3600_000, + email: 'legacy@example.com', } await persistHostAuth() await hooks.get('session_start')?.({}, context()) @@ -164,6 +214,73 @@ describe('Pi provider multi-account integration', () => { ).toHaveLength(2) }) + it('session_start safely reconciles the live canonical duplicate shape', async () => { + const accountPath = process.env.PI_ANTIGRAVITY_AUTH_FILE! + await core.mutateAccountStorage(accountPath, (current) => { + current.accounts.push( + { + refreshToken: 'legacy-a', + addedAt: 1, + lastUsed: 1, + enabled: false, + }, + { + email: 'person-b@example.com', + accountId: 'google-person-b', + refreshToken: 'account-b', + addedAt: 2, + lastUsed: 2, + enabled: true, + }, + { + email: 'person-a@example.com', + accountId: 'google-person-a', + refreshToken: 'duplicate-a', + addedAt: 3, + lastUsed: 3, + enabled: true, + }, + ) + return current + }) + hostAuth = { + refresh: 'duplicate-a|project|managed', + access: 'duplicate-access', + expires: Date.now() + 3_600_000, + email: 'person-a@example.com', + accountId: 'google-person-a', + } + await persistHostAuth() + globalThis.fetch = mock(async (input) => { + const url = String(input) + if (url.includes('oauth2.googleapis.com/token')) { + return Response.json({ + access_token: 'legacy-access', + expires_in: 3_600, + }) + } + if (url.includes('oauth2/v1/userinfo')) { + return Response.json({ + id: 'google-person-a', + email: 'person-a@example.com', + }) + } + return Response.json({ groups: [] }) + }) as unknown as typeof fetch + + await hooks.get('session_start')?.({}, context()) + + const accounts = (await core.loadAccountStorage(accountPath))!.accounts + expect(accounts).toHaveLength(2) + expect(accounts[0]).toMatchObject({ + email: 'person-a@example.com', + accountId: 'google-person-a', + refreshToken: 'duplicate-a', + enabled: false, + }) + expect(accounts[1]?.email).toBe('person-b@example.com') + }) + it('rejects a mismatched OAuth state and cancellation before exchange/persistence', async () => { await expect( provider.oauth!.login({ diff --git a/packages/pi/src/runtime.test.ts b/packages/pi/src/runtime.test.ts index 566b30ad..1fe57642 100644 --- a/packages/pi/src/runtime.test.ts +++ b/packages/pi/src/runtime.test.ts @@ -32,6 +32,7 @@ function login(index: number, token = `secret-refresh-${index}`) { return { type: 'success' as const, email: `user${index}@example.com`, + accountId: `google-user-${index}`, refresh: `${token}|project-${index}|managed-${index}`, access: `secret-access-${index}`, expires: Date.now() + 3600_000, @@ -169,10 +170,31 @@ describe('Pi shared account runtime', () => { expect(sends).toEqual(['secret-refresh-1', 'secret-refresh-2']) }) + it('does not retry a no-email account after its token rotates in one request', async () => { + const result = await pool(2) + await mutateAccountStorage(path, (current) => { + delete current.accounts[0]!.email + return current + }) + const sends: string[] = [] + await result.dispatch(model, async (_auth, account) => { + sends.push(account.parts.refreshToken) + if (sends.length > 1) return new Response('ok') + await mutateAccountStorage(path, (current) => { + current.accounts[0]!.refreshToken = 'rotated-no-email-token' + current.accounts[0]!.rateLimitResetTimes = {} + return current + }) + return new Response('', { status: 429 }) + }) + expect(sends).toEqual(['secret-refresh-1', 'secret-refresh-2']) + }) + it('stops failover when a dispatched no-email credential disappears', async () => { const result = await pool(2) await mutateAccountStorage(path, (current) => { delete current.accounts[0]!.email + delete current.accounts[0]!.accountId return current }) const send = mock(async () => { @@ -199,7 +221,10 @@ describe('Pi shared account runtime', () => { const before = 1 await mutateAccountStorage(path, (current) => { current.accounts[0]!.cachedQuotaUpdatedAt = before - if (!withEmail) delete current.accounts[0]!.email + if (!withEmail) { + delete current.accounts[0]!.email + delete current.accounts[0]!.accountId + } return current }) const response = quotaFetch.getMockImplementation()! @@ -241,6 +266,163 @@ describe('Pi shared account runtime', () => { expect(await dispatch(result)).not.toBe(0) }) + it('enriches a canonical token-only account by stable Google identity', async () => { + const result = runtime() + await result.login({ + ...login(1), + email: undefined, + refresh: 'canonical-token|project-1|managed-1', + }) + await result.setEnabled(0, false) + await result.login(login(1, 'rotated-token')) + expect((await loadAccountStorage(path))?.accounts).toEqual([ + expect.objectContaining({ + email: 'user1@example.com', + accountId: 'google-user-1', + refreshToken: 'rotated-token', + enabled: false, + }), + ]) + }) + + it('keeps only canonical A and additional B after re-authenticating A', async () => { + const result = runtime() + await result.login(login(1, 'canonical-a')) + await result.login(login(2, 'additional-b')) + await result.login(login(1, 'rotated-a')) + const accounts = (await loadAccountStorage(path))!.accounts + expect(accounts).toHaveLength(2) + expect(accounts.map((account) => account.accountId)).toEqual([ + 'google-user-1', + 'google-user-2', + ]) + expect(accounts[0]?.refreshToken).toBe('rotated-a') + }) + + it('fails closed when token-only accounts cannot be identified', async () => { + const result = new PiAccountRuntime({ + path, + refreshToken: refresh, + fetchAccountIdentity: async () => ({}), + }) + runtimes.push(result) + await mutateAccountStorage(path, (current) => { + current.accounts.push( + { refreshToken: 'unknown-1', addedAt: 1, lastUsed: 1, enabled: true }, + { refreshToken: 'unknown-2', addedAt: 2, lastUsed: 2, enabled: false }, + ) + return current + }) + await expect(result.login(login(1, 'incoming-a'))).rejects.toThrow( + 'identity is ambiguous', + ) + expect((await loadAccountStorage(path))?.accounts).toHaveLength(2) + }) + + it('reconciles the live token-only plus email-duplicate shape to the original slot', async () => { + const result = new PiAccountRuntime({ + path, + refreshToken: refresh, + fetchAccountIdentity: async (access) => + access === 'access-legacy-a' + ? { email: 'user1@example.com', accountId: 'google-user-1' } + : {}, + }) + runtimes.push(result) + await mutateAccountStorage(path, (current) => { + current.accounts.push( + { + refreshToken: 'legacy-a', + addedAt: 1, + lastUsed: 1, + enabled: false, + coolingDownUntil: Date.now() + 60_000, + cooldownReason: 'auth-failure', + }, + { + email: 'user2@example.com', + accountId: 'google-user-2', + refreshToken: 'account-b', + addedAt: 2, + lastUsed: 2, + enabled: true, + }, + { + email: 'user1@example.com', + accountId: 'google-user-1', + refreshToken: 'duplicate-a', + addedAt: 3, + lastUsed: 3, + enabled: true, + }, + ) + current.activeIndex = 2 + return current + }) + await result.describe() + getHealthTracker().recordFailure(0) + getTokenTracker().consume(0) + const healthBefore = getHealthTracker().getScore(0) + const tokensBefore = getTokenTracker().getTokens(0) + await result.login(login(1, 'current-a')) + await result.describe() + const storage = (await loadAccountStorage(path))! + expect(storage.accounts).toHaveLength(2) + expect(storage.accounts[0]).toMatchObject({ + email: 'user1@example.com', + accountId: 'google-user-1', + refreshToken: 'current-a', + enabled: false, + cooldownReason: 'auth-failure', + }) + expect(storage.accounts[1]?.email).toBe('user2@example.com') + expect(storage.activeIndex).toBe(0) + expect(getHealthTracker().getScore(0)).toBe(healthBefore) + expect(getTokenTracker().getTokens(0)).toBeCloseTo(tokensBefore, 2) + }) + + it('leaves the live duplicate shape untouched when linkage cannot be proven', async () => { + const result = new PiAccountRuntime({ + path, + refreshToken: refresh, + fetchAccountIdentity: async () => ({}), + }) + runtimes.push(result) + await mutateAccountStorage(path, (current) => { + current.accounts.push( + { refreshToken: 'unknown-a', addedAt: 1, lastUsed: 1, enabled: false }, + { + email: 'user2@example.com', + accountId: 'google-user-2', + refreshToken: 'account-b', + addedAt: 2, + lastUsed: 2, + enabled: true, + }, + { + email: 'user1@example.com', + accountId: 'google-user-1', + refreshToken: 'possible-duplicate-a', + addedAt: 3, + lastUsed: 3, + enabled: true, + }, + ) + return current + }) + + await expect(result.login(login(1, 'incoming-a'))).rejects.toThrow( + 'identity is ambiguous', + ) + const accounts = (await loadAccountStorage(path))!.accounts + expect(accounts).toHaveLength(3) + expect(accounts.map((account) => account.refreshToken)).toEqual([ + 'unknown-a', + 'account-b', + 'possible-duplicate-a', + ]) + }) + it('reloads across sessions and refreshes the selected stored credential', async () => { await pool() const next = runtime() @@ -569,6 +751,8 @@ describe('Pi shared account runtime', () => { await result.setEnabled(0, false) const next = await result.refreshHost({ ...login(1), expires: 0 }) expect(next.refresh).toStartWith('secret-refresh-2|') + expect(next.email).toBe('user2@example.com') + expect(next.accountId).toBe('google-user-2') }) it('forwards and honors the Pi OAuth refresh abort signal', async () => { @@ -617,6 +801,33 @@ describe('Pi shared account runtime', () => { expect(stored?.accounts[0]?.refreshToken).toBe('new-refresh') }) + it('session migration enriches a canonical credential from its access token', async () => { + const result = new PiAccountRuntime({ + path, + refreshToken: refresh, + fetchAccountIdentity: async (access) => + access === 'canonical-access' + ? { + email: 'canonical@example.com', + accountId: 'google-canonical', + } + : {}, + }) + runtimes.push(result) + await result.migrate({ + refresh: 'canonical-refresh|project|managed', + access: 'canonical-access', + expires: Date.now() + 3_600_000, + }) + expect((await loadAccountStorage(path))?.accounts).toEqual([ + expect.objectContaining({ + email: 'canonical@example.com', + accountId: 'google-canonical', + refreshToken: 'canonical-refresh', + }), + ]) + }) + it.each([ '{broken secret-refresh-1', '{"version":99,"accounts":[]}', diff --git a/packages/pi/src/runtime.ts b/packages/pi/src/runtime.ts index a0513e4e..38d5de44 100644 --- a/packages/pi/src/runtime.ts +++ b/packages/pi/src/runtime.ts @@ -1,8 +1,10 @@ import { + AccountIdentityAmbiguityError, AccountManager, type AccountMetadataV3, type AccountModelFamily, ANTIGRAVITY_ENDPOINT_FALLBACKS, + type AntigravityAccountIdentity, type AntigravityTokenExchangeResult, accessTokenExpired, aggregateQuota, @@ -14,6 +16,7 @@ import { ensureProjectContext, extractRateLimitBodyInfo, type FetchQuotaSummaryOptions, + fetchAntigravityAccountIdentity, fetchAvailableModels, fetchQuotaSummary, fetchWithActiveTimeout, @@ -27,6 +30,7 @@ import { parseRateLimitReason, parseRefreshParts, persistAccountPoolAtPath, + reconcileAccountIdentityAtPath, refreshAntigravityToken, resolveQuotaGroup, retryAfterMsFromResponse, @@ -46,6 +50,7 @@ function attemptKey(account: ManagedAccount): string { return defaultKeyOf({ ...account.parts, email: account.email?.trim().toLowerCase(), + accountId: account.accountId, addedAt: account.addedAt, lastUsed: account.lastUsed, }) @@ -55,6 +60,7 @@ export interface PiRuntimeOptions { path?: string pid?: number refreshToken?: typeof refreshAntigravityToken + fetchAccountIdentity?: typeof fetchAntigravityAccountIdentity } /** Host wiring only: selection/scoring, quota aggregation, cooldowns and all @@ -67,6 +73,7 @@ export class PiAccountRuntime { private manager?: AccountManager private readonly credentials = new Map() private readonly refreshToken: typeof refreshAntigravityToken + private readonly fetchAccountIdentity: typeof fetchAntigravityAccountIdentity private readonly pid: number private selected?: string private legacy?: OAuthCredentials @@ -77,6 +84,8 @@ export class PiAccountRuntime { this.settingsPath = `${this.path}.config.json` this.pid = options.pid ?? process.pid this.refreshToken = options.refreshToken ?? refreshAntigravityToken + this.fetchAccountIdentity = + options.fetchAccountIdentity ?? fetchAntigravityAccountIdentity this.quota = createQuotaManager({ keyOf: defaultKeyOf, fetchAccountQuota: async (account, signal) => { @@ -153,13 +162,23 @@ export class PiAccountRuntime { async login(result: LoginResult): Promise { if (this.legacy) await this.migrate(this.legacy) + await this.reconcileIdentity( + { + refreshToken: parseRefreshParts(result.refresh).refreshToken, + email: result.email, + accountId: result.accountId, + }, + true, + ) await persistAccountPoolAtPath(this.path, [result]) this.remember({ refresh: result.refresh, access: result.access, expires: result.expires, email: result.email, + accountId: result.accountId, }) + this.legacy = undefined } /** Import the host credential only into a missing/empty pool. An existing @@ -167,16 +186,75 @@ export class PiAccountRuntime { */ async migrate(credentials: OAuthCredentials): Promise { this.remember(credentials) + // `legacy` is a one-shot bridge from Pi's single canonical credential. + // Keeping it armed would re-run migration before every pool request. + this.legacy = undefined const parts = parseRefreshParts(credentials.refresh) if (!parts.refreshToken) return const stored = await loadAccountStorage(this.path) - if (stored?.accounts.length) return + let email = + typeof credentials.email === 'string' + ? credentials.email.trim().toLowerCase() || undefined + : undefined + let accountId = + typeof credentials.accountId === 'string' + ? credentials.accountId.trim() || undefined + : undefined + if (!email && !accountId && credentials.access) { + try { + const resolved = await this.fetchAccountIdentity(credentials.access) + email = resolved.email + accountId = resolved.accountId + } catch { + // The stored access token may be expired or temporarily unreadable. + } + } + if (!email && !accountId) { + try { + const refreshed = await this.refreshToken(parts.refreshToken) + const resolved = await this.fetchAccountIdentity(refreshed.access) + email = resolved.email + accountId = resolved.accountId + } catch { + // Preserve the token-only credential; a future explicit login can + // retry resolution and will still fail closed before any merge. + } + } + if (stored?.accounts.length) { + const identity: AntigravityAccountIdentity = { email, accountId } + if (!identity.email && !identity.accountId) return + const reconciled = await this.reconcileIdentity( + { refreshToken: parts.refreshToken, ...identity }, + true, + ) + const current = await loadAccountStorage(this.path) + const belongsToPool = + reconciled || + current?.accounts.some( + (account) => account.refreshToken === parts.refreshToken, + ) + if (belongsToPool) { + await persistAccountPoolAtPath(this.path, [ + { + type: 'success', + refresh: credentials.refresh, + access: credentials.access, + expires: credentials.expires, + email: identity.email, + accountId: identity.accountId, + projectId: parts.projectId ?? '', + }, + ]) + this.remember(credentials) + } + return + } await mutateAccountStorage(this.path, (current) => { if (current.accounts.length) return current current.accounts.push({ ...parts, - email: - typeof credentials.email === 'string' ? credentials.email : undefined, + email, + accountId, addedAt: Date.now(), lastUsed: 0, enabled: true, @@ -185,6 +263,85 @@ export class PiAccountRuntime { }) } + private async reconcileIdentity( + incoming: { + refreshToken: string + email?: string + accountId?: string + }, + failOnUnresolved: boolean, + ): Promise { + const email = incoming.email?.trim().toLowerCase() || undefined + const accountId = incoming.accountId?.trim() || undefined + if (!email && !accountId) return false + const stored = await loadAccountStorage(this.path) + if (!stored?.accounts.length) return false + const candidates = stored.accounts + .map((account, index) => ({ account, index })) + .filter(({ account }) => !account.email) + if (!candidates.length) return false + + const matching: typeof candidates = [] + let unresolved = 0 + for (const candidate of candidates) { + let resolved: AntigravityAccountIdentity + if (candidate.account.refreshToken === incoming.refreshToken) { + resolved = { email, accountId } + } else if (candidate.account.accountId) { + resolved = { accountId: candidate.account.accountId } + } else { + try { + const refreshed = await this.refreshToken( + candidate.account.refreshToken, + ) + resolved = await this.fetchAccountIdentity(refreshed.access) + } catch { + unresolved++ + continue + } + } + const resolvedEmail = resolved.email?.trim().toLowerCase() + if ( + accountId && + resolved.accountId && + accountId !== resolved.accountId && + email && + resolvedEmail === email + ) { + throw new AccountIdentityAmbiguityError( + 'Antigravity account identity is ambiguous; matching email has conflicting Google identity', + ) + } + const matches = + (!!accountId && resolved.accountId === accountId) || + (!!email && resolvedEmail === email) + if (matches) matching.push(candidate) + else if (!resolved.accountId && !resolvedEmail) unresolved++ + } + + if (!matching.length) { + if (failOnUnresolved && unresolved > 0) { + throw new AccountIdentityAmbiguityError( + 'Antigravity account identity is ambiguous; token-only accounts could not be verified, so no account was added or merged', + ) + } + return false + } + const survivor = matching.reduce((first, candidate) => + candidate.index < first.index ? candidate : first, + ) + await reconcileAccountIdentityAtPath(this.path, { + refreshToken: survivor.account.refreshToken, + email, + accountId, + }) + // Reconcile the identity enrichment while the old token is still an exact + // bridge. A subsequent OAuth upsert may rotate that token, at which point + // routing and tracker state can follow the newly stable identity. + await this.reload() + return true + } + private async reload(): Promise { const stored = (await loadAccountStorage(this.path)) ?? { version: 4 as const, @@ -209,16 +366,19 @@ export class PiAccountRuntime { } private async patch( - target: { email?: string; refreshToken: string }, + target: { email?: string; accountId?: string; refreshToken: string }, update: (account: AccountMetadataV3) => void, ): Promise { let applied = false await mutateAccountStorage(this.path, (current) => { const email = target.email?.trim().toLowerCase() + const accountId = target.accountId?.trim() const matches = current.accounts.filter((entry) => email ? entry.email?.trim().toLowerCase() === email - : entry.refreshToken === target.refreshToken, + : accountId + ? entry.accountId === accountId + : entry.refreshToken === target.refreshToken, ) if (matches.length === 1 && matches[0]) { update(matches[0]) @@ -309,6 +469,8 @@ export class PiAccountRuntime { refresh: auth.refresh, access: auth.access ?? '', expires: auth.expires ?? 0, + email: account.email ?? credentials.email, + accountId: account.accountId ?? credentials.accountId, } } catch (error) { if ( @@ -332,7 +494,11 @@ export class PiAccountRuntime { if (current) getHealthTracker().recordFailure(current.index) this.manager?.markAccountCoolingDown(account, 60_000, 'auth-failure') const applied = await this.patch( - { ...account.parts, email: account.email }, + { + ...account.parts, + email: account.email, + accountId: account.accountId, + }, (current) => { current.coolingDownUntil = Math.max( current.coolingDownUntil ?? 0, @@ -426,7 +592,11 @@ export class PiAccountRuntime { manager.updateFromAuth(account, auth) const token = account.parts.refreshToken attempted.add(attemptKey(account)) - const target = { ...account.parts, email: account.email } + const target = { + ...account.parts, + email: account.email, + accountId: account.accountId, + } signal?.throwIfAborted() manager.markAccountUsed(account.index) const located = await this.patch(target, (current) => { From 48460e85d1e7ea529bd89586430ebe3e77b3cca5 Mon Sep 17 00:00:00 2001 From: Brent Duarte Date: Thu, 10 Sep 2026 00:32:59 -0700 Subject: [PATCH 7/7] feat(pi): add live antigravity route diagnostics --- packages/pi/src/runtime.test.ts | 86 ++++++++++++++++++++++++++++++++- packages/pi/src/runtime.ts | 28 +++++++++-- 2 files changed, 108 insertions(+), 6 deletions(-) diff --git a/packages/pi/src/runtime.test.ts b/packages/pi/src/runtime.test.ts index 1fe57642..59448ace 100644 --- a/packages/pi/src/runtime.test.ts +++ b/packages/pi/src/runtime.test.ts @@ -1,4 +1,12 @@ -import { afterEach, beforeEach, describe, expect, it, mock } from 'bun:test' +import { + afterEach, + beforeEach, + describe, + expect, + it, + mock, + spyOn, +} from 'bun:test' import { mkdtemp, readFile, rm, stat, writeFile } from 'node:fs/promises' import { tmpdir } from 'node:os' import { join } from 'node:path' @@ -19,6 +27,8 @@ let directory: string let path: string const runtimes: PiAccountRuntime[] = [] let previousFetch: typeof fetch +let previousConsoleLogEnv: string | undefined +let restoreConsoleError: (() => void) | undefined const quotaFetch = mock() const refresh = mock( async (token: string): Promise => ({ @@ -59,7 +69,15 @@ async function dispatch(result: PiAccountRuntime) { return request.account.index } +function captureConsoleError() { + const errorSpy = spyOn(console, 'error').mockImplementation(() => {}) + restoreConsoleError = () => errorSpy.mockRestore() + return errorSpy +} + beforeEach(async () => { + previousConsoleLogEnv = process.env.ANTIGRAVITY_CORE_CONSOLE_LOG + delete process.env.ANTIGRAVITY_CORE_CONSOLE_LOG directory = await mkdtemp(join(tmpdir(), 'pi-pool-')) path = join(directory, 'accounts.json') initHealthTracker({}) @@ -101,6 +119,11 @@ beforeEach(async () => { }) afterEach(async () => { + restoreConsoleError?.() + restoreConsoleError = undefined + if (previousConsoleLogEnv === undefined) + delete process.env.ANTIGRAVITY_CORE_CONSOLE_LOG + else process.env.ANTIGRAVITY_CORE_CONSOLE_LOG = previousConsoleLogEnv for (const result of runtimes.splice(0)) await result.dispose() globalThis.fetch = previousFetch await rm(directory, { recursive: true, force: true }) @@ -503,6 +526,61 @@ describe('Pi shared account runtime', () => { expect(await dispatch(runtime())).not.toBe(0) }) + it('writes only dispatched account selections to stderr with redacted identities', async () => { + const result = await pool() + await writeStrategy(result.settingsPath, 'round-robin') + process.env.ANTIGRAVITY_CORE_CONSOLE_LOG = '1' + const errorSpy = captureConsoleError() + const indexes: number[] = [] + const send = mock(async (_auth: unknown, account: { index: number }) => { + indexes.push(account.index) + return indexes.length === 1 + ? new Response('', { status: 429 }) + : new Response('ok') + }) + + await result.dispatch(model, send) + + const messages = errorSpy.mock.calls.map(([message]) => String(message)) + expect(messages).toEqual([ + `[agy-route] agy${indexes[0]! + 1} u***@example.com strategy=round-robin group=gemini`, + `[agy-failover] agy${indexes[1]! + 1} u***@example.com`, + ]) + expect(messages.join('\n')).not.toContain('user1@example.com') + expect(messages.join('\n')).not.toContain('google-user') + expect(messages.join('\n')).not.toContain('secret-') + }) + + it('does not log a candidate that fails before request dispatch', async () => { + await pool() + process.env.ANTIGRAVITY_CORE_CONSOLE_LOG = '1' + const errorSpy = captureConsoleError() + const broken = new PiAccountRuntime({ + path, + pid: 0, + refreshToken: async (token) => { + if (token === 'secret-refresh-1') throw new Error('invalid grant') + return refresh(token) + }, + }) + runtimes.push(broken) + + await dispatch(broken) + + expect(errorSpy.mock.calls.map(([message]) => String(message))).toEqual([ + '[agy-route] agy2 u***@example.com strategy=hybrid group=gemini', + ]) + }) + + it('keeps route diagnostics silent when the environment flag is unset', async () => { + const result = await pool() + const errorSpy = captureConsoleError() + + await dispatch(result) + + expect(errorSpy).not.toHaveBeenCalled() + }) + it('terminates when all accounts become rate-limited, without exposing provider secrets', async () => { const result = await pool() const send = mock( @@ -627,6 +705,7 @@ describe('Pi shared account runtime', () => { it('independent OS processes reload the same pool and apply their actual PID offset', async () => { await pool() + process.env.ANTIGRAVITY_CORE_CONSOLE_LOG = '1' const script = ` import { PiAccountRuntime } from ${JSON.stringify(new URL('./runtime.ts', import.meta.url).pathname)}; const runtime = new PiAccountRuntime({ path: process.argv[1], refreshToken: async token => ({ refresh: token, access: 'test-access', expires: Date.now() + 3600000 }) }); @@ -636,6 +715,7 @@ describe('Pi shared account runtime', () => { ` const children = Array.from({ length: 3 }, () => Bun.spawn([process.execPath, '--eval', script, path], { + env: process.env, stdout: 'pipe', stderr: 'pipe', }), @@ -647,7 +727,9 @@ describe('Pi shared account runtime', () => { new Response(child.stdout).text(), new Response(child.stderr).text(), ]) - expect(errors).toBe('') + expect(errors).toMatch( + /^\[agy-route\] agy\d+ u\*\*\*@example\.com strategy=hybrid group=gemini\n$/, + ) expect(code).toBe(0) return JSON.parse(output) as { pid: number; index: number } }), diff --git a/packages/pi/src/runtime.ts b/packages/pi/src/runtime.ts index 38d5de44..d3f243af 100644 --- a/packages/pi/src/runtime.ts +++ b/packages/pi/src/runtime.ts @@ -41,6 +41,7 @@ import { readSettings } from './settings.ts' const QUOTA_REFRESH_MS = 30 * 60_000 const QUOTA_TTL_MS = computeSoftQuotaCacheTtlMs('auto', 30) +const ROUTE_DIAGNOSTIC_ENV = 'ANTIGRAVITY_CORE_CONSOLE_LOG' type LoginResult = Extract class AccountUnavailable extends Error {} @@ -56,6 +57,20 @@ function attemptKey(account: ManagedAccount): string { }) } +function redactedEmail(email?: string): string { + const match = email?.trim().match(/^(.)([^@]*)@(.+)$/) + return match ? `${match[1]}***@${match[3]}` : '(email unavailable)' +} + +function writeRouteDiagnostic(message: string): void { + if (process.env[ROUTE_DIAGNOSTIC_ENV] !== '1') return + try { + console.error(message) + } catch { + // Diagnostics must never change request dispatch behavior. + } +} + export interface PiRuntimeOptions { path?: string pid?: number @@ -541,6 +556,7 @@ export class PiAccountRuntime { // One extra selection tolerates a peer rotating a token between reload and // credential acquisition. Stable attempted identities still bound sends. const budget = (await this.reload()).getTotalAccountCount() + 1 + let dispatched = false let lastError = 'All Antigravity accounts are disabled, cooling down, or over cached quota' const family: AccountModelFamily = @@ -608,6 +624,13 @@ export class PiAccountRuntime { const consumed = getTokenTracker().consume(account.index) let response: Response try { + const accountLabel = `agy${account.index + 1} ${redactedEmail(account.email)}` + writeRouteDiagnostic( + dispatched + ? `[agy-failover] ${accountLabel}` + : `[agy-route] ${accountLabel} strategy=${config.account_selection_strategy} group=${resolveQuotaGroup('gemini', model)}`, + ) + dispatched = true response = await send(auth, account) } catch (error) { const current = this.currentAccount(account) @@ -704,10 +727,7 @@ export class PiAccountRuntime { manager .getAccounts() .map((account) => { - const email = account.email?.match(/^(.)([^@]*)@(.+)$/) - const label = email - ? `${email[1]}***@${email[3]}` - : '(email unavailable)' + const label = redactedEmail(account.email) const until = Math.max( account.coolingDownUntil ?? 0, ...Object.values(account.rateLimitResetTimes).map((n) => n ?? 0),