Add a Claude Code mod so ug --header reaches Claude

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 |
 
 </details>
 

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/<pod>" codex
+ug --header "x-databricks-traffic-id: testenv://liteswap/<pod>" 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;

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.

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<void> {
+  // 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);
+  });
+}

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.")