From ab2be5ffec8dff2ec6e8c6c97b36e380d54fe7fc Mon Sep 17 00:00:00 2001 From: Tien Le Date: Mon, 21 Sep 2026 16:24:55 +0000 Subject: [PATCH 1/4] Apply managed http_headers to Codex's provider block MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire the managed coding-agent config's `http_headers` into Codex, independent of the Claude change: `managed_state_overrides` emits `{tool}_http_headers`, and `_provider_block` merges them into the Databricks provider's `[http_headers]` table (current and legacy layouts). Admin headers win on a case-insensitive name collision, replacing ucode's own header (including routing headers) rather than emitting two spellings of one HTTP header. ug is the sole author of that table, so it is pruned wholesale before each merge and rewritten from render_overlay — dropping routing headers that no longer apply and any admin header removed from managed config, with no ug state. Adds a wire-level check to tests/test_e2e_user_agent.py that launches the real Codex binary and asserts an admin header arrives on the inference request. Note: shares the small managed_resolve plumbing and the `_header` test helper with the Claude PR (databricks/unity-gateway#772); whichever merges second resolves a trivial overlap. Co-authored-by: Isaac --- src/ucode/agents/codex.py | 32 ++++++++++++------ src/ucode/managed_resolve.py | 10 ++++++ tests/test_agent_codex.py | 62 +++++++++++++++++++++++++++++++++++ tests/test_e2e_user_agent.py | 45 +++++++++++++++++++++++-- tests/test_managed_resolve.py | 24 ++++++++++++++ 5 files changed, 161 insertions(+), 12 deletions(-) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 704a3adce..340fc5ee2 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -93,14 +93,10 @@ LEGACY_CODEX_BACKUP_PATH = APP_DIR / "codex-config.backup.toml" CODEX_MODEL_PROVIDER_NAME = "Databricks" LEGACY_CODEX_MODEL_PROVIDER_NAME = "ucode-databricks" -_MODEL_SERVICE_ROUTING_KEY_PATHS = [ - ["model_providers", CODEX_MODEL_PROVIDER_NAME, "http_headers", MODEL_PROVIDER_SERVICE_HEADER], - [ - "model_providers", - CODEX_MODEL_PROVIDER_NAME, - "http_headers", - MODEL_SERVICE_PARENT_SCHEMA_HEADER, - ], +# ug owns the whole provider http_headers table, so it is pruned before each merge and rewritten +# from render_overlay — dropping stale routing and admin headers that deep_merge cannot delete. +_PROVIDER_HTTP_HEADERS_KEY_PATHS = [ + ["model_providers", CODEX_MODEL_PROVIDER_NAME, "http_headers"], ] MINIMUM_CODEX_VERSION = (0, 145, 0) MINIMUM_CODEX_VERSION_TEXT = "0.145.0" @@ -180,6 +176,14 @@ def has_ucode_config() -> bool: ) +def _apply_managed_headers(http_headers: dict[str, str], managed: dict[str, str] | None) -> None: + """Merge admin ``managed`` headers into ``http_headers`` in place; admin wins case-insensitively.""" + for name, value in (managed or {}).items(): + for existing in [key for key in http_headers if key.casefold() == name.casefold()]: + del http_headers[existing] + http_headers[name] = value + + def _provider_block( workspace: str, databricks_profile: str | None, @@ -187,6 +191,7 @@ def _provider_block( provider: str | None = None, parent_schema: str | None = None, custom_oauth: CustomOAuthConfig | None = None, + managed_http_headers: dict[str, str] | None = None, ) -> dict: if custom_oauth: auth_argv = build_custom_auth_token_argv(workspace, custom_oauth) @@ -202,6 +207,7 @@ def _provider_block( http_headers[MODEL_SERVICE_PARENT_SCHEMA_HEADER] = parent_schema if smart_routing_v2.smart_routing_enabled(): http_headers[SMART_ROUTER_RECIPE_HEADER] = configured_router_name() + _apply_managed_headers(http_headers, managed_http_headers) return { "name": "Databricks AI Gateway", "base_url": base_url, @@ -226,6 +232,7 @@ def render_overlay( provider: str | None = None, parent_schema: str | None = None, custom_oauth: CustomOAuthConfig | None = None, + managed_http_headers: dict[str, str] | None = None, ) -> dict: overlay: dict = {"model_provider": CODEX_MODEL_PROVIDER_NAME} if model: @@ -238,6 +245,7 @@ def render_overlay( provider=provider, parent_schema=parent_schema, custom_oauth=custom_oauth, + managed_http_headers=managed_http_headers, ), } return overlay @@ -251,6 +259,7 @@ def render_legacy_overlay( provider: str | None = None, parent_schema: str | None = None, custom_oauth: CustomOAuthConfig | None = None, + managed_http_headers: dict[str, str] | None = None, ) -> dict: """Overlay for Codex CLI < 0.134.0, which only reads `~/.codex/config.toml`. @@ -271,6 +280,7 @@ def render_legacy_overlay( provider=provider, parent_schema=parent_schema, custom_oauth=custom_oauth, + managed_http_headers=managed_http_headers, ), }, } @@ -415,9 +425,10 @@ def write_tool_config( provider=provider, parent_schema=parent_schema, custom_oauth=state.get("custom_oauth"), + managed_http_headers=state.get("codex_http_headers"), ) doc = read_toml_safe(LEGACY_CODEX_CONFIG_PATH) - prune_key_paths(doc, _MODEL_SERVICE_ROUTING_KEY_PATHS) + prune_key_paths(doc, _PROVIDER_HTTP_HEADERS_KEY_PATHS) deep_merge_dict(doc, overlay) # deep_merge can't drop keys, so clear model preferences from an earlier run. profiles = doc.get("profiles") @@ -457,10 +468,11 @@ def write_tool_config( provider=provider, parent_schema=parent_schema, custom_oauth=state.get("custom_oauth"), + managed_http_headers=state.get("codex_http_headers"), ) def compose(base: dict, *, include_catalog: bool = True) -> dict: - prune_key_paths(base, _MODEL_SERVICE_ROUTING_KEY_PATHS) + prune_key_paths(base, _PROVIDER_HTTP_HEADERS_KEY_PATHS) deep_merge_dict(base, copy.deepcopy(overlay)) # deep_merge can't drop keys, so clear model preferences from an earlier run. if chosen_model is None and not smart_routing_v2.smart_routing_enabled(): diff --git a/src/ucode/managed_resolve.py b/src/ucode/managed_resolve.py index 78dd6001a..f1aaf3057 100644 --- a/src/ucode/managed_resolve.py +++ b/src/ucode/managed_resolve.py @@ -62,6 +62,12 @@ def _agent_model_config(managed: dict, tool: str) -> dict[str, object]: return _as_dict(_agent_entry(managed, tool).get("model_config")) +def _agent_http_headers(managed: dict, tool: str) -> dict[str, str]: + """Return the manifest's custom ``http_headers`` for ``tool`` (str->str only).""" + headers = _as_dict(_agent_entry(managed, tool).get("http_headers")) + return {k: v for k, v in headers.items() if isinstance(k, str) and isinstance(v, str)} + + def managed_otel_tracing_enabled(managed: dict, tool: str) -> bool: """Whether managed config enables OTLP trace export for ``tool``.""" return _agent_entry(managed, tool).get("otel_tracing_enabled") is True @@ -94,6 +100,10 @@ def managed_state_overrides(managed: dict, tool: str) -> dict[str, object]: default_model = _str(_agent_model_config(managed, tool).get("default_model")) if default_model: overrides[f"{tool}_default_model"] = default_model + if tool in ("claude", "codex"): + http_headers = _agent_http_headers(managed, tool) + if http_headers: + overrides[f"{tool}_http_headers"] = http_headers if tool in OTEL_TRACING_TOOLS and managed_otel_tracing_enabled(managed, tool): overrides[f"{tool}_otel_tracing"] = True return overrides diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index 008cd09e2..550123231 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -171,6 +171,29 @@ def test_parent_adds_discovery_header(self): headers = overlay["model_providers"]["Databricks"]["http_headers"] assert headers["Databricks-Model-Service-Parent-Schema"] == "main.default" + def test_managed_http_headers_added(self): + overlay = codex.render_overlay( + WS, managed_http_headers={"x-databricks-workspace": "eng-ml-inference"} + ) + headers = overlay["model_providers"]["Databricks"]["http_headers"] + assert headers["x-databricks-workspace"] == "eng-ml-inference" + assert "User-Agent" in headers # ucode's own header is retained alongside. + + def test_managed_http_headers_override_ucode_header(self, monkeypatch): + monkeypatch.setattr(codex, "ug_version", lambda: "0.1.0") + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.123.0") + overlay = codex.render_overlay( + WS, + provider="main.x.svc", + managed_http_headers={"databricks-model-provider-service": "admin.override"}, + ) + headers = overlay["model_providers"]["Databricks"]["http_headers"] + # Admin wins on a case-insensitive collision, leaving no duplicate spelling of the header. + assert headers == { + "User-Agent": "ucode/0.1.0 codex/0.123.0", + "databricks-model-provider-service": "admin.override", + } + class TestRenderOverlayUserAgent: def test_user_agent_set_on_provider(self, monkeypatch): @@ -338,6 +361,45 @@ def test_replaces_stale_routing_headers(self, tmp_path, monkeypatch): assert "Databricks-Model-Service-Parent-Schema" not in headers assert "Databricks-Model-Provider-Service" not in headers + def test_writes_admin_http_headers(self, tmp_path, monkeypatch): + config_path = tmp_path / ".codex" / "ucode.config.toml" + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) + monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", tmp_path / "backup.toml") + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.134.0") + monkeypatch.setattr(codex, "save_state", lambda state: None) + state = { + "workspace": WS, + "codex_models": [], + "codex_http_headers": {"x-databricks-workspace": "eng-ml-inference"}, + } + + codex.write_tool_config(state) + + headers = read_toml_safe(config_path)["model_providers"]["Databricks"]["http_headers"] + assert headers["x-databricks-workspace"] == "eng-ml-inference" # admin header applied + assert "User-Agent" in headers # ucode's own header kept + + def test_drops_admin_http_header_after_removal(self, tmp_path, monkeypatch): + config_path = tmp_path / ".codex" / "ucode.config.toml" + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) + monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", tmp_path / "backup.toml") + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.134.0") + monkeypatch.setattr(codex, "save_state", lambda state: None) + + codex.write_tool_config( + { + "workspace": WS, + "codex_models": [], + "codex_http_headers": {"x-databricks-workspace": "eng-ml-inference"}, + } + ) + # The admin removes the header from managed config; the next configure omits it entirely. + codex.write_tool_config({"workspace": WS, "codex_models": []}) + + headers = read_toml_safe(config_path)["model_providers"]["Databricks"]["http_headers"] + assert "x-databricks-workspace" not in headers # dropped on removal + assert "User-Agent" in headers + def test_legacy_replaces_stale_routing_headers(self, tmp_path, monkeypatch): config_dir = tmp_path / ".codex" legacy_path = config_dir / "config.toml" diff --git a/tests/test_e2e_user_agent.py b/tests/test_e2e_user_agent.py index d440d9691..03ff7fd1e 100644 --- a/tests/test_e2e_user_agent.py +++ b/tests/test_e2e_user_agent.py @@ -1,8 +1,9 @@ -"""End-to-end test that the User-Agent header ucode injects actually reaches the wire. +"""End-to-end test that headers ucode injects reach the wire — the agent's own +User-Agent, and admin-supplied managed ``http_headers``. We don't talk to a real Databricks workspace here — instead we stand up a tiny HTTP capture server on localhost, point each agent's *_BASE_URL at it, -launch the agent, and assert on the User-Agent the server saw. +launch the agent, and assert on the headers the server saw. The server returns a canned error so the agent itself fails; we don't care about the agent's exit code, only the headers that arrived before it bailed. @@ -135,6 +136,13 @@ def _assert_ua(req: _CapturedRequest, expected: str) -> None: assert ua == expected, f"User-Agent mismatch.\n got: {ua!r}\n expected: {expected!r}" +def _header(req: _CapturedRequest, name: str) -> str | None: + for key, value in req.headers.items(): + if key.casefold() == name.casefold(): + return value + return None + + def _run_until_first_request( cmd: list[str], env: dict[str, str], timeout: int = 20 ) -> subprocess.CompletedProcess | None: @@ -235,6 +243,39 @@ def test_user_agent_arrives_at_gateway(self, tmp_path, monkeypatch, capture_serv assert req is not None, _no_request_msg(capture_server, result) _assert_ua(req, _expected_ua("codex", "codex")) + def test_managed_http_header_arrives_at_gateway(self, tmp_path, monkeypatch, capture_server): + import ucode.config_io as config_io_mod + from ucode.agents import codex + + _require_binary("codex") + config_dir = tmp_path / "codex_home" / ".codex" + config_dir.mkdir(parents=True) + config_path = config_dir / "ucode.config.toml" + + monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) + monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", tmp_path / "codex.backup.toml") + + with pytest.MonkeyPatch().context() as mp: + mp.setattr("ucode.state.save_state", lambda s: None) + codex.write_tool_config( + { + "workspace": capture_server.base_url, + "codex_http_headers": {"x-databricks-workspace": "eng-ml-inference"}, + } + ) + + env = { + **os.environ, + "CODEX_HOME": str(config_dir), + "OPENAI_API_KEY": "test-key-not-real", + } + result = _run_until_first_request(codex.validate_cmd("codex"), env) + + req = capture_server.first_request_with_path_prefix("/ai-gateway/codex") + assert req is not None, _no_request_msg(capture_server, result) + assert _header(req, "x-databricks-workspace") == "eng-ml-inference" + class TestOpencodeUserAgent: def test_user_agent_arrives_at_gateway(self, tmp_path, monkeypatch, capture_server): diff --git a/tests/test_managed_resolve.py b/tests/test_managed_resolve.py index b82a026df..e72f13983 100644 --- a/tests/test_managed_resolve.py +++ b/tests/test_managed_resolve.py @@ -88,6 +88,30 @@ def test_resolve_state_layers_tracing_flag(self): assert resolve_state(managed, _state(), "codex")["codex_otel_tracing"] is True +class TestHttpHeaders: + def test_state_override_carries_manifest_headers(self): + for tool in ("claude", "codex"): + managed = { + "enabled_agents": {tool: {"http_headers": {"x-databricks-workspace": "eng-ml"}}} + } + assert managed_state_overrides(managed, tool)[f"{tool}_http_headers"] == { + "x-databricks-workspace": "eng-ml" + } + + def test_no_headers_adds_no_override(self): + assert "claude_http_headers" not in managed_state_overrides(MANAGED, "claude") + + def test_non_string_values_are_dropped(self): + managed = {"enabled_agents": {"claude": {"http_headers": {"ok": "v", "bad": 1}}}} + assert managed_state_overrides(managed, "claude")["claude_http_headers"] == {"ok": "v"} + + def test_resolve_state_layers_headers(self): + managed = {"enabled_agents": {"claude": {"http_headers": {"x-team": "aig"}}}} + assert resolve_state(managed, _state(), "claude")["claude_http_headers"] == { + "x-team": "aig" + } + + class TestClaudeModels: def test_proto_slots_map_to_families(self): # The manifest keeps proto spelling (`default_opus_model`); render_overlay reads `opus`. From 055dc5ddb7e43c2464e9ea3b9c66b03261dcc2ea Mon Sep 17 00:00:00 2001 From: Tien Le Date: Mon, 21 Sep 2026 17:12:57 +0000 Subject: [PATCH 2/4] Fold http_headers override into the existing claude/codex gate Removes the duplicate `if tool in ("claude", "codex")` in managed_state_overrides by grouping the http_headers override with static_models. No behavior change; the override is still emitted only for the agents whose writers consume it. Co-authored-by: Isaac --- src/ucode/managed_resolve.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/ucode/managed_resolve.py b/src/ucode/managed_resolve.py index f1aaf3057..d8b63fd48 100644 --- a/src/ucode/managed_resolve.py +++ b/src/ucode/managed_resolve.py @@ -97,13 +97,12 @@ def managed_state_overrides(managed: dict, tool: str) -> dict[str, object]: static_models = managed_static_models(managed, tool) if static_models: overrides[f"{tool}_static_models"] = static_models - default_model = _str(_agent_model_config(managed, tool).get("default_model")) - if default_model: - overrides[f"{tool}_default_model"] = default_model - if tool in ("claude", "codex"): http_headers = _agent_http_headers(managed, tool) if http_headers: overrides[f"{tool}_http_headers"] = http_headers + default_model = _str(_agent_model_config(managed, tool).get("default_model")) + if default_model: + overrides[f"{tool}_default_model"] = default_model if tool in OTEL_TRACING_TOOLS and managed_otel_tracing_enabled(managed, tool): overrides[f"{tool}_otel_tracing"] = True return overrides From 496574a7ecafc47bfbc7944c09e38250ca644e27 Mon Sep 17 00:00:00 2001 From: Tien Le Date: Mon, 21 Sep 2026 17:15:27 +0000 Subject: [PATCH 3/4] Ungate the http_headers override in managed_state_overrides Emit `{tool}_http_headers` for any tool rather than only claude/codex. Only those writers consume it today, so the override is inert for other agents, and dropping the gate keeps the resolver simple. Co-authored-by: Isaac --- src/ucode/managed_resolve.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/ucode/managed_resolve.py b/src/ucode/managed_resolve.py index d8b63fd48..907104226 100644 --- a/src/ucode/managed_resolve.py +++ b/src/ucode/managed_resolve.py @@ -97,12 +97,12 @@ def managed_state_overrides(managed: dict, tool: str) -> dict[str, object]: static_models = managed_static_models(managed, tool) if static_models: overrides[f"{tool}_static_models"] = static_models - http_headers = _agent_http_headers(managed, tool) - if http_headers: - overrides[f"{tool}_http_headers"] = http_headers default_model = _str(_agent_model_config(managed, tool).get("default_model")) if default_model: overrides[f"{tool}_default_model"] = default_model + http_headers = _agent_http_headers(managed, tool) + if http_headers: + overrides[f"{tool}_http_headers"] = http_headers if tool in OTEL_TRACING_TOOLS and managed_otel_tracing_enabled(managed, tool): overrides[f"{tool}_otel_tracing"] = True return overrides From d345df4a4ba1ddc0af67529cac8bfea12c28af8b Mon Sep 17 00:00:00 2001 From: Tien Le Date: Mon, 21 Sep 2026 18:31:04 +0000 Subject: [PATCH 4/4] Send managed http_headers on smart-routed Codex inference calls The smart-routing launch path rebuilds the provider overlay and passes it to the Codex app-server as `-c model_providers.Databricks=...` overrides, which replace the whole provider block (including its http_headers table). That overlay omitted the admin headers, so they were written to config.toml but stripped from the launched inference calls whenever smart routing was active. Thread `codex_http_headers` into the launch overlay in launch_codex, and add a regression test asserting the admin header appears in the app-server config args. Co-authored-by: Isaac --- src/ucode/smart_routing/v2.py | 1 + tests/test_codex_smart_routing_v2.py | 53 ++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+) diff --git a/src/ucode/smart_routing/v2.py b/src/ucode/smart_routing/v2.py index a79c51872..ba8848deb 100644 --- a/src/ucode/smart_routing/v2.py +++ b/src/ucode/smart_routing/v2.py @@ -594,6 +594,7 @@ def launch_codex( state.get("profile"), use_pat=bool(state.get("use_pat")), custom_oauth=(custom_oauth if custom_oauth_cli_enabled(custom_oauth) else None), + managed_http_headers=state.get("codex_http_headers"), ) catalog_path = custom_catalog_path() if catalog_path is not None: diff --git a/tests/test_codex_smart_routing_v2.py b/tests/test_codex_smart_routing_v2.py index f93c035e6..c9ca054f5 100644 --- a/tests/test_codex_smart_routing_v2.py +++ b/tests/test_codex_smart_routing_v2.py @@ -281,6 +281,59 @@ def start_interposer(*args, **kwargs): assert stopped == [True] assert processes[0].terminated is True + def test_managed_http_headers_reach_app_server_config(self, monkeypatch): + # Smart routing rebuilds the overlay and passes it to the app-server as `-c` overrides that + # replace the whole provider block, so the admin headers must be threaded through here too — + # otherwise they are written to config.toml but stripped from the launched inference calls. + processes = [] + monkeypatch.setenv(v2.ENABLE_SMART_ROUTING_ENV_VAR, "1") + monkeypatch.setenv("CODEX_HOME", "/user/codex-home") + monkeypatch.setattr(codex, "ug_version", lambda: "0.1.0") + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.148.0") + + class FakeProcess: + def __init__(self, argv, **kwargs): + self.argv = argv + processes.append(self) + + def wait(self, timeout=None): + return 0 + + def terminate(self): + pass + + def send_signal(self, _signal): + pass + + monkeypatch.setattr(v2.subprocess, "Popen", FakeProcess) + monkeypatch.setattr(v2, "get_databricks_token", lambda *_a, **_k: "token") + monkeypatch.setattr(v2, "_free_port", lambda: 41001) + monkeypatch.setattr(v2, "_wait_for_app_server", lambda port, timeout: True) + monkeypatch.setattr( + codex_interposer, + "start_interposer_thread", + lambda *_a, **_k: (41002, lambda: None), + ) + + with pytest.raises(SystemExit): + v2.launch_codex( + { + "workspace": WS, + "codex_models": ["system.ai.gpt-5-6-sol"], + "codex_http_headers": {"x-databricks-workspace": "eng-ml-inference"}, + }, + [], + binary="codex", + start_model="gpt-start", + render_overlay=codex.render_overlay, + ) + + provider_arg = next( + arg for arg in processes[0].argv if arg.startswith("model_providers.Databricks=") + ) + assert "x-databricks-workspace" in provider_arg + assert "eng-ml-inference" in provider_arg + def test_subagent_only_launch_runs_tui_directly(self, tmp_path, monkeypatch): monkeypatch.setenv(v2.ENABLE_SUBAGENT_ROUTING_ENV_VAR, "1") monkeypatch.setenv("CODEX_HOME", str(tmp_path))