diff --git a/README.md b/README.md index d64d972c..d176d10e 100644 --- a/README.md +++ b/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 ]` | Delete every skill downloaded from a schema (all bases, or one under ``) | | `ug skill remove --skills main.default.my-skill [--path ]` | 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 ` launches. Use `--enable-databricks-ai-tools` or `--disable-databricks-ai-tools` with `ug configure` to control the installation. diff --git a/src/ucode/agents/args.py b/src/ucode/agents/args.py index faca49df..a899e439 100644 --- a/src/ucode/agents/args.py +++ b/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: diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 572d9898..ce3498ac 100644 --- a/src/ucode/agents/codex.py +++ b/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, ) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 933df438..8ef4051c 100644 --- a/src/ucode/cli.py +++ b/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 `.`. 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, ) diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index a95a1fe6..43d6514a 100644 --- a/tests/test_agent_codex.py +++ b/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" diff --git a/tests/test_cli.py b/tests/test_cli.py index bff007c9..fac9b65c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -659,6 +659,47 @@ def test_codex_model_location_is_forwarded(self): assert mock_launch.call_args.kwargs["parent_schema"] == "main.default" assert mock_launch.call_args.args[1].args == [] + def test_codex_headers_are_forwarded(self): + with patch("ucode.cli._launch_tool") as mock_launch: + result = runner.invoke( + app, ["codex", "--header", "X-First: one", "--header", "X-Second: two:three"] + ) + + assert result.exit_code == 0, result.output + assert mock_launch.call_args.kwargs["headers"] == [ + "X-First: one", + "X-Second: two:three", + ] + + def test_headers_parse_values_with_colons_and_deduplicate_case_insensitively(self): + assert cli_mod._parse_custom_headers( + [ + "X-Test: first", + "x-test: second", + "X-Development-Route: route://development/test", + ] + ) == { + "x-test": "second", + "X-Development-Route": "route://development/test", + } + + @pytest.mark.parametrize( + ("value", "message"), + [ + ("missing-separator", "format `Name: value`"), + ("bad name: value", "format `Name: value`"), + ("X-Test: line\nbreak", "control characters"), + ("X-Test: safe\u0085Authorization: injected", "line separators"), + ("X-Test: safe\u2028Authorization: injected", "line separators"), + ("X-Test: safe\u2029Authorization: injected", "line separators"), + ("Authorization: secret", "protected header"), + ("Cookie: secret", "protected header"), + ], + ) + def test_invalid_codex_header_is_rejected(self, value, message): + with pytest.raises(RuntimeError, match=message): + cli_mod._parse_custom_headers([value]) + def test_codex_provider_and_model_location_are_mutually_exclusive(self): result = runner.invoke( app, diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index f5876013..261913ee 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import os import pytest @@ -84,6 +85,35 @@ def launch_v2(state, tool_args, **kwargs): ) ] + def test_codex_smart_routing_preserves_custom_header(self, monkeypatch): + captured = {} + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") + monkeypatch.setattr(codex, "_smart_routing_config_model", lambda state: "gpt-start") + monkeypatch.setattr(codex, "codex_managed_config_path", lambda: None) + + def launch_v2(state, tool_args, **kwargs): + captured.update(kwargs) + raise SystemExit(0) + + monkeypatch.setattr(v2, "launch_codex", launch_v2) + custom_headers = {"X-Development-Route": "route://development/test"} + + with pytest.raises(SystemExit): + codex.launch( + {"workspace": WS}, + [], + options=LaunchOptions( + launch_smart_routing=True, + custom_headers=tuple(custom_headers.items()), + ), + ) + + overlay = captured["render_overlay"](WS, "gpt-start") + headers = overlay["model_providers"]["Databricks"]["env_http_headers"] + env_name = headers["X-Development-Route"] + assert os.environ[env_name] == custom_headers["X-Development-Route"] + monkeypatch.delenv(env_name) + @pytest.mark.parametrize( "tool_args", [