/**
* Google Cloud Code Assist / Antigravity shared helpers.
* Adapted from oh-my-pi's google-gemini-cli + google-shared provider code for
* use inside a self-contained Prime extension.
*/
import type {
Api,
AssistantMessage,
Context,
ImageContent,
Message,
Model,
SimpleStreamOptions,
StopReason,
TextContent,
ThinkingLevel,
Tool,
ToolCall,
ToolResultMessage,
} from "@earendil-works/pi-ai";
import { normalizeSchemaForCCA } from "./schema/normalize.js";
/** pi-ai 0.9.4 removed maxRetries and maxRetryDelayMs, so the provider carries its own retry knobs. */
export interface AntigravitySimpleStreamOptions extends SimpleStreamOptions {
maxRetries?: number;
maxRetryDelayMs?: number;
}
export type GoogleThinkingLevel = "THINKING_LEVEL_UNSPECIFIED" | "MINIMAL" | "LOW" | "MEDIUM" | "HIGH";
export type FunctionCallingConfigMode = "AUTO" | "ANY" | "NONE" | "VALIDATED";
export interface GooglePart {
text?: string;
thought?: boolean;
thoughtSignature?: string;
functionCall?: { name: string; args: Record<string, unknown>; id?: string };
functionResponse?: { name: string; response: unknown; parts?: GooglePart[]; id?: string };
inlineData?: { mimeType: string; data: string };
}
export interface GoogleContent {
role: "user" | "model";
parts: GooglePart[];
}
export interface GoogleThinkingConfig {
includeThoughts?: boolean;
thinkingLevel?: GoogleThinkingLevel;
thinkingBudget?: number;
}
/** Removes unpaired Unicode surrogate characters from a string. */
export function sanitizeSurrogates(text: string): string {
return text.replace(/[\uD800-\uDBFF](?![\uDC00-\uDFFF])|(?<![\uD800-\uDBFF])[\uDC00-\uDFFF]/g, "");
}
export function headersToRecord(headers: Headers): Record<string, string> {
const result: Record<string, string> = {};
for (const [key, value] of headers.entries()) {
result[key] = value;
}
return result;
}
export function buildBaseOptions(
model: Model<Api>,
options?: AntigravitySimpleStreamOptions,
apiKey?: string,
): AntigravitySimpleStreamOptions {
return {
temperature: options?.temperature,
maxTokens: options?.maxTokens ?? (model.maxTokens > 0 ? model.maxTokens : undefined),
signal: options?.signal,
apiKey: apiKey || options?.apiKey,
transport: options?.transport,
serviceTier: options?.serviceTier,
cacheRetention: options?.cacheRetention,
sessionId: options?.sessionId,
headers: options?.headers,
onPayload: options?.onPayload,
onResponse: options?.onResponse,
timeoutMs: options?.timeoutMs,
maxRetries: options?.maxRetries,
maxRetryDelayMs: options?.maxRetryDelayMs,
metadata: options?.metadata,
};
}
export function clampReasoning(effort: ThinkingLevel | undefined): Exclude<ThinkingLevel, "xhigh" | "max"> | undefined {
return effort === "xhigh" || effort === "max" ? "high" : effort;
}
export function isThinkingPart(part: Pick<GooglePart, "thought" | "thoughtSignature">): boolean {
return part.thought === true;
}
export function retainThoughtSignature(existing: string | undefined, incoming: string | undefined): string | undefined {
if (typeof incoming === "string" && incoming.length > 0) return incoming;
return existing;
}
/** Whether this Google API model requires tool-call IDs on function calls and responses. */
export function isClaudeModel(modelId: string): boolean {
return modelId.toLowerCase().includes("claude");
}
export function requiresToolCallId(modelId: string): boolean {
return modelId.startsWith("claude-") || modelId.startsWith("gpt-oss-");
}
function getGeminiMajorVersion(modelId: string): number | undefined {
const match = modelId.toLowerCase().match(/^gemini(?:-live)?-(\d+)/);
if (!match) return undefined;
return Number.parseInt(match[1], 10);
}
function supportsMultimodalFunctionResponse(modelId: string): boolean {
const geminiMajorVersion = getGeminiMajorVersion(modelId);
if (geminiMajorVersion !== undefined) {
return geminiMajorVersion >= 3;
}
return true;
}
const base64SignaturePattern = /^[A-Za-z0-9+/]+={0,2}$/;
function isValidThoughtSignature(signature: string | undefined): boolean {
if (!signature) return false;
if (signature.length % 4 !== 0) return false;
return base64SignaturePattern.test(signature);
}
function resolveThoughtSignature(isSameProviderAndModel: boolean, signature: string | undefined): string | undefined {
return isSameProviderAndModel && isValidThoughtSignature(signature) ? signature : undefined;
}
const NON_VISION_USER_IMAGE_PLACEHOLDER = "(image omitted: model does not support images)";
const NON_VISION_TOOL_IMAGE_PLACEHOLDER = "(tool image omitted: model does not support images)";
function replaceImagesWithPlaceholder(content: (TextContent | ImageContent)[], placeholder: string): TextContent[] {
const result: TextContent[] = [];
let previousWasPlaceholder = false;
for (const block of content) {
if (block.type === "image") {
if (!previousWasPlaceholder) {
result.push({ type: "text", text: placeholder });
}
previousWasPlaceholder = true;
continue;
}
result.push(block);
previousWasPlaceholder = block.text === placeholder;
}
return result;
}
function downgradeUnsupportedImages<TApi extends Api>(messages: Message[], model: Model<TApi>): Message[] {
if (model.input.includes("image")) {
return messages;
}
return messages.map((msg) => {
if (msg.role === "user" && Array.isArray(msg.content)) {
return { ...msg, content: replaceImagesWithPlaceholder(msg.content, NON_VISION_USER_IMAGE_PLACEHOLDER) };
}
if (msg.role === "toolResult") {
return { ...msg, content: replaceImagesWithPlaceholder(msg.content, NON_VISION_TOOL_IMAGE_PLACEHOLDER) };
}
return msg;
});
}
export function transformMessages<TApi extends Api>(
messages: Message[],
model: Model<TApi>,
normalizeToolCallId?: (id: string, model: Model<TApi>, source: AssistantMessage) => string,
): Message[] {
const toolCallIdMap = new Map<string, string>();
const imageAwareMessages = downgradeUnsupportedImages(messages, model);
const transformed = imageAwareMessages.map((msg) => {
if (msg.role === "user") {
return msg;
}
if (msg.role === "toolResult") {
const normalizedId = toolCallIdMap.get(msg.toolCallId);
if (normalizedId && normalizedId !== msg.toolCallId) {
return { ...msg, toolCallId: normalizedId };
}
return msg;
}
if (msg.role === "assistant") {
const assistantMsg = msg as AssistantMessage;
const isSameModel =
assistantMsg.provider === model.provider &&
assistantMsg.api === model.api &&
assistantMsg.model === model.id;
const transformedContent = assistantMsg.content.flatMap((block) => {
if (block.type === "thinking") {
if (block.redacted) {
return isSameModel ? block : [];
}
if (isSameModel && block.thinkingSignature) return block;
if (!block.thinking || block.thinking.trim() === "") return [];
if (isSameModel) return block;
return { type: "text" as const, text: block.thinking };
}
if (block.type === "text") {
if (isSameModel) return block;
return { type: "text" as const, text: block.text };
}
if (block.type === "toolCall") {
const toolCall = block as ToolCall;
let normalizedToolCall: ToolCall = toolCall;
if (!isSameModel && toolCall.thoughtSignature) {
normalizedToolCall = { ...toolCall };
delete (normalizedToolCall as { thoughtSignature?: string }).thoughtSignature;
}
if (!isSameModel && normalizeToolCallId) {
const normalizedId = normalizeToolCallId(toolCall.id, model, assistantMsg);
if (normalizedId !== toolCall.id) {
toolCallIdMap.set(toolCall.id, normalizedId);
normalizedToolCall = { ...normalizedToolCall, id: normalizedId };
}
}
return normalizedToolCall;
}
return block;
});
return { ...assistantMsg, content: transformedContent };
}
return msg;
});
const result: Message[] = [];
let pendingToolCalls: ToolCall[] = [];
let existingToolResultIds = new Set<string>();
const insertSyntheticToolResults = () => {
if (pendingToolCalls.length > 0) {
for (const tc of pendingToolCalls) {
if (!existingToolResultIds.has(tc.id)) {
result.push({
role: "toolResult",
toolCallId: tc.id,
toolName: tc.name,
content: [{ type: "text", text: "No result provided" }],
isError: true,
timestamp: Date.now(),
} as ToolResultMessage);
}
}
pendingToolCalls = [];
existingToolResultIds = new Set();
}
};
for (let i = 0; i < transformed.length; i++) {
const msg = transformed[i];
if (msg.role === "assistant") {
insertSyntheticToolResults();
const assistantMsg = msg as AssistantMessage;
if (assistantMsg.stopReason === "error" || assistantMsg.stopReason === "aborted") {
continue;
}
const toolCalls = assistantMsg.content.filter((b) => b.type === "toolCall") as ToolCall[];
if (toolCalls.length > 0) {
pendingToolCalls = toolCalls;
existingToolResultIds = new Set();
}
result.push(msg);
} else if (msg.role === "toolResult") {
existingToolResultIds.add(msg.toolCallId);
result.push(msg);
} else if (msg.role === "user") {
insertSyntheticToolResults();
result.push(msg);
} else {
result.push(msg);
}
}
insertSyntheticToolResults();
return result;
}
/** Converts internal context to Google `Content[]`. */
export function convertMessages(model: Model<Api>, context: Context): GoogleContent[] {
const contents: GoogleContent[] = [];
const normalizeToolCallId = (id: string): string => {
if (!requiresToolCallId(model.id)) return id;
return id.replace(/[^a-zA-Z0-9_-]/g, "_").slice(0, 64);
};
const transformedMessages = transformMessages(context.messages, model, normalizeToolCallId);
for (const msg of transformedMessages) {
if (msg.role === "user") {
if (typeof msg.content === "string") {
contents.push({ role: "user", parts: [{ text: sanitizeSurrogates(msg.content) }] });
} else {
const parts: GooglePart[] = msg.content.map((item) => {
if (item.type === "text") {
return { text: sanitizeSurrogates(item.text) };
}
return { inlineData: { mimeType: item.mimeType, data: item.data } };
});
if (parts.length === 0) continue;
contents.push({ role: "user", parts });
}
} else if (msg.role === "assistant") {
const parts: GooglePart[] = [];
const isSameProviderAndModel = msg.provider === model.provider && msg.model === model.id;
for (const block of msg.content) {
if (block.type === "text") {
if (!block.text || block.text.trim() === "") continue;
const thoughtSignature = resolveThoughtSignature(isSameProviderAndModel, block.textSignature);
parts.push({ text: sanitizeSurrogates(block.text), ...(thoughtSignature && { thoughtSignature }) });
} else if (block.type === "thinking") {
if (!block.thinking || block.thinking.trim() === "") continue;
const thoughtSignature = resolveThoughtSignature(isSameProviderAndModel, block.thinkingSignature);
// Claude on Antigravity Cloud Code Assist rejects unsigned thinking parts; drop them.
if (isClaudeModel(model.id) && !thoughtSignature) continue;
// Same provider+model: always mark replayed thoughts, even unsigned ones.
// Gemini thoughts carry signatures only on functionCalls, so replaying
// unsigned thoughts as plain text makes the model emit its own thought
// summaries as visible text (thinking leak) on later turns.
if (isSameProviderAndModel) {
parts.push({
thought: true,
text: sanitizeSurrogates(block.thinking),
...(thoughtSignature && { thoughtSignature }),
});
} else {
parts.push({ text: sanitizeSurrogates(block.thinking) });
}
} else if (block.type === "toolCall") {
const thoughtSignature = resolveThoughtSignature(isSameProviderAndModel, block.thoughtSignature);
parts.push({
functionCall: {
name: block.name,
args: block.arguments ?? {},
...(requiresToolCallId(model.id) ? { id: block.id } : {}),
},
...(thoughtSignature && { thoughtSignature }),
});
}
}
if (parts.length === 0) continue;
contents.push({ role: "model", parts });
} else if (msg.role === "toolResult") {
const textContent = msg.content.filter((c): c is TextContent => c.type === "text");
const textResult = textContent.map((c) => c.text).join("\n");
const imageContent = model.input.includes("image")
? msg.content.filter((c): c is ImageContent => c.type === "image")
: [];
const hasText = textResult.length > 0;
const hasImages = imageContent.length > 0;
const modelSupportsMultimodalFunctionResponse = supportsMultimodalFunctionResponse(model.id);
const responseValue = hasText ? sanitizeSurrogates(textResult) : hasImages ? "(see attached image)" : "";
const imageParts: GooglePart[] = imageContent.map((imageBlock) => ({
inlineData: { mimeType: imageBlock.mimeType, data: imageBlock.data },
}));
const includeId = requiresToolCallId(model.id);
const functionResponsePart: GooglePart = {
functionResponse: {
name: msg.toolName,
response: msg.isError ? { error: responseValue } : { output: responseValue },
...(hasImages && modelSupportsMultimodalFunctionResponse && { parts: imageParts }),
...(includeId ? { id: msg.toolCallId } : {}),
},
};
const lastContent = contents[contents.length - 1];
if (lastContent?.role === "user" && lastContent.parts?.some((p) => p.functionResponse)) {
lastContent.parts.push(functionResponsePart);
} else {
contents.push({ role: "user", parts: [functionResponsePart] });
}
if (hasImages && !modelSupportsMultimodalFunctionResponse) {
contents.push({ role: "user", parts: [{ text: "Tool result image:" }, ...imageParts] });
}
}
}
return contents;
}
export function convertTools(
tools: Tool[],
useParameters = false,
): { functionDeclarations: Record<string, unknown>[] }[] | undefined {
if (tools.length === 0) return undefined;
return [
{
functionDeclarations: tools.map((tool) => ({
name: tool.name,
description: tool.description,
...(useParameters
? { parameters: normalizeSchemaForCCA(tool.parameters as unknown) }
: { parametersJsonSchema: tool.parameters }),
})),
},
];
}
export function mapToolChoice(choice: string): FunctionCallingConfigMode {
switch (choice) {
case "auto":
return "AUTO";
case "none":
return "NONE";
case "any":
return "ANY";
default:
return "AUTO";
}
}
export function mapStopReasonString(reason: string): StopReason {
switch (reason) {
case "STOP":
return "stop";
case "MAX_TOKENS":
return "length";
default:
return "error";
}
}