diff --git a/src/compat.ts b/src/compat.ts index 7bf2f72..9228cec 100644 --- a/src/compat.ts +++ b/src/compat.ts @@ -2153,11 +2153,11 @@ export function getEffectiveAntigravityModels(): CompatModelEntry[] { const seen = new Set(MODEL_CATALOG.map((model) => model.id.toLowerCase())); const result: CompatModelEntry[] = MODEL_CATALOG.map((model) => { const spec = getModelSpec(model.id); - const contextSpec = - getModelSpecOverride(model.id) ?? dynamicCatalog.getModelSpec(model.id); + const contextWindow = getModelSpecOverride(model.id)?.contextWindow ?? + dynamicCatalog.getModelSpec(model.id)?.contextWindow; return { ...model, - ctx: contextSpec?.contextWindow ?? model.ctx, + ctx: contextWindow ?? model.ctx, maxOutputTokens: spec.maxOutputTokens, thinkingBudget: spec.thinkingBudget, minThinkingBudget: spec.minThinkingBudget, diff --git a/src/compat/model-specs.ts b/src/compat/model-specs.ts index dd8da4b..f216fbc 100644 --- a/src/compat/model-specs.ts +++ b/src/compat/model-specs.ts @@ -9,6 +9,8 @@ export interface ModelSpec { contextWindow?: number; } +export type ModelSpecOverride = Partial; + export const DEFAULT_MODEL_SPECS: Record = { "gemini-pro-agent": { maxOutputTokens: 65535, thinkingBudget: 10001, isThinking: true, contextWindow: 1_000_000 }, "gemini-3-flash-agent": { maxOutputTokens: 65536, thinkingBudget: 10000, isThinking: true, contextWindow: 1_000_000 }, @@ -41,13 +43,13 @@ export const DEFAULT_MODEL_SPECS: Record = { "gpt-oss-120b": { maxOutputTokens: 32768, thinkingBudget: 8192, isThinking: true, contextWindow: 131_072 }, }; -let modelSpecsOverride: Record | null = null; +let modelSpecsOverride: Record | null = null; /** - * Replace the bundled model spec table with operator-provided overrides. + * Apply operator-provided partial overrides over effective model specs. * Pass `null` to restore defaults. Called once at startup from index.ts. */ -export function setModelSpecsOverride(specs: Record | null): void { +export function setModelSpecsOverride(specs: Record | null): void { modelSpecsOverride = specs && Object.keys(specs).length > 0 ? Object.fromEntries( Object.entries(specs).map(([key, spec]) => [key.toLowerCase(), spec]), @@ -56,7 +58,10 @@ export function setModelSpecsOverride(specs: Record | null): } export function getActiveModelSpecs(): Record { - return modelSpecsOverride ?? DEFAULT_MODEL_SPECS; + if (!modelSpecsOverride) return DEFAULT_MODEL_SPECS; + return Object.fromEntries( + Object.keys(modelSpecsOverride).map((model) => [model, getModelSpec(model)]), + ); } const GEMINI_MAX_OUTPUT_TOKENS = 65536; @@ -84,7 +89,7 @@ export function getStaticModelSpec(model: string): ModelSpec | undefined { return best?.spec; } -export function getModelSpecOverride(model: string): ModelSpec | undefined { +export function getModelSpecOverride(model: string): ModelSpecOverride | undefined { if (!modelSpecsOverride) return undefined; const lower = model.toLowerCase(); if (modelSpecsOverride[lower]) return modelSpecsOverride[lower]; @@ -94,19 +99,52 @@ export function getModelSpecOverride(model: string): ModelSpec | undefined { return undefined; } +function mergeModelSpec( + defaults: ModelSpec, + override: ModelSpecOverride, +): ModelSpec { + const minThinkingBudget = typeof override.minThinkingBudget === "number" && + Number.isFinite(override.minThinkingBudget) && + override.minThinkingBudget >= 0 + ? override.minThinkingBudget + : defaults.minThinkingBudget; + const contextWindow = typeof override.contextWindow === "number" && + Number.isFinite(override.contextWindow) && + override.contextWindow > 0 + ? override.contextWindow + : defaults.contextWindow; + return { + maxOutputTokens: typeof override.maxOutputTokens === "number" && + Number.isFinite(override.maxOutputTokens) && + override.maxOutputTokens > 0 + ? override.maxOutputTokens + : defaults.maxOutputTokens, + thinkingBudget: typeof override.thinkingBudget === "number" && + Number.isFinite(override.thinkingBudget) + ? override.thinkingBudget + : defaults.thinkingBudget, + ...(minThinkingBudget !== undefined ? { minThinkingBudget } : {}), + isThinking: typeof override.isThinking === "boolean" + ? override.isThinking + : defaults.isThinking, + ...(contextWindow !== undefined ? { contextWindow } : {}), + }; +} + export function getModelSpec(model: string): ModelSpec { const lower = model.toLowerCase(); - const override = getModelSpecOverride(lower); - if (override) return override; const dynamicSpec = dynamicCatalog.getModelSpec(lower); - if (dynamicSpec) return dynamicSpec; const staticSpec = getStaticModelSpec(lower); - if (staticSpec) return staticSpec; - const family = getModelFamily(model); - if (family === "claude") return { maxOutputTokens: CLAUDE_MAX_OUTPUT_TOKENS, thinkingBudget: CLAUDE_DEFAULT_THINKING_BUDGET, isThinking: true, contextWindow: 1_000_000 }; - if (family === "gemini") return { maxOutputTokens: GEMINI_MAX_OUTPUT_TOKENS, thinkingBudget: FALLBACK_THINKING_BUDGET, isThinking: true, contextWindow: 1_000_000 }; - return { maxOutputTokens: 65536, thinkingBudget: FALLBACK_THINKING_BUDGET, isThinking: false, contextWindow: 128_000 }; -} + const family = getModelFamily(model); + const defaults = dynamicSpec ?? staticSpec ?? + (family === "claude" + ? { maxOutputTokens: CLAUDE_MAX_OUTPUT_TOKENS, thinkingBudget: CLAUDE_DEFAULT_THINKING_BUDGET, isThinking: true, contextWindow: 1_000_000 } + : family === "gemini" + ? { maxOutputTokens: GEMINI_MAX_OUTPUT_TOKENS, thinkingBudget: FALLBACK_THINKING_BUDGET, isThinking: true, contextWindow: 1_000_000 } + : { maxOutputTokens: 65536, thinkingBudget: FALLBACK_THINKING_BUDGET, isThinking: false, contextWindow: 128_000 }); + const override = getModelSpecOverride(lower); + return override ? mergeModelSpec(defaults, override) : defaults; +} export function isThinkingModel(model: string): boolean { return getModelSpec(model).isThinking; diff --git a/src/config-defaults.ts b/src/config-defaults.ts index a667dde..158b0f9 100644 --- a/src/config-defaults.ts +++ b/src/config-defaults.ts @@ -55,7 +55,8 @@ export function applyConfigDefaults(config: Config): Config { idempotencyWindowMs: config.idempotencyWindowMs ?? 2000, streamRecoveryMaxRetries: safeStreamRecoveryMaxRetries(config.streamRecoveryMaxRetries), compressionMode: config.compressionMode ?? "off", -accounts: config.accounts ? config.accounts.map((account) => ({ + modelSpecs: config.modelSpecs, + accounts: config.accounts ? config.accounts.map((account) => ({ ...normalizeAccountConfig(account), tier: account.tier || "unknown", })) : [], diff --git a/src/providers/google-antigravity/catalog.ts b/src/providers/google-antigravity/catalog.ts index 96b9e96..a026cb7 100644 --- a/src/providers/google-antigravity/catalog.ts +++ b/src/providers/google-antigravity/catalog.ts @@ -10,6 +10,7 @@ // separately in `src/compat/model-specs.ts` via `maxOutputTokens`. import { dynamicCatalog } from "./dynamic-catalog.js"; +import { getModelSpecOverride } from "../../compat/model-specs.js"; export interface AntigravityModelSpec { id: string; @@ -73,16 +74,21 @@ const FALLBACK_CONTEXT_WINDOW = 128_000; * Resolve the upstream-published context window for an Antigravity model id. * * Lookup order: - * 1. Exact id match in static table (lowercased). - * 2. Dynamic catalog entry from live Antigravity endpoint. - * 3. Substring match across the table (longest key wins via the order here). - * 4. Family defaults: claude -> 1M, gemini -> 1M, gpt-oss -> 131_072. - * 5. Defensive fallback: 128_000. + * 1. Operator exact/substring override. + * 2. Exact id match in static table (lowercased). + * 3. Dynamic catalog entry from live Antigravity endpoint. + * 4. Substring match across the table (longest key wins via the order here). + * 5. Family defaults: claude -> 1M, gemini -> 1M, gpt-oss -> 131_072. + * 6. Defensive fallback: 128_000. */ export function getAntigravityContextWindow(model: string): number { if (!model) return FALLBACK_CONTEXT_WINDOW; const lower = model.toLowerCase().trim(); if (!lower) return FALLBACK_CONTEXT_WINDOW; + const overrideCtx = getModelSpecOverride(lower)?.contextWindow; + if (typeof overrideCtx === "number" && Number.isFinite(overrideCtx) && overrideCtx > 0) { + return overrideCtx; + } const exact = ANTIGRAVITY_CONTEXT_WINDOWS[lower]; if (typeof exact === "number") return exact; const dynamicCtx = dynamicCatalog.getContextWindow(lower); diff --git a/src/types.ts b/src/types.ts index 214eafc..d70530c 100644 --- a/src/types.ts +++ b/src/types.ts @@ -142,8 +142,8 @@ export interface Config { streamRecoveryMaxRetries?: number; // Prompt compression mode ("off" | "lite" | "rtk" | "rtk+lite"). compressionMode?: "off" | "lite" | "rtk" | "rtk+lite"; - // Override per-model specs used by the compat layer. Keys are model id substrings - // matched case-insensitively. When set, replaces the bundled defaults entirely. + // Partially override per-model specs used by the compat layer. Keys are model id + // substrings matched case-insensitively; omitted fields inherit effective defaults. modelSpecs?: Record; // Override model-id aliases used to translate the operator-facing name // (e.g. "gemini-3.1-pro-high") to the upstream Antigravity name @@ -224,10 +224,11 @@ export interface GoogleQuotaResponse { // Per-model thinking/output spec used by the compat layer. // Operators can override defaults via the `modelSpecs` field in accounts.json. export interface ModelSpecConfig { - maxOutputTokens: number; - thinkingBudget: number; // -1 = adaptive (model decides), >=0 = fixed + maxOutputTokens?: number; + thinkingBudget?: number; // -1 = adaptive (model decides), >=0 = fixed minThinkingBudget?: number; - isThinking: boolean; + isThinking?: boolean; + contextWindow?: number; } // Per-model quota info for an account diff --git a/src/validators.ts b/src/validators.ts index 2978010..ae15bd3 100644 --- a/src/validators.ts +++ b/src/validators.ts @@ -179,6 +179,9 @@ export function validateConfig(value: unknown): ValidationResult { if (spec.isThinking !== undefined && typeof spec.isThinking !== "boolean") { errors.push(`config.modelSpecs.${key}.isThinking must be a boolean`); } + if (spec.contextWindow !== undefined && !isPositiveNumber(spec.contextWindow)) { + errors.push(`config.modelSpecs.${key}.contextWindow must be a positive number`); + } } } } diff --git a/test/context-windows.test.ts b/test/context-windows.test.ts index e9d7239..695254d 100644 --- a/test/context-windows.test.ts +++ b/test/context-windows.test.ts @@ -7,7 +7,10 @@ import { OPENCODE_ZEN_FREE_MODELS, getOpenCodeZenContextWindow, } from "../src/providers/opencode-zen/catalog.js"; -import { getModelSpec } from "../src/compat/model-specs.js"; +import { + getModelSpec, + setModelSpecsOverride, +} from "../src/compat/model-specs.js"; /** * Verify each catalog model has a positive, finite context window. @@ -69,6 +72,30 @@ describe("official context windows", () => { assert.equal(getModelSpec("gemini-3.1-pro-low").contextWindow, 1_000_000); assert.equal(getModelSpec("gpt-oss-120b-medium").contextWindow, 131_072); }); + + it("honors an exact operator context window override", () => { + setModelSpecsOverride({ + "gemini-3.8-flash-high": { contextWindow: 222_222 }, + }); + try { + assert.equal(getModelSpec("gemini-3.8-flash-high").contextWindow, 222_222); + assert.equal(getAntigravityContextWindow("gemini-3.8-flash-high"), 222_222); + } finally { + setModelSpecsOverride(null); + } + }); + + it("honors a substring operator context window override", () => { + setModelSpecsOverride({ + "gemini-3.8": { contextWindow: 333_333 }, + }); + try { + assert.equal(getModelSpec("gemini-3.8-flash-high").contextWindow, 333_333); + assert.equal(getAntigravityContextWindow("gemini-3.8-flash-high"), 333_333); + } finally { + setModelSpecsOverride(null); + } + }); }); describe("openai codex", () => { diff --git a/test/db-store.test.ts b/test/db-store.test.ts index 3b67c90..77d4f4a 100644 --- a/test/db-store.test.ts +++ b/test/db-store.test.ts @@ -19,6 +19,8 @@ import { } from "../src/db-store.js"; import type { Config, PersistedState, TokenUsageTiered } from "../src/types.js"; import type { PersistedResponsesStore } from "../src/db-store.js"; +import { applyConfigDefaults } from "../src/config-defaults.js"; +import { validateConfig } from "../src/validators.js"; describe("db-store helpers", () => { before(async () => { @@ -96,6 +98,35 @@ describe("db-store helpers", () => { assert.equal(cached.proxyPort, 51200); }); + it("preserves partial model specs through normalization and persistence", async () => { + const modelSpecs = { + "gemini-3.8": { + maxOutputTokens: 1_000, + contextWindow: 222_222, + }, + }; + const candidate = { + accounts: [], + modelSpecs, + unknownTopLevelField: "discard me", + }; + const validation = validateConfig(candidate); + + assert.equal(validation.ok, true); + assert.ok(validation.value); + + const normalized = applyConfigDefaults(validation.value); + assert.deepEqual(normalized.modelSpecs, modelSpecs); + assert.equal("unknownTopLevelField" in normalized, false); + + await setCachedConfig(validation.value); + const cached = getCachedConfig(); + assert.ok(cached); + assert.deepEqual(cached.modelSpecs, modelSpecs); + assert.equal("unknownTopLevelField" in cached, false); + assert.deepEqual(candidate.modelSpecs, modelSpecs); + }); + it("can cache and retrieve admin token", async () => { await setCachedAdminToken("test-token-value"); assert.equal(getCachedAdminToken(), "test-token-value"); diff --git a/test/dynamic-catalog.test.ts b/test/dynamic-catalog.test.ts index f6ca4be..488ea07 100644 --- a/test/dynamic-catalog.test.ts +++ b/test/dynamic-catalog.test.ts @@ -356,6 +356,70 @@ describe("DynamicModelRegistry", () => { } }); + it("merges partial substring overrides over bundled specs without invalid payload numbers", () => { + setModelSpecsOverride({ + "gemini-3.8": { + maxOutputTokens: 1000, + isThinking: true, + }, + }); + + try { + assert.deepEqual(getModelSpec("gemini-3.8-flash-high"), { + maxOutputTokens: 1000, + thinkingBudget: -1, + isThinking: true, + contextWindow: 1_000_000, + }); + + const body = openAIToAntigravityBody({ + model: "gemini-3.8-flash-high", + messages: [{ role: "user", content: "ping" }], + }) as { request: { generationConfig?: Record } }; + const serialized = JSON.parse(JSON.stringify(body)) as typeof body; + assert.deepEqual(serialized.request.generationConfig?.thinkingConfig, { + includeThoughts: true, + }); + } finally { + setModelSpecsOverride(null); + } + }); + + it("merges partial exact overrides over runtime-discovered specs", () => { + dynamicCatalog.updateFromEndpointResponse({ + models: { + "gemini-5.0-ultra": { + maxTokens: 5_000_000, + maxOutputTokens: 100_000, + supportsThinking: true, + thinkingBudget: 50_000, + minThinkingBudget: 2_000, + quotaInfo: { remainingFraction: 1 }, + }, + }, + }); + setModelSpecsOverride({ + "gemini-5.0-ultra": { maxOutputTokens: 60_000 }, + }); + + try { + assert.deepEqual(getModelSpec("gemini-5.0-ultra"), { + maxOutputTokens: 60_000, + thinkingBudget: 50_000, + minThinkingBudget: 2_000, + isThinking: true, + contextWindow: 5_000_000, + }); + + setModelSpecsOverride({ + "gemini-5.0-ultra": { thinkingBudget: 0 }, + }); + assert.equal(getModelSpec("gemini-5.0-ultra").thinkingBudget, 0); + } finally { + setModelSpecsOverride(null); + } + }); + it("exposes dynamic context windows through getAntigravityContextWindow", () => { dynamicCatalog.updateFromEndpointResponse({ models: {