From ccee76b31eb4e7c61c36595fdf5e9581c4547c1a Mon Sep 17 00:00:00 2001
From: Jeroen Ubbink <jeroen@blackbyte.nl>
Date: Sun, 5 Jul 2026 15:39:41 +0200
Subject: [PATCH] fix(upstream): source litellm context from max_input_tokens
 only

litellm's max_tokens is the completion cap (it equals max_output_tokens for
~94% of models), not the context window, so falling back to it overstated the
context as the output limit. Take context from max_input_tokens alone; a model
that reports none falls through to the id-based estimate downstream, which is
honest about being a guess rather than mislabelling the output cap.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 routstr/upstream/pricing_resolver.py |  6 ++++-
 tests/unit/test_upstream_generic.py  | 34 ++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/routstr/upstream/pricing_resolver.py b/routstr/upstream/pricing_resolver.py
index 80b27fc8..f338c12e 100644
--- a/routstr/upstream/pricing_resolver.py
+++ b/routstr/upstream/pricing_resolver.py
@@ -103,7 +103,11 @@ def _from_litellm(model_id: str) -> ResolvedPricing | None:
     return ResolvedPricing(
         prompt=float(prompt),
         completion=float(completion),
-        context_length=_as_int(info.get("max_input_tokens") or info.get("max_tokens")),
+        # max_input_tokens is the context window; max_tokens is litellm's
+        # completion cap (it tracks max_output_tokens for ~94% of models), so
+        # it is never a context source. A missing window falls to the id-based
+        # estimate downstream rather than borrowing the output cap.
+        context_length=_as_int(info.get("max_input_tokens")),
         source="litellm",
         max_completion_tokens=_as_int(info.get("max_output_tokens")),
         input_cache_read=float(info.get("cache_read_input_token_cost") or 0.0),
diff --git a/tests/unit/test_upstream_generic.py b/tests/unit/test_upstream_generic.py
index edc448a6..28bec0ff 100644
--- a/tests/unit/test_upstream_generic.py
+++ b/tests/unit/test_upstream_generic.py
@@ -184,6 +184,40 @@ async def test_litellm_zero_price_entry_fails_closed(
     )
 
 
+@pytest.mark.asyncio
+async def test_litellm_output_cap_not_used_as_context() -> None:
+    """litellm's ``max_tokens`` is the completion cap, not the context window
+    (it tracks ``max_output_tokens`` for ~94% of models). When a model reports
+    no ``max_input_tokens``, the resolver must not smuggle the output cap in as
+    the context window; it falls back to the id-based estimate instead, while
+    ``max_tokens`` still feeds the completion limit."""
+    payload = {
+        "data": [
+            {
+                "id": "gemini/gemini-gemma-2-9b-it",
+                "object": "model",
+                "owned_by": "google",
+            },
+        ]
+    }
+
+    with _patch_models_endpoint(payload):
+        or_feed = AsyncMock(return_value=[])
+        with patch(
+            "routstr.payment.models.async_fetch_openrouter_models", or_feed
+        ):
+            models = await GenericUpstreamProvider(base_url="http://x").fetch_models()
+
+    model = _model_by_id(models, "gemini/gemini-gemma-2-9b-it")
+    assert model.enabled is True
+    # litellm gives this model max_input_tokens=None, max_tokens=8192 (an output
+    # cap). Context must come from the estimate (4096), never the 8192 cap.
+    assert model.context_length == 4096
+    # The 8192 output cap still lands where it belongs: the completion limit.
+    assert model.top_provider is not None
+    assert model.top_provider.max_completion_tokens == 8192
+
+
 # ---------------------------------------------------------------------------
 # OpenRouter fallback — litellm misses, OR carries a full payload
 # ---------------------------------------------------------------------------