From e4f26ad04ba846a95d1647acc2107b1ee82fc7d8 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 1 Sep 2026 22:36:30 +0000 Subject: [PATCH 1/2] Merge user Anthropic custom headers --- src/ucode/agents/claude.py | 41 +++++++++++++++++++++++ tests/test_agent_claude.py | 67 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 108 insertions(+) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 4fea68e8..5ef42e4b 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -193,6 +193,14 @@ def _resolve_web_search_model(state: dict) -> str | None: # Env keys ucode used to write but no longer does; stripped from the managed # settings file on every launch so stale values never linger. CLAUDE_REMOVED_ENV_KEYS = ("CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS",) +ANTHROPIC_CUSTOM_HEADERS_ENV_KEY = "ANTHROPIC_CUSTOM_HEADERS" +CLAUDE_MANAGED_CUSTOM_HEADER_NAMES = frozenset( + { + "x-databricks-use-coding-agent-mode", + "user-agent", + "databricks-model-provider-service", + } +) CLAUDE_TRACING_STOP_HOOK_SUFFIX = " autolog claude stop-hook" # Tracing is driven by an `mlflow autolog claude stop-hook` Stop hook, run by # the `mlflow` CLI on each session end. Pin to 3.11.x: 3.12 dropped the Unity @@ -626,9 +634,19 @@ def write_tool_config( # V2 installs routing hooks in a transient per-launch settings file. Persistent settings must # contain no ucode routing hooks; surgically strip legacy ones while preserving user hooks. def _compose(base: dict) -> dict: + base_env = base.get("env") + existing_custom_headers = ( + base_env.get(ANTHROPIC_CUSTOM_HEADERS_ENV_KEY) + if isinstance(base_env, dict) + else None + ) # deepcopy the overlay per file so merging into one base can't alias nested dicts into # the other (deep_merge_dict grafts overlay's own dict objects onto a base missing the key). merged = deep_merge_dict(base, copy.deepcopy(overlay)) + overlay_custom_headers = overlay["env"][ANTHROPIC_CUSTOM_HEADERS_ENV_KEY] + 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: @@ -687,6 +705,29 @@ def _compose(base: dict) -> dict: 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 rather than a map, + so a normal deep merge replaces it wholesale. Header names are case-insensitive; any existing + line with a name ucode writes is dropped before ucode's current value is appended. Non-header + lines are retained to avoid silently discarding user configuration we do not understand. + """ + + if not isinstance(existing, str) or not existing: + return ucode_headers + + preserved: list[str] = [] + for line in existing.splitlines(): + name, separator, _value = line.partition(":") + if separator and name.strip().casefold() in CLAUDE_MANAGED_CUSTOM_HEADER_NAMES: + continue + if line: + preserved.append(line) + preserved.extend(ucode_headers.splitlines()) + return "\n".join(preserved) + + 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 29088326..3d3ba971 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -382,6 +382,51 @@ def test_headers_newline_delimited(self, monkeypatch): assert "\n" in self._ua(monkeypatch) +class TestMergeAnthropicCustomHeaders: + def test_preserves_existing_user_headers(self): + existing = "X-User-Header: keep-me\nopaque-user-value" + managed = "User-Agent: ucode/1.0 claude/2.0" + + merged = claude._merge_anthropic_custom_headers(existing, managed) + + assert "X-User-Header: keep-me" in merged + assert "opaque-user-value" in merged + + def test_overrides_existing_ucode_owned_headers(self): + existing = "\n".join( + [ + "user-agent: custom-agent", + "X-Databricks-Use-Coding-Agent-Mode: false", + "Databricks-Model-Provider-Service: stale.provider.service", + ] + ) + managed = "\n".join( + [ + "x-databricks-use-coding-agent-mode: true", + "User-Agent: ucode/1.0 claude/2.0", + ] + ) + + merged = claude._merge_anthropic_custom_headers(existing, managed) + + assert "custom-agent" not in merged + assert "false" not in merged + assert "stale.provider.service" not in merged + assert merged.endswith(managed) + + def test_adds_managed_headers_when_none_exist(self): + managed = "\n".join( + [ + "x-databricks-use-coding-agent-mode: true", + "User-Agent: ucode/1.0 claude/2.0", + ] + ) + + merged = claude._merge_anthropic_custom_headers(None, managed) + + assert merged == managed + + class TestRenderOverlayWebSearchDisable: def test_settings_overlay_never_includes_mcp_servers(self): # MCP servers belong in ~/.claude.json, not settings.json. @@ -599,6 +644,28 @@ def test_managed_file_preserves_other_keys(self, monkeypatch): assert written["env"]["ANTHROPIC_BASE_URL"] assert written["apiKeyHelper"] + def test_managed_file_merges_anthropic_custom_headers(self, monkeypatch): + private_writes: list = [] + managed_writes: list = [] + existing = { + str(FAKE_MANAGED_PATH): { + "env": { + "ANTHROPIC_CUSTOM_HEADERS": "X-Enterprise-Header: retain\nUser-Agent: old" + } + } + } + self._patch(monkeypatch, private_writes, managed_writes, existing) + state = {"workspace": WS, "codex_models": []} + + claude.write_tool_config(state, "databricks-claude-sonnet-4") + + _, text = managed_writes[0] + headers = json.loads(text)["env"]["ANTHROPIC_CUSTOM_HEADERS"] + assert "X-Enterprise-Header: retain" in headers + assert "User-Agent: old" not in headers + assert "User-Agent: ucode/" in headers + assert "x-databricks-use-coding-agent-mode: true" in headers + def test_managed_file_preserves_enterprise_permission_denies(self, monkeypatch): private_writes: list = [] managed_writes: list = [] From 3b5bb53c73e9e8137a61dff470f2c3e636f6131e Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Tue, 1 Sep 2026 22:57:02 +0000 Subject: [PATCH 2/2] fix --- src/ucode/agents/claude.py | 18 ++++++----- tests/test_agent_claude.py | 62 ++++++++++++++------------------------ 2 files changed, 33 insertions(+), 47 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 5ef42e4b..3ce53822 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -636,9 +636,7 @@ def write_tool_config( def _compose(base: dict) -> dict: base_env = base.get("env") existing_custom_headers = ( - base_env.get(ANTHROPIC_CUSTOM_HEADERS_ENV_KEY) - if isinstance(base_env, dict) - else None + base_env.get(ANTHROPIC_CUSTOM_HEADERS_ENV_KEY) if isinstance(base_env, dict) else None ) # deepcopy the overlay per file so merging into one base can't alias nested dicts into # the other (deep_merge_dict grafts overlay's own dict objects onto a base missing the key). @@ -708,10 +706,16 @@ def _compose(base: dict) -> dict: 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 rather than a map, - so a normal deep merge replaces it wholesale. Header names are case-insensitive; any existing - line with a name ucode writes is dropped before ucode's current value is appended. Non-header - lines are retained to avoid silently discarding user configuration we do not understand. + 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, while + preserving all other existing headers. + 4. Combine the preserved existing headers with ucode's managed header values. + + 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: diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 3d3ba971..eea05850 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -383,48 +383,29 @@ def test_headers_newline_delimited(self, monkeypatch): class TestMergeAnthropicCustomHeaders: - def test_preserves_existing_user_headers(self): - existing = "X-User-Header: keep-me\nopaque-user-value" - managed = "User-Agent: ucode/1.0 claude/2.0" - - merged = claude._merge_anthropic_custom_headers(existing, managed) - - assert "X-User-Header: keep-me" in merged - assert "opaque-user-value" in merged - - def test_overrides_existing_ucode_owned_headers(self): - existing = "\n".join( + 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", - "X-Databricks-Use-Coding-Agent-Mode: false", - "Databricks-Model-Provider-Service: stale.provider.service", ] ) - managed = "\n".join( + headers_managed_by_ucode = "\n".join( [ "x-databricks-use-coding-agent-mode: true", "User-Agent: ucode/1.0 claude/2.0", ] ) - merged = claude._merge_anthropic_custom_headers(existing, managed) - - assert "custom-agent" not in merged - assert "false" not in merged - assert "stale.provider.service" not in merged - assert merged.endswith(managed) - - def test_adds_managed_headers_when_none_exist(self): - managed = "\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 ) - merged = claude._merge_anthropic_custom_headers(None, managed) - - assert merged == managed + assert merged_headers.splitlines() == [ + "X-User-Header: keep-me", # Preserved from existing settings. + "x-databricks-use-coding-agent-mode: true", # Newly added by ucode. + "User-Agent: ucode/1.0 claude/2.0", # From ucode; overwrites existing. + ] class TestRenderOverlayWebSearchDisable: @@ -647,24 +628,25 @@ def test_managed_file_preserves_other_keys(self, monkeypatch): def test_managed_file_merges_anthropic_custom_headers(self, monkeypatch): private_writes: list = [] managed_writes: list = [] - existing = { + existing_managed_settings = { str(FAKE_MANAGED_PATH): { - "env": { - "ANTHROPIC_CUSTOM_HEADERS": "X-Enterprise-Header: retain\nUser-Agent: old" - } + "env": {"ANTHROPIC_CUSTOM_HEADERS": "X-Enterprise-Header: retain\nUser-Agent: old"} } } - self._patch(monkeypatch, private_writes, managed_writes, existing) + self._patch(monkeypatch, private_writes, managed_writes, existing_managed_settings) + monkeypatch.setattr(claude, "ucode_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] - headers = json.loads(text)["env"]["ANTHROPIC_CUSTOM_HEADERS"] - assert "X-Enterprise-Header: retain" in headers - assert "User-Agent: old" not in headers - assert "User-Agent: ucode/" in headers - assert "x-databricks-use-coding-agent-mode: true" in headers + 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", # Newly added by ucode. + "User-Agent: ucode/1.0 claude/2.0", # From ucode; overwrites existing. + ] def test_managed_file_preserves_enterprise_permission_denies(self, monkeypatch): private_writes: list = []