diff --git a/docs/content/docs/proxy.mdx b/docs/content/docs/proxy.mdx index 78a358348..50c21b703 100644 --- a/docs/content/docs/proxy.mdx +++ b/docs/content/docs/proxy.mdx @@ -45,6 +45,7 @@ Telemetry is **off by default** (opt-in). Opt in with `HEADROOM_TELEMETRY=on` or | `--log-messages` | `false` | Store full request/response content for the live feed | | `--budget` | None | Daily budget limit in USD | | `--openai-api-url` | `https://api.openai.com` | Custom OpenAI API URL | +| `--provider-name` | Detected from `--openai-api-url` | Display name for the OpenAI-compatible upstream on the dashboard (e.g. `OpenRouter`). Well-known hosts (OpenRouter, Groq, Together, Azure OpenAI, …) are detected automatically; this overrides them. Routing and pricing are unaffected. | | `--anthropic-api-url` | Anthropic default | Custom Anthropic API URL | | `--gemini-api-url` | Gemini default | Custom Gemini API URL | | `--backend` | `anthropic` | Backend: `anthropic`, `bedrock`, `openrouter`, `anyllm`, or `litellm-` | diff --git a/headroom/cli/proxy.py b/headroom/cli/proxy.py index 48ab9140b..1b1ff08a1 100644 --- a/headroom/cli/proxy.py +++ b/headroom/cli/proxy.py @@ -803,6 +803,15 @@ def dashboard(port: int, no_open: bool) -> None: default=None, help="Custom OpenAI API URL for passthrough endpoints (env: OPENAI_TARGET_API_URL)", ) +@click.option( + "--provider-name", + default=None, + help=( + "Display name for the OpenAI-compatible upstream shown on the dashboard " + "(e.g. 'OpenRouter'). Overrides hostname detection from --openai-api-url. " + "Internal routing and pricing are unaffected." + ), +) @click.option( "--gemini-api-url", default=None, @@ -966,6 +975,7 @@ def proxy( anthropic_extra_headers: str | None, openai_extra_headers: str | None, openai_api_url: str | None, + provider_name: str | None, gemini_api_url: str | None, cloudcode_api_url: str | None, vertex_api_url: str | None, @@ -1155,6 +1165,7 @@ def proxy( anthropic_extra_headers=resolved_anthropic_extra_headers, openai_extra_headers=resolved_openai_extra_headers, openai_api_url=provider_api_overrides.openai, + provider_name=provider_name, gemini_api_url=provider_api_overrides.gemini, cloudcode_api_url=provider_api_overrides.cloudcode, vertex_api_url=provider_api_overrides.vertex, diff --git a/headroom/cli/wrap.py.orig b/headroom/cli/wrap.py.orig deleted file mode 100644 index 369ece852..000000000 --- a/headroom/cli/wrap.py.orig +++ /dev/null @@ -1,6379 +0,0 @@ -"""Wrap CLI commands to run through Headroom proxy. - -Usage: - headroom wrap claude # Start proxy + context tool + claude - headroom wrap copilot -- --model ... # Start proxy + launch GitHub Copilot CLI - headroom wrap codex # Start proxy + OpenAI Codex CLI - headroom wrap aider # Start proxy + aider - headroom wrap vibe # Start proxy + Mistral Vibe - headroom wrap cursor # Start proxy + print Cursor config instructions - headroom wrap openclaw # Install + configure OpenClaw plugin - headroom wrap claude --no-context-tool # Without CLI context-tool setup - headroom wrap claude --port 9999 # Custom proxy port - headroom wrap claude -- --model opus # Pass args to claude -""" - -from __future__ import annotations - -import errno -import importlib.util -import io -import json -import os -import re -import shutil -import signal -import socket -import subprocess -import sys -import tempfile -import time -import urllib.parse -from collections.abc import Callable -from contextlib import contextmanager -from pathlib import Path -from typing import Any, cast - -from headroom._subprocess import pid_alive, run - -# Fix Windows cp1252 encoding — box-drawing characters require UTF-8 -if sys.platform == "win32" and hasattr(sys.stdout, "buffer"): - if sys.stdout.encoding and sys.stdout.encoding.lower().replace("-", "") != "utf8": - sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace") - sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding="utf-8", errors="replace") - -import click - -from headroom import fsutil -from headroom._version import __version__ as _HEADROOM_VERSION -from headroom._version import normalize_release_version as _normalize_release_version -from headroom.agent_savings import ( - apply_agent_savings_env_defaults, -) -from headroom.copilot_auth import ( - has_oauth_auth, - resolve_client_bearer_token, - resolve_copilot_api_url, - resolve_subscription_bearer_token_details, -) -from headroom.providers.aider import build_launch_env as _build_aider_launch_env -from headroom.providers.claude import ( - REMOTE_CONTROL_BASE_URL_ENV, - TOOL_SEARCH_DEFAULT, - TOOL_SEARCH_ENV, - is_custom_anthropic_base_url, - remote_control_gate_message, -) -from headroom.providers.claude import ( - proxy_base_url as _claude_proxy_base_url, -) -from headroom.providers.codex import build_launch_env as _build_codex_launch_env -from headroom.providers.codex.install import codex_uses_chatgpt_auth -from headroom.providers.codex.threads import retag_to_headroom, retag_to_native -from headroom.providers.copilot import ( - build_launch_env as _build_copilot_launch_env, -) -from headroom.providers.copilot import ( - copilot_model_from_args as _copilot_model_from_args_impl, -) -from headroom.providers.copilot import ( - default_wire_api_for_model as _copilot_default_wire_api_for_model_impl, -) -from headroom.providers.copilot import ( - detect_running_proxy_backend as _copilot_detect_running_proxy_backend, -) -from headroom.providers.copilot import ( - is_auto_model as _is_auto_model, -) -from headroom.providers.copilot import ( - model_configured as _copilot_model_configured_impl, -) -from headroom.providers.copilot import ( - provider_key_source as _copilot_provider_key_source, -) -from headroom.providers.copilot import ( - query_proxy_config as _copilot_query_proxy_config, -) -from headroom.providers.copilot import ( - resolve_provider_type as _copilot_resolve_provider_type, -) -from headroom.providers.copilot import ( - strip_auto_model_args as _strip_auto_model_args, -) -from headroom.providers.copilot import ( - validate_configuration as _validate_copilot_configuration, -) -from headroom.providers.cursor import render_setup_lines as _render_cursor_setup_lines -from headroom.providers.mistral_vibe import build_launch_env as _build_mistral_vibe_launch_env -from headroom.providers.openclaw import ( - build_plugin_entry as _build_openclaw_plugin_entry_impl, -) -from headroom.providers.openclaw import ( - build_unwrap_entry as _build_openclaw_unwrap_entry_impl, -) -from headroom.providers.openclaw import ( - decode_entry_json as _decode_openclaw_entry_json_impl, -) -from headroom.providers.openclaw import ( - normalize_gateway_provider_ids as _normalize_openclaw_gateway_provider_ids_impl, -) -from headroom.providers.opencode import build_launch_env as _build_opencode_launch_env -from headroom.providers.opencode.config import ( - _MCP_MARKER_END, # noqa: F401 - _MCP_MARKER_START, - _PROVIDER_MARKER_END, # noqa: F401 - _PROVIDER_MARKER_START, - inject_opencode_provider_config, - opencode_config_paths, - snapshot_opencode_config_if_unwrapped, - strip_opencode_headroom_blocks, -) -from headroom.proxy.project_context import with_project_prefix as _with_project_prefix - -from .main import main - - -def _read_text(path: Path) -> str: - """Read a text file as UTF-8, falling back to the system locale encoding.""" - return fsutil.read_text(path) - - -def _write_text(path: Path, content: str) -> None: - """Write a text file as UTF-8 without translating line endings (preserves CRLF).""" - fsutil.write_text(path, content) - - -def _append_text(path: Path, content: str) -> None: - """Append to a text file as UTF-8 without translating line endings.""" - fsutil.append_text(path, content) - - -_CONTEXT_TOOL_ENV = "HEADROOM_CONTEXT_TOOL" -_CONTEXT_TOOL_RTK = "rtk" -_CONTEXT_TOOL_LEAN_CTX = "lean-ctx" -_VALID_CONTEXT_TOOLS = {_CONTEXT_TOOL_RTK, _CONTEXT_TOOL_LEAN_CTX} -_AGENT_SAVINGS_TARGET_AGENTS = {"claude", "codex", "cursor", "opencode"} -_WRAP_PROXY_TIMEOUT_ENV = "HEADROOM_WRAP_PROXY_TIMEOUT" -_WRAP_PROXY_TIMEOUT_DEFAULT_SECONDS = 45 -_WRAP_PROXY_TIMEOUT_ML_DEFAULT_SECONDS = 90 -_WRAP_PROXY_TIMEOUT_ML_MODULES = ("torch", "sentence_transformers", "spacy") - -# Issue #746: Claude Code disables on-demand tool loading (deferral) when -# ANTHROPIC_BASE_URL is a custom host and ENABLE_TOOL_SEARCH is unset, which -# inflates the local context window by tens of K tokens. Setting the env var -# when we launch Claude Code keeps deferral on. Default to "true" — defer the -# MCP/system tools for maximum context savings, matching native first-party -# behaviour (core built-ins like Read/Edit/Bash are never deferred by Claude -# Code, so the agent loop is unaffected). The key/default are shared with -# `init` and `install` via the Claude provider package to prevent drift. -_TOOL_SEARCH_ENV = TOOL_SEARCH_ENV -_TOOL_SEARCH_DEFAULT = TOOL_SEARCH_DEFAULT -_AGENT_SAVINGS_WRAP_AGENTS = {"claude", "codex", "cursor"} - -# 1M context window for `wrap claude` (#1158). Claude Code only sends the -# `context-1m` beta header — unlocking the 1M window for entitled subscription -# users — when the model id carries the `[1m]` suffix. Behind a custom -# ANTHROPIC_BASE_URL (the proxy) its `/model` picker selection does not survive, -# so `--1m` forces the suffix via ANTHROPIC_MODEL on the launched process. -_ANTHROPIC_MODEL_ENV = "ANTHROPIC_MODEL" -_CONTEXT_1M_SUFFIX = "[1m]" -# Only used when no model is otherwise selected (no ANTHROPIC_MODEL set). The -# current default Opus; the suffix logic preserves any model the user did set. -_DEFAULT_1M_MODEL = "claude-opus-4-8" - - -def _resolve_1m_model(current: str | None) -> str: - """Return the model id that makes Claude Code request the 1M window (#1158). - - Preserves a model the user already selected via ``ANTHROPIC_MODEL`` (only - appending the ``[1m]`` suffix when missing); falls back to the default Opus - when none is set. Idempotent — a value already ending in ``[1m]`` is - returned unchanged. - """ - base = (current or "").strip() or _DEFAULT_1M_MODEL - return base if base.endswith(_CONTEXT_1M_SUFFIX) else f"{base}{_CONTEXT_1M_SUFFIX}" - - -def _normalize_tool_search_mode(value: str) -> str: - """Validate an ``ENABLE_TOOL_SEARCH`` value and return it normalized. - - Mirrors the values Claude Code accepts: truthy (``true``/``1``/``yes``/ - ``on``), falsy (``false``/``0``/``no``/``off``), ``auto``, or ``auto:N`` - where ``N`` is 0-100. Raises :class:`click.ClickException` on anything else - so a typo fails loudly instead of silently leaving deferral off. - """ - normalized = value.strip().lower() - if normalized in {"true", "1", "yes", "on", "false", "0", "no", "off", "auto"}: - return normalized - if normalized.startswith("auto:"): - suffix = normalized[len("auto:") :] - if suffix.isdigit() and 0 <= int(suffix) <= 100: - return normalized - raise click.ClickException( - f"--tool-search must be one of: true, false, auto, auto:N (N 0-100); got {value!r}" - ) - - -def _configure_tool_search_env(env: dict[str, str], flag_value: str | None) -> str | None: - """Set ``ENABLE_TOOL_SEARCH`` in ``env`` so Claude Code keeps deferring tools. - - Precedence: - - 1. explicit ``--tool-search`` flag — wins (the user asked for it on the CLI), - 2. a pre-existing ``ENABLE_TOOL_SEARCH`` in the environment — respected and - left untouched (the user's own Claude Code knob), - 3. the built-in default (``true``). - - Returns the value written, or ``None`` when an existing environment value - was deliberately left in place. - """ - if flag_value is not None: - value = _normalize_tool_search_mode(flag_value) - env[_TOOL_SEARCH_ENV] = value - return value - # An empty / whitespace value counts as unset: Claude Code treats an empty - # ENABLE_TOOL_SEARCH as absent (so deferral would stay off), so we override - # it with the default rather than forwarding a no-op value. - existing = env.get(_TOOL_SEARCH_ENV) - if existing is not None and existing.strip(): - return None - env[_TOOL_SEARCH_ENV] = _TOOL_SEARCH_DEFAULT - return _TOOL_SEARCH_DEFAULT - - -def _live_wrap_module() -> Any: - """Return the current live wrap module instance.""" - return cast(Any, sys.modules[__name__]) - - -def _selected_context_tool() -> str: - """Return the configured CLI context tool. - - RTK remains the default for backward compatibility. Set - ``HEADROOM_CONTEXT_TOOL=lean-ctx`` to let lean-ctx configure the supported - coding agent instead. - """ - - raw = os.environ.get(_CONTEXT_TOOL_ENV, "").strip().lower().replace("_", "-") - if not raw: - return _CONTEXT_TOOL_RTK - if raw == "leanctx": - raw = _CONTEXT_TOOL_LEAN_CTX - if raw not in _VALID_CONTEXT_TOOLS: - raise click.ClickException( - f"{_CONTEXT_TOOL_ENV} must be one of: {', '.join(sorted(_VALID_CONTEXT_TOOLS))}" - ) - return raw - - -def _module_available(module_name: str) -> bool: - """Return whether an optional module is installed without importing it.""" - - try: - return importlib.util.find_spec(module_name) is not None - except (ImportError, ModuleNotFoundError, ValueError): - return False - - -def _ml_wrap_extras_detected() -> bool: - """Detect slow optional ML stacks without triggering their import cost.""" - - return any(_module_available(module_name) for module_name in _WRAP_PROXY_TIMEOUT_ML_MODULES) - - -def _wrap_agent_savings_profile(agent_type: str) -> str | None: - """Return the savings profile required for agent wrappers, if any.""" - - if agent_type not in _AGENT_SAVINGS_WRAP_AGENTS: - return None - return os.environ.get("HEADROOM_SAVINGS_PROFILE") or None - - -def _default_wrap_proxy_timeout_seconds() -> int: - """Return the default wrap proxy startup timeout for this environment.""" - - if _ml_wrap_extras_detected(): - return _WRAP_PROXY_TIMEOUT_ML_DEFAULT_SECONDS - return _WRAP_PROXY_TIMEOUT_DEFAULT_SECONDS - - -def _resolve_wrap_proxy_timeout_seconds() -> int: - """Resolve the wrap proxy readiness timeout from env or defaults.""" - - raw = os.environ.get(_WRAP_PROXY_TIMEOUT_ENV, "").strip() - if not raw: - return _default_wrap_proxy_timeout_seconds() - - try: - timeout_seconds = int(raw) - except ValueError: - raise RuntimeError( - f"{_WRAP_PROXY_TIMEOUT_ENV} must be a positive integer number of seconds (got {raw!r})" - ) from None - if timeout_seconds <= 0: - raise RuntimeError( - f"{_WRAP_PROXY_TIMEOUT_ENV} must be a positive integer number of seconds (got {raw!r})" - ) - return timeout_seconds - - -def _print_telemetry_notice() -> None: - """Print a telemetry notice when anonymous telemetry is enabled. - - Respects the HEADROOM_TELEMETRY and HEADROOM_TELEMETRY_WARN feature flags. - Does nothing when telemetry or warnings are disabled. - """ - from headroom.telemetry.beacon import format_telemetry_notice - - notice = format_telemetry_notice(prefix=" ") - if notice: - click.echo(notice) - - -# Proxy health check (reused from evals/suite_runner.py pattern) - - -def _check_proxy(port: int) -> bool: - """Check if Headroom proxy is running on given port.""" - try: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.settimeout(1) - s.connect(("127.0.0.1", port)) - return True - except (TimeoutError, ConnectionRefusedError, OSError): - return False - - -def _port_bind_error(port: int) -> OSError | None: - """Return the bind error for a local proxy port, or None when it is usable.""" - try: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("127.0.0.1", port)) - except OSError as exc: - return exc - except OverflowError: - return OSError(errno.EADDRNOTAVAIL, f"Port {port} out of range (0-65535)") - return None - - -def _find_available_port(start_port: int, max_attempts: int = 100) -> int: - """Find first available port >= start_port via socket.bind probe. - - Skips ports with EADDRINUSE (busy) and EACCES (reserved on Windows, - privileged on Linux) — both indicate the port can't be bound here. - Other OS errors (EADDRNOTAVAIL) propagate immediately. - Raises RuntimeError when no port is found in range. - """ - end_port = min(start_port + max_attempts, 65536) - for port in range(start_port, end_port): - error = _port_bind_error(port) - if error is None: - return port - if error.errno not in (errno.EADDRINUSE, errno.EACCES): - raise error - raise RuntimeError(f"No available port found in range {start_port}-{end_port - 1}") - - -def _get_log_path() -> Path: - """Get path for proxy log file.""" - from headroom import paths as _paths - - log_dir = _paths.log_dir() - log_dir.mkdir(parents=True, exist_ok=True) - return log_dir / "proxy.log" - - -def _get_proxy_stdio_log_path() -> Path: - """Get path for dedicated proxy stdio capture.""" - return _get_log_path().with_name("proxy-stdio.log") - - -def _start_proxy( - port: int, - *, - learn: bool = False, - memory: bool = False, - agent_type: str = "unknown", - code_graph: bool = False, - backend: str | None = None, - anyllm_provider: str | None = None, - region: str | None = None, - openai_api_url: str | None = None, - anthropic_api_url: str | None = None, - vertex_api_url: str | None = None, - clear_vertex_api_url: bool = False, - copilot_api_token: str | None = None, -) -> subprocess.Popen: - """Start Headroom proxy as a background subprocess. - - Stdout and stderr are written to a dedicated sibling file, usually - `~/.headroom/logs/proxy-stdio.log`, to avoid pipe deadlock risk without - competing with the rotating `proxy.log` runtime log. - - The caller is responsible for ensuring *port* is available - (see ``_find_available_port``). - """ - - cmd = [sys.executable, "-m", "headroom.cli", "proxy", "--port", str(port)] - - # Forward HEADROOM_MODE env var so the proxy respects the user's mode choice - headroom_mode = os.environ.get("HEADROOM_MODE") - if headroom_mode: - cmd.extend(["--mode", headroom_mode]) - - # Forward --learn flag to proxy subprocess - if learn: - cmd.append("--learn") - - # Forward --memory flag to proxy subprocess - if memory: - cmd.append("--memory") - - # Forward --code-graph flag to proxy subprocess (live file watcher) - if code_graph: - cmd.append("--code-graph") - - # Forward backend configuration to proxy subprocess - _backend = backend or os.environ.get("HEADROOM_BACKEND") - if _backend: - cmd.extend(["--backend", _backend]) - - _anyllm = anyllm_provider or os.environ.get("HEADROOM_ANYLLM_PROVIDER") - if _anyllm: - cmd.extend(["--anyllm-provider", _anyllm]) - - _region = region or os.environ.get("HEADROOM_REGION") - if _region: - cmd.extend(["--region", _region]) - - if openai_api_url: - cmd.extend(["--openai-api-url", openai_api_url]) - - if anthropic_api_url: - cmd.extend(["--anthropic-api-url", anthropic_api_url]) - - if vertex_api_url: - cmd.extend(["--vertex-api-url", vertex_api_url]) - - timeout_seconds = _resolve_wrap_proxy_timeout_seconds() - log_path = _get_log_path() - stdio_log_path = _get_proxy_stdio_log_path() - stdio_log_file = open(stdio_log_path, "a", encoding="utf-8") # noqa: SIM115 - - # Ensure proxy subprocess uses UTF-8 (Windows defaults to cp1252) - proxy_env = os.environ.copy() - proxy_env["PYTHONIOENCODING"] = "utf-8" - # Vertex AI RST_STREAMs HTTP/2 connections (error_code:2). Force HTTP/1.1 - # when wrapping a Vertex-mode client so upstream requests succeed. - if os.environ.get("CLAUDE_CODE_USE_VERTEX") or os.environ.get("ANTHROPIC_VERTEX_PROJECT_ID"): - proxy_env.setdefault("HEADROOM_HTTP2", "false") - # Tell the proxy which agent is being wrapped (for traffic learning output) - if agent_type != "unknown": - proxy_env["HEADROOM_AGENT_TYPE"] = agent_type - proxy_env.setdefault("HEADROOM_STACK", f"wrap_{agent_type}") - savings_profile = _wrap_agent_savings_profile(agent_type) - if savings_profile is not None: - apply_agent_savings_env_defaults(proxy_env, savings_profile) - if openai_api_url: - proxy_env["OPENAI_TARGET_API_URL"] = openai_api_url - if anthropic_api_url: - proxy_env["ANTHROPIC_TARGET_API_URL"] = anthropic_api_url - if clear_vertex_api_url: - proxy_env.pop("VERTEX_TARGET_API_URL", None) - if vertex_api_url: - proxy_env["VERTEX_TARGET_API_URL"] = vertex_api_url - # Pin the wrapper-validated Copilot token for this proxy instance only. - # Injected into the subprocess env here (not the parent's os.environ) so it - # never leaks into shared state. The proxy's CopilotTokenProvider honours - # GITHUB_COPILOT_API_TOKEN directly, making upstream auth deterministic. - if copilot_api_token: - proxy_env["GITHUB_COPILOT_API_TOKEN"] = copilot_api_token - if openai_api_url: - proxy_env["GITHUB_COPILOT_API_URL"] = openai_api_url - - # Detach the proxy from the launching console on Windows so an ungraceful - # close of the owning agent (closing the terminal window, taskkill, or a - # crash) cannot tree-kill the shared proxy out from under other live - # clients. Without this the proxy stays in the owner's console + Job - # object; closing that window terminates the whole tree, bypassing the - # marker-based reference counting in ``_make_cleanup`` and breaking every - # other ``headroom wrap`` instance routed through the same port. - # CREATE_NO_WINDOW — give the proxy its OWN, invisible console. - # A separate console means the parent's - # CTRL_CLOSE_EVENT never reaches it, and no - # stray console window pops up. DETACHED_PROCESS - # also isolates the console, but for a console - # subsystem exe (python.exe) it leaves the proxy - # consoleless and Windows surfaces a visible - # console window — closing that window killed - # the proxy, defeating the whole point. - # CREATE_NEW_PROCESS_GROUP — isolate from the parent's Ctrl-C - # CREATE_BREAKAWAY_FROM_JOB— survive Job kill-on-close (Windows Terminal, - # VS Code integrated terminal, conhost) - # CREATE_NO_WINDOW / DETACHED_PROCESS / CREATE_NEW_CONSOLE are mutually - # exclusive — pick exactly one. On POSIX, ``start_new_session`` already - # detaches via setsid(). ``sys.platform == "win32"`` (not ``os.name == - # "nt"``) so mypy narrows the platform and resolves the Windows-only - # ``subprocess`` constants below. - _CREATE_BREAKAWAY_FROM_JOB = 0x01000000 - creationflags = 0 - if sys.platform == "win32": - creationflags = ( - subprocess.CREATE_NO_WINDOW - | subprocess.CREATE_NEW_PROCESS_GROUP - | _CREATE_BREAKAWAY_FROM_JOB - ) - - popen_kwargs: dict[str, Any] = { - "stdout": stdio_log_file, - "stderr": stdio_log_file, - "env": proxy_env, - "start_new_session": os.name == "posix", - "creationflags": creationflags, - } - # Close the parent's copy of the stdio log handle on every exit path, - # including when BOTH spawn attempts raise. The child keeps its own - # inherited duplicate, so closing here never starves the proxy's logging. - try: - try: - proc = subprocess.Popen(cmd, **popen_kwargs) - except OSError: - # The launcher's Job object forbids breakaway. Retry without that flag; - # CREATE_NO_WINDOW still spares the proxy from console-close events. - if sys.platform == "win32": - popen_kwargs["creationflags"] = creationflags & ~_CREATE_BREAKAWAY_FROM_JOB - proc = subprocess.Popen(cmd, **popen_kwargs) - - # Wait for proxy to be ready. - # ML components (Kompress, Magika, Tree-sitter) load synchronously before - # uvicorn binds the port. On slower machines this can take 20-30 seconds. - for _i in range(timeout_seconds): - time.sleep(1) - if _check_proxy(port): - click.echo(f" Logs: {log_path}") - return proc - # Check if process died - if proc.poll() is not None: - # Read last few lines of log for error context - try: - tail = _read_text(stdio_log_path)[-500:] - except Exception: - tail = "(no log output)" - raise RuntimeError(f"Proxy exited with code {proc.returncode}: {tail}") - - proc.kill() - raise RuntimeError( - f"Proxy failed to start on port {port} within {timeout_seconds} seconds. " - f"Set {_WRAP_PROXY_TIMEOUT_ENV} to a larger number of seconds for slow startup." - ) - finally: - stdio_log_file.close() - - -def _setup_rtk(verbose: bool = False) -> Path | None: - """Ensure rtk is installed and hooks are registered.""" - from headroom.rtk import get_rtk_path - from headroom.rtk.installer import ensure_rtk, register_claude_hooks - - rtk_path = get_rtk_path() - - if rtk_path: - if verbose: - click.echo(f" rtk found at {rtk_path}") - else: - click.echo(" Downloading rtk (Rust Token Killer)...") - rtk_path = ensure_rtk() - if rtk_path: - click.echo(f" rtk installed at {rtk_path}") - else: - click.echo(" rtk download failed — continuing without it") - return None - - # Register hooks (idempotent) - if register_claude_hooks(rtk_path): - if verbose: - click.echo(" rtk hooks registered in Claude Code") - try: - linked = _ensure_rtk_on_path(rtk_path) - if linked and verbose: - click.echo(f" rtk linked onto PATH at {linked}") - except Exception as e: - if verbose: - click.echo(f" rtk PATH link skipped: {e}") - else: - click.echo(" rtk hook registration failed — continuing without it") - - return rtk_path - - -def _ensure_rtk_on_path(rtk_path: Path, path_dirs: list[str] | None = None) -> Path | None: - """Make the Headroom-managed rtk resolvable as a bare ``rtk`` on PATH. - - ``rtk init --global --auto-patch`` writes ``~/.claude/hooks/rtk-rewrite.sh``, - and ``rtk rewrite`` emits a bare ``rtk`` token at runtime that the hook feeds - back to the shell — so bare ``rtk`` has to resolve on PATH regardless of the - hook's contents. Since ``~/.headroom/bin`` (where Headroom installs rtk) is - not on PATH by default, that lookup fails and compression silently never - runs (issue #487). - - An earlier fix rewrote the generated hook to hard-code rtk's absolute path. - That mutates the hook *after* ``rtk init`` bakes in its expected SHA-256, so - rtk's integrity guard rejects it (``hook integrity check FAILED … RTK will - not execute``) and only absolutizes the hook's own ``rtk`` call — not the - bare ``rtk`` that ``rtk rewrite`` emits at runtime (issue #1631). Instead, - leave the canonical hook untouched and link the managed binary into a PATH - directory so bare ``rtk`` resolves. - - Idempotent and conservative: - * no-op if a ``rtk`` already resolves on PATH (managed or system); - * no-op on Windows (symlinks need privilege; hooks resolve differently); - * only creates/refreshes a symlink Headroom owns — never clobbers an - existing real file or foreign binary. - - Returns the link path that was created or already correct, else ``None``. - """ - if sys.platform == "win32": - return None - - # A bare `rtk` already resolves — the hook will find it, nothing to do. - if shutil.which("rtk"): - return None - - if path_dirs is None: - path_dirs = os.environ.get("PATH", "").split(os.pathsep) - - preferred = Path.home() / ".local" / "bin" - - # Prefer ~/.local/bin (conventionally on PATH), then any other PATH dir. - ordered: list[Path] = [] - if str(preferred) in path_dirs: - ordered.append(preferred) - for entry in path_dirs: - if not entry: - continue - candidate = Path(entry) - if candidate not in ordered: - ordered.append(candidate) - - target = rtk_path.resolve() - - for target_dir in ordered: - link = target_dir / "rtk" - try: - # Existing correct link — done. - if link.is_symlink() and link.resolve() == target: - return link - # Never clobber a real file or a link pointing elsewhere. - if link.exists() or link.is_symlink(): - continue - # Create ~/.local/bin on demand; other PATH dirs must already exist. - if target_dir == preferred: - target_dir.mkdir(parents=True, exist_ok=True) - if not target_dir.is_dir() or not os.access(target_dir, os.W_OK): - continue - link.symlink_to(target) - return link - except OSError: - continue - - return None - - -def _setup_lean_ctx_agent(agent: str, verbose: bool = False) -> Path | None: - """Run lean-ctx agent setup for the requested coding tool.""" - - from headroom.lean_ctx import get_lean_ctx_path - from headroom.lean_ctx.installer import ensure_lean_ctx - - lean_ctx = get_lean_ctx_path() - if not lean_ctx: - click.echo(" Downloading lean-ctx...") - lean_ctx = ensure_lean_ctx() - if not lean_ctx: - click.echo(" lean-ctx download failed — continuing without it") - return None - - try: - with tempfile.TemporaryDirectory(prefix="headroom-lean-ctx-") as setup_cwd: - # lean-ctx writes project-local files when initialized from a git - # checkout. Run from a non-project directory so setup is limited to - # home-scoped agent config such as ~/.codex or ~/.claude. - result = run( - [str(lean_ctx), "init", "--agent", agent], - capture_output=True, - text=True, - timeout=30, - cwd=setup_cwd, - ) - except Exception as e: - click.echo(f" lean-ctx setup failed — continuing without it: {e}") - return None - - if result.returncode != 0: - detail = (result.stderr or result.stdout).strip() - suffix = f": {detail}" if detail else "" - click.echo(f" lean-ctx setup failed — continuing without it{suffix}") - return None - - if verbose: - detail = result.stdout.strip() - if detail: - click.echo(f" lean-ctx configured for {agent}: {detail}") - else: - click.echo(f" lean-ctx configured for {agent}") - return lean_ctx - - -# Hook-command markers Headroom manages in Claude settings.json. unwrap drops -# any hook entry whose command contains one of these. -_HEADROOM_HOOK_MARKERS = ("rtk-rewrite", "headroom-init-claude") - -# Env vars Headroom's init/wrap inject into Claude settings.json; unwrap removes -# them. ENABLE_TOOL_SEARCH keeps Claude Code's tool deferral on behind the proxy -# (GH #746), paired with init/wrap setting it. -_HEADROOM_ENV_KEYS = ("ANTHROPIC_BASE_URL", "ENABLE_TOOL_SEARCH") - - -def _remove_claude_rtk_hooks(settings_path: Path | None = None) -> bool: - """Remove Headroom-managed entries from Claude settings.json. - - Reverses what ``headroom init claude`` and ``rtk init --auto-patch`` add: - * PreToolUse / SessionStart hooks whose command contains a Headroom marker - (``rtk-rewrite`` or ``headroom-init-claude``), and - * the ``ANTHROPIC_BASE_URL`` proxy-routing env var. - Unrelated settings and user-authored hooks are left untouched. (Previously - this only matched ``rtk-rewrite`` and returned early when no hooks existed, - so init's env + hooks survived unwrap.) - """ - - path = settings_path or (Path.home() / ".claude" / "settings.json") - if not path.exists(): - return False - - try: - payload = json.loads(_read_text(path)) - except (OSError, json.JSONDecodeError): - return False - if not isinstance(payload, dict): - return False - - changed = False - - hooks = payload.get("hooks") - if isinstance(hooks, dict): - for event, entries in list(hooks.items()): - if not isinstance(entries, list): - continue - retained_entries: list[Any] = [] - for entry in entries: - if not isinstance(entry, dict): - retained_entries.append(entry) - continue - hook_items = entry.get("hooks") - if not isinstance(hook_items, list): - retained_entries.append(entry) - continue - retained_hooks = [ - item - for item in hook_items - if not ( - isinstance(item, dict) - and any( - marker in str(item.get("command", "")).lower() - for marker in _HEADROOM_HOOK_MARKERS - ) - ) - ] - if len(retained_hooks) != len(hook_items): - changed = True - if retained_hooks: - retained_entries.append({**entry, "hooks": retained_hooks}) - elif len(retained_hooks) == len(hook_items): - retained_entries.append(entry) - else: - changed = True - if retained_entries: - hooks[event] = retained_entries - else: - del hooks[event] - changed = True - - if hooks: - payload["hooks"] = hooks - else: - payload.pop("hooks", None) - - # Remove the proxy-routing env that init/wrap injected (ANTHROPIC_BASE_URL and - # ENABLE_TOOL_SEARCH), even when no hooks remain (the early-return bug skipped - # this). List-comp, not any(), so every key is popped (no short-circuit). - env = payload.get("env") - if isinstance(env, dict): - removed_keys = [k for k in _HEADROOM_ENV_KEYS if env.pop(k, None) is not None] - if removed_keys: - changed = True - if env: - payload["env"] = env - else: - payload.pop("env", None) - - if not changed: - return False - - _write_text(path, json.dumps(payload, indent=2) + "\n") - return True - - -def _foundry_upstream_url(resource: str) -> str: - """Derive the Azure AI Foundry endpoint URL from a resource name. - - When CLAUDE_CODE_USE_FOUNDRY=1 is set, Claude Code routes requests to the - Azure AI Services endpoint it constructs from ANTHROPIC_FOUNDRY_RESOURCE. - If ANTHROPIC_FOUNDRY_BASE_URL is not already set in the environment, - we derive it here so the proxy knows where to forward compressed requests. - - Azure AI Foundry (AI Services) hosts the Anthropic-format Claude API at: - https://{resource}.services.ai.azure.com/anthropic - This matches the URL Claude Code constructs internally from ANTHROPIC_FOUNDRY_RESOURCE, - and what ANTHROPIC_FOUNDRY_BASE_URL must point to for the Anthropic SDK to reach Claude. - """ - return f"https://{resource.strip()}.services.ai.azure.com/anthropic" - - -def _foundry_proxy_url(proxy_url: str) -> str: - """Return the local proxy URL that Claude Code should use in Foundry mode. - - ANTHROPIC_FOUNDRY_BASE_URL is the full base URL the Anthropic SDK appends - /v1/messages to, so it must include the /anthropic path component to match - the Azure AI Foundry endpoint structure. _claude_proxy_base_url() returns - the bare http://127.0.0.1: — this helper appends /anthropic so the - proxy URL Claude Code receives mirrors the real Foundry URL shape. - """ - return proxy_url.rstrip("/") + "/anthropic" - - -def _vertex_target_api_url_from_claude_env(proxy_url: str) -> str | None: - """Return the Vertex upstream that the proxy should use for Claude Code.""" - explicit_target = os.environ.get("VERTEX_TARGET_API_URL", "").strip() - if explicit_target: - return ( - None - if _normalize_proxy_api_url(explicit_target) == _normalize_proxy_api_url(proxy_url) - else explicit_target - ) - - vertex_url = os.environ.get("ANTHROPIC_VERTEX_BASE_URL", "").strip() - if not vertex_url: - return None - - from headroom.providers.registry import DEFAULT_VERTEX_API_URL - - normalized_vertex_url = _normalize_proxy_api_url(vertex_url) - if normalized_vertex_url == _normalize_proxy_api_url(DEFAULT_VERTEX_API_URL): - return None - if normalized_vertex_url == _normalize_proxy_api_url(proxy_url): - return None - return vertex_url - - -def _claude_wrap_base_url_env_key(*, foundry_mode: bool = False, vertex_mode: bool = False) -> str: - if vertex_mode: - return "ANTHROPIC_VERTEX_BASE_URL" - if foundry_mode: - return "ANTHROPIC_FOUNDRY_BASE_URL" - return "ANTHROPIC_BASE_URL" - - -def _wrap_marker_path(settings_path: Path) -> Path: - """Sidecar marker path for a given settings.local.json path. - - Kept out of settings.local.json itself so Headroom's own bookkeeping never - shows up as a stray key inside a file Claude Code's config loader parses. - """ - return settings_path.parent / ".headroom_wrap_marker.json" - - -def _write_wrap_marker(settings_path: Path, *, port: int, key: str, previous: str | None) -> None: - """Best-effort record of which (pid, port, key) wrote the base_url entry. - - Lets a later wrap/doctor/unwrap invocation tell a stale leftover (writer - process is dead or its PID was recycled) from a still-live wrap session, - and recover the true prior value (issue #1768) instead of guessing. - """ - try: - ident = _proc_identity(os.getpid()) - payload = { - "pid": os.getpid(), - "start_src": ident[0] if ident else None, - "start_time": ident[1] if ident else None, - "port": port, - "key": key, - "previous": previous, - } - _write_text(_wrap_marker_path(settings_path), json.dumps(payload)) - except OSError: - pass - - -def _read_wrap_marker(settings_path: Path) -> dict[str, Any] | None: - marker = _wrap_marker_path(settings_path) - try: - rec = json.loads(_read_text(marker)) - except (OSError, ValueError): - return None - return rec if isinstance(rec, dict) else None - - -def _wrap_marker_is_stale(marker: dict[str, Any]) -> bool: - """True if ``marker`` describes a writer that is provably gone. - - Missing/invalid pid, a dead pid, or a live pid whose recorded identity no - longer matches (PID reuse) all count as stale — the entry it describes was - left behind by a wrap session that no longer exists. - """ - pid = marker.get("pid") - if not isinstance(pid, int): - return True - if not _pid_alive(pid): - return True - return _identity_mismatch(marker.get("start_src"), marker.get("start_time"), pid) - - -def _clear_wrap_marker(settings_path: Path, *, key: str) -> None: - marker = _read_wrap_marker(settings_path) - if marker is not None and marker.get("key") == key: - _wrap_marker_path(settings_path).unlink(missing_ok=True) - - -def _check_and_clear_stale_wrap_marker(settings_path: Path, *, key: str) -> str | None: - """If a stale wrap marker for ``key`` exists, restore its recorded prior - value and clear the marker. Returns the restored value, or None if there - was nothing stale to clean up. - - Called before writing a fresh base_url entry so a crashed wrap session's - leftover doesn't get treated as this session's own state to restore later. - """ - marker = _read_wrap_marker(settings_path) - if marker is None or marker.get("key") != key or not _wrap_marker_is_stale(marker): - return None - previous = marker.get("previous") - click.echo( - f"headroom: clearing stale {key} left by crashed wrap session (pid {marker.get('pid')})", - err=True, - ) - _restore_claude_wrap_base_url(previous, settings_path=settings_path, _key_override=key) - return previous - - -def _write_claude_wrap_base_url( - proxy_url: str, - *, - foundry_mode: bool = False, - vertex_mode: bool = False, - settings_path: Path | None = None, - port: int | None = None, -) -> str | None: - """Persist proxy URL into project-local settings env key for daemon child inheritance. - - Claude Code's cc-daemon pre-forks conversation workers using spawn (not - fork), so those workers read settings.json fresh rather than inheriting - the daemon's environment. Writing the mode-specific Claude base URL env - key into the project-local settings file (.claude/settings.local.json in - cwd) ensures every new conversation — including those started after the - initial launch — routes through the Headroom proxy without touching the - global user settings file or affecting sessions in other projects. Returns - the previous value so the caller can restore it on exit (issue #951). - - When ``port`` is given, also stamps a sidecar marker recording this - process's identity and the previous value, so a later crash can be - detected and self-healed (issue #1768). - """ - path = settings_path or (Path.cwd() / ".claude" / "settings.local.json") - payload: dict[str, Any] = {} - if path.exists(): - try: - payload = json.loads(_read_text(path)) - except (OSError, json.JSONDecodeError): - payload = {} - if not isinstance(payload, dict): - payload = {} - env_map = dict(payload.get("env") or {}) if isinstance(payload.get("env"), dict) else {} - key = _claude_wrap_base_url_env_key(foundry_mode=foundry_mode, vertex_mode=vertex_mode) - previous = env_map.get(key) - env_map[key] = proxy_url - payload["env"] = env_map - path.parent.mkdir(parents=True, exist_ok=True) - _write_text(path, json.dumps(payload, indent=2) + "\n") - if port is not None: - _write_wrap_marker(path, port=port, key=key, previous=previous) - return previous - - -def _restore_claude_wrap_base_url( - previous: str | None, - *, - foundry_mode: bool = False, - vertex_mode: bool = False, - settings_path: Path | None = None, - _key_override: str | None = None, -) -> None: - """Restore (or remove) the env key written by _write_claude_wrap_base_url. - - Called in both the wrap-session finally block and unwrap_claude so the - project-local settings entry is never left pointing at a dead proxy. When - ``previous`` is None the key is removed; when it has a value it is - restored — preserving any URL the project already had set. Also clears - this key's sidecar wrap marker, if any (issue #1768). - """ - path = settings_path or (Path.cwd() / ".claude" / "settings.local.json") - key = _key_override or _claude_wrap_base_url_env_key( - foundry_mode=foundry_mode, vertex_mode=vertex_mode - ) - if not path.exists(): - _clear_wrap_marker(path, key=key) - return - try: - payload = json.loads(_read_text(path)) - except (OSError, json.JSONDecodeError): - return - if not isinstance(payload, dict): - return - env_map = payload.get("env") - if not isinstance(env_map, dict): - return - if previous is None: - if key not in env_map: - _clear_wrap_marker(path, key=key) - return - del env_map[key] - if env_map: - payload["env"] = env_map - else: - payload.pop("env", None) - else: - env_map[key] = previous - payload["env"] = env_map - if payload: - _write_text(path, json.dumps(payload, indent=2) + "\n") - else: - path.unlink(missing_ok=True) - _clear_wrap_marker(path, key=key) - - -def _setup_headroom_mcp( - registrar: Any, port: int, *, verbose: bool = False, force: bool = False -) -> None: - """Register the headroom MCP server with the given agent (idempotent). - - The proxy compresses tool_result payloads and emits ``[Retrieve more: - hash=…]`` markers. Without this registration those markers point at - nothing — the agent has no ``headroom_retrieve`` tool to call. - - Generic across registrars: ``ClaudeRegistrar``, ``CodexRegistrar``, and - any future agent registrar all flow through the same setup path. - """ - from headroom.mcp_registry import build_headroom_spec, format_result - - if not registrar.detect(): - if verbose: - click.echo(f" MCP retrieve tool: {registrar.display_name} not detected — skipping") - return - - proxy_url = f"http://127.0.0.1:{port}" - spec = build_headroom_spec(proxy_url) - result = registrar.register_server(spec, force=force) - - line = format_result( - registrar.name, - result, - label="MCP retrieve tool", - verbose=verbose, - overwrite_hint=f"headroom mcp install --proxy-url {proxy_url} --force", - restart_hint=f"restart {registrar.display_name} if it was already running", - ) - if line is not None: - click.echo(line) - - -def _setup_serena_mcp( - registrar: Any, *, context: str, verbose: bool = False, force: bool = False -) -> None: - """Register Serena MCP with the given agent (idempotent). - - A prior ``headroom wrap`` may have persisted a Serena entry built from an - older spec — e.g. before ``--open-web-dashboard False`` was added to - suppress the dashboard popup (#1003). ``register_server`` returns - ``MISMATCH`` and refuses to overwrite a differing entry unless forced, so - on its own a re-wrap leaves already-wrapped users stuck on the stale spec - (and the popup) forever. When the ledger proves the entry currently in the - config is one Headroom installed, force-update it to the current spec. A - user-managed Serena (absent from our ledger) is left untouched and the - mismatch is reported as before. - """ - from headroom.mcp_registry import build_serena_spec, format_result - from headroom.mcp_registry.base import RegisterStatus - from headroom.mcp_registry.ledger import headroom_installed_matching, record_install - - if not registrar.detect(): - if verbose: - click.echo(f" Serena MCP: {registrar.display_name} not detected — skipping") - return - - if shutil.which("uvx") is None: - click.echo(" Serena MCP: uvx not found — install uv/uvx to enable Serena; skipping") - return - - spec = build_serena_spec(context) - result = registrar.register_server(spec, force=force) - - # Migrate a stale Headroom-installed entry. register_server won't overwrite - # a differing spec without force, so an older Headroom Serena entry would - # otherwise persist across re-wraps. Force-update it only when the ledger - # proves Headroom installed the entry that's currently on disk — never a - # user-managed Serena. - if ( - result.status == RegisterStatus.MISMATCH - and not force - and headroom_installed_matching(registrar.name, registrar.get_server("serena")) - ): - result = registrar.register_server(spec, force=True) - if result.status == RegisterStatus.REGISTERED: - click.echo(" Serena MCP: migrated previously-installed entry to current spec") - - if result.status == RegisterStatus.REGISTERED: - record_install(registrar.name, spec) - - line = format_result( - registrar.name, - result, - label="Serena MCP", - verbose=verbose, - overwrite_hint="update or remove the existing serena MCP entry, then rerun headroom wrap", - restart_hint=f"restart {registrar.display_name} if it was already running", - ) - if line is not None: - click.echo(line) - - -def _remove_headroom_installed_serena_mcp(registrar: Any) -> str: - """Remove Serena MCP only if the ledger proves Headroom installed it.""" - from headroom.mcp_registry.ledger import clear_install, headroom_installed_matching - - current = registrar.get_server("serena") - if not headroom_installed_matching(registrar.name, current): - return "not_headroom_owned" - if registrar.unregister_server("serena"): - clear_install(registrar.name, "serena") - return "removed" - return "failed" - - -def _disable_serena_mcp( - registrar: Any, *, verbose: bool = False, reason: str = "--no-serena" -) -> None: - """Actively disable a Headroom-installed Serena entry, not merely skip it. - - Serena used to be registered by default, so a prior ``headroom wrap`` - persists a ``serena`` entry into the agent's MCP config; the agent then - keeps launching Serena on startup. Just *skipping* registration on a later - run leaves that stale entry in place — so this removes the entry Headroom - installed. A user-managed Serena (absent from our ledger) is reported but - left untouched. ``reason`` is surfaced in the message: ``--no-serena`` when - the user opted out, or a note that tokensave is now the primary compressor. - """ - if not registrar.detect(): - if verbose: - click.echo(f" Serena MCP: {registrar.display_name} not detected — skipping") - return - - if registrar.get_server("serena") is None: - if verbose: - click.echo(f" Skipping Serena MCP ({reason})") - return - - status = _remove_headroom_installed_serena_mcp(registrar) - if status == "removed": - click.echo(f" Removed previously-installed Serena MCP ({reason})") - click.echo(f" restart {registrar.display_name} if it was already running") - elif status == "not_headroom_owned": - click.echo( - " Serena MCP is present but user-managed — leaving it in place " - "(--no-serena only removes entries Headroom installed)" - ) - else: # "failed" - click.echo( - " Serena MCP: removal failed — remove the 'serena' entry from your MCP config manually" - ) - - -# ============================================================================= -# tokensave — primary coding-task compressor (Serena is the backup) -# ============================================================================= - - -def _ensure_tokensave_binary(verbose: bool = False) -> Path | None: - """Resolve the tokensave binary, fetching the release asset if missing. - - Returns the binary path, or ``None`` when tokensave is unavailable - (offline, unsupported platform, or download failure) — the caller then - falls back to Serena. - """ - from headroom.graph.tokensave_installer import ensure_tokensave, get_tokensave_path - - existing = get_tokensave_path() - if existing: - return existing - - click.echo(" tokensave: fetching code-graph binary...") - path = ensure_tokensave() - if path: - click.echo(f" tokensave: installed at {path}") - else: - click.echo( - " tokensave: no prebuilt binary available for this platform " - "(try 'cargo install tokensave') — falling back to Serena" - ) - return path - - -def _index_tokensave_project(bin_path: Path, *, verbose: bool = False) -> None: - """Index the current project into the tokensave graph (non-fatal). - - Runs ``tokensave init`` the first time (creates ``.tokensave/``), then - ``tokensave sync`` for incremental updates. tokensave also re-checks - staleness on demand, so a failure here is logged but never blocks the - wrap — the MCP server still indexes lazily on first query. - """ - project_dir = Path.cwd() - subcommand = "sync" if (project_dir / ".tokensave").exists() else "init" - try: - result = run( - [str(bin_path), subcommand], - capture_output=True, - text=True, - timeout=60, - ) - if result.returncode == 0: - click.echo(" Code graph: indexed (tokensave)") - elif verbose: - click.echo(f" Code graph: tokensave {subcommand} failed ({result.stderr[:100]})") - except subprocess.TimeoutExpired: - click.echo(" Code graph: tokensave indexing timed out (will complete on demand)") - except Exception as e: - if verbose: - click.echo(f" Code graph: tokensave indexing skipped ({e})") - - -def _setup_tokensave_mcp(registrar: Any, *, verbose: bool = False, force: bool = False) -> bool: - """Register tokensave MCP with the given agent (idempotent). - - Returns ``True`` when tokensave is available and set up, ``False`` when the - binary is unavailable — the caller then falls back to Serena. Mirrors - :func:`_setup_serena_mcp`'s ledger-aware migration: a stale - Headroom-installed ``tokensave`` entry is force-updated to the current - spec, while a user-managed entry is left untouched. - """ - from headroom.mcp_registry import build_tokensave_spec, format_result - from headroom.mcp_registry.base import RegisterStatus - from headroom.mcp_registry.ledger import headroom_installed_matching, record_install - - if not registrar.detect(): - if verbose: - click.echo(f" tokensave MCP: {registrar.display_name} not detected — skipping") - return False - - bin_path = _ensure_tokensave_binary(verbose=verbose) - if bin_path is None: - return False - - # Warm the graph so the first query is instant (non-fatal). - _index_tokensave_project(bin_path, verbose=verbose) - - spec = build_tokensave_spec(str(bin_path)) - result = registrar.register_server(spec, force=force) - - # Migrate a stale Headroom-installed entry (e.g. an older binary path or - # pinned version), mirroring the Serena migration path. Only force-update - # when the ledger proves Headroom installed the entry on disk. - if ( - result.status == RegisterStatus.MISMATCH - and not force - and headroom_installed_matching(registrar.name, registrar.get_server("tokensave")) - ): - result = registrar.register_server(spec, force=True) - if result.status == RegisterStatus.REGISTERED: - click.echo(" tokensave MCP: migrated previously-installed entry to current spec") - - if result.status == RegisterStatus.REGISTERED: - record_install(registrar.name, spec) - - line = format_result( - registrar.name, - result, - label="tokensave MCP", - verbose=verbose, - overwrite_hint="update or remove the existing tokensave MCP entry, then rerun headroom wrap", - restart_hint=f"restart {registrar.display_name} if it was already running", - ) - if line is not None: - click.echo(line) - return True - - -def _remove_headroom_installed_tokensave_mcp(registrar: Any) -> str: - """Remove the tokensave MCP entry only if the ledger proves Headroom installed it.""" - from headroom.mcp_registry.ledger import clear_install, headroom_installed_matching - - current = registrar.get_server("tokensave") - if not headroom_installed_matching(registrar.name, current): - return "not_headroom_owned" - if registrar.unregister_server("tokensave"): - clear_install(registrar.name, "tokensave") - return "removed" - return "failed" - - -def _disable_tokensave_mcp(registrar: Any, *, verbose: bool = False) -> None: - """Make ``--no-tokensave`` actively remove a Headroom-installed tokensave entry.""" - if not registrar.detect(): - if verbose: - click.echo(f" tokensave MCP: {registrar.display_name} not detected — skipping") - return - - if registrar.get_server("tokensave") is None: - if verbose: - click.echo(" Skipping tokensave MCP (--no-tokensave)") - return - - status = _remove_headroom_installed_tokensave_mcp(registrar) - if status == "removed": - click.echo(" Removed previously-installed tokensave MCP (--no-tokensave)") - click.echo(f" restart {registrar.display_name} if it was already running") - elif status == "not_headroom_owned": - click.echo( - " tokensave MCP is present but user-managed — leaving it in place " - "(--no-tokensave only removes entries Headroom installed)" - ) - else: # "failed" - click.echo( - " tokensave MCP: removal failed — remove the 'tokensave' entry " - "from your MCP config manually" - ) - - -def _setup_coding_compressor(registrar: Any, *, serena_context: str, **kwargs: Any) -> None: - """Set up the coding-task compressor: tokensave primary, Serena backup. - - Policy (decided per the integration): - - * ``no_tokensave`` — skip/disable tokensave entirely. - * tokensave is set up by default; on success it becomes the primary - compressor and any Headroom-installed Serena entry is removed. - * Serena is the backup: registered automatically when tokensave is - unavailable (unless ``no_serena``), or forced on with ``serena=True``. - - ``kwargs`` carries the boolean flags ``serena``, ``no_serena``, - ``no_tokensave`` and the per-agent registrar ``force`` semantics. - """ - serena = bool(kwargs.get("serena")) - no_serena = bool(kwargs.get("no_serena")) - no_tokensave = bool(kwargs.get("no_tokensave")) - force = bool(kwargs.get("force")) - verbose = bool(kwargs.get("verbose")) - - tokensave_ok = False - if no_tokensave: - _disable_tokensave_mcp(registrar, verbose=verbose) - else: - tokensave_ok = _setup_tokensave_mcp(registrar, verbose=verbose, force=force) - - if serena or (not tokensave_ok and not no_serena): - _setup_serena_mcp(registrar, context=serena_context, verbose=verbose, force=force) - else: - # tokensave is primary (or Serena was explicitly disabled): drop any - # Serena entry a prior wrap installed; user-managed entries are kept. - reason = ( - "--no-serena" if no_serena else "tokensave is now the primary code-graph compressor" - ) - _disable_serena_mcp(registrar, verbose=verbose, reason=reason) - - -_CBM_MCP_SERVER_NAME = "codebase-memory-mcp" - - -def _setup_code_graph(verbose: bool = False) -> bool: - """Ensure the tokensave code graph is set up and the project indexed. - - tokensave is Headroom's primary code-graph compressor and is normally - installed by default (it builds a semantic knowledge graph the LLM can - query for call chains, definitions, and impact analysis instead of - reading whole files). ``--code-graph`` is kept for backward compatibility - and as an explicit "set up the graph and force an index now" switch, even - when tokensave registration was otherwise skipped. - - Returns True if the graph is ready, False if tokensave is unavailable. - Earlier releases backed this flag with ``codebase-memory-mcp``; that - server is no longer installed, and ``headroom unwrap`` still cleans up any - legacy ``codebase-memory-mcp`` entry a prior wrap left behind. - """ - from headroom.mcp_registry import ClaudeRegistrar - - return _setup_tokensave_mcp(ClaudeRegistrar(), verbose=verbose, force=True) - - -# rtk instructions for tools without hook support (Codex, Cursor, Aider). -# These get injected into AGENTS.md / .cursorrules so the LLM voluntarily -# uses rtk-prefixed commands. Kept concise to minimize instruction overhead. -RTK_INSTRUCTIONS_BLOCK = """\ - -# RTK (Rust Token Killer) - Token-Optimized Commands - -When running shell commands, **always prefix with `rtk`**. This reduces context -usage by 60-90% with zero behavior change. If rtk has no filter for a command, -it passes through unchanged — so it is always safe to use. - -## Key Commands -```bash -# Git (59-80% savings) -rtk git status rtk git diff rtk git log - -# Files & Search (60-75% savings) -rtk ls rtk read rtk grep -rtk find rtk diff - -# Test (90-99% savings) — shows failures only -rtk pytest tests/ rtk cargo test rtk test - -# Build & Lint (80-90% savings) — shows errors only -rtk tsc rtk lint rtk cargo build -rtk prettier --check rtk mypy rtk ruff check - -# Analysis (70-90% savings) -rtk err rtk log rtk json -rtk summary rtk deps rtk env - -# GitHub (26-87% savings) -rtk gh pr view rtk gh run list rtk gh issue list - -# Infrastructure (85% savings) -rtk docker ps rtk kubectl get rtk docker logs - -# Package managers (70-90% savings) -rtk pip list rtk pnpm install rtk npm run