fix(sse): clamp out-of-range Gemini thinking_budget and learn the cap

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<string, unknown>): 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<Record<string, unknown>> {
+  if (!body || typeof body !== "object") return [];
+  const root = body as Record<string, unknown>;
+  const configs: Array<Record<string, unknown>> = [];
+  const envelopes: unknown[] = [root.generationConfig, (root.request as Record<string, unknown> | undefined)?.generationConfig];
+  for (const env of envelopes) {
+    if (!env || typeof env !== "object") continue;
+    const tc = (env as Record<string, unknown>).thinkingConfig;
+    if (tc && typeof tc === "object") {
+      const tcr = tc as Record<string, unknown>;
+      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<string>();
+    // 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<string, unknown>
         ) 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 &&

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<string, number>();
+
+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();
+}

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(