diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index e4b519bc..17c291bf 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -28,6 +28,7 @@ MCP_USER_SCOPE, MODEL_PROVIDER_SERVICE_HEADER, MODEL_SERVICE_PARENT_SCHEMA_HEADER, + SMART_ROUTER_RECIPE_HEADER, ) from ucode.custom_oauth import ( CustomOAuthConfig, @@ -62,6 +63,7 @@ remove_smart_routing_hooks, sync_smart_routing_hooks, ) +from ucode.smart_routing.routing import configured_router_name from ucode.state import MANAGED_OVERLAY_KEY, is_tool_managed, mark_tool_managed, save_state from ucode.telemetry import agent_version, ug_version from ucode.ui import print_note, print_success, print_warning @@ -204,6 +206,7 @@ def _otel_trace_env(workspace: str) -> dict[str, str]: "user-agent", MODEL_PROVIDER_SERVICE_HEADER.casefold(), MODEL_SERVICE_PARENT_SCHEMA_HEADER.casefold(), + SMART_ROUTER_RECIPE_HEADER.casefold(), } ) # Relayed drops the user scope to deliberately omit the stale apiKeyHelper. Only applied to relayed @@ -389,6 +392,8 @@ def render_overlay( header_lines.append(f"{MODEL_PROVIDER_SERVICE_HEADER}: {provider}") elif parent_schema: header_lines.append(f"{MODEL_SERVICE_PARENT_SCHEMA_HEADER}: {parent_schema}") + if smart_routing_v2.smart_routing_enabled(): + header_lines.append(f"{SMART_ROUTER_RECIPE_HEADER}: {configured_router_name()}") # Relayed: the X-Databricks-AI-Gateway-Token swap header is added per request # by the refresh proxy, not here — a static value would go stale mid-session. custom_headers = "\n".join(header_lines) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 60491535..e2335630 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -34,6 +34,7 @@ from ucode.constants import ( MODEL_PROVIDER_SERVICE_HEADER, MODEL_SERVICE_PARENT_SCHEMA_HEADER, + SMART_ROUTER_RECIPE_HEADER, ) from ucode.custom_oauth import ( CUSTOM_OAUTH_TIMEOUT_MS, @@ -69,6 +70,7 @@ sync_smart_routing_hooks, ) from ucode.smart_routing.codex_routing import codex_model_id +from ucode.smart_routing.routing import configured_router_name from ucode.state import get_provider_service, is_tool_managed, mark_tool_managed, save_state from ucode.telemetry import agent_version, ug_version from ucode.ui import print_warning_err @@ -197,6 +199,8 @@ def _provider_block( http_headers[MODEL_PROVIDER_SERVICE_HEADER] = provider elif parent_schema: http_headers[MODEL_SERVICE_PARENT_SCHEMA_HEADER] = parent_schema + if smart_routing_v2.smart_routing_enabled(): + http_headers[SMART_ROUTER_RECIPE_HEADER] = configured_router_name() return { "name": "Databricks AI Gateway", "base_url": base_url, diff --git a/src/ucode/constants.py b/src/ucode/constants.py index c129da93..9cbef9c8 100644 --- a/src/ucode/constants.py +++ b/src/ucode/constants.py @@ -5,6 +5,8 @@ MODEL_PROVIDER_SERVICE_HEADER = "Databricks-Model-Provider-Service" MODEL_SERVICE_PARENT_SCHEMA_HEADER = "Databricks-Model-Service-Parent-Schema" +# Names the smart-router recipe (e.g. `task_v3`) in use; sent only when smart routing is enabled. +SMART_ROUTER_RECIPE_HEADER = "Databricks-Smart-Router-Recipe" # MCP server registration scopes. Claude Code supports local/project/user; the # other CLIs only take the user-scope name. Kept here (a leaf module) so both