fix: declare long-context Claude and GPT metadata

(cherry picked from commit a6359d3)

Author: dgokeeffe

src/ucode/agents/claude.py

@@ -251,8 +251,12 @@ def _maybe_add_1m_suffix(model: str) -> str:
     family, major_raw, minor_raw, _ = match.groups()
     major = int(major_raw)
     minor = int(minor_raw)
+    # 1M-token context window support: Opus from 4.6, Sonnet from 4.5 (Sonnet
+    # 4.5 gained the 1M beta window before Opus did). Without the `[1m]`
+    # suffix the gateway serves these models with the default 200k context, so
+    # gate each family at its own floor rather than a shared version.
     should_suffix = (family == "opus" and (major, minor) >= (4, 6)) or (
-        family == "sonnet" and (major, minor) >= (4, 6)
+        family == "sonnet" and (major, minor) >= (4, 5)
     )
     return f"{model}[1m]" if should_suffix else model
 

src/ucode/agents/pi.py

@@ -38,6 +38,7 @@
 from __future__ import annotations
 
 import os
+import re
 import signal
 import subprocess
 import threading
@@ -55,7 +56,9 @@
 from ucode.databricks import (
     TOKEN_REFRESH_INTERVAL_SECONDS,
     build_pi_base_urls,
+    claude_model_token_limits,
     get_databricks_token,
+    gpt_model_token_limits,
     model_is_reasoning,
     model_token_limits,
 )
@@ -117,6 +120,26 @@ def _resolve_model_selector(
     return model
 
 
+def _pi_claude_model_entry(model_id: str) -> dict:
+    """Build a Claude entry with explicit limits.
+
+    Databricks model ids do not match Pi's built-in Anthropic ids, so a bare
+    custom entry silently gets Pi's 128k context / 4k output defaults.
+    """
+    limits = claude_model_token_limits(model_id)
+    entry: dict = {
+        "id": model_id,
+        "reasoning": True,
+        "input": ["text", "image"],
+        "contextWindow": limits["context"],
+        "maxTokens": limits["output"],
+    }
+    normalized = model_id.lower().replace(".", "-")
+    if re.search(r"claude-(?:opus|sonnet)-4-(?:[6-9]|[1-9]\d)", normalized):
+        entry["compat"] = {"forceAdaptiveThinking": True}
+    return entry
+
+
 def _pi_oss_model_entry(model_id: str) -> dict:
     """Build a Pi mlflow model entry enriched from the shared limits/reasoning
     tables: `reasoning:true` for reasoning models (Pi renders their streamed
@@ -135,6 +158,23 @@ def _pi_oss_model_entry(model_id: str) -> dict:
     return entry
 
 
+def _pi_gpt_model_entry(model_id: str) -> dict:
+    """Build a Pi openai (codex) model entry with `contextWindow`/`maxTokens`
+    from `databricks.gpt_model_token_limits`. GPT ids aren't in Pi's built-in
+    catalog, so without an explicit window Pi falls back to a small default and
+    truncates long sessions."""
+    limits = gpt_model_token_limits(model_id)
+    entry: dict = {
+        "id": model_id,
+        "contextWindow": limits["context"],
+        "maxTokens": limits["output"],
+    }
+    if "gpt-5" in model_id.lower().replace(".", "-"):
+        entry["reasoning"] = True
+        entry["input"] = ["text", "image"]
+    return entry
+
+
 def render_overlay(
     model: str,
     token: str,
@@ -163,7 +203,7 @@ def render_overlay(
             # the legacy beta header instead when this is false.
             "compat": {"supportsEagerToolInputStreaming": False},
             "headers": ua_headers,
-            "models": [{"id": m} for m in claude_ids],
+            "models": [_pi_claude_model_entry(m) for m in claude_ids],
         }
         keys.append(["providers", "databricks-claude"])
     if codex_models:
@@ -173,7 +213,7 @@ def render_overlay(
             "apiKey": token,
             "authHeader": True,
             "headers": ua_headers,
-            "models": [{"id": m} for m in codex_models],
+            "models": [_pi_gpt_model_entry(m) for m in codex_models],
         }
         keys.append(["providers", "databricks-openai"])
     if gemini_models:

src/ucode/databricks.py

@@ -1226,6 +1226,76 @@ def model_token_limits(model_id: str) -> dict[str, int] | None:
     return None
 
 
+# Pi treats every custom model without explicit metadata as 128k context / 4k
+# output. Gateway ids are custom ids (not Pi's built-ins), so preserve the
+# upstream windows explicitly. Entries are ordered most-specific first after
+# normalizing dotted OpenAI ids and hyphenated Databricks ids to one form.
+# GPT-5.6 Sol/Terra/Luna support the opt-in 1.05M window; 272k is merely Pi's
+# built-in short-context pricing default, not the model's hard context limit.
+_GPT_TOKEN_LIMITS: tuple[tuple[str, dict[str, int]], ...] = (
+    ("gpt-5-6-sol", {"context": 1_050_000, "output": 128_000}),
+    ("gpt-5-6-terra", {"context": 1_050_000, "output": 128_000}),
+    ("gpt-5-6-luna", {"context": 1_050_000, "output": 128_000}),
+    ("gpt-5-5-pro", {"context": 1_050_000, "output": 128_000}),
+    ("gpt-5-4-pro", {"context": 1_050_000, "output": 128_000}),
+    ("gpt-5-5", {"context": 272_000, "output": 128_000}),
+    ("gpt-5-4-mini", {"context": 400_000, "output": 128_000}),
+    ("gpt-5-4-nano", {"context": 400_000, "output": 128_000}),
+    ("gpt-5-4", {"context": 272_000, "output": 128_000}),
+    ("gpt-5", {"context": 400_000, "output": 128_000}),
+    ("gpt-4-1", {"context": 1_047_576, "output": 32_768}),
+    ("gpt-4o", {"context": 128_000, "output": 16_384}),
+    ("gpt-4-turbo", {"context": 128_000, "output": 4_096}),
+    ("gpt-4", {"context": 8_192, "output": 8_192}),
+)
+_GPT_FALLBACK_LIMITS = {"context": 128_000, "output": 16_384}
+
+
+def _normalized_foundation_model_id(model_id: str) -> str:
+    """Strip route prefixes and normalize dotted versions to hyphens."""
+    tail = model_id.split("/")[-1]
+    if tail.startswith("system.ai."):
+        tail = tail[len("system.ai.") :]
+    if tail.startswith("databricks-"):
+        tail = tail[len("databricks-") :]
+    return tail.lower().replace(".", "-")
+
+
+def gpt_model_token_limits(model_id: str) -> dict[str, int]:
+    """Return Pi metadata limits for a GPT (codex/openai) gateway model."""
+    tail = _normalized_foundation_model_id(model_id)
+    for family, limits in _GPT_TOKEN_LIMITS:
+        if tail.startswith(family):
+            return dict(limits)
+    return dict(_GPT_FALLBACK_LIMITS)
+
+
+_CLAUDE_FALLBACK_LIMITS = {"context": 200_000, "output": 64_000}
+
+
+def claude_model_token_limits(model_id: str) -> dict[str, int]:
+    """Return Pi metadata limits for a Claude gateway model.
+
+    Claude gateway ids are custom to Pi, so omitting these fields silently
+    applies Pi's 128k/4k custom-model defaults. Current 1M families mirror
+    Anthropic's model metadata: Opus >=4.6 and Sonnet >=4.5. Sonnet 4.5 has a
+    64k output cap; later 1M models have 128k output.
+    """
+    tail = _normalized_foundation_model_id(model_id)
+    match = re.match(r"claude-(opus|sonnet|haiku)-(\d+)-(\d+)", tail)
+    if not match:
+        return dict(_CLAUDE_FALLBACK_LIMITS)
+    family, major_raw, minor_raw = match.groups()
+    version = (int(major_raw), int(minor_raw))
+    if family == "opus" and version >= (4, 6):
+        return {"context": 1_000_000, "output": 128_000}
+    if family == "sonnet" and version >= (4, 6):
+        return {"context": 1_000_000, "output": 128_000}
+    if family == "sonnet" and version >= (4, 5):
+        return {"context": 1_000_000, "output": 64_000}
+    return dict(_CLAUDE_FALLBACK_LIMITS)
+
+
 def _model_service_id(service: dict) -> str | None:
     """Extract the `system.ai.<model-name>` id from one model-service entry.
 

tests/test_agent_claude.py

@@ -48,6 +48,30 @@ def test_adds_1m_suffix_for_sonnet_4_6_and_later(self):
             overlay["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "databricks-claude-sonnet-4-7[1m]"
         )
 
+    def test_adds_1m_suffix_for_sonnet_4_5(self):
+        # Sonnet 4.5 supports the 1M context window (its 1M beta shipped before
+        # Opus's), so it must get the [1m] suffix even though it predates the
+        # Opus 4.6 floor.
+        overlay, _ = claude.render_overlay(
+            WS, "s4", claude_models={"sonnet": "databricks-claude-sonnet-4-5"}
+        )
+        assert (
+            overlay["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "databricks-claude-sonnet-4-5[1m]"
+        )
+
+    def test_does_not_add_1m_suffix_for_sonnet_4_4(self):
+        overlay, _ = claude.render_overlay(
+            WS, "s4", claude_models={"sonnet": "databricks-claude-sonnet-4-4"}
+        )
+        assert overlay["env"]["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "databricks-claude-sonnet-4-4"
+
+    def test_does_not_add_1m_suffix_for_opus_4_5(self):
+        # Opus's 1M window starts at 4.6, so 4.5 stays on the default context.
+        overlay, _ = claude.render_overlay(
+            WS, "s4", claude_models={"opus": "databricks-claude-opus-4-5"}
+        )
+        assert overlay["env"]["ANTHROPIC_DEFAULT_OPUS_MODEL"] == "databricks-claude-opus-4-5"
+
     def test_does_not_add_1m_suffix_for_haiku(self):
         overlay, _ = claude.render_overlay(
             WS, "s4", claude_models={"haiku": "databricks-claude-haiku-4-6"}

tests/test_agent_pi.py

@@ -78,6 +78,55 @@ def test_openai_provider_uses_openai_responses(self):
         assert provider["api"] == "openai-responses"
         assert provider["baseUrl"] == f"{WS}/ai-gateway/codex/v1"
 
+    def test_gpt56_sol_model_entry_pins_1m_context(self):
+        # Gateway ids are custom to Pi, so explicit metadata is required to
+        # avoid its 128k custom-model default.
+        overlay, _ = _overlay("gpt-5-6-sol", codex_models=["gpt-5-6-sol"])
+        entry = overlay["providers"]["databricks-openai"]["models"][0]
+        assert entry["id"] == "gpt-5-6-sol"
+        assert entry["contextWindow"] == 1_050_000
+        assert entry["maxTokens"] == 128_000
+        assert entry["reasoning"] is True
+        assert entry["input"] == ["text", "image"]
+
+    def test_gpt_model_entries_use_model_specific_windows(self):
+        overlay, _ = _overlay(
+            "system.ai.gpt-5-2",
+            codex_models=[
+                "system.ai.gpt-5-2",
+                "databricks-gpt-5-4-nano",
+                "databricks-gpt-5-6-sol",
+            ],
+        )
+        windows = {
+            m["id"]: m["contextWindow"] for m in overlay["providers"]["databricks-openai"]["models"]
+        }
+        assert windows == {
+            "system.ai.gpt-5-2": 400_000,
+            "databricks-gpt-5-4-nano": 400_000,
+            "databricks-gpt-5-6-sol": 1_050_000,
+        }
+
+    def test_claude_entries_pin_limits_and_capabilities(self):
+        overlay, _ = _overlay(
+            "databricks-claude-opus-4-8",
+            claude_models={
+                "opus": "databricks-claude-opus-4-8",
+                "sonnet": "system.ai.claude-sonnet-4-5",
+                "haiku": "databricks-claude-haiku-4-5",
+            },
+        )
+        entries = {m["id"]: m for m in overlay["providers"]["databricks-claude"]["models"]}
+        opus = entries["databricks-claude-opus-4-8"]
+        assert opus["contextWindow"] == 1_000_000
+        assert opus["maxTokens"] == 128_000
+        assert opus["reasoning"] is True
+        assert opus["input"] == ["text", "image"]
+        assert opus["compat"] == {"forceAdaptiveThinking": True}
+        assert entries["system.ai.claude-sonnet-4-5"]["contextWindow"] == 1_000_000
+        assert entries["system.ai.claude-sonnet-4-5"]["maxTokens"] == 64_000
+        assert entries["databricks-claude-haiku-4-5"]["contextWindow"] == 200_000
+
     def test_gemini_provider_uses_google_generative_ai(self):
         overlay, _ = _overlay("gemini-2", gemini_models=["gemini-2"])
         provider = overlay["providers"]["databricks-gemini"]

tests/test_databricks.py

@@ -202,6 +202,60 @@ def test_embedding_model_returns_none_not_fallback(self):
         assert db_mod.model_token_limits("system.ai.qwen3-embedding-0-6b") is None
 
 
+class TestGptModelTokenLimits:
+    def test_gpt56_sol_serves_1m_context_across_id_forms(self):
+        for model_id in (
+            "gpt-5.6-sol",
+            "system.ai.gpt-5-6-sol",
+            "databricks-gpt-5-6-sol",
+            "databricks-openai/gpt-5.6-sol",
+        ):
+            assert db_mod.gpt_model_token_limits(model_id) == {
+                "context": 1_050_000,
+                "output": 128_000,
+            }
+
+    def test_gpt5_windows_are_model_specific(self):
+        assert db_mod.gpt_model_token_limits("system.ai.gpt-5-2")["context"] == 400_000
+        assert db_mod.gpt_model_token_limits("databricks-gpt-5-4")["context"] == 272_000
+        assert db_mod.gpt_model_token_limits("databricks-gpt-5-4-nano")["context"] == 400_000
+        assert db_mod.gpt_model_token_limits("gpt-5.5-pro")["context"] == 1_050_000
+
+    def test_gpt41_window(self):
+        assert db_mod.gpt_model_token_limits("databricks-gpt-4-1") == {
+            "context": 1_047_576,
+            "output": 32_768,
+        }
+
+    def test_unknown_gpt_uses_conservative_floor(self):
+        assert db_mod.gpt_model_token_limits("gpt-6-turbo") == {
+            "context": 128_000,
+            "output": 16_384,
+        }
+
+
+class TestClaudeModelTokenLimits:
+    def test_1m_models_and_output_caps(self):
+        assert db_mod.claude_model_token_limits("databricks-claude-opus-4-8") == {
+            "context": 1_000_000,
+            "output": 128_000,
+        }
+        assert db_mod.claude_model_token_limits("system.ai.claude-sonnet-4-5") == {
+            "context": 1_000_000,
+            "output": 64_000,
+        }
+        assert db_mod.claude_model_token_limits("claude-sonnet-4-6[1m]") == {
+            "context": 1_000_000,
+            "output": 128_000,
+        }
+
+    def test_older_and_unknown_claude_use_200k_floor(self):
+        expected = {"context": 200_000, "output": 64_000}
+        assert db_mod.claude_model_token_limits("claude-opus-4-5") == expected
+        assert db_mod.claude_model_token_limits("claude-haiku-4-5") == expected
+        assert db_mod.claude_model_token_limits("claude-future") == expected
+
+
 class TestModelIsReasoning:
     def test_reasoning_families(self):
         assert db_mod.model_is_reasoning("system.ai.glm-5-2") is True