fix: reuse oauth singleflight in gemini discovery

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(
-    expect(String(undiciFetchMock.mock.calls[4]?.[0] || '')).toContain('/projects/project-refresh-demo/services/');
+      undiciFetchMock.mock.calls.some(([url]) => String(url || '').includes('oauth2.googleapis.com/token')),
-    expect(undiciFetchMock.mock.calls[4]?.[1]).toMatchObject({
+    ).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',

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<typeof schema.accounts.$inferSelect | null> {
-  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<ModelRefreshResult> {
   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<ModelR
         if (!shouldRetryModelDiscoveryWithOauthRefresh(error)) {
           throw error;
         }
-        const refreshedAccount = await refreshOauthAccountForModelDiscovery(discoveryAccount);
+        const refreshed = await refreshOauthAccessTokenSingleflight(discoveryAccount.id);
-        if (!refreshedAccount) {
+        if (!refreshed?.extraConfig) {
           throw error;
         }
-        discoveryAccount = refreshedAccount;
+        discoveryAccount = {
+          ...discoveryAccount,
+          accessToken: refreshed.accessToken,
+          extraConfig: refreshed.extraConfig,
+        };
         await withTimeout(
           () => withAccountProxyOverride(accountProxyUrl,
             () => validateGeminiCliOauthConnection({ account: discoveryAccount })),