Support launch-only custom request headers for Codex

README.md

@@ -395,6 +395,7 @@ 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 |
@@ -419,6 +420,9 @@ 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; they are not written to its persistent config.
+
 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/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/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,
     )
 
 

src/ucode/cli.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import os
+import re
 import shutil
 import subprocess
 from collections.abc import Iterator
@@ -2063,6 +2064,54 @@ def _smart_routing_launch_shape(tool: str, tool_args: list[str], explicit_prompt
     return tool == "claude" and tool_args[0].startswith("-")
 
 
+_HTTP_HEADER_NAME_PATTERN = re.compile(r"[!#$%&'*+\-.^_`|~0-9A-Za-z]+")
+_PROTECTED_CUSTOM_HEADER_NAMES = frozenset(
+    {
+        "authorization",
+        "connection",
+        "content-length",
+        "cookie",
+        "databricks-model-provider-service",
+        "databricks-model-service-parent-schema",
+        "host",
+        "keep-alive",
+        "proxy-authenticate",
+        "proxy-authorization",
+        "te",
+        "trailer",
+        "transfer-encoding",
+        "upgrade",
+        "user-agent",
+        "x-api-key",
+        "x-databricks-ai-gateway-token",
+        "x-databricks-use-coding-agent-mode",
+    }
+)
+
+
+def _parse_custom_headers(values: list[str] | None) -> dict[str, str]:
+    """Parse repeatable ``--header 'Name: value'`` options."""
+    parsed: dict[str, tuple[str, str]] = {}
+    for item in values or []:
+        name, separator, value = item.partition(":")
+        name = name.strip()
+        if not separator or _HTTP_HEADER_NAME_PATTERN.fullmatch(name) is None:
+            raise RuntimeError("--header must use the format `Name: value` with a valid name.")
+        value = value.strip()
+        if any(
+            ord(character) < 32 or ord(character) == 127 or character in "\u0085\u2028\u2029"
+            for character in value
+        ):
+            raise RuntimeError(
+                "--header values cannot contain control characters or line separators."
+            )
+        normalized_name = name.casefold()
+        if normalized_name in _PROTECTED_CUSTOM_HEADER_NAMES:
+            raise RuntimeError(f"--header cannot override protected header '{name}'.")
+        parsed[normalized_name] = (name, value)
+    return dict(parsed.values())
+
+
 def _launch_options(
     tool: str,
     tool_args: list[str],
@@ -2071,10 +2120,12 @@ def _launch_options(
     explicit_prompt: bool,
     user_pinned_model: str | None,
     provider: str | None,
+    custom_headers: dict[str, str] | None = None,
 ) -> LaunchOptions:
     return LaunchOptions(
         # Pinned models for providers are resolved above through the provider-specific launch path.
         user_pinned_model=user_pinned_model if provider is None else None,
+        custom_headers=tuple((custom_headers or {}).items()),
         launch_smart_routing=(
             # Smart routing is enabled globally.
             smart_routing_enabled
@@ -2125,9 +2176,11 @@ def _launch_tool(
     model: str | None = None,
     parent_schema: str | None = None,
     custom_oauth: CustomOAuthConfig | None = None,
+    headers: list[str] | None = None,
 ) -> None:
     try:
         tool = normalize_tool(tool_name)
+        custom_headers = _parse_custom_headers(headers)
         # Before any status print: a stdio-protocol subcommand owns stdout, so
         # every ug line from here on must go to stderr instead.
         if _child_owns_stdout(tool, ctx.args):
@@ -2392,6 +2445,7 @@ def _launch_tool(
             # initial/fallback model and still participates in a routed session.
             user_pinned_model=model or forwarded_model,
             provider=provider,
+            custom_headers=custom_headers,
         )
         print_success(f"Starting {TOOL_SPECS[tool]['display']}")
         with _managed_smart_routing_environment(managed, tool):
@@ -2433,6 +2487,15 @@ def _launch_tool(
     ),
 ]
 
+CustomHeaderOption = Annotated[
+    list[str] | None,
+    typer.Option(
+        "--header",
+        help="Add an HTTP header to AI Gateway requests as `Name: value`; repeatable. "
+        "Pass before any `--` separator. Credentials and transport headers are not allowed.",
+    ),
+]
+
 
 _PROMPT_SUFFIX_KEY = "ucode_explicit_prompt_suffix"
 
@@ -2590,6 +2653,7 @@ def codex_cmd(
             help="Discover model services in `<catalog>.<schema>`. Example: main.default",
         ),
     ] = None,
+    header: CustomHeaderOption = None,
     refresh: Annotated[
         bool,
         typer.Option(
@@ -2653,6 +2717,7 @@ def codex_cmd(
                 workspace_url=workspace,
                 parent_schema=model_location,
                 custom_oauth=custom_oauth,
+                headers=header,
             )
 
 

tests/test_agent_codex.py

@@ -732,6 +732,7 @@ def _patch(tmp_path, monkeypatch):
         monkeypatch.setattr(codex, "exec_or_spawn", lambda argv: launches.append(argv))
         monkeypatch.setattr(codex, "get_databricks_token", lambda workspace, profile=None: "tok")
         monkeypatch.setattr(codex, "clear_model_preferences", lambda state: False)
+        monkeypatch.setattr(codex, "codex_managed_config_path", lambda: None)
         return launches
 
     def test_sets_oauth_token(self, tmp_path, monkeypatch):
@@ -745,6 +746,46 @@ def test_sets_oauth_token(self, tmp_path, monkeypatch):
         assert os.environ["OAUTH_TOKEN"] == "fresh-token"
         assert launches[0][-1] == "--search"
 
+    def test_custom_header_is_launch_only(self, tmp_path, monkeypatch):
+        launches = self._patch(tmp_path, monkeypatch)
+        value = "route://development/test"
+
+        codex.launch(
+            {"workspace": WS},
+            [],
+            options=LaunchOptions(custom_headers=(("X-Development-Route", value),)),
+        )
+
+        env_name = f"{codex.CUSTOM_HEADER_ENV_PREFIX}_{os.getpid()}_0"
+        provider_arg = next(
+            arg for arg in launches[0] if arg.startswith("model_providers.Databricks=")
+        )
+        assert "X-Development-Route" in provider_arg
+        assert env_name in provider_arg
+        assert value not in provider_arg
+        assert os.environ[env_name] == value
+        assert "X-Development-Route" not in codex.CODEX_CONFIG_PATH.read_text(encoding="utf-8")
+        monkeypatch.delenv(env_name)
+
+    @pytest.mark.parametrize("header_table", ["http_headers", "env_http_headers"])
+    def test_custom_header_rejects_managed_header(self, tmp_path, monkeypatch, header_table):
+        self._patch(tmp_path, monkeypatch)
+        managed_path = tmp_path / "managed_config.toml"
+        managed_path.write_text(
+            f'[model_providers.Databricks.{header_table}]\nX-Test = "ENTERPRISE_VALUE"\n',
+            encoding="utf-8",
+        )
+        monkeypatch.setattr(codex, "codex_managed_config_path", lambda: managed_path)
+
+        with pytest.raises(RuntimeError, match="OS-managed Codex header"):
+            codex.launch(
+                {"workspace": WS},
+                [],
+                options=LaunchOptions(custom_headers=(("x-test", "temporary"),)),
+            )
+
+        assert "temporary" not in managed_path.read_text(encoding="utf-8")
+
     def test_provider_discovery_uses_authoritative_catalog(self, tmp_path, monkeypatch):
         launches = self._patch(tmp_path, monkeypatch)
         catalog_path = tmp_path / "models.json"
@@ -962,6 +1003,19 @@ def test_provider_rejects_managed_model_catalog(self, tmp_path, monkeypatch):
 
         assert launches == []
 
+    def test_custom_header_requires_modern_codex(self, tmp_path, monkeypatch):
+        launches = self._patch(tmp_path, monkeypatch)
+        monkeypatch.setattr(codex, "agent_version", lambda binary: "0.133.0")
+
+        with pytest.raises(RuntimeError, match="--header requires Codex 0.134.0 or newer"):
+            codex.launch(
+                {"workspace": WS},
+                [],
+                options=LaunchOptions(custom_headers=(("X-Test", "temporary"),)),
+            )
+
+        assert launches == []
+
     def test_non_provider_launch_removes_stale_provider_header(self, tmp_path, monkeypatch):
         launches = self._patch(tmp_path, monkeypatch)
         profile_path = tmp_path / "ucode.config.toml"