diff --git a/AGENTS.md b/AGENTS.md index 8ea7fb4c7..bec3d7247 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. 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/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..fb9af9d2f 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, @@ -59,6 +61,7 @@ SettingsPassthrough, apply_settings_passthrough, check_skipped_managed_overrides, + claude_header_mod_root, config_scope, current_os, managed_file_conflicts, @@ -73,8 +76,10 @@ plan_settings_passthrough, read_managed_file, reconcile_managed_file, + reconcile_privileged_files, record_settings_passthrough, record_ug_picker, + remove_privileged_tree, revert_managed_file, warn_managed_write_skipped, warn_skipped_settings_passthrough, @@ -394,12 +399,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]: @@ -1755,6 +1782,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 +1848,111 @@ 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 + + +_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. + + 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 + try: + 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 + return True + + +def _remove_header_mod() -> None: + """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_tree(root) + except (OSError, RuntimeError, subprocess.CalledProcessError) as exc: + print_warning(f"Couldn't remove 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 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 + + 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 +2055,16 @@ 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() + # 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 = 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] _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: @@ -2311,6 +2453,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..c91a1d8af --- /dev/null +++ b/src/ucode/agents/claude_header_mod/register.ts @@ -0,0 +1,55 @@ +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. + +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 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]; + // 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'); +} + +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..0ef5a83f8 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: @@ -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/src/ucode/managed_files.py b/src/ucode/managed_files.py index 5e355649c..5724c5a1d 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,31 @@ 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], *, 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] + if changed and _sudo_may_prompt(): + _print_managed_write_permission(display) + for path, text in changed: + _privileged_replace(path, text) + + +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: """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..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 @@ -31,9 +33,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 } diff --git a/tests/README.md b/tests/README.md index ed05edb9f..6482276ff 100644 --- a/tests/README.md +++ b/tests/README.md @@ -48,6 +48,13 @@ 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`, 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 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..0c82b9502 --- /dev/null +++ b/tests/native_claude/header_mod_tests/register.test.ts @@ -0,0 +1,53 @@ +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('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' }); + + 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..64d94b003 --- /dev/null +++ b/tests/test_claude_header_mod.py @@ -0,0 +1,190 @@ +"""ug's Claude header mod: install beside the managed file, enable it there, and revert it.""" + +import copy +import json +import shutil +from types import SimpleNamespace + +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, "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) + 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") + 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 + + +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_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"]} + ) + + 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 + + +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(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() + + (root / "ug-headers/.claude-plugin/plugin.json").unlink() + assert not claude.header_mod_active() diff --git a/tests/test_managed_files.py b/tests/test_managed_files.py index 4783af283..4fdd976f0 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", @@ -1002,10 +1005,9 @@ def test_real_shell_creates_missing_file(self, tmp_path, parent_exists): 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 = [] - 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 = [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