diff --git a/README.md b/README.md index d64d972c..1e4e3a89 100644 --- a/README.md +++ b/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 ]` | 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. 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 ` launches. Use `--enable-databricks-ai-tools` or `--disable-databricks-ai-tools` with `ug configure` to control the installation. diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 40813892..10bc09bd 100644 --- a/src/ucode/agents/__init__.py +++ b/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, 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/claude.py b/src/ucode/agents/claude.py index db46a8ac..7fc2ff65 100644 --- a/src/ucode/agents/claude.py +++ b/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) 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..bc2d0bf2 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): @@ -2341,6 +2394,7 @@ def _launch_tool( custom_model=None, coding_agent_config_defaults=coding_agent_config_defaults, parent_schema=parent_schema, + **({"custom_headers": custom_headers} if tool == "claude" and custom_headers else {}), ) # Relayed = a Claude subscription: forward the model to Claude Code's own flag, like `-- --model X`. should_forward_relayed_model = ( @@ -2392,6 +2446,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 +2488,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 +2654,7 @@ def codex_cmd( help="Discover model services in `.`. Example: main.default", ), ] = None, + header: CustomHeaderOption = None, refresh: Annotated[ bool, typer.Option( @@ -2653,6 +2718,7 @@ def codex_cmd( workspace_url=workspace, parent_schema=model_location, custom_oauth=custom_oauth, + headers=header, ) @@ -2679,6 +2745,7 @@ def claude_cmd( help="Discover model services in `.`. Example: main.default", ), ] = None, + header: CustomHeaderOption = None, model: Annotated[ str | None, typer.Option( @@ -2764,6 +2831,7 @@ def claude_cmd( workspace_url=workspace, parent_schema=model_location, custom_oauth=custom_oauth, + headers=header, ) diff --git a/src/ucode/state.py b/src/ucode/state.py index 3d209b51..a159c925 100644 --- a/src/ucode/state.py +++ b/src/ucode/state.py @@ -20,6 +20,7 @@ STATE_PATH = APP_DIR / "state.json" STATE_VERSION = 3 +GLOBAL_STATE_KEY = "global" # Transient key holding the developer's own values for whatever a managed config layered over them. # Present only in memory: the layered values render the agent settings files, while `save_state` # restores what's under it so `state.json` keeps recording the developer's own configuration. @@ -52,6 +53,28 @@ def load_state() -> dict: return hydrate_state(ws_state) +def load_global_state() -> dict: + """Load state shared by all configured workspaces.""" + state = load_full_state().get(GLOBAL_STATE_KEY) + return dict(state) if isinstance(state, dict) else {} + + +def save_global_state(state: dict) -> None: + """Save state shared by all configured workspaces.""" + if is_dry_run(): + return + full = load_full_state() + if state: + full[GLOBAL_STATE_KEY] = state + else: + full.pop(GLOBAL_STATE_KEY, None) + try: + APP_DIR.mkdir(parents=True, exist_ok=True) + STATE_PATH.write_text(json.dumps(full, indent=2), encoding="utf-8") + except OSError as exc: + raise RuntimeError(f"Failed to write state file: {STATE_PATH}") from exc + + def save_state(state: dict) -> None: """Save workspace state back into the per-workspace structure. @@ -247,6 +270,7 @@ def clear_state() -> None: if workspace: full.get("workspaces", {}).pop(workspace, None) full["current_workspace"] = None + full.pop(GLOBAL_STATE_KEY, None) try: APP_DIR.mkdir(parents=True, exist_ok=True) STATE_PATH.write_text(json.dumps(full, indent=2), encoding="utf-8") diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 8f484cc1..2d50a917 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -801,10 +801,18 @@ def test_strips_stale_otel_tracing_when_disabled(self, monkeypatch): class TestWriteToolConfigManagedSettings: """Every normal configuration also writes Claude Code's OS-managed settings.""" - def _patch(self, monkeypatch, private_writes, managed_writes, existing_by_path=None): + def _patch( + self, + monkeypatch, + private_writes, + managed_writes, + existing_by_path=None, + global_state=None, + ): existing_by_path = existing_by_path or {} + global_state = global_state if global_state is not None else {} monkeypatch.setattr(claude, "backup_existing_file", lambda *a, **kw: True) - # Deep-copy the seeded existing content so the compose step can't mutate the fixture. + monkeypatch.setattr( claude, "read_json_safe", @@ -816,6 +824,13 @@ def _patch(self, monkeypatch, private_writes, managed_writes, existing_by_path=N lambda path, payload: private_writes.append((str(path), payload)), ) monkeypatch.setattr(claude, "save_state", lambda state: None) + monkeypatch.setattr(claude, "load_global_state", lambda: dict(global_state)) + + def save_global_state(state): + global_state.clear() + global_state.update(state) + + monkeypatch.setattr(claude, "save_global_state", save_global_state) monkeypatch.setattr(claude, "_register_web_search_mcp", lambda *a, **kw: True) monkeypatch.setattr(claude, "managed_writes_allowed", lambda: True) # Deterministic managed path, and a mocked sudo writer so NO real sudo/`/etc` write happens. @@ -974,6 +989,87 @@ def test_managed_file_merges_anthropic_custom_headers(self, monkeypatch): "x-databricks-use-coding-agent-mode: true", # Newly added by ucode. ] + def test_writes_custom_header_and_records_it_in_existing_state(self, monkeypatch): + private_writes: list = [] + managed_writes: list = [] + global_state: dict = {} + self._patch(monkeypatch, private_writes, managed_writes, global_state=global_state) + + claude.write_tool_config( + {"workspace": WS, "codex_models": []}, + "databricks-claude-sonnet-4", + custom_headers={"X-Development-Route": "route://development/test"}, + ) + + private_headers = private_writes[0][1]["env"]["ANTHROPIC_CUSTOM_HEADERS"] + managed_headers = json.loads(managed_writes[0][1])["env"]["ANTHROPIC_CUSTOM_HEADERS"] + assert "X-Development-Route: route://development/test" in private_headers + assert "X-Development-Route: route://development/test" in managed_headers + assert global_state[claude.CLAUDE_CUSTOM_HEADERS_STATE_KEY] == { + "x-development-route": "route://development/test" + } + + def test_allows_reusing_recorded_empty_header(self, monkeypatch): + settings = {"env": {"ANTHROPIC_CUSTOM_HEADERS": "X-Development-Route:"}} + monkeypatch.setattr(claude, "read_json_safe", lambda _path: settings) + + claude._reject_custom_header_collisions( + {"x-development-route": ""}, {"x-development-route": ""} + ) + + @pytest.mark.parametrize( + ("route_headers", "remaining"), + [ + ("X-Development-Route: old", None), + ("X-Development-Route: old\nx-development-route: admin", "admin"), + ], + ) + def test_removes_only_recorded_value_from_previous_launch( + self, monkeypatch, route_headers, remaining + ): + private_writes: list = [] + managed_writes: list = [] + existing = {"env": {"ANTHROPIC_CUSTOM_HEADERS": (f"X-Enterprise: keep\n{route_headers}")}} + global_state = {claude.CLAUDE_CUSTOM_HEADERS_STATE_KEY: {"x-development-route": "old"}} + self._patch( + monkeypatch, + private_writes, + managed_writes, + { + str(claude.CLAUDE_SETTINGS_PATH): existing, + str(FAKE_MANAGED_PATH): existing, + }, + global_state, + ) + state = {"workspace": WS, "codex_models": []} + claude.write_tool_config(state, "databricks-claude-sonnet-4") + + private_headers = private_writes[0][1]["env"]["ANTHROPIC_CUSTOM_HEADERS"] + managed_headers = json.loads(managed_writes[0][1])["env"]["ANTHROPIC_CUSTOM_HEADERS"] + assert "X-Enterprise: keep" in private_headers + for headers in (private_headers, managed_headers): + assert "X-Development-Route: old" not in headers + assert ("x-development-route: admin" in headers) is bool(remaining) + assert claude.CLAUDE_CUSTOM_HEADERS_STATE_KEY not in global_state + + @pytest.mark.parametrize("source", ["private", "managed"]) + def test_rejects_existing_custom_header_without_modifying_settings(self, monkeypatch, source): + private_writes: list = [] + managed_writes: list = [] + path = claude.CLAUDE_SETTINGS_PATH if source == "private" else FAKE_MANAGED_PATH + existing = {"env": {"ANTHROPIC_CUSTOM_HEADERS": "X-Development-Route: user-value"}} + self._patch(monkeypatch, private_writes, managed_writes, {str(path): existing}) + + with pytest.raises(RuntimeError, match="cannot override existing Claude Code header"): + claude.write_tool_config( + {"workspace": WS, "codex_models": []}, + "databricks-claude-sonnet-4", + custom_headers={"x-development-route": "temporary"}, + ) + + assert private_writes == [] + assert managed_writes == [] + def test_managed_file_applies_model_default_precedence(self, monkeypatch): managed_defaults = self._write_managed_model_defaults( monkeypatch, 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..010ee2e2 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -659,6 +659,60 @@ 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_claude_headers_are_forwarded(self): + with patch("ucode.cli._launch_tool") as mock_launch: + result = runner.invoke( + app, + ["claude", "--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", [ diff --git a/tests/test_state.py b/tests/test_state.py index 33dd6aba..ed478590 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -16,8 +16,10 @@ get_provider_service, hydrate_state, load_full_state, + load_global_state, load_state, mark_tool_managed, + save_global_state, save_state, set_applied_managed_update_time, set_provider_service, @@ -142,6 +144,13 @@ def test_load_state_returns_empty_when_no_workspace(self): result = load_state() assert result == {} + def test_global_state_survives_workspace_switch(self): + save_state({"workspace": FAKE_WS}) + save_global_state({"claude_custom_headers": {"x-route": "test"}}) + save_state({"workspace": "https://other.databricks.com"}) + + assert load_global_state() == {"claude_custom_headers": {"x-route": "test"}} + # --------------------------------------------------------------------------- # clear_state @@ -151,10 +160,12 @@ def test_load_state_returns_empty_when_no_workspace(self): class TestClearState: def test_clears_current_workspace(self): save_state({"workspace": FAKE_WS, "claude_models": {}}) + save_global_state({"temporary": True}) clear_state() full = load_full_state() assert full["current_workspace"] is None assert FAKE_WS not in full.get("workspaces", {}) + assert load_global_state() == {} def test_clear_when_no_state_is_noop(self): clear_state() # should not raise