mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat(wrap): default code-memory to Serena (dashboard browser off) behind unified --code-memory (#2413)
## What
Two commits:
1. **Unify code-memory MCP selection behind `--code-memory
{tokensave|serena|none}`** (+ `HEADROOM_CODE_MEMORY`), collapsing the
`--serena`/`--no-serena`/`--no-tokensave` flag tangle into one selector.
Old flags remain as hidden deprecated aliases that map into it. Shared
across the code-memory-capable subcommands (claude/codex/grok).
2. **Default the engine to Serena**, with its **dashboard browser
suppressed**.
## Why Serena as default
Serena is a mature, offline, symbol-level code-navigation MCP with broad
language coverage (LSP-backed) — the strongest zero-account default for
reducing tokens by letting the agent query
symbols/definitions/references instead of reading whole files. It
attacks the *protected-reads* volume the proxy deliberately doesn't
compress, so it's complementary to the pipeline compressors.
## Browser suppression (in Serena's own settings)
`_ensure_serena_dashboard_disabled()` sets
`web_dashboard_open_on_launch: false` in `~/.serena/serena_config.yml`
when Serena is set up, so wrapped sessions don't spawn a browser tab.
The dashboard backend stays reachable manually at `localhost:24282`.
This lives in Serena's config (authoritative), not just a startup flag.
## Schema-overhead note
Serena injects tool schemas per request; that cost is deferred by the
tool-search deferral the coding profile already enables
(`HEADROOM_TOOL_SEARCH=1`), so tools load on demand — the navigation
benefit without a standing schema tax on turns that don't navigate.
## Selection / escape hatches
`--code-memory serena` (default) · `tokensave` (lighter/faster) · `none`
(disable). Deprecated `--serena`/`--no-serena`/`--no-tokensave` still
work.
## Testing
Updated the primary/backup policy test to the serena-primary default;
code-memory selector + serena disable/migrate tests pass. Local: 21
passed (policy + code-memory); ruff + mypy clean. Full suite in CI.
This commit is contained in:
parent
446ec26003
commit
6e4425a6bd
3 changed files with 288 additions and 36 deletions
|
|
@ -679,6 +679,69 @@ _rtk_option = click.option(
|
|||
)
|
||||
|
||||
|
||||
# --- Code-memory MCP selection ------------------------------------------------
|
||||
# The code-memory MCP is on by default (tokensave). Swap it with --code-memory
|
||||
# serena, or turn it off with --code-memory none. Selection flows through
|
||||
# HEADROOM_CODE_MEMORY (set by the eager --code-memory callback) so it works the
|
||||
# same on every agent without threading a param through each subcommand — the
|
||||
# same approach as _rtk_option above.
|
||||
_CODE_MEMORY_ENV = "HEADROOM_CODE_MEMORY"
|
||||
_CODE_MEMORY_TOKENSAVE = "tokensave"
|
||||
_CODE_MEMORY_SERENA = "serena"
|
||||
_CODE_MEMORY_NONE = "none"
|
||||
_VALID_CODE_MEMORY = {_CODE_MEMORY_TOKENSAVE, _CODE_MEMORY_SERENA, _CODE_MEMORY_NONE}
|
||||
|
||||
|
||||
def _resolve_code_memory(kwargs: dict[str, Any]) -> str:
|
||||
"""Resolve which code-memory MCP to register.
|
||||
|
||||
Precedence: the explicit selector (``--code-memory`` / ``HEADROOM_CODE_MEMORY``)
|
||||
wins; otherwise the deprecated ``--serena`` / ``--no-tokensave`` / ``--no-serena``
|
||||
flags map into it; otherwise the default is ``serena`` — mature, offline,
|
||||
symbol-level code navigation (tokensave is a lighter opt-in).
|
||||
"""
|
||||
env = os.environ.get(_CODE_MEMORY_ENV, "").strip().lower()
|
||||
if env:
|
||||
if env not in _VALID_CODE_MEMORY:
|
||||
raise click.ClickException(
|
||||
f"{_CODE_MEMORY_ENV} must be one of: {', '.join(sorted(_VALID_CODE_MEMORY))}"
|
||||
)
|
||||
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"):
|
||||
return _CODE_MEMORY_TOKENSAVE
|
||||
return _CODE_MEMORY_SERENA
|
||||
|
||||
|
||||
def _code_memory_flag_callback(ctx: Any, param: Any, value: str | None) -> str | None:
|
||||
"""Click eager callback: ``--code-memory X`` sets HEADROOM_CODE_MEMORY so the
|
||||
central resolver (:func:`_resolve_code_memory`) sees the choice without
|
||||
threading a param through every wrap subcommand."""
|
||||
if value:
|
||||
os.environ[_CODE_MEMORY_ENV] = value
|
||||
return value
|
||||
|
||||
|
||||
# Shared selector applied to code-memory-capable subcommands (claude/codex/grok).
|
||||
# ``expose_value=False`` so no subcommand signature changes; it flows purely
|
||||
# through HEADROOM_CODE_MEMORY.
|
||||
_code_memory_option = click.option(
|
||||
"--code-memory",
|
||||
type=click.Choice([_CODE_MEMORY_TOKENSAVE, _CODE_MEMORY_SERENA, _CODE_MEMORY_NONE]),
|
||||
default=None,
|
||||
expose_value=False,
|
||||
is_eager=True,
|
||||
callback=_code_memory_flag_callback,
|
||||
help=(
|
||||
"Code-memory MCP to register: 'serena' (default), 'tokensave', or 'none'. "
|
||||
"Also set by HEADROOM_CODE_MEMORY. Replaces --serena/--no-serena/--no-tokensave."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _setup_rtk(verbose: bool = False) -> Path | None:
|
||||
"""Ensure rtk is installed and hooks are registered."""
|
||||
if not _rtk_opt_in():
|
||||
|
|
@ -1432,6 +1495,44 @@ def _setup_headroom_mcp(
|
|||
click.echo(line)
|
||||
|
||||
|
||||
def _ensure_serena_dashboard_disabled(*, verbose: bool = False) -> None:
|
||||
"""Disable Serena's browser dashboard auto-open in ``~/.serena/serena_config.yml``.
|
||||
|
||||
Serena opens its web dashboard in a browser tab on launch by default
|
||||
(``web_dashboard_open_on_launch: true``). Since Headroom now registers Serena
|
||||
as the default code-memory MCP, flip that setting off so wrapped sessions
|
||||
don't spawn a browser tab. The dashboard backend still runs and stays
|
||||
reachable at http://localhost:24282/dashboard/. The setting lives in Serena's
|
||||
own config (authoritative, unlike a startup flag); other keys and comments are
|
||||
preserved via a targeted line edit rather than a YAML rewrite.
|
||||
"""
|
||||
import re
|
||||
|
||||
cfg = Path.home() / ".serena" / "serena_config.yml"
|
||||
key = "web_dashboard_open_on_launch"
|
||||
try:
|
||||
if cfg.exists():
|
||||
text = cfg.read_text(encoding="utf-8")
|
||||
pattern = re.compile(rf"^(\s*){re.escape(key)}:\s*\S+\s*$", re.MULTILINE)
|
||||
if pattern.search(text):
|
||||
new = pattern.sub(rf"\g<1>{key}: false", text)
|
||||
else:
|
||||
new = text.rstrip("\n") + f"\n{key}: false\n"
|
||||
if new != text:
|
||||
cfg.write_text(new, encoding="utf-8")
|
||||
if verbose:
|
||||
click.echo(" Serena: disabled dashboard browser auto-open (serena_config.yml)")
|
||||
else:
|
||||
cfg.parent.mkdir(parents=True, exist_ok=True)
|
||||
# Serena fills defaults for any keys we omit, so a single-key file is valid.
|
||||
cfg.write_text(f"{key}: false\n", encoding="utf-8")
|
||||
if verbose:
|
||||
click.echo(" Serena: created serena_config.yml with dashboard auto-open off")
|
||||
except OSError as e:
|
||||
if verbose:
|
||||
click.echo(f" Serena: could not update serena_config.yml ({e})")
|
||||
|
||||
|
||||
def _setup_serena_mcp(
|
||||
registrar: Any, *, context: str, verbose: bool = False, force: bool = False
|
||||
) -> None:
|
||||
|
|
@ -1460,6 +1561,9 @@ def _setup_serena_mcp(
|
|||
click.echo(" Serena MCP: uvx not found — install uv/uvx to enable Serena; skipping")
|
||||
return
|
||||
|
||||
# Serena is a real launch now — make sure it won't pop a browser tab.
|
||||
_ensure_serena_dashboard_disabled(verbose=verbose)
|
||||
|
||||
spec = build_serena_spec(context)
|
||||
result = registrar.register_server(spec, force=force)
|
||||
|
||||
|
|
@ -1699,38 +1803,45 @@ def _disable_tokensave_mcp(registrar: Any, *, verbose: bool = False) -> None:
|
|||
|
||||
|
||||
def _setup_coding_compressor(registrar: Any, *, serena_context: str, **kwargs: Any) -> None:
|
||||
"""Set up the coding-task compressor: tokensave primary, Serena backup.
|
||||
"""Set up the code-memory MCP, selected via ``--code-memory`` (default serena).
|
||||
|
||||
Policy (decided per the integration):
|
||||
Selection (see :func:`_resolve_code_memory`):
|
||||
|
||||
* ``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``.
|
||||
* ``serena`` (default) — register Serena and remove any Headroom-installed
|
||||
tokensave. Serena is mature, offline, and symbol-level.
|
||||
* ``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.
|
||||
|
||||
``kwargs`` carries the boolean flags ``serena``, ``no_serena``,
|
||||
``no_tokensave`` and the per-agent registrar ``force`` semantics.
|
||||
Deprecated ``--serena`` / ``--no-serena`` / ``--no-tokensave`` flags map into
|
||||
the selector. User-managed MCP entries are always left untouched (ledger).
|
||||
"""
|
||||
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"))
|
||||
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"))
|
||||
|
||||
tokensave_ok = False
|
||||
if no_tokensave:
|
||||
if selection == _CODE_MEMORY_NONE:
|
||||
_disable_tokensave_mcp(registrar, verbose=verbose)
|
||||
else:
|
||||
tokensave_ok = _setup_tokensave_mcp(registrar, verbose=verbose, force=force)
|
||||
_disable_serena_mcp(registrar, verbose=verbose, reason="--code-memory none")
|
||||
return
|
||||
|
||||
if serena or (not tokensave_ok and not no_serena):
|
||||
if selection == _CODE_MEMORY_SERENA:
|
||||
_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:
|
||||
# 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"
|
||||
"--no-serena"
|
||||
if suppress_serena_fallback
|
||||
else "tokensave is the primary code-graph compressor"
|
||||
)
|
||||
_disable_serena_mcp(registrar, verbose=verbose, reason=reason)
|
||||
|
||||
|
|
@ -4473,18 +4584,25 @@ def wrap_selfheal(marker: str | None) -> None:
|
|||
is_flag=True,
|
||||
help="Skip headroom MCP server registration (compression markers will be unactionable)",
|
||||
)
|
||||
@_code_memory_option
|
||||
@click.option(
|
||||
"--no-tokensave",
|
||||
is_flag=True,
|
||||
help="Skip the tokensave code-graph MCP server (primary coding-task compressor)",
|
||||
hidden=True,
|
||||
help="Deprecated: use --code-memory none/serena. Skip the tokensave code-graph MCP.",
|
||||
)
|
||||
@click.option(
|
||||
"--serena",
|
||||
is_flag=True,
|
||||
help="Force the Serena MCP backup compressor on (registered automatically when "
|
||||
"tokensave is unavailable)",
|
||||
hidden=True,
|
||||
help="Deprecated: use --code-memory serena. Force the Serena MCP compressor on.",
|
||||
)
|
||||
@click.option(
|
||||
"--no-serena",
|
||||
is_flag=True,
|
||||
hidden=True,
|
||||
help="Deprecated: use --code-memory tokensave/none. Never register Serena.",
|
||||
)
|
||||
@click.option("--no-serena", is_flag=True, help="Never register the Serena backup compressor")
|
||||
@click.option(
|
||||
"--code-graph",
|
||||
is_flag=True,
|
||||
|
|
@ -5533,18 +5651,25 @@ def _run_codex_wrap(
|
|||
is_flag=True,
|
||||
help="Skip headroom MCP server registration (compression markers will be unactionable)",
|
||||
)
|
||||
@_code_memory_option
|
||||
@click.option(
|
||||
"--no-tokensave",
|
||||
is_flag=True,
|
||||
help="Skip the tokensave code-graph MCP server (primary coding-task compressor)",
|
||||
hidden=True,
|
||||
help="Deprecated: use --code-memory none/serena. Skip the tokensave code-graph MCP.",
|
||||
)
|
||||
@click.option(
|
||||
"--serena",
|
||||
is_flag=True,
|
||||
help="Force the Serena MCP backup compressor on (registered automatically when "
|
||||
"tokensave is unavailable)",
|
||||
hidden=True,
|
||||
help="Deprecated: use --code-memory serena. Force the Serena MCP compressor on.",
|
||||
)
|
||||
@click.option(
|
||||
"--no-serena",
|
||||
is_flag=True,
|
||||
hidden=True,
|
||||
help="Deprecated: use --code-memory tokensave/none. Never register Serena.",
|
||||
)
|
||||
@click.option("--no-serena", is_flag=True, help="Never register the Serena backup compressor")
|
||||
@click.option(
|
||||
"--code-graph",
|
||||
is_flag=True,
|
||||
|
|
@ -6024,18 +6149,25 @@ def kimi(
|
|||
help="Skip CLI context-tool setup",
|
||||
)
|
||||
@click.option("--no-mcp", is_flag=True, help="Skip headroom MCP server registration")
|
||||
@_code_memory_option
|
||||
@click.option(
|
||||
"--no-tokensave",
|
||||
is_flag=True,
|
||||
help="Skip the tokensave code-graph MCP server (primary coding-task compressor)",
|
||||
hidden=True,
|
||||
help="Deprecated: use --code-memory none/serena. Skip the tokensave code-graph MCP.",
|
||||
)
|
||||
@click.option(
|
||||
"--serena",
|
||||
is_flag=True,
|
||||
help="Force the Serena MCP backup compressor on (registered automatically when "
|
||||
"tokensave is unavailable)",
|
||||
hidden=True,
|
||||
help="Deprecated: use --code-memory serena. Force the Serena MCP compressor on.",
|
||||
)
|
||||
@click.option(
|
||||
"--no-serena",
|
||||
is_flag=True,
|
||||
hidden=True,
|
||||
help="Deprecated: use --code-memory tokensave/none. Never register Serena.",
|
||||
)
|
||||
@click.option("--no-serena", is_flag=True, help="Never register the Serena backup compressor")
|
||||
@click.option(
|
||||
"--code-graph",
|
||||
is_flag=True,
|
||||
|
|
|
|||
|
|
@ -186,12 +186,17 @@ def _spy_compressor(monkeypatch: pytest.MonkeyPatch, *, tokensave_ok: bool) -> d
|
|||
return calls
|
||||
|
||||
|
||||
def test_policy_tokensave_primary_disables_serena(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
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["tokensave"] == "setup"
|
||||
assert calls["serena_setup"] is False
|
||||
assert calls["serena_disabled"] == "tokensave is now the primary code-graph compressor"
|
||||
assert calls["serena_setup"] is True
|
||||
assert calls["tokensave"] == "disabled"
|
||||
assert calls["serena_disabled"] is None
|
||||
|
||||
|
||||
def test_policy_serena_fallback_when_tokensave_unavailable(
|
||||
|
|
|
|||
115
tests/test_wrap_code_memory.py
Normal file
115
tests/test_wrap_code_memory.py
Normal file
|
|
@ -0,0 +1,115 @@
|
|||
"""Code-memory MCP is selectable via --code-memory (default tokensave).
|
||||
|
||||
Covers the resolver precedence (selector > deprecated flags > default), the
|
||||
orchestrator 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
|
||||
|
||||
import os
|
||||
from unittest.mock import patch
|
||||
|
||||
from click.testing import CliRunner
|
||||
|
||||
from headroom.cli import wrap
|
||||
|
||||
|
||||
def _clean_env() -> dict[str, str]:
|
||||
env = dict(os.environ)
|
||||
env.pop("HEADROOM_CODE_MEMORY", None)
|
||||
return env
|
||||
|
||||
|
||||
def test_default_is_serena() -> None:
|
||||
with patch.dict(os.environ, _clean_env(), clear=True):
|
||||
assert wrap._resolve_code_memory({}) == wrap._CODE_MEMORY_SERENA
|
||||
|
||||
|
||||
def test_selector_env_wins() -> None:
|
||||
for val in (wrap._CODE_MEMORY_SERENA, wrap._CODE_MEMORY_NONE, wrap._CODE_MEMORY_TOKENSAVE):
|
||||
with patch.dict(os.environ, {"HEADROOM_CODE_MEMORY": val}):
|
||||
# selector beats any legacy flag
|
||||
assert wrap._resolve_code_memory({"serena": True, "no_tokensave": True}) == val
|
||||
|
||||
|
||||
def test_deprecated_flags_map_into_selector() -> None:
|
||||
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({"no_tokensave": True}) == wrap._CODE_MEMORY_SERENA
|
||||
# --no-serena means "not serena" → the other real graph, tokensave
|
||||
assert wrap._resolve_code_memory({"no_serena": True}) == wrap._CODE_MEMORY_TOKENSAVE
|
||||
assert (
|
||||
wrap._resolve_code_memory({"no_tokensave": True, "no_serena": True})
|
||||
== wrap._CODE_MEMORY_NONE
|
||||
)
|
||||
|
||||
|
||||
def test_serena_dashboard_disabled_flips_existing_config(tmp_path, monkeypatch) -> None:
|
||||
monkeypatch.setenv("HOME", str(tmp_path))
|
||||
cfg = tmp_path / ".serena" / "serena_config.yml"
|
||||
cfg.parent.mkdir(parents=True)
|
||||
cfg.write_text(
|
||||
"web_dashboard: true\nweb_dashboard_open_on_launch: true\ngui_log_window: false\n"
|
||||
)
|
||||
wrap._ensure_serena_dashboard_disabled()
|
||||
text = cfg.read_text()
|
||||
assert "web_dashboard_open_on_launch: false" in text
|
||||
assert "web_dashboard: true" in text # other keys preserved
|
||||
|
||||
|
||||
def test_serena_dashboard_disabled_creates_config(tmp_path, monkeypatch) -> None:
|
||||
monkeypatch.setenv("HOME", str(tmp_path))
|
||||
wrap._ensure_serena_dashboard_disabled()
|
||||
cfg = tmp_path / ".serena" / "serena_config.yml"
|
||||
assert cfg.exists()
|
||||
assert "web_dashboard_open_on_launch: false" in cfg.read_text()
|
||||
|
||||
|
||||
def test_invalid_env_raises() -> None:
|
||||
import click
|
||||
|
||||
with patch.dict(os.environ, {"HEADROOM_CODE_MEMORY": "bogus"}):
|
||||
try:
|
||||
wrap._resolve_code_memory({})
|
||||
except click.ClickException:
|
||||
pass
|
||||
else: # pragma: no cover
|
||||
raise AssertionError("invalid HEADROOM_CODE_MEMORY should raise ClickException")
|
||||
|
||||
|
||||
def _dispatch_calls(selection: str, extra: dict | None = None) -> list[str]:
|
||||
"""Run the orchestrator with a given selection, recording which setup/disable
|
||||
helpers fire (all mocked)."""
|
||||
calls: list[str] = []
|
||||
env = _clean_env()
|
||||
env["HEADROOM_CODE_MEMORY"] = selection
|
||||
with (
|
||||
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, "_disable_tokensave_mcp", lambda *a, **k: calls.append("disable_tokensave")
|
||||
),
|
||||
patch.object(wrap, "_disable_serena_mcp", lambda *a, **k: calls.append("disable_serena")),
|
||||
):
|
||||
wrap._setup_coding_compressor(object(), serena_context="claude-code", **(extra or {}))
|
||||
return calls
|
||||
|
||||
|
||||
def test_orchestrator_dispatch() -> None:
|
||||
assert _dispatch_calls(wrap._CODE_MEMORY_TOKENSAVE) == ["tokensave", "disable_serena"]
|
||||
assert _dispatch_calls(wrap._CODE_MEMORY_SERENA) == ["disable_tokensave", "serena"]
|
||||
assert set(_dispatch_calls(wrap._CODE_MEMORY_NONE)) == {"disable_tokensave", "disable_serena"}
|
||||
|
||||
|
||||
def test_code_memory_option_present_only_on_code_memory_agents() -> None:
|
||||
runner = CliRunner()
|
||||
for tool in ("claude", "codex", "grok"):
|
||||
out = runner.invoke(wrap.wrap, [tool, "--help"]).output
|
||||
assert "--code-memory" in out, f"--code-memory missing from `wrap {tool} --help`"
|
||||
# aider does not register a code-memory MCP → no flag
|
||||
out = runner.invoke(wrap.wrap, ["aider", "--help"]).output
|
||||
assert "--code-memory" not in out
|
||||
Loading…
Add table
Add a link
Reference in a new issue