mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
## Description Removes both third-party CLI context tools — **rtk** and **lean-ctx** — and with them the context-tool selector itself. Headroom no longer downloads, installs or configures either one, and there is no replacement. The previous pass (#2344) gated only three entry points inside `headroom/cli/wrap.py`. That left the feature reachable in practice: | Gap | Effect | |---|---| | `scripts/install.sh:1544`, `install.ps1:1681` | Ran `rtk init --global --auto-patch` from bash/PowerShell, **bypassing the Python gate entirely** — `curl \| sh` still wrote a Claude Code `PreToolUse` hook regardless of `HEADROOM_RTK` | | `wrap.py` `_setup_context_tool_for_agent` | **`wrap openhands` was broken by default**: `rtk_required=True` met a gate returning `None` → `SystemExit(1)`. Invisible because all 8 openhands tests patched `_ensure_rtk_binary` to a fake path | | `proxy/helpers.py`, `subscription/tracker.py` | Proxy shelled out to `rtk gain` from `/stats`, the dashboard and `headroom perf`; the tracker polled it per contribution (`_RTK_WIRING_DEFAULT = "enabled"`) | | No cleanup path | Nothing removed artifacts an earlier default had installed, so a machine that once ran the old default kept rtk in the loop forever (#1669, #1955) | Also worth noting: the rtk binary download had **no SHA or signature verification** — only `rtk --version` as a smoke test. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [x] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [x] Code refactoring (no functional changes) ## Changes Made **Removed** — `headroom/rtk/` and `headroom/lean_ctx/` packages, `headroom/cli/wrap_rtk_metrics.py`, `_selected_context_tool` / `_setup_context_tool_for_agent` / `_VALID_CONTEXT_TOOLS`, the `--rtk` / `--no-rtk` / `--no-project-rtk` / `--keep-rtk` flags across all 18 wrap subcommands, `HEADROOM_RTK*`, the proxy-side `rtk gain` polling, the dashboard CLI-filtering panel (rows + all 8 `cliFiltering*` Alpine getters), `paths.rtk_path()` / `lean_ctx_path()`, the SDK path helpers, `benchmarks/rtk_loop_learn_eval.py`, and the `headroom/rtk/**` CI path filters. **Fails loudly, not silently** — `--context-tool` / `--no-context-tool` / `HEADROOM_CONTEXT_TOOL` are kept solely to error out. They live in shell profiles, aliases and CI jobs, and accepting them as a no-op would read as Headroom having quietly stopped working. The installers reject them too, which matters more than it looks: their arg parsers forward the first unknown flag **and everything after it** to the wrapped tool, so a leftover `--no-rtk` would have silently swallowed a following `--port` and then been ignored downstream. **New `headroom/context_tool_cleanup.py`** — deleting the code cannot help a machine that already ran the old default, since the hooks, binaries and injected guidance are durable on disk. `purge_context_tool_artifacts()` runs once per `wrap`/`unwrap` and removes the registered hook entries, the generated hook scripts, the Headroom-managed `~/.local/bin` symlinks, the vendored `~/.headroom/bin/{rtk,lean-ctx}` binaries, the `lean-ctx` MCP server entry and the marker-fenced instruction blocks. Deliberately conservative: idempotent, **skips** a malformed config rather than overwriting it, and only unlinks a symlink resolving inside Headroom's own bin dir so a user's own build is untouched. It reports on **stderr**, because `wrap/unwrap openclaw --prepare-only` emit machine-readable JSON on stdout as their entire contract. Skipped for `wrap selfheal` (runs from a SessionStart hook; must not race Claude Code's writer for `~/.claude.json`) and for `--help`, which must stay read-only. **Client-config hardening** (discovered while investigating a "corrupted Serena settings file" report) — `wrap.py` reset a settings file to `{}` when an existing file would not parse, then wrote that back. One hand-edited typo or a transient `EACCES`/`EINTR` on a valid file destroyed the user's `permissions`, `env` and `hooks`, on **every `headroom wrap claude`**. It now refuses to write. Separately, `fsutil.write_text` is now atomic (temp file + `fsync` + `os.replace`), fixing all 14 non-atomic client-config writes at once; it follows symlinks rather than replacing them (dotfile managers) and preserves an existing file's mode. **Deliberately kept** — `rtk` stays in the wrapper-peel list in `transforms/content_router.py`. It sits beside `sudo`/`env`/`timeout` as shell-command grammar, so `rtk cat f` is still classified as a file read for anyone running their own rtk install, which the purge intentionally leaves alone. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ ruff check headroom/ tests/ e2e/ --exclude headroom/dashboard/templates All checks passed! $ ruff format --check headroom/ tests/ e2e/ --exclude headroom/dashboard/templates 1255 files already formatted $ mypy headroom/ Success: no issues found in 508 source files $ pytest tests/test_context_tool_cleanup.py -q 11 passed $ pytest tests/test_fsutil.py -q 12 passed $ pytest tests/test_cli/test_wrap_codex.py -q # 89 tests 89 passed in 431.68s $ pytest tests/test_cli/test_wrap_opencode.py -q 39 passed in 257.46s $ pytest tests/test_cli/test_wrap_helpers.py -q 45 passed $ pytest tests/test_paths.py -q 75 passed $ pytest tests/test_cli/test_unwrap_claude.py -q 14 passed $ pytest tests/test_proxy_savings_history.py -q 39 passed $ pytest tests/test_cli/test_wrap_copilot.py -q 27 passed $ pytest tests/test_cli/test_wrap_zcode.py -q 20 passed $ pytest tests/test_subscription_tracker.py -q 9 passed $ pytest tests/test_proxy_dashboard_stats_cache.py -q 5 passed, 1 skipped ``` Repo-wide grep for 14 removed symbols (`headroom.rtk`, `headroom.lean_ctx`, `_ensure_rtk_binary`, `_selected_context_tool`, `_get_context_tool_stats`, `rtk_path`, `lean_ctx_path`, `wrap_rtk_metrics`, `HEADROOM_RTK`, `cli_tokens_avoided`, `tokens_saved_rtk`, …) across `*.py`, `*.ts`, `*.sh`, `*.ps1`, `*.yml`, `*.html`: **zero hits**. Notable test changes: `test_wrap_openhands.py` no longer patches `_ensure_rtk_binary` and asserts `wrap openhands --prepare-only` exits 0 unpatched — the regression that was previously masked. `test_wrap_continue.py` and `test_wrap_hintfile_agents.py` were removed (every test drove RTK instruction injection). A new `test_subscription_tracker.py::test_load_state_written_before_cli_context_tools_were_removed` proves a pre-removal `subscription_state.json` still loads. ## Real Behavior Proof - **Environment:** macOS 15.4 (darwin 25.4.0), Python 3.12.6, Headroom @ this branch, real `~/.headroom` and `~/.claude` on the dev machine. - **Exact command / steps and observed result:** ```text # 1. Retired flag fails loudly instead of silently no-op'ing $ headroom wrap codex --prepare-only --context-tool rtk Error: CLI context tools (rtk, lean-ctx) have been removed from Headroom: they rewrote shell commands through a third-party binary Headroom no longer manages. Drop --context-tool / --no-context-tool and unset HEADROOM_CONTEXT_TOOL; `headroom wrap` uninstalls what they left behind on first run. $ HEADROOM_CONTEXT_TOOL=lean-ctx headroom wrap codex --prepare-only Error: CLI context tools (rtk, lean-ctx) have been removed from Headroom: ... # 2. install.sh rejects the retired flags (extracted parse_wrap_args harness) ['--no-rtk', '--port', '9999'] rc=1 ERROR: CLI context tools ... Drop --no-rtk ['--context-tool=rtk'] rc=1 ERROR: CLI context tools ... Drop --context-tool $ bash -n scripts/install.sh # syntax OK # 3. Purge ran against the real machine, which had all the orphaned artifacts $ python -c "from headroom.context_tool_cleanup import purge_context_tool_artifacts; ..." removed ~/.headroom/bin/lean-ctx (51 MB) removed ~/.headroom/bin/rtk (7.7 MB) removed ~/.local/bin/rtk (symlink into ~/.headroom/bin) removed ~/.claude/hooks/rtk-rewrite.sh removed 8 lean-ctx-* hook scripts # ~/.claude.json afterwards: 90 top-level keys, 19 projects, mcpServers unchanged # → ~59 MB reclaimed, no unrelated key touched # 4. stdout stays machine-readable while the purge reports (planted a fake artifact) $ headroom wrap openclaw --prepare-only --gateway-provider-id codex >out 2>err $ cat out {"enabled":true,"config":{"proxyPort":8787,...}} # parses as JSON $ cat err Retired CLI context tool cleanup: removed /Users/tcms/.headroom/bin/rtk # 5. --help is inert (planted artifact survives), a real run purges $ headroom wrap codex --help → artifact survived: CORRECT $ headroom wrap openclaw --prepare-only → purged: CORRECT # 6. MCP purge dry-run against a copy of the real 82 KB ~/.claude.json top-level keys 90 -> 90; projects 19 -> 19; LOST keys: none all content outside mcpServers byte-identical: True ``` Dashboard rendered via the Playwright test after the panel removal: "Token Savings" shows only `Proxy 0 (0.0%)` / `Of total wire: 36.86%`, and "Token Usage" reads Before Compression → Proxy Removed → After Compression with no "Filtered (this session)" row. Nothing below the removed panel broke. - **Not tested:** Windows and Linux (macOS only) — `install.ps1` is verified by brace-balance and inspection, not executed, since no `pwsh` is available locally. The wrap e2e suite (`e2e/wrap/run.py`) was updated but not run; it needs the Docker e2e image. `serena project index` interaction is exercised in the stacked base PR. ## 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 (a CI guard enforces this) ## Additional Notes **Stacked on #2676** (`tejas/serena-config-bootstrap`) — please merge that first; this PR's base should then be retargeted to `main`, or it will read as containing that fix too. **Breaking-change migration for users:** - Drop `--rtk`, `--no-rtk`, `--no-project-rtk`, `--keep-rtk`, `--context-tool`, `--no-context-tool` from any alias, script or CI job, and unset `HEADROOM_RTK*` / `HEADROOM_CONTEXT_TOOL`. They now error rather than being ignored, so the failure is immediate and self-explaining. - Previously-installed artifacts are purged automatically on the next `wrap`/`unwrap`; no manual cleanup needed. - `headroom perf --json` no longer carries a `cli_filtering` key, and `/stats` no longer returns a `context_tool` section. **Docs:** `docs/rtk-architecture.md` deleted; RTK/lean-ctx removed from `README.md`, `docs/content/docs/{configuration,opencode,grok-build,docker-install,filesystem-contract}.mdx`, `docs/observability.md` and the matching `wiki/` pages. `REALIGNMENT/09-phase-G-rtk-observability.md` is marked SUPERSEDED rather than deleted, to keep the planning record. **Follow-ups not in scope:** `_emit_wrap_interrupted` was deleted as dead code — its only caller was the `except KeyboardInterrupt` guarding the binary download, so with no download there is nothing slow left to interrupt.
719 lines
27 KiB
Python
719 lines
27 KiB
Python
"""Direct unit tests for the shared wrap-subcommand helpers.
|
|
|
|
These helpers (`_print_wrap_banner`, `_run_proxy_only_watcher`) were
|
|
extracted to remove ~150 LOC of
|
|
copy-pasted scaffolding across the wrap subcommands (cursor / cline /
|
|
continue / goose / openhands). The wrap-*.py subcommand tests exercise
|
|
them indirectly; these tests pin the contract directly so a future
|
|
refactor that breaks one of these helpers fails *here* — at the helper
|
|
unit boundary — instead of in five different subcommand suites at
|
|
once with confusing diffs.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import errno
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import click
|
|
import pytest
|
|
from click.testing import CliRunner
|
|
|
|
from headroom import paths as paths_mod
|
|
from headroom.cli import wrap as wrap_mod
|
|
from headroom.cli.main import main
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _print_wrap_banner — centering math + box drawing.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def runner() -> CliRunner:
|
|
return CliRunner()
|
|
|
|
|
|
def _run_in_click_context(fn) -> str: # type: ignore[no-untyped-def]
|
|
"""Invoke `fn` inside a Click `runner.invoke` so `click.echo` output is captured."""
|
|
runner = CliRunner()
|
|
|
|
@click.command()
|
|
def _cmd() -> None:
|
|
fn()
|
|
|
|
result = runner.invoke(_cmd)
|
|
assert result.exit_code == 0, result.output
|
|
return result.output
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"agent",
|
|
["cline", "cursor", "continue", "goose", "openhands", "x", "a-very-long-agent-name"],
|
|
)
|
|
def test_print_wrap_banner_box_is_inner_width_chars_wide(agent: str) -> None:
|
|
"""The banner's horizontal rule should always be 47 chars between the `║` corners."""
|
|
output = _run_in_click_context(lambda: wrap_mod._print_wrap_banner(agent))
|
|
|
|
lines = [line for line in output.splitlines() if line.strip()]
|
|
assert len(lines) == 3, f"banner should be 3 non-empty lines; got {lines!r}"
|
|
top, title_line, bottom = lines
|
|
|
|
# Top and bottom are the `╔════...╗` rules with 47 equals signs between the corners.
|
|
assert top.endswith("╗")
|
|
assert bottom.endswith("╝")
|
|
assert top.count("═") == wrap_mod._WRAP_BANNER_INNER_WIDTH
|
|
assert bottom.count("═") == wrap_mod._WRAP_BANNER_INNER_WIDTH
|
|
|
|
# The middle line has the centered title.
|
|
assert title_line.startswith(" ║")
|
|
assert title_line.endswith("║")
|
|
assert f"HEADROOM WRAP: {agent.upper()}" in title_line
|
|
|
|
|
|
def test_print_wrap_banner_title_is_centered_or_near_centered() -> None:
|
|
"""Centering: pad_left and pad_right may differ by at most 1 when total padding is odd."""
|
|
output = _run_in_click_context(lambda: wrap_mod._print_wrap_banner("cline"))
|
|
|
|
lines = [line for line in output.splitlines() if line.strip()]
|
|
title_line = lines[1]
|
|
|
|
# Strip the leading " ║" and trailing "║" so we can measure spaces.
|
|
inner = title_line[3:-1]
|
|
assert len(inner) == wrap_mod._WRAP_BANNER_INNER_WIDTH
|
|
|
|
title = "HEADROOM WRAP: CLINE"
|
|
pad_left = len(inner) - len(inner.lstrip(" "))
|
|
pad_right = len(inner) - len(inner.rstrip(" "))
|
|
assert inner.strip() == title
|
|
assert abs(pad_left - pad_right) <= 1, (
|
|
f"banner not centered: pad_left={pad_left}, pad_right={pad_right}"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# wrap claude argument passthrough.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_wrap_claude_allows_claude_print_short_flag_in_passthrough_args() -> None:
|
|
"""Claude owns -p/--print; wrap claude must not parse it as --port."""
|
|
result = CliRunner().invoke(
|
|
main,
|
|
["wrap", "claude", "--prepare-only", "-p", "Say only: hello"],
|
|
)
|
|
|
|
assert result.exit_code == 0, result.output
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _run_proxy_only_watcher — must print banner, call setup callback, install
|
|
# signal handlers, and clean up. Heavily mocked since the real watcher
|
|
# blocks on `time.sleep` indefinitely.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_run_proxy_only_watcher_calls_setup_lines_callback(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""The print_setup_lines callback runs after the proxy is ready."""
|
|
# Fake proxy: a dummy object whose `.poll()` returns 0 after first iteration
|
|
# so the watcher exits cleanly via the "proxy exited unexpectedly" branch.
|
|
|
|
class _FakeProc:
|
|
def __init__(self) -> None:
|
|
self._polls = 0
|
|
|
|
def poll(self) -> int | None:
|
|
self._polls += 1
|
|
return 0 if self._polls > 1 else None
|
|
|
|
fake_proc = _FakeProc()
|
|
|
|
callback_calls: list[None] = []
|
|
|
|
def fake_setup(_port: int) -> None:
|
|
callback_calls.append(None)
|
|
|
|
monkeypatch.setattr(wrap_mod, "_ensure_proxy", lambda *a, **kw: (fake_proc, 8787))
|
|
# Replace time.sleep with a no-op so the loop spins quickly.
|
|
monkeypatch.setattr(wrap_mod.time, "sleep", lambda _s: None)
|
|
# Replace _make_cleanup to avoid side-effects on real ports/files.
|
|
monkeypatch.setattr(wrap_mod, "_make_cleanup", lambda holder, port: lambda *a, **kw: None)
|
|
# Avoid touching real signal handlers in the test process.
|
|
monkeypatch.setattr(wrap_mod.signal, "signal", lambda *a, **kw: None)
|
|
|
|
runner = CliRunner()
|
|
|
|
@click.command()
|
|
def _cmd() -> None:
|
|
wrap_mod._run_proxy_only_watcher(
|
|
agent_label="cline",
|
|
port=8787,
|
|
no_proxy=False,
|
|
learn=False,
|
|
memory=False,
|
|
agent_type="cline",
|
|
print_setup_lines=fake_setup,
|
|
)
|
|
|
|
inv = runner.invoke(_cmd)
|
|
# The watcher exits 1 when the proxy dies (our _FakeProc returns 0 on poll #2).
|
|
assert inv.exit_code == 1
|
|
assert callback_calls == [None]
|
|
# Banner is part of the helper's contract.
|
|
assert "HEADROOM WRAP: CLINE" in inv.output
|
|
# The "proxy exited unexpectedly" message is the documented exit branch.
|
|
assert "Proxy process exited unexpectedly." in inv.output
|
|
|
|
|
|
def test_run_proxy_only_watcher_keyboardinterrupt_shuts_down_cleanly(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Ctrl-C during the watcher loop prints `Shutting down...` and exits 0."""
|
|
|
|
class _FakeProc:
|
|
def poll(self) -> int | None:
|
|
return None # Proxy is healthy; loop would run forever.
|
|
|
|
sleep_calls = {"n": 0}
|
|
|
|
def raising_sleep(_s: float) -> None:
|
|
sleep_calls["n"] += 1
|
|
if sleep_calls["n"] >= 1:
|
|
raise KeyboardInterrupt
|
|
|
|
monkeypatch.setattr(wrap_mod, "_ensure_proxy", lambda *a, **kw: (_FakeProc(), 8787))
|
|
monkeypatch.setattr(wrap_mod.time, "sleep", raising_sleep)
|
|
monkeypatch.setattr(wrap_mod, "_make_cleanup", lambda holder, port: lambda *a, **kw: None)
|
|
monkeypatch.setattr(wrap_mod.signal, "signal", lambda *a, **kw: None)
|
|
|
|
runner = CliRunner()
|
|
|
|
@click.command()
|
|
def _cmd() -> None:
|
|
wrap_mod._run_proxy_only_watcher(
|
|
agent_label="cursor",
|
|
port=8787,
|
|
no_proxy=False,
|
|
learn=False,
|
|
memory=False,
|
|
agent_type="cursor",
|
|
print_setup_lines=lambda _port: None,
|
|
)
|
|
|
|
inv = runner.invoke(_cmd)
|
|
assert inv.exit_code == 0, inv.output
|
|
assert "Shutting down..." in inv.output
|
|
|
|
|
|
def test_run_proxy_only_watcher_unexpected_exception_returns_exit_1(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Unexpected exceptions in the body are caught and converted to SystemExit(1)."""
|
|
|
|
def boom(*a: Any, **kw: Any) -> None:
|
|
raise RuntimeError("boom")
|
|
|
|
monkeypatch.setattr(wrap_mod, "_ensure_proxy", boom)
|
|
monkeypatch.setattr(wrap_mod, "_make_cleanup", lambda holder, port: lambda *a, **kw: None)
|
|
monkeypatch.setattr(wrap_mod.signal, "signal", lambda *a, **kw: None)
|
|
|
|
runner = CliRunner()
|
|
|
|
@click.command()
|
|
def _cmd() -> None:
|
|
wrap_mod._run_proxy_only_watcher(
|
|
agent_label="cline",
|
|
port=8787,
|
|
no_proxy=False,
|
|
learn=False,
|
|
memory=False,
|
|
agent_type="cline",
|
|
print_setup_lines=lambda _port: None,
|
|
)
|
|
|
|
inv = runner.invoke(_cmd)
|
|
assert inv.exit_code == 1
|
|
assert "Error: boom" in inv.output
|
|
|
|
|
|
def test_run_proxy_only_watcher_calls_cleanup_on_finally(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Cleanup callable is invoked in the `finally` block regardless of exit path."""
|
|
|
|
cleanup_calls = {"n": 0}
|
|
|
|
def fake_cleanup(*a: Any, **kw: Any) -> None:
|
|
cleanup_calls["n"] += 1
|
|
|
|
def boom(*a: Any, **kw: Any) -> None:
|
|
raise RuntimeError("boom")
|
|
|
|
monkeypatch.setattr(wrap_mod, "_ensure_proxy", boom)
|
|
monkeypatch.setattr(wrap_mod, "_make_cleanup", lambda holder, port: fake_cleanup)
|
|
monkeypatch.setattr(wrap_mod.signal, "signal", lambda *a, **kw: None)
|
|
|
|
runner = CliRunner()
|
|
|
|
@click.command()
|
|
def _cmd() -> None:
|
|
wrap_mod._run_proxy_only_watcher(
|
|
agent_label="cline",
|
|
port=8787,
|
|
no_proxy=False,
|
|
learn=False,
|
|
memory=False,
|
|
agent_type="cline",
|
|
print_setup_lines=lambda _port: None,
|
|
)
|
|
|
|
inv = runner.invoke(_cmd)
|
|
assert inv.exit_code == 1
|
|
assert cleanup_calls["n"] >= 1, "cleanup must run via the finally block"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _project_name_from_cwd / _apply_project_header_env — per-project savings
|
|
# header injection for `headroom wrap claude` (issue: per-project savings).
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestApplyProjectHeaderEnv:
|
|
"""X-Headroom-Project injection into ANTHROPIC_CUSTOM_HEADERS."""
|
|
|
|
def test_sets_header_from_cwd_basename(
|
|
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
|
) -> None:
|
|
project_dir = tmp_path / "my-project"
|
|
project_dir.mkdir()
|
|
monkeypatch.chdir(project_dir)
|
|
|
|
env: dict[str, str] = {}
|
|
wrap_mod._apply_project_header_env(env)
|
|
|
|
assert env["ANTHROPIC_CUSTOM_HEADERS"] == "X-Headroom-Project: my-project"
|
|
|
|
def test_appends_to_existing_custom_headers(
|
|
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
|
) -> None:
|
|
project_dir = tmp_path / "proj"
|
|
project_dir.mkdir()
|
|
monkeypatch.chdir(project_dir)
|
|
|
|
env = {"ANTHROPIC_CUSTOM_HEADERS": "X-Custom-Trace: abc123"}
|
|
wrap_mod._apply_project_header_env(env)
|
|
|
|
# User header preserved verbatim, ours appended on a new line.
|
|
assert env["ANTHROPIC_CUSTOM_HEADERS"] == (
|
|
"X-Custom-Trace: abc123\nX-Headroom-Project: proj"
|
|
)
|
|
|
|
@pytest.mark.parametrize(
|
|
"user_value",
|
|
[
|
|
"X-Headroom-Project: their-name",
|
|
"x-headroom-project: their-name",
|
|
"X-HEADROOM-PROJECT: their-name",
|
|
"X-Other: 1\nx-Headroom-Project: their-name",
|
|
],
|
|
)
|
|
def test_existing_project_header_wins_case_insensitive(
|
|
self,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
tmp_path: Path,
|
|
user_value: str,
|
|
) -> None:
|
|
project_dir = tmp_path / "proj"
|
|
project_dir.mkdir()
|
|
monkeypatch.chdir(project_dir)
|
|
|
|
env = {"ANTHROPIC_CUSTOM_HEADERS": user_value}
|
|
wrap_mod._apply_project_header_env(env)
|
|
|
|
# Untouched: no duplicate header, user override wins.
|
|
assert env["ANTHROPIC_CUSTOM_HEADERS"] == user_value
|
|
|
|
@pytest.mark.parametrize(
|
|
"user_value",
|
|
[
|
|
"X-Headroom-Project-Id: other",
|
|
"X-Trace: mentions x-headroom-project in the value",
|
|
],
|
|
)
|
|
def test_similar_header_names_do_not_suppress_injection(
|
|
self,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
tmp_path: Path,
|
|
user_value: str,
|
|
) -> None:
|
|
project_dir = tmp_path / "proj"
|
|
project_dir.mkdir()
|
|
monkeypatch.chdir(project_dir)
|
|
|
|
env = {"ANTHROPIC_CUSTOM_HEADERS": user_value}
|
|
wrap_mod._apply_project_header_env(env)
|
|
|
|
# Only an exact header-name match counts as a user override.
|
|
assert env["ANTHROPIC_CUSTOM_HEADERS"] == (f"{user_value}\nX-Headroom-Project: proj")
|
|
|
|
def test_empty_cwd_name_sets_nothing(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""A degenerate cwd (e.g. filesystem root → empty basename) is a no-op."""
|
|
monkeypatch.setattr(wrap_mod.Path, "cwd", classmethod(lambda cls: Path("/")))
|
|
|
|
env: dict[str, str] = {}
|
|
wrap_mod._apply_project_header_env(env)
|
|
|
|
assert "ANTHROPIC_CUSTOM_HEADERS" not in env
|
|
|
|
def test_whitespace_only_cwd_name_sets_nothing(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr(wrap_mod.Path, "cwd", classmethod(lambda cls: Path("/tmp/ ")))
|
|
|
|
env: dict[str, str] = {}
|
|
wrap_mod._apply_project_header_env(env)
|
|
|
|
assert "ANTHROPIC_CUSTOM_HEADERS" not in env
|
|
|
|
def test_project_name_from_cwd_returns_basename(
|
|
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
|
) -> None:
|
|
project_dir = tmp_path / "vibe-headroom"
|
|
project_dir.mkdir()
|
|
monkeypatch.chdir(project_dir)
|
|
|
|
assert wrap_mod._project_name_from_cwd() == "vibe-headroom"
|
|
|
|
def test_non_ascii_cwd_name_is_percent_encoded(
|
|
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
|
) -> None:
|
|
"""Non-ASCII directory names must be percent-encoded for HTTP headers."""
|
|
project_dir = tmp_path / "第二大脑共享"
|
|
project_dir.mkdir()
|
|
monkeypatch.chdir(project_dir)
|
|
|
|
result = wrap_mod._project_name_from_cwd()
|
|
assert result is not None
|
|
# Must be pure ASCII so it's safe in an HTTP header value.
|
|
result.encode("ascii")
|
|
# Must round-trip back to the original name via unquote.
|
|
import urllib.parse
|
|
|
|
assert urllib.parse.unquote(result) == "第二大脑共享"
|
|
|
|
def test_non_ascii_cwd_header_is_ascii_safe(
|
|
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
|
) -> None:
|
|
"""X-Headroom-Project header value must be ASCII when cwd has non-ASCII chars."""
|
|
project_dir = tmp_path / "test-中文-项目"
|
|
project_dir.mkdir()
|
|
monkeypatch.chdir(project_dir)
|
|
|
|
env: dict[str, str] = {}
|
|
wrap_mod._apply_project_header_env(env)
|
|
|
|
header_value = env["ANTHROPIC_CUSTOM_HEADERS"]
|
|
assert header_value.startswith("X-Headroom-Project: ")
|
|
header_value.encode("ascii") # raises UnicodeEncodeError if non-ASCII
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Proxy-client reference counting
|
|
#
|
|
# The shared proxy must only be torn down by its owner once *no* other live
|
|
# wrap clients remain. Clients carry the proxy URL in ANTHROPIC_BASE_URL /
|
|
# OPENAI_BASE_URL (env, not argv), so the old `pgrep -f "127.0.0.1:<port>"`
|
|
# guard could neither see real clients nor reject unrelated processes that
|
|
# merely had the address in their command line. These tests pin the new
|
|
# marker-file contract: a per-PID file under paths.proxy_clients_dir(port).
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _FakeProxyProc:
|
|
"""Minimal stand-in for the proxy ``subprocess.Popen`` handle."""
|
|
|
|
def __init__(self) -> None:
|
|
self.terminated = False
|
|
self.killed = False
|
|
|
|
def poll(self) -> int | None:
|
|
return None # alive
|
|
|
|
def terminate(self) -> None:
|
|
self.terminated = True
|
|
|
|
def wait(self, timeout: float | None = None) -> int:
|
|
return 0
|
|
|
|
def kill(self) -> None:
|
|
self.killed = True
|
|
|
|
|
|
class TestProxyClientRefCounting:
|
|
"""Proxy lifecycle is reference-counted via marker files, not pgrep."""
|
|
|
|
PORT = 8787
|
|
|
|
@pytest.fixture
|
|
def clients_dir(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
|
|
"""Redirect ``paths.proxy_clients_dir`` into a throwaway tmp tree."""
|
|
base = tmp_path / "clients"
|
|
monkeypatch.setattr(paths_mod, "proxy_clients_dir", lambda port: base / str(port))
|
|
return base
|
|
|
|
def _write_marker(
|
|
self,
|
|
clients_dir: Path,
|
|
pid: int,
|
|
*,
|
|
identity: tuple[str, float] | None = None,
|
|
) -> Path:
|
|
marker = clients_dir / str(self.PORT) / f"{pid}.json"
|
|
marker.parent.mkdir(parents=True, exist_ok=True)
|
|
rec: dict[str, Any] = {"pid": pid, "started_at": 0}
|
|
if identity is not None:
|
|
rec["start_src"], rec["start_time"] = identity
|
|
marker.write_text(json.dumps(rec))
|
|
return marker
|
|
|
|
def test_cleanup_terminates_proxy_when_only_self_registered(self, clients_dir: Path) -> None:
|
|
"""The owner alone → no other clients → proxy is terminated on exit."""
|
|
wrap_mod._register_proxy_client(self.PORT)
|
|
proc = _FakeProxyProc()
|
|
cleanup = wrap_mod._make_cleanup([proc], self.PORT)
|
|
|
|
cleanup()
|
|
|
|
assert proc.terminated is True
|
|
# Our own marker is removed before we count.
|
|
assert wrap_mod._live_proxy_clients(self.PORT, exclude_self=False) == []
|
|
|
|
def test_cleanup_leaves_proxy_running_when_other_client_alive(self, clients_dir: Path) -> None:
|
|
"""A second live client (here: the test's parent) keeps the proxy up."""
|
|
wrap_mod._register_proxy_client(self.PORT)
|
|
other_pid = os.getppid() # alive for the duration of the test run
|
|
assert other_pid != os.getpid()
|
|
self._write_marker(clients_dir, other_pid)
|
|
|
|
proc = _FakeProxyProc()
|
|
cleanup = wrap_mod._make_cleanup([proc], self.PORT)
|
|
cleanup()
|
|
|
|
assert proc.terminated is False
|
|
|
|
def test_dead_client_marker_is_pruned_and_not_counted(self, clients_dir: Path) -> None:
|
|
"""A marker for a dead PID is pruned from disk and never counted."""
|
|
# Spawn and reap a child so its PID is reliably dead (not a zombie).
|
|
child = subprocess.Popen([sys.executable, "-c", "pass"])
|
|
child.wait()
|
|
dead_pid = child.pid
|
|
marker = self._write_marker(clients_dir, dead_pid)
|
|
|
|
live = wrap_mod._live_proxy_clients(self.PORT, exclude_self=True)
|
|
|
|
assert dead_pid not in live
|
|
assert not marker.exists()
|
|
|
|
def test_reused_pid_with_mismatched_identity_is_pruned(
|
|
self, clients_dir: Path, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
"""A *live* PID that the original client no longer owns is pruned.
|
|
|
|
Models the PID-reuse orphan path: a wrapper crashed, the OS later
|
|
recycled its PID for an unrelated long-lived process. `os.kill(pid, 0)`
|
|
succeeds, but the recorded start time no longer matches.
|
|
"""
|
|
live_pid = os.getppid() # alive, but not the process that "registered"
|
|
marker = self._write_marker(clients_dir, live_pid, identity=("psutil", 1000.0))
|
|
# The process currently holding that PID started much later → reuse.
|
|
monkeypatch.setattr(wrap_mod, "_proc_identity", lambda p: ("psutil", 9000.0))
|
|
|
|
live = wrap_mod._live_proxy_clients(self.PORT, exclude_self=True)
|
|
|
|
assert live_pid not in live
|
|
assert not marker.exists()
|
|
|
|
def test_matching_identity_within_tolerance_is_kept(
|
|
self, clients_dir: Path, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
"""Same process (start time within tolerance) is a real client — kept."""
|
|
live_pid = os.getppid()
|
|
marker = self._write_marker(clients_dir, live_pid, identity=("psutil", 1000.0))
|
|
monkeypatch.setattr(wrap_mod, "_proc_identity", lambda p: ("psutil", 1000.4))
|
|
|
|
live = wrap_mod._live_proxy_clients(self.PORT, exclude_self=True)
|
|
|
|
assert live_pid in live
|
|
assert marker.exists()
|
|
|
|
def test_identity_check_skipped_when_source_unavailable(
|
|
self, clients_dir: Path, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
"""No reuse protection (e.g. macOS w/o psutil) → fall back to existence."""
|
|
live_pid = os.getppid()
|
|
self._write_marker(clients_dir, live_pid, identity=("psutil", 1000.0))
|
|
# Start time unknowable for the live PID → must not prune a real client.
|
|
monkeypatch.setattr(wrap_mod, "_proc_identity", lambda p: None)
|
|
|
|
live = wrap_mod._live_proxy_clients(self.PORT, exclude_self=True)
|
|
|
|
assert live_pid in live
|
|
|
|
def test_non_marker_files_are_ignored(self, clients_dir: Path) -> None:
|
|
"""Stray non-numeric / non-json files don't crash or count as clients."""
|
|
d = clients_dir / str(self.PORT)
|
|
d.mkdir(parents=True, exist_ok=True)
|
|
(d / "not-a-pid.json").write_text("{}")
|
|
(d / "README.txt").write_text("ignore me")
|
|
|
|
assert wrap_mod._live_proxy_clients(self.PORT, exclude_self=True) == []
|
|
|
|
def test_cleanup_does_not_shell_out_to_pgrep(
|
|
self, clients_dir: Path, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
"""Regression: liveness is never inferred from an argv scan.
|
|
|
|
An unrelated process whose command line contains ``127.0.0.1:8787``
|
|
used to be a false positive (orphaning the proxy). The new path never
|
|
calls ``subprocess.run`` at all, so it can't be fooled by argv.
|
|
"""
|
|
wrap_mod._register_proxy_client(self.PORT)
|
|
|
|
def _no_subprocess(*args: Any, **kwargs: Any) -> None:
|
|
raise AssertionError("cleanup must not invoke subprocess.run (no pgrep)")
|
|
|
|
monkeypatch.setattr(wrap_mod.subprocess, "run", _no_subprocess)
|
|
|
|
proc = _FakeProxyProc()
|
|
cleanup = wrap_mod._make_cleanup([proc], self.PORT)
|
|
cleanup() # must not raise
|
|
|
|
assert proc.terminated is True
|
|
|
|
def test_register_then_unregister_is_idempotent(self, clients_dir: Path) -> None:
|
|
"""Register adds exactly our marker; unregister removes it; re-call is safe."""
|
|
wrap_mod._register_proxy_client(self.PORT)
|
|
all_clients = wrap_mod._live_proxy_clients(self.PORT, exclude_self=False)
|
|
assert all_clients == [os.getpid()]
|
|
|
|
wrap_mod._unregister_proxy_client(self.PORT)
|
|
assert wrap_mod._live_proxy_clients(self.PORT, exclude_self=False) == []
|
|
|
|
# Second unregister is a no-op, not an error.
|
|
wrap_mod._unregister_proxy_client(self.PORT)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _ensure_proxy — dashboard URL is surfaced even when the proxy is already up.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_ensure_proxy_already_running_prints_dashboard_url(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""When a healthy proxy is already running, the dashboard URL is printed.
|
|
|
|
Regression: the URL was only echoed on the start/restart path, so repeat
|
|
wraps (the common case) never told the user where the dashboard lives.
|
|
"""
|
|
port = 1234
|
|
monkeypatch.setattr(wrap_mod, "_find_persistent_manifest", lambda _p: None)
|
|
monkeypatch.setattr(wrap_mod, "_check_proxy", lambda _p: True)
|
|
monkeypatch.setattr(wrap_mod, "_query_proxy_health", lambda _p: {})
|
|
monkeypatch.setattr(wrap_mod, "_proxy_needs_version_restart", lambda _h: False)
|
|
monkeypatch.setattr(wrap_mod, "_proxy_health_config", lambda _h: None)
|
|
monkeypatch.setattr(wrap_mod, "_query_proxy_config", lambda _p: None)
|
|
|
|
output = _run_in_click_context(lambda: wrap_mod._ensure_proxy(port, no_proxy=False))
|
|
|
|
assert f"http://127.0.0.1:{port}/dashboard" in output
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _resolve_1m_model — 1M context window suffix logic (#1158).
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_resolve_1m_model_appends_suffix_to_user_model() -> None:
|
|
"""A model the user already selected via ANTHROPIC_MODEL is preserved, with
|
|
only the [1m] suffix appended so Claude Code requests the 1M window."""
|
|
assert wrap_mod._resolve_1m_model("claude-opus-4-1-20250805") == (
|
|
"claude-opus-4-1-20250805[1m]"
|
|
)
|
|
|
|
|
|
def test_resolve_1m_model_is_idempotent() -> None:
|
|
"""A model that already carries [1m] is returned unchanged (no double suffix)."""
|
|
assert wrap_mod._resolve_1m_model("claude-opus-4-8[1m]") == "claude-opus-4-8[1m]"
|
|
|
|
|
|
def test_resolve_1m_model_falls_back_to_default_when_unset() -> None:
|
|
"""With no model selected, fall back to the default Opus carrying [1m]."""
|
|
assert wrap_mod._resolve_1m_model(None) == "claude-opus-4-8[1m]"
|
|
assert wrap_mod._resolve_1m_model(" ") == "claude-opus-4-8[1m]"
|
|
|
|
|
|
class TestFindAvailablePort:
|
|
"""Tests for _find_available_port (Vite-style port fallback)."""
|
|
|
|
def test_port_free_returns_same(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""When port is free, returns the same port."""
|
|
monkeypatch.setattr(wrap_mod, "_port_bind_error", lambda port: None)
|
|
assert wrap_mod._find_available_port(8787) == 8787
|
|
|
|
def test_port_busy_finds_next(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""When port is busy, returns the next free port."""
|
|
|
|
def mock_bind(port: int) -> OSError | None:
|
|
if port == 8787:
|
|
return OSError(errno.EADDRINUSE, "Address in use")
|
|
return None
|
|
|
|
monkeypatch.setattr(wrap_mod, "_port_bind_error", mock_bind)
|
|
assert wrap_mod._find_available_port(8787) == 8788
|
|
|
|
def test_multiple_busy_ports(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""When multiple consecutive ports are busy, skips all of them."""
|
|
|
|
def mock_bind(port: int) -> OSError | None:
|
|
if port in (8787, 8788, 8789):
|
|
return OSError(errno.EADDRINUSE, "Address in use")
|
|
return None
|
|
|
|
monkeypatch.setattr(wrap_mod, "_port_bind_error", mock_bind)
|
|
assert wrap_mod._find_available_port(8787) == 8790
|
|
|
|
def test_propagates_unexpected_error(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""Errors other than EADDRINUSE/EACCES (e.g. EADDRNOTAVAIL) propagate."""
|
|
monkeypatch.setattr(
|
|
wrap_mod,
|
|
"_port_bind_error",
|
|
lambda port: OSError(errno.EADDRNOTAVAIL, "Address not available"),
|
|
)
|
|
with pytest.raises(OSError, match="Address not available"):
|
|
wrap_mod._find_available_port(8787)
|
|
|
|
def test_propagates_eaddrinuse_with_eacces(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""Both EADDRINUSE and EACCES are skipped (not propagated)."""
|
|
|
|
def mock_bind(port: int) -> OSError | None:
|
|
if port == 8787:
|
|
return OSError(errno.EACCES, "Permission denied")
|
|
return None
|
|
|
|
monkeypatch.setattr(wrap_mod, "_port_bind_error", mock_bind)
|
|
assert wrap_mod._find_available_port(8787) == 8788
|
|
|
|
def test_exhausts_range(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""When all ports in range are busy, raises RuntimeError."""
|
|
monkeypatch.setattr(
|
|
wrap_mod,
|
|
"_port_bind_error",
|
|
lambda port: OSError(errno.EADDRINUSE, "Address in use"),
|
|
)
|
|
with pytest.raises(RuntimeError, match="No available port found"):
|
|
wrap_mod._find_available_port(8787, max_attempts=3)
|