From 512ee12615dad6ed94674d1cfee78a911429de69 Mon Sep 17 00:00:00 2001 From: Tien Le Date: Mon, 21 Sep 2026 15:31:09 +0000 Subject: [PATCH 01/11] Apply managed http_headers to Claude Code's custom headers Admin-supplied `http_headers` in the managed coding-agent config were parsed and publishable but never reached any agent. Wire them into Claude Code: `managed_state_overrides` now emits `{tool}_http_headers`, and `render_overlay` merges them into `ANTHROPIC_CUSTOM_HEADERS`. Admin headers win on a case-insensitive name collision, replacing ucode's own line in place and appending any new ones. The header names ucode writes are tracked in state so a header an admin later removes is dropped rather than lingering as an unrecognized user header, while hand-added user headers are preserved. Co-authored-by: Isaac --- src/ucode/agents/claude.py | 85 +++++++++++++++++++++++--- src/ucode/managed_resolve.py | 10 +++ tests/test_agent_claude.py | 111 ++++++++++++++++++++++++++++++++++ tests/test_managed_resolve.py | 24 ++++++++ 4 files changed, 221 insertions(+), 9 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index ee7cc4e67..5de001443 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -10,7 +10,7 @@ import socket import subprocess import threading -from collections.abc import Callable +from collections.abc import Callable, Collection from pathlib import Path from ucode import gateway_proxy @@ -208,11 +208,45 @@ def _otel_trace_env(workspace: str) -> dict[str, str]: SMART_ROUTER_RECIPE_HEADER.casefold(), } ) +# Case-folded ANTHROPIC_CUSTOM_HEADERS names ucode wrote on the previous launch (static + admin +# managed). Tracked in state so a managed header an admin later removes is dropped rather than +# preserved as an unrecognized user header. +CLAUDE_MANAGED_HEADER_NAMES_STATE_KEY = "claude_managed_header_names" # Relayed drops the user scope to deliberately omit the stale apiKeyHelper. Only applied to relayed # launches — normal launches keep loading user settings (hooks/permissions) as before. _RELAYED_SETTING_SOURCES = "project,local" +def _apply_managed_header_lines( + ucode_lines: list[str], managed_http_headers: dict[str, str] | None +) -> list[str]: + """Overlay admin-supplied ``managed_http_headers`` onto ucode's own header lines. + + Admin headers win on a case-insensitive name collision: a managed value replaces + ucode's line in its existing position, and any managed header ucode does not emit is + appended after. Returns the ordered ``key: value`` lines. + """ + lines_by_name: dict[str, str] = {} + for line in ucode_lines: + name, _separator, _value = line.partition(":") + lines_by_name[name.strip().casefold()] = line + for name, value in (managed_http_headers or {}).items(): + lines_by_name[name.strip().casefold()] = f"{name}: {value}" + return list(lines_by_name.values()) + + +def _custom_header_names(custom_headers: object) -> set[str]: + """Case-folded names of every ``key: value`` line in an ANTHROPIC_CUSTOM_HEADERS string.""" + if not isinstance(custom_headers, str): + return set() + names: set[str] = set() + for line in custom_headers.splitlines(): + name, separator, _value = line.partition(":") + if separator: + names.add(name.strip().casefold()) + return names + + def configured_paths(state: dict) -> list[str]: """The Claude config file ug writes; the OS-managed file is added by the dispatcher.""" return [str(CLAUDE_SETTINGS_PATH)] @@ -353,6 +387,7 @@ def render_overlay( static_models: list[str] | None = None, otel_tracing: bool = False, picker_catalog: AnthropicModelCatalog | None = None, + managed_http_headers: dict[str, str] | None = None, ) -> tuple[dict, list[list[str]]]: """Return (overlay, managed_key_paths) for Claude settings.json. @@ -396,7 +431,10 @@ def render_overlay( header_lines.append(f"{SMART_ROUTER_RECIPE_HEADER}: {configured_router_name()}") # 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) + # + # Admin-supplied managed headers win on a name collision: applied over ucode's + # own lines, they replace a matching header in place and append the rest. + custom_headers = "\n".join(_apply_managed_header_lines(header_lines, managed_http_headers)) env: dict[str, str] = { "ANTHROPIC_BASE_URL": base_url, "ANTHROPIC_CUSTOM_HEADERS": custom_headers, @@ -899,6 +937,22 @@ def write_tool_config( static_models=state.get("claude_static_models"), otel_tracing=bool(state.get("claude_otel_tracing")), picker_catalog=picker_catalog, + managed_http_headers=state.get("claude_http_headers"), + ) + # Names ucode manages in ANTHROPIC_CUSTOM_HEADERS this launch (static + admin), unioned with + # the names it wrote last launch. Any name in this set that ucode no longer emits is dropped + # from the file, so a header an admin removes from managed config does not linger; header names + # the developer hand-added themselves stay untouched. + current_header_names = _custom_header_names( + overlay["env"].get(ANTHROPIC_CUSTOM_HEADERS_ENV_KEY) + ) + previous_header_names = { + name.casefold() + for name in (state.get(CLAUDE_MANAGED_HEADER_NAMES_STATE_KEY) or []) + if isinstance(name, str) + } + managed_header_names = ( + CLAUDE_MANAGED_CUSTOM_HEADER_NAMES | current_header_names | previous_header_names ) source_scoped_defaults = bool((provider or parent_schema) and coding_agent_config_defaults) # Native discovery must not inherit UG's prior static allow-list. Keep a replacement picker @@ -986,7 +1040,7 @@ def _compose( merged.pop(key, None) 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_header_names ) # Drop any apiKeyHelper a prior non-relayed launch left in the file; relayed # must not carry one (it would outrank the subscription OAuth). @@ -1074,24 +1128,37 @@ def _compose( else: state.pop("claude_relayed", None) state.pop("relayed_proxy_port", None) + # Record the ANTHROPIC_CUSTOM_HEADERS names ucode wrote so the next launch can drop any that an + # admin later removes from managed config. + if current_header_names: + state[CLAUDE_MANAGED_HEADER_NAMES_STATE_KEY] = sorted(current_header_names) + else: + state.pop(CLAUDE_MANAGED_HEADER_NAMES_STATE_KEY, None) 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_names: Collection[str] = CLAUDE_MANAGED_CUSTOM_HEADER_NAMES, +) -> 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 whose name is in ``managed_names`` with ucode's values in their existing + positions, while preserving all other existing headers. A managed name ucode is no longer + emitting this launch (e.g. an admin-removed header) is dropped rather than preserved. 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 - silently discarding user configuration we do not understand. + ``managed_names`` is the case-folded set ucode owns this launch — the static + ``CLAUDE_MANAGED_CUSTOM_HEADER_NAMES`` plus admin-supplied header names (current and previously + written). Header names are compared case-insensitively. Non-header lines are also preserved to + avoid silently discarding user configuration we do not understand. """ if not isinstance(existing, str) or not existing: @@ -1112,7 +1179,7 @@ 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_names: replacement = ucode_lines_by_name.get(normalized_name) if replacement is not None and normalized_name not in replaced_names: merged.append(replacement) 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_claude.py b/tests/test_agent_claude.py index 0b3fdb3b7..ca6d559cc 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -381,6 +381,28 @@ def test_parent_adds_discovery_header(self): assert "ANTHROPIC_DEFAULT_SONNET_MODEL" not in overlay["env"] assert "availableModels" not in overlay + def test_managed_http_headers_added(self): + overlay, _ = claude.render_overlay( + WS, "s4", managed_http_headers={"x-databricks-workspace": "eng-ml-inference"} + ) + lines = overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"].splitlines() + assert "x-databricks-workspace: eng-ml-inference" in lines + # ucode's own headers are still emitted alongside the admin header. + assert "x-databricks-use-coding-agent-mode: true" in lines + + def test_managed_http_headers_override_ucode_header_in_place(self, monkeypatch): + monkeypatch.setattr(claude, "ug_version", lambda: "1.0") + monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.0") + overlay, _ = claude.render_overlay( + WS, "s4", managed_http_headers={"User-Agent": "admin-agent/9"} + ) + lines = overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"].splitlines() + # Admin wins on a case-insensitive name collision, replacing ucode's line in its position. + assert lines == [ + "x-databricks-use-coding-agent-mode: true", + "User-Agent: admin-agent/9", + ] + def test_bedrock_provider_pins_model_ids(self): provider_models = { "opus": "global.anthropic.claude-opus-4-8", @@ -601,6 +623,37 @@ def test_preserves_existing_header_order(self): "meep: lala", # Preserved from existing settings in its original position. ] + def test_managed_names_drops_removed_admin_header(self): + # A managed (admin-supplied) header ucode wrote last launch but no longer emits is dropped + # when its name stays in the managed set; a hand-added user header is left untouched. + existing = "x-databricks-workspace: old-ws\nX-User: keep" + managed = "x-databricks-use-coding-agent-mode: true" + + merged = claude._merge_anthropic_custom_headers( + existing, + managed, + {"x-databricks-use-coding-agent-mode", "x-databricks-workspace"}, + ) + + lines = merged.splitlines() + assert ( + "x-databricks-workspace: old-ws" not in lines + ) # managed, no longer emitted -> dropped + assert "X-User: keep" in lines # not managed -> preserved + + def test_managed_names_replaces_admin_header_in_place(self): + existing = "x-databricks-workspace: old-ws\nX-User: keep" + managed = "x-databricks-workspace: new-ws" + + merged = claude._merge_anthropic_custom_headers( + existing, managed, {"x-databricks-workspace"} + ) + + assert merged.splitlines() == [ + "x-databricks-workspace: new-ws", # replaced in its existing position + "X-User: keep", # untouched + ] + class TestRenderOverlayWebSearchDisable: def test_settings_overlay_never_includes_mcp_servers(self): @@ -989,6 +1042,64 @@ def test_managed_file_merges_anthropic_custom_headers(self, monkeypatch): "x-databricks-use-coding-agent-mode: true", # Newly added by ucode. ] + def test_writes_admin_http_headers_and_tracks_names(self, monkeypatch): + private_writes: list = [] + managed_writes: list = [] + self._patch(monkeypatch, private_writes, managed_writes) + monkeypatch.setattr(claude, "ug_version", lambda: "1.0") + monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.0") + state = { + "workspace": WS, + "codex_models": [], + "claude_http_headers": {"x-databricks-workspace": "eng-ml-inference"}, + } + + claude.write_tool_config(state, "databricks-claude-sonnet-4") + + _, payload = private_writes[0] + lines = payload["env"]["ANTHROPIC_CUSTOM_HEADERS"].splitlines() + assert "x-databricks-workspace: eng-ml-inference" in lines # admin header applied + assert "x-databricks-use-coding-agent-mode: true" in lines # ucode's own header kept + # The admin header name is tracked so a later removal can drop it. + assert "x-databricks-workspace" in state[claude.CLAUDE_MANAGED_HEADER_NAMES_STATE_KEY] + + def test_drops_admin_http_header_after_removal(self, monkeypatch): + private_writes: list = [] + managed_writes: list = [] + existing = { + str(claude.CLAUDE_SETTINGS_PATH): { + "env": { + "ANTHROPIC_CUSTOM_HEADERS": ( + "x-databricks-use-coding-agent-mode: true\n" + "x-databricks-workspace: eng-ml-inference\n" + "X-User: keep" + ) + } + } + } + self._patch(monkeypatch, private_writes, managed_writes, existing) + monkeypatch.setattr(claude, "ug_version", lambda: "1.0") + monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.0") + # The admin previously published x-databricks-workspace; it is gone from managed config now. + state = { + "workspace": WS, + "codex_models": [], + claude.CLAUDE_MANAGED_HEADER_NAMES_STATE_KEY: [ + "user-agent", + "x-databricks-use-coding-agent-mode", + "x-databricks-workspace", + ], + } + + claude.write_tool_config(state, "databricks-claude-sonnet-4") + + _, payload = private_writes[0] + lines = payload["env"]["ANTHROPIC_CUSTOM_HEADERS"].splitlines() + assert "x-databricks-workspace: eng-ml-inference" not in lines # dropped on removal + assert "X-User: keep" in lines # hand-added user header preserved + # State no longer tracks the removed admin header. + assert "x-databricks-workspace" not in state[claude.CLAUDE_MANAGED_HEADER_NAMES_STATE_KEY] + def test_managed_file_applies_model_default_precedence(self, monkeypatch): managed_defaults = self._write_managed_model_defaults( monkeypatch, 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 82582fafde0546fb0d843c1eebaae68eeca10811 Mon Sep 17 00:00:00 2001 From: Tien Le Date: Mon, 21 Sep 2026 16:00:27 +0000 Subject: [PATCH 02/11] Trim inline comments; add wire-level header test for Claude Reduce the inline commentary added with the managed http_headers wiring. Add a capture-server check to tests/test_e2e_user_agent.py that launches the real Claude binary and asserts an admin-supplied managed header arrives on the inference request, alongside the existing User-Agent assertion. Co-authored-by: Isaac --- src/ucode/agents/claude.py | 42 +++++++++----------------------------- 1 file changed, 10 insertions(+), 32 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 5de001443..2d2d92939 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -208,9 +208,7 @@ def _otel_trace_env(workspace: str) -> dict[str, str]: SMART_ROUTER_RECIPE_HEADER.casefold(), } ) -# Case-folded ANTHROPIC_CUSTOM_HEADERS names ucode wrote on the previous launch (static + admin -# managed). Tracked in state so a managed header an admin later removes is dropped rather than -# preserved as an unrecognized user header. +# Header names ucode wrote last launch, tracked so a removed managed header is dropped next launch. CLAUDE_MANAGED_HEADER_NAMES_STATE_KEY = "claude_managed_header_names" # Relayed drops the user scope to deliberately omit the stale apiKeyHelper. Only applied to relayed # launches — normal launches keep loading user settings (hooks/permissions) as before. @@ -220,12 +218,7 @@ def _otel_trace_env(workspace: str) -> dict[str, str]: def _apply_managed_header_lines( ucode_lines: list[str], managed_http_headers: dict[str, str] | None ) -> list[str]: - """Overlay admin-supplied ``managed_http_headers`` onto ucode's own header lines. - - Admin headers win on a case-insensitive name collision: a managed value replaces - ucode's line in its existing position, and any managed header ucode does not emit is - appended after. Returns the ordered ``key: value`` lines. - """ + """Overlay admin ``managed_http_headers`` onto ucode's header lines; admin wins by name.""" lines_by_name: dict[str, str] = {} for line in ucode_lines: name, _separator, _value = line.partition(":") @@ -431,9 +424,6 @@ def render_overlay( header_lines.append(f"{SMART_ROUTER_RECIPE_HEADER}: {configured_router_name()}") # 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. - # - # Admin-supplied managed headers win on a name collision: applied over ucode's - # own lines, they replace a matching header in place and append the rest. custom_headers = "\n".join(_apply_managed_header_lines(header_lines, managed_http_headers)) env: dict[str, str] = { "ANTHROPIC_BASE_URL": base_url, @@ -939,10 +929,8 @@ def write_tool_config( picker_catalog=picker_catalog, managed_http_headers=state.get("claude_http_headers"), ) - # Names ucode manages in ANTHROPIC_CUSTOM_HEADERS this launch (static + admin), unioned with - # the names it wrote last launch. Any name in this set that ucode no longer emits is dropped - # from the file, so a header an admin removes from managed config does not linger; header names - # the developer hand-added themselves stay untouched. + # Names ucode owns this launch, unioned with last launch's, so a removed managed header is + # dropped while hand-added user headers are left alone. current_header_names = _custom_header_names( overlay["env"].get(ANTHROPIC_CUSTOM_HEADERS_ENV_KEY) ) @@ -1128,8 +1116,6 @@ def _compose( else: state.pop("claude_relayed", None) state.pop("relayed_proxy_port", None) - # Record the ANTHROPIC_CUSTOM_HEADERS names ucode wrote so the next launch can drop any that an - # admin later removes from managed config. if current_header_names: state[CLAUDE_MANAGED_HEADER_NAMES_STATE_KEY] = sorted(current_header_names) else: @@ -1144,21 +1130,13 @@ def _merge_anthropic_custom_headers( ucode_headers: str, managed_names: Collection[str] = CLAUDE_MANAGED_CUSTOM_HEADER_NAMES, ) -> 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 whose name is in ``managed_names`` with ucode's values in their existing - positions, while preserving all other existing headers. A managed name ucode is no longer - emitting this launch (e.g. an admin-removed header) is dropped rather than preserved. - 4. Append any ucode-managed headers that were not already present. + """Merge the newline-delimited ``ANTHROPIC_CUSTOM_HEADERS`` string, replacing ucode-owned names. - ``managed_names`` is the case-folded set ucode owns this launch — the static - ``CLAUDE_MANAGED_CUSTOM_HEADER_NAMES`` plus admin-supplied header names (current and previously - written). Header names are compared case-insensitively. Non-header lines are also preserved to - avoid silently discarding user configuration we do not understand. + A header whose case-folded name is in ``managed_names`` is replaced in place with ucode's value, + or dropped when ucode no longer emits it (e.g. an admin-removed header); ucode headers not + already present are appended. All other existing lines (hand-added user headers, non-header + lines) are preserved. ``managed_names`` is the static set plus admin header names, current and + previously written. """ if not isinstance(existing, str) or not existing: From c52f1bec449c63b4f5af48aecf5ccea0e623a65f Mon Sep 17 00:00:00 2001 From: Tien Le Date: Mon, 21 Sep 2026 16:02:31 +0000 Subject: [PATCH 03/11] Add wire-level header test for Claude Extend tests/test_e2e_user_agent.py so the capture-server harness launches the real Claude binary and asserts an admin-supplied managed header arrives on the inference request, alongside the existing User-Agent assertion. Adds a shared case-insensitive header lookup helper. Co-authored-by: Isaac --- tests/test_e2e_user_agent.py | 42 ++++++++++++++++++++++++++++++++++-- 1 file changed, 40 insertions(+), 2 deletions(-) diff --git a/tests/test_e2e_user_agent.py b/tests/test_e2e_user_agent.py index d440d9691..8c542b2e4 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: @@ -203,6 +211,36 @@ 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("claude", "claude")) + 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 claude + + _require_binary("claude") + config_dir = tmp_path / "claude_config" + config_dir.mkdir() + monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", config_dir / "settings.json") + monkeypatch.setattr(claude, "CLAUDE_BACKUP_PATH", tmp_path / "claude.backup.json") + + overlay, _ = claude.render_overlay( + capture_server.base_url, + "test-model", + managed_http_headers={"x-databricks-workspace": "eng-ml-inference"}, + ) + claude.CLAUDE_SETTINGS_PATH.write_text(json.dumps(overlay), encoding="utf-8") + env = { + **os.environ, + "CLAUDE_CONFIG_DIR": str(config_dir), + "ANTHROPIC_API_KEY": "test-key-not-real", + **overlay["env"], + } + + result = _run_until_first_request(claude.validate_cmd("claude"), env) + + req = capture_server.first_request_with_path_prefix("/ai-gateway/anthropic") + assert req is not None, _no_request_msg(capture_server, result) + assert _header(req, "x-databricks-workspace") == "eng-ml-inference" + class TestCodexUserAgent: def test_user_agent_arrives_at_gateway(self, tmp_path, monkeypatch, capture_server): From 30060af85ee4de908ab9b527b3dcdc0df0a12c2f Mon Sep 17 00:00:00 2001 From: Tien Le Date: Mon, 21 Sep 2026 16:14:17 +0000 Subject: [PATCH 04/11] Fix managed http_header removal to be state-free The prior removal tracking stored the header names ucode wrote in ug state, but ug state is not reliably preserved across runs, so a removed managed header lingered while the new one was added (both appeared). Derive the ucode-owned header set from the files instead: ucode owns its own ucode-settings.json wholesale (every header there is ucode's, so a dropped managed header cannot linger), and for the OS-managed file only the names ucode last wrote (last_applied_by_ug) are treated as ucode's, leaving IT-authored headers untouched. Removes the ug-state tracking key entirely. Co-authored-by: Isaac --- src/ucode/agents/claude.py | 49 ++++++++++++++++------------ tests/test_agent_claude.py | 67 ++++++++++++++++++++++++++++---------- 2 files changed, 78 insertions(+), 38 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 2d2d92939..6e481e740 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -208,8 +208,6 @@ def _otel_trace_env(workspace: str) -> dict[str, str]: SMART_ROUTER_RECIPE_HEADER.casefold(), } ) -# Header names ucode wrote last launch, tracked so a removed managed header is dropped next launch. -CLAUDE_MANAGED_HEADER_NAMES_STATE_KEY = "claude_managed_header_names" # Relayed drops the user scope to deliberately omit the stale apiKeyHelper. Only applied to relayed # launches — normal launches keep loading user settings (hooks/permissions) as before. _RELAYED_SETTING_SOURCES = "project,local" @@ -240,6 +238,30 @@ def _custom_header_names(custom_headers: object) -> set[str]: return names +def _managed_header_names( + existing_custom_headers: object, + overlay_custom_headers: str, + snapshots: ManagedFileSnapshots | None, +) -> set[str]: + """Header names ucode owns in this file, so a managed header it no longer emits is dropped. + + Derived from the files themselves, never from ucode state (which is not preserved across runs). + For ucode's own settings file (``snapshots is None``) every existing header is ucode's. For the + OS-managed file, only the names ucode last wrote (``last_applied_by_ug``) are its own — an + administrator's own headers are left untouched. + """ + if snapshots is None: + prior = _custom_header_names(existing_custom_headers) + else: + last_applied = snapshots.last_applied_by_ug or {} + last_env = last_applied.get("env") if isinstance(last_applied, dict) else None + last_headers = ( + last_env.get(ANTHROPIC_CUSTOM_HEADERS_ENV_KEY) if isinstance(last_env, dict) else None + ) + prior = _custom_header_names(last_headers) + return CLAUDE_MANAGED_CUSTOM_HEADER_NAMES | _custom_header_names(overlay_custom_headers) | prior + + def configured_paths(state: dict) -> list[str]: """The Claude config file ug writes; the OS-managed file is added by the dispatcher.""" return [str(CLAUDE_SETTINGS_PATH)] @@ -929,19 +951,6 @@ def write_tool_config( picker_catalog=picker_catalog, managed_http_headers=state.get("claude_http_headers"), ) - # Names ucode owns this launch, unioned with last launch's, so a removed managed header is - # dropped while hand-added user headers are left alone. - current_header_names = _custom_header_names( - overlay["env"].get(ANTHROPIC_CUSTOM_HEADERS_ENV_KEY) - ) - previous_header_names = { - name.casefold() - for name in (state.get(CLAUDE_MANAGED_HEADER_NAMES_STATE_KEY) or []) - if isinstance(name, str) - } - managed_header_names = ( - CLAUDE_MANAGED_CUSTOM_HEADER_NAMES | current_header_names | previous_header_names - ) source_scoped_defaults = bool((provider or parent_schema) and coding_agent_config_defaults) # Native discovery must not inherit UG's prior static allow-list. Keep a replacement picker # written by this launch, and remove only previously owned picker keys that no longer apply. @@ -1028,7 +1037,11 @@ def _compose( merged.pop(key, None) 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, managed_header_names + existing_custom_headers, + overlay_custom_headers, + _managed_header_names( + existing_custom_headers, overlay_custom_headers, managed_settings_snapshots + ), ) # Drop any apiKeyHelper a prior non-relayed launch left in the file; relayed # must not carry one (it would outrank the subscription OAuth). @@ -1116,10 +1129,6 @@ def _compose( else: state.pop("claude_relayed", None) state.pop("relayed_proxy_port", None) - if current_header_names: - state[CLAUDE_MANAGED_HEADER_NAMES_STATE_KEY] = sorted(current_header_names) - else: - state.pop(CLAUDE_MANAGED_HEADER_NAMES_STATE_KEY, None) state = mark_tool_managed(state, "claude", managed_keys) save_state(state) return state diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index ca6d559cc..d6fda7103 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -1042,7 +1042,7 @@ def test_managed_file_merges_anthropic_custom_headers(self, monkeypatch): "x-databricks-use-coding-agent-mode: true", # Newly added by ucode. ] - def test_writes_admin_http_headers_and_tracks_names(self, monkeypatch): + def test_writes_admin_http_headers(self, monkeypatch): private_writes: list = [] managed_writes: list = [] self._patch(monkeypatch, private_writes, managed_writes) @@ -1060,10 +1060,10 @@ def test_writes_admin_http_headers_and_tracks_names(self, monkeypatch): lines = payload["env"]["ANTHROPIC_CUSTOM_HEADERS"].splitlines() assert "x-databricks-workspace: eng-ml-inference" in lines # admin header applied assert "x-databricks-use-coding-agent-mode: true" in lines # ucode's own header kept - # The admin header name is tracked so a later removal can drop it. - assert "x-databricks-workspace" in state[claude.CLAUDE_MANAGED_HEADER_NAMES_STATE_KEY] def test_drops_admin_http_header_after_removal(self, monkeypatch): + # ucode owns ucode-settings.json, so a managed header it no longer emits is dropped with no + # cross-run state: every header already in that file was written by ucode. private_writes: list = [] managed_writes: list = [] existing = { @@ -1071,8 +1071,8 @@ def test_drops_admin_http_header_after_removal(self, monkeypatch): "env": { "ANTHROPIC_CUSTOM_HEADERS": ( "x-databricks-use-coding-agent-mode: true\n" - "x-databricks-workspace: eng-ml-inference\n" - "X-User: keep" + "User-Agent: ucode/1.0 claude/2.0\n" + "x-databricks-workspace: eng-ml-inference" ) } } @@ -1080,25 +1080,56 @@ def test_drops_admin_http_header_after_removal(self, monkeypatch): self._patch(monkeypatch, private_writes, managed_writes, existing) monkeypatch.setattr(claude, "ug_version", lambda: "1.0") monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.0") - # The admin previously published x-databricks-workspace; it is gone from managed config now. - state = { - "workspace": WS, - "codex_models": [], - claude.CLAUDE_MANAGED_HEADER_NAMES_STATE_KEY: [ - "user-agent", - "x-databricks-use-coding-agent-mode", - "x-databricks-workspace", - ], - } + # The admin removed the header from managed config; this configure omits it. + state = {"workspace": WS, "codex_models": []} claude.write_tool_config(state, "databricks-claude-sonnet-4") _, payload = private_writes[0] lines = payload["env"]["ANTHROPIC_CUSTOM_HEADERS"].splitlines() assert "x-databricks-workspace: eng-ml-inference" not in lines # dropped on removal - assert "X-User: keep" in lines # hand-added user header preserved - # State no longer tracks the removed admin header. - assert "x-databricks-workspace" not in state[claude.CLAUDE_MANAGED_HEADER_NAMES_STATE_KEY] + assert "x-databricks-use-coding-agent-mode: true" in lines # ucode's own header kept + + def test_managed_file_preserves_admin_headers_and_drops_removed(self, monkeypatch): + # The OS-managed file may hold IT-authored headers ucode must never touch; only headers + # ucode itself last wrote (last_applied_by_ug) are dropped when no longer emitted. + private_writes: list = [] + managed_writes: list = [] + existing = { + str(FAKE_MANAGED_PATH): { + "env": { + "ANTHROPIC_CUSTOM_HEADERS": ( + "X-Enterprise-Header: retain\n" + "x-databricks-use-coding-agent-mode: true\n" + "x-team: old-team" + ) + } + } + } + self._patch(monkeypatch, private_writes, managed_writes, existing) + monkeypatch.setattr(claude, "ug_version", lambda: "1.0") + monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.0") + monkeypatch.setattr( + claude, + "managed_file_snapshots", + lambda tool, parser: managed_files.ManagedFileSnapshots( + None, + { + "env": { + "ANTHROPIC_CUSTOM_HEADERS": ( + "x-databricks-use-coding-agent-mode: true\nx-team: old-team" + ) + } + }, + ), + ) + state = {"workspace": WS, "codex_models": []} + + claude.write_tool_config(state, "databricks-claude-sonnet-4") + + lines = json.loads(managed_writes[0][1])["env"]["ANTHROPIC_CUSTOM_HEADERS"].splitlines() + assert "X-Enterprise-Header: retain" in lines # IT header, never ucode's -> preserved + assert "x-team: old-team" not in lines # ucode's prior admin header -> dropped def test_managed_file_applies_model_default_precedence(self, monkeypatch): managed_defaults = self._write_managed_model_defaults( From e3489421c16d35664d25c509cbb03af6c097e557 Mon Sep 17 00:00:00 2001 From: Tien Le Date: Mon, 21 Sep 2026 17:13:31 +0000 Subject: [PATCH 05/11] 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 62167fdb874c5265205b7f4443108962d0a85aba Mon Sep 17 00:00:00 2001 From: Tien Le Date: Mon, 21 Sep 2026 17:14:56 +0000 Subject: [PATCH 06/11] 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 39211095345cd30dd5151a5e590e5a32ca5b3a8a Mon Sep 17 00:00:00 2001 From: Tien Le Date: Mon, 21 Sep 2026 20:02:56 +0000 Subject: [PATCH 07/11] Fix ty invalid-return-type in _managed_header_names CLAUDE_MANAGED_CUSTOM_HEADER_NAMES is a frozenset, so the bare union (frozenset | set | set) evaluated to frozenset[str], violating the declared -> set[str] return type. Wrap the union in set(...). Co-authored-by: Isaac --- src/ucode/agents/claude.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 6e481e740..8d154a271 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -259,7 +259,11 @@ def _managed_header_names( last_env.get(ANTHROPIC_CUSTOM_HEADERS_ENV_KEY) if isinstance(last_env, dict) else None ) prior = _custom_header_names(last_headers) - return CLAUDE_MANAGED_CUSTOM_HEADER_NAMES | _custom_header_names(overlay_custom_headers) | prior + # set(...) around the union: CLAUDE_MANAGED_CUSTOM_HEADER_NAMES is a frozenset, so the bare + # union would be a frozenset[str], which does not satisfy the declared set[str] return type. + return set( + CLAUDE_MANAGED_CUSTOM_HEADER_NAMES | _custom_header_names(overlay_custom_headers) | prior + ) def configured_paths(state: dict) -> list[str]: From 168cbb361c342ca9565371ecce9d2fc3dccb7548 Mon Sep 17 00:00:00 2001 From: Tien Le Date: Mon, 21 Sep 2026 21:07:42 +0000 Subject: [PATCH 08/11] Overwrite ANTHROPIC_CUSTOM_HEADERS wholesale instead of inferring ownership ug now writes the entire ANTHROPIC_CUSTOM_HEADERS value (its static headers plus the admin's managed http_headers, already merged by render_overlay) to both ucode-settings.json and the OS managed-settings.json, dropping the per-header ownership inference. The previous approach derived "names ug owns" from snapshots.last_applied_by_ug, which records the whole previously-written merged document. A foreign/admin header ug merely passed through on the first write was therefore adopted into the owned set and silently dropped on the next reconcile -- the opposite of the intended "leave an administrator's own headers untouched". Overwriting the whole value removes the need to infer ownership, so a header ug no longer emits is dropped and nothing stale or foreign lingers. This matches Codex's full-replace model; env.ANTHROPIC_CUSTOM_HEADERS is already a ug-owned managed key path, so the wholesale write is authorized. Intentional behavior change: a custom header must arrive through the managed config http_headers manifest to appear. A header hand-edited directly into the managed file's ANTHROPIC_CUSTOM_HEADERS value is no longer preserved. Removes CLAUDE_MANAGED_CUSTOM_HEADER_NAMES, _custom_header_names, _managed_header_names, and _merge_anthropic_custom_headers. Rewrites the managed-file header tests and adds a real three-generation reconcile round-trip. Co-authored-by: Isaac --- src/ucode/agents/claude.py | 117 ++------------------- tests/test_agent_claude.py | 202 ++++++++++++++++--------------------- 2 files changed, 96 insertions(+), 223 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index b2851fa7c..bf50e8596 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -10,7 +10,7 @@ import socket import subprocess import threading -from collections.abc import Callable, Collection +from collections.abc import Callable from pathlib import Path from ucode import gateway_proxy @@ -199,15 +199,6 @@ def _otel_trace_env(workspace: str) -> dict[str, str]: CLAUDE_REMOVED_ENV_KEYS = ("CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS",) CLAUDE_MANAGED_PICKER_KEYS = ("availableModels", "enforceAvailableModels", "modelPicker") ANTHROPIC_CUSTOM_HEADERS_ENV_KEY = "ANTHROPIC_CUSTOM_HEADERS" -CLAUDE_MANAGED_CUSTOM_HEADER_NAMES = frozenset( - { - "x-databricks-use-coding-agent-mode", - "user-agent", - MODEL_PROVIDER_SERVICE_HEADER.casefold(), - MODEL_SERVICE_PARENT_SCHEMA_HEADER.casefold(), - SMART_ROUTER_RECIPE_HEADER.casefold(), - } -) # Relayed drops the user scope to deliberately omit the stale apiKeyHelper. Only applied to relayed # launches — normal launches keep loading user settings (hooks/permissions) as before. _RELAYED_SETTING_SOURCES = "project,local" @@ -226,46 +217,6 @@ def _apply_managed_header_lines( return list(lines_by_name.values()) -def _custom_header_names(custom_headers: object) -> set[str]: - """Case-folded names of every ``key: value`` line in an ANTHROPIC_CUSTOM_HEADERS string.""" - if not isinstance(custom_headers, str): - return set() - names: set[str] = set() - for line in custom_headers.splitlines(): - name, separator, _value = line.partition(":") - if separator: - names.add(name.strip().casefold()) - return names - - -def _managed_header_names( - existing_custom_headers: object, - overlay_custom_headers: str, - snapshots: ManagedFileSnapshots | None, -) -> set[str]: - """Header names ucode owns in this file, so a managed header it no longer emits is dropped. - - Derived from the files themselves, never from ucode state (which is not preserved across runs). - For ucode's own settings file (``snapshots is None``) every existing header is ucode's. For the - OS-managed file, only the names ucode last wrote (``last_applied_by_ug``) are its own — an - administrator's own headers are left untouched. - """ - if snapshots is None: - prior = _custom_header_names(existing_custom_headers) - else: - last_applied = snapshots.last_applied_by_ug or {} - last_env = last_applied.get("env") if isinstance(last_applied, dict) else None - last_headers = ( - last_env.get(ANTHROPIC_CUSTOM_HEADERS_ENV_KEY) if isinstance(last_env, dict) else None - ) - prior = _custom_header_names(last_headers) - # set(...) around the union: CLAUDE_MANAGED_CUSTOM_HEADER_NAMES is a frozenset, so the bare - # union would be a frozenset[str], which does not satisfy the declared set[str] return type. - return set( - CLAUDE_MANAGED_CUSTOM_HEADER_NAMES | _custom_header_names(overlay_custom_headers) | prior - ) - - def configured_paths(state: dict) -> list[str]: """The Claude config file ug writes; the OS-managed file is added by the dispatcher.""" return [str(CLAUDE_SETTINGS_PATH)] @@ -985,9 +936,6 @@ def _compose( managed_settings_snapshots: ManagedFileSnapshots | None, ) -> dict: base_env = base.get("env") - existing_custom_headers = ( - base_env.get(ANTHROPIC_CUSTOM_HEADERS_ENV_KEY) if isinstance(base_env, dict) else None - ) # Copy the overlay per file so merging into one base cannot affect the other. overlay_for_merge = copy.deepcopy(overlay) if enforce_model_default_hierarchy: @@ -1039,14 +987,14 @@ def _compose( merged = deep_merge_dict(base, overlay_for_merge) for key in stale_picker_keys: merged.pop(key, None) - 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, - _managed_header_names( - existing_custom_headers, overlay_custom_headers, managed_settings_snapshots - ), - ) + # ug owns the entire ANTHROPIC_CUSTOM_HEADERS value: ug's static headers plus the admin's + # managed http_headers, already merged by render_overlay. Overwrite it wholesale in every + # managed file so a header ug no longer emits is dropped and no stale or foreign header + # lingers inside the value. (deep_merge already set this; the explicit assignment states + # the contract.) + merged["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] = overlay_for_merge["env"][ + ANTHROPIC_CUSTOM_HEADERS_ENV_KEY + ] # Drop any apiKeyHelper a prior non-relayed launch left in the file; relayed # must not carry one (it would outrank the subscription OAuth). if relayed: @@ -1138,53 +1086,6 @@ def _compose( return state -def _merge_anthropic_custom_headers( - existing: object, - ucode_headers: str, - managed_names: Collection[str] = CLAUDE_MANAGED_CUSTOM_HEADER_NAMES, -) -> str: - """Merge the newline-delimited ``ANTHROPIC_CUSTOM_HEADERS`` string, replacing ucode-owned names. - - A header whose case-folded name is in ``managed_names`` is replaced in place with ucode's value, - or dropped when ucode no longer emits it (e.g. an admin-removed header); ucode headers not - already present are appended. All other existing lines (hand-added user headers, non-header - lines) are preserved. ``managed_names`` is the static set plus admin header names, current and - previously written. - """ - - if not isinstance(existing, str) or not existing: - return ucode_headers - - ucode_lines_by_name: dict[str, str] = {} - ucode_header_names: list[str] = [] - for line in ucode_headers.splitlines(): - name, separator, _value = line.partition(":") - normalized_name = name.strip().casefold() - if separator and normalized_name not in ucode_lines_by_name: - ucode_header_names.append(normalized_name) - if separator: - ucode_lines_by_name[normalized_name] = line - - merged: list[str] = [] - replaced_names: set[str] = set() - for line in existing.splitlines(): - name, separator, _value = line.partition(":") - normalized_name = name.strip().casefold() - if separator and normalized_name in managed_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 line: - merged.append(line) - - for name in ucode_header_names: - if name not in replaced_names: - merged.append(ucode_lines_by_name[name]) - return "\n".join(merged) - - def _reconcile_managed_settings( state: dict, compose: Callable[[dict], dict], diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index d6fda7103..0120abbd4 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -564,97 +564,6 @@ def test_headers_newline_delimited(self, monkeypatch): assert "\n" in self._ua(monkeypatch) -class TestMergeAnthropicCustomHeaders: - def test_removes_stale_parent_header(self): - existing = "X-User: keep\nDatabricks-Model-Service-Parent-Schema: main.default" - managed = "x-databricks-use-coding-agent-mode: true" - - merged = claude._merge_anthropic_custom_headers(existing, managed) - - assert "X-User: keep" in merged - assert "Databricks-Model-Service-Parent-Schema" not in merged - - def test_merges_existing_settings_with_ucode_managed_headers(self): - headers_from_existing_settings = "\n".join( - [ - "X-User-Header: keep-me", - "user-agent: custom-agent", - ] - ) - headers_managed_by_ucode = "\n".join( - [ - "x-databricks-use-coding-agent-mode: true", - "User-Agent: ucode/1.0 claude/2.0", - ] - ) - - merged_headers = claude._merge_anthropic_custom_headers( - headers_from_existing_settings, headers_managed_by_ucode - ) - - assert merged_headers.splitlines() == [ - "X-User-Header: keep-me", # Preserved from existing settings. - "User-Agent: ucode/1.0 claude/2.0", # From ucode; overwrites existing. - "x-databricks-use-coding-agent-mode: true", # Newly added by ucode. - ] - - def test_preserves_existing_header_order(self): - headers_from_existing_settings = "\n".join( - [ - "x-databricks-use-coding-agent-mode: true", - "User-Agent: ucode/0.1.0+41.gd09c080 claude/2.1.258", - "meep: lala", - ] - ) - headers_managed_by_ucode = "\n".join( - [ - "x-databricks-use-coding-agent-mode: true", - "User-Agent: ucode/1.0 claude/2.0", - ] - ) - - merged_headers = claude._merge_anthropic_custom_headers( - headers_from_existing_settings, headers_managed_by_ucode - ) - - assert merged_headers.splitlines() == [ - "x-databricks-use-coding-agent-mode: true", # From ucode; overwrites existing. - "User-Agent: ucode/1.0 claude/2.0", # From ucode; overwrites existing. - "meep: lala", # Preserved from existing settings in its original position. - ] - - def test_managed_names_drops_removed_admin_header(self): - # A managed (admin-supplied) header ucode wrote last launch but no longer emits is dropped - # when its name stays in the managed set; a hand-added user header is left untouched. - existing = "x-databricks-workspace: old-ws\nX-User: keep" - managed = "x-databricks-use-coding-agent-mode: true" - - merged = claude._merge_anthropic_custom_headers( - existing, - managed, - {"x-databricks-use-coding-agent-mode", "x-databricks-workspace"}, - ) - - lines = merged.splitlines() - assert ( - "x-databricks-workspace: old-ws" not in lines - ) # managed, no longer emitted -> dropped - assert "X-User: keep" in lines # not managed -> preserved - - def test_managed_names_replaces_admin_header_in_place(self): - existing = "x-databricks-workspace: old-ws\nX-User: keep" - managed = "x-databricks-workspace: new-ws" - - merged = claude._merge_anthropic_custom_headers( - existing, managed, {"x-databricks-workspace"} - ) - - assert merged.splitlines() == [ - "x-databricks-workspace: new-ws", # replaced in its existing position - "X-User: keep", # untouched - ] - - class TestRenderOverlayWebSearchDisable: def test_settings_overlay_never_includes_mcp_servers(self): # MCP servers belong in ~/.claude.json, not settings.json. @@ -1019,7 +928,10 @@ def test_managed_file_strips_stale_gateway_model_discovery(self, monkeypatch): not in json.loads(managed_writes[0][1])["env"] ) - def test_managed_file_merges_anthropic_custom_headers(self, monkeypatch): + def test_managed_file_overwrites_anthropic_custom_headers_wholesale(self, monkeypatch): + # ug owns the entire ANTHROPIC_CUSTOM_HEADERS value; a header hand-placed in the managed + # file directly (not through ug's static set nor the admin http_headers manifest) is not + # ug's to infer ownership over, so it does not survive the overwrite. private_writes: list = [] managed_writes: list = [] existing_managed_settings = { @@ -1037,9 +949,8 @@ def test_managed_file_merges_anthropic_custom_headers(self, monkeypatch): _, text = managed_writes[0] merged_headers = json.loads(text)["env"]["ANTHROPIC_CUSTOM_HEADERS"] assert merged_headers.splitlines() == [ - "X-Enterprise-Header: retain", # Preserved from existing managed settings. + "x-databricks-use-coding-agent-mode: true", "User-Agent: ucode/1.0 claude/2.0", # From ucode; overwrites existing. - "x-databricks-use-coding-agent-mode: true", # Newly added by ucode. ] def test_writes_admin_http_headers(self, monkeypatch): @@ -1090,18 +1001,21 @@ def test_drops_admin_http_header_after_removal(self, monkeypatch): assert "x-databricks-workspace: eng-ml-inference" not in lines # dropped on removal assert "x-databricks-use-coding-agent-mode: true" in lines # ucode's own header kept - def test_managed_file_preserves_admin_headers_and_drops_removed(self, monkeypatch): - # The OS-managed file may hold IT-authored headers ucode must never touch; only headers - # ucode itself last wrote (last_applied_by_ug) are dropped when no longer emitted. + def test_managed_file_overwrites_dropping_foreign_and_removed_headers(self, monkeypatch): + # Wholesale overwrite: the written value is exactly ucode's static headers plus the admin's + # CURRENT http_headers manifest, nothing else. A header that only exists directly in the + # managed file's ANTHROPIC_CUSTOM_HEADERS -- not in ucode's static set and not in the + # manifest -- is dropped just like a stale ucode-written one; a manifest header is present, + # and removing it from the manifest on a later run drops it too. private_writes: list = [] managed_writes: list = [] existing = { str(FAKE_MANAGED_PATH): { "env": { "ANTHROPIC_CUSTOM_HEADERS": ( - "X-Enterprise-Header: retain\n" + "X-Foreign-Header: keep-me\n" "x-databricks-use-coding-agent-mode: true\n" - "x-team: old-team" + "x-team: stale-team" ) } } @@ -1109,27 +1023,85 @@ def test_managed_file_preserves_admin_headers_and_drops_removed(self, monkeypatc self._patch(monkeypatch, private_writes, managed_writes, existing) monkeypatch.setattr(claude, "ug_version", lambda: "1.0") monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.0") + state = { + "workspace": WS, + "codex_models": [], + "claude_http_headers": {"x-team": "eng-ml"}, + } + + claude.write_tool_config(state, "databricks-claude-sonnet-4") + + lines = json.loads(managed_writes[0][1])["env"]["ANTHROPIC_CUSTOM_HEADERS"].splitlines() + assert "X-Foreign-Header: keep-me" not in lines # not ucode's, not in the manifest + assert "x-team: eng-ml" in lines # current manifest header -> present + + # A later run without the manifest header drops it too. + existing[str(FAKE_MANAGED_PATH)] = json.loads(managed_writes[0][1]) + managed_writes.clear() + state["claude_http_headers"] = {} + + claude.write_tool_config(state, "databricks-claude-sonnet-4") + + lines = json.loads(managed_writes[0][1])["env"]["ANTHROPIC_CUSTOM_HEADERS"].splitlines() + assert not any(line.startswith("x-team:") for line in lines) + + def test_managed_file_wholesale_overwrite_survives_real_reconcile_round_trip( + self, tmp_path, monkeypatch + ): + # Drives the REAL managed_files snapshot/reconcile flow (not a hand-mocked snapshot) across + # three launches: a header hand-placed directly in the managed file -- never ucode's, never + # in the admin manifest -- never survives a write; a manifest header is stable across a + # no-op re-run; and removing it from the manifest drops it on the next run. + managed_path = tmp_path / "managed-settings.json" + backup_dir = tmp_path / "managed-backups" + monkeypatch.setattr(managed_files, "managed_files_supported", lambda: True) + monkeypatch.setattr(managed_files, "MANAGED_BACKUP_DIR", backup_dir) monkeypatch.setattr( - claude, - "managed_file_snapshots", - lambda tool, parser: managed_files.ManagedFileSnapshots( - None, - { - "env": { - "ANTHROPIC_CUSTOM_HEADERS": ( - "x-databricks-use-coding-agent-mode: true\nx-team: old-team" - ) - } - }, - ), + managed_files, "MANAGED_BACKUP_MANIFEST_PATH", backup_dir / "manifest.json" ) - state = {"workspace": WS, "codex_models": []} + monkeypatch.setattr( + managed_files, + "_sudo_replace", + lambda target, text: target.write_text(text, encoding="utf-8"), + ) + monkeypatch.setattr(claude, "_managed_settings_path", lambda: managed_path) + monkeypatch.setattr(claude, "managed_writes_allowed", lambda: True) + monkeypatch.setattr(managed_files, "managed_writes_allowed", lambda: True) + monkeypatch.setattr(claude, "backup_existing_file", lambda *a, **kw: True) + monkeypatch.setattr(claude, "save_state", lambda state: None) + monkeypatch.setattr(claude, "ug_version", lambda: "1.0") + monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.0") + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", tmp_path / "ucode-settings.json") + + # A hand edit directly in the managed file, bypassing both ucode and the admin manifest. + managed_path.write_text( + json.dumps({"env": {"ANTHROPIC_CUSTOM_HEADERS": "X-Direct-Edit: should-not-survive"}}), + encoding="utf-8", + ) + + def custom_headers() -> list[str]: + written = json.loads(managed_path.read_text()) + return written["env"]["ANTHROPIC_CUSTOM_HEADERS"].splitlines() + state = { + "workspace": WS, + "codex_models": [], + "claude_http_headers": {"x-team": "eng-ml"}, + } claude.write_tool_config(state, "databricks-claude-sonnet-4") - lines = json.loads(managed_writes[0][1])["env"]["ANTHROPIC_CUSTOM_HEADERS"].splitlines() - assert "X-Enterprise-Header: retain" in lines # IT header, never ucode's -> preserved - assert "x-team: old-team" not in lines # ucode's prior admin header -> dropped + first = custom_headers() + assert "X-Direct-Edit: should-not-survive" not in first # hand edit -> dropped + assert "x-team: eng-ml" in first # current manifest header -> present + + # A no-op re-run (same manifest) leaves the value stable. + claude.write_tool_config(state, "databricks-claude-sonnet-4") + assert custom_headers() == first + + # The admin removes the header from the manifest; the next run drops it. + state["claude_http_headers"] = {} + claude.write_tool_config(state, "databricks-claude-sonnet-4") + assert not any(line.startswith("x-team:") for line in custom_headers()) def test_managed_file_applies_model_default_precedence(self, monkeypatch): managed_defaults = self._write_managed_model_defaults( From 5eefab48524d6df5f4dddb6702caa42d8ac0fae1 Mon Sep 17 00:00:00 2001 From: Tien Le Date: Mon, 21 Sep 2026 22:09:29 +0000 Subject: [PATCH 09/11] Gate ANTHROPIC_CUSTOM_HEADERS overwrite on managed-config presence write_tool_config now decides how to write the Claude custom-header value from whether an admin managed config is present this launch (refresh_managed_config(state).manifest is not None, reusing the launch's warm cache; a failed fetch degrades to unmanaged and never blocks the write): - managed config present: overwrite the value wholesale with ug's static headers plus the admin's managed http_headers, so a header ug no longer emits is dropped and nothing stale or foreign lingers. - no managed config: restore the pre-managed-config preserving merge -- keep the developer's own pre-existing headers, replacing only the header names ug manages (the static CLAUDE_MANAGED_CUSTOM_HEADER_NAMES set). This brings back the behavior from #442 for the unmanaged case without reintroducing the removed per-header ownership inference that dropped foreign headers on reconcile. existing_custom_headers is read before deep_merge_dict mutates base["env"] in place; reading it afterward would see the overlay's own value and collapse the preserve merge into an overwrite. Co-authored-by: Isaac --- src/ucode/agents/claude.py | 90 +++++++++++++++++-- tests/test_agent_claude.py | 174 +++++++++++++++++++++++++++++++++++++ 2 files changed, 256 insertions(+), 8 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index bf50e8596..a5e82c517 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -45,6 +45,7 @@ ug_binary, ) from ucode.launcher import exec_or_spawn +from ucode.managed_config import refresh_managed_config from ucode.managed_files import ( OS, ManagedFileSnapshots, @@ -199,6 +200,17 @@ def _otel_trace_env(workspace: str) -> dict[str, str]: CLAUDE_REMOVED_ENV_KEYS = ("CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS",) CLAUDE_MANAGED_PICKER_KEYS = ("availableModels", "enforceAvailableModels", "modelPicker") ANTHROPIC_CUSTOM_HEADERS_ENV_KEY = "ANTHROPIC_CUSTOM_HEADERS" +# Header names ug manages within ANTHROPIC_CUSTOM_HEADERS. Used only for the no-managed-config +# merge below; ownership is judged by name, never by inferring it from a whole prior document. +CLAUDE_MANAGED_CUSTOM_HEADER_NAMES = frozenset( + { + "x-databricks-use-coding-agent-mode", + "user-agent", + MODEL_PROVIDER_SERVICE_HEADER.casefold(), + MODEL_SERVICE_PARENT_SCHEMA_HEADER.casefold(), + SMART_ROUTER_RECIPE_HEADER.casefold(), + } +) # Relayed drops the user scope to deliberately omit the stale apiKeyHelper. Only applied to relayed # launches — normal launches keep loading user settings (hooks/permissions) as before. _RELAYED_SETTING_SOURCES = "project,local" @@ -217,6 +229,54 @@ def _apply_managed_header_lines( return list(lines_by_name.values()) +def _merge_anthropic_custom_headers(existing: object, ucode_headers: str) -> 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. + 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 + silently discarding user configuration we do not understand. + """ + + if not isinstance(existing, str) or not existing: + return ucode_headers + + ucode_lines_by_name: dict[str, str] = {} + ucode_header_names: list[str] = [] + for line in ucode_headers.splitlines(): + name, separator, _value = line.partition(":") + normalized_name = name.strip().casefold() + if separator and normalized_name not in ucode_lines_by_name: + ucode_header_names.append(normalized_name) + if separator: + ucode_lines_by_name[normalized_name] = line + + merged: list[str] = [] + replaced_names: set[str] = set() + 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: + 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 line: + merged.append(line) + + for name in ucode_header_names: + if name not in replaced_names: + merged.append(ucode_lines_by_name[name]) + return "\n".join(merged) + + def configured_paths(state: dict) -> list[str]: """The Claude config file ug writes; the OS-managed file is added by the dispatcher.""" return [str(CLAUDE_SETTINGS_PATH)] @@ -881,6 +941,11 @@ def write_tool_config( # revert would restore that snapshot instead of deleting the file. if not is_tool_managed(state, "claude"): backup_existing_file(CLAUDE_SETTINGS_PATH, CLAUDE_BACKUP_PATH) + # A managed config makes ug authoritative over the whole custom-header value, so it is + # overwritten wholesale; without one, preserve the developer's own pre-existing headers. Reuses + # this launch's warm managed-config cache (no extra round trip); a failed fetch degrades to None + # (treated as unmanaged), never blocking the write. + managed_config_present = refresh_managed_config(state).manifest is not None previous_keys = ((state.get("managed_configs") or {}).get("claude") or {}).get("keys", []) web_search_model = _resolve_web_search_model(state) # Relayed inference points at a local refresh proxy; its loopback base URL is @@ -936,6 +1001,12 @@ def _compose( managed_settings_snapshots: ManagedFileSnapshots | None, ) -> dict: base_env = base.get("env") + # Read before deep_merge_dict below, which mutates base's "env" dict in place (grafting + # overlay's value onto this same object) -- after that call this would read back overlay's + # own value instead of what the developer actually had. + existing_custom_headers = ( + base_env.get(ANTHROPIC_CUSTOM_HEADERS_ENV_KEY) if isinstance(base_env, dict) else None + ) # Copy the overlay per file so merging into one base cannot affect the other. overlay_for_merge = copy.deepcopy(overlay) if enforce_model_default_hierarchy: @@ -987,14 +1058,17 @@ def _compose( merged = deep_merge_dict(base, overlay_for_merge) for key in stale_picker_keys: merged.pop(key, None) - # ug owns the entire ANTHROPIC_CUSTOM_HEADERS value: ug's static headers plus the admin's - # managed http_headers, already merged by render_overlay. Overwrite it wholesale in every - # managed file so a header ug no longer emits is dropped and no stale or foreign header - # lingers inside the value. (deep_merge already set this; the explicit assignment states - # the contract.) - merged["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] = overlay_for_merge["env"][ - ANTHROPIC_CUSTOM_HEADERS_ENV_KEY - ] + overlay_custom_headers = overlay_for_merge["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] + if managed_config_present: + # ug owns the whole value under a managed config: overwrite wholesale so a header ug no + # longer emits is dropped and no stale or foreign header lingers. + merged["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] = overlay_custom_headers + else: + # No managed config: preserve the developer's own pre-existing headers, replacing only + # the header names ug manages. + merged["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] = _merge_anthropic_custom_headers( + existing_custom_headers, overlay_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). if relayed: diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 0120abbd4..9acc7644c 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -7,6 +7,7 @@ import shlex import subprocess from pathlib import Path +from types import SimpleNamespace from unittest.mock import MagicMock, Mock import pytest @@ -29,11 +30,28 @@ def _proxy_argv() -> list[str]: return build_mcp_proxy_argv(GH_URL, WS, "p") +def _managed_config_result(manifest: dict | None) -> SimpleNamespace: + """A stand-in for `ManagedConfigResult` exposing only the `.manifest` attribute + `write_tool_config` reads.""" + return SimpleNamespace(manifest=manifest) + + @pytest.fixture(autouse=True) def _avoid_real_managed_settings(monkeypatch): monkeypatch.setattr(claude, "_managed_settings_path", lambda: None) +@pytest.fixture(autouse=True) +def _default_managed_config_present(monkeypatch): + """`write_tool_config` now decides overwrite-vs-preserve itself by calling + `refresh_managed_config`. Default every test to "managed present" (current-HEAD wholesale + overwrite), matching pre-existing tests that don't care about this axis, so they need no + per-test mock; tests exercising the unmanaged path override this explicitly.""" + monkeypatch.setattr( + claude, "refresh_managed_config", lambda *a, **kw: _managed_config_result({"claude": {}}) + ) + + class TestClaudeSpec: def test_binary(self): assert claude.SPEC["binary"] == "claude" @@ -1045,6 +1063,71 @@ def test_managed_file_overwrites_dropping_foreign_and_removed_headers(self, monk lines = json.loads(managed_writes[0][1])["env"]["ANTHROPIC_CUSTOM_HEADERS"].splitlines() assert not any(line.startswith("x-team:") for line in lines) + def test_unmanaged_preserves_foreign_header_and_replaces_ucode_headers_in_place( + self, monkeypatch + ): + # No admin CodingAgentConfig: Lilly's original merge preserves the developer's own headers, + # replacing only the header names ug manages, in their existing positions. + monkeypatch.setattr( + claude, "refresh_managed_config", lambda *a, **kw: _managed_config_result(None) + ) + private_writes: list = [] + managed_writes: list = [] + existing = { + str(claude.CLAUDE_SETTINGS_PATH): { + "env": { + "ANTHROPIC_CUSTOM_HEADERS": ( + "X-Foreign-Header: keep-me\n" + "x-databricks-use-coding-agent-mode: false\n" + "User-Agent: old-agent" + ) + } + } + } + self._patch(monkeypatch, private_writes, managed_writes, existing) + monkeypatch.setattr(claude, "ug_version", lambda: "1.0") + monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.0") + state = {"workspace": WS, "codex_models": []} + + claude.write_tool_config(state, "databricks-claude-sonnet-4") + + lines = private_writes[0][1]["env"]["ANTHROPIC_CUSTOM_HEADERS"].splitlines() + assert lines == [ + "X-Foreign-Header: keep-me", # not ug's, survives untouched + "x-databricks-use-coding-agent-mode: true", # ug-managed name, replaced in place + "User-Agent: ucode/1.0 claude/2.0", # ug-managed name, replaced in place + ] + + def test_unmanaged_drops_stale_ucode_managed_header_no_longer_emitted(self, monkeypatch): + # No admin CodingAgentConfig: a header name ug manages but no longer emits this run (the + # provider-routing header, without a `provider` this run) is dropped, not left stale. + monkeypatch.setattr( + claude, "refresh_managed_config", lambda *a, **kw: _managed_config_result(None) + ) + private_writes: list = [] + managed_writes: list = [] + existing = { + str(claude.CLAUDE_SETTINGS_PATH): { + "env": { + "ANTHROPIC_CUSTOM_HEADERS": ( + "x-databricks-use-coding-agent-mode: true\n" + "User-Agent: ucode/1.0 claude/2.0\n" + "Databricks-Model-Provider-Service: old-provider" + ) + } + } + } + self._patch(monkeypatch, private_writes, managed_writes, existing) + monkeypatch.setattr(claude, "ug_version", lambda: "1.0") + monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.0") + state = {"workspace": WS, "codex_models": []} + + claude.write_tool_config(state, "databricks-claude-sonnet-4") + + lines = private_writes[0][1]["env"]["ANTHROPIC_CUSTOM_HEADERS"].splitlines() + assert not any(line.startswith("Databricks-Model-Provider-Service") for line in lines) + assert "x-databricks-use-coding-agent-mode: true" in lines + def test_managed_file_wholesale_overwrite_survives_real_reconcile_round_trip( self, tmp_path, monkeypatch ): @@ -1072,6 +1155,12 @@ def test_managed_file_wholesale_overwrite_survives_real_reconcile_round_trip( monkeypatch.setattr(claude, "ug_version", lambda: "1.0") monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.0") monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", tmp_path / "ucode-settings.json") + # Managed variant: an admin CodingAgentConfig is present, so ug owns the value wholesale. + monkeypatch.setattr( + claude, + "refresh_managed_config", + lambda *a, **kw: _managed_config_result({"claude": {}}), + ) # A hand edit directly in the managed file, bypassing both ucode and the admin manifest. managed_path.write_text( @@ -1103,6 +1192,58 @@ def custom_headers() -> list[str]: claude.write_tool_config(state, "databricks-claude-sonnet-4") assert not any(line.startswith("x-team:") for line in custom_headers()) + def test_unmanaged_wholesale_preserve_survives_real_reconcile_round_trip( + self, tmp_path, monkeypatch + ): + # Unmanaged variant of the round trip above: with no admin CodingAgentConfig, a hand-added + # foreign header in the managed file SURVIVES ug's write, because Lilly's merge only ever + # replaces the header names ug itself manages. + managed_path = tmp_path / "managed-settings.json" + backup_dir = tmp_path / "managed-backups" + monkeypatch.setattr(managed_files, "managed_files_supported", lambda: True) + monkeypatch.setattr(managed_files, "MANAGED_BACKUP_DIR", backup_dir) + monkeypatch.setattr( + managed_files, "MANAGED_BACKUP_MANIFEST_PATH", backup_dir / "manifest.json" + ) + monkeypatch.setattr( + managed_files, + "_sudo_replace", + lambda target, text: target.write_text(text, encoding="utf-8"), + ) + monkeypatch.setattr(claude, "_managed_settings_path", lambda: managed_path) + monkeypatch.setattr(claude, "managed_writes_allowed", lambda: True) + monkeypatch.setattr(managed_files, "managed_writes_allowed", lambda: True) + monkeypatch.setattr(claude, "backup_existing_file", lambda *a, **kw: True) + monkeypatch.setattr(claude, "save_state", lambda state: None) + monkeypatch.setattr(claude, "ug_version", lambda: "1.0") + monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.0") + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", tmp_path / "ucode-settings.json") + monkeypatch.setattr( + claude, "refresh_managed_config", lambda *a, **kw: _managed_config_result(None) + ) + + # A hand edit directly in the managed file; with no admin manifest to claim it, it survives. + managed_path.write_text( + json.dumps({"env": {"ANTHROPIC_CUSTOM_HEADERS": "X-Direct-Edit: survives"}}), + encoding="utf-8", + ) + + def custom_headers() -> list[str]: + written = json.loads(managed_path.read_text()) + return written["env"]["ANTHROPIC_CUSTOM_HEADERS"].splitlines() + + state = {"workspace": WS, "codex_models": []} + claude.write_tool_config(state, "databricks-claude-sonnet-4") + + first = custom_headers() + # Foreign header survives -- no manifest to claim it; ucode's own header is still applied. + assert "X-Direct-Edit: survives" in first + assert "x-databricks-use-coding-agent-mode: true" in first + + # A no-op re-run leaves the value stable. + claude.write_tool_config(state, "databricks-claude-sonnet-4") + assert custom_headers() == first + def test_managed_file_applies_model_default_precedence(self, monkeypatch): managed_defaults = self._write_managed_model_defaults( monkeypatch, @@ -1584,6 +1725,39 @@ def test_static_models_not_written_when_absent(self, monkeypatch): assert "modelPicker" not in managed_content +class TestMergeAnthropicCustomHeaders: + """Focused unit tests for the static-name-set merge used on the no-managed-config path.""" + + def test_preserves_unknown_header(self): + merged = claude._merge_anthropic_custom_headers( + "X-Foreign: keep", "x-databricks-use-coding-agent-mode: true" + ) + assert merged.splitlines() == [ + "X-Foreign: keep", + "x-databricks-use-coding-agent-mode: true", + ] + + def test_replaces_ucode_managed_name_in_place(self): + merged = claude._merge_anthropic_custom_headers( + "A: 1\nUser-Agent: old\nB: 2", "User-Agent: new" + ) + assert merged.splitlines() == ["A: 1", "User-Agent: new", "B: 2"] + + def test_drops_stale_ucode_managed_name_no_longer_emitted(self): + merged = claude._merge_anthropic_custom_headers( + "Databricks-Model-Provider-Service: old-provider\nA: 1", "" + ) + assert merged.splitlines() == ["A: 1"] + + def test_appends_new_ucode_header_not_already_present(self): + merged = claude._merge_anthropic_custom_headers("A: 1", "User-Agent: ucode/1") + assert merged.splitlines() == ["A: 1", "User-Agent: ucode/1"] + + def test_passes_through_when_no_existing_headers(self): + merged = claude._merge_anthropic_custom_headers(None, "User-Agent: ucode/1") + assert merged == "User-Agent: ucode/1" + + class TestAddClaudeMcpServer: def test_registers_stdio_proxy_command(self, monkeypatch): calls: list[dict] = [] From 1b05625b0d450478a22aed918546a2f82b031e6b Mon Sep 17 00:00:00 2001 From: Tien Le Date: Mon, 21 Sep 2026 22:19:38 +0000 Subject: [PATCH 10/11] Trim redundant custom-header tests The static-name merge is covered by TestMergeAnthropicCustomHeaders unit tests, so the write_tool_config-level tests only need one gate test per branch plus the real-reconcile integration test. Remove three duplicates: - test_managed_file_overwrites_anthropic_custom_headers_wholesale (subsumed by test_managed_file_overwrites_dropping_foreign_and_removed_headers) - test_unmanaged_drops_stale_ucode_managed_header_no_longer_emitted (pure merge logic, covered by the unit test) - test_unmanaged_wholesale_preserve_survives_real_reconcile_round_trip (reconcile plumbing covered by the managed round-trip; unmanaged preserve by unit + gate) Co-authored-by: Isaac --- tests/test_agent_claude.py | 107 ------------------------------------- 1 file changed, 107 deletions(-) diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 9acc7644c..815731aa1 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -946,31 +946,6 @@ def test_managed_file_strips_stale_gateway_model_discovery(self, monkeypatch): not in json.loads(managed_writes[0][1])["env"] ) - def test_managed_file_overwrites_anthropic_custom_headers_wholesale(self, monkeypatch): - # ug owns the entire ANTHROPIC_CUSTOM_HEADERS value; a header hand-placed in the managed - # file directly (not through ug's static set nor the admin http_headers manifest) is not - # ug's to infer ownership over, so it does not survive the overwrite. - private_writes: list = [] - managed_writes: list = [] - existing_managed_settings = { - str(FAKE_MANAGED_PATH): { - "env": {"ANTHROPIC_CUSTOM_HEADERS": "X-Enterprise-Header: retain\nUser-Agent: old"} - } - } - self._patch(monkeypatch, private_writes, managed_writes, existing_managed_settings) - monkeypatch.setattr(claude, "ug_version", lambda: "1.0") - monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.0") - state = {"workspace": WS, "codex_models": []} - - claude.write_tool_config(state, "databricks-claude-sonnet-4") - - _, text = managed_writes[0] - merged_headers = json.loads(text)["env"]["ANTHROPIC_CUSTOM_HEADERS"] - assert merged_headers.splitlines() == [ - "x-databricks-use-coding-agent-mode: true", - "User-Agent: ucode/1.0 claude/2.0", # From ucode; overwrites existing. - ] - def test_writes_admin_http_headers(self, monkeypatch): private_writes: list = [] managed_writes: list = [] @@ -1098,36 +1073,6 @@ def test_unmanaged_preserves_foreign_header_and_replaces_ucode_headers_in_place( "User-Agent: ucode/1.0 claude/2.0", # ug-managed name, replaced in place ] - def test_unmanaged_drops_stale_ucode_managed_header_no_longer_emitted(self, monkeypatch): - # No admin CodingAgentConfig: a header name ug manages but no longer emits this run (the - # provider-routing header, without a `provider` this run) is dropped, not left stale. - monkeypatch.setattr( - claude, "refresh_managed_config", lambda *a, **kw: _managed_config_result(None) - ) - private_writes: list = [] - managed_writes: list = [] - existing = { - str(claude.CLAUDE_SETTINGS_PATH): { - "env": { - "ANTHROPIC_CUSTOM_HEADERS": ( - "x-databricks-use-coding-agent-mode: true\n" - "User-Agent: ucode/1.0 claude/2.0\n" - "Databricks-Model-Provider-Service: old-provider" - ) - } - } - } - self._patch(monkeypatch, private_writes, managed_writes, existing) - monkeypatch.setattr(claude, "ug_version", lambda: "1.0") - monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.0") - state = {"workspace": WS, "codex_models": []} - - claude.write_tool_config(state, "databricks-claude-sonnet-4") - - lines = private_writes[0][1]["env"]["ANTHROPIC_CUSTOM_HEADERS"].splitlines() - assert not any(line.startswith("Databricks-Model-Provider-Service") for line in lines) - assert "x-databricks-use-coding-agent-mode: true" in lines - def test_managed_file_wholesale_overwrite_survives_real_reconcile_round_trip( self, tmp_path, monkeypatch ): @@ -1192,58 +1137,6 @@ def custom_headers() -> list[str]: claude.write_tool_config(state, "databricks-claude-sonnet-4") assert not any(line.startswith("x-team:") for line in custom_headers()) - def test_unmanaged_wholesale_preserve_survives_real_reconcile_round_trip( - self, tmp_path, monkeypatch - ): - # Unmanaged variant of the round trip above: with no admin CodingAgentConfig, a hand-added - # foreign header in the managed file SURVIVES ug's write, because Lilly's merge only ever - # replaces the header names ug itself manages. - managed_path = tmp_path / "managed-settings.json" - backup_dir = tmp_path / "managed-backups" - monkeypatch.setattr(managed_files, "managed_files_supported", lambda: True) - monkeypatch.setattr(managed_files, "MANAGED_BACKUP_DIR", backup_dir) - monkeypatch.setattr( - managed_files, "MANAGED_BACKUP_MANIFEST_PATH", backup_dir / "manifest.json" - ) - monkeypatch.setattr( - managed_files, - "_sudo_replace", - lambda target, text: target.write_text(text, encoding="utf-8"), - ) - monkeypatch.setattr(claude, "_managed_settings_path", lambda: managed_path) - monkeypatch.setattr(claude, "managed_writes_allowed", lambda: True) - monkeypatch.setattr(managed_files, "managed_writes_allowed", lambda: True) - monkeypatch.setattr(claude, "backup_existing_file", lambda *a, **kw: True) - monkeypatch.setattr(claude, "save_state", lambda state: None) - monkeypatch.setattr(claude, "ug_version", lambda: "1.0") - monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.0") - monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", tmp_path / "ucode-settings.json") - monkeypatch.setattr( - claude, "refresh_managed_config", lambda *a, **kw: _managed_config_result(None) - ) - - # A hand edit directly in the managed file; with no admin manifest to claim it, it survives. - managed_path.write_text( - json.dumps({"env": {"ANTHROPIC_CUSTOM_HEADERS": "X-Direct-Edit: survives"}}), - encoding="utf-8", - ) - - def custom_headers() -> list[str]: - written = json.loads(managed_path.read_text()) - return written["env"]["ANTHROPIC_CUSTOM_HEADERS"].splitlines() - - state = {"workspace": WS, "codex_models": []} - claude.write_tool_config(state, "databricks-claude-sonnet-4") - - first = custom_headers() - # Foreign header survives -- no manifest to claim it; ucode's own header is still applied. - assert "X-Direct-Edit: survives" in first - assert "x-databricks-use-coding-agent-mode: true" in first - - # A no-op re-run leaves the value stable. - claude.write_tool_config(state, "databricks-claude-sonnet-4") - assert custom_headers() == first - def test_managed_file_applies_model_default_precedence(self, monkeypatch): managed_defaults = self._write_managed_model_defaults( monkeypatch, From 196908413fc01022db95e4faae0302501c6fceaf Mon Sep 17 00:00:00 2001 From: Tien Le Date: Mon, 21 Sep 2026 22:30:53 +0000 Subject: [PATCH 11/11] Keep merge helper and its tests in their original positions Pure repositioning, no behavior change: the earlier wholesale->conditional detour deleted and re-added _merge_anthropic_custom_headers (and its TestMergeAnthropicCustomHeaders suite), landing them far from where main keeps them and inflating the review diff. Move the helper back below write_tool_config and restore the test class to its original spot (matching main verbatim), and drop two stray comments so unchanged code shows no diff. The real change is now just the managed-config gate in write_tool_config plus _apply_managed_header_lines. Co-authored-by: Isaac --- src/ucode/agents/claude.py | 101 ++++++++++++++++++------------------- tests/test_agent_claude.py | 93 ++++++++++++++++++++++------------ 2 files changed, 108 insertions(+), 86 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index a5e82c517..291822ade 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -200,8 +200,6 @@ def _otel_trace_env(workspace: str) -> dict[str, str]: CLAUDE_REMOVED_ENV_KEYS = ("CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS",) CLAUDE_MANAGED_PICKER_KEYS = ("availableModels", "enforceAvailableModels", "modelPicker") ANTHROPIC_CUSTOM_HEADERS_ENV_KEY = "ANTHROPIC_CUSTOM_HEADERS" -# Header names ug manages within ANTHROPIC_CUSTOM_HEADERS. Used only for the no-managed-config -# merge below; ownership is judged by name, never by inferring it from a whole prior document. CLAUDE_MANAGED_CUSTOM_HEADER_NAMES = frozenset( { "x-databricks-use-coding-agent-mode", @@ -229,54 +227,6 @@ def _apply_managed_header_lines( return list(lines_by_name.values()) -def _merge_anthropic_custom_headers(existing: object, ucode_headers: str) -> 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. - 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 - silently discarding user configuration we do not understand. - """ - - if not isinstance(existing, str) or not existing: - return ucode_headers - - ucode_lines_by_name: dict[str, str] = {} - ucode_header_names: list[str] = [] - for line in ucode_headers.splitlines(): - name, separator, _value = line.partition(":") - normalized_name = name.strip().casefold() - if separator and normalized_name not in ucode_lines_by_name: - ucode_header_names.append(normalized_name) - if separator: - ucode_lines_by_name[normalized_name] = line - - merged: list[str] = [] - replaced_names: set[str] = set() - 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: - 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 line: - merged.append(line) - - for name in ucode_header_names: - if name not in replaced_names: - merged.append(ucode_lines_by_name[name]) - return "\n".join(merged) - - def configured_paths(state: dict) -> list[str]: """The Claude config file ug writes; the OS-managed file is added by the dispatcher.""" return [str(CLAUDE_SETTINGS_PATH)] @@ -1001,9 +951,6 @@ def _compose( managed_settings_snapshots: ManagedFileSnapshots | None, ) -> dict: base_env = base.get("env") - # Read before deep_merge_dict below, which mutates base's "env" dict in place (grafting - # overlay's value onto this same object) -- after that call this would read back overlay's - # own value instead of what the developer actually had. existing_custom_headers = ( base_env.get(ANTHROPIC_CUSTOM_HEADERS_ENV_KEY) if isinstance(base_env, dict) else None ) @@ -1160,6 +1107,54 @@ def _compose( return state +def _merge_anthropic_custom_headers(existing: object, ucode_headers: str) -> 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. + 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 + silently discarding user configuration we do not understand. + """ + + if not isinstance(existing, str) or not existing: + return ucode_headers + + ucode_lines_by_name: dict[str, str] = {} + ucode_header_names: list[str] = [] + for line in ucode_headers.splitlines(): + name, separator, _value = line.partition(":") + normalized_name = name.strip().casefold() + if separator and normalized_name not in ucode_lines_by_name: + ucode_header_names.append(normalized_name) + if separator: + ucode_lines_by_name[normalized_name] = line + + merged: list[str] = [] + replaced_names: set[str] = set() + 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: + 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 line: + merged.append(line) + + for name in ucode_header_names: + if name not in replaced_names: + merged.append(ucode_lines_by_name[name]) + return "\n".join(merged) + + def _reconcile_managed_settings( state: dict, compose: Callable[[dict], dict], diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 815731aa1..e5290adda 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -582,6 +582,66 @@ def test_headers_newline_delimited(self, monkeypatch): assert "\n" in self._ua(monkeypatch) +class TestMergeAnthropicCustomHeaders: + def test_removes_stale_parent_header(self): + existing = "X-User: keep\nDatabricks-Model-Service-Parent-Schema: main.default" + managed = "x-databricks-use-coding-agent-mode: true" + + merged = claude._merge_anthropic_custom_headers(existing, managed) + + assert "X-User: keep" in merged + assert "Databricks-Model-Service-Parent-Schema" not in merged + + def test_merges_existing_settings_with_ucode_managed_headers(self): + headers_from_existing_settings = "\n".join( + [ + "X-User-Header: keep-me", + "user-agent: custom-agent", + ] + ) + headers_managed_by_ucode = "\n".join( + [ + "x-databricks-use-coding-agent-mode: true", + "User-Agent: ucode/1.0 claude/2.0", + ] + ) + + merged_headers = claude._merge_anthropic_custom_headers( + headers_from_existing_settings, headers_managed_by_ucode + ) + + assert merged_headers.splitlines() == [ + "X-User-Header: keep-me", # Preserved from existing settings. + "User-Agent: ucode/1.0 claude/2.0", # From ucode; overwrites existing. + "x-databricks-use-coding-agent-mode: true", # Newly added by ucode. + ] + + def test_preserves_existing_header_order(self): + headers_from_existing_settings = "\n".join( + [ + "x-databricks-use-coding-agent-mode: true", + "User-Agent: ucode/0.1.0+41.gd09c080 claude/2.1.258", + "meep: lala", + ] + ) + headers_managed_by_ucode = "\n".join( + [ + "x-databricks-use-coding-agent-mode: true", + "User-Agent: ucode/1.0 claude/2.0", + ] + ) + + merged_headers = claude._merge_anthropic_custom_headers( + headers_from_existing_settings, headers_managed_by_ucode + ) + + assert merged_headers.splitlines() == [ + "x-databricks-use-coding-agent-mode: true", # From ucode; overwrites existing. + "User-Agent: ucode/1.0 claude/2.0", # From ucode; overwrites existing. + "meep: lala", # Preserved from existing settings in its original position. + ] + + class TestRenderOverlayWebSearchDisable: def test_settings_overlay_never_includes_mcp_servers(self): # MCP servers belong in ~/.claude.json, not settings.json. @@ -1618,39 +1678,6 @@ def test_static_models_not_written_when_absent(self, monkeypatch): assert "modelPicker" not in managed_content -class TestMergeAnthropicCustomHeaders: - """Focused unit tests for the static-name-set merge used on the no-managed-config path.""" - - def test_preserves_unknown_header(self): - merged = claude._merge_anthropic_custom_headers( - "X-Foreign: keep", "x-databricks-use-coding-agent-mode: true" - ) - assert merged.splitlines() == [ - "X-Foreign: keep", - "x-databricks-use-coding-agent-mode: true", - ] - - def test_replaces_ucode_managed_name_in_place(self): - merged = claude._merge_anthropic_custom_headers( - "A: 1\nUser-Agent: old\nB: 2", "User-Agent: new" - ) - assert merged.splitlines() == ["A: 1", "User-Agent: new", "B: 2"] - - def test_drops_stale_ucode_managed_name_no_longer_emitted(self): - merged = claude._merge_anthropic_custom_headers( - "Databricks-Model-Provider-Service: old-provider\nA: 1", "" - ) - assert merged.splitlines() == ["A: 1"] - - def test_appends_new_ucode_header_not_already_present(self): - merged = claude._merge_anthropic_custom_headers("A: 1", "User-Agent: ucode/1") - assert merged.splitlines() == ["A: 1", "User-Agent: ucode/1"] - - def test_passes_through_when_no_existing_headers(self): - merged = claude._merge_anthropic_custom_headers(None, "User-Agent: ucode/1") - assert merged == "User-Agent: ucode/1" - - class TestAddClaudeMcpServer: def test_registers_stdio_proxy_command(self, monkeypatch): calls: list[dict] = []