diff --git a/bin/antigravity-bridge.mjs b/bin/antigravity-bridge.mjs new file mode 100755 index 0000000000..8e8a6c6c96 --- /dev/null +++ b/bin/antigravity-bridge.mjs @@ -0,0 +1,582 @@ +#!/usr/bin/env node +/** + * OmniRoute Antigravity Bridge Proxy + * + * Intercepts Antigravity CLI and IDE requests: + * - Directs Gemini 3.8 models directly to Google backend (100% native, untouched). + * - Directs other models (Claude Sonnet 4.5/4.6, Opus, Gemini 3.7, GPT-OSS, etc.) to OmniRoute /v1/antigravity. + * - Passes all non-model Google requests (auth, onboarding, telemetry) directly to Google backend. + * - Transparently forwards all other non-target internet traffic. + */ + +import net from "node:net"; +import http from "node:http"; +import https from "node:https"; +import tls from "node:tls"; +import fs from "node:fs"; +import path from "node:path"; +import { fileURLToPath } from "node:url"; + +const PORT = parseInt(process.env.BRIDGE_PORT || "20129", 10); +const ROUTER_URL = process.env.ROUTER_URL || "http://127.0.0.1:20128/v1/antigravity"; +const ROUTER_API_KEY = + process.env.ROUTER_API_KEY || process.env.OMNIROUTE_API_KEY || "sk-omniroute-bridge-local"; + +// Connection pool agents with TCP keep-alive +const httpAgent = new http.Agent({ + keepAlive: true, + keepAliveMsecs: 60000, + maxSockets: 64, + maxFreeSockets: 16, + timeout: 120000, +}); + +const httpsAgent = new https.Agent({ + keepAlive: true, + keepAliveMsecs: 60000, + maxSockets: 64, + maxFreeSockets: 16, + timeout: 120000, +}); + +let cachedSslOptions = null; +function getSslOptions() { + if (cachedSslOptions) return cachedSslOptions; + const certDir = + process.env.CERT_DIR || path.join(process.env.HOME || process.cwd(), ".omniroute", "mitm"); + const serverKey = path.join(certDir, "server.key"); + const serverCrt = path.join(certDir, "server.crt"); + + if (!fs.existsSync(serverKey) || !fs.existsSync(serverCrt)) { + console.error("❌ Certificate files not found in", certDir); + process.exit(1); + } + + cachedSslOptions = { + key: fs.readFileSync(serverKey), + cert: fs.readFileSync(serverCrt), + }; + return cachedSslOptions; +} + +const TARGET_HOSTS = new Set([ + "cloudcode-pa.googleapis.com", + "daily-cloudcode-pa.googleapis.com", + "daily-cloudcode-pa.sandbox.googleapis.com", + "autopush-cloudcode-pa.sandbox.googleapis.com", + "preprod-daily-cloudcode-pa.sandbox.googleapis.com", + "antigravity-unleash.goog", +]); + +function isGenerationRequest(url) { + if (!url) return false; + return ( + url.includes(":generateContent") || + url.includes(":streamGenerateContent") || + url.includes("/GenerateChat") || + url.includes("/StreamGenerateChat") || + url.includes("/GenerateCode") || + url.includes("/CompleteCode") + ); +} + +function extractModel(body, url) { + if (body && typeof body === "object") { + if (typeof body.model === "string" && body.model) return body.model; + if (body.request && typeof body.request.model === "string" && body.request.model) { + return body.request.model; + } + } + if (url) { + try { + const parsed = new URL(url, "https://cloudcode-pa.googleapis.com"); + const m = parsed.searchParams.get("model"); + if (m) return m; + } catch {} + } + return null; +} + +const MODEL_ROUTING_MAP = { + // Official OmniRoute Auto Groups + "auto/best-fast": "groq/openai/gpt-oss-120b", + "auto/best-coding": "mistral/codestral-latest", + "auto/best-reasoning": "nvidia/nvidia/nemotron-3-super-120b-a12b", + "auto/best-free": "groq/qwen/qwen3.8-27b", + "auto/best-vision": "nvidia/meta/llama-3.2-90b-vision-instruct", + "auto/coding:pro": "mistral/codestral-latest", + "auto/coding:fast": "groq/openai/gpt-oss-120b", + "auto/coding:free": "groq/qwen/qwen3.8-27b", + "auto/coding:reliable": "mistral/codestral-latest", + "auto/reasoning:pro": "nvidia/nvidia/nemotron-3-super-120b-a12b", + "auto/smart": "nvidia/nvidia/nemotron-3-super-120b-a12b", + "auto/claude-sonnet": "mistral/codestral-latest", + "auto/claude-opus": "nvidia/nvidia/nemotron-3-super-120b-a12b", + "auto/gemini": "gemini/gemini-2.5-flash", + "auto/llama": "groq/openai/gpt-oss-120b", + "auto/gemma": "groq/qwen/qwen3.8-27b", + + // Human-readable Display Names (in case CLI sends displayName in envelope) + "Auto: Best Fast (OmniRoute)": "groq/openai/gpt-oss-120b", + "Auto: Best Coding (OmniRoute)": "mistral/codestral-latest", + "Auto: Best Reasoning (OmniRoute)": "nvidia/nvidia/nemotron-3-super-120b-a12b", + "Auto: Best Free (OmniRoute)": "groq/qwen/qwen3.8-27b", + "Auto: Best Vision (OmniRoute)": "nvidia/meta/llama-3.2-90b-vision-instruct", + "Auto: Coding Pro (OmniRoute)": "mistral/codestral-latest", + "Auto: Coding Fast (OmniRoute)": "groq/openai/gpt-oss-120b", + "Auto: Coding Free (OmniRoute)": "groq/qwen/qwen3.8-27b", + "Auto: Coding Reliable (OmniRoute)": "mistral/codestral-latest", + "Auto: Reasoning Pro (OmniRoute)": "nvidia/nvidia/nemotron-3-super-120b-a12b", + "Auto: Smart (OmniRoute)": "nvidia/nvidia/nemotron-3-super-120b-a12b", + "Auto: Claude Sonnet (OmniRoute)": "mistral/codestral-latest", + "Auto: Claude Opus (OmniRoute)": "nvidia/nvidia/nemotron-3-super-120b-a12b", + "Auto: Gemini (OmniRoute)": "gemini/gemini-2.5-flash", + "Auto: Llama (OmniRoute)": "groq/openai/gpt-oss-120b", + "Auto: Gemma (OmniRoute)": "groq/qwen/qwen3.8-27b", + + // Fail-safe self-healing for dead/retired models + "nvidia/deepseek-ai/deepseek-v4-pro-0813": "groq/openai/gpt-oss-120b", + "deepseek-ai/deepseek-v4-pro-0813": "groq/openai/gpt-oss-120b", + "NVIDIA: DeepSeek V4 Pro": "groq/openai/gpt-oss-120b", + "nvidia/openai/gpt-oss-120b": "groq/openai/gpt-oss-120b", + "openai/gpt-oss-120b": "groq/openai/gpt-oss-120b", + "groq/llama-3.3-70b-versatile": "groq/openai/gpt-oss-120b", + "llama-3.3-70b-versatile": "groq/openai/gpt-oss-120b", +}; + +function resolveTargetModel(model) { + if (!model) return "groq/openai/gpt-oss-120b"; + if (MODEL_ROUTING_MAP[model]) return MODEL_ROUTING_MAP[model]; + const clean = model.replace(/^models\//, "").trim(); + if (MODEL_ROUTING_MAP[clean]) return MODEL_ROUTING_MAP[clean]; + for (const [k, v] of Object.entries(MODEL_ROUTING_MAP)) { + if (k.toLowerCase() === model.toLowerCase() || k.toLowerCase() === clean.toLowerCase()) { + return v; + } + } + if ( + clean.includes("deepseek-v4-pro") || + (clean.startsWith("nvidia") && clean.includes("gpt-oss-120b")) || + clean.includes("llama-3.3-70b-versatile") + ) { + return "groq/openai/gpt-oss-120b"; + } + return clean; +} + +const OMNIROUTE_BUILTIN_GROUPS = [ + { + id: "auto/best-coding", + displayName: "Auto: Best Coding (OmniRoute)", + descriptionText: + "OmniRoute dynamic routing to the highest benchmark coding model available (Mistral Codestral)", + }, + { + id: "auto/best-reasoning", + displayName: "Auto: Best Reasoning (OmniRoute)", + descriptionText: + "OmniRoute dynamic routing to the highest benchmark reasoning model available (Nemotron 3 Super 120B)", + }, + { + id: "auto/best-fast", + displayName: "Auto: Best Fast (OmniRoute)", + descriptionText: "OmniRoute sub-second lowest latency high-throughput model (Groq LPUs)", + }, + { + id: "auto/best-vision", + displayName: "Auto: Best Vision (OmniRoute)", + descriptionText: "OmniRoute multimodal & computer vision routing", + }, + { + id: "auto/best-free", + displayName: "Auto: Best Free (OmniRoute)", + descriptionText: "OmniRoute 100% unmetered free tier model routing (Qwen 3.8 27B)", + }, + { + id: "auto/coding:pro", + displayName: "Auto: Coding Pro (OmniRoute)", + descriptionText: "OmniRoute frontier pro-tier coding model (Codestral)", + }, + { + id: "auto/coding:fast", + displayName: "Auto: Coding Fast (OmniRoute)", + descriptionText: "OmniRoute fast sub-second daily coding model (Groq 120B)", + }, + { + id: "auto/coding:free", + displayName: "Auto: Coding Free (OmniRoute)", + descriptionText: "OmniRoute zero-cost free coding model", + }, + { + id: "auto/coding:reliable", + displayName: "Auto: Coding Reliable (OmniRoute)", + descriptionText: "OmniRoute maximum uptime and reliability coding model", + }, + { + id: "auto/reasoning:pro", + displayName: "Auto: Reasoning Pro (OmniRoute)", + descriptionText: "OmniRoute deep reasoning frontier model", + }, + { + id: "auto/smart", + displayName: "Auto: Smart (OmniRoute)", + descriptionText: "OmniRoute highest intelligence general-purpose model", + }, + { + id: "auto/claude-sonnet", + displayName: "Auto: Claude Sonnet (OmniRoute)", + descriptionText: "OmniRoute automated routing across Claude Sonnet providers", + }, + { + id: "auto/claude-opus", + displayName: "Auto: Claude Opus (OmniRoute)", + descriptionText: "OmniRoute automated routing across Claude Opus providers", + }, + { + id: "auto/gemini", + displayName: "Auto: Gemini (OmniRoute)", + descriptionText: "OmniRoute automated routing across Gemini providers", + }, + { + id: "auto/llama", + displayName: "Auto: Llama (OmniRoute)", + descriptionText: "OmniRoute automated routing across Llama providers", + }, + { + id: "auto/gemma", + displayName: "Auto: Gemma (OmniRoute)", + descriptionText: "OmniRoute automated routing across Gemma providers", + }, + // Active, verified provider models + { + id: "groq/openai/gpt-oss-120b", + displayName: "Groq: GPT-OSS 120B (Ultra-Fast 0.02s)", + descriptionText: "Ultra-fast inference on Groq LPUs at sub-second speeds", + }, + { + id: "groq/qwen/qwen3.8-27b", + displayName: "Groq: Qwen 3.8 27B", + descriptionText: "High-speed Qwen 3.8 27B model on Groq", + }, + { + id: "mistral/codestral-latest", + displayName: "Mistral: Codestral Latest", + descriptionText: "Mistral flagship frontier code reasoning model", + }, + { + id: "nvidia/nvidia/nemotron-3-super-120b-a12b", + displayName: "NVIDIA: Nemotron 3 Super 120B", + descriptionText: "Nemotron 3 Super 120B Deep Reasoning model on NVIDIA NIM", + }, + { + id: "gemini/gemini-2.5-flash", + displayName: "Gemini: Gemini 2.5 Flash (AI Studio)", + descriptionText: "Google AI Studio direct Gemini 2.5 Flash route", + }, + { + id: "gemini/gemini-2.5-pro", + displayName: "Gemini: Gemini 2.5 Pro (AI Studio)", + descriptionText: "Google AI Studio direct Gemini 2.5 Pro route", + }, +]; + +const OMNIROUTE_CUSTOM_MODELS = new Set([ + ...OMNIROUTE_BUILTIN_GROUPS.map((g) => g.id), + ...Object.keys(MODEL_ROUTING_MAP), +]); + +function shouldInterceptToOmniRoute(model, url) { + if (!model) return false; + + // Never intercept non-streaming unary RPCs (Antigravity expects raw JSON/Protobuf, not SSE) + const isStreaming = + url.includes("streamGenerateContent") || + url.includes("StreamGenerateChat") || + url.includes("alt=sse"); + if (!isStreaming) return false; + + // Never intercept native Google/Gemini models (used by Antigravity core, subagents, websearch, grounding) + if (model.startsWith("gemini-") || model.startsWith("models/gemini-")) { + return false; + } + + // Never intercept native Google CloudCode PA hosted models + if ( + model === "claude-sonnet-4-6" || + model === "claude-opus-4-6" || + model === "gpt-oss-120b-medium" + ) { + return false; + } + + // Intercept any OmniRoute auto group, provider model, or mapped alias + const clean = model.replace(/^models\//, "").trim(); + if ( + clean.startsWith("auto/") || + clean.toLowerCase().includes("omniroute") || + clean.includes("/") || + OMNIROUTE_CUSTOM_MODELS.has(model) || + OMNIROUTE_CUSTOM_MODELS.has(clean) || + Boolean(MODEL_ROUTING_MAP[model]) || + Boolean(MODEL_ROUTING_MAP[clean]) + ) { + return true; + } + + return false; +} + +const internalApp = http.createServer(async (req, res) => { + const host = (req.headers.host || "cloudcode-pa.googleapis.com").split(":")[0]; + const url = req.url || "/"; + + // Collect request body + const chunks = []; + for await (const chunk of req) { + chunks.push(chunk); + } + const bodyBuffer = Buffer.concat(chunks); + + let bodyJson = null; + if (bodyBuffer.length > 0) { + try { + bodyJson = JSON.parse(bodyBuffer.toString("utf-8")); + } catch {} + } + + const model = extractModel(bodyJson, url); + const shouldIntercept = shouldInterceptToOmniRoute(model, url); + + if (shouldIntercept) { + const resolvedModel = resolveTargetModel(model); + console.log( + `[Bridge] 🔀 INTERCEPTING -> OmniRoute: "${model || "default"}" => "${resolvedModel}" (${url})` + ); + + let outgoingBuffer = bodyBuffer; + if (bodyJson) { + const cloned = JSON.parse(JSON.stringify(bodyJson)); + cloned.model = resolvedModel; + if (cloned.request && typeof cloned.request === "object") { + cloned.request.model = resolvedModel; + } + outgoingBuffer = Buffer.from(JSON.stringify(cloned), "utf-8"); + } + + // Forward to OmniRoute /v1/antigravity + try { + const forwardHeaders = { + "Content-Type": "application/json", + "Content-Length": Buffer.byteLength(outgoingBuffer), + Authorization: `Bearer ${ROUTER_API_KEY}`, + "x-omniroute-source": "agent-bridge", + "x-omniroute-agent": "antigravity", + "x-omniroute-skip-usage": "true", // Skip usage tracking for default models + }; + + const upstreamReq = http.request( + ROUTER_URL, + { + method: "POST", + headers: forwardHeaders, + agent: httpAgent, + }, + (upstreamRes) => { + res.writeHead(upstreamRes.statusCode || 200, upstreamRes.headers); + upstreamRes.pipe(res); + } + ); + upstreamReq.setNoDelay(true); + + upstreamReq.on("error", (err) => { + console.error(`[Bridge] ❌ Error forwarding to OmniRoute: ${err.message}`); + if (!res.headersSent) { + res.writeHead(502, { "Content-Type": "application/json" }); + res.end(JSON.stringify({ error: { message: `OmniRoute bridge error: ${err.message}` } })); + } + }); + + upstreamReq.write(outgoingBuffer); + upstreamReq.end(); + return; + } catch (err) { + console.error(`[Bridge] ❌ Failed to invoke OmniRoute: ${err.message}`); + } + } + + // Otherwise: Passthrough directly to Google upstream + console.log(`[Bridge] ⏩ PASSTHROUGH -> Google: ${model || "non-model"} (${url})`); + + const upstreamHeaders = { ...req.headers }; + delete upstreamHeaders["host"]; // Let https.request set the correct Host + upstreamHeaders["host"] = host; + + if (url.includes("fetchAvailableModels")) { + delete upstreamHeaders["accept-encoding"]; + } + + const googleReq = https.request( + { + hostname: host, + port: 443, + path: url, + method: req.method, + headers: upstreamHeaders, + agent: httpsAgent, + }, + (googleRes) => { + if (url.includes("fetchAvailableModels")) { + const respChunks = []; + googleRes.on("data", (chunk) => respChunks.push(chunk)); + googleRes.on("end", () => { + const respBuffer = Buffer.concat(respChunks); + let finalBuffer = respBuffer; + try { + const data = JSON.parse(respBuffer.toString("utf-8")); + if (data && data.models) { + // Inject OmniRoute built-in auto groups and models + const baseTemplate = + data.models["claude-sonnet-4-6"] || + data.models["gpt-oss-120b-medium"] || + Object.values(data.models)[0] || + {}; + + const injectedIds = []; + for (const group of OMNIROUTE_BUILTIN_GROUPS) { + data.models[group.id] = { + ...baseTemplate, + id: group.id, + name: group.id, + displayName: group.displayName, + descriptionText: group.descriptionText, + }; + injectedIds.push(group.id); + } + + // Prepend OmniRoute groups to agentModelSorts recommended group + if ( + Array.isArray(data.agentModelSorts) && + data.agentModelSorts[0]?.groups?.[0]?.modelIds + ) { + const existing = data.agentModelSorts[0].groups[0].modelIds; + data.agentModelSorts[0].groups[0].modelIds = [ + ...injectedIds, + ...existing.filter((id) => !injectedIds.includes(id)), + ]; + } + finalBuffer = Buffer.from(JSON.stringify(data), "utf-8"); + console.log( + `[Bridge] 🌟 Injected custom models into fetchAvailableModels (${finalBuffer.length} bytes)` + ); + } + } catch (err) { + console.error(`[Bridge] ⚠️ Error modifying fetchAvailableModels: ${err.message}`); + } + + const headers = { ...googleRes.headers }; + delete headers["content-length"]; + delete headers["content-encoding"]; + headers["content-length"] = String(finalBuffer.length); + res.writeHead(googleRes.statusCode || 200, headers); + res.end(finalBuffer); + }); + return; + } + + res.writeHead(googleRes.statusCode || 200, googleRes.headers); + googleRes.pipe(res); + } + ); + googleReq.setNoDelay(true); + + googleReq.on("error", (err) => { + console.error(`[Bridge] ❌ Google upstream error: ${err.message}`); + if (!res.headersSent) { + res.writeHead(502, { "Content-Type": "application/json" }); + res.end(JSON.stringify({ error: { message: `Google upstream error: ${err.message}` } })); + } + }); + + if (bodyBuffer.length > 0) { + googleReq.write(bodyBuffer); + } + googleReq.end(); +}); + +internalApp.keepAliveTimeout = 65000; +internalApp.headersTimeout = 66000; + +// Proxy server listening on HTTP port +const proxyServer = http.createServer((req, res) => { + // Plain HTTP request (non-CONNECT) + res.writeHead(200, { "Content-Type": "text/plain" }); + res.end("OmniRoute Antigravity Bridge Proxy Active\n"); +}); + +proxyServer.keepAliveTimeout = 65000; +proxyServer.headersTimeout = 66000; + +proxyServer.on("connect", (req, clientSocket, head) => { + clientSocket.setNoDelay(true); + const [targetHost, targetPortStr] = (req.url || "").split(":"); + const targetPort = parseInt(targetPortStr || "443", 10); + + if (TARGET_HOSTS.has(targetHost)) { + // Target host: Terminate TLS locally and route via internalApp + clientSocket.write("HTTP/1.1 200 Connection Established\r\n\r\n"); + + const ssl = getSslOptions(); + const tlsSocket = new tls.TLSSocket(clientSocket, { + isServer: true, + key: ssl.key, + cert: ssl.cert, + }); + tlsSocket.setNoDelay(true); + + tlsSocket.on("error", (err) => { + // Client closed or TLS error + clientSocket.destroy(); + }); + + internalApp.emit("connection", tlsSocket); + } else { + // Non-target host: Transparent raw TCP tunnel + const upstreamSocket = net.connect(targetPort, targetHost, () => { + upstreamSocket.setNoDelay(true); + clientSocket.write("HTTP/1.1 200 Connection Established\r\n\r\n"); + if (head && head.length > 0) { + upstreamSocket.write(head); + } + upstreamSocket.pipe(clientSocket); + clientSocket.pipe(upstreamSocket); + }); + + const cleanup = () => { + clientSocket.destroy(); + upstreamSocket.destroy(); + }; + + upstreamSocket.on("error", cleanup); + clientSocket.on("error", cleanup); + } +}); + +export { + resolveTargetModel, + MODEL_ROUTING_MAP, + shouldInterceptToOmniRoute, + extractModel, + OMNIROUTE_BUILTIN_GROUPS, + proxyServer, + internalApp, +}; + +const isMain = process.argv[1] && path.resolve(process.argv[1]) === fileURLToPath(import.meta.url); + +if (isMain) { + proxyServer.listen(PORT, "127.0.0.1", () => { + console.log(`🚀 OmniRoute Antigravity Bridge listening on 127.0.0.1:${PORT}`); + console.log(` Routing non-Gemini 3.8 model traffic -> ${ROUTER_URL}`); + console.log(` Preserving Gemini 3.8 native traffic -> Google`); + }); +} diff --git a/open-sse/services/emergencyFallback.ts b/open-sse/services/emergencyFallback.ts index 60e3914172..fb2ecef11f 100644 --- a/open-sse/services/emergencyFallback.ts +++ b/open-sse/services/emergencyFallback.ts @@ -36,7 +36,7 @@ export interface EmergencyFallbackConfig { export const EMERGENCY_FALLBACK_CONFIG: EmergencyFallbackConfig = { enabled: true, - provider: "nvidia", + provider: "groq", model: "openai/gpt-oss-120b", triggerOn402: true, triggerOnBudgetKeywords: true, diff --git a/src/mitm/handlers/antigravity.ts b/src/mitm/handlers/antigravity.ts index 4efdd41e38..b8d96ca3a8 100644 --- a/src/mitm/handlers/antigravity.ts +++ b/src/mitm/handlers/antigravity.ts @@ -21,6 +21,7 @@ */ import type { IncomingMessage, ServerResponse } from "node:http"; import type { AgentId } from "../types"; +import type { InterceptedRequest } from "../inspector/types"; import { MitmHandlerBase, createBoundedCollector } from "./base"; import { TOOL_RENAME_MAP } from "@omniroute/open-sse/services/claudeCodeToolRemapper"; @@ -108,7 +109,7 @@ function joinPartsText(parts: GeminiPart[] | undefined): string { export function convertGeminiToOpenAI( geminiBody: GeminiRequestBody, model: string, - stream: boolean, + stream: boolean ): OpenAIChatBody { // Unwrap the cloudcode-pa envelope (`.request`) used by the real Antigravity IDE; fall // back to the top level for the legacy `/v1beta` shape. (#4294) @@ -143,23 +144,176 @@ export function convertGeminiToOpenAI( return openaiBody; } +export interface DynamicCatalogModel { + id: string; + displayName?: string; + description?: string; +} + +/** + * Merge operator's configured dynamic models/combos into Google Antigravity's + * fetchAvailableModels response. + * + * @param catalog Raw upstream catalog response from Google Cloud Code PA + * @param dynamicModels Operator's configured combos/models to inject + */ +export function mergeAntigravityCatalog( + catalog: Record, + dynamicModels: DynamicCatalogModel[] +): Record { + if (!dynamicModels || dynamicModels.length === 0) { + return catalog; + } + + const result = { ...catalog }; + const injectedIds: string[] = []; + + if (Array.isArray(result.models)) { + const modelsArr = [...(result.models as Array>)]; + const templateModel = (modelsArr[0] as Record) || {}; + for (const m of dynamicModels) { + if (!m.id) continue; + // Collision guard: do not overwrite an existing upstream native model + if (modelsArr.some((existing) => existing.id === m.id || existing.name === m.id)) { + continue; + } + injectedIds.push(m.id); + modelsArr.push({ + ...templateModel, + id: m.id, + name: m.id, + displayName: m.displayName || m.id, + descriptionText: m.description || `OmniRoute dynamic model (${m.id})`, + }); + } + result.models = modelsArr; + } else { + const modelsObj = ( + result.models && typeof result.models === "object" + ? { ...(result.models as Record) } + : {} + ) as Record>; + + const templateModel = + modelsObj["claude-sonnet-4-6"] || + modelsObj["gemini-2.5-pro"] || + modelsObj["gemini-3.7-flash-medium"] || + Object.values(modelsObj)[0] || + {}; + + for (const m of dynamicModels) { + if (!m.id) continue; + // Collision guard: do not overwrite an existing upstream native model entry + if (modelsObj[m.id]) { + continue; + } + injectedIds.push(m.id); + modelsObj[m.id] = { + ...templateModel, + ...(typeof templateModel.id === "string" ? { id: m.id } : {}), + ...(typeof templateModel.name === "string" ? { name: m.id } : {}), + displayName: m.displayName || m.id, + descriptionText: m.description || `OmniRoute dynamic model (${m.id})`, + }; + } + result.models = modelsObj; + } + + // Prepend custom models to agentModelSorts recommended group + let sorts = Array.isArray(result.agentModelSorts) + ? [...(result.agentModelSorts as Array>)] + : []; + + if (sorts.length === 0) { + sorts = [{ groups: [{ modelIds: [] }] }]; + } + + const firstSort = { ...(sorts[0] as Record) }; + let groups = Array.isArray(firstSort.groups) + ? [...(firstSort.groups as Array>)] + : []; + if (groups.length === 0) { + groups = [{ modelIds: [] }]; + } + + const firstGroup = { ...(groups[0] as Record) }; + const existingModelIds = Array.isArray(firstGroup.modelIds) + ? (firstGroup.modelIds as string[]) + : []; + + firstGroup.modelIds = [ + ...injectedIds, + ...existingModelIds.filter((id) => !injectedIds.includes(id)), + ]; + + groups[0] = firstGroup; + firstSort.groups = groups; + sorts[0] = firstSort; + result.agentModelSorts = sorts; + + return result; +} + export class AntigravityHandler extends MitmHandlerBase { readonly agentId: AgentId = "antigravity"; + private customCatalogModels?: DynamicCatalogModel[]; + + constructor(customCatalogModels?: DynamicCatalogModel[]) { + super(); + this.customCatalogModels = customCatalogModels; + } + + /** + * Dynamically retrieve the operator's configured models / combos from the + * database repository or injected test config. + */ + async getDynamicCatalogModels(): Promise { + if (this.customCatalogModels) { + return this.customCatalogModels; + } + try { + const combosMod = await import("@/lib/db/combos").catch( + () => import("../../lib/db/combos.ts") + ); + if (typeof combosMod?.getCombos === "function") { + const combos = await combosMod.getCombos(); + if (Array.isArray(combos)) { + return combos + .filter((c: Record) => c.isActive !== false && !c.isHidden) + .map((c: Record) => { + const name = typeof c.name === "string" ? c.name.trim() : ""; + const desc = typeof c.description === "string" ? c.description.trim() : undefined; + return name ? { id: name, displayName: name, description: desc } : null; + }) + .filter((c): c is DynamicCatalogModel => Boolean(c)); + } + } + } catch { + // Ignored: return empty if DB unavailable + } + return []; + } async intercept( req: IncomingMessage, res: ServerResponse, body: Buffer, - mappedModel: string, + mappedModel: string ): Promise { const startedAt = this.now(); const intercepted = await this.hookBufferStart(req, body, mappedModel); try { + const url = req.url || ""; + if (url.includes(":fetchAvailableModels")) { + await this.interceptFetchAvailableModels(req, res, body, intercepted, startedAt); + return; + } + const geminiBody = JSON.parse(body.toString()) as GeminiRequestBody; // Streaming intent: Antigravity uses :streamGenerateContent for streaming. - const isStream = (req.url || "").includes(":streamGenerateContent"); + const isStream = url.includes(":streamGenerateContent"); const payload = convertGeminiToOpenAI(geminiBody, mappedModel, isStream); @@ -197,4 +351,71 @@ export class AntigravityHandler extends MitmHandlerBase { await this.writeError(res, err); } } + + private async interceptFetchAvailableModels( + req: IncomingMessage, + res: ServerResponse, + body: Buffer, + intercepted: InterceptedRequest, + startedAt: number + ): Promise { + const host = + (typeof req.headers.host === "string" && req.headers.host) || "cloudcode-pa.googleapis.com"; + const upstreamUrl = `https://${host}${req.url || "/v1internal:fetchAvailableModels"}`; + + const upstreamHeaders: Record = {}; + for (const [key, value] of Object.entries(req.headers)) { + if (value === undefined) continue; + const lower = key.toLowerCase(); + if ( + lower === "host" || + lower === "connection" || + lower === "content-length" || + lower === "accept-encoding" + ) { + continue; + } + upstreamHeaders[lower] = Array.isArray(value) ? value.join(", ") : value; + } + if (!upstreamHeaders["content-type"]) { + upstreamHeaders["content-type"] = "application/json"; + } + + const upstreamStart = this.now(); + const upstream = await fetch(upstreamUrl, { + method: req.method || "POST", + headers: upstreamHeaders, + body: body && body.length > 0 ? body.toString() : JSON.stringify({}), + }); + + if (!upstream.ok) { + const errText = await upstream.text().catch(() => ""); + throw new Error(`Google upstream ${upstream.status}: ${errText}`); + } + + const rawCatalog = (await upstream.json()) as Record; + const dynamicModels = await this.getDynamicCatalogModels(); + const merged = mergeAntigravityCatalog(rawCatalog, dynamicModels); + + const respText = JSON.stringify(merged); + const responseHeaders: Record = { + "content-type": "application/json; charset=utf-8", + "content-length": String(Buffer.byteLength(respText)), + }; + + if (!res.headersSent) { + res.writeHead(upstream.status, responseHeaders); + } + res.end(respText); + + const total = this.now() - startedAt; + this.hookBufferUpdate(intercepted, { + status: upstream.status, + responseHeaders, + responseBody: respText, + responseSize: Buffer.byteLength(respText), + proxyLatencyMs: upstreamStart - startedAt, + upstreamLatencyMs: total - (upstreamStart - startedAt), + }); + } } diff --git a/src/mitm/targets/antigravity.ts b/src/mitm/targets/antigravity.ts index 5d1e0cc9b8..45b5a01836 100644 --- a/src/mitm/targets/antigravity.ts +++ b/src/mitm/targets/antigravity.ts @@ -22,6 +22,7 @@ const ENDPOINTS = [ "/v1internal:streamGenerateContent", "/v1internal:loadCodeAssist", "/v1internal:onboardUser", + "/v1internal:fetchAvailableModels", ]; const INSTRUCTIONS = [ diff --git a/tests/unit/antigravity-bridge-routing.test.ts b/tests/unit/antigravity-bridge-routing.test.ts new file mode 100644 index 0000000000..20c6ea8b17 --- /dev/null +++ b/tests/unit/antigravity-bridge-routing.test.ts @@ -0,0 +1,114 @@ +import assert from "node:assert/strict"; +import test from "node:test"; + +import { + resolveTargetModel, + MODEL_ROUTING_MAP, + shouldInterceptToOmniRoute, + extractModel, + OMNIROUTE_BUILTIN_GROUPS, +} from "../../bin/antigravity-bridge.mjs"; + +test("MODEL_ROUTING_MAP contains mappings for core auto groups", () => { + assert.ok(MODEL_ROUTING_MAP["auto/best-fast"]); + assert.ok(MODEL_ROUTING_MAP["auto/best-coding"]); +}); + +test("shouldInterceptToOmniRoute preserves native Gemini models for zero Google quota leakage", () => { + const streamingUrl = + "https://cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse"; + + // Native Google Gemini models must NEVER be intercepted + assert.equal(shouldInterceptToOmniRoute("gemini-3.8-pro", streamingUrl), false); + assert.equal(shouldInterceptToOmniRoute("gemini-3.8-flash", streamingUrl), false); + assert.equal(shouldInterceptToOmniRoute("gemini-3.5-flash-lite", streamingUrl), false); + assert.equal(shouldInterceptToOmniRoute("models/gemini-2.5-pro", streamingUrl), false); + + // Native hosted models must never be intercepted + assert.equal(shouldInterceptToOmniRoute("claude-sonnet-4-6", streamingUrl), false); + assert.equal(shouldInterceptToOmniRoute("claude-opus-4-6", streamingUrl), false); + assert.equal(shouldInterceptToOmniRoute("gpt-oss-120b-medium", streamingUrl), false); +}); + +test("shouldInterceptToOmniRoute ignores non-streaming RPCs", () => { + assert.equal( + shouldInterceptToOmniRoute( + "auto/best-fast", + "https://cloudcode-pa.googleapis.com/v1internal:fetchUserInfo" + ), + false + ); + assert.equal( + shouldInterceptToOmniRoute( + "auto/best-fast", + "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist" + ), + false + ); +}); + +test("shouldInterceptToOmniRoute intercepts all OmniRoute auto groups and display names", () => { + const streamingUrl = + "https://cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse"; + + assert.equal(shouldInterceptToOmniRoute("auto/best-fast", streamingUrl), true); + assert.equal(shouldInterceptToOmniRoute("auto/best-coding", streamingUrl), true); + assert.equal(shouldInterceptToOmniRoute("auto/best-reasoning", streamingUrl), true); + assert.equal(shouldInterceptToOmniRoute("auto/best-free", streamingUrl), true); + assert.equal(shouldInterceptToOmniRoute("Auto: Best Fast (OmniRoute)", streamingUrl), true); + assert.equal(shouldInterceptToOmniRoute("Auto: Best Coding (OmniRoute)", streamingUrl), true); + assert.equal(shouldInterceptToOmniRoute("groq/openai/gpt-oss-120b", streamingUrl), true); + assert.equal(shouldInterceptToOmniRoute("mistral/codestral-latest", streamingUrl), true); +}); + +test("resolveTargetModel correctly maps OmniRoute auto groups to active providers", () => { + assert.equal(resolveTargetModel("auto/best-fast"), "groq/openai/gpt-oss-120b"); + assert.equal(resolveTargetModel("Auto: Best Fast (OmniRoute)"), "groq/openai/gpt-oss-120b"); + assert.equal(resolveTargetModel("auto/best-coding"), "mistral/codestral-latest"); + assert.equal(resolveTargetModel("Auto: Best Coding (OmniRoute)"), "mistral/codestral-latest"); + assert.equal( + resolveTargetModel("auto/best-reasoning"), + "nvidia/nvidia/nemotron-3-super-120b-a12b" + ); + assert.equal(resolveTargetModel("auto/best-free"), "groq/qwen/qwen3.8-27b"); +}); + +test("resolveTargetModel self-heals retired models and prevents upstream 410 crashes", () => { + // Deprecated/retired on NVIDIA NIM + assert.equal( + resolveTargetModel("nvidia/deepseek-ai/deepseek-v4-pro-0813"), + "groq/openai/gpt-oss-120b" + ); + assert.equal(resolveTargetModel("deepseek-ai/deepseek-v4-pro-0813"), "groq/openai/gpt-oss-120b"); + assert.equal(resolveTargetModel("NVIDIA: DeepSeek V4 Pro"), "groq/openai/gpt-oss-120b"); + assert.equal(resolveTargetModel("nvidia/openai/gpt-oss-120b"), "groq/openai/gpt-oss-120b"); +}); + +test("extractModel resolves models from envelope body and query parameters", () => { + assert.equal(extractModel({ model: "auto/best-fast" }, ""), "auto/best-fast"); + assert.equal( + extractModel({ request: { model: "Auto: Best Coding (OmniRoute)" } }, ""), + "Auto: Best Coding (OmniRoute)" + ); + assert.equal( + extractModel( + null, + "https://cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?model=auto/best-reasoning" + ), + "auto/best-reasoning" + ); +}); + +test("OMNIROUTE_BUILTIN_GROUPS provides unique IDs and valid display names", () => { + const ids = new Set(OMNIROUTE_BUILTIN_GROUPS.map((g) => g.id)); + assert.equal(ids.size, OMNIROUTE_BUILTIN_GROUPS.length); + assert.ok(ids.has("auto/best-fast")); + assert.ok(ids.has("auto/best-coding")); + assert.ok(ids.has("auto/best-reasoning")); + assert.ok(ids.has("auto/best-free")); + + for (const group of OMNIROUTE_BUILTIN_GROUPS) { + assert.ok(group.displayName.length > 0); + assert.ok(group.descriptionText.length > 0); + } +}); diff --git a/tests/unit/mitm-handler-antigravity.test.ts b/tests/unit/mitm-handler-antigravity.test.ts index ec3265d365..e997af1e8a 100644 --- a/tests/unit/mitm-handler-antigravity.test.ts +++ b/tests/unit/mitm-handler-antigravity.test.ts @@ -1,11 +1,32 @@ import test from "node:test"; import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; import { AntigravityHandler, convertGeminiToOpenAI, + mergeAntigravityCatalog, } from "../../src/mitm/handlers/antigravity.ts"; +import { ANTIGRAVITY_TARGET } from "../../src/mitm/targets/antigravity.ts"; import { runHandler } from "./_mitmHandlerHarness.ts"; +// The last test below imports src/lib/db/combos.ts, which opens the real DATA_DIR +// SQLite database. Isolate it BEFORE that import runs so `node --import tsx/esm +// tests/unit/mitm-handler-antigravity.test.ts` (this file's own documented +// Validation command) never touches the operator's real home database. +const previousDataDir = process.env.DATA_DIR; +const isolatedDataDir = fs.mkdtempSync(path.join(os.tmpdir(), "mitm-antigravity-")); +process.env.DATA_DIR = isolatedDataDir; + +test.after(async () => { + const { resetDbInstance } = await import("../../src/lib/db/core.ts"); + resetDbInstance(); + if (previousDataDir === undefined) delete process.env.DATA_DIR; + else process.env.DATA_DIR = previousDataDir; + fs.rmSync(isolatedDataDir, { recursive: true, force: true }); +}); + test("antigravity handler — forwards to OmniRoute and pipes SSE", async () => { const r = await runHandler( new AntigravityHandler(), @@ -19,12 +40,10 @@ test("antigravity handler — forwards to OmniRoute and pipes SSE", async () => }); test("antigravity handler — propagates upstream failure as 500", async () => { - const r = await runHandler( - new AntigravityHandler(), - { model: "gpt-4o" }, - "claude-3.5-sonnet", - { upstreamStatus: 500, upstreamBody: "boom" } - ); + const r = await runHandler(new AntigravityHandler(), { model: "gpt-4o" }, "claude-3.5-sonnet", { + upstreamStatus: 500, + upstreamBody: "boom", + }); assert.equal(r.status, 500); const body = r.responseChunks.join(""); // Error must NOT include raw stack trace (Hard Rule #12 sanitization). @@ -169,3 +188,269 @@ test("antigravity handler — non-streaming URL yields stream:false", async () = const forwarded = JSON.parse(r.fetchBody); assert.equal(forwarded.stream, false); }); + +test("ANTIGRAVITY_TARGET — includes fetchAvailableModels in endpointPatterns", () => { + assert.ok( + ANTIGRAVITY_TARGET.endpointPatterns.includes("/v1internal:fetchAvailableModels"), + "ANTIGRAVITY_TARGET must declare /v1internal:fetchAvailableModels in endpointPatterns" + ); +}); + +test("mergeAntigravityCatalog — merges dynamic models and prepends to agentModelSorts", () => { + const upstreamCatalog = { + models: { + "claude-sonnet-4-6": { + displayName: "Claude 3.7 Sonnet", + descriptionText: "Anthropic Claude 3.7 Sonnet", + quotaInfo: { remainingFraction: 1.0, resetTime: "2026-09-18T00:00:00Z" }, + }, + "gemini-2.5-pro": { + displayName: "Gemini 2.5 Pro", + descriptionText: "Google Gemini 2.5 Pro", + quotaInfo: { remainingFraction: 0.9, resetTime: "2026-09-18T00:00:00Z" }, + }, + }, + agentModelSorts: [ + { + groups: [ + { + modelIds: ["gemini-2.5-pro", "claude-sonnet-4-6"], + }, + ], + }, + ], + }; + + const dynamicModels = [ + { + id: "coding-titans", + displayName: "Coding Titans", + description: "Deep Architecture & Complex Logic", + }, + { + id: "speed-demons", + displayName: "Speed Demons", + description: "Sub-second Daily Coding", + }, + ]; + + const merged = mergeAntigravityCatalog(upstreamCatalog, dynamicModels); + const models = merged.models as Record>; + + // Injected models must exist and have displayName/descriptionText + assert.ok(models["coding-titans"]); + assert.equal(models["coding-titans"].displayName, "Coding Titans"); + assert.equal(models["coding-titans"].descriptionText, "Deep Architecture & Complex Logic"); + // Template properties (like quotaInfo) must be cloned + assert.deepEqual(models["coding-titans"].quotaInfo, { + remainingFraction: 1.0, + resetTime: "2026-09-18T00:00:00Z", + }); + + assert.ok(models["speed-demons"]); + assert.equal(models["speed-demons"].displayName, "Speed Demons"); + + // Native upstream models must be preserved + assert.ok(models["claude-sonnet-4-6"]); + assert.ok(models["gemini-2.5-pro"]); + + // Injected models must be prepended to agentModelSorts + const sorts = merged.agentModelSorts as Array<{ groups: Array<{ modelIds: string[] }> }>; + assert.deepEqual(sorts[0].groups[0].modelIds, [ + "coding-titans", + "speed-demons", + "gemini-2.5-pro", + "claude-sonnet-4-6", + ]); +}); + +test("mergeAntigravityCatalog — handles empty dynamicModels by returning catalog untouched", () => { + const upstreamCatalog = { + models: { "gemini-2.5-flash": { displayName: "Gemini 2.5 Flash" } }, + }; + const result = mergeAntigravityCatalog(upstreamCatalog, []); + assert.equal(result, upstreamCatalog); +}); + +test("mergeAntigravityCatalog — handles missing agentModelSorts gracefully", () => { + const upstreamCatalog = { + models: { + "gemini-2.5-flash": { displayName: "Gemini 2.5 Flash" }, + }, + }; + const dynamicModels = [{ id: "custom-combo", displayName: "Custom Combo" }]; + const merged = mergeAntigravityCatalog(upstreamCatalog, dynamicModels); + const sorts = merged.agentModelSorts as Array<{ groups: Array<{ modelIds: string[] }> }>; + assert.ok(Array.isArray(sorts)); + assert.deepEqual(sorts[0].groups[0].modelIds, ["custom-combo"]); +}); + +test("antigravity handler — intercepts fetchAvailableModels and returns merged catalog", async () => { + const upstreamCatalog = { + models: { + "claude-sonnet-4-6": { + displayName: "Claude 3.7 Sonnet", + descriptionText: "Anthropic Claude 3.7 Sonnet", + }, + "gemini-2.5-pro": { + displayName: "Gemini 2.5 Pro", + descriptionText: "Google Gemini 2.5 Pro", + }, + }, + agentModelSorts: [ + { + groups: [ + { + modelIds: ["claude-sonnet-4-6", "gemini-2.5-pro"], + }, + ], + }, + ], + }; + + const dynamicModels = [ + { + id: "coding-titans", + displayName: "Coding Titans", + description: "Deep Architecture", + }, + ]; + + const handler = new AntigravityHandler(dynamicModels); + + const r = await runHandler(handler, {}, "ag-claude-opus-4-6-thinking", { + url: "/v1internal:fetchAvailableModels", + upstreamBody: JSON.stringify(upstreamCatalog), + }); + + assert.ok(r.fetchCalled); + assert.equal(r.status, 200); + assert.ok(r.fetchUrl?.includes("/v1internal:fetchAvailableModels")); + + const responseJson = JSON.parse(r.responseChunks.join("")); + assert.ok(responseJson.models["coding-titans"]); + assert.equal(responseJson.models["coding-titans"].displayName, "Coding Titans"); + assert.ok(responseJson.models["claude-sonnet-4-6"]); + assert.equal(responseJson.agentModelSorts[0].groups[0].modelIds[0], "coding-titans"); +}); + +test("antigravity handler — propagates upstream error on fetchAvailableModels", async () => { + const handler = new AntigravityHandler(); + + const r = await runHandler(handler, {}, "ag-claude-opus-4-6-thinking", { + url: "/v1internal:fetchAvailableModels", + upstreamStatus: 502, + upstreamBody: JSON.stringify({ error: "bad gateway" }), + }); + + assert.equal(r.status, 500); + const body = r.responseChunks.join(""); + assert.ok(body.includes("mitm_error")); + assert.ok(!body.includes("at /")); +}); + +test("antigravity handler — dynamic catalog pulls configured combos from database repository", async () => { + const { createCombo, deleteCombo } = await import("../../src/lib/db/combos.ts"); + const testComboName = `test-combo-${Date.now()}`; + await createCombo({ + id: testComboName, + name: testComboName, + description: "Test dynamic combo description", + models: JSON.stringify(["google/gemini-2.5-flash"]), + strategy: "priority", + }); + + try { + const handler = new AntigravityHandler(); + const models = await handler.getDynamicCatalogModels(); + const found = models.find((m) => m.id === testComboName); + assert.ok(found, `Expected ${testComboName} to be returned from dynamic catalog models`); + assert.equal(found?.displayName, testComboName); + assert.equal(found?.description, "Test dynamic combo description"); + } finally { + await deleteCombo(testComboName); + } +}); + +test("mergeAntigravityCatalog — preserves native model metadata and avoids collision when combo shares bare model id", () => { + const upstreamCatalog = { + models: { + "gemini-2.5-pro": { + displayName: "Native Gemini 2.5 Pro", + descriptionText: "Google Official", + isNative: true, + }, + }, + agentModelSorts: [ + { + groups: [ + { + modelIds: ["gemini-2.5-pro"], + }, + ], + }, + ], + }; + + const dynamicModels = [ + { + id: "gemini-2.5-pro", + displayName: "Overwriting Combo", + description: "Should not overwrite native", + }, + { + id: "coding-titans", + displayName: "Coding Titans", + description: "Non-colliding combo", + }, + ]; + + const merged = mergeAntigravityCatalog(upstreamCatalog, dynamicModels); + const models = merged.models as Record>; + + // Native model must retain its original displayName and metadata + assert.equal(models["gemini-2.5-pro"].displayName, "Native Gemini 2.5 Pro"); + assert.equal(models["gemini-2.5-pro"].descriptionText, "Google Official"); + assert.equal(models["gemini-2.5-pro"].isNative, true); + + // Non-colliding combo must be injected + assert.ok(models["coding-titans"]); + assert.equal(models["coding-titans"].displayName, "Coding Titans"); +}); + +test("antigravity handler — filters out hidden and inactive combos from dynamic catalog", async () => { + const { createCombo, deleteCombo } = await import("../../src/lib/db/combos.ts"); + const activeComboName = `active-combo-${Date.now()}`; + const inactiveComboName = `inactive-combo-${Date.now()}`; + + await createCombo({ + id: activeComboName, + name: activeComboName, + description: "Active combo", + models: JSON.stringify(["google/gemini-2.5-flash"]), + strategy: "priority", + isActive: true, + }); + + await createCombo({ + id: inactiveComboName, + name: inactiveComboName, + description: "Inactive combo", + models: JSON.stringify(["google/gemini-2.5-flash"]), + strategy: "priority", + isActive: false, + }); + + try { + const handler = new AntigravityHandler(); + const models = await handler.getDynamicCatalogModels(); + const activeFound = models.find((m) => m.id === activeComboName); + const inactiveFound = models.find((m) => m.id === inactiveComboName); + + assert.ok(activeFound, "Active combo must be included in catalog"); + assert.equal(inactiveFound, undefined, "Inactive combo must be excluded from catalog"); + } finally { + await deleteCombo(activeComboName); + await deleteCombo(inactiveComboName); + } +});