From 7e61bc210521480e03536f2915cc8f02b5a451ea Mon Sep 17 00:00:00 2001 From: fuko2935 Date: Mon, 27 Jul 2026 00:48:50 +0300 Subject: [PATCH] fix(sse): clamp out-of-range Gemini thinking_budget and learn the cap Unregistered Gemini models (any model id containing "gemini", on any provider) previously got no thinkingBudgetCap, so a high client budget (xhigh = 131072) was forwarded verbatim and the upstream returned 400 "thinking_budget must be in the range [-1, N]". Two layers: - Proactive: capThinkingBudget now falls back to the pro-tier ceiling (32768) for Gemini-family model ids with no registry cap, so the first request is already in range (no 400->retry round-trip). - Reactive: the executor detects the range 400, clamps the nested budget (native Gemini and Antigravity envelope) to the parsed upstream max, retries the same URL once, and records the surviving cap in an in-memory per-provider+model store so subsequent requests clamp proactively (the retry cost is paid once per process, not per chat session). Existing tests updated where they locked in the previous no-cap behavior (gemini-2.5-pro max/xhigh asserted the un-sendable 131072). --- open-sse/executors/base.ts | 116 ++++++ open-sse/services/learnedThinkingCaps.ts | 113 ++++++ src/lib/modelCapabilities.ts | 45 ++- ...ap-thinking-budget-gemini-fallback.test.ts | 79 ++++ .../gemini-thinking-budget-fallback.test.ts | 338 ++++++++++++++++++ tests/unit/learned-thinking-caps.test.ts | 109 ++++++ .../unit/translator-claude-to-gemini.test.ts | 10 +- 7 files changed, 806 insertions(+), 4 deletions(-) create mode 100644 open-sse/services/learnedThinkingCaps.ts create mode 100644 tests/unit/cap-thinking-budget-gemini-fallback.test.ts create mode 100644 tests/unit/gemini-thinking-budget-fallback.test.ts create mode 100644 tests/unit/learned-thinking-caps.test.ts diff --git a/open-sse/executors/base.ts b/open-sse/executors/base.ts index 7972084afaaf..a1ed8c840983 100644 --- a/open-sse/executors/base.ts +++ b/open-sse/executors/base.ts @@ -11,6 +11,10 @@ import { detectUnsupportedParam, stripGroqUnsupportedFields, } from "../config/providerFieldStrips.ts"; +import { + recordLearnedThinkingCap, + parseThinkingBudgetMax, +} from "../services/learnedThinkingCaps.ts"; import { getParamFilterConfig, addParamToBlocklist, @@ -228,6 +232,64 @@ function hasActiveClaudeThinking(body: Record): boolean { return thinking?.type === "enabled" || thinking?.type === "adaptive"; } +/** + * Collect every `thinkingConfig` object in a transformed request body that holds + * a thinking budget, wherever the provider's envelope nests it: + * - body.generationConfig.thinkingConfig (native Gemini / openai→gemini) + * - body.request.generationConfig.thinkingConfig (Antigravity Cloud Code envelope) + * Returns only objects that actually carry a `thinkingBudget`/`thinking_budget` + * field — a request without thinking config is never mutated. + */ +function collectThinkingConfigs(body: unknown): Array> { + if (!body || typeof body !== "object") return []; + const root = body as Record; + const configs: Array> = []; + const envelopes: unknown[] = [root.generationConfig, (root.request as Record | undefined)?.generationConfig]; + for (const env of envelopes) { + if (!env || typeof env !== "object") continue; + const tc = (env as Record).thinkingConfig; + if (tc && typeof tc === "object") { + const tcr = tc as Record; + if ("thinkingBudget" in tcr || "thinking_budget" in tcr) configs.push(tcr); + } + } + return configs; +} + +/** + * Read the first thinking budget found in the body (any supported nest / naming). + * Returns null when the body carries no readable numeric budget. + */ +function readNestedThinkingBudget(body: unknown): number | null { + for (const tc of collectThinkingConfigs(body)) { + const raw = tc.thinkingBudget ?? tc.thinking_budget; + const n = Number(raw); + if (Number.isFinite(n)) return n; + } + return null; +} + +/** + * Clamp every thinking budget in the body down to `max` (only lowers; never + * raises a budget already below max). Mutates in place. Returns true when at + * least one budget was actually lowered (i.e. a retry would send a different + * body) — false means the 400 was not caused by an over-max budget we hold, so + * retrying would resend an identical body and loop. + */ +function clampNestedThinkingBudget(body: unknown, max: number): boolean { + let changed = false; + for (const tc of collectThinkingConfigs(body)) { + for (const key of ["thinkingBudget", "thinking_budget"] as const) { + const n = Number(tc[key]); + if (Number.isFinite(n) && n > max) { + tc[key] = max; + changed = true; + } + } + } + return changed; +} + /** * Strip the OmniRoute provider prefix from versioned built-in tool model * fields (e.g. `cc/claude-opus-4-8` → `claude-opus-4-8`). Versioned built-in @@ -722,6 +784,13 @@ export class BaseExecutor { // field-downgrade below, so each known field is stripped at most once across // all fallback URLs (bounded retry loop). const strippedFields = new Set(); + // Set by the thinking_budget 400 clamp-and-retry below: the upstream's + // advertised max (parsed from the error) is applied to every later + // retry/fallback URL so they don't re-hit the same 400. The clamp itself + // fires at most once per URL (guarded inline) so a persistent 400 cannot + // loop. The learned cap is also recorded process-wide via + // recordLearnedThinkingCap so future requests skip the 400 entirely. + let thinkingBudgetClampedMax: number | null = null; for (let urlIndex = 0; urlIndex < fallbackCount; urlIndex++) { const requestCredentials = withForcedResponsesUpstream( @@ -775,6 +844,13 @@ export class BaseExecutor { transformedBody as Record ) as typeof transformedBody; } + // A previous URL in this execute() already hit a thinking_budget 400 and + // recorded the upstream's max. Pre-clamp this URL's fresh transformedBody + // so it doesn't re-hit the same 400 (avoids one wasted round-trip per + // fallback URL). No-op when nothing has been learned this execute(). + if (thinkingBudgetClampedMax !== null) { + clampNestedThinkingBudget(transformedBody, thinkingBudgetClampedMax); + } try { // Timeout only covers response start; stream stalls are handled downstream. @@ -1309,6 +1385,46 @@ export class BaseExecutor { } } + // Thinking-budget 400 clamp-and-retry (Gemini-family, any provider). + // The proactive capThinkingBudget clamp misses models outside MODEL_SPECS, + // so an xhigh budget (131072) can reach the upstream and bounce with + // "thinking_budget must be in the range [-1, N]". When that happens: parse + // the advertised max, record it process-wide (so FUTURE requests clamp + // proactively via capThinkingBudget → getLearnedThinkingCap), clamp the + // nested budget in the live transformedBody, and retry the same URL once. + // Subsequent requests never pay this 400→retry round-trip. + if ( + response.status === HTTP_STATUS.BAD_REQUEST && + thinkingBudgetClampedMax === null && + transformedBody && + typeof transformedBody === "object" + ) { + const errText = await response + .clone() + .text() + .catch(() => ""); + const upstreamMax = parseThinkingBudgetMax(errText); + if (upstreamMax !== null) { + const currentBudget = readNestedThinkingBudget(transformedBody); + // Record under the budget that actually failed so the learned step + // ladder advances correctly; fall back to upstreamMax when the body + // carries no readable budget (still record so we stop re-hitting). + recordLearnedThinkingCap(this.provider, model, currentBudget ?? upstreamMax + 1); + thinkingBudgetClampedMax = upstreamMax; + if (clampNestedThinkingBudget(transformedBody, upstreamMax)) { + let retryBody = JSON.stringify(transformedBody); + if (isClaudeCodeCompatible(this.provider) || this.provider === "claude") { + retryBody = await signRequestBody(retryBody); + } + log?.info?.( + "THINKING_BUDGET", + `Upstream 400 rejected thinking_budget on ${url} — clamped to ${upstreamMax} and retrying (learned for ${this.provider}/${model})` + ); + response = await fetchWithStartTimeout(url, { ...fetchOptions, body: retryBody }); + } + } + } + // Generic reactive 400 field-downgrade; each field is stripped at most once. if ( response.status === HTTP_STATUS.BAD_REQUEST && diff --git a/open-sse/services/learnedThinkingCaps.ts b/open-sse/services/learnedThinkingCaps.ts new file mode 100644 index 000000000000..0603fc3db9a3 --- /dev/null +++ b/open-sse/services/learnedThinkingCaps.ts @@ -0,0 +1,113 @@ +/** + * Learned Thinking-Budget Caps — reactive cap memory for Gemini-family models. + * + * Proactive clamping (`capThinkingBudget` in src/lib/modelCapabilities.ts) only + * knows caps for models registered in MODEL_SPECS. Any Gemini model outside the + * registry (e.g. gemini-2.5-pro, or a Gemini id served through another provider) + * initially gets no cap, so a high client budget (xhigh = 131072) goes straight + * to the upstream and can bounce with: + * + * 400 GenerateContentRequest.generation_config.thinking_config.thinking_budget: + * thinking_budget must be in the range [-1, 65535] + * + * When base.ts's executor sees that 400 it calls `recordLearnedThinkingCap`, + * which walks the failed budget down a step ladder and stores the surviving cap + * in a module-level Map keyed "provider:model" (lowercased). Subsequent requests + * for the same provider+model read the cap via `getLearnedThinkingCap` (consulted + * by `capThinkingBudget`) so the 400→retry round-trip is paid at most once per + * process per provider+model — not once per chat session. + * + * In-memory only (per operator decision): restart resets, the first request after + * a restart may re-learn at the cost of one upstream 400. + */ + +// Known Gemini thinking-budget ceilings, highest first. 32768 is the pro-tier +// cap; 24576 is the flash-tier cap (gemini-2.5-flash in MODEL_SPECS); 8192 is a +// conservative floor that every thinking-capable Gemini accepts. A failed budget +// walks to the first step strictly below it. +const GEMINI_STEPDOWN: readonly number[] = [32768, 24576, 8192]; + +/** + * Proactive cap applied by `capThinkingBudget` (src/lib/modelCapabilities.ts) + * when a model id contains "gemini" but has no registry entry and no learned cap. + * Set to the highest rung of the step ladder so unregistered Gemini models get + * the maximum reasoning the family is known to accept, rather than letting an + * xhigh budget (131072) through to a 400. + */ +export const GEMINI_FALLBACK_THINKING_CAP: number = GEMINI_STEPDOWN[0]; + +// key: `${provider}:${model}` lowercased → highest budget known to be accepted. +const learnedCaps = new Map(); + +function buildKey(provider: string | null | undefined, model: string | null | undefined): string { + const p = typeof provider === "string" ? provider.trim().toLowerCase() : ""; + const m = typeof model === "string" ? model.trim().toLowerCase() : ""; + if (!p || !m) return ""; + return `${p}:${m}`; +} + +/** + * Return the learned cap for provider+model, or null when nothing has been + * learned yet (no upstream 400 recorded). Keyed case-insensitively. + */ +export function getLearnedThinkingCap( + provider: string | null | undefined, + model: string | null | undefined +): number | null { + const key = buildKey(provider, model); + if (!key) return null; + return learnedCaps.get(key) ?? null; +} + +/** + * Record that `failedBudget` was rejected by the upstream for provider+model and + * store the next step down as the learned cap. Returns the new cap, or null when + * the step ladder is exhausted (even 8192 failed) or the key is unusable — the + * caller should then give up clamping and surface the upstream error. + * + * Always monotonically decreases: if a cap already stored is lower than the step + * we'd compute, the stored (lower) value wins and is returned unchanged. This + * keeps concurrent failures on the same provider+model from ratcheting the cap + * back up. + */ +export function recordLearnedThinkingCap( + provider: string | null | undefined, + model: string | null | undefined, + failedBudget: number +): number | null { + const key = buildKey(provider, model); + if (!key) return null; + if (!Number.isFinite(failedBudget)) return null; + + const nextStep = GEMINI_STEPDOWN.find((step) => step < failedBudget); + if (nextStep === undefined) return null; // ladder exhausted + + const existing = learnedCaps.get(key); + if (existing !== undefined && existing <= nextStep) { + return existing; // already learned an equal-or-lower cap; keep it + } + learnedCaps.set(key, nextStep); + return nextStep; +} + +/** + * Extract the upstream-advertised maximum from a thinking_budget range error. + * Matches both snake_case and camelCase field namings: + * "thinking_budget must be in the range [-1, 65535]" → 65535 + * "thinkingBudget must be in the range [-1, 24576]" → 24576 + * Returns null when the text is not a thinking-budget range error. + */ +export function parseThinkingBudgetMax(errText: unknown): number | null { + if (typeof errText !== "string" || !errText) return null; + const match = /thinking_?budget[^\d-]*(?:must be in the range|range)[^\d-]*\[\s*-?\d+\s*,\s*(\d+)\s*\]/i.exec( + errText + ); + if (!match) return null; + const max = Number(match[1]); + return Number.isFinite(max) ? max : null; +} + +/** Test-only: clear the learned-cap Map between tests. */ +export function __test_resetLearnedThinkingCaps(): void { + learnedCaps.clear(); +} diff --git a/src/lib/modelCapabilities.ts b/src/lib/modelCapabilities.ts index 60f3d349dd50..95330b45f826 100644 --- a/src/lib/modelCapabilities.ts +++ b/src/lib/modelCapabilities.ts @@ -16,6 +16,10 @@ import { getModelContextOverride } from "@/lib/db/modelContextOverrides"; import { getModelCapabilityOverride } from "@/lib/db/modelCapabilityOverrides"; import { isVisionModelId } from "@/shared/constants/visionModels"; import { getUnsupportedParams } from "@omniroute/open-sse/config/providerRegistry.ts"; +import { + getLearnedThinkingCap, + GEMINI_FALLBACK_THINKING_CAP, +} from "@omniroute/open-sse/services/learnedThinkingCaps.ts"; const TOOL_CALLING_UNSUPPORTED_PATTERNS: string[] = [ // Specialty / non-chat surfaces must never inherit optimistic tool defaults (#8016) @@ -663,9 +667,46 @@ export function getDefaultThinkingBudget(input: CapabilityInput): number { return getResolvedModelCapabilities(input).defaultThinkingBudget; } +/** + * Clamp a requested thinking budget to the model's real ceiling. + * + * Resolution order (lowest wins): + * 1. Registry cap (MODEL_SPECS.thinkingBudgetCap) — authoritative when present. + * 2. Learned cap — a lower ceiling previously discovered via an upstream 400 + * ("thinking_budget must be in the range ...") recorded by the executor + * (open-sse/services/learnedThinkingCaps.ts). In-memory, per provider+model. + * 3. Gemini-family fallback — when the registry has no cap but the model id + * contains "gemini" (any provider: many providers host Gemini models), clamp + * to GEMINI_FALLBACK_THINKING_CAP (32768, the known pro-tier cap) instead of + * letting an xhigh budget (131072) sail through to a 400. Registered flash + * models already carry their explicit 24576 cap via rule 1, so this only + * fires for unregistered Gemini ids. + */ export function capThinkingBudget(input: CapabilityInput, budget: number): number { - const cap = getResolvedModelCapabilities(input).thinkingBudgetCap ?? budget; - return Math.min(budget, cap); + const resolved = getResolvedModelCapabilities(input); + let cap = resolved.thinkingBudgetCap; + + const modelId = resolved.model ?? resolved.rawModel ?? ""; + const modelLower = modelId.toLowerCase(); + // Learned-cap lookup needs a concrete provider key (the executor records under + // `this.provider`). When the input is a bare Gemini id, `resolved.provider` is + // null — but bare Gemini ids always route to the native Gemini provider, so + // default to "gemini". Without this a cap learned via the executor would be + // invisible to bare-model callers. Provider-qualified inputs keep their own + // provider, preserving per-provider independence. + const providerForLearned = + resolved.provider ?? (modelLower.includes("gemini") ? "gemini" : null); + + const learned = getLearnedThinkingCap(providerForLearned, modelId); + if (learned !== null) { + cap = cap === null ? learned : Math.min(cap, learned); + } + + if (cap === null && modelLower.includes("gemini")) { + cap = GEMINI_FALLBACK_THINKING_CAP; + } + + return Math.min(budget, cap ?? budget); } export function getModelContextLimit( diff --git a/tests/unit/cap-thinking-budget-gemini-fallback.test.ts b/tests/unit/cap-thinking-budget-gemini-fallback.test.ts new file mode 100644 index 000000000000..136d49dd81db --- /dev/null +++ b/tests/unit/cap-thinking-budget-gemini-fallback.test.ts @@ -0,0 +1,79 @@ +import { test, beforeEach, after } from "node:test"; +import assert from "node:assert/strict"; +import { capThinkingBudget } from "../../src/lib/modelCapabilities.ts"; +import { + recordLearnedThinkingCap, + __test_resetLearnedThinkingCaps, +} from "../../open-sse/services/learnedThinkingCaps.ts"; + +beforeEach(() => { + __test_resetLearnedThinkingCaps(); +}); + +after(() => { + __test_resetLearnedThinkingCaps(); +}); + +// ── Proactive Gemini fallback (no registry cap, no learned cap) ──────────── + +test("unregistered gemini model clamps xhigh budget to the 32768 fallback", () => { + // gemini-2.5-pro has no MODEL_SPECS entry — previously the raw 131072 sailed + // through to a 400. The gemini-substring fallback now clamps it. + assert.equal(capThinkingBudget("gemini-2.5-pro", 131072), 32768); +}); + +test("gemini fallback is provider-agnostic (model-id driven, any provider)", () => { + assert.equal(capThinkingBudget("openrouter/gemini-2.5-pro", 131072), 32768); + assert.equal(capThinkingBudget("vertex/gemini-2.5-pro", 131072), 32768); +}); + +test("gemini fallback is case-insensitive on the model id", () => { + assert.equal(capThinkingBudget("Gemini-2.5-Pro", 131072), 32768); +}); + +test("gemini fallback does not inflate a budget already below the cap", () => { + assert.equal(capThinkingBudget("gemini-2.5-pro", 1024), 1024); + assert.equal(capThinkingBudget("gemini-2.5-pro", 8192), 8192); +}); + +test("non-gemini unregistered model passes through unchanged (no fallback)", () => { + assert.equal(capThinkingBudget("some-random-model", 131072), 131072); + assert.equal(capThinkingBudget("qwen-something", 131072), 131072); +}); + +// ── Registry cap still wins for registered models ────────────────────────── + +test("registered gemini-2.5-flash keeps its explicit 24576 cap (fallback does not override)", () => { + assert.equal(capThinkingBudget("gemini-2.5-flash", 131072), 24576); +}); + +test("registered gemini-3.1-pro keeps its explicit 32768 cap", () => { + assert.equal(capThinkingBudget("gemini-3.1-pro", 131072), 32768); +}); + +// ── Learned cap integration ──────────────────────────────────────────────── + +test("a learned cap lower than the registry cap wins", () => { + // Registry says gemini-2.5-flash caps at 24576; the upstream actually rejected + // 24576, so the learned 8192 must win on the next request. Bare model id + // resolves to the native "gemini" provider for the learned-cap lookup. + recordLearnedThinkingCap("gemini", "gemini-2.5-flash", 24576); // → 8192 + assert.equal(capThinkingBudget("gemini-2.5-flash", 131072), 8192); +}); + +test("a learned cap applies to an unregistered gemini model (overrides the 32768 fallback)", () => { + // First request uses the 32768 fallback; suppose upstream rejects 32768 → + // learned 24576. Next request must use 24576, not 32768. + recordLearnedThinkingCap("openrouter", "gemini-2.5-pro", 32768); // → 24576 + assert.equal(capThinkingBudget("openrouter/gemini-2.5-pro", 131072), 24576); +}); + +test("learned caps are scoped per provider+model (no cross-contamination)", () => { + recordLearnedThinkingCap("openrouter", "gemini-2.5-pro", 32768); // → 24576 on openrouter only + // A different provider serving the same model id is unaffected (still fallback). + assert.equal(capThinkingBudget("vertex/gemini-2.5-pro", 131072), 32768); + // The native-gemini provider key is independent too. + assert.equal(capThinkingBudget("gemini/gemini-2.5-pro", 131072), 32768); + // A different model on the same provider keeps its own (registry) cap. + assert.equal(capThinkingBudget("openrouter/gemini-2.5-flash", 131072), 24576); // registry cap +}); diff --git a/tests/unit/gemini-thinking-budget-fallback.test.ts b/tests/unit/gemini-thinking-budget-fallback.test.ts new file mode 100644 index 000000000000..b0aae2ee38f0 --- /dev/null +++ b/tests/unit/gemini-thinking-budget-fallback.test.ts @@ -0,0 +1,338 @@ +import { test, after, beforeEach } from "node:test"; +import assert from "node:assert/strict"; +import { BaseExecutor } from "../../open-sse/executors/base.ts"; +import { openaiToGeminiRequest } from "../../open-sse/translator/request/openai-to-gemini.ts"; +import { + getLearnedThinkingCap, + __test_resetLearnedThinkingCaps, +} from "../../open-sse/services/learnedThinkingCaps.ts"; + +const GEMINI_400_BODY = JSON.stringify({ + error: { + code: 400, + message: + "* GenerateContentRequest.generation_config.thinking_config.thinking_budget: " + + "thinking_budget must be in the range [-1, 65535]", + status: "INVALID_ARGUMENT", + }, +}); + +// Passthrough executor: returns the body unchanged so we assert on exactly what +// base.ts sends upstream. +class SimpleExecutor extends BaseExecutor { + constructor() { + super("gemini", { + baseUrls: ["https://generativelanguage.example/v1/models"], + }); + } + async transformRequest(_model: string, body: Record) { + return { ...body }; + } +} + +beforeEach(() => { + __test_resetLearnedThinkingCaps(); +}); + +after(() => { + __test_resetLearnedThinkingCaps(); +}); + +// ── Reactive clamp-and-retry ─────────────────────────────────────────────── + +test("400 'thinking_budget must be in the range' clamps nested budget and retries once", async () => { + const executor = new SimpleExecutor(); + const originalFetch = globalThis.fetch; + const capturedBodies: Record[] = []; + + globalThis.fetch = async (_url: string | URL | Request, init: RequestInit = {}) => { + const body = JSON.parse(String(init.body)); + capturedBodies.push(body); + if (capturedBodies.length === 1) { + return new Response(GEMINI_400_BODY, { + status: 400, + headers: { "Content-Type": "application/json" }, + }); + } + return new Response(JSON.stringify({ ok: true }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + }; + + try { + const result = await executor.execute({ + model: "gemini-2.5-pro", + body: { + contents: [{ role: "user", parts: [{ text: "hi" }] }], + generationConfig: { thinkingConfig: { thinkingBudget: 131072 } }, + }, + stream: false, + credentials: {}, + }); + + assert.equal(capturedBodies.length, 2, "fetch should be called exactly twice"); + assert.equal( + capturedBodies[0].generationConfig.thinkingConfig.thinkingBudget, + 131072, + "first request sends the raw over-max budget" + ); + assert.equal( + capturedBodies[1].generationConfig.thinkingConfig.thinkingBudget, + 65535, + "retry clamps the budget to the upstream-advertised max" + ); + assert.equal(result.response.status, 200); + } finally { + globalThis.fetch = originalFetch; + } +}); + +test("the learned cap is recorded process-wide after a 400 (provider+model key)", async () => { + const executor = new SimpleExecutor(); + const originalFetch = globalThis.fetch; + + globalThis.fetch = async () => + new Response(GEMINI_400_BODY, { + status: 400, + headers: { "Content-Type": "application/json" }, + }); + + try { + await executor + .execute({ + model: "gemini-2.5-pro", + body: { generationConfig: { thinkingConfig: { thinkingBudget: 131072 } } }, + stream: false, + credentials: {}, + }) + .catch(() => {}); + } finally { + globalThis.fetch = originalFetch; + } + + assert.equal( + getLearnedThinkingCap("gemini", "gemini-2.5-pro"), + 32768, + "a cap must be learned so future requests clamp proactively (first step below 131072)" + ); +}); + +test("Antigravity envelope: budget nested under request.generationConfig is clamped too", async () => { + class AntigravityLikeExecutor extends BaseExecutor { + constructor() { + super("antigravity", { baseUrls: ["https://antigravity.example/v1"] }); + } + async transformRequest(_model: string, body: Record) { + return { ...body }; + } + } + const executor = new AntigravityLikeExecutor(); + const originalFetch = globalThis.fetch; + const capturedBodies: Record[] = []; + + globalThis.fetch = async (_url: string | URL | Request, init: RequestInit = {}) => { + const body = JSON.parse(String(init.body)); + capturedBodies.push(body); + if (capturedBodies.length === 1) { + return new Response(GEMINI_400_BODY, { + status: 400, + headers: { "Content-Type": "application/json" }, + }); + } + return new Response(JSON.stringify({ ok: true }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + }; + + try { + const result = await executor.execute({ + model: "gemini-3.1-pro", + body: { + request: { + contents: [{ role: "user", parts: [{ text: "hi" }] }], + generationConfig: { thinkingConfig: { thinkingBudget: 131072 } }, + }, + }, + stream: false, + credentials: {}, + }); + + assert.equal(capturedBodies.length, 2); + assert.equal( + capturedBodies[1].request.generationConfig.thinkingConfig.thinkingBudget, + 65535, + "retry clamps the Antigravity-envelope budget" + ); + assert.equal(result.response.status, 200); + } finally { + globalThis.fetch = originalFetch; + } +}); + +// ── Guards ───────────────────────────────────────────────────────────────── + +test("no retry when the body carries no thinking budget (nothing to clamp → would loop)", async () => { + const executor = new SimpleExecutor(); + const originalFetch = globalThis.fetch; + let callCount = 0; + + globalThis.fetch = async () => { + callCount++; + return new Response(GEMINI_400_BODY, { + status: 400, + headers: { "Content-Type": "application/json" }, + }); + }; + + try { + const result = await executor.execute({ + model: "gemini-2.5-pro", + body: { contents: [{ role: "user", parts: [{ text: "hi" }] }] }, + stream: false, + credentials: {}, + }); + assert.equal(callCount, 1, "no spurious retry when there is no budget to lower"); + assert.equal(result.response.status, 400); + } finally { + globalThis.fetch = originalFetch; + } +}); + +test("no retry when the budget is already at/below the upstream max", async () => { + const executor = new SimpleExecutor(); + const originalFetch = globalThis.fetch; + let callCount = 0; + + globalThis.fetch = async () => { + callCount++; + return new Response(GEMINI_400_BODY, { + status: 400, + headers: { "Content-Type": "application/json" }, + }); + }; + + try { + const result = await executor.execute({ + model: "gemini-2.5-pro", + body: { generationConfig: { thinkingConfig: { thinkingBudget: 8192 } } }, + stream: false, + credentials: {}, + }); + // 8192 < 65535 → clamping would not change the body → a retry would resend + // an identical request and 400 again. Guard must skip it. + assert.equal(callCount, 1, "no retry when clamping would not change the body"); + assert.equal(result.response.status, 400); + } finally { + globalThis.fetch = originalFetch; + } +}); + +test("no retry for a non-thinking-budget 400", async () => { + const executor = new SimpleExecutor(); + const originalFetch = globalThis.fetch; + let callCount = 0; + + globalThis.fetch = async () => { + callCount++; + return new Response(JSON.stringify({ error: "some unrelated upstream error" }), { + status: 400, + headers: { "Content-Type": "application/json" }, + }); + }; + + try { + const result = await executor.execute({ + model: "gemini-2.5-pro", + body: { generationConfig: { thinkingConfig: { thinkingBudget: 131072 } } }, + stream: false, + credentials: {}, + }); + assert.equal(callCount, 1, "unrelated 400 must not trigger the budget retry"); + assert.equal(result.response.status, 400); + } finally { + globalThis.fetch = originalFetch; + } +}); + +test("clamp fires at most once per execute() even if the retry also 400s", async () => { + const executor = new SimpleExecutor(); + const originalFetch = globalThis.fetch; + let callCount = 0; + + globalThis.fetch = async () => { + callCount++; + return new Response(GEMINI_400_BODY, { + status: 400, + headers: { "Content-Type": "application/json" }, + }); + }; + + try { + const result = await executor.execute({ + model: "gemini-2.5-pro", + body: { generationConfig: { thinkingConfig: { thinkingBudget: 131072 } } }, + stream: false, + credentials: {}, + }); + assert.equal(callCount, 2, "original + exactly one clamp retry, no loop"); + assert.equal(result.response.status, 400); + } finally { + globalThis.fetch = originalFetch; + } +}); + +// ── End-to-end: the user's actual scenario (xhigh on an unregistered model) ── + +test("E2E: reasoning_effort=xhigh on an unregistered Gemini model is clamped proactively — first request succeeds, no double round-trip", async () => { + // The user's report: Claude Code's xhigh/ultracode sends a huge budget that an + // unregistered Gemini model (e.g. gemini-2.5-pro, absent from MODEL_SPECS) + // used to forward verbatim → upstream 400 → failed request. The proactive + // gemini fallback in capThinkingBudget must clamp it BEFORE the first fetch so + // the request succeeds in a single round-trip (no 400→retry latency). + class RealTranslatorExecutor extends BaseExecutor { + constructor() { + super("gemini", { baseUrls: ["https://generativelanguage.example/v1/models"] }); + } + async transformRequest(model: string, body: Record) { + return openaiToGeminiRequest(model, body, false) as Record; + } + } + const executor = new RealTranslatorExecutor(); + const originalFetch = globalThis.fetch; + const capturedBodies: Record[] = []; + + globalThis.fetch = async (_url: string | URL | Request, init: RequestInit = {}) => { + const body = JSON.parse(String(init.body)); + capturedBodies.push(body); + // First request succeeds — the budget must already be within range. + return new Response(JSON.stringify({ ok: true }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + }; + + try { + const result = await executor.execute({ + model: "gemini-2.5-pro", // NOT in MODEL_SPECS + body: { + messages: [{ role: "user", content: "hi" }], + reasoning_effort: "xhigh", // would have sent 131072 before the fix + }, + stream: false, + credentials: {}, + }); + + assert.equal(capturedBodies.length, 1, "single request — no 400→retry double round-trip"); + const budget = capturedBodies[0].generationConfig.thinkingConfig.thinkingBudget; + assert.ok( + budget <= 32768, + `budget must be proactively clamped to the pro-tier ceiling, got ${budget}` + ); + assert.equal(budget, 32768); + assert.equal(result.response.status, 200); + } finally { + globalThis.fetch = originalFetch; + } +}); diff --git a/tests/unit/learned-thinking-caps.test.ts b/tests/unit/learned-thinking-caps.test.ts new file mode 100644 index 000000000000..e0cc3e1198ac --- /dev/null +++ b/tests/unit/learned-thinking-caps.test.ts @@ -0,0 +1,109 @@ +import { test, after, beforeEach } from "node:test"; +import assert from "node:assert/strict"; +import { + getLearnedThinkingCap, + recordLearnedThinkingCap, + parseThinkingBudgetMax, + __test_resetLearnedThinkingCaps, +} from "../../open-sse/services/learnedThinkingCaps.ts"; + +beforeEach(() => { + __test_resetLearnedThinkingCaps(); +}); + +after(() => { + __test_resetLearnedThinkingCaps(); +}); + +// ── parseThinkingBudgetMax ───────────────────────────────────────────────── + +test("parseThinkingBudgetMax extracts 65535 from the real Gemini error body", () => { + const err = + "* GenerateContentRequest.generation_config.thinking_config.thinking_budget: " + + "thinking_budget must be in the range [-1, 65535]"; + assert.equal(parseThinkingBudgetMax(err), 65535); +}); + +test("parseThinkingBudgetMax handles camelCase thinkingBudget variant", () => { + const err = "generationConfig.thinkingConfig.thinkingBudget must be in the range [-1, 24576]"; + assert.equal(parseThinkingBudgetMax(err), 24576); +}); + +test("parseThinkingBudgetMax returns null for unrelated 400 text", () => { + assert.equal(parseThinkingBudgetMax("some random upstream error"), null); + assert.equal(parseThinkingBudgetMax(""), null); + assert.equal(parseThinkingBudgetMax(null), null); + assert.equal(parseThinkingBudgetMax(undefined), null); +}); + +test("parseThinkingBudgetMax returns null when thinking_budget mentioned but no range", () => { + assert.equal(parseThinkingBudgetMax("thinking_budget is deprecated"), null); +}); + +// ── recordLearnedThinkingCap step ladder ─────────────────────────────────── + +test("recordLearnedThinkingCap steps 131072 → 32768 on first failure", () => { + const next = recordLearnedThinkingCap("gemini", "gemini-2.5-pro", 131072); + assert.equal(next, 32768); + assert.equal(getLearnedThinkingCap("gemini", "gemini-2.5-pro"), 32768); +}); + +test("recordLearnedThinkingCap steps 32768 → 24576 on second failure", () => { + recordLearnedThinkingCap("gemini", "gemini-2.5-pro", 131072); + const next = recordLearnedThinkingCap("gemini", "gemini-2.5-pro", 32768); + assert.equal(next, 24576); + assert.equal(getLearnedThinkingCap("gemini", "gemini-2.5-pro"), 24576); +}); + +test("recordLearnedThinkingCap steps 24576 → 8192 on third failure", () => { + recordLearnedThinkingCap("gemini", "gemini-2.5-pro", 131072); + recordLearnedThinkingCap("gemini", "gemini-2.5-pro", 32768); + const next = recordLearnedThinkingCap("gemini", "gemini-2.5-pro", 24576); + assert.equal(next, 8192); + assert.equal(getLearnedThinkingCap("gemini", "gemini-2.5-pro"), 8192); +}); + +test("recordLearnedThinkingCap returns null when steps exhausted (8192 failed)", () => { + recordLearnedThinkingCap("gemini", "gemini-2.5-pro", 131072); + recordLearnedThinkingCap("gemini", "gemini-2.5-pro", 32768); + recordLearnedThinkingCap("gemini", "gemini-2.5-pro", 24576); + const next = recordLearnedThinkingCap("gemini", "gemini-2.5-pro", 8192); + assert.equal(next, null); +}); + +test("recordLearnedThinkingCap picks the first step strictly below failedBudget", () => { + // A budget of 30000 failing should skip 32768 (not below) and land on 24576. + const next = recordLearnedThinkingCap("gemini", "gemini-3.1-pro", 30000); + assert.equal(next, 24576); +}); + +// ── getLearnedThinkingCap key handling ───────────────────────────────────── + +test("getLearnedThinkingCap returns null for unknown provider+model", () => { + assert.equal(getLearnedThinkingCap("gemini", "unknown-model"), null); +}); + +test("getLearnedThinkingCap is keyed case-insensitively on provider+model", () => { + recordLearnedThinkingCap("Gemini", "Gemini-2.5-Pro", 131072); + assert.equal(getLearnedThinkingCap("gemini", "gemini-2.5-pro"), 32768); + assert.equal(getLearnedThinkingCap("GEMINI", "GEMINI-2.5-PRO"), 32768); +}); + +test("different providers for the same model id have independent caps", () => { + recordLearnedThinkingCap("gemini", "gemini-2.5-pro", 131072); + assert.equal(getLearnedThinkingCap("openrouter", "gemini-2.5-pro"), null); +}); + +test("different models for the same provider have independent caps", () => { + recordLearnedThinkingCap("gemini", "gemini-2.5-pro", 131072); + assert.equal(getLearnedThinkingCap("gemini", "gemini-2.5-flash"), null); +}); + +test("handles empty/null provider or model gracefully", () => { + assert.equal(getLearnedThinkingCap("", "gemini-2.5-pro"), null); + assert.equal(getLearnedThinkingCap("gemini", ""), null); + assert.equal(getLearnedThinkingCap(null, "gemini-2.5-pro"), null); + assert.equal(getLearnedThinkingCap("gemini", null), null); + assert.equal(recordLearnedThinkingCap("", "gemini-2.5-pro", 131072), null); + assert.equal(recordLearnedThinkingCap("gemini", "", 131072), null); +}); diff --git a/tests/unit/translator-claude-to-gemini.test.ts b/tests/unit/translator-claude-to-gemini.test.ts index 72a453e9e10e..d82e6b427630 100644 --- a/tests/unit/translator-claude-to-gemini.test.ts +++ b/tests/unit/translator-claude-to-gemini.test.ts @@ -239,12 +239,18 @@ test("Claude -> Gemini handles empty bodies without producing invalid content", }); test("Claude -> Gemini maps output_config.effort to thinkingConfig when thinking absent", () => { + // NOTE: max/xhigh previously asserted 131072, but that locked in the OLD + // no-cap behavior — gemini-2.5-pro is unregistered, so the raw budget sailed + // to the upstream and 400'd ("thinking_budget must be in the range"). The + // gemini-substring fallback in capThinkingBudget now clamps unregistered + // Gemini models to the pro-tier ceiling 32768, which is what the upstream + // actually accepts. low/medium/high are below the cap and stay unchanged. const cases: Array<{ effort: string; expected: number }> = [ { effort: "low", expected: 1024 }, { effort: "medium", expected: 10240 }, { effort: "high", expected: 32768 }, - { effort: "max", expected: 131072 }, - { effort: "xhigh", expected: 131072 }, + { effort: "max", expected: 32768 }, + { effort: "xhigh", expected: 32768 }, ]; for (const { effort, expected } of cases) {