Merge pull request #234 from cita-777/codex/gemini-oauth-hardening

Fix Gemini OAuth refresh and failure coverage
2026-03-21 20:13:08 +08:00
parent e09fb43b33 b7a1e523e7
commit 4c14f9c9b0
5 changed files with 590 additions and 71 deletions
@@ -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';
@@ -406,7 +407,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);
@@ -443,71 +444,92 @@ export async function geminiProxyRoute(app: FastifyInstance) {
            ? buildGeminiCliActionPath({ isStreamAction, isCountTokensAction })
            : resolveUpstreamPath(apiVersion, actualModelAction);
          const query = new URLSearchParams(request.query as Record<string, string>).toString();
-          const requestBody = isInternalGemini
-            ? (
-              isCountTokensAction
-                ? { request: normalizedBody }
-                : wrapGeminiCliRequest({
-                  modelName: actualModel,
-                  projectId: oauth?.projectId || '',
-                  request: normalizedBody as Record<string, unknown>,
-                })
-            )
-            : 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<string, unknown>,
-              }),
-            }
-            : {
-              '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<string, unknown>,
-                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<string, unknown>,
+                  })
+              )
+              : 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<string, unknown>,
+                }),
+              }
+              : {
+                '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<string, unknown>,
+                  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;