Apply managed http_headers to Codex's provider block

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,13 +176,22 @@ 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,
     use_pat: bool = False,
     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():

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,9 @@ 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
+    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

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:

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"

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))

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):

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`.