Skip to content
5 changes: 5 additions & 0 deletions src/ucode/agents/claude.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
4 changes: 4 additions & 0 deletions src/ucode/agents/codex.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
2 changes: 2 additions & 0 deletions src/ucode/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading