From 288f8539ccd4716d16a8333a7c450787001338ce Mon Sep 17 00:00:00 2001 From: AarushiShah-db Date: Thu, 28 May 2026 20:09:35 +0000 Subject: [PATCH 1/4] claude: append [1m] to Opus >=4.6 and Sonnet 4.6 model ids when writing settings.json; update tests --- src/ucode/agents/claude.py | 24 +++++++++++++++++++++--- tests/test_agent_claude.py | 28 ++++++++++++++++++++++++---- 2 files changed, 45 insertions(+), 7 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 91793c7..01dd33b 100644 --- a/src/ucode/agents/claude.py +++ b/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,7 +105,7 @@ 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", @@ -111,9 +113,9 @@ def render_overlay( } 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 diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 8d0526e..4ef2f83 100644 --- a/tests/test_agent_claude.py +++ b/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(self): + overlay, _ = claude.render_overlay(WS, "databricks-claude-sonnet-4-6") + assert overlay["env"]["ANTHROPIC_MODEL"] == "databricks-claude-sonnet-4-6[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"} From 084ee1856ed24d4def7e094a863274435f763bd3 Mon Sep 17 00:00:00 2001 From: AarushiShah-db Date: Thu, 28 May 2026 20:14:28 +0000 Subject: [PATCH 2/4] claude: allow opus 4.8 discovery --- src/ucode/databricks.py | 9 --------- tests/test_databricks.py | 17 +++++++++++++++++ 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index d296c0e..5d9b3dd 100644 --- a/src/ucode/databricks.py +++ b/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. @@ -830,7 +822,6 @@ def discover_claude_models(workspace: str, token: str) -> tuple[dict[str, str], 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 ] result: dict[str, str] = {} diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 49306cd..fffdba5 100644 --- a/tests/test_databricks.py +++ b/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) From 93d251fc5ef857ff68bc8ef1ff475e5c6d35fa8e Mon Sep 17 00:00:00 2001 From: AarushiShah-db Date: Thu, 28 May 2026 20:15:58 +0000 Subject: [PATCH 3/4] chore: format databricks --- src/ucode/databricks.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 5d9b3dd..5ed537c 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -820,8 +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") + if isinstance(m.get("id"), str) and not m["id"].endswith("-anthropic") ] result: dict[str, str] = {} From a1adabb5580b2e1889614cd8f9653f4a9217f17a Mon Sep 17 00:00:00 2001 From: AarushiShah-db Date: Thu, 28 May 2026 20:23:21 +0000 Subject: [PATCH 4/4] claude: apply 1m suffix to future sonnet models --- src/ucode/agents/claude.py | 2 +- tests/test_agent_claude.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 01dd33b..0fe4ec7 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -142,7 +142,7 @@ def _maybe_add_1m_suffix(model: str) -> str: 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) + family == "sonnet" and (major, minor) >= (4, 6) ) return f"{model}[1m]" if should_suffix else model diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 4ef2f83..ea33c63 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -29,9 +29,9 @@ 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(self): - overlay, _ = claude.render_overlay(WS, "databricks-claude-sonnet-4-6") - assert overlay["env"]["ANTHROPIC_MODEL"] == "databricks-claude-sonnet-4-6[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")