diff --git a/README.md b/README.md index 661ab410..000a662f 100644 --- a/README.md +++ b/README.md @@ -54,7 +54,9 @@ ucode claude --enable-smart-routing ``` The flag applies only to that launch; later launches use normal model selection unless the flag is -passed again. +passed again. Smart routing uses the `task_v1` router by default. Power users can select another +router for a launch by setting `SMART_ROUTER_NAME`, for example +`SMART_ROUTER_NAME=task_v2 ucode codex --enable-smart-routing`. To configure all tools at once: diff --git a/src/ucode/smart_routing/claude_routing.py b/src/ucode/smart_routing/claude_routing.py index b755bb4d..2cd38804 100644 --- a/src/ucode/smart_routing/claude_routing.py +++ b/src/ucode/smart_routing/claude_routing.py @@ -80,7 +80,7 @@ def request_routing_decision( *, timeout: float = REQUEST_TIMEOUT_S, ) -> tuple[RoutingDecision | None, str | None]: - """Ask the workspace ``task_v1`` router for a servable Claude model. + """Ask the router for a servable Claude model. Offers the full ``cc`` menu; resolves the router's pick back to the workspace's routable id (e.g. ``system.ai.claude-opus-4-8``). @@ -96,6 +96,7 @@ def request_routing_decision( task, [(arm, "claude") for arm in CLAUDE_ROUTE_ARMS], lambda raw_model: available.get(_normalize_model(raw_model)), + router_name=routing.configured_router_name(), timeout=timeout, ) diff --git a/src/ucode/smart_routing/codex_routing.py b/src/ucode/smart_routing/codex_routing.py index b0e7cf18..ea048ef9 100644 --- a/src/ucode/smart_routing/codex_routing.py +++ b/src/ucode/smart_routing/codex_routing.py @@ -1,7 +1,7 @@ """Databricks AI Gateway routing helpers for Codex sessions and subagents. Codex-specific configuration on top of the shared :mod:`ucode.smart_routing.routing` -core: the workspace-backed ``task_v1`` route options, the ``spawn_agent`` tool +core: the workspace-backed route options, the ``spawn_agent`` tool detector, the Codex model-id translation, and the artifact paths. """ @@ -43,18 +43,19 @@ def request_routing_decision( timeout: float = REQUEST_TIMEOUT_S, log: Callable[[str], None] | None = None, ) -> tuple[RoutingDecision | None, str | None]: - """Ask the workspace ``task_v1`` router for a servable Codex model.""" + """Ask the router for a servable Codex model.""" available = {_normalize_model(model): model for model in available_models} route_options = [(model, "codex") for model in available] if not route_options: return None, "no cached model services are available" + router_name = routing.configured_router_name() if log is not None: payload = { "route_options": [ {"model": model, "harness": harness} for model, harness in route_options ], "task": {"prompt": task}, - "route_selector": {"router_name": ROUTER_NAME}, + "route_selector": {"router_name": router_name}, } url = workspace.rstrip("/") + ROUTING_PATH log(f"[ROUTE] request POST {url}: {json.dumps(payload, separators=(',', ':'))}") @@ -64,12 +65,13 @@ def request_routing_decision( task, route_options, lambda raw_model: available.get(_normalize_model(raw_model)), + router_name=router_name, timeout=timeout, ) def resolve_routed_model(raw_model: str, available_models: list[str]) -> str | None: - """Map a ``task_v1`` arm to a model the configured workspace can serve.""" + """Map a router arm to a model the configured workspace can serve.""" normalized = {_normalize_model(model): model for model in available_models} return normalized.get(_normalize_model(raw_model)) diff --git a/src/ucode/smart_routing/routing.py b/src/ucode/smart_routing/routing.py index e27401fd..0fae8be8 100644 --- a/src/ucode/smart_routing/routing.py +++ b/src/ucode/smart_routing/routing.py @@ -1,7 +1,7 @@ """Shared AI Gateway routing helpers for coding-agent sessions and subagents. Both the Codex and Claude Code integrations route through the workspace's -``task_v1`` router at ``/ai-gateway/routing/v1/routes:select``. The +configured router at ``/ai-gateway/routing/v1/routes:select``. The harness-agnostic mechanics live here — the gateway call, the decision shape, model-name normalization, and the canary/audit/decision bookkeeping. Each harness module (``codex_routing`` / ``claude_routing``) supplies its own route @@ -11,6 +11,7 @@ from __future__ import annotations import json +import os import time import urllib.error import urllib.request @@ -21,6 +22,7 @@ from typing import Any ROUTER_NAME = "task_v1" +ROUTER_NAME_ENV_VAR = "SMART_ROUTER_NAME" ROUTING_PATH = "/ai-gateway/routing/v1/routes:select" REQUEST_TIMEOUT_S = 30.0 SUBAGENT_ROUTING_DISCLAIMER = ( @@ -97,6 +99,11 @@ def normalize_model(model: str) -> str: return tail.lower() +def configured_router_name() -> str: + """Return the environment-selected router, falling back to ``task_v1``.""" + return os.environ.get(ROUTER_NAME_ENV_VAR, "").strip() or ROUTER_NAME + + def select_route( workspace: str, token: str, @@ -104,7 +111,7 @@ def select_route( route_options: Iterable[tuple[str, str | None]], resolve: Callable[[str], str | None], *, - router_name: str = ROUTER_NAME, + router_name: str, timeout: float = REQUEST_TIMEOUT_S, ) -> tuple[RoutingDecision | None, str | None]: """POST one ``routes:select`` request and resolve the router's pick. diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index f28cea01..10c9ff31 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -191,6 +191,7 @@ def _request_claude_routing_decision( prompt, [(model, "claude") for model in available], 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_routing.py b/tests/test_claude_routing.py index 9c9d5fb7..83c66e0d 100644 --- a/tests/test_claude_routing.py +++ b/tests/test_claude_routing.py @@ -25,6 +25,7 @@ def read(self) -> bytes: def test_routes_with_task_v1_claude_menu(monkeypatch): + monkeypatch.delenv("SMART_ROUTER_NAME", raising=False) captured = {} def fake_urlopen(request, timeout): diff --git a/tests/test_codex_routing.py b/tests/test_codex_routing.py index 8a41e2a6..097d8969 100644 --- a/tests/test_codex_routing.py +++ b/tests/test_codex_routing.py @@ -39,6 +39,7 @@ def read(self) -> bytes: def test_routes_with_models_from_stored_state(monkeypatch): + monkeypatch.delenv("SMART_ROUTER_NAME", raising=False) captured = {} task = "Refactor the parser" + "x" * 5000 @@ -81,6 +82,31 @@ def fake_urlopen(request, timeout): } +def test_router_name_can_be_selected_with_environment_variable(monkeypatch): + captured = {} + monkeypatch.setenv("SMART_ROUTER_NAME", " task_v2 ") + + def fake_urlopen(request, timeout): + captured["body"] = json.loads(request.data) + return _Response({"route_selection": [{"route_option": {"model": "gpt-5-6-sol"}}]}) + + monkeypatch.setattr(codex_routing.urllib.request, "urlopen", fake_urlopen) + + decision, error = codex_routing.request_routing_decision( + WS, "token", "Refactor the parser", ["system.ai.gpt-5-6-sol"] + ) + + assert error is None + assert decision is not None + assert captured["body"]["route_selector"] == {"router_name": "task_v2"} + + +def test_blank_router_name_environment_variable_uses_default(monkeypatch): + monkeypatch.setenv("SMART_ROUTER_NAME", " ") + + assert codex_routing.routing.configured_router_name() == "task_v1" + + def test_router_model_is_not_substituted_when_exact_model_is_unavailable(): model = codex_routing.resolve_routed_model( "gpt-5-6-luna", diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index fcfcd1ea..58e72126 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -462,15 +462,17 @@ def test_router_failure_keeps_original_model(self): def test_routing_request_uses_models_prompt_and_same_token(monkeypatch): + monkeypatch.delenv("SMART_ROUTER_NAME", raising=False) captured = {} logged = [] - def select_route(workspace, token, task, route_options, resolve, *, timeout): + def select_route(workspace, token, task, route_options, resolve, *, router_name, timeout): captured.update( workspace=workspace, token=token, task=task, route_options=list(route_options), + router_name=router_name, timeout=timeout, ) return ( @@ -503,6 +505,7 @@ def select_route(workspace, token, task, route_options, resolve, *, timeout): "workspace": WS, "token": "same-oauth-token", "task": "Fix the parser", + "router_name": "task_v1", "timeout": codex_routing.REQUEST_TIMEOUT_S, "route_options": [ ("kimi-k3-neo", "codex"),