[codex] align antigravity special-model non-stream path with CPA (#444)
* 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:
@@ -1,10 +1,12 @@
|
|||||||
import { createHash, randomUUID } from 'node:crypto';
|
import { createHash, randomUUID } from 'node:crypto';
|
||||||
import { Response } from 'undici';
|
import { Response } from 'undici';
|
||||||
|
import { geminiGenerateContentTransformer } from '../../transformers/gemini/generate-content/index.js';
|
||||||
import type { RuntimeDispatchInput, RuntimeExecutor, RuntimeResponse } from './types.js';
|
import type { RuntimeDispatchInput, RuntimeExecutor, RuntimeResponse } from './types.js';
|
||||||
import {
|
import {
|
||||||
asTrimmedString,
|
asTrimmedString,
|
||||||
materializeErrorResponse,
|
materializeErrorResponse,
|
||||||
performFetch,
|
performFetch,
|
||||||
|
readRuntimeResponseText,
|
||||||
withRequestBody,
|
withRequestBody,
|
||||||
} from './types.js';
|
} from './types.js';
|
||||||
|
|
||||||
@@ -133,6 +135,68 @@ function antigravityNoCapacityRetryDelay(attempt: number): number {
|
|||||||
return Math.min((attempt + 1) * 250, 2000);
|
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> {
|
async function waitMs(ms: number): Promise<void> {
|
||||||
if (ms <= 0) return;
|
if (ms <= 0) return;
|
||||||
await new Promise((resolve) => setTimeout(resolve, ms));
|
await new Promise((resolve) => setTimeout(resolve, ms));
|
||||||
@@ -147,6 +211,8 @@ export const antigravityExecutor: RuntimeExecutor = {
|
|||||||
input.request.runtime?.action,
|
input.request.runtime?.action,
|
||||||
);
|
);
|
||||||
const baseAttempts = 3;
|
const baseAttempts = 3;
|
||||||
|
const useStreamEndpoint = antigravityUsesStreamEndpoint(input.request);
|
||||||
|
const aggregateStreamResponse = antigravityShouldAggregateStreamResponse(input.request);
|
||||||
let lastResponse: RuntimeResponse | null = null;
|
let lastResponse: RuntimeResponse | null = null;
|
||||||
|
|
||||||
attemptLoop:
|
attemptLoop:
|
||||||
@@ -156,7 +222,7 @@ export const antigravityExecutor: RuntimeExecutor = {
|
|||||||
const minimalHeaders: Record<string, string> = {
|
const minimalHeaders: Record<string, string> = {
|
||||||
Authorization: input.request.headers.Authorization || input.request.headers.authorization || '',
|
Authorization: input.request.headers.Authorization || input.request.headers.authorization || '',
|
||||||
'Content-Type': 'application/json',
|
'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',
|
'User-Agent': 'antigravity/1.19.6 darwin/arm64',
|
||||||
};
|
};
|
||||||
let response: RuntimeResponse;
|
let response: RuntimeResponse;
|
||||||
@@ -175,7 +241,9 @@ export const antigravityExecutor: RuntimeExecutor = {
|
|||||||
}
|
}
|
||||||
throw error;
|
throw error;
|
||||||
}
|
}
|
||||||
if (response.ok) return response;
|
if (response.ok) {
|
||||||
|
return materializeAntigravitySuccessResponse(response, aggregateStreamResponse);
|
||||||
|
}
|
||||||
|
|
||||||
const errorResponse = await materializeErrorResponse(response);
|
const errorResponse = await materializeErrorResponse(response);
|
||||||
const errorText = await errorResponse.text().catch(() => '');
|
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 || '',
|
Authorization: input.request.headers.Authorization || input.request.headers.authorization || '',
|
||||||
'Content-Type': 'application/json',
|
'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',
|
'User-Agent': 'antigravity/1.19.6 darwin/arm64',
|
||||||
}));
|
}));
|
||||||
|
if (fallbackResponse.ok) {
|
||||||
|
return materializeAntigravitySuccessResponse(fallbackResponse, aggregateStreamResponse);
|
||||||
|
}
|
||||||
|
return fallbackResponse;
|
||||||
},
|
},
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -1,4 +1,5 @@
|
|||||||
import type { PreparedProviderRequest, PrepareProviderRequestInput, ProviderAction, ProviderProfile } from './types.js';
|
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';
|
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() : '';
|
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 {
|
function resolvePath(action: ProviderAction): string {
|
||||||
if (action === 'countTokens') return '/v1internal:countTokens';
|
if (action === 'countTokens') return '/v1internal:countTokens';
|
||||||
if (action === 'streamGenerateContent') return '/v1internal:streamGenerateContent?alt=sse';
|
if (action === 'streamGenerateContent') return '/v1internal:streamGenerateContent?alt=sse';
|
||||||
@@ -20,7 +16,7 @@ function resolvePath(action: ProviderAction): string {
|
|||||||
export const antigravityProviderProfile: ProviderProfile = {
|
export const antigravityProviderProfile: ProviderProfile = {
|
||||||
id: 'antigravity',
|
id: 'antigravity',
|
||||||
prepareRequest(input: PrepareProviderRequestInput): PreparedProviderRequest {
|
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);
|
const projectId = asTrimmedString(input.oauthProjectId);
|
||||||
return {
|
return {
|
||||||
path: resolvePath(action),
|
path: resolvePath(action),
|
||||||
@@ -38,7 +34,7 @@ export const antigravityProviderProfile: ProviderProfile = {
|
|||||||
runtime: {
|
runtime: {
|
||||||
executor: 'antigravity',
|
executor: 'antigravity',
|
||||||
modelName: input.modelName,
|
modelName: input.modelName,
|
||||||
stream: action === 'streamGenerateContent',
|
stream: input.stream,
|
||||||
oauthProjectId: projectId,
|
oauthProjectId: projectId,
|
||||||
action,
|
action,
|
||||||
},
|
},
|
||||||
|
|||||||
@@ -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');
|
||||||
|
});
|
||||||
|
});
|
||||||
@@ -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';
|
||||||
|
}
|
||||||
@@ -30,6 +30,7 @@ import {
|
|||||||
} from '../../transformers/gemini/generate-content/index.js';
|
} from '../../transformers/gemini/generate-content/index.js';
|
||||||
import { createChatEndpointStrategy } from '../../transformers/shared/chatEndpointStrategy.js';
|
import { createChatEndpointStrategy } from '../../transformers/shared/chatEndpointStrategy.js';
|
||||||
import { normalizeUpstreamFinalResponse } from '../../transformers/shared/normalized.js';
|
import { normalizeUpstreamFinalResponse } from '../../transformers/shared/normalized.js';
|
||||||
|
import { resolveAntigravityProviderAction } from '../providers/antigravityRuntime.js';
|
||||||
import {
|
import {
|
||||||
createGeminiCliStreamReader,
|
createGeminiCliStreamReader,
|
||||||
unwrapGeminiCliPayload,
|
unwrapGeminiCliPayload,
|
||||||
@@ -94,11 +95,10 @@ function isInternalGeminiPlatform(platform: unknown): boolean {
|
|||||||
}
|
}
|
||||||
|
|
||||||
function buildGeminiCliActionPath(input: {
|
function buildGeminiCliActionPath(input: {
|
||||||
isStreamAction: boolean;
|
action: 'generateContent' | 'streamGenerateContent' | 'countTokens';
|
||||||
isCountTokensAction: boolean;
|
|
||||||
}) {
|
}) {
|
||||||
if (input.isCountTokensAction) return '/v1internal:countTokens';
|
if (input.action === 'countTokens') return '/v1internal:countTokens';
|
||||||
if (input.isStreamAction) return '/v1internal:streamGenerateContent?alt=sse';
|
if (input.action === 'streamGenerateContent') return '/v1internal:streamGenerateContent?alt=sse';
|
||||||
return '/v1internal:generateContent';
|
return '/v1internal:generateContent';
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -656,8 +656,19 @@ export async function geminiProxyRoute(app: FastifyInstance) {
|
|||||||
/^models\/[^:]+/,
|
/^models\/[^:]+/,
|
||||||
`models/${actualModel}`,
|
`models/${actualModel}`,
|
||||||
);
|
);
|
||||||
|
const internalGeminiAction = isCountTokensAction
|
||||||
|
? 'countTokens'
|
||||||
|
: (
|
||||||
|
isGeminiCli
|
||||||
|
? (isStreamAction ? 'streamGenerateContent' : 'generateContent')
|
||||||
|
: resolveAntigravityProviderAction(
|
||||||
|
isStreamAction ? 'streamGenerateContent' : 'generateContent',
|
||||||
|
isStreamAction,
|
||||||
|
actualModel,
|
||||||
|
)
|
||||||
|
);
|
||||||
upstreamPath = isInternalGemini
|
upstreamPath = isInternalGemini
|
||||||
? buildGeminiCliActionPath({ isStreamAction, isCountTokensAction })
|
? buildGeminiCliActionPath({ action: internalGeminiAction })
|
||||||
: resolveUpstreamPath(apiVersion, actualModelAction);
|
: resolveUpstreamPath(apiVersion, actualModelAction);
|
||||||
const query = new URLSearchParams(request.query as Record<string, string>).toString();
|
const query = new URLSearchParams(request.query as Record<string, string>).toString();
|
||||||
const channelProxyUrl = resolveChannelProxyUrl(selected.site, selected.account.extraConfig);
|
const channelProxyUrl = resolveChannelProxyUrl(selected.site, selected.account.extraConfig);
|
||||||
@@ -676,7 +687,7 @@ export async function geminiProxyRoute(app: FastifyInstance) {
|
|||||||
const requestHeaders = isInternalGemini
|
const requestHeaders = isInternalGemini
|
||||||
? {
|
? {
|
||||||
'Content-Type': 'application/json',
|
'Content-Type': 'application/json',
|
||||||
...(isStreamAction ? { Accept: 'text/event-stream' } : {}),
|
...(internalGeminiAction === 'streamGenerateContent' ? { Accept: 'text/event-stream' } : {}),
|
||||||
Authorization: `Bearer ${selected.tokenValue}`,
|
Authorization: `Bearer ${selected.tokenValue}`,
|
||||||
...buildOauthProviderHeaders({
|
...buildOauthProviderHeaders({
|
||||||
account: selected.account,
|
account: selected.account,
|
||||||
@@ -720,9 +731,7 @@ export async function geminiProxyRoute(app: FastifyInstance) {
|
|||||||
modelName: actualModel,
|
modelName: actualModel,
|
||||||
stream: isStreamAction,
|
stream: isStreamAction,
|
||||||
oauthProjectId: oauth?.projectId || null,
|
oauthProjectId: oauth?.projectId || null,
|
||||||
action: isCountTokensAction
|
action: internalGeminiAction,
|
||||||
? 'countTokens'
|
|
||||||
: (isStreamAction ? 'streamGenerateContent' : 'generateContent'),
|
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
buildInit: async (_requestUrl, requestForFetch) => withSiteRecordProxyRequestInit(selected.site, {
|
buildInit: async (_requestUrl, requestForFetch) => withSiteRecordProxyRequestInit(selected.site, {
|
||||||
|
|||||||
@@ -1144,7 +1144,7 @@ describe('gemini native proxy routes', () => {
|
|||||||
expect(targetUrl).toBe('https://api.openai.com/v1/responses');
|
expect(targetUrl).toBe('https://api.openai.com/v1/responses');
|
||||||
});
|
});
|
||||||
|
|
||||||
it('routes Gemini native generateContent requests to antigravity upstreams through the internal content endpoint', async () => {
|
it('routes Gemini native generateContent requests to antigravity special models through the internal stream endpoint and aggregates back to Gemini JSON', async () => {
|
||||||
selectChannelMock.mockReturnValue({
|
selectChannelMock.mockReturnValue({
|
||||||
channel: { id: 43, routeId: 22 },
|
channel: { id: 43, routeId: 22 },
|
||||||
site: { id: 79, name: 'antigravity-site', url: 'https://cloudcode-pa.googleapis.com', platform: 'antigravity' },
|
site: { id: 79, name: 'antigravity-site', url: 'https://cloudcode-pa.googleapis.com', platform: 'antigravity' },
|
||||||
@@ -1163,28 +1163,16 @@ describe('gemini native proxy routes', () => {
|
|||||||
tokenValue: 'antigravity-access-token',
|
tokenValue: 'antigravity-access-token',
|
||||||
actualModel: 'gemini-3-pro-preview',
|
actualModel: 'gemini-3-pro-preview',
|
||||||
});
|
});
|
||||||
fetchMock.mockResolvedValue(new Response(JSON.stringify({
|
fetchMock.mockResolvedValue(new Response([
|
||||||
response: {
|
'data: {"response":{"responseId":"antigravity-response-1","modelVersion":"gemini-3-pro-preview","candidates":[{"content":{"role":"model","parts":[{"text":"hello "}]},"index":0}]}}',
|
||||||
responseId: 'antigravity-response-1',
|
'',
|
||||||
modelVersion: 'gemini-3-pro-preview',
|
'data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"text":"from antigravity"}]},"finishReason":"STOP","index":0}],"usageMetadata":{"promptTokenCount":8,"candidatesTokenCount":4,"totalTokenCount":12}}}',
|
||||||
candidates: [
|
'',
|
||||||
{
|
'data: [DONE]',
|
||||||
content: {
|
'',
|
||||||
role: 'model',
|
].join('\n'), {
|
||||||
parts: [{ text: 'hello from antigravity' }],
|
|
||||||
},
|
|
||||||
finishReason: 'STOP',
|
|
||||||
},
|
|
||||||
],
|
|
||||||
usageMetadata: {
|
|
||||||
promptTokenCount: 8,
|
|
||||||
candidatesTokenCount: 4,
|
|
||||||
totalTokenCount: 12,
|
|
||||||
},
|
|
||||||
},
|
|
||||||
}), {
|
|
||||||
status: 200,
|
status: 200,
|
||||||
headers: { 'content-type': 'application/json' },
|
headers: { 'content-type': 'text/event-stream; charset=utf-8' },
|
||||||
}));
|
}));
|
||||||
|
|
||||||
const response = await app.inject({
|
const response = await app.inject({
|
||||||
@@ -1205,11 +1193,11 @@ describe('gemini native proxy routes', () => {
|
|||||||
|
|
||||||
expect(response.statusCode).toBe(200);
|
expect(response.statusCode).toBe(200);
|
||||||
const [targetUrl, requestInit] = fetchMock.mock.calls[0] as [string, RequestInit];
|
const [targetUrl, requestInit] = fetchMock.mock.calls[0] as [string, RequestInit];
|
||||||
expect(targetUrl).toBe('https://daily-cloudcode-pa.googleapis.com/v1internal:generateContent');
|
expect(targetUrl).toBe('https://daily-cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse');
|
||||||
expect(requestInit.headers).toMatchObject({
|
expect(requestInit.headers).toMatchObject({
|
||||||
Authorization: 'Bearer antigravity-access-token',
|
Authorization: 'Bearer antigravity-access-token',
|
||||||
'Content-Type': 'application/json',
|
'Content-Type': 'application/json',
|
||||||
Accept: 'application/json',
|
Accept: 'text/event-stream',
|
||||||
'User-Agent': 'antigravity/1.19.6 darwin/arm64',
|
'User-Agent': 'antigravity/1.19.6 darwin/arm64',
|
||||||
});
|
});
|
||||||
const upstreamBody = JSON.parse(String(requestInit.body));
|
const upstreamBody = JSON.parse(String(requestInit.body));
|
||||||
|
|||||||
@@ -17,7 +17,7 @@ describe('dispatchRuntimeRequest', () => {
|
|||||||
fetchMock.mockReset();
|
fetchMock.mockReset();
|
||||||
});
|
});
|
||||||
|
|
||||||
it('routes antigravity runtime requests through daily then sandbox base urls and rewrites the payload fingerprint', async () => {
|
it('routes standard antigravity runtime requests through daily then sandbox base urls and rewrites the payload fingerprint', async () => {
|
||||||
const { dispatchRuntimeRequest } = await import('./runtimeExecutor.js');
|
const { dispatchRuntimeRequest } = await import('./runtimeExecutor.js');
|
||||||
const request: BuiltEndpointRequest = {
|
const request: BuiltEndpointRequest = {
|
||||||
endpoint: 'chat',
|
endpoint: 'chat',
|
||||||
@@ -29,7 +29,7 @@ describe('dispatchRuntimeRequest', () => {
|
|||||||
},
|
},
|
||||||
body: {
|
body: {
|
||||||
project: 'project-demo',
|
project: 'project-demo',
|
||||||
model: 'gemini-3-pro-preview',
|
model: 'gemini-2.5-pro',
|
||||||
request: {
|
request: {
|
||||||
contents: [
|
contents: [
|
||||||
{
|
{
|
||||||
@@ -41,7 +41,7 @@ describe('dispatchRuntimeRequest', () => {
|
|||||||
},
|
},
|
||||||
runtime: {
|
runtime: {
|
||||||
executor: 'antigravity',
|
executor: 'antigravity',
|
||||||
modelName: 'gemini-3-pro-preview',
|
modelName: 'gemini-2.5-pro',
|
||||||
stream: false,
|
stream: false,
|
||||||
},
|
},
|
||||||
};
|
};
|
||||||
@@ -86,7 +86,7 @@ describe('dispatchRuntimeRequest', () => {
|
|||||||
const upstreamBody = JSON.parse(String(firstInit.body));
|
const upstreamBody = JSON.parse(String(firstInit.body));
|
||||||
expect(upstreamBody).toMatchObject({
|
expect(upstreamBody).toMatchObject({
|
||||||
project: 'project-demo',
|
project: 'project-demo',
|
||||||
model: 'gemini-3-pro-preview',
|
model: 'gemini-2.5-pro',
|
||||||
userAgent: 'antigravity',
|
userAgent: 'antigravity',
|
||||||
requestType: 'agent',
|
requestType: 'agent',
|
||||||
request: {
|
request: {
|
||||||
@@ -102,6 +102,169 @@ describe('dispatchRuntimeRequest', () => {
|
|||||||
expect(upstreamBody.requestId).toMatch(/^agent-[0-9a-f-]{36}$/i);
|
expect(upstreamBody.requestId).toMatch(/^agent-[0-9a-f-]{36}$/i);
|
||||||
});
|
});
|
||||||
|
|
||||||
|
it('routes antigravity gemini-3-pro non-stream requests through the stream endpoint and aggregates SSE back to JSON', async () => {
|
||||||
|
const { dispatchRuntimeRequest } = await import('./runtimeExecutor.js');
|
||||||
|
const request: BuiltEndpointRequest = {
|
||||||
|
endpoint: 'chat',
|
||||||
|
path: '/v1internal:streamGenerateContent?alt=sse',
|
||||||
|
headers: {
|
||||||
|
Authorization: 'Bearer antigravity-token',
|
||||||
|
'Content-Type': 'application/json',
|
||||||
|
},
|
||||||
|
body: {
|
||||||
|
project: 'project-demo',
|
||||||
|
model: 'gemini-3-pro-low',
|
||||||
|
request: {
|
||||||
|
contents: [
|
||||||
|
{
|
||||||
|
role: 'user',
|
||||||
|
parts: [{ text: 'hello runtime executor' }],
|
||||||
|
},
|
||||||
|
],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
runtime: {
|
||||||
|
executor: 'antigravity',
|
||||||
|
modelName: 'gemini-3-pro-low',
|
||||||
|
stream: false,
|
||||||
|
action: 'streamGenerateContent',
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
fetchMock.mockResolvedValueOnce(new Response([
|
||||||
|
'data: {"response":{"responseId":"antigravity-stream-1","modelVersion":"gemini-3-pro-low","candidates":[{"content":{"role":"model","parts":[{"text":"hello "}]},"index":0}]}}',
|
||||||
|
'',
|
||||||
|
'data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"text":"from antigravity"}]},"finishReason":"STOP","index":0}],"usageMetadata":{"promptTokenCount":5,"candidatesTokenCount":4,"totalTokenCount":9}}}',
|
||||||
|
'',
|
||||||
|
'data: [DONE]',
|
||||||
|
'',
|
||||||
|
].join('\n'), {
|
||||||
|
status: 200,
|
||||||
|
headers: { 'content-type': 'text/event-stream; charset=utf-8' },
|
||||||
|
}) as unknown as Awaited<ReturnType<typeof fetch>>);
|
||||||
|
|
||||||
|
const response = await dispatchRuntimeRequest({
|
||||||
|
siteUrl: 'https://cloudcode-pa.googleapis.com',
|
||||||
|
request,
|
||||||
|
buildInit: (_url, nextRequest) => ({
|
||||||
|
method: 'POST',
|
||||||
|
headers: nextRequest.headers,
|
||||||
|
body: JSON.stringify(nextRequest.body),
|
||||||
|
}),
|
||||||
|
});
|
||||||
|
|
||||||
|
expect(response.ok).toBe(true);
|
||||||
|
expect(fetchMock).toHaveBeenCalledTimes(1);
|
||||||
|
expect(fetchMock.mock.calls[0]?.[0]).toBe('https://daily-cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse');
|
||||||
|
|
||||||
|
const requestInit = fetchMock.mock.calls[0]?.[1] as RequestInit;
|
||||||
|
expect(requestInit.headers).toMatchObject({
|
||||||
|
Authorization: 'Bearer antigravity-token',
|
||||||
|
'Content-Type': 'application/json',
|
||||||
|
Accept: 'text/event-stream',
|
||||||
|
'User-Agent': 'antigravity/1.19.6 darwin/arm64',
|
||||||
|
});
|
||||||
|
|
||||||
|
const upstreamBody = JSON.parse(String(requestInit.body));
|
||||||
|
expect(upstreamBody).toMatchObject({
|
||||||
|
project: 'project-demo',
|
||||||
|
model: 'gemini-3-pro-low',
|
||||||
|
userAgent: 'antigravity',
|
||||||
|
requestType: 'agent',
|
||||||
|
});
|
||||||
|
|
||||||
|
expect(response.headers.get('content-type')).toContain('application/json');
|
||||||
|
await expect(response.json()).resolves.toEqual({
|
||||||
|
responseId: 'antigravity-stream-1',
|
||||||
|
modelVersion: 'gemini-3-pro-low',
|
||||||
|
candidates: [
|
||||||
|
{
|
||||||
|
index: 0,
|
||||||
|
content: {
|
||||||
|
role: 'model',
|
||||||
|
parts: [{ text: 'hello from antigravity' }],
|
||||||
|
},
|
||||||
|
finishReason: 'STOP',
|
||||||
|
},
|
||||||
|
],
|
||||||
|
usageMetadata: {
|
||||||
|
promptTokenCount: 5,
|
||||||
|
candidatesTokenCount: 4,
|
||||||
|
totalTokenCount: 9,
|
||||||
|
},
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
it('aggregates an unterminated trailing SSE event for antigravity non-stream stream-endpoint requests', async () => {
|
||||||
|
const { dispatchRuntimeRequest } = await import('./runtimeExecutor.js');
|
||||||
|
const request: BuiltEndpointRequest = {
|
||||||
|
endpoint: 'chat',
|
||||||
|
path: '/v1internal:streamGenerateContent?alt=sse',
|
||||||
|
headers: {
|
||||||
|
Authorization: 'Bearer antigravity-token',
|
||||||
|
'Content-Type': 'application/json',
|
||||||
|
},
|
||||||
|
body: {
|
||||||
|
project: 'project-demo',
|
||||||
|
model: 'gemini-3-pro-low',
|
||||||
|
request: {
|
||||||
|
contents: [
|
||||||
|
{
|
||||||
|
role: 'user',
|
||||||
|
parts: [{ text: 'hello runtime executor' }],
|
||||||
|
},
|
||||||
|
],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
runtime: {
|
||||||
|
executor: 'antigravity',
|
||||||
|
modelName: 'gemini-3-pro-low',
|
||||||
|
stream: false,
|
||||||
|
action: 'streamGenerateContent',
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
fetchMock.mockResolvedValueOnce(new Response([
|
||||||
|
'data: {"response":{"responseId":"antigravity-stream-tail","modelVersion":"gemini-3-pro-low","candidates":[{"content":{"role":"model","parts":[{"text":"hello "}]},"index":0}]}}',
|
||||||
|
'',
|
||||||
|
'data: {"response":{"candidates":[{"content":{"role":"model","parts":[{"text":"tail"}]},"finishReason":"STOP","index":0}],"usageMetadata":{"promptTokenCount":7,"candidatesTokenCount":5,"totalTokenCount":12}}}',
|
||||||
|
].join('\n'), {
|
||||||
|
status: 200,
|
||||||
|
headers: { 'content-type': 'text/event-stream; charset=utf-8' },
|
||||||
|
}) as unknown as Awaited<ReturnType<typeof fetch>>);
|
||||||
|
|
||||||
|
const response = await dispatchRuntimeRequest({
|
||||||
|
siteUrl: 'https://cloudcode-pa.googleapis.com',
|
||||||
|
request,
|
||||||
|
buildInit: (_url, nextRequest) => ({
|
||||||
|
method: 'POST',
|
||||||
|
headers: nextRequest.headers,
|
||||||
|
body: JSON.stringify(nextRequest.body),
|
||||||
|
}),
|
||||||
|
});
|
||||||
|
|
||||||
|
expect(response.ok).toBe(true);
|
||||||
|
await expect(response.json()).resolves.toEqual({
|
||||||
|
responseId: 'antigravity-stream-tail',
|
||||||
|
modelVersion: 'gemini-3-pro-low',
|
||||||
|
candidates: [
|
||||||
|
{
|
||||||
|
index: 0,
|
||||||
|
content: {
|
||||||
|
role: 'model',
|
||||||
|
parts: [{ text: 'hello tail' }],
|
||||||
|
},
|
||||||
|
finishReason: 'STOP',
|
||||||
|
},
|
||||||
|
],
|
||||||
|
usageMetadata: {
|
||||||
|
promptTokenCount: 7,
|
||||||
|
candidatesTokenCount: 5,
|
||||||
|
totalTokenCount: 12,
|
||||||
|
},
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
it('keeps gemini-cli countTokens payload lean while forcing a model-aware user agent', async () => {
|
it('keeps gemini-cli countTokens payload lean while forcing a model-aware user agent', async () => {
|
||||||
const { dispatchRuntimeRequest } = await import('./runtimeExecutor.js');
|
const { dispatchRuntimeRequest } = await import('./runtimeExecutor.js');
|
||||||
const request: BuiltEndpointRequest = {
|
const request: BuiltEndpointRequest = {
|
||||||
@@ -181,7 +344,7 @@ describe('dispatchRuntimeRequest', () => {
|
|||||||
},
|
},
|
||||||
body: {
|
body: {
|
||||||
project: 'project-demo',
|
project: 'project-demo',
|
||||||
model: 'gemini-3-pro-preview',
|
model: 'gemini-2.5-pro',
|
||||||
request: {
|
request: {
|
||||||
contents: [
|
contents: [
|
||||||
{
|
{
|
||||||
@@ -193,7 +356,7 @@ describe('dispatchRuntimeRequest', () => {
|
|||||||
},
|
},
|
||||||
runtime: {
|
runtime: {
|
||||||
executor: 'antigravity',
|
executor: 'antigravity',
|
||||||
modelName: 'gemini-3-pro-preview',
|
modelName: 'gemini-2.5-pro',
|
||||||
stream: false,
|
stream: false,
|
||||||
},
|
},
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -1211,11 +1211,18 @@ describe('buildUpstreamEndpointRequest', () => {
|
|||||||
},
|
},
|
||||||
});
|
});
|
||||||
|
|
||||||
expect(request.path).toBe('/v1internal:generateContent');
|
expect(request.path).toBe('/v1internal:streamGenerateContent?alt=sse');
|
||||||
expect(request.headers.Authorization).toBe('Bearer oauth-access-token');
|
expect(request.headers.Authorization).toBe('Bearer oauth-access-token');
|
||||||
|
expect(request.headers.Accept).toBe('text/event-stream');
|
||||||
expect(request.headers['User-Agent']).toBe('antigravity/1.19.6 darwin/arm64');
|
expect(request.headers['User-Agent']).toBe('antigravity/1.19.6 darwin/arm64');
|
||||||
expect(request.headers['X-Goog-Api-Client']).toBeUndefined();
|
expect(request.headers['X-Goog-Api-Client']).toBeUndefined();
|
||||||
expect(request.headers['Client-Metadata']).toBeUndefined();
|
expect(request.headers['Client-Metadata']).toBeUndefined();
|
||||||
|
expect(request.runtime).toMatchObject({
|
||||||
|
executor: 'antigravity',
|
||||||
|
modelName: 'gemini-3-pro-preview',
|
||||||
|
stream: false,
|
||||||
|
action: 'streamGenerateContent',
|
||||||
|
});
|
||||||
expect(request.body).toEqual({
|
expect(request.body).toEqual({
|
||||||
project: 'project-demo',
|
project: 'project-demo',
|
||||||
model: 'gemini-3-pro-preview',
|
model: 'gemini-3-pro-preview',
|
||||||
|
|||||||
Reference in New Issue
Block a user