From dd48550ffaee3672af2c405ffc983770f0014f8d Mon Sep 17 00:00:00 2001 From: Cita Date: Sat, 21 Mar 2026 18:38:00 +0800 Subject: [PATCH] Fix Gemini OAuth refresh and failure coverage --- .../proxy-core/surfaces/geminiSurface.ts | 148 ++++++++------- src/server/routes/api/oauth.test.ts | 162 ++++++++++++++++ src/server/routes/proxy/gemini.test.ts | 175 ++++++++++++++++++ .../services/modelService.discovery.test.ts | 115 ++++++++++++ src/server/services/modelService.ts | 105 ++++++++++- 5 files changed, 634 insertions(+), 71 deletions(-) diff --git a/src/server/proxy-core/surfaces/geminiSurface.ts b/src/server/proxy-core/surfaces/geminiSurface.ts index a77d6e6..c2e97f4 100644 --- a/src/server/proxy-core/surfaces/geminiSurface.ts +++ b/src/server/proxy-core/surfaces/geminiSurface.ts @@ -9,6 +9,7 @@ import { isModelAllowedByPolicyOrAllowedRoutes } from '../../services/downstream import { tokenRouter } from '../../services/tokenRouter.js'; import { buildOauthProviderHeaders } from '../../services/oauth/service.js'; import { getOauthInfoFromExtraConfig } from '../../services/oauth/oauthAccount.js'; +import { refreshOauthAccessTokenSingleflight } from '../../services/oauth/refreshSingleflight.js'; import { resolveChannelProxyUrl, withSiteRecordProxyRequestInit } from '../../services/siteProxy.js'; import { refreshModelsAndRebuildRoutes } from '../../services/modelService.js'; import { getDownstreamRoutingPolicy } from '../../routes/proxy/downstreamPolicy.js'; @@ -401,7 +402,7 @@ export async function geminiProxyRoute(app: FastifyInstance) { isGeminiCliDownstream ? omitGeminiCliModelField(request.body) : (request.body || {}), actualModel, ); - const oauth = getOauthInfoFromExtraConfig(selected.account.extraConfig); + let oauth = getOauthInfoFromExtraConfig(selected.account.extraConfig); const isGeminiCli = isGeminiCliPlatform(selected.site.platform); const isInternalGemini = isInternalGeminiPlatform(selected.site.platform); const isDirectGeminiFamily = isDirectGeminiFamilyPlatform(selected.site.platform); @@ -438,71 +439,92 @@ export async function geminiProxyRoute(app: FastifyInstance) { ? buildGeminiCliActionPath({ isStreamAction, isCountTokensAction }) : resolveUpstreamPath(apiVersion, actualModelAction); const query = new URLSearchParams(request.query as Record).toString(); - const requestBody = isInternalGemini - ? ( - isCountTokensAction - ? { request: normalizedBody } - : wrapGeminiCliRequest({ - modelName: actualModel, - projectId: oauth?.projectId || '', - request: normalizedBody as Record, - }) - ) - : normalizedBody; - const requestHeaders = isInternalGemini - ? { - 'Content-Type': 'application/json', - ...(isStreamAction ? { Accept: 'text/event-stream' } : {}), - Authorization: `Bearer ${selected.tokenValue}`, - ...buildOauthProviderHeaders({ - extraConfig: typeof selected.account.extraConfig === 'string' ? selected.account.extraConfig : null, - downstreamHeaders: request.headers as Record, - }), - } - : { - 'Content-Type': 'application/json', - }; - const targetUrl = isInternalGemini - ? `${selected.site.url}${upstreamPath}` - : geminiGenerateContentTransformer.resolveActionUrl( - selected.site.url, - apiVersion, - actualModelAction, - selected.tokenValue, - query, - ); const channelProxyUrl = resolveChannelProxyUrl(selected.site, selected.account.extraConfig); - const upstream = isInternalGemini - ? await dispatchRuntimeRequest({ - siteUrl: selected.site.url, - targetUrl, - request: { - endpoint: 'chat', - path: upstreamPath, - headers: requestHeaders, - body: requestBody as Record, - runtime: { - executor: isGeminiCli ? 'gemini-cli' : 'antigravity', - modelName: actualModel, - stream: isStreamAction, - oauthProjectId: oauth?.projectId || null, - action: isCountTokensAction - ? 'countTokens' - : (isStreamAction ? 'streamGenerateContent' : 'generateContent'), + const dispatchSelectedRequest = async () => { + const requestBody = isInternalGemini + ? ( + isCountTokensAction + ? { request: normalizedBody } + : wrapGeminiCliRequest({ + modelName: actualModel, + projectId: oauth?.projectId || '', + request: normalizedBody as Record, + }) + ) + : normalizedBody; + const requestHeaders = isInternalGemini + ? { + 'Content-Type': 'application/json', + ...(isStreamAction ? { Accept: 'text/event-stream' } : {}), + Authorization: `Bearer ${selected.tokenValue}`, + ...buildOauthProviderHeaders({ + extraConfig: typeof selected.account.extraConfig === 'string' ? selected.account.extraConfig : null, + downstreamHeaders: request.headers as Record, + }), + } + : { + 'Content-Type': 'application/json', + }; + const targetUrl = isInternalGemini + ? `${selected.site.url}${upstreamPath}` + : geminiGenerateContentTransformer.resolveActionUrl( + selected.site.url, + apiVersion, + actualModelAction, + selected.tokenValue, + query, + ); + + return isInternalGemini + ? dispatchRuntimeRequest({ + siteUrl: selected.site.url, + targetUrl, + request: { + endpoint: 'chat', + path: upstreamPath, + headers: requestHeaders, + body: requestBody as Record, + runtime: { + executor: isGeminiCli ? 'gemini-cli' : 'antigravity', + modelName: actualModel, + stream: isStreamAction, + oauthProjectId: oauth?.projectId || null, + action: isCountTokensAction + ? 'countTokens' + : (isStreamAction ? 'streamGenerateContent' : 'generateContent'), + }, }, - }, - buildInit: async (_requestUrl, requestForFetch) => withSiteRecordProxyRequestInit(selected.site, { + buildInit: async (_requestUrl, requestForFetch) => withSiteRecordProxyRequestInit(selected.site, { + method: 'POST', + headers: requestForFetch.headers, + body: JSON.stringify(requestForFetch.body), + }, channelProxyUrl), + }) + : fetch(targetUrl, { method: 'POST', - headers: requestForFetch.headers, - body: JSON.stringify(requestForFetch.body), - }, channelProxyUrl), - }) - : await fetch(targetUrl, { - method: 'POST', - headers: requestHeaders, - body: JSON.stringify(requestBody), - }); - const contentType = upstream.headers.get('content-type') || 'application/json'; + headers: requestHeaders, + body: JSON.stringify(requestBody), + }); + }; + + let upstream = await dispatchSelectedRequest(); + let contentType = upstream.headers.get('content-type') || 'application/json'; + if (upstream.status === 401 && oauth) { + try { + const refreshed = await refreshOauthAccessTokenSingleflight(selected.account.id); + selected.tokenValue = refreshed.accessToken; + selected.account = { + ...selected.account, + accessToken: refreshed.accessToken, + extraConfig: refreshed.extraConfig ?? selected.account.extraConfig, + }; + oauth = getOauthInfoFromExtraConfig(selected.account.extraConfig); + upstream = await dispatchSelectedRequest(); + contentType = upstream.headers.get('content-type') || 'application/json'; + } catch { + // Preserve the original 401 response when refresh fails. + } + } if (!upstream.ok) { lastStatus = upstream.status; lastContentType = contentType; diff --git a/src/server/routes/api/oauth.test.ts b/src/server/routes/api/oauth.test.ts index d00d588..2506b93 100644 --- a/src/server/routes/api/oauth.test.ts +++ b/src/server/routes/api/oauth.test.ts @@ -789,6 +789,168 @@ describe('oauth routes', { timeout: 15_000 }, () => { expect(String(fetchMock.mock.calls[7]?.[0] || '')).toContain('/projects/gen-lang-client-0123456789/services/cloudaicompanion.googleapis.com'); }); + it('surfaces Cloud AI API enable failures during Gemini CLI oauth setup and keeps the account rolled back', async () => { + fetchMock + .mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => ({ + access_token: 'gemini-access-token', + refresh_token: 'gemini-refresh-token', + expires_in: 3600, + token_type: 'Bearer', + scope: 'cloud-platform', + }), + text: async () => JSON.stringify({ ok: true }), + }) + .mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => ({ + projects: [ + { projectId: 'project-enable-failure' }, + ], + }), + text: async () => JSON.stringify({ ok: true }), + }) + .mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => ({ + cloudaicompanionProject: { + id: 'project-enable-failure', + }, + allowedTiers: [ + { id: 'legacy-tier', isDefault: true }, + ], + }), + text: async () => JSON.stringify({ ok: true }), + }) + .mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => ({ + done: true, + response: { + cloudaicompanionProject: { + id: 'project-enable-failure', + }, + }, + }), + text: async () => JSON.stringify({ done: true }), + }) + .mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => ({ state: 'DISABLED' }), + text: async () => JSON.stringify({ state: 'DISABLED' }), + }) + .mockResolvedValueOnce({ + ok: false, + status: 403, + json: async () => ({ error: { message: 'Permission denied for project-enable-failure' } }), + text: async () => JSON.stringify({ error: { message: 'Permission denied for project-enable-failure' } }), + }); + + const startResponse = await app.inject({ + method: 'POST', + url: '/api/oauth/providers/gemini-cli/start', + headers: { + host: 'metapi.example', + 'x-forwarded-proto': 'https', + }, + }); + expect(startResponse.statusCode).toBe(200); + const startBody = startResponse.json() as { state: string }; + + const callbackResponse = await app.inject({ + method: 'POST', + url: `/api/oauth/sessions/${encodeURIComponent(startBody.state)}/manual-callback`, + payload: { + callbackUrl: `http://localhost:8085/oauth2callback?state=${encodeURIComponent(startBody.state)}&code=gemini-oauth-code-enable-failure`, + }, + }); + expect(callbackResponse.statusCode).toBe(500); + expect(callbackResponse.json()).toMatchObject({ + message: expect.stringContaining('project activation required'), + }); + + const sessionResponse = await app.inject({ + method: 'GET', + url: `/api/oauth/sessions/${startBody.state}`, + }); + expect(sessionResponse.statusCode).toBe(200); + expect(sessionResponse.json()).toMatchObject({ + provider: 'gemini-cli', + status: 'error', + error: expect.stringContaining('Permission denied for project-enable-failure'), + }); + + const accounts = await db.select().from(schema.accounts).all(); + expect(accounts).toEqual([]); + }); + + it('fails Gemini CLI oauth setup when the Google account has no available Cloud projects', async () => { + fetchMock + .mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => ({ + access_token: 'gemini-access-token', + refresh_token: 'gemini-refresh-token', + expires_in: 3600, + token_type: 'Bearer', + scope: 'cloud-platform', + }), + text: async () => JSON.stringify({ ok: true }), + }) + .mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => ({ + projects: [], + }), + text: async () => JSON.stringify({ ok: true }), + }); + + const startResponse = await app.inject({ + method: 'POST', + url: '/api/oauth/providers/gemini-cli/start', + headers: { + host: 'metapi.example', + 'x-forwarded-proto': 'https', + }, + }); + expect(startResponse.statusCode).toBe(200); + const startBody = startResponse.json() as { state: string }; + + const callbackResponse = await app.inject({ + method: 'POST', + url: `/api/oauth/sessions/${encodeURIComponent(startBody.state)}/manual-callback`, + payload: { + callbackUrl: `http://localhost:8085/oauth2callback?state=${encodeURIComponent(startBody.state)}&code=gemini-oauth-code-no-projects`, + }, + }); + expect(callbackResponse.statusCode).toBe(500); + expect(callbackResponse.json()).toMatchObject({ + message: 'no Google Cloud projects available for this account', + }); + + const sessionResponse = await app.inject({ + method: 'GET', + url: `/api/oauth/sessions/${startBody.state}`, + }); + expect(sessionResponse.statusCode).toBe(200); + expect(sessionResponse.json()).toMatchObject({ + provider: 'gemini-cli', + status: 'error', + error: 'no Google Cloud projects available for this account', + }); + + const accounts = await db.select().from(schema.accounts).all(); + expect(accounts).toEqual([]); + }); + it('lists oauth connection health metadata and supports deleting the connection', async () => { const site = await db.insert(schema.sites).values({ name: 'ChatGPT Codex OAuth', diff --git a/src/server/routes/proxy/gemini.test.ts b/src/server/routes/proxy/gemini.test.ts index 42931da..07d13f1 100644 --- a/src/server/routes/proxy/gemini.test.ts +++ b/src/server/routes/proxy/gemini.test.ts @@ -7,6 +7,7 @@ import { resetUpstreamEndpointRuntimeState } from './upstreamEndpoint.js'; const fetchMock = vi.fn(); const fetchModelPricingCatalogMock = vi.fn(); const refreshModelsAndRebuildRoutesMock = vi.fn(); +const refreshOauthAccessTokenSingleflightMock = vi.fn(); const selectChannelMock = vi.fn(); const selectNextChannelMock = vi.fn(); const recordSuccessMock = vi.fn(); @@ -43,6 +44,7 @@ function createDbSelectChain() { vi.mock('undici', () => ({ fetch: (...args: unknown[]) => fetchMock(...args), + Response: globalThis.Response, })); vi.mock('../../services/modelPricingService.js', () => ({ @@ -53,6 +55,10 @@ vi.mock('../../services/modelService.js', () => ({ refreshModelsAndRebuildRoutes: (...args: unknown[]) => refreshModelsAndRebuildRoutesMock(...args), })); +vi.mock('../../services/oauth/refreshSingleflight.js', () => ({ + refreshOauthAccessTokenSingleflight: (...args: unknown[]) => refreshOauthAccessTokenSingleflightMock(...args), +})); + vi.mock('../../services/tokenRouter.js', () => ({ tokenRouter: { selectChannel: (...args: unknown[]) => selectChannelMock(...args), @@ -154,6 +160,7 @@ describe('gemini native proxy routes', () => { fetchMock.mockReset(); fetchModelPricingCatalogMock.mockReset(); refreshModelsAndRebuildRoutesMock.mockReset(); + refreshOauthAccessTokenSingleflightMock.mockReset(); selectChannelMock.mockReset(); selectNextChannelMock.mockReset(); recordSuccessMock.mockReset(); @@ -514,6 +521,174 @@ describe('gemini native proxy routes', () => { }); }); + it('refreshes gemini-cli oauth token and retries the same internal request on 401', async () => { + selectChannelMock.mockReturnValue({ + channel: { id: 31, routeId: 22 }, + site: { id: 66, name: 'gemini-cli-site', url: 'https://cloudcode-pa.googleapis.com', platform: 'gemini-cli' }, + account: { + id: 36, + username: 'gemini-cli-user@example.com', + accessToken: 'oauth-access-token', + extraConfig: JSON.stringify({ + credentialMode: 'session', + oauth: { + provider: 'gemini-cli', + email: 'gemini-cli-user@example.com', + projectId: 'project-before-refresh', + refreshToken: 'gemini-refresh-token', + }, + }), + }, + tokenName: 'default', + tokenValue: 'oauth-access-token', + actualModel: 'gemini-2.5-pro', + }); + refreshOauthAccessTokenSingleflightMock.mockResolvedValue({ + accountId: 36, + accessToken: 'refreshed-access-token', + extraConfig: JSON.stringify({ + credentialMode: 'session', + oauth: { + provider: 'gemini-cli', + email: 'gemini-cli-user@example.com', + projectId: 'project-after-refresh', + refreshToken: 'gemini-refresh-token', + }, + }), + }); + fetchMock + .mockResolvedValueOnce(new Response(JSON.stringify({ + error: { message: 'token expired' }, + }), { + status: 401, + headers: { 'content-type': 'application/json' }, + })) + .mockResolvedValueOnce(new Response(JSON.stringify({ + response: { + candidates: [ + { + content: { + parts: [{ text: 'hello after refresh' }], + role: 'model', + }, + finishReason: 'STOP', + }, + ], + }, + }), { + status: 200, + headers: { 'content-type': 'application/json' }, + })); + + const response = await app.inject({ + method: 'POST', + url: '/v1beta/models/gemini-2.5-pro:generateContent', + headers: { + authorization: 'Bearer sk-managed-gemini', + }, + payload: { + contents: [ + { + role: 'user', + parts: [{ text: 'hello' }], + }, + ], + }, + }); + + expect(response.statusCode).toBe(200); + expect(refreshOauthAccessTokenSingleflightMock).toHaveBeenCalledWith(36); + expect(fetchMock).toHaveBeenCalledTimes(2); + expect(selectNextChannelMock).not.toHaveBeenCalled(); + + const [, firstRequestInit] = fetchMock.mock.calls[0] as [string, RequestInit]; + const [, secondRequestInit] = fetchMock.mock.calls[1] as [string, RequestInit]; + expect(firstRequestInit.headers).toMatchObject({ + Authorization: 'Bearer oauth-access-token', + }); + expect(secondRequestInit.headers).toMatchObject({ + Authorization: 'Bearer refreshed-access-token', + }); + expect(JSON.parse(String(secondRequestInit.body))).toEqual({ + project: 'project-after-refresh', + model: 'gemini-2.5-pro', + request: { + contents: [ + { + role: 'user', + parts: [{ text: 'hello' }], + }, + ], + }, + }); + expect(response.json()).toEqual({ + responseId: '', + modelVersion: '', + candidates: [ + { + index: 0, + content: { + parts: [{ text: 'hello after refresh' }], + role: 'model', + }, + finishReason: 'STOP', + }, + ], + }); + }); + + it('returns a server error when gemini-cli oauth project metadata is missing at runtime', async () => { + selectChannelMock.mockReturnValue({ + channel: { id: 31, routeId: 22 }, + site: { id: 66, name: 'gemini-cli-site', url: 'https://cloudcode-pa.googleapis.com', platform: 'gemini-cli' }, + account: { + id: 36, + username: 'gemini-cli-user@example.com', + accessToken: 'oauth-access-token', + extraConfig: JSON.stringify({ + credentialMode: 'session', + oauth: { + provider: 'gemini-cli', + email: 'gemini-cli-user@example.com', + }, + }), + }, + tokenName: 'default', + tokenValue: 'oauth-access-token', + actualModel: 'gemini-2.5-pro', + }); + selectNextChannelMock.mockReturnValue(null); + + const response = await app.inject({ + method: 'POST', + url: '/v1beta/models/gemini-2.5-pro:generateContent', + headers: { + authorization: 'Bearer sk-managed-gemini', + }, + payload: { + contents: [ + { + role: 'user', + parts: [{ text: 'hello' }], + }, + ], + }, + }); + + expect(response.statusCode).toBe(500); + expect(fetchMock).not.toHaveBeenCalled(); + expect(recordFailureMock).toHaveBeenCalledWith(31, { + status: 500, + errorText: 'Gemini CLI OAuth project is missing', + }); + expect(JSON.parse(response.body)).toEqual({ + error: { + message: 'Gemini CLI OAuth project is missing', + type: 'server_error', + }, + }); + }); + it('routes Gemini native generateContent requests to openai upstreams and serializes the response back to Gemini shape', async () => { selectChannelMock.mockReturnValue({ channel: { id: 41, routeId: 22 }, diff --git a/src/server/services/modelService.discovery.test.ts b/src/server/services/modelService.discovery.test.ts index 3cef888..502076e 100644 --- a/src/server/services/modelService.discovery.test.ts +++ b/src/server/services/modelService.discovery.test.ts @@ -719,6 +719,121 @@ describe('refreshModelsForAccount credential discovery', () => { }); }); + it('refreshes gemini oauth access token during validation after a 401 and persists the refreshed credentials', async () => { + getApiTokenMock.mockResolvedValue(null); + getModelsMock.mockRejectedValue(new Error('gemini oauth validation should not call adapter.getModels')); + undiciFetchMock + .mockResolvedValueOnce({ + ok: false, + status: 401, + json: async () => ({ error: 'expired' }), + text: async () => 'expired', + }) + .mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => ({ + access_token: 'gemini-access-token-refreshed', + refresh_token: 'gemini-refresh-token-next', + expires_in: 3600, + token_type: 'Bearer', + scope: 'cloud-platform', + }), + text: async () => JSON.stringify({ ok: true }), + }) + .mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => ({ state: 'ENABLED' }), + text: async () => JSON.stringify({ state: 'ENABLED' }), + }) + .mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => ({ email: 'gemini-refreshed-user@example.com' }), + text: async () => JSON.stringify({ email: 'gemini-refreshed-user@example.com' }), + }) + .mockResolvedValueOnce({ + ok: true, + status: 200, + json: async () => ({ state: 'ENABLED' }), + text: async () => JSON.stringify({ state: 'ENABLED' }), + }); + + const site = await db.insert(schema.sites).values({ + name: 'gemini-refresh-site', + url: 'https://gemini.example.com', + platform: 'gemini', + status: 'active', + }).returning().get(); + + const account = await db.insert(schema.accounts).values({ + siteId: site.id, + username: 'gemini-user@example.com', + accessToken: 'gemini-access-token-expired', + apiToken: null, + status: 'active', + oauthProvider: 'gemini-cli', + oauthAccountKey: 'gemini-user@example.com', + oauthProjectId: 'project-refresh-demo', + extraConfig: JSON.stringify({ + credentialMode: 'session', + oauth: { + provider: 'gemini-cli', + email: 'gemini-user@example.com', + accountId: 'gemini-user@example.com', + accountKey: 'gemini-user@example.com', + projectId: 'project-refresh-demo', + refreshToken: 'gemini-refresh-token', + }, + }), + }).returning().get(); + + const result = await refreshModelsForAccount(account.id); + + expect(result).toMatchObject({ + accountId: account.id, + refreshed: true, + status: 'success', + modelCount: expect.any(Number), + discoveredByCredential: true, + }); + expect(undiciFetchMock).toHaveBeenCalledTimes(5); + expect(String(undiciFetchMock.mock.calls[0]?.[0] || '')).toContain('/projects/project-refresh-demo/services/'); + expect(undiciFetchMock.mock.calls[0]?.[1]).toMatchObject({ + method: 'GET', + headers: expect.objectContaining({ + Authorization: 'Bearer gemini-access-token-expired', + }), + }); + expect(String(undiciFetchMock.mock.calls[1]?.[0] || '')).toContain('oauth2.googleapis.com/token'); + expect(String(undiciFetchMock.mock.calls[4]?.[0] || '')).toContain('/projects/project-refresh-demo/services/'); + expect(undiciFetchMock.mock.calls[4]?.[1]).toMatchObject({ + method: 'GET', + headers: expect.objectContaining({ + Authorization: 'Bearer gemini-access-token-refreshed', + }), + }); + + const latest = await db.select().from(schema.accounts) + .where(eq(schema.accounts.id, account.id)) + .get(); + expect(latest).toMatchObject({ + accessToken: 'gemini-access-token-refreshed', + oauthProvider: 'gemini-cli', + oauthAccountKey: 'gemini-refreshed-user@example.com', + oauthProjectId: 'project-refresh-demo', + }); + const parsed = JSON.parse(latest!.extraConfig || '{}'); + expect(parsed.oauth).toMatchObject({ + provider: 'gemini-cli', + email: 'gemini-refreshed-user@example.com', + refreshToken: 'gemini-refresh-token-next', + projectId: 'project-refresh-demo', + modelDiscoveryStatus: 'healthy', + }); + }); + it('discovers antigravity oauth models via fetchAvailableModels fallback using the oauth project id', async () => { getApiTokenMock.mockResolvedValue(null); getModelsMock.mockRejectedValue(new Error('antigravity oauth discovery should not call adapter.getModels')); diff --git a/src/server/services/modelService.ts b/src/server/services/modelService.ts index 9f5dbbf..c657353 100644 --- a/src/server/services/modelService.ts +++ b/src/server/services/modelService.ts @@ -24,6 +24,7 @@ import { withAccountProxyOverride, withExplicitProxyRequestInit, withSiteRecordP import { getCodexOauthInfoFromExtraConfig, isCodexPlatform } from './oauth/codexAccount.js'; import { buildOauthInfo, getOauthInfoFromExtraConfig } from './oauth/oauthAccount.js'; import { CLAUDE_DEFAULT_ANTHROPIC_VERSION } from './oauth/claudeProvider.js'; +import { getOAuthProviderDefinition } from './oauth/providers.js'; import { ANTIGRAVITY_DAILY_UPSTREAM_BASE_URL, ANTIGRAVITY_MODELS_USER_AGENT, @@ -481,6 +482,76 @@ function buildSuccessfulRefreshResult(input: { }; } +function shouldRetryModelDiscoveryWithOauthRefresh(error: unknown): boolean { + const message = ((error as { message?: string })?.message || '').toLowerCase(); + return message.includes('http 401') + || message.includes('unauthorized') + || message.includes('unauthenticated'); +} + +async function refreshOauthAccountForModelDiscovery( + account: typeof schema.accounts.$inferSelect, +): Promise { + const oauth = getOauthInfoFromExtraConfig(account.extraConfig); + if (!oauth?.provider || !oauth.refreshToken) { + return null; + } + + const definition = getOAuthProviderDefinition(oauth.provider); + if (!definition) { + return null; + } + + const refreshed = await definition.refreshAccessToken({ + refreshToken: oauth.refreshToken, + oauth: { + projectId: oauth.projectId, + providerData: oauth.providerData, + }, + }); + const nextOauth = buildOauthInfo(account.extraConfig, { + provider: oauth.provider, + accountId: refreshed.accountId || oauth.accountId, + accountKey: refreshed.accountKey || oauth.accountKey || refreshed.accountId || oauth.accountId, + email: refreshed.email || oauth.email, + planType: refreshed.planType || oauth.planType, + projectId: refreshed.projectId || oauth.projectId, + refreshToken: refreshed.refreshToken || oauth.refreshToken, + tokenExpiresAt: refreshed.tokenExpiresAt || oauth.tokenExpiresAt, + idToken: refreshed.idToken || oauth.idToken, + providerData: { + ...(oauth.providerData || {}), + ...(refreshed.providerData || {}), + }, + }); + const extraConfig = mergeAccountExtraConfig(account.extraConfig, { + credentialMode: 'session', + oauth: nextOauth, + }); + const updatedAt = new Date().toISOString(); + + await db.update(schema.accounts).set({ + accessToken: refreshed.accessToken, + oauthProvider: oauth.provider, + oauthAccountKey: nextOauth.accountKey || nextOauth.accountId || null, + oauthProjectId: nextOauth.projectId || null, + extraConfig, + status: 'active', + updatedAt, + }).where(eq(schema.accounts.id, account.id)).run(); + + return { + ...account, + accessToken: refreshed.accessToken, + oauthProvider: oauth.provider, + oauthAccountKey: nextOauth.accountKey || nextOauth.accountId || null, + oauthProjectId: nextOauth.projectId || null, + extraConfig, + status: 'active', + updatedAt, + }; +} + export async function refreshModelsForAccount(accountId: number): Promise { const row = await db.select().from(schema.accounts) .innerJoin(schema.sites, eq(schema.accounts.siteId, schema.sites.id)) @@ -661,13 +732,31 @@ export async function refreshModelsForAccount(accountId: number): Promise withAccountProxyOverride(accountProxyUrl, - () => validateGeminiCliOauthConnection({ account })), - MODEL_DISCOVERY_TIMEOUT_MS, - `gemini cli oauth validation timeout (${Math.round(MODEL_DISCOVERY_TIMEOUT_MS / 1000)}s)`, - ); + try { + await withTimeout( + () => withAccountProxyOverride(accountProxyUrl, + () => validateGeminiCliOauthConnection({ account: discoveryAccount })), + MODEL_DISCOVERY_TIMEOUT_MS, + `gemini cli oauth validation timeout (${Math.round(MODEL_DISCOVERY_TIMEOUT_MS / 1000)}s)`, + ); + } catch (error) { + if (!shouldRetryModelDiscoveryWithOauthRefresh(error)) { + throw error; + } + const refreshedAccount = await refreshOauthAccountForModelDiscovery(discoveryAccount); + if (!refreshedAccount) { + throw error; + } + discoveryAccount = refreshedAccount; + await withTimeout( + () => withAccountProxyOverride(accountProxyUrl, + () => validateGeminiCliOauthConnection({ account: discoveryAccount })), + MODEL_DISCOVERY_TIMEOUT_MS, + `gemini cli oauth validation timeout (${Math.round(MODEL_DISCOVERY_TIMEOUT_MS / 1000)}s)`, + ); + } await db.insert(schema.modelAvailability).values( GEMINI_CLI_STATIC_MODELS.map((modelName) => ({ accountId, @@ -678,7 +767,7 @@ export async function refreshModelsForAccount(accountId: number): Promise