diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 602c2538..f7333140 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -196,6 +196,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 @@ -629,9 +637,17 @@ 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: @@ -690,6 +706,54 @@ 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. 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 c76b46e1..bd504df4 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -410,6 +410,57 @@ def test_headers_newline_delimited(self, monkeypatch): assert "\n" in self._ua(monkeypatch) +class TestMergeAnthropicCustomHeaders: + 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. @@ -627,6 +678,29 @@ 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_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, "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] + merged_headers = json.loads(text)["env"]["ANTHROPIC_CUSTOM_HEADERS"] + assert merged_headers.splitlines() == [ + "X-Enterprise-Header: retain", # Preserved from existing managed 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_managed_file_preserves_enterprise_permission_denies(self, monkeypatch): private_writes: list = [] managed_writes: list = []