From 9f636a21980784db000b61cf1af0918eb1be9bed Mon Sep 17 00:00:00 2001 From: Pooja Date: Wed, 16 Sep 2026 04:15:49 +0000 Subject: [PATCH 1/3] Send smart-router recipe as an inference request header Stamp the configured smart-router recipe (e.g. task_v3) on every Claude Code and Codex inference request via each agent's custom-header channel, so the gateway can attribute inference traffic to the router recipe in use. Only sent when smart routing is enabled. Co-authored-by: Isaac --- src/ucode/agents/claude.py | 7 +++++++ src/ucode/agents/codex.py | 6 ++++++ src/ucode/constants.py | 2 ++ 3 files changed, 15 insertions(+) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 679779ea3..c3270346a 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -30,6 +30,7 @@ MCP_USER_SCOPE, MODEL_PROVIDER_SERVICE_HEADER, MODEL_SERVICE_PARENT_SCHEMA_HEADER, + SMART_ROUTER_RECIPE_HEADER, ) from ucode.custom_oauth import CustomOAuthConfig, build_custom_auth_shell_command from ucode.databricks import ( @@ -58,6 +59,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.tracing import tracing_env @@ -181,6 +183,7 @@ def _resolve_web_search_model(state: dict) -> str | None: "user-agent", MODEL_PROVIDER_SERVICE_HEADER.casefold(), MODEL_SERVICE_PARENT_SCHEMA_HEADER.casefold(), + SMART_ROUTER_RECIPE_HEADER.casefold(), } ) CLAUDE_TRACING_STOP_HOOK_SUFFIX = " autolog claude stop-hook" @@ -379,6 +382,10 @@ 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}") + # Stamp the router recipe on every inference request so the gateway can attribute + # traffic to it. + if smart_routing_v2.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 5113beaa0..5bb68f4a0 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, @@ -68,6 +69,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 @@ -196,6 +198,10 @@ def _provider_block( http_headers[MODEL_PROVIDER_SERVICE_HEADER] = provider elif parent_schema: http_headers[MODEL_SERVICE_PARENT_SCHEMA_HEADER] = parent_schema + # Stamp the router recipe on every inference request so the gateway can attribute + # traffic to it. + if smart_routing_v2.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 c129da93e..e2ca157d1 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 routing is on. +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 From 08c4152a885d2b63eea6c5177b35c7caf563a7c1 Mon Sep 17 00:00:00 2001 From: Pooja Date: Wed, 16 Sep 2026 17:05:26 +0000 Subject: [PATCH 2/3] address feedback --- src/ucode/agents/claude.py | 2 -- src/ucode/agents/codex.py | 2 -- src/ucode/constants.py | 2 +- 3 files changed, 1 insertion(+), 5 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index c3270346a..1602f64df 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -382,8 +382,6 @@ 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}") - # Stamp the router recipe on every inference request so the gateway can attribute - # traffic to it. if smart_routing_v2.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 diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 5bb68f4a0..fc52090d3 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -198,8 +198,6 @@ def _provider_block( http_headers[MODEL_PROVIDER_SERVICE_HEADER] = provider elif parent_schema: http_headers[MODEL_SERVICE_PARENT_SCHEMA_HEADER] = parent_schema - # Stamp the router recipe on every inference request so the gateway can attribute - # traffic to it. if smart_routing_v2.enabled(): http_headers[SMART_ROUTER_RECIPE_HEADER] = configured_router_name() return { diff --git a/src/ucode/constants.py b/src/ucode/constants.py index e2ca157d1..9cbef9c86 100644 --- a/src/ucode/constants.py +++ b/src/ucode/constants.py @@ -5,7 +5,7 @@ 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 routing is on. +# 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 From b938504f43e3c9f6c4a9870517480237f470d99c Mon Sep 17 00:00:00 2001 From: Pooja Date: Wed, 16 Sep 2026 19:39:44 +0000 Subject: [PATCH 3/3] Use renamed smart_routing_enabled() API after merge main renamed smart_routing_v2.enabled() to smart_routing_enabled(); update the two recipe-header call sites added on this branch. Co-authored-by: Isaac --- src/ucode/agents/claude.py | 2 +- src/ucode/agents/codex.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 19c86899d..4477ea0e9 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -413,7 +413,7 @@ 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.enabled(): + 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. diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 1811ee706..68d95dbea 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -199,7 +199,7 @@ 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.enabled(): + if smart_routing_v2.smart_routing_enabled(): http_headers[SMART_ROUTER_RECIPE_HEADER] = configured_router_name() return { "name": "Databricks AI Gateway",