diff --git a/src/server/proxy-core/surfaces/chatSurface.ts b/src/server/proxy-core/surfaces/chatSurface.ts index b6a7c07..fdea9b3 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/proxy-core/surfaces/geminiSurface.ts b/src/server/proxy-core/surfaces/geminiSurface.ts index a77d6e6..8e0ff6d 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'; @@ -367,6 +368,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' }, @@ -401,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); @@ -438,71 +444,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; @@ -857,11 +884,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/api/oauth.test.ts b/src/server/routes/api/oauth.test.ts index 85743fb..6358f03 100644 --- a/src/server/routes/api/oauth.test.ts +++ b/src/server/routes/api/oauth.test.ts @@ -1058,6 +1058,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/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/gemini.test.ts b/src/server/routes/proxy/gemini.test.ts index 42931da..8fe2d2d 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 }, @@ -602,6 +777,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 b3f47a9..21888a5 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..6c91a0f 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, @@ -1521,32 +1542,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 42ec938..b54cc08 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(); @@ -614,6 +620,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')); @@ -796,6 +874,130 @@ describe('refreshModelsForAccount credential discovery', () => { }); }); + 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, + status: 401, + json: async () => ({ error: 'expired' }), + text: async () => 'expired', + }) + .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(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', + headers: expect.objectContaining({ + Authorization: 'Bearer gemini-access-token-expired', + }), + }); + 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', + }), + }); + + 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'));