/**
* 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,
});
}