akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

refactor(wrap): retire tokensave; Serena is the code-memory MCP (#2499)

Browse source 
## Description

`tokensave` was a **downloaded third-party Rust binary**
(`aovestdipaperino/tokensave`) that `headroom wrap` registered as a
code-graph MCP server. This removes it entirely and standardises on
**Serena** as the code-memory MCP — which was already the default in
`wrap`. Serena runs on demand via `uvx`, so Headroom no longer downloads
or executes a binary of its own for code memory.

The change is a *removal + safe transition*, not a behaviour flip:
Serena was already the default, so existing users move over
automatically. This PR also folds in a small README repositioning
(Headroom = the proxy; Serena is the recommended companion; RTK/lean-ctx
are third-party tools we don't control), since it's the same
tooling-stack story.

Closes #

## Type of Change

- [ ] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [x] Documentation update
- [x] Code refactoring (no functional changes)

## Changes Made

- **Removed the external tool:** `headroom/graph/tokensave_installer.py`
(the download-and-execute path), `build_tokensave_spec`, and the
`_ensure_tokensave_binary` / `_index_tokensave_project` /
`_setup_tokensave_mcp` helpers; dropped the dead `_setup_code_graph`.
- **`--code-memory`** now offers `serena` (default) or `none` — the
`tokensave` choice is gone. The strands `HeadroomBundle` uses Serena as
its (default-on) code-memory MCP.
- **Tombstone / transition (ledger-verified):** `headroom wrap` **and**
`headroom unwrap` remove a previously Headroom-installed `tokensave` MCP
entry so upgrading users stop launching it, and print that the leftover
`~/.local/bin/tokensave` binary and `.tokensave/` folders are safe to
delete. A user-managed `tokensave` entry is left untouched. Mirrors the
existing `codebase-memory-mcp` retirement.
- **Graceful for existing users:** `HEADROOM_CODE_MEMORY=tokensave` and
`--no-tokensave` resolve to Serena instead of erroring; `--no-serena`
now means "no code memory". **No state migration needed** — both tools'
indexes are regenerable caches of the source, so Serena simply
re-indexes.
- **Docs/README:** replaced the "tokensave binary trust model" section
with a Serena note + an "Upgrading from tokensave?" callout; retired the
RTK "first-class part of our stack" framing.
- **Tests:** deleted the tokensave-only test files
(`test_graph_tokensave.py`, `test_cli/test_tokensave_helpers.py`,
`test_cli/test_tokensave_setup.py`), rewrote `test_wrap_code_memory.py`
for the new resolver/dispatch, fixed a codex test that patched a removed
symbol.

Net: **+102 / −1187 lines.**

## Testing

- [x] Unit tests pass (`pytest`) — targeted to the affected areas
- [x] Linting passes (`ruff check`)
- [x] Type checking passes (`mypy`)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
$ ruff check headroom/cli/wrap.py headroom/mcp_registry/ headroom/integrations/strands/ tests/test_wrap_code_memory.py tests/test_cli/conftest.py tests/test_cli/test_wrap_codex.py
All checks passed!

$ mypy headroom/cli/wrap.py headroom/mcp_registry headroom/integrations/strands/bundle.py
Success: no issues found in 12 source files

$ pytest tests/test_wrap_code_memory.py tests/test_cli/test_wrap_codex.py \
         tests/test_cli/test_wrap_claude_vertex_proxy_env.py \
         tests/test_cli/test_wrap_claude_finally_unbound.py -q
======================== 120 passed in 97.86s (0:01:37) ========================

$ pytest tests/test_cli --collect-only -q
========================= 713 tests collected in 1.82s =========================   # no import errors after symbol removal
```

## Real Behavior Proof

- **Environment:** macOS (darwin), Python 3.12.6, local `.venv`, on
branch `tejas/remove-tokensave`.
- **Exact command / steps:**
- `python -c "from headroom.integrations.strands.bundle import
HeadroomBundle; b=HeadroomBundle(enable_headroom_mcp=False,
enable_serena_mcp=False); print(len(b.tools))"` → confirms the module
imports after `build_tokensave_spec` removal (the import that my change
would otherwise break).
- CliRunner-driven `wrap codex --prepare-only` (in `test_wrap_codex.py`)
writes `[mcp_servers.serena]` (with `command = "uvx"`, `"--context",
"codex"`) to the codex config and **no** tokensave entry.
- `_resolve_code_memory` unit tests confirm: default → `serena`;
`HEADROOM_CODE_MEMORY=tokensave` → `serena`; `--no-serena` → `none`;
`--code-memory bogus` → `ClickException`.
- **Observed result:** import OK (`tools: 0`); Serena registered,
tokensave absent; resolver behaves as above; 120/120 tests pass.
- **Not tested:** a live `headroom wrap` against a real agent on a
machine with a *previously-installed* tokensave MCP entry — the
tombstone-removal path is covered by unit tests with a fake
registrar/ledger, not an end-to-end run.

## 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
- [x] I did **not** edit `CHANGELOG.md` — it is generated by
release-please from my Conventional Commit PR title

## Additional Notes

- `--no-tokensave` / `--serena` / `--no-serena` are retained as hidden,
deprecated flags (no-ops or mapped) so existing scripts don't break.
- `--code-graph` is unchanged — it's the proxy's live file-watcher flag
and was never the tokensave MCP; only the dead tokensave hook behind it
was removed.
- `CHANGELOG.md` intentionally left untouched.

🤖 Generated with [Claude Code](https://claude.com/claude-code)
This commit is contained in:
Tejas Chopra 2026-07-22 20:59:24 -07:00 • committed by GitHub
parent 5bd2266f16
commit 5d23a0aec2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 102 additions and 1187 deletions
Download patch file Download diff file Expand all files Collapse all files

4
README.md
View file

@ -105,7 +105,7 @@ headroom perf
headroom dashboard # live savings dashboard (proxy must be running) headroom dashboard # live savings dashboard (proxy must be running)
``` ```
To use headroom, it is recommended you launch a wrapped agent session each time so that all necessary setup is completed. When wrapping a coding agent, headroom starts a local proxy, sets up an MCP server that provides tools such as rtk and tokensave, and launches a coding agent session configured to proxy requests to headroom. To use headroom, it is recommended you launch a wrapped agent session each time so that all necessary setup is completed. When wrapping a coding agent, headroom starts a local proxy, installs **Serena** for semantic code navigation, and launches a coding agent session configured to proxy requests through headroom.
The `headroom` CLI ships **only** via the PyPI package. The npm `headroom-ai` is the TypeScript SDK — a library you import (`import { compress } from 'headroom-ai'`), not a CLI, so it provides no `headroom` command. The `headroom` CLI ships **only** via the PyPI package. The npm `headroom-ai` is the TypeScript SDK — a library you import (`import { compress } from 'headroom-ai'`), not a CLI, so it provides no `headroom` command.
@ -538,7 +538,7 @@ Headroom runs **locally**, covers **every** content type, works with every major
| [Compresr](https://compresr.ai), [Token Co.](https://thetokencompany.ai) | Text sent to their API | Hosted API call | No | No | | [Compresr](https://compresr.ai), [Token Co.](https://thetokencompany.ai) | Text sent to their API | Hosted API call | No | No |
| OpenAI Compaction | Conversation history | Provider-native | No | No | | OpenAI Compaction | Conversation history | Provider-native | No | No |
> **Attribution.** Headroom ships with the excellent [RTK](https://github.com/rtk-ai/rtk) binary for shell-output rewriting — `git show --short`, scoped `ls`, summarized installers. Huge thanks to the RTK team; their tool is a first-class part of our stack, and Headroom compresses everything downstream of it. Headroom can also use [lean-ctx](https://github.com/yvgude/lean-ctx) as the selected CLI context tool; set `HEADROOM_CONTEXT_TOOL=lean-ctx` before running `headroom wrap ...`. > **Stack & integrations.** Headroom is the **proxy** — that's what we build and offer, and it compresses everything flowing through it no matter what sits upstream. Our recommended companion is **[Serena](https://github.com/oraios/serena)** (installed by default when you wrap an agent) for semantic code navigation — plus **Ponytail** if you want leaner model output. Everything else is your call: Headroom vendors the third-party [RTK](https://github.com/rtk-ai/rtk) and [lean-ctx](https://github.com/yvgude/lean-ctx) binaries for shell-output rewriting, but we don't own or control either project — swap between them with `HEADROOM_CONTEXT_TOOL`, or turn them off. You're free to attach your own tooling too — code-memory MCP, Graphify, Caveman, or any MCP server — and Headroom compresses downstream of all of it.
## Contributing ## Contributing

10
docs/content/docs/proxy.mdx
View file

@ -71,15 +71,13 @@ Avoid setting process-wide variables such as `HTTP_PROXY`, `HTTPS_PROXY`, `ALL_P
| `--intercept-tool-results` | `false` | Opt into tool-result interceptors such as ast-grep Read outlining. | | `--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. | | `--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-aware` / `--no-code-aware` | disabled | Enable or disable AST-based code compression. Requires `headroom-ai[code]`. |
| `--code-graph` | `false` | Force a tokensave code-graph index of the current project (tokensave is the default coding-task compressor registered by `headroom wrap`). | | `--code-graph` | `false` | Enable the proxy's live code-graph file watcher for the current project. |
#### tokensave binary trust model #### Code-memory MCP (Serena)
`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. `headroom wrap` registers **[Serena](https://github.com/oraios/serena)** as the code-memory MCP for semantic, symbol-level code navigation. Serena runs on demand via `uvx` — Headroom downloads and executes no binary of its own — and indexes the current project locally. Pass `--code-memory none` to register no code-memory MCP.
- Set `HEADROOM_BINARIES_OFFLINE=1` to never reach the network — `wrap` then uses an already-installed tokensave or falls back to Serena. > **Upgrading from tokensave?** Earlier releases registered a `tokensave` MCP server (a downloaded Rust binary). tokensave has been retired in favour of Serena. On your next `headroom wrap` / `headroom unwrap`, Headroom removes the `tokensave` MCP entry it installed and switches you to Serena — nothing to migrate, since both are just indexes rebuilt from your source. The leftover `tokensave` binary in `~/.local/bin` and any `.tokensave/` folders are unused and safe to delete.
- `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. 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.

269
headroom/cli/wrap.py
View file

@ -741,39 +741,38 @@ _serena_instructions_option = click.option(
# --- Code-memory MCP selection ------------------------------------------------ # --- Code-memory MCP selection ------------------------------------------------
# The code-memory MCP is on by default (tokensave). Swap it with --code-memory # The code-memory MCP is Serena by default; turn it off with --code-memory none.
# serena, or turn it off with --code-memory none. Selection flows through # Selection flows through HEADROOM_CODE_MEMORY (set by the eager --code-memory
# HEADROOM_CODE_MEMORY (set by the eager --code-memory callback) so it works the # callback) so it works the same on every agent without threading a param
# same on every agent without threading a param through each subcommand — the # through each subcommand — the same approach as _rtk_option above.
# same approach as _rtk_option above.
_CODE_MEMORY_ENV = "HEADROOM_CODE_MEMORY" _CODE_MEMORY_ENV = "HEADROOM_CODE_MEMORY"
_CODE_MEMORY_TOKENSAVE = "tokensave"
_CODE_MEMORY_SERENA = "serena" _CODE_MEMORY_SERENA = "serena"
_CODE_MEMORY_NONE = "none" _CODE_MEMORY_NONE = "none"
_VALID_CODE_MEMORY = {_CODE_MEMORY_TOKENSAVE, _CODE_MEMORY_SERENA, _CODE_MEMORY_NONE} _VALID_CODE_MEMORY = {_CODE_MEMORY_SERENA, _CODE_MEMORY_NONE}
def _resolve_code_memory(kwargs: dict[str, Any]) -> str: def _resolve_code_memory(kwargs: dict[str, Any]) -> str:
"""Resolve which code-memory MCP to register. """Resolve which code-memory MCP to register.
Precedence: the explicit selector (``--code-memory`` / ``HEADROOM_CODE_MEMORY``) Precedence: the explicit selector (``--code-memory`` / ``HEADROOM_CODE_MEMORY``)
wins; otherwise the deprecated ``--serena`` / ``--no-tokensave`` / ``--no-serena`` wins; otherwise the deprecated ``--serena`` / ``--no-serena`` flags map into
flags map into it; otherwise the default is ``serena`` — mature, offline, it; otherwise the default is ``serena`` — mature, offline, symbol-level code
symbol-level code navigation (tokensave is a lighter opt-in). navigation. The retired ``tokensave`` option is accepted gracefully: an
explicit ``tokensave`` selector (or the deprecated ``--no-tokensave`` flag)
now resolves to Serena.
""" """
env = os.environ.get(_CODE_MEMORY_ENV, "").strip().lower() env = os.environ.get(_CODE_MEMORY_ENV, "").strip().lower()
if env == "tokensave":
click.echo(" Note: the tokensave code-memory option was retired — using Serena instead.")
return _CODE_MEMORY_SERENA
if env: if env:
if env not in _VALID_CODE_MEMORY: if env not in _VALID_CODE_MEMORY:
raise click.ClickException( raise click.ClickException(
f"{_CODE_MEMORY_ENV} must be one of: {', '.join(sorted(_VALID_CODE_MEMORY))}" f"{_CODE_MEMORY_ENV} must be one of: {', '.join(sorted(_VALID_CODE_MEMORY))}"
) )
return env return env
if kwargs.get("serena"):
return _CODE_MEMORY_SERENA
if kwargs.get("no_tokensave"):
return _CODE_MEMORY_NONE if kwargs.get("no_serena") else _CODE_MEMORY_SERENA
if kwargs.get("no_serena"): if kwargs.get("no_serena"):
return _CODE_MEMORY_TOKENSAVE return _CODE_MEMORY_NONE
return _CODE_MEMORY_SERENA return _CODE_MEMORY_SERENA
@ -791,14 +790,14 @@ def _code_memory_flag_callback(ctx: Any, param: Any, value: str | None) -> str |
# through HEADROOM_CODE_MEMORY. # through HEADROOM_CODE_MEMORY.
_code_memory_option = click.option( _code_memory_option = click.option(
"--code-memory", "--code-memory",
type=click.Choice([_CODE_MEMORY_TOKENSAVE, _CODE_MEMORY_SERENA, _CODE_MEMORY_NONE]), type=click.Choice([_CODE_MEMORY_SERENA, _CODE_MEMORY_NONE]),
default=None, default=None,
expose_value=False, expose_value=False,
is_eager=True, is_eager=True,
callback=_code_memory_flag_callback, callback=_code_memory_flag_callback,
help=( help=(
"Code-memory MCP to register: 'serena' (default), 'tokensave', or 'none'. " "Code-memory MCP to register: 'serena' (default) or 'none'. "
"Also set by HEADROOM_CODE_MEMORY. Replaces --serena/--no-serena/--no-tokensave." "Also set by HEADROOM_CODE_MEMORY. Replaces --serena/--no-serena."
), ),
) )
@ -1800,7 +1799,7 @@ def _index_serena_project(*, verbose: bool = False) -> None:
start the MCP server) in the project directory so the first symbol query is start the MCP server) in the project directory so the first symbol query is
not paying for a cold index. Timeout-guarded and best-effort: Serena also not paying for a cold index. Timeout-guarded and best-effort: Serena also
indexes lazily on demand, so a failure or timeout here never blocks the indexes lazily on demand, so a failure or timeout here never blocks the
wrap. Mirrors :func:`_index_tokensave_project`. wrap.
""" """
if shutil.which("uvx") is None: if shutil.which("uvx") is None:
if verbose: if verbose:
@ -1927,8 +1926,8 @@ def _disable_serena_mcp(
keeps launching Serena on startup. Just *skipping* registration on a later keeps launching Serena on startup. Just *skipping* registration on a later
run leaves that stale entry in place — so this removes the entry Headroom 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 installed. A user-managed Serena (absent from our ledger) is reported but
left untouched. ``reason`` is surfaced in the message: ``--no-serena`` when left untouched. ``reason`` is surfaced in the message (e.g. ``--no-serena``
the user opted out, or a note that tokensave is now the primary compressor. or ``--code-memory none`` when the user opted out).
""" """
if not registrar.detect(): if not registrar.detect():
if verbose: if verbose:
@ -1956,119 +1955,11 @@ def _disable_serena_mcp(
# ============================================================================= # =============================================================================
# tokensave — primary coding-task compressor (Serena is the backup) # tokensave — retired; Serena replaced it. The helpers below only clean up a
# tokensave entry a prior release installed, so upgrading users stop launching it.
# ============================================================================= # =============================================================================
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: def _remove_headroom_installed_tokensave_mcp(registrar: Any) -> str:
"""Remove the tokensave MCP entry only if the ledger proves Headroom installed it.""" """Remove the tokensave MCP entry only if the ledger proves Headroom installed it."""
from headroom.mcp_registry.ledger import clear_install, headroom_installed_matching from headroom.mcp_registry.ledger import clear_install, headroom_installed_matching
@ -2083,25 +1974,33 @@ def _remove_headroom_installed_tokensave_mcp(registrar: Any) -> str:
def _disable_tokensave_mcp(registrar: Any, *, verbose: bool = False) -> None: def _disable_tokensave_mcp(registrar: Any, *, verbose: bool = False) -> None:
"""Make ``--no-tokensave`` actively remove a Headroom-installed tokensave entry.""" """Remove a Headroom-installed tokensave MCP entry left by a prior release.
tokensave was retired in favour of Serena. On upgrade we actively remove the
stale ``tokensave`` entry so the agent stops launching it, and point the user
at the leftover on-disk artifacts (we never delete files for them). A
user-managed entry (absent from our ledger) is reported but left in place.
"""
if not registrar.detect(): if not registrar.detect():
if verbose: if verbose:
click.echo(f" tokensave MCP: {registrar.display_name} not detected — skipping") click.echo(f" tokensave MCP: {registrar.display_name} not detected — skipping")
return return
if registrar.get_server("tokensave") is None: if registrar.get_server("tokensave") is None:
if verbose:
click.echo(" Skipping tokensave MCP (--no-tokensave)")
return return
status = _remove_headroom_installed_tokensave_mcp(registrar) status = _remove_headroom_installed_tokensave_mcp(registrar)
if status == "removed": if status == "removed":
click.echo(" Removed previously-installed tokensave MCP (--no-tokensave)") click.echo(" Removed retired tokensave MCP (replaced by Serena)")
click.echo(f" restart {registrar.display_name} if it was already running") click.echo(f" restart {registrar.display_name} if it was already running")
click.echo(
" leftover files are safe to delete: the 'tokensave' binary in "
"~/.local/bin and any '.tokensave/' folder in your projects"
)
elif status == "not_headroom_owned": elif status == "not_headroom_owned":
click.echo( click.echo(
" tokensave MCP is present but user-managed — leaving it in place " " tokensave MCP is present but user-managed — leaving it in place "
"(--no-tokensave only removes entries Headroom installed)" "(Headroom only removes entries it installed)"
) )
else: # "failed" else: # "failed"
click.echo( click.echo(
@ -2115,68 +2014,31 @@ def _setup_coding_compressor(registrar: Any, *, serena_context: str, **kwargs: A
Selection (see :func:`_resolve_code_memory`): Selection (see :func:`_resolve_code_memory`):
* ``serena`` (default) — register Serena and remove any Headroom-installed * ``serena`` (default) — register Serena (mature, offline, symbol-level).
tokensave. Serena is mature, offline, and symbol-level. * ``none`` — register nothing.
* ``tokensave`` — register tokensave (lighter/faster); Serena is registered
automatically only as a backup when tokensave is unavailable (unless the
deprecated ``--no-serena`` suppressed the fallback).
* ``none`` — remove both Headroom-installed entries.
Deprecated ``--serena`` / ``--no-serena`` / ``--no-tokensave`` flags map into Either way, any Headroom-installed ``tokensave`` entry from a prior release
the selector. User-managed MCP entries are always left untouched (ledger). is removed (tokensave was retired in favour of Serena). The deprecated
``--serena`` / ``--no-serena`` flags map into the selector; user-managed MCP
entries are always left untouched (ledger).
""" """
force = bool(kwargs.get("force")) force = bool(kwargs.get("force"))
verbose = bool(kwargs.get("verbose")) verbose = bool(kwargs.get("verbose"))
selection = _resolve_code_memory(kwargs) selection = _resolve_code_memory(kwargs)
# Deprecated --no-serena: in tokensave mode, don't auto-fall back to Serena.
suppress_serena_fallback = bool(kwargs.get("no_serena")) # Retire any tokensave entry a prior release installed, whatever the selection.
_disable_tokensave_mcp(registrar, verbose=verbose)
if selection == _CODE_MEMORY_NONE: if selection == _CODE_MEMORY_NONE:
_disable_tokensave_mcp(registrar, verbose=verbose)
_disable_serena_mcp(registrar, verbose=verbose, reason="--code-memory none") _disable_serena_mcp(registrar, verbose=verbose, reason="--code-memory none")
return return
if selection == _CODE_MEMORY_SERENA: _setup_serena_mcp(registrar, context=serena_context, verbose=verbose, force=force)
_disable_tokensave_mcp(registrar, verbose=verbose)
_setup_serena_mcp(registrar, context=serena_context, verbose=verbose, force=force)
return
# tokensave (explicit opt-in): register it; Serena is the automatic backup.
tokensave_ok = _setup_tokensave_mcp(registrar, verbose=verbose, force=force)
if not tokensave_ok and not suppress_serena_fallback:
_setup_serena_mcp(registrar, context=serena_context, verbose=verbose, force=force)
else:
reason = (
"--no-serena"
if suppress_serena_fallback
else "tokensave is the primary code-graph compressor"
)
_disable_serena_mcp(registrar, verbose=verbose, reason=reason)
_CBM_MCP_SERVER_NAME = "codebase-memory-mcp" _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). # rtk instructions for tools without hook support (Codex, Cursor, Aider).
# These get injected into AGENTS.md / .cursorrules so the LLM voluntarily # These get injected into AGENTS.md / .cursorrules so the LLM voluntarily
# uses rtk-prefixed commands. Kept concise to minimize instruction overhead. # uses rtk-prefixed commands. Kept concise to minimize instruction overhead.
@ -4599,9 +4461,6 @@ def _launch_tool(
if configure_launch is not None: if configure_launch is not None:
args, env, env_vars_display = configure_launch(actual_port, args, env, env_vars_display) args, env, env_vars_display = configure_launch(actual_port, args, env, env_vars_display)
if code_graph:
_setup_code_graph(verbose=False)
click.echo() click.echo()
click.echo(f" Launching {tool_label} (API routed through Headroom)...") click.echo(f" Launching {tool_label} (API routed through Headroom)...")
for var in env_vars_display: for var in env_vars_display:
@ -4900,7 +4759,7 @@ def wrap_selfheal(marker: str | None) -> None:
"--no-tokensave", "--no-tokensave",
is_flag=True, is_flag=True,
hidden=True, hidden=True,
help="Deprecated: use --code-memory none/serena. Skip the tokensave code-graph MCP.", help="Deprecated and ignored: tokensave was retired; Serena is the default code memory.",
) )
@click.option( @click.option(
"--serena", "--serena",
@ -4912,12 +4771,12 @@ def wrap_selfheal(marker: str | None) -> None:
"--no-serena", "--no-serena",
is_flag=True, is_flag=True,
hidden=True, hidden=True,
help="Deprecated: use --code-memory tokensave/none. Never register Serena.", help="Deprecated: use --code-memory none. Register no code-memory MCP.",
) )
@click.option( @click.option(
"--code-graph", "--code-graph",
is_flag=True, is_flag=True,
help="Force a tokensave code-graph index now (tokensave is the default compressor)", help="Enable the proxy's live code-graph file watcher for the current project.",
) )
@click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)") @click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)")
@click.option( @click.option(
@ -4991,17 +4850,14 @@ def claude(
\b \b
Examples: Examples:
headroom wrap claude # Start everything headroom wrap claude # Start everything (Serena code memory)
headroom wrap claude --memory # With persistent memory headroom wrap claude --memory # With persistent memory
headroom wrap claude --resume <id> # Resume a session headroom wrap claude --resume <id> # Resume a session
headroom wrap claude -- -p # Claude in print mode 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 --context-tool # Enable CLI context-tool setup
headroom wrap claude --no-context-tool # Skip 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-mcp # Skip MCP retrieve tool registration
headroom wrap claude --no-serena # Never register the Serena backup headroom wrap claude --code-memory none # No code-memory MCP
headroom wrap claude --1m # Preserve the 1M context window headroom wrap claude --1m # Preserve the 1M context window
""" """
# RTK/context-tool is opt-in (off by default): --context-tool (legacy) and # RTK/context-tool is opt-in (off by default): --context-tool (legacy) and
@ -5156,7 +5012,7 @@ def claude(
elif verbose: elif verbose:
click.echo(" Skipping MCP retrieve tool (--no-mcp)") click.echo(" Skipping MCP retrieve tool (--no-mcp)")
# Coding-task compressor: tokensave primary, Serena backup. # Coding-task compressor: Serena (retires any legacy tokensave entry).
from headroom.mcp_registry import ClaudeRegistrar from headroom.mcp_registry import ClaudeRegistrar
_setup_coding_compressor( _setup_coding_compressor(
@ -5168,9 +5024,6 @@ def claude(
verbose=verbose, verbose=verbose,
) )
if code_graph:
_setup_code_graph(verbose=verbose)
proxy_url = _claude_proxy_base_url(actual_port) proxy_url = _claude_proxy_base_url(actual_port)
click.echo() click.echo()
click.echo(" Launching Claude Code (API routed through Headroom)...") click.echo(" Launching Claude Code (API routed through Headroom)...")
@ -5800,8 +5653,8 @@ def _prepare_codex_wrap_state(
elif verbose: elif verbose:
click.echo(" Skipping MCP retrieve tool (--no-mcp)") click.echo(" Skipping MCP retrieve tool (--no-mcp)")
# Coding-task compressor: tokensave primary, Serena backup. Codex starts # Coding-task compressor: Serena (retires any legacy tokensave entry). Codex
# long-lived MCP subprocesses from config.toml, so force re-registration. # starts long-lived MCP subprocesses from config.toml, so force re-registration.
from headroom.mcp_registry import CodexRegistrar from headroom.mcp_registry import CodexRegistrar
_setup_coding_compressor( _setup_coding_compressor(
@ -5978,7 +5831,7 @@ def _run_codex_wrap(
"--no-tokensave", "--no-tokensave",
is_flag=True, is_flag=True,
hidden=True, hidden=True,
help="Deprecated: use --code-memory none/serena. Skip the tokensave code-graph MCP.", help="Deprecated and ignored: tokensave was retired; Serena is the default code memory.",
) )
@click.option( @click.option(
"--serena", "--serena",
@ -5990,12 +5843,12 @@ def _run_codex_wrap(
"--no-serena", "--no-serena",
is_flag=True, is_flag=True,
hidden=True, hidden=True,
help="Deprecated: use --code-memory tokensave/none. Never register Serena.", help="Deprecated: use --code-memory none. Register no code-memory MCP.",
) )
@click.option( @click.option(
"--code-graph", "--code-graph",
is_flag=True, is_flag=True,
help="Force a tokensave code-graph index now (tokensave is the default compressor)", help="Enable the proxy's live code-graph file watcher for the current project.",
) )
@click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)") @click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)")
@click.option( @click.option(
@ -6051,9 +5904,7 @@ def codex(
headroom wrap codex -- "fix the bug" # Pass prompt to 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-context-tool # Skip CLI context-tool setup
headroom wrap codex --no-mcp # Skip MCP retrieve tool registration headroom wrap codex --no-mcp # Skip MCP retrieve tool registration
headroom wrap codex --no-tokensave # Skip tokensave; fall back to Serena headroom wrap codex --code-memory none # No code-memory MCP
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 --port 9999 # Custom proxy port
headroom wrap codex --backend anyllm --anyllm-provider groq headroom wrap codex --backend anyllm --anyllm-provider groq
""" """
@ -6477,7 +6328,7 @@ def kimi(
"--no-tokensave", "--no-tokensave",
is_flag=True, is_flag=True,
hidden=True, hidden=True,
help="Deprecated: use --code-memory none/serena. Skip the tokensave code-graph MCP.", help="Deprecated and ignored: tokensave was retired; Serena is the default code memory.",
) )
@click.option( @click.option(
"--serena", "--serena",
@ -6489,12 +6340,12 @@ def kimi(
"--no-serena", "--no-serena",
is_flag=True, is_flag=True,
hidden=True, hidden=True,
help="Deprecated: use --code-memory tokensave/none. Never register Serena.", help="Deprecated: use --code-memory none. Register no code-memory MCP.",
) )
@click.option( @click.option(
"--code-graph", "--code-graph",
is_flag=True, is_flag=True,
help="Force a tokensave code-graph index now (tokensave is the default compressor)", help="Enable the proxy's live code-graph file watcher for the current project.",
) )
@click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)") @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("--learn", is_flag=True, help="Enable live traffic learning")

256
headroom/graph/tokensave_installer.py
View file

@ -1,256 +0,0 @@
"""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

2
headroom/integrations/strands/__init__.py
View file

@ -90,6 +90,6 @@ __all__ = [
# Provider detection # Provider detection
"get_headroom_provider", "get_headroom_provider",
"get_model_name_from_strands", "get_model_name_from_strands",
# One-helper MCP + hook wiring (Headroom + tokensave/Serena + RTK-equivalent) # One-helper MCP + hook wiring (Headroom + Serena + RTK-equivalent)
"HeadroomBundle", "HeadroomBundle",
] ]

45
headroom/integrations/strands/bundle.py
View file

@ -11,14 +11,11 @@ Strands-native primitives:
needs the original; Strands' MCP dispatcher resolves it via this needs the original; Strands' MCP dispatcher resolves it via this
server. Works identically in streaming and non-streaming. server. Works identically in streaming and non-streaming.
* **tokensave MCP** — the primary coding-task compressor: a local * **Serena MCP** — the coding-task compressor (symbol search,
semantic code-graph server (``tokensave serve``) the agent queries references, call chains, impact analysis) so the agent queries
for symbols, call chains, and impact analysis instead of reading the code graph instead of reading whole files. Auto-installed via
whole files. Requires the ``tokensave`` binary on PATH. ``uvx`` on first launch. On by default; disable with
``enable_serena_mcp=False``.
* **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. * **HeadroomHookProvider** — the RTK-equivalent for Strands.
Compresses tool outputs in-place via ``AfterToolCallEvent`` so Compresses tool outputs in-place via ``AfterToolCallEvent`` so
@ -83,7 +80,6 @@ from headroom.mcp_registry.install import (
DEFAULT_PROXY_URL, DEFAULT_PROXY_URL,
build_headroom_spec, build_headroom_spec,
build_serena_spec, build_serena_spec,
build_tokensave_spec,
) )
from .hooks import HeadroomHookProvider from .hooks import HeadroomHookProvider
@ -110,10 +106,6 @@ def _make_headroom_client(proxy_url: str) -> MCPClient:
return _client_for(build_headroom_spec(proxy_url)) 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: def _make_serena_client(context: str) -> MCPClient:
return _client_for(build_serena_spec(context)) return _client_for(build_serena_spec(context))
@ -128,12 +120,8 @@ class HeadroomBundle:
(``http://127.0.0.1:8787``). (``http://127.0.0.1:8787``).
serena_context: Serena context label. Default ``"ide-assistant"``. serena_context: Serena context label. Default ``"ide-assistant"``.
enable_headroom_mcp: Include the Headroom MCP server. Default True. enable_headroom_mcp: Include the Headroom MCP server. Default True.
enable_tokensave_mcp: Include the tokensave MCP server — the primary enable_serena_mcp: Include the Serena MCP server — the coding-task
coding-task compressor. Default True. Requires the ``tokensave`` compressor. Default True. Adds a ``uvx`` first-launch download.
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 enable_hooks: Include :class:`HeadroomHookProvider` for in-place
tool-output compression (the RTK-equivalent for Strands). tool-output compression (the RTK-equivalent for Strands).
Default True. Default True.
@ -150,10 +138,8 @@ class HeadroomBundle:
proxy_url: str = DEFAULT_PROXY_URL proxy_url: str = DEFAULT_PROXY_URL
serena_context: str = DEFAULT_SERENA_CONTEXT serena_context: str = DEFAULT_SERENA_CONTEXT
enable_headroom_mcp: bool = True enable_headroom_mcp: bool = True
# tokensave is the primary coding-task compressor; Serena is the backup # Serena is the coding-task compressor (symbol-level code navigation).
# and stays off unless explicitly enabled. enable_serena_mcp: bool = True
enable_tokensave_mcp: bool = True
enable_serena_mcp: bool = False
# The proxy is the single source of truth for compression — it sees # The proxy is the single source of truth for compression — it sees
# the full message list, owns CompressionPolicy, owns PrefixCacheTracker, # the full message list, owns CompressionPolicy, owns PrefixCacheTracker,
# and places `cache_control` breakpoints. The in-process hook # and places `cache_control` breakpoints. The in-process hook
@ -166,7 +152,6 @@ class HeadroomBundle:
config: HeadroomConfig | None = None config: HeadroomConfig | None = None
_headroom_mcp: MCPClient | None = field(default=None, init=False, repr=False, compare=False) _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) _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) _hook: HeadroomHookProvider | None = field(default=None, init=False, repr=False, compare=False)
@ -177,13 +162,10 @@ class HeadroomBundle:
"HeadroomBundle: Headroom MCP client constructed (proxy_url=%s)", "HeadroomBundle: Headroom MCP client constructed (proxy_url=%s)",
self.proxy_url, 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: if self.enable_serena_mcp:
self._serena_mcp = _make_serena_client(self.serena_context) self._serena_mcp = _make_serena_client(self.serena_context)
logger.info( logger.info(
"HeadroomBundle: Serena MCP client constructed (backup, context=%s)", "HeadroomBundle: Serena MCP client constructed (code memory, context=%s)",
self.serena_context, self.serena_context,
) )
if self.enable_hooks: if self.enable_hooks:
@ -200,8 +182,6 @@ class HeadroomBundle:
out: list[Any] = [] out: list[Any] = []
if self._headroom_mcp is not None: if self._headroom_mcp is not None:
out.append(self._headroom_mcp) out.append(self._headroom_mcp)
if self._tokensave_mcp is not None:
out.append(self._tokensave_mcp)
if self._serena_mcp is not None: if self._serena_mcp is not None:
out.append(self._serena_mcp) out.append(self._serena_mcp)
return out return out
@ -216,11 +196,6 @@ class HeadroomBundle:
"""Direct handle to the Headroom MCPClient (for advanced callers).""" """Direct handle to the Headroom MCPClient (for advanced callers)."""
return self._headroom_mcp 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 @property
def serena_mcp(self) -> MCPClient | None: def serena_mcp(self) -> MCPClient | None:
"""Direct handle to the Serena MCPClient (for advanced callers).""" """Direct handle to the Serena MCPClient (for advanced callers)."""

2
headroom/mcp_registry/__init__.py
View file

@ -22,7 +22,6 @@ from .install import (
DEFAULT_PROXY_URL, DEFAULT_PROXY_URL,
build_headroom_spec, build_headroom_spec,
build_serena_spec, build_serena_spec,
build_tokensave_spec,
get_all_registrars, get_all_registrars,
install_everywhere, install_everywhere,
) )
@ -43,7 +42,6 @@ __all__ = [
"build_headroom_spec", "build_headroom_spec",
"build_serena_spec", "build_serena_spec",
"build_server_json", "build_server_json",
"build_tokensave_spec",
"format_result", "format_result",
"format_results", "format_results",
"get_all_registrars", "get_all_registrars",

16
headroom/mcp_registry/install.py
View file

@ -72,22 +72,6 @@ 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( def install_everywhere(
proxy_url: str = DEFAULT_PROXY_URL, proxy_url: str = DEFAULT_PROXY_URL,
*, *,

13
tests/test_cli/conftest.py
View file

@ -1,13 +1,10 @@
"""Shared fixtures for the CLI test suite. """Shared fixtures for the CLI test suite.
tokensave is now the primary coding-task compressor, so a default ``headroom wrap`` may fetch helper binaries (e.g. rtk, lean-ctx) over the
``headroom wrap`` tries to fetch the tokensave release binary. Force offline network via ``headroom.binaries``. Force offline across CLI tests so a missing
across CLI tests so a missing binary resolves to ``None`` (→ Serena fallback) binary resolves locally instead of reaching out to GitHub releases. Tests that
instead of reaching out to GitHub releases. Tests that exercise the exercise a binary-present path patch the relevant resolver directly and are
tokensave-present path patch ``_ensure_tokensave_binary`` / ``ensure_tokensave`` unaffected by this guard.
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 from __future__ import annotations

112
tests/test_cli/test_tokensave_helpers.py
View file

@ -1,112 +0,0 @@
"""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

232
tests/test_cli/test_tokensave_setup.py
View file

@ -1,232 +0,0 @@
"""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_serena_primary_by_default_disables_tokensave(
monkeypatch: pytest.MonkeyPatch,
) -> None:
# Serena is now the default code-memory engine: with no explicit selection it
# is set up and any Headroom-installed tokensave entry is disabled (the
# inverse of the old tokensave-primary policy).
calls = _spy_compressor(monkeypatch, tokensave_ok=True)
wrap_cli._setup_coding_compressor(_FakeRegistrar(), serena_context="claude-code")
assert calls["serena_setup"] is True
assert calls["tokensave"] == "disabled"
assert calls["serena_disabled"] is None
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

9
tests/test_cli/test_wrap_codex.py
View file

@ -1650,13 +1650,8 @@ def test_wrap_codex_prepare_only_registers_serena_when_uvx_exists(
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None): with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
with patch("headroom.cli.wrap.shutil.which", side_effect=fake_which): with patch("headroom.cli.wrap.shutil.which", side_effect=fake_which):
# tokensave is the primary code-graph compressor; Serena is only # Serena is the code-memory MCP; assert it lands in the codex config.
# the backup, registered when tokensave is unavailable. Force it result = runner.invoke(main, ["wrap", "codex", "--prepare-only"])
# unavailable so this test deterministically exercises the Serena
# path regardless of whether a real tokensave binary was installed
# in the shared bin dir by an earlier test in the suite.
with patch("headroom.cli.wrap._ensure_tokensave_binary", return_value=None):
result = runner.invoke(main, ["wrap", "codex", "--prepare-only"])
assert result.exit_code == 0, result.output assert result.exit_code == 0, result.output
content = config_file.read_text(encoding="utf-8") content = config_file.read_text(encoding="utf-8")

284
tests/test_graph_tokensave.py
View file

@ -1,284 +0,0 @@
"""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

35
tests/test_wrap_code_memory.py
View file

@ -1,8 +1,9 @@
"""Code-memory MCP is selectable via --code-memory (default tokensave). """Code-memory MCP is selectable via --code-memory (default serena).
Covers the resolver precedence (selector > deprecated flags > default), the Covers the resolver precedence (selector > deprecated flags > default), the
orchestrator dispatch for each selection, and that --code-memory is exposed on graceful retirement of the removed ``tokensave`` option, the orchestrator
the code-memory-capable subcommands (claude/codex/grok) but not others. dispatch for each selection, and that --code-memory is exposed on the
code-memory-capable subcommands (claude/codex/grok) but not others.
""" """
from __future__ import annotations from __future__ import annotations
@ -10,6 +11,7 @@ from __future__ import annotations
import os import os
from unittest.mock import patch from unittest.mock import patch
import click
from click.testing import CliRunner from click.testing import CliRunner
from headroom.cli import wrap from headroom.cli import wrap
@ -27,22 +29,26 @@ def test_default_is_serena() -> None:
def test_selector_env_wins() -> None: def test_selector_env_wins() -> None:
for val in (wrap._CODE_MEMORY_SERENA, wrap._CODE_MEMORY_NONE, wrap._CODE_MEMORY_TOKENSAVE): for val in (wrap._CODE_MEMORY_SERENA, wrap._CODE_MEMORY_NONE):
with patch.dict(os.environ, {"HEADROOM_CODE_MEMORY": val}): with patch.dict(os.environ, {"HEADROOM_CODE_MEMORY": val}):
# selector beats any legacy flag # selector beats any legacy flag
assert wrap._resolve_code_memory({"serena": True, "no_tokensave": True}) == val assert wrap._resolve_code_memory({"serena": True, "no_serena": True}) == val
def test_deprecated_flags_map_into_selector() -> None: def test_deprecated_flags_map_into_selector() -> None:
with patch.dict(os.environ, _clean_env(), clear=True): with patch.dict(os.environ, _clean_env(), clear=True):
assert wrap._resolve_code_memory({"serena": True}) == wrap._CODE_MEMORY_SERENA assert wrap._resolve_code_memory({"serena": True}) == wrap._CODE_MEMORY_SERENA
# tokensave is retired: --no-tokensave is now a no-op → default serena
assert wrap._resolve_code_memory({"no_tokensave": True}) == wrap._CODE_MEMORY_SERENA assert wrap._resolve_code_memory({"no_tokensave": True}) == wrap._CODE_MEMORY_SERENA
# --no-serena means "not serena" → the other real graph, tokensave # --no-serena means "no code memory" now that tokensave is gone
assert wrap._resolve_code_memory({"no_serena": True}) == wrap._CODE_MEMORY_TOKENSAVE assert wrap._resolve_code_memory({"no_serena": True}) == wrap._CODE_MEMORY_NONE
assert (
wrap._resolve_code_memory({"no_tokensave": True, "no_serena": True})
== wrap._CODE_MEMORY_NONE def test_retired_tokensave_selector_maps_to_serena() -> None:
) # An explicit HEADROOM_CODE_MEMORY=tokensave (or --code-memory tokensave from
# an old script) degrades gracefully to Serena instead of erroring.
with patch.dict(os.environ, {"HEADROOM_CODE_MEMORY": "tokensave"}):
assert wrap._resolve_code_memory({}) == wrap._CODE_MEMORY_SERENA
def test_serena_dashboard_disabled_flips_existing_config(tmp_path, monkeypatch) -> None: def test_serena_dashboard_disabled_flips_existing_config(tmp_path, monkeypatch) -> None:
@ -67,8 +73,6 @@ def test_serena_dashboard_disabled_creates_config(tmp_path, monkeypatch) -> None
def test_invalid_env_raises() -> None: def test_invalid_env_raises() -> None:
import click
with patch.dict(os.environ, {"HEADROOM_CODE_MEMORY": "bogus"}): with patch.dict(os.environ, {"HEADROOM_CODE_MEMORY": "bogus"}):
try: try:
wrap._resolve_code_memory({}) wrap._resolve_code_memory({})
@ -86,9 +90,6 @@ def _dispatch_calls(selection: str, extra: dict | None = None) -> list[str]:
env["HEADROOM_CODE_MEMORY"] = selection env["HEADROOM_CODE_MEMORY"] = selection
with ( with (
patch.dict(os.environ, env, clear=True), patch.dict(os.environ, env, clear=True),
patch.object(
wrap, "_setup_tokensave_mcp", lambda *a, **k: (calls.append("tokensave"), True)[1]
),
patch.object(wrap, "_setup_serena_mcp", lambda *a, **k: calls.append("serena")), patch.object(wrap, "_setup_serena_mcp", lambda *a, **k: calls.append("serena")),
patch.object( patch.object(
wrap, "_disable_tokensave_mcp", lambda *a, **k: calls.append("disable_tokensave") wrap, "_disable_tokensave_mcp", lambda *a, **k: calls.append("disable_tokensave")
@ -100,7 +101,7 @@ def _dispatch_calls(selection: str, extra: dict | None = None) -> list[str]:
def test_orchestrator_dispatch() -> None: def test_orchestrator_dispatch() -> None:
assert _dispatch_calls(wrap._CODE_MEMORY_TOKENSAVE) == ["tokensave", "disable_serena"] # A legacy tokensave entry is always retired first, then the selection applies.
assert _dispatch_calls(wrap._CODE_MEMORY_SERENA) == ["disable_tokensave", "serena"] assert _dispatch_calls(wrap._CODE_MEMORY_SERENA) == ["disable_tokensave", "serena"]
assert set(_dispatch_calls(wrap._CODE_MEMORY_NONE)) == {"disable_tokensave", "disable_serena"} assert set(_dispatch_calls(wrap._CODE_MEMORY_NONE)) == {"disable_tokensave", "disable_serena"}