diff --git a/AGENTS.md b/AGENTS.md index aee1fb178..8ea7fb4c7 100644 --- a/AGENTS.md +++ b/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 | diff --git a/README.md b/README.md index 2781f8686..1aecc31b5 100644 --- a/README.md +++ b/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/" 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. diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 5c5794ef5..20431ab30 100644 --- a/src/ucode/agents/codex.py +++ b/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") diff --git a/src/ucode/cli.py b/src/ucode/cli.py index b084e1e0d..cd747561e 100644 --- a/src/ucode/cli.py +++ b/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: diff --git a/src/ucode/constants.py b/src/ucode/constants.py index 3a5d7c385..6764ba84b 100644 --- a/src/ucode/constants.py +++ b/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 diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index f916243d9..0d426819e 100644 --- a/src/ucode/databricks.py +++ b/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 {}) diff --git a/src/ucode/mcp_oauth.py b/src/ucode/mcp_oauth.py index 9335f8f2f..f2cf255d5 100644 --- a/src/ucode/mcp_oauth.py +++ b/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 diff --git a/src/ucode/mcp_proxy.py b/src/ucode/mcp_proxy.py index f870cd2e7..a290568a5 100644 --- a/src/ucode/mcp_proxy.py +++ b/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, diff --git a/src/ucode/mcp_web_search.py b/src/ucode/mcp_web_search.py index bce4a9eaa..b9ee6ebcb 100644 --- a/src/ucode/mcp_web_search.py +++ b/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", diff --git a/src/ucode/smart_routing/routing.py b/src/ucode/smart_routing/routing.py index 2356fe881..344ffeffe 100644 --- a/src/ucode/smart_routing/routing.py +++ b/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", }, diff --git a/src/ucode/ug_headers.py b/src/ucode/ug_headers.py new file mode 100644 index 000000000..eaf3be1fb --- /dev/null +++ b/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()) diff --git a/tests/README.md b/tests/README.md index 42e751386..1151ae85e 100644 --- a/tests/README.md +++ b/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 diff --git a/tests/conftest.py b/tests/conftest.py index fb5901ddd..1418f7058 100644 --- a/tests/conftest.py +++ b/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 diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index b76a4dd6b..3ae409d38 100644 --- a/tests/test_agent_codex.py +++ b/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" diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 3ed2ddbe5..d8d9ec0eb 100644 --- a/tests/test_databricks.py +++ b/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) diff --git a/tests/test_ug_headers.py b/tests/test_ug_headers.py new file mode 100644 index 000000000..1d574c7cb --- /dev/null +++ b/tests/test_ug_headers.py @@ -0,0 +1,66 @@ +"""Tests for the `ug --header` / `UG_HEADERS` parser.""" + +from __future__ import annotations + +import os + +import pytest +from typer.testing import CliRunner + +from ucode.cli import app +from ucode.ug_headers import ug_headers + + +def test_ug_headers_parse_one_per_line(monkeypatch): + monkeypatch.setenv("UG_HEADERS", "x-databricks-traffic-id: testenv://liteswap/a\n\nX-B: 2 ") + assert ug_headers() == {"x-databricks-traffic-id": "testenv://liteswap/a", "X-B": "2"} + + +@pytest.mark.parametrize( + "line", + [ + "not a header", + "x-databricks-traffic-id: ", + "bad name: 1", + "X-Control: a\x01b", + "X-Unicode: 日本", + ], +) +def test_ug_headers_skip_invalid_line(monkeypatch, line): + monkeypatch.setenv("UG_HEADERS", f"{line}\nX-Ok: 1") + assert ug_headers() == {"X-Ok": "1"} + + +def test_ug_headers_keep_tab_and_latin1_values(monkeypatch): + monkeypatch.setenv("UG_HEADERS", "X-Tab: a\tb\nX-Latin: caf\xe9") + assert ug_headers() == {"X-Tab": "a\tb", "X-Latin": "caf\xe9"} + + +def test_later_line_replaces_same_name_in_any_case(monkeypatch): + monkeypatch.setenv("UG_HEADERS", "X-A: 1\nx-a: 2") + assert ug_headers() == {"x-a": "2"} + + +def test_header_flag_appends_to_ug_headers(monkeypatch): + monkeypatch.setenv("UG_HEADERS", "X-Env: 1") + result = CliRunner().invoke( + app, ["--header", "X-Flag: 2", "--header", "Authorization Bearer secret-tok", "status"] + ) + assert result.exit_code == 1 + assert "--header #2: expected `Name: value`" in result.output + # The value may be a token, so the error must not echo it. + assert "secret-tok" not in result.output + + # The rejected launch exits ug in real use; reset what it left in this shared process. + monkeypatch.setenv("UG_HEADERS", "X-Env: 1") + # `--help` on the subcommand still runs ug's own callback, which sets the variable, first. + result = CliRunner().invoke(app, ["--header", "X-Flag: 2", "status", "--help"]) + assert result.exit_code == 0, result.output + assert os.environ["UG_HEADERS"] == "X-Env: 1\nX-Flag: 2" + + +def test_bad_inherited_line_never_blocks_a_command(monkeypatch): + monkeypatch.setenv("UG_HEADERS", "Authorization Bearer secret-tok\nX-Ok: 1") + result = CliRunner().invoke(app, ["revert", "--help"]) + assert result.exit_code == 0, result.output + assert "secret-tok" not in result.output