From 500ba2578759331d6e48c10f13f97654d8e70462 Mon Sep 17 00:00:00 2001 From: Cita Date: Sat, 21 Mar 2026 18:37:36 +0800 Subject: [PATCH 1/4] fix: bypass multimodal endpoint memory for gemini --- .../proxy-core/surfaces/geminiSurface.ts | 22 ++++-- src/server/routes/proxy/gemini.test.ts | 75 ++++++++++++++++++ .../routes/proxy/upstreamEndpoint.test.ts | 76 +++++++++++++++++++ src/server/routes/proxy/upstreamEndpoint.ts | 47 ++++++++---- 4 files changed, 202 insertions(+), 18 deletions(-) diff --git a/src/server/proxy-core/surfaces/geminiSurface.ts b/src/server/proxy-core/surfaces/geminiSurface.ts index a77d6e6..112ddbf 100644 --- a/src/server/proxy-core/surfaces/geminiSurface.ts +++ b/src/server/proxy-core/surfaces/geminiSurface.ts @@ -367,6 +367,11 @@ export async function geminiProxyRoute(app: FastifyInstance) { const isCountTokensAction = isGeminiCliDownstream ? options?.action === 'countTokens' : modelActionPath.endsWith(':countTokens'); + const rawUrl = request.raw.url || request.url || ''; + const wantsSseEnvelope = ( + isStreamAction + && (isGeminiCliDownstream || /(?:^|[?&])alt=sse(?:&|$)/i.test(rawUrl)) + ); if (!requestedModel) { return reply.code(400).send({ error: { message: 'Gemini model path is required', type: 'invalid_request_error' }, @@ -857,11 +862,18 @@ export async function geminiProxyRoute(app: FastifyInstance) { parsedUsage.completionTokens, parsedUsage.totalTokens, ); - return reply.code(upstream.status).send( - isGeminiCliDownstream - ? { response: geminiResponse } - : geminiResponse, - ); + const downstreamPayload = isGeminiCliDownstream + ? { response: geminiResponse } + : geminiResponse; + if (wantsSseEnvelope) { + // Some compatibility upstreams finish stream requests with a final JSON payload. + // Preserve Gemini streaming UX by wrapping that terminal payload back into one SSE event. + return reply + .code(upstream.status) + .type('text/event-stream; charset=utf-8') + .send(geminiGenerateContentTransformer.stream.serializeSsePayload(downstreamPayload)); + } + return reply.code(upstream.status).send(downstreamPayload); } catch (error) { lastStatus = 502; lastContentType = 'application/json'; diff --git a/src/server/routes/proxy/gemini.test.ts b/src/server/routes/proxy/gemini.test.ts index 42931da..dfb6b61 100644 --- a/src/server/routes/proxy/gemini.test.ts +++ b/src/server/routes/proxy/gemini.test.ts @@ -602,6 +602,81 @@ describe('gemini native proxy routes', () => { }); }); + it('serializes non-streaming generic upstream JSON into Gemini SSE when alt=sse is requested', async () => { + selectChannelMock.mockReturnValue({ + channel: { id: 41, routeId: 22 }, + site: { id: 77, name: 'openai-site', url: 'https://api.openai.com', platform: 'openai' }, + account: { id: 37, username: 'openai-user@example.com' }, + tokenName: 'default', + tokenValue: 'openai-access-token', + actualModel: 'gpt-4.1', + }); + fetchMock.mockResolvedValue(new Response(JSON.stringify({ + id: 'chatcmpl-openai-stream-1', + object: 'chat.completion', + created: 1_742_160_003, + model: 'gpt-4.1', + choices: [ + { + index: 0, + message: { + role: 'assistant', + content: 'hello from openai stream fallback', + }, + finish_reason: 'stop', + }, + ], + usage: { + prompt_tokens: 5, + completion_tokens: 6, + total_tokens: 11, + }, + }), { + status: 200, + headers: { 'content-type': 'application/json' }, + })); + + const response = await app.inject({ + method: 'POST', + url: '/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse', + headers: { + authorization: 'Bearer sk-managed-gemini', + }, + payload: { + contents: [ + { + role: 'user', + parts: [{ text: 'hello' }], + }, + ], + }, + }); + + expect(response.statusCode).toBe(200); + expect(response.headers['content-type']).toContain('text/event-stream'); + expect(parseSsePayloads(response.body)).toEqual([ + { + responseId: 'chatcmpl-openai-stream-1', + modelVersion: 'gpt-4.1', + candidates: [ + { + index: 0, + content: { + role: 'model', + parts: [{ text: 'hello from openai stream fallback' }], + }, + finishReason: 'STOP', + }, + ], + usageMetadata: { + promptTokenCount: 5, + candidatesTokenCount: 6, + totalTokenCount: 11, + }, + }, + ]); + }); + it('exposes GeminiCLI downstream generateContent endpoint and wraps the downstream response payload', async () => { selectChannelMock.mockReturnValue({ channel: { id: 42, routeId: 22 }, diff --git a/src/server/routes/proxy/upstreamEndpoint.test.ts b/src/server/routes/proxy/upstreamEndpoint.test.ts index 7ca1a04..007cf79 100644 --- a/src/server/routes/proxy/upstreamEndpoint.test.ts +++ b/src/server/routes/proxy/upstreamEndpoint.test.ts @@ -196,6 +196,82 @@ describe('resolveUpstreamEndpointCandidates', () => { expect(order).toEqual(['responses', 'messages', 'chat']); }); + it('does not apply runtime endpoint memory to image attachments', async () => { + recordUpstreamEndpointSuccess({ + siteId: baseContext.site.id, + endpoint: 'responses', + downstreamFormat: 'openai', + modelName: 'gpt-5.3', + requestCapabilities: { + conversationFileSummary: { + hasImage: true, + hasAudio: false, + hasDocument: false, + hasRemoteDocumentUrl: false, + }, + }, + }); + + const order = await resolveUpstreamEndpointCandidates( + { + ...baseContext, + site: { ...baseContext.site, platform: 'new-api' }, + }, + 'gpt-5.3', + 'openai', + undefined, + { + conversationFileSummary: { + hasImage: true, + hasAudio: false, + hasDocument: false, + hasRemoteDocumentUrl: false, + }, + }, + ); + + expect(order).toEqual(['chat', 'messages', 'responses']); + }); + + it('does not apply runtime endpoint memory to document attachments', async () => { + recordUpstreamEndpointSuccess({ + siteId: baseContext.site.id, + endpoint: 'messages', + downstreamFormat: 'openai', + modelName: 'gpt-5.3', + requestCapabilities: { + hasNonImageFileInput: true, + conversationFileSummary: { + hasImage: false, + hasAudio: false, + hasDocument: true, + hasRemoteDocumentUrl: false, + }, + }, + }); + + const order = await resolveUpstreamEndpointCandidates( + { + ...baseContext, + site: { ...baseContext.site, platform: 'new-api' }, + }, + 'gpt-5.3', + 'openai', + undefined, + { + hasNonImageFileInput: true, + conversationFileSummary: { + hasImage: false, + hasAudio: false, + hasDocument: true, + hasRemoteDocumentUrl: false, + }, + }, + ); + + expect(order).toEqual(['responses', 'messages', 'chat']); + }); + it('remembers the last successful endpoint per site capability profile', async () => { recordUpstreamEndpointSuccess({ siteId: baseContext.site.id, diff --git a/src/server/routes/proxy/upstreamEndpoint.ts b/src/server/routes/proxy/upstreamEndpoint.ts index e388cdd..c79939d 100644 --- a/src/server/routes/proxy/upstreamEndpoint.ts +++ b/src/server/routes/proxy/upstreamEndpoint.ts @@ -41,6 +41,8 @@ export type EndpointPreference = DownstreamFormat | 'responses'; type EndpointCapabilityProfile = { preferMessagesForClaudeModel: boolean; + hasImageInput: boolean; + hasAudioInput: boolean; hasNonImageFileInput: boolean; hasRemoteDocumentUrl: boolean; wantsNativeResponsesReasoning: boolean; @@ -711,22 +713,35 @@ function buildEndpointCapabilityProfile(input?: { wantsNativeResponsesReasoning?: boolean; }; }): EndpointCapabilityProfile { + const conversationFileSummary = input?.requestCapabilities?.conversationFileSummary; return { preferMessagesForClaudeModel: ( isClaudeFamilyModel(asTrimmedString(input?.modelName)) || isClaudeFamilyModel(asTrimmedString(input?.requestedModelHint)) ), + hasImageInput: conversationFileSummary?.hasImage === true, + hasAudioInput: conversationFileSummary?.hasAudio === true, hasNonImageFileInput: ( - input?.requestCapabilities?.conversationFileSummary?.hasDocument === true + conversationFileSummary?.hasDocument === true || input?.requestCapabilities?.hasNonImageFileInput === true ), hasRemoteDocumentUrl: ( - input?.requestCapabilities?.conversationFileSummary?.hasRemoteDocumentUrl === true + conversationFileSummary?.hasRemoteDocumentUrl === true ), wantsNativeResponsesReasoning: input?.requestCapabilities?.wantsNativeResponsesReasoning === true, }; } +function shouldUseEndpointRuntimeMemory(capabilityProfile: EndpointCapabilityProfile): boolean { + // Attachment-capable requests are not protocol-equivalent across chat/messages/responses. + // A transient 200 on one endpoint should not bias later multimodal requests onto a lossy path. + return ( + !capabilityProfile.hasImageInput + && !capabilityProfile.hasAudioInput + && !capabilityProfile.hasNonImageFileInput + ); +} + function buildEndpointRuntimeStateKey(input: { siteId: number; downstreamFormat: EndpointPreference; @@ -866,17 +881,19 @@ export function recordUpstreamEndpointSuccess(input: { wantsNativeResponsesReasoning?: boolean; }; }): void { + const capabilityProfile = buildEndpointCapabilityProfile({ + modelName: input.modelName, + requestedModelHint: input.requestedModelHint, + requestCapabilities: input.requestCapabilities, + }); + if (!shouldUseEndpointRuntimeMemory(capabilityProfile)) return; if (!shouldRememberSuccessfulEndpoint(input)) return; const nowMs = Date.now(); const key = buildEndpointRuntimeStateKey({ siteId: input.siteId, downstreamFormat: input.downstreamFormat, - capabilityProfile: buildEndpointCapabilityProfile({ - modelName: input.modelName, - requestedModelHint: input.requestedModelHint, - requestCapabilities: input.requestCapabilities, - }), + capabilityProfile, }); const state = getOrCreateEndpointRuntimeState(key, nowMs); state.preferredEndpoint = input.endpoint; @@ -898,17 +915,19 @@ export function recordUpstreamEndpointFailure(input: { wantsNativeResponsesReasoning?: boolean; }; }): void { + const capabilityProfile = buildEndpointCapabilityProfile({ + modelName: input.modelName, + requestedModelHint: input.requestedModelHint, + requestCapabilities: input.requestCapabilities, + }); + if (!shouldUseEndpointRuntimeMemory(capabilityProfile)) return; if (!shouldBlockEndpointByError(input.status, input.errorText)) return; const nowMs = Date.now(); const key = buildEndpointRuntimeStateKey({ siteId: input.siteId, downstreamFormat: input.downstreamFormat, - capabilityProfile: buildEndpointCapabilityProfile({ - modelName: input.modelName, - requestedModelHint: input.requestedModelHint, - requestCapabilities: input.requestCapabilities, - }), + capabilityProfile, }); const state = getOrCreateEndpointRuntimeState(key, nowMs); state.blockedUntilMsByEndpoint[input.endpoint] = nowMs + ENDPOINT_RUNTIME_BLOCK_TTL_MS; @@ -1009,7 +1028,9 @@ export async function resolveUpstreamEndpointCandidates( capabilityProfile, }); const applyRuntimePreference = (candidates: UpstreamEndpoint[]) => ( - applyEndpointRuntimePreference(candidates, runtimeStateKey) + shouldUseEndpointRuntimeMemory(capabilityProfile) + ? applyEndpointRuntimePreference(candidates, runtimeStateKey) + : candidates ); const conversationFileSummary = requestCapabilities?.conversationFileSummary ?? { hasImage: false, From dd48550ffaee3672af2c405ffc983770f0014f8d Mon Sep 17 00:00:00 2001 From: Cita Date: Sat, 21 Mar 2026 18:38:00 +0800 Subject: [PATCH 2/4] 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 Date: Sat, 21 Mar 2026 18:42:28 +0800 Subject: [PATCH 3/4] Fix Claude count tokens routing and model discovery --- src/server/proxy-core/surfaces/chatSurface.ts | 14 +++- .../routes/proxy/chat.count-tokens.test.ts | 40 +++++++++++ src/server/routes/proxy/upstreamEndpoint.ts | 64 ++++++++++++----- .../services/modelService.discovery.test.ts | 72 +++++++++++++++++++ src/server/services/modelService.ts | 57 ++++++++------- 5 files changed, 201 insertions(+), 46 deletions(-) diff --git a/src/server/proxy-core/surfaces/chatSurface.ts b/src/server/proxy-core/surfaces/chatSurface.ts index b1809ef..9986508 100644 --- a/src/server/proxy-core/surfaces/chatSurface.ts +++ b/src/server/proxy-core/surfaces/chatSurface.ts @@ -869,7 +869,17 @@ export async function handleClaudeCountTokensSurfaceRequest( } excludeChannelIds.push(selected.channel.id); - if (String(selected.site.platform || '').trim().toLowerCase() !== 'claude') { + const modelName = selected.actualModel || requestedModel; + const endpointCandidates = await resolveUpstreamEndpointCandidates( + { + site: selected.site, + account: selected.account, + }, + modelName, + 'claude', + requestedModel, + ); + if (!endpointCandidates.includes('messages')) { if (retryCount < MAX_RETRIES) { retryCount += 1; continue; @@ -881,8 +891,6 @@ export async function handleClaudeCountTokensSurfaceRequest( }, }); } - - const modelName = selected.actualModel || requestedModel; const oauth = getOauthInfoFromExtraConfig(selected.account.extraConfig); const startTime = Date.now(); diff --git a/src/server/routes/proxy/chat.count-tokens.test.ts b/src/server/routes/proxy/chat.count-tokens.test.ts index 831485b..9a0754a 100644 --- a/src/server/routes/proxy/chat.count-tokens.test.ts +++ b/src/server/routes/proxy/chat.count-tokens.test.ts @@ -156,4 +156,44 @@ describe('claude count_tokens proxy route', () => { }, ]); }); + + it('supports /v1/messages/count_tokens for openai-platform gateways that expose Claude messages endpoints', async () => { + selectChannelMock.mockReturnValue({ + channel: { id: 12, routeId: 23 }, + site: { name: 'gateway-site', url: 'https://gateway.example.com', platform: 'openai' }, + account: { id: 34, username: 'gateway-user@example.com' }, + tokenName: 'default', + tokenValue: 'sk-gateway', + actualModel: 'claude-sonnet-4-5-20250929', + }); + fetchMock.mockResolvedValue(new Response(JSON.stringify({ + input_tokens: 9, + }), { + status: 200, + headers: { 'content-type': 'application/json' }, + })); + + const response = await app.inject({ + method: 'POST', + url: '/v1/messages/count_tokens', + payload: { + model: 'claude-sonnet-4-5-20250929', + messages: [ + { + role: 'user', + content: [{ type: 'text', text: 'count through a compatible gateway' }], + }, + ], + }, + }); + + expect(response.statusCode).toBe(200); + expect(response.json()).toEqual({ input_tokens: 9 }); + expect(fetchMock).toHaveBeenCalledTimes(1); + + const [targetUrl, options] = fetchMock.mock.calls[0] as [string, any]; + expect(targetUrl).toBe('https://gateway.example.com/v1/messages/count_tokens?beta=true'); + expect(options.headers['x-api-key']).toBe('sk-gateway'); + expect(options.headers['anthropic-version']).toBe('2023-06-01'); + }); }); diff --git a/src/server/routes/proxy/upstreamEndpoint.ts b/src/server/routes/proxy/upstreamEndpoint.ts index e388cdd..94e798a 100644 --- a/src/server/routes/proxy/upstreamEndpoint.ts +++ b/src/server/routes/proxy/upstreamEndpoint.ts @@ -1521,32 +1521,60 @@ export function buildClaudeCountTokensUpstreamRequest(input: { delete sanitizedBody.maxTokens; delete sanitizedBody.stream; const providerProfile = resolveProviderProfile(sitePlatform); - if (providerProfile?.id !== 'claude') { - throw new Error(`missing claude provider profile for platform: ${sitePlatform || 'unknown'}`); + const effectiveClaudeHeaders = { + ...claudeHeaders, + ...(betas.length > 0 ? { 'anthropic-beta': betas.join(',') } : {}), + }; + + if (providerProfile?.id === 'claude') { + const prepared = providerProfile.prepareRequest({ + endpoint: 'messages', + modelName: input.modelName, + stream: false, + tokenValue: input.tokenValue, + oauthProvider: input.oauthProvider, + sitePlatform, + baseHeaders: { + 'Content-Type': 'application/json', + }, + claudeHeaders: effectiveClaudeHeaders, + body: sanitizedBody, + action: 'countTokens', + }); + + return { + path: prepared.path, + headers: prepared.headers, + body: prepared.body, + runtime: { + executor: 'claude', + modelName: input.modelName, + stream: false, + action: 'countTokens', + }, + }; } - const prepared = providerProfile.prepareRequest({ - endpoint: 'messages', - modelName: input.modelName, - stream: false, - tokenValue: input.tokenValue, - oauthProvider: input.oauthProvider, - sitePlatform, + const anthropicVersion = ( + effectiveClaudeHeaders['anthropic-version'] + || '2023-06-01' + ); + const isClaudeOauthUpstream = sitePlatform === 'claude' && input.oauthProvider === 'claude'; + const headers = buildClaudeRuntimeHeaders({ baseHeaders: { 'Content-Type': 'application/json', }, - claudeHeaders: { - ...claudeHeaders, - ...(betas.length > 0 ? { 'anthropic-beta': betas.join(',') } : {}), - }, - body: sanitizedBody, - action: 'countTokens', + claudeHeaders: effectiveClaudeHeaders, + anthropicVersion, + stream: false, + isClaudeOauthUpstream, + tokenValue: input.tokenValue, }); return { - path: prepared.path, - headers: prepared.headers, - body: prepared.body, + path: '/v1/messages/count_tokens?beta=true', + headers, + body: sanitizedBody, runtime: { executor: 'claude', modelName: input.modelName, diff --git a/src/server/services/modelService.discovery.test.ts b/src/server/services/modelService.discovery.test.ts index 3cef888..0efb7a3 100644 --- a/src/server/services/modelService.discovery.test.ts +++ b/src/server/services/modelService.discovery.test.ts @@ -537,6 +537,78 @@ describe('refreshModelsForAccount credential discovery', () => { ]); }); + it('discovers Claude OAuth models from the upstream /v1/models response', async () => { + getApiTokenMock.mockResolvedValue(null); + getModelsMock.mockRejectedValue(new Error('claude oauth discovery should not call adapter.getModels')); + undiciFetchMock.mockResolvedValue({ + ok: true, + status: 200, + json: async () => ({ + data: [ + { id: 'claude-3-7-sonnet-latest' }, + { id: 'claude-opus-4-1-20250805' }, + ], + }), + text: async () => JSON.stringify({ ok: true }), + }); + + const site = await db.insert(schema.sites).values({ + name: 'claude-site', + url: 'https://api.anthropic.com', + platform: 'claude', + status: 'active', + }).returning().get(); + + const account = await db.insert(schema.accounts).values({ + siteId: site.id, + username: 'claude-user@example.com', + accessToken: 'claude-oauth-token', + apiToken: null, + status: 'active', + extraConfig: JSON.stringify({ + credentialMode: 'session', + oauth: { + provider: 'claude', + email: 'claude-user@example.com', + accountKey: 'claude-user@example.com', + }, + }), + }).returning().get(); + + const result = await refreshModelsForAccount(account.id); + + expect(result).toMatchObject({ + accountId: account.id, + refreshed: true, + status: 'success', + errorCode: null, + tokenScanned: 0, + discoveredByCredential: true, + discoveredApiToken: false, + modelCount: 2, + modelsPreview: ['claude-3-7-sonnet-latest', 'claude-opus-4-1-20250805'], + }); + expect(getModelsMock).not.toHaveBeenCalled(); + expect(undiciFetchMock).toHaveBeenCalledTimes(1); + expect(String(undiciFetchMock.mock.calls[0]?.[0] || '')).toBe('https://api.anthropic.com/v1/models'); + expect(undiciFetchMock.mock.calls[0]?.[1]).toMatchObject({ + method: 'GET', + headers: expect.objectContaining({ + Authorization: 'Bearer claude-oauth-token', + Accept: 'application/json', + 'anthropic-version': '2023-06-01', + }), + }); + + const rows = await db.select().from(schema.modelAvailability) + .where(eq(schema.modelAvailability.accountId, account.id)) + .all(); + expect(rows.map((row) => row.modelName).sort()).toEqual([ + 'claude-3-7-sonnet-latest', + 'claude-opus-4-1-20250805', + ]); + }); + it('marks codex oauth account abnormal when upstream cloud discovery fails', async () => { getApiTokenMock.mockResolvedValue(null); getModelsMock.mockRejectedValue(new Error('codex plan discovery should not call adapter.getModels')); diff --git a/src/server/services/modelService.ts b/src/server/services/modelService.ts index 9f5dbbf..fc216d4 100644 --- a/src/server/services/modelService.ts +++ b/src/server/services/modelService.ts @@ -40,18 +40,6 @@ const API_TOKEN_DISCOVERY_TIMEOUT_MS = 8_000; const MODEL_DISCOVERY_TIMEOUT_MS = 12_000; const MODEL_REFRESH_BATCH_SIZE = 3; const CODEX_MODELS_CLIENT_VERSION = '1.0.0'; -const CLAUDE_STATIC_MODELS = [ - 'claude-haiku-4-5-20251001', - 'claude-sonnet-4-5-20250929', - 'claude-sonnet-4-6', - 'claude-opus-4-6', - 'claude-opus-4-5-20251101', - 'claude-opus-4-1-20250805', - 'claude-opus-4-20250514', - 'claude-sonnet-4-20250514', - 'claude-3-7-sonnet-20250219', - 'claude-3-5-haiku-20241022', -]; const GEMINI_CLI_STATIC_MODELS = [ 'gemini-2.5-pro', 'gemini-2.5-flash', @@ -190,6 +178,21 @@ function extractCodexModelIds(payload: unknown): string[] { }); } +function extractClaudeModelIds(payload: unknown): string[] { + if (!payload || typeof payload !== 'object' || Array.isArray(payload)) return []; + const record = payload as { data?: unknown }; + if (!Array.isArray(record.data)) return []; + + return record.data.flatMap((item) => { + if (!item || typeof item !== 'object' || Array.isArray(item)) return []; + const value = item as { id?: unknown; name?: unknown }; + const id = typeof value.id === 'string' + ? value.id.trim() + : (typeof value.name === 'string' ? value.name.trim() : ''); + return id ? [id] : []; + }); +} + function normalizeBaseUrl(baseUrl: string): string { return (baseUrl || '').replace(/\/+$/, ''); } @@ -284,10 +287,10 @@ async function discoverCodexModelsFromCloud(input: { return normalizeModels(extractCodexModelIds(await response.json())); } -async function validateClaudeOauthConnection(input: { +async function discoverClaudeModelsFromCloud(input: { site: typeof schema.sites.$inferSelect; account: typeof schema.accounts.$inferSelect; -}): Promise { +}): Promise { const accessToken = (input.account.accessToken || '').trim(); if (!accessToken) { throw new Error('claude oauth access token missing'); @@ -305,8 +308,9 @@ async function validateClaudeOauthConnection(input: { ); if (!response.ok) { const text = await response.text().catch(() => ''); - throw new Error(`HTTP ${response.status}: ${text || 'claude oauth validation failed'}`); + throw new Error(`HTTP ${response.status}: ${text || 'claude oauth model discovery failed'}`); } + return normalizeModels(extractClaudeModelIds(await response.json())); } async function validateGeminiCliOauthConnection(input: { @@ -595,14 +599,17 @@ export async function refreshModelsForAccount(accountId: number): Promise withAccountProxyOverride(accountProxyUrl, - () => validateClaudeOauthConnection({ site, account })), + () => discoverClaudeModelsFromCloud({ site, account })), MODEL_DISCOVERY_TIMEOUT_MS, - `claude oauth validation timeout (${Math.round(MODEL_DISCOVERY_TIMEOUT_MS / 1000)}s)`, + `claude oauth model discovery timeout (${Math.round(MODEL_DISCOVERY_TIMEOUT_MS / 1000)}s)`, ); + if (claudeModels.length === 0) { + throw new Error('未获取到可用模型'); + } await db.insert(schema.modelAvailability).values( - CLAUDE_STATIC_MODELS.map((modelName) => ({ + claudeModels.map((modelName) => ({ accountId, modelName, available: true, @@ -614,26 +621,26 @@ export async function refreshModelsForAccount(accountId: number): Promise Date: Sat, 21 Mar 2026 19:29:33 +0800 Subject: [PATCH 4/4] 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 })),