Add ug --header / UG_HEADERS for per-launch request headers

AGENTS.md

@@ -174,7 +174,7 @@ Fields live in `~/.codex/ucode.config.toml` and `/etc/codex/managed_config.toml`
 | `model_provider` | Create/replace | Create/replace | Gateway provider (`Databricks`) |
 | `model` | Ignore | Create/replace | The config's default model |
 | `model_providers.Databricks` | Merge | Merge | Provider block: name, gateway base URL, wire API, `ug` auth command; other keys left alone |
-| `http_headers` | Merge | Merge | In `[model_providers.Databricks]`; merge `ug`'s routing headers by name, admin headers added under managed config |
+| `http_headers` | Merge | Merge | In `[model_providers.Databricks]`; merge `ug`'s routing headers by name, admin headers added under managed config. `UG_HEADERS` names not already present are added to the launch's `-c` flags only, never a file |
 | `model_catalog_json` | Create/replace | Create/replace | In `~/.codex/config.toml`; `ug`'s own catalog reference, for a static model list |
 | `otel.trace_exporter` | Ignore | Create/replace | In `~/.codex/config.toml` when the config enables tracing, for Codex desktop; another exporter and other telemetry settings left alone |
 | `mcp_servers` | Ignore | Merge | Managed file; add/update the config's MCP server entries, other entries left alone |

README.md

@@ -140,6 +140,18 @@ precedence; later launches use whatever files are present.
 The choice is saved for launches after a successful configure; `ug configure` without `--user`
 and `ug revert` return to `--global`.
 
+### Extra headers for one session
+
+```bash
+ug --header "x-databricks-traffic-id: testenv://liteswap/<pod>" codex
+UG_HEADERS="X-Team: eng" ug codex
+```
+
+`--header` (repeatable) and `UG_HEADERS` (one `Name: value` per line) add headers to that launch's
+Databricks requests only: Codex's gateway requests, through the launch's `-c` provider headers, and
+ug's own requests (workspace APIs, MCP proxy, web search, smart routing). A header already sent
+keeps its value. Nothing is written to a config file.
+
 Available coding agents are `codex`, `claude`, `gemini`, `opencode`,
 `copilot`, and `pi`. `cursor` can be included in `--agents` for MCP-only setup;
 Cursor models still run through your Cursor account.

src/ucode/agents/codex.py

@@ -95,6 +95,7 @@
 from ucode.smart_routing.routing import configured_router_name
 from ucode.state import get_provider_service, is_tool_managed, mark_tool_managed, save_state
 from ucode.telemetry import agent_version, ucode_user_agent, ug_version
+from ucode.ug_headers import ug_headers
 from ucode.ui import print_warning_err
 
 from .args import LaunchOptions
@@ -1235,9 +1236,21 @@ def _export_launch_routing_headers(config: dict) -> None:
 def _render_launch_overlay(*args, **kwargs) -> dict:
     overlay = render_overlay(*args, **kwargs)
     _export_launch_routing_headers(overlay)
+    _add_launch_headers(overlay)
     return overlay
 
 
+def _add_launch_headers(config: dict) -> None:
+    """Add this launch's ``ug --header`` / ``UG_HEADERS`` to the provider's ``-c`` headers.
+
+    A header ug or the admin already sends keeps its value. The headers reach only this launch's
+    ``-c`` flags, never a config file."""
+    present = {name.casefold() for name in _managed_provider_headers(config) or {}}
+    for name, value in ug_headers().items():
+        if name.casefold() not in present:
+            _set_routing_header(config, name, value)
+
+
 def _set_routing_header(config: dict, header: str, value: str | None) -> None:
     model_providers = config.get("model_providers")
     if not isinstance(model_providers, dict):
@@ -1534,6 +1547,7 @@ def launch(
     _set_provider_header(profile_doc, provider)
     _set_parent_schema_header(profile_doc, parent_schema if not provider else None)
     _export_launch_routing_headers(profile_doc)
+    _add_launch_headers(profile_doc)
     updating = tool_args[:1] == ["update"]
     if updating and _is_ucode_catalog_reference(profile_doc.get("model_catalog_json")):
         profile_doc.pop("model_catalog_json")

src/ucode/cli.py

@@ -54,6 +54,7 @@
     AGENT_CODEX,
     ISAAC_SMART_ROUTER_FEEDBACK_ENV_VAR,
     SMART_ROUTING_ENV_KEYS,
+    UG_HEADERS_ENV_VAR,
 )
 from ucode.custom_oauth import (
     CUSTOM_OAUTH_CLI_ENV_VAR,
@@ -182,6 +183,7 @@
     workspace_self_managed_agents,
 )
 from ucode.string_utils import is_valid_catalog_schema
+from ucode.ug_headers import parse_header_line
 from ucode.ui import (
     console,
     heading,
@@ -3760,8 +3762,24 @@ def default(
     ] = False,
     skip_preflight: SkipPreflightOption = False,
     workspace: WorkspaceOption = None,
+    header: Annotated[
+        list[str] | None,
+        typer.Option(
+            "--header",
+            help="Extra `Name: value` header for this launch's Codex and ug Databricks requests "
+            "(repeatable). A header already sent keeps its value.",
+        ),
+    ] = None,
 ) -> None:
     """Configure and launch coding agents through Databricks AI Gateway."""
+    for number, line in enumerate(header or [], 1):
+        if parse_header_line(line) is None:
+            # Name only the flag's position: the value may be a token.
+            print_err(f"--header #{number}: expected `Name: value` with a valid name and value.")
+            raise typer.Exit(1)
+    if header:
+        existing = os.environ.get(UG_HEADERS_ENV_VAR, "")
+        os.environ[UG_HEADERS_ENV_VAR] = "\n".join(filter(None, [existing, *header]))
     if ctx.invoked_subcommand != "revert":
         set_config_scope(load_config_scope())
     if ctx.invoked_subcommand is not None:

src/ucode/constants.py

@@ -40,6 +40,8 @@
 # reads the env var lives in ``ucode.agents.request_tags`` (it needs ``ui``).
 AI_GATEWAY_REQUEST_TAGS_HEADER = "Databricks-Ai-Gateway-Request-Tags"
 AI_GATEWAY_REQUEST_TAGS_ENV_VAR = "AI_GATEWAY_REQUEST_TAGS"
+# Extra Databricks request headers, one `Name: value` per line; `ug --header` adds to it.
+UG_HEADERS_ENV_VAR = "UG_HEADERS"
 
 
 # MCP server registration scopes. Claude Code supports local/project/user; the

src/ucode/databricks.py

@@ -43,6 +43,7 @@
 from ucode.os_compatibility import subprocess_cross_os
 from ucode.os_compatibility.file_lock_cross_os import acquire_exclusive_file_lock
 from ucode.telemetry import ucode_user_agent, ug_version
+from ucode.ug_headers import ug_headers
 from ucode.ui import (
     err_console,
     normalize_workspace_url,
@@ -339,7 +340,7 @@ def _http_get_bytes(
     if max_retries < 0:
         raise ValueError("max_retries must be non-negative")
 
-    request_headers = {"Authorization": f"Bearer {token}"}
+    request_headers = {**ug_headers(), "Authorization": f"Bearer {token}"}
     request_headers.update(headers or {})
     request = urllib_request.Request(url, headers=request_headers)
     for attempt in range(max_retries + 1):
@@ -443,7 +444,11 @@ def _http_send_json(
     empty body there is the expected result, not a decode failure.
     """
     body_bytes = json.dumps(payload).encode("utf-8") if payload is not None else None
-    request_headers = {"Authorization": f"Bearer {token}", "Accept": "application/json"}
+    request_headers = {
+        **ug_headers(),
+        "Authorization": f"Bearer {token}",
+        "Accept": "application/json",
+    }
     if body_bytes is not None:
         request_headers["Content-Type"] = "application/json"
     request_headers.update(headers or {})

src/ucode/mcp_oauth.py

@@ -18,6 +18,7 @@
 from urllib.parse import urlencode
 
 from ucode.config_io import APP_DIR, read_json_safe, write_json_file
+from ucode.ug_headers import ug_headers
 
 # Claude Code's published URL client ID (Client ID Metadata Document). The
 # workspace must recognize it before ug registers it; merely serving the
@@ -77,7 +78,7 @@ def _probe_oauth_client(workspace: str, client_id: str) -> bool | None:
         f"{workspace.rstrip('/')}/oidc/v1/token",
         data=body,
         method="POST",
-        headers={"Content-Type": "application/x-www-form-urlencoded"},
+        headers={**ug_headers(), "Content-Type": "application/x-www-form-urlencoded"},
     )
     try:
         urllib.request.urlopen(request, timeout=10)  # noqa: S310 - fixed https workspace URL

src/ucode/mcp_proxy.py

@@ -52,6 +52,7 @@
     mcp_service_needs_connection_login,
 )
 from ucode.mcp_connection_login import connection_from_url, run_connection_login
+from ucode.ug_headers import ug_headers
 
 # Exit code used when the proxy cannot continue. MCP clients surface a non-zero
 # exit far more usefully than a timeout, so bail out instead of hanging.
@@ -521,6 +522,7 @@ async def _run(url: str, workspace: str, profile: str | None, use_pat: bool = Fa
     # yields a (read, write) pair in both.
     async with httpx.AsyncClient(
         auth=auth,
+        headers=ug_headers(),
         timeout=httpx.Timeout(
             connect=30.0,
             read=300.0,

src/ucode/mcp_web_search.py

@@ -22,6 +22,7 @@
 from urllib import request as urllib_request
 
 from ucode.databricks import get_databricks_token
+from ucode.ug_headers import ug_headers
 
 PROTOCOL_VERSION = "2024-11-05"
 SERVER_NAME = "ucode-web-search"
@@ -145,6 +146,7 @@ def _call_responses_api(query: str) -> dict[str, Any]:
         data=body,
         method="POST",
         headers={
+            **ug_headers(),
             "Authorization": f"Bearer {token}",
             "Content-Type": "application/json",
             "Accept": "application/json",

src/ucode/smart_routing/routing.py

@@ -25,6 +25,7 @@
 from typing import Any
 
 from ucode.smart_routing.rationale import format_rationale
+from ucode.ug_headers import ug_headers
 
 ROUTER_NAME = "task_v3"
 ROUTER_NAME_ENV_VAR = "SMART_ROUTER_NAME"
@@ -239,6 +240,7 @@ def select_route(
         workspace.rstrip("/") + ROUTING_PATH,
         data=json.dumps(body).encode("utf-8"),
         headers={
+            **ug_headers(),
             "Authorization": f"Bearer {token}",
             "Content-Type": "application/json",
         },

src/ucode/ug_headers.py

@@ -0,0 +1,42 @@
+"""Extra request headers from ``ug --header`` / ``UG_HEADERS``.
+
+Lives outside ``ucode.agents`` so ``ucode.databricks``, which that package imports, can import it.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+
+from ucode.constants import UG_HEADERS_ENV_VAR
+
+_HEADER_NAME_RE = re.compile(r"[!#$%&'*+.^_`|~0-9A-Za-z-]+")
+
+
+def _is_sendable_value(value: str) -> bool:
+    """No control characters, and only Latin-1, which urllib and httpx encode header values as."""
+    return all(char == "\t" or " " <= char <= "~" or "\x80" <= char <= "\xff" for char in value)
+
+
+def parse_header_line(line: str) -> tuple[str, str] | None:
+    """``(name, value)`` for one ``Name: value`` line, or ``None`` when it isn't a valid header."""
+    name, separator, value = line.partition(":")
+    name, value = name.strip(), value.strip()
+    if not separator or not _HEADER_NAME_RE.fullmatch(name) or not value:
+        return None
+    return (name, value) if _is_sendable_value(value) else None
+
+
+def ug_headers() -> dict[str, str]:
+    """``UG_HEADERS`` as ``{name: value}``, one ``Name: value`` per line.
+
+    Invalid lines are skipped, so a bad inherited value never blocks a command. A later line
+    replaces an earlier one with the same name in any case, so ``--header`` overrides an inherited
+    value. Spread it first into a urllib header dict: urllib capitalizes names and keeps the last
+    value, so ug's own headers listed after it win in any case."""
+    headers: dict[str, tuple[str, str]] = {}
+    for line in os.environ.get(UG_HEADERS_ENV_VAR, "").splitlines():
+        parsed = parse_header_line(line)
+        if parsed:
+            headers[parsed[0].casefold()] = parsed
+    return dict(headers.values())

tests/README.md

@@ -43,6 +43,11 @@ TypeScript event-forwarding and launch-option tests in CI. The existing Claude
 skill-toggle integration journey checks the live recipe payload field across on/off/on.
 See [the mod guide](../docs/claude-mod.md).
 
+`test_ug_headers.py` covers the `UG_HEADERS` parser and the `--header` flag, `test_databricks.py`
+ug's own requests carrying them, and `test_agent_codex.py` `UG_HEADERS` reaching a Codex launch's
+`-c` headers (plain and smart-routed) without writing its profile. No automated test checks the
+header at the gateway.
+
 The required Isaac feedback form and background delivery notifications need manual
 validation using the released Isaac skill and existing capture/retry CLI. Check
 category selection, empty and inline feedback, native retries, the off form's

tests/conftest.py

@@ -77,6 +77,8 @@ def reject_privileged_write(path, _desired_text):
     monkeypatch.delenv("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", raising=False)
     monkeypatch.delenv("UCODE_SESSION_ENV_FILE", raising=False)
     monkeypatch.delenv("UCODE_SMART_ROUTER_PYTHON", raising=False)
+    # `ug --header` exports this to the agent, so an outer session's value would leak into launches.
+    monkeypatch.delenv("UG_HEADERS", raising=False)
     # When the suite runs inside a ug smart-routing session, these ambient env vars
     # make smart_routing_enabled()/configured_router_name() read the host's live
     # routing state, leaking a Databricks-Smart-Router-Recipe header into

tests/test_agent_codex.py

@@ -1489,6 +1489,38 @@ def test_non_provider_launch_removes_stale_provider_header(self, tmp_path, monke
         )
         assert "Databricks-Model-Provider-Service" not in provider_arg
 
+    def test_launch_adds_ug_headers_without_writing_the_profile(self, tmp_path, monkeypatch):
+        launches = self._patch(tmp_path, monkeypatch)
+        profile_path = tmp_path / "ucode.config.toml"
+        before = profile_path.read_text(encoding="utf-8")
+        monkeypatch.setenv("UG_HEADERS", "x-databricks-traffic-id: testenv://liteswap/a")
+
+        codex.launch({"workspace": WS}, [], options=LaunchOptions())
+
+        provider_arg = next(
+            arg for arg in launches[0] if arg.startswith("model_providers.Databricks=")
+        )
+        assert "x-databricks-traffic-id" in provider_arg
+        assert "testenv://liteswap/a" in provider_arg
+        assert profile_path.read_text(encoding="utf-8") == before
+
+    def test_smart_routing_overlay_adds_ug_headers_keeping_existing_ones(self, monkeypatch):
+        monkeypatch.setenv("UG_HEADERS", "x-team: other\nX-New: 1")
+        monkeypatch.setattr(
+            codex,
+            "render_overlay",
+            lambda *a, **k: {
+                "model_providers": {"Databricks": {"http_headers": {"X-Team": "eng"}}}
+            },
+        )
+
+        overlay = codex._render_launch_overlay(WS)
+
+        assert overlay["model_providers"]["Databricks"]["http_headers"] == {
+            "X-Team": "eng",
+            "X-New": "1",
+        }
+
     def test_provider_discovery_uses_custom_oauth_token(self, tmp_path, monkeypatch):
         self._patch(tmp_path, monkeypatch)
         catalog_path = tmp_path / "models.json"

tests/test_databricks.py

@@ -60,6 +60,8 @@
 class _FakeResponse:
     """Minimal urlopen context manager returning a JSON body, or raw bytes as given."""
 
+    status = 200
+
     def __init__(self, payload: dict | bytes):
         self._body = payload if isinstance(payload, bytes) else json.dumps(payload).encode("utf-8")
 
@@ -3975,6 +3977,45 @@ def fake_urlopen(request, timeout=None):
         assert reason is None
 
 
+class TestHttpUgHeaders:
+    def test_get_and_send_carry_ug_headers(self, monkeypatch):
+        monkeypatch.setenv("UG_HEADERS", "x-databricks-traffic-id: testenv://liteswap/abc")
+        requests = []
+
+        def fake_urlopen(request, timeout=None):
+            requests.append(request)
+            return _FakeResponse({})
+
+        monkeypatch.setattr(db_mod.urllib_request, "urlopen", fake_urlopen)
+
+        db_mod._http_get_json(f"{WS}/get", "tok")
+        db_mod._http_post_json(f"{WS}/post", "tok", {})
+
+        assert [request.get_header("X-databricks-traffic-id") for request in requests] == [
+            "testenv://liteswap/abc",
+            "testenv://liteswap/abc",
+        ]
+
+    def test_ug_headers_never_replace_auth_or_content_type(self, monkeypatch):
+        monkeypatch.setenv("UG_HEADERS", "authorization: other\ncontent-type: text/plain")
+        requests = []
+
+        def fake_urlopen(request, timeout=None):
+            requests.append(request)
+            return _FakeResponse({})
+
+        monkeypatch.setattr(db_mod.urllib_request, "urlopen", fake_urlopen)
+
+        db_mod._http_get_json(f"{WS}/get", "tok")
+        db_mod._http_post_json(f"{WS}/post", "tok", {})
+
+        assert [request.get_header("Authorization") for request in requests] == [
+            "Bearer tok",
+            "Bearer tok",
+        ]
+        assert requests[1].get_header("Content-type") == "application/json"
+
+
 class TestParseDatabricksCliVersion:
     def test_parses_standard_format(self):
         assert _parse_databricks_cli_version("Databricks CLI v0.299.2") == (0, 299, 2)