diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index f081cf9b..331f6690 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -786,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/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/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 9b8a29eb..2a20f342 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,47 @@ _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, + ) + + # 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 + 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 +390,8 @@ 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) + # modelPicker takes priority over model discovery. + 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_agent_claude.py b/tests/test_agent_claude.py index 17532fa3..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,6 +755,33 @@ def test_managed_file_preserves_other_keys(self, monkeypatch): assert written["env"]["ANTHROPIC_BASE_URL"] assert written["apiKeyHelper"] + def test_managed_file_updates_gateway_settings_without_changing_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, + "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": []} + + claude.write_tool_config(state, "databricks-claude-sonnet-4") + + 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 = [] managed_writes: list = [] diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index ad0429d9..49e3f59e 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -15,6 +15,67 @@ 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", 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": [