fix: share Claude context capability policy

Centralize Claude family/version parsing so Pi metadata, adaptive-thinking compatibility, and Claude Code's [1m] selector cannot drift. Cover Sonnet 4.5, Opus 4.6, future major versions, Fable fallback, and prefixed model IDs.

Author: dgokeeffe

src/ucode/agents/claude.py

@@ -23,6 +23,7 @@
 from ucode.databricks import (
     build_auth_shell_command,
     build_tool_base_url,
+    claude_model_supports_1m,
     get_databricks_token,
 )
 from ucode.launcher import exec_or_spawn
@@ -65,11 +66,6 @@ def _resolve_web_search_model(state: dict) -> str | None:
 
 
 WEB_SEARCH_MCP_NAME = "web_search"
-# Matches both the AI Gateway form (`databricks-claude-opus-4-8`) and the UC
-# model-services form (`system.ai.claude-opus-4-8`).
-_CLAUDE_MODEL_RE = re.compile(
-    r"^(?:system\.ai\.)?(?:databricks-)?claude-(opus|sonnet)-(\d+)-(\d+)(.*)$"
-)
 
 # Env keys the MLflow Stop hook reads to route traces. Written into the
 # settings `env` block alongside the hook itself.
@@ -242,19 +238,9 @@ def render_overlay(
 
 
 def _maybe_add_1m_suffix(model: str) -> str:
-    if model.endswith("[1m]"):
-        return model
-    match = _CLAUDE_MODEL_RE.match(model)
-    if not match:
+    if model.endswith("[1m]") or not claude_model_supports_1m(model):
         return model
-
-    family, major_raw, minor_raw, _ = match.groups()
-    major = int(major_raw)
-    minor = int(minor_raw)
-    should_suffix = (family == "opus" and (major, minor) >= (4, 6)) or (
-        family == "sonnet" and (major, minor) >= (4, 6)
-    )
-    return f"{model}[1m]" if should_suffix else model
+    return f"{model}[1m]"
 
 
 def _register_web_search_mcp(workspace: str, search_model: str, profile: str | None = None) -> bool:

src/ucode/agents/pi.py

@@ -51,7 +51,9 @@
 from ucode.databricks import (
     TOKEN_REFRESH_INTERVAL_SECONDS,
     build_pi_base_urls,
+    claude_model_capabilities,
     get_databricks_token,
+    gpt_model_token_limits,
     model_is_reasoning,
     model_token_limits,
 )
@@ -113,6 +115,25 @@ 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.
+    """
+    capabilities = claude_model_capabilities(model_id)
+    entry: dict = {
+        "id": model_id,
+        "reasoning": True,
+        "input": ["text", "image"],
+        "contextWindow": capabilities.context,
+        "maxTokens": capabilities.output,
+    }
+    if capabilities.force_adaptive_thinking:
+        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
@@ -131,6 +152,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,
@@ -159,7 +197,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:
@@ -169,7 +207,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

@@ -22,6 +22,7 @@
 from concurrent.futures import (
     TimeoutError as FutureTimeoutError,
 )
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Literal, cast, overload
 from urllib import error as urllib_error
@@ -1232,6 +1233,106 @@ 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)
+
+
+@dataclass(frozen=True)
+class ClaudeModelCapabilities:
+    context: int
+    output: int
+    supports_1m: bool = False
+    force_adaptive_thinking: bool = False
+
+
+_CLAUDE_FALLBACK_CAPABILITIES = ClaudeModelCapabilities(context=200_000, output=64_000)
+_CLAUDE_MODEL_RE = re.compile(r"^claude-(fable|opus|sonnet|haiku)-(\d+)(?:-(\d+))?")
+
+
+def claude_model_capabilities(model_id: str) -> ClaudeModelCapabilities:
+    """Return the shared Claude capability policy for every agent.
+
+    Opus gained the opt-in 1M window in 4.6; Sonnet gained it in 4.5.
+    Sonnet 4.5 retains a 64k output cap, while later 1M tiers use 128k.
+    Opus 4.5, Haiku, Fable, and unrecognized ids intentionally use the
+    conservative 200k fallback until separately verified.
+    """
+    tail = _normalized_foundation_model_id(model_id)
+    match = _CLAUDE_MODEL_RE.match(tail)
+    if not match:
+        return _CLAUDE_FALLBACK_CAPABILITIES
+    family, major_raw, minor_raw = match.groups()
+    version = (int(major_raw), int(minor_raw or 0))
+    if family == "opus" and version >= (4, 6):
+        return ClaudeModelCapabilities(
+            context=1_000_000,
+            output=128_000,
+            supports_1m=True,
+            force_adaptive_thinking=True,
+        )
+    if family == "sonnet" and version >= (4, 6):
+        return ClaudeModelCapabilities(
+            context=1_000_000,
+            output=128_000,
+            supports_1m=True,
+            force_adaptive_thinking=True,
+        )
+    if family == "sonnet" and version >= (4, 5):
+        return ClaudeModelCapabilities(context=1_000_000, output=64_000, supports_1m=True)
+    return _CLAUDE_FALLBACK_CAPABILITIES
+
+
+def claude_model_supports_1m(model_id: str) -> bool:
+    """Whether Claude Code should request the model's opt-in ``[1m]`` tier."""
+    return claude_model_capabilities(model_id).supports_1m
+
+
+def claude_model_token_limits(model_id: str) -> dict[str, int]:
+    """Return Pi metadata limits from the shared Claude capability policy."""
+    capabilities = claude_model_capabilities(model_id)
+    return {"context": capabilities.context, "output": capabilities.output}
+
+
 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"}
@@ -72,6 +96,19 @@ def test_no_1m_suffix_for_model_services_haiku(self):
         )
         assert overlay["env"]["ANTHROPIC_DEFAULT_HAIKU_MODEL"] == "system.ai.claude-haiku-4-6"
 
+    @pytest.mark.parametrize(
+        ("model_id", "expected"),
+        [
+            ("system.ai.claude-opus-5", "system.ai.claude-opus-5[1m]"),
+            ("databricks-claude-sonnet-5", "databricks-claude-sonnet-5[1m]"),
+            ("system.ai.claude-opus-4-5", "system.ai.claude-opus-4-5"),
+            ("system.ai.claude-fable-5", "system.ai.claude-fable-5"),
+            ("not-a-claude-model", "not-a-claude-model"),
+        ],
+    )
+    def test_suffix_uses_shared_capability_policy(self, model_id, expected):
+        assert claude._maybe_add_1m_suffix(model_id) == expected
+
     def test_sets_anthropic_base_url(self):
         overlay, _ = claude.render_overlay(WS, "s4")
         assert overlay["env"]["ANTHROPIC_BASE_URL"] == f"{WS}/ai-gateway/anthropic"

tests/test_agent_pi.py

@@ -78,6 +78,57 @@ 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",
+                "fable": "system.ai.claude-fable-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
+        assert entries["system.ai.claude-fable-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"]