fix: price discovered models by family and keep retry knobs on the extension options type
1234567891011121314151617181920212223242526272829303132333435363738import assert from "node:assert/strict";import { resolveAntigravityModelCost, STATIC_PROVIDER_MODELS } from "./models.js";
const cases: Array<[string, { input: number; output: number; cacheRead: number; cacheWrite: number }]> = [ // Static catalog models keep existing list prices ["gemini-3.7-flash-low", { input: 0.75, output: 3.75, cacheRead: 0.075, cacheWrite: 0 }], ["gemini-2.5-flash", { input: 0.3, output: 2.5, cacheRead: 0.03, cacheWrite: 0 }], ["gemini-2.5-pro", { input: 1.25, output: 10, cacheRead: 0.125, cacheWrite: 0 }], ["gemini-3.1-flash-lite", { input: 0.25, output: 1.5, cacheRead: 0.025, cacheWrite: 0 }], ["claude-opus-4-6-thinking", { input: 5, output: 25, cacheRead: 0.5, cacheWrite: 6.25 }], ["claude-sonnet-4-6", { input: 3, output: 15, cacheRead: 0.3, cacheWrite: 3.75 }], ["gpt-oss-120b-medium", { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }],
// Discovered models resolve to intended family list prices ["gemini-3.8-flash-low", { input: 0.75, output: 3.75, cacheRead: 0.075, cacheWrite: 0 }], ["gemini-3.8-flash-tiered", { input: 0.75, output: 3.75, cacheRead: 0.075, cacheWrite: 0 }], ["gemini-3.8-flash-lite", { input: 0.25, output: 1.5, cacheRead: 0.025, cacheWrite: 0 }], ["gemini-3.8-pro", { input: 2, output: 12, cacheRead: 0.2, cacheWrite: 0 }], ["claude-opus-4-7-thinking", { input: 5, output: 25, cacheRead: 0.5, cacheWrite: 6.25 }], ["claude-sonnet-4-7", { input: 3, output: 15, cacheRead: 0.3, cacheWrite: 3.75 }], ["claude-haiku-4-5", { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }],
// Unknown models fall through to zero ["unknown-model-xyz", { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }],];
for (const [id, expected] of cases) { const cost = resolveAntigravityModelCost(id); assert.deepEqual(cost, expected, `Mismatch for ${id}`);}
// Verify static provider configs retain identical costsfor (const model of STATIC_PROVIDER_MODELS) { const expected = resolveAntigravityModelCost(model.id); assert.deepEqual(model.cost, expected, `Mismatch in static provider model ${model.id}`);}
console.log(`OK: verified ${cases.length} resolution cases and ${STATIC_PROVIDER_MODELS.length} static provider models`);1 unmodified line234567891088 unmodified lines991001011021021031041051061077 unmodified lines1 unmodified line * Antigravity model discovery via the Cloud Code Assist fetchAvailableModels * endpoint. Copied from oh-my-pi's catalog discovery/antigravity module. */import type { ProviderModelConfig } from "@earendil-works/pi-coding-agent";import { resolveAntigravityModelCost } from "./models.js";import { ANTIGRAVITY_ENDPOINT_FALLBACKS, ensureAntigravityVersion, getAntigravityUserAgent } from "./wire.js";
const FETCH_AVAILABLE_MODELS_PATH = "/v1internal:fetchAvailableModels";const DEFAULT_CONTEXT_WINDOW = 200_000;88 unmodified lines id: modelId, name: model.displayName || modelId, reasoning: model.supportsThinking === true, input: supportsImages ? ["text", "image"] : ["text"], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, cost: resolveAntigravityModelCost(modelId), contextWindow: toPositiveNumber(model.maxTokens, DEFAULT_CONTEXT_WINDOW), maxTokens: toPositiveNumber(model.maxOutputTokens, DEFAULT_MAX_TOKENS), }); }7 unmodified lines10 unmodified lines111213141515161718192021222324252627282930313229 unmodified lines6263646561666768697071727374366 unmodified linesExpand all10 unmodified lines Message, Model, SimpleStreamOptions, StopReason, StreamOptions, 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";
29 unmodified lines } return result;}
export function buildBaseOptions(model: Model<Api>, options?: SimpleStreamOptions, apiKey?: string): StreamOptions {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,366 unmodified linesExpand all12 unmodified lines1314151617171819202122232463 unmodified lines88899091919293949512 unmodified lines */import { readFileSync } from "node:fs";import { homedir } from "node:os";import { join } from "node:path";import type { Api, AssistantMessageEventStream, Context, Model, SimpleStreamOptions } from "@earendil-works/pi-ai";import type { Api, AssistantMessageEventStream, Context, Model } from "@earendil-works/pi-ai";import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";import { fetchAntigravityModels } from "./discovery.js";import type { AntigravitySimpleStreamOptions } from "./helpers.js";import { STATIC_PROVIDER_MODELS } from "./models.js";import { antigravityOAuth } from "./oauth.js";import { streamSimpleAntigravity } from "./stream.js";import { ANTIGRAVITY_DAILY_ENDPOINT } from "./wire.js";63 unmodified lines baseUrl, api: PROVIDER_API, models, oauth: antigravityOAuth, streamSimple: (model: Model<Api>, context: Context, options?: SimpleStreamOptions): AssistantMessageEventStream => streamSimple: (model: Model<Api>, context: Context, options?: AntigravitySimpleStreamOptions): AssistantMessageEventStream => streamSimpleAntigravity(model, context, options), });}41 unmodified lines424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939453959697989921 unmodified lines121122123124831251261271281291 unmodified line131132133134931351361371381391 unmodified line1411421431441031451461471481491 unmodified line1511521531541131551561571581591 unmodified line1611621631641231651661671681691 unmodified line1711721731741331751761771781791 unmodified line1811821831841431851861871881891 unmodified line1911921931941531951961971981991 unmodified line2012022032041632052062072082091 unmodified line2112122132141732152162172182196 unmodified lines2262272282291882302312322332346 unmodified lines2412422432442032452462472482496 unmodified lines2562572582592182602612622632646 unmodified lines2712722732742332752762772782796 unmodified lines2862872882892482902912922932946 unmodified lines3013023033042633053063073083096 unmodified lines31631731831927832032132232332484 unmodified lines41 unmodified linesexport function identityClassFor(modelId: string = ""): AntigravityModelProfile["identityClass"] { return (modelId || "").startsWith("claude-") ? "anthropic" : "gemini";}
export interface AntigravityModelCost { input: number; output: number; cacheRead: number; cacheWrite: number;}
export const ANTIGRAVITY_MODEL_COSTS = { "claude-opus": { input: 5, output: 25, cacheRead: 0.5, cacheWrite: 6.25 }, "claude-sonnet": { input: 3, output: 15, cacheRead: 0.3, cacheWrite: 3.75 }, "claude-haiku": { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, "gemini-flash": { input: 0.75, output: 3.75, cacheRead: 0.075, cacheWrite: 0 }, "gemini-flash-lite": { input: 0.25, output: 1.5, cacheRead: 0.025, cacheWrite: 0 }, "gemini-pro": { input: 2, output: 12, cacheRead: 0.2, cacheWrite: 0 }, "gpt-oss": { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, "gemini-2.5-flash": { input: 0.3, output: 2.5, cacheRead: 0.03, cacheWrite: 0 }, "gemini-2.5-flash-lite": { input: 0.1, output: 0.4, cacheRead: 0.01, cacheWrite: 0 }, "gemini-2.5-pro": { input: 1.25, output: 10, cacheRead: 0.125, cacheWrite: 0 }, "gemini-3.5-flash-extra-low": { input: 1.5, output: 9, cacheRead: 0.15, cacheWrite: 0 }, "gemini-3.5-flash-low": { input: 1.5, output: 9, cacheRead: 0.15, cacheWrite: 0 }, "gemini-3-flash-agent": { input: 0.5, output: 3, cacheRead: 0.05, cacheWrite: 0 },} as const satisfies Record<string, AntigravityModelCost>;
const ZERO_COST: AntigravityModelCost = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 };
export function resolveAntigravityModelCost(modelId: string = ""): AntigravityModelCost { if (modelId in ANTIGRAVITY_MODEL_COSTS) { return ANTIGRAVITY_MODEL_COSTS[modelId as keyof typeof ANTIGRAVITY_MODEL_COSTS]; } const lower = (modelId || "").toLowerCase(); if (lower.startsWith("claude-") && lower.includes("opus")) return ANTIGRAVITY_MODEL_COSTS["claude-opus"]; if (lower.startsWith("claude-") && lower.includes("sonnet")) return ANTIGRAVITY_MODEL_COSTS["claude-sonnet"]; if (lower.startsWith("claude-") && lower.includes("haiku")) return ANTIGRAVITY_MODEL_COSTS["claude-haiku"]; if (lower.startsWith("gpt-oss")) return ANTIGRAVITY_MODEL_COSTS["gpt-oss"]; if (lower.startsWith("gemini-") && lower.includes("flash-lite")) return ANTIGRAVITY_MODEL_COSTS["gemini-flash-lite"]; if (lower.startsWith("gemini-") && (lower.includes("flash") || lower.includes("flash-agent"))) { return ANTIGRAVITY_MODEL_COSTS["gemini-flash"]; } if (lower.startsWith("gemini-") && lower.includes("pro")) return ANTIGRAVITY_MODEL_COSTS["gemini-pro"]; return ZERO_COST;}
interface AntigravityModelDef { id: string; name: string; reasoning: boolean; input: ("text" | "image")[]; contextWindow: number; maxTokens: number; cost: { input: number; output: number; cacheRead: number; cacheWrite: number }; cost: AntigravityModelCost; thinking: AntigravityModelThinking;}
const THINK_BUDGET_DEFAULT: AntigravityModelThinking = {21 unmodified lines reasoning: true, input: ["text", "image"], contextWindow: 200000, maxTokens: 64000, cost: { input: 5, output: 25, cacheRead: 0.5, cacheWrite: 6.25 }, cost: resolveAntigravityModelCost("claude-opus-4-5-thinking"), thinking: { ...THINK_BUDGET_DEFAULT }, }, { id: "claude-opus-4-6-thinking",1 unmodified line reasoning: true, input: ["text", "image"], contextWindow: 250000, maxTokens: 64000, cost: { input: 5, output: 25, cacheRead: 0.5, cacheWrite: 6.25 }, cost: resolveAntigravityModelCost("claude-opus-4-6-thinking"), thinking: { ...THINK_BUDGET_DEFAULT }, }, { id: "claude-sonnet-4-5",1 unmodified line reasoning: true, input: ["text", "image"], contextWindow: 1000000, maxTokens: 64000, cost: { input: 3, output: 15, cacheRead: 0.3, cacheWrite: 3.75 }, cost: resolveAntigravityModelCost("claude-sonnet-4-5"), thinking: { ...THINK_BUDGET_DEFAULT }, }, { id: "claude-sonnet-4-6",1 unmodified line reasoning: true, input: ["text", "image"], contextWindow: 250000, maxTokens: 64000, cost: { input: 3, output: 15, cacheRead: 0.3, cacheWrite: 3.75 }, cost: resolveAntigravityModelCost("claude-sonnet-4-6"), thinking: { ...THINK_BUDGET_DEFAULT }, }, { id: "gemini-2.5-flash",1 unmodified line reasoning: true, input: ["text", "image"], contextWindow: 1048576, maxTokens: 65535, cost: { input: 0.3, output: 2.5, cacheRead: 0.03, cacheWrite: 0 }, cost: resolveAntigravityModelCost("gemini-2.5-flash"), thinking: { ...THINK_BUDGET_DEFAULT }, }, { id: "gemini-2.5-flash-lite",1 unmodified line reasoning: true, input: ["text", "image"], contextWindow: 1048576, maxTokens: 65535, cost: { input: 0.1, output: 0.4, cacheRead: 0.01, cacheWrite: 0 }, cost: resolveAntigravityModelCost("gemini-2.5-flash-lite"), thinking: { ...THINK_BUDGET_DEFAULT }, }, { id: "gemini-2.5-pro",1 unmodified line reasoning: true, input: ["text", "image"], contextWindow: 1048576, maxTokens: 65536, cost: { input: 1.25, output: 10, cacheRead: 0.125, cacheWrite: 0 }, cost: resolveAntigravityModelCost("gemini-2.5-pro"), thinking: { ...THINK_BUDGET_DEFAULT, requiresEffort: true }, }, { id: "gemini-3-pro-low",1 unmodified line reasoning: true, input: ["text", "image"], contextWindow: 1048576, maxTokens: 65535, cost: { input: 2, output: 12, cacheRead: 0.2, cacheWrite: 0 }, cost: resolveAntigravityModelCost("gemini-3-pro-low"), thinking: { mode: "google-level", suppressWhenOff: true, requiresEffort: true }, }, { id: "gemini-3.1-flash-lite",1 unmodified line reasoning: true, input: ["text", "image"], contextWindow: 1048576, maxTokens: 65535, cost: { input: 0.25, output: 1.5, cacheRead: 0.025, cacheWrite: 0 }, cost: resolveAntigravityModelCost("gemini-3.1-flash-lite"), thinking: { mode: "google-level", suppressWhenOff: false, requiresEffort: true }, }, { id: "gemini-3.1-pro-low",1 unmodified line reasoning: true, input: ["text", "image"], contextWindow: 1048576, maxTokens: 65535, cost: { input: 2, output: 12, cacheRead: 0.2, cacheWrite: 0 }, cost: resolveAntigravityModelCost("gemini-3.1-pro-low"), thinking: { mode: "budget", effortBudgets: GEMINI_PRO_BUDGETS, suppressWhenOff: true,6 unmodified lines reasoning: true, input: ["text", "image"], contextWindow: 1048576, maxTokens: 65535, cost: { input: 2, output: 12, cacheRead: 0.2, cacheWrite: 0 }, cost: resolveAntigravityModelCost("gemini-pro-agent"), thinking: { mode: "budget", effortBudgets: GEMINI_PRO_BUDGETS, suppressWhenOff: true,6 unmodified lines reasoning: true, input: ["text", "image"], contextWindow: 1048576, maxTokens: 65536, cost: { input: 1.5, output: 9, cacheRead: 0.15, cacheWrite: 0 }, cost: resolveAntigravityModelCost("gemini-3.5-flash-extra-low"), thinking: { mode: "budget", effortBudgets: GEMINI_FLASH_BUDGETS, suppressWhenOff: true,6 unmodified lines reasoning: true, input: ["text", "image"], contextWindow: 1048576, maxTokens: 65536, cost: { input: 1.5, output: 9, cacheRead: 0.15, cacheWrite: 0 }, cost: resolveAntigravityModelCost("gemini-3.5-flash-low"), thinking: { mode: "budget", effortBudgets: GEMINI_FLASH_BUDGETS, suppressWhenOff: true,6 unmodified lines reasoning: true, input: ["text", "image"], contextWindow: 1048576, maxTokens: 65536, cost: { input: 0.5, output: 3, cacheRead: 0.05, cacheWrite: 0 }, cost: resolveAntigravityModelCost("gemini-3-flash-agent"), thinking: { mode: "budget", effortBudgets: GEMINI_FLASH_BUDGETS, suppressWhenOff: true,6 unmodified lines reasoning: true, input: ["text", "image"], contextWindow: 1048576, maxTokens: 65536, cost: { input: 0.75, output: 3.75, cacheRead: 0.075, cacheWrite: 0 }, cost: resolveAntigravityModelCost("gemini-3.6-flash-low"), thinking: { mode: "google-level", suppressWhenOff: false, requiresEffort: true,6 unmodified lines reasoning: true, input: ["text", "image"], contextWindow: 1048576, maxTokens: 65536, cost: { input: 0.75, output: 3.75, cacheRead: 0.075, cacheWrite: 0 }, cost: resolveAntigravityModelCost("gemini-3.7-flash-low"), thinking: { mode: "google-level", suppressWhenOff: false, requiresEffort: true,6 unmodified lines reasoning: true, input: ["text"], contextWindow: 131072, maxTokens: 32768, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, cost: resolveAntigravityModelCost("gpt-oss-120b-medium"), thinking: { ...THINK_BUDGET_DEFAULT }, },];
84 unmodified lines12 unmodified lines13141516171718192021222324252627354 unmodified linesExpand all382383384385386386387388389390622 unmodified linesExpand all10131014101510161017101710181019102010218 unmodified lines12 unmodified lines type Context, calculateCost, createAssistantMessageEventStream, type Model, type SimpleStreamOptions, type TextContent, type ThinkingContent, type ThinkingLevel, type ToolCall,} from "@earendil-works/pi-ai";import { type AntigravitySimpleStreamOptions, buildBaseOptions, clampReasoning, convertMessages, convertTools,354 unmodified linesExpand all labels.used_claude_conservative = usageLabel; return { sessionId: state.sessionId, requestId, labels };}
export interface AntigravityStreamOptions extends SimpleStreamOptions {export interface AntigravityStreamOptions extends AntigravitySimpleStreamOptions { toolChoice?: "auto" | "none" | "any"; topP?: number; topK?: number; minP?: number;622 unmodified linesExpand all
export function streamSimpleAntigravity( model: Model<Api>, context: Context, options?: SimpleStreamOptions, options?: AntigravitySimpleStreamOptions,): AssistantMessageEventStream { const apiKey = options?.apiKey; if (!apiKey) { throw new Error("Antigravity requires OAuth authentication. Use /login google-antigravity to authenticate.");8 unmodified lines