0ef5fcb1c5
Security / Dependency audit (pip-audit) (push) Has been cancelled
Security / CodeQL (javascript-typescript) (push) Has been cancelled
Security / CodeQL (python) (push) Has been cancelled
Security / Secret scan (gitleaks) (push) Has been cancelled
rust / test (ubuntu) (push) Has been cancelled
rust / simulator e2e (macos-latest) (push) Has been cancelled
rust / simulator e2e (ubuntu-latest) (push) Has been cancelled
rust / simulator e2e (windows-latest) (push) Has been cancelled
rust / wheels (aarch64-apple-darwin) (push) Has been cancelled
rust / wheels (x86_64-unknown-linux-gnu) (push) Has been cancelled
rust / wheels (x86_64-apple-darwin) (push) Has been cancelled
rust / audit (push) Has been cancelled
rust / parity (nightly, allowed to fail during Phase 0) (push) Has been cancelled
CI / commitlint (push) Has been skipped
Dev Containers / validate (.devcontainer/devcontainer.json, default) (push) Failing after 0s
Dev Containers / validate (.devcontainer/memory-stack/devcontainer.json, memory-stack) (push) Failing after 0s
Dev Containers / validate-worktree (push) Failing after 0s
CI / changes (push) Failing after 4s
Deploy Documentation / validate (push) Has been skipped
Deploy Documentation / deploy (push) Failing after 1s
Init Native E2E / init-native (ubuntu-latest, claude) (push) Failing after 1s
Init Native E2E / init-native (ubuntu-latest, codex) (push) Failing after 1s
Install Native E2E / install-native (ubuntu-latest) (push) Failing after 1s
OpenCode Plugin / typecheck + build + test (push) Failing after 1s
Init Native E2E / init-native (ubuntu-latest, copilot) (push) Failing after 1s
Release Please / release-please (push) Failing after 1s
Wrap E2E / docker-wrap-e2e (push) Failing after 1s
Wrap Native E2E / wrap-native (ubuntu-latest) (push) Failing after 1s
Init E2E / docker-init-e2e (push) Failing after 4s
Merge Conflicts / merge-conflicts (push) Failing after 4s
CI / lint (push) Has been cancelled
CI / build-wheel (push) Has been cancelled
CI / build-wheel-windows (push) Has been cancelled
CI / prefetch-model (push) Has been cancelled
CI / test-dashboard-ui (push) Has been cancelled
CI / test (1) (push) Has been cancelled
CI / test (2) (push) Has been cancelled
CI / test (3) (push) Has been cancelled
CI / test (4) (push) Has been cancelled
CI / test-extras (push) Has been cancelled
CI / test-agno (push) Has been cancelled
CI / build (push) Has been cancelled
CI / workflow-validation (push) Has been cancelled
CI / docker-native-e2e (push) Has been cancelled
CI / windows-native-wrapper (push) Has been cancelled
CI / macos-native-wrapper (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code-nonroot name:code-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code-slim name:code-slim]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code-slim-nonroot name:code-slim-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-nonroot name:nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-slim name:slim]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-slim-nonroot name:slim-nonroot]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime name:]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-code name:code]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-code-nonroot name:code-nonroot]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-code-slim name:code-slim]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-code-slim-nonroot name:code-slim-nonroot]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-nonroot name:nonroot]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-slim name:slim]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-slim-nonroot name:slim-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime name:]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code name:code]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code-nonroot name:code-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code-slim name:code-slim]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code-slim-nonroot name:code-slim-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-nonroot name:nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-slim name:slim]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-slim-nonroot name:slim-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime name:]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code name:code]) (push) Has been cancelled
Docker / promote-latest (push) Has been cancelled
Init Native E2E / init-native (macos-latest, claude) (push) Has been cancelled
Init Native E2E / init-native (macos-latest, codex) (push) Has been cancelled
Init Native E2E / init-native (macos-latest, copilot) (push) Has been cancelled
Install Native E2E / install-native (macos-latest) (push) Has been cancelled
Wrap Native E2E / wrap-native (macos-latest) (push) Has been cancelled
6353 lines
251 KiB
Python
6353 lines
251 KiB
Python
"""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:<port> — 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 = """\
|
|
<!-- headroom:rtk-instructions -->
|
|
# 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 <path> rtk read <file> rtk grep <pattern>
|
|
rtk find <pattern> rtk diff <file>
|
|
|
|
# Test (90-99% savings) — shows failures only
|
|
rtk pytest tests/ rtk cargo test rtk test <cmd>
|
|
|
|
# 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 <cmd> rtk log <file> rtk json <file>
|
|
rtk summary <cmd> rtk deps rtk env
|
|
|
|
# GitHub (26-87% savings)
|
|
rtk gh pr view <n> rtk gh run list rtk gh issue list
|
|
|
|
# Infrastructure (85% savings)
|
|
rtk docker ps rtk kubectl get rtk docker logs <c>
|
|
|
|
# Package managers (70-90% savings)
|
|
rtk pip list rtk pnpm install rtk npm run <script>
|
|
```
|
|
|
|
## Rules
|
|
- In command chains, prefix each segment: `rtk git add . && rtk git commit -m "msg"`
|
|
- For debugging, use raw command without rtk prefix
|
|
- `rtk proxy <cmd>` runs command without filtering but tracks usage
|
|
<!-- /headroom:rtk-instructions -->
|
|
"""
|
|
|
|
# Marker used to detect if instructions are already injected
|
|
_RTK_MARKER = "<!-- headroom:rtk-instructions -->"
|
|
|
|
# Memory MCP markers
|
|
_MEMORY_MCP_MARKER = "# --- Headroom memory MCP (auto-injected) ---"
|
|
_MEMORY_MCP_END = "# --- end Headroom memory ---"
|
|
_MEMORY_AGENTS_MARKER = "<!-- headroom:memory-instructions -->"
|
|
|
|
# Codex config injection markers
|
|
_CODEX_TOP_LEVEL_MARKER = "# --- Headroom proxy (auto-injected by headroom wrap codex) ---"
|
|
_CODEX_END_MARKER = "# --- end Headroom ---"
|
|
_CODEX_MCP_MARKER = "# --- Headroom MCP server ---"
|
|
_CODEX_MCP_END = "# --- end Headroom MCP server ---"
|
|
# File name used for the pre-wrap snapshot of the Codex config file. The
|
|
# snapshot lets `headroom unwrap codex` restore the exact prior state, even
|
|
# if the user had their own `model_provider` / `[model_providers.*]` config
|
|
# before running wrap.
|
|
_CODEX_CONFIG_BACKUP_SUFFIX = ".headroom-backup"
|
|
|
|
|
|
def _codex_home_dir() -> Path:
|
|
"""Return Codex's config directory, respecting ``CODEX_HOME`` when set."""
|
|
codex_home = os.environ.get("CODEX_HOME")
|
|
if codex_home:
|
|
return Path(codex_home).expanduser()
|
|
return Path.home() / ".codex"
|
|
|
|
|
|
@contextmanager
|
|
def _codex_session_home_overlay() -> Any:
|
|
"""Seed a temporary Codex home from the active home and point the process at it."""
|
|
source_home = _codex_home_dir()
|
|
original_codex_home = os.environ.get("CODEX_HOME")
|
|
|
|
with tempfile.TemporaryDirectory(prefix="headroom-codex-home-") as tmp_dir:
|
|
session_home = Path(tmp_dir)
|
|
if source_home.exists():
|
|
shutil.copytree(source_home, session_home, dirs_exist_ok=True)
|
|
|
|
os.environ["CODEX_HOME"] = str(session_home)
|
|
try:
|
|
yield session_home
|
|
finally:
|
|
if original_codex_home is None:
|
|
os.environ.pop("CODEX_HOME", None)
|
|
else:
|
|
os.environ["CODEX_HOME"] = original_codex_home
|
|
|
|
|
|
def _codex_config_paths() -> tuple[Path, Path]:
|
|
"""Return ``(config_file, backup_file)`` paths for the Codex TOML config."""
|
|
config_dir = _codex_home_dir()
|
|
config_file = config_dir / "config.toml"
|
|
backup_file = config_dir / f"config.toml{_CODEX_CONFIG_BACKUP_SUFFIX}"
|
|
return config_file, backup_file
|
|
|
|
|
|
def _strip_codex_headroom_blocks(
|
|
content: str,
|
|
*,
|
|
remove_mcp: bool = False,
|
|
remove_named_mcp: bool = True,
|
|
) -> str:
|
|
"""Remove all Headroom-managed blocks from a Codex ``config.toml`` string.
|
|
|
|
Returns the cleaned content. Safe to call on content that never contained
|
|
any markers — it will be returned effectively unchanged (only trailing
|
|
whitespace is normalized).
|
|
"""
|
|
import re
|
|
|
|
def _remove_marker_span(text: str, start_marker: str, end_marker: str) -> str:
|
|
while start_marker in text and end_marker in text:
|
|
start = text.index(start_marker)
|
|
end_idx = text.index(end_marker, start)
|
|
if end_idx < start:
|
|
break
|
|
end = end_idx + len(end_marker)
|
|
text = text[:start].rstrip("\n") + "\n" + text[end:].lstrip("\n")
|
|
text = text.replace(start_marker + "\n", "")
|
|
text = text.replace(end_marker + "\n", "")
|
|
return text
|
|
|
|
# Remove any top-level-marker → end-marker span, possibly repeated.
|
|
content = _remove_marker_span(content, _CODEX_TOP_LEVEL_MARKER, _CODEX_END_MARKER)
|
|
|
|
if remove_mcp:
|
|
# Remove Headroom-managed MCP blocks written by `wrap codex`.
|
|
content = _remove_marker_span(content, _CODEX_MCP_MARKER, _CODEX_MCP_END)
|
|
if remove_named_mcp:
|
|
content = re.sub(
|
|
r"(?ms)^# --- Headroom MCP server: [^\n]+ ---\n.*?"
|
|
r"^# --- end Headroom MCP server: [^\n]+ ---\n?",
|
|
"",
|
|
content,
|
|
)
|
|
content = _remove_marker_span(content, _MEMORY_MCP_MARKER, _MEMORY_MCP_END)
|
|
|
|
# Strip any leftover top-level keys that older (or crashed) versions of
|
|
# `wrap codex` may have written outside the marker block.
|
|
content = re.sub(r'(?m)^[ \t]*model_provider[ \t]*=[ \t]*"headroom"[ \t]*\r?\n', "", content)
|
|
content = re.sub(
|
|
r'(?m)^[ \t]*openai_base_url[ \t]*=[ \t]*"http://127\.0\.0\.1:\d+/v1"[ \t]*\r?\n',
|
|
"",
|
|
content,
|
|
)
|
|
|
|
# Strip any orphaned `[model_providers.headroom]` table with the fields we
|
|
# write. We only remove it if the table is recognisably ours (base_url
|
|
# mentions localhost and a Headroom proxy port). This protects users who
|
|
# happen to have a differently configured `headroom` provider.
|
|
orphan_headroom_table = re.compile(
|
|
r"(?ms)^\[model_providers\.headroom\][^\[]*?"
|
|
r'base_url[ \t]*=[ \t]*"http://127\.0\.0\.1:\d+/v1"[^\[]*?'
|
|
r"(?=^\[|\Z)"
|
|
)
|
|
content = orphan_headroom_table.sub("", content)
|
|
|
|
return content.lstrip("\n").rstrip() + "\n" if content.strip() else ""
|
|
|
|
|
|
# Top-level bare keys we redirect to headroom values when the user already
|
|
# has them set. Match the entire line (including any trailing comment) so
|
|
# we can rewrite it cleanly. Bare keys must precede any [section] in TOML,
|
|
# so a `^` anchor combined with `^[ \t]*key` is sufficient — table lines
|
|
# start with `[`, not with the key name.
|
|
_REDIRECTABLE_KEYS: tuple[str, ...] = ("model_provider", "openai_base_url")
|
|
|
|
|
|
def _strip_existing_codex_headroom_provider_table(content: str) -> str:
|
|
"""Remove a pre-existing ``[model_providers.headroom]`` table before wrap."""
|
|
if "[model_providers.headroom]" not in content:
|
|
return content
|
|
|
|
import re # local import to match surrounding helper convention
|
|
|
|
provider_table = re.compile(
|
|
r"(?ms)^[ \t]*\[model_providers\.headroom\][^\n]*\n.*?(?=^[ \t]*\[|\Z)"
|
|
)
|
|
content = provider_table.sub("", content)
|
|
return content.lstrip("\n").rstrip() + "\n" if content.strip() else ""
|
|
|
|
|
|
def _redirect_existing_top_level_keys(content: str, port: int) -> str:
|
|
"""Rewrite user-defined top-level keys so wrap does not create duplicates.
|
|
|
|
Codex's ``config.toml`` rejects duplicate top-level keys (TOML spec),
|
|
which would break ``codex`` startup after ``headroom wrap codex`` runs
|
|
on a config that already declares its own ``model_provider`` or
|
|
``openai_base_url``.
|
|
|
|
For each redirectable key, if the user's line already sets it, replace
|
|
the value with the headroom one and append ``# was: <original-value>``
|
|
so the user can still see and recover their previous setting. The
|
|
snapshot taken in ``_snapshot_codex_config_if_unwrapped`` ensures the
|
|
pre-wrap file can be restored byte-for-byte on ``headroom unwrap
|
|
codex``.
|
|
|
|
Returns the modified content. If no redirectable keys are present,
|
|
the content is returned unchanged and the caller should fall back to
|
|
prepending the marker-delimited top-level block (current behavior).
|
|
"""
|
|
import re # local import to match the module's existing convention
|
|
|
|
if not content.strip():
|
|
return content
|
|
|
|
def _make_replacer(current_key: str, current_port: int) -> Callable[[re.Match[str]], str]:
|
|
def _replace(match: re.Match[str]) -> str:
|
|
original_value = match.group("value")
|
|
if current_key == "model_provider":
|
|
new_value = "headroom"
|
|
else: # openai_base_url
|
|
new_value = f"http://127.0.0.1:{current_port}/v1"
|
|
if original_value == new_value:
|
|
return match.group(0)
|
|
# Keep the user's original value in a trailing comment so they
|
|
# can see what was changed. This is metadata, not a TOML
|
|
# duplicate.
|
|
return f'{current_key} = "{new_value}" # was: {original_value}'
|
|
|
|
return _replace
|
|
|
|
redirected = content
|
|
for key in _REDIRECTABLE_KEYS:
|
|
pattern = re.compile(rf'(?m)^[ \t]*{re.escape(key)}[ \t]*=[ \t]*"(?P<value>[^"\n]*)"[^\n]*')
|
|
redirected = pattern.sub(_make_replacer(key, port), redirected, count=1)
|
|
return redirected
|
|
|
|
|
|
def _has_redirectable_top_level_key(content: str, key: str) -> bool:
|
|
"""Return True if ``content`` declares ``key = "..."`` as a top-level key."""
|
|
import re # local import to match the module's existing convention
|
|
|
|
pattern = re.compile(rf'(?m)^[ \t]*{key}[ \t]*=[ \t]*"[^"\n]*"')
|
|
return pattern.search(content) is not None
|
|
|
|
|
|
def _codex_config_has_headroom_markers(content: str) -> bool:
|
|
"""Return whether a Codex config already contains wrap-owned markers."""
|
|
managed_markers = (
|
|
_CODEX_TOP_LEVEL_MARKER,
|
|
_CODEX_END_MARKER,
|
|
_CODEX_MCP_MARKER,
|
|
_MEMORY_MCP_MARKER,
|
|
)
|
|
return any(marker in content for marker in managed_markers)
|
|
|
|
|
|
def _snapshot_codex_config_if_unwrapped(config_file: Path, backup_file: Path) -> None:
|
|
"""Snapshot ``config.toml`` to ``backup_file`` before the first injection.
|
|
|
|
Called as the first step of every Headroom injection into Codex's
|
|
``config.toml``. Guarantees that ``headroom unwrap codex`` can restore the
|
|
user's original file byte-for-byte.
|
|
|
|
Rules:
|
|
|
|
* If the backup already exists, leave it alone — we only snapshot the
|
|
*pre-wrap* state, so running wrap repeatedly must not clobber it.
|
|
* If the config file doesn't exist yet, there's nothing to back up; unwrap
|
|
will remove the file entirely instead of restoring a snapshot.
|
|
* If the config already contains any Headroom-managed Codex marker, a wrap
|
|
run is already active: do not snapshot the injected state.
|
|
"""
|
|
if backup_file.exists():
|
|
return
|
|
if not config_file.exists():
|
|
return
|
|
try:
|
|
content = _read_text(config_file)
|
|
except OSError:
|
|
return
|
|
if _codex_config_has_headroom_markers(content):
|
|
return
|
|
backup_file.parent.mkdir(parents=True, exist_ok=True)
|
|
shutil.copy2(config_file, backup_file)
|
|
|
|
|
|
def _ensure_rtk_binary(verbose: bool = False) -> Path | None:
|
|
"""Ensure rtk binary is installed (download if needed). No hook registration."""
|
|
from headroom.rtk import get_rtk_path
|
|
from headroom.rtk.installer import ensure_rtk
|
|
|
|
rtk_path = get_rtk_path()
|
|
|
|
if rtk_path:
|
|
if verbose:
|
|
click.echo(f" rtk found at {rtk_path}")
|
|
return rtk_path
|
|
|
|
click.echo(" Downloading rtk (Rust Token Killer)...")
|
|
rtk_path = ensure_rtk()
|
|
if rtk_path:
|
|
click.echo(f" rtk installed at {rtk_path}")
|
|
return rtk_path
|
|
|
|
click.echo(" rtk download failed — continuing without it")
|
|
return None
|
|
|
|
|
|
def _prepare_wrap_rtk(verbose: bool = False, *, label: str | None = None) -> Path | None:
|
|
"""Ensure rtk is present for host-bridged wrap flows without host-specific setup."""
|
|
if label:
|
|
click.echo(f" Preparing rtk for {label}...")
|
|
return _ensure_rtk_binary(verbose=verbose)
|
|
|
|
|
|
# Canonical casing for the proxy's per-project savings header (matched
|
|
# case-insensitively by headroom.proxy.project_context.PROJECT_HEADER).
|
|
_PROJECT_HEADER_NAME = "X-Headroom-Project"
|
|
|
|
|
|
def _project_name_from_cwd() -> str | None:
|
|
"""Project label for X-Headroom-Project: basename of the launch directory.
|
|
|
|
Non-ASCII characters are percent-encoded (RFC 3986) so the header value
|
|
stays within the visible-ASCII range required by RFC 7230. The proxy
|
|
decodes the value in sanitize_project_name before storing it.
|
|
"""
|
|
name = Path.cwd().name.strip()
|
|
if not name:
|
|
return None
|
|
return urllib.parse.quote(name, safe="-_.() ")
|
|
|
|
|
|
def _apply_project_header_env(env: dict[str, str]) -> None:
|
|
"""Inject X-Headroom-Project into ``ANTHROPIC_CUSTOM_HEADERS``.
|
|
|
|
Claude Code reads ``ANTHROPIC_CUSTOM_HEADERS`` as newline-separated
|
|
``Name: value`` lines and attaches them to every API request; the
|
|
Headroom proxy uses the X-Headroom-Project header for per-project
|
|
savings attribution. An existing user-supplied x-headroom-project
|
|
header (any casing) always wins — we never duplicate or overwrite it,
|
|
and any other user headers are preserved by appending.
|
|
"""
|
|
project = _project_name_from_cwd()
|
|
if not project:
|
|
return
|
|
header_line = f"{_PROJECT_HEADER_NAME}: {project}"
|
|
existing = env.get("ANTHROPIC_CUSTOM_HEADERS")
|
|
if existing:
|
|
for line in existing.splitlines():
|
|
name = line.split(":", 1)[0].strip()
|
|
if name.lower() == _PROJECT_HEADER_NAME.lower():
|
|
return # user override wins
|
|
env["ANTHROPIC_CUSTOM_HEADERS"] = f"{existing}\n{header_line}"
|
|
else:
|
|
env["ANTHROPIC_CUSTOM_HEADERS"] = header_line
|
|
|
|
|
|
# Codex's own built-in providers plus Headroom's injected one — never treated
|
|
# as a "custom upstream to preserve" by _detect_custom_codex_upstream_base_url.
|
|
_CODEX_BUILTIN_PROVIDER_NAMES = frozenset({"openai", "anthropic", "azure", "headroom"})
|
|
|
|
# Header carrying a preserved custom upstream (freemodel.dev, LiteLLM, vLLM,
|
|
# ...) so the proxy forwards to it instead of the hardcoded OpenAI default.
|
|
# Codex's env_http_headers only accepts an env-var *name* per header (not a
|
|
# literal value), so the detected URL is exported into this env var by the
|
|
# `wrap codex` launch path — see its use in `codex()` below.
|
|
_UPSTREAM_BASE_URL_HEADER_NAME = "X-Headroom-Base-Url"
|
|
_UPSTREAM_BASE_URL_ENV_VAR = "HEADROOM_CODEX_UPSTREAM_BASE_URL"
|
|
|
|
|
|
def _codex_custom_provider_base_urls(content: str) -> dict[str, str]:
|
|
"""Return ``{provider_name: base_url}`` for user-declared custom providers.
|
|
|
|
Excludes Codex's built-ins (``openai``/``anthropic``/``azure``) and
|
|
Headroom's own ``headroom`` table. A table with no ``base_url`` line, or
|
|
one already pointing at Headroom's own localhost proxy (a leftover from a
|
|
prior wrap this pass hasn't stripped yet), is excluded too.
|
|
"""
|
|
import re
|
|
|
|
tables: dict[str, str] = {}
|
|
for match in re.finditer(
|
|
r"(?ms)^[ \t]*\[model_providers\.(?P<name>[^\]\s]+)\][ \t]*\n"
|
|
r"(?P<body>.*?)(?=^[ \t]*\[|\Z)",
|
|
content,
|
|
):
|
|
name = match.group("name")
|
|
if name in _CODEX_BUILTIN_PROVIDER_NAMES:
|
|
continue
|
|
base_match = re.search(
|
|
r'(?m)^[ \t]*base_url[ \t]*=[ \t]*"(?P<url>[^"\n]*)"', match.group("body")
|
|
)
|
|
if not base_match:
|
|
continue
|
|
url = base_match.group("url").strip().rstrip("/")
|
|
if not url or url.startswith(("http://127.0.0.1", "http://localhost")):
|
|
continue
|
|
tables[name] = url
|
|
return tables
|
|
|
|
|
|
def _detect_custom_codex_upstream_base_url(content: str) -> str | None:
|
|
"""Return a user-configured custom provider ``base_url`` to preserve, if any.
|
|
|
|
Codex lets users declare OpenAI-compatible gateways (LiteLLM, vLLM,
|
|
freemodel.dev, ...) under ``[model_providers.<name>]`` and select one via
|
|
the top-level ``model_provider`` key. Before this, ``headroom wrap codex``
|
|
unconditionally pointed the proxy's upstream OpenAI route at
|
|
``api.openai.com``, silently discarding that selection — the user's
|
|
gateway API key then gets sent to OpenAI, which rejects it (#1614).
|
|
|
|
If the top-level ``model_provider`` names one of the detected custom
|
|
tables, that selection wins unambiguously. This also covers the
|
|
already-wrapped case: once wrap has run once, the top-level key reads
|
|
``model_provider = "headroom" # was: <original>`` (see
|
|
``_redirect_existing_top_level_keys``), so the original selection is
|
|
recovered from that trailing comment on re-wrap / port changes.
|
|
|
|
Falls back to the sole candidate when exactly one custom table exists
|
|
and there is no (or no matching) top-level selection — the common case
|
|
from the bug report, where the table is declared but selection happens
|
|
via ``--profile`` rather than a static top-level key. Returns ``None``
|
|
when there are multiple, un-selected candidates (ambiguous — guessing
|
|
wrong is worse than the prior default behavior) or none at all.
|
|
"""
|
|
import re
|
|
|
|
candidates = _codex_custom_provider_base_urls(content)
|
|
if not candidates:
|
|
return None
|
|
|
|
selected = re.search(
|
|
r'(?m)^[ \t]*model_provider[ \t]*=[ \t]*"(?P<name>[^"\n]*)"'
|
|
r"(?:[ \t]*#[ \t]*was:[ \t]*(?P<was>[^\r\n]*))?",
|
|
content,
|
|
)
|
|
if selected:
|
|
was = (selected.group("was") or "").strip()
|
|
chosen = was or selected.group("name")
|
|
if chosen in candidates:
|
|
return candidates[chosen]
|
|
|
|
if len(candidates) == 1:
|
|
return next(iter(candidates.values()))
|
|
return None
|
|
|
|
|
|
def _inject_codex_provider_config(port: int) -> str | None:
|
|
"""Inject a Headroom model provider into Codex's config.toml.
|
|
|
|
Two keys need to be in effect for the proxy to route all traffic:
|
|
|
|
* ``model_provider = "headroom"`` — selects the custom provider for
|
|
API-key mode traffic.
|
|
* ``openai_base_url = "http://127.0.0.1:{port}/v1"`` — overrides the
|
|
built-in ``openai`` provider's base URL. This is the critical key for
|
|
**subscription (ChatGPT plan) users**: Codex detects subscription auth
|
|
and routes through the built-in ``openai`` provider regardless of
|
|
``model_provider``, so without this override it bypasses the proxy and
|
|
hits ``https://chatgpt.com/backend-api/codex`` directly.
|
|
|
|
If the user has not already declared these top-level keys, they are
|
|
added in a marker-delimited block at the top of the file. If the
|
|
user *has* declared one or both, the existing lines are rewritten
|
|
in place to the headroom values (with the previous value kept in a
|
|
``# was: …`` trailing comment) so the resulting file stays TOML-valid
|
|
— TOML rejects duplicate top-level keys, which would break
|
|
``codex`` startup.
|
|
|
|
Safe to call multiple times — the injected block is fully replaced on
|
|
each call, so re-running with a different ``port`` updates the config.
|
|
Before the first injection, the pre-wrap file is snapshotted to
|
|
``config.toml.headroom-backup`` so ``headroom unwrap codex``
|
|
can restore it byte-for-byte.
|
|
|
|
Returns the custom upstream ``base_url`` preserved from an existing
|
|
``[model_providers.*]`` table, if one was detected (#1614); ``None``
|
|
otherwise. Callers that go on to launch Codex should export this value
|
|
into ``HEADROOM_CODEX_UPSTREAM_BASE_URL`` (the injected
|
|
``env_http_headers`` entry maps it to the ``X-Headroom-Base-Url`` header,
|
|
which the proxy's OpenAI HTTP handlers honor over the hardcoded
|
|
``api.openai.com`` default) — see its use in ``codex()`` below.
|
|
"""
|
|
config_file, backup_file = _codex_config_paths()
|
|
config_dir = config_file.parent
|
|
|
|
# Detect an existing custom OpenAI-compatible provider BEFORE building the
|
|
# injected block below, so it can be preserved as the upstream the proxy
|
|
# forwards to instead of silently rerouting to api.openai.com (#1614).
|
|
# Best-effort: any read/parse failure just means nothing is preserved,
|
|
# matching prior behavior.
|
|
custom_upstream_base_url: str | None = None
|
|
if config_file.exists():
|
|
try:
|
|
custom_upstream_base_url = _detect_custom_codex_upstream_base_url(
|
|
_read_text(config_file)
|
|
)
|
|
except OSError:
|
|
custom_upstream_base_url = None
|
|
|
|
# The injected content is split into two self-contained, marker-delimited
|
|
# blocks: a top-level key block (at the start of the file, because bare
|
|
# TOML keys must precede any [section]) and a provider-table block (at
|
|
# the end). Each block has its own matching begin/end marker pair so
|
|
# stripping them is unambiguous and never consumes user content that
|
|
# happens to sit between the two. The top-level block is built
|
|
# dynamically below — it contains only keys the user has not already
|
|
# declared (we rewrite the existing ones in place to avoid TOML
|
|
# duplicate-key errors).
|
|
# Emit requires_openai_auth only for ChatGPT-OAuth users (restores the
|
|
# account menu); omitting it for API-key users avoids forcing an OAuth
|
|
# login (#406).
|
|
requires_openai_auth = (
|
|
"requires_openai_auth = true\n" if codex_uses_chatgpt_auth(config_dir / "auth.json") else ""
|
|
)
|
|
# Per-project savings: Codex sends the X-Headroom-Project header only
|
|
# when the mapped env var (HEADROOM_PROJECT, set by `headroom wrap
|
|
# codex`) exists at Codex runtime. When a custom upstream was detected,
|
|
# add a second entry so Codex also sends X-Headroom-Base-Url — the proxy
|
|
# forwards there instead of api.openai.com (#1614).
|
|
env_http_headers_map = {_PROJECT_HEADER_NAME: "HEADROOM_PROJECT"}
|
|
if custom_upstream_base_url:
|
|
env_http_headers_map[_UPSTREAM_BASE_URL_HEADER_NAME] = _UPSTREAM_BASE_URL_ENV_VAR
|
|
env_http_headers_toml = ", ".join(f'"{k}" = "{v}"' for k, v in env_http_headers_map.items())
|
|
provider_section = (
|
|
f"{_CODEX_TOP_LEVEL_MARKER}\n"
|
|
"[model_providers.headroom]\n"
|
|
'name = "OpenAI via Headroom proxy"\n'
|
|
f'base_url = "http://127.0.0.1:{port}/v1"\n'
|
|
f"supports_websockets = true\n"
|
|
f"{requires_openai_auth}"
|
|
# Inline table keeps the key inside this section so
|
|
# _strip_codex_headroom_blocks removes it with the rest of the block.
|
|
f"env_http_headers = {{ {env_http_headers_toml} }}\n"
|
|
f"{_CODEX_END_MARKER}\n"
|
|
)
|
|
|
|
# The two redirectable keys and their headroom target values.
|
|
_REDIRECT_TARGETS = {
|
|
"model_provider": "headroom",
|
|
"openai_base_url": f"http://127.0.0.1:{port}/v1",
|
|
}
|
|
|
|
def _build_top_level_block(user_content: str) -> str:
|
|
"""Build a marker-delimited block containing only the keys the user
|
|
has not already declared at the top level. For keys the user
|
|
*has* declared, the in-place rewrite below handles them.
|
|
"""
|
|
lines = [_CODEX_TOP_LEVEL_MARKER]
|
|
for key, value in _REDIRECT_TARGETS.items():
|
|
if _has_redirectable_top_level_key(user_content, key):
|
|
continue
|
|
lines.append(f'{key} = "{value}"')
|
|
if len(lines) == 1:
|
|
# User already declared every redirectable key — no marker
|
|
# block needed (it would be empty).
|
|
return ""
|
|
lines.append(_CODEX_END_MARKER)
|
|
return "\n".join(lines) + "\n"
|
|
|
|
try:
|
|
config_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Snapshot the pre-wrap state before touching anything. No-op if the
|
|
# config is already wrapped, is missing, or we've already snapshotted.
|
|
_snapshot_codex_config_if_unwrapped(config_file, backup_file)
|
|
|
|
if config_file.exists():
|
|
content = _read_text(config_file)
|
|
# Remove any prior Headroom-managed blocks before re-injecting so
|
|
# the operation is idempotent and supports port changes.
|
|
content = _strip_codex_headroom_blocks(content)
|
|
content = _strip_existing_codex_headroom_provider_table(content)
|
|
|
|
# Bare top-level keys must precede any [section] in TOML, and
|
|
# TOML rejects duplicate top-level keys. Rewrite any existing
|
|
# top-level ``model_provider`` / ``openai_base_url`` in place
|
|
# to the headroom values; for keys the user has not declared,
|
|
# add them in a marker-delimited block at the top of the
|
|
# file. The original values are kept in a trailing ``# was:
|
|
# <value>`` comment, and the snapshot mechanism guarantees
|
|
# byte-for-byte restoration on unwrap.
|
|
user_content = content.strip()
|
|
if user_content:
|
|
redirected = _redirect_existing_top_level_keys(user_content, port)
|
|
top_block = _build_top_level_block(user_content)
|
|
if top_block:
|
|
content = top_block + "\n" + redirected + "\n\n" + provider_section
|
|
else:
|
|
content = redirected + "\n\n" + provider_section
|
|
else:
|
|
# Empty user content — no keys to rewrite in place; emit
|
|
# the full marker block with both redirectable keys.
|
|
content = (
|
|
f"{_CODEX_TOP_LEVEL_MARKER}\n"
|
|
f'model_provider = "{_REDIRECT_TARGETS["model_provider"]}"\n'
|
|
f'openai_base_url = "{_REDIRECT_TARGETS["openai_base_url"]}"\n'
|
|
f"{_CODEX_END_MARKER}\n"
|
|
f"\n{provider_section}"
|
|
)
|
|
else:
|
|
# No config file yet — same as the empty-content path.
|
|
content = (
|
|
f"{_CODEX_TOP_LEVEL_MARKER}\n"
|
|
f'model_provider = "{_REDIRECT_TARGETS["model_provider"]}"\n'
|
|
f'openai_base_url = "{_REDIRECT_TARGETS["openai_base_url"]}"\n'
|
|
f"{_CODEX_END_MARKER}\n"
|
|
f"\n{provider_section}"
|
|
)
|
|
|
|
_write_text(config_file, content)
|
|
click.echo(f" Codex config: injected Headroom provider (WS + HTTP) into {config_file}")
|
|
if custom_upstream_base_url:
|
|
click.echo(
|
|
f" Codex config: preserving existing custom upstream "
|
|
f"{custom_upstream_base_url} (from a pre-existing [model_providers.*] "
|
|
"base_url)"
|
|
)
|
|
# Pull existing native threads into the headroom-provider menu so Codex's
|
|
# history list stays whole once it routes through Headroom. Best-effort.
|
|
retag_to_headroom(_codex_home_dir())
|
|
except Exception as e:
|
|
click.echo(f" Warning: could not update Codex config: {e}")
|
|
return None
|
|
|
|
return custom_upstream_base_url
|
|
|
|
|
|
def _restore_codex_provider_config() -> tuple[str, Path]:
|
|
"""Undo ``_inject_codex_provider_config`` for the active Codex config file.
|
|
|
|
Returns a tuple of ``(status, config_file)`` where status is one of:
|
|
|
|
* ``"restored"`` — a pre-wrap backup existed and was restored; backup
|
|
file has been removed.
|
|
* ``"cleaned"`` — no backup existed, but the Headroom-managed block was
|
|
found and stripped out (preserving surrounding user content).
|
|
* ``"removed"`` — the config file only contained Headroom-managed
|
|
content (created by wrap) and has been deleted.
|
|
* ``"noop"`` — nothing to undo; no Headroom marker and no backup.
|
|
"""
|
|
config_file, backup_file = _codex_config_paths()
|
|
|
|
# Case 1: pre-wrap snapshot exists — restore it exactly.
|
|
if backup_file.exists():
|
|
shutil.copy2(backup_file, config_file)
|
|
backup_file.unlink()
|
|
return "restored", config_file
|
|
|
|
# Case 2: no backup, but config file exists and has markers — strip them.
|
|
if config_file.exists():
|
|
original = _read_text(config_file)
|
|
if _codex_config_has_headroom_markers(original):
|
|
# Without a backup, only remove named MCP blocks when this file
|
|
# also carries wrap-owned provider markers from a full wrap.
|
|
remove_named_mcp = any(
|
|
marker in original
|
|
for marker in (
|
|
_CODEX_TOP_LEVEL_MARKER,
|
|
_CODEX_END_MARKER,
|
|
_CODEX_MCP_MARKER,
|
|
_CODEX_MCP_END,
|
|
)
|
|
)
|
|
cleaned = _strip_codex_headroom_blocks(
|
|
original,
|
|
remove_mcp=True,
|
|
remove_named_mcp=remove_named_mcp,
|
|
)
|
|
if not cleaned.strip():
|
|
# Nothing left but Headroom content — remove the file entirely
|
|
# so Codex falls back to its default config.
|
|
config_file.unlink()
|
|
return "removed", config_file
|
|
_write_text(config_file, cleaned)
|
|
return "cleaned", config_file
|
|
|
|
# Nothing to undo.
|
|
return "noop", config_file
|
|
|
|
|
|
def _emit_wrap_interrupted(agent: str, marker_path: Path | None) -> None:
|
|
"""Log a clear interruption message after a partial wrap setup.
|
|
|
|
Called when a wrap subcommand catches ``KeyboardInterrupt`` between marker
|
|
injection and proxy startup. The marker file (if any) is left on disk —
|
|
re-running the same ``headroom wrap <agent>`` command is idempotent and
|
|
safe.
|
|
"""
|
|
if marker_path is not None:
|
|
click.echo(
|
|
f"\n Wrap was interrupted; marker file at {marker_path} is on "
|
|
f"disk. Rerun `headroom wrap {agent}` to retry — it's idempotent."
|
|
)
|
|
else:
|
|
click.echo(
|
|
f"\n Wrap was interrupted before any on-disk changes. Rerun "
|
|
f"`headroom wrap {agent}` to retry — it's idempotent."
|
|
)
|
|
|
|
|
|
_WRAP_BANNER_INNER_WIDTH = 47
|
|
|
|
|
|
def _print_wrap_banner(agent: str) -> None:
|
|
"""Print a centered ``HEADROOM WRAP: <AGENT>`` banner.
|
|
|
|
Every Pattern-B wrap subcommand (proxy-only + watcher loop) used to
|
|
inline this 3-line box by hand with hand-padded spaces, which made
|
|
title-length changes silently miscenter the title. Compute padding
|
|
here so adding a 9th agent just works.
|
|
"""
|
|
title = f"HEADROOM WRAP: {agent.upper()}"
|
|
pad_total = _WRAP_BANNER_INNER_WIDTH - len(title)
|
|
pad_left = pad_total // 2
|
|
pad_right = pad_total - pad_left
|
|
click.echo()
|
|
click.echo(" ╔" + "═" * _WRAP_BANNER_INNER_WIDTH + "╗")
|
|
click.echo(f" ║{' ' * pad_left}{title}{' ' * pad_right}║")
|
|
click.echo(" ╚" + "═" * _WRAP_BANNER_INNER_WIDTH + "╝")
|
|
click.echo()
|
|
|
|
|
|
def _setup_context_tool_for_agent(
|
|
*,
|
|
agent: str,
|
|
agent_display: str,
|
|
marker_path: Path | None,
|
|
on_rtk_ready: Callable[[Path], object] | None = None,
|
|
rtk_required: bool = False,
|
|
verbose: bool = False,
|
|
) -> Path | None:
|
|
"""Run the rtk-or-lean-ctx context-tool setup with KeyboardInterrupt handling.
|
|
|
|
Replaces the ``try / except KeyboardInterrupt / rtk-vs-lean-ctx fork``
|
|
each wrap subcommand was inlining. Returns the rtk binary path if rtk
|
|
mode was selected and the install succeeded; otherwise ``None``.
|
|
|
|
Args:
|
|
agent: Internal agent name (used by ``_setup_lean_ctx_agent`` and
|
|
the interrupt message).
|
|
agent_display: User-facing capitalized name for the echo lines
|
|
(e.g. ``"Cline"``, ``"OpenHands"``).
|
|
marker_path: Optional path to report on Ctrl+C interruption. Pass
|
|
``None`` for env-only agents that never touch disk (openhands).
|
|
on_rtk_ready: Optional callback invoked with the rtk binary path
|
|
when rtk install succeeds. Use to inject into a marker file
|
|
(e.g. ``.clinerules``, ``.goosehints``, ``config.json``).
|
|
rtk_required: If ``True`` and rtk install fails, raise
|
|
``SystemExit(1)`` instead of silently falling through. Use this
|
|
for env-var-only agents (openhands) where there is no
|
|
fallback marker file to write to.
|
|
verbose: Forwarded to the rtk/lean-ctx installers.
|
|
"""
|
|
try:
|
|
if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX:
|
|
click.echo(f" Setting up lean-ctx for {agent_display}...")
|
|
_setup_lean_ctx_agent(agent, verbose=verbose)
|
|
return None
|
|
click.echo(f" Setting up rtk for {agent_display}...")
|
|
rtk_path = _ensure_rtk_binary(verbose=verbose)
|
|
if not rtk_path:
|
|
if rtk_required:
|
|
click.echo(
|
|
" Error: rtk install failed; refusing to inject "
|
|
f"context-tool guidance for {agent_display} without rtk. "
|
|
"Install rtk manually and re-run, or pass --no-context-tool "
|
|
"to skip rtk."
|
|
)
|
|
raise SystemExit(1)
|
|
return None
|
|
if on_rtk_ready is not None:
|
|
on_rtk_ready(rtk_path)
|
|
return rtk_path
|
|
except KeyboardInterrupt:
|
|
_emit_wrap_interrupted(
|
|
agent, marker_path if (marker_path and marker_path.exists()) else None
|
|
)
|
|
raise SystemExit(130) from None
|
|
|
|
|
|
def _run_proxy_only_watcher(
|
|
*,
|
|
agent_label: str,
|
|
port: int,
|
|
no_proxy: bool,
|
|
learn: bool,
|
|
memory: bool,
|
|
agent_type: str,
|
|
print_setup_lines: Callable[[int], None],
|
|
) -> None:
|
|
"""Shared scaffolding for proxy-only wrap subcommands (no child binary launch).
|
|
|
|
Pattern-B subcommands (cursor / cline / continue) all start the proxy,
|
|
print agent-specific configuration instructions, then block until
|
|
Ctrl+C. This helper unifies that lifecycle so the per-agent diff is
|
|
just the ``print_setup_lines`` callback.
|
|
|
|
The Pattern-A subcommands (aider / copilot / codex / goose / openhands)
|
|
launch a child binary via ``_launch_tool`` instead and never come
|
|
through here. ``_launch_tool`` owns the proxy lifecycle on that path.
|
|
"""
|
|
proxy_holder: list[subprocess.Popen | None] = [None]
|
|
port_holder: list[int] = [port]
|
|
cleanup = _make_cleanup(proxy_holder, port_holder)
|
|
signal.signal(signal.SIGINT, cleanup)
|
|
signal.signal(signal.SIGTERM, cleanup)
|
|
|
|
try:
|
|
_print_wrap_banner(agent_label)
|
|
_register_proxy_client(port)
|
|
proxy_holder[0], actual_port = _ensure_proxy(
|
|
port, no_proxy, learn=learn, memory=memory, agent_type=agent_type
|
|
)
|
|
if actual_port != port:
|
|
_unregister_proxy_client(port)
|
|
_register_proxy_client(actual_port)
|
|
port_holder[0] = actual_port
|
|
_push_runtime_env(actual_port, no_proxy)
|
|
click.echo()
|
|
print_setup_lines(actual_port)
|
|
click.echo()
|
|
click.echo(" Press Ctrl+C to stop the proxy.")
|
|
click.echo()
|
|
|
|
try:
|
|
while True:
|
|
time.sleep(1)
|
|
proc = proxy_holder[0]
|
|
if proc and proc.poll() is not None:
|
|
click.echo(" Proxy process exited unexpectedly.")
|
|
raise SystemExit(1)
|
|
except KeyboardInterrupt:
|
|
click.echo("\n Shutting down...")
|
|
except SystemExit:
|
|
raise
|
|
except Exception as e:
|
|
click.echo(f" Error: {e}")
|
|
raise SystemExit(1) from e
|
|
finally:
|
|
cleanup()
|
|
|
|
|
|
def _inject_rtk_instructions(file_path: Path, verbose: bool = False) -> bool:
|
|
"""Inject rtk instructions into a file (AGENTS.md, .cursorrules, etc.).
|
|
|
|
Idempotent — skips if marker already present. Appends to existing content.
|
|
Returns True if instructions were written.
|
|
"""
|
|
if file_path.exists():
|
|
existing = _read_text(file_path)
|
|
if _RTK_MARKER in existing:
|
|
if verbose:
|
|
click.echo(f" rtk instructions already in {file_path.name}")
|
|
return True
|
|
# Append to existing file
|
|
_append_text(file_path, "\n\n" + RTK_INSTRUCTIONS_BLOCK)
|
|
else:
|
|
file_path.parent.mkdir(parents=True, exist_ok=True)
|
|
_write_text(file_path, RTK_INSTRUCTIONS_BLOCK)
|
|
|
|
click.echo(f" rtk instructions injected into {file_path}")
|
|
return True
|
|
|
|
|
|
def _remove_rtk_instructions(file_path: Path) -> bool:
|
|
"""Remove Headroom's marker-fenced rtk guidance from an instruction file."""
|
|
if not file_path.exists():
|
|
return False
|
|
|
|
content = _read_text(file_path)
|
|
end_marker = "<!-- /headroom:rtk-instructions -->"
|
|
start = content.find(_RTK_MARKER)
|
|
if start < 0:
|
|
return False
|
|
|
|
end = content.find(end_marker, start)
|
|
if end < 0:
|
|
return False
|
|
end += len(end_marker)
|
|
prefix = content[:start].rstrip()
|
|
suffix = content[end:].lstrip("\r\n")
|
|
cleaned = "\n\n".join(part for part in (prefix, suffix) if part)
|
|
if cleaned:
|
|
cleaned = cleaned.rstrip() + "\n"
|
|
|
|
if cleaned:
|
|
_write_text(file_path, cleaned)
|
|
else:
|
|
file_path.unlink()
|
|
return True
|
|
|
|
|
|
def _inject_memory_mcp_config(user_id: str) -> None:
|
|
"""Register headroom memory as an MCP server in Codex's config.toml.
|
|
|
|
Idempotent — replaces existing section if present.
|
|
"""
|
|
import sys
|
|
|
|
config_file, _ = _codex_config_paths()
|
|
config_dir = config_file.parent
|
|
|
|
# Use forward slashes in TOML paths (works on all platforms, avoids
|
|
# backslash escaping issues on Windows)
|
|
python_bin = sys.executable.replace("\\", "/")
|
|
mcp_section = (
|
|
f"\n{_MEMORY_MCP_MARKER}\n"
|
|
f"[mcp_servers.headroom_memory]\n"
|
|
f'command = "{python_bin}"\n'
|
|
f'args = ["-m", "headroom.memory.mcp_server", "--user", "{user_id}"]\n'
|
|
f"startup_timeout_sec = 30\n"
|
|
f"tool_timeout_sec = 30\n"
|
|
f"{_MEMORY_MCP_END}\n"
|
|
)
|
|
|
|
try:
|
|
config_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Snapshot pre-wrap state before touching config.toml so `unwrap codex`
|
|
# can fully restore it even when only `--memory` (not a full provider
|
|
# injection) was used.
|
|
_, backup_file = _codex_config_paths()
|
|
_snapshot_codex_config_if_unwrapped(config_file, backup_file)
|
|
|
|
if config_file.exists():
|
|
content = _read_text(config_file)
|
|
if _MEMORY_MCP_MARKER in content:
|
|
start = content.index(_MEMORY_MCP_MARKER)
|
|
end = content.index(_MEMORY_MCP_END) + len(_MEMORY_MCP_END)
|
|
content = content[:start].rstrip("\n") + mcp_section + content[end:].lstrip("\n")
|
|
else:
|
|
content = content.rstrip() + "\n" + mcp_section
|
|
else:
|
|
content = mcp_section
|
|
|
|
_write_text(config_file, content)
|
|
click.echo(f" Memory MCP: registered in {config_file}")
|
|
except Exception as e:
|
|
click.echo(f" Warning: could not register memory MCP: {e}")
|
|
|
|
|
|
def _inject_memory_agents_md(file_path: Path) -> bool:
|
|
"""Inject memory usage guidance into AGENTS.md.
|
|
|
|
Idempotent — skips if marker already present.
|
|
"""
|
|
memory_block = (
|
|
f"{_MEMORY_AGENTS_MARKER}\n"
|
|
"## Memory\n\n"
|
|
"Use the `headroom_memory` MCP server for persistent cross-session knowledge.\n\n"
|
|
"**Before** answering questions about prior decisions, conventions, project context,\n"
|
|
"architecture, user preferences, org info, codenames, debugging history, or anything\n"
|
|
"from past sessions — call `memory_search` first.\n\n"
|
|
"**After** making durable decisions, discovering conventions, or learning important\n"
|
|
"facts — call `memory_save` to persist them for future sessions.\n\n"
|
|
"Memory is your first source of truth for anything not visible in the current conversation.\n"
|
|
)
|
|
|
|
if file_path.exists():
|
|
existing = _read_text(file_path)
|
|
if _MEMORY_AGENTS_MARKER in existing:
|
|
return True # Already injected
|
|
_append_text(file_path, "\n\n" + memory_block)
|
|
else:
|
|
file_path.parent.mkdir(parents=True, exist_ok=True)
|
|
_write_text(file_path, memory_block)
|
|
|
|
click.echo(f" Memory guidance injected into {file_path.name}")
|
|
return True
|
|
|
|
|
|
def _apply_rtk_to_systemmessage_field(
|
|
container: dict[str, Any],
|
|
location_label: str,
|
|
verbose: bool = False,
|
|
) -> tuple[bool, bool]:
|
|
"""Apply the RTK block to ``container["systemMessage"]`` in place.
|
|
|
|
Returns ``(changed, ok)``:
|
|
|
|
* ``changed`` is ``True`` if the field was written (or rewritten) on this
|
|
call. ``False`` for idempotent skips and for refusals.
|
|
* ``ok`` is ``True`` for "RTK guidance is now present (or refused-safely)",
|
|
``False`` only for refusals where the user must intervene. Callers
|
|
surface ``not ok`` as a warning to the user.
|
|
|
|
Refusal cases (loud, no silent overwrite):
|
|
|
|
* ``systemMessage`` exists and is **not a string** (dict / list / number).
|
|
We never clobber user data of an unknown shape. The user must remove or
|
|
clear the field before re-running.
|
|
"""
|
|
existing_msg = container.get("systemMessage")
|
|
|
|
if isinstance(existing_msg, str) and _RTK_MARKER in existing_msg:
|
|
if verbose:
|
|
click.echo(f" rtk instructions already in {location_label}")
|
|
return False, True
|
|
|
|
if existing_msg is None or (isinstance(existing_msg, str) and not existing_msg.strip()):
|
|
container["systemMessage"] = RTK_INSTRUCTIONS_BLOCK
|
|
return True, True
|
|
|
|
if isinstance(existing_msg, str):
|
|
container["systemMessage"] = existing_msg.rstrip() + "\n\n" + RTK_INSTRUCTIONS_BLOCK
|
|
return True, True
|
|
|
|
# Non-string, non-null value present — refuse loudly. We will not clobber
|
|
# user data of unknown shape.
|
|
click.echo(
|
|
f" Warning: {location_label} systemMessage is not a string "
|
|
f"(type={type(existing_msg).__name__}); refusing to overwrite. "
|
|
"To opt in, remove or clear the existing systemMessage value and re-run."
|
|
)
|
|
return False, False
|
|
|
|
|
|
def _inject_continue_rtk_systemmessage(config_file: Path, verbose: bool = False) -> bool:
|
|
"""Inject the rtk instructions block into Continue's ``.continue/config.json``.
|
|
|
|
Continue's schema supports both a top-level ``systemMessage`` string and a
|
|
per-model ``systemMessage`` on each entry in the ``models`` array. The
|
|
per-model value, when set, overrides the top-level one — so users with
|
|
per-model configs would otherwise silently get no RTK guidance. This
|
|
helper writes the RTK block into **every** ``systemMessage`` site:
|
|
|
|
* top-level ``systemMessage``
|
|
* each ``models[i].systemMessage`` where ``models[i]`` is a dict
|
|
|
|
The RTK marker (``<!-- headroom:rtk-instructions -->``) is the idempotency
|
|
token: if a prior ``systemMessage`` already contains the marker we leave
|
|
that site alone. If the existing value is a non-empty string we append
|
|
with a separator. If the existing value is **non-string** (dict / list /
|
|
number) we refuse loudly and leave it untouched — we do not clobber user
|
|
data of unknown shape. To opt in to overwrite, the user must clear the
|
|
existing value first.
|
|
|
|
The config file is read/written as JSON. Malformed JSON is left untouched
|
|
and the helper returns ``False``. Note: Continue's modern config is
|
|
YAML-first; users on the YAML schema should configure systemMessage
|
|
through that file instead — this helper only handles the JSON variant.
|
|
|
|
Returns ``True`` if injection succeeded (or was already idempotent at
|
|
every site); ``False`` if any site refused or the file was malformed.
|
|
"""
|
|
if config_file.exists():
|
|
try:
|
|
content = _read_text(config_file)
|
|
except OSError as exc:
|
|
click.echo(f" Warning: could not read {config_file}: {exc}")
|
|
return False
|
|
if not content.strip():
|
|
data: dict[str, Any] = {}
|
|
else:
|
|
try:
|
|
parsed = json.loads(content)
|
|
except json.JSONDecodeError as exc:
|
|
click.echo(
|
|
f" Warning: {config_file} is not valid JSON ({exc.msg}); "
|
|
"not modifying — fix the file manually before re-running."
|
|
)
|
|
return False
|
|
if not isinstance(parsed, dict):
|
|
click.echo(
|
|
f" Warning: {config_file} top-level value is not an object; "
|
|
"Continue expects a JSON object — leaving file untouched."
|
|
)
|
|
return False
|
|
data = parsed
|
|
else:
|
|
data = {}
|
|
|
|
any_changed = False
|
|
all_ok = True
|
|
|
|
# 1. Top-level systemMessage.
|
|
changed, ok = _apply_rtk_to_systemmessage_field(
|
|
data, location_label=f"{config_file.name} (top-level)", verbose=verbose
|
|
)
|
|
any_changed = any_changed or changed
|
|
all_ok = all_ok and ok
|
|
|
|
# 2. Per-model systemMessage. Continue's models[] entry overrides the
|
|
# top-level value when set, so we must visit each one.
|
|
models = data.get("models")
|
|
if isinstance(models, list):
|
|
for idx, model in enumerate(models):
|
|
if not isinstance(model, dict):
|
|
continue
|
|
label = f"{config_file.name} models[{idx}]"
|
|
if isinstance(model.get("title"), str):
|
|
label = f"{config_file.name} models[{idx}] ({model['title']})"
|
|
changed_i, ok_i = _apply_rtk_to_systemmessage_field(
|
|
model, location_label=label, verbose=verbose
|
|
)
|
|
any_changed = any_changed or changed_i
|
|
all_ok = all_ok and ok_i
|
|
|
|
if any_changed:
|
|
config_file.parent.mkdir(parents=True, exist_ok=True)
|
|
_write_text(config_file, json.dumps(data, indent=2) + "\n")
|
|
click.echo(f" rtk instructions injected into {config_file}")
|
|
elif all_ok and verbose:
|
|
# Idempotent re-run with no refusals — nothing to do.
|
|
click.echo(f" rtk instructions already present in {config_file.name}")
|
|
|
|
return all_ok
|
|
|
|
|
|
def _resolve_copilot_provider_type(backend: str | None, provider_type: str) -> str:
|
|
"""Resolve Copilot BYOK provider type for the current proxy backend."""
|
|
return _copilot_resolve_provider_type(backend, provider_type)
|
|
|
|
|
|
def _query_proxy_config(port: int) -> dict[str, Any] | None:
|
|
"""Query the running proxy's feature configuration via /health.
|
|
|
|
Returns a dict with keys like backend, optimize, cache, rate_limit,
|
|
memory, learn, code_graph, pid. Returns None if unreachable or the
|
|
response lacks a config block.
|
|
"""
|
|
return _copilot_query_proxy_config(port)
|
|
|
|
|
|
def _query_proxy_health(port: int) -> dict[str, Any] | None:
|
|
"""Query the running proxy's full /health payload."""
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
url = f"http://127.0.0.1:{port}/health"
|
|
try:
|
|
with urllib.request.urlopen(url, timeout=2) as response:
|
|
payload = json.loads(response.read().decode("utf-8"))
|
|
except (OSError, urllib.error.URLError, ValueError, json.JSONDecodeError):
|
|
return None
|
|
return payload if isinstance(payload, dict) else None
|
|
|
|
|
|
def _proxy_health_config(payload: dict[str, Any] | None) -> dict[str, Any] | None:
|
|
"""Extract the config block from a Headroom /health payload."""
|
|
if payload is None:
|
|
return None
|
|
config = payload.get("config")
|
|
return config if isinstance(config, dict) else None
|
|
|
|
|
|
def _env_bool_value(value: str) -> bool:
|
|
return value.strip().lower() in {"1", "true", "yes", "on"}
|
|
|
|
|
|
def _agent_savings_config_mismatches(
|
|
running_config: dict[str, Any],
|
|
agent_type: str,
|
|
) -> list[str]:
|
|
"""Return restart reasons when a running proxy lacks target agent savings."""
|
|
|
|
if agent_type not in _AGENT_SAVINGS_TARGET_AGENTS:
|
|
return []
|
|
|
|
if _wrap_agent_savings_profile(agent_type) is None:
|
|
return []
|
|
|
|
desired_env = os.environ.copy()
|
|
apply_agent_savings_env_defaults(desired_env)
|
|
checks: tuple[tuple[str, str, str, str], ...] = (
|
|
("HEADROOM_SAVINGS_PROFILE", "savings_profile", "savings-profile", "str"),
|
|
("HEADROOM_TARGET_RATIO", "target_ratio", "target-ratio", "float"),
|
|
(
|
|
"HEADROOM_COMPRESS_USER_MESSAGES",
|
|
"compress_user_messages",
|
|
"compress-user-messages",
|
|
"bool",
|
|
),
|
|
(
|
|
"HEADROOM_COMPRESS_SYSTEM_MESSAGES",
|
|
"compress_system_messages",
|
|
"compress-system-messages",
|
|
"bool",
|
|
),
|
|
("HEADROOM_PROTECT_RECENT", "protect_recent", "protect-recent", "int"),
|
|
(
|
|
"HEADROOM_PROTECT_ANALYSIS_CONTEXT",
|
|
"protect_analysis_context",
|
|
"protect-analysis-context",
|
|
"bool",
|
|
),
|
|
("HEADROOM_MIN_TOKENS", "min_tokens_to_crush", "min-tokens", "int"),
|
|
("HEADROOM_MAX_ITEMS", "max_items_after_crush", "max-items", "int"),
|
|
(
|
|
"HEADROOM_SMART_CRUSHER_COMPACTION",
|
|
"smart_crusher_with_compaction",
|
|
"smart-crusher-compaction",
|
|
"bool",
|
|
),
|
|
("HEADROOM_ACCURACY_GUARD", "accuracy_guard", "accuracy-guard", "str"),
|
|
)
|
|
|
|
mismatches: list[str] = []
|
|
for env_key, config_key, label, value_type in checks:
|
|
expected = desired_env.get(env_key)
|
|
if expected is None:
|
|
continue
|
|
actual = running_config.get(config_key)
|
|
try:
|
|
if value_type == "float":
|
|
matches = actual is not None and abs(float(actual) - float(expected)) < 1e-9
|
|
elif value_type == "int":
|
|
matches = actual is not None and int(actual) == int(expected)
|
|
elif value_type == "bool":
|
|
matches = actual is not None and bool(actual) is _env_bool_value(expected)
|
|
else:
|
|
matches = str(actual or "").strip().lower() == expected.strip().lower()
|
|
except (TypeError, ValueError):
|
|
matches = False
|
|
if not matches:
|
|
mismatches.append(label)
|
|
|
|
return mismatches
|
|
|
|
|
|
def _proxy_active_session_count(payload: dict[str, Any] | None) -> int:
|
|
"""Return active session count from /health runtime metadata."""
|
|
if payload is None:
|
|
return 0
|
|
runtime = payload.get("runtime")
|
|
if not isinstance(runtime, dict):
|
|
return 0
|
|
websocket_sessions = runtime.get("websocket_sessions")
|
|
if not isinstance(websocket_sessions, dict):
|
|
return 0
|
|
counts = []
|
|
for key in ("active_sessions", "active_relay_tasks"):
|
|
value = websocket_sessions.get(key, 0)
|
|
if isinstance(value, int):
|
|
counts.append(value)
|
|
return max(counts, default=0)
|
|
|
|
|
|
def _normalize_proxy_api_url(url: object) -> str | None:
|
|
"""Normalize configured upstream URLs for running-proxy comparisons."""
|
|
if not isinstance(url, str):
|
|
return None
|
|
normalized = url.strip().rstrip("/")
|
|
if normalized.endswith("/v1"):
|
|
normalized = normalized[:-3]
|
|
return normalized or None
|
|
|
|
|
|
def _proxy_version(payload: dict[str, Any] | None) -> str | None:
|
|
"""Return the running proxy version when it exposes one."""
|
|
if payload is None:
|
|
return None
|
|
version = payload.get("version")
|
|
return version if isinstance(version, str) and version else None
|
|
|
|
|
|
def _proxy_needs_version_restart(payload: dict[str, Any] | None) -> bool:
|
|
"""Return True when a running Headroom proxy uses a different package version."""
|
|
running_version = _proxy_version(payload)
|
|
running_release = _normalize_release_version(running_version)
|
|
current_release = _normalize_release_version(_HEADROOM_VERSION)
|
|
return (
|
|
running_release is not None
|
|
and current_release is not None
|
|
and running_release != current_release
|
|
)
|
|
|
|
|
|
def _detect_running_proxy_backend(port: int) -> str | None:
|
|
"""Read the backend of an already-running proxy from its health endpoint."""
|
|
return _copilot_detect_running_proxy_backend(port)
|
|
|
|
|
|
def _kill_proxy_by_pid(pid: int, port: int) -> bool:
|
|
"""Terminate a proxy process by PID and wait for the port to free up.
|
|
|
|
Sends SIGTERM first, falls back to SIGKILL after 5 seconds.
|
|
Returns True if the port is free afterwards, False otherwise.
|
|
"""
|
|
try:
|
|
os.kill(pid, signal.SIGTERM)
|
|
except ProcessLookupError:
|
|
pass # Already gone
|
|
except PermissionError:
|
|
click.echo(f" Warning: No permission to kill proxy PID {pid}")
|
|
return False
|
|
|
|
# Wait for port to free (up to 5 seconds)
|
|
for _ in range(50):
|
|
time.sleep(0.1)
|
|
if not _check_proxy(port):
|
|
return True
|
|
|
|
# SIGTERM didn't work — escalate to SIGKILL (Unix) or terminate (Windows)
|
|
try:
|
|
_kill_signal = getattr(signal, "SIGKILL", signal.SIGTERM)
|
|
os.kill(pid, _kill_signal)
|
|
except (ProcessLookupError, PermissionError):
|
|
pass
|
|
|
|
for _ in range(20):
|
|
time.sleep(0.1)
|
|
if not _check_proxy(port):
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
def _stop_local_proxy_for_unwrap(port: int) -> str:
|
|
"""Stop a local Headroom proxy for durable unwrap commands.
|
|
|
|
Returns a status string:
|
|
* ``"stopped"``: a Headroom proxy was identified and stopped.
|
|
* ``"not_running"``: nothing is listening on the requested port.
|
|
* ``"unidentified"``: something is listening, but it did not expose
|
|
Headroom's health/config payload, so we did not kill it.
|
|
* ``"no_pid"``: the service looked like Headroom but did not expose a PID.
|
|
* ``"failed"``: a PID was found but the port stayed bound after stop.
|
|
"""
|
|
|
|
if not _check_proxy(port):
|
|
return "not_running"
|
|
|
|
running_config = _query_proxy_config(port)
|
|
if running_config is None:
|
|
return "unidentified"
|
|
|
|
proxy_pid = running_config.get("pid")
|
|
if proxy_pid is None:
|
|
return "no_pid"
|
|
|
|
try:
|
|
pid = int(proxy_pid)
|
|
except (TypeError, ValueError):
|
|
return "no_pid"
|
|
|
|
return "stopped" if _kill_proxy_by_pid(pid, port) else "failed"
|
|
|
|
|
|
def _echo_unwrap_proxy_stop_status(status: str, port: int) -> None:
|
|
"""Print a human-readable proxy stop result for unwrap commands."""
|
|
|
|
if status == "stopped":
|
|
click.echo(f" Stopped local Headroom proxy on port {port}.")
|
|
elif status == "not_running":
|
|
click.echo(f" No local Headroom proxy detected on port {port}.")
|
|
elif status == "unidentified":
|
|
click.echo(
|
|
f" Warning: port {port} is in use, but it did not look like Headroom; left it running."
|
|
)
|
|
elif status == "no_pid":
|
|
click.echo(
|
|
f" Warning: Headroom proxy on port {port} did not expose a PID; left it running."
|
|
)
|
|
else:
|
|
click.echo(f" Warning: failed to stop Headroom proxy on port {port}; stop it manually.")
|
|
|
|
|
|
def _find_persistent_manifest(port: int) -> Any:
|
|
"""Return a matching persistent deployment manifest for the requested port."""
|
|
from headroom.install.state import list_manifests
|
|
|
|
manifests = [manifest for manifest in list_manifests() if manifest.port == port]
|
|
manifests.sort(key=lambda manifest: (manifest.profile != "default", manifest.profile))
|
|
return manifests[0] if manifests else None
|
|
|
|
|
|
def _recover_persistent_proxy(port: int) -> bool:
|
|
"""Start or recover a matching persistent deployment for the requested port."""
|
|
from headroom.install.health import probe_ready
|
|
from headroom.install.models import InstallPreset, SupervisorKind
|
|
from headroom.install.runtime import start_detached_agent, start_persistent_docker, wait_ready
|
|
from headroom.install.supervisors import start_supervisor
|
|
|
|
manifest = _find_persistent_manifest(port)
|
|
if manifest is None:
|
|
return False
|
|
|
|
if probe_ready(manifest.health_url):
|
|
click.echo(f" Reusing persistent deployment '{manifest.profile}' on port {port}")
|
|
return True
|
|
|
|
if manifest.supervisor_kind == SupervisorKind.TASK.value:
|
|
click.echo(
|
|
f" Warning: task-based deployment '{manifest.profile}' cannot be auto-recovered via wrap"
|
|
)
|
|
return False
|
|
|
|
click.echo(f" Recovering persistent deployment '{manifest.profile}' on port {port}...")
|
|
try:
|
|
if manifest.preset == InstallPreset.PERSISTENT_DOCKER.value:
|
|
start_persistent_docker(manifest)
|
|
elif manifest.supervisor_kind == SupervisorKind.SERVICE.value:
|
|
start_supervisor(manifest)
|
|
else:
|
|
start_detached_agent(manifest.profile)
|
|
except Exception as exc:
|
|
click.echo(
|
|
f" Warning: could not recover persistent deployment '{manifest.profile}': {exc}"
|
|
)
|
|
return False
|
|
|
|
if wait_ready(manifest, timeout_seconds=45):
|
|
click.echo(f" Recovered persistent deployment '{manifest.profile}' on port {port}")
|
|
return True
|
|
|
|
click.echo(f" Warning: persistent deployment '{manifest.profile}' did not become ready")
|
|
return False
|
|
|
|
|
|
def _restart_persistent_proxy(manifest: Any, port: int) -> bool:
|
|
"""Restart a persistent deployment after an idle stale-version detection."""
|
|
from headroom.install.models import InstallPreset, SupervisorKind
|
|
from headroom.install.runtime import (
|
|
start_detached_agent,
|
|
start_persistent_docker,
|
|
stop_runtime,
|
|
wait_ready,
|
|
)
|
|
from headroom.install.supervisors import start_supervisor
|
|
|
|
click.echo(
|
|
f" Restarting persistent deployment '{manifest.profile}' "
|
|
f"with Headroom {_HEADROOM_VERSION}..."
|
|
)
|
|
try:
|
|
if manifest.preset == InstallPreset.PERSISTENT_DOCKER.value:
|
|
stop_runtime(manifest)
|
|
start_persistent_docker(manifest)
|
|
elif manifest.supervisor_kind == SupervisorKind.SERVICE.value:
|
|
# start_supervisor performs the platform-native restart operation:
|
|
# systemd restart, launchctl kickstart -k, or sc.exe start.
|
|
start_supervisor(manifest)
|
|
else:
|
|
stop_runtime(manifest)
|
|
start_detached_agent(manifest.profile)
|
|
except Exception as exc:
|
|
click.echo(
|
|
f" Warning: could not restart persistent deployment '{manifest.profile}': {exc}"
|
|
)
|
|
return False
|
|
|
|
if wait_ready(manifest, timeout_seconds=45):
|
|
click.echo(f" Restarted persistent deployment '{manifest.profile}' on port {port}")
|
|
return True
|
|
|
|
click.echo(f" Warning: persistent deployment '{manifest.profile}' did not become ready")
|
|
return False
|
|
|
|
|
|
def _copilot_model_configured(copilot_args: tuple[str, ...], env: dict[str, str]) -> bool:
|
|
"""Return True when Copilot BYOK model selection is configured."""
|
|
return _copilot_model_configured_impl(copilot_args, env)
|
|
|
|
|
|
def _copilot_model_from_args(copilot_args: tuple[str, ...], env: dict[str, str]) -> str | None:
|
|
"""Resolve the Copilot model from command-line args or environment."""
|
|
return _copilot_model_from_args_impl(copilot_args, env)
|
|
|
|
|
|
def _copilot_default_wire_api_for_model(model: str | None) -> str:
|
|
"""Return the default OpenAI-compatible wire API for a Copilot model."""
|
|
return _copilot_default_wire_api_for_model_impl(model)
|
|
|
|
|
|
def _should_use_copilot_oauth(
|
|
*,
|
|
backend: str | None,
|
|
provider_type: str,
|
|
env: dict[str, str],
|
|
force_subscription: bool = False,
|
|
) -> bool:
|
|
"""Prefer a reusable Copilot OAuth session when the requested routing supports it."""
|
|
if force_subscription:
|
|
return True
|
|
if env.get("COPILOT_PROVIDER_API_KEY") or env.get("COPILOT_PROVIDER_BEARER_TOKEN"):
|
|
return False
|
|
if provider_type == "anthropic":
|
|
return False
|
|
|
|
effective_backend = backend or os.environ.get("HEADROOM_BACKEND")
|
|
if effective_backend not in (None, "", "anthropic"):
|
|
return False
|
|
|
|
return has_oauth_auth()
|
|
|
|
|
|
def _push_runtime_env(port: int, no_proxy: bool) -> None:
|
|
"""Hot-sync this session's live env knobs to the proxy on ``port``.
|
|
|
|
Live knobs (the output-shaper family, the ast-grep read threshold) are read
|
|
from the *proxy's* process environment. A proxy we reused — rather than
|
|
started — would otherwise ignore values exported in this shell, since its
|
|
environment was snapshotted when it first launched. Pushing them to
|
|
``/admin/runtime-env`` applies them in memory with no disruptive restart.
|
|
|
|
Best-effort: a silent no-op when nothing is explicitly set, when there is no
|
|
proxy (``--no-proxy``), when the proxy is unreachable, or when it predates
|
|
the endpoint (older build returns 404).
|
|
"""
|
|
if no_proxy:
|
|
return
|
|
from headroom.proxy import runtime_env as _rt
|
|
|
|
payload = _rt.explicit_env(os.environ)
|
|
if not payload:
|
|
return
|
|
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
request = urllib.request.Request(
|
|
f"http://127.0.0.1:{port}/admin/runtime-env",
|
|
data=json.dumps(payload).encode("utf-8"),
|
|
method="POST",
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(request, timeout=2) as response:
|
|
response.read()
|
|
except (OSError, urllib.error.URLError, ValueError):
|
|
return
|
|
click.echo(f" Synced output settings to proxy: {', '.join(sorted(payload))}")
|
|
|
|
|
|
def _ensure_proxy(
|
|
port: int,
|
|
no_proxy: bool,
|
|
*,
|
|
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,
|
|
) -> tuple[subprocess.Popen | None, int]:
|
|
"""Start or verify proxy. Returns (process_handle, actual_port)."""
|
|
helpers = _live_wrap_module()
|
|
# --no-proxy reuses an already-running proxy, so backend/region/provider
|
|
# flags (which only apply when we start one) would be silently dropped.
|
|
if no_proxy and (backend or anyllm_provider or region):
|
|
click.echo(
|
|
" Warning: --backend/--region/--anyllm-provider have no effect with --no-proxy "
|
|
"(reusing the existing proxy)."
|
|
)
|
|
if not no_proxy:
|
|
manifest = helpers._find_persistent_manifest(port)
|
|
if manifest is not None:
|
|
from headroom.install.health import probe_ready
|
|
|
|
if probe_ready(manifest.health_url):
|
|
health_payload = helpers._query_proxy_health(port)
|
|
if helpers._proxy_needs_version_restart(health_payload):
|
|
running_version = helpers._proxy_version(health_payload) or "unknown"
|
|
active_sessions = helpers._proxy_active_session_count(health_payload)
|
|
other_wrappers = helpers._live_proxy_clients(port, exclude_self=True)
|
|
if active_sessions > 0 or other_wrappers:
|
|
detail = (
|
|
f"{active_sessions} active session(s)"
|
|
if active_sessions > 0
|
|
else f"{len(other_wrappers)} attached wrapper(s)"
|
|
)
|
|
click.echo(
|
|
f" Proxy on port {port} is running Headroom {running_version}; "
|
|
f"current CLI is {_HEADROOM_VERSION}."
|
|
)
|
|
click.echo(
|
|
f" Leaving it running because {detail} "
|
|
"are still attached; it will be restarted when idle."
|
|
)
|
|
return None, port
|
|
if helpers._restart_persistent_proxy(manifest, port):
|
|
return None, port
|
|
raise click.ClickException(
|
|
f"Persistent deployment '{manifest.profile}' on port {port} "
|
|
f"is running stale Headroom {running_version} and could not be restarted."
|
|
)
|
|
# Check if the running proxy has the features we need.
|
|
# Without this, a persistent deployment started for one use case
|
|
# (e.g. --backend anthropic) would be silently reused for another
|
|
# (e.g. --subscription --provider-type openai) causing auth failures.
|
|
running_config = helpers._proxy_health_config(health_payload)
|
|
if running_config is None:
|
|
running_config = helpers._query_proxy_config(port)
|
|
if running_config is not None:
|
|
missing = []
|
|
if memory and not running_config.get("memory"):
|
|
missing.append("memory")
|
|
if learn and not running_config.get("learn"):
|
|
missing.append("learn")
|
|
if code_graph and not running_config.get("code_graph"):
|
|
missing.append("code_graph")
|
|
if openai_api_url:
|
|
running_openai_url = _normalize_proxy_api_url(
|
|
running_config.get("openai_api_url")
|
|
)
|
|
requested_openai_url = _normalize_proxy_api_url(openai_api_url)
|
|
if running_openai_url != requested_openai_url:
|
|
missing.append("openai-api-url")
|
|
if not missing:
|
|
click.echo(f" Proxy already running on port {port}")
|
|
click.echo(f" Dashboard: http://127.0.0.1:{port}/dashboard")
|
|
return None, port
|
|
# Features mismatch or config unavailable — fall through to
|
|
# the non-persistent path which handles proxy restart.
|
|
else:
|
|
if helpers._recover_persistent_proxy(port):
|
|
# If the caller requested feature-sensitive config (e.g.
|
|
# openai_api_url for Copilot subscription), continue into
|
|
# the shared running-proxy checks below so mismatch-driven
|
|
# restart logic can run. For plain recover-only calls,
|
|
# preserve the historical fast return.
|
|
if not any((memory, learn, code_graph, openai_api_url)):
|
|
return None, port
|
|
if not helpers._check_proxy(port):
|
|
return None, port
|
|
|
|
# A freshly recovered persistent proxy may not expose
|
|
# a full config payload yet. In feature-sensitive flows
|
|
# (e.g. Copilot subscription), treat missing or mismatched
|
|
# config as restart-required and refresh the persistent
|
|
# deployment directly instead of silently reusing it.
|
|
health_payload = helpers._query_proxy_health(port)
|
|
running_config = helpers._proxy_health_config(health_payload)
|
|
if running_config is None:
|
|
running_config = helpers._query_proxy_config(port)
|
|
|
|
if running_config is None:
|
|
click.echo(
|
|
f" Recovered persistent deployment '{manifest.profile}' "
|
|
"did not expose config; restarting with requested features..."
|
|
)
|
|
if helpers._restart_persistent_proxy(manifest, port):
|
|
return None, port
|
|
raise click.ClickException(
|
|
f"Persistent deployment '{manifest.profile}' on port {port} "
|
|
"could not be restarted after recovery."
|
|
)
|
|
|
|
missing = []
|
|
if memory and not running_config.get("memory"):
|
|
missing.append("memory")
|
|
if learn and not running_config.get("learn"):
|
|
missing.append("learn")
|
|
if code_graph and not running_config.get("code_graph"):
|
|
missing.append("code-graph")
|
|
if openai_api_url:
|
|
running_openai_url = _normalize_proxy_api_url(
|
|
running_config.get("openai_api_url")
|
|
)
|
|
requested_openai_url = _normalize_proxy_api_url(openai_api_url)
|
|
if running_openai_url != requested_openai_url:
|
|
missing.append("openai-api-url")
|
|
|
|
if missing:
|
|
flags_str = ", ".join(f"--{f}" for f in missing)
|
|
click.echo(
|
|
f" Recovered persistent deployment '{manifest.profile}' is missing: "
|
|
f"{flags_str}; restarting..."
|
|
)
|
|
if helpers._restart_persistent_proxy(manifest, port):
|
|
return None, port
|
|
raise click.ClickException(
|
|
f"Persistent deployment '{manifest.profile}' on port {port} "
|
|
"could not be restarted with requested features."
|
|
)
|
|
return None, port
|
|
elif helpers._check_proxy(port):
|
|
raise click.ClickException(
|
|
f"Persistent deployment '{manifest.profile}' on port {port} is not healthy."
|
|
)
|
|
click.echo(
|
|
f" Warning: persistent deployment '{manifest.profile}' on port {port} "
|
|
"is stale; starting a fresh proxy instead."
|
|
)
|
|
|
|
if helpers._check_proxy(port):
|
|
# Proxy is running — check if it has the features we need
|
|
needs_restart = False
|
|
health_payload = helpers._query_proxy_health(port)
|
|
running_config = helpers._proxy_health_config(health_payload)
|
|
if running_config is None:
|
|
running_config = helpers._query_proxy_config(port)
|
|
|
|
if helpers._proxy_needs_version_restart(health_payload):
|
|
running_version = helpers._proxy_version(health_payload) or "unknown"
|
|
active_sessions = helpers._proxy_active_session_count(health_payload)
|
|
other_wrappers = helpers._live_proxy_clients(port, exclude_self=True)
|
|
if active_sessions > 0 or other_wrappers:
|
|
# active_sessions only counts Codex WebSocket relay; the
|
|
# marker list also covers HTTP wrap clients. Either means a
|
|
# live session is attached, so don't restart the shared
|
|
# proxy out from under it — defer until idle.
|
|
detail = (
|
|
f"{active_sessions} active session(s)"
|
|
if active_sessions > 0
|
|
else f"{len(other_wrappers)} attached wrapper(s)"
|
|
)
|
|
click.echo(
|
|
f" Proxy on port {port} is running Headroom {running_version}; "
|
|
f"current CLI is {_HEADROOM_VERSION}."
|
|
)
|
|
click.echo(
|
|
f" Leaving it running because {detail} "
|
|
"are still attached; it will be restarted when idle."
|
|
)
|
|
return None, port
|
|
|
|
click.echo(
|
|
f" Proxy on port {port} is running Headroom {running_version}; "
|
|
f"restarting with {_HEADROOM_VERSION}..."
|
|
)
|
|
proxy_pid = running_config.get("pid") if running_config is not None else None
|
|
if proxy_pid is None:
|
|
raise click.ClickException(
|
|
f"Proxy on port {port} is stale but did not expose a PID. "
|
|
"Stop it manually and retry."
|
|
)
|
|
if not helpers._kill_proxy_by_pid(int(proxy_pid), port):
|
|
raise click.ClickException(
|
|
f"Failed to stop stale proxy (PID {proxy_pid}) on port {port}. "
|
|
"Stop it manually and retry."
|
|
)
|
|
needs_restart = True
|
|
|
|
if running_config is not None:
|
|
missing = []
|
|
if memory and not running_config.get("memory"):
|
|
missing.append("memory")
|
|
if learn and not running_config.get("learn"):
|
|
missing.append("learn")
|
|
if code_graph and not running_config.get("code_graph"):
|
|
missing.append("code_graph")
|
|
expected_savings_profile = helpers._wrap_agent_savings_profile(agent_type)
|
|
if (
|
|
expected_savings_profile is not None
|
|
and running_config.get("savings_profile") != expected_savings_profile
|
|
):
|
|
missing.append("savings-profile")
|
|
if openai_api_url:
|
|
running_openai_url = _normalize_proxy_api_url(
|
|
running_config.get("openai_api_url")
|
|
)
|
|
requested_openai_url = _normalize_proxy_api_url(openai_api_url)
|
|
if running_openai_url != requested_openai_url:
|
|
missing.append("openai-api-url")
|
|
if vertex_api_url or clear_vertex_api_url:
|
|
running_vertex_url = _normalize_proxy_api_url(
|
|
running_config.get("vertex_api_url")
|
|
)
|
|
requested_vertex_url = _normalize_proxy_api_url(vertex_api_url)
|
|
if running_vertex_url != requested_vertex_url:
|
|
missing.append("vertex-api-url")
|
|
|
|
if missing:
|
|
flags_str = ", ".join(
|
|
f if f.startswith("--") else f"--{f.replace('_', '-')}" for f in missing
|
|
)
|
|
other_wrappers = helpers._live_proxy_clients(port, exclude_self=True)
|
|
if other_wrappers:
|
|
# Another wrapper is attached to this proxy; restarting it
|
|
# to add flags would drop their in-flight requests. Reuse
|
|
# the running proxy as-is rather than disrupt them.
|
|
click.echo(
|
|
f" Proxy on port {port} is missing: {flags_str}, but "
|
|
f"{len(other_wrappers)} other wrapper(s) are attached."
|
|
)
|
|
click.echo(
|
|
" Leaving it running to avoid disrupting them; this "
|
|
"session will use the existing proxy as-is."
|
|
)
|
|
else:
|
|
needs_restart = True
|
|
click.echo(f" Proxy on port {port} is missing: {flags_str}")
|
|
click.echo(" Restarting proxy with upgraded configuration...")
|
|
|
|
# Merge: keep features the running proxy already has
|
|
memory = memory or bool(running_config.get("memory"))
|
|
learn = learn or bool(running_config.get("learn"))
|
|
code_graph = code_graph or bool(running_config.get("code_graph"))
|
|
|
|
proxy_pid = running_config.get("pid")
|
|
if proxy_pid is not None:
|
|
if not helpers._kill_proxy_by_pid(int(proxy_pid), port):
|
|
raise click.ClickException(
|
|
f"Failed to stop existing proxy (PID {proxy_pid}) on port {port}. "
|
|
"Stop it manually and retry."
|
|
)
|
|
else:
|
|
click.echo(
|
|
" Warning: Running proxy does not expose PID. "
|
|
"Cannot restart automatically."
|
|
)
|
|
click.echo(
|
|
f" Please stop the proxy on port {port} manually "
|
|
f"and rerun with {flags_str}."
|
|
)
|
|
return None, port
|
|
|
|
if not needs_restart:
|
|
click.echo(f" Proxy already running on port {port}")
|
|
click.echo(f" Dashboard: http://127.0.0.1:{port}/dashboard")
|
|
return None, port
|
|
|
|
# Start (or restart) the proxy with the requested flags
|
|
# Find an available port (port may be busy from a stale proxy).
|
|
try:
|
|
actual_port = helpers._find_available_port(port)
|
|
except OSError as e:
|
|
raise click.ClickException(f"Port {port} is unavailable: {e}") from e
|
|
except RuntimeError as e:
|
|
raise click.ClickException(str(e)) from e
|
|
|
|
if actual_port != port:
|
|
click.echo(f" Port {port} is in use, using port {actual_port} instead.")
|
|
|
|
click.echo(f" Starting Headroom proxy on port {actual_port}...")
|
|
try:
|
|
proc = cast(
|
|
subprocess.Popen[Any],
|
|
_live_wrap_module()._start_proxy(
|
|
actual_port,
|
|
learn=learn,
|
|
memory=memory,
|
|
agent_type=agent_type,
|
|
code_graph=code_graph,
|
|
backend=backend,
|
|
anyllm_provider=anyllm_provider,
|
|
region=region,
|
|
openai_api_url=openai_api_url,
|
|
anthropic_api_url=anthropic_api_url,
|
|
vertex_api_url=vertex_api_url,
|
|
clear_vertex_api_url=clear_vertex_api_url,
|
|
copilot_api_token=copilot_api_token,
|
|
),
|
|
)
|
|
click.echo(f" Proxy ready on http://127.0.0.1:{actual_port}")
|
|
click.echo(f" Dashboard: http://127.0.0.1:{actual_port}/dashboard")
|
|
return proc, actual_port
|
|
except RuntimeError as e:
|
|
click.echo(f" Error: {e}")
|
|
raise SystemExit(1) from e
|
|
else:
|
|
if not helpers._check_proxy(port):
|
|
click.echo(f" Warning: No proxy detected on port {port}")
|
|
elif vertex_api_url or clear_vertex_api_url:
|
|
health_payload = helpers._query_proxy_health(port)
|
|
running_config = helpers._proxy_health_config(health_payload)
|
|
if running_config is None:
|
|
running_config = helpers._query_proxy_config(port)
|
|
running_vertex_url = (
|
|
_normalize_proxy_api_url(running_config.get("vertex_api_url"))
|
|
if running_config is not None
|
|
else None
|
|
)
|
|
requested_vertex_url = _normalize_proxy_api_url(vertex_api_url)
|
|
if running_vertex_url != requested_vertex_url:
|
|
click.echo(
|
|
" Warning: --no-proxy is set, but the running proxy does not "
|
|
"advertise the requested Vertex target. Requests may still go "
|
|
"to the proxy's existing Vertex upstream."
|
|
)
|
|
return None, port
|
|
|
|
|
|
def _client_marker_path(port: int) -> Path:
|
|
"""Path to this process's wrap-client marker for ``port``."""
|
|
from headroom import paths as _paths
|
|
|
|
d = _paths.proxy_clients_dir(port)
|
|
d.mkdir(parents=True, exist_ok=True)
|
|
return d / f"{os.getpid()}.json"
|
|
|
|
|
|
def _proc_identity(pid: int) -> tuple[str, float] | None:
|
|
"""Best-effort ``(source, start_time)`` identity for a PID.
|
|
|
|
Used to defeat PID reuse: a marker is only trusted while the live PID is
|
|
*the same process* that wrote it. Returns ``None`` when start time can't be
|
|
determined (e.g. macOS without psutil), in which case callers fall back to
|
|
existence-only liveness — no regression, just no reuse protection there.
|
|
|
|
The ``source`` tag ("psutil" vs "proc") guards against comparing values in
|
|
different units; we only compare like-for-like.
|
|
"""
|
|
try:
|
|
import psutil # type: ignore[import-untyped] # optional dependency; portable when present
|
|
|
|
return ("psutil", psutil.Process(pid).create_time())
|
|
except Exception:
|
|
pass
|
|
# Linux fallback: field 22 of /proc/<pid>/stat is starttime in clock ticks
|
|
# since boot — a stable per-process value. `comm` (field 2) may contain
|
|
# spaces/parens, so split after the final ')'.
|
|
try:
|
|
with open(f"/proc/{pid}/stat", "rb") as fh:
|
|
fields = fh.read().rpartition(b")")[2].split()
|
|
return ("proc", float(fields[19]))
|
|
except (OSError, IndexError, ValueError):
|
|
return None
|
|
|
|
|
|
def _register_proxy_client(port: int) -> None:
|
|
"""Register this wrap process as a live client of the shared proxy.
|
|
|
|
Best-effort: a failed write just means our marker is missing, and the
|
|
liveness pruning in :func:`_live_proxy_clients` is the real safety net.
|
|
"""
|
|
try:
|
|
payload: dict[str, Any] = {"pid": os.getpid(), "started_at": time.time()}
|
|
ident = _proc_identity(os.getpid())
|
|
if ident is not None:
|
|
payload["start_src"], payload["start_time"] = ident
|
|
_write_text(_client_marker_path(port), json.dumps(payload))
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def _unregister_proxy_client(port: int) -> None:
|
|
"""Remove this process's client marker (idempotent)."""
|
|
try:
|
|
_client_marker_path(port).unlink(missing_ok=True)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def _pid_alive(pid: int) -> bool:
|
|
"""Return True if ``pid`` names a live process.
|
|
|
|
Thin wrapper over the shared Windows-safe helper so the marker-cleanup path
|
|
and the install/runtime status path use one liveness probe (see #1544).
|
|
"""
|
|
return pid_alive(pid)
|
|
|
|
|
|
def _identity_mismatch(src: Any, recorded: Any, pid: int) -> bool:
|
|
"""True only if ``pid``'s current identity *provably* differs from the
|
|
recorded ``(src, recorded)`` identity (i.e. the PID was recycled).
|
|
|
|
Conservative by design: any uncertainty (unknown/legacy identity, unknown
|
|
start time, mismatched source) returns ``False`` — never claim a mismatch
|
|
without proof, since the caller uses this to decide whether to trust or
|
|
discard state tied to a live PID.
|
|
"""
|
|
if not isinstance(src, str) or not isinstance(recorded, int | float):
|
|
return False # legacy / identity-less record — can't tell
|
|
ident = _proc_identity(pid)
|
|
if ident is None or ident[0] != src:
|
|
return False # can't compare like-for-like — don't claim mismatch
|
|
# Start times are stable per process; >1s apart means a different process.
|
|
return abs(ident[1] - float(recorded)) > 1.0
|
|
|
|
|
|
def _marker_pid_reused(marker: Path, pid: int) -> bool:
|
|
"""True only if the live ``pid`` is *provably* a different process than the
|
|
one that wrote ``marker`` (i.e. the PID was recycled after a crash).
|
|
"""
|
|
try:
|
|
rec = json.loads(_read_text(marker))
|
|
except (OSError, ValueError):
|
|
return False
|
|
return _identity_mismatch(rec.get("start_src"), rec.get("start_time"), pid)
|
|
|
|
|
|
def _live_proxy_clients(port: int, *, exclude_self: bool = True) -> list[int]:
|
|
"""Live wrap-client PIDs for ``port``, pruning stale markers as we go."""
|
|
from headroom import paths as _paths
|
|
|
|
d = _paths.proxy_clients_dir(port)
|
|
if not d.exists():
|
|
return []
|
|
me = os.getpid()
|
|
live: list[int] = []
|
|
for marker in d.glob("*.json"):
|
|
try:
|
|
pid = int(marker.stem)
|
|
except ValueError:
|
|
continue
|
|
# Stale if the PID is gone, or recycled by an unrelated process.
|
|
if not _pid_alive(pid) or _marker_pid_reused(marker, pid):
|
|
try:
|
|
marker.unlink(missing_ok=True)
|
|
except OSError:
|
|
pass
|
|
continue
|
|
if not (exclude_self and pid == me):
|
|
live.append(pid)
|
|
return live
|
|
|
|
|
|
def _make_cleanup(proxy_proc_holder: list, port: int | list[int] = 8787) -> Any:
|
|
"""Create a cleanup function that terminates the proxy on exit.
|
|
|
|
Only kills the proxy when no other live headroom-wrapped clients remain,
|
|
tracked via per-PID marker files in ``paths.proxy_clients_dir(port)``.
|
|
|
|
``port`` can be an ``int`` or a ``list[int]``. When a port fallback occurs
|
|
(``_ensure_proxy`` ups the port because the requested one is busy), the
|
|
caller can update ``port[0]`` in-place and the closure picks it up.
|
|
"""
|
|
|
|
def _other_clients_exist() -> bool:
|
|
p = port[0] if isinstance(port, list) else port
|
|
return len(_live_proxy_clients(p, exclude_self=True)) > 0
|
|
|
|
def cleanup(signum: int | None = None, frame: Any = None) -> None:
|
|
p = port[0] if isinstance(port, list) else port
|
|
_unregister_proxy_client(p)
|
|
proc = proxy_proc_holder[0] if proxy_proc_holder else None
|
|
if proc and proc.poll() is None:
|
|
if _other_clients_exist():
|
|
# Other clients still using the proxy — leave it running.
|
|
return
|
|
proc.terminate()
|
|
try:
|
|
proc.wait(timeout=5)
|
|
except subprocess.TimeoutExpired:
|
|
proc.kill()
|
|
|
|
return cleanup
|
|
|
|
|
|
def _ignore_child_sigint(signum: int | None = None, frame: Any = None) -> None:
|
|
"""Keep the wrapper alive when Ctrl-C is intended for the child CLI."""
|
|
|
|
return None
|
|
|
|
|
|
def _launch_tool(
|
|
binary: str,
|
|
args: tuple,
|
|
env: dict[str, str],
|
|
port: int,
|
|
no_proxy: bool,
|
|
tool_label: str,
|
|
env_vars_display: list[str],
|
|
*,
|
|
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,
|
|
copilot_api_token: str | None = None,
|
|
) -> None:
|
|
"""Common logic: start proxy, launch tool, clean up."""
|
|
proxy_holder: list[subprocess.Popen | None] = [None]
|
|
port_holder: list[int] = [port]
|
|
cleanup = _make_cleanup(proxy_holder, port_holder)
|
|
signal.signal(signal.SIGINT, _ignore_child_sigint)
|
|
signal.signal(signal.SIGTERM, cleanup)
|
|
|
|
try:
|
|
click.echo()
|
|
padded = f"HEADROOM WRAP: {tool_label}".center(47)
|
|
click.echo(" ╔═══════════════════════════════════════════════╗")
|
|
click.echo(f" ║{padded}║")
|
|
click.echo(" ╚═══════════════════════════════════════════════╝")
|
|
click.echo()
|
|
|
|
_register_proxy_client(port)
|
|
proxy_holder[0], actual_port = _ensure_proxy(
|
|
port,
|
|
no_proxy,
|
|
learn=learn,
|
|
memory=memory,
|
|
agent_type=agent_type,
|
|
code_graph=code_graph,
|
|
backend=backend,
|
|
anyllm_provider=anyllm_provider,
|
|
region=region,
|
|
openai_api_url=openai_api_url,
|
|
copilot_api_token=copilot_api_token,
|
|
)
|
|
if actual_port != port:
|
|
_unregister_proxy_client(port)
|
|
_register_proxy_client(actual_port)
|
|
port_holder[0] = actual_port
|
|
_push_runtime_env(actual_port, no_proxy)
|
|
|
|
# If port fell back, update env URLs to point at the actual port
|
|
if actual_port != port:
|
|
for k, v in dict(env).items():
|
|
env[k] = v.replace(f"127.0.0.1:{port}", f"127.0.0.1:{actual_port}")
|
|
|
|
if code_graph:
|
|
_setup_code_graph(verbose=False)
|
|
|
|
click.echo()
|
|
click.echo(f" Launching {tool_label} (API routed through Headroom)...")
|
|
for var in env_vars_display:
|
|
click.echo(f" {var}")
|
|
if args:
|
|
click.echo(f" Extra args: {' '.join(args)}")
|
|
_print_telemetry_notice()
|
|
click.echo()
|
|
|
|
result = subprocess.run([binary, *args], env=env)
|
|
raise SystemExit(result.returncode)
|
|
|
|
except SystemExit:
|
|
raise
|
|
except Exception as e:
|
|
click.echo(f" Error: {e}")
|
|
raise SystemExit(1) from e
|
|
finally:
|
|
cleanup()
|
|
|
|
|
|
def _run_checked(
|
|
cmd: list[str],
|
|
*,
|
|
cwd: Path | None = None,
|
|
action: str,
|
|
) -> subprocess.CompletedProcess[str]:
|
|
"""Run subprocess and raise a ClickException with actionable context on failure."""
|
|
try:
|
|
return run(
|
|
cmd,
|
|
cwd=str(cwd) if cwd else None,
|
|
check=True,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
except FileNotFoundError as e:
|
|
raise click.ClickException(f"{action} failed: command not found: {cmd[0]}") from e
|
|
except subprocess.CalledProcessError as e:
|
|
stderr = (e.stderr or "").strip()
|
|
stdout = (e.stdout or "").strip()
|
|
details = stderr or stdout or f"exit code {e.returncode}"
|
|
raise click.ClickException(f"{action} failed: {details}") from e
|
|
|
|
|
|
def _resolve_openclaw_extensions_dir(openclaw_bin: str) -> Path:
|
|
"""Resolve OpenClaw extension root from active config file path."""
|
|
result = _run_checked([openclaw_bin, "config", "file"], action="openclaw config file")
|
|
lines = result.stdout.strip().splitlines()
|
|
config_path_str = lines[-1].strip() if lines else ""
|
|
if not config_path_str:
|
|
raise click.ClickException(
|
|
"Unable to resolve OpenClaw config path from `openclaw config file`."
|
|
)
|
|
config_path = Path(config_path_str).expanduser()
|
|
return config_path.parent / "extensions"
|
|
|
|
|
|
def _normalize_openclaw_gateway_provider_ids(provider_ids: tuple[str, ...] | None) -> list[str]:
|
|
"""Normalize configured OpenClaw provider ids, defaulting to openai-codex."""
|
|
return _normalize_openclaw_gateway_provider_ids_impl(provider_ids)
|
|
|
|
|
|
def _read_openclaw_config_value(openclaw_bin: str, path: str) -> Any | None:
|
|
"""Read an OpenClaw config value when present, returning None on missing paths."""
|
|
result = run(
|
|
[openclaw_bin, "config", "get", path],
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
if result.returncode != 0:
|
|
return None
|
|
|
|
output = result.stdout.strip()
|
|
if not output:
|
|
return None
|
|
|
|
try:
|
|
return json.loads(output)
|
|
except json.JSONDecodeError:
|
|
return output
|
|
|
|
|
|
def _decode_openclaw_entry_json(raw_value: str | None) -> Any | None:
|
|
"""Decode a JSON payload captured from `openclaw config get` when available."""
|
|
return _decode_openclaw_entry_json_impl(raw_value)
|
|
|
|
|
|
def _build_openclaw_plugin_entry(
|
|
*,
|
|
existing_entry: Any,
|
|
proxy_port: int,
|
|
startup_timeout_ms: int,
|
|
python_path: str | None,
|
|
no_auto_start: bool,
|
|
gateway_provider_ids: tuple[str, ...] | None,
|
|
enabled: bool,
|
|
) -> dict[str, object]:
|
|
"""Merge managed Headroom plugin settings with any existing entry payload."""
|
|
return _build_openclaw_plugin_entry_impl(
|
|
existing_entry=existing_entry,
|
|
proxy_port=proxy_port,
|
|
startup_timeout_ms=startup_timeout_ms,
|
|
python_path=python_path,
|
|
no_auto_start=no_auto_start,
|
|
gateway_provider_ids=gateway_provider_ids,
|
|
enabled=enabled,
|
|
)
|
|
|
|
|
|
def _build_openclaw_unwrap_entry(existing_entry: Any) -> dict[str, object]:
|
|
"""Disable the managed plugin while preserving unrelated user config."""
|
|
return _build_openclaw_unwrap_entry_impl(existing_entry)
|
|
|
|
|
|
def _write_openclaw_plugin_entry(openclaw_bin: str, entry: dict[str, object]) -> None:
|
|
"""Persist the Headroom plugin config entry."""
|
|
_run_checked(
|
|
[
|
|
openclaw_bin,
|
|
"config",
|
|
"set",
|
|
"plugins.entries.headroom",
|
|
json.dumps(entry, separators=(",", ":")),
|
|
"--strict-json",
|
|
],
|
|
action="openclaw config set plugins.entries.headroom",
|
|
)
|
|
|
|
|
|
def _set_openclaw_context_engine_slot(openclaw_bin: str, engine_id: str) -> None:
|
|
"""Persist the selected OpenClaw context engine slot."""
|
|
_run_checked(
|
|
[
|
|
openclaw_bin,
|
|
"config",
|
|
"set",
|
|
"plugins.slots.contextEngine",
|
|
json.dumps(engine_id),
|
|
"--strict-json",
|
|
],
|
|
action="openclaw config set plugins.slots.contextEngine",
|
|
)
|
|
|
|
|
|
def _restart_or_start_openclaw_gateway(openclaw_bin: str) -> tuple[str, str]:
|
|
"""Restart the gateway when running, otherwise start it."""
|
|
restart_result = run(
|
|
[openclaw_bin, "gateway", "restart"],
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
if restart_result.returncode == 0:
|
|
output = restart_result.stdout.strip() or restart_result.stderr.strip()
|
|
return "restarted", output
|
|
|
|
start_result = _run_checked(
|
|
[openclaw_bin, "gateway", "start"],
|
|
action="openclaw gateway start",
|
|
)
|
|
output = start_result.stdout.strip() or start_result.stderr.strip()
|
|
return "started", output
|
|
|
|
|
|
def _copy_openclaw_plugin_into_extensions(
|
|
*,
|
|
plugin_dir: Path,
|
|
openclaw_bin: str,
|
|
) -> Path:
|
|
"""Fallback install path when `openclaw plugins install` is blocked on linked source."""
|
|
dist_dir = plugin_dir / "dist"
|
|
if not dist_dir.exists():
|
|
raise click.ClickException(
|
|
f"Plugin dist folder missing at {dist_dir}. Build the plugin first."
|
|
)
|
|
hook_shim_dir = plugin_dir / "hook-shim"
|
|
if not hook_shim_dir.exists():
|
|
raise click.ClickException(
|
|
f"Plugin hook-shim folder missing at {hook_shim_dir}. Build the plugin first."
|
|
)
|
|
|
|
extensions_dir = _resolve_openclaw_extensions_dir(openclaw_bin)
|
|
target_dir = extensions_dir / "headroom"
|
|
target_dist = target_dir / "dist"
|
|
target_hook_shim = target_dir / "hook-shim"
|
|
target_dir.mkdir(parents=True, exist_ok=True)
|
|
if target_dist.exists():
|
|
shutil.rmtree(target_dist)
|
|
if target_hook_shim.exists():
|
|
shutil.rmtree(target_hook_shim)
|
|
shutil.copytree(dist_dir, target_dist)
|
|
shutil.copytree(hook_shim_dir, target_hook_shim)
|
|
|
|
for filename in ("openclaw.plugin.json", "package.json", "README.md"):
|
|
source = plugin_dir / filename
|
|
if source.exists():
|
|
shutil.copy2(source, target_dir / filename)
|
|
|
|
return target_dir
|
|
|
|
|
|
@main.group()
|
|
def wrap() -> None:
|
|
"""Wrap CLI tools to run through Headroom.
|
|
|
|
\b
|
|
Starts a Headroom proxy, configures the environment, and launches
|
|
the target tool so all API calls route through Headroom automatically.
|
|
|
|
\b
|
|
Supported tools (one Click subcommand per tool):
|
|
headroom wrap claude # Claude Code (Anthropic)
|
|
headroom wrap codex # OpenAI Codex CLI
|
|
headroom wrap copilot -- --model claude-sonnet-4-20250514
|
|
headroom wrap aider # Aider
|
|
headroom wrap vibe # Mistral Vibe
|
|
headroom wrap cursor # Cursor (prints config instructions)
|
|
headroom wrap cline # Cline (VS Code; prints config instructions)
|
|
headroom wrap continue # Continue (VS Code/JetBrains; injects systemMessage)
|
|
headroom wrap goose # Goose (Block) CLI
|
|
headroom wrap openhands # OpenHands CLI
|
|
headroom wrap openclaw # OpenClaw plugin bootstrap
|
|
headroom wrap opencode # OpenCode CLI
|
|
|
|
\b
|
|
`wrap` vs `proxy`:
|
|
- `headroom wrap <tool>` — convenience: starts the proxy for you,
|
|
sets the right env vars, and launches the wrapped CLI.
|
|
- `headroom proxy` — just the proxy. Use this with any
|
|
OpenAI/Anthropic-compatible client by setting
|
|
ANTHROPIC_BASE_URL / OPENAI_BASE_URL yourself.
|
|
|
|
\b
|
|
`openclaw` is a separate tool — different from opencode.
|
|
"""
|
|
|
|
|
|
@main.group()
|
|
def unwrap() -> None:
|
|
"""Undo durable Headroom wrapping for supported tools."""
|
|
|
|
|
|
# =============================================================================
|
|
# Claude Code
|
|
# =============================================================================
|
|
|
|
|
|
@wrap.command(context_settings={"ignore_unknown_options": True})
|
|
@click.option(
|
|
"--port", "-p", default=8787, type=click.IntRange(1, 65535), help="Proxy port (default: 8787)"
|
|
)
|
|
@click.option(
|
|
"--no-context-tool",
|
|
"--no-rtk",
|
|
"no_rtk",
|
|
is_flag=True,
|
|
help="Skip CLI context-tool setup",
|
|
)
|
|
@click.option(
|
|
"--context-tool",
|
|
"context_tool",
|
|
is_flag=True,
|
|
help="Enable CLI context-tool setup",
|
|
)
|
|
@click.option(
|
|
"--no-mcp",
|
|
is_flag=True,
|
|
help="Skip headroom MCP server registration (compression markers will be unactionable)",
|
|
)
|
|
@click.option(
|
|
"--no-tokensave",
|
|
is_flag=True,
|
|
help="Skip the tokensave code-graph MCP server (primary coding-task compressor)",
|
|
)
|
|
@click.option(
|
|
"--serena",
|
|
is_flag=True,
|
|
help="Force the Serena MCP backup compressor on (registered automatically when "
|
|
"tokensave is unavailable)",
|
|
)
|
|
@click.option("--no-serena", is_flag=True, help="Never register the Serena backup compressor")
|
|
@click.option(
|
|
"--code-graph",
|
|
is_flag=True,
|
|
help="Force a tokensave code-graph index now (tokensave is the default compressor)",
|
|
)
|
|
@click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)")
|
|
@click.option(
|
|
"--learn", is_flag=True, help="Enable live traffic learning (patterns saved to MEMORY.md)"
|
|
)
|
|
@click.option("--memory", is_flag=True, help="Enable persistent cross-session memory")
|
|
@click.option(
|
|
"--tool-search",
|
|
"tool_search",
|
|
default=None,
|
|
metavar="MODE",
|
|
help=(
|
|
"Keep Claude Code's on-demand tool loading (deferral) active through the "
|
|
"proxy. MODE is true (default), auto, auto:N, or false. Without it, a "
|
|
"custom ANTHROPIC_BASE_URL makes Claude Code load every tool schema "
|
|
"eagerly, inflating local context (issue #746). A pre-set "
|
|
"ENABLE_TOOL_SEARCH env var is respected."
|
|
),
|
|
)
|
|
@click.option(
|
|
"--backend",
|
|
default=None,
|
|
help="API backend for the proxy: 'anthropic' (default), 'litellm-vertex_ai', etc. "
|
|
"(env: HEADROOM_BACKEND). For Vertex, prefer CLAUDE_CODE_USE_VERTEX=1 (native, "
|
|
"keeps your GCP auth) over a litellm backend.",
|
|
)
|
|
@click.option(
|
|
"--region",
|
|
default=None,
|
|
help="Cloud region for Vertex/Bedrock backends (env: HEADROOM_REGION).",
|
|
)
|
|
@click.option(
|
|
"--1m",
|
|
"context_1m",
|
|
is_flag=True,
|
|
help=(
|
|
"Preserve the 1M context window. Behind a custom ANTHROPIC_BASE_URL "
|
|
"Claude Code drops the context-1m beta header and caps at 200k; this "
|
|
"sets ANTHROPIC_MODEL=<opus>[1m] on the launched process so the 1M "
|
|
"window activates through the proxy (issue #1158)."
|
|
),
|
|
)
|
|
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
|
|
@click.option("--prepare-only", is_flag=True, hidden=True)
|
|
@click.argument("claude_args", nargs=-1, type=click.UNPROCESSED)
|
|
def claude(
|
|
port: int,
|
|
no_rtk: bool,
|
|
context_tool: bool,
|
|
no_mcp: bool,
|
|
no_tokensave: bool,
|
|
serena: bool,
|
|
no_serena: bool,
|
|
code_graph: bool,
|
|
no_proxy: bool,
|
|
learn: bool,
|
|
memory: bool,
|
|
tool_search: str | None,
|
|
backend: str | None,
|
|
region: str | None,
|
|
context_1m: bool,
|
|
verbose: bool,
|
|
prepare_only: bool,
|
|
claude_args: tuple,
|
|
) -> None:
|
|
"""Launch Claude Code through Headroom proxy.
|
|
|
|
\b
|
|
Sets ANTHROPIC_BASE_URL to route all Anthropic API calls through Headroom.
|
|
All unknown flags are passed through to claude (e.g. --resume, --model).
|
|
|
|
\b
|
|
Examples:
|
|
headroom wrap claude # Start everything
|
|
headroom wrap claude --memory # With persistent memory
|
|
headroom wrap claude --resume <id> # Resume a session
|
|
headroom wrap claude -- -p # Claude in print mode
|
|
headroom wrap claude # tokensave code graph (primary)
|
|
headroom wrap claude --no-tokensave # Skip tokensave; fall back to Serena
|
|
headroom wrap claude --serena # Also register the Serena backup
|
|
headroom wrap claude --context-tool # Enable CLI context-tool setup
|
|
headroom wrap claude --no-context-tool # Skip CLI context-tool setup
|
|
headroom wrap claude --no-mcp # Skip MCP retrieve tool registration
|
|
headroom wrap claude --no-serena # Never register the Serena backup
|
|
headroom wrap claude --1m # Preserve the 1M context window
|
|
"""
|
|
setup_context_tool = context_tool and not no_rtk
|
|
if prepare_only:
|
|
if setup_context_tool:
|
|
if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX:
|
|
_setup_lean_ctx_agent("claude", verbose=verbose)
|
|
else:
|
|
_prepare_wrap_rtk(verbose=verbose, label="Claude")
|
|
return
|
|
|
|
claude_bin = shutil.which("claude")
|
|
if not claude_bin:
|
|
click.echo("Error: 'claude' not found in PATH.")
|
|
click.echo("Install Claude Code: https://docs.anthropic.com/en/docs/claude-code")
|
|
raise SystemExit(1)
|
|
|
|
# Validate --tool-search up front so a typo fails before we start the proxy.
|
|
if tool_search is not None:
|
|
tool_search = _normalize_tool_search_mode(tool_search)
|
|
|
|
# Setup rtk before launching (Claude-specific)
|
|
proxy_holder: list[subprocess.Popen | None] = [None]
|
|
_saved_base_url: list[str | None] = [None] # previous settings.json value for restore
|
|
_settings_foundry: list[bool] = [False]
|
|
port_holder: list[int] = [port]
|
|
_settings_vertex: list[bool] = [False]
|
|
cleanup = _make_cleanup(proxy_holder, port_holder)
|
|
signal.signal(signal.SIGINT, _ignore_child_sigint)
|
|
signal.signal(signal.SIGTERM, cleanup)
|
|
if hasattr(signal, "SIGHUP"):
|
|
# Terminal close / tmux kill-session sends SIGHUP, not SIGTERM — without
|
|
# this, the finally block's base_url restore never runs (issue #1768).
|
|
signal.signal(signal.SIGHUP, cleanup)
|
|
|
|
# Memory sync BEFORE proxy startup — sync headroom DB ↔ Claude's files
|
|
if memory:
|
|
try:
|
|
mem_dir = Path.cwd() / ".headroom"
|
|
mem_dir.mkdir(parents=True, exist_ok=True)
|
|
_sync_db = str(mem_dir / "memory.db")
|
|
_sync_user = os.environ.get("USER", os.environ.get("USERNAME", "default"))
|
|
|
|
click.echo(f" Syncing memory (user={_sync_user})...")
|
|
sync_result = run(
|
|
[
|
|
sys.executable,
|
|
"-m",
|
|
"headroom.memory.sync",
|
|
"--db",
|
|
_sync_db,
|
|
"--user",
|
|
_sync_user,
|
|
"--agent",
|
|
"claude",
|
|
"--force",
|
|
],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=30,
|
|
)
|
|
if sync_result.returncode == 0 and sync_result.stdout.strip():
|
|
import json as _json
|
|
|
|
stats = _json.loads(sync_result.stdout.strip().split("\n")[-1])
|
|
imp, exp, ms = stats["imported"], stats["exported"], stats["ms"]
|
|
if imp or exp:
|
|
click.echo(f" Memory synced: {imp} imported, {exp} exported ({ms}ms)")
|
|
else:
|
|
click.echo(f" Memory: up to date ({ms}ms)")
|
|
elif sync_result.returncode != 0:
|
|
click.echo(f" Warning: memory sync error: {sync_result.stderr[-200:]}")
|
|
except Exception as e:
|
|
click.echo(f" Warning: memory sync failed: {e}")
|
|
|
|
try:
|
|
click.echo()
|
|
click.echo(" ╔═══════════════════════════════════════════════╗")
|
|
click.echo(" ║ HEADROOM WRAP: CLAUDE ║")
|
|
click.echo(" ╚═══════════════════════════════════════════════╝")
|
|
click.echo()
|
|
|
|
# Detect Foundry mode: Claude Code uses ANTHROPIC_FOUNDRY_BASE_URL instead of
|
|
# ANTHROPIC_BASE_URL when CLAUDE_CODE_USE_FOUNDRY=1 is set.
|
|
# Users typically set ANTHROPIC_FOUNDRY_RESOURCE (the resource name) rather
|
|
# than the full ANTHROPIC_FOUNDRY_BASE_URL. When the URL is absent we derive
|
|
# it from the resource name so the proxy has an upstream to forward to.
|
|
foundry_upstream = None
|
|
if os.environ.get("CLAUDE_CODE_USE_FOUNDRY"):
|
|
foundry_upstream = os.environ.get("ANTHROPIC_FOUNDRY_BASE_URL")
|
|
if not foundry_upstream:
|
|
resource = os.environ.get("ANTHROPIC_FOUNDRY_RESOURCE", "").strip()
|
|
if resource:
|
|
foundry_upstream = _foundry_upstream_url(resource)
|
|
|
|
# Detect Vertex mode: with CLAUDE_CODE_USE_VERTEX=1, Claude Code IGNORES
|
|
# ANTHROPIC_BASE_URL and authenticates to Google Vertex with GCP ADC. The
|
|
# documented way to route its Vertex :rawPredict / :streamRawPredict
|
|
# traffic through a gateway is ANTHROPIC_VERTEX_BASE_URL. Point it at
|
|
# Headroom and the proxy compresses the request, then forwards to the
|
|
# real regional Vertex host (derived per-request from the path's
|
|
# location) using Claude Code's own ADC token — no API key, no creds held
|
|
# by Headroom. This is the turnkey Vertex compression path.
|
|
use_vertex = bool(os.environ.get("CLAUDE_CODE_USE_VERTEX"))
|
|
proxy_url = _claude_proxy_base_url(port)
|
|
vertex_upstream = _vertex_target_api_url_from_claude_env(proxy_url) if use_vertex else None
|
|
|
|
_register_proxy_client(port)
|
|
proxy_holder[0], actual_port = _ensure_proxy(
|
|
port,
|
|
no_proxy,
|
|
learn=learn,
|
|
memory=memory,
|
|
agent_type="claude",
|
|
code_graph=code_graph,
|
|
backend=backend,
|
|
region=region,
|
|
anthropic_api_url=foundry_upstream,
|
|
vertex_api_url=vertex_upstream,
|
|
clear_vertex_api_url=use_vertex and vertex_upstream is None,
|
|
)
|
|
if actual_port != port:
|
|
_unregister_proxy_client(port)
|
|
_register_proxy_client(actual_port)
|
|
port_holder[0] = actual_port
|
|
_push_runtime_env(actual_port, no_proxy)
|
|
|
|
if setup_context_tool:
|
|
if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX:
|
|
click.echo(" Setting up lean-ctx...")
|
|
_setup_lean_ctx_agent("claude", verbose=verbose)
|
|
else:
|
|
click.echo(" Setting up rtk...")
|
|
_setup_rtk(verbose=verbose)
|
|
elif verbose:
|
|
click.echo(" Skipping CLI context tool (--no-context-tool)")
|
|
|
|
if not no_mcp:
|
|
from headroom.mcp_registry import ClaudeRegistrar
|
|
|
|
_setup_headroom_mcp(ClaudeRegistrar(), actual_port, verbose=verbose)
|
|
elif verbose:
|
|
click.echo(" Skipping MCP retrieve tool (--no-mcp)")
|
|
|
|
# Coding-task compressor: tokensave primary, Serena backup.
|
|
from headroom.mcp_registry import ClaudeRegistrar
|
|
|
|
_setup_coding_compressor(
|
|
ClaudeRegistrar(),
|
|
serena_context="claude-code",
|
|
serena=serena,
|
|
no_serena=no_serena,
|
|
no_tokensave=no_tokensave,
|
|
verbose=verbose,
|
|
)
|
|
|
|
if code_graph:
|
|
_setup_code_graph(verbose=verbose)
|
|
|
|
proxy_url = _claude_proxy_base_url(actual_port)
|
|
click.echo()
|
|
click.echo(" Launching Claude Code (API routed through Headroom)...")
|
|
if use_vertex:
|
|
click.echo(
|
|
f" Vertex mode: ANTHROPIC_VERTEX_BASE_URL={proxy_url} "
|
|
"→ compress, then forward to Vertex with your GCP ADC token"
|
|
)
|
|
elif foundry_upstream:
|
|
click.echo(
|
|
f" Foundry mode: ANTHROPIC_FOUNDRY_BASE_URL={_foundry_proxy_url(proxy_url)} → upstream {foundry_upstream}"
|
|
)
|
|
else:
|
|
click.echo(f" ANTHROPIC_BASE_URL={proxy_url}")
|
|
if is_custom_anthropic_base_url(proxy_url):
|
|
click.echo(
|
|
" "
|
|
+ remote_control_gate_message(
|
|
f"the wrapped Claude session's {REMOTE_CONTROL_BASE_URL_ENV}"
|
|
)
|
|
)
|
|
if claude_args:
|
|
click.echo(f" Extra args: {' '.join(claude_args)}")
|
|
_print_telemetry_notice()
|
|
click.echo()
|
|
|
|
env = os.environ.copy()
|
|
if use_vertex:
|
|
# Claude Code stays in Vertex mode (keeps CLAUDE_CODE_USE_VERTEX,
|
|
# ANTHROPIC_VERTEX_PROJECT_ID, CLOUD_ML_REGION, ADC — all inherited);
|
|
# we only redirect its Vertex endpoint to Headroom.
|
|
env["ANTHROPIC_VERTEX_BASE_URL"] = proxy_url
|
|
elif foundry_upstream:
|
|
# ANTHROPIC_FOUNDRY_BASE_URL is the base URL the Anthropic SDK
|
|
# appends /v1/messages to. The real Foundry URL includes /anthropic,
|
|
# so the proxy URL must mirror that structure.
|
|
env["ANTHROPIC_FOUNDRY_BASE_URL"] = _foundry_proxy_url(proxy_url)
|
|
else:
|
|
env["ANTHROPIC_BASE_URL"] = proxy_url
|
|
|
|
# Issue #951: write to settings.json so daemon-spawned conversation
|
|
# workers (which read settings.json fresh rather than inheriting the
|
|
# daemon's environment) also route through Headroom.
|
|
_settings_vertex[0] = bool(use_vertex)
|
|
_settings_foundry[0] = bool(foundry_upstream) and not _settings_vertex[0]
|
|
_wrap_settings_path = Path.cwd() / ".claude" / "settings.local.json"
|
|
_check_and_clear_stale_wrap_marker(
|
|
_wrap_settings_path,
|
|
key=_claude_wrap_base_url_env_key(
|
|
foundry_mode=_settings_foundry[0], vertex_mode=_settings_vertex[0]
|
|
),
|
|
)
|
|
_saved_base_url[0] = _write_claude_wrap_base_url(
|
|
(
|
|
_foundry_proxy_url(proxy_url)
|
|
if _settings_foundry[0]
|
|
else env["ANTHROPIC_VERTEX_BASE_URL"]
|
|
if _settings_vertex[0]
|
|
else proxy_url
|
|
),
|
|
foundry_mode=_settings_foundry[0],
|
|
vertex_mode=_settings_vertex[0],
|
|
settings_path=_wrap_settings_path,
|
|
port=port,
|
|
)
|
|
|
|
# Per-project savings attribution: tag every request with the launch
|
|
# directory's name via X-Headroom-Project (user override wins).
|
|
_apply_project_header_env(env)
|
|
|
|
# Issue #746: keep Claude Code's on-demand tool loading on through the
|
|
# proxy so tool schemas are not eagerly materialized into local context.
|
|
_tool_search_value = _configure_tool_search_env(env, tool_search)
|
|
if _tool_search_value is not None:
|
|
click.echo(
|
|
f" {_TOOL_SEARCH_ENV}={_tool_search_value} "
|
|
"(on-demand tool loading kept on; issue #746)"
|
|
)
|
|
elif verbose:
|
|
click.echo(
|
|
f" {_TOOL_SEARCH_ENV}={env.get(_TOOL_SEARCH_ENV)} "
|
|
"(using your existing environment value)"
|
|
)
|
|
|
|
# Issue #1158: opt-in 1M context window. Claude Code only sends the
|
|
# context-1m beta header when the model id carries the [1m] suffix, so
|
|
# force it via ANTHROPIC_MODEL on the launched process.
|
|
if context_1m:
|
|
env[_ANTHROPIC_MODEL_ENV] = _resolve_1m_model(env.get(_ANTHROPIC_MODEL_ENV))
|
|
click.echo(
|
|
f" {_ANTHROPIC_MODEL_ENV}={env[_ANTHROPIC_MODEL_ENV]} "
|
|
"(1M context window; issue #1158)"
|
|
)
|
|
|
|
result = subprocess.run([claude_bin, *claude_args], env=env)
|
|
raise SystemExit(result.returncode)
|
|
|
|
except SystemExit:
|
|
raise
|
|
except Exception as e:
|
|
click.echo(f" Error: {e}")
|
|
raise SystemExit(1) from e
|
|
finally:
|
|
_restore_claude_wrap_base_url(
|
|
_saved_base_url[0],
|
|
foundry_mode=_settings_foundry[0],
|
|
vertex_mode=_settings_vertex[0],
|
|
settings_path=_wrap_settings_path,
|
|
)
|
|
cleanup()
|
|
|
|
|
|
# =============================================================================
|
|
# Claude Code (unwrap)
|
|
# =============================================================================
|
|
|
|
|
|
@unwrap.command("claude")
|
|
@click.option(
|
|
"--port", "-p", default=8787, type=click.IntRange(1, 65535), help="Proxy port (default: 8787)"
|
|
)
|
|
@click.option("--no-stop-proxy", is_flag=True, help="Do not stop the local Headroom proxy")
|
|
@click.option("--keep-mcp", is_flag=True, help="Keep Headroom MCP registrations")
|
|
@click.option("--keep-rtk", is_flag=True, help="Keep rtk Claude hooks")
|
|
def unwrap_claude(
|
|
port: int,
|
|
no_stop_proxy: bool,
|
|
keep_mcp: bool,
|
|
keep_rtk: bool,
|
|
) -> None:
|
|
"""Undo durable setup from ``headroom wrap claude``."""
|
|
click.echo()
|
|
click.echo(" ╔═══════════════════════════════════════════════╗")
|
|
click.echo(" ║ HEADROOM UNWRAP: CLAUDE ║")
|
|
click.echo(" ╚═══════════════════════════════════════════════╝")
|
|
click.echo()
|
|
|
|
if not keep_mcp:
|
|
from headroom.mcp_registry import ClaudeRegistrar
|
|
|
|
registrar = ClaudeRegistrar()
|
|
if registrar.detect():
|
|
removed_headroom = registrar.unregister_server("headroom")
|
|
removed_code_graph = registrar.unregister_server(_CBM_MCP_SERVER_NAME)
|
|
tokensave_status = _remove_headroom_installed_tokensave_mcp(registrar)
|
|
serena_status = _remove_headroom_installed_serena_mcp(registrar)
|
|
if removed_headroom:
|
|
click.echo(" Removed Headroom MCP retrieve tool from Claude.")
|
|
else:
|
|
click.echo(" Headroom MCP retrieve tool was not registered in Claude.")
|
|
if removed_code_graph:
|
|
click.echo(" Removed legacy codebase-memory-mcp code graph server from Claude.")
|
|
if tokensave_status == "removed":
|
|
click.echo(" Removed Headroom-installed tokensave MCP server from Claude.")
|
|
elif tokensave_status == "failed":
|
|
click.echo(
|
|
" tokensave MCP server matched Headroom ledger but could not be removed."
|
|
)
|
|
if serena_status == "removed":
|
|
click.echo(" Removed Headroom-installed Serena MCP server from Claude.")
|
|
elif serena_status == "failed":
|
|
click.echo(" Serena MCP server matched Headroom ledger but could not be removed.")
|
|
else:
|
|
click.echo(" Claude Code not detected; skipped MCP cleanup.")
|
|
else:
|
|
click.echo(" Kept Claude MCP registrations (--keep-mcp).")
|
|
|
|
if not keep_rtk:
|
|
if _remove_claude_rtk_hooks():
|
|
click.echo(" Removed rtk Claude hook from settings.json.")
|
|
else:
|
|
click.echo(" No rtk Claude hook found in settings.json.")
|
|
else:
|
|
click.echo(" Kept rtk Claude hooks (--keep-rtk).")
|
|
|
|
_unwrap_settings_path = Path.cwd() / ".claude" / "settings.local.json"
|
|
for _foundry, _vertex in ((False, False), (True, False), (False, True)):
|
|
_key = _claude_wrap_base_url_env_key(foundry_mode=_foundry, vertex_mode=_vertex)
|
|
_marker = _read_wrap_marker(_unwrap_settings_path)
|
|
_prior = (
|
|
_marker.get("previous") if _marker is not None and _marker.get("key") == _key else None
|
|
)
|
|
_restore_claude_wrap_base_url(
|
|
_prior,
|
|
foundry_mode=_foundry,
|
|
vertex_mode=_vertex,
|
|
settings_path=_unwrap_settings_path,
|
|
)
|
|
|
|
click.echo()
|
|
click.echo("✓ Claude is no longer durably wrapped by Headroom.")
|
|
if not no_stop_proxy:
|
|
_echo_unwrap_proxy_stop_status(_stop_local_proxy_for_unwrap(port), port)
|
|
click.echo()
|
|
|
|
|
|
# =============================================================================
|
|
# GitHub Copilot CLI
|
|
# =============================================================================
|
|
|
|
|
|
@wrap.command(context_settings={"ignore_unknown_options": True})
|
|
@click.option(
|
|
"--port", "-p", default=8787, type=click.IntRange(1, 65535), help="Proxy port (default: 8787)"
|
|
)
|
|
@click.option(
|
|
"--no-context-tool",
|
|
"--no-rtk",
|
|
"no_rtk",
|
|
is_flag=True,
|
|
help="Skip CLI context-tool setup",
|
|
)
|
|
@click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)")
|
|
@click.option(
|
|
"--backend",
|
|
default=None,
|
|
help="API backend for the proxy: 'anthropic', 'anyllm', 'litellm-vertex', etc. (env: HEADROOM_BACKEND)",
|
|
)
|
|
@click.option(
|
|
"--anyllm-provider",
|
|
default=None,
|
|
help="Provider for any-llm backend: openai, mistral, groq, etc. (env: HEADROOM_ANYLLM_PROVIDER)",
|
|
)
|
|
@click.option(
|
|
"--region", default=None, help="Cloud region for Bedrock/Vertex (env: HEADROOM_REGION)"
|
|
)
|
|
@click.option(
|
|
"--provider-type",
|
|
type=click.Choice(["auto", "anthropic", "openai"]),
|
|
default="auto",
|
|
show_default=True,
|
|
help="Copilot BYOK provider mode. 'auto' uses anthropic for the default proxy backend and openai for translated backends.",
|
|
)
|
|
@click.option(
|
|
"--wire-api",
|
|
type=click.Choice(["completions", "responses"]),
|
|
default=None,
|
|
help="OpenAI-compatible Copilot wire API. Defaults to 'completions' when provider-type resolves to openai.",
|
|
)
|
|
@click.option(
|
|
"--subscription",
|
|
is_flag=True,
|
|
help=(
|
|
"Experimental: route GitHub-authenticated Copilot CLI traffic through Headroom "
|
|
"without requiring a provider API key."
|
|
),
|
|
)
|
|
@click.option("--memory", is_flag=True, help="Enable persistent cross-session memory")
|
|
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
|
|
@click.argument("copilot_args", nargs=-1, type=click.UNPROCESSED)
|
|
def copilot(
|
|
port: int,
|
|
no_rtk: bool,
|
|
no_proxy: bool,
|
|
backend: str | None,
|
|
anyllm_provider: str | None,
|
|
region: str | None,
|
|
provider_type: str,
|
|
wire_api: str | None,
|
|
subscription: bool,
|
|
memory: bool,
|
|
verbose: bool,
|
|
copilot_args: tuple[str, ...],
|
|
) -> None:
|
|
"""Launch GitHub Copilot CLI through Headroom proxy.
|
|
|
|
\b
|
|
Configures Copilot CLI BYOK provider variables so Copilot routes through
|
|
the local Headroom proxy. In auto mode, the wrapper uses Anthropic-style
|
|
routing for the stock proxy backend and OpenAI-compatible routing for
|
|
translated backends such as any-llm and LiteLLM.
|
|
|
|
\b
|
|
Examples:
|
|
headroom wrap copilot -- --model claude-sonnet-4-20250514
|
|
headroom wrap copilot --backend anyllm --anyllm-provider groq -- --model gpt-4o
|
|
headroom wrap copilot --provider-type openai --wire-api responses -- --model gpt-5.4
|
|
headroom wrap copilot --subscription -- --model gpt-4.1
|
|
headroom wrap copilot --no-context-tool -- --prompt "explain this file"
|
|
|
|
\b
|
|
Copilot hosted API (--subscription and the implicit OAuth path) routes to the
|
|
generic host https://api.githubcopilot.com, which serves the full model set.
|
|
Enterprise / data-residency accounts provisioned on a dedicated host pin it
|
|
explicitly with GITHUB_COPILOT_API_URL (the override flows through to upstream).
|
|
See TESTING-copilot-subscription.md for details.
|
|
"""
|
|
copilot_bin = shutil.which("copilot")
|
|
if not copilot_bin:
|
|
click.echo("Error: 'copilot' not found in PATH.")
|
|
click.echo(
|
|
"Install GitHub Copilot CLI: "
|
|
"https://docs.github.com/en/copilot/how-tos/copilot-cli/set-up-copilot-cli/install-copilot-cli"
|
|
)
|
|
raise SystemExit(1)
|
|
|
|
effective_backend = backend or os.environ.get("HEADROOM_BACKEND")
|
|
if _check_proxy(port):
|
|
running_backend = _detect_running_proxy_backend(port)
|
|
if effective_backend and running_backend and effective_backend != running_backend:
|
|
raise click.ClickException(
|
|
f"Proxy already running on port {port} with backend '{running_backend}'. "
|
|
f"Stop it or rerun with --backend {running_backend}."
|
|
)
|
|
effective_backend = running_backend or effective_backend
|
|
|
|
effective_provider_type = _resolve_copilot_provider_type(effective_backend, provider_type)
|
|
if subscription:
|
|
if effective_backend not in (None, "", "anthropic"):
|
|
raise click.ClickException(
|
|
"--subscription routes to GitHub Copilot's hosted API and cannot be combined "
|
|
"with translated backends such as anyllm or litellm-*."
|
|
)
|
|
if provider_type == "anthropic":
|
|
raise click.ClickException(
|
|
"--subscription uses Copilot's OpenAI-compatible hosted API path; "
|
|
"do not combine it with --provider-type anthropic."
|
|
)
|
|
effective_provider_type = "openai"
|
|
_validate_copilot_configuration(
|
|
provider_type=effective_provider_type,
|
|
wire_api=wire_api,
|
|
backend=effective_backend,
|
|
)
|
|
|
|
if not no_rtk:
|
|
if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX:
|
|
click.echo(" Setting up lean-ctx for Copilot...")
|
|
_setup_lean_ctx_agent("copilot", verbose=verbose)
|
|
else:
|
|
click.echo(" Setting up rtk for Copilot...")
|
|
rtk_path = _ensure_rtk_binary(verbose=verbose)
|
|
if rtk_path:
|
|
copilot_instructions = Path.cwd() / ".github" / "copilot-instructions.md"
|
|
_inject_rtk_instructions(copilot_instructions, verbose=verbose)
|
|
|
|
env = os.environ.copy()
|
|
openai_api_url: str | None = None
|
|
copilot_proxy_token: str | None = None
|
|
subscription_resolution = None
|
|
if _should_use_copilot_oauth(
|
|
backend=effective_backend,
|
|
provider_type=provider_type,
|
|
env=env,
|
|
force_subscription=subscription,
|
|
):
|
|
if subscription:
|
|
subscription_resolution = resolve_subscription_bearer_token_details()
|
|
client_bearer = (
|
|
subscription_resolution.token if subscription_resolution is not None else None
|
|
)
|
|
else:
|
|
client_bearer = resolve_client_bearer_token()
|
|
if not client_bearer:
|
|
raise click.ClickException(
|
|
"GitHub Copilot subscription mode requires a reusable GitHub/Copilot bearer "
|
|
"token, but none could be resolved. Run `headroom copilot-auth login` first, or set "
|
|
"GITHUB_COPILOT_TOKEN / GITHUB_COPILOT_GITHUB_TOKEN."
|
|
)
|
|
|
|
selected_model = _copilot_model_from_args(copilot_args, env)
|
|
|
|
# ``--model auto`` is a Copilot-internal routing token that the BYOK
|
|
# API rejects with ``400 The requested model is not supported``. In
|
|
# subscription/OAuth mode we route to the real Copilot hosted API, so
|
|
# Copilot's own native auto-selection works fine — we just need to
|
|
# strip the ``--model auto`` flag before launch so Copilot doesn't
|
|
# forward it to the provider endpoint.
|
|
if _is_auto_model(selected_model):
|
|
copilot_args = _strip_auto_model_args(copilot_args)
|
|
selected_model = None
|
|
click.echo(
|
|
" Note: '--model auto' is not forwarded to the Copilot API "
|
|
"(it would cause a 400). Removed it; Copilot will use its own "
|
|
"automatic model selection."
|
|
)
|
|
|
|
effective_wire_api = wire_api or (
|
|
_copilot_default_wire_api_for_model(selected_model) if subscription else "completions"
|
|
)
|
|
env["COPILOT_PROVIDER_TYPE"] = "openai"
|
|
# Per-project savings: the Copilot CLI cannot send custom headers, so
|
|
# the project rides as a /p/<name> base-URL prefix the proxy strips.
|
|
env["COPILOT_PROVIDER_BASE_URL"] = _with_project_prefix(
|
|
f"http://127.0.0.1:{port}/v1", _project_name_from_cwd()
|
|
)
|
|
env["COPILOT_PROVIDER_WIRE_API"] = effective_wire_api
|
|
env["COPILOT_PROVIDER_BEARER_TOKEN"] = client_bearer
|
|
env["GITHUB_COPILOT_USE_TOKEN_EXCHANGE"] = "false"
|
|
env.pop("COPILOT_PROVIDER_API_KEY", None)
|
|
# Hand the exact token we resolved (and, for --subscription, validated
|
|
# against GitHub) to the proxy explicitly via copilot_proxy_token below.
|
|
# The proxy pins it as GITHUB_COPILOT_API_TOKEN, so upstream auth is
|
|
# deterministic instead of the proxy re-running unvalidated discovery
|
|
# (read_cached_oauth_token returns the *first* candidate, which may not
|
|
# be the one the wrapper approved → environment-dependent 401s). Passing
|
|
# it as a launch argument — rather than mutating this process's global
|
|
# os.environ — keeps the token off shared state and out of unrelated
|
|
# code paths.
|
|
copilot_proxy_token = client_bearer
|
|
env_vars_display = [
|
|
"COPILOT_PROVIDER_TYPE=openai",
|
|
f"COPILOT_PROVIDER_BASE_URL={env['COPILOT_PROVIDER_BASE_URL']}",
|
|
f"COPILOT_PROVIDER_WIRE_API={effective_wire_api}",
|
|
(
|
|
"COPILOT_AUTH_MODE=github-subscription-experimental"
|
|
if subscription
|
|
else "COPILOT_AUTH_MODE=github-oauth"
|
|
),
|
|
]
|
|
# Non-subscription OAuth keeps upstream's generic-host policy from
|
|
# #610. Subscription mode can use the endpoint returned by the Copilot
|
|
# token exchange, which is how Business accounts advertise their API
|
|
# host without requiring users to configure it manually.
|
|
openai_api_url = (
|
|
subscription_resolution.api_url
|
|
if subscription_resolution is not None
|
|
else resolve_copilot_api_url(client_bearer)
|
|
)
|
|
env["GITHUB_COPILOT_API_URL"] = openai_api_url
|
|
env["OPENAI_TARGET_API_URL"] = openai_api_url
|
|
env_vars_display.append(f"COPILOT_PROVIDER_API_URL={openai_api_url}")
|
|
else:
|
|
env, env_vars_display = _build_copilot_launch_env(
|
|
port=port,
|
|
provider_type=effective_provider_type,
|
|
wire_api=wire_api,
|
|
environ=env,
|
|
project=_project_name_from_cwd(),
|
|
)
|
|
|
|
if not env.get("COPILOT_PROVIDER_API_KEY"):
|
|
src = _copilot_provider_key_source(effective_provider_type)
|
|
click.echo(
|
|
f"\n Error: Copilot BYOK mode requires a provider API key.\n"
|
|
f" `headroom wrap copilot` uses Copilot's BYOK mode, which bypasses GitHub's\n"
|
|
f" Copilot API and routes requests directly to the model provider through the\n"
|
|
f" Headroom proxy. A GitHub Copilot subscription alone is not sufficient.\n\n"
|
|
f" Set one of:\n"
|
|
f" export {src}=sk-... # recommended\n"
|
|
f" export COPILOT_PROVIDER_API_KEY=sk-... # also works\n"
|
|
)
|
|
raise SystemExit(1)
|
|
|
|
if not subscription and not _copilot_model_configured(copilot_args, env):
|
|
# Distinguish between "--model auto" (wrong model for BYOK) and
|
|
# genuinely missing model (no --model flag at all).
|
|
raw_model = _copilot_model_from_args(copilot_args, env)
|
|
if _is_auto_model(raw_model):
|
|
click.echo(
|
|
" Error: '--model auto' is not supported in Copilot BYOK mode.\n"
|
|
" BYOK routes to an external provider (Anthropic/OpenAI) which\n"
|
|
" does not recognise 'auto' as a model name — the request will\n"
|
|
" fail with a 400 error.\n"
|
|
" Options:\n"
|
|
" • Use a concrete model: --model gpt-4o\n"
|
|
" • Use subscription mode for native auto-routing:\n"
|
|
" headroom wrap copilot --subscription -- --model auto"
|
|
)
|
|
raise SystemExit(1)
|
|
else:
|
|
click.echo(
|
|
" Note: Copilot BYOK requires a model. Pass `--model <name>` "
|
|
"or set `COPILOT_MODEL` / `COPILOT_PROVIDER_MODEL_ID`."
|
|
)
|
|
|
|
_launch_tool(
|
|
binary=copilot_bin,
|
|
args=copilot_args,
|
|
env=env,
|
|
port=port,
|
|
no_proxy=no_proxy,
|
|
tool_label="COPILOT",
|
|
env_vars_display=env_vars_display,
|
|
learn=False,
|
|
memory=memory,
|
|
agent_type="copilot",
|
|
backend=backend,
|
|
anyllm_provider=anyllm_provider,
|
|
region=region,
|
|
openai_api_url=openai_api_url,
|
|
copilot_api_token=copilot_proxy_token,
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# GitHub Copilot CLI (unwrap)
|
|
# =============================================================================
|
|
|
|
|
|
@unwrap.command("copilot")
|
|
@click.option(
|
|
"--port", "-p", default=8787, type=click.IntRange(1, 65535), help="Proxy port (default: 8787)"
|
|
)
|
|
@click.option("--no-stop-proxy", is_flag=True, help="Do not stop the local Headroom proxy")
|
|
def unwrap_copilot(port: int, no_stop_proxy: bool) -> None:
|
|
"""Undo durable setup from ``headroom wrap copilot``."""
|
|
instructions = Path.cwd() / ".github" / "copilot-instructions.md"
|
|
if _remove_rtk_instructions(instructions):
|
|
click.echo(" Removed Headroom rtk instructions from Copilot.")
|
|
else:
|
|
click.echo(" No Headroom rtk instructions found for Copilot.")
|
|
|
|
if not no_stop_proxy:
|
|
_echo_unwrap_proxy_stop_status(_stop_local_proxy_for_unwrap(port), port)
|
|
|
|
|
|
# =============================================================================
|
|
# OpenAI Codex CLI
|
|
# =============================================================================
|
|
|
|
|
|
def _prepare_codex_wrap_state(
|
|
*,
|
|
port: int,
|
|
no_rtk: bool,
|
|
no_mcp: bool,
|
|
no_tokensave: bool,
|
|
serena: bool,
|
|
no_serena: bool,
|
|
memory: bool,
|
|
verbose: bool,
|
|
rtk_home: Path | None = None,
|
|
) -> None:
|
|
"""Prepare the active Codex home for a wrap or prepare-only invocation."""
|
|
# Snapshot Codex config.toml BEFORE any wrap-time mutation so
|
|
# `headroom unwrap codex` can restore the user's pre-wrap state
|
|
# byte-for-byte. The snapshot is a no-op if the backup already exists
|
|
# or if the file already has Headroom markers, so this is safe to
|
|
# call repeatedly. Crucially this must run before MCP install, which
|
|
# writes its marker block to the same file.
|
|
_codex_config_file, _codex_backup_file = _codex_config_paths()
|
|
_snapshot_codex_config_if_unwrapped(_codex_config_file, _codex_backup_file)
|
|
|
|
# Setup CLI context tool for Codex.
|
|
if not no_rtk:
|
|
if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX:
|
|
click.echo(" Setting up lean-ctx for Codex...")
|
|
_setup_lean_ctx_agent("codex", verbose=verbose)
|
|
else:
|
|
click.echo(" Setting up rtk for Codex...")
|
|
rtk_path = _ensure_rtk_binary(verbose=verbose)
|
|
if rtk_path:
|
|
# Keep RTK guidance local to the user's Codex configuration.
|
|
global_agents = (rtk_home or _codex_home_dir()) / "AGENTS.md"
|
|
_inject_rtk_instructions(global_agents, verbose=verbose)
|
|
|
|
# Register headroom MCP server in Codex config.toml so Codex can
|
|
# call headroom_retrieve on compression markers from the proxy.
|
|
if not no_mcp:
|
|
from headroom.mcp_registry import CodexRegistrar
|
|
|
|
# Codex starts a long-lived local MCP subprocess from config.toml.
|
|
# If a previous wrap used another port, retrieval can silently point
|
|
# at the wrong proxy while model traffic uses the right one.
|
|
_setup_headroom_mcp(CodexRegistrar(), port, verbose=verbose, force=True)
|
|
elif verbose:
|
|
click.echo(" Skipping MCP retrieve tool (--no-mcp)")
|
|
|
|
# Coding-task compressor: tokensave primary, Serena backup. Codex starts
|
|
# long-lived MCP subprocesses from config.toml, so force re-registration.
|
|
from headroom.mcp_registry import CodexRegistrar
|
|
|
|
_setup_coding_compressor(
|
|
CodexRegistrar(),
|
|
serena_context="codex",
|
|
serena=serena,
|
|
no_serena=no_serena,
|
|
no_tokensave=no_tokensave,
|
|
verbose=verbose,
|
|
force=True,
|
|
)
|
|
|
|
# Setup memory MCP server for Codex (native tool integration)
|
|
if memory:
|
|
click.echo(" Setting up memory for Codex...")
|
|
mem_dir = Path.cwd() / ".headroom"
|
|
mem_dir.mkdir(parents=True, exist_ok=True)
|
|
db_path = str(mem_dir / "memory.db")
|
|
mem_user = os.environ.get("USER", os.environ.get("USERNAME", "default"))
|
|
|
|
# Register MCP server in Codex config
|
|
_inject_memory_mcp_config(mem_user)
|
|
|
|
# Inject memory guidance into project AGENTS.md
|
|
agents_md = Path.cwd() / "AGENTS.md"
|
|
_inject_memory_agents_md(agents_md)
|
|
|
|
# Sync Claude's memories → DB so MCP search finds them
|
|
try:
|
|
import asyncio
|
|
|
|
from headroom.memory.sync import _build_sync_backend, sync_import
|
|
from headroom.memory.sync_adapters.claude_code import (
|
|
ClaudeCodeAdapter,
|
|
get_claude_memory_dir,
|
|
)
|
|
|
|
claude_memory_dir = get_claude_memory_dir()
|
|
|
|
async def _import_claude_memories() -> int:
|
|
backend = _build_sync_backend(db_path)
|
|
await backend._ensure_initialized()
|
|
adapter = ClaudeCodeAdapter(claude_memory_dir)
|
|
count = await sync_import(backend, adapter, mem_user)
|
|
await backend.close()
|
|
return count
|
|
|
|
imported = asyncio.run(_import_claude_memories())
|
|
if imported:
|
|
click.echo(f" Memory: imported {imported} memories from Claude")
|
|
except Exception as e:
|
|
click.echo(f" Warning: Claude memory import failed: {e}")
|
|
|
|
# Inject Headroom provider into Codex config so WebSocket traffic also
|
|
# routes through the proxy. Codex ignores OPENAI_BASE_URL for its WS
|
|
# transport unless a custom provider declares supports_websockets = true.
|
|
# NOTE: this must run BEFORE _inject_memory_mcp_config because it rewrites
|
|
# the config file. Re-inject MCP config after if memory is enabled.
|
|
_inject_codex_provider_config(port)
|
|
if memory:
|
|
_inject_memory_mcp_config(os.environ.get("USER", os.environ.get("USERNAME", "default")))
|
|
|
|
|
|
def _run_codex_wrap(
|
|
*,
|
|
port: int,
|
|
no_rtk: bool,
|
|
no_mcp: bool,
|
|
no_tokensave: bool,
|
|
serena: bool,
|
|
no_serena: bool,
|
|
code_graph: bool,
|
|
no_proxy: bool,
|
|
learn: bool,
|
|
memory: bool,
|
|
backend: str | None,
|
|
anyllm_provider: str | None,
|
|
region: str | None,
|
|
verbose: bool,
|
|
prepare_only: bool,
|
|
codex_args: tuple,
|
|
) -> None:
|
|
"""Execute the Codex wrap flow with the session overlay when launching."""
|
|
if prepare_only:
|
|
_prepare_codex_wrap_state(
|
|
port=port,
|
|
no_rtk=no_rtk,
|
|
no_mcp=no_mcp,
|
|
no_tokensave=no_tokensave,
|
|
serena=serena,
|
|
no_serena=no_serena,
|
|
memory=memory,
|
|
verbose=verbose,
|
|
)
|
|
return
|
|
|
|
codex_bin = shutil.which("codex")
|
|
if not codex_bin:
|
|
click.echo("Error: 'codex' not found in PATH.")
|
|
click.echo("Install Codex CLI: npm install -g @openai/codex")
|
|
raise SystemExit(1)
|
|
|
|
active_codex_home = _codex_home_dir()
|
|
with _codex_session_home_overlay() as session_codex_home:
|
|
_prepare_codex_wrap_state(
|
|
port=port,
|
|
no_rtk=no_rtk,
|
|
no_mcp=no_mcp,
|
|
no_tokensave=no_tokensave,
|
|
serena=serena,
|
|
no_serena=no_serena,
|
|
memory=memory,
|
|
verbose=verbose,
|
|
rtk_home=active_codex_home,
|
|
)
|
|
|
|
env, env_vars_display = _build_codex_launch_env(port, os.environ)
|
|
|
|
# Per-project savings attribution: the injected provider config maps the
|
|
# X-Headroom-Project header to HEADROOM_PROJECT via env_http_headers, so
|
|
# Codex sends it only when this var is set. A user-set value wins.
|
|
_codex_project = _project_name_from_cwd()
|
|
if _codex_project and "HEADROOM_PROJECT" not in env:
|
|
env["HEADROOM_PROJECT"] = _codex_project
|
|
|
|
env["CODEX_HOME"] = str(session_codex_home)
|
|
|
|
_launch_tool(
|
|
binary=codex_bin,
|
|
args=codex_args,
|
|
env=env,
|
|
port=port,
|
|
no_proxy=no_proxy,
|
|
tool_label="CODEX",
|
|
env_vars_display=env_vars_display,
|
|
learn=learn,
|
|
memory=memory,
|
|
agent_type="codex",
|
|
code_graph=code_graph,
|
|
backend=backend,
|
|
anyllm_provider=anyllm_provider,
|
|
region=region,
|
|
)
|
|
|
|
|
|
@wrap.command(context_settings={"ignore_unknown_options": True})
|
|
@click.option(
|
|
"--port", "-p", default=8787, type=click.IntRange(1, 65535), help="Proxy port (default: 8787)"
|
|
)
|
|
@click.option(
|
|
"--no-context-tool",
|
|
"--no-rtk",
|
|
"no_rtk",
|
|
is_flag=True,
|
|
help="Skip CLI context-tool setup",
|
|
)
|
|
@click.option(
|
|
"--no-mcp",
|
|
is_flag=True,
|
|
help="Skip headroom MCP server registration (compression markers will be unactionable)",
|
|
)
|
|
@click.option(
|
|
"--no-tokensave",
|
|
is_flag=True,
|
|
help="Skip the tokensave code-graph MCP server (primary coding-task compressor)",
|
|
)
|
|
@click.option(
|
|
"--serena",
|
|
is_flag=True,
|
|
help="Force the Serena MCP backup compressor on (registered automatically when "
|
|
"tokensave is unavailable)",
|
|
)
|
|
@click.option("--no-serena", is_flag=True, help="Never register the Serena backup compressor")
|
|
@click.option(
|
|
"--code-graph",
|
|
is_flag=True,
|
|
help="Force a tokensave code-graph index now (tokensave is the default compressor)",
|
|
)
|
|
@click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)")
|
|
@click.option(
|
|
"--learn", is_flag=True, help="Enable live traffic learning (patterns saved to AGENTS.md)"
|
|
)
|
|
@click.option(
|
|
"--backend",
|
|
default=None,
|
|
help="API backend for the proxy: 'anthropic', 'anyllm', 'litellm-vertex', etc. (env: HEADROOM_BACKEND)",
|
|
)
|
|
@click.option(
|
|
"--anyllm-provider",
|
|
default=None,
|
|
help="Provider for any-llm backend: openai, mistral, groq, etc. (env: HEADROOM_ANYLLM_PROVIDER)",
|
|
)
|
|
@click.option(
|
|
"--region", default=None, help="Cloud region for Bedrock/Vertex (env: HEADROOM_REGION)"
|
|
)
|
|
@click.option("--memory", is_flag=True, help="Enable persistent cross-session memory")
|
|
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
|
|
@click.option("--prepare-only", is_flag=True, hidden=True)
|
|
@click.argument("codex_args", nargs=-1, type=click.UNPROCESSED)
|
|
def codex(
|
|
port: int,
|
|
no_rtk: bool,
|
|
no_mcp: bool,
|
|
no_tokensave: bool,
|
|
serena: bool,
|
|
no_serena: bool,
|
|
code_graph: bool,
|
|
no_proxy: bool,
|
|
learn: bool,
|
|
memory: bool,
|
|
backend: str | None,
|
|
anyllm_provider: str | None,
|
|
region: str | None,
|
|
verbose: bool,
|
|
prepare_only: bool,
|
|
codex_args: tuple,
|
|
) -> None:
|
|
"""Launch OpenAI Codex CLI through Headroom proxy.
|
|
|
|
\b
|
|
Sets OPENAI_BASE_URL to route all OpenAI API calls through Headroom.
|
|
Sets up the selected CLI context tool so Codex uses token-optimized
|
|
commands (60-90% savings on shell output). Also
|
|
registers the headroom MCP server in the active Codex config file
|
|
so Codex can call ``headroom_retrieve`` on compression markers.
|
|
|
|
\b
|
|
Examples:
|
|
headroom wrap codex # Start proxy + context tool + mcp + codex
|
|
headroom wrap codex -- "fix the bug" # Pass prompt to codex
|
|
headroom wrap codex --no-context-tool # Skip CLI context-tool setup
|
|
headroom wrap codex --no-mcp # Skip MCP retrieve tool registration
|
|
headroom wrap codex --no-tokensave # Skip tokensave; fall back to Serena
|
|
headroom wrap codex --serena # Also register the Serena backup
|
|
headroom wrap codex --no-serena # Never register the Serena backup
|
|
headroom wrap codex --port 9999 # Custom proxy port
|
|
headroom wrap codex --backend anyllm --anyllm-provider groq
|
|
"""
|
|
return _run_codex_wrap(
|
|
port=port,
|
|
no_rtk=no_rtk,
|
|
no_mcp=no_mcp,
|
|
no_tokensave=no_tokensave,
|
|
serena=serena,
|
|
no_serena=no_serena,
|
|
code_graph=code_graph,
|
|
no_proxy=no_proxy,
|
|
learn=learn,
|
|
memory=memory,
|
|
backend=backend,
|
|
anyllm_provider=anyllm_provider,
|
|
region=region,
|
|
verbose=verbose,
|
|
prepare_only=prepare_only,
|
|
codex_args=codex_args,
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# Aider
|
|
# =============================================================================
|
|
|
|
|
|
@wrap.command(context_settings={"ignore_unknown_options": True})
|
|
@click.option(
|
|
"--port", "-p", default=8787, type=click.IntRange(1, 65535), help="Proxy port (default: 8787)"
|
|
)
|
|
@click.option(
|
|
"--no-context-tool",
|
|
"--no-rtk",
|
|
"no_rtk",
|
|
is_flag=True,
|
|
help="Skip CLI context-tool setup",
|
|
)
|
|
@click.option(
|
|
"--code-graph",
|
|
is_flag=True,
|
|
help="Enable code graph indexing via codebase-memory-mcp (optional)",
|
|
)
|
|
@click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)")
|
|
@click.option("--learn", is_flag=True, help="Enable live traffic learning")
|
|
@click.option("--memory", is_flag=True, help="Enable persistent cross-session memory")
|
|
@click.option(
|
|
"--backend", default=None, help="API backend: 'anthropic', 'anyllm', 'litellm-vertex', etc."
|
|
)
|
|
@click.option("--anyllm-provider", default=None, help="Provider for any-llm backend")
|
|
@click.option("--region", default=None, help="Cloud region for Bedrock/Vertex")
|
|
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
|
|
@click.option("--prepare-only", is_flag=True, hidden=True)
|
|
@click.argument("aider_args", nargs=-1, type=click.UNPROCESSED)
|
|
def aider(
|
|
port: int,
|
|
no_rtk: bool,
|
|
code_graph: bool,
|
|
no_proxy: bool,
|
|
learn: bool,
|
|
memory: bool,
|
|
backend: str | None,
|
|
anyllm_provider: str | None,
|
|
region: str | None,
|
|
verbose: bool,
|
|
prepare_only: bool,
|
|
aider_args: tuple,
|
|
) -> None:
|
|
"""Launch aider through Headroom proxy.
|
|
|
|
\b
|
|
Sets OPENAI_API_BASE to route all API calls through Headroom.
|
|
Sets up the selected CLI context tool so aider uses token-optimized commands.
|
|
|
|
\b
|
|
Examples:
|
|
headroom wrap aider # Start proxy + context tool + aider
|
|
headroom wrap aider -- --model gpt-4o # Use GPT-4o
|
|
headroom wrap aider -- --model claude-sonnet-4 # Use Claude
|
|
headroom wrap aider --no-context-tool # Skip CLI context-tool setup
|
|
headroom wrap aider --backend litellm-vertex --region us-central1
|
|
"""
|
|
# Setup CLI context tool for aider.
|
|
if not no_rtk:
|
|
if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX:
|
|
click.echo(" Setting up lean-ctx for aider...")
|
|
_setup_lean_ctx_agent("aider", verbose=verbose)
|
|
else:
|
|
click.echo(" Setting up rtk for aider...")
|
|
rtk_path = _ensure_rtk_binary(verbose=verbose)
|
|
if rtk_path:
|
|
# aider reads CONVENTIONS.md from project root
|
|
conventions = Path.cwd() / "CONVENTIONS.md"
|
|
_inject_rtk_instructions(conventions, verbose=verbose)
|
|
|
|
if prepare_only:
|
|
return
|
|
|
|
aider_bin = shutil.which("aider")
|
|
if not aider_bin:
|
|
click.echo("Error: 'aider' not found in PATH.")
|
|
click.echo("Install aider: pip install aider-chat")
|
|
raise SystemExit(1)
|
|
|
|
env, env_vars_display = _build_aider_launch_env(
|
|
port, os.environ, project=_project_name_from_cwd()
|
|
)
|
|
|
|
_launch_tool(
|
|
binary=aider_bin,
|
|
args=aider_args,
|
|
env=env,
|
|
port=port,
|
|
no_proxy=no_proxy,
|
|
tool_label="AIDER",
|
|
env_vars_display=env_vars_display,
|
|
learn=learn,
|
|
memory=memory,
|
|
agent_type="aider",
|
|
code_graph=code_graph,
|
|
backend=backend,
|
|
anyllm_provider=anyllm_provider,
|
|
region=region,
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# Mistral Vibe
|
|
# =============================================================================
|
|
|
|
|
|
@wrap.command(context_settings={"ignore_unknown_options": True})
|
|
@click.option(
|
|
"--port", "-p", default=8787, type=click.IntRange(1, 65535), help="Proxy port (default: 8787)"
|
|
)
|
|
@click.option(
|
|
"--no-context-tool",
|
|
"--no-rtk",
|
|
"no_rtk",
|
|
is_flag=True,
|
|
help="Skip CLI context-tool setup (no effect for vibe)",
|
|
)
|
|
@click.option(
|
|
"--code-graph",
|
|
is_flag=True,
|
|
help="Enable code graph indexing via codebase-memory-mcp (optional)",
|
|
)
|
|
@click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)")
|
|
@click.option("--learn", is_flag=True, help="Enable live traffic learning")
|
|
@click.option("--memory", is_flag=True, help="Enable persistent cross-session memory")
|
|
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
|
|
@click.option("--prepare-only", is_flag=True, hidden=True)
|
|
@click.argument("vibe_args", nargs=-1, type=click.UNPROCESSED)
|
|
def vibe(
|
|
port: int,
|
|
no_rtk: bool,
|
|
code_graph: bool,
|
|
no_proxy: bool,
|
|
learn: bool,
|
|
memory: bool,
|
|
verbose: bool,
|
|
prepare_only: bool,
|
|
vibe_args: tuple,
|
|
) -> None:
|
|
"""Launch Mistral Vibe through Headroom proxy.
|
|
|
|
\b
|
|
Sets VIBE_PROVIDERS to route all Mistral API calls through Headroom.
|
|
|
|
\b
|
|
Examples:
|
|
headroom wrap vibe # Start proxy + vibe
|
|
headroom wrap vibe -- "fix the bug" # Pass prompt to vibe
|
|
headroom wrap vibe --port 9999 # Custom proxy port
|
|
headroom wrap vibe --no-context-tool # Skip CLI context-tool setup
|
|
"""
|
|
if prepare_only:
|
|
return
|
|
|
|
vibe_bin = shutil.which("vibe")
|
|
if not vibe_bin:
|
|
click.echo("Error: 'vibe' not found in PATH.")
|
|
click.echo("Install Mistral Vibe: https://github.com/mistralai/mistral-vibe")
|
|
raise SystemExit(1)
|
|
|
|
env, env_vars_display = _build_mistral_vibe_launch_env(
|
|
port, os.environ, project=_project_name_from_cwd()
|
|
)
|
|
|
|
_launch_tool(
|
|
binary=vibe_bin,
|
|
args=vibe_args,
|
|
env=env,
|
|
port=port,
|
|
no_proxy=no_proxy,
|
|
tool_label="VIBE",
|
|
env_vars_display=env_vars_display,
|
|
learn=learn,
|
|
memory=memory,
|
|
agent_type="vibe",
|
|
code_graph=code_graph,
|
|
openai_api_url="https://api.mistral.ai",
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# Cursor
|
|
# =============================================================================
|
|
|
|
|
|
@wrap.command(context_settings={"ignore_unknown_options": True})
|
|
@click.option(
|
|
"--port", "-p", default=8787, type=click.IntRange(1, 65535), help="Proxy port (default: 8787)"
|
|
)
|
|
@click.option(
|
|
"--no-context-tool",
|
|
"--no-rtk",
|
|
"no_rtk",
|
|
is_flag=True,
|
|
help="Skip CLI context-tool setup",
|
|
)
|
|
@click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)")
|
|
@click.option(
|
|
"--learn", is_flag=True, help="Enable live traffic learning (patterns saved to .cursor/rules/)"
|
|
)
|
|
@click.option("--memory", is_flag=True, help="Enable persistent cross-session memory")
|
|
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
|
|
@click.option("--prepare-only", is_flag=True, hidden=True)
|
|
def cursor(
|
|
port: int,
|
|
no_rtk: bool,
|
|
no_proxy: bool,
|
|
learn: bool,
|
|
memory: bool,
|
|
verbose: bool,
|
|
prepare_only: bool,
|
|
) -> None:
|
|
"""Start Headroom proxy for use with Cursor.
|
|
|
|
\b
|
|
Cursor reads its API configuration from its settings UI, not from
|
|
environment variables. This command starts the proxy, sets up the selected
|
|
CLI context tool, and prints the Cursor settings.
|
|
|
|
\b
|
|
After running this command, open Cursor and configure:
|
|
Settings > Models > OpenAI API Key > Advanced > Override Base URL
|
|
|
|
\b
|
|
Example:
|
|
headroom wrap cursor # Start proxy + context-tool instructions
|
|
headroom wrap cursor --no-context-tool # Proxy only, no CLI context tool
|
|
headroom wrap cursor --port 9999 # Custom proxy port
|
|
"""
|
|
cursorrules: Path | None = Path.cwd() / ".cursorrules" if not no_rtk else None
|
|
cursor_hook_registered = False
|
|
if not no_rtk:
|
|
|
|
def _register_cursor_hook(rtk_path: Path) -> None:
|
|
# rtk registers a native hook for Cursor (`rtk init --agent cursor`),
|
|
# same mechanism as Claude Code. Prefer that over injecting the
|
|
# RTK_INSTRUCTIONS_BLOCK text into .cursorrules — a silent hook makes
|
|
# the custom-rules text redundant guidance (GH #756).
|
|
nonlocal cursor_hook_registered
|
|
from headroom.rtk.installer import register_agent_hooks
|
|
|
|
# rtk may exit 0 without writing hooks.json (e.g. an rtk build that
|
|
# doesn't support --agent cursor), so trust the file, not the exit
|
|
# code: only skip the .cursorrules fallback if the native hook is
|
|
# actually on disk (GH #756).
|
|
cursor_hooks_json = Path.home() / ".cursor" / "hooks.json"
|
|
if register_agent_hooks(rtk_path, agent="cursor") and cursor_hooks_json.is_file():
|
|
cursor_hook_registered = True
|
|
else:
|
|
_inject_rtk_instructions(cast(Path, cursorrules), verbose=verbose)
|
|
|
|
_setup_context_tool_for_agent(
|
|
agent="cursor",
|
|
agent_display="Cursor",
|
|
marker_path=cursorrules,
|
|
on_rtk_ready=_register_cursor_hook,
|
|
verbose=verbose,
|
|
)
|
|
|
|
if prepare_only:
|
|
return
|
|
|
|
def _print_cursor_setup(actual_port: int) -> None:
|
|
for line in _render_cursor_setup_lines(actual_port, project=_project_name_from_cwd()):
|
|
click.echo(line)
|
|
if not no_rtk:
|
|
click.echo()
|
|
if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX:
|
|
click.echo(" lean-ctx configured for Cursor")
|
|
elif cursor_hook_registered:
|
|
click.echo(" rtk hook registered for Cursor")
|
|
else:
|
|
click.echo(" rtk instructions injected into .cursorrules")
|
|
click.echo(" Cursor will use token-optimized commands automatically.")
|
|
|
|
_run_proxy_only_watcher(
|
|
agent_label="cursor",
|
|
port=port,
|
|
no_proxy=no_proxy,
|
|
learn=learn,
|
|
memory=memory,
|
|
agent_type="cursor",
|
|
print_setup_lines=_print_cursor_setup,
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# Cline (VS Code extension)
|
|
# =============================================================================
|
|
|
|
|
|
@wrap.command(context_settings={"ignore_unknown_options": True})
|
|
@click.option(
|
|
"--port", "-p", default=8787, type=click.IntRange(1, 65535), help="Proxy port (default: 8787)"
|
|
)
|
|
@click.option(
|
|
"--no-context-tool",
|
|
"--no-rtk",
|
|
"no_rtk",
|
|
is_flag=True,
|
|
help="Skip CLI context-tool setup",
|
|
)
|
|
@click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)")
|
|
@click.option("--learn", is_flag=True, help="Enable live traffic learning")
|
|
@click.option("--memory", is_flag=True, help="Enable persistent cross-session memory")
|
|
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
|
|
@click.option("--prepare-only", is_flag=True, hidden=True)
|
|
def cline(
|
|
port: int,
|
|
no_rtk: bool,
|
|
no_proxy: bool,
|
|
learn: bool,
|
|
memory: bool,
|
|
verbose: bool,
|
|
prepare_only: bool,
|
|
) -> None:
|
|
"""Start Headroom proxy for use with Cline (VS Code extension).
|
|
|
|
\b
|
|
Cline is a VS Code extension that reads its API configuration from the
|
|
VS Code settings UI, not from environment variables. This command starts
|
|
the proxy, sets up the selected CLI context tool (injecting RTK guidance
|
|
into .clinerules at the project root), and prints the Cline settings the
|
|
user should configure.
|
|
|
|
\b
|
|
After running this command, open Cline's settings in VS Code and configure
|
|
the API Base URL to point at the local Headroom proxy.
|
|
|
|
\b
|
|
Uninstall: there is no ``headroom unwrap cline`` subcommand. To remove the
|
|
injected guidance, hand-edit ``.clinerules`` at the project root and
|
|
delete everything between ``<!-- headroom:rtk-instructions -->`` and
|
|
``<!-- /headroom:rtk-instructions -->`` (inclusive). If ``lean-ctx`` mode
|
|
is selected, the lean-ctx agent name ``cline`` may not be recognized by
|
|
the local lean-ctx binary; a warning is printed in that case and setup
|
|
is skipped silently.
|
|
|
|
\b
|
|
Examples:
|
|
headroom wrap cline # Start proxy + .clinerules instructions
|
|
headroom wrap cline --no-context-tool # Proxy only, no CLI context tool
|
|
headroom wrap cline --port 9999 # Custom proxy port
|
|
"""
|
|
# Pre-compute the marker path so the KeyboardInterrupt handler can report
|
|
# its location even if the interrupt fires before _inject_rtk_instructions
|
|
# returns (e.g., during the inner _ensure_rtk_binary download).
|
|
clinerules: Path | None = Path.cwd() / ".clinerules" if not no_rtk else None
|
|
if not no_rtk:
|
|
_setup_context_tool_for_agent(
|
|
agent="cline",
|
|
agent_display="Cline",
|
|
marker_path=clinerules,
|
|
on_rtk_ready=lambda _rtk: _inject_rtk_instructions(
|
|
cast(Path, clinerules), verbose=verbose
|
|
),
|
|
verbose=verbose,
|
|
)
|
|
|
|
if prepare_only:
|
|
return
|
|
|
|
def _print_cline_setup(actual_port: int) -> None:
|
|
anthropic_base = _claude_proxy_base_url(actual_port)
|
|
openai_base = f"http://127.0.0.1:{actual_port}/v1"
|
|
click.echo(" Configure Cline in VS Code:")
|
|
click.echo(" Settings > Cline > API Provider")
|
|
click.echo(f" Anthropic Base URL: {anthropic_base}")
|
|
click.echo(f" OpenAI Compatible Base URL: {openai_base}")
|
|
if not no_rtk:
|
|
click.echo()
|
|
if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX:
|
|
click.echo(" lean-ctx configured for Cline")
|
|
else:
|
|
click.echo(" rtk instructions injected into .clinerules")
|
|
click.echo(" Cline will use token-optimized commands automatically.")
|
|
|
|
_run_proxy_only_watcher(
|
|
agent_label="cline",
|
|
port=port,
|
|
no_proxy=no_proxy,
|
|
learn=learn,
|
|
memory=memory,
|
|
agent_type="cline",
|
|
print_setup_lines=_print_cline_setup,
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# Continue (VS Code / JetBrains extension)
|
|
# =============================================================================
|
|
|
|
|
|
@wrap.command("continue", context_settings={"ignore_unknown_options": True})
|
|
@click.option(
|
|
"--port", "-p", default=8787, type=click.IntRange(1, 65535), help="Proxy port (default: 8787)"
|
|
)
|
|
@click.option(
|
|
"--no-context-tool",
|
|
"--no-rtk",
|
|
"no_rtk",
|
|
is_flag=True,
|
|
help="Skip CLI context-tool setup",
|
|
)
|
|
@click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)")
|
|
@click.option("--learn", is_flag=True, help="Enable live traffic learning")
|
|
@click.option("--memory", is_flag=True, help="Enable persistent cross-session memory")
|
|
@click.option(
|
|
"--config",
|
|
"config_path",
|
|
type=click.Path(path_type=Path, file_okay=True, dir_okay=False),
|
|
default=None,
|
|
help="Path to Continue config.json (default: ./.continue/config.json)",
|
|
)
|
|
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
|
|
@click.option("--prepare-only", is_flag=True, hidden=True)
|
|
def continue_dev(
|
|
port: int,
|
|
no_rtk: bool,
|
|
no_proxy: bool,
|
|
learn: bool,
|
|
memory: bool,
|
|
config_path: Path | None,
|
|
verbose: bool,
|
|
prepare_only: bool,
|
|
) -> None:
|
|
"""Start Headroom proxy for use with Continue (VS Code / JetBrains).
|
|
|
|
\b
|
|
Continue reads its model configuration from .continue/config.json (a JSON
|
|
document with a top-level ``systemMessage`` and a ``models`` array). This
|
|
command starts the proxy, sets up the selected CLI context tool by
|
|
extending ``systemMessage`` with RTK guidance, and prints the per-model
|
|
``apiBase`` the user should configure manually.
|
|
|
|
\b
|
|
Continue is an IDE extension — its API base URL is configured per-model
|
|
in config.json (or via the IDE UI), not via environment variables. The
|
|
config file is overridable via --config.
|
|
|
|
\b
|
|
Note: Continue's modern config is YAML-first (``.continue/config.yaml``).
|
|
This helper only writes the JSON variant. Users on the YAML schema should
|
|
configure ``systemMessage`` through that file by hand.
|
|
|
|
\b
|
|
Per-model handling: Continue overrides top-level ``systemMessage`` with
|
|
per-model ``systemMessage`` when set, so this command also injects into
|
|
each ``models[i].systemMessage`` if the ``models`` array is present.
|
|
Existing non-string ``systemMessage`` values are NEVER overwritten — the
|
|
command warns loudly and leaves them in place. To opt in, clear the
|
|
existing value first.
|
|
|
|
\b
|
|
Uninstall: there is no ``headroom unwrap continue`` subcommand. To remove
|
|
the injected guidance, hand-edit ``.continue/config.json`` and delete
|
|
everything between ``<!-- headroom:rtk-instructions -->`` and
|
|
``<!-- /headroom:rtk-instructions -->`` (inclusive) from every
|
|
``systemMessage`` field — both top-level and inside ``models[*]``. If
|
|
``lean-ctx`` mode is selected, the lean-ctx agent name ``continue`` may
|
|
not be recognized by the local lean-ctx binary; a warning is printed in
|
|
that case and setup is skipped silently.
|
|
|
|
\b
|
|
Examples:
|
|
headroom wrap continue # Start proxy + inject systemMessage
|
|
headroom wrap continue --no-context-tool # Proxy only
|
|
headroom wrap continue --port 9999 # Custom proxy port
|
|
headroom wrap continue --config path/to/config.json
|
|
"""
|
|
config_file = config_path or (Path.cwd() / ".continue" / "config.json")
|
|
|
|
if not no_rtk:
|
|
_setup_context_tool_for_agent(
|
|
agent="continue",
|
|
agent_display="Continue",
|
|
marker_path=config_file,
|
|
on_rtk_ready=lambda _rtk: _inject_continue_rtk_systemmessage(
|
|
config_file, verbose=verbose
|
|
),
|
|
verbose=verbose,
|
|
)
|
|
|
|
if prepare_only:
|
|
return
|
|
|
|
def _print_continue_setup(actual_port: int) -> None:
|
|
anthropic_base = _claude_proxy_base_url(actual_port)
|
|
openai_base = f"http://127.0.0.1:{actual_port}/v1"
|
|
click.echo(" Configure Continue in your IDE:")
|
|
click.echo(f" Edit {config_file} and set, per model:")
|
|
click.echo(f' "apiBase": "{openai_base}" # OpenAI-compatible models')
|
|
click.echo(f' "apiBase": "{anthropic_base}" # Anthropic models')
|
|
if not no_rtk:
|
|
click.echo()
|
|
if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX:
|
|
click.echo(" lean-ctx configured for Continue")
|
|
else:
|
|
click.echo(f" rtk instructions injected into {config_file.name} systemMessage")
|
|
click.echo(" Continue will use token-optimized commands automatically.")
|
|
|
|
_run_proxy_only_watcher(
|
|
agent_label="continue",
|
|
port=port,
|
|
no_proxy=no_proxy,
|
|
learn=learn,
|
|
memory=memory,
|
|
agent_type="continue",
|
|
print_setup_lines=_print_continue_setup,
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# Goose (Block)
|
|
# =============================================================================
|
|
|
|
|
|
@wrap.command(context_settings={"ignore_unknown_options": True})
|
|
@click.option(
|
|
"--port", "-p", default=8787, type=click.IntRange(1, 65535), help="Proxy port (default: 8787)"
|
|
)
|
|
@click.option(
|
|
"--no-context-tool",
|
|
"--no-rtk",
|
|
"no_rtk",
|
|
is_flag=True,
|
|
help="Skip CLI context-tool setup",
|
|
)
|
|
@click.option(
|
|
"--code-graph",
|
|
is_flag=True,
|
|
help="Enable code graph indexing via codebase-memory-mcp (optional)",
|
|
)
|
|
@click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)")
|
|
@click.option("--learn", is_flag=True, help="Enable live traffic learning")
|
|
@click.option("--memory", is_flag=True, help="Enable persistent cross-session memory")
|
|
@click.option(
|
|
"--backend", default=None, help="API backend: 'anthropic', 'anyllm', 'litellm-vertex', etc."
|
|
)
|
|
@click.option("--anyllm-provider", default=None, help="Provider for any-llm backend")
|
|
@click.option("--region", default=None, help="Cloud region for Bedrock/Vertex")
|
|
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
|
|
@click.option("--prepare-only", is_flag=True, hidden=True)
|
|
@click.argument("goose_args", nargs=-1, type=click.UNPROCESSED)
|
|
def goose(
|
|
port: int,
|
|
no_rtk: bool,
|
|
code_graph: bool,
|
|
no_proxy: bool,
|
|
learn: bool,
|
|
memory: bool,
|
|
backend: str | None,
|
|
anyllm_provider: str | None,
|
|
region: str | None,
|
|
verbose: bool,
|
|
prepare_only: bool,
|
|
goose_args: tuple,
|
|
) -> None:
|
|
"""Launch Goose (Block) CLI through Headroom proxy.
|
|
|
|
\b
|
|
Sets OPENAI_BASE_URL and ANTHROPIC_BASE_URL to route Goose's API calls
|
|
through Headroom. Sets up the selected CLI context tool by injecting RTK
|
|
guidance into .goosehints at the project root (Goose reads this file as
|
|
extra system context).
|
|
|
|
\b
|
|
Uninstall: there is no ``headroom unwrap goose`` subcommand. To remove the
|
|
injected guidance, hand-edit ``.goosehints`` at the project root and
|
|
delete everything between ``<!-- headroom:rtk-instructions -->`` and
|
|
``<!-- /headroom:rtk-instructions -->`` (inclusive). If ``lean-ctx`` mode
|
|
is selected, the lean-ctx agent name ``goose`` may not be recognized by
|
|
the local lean-ctx binary; a warning is printed in that case and setup
|
|
is skipped silently.
|
|
|
|
\b
|
|
Examples:
|
|
headroom wrap goose # Start proxy + context tool + goose
|
|
headroom wrap goose -- session # Start a Goose session
|
|
headroom wrap goose -- --provider anthropic # Pass args to goose
|
|
headroom wrap goose --no-context-tool # Skip CLI context-tool setup
|
|
"""
|
|
# Goose reads .goosehints from the project root as extra context.
|
|
# Pre-compute the marker path so the KeyboardInterrupt handler can report
|
|
# its location even if the interrupt fires before _inject_rtk_instructions
|
|
# returns (e.g., during the inner _ensure_rtk_binary download).
|
|
goosehints: Path | None = Path.cwd() / ".goosehints" if not no_rtk else None
|
|
if not no_rtk:
|
|
_setup_context_tool_for_agent(
|
|
agent="goose",
|
|
agent_display="Goose",
|
|
marker_path=goosehints,
|
|
on_rtk_ready=lambda _rtk: _inject_rtk_instructions(
|
|
cast(Path, goosehints), verbose=verbose
|
|
),
|
|
verbose=verbose,
|
|
)
|
|
|
|
if prepare_only:
|
|
return
|
|
|
|
goose_bin = shutil.which("goose")
|
|
if not goose_bin:
|
|
click.echo("Error: 'goose' not found in PATH.")
|
|
click.echo("Install Goose: https://block.github.io/goose/")
|
|
raise SystemExit(1)
|
|
|
|
# Goose accepts OpenAI- and Anthropic-compatible providers; route both.
|
|
env = os.environ.copy()
|
|
openai_base = f"http://127.0.0.1:{port}/v1"
|
|
anthropic_base = _claude_proxy_base_url(port)
|
|
env["OPENAI_BASE_URL"] = openai_base
|
|
env["OPENAI_API_BASE"] = openai_base
|
|
env["ANTHROPIC_BASE_URL"] = anthropic_base
|
|
env_vars_display = [
|
|
f"OPENAI_BASE_URL={openai_base}",
|
|
f"ANTHROPIC_BASE_URL={anthropic_base}",
|
|
]
|
|
|
|
_launch_tool(
|
|
binary=goose_bin,
|
|
args=goose_args,
|
|
env=env,
|
|
port=port,
|
|
no_proxy=no_proxy,
|
|
tool_label="GOOSE",
|
|
env_vars_display=env_vars_display,
|
|
learn=learn,
|
|
memory=memory,
|
|
agent_type="goose",
|
|
code_graph=code_graph,
|
|
backend=backend,
|
|
anyllm_provider=anyllm_provider,
|
|
region=region,
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# OpenHands
|
|
# =============================================================================
|
|
|
|
|
|
@wrap.command(context_settings={"ignore_unknown_options": True})
|
|
@click.option(
|
|
"--port", "-p", default=8787, type=click.IntRange(1, 65535), help="Proxy port (default: 8787)"
|
|
)
|
|
@click.option(
|
|
"--no-context-tool",
|
|
"--no-rtk",
|
|
"no_rtk",
|
|
is_flag=True,
|
|
help="Skip CLI context-tool setup",
|
|
)
|
|
@click.option(
|
|
"--code-graph",
|
|
is_flag=True,
|
|
help="Enable code graph indexing via codebase-memory-mcp (optional)",
|
|
)
|
|
@click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)")
|
|
@click.option("--learn", is_flag=True, help="Enable live traffic learning")
|
|
@click.option("--memory", is_flag=True, help="Enable persistent cross-session memory")
|
|
@click.option(
|
|
"--backend", default=None, help="API backend: 'anthropic', 'anyllm', 'litellm-vertex', etc."
|
|
)
|
|
@click.option("--anyllm-provider", default=None, help="Provider for any-llm backend")
|
|
@click.option("--region", default=None, help="Cloud region for Bedrock/Vertex")
|
|
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
|
|
@click.option("--prepare-only", is_flag=True, hidden=True)
|
|
@click.argument("openhands_args", nargs=-1, type=click.UNPROCESSED)
|
|
def openhands(
|
|
port: int,
|
|
no_rtk: bool,
|
|
code_graph: bool,
|
|
no_proxy: bool,
|
|
learn: bool,
|
|
memory: bool,
|
|
backend: str | None,
|
|
anyllm_provider: str | None,
|
|
region: str | None,
|
|
verbose: bool,
|
|
prepare_only: bool,
|
|
openhands_args: tuple,
|
|
) -> None:
|
|
"""Launch OpenHands CLI through Headroom proxy.
|
|
|
|
\b
|
|
Sets OPENAI_BASE_URL / ANTHROPIC_BASE_URL to route OpenHands' API calls
|
|
through Headroom. Instructions are injected via the
|
|
``OPENHANDS_INSTRUCTIONS`` environment variable at launch time so the
|
|
on-disk OpenHands config is left untouched.
|
|
|
|
\b
|
|
The ``OPENHANDS_INSTRUCTIONS`` value injected by this command contains the
|
|
``<!-- headroom:rtk-instructions -->`` marker. To uninstall, simply do not
|
|
set ``OPENHANDS_INSTRUCTIONS`` in the parent shell — this command never
|
|
writes to disk, so nothing to clean up. If ``lean-ctx`` mode is selected,
|
|
the lean-ctx agent name ``openhands`` may not be recognized by the local
|
|
lean-ctx binary; a warning is printed in that case and rtk-style guidance
|
|
falls through.
|
|
|
|
\b
|
|
Examples:
|
|
headroom wrap openhands # Start proxy + context tool + openhands
|
|
headroom wrap openhands -- --task ... # Pass args to openhands
|
|
headroom wrap openhands --no-context-tool
|
|
"""
|
|
# openhands never writes to disk — its rtk guidance ships via the
|
|
# OPENHANDS_INSTRUCTIONS env var below — so marker_path is None and
|
|
# rtk_required gates the env-only path: without an rtk binary there
|
|
# is no fallback marker file to fall through to.
|
|
rtk_path: Path | None = None
|
|
if not no_rtk:
|
|
rtk_path = _setup_context_tool_for_agent(
|
|
agent="openhands",
|
|
agent_display="OpenHands",
|
|
marker_path=None,
|
|
rtk_required=True,
|
|
verbose=verbose,
|
|
)
|
|
|
|
if prepare_only:
|
|
return
|
|
|
|
openhands_bin = shutil.which("openhands")
|
|
if not openhands_bin:
|
|
click.echo("Error: 'openhands' not found in PATH.")
|
|
click.echo("Install OpenHands: https://docs.all-hands.dev/")
|
|
raise SystemExit(1)
|
|
|
|
env = os.environ.copy()
|
|
openai_base = f"http://127.0.0.1:{port}/v1"
|
|
anthropic_base = _claude_proxy_base_url(port)
|
|
env["OPENAI_BASE_URL"] = openai_base
|
|
env["OPENAI_API_BASE"] = openai_base
|
|
env["ANTHROPIC_BASE_URL"] = anthropic_base
|
|
# Also set LLM_BASE_URL for OpenHands' generic LLM provider config.
|
|
env["LLM_BASE_URL"] = openai_base
|
|
if not no_rtk and rtk_path:
|
|
# Inject rtk guidance via env var so OpenHands picks it up as the
|
|
# session's instruction prefix. Appending instead of overwriting
|
|
# any pre-existing OPENHANDS_INSTRUCTIONS so user-supplied content
|
|
# is preserved. The marker check guards against double-injection
|
|
# when the user inherits an env var that already has the rtk block.
|
|
existing_instructions = env.get("OPENHANDS_INSTRUCTIONS", "")
|
|
if _RTK_MARKER in existing_instructions:
|
|
# Already injected — pre-existing env var contains marker.
|
|
pass
|
|
elif existing_instructions.strip():
|
|
env["OPENHANDS_INSTRUCTIONS"] = (
|
|
existing_instructions.rstrip() + "\n\n" + RTK_INSTRUCTIONS_BLOCK
|
|
)
|
|
else:
|
|
env["OPENHANDS_INSTRUCTIONS"] = RTK_INSTRUCTIONS_BLOCK
|
|
|
|
env_vars_display = [
|
|
f"OPENAI_BASE_URL={openai_base}",
|
|
f"ANTHROPIC_BASE_URL={anthropic_base}",
|
|
f"LLM_BASE_URL={openai_base}",
|
|
]
|
|
if not no_rtk and "OPENHANDS_INSTRUCTIONS" in env:
|
|
env_vars_display.append("OPENHANDS_INSTRUCTIONS=<rtk instructions injected>")
|
|
|
|
_launch_tool(
|
|
binary=openhands_bin,
|
|
args=openhands_args,
|
|
env=env,
|
|
port=port,
|
|
no_proxy=no_proxy,
|
|
tool_label="OPENHANDS",
|
|
env_vars_display=env_vars_display,
|
|
learn=learn,
|
|
memory=memory,
|
|
agent_type="openhands",
|
|
code_graph=code_graph,
|
|
backend=backend,
|
|
anyllm_provider=anyllm_provider,
|
|
region=region,
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# OpenClaw
|
|
# =============================================================================
|
|
|
|
|
|
@wrap.command("openclaw")
|
|
@click.option(
|
|
"--plugin-path",
|
|
type=click.Path(path_type=Path, file_okay=False, dir_okay=True),
|
|
default=None,
|
|
help="Path to local OpenClaw plugin source directory (advanced/dev override)",
|
|
)
|
|
@click.option(
|
|
"--plugin-spec",
|
|
default="headroom-ai/openclaw",
|
|
show_default=True,
|
|
help="NPM plugin spec for OpenClaw install (used when --plugin-path is omitted)",
|
|
)
|
|
@click.option(
|
|
"--skip-build",
|
|
is_flag=True,
|
|
help="Skip npm install/build in local source mode (--plugin-path)",
|
|
)
|
|
@click.option(
|
|
"--copy",
|
|
is_flag=True,
|
|
help="Install by copying plugin path instead of using --link",
|
|
)
|
|
@click.option(
|
|
"--proxy-port", default=8787, type=click.IntRange(1, 65535), help="Headroom proxy port"
|
|
)
|
|
@click.option("--startup-timeout-ms", default=20000, type=int, help="Proxy startup timeout")
|
|
@click.option(
|
|
"--gateway-provider-id",
|
|
"gateway_provider_ids",
|
|
multiple=True,
|
|
help="OpenClaw provider id to route through Headroom (repeatable; default: openai-codex)",
|
|
)
|
|
@click.option(
|
|
"--python-path",
|
|
default=None,
|
|
help="Optional Python executable for proxy launcher fallback",
|
|
)
|
|
@click.option(
|
|
"--no-auto-start",
|
|
is_flag=True,
|
|
help="Disable plugin auto-start of local headroom proxy",
|
|
)
|
|
@click.option(
|
|
"--no-restart",
|
|
is_flag=True,
|
|
help="Do not restart OpenClaw gateway at the end",
|
|
)
|
|
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
|
|
@click.option("--prepare-only", is_flag=True, hidden=True)
|
|
@click.option("--existing-entry-json", default=None, hidden=True)
|
|
def openclaw(
|
|
plugin_path: Path | None,
|
|
plugin_spec: str,
|
|
skip_build: bool,
|
|
copy: bool,
|
|
proxy_port: int,
|
|
startup_timeout_ms: int,
|
|
gateway_provider_ids: tuple[str, ...],
|
|
python_path: str | None,
|
|
no_auto_start: bool,
|
|
no_restart: bool,
|
|
verbose: bool,
|
|
prepare_only: bool,
|
|
existing_entry_json: str | None,
|
|
) -> None:
|
|
"""Install and configure Headroom OpenClaw plugin in one command.
|
|
|
|
\b
|
|
What this command does:
|
|
1. Installs OpenClaw plugin from npm (or local --plugin-path)
|
|
2. Builds plugin source if --plugin-path is used
|
|
3. Writes minimal plugin config and sets contextEngine slot
|
|
4. Validates config
|
|
5. Restarts OpenClaw gateway (unless --no-restart)
|
|
|
|
\b
|
|
Example:
|
|
headroom wrap openclaw
|
|
headroom wrap openclaw --plugin-path C:\\git\\headroom\\plugins\\openclaw
|
|
"""
|
|
if prepare_only:
|
|
entry = _build_openclaw_plugin_entry(
|
|
existing_entry=_decode_openclaw_entry_json(existing_entry_json),
|
|
proxy_port=proxy_port,
|
|
startup_timeout_ms=startup_timeout_ms,
|
|
python_path=python_path,
|
|
no_auto_start=no_auto_start,
|
|
gateway_provider_ids=gateway_provider_ids,
|
|
enabled=True,
|
|
)
|
|
click.echo(json.dumps(entry, separators=(",", ":")))
|
|
return
|
|
|
|
openclaw_bin = shutil.which("openclaw")
|
|
if not openclaw_bin:
|
|
raise click.ClickException("'openclaw' not found in PATH. Install OpenClaw CLI first.")
|
|
|
|
plugin_dir = plugin_path.resolve() if plugin_path else None
|
|
local_source_mode = plugin_dir is not None
|
|
if plugin_dir:
|
|
if not plugin_dir.exists():
|
|
raise click.ClickException(f"Plugin path not found: {plugin_dir}.")
|
|
if not (plugin_dir / "package.json").exists():
|
|
raise click.ClickException(f"Invalid plugin path (missing package.json): {plugin_dir}")
|
|
if not (plugin_dir / "openclaw.plugin.json").exists():
|
|
raise click.ClickException(
|
|
f"Invalid plugin path (missing openclaw.plugin.json): {plugin_dir}"
|
|
)
|
|
|
|
npm_bin = shutil.which("npm")
|
|
if local_source_mode and not skip_build and not npm_bin:
|
|
raise click.ClickException(
|
|
"'npm' not found in PATH. Install Node/npm or rerun with --skip-build."
|
|
)
|
|
|
|
click.echo()
|
|
click.echo(" ╔═══════════════════════════════════════════════╗")
|
|
click.echo(" ║ HEADROOM WRAP: OPENCLAW ║")
|
|
click.echo(" ╚═══════════════════════════════════════════════╝")
|
|
click.echo()
|
|
if local_source_mode:
|
|
click.echo(f" Plugin source: local ({plugin_dir})")
|
|
else:
|
|
click.echo(f" Plugin source: npm ({plugin_spec})")
|
|
|
|
if local_source_mode and not skip_build:
|
|
click.echo(" Building OpenClaw plugin (npm install + npm run build)...")
|
|
_run_checked([npm_bin or "npm", "install"], cwd=plugin_dir, action="npm install")
|
|
_run_checked([npm_bin or "npm", "run", "build"], cwd=plugin_dir, action="npm run build")
|
|
elif not local_source_mode and skip_build:
|
|
click.echo(" Skipping build: npm install mode does not build local source.")
|
|
|
|
effective_python_path = python_path
|
|
if effective_python_path is None and not no_auto_start and sys.executable:
|
|
effective_python_path = sys.executable
|
|
|
|
existing_entry = _read_openclaw_config_value(openclaw_bin, "plugins.entries.headroom")
|
|
entry = _build_openclaw_plugin_entry(
|
|
existing_entry=existing_entry,
|
|
proxy_port=proxy_port,
|
|
startup_timeout_ms=startup_timeout_ms,
|
|
python_path=effective_python_path,
|
|
no_auto_start=no_auto_start,
|
|
gateway_provider_ids=gateway_provider_ids,
|
|
enabled=True,
|
|
)
|
|
|
|
click.echo(" Writing plugin configuration...")
|
|
_write_openclaw_plugin_entry(openclaw_bin, entry)
|
|
|
|
install_cmd = [
|
|
openclaw_bin,
|
|
"plugins",
|
|
"install",
|
|
"--dangerously-force-unsafe-install",
|
|
]
|
|
if local_source_mode:
|
|
if copy:
|
|
install_cmd.append(str(plugin_dir))
|
|
install_cwd = None
|
|
else:
|
|
install_cmd.extend(["--link", "."])
|
|
install_cwd = plugin_dir
|
|
else:
|
|
install_cmd.append(plugin_spec)
|
|
install_cwd = None
|
|
|
|
click.echo(" Installing OpenClaw plugin with required unsafe-install flag...")
|
|
install_result = run(
|
|
install_cmd,
|
|
cwd=str(install_cwd) if install_cwd else None,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
if install_result.returncode != 0:
|
|
combined_error = "\n".join(
|
|
x for x in [install_result.stderr.strip(), install_result.stdout.strip()] if x
|
|
)
|
|
plugin_already_exists = "plugin already exists" in combined_error.lower()
|
|
linked_install_bug = (
|
|
"also not a valid hook pack" in combined_error.lower()
|
|
and "--dangerously-force-unsafe-install" in " ".join(install_cmd)
|
|
)
|
|
if plugin_already_exists:
|
|
click.echo(" Plugin already installed; continuing with configuration/update steps.")
|
|
elif linked_install_bug and local_source_mode and plugin_dir is not None:
|
|
click.echo(
|
|
" OpenClaw linked-path install bug detected; applying extension-path fallback..."
|
|
)
|
|
target_dir = _copy_openclaw_plugin_into_extensions(
|
|
plugin_dir=plugin_dir,
|
|
openclaw_bin=openclaw_bin,
|
|
)
|
|
click.echo(f" Fallback plugin copy completed: {target_dir}")
|
|
else:
|
|
details = combined_error or f"exit code {install_result.returncode}"
|
|
raise click.ClickException(f"openclaw plugins install failed: {details}")
|
|
elif verbose and install_result.stdout.strip():
|
|
click.echo(install_result.stdout.strip())
|
|
|
|
_set_openclaw_context_engine_slot(openclaw_bin, "headroom")
|
|
_run_checked(
|
|
[openclaw_bin, "config", "validate"],
|
|
action="openclaw config validate",
|
|
)
|
|
|
|
if no_restart:
|
|
click.echo(" Skipping gateway restart (--no-restart).")
|
|
click.echo(
|
|
" Run `openclaw gateway restart` (or `openclaw gateway start`) to apply plugin changes."
|
|
)
|
|
else:
|
|
click.echo(" Applying plugin changes to OpenClaw gateway...")
|
|
gateway_action, gateway_output = _restart_or_start_openclaw_gateway(openclaw_bin)
|
|
click.echo(f" Gateway {gateway_action}.")
|
|
if verbose and gateway_output:
|
|
click.echo(gateway_output)
|
|
|
|
inspect_result = _run_checked(
|
|
[openclaw_bin, "plugins", "inspect", "headroom"],
|
|
action="openclaw plugins inspect headroom",
|
|
)
|
|
if verbose and inspect_result.stdout.strip():
|
|
click.echo(inspect_result.stdout.strip())
|
|
|
|
click.echo()
|
|
click.echo("✓ OpenClaw is configured to use Headroom context compression.")
|
|
click.echo(" Plugin: headroom")
|
|
click.echo(" Slot: plugins.slots.contextEngine = headroom")
|
|
click.echo()
|
|
|
|
|
|
# =============================================================================
|
|
# OpenCode
|
|
# =============================================================================
|
|
|
|
|
|
@wrap.command(context_settings={"ignore_unknown_options": True})
|
|
@click.option(
|
|
"--port", "-p", default=8787, type=click.IntRange(1, 65535), help="Proxy port (default: 8787)"
|
|
)
|
|
@click.option(
|
|
"--no-context-tool",
|
|
"--no-rtk",
|
|
"no_rtk",
|
|
is_flag=True,
|
|
help="Skip CLI context-tool setup",
|
|
)
|
|
@click.option("--no-mcp", is_flag=True, help="Skip headroom MCP server registration")
|
|
@click.option("--no-serena", is_flag=True, help="Skip Serena MCP server registration")
|
|
@click.option(
|
|
"--code-graph",
|
|
is_flag=True,
|
|
help="Enable code graph indexing via codebase-memory-mcp (optional)",
|
|
)
|
|
@click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)")
|
|
@click.option("--learn", is_flag=True, help="Enable live traffic learning")
|
|
@click.option("--memory", is_flag=True, help="Enable persistent cross-session memory")
|
|
@click.option(
|
|
"--backend", default=None, help="API backend: 'anthropic', 'anyllm', 'litellm-vertex', etc."
|
|
)
|
|
@click.option("--anyllm-provider", default=None, help="Provider for any-llm backend")
|
|
@click.option("--region", default=None, help="Cloud region for Bedrock/Vertex")
|
|
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
|
|
@click.option("--prepare-only", is_flag=True, hidden=True)
|
|
@click.argument("opencode_args", nargs=-1, type=click.UNPROCESSED)
|
|
def opencode(
|
|
port: int,
|
|
no_rtk: bool,
|
|
no_mcp: bool,
|
|
no_serena: bool,
|
|
code_graph: bool,
|
|
no_proxy: bool,
|
|
learn: bool,
|
|
memory: bool,
|
|
backend: str | None,
|
|
anyllm_provider: str | None,
|
|
region: str | None,
|
|
verbose: bool,
|
|
prepare_only: bool,
|
|
opencode_args: tuple,
|
|
) -> None:
|
|
"""Launch OpenCode through Headroom proxy.
|
|
|
|
\b
|
|
Sets OPENCODE_CONFIG_CONTENT to route all OpenCode API calls through
|
|
Headroom. Configures a headroom provider via @ai-sdk/openai-compatible.
|
|
Also sets OPENAI_BASE_URL and ANTHROPIC_BASE_URL as fallbacks.
|
|
|
|
\b
|
|
Examples:
|
|
headroom wrap opencode # Start proxy + context tool + opencode
|
|
headroom wrap opencode -- "fix the bug" # Pass prompt to opencode
|
|
headroom wrap opencode --no-context-tool # Skip CLI context-tool setup
|
|
headroom wrap opencode --no-mcp # Skip MCP retrieve tool registration
|
|
headroom wrap opencode --no-serena # Skip Serena MCP registration
|
|
headroom wrap opencode --port 9999 # Custom proxy port
|
|
headroom wrap opencode --backend anyllm --anyllm-provider groq
|
|
"""
|
|
# Snapshot OpenCode config.json BEFORE any wrap-time mutation so
|
|
# `headroom unwrap opencode` can restore the user's pre-wrap state.
|
|
_opencode_config_file, _opencode_backup_file = opencode_config_paths()
|
|
snapshot_opencode_config_if_unwrapped(_opencode_config_file, _opencode_backup_file)
|
|
|
|
# Setup CLI context tool for OpenCode.
|
|
if not no_rtk:
|
|
if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX:
|
|
click.echo(" Setting up lean-ctx for OpenCode...")
|
|
_setup_lean_ctx_agent("opencode", verbose=verbose)
|
|
else:
|
|
click.echo(" Setting up rtk for OpenCode...")
|
|
rtk_path = _ensure_rtk_binary(verbose=verbose)
|
|
if rtk_path:
|
|
# Inject into project AGENTS.md
|
|
project_agents = Path.cwd() / "AGENTS.md"
|
|
_inject_rtk_instructions(project_agents, verbose=verbose)
|
|
# Inject into global OpenCode AGENTS.md
|
|
global_agents = _opencode_home_dir() / "AGENTS.md"
|
|
_inject_rtk_instructions(global_agents, verbose=verbose)
|
|
|
|
# Register headroom MCP server in OpenCode config so OpenCode can
|
|
# call headroom_retrieve on compression markers from the proxy.
|
|
if not no_mcp:
|
|
from headroom.mcp_registry import OpencodeRegistrar
|
|
|
|
_setup_headroom_mcp(OpencodeRegistrar(), port, verbose=verbose, force=True)
|
|
elif verbose:
|
|
click.echo(" Skipping MCP retrieve tool (--no-mcp)")
|
|
|
|
if not no_serena:
|
|
from headroom.mcp_registry import OpencodeRegistrar
|
|
|
|
# Serena ships no "opencode" context (only agent/codex/claude-code/ide/…);
|
|
# passing --context opencode crashes Serena on launch (#1549/#1572). Use
|
|
# the generic "agent" context, which OpenCode is.
|
|
_setup_serena_mcp(OpencodeRegistrar(), context="agent", verbose=verbose, force=True)
|
|
else:
|
|
from headroom.mcp_registry import OpencodeRegistrar
|
|
|
|
_disable_serena_mcp(OpencodeRegistrar(), verbose=verbose)
|
|
|
|
# Setup memory MCP server for OpenCode (native tool integration)
|
|
if memory:
|
|
click.echo(" Setting up memory for OpenCode...")
|
|
mem_dir = Path.cwd() / ".headroom"
|
|
mem_dir.mkdir(parents=True, exist_ok=True)
|
|
mem_user = os.environ.get("USER", os.environ.get("USERNAME", "default"))
|
|
_inject_memory_mcp_config(mem_user)
|
|
agents_md = Path.cwd() / "AGENTS.md"
|
|
_inject_memory_agents_md(agents_md)
|
|
|
|
if prepare_only:
|
|
inject_opencode_provider_config(port)
|
|
return
|
|
|
|
opencode_bin = shutil.which("opencode")
|
|
if not opencode_bin:
|
|
click.echo("Error: 'opencode' not found in PATH.")
|
|
click.echo("Install OpenCode: https://opencode.ai")
|
|
raise SystemExit(1)
|
|
|
|
# Register our proxy client marker BEFORE _ensure_proxy so that another
|
|
# wrapper's cleanup sees us as an active client and doesn't terminate a
|
|
# shared proxy during the startup gap.
|
|
_register_proxy_client(port)
|
|
|
|
# Resolve port before config injection so the provider block and MCP
|
|
# URL both point at the port the proxy will actually be on.
|
|
_opencode_proxy, actual_port = _ensure_proxy(
|
|
port,
|
|
no_proxy,
|
|
learn=learn,
|
|
memory=memory,
|
|
agent_type="opencode",
|
|
code_graph=code_graph,
|
|
backend=backend,
|
|
anyllm_provider=anyllm_provider,
|
|
region=region,
|
|
)
|
|
|
|
# If the proxy fell back to a different port, move our marker so
|
|
# cleanup tracking stays accurate and update MCP config.
|
|
if actual_port != port:
|
|
_unregister_proxy_client(port)
|
|
_register_proxy_client(actual_port)
|
|
if not no_mcp:
|
|
from headroom.mcp_registry import OpencodeRegistrar
|
|
|
|
_setup_headroom_mcp(OpencodeRegistrar(), actual_port, verbose=verbose, force=True)
|
|
|
|
env, env_vars_display = _build_opencode_launch_env(
|
|
actual_port, os.environ, project=_project_name_from_cwd(), include_mcp=not no_mcp
|
|
)
|
|
|
|
# Inject Headroom provider into OpenCode config so traffic routes through proxy.
|
|
inject_opencode_provider_config(actual_port)
|
|
if memory:
|
|
mem_dir = Path.cwd() / ".headroom"
|
|
_inject_memory_mcp_config(
|
|
os.environ.get("USER", os.environ.get("USERNAME", "default")),
|
|
)
|
|
|
|
# Proxy already started by _ensure_proxy above; tell _launch_tool to
|
|
# skip duplicate startup.
|
|
try:
|
|
_launch_tool(
|
|
binary=opencode_bin,
|
|
args=opencode_args,
|
|
env=env,
|
|
port=actual_port,
|
|
no_proxy=True,
|
|
tool_label="OPENCODE",
|
|
env_vars_display=env_vars_display,
|
|
learn=learn,
|
|
memory=memory,
|
|
agent_type="opencode",
|
|
code_graph=code_graph,
|
|
backend=backend,
|
|
anyllm_provider=anyllm_provider,
|
|
region=region,
|
|
)
|
|
finally:
|
|
if _opencode_proxy and _opencode_proxy.poll() is None:
|
|
_other = _live_proxy_clients(actual_port, exclude_self=True)
|
|
if not _other:
|
|
_opencode_proxy.terminate()
|
|
try:
|
|
_opencode_proxy.wait(timeout=5)
|
|
except subprocess.TimeoutExpired:
|
|
_opencode_proxy.kill()
|
|
|
|
|
|
def _opencode_home_dir() -> Path:
|
|
"""Return the OpenCode home/config directory."""
|
|
env_path = os.environ.get("OPENCODE_HOME", "").strip()
|
|
if env_path:
|
|
return Path(env_path).expanduser()
|
|
return Path.home() / ".config" / "opencode"
|
|
|
|
|
|
# =============================================================================
|
|
# OpenCode (unwrap)
|
|
# =============================================================================
|
|
|
|
|
|
@unwrap.command("opencode")
|
|
@click.option(
|
|
"--port", "-p", default=8787, type=click.IntRange(1, 65535), help="Proxy port (default: 8787)"
|
|
)
|
|
@click.option("--no-stop-proxy", is_flag=True, help="Do not stop the local Headroom proxy")
|
|
def unwrap_opencode(port: int, no_stop_proxy: bool) -> None:
|
|
"""Undo ``headroom wrap opencode`` edits to the active OpenCode config file.
|
|
|
|
Behaviour:
|
|
|
|
* If a pre-wrap backup (``opencode.json.headroom-backup``) exists, the
|
|
original file is restored byte-for-byte and the backup is removed.
|
|
* Otherwise, if the config file still contains the Headroom-managed
|
|
block, that block is stripped out and the rest of the file is
|
|
preserved.
|
|
* If the config only ever contained Headroom-written content, the file
|
|
is removed entirely so OpenCode falls back to its defaults.
|
|
* If neither a backup nor a Headroom block is present, this is a safe
|
|
no-op.
|
|
"""
|
|
click.echo()
|
|
click.echo(" ╔═══════════════════════════════════════════════╗")
|
|
click.echo(" ║ HEADROOM UNWRAP: OPENCODE ║")
|
|
click.echo(" ╚═══════════════════════════════════════════════╝")
|
|
click.echo()
|
|
|
|
config_file, backup_file = opencode_config_paths()
|
|
|
|
if backup_file.exists():
|
|
try:
|
|
shutil.copy2(backup_file, config_file)
|
|
backup_file.unlink()
|
|
click.echo(f" Restored prior {config_file} from pre-wrap backup.")
|
|
status = "restored"
|
|
except OSError as exc:
|
|
raise click.ClickException(
|
|
f"could not restore OpenCode config from backup: {exc}"
|
|
) from exc
|
|
elif config_file.exists():
|
|
content = _read_text(config_file)
|
|
if _PROVIDER_MARKER_START in content or _MCP_MARKER_START in content:
|
|
cleaned = strip_opencode_headroom_blocks(content)
|
|
if cleaned.strip():
|
|
_write_text(config_file, cleaned + "\n")
|
|
click.echo(f" Removed Headroom block from {config_file}; other content preserved.")
|
|
status = "cleaned"
|
|
else:
|
|
config_file.unlink()
|
|
click.echo(f" Removed {config_file} (contained only Headroom-written config).")
|
|
status = "removed"
|
|
else:
|
|
click.echo(f" Nothing to undo: {config_file} has no Headroom wrap markers.")
|
|
status = "noop"
|
|
else:
|
|
click.echo(f" Nothing to undo: {config_file} does not exist.")
|
|
status = "noop"
|
|
|
|
# Remove Serena MCP if it was installed by Headroom.
|
|
# Also remove the headroom MCP server itself.
|
|
from headroom.mcp_registry import OpencodeRegistrar
|
|
|
|
opencode_registrar = OpencodeRegistrar()
|
|
if opencode_registrar.detect():
|
|
if opencode_registrar.unregister_server("headroom"):
|
|
click.echo(" Removed Headroom MCP server from OpenCode.")
|
|
serena_status = _remove_headroom_installed_serena_mcp(opencode_registrar)
|
|
if serena_status == "removed":
|
|
click.echo(" Removed Headroom-installed Serena MCP server from OpenCode.")
|
|
elif serena_status == "failed":
|
|
click.echo(" Serena MCP server matched Headroom ledger but could not be removed.")
|
|
|
|
click.echo()
|
|
click.echo("✓ OpenCode is no longer routed through the Headroom proxy.")
|
|
if not no_stop_proxy and status != "noop":
|
|
_echo_unwrap_proxy_stop_status(_stop_local_proxy_for_unwrap(port), port)
|
|
click.echo()
|
|
|
|
|
|
@unwrap.command("openclaw")
|
|
@click.option(
|
|
"--proxy-port", default=8787, type=click.IntRange(1, 65535), help="Headroom proxy port"
|
|
)
|
|
@click.option("--no-stop-proxy", is_flag=True, help="Do not stop the local Headroom proxy")
|
|
@click.option("--no-restart", is_flag=True, help="Do not restart OpenClaw gateway at the end")
|
|
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
|
|
@click.option("--prepare-only", is_flag=True, hidden=True)
|
|
@click.option("--existing-entry-json", default=None, hidden=True)
|
|
def unwrap_openclaw(
|
|
proxy_port: int,
|
|
no_stop_proxy: bool,
|
|
no_restart: bool,
|
|
verbose: bool,
|
|
prepare_only: bool,
|
|
existing_entry_json: str | None,
|
|
) -> None:
|
|
"""Disable the Headroom OpenClaw plugin and restore the legacy engine slot."""
|
|
if prepare_only:
|
|
click.echo(
|
|
json.dumps(
|
|
_build_openclaw_unwrap_entry(_decode_openclaw_entry_json(existing_entry_json)),
|
|
separators=(",", ":"),
|
|
)
|
|
)
|
|
return
|
|
|
|
openclaw_bin = shutil.which("openclaw")
|
|
if not openclaw_bin:
|
|
raise click.ClickException("'openclaw' not found in PATH. Install OpenClaw CLI first.")
|
|
|
|
click.echo()
|
|
click.echo(" ╔═══════════════════════════════════════════════╗")
|
|
click.echo(" ║ HEADROOM UNWRAP: OPENCLAW ║")
|
|
click.echo(" ╚═══════════════════════════════════════════════╝")
|
|
click.echo()
|
|
click.echo(" Disabling Headroom plugin and removing engine mapping...")
|
|
|
|
existing_entry = _read_openclaw_config_value(openclaw_bin, "plugins.entries.headroom")
|
|
entry = _build_openclaw_unwrap_entry(existing_entry)
|
|
_write_openclaw_plugin_entry(openclaw_bin, entry)
|
|
_set_openclaw_context_engine_slot(openclaw_bin, "legacy")
|
|
_run_checked(
|
|
[openclaw_bin, "config", "validate"],
|
|
action="openclaw config validate",
|
|
)
|
|
|
|
if no_restart:
|
|
click.echo(" Skipping gateway restart (--no-restart).")
|
|
click.echo(
|
|
" Run `openclaw gateway restart` (or `openclaw gateway start`) to apply unwrap changes."
|
|
)
|
|
else:
|
|
click.echo(" Applying unwrap changes to OpenClaw gateway...")
|
|
gateway_action, gateway_output = _restart_or_start_openclaw_gateway(openclaw_bin)
|
|
click.echo(f" Gateway {gateway_action}.")
|
|
if verbose and gateway_output:
|
|
click.echo(gateway_output)
|
|
|
|
if verbose:
|
|
inspect_result = _run_checked(
|
|
[openclaw_bin, "plugins", "inspect", "headroom"],
|
|
action="openclaw plugins inspect headroom",
|
|
)
|
|
if inspect_result.stdout.strip():
|
|
click.echo(inspect_result.stdout.strip())
|
|
|
|
click.echo()
|
|
click.echo("✓ OpenClaw Headroom wrap removed.")
|
|
click.echo(" Plugin: headroom (installed, disabled)")
|
|
click.echo(" Slot: plugins.slots.contextEngine = legacy")
|
|
if not no_stop_proxy:
|
|
_echo_unwrap_proxy_stop_status(_stop_local_proxy_for_unwrap(proxy_port), proxy_port)
|
|
click.echo()
|
|
|
|
|
|
# =============================================================================
|
|
# OpenAI Codex CLI (unwrap)
|
|
# =============================================================================
|
|
|
|
|
|
@unwrap.command("codex")
|
|
@click.option(
|
|
"--port", "-p", default=8787, type=click.IntRange(1, 65535), help="Proxy port (default: 8787)"
|
|
)
|
|
@click.option("--no-stop-proxy", is_flag=True, help="Do not stop the local Headroom proxy")
|
|
def unwrap_codex(port: int, no_stop_proxy: bool) -> None:
|
|
"""Undo ``headroom wrap codex`` edits to the active Codex config file.
|
|
|
|
Behaviour:
|
|
|
|
* If a pre-wrap backup (``config.toml.headroom-backup``) exists, the
|
|
original file is restored byte-for-byte and the backup is removed.
|
|
* Otherwise, if the config file still contains the Headroom-managed
|
|
block, that block is stripped out and the rest of the file is
|
|
preserved.
|
|
* If the config only ever contained Headroom-written content, the file
|
|
is removed entirely so Codex falls back to its defaults.
|
|
* If neither a backup nor a Headroom block is present, this is a safe
|
|
no-op (the user either never wrapped that config, or already unwrapped
|
|
it). When ``CODEX_HOME`` is unset, print a warning hint because Headroom
|
|
may be looking at the default config while Codex was wrapped with a
|
|
custom home.
|
|
"""
|
|
click.echo()
|
|
click.echo(" ╔═══════════════════════════════════════════════╗")
|
|
click.echo(" ║ HEADROOM UNWRAP: CODEX ║")
|
|
click.echo(" ╚═══════════════════════════════════════════════╝")
|
|
click.echo()
|
|
|
|
try:
|
|
status, config_file = _restore_codex_provider_config()
|
|
except Exception as e: # pragma: no cover - filesystem-level errors
|
|
raise click.ClickException(f"could not unwrap Codex config: {e}") from e
|
|
|
|
if status == "restored":
|
|
click.echo(f" Restored prior {config_file} from pre-wrap backup.")
|
|
elif status == "cleaned":
|
|
click.echo(f" Removed Headroom block from {config_file}; other content preserved.")
|
|
elif status == "removed":
|
|
click.echo(f" Removed {config_file} (contained only Headroom-written config).")
|
|
else:
|
|
if not os.environ.get("CODEX_HOME"):
|
|
click.echo(
|
|
" Warning: found no Headroom wrap markers in the default Codex config. "
|
|
"If you wrapped Codex with CODEX_HOME, rerun unwrap with the same "
|
|
"environment variable, e.g. CODEX_HOME=/path/to/codex-home "
|
|
"headroom unwrap codex."
|
|
)
|
|
click.echo(f" Nothing to undo: {config_file} has no Headroom wrap markers.")
|
|
|
|
# tokensave and Serena are each written as their own [mcp_servers.<name>]
|
|
# table with Headroom markers, separate from the provider block handled
|
|
# above — a "cleaned" restore leaves them behind. Remove them explicitly
|
|
# (only if we installed them), mirroring unwrap_claude. Runs after the
|
|
# restore so a backup-restore that already dropped them is a safe no-op.
|
|
from headroom.mcp_registry import CodexRegistrar
|
|
|
|
codex_registrar = CodexRegistrar()
|
|
if codex_registrar.detect():
|
|
tokensave_status = _remove_headroom_installed_tokensave_mcp(codex_registrar)
|
|
if tokensave_status == "removed":
|
|
click.echo(" Removed Headroom-installed tokensave MCP server from Codex.")
|
|
elif tokensave_status == "failed":
|
|
click.echo(" tokensave MCP server matched Headroom ledger but could not be removed.")
|
|
|
|
serena_status = _remove_headroom_installed_serena_mcp(codex_registrar)
|
|
if serena_status == "removed":
|
|
click.echo(" Removed Headroom-installed Serena MCP server from Codex.")
|
|
elif serena_status == "failed":
|
|
click.echo(" Serena MCP server matched Headroom ledger but could not be removed.")
|
|
|
|
# `wrap codex` injects the marker-fenced rtk guidance into the Codex global
|
|
# AGENTS.md (`_codex_home_dir() / "AGENTS.md"`); that block is durable state
|
|
# the config restore above does not touch. Without removing it, a plain
|
|
# `codex` launch keeps following Headroom's "prefix shell commands with rtk"
|
|
# instruction and fails when the managed rtk binary is off PATH. Mirror what
|
|
# unwrap_copilot already does. Best-effort and unconditional, like the MCP
|
|
# cleanup above.
|
|
if _remove_rtk_instructions(_codex_home_dir() / "AGENTS.md"):
|
|
click.echo(" Removed Headroom rtk instructions from Codex AGENTS.md.")
|
|
|
|
if status in {"restored", "cleaned", "removed"}:
|
|
# Hand the threads back to the native-provider menu so the full history
|
|
# stays visible once Codex no longer routes through Headroom. Best-effort.
|
|
retag_to_native(_codex_home_dir())
|
|
|
|
click.echo()
|
|
click.echo("✓ Codex is no longer routed through the Headroom proxy.")
|
|
if not no_stop_proxy and status != "noop":
|
|
_echo_unwrap_proxy_stop_status(_stop_local_proxy_for_unwrap(port), port)
|
|
click.echo()
|