Nvidia/BossAI
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[codex] align antigravity special-model non-stream path with CPA (#444)

Browse Source 
* align antigravity special-model non-stream path with cpa

* fix antigravity tail aggregation review feedback

* fix antigravity fallback aggregation review feedback
This commit is contained in:
Cita
2026-04-06 15:03:55 +08:00
committed by GitHub
parent 448f488f95
commit 42f9049aff
8 changed files with 348 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/server/proxy-core/executors/antigravityExecutor.ts
+78 -4
View File
@@ -1,10 +1,12 @@
import { createHash, randomUUID } from 'node:crypto';
import { Response } from 'undici';
import { geminiGenerateContentTransformer } from '../../transformers/gemini/generate-content/index.js';
import type { RuntimeDispatchInput, RuntimeExecutor, RuntimeResponse } from './types.js';
import {
asTrimmedString,
materializeErrorResponse,
performFetch,
readRuntimeResponseText,
withRequestBody,
} from './types.js';
@@ -133,6 +135,68 @@ function antigravityNoCapacityRetryDelay(attempt: number): number {
return Math.min((attempt + 1) * 250, 2000);
}
function isRecord(value: unknown): value is Record<string, unknown> {
return !!value && typeof value === 'object' && !Array.isArray(value);
}
function unwrapAntigravityAggregatePayload(payload: unknown): unknown {
if (!isRecord(payload)) return payload;
return payload.response !== undefined ? payload.response : payload;
}
function antigravityUsesStreamEndpoint(request: RuntimeDispatchInput['request']): boolean {
return request.runtime?.action === 'streamGenerateContent'
|| request.path.includes(':streamGenerateContent');
}
function antigravityShouldAggregateStreamResponse(request: RuntimeDispatchInput['request']): boolean {
return antigravityUsesStreamEndpoint(request) && request.runtime?.stream !== true;
}
async function materializeAntigravitySuccessResponse(
response: RuntimeResponse,
aggregateStreamResponse: boolean,
): Promise<RuntimeResponse> {
if (!aggregateStreamResponse) return response;
const rawText = await readRuntimeResponseText(response);
const contentType = response.headers.get('content-type');
const parsed = geminiGenerateContentTransformer.stream.parseGeminiStreamPayload(rawText, contentType);
const aggregateState = geminiGenerateContentTransformer.stream.createAggregateState();
for (const event of parsed.events) {
geminiGenerateContentTransformer.stream.applyAggregate(
aggregateState,
unwrapAntigravityAggregatePayload(event),
);
}
if (parsed.rest.trim()) {
const trailingPayload = geminiGenerateContentTransformer.stream.parseGeminiStreamPayload(
`${parsed.rest}\n\n`,
'text/event-stream',
);
for (const event of trailingPayload.events) {
geminiGenerateContentTransformer.stream.applyAggregate(
aggregateState,
unwrapAntigravityAggregatePayload(event),
);
}
}
const headers = new Headers(response.headers);
headers.set('content-type', 'application/json; charset=utf-8');
headers.delete('content-encoding');
headers.delete('content-length');
return new Response(
JSON.stringify(geminiGenerateContentTransformer.stream.serializeAggregateJsonPayload(aggregateState)),
{
status: response.status,
headers,
},
);
}
async function waitMs(ms: number): Promise<void> {
if (ms <= 0) return;
await new Promise((resolve) => setTimeout(resolve, ms));
@@ -147,6 +211,8 @@ export const antigravityExecutor: RuntimeExecutor = {
input.request.runtime?.action,
);
const baseAttempts = 3;
const useStreamEndpoint = antigravityUsesStreamEndpoint(input.request);
const aggregateStreamResponse = antigravityShouldAggregateStreamResponse(input.request);
let lastResponse: RuntimeResponse | null = null;
attemptLoop:
@@ -156,7 +222,7 @@ export const antigravityExecutor: RuntimeExecutor = {
const minimalHeaders: Record<string, string> = {
Authorization: input.request.headers.Authorization || input.request.headers.authorization || '',
'Content-Type': 'application/json',
Accept: input.request.runtime?.stream ? 'text/event-stream' : 'application/json',
Accept: useStreamEndpoint ? 'text/event-stream' : 'application/json',
'User-Agent': 'antigravity/1.19.6 darwin/arm64',
};
let response: RuntimeResponse;
@@ -175,7 +241,9 @@ export const antigravityExecutor: RuntimeExecutor = {
}
throw error;
}
if (response.ok) return response;
if (response.ok) {
return materializeAntigravitySuccessResponse(response, aggregateStreamResponse);
}
const errorResponse = await materializeErrorResponse(response);
const errorText = await errorResponse.text().catch(() => '');
@@ -202,11 +270,17 @@ export const antigravityExecutor: RuntimeExecutor = {
}
}
return lastResponse || performFetch(input, withRequestBody(input.request, runtimeBody, {
if (lastResponse) return lastResponse;
const fallbackResponse = await performFetch(input, withRequestBody(input.request, runtimeBody, {
Authorization: input.request.headers.Authorization || input.request.headers.authorization || '',
'Content-Type': 'application/json',
Accept: input.request.runtime?.stream ? 'text/event-stream' : 'application/json',
Accept: useStreamEndpoint ? 'text/event-stream' : 'application/json',
'User-Agent': 'antigravity/1.19.6 darwin/arm64',
}));
if (fallbackResponse.ok) {
return materializeAntigravitySuccessResponse(fallbackResponse, aggregateStreamResponse);
}
return fallbackResponse;
},
};
src/server/proxy-core/providers/antigravityProviderProfile.ts
+3 -7
View File
@@ -1,4 +1,5 @@
import type { PreparedProviderRequest, PrepareProviderRequestInput, ProviderAction, ProviderProfile } from './types.js';
import { resolveAntigravityProviderAction } from './antigravityRuntime.js';
const ANTIGRAVITY_RUNTIME_USER_AGENT = 'antigravity/1.19.6 darwin/arm64';
@@ -6,11 +7,6 @@ function asTrimmedString(value: unknown): string {
return typeof value === 'string' ? value.trim() : '';
}
function resolveAction(action: ProviderAction | undefined, stream: boolean): ProviderAction {
if (action) return action;
return stream ? 'streamGenerateContent' : 'generateContent';
}
function resolvePath(action: ProviderAction): string {
if (action === 'countTokens') return '/v1internal:countTokens';
if (action === 'streamGenerateContent') return '/v1internal:streamGenerateContent?alt=sse';
@@ -20,7 +16,7 @@ function resolvePath(action: ProviderAction): string {
export const antigravityProviderProfile: ProviderProfile = {
id: 'antigravity',
prepareRequest(input: PrepareProviderRequestInput): PreparedProviderRequest {
const action = resolveAction(input.action, input.stream);
const action = resolveAntigravityProviderAction(input.action, input.stream, input.modelName);
const projectId = asTrimmedString(input.oauthProjectId);
return {
path: resolvePath(action),
@@ -38,7 +34,7 @@ export const antigravityProviderProfile: ProviderProfile = {
runtime: {
executor: 'antigravity',
modelName: input.modelName,
stream: action === 'streamGenerateContent',
stream: input.stream,
oauthProjectId: projectId,
action,
},
src/server/proxy-core/providers/antigravityRuntime.test.ts
+39
View File
@@ -0,0 +1,39 @@
import { describe, expect, it } from 'vitest';
import {
resolveAntigravityProviderAction,
shouldUseAntigravityStreamAction,
} from './antigravityRuntime.js';
describe('shouldUseAntigravityStreamAction', () => {
it('returns true for Claude-family model names', () => {
expect(shouldUseAntigravityStreamAction('claude')).toBe(true);
expect(shouldUseAntigravityStreamAction('claude-2')).toBe(true);
expect(shouldUseAntigravityStreamAction('claude-instant')).toBe(true);
expect(shouldUseAntigravityStreamAction('CLAUDE-OPUS-4-1')).toBe(true);
});
it('returns true for Gemini special-model names', () => {
expect(shouldUseAntigravityStreamAction('gemini-3-pro')).toBe(true);
expect(shouldUseAntigravityStreamAction('gemini-3-pro-latest')).toBe(true);
expect(shouldUseAntigravityStreamAction('gemini-3.1-flash-image')).toBe(true);
});
it('returns false for models outside the Claude-family heuristic', () => {
expect(shouldUseAntigravityStreamAction('gemini-2.5-pro')).toBe(false);
expect(shouldUseAntigravityStreamAction('gpt-5')).toBe(false);
});
});
describe('resolveAntigravityProviderAction', () => {
it('routes Claude-family non-stream requests through streamGenerateContent', () => {
expect(resolveAntigravityProviderAction('generateContent', false, 'claude-opus-4-1'))
.toBe('streamGenerateContent');
});
it('routes Gemini special-model non-stream requests through streamGenerateContent', () => {
expect(resolveAntigravityProviderAction('generateContent', false, 'gemini-3-pro-latest'))
.toBe('streamGenerateContent');
expect(resolveAntigravityProviderAction('generateContent', false, 'gemini-3.1-flash-image'))
.toBe('streamGenerateContent');
});
});
src/server/proxy-core/providers/antigravityRuntime.ts
+21
View File
@@ -0,0 +1,21 @@
import type { ProviderAction } from './types.js';
export function shouldUseAntigravityStreamAction(modelName: string): boolean {
const normalizedModel = modelName.toLowerCase();
return normalizedModel.includes('claude')
|| normalizedModel.includes('gemini-3-pro')
|| normalizedModel.includes('gemini-3.1-flash-image');
}
export function resolveAntigravityProviderAction(
action: ProviderAction | undefined,
stream: boolean,
modelName: string,
): ProviderAction {
if (action === 'countTokens') return action;
if (action === 'streamGenerateContent') return action;
if (stream || shouldUseAntigravityStreamAction(modelName)) {
return 'streamGenerateContent';
}
return action || 'generateContent';
}
src/server/proxy-core/surfaces/geminiSurface.ts
+18 -9
View File
@@ -30,6 +30,7 @@ import {
} from '../../transformers/gemini/generate-content/index.js';
import { createChatEndpointStrategy } from '../../transformers/shared/chatEndpointStrategy.js';
import { normalizeUpstreamFinalResponse } from '../../transformers/shared/normalized.js';
import { resolveAntigravityProviderAction } from '../providers/antigravityRuntime.js';
import {
createGeminiCliStreamReader,
unwrapGeminiCliPayload,
@@ -94,11 +95,10 @@ function isInternalGeminiPlatform(platform: unknown): boolean {
}
function buildGeminiCliActionPath(input: {
isStreamAction: boolean;
isCountTokensAction: boolean;
action: 'generateContent' | 'streamGenerateContent' | 'countTokens';
}) {
if (input.isCountTokensAction) return '/v1internal:countTokens';
if (input.isStreamAction) return '/v1internal:streamGenerateContent?alt=sse';
if (input.action === 'countTokens') return '/v1internal:countTokens';
if (input.action === 'streamGenerateContent') return '/v1internal:streamGenerateContent?alt=sse';
return '/v1internal:generateContent';
}
@@ -656,8 +656,19 @@ export async function geminiProxyRoute(app: FastifyInstance) {
/^models\/[^:]+/,
`models/${actualModel}`,
);
const internalGeminiAction = isCountTokensAction
? 'countTokens'
: (
isGeminiCli
? (isStreamAction ? 'streamGenerateContent' : 'generateContent')
: resolveAntigravityProviderAction(
isStreamAction ? 'streamGenerateContent' : 'generateContent',
isStreamAction,
actualModel,
)
);
upstreamPath = isInternalGemini
? buildGeminiCliActionPath({ isStreamAction, isCountTokensAction })
? buildGeminiCliActionPath({ action: internalGeminiAction })
: resolveUpstreamPath(apiVersion, actualModelAction);
const query = new URLSearchParams(request.query as Record<string, string>).toString();
const channelProxyUrl = resolveChannelProxyUrl(selected.site, selected.account.extraConfig);
@@ -676,7 +687,7 @@ export async function geminiProxyRoute(app: FastifyInstance) {
const requestHeaders = isInternalGemini
? {
'Content-Type': 'application/json',
...(isStreamAction ? { Accept: 'text/event-stream' } : {}),
...(internalGeminiAction === 'streamGenerateContent' ? { Accept: 'text/event-stream' } : {}),
Authorization: `Bearer ${selected.tokenValue}`,
...buildOauthProviderHeaders({
account: selected.account,
@@ -720,9 +731,7 @@ export async function geminiProxyRoute(app: FastifyInstance) {
modelName: actualModel,
stream: isStreamAction,
oauthProjectId: oauth?.projectId || null,
action: isCountTokensAction
? 'countTokens'
: (isStreamAction ? 'streamGenerateContent' : 'generateContent'),
action: internalGeminiAction,
},
},
buildInit: async (_requestUrl, requestForFetch) => withSiteRecordProxyRequestInit(selected.site, {