From 09641eb062a44d8789d00bcfe4be1882032b699a Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Thu, 3 Sep 2026 19:21:05 +0000 Subject: [PATCH 1/6] render display name in box --- src/ucode/databricks.py | 18 +++++++++++++++--- src/ucode/smart_routing/claude_pty.py | 19 +++++++++++++++---- src/ucode/smart_routing/v2.py | 11 ++++++++--- tests/test_claude_smart_routing_v2.py | 23 ++++++++++++++++++++--- tests/test_databricks.py | 15 +++++++++++++++ 5 files changed, 73 insertions(+), 13 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index f4138719..c708067e 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2874,12 +2874,21 @@ 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. """ + model_ids, _display_names, reason = list_anthropic_model_catalog(workspace, token) + return model_ids, reason + + +def list_anthropic_model_catalog( + workspace: str, token: str +) -> tuple[list[str], dict[str, str], str | None]: + """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 [], {}, 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 +2897,12 @@ 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 model_ids, display_names, None + return [], {}, "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..a7527baf 100644 --- a/src/ucode/smart_routing/claude_pty.py +++ b/src/ucode/smart_routing/claude_pty.py @@ -158,16 +158,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], tuple[str, str] | tuple[str, str, str]], on_blocked_prompt: Callable[[str, str], None], stop: threading.Event, *, @@ -214,12 +219,18 @@ def serve() -> None: and prompt.strip() and not is_command ): - model, rationale = route_prompt(prompt) + routed = route_prompt(prompt) + if len(routed) == 2: + model, rationale = routed + display_model = model + else: + model, display_model, rationale = routed if valid_model_name(model): claimed = True response = { "action": "block", "model": model, + "display_model": display_model, "rationale": rationale, } blocked = (prompt, model) @@ -269,7 +280,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], tuple[str, str] | tuple[str, str, str]], 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..a653a28e 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,7 +348,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" - model_ids, discovery_error = list_anthropic_models(workspace, token) + model_ids, display_names, discovery_error = list_anthropic_model_catalog(workspace, token) if not model_ids: raise RuntimeError(discovery_error or "Anthropic models endpoint returned no Claude models") @@ -381,9 +382,13 @@ def launch_claude( model_setting = _ClaudeModelSettingGuard(user_settings_path) - def route_prompt(prompt: str) -> tuple[str, str]: + def route_prompt(prompt: str) -> tuple[str, str, str]: decision = _route_claude_prompt(state, token, prompt, model_ids) - return model_name(decision.model), decision.rationale + return ( + model_name(decision.model), + display_names.get(decision.model, decision.model), + 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..e9fcbf25 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -47,6 +47,18 @@ 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]] = [] @@ -102,9 +114,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", + "list_anthropic_model_catalog", lambda *_args: ( ["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5"], + {"system.ai.claude-sonnet-5": "Claude Sonnet 5"}, None, ), ) @@ -199,7 +212,9 @@ 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: (["opus"], {}, None) + ) def fake_run(_argv, **_kwargs): user_settings.write_text(json.dumps({"model": "user-selected"})) @@ -228,7 +243,9 @@ 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: (["opus"], {}, None) + ) def fake_run(_argv, **kwargs): routed_model = "system.ai.claude-opus-4-8" diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 74eecb84..1ed0d664 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)) + + models, display_names, reason = db_mod.list_anthropic_model_catalog(WS, "token") + + assert models == ["system.ai.glm-5-3-flash", "opaque-model-id"] + assert display_names == {"system.ai.glm-5-3-flash": "GLM 5.3 Flash"} + assert reason is None + def test_selects_opus_4_8_when_advertised(self, monkeypatch): payload = { "data": [ From 0765be9c4d82d7bb43bb6b1f3ef08c244653dbca Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Thu, 3 Sep 2026 19:22:24 +0000 Subject: [PATCH 2/6] oop --- src/ucode/smart_routing/claude_pty.py | 3 ++- tests/test_claude_smart_routing_v2.py | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/ucode/smart_routing/claude_pty.py b/src/ucode/smart_routing/claude_pty.py index a7527baf..1841497c 100644 --- a/src/ucode/smart_routing/claude_pty.py +++ b/src/ucode/smart_routing/claude_pty.py @@ -230,9 +230,10 @@ def serve() -> None: response = { "action": "block", "model": model, - "display_model": display_model, "rationale": rationale, } + if display_model != model: + 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}") diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index e9fcbf25..9a5a834f 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -170,6 +170,7 @@ def fake_run(argv, **kwargs): assert captured["argv"][-1] == "--debug" assert captured["routed_model"] == ( "system.ai.claude-sonnet-5[1m]", + "Claude Sonnet 5", "Selected for the parser task.", ) assert {definition["model"] for definition in captured["agents"].values()} == { From 9ddc2e2fcc467f5ef77822f54b64457247af5338 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Thu, 3 Sep 2026 19:22:42 +0000 Subject: [PATCH 3/6] ruff --- tests/test_claude_smart_routing_v2.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index 9a5a834f..86413260 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -213,9 +213,7 @@ 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_model_catalog", lambda *_args: (["opus"], {}, None) - ) + monkeypatch.setattr(v2, "list_anthropic_model_catalog", lambda *_args: (["opus"], {}, None)) def fake_run(_argv, **_kwargs): user_settings.write_text(json.dumps({"model": "user-selected"})) @@ -244,9 +242,7 @@ 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_model_catalog", lambda *_args: (["opus"], {}, None) - ) + monkeypatch.setattr(v2, "list_anthropic_model_catalog", lambda *_args: (["opus"], {}, None)) def fake_run(_argv, **kwargs): routed_model = "system.ai.claude-opus-4-8" From d38b9aac248b378dbd5b82ffa1cb680649e2fbf8 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Thu, 3 Sep 2026 19:33:00 +0000 Subject: [PATCH 4/6] update --- src/ucode/databricks.py | 22 ++++++++++++++++------ src/ucode/smart_routing/v2.py | 11 +++++++---- tests/test_claude_smart_routing_v2.py | 12 ++++++++---- tests/test_databricks.py | 10 ++++++---- 4 files changed, 37 insertions(+), 18 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index c708067e..d9a6a5c6 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,17 +2884,17 @@ 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. """ - model_ids, _display_names, reason = list_anthropic_model_catalog(workspace, token) - return model_ids, reason + catalog = list_anthropic_model_catalog(workspace, token) + return catalog.model_ids, catalog.error_msg def list_anthropic_model_catalog( workspace: str, token: str -) -> tuple[list[str], dict[str, str], str | None]: +) -> 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([], {}, reason) data = cast(dict, payload) if isinstance(payload, dict) else {} model_ids: list[str] = [] @@ -2901,8 +2911,8 @@ def list_anthropic_model_catalog( if isinstance(display_name, str) and display_name: display_names[model_id] = display_name if model_ids: - return model_ids, display_names, None - return [], {}, "AI Gateway returned no Anthropic model ids" + return AnthropicModelCatalog(model_ids, display_names) + return AnthropicModelCatalog([], {}, "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/v2.py b/src/ucode/smart_routing/v2.py index a653a28e..1ff563b8 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -348,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, display_names, discovery_error = list_anthropic_model_catalog(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" @@ -386,7 +389,7 @@ def route_prompt(prompt: str) -> tuple[str, str, str]: decision = _route_claude_prompt(state, token, prompt, model_ids) return ( model_name(decision.model), - display_names.get(decision.model, decision.model), + catalog.model_id_to_display_name.get(decision.model, decision.model), decision.rationale, ) diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index 86413260..d7ace400 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 @@ -115,10 +116,9 @@ def test_restores_model_captured_immediately_before_switch(self, tmp_path, monke monkeypatch.setattr( v2, "list_anthropic_model_catalog", - lambda *_args: ( + lambda *_args: AnthropicModelCatalog( ["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5"], {"system.ai.claude-sonnet-5": "Claude Sonnet 5"}, - None, ), ) monkeypatch.setattr( @@ -213,7 +213,9 @@ 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_model_catalog", lambda *_args: (["opus"], {}, None)) + monkeypatch.setattr( + v2, "list_anthropic_model_catalog", lambda *_args: AnthropicModelCatalog(["opus"], {}) + ) def fake_run(_argv, **_kwargs): user_settings.write_text(json.dumps({"model": "user-selected"})) @@ -242,7 +244,9 @@ 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_model_catalog", lambda *_args: (["opus"], {}, None)) + monkeypatch.setattr( + v2, "list_anthropic_model_catalog", lambda *_args: AnthropicModelCatalog(["opus"], {}) + ) def fake_run(_argv, **kwargs): routed_model = "system.ai.claude-opus-4-8" diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 1ed0d664..f6ced7be 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -204,11 +204,13 @@ def test_lists_anthropic_display_names_with_model_ids(self, monkeypatch): } monkeypatch.setattr(db_mod, "_http_get_json", lambda *_args, **_kwargs: (payload, None)) - models, display_names, reason = db_mod.list_anthropic_model_catalog(WS, "token") + catalog = db_mod.list_anthropic_model_catalog(WS, "token") - assert models == ["system.ai.glm-5-3-flash", "opaque-model-id"] - assert display_names == {"system.ai.glm-5-3-flash": "GLM 5.3 Flash"} - assert reason is None + 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 = { From 58b494fdd4c816b4e72db70fa954e155d634918c Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Thu, 3 Sep 2026 19:41:06 +0000 Subject: [PATCH 5/6] update --- src/ucode/databricks.py | 14 ++++++++--- src/ucode/smart_routing/claude_pty.py | 25 ++++++++++++------- src/ucode/smart_routing/v2.py | 10 ++++---- tests/test_claude_smart_routing_v2.py | 35 ++++++++++++++++++--------- 4 files changed, 56 insertions(+), 28 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index d9a6a5c6..0ed57d03 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2894,7 +2894,9 @@ def list_anthropic_model_catalog( """Return advertised Anthropic model ids and their optional display names.""" payload, reason = _get_anthropic_models_json(workspace, token) if payload is None: - return AnthropicModelCatalog([], {}, 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] = [] @@ -2911,8 +2913,14 @@ def list_anthropic_model_catalog( if isinstance(display_name, str) and display_name: display_names[model_id] = display_name if model_ids: - return AnthropicModelCatalog(model_ids, display_names) - return AnthropicModelCatalog([], {}, "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 1841497c..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.""" @@ -172,7 +182,7 @@ def first_prompt_hook_output(response: dict | None) -> dict | None: def serve_first_prompt_socket( path: Path, - route_prompt: Callable[[str], tuple[str, str] | tuple[str, str, str]], + route_prompt: Callable[[str], FirstPromptRoute], on_blocked_prompt: Callable[[str, str], None], stop: threading.Event, *, @@ -220,11 +230,9 @@ def serve() -> None: and not is_command ): routed = route_prompt(prompt) - if len(routed) == 2: - model, rationale = routed - display_model = model - else: - model, display_model, rationale = routed + model = routed.model + display_model = routed.display_model + rationale = routed.rationale if valid_model_name(model): claimed = True response = { @@ -232,8 +240,7 @@ def serve() -> None: "model": model, "rationale": rationale, } - if display_model != model: - response["display_model"] = display_model + 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}") @@ -281,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] | tuple[str, 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 1ff563b8..9b8a29eb 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -385,12 +385,12 @@ def launch_claude( model_setting = _ClaudeModelSettingGuard(user_settings_path) - def route_prompt(prompt: str) -> tuple[str, str, str]: + def route_prompt(prompt: str) -> claude_pty.FirstPromptRoute: decision = _route_claude_prompt(state, token, prompt, model_ids) - return ( - model_name(decision.model), - catalog.model_id_to_display_name.get(decision.model, 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( diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index d7ace400..ad0429d9 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -66,7 +66,9 @@ def test_blocks_once_then_allows_replay(self, tmp_path): 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, ) @@ -83,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"} @@ -117,8 +120,8 @@ def test_restores_model_captured_immediately_before_switch(self, tmp_path, monke v2, "list_anthropic_model_catalog", lambda *_args: AnthropicModelCatalog( - ["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5"], - {"system.ai.claude-sonnet-5": "Claude Sonnet 5"}, + 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( @@ -168,10 +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]", - "Claude Sonnet 5", - "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", @@ -214,7 +217,9 @@ def test_does_not_restore_when_wrapper_never_switches(self, tmp_path, monkeypatc 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_model_catalog", lambda *_args: AnthropicModelCatalog(["opus"], {}) + v2, + "list_anthropic_model_catalog", + lambda *_args: AnthropicModelCatalog(model_ids=["opus"], model_id_to_display_name={}), ) def fake_run(_argv, **_kwargs): @@ -245,7 +250,9 @@ def test_restores_after_routed_model_persists_and_preserves_later_choice( 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_model_catalog", lambda *_args: AnthropicModelCatalog(["opus"], {}) + v2, + "list_anthropic_model_catalog", + lambda *_args: AnthropicModelCatalog(model_ids=["opus"], model_id_to_display_name={}), ) def fake_run(_argv, **kwargs): @@ -473,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", ) @@ -531,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"), ) From 62c20051b33f9a118b8f89b0c34a3d6886695f25 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Thu, 3 Sep 2026 19:51:58 +0000 Subject: [PATCH 6/6] hi --- src/ucode/databricks.py | 12 +++--------- tests/test_databricks.py | 4 +--- 2 files changed, 4 insertions(+), 12 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 0ed57d03..ee54d153 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2888,15 +2888,11 @@ def list_anthropic_models(workspace: str, token: str) -> tuple[list[str], str | return catalog.model_ids, catalog.error_msg -def list_anthropic_model_catalog( - workspace: str, token: str -) -> AnthropicModelCatalog: +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 AnthropicModelCatalog( - model_ids=[], model_id_to_display_name={}, error_msg=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] = [] @@ -2913,9 +2909,7 @@ def list_anthropic_model_catalog( if isinstance(display_name, str) and display_name: display_names[model_id] = display_name if model_ids: - return AnthropicModelCatalog( - model_ids=model_ids, model_id_to_display_name=display_names - ) + return AnthropicModelCatalog(model_ids=model_ids, model_id_to_display_name=display_names) return AnthropicModelCatalog( model_ids=[], model_id_to_display_name={}, diff --git a/tests/test_databricks.py b/tests/test_databricks.py index f6ced7be..7debda88 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -207,9 +207,7 @@ def test_lists_anthropic_display_names_with_model_ids(self, monkeypatch): 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.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):