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,9 +51,12 @@
 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,
+    preferred_gpt_model,
 )
 from ucode.state import mark_tool_managed, save_state
 from ucode.telemetry import agent_version, ucode_version
@@ -113,6 +116,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 +153,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 +198,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 +208,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:
@@ -259,9 +298,9 @@ def default_model(state: dict) -> str | None:
     for family in ("opus", "sonnet", "haiku"):
         if claude_models.get(family):
             return claude_models[family]
-    codex_models = state.get("codex_models") or []
-    if codex_models:
-        return codex_models[0]
+    codex_model = preferred_gpt_model(state.get("codex_models") or [])
+    if codex_model:
+        return codex_model
     gemini_models = state.get("gemini_models") or []
     if gemini_models:
         return gemini_models[0]

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
@@ -1235,6 +1236,136 @@ 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 case-insensitively and normalize dotted versions."""
+    tail = model_id.split("/")[-1].lower()
+    if tail.startswith("system.ai."):
+        tail = tail[len("system.ai.") :]
+    if tail.startswith("databricks-"):
+        tail = tail[len("databricks-") :]
+    return tail.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 == family or tail.startswith(f"{family}-"):
+            return dict(limits)
+    return dict(_GPT_FALLBACK_LIMITS)
+
+
+def preferred_gpt_model(model_ids: list[str]) -> str | None:
+    """Prefer the newest numeric GPT id, then any other Responses endpoint.
+
+    ``gpt-oss`` is chat-completions-only and must never be selected for the
+    Responses route, even if stale state supplies it here.
+    """
+    eligible = [
+        model_id
+        for model_id in model_ids
+        if not _normalized_foundation_model_id(model_id).startswith("gpt-oss")
+    ]
+    numeric_gpt = [
+        model_id
+        for model_id in eligible
+        if re.match(r"^gpt-\d(?:-|$)", _normalized_foundation_model_id(model_id))
+    ]
+    if numeric_gpt:
+        return min(
+            numeric_gpt,
+            key=lambda model_id: model_version_sort_key(_normalized_foundation_model_id(model_id)),
+        )
+    return eligible[0] if eligible else None
+
+
+@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's verified 1M tiers retain a conservative 64k output cap. Fable 5
+    is 1M by default (so it needs no ``[1m]`` suffix) with a 128k output cap.
+    Opus 4.5, Haiku, and unrecognized ids use the conservative 200k fallback.
+    """
+    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=64_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)
+    if family == "fable" and version >= (5, 0):
+        return ClaudeModelCapabilities(
+            context=1_000_000,
+            output=128_000,
+            force_adaptive_thinking=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"