diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 10c9ff31..84b019b3 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -3,6 +3,7 @@ import hashlib import json import os +import re import signal import socket import subprocess @@ -50,6 +51,10 @@ "Complete the delegated task exactly as requested. Follow the parent agent's instructions and " "return a concise report of your findings or changes." ) +# Keep this pattern in sync with the server-side Anthropic model prefixing logic. The prefix is +# needed because Anthropic omits models from its catalog unless the model id contains "anthropic" +# or "claude". +_ANTHROPIC_AIGW_MODEL_RE = re.compile(r"^anthropic-aigw-[0-9a-fA-F]{8}-(.+)$") def enabled() -> bool: @@ -107,6 +112,13 @@ def _canonical_claude_models(model_ids: list[str]) -> list[str]: ) +def _claude_router_model_id(model: str) -> str: + """Unwrap an Anthropic gateway id, then apply standard model normalization.""" + if match := _ANTHROPIC_AIGW_MODEL_RE.fullmatch(model): + model = match.group(1) + return routing.normalize_model(model) + + def _claude_model_overrides(model_ids: list[str]) -> dict[str, str]: overrides: dict[str, str] = {} prefix = "system.ai." @@ -182,14 +194,15 @@ def _request_claude_routing_decision( ) -> tuple[routing.RoutingDecision | None, str | None]: available: dict[str, str] = {} for model in _canonical_claude_models(model_ids): - available.setdefault(routing.normalize_model(model), model) + available.setdefault(_claude_router_model_id(model), model) if not available: return None, "Anthropic models endpoint returned no Claude models" + route_options = [(model, "claude") for model in available] return routing.select_route( workspace, token, prompt, - [(model, "claude") for model in available], + route_options, lambda selected: available.get(routing.normalize_model(selected)), router_name=routing.configured_router_name(), timeout=CLAUDE_ROUTE_SELECTION_TIMEOUT_S, diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index cbc4af43..8d722270 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -261,6 +261,54 @@ def fake_run(_argv, **kwargs): class TestSubagentRouting: + @pytest.mark.parametrize( + ("model", "expected"), + [ + ("anthropic-aigw-73ea02b2-system.ai.glm-5-2", "glm-5-2"), + ( + "anthropic-aigw-73ea02b-system.ai.glm-5-2", + "anthropic-aigw-73ea02b-system.ai.glm-5-2", + ), + ("system.ai.claude-opus-4-8", "claude-opus-4-8"), + ], + ) + def test_normalizes_router_model_id(self, model, expected): + assert v2._claude_router_model_id(model) == expected + + def test_routes_anthropic_gateway_alias_by_embedded_model_id(self, monkeypatch): + gateway_alias = "anthropic-aigw-73ea02b2-system.ai.glm-5-2" + captured = {} + monkeypatch.setenv("SMART_ROUTER_NAME", "task_v2") + + def fake_select(workspace, token, task, route_options, resolve, **kwargs): + captured["route_options"] = list(route_options) + captured["router_name"] = kwargs["router_name"] + return ( + routing.RoutingDecision( + model=resolve("glm-5-2"), + raw_model="glm-5-2", + ), + None, + ) + + monkeypatch.setattr(routing, "select_route", fake_select) + decision, error = v2._request_claude_routing_decision( + "https://example.com", + "secret-token", + "inspect the parser", + ["system.ai.claude-opus-4-8", gateway_alias], + ) + + assert error is None + assert decision.model == gateway_alias + assert captured == { + "route_options": [ + ("claude-opus-4-8", "claude"), + ("glm-5-2", "claude"), + ], + "router_name": "task_v2", + } + def test_routes_agent_prompt_with_initialized_model_menu(self, tmp_path, monkeypatch): captured = {} decisions_path = tmp_path / "decisions.jsonl"