boos/BoosAPI
2
0
Fork 2
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: address gemini native and route coverage regressions

Browse Source
This commit is contained in:
ciat-777
2026-03-15 21:47:00 +08:00
parent a64a51aa0c
commit 4b42751953
15 changed files with 1006 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/server/routes/proxy/gemini.test.ts
+194 -1
View File
@@ -6,10 +6,19 @@ import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vites
const fetchMock = vi.fn();
const selectChannelMock = vi.fn();
const selectNextChannelMock = vi.fn();
const recordSuccessMock = vi.fn();
const recordFailureMock = vi.fn();
const explainSelectionMock = vi.fn();
const invalidateTokenRouterCacheMock = vi.fn();
const authorizeDownstreamTokenMock = vi.fn();
const consumeManagedKeyRequestMock = vi.fn();
const isModelAllowedByPolicyOrAllowedRoutesMock = vi.fn();
const dbInsertValuesMock = vi.fn((_values?: unknown) => ({
run: () => undefined,
}));
const dbInsertMock = vi.fn((_table?: unknown) => ({
values: (values: unknown) => dbInsertValuesMock(values),
}));
vi.mock('undici', () => ({
fetch: (...args: unknown[]) => fetchMock(...args),
@@ -19,8 +28,9 @@ vi.mock('../../services/tokenRouter.js', () => ({
tokenRouter: {
selectChannel: (...args: unknown[]) => selectChannelMock(...args),
selectNextChannel: (...args: unknown[]) => selectNextChannelMock(...args),
recordSuccess: (...args: unknown[]) => recordSuccessMock(...args),
recordFailure: (...args: unknown[]) => recordFailureMock(...args),
explainSelection: vi.fn(async () => ({ selectedChannelId: 11 })),
explainSelection: (...args: unknown[]) => explainSelectionMock(...args),
},
invalidateTokenRouterCache: (...args: unknown[]) => invalidateTokenRouterCacheMock(...args),
}));
@@ -28,6 +38,16 @@ vi.mock('../../services/tokenRouter.js', () => ({
vi.mock('../../services/downstreamApiKeyService.js', () => ({
authorizeDownstreamToken: (...args: unknown[]) => authorizeDownstreamTokenMock(...args),
consumeManagedKeyRequest: (...args: unknown[]) => consumeManagedKeyRequestMock(...args),
isModelAllowedByPolicyOrAllowedRoutes: (...args: unknown[]) => isModelAllowedByPolicyOrAllowedRoutesMock(...args),
}));
vi.mock('../../db/index.js', () => ({
db: {
insert: (arg: unknown) => dbInsertMock(arg),
},
schema: {
proxyLogs: {},
},
}));
function parseSsePayloads(body: string): Array<Record<string, unknown>> {
@@ -82,9 +102,14 @@ describe('gemini native proxy routes', () => {
fetchMock.mockReset();
selectChannelMock.mockReset();
selectNextChannelMock.mockReset();
recordSuccessMock.mockReset();
recordFailureMock.mockReset();
explainSelectionMock.mockReset();
authorizeDownstreamTokenMock.mockReset();
consumeManagedKeyRequestMock.mockReset();
isModelAllowedByPolicyOrAllowedRoutesMock.mockReset();
dbInsertMock.mockClear();
dbInsertValuesMock.mockClear();
authorizeDownstreamTokenMock.mockResolvedValue({
ok: true,
@@ -102,7 +127,10 @@ describe('gemini native proxy routes', () => {
actualModel: 'gemini-2.5-flash',
});
selectNextChannelMock.mockReturnValue(null);
recordSuccessMock.mockResolvedValue(undefined);
recordFailureMock.mockResolvedValue(undefined);
explainSelectionMock.mockResolvedValue({ selectedChannelId: 11 });
isModelAllowedByPolicyOrAllowedRoutesMock.mockResolvedValue(true);
});
afterAll(async () => {
@@ -241,6 +269,108 @@ describe('gemini native proxy routes', () => {
});
});
it('filters Gemini native model list by downstream managed-key policy', async () => {
authorizeDownstreamTokenMock.mockResolvedValue({
ok: true,
source: 'managed',
token: 'sk-managed-gemini',
key: { id: 91 },
policy: { supportedModels: ['gemini-2.5-flash'], allowedRouteIds: [], siteWeightMultipliers: {} },
});
isModelAllowedByPolicyOrAllowedRoutesMock.mockImplementation(async (modelName: string) => modelName === 'gemini-2.5-flash');
explainSelectionMock.mockImplementation(async (modelName: string) => (
modelName === 'gemini-2.5-flash'
? { selectedChannelId: 11 }
: { selectedChannelId: undefined }
));
fetchMock.mockResolvedValue(new Response(JSON.stringify({
models: [
{ name: 'models/gemini-2.5-flash', displayName: 'Gemini 2.5 Flash' },
{ name: 'models/gemini-2.0-flash', displayName: 'Gemini 2.0 Flash' },
],
}), {
status: 200,
headers: { 'content-type': 'application/json' },
}));
const response = await app.inject({
method: 'GET',
url: '/v1beta/models',
headers: {
authorization: 'Bearer sk-managed-gemini',
},
});
expect(response.statusCode).toBe(200);
expect(response.json()).toEqual({
models: [
{
name: 'models/gemini-2.5-flash',
displayName: 'Gemini 2.5 Flash',
},
],
});
expect(isModelAllowedByPolicyOrAllowedRoutesMock).toHaveBeenCalledWith('gemini-2.5-flash', { supportedModels: ['gemini-2.5-flash'], allowedRouteIds: [], siteWeightMultipliers: {} });
expect(isModelAllowedByPolicyOrAllowedRoutesMock).toHaveBeenCalledWith('gemini-2.0-flash', { supportedModels: ['gemini-2.5-flash'], allowedRouteIds: [], siteWeightMultipliers: {} });
});
it('writes a proxy log row for successful native generateContent requests', async () => {
fetchMock.mockResolvedValue(new Response(JSON.stringify({
candidates: [
{
content: {
parts: [{ text: 'hello from gemini' }],
role: 'model',
},
finishReason: 'STOP',
},
],
usageMetadata: {
promptTokenCount: 10,
candidatesTokenCount: 5,
totalTokenCount: 15,
},
}), {
status: 200,
headers: { 'content-type': 'application/json' },
}));
const response = await app.inject({
method: 'POST',
url: '/v1beta/models/gemini-2.5-flash:generateContent',
headers: {
'x-goog-api-key': 'sk-managed-gemini',
},
payload: {
contents: [
{
role: 'user',
parts: [{ text: 'hello' }],
},
],
},
});
expect(response.statusCode).toBe(200);
expect(recordSuccessMock).toHaveBeenCalledWith(11, expect.any(Number), 0);
expect(dbInsertMock).toHaveBeenCalledTimes(1);
expect(dbInsertValuesMock).toHaveBeenCalledWith(expect.objectContaining({
routeId: 22,
channelId: 11,
accountId: 33,
modelRequested: 'gemini-2.5-flash',
modelActual: 'gemini-2.5-flash',
status: 'success',
httpStatus: 200,
retryCount: 0,
promptTokens: 10,
completionTokens: 5,
totalTokens: 15,
errorMessage: '[downstream:/v1beta/models/gemini-2.5-flash:generateContent] [upstream:/v1beta/models/gemini-2.5-flash:generateContent]',
createdAt: expect.any(String),
}));
});
it('forwards explicit gemini version paths through transformer-owned parsing helpers', async () => {
fetchMock.mockResolvedValue(new Response(JSON.stringify({
candidates: [
@@ -842,4 +972,67 @@ describe('gemini native proxy routes', () => {
expect(recordFailureMock).toHaveBeenCalledWith(11);
expect(response.body).toContain('hello from fallback sse');
});
it('writes failed and successful proxy log rows for Gemini-native stream retries', async () => {
selectNextChannelMock.mockReturnValue({
channel: { id: 12, routeId: 22 },
site: { id: 45, name: 'gemini-site-2', url: 'https://generativelanguage.googleapis.com', platform: 'gemini' },
account: { id: 34, username: 'demo-user-2' },
tokenName: 'fallback',
tokenValue: 'gemini-key-2',
actualModel: 'gemini-2.5-flash',
});
const encoder = new TextEncoder();
const upstreamBody = new ReadableStream<Uint8Array>({
start(controller) {
controller.enqueue(encoder.encode('data: {"responseId":"resp-fallback","candidates":[{"content":{"role":"model","parts":[{"text":"hello from fallback sse"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":11,"candidatesTokenCount":6,"totalTokenCount":17}}\r\n\r\n'));
controller.close();
},
});
fetchMock
.mockResolvedValueOnce(new Response(JSON.stringify({
error: { message: 'upstream unavailable' },
}), {
status: 500,
headers: { 'content-type': 'application/json' },
}))
.mockResolvedValueOnce(new Response(upstreamBody, {
status: 200,
headers: { 'content-type': 'text/event-stream; charset=utf-8' },
}));
const response = await app.inject({
method: 'POST',
url: '/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse',
headers: {
'x-goog-api-key': 'sk-managed-gemini',
},
payload: {
contents: [{ role: 'user', parts: [{ text: 'hello' }] }],
},
});
expect(response.statusCode).toBe(200);
expect(dbInsertMock).toHaveBeenCalledTimes(2);
expect(dbInsertValuesMock).toHaveBeenNthCalledWith(1, expect.objectContaining({
channelId: 11,
status: 'failed',
httpStatus: 500,
retryCount: 0,
errorMessage: '[downstream:/v1beta/models/gemini-2.5-flash:streamGenerateContent] [upstream:/v1beta/models/gemini-2.5-flash:streamGenerateContent] {\"error\":{\"message\":\"upstream unavailable\"}}',
}));
expect(dbInsertValuesMock).toHaveBeenNthCalledWith(2, expect.objectContaining({
channelId: 12,
status: 'success',
httpStatus: 200,
retryCount: 1,
promptTokens: 11,
completionTokens: 6,
totalTokens: 17,
errorMessage: '[downstream:/v1beta/models/gemini-2.5-flash:streamGenerateContent] [upstream:/v1beta/models/gemini-2.5-flash:streamGenerateContent]',
}));
expect(recordSuccessMock).toHaveBeenCalledWith(12, expect.any(Number), 0);
});
});
src/server/routes/proxy/gemini.ts
+206 -7
View File
@@ -1,8 +1,13 @@
import { FastifyInstance, FastifyReply, FastifyRequest } from 'fastify';
import { TextDecoder } from 'node:util';
import { fetch } from 'undici';
import { db, schema } from '../../db/index.js';
import { formatUtcSqlDateTime } from '../../services/localTimeService.js';
import { parseProxyUsage } from '../../services/proxyUsageParser.js';
import { isModelAllowedByPolicyOrAllowedRoutes } from '../../services/downstreamApiKeyService.js';
import { tokenRouter } from '../../services/tokenRouter.js';
import { getDownstreamRoutingPolicy } from './downstreamPolicy.js';
import { composeProxyLogMessage } from './logPathMeta.js';
import {
geminiGenerateContentTransformer,
} from '../../transformers/gemini/generate-content/index.js';
@@ -14,6 +19,11 @@ const GEMINI_MODEL_PROBES = [
'gemini-1.5-flash',
'gemini-pro',
];
const EMPTY_PROXY_USAGE = {
promptTokens: 0,
completionTokens: 0,
totalTokens: 0,
};
async function selectGeminiChannel(request: FastifyRequest) {
const policy = getDownstreamRoutingPolicy(request);
@@ -33,6 +43,106 @@ async function selectNextGeminiProbeChannel(request: FastifyRequest, excludeChan
return null;
}
function resolveDownstreamPath(request: FastifyRequest): string {
const rawUrl = request.raw.url || request.url || '';
const withoutQuery = rawUrl.split('?')[0] || rawUrl;
return withoutQuery || '/v1beta/models';
}
function resolveUpstreamPath(apiVersion: string, modelActionPath: string): string {
const normalizedVersion = apiVersion.replace(/^\/+/, '');
const normalizedAction = modelActionPath.replace(/^\/+/, '');
return `/${normalizedVersion}/${normalizedAction}`;
}
function hasDownstreamModelRestrictions(policy: { supportedModels?: unknown; allowedRouteIds?: unknown }): boolean {
const supportedModels = Array.isArray(policy.supportedModels) ? policy.supportedModels : [];
const allowedRouteIds = Array.isArray(policy.allowedRouteIds) ? policy.allowedRouteIds : [];
return supportedModels.length > 0 || allowedRouteIds.length > 0;
}
function extractGeminiListedModelName(item: unknown): string {
if (!item || typeof item !== 'object') return '';
const rawName = typeof (item as { name?: unknown }).name === 'string'
? (item as { name: string }).name.trim()
: '';
if (!rawName) return '';
return rawName.startsWith('models/') ? rawName.slice('models/'.length) : rawName;
}
async function filterGeminiListedModelsForPolicy(
payload: unknown,
request: FastifyRequest,
): Promise<unknown> {
if (!payload || typeof payload !== 'object' || !Array.isArray((payload as { models?: unknown[] }).models)) {
return payload;
}
const policy = getDownstreamRoutingPolicy(request);
if (!hasDownstreamModelRestrictions(policy)) {
return payload;
}
const filteredModels: unknown[] = [];
for (const item of (payload as { models: unknown[] }).models) {
const modelName = extractGeminiListedModelName(item);
if (!modelName) continue;
if (!await isModelAllowedByPolicyOrAllowedRoutes(modelName, policy)) continue;
const decision = await tokenRouter.explainSelection?.(modelName, [], policy);
if (decision && typeof decision.selectedChannelId !== 'number') continue;
filteredModels.push(item);
}
return {
...(payload as Record<string, unknown>),
models: filteredModels,
};
}
async function logProxy(
selected: any,
modelRequested: string,
status: string,
httpStatus: number,
latencyMs: number,
errorMessage: string | null,
retryCount: number,
downstreamPath: string,
upstreamPath: string | null,
promptTokens = 0,
completionTokens = 0,
totalTokens = 0,
) {
try {
const createdAt = formatUtcSqlDateTime(new Date());
const normalizedErrorMessage = composeProxyLogMessage({
downstreamPath,
upstreamPath,
errorMessage,
});
await db.insert(schema.proxyLogs).values({
routeId: selected.channel.routeId,
channelId: selected.channel.id,
accountId: selected.account.id,
modelRequested,
modelActual: selected.actualModel || modelRequested,
status,
httpStatus,
latencyMs,
promptTokens,
completionTokens,
totalTokens,
estimatedCost: 0,
billingDetails: null,
errorMessage: normalizedErrorMessage,
retryCount,
createdAt,
}).run();
} catch (error) {
console.warn('[proxy/gemini] failed to write proxy log', error);
}
}
export async function geminiProxyRoute(app: FastifyInstance) {
const listModels = async (request: FastifyRequest, reply: FastifyReply) => {
const apiVersion = geminiGenerateContentTransformer.resolveProxyApiVersion(
@@ -72,7 +182,9 @@ export async function geminiProxyRoute(app: FastifyInstance) {
}
try {
return reply.code(upstream.status).send(JSON.parse(text));
const parsed = JSON.parse(text);
const filtered = await filterGeminiListedModelsForPolicy(parsed, request);
return reply.code(upstream.status).send(filtered);
} catch {
return reply.code(upstream.status).type(upstream.headers.get('content-type') || 'application/json').send(text);
}
@@ -107,6 +219,7 @@ export async function geminiProxyRoute(app: FastifyInstance) {
}
const policy = getDownstreamRoutingPolicy(request);
const downstreamPath = resolveDownstreamPath(request);
const excludeChannelIds: number[] = [];
let retryCount = 0;
let lastStatus = 503;
@@ -132,7 +245,9 @@ export async function geminiProxyRoute(app: FastifyInstance) {
/^models\/[^:]+/,
`models/${selected.actualModel || requestedModel}`,
);
const upstreamPath = resolveUpstreamPath(apiVersion, actualModelAction);
const query = new URLSearchParams(request.query as Record<string, string>).toString();
const startTime = Date.now();
try {
const upstream = await fetch(
geminiGenerateContentTransformer.resolveActionUrl(
@@ -156,6 +271,17 @@ export async function geminiProxyRoute(app: FastifyInstance) {
lastContentType = contentType;
lastText = await upstream.text();
await tokenRouter.recordFailure?.(selected.channel.id);
await logProxy(
selected,
requestedModel,
'failed',
lastStatus,
Date.now() - startTime,
lastText,
retryCount,
downstreamPath,
upstreamPath,
);
if (retryCount < MAX_RETRIES) {
retryCount += 1;
continue;
@@ -174,6 +300,19 @@ export async function geminiProxyRoute(app: FastifyInstance) {
reply.raw.setHeader('Content-Type', contentType || 'text/event-stream');
const reader = upstream.body?.getReader();
if (!reader) {
const latency = Date.now() - startTime;
await tokenRouter.recordSuccess?.(selected.channel.id, latency, 0);
await logProxy(
selected,
requestedModel,
'success',
upstream.status,
latency,
null,
retryCount,
downstreamPath,
upstreamPath,
);
reply.raw.end();
return;
}
@@ -209,21 +348,70 @@ export async function geminiProxyRoute(app: FastifyInstance) {
reader.releaseLock();
reply.raw.end();
}
const parsedUsage = parseProxyUsage(aggregateState);
const latency = Date.now() - startTime;
await tokenRouter.recordSuccess?.(selected.channel.id, latency, 0);
await logProxy(
selected,
requestedModel,
'success',
upstream.status,
latency,
null,
retryCount,
downstreamPath,
upstreamPath,
parsedUsage.promptTokens,
parsedUsage.completionTokens,
parsedUsage.totalTokens,
);
return;
}
const text = await upstream.text();
const aggregateState = geminiGenerateContentTransformer.stream.createAggregateState();
let parsedUsage = EMPTY_PROXY_USAGE;
try {
const parsed = JSON.parse(text);
const aggregateState = geminiGenerateContentTransformer.stream.createAggregateState();
const responsePayload = geminiGenerateContentTransformer.stream.serializeUpstreamJsonPayload(
aggregateState,
parsed,
isStreamAction,
);
parsedUsage = parseProxyUsage(aggregateState);
const latency = Date.now() - startTime;
await tokenRouter.recordSuccess?.(selected.channel.id, latency, 0);
await logProxy(
selected,
requestedModel,
'success',
upstream.status,
latency,
null,
retryCount,
downstreamPath,
upstreamPath,
parsedUsage.promptTokens,
parsedUsage.completionTokens,
parsedUsage.totalTokens,
);
return reply.code(upstream.status).send(
geminiGenerateContentTransformer.stream.serializeUpstreamJsonPayload(
aggregateState,
parsed,
isStreamAction,
),
responsePayload,
);
} catch {
const latency = Date.now() - startTime;
await tokenRouter.recordSuccess?.(selected.channel.id, latency, 0);
await logProxy(
selected,
requestedModel,
'success',
upstream.status,
latency,
null,
retryCount,
downstreamPath,
upstreamPath,
);
return reply.code(upstream.status).type(contentType || 'application/json').send(text);
}
} catch (error) {
@@ -236,6 +424,17 @@ export async function geminiProxyRoute(app: FastifyInstance) {
},
});
await tokenRouter.recordFailure?.(selected.channel.id);
await logProxy(
selected,
requestedModel,
'failed',
0,
Date.now() - startTime,
error instanceof Error ? error.message : 'Gemini upstream request failed',
retryCount,
downstreamPath,
upstreamPath,
);
if (retryCount < MAX_RETRIES) {
retryCount += 1;
continue;