From 9108bad0aeb46b625419c920bfb2a031aea7fc30 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Fri, 11 Sep 2026 02:18:36 +0000 Subject: [PATCH 1/5] update --- src/ucode/cli.py | 67 ++++++++---- src/ucode/smart_routing/claude_routing.py | 44 +++++++- src/ucode/smart_routing/codex_interposer.py | 8 +- src/ucode/smart_routing/codex_routing.py | 46 +++++--- src/ucode/smart_routing/routing.py | 112 ++++++++++++++++++-- src/ucode/smart_routing/v2.py | 105 +++++++++++++++++- tests/test_claude_routing.py | 91 ++++++++++++++++ tests/test_claude_smart_routing_v2.py | 85 +++++++++++++-- tests/test_cli.py | 19 ++++ tests/test_codex_routing.py | 77 ++++++++++++++ tests/test_codex_smart_routing_v2.py | 12 ++- 11 files changed, 602 insertions(+), 64 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index ef4e2231..aeba7c5d 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1743,6 +1743,26 @@ def _smart_routing_v2_flag(enabled: bool) -> Iterator[None]: os.environ[smart_routing_v2.ENV_VAR] = previous +@contextmanager +def _disable_smart_routing_for_subcommand(tool: str, ctx: typer.Context) -> Iterator[None]: + """Keep native agent subcommands out of every smart-routing path. + + The environment flag is also consulted during bootstrap/version checks, + before the final launch options are built. Native positional subcommands + must therefore suppress the flag for the whole ucode launch flow. An + explicit prompt after `--` remains eligible for routing. + """ + if _smart_routing_launch_shape(tool, ctx.args, _has_explicit_prompt(ctx)): + yield + return + previous = os.environ.pop(smart_routing_v2.ENV_VAR, None) + try: + yield + finally: + if previous is not None: + os.environ[smart_routing_v2.ENV_VAR] = previous + + def _migrate_legacy_smart_routing(state: dict) -> dict: """Remove the former persisted opt-in and its permanent routing hooks.""" if smart_routing_v2.LEGACY_STATE_KEY not in state: @@ -1921,6 +1941,13 @@ def _should_launch_smart_routing( ) -> bool: if model is not None or has_explicit_model_arg(tool_args): return False + return _smart_routing_launch_shape(tool, tool_args, explicit_prompt) + + +def _smart_routing_launch_shape( + tool: str, tool_args: list[str], explicit_prompt: bool +) -> bool: + """Whether the forwarded arguments represent an interactive launch.""" if not tool_args or explicit_prompt: return True return tool == "claude" and tool_args[0].startswith("-") @@ -2473,15 +2500,16 @@ def codex_cmd( print_success("Codex smart routing disabled; ug routing hooks removed") return with _smart_routing_v2_flag(enable_smart_routing_flag): - _launch_tool( - "codex", - ctx, - provider=provider, - refresh=refresh, - skip_preflight=skip_preflight, - workspace_url=workspace, - custom_oauth=custom_oauth, - ) + with _disable_smart_routing_for_subcommand("codex", ctx): + _launch_tool( + "codex", + ctx, + provider=provider, + refresh=refresh, + skip_preflight=skip_preflight, + workspace_url=workspace, + custom_oauth=custom_oauth, + ) @app.command( @@ -2574,16 +2602,17 @@ def claude_cmd( if enable_model_discovery: os.environ[claude_agent.GATEWAY_MODEL_DISCOVERY_ENV_VAR] = "1" with _smart_routing_v2_flag(enable_smart_routing_flag): - _launch_tool( - "claude", - ctx, - provider=provider, - model=model, - refresh=refresh, - skip_preflight=skip_preflight, - workspace_url=workspace, - custom_oauth=custom_oauth, - ) + with _disable_smart_routing_for_subcommand("claude", ctx): + _launch_tool( + "claude", + ctx, + provider=provider, + model=model, + refresh=refresh, + skip_preflight=skip_preflight, + workspace_url=workspace, + custom_oauth=custom_oauth, + ) @app.command("gemini", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) diff --git a/src/ucode/smart_routing/claude_routing.py b/src/ucode/smart_routing/claude_routing.py index cbdd5cf2..58ac837d 100644 --- a/src/ucode/smart_routing/claude_routing.py +++ b/src/ucode/smart_routing/claude_routing.py @@ -9,12 +9,17 @@ from __future__ import annotations +import os + # Re-exported so tests can patch the shared ``urlopen`` seam via # ``claude_routing.urllib.request`` — the call lives in ``routing``, but Python # modules are singletons so patching this name patches the one call site. import urllib.request # noqa: F401 +from collections.abc import Callable, Mapping +from pathlib import Path from typing import Any +from ucode import config_io from ucode.config_io import APP_DIR from ucode.smart_routing import routing from ucode.smart_routing.routing import RoutingDecision @@ -29,10 +34,16 @@ CANARY_PATH = APP_DIR / "claude-smart-routing-canary.json" AUDIT_PATH = APP_DIR / "claude-smart-routing-audit.jsonl" DECISIONS_PATH = APP_DIR / "claude-smart-routing-decisions.jsonl" +REQUESTS_LOG_FILENAME = "claude-smart-routing-requests.jsonl" _normalize_model = routing.normalize_model +def request_log_path() -> Path: + """Return the Claude Code smart-routing request log path.""" + return config_io.APP_DIR / REQUESTS_LOG_FILENAME + + # Claude Code CLI options that consume a following value (from `claude --help`); # their values must not be mistaken for the seed prompt. Options whose value is # optional (`-c`/`-d`/`-r`/`-w`) are intentionally omitted — treating them as @@ -79,6 +90,8 @@ def request_routing_decision( available_models: list[str], *, timeout: float = REQUEST_TIMEOUT_S, + log: Callable[[str], None] | None = None, + extra_headers: Mapping[str, str] | None = None, ) -> tuple[RoutingDecision | None, str | None]: """Ask the router for a servable Claude model. @@ -90,14 +103,26 @@ def request_routing_decision( if missing: return None, f"required Claude routing models are unavailable: {', '.join(missing)}" + route_options = [(arm, "claude") for arm in CLAUDE_ROUTE_ARMS] + router_name = routing.configured_router_name() + headers = routing.route_request_headers(token, extra_headers) + routing.log_route_request( + workspace, + routing.route_request_body(task, route_options, router_name=router_name), + headers=headers, + log=log, + request_log_path=request_log_path(), + ) + select_kwargs: dict[str, Any] = {"router_name": router_name, "timeout": timeout} + if extra_headers: + select_kwargs["extra_headers"] = extra_headers return routing.select_route( workspace, token, task, - [(arm, "claude") for arm in CLAUDE_ROUTE_ARMS], + route_options, lambda raw_model: available.get(_normalize_model(raw_model)), - router_name=routing.configured_router_name(), - timeout=timeout, + **select_kwargs, ) @@ -109,6 +134,7 @@ def route_pre_tool_use( available_models: list[str], timeout: float = REQUEST_TIMEOUT_S, audit_decision: bool = False, + extra_headers: Mapping[str, str] | None = None, ) -> dict[str, Any] | None: """Route one Claude Code ``Agent`` (subagent-spawn) call, rewriting its model.""" record = None @@ -121,7 +147,15 @@ def record(payload, task, decision, requested): payload, is_spawn_agent=is_spawn_agent_tool, decision_fn=lambda task: request_routing_decision( - workspace, token, task, available_models, timeout=timeout + workspace, + token, + task, + available_models, + timeout=timeout, + extra_headers=extra_headers + or routing.route_forward_headers_from_lines( + os.environ.get("ANTHROPIC_CUSTOM_HEADERS") + ), ), default_task_label="Claude Code subagent task", model_id_mapper=_claude_model_id, @@ -148,7 +182,7 @@ def record_subagent_start(payload: dict[str, Any]) -> dict[str, Any]: def clear_routing_artifacts() -> None: """Remove ucode-owned routing canary and audit files.""" - routing.clear_artifacts((CANARY_PATH, AUDIT_PATH, DECISIONS_PATH)) + routing.clear_artifacts((CANARY_PATH, AUDIT_PATH, DECISIONS_PATH, request_log_path())) def _claude_model_id(model: str) -> str: diff --git a/src/ucode/smart_routing/codex_interposer.py b/src/ucode/smart_routing/codex_interposer.py index 5d39b5d1..e71f1ddf 100644 --- a/src/ucode/smart_routing/codex_interposer.py +++ b/src/ucode/smart_routing/codex_interposer.py @@ -6,7 +6,7 @@ import threading import time import uuid -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass, replace from pathlib import Path @@ -210,6 +210,7 @@ async def _handle_tui( available_models: list[str] | None = None, workspace: str | None = None, token_provider: TokenProvider | None = None, + route_headers: Mapping[str, str] | None = None, switch_message_fn: SwitchMessageFn | None = None, ) -> None: path = getattr(getattr(tui, "request", None), "path", "/") or "/" @@ -229,6 +230,7 @@ def route_decision(prompt: str): prompt, list(available_models or []), log=log, + extra_headers=route_headers, ) sess = _Session( @@ -299,6 +301,7 @@ async def _serve( available_models: list[str] | None = None, workspace: str | None = None, token_provider: TokenProvider | None = None, + route_headers: Mapping[str, str] | None = None, switch_message_fn: SwitchMessageFn | None = None, ): async def handler(tui): @@ -312,6 +315,7 @@ async def handler(tui): available_models, workspace, token_provider, + route_headers, switch_message_fn, ) except Exception as exc: # noqa: BLE001 @@ -331,6 +335,7 @@ def start_interposer_thread( available_models: list[str] | None = None, workspace: str | None = None, token_provider: TokenProvider | None = None, + route_headers: Mapping[str, str] | None = None, switch_message_fn: SwitchMessageFn | None = None, switch_message: str | None = None, log_path: Path | None = None, @@ -364,6 +369,7 @@ def run() -> None: available_models, workspace, token_provider, + route_headers, switch_message_fn, ) ) diff --git a/src/ucode/smart_routing/codex_routing.py b/src/ucode/smart_routing/codex_routing.py index ea048ef9..fb0d2853 100644 --- a/src/ucode/smart_routing/codex_routing.py +++ b/src/ucode/smart_routing/codex_routing.py @@ -7,16 +7,17 @@ from __future__ import annotations -import json import re # Re-exported so tests can patch the shared ``urlopen`` seam via # ``codex_routing.urllib.request`` — the actual call lives in ``routing``, but # Python modules are singletons so patching this name patches the one call site. import urllib.request # noqa: F401 -from collections.abc import Callable +from collections.abc import Callable, Mapping +from pathlib import Path from typing import Any +from ucode import config_io from ucode.config_io import APP_DIR from ucode.smart_routing import routing from ucode.smart_routing.routing import RoutingDecision @@ -28,12 +29,18 @@ CANARY_PATH = APP_DIR / "codex-smart-routing-canary.json" AUDIT_PATH = APP_DIR / "codex-smart-routing-audit.jsonl" DECISIONS_PATH = APP_DIR / "codex-smart-routing-decisions.jsonl" +REQUESTS_LOG_FILENAME = "codex-smart-routing-requests.jsonl" _GPT_RE = re.compile(r"gpt-(\d+)(?:[.-](\d+))?(?:[.-](\d+))?(-.+|[a-z].*)?") _normalize_model = routing.normalize_model +def request_log_path() -> Path: + """Return the Codex smart-routing request log path.""" + return config_io.APP_DIR / REQUESTS_LOG_FILENAME + + def request_routing_decision( workspace: str, token: str, @@ -42,6 +49,7 @@ def request_routing_decision( *, timeout: float = REQUEST_TIMEOUT_S, log: Callable[[str], None] | None = None, + extra_headers: Mapping[str, str] | None = None, ) -> tuple[RoutingDecision | None, str | None]: """Ask the router for a servable Codex model.""" available = {_normalize_model(model): model for model in available_models} @@ -49,24 +57,24 @@ def request_routing_decision( 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}, - } - url = workspace.rstrip("/") + ROUTING_PATH - log(f"[ROUTE] request POST {url}: {json.dumps(payload, separators=(',', ':'))}") + headers = routing.route_request_headers(token, extra_headers) + routing.log_route_request( + workspace, + routing.route_request_body(task, route_options, router_name=router_name), + headers=headers, + log=log, + request_log_path=request_log_path(), + ) + select_kwargs: dict[str, Any] = {"router_name": router_name, "timeout": timeout} + if extra_headers: + select_kwargs["extra_headers"] = extra_headers return routing.select_route( workspace, token, task, route_options, lambda raw_model: available.get(_normalize_model(raw_model)), - router_name=router_name, - timeout=timeout, + **select_kwargs, ) @@ -84,6 +92,7 @@ def route_pre_tool_use( available_models: list[str], timeout: float = REQUEST_TIMEOUT_S, audit_decision: bool = False, + extra_headers: Mapping[str, str] | None = None, ) -> dict[str, Any] | None: """Route one Codex ``spawn_agent`` call and rewrite its model.""" record = None @@ -96,7 +105,12 @@ def record(payload, task, decision, requested): payload, is_spawn_agent=is_spawn_agent_tool, decision_fn=lambda task: request_routing_decision( - workspace, token, task, available_models, timeout=timeout + workspace, + token, + task, + available_models, + timeout=timeout, + extra_headers=extra_headers, ), default_task_label="Codex subagent task", model_id_mapper=codex_model_id, @@ -124,7 +138,7 @@ def record_subagent_start(payload: dict[str, Any]) -> dict[str, Any]: def clear_routing_artifacts() -> None: """Remove ucode-owned routing canary and audit files.""" - routing.clear_artifacts((CANARY_PATH, AUDIT_PATH, DECISIONS_PATH)) + routing.clear_artifacts((CANARY_PATH, AUDIT_PATH, DECISIONS_PATH, request_log_path())) def _parse_gpt(model: str) -> tuple[int, int, int, str] | None: diff --git a/src/ucode/smart_routing/routing.py b/src/ucode/smart_routing/routing.py index 5ad46fe3..d1c149eb 100644 --- a/src/ucode/smart_routing/routing.py +++ b/src/ucode/smart_routing/routing.py @@ -16,7 +16,7 @@ import urllib.error import urllib.request import uuid -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Mapping from dataclasses import dataclass from pathlib import Path from typing import Any @@ -25,6 +25,22 @@ ROUTER_NAME_ENV_VAR = "SMART_ROUTER_NAME" ROUTING_PATH = "/ai-gateway/routing/v1/routes:select" REQUEST_TIMEOUT_S = 30.0 +ROUTE_FORWARD_HEADER_NAMES = frozenset( + { + "databricks-ai-gateway-request-tags", + "user-agent", + "x-databricks-traffic-id", + "x-databricks-use-coding-agent-mode", + } +) +SENSITIVE_HEADER_NAMES = frozenset( + { + "authorization", + "cookie", + "proxy-authorization", + "x-databricks-ai-gateway-token", + } +) SUBAGENT_ROUTING_DISCLAIMER = ( "Spawned subagents are routed independently based on their own complexity." ) @@ -104,6 +120,81 @@ def configured_router_name() -> str: return os.environ.get(ROUTER_NAME_ENV_VAR, "").strip() or ROUTER_NAME +def route_request_body( + task: str, + route_options: Iterable[tuple[str, str | None]], + *, + router_name: str, +) -> dict[str, Any]: + """Return the JSON body sent to ``routes:select``.""" + return { + "route_options": [{"model": model, "harness": harness} for model, harness in route_options], + "task": {"prompt": task}, + "route_selector": {"router_name": router_name}, + } + + +def route_forward_headers_from_lines(headers: object) -> dict[str, str]: + """Parse newline-delimited gateway headers to forward to ``routes:select``.""" + if not isinstance(headers, str): + return {} + forwarded: dict[str, str] = {} + for line in headers.splitlines(): + name, separator, value = line.partition(":") + normalized = name.strip().casefold() + if not separator or normalized not in ROUTE_FORWARD_HEADER_NAMES: + continue + clean_name = name.strip() + clean_value = value.strip() + if clean_name and clean_value: + forwarded[clean_name] = clean_value + return forwarded + + +def route_request_headers( + token: str, + extra_headers: Mapping[str, str] | None = None, +) -> dict[str, str]: + """Return HTTP headers for ``routes:select`` without allowing auth overrides.""" + headers: dict[str, str] = {} + for name, value in (extra_headers or {}).items(): + normalized = name.strip().casefold() + if normalized not in ROUTE_FORWARD_HEADER_NAMES: + continue + clean_name = name.strip() + clean_value = value.strip() + if clean_name and clean_value: + headers[clean_name] = clean_value + headers["Authorization"] = f"Bearer {token}" + headers["Content-Type"] = "application/json" + return headers + + +def log_route_request( + workspace: str, + body: dict[str, Any], + *, + headers: Mapping[str, str] | None = None, + log: Callable[[str], None] | None = None, + request_log_path: Path | None = None, +) -> None: + """Record an outgoing route-selection request without credentials.""" + url = workspace.rstrip("/") + ROUTING_PATH + if log is not None: + log(f"[ROUTE] request POST {url}: {json.dumps(body, separators=(',', ':'))}") + if request_log_path is not None: + _append_jsonl( + request_log_path, + { + "at": time.time(), + "method": "POST", + "url": url, + "headers": _redacted_headers(headers or {}), + "body": body, + }, + ) + + def select_route( workspace: str, token: str, @@ -113,6 +204,7 @@ def select_route( *, router_name: str, timeout: float = REQUEST_TIMEOUT_S, + extra_headers: Mapping[str, str] | None = None, ) -> tuple[RoutingDecision | None, str | None]: """POST one ``routes:select`` request and resolve the router's pick. @@ -122,18 +214,11 @@ def select_route( None when the arm is unservable. Returns ``(decision, error)``; a failed call yields ``(None, reason)`` so callers can fail open. """ - body = { - "route_options": [{"model": model, "harness": harness} for model, harness in route_options], - "task": {"prompt": task}, - "route_selector": {"router_name": router_name}, - } + body = route_request_body(task, route_options, router_name=router_name) request = urllib.request.Request( workspace.rstrip("/") + ROUTING_PATH, data=json.dumps(body).encode("utf-8"), - headers={ - "Authorization": f"Bearer {token}", - "Content-Type": "application/json", - }, + headers=route_request_headers(token, extra_headers), method="POST", ) try: @@ -342,6 +427,13 @@ def _selected_model(payload: Any) -> str | None: return model if isinstance(model, str) and model else None +def _redacted_headers(headers: Mapping[str, str]) -> dict[str, str]: + redacted: dict[str, str] = {} + for name, value in headers.items(): + redacted[name] = "[REDACTED]" if name.strip().casefold() in SENSITIVE_HEADER_NAMES else value + return redacted + + def _pending_decision( decisions_path: Path, audit_path: Path, session_id: Any, actual_model: Any ) -> dict[str, Any] | None: diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index e8ce72c6..41caa3a7 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -38,6 +38,7 @@ LEGACY_STATE_KEY = "smart_routing_enabled" CODEX_INTERPOSER_LOG = APP_DIR / "codex-v2-interposer.log" +CODEX_CUSTOM_HEADERS_ENV_VAR = "CODEX_CUSTOM_HEADERS" CLAUDE_TARGET_MODEL = "system.ai.claude-sonnet-4-6[1m]" # TODO(lilly): replace with smart router. CLAUDE_PTY_LOG = APP_DIR / "claude-v2-pty.log" @@ -239,6 +240,9 @@ def _request_claude_routing_decision( token: str, prompt: str, model_ids: list[str], + *, + log: Callable[[str], None] | None = None, + extra_headers: dict[str, str] | None = None, ) -> tuple[routing.RoutingDecision | None, str | None]: available: dict[str, str] = {} for model in _canonical_claude_models(model_ids): @@ -246,14 +250,28 @@ def _request_claude_routing_decision( if not available: return None, "Anthropic models endpoint returned no Claude models" route_options = [(model, "claude") for model in available] + router_name = routing.configured_router_name() + headers = routing.route_request_headers(token, extra_headers) + routing.log_route_request( + workspace, + routing.route_request_body(prompt, route_options, router_name=router_name), + headers=headers, + log=log, + request_log_path=claude_routing.request_log_path(), + ) + select_kwargs: dict[str, object] = { + "router_name": router_name, + "timeout": CLAUDE_ROUTE_SELECTION_TIMEOUT_S, + } + if extra_headers: + select_kwargs["extra_headers"] = extra_headers return routing.select_route( workspace, token, prompt, route_options, lambda selected: available.get(_claude_router_model_id(selected)), - router_name=routing.configured_router_name(), - timeout=CLAUDE_ROUTE_SELECTION_TIMEOUT_S, + **select_kwargs, ) @@ -262,6 +280,9 @@ def _route_claude_prompt( token: str, prompt: str, model_ids: list[str] | None = None, + *, + log: Callable[[str], None] | None = None, + extra_headers: dict[str, str] | None = None, ) -> routing.RoutingDecision: workspace = state.get("workspace") if not isinstance(workspace, str): @@ -273,7 +294,14 @@ def _route_claude_prompt( raise RuntimeError( discovery_error or "Anthropic models endpoint returned no Claude models" ) - decision, error = _request_claude_routing_decision(workspace, token, prompt, model_ids) + decision, error = _request_claude_routing_decision( + workspace, + token, + prompt, + model_ids, + log=log, + extra_headers=extra_headers, + ) if decision is None: raise RuntimeError(error or "router returned no Claude model selection") return decision @@ -286,13 +314,21 @@ def route_claude_pre_tool_use( token: str, available_models: list[str], audit_decision: bool = False, + extra_headers: dict[str, str] | None = None, ) -> dict | None: """Route a Claude Agent call through a transient exact-model agent definition.""" route = routing.resolve_spawn_route( payload, is_spawn_agent=claude_routing.is_spawn_agent_tool, decision_fn=lambda task: _request_claude_routing_decision( - workspace, token, task, available_models + workspace, + token, + task, + available_models, + extra_headers=extra_headers + or routing.route_forward_headers_from_lines( + os.environ.get("ANTHROPIC_CUSTOM_HEADERS") + ), ), default_task_label="Claude Code subagent task", model_id_mapper=lambda model: model, @@ -416,6 +452,10 @@ def launch_claude( raise RuntimeError("Claude settings 'env' must be an object for smart routing.") env.pop("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", None) env[FIRST_PROMPT_SOCKET_ENV] = str(socket_path) + route_headers = { + **routing.route_forward_headers_from_lines(os.environ.get("ANTHROPIC_CUSTOM_HEADERS")), + **routing.route_forward_headers_from_lines(env.get("ANTHROPIC_CUSTOM_HEADERS")), + } model_overrides = settings.setdefault("modelOverrides", {}) if not isinstance(model_overrides, dict): raise RuntimeError("Claude settings 'modelOverrides' must be an object for smart routing.") @@ -433,8 +473,23 @@ def launch_claude( model_setting = _ClaudeModelSettingGuard(user_settings_path) + def log_route_request(message: str) -> None: + try: + CLAUDE_PTY_LOG.parent.mkdir(parents=True, exist_ok=True) + with open(CLAUDE_PTY_LOG, "a", encoding="utf-8") as handle: + handle.write(f"{time.strftime('%H:%M:%S')} {message}\n") + except OSError: + pass + def route_prompt(prompt: str) -> claude_pty.FirstPromptRoute: - decision = _route_claude_prompt(state, token, prompt, model_ids) + decision = _route_claude_prompt( + state, + token, + prompt, + model_ids, + log=log_route_request, + extra_headers=route_headers, + ) return claude_pty.FirstPromptRoute( model=model_name(_unwrapped_claude_model_id(decision.model)), display_model=catalog.model_id_to_display_name.get(decision.model, decision.model), @@ -486,6 +541,43 @@ def _v2_pre_tool_use_hooks(state: dict, available_models: list[str]) -> list[dic ) +def _codex_route_headers(overlay: dict) -> dict[str, str]: + providers = overlay.get("model_providers") + if not isinstance(providers, dict): + return {} + provider = providers.get("ucode-databricks") + if not isinstance(provider, dict): + return {} + headers = provider.get("http_headers") + if not isinstance(headers, dict): + return {} + return { + str(name): str(value) + for name, value in headers.items() + if isinstance(name, str) and isinstance(value, str) + } + + +def _apply_codex_custom_headers(overlay: dict) -> dict[str, str]: + """Add shell-provided gateway headers to the Codex provider overlay.""" + custom_headers = routing.route_forward_headers_from_lines( + os.environ.get(CODEX_CUSTOM_HEADERS_ENV_VAR) + ) + if not custom_headers: + return overlay + providers = overlay.get("model_providers") + if not isinstance(providers, dict): + return overlay + provider = providers.get("ucode-databricks") + if not isinstance(provider, dict): + return overlay + headers = provider.setdefault("http_headers", {}) + if not isinstance(headers, dict): + return overlay + headers.update(custom_headers) + return overlay + + def launch_codex( state: dict, tool_args: list[str], @@ -518,6 +610,8 @@ def launch_codex( state.get("profile"), use_pat=bool(state.get("use_pat")), ) + _apply_codex_custom_headers(overlay) + route_headers = _codex_route_headers(overlay) overlay["hooks"] = { "PreToolUse": _v2_pre_tool_use_hooks(state, available_models), } @@ -546,6 +640,7 @@ def launch_codex( available_models=available_models, workspace=workspace, token_provider=lambda: get_databricks_token(workspace, profile), + route_headers=route_headers, switch_message_fn=format_routing_notice, log_path=CODEX_INTERPOSER_LOG, ) diff --git a/tests/test_claude_routing.py b/tests/test_claude_routing.py index ac41726d..6345daaa 100644 --- a/tests/test_claude_routing.py +++ b/tests/test_claude_routing.py @@ -70,6 +70,97 @@ def fake_urlopen(request, timeout): } +def test_routes_select_request_is_logged(monkeypatch, tmp_path): + monkeypatch.delenv("SMART_ROUTER_NAME", raising=False) + monkeypatch.setattr(claude_routing.config_io, "APP_DIR", tmp_path) + logged = [] + + monkeypatch.setattr( + claude_routing.urllib.request, + "urlopen", + lambda request, timeout: _Response( + {"route_selection": [{"route_option": {"model": "claude-sonnet-5"}}]} + ), + ) + + decision, error = claude_routing.request_routing_decision( + WS, + "secret-token", + "Map the codebase", + ["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5"], + log=logged.append, + ) + + assert error is None + assert decision is not None + assert len(logged) == 1 + assert logged[0].startswith(f"[ROUTE] request POST {WS}/ai-gateway/routing/v1/routes:select: ") + logged_body = json.loads(logged[0].split(": ", 1)[1]) + record = json.loads((tmp_path / claude_routing.REQUESTS_LOG_FILENAME).read_text()) + assert record["method"] == "POST" + assert record["url"] == f"{WS}/ai-gateway/routing/v1/routes:select" + assert logged_body == record["body"] == { + "route_options": [ + {"model": "claude-opus-4-8", "harness": "claude"}, + {"model": "claude-sonnet-5", "harness": "claude"}, + ], + "task": {"prompt": "Map the codebase"}, + "route_selector": {"router_name": claude_routing.ROUTER_NAME}, + } + assert "secret-token" not in logged[0] + assert "secret-token" not in (tmp_path / claude_routing.REQUESTS_LOG_FILENAME).read_text() + + +def test_routes_select_forwards_gateway_headers_without_auth_override(monkeypatch, tmp_path): + monkeypatch.setattr(claude_routing.config_io, "APP_DIR", tmp_path) + captured = {} + + def fake_urlopen(request, timeout): + captured["headers"] = {key.casefold(): value for key, value in request.headers.items()} + return _Response({"route_selection": [{"route_option": {"model": "claude-sonnet-5"}}]}) + + monkeypatch.setattr(claude_routing.urllib.request, "urlopen", fake_urlopen) + + decision, error = claude_routing.request_routing_decision( + WS, + "real-token", + "Map the codebase", + ["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5"], + extra_headers=claude_routing.routing.route_forward_headers_from_lines( + "\n".join( + [ + "x-databricks-traffic-id: testenv://liteswap/arnav-r315-task-v3", + "x-databricks-use-coding-agent-mode: true", + 'Databricks-Ai-Gateway-Request-Tags: {"source":"isaac-cli"}', + "Authorization: Bearer wrong-token", + "X-Ignored: no", + ] + ) + ), + ) + + assert error is None + assert decision is not None + assert captured["headers"]["authorization"] == "Bearer real-token" + assert ( + captured["headers"]["x-databricks-traffic-id"] + == "testenv://liteswap/arnav-r315-task-v3" + ) + assert captured["headers"]["x-databricks-use-coding-agent-mode"] == "true" + assert ( + captured["headers"]["databricks-ai-gateway-request-tags"] + == '{"source":"isaac-cli"}' + ) + assert "x-ignored" not in captured["headers"] + record = json.loads((tmp_path / claude_routing.REQUESTS_LOG_FILENAME).read_text()) + assert record["headers"]["Authorization"] == "[REDACTED]" + assert ( + record["headers"]["x-databricks-traffic-id"] + == "testenv://liteswap/arnav-r315-task-v3" + ) + assert "wrong-token" not in (tmp_path / claude_routing.REQUESTS_LOG_FILENAME).read_text() + + def test_missing_arm_short_circuits_without_calling_router(monkeypatch): def fail(*args, **kwargs): raise AssertionError("router must not be called when an arm is missing") diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index 2a2c8eaa..5f2075d7 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -172,7 +172,21 @@ def test_strips_gateway_prefix_for_interposer(self): def test_restores_model_captured_immediately_before_switch(self, tmp_path, monkeypatch): ucode_settings = tmp_path / "ucode-settings.json" user_settings = tmp_path / "settings.json" - ucode_settings.write_text(json.dumps({"env": {"ANTHROPIC_BASE_URL": "https://gw"}})) + ucode_settings.write_text( + json.dumps( + { + "env": { + "ANTHROPIC_BASE_URL": "https://gw", + "ANTHROPIC_CUSTOM_HEADERS": "\n".join( + [ + "x-databricks-traffic-id: testenv://liteswap/arnav-r315-task-v3", + 'Databricks-Ai-Gateway-Request-Tags: {"source":"isaac-cli"}', + ] + ), + } + } + ) + ) user_settings.write_text(json.dumps({"model": "opus", "theme": "dark"})) monkeypatch.setattr(claude, "APP_DIR", tmp_path) monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", ucode_settings) @@ -189,15 +203,17 @@ def test_restores_model_captured_immediately_before_switch(self, tmp_path, monke model_id_to_display_name={"system.ai.claude-sonnet-5": "Claude Sonnet 5"}, ), ) - monkeypatch.setattr( - v2, - "_route_claude_prompt", - lambda *_args: v2.routing.RoutingDecision( + routed_call = {} + + def route_claude_prompt(*_args, **kwargs): + routed_call.update(kwargs) + return v2.routing.RoutingDecision( model="system.ai.claude-sonnet-5", raw_model="claude-sonnet-5", rationale="Selected for the parser task.", - ), - ) + ) + + monkeypatch.setattr(v2, "_route_claude_prompt", route_claude_prompt) captured: dict = {} def fake_run(argv, **kwargs): @@ -241,6 +257,12 @@ def fake_run(argv, **kwargs): display_model="Claude Sonnet 5", rationale="Selected for the parser task.", ) + assert routed_call["extra_headers"]["x-databricks-traffic-id"] == ( + "testenv://liteswap/arnav-r315-task-v3" + ) + assert routed_call["extra_headers"]["Databricks-Ai-Gateway-Request-Tags"] == ( + '{"source":"isaac-cli"}' + ) assert {definition["model"] for definition in captured["agents"].values()} == { "system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5", @@ -399,10 +421,55 @@ def fake_select(workspace, token, task, route_options, resolve, **kwargs): "router_name": "task_v2", } + def test_logs_v2_routes_select_request(self, tmp_path, monkeypatch): + monkeypatch.setattr(v2.claude_routing.config_io, "APP_DIR", tmp_path) + monkeypatch.setenv("SMART_ROUTER_NAME", "task_v2") + logged = [] + + def fake_select(workspace, token, task, route_options, resolve, **kwargs): + return ( + routing.RoutingDecision( + model=resolve("claude-sonnet-5"), + raw_model="claude-sonnet-5", + ), + 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", "system.ai.claude-sonnet-5"], + log=logged.append, + ) + + assert error is None + assert decision is not None + assert len(logged) == 1 + assert logged[0].startswith( + "[ROUTE] request POST https://example.com/ai-gateway/routing/v1/routes:select: " + ) + record = json.loads((tmp_path / v2.claude_routing.REQUESTS_LOG_FILENAME).read_text()) + assert json.loads(logged[0].split(": ", 1)[1]) == record["body"] == { + "route_options": [ + {"model": "claude-opus-4-8", "harness": "claude"}, + {"model": "claude-sonnet-5", "harness": "claude"}, + ], + "task": {"prompt": "inspect the parser"}, + "route_selector": {"router_name": "task_v2"}, + } + assert "secret-token" not in logged[0] + assert "secret-token" not in (tmp_path / v2.claude_routing.REQUESTS_LOG_FILENAME).read_text() + def test_routes_agent_prompt_with_initialized_model_menu(self, tmp_path, monkeypatch): captured = {} decisions_path = tmp_path / "decisions.jsonl" monkeypatch.setattr(v2.claude_routing, "DECISIONS_PATH", decisions_path) + monkeypatch.setenv( + "ANTHROPIC_CUSTOM_HEADERS", + "x-databricks-traffic-id: testenv://liteswap/arnav-r315-task-v3", + ) def fake_select(workspace, token, task, route_options, resolve, **kwargs): captured.update( @@ -410,6 +477,7 @@ def fake_select(workspace, token, task, route_options, resolve, **kwargs): token=token, task=task, route_options=list(route_options), + extra_headers=kwargs.get("extra_headers"), ) return ( routing.RoutingDecision( @@ -442,6 +510,9 @@ def fake_select(workspace, token, task, route_options, resolve, **kwargs): ("claude-opus-4-8", "claude"), ("claude-sonnet-5", "claude"), ], + "extra_headers": { + "x-databricks-traffic-id": "testenv://liteswap/arnav-r315-task-v3" + }, } updated_input = output["hookSpecificOutput"]["updatedInput"] assert "model" not in updated_input diff --git a/tests/test_cli.py b/tests/test_cli.py index 00a01baf..38c22702 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -451,6 +451,25 @@ def test_codex_enable_smart_routing_is_consumed_by_ucode(self): assert "ENABLE_SMART_ROUTING_V2" not in os.environ assert mock_launch.call_args.args[1].args == [] + @pytest.mark.parametrize("tool, subcommand", [("codex", "app"), ("claude", "update")]) + def test_native_subcommand_suppresses_inherited_smart_routing( + self, monkeypatch, tool, subcommand + ): + monkeypatch.setenv("ENABLE_SMART_ROUTING_V2", "1") + observed = [] + + with patch( + "ucode.cli._launch_tool", + side_effect=lambda *_args, **_kwargs: observed.append( + os.environ.get("ENABLE_SMART_ROUTING_V2") + ), + ): + result = runner.invoke(app, [tool, subcommand]) + + assert result.exit_code == 0, result.output + assert observed == [None] + assert os.environ["ENABLE_SMART_ROUTING_V2"] == "1" + def test_claude_enable_smart_routing_forwards_positional_prompt_to_v2(self): captured = [] diff --git a/tests/test_codex_routing.py b/tests/test_codex_routing.py index 683d016b..ffa194cb 100644 --- a/tests/test_codex_routing.py +++ b/tests/test_codex_routing.py @@ -82,6 +82,83 @@ def fake_urlopen(request, timeout): } +def test_routes_select_request_is_logged(monkeypatch, tmp_path): + monkeypatch.delenv("SMART_ROUTER_NAME", raising=False) + monkeypatch.setattr(codex_routing.config_io, "APP_DIR", tmp_path) + logged = [] + + monkeypatch.setattr( + codex_routing.urllib.request, + "urlopen", + lambda request, timeout: _Response( + {"route_selection": [{"route_option": {"model": "gpt-5-6-sol"}}]} + ), + ) + + decision, error = codex_routing.request_routing_decision( + WS, + "secret-token", + "Fix the parser", + ["system.ai.gpt-5-6-sol"], + log=logged.append, + ) + + assert error is None + assert decision is not None + assert len(logged) == 1 + assert logged[0].startswith(f"[ROUTE] request POST {WS}/ai-gateway/routing/v1/routes:select: ") + logged_body = json.loads(logged[0].split(": ", 1)[1]) + record = json.loads((tmp_path / codex_routing.REQUESTS_LOG_FILENAME).read_text()) + assert record["method"] == "POST" + assert record["url"] == f"{WS}/ai-gateway/routing/v1/routes:select" + assert logged_body == record["body"] == { + "route_options": [{"model": "gpt-5-6-sol", "harness": "codex"}], + "task": {"prompt": "Fix the parser"}, + "route_selector": {"router_name": codex_routing.routing.ROUTER_NAME}, + } + assert "secret-token" not in logged[0] + assert "secret-token" not in (tmp_path / codex_routing.REQUESTS_LOG_FILENAME).read_text() + + +def test_routes_select_forwards_gateway_headers_without_auth_override(monkeypatch, tmp_path): + monkeypatch.setattr(codex_routing.config_io, "APP_DIR", tmp_path) + captured = {} + + def fake_urlopen(request, timeout): + captured["headers"] = {key.casefold(): value for key, value in request.headers.items()} + 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, + "real-token", + "Fix the parser", + ["system.ai.gpt-5-6-sol"], + extra_headers={ + "x-databricks-traffic-id": "testenv://liteswap/arnav-r315-task-v3", + "Authorization": "Bearer wrong-token", + "X-Ignored": "no", + }, + ) + + assert error is None + assert decision is not None + assert captured["headers"]["authorization"] == "Bearer real-token" + assert ( + captured["headers"]["x-databricks-traffic-id"] + == "testenv://liteswap/arnav-r315-task-v3" + ) + assert "x-ignored" not in captured["headers"] + record = json.loads((tmp_path / codex_routing.REQUESTS_LOG_FILENAME).read_text()) + assert record["headers"]["Authorization"] == "[REDACTED]" + assert ( + record["headers"]["x-databricks-traffic-id"] + == "testenv://liteswap/arnav-r315-task-v3" + ) + assert "wrong-token" not in (tmp_path / codex_routing.REQUESTS_LOG_FILENAME).read_text() + + def test_router_name_can_be_overridden_with_environment_variable(monkeypatch): captured = {} monkeypatch.setenv("SMART_ROUTER_NAME", " custom_router ") diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index fbc61e7d..430ca1c8 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -167,6 +167,13 @@ def start_interposer(*args, **kwargs): monkeypatch.setattr(codex_interposer, "start_interposer_thread", start_interposer) + def render_overlay(*args, **kwargs): + overlay = codex.render_overlay(*args, **kwargs) + overlay["model_providers"]["ucode-databricks"]["http_headers"][ + "x-databricks-traffic-id" + ] = "testenv://liteswap/arnav-r315-task-v3" + return overlay + with pytest.raises(SystemExit) as exc: v2.launch_codex( { @@ -178,7 +185,7 @@ def start_interposer(*args, **kwargs): ["--search"], binary="codex", start_model="gpt-start", - render_overlay=codex.render_overlay, + render_overlay=render_overlay, ) assert exc.value.code == 7 @@ -221,6 +228,9 @@ def start_interposer(*args, **kwargs): "system.ai.glm-5-2", ] assert interposer_args["kwargs"]["workspace"] == WS + assert interposer_args["kwargs"]["route_headers"]["x-databricks-traffic-id"] == ( + "testenv://liteswap/arnav-r315-task-v3" + ) assert token_calls == [(WS, "myprof")] assert interposer_args["kwargs"]["token_provider"]() == "token-2" assert token_calls == [(WS, "myprof"), (WS, "myprof")] From 8b8f1708804122c171fa0e785fd62222e099cf77 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Fri, 11 Sep 2026 02:23:06 +0000 Subject: [PATCH 2/5] update --- src/ucode/cli.py | 2 +- src/ucode/smart_routing/claude_routing.py | 22 +--------- src/ucode/smart_routing/codex_routing.py | 21 +--------- src/ucode/smart_routing/routing.py | 25 ----------- src/ucode/smart_routing/v2.py | 22 +--------- tests/test_claude_routing.py | 51 +---------------------- tests/test_claude_smart_routing_v2.py | 41 ------------------ tests/test_codex_routing.py | 48 +-------------------- 8 files changed, 8 insertions(+), 224 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index aeba7c5d..eb82d14d 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1744,7 +1744,7 @@ def _smart_routing_v2_flag(enabled: bool) -> Iterator[None]: @contextmanager -def _disable_smart_routing_for_subcommand(tool: str, ctx: typer.Context) -> Iterator[None]: +def _disable_smart_routing_for_subcommand(tool: str, ctx: Any) -> Iterator[None]: """Keep native agent subcommands out of every smart-routing path. The environment flag is also consulted during bootstrap/version checks, diff --git a/src/ucode/smart_routing/claude_routing.py b/src/ucode/smart_routing/claude_routing.py index 58ac837d..ced234af 100644 --- a/src/ucode/smart_routing/claude_routing.py +++ b/src/ucode/smart_routing/claude_routing.py @@ -15,11 +15,9 @@ # ``claude_routing.urllib.request`` — the call lives in ``routing``, but Python # modules are singletons so patching this name patches the one call site. import urllib.request # noqa: F401 -from collections.abc import Callable, Mapping -from pathlib import Path +from collections.abc import Mapping from typing import Any -from ucode import config_io from ucode.config_io import APP_DIR from ucode.smart_routing import routing from ucode.smart_routing.routing import RoutingDecision @@ -34,16 +32,9 @@ CANARY_PATH = APP_DIR / "claude-smart-routing-canary.json" AUDIT_PATH = APP_DIR / "claude-smart-routing-audit.jsonl" DECISIONS_PATH = APP_DIR / "claude-smart-routing-decisions.jsonl" -REQUESTS_LOG_FILENAME = "claude-smart-routing-requests.jsonl" - _normalize_model = routing.normalize_model -def request_log_path() -> Path: - """Return the Claude Code smart-routing request log path.""" - return config_io.APP_DIR / REQUESTS_LOG_FILENAME - - # Claude Code CLI options that consume a following value (from `claude --help`); # their values must not be mistaken for the seed prompt. Options whose value is # optional (`-c`/`-d`/`-r`/`-w`) are intentionally omitted — treating them as @@ -90,7 +81,6 @@ def request_routing_decision( available_models: list[str], *, timeout: float = REQUEST_TIMEOUT_S, - log: Callable[[str], None] | None = None, extra_headers: Mapping[str, str] | None = None, ) -> tuple[RoutingDecision | None, str | None]: """Ask the router for a servable Claude model. @@ -105,14 +95,6 @@ def request_routing_decision( route_options = [(arm, "claude") for arm in CLAUDE_ROUTE_ARMS] router_name = routing.configured_router_name() - headers = routing.route_request_headers(token, extra_headers) - routing.log_route_request( - workspace, - routing.route_request_body(task, route_options, router_name=router_name), - headers=headers, - log=log, - request_log_path=request_log_path(), - ) select_kwargs: dict[str, Any] = {"router_name": router_name, "timeout": timeout} if extra_headers: select_kwargs["extra_headers"] = extra_headers @@ -182,7 +164,7 @@ def record_subagent_start(payload: dict[str, Any]) -> dict[str, Any]: def clear_routing_artifacts() -> None: """Remove ucode-owned routing canary and audit files.""" - routing.clear_artifacts((CANARY_PATH, AUDIT_PATH, DECISIONS_PATH, request_log_path())) + routing.clear_artifacts((CANARY_PATH, AUDIT_PATH, DECISIONS_PATH)) def _claude_model_id(model: str) -> str: diff --git a/src/ucode/smart_routing/codex_routing.py b/src/ucode/smart_routing/codex_routing.py index fb0d2853..d15287bc 100644 --- a/src/ucode/smart_routing/codex_routing.py +++ b/src/ucode/smart_routing/codex_routing.py @@ -13,11 +13,9 @@ # ``codex_routing.urllib.request`` — the actual call lives in ``routing``, but # Python modules are singletons so patching this name patches the one call site. import urllib.request # noqa: F401 -from collections.abc import Callable, Mapping -from pathlib import Path +from collections.abc import Mapping from typing import Any -from ucode import config_io from ucode.config_io import APP_DIR from ucode.smart_routing import routing from ucode.smart_routing.routing import RoutingDecision @@ -29,18 +27,12 @@ CANARY_PATH = APP_DIR / "codex-smart-routing-canary.json" AUDIT_PATH = APP_DIR / "codex-smart-routing-audit.jsonl" DECISIONS_PATH = APP_DIR / "codex-smart-routing-decisions.jsonl" -REQUESTS_LOG_FILENAME = "codex-smart-routing-requests.jsonl" _GPT_RE = re.compile(r"gpt-(\d+)(?:[.-](\d+))?(?:[.-](\d+))?(-.+|[a-z].*)?") _normalize_model = routing.normalize_model -def request_log_path() -> Path: - """Return the Codex smart-routing request log path.""" - return config_io.APP_DIR / REQUESTS_LOG_FILENAME - - def request_routing_decision( workspace: str, token: str, @@ -48,7 +40,6 @@ def request_routing_decision( available_models: list[str], *, timeout: float = REQUEST_TIMEOUT_S, - log: Callable[[str], None] | None = None, extra_headers: Mapping[str, str] | None = None, ) -> tuple[RoutingDecision | None, str | None]: """Ask the router for a servable Codex model.""" @@ -57,14 +48,6 @@ def request_routing_decision( if not route_options: return None, "no cached model services are available" router_name = routing.configured_router_name() - headers = routing.route_request_headers(token, extra_headers) - routing.log_route_request( - workspace, - routing.route_request_body(task, route_options, router_name=router_name), - headers=headers, - log=log, - request_log_path=request_log_path(), - ) select_kwargs: dict[str, Any] = {"router_name": router_name, "timeout": timeout} if extra_headers: select_kwargs["extra_headers"] = extra_headers @@ -138,7 +121,7 @@ def record_subagent_start(payload: dict[str, Any]) -> dict[str, Any]: def clear_routing_artifacts() -> None: """Remove ucode-owned routing canary and audit files.""" - routing.clear_artifacts((CANARY_PATH, AUDIT_PATH, DECISIONS_PATH, request_log_path())) + routing.clear_artifacts((CANARY_PATH, AUDIT_PATH, DECISIONS_PATH)) def _parse_gpt(model: str) -> tuple[int, int, int, str] | None: diff --git a/src/ucode/smart_routing/routing.py b/src/ucode/smart_routing/routing.py index d1c149eb..ef164352 100644 --- a/src/ucode/smart_routing/routing.py +++ b/src/ucode/smart_routing/routing.py @@ -170,31 +170,6 @@ def route_request_headers( return headers -def log_route_request( - workspace: str, - body: dict[str, Any], - *, - headers: Mapping[str, str] | None = None, - log: Callable[[str], None] | None = None, - request_log_path: Path | None = None, -) -> None: - """Record an outgoing route-selection request without credentials.""" - url = workspace.rstrip("/") + ROUTING_PATH - if log is not None: - log(f"[ROUTE] request POST {url}: {json.dumps(body, separators=(',', ':'))}") - if request_log_path is not None: - _append_jsonl( - request_log_path, - { - "at": time.time(), - "method": "POST", - "url": url, - "headers": _redacted_headers(headers or {}), - "body": body, - }, - ) - - def select_route( workspace: str, token: str, diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 41caa3a7..935b0796 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -241,7 +241,6 @@ def _request_claude_routing_decision( prompt: str, model_ids: list[str], *, - log: Callable[[str], None] | None = None, extra_headers: dict[str, str] | None = None, ) -> tuple[routing.RoutingDecision | None, str | None]: available: dict[str, str] = {} @@ -251,14 +250,6 @@ def _request_claude_routing_decision( return None, "Anthropic models endpoint returned no Claude models" route_options = [(model, "claude") for model in available] router_name = routing.configured_router_name() - headers = routing.route_request_headers(token, extra_headers) - routing.log_route_request( - workspace, - routing.route_request_body(prompt, route_options, router_name=router_name), - headers=headers, - log=log, - request_log_path=claude_routing.request_log_path(), - ) select_kwargs: dict[str, object] = { "router_name": router_name, "timeout": CLAUDE_ROUTE_SELECTION_TIMEOUT_S, @@ -281,7 +272,6 @@ def _route_claude_prompt( prompt: str, model_ids: list[str] | None = None, *, - log: Callable[[str], None] | None = None, extra_headers: dict[str, str] | None = None, ) -> routing.RoutingDecision: workspace = state.get("workspace") @@ -299,7 +289,6 @@ def _route_claude_prompt( token, prompt, model_ids, - log=log, extra_headers=extra_headers, ) if decision is None: @@ -473,21 +462,12 @@ def launch_claude( model_setting = _ClaudeModelSettingGuard(user_settings_path) - def log_route_request(message: str) -> None: - try: - CLAUDE_PTY_LOG.parent.mkdir(parents=True, exist_ok=True) - with open(CLAUDE_PTY_LOG, "a", encoding="utf-8") as handle: - handle.write(f"{time.strftime('%H:%M:%S')} {message}\n") - except OSError: - pass - def route_prompt(prompt: str) -> claude_pty.FirstPromptRoute: decision = _route_claude_prompt( state, token, prompt, model_ids, - log=log_route_request, extra_headers=route_headers, ) return claude_pty.FirstPromptRoute( @@ -498,7 +478,7 @@ def route_prompt(prompt: str) -> claude_pty.FirstPromptRoute: print_note( "Smart routing v2: the first submitted prompt will select Claude Code's " - f"model; log: {CLAUDE_PTY_LOG}." + "model." ) try: returncode = claude_pty.run_claude_pty( diff --git a/tests/test_claude_routing.py b/tests/test_claude_routing.py index 6345daaa..54483afd 100644 --- a/tests/test_claude_routing.py +++ b/tests/test_claude_routing.py @@ -70,49 +70,7 @@ def fake_urlopen(request, timeout): } -def test_routes_select_request_is_logged(monkeypatch, tmp_path): - monkeypatch.delenv("SMART_ROUTER_NAME", raising=False) - monkeypatch.setattr(claude_routing.config_io, "APP_DIR", tmp_path) - logged = [] - - monkeypatch.setattr( - claude_routing.urllib.request, - "urlopen", - lambda request, timeout: _Response( - {"route_selection": [{"route_option": {"model": "claude-sonnet-5"}}]} - ), - ) - - decision, error = claude_routing.request_routing_decision( - WS, - "secret-token", - "Map the codebase", - ["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5"], - log=logged.append, - ) - - assert error is None - assert decision is not None - assert len(logged) == 1 - assert logged[0].startswith(f"[ROUTE] request POST {WS}/ai-gateway/routing/v1/routes:select: ") - logged_body = json.loads(logged[0].split(": ", 1)[1]) - record = json.loads((tmp_path / claude_routing.REQUESTS_LOG_FILENAME).read_text()) - assert record["method"] == "POST" - assert record["url"] == f"{WS}/ai-gateway/routing/v1/routes:select" - assert logged_body == record["body"] == { - "route_options": [ - {"model": "claude-opus-4-8", "harness": "claude"}, - {"model": "claude-sonnet-5", "harness": "claude"}, - ], - "task": {"prompt": "Map the codebase"}, - "route_selector": {"router_name": claude_routing.ROUTER_NAME}, - } - assert "secret-token" not in logged[0] - assert "secret-token" not in (tmp_path / claude_routing.REQUESTS_LOG_FILENAME).read_text() - - -def test_routes_select_forwards_gateway_headers_without_auth_override(monkeypatch, tmp_path): - monkeypatch.setattr(claude_routing.config_io, "APP_DIR", tmp_path) +def test_routes_select_forwards_gateway_headers_without_auth_override(monkeypatch): captured = {} def fake_urlopen(request, timeout): @@ -152,13 +110,6 @@ def fake_urlopen(request, timeout): == '{"source":"isaac-cli"}' ) assert "x-ignored" not in captured["headers"] - record = json.loads((tmp_path / claude_routing.REQUESTS_LOG_FILENAME).read_text()) - assert record["headers"]["Authorization"] == "[REDACTED]" - assert ( - record["headers"]["x-databricks-traffic-id"] - == "testenv://liteswap/arnav-r315-task-v3" - ) - assert "wrong-token" not in (tmp_path / claude_routing.REQUESTS_LOG_FILENAME).read_text() def test_missing_arm_short_circuits_without_calling_router(monkeypatch): diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index 5f2075d7..7ba21e22 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -421,47 +421,6 @@ def fake_select(workspace, token, task, route_options, resolve, **kwargs): "router_name": "task_v2", } - def test_logs_v2_routes_select_request(self, tmp_path, monkeypatch): - monkeypatch.setattr(v2.claude_routing.config_io, "APP_DIR", tmp_path) - monkeypatch.setenv("SMART_ROUTER_NAME", "task_v2") - logged = [] - - def fake_select(workspace, token, task, route_options, resolve, **kwargs): - return ( - routing.RoutingDecision( - model=resolve("claude-sonnet-5"), - raw_model="claude-sonnet-5", - ), - 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", "system.ai.claude-sonnet-5"], - log=logged.append, - ) - - assert error is None - assert decision is not None - assert len(logged) == 1 - assert logged[0].startswith( - "[ROUTE] request POST https://example.com/ai-gateway/routing/v1/routes:select: " - ) - record = json.loads((tmp_path / v2.claude_routing.REQUESTS_LOG_FILENAME).read_text()) - assert json.loads(logged[0].split(": ", 1)[1]) == record["body"] == { - "route_options": [ - {"model": "claude-opus-4-8", "harness": "claude"}, - {"model": "claude-sonnet-5", "harness": "claude"}, - ], - "task": {"prompt": "inspect the parser"}, - "route_selector": {"router_name": "task_v2"}, - } - assert "secret-token" not in logged[0] - assert "secret-token" not in (tmp_path / v2.claude_routing.REQUESTS_LOG_FILENAME).read_text() - def test_routes_agent_prompt_with_initialized_model_menu(self, tmp_path, monkeypatch): captured = {} decisions_path = tmp_path / "decisions.jsonl" diff --git a/tests/test_codex_routing.py b/tests/test_codex_routing.py index ffa194cb..17c95fee 100644 --- a/tests/test_codex_routing.py +++ b/tests/test_codex_routing.py @@ -82,46 +82,7 @@ def fake_urlopen(request, timeout): } -def test_routes_select_request_is_logged(monkeypatch, tmp_path): - monkeypatch.delenv("SMART_ROUTER_NAME", raising=False) - monkeypatch.setattr(codex_routing.config_io, "APP_DIR", tmp_path) - logged = [] - - monkeypatch.setattr( - codex_routing.urllib.request, - "urlopen", - lambda request, timeout: _Response( - {"route_selection": [{"route_option": {"model": "gpt-5-6-sol"}}]} - ), - ) - - decision, error = codex_routing.request_routing_decision( - WS, - "secret-token", - "Fix the parser", - ["system.ai.gpt-5-6-sol"], - log=logged.append, - ) - - assert error is None - assert decision is not None - assert len(logged) == 1 - assert logged[0].startswith(f"[ROUTE] request POST {WS}/ai-gateway/routing/v1/routes:select: ") - logged_body = json.loads(logged[0].split(": ", 1)[1]) - record = json.loads((tmp_path / codex_routing.REQUESTS_LOG_FILENAME).read_text()) - assert record["method"] == "POST" - assert record["url"] == f"{WS}/ai-gateway/routing/v1/routes:select" - assert logged_body == record["body"] == { - "route_options": [{"model": "gpt-5-6-sol", "harness": "codex"}], - "task": {"prompt": "Fix the parser"}, - "route_selector": {"router_name": codex_routing.routing.ROUTER_NAME}, - } - assert "secret-token" not in logged[0] - assert "secret-token" not in (tmp_path / codex_routing.REQUESTS_LOG_FILENAME).read_text() - - -def test_routes_select_forwards_gateway_headers_without_auth_override(monkeypatch, tmp_path): - monkeypatch.setattr(codex_routing.config_io, "APP_DIR", tmp_path) +def test_routes_select_forwards_gateway_headers_without_auth_override(monkeypatch): captured = {} def fake_urlopen(request, timeout): @@ -150,13 +111,6 @@ def fake_urlopen(request, timeout): == "testenv://liteswap/arnav-r315-task-v3" ) assert "x-ignored" not in captured["headers"] - record = json.loads((tmp_path / codex_routing.REQUESTS_LOG_FILENAME).read_text()) - assert record["headers"]["Authorization"] == "[REDACTED]" - assert ( - record["headers"]["x-databricks-traffic-id"] - == "testenv://liteswap/arnav-r315-task-v3" - ) - assert "wrong-token" not in (tmp_path / codex_routing.REQUESTS_LOG_FILENAME).read_text() def test_router_name_can_be_overridden_with_environment_variable(monkeypatch): From 0da85c7743aac1f9e8cc50fbdc11794d7cda14f6 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Fri, 11 Sep 2026 02:29:17 +0000 Subject: [PATCH 3/5] update --- src/ucode/smart_routing/claude_routing.py | 11 ---- src/ucode/smart_routing/codex_interposer.py | 8 +-- src/ucode/smart_routing/codex_routing.py | 6 -- src/ucode/smart_routing/routing.py | 67 ++------------------- src/ucode/smart_routing/v2.py | 58 ------------------ tests/test_claude_routing.py | 42 ------------- tests/test_claude_smart_routing_v2.py | 15 ----- tests/test_codex_routing.py | 31 ---------- tests/test_codex_smart_routing_v2.py | 19 ------ 9 files changed, 6 insertions(+), 251 deletions(-) diff --git a/src/ucode/smart_routing/claude_routing.py b/src/ucode/smart_routing/claude_routing.py index ced234af..62687bd0 100644 --- a/src/ucode/smart_routing/claude_routing.py +++ b/src/ucode/smart_routing/claude_routing.py @@ -9,13 +9,10 @@ from __future__ import annotations -import os - # Re-exported so tests can patch the shared ``urlopen`` seam via # ``claude_routing.urllib.request`` — the call lives in ``routing``, but Python # modules are singletons so patching this name patches the one call site. import urllib.request # noqa: F401 -from collections.abc import Mapping from typing import Any from ucode.config_io import APP_DIR @@ -81,7 +78,6 @@ def request_routing_decision( available_models: list[str], *, timeout: float = REQUEST_TIMEOUT_S, - extra_headers: Mapping[str, str] | None = None, ) -> tuple[RoutingDecision | None, str | None]: """Ask the router for a servable Claude model. @@ -96,8 +92,6 @@ def request_routing_decision( route_options = [(arm, "claude") for arm in CLAUDE_ROUTE_ARMS] router_name = routing.configured_router_name() select_kwargs: dict[str, Any] = {"router_name": router_name, "timeout": timeout} - if extra_headers: - select_kwargs["extra_headers"] = extra_headers return routing.select_route( workspace, token, @@ -116,7 +110,6 @@ def route_pre_tool_use( available_models: list[str], timeout: float = REQUEST_TIMEOUT_S, audit_decision: bool = False, - extra_headers: Mapping[str, str] | None = None, ) -> dict[str, Any] | None: """Route one Claude Code ``Agent`` (subagent-spawn) call, rewriting its model.""" record = None @@ -134,10 +127,6 @@ def record(payload, task, decision, requested): task, available_models, timeout=timeout, - extra_headers=extra_headers - or routing.route_forward_headers_from_lines( - os.environ.get("ANTHROPIC_CUSTOM_HEADERS") - ), ), default_task_label="Claude Code subagent task", model_id_mapper=_claude_model_id, diff --git a/src/ucode/smart_routing/codex_interposer.py b/src/ucode/smart_routing/codex_interposer.py index e71f1ddf..5d39b5d1 100644 --- a/src/ucode/smart_routing/codex_interposer.py +++ b/src/ucode/smart_routing/codex_interposer.py @@ -6,7 +6,7 @@ import threading import time import uuid -from collections.abc import Callable, Mapping +from collections.abc import Callable from dataclasses import dataclass, replace from pathlib import Path @@ -210,7 +210,6 @@ async def _handle_tui( available_models: list[str] | None = None, workspace: str | None = None, token_provider: TokenProvider | None = None, - route_headers: Mapping[str, str] | None = None, switch_message_fn: SwitchMessageFn | None = None, ) -> None: path = getattr(getattr(tui, "request", None), "path", "/") or "/" @@ -230,7 +229,6 @@ def route_decision(prompt: str): prompt, list(available_models or []), log=log, - extra_headers=route_headers, ) sess = _Session( @@ -301,7 +299,6 @@ async def _serve( available_models: list[str] | None = None, workspace: str | None = None, token_provider: TokenProvider | None = None, - route_headers: Mapping[str, str] | None = None, switch_message_fn: SwitchMessageFn | None = None, ): async def handler(tui): @@ -315,7 +312,6 @@ async def handler(tui): available_models, workspace, token_provider, - route_headers, switch_message_fn, ) except Exception as exc: # noqa: BLE001 @@ -335,7 +331,6 @@ def start_interposer_thread( available_models: list[str] | None = None, workspace: str | None = None, token_provider: TokenProvider | None = None, - route_headers: Mapping[str, str] | None = None, switch_message_fn: SwitchMessageFn | None = None, switch_message: str | None = None, log_path: Path | None = None, @@ -369,7 +364,6 @@ def run() -> None: available_models, workspace, token_provider, - route_headers, switch_message_fn, ) ) diff --git a/src/ucode/smart_routing/codex_routing.py b/src/ucode/smart_routing/codex_routing.py index d15287bc..5a68aa73 100644 --- a/src/ucode/smart_routing/codex_routing.py +++ b/src/ucode/smart_routing/codex_routing.py @@ -13,7 +13,6 @@ # ``codex_routing.urllib.request`` — the actual call lives in ``routing``, but # Python modules are singletons so patching this name patches the one call site. import urllib.request # noqa: F401 -from collections.abc import Mapping from typing import Any from ucode.config_io import APP_DIR @@ -40,7 +39,6 @@ def request_routing_decision( available_models: list[str], *, timeout: float = REQUEST_TIMEOUT_S, - extra_headers: Mapping[str, str] | None = None, ) -> tuple[RoutingDecision | None, str | None]: """Ask the router for a servable Codex model.""" available = {_normalize_model(model): model for model in available_models} @@ -49,8 +47,6 @@ def request_routing_decision( return None, "no cached model services are available" router_name = routing.configured_router_name() select_kwargs: dict[str, Any] = {"router_name": router_name, "timeout": timeout} - if extra_headers: - select_kwargs["extra_headers"] = extra_headers return routing.select_route( workspace, token, @@ -75,7 +71,6 @@ def route_pre_tool_use( available_models: list[str], timeout: float = REQUEST_TIMEOUT_S, audit_decision: bool = False, - extra_headers: Mapping[str, str] | None = None, ) -> dict[str, Any] | None: """Route one Codex ``spawn_agent`` call and rewrite its model.""" record = None @@ -93,7 +88,6 @@ def record(payload, task, decision, requested): task, available_models, timeout=timeout, - extra_headers=extra_headers, ), default_task_label="Codex subagent task", model_id_mapper=codex_model_id, diff --git a/src/ucode/smart_routing/routing.py b/src/ucode/smart_routing/routing.py index ef164352..94683b5c 100644 --- a/src/ucode/smart_routing/routing.py +++ b/src/ucode/smart_routing/routing.py @@ -16,7 +16,7 @@ import urllib.error import urllib.request import uuid -from collections.abc import Callable, Iterable, Mapping +from collections.abc import Callable, Iterable from dataclasses import dataclass from pathlib import Path from typing import Any @@ -25,22 +25,6 @@ ROUTER_NAME_ENV_VAR = "SMART_ROUTER_NAME" ROUTING_PATH = "/ai-gateway/routing/v1/routes:select" REQUEST_TIMEOUT_S = 30.0 -ROUTE_FORWARD_HEADER_NAMES = frozenset( - { - "databricks-ai-gateway-request-tags", - "user-agent", - "x-databricks-traffic-id", - "x-databricks-use-coding-agent-mode", - } -) -SENSITIVE_HEADER_NAMES = frozenset( - { - "authorization", - "cookie", - "proxy-authorization", - "x-databricks-ai-gateway-token", - } -) SUBAGENT_ROUTING_DISCLAIMER = ( "Spawned subagents are routed independently based on their own complexity." ) @@ -134,42 +118,6 @@ def route_request_body( } -def route_forward_headers_from_lines(headers: object) -> dict[str, str]: - """Parse newline-delimited gateway headers to forward to ``routes:select``.""" - if not isinstance(headers, str): - return {} - forwarded: dict[str, str] = {} - for line in headers.splitlines(): - name, separator, value = line.partition(":") - normalized = name.strip().casefold() - if not separator or normalized not in ROUTE_FORWARD_HEADER_NAMES: - continue - clean_name = name.strip() - clean_value = value.strip() - if clean_name and clean_value: - forwarded[clean_name] = clean_value - return forwarded - - -def route_request_headers( - token: str, - extra_headers: Mapping[str, str] | None = None, -) -> dict[str, str]: - """Return HTTP headers for ``routes:select`` without allowing auth overrides.""" - headers: dict[str, str] = {} - for name, value in (extra_headers or {}).items(): - normalized = name.strip().casefold() - if normalized not in ROUTE_FORWARD_HEADER_NAMES: - continue - clean_name = name.strip() - clean_value = value.strip() - if clean_name and clean_value: - headers[clean_name] = clean_value - headers["Authorization"] = f"Bearer {token}" - headers["Content-Type"] = "application/json" - return headers - - def select_route( workspace: str, token: str, @@ -179,7 +127,6 @@ def select_route( *, router_name: str, timeout: float = REQUEST_TIMEOUT_S, - extra_headers: Mapping[str, str] | None = None, ) -> tuple[RoutingDecision | None, str | None]: """POST one ``routes:select`` request and resolve the router's pick. @@ -193,7 +140,10 @@ def select_route( request = urllib.request.Request( workspace.rstrip("/") + ROUTING_PATH, data=json.dumps(body).encode("utf-8"), - headers=route_request_headers(token, extra_headers), + headers={ + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + }, method="POST", ) try: @@ -402,13 +352,6 @@ def _selected_model(payload: Any) -> str | None: return model if isinstance(model, str) and model else None -def _redacted_headers(headers: Mapping[str, str]) -> dict[str, str]: - redacted: dict[str, str] = {} - for name, value in headers.items(): - redacted[name] = "[REDACTED]" if name.strip().casefold() in SENSITIVE_HEADER_NAMES else value - return redacted - - def _pending_decision( decisions_path: Path, audit_path: Path, session_id: Any, actual_model: Any ) -> dict[str, Any] | None: diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 935b0796..3db3f411 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -38,7 +38,6 @@ LEGACY_STATE_KEY = "smart_routing_enabled" CODEX_INTERPOSER_LOG = APP_DIR / "codex-v2-interposer.log" -CODEX_CUSTOM_HEADERS_ENV_VAR = "CODEX_CUSTOM_HEADERS" CLAUDE_TARGET_MODEL = "system.ai.claude-sonnet-4-6[1m]" # TODO(lilly): replace with smart router. CLAUDE_PTY_LOG = APP_DIR / "claude-v2-pty.log" @@ -240,8 +239,6 @@ def _request_claude_routing_decision( token: str, prompt: str, model_ids: list[str], - *, - extra_headers: dict[str, str] | None = None, ) -> tuple[routing.RoutingDecision | None, str | None]: available: dict[str, str] = {} for model in _canonical_claude_models(model_ids): @@ -254,8 +251,6 @@ def _request_claude_routing_decision( "router_name": router_name, "timeout": CLAUDE_ROUTE_SELECTION_TIMEOUT_S, } - if extra_headers: - select_kwargs["extra_headers"] = extra_headers return routing.select_route( workspace, token, @@ -271,8 +266,6 @@ def _route_claude_prompt( token: str, prompt: str, model_ids: list[str] | None = None, - *, - extra_headers: dict[str, str] | None = None, ) -> routing.RoutingDecision: workspace = state.get("workspace") if not isinstance(workspace, str): @@ -289,7 +282,6 @@ def _route_claude_prompt( token, prompt, model_ids, - extra_headers=extra_headers, ) if decision is None: raise RuntimeError(error or "router returned no Claude model selection") @@ -303,7 +295,6 @@ def route_claude_pre_tool_use( token: str, available_models: list[str], audit_decision: bool = False, - extra_headers: dict[str, str] | None = None, ) -> dict | None: """Route a Claude Agent call through a transient exact-model agent definition.""" route = routing.resolve_spawn_route( @@ -314,10 +305,6 @@ def route_claude_pre_tool_use( token, task, available_models, - extra_headers=extra_headers - or routing.route_forward_headers_from_lines( - os.environ.get("ANTHROPIC_CUSTOM_HEADERS") - ), ), default_task_label="Claude Code subagent task", model_id_mapper=lambda model: model, @@ -441,10 +428,6 @@ def launch_claude( raise RuntimeError("Claude settings 'env' must be an object for smart routing.") env.pop("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", None) env[FIRST_PROMPT_SOCKET_ENV] = str(socket_path) - route_headers = { - **routing.route_forward_headers_from_lines(os.environ.get("ANTHROPIC_CUSTOM_HEADERS")), - **routing.route_forward_headers_from_lines(env.get("ANTHROPIC_CUSTOM_HEADERS")), - } model_overrides = settings.setdefault("modelOverrides", {}) if not isinstance(model_overrides, dict): raise RuntimeError("Claude settings 'modelOverrides' must be an object for smart routing.") @@ -468,7 +451,6 @@ def route_prompt(prompt: str) -> claude_pty.FirstPromptRoute: token, prompt, model_ids, - extra_headers=route_headers, ) return claude_pty.FirstPromptRoute( model=model_name(_unwrapped_claude_model_id(decision.model)), @@ -521,43 +503,6 @@ def _v2_pre_tool_use_hooks(state: dict, available_models: list[str]) -> list[dic ) -def _codex_route_headers(overlay: dict) -> dict[str, str]: - providers = overlay.get("model_providers") - if not isinstance(providers, dict): - return {} - provider = providers.get("ucode-databricks") - if not isinstance(provider, dict): - return {} - headers = provider.get("http_headers") - if not isinstance(headers, dict): - return {} - return { - str(name): str(value) - for name, value in headers.items() - if isinstance(name, str) and isinstance(value, str) - } - - -def _apply_codex_custom_headers(overlay: dict) -> dict[str, str]: - """Add shell-provided gateway headers to the Codex provider overlay.""" - custom_headers = routing.route_forward_headers_from_lines( - os.environ.get(CODEX_CUSTOM_HEADERS_ENV_VAR) - ) - if not custom_headers: - return overlay - providers = overlay.get("model_providers") - if not isinstance(providers, dict): - return overlay - provider = providers.get("ucode-databricks") - if not isinstance(provider, dict): - return overlay - headers = provider.setdefault("http_headers", {}) - if not isinstance(headers, dict): - return overlay - headers.update(custom_headers) - return overlay - - def launch_codex( state: dict, tool_args: list[str], @@ -590,8 +535,6 @@ def launch_codex( state.get("profile"), use_pat=bool(state.get("use_pat")), ) - _apply_codex_custom_headers(overlay) - route_headers = _codex_route_headers(overlay) overlay["hooks"] = { "PreToolUse": _v2_pre_tool_use_hooks(state, available_models), } @@ -620,7 +563,6 @@ def launch_codex( available_models=available_models, workspace=workspace, token_provider=lambda: get_databricks_token(workspace, profile), - route_headers=route_headers, switch_message_fn=format_routing_notice, log_path=CODEX_INTERPOSER_LOG, ) diff --git a/tests/test_claude_routing.py b/tests/test_claude_routing.py index 54483afd..ac41726d 100644 --- a/tests/test_claude_routing.py +++ b/tests/test_claude_routing.py @@ -70,48 +70,6 @@ def fake_urlopen(request, timeout): } -def test_routes_select_forwards_gateway_headers_without_auth_override(monkeypatch): - captured = {} - - def fake_urlopen(request, timeout): - captured["headers"] = {key.casefold(): value for key, value in request.headers.items()} - return _Response({"route_selection": [{"route_option": {"model": "claude-sonnet-5"}}]}) - - monkeypatch.setattr(claude_routing.urllib.request, "urlopen", fake_urlopen) - - decision, error = claude_routing.request_routing_decision( - WS, - "real-token", - "Map the codebase", - ["system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5"], - extra_headers=claude_routing.routing.route_forward_headers_from_lines( - "\n".join( - [ - "x-databricks-traffic-id: testenv://liteswap/arnav-r315-task-v3", - "x-databricks-use-coding-agent-mode: true", - 'Databricks-Ai-Gateway-Request-Tags: {"source":"isaac-cli"}', - "Authorization: Bearer wrong-token", - "X-Ignored: no", - ] - ) - ), - ) - - assert error is None - assert decision is not None - assert captured["headers"]["authorization"] == "Bearer real-token" - assert ( - captured["headers"]["x-databricks-traffic-id"] - == "testenv://liteswap/arnav-r315-task-v3" - ) - assert captured["headers"]["x-databricks-use-coding-agent-mode"] == "true" - assert ( - captured["headers"]["databricks-ai-gateway-request-tags"] - == '{"source":"isaac-cli"}' - ) - assert "x-ignored" not in captured["headers"] - - def test_missing_arm_short_circuits_without_calling_router(monkeypatch): def fail(*args, **kwargs): raise AssertionError("router must not be called when an arm is missing") diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index 7ba21e22..8f37a80a 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -257,12 +257,6 @@ def fake_run(argv, **kwargs): display_model="Claude Sonnet 5", rationale="Selected for the parser task.", ) - assert routed_call["extra_headers"]["x-databricks-traffic-id"] == ( - "testenv://liteswap/arnav-r315-task-v3" - ) - assert routed_call["extra_headers"]["Databricks-Ai-Gateway-Request-Tags"] == ( - '{"source":"isaac-cli"}' - ) assert {definition["model"] for definition in captured["agents"].values()} == { "system.ai.claude-opus-4-8", "system.ai.claude-sonnet-5", @@ -425,18 +419,12 @@ def test_routes_agent_prompt_with_initialized_model_menu(self, tmp_path, monkeyp captured = {} decisions_path = tmp_path / "decisions.jsonl" monkeypatch.setattr(v2.claude_routing, "DECISIONS_PATH", decisions_path) - monkeypatch.setenv( - "ANTHROPIC_CUSTOM_HEADERS", - "x-databricks-traffic-id: testenv://liteswap/arnav-r315-task-v3", - ) - def fake_select(workspace, token, task, route_options, resolve, **kwargs): captured.update( workspace=workspace, token=token, task=task, route_options=list(route_options), - extra_headers=kwargs.get("extra_headers"), ) return ( routing.RoutingDecision( @@ -469,9 +457,6 @@ def fake_select(workspace, token, task, route_options, resolve, **kwargs): ("claude-opus-4-8", "claude"), ("claude-sonnet-5", "claude"), ], - "extra_headers": { - "x-databricks-traffic-id": "testenv://liteswap/arnav-r315-task-v3" - }, } updated_input = output["hookSpecificOutput"]["updatedInput"] assert "model" not in updated_input diff --git a/tests/test_codex_routing.py b/tests/test_codex_routing.py index 17c95fee..683d016b 100644 --- a/tests/test_codex_routing.py +++ b/tests/test_codex_routing.py @@ -82,37 +82,6 @@ def fake_urlopen(request, timeout): } -def test_routes_select_forwards_gateway_headers_without_auth_override(monkeypatch): - captured = {} - - def fake_urlopen(request, timeout): - captured["headers"] = {key.casefold(): value for key, value in request.headers.items()} - 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, - "real-token", - "Fix the parser", - ["system.ai.gpt-5-6-sol"], - extra_headers={ - "x-databricks-traffic-id": "testenv://liteswap/arnav-r315-task-v3", - "Authorization": "Bearer wrong-token", - "X-Ignored": "no", - }, - ) - - assert error is None - assert decision is not None - assert captured["headers"]["authorization"] == "Bearer real-token" - assert ( - captured["headers"]["x-databricks-traffic-id"] - == "testenv://liteswap/arnav-r315-task-v3" - ) - assert "x-ignored" not in captured["headers"] - - def test_router_name_can_be_overridden_with_environment_variable(monkeypatch): captured = {} monkeypatch.setenv("SMART_ROUTER_NAME", " custom_router ") diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index 430ca1c8..d3a99b02 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -228,9 +228,6 @@ def render_overlay(*args, **kwargs): "system.ai.glm-5-2", ] assert interposer_args["kwargs"]["workspace"] == WS - assert interposer_args["kwargs"]["route_headers"]["x-databricks-traffic-id"] == ( - "testenv://liteswap/arnav-r315-task-v3" - ) assert token_calls == [(WS, "myprof")] assert interposer_args["kwargs"]["token_provider"]() == "token-2" assert token_calls == [(WS, "myprof"), (WS, "myprof")] @@ -569,7 +566,6 @@ def test_rewrites_nested_collaboration_mode_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, *, router_name, timeout): captured.update( @@ -601,7 +597,6 @@ def select_route(workspace, token, task, route_options, resolve, *, router_name, "system.ai.gpt-5-6-luna", "system.ai.glm-5-2", ], - log=logged.append, ) assert reason is None @@ -619,17 +614,3 @@ def select_route(workspace, token, task, route_options, resolve, *, router_name, ("glm-5-2", "codex"), ], } - assert len(logged) == 1 - assert logged[0].startswith(f"[ROUTE] request POST {WS}/ai-gateway/routing/v1/routes:select: ") - request_payload = json.loads(logged[0].split(": ", 1)[1]) - assert request_payload == { - "route_options": [ - {"model": "kimi-k3-neo", "harness": "codex"}, - {"model": "gpt-5-6-sol", "harness": "codex"}, - {"model": "gpt-5-6-luna", "harness": "codex"}, - {"model": "glm-5-2", "harness": "codex"}, - ], - "task": {"prompt": "Fix the parser"}, - "route_selector": {"router_name": codex_routing.routing.ROUTER_NAME}, - } - assert "same-oauth-token" not in logged[0] From 634c784f4fa9a8b52e00c4306c63c9e3cc4f874c Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Fri, 11 Sep 2026 02:36:29 +0000 Subject: [PATCH 4/5] Remove unrelated smart routing changes --- src/ucode/smart_routing/claude_routing.py | 15 ++++------- src/ucode/smart_routing/codex_routing.py | 23 ++++++++++++----- src/ucode/smart_routing/routing.py | 20 ++++----------- src/ucode/smart_routing/v2.py | 29 +++++---------------- tests/test_claude_smart_routing_v2.py | 31 ++++++----------------- tests/test_codex_smart_routing_v2.py | 25 ++++++++++++------ 6 files changed, 57 insertions(+), 86 deletions(-) diff --git a/src/ucode/smart_routing/claude_routing.py b/src/ucode/smart_routing/claude_routing.py index 62687bd0..cbdd5cf2 100644 --- a/src/ucode/smart_routing/claude_routing.py +++ b/src/ucode/smart_routing/claude_routing.py @@ -29,6 +29,7 @@ CANARY_PATH = APP_DIR / "claude-smart-routing-canary.json" AUDIT_PATH = APP_DIR / "claude-smart-routing-audit.jsonl" DECISIONS_PATH = APP_DIR / "claude-smart-routing-decisions.jsonl" + _normalize_model = routing.normalize_model @@ -89,16 +90,14 @@ def request_routing_decision( if missing: return None, f"required Claude routing models are unavailable: {', '.join(missing)}" - route_options = [(arm, "claude") for arm in CLAUDE_ROUTE_ARMS] - router_name = routing.configured_router_name() - select_kwargs: dict[str, Any] = {"router_name": router_name, "timeout": timeout} return routing.select_route( workspace, token, task, - route_options, + [(arm, "claude") for arm in CLAUDE_ROUTE_ARMS], lambda raw_model: available.get(_normalize_model(raw_model)), - **select_kwargs, + router_name=routing.configured_router_name(), + timeout=timeout, ) @@ -122,11 +121,7 @@ def record(payload, task, decision, requested): payload, is_spawn_agent=is_spawn_agent_tool, decision_fn=lambda task: request_routing_decision( - workspace, - token, - task, - available_models, - timeout=timeout, + workspace, token, task, available_models, timeout=timeout ), default_task_label="Claude Code subagent task", model_id_mapper=_claude_model_id, diff --git a/src/ucode/smart_routing/codex_routing.py b/src/ucode/smart_routing/codex_routing.py index 5a68aa73..ea048ef9 100644 --- a/src/ucode/smart_routing/codex_routing.py +++ b/src/ucode/smart_routing/codex_routing.py @@ -7,12 +7,14 @@ from __future__ import annotations +import json import re # Re-exported so tests can patch the shared ``urlopen`` seam via # ``codex_routing.urllib.request`` — the actual call lives in ``routing``, but # Python modules are singletons so patching this name patches the one call site. import urllib.request # noqa: F401 +from collections.abc import Callable from typing import Any from ucode.config_io import APP_DIR @@ -39,6 +41,7 @@ def request_routing_decision( available_models: list[str], *, timeout: float = REQUEST_TIMEOUT_S, + log: Callable[[str], None] | None = None, ) -> tuple[RoutingDecision | None, str | None]: """Ask the router for a servable Codex model.""" available = {_normalize_model(model): model for model in available_models} @@ -46,14 +49,24 @@ def request_routing_decision( if not route_options: return None, "no cached model services are available" router_name = routing.configured_router_name() - select_kwargs: dict[str, Any] = {"router_name": router_name, "timeout": timeout} + 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}, + } + url = workspace.rstrip("/") + ROUTING_PATH + log(f"[ROUTE] request POST {url}: {json.dumps(payload, separators=(',', ':'))}") return routing.select_route( workspace, token, task, route_options, lambda raw_model: available.get(_normalize_model(raw_model)), - **select_kwargs, + router_name=router_name, + timeout=timeout, ) @@ -83,11 +96,7 @@ def record(payload, task, decision, requested): payload, is_spawn_agent=is_spawn_agent_tool, decision_fn=lambda task: request_routing_decision( - workspace, - token, - task, - available_models, - timeout=timeout, + workspace, token, task, available_models, timeout=timeout ), default_task_label="Codex subagent task", model_id_mapper=codex_model_id, diff --git a/src/ucode/smart_routing/routing.py b/src/ucode/smart_routing/routing.py index 94683b5c..5ad46fe3 100644 --- a/src/ucode/smart_routing/routing.py +++ b/src/ucode/smart_routing/routing.py @@ -104,20 +104,6 @@ def configured_router_name() -> str: return os.environ.get(ROUTER_NAME_ENV_VAR, "").strip() or ROUTER_NAME -def route_request_body( - task: str, - route_options: Iterable[tuple[str, str | None]], - *, - router_name: str, -) -> dict[str, Any]: - """Return the JSON body sent to ``routes:select``.""" - return { - "route_options": [{"model": model, "harness": harness} for model, harness in route_options], - "task": {"prompt": task}, - "route_selector": {"router_name": router_name}, - } - - def select_route( workspace: str, token: str, @@ -136,7 +122,11 @@ def select_route( None when the arm is unservable. Returns ``(decision, error)``; a failed call yields ``(None, reason)`` so callers can fail open. """ - body = route_request_body(task, route_options, router_name=router_name) + body = { + "route_options": [{"model": model, "harness": harness} for model, harness in route_options], + "task": {"prompt": task}, + "route_selector": {"router_name": router_name}, + } request = urllib.request.Request( workspace.rstrip("/") + ROUTING_PATH, data=json.dumps(body).encode("utf-8"), diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index 3db3f411..e8ce72c6 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -246,18 +246,14 @@ def _request_claude_routing_decision( if not available: return None, "Anthropic models endpoint returned no Claude models" route_options = [(model, "claude") for model in available] - router_name = routing.configured_router_name() - select_kwargs: dict[str, object] = { - "router_name": router_name, - "timeout": CLAUDE_ROUTE_SELECTION_TIMEOUT_S, - } return routing.select_route( workspace, token, prompt, route_options, lambda selected: available.get(_claude_router_model_id(selected)), - **select_kwargs, + router_name=routing.configured_router_name(), + timeout=CLAUDE_ROUTE_SELECTION_TIMEOUT_S, ) @@ -277,12 +273,7 @@ def _route_claude_prompt( raise RuntimeError( discovery_error or "Anthropic models endpoint returned no Claude models" ) - decision, error = _request_claude_routing_decision( - workspace, - token, - prompt, - model_ids, - ) + decision, error = _request_claude_routing_decision(workspace, token, prompt, model_ids) if decision is None: raise RuntimeError(error or "router returned no Claude model selection") return decision @@ -301,10 +292,7 @@ def route_claude_pre_tool_use( payload, is_spawn_agent=claude_routing.is_spawn_agent_tool, decision_fn=lambda task: _request_claude_routing_decision( - workspace, - token, - task, - available_models, + workspace, token, task, available_models ), default_task_label="Claude Code subagent task", model_id_mapper=lambda model: model, @@ -446,12 +434,7 @@ def launch_claude( model_setting = _ClaudeModelSettingGuard(user_settings_path) def route_prompt(prompt: str) -> claude_pty.FirstPromptRoute: - decision = _route_claude_prompt( - state, - token, - prompt, - model_ids, - ) + decision = _route_claude_prompt(state, token, prompt, model_ids) return claude_pty.FirstPromptRoute( model=model_name(_unwrapped_claude_model_id(decision.model)), display_model=catalog.model_id_to_display_name.get(decision.model, decision.model), @@ -460,7 +443,7 @@ def route_prompt(prompt: str) -> claude_pty.FirstPromptRoute: print_note( "Smart routing v2: the first submitted prompt will select Claude Code's " - "model." + f"model; log: {CLAUDE_PTY_LOG}." ) try: returncode = claude_pty.run_claude_pty( diff --git a/tests/test_claude_smart_routing_v2.py b/tests/test_claude_smart_routing_v2.py index 8f37a80a..2a2c8eaa 100644 --- a/tests/test_claude_smart_routing_v2.py +++ b/tests/test_claude_smart_routing_v2.py @@ -172,21 +172,7 @@ def test_strips_gateway_prefix_for_interposer(self): def test_restores_model_captured_immediately_before_switch(self, tmp_path, monkeypatch): ucode_settings = tmp_path / "ucode-settings.json" user_settings = tmp_path / "settings.json" - ucode_settings.write_text( - json.dumps( - { - "env": { - "ANTHROPIC_BASE_URL": "https://gw", - "ANTHROPIC_CUSTOM_HEADERS": "\n".join( - [ - "x-databricks-traffic-id: testenv://liteswap/arnav-r315-task-v3", - 'Databricks-Ai-Gateway-Request-Tags: {"source":"isaac-cli"}', - ] - ), - } - } - ) - ) + ucode_settings.write_text(json.dumps({"env": {"ANTHROPIC_BASE_URL": "https://gw"}})) user_settings.write_text(json.dumps({"model": "opus", "theme": "dark"})) monkeypatch.setattr(claude, "APP_DIR", tmp_path) monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", ucode_settings) @@ -203,17 +189,15 @@ def test_restores_model_captured_immediately_before_switch(self, tmp_path, monke model_id_to_display_name={"system.ai.claude-sonnet-5": "Claude Sonnet 5"}, ), ) - routed_call = {} - - def route_claude_prompt(*_args, **kwargs): - routed_call.update(kwargs) - return v2.routing.RoutingDecision( + monkeypatch.setattr( + v2, + "_route_claude_prompt", + lambda *_args: v2.routing.RoutingDecision( model="system.ai.claude-sonnet-5", raw_model="claude-sonnet-5", rationale="Selected for the parser task.", - ) - - monkeypatch.setattr(v2, "_route_claude_prompt", route_claude_prompt) + ), + ) captured: dict = {} def fake_run(argv, **kwargs): @@ -419,6 +403,7 @@ def test_routes_agent_prompt_with_initialized_model_menu(self, tmp_path, monkeyp captured = {} decisions_path = tmp_path / "decisions.jsonl" monkeypatch.setattr(v2.claude_routing, "DECISIONS_PATH", decisions_path) + def fake_select(workspace, token, task, route_options, resolve, **kwargs): captured.update( workspace=workspace, diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index d3a99b02..fbc61e7d 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -167,13 +167,6 @@ def start_interposer(*args, **kwargs): monkeypatch.setattr(codex_interposer, "start_interposer_thread", start_interposer) - def render_overlay(*args, **kwargs): - overlay = codex.render_overlay(*args, **kwargs) - overlay["model_providers"]["ucode-databricks"]["http_headers"][ - "x-databricks-traffic-id" - ] = "testenv://liteswap/arnav-r315-task-v3" - return overlay - with pytest.raises(SystemExit) as exc: v2.launch_codex( { @@ -185,7 +178,7 @@ def render_overlay(*args, **kwargs): ["--search"], binary="codex", start_model="gpt-start", - render_overlay=render_overlay, + render_overlay=codex.render_overlay, ) assert exc.value.code == 7 @@ -566,6 +559,7 @@ def test_rewrites_nested_collaboration_mode_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, *, router_name, timeout): captured.update( @@ -597,6 +591,7 @@ def select_route(workspace, token, task, route_options, resolve, *, router_name, "system.ai.gpt-5-6-luna", "system.ai.glm-5-2", ], + log=logged.append, ) assert reason is None @@ -614,3 +609,17 @@ def select_route(workspace, token, task, route_options, resolve, *, router_name, ("glm-5-2", "codex"), ], } + assert len(logged) == 1 + assert logged[0].startswith(f"[ROUTE] request POST {WS}/ai-gateway/routing/v1/routes:select: ") + request_payload = json.loads(logged[0].split(": ", 1)[1]) + assert request_payload == { + "route_options": [ + {"model": "kimi-k3-neo", "harness": "codex"}, + {"model": "gpt-5-6-sol", "harness": "codex"}, + {"model": "gpt-5-6-luna", "harness": "codex"}, + {"model": "glm-5-2", "harness": "codex"}, + ], + "task": {"prompt": "Fix the parser"}, + "route_selector": {"router_name": codex_routing.routing.ROUTER_NAME}, + } + assert "same-oauth-token" not in logged[0] From 14ee97e2e1557ee2f6bdce082465a2fc42c47353 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Fri, 11 Sep 2026 02:39:51 +0000 Subject: [PATCH 5/5] hi --- src/ucode/cli.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index eb82d14d..21a853f2 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1944,9 +1944,7 @@ def _should_launch_smart_routing( return _smart_routing_launch_shape(tool, tool_args, explicit_prompt) -def _smart_routing_launch_shape( - tool: str, tool_args: list[str], explicit_prompt: bool -) -> bool: +def _smart_routing_launch_shape(tool: str, tool_args: list[str], explicit_prompt: bool) -> bool: """Whether the forwarded arguments represent an interactive launch.""" if not tool_args or explicit_prompt: return True