From 6b53df0147989dac42a76ca53039fc7a714d495c Mon Sep 17 00:00:00 2001 From: Tam Nhu Tran Date: Thu, 16 Apr 2026 14:15:03 -0400 Subject: [PATCH] fix(cliproxy): align delegated gemini auth recovery --- src/cliproxy/quota-fetcher-gemini-cli.ts | 140 ++++++++++++++---- .../cliproxy/quota-fetcher-gemini-cli.test.ts | 16 +- 2 files changed, 123 insertions(+), 33 deletions(-) diff --git a/src/cliproxy/quota-fetcher-gemini-cli.ts b/src/cliproxy/quota-fetcher-gemini-cli.ts index bcbd7887..ad88b9c9 100644 --- a/src/cliproxy/quota-fetcher-gemini-cli.ts +++ b/src/cliproxy/quota-fetcher-gemini-cli.ts @@ -14,6 +14,7 @@ import { buildGeminiCliBucketsFromParsedBuckets, type GeminiCliParsedBucket, } from './gemini-cli-quota-normalizer'; +import { mapExternalProviderName } from './provider-capabilities'; import { buildManagementHeaders, buildProxyUrl, getProxyTarget } from './proxy-target-resolver'; import type { GeminiCliQuotaResult, GeminiCliBucket } from './quota-types'; import { @@ -33,6 +34,7 @@ const GEMINI_CLI_ERROR_DETAIL_MAX_LENGTH = 320; const GEMINI_CLI_ERROR_DETAIL_TRUNCATION_SUFFIX = '...[truncated]'; const GEMINI_CLI_G1_CREDIT_TYPE = 'GOOGLE_ONE_AI'; const MANAGEMENT_API_TIMEOUT_MS = 5000; +const SECONDARY_REQUEST_TIMEOUT_MS = 2000; /** Auth data extracted from Gemini CLI auth file */ interface GeminiCliAuthData { @@ -114,6 +116,20 @@ interface ManagedResponse { viaManagement: boolean; } +interface ManagedGeminiAuthContext { + authIndexLookupPromise?: Promise; +} + +interface ManagedGeminiAuthLookupResult { + authIndex: string | number | null; + unavailable: boolean; +} + +interface ManagedGeminiRequestResult { + response: ManagedResponse | null; + unavailable: boolean; +} + function getRemainingTimeoutMs(deadlineMs: number): number { return Math.max(1, deadlineMs - Date.now()); } @@ -214,8 +230,8 @@ async function readManagedResponse( } function isGeminiAuthFileForAccount(file: ManagementAuthFile, accountId: string): boolean { - const provider = normalizeStringValue(file.provider ?? file.type); - if (provider !== 'gemini') { + const rawProvider = normalizeStringValue(file.provider ?? file.type); + if (!rawProvider || mapExternalProviderName(rawProvider) !== 'gemini') { return false; } @@ -242,7 +258,7 @@ function isGeminiAuthFileForAccount(file: ManagementAuthFile, accountId: string) async function findManagedGeminiAuthIndex( accountId: string, timeoutMs: number -): Promise { +): Promise { const target = getProxyTarget(); const controller = new AbortController(); const timeoutId = setTimeout(() => controller.abort(), timeoutMs); @@ -255,15 +271,35 @@ async function findManagedGeminiAuthIndex( clearTimeout(timeoutId); if (!response.ok) { - return null; + return { authIndex: null, unavailable: true }; } const data = (await response.json()) as { files?: ManagementAuthFile[] }; const match = data.files?.find((file) => isGeminiAuthFileForAccount(file, accountId)); - return match?.auth_index ?? null; + return { authIndex: match?.auth_index ?? null, unavailable: false }; } catch { clearTimeout(timeoutId); - return null; + return { authIndex: null, unavailable: true }; + } +} + +async function getManagedGeminiAuthIndex( + accountId: string, + timeoutMs: number, + context?: ManagedGeminiAuthContext +): Promise { + if (!context) { + return await findManagedGeminiAuthIndex(accountId, timeoutMs); + } + + context.authIndexLookupPromise ??= findManagedGeminiAuthIndex(accountId, timeoutMs); + return await context.authIndexLookupPromise; +} + +class GeminiManagedAuthUnavailableError extends Error { + constructor() { + super('CLIProxy managed Gemini auth is temporarily unavailable'); + this.name = 'GeminiManagedAuthUnavailableError'; } } @@ -271,16 +307,27 @@ async function performManagedGeminiRequest( accountId: string, url: string, body: string, - timeoutMs: number -): Promise { - const authIndex = await findManagedGeminiAuthIndex(accountId, timeoutMs); + timeoutMs: number, + authContext?: ManagedGeminiAuthContext +): Promise { + const deadlineMs = Date.now() + timeoutMs; + const lookupResult = await getManagedGeminiAuthIndex( + accountId, + getRemainingTimeoutMs(deadlineMs), + authContext + ); + if (lookupResult.unavailable) { + return { response: null, unavailable: true }; + } + + const authIndex = lookupResult.authIndex; if (authIndex === null || authIndex === undefined) { - return null; + return { response: null, unavailable: false }; } const target = getProxyTarget(); const controller = new AbortController(); - const timeoutId = setTimeout(() => controller.abort(), timeoutMs); + const timeoutId = setTimeout(() => controller.abort(), getRemainingTimeoutMs(deadlineMs)); try { const response = await fetch(buildProxyUrl(target, '/v0/management/api-call'), { @@ -303,20 +350,23 @@ async function performManagedGeminiRequest( clearTimeout(timeoutId); if (!response.ok) { - return null; + return { response: null, unavailable: true }; } const apiResponse = (await response.json()) as ManagementApiCallResponse; const bodyText = typeof apiResponse.body === 'string' ? apiResponse.body : ''; return { - status: typeof apiResponse.status_code === 'number' ? apiResponse.status_code : 500, - bodyText, - json: safeParseJson(bodyText), - viaManagement: true, + response: { + status: typeof apiResponse.status_code === 'number' ? apiResponse.status_code : 500, + bodyText, + json: safeParseJson(bodyText), + viaManagement: true, + }, + unavailable: false, }; } catch { clearTimeout(timeoutId); - return null; + return { response: null, unavailable: true }; } } @@ -325,10 +375,11 @@ async function performGeminiCliRequest( accessToken: string, url: string, body: string, - preferManagement = false + preferManagement = false, + authContext?: ManagedGeminiAuthContext ): Promise { - const deadlineMs = Date.now() + MANAGEMENT_API_TIMEOUT_MS; let managementAttempted = false; + let managementUnavailable = false; if (preferManagement) { managementAttempted = true; @@ -336,15 +387,20 @@ async function performGeminiCliRequest( accountId, url, body, - getRemainingTimeoutMs(deadlineMs) + MANAGEMENT_API_TIMEOUT_MS, + authContext ); - if (managedResult) { - return managedResult; + managementUnavailable = managedResult.unavailable; + if (managedResult.response) { + return managedResult.response; } } const controller = new AbortController(); - const timeoutId = setTimeout(() => controller.abort(), getRemainingTimeoutMs(deadlineMs)); + const timeoutId = setTimeout( + () => controller.abort(), + managementAttempted ? SECONDARY_REQUEST_TIMEOUT_MS : MANAGEMENT_API_TIMEOUT_MS + ); try { const response = await fetch(url, { @@ -364,6 +420,9 @@ async function performGeminiCliRequest( } if (managementAttempted) { + if (managementUnavailable) { + throw new GeminiManagedAuthUnavailableError(); + } return directResult; } @@ -371,9 +430,16 @@ async function performGeminiCliRequest( accountId, url, body, - getRemainingTimeoutMs(deadlineMs) + SECONDARY_REQUEST_TIMEOUT_MS, + authContext ); - return managedResult ?? directResult; + if (managedResult.response) { + return managedResult.response; + } + if (managedResult.unavailable) { + throw new GeminiManagedAuthUnavailableError(); + } + return directResult; } catch (error) { clearTimeout(timeoutId); throw error; @@ -524,7 +590,8 @@ async function fetchGeminiCliSupplementary( accountId: string, accessToken: string, projectId: string, - verbose: boolean + verbose: boolean, + authContext?: ManagedGeminiAuthContext ): Promise { const requestBody = JSON.stringify({ cloudaicompanionProject: projectId, @@ -541,7 +608,9 @@ async function fetchGeminiCliSupplementary( accountId, accessToken, GEMINI_CLI_CODE_ASSIST_URL, - requestBody + requestBody, + false, + authContext ); if (response.status !== 200) { @@ -890,11 +959,13 @@ async function fetchWithAuthData( }); } + const authContext: ManagedGeminiAuthContext = {}; const supplementaryPromise = fetchGeminiCliSupplementary( accountId, authData.accessToken, authData.projectId, - verbose + verbose, + authContext ); const requestBody = JSON.stringify({ project: authData.projectId }); @@ -904,7 +975,8 @@ async function fetchWithAuthData( authData.accessToken, GEMINI_CLI_QUOTA_URL, requestBody, - authData.isExpired + authData.isExpired, + authContext ); if (verbose) { @@ -952,6 +1024,16 @@ async function fetchWithAuthData( accountId, }; } catch (err) { + if (err instanceof GeminiManagedAuthUnavailableError) { + return buildGeminiCliFailureResult(accountId, authData.projectId, { + error: 'Gemini delegated auth refresh is temporarily unavailable', + errorCode: 'managed_auth_unavailable', + errorDetail: err.message, + actionHint: 'Retry later. CLIProxy management could not refresh this Gemini account.', + retryable: true, + }); + } + const errorMsg = err instanceof Error && err.name === 'AbortError' ? 'Request timeout' diff --git a/tests/unit/cliproxy/quota-fetcher-gemini-cli.test.ts b/tests/unit/cliproxy/quota-fetcher-gemini-cli.test.ts index 7fe6aa86..63c4672b 100644 --- a/tests/unit/cliproxy/quota-fetcher-gemini-cli.test.ts +++ b/tests/unit/cliproxy/quota-fetcher-gemini-cli.test.ts @@ -484,6 +484,12 @@ describe('Gemini CLI Quota Fetcher', () => { writeActiveGeminiAccount('reauth@example.com'); mockFetch([ + { + url: MANAGEMENT_AUTH_FILES_URL, + response: { + files: [], + }, + }, { url: GEMINI_QUOTA_URL, method: 'POST', @@ -670,7 +676,7 @@ describe('Gemini CLI Quota Fetcher', () => { files: [ { auth_index: 'target-auth-index', - provider: 'gemini', + provider: 'gemini-cli', email: 'target@example.com', name: 'target@example.com-gen-lang-client-target-project.json', }, @@ -781,7 +787,7 @@ describe('Gemini CLI Quota Fetcher', () => { files: [ { auth_index: 'retry-auth-index', - provider: 'gemini', + provider: 'gemini-cli', email: 'retry@example.com', name: 'retry@example.com-gen-lang-client-retry-project.json', }, @@ -825,7 +831,7 @@ describe('Gemini CLI Quota Fetcher', () => { } }); - it('does not retry the management API twice when the preferred managed path already failed', async () => { + it('reports a retryable management failure instead of reauth when delegated auth is unavailable', async () => { writeGeminiToken( { type: 'gemini', @@ -890,7 +896,9 @@ describe('Gemini CLI Quota Fetcher', () => { const result = await fetchGeminiCliQuota('managed-failure@example.com'); expect(result.success).toBe(false); - expect(result.needsReauth).toBe(true); + expect(result.needsReauth).toBeUndefined(); + expect(result.retryable).toBe(true); + expect(result.errorCode).toBe('managed_auth_unavailable'); expect(directQuotaAttempt).toBe(1); expect(managedLookupAttempt).toBe(1); expect(managedRequestAttempt).toBe(0);