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..907104226 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,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 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_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_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)) 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`.