diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index f4138719..ee54d153 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -24,6 +24,7 @@ from concurrent.futures import ( TimeoutError as FutureTimeoutError, ) +from dataclasses import dataclass from decimal import Decimal, InvalidOperation from pathlib import Path from typing import Literal, NamedTuple, NoReturn, cast, overload @@ -70,6 +71,15 @@ _ANTHROPIC_MODEL_DISCOVERY_SETUP_MAX_RETRIES = 2 +@dataclass(frozen=True) +class AnthropicModelCatalog: + """Models advertised by AI Gateway's Anthropic endpoint.""" + + model_ids: list[str] + model_id_to_display_name: dict[str, str] + error_msg: str | None = None + + def _debug_enabled() -> bool: return os.environ.get("UCODE_DEBUG") == "1" @@ -2874,12 +2884,19 @@ def list_anthropic_models(workspace: str, token: str) -> tuple[list[str], str | using that mode must not apply ucode's legacy ``databricks-claude-*`` family validation. """ + catalog = list_anthropic_model_catalog(workspace, token) + return catalog.model_ids, catalog.error_msg + + +def list_anthropic_model_catalog(workspace: str, token: str) -> AnthropicModelCatalog: + """Return advertised Anthropic model ids and their optional display names.""" payload, reason = _get_anthropic_models_json(workspace, token) if payload is None: - return [], reason + return AnthropicModelCatalog(model_ids=[], model_id_to_display_name={}, error_msg=reason) data = cast(dict, payload) if isinstance(payload, dict) else {} model_ids: list[str] = [] + display_names: dict[str, str] = {} seen: set[str] = set() for model in data.get("data", []): if not isinstance(model, dict): @@ -2888,9 +2905,16 @@ def list_anthropic_models(workspace: str, token: str) -> tuple[list[str], str | if isinstance(model_id, str) and model_id and model_id not in seen: seen.add(model_id) model_ids.append(model_id) + display_name = model.get("display_name") + if isinstance(display_name, str) and display_name: + display_names[model_id] = display_name if model_ids: - return model_ids, None - return [], "AI Gateway returned no Anthropic model ids" + return AnthropicModelCatalog(model_ids=model_ids, model_id_to_display_name=display_names) + return AnthropicModelCatalog( + model_ids=[], + model_id_to_display_name={}, + error_msg="AI Gateway returned no Anthropic model ids", + ) def discover_claude_models(workspace: str, token: str) -> tuple[dict[str, str], str | None]: diff --git a/src/ucode/smart_routing/claude_pty.py b/src/ucode/smart_routing/claude_pty.py index 933b7cef..88d07f0b 100644 --- a/src/ucode/smart_routing/claude_pty.py +++ b/src/ucode/smart_routing/claude_pty.py @@ -23,6 +23,7 @@ import time import tty from collections.abc import Callable +from dataclasses import dataclass from pathlib import Path MAX_MODEL_NAME_LEN = 200 @@ -56,6 +57,15 @@ def valid_model_name(name: object) -> bool: ) +@dataclass(frozen=True) +class FirstPromptRoute: + """The model switch and user-facing label for a first-prompt route.""" + + model: str + display_model: str + rationale: str + + class ConfirmationState: """Detect and accept Claude's optional cache-cost confirmation dialog.""" @@ -158,16 +168,21 @@ def first_prompt_hook_output(response: dict | None) -> dict | None: if not valid_model_name(model): return None assert isinstance(model, str) + display_model = response.get("display_model", model) + if not isinstance(display_model, str) or not display_model: + display_model = model rationale = response.get("rationale") return { "decision": "block", - "reason": format_routing_notice(model, rationale if isinstance(rationale, str) else None), + "reason": format_routing_notice( + display_model, rationale if isinstance(rationale, str) else None + ), } def serve_first_prompt_socket( path: Path, - route_prompt: Callable[[str], tuple[str, str]], + route_prompt: Callable[[str], FirstPromptRoute], on_blocked_prompt: Callable[[str, str], None], stop: threading.Event, *, @@ -214,7 +229,10 @@ def serve() -> None: and prompt.strip() and not is_command ): - model, rationale = route_prompt(prompt) + routed = route_prompt(prompt) + model = routed.model + display_model = routed.display_model + rationale = routed.rationale if valid_model_name(model): claimed = True response = { @@ -222,6 +240,7 @@ def serve() -> None: "model": model, "rationale": rationale, } + response["display_model"] = display_model blocked = (prompt, model) except Exception as exc: # noqa: BLE001 - hooks must fail open log(f"[ERR] first-prompt request: {exc!r}") @@ -269,7 +288,7 @@ def sync_winsize(master_fd: int, stdin_fd: int = 0) -> None: def run_claude_pty( argv: list[str], *, - route_prompt: Callable[[str], tuple[str, str]], + route_prompt: Callable[[str], FirstPromptRoute], socket_path: Path, prepare_model_switch: Callable[[str], None] = lambda _model: None, model_switch_persisted: Callable[[], bool] = lambda: True, diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 84b019b3..9b8a29eb 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -21,6 +21,7 @@ from ucode.databricks import ( build_auth_token_argv, get_databricks_token, + list_anthropic_model_catalog, list_anthropic_models, ) from ucode.smart_routing import claude_routing, codex_interposer, routing @@ -347,9 +348,12 @@ 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" - model_ids, discovery_error = list_anthropic_models(workspace, token) - if not model_ids: - raise RuntimeError(discovery_error or "Anthropic models endpoint returned no Claude models") + catalog = list_anthropic_model_catalog(workspace, token) + if not catalog.model_ids: + raise RuntimeError( + catalog.error_msg or "Anthropic models endpoint returned no Claude models" + ) + model_ids = catalog.model_ids run_id = f"{os.getpid()}-{uuid.uuid4().hex[:8]}" socket_path = APP_DIR / f"claude-v2-{run_id}.sock" @@ -381,9 +385,13 @@ def launch_claude( model_setting = _ClaudeModelSettingGuard(user_settings_path) - def route_prompt(prompt: str) -> tuple[str, str]: + def route_prompt(prompt: str) -> claude_pty.FirstPromptRoute: decision = _route_claude_prompt(state, token, prompt, model_ids) - return model_name(decision.model), decision.rationale + return claude_pty.FirstPromptRoute( + model=model_name(decision.model), + display_model=catalog.model_id_to_display_name.get(decision.model, decision.model), + rationale=decision.rationale, + ) print_note( "Smart routing v2: the first submitted prompt will select Claude Code's " diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index 8d722270..ad0429d9 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -11,6 +11,7 @@ import pytest from ucode.agents import claude +from ucode.databricks import AnthropicModelCatalog from ucode.smart_routing import claude_hooks, claude_pty, routing, v2 @@ -47,13 +48,27 @@ def test_omits_reason_when_router_returns_none(self): assert "Reason" not in result["reason"] + def test_displays_catalog_name_while_retaining_routable_model(self): + result = claude_pty.first_prompt_hook_output( + { + "action": "block", + "model": "anthropic-aigw-77df06ea-system.ai.glm-5-3-flash", + "display_model": "GLM 5.3 Flash", + } + ) + + assert "Selected Model : GLM 5.3 Flash" in result["reason"] + assert "anthropic-aigw-77df06ea" not in result["reason"] + def test_blocks_once_then_allows_replay(self, tmp_path): socket_path = tmp_path / "first.sock" blocked: list[tuple[str, str]] = [] stop = threading.Event() claude_pty.serve_first_prompt_socket( socket_path, - lambda _prompt: ("sonnet", "Selected for a narrow task."), + lambda _prompt: claude_pty.FirstPromptRoute( + model="sonnet", display_model="sonnet", rationale="Selected for a narrow task." + ), lambda prompt, model: blocked.append((prompt, model)), stop, ) @@ -70,6 +85,7 @@ def test_blocks_once_then_allows_replay(self, tmp_path): assert first == { "action": "block", "model": "sonnet", + "display_model": "sonnet", "rationale": "Selected for a narrow task.", } assert replay == {"action": "allow"} @@ -102,10 +118,10 @@ def test_restores_model_captured_immediately_before_switch(self, tmp_path, monke monkeypatch.setattr(v2, "build_auth_token_argv", lambda *_args, **_kwargs: ["ucode"]) monkeypatch.setattr( v2, - "list_anthropic_models", - lambda *_args: ( - ["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5"], - None, + "list_anthropic_model_catalog", + lambda *_args: AnthropicModelCatalog( + model_ids=["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5"], + model_id_to_display_name={"system.ai.claude-sonnet-5": "Claude Sonnet 5"}, ), ) monkeypatch.setattr( @@ -155,9 +171,10 @@ def fake_run(argv, **kwargs): assert exc.value.code == 0 assert captured["argv"][3:5] == ["--model", "opus"] assert captured["argv"][-1] == "--debug" - assert captured["routed_model"] == ( - "system.ai.claude-sonnet-5[1m]", - "Selected for the parser task.", + assert captured["routed_model"] == claude_pty.FirstPromptRoute( + model="system.ai.claude-sonnet-5[1m]", + display_model="Claude Sonnet 5", + rationale="Selected for the parser task.", ) assert {definition["model"] for definition in captured["agents"].values()} == { "system.ai.claude-opus-4-8", @@ -199,7 +216,11 @@ def test_does_not_restore_when_wrapper_never_switches(self, tmp_path, monkeypatc monkeypatch.setattr(v2, "APP_DIR", tmp_path) monkeypatch.setattr(v2, "get_databricks_token", lambda *_args, **_kwargs: "token") monkeypatch.setattr(v2, "build_auth_token_argv", lambda *_args, **_kwargs: ["ucode"]) - monkeypatch.setattr(v2, "list_anthropic_models", lambda *_args: (["opus"], None)) + monkeypatch.setattr( + v2, + "list_anthropic_model_catalog", + lambda *_args: AnthropicModelCatalog(model_ids=["opus"], model_id_to_display_name={}), + ) def fake_run(_argv, **_kwargs): user_settings.write_text(json.dumps({"model": "user-selected"})) @@ -228,7 +249,11 @@ def test_restores_after_routed_model_persists_and_preserves_later_choice( monkeypatch.setattr(v2, "APP_DIR", tmp_path) monkeypatch.setattr(v2, "get_databricks_token", lambda *_args, **_kwargs: "token") monkeypatch.setattr(v2, "build_auth_token_argv", lambda *_args, **_kwargs: ["ucode"]) - monkeypatch.setattr(v2, "list_anthropic_models", lambda *_args: (["opus"], None)) + monkeypatch.setattr( + v2, + "list_anthropic_model_catalog", + lambda *_args: AnthropicModelCatalog(model_ids=["opus"], model_id_to_display_name={}), + ) def fake_run(_argv, **kwargs): routed_model = "system.ai.claude-opus-4-8" @@ -455,7 +480,9 @@ def is_alive(): with pytest.raises(RuntimeError, match="Claude was not launched"): claude_pty.run_claude_pty( ["claude"], - route_prompt=lambda _prompt: ("sonnet", ""), + route_prompt=lambda _prompt: claude_pty.FirstPromptRoute( + model="sonnet", display_model="sonnet", rationale="" + ), socket_path=tmp_path / "missing.sock", ) @@ -513,7 +540,11 @@ def read_until(suffix): str(capture), str(restored), ], - route_prompt=lambda _prompt: ("system.ai.claude-sonnet-5", ""), + route_prompt=lambda _prompt: claude_pty.FirstPromptRoute( + model="system.ai.claude-sonnet-5", + display_model="system.ai.claude-sonnet-5", + rationale="", + ), socket_path=socket_path, restore_model_setting=lambda: restored.write_text("restored"), ) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 74eecb84..7debda88 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -195,6 +195,21 @@ def fake_get(url, token, **kwargs): {"max_retries": 2}, ) + def test_lists_anthropic_display_names_with_model_ids(self, monkeypatch): + payload = { + "data": [ + {"id": "system.ai.glm-5-3-flash", "display_name": "GLM 5.3 Flash"}, + {"id": "opaque-model-id"}, + ] + } + monkeypatch.setattr(db_mod, "_http_get_json", lambda *_args, **_kwargs: (payload, None)) + + catalog = db_mod.list_anthropic_model_catalog(WS, "token") + + assert catalog.model_ids == ["system.ai.glm-5-3-flash", "opaque-model-id"] + assert catalog.model_id_to_display_name == {"system.ai.glm-5-3-flash": "GLM 5.3 Flash"} + assert catalog.error_msg is None + def test_selects_opus_4_8_when_advertised(self, monkeypatch): payload = { "data": [