mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat(wrap): make tokensave the primary coding-task compressor, Serena the backup (#1230)
## Description Makes **tokensave** ([github.com/aovestdipaperino/tokensave](https://github.com/aovestdipaperino/tokensave)) the **primary coding-task compressor** that `headroom wrap` installs, and demotes **Serena** to a **backup**. tokensave is a local semantic code-graph MCP server (`tokensave serve`): the agent queries it for symbols, call chains, and impact analysis instead of grepping/reading whole files — the same role Serena filled, but as a pre-indexed graph. Serena now only registers when tokensave is unavailable (or when forced with `--serena`). Closes # ## Type of Change - [ ] Bug fix (non-breaking change that fixes an issue) - [x] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/graph/tokensave_installer.py` (new): fetch the prebuilt tokensave release binary for the platform (release-binary only — no `cargo` compile at wrap time); honors `HEADROOM_BINARIES_OFFLINE`; returns `None` (→ Serena) when no asset exists (e.g. x86_64 macOS) or the download fails. - `mcp_registry`: `build_tokensave_spec()`; registration/disable/migrate go through the existing `ServerSpec` + ownership-ledger flow, identical to Serena. - `cli/wrap.py`: new `_setup_coding_compressor` primary/backup policy; tokensave setup/disable/migrate/index helpers. New flags `--no-tokensave` (skip primary) and `--serena` (force backup on); `--no-serena` now means "never register the backup". Default wrap removes a previously Headroom-installed Serena entry once tokensave is primary (user-managed entries preserved). `--code-graph` repointed to tokensave; the legacy `codebase-memory-mcp` install path is dropped (unwrap still cleans up legacy entries). `unwrap claude|codex` remove a ledger-owned tokensave entry. - Strands `HeadroomBundle`: `enable_tokensave_mcp=True` (primary); `enable_serena_mcp` now defaults `False` (backup). - `docs/content/docs/proxy.mdx`: `--code-graph` description updated from codebase-memory-mcp to tokensave. - Tests: tokensave installer (incl. error paths), register/disable/migrate, primary/backup policy, and the binary-resolution/indexing helpers. A scoped `tests/test_cli/conftest.py` offline guard keeps the CLI suite hermetic. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text $ uv run pytest -q tests/test_graph_tokensave.py tests/test_cli/test_tokensave_setup.py tests/test_cli/test_tokensave_helpers.py 41 passed $ uv run pytest -q tests/test_cli/ tests/test_graph.py tests/test_graph_tokensave.py 421 passed # full CLI + graph suites, incl. all pre-existing Serena/unwrap/registry tests $ uv run pytest -q tests/test_mcp_registry/ tests/test_proxy_healthchecks.py passed $ uv run ruff format --check headroom/ tests/ # 822 files already formatted $ uv run ruff check <changed files> # All checks passed! $ uv run mypy headroom/graph/tokensave_installer.py headroom/mcp_registry/install.py Success: no issues found in 2 source files # Coverage on new module headroom/graph/tokensave_installer.py 99% ``` ## Real Behavior Proof - Environment: macOS (darwin arm64), Python 3.14, `uv` dev env; tokensave 7.0.2 binary present on PATH and exercised against this repo's `.tokensave/` graph during development. The installer pins release **v7.0.2** (SHA-256-verified) across macOS arm64, Linux aarch64/x86_64, and Windows x86_64/aarch64. - Exact command / steps: `headroom wrap claude` registers `tokensave serve` as the primary MCP code-graph server and indexes the project; with the binary removed from PATH and `HEADROOM_BINARIES_OFFLINE=1`, the same command falls back to registering Serena. Behavior is pinned by the unit tests (binary-present → tokensave registered + Serena entry removed; binary-absent → Serena fallback; `--serena` forces backup on; `--no-serena` suppresses it; `--no-tokensave` disables primary). - Observed result: tokensave registered as primary on the binary-present path; Serena registered on the unavailable path; unwrap removes only ledger-owned entries. - Not tested: live end-to-end agent session inside Claude Code / Codex against a real provider API; Windows/Linux release-asset download (covered by unit tests with mocked archives, not a live fetch). ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable ## Additional Notes - CHANGELOG is left untouched: this repo generates it via release-please from Conventional Commits, so a manual edit is N/A. - `strands/bundle.py` shows 0% patch coverage because that module hard-imports the optional `strands` SDK, which CI does not install (the pre-existing `_make_serena_client` was likewise uncovered) — not a regression. - A `test (3)` shard failure on `headroom.memory.bridge` is a pre-existing offline-CI flake (cannot reach huggingface.co); it touches no file in this PR and the scoped offline guard only applies under `tests/test_cli/`. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Paperinik
GitHub
parent
c6c921a7c1
commit
dca9853ed9
13 changed files with 1512 additions and 579 deletions
1
.gitignore
vendored
1
.gitignore
vendored
|
|
@ -252,3 +252,4 @@ uv.lock
|
|||
# package shadows the maturin overlay on sys.path.
|
||||
/headroom/_core.*.so
|
||||
/headroom/_core.so
|
||||
.tokensave
|
||||
|
|
|
|||
664
Cargo.lock
generated
664
Cargo.lock
generated
File diff suppressed because it is too large
Load diff
|
|
@ -59,7 +59,15 @@ Telemetry is **off by default** (opt-in). Opt in with `HEADROOM_TELEMETRY=on` or
|
|||
| `--intercept-tool-results` | `false` | Opt into tool-result interceptors such as ast-grep Read outlining. |
|
||||
| `--no-read-lifecycle` | `false` | Disable stale/superseded Read-output compression. |
|
||||
| `--code-aware` / `--no-code-aware` | disabled | Enable or disable AST-based code compression. Requires `headroom-ai[code]`. |
|
||||
| `--code-graph` | `false` | Index the current project and watch files via codebase-memory-mcp. |
|
||||
| `--code-graph` | `false` | Force a tokensave code-graph index of the current project (tokensave is the default coding-task compressor registered by `headroom wrap`). |
|
||||
|
||||
#### tokensave binary trust model
|
||||
|
||||
`headroom wrap` registers **tokensave** (a local code-graph MCP server) as the default coding-task compressor. tokensave ships as a single prebuilt Rust binary, so `wrap` downloads the release asset for your platform from GitHub and runs it locally. Because the binary is executed, every supported asset is **pinned to a SHA-256 digest in Headroom** (`headroom/graph/tokensave_installer.py`); the downloaded bytes are verified against that digest before extraction, and a mismatch aborts the install (Headroom falls back to the Serena backup) rather than running unverified code.
|
||||
|
||||
- Set `HEADROOM_BINARIES_OFFLINE=1` to never reach the network — `wrap` then uses an already-installed tokensave or falls back to Serena.
|
||||
- `HEADROOM_TOKENSAVE_VERSION` overrides the pinned release tag. Since an overridden version has no pinned digest, the download is **refused** unless you also set `HEADROOM_TOKENSAVE_ALLOW_UNVERIFIED=1`.
|
||||
- Pass `--no-tokensave` to skip the primary compressor entirely, or `--serena` to force the Serena backup on.
|
||||
|
||||
By default, the proxy uses the shared **ContentRouter** pipeline. It routes text, logs, JSON, code, images, and tool outputs through the currently enabled compressors and preserves reversible CCR markers where applicable.
|
||||
|
||||
|
|
|
|||
|
|
@ -866,15 +866,18 @@ def _remove_headroom_installed_serena_mcp(registrar: Any) -> str:
|
|||
return "failed"
|
||||
|
||||
|
||||
def _disable_serena_mcp(registrar: Any, *, verbose: bool = False) -> None:
|
||||
"""Make ``--no-serena`` actively disable Serena, not merely skip adding it.
|
||||
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 is 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
|
||||
``--no-serena`` run leaves that stale entry in place — so the flag has to
|
||||
remove the entry Headroom installed. A user-managed Serena (absent from
|
||||
our ledger) is reported but left untouched.
|
||||
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:
|
||||
|
|
@ -883,12 +886,12 @@ def _disable_serena_mcp(registrar: Any, *, verbose: bool = False) -> None:
|
|||
|
||||
if registrar.get_server("serena") is None:
|
||||
if verbose:
|
||||
click.echo(" Skipping Serena MCP (--no-serena)")
|
||||
click.echo(f" Skipping Serena MCP ({reason})")
|
||||
return
|
||||
|
||||
status = _remove_headroom_installed_serena_mcp(registrar)
|
||||
if status == "removed":
|
||||
click.echo(" Removed previously-installed Serena MCP (--no-serena)")
|
||||
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(
|
||||
|
|
@ -901,122 +904,219 @@ def _disable_serena_mcp(registrar: Any, *, verbose: bool = False) -> None:
|
|||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# 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 _register_cbm_mcp_server(cbm_bin: str) -> None:
|
||||
"""Register codebase-memory-mcp as an MCP server in Claude Code.
|
||||
|
||||
Uses ``claude mcp add`` so the tools appear in ``/mcp`` automatically.
|
||||
Idempotent — skips if already registered.
|
||||
"""
|
||||
claude_cli = shutil.which("claude")
|
||||
if not claude_cli:
|
||||
return
|
||||
|
||||
# Check if already registered
|
||||
check = run(
|
||||
[claude_cli, "mcp", "get", _CBM_MCP_SERVER_NAME],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
if check.returncode == 0:
|
||||
return # Already registered
|
||||
|
||||
result = run(
|
||||
[claude_cli, "mcp", "add", _CBM_MCP_SERVER_NAME, "-s", "user", "--", cbm_bin],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
if result.returncode == 0:
|
||||
click.echo(f" Code graph: registered {_CBM_MCP_SERVER_NAME} MCP server")
|
||||
else:
|
||||
pass # Non-critical — tools won't appear in /mcp but graph still works
|
||||
|
||||
|
||||
def _setup_code_graph(verbose: bool = False) -> bool:
|
||||
"""Ensure codebase-memory-mcp is installed, registered as MCP server, and project is indexed.
|
||||
"""Ensure the tokensave code graph is set up and the project indexed.
|
||||
|
||||
codebase-memory-mcp builds a knowledge graph of the codebase using
|
||||
tree-sitter, enabling the LLM to query code structure (call chains,
|
||||
function definitions, impact analysis) instead of reading entire files.
|
||||
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.
|
||||
|
||||
Steps:
|
||||
1. Download the binary if not already present.
|
||||
2. Register as an MCP server in Claude Code (``claude mcp add``).
|
||||
3. Index the current project (fast, idempotent).
|
||||
|
||||
With Claude Code's MCP Tool Search, the 14 graph tools add ~200 tokens
|
||||
overhead per request (not the full ~1,915) — they're lazy-loaded.
|
||||
|
||||
Returns True if graph is ready, False if setup failed.
|
||||
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.graph.installer import ensure_cbm, get_cbm_path
|
||||
from headroom.mcp_registry import ClaudeRegistrar
|
||||
|
||||
cbm_path = get_cbm_path()
|
||||
if not cbm_path:
|
||||
click.echo(" Code graph: downloading codebase-memory-mcp...")
|
||||
cbm_path = ensure_cbm()
|
||||
if cbm_path:
|
||||
click.echo(f" Code graph: installed at {cbm_path}")
|
||||
else:
|
||||
click.echo(" Code graph: download failed — skipping")
|
||||
return False
|
||||
|
||||
cbm_bin = str(cbm_path)
|
||||
|
||||
# Register as MCP server so tools appear in /mcp
|
||||
_register_cbm_mcp_server(cbm_bin)
|
||||
|
||||
# Index current project (fast — ~1s for most repos, idempotent)
|
||||
project_dir = str(Path.cwd())
|
||||
try:
|
||||
result = run(
|
||||
[
|
||||
cbm_bin,
|
||||
"cli",
|
||||
"index_repository",
|
||||
json.dumps({"repo_path": project_dir, "mode": "fast"}),
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
)
|
||||
if result.returncode == 0:
|
||||
# Parse node/edge counts from output
|
||||
for line in result.stdout.splitlines():
|
||||
if '"nodes"' in line:
|
||||
try:
|
||||
# Parse the JSON response to extract node/edge counts
|
||||
import re
|
||||
|
||||
m_nodes = re.search(r'"nodes":(\d+)', line)
|
||||
m_edges = re.search(r'"edges":(\d+)', line)
|
||||
if m_nodes and m_edges:
|
||||
nodes = int(m_nodes.group(1))
|
||||
edges = int(m_edges.group(1))
|
||||
click.echo(
|
||||
f" Code graph: indexed ({nodes:,} symbols, "
|
||||
f"{edges:,} relationships)"
|
||||
)
|
||||
else:
|
||||
click.echo(" Code graph: indexed")
|
||||
except (ValueError, AttributeError):
|
||||
click.echo(" Code graph: indexed")
|
||||
return True
|
||||
click.echo(" Code graph: indexed")
|
||||
return True
|
||||
else:
|
||||
if verbose:
|
||||
click.echo(f" Code graph: indexing failed ({result.stderr[:100]})")
|
||||
return False
|
||||
except subprocess.TimeoutExpired:
|
||||
click.echo(" Code graph: indexing timed out (will complete in background)")
|
||||
return False
|
||||
except Exception as e:
|
||||
if verbose:
|
||||
click.echo(f" Code graph: setup failed ({e})")
|
||||
return False
|
||||
return _setup_tokensave_mcp(ClaudeRegistrar(), verbose=verbose, force=True)
|
||||
|
||||
|
||||
# rtk instructions for tools without hook support (Codex, Cursor, Aider).
|
||||
|
|
@ -3073,11 +3173,22 @@ def unwrap() -> None:
|
|||
is_flag=True,
|
||||
help="Skip headroom MCP server registration (compression markers will be unactionable)",
|
||||
)
|
||||
@click.option("--no-serena", is_flag=True, help="Skip Serena MCP server registration")
|
||||
@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="Enable code graph indexing via codebase-memory-mcp (optional)",
|
||||
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(
|
||||
|
|
@ -3127,6 +3238,8 @@ def claude(
|
|||
port: int,
|
||||
no_rtk: bool,
|
||||
no_mcp: bool,
|
||||
no_tokensave: bool,
|
||||
serena: bool,
|
||||
no_serena: bool,
|
||||
code_graph: bool,
|
||||
no_proxy: bool,
|
||||
|
|
@ -3152,10 +3265,12 @@ def claude(
|
|||
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 --code-graph # With code graph intelligence
|
||||
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 --no-context-tool # Skip CLI context-tool setup
|
||||
headroom wrap claude --no-mcp # Skip MCP retrieve tool registration
|
||||
headroom wrap claude --no-serena # Skip Serena MCP registration
|
||||
headroom wrap claude --no-serena # Never register the Serena backup
|
||||
headroom wrap claude --1m # Preserve the 1M context window
|
||||
"""
|
||||
if prepare_only:
|
||||
|
|
@ -3285,14 +3400,17 @@ def claude(
|
|||
elif verbose:
|
||||
click.echo(" Skipping MCP retrieve tool (--no-mcp)")
|
||||
|
||||
if not no_serena:
|
||||
from headroom.mcp_registry import ClaudeRegistrar
|
||||
# Coding-task compressor: tokensave primary, Serena backup.
|
||||
from headroom.mcp_registry import ClaudeRegistrar
|
||||
|
||||
_setup_serena_mcp(ClaudeRegistrar(), context="claude-code", verbose=verbose)
|
||||
else:
|
||||
from headroom.mcp_registry import ClaudeRegistrar
|
||||
|
||||
_disable_serena_mcp(ClaudeRegistrar(), verbose=verbose)
|
||||
_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)
|
||||
|
|
@ -3410,13 +3528,20 @@ def unwrap_claude(
|
|||
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 code graph MCP server from Claude.")
|
||||
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":
|
||||
|
|
@ -3770,11 +3895,22 @@ def unwrap_copilot(port: int, no_stop_proxy: bool) -> None:
|
|||
is_flag=True,
|
||||
help="Skip headroom MCP server registration (compression markers will be unactionable)",
|
||||
)
|
||||
@click.option("--no-serena", is_flag=True, help="Skip Serena MCP server registration")
|
||||
@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="Enable code graph indexing via codebase-memory-mcp (optional)",
|
||||
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(
|
||||
|
|
@ -3801,6 +3937,8 @@ def codex(
|
|||
port: int,
|
||||
no_rtk: bool,
|
||||
no_mcp: bool,
|
||||
no_tokensave: bool,
|
||||
serena: bool,
|
||||
no_serena: bool,
|
||||
code_graph: bool,
|
||||
no_proxy: bool,
|
||||
|
|
@ -3828,7 +3966,9 @@ def 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-serena # Skip Serena MCP 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
|
||||
"""
|
||||
|
|
@ -3866,14 +4006,19 @@ def codex(
|
|||
elif verbose:
|
||||
click.echo(" Skipping MCP retrieve tool (--no-mcp)")
|
||||
|
||||
if not no_serena:
|
||||
from headroom.mcp_registry import CodexRegistrar
|
||||
# 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_serena_mcp(CodexRegistrar(), context="codex", verbose=verbose, force=True)
|
||||
else:
|
||||
from headroom.mcp_registry import CodexRegistrar
|
||||
|
||||
_disable_serena_mcp(CodexRegistrar(), verbose=verbose)
|
||||
_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:
|
||||
|
|
@ -5317,15 +5462,21 @@ def unwrap_codex(port: int, no_stop_proxy: bool) -> None:
|
|||
)
|
||||
click.echo(f" Nothing to undo: {config_file} has no Headroom wrap markers.")
|
||||
|
||||
# Serena is written as its own [mcp_servers.serena] table with Headroom
|
||||
# markers, separate from the provider block handled above — a "cleaned"
|
||||
# restore leaves it behind. Remove it explicitly (only if we installed it),
|
||||
# mirroring unwrap_claude. Runs after the restore so a backup-restore that
|
||||
# already dropped Serena makes this a safe no-op.
|
||||
# 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.")
|
||||
|
|
|
|||
256
headroom/graph/tokensave_installer.py
Normal file
256
headroom/graph/tokensave_installer.py
Normal file
|
|
@ -0,0 +1,256 @@
|
|||
"""Download and install the ``tokensave`` binary from GitHub releases.
|
||||
|
||||
tokensave (https://github.com/aovestdipaperino/tokensave) is the primary
|
||||
coding-task compressor: a local semantic code-graph MCP server. It is a
|
||||
single self-contained Rust binary, so — like ``codebase-memory-mcp`` and
|
||||
``rtk`` — Headroom fetches the prebuilt release asset for the current
|
||||
platform, caches it under ``~/.local/bin``, and registers it as an MCP
|
||||
server.
|
||||
|
||||
Release-binary only. tokensave is also published to crates.io
|
||||
(``cargo install tokensave``), but we never shell out to cargo here: a
|
||||
multi-minute compile is the wrong thing to trigger from ``headroom wrap``.
|
||||
When no prebuilt asset exists for the platform (e.g. x86_64 macOS, which
|
||||
tokensave does not currently publish) or the download fails, this module
|
||||
returns ``None`` and the caller falls back to Serena, the backup compressor.
|
||||
|
||||
Supply-chain integrity:
|
||||
Because ``headroom wrap`` downloads and then *executes* this binary by
|
||||
default, every release asset is pinned to a SHA-256 digest in
|
||||
``TOKENSAVE_ASSET_DIGESTS`` below. The downloaded bytes are verified
|
||||
against the pinned digest before the archive is unpacked; a mismatch
|
||||
aborts the install (→ Serena fallback) rather than running unverified
|
||||
code. When ``HEADROOM_TOKENSAVE_VERSION`` overrides the pinned tag there
|
||||
is no pinned digest, so the download is refused unless the operator
|
||||
explicitly opts out of verification via
|
||||
``HEADROOM_TOKENSAVE_ALLOW_UNVERIFIED=1``.
|
||||
|
||||
Env vars:
|
||||
HEADROOM_BINARIES_OFFLINE if set, never reach the network (returns
|
||||
the already-installed binary or ``None``).
|
||||
HEADROOM_TOKENSAVE_VERSION override the pinned release tag.
|
||||
HEADROOM_TOKENSAVE_ALLOW_UNVERIFIED permit installing an asset that has
|
||||
no pinned digest (only relevant when the
|
||||
version is overridden).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
import platform
|
||||
import stat
|
||||
import tarfile
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from urllib.request import urlopen
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
#: Pinned release. Override with HEADROOM_TOKENSAVE_VERSION.
|
||||
TOKENSAVE_VERSION = "v7.0.2"
|
||||
TOKENSAVE_REPO = "aovestdipaperino/tokensave"
|
||||
TOKENSAVE_BIN_DIR = Path.home() / ".local" / "bin"
|
||||
TOKENSAVE_BIN_NAME = "tokensave"
|
||||
|
||||
GITHUB_RELEASE_URL = f"https://github.com/{TOKENSAVE_REPO}/releases/download"
|
||||
|
||||
#: SHA-256 of each pinned release asset, keyed by asset filename. The binary
|
||||
#: is downloaded and executed by default, so its bytes are verified against
|
||||
#: this map before extraction. Regenerate when bumping TOKENSAVE_VERSION:
|
||||
#: for f in <assets>; do curl -sL <url>/$f | shasum -a 256; done
|
||||
TOKENSAVE_ASSET_DIGESTS: dict[str, str] = {
|
||||
"tokensave-v7.0.2-aarch64-macos.tar.gz": (
|
||||
"6d0e07aba5b63df278409feabea54bdd0da82ec63d633cd975ea353773c4efee"
|
||||
),
|
||||
"tokensave-v7.0.2-aarch64-linux.tar.gz": (
|
||||
"69c88d0617036d44f2620f5779cd8578fad77664c2373d64de632b8e346ad334"
|
||||
),
|
||||
"tokensave-v7.0.2-x86_64-linux.tar.gz": (
|
||||
"d35519fe698a24d2e2bb5622e94b3bdb4794dc1e36acffc980260b50afb40460"
|
||||
),
|
||||
"tokensave-v7.0.2-x86_64-windows.zip": (
|
||||
"85f90d358c5f4713b5ac7274f4fa46e985fabc5b76c843ea8456b0d74e1cdd02"
|
||||
),
|
||||
"tokensave-v7.0.2-aarch64-windows.zip": (
|
||||
"8706d0d64f429ba7fe58deec9fef319956306797bded476cab4132e71705e8b0"
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _pinned_version() -> str:
|
||||
return os.environ.get("HEADROOM_TOKENSAVE_VERSION", "").strip() or TOKENSAVE_VERSION
|
||||
|
||||
|
||||
def _detect_asset(version: str) -> tuple[str, str] | None:
|
||||
"""Return ``(asset_filename, archive_kind)`` for this platform.
|
||||
|
||||
``archive_kind`` is ``"tar.gz"`` or ``"zip"``. Returns ``None`` when
|
||||
tokensave publishes no prebuilt asset for the current platform (the
|
||||
caller then falls back to Serena). Release assets are named
|
||||
``tokensave-<version>-<arch>-<os>.<ext>``.
|
||||
"""
|
||||
system = platform.system().lower()
|
||||
machine = platform.machine().lower()
|
||||
|
||||
if system == "darwin":
|
||||
if machine == "arm64":
|
||||
return f"tokensave-{version}-aarch64-macos.tar.gz", "tar.gz"
|
||||
# No x86_64-macos release asset is published — fall back to Serena.
|
||||
return None
|
||||
if system == "linux":
|
||||
arch = "aarch64" if machine in ("aarch64", "arm64") else "x86_64"
|
||||
return f"tokensave-{version}-{arch}-linux.tar.gz", "tar.gz"
|
||||
if system == "windows":
|
||||
arch = "aarch64" if machine in ("aarch64", "arm64") else "x86_64"
|
||||
return f"tokensave-{version}-{arch}-windows.zip", "zip"
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _verify_asset_digest(filename: str, data: bytes) -> None:
|
||||
"""Verify downloaded bytes against the pinned SHA-256 digest.
|
||||
|
||||
Raises ``RuntimeError`` on a digest mismatch, or when the asset has no
|
||||
pinned digest (i.e. a version override) unless the operator has set
|
||||
``HEADROOM_TOKENSAVE_ALLOW_UNVERIFIED``.
|
||||
"""
|
||||
expected = TOKENSAVE_ASSET_DIGESTS.get(filename)
|
||||
if expected is None:
|
||||
if os.environ.get("HEADROOM_TOKENSAVE_ALLOW_UNVERIFIED"):
|
||||
logger.warning(
|
||||
"tokensave asset %s has no pinned digest; installing unverified "
|
||||
"(HEADROOM_TOKENSAVE_ALLOW_UNVERIFIED is set)",
|
||||
filename,
|
||||
)
|
||||
return
|
||||
raise RuntimeError(
|
||||
f"no pinned SHA-256 digest for tokensave asset {filename!r}; refusing to "
|
||||
"install unverified. Set HEADROOM_TOKENSAVE_ALLOW_UNVERIFIED=1 to override."
|
||||
)
|
||||
actual = hashlib.sha256(data).hexdigest()
|
||||
if actual != expected:
|
||||
raise RuntimeError(
|
||||
f"tokensave asset {filename!r} failed integrity check: "
|
||||
f"expected sha256 {expected}, got {actual}"
|
||||
)
|
||||
logger.debug("Verified tokensave asset %s (sha256 %s)", filename, actual)
|
||||
|
||||
|
||||
def get_tokensave_path() -> Path | None:
|
||||
"""Find the tokensave binary on PATH or in our install dir; else ``None``."""
|
||||
import shutil
|
||||
|
||||
found = shutil.which(TOKENSAVE_BIN_NAME)
|
||||
if found:
|
||||
return Path(found)
|
||||
|
||||
for name in (TOKENSAVE_BIN_NAME, f"{TOKENSAVE_BIN_NAME}.exe"):
|
||||
installed = TOKENSAVE_BIN_DIR / name
|
||||
if installed.exists() and installed.is_file():
|
||||
return installed
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def download_tokensave(version: str | None = None) -> Path:
|
||||
"""Download and unpack the tokensave release binary. Returns its path.
|
||||
|
||||
Raises ``RuntimeError`` when no asset exists for this platform, or when
|
||||
the download / extraction / verification fails.
|
||||
"""
|
||||
version = version or _pinned_version()
|
||||
asset = _detect_asset(version)
|
||||
if asset is None:
|
||||
raise RuntimeError(
|
||||
f"no prebuilt tokensave asset for {platform.system()} {platform.machine()}"
|
||||
)
|
||||
filename, kind = asset
|
||||
url = f"{GITHUB_RELEASE_URL}/{version}/{filename}"
|
||||
|
||||
TOKENSAVE_BIN_DIR.mkdir(parents=True, exist_ok=True)
|
||||
bin_name = f"{TOKENSAVE_BIN_NAME}.exe" if kind == "zip" else TOKENSAVE_BIN_NAME
|
||||
target_path = TOKENSAVE_BIN_DIR / bin_name
|
||||
|
||||
logger.info("Downloading tokensave %s for %s ...", version, filename)
|
||||
|
||||
try:
|
||||
if not url.startswith(("http://", "https://")):
|
||||
raise ValueError(f"Invalid URL: {url}")
|
||||
with urlopen(url, timeout=60) as response: # noqa: S310
|
||||
data = response.read()
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to download tokensave from {url}: {e}") from e
|
||||
|
||||
_verify_asset_digest(filename, data)
|
||||
|
||||
try:
|
||||
if kind == "tar.gz":
|
||||
with tarfile.open(fileobj=io.BytesIO(data), mode="r:gz") as tar:
|
||||
for member in tar.getmembers():
|
||||
if member.name == TOKENSAVE_BIN_NAME or member.name.endswith(
|
||||
f"/{TOKENSAVE_BIN_NAME}"
|
||||
):
|
||||
member.name = target_path.name
|
||||
tar.extract(member, TOKENSAVE_BIN_DIR)
|
||||
break
|
||||
else:
|
||||
raise RuntimeError("tokensave binary not found in archive")
|
||||
else: # zip
|
||||
with zipfile.ZipFile(io.BytesIO(data)) as zf:
|
||||
for name in zf.namelist():
|
||||
if name.endswith(f"{TOKENSAVE_BIN_NAME}.exe") or name.endswith(
|
||||
f"/{TOKENSAVE_BIN_NAME}"
|
||||
):
|
||||
with zf.open(name) as src, open(target_path, "wb") as dst:
|
||||
dst.write(src.read())
|
||||
break
|
||||
else:
|
||||
raise RuntimeError("tokensave binary not found in archive")
|
||||
except (tarfile.TarError, zipfile.BadZipFile) as e:
|
||||
raise RuntimeError(f"Failed to extract tokensave archive: {e}") from e
|
||||
|
||||
if kind != "zip":
|
||||
target_path.chmod(target_path.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH)
|
||||
|
||||
try:
|
||||
from headroom._subprocess import run as _run
|
||||
|
||||
result = _run(
|
||||
[str(target_path), "--version"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=10,
|
||||
)
|
||||
if result.returncode == 0:
|
||||
logger.info("Installed tokensave: %s", result.stdout.strip())
|
||||
else:
|
||||
logger.warning("tokensave installed but version check failed")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return target_path
|
||||
|
||||
|
||||
def ensure_tokensave(version: str | None = None) -> Path | None:
|
||||
"""Ensure tokensave is available, downloading the release binary if needed.
|
||||
|
||||
Returns the binary path, or ``None`` when the binary is absent and cannot
|
||||
be fetched (offline, unsupported platform, or download failure). Callers
|
||||
treat ``None`` as "tokensave unavailable → fall back to Serena".
|
||||
"""
|
||||
existing = get_tokensave_path()
|
||||
if existing:
|
||||
return existing
|
||||
|
||||
if os.environ.get("HEADROOM_BINARIES_OFFLINE"):
|
||||
logger.info("tokensave not installed and HEADROOM_BINARIES_OFFLINE set — skipping download")
|
||||
return None
|
||||
|
||||
try:
|
||||
return download_tokensave(version)
|
||||
except RuntimeError as e:
|
||||
logger.warning("Could not install tokensave: %s", e)
|
||||
return None
|
||||
|
|
@ -90,6 +90,6 @@ __all__ = [
|
|||
# Provider detection
|
||||
"get_headroom_provider",
|
||||
"get_model_name_from_strands",
|
||||
# One-helper MCP + hook wiring (Headroom + Serena + RTK-equivalent)
|
||||
# One-helper MCP + hook wiring (Headroom + tokensave/Serena + RTK-equivalent)
|
||||
"HeadroomBundle",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -11,8 +11,14 @@ Strands-native primitives:
|
|||
needs the original; Strands' MCP dispatcher resolves it via this
|
||||
server. Works identically in streaming and non-streaming.
|
||||
|
||||
* **Serena MCP** — semantic code intelligence (symbol search,
|
||||
references, etc.). Auto-installed via ``uvx`` on first launch.
|
||||
* **tokensave MCP** — the primary coding-task compressor: a local
|
||||
semantic code-graph server (``tokensave serve``) the agent queries
|
||||
for symbols, call chains, and impact analysis instead of reading
|
||||
whole files. Requires the ``tokensave`` binary on PATH.
|
||||
|
||||
* **Serena MCP** — the backup coding-task compressor (symbol search,
|
||||
references, etc.), auto-installed via ``uvx`` on first launch.
|
||||
Off by default; enable with ``enable_serena_mcp=True``.
|
||||
|
||||
* **HeadroomHookProvider** — the RTK-equivalent for Strands.
|
||||
Compresses tool outputs in-place via ``AfterToolCallEvent`` so
|
||||
|
|
@ -77,6 +83,7 @@ from headroom.mcp_registry.install import (
|
|||
DEFAULT_PROXY_URL,
|
||||
build_headroom_spec,
|
||||
build_serena_spec,
|
||||
build_tokensave_spec,
|
||||
)
|
||||
|
||||
from .hooks import HeadroomHookProvider
|
||||
|
|
@ -89,8 +96,7 @@ logger = logging.getLogger(__name__)
|
|||
DEFAULT_SERENA_CONTEXT = "ide-assistant"
|
||||
|
||||
|
||||
def _make_headroom_client(proxy_url: str) -> MCPClient:
|
||||
spec = build_headroom_spec(proxy_url)
|
||||
def _client_for(spec: Any) -> MCPClient:
|
||||
params = StdioServerParameters(
|
||||
command=spec.command,
|
||||
args=list(spec.args),
|
||||
|
|
@ -100,14 +106,16 @@ def _make_headroom_client(proxy_url: str) -> MCPClient:
|
|||
return MCPClient(partial(stdio_client, params))
|
||||
|
||||
|
||||
def _make_headroom_client(proxy_url: str) -> MCPClient:
|
||||
return _client_for(build_headroom_spec(proxy_url))
|
||||
|
||||
|
||||
def _make_tokensave_client() -> MCPClient:
|
||||
return _client_for(build_tokensave_spec())
|
||||
|
||||
|
||||
def _make_serena_client(context: str) -> MCPClient:
|
||||
spec = build_serena_spec(context)
|
||||
params = StdioServerParameters(
|
||||
command=spec.command,
|
||||
args=list(spec.args),
|
||||
env=dict(spec.env) if spec.env else None,
|
||||
)
|
||||
return MCPClient(partial(stdio_client, params))
|
||||
return _client_for(build_serena_spec(context))
|
||||
|
||||
|
||||
@dataclass
|
||||
|
|
@ -120,8 +128,12 @@ class HeadroomBundle:
|
|||
(``http://127.0.0.1:8787``).
|
||||
serena_context: Serena context label. Default ``"ide-assistant"``.
|
||||
enable_headroom_mcp: Include the Headroom MCP server. Default True.
|
||||
enable_serena_mcp: Include the Serena MCP server. Default True.
|
||||
Disabling skips the ``uvx`` first-launch download entirely.
|
||||
enable_tokensave_mcp: Include the tokensave MCP server — the primary
|
||||
coding-task compressor. Default True. Requires the ``tokensave``
|
||||
binary on PATH (``tokensave serve``).
|
||||
enable_serena_mcp: Include the Serena MCP server — the backup
|
||||
coding-task compressor. Default False (tokensave is primary).
|
||||
Enabling adds the ``uvx`` first-launch download.
|
||||
enable_hooks: Include :class:`HeadroomHookProvider` for in-place
|
||||
tool-output compression (the RTK-equivalent for Strands).
|
||||
Default True.
|
||||
|
|
@ -138,7 +150,10 @@ class HeadroomBundle:
|
|||
proxy_url: str = DEFAULT_PROXY_URL
|
||||
serena_context: str = DEFAULT_SERENA_CONTEXT
|
||||
enable_headroom_mcp: bool = True
|
||||
enable_serena_mcp: bool = True
|
||||
# tokensave is the primary coding-task compressor; Serena is the backup
|
||||
# and stays off unless explicitly enabled.
|
||||
enable_tokensave_mcp: bool = True
|
||||
enable_serena_mcp: bool = False
|
||||
# The proxy is the single source of truth for compression — it sees
|
||||
# the full message list, owns CompressionPolicy, owns PrefixCacheTracker,
|
||||
# and places `cache_control` breakpoints. The in-process hook
|
||||
|
|
@ -151,6 +166,7 @@ class HeadroomBundle:
|
|||
config: HeadroomConfig | None = None
|
||||
|
||||
_headroom_mcp: MCPClient | None = field(default=None, init=False, repr=False, compare=False)
|
||||
_tokensave_mcp: MCPClient | None = field(default=None, init=False, repr=False, compare=False)
|
||||
_serena_mcp: MCPClient | None = field(default=None, init=False, repr=False, compare=False)
|
||||
_hook: HeadroomHookProvider | None = field(default=None, init=False, repr=False, compare=False)
|
||||
|
||||
|
|
@ -161,10 +177,13 @@ class HeadroomBundle:
|
|||
"HeadroomBundle: Headroom MCP client constructed (proxy_url=%s)",
|
||||
self.proxy_url,
|
||||
)
|
||||
if self.enable_tokensave_mcp:
|
||||
self._tokensave_mcp = _make_tokensave_client()
|
||||
logger.info("HeadroomBundle: tokensave MCP client constructed (primary)")
|
||||
if self.enable_serena_mcp:
|
||||
self._serena_mcp = _make_serena_client(self.serena_context)
|
||||
logger.info(
|
||||
"HeadroomBundle: Serena MCP client constructed (context=%s)",
|
||||
"HeadroomBundle: Serena MCP client constructed (backup, context=%s)",
|
||||
self.serena_context,
|
||||
)
|
||||
if self.enable_hooks:
|
||||
|
|
@ -181,6 +200,8 @@ class HeadroomBundle:
|
|||
out: list[Any] = []
|
||||
if self._headroom_mcp is not None:
|
||||
out.append(self._headroom_mcp)
|
||||
if self._tokensave_mcp is not None:
|
||||
out.append(self._tokensave_mcp)
|
||||
if self._serena_mcp is not None:
|
||||
out.append(self._serena_mcp)
|
||||
return out
|
||||
|
|
@ -195,6 +216,11 @@ class HeadroomBundle:
|
|||
"""Direct handle to the Headroom MCPClient (for advanced callers)."""
|
||||
return self._headroom_mcp
|
||||
|
||||
@property
|
||||
def tokensave_mcp(self) -> MCPClient | None:
|
||||
"""Direct handle to the tokensave MCPClient (for advanced callers)."""
|
||||
return self._tokensave_mcp
|
||||
|
||||
@property
|
||||
def serena_mcp(self) -> MCPClient | None:
|
||||
"""Direct handle to the Serena MCPClient (for advanced callers)."""
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from .install import (
|
|||
DEFAULT_PROXY_URL,
|
||||
build_headroom_spec,
|
||||
build_serena_spec,
|
||||
build_tokensave_spec,
|
||||
get_all_registrars,
|
||||
install_everywhere,
|
||||
)
|
||||
|
|
@ -38,6 +39,7 @@ __all__ = [
|
|||
"any_succeeded",
|
||||
"build_headroom_spec",
|
||||
"build_serena_spec",
|
||||
"build_tokensave_spec",
|
||||
"format_result",
|
||||
"format_results",
|
||||
"get_all_registrars",
|
||||
|
|
|
|||
|
|
@ -68,6 +68,22 @@ def build_serena_spec(context: str) -> ServerSpec:
|
|||
)
|
||||
|
||||
|
||||
def build_tokensave_spec(binary: str = "tokensave") -> ServerSpec:
|
||||
"""Construct the canonical tokensave MCP server spec.
|
||||
|
||||
tokensave (https://github.com/aovestdipaperino/tokensave) is the primary
|
||||
coding-task compressor — a local semantic code-graph server launched as
|
||||
``tokensave serve`` over stdio. ``binary`` is the command the agent runs;
|
||||
pass an absolute path when tokensave was fetched to ``~/.local/bin`` and
|
||||
is not on the agent's PATH, or leave the default when it is on PATH.
|
||||
"""
|
||||
return ServerSpec(
|
||||
name="tokensave",
|
||||
command=binary,
|
||||
args=("serve",),
|
||||
)
|
||||
|
||||
|
||||
def install_everywhere(
|
||||
proxy_url: str = DEFAULT_PROXY_URL,
|
||||
*,
|
||||
|
|
|
|||
20
tests/test_cli/conftest.py
Normal file
20
tests/test_cli/conftest.py
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
"""Shared fixtures for the CLI test suite.
|
||||
|
||||
tokensave is now the primary coding-task compressor, so a default
|
||||
``headroom wrap`` tries to fetch the tokensave release binary. Force offline
|
||||
across CLI tests so a missing binary resolves to ``None`` (→ Serena fallback)
|
||||
instead of reaching out to GitHub releases. Tests that exercise the
|
||||
tokensave-present path patch ``_ensure_tokensave_binary`` / ``ensure_tokensave``
|
||||
directly and are unaffected by this guard. This env only gates the new
|
||||
tokensave installer (``headroom.graph.tokensave_installer``); rtk and
|
||||
codebase-memory-mcp installers do not read it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _tokensave_offline(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("HEADROOM_BINARIES_OFFLINE", "1")
|
||||
112
tests/test_cli/test_tokensave_helpers.py
Normal file
112
tests/test_cli/test_tokensave_helpers.py
Normal file
|
|
@ -0,0 +1,112 @@
|
|||
"""Coverage for the tokensave binary-resolution and indexing helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.cli import wrap as wrap_cli
|
||||
from headroom.graph import tokensave_installer as ts
|
||||
|
||||
_FAKE_BIN = Path("/usr/local/bin/tokensave")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _ensure_tokensave_binary
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_ensure_binary_returns_existing_without_fetch(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(ts, "get_tokensave_path", lambda: _FAKE_BIN)
|
||||
|
||||
def _should_not_run(*a, **k):
|
||||
raise AssertionError("must not download when binary already present")
|
||||
|
||||
monkeypatch.setattr(ts, "ensure_tokensave", _should_not_run)
|
||||
assert wrap_cli._ensure_tokensave_binary() == _FAKE_BIN
|
||||
|
||||
|
||||
def test_ensure_binary_fetches_when_absent(
|
||||
monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
|
||||
) -> None:
|
||||
monkeypatch.setattr(ts, "get_tokensave_path", lambda: None)
|
||||
monkeypatch.setattr(ts, "ensure_tokensave", lambda: _FAKE_BIN)
|
||||
assert wrap_cli._ensure_tokensave_binary() == _FAKE_BIN
|
||||
assert "installed at" in capsys.readouterr().out
|
||||
|
||||
|
||||
def test_ensure_binary_none_prints_fallback(
|
||||
monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
|
||||
) -> None:
|
||||
monkeypatch.setattr(ts, "get_tokensave_path", lambda: None)
|
||||
monkeypatch.setattr(ts, "ensure_tokensave", lambda: None)
|
||||
assert wrap_cli._ensure_tokensave_binary() is None
|
||||
assert "falling back to Serena" in capsys.readouterr().out
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _index_tokensave_project
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _patch_run(monkeypatch: pytest.MonkeyPatch, result):
|
||||
calls: list[list[str]] = []
|
||||
|
||||
def fake_run(cmd, **kwargs):
|
||||
calls.append(cmd)
|
||||
if isinstance(result, Exception):
|
||||
raise result
|
||||
return result
|
||||
|
||||
monkeypatch.setattr(wrap_cli.subprocess, "run", fake_run)
|
||||
return calls
|
||||
|
||||
|
||||
def test_index_runs_init_when_no_db(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path, capsys: pytest.CaptureFixture[str]
|
||||
) -> None:
|
||||
monkeypatch.chdir(tmp_path)
|
||||
calls = _patch_run(monkeypatch, SimpleNamespace(returncode=0, stdout="", stderr=""))
|
||||
wrap_cli._index_tokensave_project(_FAKE_BIN)
|
||||
assert calls == [[str(_FAKE_BIN), "init"]]
|
||||
assert "Code graph: indexed (tokensave)" in capsys.readouterr().out
|
||||
|
||||
|
||||
def test_index_runs_sync_when_db_exists(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
||||
(tmp_path / ".tokensave").mkdir()
|
||||
monkeypatch.chdir(tmp_path)
|
||||
calls = _patch_run(monkeypatch, SimpleNamespace(returncode=0, stdout="", stderr=""))
|
||||
wrap_cli._index_tokensave_project(_FAKE_BIN)
|
||||
assert calls == [[str(_FAKE_BIN), "sync"]]
|
||||
|
||||
|
||||
def test_index_nonzero_is_nonfatal(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path, capsys: pytest.CaptureFixture[str]
|
||||
) -> None:
|
||||
monkeypatch.chdir(tmp_path)
|
||||
_patch_run(monkeypatch, SimpleNamespace(returncode=1, stdout="", stderr="boom"))
|
||||
wrap_cli._index_tokensave_project(_FAKE_BIN, verbose=True)
|
||||
assert "init failed" in capsys.readouterr().out
|
||||
|
||||
|
||||
def test_index_timeout_is_nonfatal(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path, capsys: pytest.CaptureFixture[str]
|
||||
) -> None:
|
||||
import subprocess
|
||||
|
||||
monkeypatch.chdir(tmp_path)
|
||||
_patch_run(monkeypatch, subprocess.TimeoutExpired(cmd="tokensave", timeout=60))
|
||||
wrap_cli._index_tokensave_project(_FAKE_BIN)
|
||||
assert "timed out" in capsys.readouterr().out
|
||||
|
||||
|
||||
def test_index_exception_is_nonfatal(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path, capsys: pytest.CaptureFixture[str]
|
||||
) -> None:
|
||||
monkeypatch.chdir(tmp_path)
|
||||
_patch_run(monkeypatch, FileNotFoundError("no binary"))
|
||||
# Must not raise even when the binary is missing.
|
||||
wrap_cli._index_tokensave_project(_FAKE_BIN, verbose=True)
|
||||
assert "indexing skipped" in capsys.readouterr().out
|
||||
227
tests/test_cli/test_tokensave_setup.py
Normal file
227
tests/test_cli/test_tokensave_setup.py
Normal file
|
|
@ -0,0 +1,227 @@
|
|||
"""tokensave is the primary coding-task compressor; Serena is the backup.
|
||||
|
||||
These tests pin the wrap-time policy in :func:`_setup_coding_compressor` and
|
||||
the tokensave register/disable/migrate helpers, mirroring the Serena tests.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.cli import wrap as wrap_cli
|
||||
from headroom.mcp_registry import build_tokensave_spec
|
||||
from headroom.mcp_registry.base import RegisterResult, RegisterStatus, ServerSpec
|
||||
from headroom.mcp_registry.ledger import headroom_installed_matching, record_install
|
||||
|
||||
_FAKE_BIN = Path("/usr/local/bin/tokensave")
|
||||
|
||||
|
||||
def _equivalent(a: ServerSpec, b: ServerSpec) -> bool:
|
||||
return (a.command, tuple(a.args), dict(a.env)) == (b.command, tuple(b.args), dict(b.env))
|
||||
|
||||
|
||||
class _FakeRegistrar:
|
||||
"""Registrar mirroring real ``register_server`` overwrite semantics."""
|
||||
|
||||
def __init__(self, name: str = "claude", *, detected: bool = True, server=None):
|
||||
self.name = name
|
||||
self.display_name = name.capitalize()
|
||||
self._detected = detected
|
||||
self._server = server
|
||||
self.force_calls: list[bool] = []
|
||||
self.unregistered: list[str] = []
|
||||
|
||||
def detect(self) -> bool:
|
||||
return self._detected
|
||||
|
||||
def get_server(self, server_name: str):
|
||||
return self._server if server_name == "tokensave" else None
|
||||
|
||||
def register_server(self, spec: ServerSpec, *, force: bool = False) -> RegisterResult:
|
||||
self.force_calls.append(force)
|
||||
if self._server is not None and not _equivalent(self._server, spec) and not force:
|
||||
return RegisterResult(RegisterStatus.MISMATCH, "differs")
|
||||
self._server = spec
|
||||
return RegisterResult(RegisterStatus.REGISTERED, "ok")
|
||||
|
||||
def unregister_server(self, server_name: str) -> bool:
|
||||
self.unregistered.append(server_name)
|
||||
self._server = None
|
||||
return True
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _workspace(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
||||
monkeypatch.setenv("HEADROOM_WORKSPACE_DIR", str(tmp_path / ".headroom"))
|
||||
# Never touch the network or run the real binary during these unit tests.
|
||||
monkeypatch.setattr(wrap_cli, "_index_tokensave_project", lambda *a, **k: None)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _setup_tokensave_mcp
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_setup_registers_and_records_when_binary_available(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(wrap_cli, "_ensure_tokensave_binary", lambda verbose=False: _FAKE_BIN)
|
||||
registrar = _FakeRegistrar()
|
||||
|
||||
assert wrap_cli._setup_tokensave_mcp(registrar) is True
|
||||
assert registrar._server is not None
|
||||
assert registrar._server.name == "tokensave"
|
||||
assert registrar._server.command == str(_FAKE_BIN)
|
||||
# Ledger now proves Headroom owns the entry.
|
||||
assert headroom_installed_matching("claude", registrar.get_server("tokensave"))
|
||||
|
||||
|
||||
def test_setup_returns_false_when_binary_unavailable(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(wrap_cli, "_ensure_tokensave_binary", lambda verbose=False: None)
|
||||
registrar = _FakeRegistrar()
|
||||
|
||||
assert wrap_cli._setup_tokensave_mcp(registrar) is False
|
||||
assert registrar._server is None # nothing registered
|
||||
|
||||
|
||||
def test_setup_skips_when_agent_not_detected(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
sentinel = {"called": False}
|
||||
|
||||
def _should_not_run(verbose=False):
|
||||
sentinel["called"] = True
|
||||
return _FAKE_BIN
|
||||
|
||||
monkeypatch.setattr(wrap_cli, "_ensure_tokensave_binary", _should_not_run)
|
||||
registrar = _FakeRegistrar(detected=False)
|
||||
|
||||
assert wrap_cli._setup_tokensave_mcp(registrar) is False
|
||||
assert sentinel["called"] is False # never even fetched the binary
|
||||
|
||||
|
||||
def test_setup_migrates_stale_headroom_entry(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(wrap_cli, "_ensure_tokensave_binary", lambda verbose=False: _FAKE_BIN)
|
||||
# A stale Headroom-installed entry (different binary path) is on disk.
|
||||
stale = build_tokensave_spec("/old/path/tokensave")
|
||||
record_install("claude", stale)
|
||||
registrar = _FakeRegistrar(server=stale)
|
||||
|
||||
assert wrap_cli._setup_tokensave_mcp(registrar) is True
|
||||
# Force-updated to the current spec.
|
||||
assert registrar.force_calls[-1] is True
|
||||
assert registrar._server.command == str(_FAKE_BIN)
|
||||
|
||||
|
||||
def test_setup_preserves_user_managed_mismatch(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(wrap_cli, "_ensure_tokensave_binary", lambda verbose=False: _FAKE_BIN)
|
||||
# User-managed entry (NOT in ledger) that differs from our spec.
|
||||
user = ServerSpec(name="tokensave", command="/custom/tokensave", args=("serve",))
|
||||
registrar = _FakeRegistrar(server=user)
|
||||
|
||||
wrap_cli._setup_tokensave_mcp(registrar)
|
||||
# Never force-overwrote a user-managed entry.
|
||||
assert True not in registrar.force_calls
|
||||
assert registrar._server.command == "/custom/tokensave"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _disable_tokensave_mcp
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_disable_removes_headroom_installed(capsys: pytest.CaptureFixture[str]) -> None:
|
||||
spec = build_tokensave_spec(str(_FAKE_BIN))
|
||||
record_install("claude", spec)
|
||||
registrar = _FakeRegistrar(server=spec)
|
||||
|
||||
wrap_cli._disable_tokensave_mcp(registrar, verbose=True)
|
||||
|
||||
assert registrar.unregistered == ["tokensave"]
|
||||
assert "Removed previously-installed tokensave MCP" in capsys.readouterr().out
|
||||
|
||||
|
||||
def test_disable_preserves_user_managed(capsys: pytest.CaptureFixture[str]) -> None:
|
||||
user = ServerSpec(name="tokensave", command="/custom/tokensave")
|
||||
registrar = _FakeRegistrar(server=user)
|
||||
|
||||
wrap_cli._disable_tokensave_mcp(registrar, verbose=True)
|
||||
|
||||
assert registrar.unregistered == []
|
||||
assert "user-managed" in capsys.readouterr().out
|
||||
|
||||
|
||||
def test_disable_noop_when_absent(capsys: pytest.CaptureFixture[str]) -> None:
|
||||
registrar = _FakeRegistrar(server=None)
|
||||
wrap_cli._disable_tokensave_mcp(registrar, verbose=True)
|
||||
assert registrar.unregistered == []
|
||||
assert "Skipping tokensave MCP" in capsys.readouterr().out
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _setup_coding_compressor — primary/backup policy
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _spy_compressor(monkeypatch: pytest.MonkeyPatch, *, tokensave_ok: bool) -> dict:
|
||||
calls: dict[str, object] = {"serena_setup": False, "serena_disabled": None, "tokensave": None}
|
||||
|
||||
def fake_setup_tokensave(reg, *, verbose=False, force=False):
|
||||
calls["tokensave"] = "setup"
|
||||
return tokensave_ok
|
||||
|
||||
def fake_disable_tokensave(reg, *, verbose=False):
|
||||
calls["tokensave"] = "disabled"
|
||||
|
||||
def fake_setup_serena(reg, *, context, verbose=False, force=False):
|
||||
calls["serena_setup"] = True
|
||||
|
||||
def fake_disable_serena(reg, *, verbose=False, reason="--no-serena"):
|
||||
calls["serena_disabled"] = reason
|
||||
|
||||
monkeypatch.setattr(wrap_cli, "_setup_tokensave_mcp", fake_setup_tokensave)
|
||||
monkeypatch.setattr(wrap_cli, "_disable_tokensave_mcp", fake_disable_tokensave)
|
||||
monkeypatch.setattr(wrap_cli, "_setup_serena_mcp", fake_setup_serena)
|
||||
monkeypatch.setattr(wrap_cli, "_disable_serena_mcp", fake_disable_serena)
|
||||
return calls
|
||||
|
||||
|
||||
def test_policy_tokensave_primary_disables_serena(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
calls = _spy_compressor(monkeypatch, tokensave_ok=True)
|
||||
wrap_cli._setup_coding_compressor(_FakeRegistrar(), serena_context="claude-code")
|
||||
assert calls["tokensave"] == "setup"
|
||||
assert calls["serena_setup"] is False
|
||||
assert calls["serena_disabled"] == "tokensave is now the primary code-graph compressor"
|
||||
|
||||
|
||||
def test_policy_serena_fallback_when_tokensave_unavailable(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
calls = _spy_compressor(monkeypatch, tokensave_ok=False)
|
||||
wrap_cli._setup_coding_compressor(_FakeRegistrar(), serena_context="claude-code")
|
||||
assert calls["serena_setup"] is True
|
||||
|
||||
|
||||
def test_policy_force_serena_even_when_tokensave_ok(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
calls = _spy_compressor(monkeypatch, tokensave_ok=True)
|
||||
wrap_cli._setup_coding_compressor(_FakeRegistrar(), serena_context="claude-code", serena=True)
|
||||
assert calls["serena_setup"] is True
|
||||
|
||||
|
||||
def test_policy_no_serena_suppresses_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
calls = _spy_compressor(monkeypatch, tokensave_ok=False)
|
||||
wrap_cli._setup_coding_compressor(
|
||||
_FakeRegistrar(), serena_context="claude-code", no_serena=True
|
||||
)
|
||||
assert calls["serena_setup"] is False
|
||||
assert calls["serena_disabled"] == "--no-serena"
|
||||
|
||||
|
||||
def test_policy_no_tokensave_disables_and_falls_back(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
calls = _spy_compressor(monkeypatch, tokensave_ok=True)
|
||||
wrap_cli._setup_coding_compressor(
|
||||
_FakeRegistrar(), serena_context="claude-code", no_tokensave=True
|
||||
)
|
||||
assert calls["tokensave"] == "disabled"
|
||||
# tokensave disabled → treated as unavailable → Serena fallback registers.
|
||||
assert calls["serena_setup"] is True
|
||||
284
tests/test_graph_tokensave.py
Normal file
284
tests/test_graph_tokensave.py
Normal file
|
|
@ -0,0 +1,284 @@
|
|||
"""Tests for the tokensave release-binary installer."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import tarfile
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.graph import tokensave_installer as ts
|
||||
|
||||
|
||||
def _tar_archive(member_name: str = ts.TOKENSAVE_BIN_NAME) -> bytes:
|
||||
payload = io.BytesIO()
|
||||
with tarfile.open(fileobj=payload, mode="w:gz") as tar:
|
||||
data = b"#!/bin/sh\necho version\n"
|
||||
info = tarfile.TarInfo(name=member_name)
|
||||
info.size = len(data)
|
||||
tar.addfile(info, io.BytesIO(data))
|
||||
return payload.getvalue()
|
||||
|
||||
|
||||
def _zip_archive(member_name: str = "tokensave.exe") -> bytes:
|
||||
payload = io.BytesIO()
|
||||
with zipfile.ZipFile(payload, "w") as zf:
|
||||
zf.writestr(member_name, b"binary")
|
||||
return payload.getvalue()
|
||||
|
||||
|
||||
class FakeResponse:
|
||||
def __init__(self, data: bytes) -> None:
|
||||
self._data = data
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb) -> None:
|
||||
return None
|
||||
|
||||
def read(self) -> bytes:
|
||||
return self._data
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("system", "machine", "expected"),
|
||||
[
|
||||
("darwin", "arm64", ("tokensave-v9-aarch64-macos.tar.gz", "tar.gz")),
|
||||
("linux", "aarch64", ("tokensave-v9-aarch64-linux.tar.gz", "tar.gz")),
|
||||
("linux", "arm64", ("tokensave-v9-aarch64-linux.tar.gz", "tar.gz")),
|
||||
("linux", "x86_64", ("tokensave-v9-x86_64-linux.tar.gz", "tar.gz")),
|
||||
("windows", "amd64", ("tokensave-v9-x86_64-windows.zip", "zip")),
|
||||
("windows", "arm64", ("tokensave-v9-aarch64-windows.zip", "zip")),
|
||||
],
|
||||
)
|
||||
def test_detect_asset_variants(monkeypatch, system, machine, expected) -> None:
|
||||
monkeypatch.setattr(ts.platform, "system", lambda: system)
|
||||
monkeypatch.setattr(ts.platform, "machine", lambda: machine)
|
||||
assert ts._detect_asset("v9") == expected
|
||||
|
||||
|
||||
def test_detect_asset_returns_none_for_intel_mac_and_unknown(monkeypatch) -> None:
|
||||
monkeypatch.setattr(ts.platform, "system", lambda: "darwin")
|
||||
monkeypatch.setattr(ts.platform, "machine", lambda: "x86_64")
|
||||
assert ts._detect_asset("v9") is None # no x86_64-macos asset is published
|
||||
|
||||
monkeypatch.setattr(ts.platform, "system", lambda: "solaris")
|
||||
monkeypatch.setattr(ts.platform, "machine", lambda: "sparc")
|
||||
assert ts._detect_asset("v9") is None
|
||||
|
||||
|
||||
def test_get_tokensave_path_prefers_path_then_install_dir(monkeypatch, tmp_path: Path) -> None:
|
||||
on_path = tmp_path / "on-path"
|
||||
installed = tmp_path / ts.TOKENSAVE_BIN_NAME
|
||||
installed.write_text("bin")
|
||||
monkeypatch.setattr(ts, "TOKENSAVE_BIN_DIR", tmp_path)
|
||||
monkeypatch.setattr("shutil.which", lambda name: str(on_path))
|
||||
assert ts.get_tokensave_path() == on_path
|
||||
|
||||
monkeypatch.setattr("shutil.which", lambda name: None)
|
||||
assert ts.get_tokensave_path() == installed
|
||||
|
||||
installed.unlink()
|
||||
assert ts.get_tokensave_path() is None
|
||||
|
||||
|
||||
def test_ensure_offline_returns_none_when_absent(monkeypatch, tmp_path: Path) -> None:
|
||||
monkeypatch.setattr(ts, "TOKENSAVE_BIN_DIR", tmp_path)
|
||||
monkeypatch.setattr("shutil.which", lambda name: None)
|
||||
monkeypatch.setenv("HEADROOM_BINARIES_OFFLINE", "1")
|
||||
|
||||
def _boom(*a, **k):
|
||||
raise AssertionError("download must not run when offline")
|
||||
|
||||
monkeypatch.setattr(ts, "download_tokensave", _boom)
|
||||
assert ts.ensure_tokensave() is None
|
||||
|
||||
|
||||
def test_ensure_returns_existing_without_download(monkeypatch, tmp_path: Path) -> None:
|
||||
existing = tmp_path / ts.TOKENSAVE_BIN_NAME
|
||||
existing.write_text("bin")
|
||||
monkeypatch.setattr(ts, "get_tokensave_path", lambda: existing)
|
||||
|
||||
def _boom(*a, **k):
|
||||
raise AssertionError("download must not run when binary present")
|
||||
|
||||
monkeypatch.setattr(ts, "download_tokensave", _boom)
|
||||
assert ts.ensure_tokensave() == existing
|
||||
|
||||
|
||||
def test_ensure_returns_none_on_unsupported_platform(monkeypatch, tmp_path: Path) -> None:
|
||||
monkeypatch.setattr(ts, "get_tokensave_path", lambda: None)
|
||||
monkeypatch.delenv("HEADROOM_BINARIES_OFFLINE", raising=False)
|
||||
monkeypatch.setattr(ts.platform, "system", lambda: "darwin")
|
||||
monkeypatch.setattr(ts.platform, "machine", lambda: "x86_64") # no asset
|
||||
assert ts.ensure_tokensave() is None
|
||||
|
||||
|
||||
def test_download_tokensave_tarball(monkeypatch, tmp_path: Path) -> None:
|
||||
monkeypatch.setattr(ts, "TOKENSAVE_BIN_DIR", tmp_path)
|
||||
monkeypatch.setattr(ts.platform, "system", lambda: "linux")
|
||||
monkeypatch.setattr(ts.platform, "machine", lambda: "x86_64")
|
||||
# Synthetic archive bytes won't match the pinned digest; this test covers
|
||||
# extraction, not integrity, so opt out of verification explicitly.
|
||||
monkeypatch.setenv("HEADROOM_TOKENSAVE_ALLOW_UNVERIFIED", "1")
|
||||
monkeypatch.setattr(ts, "urlopen", lambda url, timeout=60: FakeResponse(_tar_archive()))
|
||||
monkeypatch.setattr(
|
||||
"subprocess.run", lambda *a, **k: SimpleNamespace(returncode=0, stdout="tokensave 6\n")
|
||||
)
|
||||
path = ts.download_tokensave(version="v0.0.0-test")
|
||||
assert path == tmp_path / ts.TOKENSAVE_BIN_NAME
|
||||
assert path.exists()
|
||||
|
||||
|
||||
def test_download_tokensave_zip_windows(monkeypatch, tmp_path: Path) -> None:
|
||||
monkeypatch.setattr(ts, "TOKENSAVE_BIN_DIR", tmp_path)
|
||||
monkeypatch.setattr(ts.platform, "system", lambda: "windows")
|
||||
monkeypatch.setattr(ts.platform, "machine", lambda: "amd64")
|
||||
monkeypatch.setenv("HEADROOM_TOKENSAVE_ALLOW_UNVERIFIED", "1")
|
||||
monkeypatch.setattr(ts, "urlopen", lambda url, timeout=60: FakeResponse(_zip_archive()))
|
||||
monkeypatch.setattr(
|
||||
"subprocess.run", lambda *a, **k: SimpleNamespace(returncode=0, stdout="tokensave 6\n")
|
||||
)
|
||||
path = ts.download_tokensave(version="v0.0.0-test")
|
||||
assert path == tmp_path / "tokensave.exe"
|
||||
assert path.exists()
|
||||
|
||||
|
||||
def test_download_raises_for_unsupported_platform(monkeypatch, tmp_path: Path) -> None:
|
||||
monkeypatch.setattr(ts, "TOKENSAVE_BIN_DIR", tmp_path)
|
||||
monkeypatch.setattr(ts.platform, "system", lambda: "darwin")
|
||||
monkeypatch.setattr(ts.platform, "machine", lambda: "x86_64")
|
||||
with pytest.raises(RuntimeError, match="no prebuilt tokensave asset"):
|
||||
ts.download_tokensave(version="v7.0.0")
|
||||
|
||||
|
||||
def test_download_wraps_network_failure(monkeypatch, tmp_path: Path) -> None:
|
||||
monkeypatch.setattr(ts, "TOKENSAVE_BIN_DIR", tmp_path)
|
||||
monkeypatch.setattr(ts.platform, "system", lambda: "linux")
|
||||
monkeypatch.setattr(ts.platform, "machine", lambda: "x86_64")
|
||||
|
||||
def _boom(url, timeout=60):
|
||||
raise OSError("connection refused")
|
||||
|
||||
monkeypatch.setattr(ts, "urlopen", _boom)
|
||||
with pytest.raises(RuntimeError, match="Failed to download tokensave"):
|
||||
ts.download_tokensave(version="v7.0.0")
|
||||
|
||||
|
||||
def test_download_raises_when_binary_missing_from_tarball(monkeypatch, tmp_path: Path) -> None:
|
||||
monkeypatch.setattr(ts, "TOKENSAVE_BIN_DIR", tmp_path)
|
||||
monkeypatch.setattr(ts.platform, "system", lambda: "linux")
|
||||
monkeypatch.setattr(ts.platform, "machine", lambda: "x86_64")
|
||||
# Archive contains an unrelated member, not the tokensave binary.
|
||||
monkeypatch.setenv("HEADROOM_TOKENSAVE_ALLOW_UNVERIFIED", "1")
|
||||
monkeypatch.setattr(
|
||||
ts, "urlopen", lambda url, timeout=60: FakeResponse(_tar_archive("README.md"))
|
||||
)
|
||||
with pytest.raises(RuntimeError, match="binary not found in archive"):
|
||||
ts.download_tokensave(version="v0.0.0-test")
|
||||
|
||||
|
||||
def test_download_raises_when_binary_missing_from_zip(monkeypatch, tmp_path: Path) -> None:
|
||||
monkeypatch.setattr(ts, "TOKENSAVE_BIN_DIR", tmp_path)
|
||||
monkeypatch.setattr(ts.platform, "system", lambda: "windows")
|
||||
monkeypatch.setattr(ts.platform, "machine", lambda: "amd64")
|
||||
monkeypatch.setenv("HEADROOM_TOKENSAVE_ALLOW_UNVERIFIED", "1")
|
||||
monkeypatch.setattr(
|
||||
ts, "urlopen", lambda url, timeout=60: FakeResponse(_zip_archive("notes.txt"))
|
||||
)
|
||||
with pytest.raises(RuntimeError, match="binary not found in archive"):
|
||||
ts.download_tokensave(version="v0.0.0-test")
|
||||
|
||||
|
||||
def test_download_tolerates_failed_version_check(monkeypatch, tmp_path: Path) -> None:
|
||||
monkeypatch.setattr(ts, "TOKENSAVE_BIN_DIR", tmp_path)
|
||||
monkeypatch.setattr(ts.platform, "system", lambda: "linux")
|
||||
monkeypatch.setattr(ts.platform, "machine", lambda: "x86_64")
|
||||
monkeypatch.setenv("HEADROOM_TOKENSAVE_ALLOW_UNVERIFIED", "1")
|
||||
monkeypatch.setattr(ts, "urlopen", lambda url, timeout=60: FakeResponse(_tar_archive()))
|
||||
# Non-zero return code and a raising probe must both be non-fatal.
|
||||
monkeypatch.setattr(
|
||||
"subprocess.run", lambda *a, **k: SimpleNamespace(returncode=1, stdout="", stderr="x")
|
||||
)
|
||||
assert ts.download_tokensave(version="v0.0.0-test") == tmp_path / ts.TOKENSAVE_BIN_NAME
|
||||
|
||||
monkeypatch.setattr(
|
||||
"subprocess.run", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("probe boom"))
|
||||
)
|
||||
assert ts.download_tokensave(version="v0.0.0-test") == tmp_path / ts.TOKENSAVE_BIN_NAME
|
||||
|
||||
|
||||
def test_verify_asset_digest_accepts_matching_hash(monkeypatch) -> None:
|
||||
import hashlib
|
||||
|
||||
data = b"some-release-bytes"
|
||||
digest = hashlib.sha256(data).hexdigest()
|
||||
monkeypatch.setattr(ts, "TOKENSAVE_ASSET_DIGESTS", {"asset.tar.gz": digest})
|
||||
# No exception => verification passed.
|
||||
ts._verify_asset_digest("asset.tar.gz", data)
|
||||
|
||||
|
||||
def test_verify_asset_digest_rejects_mismatch(monkeypatch) -> None:
|
||||
monkeypatch.setattr(ts, "TOKENSAVE_ASSET_DIGESTS", {"asset.tar.gz": "00" * 32})
|
||||
with pytest.raises(RuntimeError, match="failed integrity check"):
|
||||
ts._verify_asset_digest("asset.tar.gz", b"tampered")
|
||||
|
||||
|
||||
def test_verify_asset_digest_refuses_unpinned_without_optout(monkeypatch) -> None:
|
||||
monkeypatch.setattr(ts, "TOKENSAVE_ASSET_DIGESTS", {})
|
||||
monkeypatch.delenv("HEADROOM_TOKENSAVE_ALLOW_UNVERIFIED", raising=False)
|
||||
with pytest.raises(RuntimeError, match="no pinned SHA-256 digest"):
|
||||
ts._verify_asset_digest("unknown.tar.gz", b"bytes")
|
||||
|
||||
|
||||
def test_verify_asset_digest_allows_unpinned_with_optout(monkeypatch) -> None:
|
||||
monkeypatch.setattr(ts, "TOKENSAVE_ASSET_DIGESTS", {})
|
||||
monkeypatch.setenv("HEADROOM_TOKENSAVE_ALLOW_UNVERIFIED", "1")
|
||||
ts._verify_asset_digest("unknown.tar.gz", b"bytes") # no exception
|
||||
|
||||
|
||||
def test_download_aborts_on_digest_mismatch(monkeypatch, tmp_path: Path) -> None:
|
||||
monkeypatch.setattr(ts, "TOKENSAVE_BIN_DIR", tmp_path)
|
||||
monkeypatch.setattr(ts.platform, "system", lambda: "linux")
|
||||
monkeypatch.setattr(ts.platform, "machine", lambda: "x86_64")
|
||||
monkeypatch.delenv("HEADROOM_TOKENSAVE_ALLOW_UNVERIFIED", raising=False)
|
||||
# Pin a digest that the synthetic archive cannot match.
|
||||
monkeypatch.setattr(
|
||||
ts, "TOKENSAVE_ASSET_DIGESTS", {"tokensave-v7.0.0-x86_64-linux.tar.gz": "00" * 32}
|
||||
)
|
||||
monkeypatch.setattr(ts, "urlopen", lambda url, timeout=60: FakeResponse(_tar_archive()))
|
||||
with pytest.raises(RuntimeError, match="failed integrity check"):
|
||||
ts.download_tokensave(version="v7.0.0")
|
||||
# The unverified binary must not have been written.
|
||||
assert not (tmp_path / ts.TOKENSAVE_BIN_NAME).exists()
|
||||
|
||||
|
||||
def test_download_honors_invalid_url_scheme(monkeypatch, tmp_path: Path) -> None:
|
||||
monkeypatch.setattr(ts, "TOKENSAVE_BIN_DIR", tmp_path)
|
||||
monkeypatch.setattr(ts.platform, "system", lambda: "linux")
|
||||
monkeypatch.setattr(ts.platform, "machine", lambda: "x86_64")
|
||||
monkeypatch.setattr(ts, "GITHUB_RELEASE_URL", "ftp://example.test/releases")
|
||||
with pytest.raises(RuntimeError, match="Failed to download tokensave"):
|
||||
ts.download_tokensave(version="v7.0.0")
|
||||
|
||||
|
||||
def test_ensure_returns_none_when_download_fails(monkeypatch, tmp_path: Path) -> None:
|
||||
monkeypatch.setattr(ts, "get_tokensave_path", lambda: None)
|
||||
monkeypatch.delenv("HEADROOM_BINARIES_OFFLINE", raising=False)
|
||||
|
||||
def _raise(version=None):
|
||||
raise RuntimeError("download failed")
|
||||
|
||||
monkeypatch.setattr(ts, "download_tokensave", _raise)
|
||||
assert ts.ensure_tokensave() is None
|
||||
|
||||
|
||||
def test_pinned_version_env_override(monkeypatch) -> None:
|
||||
monkeypatch.setenv("HEADROOM_TOKENSAVE_VERSION", "v9.9.9")
|
||||
assert ts._pinned_version() == "v9.9.9"
|
||||
monkeypatch.delenv("HEADROOM_TOKENSAVE_VERSION", raising=False)
|
||||
assert ts._pinned_version() == ts.TOKENSAVE_VERSION
|
||||
Loading…
Add table
Add a link
Reference in a new issue