From 57e6d51b0620689ada5f9f4b1656a8c71e886624 Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Thu, 22 Oct 2026 20:02:48 +0000 Subject: [PATCH 1/5] Add a Claude Code mod so `ug --header` reaches Claude Claude takes ANTHROPIC_CUSTOM_HEADERS whole from managed settings, so a header from `ug --header` / UG_HEADERS in ug's own settings file or in the environment never reached the gateway, and writing it into the machine-wide file would leak it into every session. `ug configure` (and an interactive launch) now installs a small Claude mod in a root-owned `ucode-plugins` folder beside the managed settings, on Linux and macOS, through the existing sudo worker; its four files join the allowlist. The managed file enables it as the organization's mod, ahead of the built-in guard, the one tier allowed to change a variable managed settings set. Before each request the mod adds the session's UG_HEADERS lines to ANTHROPIC_CUSTOM_HEADERS in that Claude process only; a header already present keeps its value. Revert removes the managed entries; a launch warns when the mod isn't active. The sudo script now creates each missing parent folder root-owned and 0755: `mkdir -p` left the folders above the innermost one with the caller's umask, so under 027 neither ug nor Claude could read the mod. Co-authored-by: Isaac --- AGENTS.md | 1 + README.md | 13 ++- src/ucode/agents/claude.py | 100 ++++++++++++++++- .../agents/claude_header_mod/register.ts | 37 ++++++ src/ucode/cli.py | 4 +- src/ucode/managed_files.py | 29 +++++ src/ucode/managed_replace.sh | 25 ++++- tests/README.md | 6 + tests/conftest.py | 3 + .../header_mod_tests/register.test.ts | 45 ++++++++ tests/native_claude/test_mod.py | 19 ++++ tests/test_claude_header_mod.py | 106 ++++++++++++++++++ tests/test_managed_files.py | 15 ++- 13 files changed, 386 insertions(+), 17 deletions(-) create mode 100644 src/ucode/agents/claude_header_mod/register.ts create mode 100644 tests/native_claude/header_mod_tests/register.test.ts create mode 100644 tests/test_claude_header_mod.py diff --git a/AGENTS.md b/AGENTS.md index 8ea7fb4c7..448699ae5 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -135,6 +135,7 @@ Fields live in `~/.claude/ucode-settings.json` and the OS-managed settings file | Smart-routing hooks | Merge | Merge | `PreToolUse`, `SessionStart`, `SubagentStart`; only `ug`'s own marked handlers, other hooks left alone | | Admin `agent_native_settings` | Ignore | Create/replace | Managed file only; copied as-is. Keys `ug` owns are skipped, `permissions.deny` and hooks are merged, and `ug`'s MCP servers are added to an admin `allowedMcpServers`. Removed when the admin drops them, unless edited by hand | | Smart Router Orchestrator hooks | Merge | Merge | Launch-only `UserPromptSubmit` and compact `SessionStart` handlers for smart-routed sessions with `ENABLE_SMART_ROUTER_ORCHESTRATOR=1`; read the same session controls as routing | +| Header mod | Merge | Merge | Managed file only, Linux and macOS: `extraKnownMarketplaces.ucode-managed`, `enabledPlugins["ug-headers@ucode-managed"]`, and `ug-headers@ucode-managed` first in `prependPlugins` (an unset list becomes the mod plus `sec-default@builtin`). The mod's files live in a root-owned `ucode-plugins` folder beside the managed file and add `UG_HEADERS` to each request; revert removes the entries and leaves the files | diff --git a/README.md b/README.md index 01c122ce9..626e4e848 100644 --- a/README.md +++ b/README.md @@ -145,14 +145,19 @@ machine-wide; use `--global` for it. ### Extra headers for one session ```bash -ug --header "x-databricks-traffic-id: testenv://liteswap/" codex +ug --header "x-databricks-traffic-id: testenv://liteswap/" claude UG_HEADERS="X-Team: eng" ug codex ``` `--header` (repeatable) and `UG_HEADERS` (one `Name: value` per line) add headers to that launch's -Databricks requests only: Codex's gateway requests, through the launch's `-c` provider headers, and -ug's own requests (workspace APIs, MCP proxy, web search, smart routing). A header already sent -keeps its value. Nothing is written to a config file. +Databricks requests only: Claude Code's and Codex's gateway requests, and ug's own requests +(workspace APIs, MCP proxy, web search, smart routing). A header already sent keeps its value. +Nothing is written to a config file. + +- **Codex**: added to the launch's `-c` provider headers. +- **Claude Code**: added by a mod, even when managed settings set `ANTHROPIC_CUSTOM_HEADERS`. + `ug configure` installs it beside the managed settings (Linux and macOS, with sudo, not under + `--user`); without it the launch warns and Claude doesn't send them. Available coding agents are `codex`, `claude`, `gemini`, `opencode`, `copilot`, and `pi`. `cursor` can be included in `--agents` for MCP-only setup; diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 250439294..29ee5ef1a 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -21,6 +21,7 @@ ToolSpec, backup_existing_file, deep_merge_dict, + is_dry_run, prune_key_paths, read_json_safe, write_json_file, @@ -33,6 +34,7 @@ MODEL_PROVIDER_SERVICE_HEADER, MODEL_SERVICE_PARENT_SCHEMA_HEADER, SMART_ROUTER_RECIPE_HEADER, + UG_HEADERS_ENV_VAR, ) from ucode.custom_oauth import ( CustomOAuthConfig, @@ -53,12 +55,14 @@ from ucode.launcher import exec_or_spawn from ucode.managed_config import refresh_managed_config from ucode.managed_files import ( + CLAUDE_HEADER_MOD_FILES, OS, ManagedFileSnapshots, ManagedFileWriteUnavailable, SettingsPassthrough, apply_settings_passthrough, check_skipped_managed_overrides, + claude_header_mod_root, config_scope, current_os, managed_file_conflicts, @@ -73,6 +77,7 @@ plan_settings_passthrough, read_managed_file, reconcile_managed_file, + reconcile_privileged_files, record_settings_passthrough, record_ug_picker, revert_managed_file, @@ -1755,6 +1760,7 @@ def _required_custom_headers(value: object) -> dict[str, str] | None: def _managed_settings_conflicts( existing: dict, desired: dict, owned_paths: list[list[str]] ) -> list[str]: + owned_paths = [path for path in owned_paths if path not in CLAUDE_HEADER_MOD_PATHS] conflicts = managed_file_conflicts(existing, desired, owned_paths) header_path = f"env.{ANTHROPIC_CUSTOM_HEADERS_ENV_KEY}" if header_path in conflicts: @@ -1820,6 +1826,87 @@ def _allow_ug_mcp_servers(settings: dict, admin_allowlist: list | None, names: s return settings +CLAUDE_HEADER_MOD_MARKETPLACE = "ucode-managed" +CLAUDE_HEADER_MOD_ID = f"ug-headers@{CLAUDE_HEADER_MOD_MARKETPLACE}" +_CLAUDE_GUARD_MOD_ID = "sec-default@builtin" +# The header mod's managed-file entries. Revert removes them; they never count as conflicts, since +# nothing in ug's own settings file can be overridden by them. +CLAUDE_HEADER_MOD_PATHS = [ + ["extraKnownMarketplaces", CLAUDE_HEADER_MOD_MARKETPLACE], + ["enabledPlugins", CLAUDE_HEADER_MOD_ID], + ["prependPlugins"], +] + + +def _header_mod_files(root: Path) -> dict[Path, str]: + """The marketplace and plugin ug installs at ``root`` for its Claude header mod.""" + description = "Adds `ug --header` / UG_HEADERS to Claude's gateway requests." + docs = { + ".claude-plugin/marketplace.json": { + "name": CLAUDE_HEADER_MOD_MARKETPLACE, + "owner": {"name": "Unity Gateway"}, + # A relative source loads in place, which Claude requires of an organization's mod. + "plugins": [ + {"name": "ug-headers", "source": "./ug-headers", "description": description} + ], + }, + "ug-headers/.claude-plugin/plugin.json": {"name": "ug-headers", "description": description}, + "ug-headers/hooks/hooks.json": {"modules": ["./register.ts"]}, + } + files = {root / name: json.dumps(doc, indent=2) + "\n" for name, doc in docs.items()} + mod = Path(__file__).with_name("claude_header_mod") / "register.ts" + files[root / "ug-headers/hooks/register.ts"] = mod.read_text(encoding="utf-8") + return files + + +def _enable_header_mod(settings: dict, root: Path | None) -> dict: + """Enable ug's header mod in the managed file and run it ahead of every user's mod.""" + if root is None: + return settings + marketplaces = settings.setdefault("extraKnownMarketplaces", {}) + plugins = settings.setdefault("enabledPlugins", {}) + if not isinstance(marketplaces, dict) or not isinstance(plugins, dict): + return settings + marketplaces[CLAUDE_HEADER_MOD_MARKETPLACE] = { + "source": {"source": "directory", "path": str(root)} + } + plugins[CLAUDE_HEADER_MOD_ID] = True + # An unset list means the built-in guard alone, which a set list has to name to keep it. + prepend = settings.get("prependPlugins") + prepend = prepend if isinstance(prepend, list) else [_CLAUDE_GUARD_MOD_ID] + settings["prependPlugins"] = [ + CLAUDE_HEADER_MOD_ID, + *(entry for entry in prepend if entry != CLAUDE_HEADER_MOD_ID), + ] + return settings + + +def _install_header_mod(root: Path | None) -> None: + """Install or update the header mod's files; a failure only costs ``ug --header`` for Claude.""" + if root is None or is_dry_run(): + return + try: + reconcile_privileged_files(_header_mod_files(root)) + except (OSError, RuntimeError, subprocess.CalledProcessError) as exc: + print_warning(f"Couldn't install ug's Claude header mod at {root}: {exc}") + + +def header_mod_active() -> bool: + """Whether ug's header mod is installed and enabled, so ``UG_HEADERS`` reaches Claude.""" + root = claude_header_mod_root() + path = _managed_settings_path() + if root is None or path is None: + return False + try: + if not (root / CLAUDE_HEADER_MOD_FILES[-1]).is_file(): + return False + settings = _parse_managed_settings(read_managed_file(path) or "{}") + except (OSError, RuntimeError): + return False + plugins = settings.get("enabledPlugins") + return isinstance(plugins, dict) and plugins.get(CLAUDE_HEADER_MOD_ID) is True + + def _is_shared_claude_list(path: list[str]) -> bool: """Lists IT and ug also write to in the managed file, so the admin's entries merge into them.""" return path == CLAUDE_PERMISSIONS_DENY_PATH or (len(path) == 2 and path[0] == "hooks") @@ -1922,7 +2009,10 @@ def _reconcile_managed_settings( "ucode did not modify the file. Repair it or contact your administrator." ) from exc managed_before = copy.deepcopy(existing) - desired_settings = compose(existing) + header_mod_root = claude_header_mod_root() + desired_settings = _enable_header_mod(compose(existing), header_mod_root) + if header_mod_root is not None: + owned_paths = [*owned_paths, *CLAUDE_HEADER_MOD_PATHS] _preserve_permission_denies(managed_before, desired_settings, withdrawn=withdrawn_denies or []) conflicts = _managed_settings_conflicts(managed_before, desired_settings, owned_paths) if not managed_writes_allowed() and not conflicts: @@ -1950,6 +2040,8 @@ def _reconcile_managed_settings( mark_managed_file_verified(state, "claude", path) record_ug_picker("claude", {key: desired_settings[key] for key in picker_keys}) record_settings_passthrough("claude", passthrough_leaves or []) + # After the settings write, so a first install reuses its sudo session and password notice. + _install_header_mod(header_mod_root) def _preserve_permission_denies( @@ -2311,6 +2403,12 @@ def launch( tool_args = _external_web_search_args(state, tool_args) binary = SPEC["binary"] workspace = state.get("workspace") + if os.environ.get(UG_HEADERS_ENV_VAR, "").strip() and not header_mod_active(): + print_warning( + "Claude Code won't send your `ug --header` / UG_HEADERS headers: ug's header mod isn't " + "installed. Run `ug configure` from an interactive terminal on Linux or macOS, without " + "UG_SKIP_OS_SETTINGS." + ) if workspace and os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1": # Discovery is launch-scoped. Pass it in the process environment rather # than persisting it in Claude's private or OS-managed settings. diff --git a/src/ucode/agents/claude_header_mod/register.ts b/src/ucode/agents/claude_header_mod/register.ts new file mode 100644 index 000000000..0d56fadcd --- /dev/null +++ b/src/ucode/agents/claude_header_mod/register.ts @@ -0,0 +1,37 @@ +import type { EngineInterface, On } from 'claude-code'; + +// Adds the launch's `UG_HEADERS` lines (`ug --header`) to `ANTHROPIC_CUSTOM_HEADERS` for this +// process only. ug installs it as the organization's mod, the one tier allowed to change a variable +// managed settings set. A header whose name is already present is left alone, so neither ug's nor an +// admin's header can be replaced from a launch. +export function headerLines(base: string | undefined, extra: string | undefined): string | undefined { + const added = (extra ?? '').split('\n').filter((line) => line.trim()); + if (!added.length) return undefined; + const lines = (base ?? '').split('\n').filter((line) => line.trim()); + const name = (line: string) => line.split(':', 1)[0].trim().toLowerCase(); + const present = new Set(lines.map(name)); + const merged = [...lines, ...added.filter((line) => !present.has(name(line)))]; + return merged.length === lines.length ? undefined : merged.join('\n'); +} + +async function apply($: EngineInterface): Promise { + // Claude's capability scanner requires literal environment-variable names. + const merged = headerLines( + await $.env.get('ANTHROPIC_CUSTOM_HEADERS'), + await $.env.get('UG_HEADERS'), + ); + if (merged !== undefined) await $.env.set('ANTHROPIC_CUSTOM_HEADERS', merged); +} + +export function register(on: On): void { + on('session.start', async ($, e, next) => { + await apply($); + return next(e); + }); + // Claude rebuilds its client per request and a settings reload resets the variable, so add the + // headers again before each one. + on('turn.step', async function* ($, e, next) { + await apply($); + return yield* next(e); + }); +} diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 0dc1cc36c..b75c10663 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -3841,8 +3841,8 @@ def default( list[str] | None, typer.Option( "--header", - help="Extra `Name: value` header for this launch's Codex and ug Databricks requests " - "(repeatable). A header already sent keeps its value.", + help="Extra `Name: value` header for this launch's Claude Code, Codex, and ug " + "Databricks requests (repeatable). A header already sent keeps its value.", ), ] = None, ) -> None: diff --git a/src/ucode/managed_files.py b/src/ucode/managed_files.py index 5e355649c..e942b2a03 100644 --- a/src/ucode/managed_files.py +++ b/src/ucode/managed_files.py @@ -68,12 +68,27 @@ class OS(Enum): MANAGED_FILES_SUPPORTED_OSES = (OS.LINUX, OS.MACOS, OS.WINDOWS) +# ug's root-owned Claude plugin marketplace, holding the mod that adds `UG_HEADERS` to each request. +# Claude treats a mod as the organization's only when it loads from a directory marketplace that +# managed settings name, so it lives beside the managed file rather than in ug's own install. +CLAUDE_HEADER_MOD_ROOTS = { + OS.LINUX: Path("/etc/claude-code/ucode-plugins"), + OS.MACOS: Path("/Library/Application Support/ClaudeCode/ucode-plugins"), +} +CLAUDE_HEADER_MOD_FILES = ( + ".claude-plugin/marketplace.json", + "ug-headers/.claude-plugin/plugin.json", + "ug-headers/hooks/hooks.json", + "ug-headers/hooks/register.ts", +) + _SUDO_REPLACE_TARGETS = { OS.LINUX: frozenset( { Path("/etc/claude-code/managed-settings.json"), Path("/etc/codex/managed_config.toml"), UNITY_GATEWAY_MANAGED_SETTINGS_PATH, + *(CLAUDE_HEADER_MOD_ROOTS[OS.LINUX] / name for name in CLAUDE_HEADER_MOD_FILES), } ), OS.MACOS: frozenset( @@ -81,6 +96,7 @@ class OS(Enum): Path("/Library/Application Support/ClaudeCode/managed-settings.json"), Path("/etc/codex/managed_config.toml"), UNITY_GATEWAY_MANAGED_SETTINGS_PATH, + *(CLAUDE_HEADER_MOD_ROOTS[OS.MACOS] / name for name in CLAUDE_HEADER_MOD_FILES), } ), } @@ -136,6 +152,19 @@ def managed_files_supported() -> bool: return current_os() in MANAGED_FILES_SUPPORTED_OSES +def claude_header_mod_root() -> Path | None: + """Where ug installs its Claude header mod, or ``None`` where it installs none.""" + # shortcut: no Windows install (no allowlisted admin folder), add one when Windows needs it. + return None if managed_settings_skipped() else CLAUDE_HEADER_MOD_ROOTS.get(current_os()) + + +def reconcile_privileged_files(files: dict[Path, str]) -> None: + """Write each allowlisted root-owned file whose contents differ, through the sudo worker.""" + for path, text in files.items(): + if read_managed_file(path) != text: + _privileged_replace(path, text) + + def windows_claude_managed_settings_path() -> Path: """The only machine-wide settings file Claude Code reads on Windows.""" return system_cross_os.windows_program_files_dir() / "ClaudeCode" / "managed-settings.json" diff --git a/src/ucode/managed_replace.sh b/src/ucode/managed_replace.sh index 8e7170217..1f372cdb0 100644 --- a/src/ucode/managed_replace.sh +++ b/src/ucode/managed_replace.sh @@ -31,9 +31,17 @@ target_is_allowed() { "linux:/etc/claude-code/managed-settings.json"|\ "linux:/etc/codex/managed_config.toml"|\ "linux:/etc/unitygateway/managed-settings.json"|\ + "linux:/etc/claude-code/ucode-plugins/.claude-plugin/marketplace.json"|\ + "linux:/etc/claude-code/ucode-plugins/ug-headers/.claude-plugin/plugin.json"|\ + "linux:/etc/claude-code/ucode-plugins/ug-headers/hooks/hooks.json"|\ + "linux:/etc/claude-code/ucode-plugins/ug-headers/hooks/register.ts"|\ "macos:/Library/Application Support/ClaudeCode/managed-settings.json"|\ "macos:/etc/codex/managed_config.toml"|\ - "macos:/etc/unitygateway/managed-settings.json") return 0 ;; + "macos:/etc/unitygateway/managed-settings.json"|\ + "macos:/Library/Application Support/ClaudeCode/ucode-plugins/.claude-plugin/marketplace.json"|\ + "macos:/Library/Application Support/ClaudeCode/ucode-plugins/ug-headers/.claude-plugin/plugin.json"|\ + "macos:/Library/Application Support/ClaudeCode/ucode-plugins/ug-headers/hooks/hooks.json"|\ + "macos:/Library/Application Support/ClaudeCode/ucode-plugins/ug-headers/hooks/register.ts") return 0 ;; *) return 1 ;; esac } @@ -132,10 +140,19 @@ replace_one() ( exit 1 fi + # `mkdir -p` would leave the folders above `$parent` with the caller's umask (027 makes them + # unreadable to the agents), so create each missing folder root-owned and world-readable. + make_parent() { + [ -n "$1" ] && [ ! -d "$1" ] || return 0 + if [ -L "$1" ] || [ -e "$1" ]; then + return 1 + fi + make_parent "${1%/*}" || return 1 + mkdir "$1" && chown 0:0 "$1" && chmod 0755 "$1" + } + if [ ! -d "$parent" ]; then - run_step mkdir -p "$parent" - run_step chown 0:0 "$parent" - run_step chmod 0755 "$parent" + run_step make_parent "$parent" fi if [ "$target" = /etc/unitygateway/managed-settings.json ]; then run_step chown 0:0 "$parent" diff --git a/tests/README.md b/tests/README.md index ed05edb9f..b5fe813cf 100644 --- a/tests/README.md +++ b/tests/README.md @@ -48,6 +48,12 @@ ug's own requests carrying them, and `test_agent_codex.py` `UG_HEADERS` reaching `-c` headers (plain and smart-routed) without writing its profile. No automated test checks the header at the gateway. +`test_claude_header_mod.py` covers installing the Claude header mod's files and managed entries +(no rewrite when unchanged), keeping an admin's `prependPlugins`, revert, the sudo allowlist, and an +unreadable mod folder. `native_claude/test_mod.py` validates the mod and runs its merge tests +(`native_claude/header_mod_tests`). Claude re-reading `ANTHROPIC_CUSTOM_HEADERS` per request was +checked by hand on Claude 2.1.304; no automated test sends a request through the installed mod. + The required Isaac feedback form and background delivery notifications need manual validation using the released Isaac skill and existing capture/retry CLI. Check category selection, empty and inline feedback, native retries, the off form's diff --git a/tests/conftest.py b/tests/conftest.py index f722eac5f..88c0136d4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -37,6 +37,7 @@ def _isolate_ucode_state(tmp_path, monkeypatch): import ucode.os_compatibility.subprocess_cross_os as subprocess_cross_os_mod import ucode.state as state_mod from ucode import desktop_auth as desktop_auth_mod + from ucode.agents import claude as claude_mod from ucode.agents import codex as codex_mod from ucode.agents.desktop import claude as desktop_mod @@ -83,6 +84,8 @@ def install_test_policy(path, desired_text, *, interactive=False, allow_sudo=Tru ) monkeypatch.setattr(managed_files_mod, "_config_scope", "global") monkeypatch.setattr(codex_mod, "codex_managed_config_path", lambda: None) + # The header mod installs beside the real managed file; tests opt in with a tmp root. + monkeypatch.setattr(claude_mod, "claude_header_mod_root", lambda: None) monkeypatch.setattr( codex_mod, "CODEX_MODEL_CATALOG_PATH", state_dir / "codex-model-catalog.json" ) diff --git a/tests/native_claude/header_mod_tests/register.test.ts b/tests/native_claude/header_mod_tests/register.test.ts new file mode 100644 index 000000000..68eddad62 --- /dev/null +++ b/tests/native_claude/header_mod_tests/register.test.ts @@ -0,0 +1,45 @@ +import type { On } from 'claude-code'; +import { expect, test } from 'claude-code/testing'; +import { headerLines } from '../hooks/register.ts'; + +const start = { surface: 'terminal', isInteractive: true, cwd: '/work' } as const; + +function setup(on: On, env: Record) { + on('env.get', ($, e) => ({ value: env[e.name] })); + on('env.set', ($, e) => { + env[e.name] = e.value; + return { value: undefined }; + }); + // Stands in for Claude's own handling beneath the mod. + on('session.start', ($, e) => ({ cwd: e.cwd })); + return env; +} + +test('session start adds the launch headers after the managed ones', async ($, on) => { + const env = setup(on, { + ANTHROPIC_CUSTOM_HEADERS: 'x-databricks-use-coding-agent-mode: true', + UG_HEADERS: 'x-databricks-traffic-id: testenv://liteswap/pod-1', + }); + + await $.session.start(start); + + expect(env.ANTHROPIC_CUSTOM_HEADERS).toBe( + 'x-databricks-use-coding-agent-mode: true\nx-databricks-traffic-id: testenv://liteswap/pod-1', + ); +}); + +test('a header already sent keeps its value, and reapplying changes nothing', () => { + const base = 'X-Team: eng\nUser-Agent: ucode/1'; + const merged = headerLines(base, 'x-team: other\nX-New: 1'); + expect(merged).toBe('X-Team: eng\nUser-Agent: ucode/1\nX-New: 1'); + expect(headerLines(merged, 'x-team: other\nX-New: 1')).toBeUndefined(); +}); + +test('without launch headers the variable is left alone', async ($, on) => { + const env = setup(on, { ANTHROPIC_CUSTOM_HEADERS: 'X-Team: eng' }); + + await $.session.start(start); + + expect(env.ANTHROPIC_CUSTOM_HEADERS).toBe('X-Team: eng'); + expect(headerLines(undefined, '\n')).toBeUndefined(); +}); diff --git a/tests/native_claude/test_mod.py b/tests/native_claude/test_mod.py index 166db82a4..139cc6f2e 100644 --- a/tests/native_claude/test_mod.py +++ b/tests/native_claude/test_mod.py @@ -12,6 +12,7 @@ import pytest from session import Session +from ucode.agents.claude import _header_mod_files from ucode.smart_routing import config from ucode.smart_routing.v2 import _claude_mod_launch_options, _write_routed_claude_plugin @@ -53,6 +54,24 @@ def test_mod_validation_and_event_forwarding(native_binary, tmp_path): assert "./register.ts hooks: session.start" in result.stdout +def test_header_mod_validation_and_merge(native_binary, tmp_path): + for path, text in _header_mod_files(tmp_path).items(): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(text) + plugin = tmp_path / "ug-headers" + shutil.copytree(Path(__file__).with_name("header_mod_tests"), plugin / "tests") + for command in ("validate", "test"): + result = subprocess.run( + [native_binary, "plugin", command, str(plugin)], + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, result.stdout + result.stderr + if command == "validate": + assert "env writes: ANTHROPIC_CUSTOM_HEADERS" in result.stdout + + def test_native_options_are_launch_local(native_binary, tmp_path): env = { "SMART_ROUTER_CONFIG_VERSION": config.SUBAGENT_ORCH_V0_CLAUDE_ONLY, diff --git a/tests/test_claude_header_mod.py b/tests/test_claude_header_mod.py new file mode 100644 index 000000000..3c9c4cfb1 --- /dev/null +++ b/tests/test_claude_header_mod.py @@ -0,0 +1,106 @@ +"""ug's Claude header mod: install beside the managed file, enable it there, and revert it.""" + +import copy +import json + +from ucode import managed_files +from ucode.agents import claude +from ucode.managed_files import CLAUDE_HEADER_MOD_FILES, CLAUDE_HEADER_MOD_ROOTS, OS + +WS = "https://example.cloud.databricks.com" + + +def _managed_env(tmp_path, monkeypatch, managed_doc=None): + """Real reconcile flow, with privileged writes landing in ``tmp_path`` and counted.""" + managed_path = tmp_path / "managed-settings.json" + if managed_doc is not None: + managed_path.write_text(json.dumps(managed_doc), encoding="utf-8") + root = tmp_path / "ucode-plugins" + backup_dir = tmp_path / "managed-backups" + writes: list[str] = [] + monkeypatch.setattr(managed_files, "managed_files_supported", lambda: True) + monkeypatch.setattr(managed_files, "current_os", lambda: OS.LINUX) + monkeypatch.setattr(managed_files, "MANAGED_BACKUP_DIR", backup_dir) + monkeypatch.setattr(managed_files, "MANAGED_BACKUP_MANIFEST_PATH", backup_dir / "manifest.json") + + def _write(target, text): + writes.append(str(target)) + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(text, encoding="utf-8") + + monkeypatch.setattr(managed_files, "_sudo_replace", _write) + monkeypatch.setattr(managed_files, "_sudo_remove", lambda path: path.unlink()) + monkeypatch.setattr(claude, "_managed_settings_path", lambda: managed_path) + monkeypatch.setattr(claude, "claude_header_mod_root", lambda: root) + monkeypatch.setattr(claude, "managed_writes_allowed", lambda: True) + monkeypatch.setattr(managed_files.sys.stdin, "isatty", 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") + return managed_path, root, writes + + +def test_configure_installs_and_enables_mod_then_stays_put(tmp_path, monkeypatch): + managed_path, root, writes = _managed_env(tmp_path, monkeypatch) + state = {"workspace": WS, "codex_models": []} + + claude.write_tool_config(state, "databricks-claude-sonnet-4") + + for path, text in claude._header_mod_files(root).items(): + assert path.read_text() == text + doc = json.loads(managed_path.read_text()) + assert doc["extraKnownMarketplaces"]["ucode-managed"] == { + "source": {"source": "directory", "path": str(root)} + } + assert doc["enabledPlugins"] == {"ug-headers@ucode-managed": True} + # The built-in guard is the default list, so naming the mod there must keep it. + assert doc["prependPlugins"] == ["ug-headers@ucode-managed", "sec-default@builtin"] + assert claude.header_mod_active() + + count = len(writes) + claude.write_tool_config(state, "databricks-claude-sonnet-4") + assert len(writes) == count # unchanged files are never rewritten through sudo + + +def test_admin_prepend_list_is_kept_and_revert_restores_it(tmp_path, monkeypatch): + managed_path, _root, _writes = _managed_env( + tmp_path, monkeypatch, {"prependPlugins": ["acme-guard@acme"]} + ) + + claude.write_tool_config({"workspace": WS, "codex_models": []}, "databricks-claude-sonnet-4") + assert json.loads(managed_path.read_text())["prependPlugins"] == [ + "ug-headers@ucode-managed", + "acme-guard@acme", + ] + + claude.revert_managed_settings() + assert json.loads(managed_path.read_text()) == {"prependPlugins": ["acme-guard@acme"]} + + +def test_mod_entries_never_count_as_conflicts(): + existing = {"prependPlugins": ["acme-guard@acme"], "enabledPlugins": {"x@y": True}} + desired = claude._enable_header_mod(copy.deepcopy(existing), CLAUDE_HEADER_MOD_ROOTS[OS.LINUX]) + assert ( + claude._managed_settings_conflicts(existing, desired, claude.CLAUDE_HEADER_MOD_PATHS) == [] + ) + + +def test_mod_files_are_on_the_sudo_allowlist(): + for os_family, root in CLAUDE_HEADER_MOD_ROOTS.items(): + files = set(claude._header_mod_files(root)) + assert files == {root / name for name in CLAUDE_HEADER_MOD_FILES} + assert files <= managed_files._SUDO_REPLACE_TARGETS[os_family] + + +def test_unreadable_mod_folder_reads_as_inactive(tmp_path, monkeypatch): + # A folder the user can't enter (an earlier install under umask 027) must warn, not crash. + monkeypatch.setattr(claude, "claude_header_mod_root", lambda: tmp_path) + monkeypatch.setattr(claude, "_managed_settings_path", lambda: tmp_path / "managed.json") + + def denied(self): + raise PermissionError(13, "Permission denied", str(self)) + + monkeypatch.setattr(claude.Path, "is_file", denied) + assert claude.header_mod_active() is False diff --git a/tests/test_managed_files.py b/tests/test_managed_files.py index 4783af283..a65626d31 100644 --- a/tests/test_managed_files.py +++ b/tests/test_managed_files.py @@ -958,12 +958,13 @@ def test_real_shell_preserves_file_on_failed_step(self, tmp_path, failed_step): assert target.stat().st_mode & 0o777 == 0o640 assert not list(tmp_path.glob(".managed_config.toml.ucode.*")) - @pytest.mark.parametrize("parent_exists", [True, False]) - def test_real_shell_creates_missing_file(self, tmp_path, parent_exists): + @pytest.mark.parametrize("missing_parents", [0, 1, 2]) + def test_real_shell_creates_missing_file(self, tmp_path, missing_parents): """The new-file branch installs content, mode, and (root) ownership from scratch.""" source = tmp_path / "source.json" source.write_text("managed settings\n", encoding="utf-8") - parent = tmp_path if parent_exists else tmp_path / "codex" + created = [tmp_path / "codex", tmp_path / "codex" / "nested"][:missing_parents] + parent = created[-1] if created else tmp_path target = parent / "managed_config.toml" chown_log = tmp_path / "chown.log" # Unit-only relocation: production still accepts only its fixed machine-wide targets. @@ -978,6 +979,8 @@ def test_real_shell_creates_missing_file(self, tmp_path, parent_exists): } """ script = script.replace('\ncase "$mode" in', shim_chown + '\ncase "$mode" in') + # A restrictive umask must not leave any created folder unreadable to the agents. + script = "umask 027\n" + script request = " ".join( [ "REPLACE", @@ -1003,9 +1006,9 @@ def test_real_shell_creates_missing_file(self, tmp_path, parent_exists): assert target.read_text(encoding="utf-8") == "managed settings\n" assert target.stat().st_mode & 0o777 == 0o644 expected_chown_calls = [] - if not parent_exists: - assert parent.stat().st_mode & 0o777 == 0o755 - expected_chown_calls.append(f"0:0 {parent}") + for folder in created: + assert folder.stat().st_mode & 0o777 == 0o755 + expected_chown_calls.append(f"0:0 {folder}") chown_calls = chown_log.read_text(encoding="utf-8").splitlines() assert chown_calls[: len(expected_chown_calls)] == expected_chown_calls assert len(chown_calls) == len(expected_chown_calls) + 1 From 9a2b1a63fe398f363798e47b867460603b3a645b Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Thu, 22 Oct 2026 20:36:38 +0000 Subject: [PATCH 2/5] Fix the header mod folders' umask with one `umask 022` The sudo script created the mod's parent folders with the caller's umask (027 left them unreadable to ug and Claude). Setting `umask 022` in the root script fixes that for every folder `mkdir -p` makes, so the per-folder creation helper isn't needed. Co-authored-by: Isaac --- src/ucode/managed_replace.sh | 17 +++++------------ tests/test_managed_files.py | 3 +-- 2 files changed, 6 insertions(+), 14 deletions(-) diff --git a/src/ucode/managed_replace.sh b/src/ucode/managed_replace.sh index 1f372cdb0..9a92a231e 100644 --- a/src/ucode/managed_replace.sh +++ b/src/ucode/managed_replace.sh @@ -10,6 +10,8 @@ set -u set -f PATH=/usr/bin:/bin:/usr/sbin:/sbin export PATH +# Folders `mkdir -p` creates must stay readable to the agents whatever the caller's umask. +umask 022 mode=$1 platform=$2 @@ -140,19 +142,10 @@ replace_one() ( exit 1 fi - # `mkdir -p` would leave the folders above `$parent` with the caller's umask (027 makes them - # unreadable to the agents), so create each missing folder root-owned and world-readable. - make_parent() { - [ -n "$1" ] && [ ! -d "$1" ] || return 0 - if [ -L "$1" ] || [ -e "$1" ]; then - return 1 - fi - make_parent "${1%/*}" || return 1 - mkdir "$1" && chown 0:0 "$1" && chmod 0755 "$1" - } - if [ ! -d "$parent" ]; then - run_step make_parent "$parent" + run_step mkdir -p "$parent" + run_step chown 0:0 "$parent" + run_step chmod 0755 "$parent" fi if [ "$target" = /etc/unitygateway/managed-settings.json ]; then run_step chown 0:0 "$parent" diff --git a/tests/test_managed_files.py b/tests/test_managed_files.py index a65626d31..4fdd976f0 100644 --- a/tests/test_managed_files.py +++ b/tests/test_managed_files.py @@ -1005,10 +1005,9 @@ def test_real_shell_creates_missing_file(self, tmp_path, missing_parents): assert result.stdout == "OK 1\n" assert target.read_text(encoding="utf-8") == "managed settings\n" assert target.stat().st_mode & 0o777 == 0o644 - expected_chown_calls = [] for folder in created: assert folder.stat().st_mode & 0o777 == 0o755 - expected_chown_calls.append(f"0:0 {folder}") + expected_chown_calls = [f"0:0 {parent}"] if created else [] chown_calls = chown_log.read_text(encoding="utf-8").splitlines() assert chown_calls[: len(expected_chown_calls)] == expected_chown_calls assert len(chown_calls) == len(expected_chown_calls) + 1 From e403dafdfacaaace5bd1df1189c5821a0de08de6 Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Fri, 23 Oct 2026 15:45:30 +0000 Subject: [PATCH 3/5] Enable the header mod only once installed, and clean it up on revert Address review feedback on the Claude header mod: - Install the mod's files before the managed settings name them, and add the entries only when every file is in place (taking ug's entries back out otherwise), so a failed install never enables a missing plugin for every user on the host. The install now shows ug's sudo notice itself. - `ug revert` removes the mod's root-owned files and the folders they leave empty. When an admin has extended `prependPlugins` since configure, revert keeps `sec-default@builtin`, since a non-empty list without it turns the built-in guard off. - `header_mod_active` checks every mod file and all three managed entries. - The mod skips malformed and repeated UG_HEADERS lines and reads CRLF header strings, for launches that don't go through ug. Co-authored-by: Isaac --- AGENTS.md | 2 +- src/ucode/agents/claude.py | 113 +++++++++++++++--- .../agents/claude_header_mod/register.ts | 26 +++- src/ucode/managed_files.py | 40 ++++++- tests/README.md | 5 +- tests/conftest.py | 3 + .../header_mod_tests/register.test.ts | 8 ++ tests/test_claude_header_mod.py | 61 ++++++++++ 8 files changed, 230 insertions(+), 28 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 448699ae5..7d217072e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -135,7 +135,7 @@ Fields live in `~/.claude/ucode-settings.json` and the OS-managed settings file | Smart-routing hooks | Merge | Merge | `PreToolUse`, `SessionStart`, `SubagentStart`; only `ug`'s own marked handlers, other hooks left alone | | Admin `agent_native_settings` | Ignore | Create/replace | Managed file only; copied as-is. Keys `ug` owns are skipped, `permissions.deny` and hooks are merged, and `ug`'s MCP servers are added to an admin `allowedMcpServers`. Removed when the admin drops them, unless edited by hand | | Smart Router Orchestrator hooks | Merge | Merge | Launch-only `UserPromptSubmit` and compact `SessionStart` handlers for smart-routed sessions with `ENABLE_SMART_ROUTER_ORCHESTRATOR=1`; read the same session controls as routing | -| Header mod | Merge | Merge | Managed file only, Linux and macOS: `extraKnownMarketplaces.ucode-managed`, `enabledPlugins["ug-headers@ucode-managed"]`, and `ug-headers@ucode-managed` first in `prependPlugins` (an unset list becomes the mod plus `sec-default@builtin`). The mod's files live in a root-owned `ucode-plugins` folder beside the managed file and add `UG_HEADERS` to each request; revert removes the entries and leaves the files | +| Header mod | Merge | Merge | Managed file only, Linux and macOS: `extraKnownMarketplaces.ucode-managed`, `enabledPlugins["ug-headers@ucode-managed"]`, and `ug-headers@ucode-managed` first in `prependPlugins` (an unset list becomes the mod plus `sec-default@builtin`). The mod's files live in a root-owned `ucode-plugins` folder beside the managed file and add `UG_HEADERS` to each request. The entries are written only once the files are in place; revert removes both, keeping `sec-default@builtin` when an admin has extended the list | diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 29ee5ef1a..15abb4029 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -55,7 +55,6 @@ from ucode.launcher import exec_or_spawn from ucode.managed_config import refresh_managed_config from ucode.managed_files import ( - CLAUDE_HEADER_MOD_FILES, OS, ManagedFileSnapshots, ManagedFileWriteUnavailable, @@ -75,11 +74,13 @@ managed_writes_allowed, mark_managed_file_verified, plan_settings_passthrough, + privileged_files_current, read_managed_file, reconcile_managed_file, reconcile_privileged_files, record_settings_passthrough, record_ug_picker, + remove_privileged_files, revert_managed_file, warn_managed_write_skipped, warn_skipped_settings_passthrough, @@ -399,12 +400,34 @@ def managed_settings_status(state: dict) -> tuple[Path | None, str, str]: def revert_managed_settings() -> str: - return revert_managed_file( - "claude", - display="Claude Code", - parser=_parse_managed_settings, - dumper=_dump_managed_settings, + # ug lists the built-in guard to keep it when it creates `prependPlugins`, so revert takes it + # out with ug's entries. When an admin has since added to the list, put the guard back: a + # non-empty list without it turns the guard off. + guard_listed = _CLAUDE_GUARD_MOD_ID in _managed_prepend_plugins() + + def dump(settings: dict) -> str: + prepend = settings.get("prependPlugins") + if guard_listed and isinstance(prepend, list) and prepend: + if _CLAUDE_GUARD_MOD_ID not in prepend: + settings["prependPlugins"] = [_CLAUDE_GUARD_MOD_ID, *prepend] + return _dump_managed_settings(settings) + + status = revert_managed_file( + "claude", display="Claude Code", parser=_parse_managed_settings, dumper=dump ) + _remove_header_mod() + return status + + +def _managed_prepend_plugins() -> list: + """The managed file's ``prependPlugins`` list, or an empty one.""" + path = _managed_settings_path() + try: + text = read_managed_file(path) if path is not None else None + prepend = _parse_managed_settings(text).get("prependPlugins") if text else None + except RuntimeError: + return [] + return prepend if isinstance(prepend, list) else [] def _managed_relayed_conflicts(path: Path) -> list[str]: @@ -1859,10 +1882,15 @@ def _header_mod_files(root: Path) -> dict[Path, str]: return files -def _enable_header_mod(settings: dict, root: Path | None) -> dict: - """Enable ug's header mod in the managed file and run it ahead of every user's mod.""" +def _enable_header_mod(settings: dict, root: Path | None, *, installed: bool = True) -> dict: + """Enable ug's header mod in the managed file and run it ahead of every user's mod. + + Without its files in place (``installed`` false), take ug's entries back out instead, so the + managed file never names a plugin that isn't there.""" if root is None: return settings + if not installed: + return _disable_header_mod(settings) marketplaces = settings.setdefault("extraKnownMarketplaces", {}) plugins = settings.setdefault("enabledPlugins", {}) if not isinstance(marketplaces, dict) or not isinstance(plugins, dict): @@ -1881,14 +1909,54 @@ def _enable_header_mod(settings: dict, root: Path | None) -> dict: return settings -def _install_header_mod(root: Path | None) -> None: - """Install or update the header mod's files; a failure only costs ``ug --header`` for Claude.""" +def _disable_header_mod(settings: dict) -> dict: + """Remove ug's header-mod entries, restoring the default plugin list when only it was set.""" + for key, name in ( + ("extraKnownMarketplaces", CLAUDE_HEADER_MOD_MARKETPLACE), + ("enabledPlugins", CLAUDE_HEADER_MOD_ID), + ): + entries = settings.get(key) + if isinstance(entries, dict) and name in entries: + del entries[name] + if not entries: + del settings[key] + prepend = settings.get("prependPlugins") + if isinstance(prepend, list) and CLAUDE_HEADER_MOD_ID in prepend: + rest = [entry for entry in prepend if entry != CLAUDE_HEADER_MOD_ID] + if rest in ([], [_CLAUDE_GUARD_MOD_ID]): + del settings["prependPlugins"] + else: + settings["prependPlugins"] = rest + return settings + + +def _install_header_mod(root: Path | None) -> bool: + """Install or update the header mod's files, returning whether they are all in place. + + A failure only costs ``ug --header`` for Claude; the caller then leaves the mod disabled.""" + if root is None: + return False + files = _header_mod_files(root) + if is_dry_run() or not managed_writes_allowed(): + # This run can't write as root: the mod counts as installed only if a run before did it. + return privileged_files_current(files) + try: + reconcile_privileged_files(files, display="Claude Code") + except (OSError, RuntimeError, subprocess.CalledProcessError) as exc: + print_warning(f"Couldn't install ug's Claude header mod at {root}: {exc}") + return False + return True + + +def _remove_header_mod() -> None: + """Remove the header mod's root-owned files and folders, for ``ug revert``.""" + root = claude_header_mod_root() if root is None or is_dry_run(): return try: - reconcile_privileged_files(_header_mod_files(root)) + remove_privileged_files(list(_header_mod_files(root)), root) except (OSError, RuntimeError, subprocess.CalledProcessError) as exc: - print_warning(f"Couldn't install ug's Claude header mod at {root}: {exc}") + print_warning(f"Couldn't remove ug's Claude header mod at {root}: {exc}") def header_mod_active() -> bool: @@ -1898,13 +1966,22 @@ def header_mod_active() -> bool: if root is None or path is None: return False try: - if not (root / CLAUDE_HEADER_MOD_FILES[-1]).is_file(): + if not all(file.is_file() for file in _header_mod_files(root)): return False settings = _parse_managed_settings(read_managed_file(path) or "{}") except (OSError, RuntimeError): return False plugins = settings.get("enabledPlugins") - return isinstance(plugins, dict) and plugins.get(CLAUDE_HEADER_MOD_ID) is True + marketplaces = settings.get("extraKnownMarketplaces") + prepend = settings.get("prependPlugins") + return ( + isinstance(plugins, dict) + and plugins.get(CLAUDE_HEADER_MOD_ID) is True + and isinstance(marketplaces, dict) + and CLAUDE_HEADER_MOD_MARKETPLACE in marketplaces + and isinstance(prepend, list) + and CLAUDE_HEADER_MOD_ID in prepend + ) def _is_shared_claude_list(path: list[str]) -> bool: @@ -2010,7 +2087,11 @@ def _reconcile_managed_settings( ) from exc managed_before = copy.deepcopy(existing) header_mod_root = claude_header_mod_root() - desired_settings = _enable_header_mod(compose(existing), header_mod_root) + # Install the files before the settings name them, so a failed install never enables a + # plugin that isn't there. + desired_settings = _enable_header_mod( + compose(existing), header_mod_root, installed=_install_header_mod(header_mod_root) + ) if header_mod_root is not None: owned_paths = [*owned_paths, *CLAUDE_HEADER_MOD_PATHS] _preserve_permission_denies(managed_before, desired_settings, withdrawn=withdrawn_denies or []) @@ -2040,8 +2121,6 @@ def _reconcile_managed_settings( mark_managed_file_verified(state, "claude", path) record_ug_picker("claude", {key: desired_settings[key] for key in picker_keys}) record_settings_passthrough("claude", passthrough_leaves or []) - # After the settings write, so a first install reuses its sudo session and password notice. - _install_header_mod(header_mod_root) def _preserve_permission_denies( diff --git a/src/ucode/agents/claude_header_mod/register.ts b/src/ucode/agents/claude_header_mod/register.ts index 0d56fadcd..c91a1d8af 100644 --- a/src/ucode/agents/claude_header_mod/register.ts +++ b/src/ucode/agents/claude_header_mod/register.ts @@ -4,13 +4,31 @@ import type { EngineInterface, On } from 'claude-code'; // process only. ug installs it as the organization's mod, the one tier allowed to change a variable // managed settings set. A header whose name is already present is left alone, so neither ug's nor an // admin's header can be replaced from a launch. + +const HEADER_NAME = /^[!#$%&'*+.^_`|~0-9A-Za-z-]+$/; + +// `Name: value` with an HTTP token name, a non-empty value, and no control characters but tab. +function isHeaderLine(line: string): boolean { + const colon = line.indexOf(':'); + return ( + colon > 0 && + HEADER_NAME.test(line.slice(0, colon)) && + line.slice(colon + 1).trim() !== '' && + !/[\x00-\x08\x0a-\x1f\x7f]/.test(line) + ); +} + export function headerLines(base: string | undefined, extra: string | undefined): string | undefined { - const added = (extra ?? '').split('\n').filter((line) => line.trim()); - if (!added.length) return undefined; - const lines = (base ?? '').split('\n').filter((line) => line.trim()); + const lines = (base ?? '').split(/\r?\n/).filter((line) => line.trim()); const name = (line: string) => line.split(':', 1)[0].trim().toLowerCase(); const present = new Set(lines.map(name)); - const merged = [...lines, ...added.filter((line) => !present.has(name(line)))]; + const merged = [...lines]; + // Malformed lines are skipped: launched without ug, nothing else has checked them. + for (const line of (extra ?? '').split(/\r?\n/).map((l) => l.trim())) { + if (!isHeaderLine(line) || present.has(name(line))) continue; + present.add(name(line)); + merged.push(line); + } return merged.length === lines.length ? undefined : merged.join('\n'); } diff --git a/src/ucode/managed_files.py b/src/ucode/managed_files.py index e942b2a03..8883d7711 100644 --- a/src/ucode/managed_files.py +++ b/src/ucode/managed_files.py @@ -158,11 +158,43 @@ def claude_header_mod_root() -> Path | None: return None if managed_settings_skipped() else CLAUDE_HEADER_MOD_ROOTS.get(current_os()) -def reconcile_privileged_files(files: dict[Path, str]) -> None: +def privileged_files_current(files: dict[Path, str]) -> bool: + """Whether every root-owned file already holds its contents.""" + try: + return all(read_managed_file(path) == text for path, text in files.items()) + except RuntimeError: + return False + + +def reconcile_privileged_files(files: dict[Path, str], *, display: str) -> None: """Write each allowlisted root-owned file whose contents differ, through the sudo worker.""" - for path, text in files.items(): - if read_managed_file(path) != text: - _privileged_replace(path, text) + changed = [(path, text) for path, text in files.items() if read_managed_file(path) != text] + if changed and _sudo_may_prompt(): + _print_managed_write_permission(display) + for path, text in changed: + _privileged_replace(path, text) + + +def remove_privileged_files(paths: list[Path], root: Path) -> None: + """Remove root-owned files ug installed, then the folders under ``root`` they leave empty.""" + for path in paths: + if path.is_file() or path.is_symlink(): + _privileged_remove(path) + folders = { + parent for path in paths for parent in path.parents if root in (parent, *parent.parents) + } + for folder in sorted(folders, key=lambda f: len(f.parts), reverse=True): + if folder.is_dir() and not folder.is_symlink() and not any(folder.iterdir()): + _privileged_rmdir(folder) + + +def _privileged_rmdir(folder: Path) -> None: + if current_os() is OS.WINDOWS: + folder.rmdir() + return + subprocess_cross_os.run( + _sudo_command("rmdir", str(folder)), capture_output=True, text=True, check=True + ) def windows_claude_managed_settings_path() -> Path: diff --git a/tests/README.md b/tests/README.md index b5fe813cf..7d205dfdb 100644 --- a/tests/README.md +++ b/tests/README.md @@ -49,8 +49,9 @@ ug's own requests carrying them, and `test_agent_codex.py` `UG_HEADERS` reaching header at the gateway. `test_claude_header_mod.py` covers installing the Claude header mod's files and managed entries -(no rewrite when unchanged), keeping an admin's `prependPlugins`, revert, the sudo allowlist, and an -unreadable mod folder. `native_claude/test_mod.py` validates the mod and runs its merge tests +(no rewrite when unchanged), keeping an admin's `prependPlugins`, a failed install never enabling +the plugin, revert removing the files and keeping the guard, the sudo allowlist, `header_mod_active` +needing every file and entry, and an unreadable mod folder. `native_claude/test_mod.py` validates the mod and runs its merge tests (`native_claude/header_mod_tests`). Claude re-reading `ANTHROPIC_CUSTOM_HEADERS` per request was checked by hand on Claude 2.1.304; no automated test sends a request through the installed mod. diff --git a/tests/conftest.py b/tests/conftest.py index 88c0136d4..c32948cf7 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -98,6 +98,9 @@ def reject_privileged_write(path, _desired_text): ) monkeypatch.setattr(managed_files_mod, "_sudo_replace", reject_privileged_write) + monkeypatch.setattr( + managed_files_mod, "_privileged_rmdir", lambda path: reject_privileged_write(path, "") + ) # Windows CI runners are elevated; never let an unmocked test write C:\Program Files. monkeypatch.setattr(managed_files_mod.system_cross_os, "is_admin", lambda: False) monkeypatch.setattr(managed_files_mod, "_windows_replace", reject_privileged_write) diff --git a/tests/native_claude/header_mod_tests/register.test.ts b/tests/native_claude/header_mod_tests/register.test.ts index 68eddad62..0c82b9502 100644 --- a/tests/native_claude/header_mod_tests/register.test.ts +++ b/tests/native_claude/header_mod_tests/register.test.ts @@ -35,6 +35,14 @@ test('a header already sent keeps its value, and reapplying changes nothing', () expect(headerLines(merged, 'x-team: other\nX-New: 1')).toBeUndefined(); }); +test('malformed and repeated launch lines are skipped, and CRLF bases are read', () => { + const extra = 'no colon\nbad name: 1\nX-Empty: \nX-Ctl: a\u0001b\nX-New: 1\nx-new: 2'; + expect(headerLines('X-Team: eng\r\nUser-Agent: ucode/1', extra)).toBe( + 'X-Team: eng\nUser-Agent: ucode/1\nX-New: 1', + ); + expect(headerLines('X-Team: eng', 'no colon\nbad name: 1')).toBeUndefined(); +}); + test('without launch headers the variable is left alone', async ($, on) => { const env = setup(on, { ANTHROPIC_CUSTOM_HEADERS: 'X-Team: eng' }); diff --git a/tests/test_claude_header_mod.py b/tests/test_claude_header_mod.py index 3c9c4cfb1..16c9fb831 100644 --- a/tests/test_claude_header_mod.py +++ b/tests/test_claude_header_mod.py @@ -30,6 +30,7 @@ def _write(target, text): monkeypatch.setattr(managed_files, "_sudo_replace", _write) monkeypatch.setattr(managed_files, "_sudo_remove", lambda path: path.unlink()) + monkeypatch.setattr(managed_files, "_privileged_rmdir", lambda folder: folder.rmdir()) monkeypatch.setattr(claude, "_managed_settings_path", lambda: managed_path) monkeypatch.setattr(claude, "claude_header_mod_root", lambda: root) monkeypatch.setattr(claude, "managed_writes_allowed", lambda: True) @@ -104,3 +105,63 @@ def denied(self): monkeypatch.setattr(claude.Path, "is_file", denied) assert claude.header_mod_active() is False + + +def test_failed_install_never_enables_a_missing_plugin(tmp_path, monkeypatch): + managed_path, root, _writes = _managed_env(tmp_path, monkeypatch) + write = managed_files._sudo_replace + + def refuse_mod_files(target, text): + if root in target.parents: + raise PermissionError("denied") + write(target, text) + + monkeypatch.setattr(managed_files, "_sudo_replace", refuse_mod_files) + + claude.write_tool_config({"workspace": WS, "codex_models": []}, "databricks-claude-sonnet-4") + + doc = json.loads(managed_path.read_text()) + assert not {"extraKnownMarketplaces", "enabledPlugins", "prependPlugins"} & set(doc) + assert not claude.header_mod_active() + + +def test_revert_removes_the_mod_files_and_folders(tmp_path, monkeypatch): + _managed_path, root, _writes = _managed_env(tmp_path, monkeypatch) + claude.write_tool_config({"workspace": WS, "codex_models": []}, "databricks-claude-sonnet-4") + assert (root / "ug-headers/hooks/register.ts").is_file() + + claude.revert_managed_settings() + + assert not root.exists() + + +def test_revert_keeps_the_guard_when_an_admin_extended_the_list(tmp_path, monkeypatch): + managed_path, _root, _writes = _managed_env(tmp_path, monkeypatch) + claude.write_tool_config({"workspace": WS, "codex_models": []}, "databricks-claude-sonnet-4") + doc = json.loads(managed_path.read_text()) + doc["prependPlugins"].append("acme-guard@acme") + managed_path.write_text(json.dumps(doc), encoding="utf-8") + + claude.revert_managed_settings() + + # A non-empty list without the guard would turn it off, so revert keeps it. + assert json.loads(managed_path.read_text())["prependPlugins"] == [ + "sec-default@builtin", + "acme-guard@acme", + ] + + +def test_header_mod_active_needs_every_file_and_entry(tmp_path, monkeypatch): + managed_path, root, _writes = _managed_env(tmp_path, monkeypatch) + claude.write_tool_config({"workspace": WS, "codex_models": []}, "databricks-claude-sonnet-4") + assert claude.header_mod_active() + + enabled = managed_path.read_text() + doc = json.loads(enabled) + doc["prependPlugins"] = ["sec-default@builtin"] + managed_path.write_text(json.dumps(doc), encoding="utf-8") + assert not claude.header_mod_active() + + managed_path.write_text(enabled, encoding="utf-8") + (root / "ug-headers/.claude-plugin/plugin.json").unlink() + assert not claude.header_mod_active() From c2d80fdf4ea985c247f24d11893c9057204c2320 Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Fri, 23 Oct 2026 16:10:44 +0000 Subject: [PATCH 4/5] Trim the header mod review fixes to the minimal change Enable the mod only when this run installed its files, leaving the managed file as is otherwise, instead of taking ug's entries back out; a run that can't write as root no longer needs to check the files. Remove the mod's folder on revert with one `rm -rf` of ug's fixed `ucode-plugins` path. `header_mod_active` checks every file and `enabledPlugins`. Co-authored-by: Isaac --- src/ucode/agents/claude.py | 67 +++++++-------------------------- src/ucode/managed_files.py | 36 ++++-------------- tests/README.md | 2 +- tests/conftest.py | 3 -- tests/test_claude_header_mod.py | 14 ++----- 5 files changed, 26 insertions(+), 96 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 15abb4029..912b49373 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -74,13 +74,12 @@ managed_writes_allowed, mark_managed_file_verified, plan_settings_passthrough, - privileged_files_current, read_managed_file, reconcile_managed_file, reconcile_privileged_files, record_settings_passthrough, record_ug_picker, - remove_privileged_files, + remove_privileged_tree, revert_managed_file, warn_managed_write_skipped, warn_skipped_settings_passthrough, @@ -1882,15 +1881,10 @@ def _header_mod_files(root: Path) -> dict[Path, str]: return files -def _enable_header_mod(settings: dict, root: Path | None, *, installed: bool = True) -> dict: - """Enable ug's header mod in the managed file and run it ahead of every user's mod. - - Without its files in place (``installed`` false), take ug's entries back out instead, so the - managed file never names a plugin that isn't there.""" +def _enable_header_mod(settings: dict, root: Path | None) -> dict: + """Enable ug's header mod in the managed file and run it ahead of every user's mod.""" if root is None: return settings - if not installed: - return _disable_header_mod(settings) marketplaces = settings.setdefault("extraKnownMarketplaces", {}) plugins = settings.setdefault("enabledPlugins", {}) if not isinstance(marketplaces, dict) or not isinstance(plugins, dict): @@ -1909,39 +1903,14 @@ def _enable_header_mod(settings: dict, root: Path | None, *, installed: bool = T return settings -def _disable_header_mod(settings: dict) -> dict: - """Remove ug's header-mod entries, restoring the default plugin list when only it was set.""" - for key, name in ( - ("extraKnownMarketplaces", CLAUDE_HEADER_MOD_MARKETPLACE), - ("enabledPlugins", CLAUDE_HEADER_MOD_ID), - ): - entries = settings.get(key) - if isinstance(entries, dict) and name in entries: - del entries[name] - if not entries: - del settings[key] - prepend = settings.get("prependPlugins") - if isinstance(prepend, list) and CLAUDE_HEADER_MOD_ID in prepend: - rest = [entry for entry in prepend if entry != CLAUDE_HEADER_MOD_ID] - if rest in ([], [_CLAUDE_GUARD_MOD_ID]): - del settings["prependPlugins"] - else: - settings["prependPlugins"] = rest - return settings - - def _install_header_mod(root: Path | None) -> bool: - """Install or update the header mod's files, returning whether they are all in place. + """Install or update the header mod's files, returning whether this run did. - A failure only costs ``ug --header`` for Claude; the caller then leaves the mod disabled.""" - if root is None: + A failure only costs ``ug --header`` for Claude; the caller then doesn't enable the mod.""" + if root is None or is_dry_run() or not managed_writes_allowed(): return False - files = _header_mod_files(root) - if is_dry_run() or not managed_writes_allowed(): - # This run can't write as root: the mod counts as installed only if a run before did it. - return privileged_files_current(files) try: - reconcile_privileged_files(files, display="Claude Code") + reconcile_privileged_files(_header_mod_files(root), display="Claude Code") except (OSError, RuntimeError, subprocess.CalledProcessError) as exc: print_warning(f"Couldn't install ug's Claude header mod at {root}: {exc}") return False @@ -1949,12 +1918,12 @@ def _install_header_mod(root: Path | None) -> bool: def _remove_header_mod() -> None: - """Remove the header mod's root-owned files and folders, for ``ug revert``.""" + """Remove the header mod's root-owned folder, for ``ug revert``.""" root = claude_header_mod_root() if root is None or is_dry_run(): return try: - remove_privileged_files(list(_header_mod_files(root)), root) + remove_privileged_tree(root) except (OSError, RuntimeError, subprocess.CalledProcessError) as exc: print_warning(f"Couldn't remove ug's Claude header mod at {root}: {exc}") @@ -1972,16 +1941,7 @@ def header_mod_active() -> bool: except (OSError, RuntimeError): return False plugins = settings.get("enabledPlugins") - marketplaces = settings.get("extraKnownMarketplaces") - prepend = settings.get("prependPlugins") - return ( - isinstance(plugins, dict) - and plugins.get(CLAUDE_HEADER_MOD_ID) is True - and isinstance(marketplaces, dict) - and CLAUDE_HEADER_MOD_MARKETPLACE in marketplaces - and isinstance(prepend, list) - and CLAUDE_HEADER_MOD_ID in prepend - ) + return isinstance(plugins, dict) and plugins.get(CLAUDE_HEADER_MOD_ID) is True def _is_shared_claude_list(path: list[str]) -> bool: @@ -2088,10 +2048,9 @@ def _reconcile_managed_settings( managed_before = copy.deepcopy(existing) header_mod_root = claude_header_mod_root() # Install the files before the settings name them, so a failed install never enables a - # plugin that isn't there. - desired_settings = _enable_header_mod( - compose(existing), header_mod_root, installed=_install_header_mod(header_mod_root) - ) + # plugin that isn't there; a run that can't write as root leaves any earlier entries as is. + installed = _install_header_mod(header_mod_root) + desired_settings = _enable_header_mod(compose(existing), header_mod_root if installed else None) if header_mod_root is not None: owned_paths = [*owned_paths, *CLAUDE_HEADER_MOD_PATHS] _preserve_permission_denies(managed_before, desired_settings, withdrawn=withdrawn_denies or []) diff --git a/src/ucode/managed_files.py b/src/ucode/managed_files.py index 8883d7711..5724c5a1d 100644 --- a/src/ucode/managed_files.py +++ b/src/ucode/managed_files.py @@ -158,14 +158,6 @@ def claude_header_mod_root() -> Path | None: return None if managed_settings_skipped() else CLAUDE_HEADER_MOD_ROOTS.get(current_os()) -def privileged_files_current(files: dict[Path, str]) -> bool: - """Whether every root-owned file already holds its contents.""" - try: - return all(read_managed_file(path) == text for path, text in files.items()) - except RuntimeError: - return False - - def reconcile_privileged_files(files: dict[Path, str], *, display: str) -> None: """Write each allowlisted root-owned file whose contents differ, through the sudo worker.""" changed = [(path, text) for path, text in files.items() if read_managed_file(path) != text] @@ -175,26 +167,14 @@ def reconcile_privileged_files(files: dict[Path, str], *, display: str) -> None: _privileged_replace(path, text) -def remove_privileged_files(paths: list[Path], root: Path) -> None: - """Remove root-owned files ug installed, then the folders under ``root`` they leave empty.""" - for path in paths: - if path.is_file() or path.is_symlink(): - _privileged_remove(path) - folders = { - parent for path in paths for parent in path.parents if root in (parent, *parent.parents) - } - for folder in sorted(folders, key=lambda f: len(f.parts), reverse=True): - if folder.is_dir() and not folder.is_symlink() and not any(folder.iterdir()): - _privileged_rmdir(folder) - - -def _privileged_rmdir(folder: Path) -> None: - if current_os() is OS.WINDOWS: - folder.rmdir() - return - subprocess_cross_os.run( - _sudo_command("rmdir", str(folder)), capture_output=True, text=True, check=True - ) +def remove_privileged_tree(root: Path) -> None: + """Remove a root-owned folder ug owns entirely, such as its header mod's, through sudo.""" + if root.is_symlink(): + raise RuntimeError(f"Refusing to remove {root} through a symlink.") + if root.exists(): + subprocess_cross_os.run( + _sudo_command("rm", "-rf", "--", str(root)), capture_output=True, text=True, check=True + ) def windows_claude_managed_settings_path() -> Path: diff --git a/tests/README.md b/tests/README.md index 7d205dfdb..6482276ff 100644 --- a/tests/README.md +++ b/tests/README.md @@ -51,7 +51,7 @@ header at the gateway. `test_claude_header_mod.py` covers installing the Claude header mod's files and managed entries (no rewrite when unchanged), keeping an admin's `prependPlugins`, a failed install never enabling the plugin, revert removing the files and keeping the guard, the sudo allowlist, `header_mod_active` -needing every file and entry, and an unreadable mod folder. `native_claude/test_mod.py` validates the mod and runs its merge tests +needing every file, and an unreadable mod folder. `native_claude/test_mod.py` validates the mod and runs its merge tests (`native_claude/header_mod_tests`). Claude re-reading `ANTHROPIC_CUSTOM_HEADERS` per request was checked by hand on Claude 2.1.304; no automated test sends a request through the installed mod. diff --git a/tests/conftest.py b/tests/conftest.py index c32948cf7..88c0136d4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -98,9 +98,6 @@ def reject_privileged_write(path, _desired_text): ) monkeypatch.setattr(managed_files_mod, "_sudo_replace", reject_privileged_write) - monkeypatch.setattr( - managed_files_mod, "_privileged_rmdir", lambda path: reject_privileged_write(path, "") - ) # Windows CI runners are elevated; never let an unmocked test write C:\Program Files. monkeypatch.setattr(managed_files_mod.system_cross_os, "is_admin", lambda: False) monkeypatch.setattr(managed_files_mod, "_windows_replace", reject_privileged_write) diff --git a/tests/test_claude_header_mod.py b/tests/test_claude_header_mod.py index 16c9fb831..2b53697da 100644 --- a/tests/test_claude_header_mod.py +++ b/tests/test_claude_header_mod.py @@ -2,6 +2,7 @@ import copy import json +import shutil from ucode import managed_files from ucode.agents import claude @@ -30,7 +31,7 @@ def _write(target, text): monkeypatch.setattr(managed_files, "_sudo_replace", _write) monkeypatch.setattr(managed_files, "_sudo_remove", lambda path: path.unlink()) - monkeypatch.setattr(managed_files, "_privileged_rmdir", lambda folder: folder.rmdir()) + monkeypatch.setattr(claude, "remove_privileged_tree", shutil.rmtree) monkeypatch.setattr(claude, "_managed_settings_path", lambda: managed_path) monkeypatch.setattr(claude, "claude_header_mod_root", lambda: root) monkeypatch.setattr(claude, "managed_writes_allowed", lambda: True) @@ -151,17 +152,10 @@ def test_revert_keeps_the_guard_when_an_admin_extended_the_list(tmp_path, monkey ] -def test_header_mod_active_needs_every_file_and_entry(tmp_path, monkeypatch): - managed_path, root, _writes = _managed_env(tmp_path, monkeypatch) +def test_header_mod_active_needs_every_file(tmp_path, monkeypatch): + _managed_path, root, _writes = _managed_env(tmp_path, monkeypatch) claude.write_tool_config({"workspace": WS, "codex_models": []}, "databricks-claude-sonnet-4") assert claude.header_mod_active() - enabled = managed_path.read_text() - doc = json.loads(enabled) - doc["prependPlugins"] = ["sec-default@builtin"] - managed_path.write_text(json.dumps(doc), encoding="utf-8") - assert not claude.header_mod_active() - - managed_path.write_text(enabled, encoding="utf-8") (root / "ug-headers/.claude-plugin/plugin.json").unlink() assert not claude.header_mod_active() From fb7ad1a30210287bd5d08c76cc916d0f07da84dc Mon Sep 17 00:00:00 2001 From: Andy Xu Date: Fri, 23 Oct 2026 20:12:13 +0000 Subject: [PATCH 5/5] Add the Claude header mod only on configure or a launch with headers A plain launch re-added the mod's managed-settings entries whenever another writer of the file dropped them, so every session rewrote it through sudo (CUJ 8). `ug configure` and launches with UG_HEADERS still install it. Also stub refresh_managed_config in the header mod tests, which reached the databricks CLI on CI. Co-authored-by: Isaac --- AGENTS.md | 2 +- src/ucode/agents/claude.py | 14 +++++++++++++- src/ucode/cli.py | 1 + tests/test_claude_header_mod.py | 29 +++++++++++++++++++++++++++++ 4 files changed, 44 insertions(+), 2 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 7d217072e..bec3d7247 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -135,7 +135,7 @@ Fields live in `~/.claude/ucode-settings.json` and the OS-managed settings file | Smart-routing hooks | Merge | Merge | `PreToolUse`, `SessionStart`, `SubagentStart`; only `ug`'s own marked handlers, other hooks left alone | | Admin `agent_native_settings` | Ignore | Create/replace | Managed file only; copied as-is. Keys `ug` owns are skipped, `permissions.deny` and hooks are merged, and `ug`'s MCP servers are added to an admin `allowedMcpServers`. Removed when the admin drops them, unless edited by hand | | Smart Router Orchestrator hooks | Merge | Merge | Launch-only `UserPromptSubmit` and compact `SessionStart` handlers for smart-routed sessions with `ENABLE_SMART_ROUTER_ORCHESTRATOR=1`; read the same session controls as routing | -| Header mod | Merge | Merge | Managed file only, Linux and macOS: `extraKnownMarketplaces.ucode-managed`, `enabledPlugins["ug-headers@ucode-managed"]`, and `ug-headers@ucode-managed` first in `prependPlugins` (an unset list becomes the mod plus `sec-default@builtin`). The mod's files live in a root-owned `ucode-plugins` folder beside the managed file and add `UG_HEADERS` to each request. The entries are written only once the files are in place; revert removes both, keeping `sec-default@builtin` when an admin has extended the list | +| Header mod | Merge | Merge | Managed file only, Linux and macOS: `extraKnownMarketplaces.ucode-managed`, `enabledPlugins["ug-headers@ucode-managed"]`, and `ug-headers@ucode-managed` first in `prependPlugins` (an unset list becomes the mod plus `sec-default@builtin`). The mod's files live in a root-owned `ucode-plugins` folder beside the managed file and add `UG_HEADERS` to each request. Added by `ug configure` or a launch with `UG_HEADERS`, never by a plain launch, so another writer of the file that drops them isn't fought on every launch. The entries are written only once the files are in place; revert removes both, keeping `sec-default@builtin` when an admin has extended the list | diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 912b49373..fb9af9d2f 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -1903,6 +1903,15 @@ def _enable_header_mod(settings: dict, root: Path | None) -> dict: return settings +_header_mod_requested = False + + +def request_header_mod() -> None: + """Have this run add the header mod even without headers, as ``ug configure`` does.""" + global _header_mod_requested + _header_mod_requested = True + + def _install_header_mod(root: Path | None) -> bool: """Install or update the header mod's files, returning whether this run did. @@ -2047,9 +2056,12 @@ def _reconcile_managed_settings( ) from exc managed_before = copy.deepcopy(existing) header_mod_root = claude_header_mod_root() + # Only `ug configure` and a launch with headers add the mod: another writer of this file + # drops entries it doesn't know, and re-adding them would rewrite it on every launch. + wants_mod = _header_mod_requested or bool(os.environ.get(UG_HEADERS_ENV_VAR, "").strip()) # Install the files before the settings name them, so a failed install never enables a # plugin that isn't there; a run that can't write as root leaves any earlier entries as is. - installed = _install_header_mod(header_mod_root) + installed = wants_mod and _install_header_mod(header_mod_root) desired_settings = _enable_header_mod(compose(existing), header_mod_root if installed else None) if header_mod_root is not None: owned_paths = [*owned_paths, *CLAUDE_HEADER_MOD_PATHS] diff --git a/src/ucode/cli.py b/src/ucode/cli.py index b75c10663..0ef5a83f8 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -4440,6 +4440,7 @@ def configure( raise RuntimeError("Use either --global or --user, not both.") scope = "user" if user_scope else "global" set_config_scope(scope) + claude_agent.request_header_mod() custom_oauth = _custom_oauth_config(client_id, redirect_url, scopes) if custom_oauth is not None and use_pat: raise RuntimeError("--client-id cannot be combined with --use-pat.") diff --git a/tests/test_claude_header_mod.py b/tests/test_claude_header_mod.py index 2b53697da..64d94b003 100644 --- a/tests/test_claude_header_mod.py +++ b/tests/test_claude_header_mod.py @@ -3,6 +3,7 @@ import copy import json import shutil +from types import SimpleNamespace from ucode import managed_files from ucode.agents import claude @@ -41,6 +42,12 @@ def _write(target, text): 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: SimpleNamespace(manifest={"claude": {}}) + ) + # As `ug configure` runs it. + monkeypatch.setattr(claude, "_header_mod_requested", True) + monkeypatch.delenv("UG_HEADERS", raising=False) return managed_path, root, writes @@ -66,6 +73,28 @@ def test_configure_installs_and_enables_mod_then_stays_put(tmp_path, monkeypatch assert len(writes) == count # unchanged files are never rewritten through sudo +def test_launch_adds_the_mod_only_with_headers_and_keeps_configures(tmp_path, monkeypatch): + managed_path, root, writes = _managed_env(tmp_path, monkeypatch) + monkeypatch.setattr(claude, "_header_mod_requested", False) + state = {"workspace": WS, "codex_models": []} + + # Another writer of the managed file drops entries it doesn't know; a plain launch mustn't + # add them back every time. + claude.write_tool_config(state, "databricks-claude-sonnet-4") + assert not root.exists() + assert "enabledPlugins" not in json.loads(managed_path.read_text()) + + monkeypatch.setenv("UG_HEADERS", "x-databricks-traffic-id: testenv://liteswap/unit") + claude.write_tool_config(state, "databricks-claude-sonnet-4") + assert claude.header_mod_active() + + # Once added, a plain launch leaves the mod as it is. + monkeypatch.delenv("UG_HEADERS") + count = len(writes) + claude.write_tool_config(state, "databricks-claude-sonnet-4") + assert claude.header_mod_active() and len(writes) == count + + def test_admin_prepend_list_is_kept_and_revert_restores_it(tmp_path, monkeypatch): managed_path, _root, _writes = _managed_env( tmp_path, monkeypatch, {"prependPlugins": ["acme-guard@acme"]}