Something went wrong. Try again.
antigravity oauth for prime-agent ported from omp
Something went wrong. Try again.
prime-antigravity stream.ts
33 kB · 1029 lines
TypeScript
at main
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030/** * Antigravity stream provider. * * Wire protocol, headers, retry/failover behavior, and the per-request * spoofing envelope (sessionId/requestId/labels) are copied from oh-my-pi's * google-gemini-cli provider for the google-antigravity provider. */import { createHash, randomUUID } from "node:crypto";import { type Api, type AssistantMessage, type AssistantMessageEventStream, type Context, calculateCost, createAssistantMessageEventStream, type Model, type TextContent, type ThinkingContent, type ThinkingLevel, type ToolCall,} from "@earendil-works/pi-ai";import { type AntigravitySimpleStreamOptions, buildBaseOptions, clampReasoning, convertMessages, convertTools, type FunctionCallingConfigMode, type GoogleContent, type GoogleThinkingConfig, type GoogleThinkingLevel, headersToRecord, isThinkingPart, mapStopReasonString, mapToolChoice, retainThoughtSignature, sanitizeSurrogates,} from "./helpers.js";import { type AntigravityModelProfile, getAntigravityProfile } from "./models.js";import { ANTIGRAVITY_DAILY_ENDPOINT, ANTIGRAVITY_ENDPOINT_FALLBACKS, ANTIGRAVITY_SANDBOX_ENDPOINT, CLAUDE_THINKING_BETA_HEADER, getAntigravityModelWireProfile, getAntigravityUserAgent,} from "./wire.js";
const GOOGLE_THINKING: Partial<Record<ThinkingLevel, number>> = { minimal: 1024, low: 4096, medium: 8192, high: 16384, xhigh: 24575, max: 32768,};
const MIN_OUTPUT_TOKENS = 1024;
/** Map a normalized effort to Google's thinkingLevel enum for google-level models. */function mapEffortToGoogleLevel(effort: ThinkingLevel, minimalIsLow: boolean): GoogleThinkingLevel { if (effort === "minimal") return minimalIsLow ? "LOW" : "MINIMAL"; switch (effort) { case "low": return "LOW"; case "medium": return "MEDIUM"; default: return "HIGH"; }}
const FORCED_TOOL_DIRECTIVE = "TOOL-ONLY TURN. This turn accepts a tool call and nothing else; a text reply here is discarded unread and you will be re-prompted. Emit the tool call now.\n";
const MAX_RETRIES = 3;const BASE_DELAY_MS = 1000;const MAX_EMPTY_STREAM_RETRIES = 2;const EMPTY_STREAM_BASE_DELAY_MS = 500;const RATE_LIMIT_BUDGET_MS = 5 * 60 * 1000;
let toolCallCounter = 0;
interface CloudCodeAssistRequest { project: string; model: string; request: { contents: GoogleContent[]; sessionId?: string; systemInstruction?: { role?: string; parts: { text: string }[] }; generationConfig?: { maxOutputTokens?: number; temperature?: number; topP?: number; topK?: number; minP?: number; presencePenalty?: number; repetitionPenalty?: number; thinkingConfig?: GoogleThinkingConfig; }; tools?: { functionDeclarations: Record<string, unknown>[] }[] | undefined; toolConfig?: { functionCallingConfig: { mode: FunctionCallingConfigMode; allowedFunctionNames?: string[]; }; }; labels?: Record<string, string>; }; requestType?: string; userAgent?: string; requestId?: string;}
interface CloudCodeAssistResponseChunk { response?: { candidates?: Array<{ content?: { role: string; parts?: Array<{ text?: string; thought?: boolean; thoughtSignature?: string; functionCall?: { name: string; args: Record<string, unknown>; id?: string }; }>; }; finishReason?: string; }>; usageMetadata?: { promptTokenCount?: number; candidatesTokenCount?: number; thoughtsTokenCount?: number; totalTokenCount?: number; cachedContentTokenCount?: number; }; modelVersion?: string; responseId?: string; promptFeedback?: { blockReason?: string; blockReasonMessage?: string }; }; error?: { code?: number; message?: string; status?: string }; traceId?: string;}
interface ParsedCredentials { accessToken: string; projectId: string; email?: string;}
function parseCredentials(apiKeyRaw: string): ParsedCredentials { let raw: unknown; try { raw = JSON.parse(apiKeyRaw); } catch { throw new Error("Invalid Google Cloud Code Assist credentials. Use /login to re-authenticate."); } if (typeof raw !== "object" || raw === null) { throw new Error("Invalid Google Cloud Code Assist credentials. Use /login to re-authenticate."); } const parsed = raw as { token?: unknown; projectId?: unknown; project_id?: unknown; email?: unknown }; const token = typeof parsed.token === "string" ? parsed.token : undefined; const projectId = typeof parsed.projectId === "string" ? parsed.projectId : typeof parsed.project_id === "string" ? parsed.project_id : undefined; const email = typeof parsed.email === "string" && parsed.email.length > 0 ? parsed.email : undefined; if (!token || !projectId) { throw new Error("Missing token or projectId in Google Cloud credentials. Use /login to re-authenticate."); } return { accessToken: token, projectId, email };}
function sleep(ms: number, signal?: AbortSignal): Promise<void> { return new Promise((resolve, reject) => { if (signal?.aborted) { reject(new Error("Request was aborted")); return; } const timeout = setTimeout(resolve, ms); signal?.addEventListener("abort", () => { clearTimeout(timeout); reject(new Error("Request was aborted")); }); });}
function extractRetryDelay(errorText: string, response?: Response | Headers): number | undefined { const normalizeDelay = (ms: number): number | undefined => (ms > 0 ? Math.ceil(ms + 1000) : undefined); const headers = response instanceof Headers ? response : response?.headers; if (headers) { const retryAfter = headers.get("retry-after"); if (retryAfter) { const retryAfterSeconds = Number(retryAfter); if (Number.isFinite(retryAfterSeconds)) { const delay = normalizeDelay(retryAfterSeconds * 1000); if (delay !== undefined) return delay; } const retryAfterDate = new Date(retryAfter); const retryAfterMs = retryAfterDate.getTime(); if (!Number.isNaN(retryAfterMs)) { const delay = normalizeDelay(retryAfterMs - Date.now()); if (delay !== undefined) return delay; } } const rateLimitReset = headers.get("x-ratelimit-reset"); if (rateLimitReset) { const resetSeconds = Number.parseInt(rateLimitReset, 10); if (!Number.isNaN(resetSeconds)) { const delay = normalizeDelay(resetSeconds * 1000 - Date.now()); if (delay !== undefined) return delay; } } const rateLimitResetAfter = headers.get("x-ratelimit-reset-after"); if (rateLimitResetAfter) { const resetAfterSeconds = Number(rateLimitResetAfter); if (Number.isFinite(resetAfterSeconds)) { const delay = normalizeDelay(resetAfterSeconds * 1000); if (delay !== undefined) return delay; } } }
const durationMatch = errorText.match(/reset after (?:(\d+)h)?(?:(\d+)m)?(\d+(?:\.\d+)?)s/i); if (durationMatch) { const hours = durationMatch[1] ? parseInt(durationMatch[1], 10) : 0; const minutes = durationMatch[2] ? parseInt(durationMatch[2], 10) : 0; const seconds = parseFloat(durationMatch[3]); if (!Number.isNaN(seconds)) { const totalMs = ((hours * 60 + minutes) * 60 + seconds) * 1000; const delay = normalizeDelay(totalMs); if (delay !== undefined) return delay; } }
const retryInMatch = errorText.match(/Please retry in ([0-9.]+)(ms|s)/i); if (retryInMatch?.[1]) { const value = parseFloat(retryInMatch[1]); if (!Number.isNaN(value) && value > 0) { const ms = retryInMatch[2].toLowerCase() === "ms" ? value : value * 1000; const delay = normalizeDelay(ms); if (delay !== undefined) return delay; } }
const retryDelayMatch = errorText.match(/"retryDelay":\s*"([0-9.]+)(ms|s)"/i); if (retryDelayMatch?.[1]) { const value = parseFloat(retryDelayMatch[1]); if (!Number.isNaN(value) && value > 0) { const ms = retryDelayMatch[2].toLowerCase() === "ms" ? value : value * 1000; const delay = normalizeDelay(ms); if (delay !== undefined) return delay; } }
return undefined;}
function isRetryableError(status: number, errorText: string): boolean { if (status === 408 || status === 429 || status >= 500) { return true; } return /resource.?exhausted|rate.?limit|overloaded|service.?unavailable|other.?side.?closed/i.test(errorText);}
function extractErrorMessage(errorText: string): string { try { const parsed = JSON.parse(errorText) as { error?: { message?: string } }; if (parsed.error?.message) { return parsed.error.message; } } catch { // Not JSON, return as-is } return errorText;}
function extractGoogleValidationUrl(errorBody: string): string | undefined { if (!errorBody.includes("VALIDATION_REQUIRED")) return undefined; const start = errorBody.indexOf("{"); if (start === -1) return undefined; try { const parsed = JSON.parse(errorBody.slice(start)) as { error?: { details?: Array<{ reason?: string; metadata?: { validation_url?: string } }> }; }; const detail = parsed.error?.details?.find( (d) => d.reason === "VALIDATION_REQUIRED" && typeof d.metadata?.validation_url === "string", ); return detail?.metadata?.validation_url; } catch { return undefined; }}
function formatAntigravityApiError(status: number, errorText: string, email?: string): string { const validationUrl = extractGoogleValidationUrl(errorText); if (validationUrl) { const account = email ? ` for ${email}` : ""; return `Cloud Code Assist API error (${status}): Account verification required${account}. Visit ${validationUrl} to continue, then retry your request.`; } return `Cloud Code Assist API error (${status}): ${extractErrorMessage(errorText)}`;}
const INT63_MASK = (1n << 63n) - 1n;function formatSignedDecimalSessionId(value: bigint): string { return `-${value.toString()}`;}
function deriveSignedDecimalFromHash(text: string): string { const digest = createHash("sha256").update(text).digest(); let value = 0n; for (let index = 0; index < 8; index += 1) { value = (value << 8n) | BigInt(digest[index] ?? 0); } return formatSignedDecimalSessionId(value & INT63_MASK);}
function getFirstUserText(context: Context): string | undefined { for (const message of context.messages) { if (message.role !== "user") continue; if (typeof message.content === "string") return message.content; if (Array.isArray(message.content)) { const firstTextPart = message.content.find((item): item is TextContent => item.type === "text"); return firstTextPart?.text; } return undefined; } return undefined;}
function deriveAntigravitySessionId(context: Context): string { const text = getFirstUserText(context); if (text && text.trim().length > 0) { return deriveSignedDecimalFromHash(text); } const digest = createHash("sha256").update(randomUUID()).digest(); let value = 0n; for (let index = 0; index < 8; index += 1) { value = (value << 8n) | BigInt(digest[index] ?? 0); } return formatSignedDecimalSessionId(value & INT63_MASK);}
interface AntigravitySessionState { agentId: string; trajectoryId: string; sessionId: string; stepIndex: number; lastExecutionId?: string;}
const sessionStates = new Map<string, AntigravitySessionState>();
function getSessionState(key: string, context: Context): AntigravitySessionState { const existing = sessionStates.get(key); if (existing) return existing; const state: AntigravitySessionState = { agentId: randomUUID(), trajectoryId: randomUUID(), sessionId: deriveAntigravitySessionId(context), stepIndex: 1, }; sessionStates.set(key, state); return state;}
function buildAntigravityRequestEnvelope( profile: AntigravityModelProfile, state: AntigravitySessionState,): { sessionId: string; requestId: string; labels: Record<string, string> } { state.stepIndex = state.stepIndex + 1; const step = state.stepIndex; const requestId = `agent/${state.agentId}/${Date.now()}/${state.trajectoryId}/${step}`; const labels: Record<string, string> = {}; if (state.lastExecutionId) labels.last_execution_id = state.lastExecutionId; labels.last_step_index = String(step - 1); if (profile.modelEnum !== undefined) labels.model_enum = profile.modelEnum; labels.trajectory_id = state.trajectoryId; const usageLabel = String(profile.identityClass === "anthropic"); labels.used_claude = usageLabel; labels.used_claude_conservative = usageLabel; return { sessionId: state.sessionId, requestId, labels };}
export interface AntigravityStreamOptions extends AntigravitySimpleStreamOptions { toolChoice?: "auto" | "none" | "any"; topP?: number; topK?: number; minP?: number; presencePenalty?: number; repetitionPenalty?: number;}
export function buildRequest( model: Model<Api>, context: Context, projectId: string, options: AntigravityStreamOptions = {}, profile: AntigravityModelProfile = getAntigravityProfile(model.id), state: AntigravitySessionState = getSessionState(options?.sessionId ?? "default", context),): CloudCodeAssistRequest { const contents = convertMessages(model, context);
const generationConfig: CloudCodeAssistRequest["request"]["generationConfig"] = {}; if (options.temperature !== undefined) { generationConfig.temperature = options.temperature; } if (options.maxTokens !== undefined) { generationConfig.maxOutputTokens = options.maxTokens; } if (options.topP !== undefined) { generationConfig.topP = options.topP; } if (options.topK !== undefined) { generationConfig.topK = options.topK; } if (options.minP !== undefined) { generationConfig.minP = options.minP; } if (options.presencePenalty !== undefined) { generationConfig.presencePenalty = options.presencePenalty; } if (options.repetitionPenalty !== undefined) { generationConfig.repetitionPenalty = options.repetitionPenalty; }
const thinking = profile.thinking; if (model.reasoning && options.reasoning && options.reasoning !== "off") { const effort = clampReasoning(options.reasoning)!; if (thinking.mode === "google-level") { generationConfig.thinkingConfig = { includeThoughts: true, thinkingLevel: mapEffortToGoogleLevel(effort, thinking.minimalLevelIsLow ?? false), }; } else { const customBudgets = options.thinkingBudgets; let thinkingBudget = customBudgets?.[effort] ?? thinking.effortBudgets?.[effort] ?? GOOGLE_THINKING[effort] ?? 0; const maxTokens = Math.min((options.maxTokens ?? model.maxTokens) + thinkingBudget, model.maxTokens); if (maxTokens <= thinkingBudget) { thinkingBudget = Math.max(0, maxTokens - MIN_OUTPUT_TOKENS); } if (thinkingBudget > 0) { generationConfig.thinkingConfig = { includeThoughts: true, thinkingBudget }; generationConfig.maxOutputTokens = maxTokens; } } } else if (model.reasoning && thinking.suppressWhenOff) { // Omit thinkingConfig when off is not special-cased; otherwise CCA // re-applies the per-id baked server default and bills hidden reasoning. generationConfig.thinkingConfig = { includeThoughts: false, ...(thinking.mode === "google-level" ? { thinkingLevel: "MINIMAL" } : { thinkingBudget: 0 }), }; }
const request: CloudCodeAssistRequest["request"] = { contents };
if (context.systemPrompt) { request.systemInstruction = { role: "user", parts: [{ text: sanitizeSurrogates(context.systemPrompt) }], }; }
if (context.tools && context.tools.length > 0) { // Antigravity / Cloud Code Assist requires the legacy `parameters` field // regardless of model family. request.tools = convertTools(context.tools, true); if (options.toolChoice) { request.toolConfig = { functionCallingConfig: { mode: mapToolChoice(options.toolChoice), }, }; // Cloud Code Assist drops toolConfig on Antigravity's Gemini routes: // the backend answers in text under mode ANY and still emits calls // under NONE. Claude routes implement it, so only Gemini needs the // forced choice restated in the transcript. if (profile.identityClass !== "anthropic" && request.toolConfig?.functionCallingConfig.mode === "ANY") { contents.push({ role: "user", parts: [{ text: FORCED_TOOL_DIRECTIVE }] }); } } // Antigravity's default tool mode is VALIDATED. if (!request.toolConfig) { request.toolConfig = { functionCallingConfig: { mode: "VALIDATED" }, }; } }
// Claude on Antigravity always forces VALIDATED, even with no tools declared. if (profile.identityClass === "anthropic") { request.toolConfig = { functionCallingConfig: { mode: "VALIDATED" }, }; }
const wireModelId = profile.wireModelId; const wireProfile = getAntigravityModelWireProfile(wireModelId); if (wireProfile) { generationConfig.maxOutputTokens = wireProfile.maxOutputTokens; }
const envelope = buildAntigravityRequestEnvelope(profile, state); request.labels = envelope.labels; if (Object.keys(generationConfig).length > 0) { request.generationConfig = generationConfig; } request.sessionId = envelope.sessionId;
return { project: projectId, requestId: envelope.requestId, request, model: wireModelId, userAgent: "antigravity", requestType: "agent", };}
export function streamAntigravity( model: Model<Api>, context: Context, options?: AntigravityStreamOptions,): AssistantMessageEventStream { const stream = createAssistantMessageEventStream();
(async () => { const output: AssistantMessage = { role: "assistant", content: [], api: model.api, provider: model.provider, model: model.id, usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "stop", timestamp: Date.now(), };
try { const apiKeyRaw = options?.apiKey; if (!apiKeyRaw) { throw new Error( "Antigravity requires OAuth authentication. Use /login google-antigravity to authenticate.", ); } const credentials = parseCredentials(apiKeyRaw); const profile = getAntigravityProfile(model.id); const state = getSessionState(options?.sessionId ?? "default", context);
const baseUrl = model.baseUrl?.trim(); const cleanUrl = baseUrl ? baseUrl.replace(/\/+$/, "") : ""; let endpoints: string[]; if (cleanUrl && cleanUrl !== ANTIGRAVITY_DAILY_ENDPOINT && cleanUrl !== ANTIGRAVITY_SANDBOX_ENDPOINT) { endpoints = [cleanUrl]; } else { endpoints = [...ANTIGRAVITY_ENDPOINT_FALLBACKS]; }
let requestBody = buildRequest(model, context, credentials.projectId, options ?? {}, profile, state); const nextRequestBody = await options?.onPayload?.(requestBody, model); if (nextRequestBody !== undefined) { requestBody = nextRequestBody as CloudCodeAssistRequest; }
const requestHeaders: Record<string, string> = { Authorization: `Bearer ${credentials.accessToken}`, "Content-Type": "application/json", Accept: "text/event-stream", "User-Agent": getAntigravityUserAgent(), ...(profile.identityClass === "anthropic" && model.reasoning ? { "anthropic-beta": CLAUDE_THINKING_BETA_HEADER } : {}), ...options?.headers, }; const requestBodyJson = JSON.stringify(requestBody);
let response: Response | undefined; let lastError: Error | undefined; let requestUrl: string | undefined; let endpointIndex = 0;
for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) { if (options?.signal?.aborted) { throw new Error("Request was aborted"); }
try { const endpoint = endpoints[endpointIndex]; requestUrl = `${endpoint}/v1internal:streamGenerateContent?alt=sse`; response = await fetch(requestUrl, { method: "POST", headers: requestHeaders, body: requestBodyJson, signal: options?.signal, }); await options?.onResponse?.( { status: response.status, headers: headersToRecord(response.headers) }, model, );
if (response.ok) { break; }
const errorText = await response.text();
if ((response.status === 403 || response.status === 404) && endpointIndex < endpoints.length - 1) { endpointIndex++; continue; }
if (attempt < MAX_RETRIES && isRetryableError(response.status, errorText)) { if (endpointIndex < endpoints.length - 1) { endpointIndex++; } const serverDelay = extractRetryDelay(errorText, response); const delayMs = serverDelay ?? BASE_DELAY_MS * 2 ** attempt; const maxDelayMs = options?.maxRetryDelayMs ?? RATE_LIMIT_BUDGET_MS; if (maxDelayMs > 0 && serverDelay && serverDelay > maxDelayMs) { const delaySeconds = Math.ceil(serverDelay / 1000); throw new Error( `Server requested ${delaySeconds}s retry delay (max: ${Math.ceil(maxDelayMs / 1000)}s). ${extractErrorMessage(errorText)}`, ); } await sleep(delayMs, options?.signal); continue; }
throw new Error(formatAntigravityApiError(response.status, errorText, credentials.email)); } catch (error) { if (error instanceof Error) { if (error.name === "AbortError" || error.message === "Request was aborted") { throw new Error("Request was aborted"); } } lastError = error instanceof Error ? error : new Error(String(error)); if (lastError.message === "fetch failed" && lastError.cause instanceof Error) { lastError = new Error(`Network error: ${lastError.cause.message}`); } if (attempt < MAX_RETRIES) { const delayMs = BASE_DELAY_MS * 2 ** attempt; await sleep(delayMs, options?.signal); continue; } throw lastError; } }
if (!response || !response.ok) { throw lastError ?? new Error("Failed to get response after retries"); }
let started = false; const ensureStarted = () => { if (!started) { stream.push({ type: "start", partial: output }); started = true; } };
const resetOutput = () => { output.content = []; output.usage = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }; output.stopReason = "stop"; output.errorMessage = undefined; output.timestamp = Date.now(); started = false; };
const streamResponse = async (activeResponse: Response): Promise<boolean> => { if (!activeResponse.body) { throw new Error("No response body"); }
let hasContent = false; let currentBlock: TextContent | ThinkingContent | null = null; const blocks = output.content; const blockIndex = () => blocks.length - 1;
const reader = activeResponse.body.getReader(); const decoder = new TextDecoder(); let buffer = "";
const abortHandler = () => { void reader.cancel().catch(() => {}); }; options?.signal?.addEventListener("abort", abortHandler);
try { while (true) { if (options?.signal?.aborted) { throw new Error("Request was aborted"); }
const { done, value } = await reader.read(); if (done) break;
buffer += decoder.decode(value, { stream: true }); const lines = buffer.split("\n"); buffer = lines.pop() || "";
for (const line of lines) { if (!line.startsWith("data:")) continue;
const jsonStr = line.slice(5).trim(); if (!jsonStr) continue;
let chunk: CloudCodeAssistResponseChunk; try { chunk = JSON.parse(jsonStr); } catch { continue; }
if (chunk.error) { const detail = chunk.error.message || chunk.error.status || "unknown error"; throw new Error(`Cloud Code Assist stream error: ${detail}`); }
const responseData = chunk.response; if (!responseData) continue; output.responseId ||= responseData.responseId;
if (!responseData.candidates?.length && responseData.promptFeedback?.blockReason) { const detail = responseData.promptFeedback.blockReasonMessage; throw new Error( `Request blocked by Google (${responseData.promptFeedback.blockReason})${detail ? `: ${detail}` : ""}`, ); }
const candidate = responseData.candidates?.[0]; if (candidate?.content?.parts) { for (const part of candidate.content.parts) { if (part.text !== undefined && part.text !== "") { hasContent = true; const isThinking = isThinkingPart(part); if ( !currentBlock || (isThinking && currentBlock.type !== "thinking") || (!isThinking && currentBlock.type !== "text") ) { if (currentBlock) { if (currentBlock.type === "text") { stream.push({ type: "text_end", contentIndex: blockIndex(), content: currentBlock.text, partial: output, }); } else { stream.push({ type: "thinking_end", contentIndex: blockIndex(), content: currentBlock.thinking, partial: output, }); } } if (isThinking) { currentBlock = { type: "thinking", thinking: "", thinkingSignature: undefined }; output.content.push(currentBlock); ensureStarted(); stream.push({ type: "thinking_start", contentIndex: blockIndex(), partial: output, }); } else { currentBlock = { type: "text", text: "" }; output.content.push(currentBlock); ensureStarted(); stream.push({ type: "text_start", contentIndex: blockIndex(), partial: output }); } } if (currentBlock.type === "thinking") { currentBlock.thinking += part.text; currentBlock.thinkingSignature = retainThoughtSignature( currentBlock.thinkingSignature, part.thoughtSignature, ); stream.push({ type: "thinking_delta", contentIndex: blockIndex(), delta: part.text, partial: output, }); } else { currentBlock.text += part.text; currentBlock.textSignature = retainThoughtSignature( currentBlock.textSignature, part.thoughtSignature, ); stream.push({ type: "text_delta", contentIndex: blockIndex(), delta: part.text, partial: output, }); } } else if (part.text === "" && part.thoughtSignature && !part.functionCall) { // Terminal thought signatures on empty text parts (e.g. at stream end). if (currentBlock) { if (currentBlock.type === "thinking") { currentBlock.thinkingSignature = retainThoughtSignature( currentBlock.thinkingSignature, part.thoughtSignature, ); } else { currentBlock.textSignature = retainThoughtSignature( currentBlock.textSignature, part.thoughtSignature, ); } } }
if (part.functionCall) { hasContent = true; if (currentBlock) { if (currentBlock.type === "text") { stream.push({ type: "text_end", contentIndex: blockIndex(), content: currentBlock.text, partial: output, }); } else { stream.push({ type: "thinking_end", contentIndex: blockIndex(), content: currentBlock.thinking, partial: output, }); } currentBlock = null; }
const providedId = part.functionCall.id; const needsNewId = !providedId || output.content.some((b) => b.type === "toolCall" && b.id === providedId); const toolCallId = needsNewId ? `${part.functionCall.name}_${Date.now()}_${++toolCallCounter}` : providedId;
const toolCall: ToolCall = { type: "toolCall", id: toolCallId, name: part.functionCall.name || "", arguments: (part.functionCall.args as Record<string, unknown>) ?? {}, ...(part.thoughtSignature && { thoughtSignature: part.thoughtSignature }), };
output.content.push(toolCall); ensureStarted(); stream.push({ type: "toolcall_start", contentIndex: blockIndex(), partial: output }); stream.push({ type: "toolcall_delta", contentIndex: blockIndex(), delta: JSON.stringify(toolCall.arguments), partial: output, }); stream.push({ type: "toolcall_end", contentIndex: blockIndex(), toolCall, partial: output, }); } } }
if (candidate?.finishReason) { output.stopReason = mapStopReasonString(candidate.finishReason); if (output.content.some((b) => b.type === "toolCall")) { output.stopReason = "toolUse"; } }
if (responseData.usageMetadata) { const promptTokens = responseData.usageMetadata.promptTokenCount || 0; const cacheReadTokens = responseData.usageMetadata.cachedContentTokenCount || 0; const thinkingTokens = responseData.usageMetadata.thoughtsTokenCount || 0; output.usage = { input: promptTokens - cacheReadTokens, output: (responseData.usageMetadata.candidatesTokenCount || 0) + thinkingTokens, cacheRead: cacheReadTokens, cacheWrite: 0, totalTokens: responseData.usageMetadata.totalTokenCount || 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0, }, }; calculateCost(model, output.usage); } } } } finally { options?.signal?.removeEventListener("abort", abortHandler); }
if (currentBlock) { if (currentBlock.type === "text") { stream.push({ type: "text_end", contentIndex: blockIndex(), content: currentBlock.text, partial: output, }); } else { stream.push({ type: "thinking_end", contentIndex: blockIndex(), content: currentBlock.thinking, partial: output, }); } }
return hasContent; };
let receivedContent = false; let currentResponse = response;
for (let emptyAttempt = 0; emptyAttempt <= MAX_EMPTY_STREAM_RETRIES; emptyAttempt++) { if (options?.signal?.aborted) { throw new Error("Request was aborted"); }
if (emptyAttempt > 0) { const backoffMs = EMPTY_STREAM_BASE_DELAY_MS * 2 ** (emptyAttempt - 1); await sleep(backoffMs, options?.signal);
if (!requestUrl) { throw new Error("Missing request URL"); }
currentResponse = await fetch(requestUrl, { method: "POST", headers: requestHeaders, body: requestBodyJson, signal: options?.signal, }); await options?.onResponse?.( { status: currentResponse.status, headers: headersToRecord(currentResponse.headers) }, model, );
if (!currentResponse.ok) { const retryErrorText = await currentResponse.text(); throw new Error(`Cloud Code Assist API error (${currentResponse.status}): ${retryErrorText}`); } }
const streamed = await streamResponse(currentResponse); if (streamed) { receivedContent = true; break; }
if (emptyAttempt < MAX_EMPTY_STREAM_RETRIES) { resetOutput(); } }
if (!receivedContent) { throw new Error("Cloud Code Assist API returned an empty response"); }
if (options?.signal?.aborted) { throw new Error("Request was aborted"); }
if (output.stopReason === "aborted" || output.stopReason === "error") { throw new Error(output.errorMessage ?? "An unknown error occurred"); }
state.lastExecutionId = output.responseId; stream.push({ type: "done", reason: output.stopReason, message: output }); stream.end(); } catch (error) { output.stopReason = options?.signal?.aborted ? "aborted" : "error"; output.errorMessage = error instanceof Error ? error.message : JSON.stringify(error); stream.push({ type: "error", reason: output.stopReason, error: output }); stream.end(); } })();
return stream;}
export function streamSimpleAntigravity( model: Model<Api>, context: Context, options?: AntigravitySimpleStreamOptions,): AssistantMessageEventStream { const apiKey = options?.apiKey; if (!apiKey) { throw new Error("Antigravity requires OAuth authentication. Use /login google-antigravity to authenticate."); } const base = buildBaseOptions(model, options, apiKey); return streamAntigravity(model, context, { ...base, reasoning: options?.reasoning, thinkingBudgets: options?.thinkingBudgets, });}