Append [1m] to Opus >=4.6 and Sonnet 4.6 model for 1m context length support (#118)

* claude: append [1m] to Opus >=4.6 and Sonnet 4.6 model ids when writing settings.json; update tests

* claude: allow opus 4.8 discovery

* chore: format databricks

* claude: apply 1m suffix to future sonnet models

Author: AarushiShah-db

src/ucode/agents/claude.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import os
+import re
 import shutil
 from pathlib import Path
 
@@ -58,6 +59,7 @@ def _resolve_web_search_model(state: dict) -> str | None:
 
 
 WEB_SEARCH_MCP_NAME = "web_search"
+_CLAUDE_MODEL_RE = re.compile(r"^databricks-claude-(opus|sonnet)-(\d+)-(\d+)(.*)$")
 
 
 def _web_search_mcp_entry(workspace: str, search_model: str, profile: str | None = None) -> dict:
@@ -103,17 +105,17 @@ def render_overlay(
         ]
     )
     env: dict[str, str] = {
-        "ANTHROPIC_MODEL": model,
+        "ANTHROPIC_MODEL": _maybe_add_1m_suffix(model),
         "ANTHROPIC_BASE_URL": base_url,
         "ANTHROPIC_CUSTOM_HEADERS": custom_headers,
         "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1",
         "CLAUDE_CODE_API_KEY_HELPER_TTL_MS": "900000",
     }
     if claude_models:
         if claude_models.get("opus"):
-            env["ANTHROPIC_DEFAULT_OPUS_MODEL"] = claude_models["opus"]
+            env["ANTHROPIC_DEFAULT_OPUS_MODEL"] = _maybe_add_1m_suffix(claude_models["opus"])
         if claude_models.get("sonnet"):
-            env["ANTHROPIC_DEFAULT_SONNET_MODEL"] = claude_models["sonnet"]
+            env["ANTHROPIC_DEFAULT_SONNET_MODEL"] = _maybe_add_1m_suffix(claude_models["sonnet"])
         if claude_models.get("haiku"):
             env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] = claude_models["haiku"]
     overlay: dict = {"apiKeyHelper": build_auth_shell_command(workspace, profile), "env": env}
@@ -129,6 +131,22 @@ def render_overlay(
     return overlay, keys
 
 
+def _maybe_add_1m_suffix(model: str) -> str:
+    if model.endswith("[1m]"):
+        return model
+    match = _CLAUDE_MODEL_RE.match(model)
+    if not match:
+        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
+
+
 def _register_web_search_mcp(workspace: str, search_model: str, profile: str | None = None) -> bool:
     """Register (or replace) the web_search MCP server in Claude Code's user
     scope via `claude mcp add-json`. Removes any prior entry first so re-runs

src/ucode/databricks.py

@@ -804,14 +804,6 @@ def build_auth_shell_command(workspace: str, profile: str | None = None) -> str:
     )
 
 
-# Model ids the AI Gateway advertises but currently rejects at launch with
-# "The provided model identifier is invalid." Filter them out at discovery so
-# agent harnesses don't try to boot with them.
-_CLAUDE_MODEL_DISCOVERY_DENYLIST = {
-    "databricks-claude-opus-4-8",
-}
-
-
 def discover_claude_models(workspace: str, token: str) -> tuple[dict[str, str], str | None]:
     """Discover Claude families on this workspace's AI Gateway.
 
@@ -828,9 +820,7 @@ def discover_claude_models(workspace: str, token: str) -> tuple[dict[str, str],
     raw_ids = [
         m["id"]
         for m in data.get("data", [])
-        if isinstance(m.get("id"), str)
-        and not m["id"].endswith("-anthropic")
-        and m["id"] not in _CLAUDE_MODEL_DISCOVERY_DENYLIST
+        if isinstance(m.get("id"), str) and not m["id"].endswith("-anthropic")
     ]
 
     result: dict[str, str] = {}

tests/test_agent_claude.py

@@ -25,6 +25,22 @@ def test_sets_anthropic_model(self):
         overlay, _ = claude.render_overlay(WS, "databricks-claude-sonnet-4")
         assert overlay["env"]["ANTHROPIC_MODEL"] == "databricks-claude-sonnet-4"
 
+    def test_adds_1m_suffix_for_opus_4_6_and_later(self):
+        overlay, _ = claude.render_overlay(WS, "databricks-claude-opus-4-7")
+        assert overlay["env"]["ANTHROPIC_MODEL"] == "databricks-claude-opus-4-7[1m]"
+
+    def test_adds_1m_suffix_for_sonnet_4_6_and_later(self):
+        overlay, _ = claude.render_overlay(WS, "databricks-claude-sonnet-4-7")
+        assert overlay["env"]["ANTHROPIC_MODEL"] == "databricks-claude-sonnet-4-7[1m]"
+
+    def test_does_not_add_1m_suffix_for_other_models(self):
+        overlay, _ = claude.render_overlay(WS, "databricks-claude-haiku-4-6")
+        assert overlay["env"]["ANTHROPIC_MODEL"] == "databricks-claude-haiku-4-6"
+
+    def test_does_not_duplicate_1m_suffix(self):
+        overlay, _ = claude.render_overlay(WS, "databricks-claude-opus-4-7[1m]")
+        assert overlay["env"]["ANTHROPIC_MODEL"] == "databricks-claude-opus-4-7[1m]"
+
     def test_sets_anthropic_base_url(self):
         overlay, _ = claude.render_overlay(WS, "s4")
         assert overlay["env"]["ANTHROPIC_BASE_URL"] == f"{WS}/ai-gateway/anthropic"
@@ -43,12 +59,16 @@ def test_sets_api_key_helper(self):
         assert WS in overlay["apiKeyHelper"]
 
     def test_model_overrides_when_all_provided(self):
-        models = {"sonnet": "s4", "opus": "o4", "haiku": "h4"}
+        models = {
+            "sonnet": "databricks-claude-sonnet-4-6",
+            "opus": "databricks-claude-opus-4-7",
+            "haiku": "databricks-claude-haiku-4-6",
+        }
         overlay, _ = claude.render_overlay(WS, "s4", claude_models=models)
         env = overlay["env"]
-        assert env["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "s4"
-        assert env["ANTHROPIC_DEFAULT_OPUS_MODEL"] == "o4"
-        assert env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] == "h4"
+        assert env["ANTHROPIC_DEFAULT_SONNET_MODEL"] == "databricks-claude-sonnet-4-6[1m]"
+        assert env["ANTHROPIC_DEFAULT_OPUS_MODEL"] == "databricks-claude-opus-4-7[1m]"
+        assert env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] == "databricks-claude-haiku-4-6"
 
     def test_model_overrides_partial(self):
         models = {"sonnet": "s4"}

tests/test_databricks.py

@@ -114,6 +114,23 @@ def test_codex_url_format(self):
         assert urls["codex"] == f"{WS}/ai-gateway/codex/v1"
 
 
+class TestDiscoverClaudeModels:
+    def test_selects_opus_4_8_when_advertised(self, monkeypatch):
+        payload = {
+            "data": [
+                {"id": "databricks-claude-opus-4-7"},
+                {"id": "databricks-claude-opus-4-8"},
+                {"id": "databricks-claude-sonnet-4-6"},
+            ]
+        }
+        monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None))
+
+        models, reason = db_mod.discover_claude_models(WS, "token")
+
+        assert reason is None
+        assert models["opus"] == "databricks-claude-opus-4-8"
+
+
 class TestBuildAuthShellCommand:
     def test_contains_workspace(self):
         cmd = build_auth_shell_command(WS)