From 8d9f6203545ff6b13256fe8c94fc2b4db6260ad7 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Sat, 5 Sep 2026 01:20:13 +0000 Subject: [PATCH 1/5] add in the picker --- src/ucode/smart_routing/v2.py | 45 ++++++++++++++++++- tests/test_claude_smart_routing_v2.py | 63 +++++++++++++++++++++++++++ 2 files changed, 107 insertions(+), 1 deletion(-) diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 9b8a29eb..6f50dc89 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -19,6 +19,7 @@ from ucode.config_io import APP_DIR, read_json_safe, read_toml_safe, write_json_file from ucode.constants import LOOPBACK_HOST from ucode.databricks import ( + AnthropicModelCatalog, build_auth_token_argv, get_databricks_token, list_anthropic_model_catalog, @@ -58,6 +59,45 @@ _ANTHROPIC_AIGW_MODEL_RE = re.compile(r"^anthropic-aigw-[0-9a-fA-F]{8}-(.+)$") +def _model_picker_catalog() -> AnthropicModelCatalog | None: + """Read model-picker rows using the managed-settings then ucode-settings waterfall. + + A managed picker is authoritative for smart routing: its rows are the models the + administrator exposed, so there is no need to query the gateway catalog first. + """ + try: + from ucode.agents.claude import ( + CLAUDE_SETTINGS_PATH, + CLAUDE_USER_SETTINGS_PATH, + _managed_settings_path, + ) + + paths = [_managed_settings_path(), CLAUDE_SETTINGS_PATH, CLAUDE_USER_SETTINGS_PATH] + except (ImportError, OSError): + return None + for path in paths: + if path is None or not path.is_file(): + continue + settings = read_json_safe(path) + picker_settings = settings.get("modelPicker") if isinstance(settings, dict) else None + picker = picker_settings.get("options") if isinstance(picker_settings, dict) else None + if not isinstance(picker, list): + continue + model_ids: list[str] = [] + seen: set[str] = set() + for row in picker: + if not isinstance(row, dict) or not isinstance(row.get("model"), str): + continue + model_id = row["model"].strip() + if not model_id or model_id in seen: + continue + seen.add(model_id) + model_ids.append(model_id) + if model_ids: + return AnthropicModelCatalog(model_ids, {}) + return None + + def enabled() -> bool: return os.environ.get(ENV_VAR) == "1" @@ -348,7 +388,10 @@ def launch_claude( os.environ[OAUTH_TOKEN_ENV_VAR] = token os.environ[GATEWAY_MODEL_DISCOVERY_ENV_VAR] = "1" os.environ["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] = "1" - catalog = list_anthropic_model_catalog(workspace, token) + # Enterprise managed settings may provide the authoritative model picker. Prefer + # those IDs so routing follows the administrator's curated list and avoids an + # unnecessary /anthropic/v1/models request. + catalog = _model_picker_catalog() or list_anthropic_model_catalog(workspace, token) if not catalog.model_ids: raise RuntimeError( catalog.error_msg or "Anthropic models endpoint returned no Claude models" diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index ad0429d9..02ae7473 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -15,6 +15,69 @@ from ucode.smart_routing import claude_hooks, claude_pty, routing, v2 +class TestManagedModelPicker: + def test_reads_model_ids_from_managed_picker(self, tmp_path, monkeypatch): + path = tmp_path / "managed-settings.json" + path.write_text( + json.dumps( + { + "modelPicker": { + "options": [ + {"model": "system.ai.claude-opus-4-8", "label": "Opus"}, + {"model": "system.ai.claude-sonnet-5", "label": "Sonnet"}, + ] + } + } + ) + ) + monkeypatch.setattr(claude, "_managed_settings_path", lambda: path) + + catalog = v2._model_picker_catalog() + + assert catalog is not None + assert catalog.model_ids == ["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5"] + assert catalog.model_id_to_display_name == {} + + def test_ignores_empty_or_missing_picker(self, tmp_path, monkeypatch): + path = tmp_path / "managed-settings.json" + path.write_text(json.dumps({"env": {}})) + monkeypatch.setattr(claude, "_managed_settings_path", lambda: path) + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", tmp_path / "ucode-settings.json") + + assert v2._model_picker_catalog() is None + + def test_falls_back_to_ucode_settings_picker(self, tmp_path, monkeypatch): + managed = tmp_path / "managed-settings.json" + managed.write_text(json.dumps({"env": {}})) + ucode_settings = tmp_path / "ucode-settings.json" + ucode_settings.write_text( + json.dumps({"modelPicker": {"options": [{"model": "system.ai.claude-opus-5"}]}}) + ) + monkeypatch.setattr(claude, "_managed_settings_path", lambda: managed) + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", ucode_settings) + + catalog = v2._model_picker_catalog() + + assert catalog is not None + assert catalog.model_ids == ["system.ai.claude-opus-5"] + + def test_falls_back_to_user_settings_picker(self, tmp_path, monkeypatch): + user_settings = tmp_path / "settings.json" + user_settings.write_text( + json.dumps( + {"modelPicker": {"options": [{"model": "system.ai.claude-sonnet-5"}]}} + ) + ) + monkeypatch.setattr(claude, "_managed_settings_path", lambda: tmp_path / "missing-managed") + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", tmp_path / "missing-ucode") + monkeypatch.setattr(claude, "CLAUDE_USER_SETTINGS_PATH", user_settings) + + catalog = v2._model_picker_catalog() + + assert catalog is not None + assert catalog.model_ids == ["system.ai.claude-sonnet-5"] + + class TestDirectModelCommand: @pytest.mark.parametrize( "name", From fd1719a2f1e81b320b8917d8ae602fc2b4ed1a26 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Sat, 5 Sep 2026 01:23:08 +0000 Subject: [PATCH 2/5] Document model picker discovery precedence --- src/ucode/smart_routing/v2.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 6f50dc89..2a20f342 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -72,6 +72,8 @@ def _model_picker_catalog() -> AnthropicModelCatalog | None: _managed_settings_path, ) + # Hierarchy: managed settings, CLI-supplied settings (ucode-settings.json), local user + # settings, based on the modelPicker scope documented at https://code.claude.com/docs/en/settings-reference#modelpicker. paths = [_managed_settings_path(), CLAUDE_SETTINGS_PATH, CLAUDE_USER_SETTINGS_PATH] except (ImportError, OSError): return None @@ -388,9 +390,7 @@ def launch_claude( os.environ[OAUTH_TOKEN_ENV_VAR] = token os.environ[GATEWAY_MODEL_DISCOVERY_ENV_VAR] = "1" os.environ["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] = "1" - # Enterprise managed settings may provide the authoritative model picker. Prefer - # those IDs so routing follows the administrator's curated list and avoids an - # unnecessary /anthropic/v1/models request. + # modelPicker takes priority over model discovery. catalog = _model_picker_catalog() or list_anthropic_model_catalog(workspace, token) if not catalog.model_ids: raise RuntimeError( From 7f1d0edf37e9b6a910334e00395319d571bde4e5 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Sat, 5 Sep 2026 01:35:59 +0000 Subject: [PATCH 3/5] update --- src/ucode/agents/claude.py | 3 +++ tests/test_agent_claude.py | 19 +++++++++++++++++++ tests/test_claude_smart_routing_v2.py | 4 +--- 3 files changed, 23 insertions(+), 3 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 01b3a1f9..6c3bffd5 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -707,6 +707,9 @@ def _compose(base: dict, *, enforce_model_default_hierarchy: bool) -> dict: else: target_env[key] = selected_default_model merged = deep_merge_dict(base, overlay_for_merge) + if enforce_model_default_hierarchy and "modelPicker" in base: + # Claude's managed modelPicker is administrator-owned; ucode must leave it intact. + merged["modelPicker"] = copy.deepcopy(base["modelPicker"]) overlay_custom_headers = overlay_for_merge["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] merged["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] = _merge_anthropic_custom_headers( existing_custom_headers, overlay_custom_headers diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 61f07438..ffd1f994 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -762,6 +762,25 @@ def test_managed_file_preserves_other_keys(self, monkeypatch): assert written["env"]["ANTHROPIC_BASE_URL"] assert written["apiKeyHelper"] + def test_managed_file_preserves_model_picker(self, monkeypatch): + private_writes: list = [] + managed_writes: list = [] + picker = { + "replaceBuiltInOptions": True, + "options": [ + {"model": "system.ai.claude-opus-4-8"}, + {"model": "system.ai.glm-5-2"}, + ], + } + existing = {str(FAKE_MANAGED_PATH): {"modelPicker": picker}} + self._patch(monkeypatch, private_writes, managed_writes, existing) + state = {"workspace": WS, "codex_models": []} + + claude.write_tool_config(state, "databricks-claude-sonnet-4") + + written = json.loads(managed_writes[0][1]) + assert written["modelPicker"] == picker + def test_managed_file_strips_stale_gateway_model_discovery(self, monkeypatch): private_writes: list = [] managed_writes: list = [] diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index 02ae7473..49e3f59e 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -64,9 +64,7 @@ def test_falls_back_to_ucode_settings_picker(self, tmp_path, monkeypatch): def test_falls_back_to_user_settings_picker(self, tmp_path, monkeypatch): user_settings = tmp_path / "settings.json" user_settings.write_text( - json.dumps( - {"modelPicker": {"options": [{"model": "system.ai.claude-sonnet-5"}]}} - ) + json.dumps({"modelPicker": {"options": [{"model": "system.ai.claude-sonnet-5"}]}}) ) monkeypatch.setattr(claude, "_managed_settings_path", lambda: tmp_path / "missing-managed") monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", tmp_path / "missing-ucode") From 92d76fbba8488245255129381434c81fe6fa44e5 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 9 Sep 2026 15:08:52 +0000 Subject: [PATCH 4/5] update --- src/ucode/agents/claude.py | 6 +++--- tests/test_agent_claude.py | 14 ++++++++++++-- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index c6aa03c0..331f6690 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -653,9 +653,6 @@ def _compose(base: dict, *, enforce_model_default_hierarchy: bool) -> dict: else: target_env[key] = selected_default_model merged = deep_merge_dict(base, overlay_for_merge) - if enforce_model_default_hierarchy and "modelPicker" in base: - # Claude's managed modelPicker is administrator-owned; ucode must leave it intact. - merged["modelPicker"] = copy.deepcopy(base["modelPicker"]) overlay_custom_headers = overlay_for_merge["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] merged["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] = _merge_anthropic_custom_headers( existing_custom_headers, overlay_custom_headers @@ -789,6 +786,9 @@ def _reconcile_managed_settings( configuration mirrors ucode's settings there. The same compose operation that produced the private file is applied to the existing managed file, preserving unrelated IT-authored keys. + `ug configure` updates gateway-owned fields in this file, but does not generate or modify + the `modelPicker` object; an existing picker is retained by the merge. + Relayed launches are skipped: they depend on a per-session loopback refresh proxy that only runs during `ucode claude`, so a bare `claude` could not reach the gateway anyway. """ diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 5d9af510..ac5ea339 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -738,6 +738,8 @@ def test_writes_managed_file_by_default(self, monkeypatch): # Private file still written; managed file written too. assert str(claude.CLAUDE_SETTINGS_PATH) in [p for p, _ in private_writes] assert [p for p, _ in managed_writes] == [str(FAKE_MANAGED_PATH)] + assert "modelPicker" not in private_writes[0][1] + assert "modelPicker" not in json.loads(managed_writes[0][1]) def test_managed_file_preserves_other_keys(self, monkeypatch): private_writes: list = [] @@ -753,7 +755,7 @@ def test_managed_file_preserves_other_keys(self, monkeypatch): assert written["env"]["ANTHROPIC_BASE_URL"] assert written["apiKeyHelper"] - def test_managed_file_preserves_model_picker(self, monkeypatch): + def test_managed_file_updates_gateway_settings_without_changing_model_picker(self, monkeypatch): private_writes: list = [] managed_writes: list = [] picker = { @@ -763,7 +765,14 @@ def test_managed_file_preserves_model_picker(self, monkeypatch): {"model": "system.ai.glm-5-2"}, ], } - existing = {str(FAKE_MANAGED_PATH): {"modelPicker": picker}} + existing = { + str(FAKE_MANAGED_PATH): { + "modelPicker": picker, + "env": { + "ANTHROPIC_BASE_URL": "https://old-workspace.databricks.com/ai-gateway/anthropic" + }, + } + } self._patch(monkeypatch, private_writes, managed_writes, existing) state = {"workspace": WS, "codex_models": []} @@ -771,6 +780,7 @@ def test_managed_file_preserves_model_picker(self, monkeypatch): written = json.loads(managed_writes[0][1]) assert written["modelPicker"] == picker + assert written["env"]["ANTHROPIC_BASE_URL"] == f"{WS}/ai-gateway/anthropic" def test_managed_file_strips_stale_gateway_model_discovery(self, monkeypatch): private_writes: list = [] From b5c4de158216b0aa3bff47c933c8c8448a6b05e4 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Wed, 9 Sep 2026 15:26:15 +0000 Subject: [PATCH 5/5] fix test --- src/ucode/databricks.py | 7 ++++--- tests/test_databricks.py | 18 ++++++++++++++++++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index fd656d72..0e942a80 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1505,7 +1505,8 @@ def build_auth_shell_command( # Claude model families ucode buckets, newest tier first. Each maps to a # Claude Code family alias (ANTHROPIC_DEFAULT__MODEL). Add an entry to # support a new family in both discovery paths (`claude--*` via the -# model-services listing and `databricks-claude--*` via the AI Gateway). +# model-services listing and either `databricks-claude--*` or +# `system.ai.claude--*` via the AI Gateway). ANTHROPIC_FAMILIES = ("fable", "opus", "sonnet", "haiku") @@ -2938,7 +2939,7 @@ def discover_claude_models(workspace: str, token: str) -> tuple[dict[str, str], result: dict[str, str] = {} for family in ANTHROPIC_FAMILIES: candidates = sorted( - [m for m in raw_ids if f"databricks-claude-{family}-" in m], + [m for m in raw_ids if f"claude-{family}-" in m], reverse=True, ) if candidates: @@ -2953,7 +2954,7 @@ def discover_claude_models(workspace: str, token: str) -> tuple[dict[str, str], families = ",".join(ANTHROPIC_FAMILIES) return {}, ( "AI Gateway returned model ids but none matched " - f"`databricks-claude-{{{families}}}-*` (got: {sample})" + f"`*-claude-{{{families}}}-*` (got: {sample})" ) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 1f280e5a..bbc07bd8 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -225,6 +225,24 @@ def test_selects_opus_4_8_when_advertised(self, monkeypatch): assert reason is None assert models["opus"] == "databricks-claude-opus-4-8" + def test_buckets_system_ai_claude_models(self, monkeypatch): + payload = { + "data": [ + {"id": "system.ai.claude-opus-4-8"}, + {"id": "system.ai.claude-sonnet-4-6"}, + {"id": "system.ai.glm-5-3-flash"}, + ] + } + monkeypatch.setattr(db_mod, "_http_get_json", lambda *_args, **_kwargs: (payload, None)) + + models, reason = db_mod.discover_claude_models(WS, "token") + + assert reason is None + assert models == { + "opus": "system.ai.claude-opus-4-8", + "sonnet": "system.ai.claude-sonnet-4-6", + } + def test_buckets_fable_family(self, monkeypatch): payload = { "data": [