diff --git a/src/cliproxy/quota-fetcher-gemini-cli.ts b/src/cliproxy/quota-fetcher-gemini-cli.ts index 3fa5987d..bcbd7887 100644 --- a/src/cliproxy/quota-fetcher-gemini-cli.ts +++ b/src/cliproxy/quota-fetcher-gemini-cli.ts @@ -114,6 +114,10 @@ interface ManagedResponse { viaManagement: boolean; } +function getRemainingTimeoutMs(deadlineMs: number): number { + return Math.max(1, deadlineMs - Date.now()); +} + /** * Extract project ID from account field * Input: "user@example.com (cloudaicompanion-abc-123)" @@ -235,10 +239,13 @@ function isGeminiAuthFileForAccount(file: ManagementAuthFile, accountId: string) ); } -async function findManagedGeminiAuthIndex(accountId: string): Promise { +async function findManagedGeminiAuthIndex( + accountId: string, + timeoutMs: number +): Promise { const target = getProxyTarget(); const controller = new AbortController(); - const timeoutId = setTimeout(() => controller.abort(), MANAGEMENT_API_TIMEOUT_MS); + const timeoutId = setTimeout(() => controller.abort(), timeoutMs); try { const response = await fetch(buildProxyUrl(target, '/v0/management/auth-files'), { @@ -263,16 +270,17 @@ async function findManagedGeminiAuthIndex(accountId: string): Promise { - const authIndex = await findManagedGeminiAuthIndex(accountId); + const authIndex = await findManagedGeminiAuthIndex(accountId, timeoutMs); if (authIndex === null || authIndex === undefined) { return null; } const target = getProxyTarget(); const controller = new AbortController(); - const timeoutId = setTimeout(() => controller.abort(), MANAGEMENT_API_TIMEOUT_MS); + const timeoutId = setTimeout(() => controller.abort(), timeoutMs); try { const response = await fetch(buildProxyUrl(target, '/v0/management/api-call'), { @@ -319,15 +327,24 @@ async function performGeminiCliRequest( body: string, preferManagement = false ): Promise { + const deadlineMs = Date.now() + MANAGEMENT_API_TIMEOUT_MS; + let managementAttempted = false; + if (preferManagement) { - const managedResult = await performManagedGeminiRequest(accountId, url, body); + managementAttempted = true; + const managedResult = await performManagedGeminiRequest( + accountId, + url, + body, + getRemainingTimeoutMs(deadlineMs) + ); if (managedResult) { return managedResult; } } const controller = new AbortController(); - const timeoutId = setTimeout(() => controller.abort(), MANAGEMENT_API_TIMEOUT_MS); + const timeoutId = setTimeout(() => controller.abort(), getRemainingTimeoutMs(deadlineMs)); try { const response = await fetch(url, { @@ -346,7 +363,16 @@ async function performGeminiCliRequest( return directResult; } - const managedResult = await performManagedGeminiRequest(accountId, url, body); + if (managementAttempted) { + return directResult; + } + + const managedResult = await performManagedGeminiRequest( + accountId, + url, + body, + getRemainingTimeoutMs(deadlineMs) + ); return managedResult ?? directResult; } catch (error) { clearTimeout(timeoutId); diff --git a/tests/unit/cliproxy/quota-fetcher-gemini-cli.test.ts b/tests/unit/cliproxy/quota-fetcher-gemini-cli.test.ts index f4481a43..7fe6aa86 100644 --- a/tests/unit/cliproxy/quota-fetcher-gemini-cli.test.ts +++ b/tests/unit/cliproxy/quota-fetcher-gemini-cli.test.ts @@ -714,6 +714,7 @@ describe('Gemini CLI Quota Fetcher', () => { expect(managedLookupRequest.url).toBe(MANAGEMENT_AUTH_FILES_URL); expect(managedQuotaRequest.url).toBe(MANAGEMENT_API_CALL_URL); expect(managedQuotaRequest.body).toContain('"auth_index":"target-auth-index"'); + expect(managedQuotaRequest.body).toContain('"Authorization":"Bearer $TOKEN$"'); expect(managedQuotaRequest.body).not.toContain('default-refresh-token'); }); @@ -824,6 +825,80 @@ describe('Gemini CLI Quota Fetcher', () => { } }); + it('does not retry the management API twice when the preferred managed path already failed', async () => { + writeGeminiToken( + { + type: 'gemini', + email: 'managed-failure@example.com', + project_id: 'managed-failure-project', + token: { + access_token: 'expired-token', + refresh_token: 'refresh-token', + expiry: Date.now() - 1000, + }, + }, + 'gemini-managed-failure.json' + ); + + mockFetch([ + { + url: GEMINI_CODE_ASSIST_URL, + method: 'POST', + status: 503, + response: { error: { message: 'supplementary unavailable' } }, + }, + ]); + + const originalFetch = globalThis.fetch; + let directQuotaAttempt = 0; + let managedLookupAttempt = 0; + let managedRequestAttempt = 0; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const url = + typeof input === 'string' ? input : input instanceof URL ? input.href : input.url; + + if (url === MANAGEMENT_AUTH_FILES_URL) { + managedLookupAttempt += 1; + return new Response('lookup unavailable', { status: 503 }); + } + + if (url === MANAGEMENT_API_CALL_URL) { + managedRequestAttempt += 1; + return new Response('should not be called', { status: 500 }); + } + + if (url === GEMINI_QUOTA_URL) { + directQuotaAttempt += 1; + return new Response( + JSON.stringify({ + error: { + message: 'Session expired', + status: 'UNAUTHENTICATED', + }, + }), + { + status: 401, + headers: { 'Content-Type': 'application/json' }, + } + ); + } + + return originalFetch(input, init); + }) as typeof fetch; + + try { + const result = await fetchGeminiCliQuota('managed-failure@example.com'); + + expect(result.success).toBe(false); + expect(result.needsReauth).toBe(true); + expect(directQuotaAttempt).toBe(1); + expect(managedLookupAttempt).toBe(1); + expect(managedRequestAttempt).toBe(0); + } finally { + globalThis.fetch = originalFetch; + } + }); + it('classifies model capacity exhaustion separately from generic rate limits', async () => { writeActiveGeminiAccount('capacity@example.com');