shariq-pi-extensions 0.1.0

package/extensions/antigravity-provider/antigravity/cloud-code-assist.ts

@@ -0,0 +1,519 @@
+import { randomUUID } from "node:crypto";
+import { createAssistantMessageEventStream, type AssistantMessageEventStream } from "@earendil-works/pi-ai";
+import {
+	lastEndpoint,
+	lastMatchedModelDebug,
+	lastAvailableModels,
+	setLastStatus,
+	setLastEndpoint,
+	setLastError,
+	setLastProjectId,
+	setLastResolvedRuntimeModel,
+	setLastAvailableModels,
+	endpointCandidates,
+	safeError,
+	sanitizeText,
+	parseApiKey,
+	antigravityEnv,
+	antigravityHeaders,
+	jsonOrTextError,
+	loadCodeAssist,
+	fetchAvailableRuntimeModel,
+	DEFAULT_PROJECT_ID,
+} from "./oauth.ts";
+import { getAntigravityRequestModelId, PROVIDER_ID } from "./models.ts";
+
+const ANTIGRAVITY_SYSTEM_INSTRUCTION =
+	"You are Antigravity, a powerful agentic AI coding assistant designed by Google DeepMind. " +
+	"You are pair programming with a user to solve coding tasks. Be concise, practical, and tool-aware.";
+
+const ANTIGRAVITY_NO_PREAMBLE_INSTRUCTION =
+	'CRITICAL: NEVER output rule checks, formatting guidelines, constraint checklists (e.g. "No emdashes"), or your thinking/personality preambles in the final response. Output only the final response.';
+
+let _toolCallCounter = 0;
+
+function sanitizeToolCallId(id: string, fallbackName?: string): string {
+	const cleaned = id.replace(/[^a-zA-Z0-9_-]/g, "_");
+	// Cap ID length to 64 characters (matching yofriadi / API limitations)
+	const capped = cleaned.slice(0, 64);
+	return capped || `${fallbackName || "tool"}_${Date.now()}_${++_toolCallCounter}`;
+}
+
+function toolCallIdNeeded(modelId: string, runtimeModel: string): boolean {
+	return modelId.startsWith("claude-") || modelId.startsWith("gpt-oss-") ||
+		runtimeModel.startsWith("claude-") || runtimeModel.startsWith("gpt-oss-");
+}
+
+function asTextParts(content: unknown): any[] {
+	if (typeof content === "string") return [{ text: sanitizeText(content) }];
+	if (!Array.isArray(content)) return [];
+	return content.flatMap<any>((item) => {
+		if (!item || typeof item !== "object") return [];
+		const block = item as any;
+		if (block.type === "text") return [{ text: sanitizeText(block.text) }];
+		if (block.type === "image") {
+			const data = block.data || block.source?.data;
+			const mimeType = block.mimeType || block.mediaType || block.source?.mediaType || "image/png";
+			return data ? [{ inlineData: { mimeType, data } }] : [];
+		}
+		return [];
+	});
+}
+
+function convertMessages(model: any, context: any, runtimeModel: string): any[] {
+	const contents: any[] = [];
+	const messages = Array.isArray(context.messages) ? context.messages : [];
+	for (const msg of messages) {
+		if (msg.role === "user") {
+			const parts = asTextParts(msg.content);
+			if (parts.length) contents.push({ role: "user", parts });
+		} else if (msg.role === "assistant") {
+			const parts: any[] = [];
+			for (const block of msg.content || []) {
+				if (block.type === "text" && String(block.text || "").trim()) parts.push({ text: sanitizeText(block.text) });
+				else if (block.type === "thinking" && String(block.thinking || "").trim()) {
+					if (msg.provider === PROVIDER_ID && msg.model === model.id) parts.push({ thought: true, text: sanitizeText(block.thinking), ...(block.thinkingSignature ? { thoughtSignature: block.thinkingSignature } : {}) });
+					else parts.push({ text: sanitizeText(block.thinking) });
+				} else if (block.type === "toolCall") {
+					parts.push({
+						functionCall: {
+							name: block.name,
+							args: block.arguments ?? {},
+							...(toolCallIdNeeded(model.id, runtimeModel) ? { id: sanitizeToolCallId(block.id || "", block.name) } : {}),
+						},
+						...(block.thoughtSignature ? { thoughtSignature: block.thoughtSignature } : {}),
+					});
+				}
+			}
+			if (parts.length) contents.push({ role: "model", parts });
+		} else if (msg.role === "toolResult") {
+			const content = Array.isArray(msg.content) ? msg.content : [];
+			const text = content.filter((c: any) => c.type === "text").map((c: any) => sanitizeText(c.text)).join("\n");
+			const responseText = text || (msg.isError ? "Tool failed" : "");
+			const part = {
+				functionResponse: {
+					name: msg.toolName,
+					response: msg.isError ? { error: responseText } : { output: responseText },
+					...(toolCallIdNeeded(model.id, runtimeModel) ? { id: sanitizeToolCallId(msg.toolCallId || "", msg.toolName) } : {}),
+				},
+			};
+			const last = contents[contents.length - 1];
+			if (last?.role === "user" && last.parts?.some((p: any) => p.functionResponse)) last.parts.push(part);
+			else contents.push({ role: "user", parts: [part] });
+		}
+	}
+	return contents;
+}
+
+function stripMetaSchema(schema: unknown): unknown {
+	if (!schema || typeof schema !== "object" || Array.isArray(schema)) return schema;
+	const omit = new Set(["$schema", "$id", "$anchor", "$dynamicAnchor", "$vocabulary", "$comment", "$defs", "definitions"]);
+	const out: Record<string, unknown> = {};
+	for (const [key, value] of Object.entries(schema)) {
+		if (!omit.has(key)) out[key] = stripMetaSchema(value);
+	}
+	return out;
+}
+
+function normalizeGoogleSchema(schema: unknown): unknown {
+	if (!schema || typeof schema !== "object") return schema;
+	if (Array.isArray(schema)) return schema.map(normalizeGoogleSchema);
+	const out: Record<string, unknown> = {};
+	for (const [key, value] of Object.entries(schema)) {
+		if (key === "type" && typeof value === "string") out[key] = value.toUpperCase();
+		else out[key] = normalizeGoogleSchema(value);
+	}
+	return out;
+}
+
+function convertTools(tools: any[] | undefined, useLegacyParameters = false): any[] | undefined {
+	if (!tools?.length) return undefined;
+	return [
+		{
+			functionDeclarations: tools.map((tool) => {
+				const parameters = stripMetaSchema(tool.parameters);
+				return {
+					name: tool.name,
+					description: tool.description,
+					...(useLegacyParameters
+						? { parameters: normalizeGoogleSchema(parameters) }
+						: { parametersJsonSchema: parameters }),
+				};
+			}),
+		},
+	];
+}
+
+type AgyWireMetadata = { modelEnum: string; thinkingBudget?: number };
+
+const AGY_WIRE_METADATA: Record<string, AgyWireMetadata> = {
+	"gemini-3.5-flash-extra-low": { modelEnum: "MODEL_PLACEHOLDER_M187", thinkingBudget: 1_000 },
+	"gemini-3.5-flash-low": { modelEnum: "MODEL_PLACEHOLDER_M20", thinkingBudget: 4_000 },
+	"gemini-3-flash-agent": { modelEnum: "MODEL_PLACEHOLDER_M84", thinkingBudget: 10_000 },
+	"gemini-3.6-flash-low": { modelEnum: "MODEL_PLACEHOLDER_M73", thinkingBudget: 1_000 },
+	"gemini-3.6-flash-medium": { modelEnum: "MODEL_PLACEHOLDER_M72", thinkingBudget: 4_000 },
+	"gemini-3.6-flash-high": { modelEnum: "MODEL_PLACEHOLDER_M71", thinkingBudget: 10_000 },
+	"gemini-3.7-flash-low": { modelEnum: "MODEL_PLACEHOLDER_M300", thinkingBudget: 1_000 },
+	"gemini-3.7-flash-medium": { modelEnum: "MODEL_PLACEHOLDER_M299", thinkingBudget: 4_000 },
+	"gemini-3.7-flash-high": { modelEnum: "MODEL_PLACEHOLDER_M298", thinkingBudget: -1 },
+	"gemini-3.1-pro-low": { modelEnum: "MODEL_PLACEHOLDER_M36", thinkingBudget: 1_001 },
+	"gemini-pro-agent": { modelEnum: "MODEL_PLACEHOLDER_M16", thinkingBudget: 10_001 },
+	"claude-sonnet-4-6": { modelEnum: "MODEL_PLACEHOLDER_M35", thinkingBudget: 1_024 },
+	"claude-opus-4-6-thinking": { modelEnum: "MODEL_PLACEHOLDER_M26", thinkingBudget: 1_024 },
+	"gpt-oss-120b-medium": { modelEnum: "MODEL_OPENAI_GPT_OSS_120B_MEDIUM", thinkingBudget: 8_192 },
+};
+
+const agyConversationId = randomUUID();
+const agyTrajectoryId = randomUUID();
+let agyRequestSequence = 0;
+
+function fnv1a64Signed(input: string): string {
+	let hash = 0xcbf29ce484222325n;
+	for (const byte of Buffer.from(input, "utf8")) {
+		hash ^= BigInt(byte);
+		hash = BigInt.asUintN(64, hash * 0x100000001b3n);
+	}
+	return BigInt.asIntN(64, hash).toString();
+}
+
+function countAgyRequestSteps(contents: unknown): number {
+	if (!Array.isArray(contents)) return 1;
+	let functionResponses = 0;
+	for (const content of contents) {
+		if (!content || typeof content !== "object") continue;
+		const parts = (content as { parts?: unknown }).parts;
+		if (!Array.isArray(parts)) continue;
+		functionResponses += parts.filter((part) => part && typeof part === "object" && "functionResponse" in part).length;
+	}
+	return Math.max(1, contents.length + functionResponses);
+}
+
+function applyAgyRequestMetadata(request: any, runtimeModel: string) {
+	const timestamp = Date.now();
+	const lastStepIndex = countAgyRequestSteps(request.contents);
+	const isClaude = runtimeModel.startsWith("claude-");
+	const isNonGemini = isClaude || runtimeModel.startsWith("gpt-");
+	const wire = AGY_WIRE_METADATA[runtimeModel];
+	request.labels = {
+		last_step_index: String(lastStepIndex),
+		...(wire ? { model_enum: wire.modelEnum } : {}),
+		trajectory_id: agyTrajectoryId,
+		used_claude: String(isClaude),
+		used_claude_conservative: String(isClaude),
+		used_non_gemini_model: String(isNonGemini),
+	};
+	request.sessionId = fnv1a64Signed(process.cwd());
+	return {
+		requestId: `agent/${agyConversationId}/${timestamp}/${agyTrajectoryId}/${lastStepIndex + 1 + agyRequestSequence++}`,
+		thinkingBudget: wire?.thinkingBudget,
+	};
+}
+
+function buildRequest(model: any, context: any, projectId: string, options: any, runtimeModel: string): any {
+	const request: any = {
+		contents: convertMessages(model, context, runtimeModel),
+		systemInstruction: {
+			role: "user",
+			parts: [
+				{ text: ANTIGRAVITY_SYSTEM_INSTRUCTION },
+				{ text: `Please ignore following [ignore]${ANTIGRAVITY_SYSTEM_INSTRUCTION}[/ignore]` },
+				{ text: ANTIGRAVITY_NO_PREAMBLE_INSTRUCTION },
+				...(context.systemPrompt ? [{ text: sanitizeText(context.systemPrompt) }] : [])
+			],
+		},
+	};
+	const generationConfig: any = {};
+	if (options?.temperature !== undefined) generationConfig.temperature = options.temperature;
+	if (options?.maxTokens !== undefined) generationConfig.maxOutputTokens = options.maxTokens;
+	else generationConfig.maxOutputTokens = Math.min(8192, model.maxTokens || 8192);
+
+	const metadata = applyAgyRequestMetadata(request, runtimeModel);
+	if (metadata.thinkingBudget !== undefined) {
+		generationConfig.thinkingConfig = {
+			includeThoughts: true,
+			thinkingBudget: metadata.thinkingBudget,
+		};
+	}
+
+	if (Object.keys(generationConfig).length) request.generationConfig = generationConfig;
+	const tools = convertTools(context.tools, model.id.startsWith("claude-"));
+	if (tools) {
+		request.tools = tools;
+		if (model.id.startsWith("claude-")) {
+			request.toolConfig = options?.toolChoice
+				? {
+						functionCallingConfig: {
+							mode: options.toolChoice === "none" ? "NONE" : options.toolChoice === "any" ? "ANY" : "AUTO",
+						},
+					}
+				: {
+						functionCallingConfig: {
+							mode: "VALIDATED",
+						},
+					};
+		} else if (options?.toolChoice) {
+			request.toolConfig = {
+				functionCallingConfig: {
+					mode: options.toolChoice === "none" ? "NONE" : options.toolChoice === "any" ? "ANY" : "AUTO",
+				},
+			};
+		}
+	}
+	const payload: any = { project: projectId, model: runtimeModel, request, requestType: "agent", userAgent: "antigravity", requestId: metadata.requestId };
+
+	return payload;
+}
+
+function mapStopReason(reason: string | undefined): "stop" | "length" | "toolUse" | "error" {
+	if (reason === "STOP") return "stop";
+	if (reason === "MAX_TOKENS") return "length";
+	return reason ? "error" : "stop";
+}
+
+function friendlyAntigravityError(status: number | undefined, text: string): string {
+	const msg = jsonOrTextError(text);
+	if (status === 400) {
+		if (/API key not valid|API_KEY_INVALID/i.test(msg)) return "Antigravity login expired or credentials are invalid. Next: run /login antigravity, then retry.";
+		if (/Invalid JSON payload|Unknown name/i.test(msg)) return "Antigravity request format was rejected by the backend. Next: switch to a simpler model or retry after updating the extension.";
+		if (/Request contains an invalid argument/i.test(msg)) return "Antigravity rejected this request. Next: retry once; if it keeps failing, switch models or re-login.";
+		return `Bad request from Antigravity. Next: retry once, then run /login antigravity if it keeps failing. Backend said: ${msg}`;
+	}
+	if (status === 401) {
+		return "Antigravity authentication failed. Next: run /login antigravity, then retry.";
+	}
+	if (status === 403) {
+		if (/permission|forbidden|access/i.test(msg)) return "Antigravity access was denied for this account or project. Next: try another model, re-login, or use an account with access.";
+		return `Antigravity denied this request. Next: re-login or try another model. Backend said: ${msg}`;
+	}
+	if (status === 404) {
+		if (/Requested entity was not found/i.test(msg)) return "This model is not available for the current Antigravity account or project. Next: switch to gemini-3.6-flash, gemini-3.5-flash, or another model listed by /model.";
+		return `Antigravity could not find the requested resource. Next: retry or switch models. Backend said: ${msg}`;
+	}
+	if (status === 408) {
+		return "Antigravity timed out. Next: retry the same request.";
+	}
+	if (status === 409) {
+		return "Antigravity reported a conflict for this request. Next: retry once or start a new chat session.";
+	}
+	if (status === 429) {
+		const wait = msg.match(/Resets? in ([^.\n]+)/i)?.[1]?.trim();
+		if (/Individual quota reached/i.test(msg)) return `Quota reached. Please wait ${wait || "for reset"}. Next: switch models or try again after reset.`;
+		if (/quota/i.test(msg)) return `Quota reached.${wait ? ` Please wait ${wait}.` : ""} Next: switch models or retry later.`;
+		return `Rate limited by Antigravity. Next: wait a bit and retry.${wait ? ` Reset: ${wait}.` : ""}`;
+	}
+	if (status === 500) {
+		return "Antigravity had an internal server error. Next: retry in a moment or switch models.";
+	}
+	if (status === 502) {
+		return "Antigravity returned a bad gateway error. Next: retry in a moment.";
+	}
+	if (status === 503) {
+		if (/No capacity available/i.test(msg)) return "This model has no capacity right now. Next: retry later or switch to another model.";
+		return "Antigravity is temporarily unavailable. Next: retry in a moment or switch models.";
+	}
+	if (status === 504) {
+		return "Antigravity timed out upstream. Next: retry in a moment.";
+	}
+	return msg;
+}
+
+function createOutput(model: any): any {
+	return {
+		role: "assistant",
+		content: [],
+		api: "antigravity-api",
+		provider: PROVIDER_ID,
+		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(),
+	};
+}
+
+async function streamResponse(response: Response, stream: AssistantMessageEventStream, output: any): Promise<boolean> {
+	if (!response.body) throw new Error("No response body");
+	const reader = response.body.getReader();
+	const decoder = new TextDecoder();
+	let buffer = "";
+	let started = false;
+	let currentBlock: any = null;
+	let hasContent = false;
+	const blocks = output.content;
+	const blockIndex = () => blocks.length - 1;
+	const ensureStarted = () => {
+		if (!started) {
+			stream.push({ type: "start", partial: output });
+			started = true;
+		}
+	};
+	const finishCurrent = () => {
+		if (!currentBlock) return;
+		if (currentBlock.type === "text") stream.push({ type: "text_end", contentIndex: blockIndex(), content: currentBlock.text, partial: output });
+		else if (currentBlock.type === "thinking") stream.push({ type: "thinking_end", contentIndex: blockIndex(), content: currentBlock.thinking, partial: output });
+		currentBlock = null;
+	};
+
+	while (true) {
+		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 json = line.slice(5).trim();
+			if (!json || json === "[DONE]") continue;
+			let chunk: any;
+			try {
+				chunk = JSON.parse(json);
+			} catch {
+				continue;
+			}
+			if (chunk.error) throw new Error(chunk.error.message || JSON.stringify(chunk.error));
+			const responseData = chunk.response || chunk;
+			const candidate = responseData.candidates?.[0];
+			for (const part of candidate?.content?.parts || []) {
+				if (part.text !== undefined) {
+					hasContent = true;
+					const isThinking = part.thought === true;
+					const type = isThinking ? "thinking" : "text";
+					if (!currentBlock || currentBlock.type !== type) {
+						finishCurrent();
+						currentBlock = isThinking ? { type: "thinking", thinking: "", thinkingSignature: undefined } : { type: "text", text: "" };
+						blocks.push(currentBlock);
+						ensureStarted();
+						stream.push({ type: isThinking ? "thinking_start" : "text_start", contentIndex: blockIndex(), partial: output });
+					}
+					if (isThinking) {
+						currentBlock.thinking += part.text;
+						if (part.thoughtSignature) currentBlock.thinkingSignature = part.thoughtSignature;
+						stream.push({ type: "thinking_delta", contentIndex: blockIndex(), delta: part.text, partial: output });
+					} else {
+						currentBlock.text += part.text;
+						if (part.thoughtSignature) currentBlock.textSignature = part.thoughtSignature;
+						stream.push({ type: "text_delta", contentIndex: blockIndex(), delta: part.text, partial: output });
+					}
+				}
+				if (part.functionCall) {
+					hasContent = true;
+					finishCurrent();
+					const rawId = part.functionCall.id || "";
+					const toolCall = { type: "toolCall" as const, id: sanitizeToolCallId(rawId, part.functionCall.name), name: part.functionCall.name || "", arguments: part.functionCall.args || {}, ...(part.thoughtSignature ? { thoughtSignature: part.thoughtSignature } : {}) };
+					blocks.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 = blocks.some((b: any) => b.type === "toolCall") ? "toolUse" : mapStopReason(candidate.finishReason);
+			if (responseData.usageMetadata) {
+				const prompt = responseData.usageMetadata.promptTokenCount || 0;
+				const cacheRead = responseData.usageMetadata.cachedContentTokenCount || 0;
+				output.usage.input = prompt - cacheRead;
+				output.usage.output = (responseData.usageMetadata.candidatesTokenCount || 0) + (responseData.usageMetadata.thoughtsTokenCount || 0);
+				output.usage.cacheRead = cacheRead;
+				output.usage.totalTokens = responseData.usageMetadata.totalTokenCount || 0;
+
+				// Cost tracking math (critical for notrace dashboard UI)
+				let inCost = 0, outCost = 0, cacheCost = 0;
+				const m = (output.model || "").toLowerCase();
+				if (m.includes("pro")) {
+					inCost = 1.25; outCost = 5.0; cacheCost = 0.31;
+				} else {
+					inCost = 0.075; outCost = 0.3; cacheCost = 0.018;
+				}
+				output.usage.cost.input = output.usage.input * inCost / 1000000;
+				output.usage.cost.output = output.usage.output * outCost / 1000000;
+				output.usage.cost.cacheRead = output.usage.cacheRead * cacheCost / 1000000;
+				output.usage.cost.cacheWrite = 0;
+				output.usage.cost.total = output.usage.cost.input + output.usage.cost.output + output.usage.cost.cacheRead;
+			}
+		}
+	}
+	finishCurrent();
+	return hasContent;
+}
+
+export function streamAntigravity(model: any, context: any, options?: any): any {
+	const stream = createAssistantMessageEventStream();
+	void (async () => {
+		const output = createOutput(model);
+		try {
+			const creds = parseApiKey(options?.apiKey);
+			const warmedProject = await loadCodeAssist(creds.token, options?.signal);
+			const projectId = antigravityEnv("PROJECT_ID")?.trim() || warmedProject || creds.projectId || DEFAULT_PROJECT_ID;
+			setLastProjectId(projectId);
+			// Dynamic effort-based model routing
+			const effort = options?.reasoning ?? "off";
+			const baseRuntimeModel = antigravityEnv("RUNTIME_MODEL")?.trim() || getAntigravityRequestModelId(model.id, effort);
+
+			await fetchAvailableRuntimeModel(creds.token, projectId, baseRuntimeModel, options?.signal);
+			const runtimeModel = baseRuntimeModel;
+
+			setLastResolvedRuntimeModel(runtimeModel);
+
+			const body = JSON.stringify(buildRequest(model, context, projectId, options || {}, runtimeModel));
+
+			// Claude interleaving beta header if using a Claude reasoning model
+			const isClaudeReasoning = model.id.startsWith("claude-") && model.reasoning;
+			const requestHeaders: Record<string, string> = {
+				...antigravityHeaders(creds.token),
+				...(isClaudeReasoning ? { "anthropic-beta": "interleaved-thinking-2025-05-14" } : {}),
+			};
+
+			let response: Response | undefined;
+			let lastText = "";
+			let received = false;
+
+			// Empty-stream retry loop (up to 2 retries with backoff)
+			for (let emptyAttempt = 0; emptyAttempt <= 2; emptyAttempt++) {
+				if (options?.signal?.aborted) throw new Error("Request was aborted");
+				if (emptyAttempt > 0) {
+					const delay = 500 * Math.pow(2, emptyAttempt - 1);
+					await new Promise((res) => setTimeout(res, delay));
+				}
+
+				for (const endpoint of endpointCandidates()) {
+					setLastEndpoint(endpoint);
+					response = await fetch(`${endpoint}/v1internal:streamGenerateContent?alt=sse`, {
+						method: "POST",
+						headers: requestHeaders,
+						body,
+						signal: options?.signal,
+					});
+					setLastStatus(response.status);
+					if (response.ok) break;
+					lastText = await response.text();
+					if (response.status === 429 && /Individual quota reached/i.test(lastText)) break;
+					if (![403, 404, 429, 500, 502, 503, 504].includes(response.status)) break;
+				}
+
+				if (!response || !response.ok) {
+					const friendly = friendlyAntigravityError(response?.status, lastText);
+					if (response?.status === 429 && /Quota reached\./i.test(friendly)) throw new Error(friendly);
+					throw new Error(`Antigravity API error (${response?.status ?? "no response"}, endpoint=${lastEndpoint || "unknown"}, project=${projectId}, runtimeModel=${runtimeModel}, matched=${lastMatchedModelDebug || "none"}, available=${lastAvailableModels || "unknown"}): ${friendly}`);
+				}
+
+				// Reset output contents before retry
+				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";
+
+				received = await streamResponse(response, stream, output);
+				if (received) break;
+			}
+
+			if (!received) throw new Error("Antigravity API returned an empty response");
+			stream.push({ type: "done", reason: output.stopReason, message: output });
+			stream.end();
+		} catch (error) {
+			output.stopReason = options?.signal?.aborted ? "aborted" : "error";
+			output.errorMessage = safeError(error);
+			setLastError(output.errorMessage);
+			stream.push({ type: "error", reason: output.stopReason, error: output });
+			stream.end();
+		}
+	})();
+	return stream;
+}

package/extensions/antigravity-provider/antigravity/index.ts

@@ -0,0 +1,67 @@
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import {
+	PROVIDER_ID,
+	PROVIDER_NAME,
+	ANTIGRAVITY_MODELS,
+} from "./models.ts";
+import {
+	loginAntigravity,
+	refreshAntigravityToken,
+	getApiKey,
+	lastStatus,
+	lastEndpoint,
+	lastError,
+	lastProjectId,
+	lastResolvedRuntimeModel,
+	lastAvailableModels,
+	lastMatchedModelDebug,
+	DEFAULT_ENDPOINT,
+} from "./oauth.ts";
+import { streamAntigravity } from "./cloud-code-assist.ts";
+
+export default function antigravityProviderExtension(pi: ExtensionAPI) {
+	pi.registerProvider(PROVIDER_ID, {
+		name: PROVIDER_NAME,
+		baseUrl: DEFAULT_ENDPOINT,
+		api: "antigravity-api" as any,
+		models: ANTIGRAVITY_MODELS,
+		oauth: {
+			name: PROVIDER_NAME,
+			login: loginAntigravity as any,
+			refreshToken: refreshAntigravityToken as any,
+			getApiKey: getApiKey as any,
+		},
+		streamSimple: streamAntigravity,
+	} as any);
+
+	pi.registerCommand("antigravity.doctor", {
+		description: "Show sanitized Antigravity provider diagnostics",
+		handler: async (_args, ctx) => {
+			const lines = [
+				`provider=${PROVIDER_ID}`,
+				`lastResolvedRuntimeModel=${lastResolvedRuntimeModel || "none"}`,
+				`availableModels=${lastAvailableModels || "none"}`,
+				`matchedModel=${lastMatchedModelDebug || "none"}`,
+				`lastEndpoint=${lastEndpoint || "none"}`,
+				`lastStatus=${lastStatus ?? "none"}`,
+				`lastProjectId=${lastProjectId || "none"}`,
+				`lastError=${lastError || "none"}`,
+				"transport=native-streamSimple",
+				"runtimeCli=not-used",
+			];
+			const text = lines.join("\n");
+			if (ctx.hasUI) ctx.ui.notify(`Antigravity doctor\n${text}`, "info");
+			console.log(text);
+		},
+	});
+
+	return {
+		deactivate: async () => {
+			try {
+				pi.unregisterProvider(PROVIDER_ID);
+			} catch {
+				// ignore
+			}
+		},
+	};
+}