From b7a1e523e71e8bf86d9cbeca2be69b343ab1b303 Mon Sep 17 00:00:00 2001 From: Cita Date: Sat, 21 Mar 2026 19:29:33 +0800 Subject: [PATCH] fix: reuse oauth singleflight in gemini discovery --- .../services/modelService.discovery.test.ts | 73 +++++++++++------- src/server/services/modelService.ts | 75 ++----------------- 2 files changed, 52 insertions(+), 96 deletions(-) diff --git a/src/server/services/modelService.discovery.test.ts b/src/server/services/modelService.discovery.test.ts index 502076e..de1b500 100644 --- a/src/server/services/modelService.discovery.test.ts +++ b/src/server/services/modelService.discovery.test.ts @@ -8,6 +8,7 @@ const getApiTokenMock = vi.fn(); const getModelsMock = vi.fn(); const undiciFetchMock = vi.fn(); const proxyAgentCtorMock = vi.fn(); +const refreshOauthAccessTokenSingleflightMock = vi.fn(); class MockProxyAgent { readonly proxyUrl: string; @@ -33,6 +34,10 @@ vi.mock('undici', () => ({ Agent: MockAgent, })); +vi.mock('./oauth/refreshSingleflight.js', () => ({ + refreshOauthAccessTokenSingleflight: (...args: unknown[]) => refreshOauthAccessTokenSingleflightMock(...args), +})); + type DbModule = typeof import('../db/index.js'); type ModelServiceModule = typeof import('./modelService.js'); @@ -62,6 +67,7 @@ describe('refreshModelsForAccount credential discovery', () => { getModelsMock.mockReset(); undiciFetchMock.mockReset(); proxyAgentCtorMock.mockReset(); + refreshOauthAccessTokenSingleflightMock.mockReset(); await db.delete(schema.routeChannels).run(); await db.delete(schema.tokenRoutes).run(); @@ -719,9 +725,39 @@ describe('refreshModelsForAccount credential discovery', () => { }); }); - it('refreshes gemini oauth access token during validation after a 401 and persists the refreshed credentials', async () => { + it('refreshes gemini oauth access token during validation through singleflight and reuses the refreshed account state', async () => { getApiTokenMock.mockResolvedValue(null); getModelsMock.mockRejectedValue(new Error('gemini oauth validation should not call adapter.getModels')); + refreshOauthAccessTokenSingleflightMock.mockImplementation(async (accountId: number) => { + const refreshedExtraConfig = JSON.stringify({ + credentialMode: 'session', + oauth: { + provider: 'gemini-cli', + email: 'gemini-refreshed-user@example.com', + accountId: 'gemini-refreshed-user@example.com', + accountKey: 'gemini-refreshed-user@example.com', + projectId: 'project-refresh-demo', + refreshToken: 'gemini-refresh-token-next', + }, + }); + + await db.update(schema.accounts).set({ + accessToken: 'gemini-access-token-refreshed', + oauthProvider: 'gemini-cli', + oauthAccountKey: 'gemini-refreshed-user@example.com', + oauthProjectId: 'project-refresh-demo', + extraConfig: refreshedExtraConfig, + status: 'active', + updatedAt: '2026-03-21T00:00:00.000Z', + }).where(eq(schema.accounts.id, accountId)).run(); + + return { + accountId, + accessToken: 'gemini-access-token-refreshed', + accountKey: 'gemini-refreshed-user@example.com', + extraConfig: refreshedExtraConfig, + }; + }); undiciFetchMock .mockResolvedValueOnce({ ok: false, @@ -729,30 +765,6 @@ describe('refreshModelsForAccount credential discovery', () => { 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, @@ -798,7 +810,8 @@ describe('refreshModelsForAccount credential discovery', () => { modelCount: expect.any(Number), discoveredByCredential: true, }); - expect(undiciFetchMock).toHaveBeenCalledTimes(5); + expect(refreshOauthAccessTokenSingleflightMock).toHaveBeenCalledWith(account.id); + expect(undiciFetchMock).toHaveBeenCalledTimes(2); expect(String(undiciFetchMock.mock.calls[0]?.[0] || '')).toContain('/projects/project-refresh-demo/services/'); expect(undiciFetchMock.mock.calls[0]?.[1]).toMatchObject({ method: 'GET', @@ -806,9 +819,11 @@ describe('refreshModelsForAccount credential discovery', () => { 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({ + expect( + undiciFetchMock.mock.calls.some(([url]) => String(url || '').includes('oauth2.googleapis.com/token')), + ).toBe(false); + expect(String(undiciFetchMock.mock.calls[1]?.[0] || '')).toContain('/projects/project-refresh-demo/services/'); + expect(undiciFetchMock.mock.calls[1]?.[1]).toMatchObject({ method: 'GET', headers: expect.objectContaining({ Authorization: 'Bearer gemini-access-token-refreshed', diff --git a/src/server/services/modelService.ts b/src/server/services/modelService.ts index c657353..b31fcf8 100644 --- a/src/server/services/modelService.ts +++ b/src/server/services/modelService.ts @@ -24,7 +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 { refreshOauthAccessTokenSingleflight } from './oauth/refreshSingleflight.js'; import { ANTIGRAVITY_DAILY_UPSTREAM_BASE_URL, ANTIGRAVITY_MODELS_USER_AGENT, @@ -489,69 +489,6 @@ function shouldRetryModelDiscoveryWithOauthRefresh(error: unknown): boolean { || 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)) @@ -745,11 +682,15 @@ export async function refreshModelsForAccount(accountId: number): Promise withAccountProxyOverride(accountProxyUrl, () => validateGeminiCliOauthConnection({ account: discoveryAccount })),