Support temporary global custom request headers for Claude Code

README.md

@@ -395,9 +395,11 @@ The output looks like:
 | `ug codex --enable-smart-routing` | Enable AI Gateway routing for Codex sessions and subagents |
 | `ug codex --refresh` | Re-check Databricks, refresh models/configuration, and launch Codex |
 | `ug codex --model-location main.default` | Discover model services in the specified catalog and schema |
+| `ug codex --header 'X-Development-Route: test-target'` | Add a launch-only custom header to Codex gateway requests |
 | `ug claude --enable-smart-routing` | Enable AI Gateway routing for Claude Code sessions and subagents |
 | `ug claude --refresh` | Re-check Databricks, refresh models/configuration, and launch Claude Code |
 | `ug claude --model-location main.default` | Discover model services in the specified catalog and schema |
+| `ug claude --header 'X-Development-Route: test-target'` | Set a temporary global header for Claude Code gateway requests |
 | `ug configure --agents claude,codex,pi` | Configure the requested agents that are available; skip the rest with a warning |
 | `ug configure --agents claude --mcp system.ai.slack` | Configure an agent and register its Databricks MCP server(s) in one command |
 | `ug mcp add --location system.ai` | Register a schema's MCP servers, keeping any already configured (additive; never removes) |
@@ -419,6 +421,10 @@ The output looks like:
 | `ug skill remove --location main.default [--path <dir>]` | Delete every skill downloaded from a schema (all bases, or one under `<dir>`) |
 | `ug skill remove --skills main.default.my-skill [--path <dir>]` | Delete named downloaded skills by fully-qualified name (comma-separated; may span schemas; `--path` limits to one base) |
 
+`--header` is repeatable and intended for non-secret development routing values. Codex reads the
+values from launch-only environment variables. Claude Code requires them in global settings, so
+they apply to all Claude sessions until the next successful launch or configuration without them.
+
 Databricks AI Tools are installed only by `ug configure`, never by `ug <agent>` launches.
 Use `--enable-databricks-ai-tools` or `--disable-databricks-ai-tools` with `ug configure` to
 control the installation.

src/ucode/agents/__init__.py

@@ -394,6 +394,7 @@ def configure_tool(
     custom_model: str | None = None,
     coding_agent_config_defaults: dict[str, str] | None = None,
     parent_schema: str | None = None,
+    custom_headers: dict[str, str] | None = None,
 ) -> dict:
     result: dict | tuple[dict, str]
     if tool == "codex":
@@ -415,6 +416,7 @@ def configure_tool(
             custom_model=custom_model,
             coding_agent_config_defaults=coding_agent_config_defaults,
             parent_schema=parent_schema,
+            custom_headers=custom_headers,
         )
     else:
         # Every tool in this branch needs a model — including gemini under a provider,

src/ucode/agents/args.py

@@ -11,6 +11,7 @@ class LaunchOptions:
 
     launch_smart_routing: bool = False
     user_pinned_model: str | None = None
+    custom_headers: tuple[tuple[str, str], ...] = ()
 
 
 def explicit_model_arg_value(tool_args: list[str]) -> str | None:

src/ucode/agents/claude.py

@@ -11,7 +11,7 @@
 import socket
 import subprocess
 import threading
-from collections.abc import Callable
+from collections.abc import Callable, Collection
 from pathlib import Path
 from typing import cast
 
@@ -60,7 +60,14 @@
     remove_smart_routing_hooks,
     sync_smart_routing_hooks,
 )
-from ucode.state import MANAGED_OVERLAY_KEY, is_tool_managed, mark_tool_managed, save_state
+from ucode.state import (
+    MANAGED_OVERLAY_KEY,
+    is_tool_managed,
+    load_global_state,
+    mark_tool_managed,
+    save_global_state,
+    save_state,
+)
 from ucode.telemetry import agent_version, ug_version
 from ucode.tracing import tracing_env
 from ucode.ui import print_note, print_success, print_warning
@@ -76,6 +83,7 @@
 # The default model is stored in Claude's default user settings, not the ucode settings.
 CLAUDE_USER_SETTINGS_PATH = CLAUDE_CONFIG_DIR / "settings.json"
 CLAUDE_BACKUP_PATH = APP_DIR / "claude-ucode-settings.backup.json"
+CLAUDE_CUSTOM_HEADERS_STATE_KEY = "claude_custom_headers"
 WEB_SEARCH_MCP_STATE_KEY = "claude_web_search_mcp"
 MINIMUM_CLAUDE_VERSION = (2, 1, 248)
 MINIMUM_CLAUDE_VERSION_TEXT = "2.1.248"
@@ -371,6 +379,7 @@ def render_overlay(
     parent_schema: str | None = None,
     static_models: list[str] | None = None,
     otel_tracing: bool = False,
+    custom_headers: dict[str, str] | None = None,
 ) -> tuple[dict, list[list[str]]]:
     """Return (overlay, managed_key_paths) for Claude settings.json.
 
@@ -410,12 +419,13 @@ 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}")
+    header_lines.extend(f"{name}: {value}" for name, value in (custom_headers or {}).items())
     # 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)
+    rendered_custom_headers = "\n".join(header_lines)
     env: dict[str, str] = {
         "ANTHROPIC_BASE_URL": base_url,
-        "ANTHROPIC_CUSTOM_HEADERS": custom_headers,
+        "ANTHROPIC_CUSTOM_HEADERS": rendered_custom_headers,
         "CLAUDE_CODE_API_KEY_HELPER_TTL_MS": "900000",
         # 1h prompt caching needs the extended-cache-ttl beta header, which
         # Claude Code only sends when experimental betas are enabled — so we must
@@ -721,6 +731,65 @@ def disable_smart_routing(state: dict) -> bool:
     return changed
 
 
+def _custom_headers_from_settings(settings: dict) -> dict[str, set[str]]:
+    env = settings.get("env")
+    value = env.get(ANTHROPIC_CUSTOM_HEADERS_ENV_KEY) if isinstance(env, dict) else None
+    if not isinstance(value, str):
+        return {}
+    headers: dict[str, set[str]] = {}
+    for line in value.splitlines():
+        name, separator, header_value = line.partition(":")
+        if separator:
+            headers.setdefault(name.strip().casefold(), set()).add(header_value.strip())
+    return headers
+
+
+def _recorded_custom_headers() -> dict[str, str]:
+    recorded = load_global_state().get(CLAUDE_CUSTOM_HEADERS_STATE_KEY)
+    if not isinstance(recorded, dict):
+        return {}
+    return {
+        name.casefold(): value
+        for name, value in recorded.items()
+        if isinstance(name, str) and isinstance(value, str)
+    }
+
+
+def _reject_custom_header_collisions(
+    custom_headers: dict[str, str], previous_custom_headers: dict[str, str]
+) -> None:
+    if not custom_headers:
+        return
+
+    settings_sources = [read_json_safe(CLAUDE_SETTINGS_PATH)]
+    managed_path = _managed_settings_path()
+    if managed_path is not None:
+        managed_text = read_managed_file(managed_path)
+        if managed_text is not None:
+            try:
+                settings_sources.append(_parse_managed_settings(managed_text))
+            except RuntimeError as exc:
+                raise RuntimeError(
+                    f"Cannot safely inspect Claude Code managed settings at {managed_path}: {exc}."
+                ) from exc
+
+    conflicts: set[str] = set()
+    for settings in settings_sources:
+        existing = _custom_headers_from_settings(settings)
+        for name in custom_headers:
+            previous_value = previous_custom_headers.get(name)
+            if existing.get(name, set()) - (
+                {previous_value} if previous_value is not None else set()
+            ):
+                conflicts.add(name)
+    if conflicts:
+        names = ", ".join(sorted(conflicts))
+        raise RuntimeError(
+            f"--header cannot override existing Claude Code header(s): {names}. "
+            "Remove the existing setting or use a different header name."
+        )
+
+
 def write_tool_config(
     state: dict,
     model: str | None,
@@ -731,7 +800,19 @@ def write_tool_config(
     custom_model: str | None = None,
     coding_agent_config_defaults: dict[str, str] | None = None,
     parent_schema: str | None = None,
+    custom_headers: dict[str, str] | None = None,
 ) -> dict:
+    previous_custom_headers = _recorded_custom_headers()
+    current_custom_headers = {
+        name.casefold(): value for name, value in (custom_headers or {}).items()
+    }
+    _reject_custom_header_collisions(current_custom_headers, previous_custom_headers)
+    if previous_custom_headers or current_custom_headers:
+        global_state = load_global_state()
+        global_state[CLAUDE_CUSTOM_HEADERS_STATE_KEY] = (
+            previous_custom_headers | current_custom_headers
+        )
+        save_global_state(global_state)
     # Back up only a file that predates ucode's management of the tool. A
     # re-configure would otherwise snapshot ucode's own generated file, and
     # revert would restore that snapshot instead of deleting the file.
@@ -758,6 +839,7 @@ def write_tool_config(
         parent_schema=parent_schema,
         static_models=state.get("claude_static_models"),
         otel_tracing=bool(state.get("claude_otel_tracing")),
+        custom_headers=custom_headers,
     )
     tracing_env_vars = tracing_env(state, "claude")
     stop_hook_command = claude_tracing_stop_hook_command() if tracing_env_vars else None
@@ -821,10 +903,16 @@ def _compose(base: dict, *, enforce_model_default_hierarchy: bool) -> dict:
                     target_env.pop(key, None)
                 else:
                     target_env[key] = selected_default_model
+        managed_custom_header_names = (
+            CLAUDE_MANAGED_CUSTOM_HEADER_NAMES | current_custom_headers.keys()
+        )
         merged = deep_merge_dict(base, overlay_for_merge)
         overlay_custom_headers = overlay_for_merge["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY]
         merged["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] = _merge_anthropic_custom_headers(
-            existing_custom_headers, overlay_custom_headers
+            existing_custom_headers,
+            overlay_custom_headers,
+            managed_custom_header_names,
+            previous_custom_headers,
         )
         # Drop any apiKeyHelper a prior non-relayed launch left in the file; relayed
         # must not carry one (it would outrank the subscription OAuth).
@@ -895,20 +983,31 @@ def _compose(base: dict, *, enforce_model_default_hierarchy: bool) -> dict:
     else:
         state.pop("claude_relayed", None)
         state.pop("relayed_proxy_port", None)
+    global_state = load_global_state()
+    if current_custom_headers:
+        global_state[CLAUDE_CUSTOM_HEADERS_STATE_KEY] = current_custom_headers
+    else:
+        global_state.pop(CLAUDE_CUSTOM_HEADERS_STATE_KEY, None)
+    save_global_state(global_state)
     state = mark_tool_managed(state, "claude", managed_keys)
     save_state(state)
     return state
 
 
-def _merge_anthropic_custom_headers(existing: object, ucode_headers: str) -> str:
+def _merge_anthropic_custom_headers(
+    existing: object,
+    ucode_headers: str,
+    managed_header_names: Collection[str] = CLAUDE_MANAGED_CUSTOM_HEADER_NAMES,
+    removable_header_values: dict[str, str] | None = None,
+) -> str:
     """Preserve user headers while replacing the header names managed by ucode.
 
     Claude's ``ANTHROPIC_CUSTOM_HEADERS`` value is a newline-delimited string. To merge it, we:
 
     1. Split the existing custom headers by newline into individual header items.
     2. Split each item on ``:`` to identify its header name.
-    3. Replace headers in ``CLAUDE_MANAGED_CUSTOM_HEADER_NAMES`` with ucode's values in their
-       existing positions, while preserving all other existing headers.
+    3. Replace headers in ``managed_header_names`` with ucode's values and remove
+       exact values previously recorded as temporary, while preserving all others.
     4. Append any ucode-managed headers that were not already present.
 
     Header names are compared case-insensitively. Non-header lines are also preserved to avoid
@@ -933,12 +1032,18 @@ def _merge_anthropic_custom_headers(existing: object, ucode_headers: str) -> str
     for line in existing.splitlines():
         name, separator, _value = line.partition(":")
         normalized_name = name.strip().casefold()
-        if separator and normalized_name in CLAUDE_MANAGED_CUSTOM_HEADER_NAMES:
+        if separator and normalized_name in managed_header_names:
             replacement = ucode_lines_by_name.get(normalized_name)
             if replacement is not None and normalized_name not in replaced_names:
                 merged.append(replacement)
                 replaced_names.add(normalized_name)
             continue
+        if (
+            separator
+            and removable_header_values
+            and removable_header_values.get(normalized_name) == _value.strip()
+        ):
+            continue
         if line:
             merged.append(line)
 

src/ucode/agents/codex.py

@@ -85,6 +85,7 @@
 LEGACY_CODEX_BACKUP_PATH = APP_DIR / "codex-config.backup.toml"
 CODEX_MODEL_PROVIDER_NAME = "Databricks"
 LEGACY_CODEX_MODEL_PROVIDER_NAME = "ucode-databricks"
+CUSTOM_HEADER_ENV_PREFIX = "UCODE_CUSTOM_HEADER"
 _MODEL_SERVICE_ROUTING_KEY_PATHS = [
     ["model_providers", CODEX_MODEL_PROVIDER_NAME, "http_headers", MODEL_PROVIDER_SERVICE_HEADER],
     [
@@ -503,6 +504,61 @@ def _parse_managed_config(text: str) -> dict:
         raise RuntimeError(f"invalid TOML: {exc}") from exc
 
 
+def _managed_header_names() -> set[str]:
+    """Return header names fixed by OS-managed config."""
+    path = codex_managed_config_path()
+    if path is None:
+        return set()
+    text = read_managed_file(path)
+    if text is None:
+        return set()
+    try:
+        doc = _parse_managed_config(text)
+    except RuntimeError:
+        # Configuration validates the managed file before launch and reports the
+        # full parse error there.
+        return set()
+    providers = doc.get("model_providers")
+    provider = providers.get(CODEX_MODEL_PROVIDER_NAME) if isinstance(providers, dict) else None
+    if not isinstance(provider, dict):
+        return set()
+    names: set[str] = set()
+    for table_name in ("http_headers", "env_http_headers"):
+        headers = provider.get(table_name)
+        if isinstance(headers, dict):
+            names.update(str(name).casefold() for name in headers)
+    return names
+
+
+def _with_custom_headers(doc: dict, custom_headers: dict[str, str]) -> dict:
+    """Return a launch-only config that reads custom header values from the environment."""
+    if not custom_headers:
+        return doc
+    blocked_names = _managed_header_names() & {name.casefold() for name in custom_headers}
+    if blocked_names:
+        names = ", ".join(sorted(blocked_names))
+        raise RuntimeError(f"--header cannot override OS-managed Codex header(s): {names}.")
+
+    launch_doc = copy.deepcopy(doc)
+    providers = launch_doc.get("model_providers")
+    provider = providers.get(CODEX_MODEL_PROVIDER_NAME) if isinstance(providers, dict) else None
+    if not isinstance(provider, dict):
+        raise RuntimeError("Codex's Databricks model provider configuration is missing or invalid.")
+    env_headers = provider.setdefault("env_http_headers", {})
+    if not isinstance(env_headers, dict):
+        raise RuntimeError("Codex's Databricks environment-header configuration is invalid.")
+
+    requested_names = {name.casefold() for name in custom_headers}
+    for existing_name in list(env_headers):
+        if str(existing_name).casefold() in requested_names:
+            env_headers.pop(existing_name, None)
+    for index, (name, value) in enumerate(custom_headers.items()):
+        env_name = f"{CUSTOM_HEADER_ENV_PREFIX}_{os.getpid()}_{index}"
+        os.environ[env_name] = value
+        env_headers[name] = env_name
+    return launch_doc
+
+
 def managed_config_is_current(state: dict) -> bool:
     path = codex_managed_config_path()
     if path is None:
@@ -747,8 +803,9 @@ def launch(
     *,
     options: LaunchOptions,
 ) -> None:
+    custom_headers = dict(options.custom_headers)
     if options.launch_smart_routing:
-        _launch_smart_routing(state, tool_args)
+        _launch_smart_routing(state, tool_args, custom_headers=custom_headers)
         return
     clear_model_preferences(state)
     binary = SPEC["binary"]
@@ -775,6 +832,8 @@ def launch(
         if state.get("codex_otel_tracing"):
             otel_args = codex_config_args(_otel_overlay(workspace, token))
     if _use_legacy_layout():
+        if custom_headers:
+            raise RuntimeError(f"--header requires Codex {MINIMUM_CODEX_VERSION_TEXT} or newer.")
         print_warning_err(
             f"Codex {agent_version(binary)} is outdated. Upgrade Codex to "
             f"{MINIMUM_CODEX_VERSION_TEXT} or newer, then run `codex --version` to verify "
@@ -826,10 +885,13 @@ def launch(
                 slugs = catalog_slugs(catalog)
                 if slugs:
                     profile_doc["model"] = slugs[0]
+    profile_doc = _with_custom_headers(profile_doc, custom_headers)
     exec_or_spawn([binary, *codex_config_args(profile_doc), *otel_args, *tool_args])
 
 
-def _launch_smart_routing(state: dict, tool_args: list[str]) -> None:
+def _launch_smart_routing(
+    state: dict, tool_args: list[str], *, custom_headers: dict[str, str] | None = None
+) -> None:
     """Launch the Codex TUI through the smart-routing interposer."""
     binary = SPEC["binary"]
     version_text = agent_version(binary)
@@ -848,12 +910,20 @@ def _launch_smart_routing(state: dict, tool_args: list[str]) -> None:
         or (codex_model_id(models[0]) if models else None)
         or APP_SERVER_SMART_ROUTING_STARTING_MODEL
     )
+    if custom_headers:
+
+        def routed_render_overlay(*args, **kwargs) -> dict:
+            return _with_custom_headers(render_overlay(*args, **kwargs), custom_headers)
+
+    else:
+        routed_render_overlay = render_overlay
+
     smart_routing_v2.launch_codex(
         state,
         tool_args,
         binary=binary,
         start_model=start_model,
-        render_overlay=render_overlay,
+        render_overlay=routed_render_overlay,
     )