mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix: remove rtk and lean-ctx CLI context tools (#2677)
## 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.
This commit is contained in:
parent
759209cff3
commit
e0ce4b1d48
101 changed files with 1546 additions and 8746 deletions
1
.github/workflows/ci.yml
vendored
1
.github/workflows/ci.yml
vendored
|
|
@ -70,7 +70,6 @@ jobs:
|
|||
- 'headroom/cli/**'
|
||||
- 'headroom/install/**'
|
||||
- 'headroom/providers/**'
|
||||
- 'headroom/rtk/**'
|
||||
- 'crates/**'
|
||||
- '**/*.rs'
|
||||
- 'Cargo.toml'
|
||||
|
|
|
|||
1
.github/workflows/wrap-e2e.yml
vendored
1
.github/workflows/wrap-e2e.yml
vendored
|
|
@ -9,7 +9,6 @@ on:
|
|||
paths:
|
||||
- 'headroom/cli/**'
|
||||
- 'headroom/providers/**'
|
||||
- 'headroom/rtk/**'
|
||||
- 'crates/**'
|
||||
- 'docker/**'
|
||||
- 'Dockerfile'
|
||||
|
|
|
|||
1
.github/workflows/wrap-native-e2e.yml
vendored
1
.github/workflows/wrap-native-e2e.yml
vendored
|
|
@ -16,7 +16,6 @@ on:
|
|||
paths:
|
||||
- "headroom/cli/**"
|
||||
- "headroom/providers/**"
|
||||
- "headroom/rtk/**"
|
||||
- "tests/test_cli/test_wrap_bridge.py"
|
||||
- ".github/actions/headroom-e2e-setup/**"
|
||||
- ".github/workflows/wrap-native-e2e.yml"
|
||||
|
|
|
|||
|
|
@ -533,12 +533,10 @@ Headroom runs **locally**, covers **every** content type, works with every major
|
|||
| | Scope | Deploy | Local | Reversible |
|
||||
|------------------------------------------------------------------------------|------------------------------------------------|------------------------------------|:-----:|:----------:|
|
||||
| **Headroom** | All context — tools, RAG, logs, files, history | Proxy · library · middleware · MCP | Yes | Yes |
|
||||
| [RTK](https://github.com/rtk-ai/rtk) | CLI command outputs | CLI wrapper | Yes | No |
|
||||
| [lean-ctx](https://github.com/yvgude/lean-ctx) | Tool output, files, shell, history | Proxy · library · middleware · MCP · CLI | Yes | Yes |
|
||||
| [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 |
|
||||
|
||||
> **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.
|
||||
> **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: you're free to attach your own tooling — code-memory MCP, Graphify, Caveman, or any MCP server — and Headroom compresses downstream of all of it.
|
||||
|
||||
## Contributing
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,14 @@
|
|||
# Phase G — RTK Breadth + Observability
|
||||
|
||||
> **SUPERSEDED.** RTK and lean-ctx were removed from Headroom entirely: the
|
||||
> `headroom/rtk/` and `headroom/lean_ctx/` packages, all `--rtk` / `--context-tool`
|
||||
> flags, the wrap-side hooks and hint-file injection, and the proxy-side `rtk gain`
|
||||
> polling are all gone, and `headroom/context_tool_cleanup.py` uninstalls what
|
||||
> earlier versions left on disk. The RTK-specific plan below is historical; the
|
||||
> non-RTK observability items (cache-hit rate, compression ratio, token
|
||||
> validation) were kept. `docs/rtk-architecture.md`, referenced throughout this
|
||||
> document, was deleted with the feature.
|
||||
|
||||
**Goal:** Extend RTK coverage to more wrap-CLI agents; close the dead `tokens_saved_rtk` data plane; add per-invocation RTK metrics; add the cache-hit-rate, compression-ratio, token-validation observability surface that's missing today.
|
||||
|
||||
**Calendar:** 1 week.
|
||||
|
|
|
|||
|
|
@ -92,7 +92,11 @@ No re-scoping needed; revisit after Phase D lands.
|
|||
|
||||
## Q9. RTK proxy-side invocation — ever revisit?
|
||||
|
||||
**Recommendation:** **No, document the decision in `docs/rtk-architecture.md`** (Phase G PR-G3). The argument:
|
||||
**Resolved — moot.** RTK was removed from Headroom outright (see
|
||||
`09-phase-G-rtk-observability.md`), so there is no proxy-side invocation to
|
||||
revisit. The original recommendation was "no, document the decision in
|
||||
`docs/rtk-architecture.md`" (that doc was deleted with the feature). The argument
|
||||
is kept because reasons 1–3 apply to any future shell-output rewriter:
|
||||
1. Cache hot zone risk: shell-out + buffer per tool result is correctness-fragile.
|
||||
2. Parallel implementation: `crates/headroom-core/src/transforms/log_compressor.rs` covers post-hoc log/output compression; RTK rewrites *commands* (different value).
|
||||
3. RTK itself is a third-party binary the team doesn't control; an upstream version change silently busts cache.
|
||||
|
|
|
|||
|
|
@ -1,287 +0,0 @@
|
|||
"""RTK-loop eval — does Headroom Learn catch a loop and write a guardrail that
|
||||
would prevent it recurring?
|
||||
|
||||
This is the agentic eval for the loop-weighting work. It runs in two phases:
|
||||
|
||||
Phase 1 — TRIGGER + LEARN
|
||||
Reproduce an RTK re-fetch loop (a grep whose RTK-truncated output forces the
|
||||
agent to re-run larger-limit variants), run it through ``SessionAnalyzer``,
|
||||
and SCORE the resulting guardrail:
|
||||
• produced — a loop guardrail was emitted at all
|
||||
• ranked_first — it outranks the one-off rules (the weighting works)
|
||||
• names_command — the rule identifies the command that looped
|
||||
• prescribes_fix — the rule says how to avoid it (fetch full output once)
|
||||
• weight_reflects — its savings estimate >= the MEASURED wasted tokens
|
||||
|
||||
Phase 2 — GUARDRAIL HOLDS
|
||||
Inject that guardrail as a prior learned pattern, then feed a session where
|
||||
the agent FOLLOWED it (one full-output fetch, no loop). Re-run the analyzer
|
||||
and assert NO new loop guardrail is produced for that command — i.e. once
|
||||
the rule exists and is honored, the loop does not re-trigger and Learn does
|
||||
not need to relearn it.
|
||||
|
||||
Runs deterministically by default (a stubbed analyzer LLM so CI is hermetic).
|
||||
With ``--real`` it drives the real analyzer LLM and scores the actually-generated
|
||||
rule, using an API key (ANTHROPIC/OPENAI/GEMINI) or an installed CLI backend.
|
||||
|
||||
Usage:
|
||||
python benchmarks/rtk_loop_learn_eval.py # deterministic
|
||||
python benchmarks/rtk_loop_learn_eval.py --real # real LLM (API key)
|
||||
HEADROOM_LEARN_CLI=claude python benchmarks/rtk_loop_learn_eval.py --real # via CLI
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
from contextlib import nullcontext
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
# Allow running as a plain script from the repo root.
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
from headroom.learn.analyzer import SessionAnalyzer # noqa: E402
|
||||
from headroom.learn.fixtures import rtk_refetch_loop_session # noqa: E402
|
||||
from headroom.learn.loops import detect_loops # noqa: E402
|
||||
from headroom.learn.models import ( # noqa: E402
|
||||
ProjectInfo,
|
||||
SessionData,
|
||||
ToolCall,
|
||||
)
|
||||
|
||||
REPETITIONS = 6
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Deterministic LLM stub — stands in for the analyzer's _call_llm in CI.
|
||||
# It mimics a competent model: emits the loop guardrail (under-estimating its
|
||||
# savings, so the weighting layer has real work to do) plus a one-off rule the
|
||||
# model would naively rank higher. In Phase 2 it emits NO loop rule, because a
|
||||
# non-looping guarded session gives it nothing to relearn.
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _stub_llm_phase1(digest: str, model: str) -> dict:
|
||||
return {
|
||||
"context_file_rules": [
|
||||
{
|
||||
"section": "Use uv for Python",
|
||||
"content": "Use `uv run python` instead of `python3`.",
|
||||
"estimated_tokens_saved": 900, # model rates the one-off high
|
||||
"evidence_count": 2,
|
||||
},
|
||||
{
|
||||
"section": "Avoid grep TimeoutError re-fetch loop",
|
||||
"content": (
|
||||
"When searching logs for TimeoutError, capture the full "
|
||||
"result once (grep into a file and read it) instead of "
|
||||
"re-running grep with larger `head` limits."
|
||||
),
|
||||
"estimated_tokens_saved": 150, # simulated low estimate (stub value, not a real-model figure)
|
||||
"evidence_count": 1,
|
||||
},
|
||||
],
|
||||
"memory_file_rules": [],
|
||||
}
|
||||
|
||||
|
||||
def _stub_llm_phase2(digest: str, model: str) -> dict:
|
||||
# Guarded, non-looping session → nothing new to learn about the grep.
|
||||
return {"context_file_rules": [], "memory_file_rules": []}
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Scoring
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@dataclass
|
||||
class Scorecard:
|
||||
checks: dict[str, bool] = field(default_factory=dict)
|
||||
notes: dict[str, str] = field(default_factory=dict)
|
||||
|
||||
def add(self, name: str, passed: bool, note: str = "") -> None:
|
||||
self.checks[name] = passed
|
||||
if note:
|
||||
self.notes[name] = note
|
||||
|
||||
@property
|
||||
def passed(self) -> bool:
|
||||
return all(self.checks.values())
|
||||
|
||||
def render(self) -> str:
|
||||
width = max(len(k) for k in self.checks)
|
||||
lines = []
|
||||
for name, ok in self.checks.items():
|
||||
mark = "PASS" if ok else "FAIL"
|
||||
note = f" ({self.notes[name]})" if name in self.notes else ""
|
||||
lines.append(f" [{mark}] {name.ljust(width)}{note}")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _guarded_session() -> SessionData:
|
||||
"""A session where the agent followed the guardrail: one full-output fetch,
|
||||
no re-fetch loop."""
|
||||
return SessionData(
|
||||
session_id="guarded",
|
||||
tool_calls=[
|
||||
ToolCall(
|
||||
name="Bash",
|
||||
tool_call_id="tc_0",
|
||||
input_data={"command": "grep -rn 'TimeoutError' logs/ > /tmp/hits.txt"},
|
||||
output="(wrote 1240 matches to /tmp/hits.txt)",
|
||||
is_error=False,
|
||||
msg_index=0,
|
||||
output_bytes=40,
|
||||
),
|
||||
ToolCall(
|
||||
name="Read",
|
||||
tool_call_id="tc_1",
|
||||
input_data={"file_path": "/tmp/hits.txt"},
|
||||
output="logs/app.log:42: TimeoutError ...",
|
||||
is_error=False,
|
||||
msg_index=1,
|
||||
output_bytes=8000,
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def run_eval(*, use_real_llm: bool) -> Scorecard:
|
||||
project = ProjectInfo(
|
||||
name="rtk-loop-eval",
|
||||
project_path=Path("/tmp/rtk-loop-eval"),
|
||||
data_path=Path("/tmp/rtk-loop-eval-data"),
|
||||
)
|
||||
card = Scorecard()
|
||||
|
||||
# ---- Phase 1: trigger + learn -----------------------------------------
|
||||
loop_session = rtk_refetch_loop_session(repetitions=REPETITIONS)
|
||||
loops = detect_loops([loop_session])
|
||||
measured_waste = loops[0].wasted_tokens if loops else 0
|
||||
card.add("loop_detected", bool(loops), f"{len(loops)} loop(s), ~{measured_waste:,} tok wasted")
|
||||
|
||||
analyzer = SessionAnalyzer(model=None if use_real_llm else "stub")
|
||||
phase1_ctx = (
|
||||
nullcontext()
|
||||
if use_real_llm
|
||||
else patch("headroom.learn.analyzer._call_llm", _stub_llm_phase1)
|
||||
)
|
||||
with phase1_ctx:
|
||||
result = analyzer.analyze(project, [loop_session])
|
||||
|
||||
recs = result.recommendations
|
||||
loop_recs = [r for r in recs if r.is_loop_guardrail]
|
||||
card.add("guardrail_produced", bool(loop_recs))
|
||||
|
||||
top = recs[0] if recs else None
|
||||
card.add(
|
||||
"ranked_first",
|
||||
bool(top and top.is_loop_guardrail),
|
||||
"" if (top and top.is_loop_guardrail) else "loop rule did not rank #1",
|
||||
)
|
||||
|
||||
guardrail = loop_recs[0] if loop_recs else None
|
||||
text = (guardrail.section + " " + guardrail.content).lower() if guardrail else ""
|
||||
# The rule must identify the LOOPING COMMAND (grep + its output-limit shape),
|
||||
# not the incidental search string — a good fix generalizes beyond it. (The
|
||||
# real-LLM run surfaced this: the model wrote a general "grepping logs / `head
|
||||
# -N` limits" rule and never echoed "TimeoutError", which an earlier
|
||||
# literal-match check wrongly failed.)
|
||||
card.add(
|
||||
"names_command",
|
||||
"grep" in text and any(k in text for k in ("head", "log", "limit")),
|
||||
)
|
||||
card.add(
|
||||
"prescribes_fix",
|
||||
any(k in text for k in ("full", "once", "into a file", "instead", "limit")),
|
||||
)
|
||||
card.add(
|
||||
"weight_reflects_waste",
|
||||
bool(guardrail and guardrail.estimated_tokens_saved >= measured_waste),
|
||||
""
|
||||
if (guardrail and guardrail.estimated_tokens_saved >= measured_waste)
|
||||
else f"savings {getattr(guardrail, 'estimated_tokens_saved', 0)} < waste {measured_waste}",
|
||||
)
|
||||
|
||||
# ---- Phase 2: guardrail holds -----------------------------------------
|
||||
# Inject the produced guardrail as a prior pattern via the project's
|
||||
# context file, then analyze a guarded (non-looping) session.
|
||||
held = True
|
||||
note = ""
|
||||
if guardrail:
|
||||
ctx_path = Path("/tmp/rtk-loop-eval-CLAUDE.md")
|
||||
ctx_path.write_text(
|
||||
"<!-- headroom:learn:start -->\n"
|
||||
f"### {guardrail.section}\n{guardrail.content}\n"
|
||||
"<!-- headroom:learn:end -->\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
project.context_file = ctx_path
|
||||
phase2_ctx = (
|
||||
nullcontext()
|
||||
if use_real_llm
|
||||
else patch("headroom.learn.analyzer._call_llm", _stub_llm_phase2)
|
||||
)
|
||||
with phase2_ctx:
|
||||
held_result = analyzer.analyze(project, [_guarded_session()])
|
||||
# No NEW loop guardrail should be needed for the (now-guarded) grep.
|
||||
new_loop_rules = [
|
||||
r
|
||||
for r in held_result.recommendations
|
||||
if r.is_loop_guardrail and "grep" in (r.section + r.content).lower()
|
||||
]
|
||||
held = not new_loop_rules
|
||||
note = "" if held else f"{len(new_loop_rules)} new grep loop rule(s) re-emitted"
|
||||
else:
|
||||
held = False
|
||||
note = "no guardrail from phase 1 to test"
|
||||
card.add("guardrail_holds", held, note)
|
||||
|
||||
return card
|
||||
|
||||
|
||||
def _real_backend_available() -> bool:
|
||||
"""True when the analyzer can reach a real LLM — API key or installed CLI."""
|
||||
import shutil
|
||||
|
||||
if any(os.environ.get(k) for k in ("ANTHROPIC_API_KEY", "OPENAI_API_KEY", "GEMINI_API_KEY")):
|
||||
return True
|
||||
return any(shutil.which(cli) for cli in ("claude", "gemini", "codex"))
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--real",
|
||||
action="store_true",
|
||||
help="Drive the real analyzer LLM — needs an API key (ANTHROPIC_API_KEY / "
|
||||
"OPENAI_API_KEY / GEMINI_API_KEY) or an installed CLI backend "
|
||||
"(claude / gemini / codex; force one with HEADROOM_LEARN_CLI=claude).",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.real and not _real_backend_available():
|
||||
print(
|
||||
"--real needs an LLM backend (API key or claude/gemini/codex CLI); "
|
||||
"falling back to deterministic mode.\n"
|
||||
)
|
||||
args.real = False
|
||||
|
||||
mode = "REAL LLM" if args.real else "deterministic stub"
|
||||
print(f"RTK-loop eval — mode: {mode}\n")
|
||||
card = run_eval(use_real_llm=args.real)
|
||||
print(card.render())
|
||||
print()
|
||||
if card.passed:
|
||||
print("RESULT: PASS — loop caught, guardrail ranked first, and it holds.")
|
||||
return 0
|
||||
print("RESULT: FAIL — see failed checks above.")
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
|
|
@ -98,11 +98,11 @@ pub const METRIC_PROXY_RESPONSE_STATUS_COUNT_TOTAL_HELP: &str =
|
|||
// for `proxy_image_generation_call_log_redacted_total`,
|
||||
// `wrap_rtk_invocations_total`, and `wrap_rtk_tokens_saved_per_session`
|
||||
// were removed because the underlying counters had no production
|
||||
// emit site on the Rust side. The same metrics are exported by the
|
||||
// Python proxy (`headroom/proxy/prometheus_metrics.py`) which is the
|
||||
// natural owner: image redaction is a Python-proxy operation and RTK
|
||||
// invocation tracking lives in the wrap CLI, both Python-side
|
||||
// surfaces. See `docs/observability.md`.
|
||||
// emit site on the Rust side. Image redaction is exported by the
|
||||
// Python proxy (`headroom/proxy/prometheus_metrics.py`), its natural
|
||||
// owner. The two `wrap_rtk_*` names are gone for good: the rtk
|
||||
// integration they measured has been removed from Headroom.
|
||||
// See `docs/observability.md`.
|
||||
|
||||
// ---------- shared label keys ----------
|
||||
|
||||
|
|
|
|||
|
|
@ -283,10 +283,10 @@ pub fn record_response_status(status: &str, reason: Option<&str>, request_id: &s
|
|||
// Phase G PR-G3 remediation (C3 + C4): the image-redacted counter
|
||||
// and the wrap_rtk_invocations counter were originally registered
|
||||
// here but neither had a production emit site that crossed the
|
||||
// Python/Rust boundary. Both have moved Python-side
|
||||
// (`headroom.proxy.request_logger::redactions_total` and
|
||||
// `headroom.cli.wrap_rtk_metrics::rtk_invocation_counts`) and the
|
||||
// Python proxy's `/metrics` exporter surfaces them — see
|
||||
// Python/Rust boundary. The image-redacted counter moved Python-side
|
||||
// (`headroom.proxy.request_logger::redactions_total`) and the Python
|
||||
// proxy's `/metrics` exporter surfaces it; the RTK counter is gone
|
||||
// entirely along with the rtk integration itself — see
|
||||
// `docs/observability.md` for the placement decision. Keeping a
|
||||
// dead Rust counter would (a) violate the "no dead metrics
|
||||
// registered" review finding and (b) mislead Phase H canary
|
||||
|
|
|
|||
|
|
@ -38,8 +38,6 @@ services:
|
|||
# if you want to use a custom OpenAI-compatible API endpoint,
|
||||
# uncomment and set the following line with the desired URL
|
||||
# - OPENAI_TARGET_API_URL=https://api.x.ai
|
||||
# CLI-filtering dashboard figures require the `rtk` binary inside this
|
||||
# container; it is not installed by this image. See docs/content/docs/docker-install.mdx.
|
||||
ports:
|
||||
- "8787:8787"
|
||||
volumes:
|
||||
|
|
|
|||
|
|
@ -13,8 +13,6 @@ services:
|
|||
# the Docker bind-mount source and is intentionally different.
|
||||
HEADROOM_WORKSPACE_DIR: /tmp/headroom-home/.headroom
|
||||
HEADROOM_CONFIG_DIR: /tmp/headroom-home/.headroom/config
|
||||
# CLI-filtering dashboard figures require the `rtk` binary inside this
|
||||
# container; it is not installed by this image. See docs/content/docs/docker-install.mdx.
|
||||
volumes:
|
||||
- ${HEADROOM_WORKSPACE:-.}:/workspace
|
||||
- ${HEADROOM_HOST_HOME:?set HEADROOM_HOST_HOME}/.headroom:/tmp/headroom-home/.headroom
|
||||
|
|
@ -35,8 +33,6 @@ services:
|
|||
# above for rationale.
|
||||
HEADROOM_WORKSPACE_DIR: /tmp/headroom-home/.headroom
|
||||
HEADROOM_CONFIG_DIR: /tmp/headroom-home/.headroom/config
|
||||
# CLI-filtering dashboard figures require the `rtk` binary inside this
|
||||
# container; it is not installed by this image. See docs/content/docs/docker-install.mdx.
|
||||
ports:
|
||||
- "${HEADROOM_PORT:-8787}:${HEADROOM_PORT:-8787}"
|
||||
volumes:
|
||||
|
|
|
|||
|
|
@ -5,27 +5,8 @@ description: All configuration options for the Headroom Python and TypeScript SD
|
|||
|
||||
Headroom can be configured via the SDK constructor, proxy command line, environment variables, or per-request overrides.
|
||||
|
||||
## CLI Context Tool
|
||||
|
||||
`headroom wrap ...` uses RTK for local shell-output filtering by default.
|
||||
Set `HEADROOM_CONTEXT_TOOL=lean-ctx` to have wrap commands install or reuse
|
||||
`lean-ctx` and run `lean-ctx init --agent <tool>` instead of RTK setup.
|
||||
|
||||
```bash
|
||||
export HEADROOM_CONTEXT_TOOL=lean-ctx
|
||||
headroom wrap claude
|
||||
headroom wrap codex --prepare-only
|
||||
```
|
||||
|
||||
Supported values are `rtk` and `lean-ctx`; unset defaults to `rtk`.
|
||||
|
||||
If Codex history disappeared after using an older wrapper, see [Recover Codex State](/docs/codex-recovery) before wrapping Codex again.
|
||||
|
||||
The proxy reads RTK lifetime savings with global scope by default so a shared
|
||||
daemon reports savings across the operator's projects. Set
|
||||
`HEADROOM_RTK_GAIN_SCOPE=project` to query `rtk gain --project` from the
|
||||
proxy process working directory.
|
||||
|
||||
## SDK Modes (`default_mode` / `headroom_mode`)
|
||||
|
||||
These modes apply to SDK usage via `HeadroomClient(default_mode=...)` or per-request `headroom_mode=...`. They are **not** the same as the proxy `--mode` flag.
|
||||
|
|
|
|||
|
|
@ -75,7 +75,7 @@ docker run --rm -it \
|
|||
`wrap` is host-oriented in Docker-native mode:
|
||||
|
||||
- the wrapper starts the Headroom proxy in Docker
|
||||
- container-side prep writes Headroom config, memory, and selected CLI context-tool setup into mounted host files
|
||||
- container-side prep writes Headroom config and memory into mounted host files
|
||||
- the target CLI itself is launched on the host by the wrapper
|
||||
|
||||
Supported host wrap flows:
|
||||
|
|
@ -191,9 +191,6 @@ That keeps provider auth and runtime config working without maintaining a separa
|
|||
- The install scripts are idempotent: rerunning them refreshes the wrapper and image without duplicating shell profile blocks.
|
||||
- For persistent service and task installs, use the Python-native `headroom install ...` workflow — see [Persistent Installs](/docs/persistent-installs).
|
||||
- For Docker-native `headroom install ...`, the wrapper persists its profile manifest under `~/.headroom/deploy/<profile>/`.
|
||||
- The `rtk` binary is not bundled in the Docker image. Dashboard CLI-filtering
|
||||
savings figures show as "not installed" (not `0`) until `rtk` is installed
|
||||
inside the container.
|
||||
|
||||
## Next steps
|
||||
|
||||
|
|
|
|||
|
|
@ -54,8 +54,6 @@ Examples:
|
|||
| Memory bridge state | `${WORKSPACE_DIR}/bridge_state.json` | — |
|
||||
| Proxy log directory | `${WORKSPACE_DIR}/logs/` | — |
|
||||
| HTTP 400 debug dumps | `${WORKSPACE_DIR}/logs/debug_400/` | — |
|
||||
| Vendored `rtk` binary | `${WORKSPACE_DIR}/bin/rtk[.exe]` | — |
|
||||
| Vendored `lean-ctx` binary | `${WORKSPACE_DIR}/bin/lean-ctx[.exe]` | — |
|
||||
| Deployment profiles | `${WORKSPACE_DIR}/deploy/` | — |
|
||||
| Beacon lock file | `${WORKSPACE_DIR}/.beacon_lock_<port>` | — |
|
||||
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ title: Grok Build Integration
|
|||
description: Route Grok Build traffic through Headroom for token compression and per-project savings attribution.
|
||||
---
|
||||
|
||||
Use `headroom wrap grok-build` to route Grok Build LLM traffic through the local Headroom proxy. The wrapper starts or reuses the proxy, injects a reversible `[model.grok-build]` override into `~/.grok/config.toml` (or `$GROK_HOME/config.toml`), optionally sets up RTK or `lean-ctx`, and prints next steps for launching `grok`.
|
||||
Use `headroom wrap grok-build` to route Grok Build LLM traffic through the local Headroom proxy. The wrapper starts or reuses the proxy, injects a reversible `[model.grok-build]` override into `~/.grok/config.toml` (or `$GROK_HOME/config.toml`), and prints next steps for launching `grok`.
|
||||
|
||||
## Quick Start
|
||||
|
||||
|
|
@ -30,7 +30,6 @@ headroom unwrap grok-build
|
|||
| Proxy | Starts the Headroom proxy unless `--no-proxy` is set |
|
||||
| Model config | Writes or updates `[model.grok-build] base_url` in Grok's `config.toml`, pointing at `http://127.0.0.1:<port>/v1` (with optional `/p/<project>` prefix for savings attribution) |
|
||||
| Existing config | If you already have a `[model.grok-build]` table, Headroom rewrites `base_url` in place instead of appending a duplicate table (invalid TOML) |
|
||||
| Context tool | Injects RTK or `lean-ctx` guidance into project `AGENTS.md` unless `--no-context-tool` is set |
|
||||
| MCP install | `headroom mcp install` can register Headroom MCP via `GrokRegistrar` |
|
||||
| Backup | Snapshots `config.toml` to `config.toml.headroom-backup` before the first injection |
|
||||
|
||||
|
|
@ -39,7 +38,6 @@ headroom unwrap grok-build
|
|||
```bash
|
||||
headroom wrap grok-build \
|
||||
--port 8787 \ # Proxy port (default: 8787)
|
||||
--no-context-tool \ # Skip RTK / lean-ctx setup
|
||||
--no-proxy \ # Use an existing proxy instead of starting one
|
||||
--learn \ # Enable live traffic learning
|
||||
--memory # Enable persistent memory
|
||||
|
|
@ -50,7 +48,6 @@ headroom wrap grok-build \
|
|||
| Variable | Description |
|
||||
|---|---|
|
||||
| `GROK_HOME` | Override Grok config directory (default: `~/.grok`) |
|
||||
| `HEADROOM_CONTEXT_TOOL` | Set to `lean-ctx` to use lean-ctx instead of RTK |
|
||||
| `XAI_API_KEY` | Grok API key (also accepts `GROK_CODE_XAI_API_KEY`) |
|
||||
|
||||
## Persistent Install
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ title: OpenCode Integration
|
|||
description: Route OpenCode traffic through Headroom for token compression, MCP tools, and cached model access. One command to wrap, one to unwrap.
|
||||
---
|
||||
|
||||
Use `headroom wrap opencode` to route OpenCode LLM traffic through the Headroom proxy with a single command. The wrapper starts or reuses the proxy, writes OpenCode config, injects Headroom MCP tools, adds RTK context filtering, and launches OpenCode with the generated config.
|
||||
Use `headroom wrap opencode` to route OpenCode LLM traffic through the Headroom proxy with a single command. The wrapper starts or reuses the proxy, writes OpenCode config, injects Headroom MCP tools, and launches OpenCode with the generated config.
|
||||
|
||||
The `headroom-opencode` npm package also exports a native OpenCode plugin. The plugin can be used directly from OpenCode config when you want in-process transport interception plus the Headroom retrieve tool.
|
||||
|
||||
|
|
@ -27,7 +27,6 @@ headroom unwrap opencode
|
|||
| Provider injection | Writes a `headroom` provider using `@ai-sdk/openai-compatible` into `opencode.json`, pointing at `http://127.0.0.1:<port>/v1` |
|
||||
| Runtime env | Sets `OPENCODE_CONFIG_CONTENT` with provider, plugin, and optional local MCP config so OpenCode picks up Headroom at launch |
|
||||
| Provider compatibility | Leaves `OPENAI_BASE_URL` and `ANTHROPIC_BASE_URL` untouched so OpenCode `/connect` providers keep their own routing |
|
||||
| Context tool | Injects RTK (or `lean-ctx`) instructions into `~/.config/opencode/AGENTS.md` and project `AGENTS.md` |
|
||||
| MCP setup | Registers the Headroom MCP server (`headroom_compress`, `headroom_retrieve`, `headroom_stats`) |
|
||||
| Serena MCP | Optionally registers Serena code graph tools (`--no-serena` to skip) |
|
||||
| Backup | Snapshots `opencode.json` to `opencode.json.headroom-backup` before making any changes |
|
||||
|
|
@ -38,7 +37,6 @@ headroom unwrap opencode
|
|||
```bash
|
||||
headroom wrap opencode \
|
||||
--port 8787 \
|
||||
--no-rtk \
|
||||
--no-mcp \
|
||||
--no-serena \
|
||||
--code-graph \
|
||||
|
|
@ -71,7 +69,6 @@ The default model is `headroom/claude-sonnet-4-6`. Change it in `opencode.json`
|
|||
|---|---|
|
||||
| `OPENCODE_CONFIG_CONTENT` | JSON payload with provider, plugin, and optional local MCP config injected by `wrap` |
|
||||
| `HEADROOM_PROXY_URL` | Proxy URL passed to Headroom MCP when a non-default port is used, and to the native plugin when configured |
|
||||
| `HEADROOM_CONTEXT_TOOL` | Set to `lean-ctx` to use lean-ctx instead of RTK |
|
||||
|
||||
## Failure Learning
|
||||
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
The Headroom Rust proxy exposes Prometheus-format metrics on the
|
||||
`/metrics` endpoint of every running proxy instance. The metric
|
||||
catalogue below covers Phase D (Bedrock route instrumentation) and
|
||||
Phase G PR-G3 (per-invocation RTK + proxy-wide observability).
|
||||
Phase G PR-G3 (proxy-wide observability).
|
||||
|
||||
All metric names + label keys are constants in
|
||||
`crates/headroom-proxy/src/observability/metric_names.rs`, so any
|
||||
|
|
@ -59,17 +59,6 @@ intentional byte mutations do not trip the alarm.
|
|||
| `proxy_service_tier_count_total` | Counter | `tier` | Service-tier distribution observed at the proxy. |
|
||||
| `proxy_response_status_count_total` | Counter | `status` | Terminal status distribution (`completed`, `incomplete`, `failed`, `cancelled`, `in_progress`). |
|
||||
|
||||
#### Wrap CLI / RTK (Python-side)
|
||||
|
||||
| Name | Type | Labels | Purpose |
|
||||
|------|------|--------|---------|
|
||||
| `wrap_rtk_invocations_total` | Counter | `tool` | RTK invocations observed via the wrap-CLI tail. Surfaced via the Python proxy's `/metrics` exporter; the wrap CLI bumps `headroom.cli.wrap_rtk_metrics.record_rtk_invocation(...)`. |
|
||||
|
||||
> **C4 remediation:** This counter is Python-side because RTK is
|
||||
> wrapped by `headroom wrap` (Python CLI) and the wrap-side tail
|
||||
> is the natural emit site. The Rust proxy previously held a dead
|
||||
> counter for this metric; that has been removed.
|
||||
|
||||
#### Image log redaction (Python-side)
|
||||
|
||||
| Name | Type | Labels | Purpose |
|
||||
|
|
@ -141,9 +130,6 @@ sum by (strategy) (rate(proxy_compression_rejected_by_token_check_total{strategy
|
|||
# Upstream rate-limit headroom (smaller = closer to throttle).
|
||||
proxy_rate_limit_remaining_tokens{provider="anthropic"}
|
||||
|
||||
# RTK invocation rate (Python-side).
|
||||
sum by (tool) (rate(wrap_rtk_invocations_total{tool!="__init__"}[5m]))
|
||||
|
||||
# Image-redaction rate (Python-side).
|
||||
rate(proxy_image_generation_call_log_redacted_total[5m])
|
||||
```
|
||||
|
|
@ -259,16 +245,12 @@ Every label vocabulary is bounded by code, not customer input:
|
|||
`"other"` and a `tracing::warn!` is emitted so wire-format drift
|
||||
surfaces loudly in logs.
|
||||
- `status`: 5-variant enum.
|
||||
- `tool` (Python-side `wrap_rtk_invocations_total`): bounded by the
|
||||
set of tools the wrap CLI rewrites, captured by
|
||||
`headroom.cli.wrap_rtk_metrics`.
|
||||
|
||||
There is no code path where a malicious client can drive label
|
||||
cardinality unbounded.
|
||||
|
||||
## See also
|
||||
|
||||
- `docs/rtk-architecture.md` — why RTK lives wrap-side, not proxy-side.
|
||||
- `crates/headroom-proxy/src/observability/` — implementation.
|
||||
- `REALIGNMENT/09-phase-G-rtk-observability.md` — spec.
|
||||
- `REALIGNMENT/10-phase-H-python-retirement.md` — H1 acceptance gate.
|
||||
|
|
|
|||
|
|
@ -1,122 +0,0 @@
|
|||
# RTK architecture — why wrap-CLI only
|
||||
|
||||
**Status:** decided. Locked at Phase G PR-G3 (2026-05).
|
||||
**Owner:** Headroom realignment.
|
||||
|
||||
## TL;DR
|
||||
|
||||
**RTK is a wrap-CLI hook, not a proxy-side compressor.** The Headroom
|
||||
proxy does NOT invoke RTK on tool-result content. Future contributors
|
||||
who consider moving RTK into the proxy hot path: read this doc first.
|
||||
|
||||
## Background
|
||||
|
||||
RTK (Realtime Token Kompress) rewrites shell **commands** at exec
|
||||
time so that a `git diff` or `grep` invocation emits a more
|
||||
compressed output before the agent ever ingests it. RTK runs in the
|
||||
wrap-CLI tail — `headroom wrap claude`, `headroom wrap codex`, etc.
|
||||
— where it installs a `~/.rtk/bin/rtk` shim ahead of the agent CLI
|
||||
and intercepts shelled-out subprocesses.
|
||||
|
||||
It surfaces value in two places:
|
||||
1. **Tokens saved per invocation** — measured by `rtk gain --format json`.
|
||||
2. **Tokens saved per session** — aggregated at wrap-session end.
|
||||
|
||||
Both signals feed `wrap_rtk_invocations_total` and
|
||||
`wrap_rtk_tokens_saved_per_session` (registered by the Rust proxy's
|
||||
observability surface so a single `/metrics` scrape exposes the full
|
||||
picture).
|
||||
|
||||
## Proxy-side RTK was considered and rejected
|
||||
|
||||
At Phase G scoping, three reviewers floated the idea of invoking
|
||||
RTK on the **proxy** side: when a `tool_result` block flows
|
||||
upstream, dispatch it through RTK to shrink the content before it
|
||||
hits the model.
|
||||
|
||||
**Decision: rejected.** Three load-bearing reasons.
|
||||
|
||||
### 1. Cache hot zone risk
|
||||
|
||||
The proxy's Phase B cache-safety contract pins `tool_result`
|
||||
content as part of the cache hot zone. Compression there bursts
|
||||
the prompt cache because the rewritten bytes diverge from the
|
||||
canonical wire bytes the upstream cached. Phase B PR-B2 → PR-B7
|
||||
spent ~3000 LOC carving the live-zone-only surface specifically
|
||||
to prevent this class of cache-invalidation. Inserting RTK
|
||||
proxy-side would re-introduce it.
|
||||
|
||||
### 2. Parallel implementation with `log_compressor.rs`
|
||||
|
||||
The Rust proxy already has a `crates/headroom-core/src/transforms/log_compressor.rs`
|
||||
that compresses **tool output text** in the live zone. It uses the
|
||||
same heuristics RTK uses (whitespace de-dup, line de-dup,
|
||||
file-listing collapse) but invoked at the proxy's per-block
|
||||
dispatcher rather than at the shell exec boundary. Adding RTK
|
||||
proxy-side would mean two implementations of the same compression
|
||||
in the same hot path; "no silent fallbacks, no parallel impls" is
|
||||
explicit project policy.
|
||||
|
||||
### 3. Command-rewrite vs output-rewrite — different value propositions
|
||||
|
||||
RTK rewrites **commands** before they execute. The
|
||||
`git log --oneline` you typed becomes `git log --oneline -n 50`
|
||||
because RTK has learned that the first 50 commits are usually
|
||||
enough context. That's a fundamentally different mechanism from
|
||||
compressing the **output** of an unmodified command. A proxy-side
|
||||
invocation would skip the command-rewrite half — the half that
|
||||
generates the largest savings on heavy shell workloads — and only
|
||||
catch the output side, which is already covered by
|
||||
`log_compressor` and `code_compressor`.
|
||||
|
||||
## What the proxy does provide
|
||||
|
||||
Per Phase G PR-G3, the proxy exposes RTK-derived metrics via its
|
||||
registry:
|
||||
|
||||
- `wrap_rtk_invocations_total{tool}` — driven by the wrap-CLI
|
||||
polling `rtk gain --format json` and incrementing the registered
|
||||
counter by the delta since last poll.
|
||||
- `wrap_rtk_tokens_saved_per_session` — emitted at wrap-session
|
||||
close.
|
||||
|
||||
This keeps the operator dashboard single-pane-of-glass without
|
||||
re-implementing RTK inside the proxy.
|
||||
|
||||
## What the wrap CLI does
|
||||
|
||||
Every `headroom wrap <agent>` subcommand:
|
||||
|
||||
1. Ensures the RTK binary is installed via `_ensure_rtk_binary()`.
|
||||
2. Injects the `<!-- headroom:rtk-instructions -->` block into the
|
||||
agent's instruction file (e.g. `AGENTS.md`, `.cursorrules`).
|
||||
3. Spawns the proxy and the agent CLI side-by-side.
|
||||
4. Polls `rtk gain --format json` on a 5-second memoization window
|
||||
and feeds the delta into the proxy's metric registry.
|
||||
|
||||
See `headroom/cli/wrap/` for the per-agent shims.
|
||||
|
||||
## Re-litigation policy
|
||||
|
||||
A change to this architecture should:
|
||||
|
||||
1. Quote the live-zone-only contract from
|
||||
`REALIGNMENT/04-phase-B-live-zone.md` and explain why the
|
||||
cache-burst risk is acceptable.
|
||||
2. Show measurements (not estimates) that proxy-side RTK adds value
|
||||
beyond `log_compressor.rs` on real production traffic.
|
||||
3. Have an exit ramp: a CLI flag to disable proxy-side RTK without
|
||||
reverting the wrap-CLI integration.
|
||||
|
||||
Without all three, treat the proposal as a regression and link this
|
||||
doc.
|
||||
|
||||
## References
|
||||
|
||||
- `REALIGNMENT/09-phase-G-rtk-observability.md` — Phase G plan.
|
||||
- `REALIGNMENT/04-phase-B-live-zone.md` — cache hot-zone contract.
|
||||
- `headroom/cli/wrap/` — wrap-CLI implementation.
|
||||
- `crates/headroom-core/src/transforms/log_compressor.rs` — the
|
||||
proxy-side log compressor RTK would parallel.
|
||||
- 2026-05-01 user direction message archived in
|
||||
`project_compression_realignment_2026_05` memory note.
|
||||
107
e2e/wrap/run.py
107
e2e/wrap/run.py
|
|
@ -21,7 +21,6 @@ import httpx
|
|||
REPO_ROOT = Path("/workspace")
|
||||
PLUGIN_DIR = REPO_ROOT / "plugins" / "openclaw"
|
||||
SDK_DIR = REPO_ROOT / "sdk" / "typescript"
|
||||
RTK_MARKER = "<!-- headroom:rtk-instructions -->"
|
||||
PROXY_PORT = 28887
|
||||
CODEX_PORT = 28888
|
||||
AIDER_PORT = 28889
|
||||
|
|
@ -368,25 +367,10 @@ def create_shims(shim_dir: Path) -> None:
|
|||
raise SystemExit(0)
|
||||
"""
|
||||
)
|
||||
rtk_shim = textwrap.dedent(
|
||||
"""\
|
||||
#!/usr/bin/env python3
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
if "--version" in sys.argv:
|
||||
print("rtk e2e-shim")
|
||||
else:
|
||||
print("rtk shim")
|
||||
raise SystemExit(0)
|
||||
"""
|
||||
)
|
||||
write_executable(shim_dir / "claude", generic_shim)
|
||||
write_executable(shim_dir / "codex", codex_shim)
|
||||
write_executable(shim_dir / "aider", generic_shim)
|
||||
write_executable(shim_dir / "opencode", generic_shim)
|
||||
write_executable(shim_dir / "rtk", rtk_shim)
|
||||
|
||||
|
||||
def start_mock_server(port: int) -> tuple[MockOpenAIServer, threading.Thread]:
|
||||
|
|
@ -569,16 +553,6 @@ def verify_codex_wrap(
|
|||
cwd=project_dir,
|
||||
timeout=120,
|
||||
)
|
||||
# RTK guidance for Codex is global-only (#1240): it is injected into
|
||||
# ~/.codex/AGENTS.md, never a project-level AGENTS.md. A project AGENTS.md is
|
||||
# written only when `wrap codex --memory` is used (for memory guidance), which
|
||||
# this scenario does not exercise.
|
||||
global_agents = Path(base_env["HOME"]) / ".codex" / "AGENTS.md"
|
||||
assert_true(global_agents.exists(), "Codex wrap should create ~/.codex/AGENTS.md")
|
||||
assert_true(
|
||||
RTK_MARKER in global_agents.read_text(encoding="utf-8"), "Missing global RTK marker"
|
||||
)
|
||||
|
||||
config_path = Path(base_env["HOME"]) / ".codex" / "config.toml"
|
||||
assert_true(
|
||||
config_path.exists(),
|
||||
|
|
@ -688,13 +662,6 @@ def verify_aider_wrap(base_env: dict[str, str], project_dir: Path, log_dir: Path
|
|||
cwd=project_dir,
|
||||
timeout=120,
|
||||
)
|
||||
conventions = project_dir / "CONVENTIONS.md"
|
||||
assert_true(conventions.exists(), "Aider wrap should create CONVENTIONS.md")
|
||||
assert_true(
|
||||
RTK_MARKER in conventions.read_text(encoding="utf-8"),
|
||||
"Aider wrap should inject RTK instructions",
|
||||
)
|
||||
|
||||
entries = read_jsonl(log_dir / "aider.jsonl")
|
||||
assert_true(len(entries) > 0, "Aider shim should have been invoked")
|
||||
env_vars = entries[-1]["env"]
|
||||
|
|
@ -745,84 +712,52 @@ def verify_cursor_wrap(base_env: dict[str, str], project_dir: Path) -> None:
|
|||
"Cursor wrap should print the Anthropic base URL override",
|
||||
)
|
||||
wait_for_http(f"http://127.0.0.1:{port}/health", timeout=15)
|
||||
# rtk registers a native Cursor hook (rtk init --agent cursor) when it
|
||||
# can (~/.cursor exists); headroom only falls back to injecting
|
||||
# .cursorrules text if that registration fails (GH #756). Accept
|
||||
# either outcome rather than assuming the fallback path.
|
||||
cursorrules = project_dir / ".cursorrules"
|
||||
cursor_hooks_json = Path(base_env["HOME"]) / ".cursor" / "hooks.json"
|
||||
native_hook_registered = (
|
||||
cursor_hooks_json.exists() and "rtk" in cursor_hooks_json.read_text(encoding="utf-8")
|
||||
)
|
||||
if not native_hook_registered:
|
||||
assert_true(
|
||||
cursorrules.exists(),
|
||||
"Cursor wrap should create .cursorrules when the native rtk hook is unavailable",
|
||||
)
|
||||
assert_true(
|
||||
RTK_MARKER in cursorrules.read_text(encoding="utf-8"),
|
||||
"Cursor wrap should inject RTK instructions",
|
||||
)
|
||||
finally:
|
||||
stop_process(proc)
|
||||
|
||||
|
||||
def verify_cline_wrap(base_env: dict[str, str], project_dir: Path) -> None:
|
||||
"""Smoke test: `wrap cline --prepare-only` writes RTK guidance to .clinerules."""
|
||||
"""Smoke test: `wrap cline --prepare-only` exits clean.
|
||||
|
||||
These three wraps used to be verified by the hint-file guidance they wrote.
|
||||
With the CLI context tools removed they produce no on-disk artifact, so the
|
||||
remaining assertion is that the prepare path still runs without crashing —
|
||||
``run`` raises on a non-zero exit.
|
||||
"""
|
||||
run(
|
||||
["headroom", "wrap", "cline", "--prepare-only", "--port", str(CLINE_PORT)],
|
||||
env=base_env,
|
||||
cwd=project_dir,
|
||||
timeout=60,
|
||||
)
|
||||
clinerules = project_dir / ".clinerules"
|
||||
assert_true(clinerules.exists(), "Cline wrap should create .clinerules")
|
||||
assert_true(
|
||||
RTK_MARKER in clinerules.read_text(encoding="utf-8"),
|
||||
"Cline wrap should inject RTK instructions",
|
||||
)
|
||||
|
||||
|
||||
def verify_continue_wrap(base_env: dict[str, str], project_dir: Path) -> None:
|
||||
"""Smoke test: `wrap continue --prepare-only` injects RTK into .continue/config.json."""
|
||||
"""Smoke test: `wrap continue --prepare-only` exits clean (see verify_cline_wrap)."""
|
||||
run(
|
||||
["headroom", "wrap", "continue", "--prepare-only", "--port", str(CONTINUE_PORT)],
|
||||
env=base_env,
|
||||
cwd=project_dir,
|
||||
timeout=60,
|
||||
)
|
||||
config_file = project_dir / ".continue" / "config.json"
|
||||
assert_true(config_file.exists(), "Continue wrap should create .continue/config.json")
|
||||
data = json.loads(config_file.read_text(encoding="utf-8"))
|
||||
system_message = data.get("systemMessage", "")
|
||||
assert_true(
|
||||
RTK_MARKER in system_message,
|
||||
"Continue wrap should inject RTK instructions into systemMessage",
|
||||
)
|
||||
|
||||
|
||||
def verify_goose_wrap(base_env: dict[str, str], project_dir: Path) -> None:
|
||||
"""Smoke test: `wrap goose --prepare-only` writes RTK guidance to .goosehints."""
|
||||
"""Smoke test: `wrap goose --prepare-only` exits clean (see verify_cline_wrap)."""
|
||||
run(
|
||||
["headroom", "wrap", "goose", "--prepare-only", "--port", str(GOOSE_PORT)],
|
||||
env=base_env,
|
||||
cwd=project_dir,
|
||||
timeout=60,
|
||||
)
|
||||
goosehints = project_dir / ".goosehints"
|
||||
assert_true(goosehints.exists(), "Goose wrap should create .goosehints")
|
||||
assert_true(
|
||||
RTK_MARKER in goosehints.read_text(encoding="utf-8"),
|
||||
"Goose wrap should inject RTK instructions",
|
||||
)
|
||||
|
||||
|
||||
def verify_openhands_wrap(base_env: dict[str, str], project_dir: Path) -> None:
|
||||
"""Smoke test: `wrap openhands --prepare-only` exits clean and ensures rtk is present.
|
||||
"""Smoke test: `wrap openhands --prepare-only` exits clean.
|
||||
|
||||
OpenHands wires instructions via the OPENHANDS_INSTRUCTIONS env var at launch
|
||||
time (no on-disk artifact), so --prepare-only just exercises the rtk-binary
|
||||
setup path. The env-var wiring is covered by the unit tests.
|
||||
This one is a real regression guard: openhands used to *require* the rtk
|
||||
binary, so once rtk became opt-in the default path exited 1. Nothing is
|
||||
stubbed here, so a reintroduced hard dependency fails the run.
|
||||
"""
|
||||
run(
|
||||
["headroom", "wrap", "openhands", "--prepare-only", "--port", str(OPENHANDS_PORT)],
|
||||
|
|
@ -953,9 +888,6 @@ def main() -> None:
|
|||
"PATH": f"{shim_dir}{os.pathsep}{base_env['PATH']}",
|
||||
"HEADROOM_E2E_LOG_DIR": str(log_dir),
|
||||
"OPENAI_TARGET_API_URL": "http://127.0.0.1:19001/v1",
|
||||
# RTK is opt-in (off by default). These wrap smoke tests assert
|
||||
# RTK-instruction injection, so exercise the RTK-on path.
|
||||
"HEADROOM_RTK": "1",
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -987,19 +919,6 @@ def verify_opencode_wrap(base_env: dict[str, str], project_dir: Path, log_dir: P
|
|||
cwd=project_dir,
|
||||
timeout=120,
|
||||
)
|
||||
global_agents = Path(base_env["HOME"]) / ".config" / "opencode" / "AGENTS.md"
|
||||
project_agents = project_dir / "AGENTS.md"
|
||||
assert_true(global_agents.exists(), "Opencode wrap should create ~/.config/opencode/AGENTS.md")
|
||||
assert_true(project_agents.exists(), "Opencode wrap should create project AGENTS.md")
|
||||
assert_true(
|
||||
RTK_MARKER in global_agents.read_text(encoding="utf-8"),
|
||||
"Missing RTK marker in global AGENTS.md",
|
||||
)
|
||||
assert_true(
|
||||
RTK_MARKER in project_agents.read_text(encoding="utf-8"),
|
||||
"Missing RTK marker in project AGENTS.md",
|
||||
)
|
||||
|
||||
entries = read_jsonl(log_dir / "opencode.jsonl")
|
||||
assert_true(len(entries) > 0, "Opencode shim should have been invoked")
|
||||
env_vars = entries[-1]["env"]
|
||||
|
|
|
|||
|
|
@ -1,8 +1,7 @@
|
|||
"""Codex transcript audit — read-pattern analysis for shell-based clients.
|
||||
|
||||
Codex has no structured Read tool: it reads files through shell commands
|
||||
(``cat``, ``sed -n 'a,bp'``, ``head``/``tail``, ``nl``) — frequently
|
||||
wrapped by rtk (``rtk read <file>``, ``rtk proxy <cmd>``). This module
|
||||
(``cat``, ``sed -n 'a,bp'``, ``head``/``tail``, ``nl``). This module
|
||||
classifies ``exec_command`` calls in Codex session transcripts
|
||||
(``~/.codex/sessions/**/*.jsonl``) and measures the read pattern so the
|
||||
read-maturation mechanism can be sized for Codex workloads.
|
||||
|
|
@ -23,8 +22,8 @@ from collections import Counter
|
|||
from dataclasses import asdict, dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
# Programs whose output is file content. "read" is rtk's read command.
|
||||
_READ_PROGS = frozenset({"cat", "sed", "head", "tail", "nl", "bat", "more", "read"})
|
||||
# Programs whose output is file content.
|
||||
_READ_PROGS = frozenset({"cat", "sed", "head", "tail", "nl", "bat", "more"})
|
||||
_SEARCH_PROGS = frozenset({"rg", "grep", "ugrep", "ag", "fd", "find"})
|
||||
_BUILD_PROGS = frozenset({"python", "python3", "pytest", "cargo", "npm", "make", "uv", "ruff"})
|
||||
_RANGE_RE = re.compile(r"^\d+([,:-]\d+)?p?$")
|
||||
|
|
@ -57,19 +56,6 @@ class CodexAuditReport:
|
|||
return asdict(self)
|
||||
|
||||
|
||||
def strip_wrappers(cmd: str) -> str:
|
||||
"""Peel rtk wrappers: ``rtk <cmd>`` and ``rtk proxy <cmd>``."""
|
||||
c = cmd.strip()
|
||||
while True:
|
||||
if c.startswith("rtk "):
|
||||
c = c[4:].strip()
|
||||
continue
|
||||
if c.startswith("proxy "):
|
||||
c = c[6:].strip()
|
||||
continue
|
||||
return c
|
||||
|
||||
|
||||
def _resolve_path(path: str | None, workdir: str = "") -> str | None:
|
||||
if not path:
|
||||
return None
|
||||
|
|
@ -129,7 +115,7 @@ def classify_command(cmd: str, workdir: str = "") -> tuple[str, str | None, bool
|
|||
Categories: read, search, git, edit, build/test, compound, other.
|
||||
For reads and edits, the path is resolved against ``workdir`` when relative.
|
||||
"""
|
||||
c = strip_wrappers(cmd)
|
||||
c = cmd.strip()
|
||||
try:
|
||||
toks = shlex.split(c)
|
||||
except ValueError:
|
||||
|
|
|
|||
|
|
@ -177,7 +177,7 @@ def simulate_codex_maturation(root: Path) -> MaturationSimReport:
|
|||
"""Run the maturation simulation over Codex shell-based transcripts.
|
||||
|
||||
Codex has no structured ``Read`` tool. It reads files through
|
||||
``exec_command`` calls such as ``cat``, ``sed -n``, and ``rtk read``.
|
||||
``exec_command`` calls such as ``cat``, ``sed -n``, and ``head``.
|
||||
This mirrors ``simulate_maturation`` with the Codex command classifier
|
||||
so ``headroom audit-reads --codex --simulate-maturation`` sizes the
|
||||
same read-maturation policy from Codex traffic instead of returning an
|
||||
|
|
|
|||
|
|
@ -32,7 +32,7 @@ from .main import main
|
|||
"--codex",
|
||||
"codex_mode",
|
||||
is_flag=True,
|
||||
help="Audit Codex transcripts instead (shell-based reads: cat/sed/rtk read). "
|
||||
help="Audit Codex transcripts instead (shell-based reads: cat/sed/head). "
|
||||
"Default path becomes ~/.codex/sessions.",
|
||||
)
|
||||
def audit_reads_cmd(
|
||||
|
|
|
|||
|
|
@ -66,11 +66,6 @@ warnings.filterwarnings("ignore", category=UserWarning, module="huggingface_hub"
|
|||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_CONTEXT_TOOL_ENV = "HEADROOM_CONTEXT_TOOL"
|
||||
_CONTEXT_TOOL_RTK = "rtk"
|
||||
_CONTEXT_TOOL_LEAN_CTX = "lean-ctx"
|
||||
_VALID_CONTEXT_TOOLS = {_CONTEXT_TOOL_RTK, _CONTEXT_TOOL_LEAN_CTX}
|
||||
|
||||
|
||||
def _get_env_bool(name: str, default: bool) -> bool:
|
||||
val = os.environ.get(name)
|
||||
|
|
@ -117,19 +112,6 @@ def _get_env_float_optional(name: str) -> float | None:
|
|||
raise click.ClickException(f"{name} must be a number, got {val!r}") from None
|
||||
|
||||
|
||||
def _selected_context_tool() -> str:
|
||||
raw = os.environ.get(_CONTEXT_TOOL_ENV, "").strip().lower().replace("_", "-")
|
||||
if not raw:
|
||||
return _CONTEXT_TOOL_RTK
|
||||
if raw == "leanctx":
|
||||
raw = _CONTEXT_TOOL_LEAN_CTX
|
||||
if raw not in _VALID_CONTEXT_TOOLS:
|
||||
raise click.ClickException(
|
||||
f"{_CONTEXT_TOOL_ENV} must be one of: {', '.join(sorted(_VALID_CONTEXT_TOOLS))}"
|
||||
)
|
||||
return raw
|
||||
|
||||
|
||||
@main.command()
|
||||
@click.option(
|
||||
"--port",
|
||||
|
|
@ -1460,7 +1442,6 @@ Memory (Multi-Provider):
|
|||
from headroom.proxy.server import _get_code_aware_banner_status
|
||||
|
||||
code_aware_line = f" Code-Aware: {_get_code_aware_banner_status(config)}"
|
||||
context_tool_line = f" Context Tool: {_selected_context_tool()}"
|
||||
|
||||
# Performance tuning section — only shown when at least one tuning var is active.
|
||||
_embed_socket = os.environ.get("HEADROOM_EMBEDDING_SERVER_SOCKET") or (
|
||||
|
|
@ -1490,7 +1471,6 @@ Starting proxy server...
|
|||
Memory: {memory_status}
|
||||
License: {license_status}
|
||||
{code_aware_line}
|
||||
{context_tool_line}
|
||||
{extensions_line}
|
||||
{security_line}
|
||||
{stateless_line}{telemetry_line}
|
||||
|
|
|
|||
1404
headroom/cli/wrap.py
1404
headroom/cli/wrap.py
File diff suppressed because it is too large
Load diff
|
|
@ -1,73 +0,0 @@
|
|||
"""RTK invocation metrics for the wrap CLI.
|
||||
|
||||
Phase G PR-G3 remediation (C4): RTK lives wrap-side, not proxy-side
|
||||
(see ``docs/rtk-architecture.md``). The wrap CLI tails
|
||||
``rtk gain --format json`` and bumps a process-local counter keyed
|
||||
by rewritten command name (`git`, `ls`, `cargo`, ...). The Python
|
||||
proxy's ``/metrics`` endpoint then surfaces the counter as
|
||||
``wrap_rtk_invocations_total{tool=...}`` for fleet-wide scrape.
|
||||
|
||||
The counter primitives live here (not in ``wrap.py``) so the
|
||||
proxy's prometheus exporter can import them without dragging in the
|
||||
full ``wrap.py`` module — that module owns subprocess-level CLI
|
||||
spawning and is heavyweight to import at proxy startup.
|
||||
|
||||
Per realignment build-constraint "no silent fallbacks":
|
||||
``record_rtk_invocation`` raises on a non-string tool name rather
|
||||
than coercing; a caller passing the wrong type is a bug, not a
|
||||
runtime fallback condition.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from collections import defaultdict
|
||||
from collections.abc import Mapping
|
||||
|
||||
# Module-level counter — process-local. Multiple worker processes
|
||||
# (uvicorn workers) each maintain their own; the Python proxy already
|
||||
# documents this in ``docs/observability.md``. Reset is exposed for
|
||||
# tests; production code never reaches for it.
|
||||
_rtk_invocation_counts: dict[str, int] = defaultdict(int)
|
||||
_lock = threading.Lock()
|
||||
|
||||
|
||||
def record_rtk_invocation(tool: str, delta: int = 1) -> None:
|
||||
"""Record one (or `delta`) RTK invocation(s) for the given tool.
|
||||
|
||||
`tool` is the rewritten command name as observed in the
|
||||
``rtk gain --format json`` output (e.g. ``"git"``, ``"ls"``,
|
||||
``"cargo"``). The counter is keyed verbatim.
|
||||
|
||||
`delta` defaults to 1 for the common "one invocation seen" path
|
||||
but accepts arbitrary positive deltas so the wrap tail can bump
|
||||
by a JSON-reported batch count.
|
||||
|
||||
Raises:
|
||||
TypeError: if `tool` is not a `str` or `delta` is not an `int`.
|
||||
ValueError: if `delta` is negative.
|
||||
"""
|
||||
if not isinstance(tool, str):
|
||||
raise TypeError(f"tool must be a str, got {type(tool).__name__}")
|
||||
if not isinstance(delta, int):
|
||||
raise TypeError(f"delta must be an int, got {type(delta).__name__}")
|
||||
if delta < 0:
|
||||
raise ValueError(f"delta must be non-negative, got {delta}")
|
||||
with _lock:
|
||||
_rtk_invocation_counts[tool] += delta
|
||||
|
||||
|
||||
def rtk_invocation_counts() -> Mapping[str, int]:
|
||||
"""Return a snapshot of the current invocation counts.
|
||||
|
||||
Returns a plain dict (not the defaultdict) so callers cannot
|
||||
accidentally pollute the counter map by reading absent keys.
|
||||
"""
|
||||
with _lock:
|
||||
return dict(_rtk_invocation_counts)
|
||||
|
||||
|
||||
def reset_rtk_invocations() -> None:
|
||||
"""Reset the counter map. Test-only — never called from production."""
|
||||
with _lock:
|
||||
_rtk_invocation_counts.clear()
|
||||
383
headroom/context_tool_cleanup.py
Normal file
383
headroom/context_tool_cleanup.py
Normal file
|
|
@ -0,0 +1,383 @@
|
|||
"""Undo the retired CLI context tools on machines that already have them.
|
||||
|
||||
Headroom used to download ``rtk`` (Rust Token Killer) and ``lean-ctx``, register
|
||||
their Claude Code ``PreToolUse`` hooks, symlink the managed binaries onto
|
||||
``PATH``, and inject a marker-fenced instruction block ("always prefix shell
|
||||
commands with ``rtk``" / "prefer ``ctx_read`` over ``Read``") into a dozen agent
|
||||
hint files. Both integrations have been removed from Headroom.
|
||||
|
||||
Deleting the code is not enough: everything above is *durable state on the
|
||||
user's disk*. Left alone, the Claude hooks keep rewriting every Bash command
|
||||
through binaries Headroom no longer manages, and the injected guidance keeps
|
||||
telling agents to use tools that may not resolve. So ``headroom wrap`` /
|
||||
``headroom unwrap`` call :func:`purge_context_tool_artifacts` once per run to
|
||||
remove what earlier versions installed.
|
||||
|
||||
Everything here is idempotent, best-effort and deliberately conservative:
|
||||
|
||||
* only files Headroom installed (or caused a context tool to install) are
|
||||
deleted;
|
||||
* ``~/.local/bin/{rtk,lean-ctx}`` is unlinked only when it is a symlink into
|
||||
Headroom's own bin directory — a user's own build is never touched;
|
||||
* a JSON config that does not parse is reported and **skipped**, never
|
||||
overwritten (a hand-edited typo must not cost the user their settings);
|
||||
* the tools' own backups of *config* files (``~/.claude.json.lean-ctx.bak`` and
|
||||
friends) are left in place — they hold the user's real settings history. Only
|
||||
backups of the hook scripts being deleted are cleaned up.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from headroom import fsutil, paths
|
||||
|
||||
# Substrings that identify a hook command as one of the retired tools'.
|
||||
# ``rtk init --auto-patch`` writes ``~/.claude/hooks/rtk-rewrite.sh`` and shells
|
||||
# out to ``rtk rewrite``; ``lean-ctx init --agent`` writes ``lean-ctx-rewrite.sh``
|
||||
# / ``lean-ctx-redirect.sh`` plus ``-native`` variants that exec ``lean-ctx hook
|
||||
# rewrite``. All are specific enough that a user-authored hook cannot match by
|
||||
# accident.
|
||||
_HOOK_COMMAND_MARKERS = (
|
||||
"rtk-rewrite",
|
||||
"rtk rewrite",
|
||||
"lean-ctx-rewrite",
|
||||
"lean-ctx-redirect",
|
||||
"lean-ctx hook",
|
||||
)
|
||||
|
||||
# Fence Headroom wrapped its injected guidance in. Both tools shared it — the
|
||||
# lean-ctx block was written inside the same ``rtk-instructions`` markers.
|
||||
_FENCE_START = "<!-- headroom:rtk-instructions -->"
|
||||
_FENCE_END = "<!-- /headroom:rtk-instructions -->"
|
||||
|
||||
# Managed binary names (Windows ships the .exe).
|
||||
_BINARY_NAMES = ("rtk", "rtk.exe", "lean-ctx", "lean-ctx.exe")
|
||||
|
||||
# Hook scripts the tools generated, relative to ``~/.claude/hooks``. Each is
|
||||
# also cleaned up in its ``.lean-ctx.bak`` form: lean-ctx copies the previous
|
||||
# script aside on every re-init, so the backups pile up alongside the originals.
|
||||
_HOOK_SCRIPTS = (
|
||||
"rtk-rewrite.sh",
|
||||
".rtk-hook.sha256",
|
||||
"lean-ctx-rewrite.sh",
|
||||
"lean-ctx-redirect.sh",
|
||||
"lean-ctx-rewrite-native",
|
||||
"lean-ctx-redirect-native",
|
||||
)
|
||||
|
||||
# MCP server entries the tools registered, and the config files holding them.
|
||||
# lean-ctx registers itself as an MCP server during ``lean-ctx init``; rtk never
|
||||
# did, but it is matched too so a stale hand-added entry is cleaned up as well.
|
||||
_MCP_SERVER_NAMES = ("lean-ctx", "lean_ctx", "rtk")
|
||||
|
||||
|
||||
def purge_context_tool_artifacts() -> list[str]:
|
||||
"""Remove every rtk / lean-ctx artifact an earlier Headroom version installed.
|
||||
|
||||
Returns human-readable descriptions of what was removed — plus a line for
|
||||
any config that had to be skipped because the user must fix it by hand. An
|
||||
empty list means there was nothing to do, which is the steady state after
|
||||
the first run.
|
||||
"""
|
||||
home = Path.home()
|
||||
project = Path.cwd()
|
||||
report: list[str] = []
|
||||
|
||||
# 1. Hook registrations (Claude Code's settings.json, Cursor's hooks.json).
|
||||
for config in (home / ".claude" / "settings.json", home / ".cursor" / "hooks.json"):
|
||||
report += _purge_hook_config(config)
|
||||
|
||||
# 2. The generated hook scripts, their integrity digests and stale backups.
|
||||
hooks_dir = home / ".claude" / "hooks"
|
||||
report += _remove_files(
|
||||
*(hooks_dir / name for name in _HOOK_SCRIPTS),
|
||||
*(hooks_dir / f"{name}.lean-ctx.bak" for name in _HOOK_SCRIPTS),
|
||||
)
|
||||
|
||||
# 3. The PATH symlinks, then the managed binaries they pointed at.
|
||||
for name in ("rtk", "lean-ctx"):
|
||||
report += _remove_managed_path_link(home / ".local" / "bin" / name)
|
||||
report += _remove_files(*(paths.bin_dir() / name for name in _BINARY_NAMES))
|
||||
|
||||
# 4. MCP server registrations (lean-ctx registers itself during init).
|
||||
report += _purge_mcp_entries(home / ".claude.json", "mcpServers")
|
||||
report += _purge_mcp_entries(_opencode_home(home) / "opencode.json", "mcp")
|
||||
|
||||
# 5. Marker-fenced guidance in every hint file the wrap harnesses wrote to.
|
||||
for hint_file in _instruction_files(home, project):
|
||||
report += _purge_fenced_block(hint_file)
|
||||
report += _purge_continue_system_messages(project / ".continue" / "config.json")
|
||||
|
||||
return report
|
||||
|
||||
|
||||
def _instruction_files(home: Path, project: Path) -> list[Path]:
|
||||
"""Hint files the wrap subcommands injected the context-tool block into.
|
||||
|
||||
De-duplicated: ``CODEX_HOME`` / ``OPENCODE_HOME`` can point at the project
|
||||
directory, and a file must not be reported twice.
|
||||
"""
|
||||
candidates = (
|
||||
home / ".claude" / "CLAUDE.md",
|
||||
_codex_home(home) / "AGENTS.md",
|
||||
_opencode_home(home) / "AGENTS.md",
|
||||
project / "AGENTS.md",
|
||||
project / ".github" / "copilot-instructions.md",
|
||||
# Shared by `wrap aider` and `wrap openclaude`.
|
||||
project / "CONVENTIONS.md",
|
||||
project / ".clinerules",
|
||||
project / ".goosehints",
|
||||
project / ".cursorrules",
|
||||
)
|
||||
return list(dict.fromkeys(candidates))
|
||||
|
||||
|
||||
def _codex_home(home: Path) -> Path:
|
||||
"""Codex's config directory, respecting ``CODEX_HOME``."""
|
||||
configured = os.environ.get("CODEX_HOME", "").strip()
|
||||
return Path(configured).expanduser() if configured else home / ".codex"
|
||||
|
||||
|
||||
def _opencode_home(home: Path) -> Path:
|
||||
"""OpenCode's config directory, respecting ``OPENCODE_HOME``."""
|
||||
configured = os.environ.get("OPENCODE_HOME", "").strip()
|
||||
return Path(configured).expanduser() if configured else home / ".config" / "opencode"
|
||||
|
||||
|
||||
# --- hook registrations -------------------------------------------------------
|
||||
|
||||
|
||||
def _references_context_tool(entry: Any) -> bool:
|
||||
"""Whether a hook entry's command is one a retired context tool registered."""
|
||||
if not isinstance(entry, dict):
|
||||
return False
|
||||
command = str(entry.get("command", "")).lower()
|
||||
return any(marker in command for marker in _HOOK_COMMAND_MARKERS)
|
||||
|
||||
|
||||
def _prune_hooks(hooks: Any) -> tuple[Any, bool]:
|
||||
"""Drop retired-tool entries from a ``hooks`` mapping; return ``(pruned, changed)``.
|
||||
|
||||
Handles both shapes Headroom's installers produced: Claude Code nests
|
||||
``hooks.<Event>[].hooks[].command`` while Cursor's ``hooks.json`` puts the
|
||||
command directly on the event entry. Anything unrecognised is passed
|
||||
through untouched, and user-authored entries always survive.
|
||||
"""
|
||||
if not isinstance(hooks, dict):
|
||||
return hooks, False
|
||||
|
||||
changed = False
|
||||
pruned: dict[str, Any] = {}
|
||||
for event, entries in hooks.items():
|
||||
if not isinstance(entries, list):
|
||||
pruned[event] = entries
|
||||
continue
|
||||
|
||||
retained: list[Any] = []
|
||||
for entry in entries:
|
||||
# Cursor shape: the command sits on the entry itself.
|
||||
if _references_context_tool(entry):
|
||||
changed = True
|
||||
continue
|
||||
# Claude shape: a matcher entry holding a list of hooks.
|
||||
inner = entry.get("hooks") if isinstance(entry, dict) else None
|
||||
if isinstance(inner, list):
|
||||
kept_inner = [item for item in inner if not _references_context_tool(item)]
|
||||
if len(kept_inner) != len(inner):
|
||||
changed = True
|
||||
if not kept_inner:
|
||||
# The matcher existed only to hold the retired tool's hook.
|
||||
continue
|
||||
entry = {**entry, "hooks": kept_inner}
|
||||
retained.append(entry)
|
||||
|
||||
if not entries:
|
||||
pruned[event] = entries
|
||||
elif retained:
|
||||
pruned[event] = retained
|
||||
else:
|
||||
# Every entry for this event was ours — drop the empty event.
|
||||
changed = True
|
||||
|
||||
return pruned, changed
|
||||
|
||||
|
||||
def _purge_hook_config(path: Path) -> list[str]:
|
||||
"""Remove retired-tool hook registrations from a JSON hook config."""
|
||||
if not path.is_file():
|
||||
return []
|
||||
try:
|
||||
payload = json.loads(fsutil.read_text(path))
|
||||
except (OSError, json.JSONDecodeError) as exc:
|
||||
return [f"skipped {path} (unreadable JSON: {exc}) — remove any stale hook by hand"]
|
||||
if not isinstance(payload, dict):
|
||||
return [f"skipped {path} (not a JSON object) — remove any stale hook by hand"]
|
||||
|
||||
hooks, changed = _prune_hooks(payload.get("hooks"))
|
||||
if not changed:
|
||||
return []
|
||||
if hooks:
|
||||
payload["hooks"] = hooks
|
||||
else:
|
||||
payload.pop("hooks", None)
|
||||
fsutil.write_text(path, json.dumps(payload, indent=2) + "\n")
|
||||
return [f"removed the retired context-tool hook from {path}"]
|
||||
|
||||
|
||||
# --- MCP registrations --------------------------------------------------------
|
||||
|
||||
|
||||
def _purge_mcp_entries(path: Path, container_key: str) -> list[str]:
|
||||
"""Drop retired-tool MCP server entries from a client config.
|
||||
|
||||
``lean-ctx init`` registers lean-ctx as an MCP server in the harness's own
|
||||
config — Claude Code keeps them under ``mcpServers``, OpenCode under ``mcp``.
|
||||
Only the exactly-named entries are removed; every other server, and every
|
||||
unrelated top-level key, is preserved byte-for-byte.
|
||||
"""
|
||||
if not path.is_file():
|
||||
return []
|
||||
try:
|
||||
payload = json.loads(fsutil.read_text(path))
|
||||
except (OSError, json.JSONDecodeError) as exc:
|
||||
return [f"skipped {path} (unreadable JSON: {exc}) — remove any stale MCP entry by hand"]
|
||||
if not isinstance(payload, dict):
|
||||
return [f"skipped {path} (not a JSON object) — remove any stale MCP entry by hand"]
|
||||
|
||||
servers = payload.get(container_key)
|
||||
if not isinstance(servers, dict):
|
||||
return []
|
||||
removed = [name for name in _MCP_SERVER_NAMES if name in servers]
|
||||
if not removed:
|
||||
return []
|
||||
for name in removed:
|
||||
servers.pop(name, None)
|
||||
fsutil.write_text(path, json.dumps(payload, indent=2) + "\n")
|
||||
return [f"removed the {', '.join(removed)} MCP server entry from {path}"]
|
||||
|
||||
|
||||
# --- binaries -----------------------------------------------------------------
|
||||
|
||||
|
||||
def _remove_files(*targets: Path) -> list[str]:
|
||||
"""Delete Headroom-managed files, reporting each removal."""
|
||||
report: list[str] = []
|
||||
for target in targets:
|
||||
if not target.exists() and not target.is_symlink():
|
||||
continue
|
||||
try:
|
||||
target.unlink()
|
||||
except OSError as exc:
|
||||
report.append(f"could not remove {target} ({exc})")
|
||||
continue
|
||||
report.append(f"removed {target}")
|
||||
return report
|
||||
|
||||
|
||||
def _remove_managed_path_link(link: Path) -> list[str]:
|
||||
"""Unlink ``link`` only if it is a symlink into Headroom's bin directory.
|
||||
|
||||
A real file, or a symlink to a user's own build, is left alone — the PATH
|
||||
entry is not ours to reclaim unless we created it.
|
||||
"""
|
||||
if not link.is_symlink():
|
||||
return []
|
||||
try:
|
||||
target = link.resolve()
|
||||
managed_dir = paths.bin_dir().resolve()
|
||||
except OSError:
|
||||
return []
|
||||
if managed_dir not in target.parents:
|
||||
return []
|
||||
return _remove_files(link)
|
||||
|
||||
|
||||
# --- marker-fenced guidance ---------------------------------------------------
|
||||
|
||||
|
||||
def _strip_fences(content: str) -> str | None:
|
||||
"""Remove every fenced context-tool block from ``content``.
|
||||
|
||||
Returns the cleaned text (``""`` when nothing but the block was in the
|
||||
file), or ``None`` when no complete fence was present.
|
||||
"""
|
||||
cleaned = content
|
||||
found = False
|
||||
while True:
|
||||
start = cleaned.find(_FENCE_START)
|
||||
if start < 0:
|
||||
break
|
||||
end = cleaned.find(_FENCE_END, start)
|
||||
if end < 0:
|
||||
break
|
||||
end += len(_FENCE_END)
|
||||
prefix = cleaned[:start].rstrip()
|
||||
suffix = cleaned[end:].lstrip("\n")
|
||||
cleaned = "\n\n".join(part for part in (prefix, suffix) if part)
|
||||
found = True
|
||||
|
||||
if not found:
|
||||
return None
|
||||
return cleaned.rstrip() + "\n" if cleaned.strip() else ""
|
||||
|
||||
|
||||
def _purge_fenced_block(path: Path) -> list[str]:
|
||||
"""Strip the context-tool guidance block from a hint file (deleting an empty file)."""
|
||||
if not path.is_file():
|
||||
return []
|
||||
try:
|
||||
content = fsutil.read_text(path)
|
||||
except OSError:
|
||||
return []
|
||||
cleaned = _strip_fences(content)
|
||||
if cleaned is None:
|
||||
return []
|
||||
if cleaned:
|
||||
fsutil.write_text(path, cleaned)
|
||||
return [f"removed the context-tool guidance block from {path}"]
|
||||
return _remove_files(path)
|
||||
|
||||
|
||||
def _purge_continue_system_messages(path: Path) -> list[str]:
|
||||
"""Strip the context-tool block from every ``systemMessage`` in Continue's config.
|
||||
|
||||
Continue supports a top-level ``systemMessage`` plus a per-model override on
|
||||
each ``models[]`` entry, and wrap injected into all of them. A field left
|
||||
empty by the strip is removed so the config returns to "not set".
|
||||
"""
|
||||
if not path.is_file():
|
||||
return []
|
||||
try:
|
||||
payload = json.loads(fsutil.read_text(path))
|
||||
except (OSError, json.JSONDecodeError) as exc:
|
||||
return [f"skipped {path} (unreadable JSON: {exc}) — remove any stale guidance by hand"]
|
||||
if not isinstance(payload, dict):
|
||||
return [f"skipped {path} (not a JSON object) — remove any stale guidance by hand"]
|
||||
|
||||
containers: list[dict[str, Any]] = [payload]
|
||||
models = payload.get("models")
|
||||
if isinstance(models, list):
|
||||
containers += [model for model in models if isinstance(model, dict)]
|
||||
|
||||
changed = False
|
||||
for container in containers:
|
||||
message = container.get("systemMessage")
|
||||
if not isinstance(message, str):
|
||||
continue
|
||||
cleaned = _strip_fences(message)
|
||||
if cleaned is None:
|
||||
continue
|
||||
changed = True
|
||||
if cleaned.strip():
|
||||
container["systemMessage"] = cleaned
|
||||
else:
|
||||
container.pop("systemMessage", None)
|
||||
|
||||
if not changed:
|
||||
return []
|
||||
fsutil.write_text(path, json.dumps(payload, indent=2) + "\n")
|
||||
return [f"removed the context-tool guidance from {path}"]
|
||||
|
|
@ -222,8 +222,6 @@
|
|||
</div>
|
||||
<div class="mt-1 text-xs text-gray-500 leading-relaxed">
|
||||
<span x-text="'Proxy ' + formatNumber(stats.tokens?.proxy_compression_saved || 0) + ' (' + proxyShareOfTotal.toFixed(1) + '%)'"></span>
|
||||
<span class="mx-1 text-gray-600">/</span>
|
||||
<span x-text="cliFilteringAvailable ? (cliFilteringLabel + ' ' + formatNumber(cliFilteringSaved) + ' this session (' + cliFilteringSessionPctDisplay.toFixed(1) + '%)') : (cliFilteringLabel + ' not installed')"></span>
|
||||
</div>
|
||||
<div class="mt-1 text-xs text-gray-600 leading-relaxed">
|
||||
<span x-text="'Of total wire: ' + (stats.tokens?.savings_percent || 0).toFixed(2) + '%'" title="Savings as fraction of all input tokens including frozen prefix"></span>
|
||||
|
|
@ -389,14 +387,6 @@
|
|||
<span class="text-sm text-gray-400">Before Compression</span>
|
||||
<span class="font-mono text-sm" x-text="formatNumber(stats.tokens?.total_before_compression || 0)"></span>
|
||||
</div>
|
||||
<div class="flex justify-between items-center" x-show="cliFilteringAvailable">
|
||||
<span class="text-sm text-gray-400" x-text="cliFilteringLabel + ' Filtered (this session)'"></span>
|
||||
<span class="font-mono text-sm text-emerald-400" x-text="formatNumber(cliFilteringSaved)"></span>
|
||||
</div>
|
||||
<div class="flex justify-between items-center" x-show="!cliFilteringAvailable">
|
||||
<span class="text-sm text-gray-400" x-text="cliFilteringLabel + ' Filtered (this session)'"></span>
|
||||
<span class="font-mono text-sm text-gray-600">not installed</span>
|
||||
</div>
|
||||
<div class="flex justify-between items-center">
|
||||
<span class="text-sm text-gray-400">Proxy Removed</span>
|
||||
<span class="font-mono text-sm text-accent" x-text="formatNumber(stats.tokens?.proxy_compression_saved || 0)"></span>
|
||||
|
|
@ -1466,17 +1456,6 @@
|
|||
<div class="mt-2 text-xs text-gray-500">Based on persisted weekly buckets</div>
|
||||
</div>
|
||||
|
||||
<template x-if="historyStats.cli_filtering && historyCliFilteringAvailable">
|
||||
<div class="bg-surface rounded-lg p-4 border border-border">
|
||||
<div class="text-xs text-gray-500 uppercase tracking-wide mb-1"
|
||||
x-text="(historyStats.cli_filtering?.label || cliFilteringLabel) + ' Lifetime Saved'"></div>
|
||||
<div class="flex items-baseline gap-2">
|
||||
<span class="text-3xl font-light tabular-nums text-cyan-400"
|
||||
x-text="formatNumber(historyStats.cli_filtering?.lifetime?.tokens_saved || 0)"></span>
|
||||
</div>
|
||||
<div class="mt-2 text-xs text-gray-500">CLI output filtering (lifetime)</div>
|
||||
</div>
|
||||
</template>
|
||||
</div>
|
||||
|
||||
<template x-if="hasHistoricalData">
|
||||
|
|
@ -2630,49 +2609,6 @@
|
|||
return (this.stats.tokens?.proxy_compression_saved || 0) / total * 100;
|
||||
},
|
||||
|
||||
get cliFilteringLabel() {
|
||||
const raw = this.stats.savings?.by_layer?.cli_filtering?.label
|
||||
|| this.stats.context_tool?.label
|
||||
|| this.stats.context_tool?.configured
|
||||
|| 'Context Tool';
|
||||
if (String(raw).toLowerCase() === 'lean-ctx') return 'Lean-ctx';
|
||||
return String(raw);
|
||||
},
|
||||
|
||||
get cliFilteringSaved() {
|
||||
return this.stats.tokens?.cli_filtering_saved
|
||||
?? this.stats.tokens?.cli_tokens_avoided
|
||||
?? this.stats.tokens?.rtk_saved
|
||||
?? 0;
|
||||
},
|
||||
|
||||
get cliFilteringShareOfTotal() {
|
||||
const total = this.compressionTotalBefore;
|
||||
if (total <= 0) return 0;
|
||||
return this.cliFilteringSaved / total * 100;
|
||||
},
|
||||
|
||||
get cliFilteringLifetime() {
|
||||
return this.stats.savings?.by_layer?.cli_filtering?.lifetime?.tokens_saved ?? 0;
|
||||
},
|
||||
|
||||
get cliFilteringSessionPctDisplay() {
|
||||
const p = this.stats.savings?.by_layer?.cli_filtering?.session_savings_pct;
|
||||
return (p === null || p === undefined) ? this.cliFilteringShareOfTotal : p;
|
||||
},
|
||||
|
||||
get cliFilteringAvailable() {
|
||||
const ct = this.stats.context_tool;
|
||||
if (ct && typeof ct.available === 'boolean') return ct.available;
|
||||
return true;
|
||||
},
|
||||
|
||||
get historyCliFilteringAvailable() {
|
||||
const hf = this.historyStats?.cli_filtering;
|
||||
if (hf && typeof hf.available === 'boolean') return hf.available;
|
||||
return true;
|
||||
},
|
||||
|
||||
// --- Headline savings percent ---
|
||||
//
|
||||
get headlineSavingsPercent() {
|
||||
|
|
|
|||
|
|
@ -22,6 +22,8 @@ from __future__ import annotations
|
|||
|
||||
import locale
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
# Sentinel so ``default=None`` can be a real return value if a caller wants it.
|
||||
|
|
@ -65,14 +67,47 @@ def read_text(path: str | os.PathLike[str], *, default: object = _RAISE) -> str:
|
|||
|
||||
|
||||
def write_text(path: str | os.PathLike[str], content: str) -> None:
|
||||
"""Write text as UTF-8 without translating line endings.
|
||||
"""Write text as UTF-8 without translating line endings, atomically.
|
||||
|
||||
``newline=""`` disables the platform ``\\n`` → ``\\r\\n`` rewrite, so the
|
||||
bytes written match ``content`` exactly and existing ``\\r\\n`` endings are
|
||||
never doubled.
|
||||
|
||||
The write goes to a temp file in the same directory, is fsynced, then moved
|
||||
into place with :func:`os.replace` (atomic on POSIX and Windows). Opening a
|
||||
config file ``"w"`` truncates it to zero bytes *before* the new content is
|
||||
written, so a crash, SIGKILL, OOM or ENOSPC mid-write left the user with a
|
||||
half-written ``~/.claude.json`` / ``settings.local.json``. Callers now see
|
||||
either the old file or the complete new one.
|
||||
|
||||
A symlink is followed rather than replaced (dotfile managers symlink these
|
||||
configs), and an existing file's mode is preserved — ``mkstemp`` creates
|
||||
0600, which would otherwise silently tighten a 0644 config.
|
||||
|
||||
Known tradeoff: this needs a writable *parent directory*, whereas an in-place
|
||||
truncate only needed a writable *file*, so a read-only directory holding a
|
||||
writable config now raises ``PermissionError``. That is deliberate — there is
|
||||
no in-place fallback, because falling back on ``OSError`` would turn the
|
||||
ENOSPC case into "file truncated, rewrite failed", the exact data loss this
|
||||
function exists to prevent. A ``PermissionError`` naming the directory is
|
||||
both rarer and actionable.
|
||||
"""
|
||||
with Path(path).open("w", encoding="utf-8", newline="") as f:
|
||||
f.write(content)
|
||||
target = Path(path)
|
||||
if target.is_symlink():
|
||||
target = Path(os.path.realpath(target))
|
||||
fd, tmp_name = tempfile.mkstemp(dir=target.parent, prefix=f".{target.name}.", suffix=".tmp")
|
||||
tmp_path = Path(tmp_name)
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8", newline="") as f:
|
||||
f.write(content)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
if target.exists():
|
||||
shutil.copymode(target, tmp_path)
|
||||
os.replace(tmp_path, target)
|
||||
except BaseException:
|
||||
tmp_path.unlink(missing_ok=True)
|
||||
raise
|
||||
|
||||
|
||||
def append_text(path: str | os.PathLike[str], content: str) -> None:
|
||||
|
|
|
|||
|
|
@ -90,6 +90,6 @@ __all__ = [
|
|||
# Provider detection
|
||||
"get_headroom_provider",
|
||||
"get_model_name_from_strands",
|
||||
# One-helper MCP + hook wiring (Headroom + Serena + RTK-equivalent)
|
||||
# One-helper MCP + hook wiring (Headroom MCP + Serena + tool-output hooks)
|
||||
"HeadroomBundle",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -17,8 +17,8 @@ Strands-native primitives:
|
|||
``uvx`` on first launch. On by default; disable with
|
||||
``enable_serena_mcp=False``.
|
||||
|
||||
* **HeadroomHookProvider** — the RTK-equivalent for Strands.
|
||||
Compresses tool outputs in-place via ``AfterToolCallEvent`` so
|
||||
* **HeadroomHookProvider** — the in-process tool-output compressor
|
||||
for Strands. Compresses tool outputs in-place via ``AfterToolCallEvent`` so
|
||||
verbose JSON / log / search outputs are shrunk before they
|
||||
pollute the agent's context.
|
||||
|
||||
|
|
@ -123,8 +123,7 @@ class HeadroomBundle:
|
|||
enable_serena_mcp: Include the Serena MCP server — the coding-task
|
||||
compressor. Default True. Adds a ``uvx`` first-launch download.
|
||||
enable_hooks: Include :class:`HeadroomHookProvider` for in-place
|
||||
tool-output compression (the RTK-equivalent for Strands).
|
||||
Default True.
|
||||
tool-output compression. Default True.
|
||||
config: Optional :class:`HeadroomConfig` passed to
|
||||
:class:`HeadroomHookProvider`. Default uses framework
|
||||
defaults.
|
||||
|
|
|
|||
|
|
@ -1,46 +0,0 @@
|
|||
"""lean-ctx integration for Headroom.
|
||||
|
||||
lean-ctx configures supported coding agents to route tool output through its
|
||||
context-filtering layer. Headroom downloads and manages the lean-ctx binary.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import platform
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
from headroom import paths as _paths
|
||||
|
||||
LEAN_CTX_VERSION = "v3.4.7"
|
||||
LEAN_CTX_BIN_DIR = _paths.bin_dir()
|
||||
_LEAN_CTX_NAME = "lean-ctx.exe" if platform.system() == "Windows" else "lean-ctx"
|
||||
LEAN_CTX_BIN_PATH = _paths.lean_ctx_path()
|
||||
|
||||
|
||||
def _managed_lean_ctx_candidates() -> list[Path]:
|
||||
"""Return known Headroom-managed lean-ctx binary paths."""
|
||||
candidates = [LEAN_CTX_BIN_DIR / _LEAN_CTX_NAME]
|
||||
for name in ("lean-ctx", "lean-ctx.exe"):
|
||||
path = LEAN_CTX_BIN_DIR / name
|
||||
if path not in candidates:
|
||||
candidates.append(path)
|
||||
return candidates
|
||||
|
||||
|
||||
def get_lean_ctx_path() -> Path | None:
|
||||
"""Get path to lean-ctx binary — check PATH first, then ~/.headroom/bin/."""
|
||||
system_lean_ctx = shutil.which("lean-ctx")
|
||||
if system_lean_ctx:
|
||||
return Path(system_lean_ctx)
|
||||
|
||||
for candidate in _managed_lean_ctx_candidates():
|
||||
if candidate.exists() and candidate.is_file():
|
||||
return candidate
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def is_lean_ctx_installed() -> bool:
|
||||
"""Check if lean-ctx is available."""
|
||||
return get_lean_ctx_path() is not None
|
||||
|
|
@ -1,179 +0,0 @@
|
|||
"""Download and install lean-ctx binary from GitHub releases."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
import platform
|
||||
import stat
|
||||
import subprocess
|
||||
import tarfile
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from urllib.request import urlopen
|
||||
|
||||
from headroom._subprocess import run
|
||||
|
||||
from . import LEAN_CTX_BIN_DIR, LEAN_CTX_VERSION
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
GITHUB_RELEASE_URL = "https://github.com/yvgude/lean-ctx/releases/download"
|
||||
|
||||
|
||||
def _detect_runtime_target_triple() -> str:
|
||||
"""Detect platform and return the lean-ctx release target triple."""
|
||||
system = platform.system()
|
||||
machine = platform.machine()
|
||||
|
||||
if system == "Darwin":
|
||||
arch = "aarch64" if machine == "arm64" else "x86_64"
|
||||
return f"{arch}-apple-darwin"
|
||||
if system == "Linux":
|
||||
arch = "aarch64" if machine == "aarch64" else "x86_64"
|
||||
suffix = "unknown-linux-musl" if _is_musl() else "unknown-linux-gnu"
|
||||
return f"{arch}-{suffix}"
|
||||
if system == "Windows":
|
||||
return "x86_64-pc-windows-msvc"
|
||||
|
||||
raise RuntimeError(f"Unsupported platform: {system} {machine}")
|
||||
|
||||
|
||||
def _is_musl() -> bool:
|
||||
try:
|
||||
result = run(
|
||||
["ldd", "--version"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=2,
|
||||
check=False,
|
||||
)
|
||||
return "musl" in (result.stdout + result.stderr).lower()
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def _get_target_triple() -> str:
|
||||
"""Return the requested lean-ctx target triple, honoring explicit overrides."""
|
||||
return _get_explicit_target_triple() or _detect_runtime_target_triple()
|
||||
|
||||
|
||||
def _get_explicit_target_triple() -> str:
|
||||
"""Return the explicitly requested lean-ctx target triple, if any."""
|
||||
return (
|
||||
os.environ.get("HEADROOM_LEAN_CTX_TARGET", "").strip()
|
||||
or os.environ.get("LEAN_CTX_TARGET", "").strip()
|
||||
)
|
||||
|
||||
|
||||
def _binary_name_for_target(target: str) -> str:
|
||||
"""Return the expected binary name for a target triple."""
|
||||
return "lean-ctx.exe" if "windows" in target else "lean-ctx"
|
||||
|
||||
|
||||
def _should_verify_target(target: str) -> bool:
|
||||
"""Verify runtime-detected targets; explicit overrides may be cross-target."""
|
||||
if _get_explicit_target_triple():
|
||||
return False
|
||||
return target == _detect_runtime_target_triple()
|
||||
|
||||
|
||||
def _get_download_url(version: str) -> tuple[str, str]:
|
||||
"""Get download URL and extension for this platform."""
|
||||
target = _get_target_triple()
|
||||
ext = "zip" if "windows" in target else "tar.gz"
|
||||
url = f"{GITHUB_RELEASE_URL}/{version}/lean-ctx-{target}.{ext}"
|
||||
return url, ext
|
||||
|
||||
|
||||
def download_lean_ctx(version: str | None = None) -> Path:
|
||||
"""Download lean-ctx binary from GitHub releases."""
|
||||
version = version or LEAN_CTX_VERSION
|
||||
target = _get_target_triple()
|
||||
url, ext = _get_download_url(version)
|
||||
target_path = LEAN_CTX_BIN_DIR / _binary_name_for_target(target)
|
||||
|
||||
LEAN_CTX_BIN_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
logger.info("Downloading lean-ctx %s from %s ...", version, url)
|
||||
|
||||
try:
|
||||
if not url.startswith(("http://", "https://")):
|
||||
raise ValueError(f"Invalid URL scheme in {url}")
|
||||
try:
|
||||
with urlopen(url, timeout=30) as response:
|
||||
data = response.read()
|
||||
except Exception as download_err:
|
||||
if "CERTIFICATE_VERIFY_FAILED" in str(download_err):
|
||||
raise RuntimeError(
|
||||
"TLS verification failed downloading lean-ctx; "
|
||||
"fix the local trust store and retry."
|
||||
) from download_err
|
||||
raise
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to download lean-ctx from {url}: {e}") from e
|
||||
|
||||
try:
|
||||
if ext == "tar.gz":
|
||||
with tarfile.open(fileobj=io.BytesIO(data), mode="r:gz") as tar:
|
||||
for member in tar.getmembers():
|
||||
if member.name.endswith("/lean-ctx") or member.name == "lean-ctx":
|
||||
member.name = target_path.name
|
||||
tar.extract(member, LEAN_CTX_BIN_DIR)
|
||||
break
|
||||
else:
|
||||
raise RuntimeError("lean-ctx binary not found in archive")
|
||||
elif ext == "zip":
|
||||
with zipfile.ZipFile(io.BytesIO(data)) as zf:
|
||||
for name in zf.namelist():
|
||||
if name.endswith("lean-ctx.exe") or name.endswith("/lean-ctx"):
|
||||
with zf.open(name) as src, open(target_path, "wb") as dst:
|
||||
dst.write(src.read())
|
||||
break
|
||||
else:
|
||||
raise RuntimeError("lean-ctx binary not found in archive")
|
||||
except (tarfile.TarError, zipfile.BadZipFile) as e:
|
||||
raise RuntimeError(f"Failed to extract lean-ctx archive: {e}") from e
|
||||
|
||||
if "windows" not in target:
|
||||
target_path.chmod(target_path.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH)
|
||||
|
||||
if _should_verify_target(target):
|
||||
try:
|
||||
result = run(
|
||||
[str(target_path), "--version"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(f"lean-ctx verification failed: {result.stderr}")
|
||||
logger.info("lean-ctx installed: %s", result.stdout.strip())
|
||||
except FileNotFoundError as e:
|
||||
raise RuntimeError("lean-ctx binary not found after extraction") from e
|
||||
except subprocess.TimeoutExpired as e:
|
||||
raise RuntimeError("lean-ctx verification timed out") from e
|
||||
else:
|
||||
logger.info(
|
||||
"lean-ctx installed for target %s at %s (verification skipped)",
|
||||
target,
|
||||
target_path,
|
||||
)
|
||||
|
||||
return target_path
|
||||
|
||||
|
||||
def ensure_lean_ctx(version: str | None = None) -> Path | None:
|
||||
"""Ensure lean-ctx is installed — download if needed."""
|
||||
from . import get_lean_ctx_path
|
||||
|
||||
existing = get_lean_ctx_path()
|
||||
if existing:
|
||||
return existing
|
||||
|
||||
try:
|
||||
return download_lean_ctx(version)
|
||||
except RuntimeError as e:
|
||||
logger.warning("Could not install lean-ctx: %s", e)
|
||||
return None
|
||||
|
|
@ -178,7 +178,7 @@ class SessionAnalyzer:
|
|||
total_failures=len(failed_calls),
|
||||
)
|
||||
|
||||
# Detect loops up front: an RTK re-fetch loop has NO failed calls
|
||||
# Detect loops up front: a re-fetch loop has NO failed calls
|
||||
# (each truncated command succeeds), so it must be a first-class reason
|
||||
# to analyze — otherwise the guard below would skip the most expensive
|
||||
# waste pattern whenever a session has no failures and no events.
|
||||
|
|
@ -423,7 +423,7 @@ Your job is to identify patterns that, if documented, would PREVENT TOKEN WASTE
|
|||
Focus on (in priority order):
|
||||
1. **Loops (HIGHEST PRIORITY)** — patterns that REPEATED within a session. If the
|
||||
digest has a "Detected Loops" section, every loop there MUST get a guardrail
|
||||
rule, because loop waste scales with repetition. This includes RTK re-fetch
|
||||
rule, because loop waste scales with repetition. This includes re-fetch
|
||||
loops: a command whose output was truncated, so the agent re-ran variants of
|
||||
it to fetch more. The fix names the command and prescribes getting the full
|
||||
output up front (e.g., "read the whole file" / "raise the output limit for X").
|
||||
|
|
|
|||
|
|
@ -1,13 +1,12 @@
|
|||
"""Synthetic session fixtures that reproduce known waste patterns.
|
||||
|
||||
These build :class:`SessionData` shaped like the real patterns Headroom Learn
|
||||
must catch, so both unit tests and the RTK-loop eval (``benchmarks/
|
||||
rtk_loop_learn_eval.py``) drive the analyzer from one source of truth instead
|
||||
must catch, so unit tests drive the analyzer from one source of truth instead
|
||||
of hand-mocking calls inline.
|
||||
|
||||
The headline fixture is the **RTK re-fetch loop**. RTK truncates a shell
|
||||
command's output; when the truncation drops what the agent needed, the agent
|
||||
re-runs a *variant* to fetch more. Critically these calls SUCCEED
|
||||
The headline fixture is the **re-fetch loop**. A shell command's output is
|
||||
truncated by an output limit; when the truncation drops what the agent needed,
|
||||
the agent re-runs a *variant* to fetch more. Critically these calls SUCCEED
|
||||
(``is_error=False``) — the loop is invisible to failure-only analysis, which
|
||||
is exactly why it was historically under-weighted.
|
||||
"""
|
||||
|
|
@ -47,25 +46,25 @@ def _tc(
|
|||
)
|
||||
|
||||
|
||||
def rtk_refetch_loop_session(
|
||||
session_id: str = "rtk-loop",
|
||||
def refetch_loop_session(
|
||||
session_id: str = "refetch-loop",
|
||||
*,
|
||||
repetitions: int = 5,
|
||||
bytes_per_call: int = 4000,
|
||||
) -> SessionData:
|
||||
"""A session where RTK truncation forces repeated re-fetches of one command.
|
||||
"""A session where output truncation forces repeated re-fetches of one command.
|
||||
|
||||
The agent greps a large log; RTK rewrites each invocation with an output
|
||||
limit. Each call succeeds but returns a truncated window, so the agent
|
||||
bumps the limit / shifts the window and re-runs — ``repetitions`` times.
|
||||
None of the calls error. The fix a good guardrail should produce: fetch the
|
||||
full result up front (e.g., disable RTK truncation for this command, or
|
||||
grep into a file and read it once).
|
||||
The agent greps a large log with an output limit on each invocation. Each
|
||||
call succeeds but returns a truncated window, so the agent bumps the limit
|
||||
/ shifts the window and re-runs — ``repetitions`` times. None of the calls
|
||||
error. The fix a good guardrail should produce: fetch the full result up
|
||||
front (e.g., raise the output limit for this command, or grep into a file
|
||||
and read it once).
|
||||
"""
|
||||
calls: list[ToolCall] = []
|
||||
limit = 50
|
||||
for i in range(repetitions):
|
||||
# Same base command; only the output-limit varies — the RTK signature.
|
||||
# Same base command; only the output-limit varies — the loop signature.
|
||||
command = f"grep -rn 'TimeoutError' logs/ | head -{limit}"
|
||||
output = "logs/app.log:" + ("x" * (bytes_per_call - 20)) + "\n(truncated)"
|
||||
calls.append(_tc("Bash", command, output, msg_index=i * 2))
|
||||
|
|
|
|||
|
|
@ -7,13 +7,12 @@ being a one-time cost. Two loop shapes matter:
|
|||
1. **Error loops** — the same call fails, the agent retries, it fails again
|
||||
(e.g. a wrong path read N times). Every repetition is pure waste.
|
||||
|
||||
2. **RTK re-fetch loops** — RTK (Realtime Token Kompress) rewrites a shell
|
||||
command to truncate its output (``grep foo`` → ``grep foo | head -50``).
|
||||
When the truncation drops what the agent needed, the agent re-runs a
|
||||
*variant* of the same command to fetch more (``head -100``, a new offset,
|
||||
a narrower pattern). Each call succeeds (``is_error=False``) but returns
|
||||
insufficient output, so the loop is invisible to failure-only analysis.
|
||||
See ``docs/rtk-architecture.md`` for why RTK truncates commands.
|
||||
2. **Re-fetch loops** — a shell command's output is limited or truncated
|
||||
(``grep foo | head -50``). When the limit drops what the agent needed, the
|
||||
agent re-runs a *variant* of the same command to fetch more (``head -100``,
|
||||
a new offset, a narrower pattern). Each call succeeds (``is_error=False``)
|
||||
but returns insufficient output, so the loop is invisible to failure-only
|
||||
analysis.
|
||||
|
||||
This module collapses such variants to a canonical signature, counts the
|
||||
repetitions, and measures the wasted tokens so the analyzer can (a) surface
|
||||
|
|
@ -41,7 +40,7 @@ DEFAULT_MIN_OCCURRENCES = 3
|
|||
# estimate. The analyzer's digest builder uses the same 4:1 approximation.
|
||||
_BYTES_PER_TOKEN = 4
|
||||
|
||||
# Pagination / output-limiting fragments that vary between RTK re-fetch
|
||||
# Pagination / output-limiting fragments that vary between re-fetch
|
||||
# attempts but do NOT change which command is being run. Stripping these is
|
||||
# what collapses ``grep foo | head -50`` and ``grep foo | head -100`` to one
|
||||
# signature. Order-independent: applied as a global substitution.
|
||||
|
|
@ -83,14 +82,14 @@ class LoopPattern:
|
|||
|
||||
@property
|
||||
def kind(self) -> str:
|
||||
return "error-loop" if self.is_error_loop else "rtk-refetch-loop"
|
||||
return "error-loop" if self.is_error_loop else "refetch-loop"
|
||||
|
||||
|
||||
def _canonical_signature(tc: ToolCall) -> str:
|
||||
"""Collapse a tool call to a signature stable across re-fetch variants.
|
||||
|
||||
For shell commands this strips pagination/limit fragments and bare
|
||||
integers so RTK truncation variants of the same command map together.
|
||||
integers so output-limit variants of the same command map together.
|
||||
For other tools the input summary is normalized on whitespace only.
|
||||
"""
|
||||
raw = tc.input_summary.strip()
|
||||
|
|
@ -140,7 +139,7 @@ def detect_loops(
|
|||
wasted = sum(_tokens(c) for c in calls)
|
||||
else:
|
||||
# Re-fetch loop: the first call is legitimate; the N-1 follow-ups
|
||||
# are the redundant re-fetches RTK truncation provoked.
|
||||
# are the redundant re-fetches the output truncation provoked.
|
||||
per_call = sorted((_tokens(c) for c in calls), reverse=True)
|
||||
wasted = sum(per_call[1:])
|
||||
loops.append(
|
||||
|
|
|
|||
|
|
@ -78,10 +78,6 @@ _DEBUG_400_DIR = "debug_400"
|
|||
_CODEX_WIRE_DEBUG_DIR = "codex_wire"
|
||||
_BIN_DIR = "bin"
|
||||
_PROXY_CLIENTS_DIR = "clients"
|
||||
_RTK_UNIX = "rtk"
|
||||
_RTK_WIN = "rtk.exe"
|
||||
_LEAN_CTX_UNIX = "lean-ctx"
|
||||
_LEAN_CTX_WIN = "lean-ctx.exe"
|
||||
_DEPLOY_DIR = "deploy"
|
||||
_PLUGINS_DIR = "plugins"
|
||||
|
||||
|
|
@ -334,20 +330,6 @@ def proxy_clients_dir(port: int) -> Path:
|
|||
return workspace_dir() / _PROXY_CLIENTS_DIR / str(port)
|
||||
|
||||
|
||||
def rtk_path() -> Path:
|
||||
"""Return the path to the vendored ``rtk`` binary."""
|
||||
|
||||
name = _RTK_WIN if os.name == "nt" else _RTK_UNIX
|
||||
return bin_dir() / name
|
||||
|
||||
|
||||
def lean_ctx_path() -> Path:
|
||||
"""Return the path to the vendored ``lean-ctx`` binary."""
|
||||
|
||||
name = _LEAN_CTX_WIN if os.name == "nt" else _LEAN_CTX_UNIX
|
||||
return bin_dir() / name
|
||||
|
||||
|
||||
def deploy_root() -> Path:
|
||||
"""Return the root directory for persistent deployment profiles."""
|
||||
|
||||
|
|
@ -448,8 +430,6 @@ __all__ = [
|
|||
"codex_wire_debug_dir",
|
||||
"bin_dir",
|
||||
"proxy_clients_dir",
|
||||
"rtk_path",
|
||||
"lean_ctx_path",
|
||||
"deploy_root",
|
||||
"beacon_lock_path",
|
||||
"models_config_path",
|
||||
|
|
|
|||
|
|
@ -443,72 +443,15 @@ def parse_log_files(last_n_hours: float = 168.0) -> PerfReport:
|
|||
return report
|
||||
|
||||
|
||||
def _context_tool_lifetime_savings() -> dict | None:
|
||||
"""Lifetime savings from the configured CLI context tool (RTK / lean-ctx).
|
||||
|
||||
``perf`` reports a windowed view of the proxy's *compression* logs. The CLI
|
||||
context tool (RTK) keeps its own lifetime counter that never lands in
|
||||
``proxy.log``, so without this it stays invisible in ``headroom perf`` even
|
||||
when it dwarfs proxy-side savings. Lifetime (not session) is the right scope
|
||||
here: ``perf`` is a one-shot CLI, so the proxy-session baseline ``/stats``
|
||||
subtracts is meaningless out of process.
|
||||
|
||||
Best-effort: returns ``None`` when no tool is installed or its stats cannot
|
||||
be read, so the report degrades to proxy-only rather than erroring.
|
||||
"""
|
||||
try:
|
||||
from headroom.proxy.helpers import _get_context_tool_stats
|
||||
|
||||
stats = _get_context_tool_stats()
|
||||
except Exception:
|
||||
return None
|
||||
if not stats or not stats.get("installed", False):
|
||||
return None
|
||||
lifetime = stats.get("lifetime") or {}
|
||||
tokens_saved = int(lifetime.get("tokens_saved", 0) or 0)
|
||||
if tokens_saved <= 0:
|
||||
return None
|
||||
return {
|
||||
"tool": str(stats.get("tool", "rtk")),
|
||||
"label": str(stats.get("label", "RTK")),
|
||||
"tokens_saved": tokens_saved,
|
||||
"commands": int(lifetime.get("commands", 0) or 0),
|
||||
"savings_pct": round(float(lifetime.get("savings_pct", 0.0) or 0.0), 1),
|
||||
}
|
||||
|
||||
|
||||
def _cli_filtering_report_lines() -> list[str]:
|
||||
"""Render the context-tool (RTK) lifetime savings section, or [] if absent."""
|
||||
cli = _context_tool_lifetime_savings()
|
||||
if not cli:
|
||||
return []
|
||||
return [
|
||||
f"{cli['label']} CLI Filtering (lifetime, all-time)",
|
||||
"-" * 40,
|
||||
f" Tokens saved: {cli['tokens_saved']:,} ({cli['savings_pct']:.1f}%)",
|
||||
f" Commands: {cli['commands']:,}",
|
||||
f" Note: {cli['label']}'s own lifetime counter — not limited to the --hours window.",
|
||||
"",
|
||||
]
|
||||
|
||||
|
||||
def format_report(report: PerfReport) -> str:
|
||||
"""Format a PerfReport into a human-readable string."""
|
||||
lines: list[str] = []
|
||||
cli_filtering_lines = _cli_filtering_report_lines()
|
||||
|
||||
if not report.perf_records and not report.router_records:
|
||||
if cli_filtering_lines:
|
||||
# RTK savings are independent of proxy logs — surface them even when
|
||||
# there is no proxy traffic in the window.
|
||||
lines.append("No proxy performance data in ~/.headroom/logs/ for this window.")
|
||||
lines.append("")
|
||||
lines.extend(cli_filtering_lines)
|
||||
else:
|
||||
lines.append("No performance data found in ~/.headroom/logs/")
|
||||
lines.append("")
|
||||
lines.append("Start the proxy to begin collecting data:")
|
||||
lines.append(" headroom proxy")
|
||||
lines.append("No performance data found in ~/.headroom/logs/")
|
||||
lines.append("")
|
||||
lines.append("Start the proxy to begin collecting data:")
|
||||
lines.append(" headroom proxy")
|
||||
return "\n".join(lines)
|
||||
|
||||
# Header
|
||||
|
|
@ -759,10 +702,6 @@ def format_report(report: PerfReport) -> str:
|
|||
lines.append(f" {i}. {rec}")
|
||||
lines.append("")
|
||||
|
||||
# CLI context-tool (RTK) lifetime savings — its own counter never reaches
|
||||
# proxy.log, so surface it here or it stays invisible in `headroom perf`.
|
||||
lines.extend(cli_filtering_lines)
|
||||
|
||||
# Footer
|
||||
lines.append(
|
||||
f"Log files: {report.log_files_read} | Lines parsed: {report.total_lines_parsed:,}"
|
||||
|
|
@ -1106,9 +1045,6 @@ def build_perf_summary(report: PerfReport) -> dict:
|
|||
"throughput": calculate_throughput(report),
|
||||
"log_files_read": report.log_files_read,
|
||||
"total_lines_parsed": report.total_lines_parsed,
|
||||
# RTK/CLI context-tool lifetime savings (its own counter, not in
|
||||
# proxy.log) — None when no tool is installed. Mirrors the text report.
|
||||
"cli_filtering": _context_tool_lifetime_savings(),
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -382,21 +382,15 @@ def build_prefix_cache_stats(
|
|||
def merge_cost_stats(
|
||||
cost_stats: dict | None,
|
||||
cache_stats: dict,
|
||||
cli_tokens_avoided: int = 0,
|
||||
) -> dict | None:
|
||||
"""Merge compression, cache, and CLI savings into cost stats.
|
||||
"""Merge compression and cache savings into cost stats.
|
||||
|
||||
Each savings layer is reported separately with its own scope:
|
||||
- savings_usd: compression savings at model list price (monotonic)
|
||||
- cache_savings_usd: prefix cache discount from provider (separate)
|
||||
- cli_tokens_avoided: tokens filtered by the selected CLI context tool
|
||||
(token count only, no $ estimate)
|
||||
|
||||
The dollar metric (savings_usd) remains ONLY proxy compression savings
|
||||
priced at the model's published input rate. CLI filtering is folded into
|
||||
the dashboard's compression token total, but it has no reliable
|
||||
model-specific dollar estimate because those tokens never reached the
|
||||
proxy request.
|
||||
priced at the model's published input rate.
|
||||
Prefix cache savings stay separate because they are a provider discount,
|
||||
not token removal. This avoids the non-monotonic moving-average repricing
|
||||
bug (#83).
|
||||
|
|
@ -412,10 +406,6 @@ def merge_cost_stats(
|
|||
"savings_usd": round(compression_savings, 4),
|
||||
"compression_savings_usd": round(compression_savings, 4),
|
||||
"cache_savings_usd": round(cache_net, 4),
|
||||
"cli_tokens_avoided": cli_tokens_avoided,
|
||||
"cli_filtering_tokens_avoided": cli_tokens_avoided,
|
||||
"cli_tokens_included_in_compression": True,
|
||||
"cli_filtering_tokens_included_in_compression": True,
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -474,7 +464,6 @@ def build_session_summary(
|
|||
proxy: Any,
|
||||
metrics: Any,
|
||||
prefix_cache_stats: dict,
|
||||
cli_tokens_avoided: int,
|
||||
total_tokens_before: int,
|
||||
) -> dict[str, Any]:
|
||||
"""Build a human-readable session summary from metrics and request logs.
|
||||
|
|
@ -559,8 +548,7 @@ def build_session_summary(
|
|||
best_detail = f"{best['original']:,} → {best['optimized']:,} tokens"
|
||||
|
||||
# Cost summary — dollar savings are proxy-compression only at model list
|
||||
# price. CLI filtering tokens are counted in token savings but have no
|
||||
# model-specific price because they never reached the proxy request.
|
||||
# price.
|
||||
cost_stats = proxy.cost_tracker.stats() if proxy.cost_tracker else {}
|
||||
cost_with = cost_stats.get("cost_with_headroom_usd", 0.0)
|
||||
compression_savings = cost_stats.get("savings_usd", 0.0)
|
||||
|
|
@ -585,22 +573,12 @@ def build_session_summary(
|
|||
"best_compression_pct": best_compression,
|
||||
"best_detail": best_detail,
|
||||
"total_tokens_removed": metrics.tokens_saved_total,
|
||||
"cli_filtering_tokens_avoided": cli_tokens_avoided,
|
||||
"total_tokens_saved_with_cli_filtering": (
|
||||
metrics.tokens_saved_total + cli_tokens_avoided
|
||||
),
|
||||
"total_tokens_before_with_cli_filtering": total_tokens_before,
|
||||
"rtk_tokens_avoided": cli_tokens_avoided,
|
||||
"total_tokens_saved_with_rtk": metrics.tokens_saved_total + cli_tokens_avoided,
|
||||
"total_tokens_before_with_rtk": total_tokens_before,
|
||||
"total_tokens_before": total_tokens_before,
|
||||
# Tool-schema deferral / turn-hook tool shrink, tracked apart from
|
||||
# message compression. New fields (existing ones stay message+CLI only
|
||||
# for backward compat) so consumers can see the full picture.
|
||||
# message compression, so consumers can see the full picture.
|
||||
"tool_schema_tokens_saved": getattr(metrics, "tool_search_saved_total", 0),
|
||||
"total_tokens_saved_all_layers": (
|
||||
metrics.tokens_saved_total
|
||||
+ cli_tokens_avoided
|
||||
+ getattr(metrics, "tool_search_saved_total", 0)
|
||||
metrics.tokens_saved_total + getattr(metrics, "tool_search_saved_total", 0)
|
||||
),
|
||||
},
|
||||
"uncompressed_requests": {k: v for k, v in uncompressed_reasons.items() if v > 0},
|
||||
|
|
@ -612,16 +590,6 @@ def build_session_summary(
|
|||
"breakdown": {
|
||||
"cache_savings_usd": round(cache_net, 2),
|
||||
"compression_savings_usd": round(compression_savings, 2),
|
||||
"cli_filtering_savings_usd": None,
|
||||
"cli_filtering_savings_note": (
|
||||
"CLI filtering tokens are included in token savings only; "
|
||||
"dollar savings use proxy compression tokens at model list price."
|
||||
),
|
||||
"rtk_savings_usd": None,
|
||||
"rtk_savings_note": (
|
||||
"CLI filtering tokens are included in token savings only; dollar savings "
|
||||
"use proxy compression tokens at model list price."
|
||||
),
|
||||
},
|
||||
},
|
||||
}
|
||||
|
|
|
|||
|
|
@ -23,7 +23,6 @@ from pathlib import Path
|
|||
from typing import TYPE_CHECKING, Any, Literal, cast
|
||||
|
||||
from headroom import paths as _paths
|
||||
from headroom._subprocess import run
|
||||
from headroom.proxy import (
|
||||
diagnostic_decode_policy,
|
||||
memory_injection_mode_policy,
|
||||
|
|
@ -499,37 +498,6 @@ def append_text_to_latest_user_input_item(
|
|||
return body_input, 0
|
||||
|
||||
|
||||
_CONTEXT_TOOL_ENV = "HEADROOM_CONTEXT_TOOL"
|
||||
_CONTEXT_TOOL_RTK = "rtk"
|
||||
_CONTEXT_TOOL_LEAN_CTX = "lean-ctx"
|
||||
_RTK_GAIN_SCOPE_ENV = "HEADROOM_RTK_GAIN_SCOPE"
|
||||
_RTK_GAIN_SCOPE_GLOBAL = "global"
|
||||
_RTK_GAIN_SCOPE_PROJECT = "project"
|
||||
_RTK_GAIN_SCOPES = {_RTK_GAIN_SCOPE_GLOBAL, _RTK_GAIN_SCOPE_PROJECT}
|
||||
|
||||
RTK_STATS_CACHE_TTL_SECONDS = float(os.environ.get("HEADROOM_CONTEXT_TOOL_STATS_TTL_SECONDS", "60"))
|
||||
CONTEXT_TOOL_STATS_CACHE_TTL_SECONDS = RTK_STATS_CACHE_TTL_SECONDS
|
||||
_context_tool_stats_cache_lock = threading.Lock()
|
||||
_context_tool_stats_cache: dict[str, Any] = {
|
||||
"expires_at": 0.0,
|
||||
"has_value": False,
|
||||
"tool": None,
|
||||
"value": None,
|
||||
}
|
||||
_context_tool_session_baseline: dict[str, Any] = {
|
||||
"initialized": False,
|
||||
"tool": None,
|
||||
"total_commands": 0,
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
"tokens_saved": 0,
|
||||
"total_time_ms": 0,
|
||||
"captured_at": 0.0,
|
||||
}
|
||||
_rtk_stats_cache_lock = _context_tool_stats_cache_lock
|
||||
_rtk_stats_cache = _context_tool_stats_cache
|
||||
_rtk_session_baseline = _context_tool_session_baseline
|
||||
|
||||
# Maximum request body size (100MB - increased to support image-heavy requests)
|
||||
MAX_REQUEST_BODY_SIZE = 100 * 1024 * 1024
|
||||
|
||||
|
|
@ -1010,529 +978,6 @@ def _setup_file_logging() -> None:
|
|||
pass
|
||||
|
||||
|
||||
def _selected_context_tool() -> str:
|
||||
raw = os.environ.get(_CONTEXT_TOOL_ENV, _CONTEXT_TOOL_RTK).strip().lower()
|
||||
normalized = raw.replace("_", "-")
|
||||
if normalized in ("leanctx", _CONTEXT_TOOL_LEAN_CTX):
|
||||
return _CONTEXT_TOOL_LEAN_CTX
|
||||
return _CONTEXT_TOOL_RTK
|
||||
|
||||
|
||||
def _context_tool_label(tool: str) -> str:
|
||||
if tool == _CONTEXT_TOOL_LEAN_CTX:
|
||||
return "lean-ctx"
|
||||
return "RTK"
|
||||
|
||||
|
||||
def _context_tool_default_scope(tool: str) -> str:
|
||||
if tool == _CONTEXT_TOOL_LEAN_CTX:
|
||||
return "local"
|
||||
return _RTK_GAIN_SCOPE_GLOBAL
|
||||
|
||||
|
||||
def _rtk_gain_scope() -> str:
|
||||
raw = os.environ.get(_RTK_GAIN_SCOPE_ENV, "").strip().lower()
|
||||
if not raw:
|
||||
return _RTK_GAIN_SCOPE_GLOBAL
|
||||
if raw in _RTK_GAIN_SCOPES:
|
||||
return raw
|
||||
|
||||
logger.warning(
|
||||
"event=rtk_gain_scope_invalid env=%s value=%r default=%s",
|
||||
_RTK_GAIN_SCOPE_ENV,
|
||||
raw,
|
||||
_RTK_GAIN_SCOPE_GLOBAL,
|
||||
)
|
||||
return _RTK_GAIN_SCOPE_GLOBAL
|
||||
|
||||
|
||||
def _rtk_gain_command(rtk_path: Any, scope: str) -> list[str]:
|
||||
command = [str(rtk_path), "gain"]
|
||||
if scope == _RTK_GAIN_SCOPE_PROJECT:
|
||||
command.append("--project")
|
||||
command.extend(["--format", "json"])
|
||||
return command
|
||||
|
||||
|
||||
def _coerce_int(value: Any, default: int = 0) -> int:
|
||||
try:
|
||||
return int(value or 0)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
|
||||
def _coerce_float(value: Any, default: float = 0.0) -> float:
|
||||
try:
|
||||
return float(value or 0.0)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
|
||||
def _first_value(mapping: dict[str, Any], keys: tuple[str, ...], default: Any = 0) -> Any:
|
||||
for key in keys:
|
||||
if key in mapping and mapping[key] is not None:
|
||||
return mapping[key]
|
||||
return default
|
||||
|
||||
|
||||
def _context_tool_summary_payload(
|
||||
*,
|
||||
tool: str,
|
||||
installed: bool,
|
||||
scope: str | None = None,
|
||||
summary: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Normalize RTK/lean-ctx lifetime gain output into one schema.
|
||||
|
||||
Both tools expose cumulative counters, but field names vary slightly.
|
||||
Headroom computes session values by subtracting a startup baseline, so
|
||||
keeping raw input/output counters is necessary for a truthful session
|
||||
savings percentage.
|
||||
"""
|
||||
|
||||
summary = summary or {}
|
||||
input_tokens = _coerce_int(
|
||||
_first_value(
|
||||
summary,
|
||||
(
|
||||
"total_input",
|
||||
"total_input_tokens",
|
||||
"input_tokens",
|
||||
"tokens_input",
|
||||
"totalBefore",
|
||||
),
|
||||
)
|
||||
)
|
||||
output_tokens = _coerce_int(
|
||||
_first_value(
|
||||
summary,
|
||||
(
|
||||
"total_output",
|
||||
"total_output_tokens",
|
||||
"output_tokens",
|
||||
"tokens_output",
|
||||
"totalAfter",
|
||||
),
|
||||
)
|
||||
)
|
||||
tokens_saved = _coerce_int(
|
||||
_first_value(
|
||||
summary,
|
||||
(
|
||||
"total_saved",
|
||||
"tokens_saved",
|
||||
"total_tokens_saved",
|
||||
"saved_tokens",
|
||||
"totalSaved",
|
||||
),
|
||||
)
|
||||
)
|
||||
if tokens_saved <= 0 and input_tokens > 0 and output_tokens >= 0:
|
||||
tokens_saved = max(input_tokens - output_tokens, 0)
|
||||
if input_tokens <= 0 and tokens_saved > 0 and output_tokens >= 0:
|
||||
input_tokens = tokens_saved + output_tokens
|
||||
|
||||
lifetime_savings_pct = _coerce_float(
|
||||
_first_value(
|
||||
summary,
|
||||
(
|
||||
"avg_savings_pct",
|
||||
"average_savings_pct",
|
||||
"savings_pct",
|
||||
"savings_percent",
|
||||
"avgSavingsPct",
|
||||
),
|
||||
0.0,
|
||||
)
|
||||
)
|
||||
if lifetime_savings_pct <= 0 and input_tokens > 0:
|
||||
lifetime_savings_pct = (tokens_saved / input_tokens) * 100.0
|
||||
|
||||
return {
|
||||
"tool": tool,
|
||||
"label": _context_tool_label(tool),
|
||||
"installed": installed,
|
||||
"scope": scope or _context_tool_default_scope(tool),
|
||||
"total_commands": _coerce_int(
|
||||
_first_value(
|
||||
summary,
|
||||
(
|
||||
"total_commands",
|
||||
"commands",
|
||||
"command_count",
|
||||
"totalCommandCount",
|
||||
),
|
||||
)
|
||||
),
|
||||
"input_tokens": input_tokens,
|
||||
"output_tokens": output_tokens,
|
||||
"tokens_saved": tokens_saved,
|
||||
# Backward-compatible name. See `lifetime_avg_savings_pct` and
|
||||
# `session_savings_pct` below for explicit scopes.
|
||||
"avg_savings_pct": lifetime_savings_pct,
|
||||
"lifetime_avg_savings_pct": lifetime_savings_pct,
|
||||
"total_time_ms": _coerce_int(
|
||||
_first_value(summary, ("total_time_ms", "time_ms", "totalTimeMs"))
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _context_tool_zero_payload(
|
||||
*,
|
||||
tool: str,
|
||||
installed: bool,
|
||||
scope: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
return _context_tool_summary_payload(
|
||||
tool=tool,
|
||||
installed=installed,
|
||||
scope=scope,
|
||||
summary={},
|
||||
)
|
||||
|
||||
|
||||
def _read_rtk_lifetime_stats() -> dict[str, Any] | None:
|
||||
"""Read rtk's lifetime stats using the configured gain scope."""
|
||||
|
||||
from headroom.rtk import get_rtk_path
|
||||
|
||||
scope = _rtk_gain_scope()
|
||||
rtk_path = get_rtk_path()
|
||||
if not rtk_path:
|
||||
return _context_tool_zero_payload(
|
||||
tool=_CONTEXT_TOOL_RTK,
|
||||
installed=False,
|
||||
scope=scope,
|
||||
)
|
||||
|
||||
try:
|
||||
result = run(
|
||||
_rtk_gain_command(rtk_path, scope),
|
||||
capture_output=True,
|
||||
text=True,
|
||||
# rtk output is UTF-8 (emoji etc.); without this, Windows decodes
|
||||
# with cp1252 and the reader thread dies with UnicodeDecodeError.
|
||||
encoding="utf-8",
|
||||
errors="replace",
|
||||
timeout=5,
|
||||
)
|
||||
if result.returncode == 0 and result.stdout.strip():
|
||||
data = json.loads(result.stdout)
|
||||
summary = data.get("summary", {})
|
||||
payload = _context_tool_summary_payload(
|
||||
tool=_CONTEXT_TOOL_RTK,
|
||||
installed=True,
|
||||
scope=scope,
|
||||
summary=summary if isinstance(summary, dict) else {},
|
||||
)
|
||||
else:
|
||||
# A failed read is "no data", never a zero counter — a synthetic
|
||||
# zero here re-pins the session baseline and inflates session
|
||||
# savings by the tool's whole lifetime on recovery.
|
||||
stderr_excerpt = (result.stderr or "")[:200]
|
||||
logger.warning(
|
||||
"event=rtk_stats_subprocess_failed reason=non_zero_exit rc=%s stderr=%r",
|
||||
result.returncode,
|
||||
stderr_excerpt,
|
||||
)
|
||||
return None
|
||||
except Exception as exc:
|
||||
# Reason is the exception class name (without payload — RTK
|
||||
# exceptions can carry filesystem paths).
|
||||
logger.warning(
|
||||
"event=rtk_stats_subprocess_failed reason=%s error=%s",
|
||||
type(exc).__name__,
|
||||
exc,
|
||||
)
|
||||
return None
|
||||
|
||||
return payload
|
||||
|
||||
|
||||
def _read_lean_ctx_lifetime_stats() -> dict[str, Any] | None:
|
||||
"""Read lean-ctx's current project-level lifetime stats."""
|
||||
|
||||
from headroom.lean_ctx import get_lean_ctx_path
|
||||
|
||||
lean_ctx_path = get_lean_ctx_path()
|
||||
if not lean_ctx_path:
|
||||
return _context_tool_zero_payload(tool=_CONTEXT_TOOL_LEAN_CTX, installed=False)
|
||||
|
||||
try:
|
||||
result = run(
|
||||
[str(lean_ctx_path), "gain", "--json"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
# UTF-8 regardless of the Windows console code page (cp1252).
|
||||
encoding="utf-8",
|
||||
errors="replace",
|
||||
timeout=5,
|
||||
)
|
||||
# Failed reads return None ("no data") — mirrors the rtk reader so
|
||||
# the baseline logic never sees synthetic zeros from either tool.
|
||||
if result.returncode != 0 or not result.stdout.strip():
|
||||
logger.warning(
|
||||
"event=lean_ctx_stats_subprocess_failed reason=non_zero_exit rc=%s",
|
||||
result.returncode,
|
||||
)
|
||||
return None
|
||||
|
||||
data = json.loads(result.stdout)
|
||||
summary = data.get("summary", data) if isinstance(data, dict) else {}
|
||||
if not isinstance(summary, dict):
|
||||
logger.warning("event=lean_ctx_stats_subprocess_failed reason=bad_payload")
|
||||
return None
|
||||
|
||||
return _context_tool_summary_payload(
|
||||
tool=_CONTEXT_TOOL_LEAN_CTX,
|
||||
installed=True,
|
||||
summary=summary,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"event=lean_ctx_stats_subprocess_failed reason=%s",
|
||||
type(exc).__name__,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _read_context_tool_lifetime_stats(tool: str) -> dict[str, Any] | None:
|
||||
if tool == _CONTEXT_TOOL_LEAN_CTX:
|
||||
return _read_lean_ctx_lifetime_stats()
|
||||
return _read_rtk_lifetime_stats()
|
||||
|
||||
|
||||
async def initialize_context_tool_session_baseline() -> None:
|
||||
"""Pin the current context-tool counters as the proxy-session baseline."""
|
||||
|
||||
tool = _selected_context_tool()
|
||||
payload = await asyncio.to_thread(_read_context_tool_lifetime_stats, tool)
|
||||
with _context_tool_stats_cache_lock:
|
||||
if payload is None or not payload.get("installed", False):
|
||||
# Failed or tool-absent read: defer the pin to the first
|
||||
# successful read (guarded lazy-init) — pinning zeros here would
|
||||
# inflate session savings by the tool's whole lifetime once it
|
||||
# recovers or gets installed.
|
||||
_context_tool_session_baseline.update(
|
||||
{
|
||||
"initialized": False,
|
||||
"tool": tool,
|
||||
"total_commands": 0,
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
"tokens_saved": 0,
|
||||
"total_time_ms": 0,
|
||||
"captured_at": time.time(),
|
||||
}
|
||||
)
|
||||
else:
|
||||
_context_tool_session_baseline.update(
|
||||
{
|
||||
"initialized": True,
|
||||
"tool": tool,
|
||||
"total_commands": int(payload.get("total_commands", 0) or 0),
|
||||
"input_tokens": int(payload.get("input_tokens", 0) or 0),
|
||||
"output_tokens": int(payload.get("output_tokens", 0) or 0),
|
||||
"tokens_saved": int(payload.get("tokens_saved", 0) or 0),
|
||||
"total_time_ms": int(payload.get("total_time_ms", 0) or 0),
|
||||
"captured_at": time.time(),
|
||||
}
|
||||
)
|
||||
_context_tool_stats_cache.update(
|
||||
{
|
||||
"expires_at": 0.0,
|
||||
"has_value": False,
|
||||
"tool": None,
|
||||
"value": None,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def initialize_rtk_session_baseline() -> None:
|
||||
"""Backward-compatible alias for initialize_context_tool_session_baseline."""
|
||||
|
||||
await initialize_context_tool_session_baseline()
|
||||
|
||||
|
||||
def _get_context_tool_stats() -> dict[str, Any] | None:
|
||||
"""Get context-tool savings for the current Headroom proxy session.
|
||||
|
||||
RTK and lean-ctx persist project-level lifetime counters. Dashboard stats
|
||||
should be session-local, so we subtract the counter snapshot captured at
|
||||
proxy startup instead of resetting the tool's own history.
|
||||
"""
|
||||
|
||||
tool = _selected_context_tool()
|
||||
now = time.monotonic()
|
||||
with _context_tool_stats_cache_lock:
|
||||
cached_value = cast(dict[str, Any] | None, _context_tool_stats_cache["value"])
|
||||
if (
|
||||
_context_tool_stats_cache["has_value"]
|
||||
and now < float(_context_tool_stats_cache["expires_at"])
|
||||
and _context_tool_stats_cache.get("tool") == tool
|
||||
):
|
||||
return cached_value
|
||||
|
||||
payload = _read_context_tool_lifetime_stats(tool)
|
||||
with _context_tool_stats_cache_lock:
|
||||
# Baseline mutations only happen on successful reads from an
|
||||
# installed tool — a failed read (None) or a tool-absent zero payload
|
||||
# must never pin or re-pin, or session deltas inflate by the whole
|
||||
# lifetime when the tool comes back.
|
||||
tool_installed = payload is not None and bool(payload.get("installed", False))
|
||||
if (
|
||||
payload is not None
|
||||
and tool_installed
|
||||
and (
|
||||
not _context_tool_session_baseline["initialized"]
|
||||
or _context_tool_session_baseline.get("tool") != tool
|
||||
)
|
||||
):
|
||||
_context_tool_session_baseline.update(
|
||||
{
|
||||
"initialized": True,
|
||||
"tool": tool,
|
||||
"total_commands": int(payload.get("total_commands", 0) or 0),
|
||||
"input_tokens": int(payload.get("input_tokens", 0) or 0),
|
||||
"output_tokens": int(payload.get("output_tokens", 0) or 0),
|
||||
"tokens_saved": int(payload.get("tokens_saved", 0) or 0),
|
||||
"total_time_ms": int(payload.get("total_time_ms", 0) or 0),
|
||||
"captured_at": time.time(),
|
||||
}
|
||||
)
|
||||
|
||||
if payload is not None:
|
||||
lifetime_total_commands = int(payload.get("total_commands", 0) or 0)
|
||||
lifetime_input_tokens = int(payload.get("input_tokens", 0) or 0)
|
||||
lifetime_output_tokens = int(payload.get("output_tokens", 0) or 0)
|
||||
lifetime_tokens_saved = int(payload.get("tokens_saved", 0) or 0)
|
||||
lifetime_total_time_ms = int(payload.get("total_time_ms", 0) or 0)
|
||||
baseline_total_commands = int(_context_tool_session_baseline["total_commands"])
|
||||
baseline_input_tokens = int(_context_tool_session_baseline["input_tokens"])
|
||||
baseline_output_tokens = int(_context_tool_session_baseline["output_tokens"])
|
||||
baseline_tokens_saved = int(_context_tool_session_baseline["tokens_saved"])
|
||||
baseline_total_time_ms = int(_context_tool_session_baseline["total_time_ms"])
|
||||
# A tool-absent payload carries zero counters that are not a
|
||||
# genuine external reset — only successful installed reads may
|
||||
# re-pin the baseline.
|
||||
counter_reset_detected = tool_installed and (
|
||||
lifetime_total_commands < baseline_total_commands
|
||||
or lifetime_input_tokens < baseline_input_tokens
|
||||
or lifetime_output_tokens < baseline_output_tokens
|
||||
or lifetime_tokens_saved < baseline_tokens_saved
|
||||
or lifetime_total_time_ms < baseline_total_time_ms
|
||||
)
|
||||
if counter_reset_detected:
|
||||
baseline_total_commands = lifetime_total_commands
|
||||
baseline_input_tokens = lifetime_input_tokens
|
||||
baseline_output_tokens = lifetime_output_tokens
|
||||
baseline_tokens_saved = lifetime_tokens_saved
|
||||
baseline_total_time_ms = lifetime_total_time_ms
|
||||
_context_tool_session_baseline.update(
|
||||
{
|
||||
"total_commands": baseline_total_commands,
|
||||
"input_tokens": baseline_input_tokens,
|
||||
"output_tokens": baseline_output_tokens,
|
||||
"tokens_saved": baseline_tokens_saved,
|
||||
"total_time_ms": baseline_total_time_ms,
|
||||
"captured_at": time.time(),
|
||||
}
|
||||
)
|
||||
|
||||
session_total_commands = max(lifetime_total_commands - baseline_total_commands, 0)
|
||||
session_input_tokens = max(lifetime_input_tokens - baseline_input_tokens, 0)
|
||||
session_output_tokens = max(lifetime_output_tokens - baseline_output_tokens, 0)
|
||||
session_tokens_saved = max(lifetime_tokens_saved - baseline_tokens_saved, 0)
|
||||
session_total_time_ms = max(lifetime_total_time_ms - baseline_total_time_ms, 0)
|
||||
session_savings_pct = (
|
||||
round(session_tokens_saved / session_input_tokens * 100.0, 4)
|
||||
if session_input_tokens > 0
|
||||
else None
|
||||
)
|
||||
session_avg_time_ms = (
|
||||
round(session_total_time_ms / session_total_commands, 2)
|
||||
if session_total_commands > 0 and session_total_time_ms > 0
|
||||
else None
|
||||
)
|
||||
lifetime_savings_pct = float(payload.get("lifetime_avg_savings_pct", 0.0) or 0.0)
|
||||
|
||||
payload = {
|
||||
**payload,
|
||||
"tool": tool,
|
||||
"label": _context_tool_label(tool),
|
||||
# Backward-compatible session-delta fields.
|
||||
"total_commands": session_total_commands,
|
||||
"input_tokens": session_input_tokens,
|
||||
"output_tokens": session_output_tokens,
|
||||
"tokens_saved": session_tokens_saved,
|
||||
"total_time_ms": session_total_time_ms,
|
||||
"session_savings_pct": session_savings_pct,
|
||||
"session_avg_time_ms": session_avg_time_ms,
|
||||
# Keep old field for compatibility, but declare its scope.
|
||||
"avg_savings_pct": lifetime_savings_pct,
|
||||
"avg_savings_pct_scope": "lifetime",
|
||||
"lifetime_avg_savings_pct": lifetime_savings_pct,
|
||||
"lifetime_total_commands": lifetime_total_commands,
|
||||
"lifetime_input_tokens": lifetime_input_tokens,
|
||||
"lifetime_output_tokens": lifetime_output_tokens,
|
||||
"lifetime_tokens_saved": lifetime_tokens_saved,
|
||||
"lifetime_total_time_ms": lifetime_total_time_ms,
|
||||
"session_baseline_total_commands": baseline_total_commands,
|
||||
"session_baseline_input_tokens": baseline_input_tokens,
|
||||
"session_baseline_output_tokens": baseline_output_tokens,
|
||||
"session_baseline_tokens_saved": baseline_tokens_saved,
|
||||
"session_baseline_total_time_ms": baseline_total_time_ms,
|
||||
"session_baseline_captured_at": _context_tool_session_baseline.get(
|
||||
"captured_at", 0.0
|
||||
),
|
||||
"session": {
|
||||
"commands": session_total_commands,
|
||||
"input_tokens": session_input_tokens,
|
||||
"output_tokens": session_output_tokens,
|
||||
"tokens_saved": session_tokens_saved,
|
||||
"savings_pct": session_savings_pct,
|
||||
"total_time_ms": session_total_time_ms,
|
||||
"avg_time_ms": session_avg_time_ms,
|
||||
},
|
||||
"lifetime": {
|
||||
"commands": lifetime_total_commands,
|
||||
"input_tokens": lifetime_input_tokens,
|
||||
"output_tokens": lifetime_output_tokens,
|
||||
"tokens_saved": lifetime_tokens_saved,
|
||||
"savings_pct": lifetime_savings_pct,
|
||||
"total_time_ms": lifetime_total_time_ms,
|
||||
},
|
||||
"baseline": {
|
||||
"commands": baseline_total_commands,
|
||||
"input_tokens": baseline_input_tokens,
|
||||
"output_tokens": baseline_output_tokens,
|
||||
"tokens_saved": baseline_tokens_saved,
|
||||
"total_time_ms": baseline_total_time_ms,
|
||||
"captured_at": _context_tool_session_baseline.get("captured_at", 0.0),
|
||||
},
|
||||
"sampled_at": time.time(),
|
||||
"sample_ttl_seconds": CONTEXT_TOOL_STATS_CACHE_TTL_SECONDS,
|
||||
"refresh_interval_seconds": CONTEXT_TOOL_STATS_CACHE_TTL_SECONDS,
|
||||
"counter_reset_detected": counter_reset_detected,
|
||||
}
|
||||
|
||||
_context_tool_stats_cache.update(
|
||||
{
|
||||
"expires_at": time.monotonic() + CONTEXT_TOOL_STATS_CACHE_TTL_SECONDS,
|
||||
"has_value": True,
|
||||
"tool": tool,
|
||||
"value": payload,
|
||||
}
|
||||
)
|
||||
return payload
|
||||
|
||||
|
||||
def _get_rtk_stats() -> dict[str, Any] | None:
|
||||
"""Backward-compatible alias for selected context-tool stats."""
|
||||
|
||||
return _get_context_tool_stats()
|
||||
|
||||
|
||||
def is_anthropic_auth(headers: dict[str, str]) -> bool:
|
||||
"""Detect Anthropic auth signals in request headers."""
|
||||
if headers.get("x-api-key") or headers.get("anthropic-version"):
|
||||
|
|
|
|||
|
|
@ -1554,32 +1554,4 @@ class PrometheusMetrics:
|
|||
value=redactions_total(),
|
||||
)
|
||||
|
||||
# Phase G PR-G3 remediation (C4): RTK invocations counter
|
||||
# also lives Python-side. RTK is wrapped by the
|
||||
# `headroom wrap` CLI (headroom.cli.wrap); the proxy
|
||||
# observes invocation counts via a process-local tracker
|
||||
# the wrap tail bumps. The Rust proxy previously held a
|
||||
# dead counter for this; that's been removed.
|
||||
from headroom.cli.wrap_rtk_metrics import rtk_invocation_counts
|
||||
|
||||
counts = rtk_invocation_counts()
|
||||
lines.extend(
|
||||
[
|
||||
"# HELP wrap_rtk_invocations_total RTK invocations observed via the wrap CLI tail",
|
||||
"# TYPE wrap_rtk_invocations_total counter",
|
||||
]
|
||||
)
|
||||
if not counts:
|
||||
# Emit a zero-row under the sentinel tool name so
|
||||
# the family advertises HELP/TYPE on a fresh boot
|
||||
# and dashboards can probe it before any RTK
|
||||
# invocation has happened. Matches the Rust side's
|
||||
# H3 force-zero contract.
|
||||
lines.append('wrap_rtk_invocations_total{tool="__init__"} 0')
|
||||
else:
|
||||
for tool, count in counts.items():
|
||||
safe_tool = _escape_label_value(str(tool))
|
||||
lines.append(f'wrap_rtk_invocations_total{{tool="{safe_tool}"}} {count}')
|
||||
lines.append("")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
|
|
|||
|
|
@ -133,12 +133,9 @@ from headroom.proxy.helpers import (
|
|||
MAX_REQUEST_BODY_SIZE, # noqa: F401
|
||||
MAX_SSE_BUFFER_SIZE, # noqa: F401
|
||||
RETRYABLE_OVERLOAD_STATUSES,
|
||||
_get_context_tool_stats,
|
||||
_get_image_compressor, # noqa: F401
|
||||
_get_rtk_stats, # noqa: F401
|
||||
_read_request_json, # noqa: F401
|
||||
_setup_file_logging, # noqa: F401
|
||||
initialize_context_tool_session_baseline,
|
||||
is_anthropic_auth, # noqa: F401
|
||||
jitter_delay_ms,
|
||||
resolve_display_provider,
|
||||
|
|
@ -2484,7 +2481,6 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
app.state.started_at = time.time()
|
||||
app.state.ready = False
|
||||
app.state.startup_error = None
|
||||
await initialize_context_tool_session_baseline()
|
||||
|
||||
try:
|
||||
try:
|
||||
|
|
@ -3640,39 +3636,11 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
# Build prefix cache stats once (used in both prefix_cache and cost)
|
||||
prefix_cache_stats = _build_prefix_cache_stats(m, proxy.cost_tracker)
|
||||
|
||||
# Fetch CLI filtering savings from the selected context tool. These
|
||||
# tokens are avoided before they reach model context.
|
||||
cli_filtering_stats = await asyncio.to_thread(_get_context_tool_stats)
|
||||
cli_filtering_tool = (
|
||||
str(cli_filtering_stats.get("tool", "rtk")) if cli_filtering_stats else "rtk"
|
||||
)
|
||||
cli_filtering_label = (
|
||||
str(cli_filtering_stats.get("label", "RTK")) if cli_filtering_stats else "RTK"
|
||||
)
|
||||
cli_tokens_avoided = (
|
||||
cli_filtering_stats.get("tokens_saved", 0) if cli_filtering_stats else 0
|
||||
)
|
||||
cli_filtering_session = (
|
||||
cli_filtering_stats.get("session", {}) if cli_filtering_stats else {}
|
||||
)
|
||||
cli_filtering_lifetime = (
|
||||
cli_filtering_stats.get("lifetime", {}) if cli_filtering_stats else {}
|
||||
)
|
||||
rtk_tokens_avoided = cli_tokens_avoided if cli_filtering_tool == "rtk" else 0
|
||||
lean_ctx_tokens_avoided = cli_tokens_avoided if cli_filtering_tool == "lean-ctx" else 0
|
||||
cli_filtering_available = bool(
|
||||
cli_filtering_stats and cli_filtering_stats.get("installed", False)
|
||||
)
|
||||
|
||||
# Calculate total tokens before Headroom-side reduction. Proxy
|
||||
# compression and the configured context tool both remove tokens before
|
||||
# they reach model context, so dashboard-facing savings combines them.
|
||||
# Calculate total tokens before Headroom-side reduction.
|
||||
proxy_compression_tokens = m.tokens_saved_total
|
||||
# "All layers" must include tool-schema deferral (the tool_search layer
|
||||
# enumerated in by_layer below) — otherwise the advertised total omits it.
|
||||
all_layers_tokens_saved = (
|
||||
proxy_compression_tokens + cli_tokens_avoided + m.tool_search_saved_total
|
||||
)
|
||||
all_layers_tokens_saved = proxy_compression_tokens + m.tool_search_saved_total
|
||||
total_tokens_before = m.tokens_input_total + all_layers_tokens_saved
|
||||
proxy_total_before_compression = m.tokens_input_total + proxy_compression_tokens
|
||||
# `attempted_input_tokens` is the compressible-only denominator
|
||||
|
|
@ -3703,9 +3671,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
)
|
||||
|
||||
# Build human-readable summary
|
||||
summary = _build_session_summary(
|
||||
proxy, m, prefix_cache_stats, cli_tokens_avoided, total_tokens_before
|
||||
)
|
||||
summary = _build_session_summary(proxy, m, prefix_cache_stats, total_tokens_before)
|
||||
# DEBUG: log the summary payload for external upsert consumers
|
||||
try:
|
||||
logger.debug("/stats summary data: %r", summary)
|
||||
|
|
@ -3804,47 +3770,11 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
"total_tokens": total_tokens_all_layers,
|
||||
"per_project": persistent_savings.get("projects", {}),
|
||||
"by_layer": {
|
||||
"cli_filtering": {
|
||||
"tool": cli_filtering_tool,
|
||||
"label": cli_filtering_label,
|
||||
"available": cli_filtering_available,
|
||||
"tokens": cli_tokens_avoided,
|
||||
"tokens_saved": cli_tokens_avoided,
|
||||
"session": cli_filtering_session,
|
||||
"lifetime": cli_filtering_lifetime,
|
||||
"session_savings_pct": (
|
||||
cli_filtering_stats.get("session_savings_pct")
|
||||
if cli_filtering_stats
|
||||
else None
|
||||
),
|
||||
"lifetime_savings_pct": (
|
||||
cli_filtering_stats.get("lifetime_avg_savings_pct")
|
||||
if cli_filtering_stats
|
||||
else None
|
||||
),
|
||||
"refresh_interval_seconds": (
|
||||
cli_filtering_stats.get("refresh_interval_seconds")
|
||||
if cli_filtering_stats
|
||||
else None
|
||||
),
|
||||
"included_in": "tokens.saved",
|
||||
"description": (
|
||||
f"Tokens avoided by CLI output filtering ({cli_filtering_label}) "
|
||||
"before reaching context. "
|
||||
"Included in dashboard token savings, but not in dollar savings."
|
||||
),
|
||||
},
|
||||
"compression": {
|
||||
"tokens": proxy_compression_tokens,
|
||||
"proxy_tokens": proxy_compression_tokens,
|
||||
"cli_filtering_tokens": cli_tokens_avoided,
|
||||
"rtk_tokens": rtk_tokens_avoided,
|
||||
"lean_ctx_tokens": lean_ctx_tokens_avoided,
|
||||
"all_layers_tokens": all_layers_tokens_saved,
|
||||
"description": (
|
||||
"Tokens removed by Headroom proxy compression. "
|
||||
"Dashboard token savings also includes CLI context-tool filtering."
|
||||
),
|
||||
"description": ("Tokens removed by Headroom proxy compression."),
|
||||
},
|
||||
"prefix_cache": {
|
||||
"discount_usd": round(cache_net_usd, 4),
|
||||
|
|
@ -3896,10 +3826,6 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
"output_reduction": output_reduction,
|
||||
"saved": all_layers_tokens_saved,
|
||||
"proxy_compression_saved": proxy_compression_tokens,
|
||||
"cli_filtering_saved": cli_tokens_avoided,
|
||||
"rtk_saved": rtk_tokens_avoided,
|
||||
"lean_ctx_saved": lean_ctx_tokens_avoided,
|
||||
"cli_tokens_avoided": cli_tokens_avoided,
|
||||
"proxy_total_before_compression": proxy_total_before_compression,
|
||||
"total_before_compression": total_tokens_before,
|
||||
"all_layers_saved": all_layers_tokens_saved,
|
||||
|
|
@ -4073,7 +3999,6 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
"cost": _merge_cost_stats(
|
||||
proxy.cost_tracker.stats() if proxy.cost_tracker else None,
|
||||
prefix_cache_stats,
|
||||
cli_tokens_avoided=cli_tokens_avoided,
|
||||
),
|
||||
"compression": {
|
||||
"ccr_entries": compression_stats.get("entry_count", 0),
|
||||
|
|
@ -4109,13 +4034,6 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
),
|
||||
},
|
||||
"toin": get_toin().get_stats(),
|
||||
"context_tool": {
|
||||
"configured": cli_filtering_tool,
|
||||
"label": cli_filtering_label,
|
||||
"available": cli_filtering_available,
|
||||
"stats": cli_filtering_stats,
|
||||
},
|
||||
"cli_filtering": cli_filtering_stats,
|
||||
"proxy_inbound": proxy.metrics.inbound_snapshot(),
|
||||
"cache": await proxy.cache.stats() if proxy.cache else None,
|
||||
"rate_limiter": await proxy.rate_limiter.stats() if proxy.rate_limiter else None,
|
||||
|
|
@ -4237,7 +4155,6 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
await proxy.metrics.reset_runtime()
|
||||
if proxy.cost_tracker:
|
||||
proxy.cost_tracker.reset_runtime()
|
||||
await initialize_context_tool_session_baseline()
|
||||
async with _stats_snapshot_lock:
|
||||
_stats_snapshot["value"] = None
|
||||
_stats_snapshot["expires_at"] = 0.0
|
||||
|
|
@ -4249,17 +4166,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
series: Literal["history", "hourly", "daily", "weekly", "monthly"] = "history",
|
||||
history_mode: Literal["compact", "full", "none"] = "compact",
|
||||
):
|
||||
"""Get durable proxy compression history plus display-session state.
|
||||
|
||||
The JSON payload also carries a ``cli_filtering`` key with live RTK
|
||||
stats. This is a curated subset (``tool``, ``label``, ``available``,
|
||||
``lifetime``, ``session``) tailored to the Historical tab, not the
|
||||
full ``_get_context_tool_stats()`` payload that ``/stats`` exposes.
|
||||
It is ``None`` only when the stats read hard-fails; when the tool is
|
||||
merely absent, ``cli_filtering`` stays populated with
|
||||
``available: False`` and zeroed counters so the tab can distinguish
|
||||
"not installed" from "installed, no data yet."
|
||||
"""
|
||||
"""Get durable proxy compression history plus display-session state."""
|
||||
if format == "csv":
|
||||
filename = f"headroom-stats-history-{series}.csv"
|
||||
return Response(
|
||||
|
|
@ -4268,31 +4175,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
headers={"Content-Disposition": f'attachment; filename="{filename}"'},
|
||||
)
|
||||
|
||||
history = proxy.metrics.savings_tracker.history_response(history_mode=history_mode)
|
||||
|
||||
# Augment with live RTK/cli-filtering lifetime stats so the Historical
|
||||
# tab can display them. These live in the context-tool's own stats file
|
||||
# and survive proxy restarts — exactly what the Historical tab needs.
|
||||
# Best-effort: if the RTK stats file can't be read (missing, parse error,
|
||||
# IO), fall back to None so the Historical tab stays available instead of
|
||||
# 500ing. The tab hides the card when cli_filtering is None.
|
||||
try:
|
||||
cli_stats = await asyncio.to_thread(_get_context_tool_stats)
|
||||
except Exception:
|
||||
logger.debug("stats-history: RTK stats unavailable", exc_info=True)
|
||||
cli_stats = None
|
||||
if cli_stats:
|
||||
history["cli_filtering"] = {
|
||||
"tool": str(cli_stats.get("tool", "rtk")),
|
||||
"label": str(cli_stats.get("label", "RTK")),
|
||||
"available": bool(cli_stats.get("installed", False)),
|
||||
"lifetime": cli_stats.get("lifetime", {}),
|
||||
"session": cli_stats.get("session", {}),
|
||||
}
|
||||
else:
|
||||
history["cli_filtering"] = None
|
||||
|
||||
return history
|
||||
return proxy.metrics.savings_tracker.history_response(history_mode=history_mode)
|
||||
|
||||
@app.get("/transformations/feed", dependencies=[Depends(_require_loopback)])
|
||||
async def transformations_feed(limit: int = 20):
|
||||
|
|
|
|||
|
|
@ -1,48 +0,0 @@
|
|||
"""rtk (Rust Token Killer) integration for Headroom.
|
||||
|
||||
rtk compresses CLI output (test results, git diffs, log dumps) before it
|
||||
enters the LLM context window. Headroom downloads and manages the rtk binary.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import platform
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
from headroom import paths as _paths
|
||||
|
||||
RTK_VERSION = "v0.42.4"
|
||||
RTK_BIN_DIR = _paths.bin_dir()
|
||||
_RTK_NAME = "rtk.exe" if platform.system() == "Windows" else "rtk"
|
||||
RTK_BIN_PATH = RTK_BIN_DIR / _RTK_NAME
|
||||
|
||||
|
||||
def _managed_rtk_candidates() -> list[Path]:
|
||||
"""Return known Headroom-managed rtk binary paths."""
|
||||
candidates = [RTK_BIN_DIR / _RTK_NAME]
|
||||
for name in ("rtk", "rtk.exe"):
|
||||
path = RTK_BIN_DIR / name
|
||||
if path not in candidates:
|
||||
candidates.append(path)
|
||||
return candidates
|
||||
|
||||
|
||||
def get_rtk_path() -> Path | None:
|
||||
"""Get path to rtk binary — check PATH first, then ~/.headroom/bin/."""
|
||||
# Check if rtk is already in PATH (e.g., installed via brew)
|
||||
system_rtk = shutil.which("rtk")
|
||||
if system_rtk:
|
||||
return Path(system_rtk)
|
||||
|
||||
# Check Headroom-managed install
|
||||
for candidate in _managed_rtk_candidates():
|
||||
if candidate.exists() and candidate.is_file():
|
||||
return candidate
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def is_rtk_installed() -> bool:
|
||||
"""Check if rtk is available."""
|
||||
return get_rtk_path() is not None
|
||||
|
|
@ -1,247 +0,0 @@
|
|||
"""Download and install rtk binary from GitHub releases."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
import platform
|
||||
import stat
|
||||
import subprocess
|
||||
import tarfile
|
||||
import tempfile
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from urllib.request import urlopen
|
||||
|
||||
from headroom._subprocess import run
|
||||
|
||||
from . import RTK_BIN_DIR, RTK_BIN_PATH, RTK_VERSION
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
GITHUB_RELEASE_URL = "https://github.com/rtk-ai/rtk/releases/download"
|
||||
|
||||
|
||||
def _detect_runtime_target_triple() -> str:
|
||||
"""Detect platform and return the rtk release target triple."""
|
||||
system = platform.system()
|
||||
machine = platform.machine()
|
||||
|
||||
if system == "Darwin":
|
||||
arch = "aarch64" if machine == "arm64" else "x86_64"
|
||||
return f"{arch}-apple-darwin"
|
||||
elif system == "Linux":
|
||||
arch = "aarch64" if machine == "aarch64" else "x86_64"
|
||||
suffix = "unknown-linux-gnu" if arch == "aarch64" else "unknown-linux-musl"
|
||||
return f"{arch}-{suffix}"
|
||||
elif system == "Windows":
|
||||
return "x86_64-pc-windows-msvc"
|
||||
|
||||
raise RuntimeError(f"Unsupported platform: {system} {machine}")
|
||||
|
||||
|
||||
def _get_target_triple() -> str:
|
||||
"""Return the requested rtk target triple, honoring explicit overrides."""
|
||||
return os.environ.get("HEADROOM_RTK_TARGET", "").strip() or _detect_runtime_target_triple()
|
||||
|
||||
|
||||
def _binary_name_for_target(target: str) -> str:
|
||||
"""Return the expected binary name for a target triple."""
|
||||
return "rtk.exe" if "windows" in target else "rtk"
|
||||
|
||||
|
||||
def _should_verify_target(target: str) -> bool:
|
||||
"""Verify only when the requested target matches the current runtime."""
|
||||
return target == _detect_runtime_target_triple()
|
||||
|
||||
|
||||
def _get_download_url(version: str) -> tuple[str, str]:
|
||||
"""Get download URL and extension for this platform.
|
||||
|
||||
Returns (url, extension) where extension is 'tar.gz' or 'zip'.
|
||||
"""
|
||||
target = _get_target_triple()
|
||||
|
||||
if "windows" in target:
|
||||
ext = "zip"
|
||||
else:
|
||||
ext = "tar.gz"
|
||||
|
||||
url = f"{GITHUB_RELEASE_URL}/{version}/rtk-{target}.{ext}"
|
||||
return url, ext
|
||||
|
||||
|
||||
def download_rtk(version: str | None = None) -> Path:
|
||||
"""Download rtk binary from GitHub releases.
|
||||
|
||||
Args:
|
||||
version: Version to download (e.g., "v0.42.4"). Defaults to pinned version.
|
||||
|
||||
Returns:
|
||||
Path to the installed binary.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If download or extraction fails.
|
||||
"""
|
||||
version = version or RTK_VERSION
|
||||
target = _get_target_triple()
|
||||
url, ext = _get_download_url(version)
|
||||
target_path = RTK_BIN_DIR / _binary_name_for_target(target)
|
||||
|
||||
RTK_BIN_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
logger.info("Downloading rtk %s from %s ...", version, url)
|
||||
|
||||
try:
|
||||
# Validate URL scheme to prevent B310 warning
|
||||
if not url.startswith(("http://", "https://")):
|
||||
raise ValueError(f"Invalid URL scheme in {url}")
|
||||
|
||||
# Fail closed on TLS errors rather than executing an unverifiable download.
|
||||
try:
|
||||
with urlopen(url, timeout=30) as response:
|
||||
data = response.read()
|
||||
except Exception as download_err:
|
||||
if "CERTIFICATE_VERIFY_FAILED" in str(download_err):
|
||||
raise RuntimeError(
|
||||
"TLS verification failed downloading rtk; fix the local trust store and retry."
|
||||
) from download_err
|
||||
raise
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to download rtk from {url}: {e}") from e
|
||||
|
||||
# Extract binary
|
||||
try:
|
||||
if ext == "tar.gz":
|
||||
with tarfile.open(fileobj=io.BytesIO(data), mode="r:gz") as tar:
|
||||
# Find the rtk binary inside the archive
|
||||
for member in tar.getmembers():
|
||||
if member.name.endswith("/rtk") or member.name == "rtk":
|
||||
member.name = target_path.name # Flatten path
|
||||
tar.extract(member, RTK_BIN_DIR)
|
||||
break
|
||||
else:
|
||||
raise RuntimeError("rtk binary not found in archive")
|
||||
elif ext == "zip":
|
||||
with zipfile.ZipFile(io.BytesIO(data)) as zf:
|
||||
for name in zf.namelist():
|
||||
if name.endswith("rtk.exe") or name.endswith("/rtk"):
|
||||
with zf.open(name) as src, open(target_path, "wb") as dst:
|
||||
dst.write(src.read())
|
||||
break
|
||||
else:
|
||||
raise RuntimeError("rtk binary not found in archive")
|
||||
except (tarfile.TarError, zipfile.BadZipFile) as e:
|
||||
raise RuntimeError(f"Failed to extract rtk archive: {e}") from e
|
||||
|
||||
# Make executable (skip on Windows — no Unix permissions)
|
||||
if "windows" not in target:
|
||||
target_path.chmod(target_path.stat().st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH)
|
||||
|
||||
if _should_verify_target(target):
|
||||
try:
|
||||
result = run(
|
||||
[str(target_path), "--version"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(f"rtk verification failed: {result.stderr}")
|
||||
logger.info("rtk installed: %s", result.stdout.strip())
|
||||
except FileNotFoundError as e:
|
||||
raise RuntimeError("rtk binary not found after extraction") from e
|
||||
except subprocess.TimeoutExpired as e:
|
||||
raise RuntimeError("rtk verification timed out") from e
|
||||
else:
|
||||
logger.info("rtk installed for target %s at %s (verification skipped)", target, target_path)
|
||||
|
||||
return target_path
|
||||
|
||||
|
||||
# Agents rtk registers a *native* hook for via `rtk init --agent <name>`.
|
||||
# For these, headroom must not also inject the RTK_INSTRUCTIONS_BLOCK text
|
||||
# into a rules/instructions file — that duplicates guidance rtk's own hook
|
||||
# already provides silently (GH #756).
|
||||
RTK_NATIVE_HOOK_AGENTS = frozenset(
|
||||
{"claude", "cursor", "windsurf", "cline", "kilocode", "antigravity", "pi", "hermes"}
|
||||
)
|
||||
|
||||
|
||||
def register_claude_hooks(rtk_path: Path | None = None) -> bool:
|
||||
"""Register rtk hooks in Claude Code settings.
|
||||
|
||||
Runs `rtk init --global` which adds a PreToolUse hook to
|
||||
~/.claude/settings.json that rewrites Bash commands through rtk.
|
||||
|
||||
Returns True if hooks were registered successfully.
|
||||
"""
|
||||
return register_agent_hooks(rtk_path, agent="claude")
|
||||
|
||||
|
||||
def register_agent_hooks(rtk_path: Path | None = None, *, agent: str = "claude") -> bool:
|
||||
"""Register rtk's native hook for ``agent`` via ``rtk init --agent``.
|
||||
|
||||
Only agents in ``RTK_NATIVE_HOOK_AGENTS`` support this; callers must not
|
||||
invoke this for agents rtk has no native hook for (rtk itself will just
|
||||
reject the ``--agent`` value).
|
||||
|
||||
Returns True if hooks were registered successfully.
|
||||
"""
|
||||
rtk_path = rtk_path or RTK_BIN_PATH
|
||||
args = [str(rtk_path), "init", "--global", "--auto-patch"]
|
||||
if agent != "claude":
|
||||
args += ["--agent", agent]
|
||||
|
||||
# Capture output to a temp file rather than pipes: `rtk init` may fork a
|
||||
# background process that inherits our stdout/stderr, and a piped
|
||||
# `subprocess.run` drains those pipes until EOF — which never arrives while
|
||||
# the daemon holds them open, so it blocks to the timeout even though
|
||||
# `rtk init` itself exited and already registered the hooks. A file fd has
|
||||
# no such reader, so we wait only on the direct child. stdin is DEVNULL so a
|
||||
# stray prompt can never block either.
|
||||
try:
|
||||
with tempfile.TemporaryFile(mode="w+", encoding="utf-8", errors="replace") as out:
|
||||
try:
|
||||
result = subprocess.run(
|
||||
args,
|
||||
stdin=subprocess.DEVNULL,
|
||||
stdout=out,
|
||||
stderr=out,
|
||||
timeout=10,
|
||||
)
|
||||
except subprocess.TimeoutExpired:
|
||||
# Read the temp file while it is still open — the outer handler
|
||||
# runs after the `with` closes it, so any captured diagnostics
|
||||
# would be gone by then.
|
||||
out.seek(0)
|
||||
logger.warning("rtk init timed out: %s", out.read().strip())
|
||||
return False
|
||||
if result.returncode == 0:
|
||||
logger.info("rtk hooks registered for %s", agent)
|
||||
return True
|
||||
out.seek(0)
|
||||
logger.warning("rtk init failed: %s", out.read().strip())
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.warning("Failed to register rtk hooks: %s", e)
|
||||
return False
|
||||
|
||||
|
||||
def ensure_rtk(version: str | None = None) -> Path | None:
|
||||
"""Ensure rtk is installed — download if needed.
|
||||
|
||||
Returns path to rtk binary, or None if installation failed.
|
||||
"""
|
||||
from . import get_rtk_path
|
||||
|
||||
existing = get_rtk_path()
|
||||
if existing:
|
||||
return existing
|
||||
|
||||
try:
|
||||
return download_rtk(version)
|
||||
except RuntimeError as e:
|
||||
logger.warning("Could not install rtk: %s", e)
|
||||
return None
|
||||
|
|
@ -383,43 +383,30 @@ class HeadroomContribution:
|
|||
tokens_saved_compression: int = 0
|
||||
"""Input tokens removed by proxy compression."""
|
||||
|
||||
tokens_saved_cli_filtering: int = 0
|
||||
"""Tokens avoided by the selected CLI context tool before reaching context."""
|
||||
|
||||
tokens_saved_rtk: int = 0
|
||||
"""Deprecated alias for CLI filtering tokens from older persisted state."""
|
||||
|
||||
tokens_saved_cache_reads: int = 0
|
||||
"""Input tokens served from Anthropic prefix-cache (discounted reads)."""
|
||||
|
||||
compression_savings_usd: float = 0.0
|
||||
cache_savings_usd: float = 0.0
|
||||
|
||||
def cli_filtering_saved(self) -> int:
|
||||
return max(self.tokens_saved_cli_filtering, self.tokens_saved_rtk)
|
||||
|
||||
def total_saved(self) -> int:
|
||||
return (
|
||||
self.tokens_saved_compression
|
||||
+ self.cli_filtering_saved()
|
||||
+ self.tokens_saved_cache_reads
|
||||
)
|
||||
return self.tokens_saved_compression + self.tokens_saved_cache_reads
|
||||
|
||||
def compression_saved(self) -> int:
|
||||
"""Tokens removed before model context by compression plus CLI filtering."""
|
||||
"""Tokens removed before model context by proxy compression."""
|
||||
|
||||
return self.tokens_saved_compression + self.cli_filtering_saved()
|
||||
return self.tokens_saved_compression
|
||||
|
||||
def total_savings_usd(self) -> float:
|
||||
return self.compression_savings_usd + self.cache_savings_usd
|
||||
|
||||
def raw_without_headroom(self) -> int:
|
||||
return self.tokens_submitted + self.tokens_saved_compression + self.cli_filtering_saved()
|
||||
return self.tokens_submitted + self.tokens_saved_compression
|
||||
|
||||
def efficiency_pct(self) -> float:
|
||||
# Fraction of the pre-Headroom input that compression + CLI filtering
|
||||
# removed. Use compression_saved() (which excludes cache reads) as the
|
||||
# numerator: raw_without_headroom() also excludes cache reads, so mixing
|
||||
# Fraction of the pre-Headroom input that compression removed. Use
|
||||
# compression_saved() (which excludes cache reads) as the numerator:
|
||||
# raw_without_headroom() also excludes cache reads, so mixing
|
||||
# total_saved() (which adds tokens_saved_cache_reads) into this ratio
|
||||
# made it inconsistent with its own denominator and let efficiency exceed
|
||||
# 100% (e.g. submitted=100, cache_reads=1000 -> 1000%). Cache reads are a
|
||||
|
|
@ -431,22 +418,17 @@ class HeadroomContribution:
|
|||
return round(self.compression_saved() / raw * 100, 1)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
# Older persisted payloads also carry ``cli_filtering`` / ``rtk`` /
|
||||
# ``cli_filtering_raw`` / ``rtk_raw`` keys from the retired CLI
|
||||
# context-tool layer. They are simply no longer written or read;
|
||||
# SubscriptionTracker loads state key-by-key with ``dict.get`` (never
|
||||
# ``HeadroomContribution(**payload)``), so an on-disk payload carrying
|
||||
# those extra keys still loads without raising.
|
||||
return {
|
||||
"tokens_submitted": self.tokens_submitted,
|
||||
"tokens_saved": {
|
||||
"compression": self.compression_saved(),
|
||||
"proxy_compression": self.tokens_saved_compression,
|
||||
"cli_filtering": self.cli_filtering_saved(),
|
||||
"rtk": self.cli_filtering_saved(),
|
||||
# PR-G2 (Realignment) — raw counters, distinct from the
|
||||
# dashboard-facing ``cli_filtering`` / ``rtk`` keys (which
|
||||
# both report ``max(cli_filtering, rtk)`` for legacy
|
||||
# display). Persisted so the tracker can round-trip each
|
||||
# counter independently — the bug PR-G2 retires is that
|
||||
# ``tokens_saved_rtk`` and ``tokens_saved_cli_filtering``
|
||||
# used to be identical.
|
||||
"cli_filtering_raw": self.tokens_saved_cli_filtering,
|
||||
"rtk_raw": self.tokens_saved_rtk,
|
||||
"cache_reads": self.tokens_saved_cache_reads,
|
||||
"total": self.total_saved(),
|
||||
},
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ import threading
|
|||
import time
|
||||
from datetime import timedelta
|
||||
from pathlib import Path
|
||||
from typing import Any, cast
|
||||
from typing import Any
|
||||
|
||||
from headroom import paths as _paths
|
||||
from headroom.subscription.base import QuotaTracker
|
||||
|
|
@ -53,63 +53,6 @@ _PERSIST_FILE_ENV = _paths.HEADROOM_SUBSCRIPTION_STATE_PATH_ENV
|
|||
_DEFAULT_PERSIST_DIR = ".headroom"
|
||||
_DEFAULT_PERSIST_FILE = "subscription_state.json"
|
||||
|
||||
# PR-G2 (Realignment) — RTK savings wiring.
|
||||
#
|
||||
# Operators can disable the RTK polling from inside ``update_contribution``
|
||||
# without uninstalling the binary or unsetting ``HEADROOM_CONTEXT_TOOL``.
|
||||
# Used for diagnostics and for environments where RTK is intentionally
|
||||
# excluded from headroom accounting (e.g. shadow tests).
|
||||
#
|
||||
# Loud / configurable / no silent fallback: unknown values raise loudly via
|
||||
# the parser below — they do not silently default to ``enabled``.
|
||||
_RTK_WIRING_ENV = "HEADROOM_RTK_WIRING"
|
||||
_RTK_WIRING_DEFAULT = "enabled"
|
||||
_RTK_WIRING_ALLOWED = ("enabled", "disabled")
|
||||
|
||||
# PR-G2 remediation (C3) — multi-worker poll ownership.
|
||||
#
|
||||
# Each uvicorn worker independently runs ``configure_subscription_tracker``,
|
||||
# so each worker would poll RTK and add the same delta to its own
|
||||
# ``c.tokens_saved_rtk``. The persisted state is shared by atomic
|
||||
# os.replace, but the in-memory counters diverge per worker — and any
|
||||
# dashboard hitting a non-owner worker would see drifted values.
|
||||
#
|
||||
# The owner-election strategy mirrors the beacon's file-lock pattern in
|
||||
# ``headroom/proxy/server.py``: a non-blocking ``fcntl.flock`` on
|
||||
# ``HEADROOM_RTK_POLL_LOCK`` (default under the workspace dir). Only the
|
||||
# lock holder polls; non-owners return 0 from ``_poll_rtk_delta`` and
|
||||
# delegate to whatever the owner writes to the shared state file.
|
||||
#
|
||||
# Loud / no silent fallback: when ``fcntl`` is unavailable (Windows), every
|
||||
# worker polls — but the explicit ``WindowsNoLockMode`` log line surfaces
|
||||
# the choice so operators see it in startup logs.
|
||||
_RTK_POLL_LOCK_ENV = "HEADROOM_RTK_POLL_LOCK"
|
||||
|
||||
|
||||
def _rtk_wiring_mode() -> str:
|
||||
"""Return ``enabled`` or ``disabled``. Raises on unknown values.
|
||||
|
||||
Read at call-time so operators can flip the env var without a restart.
|
||||
"""
|
||||
raw = os.environ.get(_RTK_WIRING_ENV, "").strip().lower()
|
||||
if not raw:
|
||||
return _RTK_WIRING_DEFAULT
|
||||
if raw in _RTK_WIRING_ALLOWED:
|
||||
return raw
|
||||
raise ValueError(f"Invalid {_RTK_WIRING_ENV}={raw!r}; expected one of {_RTK_WIRING_ALLOWED}")
|
||||
|
||||
|
||||
def _validate_rtk_env_at_startup() -> None:
|
||||
"""Validate RTK env vars eagerly at proxy startup.
|
||||
|
||||
Raises ``ValueError`` loudly if ``HEADROOM_RTK_WIRING`` is set to an
|
||||
invalid value. PR-G2 remediation (H1): previously a typo at startup
|
||||
would silently default to enabled but get swallowed at every
|
||||
``update_contribution`` call — fail loudly here instead.
|
||||
"""
|
||||
_rtk_wiring_mode()
|
||||
|
||||
|
||||
# Singleton on-demand poll floor (seconds): the dashboard may request a fresh
|
||||
# poll if the cached snapshot is stale, but we cap how often we will actually
|
||||
# hit Anthropic to avoid 429s / OAuth-token flagging. Bounded across users.
|
||||
|
|
@ -186,34 +129,6 @@ class SubscriptionTracker(QuotaTracker):
|
|||
self._current_token: str | None = None
|
||||
self._full_tokens: dict[str, int] = {} # token_prefix -> count of requests
|
||||
|
||||
# PR-G2 (Realignment) — most recent session-incremental ``tokens_saved``
|
||||
# observed in ``_get_rtk_stats()['session']['tokens_saved']``. This
|
||||
# field is de-baselined by the helper (see
|
||||
# :func:`~headroom.proxy.helpers._get_context_tool_stats`) so it
|
||||
# already represents savings accumulated since the proxy session
|
||||
# baseline was pinned at startup.
|
||||
#
|
||||
# PR-G2 remediation (C1): previously we read
|
||||
# ``lifetime_tokens_saved`` (the raw monotonic counter from
|
||||
# ``rtk gain --project``), which on the very first poll emitted the
|
||||
# entire pre-Headroom RTK history as one fake delta. Switching to
|
||||
# the session-incremental field dissolves both the first-poll
|
||||
# phantom and the post-restart phantom (the helper rebaselines
|
||||
# session counters at every proxy startup, so a fresh process sees
|
||||
# ``session.tokens_saved == 0`` until new RTK invocations land).
|
||||
#
|
||||
# Monotonic non-decreasing within a proxy session: only advances on
|
||||
# positive delta and on explicit counter-reset detection (session
|
||||
# value drops below the last seen value).
|
||||
self._last_rtk_tokens_saved: int = 0
|
||||
|
||||
# PR-G2 remediation (C3) — owner-election state for multi-worker
|
||||
# poll deduplication. None means we haven't tried to elect yet; True
|
||||
# means this worker holds the lock and polls; False means another
|
||||
# worker owns the lock and we skip polling.
|
||||
self._rtk_poll_owner: bool | None = None
|
||||
self._rtk_poll_lock_fd: Any = None
|
||||
|
||||
self._stop_event: asyncio.Event | None = None
|
||||
self._poll_task: asyncio.Task[None] | None = None
|
||||
|
||||
|
|
@ -258,9 +173,6 @@ class SubscriptionTracker(QuotaTracker):
|
|||
except (asyncio.TimeoutError, asyncio.CancelledError):
|
||||
self._poll_task.cancel()
|
||||
self._persist_state()
|
||||
# PR-G2 remediation (C3): release the poll lock so a subsequent
|
||||
# process / worker restart can re-elect the owner.
|
||||
self._release_rtk_poll_lock()
|
||||
logger.info("Subscription tracker stopped")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
|
|
@ -290,8 +202,6 @@ class SubscriptionTracker(QuotaTracker):
|
|||
*,
|
||||
tokens_submitted: int = 0,
|
||||
tokens_saved_compression: int = 0,
|
||||
tokens_saved_cli_filtering: int | None = None,
|
||||
tokens_saved_rtk: int | None = None,
|
||||
tokens_saved_cache_reads: int = 0,
|
||||
compression_savings_usd: float = 0.0,
|
||||
cache_savings_usd: float = 0.0,
|
||||
|
|
@ -299,257 +209,15 @@ class SubscriptionTracker(QuotaTracker):
|
|||
"""Update headroom contribution counters for the current session window.
|
||||
|
||||
Called after each proxy request completes with the actual token deltas.
|
||||
|
||||
PR-G2 (Realignment) — ``tokens_saved_rtk`` is now sourced from RTK's
|
||||
own stats endpoint (``rtk gain --format json`` via
|
||||
:func:`headroom.proxy.helpers._get_rtk_stats`) when the caller does
|
||||
not pass an explicit value. The tracker computes the delta against
|
||||
the last per-session ``tokens_saved`` it observed and feeds only the
|
||||
delta into the contribution counter.
|
||||
|
||||
PR-G2 remediation (C1): the RTK source is the SESSION-incremental
|
||||
``session.tokens_saved`` field of the helper payload, NOT the raw
|
||||
``lifetime_tokens_saved`` counter. The helper de-baselines per
|
||||
proxy session, so the first poll after process startup correctly
|
||||
reads 0 instead of the entire pre-Headroom RTK history.
|
||||
|
||||
Args:
|
||||
tokens_saved_cli_filtering: Explicit per-call CLI-filtering
|
||||
contribution. PR-G2 remediation (M5): ``None`` means "caller
|
||||
omitted, default 0" (mirrors the ``tokens_saved_rtk is
|
||||
None`` semantic). An explicit ``0`` is honored verbatim.
|
||||
tokens_saved_rtk: Explicit override for tokens saved by RTK on
|
||||
this call. If ``None`` (the default), the tracker polls RTK
|
||||
stats itself and writes the per-call delta. If passed
|
||||
(including ``0``), the override is used verbatim.
|
||||
"""
|
||||
# Polled outside the lock so the subprocess call can't deadlock the
|
||||
# event loop or contend with concurrent ``notify_active`` callers.
|
||||
if tokens_saved_rtk is None:
|
||||
tokens_saved_rtk = self._poll_rtk_delta()
|
||||
|
||||
with self._lock:
|
||||
c = self._state.contribution
|
||||
# PR-G2 remediation (M5): explicit None-guard so callers can
|
||||
# pass ``0`` and have it honored without colliding with the
|
||||
# default "caller omitted" semantic.
|
||||
if tokens_saved_cli_filtering is None:
|
||||
cli_filtering = 0
|
||||
else:
|
||||
cli_filtering = tokens_saved_cli_filtering
|
||||
c.tokens_submitted += max(tokens_submitted, 0)
|
||||
c.tokens_saved_compression += max(tokens_saved_compression, 0)
|
||||
c.tokens_saved_cli_filtering += max(cli_filtering, 0)
|
||||
c.tokens_saved_rtk += max(tokens_saved_rtk, 0)
|
||||
c.tokens_saved_cache_reads += max(tokens_saved_cache_reads, 0)
|
||||
c.compression_savings_usd += max(compression_savings_usd, 0.0)
|
||||
c.cache_savings_usd += max(cache_savings_usd, 0.0)
|
||||
|
||||
def _poll_rtk_delta(self) -> int:
|
||||
"""Return the delta of the session-scoped RTK ``tokens_saved`` since last poll.
|
||||
|
||||
Implementation of PR-G2 data-plane wiring. Calls
|
||||
:func:`headroom.proxy.helpers._get_rtk_stats` and reads the
|
||||
session-incremental ``session.tokens_saved`` field (de-baselined per
|
||||
proxy session by the helper), then diffs against
|
||||
``self._last_rtk_tokens_saved``.
|
||||
|
||||
Returns ``0`` (never negative) when:
|
||||
- ``HEADROOM_RTK_WIRING=disabled`` — operator opt-out.
|
||||
- ``_get_rtk_stats()`` returns ``None`` — RTK not selected, or the
|
||||
stat read failed this poll ("no data"); explicit zero contribution
|
||||
is the right answer and the high-water mark is preserved.
|
||||
- ``_get_rtk_stats()`` raises — transient error, logged loudly.
|
||||
- The session counter regressed (RTK reset / new project) — that
|
||||
path also re-baselines ``_last_rtk_tokens_saved`` to the new
|
||||
(smaller) value so subsequent polls return correct deltas.
|
||||
- PR-G2 remediation (C3): another worker holds the RTK poll lock —
|
||||
this worker skips polling so we don't double-count.
|
||||
|
||||
Otherwise advances ``self._last_rtk_tokens_saved`` to the new
|
||||
session total and returns the positive delta.
|
||||
|
||||
PR-G2 remediation (H1): ``HEADROOM_RTK_WIRING`` is now validated at
|
||||
startup via :func:`_validate_rtk_env_at_startup`; an invalid value
|
||||
raises at proxy boot rather than being silently swallowed here.
|
||||
We still catch + structured-log per-call as a defence-in-depth so
|
||||
that an env var flipped to garbage after startup is at least loud
|
||||
in the logs.
|
||||
"""
|
||||
try:
|
||||
wiring_mode = _rtk_wiring_mode()
|
||||
except ValueError as exc:
|
||||
# PR-G2 remediation (H1): elevate to ERROR — this is config
|
||||
# corruption, not a transient runtime hiccup. The bad value
|
||||
# should have been caught at startup but a rotation could flip
|
||||
# it mid-run; either way the operator must see this.
|
||||
logger.error(
|
||||
"event=subscription_rtk_invalid_env error=%s",
|
||||
exc,
|
||||
)
|
||||
return 0
|
||||
if wiring_mode == "disabled":
|
||||
return 0
|
||||
|
||||
# PR-G2 remediation (C3): only the lock-holder worker polls. We try
|
||||
# once per tracker instance and cache the verdict; the lock is
|
||||
# released on tracker stop.
|
||||
if not self._try_acquire_rtk_poll_lock():
|
||||
return 0
|
||||
|
||||
try:
|
||||
# Local import keeps the tracker module decoupled from the proxy
|
||||
# helper at import time (helpers.py imports many heavy deps).
|
||||
from headroom.proxy.helpers import _get_rtk_stats
|
||||
except Exception as exc: # pragma: no cover — defensive
|
||||
logger.warning(
|
||||
"event=subscription_rtk_helper_import_failed error=%s",
|
||||
exc,
|
||||
)
|
||||
return 0
|
||||
|
||||
try:
|
||||
stats = _get_rtk_stats()
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"event=subscription_rtk_stats_fetch_failed error=%s",
|
||||
exc,
|
||||
)
|
||||
return 0
|
||||
|
||||
if stats is None:
|
||||
logger.info(
|
||||
"event=subscription_rtk_stats_unavailable wiring=%s",
|
||||
wiring_mode,
|
||||
)
|
||||
return 0
|
||||
|
||||
# PR-G2 remediation (C1): read the SESSION-incremental field, not
|
||||
# the raw lifetime counter. The helper de-baselines per proxy
|
||||
# session at startup, so ``session.tokens_saved`` already excludes
|
||||
# the pre-Headroom RTK history. Falls back to the top-level
|
||||
# ``tokens_saved`` (which is also session-scoped in the canonical
|
||||
# payload; see ``_get_context_tool_stats``) and finally to 0 for
|
||||
# not-installed zero payloads (failed reads arrive as ``None`` and
|
||||
# returned above).
|
||||
session_payload = stats.get("session")
|
||||
if isinstance(session_payload, dict) and "tokens_saved" in session_payload:
|
||||
current_total_raw = session_payload.get("tokens_saved", 0)
|
||||
else:
|
||||
current_total_raw = stats.get("tokens_saved", 0)
|
||||
try:
|
||||
current_total = int(current_total_raw or 0)
|
||||
except (TypeError, ValueError) as exc:
|
||||
logger.warning(
|
||||
"event=subscription_rtk_stats_coerce_failed value=%r error=%s",
|
||||
current_total_raw,
|
||||
exc,
|
||||
)
|
||||
return 0
|
||||
|
||||
with self._lock:
|
||||
last = self._last_rtk_tokens_saved
|
||||
if current_total < last:
|
||||
# Counter regressed: helper rebaselined (session reset) or
|
||||
# RTK rebuilt its DB. Re-baseline silently — losing one
|
||||
# delta is preferable to reporting a giant negative number.
|
||||
logger.info(
|
||||
"event=subscription_rtk_counter_regressed previous=%d current=%d",
|
||||
last,
|
||||
current_total,
|
||||
)
|
||||
self._last_rtk_tokens_saved = current_total
|
||||
return 0
|
||||
delta = current_total - last
|
||||
if delta > 0:
|
||||
self._last_rtk_tokens_saved = current_total
|
||||
return delta
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Multi-worker poll ownership (PR-G2 remediation C3)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _try_acquire_rtk_poll_lock(self) -> bool:
|
||||
"""Try to acquire the RTK poll file lock (non-blocking).
|
||||
|
||||
Returns ``True`` if this worker owns the lock and should poll RTK.
|
||||
Caches the verdict so we don't pay the syscall on every call. The
|
||||
lock is released in :meth:`_release_rtk_poll_lock`, called from
|
||||
:meth:`stop`.
|
||||
|
||||
Mirrors the beacon's ``_try_acquire_beacon_lock`` pattern in
|
||||
``headroom/proxy/server.py`` (fcntl.flock, LOCK_EX | LOCK_NB).
|
||||
"""
|
||||
if self._rtk_poll_owner is not None:
|
||||
return self._rtk_poll_owner
|
||||
|
||||
try:
|
||||
import fcntl
|
||||
except ImportError:
|
||||
# Platform without fcntl (Windows). Every worker polls; log
|
||||
# loudly so the operator knows the multi-worker invariant is
|
||||
# weaker on this platform.
|
||||
logger.warning("event=subscription_rtk_poll_lock_unavailable platform=no-fcntl")
|
||||
self._rtk_poll_owner = True
|
||||
return True
|
||||
|
||||
lock_path = self._rtk_poll_lock_path()
|
||||
fd = None
|
||||
try:
|
||||
lock_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
fd = open(lock_path, "w") # noqa: SIM115
|
||||
fcntl_any = cast(Any, fcntl)
|
||||
fcntl_any.flock(fd, fcntl_any.LOCK_EX | fcntl_any.LOCK_NB)
|
||||
fd.write(str(os.getpid()))
|
||||
fd.flush()
|
||||
self._rtk_poll_lock_fd = fd
|
||||
self._rtk_poll_owner = True
|
||||
logger.info(
|
||||
"event=subscription_rtk_poll_lock_acquired pid=%d path=%s",
|
||||
os.getpid(),
|
||||
lock_path,
|
||||
)
|
||||
return True
|
||||
except OSError:
|
||||
if fd is not None:
|
||||
fd.close()
|
||||
self._rtk_poll_owner = False
|
||||
logger.info(
|
||||
"event=subscription_rtk_poll_lock_skipped pid=%d path=%s",
|
||||
os.getpid(),
|
||||
lock_path,
|
||||
)
|
||||
return False
|
||||
|
||||
def _release_rtk_poll_lock(self) -> None:
|
||||
"""Release the RTK poll file lock; safe to call repeatedly."""
|
||||
fd = self._rtk_poll_lock_fd
|
||||
if fd is None:
|
||||
return
|
||||
try:
|
||||
import fcntl
|
||||
|
||||
fcntl_any = cast(Any, fcntl)
|
||||
fcntl_any.flock(fd, fcntl_any.LOCK_UN)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
fd.close()
|
||||
except Exception:
|
||||
pass
|
||||
self._rtk_poll_lock_fd = None
|
||||
try:
|
||||
self._rtk_poll_lock_path().unlink(missing_ok=True)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _rtk_poll_lock_path(self) -> Path:
|
||||
"""Return the path to the RTK poll lock file."""
|
||||
override = os.environ.get(_RTK_POLL_LOCK_ENV, "").strip()
|
||||
if override:
|
||||
return Path(override)
|
||||
return self._persist_path.parent / ".rtk_poll_lock"
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# State access
|
||||
# ------------------------------------------------------------------
|
||||
|
|
@ -824,37 +492,10 @@ class SubscriptionTracker(QuotaTracker):
|
|||
c.tokens_saved_compression = int(
|
||||
saved.get("proxy_compression", saved.get("compression", 0))
|
||||
)
|
||||
# PR-G2 (Realignment) — prefer the raw counters when present
|
||||
# (new format). For backward compatibility with state written
|
||||
# before PR-G2 we fall back to the dashboard-aliased keys.
|
||||
cli_filtering = int(
|
||||
saved.get(
|
||||
"cli_filtering_raw",
|
||||
saved.get("cli_filtering", saved.get("rtk", 0)),
|
||||
)
|
||||
)
|
||||
c.tokens_saved_cli_filtering = cli_filtering
|
||||
# PR-G2 remediation (M2): legacy state files written before this
|
||||
# PR have no ``rtk_raw`` key — but pre-G2 the ``rtk`` field
|
||||
# silently mirrored ``cli_filtering``. Treat the legacy ``rtk``
|
||||
# field as authoritative for the rtk_raw counter to avoid
|
||||
# zeroing historical accumulation. New format writes ``rtk_raw``
|
||||
# explicitly; legacy writes had only ``rtk`` (aliased) and we
|
||||
# honor it on read.
|
||||
is_legacy_state = "rtk_raw" not in saved
|
||||
if is_legacy_state:
|
||||
# Pre-G2 semantic: ``rtk`` == ``cli_filtering``. Carry the
|
||||
# accumulated value forward instead of silently zeroing it.
|
||||
legacy_rtk_value = int(saved.get("rtk", cli_filtering))
|
||||
c.tokens_saved_rtk = legacy_rtk_value
|
||||
logger.info(
|
||||
"event=subscription_state_legacy_load "
|
||||
"migrated_rtk_raw_from_cli_filtering=%d path=%s",
|
||||
legacy_rtk_value,
|
||||
self._persist_path,
|
||||
)
|
||||
else:
|
||||
c.tokens_saved_rtk = int(saved.get("rtk_raw", 0))
|
||||
# State written before the CLI context tools were removed also carries
|
||||
# ``cli_filtering`` / ``cli_filtering_raw`` / ``rtk`` / ``rtk_raw``
|
||||
# counters. Those tools are gone, so the keys are simply ignored —
|
||||
# read must stay tolerant of them rather than raising on an old file.
|
||||
c.tokens_saved_cache_reads = int(saved.get("cache_reads", 0))
|
||||
savings_usd = contrib.get("savings_usd", {})
|
||||
c.compression_savings_usd = float(savings_usd.get("compression", 0.0))
|
||||
|
|
@ -964,14 +605,7 @@ def configure_subscription_tracker(
|
|||
persist_path: Path | None = None,
|
||||
client: SubscriptionClient | None = None,
|
||||
) -> SubscriptionTracker:
|
||||
"""Create (or return existing) global tracker singleton.
|
||||
|
||||
PR-G2 remediation (H1): validates RTK-related env vars eagerly here so
|
||||
a typo (``HEADROOM_RTK_WIRING=enabld``) crashes the proxy at startup
|
||||
instead of being silently swallowed at every ``update_contribution``
|
||||
call.
|
||||
"""
|
||||
_validate_rtk_env_at_startup()
|
||||
"""Create (or return existing) global tracker singleton."""
|
||||
global _tracker_instance
|
||||
with _tracker_lock:
|
||||
if _tracker_instance is None:
|
||||
|
|
|
|||
|
|
@ -87,11 +87,6 @@ function Require-Command {
|
|||
}
|
||||
}
|
||||
|
||||
function Get-RtkTarget {
|
||||
$arch = if ($env:PROCESSOR_ARCHITECTURE -match 'ARM64') { 'aarch64' } else { 'x86_64' }
|
||||
return "${arch}-pc-windows-msvc"
|
||||
}
|
||||
|
||||
function Ensure-HostDirs {
|
||||
foreach ($dir in @(
|
||||
(Join-Path $HostHome '.headroom'),
|
||||
|
|
@ -825,52 +820,6 @@ function Parse-InstallProfileArgs {
|
|||
return $profile
|
||||
}
|
||||
|
||||
function Invoke-ClaudeRtkInit {
|
||||
$rtkPath = Join-Path $HostHome '.headroom\bin\rtk.exe'
|
||||
if (-not (Test-Path $rtkPath)) {
|
||||
Write-Warning "rtk was not installed at $rtkPath; Claude hooks were not registered"
|
||||
return
|
||||
}
|
||||
|
||||
try {
|
||||
& $rtkPath init --global --auto-patch | Out-Null
|
||||
} catch {
|
||||
Write-Warning "Failed to register Claude hooks with rtk; continuing without hook registration"
|
||||
}
|
||||
}
|
||||
|
||||
function Get-ContextTool {
|
||||
$value = $env:HEADROOM_CONTEXT_TOOL
|
||||
if ([string]::IsNullOrWhiteSpace($value)) {
|
||||
return 'rtk'
|
||||
}
|
||||
|
||||
$value = $value.Trim().ToLowerInvariant().Replace('_', '-')
|
||||
if ($value -eq 'leanctx') {
|
||||
return 'lean-ctx'
|
||||
}
|
||||
if ($value -ne 'rtk' -and $value -ne 'lean-ctx') {
|
||||
Fail 'HEADROOM_CONTEXT_TOOL must be one of: lean-ctx, rtk'
|
||||
}
|
||||
return $value
|
||||
}
|
||||
|
||||
function Invoke-LeanCtxInit {
|
||||
param([string]$Agent)
|
||||
|
||||
$cmd = Get-Command lean-ctx -ErrorAction SilentlyContinue
|
||||
if (-not $cmd) {
|
||||
Write-Warning "lean-ctx is not installed on PATH; $Agent lean-ctx setup was skipped"
|
||||
return
|
||||
}
|
||||
|
||||
try {
|
||||
& $cmd.Source init --agent $Agent | Out-Null
|
||||
} catch {
|
||||
Write-Warning "Failed to initialize lean-ctx for $Agent; continuing without lean-ctx setup"
|
||||
}
|
||||
}
|
||||
|
||||
function Invoke-WithTemporaryEnv {
|
||||
param(
|
||||
[hashtable]$Environment,
|
||||
|
|
@ -1399,7 +1348,6 @@ function Parse-WrapArgs {
|
|||
$known = New-Object System.Collections.Generic.List[string]
|
||||
$hostArgs = New-Object System.Collections.Generic.List[string]
|
||||
$port = 8787
|
||||
$noRtk = $false
|
||||
$noProxy = $false
|
||||
$learn = $false
|
||||
$backend = $null
|
||||
|
|
@ -1431,12 +1379,6 @@ function Parse-WrapArgs {
|
|||
$i += 1
|
||||
continue
|
||||
}
|
||||
'^--no-rtk$' {
|
||||
$noRtk = $true
|
||||
$known.Add($arg)
|
||||
$i += 1
|
||||
continue
|
||||
}
|
||||
'^--no-proxy$' {
|
||||
$noProxy = $true
|
||||
$known.Add($arg)
|
||||
|
|
@ -1496,6 +1438,13 @@ function Parse-WrapArgs {
|
|||
$i += 1
|
||||
continue
|
||||
}
|
||||
'^--rtk$|^--no-rtk$|^--no-project-rtk$|^--keep-rtk$|^--context-tool$|^--context-tool=|^--no-context-tool$' {
|
||||
# Retired CLI context tools (rtk, lean-ctx). Reject explicitly: the
|
||||
# default branch below forwards the first unknown flag AND everything
|
||||
# after it to the wrapped tool, so a leftover --no-rtk in a script
|
||||
# would silently swallow a following --port and be ignored downstream.
|
||||
Fail "CLI context tools (rtk, lean-ctx) have been removed from Headroom. Drop $arg and unset HEADROOM_CONTEXT_TOOL; 'headroom wrap' uninstalls what they left behind on first run."
|
||||
}
|
||||
default {
|
||||
for ($j = $i; $j -lt $Arguments.Count; $j++) {
|
||||
$hostArgs.Add($Arguments[$j])
|
||||
|
|
@ -1509,7 +1458,6 @@ function Parse-WrapArgs {
|
|||
KnownArgs = $known.ToArray()
|
||||
HostArgs = $hostArgs.ToArray()
|
||||
Port = $port
|
||||
NoRtk = $noRtk
|
||||
NoProxy = $noProxy
|
||||
Learn = $learn
|
||||
Backend = $backend
|
||||
|
|
@ -1528,8 +1476,6 @@ function Invoke-PrepareOnly {
|
|||
$dockerArgs.AddRange([string[]]@('run','--rm'))
|
||||
Add-TtyArgs -ArgsList $dockerArgs
|
||||
$dockerArgs.AddRange((Get-SharedDockerArgs))
|
||||
$dockerArgs.Add('--env')
|
||||
$dockerArgs.Add("HEADROOM_RTK_TARGET=$(Get-RtkTarget)")
|
||||
$dockerArgs.Add('--entrypoint')
|
||||
$dockerArgs.Add('headroom')
|
||||
$dockerArgs.Add($HeadroomImage)
|
||||
|
|
@ -1647,7 +1593,6 @@ switch ($args[0]) {
|
|||
}
|
||||
|
||||
$parsed = Parse-WrapArgs -Arguments $wrapArgs
|
||||
$contextTool = Get-ContextTool
|
||||
$proxyArgs = New-Object System.Collections.Generic.List[string]
|
||||
if ($parsed.Learn) { $proxyArgs.Add('--learn') }
|
||||
if ($parsed.Backend) { $proxyArgs.AddRange([string[]]@('--backend', $parsed.Backend)) }
|
||||
|
|
@ -1667,20 +1612,10 @@ switch ($args[0]) {
|
|||
if (-not $parsed.NoProxy) {
|
||||
$prepareArgs.Add('--no-proxy')
|
||||
}
|
||||
if ((-not $parsed.NoRtk) -and $contextTool -eq 'lean-ctx') {
|
||||
$prepareArgs.Add('--no-rtk')
|
||||
}
|
||||
Invoke-PrepareOnly -Tool $tool -KnownArgs $prepareArgs.ToArray()
|
||||
|
||||
if ((-not $parsed.NoRtk) -and $contextTool -eq 'lean-ctx') {
|
||||
Invoke-LeanCtxInit -Agent $tool
|
||||
}
|
||||
|
||||
switch ($tool) {
|
||||
'claude' {
|
||||
if ((-not $parsed.NoRtk) -and $contextTool -eq 'rtk') {
|
||||
Invoke-ClaudeRtkInit
|
||||
}
|
||||
$exitCode = Invoke-WithTemporaryEnv -Environment @{ ANTHROPIC_BASE_URL = "http://127.0.0.1:$($parsed.Port)" } -Command 'claude' -Arguments $parsed.HostArgs
|
||||
exit $exitCode
|
||||
}
|
||||
|
|
|
|||
|
|
@ -82,33 +82,6 @@ require_cmd() {
|
|||
command -v "$1" >/dev/null 2>&1 || die "Missing required command: $1"
|
||||
}
|
||||
|
||||
detect_rtk_target() {
|
||||
local system
|
||||
local machine
|
||||
system="$(uname -s)"
|
||||
machine="$(uname -m)"
|
||||
|
||||
case "${system}" in
|
||||
Darwin)
|
||||
if [[ "${machine}" == "arm64" ]]; then
|
||||
printf 'aarch64-apple-darwin'
|
||||
else
|
||||
printf 'x86_64-apple-darwin'
|
||||
fi
|
||||
;;
|
||||
Linux)
|
||||
if [[ "${machine}" == "aarch64" ]]; then
|
||||
printf 'aarch64-unknown-linux-gnu'
|
||||
else
|
||||
printf 'x86_64-unknown-linux-musl'
|
||||
fi
|
||||
;;
|
||||
*)
|
||||
die "Unsupported host platform for Docker-native wrapper: ${system}/${machine}"
|
||||
;;
|
||||
esac
|
||||
}
|
||||
|
||||
ensure_host_dirs() {
|
||||
mkdir -p \
|
||||
"${HEADROOM_HOST_HOME}/.headroom" \
|
||||
|
|
@ -785,66 +758,20 @@ parse_install_profile_arg() {
|
|||
done
|
||||
}
|
||||
|
||||
run_claude_rtk_init() {
|
||||
local rtk_bin="${HEADROOM_HOST_HOME}/.headroom/bin/rtk"
|
||||
if [[ ! -x "${rtk_bin}" ]]; then
|
||||
warn "rtk was not installed at ${rtk_bin}; Claude hooks were not registered"
|
||||
return
|
||||
fi
|
||||
|
||||
if ! "${rtk_bin}" init --global --auto-patch >/dev/null 2>&1; then
|
||||
warn "Failed to register Claude hooks with rtk; continuing without hook registration"
|
||||
fi
|
||||
}
|
||||
|
||||
selected_context_tool() {
|
||||
local value="${HEADROOM_CONTEXT_TOOL:-rtk}"
|
||||
value="${value,,}"
|
||||
value="${value//_/-}"
|
||||
if [[ -z "${value}" ]]; then
|
||||
value="rtk"
|
||||
elif [[ "${value}" == "leanctx" ]]; then
|
||||
value="lean-ctx"
|
||||
fi
|
||||
|
||||
case "${value}" in
|
||||
rtk|lean-ctx)
|
||||
printf '%s\n' "${value}"
|
||||
;;
|
||||
*)
|
||||
die "HEADROOM_CONTEXT_TOOL must be one of: lean-ctx, rtk"
|
||||
;;
|
||||
esac
|
||||
}
|
||||
|
||||
run_lean_ctx_init() {
|
||||
local agent="$1"
|
||||
if ! command -v lean-ctx >/dev/null 2>&1; then
|
||||
warn "lean-ctx is not installed on PATH; ${agent} lean-ctx setup was skipped"
|
||||
return
|
||||
fi
|
||||
|
||||
if ! lean-ctx init --agent "${agent}" >/dev/null 2>&1; then
|
||||
warn "Failed to initialize lean-ctx for ${agent}; continuing without lean-ctx setup"
|
||||
fi
|
||||
}
|
||||
|
||||
parse_wrap_args() {
|
||||
local -n out_known=$1
|
||||
local -n out_host=$2
|
||||
local -n out_port=$3
|
||||
local -n out_no_rtk=$4
|
||||
local -n out_no_proxy=$5
|
||||
local -n out_learn=$6
|
||||
local -n out_backend=$7
|
||||
local -n out_anyllm=$8
|
||||
local -n out_region=$9
|
||||
shift 9
|
||||
local -n out_no_proxy=$4
|
||||
local -n out_learn=$5
|
||||
local -n out_backend=$6
|
||||
local -n out_anyllm=$7
|
||||
local -n out_region=$8
|
||||
shift 8
|
||||
|
||||
out_known=()
|
||||
out_host=()
|
||||
out_port=8787
|
||||
out_no_rtk=0
|
||||
out_no_proxy=0
|
||||
out_learn=0
|
||||
out_backend=""
|
||||
|
|
@ -871,11 +798,6 @@ parse_wrap_args() {
|
|||
out_known+=("$1")
|
||||
shift
|
||||
;;
|
||||
--no-rtk)
|
||||
out_no_rtk=1
|
||||
out_known+=("$1")
|
||||
shift
|
||||
;;
|
||||
--no-proxy)
|
||||
out_no_proxy=1
|
||||
out_known+=("$1")
|
||||
|
|
@ -923,6 +845,13 @@ parse_wrap_args() {
|
|||
out_known+=("$1")
|
||||
shift
|
||||
;;
|
||||
--rtk|--no-rtk|--no-project-rtk|--keep-rtk|--context-tool|--no-context-tool|--context-tool=*)
|
||||
# Retired CLI context tools (rtk, lean-ctx). Reject explicitly: the
|
||||
# catch-all below forwards the first unknown flag AND everything after it
|
||||
# to the wrapped tool, so a leftover --no-rtk in a script would silently
|
||||
# swallow a following --port and be ignored by the wrapped CLI.
|
||||
die "CLI context tools (rtk, lean-ctx) have been removed from Headroom. Drop $1 and unset HEADROOM_CONTEXT_TOOL; 'headroom wrap' uninstalls what they left behind on first run."
|
||||
;;
|
||||
*)
|
||||
out_host+=("$@")
|
||||
break
|
||||
|
|
@ -939,7 +868,6 @@ run_prepare_only() {
|
|||
args=(docker run --rm)
|
||||
append_tty_args args
|
||||
append_common_container_args args
|
||||
args+=(--env "HEADROOM_RTK_TARGET=$(detect_rtk_target)")
|
||||
args+=(--entrypoint headroom "${HEADROOM_IMAGE}" wrap "${tool}" --prepare-only "$@")
|
||||
"${args[@]}"
|
||||
}
|
||||
|
|
@ -1502,9 +1430,8 @@ main() {
|
|||
return
|
||||
fi
|
||||
|
||||
local known_args host_args port no_rtk no_proxy learn backend anyllm region context_tool
|
||||
parse_wrap_args known_args host_args port no_rtk no_proxy learn backend anyllm region "$@"
|
||||
context_tool="$(selected_context_tool)"
|
||||
local known_args host_args port no_proxy learn backend anyllm region
|
||||
parse_wrap_args known_args host_args port no_proxy learn backend anyllm region "$@"
|
||||
|
||||
local proxy_args=()
|
||||
if [[ "${learn}" -eq 1 ]]; then
|
||||
|
|
@ -1530,20 +1457,10 @@ main() {
|
|||
if [[ "${no_proxy}" -eq 0 ]]; then
|
||||
prep_args+=(--no-proxy)
|
||||
fi
|
||||
if [[ "${no_rtk}" -eq 0 && "${context_tool}" == "lean-ctx" ]]; then
|
||||
prep_args+=(--no-rtk)
|
||||
fi
|
||||
run_prepare_only "${tool}" "${prep_args[@]}"
|
||||
|
||||
if [[ "${no_rtk}" -eq 0 && "${context_tool}" == "lean-ctx" ]]; then
|
||||
run_lean_ctx_init "${tool}"
|
||||
fi
|
||||
|
||||
case "${tool}" in
|
||||
claude)
|
||||
if [[ "${no_rtk}" -eq 0 && "${context_tool}" == "rtk" ]]; then
|
||||
run_claude_rtk_init
|
||||
fi
|
||||
ANTHROPIC_BASE_URL="http://127.0.0.1:${port}" run_host_tool claude "${host_args[@]}"
|
||||
;;
|
||||
codex)
|
||||
|
|
|
|||
|
|
@ -131,7 +131,6 @@ export {
|
|||
proxyLogPath,
|
||||
debug400Dir,
|
||||
binDir,
|
||||
rtkPath,
|
||||
deployRoot,
|
||||
beaconLockPath,
|
||||
modelsConfigPath,
|
||||
|
|
|
|||
|
|
@ -224,12 +224,6 @@ export function binDir(): string {
|
|||
return joinPath(workspaceDir(), "bin");
|
||||
}
|
||||
|
||||
export function rtkPath(): string {
|
||||
if (!isNode()) return "";
|
||||
const name = process.platform === "win32" ? "rtk.exe" : "rtk";
|
||||
return joinPath(binDir(), name);
|
||||
}
|
||||
|
||||
export function deployRoot(): string {
|
||||
if (!isNode()) return "";
|
||||
return joinPath(workspaceDir(), "deploy");
|
||||
|
|
|
|||
|
|
@ -27,7 +27,6 @@ import {
|
|||
pluginConfigDir,
|
||||
pluginWorkspaceDir,
|
||||
proxyLogPath,
|
||||
rtkPath,
|
||||
savingsPath,
|
||||
sessionStatsPath,
|
||||
subscriptionStatePath,
|
||||
|
|
@ -257,12 +256,6 @@ describe("derived-only resources", () => {
|
|||
expect(binDir()).toBe(path.join(os.homedir(), ".headroom", "bin"));
|
||||
});
|
||||
|
||||
it("rtkPath ends with rtk or rtk.exe", () => {
|
||||
const p = rtkPath();
|
||||
const expected = process.platform === "win32" ? "rtk.exe" : "rtk";
|
||||
expect(path.basename(p)).toBe(expected);
|
||||
});
|
||||
|
||||
it("deployRoot", () => {
|
||||
expect(deployRoot()).toBe(path.join(os.homedir(), ".headroom", "deploy"));
|
||||
});
|
||||
|
|
@ -347,10 +340,6 @@ describe("derived-only helpers follow workspace env", () => {
|
|||
it("binDir", () => {
|
||||
expect(binDir()).toBe(path.join("/tmp/alt_ws", "bin"));
|
||||
});
|
||||
it("rtkPath", () => {
|
||||
const expected = process.platform === "win32" ? "rtk.exe" : "rtk";
|
||||
expect(rtkPath()).toBe(path.join("/tmp/alt_ws", "bin", expected));
|
||||
});
|
||||
it("deployRoot", () => {
|
||||
expect(deployRoot()).toBe(path.join("/tmp/alt_ws", "deploy"));
|
||||
});
|
||||
|
|
|
|||
|
|
@ -6,25 +6,20 @@ import json
|
|||
|
||||
import pytest
|
||||
|
||||
from headroom.audit.codex import audit_codex, classify_command, render_codex_text, strip_wrappers
|
||||
from headroom.audit.codex import audit_codex, classify_command, render_codex_text
|
||||
|
||||
|
||||
class TestClassifier:
|
||||
def test_strip_wrappers(self):
|
||||
assert strip_wrappers("rtk cat foo.py") == "cat foo.py"
|
||||
assert strip_wrappers("rtk proxy sed -n '1,20p' foo.py") == "sed -n '1,20p' foo.py"
|
||||
assert strip_wrappers("git status") == "git status"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("cmd", "category", "partial"),
|
||||
[
|
||||
("cat src/foo.py", "read", False),
|
||||
("sed -n '1,200p' src/foo.py", "read", True),
|
||||
("rtk read src/foo.py --lines 10-50", "read", True),
|
||||
("bat src/foo.py --lines 10-50", "read", True),
|
||||
("head -50 src/foo.py", "read", True),
|
||||
("nl headroom/config.py", "read", False),
|
||||
("rg -n 'def apply' headroom/", "search", False),
|
||||
("rtk grep -n pattern .", "search", False),
|
||||
("grep -rn pattern .", "search", False),
|
||||
("git diff HEAD~1", "git", False),
|
||||
("apply_patch <<'EOF'\n*** Begin Patch\nEOF", "edit", False),
|
||||
("pytest tests/ -x -q", "build/test", False),
|
||||
|
|
@ -100,7 +95,7 @@ def codex_dir(tmp_path):
|
|||
_output("c2", content[:500]),
|
||||
_call("c3", "rg -n 'def ' src/"),
|
||||
_output("c3", "src/foo.py:1:def x():"),
|
||||
_call("c4", "rtk read src/bar.py --lines 1-50"),
|
||||
_call("c4", "bat src/bar.py --lines 1-50"),
|
||||
_output("c4", "bar content " * 10),
|
||||
]
|
||||
sessions = tmp_path / "sessions" / "2026" / "06"
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
"""Shared fixtures for the CLI test suite.
|
||||
|
||||
``headroom wrap`` may fetch helper binaries (e.g. rtk, lean-ctx) over the
|
||||
network via ``headroom.binaries``. Force offline across CLI tests so a missing
|
||||
binary resolves locally instead of reaching out to GitHub releases. Tests that
|
||||
``headroom wrap`` may fetch helper binaries (e.g. tokensave) over the network
|
||||
via ``headroom.binaries``. Force offline across CLI tests so a missing binary
|
||||
resolves locally instead of reaching out to GitHub releases. Tests that
|
||||
exercise a binary-present path patch the relevant resolver directly and are
|
||||
unaffected by this guard.
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -53,9 +53,8 @@ def test_subcommand_verbose_flag_still_works() -> None:
|
|||
|
||||
with patch("headroom.cli.wrap.shutil.which", return_value="claude"):
|
||||
with patch("headroom.cli.wrap._ensure_proxy", return_value=(None, 8787)):
|
||||
with patch("headroom.cli.wrap._setup_rtk", return_value=None):
|
||||
with patch("headroom.cli.wrap.subprocess.run", return_value=completed):
|
||||
result = runner.invoke(main, ["wrap", "claude", "-v"])
|
||||
with patch("headroom.cli.wrap.subprocess.run", return_value=completed):
|
||||
result = runner.invoke(main, ["wrap", "claude", "-v"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "HEADROOM WRAP: CLAUDE" in result.output
|
||||
|
|
|
|||
|
|
@ -21,7 +21,7 @@ def _no_persistent_manifest(monkeypatch: pytest.MonkeyPatch) -> None:
|
|||
monkeypatch.setattr(wrap_cli, "_find_persistent_manifest", lambda _port: None)
|
||||
|
||||
|
||||
def test_remove_claude_rtk_hooks_preserves_unrelated_hooks(tmp_path: Path) -> None:
|
||||
def test_remove_claude_managed_hooks_preserves_unrelated_hooks(tmp_path: Path) -> None:
|
||||
settings = tmp_path / "settings.json"
|
||||
settings.write_text(
|
||||
json.dumps(
|
||||
|
|
@ -34,7 +34,9 @@ def test_remove_claude_rtk_hooks_preserves_unrelated_hooks(tmp_path: Path) -> No
|
|||
"hooks": [
|
||||
{
|
||||
"type": "command",
|
||||
"command": "/Users/test/.claude/hooks/rtk-rewrite.sh",
|
||||
"command": (
|
||||
"headroom init hook ensure --marker headroom-init-claude"
|
||||
),
|
||||
},
|
||||
{"type": "command", "command": "echo keep"},
|
||||
],
|
||||
|
|
@ -50,7 +52,7 @@ def test_remove_claude_rtk_hooks_preserves_unrelated_hooks(tmp_path: Path) -> No
|
|||
encoding="utf-8",
|
||||
)
|
||||
|
||||
assert wrap_cli._remove_claude_rtk_hooks(settings) is True
|
||||
assert wrap_cli._remove_claude_managed_hooks(settings) is True
|
||||
|
||||
payload = json.loads(settings.read_text(encoding="utf-8"))
|
||||
pre_tool_hooks = payload["hooks"]["PreToolUse"][0]["hooks"]
|
||||
|
|
@ -58,7 +60,7 @@ def test_remove_claude_rtk_hooks_preserves_unrelated_hooks(tmp_path: Path) -> No
|
|||
assert payload["hooks"]["SessionStart"][0]["hooks"][0]["command"] == "keep"
|
||||
|
||||
|
||||
def test_unwrap_claude_removes_mcp_rtk_and_stops_proxy(
|
||||
def test_unwrap_claude_removes_mcp_purges_retired_hook_and_stops_proxy(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
|
|
@ -116,6 +118,8 @@ def test_unwrap_claude_removes_mcp_rtk_and_stops_proxy(
|
|||
assert unregistered == ["headroom", "codebase-memory-mcp"]
|
||||
assert stopped == [9999]
|
||||
assert "Stopped local Headroom proxy on port 9999" in result.output
|
||||
# The leftover retired context-tool hook is purged end-to-end by unwrap
|
||||
# (via purge_context_tool_artifacts), leaving no hooks behind.
|
||||
assert "hooks" not in json.loads(settings.read_text(encoding="utf-8"))
|
||||
|
||||
|
||||
|
|
@ -146,7 +150,7 @@ def test_unwrap_claude_preserves_user_managed_serena(
|
|||
|
||||
with (
|
||||
patch("headroom.mcp_registry.ClaudeRegistrar", return_value=Registrar()),
|
||||
patch("headroom.cli.wrap._remove_claude_rtk_hooks", return_value=False),
|
||||
patch("headroom.cli.wrap._remove_claude_managed_hooks", return_value=False),
|
||||
patch("headroom.cli.wrap._stop_local_proxy_for_unwrap"),
|
||||
):
|
||||
result = runner.invoke(main, ["unwrap", "claude"])
|
||||
|
|
@ -186,7 +190,7 @@ def test_unwrap_claude_removes_headroom_installed_serena(
|
|||
|
||||
with (
|
||||
patch("headroom.mcp_registry.ClaudeRegistrar", return_value=Registrar()),
|
||||
patch("headroom.cli.wrap._remove_claude_rtk_hooks", return_value=False),
|
||||
patch("headroom.cli.wrap._remove_claude_managed_hooks", return_value=False),
|
||||
patch("headroom.cli.wrap._stop_local_proxy_for_unwrap"),
|
||||
):
|
||||
result = runner.invoke(main, ["unwrap", "claude"])
|
||||
|
|
@ -201,17 +205,16 @@ def test_unwrap_claude_keep_flags_skip_cleanup(
|
|||
) -> None:
|
||||
with (
|
||||
patch("headroom.mcp_registry.ClaudeRegistrar") as registrar,
|
||||
patch("headroom.cli.wrap._remove_claude_rtk_hooks") as remove_rtk,
|
||||
patch("headroom.cli.wrap._remove_claude_managed_hooks", return_value=False),
|
||||
patch("headroom.cli.wrap._stop_local_proxy_for_unwrap") as stop_proxy,
|
||||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["unwrap", "claude", "--keep-mcp", "--keep-rtk", "--no-stop-proxy"],
|
||||
["unwrap", "claude", "--keep-mcp", "--no-stop-proxy"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
registrar.assert_not_called()
|
||||
remove_rtk.assert_not_called()
|
||||
stop_proxy.assert_not_called()
|
||||
|
||||
|
||||
|
|
@ -224,7 +227,7 @@ def test_unwrap_claude_restores_all_base_url_modes(runner: CliRunner) -> None:
|
|||
with patch("headroom.cli.wrap._restore_claude_wrap_base_url", side_effect=restore_base_url):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["unwrap", "claude", "--keep-mcp", "--keep-rtk", "--no-stop-proxy"],
|
||||
["unwrap", "claude", "--keep-mcp", "--no-stop-proxy"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
|
|
@ -280,7 +283,7 @@ def test_unwrap_claude_stops_claude_owned_persistent_deployment(
|
|||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["unwrap", "claude", "--keep-mcp", "--keep-rtk", "--port", "8787"],
|
||||
["unwrap", "claude", "--keep-mcp", "--port", "8787"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
|
|
@ -307,7 +310,7 @@ def test_unwrap_claude_reports_ambiguous_same_port_persistent_deployment(
|
|||
with patch("headroom.cli.wrap._stop_local_proxy_for_unwrap") as stop_local:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["unwrap", "claude", "--keep-mcp", "--keep-rtk", "--port", "8787"],
|
||||
["unwrap", "claude", "--keep-mcp", "--port", "8787"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
|
|
@ -326,7 +329,7 @@ def test_unwrap_claude_warns_about_same_port_inherited_env(
|
|||
with patch("headroom.cli.wrap._stop_local_proxy_for_unwrap", return_value="stopped"):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["unwrap", "claude", "--keep-mcp", "--keep-rtk", "--port", "8787"],
|
||||
["unwrap", "claude", "--keep-mcp", "--port", "8787"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
|
|
@ -343,7 +346,7 @@ def test_unwrap_claude_ignores_malformed_inherited_env_port(
|
|||
with patch("headroom.cli.wrap._stop_local_proxy_for_unwrap", return_value="stopped"):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["unwrap", "claude", "--keep-mcp", "--keep-rtk", "--port", "8787"],
|
||||
["unwrap", "claude", "--keep-mcp", "--port", "8787"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
|
|
@ -351,7 +354,7 @@ def test_unwrap_claude_ignores_malformed_inherited_env_port(
|
|||
assert "Claude is no longer durably wrapped by Headroom." in result.output
|
||||
|
||||
|
||||
def test_remove_claude_rtk_hooks_removes_init_hooks_and_env(tmp_path: Path) -> None:
|
||||
def test_remove_claude_managed_hooks_removes_init_hooks_and_env(tmp_path: Path) -> None:
|
||||
settings = tmp_path / "settings.json"
|
||||
settings.write_text(
|
||||
json.dumps(
|
||||
|
|
@ -393,7 +396,7 @@ def test_remove_claude_rtk_hooks_removes_init_hooks_and_env(tmp_path: Path) -> N
|
|||
encoding="utf-8",
|
||||
)
|
||||
|
||||
assert wrap_cli._remove_claude_rtk_hooks(settings) is True
|
||||
assert wrap_cli._remove_claude_managed_hooks(settings) is True
|
||||
|
||||
payload = json.loads(settings.read_text(encoding="utf-8"))
|
||||
# ANTHROPIC_BASE_URL stripped; unrelated env var preserved
|
||||
|
|
@ -407,7 +410,7 @@ def test_remove_claude_rtk_hooks_removes_init_hooks_and_env(tmp_path: Path) -> N
|
|||
assert payload["model"] == "opus"
|
||||
|
||||
|
||||
def test_remove_claude_rtk_hooks_strips_env_without_hooks(tmp_path: Path) -> None:
|
||||
def test_remove_claude_managed_hooks_strips_env_without_hooks(tmp_path: Path) -> None:
|
||||
# Regression: unwrap previously returned early when no hooks existed,
|
||||
# leaving init's ANTHROPIC_BASE_URL behind in settings.json.
|
||||
settings = tmp_path / "settings.json"
|
||||
|
|
@ -416,13 +419,13 @@ def test_remove_claude_rtk_hooks_strips_env_without_hooks(tmp_path: Path) -> Non
|
|||
encoding="utf-8",
|
||||
)
|
||||
|
||||
assert wrap_cli._remove_claude_rtk_hooks(settings) is True
|
||||
assert wrap_cli._remove_claude_managed_hooks(settings) is True
|
||||
|
||||
payload = json.loads(settings.read_text(encoding="utf-8"))
|
||||
assert "env" not in payload # emptied env dict is dropped
|
||||
|
||||
|
||||
def test_remove_claude_rtk_hooks_noop_when_nothing_managed(tmp_path: Path) -> None:
|
||||
def test_remove_claude_managed_hooks_noop_when_nothing_managed(tmp_path: Path) -> None:
|
||||
settings = tmp_path / "settings.json"
|
||||
original = {
|
||||
"model": "opus",
|
||||
|
|
@ -435,12 +438,12 @@ def test_remove_claude_rtk_hooks_noop_when_nothing_managed(tmp_path: Path) -> No
|
|||
}
|
||||
settings.write_text(json.dumps(original) + "\n", encoding="utf-8")
|
||||
|
||||
assert wrap_cli._remove_claude_rtk_hooks(settings) is False
|
||||
assert wrap_cli._remove_claude_managed_hooks(settings) is False
|
||||
# nothing managed -> file untouched
|
||||
assert json.loads(settings.read_text(encoding="utf-8")) == original
|
||||
|
||||
|
||||
def test_remove_claude_rtk_hooks_strips_enable_tool_search(tmp_path: Path) -> None:
|
||||
def test_remove_claude_managed_hooks_strips_enable_tool_search(tmp_path: Path) -> None:
|
||||
# unwrap must remove BOTH env vars init writes (ANTHROPIC_BASE_URL +
|
||||
# ENABLE_TOOL_SEARCH, GH #746), leaving user-set vars intact.
|
||||
settings = tmp_path / "settings.json"
|
||||
|
|
@ -458,7 +461,7 @@ def test_remove_claude_rtk_hooks_strips_enable_tool_search(tmp_path: Path) -> No
|
|||
encoding="utf-8",
|
||||
)
|
||||
|
||||
assert wrap_cli._remove_claude_rtk_hooks(settings) is True
|
||||
assert wrap_cli._remove_claude_managed_hooks(settings) is True
|
||||
|
||||
payload = json.loads(settings.read_text(encoding="utf-8"))
|
||||
assert payload["env"] == {"KEEP": "1"}
|
||||
|
|
|
|||
|
|
@ -33,7 +33,7 @@ def test_wrap_aider_sets_provider_envs(
|
|||
|
||||
with patch("headroom.cli.wrap.shutil.which", return_value="aider"):
|
||||
with patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool):
|
||||
result = runner.invoke(main, ["wrap", "aider", "--no-rtk", "--", "--model", "gpt-4o"])
|
||||
result = runner.invoke(main, ["wrap", "aider", "--", "--model", "gpt-4o"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
env = captured["env"]
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
|
@ -11,14 +10,16 @@ import pytest
|
|||
from click.testing import CliRunner
|
||||
|
||||
from headroom.cli.main import main
|
||||
from headroom.cli.wrap import _setup_lean_ctx_agent
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _default_context_tool(monkeypatch) -> None:
|
||||
def _no_retired_context_tool_env(monkeypatch) -> None:
|
||||
"""Keep a developer's exported HEADROOM_CONTEXT_TOOL from failing every test.
|
||||
|
||||
The var is now rejected outright, so leaving it set in the ambient
|
||||
environment would abort each wrap invocation below.
|
||||
"""
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
monkeypatch.delenv("LEAN_CTX_AGENT", raising=False)
|
||||
monkeypatch.delenv("LEAN_CTX_DATA_DIR", raising=False)
|
||||
|
||||
|
||||
def _set_test_home(monkeypatch, tmp_path: Path) -> None:
|
||||
|
|
@ -30,61 +31,18 @@ def _set_test_home(monkeypatch, tmp_path: Path) -> None:
|
|||
def test_wrap_claude_prepare_only_skips_host_binary_lookup() -> None:
|
||||
runner = CliRunner()
|
||||
|
||||
with patch("headroom.cli.wrap._prepare_wrap_rtk") as prepare_rtk:
|
||||
with patch("headroom.cli.wrap.shutil.which") as which_mock:
|
||||
result = runner.invoke(main, ["wrap", "claude", "--prepare-only", "--context-tool"])
|
||||
with patch("headroom.cli.wrap.shutil.which") as which_mock:
|
||||
result = runner.invoke(main, ["wrap", "claude", "--prepare-only"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
prepare_rtk.assert_called_once()
|
||||
which_mock.assert_not_called()
|
||||
|
||||
|
||||
def test_wrap_claude_prepare_only_uses_lean_ctx_when_configured(monkeypatch) -> None:
|
||||
runner = CliRunner()
|
||||
monkeypatch.setenv("HEADROOM_CONTEXT_TOOL", "lean-ctx")
|
||||
|
||||
with patch("headroom.cli.wrap._prepare_wrap_rtk") as prepare_rtk:
|
||||
with patch(
|
||||
"headroom.cli.wrap._setup_lean_ctx_agent",
|
||||
return_value=Path("lean-ctx"),
|
||||
) as setup:
|
||||
result = runner.invoke(main, ["wrap", "claude", "--prepare-only", "--context-tool"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
prepare_rtk.assert_not_called()
|
||||
setup.assert_called_once_with("claude", verbose=False)
|
||||
|
||||
|
||||
def test_setup_lean_ctx_agent_runs_outside_project_root(monkeypatch, tmp_path: Path) -> None:
|
||||
project_root = tmp_path / "project"
|
||||
project_root.mkdir()
|
||||
(project_root / ".git").mkdir()
|
||||
lean_ctx = tmp_path / "lean-ctx"
|
||||
lean_ctx.write_text("#!/bin/sh\n", encoding="utf-8")
|
||||
calls: list[dict] = []
|
||||
|
||||
def fake_run(*args, **kwargs):
|
||||
calls.append({"args": args, "kwargs": kwargs})
|
||||
return subprocess.CompletedProcess(args[0], 0, stdout="", stderr="")
|
||||
|
||||
monkeypatch.chdir(project_root)
|
||||
monkeypatch.setattr("headroom.lean_ctx.get_lean_ctx_path", lambda: lean_ctx)
|
||||
monkeypatch.setattr("headroom.cli.wrap.subprocess.run", fake_run)
|
||||
|
||||
assert _setup_lean_ctx_agent("codex") == lean_ctx
|
||||
|
||||
assert calls
|
||||
cwd = Path(calls[0]["kwargs"]["cwd"])
|
||||
assert cwd != project_root
|
||||
assert project_root not in cwd.parents
|
||||
|
||||
|
||||
def test_wrap_codex_prepare_only_updates_config(monkeypatch, tmp_path: Path) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
runner = CliRunner()
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
||||
result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"])
|
||||
result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
config_file = tmp_path / ".codex" / "config.toml"
|
||||
|
|
@ -103,7 +61,7 @@ def test_wrap_grok_build_uses_actual_proxy_port(monkeypatch, tmp_path: Path) ->
|
|||
|
||||
monkeypatch.setattr("headroom.cli.wrap._run_proxy_only_watcher", fake_watcher)
|
||||
|
||||
result = runner.invoke(main, ["wrap", "grok-build", "--no-context-tool", "--port", "8787"])
|
||||
result = runner.invoke(main, ["wrap", "grok-build", "--port", "8787"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
config_file = tmp_path / ".grok" / "config.toml"
|
||||
|
|
@ -115,134 +73,41 @@ def test_wrap_grok_build_uses_actual_proxy_port(monkeypatch, tmp_path: Path) ->
|
|||
assert "http://127.0.0.1:8787/" not in result.output
|
||||
|
||||
|
||||
def test_wrap_codex_prepare_only_uses_lean_ctx_when_configured(monkeypatch, tmp_path: Path) -> None:
|
||||
def test_wrap_rejects_retired_context_tool_flag(monkeypatch, tmp_path: Path) -> None:
|
||||
"""A surviving --context-tool must fail loudly, not be silently ignored.
|
||||
|
||||
rtk / lean-ctx are gone, but the flag lives on in shell profiles, scripts and
|
||||
CI jobs. Accepting it as a no-op would look like Headroom had quietly stopped
|
||||
filtering; the user needs to be told the feature was removed.
|
||||
"""
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
runner = CliRunner()
|
||||
|
||||
with runner.isolated_filesystem(temp_dir=str(tmp_path)):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "codex", "--prepare-only", "--no-context-tool", "--no-mcp", "--no-serena"],
|
||||
)
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "have been removed from Headroom" in result.output
|
||||
|
||||
|
||||
def test_wrap_rejects_retired_context_tool_env(monkeypatch, tmp_path: Path) -> None:
|
||||
"""An exported HEADROOM_CONTEXT_TOOL fails too, with the same message.
|
||||
|
||||
The env var is the form most likely to be left behind in a shell rc, where
|
||||
it would otherwise never surface.
|
||||
"""
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
monkeypatch.setenv("HEADROOM_CONTEXT_TOOL", "lean-ctx")
|
||||
runner = CliRunner()
|
||||
|
||||
with runner.isolated_filesystem(temp_dir=str(tmp_path)):
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary") as ensure_rtk:
|
||||
with patch(
|
||||
"headroom.cli.wrap._setup_lean_ctx_agent",
|
||||
return_value=Path("lean-ctx"),
|
||||
) as setup:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "codex", "--prepare-only", "--no-mcp", "--no-serena"],
|
||||
)
|
||||
result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--no-mcp", "--no-serena"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
ensure_rtk.assert_not_called()
|
||||
setup.assert_called_once_with("codex", verbose=False)
|
||||
assert not Path("AGENTS.md").exists()
|
||||
|
||||
|
||||
def test_wrap_codex_prepare_only_accepts_no_context_tool_alias(monkeypatch, tmp_path: Path) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
monkeypatch.setenv("HEADROOM_CONTEXT_TOOL", "lean-ctx")
|
||||
runner = CliRunner()
|
||||
|
||||
with runner.isolated_filesystem(temp_dir=str(tmp_path)):
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary") as ensure_rtk:
|
||||
with patch("headroom.cli.wrap._setup_lean_ctx_agent") as setup:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"wrap",
|
||||
"codex",
|
||||
"--prepare-only",
|
||||
"--no-context-tool",
|
||||
"--no-mcp",
|
||||
"--no-serena",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
ensure_rtk.assert_not_called()
|
||||
setup.assert_not_called()
|
||||
|
||||
|
||||
def test_wrap_aider_prepare_only_injects_conventions(monkeypatch, tmp_path: Path) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
monkeypatch.setenv("HEADROOM_RTK", "1") # RTK is opt-in; exercise the RTK-on path
|
||||
runner = CliRunner()
|
||||
|
||||
with runner.isolated_filesystem(temp_dir=str(tmp_path)):
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=Path("rtk")):
|
||||
result = runner.invoke(main, ["wrap", "aider", "--prepare-only"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
conventions = Path("CONVENTIONS.md")
|
||||
assert conventions.exists()
|
||||
assert "headroom:rtk-instructions" in conventions.read_text(encoding="utf-8")
|
||||
|
||||
|
||||
def test_wrap_cursor_prepare_only_registers_native_hook(monkeypatch, tmp_path: Path) -> None:
|
||||
# GH #756: when rtk's own `--agent cursor` hook registers successfully,
|
||||
# headroom must not also inject RTK_INSTRUCTIONS_BLOCK into .cursorrules.
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
monkeypatch.setenv("HEADROOM_RTK", "1") # RTK is opt-in; exercise the RTK-on path
|
||||
runner = CliRunner()
|
||||
|
||||
# headroom trusts the on-disk hook, not rtk's exit code, so simulate rtk
|
||||
# actually writing ~/.cursor/hooks.json when registration succeeds.
|
||||
def _register(_rtk_path, *, agent):
|
||||
hooks = tmp_path / ".cursor" / "hooks.json"
|
||||
hooks.parent.mkdir(parents=True, exist_ok=True)
|
||||
hooks.write_text('{"hooks": {"preToolUse": [{"command": "rtk hook cursor"}]}}')
|
||||
return True
|
||||
|
||||
with runner.isolated_filesystem(temp_dir=str(tmp_path)):
|
||||
with (
|
||||
patch("headroom.cli.wrap._ensure_rtk_binary", return_value=Path("rtk")),
|
||||
patch("headroom.rtk.installer.register_agent_hooks", side_effect=_register) as register,
|
||||
):
|
||||
result = runner.invoke(main, ["wrap", "cursor", "--prepare-only"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
register.assert_called_once_with(Path("rtk"), agent="cursor")
|
||||
assert not Path(".cursorrules").exists()
|
||||
|
||||
|
||||
def test_wrap_cursor_prepare_only_falls_back_to_cursorrules_when_hook_fails(
|
||||
monkeypatch, tmp_path: Path
|
||||
) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
monkeypatch.setenv("HEADROOM_RTK", "1") # RTK is opt-in; exercise the RTK-on path
|
||||
runner = CliRunner()
|
||||
|
||||
with runner.isolated_filesystem(temp_dir=str(tmp_path)):
|
||||
with (
|
||||
patch("headroom.cli.wrap._ensure_rtk_binary", return_value=Path("rtk")),
|
||||
patch("headroom.rtk.installer.register_agent_hooks", return_value=False),
|
||||
):
|
||||
result = runner.invoke(main, ["wrap", "cursor", "--prepare-only"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
cursorrules = Path(".cursorrules")
|
||||
assert cursorrules.exists()
|
||||
assert "headroom:rtk-instructions" in cursorrules.read_text(encoding="utf-8")
|
||||
|
||||
|
||||
def test_wrap_cursor_prepare_only_uses_lean_ctx_when_configured(
|
||||
monkeypatch, tmp_path: Path
|
||||
) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
monkeypatch.setenv("HEADROOM_CONTEXT_TOOL", "lean-ctx")
|
||||
runner = CliRunner()
|
||||
|
||||
with runner.isolated_filesystem(temp_dir=str(tmp_path)):
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary") as ensure_rtk:
|
||||
with patch(
|
||||
"headroom.cli.wrap._setup_lean_ctx_agent",
|
||||
return_value=Path("lean-ctx"),
|
||||
) as setup:
|
||||
result = runner.invoke(main, ["wrap", "cursor", "--prepare-only"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
ensure_rtk.assert_not_called()
|
||||
setup.assert_called_once_with("cursor", verbose=False)
|
||||
assert not Path(".cursorrules").exists()
|
||||
assert result.exit_code != 0
|
||||
assert "have been removed from Headroom" in result.output
|
||||
|
||||
|
||||
def test_wrap_openclaw_prepare_only_emits_config_without_python_default() -> None:
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ from __future__ import annotations
|
|||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import click
|
||||
import pytest
|
||||
|
||||
from headroom.cli import wrap as wrap_cli
|
||||
|
|
@ -157,21 +158,50 @@ def test_restore_noop_when_file_corrupt(tmp_path: Path) -> None:
|
|||
wrap_cli._restore_claude_wrap_base_url(None, settings_path=path) # must not raise
|
||||
|
||||
|
||||
def test_write_recovers_from_corrupt_file(tmp_path: Path) -> None:
|
||||
def test_write_refuses_to_clobber_a_corrupt_file(tmp_path: Path) -> None:
|
||||
"""A file that will not parse is DATA, not a blank slate — never overwrite it.
|
||||
|
||||
This previously "recovered" by resetting the payload to ``{}`` and writing
|
||||
that back, so a single hand-edited typo (or a transient read error) silently
|
||||
destroyed the user's whole settings file — permissions, env and hooks — on
|
||||
every ``headroom wrap claude``. Refusing leaves the file for the user to fix.
|
||||
"""
|
||||
path = _settings(tmp_path)
|
||||
path.parent.mkdir(parents=True)
|
||||
path.write_text("not valid json {{{{", encoding="utf-8")
|
||||
prev = wrap_cli._write_claude_wrap_base_url("http://127.0.0.1:8787", settings_path=path)
|
||||
assert prev is None # treated as fresh
|
||||
payload = json.loads(path.read_text(encoding="utf-8"))
|
||||
assert payload["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:8787"
|
||||
original = '{"permissions": {"allow": ["Bash"]}, oops'
|
||||
path.write_text(original, encoding="utf-8")
|
||||
|
||||
with pytest.raises(click.ClickException, match="not valid JSON"):
|
||||
wrap_cli._write_claude_wrap_base_url("http://127.0.0.1:8787", settings_path=path)
|
||||
|
||||
assert path.read_text(encoding="utf-8") == original # untouched
|
||||
|
||||
|
||||
def test_write_recovers_from_non_dict_payload(tmp_path: Path) -> None:
|
||||
def test_write_refuses_non_dict_payload(tmp_path: Path) -> None:
|
||||
path = _settings(tmp_path)
|
||||
path.parent.mkdir(parents=True)
|
||||
path.write_text("[1, 2, 3]", encoding="utf-8") # valid JSON but not a dict
|
||||
original = "[1, 2, 3]" # valid JSON but not a settings object
|
||||
path.write_text(original, encoding="utf-8")
|
||||
|
||||
with pytest.raises(click.ClickException, match="does not contain a JSON object"):
|
||||
wrap_cli._write_claude_wrap_base_url("http://127.0.0.1:8787", settings_path=path)
|
||||
|
||||
assert path.read_text(encoding="utf-8") == original # untouched
|
||||
|
||||
|
||||
def test_write_recovers_from_an_empty_file(tmp_path: Path) -> None:
|
||||
"""An empty file has no settings to lose, so recover rather than strand the user.
|
||||
|
||||
A zero-byte settings.json is the classic residue of an interrupted
|
||||
non-atomic write, so this is the one case where treating the file as fresh
|
||||
is both safe and the helpful thing to do.
|
||||
"""
|
||||
path = _settings(tmp_path)
|
||||
path.parent.mkdir(parents=True)
|
||||
path.write_text(" \n", encoding="utf-8")
|
||||
|
||||
prev = wrap_cli._write_claude_wrap_base_url("http://127.0.0.1:8787", settings_path=path)
|
||||
|
||||
assert prev is None
|
||||
payload = json.loads(path.read_text(encoding="utf-8"))
|
||||
assert payload["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:8787"
|
||||
|
|
|
|||
|
|
@ -59,7 +59,7 @@ def test_finally_survives_early_proxy_start_failure(
|
|||
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "claude", "--no-context-tool", "--no-mcp", "--no-tokensave", "--no-serena"],
|
||||
["wrap", "claude", "--no-mcp", "--no-tokensave", "--no-serena"],
|
||||
)
|
||||
|
||||
# The finally must complete: no UnboundLocalError masking the real failure,
|
||||
|
|
|
|||
|
|
@ -102,7 +102,6 @@ def _invoke_wrap_claude(
|
|||
[
|
||||
"wrap",
|
||||
"claude",
|
||||
"--no-context-tool",
|
||||
"--no-mcp",
|
||||
"--no-tokensave",
|
||||
"--no-serena",
|
||||
|
|
|
|||
|
|
@ -464,45 +464,6 @@ class TestInjectAndRestoreRoundTrip:
|
|||
assert status == "restored"
|
||||
assert config_file.read_text(encoding="utf-8") == malformed
|
||||
|
||||
def test_unwrap_removes_rtk_block_from_global_agents(
|
||||
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
"""`wrap codex` injects the rtk block into the Codex global AGENTS.md;
|
||||
`unwrap codex` must take it back out (regression for #1421)."""
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
monkeypatch.setenv("HEADROOM_RTK", "1")
|
||||
codex_home = tmp_path / ".codex"
|
||||
codex_home.mkdir()
|
||||
agents = codex_home / "AGENTS.md"
|
||||
wrap_mod._inject_rtk_instructions(agents)
|
||||
assert wrap_mod._RTK_MARKER in agents.read_text(encoding="utf-8")
|
||||
|
||||
wrap_mod.unwrap_codex.callback(port=8787, no_stop_proxy=True)
|
||||
|
||||
remaining = agents.read_text(encoding="utf-8") if agents.exists() else ""
|
||||
assert wrap_mod._RTK_MARKER not in remaining
|
||||
|
||||
def test_unwrap_preserves_user_content_in_global_agents(
|
||||
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
"""Only the marker-fenced rtk block is removed; the user's own AGENTS.md
|
||||
prose survives the unwrap."""
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
monkeypatch.setenv("HEADROOM_RTK", "1")
|
||||
codex_home = tmp_path / ".codex"
|
||||
codex_home.mkdir()
|
||||
agents = codex_home / "AGENTS.md"
|
||||
agents.write_text("# My project rules\n\nAlways write tests.\n", encoding="utf-8")
|
||||
wrap_mod._inject_rtk_instructions(agents)
|
||||
assert wrap_mod._RTK_MARKER in agents.read_text(encoding="utf-8")
|
||||
|
||||
wrap_mod.unwrap_codex.callback(port=8787, no_stop_proxy=True)
|
||||
|
||||
remaining = agents.read_text(encoding="utf-8")
|
||||
assert wrap_mod._RTK_MARKER not in remaining
|
||||
assert "# My project rules" in remaining
|
||||
assert "Always write tests." in remaining
|
||||
|
||||
def test_unwrap_is_safe_when_no_global_agents(
|
||||
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
|
|
@ -554,8 +515,7 @@ class TestWrapRetagsThreadProviders:
|
|||
self._seed_threads(gui_db, [("a", "openai"), ("b", "headroom"), ("c", "anthropic")])
|
||||
self._seed_threads(cli_db, [("d", "openai")])
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
||||
wrap_result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"])
|
||||
wrap_result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"])
|
||||
assert wrap_result.exit_code == 0, wrap_result.output
|
||||
# Native threads are now visible under the headroom provider menu;
|
||||
# third-party providers are left untouched.
|
||||
|
|
@ -1021,8 +981,7 @@ def test_wrap_codex_prepare_only_creates_backup_and_config(
|
|||
original = 'model_provider = "openai"\n'
|
||||
config_file.write_text(original, encoding="utf-8")
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
||||
result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"])
|
||||
result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert 'model_provider = "headroom"' in config_file.read_text(encoding="utf-8")
|
||||
|
|
@ -1031,6 +990,34 @@ def test_wrap_codex_prepare_only_creates_backup_and_config(
|
|||
assert backup.read_text(encoding="utf-8") == original
|
||||
|
||||
|
||||
def test_wrap_codex_registers_mcp_when_codex_home_does_not_exist_yet(
|
||||
runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
"""MCP must register on a machine where Codex was installed but never launched.
|
||||
|
||||
``CodexRegistrar.detect()`` is just ``~/.codex`` being a directory, and that
|
||||
directory used to be created as a side effect of writing the rtk guidance
|
||||
into ``$CODEX_HOME/AGENTS.md``. Once the CLI context tools were removed,
|
||||
nothing created it, so detect() said "Codex not detected" and Headroom
|
||||
silently skipped MCP registration — leaving every compression marker the
|
||||
proxy emits unresolvable, with no error shown.
|
||||
|
||||
Every other codex test pre-creates ``~/.codex``, which is exactly why none of
|
||||
them caught it; this one deliberately does not.
|
||||
"""
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
codex_dir = tmp_path / ".codex"
|
||||
assert not codex_dir.exists() # the whole point
|
||||
|
||||
result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
config = codex_dir / "config.toml"
|
||||
assert config.exists(), "wrap codex did not persist config in the durable Codex home"
|
||||
assert "[mcp_servers.headroom]" in config.read_text(encoding="utf-8")
|
||||
assert "not detected" not in result.output
|
||||
|
||||
|
||||
def test_wrap_codex_prepare_only_respects_codex_home(
|
||||
runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
|
|
@ -1039,11 +1026,10 @@ def test_wrap_codex_prepare_only_respects_codex_home(
|
|||
codex_home.mkdir()
|
||||
monkeypatch.setenv("CODEX_HOME", str(codex_home))
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "codex", "--prepare-only", "--no-serena", "--port", "8787"],
|
||||
)
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "codex", "--prepare-only", "--no-serena", "--port", "8787"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
config_file = codex_home / "config.toml"
|
||||
|
|
@ -1090,23 +1076,22 @@ def test_wrap_codex_launch_uses_durable_codex_home(
|
|||
rollout.parent.mkdir(parents=True)
|
||||
rollout.write_text('{"type":"session_meta"}\n', encoding="utf-8")
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
||||
with patch(
|
||||
"headroom.cli.wrap.shutil.which",
|
||||
side_effect=lambda cmd: "/fake/codex" if cmd == "codex" else None,
|
||||
):
|
||||
with patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"wrap",
|
||||
"codex",
|
||||
"--port",
|
||||
"8787",
|
||||
"--no-tokensave",
|
||||
"--no-serena",
|
||||
],
|
||||
)
|
||||
with patch(
|
||||
"headroom.cli.wrap.shutil.which",
|
||||
side_effect=lambda cmd: "/fake/codex" if cmd == "codex" else None,
|
||||
):
|
||||
with patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"wrap",
|
||||
"codex",
|
||||
"--port",
|
||||
"8787",
|
||||
"--no-tokensave",
|
||||
"--no-serena",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert launch_env["CODEX_HOME"] == str(codex_home)
|
||||
|
|
@ -1227,25 +1212,23 @@ def test_wrap_codex_rejects_custom_provider_without_upstream_base_url(
|
|||
kwargs["env_vars_display"],
|
||||
)
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
||||
with patch(
|
||||
"headroom.cli.wrap.shutil.which",
|
||||
side_effect=lambda cmd: "/fake/codex" if cmd == "codex" else None,
|
||||
):
|
||||
with patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"wrap",
|
||||
"codex",
|
||||
"--port",
|
||||
"8787",
|
||||
"--no-rtk",
|
||||
"--no-mcp",
|
||||
"--no-tokensave",
|
||||
"--no-serena",
|
||||
],
|
||||
)
|
||||
with patch(
|
||||
"headroom.cli.wrap.shutil.which",
|
||||
side_effect=lambda cmd: "/fake/codex" if cmd == "codex" else None,
|
||||
):
|
||||
with patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"wrap",
|
||||
"codex",
|
||||
"--port",
|
||||
"8787",
|
||||
"--no-mcp",
|
||||
"--no-tokensave",
|
||||
"--no-serena",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "custom provider 'company' has no upstream base_url" in result.output
|
||||
|
|
@ -1276,101 +1259,29 @@ def test_wrap_codex_routes_model_provider_selected_by_config_argument(
|
|||
)
|
||||
configured_env.update(env)
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
||||
with patch(
|
||||
"headroom.cli.wrap.shutil.which",
|
||||
side_effect=lambda cmd: "/fake/codex" if cmd == "codex" else None,
|
||||
):
|
||||
with patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"wrap",
|
||||
"codex",
|
||||
"--no-rtk",
|
||||
"--no-mcp",
|
||||
"--no-tokensave",
|
||||
"--no-serena",
|
||||
"--",
|
||||
"--config",
|
||||
'model_provider="company"',
|
||||
],
|
||||
)
|
||||
with patch(
|
||||
"headroom.cli.wrap.shutil.which",
|
||||
side_effect=lambda cmd: "/fake/codex" if cmd == "codex" else None,
|
||||
):
|
||||
with patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"wrap",
|
||||
"codex",
|
||||
"--no-mcp",
|
||||
"--no-tokensave",
|
||||
"--no-serena",
|
||||
"--",
|
||||
"--config",
|
||||
'model_provider="company"',
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert configured_env[wrap_mod._UPSTREAM_BASE_URL_ENV_VAR] == ("https://api.example.test/v1")
|
||||
|
||||
|
||||
def test_wrap_codex_injects_rtk_globally_without_changing_project_agents(
|
||||
runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
monkeypatch.setenv("HEADROOM_RTK", "1")
|
||||
project_dir = tmp_path / "project"
|
||||
project_dir.mkdir()
|
||||
project_agents = project_dir / "AGENTS.md"
|
||||
original = "# Project instructions\n\nUse the repository conventions.\n"
|
||||
project_agents.write_text(original, encoding="utf-8")
|
||||
original_bytes = project_agents.read_bytes()
|
||||
monkeypatch.chdir(project_dir)
|
||||
|
||||
with patch(
|
||||
"headroom.cli.wrap._ensure_rtk_binary",
|
||||
return_value=tmp_path / "rtk",
|
||||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"wrap",
|
||||
"codex",
|
||||
"--prepare-only",
|
||||
"--no-mcp",
|
||||
"--no-serena",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert project_agents.read_bytes() == original_bytes
|
||||
global_agents = tmp_path / ".codex" / "AGENTS.md"
|
||||
assert wrap_mod._RTK_MARKER.encode() in global_agents.read_bytes()
|
||||
|
||||
|
||||
def test_wrap_codex_launch_injects_rtk_globally_without_changing_project_agents(
|
||||
runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
monkeypatch.setenv("HEADROOM_RTK", "1")
|
||||
project_dir = tmp_path / "project"
|
||||
project_dir.mkdir()
|
||||
project_agents = project_dir / "AGENTS.md"
|
||||
original = "# Project instructions\n\nUse the repository conventions.\n"
|
||||
project_agents.write_text(original, encoding="utf-8")
|
||||
original_bytes = project_agents.read_bytes()
|
||||
monkeypatch.chdir(project_dir)
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=tmp_path / "rtk"):
|
||||
with patch(
|
||||
"headroom.cli.wrap.shutil.which",
|
||||
side_effect=lambda cmd: "/fake/codex" if cmd == "codex" else None,
|
||||
):
|
||||
with patch("headroom.cli.wrap._launch_tool"):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"wrap",
|
||||
"codex",
|
||||
"--no-mcp",
|
||||
"--no-serena",
|
||||
"--no-tokensave",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert project_agents.read_bytes() == original_bytes
|
||||
global_agents = tmp_path / ".codex" / "AGENTS.md"
|
||||
assert wrap_mod._RTK_MARKER.encode() in global_agents.read_bytes()
|
||||
|
||||
|
||||
def test_unwrap_codex_without_codex_home_warns_on_ambiguous_noop(
|
||||
runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
|
|
@ -1379,19 +1290,18 @@ def test_unwrap_codex_without_codex_home_warns_on_ambiguous_noop(
|
|||
codex_home.mkdir()
|
||||
monkeypatch.setenv("CODEX_HOME", str(codex_home))
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
||||
wrap_result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"wrap",
|
||||
"codex",
|
||||
"--prepare-only",
|
||||
"--no-mcp",
|
||||
"--no-serena",
|
||||
"--port",
|
||||
"8787",
|
||||
],
|
||||
)
|
||||
wrap_result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"wrap",
|
||||
"codex",
|
||||
"--prepare-only",
|
||||
"--no-mcp",
|
||||
"--no-serena",
|
||||
"--port",
|
||||
"8787",
|
||||
],
|
||||
)
|
||||
|
||||
assert wrap_result.exit_code == 0, wrap_result.output
|
||||
config_file = codex_home / "config.toml"
|
||||
|
|
@ -1556,8 +1466,7 @@ def test_wrap_codex_prepare_only_updates_stale_mcp_proxy_url(
|
|||
encoding="utf-8",
|
||||
)
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
||||
result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"])
|
||||
result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
content = config_file.read_text(encoding="utf-8")
|
||||
|
|
@ -1603,28 +1512,27 @@ def test_wrap_codex_memory_prepare_only_uses_local_db_without_persisting_it(
|
|||
imported_users.append(user_id)
|
||||
return 0
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
||||
with patch("headroom.memory.sync._build_sync_backend", side_effect=fake_build_sync_backend):
|
||||
with patch("headroom.memory.sync.sync_import", side_effect=fake_sync_import):
|
||||
with patch("headroom.memory.sync._build_sync_backend", side_effect=fake_build_sync_backend):
|
||||
with patch("headroom.memory.sync.sync_import", side_effect=fake_sync_import):
|
||||
with patch(
|
||||
"headroom.memory.sync_adapters.claude_code.ClaudeCodeAdapter",
|
||||
FakeClaudeCodeAdapter,
|
||||
):
|
||||
with patch(
|
||||
"headroom.memory.sync_adapters.claude_code.ClaudeCodeAdapter",
|
||||
FakeClaudeCodeAdapter,
|
||||
"headroom.memory.sync_adapters.claude_code.get_claude_memory_dir",
|
||||
return_value=tmp_path / "claude-memory",
|
||||
):
|
||||
with patch(
|
||||
"headroom.memory.sync_adapters.claude_code.get_claude_memory_dir",
|
||||
return_value=tmp_path / "claude-memory",
|
||||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"wrap",
|
||||
"codex",
|
||||
"--memory",
|
||||
"--prepare-only",
|
||||
"--no-mcp",
|
||||
"--no-serena",
|
||||
],
|
||||
)
|
||||
result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"wrap",
|
||||
"codex",
|
||||
"--memory",
|
||||
"--prepare-only",
|
||||
"--no-mcp",
|
||||
"--no-serena",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert backend_paths == [str(project_dir / ".headroom" / "memory.db")]
|
||||
|
|
@ -1648,10 +1556,9 @@ def test_wrap_codex_prepare_only_registers_serena_when_uvx_exists(
|
|||
return "/usr/local/bin/uvx"
|
||||
return None
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
||||
with patch("headroom.cli.wrap.shutil.which", side_effect=fake_which):
|
||||
# Serena is the code-memory MCP; assert it lands in the codex config.
|
||||
result = runner.invoke(main, ["wrap", "codex", "--prepare-only"])
|
||||
with patch("headroom.cli.wrap.shutil.which", side_effect=fake_which):
|
||||
# Serena is the code-memory MCP; assert it lands in the codex config.
|
||||
result = runner.invoke(main, ["wrap", "codex", "--prepare-only"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
content = config_file.read_text(encoding="utf-8")
|
||||
|
|
@ -1667,8 +1574,7 @@ def test_wrap_codex_prepare_only_no_serena_skips_serena(
|
|||
config_file = tmp_path / ".codex" / "config.toml"
|
||||
config_file.parent.mkdir(parents=True)
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
||||
result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--no-serena"])
|
||||
result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--no-serena"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "[mcp_servers.serena]" not in config_file.read_text(encoding="utf-8")
|
||||
|
|
@ -1690,8 +1596,7 @@ def test_unwrap_codex_restores_prior_config_end_to_end(
|
|||
)
|
||||
config_file.write_text(original, encoding="utf-8")
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
||||
wrap_result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"])
|
||||
wrap_result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"])
|
||||
assert wrap_result.exit_code == 0, wrap_result.output
|
||||
assert 'model_provider = "headroom"' in config_file.read_text(encoding="utf-8")
|
||||
|
||||
|
|
@ -1746,28 +1651,27 @@ def test_wrap_codex_memory_prepare_only_unwrap_removes_memory_mcp_without_prior_
|
|||
async def fake_sync_import(backend: FakeBackend, adapter: object, user_id: str) -> int:
|
||||
return 0
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
||||
with patch("headroom.memory.sync._build_sync_backend", return_value=FakeBackend()):
|
||||
with patch("headroom.memory.sync.sync_import", side_effect=fake_sync_import):
|
||||
with patch("headroom.memory.sync._build_sync_backend", return_value=FakeBackend()):
|
||||
with patch("headroom.memory.sync.sync_import", side_effect=fake_sync_import):
|
||||
with patch(
|
||||
"headroom.memory.sync_adapters.claude_code.ClaudeCodeAdapter",
|
||||
autospec=True,
|
||||
):
|
||||
with patch(
|
||||
"headroom.memory.sync_adapters.claude_code.ClaudeCodeAdapter",
|
||||
autospec=True,
|
||||
"headroom.memory.sync_adapters.claude_code.get_claude_memory_dir",
|
||||
return_value=tmp_path / "claude-memory",
|
||||
):
|
||||
with patch(
|
||||
"headroom.memory.sync_adapters.claude_code.get_claude_memory_dir",
|
||||
return_value=tmp_path / "claude-memory",
|
||||
):
|
||||
wrap_result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"wrap",
|
||||
"codex",
|
||||
"--memory",
|
||||
"--prepare-only",
|
||||
"--no-mcp",
|
||||
"--no-serena",
|
||||
],
|
||||
)
|
||||
wrap_result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"wrap",
|
||||
"codex",
|
||||
"--memory",
|
||||
"--prepare-only",
|
||||
"--no-mcp",
|
||||
"--no-serena",
|
||||
],
|
||||
)
|
||||
|
||||
assert wrap_result.exit_code == 0, wrap_result.output
|
||||
config_file = tmp_path / ".codex" / "config.toml"
|
||||
|
|
@ -1806,22 +1710,21 @@ def test_wrap_codex_memory_launch_failure_unwrap_cleans_memory_only_config(
|
|||
def fake_which(cmd: str) -> str | None:
|
||||
return None if cmd == "codex" else shutil.which(cmd)
|
||||
|
||||
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.memory.sync._build_sync_backend", return_value=FakeBackend()):
|
||||
with patch("headroom.memory.sync.sync_import", side_effect=fake_sync_import):
|
||||
with patch("headroom.cli.wrap.shutil.which", side_effect=fake_which):
|
||||
with patch("headroom.memory.sync._build_sync_backend", return_value=FakeBackend()):
|
||||
with patch("headroom.memory.sync.sync_import", side_effect=fake_sync_import):
|
||||
with patch(
|
||||
"headroom.memory.sync_adapters.claude_code.ClaudeCodeAdapter",
|
||||
autospec=True,
|
||||
):
|
||||
with patch(
|
||||
"headroom.memory.sync_adapters.claude_code.ClaudeCodeAdapter",
|
||||
autospec=True,
|
||||
"headroom.memory.sync_adapters.claude_code.get_claude_memory_dir",
|
||||
return_value=tmp_path / "claude-memory",
|
||||
):
|
||||
with patch(
|
||||
"headroom.memory.sync_adapters.claude_code.get_claude_memory_dir",
|
||||
return_value=tmp_path / "claude-memory",
|
||||
):
|
||||
wrap_result = runner.invoke(
|
||||
main,
|
||||
["wrap", "codex", "--memory", "--no-mcp", "--no-serena"],
|
||||
)
|
||||
wrap_result = runner.invoke(
|
||||
main,
|
||||
["wrap", "codex", "--memory", "--no-mcp", "--no-serena"],
|
||||
)
|
||||
|
||||
assert wrap_result.exit_code == 1
|
||||
config_file = tmp_path / ".codex" / "config.toml"
|
||||
|
|
@ -1883,8 +1786,7 @@ def test_unwrap_codex_removes_headroom_only_config_file(
|
|||
) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
||||
wrap_result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"])
|
||||
wrap_result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"])
|
||||
assert wrap_result.exit_code == 0, wrap_result.output
|
||||
|
||||
config_file = tmp_path / ".codex" / "config.toml"
|
||||
|
|
@ -1905,8 +1807,7 @@ def test_unwrap_codex_preserves_unrelated_sections(
|
|||
original = '[mcp_servers.local_thing]\ncommand = "/usr/local/bin/thing"\nargs = ["--serve"]\n'
|
||||
config_file.write_text(original, encoding="utf-8")
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
||||
runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"])
|
||||
runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"])
|
||||
|
||||
result = runner.invoke(main, ["unwrap", "codex"])
|
||||
assert result.exit_code == 0, result.output
|
||||
|
|
@ -2003,7 +1904,7 @@ class TestCodexPortResolution:
|
|||
runner = CliRunner()
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "codex", "--port", "8787", "--no-rtk", "--no-mcp", "--no-serena"],
|
||||
["wrap", "codex", "--port", "8787", "--no-mcp", "--no-serena"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, f"CLI failed: {result.output}"
|
||||
|
|
@ -2047,7 +1948,6 @@ class TestCodexLaunchExportsCustomUpstream:
|
|||
|
||||
wrap_mod._run_codex_wrap(
|
||||
port=8787,
|
||||
no_rtk=True,
|
||||
no_mcp=True,
|
||||
no_tokensave=True,
|
||||
serena=False,
|
||||
|
|
|
|||
|
|
@ -1,316 +0,0 @@
|
|||
"""Tests for `headroom wrap continue` command (PR-G1, Phase G)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from click.testing import CliRunner
|
||||
|
||||
from headroom.cli import wrap as wrap_mod
|
||||
from headroom.cli.main import main
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def runner() -> CliRunner:
|
||||
return CliRunner()
|
||||
|
||||
|
||||
def test_inject_continue_rtk_systemmessage_new_file(tmp_path: Path) -> None:
|
||||
"""Writing into a non-existent config.json creates parents + sets systemMessage."""
|
||||
config_file = tmp_path / ".continue" / "config.json"
|
||||
assert not config_file.exists()
|
||||
|
||||
assert wrap_mod._inject_continue_rtk_systemmessage(config_file) is True
|
||||
|
||||
data = json.loads(config_file.read_text())
|
||||
assert wrap_mod._RTK_MARKER in data["systemMessage"]
|
||||
|
||||
|
||||
def test_inject_continue_rtk_systemmessage_preserves_existing_keys(tmp_path: Path) -> None:
|
||||
"""Pre-existing keys are not touched; per-model entries get systemMessage."""
|
||||
config_file = tmp_path / ".continue" / "config.json"
|
||||
config_file.parent.mkdir(parents=True)
|
||||
config_file.write_text(json.dumps({"models": [{"title": "GPT-4o", "provider": "openai"}]}))
|
||||
|
||||
wrap_mod._inject_continue_rtk_systemmessage(config_file)
|
||||
|
||||
data = json.loads(config_file.read_text())
|
||||
# Pre-existing fields on the model entry are preserved verbatim.
|
||||
assert data["models"][0]["title"] == "GPT-4o"
|
||||
assert data["models"][0]["provider"] == "openai"
|
||||
# Top-level systemMessage is set.
|
||||
assert wrap_mod._RTK_MARKER in data["systemMessage"]
|
||||
# Per-model systemMessage is also populated (Continue overrides top-level
|
||||
# with per-model when set, so we must visit each model).
|
||||
assert wrap_mod._RTK_MARKER in data["models"][0]["systemMessage"]
|
||||
|
||||
|
||||
def test_inject_continue_rtk_systemmessage_appends_to_existing_message(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Pre-existing systemMessage content is preserved; rtk block is appended."""
|
||||
config_file = tmp_path / ".continue" / "config.json"
|
||||
config_file.parent.mkdir(parents=True)
|
||||
existing_msg = "You are a helpful assistant."
|
||||
config_file.write_text(json.dumps({"systemMessage": existing_msg}))
|
||||
|
||||
wrap_mod._inject_continue_rtk_systemmessage(config_file)
|
||||
|
||||
data = json.loads(config_file.read_text())
|
||||
assert data["systemMessage"].startswith(existing_msg)
|
||||
assert wrap_mod._RTK_MARKER in data["systemMessage"]
|
||||
|
||||
|
||||
def test_inject_continue_rtk_systemmessage_idempotent(tmp_path: Path) -> None:
|
||||
"""Re-injection must not duplicate the marker."""
|
||||
config_file = tmp_path / ".continue" / "config.json"
|
||||
|
||||
wrap_mod._inject_continue_rtk_systemmessage(config_file)
|
||||
wrap_mod._inject_continue_rtk_systemmessage(config_file)
|
||||
|
||||
data = json.loads(config_file.read_text())
|
||||
assert data["systemMessage"].count(wrap_mod._RTK_MARKER) == 1
|
||||
|
||||
|
||||
def test_inject_continue_rtk_systemmessage_refuses_invalid_json(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Malformed JSON must be left untouched and the helper must return False."""
|
||||
config_file = tmp_path / ".continue" / "config.json"
|
||||
config_file.parent.mkdir(parents=True)
|
||||
malformed = '{ "models": [ this is not valid json'
|
||||
config_file.write_text(malformed)
|
||||
|
||||
result = wrap_mod._inject_continue_rtk_systemmessage(config_file)
|
||||
|
||||
assert result is False
|
||||
assert config_file.read_text() == malformed
|
||||
|
||||
|
||||
def test_inject_continue_rtk_systemmessage_refuses_non_object_root(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""A JSON array at the root is not a valid Continue config; leave untouched."""
|
||||
config_file = tmp_path / ".continue" / "config.json"
|
||||
config_file.parent.mkdir(parents=True)
|
||||
config_file.write_text("[]")
|
||||
|
||||
result = wrap_mod._inject_continue_rtk_systemmessage(config_file)
|
||||
|
||||
assert result is False
|
||||
assert config_file.read_text() == "[]"
|
||||
|
||||
|
||||
def test_wrap_continue_prepare_only_injects_systemmessage(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""`wrap continue --prepare-only` injects into ./.continue/config.json by default."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "continue", "--prepare-only"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
config_file = tmp_path / ".continue" / "config.json"
|
||||
assert config_file.exists()
|
||||
data = json.loads(config_file.read_text())
|
||||
assert wrap_mod._RTK_MARKER in data["systemMessage"]
|
||||
|
||||
|
||||
def test_wrap_continue_respects_custom_config_path(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""--config writes to the user-specified path, not the cwd default."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
custom_config = tmp_path / "custom" / "my-continue.json"
|
||||
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "continue", "--prepare-only", "--config", str(custom_config)],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert custom_config.exists()
|
||||
assert not (tmp_path / ".continue" / "config.json").exists()
|
||||
data = json.loads(custom_config.read_text())
|
||||
assert wrap_mod._RTK_MARKER in data["systemMessage"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# H1: non-string systemMessage must NOT be silently clobbered.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"non_string_value",
|
||||
[
|
||||
{"role": "system", "content": "You are helpful."}, # dict
|
||||
["You are helpful.", "Respond in JSON."], # list
|
||||
42, # int
|
||||
],
|
||||
ids=["dict", "list", "int"],
|
||||
)
|
||||
def test_inject_continue_rtk_systemmessage_refuses_non_string_top_level(
|
||||
tmp_path: Path,
|
||||
non_string_value: object,
|
||||
) -> None:
|
||||
"""A non-string top-level systemMessage must NOT be overwritten."""
|
||||
config_file = tmp_path / ".continue" / "config.json"
|
||||
config_file.parent.mkdir(parents=True)
|
||||
original_payload = {"systemMessage": non_string_value, "other": "untouched"}
|
||||
config_file.write_text(json.dumps(original_payload))
|
||||
original_bytes = config_file.read_bytes()
|
||||
|
||||
result = wrap_mod._inject_continue_rtk_systemmessage(config_file)
|
||||
|
||||
assert result is False, "must report refusal when user data would be clobbered"
|
||||
# File must be byte-identical to before.
|
||||
assert config_file.read_bytes() == original_bytes
|
||||
data = json.loads(config_file.read_text())
|
||||
assert data["systemMessage"] == non_string_value
|
||||
assert data["other"] == "untouched"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"non_string_value",
|
||||
[
|
||||
{"role": "system", "content": "Per-model system."},
|
||||
["List", "of", "strings"],
|
||||
7,
|
||||
],
|
||||
ids=["dict", "list", "int"],
|
||||
)
|
||||
def test_inject_continue_rtk_systemmessage_refuses_non_string_per_model(
|
||||
tmp_path: Path,
|
||||
non_string_value: object,
|
||||
) -> None:
|
||||
"""A non-string per-model systemMessage must NOT be overwritten."""
|
||||
config_file = tmp_path / ".continue" / "config.json"
|
||||
config_file.parent.mkdir(parents=True)
|
||||
original_payload = {
|
||||
"models": [
|
||||
{"title": "GPT-4o", "provider": "openai", "systemMessage": non_string_value},
|
||||
],
|
||||
}
|
||||
config_file.write_text(json.dumps(original_payload))
|
||||
|
||||
result = wrap_mod._inject_continue_rtk_systemmessage(config_file)
|
||||
|
||||
assert result is False, "must report refusal when per-model user data would be clobbered"
|
||||
data = json.loads(config_file.read_text())
|
||||
# The non-string per-model value must be preserved.
|
||||
assert data["models"][0]["systemMessage"] == non_string_value
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# M2: per-model systemMessage handling.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_inject_continue_rtk_systemmessage_visits_each_model(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Each models[i].systemMessage gets the RTK block."""
|
||||
config_file = tmp_path / ".continue" / "config.json"
|
||||
config_file.parent.mkdir(parents=True)
|
||||
config_file.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"models": [
|
||||
{"title": "A", "systemMessage": "user value"},
|
||||
{"title": "B"}, # no systemMessage yet
|
||||
],
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
assert wrap_mod._inject_continue_rtk_systemmessage(config_file) is True
|
||||
|
||||
data = json.loads(config_file.read_text())
|
||||
# Pre-existing per-model systemMessage is preserved + RTK block appended.
|
||||
assert "user value" in data["models"][0]["systemMessage"]
|
||||
assert wrap_mod._RTK_MARKER in data["models"][0]["systemMessage"]
|
||||
# Model with no systemMessage gets the RTK block fresh.
|
||||
assert wrap_mod._RTK_MARKER in data["models"][1]["systemMessage"]
|
||||
# Top-level also populated.
|
||||
assert wrap_mod._RTK_MARKER in data["systemMessage"]
|
||||
|
||||
|
||||
def test_inject_continue_rtk_systemmessage_per_model_idempotent(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Re-running must not duplicate per-model RTK blocks."""
|
||||
config_file = tmp_path / ".continue" / "config.json"
|
||||
config_file.parent.mkdir(parents=True)
|
||||
config_file.write_text(
|
||||
json.dumps({"models": [{"title": "A", "systemMessage": "user"}]}),
|
||||
)
|
||||
|
||||
wrap_mod._inject_continue_rtk_systemmessage(config_file)
|
||||
wrap_mod._inject_continue_rtk_systemmessage(config_file)
|
||||
|
||||
data = json.loads(config_file.read_text())
|
||||
assert data["models"][0]["systemMessage"].count(wrap_mod._RTK_MARKER) == 1
|
||||
assert data["systemMessage"].count(wrap_mod._RTK_MARKER) == 1
|
||||
|
||||
|
||||
def test_inject_continue_rtk_systemmessage_skips_non_dict_model_entries(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""A models[] entry that isn't a dict must be left untouched."""
|
||||
config_file = tmp_path / ".continue" / "config.json"
|
||||
config_file.parent.mkdir(parents=True)
|
||||
config_file.write_text(
|
||||
json.dumps({"models": ["just a string entry", {"title": "B"}]}),
|
||||
)
|
||||
|
||||
wrap_mod._inject_continue_rtk_systemmessage(config_file)
|
||||
|
||||
data = json.loads(config_file.read_text())
|
||||
# Non-dict entry preserved verbatim.
|
||||
assert data["models"][0] == "just a string entry"
|
||||
# Dict entry got the RTK block.
|
||||
assert wrap_mod._RTK_MARKER in data["models"][1]["systemMessage"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# M4: Ctrl-C during prelude emits a clear message.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_wrap_continue_keyboardinterrupt_during_prelude_emits_clear_message(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Ctrl-C between marker injection and proxy start must signal clearly."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
|
||||
config_file = tmp_path / ".continue" / "config.json"
|
||||
|
||||
def raise_kbd_after_inject(*args, **kwargs): # noqa: ANN002, ANN003
|
||||
# Simulate the user hitting Ctrl-C right after we wrote config.json
|
||||
# but before the proxy started.
|
||||
config_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
config_file.write_text('{"systemMessage": "marker block"}')
|
||||
raise KeyboardInterrupt
|
||||
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", side_effect=raise_kbd_after_inject):
|
||||
result = runner.invoke(main, ["wrap", "continue", "--prepare-only"])
|
||||
|
||||
assert result.exit_code == 130
|
||||
assert "interrupted" in result.output.lower()
|
||||
assert "idempotent" in result.output.lower()
|
||||
assert config_file.exists()
|
||||
assert str(config_file) in result.output
|
||||
|
|
@ -22,9 +22,9 @@ def _expected_project_prefix() -> str:
|
|||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _enable_rtk(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# RTK is opt-in (off by default); these tests exercise the RTK-on injection path.
|
||||
monkeypatch.setenv("HEADROOM_RTK", "1")
|
||||
def _no_retired_context_tool_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""A developer's exported HEADROOM_CONTEXT_TOOL would abort every wrap below."""
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -94,13 +94,13 @@ def wrap_modules(monkeypatch: pytest.MonkeyPatch) -> tuple[types.ModuleType, cli
|
|||
headroom_pkg.cli = saved_headroom_cli_attr
|
||||
|
||||
|
||||
def test_wrap_copilot_auto_anthropic_injects_instructions(
|
||||
def test_wrap_copilot_auto_anthropic_sets_provider_env(
|
||||
runner: CliRunner,
|
||||
wrap_modules: tuple[types.ModuleType, click.Group],
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
wrap_cli, main = wrap_modules
|
||||
_wrap_cli, main = wrap_modules
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test-dummy")
|
||||
captured: dict[str, object] = {}
|
||||
|
|
@ -111,7 +111,6 @@ def test_wrap_copilot_auto_anthropic_injects_instructions(
|
|||
with (
|
||||
patch("headroom.cli.wrap.shutil.which", return_value="copilot"),
|
||||
patch("headroom.cli.wrap.has_oauth_auth", return_value=False),
|
||||
patch("headroom.cli.wrap._ensure_rtk_binary", return_value=Path("/tmp/rtk")),
|
||||
patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool),
|
||||
):
|
||||
result = runner.invoke(
|
||||
|
|
@ -120,12 +119,6 @@ def test_wrap_copilot_auto_anthropic_injects_instructions(
|
|||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
instructions = tmp_path / ".github" / "copilot-instructions.md"
|
||||
assert instructions.exists()
|
||||
content = instructions.read_text(encoding="utf-8")
|
||||
assert wrap_cli._RTK_MARKER in content
|
||||
assert "RTK (Rust Token Killer)" in content
|
||||
|
||||
env = captured["env"]
|
||||
assert isinstance(env, dict)
|
||||
assert env["COPILOT_PROVIDER_TYPE"] == "anthropic"
|
||||
|
|
@ -158,7 +151,6 @@ def test_wrap_copilot_openai_backend_sets_completions_env(
|
|||
[
|
||||
"wrap",
|
||||
"copilot",
|
||||
"--no-rtk",
|
||||
"--backend",
|
||||
"anyllm",
|
||||
"--anyllm-provider",
|
||||
|
|
@ -205,7 +197,6 @@ def test_wrap_copilot_byok_rejects_auto_model_before_launch(
|
|||
"copilot",
|
||||
"--provider-type",
|
||||
"openai",
|
||||
"--no-context-tool",
|
||||
"--",
|
||||
"--model",
|
||||
"auto",
|
||||
|
|
@ -238,7 +229,7 @@ def test_wrap_copilot_auto_detects_running_proxy_backend(
|
|||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "copilot", "--no-rtk", "--", "--model", "gpt-4o"],
|
||||
["wrap", "copilot", "--", "--model", "gpt-4o"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
|
|
@ -269,7 +260,7 @@ def test_wrap_copilot_prefers_existing_oauth_session(
|
|||
with patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "copilot", "--no-rtk", "--", "--model", "claude-sonnet-4.6"],
|
||||
["wrap", "copilot", "--", "--model", "claude-sonnet-4.6"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
|
|
@ -315,7 +306,7 @@ def test_wrap_copilot_subscription_uses_github_auth_without_provider_key(
|
|||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "copilot", "--subscription", "--no-rtk"],
|
||||
["wrap", "copilot", "--subscription"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
|
|
@ -355,7 +346,7 @@ def test_wrap_copilot_subscription_defaults_to_responses_for_reasoning_model(
|
|||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "copilot", "--subscription", "--no-rtk", "--", "--model", "gpt-5.4"],
|
||||
["wrap", "copilot", "--subscription", "--", "--model", "gpt-5.4"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
|
|
@ -395,7 +386,7 @@ def test_wrap_copilot_subscription_keeps_gpt4_on_completions(
|
|||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "copilot", "--subscription", "--no-rtk", "--", "--model", "gpt-4.1"],
|
||||
["wrap", "copilot", "--subscription", "--", "--model", "gpt-4.1"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
|
|
@ -434,7 +425,6 @@ def test_wrap_copilot_subscription_allows_explicit_responses_wire_api(
|
|||
"--subscription",
|
||||
"--wire-api",
|
||||
"responses",
|
||||
"--no-rtk",
|
||||
"--",
|
||||
"--model",
|
||||
"gpt-5.4",
|
||||
|
|
@ -487,7 +477,7 @@ def test_wrap_copilot_subscription_pins_validated_token_for_proxy(
|
|||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "copilot", "--subscription", "--no-rtk"],
|
||||
["wrap", "copilot", "--subscription"],
|
||||
env={
|
||||
"GITHUB_COPILOT_API_TOKEN": "stale-parent-token",
|
||||
"GITHUB_COPILOT_REFRESH_OAUTH_TOKEN": "stale-parent-refresh",
|
||||
|
|
@ -525,7 +515,7 @@ def test_wrap_copilot_subscription_requires_reusable_auth(
|
|||
patch("headroom.cli.wrap.shutil.which", return_value="copilot"),
|
||||
patch("headroom.cli.wrap.resolve_subscription_bearer_token_details", return_value=None),
|
||||
):
|
||||
result = runner.invoke(main, ["wrap", "copilot", "--subscription", "--no-rtk"])
|
||||
result = runner.invoke(main, ["wrap", "copilot", "--subscription"])
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "subscription mode requires a reusable GitHub/Copilot bearer token" in result.output
|
||||
|
|
@ -540,7 +530,7 @@ def test_wrap_copilot_subscription_rejects_translated_backend(
|
|||
with patch("headroom.cli.wrap.shutil.which", return_value="copilot"):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "copilot", "--subscription", "--backend", "anyllm", "--no-rtk"],
|
||||
["wrap", "copilot", "--subscription", "--backend", "anyllm"],
|
||||
)
|
||||
|
||||
assert result.exit_code != 0
|
||||
|
|
@ -555,7 +545,7 @@ def test_wrap_copilot_subscription_rejects_anthropic_provider_type(
|
|||
with patch("headroom.cli.wrap.shutil.which", return_value="copilot"):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "copilot", "--subscription", "--provider-type", "anthropic", "--no-rtk"],
|
||||
["wrap", "copilot", "--subscription", "--provider-type", "anthropic"],
|
||||
)
|
||||
|
||||
assert result.exit_code != 0
|
||||
|
|
@ -589,7 +579,6 @@ def test_wrap_copilot_translated_backend_still_requires_byok(
|
|||
[
|
||||
"wrap",
|
||||
"copilot",
|
||||
"--no-rtk",
|
||||
"--backend",
|
||||
"anyllm",
|
||||
"--",
|
||||
|
|
@ -669,7 +658,7 @@ def test_wrap_copilot_clears_stale_wire_api_in_anthropic_mode(
|
|||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "copilot", "--no-rtk", "--", "--model", "claude-sonnet-4-20250514"],
|
||||
["wrap", "copilot", "--", "--model", "claude-sonnet-4-20250514"],
|
||||
env={
|
||||
"COPILOT_PROVIDER_WIRE_API": "responses",
|
||||
"ANTHROPIC_API_KEY": "sk-test-dummy",
|
||||
|
|
@ -696,20 +685,19 @@ def test_wrap_copilot_fails_when_binary_missing(
|
|||
assert "Install GitHub Copilot CLI" in result.output
|
||||
|
||||
|
||||
def test_unwrap_copilot_removes_rtk_instructions_and_stops_proxy(
|
||||
def test_unwrap_copilot_stops_proxy(
|
||||
runner: CliRunner,
|
||||
wrap_modules: tuple[types.ModuleType, click.Group],
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
wrap_cli, main = wrap_modules
|
||||
"""`unwrap copilot` stops the local proxy on the requested port.
|
||||
|
||||
Copilot is env-var wrapped, so there is no config to restore — stopping the
|
||||
proxy (and reporting it) is the whole contract.
|
||||
"""
|
||||
_wrap_cli, main = wrap_modules
|
||||
monkeypatch.chdir(tmp_path)
|
||||
instructions = tmp_path / ".github" / "copilot-instructions.md"
|
||||
instructions.parent.mkdir()
|
||||
instructions.write_text(
|
||||
"Keep user guidance.\n\n" + wrap_cli.RTK_INSTRUCTIONS_BLOCK,
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
with patch(
|
||||
"headroom.cli.wrap._stop_local_proxy_for_unwrap",
|
||||
|
|
@ -718,93 +706,27 @@ def test_unwrap_copilot_removes_rtk_instructions_and_stops_proxy(
|
|||
result = runner.invoke(main, ["unwrap", "copilot", "--port", "9999"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert instructions.read_text(encoding="utf-8") == "Keep user guidance.\n"
|
||||
stop_proxy.assert_called_once_with(9999)
|
||||
assert "Removed Headroom rtk instructions from Copilot." in result.output
|
||||
assert "Stopped local Headroom proxy on port 9999" in result.output
|
||||
|
||||
|
||||
def test_unwrap_copilot_preserves_instructions_after_rtk_block(
|
||||
def test_unwrap_copilot_leaves_user_instruction_file_untouched(
|
||||
runner: CliRunner,
|
||||
wrap_modules: tuple[types.ModuleType, click.Group],
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
wrap_cli, main = wrap_modules
|
||||
monkeypatch.chdir(tmp_path)
|
||||
instructions = tmp_path / ".github" / "copilot-instructions.md"
|
||||
instructions.parent.mkdir()
|
||||
instructions.write_text(
|
||||
wrap_cli.RTK_INSTRUCTIONS_BLOCK + "\nKeep trailing guidance.\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
result = runner.invoke(main, ["unwrap", "copilot", "--no-stop-proxy"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert instructions.read_text(encoding="utf-8") == "Keep trailing guidance.\n"
|
||||
|
||||
|
||||
def test_unwrap_copilot_leaves_malformed_marker_content_unchanged(
|
||||
runner: CliRunner,
|
||||
wrap_modules: tuple[types.ModuleType, click.Group],
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
wrap_cli, main = wrap_modules
|
||||
monkeypatch.chdir(tmp_path)
|
||||
instructions = tmp_path / ".github" / "copilot-instructions.md"
|
||||
instructions.parent.mkdir()
|
||||
content = f"<!-- /headroom:rtk-instructions -->\nKeep user guidance.\n{wrap_cli._RTK_MARKER}\n"
|
||||
instructions.write_text(content, encoding="utf-8")
|
||||
|
||||
result = runner.invoke(main, ["unwrap", "copilot", "--no-stop-proxy"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert instructions.read_text(encoding="utf-8") == content
|
||||
assert "No Headroom rtk instructions found for Copilot." in result.output
|
||||
|
||||
|
||||
def test_unwrap_copilot_deletes_generated_only_instruction_file(
|
||||
runner: CliRunner,
|
||||
wrap_modules: tuple[types.ModuleType, click.Group],
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
wrap_cli, main = wrap_modules
|
||||
monkeypatch.chdir(tmp_path)
|
||||
instructions = tmp_path / ".github" / "copilot-instructions.md"
|
||||
instructions.parent.mkdir()
|
||||
instructions.write_text(wrap_cli.RTK_INSTRUCTIONS_BLOCK, encoding="utf-8")
|
||||
|
||||
result = runner.invoke(main, ["unwrap", "copilot", "--no-stop-proxy"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert not instructions.exists()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("create_user_file", [False, True])
|
||||
def test_unwrap_copilot_is_noop_without_managed_instructions(
|
||||
runner: CliRunner,
|
||||
wrap_modules: tuple[types.ModuleType, click.Group],
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
create_user_file: bool,
|
||||
) -> None:
|
||||
"""A user-authored copilot-instructions.md is never rewritten or deleted."""
|
||||
_wrap_cli, main = wrap_modules
|
||||
monkeypatch.chdir(tmp_path)
|
||||
instructions = tmp_path / ".github" / "copilot-instructions.md"
|
||||
if create_user_file:
|
||||
instructions.parent.mkdir()
|
||||
instructions.write_text("Keep user guidance.\n", encoding="utf-8")
|
||||
instructions.parent.mkdir()
|
||||
instructions.write_text("Keep user guidance.\n", encoding="utf-8")
|
||||
|
||||
result = runner.invoke(main, ["unwrap", "copilot", "--no-stop-proxy"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert instructions.exists() is create_user_file
|
||||
if create_user_file:
|
||||
assert instructions.read_text(encoding="utf-8") == "Keep user guidance.\n"
|
||||
assert "No Headroom rtk instructions found for Copilot." in result.output
|
||||
assert instructions.read_text(encoding="utf-8") == "Keep user guidance.\n"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -870,7 +792,7 @@ def test_wrap_copilot_oauth_keeps_generic_endpoint_when_account_advertised(
|
|||
patch("headroom.copilot_auth._fetch_copilot_user_info", return_value=_ACCOUNT_USER_INFO),
|
||||
patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool),
|
||||
):
|
||||
result = runner.invoke(main, ["wrap", "copilot", "--no-rtk", "--", "--model", "gpt-5.4"])
|
||||
result = runner.invoke(main, ["wrap", "copilot", "--", "--model", "gpt-5.4"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
env = captured["env"]
|
||||
|
|
@ -903,7 +825,7 @@ def test_wrap_copilot_oauth_honors_api_url_override(
|
|||
patch("headroom.copilot_auth._fetch_copilot_user_info", return_value=_ACCOUNT_USER_INFO),
|
||||
patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool),
|
||||
):
|
||||
result = runner.invoke(main, ["wrap", "copilot", "--no-rtk", "--", "--model", "gpt-5.4"])
|
||||
result = runner.invoke(main, ["wrap", "copilot", "--", "--model", "gpt-5.4"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
env = captured["env"]
|
||||
|
|
@ -940,7 +862,7 @@ def test_wrap_copilot_byok_never_resolves_copilot_endpoint(
|
|||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "copilot", "--no-rtk", "--provider-type", "openai", "--", "--model", "gpt-4o"],
|
||||
["wrap", "copilot", "--provider-type", "openai", "--", "--model", "gpt-4o"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
|
|
@ -976,7 +898,7 @@ def test_wrap_copilot_subscription_uses_resolved_subscription_endpoint(
|
|||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "copilot", "--subscription", "--no-rtk", "--", "--model", "gpt-5.4"],
|
||||
["wrap", "copilot", "--subscription", "--", "--model", "gpt-5.4"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
|
|
@ -1027,7 +949,7 @@ def test_wrap_copilot_subscription_normalizes_enterprise_host(
|
|||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "copilot", "--subscription", "--no-rtk", "--", "--model", "gpt-5.4"],
|
||||
["wrap", "copilot", "--subscription", "--", "--model", "gpt-5.4"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
|
|
@ -1069,7 +991,7 @@ def test_wrap_copilot_subscription_honors_api_url_override(
|
|||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "copilot", "--subscription", "--no-rtk", "--", "--model", "gpt-5.4"],
|
||||
["wrap", "copilot", "--subscription", "--", "--model", "gpt-5.4"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
|
|
|
|||
|
|
@ -12,16 +12,10 @@ from pathlib import Path
|
|||
|
||||
import pytest
|
||||
|
||||
from headroom.cli.wrap import (
|
||||
_MEMORY_AGENTS_MARKER,
|
||||
_RTK_MARKER,
|
||||
_inject_memory_agents_md,
|
||||
_inject_rtk_instructions,
|
||||
)
|
||||
from headroom.cli.wrap import _MEMORY_AGENTS_MARKER, _inject_memory_agents_md
|
||||
|
||||
# (inject_fn, marker) for the two prose injectors that share the bug.
|
||||
# (inject_fn, marker) for the prose injectors that shared the bug.
|
||||
INJECTORS = [
|
||||
pytest.param(_inject_rtk_instructions, _RTK_MARKER, id="rtk"),
|
||||
pytest.param(_inject_memory_agents_md, _MEMORY_AGENTS_MARKER, id="memory_agents"),
|
||||
]
|
||||
|
||||
|
|
@ -31,11 +25,7 @@ _EXISTING = "Be in “happy places” — really.\n".encode() + b"legacy \x9d by
|
|||
|
||||
|
||||
@pytest.mark.parametrize("inject, marker", INJECTORS)
|
||||
def test_inject_appends_into_file_with_non_ascii_and_stray_byte(
|
||||
inject, marker, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
if marker == _RTK_MARKER:
|
||||
monkeypatch.setenv("HEADROOM_RTK", "1")
|
||||
def test_inject_appends_into_file_with_non_ascii_and_stray_byte(inject, marker, tmp_path: Path):
|
||||
target = tmp_path / "AGENTS.md"
|
||||
target.write_bytes(_EXISTING)
|
||||
|
||||
|
|
@ -49,11 +39,7 @@ def test_inject_appends_into_file_with_non_ascii_and_stray_byte(
|
|||
|
||||
|
||||
@pytest.mark.parametrize("inject, marker", INJECTORS)
|
||||
def test_inject_creates_file_when_absent(
|
||||
inject, marker, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
if marker == _RTK_MARKER:
|
||||
monkeypatch.setenv("HEADROOM_RTK", "1")
|
||||
def test_inject_creates_file_when_absent(inject, marker, tmp_path: Path):
|
||||
target = tmp_path / "nested" / "AGENTS.md"
|
||||
|
||||
assert inject(target) is True
|
||||
|
|
@ -61,9 +47,7 @@ def test_inject_creates_file_when_absent(
|
|||
|
||||
|
||||
@pytest.mark.parametrize("inject, marker", INJECTORS)
|
||||
def test_inject_is_idempotent(inject, marker, tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
if marker == _RTK_MARKER:
|
||||
monkeypatch.setenv("HEADROOM_RTK", "1")
|
||||
def test_inject_is_idempotent(inject, marker, tmp_path: Path):
|
||||
target = tmp_path / "AGENTS.md"
|
||||
target.write_bytes(_EXISTING)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,11 +1,7 @@
|
|||
"""Tests for `headroom wrap goose` command (PR-G1, Phase G).
|
||||
|
||||
Hint-file injection tests (.goosehints idempotency, no-context-tool,
|
||||
existing-content preservation, Ctrl-C handling) live in
|
||||
`test_wrap_hintfile_agents.py` — the shared parameterized file that
|
||||
covers `wrap cline` too. This file keeps only goose-specific behavior:
|
||||
the OPENAI/ANTHROPIC env-var fan-out for the child binary launch, and
|
||||
the goose-binary-not-found error path.
|
||||
Covers goose-specific behavior: the OPENAI/ANTHROPIC env-var fan-out for
|
||||
the child binary launch, and the goose-binary-not-found error path.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -41,8 +37,7 @@ def test_wrap_goose_sets_provider_envs(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="goose"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "goose", "--port", "9000", "--", "session"])
|
||||
result = runner.invoke(main, ["wrap", "goose", "--port", "9000", "--", "session"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
env = captured["env"]
|
||||
|
|
@ -65,8 +60,7 @@ def test_wrap_goose_missing_binary_errors_clearly(
|
|||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value=None):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "goose"])
|
||||
result = runner.invoke(main, ["wrap", "goose"])
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert "'goose' not found in PATH" in result.output
|
||||
|
|
|
|||
|
|
@ -32,13 +32,10 @@ def test_wrap_grok_sets_proxy_env(
|
|||
captured.update(kwargs)
|
||||
|
||||
with patch("headroom.cli.wrap.shutil.which", return_value="grok"):
|
||||
with patch("headroom.cli.wrap._setup_context_tool_for_agent"):
|
||||
with patch("headroom.cli.wrap._setup_headroom_mcp"):
|
||||
with patch("headroom.cli.wrap._setup_coding_compressor"):
|
||||
with patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool):
|
||||
result = runner.invoke(
|
||||
main, ["wrap", "grok", "--no-rtk", "--no-mcp", "--", "-p", "hello"]
|
||||
)
|
||||
with patch("headroom.cli.wrap._setup_headroom_mcp"):
|
||||
with patch("headroom.cli.wrap._setup_coding_compressor"):
|
||||
with patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool):
|
||||
result = runner.invoke(main, ["wrap", "grok", "--no-mcp", "--", "-p", "hello"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
env = captured["env"]
|
||||
|
|
@ -56,10 +53,9 @@ def test_wrap_grok_missing_binary_exits(
|
|||
monkeypatch.chdir(tmp_path)
|
||||
|
||||
with patch("headroom.cli.wrap.shutil.which", return_value=None):
|
||||
with patch("headroom.cli.wrap._setup_context_tool_for_agent"):
|
||||
with patch("headroom.cli.wrap._setup_headroom_mcp"):
|
||||
with patch("headroom.cli.wrap._setup_coding_compressor"):
|
||||
result = runner.invoke(main, ["wrap", "grok", "--no-rtk", "--no-mcp"])
|
||||
with patch("headroom.cli.wrap._setup_headroom_mcp"):
|
||||
with patch("headroom.cli.wrap._setup_coding_compressor"):
|
||||
result = runner.invoke(main, ["wrap", "grok", "--no-mcp"])
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert "grok" in result.output.lower()
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
"""Direct unit tests for the shared wrap-subcommand helpers.
|
||||
|
||||
These helpers (`_print_wrap_banner`, `_setup_context_tool_for_agent`,
|
||||
`_run_proxy_only_watcher`) were extracted to remove ~150 LOC of
|
||||
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
|
||||
|
|
@ -19,7 +19,6 @@ import subprocess
|
|||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import click
|
||||
import pytest
|
||||
|
|
@ -97,31 +96,10 @@ def test_print_wrap_banner_title_is_centered_or_near_centered() -> None:
|
|||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _setup_context_tool_for_agent — all five branches:
|
||||
# 1. lean-ctx mode → calls _setup_lean_ctx_agent, returns None
|
||||
# 2. rtk install success → calls on_rtk_ready, returns rtk_path
|
||||
# 3. rtk install fail + rtk_required=False → returns None silently
|
||||
# 4. rtk install fail + rtk_required=True → SystemExit(1)
|
||||
# 5. KeyboardInterrupt → _emit_wrap_interrupted, SystemExit(130)
|
||||
# wrap claude argument passthrough.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_claude_context_tool_is_opt_in_for_prepare_only(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Claude skips context-tool setup unless the positive flag is passed."""
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
runner = CliRunner()
|
||||
|
||||
with patch.object(wrap_mod, "_prepare_wrap_rtk") as prepare_rtk:
|
||||
default = runner.invoke(main, ["wrap", "claude", "--prepare-only"])
|
||||
opt_in = runner.invoke(main, ["wrap", "claude", "--prepare-only", "--context-tool"])
|
||||
|
||||
assert default.exit_code == 0, default.output
|
||||
assert opt_in.exit_code == 0, opt_in.output
|
||||
assert prepare_rtk.call_count == 1
|
||||
|
||||
|
||||
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(
|
||||
|
|
@ -132,197 +110,6 @@ def test_wrap_claude_allows_claude_print_short_flag_in_passthrough_args() -> Non
|
|||
assert result.exit_code == 0, result.output
|
||||
|
||||
|
||||
def test_claude_context_tool_opt_in_preserves_lean_ctx_selection(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""The positive flag enables the configured lean-ctx installer."""
|
||||
monkeypatch.setenv("HEADROOM_CONTEXT_TOOL", "lean-ctx")
|
||||
runner = CliRunner()
|
||||
|
||||
with patch.object(wrap_mod, "_setup_lean_ctx_agent") as setup_lean_ctx:
|
||||
result = runner.invoke(main, ["wrap", "claude", "--prepare-only", "--context-tool"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
setup_lean_ctx.assert_called_once_with("claude", verbose=False)
|
||||
|
||||
|
||||
def test_claude_no_context_tool_wins_over_context_tool(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""The legacy opt-out remains authoritative when both flags are supplied."""
|
||||
runner = CliRunner()
|
||||
|
||||
with patch.object(wrap_mod, "_prepare_wrap_rtk") as prepare_rtk:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "claude", "--prepare-only", "--context-tool", "--no-context-tool"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
prepare_rtk.assert_not_called()
|
||||
|
||||
|
||||
def test_non_claude_context_tool_setup_remains_default(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Copilot still sets up RTK without a new positive opt-in flag."""
|
||||
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test-dummy")
|
||||
|
||||
with (
|
||||
patch.object(wrap_mod.shutil, "which", return_value="copilot"),
|
||||
patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")) as ensure_rtk,
|
||||
patch.object(wrap_mod, "_launch_tool"),
|
||||
):
|
||||
result = CliRunner().invoke(
|
||||
main,
|
||||
["wrap", "copilot", "--no-proxy", "--", "--model", "claude-sonnet-4-20250514"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
ensure_rtk.assert_called_once_with(verbose=False)
|
||||
|
||||
|
||||
def test_setup_context_tool_lean_ctx_calls_lean_ctx_setup(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""When HEADROOM_CONTEXT_TOOL=lean-ctx, helper calls _setup_lean_ctx_agent."""
|
||||
monkeypatch.setenv("HEADROOM_CONTEXT_TOOL", "lean-ctx")
|
||||
called_with: dict[str, Any] = {}
|
||||
|
||||
def fake_lean_ctx(agent: str, verbose: bool = False) -> Path | None:
|
||||
called_with["agent"] = agent
|
||||
called_with["verbose"] = verbose
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(wrap_mod, "_setup_lean_ctx_agent", fake_lean_ctx)
|
||||
|
||||
runner = CliRunner()
|
||||
|
||||
@click.command()
|
||||
def _cmd() -> None:
|
||||
result = wrap_mod._setup_context_tool_for_agent(
|
||||
agent="cline",
|
||||
agent_display="Cline",
|
||||
marker_path=None,
|
||||
)
|
||||
assert result is None
|
||||
|
||||
inv = runner.invoke(_cmd)
|
||||
assert inv.exit_code == 0, inv.output
|
||||
assert called_with == {"agent": "cline", "verbose": False}
|
||||
|
||||
|
||||
def test_setup_context_tool_rtk_success_calls_on_rtk_ready(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
"""rtk install success → on_rtk_ready receives the rtk binary path."""
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
fake_rtk = Path("/tmp/rtk-fake")
|
||||
received: list[Path] = []
|
||||
|
||||
monkeypatch.setattr(wrap_mod, "_ensure_rtk_binary", lambda verbose=False: fake_rtk)
|
||||
|
||||
runner = CliRunner()
|
||||
|
||||
@click.command()
|
||||
def _cmd() -> None:
|
||||
result = wrap_mod._setup_context_tool_for_agent(
|
||||
agent="cline",
|
||||
agent_display="Cline",
|
||||
marker_path=tmp_path / ".clinerules",
|
||||
on_rtk_ready=lambda rtk: received.append(rtk),
|
||||
)
|
||||
assert result == fake_rtk
|
||||
|
||||
inv = runner.invoke(_cmd)
|
||||
assert inv.exit_code == 0, inv.output
|
||||
assert received == [fake_rtk]
|
||||
|
||||
|
||||
def test_setup_context_tool_rtk_failure_with_not_required_returns_none(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""rtk install failure + rtk_required=False → silent fall-through, None."""
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
monkeypatch.setattr(wrap_mod, "_ensure_rtk_binary", lambda verbose=False: None)
|
||||
|
||||
on_rtk_called = False
|
||||
|
||||
def _should_not_be_called(_rtk: Path) -> None:
|
||||
nonlocal on_rtk_called
|
||||
on_rtk_called = True
|
||||
|
||||
runner = CliRunner()
|
||||
|
||||
@click.command()
|
||||
def _cmd() -> None:
|
||||
result = wrap_mod._setup_context_tool_for_agent(
|
||||
agent="cursor",
|
||||
agent_display="Cursor",
|
||||
marker_path=None,
|
||||
on_rtk_ready=_should_not_be_called,
|
||||
rtk_required=False,
|
||||
)
|
||||
assert result is None
|
||||
|
||||
inv = runner.invoke(_cmd)
|
||||
assert inv.exit_code == 0, inv.output
|
||||
assert not on_rtk_called, "on_rtk_ready should not be called when rtk install fails"
|
||||
|
||||
|
||||
def test_setup_context_tool_rtk_failure_with_required_exits_1(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""rtk install failure + rtk_required=True → SystemExit(1) with refusal message."""
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
monkeypatch.setattr(wrap_mod, "_ensure_rtk_binary", lambda verbose=False: None)
|
||||
|
||||
runner = CliRunner()
|
||||
|
||||
@click.command()
|
||||
def _cmd() -> None:
|
||||
wrap_mod._setup_context_tool_for_agent(
|
||||
agent="openhands",
|
||||
agent_display="OpenHands",
|
||||
marker_path=None,
|
||||
rtk_required=True,
|
||||
)
|
||||
|
||||
inv = runner.invoke(_cmd)
|
||||
assert inv.exit_code == 1, inv.output
|
||||
assert "rtk install failed" in inv.output
|
||||
assert "refusing to inject" in inv.output
|
||||
|
||||
|
||||
def test_setup_context_tool_keyboardinterrupt_emits_interrupted_and_exits_130(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
"""KeyboardInterrupt during setup → _emit_wrap_interrupted, SystemExit(130)."""
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
|
||||
marker = tmp_path / ".clinerules"
|
||||
marker.write_text("pre-existing")
|
||||
|
||||
def raise_kbd(verbose: bool = False) -> Path | None:
|
||||
raise KeyboardInterrupt
|
||||
|
||||
monkeypatch.setattr(wrap_mod, "_ensure_rtk_binary", raise_kbd)
|
||||
|
||||
runner = CliRunner()
|
||||
|
||||
@click.command()
|
||||
def _cmd() -> None:
|
||||
wrap_mod._setup_context_tool_for_agent(
|
||||
agent="cline",
|
||||
agent_display="Cline",
|
||||
marker_path=marker,
|
||||
)
|
||||
|
||||
inv = runner.invoke(_cmd)
|
||||
assert inv.exit_code == 130
|
||||
assert "interrupted" in inv.output.lower()
|
||||
assert "idempotent" in inv.output.lower()
|
||||
assert str(marker) in inv.output
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _run_proxy_only_watcher — must print banner, call setup callback, install
|
||||
# signal handlers, and clean up. Heavily mocked since the real watcher
|
||||
|
|
|
|||
|
|
@ -1,188 +0,0 @@
|
|||
"""Shared hint-file agent tests for `headroom wrap {cline,goose}` (PR-G1).
|
||||
|
||||
Cline and Goose are different wrap patterns (cline is proxy-only watcher,
|
||||
goose launches a child binary) but both inject the RTK guidance into a
|
||||
*hint file* at the project root — `.clinerules` and `.goosehints` —
|
||||
through the same code path (`_inject_rtk_instructions` via the shared
|
||||
`_setup_context_tool_for_agent` helper).
|
||||
|
||||
That hint-file plumbing is the same for both, so the tests covering it
|
||||
parametrize over `(agent, hint_filename)` here. Agent-specific behavior
|
||||
(goose's env-var fan-out, goose's binary discovery, cline's IDE setup
|
||||
print-out) lives in `test_wrap_goose.py` / `test_wrap_cline.py`
|
||||
respectively — those files keep only what genuinely diverges per agent.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from click.testing import CliRunner
|
||||
|
||||
from headroom.cli import wrap as wrap_mod
|
||||
from headroom.cli.main import main
|
||||
|
||||
# (subcommand, hint-file basename) — used by every test below.
|
||||
HINTFILE_AGENTS = [
|
||||
pytest.param("cline", ".clinerules", id="cline"),
|
||||
pytest.param("goose", ".goosehints", id="goose"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _enable_rtk(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# RTK is opt-in (off by default); these tests exercise the RTK-on injection path.
|
||||
monkeypatch.setenv("HEADROOM_RTK", "1")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def runner() -> CliRunner:
|
||||
return CliRunner()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("agent,hintfile", HINTFILE_AGENTS)
|
||||
def test_prepare_only_injects_rtk_into_hintfile(
|
||||
agent: str,
|
||||
hintfile: str,
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""`wrap <agent> --prepare-only` writes the RTK block to the hint file at cwd."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", agent, "--prepare-only"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
marker = tmp_path / hintfile
|
||||
assert marker.exists(), f"{hintfile} should be created"
|
||||
content = marker.read_text(encoding="utf-8")
|
||||
assert wrap_mod._RTK_MARKER in content
|
||||
assert "RTK (Rust Token Killer)" in content
|
||||
|
||||
|
||||
@pytest.mark.parametrize("agent,hintfile", HINTFILE_AGENTS)
|
||||
def test_prepare_only_idempotent_no_duplicate_block(
|
||||
agent: str,
|
||||
hintfile: str,
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Running prepare-only twice must not duplicate the RTK block in the hint file."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
runner.invoke(main, ["wrap", agent, "--prepare-only"])
|
||||
runner.invoke(main, ["wrap", agent, "--prepare-only"])
|
||||
|
||||
content = (tmp_path / hintfile).read_text(encoding="utf-8")
|
||||
assert content.count(wrap_mod._RTK_MARKER) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("agent,hintfile", HINTFILE_AGENTS)
|
||||
def test_no_context_tool_does_not_create_hintfile(
|
||||
agent: str,
|
||||
hintfile: str,
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""--no-context-tool must not create the hint file and must not invoke rtk."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary") as ensure:
|
||||
result = runner.invoke(main, ["wrap", agent, "--prepare-only", "--no-context-tool"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert not (tmp_path / hintfile).exists()
|
||||
ensure.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("agent,hintfile", HINTFILE_AGENTS)
|
||||
def test_preserves_existing_hintfile_content(
|
||||
agent: str,
|
||||
hintfile: str,
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Pre-existing hint-file content must be preserved when RTK is appended."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
marker_path = tmp_path / hintfile
|
||||
original = "# Project conventions\n\nAlways use Python 3.12.\n"
|
||||
marker_path.write_text(original, encoding="utf-8")
|
||||
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", agent, "--prepare-only"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
content = marker_path.read_text(encoding="utf-8")
|
||||
assert "Always use Python 3.12." in content
|
||||
assert wrap_mod._RTK_MARKER in content
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# M4: Ctrl-C during prelude emits a clear "interrupted, marker may be on disk"
|
||||
# message and exits non-zero (130, the conventional shell signal-130 code).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("agent,hintfile", HINTFILE_AGENTS)
|
||||
def test_keyboardinterrupt_during_prelude_emits_clear_message(
|
||||
agent: str,
|
||||
hintfile: str,
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Ctrl-C after marker injection but before proxy startup must report clearly."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
|
||||
def raise_kbd(*args, **kwargs): # noqa: ANN002, ANN003
|
||||
# Simulate the user hitting Ctrl-C right after the prelude wrote the
|
||||
# hint-file marker but before _ensure_proxy returns. We trigger via
|
||||
# _ensure_rtk_binary side-effect so the marker exists on disk.
|
||||
marker_path = tmp_path / hintfile
|
||||
marker_path.write_text(wrap_mod.RTK_INSTRUCTIONS_BLOCK, encoding="utf-8")
|
||||
raise KeyboardInterrupt
|
||||
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", side_effect=raise_kbd):
|
||||
result = runner.invoke(main, ["wrap", agent, "--prepare-only"])
|
||||
|
||||
assert result.exit_code == 130
|
||||
assert "interrupted" in result.output.lower()
|
||||
assert "idempotent" in result.output.lower()
|
||||
assert (tmp_path / hintfile).exists()
|
||||
assert hintfile in result.output
|
||||
|
||||
|
||||
@pytest.mark.parametrize("agent,hintfile", HINTFILE_AGENTS)
|
||||
def test_inject_rtk_handles_utf8_content(
|
||||
agent: str,
|
||||
hintfile: str,
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Existing hint files with non-ASCII UTF-8 content must not crash (#1126)."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
marker_path = tmp_path / hintfile
|
||||
original = "# Instructions\n\nUse “smart quotes” and an em dash — here.\n"
|
||||
marker_path.write_text(original, encoding="utf-8")
|
||||
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", agent, "--prepare-only"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
content = marker_path.read_text(encoding="utf-8")
|
||||
assert "“smart quotes”" in content
|
||||
assert wrap_mod._RTK_MARKER in content
|
||||
|
|
@ -17,7 +17,6 @@ import yaml
|
|||
from click.testing import CliRunner
|
||||
|
||||
from headroom.cli.main import main
|
||||
from headroom.cli.wrap import _inject_rtk_instructions
|
||||
from headroom.providers.omp import (
|
||||
MANAGED_MARKER,
|
||||
backup_path,
|
||||
|
|
@ -186,7 +185,7 @@ def test_build_launch_env_passes_env_through_and_emits_display(omp_home: Path) -
|
|||
|
||||
def test_wrap_omp_missing_binary_exits_with_install_hint(runner: CliRunner, omp_home: Path) -> None:
|
||||
with patch("headroom.cli.wrap.shutil.which", return_value=None):
|
||||
result = runner.invoke(main, ["wrap", "omp", "--no-rtk"])
|
||||
result = runner.invoke(main, ["wrap", "omp"])
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert "npm install -g @oh-my-pi/pi-coding-agent" in result.output
|
||||
|
|
@ -208,7 +207,7 @@ def test_wrap_omp_happy_path_injects_before_launch(runner: CliRunner, omp_home:
|
|||
patch("headroom.cli.wrap.shutil.which", return_value="omp"),
|
||||
patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool),
|
||||
):
|
||||
result = runner.invoke(main, ["wrap", "omp", "--no-rtk", "--", "-p", "fix the bug"])
|
||||
result = runner.invoke(main, ["wrap", "omp", "--", "-p", "fix the bug"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert captured["tool_label"] == "OMP"
|
||||
|
|
@ -226,33 +225,18 @@ def test_wrap_omp_happy_path_injects_before_launch(runner: CliRunner, omp_home:
|
|||
assert f"models.yml: providers.anthropic.baseUrl={base_url}" in display
|
||||
|
||||
|
||||
def test_wrap_omp_no_rtk_skips_agents_md(runner: CliRunner, omp_home: Path, tmp_path: Path) -> None:
|
||||
with (
|
||||
patch("headroom.cli.wrap.shutil.which", return_value="omp"),
|
||||
patch("headroom.cli.wrap._launch_tool"),
|
||||
):
|
||||
result = runner.invoke(main, ["wrap", "omp", "--no-rtk"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert not (tmp_path / "AGENTS.md").exists()
|
||||
|
||||
|
||||
def test_wrap_omp_rtk_injects_into_cwd_agents_md(
|
||||
runner: CliRunner, omp_home: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
def test_wrap_omp_does_not_write_agents_md(
|
||||
runner: CliRunner, omp_home: Path, tmp_path: Path
|
||||
) -> None:
|
||||
monkeypatch.setenv("HEADROOM_CONTEXT_TOOL", "rtk")
|
||||
monkeypatch.setenv("HEADROOM_RTK", "1")
|
||||
"""`wrap omp` redirects via models.yml only; it never authors AGENTS.md."""
|
||||
with (
|
||||
patch("headroom.cli.wrap.shutil.which", return_value="omp"),
|
||||
patch("headroom.cli.wrap._launch_tool"),
|
||||
patch("headroom.cli.wrap._ensure_rtk_binary", return_value=tmp_path / "rtk"),
|
||||
):
|
||||
result = runner.invoke(main, ["wrap", "omp"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
agents_md = tmp_path / "AGENTS.md"
|
||||
assert agents_md.exists()
|
||||
assert "headroom:rtk-instructions" in agents_md.read_text(encoding="utf-8")
|
||||
assert not (tmp_path / "AGENTS.md").exists()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -260,18 +244,11 @@ def test_wrap_omp_rtk_injects_into_cwd_agents_md(
|
|||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_unwrap_omp_restored_and_cleans_agents_md(
|
||||
runner: CliRunner, omp_home: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setenv("HEADROOM_RTK", "1")
|
||||
def test_unwrap_omp_restores_pristine_and_stops_proxy(runner: CliRunner, omp_home: Path) -> None:
|
||||
original = "providers:\n anthropic:\n apiKey: sk-user-secret\n"
|
||||
omp_home.write_bytes(original.encode("utf-8"))
|
||||
inject_models_override(8787, "proj")
|
||||
|
||||
agents_md = tmp_path / "AGENTS.md"
|
||||
agents_md.write_text("# My project rules\n\nBe nice.\n", encoding="utf-8")
|
||||
_inject_rtk_instructions(agents_md)
|
||||
|
||||
stopped: list[int] = []
|
||||
with patch(
|
||||
"headroom.cli.wrap._stop_local_proxy_for_unwrap",
|
||||
|
|
@ -284,11 +261,6 @@ def test_unwrap_omp_restored_and_cleans_agents_md(
|
|||
assert omp_home.read_bytes() == original.encode("utf-8")
|
||||
assert not backup_path(omp_home).exists()
|
||||
|
||||
# Only the marker-fenced rtk block is scrubbed; user content survives.
|
||||
remaining = agents_md.read_text(encoding="utf-8")
|
||||
assert "headroom:rtk-instructions" not in remaining
|
||||
assert "Be nice." in remaining
|
||||
|
||||
# A real restore (not a noop) attempts to stop the proxy on the given port.
|
||||
assert stopped == [8787]
|
||||
|
||||
|
|
|
|||
|
|
@ -9,27 +9,16 @@ from urllib.parse import quote
|
|||
import pytest
|
||||
from click.testing import CliRunner
|
||||
|
||||
import headroom.cli.wrap as wrap_cli
|
||||
from headroom.cli.main import main
|
||||
|
||||
OPENCLAUDE_BINARY = "openclaude"
|
||||
OPENCLAUDE_INSTRUCTIONS_FILE = "CONVENTIONS.md"
|
||||
OPENCLAUDE_MODEL_ARG = "gpt-4o"
|
||||
RTK_BINARY = "rtk"
|
||||
UTF8_ENCODING = "utf-8"
|
||||
WINDOWS_DEFAULT_TEXT_ENCODING = "cp1252"
|
||||
|
||||
|
||||
def _expected_project_prefix() -> str:
|
||||
return f"/p/{quote(Path.cwd().name, safe='')}"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _enable_rtk(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# RTK is opt-in (off by default); these tests exercise the RTK-on injection path.
|
||||
monkeypatch.setenv("HEADROOM_RTK", "1")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def runner() -> CliRunner:
|
||||
return CliRunner()
|
||||
|
|
@ -48,7 +37,7 @@ def test_wrap_openclaude_routes_proxy_envs(
|
|||
with patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "openclaude", "--no-rtk", "--", "--model", OPENCLAUDE_MODEL_ARG],
|
||||
["wrap", "openclaude", "--", "--model", OPENCLAUDE_MODEL_ARG],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
|
|
@ -62,56 +51,11 @@ def test_wrap_openclaude_routes_proxy_envs(
|
|||
assert captured["args"] == ("--model", OPENCLAUDE_MODEL_ARG)
|
||||
|
||||
|
||||
def test_wrap_openclaude_default_rtk_injects_instructions(
|
||||
runner: CliRunner, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.chdir(tmp_path)
|
||||
original_write_text = Path.write_text
|
||||
|
||||
def write_text_with_windows_default(
|
||||
self: Path,
|
||||
data: str,
|
||||
encoding: str | None = None,
|
||||
errors: str | None = None,
|
||||
newline: str | None = None,
|
||||
) -> int:
|
||||
if encoding is None:
|
||||
encoded = data.encode(WINDOWS_DEFAULT_TEXT_ENCODING, errors=errors or "strict")
|
||||
self.write_bytes(encoded)
|
||||
return len(data)
|
||||
return original_write_text(
|
||||
self,
|
||||
data,
|
||||
encoding=encoding,
|
||||
errors=errors,
|
||||
newline=newline,
|
||||
)
|
||||
|
||||
def fake_launch_tool(**_kwargs): # noqa: ANN003
|
||||
return None
|
||||
|
||||
with patch("headroom.cli.wrap.shutil.which", return_value=OPENCLAUDE_BINARY):
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=tmp_path / RTK_BINARY):
|
||||
with patch.object(
|
||||
Path,
|
||||
"write_text",
|
||||
autospec=True,
|
||||
side_effect=write_text_with_windows_default,
|
||||
):
|
||||
with patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool):
|
||||
result = runner.invoke(main, ["wrap", "openclaude"])
|
||||
|
||||
instructions = tmp_path / OPENCLAUDE_INSTRUCTIONS_FILE
|
||||
assert result.exit_code == 0, result.output
|
||||
assert instructions.exists()
|
||||
assert wrap_cli._RTK_MARKER in instructions.read_text(encoding=UTF8_ENCODING)
|
||||
|
||||
|
||||
def test_wrap_openclaude_missing_binary_errors(
|
||||
runner: CliRunner, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.chdir(tmp_path)
|
||||
with patch("headroom.cli.wrap.shutil.which", return_value=None):
|
||||
result = runner.invoke(main, ["wrap", "openclaude", "--no-rtk"])
|
||||
result = runner.invoke(main, ["wrap", "openclaude"])
|
||||
assert result.exit_code == 1
|
||||
assert "openclaude" in result.output.lower()
|
||||
|
|
|
|||
|
|
@ -15,9 +15,9 @@ from headroom.copilot_auth import CopilotSubscriptionTokenResolution
|
|||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _enable_rtk(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# RTK is opt-in (off by default); these tests exercise the RTK-on injection path.
|
||||
monkeypatch.setenv("HEADROOM_RTK", "1")
|
||||
def _no_retired_context_tool_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""A developer's exported HEADROOM_CONTEXT_TOOL would abort every wrap below."""
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -120,7 +120,6 @@ def test_wrap_opencode_copilot_subscription_normalizes_enterprise_host_and_hando
|
|||
"wrap",
|
||||
"opencode",
|
||||
"--copilot-subscription",
|
||||
"--no-rtk",
|
||||
"--no-mcp",
|
||||
"--no-serena",
|
||||
],
|
||||
|
|
@ -173,7 +172,7 @@ def test_wrap_opencode_copilot_subscription_rejects_incompatible_modes(
|
|||
with patch.object(wrap_mod, "_ensure_proxy", side_effect=AssertionError("proxy launched")):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "opencode", "--copilot-subscription", "--no-rtk", "--no-mcp", *extra_args],
|
||||
["wrap", "opencode", "--copilot-subscription", "--no-mcp", *extra_args],
|
||||
)
|
||||
assert result.exit_code == 1
|
||||
assert message in result.output
|
||||
|
|
@ -192,7 +191,7 @@ def test_wrap_opencode_copilot_subscription_rejects_headroom_backend_env(
|
|||
with patch.object(wrap_mod, "_ensure_proxy", side_effect=AssertionError("proxy launched")):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "opencode", "--copilot-subscription", "--no-rtk", "--no-mcp"],
|
||||
["wrap", "opencode", "--copilot-subscription", "--no-mcp"],
|
||||
)
|
||||
assert result.exit_code == 1
|
||||
assert "translated backends" in result.output
|
||||
|
|
@ -216,7 +215,7 @@ def test_wrap_opencode_copilot_subscription_requires_login_before_launch(
|
|||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "opencode", "--copilot-subscription", "--no-rtk", "--no-mcp"],
|
||||
["wrap", "opencode", "--copilot-subscription", "--no-mcp"],
|
||||
)
|
||||
assert result.exit_code == 1
|
||||
assert "headroom copilot-auth login" in result.output
|
||||
|
|
@ -272,7 +271,6 @@ def test_wrap_opencode_copilot_subscription_cleans_up_proxy_on_config_failure(
|
|||
"wrap",
|
||||
"opencode",
|
||||
"--copilot-subscription",
|
||||
"--no-rtk",
|
||||
"--no-mcp",
|
||||
"--no-serena",
|
||||
],
|
||||
|
|
@ -304,11 +302,10 @@ def test_wrap_opencode_sets_config_content_env(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "opencode", "--port", "9000", "--no-mcp", "--", "--model", "gpt-4o"],
|
||||
)
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "opencode", "--port", "9000", "--no-mcp", "--", "--model", "gpt-4o"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
env = captured["env"]
|
||||
|
|
@ -342,8 +339,7 @@ def test_wrap_opencode_does_not_add_base_url_env_vars(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
env = captured["env"]
|
||||
|
|
@ -362,8 +358,7 @@ def test_wrap_opencode_missing_binary_errors_clearly(
|
|||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value=None):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "opencode"])
|
||||
result = runner.invoke(main, ["wrap", "opencode"])
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert "'opencode' not found in PATH" in result.output
|
||||
|
|
@ -380,8 +375,7 @@ def test_wrap_opencode_prepare_only_injects_config(
|
|||
_set_test_home(monkeypatch, tmp_path)
|
||||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--prepare-only"])
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--prepare-only"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
config_file = tmp_path / ".config" / "opencode" / "opencode.json"
|
||||
|
|
@ -400,8 +394,7 @@ def test_wrap_opencode_prepare_only_registers_serena_with_agent_context(
|
|||
_set_test_home(monkeypatch, tmp_path)
|
||||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--prepare-only"])
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--prepare-only"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
config_file = tmp_path / ".config" / "opencode" / "opencode.json"
|
||||
|
|
@ -427,8 +420,7 @@ def test_wrap_opencode_no_mcp_skips_mcp_injection(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
env = captured["env"]
|
||||
|
|
@ -456,8 +448,7 @@ def test_wrap_opencode_injects_mcp_by_default(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000"])
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
env = captured["env"]
|
||||
|
|
@ -471,120 +462,6 @@ def test_wrap_opencode_injects_mcp_by_default(
|
|||
}
|
||||
|
||||
|
||||
def test_wrap_opencode_injects_rtk_into_agents_md(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""RTK instructions are injected into global and project AGENTS.md."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=SystemExit(0)):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
global_agents = tmp_path / ".config" / "opencode" / "AGENTS.md"
|
||||
project_agents = tmp_path / "AGENTS.md"
|
||||
assert global_agents.exists(), "Global AGENTS.md should be created"
|
||||
assert project_agents.exists(), "Project AGENTS.md should be created"
|
||||
assert wrap_mod._RTK_MARKER in global_agents.read_text(encoding="utf-8")
|
||||
assert wrap_mod._RTK_MARKER in project_agents.read_text(encoding="utf-8")
|
||||
|
||||
|
||||
def test_unwrap_opencode_removes_rtk_from_agents_md(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""unwrap opencode removes the rtk block that wrap opencode injected into both
|
||||
the project and global AGENTS.md — mirroring unwrap_codex / unwrap_copilot."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=SystemExit(0)):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
|
||||
global_agents = tmp_path / ".config" / "opencode" / "AGENTS.md"
|
||||
project_agents = tmp_path / "AGENTS.md"
|
||||
assert wrap_mod._RTK_MARKER in global_agents.read_text(encoding="utf-8")
|
||||
assert wrap_mod._RTK_MARKER in project_agents.read_text(encoding="utf-8")
|
||||
|
||||
with patch.object(wrap_mod, "_stop_local_proxy_for_unwrap", return_value="stopped"):
|
||||
result = runner.invoke(main, ["unwrap", "opencode"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
|
||||
# Both rtk blocks are gone after unwrap (previously left behind). A file that
|
||||
# held only the rtk block is removed entirely by _remove_rtk_instructions, so
|
||||
# treat a missing file as "block gone".
|
||||
def _rtk_absent(path: Path) -> bool:
|
||||
return not path.exists() or wrap_mod._RTK_MARKER not in path.read_text(encoding="utf-8")
|
||||
|
||||
assert _rtk_absent(global_agents)
|
||||
assert _rtk_absent(project_agents)
|
||||
|
||||
|
||||
def test_wrap_opencode_no_project_rtk_only_skips_project_agents_md(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
project_agents = tmp_path / "AGENTS.md"
|
||||
project_agents.write_text("# Team instructions\n", encoding="utf-8")
|
||||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=SystemExit(0)):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"wrap",
|
||||
"opencode",
|
||||
"--no-project-rtk",
|
||||
"--no-proxy",
|
||||
"--port",
|
||||
"9000",
|
||||
"--no-mcp",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert project_agents.read_text(encoding="utf-8") == "# Team instructions\n"
|
||||
global_agents = tmp_path / ".config" / "opencode" / "AGENTS.md"
|
||||
assert wrap_mod._RTK_MARKER in global_agents.read_text(encoding="utf-8")
|
||||
|
||||
|
||||
def test_wrap_opencode_idempotent_no_duplicate_block(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Running wrap twice must not duplicate the RTK block in AGENTS.md."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=SystemExit(0)):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
|
||||
project_agents = tmp_path / "AGENTS.md"
|
||||
content = project_agents.read_text(encoding="utf-8")
|
||||
assert content.count(wrap_mod._RTK_MARKER) == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unwrap opencode
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -691,8 +568,7 @@ def test_wrap_opencode_preserves_existing_user_providers(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=SystemExit(0)):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
config = json.loads(config_file.read_text(encoding="utf-8"))
|
||||
|
|
@ -712,9 +588,8 @@ def test_wrap_opencode_port_change_updates_existing_config(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=SystemExit(0)):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
runner.invoke(main, ["wrap", "opencode", "--port", "9001", "--no-mcp"])
|
||||
runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
runner.invoke(main, ["wrap", "opencode", "--port", "9001", "--no-mcp"])
|
||||
|
||||
config_file = tmp_path / ".config" / "opencode" / "opencode.json"
|
||||
config = json.loads(config_file.read_text(encoding="utf-8"))
|
||||
|
|
@ -739,8 +614,7 @@ def test_wrap_opencode_handles_malformed_config_file(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=SystemExit(0)):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert backup_file.exists(), "backup must be created before overwriting"
|
||||
|
|
@ -768,8 +642,7 @@ def test_wrap_opencode_handles_empty_config_file(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=SystemExit(0)):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
config = json.loads(config_file.read_text(encoding="utf-8"))
|
||||
|
|
@ -791,20 +664,19 @@ def test_wrap_opencode_handles_config_dir_missing(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=SystemExit(0)):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert config_dir.exists()
|
||||
assert (config_dir / "opencode.json").exists()
|
||||
|
||||
|
||||
def test_wrap_opencode_rtk_preserves_existing_agents_md(
|
||||
def test_wrap_opencode_leaves_agents_md_untouched(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""RTK injection appends to AGENTS.md without removing existing content."""
|
||||
"""`wrap opencode` never rewrites an existing AGENTS.md."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
|
|
@ -814,39 +686,11 @@ def test_wrap_opencode_rtk_preserves_existing_agents_md(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=SystemExit(0)):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
content = (tmp_path / "AGENTS.md").read_text(encoding="utf-8")
|
||||
assert existing_content in content
|
||||
assert wrap_mod._RTK_MARKER in content
|
||||
|
||||
|
||||
def test_wrap_opencode_no_rtk_leaves_agents_md_untouched(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""`--no-rtk` flag leaves existing AGENTS.md untouched."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
|
||||
existing_content = "# My custom rules\nUse spaces, not tabs."
|
||||
(tmp_path / "AGENTS.md").write_text(existing_content)
|
||||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=SystemExit(0)):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(
|
||||
main, ["wrap", "opencode", "--port", "9000", "--no-rtk", "--no-mcp"]
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
content = (tmp_path / "AGENTS.md").read_text(encoding="utf-8")
|
||||
assert content == existing_content, "--no-rtk modified AGENTS.md"
|
||||
assert wrap_mod._RTK_MARKER not in content
|
||||
assert content == existing_content, "wrap opencode modified AGENTS.md"
|
||||
|
||||
|
||||
def test_wrap_opencode_respects_opencode_config_env(
|
||||
|
|
@ -864,8 +708,7 @@ def test_wrap_opencode_respects_opencode_config_env(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=SystemExit(0)):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert custom_config.exists()
|
||||
|
|
@ -895,8 +738,7 @@ def test_wrap_opencode_headroom_project_from_cwd(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
env = captured["env"]
|
||||
|
|
@ -921,8 +763,7 @@ def test_wrap_opencode_respects_existing_headroom_project(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
env = captured["env"]
|
||||
|
|
@ -945,8 +786,7 @@ def test_wrap_opencode_config_merges_existing_model(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=SystemExit(0)):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
config = json.loads(config_file.read_text(encoding="utf-8"))
|
||||
|
|
@ -1038,8 +878,7 @@ def test_wrap_unwrap_rewrap_is_idempotent(
|
|||
# First wrap
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=SystemExit(0)):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
|
||||
# Unwrap
|
||||
with patch.object(wrap_mod, "_stop_local_proxy_for_unwrap", return_value="stopped"):
|
||||
|
|
@ -1053,8 +892,7 @@ def test_wrap_unwrap_rewrap_is_idempotent(
|
|||
# Re-wrap
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=SystemExit(0)):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
runner.invoke(main, ["wrap", "opencode", "--port", "9001", "--no-mcp"])
|
||||
runner.invoke(main, ["wrap", "opencode", "--port", "9001", "--no-mcp"])
|
||||
|
||||
# After re-wrap, headroom should be back, model unchanged
|
||||
after_rewrap = json.loads(config_file.read_text(encoding="utf-8"))
|
||||
|
|
@ -1103,8 +941,7 @@ def test_wrap_opencode_no_arguments_is_valid(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--no-mcp"])
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--no-mcp"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert captured["tool_label"] == "OPENCODE"
|
||||
|
|
@ -1123,10 +960,9 @@ def test_wrap_opencode_with_memory_flag(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=SystemExit(0)):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(
|
||||
main, ["wrap", "opencode", "--port", "9000", "--memory", "--no-mcp"]
|
||||
)
|
||||
result = runner.invoke(
|
||||
main, ["wrap", "opencode", "--port", "9000", "--memory", "--no-mcp"]
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
|
||||
|
|
@ -1143,21 +979,20 @@ def test_wrap_opencode_with_backend_and_anyllm_provider(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=SystemExit(0)):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"wrap",
|
||||
"opencode",
|
||||
"--port",
|
||||
"9000",
|
||||
"--backend",
|
||||
"anyllm",
|
||||
"--anyllm-provider",
|
||||
"groq",
|
||||
"--no-mcp",
|
||||
],
|
||||
)
|
||||
result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"wrap",
|
||||
"opencode",
|
||||
"--port",
|
||||
"9000",
|
||||
"--backend",
|
||||
"anyllm",
|
||||
"--anyllm-provider",
|
||||
"groq",
|
||||
"--no-mcp",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
|
||||
|
|
@ -1174,10 +1009,9 @@ def test_wrap_opencode_with_no_proxy(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=SystemExit(0)):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(
|
||||
main, ["wrap", "opencode", "--port", "9000", "--no-proxy", "--no-mcp"]
|
||||
)
|
||||
result = runner.invoke(
|
||||
main, ["wrap", "opencode", "--port", "9000", "--no-proxy", "--no-mcp"]
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
|
||||
|
|
@ -1194,10 +1028,9 @@ def test_wrap_opencode_with_verbose_flag(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=SystemExit(0)):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(
|
||||
main, ["wrap", "opencode", "--port", "9000", "--verbose", "--no-mcp"]
|
||||
)
|
||||
result = runner.invoke(
|
||||
main, ["wrap", "opencode", "--port", "9000", "--verbose", "--no-mcp"]
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
|
||||
|
|
@ -1207,7 +1040,7 @@ def test_wrap_opencode_respects_opencode_home_env(
|
|||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""OPENCODE_HOME env var controls where AGENTS.md is written."""
|
||||
"""OPENCODE_HOME env var controls where opencode.json is written."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
custom_home = str(tmp_path / "custom-opencode-home")
|
||||
|
|
@ -1216,12 +1049,10 @@ def test_wrap_opencode_respects_opencode_home_env(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="opencode"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=SystemExit(0)):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
result = runner.invoke(main, ["wrap", "opencode", "--port", "9000", "--no-mcp"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
agents_md = Path(custom_home) / "AGENTS.md"
|
||||
assert agents_md.exists()
|
||||
assert (Path(custom_home) / "opencode.json").exists()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -17,6 +17,25 @@ def runner() -> CliRunner:
|
|||
return CliRunner()
|
||||
|
||||
|
||||
def test_wrap_openhands_prepare_only_succeeds_unpatched(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""`wrap openhands --prepare-only` must succeed with nothing patched.
|
||||
|
||||
The subcommand used to hard-fail by default because a CLI context tool was
|
||||
a required dependency of the prelude. Nothing is fetched or written now, so
|
||||
the bare invocation has to exit 0 on its own.
|
||||
"""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
|
||||
result = runner.invoke(main, ["wrap", "openhands", "--prepare-only"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
|
||||
|
||||
def test_wrap_openhands_sets_provider_envs(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
|
|
@ -33,10 +52,9 @@ def test_wrap_openhands_sets_provider_envs(
|
|||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="openhands"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(
|
||||
main, ["wrap", "openhands", "--port", "9000", "--", "--task", "demo"]
|
||||
)
|
||||
result = runner.invoke(
|
||||
main, ["wrap", "openhands", "--port", "9000", "--", "--task", "demo"]
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
env = captured["env"]
|
||||
|
|
@ -50,88 +68,6 @@ def test_wrap_openhands_sets_provider_envs(
|
|||
assert captured["args"] == ("--task", "demo")
|
||||
|
||||
|
||||
def test_wrap_openhands_injects_rtk_via_env_var(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""OPENHANDS_INSTRUCTIONS env var must contain the RTK block at launch."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
monkeypatch.delenv("OPENHANDS_INSTRUCTIONS", raising=False)
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_launch_tool(**kwargs): # noqa: ANN003
|
||||
captured.update(kwargs)
|
||||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="openhands"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "openhands"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
env = captured["env"]
|
||||
assert isinstance(env, dict)
|
||||
instructions = env.get("OPENHANDS_INSTRUCTIONS", "")
|
||||
assert wrap_mod._RTK_MARKER in instructions
|
||||
assert "RTK (Rust Token Killer)" in instructions
|
||||
|
||||
|
||||
def test_wrap_openhands_preserves_existing_openhands_instructions(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Pre-existing OPENHANDS_INSTRUCTIONS env content is preserved, rtk is appended."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
monkeypatch.setenv("OPENHANDS_INSTRUCTIONS", "Prefer typed Python.")
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_launch_tool(**kwargs): # noqa: ANN003
|
||||
captured.update(kwargs)
|
||||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="openhands"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "openhands"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
env = captured["env"]
|
||||
instructions = env.get("OPENHANDS_INSTRUCTIONS", "")
|
||||
assert "Prefer typed Python." in instructions
|
||||
assert wrap_mod._RTK_MARKER in instructions
|
||||
|
||||
|
||||
def test_wrap_openhands_idempotent_already_injected(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""If OPENHANDS_INSTRUCTIONS already contains the marker, do not re-append."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
pre_existing = "Prefer typed Python.\n\n" + wrap_mod.RTK_INSTRUCTIONS_BLOCK
|
||||
monkeypatch.setenv("OPENHANDS_INSTRUCTIONS", pre_existing)
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_launch_tool(**kwargs): # noqa: ANN003
|
||||
captured.update(kwargs)
|
||||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="openhands"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "openhands"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
env = captured["env"]
|
||||
instructions = env.get("OPENHANDS_INSTRUCTIONS", "")
|
||||
assert instructions.count(wrap_mod._RTK_MARKER) == 1
|
||||
|
||||
|
||||
def test_wrap_openhands_missing_binary_errors_clearly(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
|
|
@ -142,117 +78,7 @@ def test_wrap_openhands_missing_binary_errors_clearly(
|
|||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value=None):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "openhands"])
|
||||
result = runner.invoke(main, ["wrap", "openhands"])
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert "'openhands' not found in PATH" in result.output
|
||||
|
||||
|
||||
def test_wrap_openhands_no_context_tool_does_not_inject(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""--no-context-tool must skip OPENHANDS_INSTRUCTIONS injection."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("OPENHANDS_INSTRUCTIONS", raising=False)
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_launch_tool(**kwargs): # noqa: ANN003
|
||||
captured.update(kwargs)
|
||||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="openhands"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary") as ensure:
|
||||
result = runner.invoke(main, ["wrap", "openhands", "--no-context-tool"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
ensure.assert_not_called()
|
||||
env = captured["env"]
|
||||
assert isinstance(env, dict)
|
||||
assert "OPENHANDS_INSTRUCTIONS" not in env or env["OPENHANDS_INSTRUCTIONS"] == ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# M3: rtk install failure must fail loudly — no silent fallback to env
|
||||
# injection without rtk on disk.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_wrap_openhands_rtk_install_failure_aborts_loudly(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""If rtk install fails, command must exit non-zero with a clear error."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
monkeypatch.delenv("OPENHANDS_INSTRUCTIONS", raising=False)
|
||||
|
||||
launch_called: list[bool] = []
|
||||
|
||||
def fake_launch_tool(**kwargs): # noqa: ANN003
|
||||
launch_called.append(True)
|
||||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="openhands"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=None):
|
||||
result = runner.invoke(main, ["wrap", "openhands"])
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert "rtk install failed" in result.output
|
||||
assert "--no-context-tool" in result.output
|
||||
# _launch_tool must NOT have been invoked when rtk install fails.
|
||||
assert launch_called == []
|
||||
|
||||
|
||||
def test_wrap_openhands_rtk_install_failure_with_no_context_tool_still_launches(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""--no-context-tool bypasses rtk entirely — should still launch."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("OPENHANDS_INSTRUCTIONS", raising=False)
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_launch_tool(**kwargs): # noqa: ANN003
|
||||
captured.update(kwargs)
|
||||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="openhands"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool):
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=None) as ensure:
|
||||
result = runner.invoke(main, ["wrap", "openhands", "--no-context-tool"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
# rtk should never have been queried.
|
||||
ensure.assert_not_called()
|
||||
env = captured["env"]
|
||||
assert "OPENHANDS_INSTRUCTIONS" not in env or env["OPENHANDS_INSTRUCTIONS"] == ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# M4: Ctrl-C during prelude emits a clear "no on-disk changes" message.
|
||||
# openhands never writes to disk (env-var injection only).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_wrap_openhands_keyboardinterrupt_during_prelude_emits_clear_message(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Ctrl-C during the prelude must signal cleanly with no on-disk artifact."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
monkeypatch.delenv("OPENHANDS_INSTRUCTIONS", raising=False)
|
||||
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", side_effect=KeyboardInterrupt):
|
||||
result = runner.invoke(main, ["wrap", "openhands"])
|
||||
|
||||
assert result.exit_code == 130
|
||||
assert "interrupted" in result.output.lower()
|
||||
assert "idempotent" in result.output.lower()
|
||||
|
|
|
|||
|
|
@ -1,113 +0,0 @@
|
|||
"""Phase G PR-G3 remediation (C4) — wrap-CLI RTK metrics primitive.
|
||||
|
||||
The Rust proxy previously held a dead `wrap_rtk_invocations_total`
|
||||
counter. C4 remediation moved it Python-side because the wrap CLI
|
||||
(headroom.cli.wrap) is where RTK invocations are actually counted.
|
||||
These tests cover the counter primitives in isolation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.cli.wrap_rtk_metrics import (
|
||||
record_rtk_invocation,
|
||||
reset_rtk_invocations,
|
||||
rtk_invocation_counts,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_between_tests():
|
||||
"""Reset the module-level counter map between tests so each
|
||||
test owns a clean slate."""
|
||||
reset_rtk_invocations()
|
||||
yield
|
||||
reset_rtk_invocations()
|
||||
|
||||
|
||||
def test_record_increments_default_delta_one():
|
||||
record_rtk_invocation("git")
|
||||
counts = rtk_invocation_counts()
|
||||
assert counts == {"git": 1}
|
||||
|
||||
|
||||
def test_record_accumulates_per_tool():
|
||||
record_rtk_invocation("git")
|
||||
record_rtk_invocation("git")
|
||||
record_rtk_invocation("ls")
|
||||
record_rtk_invocation("cargo")
|
||||
record_rtk_invocation("cargo")
|
||||
record_rtk_invocation("cargo")
|
||||
counts = rtk_invocation_counts()
|
||||
assert counts == {"git": 2, "ls": 1, "cargo": 3}
|
||||
|
||||
|
||||
def test_record_with_explicit_delta():
|
||||
record_rtk_invocation("git", delta=5)
|
||||
record_rtk_invocation("git", delta=2)
|
||||
counts = rtk_invocation_counts()
|
||||
assert counts == {"git": 7}
|
||||
|
||||
|
||||
def test_record_zero_delta_is_noop_record():
|
||||
# delta=0 is legal — caller may want to "touch" the counter to
|
||||
# ensure the key exists before later increments.
|
||||
record_rtk_invocation("git", delta=0)
|
||||
counts = rtk_invocation_counts()
|
||||
assert counts == {"git": 0}
|
||||
|
||||
|
||||
def test_record_rejects_negative_delta():
|
||||
with pytest.raises(ValueError, match="must be non-negative"):
|
||||
record_rtk_invocation("git", delta=-1)
|
||||
|
||||
|
||||
def test_record_rejects_non_string_tool():
|
||||
with pytest.raises(TypeError, match="tool must be a str"):
|
||||
record_rtk_invocation(123, delta=1) # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_record_rejects_non_int_delta():
|
||||
with pytest.raises(TypeError, match="delta must be an int"):
|
||||
record_rtk_invocation("git", delta="1") # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_counts_returns_snapshot_not_view():
|
||||
# The returned mapping must be a plain dict copy, not the
|
||||
# internal defaultdict — otherwise callers could pollute the
|
||||
# counter map by reading absent keys.
|
||||
record_rtk_invocation("git")
|
||||
counts = rtk_invocation_counts()
|
||||
# Reading a key that's not present must not add it to the
|
||||
# internal map.
|
||||
_ = counts.get("nonexistent_tool", 0)
|
||||
counts2 = rtk_invocation_counts()
|
||||
assert "nonexistent_tool" not in counts2
|
||||
|
||||
|
||||
def test_thread_safe_concurrent_increments():
|
||||
# 10 threads each bumping `git` 100 times: final count must be
|
||||
# exactly 1000. The threading.Lock guards the dict update so
|
||||
# races are impossible.
|
||||
def worker():
|
||||
for _ in range(100):
|
||||
record_rtk_invocation("git")
|
||||
|
||||
threads = [threading.Thread(target=worker) for _ in range(10)]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
counts = rtk_invocation_counts()
|
||||
assert counts == {"git": 1000}
|
||||
|
||||
|
||||
def test_reset_clears_counts():
|
||||
record_rtk_invocation("git", delta=42)
|
||||
record_rtk_invocation("ls", delta=7)
|
||||
assert rtk_invocation_counts() != {}
|
||||
reset_rtk_invocations()
|
||||
assert rtk_invocation_counts() == {}
|
||||
|
|
@ -1,127 +0,0 @@
|
|||
"""Tests for ``_ensure_rtk_on_path``.
|
||||
|
||||
``rtk init --global --auto-patch`` writes ``~/.claude/hooks/rtk-rewrite.sh``,
|
||||
and ``rtk rewrite`` emits a bare ``rtk`` token at runtime that the hook feeds
|
||||
back to the shell — so bare ``rtk`` must resolve on PATH. Since
|
||||
``~/.headroom/bin`` is not on PATH by default, that lookup fails and token
|
||||
compression never runs (issue #487).
|
||||
|
||||
The earlier fix rewrote the generated hook to hard-code rtk's absolute path,
|
||||
but that mutates the hook after ``rtk init`` bakes in its expected SHA-256, so
|
||||
rtk's integrity guard rejects it (issue #1631). ``_ensure_rtk_on_path`` instead
|
||||
leaves the canonical hook untouched and links the managed binary into a PATH
|
||||
directory so bare ``rtk`` resolves.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.cli import wrap
|
||||
from headroom.cli.wrap import _ensure_rtk_on_path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def rtk_binary(tmp_path: Path) -> Path:
|
||||
managed = tmp_path / ".headroom" / "bin" / "rtk"
|
||||
managed.parent.mkdir(parents=True)
|
||||
managed.write_text("#!/bin/sh\n")
|
||||
managed.chmod(0o755)
|
||||
return managed
|
||||
|
||||
|
||||
def test_noop_when_rtk_already_on_path(rtk_binary: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(wrap.sys, "platform", "linux")
|
||||
monkeypatch.setattr(wrap.shutil, "which", lambda _cmd: "/usr/bin/rtk")
|
||||
|
||||
assert _ensure_rtk_on_path(rtk_binary, path_dirs=["/usr/bin"]) is None
|
||||
|
||||
|
||||
def test_noop_on_windows(rtk_binary: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(wrap.sys, "platform", "win32")
|
||||
|
||||
assert _ensure_rtk_on_path(rtk_binary, path_dirs=["C:\\bin"]) is None
|
||||
|
||||
|
||||
def test_links_into_path_dir_when_missing(
|
||||
rtk_binary: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(wrap.sys, "platform", "linux")
|
||||
monkeypatch.setattr(wrap.shutil, "which", lambda _cmd: None)
|
||||
bindir = tmp_path / "path-bin"
|
||||
bindir.mkdir()
|
||||
|
||||
link = _ensure_rtk_on_path(rtk_binary, path_dirs=[str(bindir)])
|
||||
|
||||
assert link == bindir / "rtk"
|
||||
assert link.is_symlink()
|
||||
assert link.resolve() == rtk_binary.resolve()
|
||||
|
||||
|
||||
def test_prefers_local_bin(
|
||||
rtk_binary: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(wrap.sys, "platform", "linux")
|
||||
monkeypatch.setattr(wrap.shutil, "which", lambda _cmd: None)
|
||||
home = tmp_path / "home"
|
||||
monkeypatch.setattr(wrap.Path, "home", classmethod(lambda _cls: home))
|
||||
other = tmp_path / "other-bin"
|
||||
other.mkdir()
|
||||
local_bin = home / ".local" / "bin"
|
||||
|
||||
# ~/.local/bin does not exist yet but is on PATH — it is created on demand
|
||||
# and preferred over the other writable dir.
|
||||
link = _ensure_rtk_on_path(rtk_binary, path_dirs=[str(other), str(local_bin)])
|
||||
|
||||
assert link == local_bin / "rtk"
|
||||
assert link.is_symlink()
|
||||
|
||||
|
||||
def test_idempotent_second_run(
|
||||
rtk_binary: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(wrap.sys, "platform", "linux")
|
||||
monkeypatch.setattr(wrap.shutil, "which", lambda _cmd: None)
|
||||
bindir = tmp_path / "path-bin"
|
||||
bindir.mkdir()
|
||||
|
||||
first = _ensure_rtk_on_path(rtk_binary, path_dirs=[str(bindir)])
|
||||
second = _ensure_rtk_on_path(rtk_binary, path_dirs=[str(bindir)])
|
||||
|
||||
assert first == second == bindir / "rtk"
|
||||
assert second.resolve() == rtk_binary.resolve()
|
||||
|
||||
|
||||
def test_does_not_clobber_existing_file(
|
||||
rtk_binary: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(wrap.sys, "platform", "linux")
|
||||
monkeypatch.setattr(wrap.shutil, "which", lambda _cmd: None)
|
||||
occupied = tmp_path / "occupied-bin"
|
||||
occupied.mkdir()
|
||||
foreign = occupied / "rtk"
|
||||
foreign.write_text("#!/bin/sh\n# a different rtk\n")
|
||||
fallback = tmp_path / "fallback-bin"
|
||||
fallback.mkdir()
|
||||
|
||||
link = _ensure_rtk_on_path(rtk_binary, path_dirs=[str(occupied), str(fallback)])
|
||||
|
||||
# The real file is left untouched; the link lands in the next writable dir.
|
||||
assert foreign.read_text() == "#!/bin/sh\n# a different rtk\n"
|
||||
assert not foreign.is_symlink()
|
||||
assert link == fallback / "rtk"
|
||||
assert link.is_symlink()
|
||||
|
||||
|
||||
def test_noop_when_no_writable_path_dir(
|
||||
rtk_binary: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(wrap.sys, "platform", "linux")
|
||||
monkeypatch.setattr(wrap.shutil, "which", lambda _cmd: None)
|
||||
home = tmp_path / "home"
|
||||
monkeypatch.setattr(wrap.Path, "home", classmethod(lambda _cls: home))
|
||||
|
||||
# Only a non-existent, non-preferred dir on PATH — nothing to link into.
|
||||
assert _ensure_rtk_on_path(rtk_binary, path_dirs=[str(tmp_path / "ghost")]) is None
|
||||
|
|
@ -23,7 +23,7 @@ def _opt_in(monkeypatch: pytest.MonkeyPatch) -> None:
|
|||
"""Enable the opt-in gate so injection actually writes.
|
||||
|
||||
Instruction injection rewrites the user's CLAUDE.md/AGENTS.md, so it is
|
||||
off by default (mirrors RTK). Tests that exercise the write path must opt in via
|
||||
off by default. Tests that exercise the write path must opt in via
|
||||
``HEADROOM_SERENA_INSTRUCTIONS``.
|
||||
"""
|
||||
monkeypatch.setenv("HEADROOM_SERENA_INSTRUCTIONS", "1")
|
||||
|
|
|
|||
|
|
@ -243,39 +243,3 @@ def test_wrap_vibe_providers_json_structure(
|
|||
assert "api_base" in providers[0]
|
||||
assert providers[0]["browser_auth_base_url"] == "https://console.mistral.ai"
|
||||
assert providers[0]["browser_auth_api_base_url"] == "https://console.mistral.ai/api"
|
||||
|
||||
|
||||
def test_wrap_vibe_no_context_tool(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""--no-context-tool and --no-rtk flags are accepted and not passed to vibe."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
def fake_launch_tool(**kwargs: Any) -> None: # noqa: ANN003
|
||||
captured.update(kwargs)
|
||||
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="vibe"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool):
|
||||
with patch.object(wrap_mod, "_project_name_from_cwd", return_value=None):
|
||||
# Test --no-context-tool
|
||||
result = runner.invoke(main, ["wrap", "vibe", "--no-context-tool", "--", "test"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert captured["args"] == ("test",)
|
||||
assert "--no-context-tool" not in captured["args"]
|
||||
|
||||
captured.clear()
|
||||
with patch.object(wrap_mod.shutil, "which", return_value="vibe"):
|
||||
with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool):
|
||||
with patch.object(wrap_mod, "_project_name_from_cwd", return_value=None):
|
||||
# Test --no-rtk
|
||||
result = runner.invoke(main, ["wrap", "vibe", "--no-rtk", "--", "test"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert captured["args"] == ("test",)
|
||||
assert "--no-rtk" not in captured["args"]
|
||||
|
|
|
|||
|
|
@ -2,8 +2,8 @@
|
|||
|
||||
ZCode is a desktop Electron app (zcode.z.ai) with no CLI binary. The wrap
|
||||
command follows the Pattern-B (proxy-only watcher) approach: it starts the
|
||||
proxy, injects RTK guidance into AGENTS.md at the project root, and prints
|
||||
the ZCode settings the user should configure in the app's settings UI.
|
||||
proxy and prints the ZCode settings the user should configure in the app's
|
||||
settings UI.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -18,96 +18,11 @@ from headroom.cli import wrap as wrap_mod
|
|||
from headroom.cli.main import main
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _enable_rtk(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# RTK is opt-in (off by default); these tests exercise the RTK-on injection path.
|
||||
monkeypatch.setenv("HEADROOM_RTK", "1")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def runner() -> CliRunner:
|
||||
return CliRunner()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Wrap: --prepare-only RTK injection into AGENTS.md
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_prepare_only_injects_rtk_into_agents_md(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""``wrap zcode --prepare-only`` writes the RTK block to AGENTS.md at cwd."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "zcode", "--prepare-only"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
marker = tmp_path / "AGENTS.md"
|
||||
assert marker.exists(), "AGENTS.md should be created"
|
||||
content = marker.read_text(encoding="utf-8")
|
||||
assert wrap_mod._RTK_MARKER in content
|
||||
assert "RTK (Rust Token Killer)" in content
|
||||
|
||||
|
||||
def test_prepare_only_idempotent_no_duplicate_block(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Running prepare-only twice must not duplicate the RTK block in AGENTS.md."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
runner.invoke(main, ["wrap", "zcode", "--prepare-only"])
|
||||
runner.invoke(main, ["wrap", "zcode", "--prepare-only"])
|
||||
|
||||
content = (tmp_path / "AGENTS.md").read_text(encoding="utf-8")
|
||||
assert content.count(wrap_mod._RTK_MARKER) == 1
|
||||
|
||||
|
||||
def test_no_context_tool_does_not_create_agents_md(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""--no-context-tool must not create AGENTS.md and must not invoke rtk."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary") as ensure:
|
||||
result = runner.invoke(main, ["wrap", "zcode", "--prepare-only", "--no-context-tool"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert not (tmp_path / "AGENTS.md").exists()
|
||||
ensure.assert_not_called()
|
||||
|
||||
|
||||
def test_preserves_existing_agents_md_content(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Pre-existing AGENTS.md content must be preserved when RTK is appended."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
agents_md = tmp_path / "AGENTS.md"
|
||||
original = "# Project conventions\n\nAlways use Python 3.12.\n"
|
||||
agents_md.write_text(original, encoding="utf-8")
|
||||
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
result = runner.invoke(main, ["wrap", "zcode", "--prepare-only"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
content = agents_md.read_text(encoding="utf-8")
|
||||
assert "Always use Python 3.12." in content
|
||||
assert wrap_mod._RTK_MARKER in content
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Wrap: setup instructions output
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -122,16 +37,13 @@ def test_wrap_prints_proxy_urls(
|
|||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
|
||||
fake_rtk = Path("/tmp/rtk")
|
||||
|
||||
def fake_watcher(**kwargs): # noqa: ANN003
|
||||
print_fn = kwargs.get("print_setup_lines")
|
||||
if callable(print_fn):
|
||||
print_fn(kwargs["port"])
|
||||
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=fake_rtk):
|
||||
with patch.object(wrap_mod, "_run_proxy_only_watcher", side_effect=fake_watcher):
|
||||
result = runner.invoke(main, ["wrap", "zcode", "--port", "9000"])
|
||||
with patch.object(wrap_mod, "_run_proxy_only_watcher", side_effect=fake_watcher):
|
||||
result = runner.invoke(main, ["wrap", "zcode", "--port", "9000"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "http://127.0.0.1:9000/v1" in result.output
|
||||
|
|
@ -139,85 +51,6 @@ def test_wrap_prints_proxy_urls(
|
|||
assert "Settings > Model Settings" in result.output
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unwrap: RTK removal from AGENTS.md
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_unwrap_removes_rtk_from_agents_md(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""``unwrap zcode`` removes RTK instructions from AGENTS.md."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
agents_md = tmp_path / "AGENTS.md"
|
||||
agents_md.write_text(
|
||||
"# Project\n\nSome content.\n\n" + wrap_mod.RTK_INSTRUCTIONS_BLOCK + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
with patch.object(wrap_mod, "_stop_local_proxy_for_unwrap", return_value="stopped"):
|
||||
result = runner.invoke(main, ["unwrap", "zcode"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "Removed Headroom rtk instructions" in result.output
|
||||
content = agents_md.read_text(encoding="utf-8")
|
||||
assert wrap_mod._RTK_MARKER not in content
|
||||
assert "Some content." in content
|
||||
|
||||
|
||||
def test_unwrap_deletes_empty_agents_md(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""``unwrap zcode`` deletes AGENTS.md if it only contained RTK instructions."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
agents_md = tmp_path / "AGENTS.md"
|
||||
agents_md.write_text(wrap_mod.RTK_INSTRUCTIONS_BLOCK + "\n", encoding="utf-8")
|
||||
|
||||
with patch.object(wrap_mod, "_stop_local_proxy_for_unwrap", return_value="stopped"):
|
||||
result = runner.invoke(main, ["unwrap", "zcode"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert not agents_md.exists(), "AGENTS.md should be deleted when only RTK content"
|
||||
|
||||
|
||||
def test_unwrap_noop_when_no_markers(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""``unwrap zcode`` is a safe no-op when AGENTS.md has no Headroom markers."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
agents_md = tmp_path / "AGENTS.md"
|
||||
agents_md.write_text("# Project\n\nSome content.\n", encoding="utf-8")
|
||||
|
||||
with patch.object(wrap_mod, "_stop_local_proxy_for_unwrap", return_value="stopped"):
|
||||
result = runner.invoke(main, ["unwrap", "zcode"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "Nothing to undo" in result.output
|
||||
content = agents_md.read_text(encoding="utf-8")
|
||||
assert content == "# Project\n\nSome content.\n"
|
||||
|
||||
|
||||
def test_unwrap_noop_when_no_agents_md(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""``unwrap zcode`` is a safe no-op when AGENTS.md does not exist."""
|
||||
monkeypatch.chdir(tmp_path)
|
||||
|
||||
with patch.object(wrap_mod, "_stop_local_proxy_for_unwrap", return_value="stopped"):
|
||||
result = runner.invoke(main, ["unwrap", "zcode"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "Nothing to undo" in result.output
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Runtime: proxy targets
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -471,10 +304,9 @@ def test_wrap_zcode_passes_upstream_to_watcher(
|
|||
provider_name="Z.ai", base_url="https://api.z.ai/api/anthropic", kind="anthropic"
|
||||
)
|
||||
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
with patch.object(wrap_mod, "_detect_zcode_upstream", return_value=fake_upstream):
|
||||
with patch.object(wrap_mod, "_run_proxy_only_watcher", side_effect=fake_watcher):
|
||||
runner.invoke(main, ["wrap", "zcode", "--port", "9000"])
|
||||
with patch.object(wrap_mod, "_detect_zcode_upstream", return_value=fake_upstream):
|
||||
with patch.object(wrap_mod, "_run_proxy_only_watcher", side_effect=fake_watcher):
|
||||
runner.invoke(main, ["wrap", "zcode", "--port", "9000"])
|
||||
|
||||
assert captured.get("anthropic_api_url") == "https://api.z.ai/api/anthropic"
|
||||
assert captured.get("openai_api_url") is None
|
||||
|
|
@ -501,10 +333,9 @@ def test_wrap_zcode_passes_openai_upstream(
|
|||
provider_name="Custom", base_url="https://my-api.example.com/v1", kind="openai"
|
||||
)
|
||||
|
||||
with patch.object(wrap_mod, "_ensure_rtk_binary", return_value=Path("/tmp/rtk")):
|
||||
with patch.object(wrap_mod, "_detect_zcode_upstream", return_value=fake_upstream):
|
||||
with patch.object(wrap_mod, "_run_proxy_only_watcher", side_effect=fake_watcher):
|
||||
runner.invoke(main, ["wrap", "zcode", "--port", "9000"])
|
||||
with patch.object(wrap_mod, "_detect_zcode_upstream", return_value=fake_upstream):
|
||||
with patch.object(wrap_mod, "_run_proxy_only_watcher", side_effect=fake_watcher):
|
||||
runner.invoke(main, ["wrap", "zcode", "--port", "9000"])
|
||||
|
||||
assert captured.get("anthropic_api_url") is None
|
||||
assert captured.get("openai_api_url") == "https://my-api.example.com/v1"
|
||||
|
|
|
|||
218
tests/test_context_tool_cleanup.py
Normal file
218
tests/test_context_tool_cleanup.py
Normal file
|
|
@ -0,0 +1,218 @@
|
|||
"""The retired rtk / lean-ctx integrations must be uninstalled, not just unshipped.
|
||||
|
||||
Deleting the integration code does nothing for a machine that already ran the
|
||||
old default — the Claude ``PreToolUse`` hook, the vendored binaries, the MCP
|
||||
registration and the injected hint-file guidance are all durable on disk. These
|
||||
tests pin the two properties that make the cleanup safe to run unattended on
|
||||
every ``wrap``: it removes everything Headroom put there, and it touches nothing
|
||||
else.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom import context_tool_cleanup, paths
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def home(monkeypatch, tmp_path):
|
||||
"""Point HOME, cwd and Headroom's bin dir at a scratch tree."""
|
||||
monkeypatch.setattr("pathlib.Path.home", lambda: tmp_path)
|
||||
monkeypatch.setattr(paths, "bin_dir", lambda: tmp_path / ".headroom" / "bin")
|
||||
monkeypatch.delenv("CODEX_HOME", raising=False)
|
||||
monkeypatch.delenv("OPENCODE_HOME", raising=False)
|
||||
project = tmp_path / "project"
|
||||
project.mkdir()
|
||||
monkeypatch.chdir(project)
|
||||
return tmp_path
|
||||
|
||||
|
||||
def _write(path, content):
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(content)
|
||||
return path
|
||||
|
||||
|
||||
def test_removes_hooks_for_both_tools_but_keeps_user_hooks(home):
|
||||
settings = _write(
|
||||
home / ".claude" / "settings.json",
|
||||
json.dumps(
|
||||
{
|
||||
"permissions": {"allow": ["Bash"]},
|
||||
"hooks": {
|
||||
"PreToolUse": [
|
||||
{
|
||||
"hooks": [
|
||||
{"type": "command", "command": "~/.claude/hooks/rtk-rewrite.sh"}
|
||||
]
|
||||
},
|
||||
{"hooks": [{"type": "command", "command": "lean-ctx hook rewrite"}]},
|
||||
{"hooks": [{"type": "command", "command": "my-own-linter --check"}]},
|
||||
],
|
||||
"SessionStart": [{"hooks": [{"type": "command", "command": "echo hi"}]}],
|
||||
},
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
report = context_tool_cleanup.purge_context_tool_artifacts()
|
||||
|
||||
payload = json.loads(settings.read_text())
|
||||
commands = [
|
||||
item["command"] for entry in payload["hooks"]["PreToolUse"] for item in entry["hooks"]
|
||||
]
|
||||
assert commands == ["my-own-linter --check"]
|
||||
# Unrelated events and unrelated top-level keys survive untouched.
|
||||
assert payload["hooks"]["SessionStart"][0]["hooks"][0]["command"] == "echo hi"
|
||||
assert payload["permissions"] == {"allow": ["Bash"]}
|
||||
assert any("hook" in line for line in report)
|
||||
|
||||
|
||||
def test_removes_binaries_hook_scripts_and_backups(home):
|
||||
bin_dir = home / ".headroom" / "bin"
|
||||
rtk = _write(bin_dir / "rtk", "binary")
|
||||
lean = _write(bin_dir / "lean-ctx", "binary")
|
||||
script = _write(home / ".claude" / "hooks" / "lean-ctx-rewrite.sh", "#!/bin/sh\n")
|
||||
backup = _write(home / ".claude" / "hooks" / "lean-ctx-rewrite.sh.lean-ctx.bak", "#!/bin/sh\n")
|
||||
|
||||
context_tool_cleanup.purge_context_tool_artifacts()
|
||||
|
||||
assert not rtk.exists()
|
||||
assert not lean.exists()
|
||||
assert not script.exists()
|
||||
assert not backup.exists()
|
||||
|
||||
|
||||
def test_leaves_a_users_own_binary_on_path_alone(home):
|
||||
"""A real file in ~/.local/bin is not ours to reclaim — only our symlink is."""
|
||||
own = _write(home / ".local" / "bin" / "lean-ctx", "my own build")
|
||||
managed = _write(home / ".headroom" / "bin" / "rtk", "binary")
|
||||
link = home / ".local" / "bin" / "rtk"
|
||||
link.symlink_to(managed)
|
||||
|
||||
context_tool_cleanup.purge_context_tool_artifacts()
|
||||
|
||||
assert own.exists() and own.read_text() == "my own build"
|
||||
assert not link.exists()
|
||||
|
||||
|
||||
def test_removes_mcp_entry_and_preserves_siblings(home):
|
||||
config = _write(
|
||||
home / ".claude.json",
|
||||
json.dumps(
|
||||
{
|
||||
"projects": {"/some/path": {"history": []}},
|
||||
"mcpServers": {
|
||||
"lean-ctx": {"command": "lean-ctx", "args": ["mcp"]},
|
||||
"headroom": {"command": "headroom", "args": ["mcp"]},
|
||||
},
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
context_tool_cleanup.purge_context_tool_artifacts()
|
||||
|
||||
payload = json.loads(config.read_text())
|
||||
assert list(payload["mcpServers"]) == ["headroom"]
|
||||
assert payload["projects"] == {"/some/path": {"history": []}}
|
||||
|
||||
|
||||
def test_strips_guidance_fence_but_keeps_surrounding_prose(home):
|
||||
agents = _write(
|
||||
home / "project" / "AGENTS.md",
|
||||
"# My project\n\nMy own notes.\n\n"
|
||||
"<!-- headroom:rtk-instructions -->\nAlways prefix with rtk.\n"
|
||||
"<!-- /headroom:rtk-instructions -->\n",
|
||||
)
|
||||
|
||||
context_tool_cleanup.purge_context_tool_artifacts()
|
||||
|
||||
content = agents.read_text()
|
||||
assert "rtk" not in content
|
||||
assert "My own notes." in content
|
||||
assert content.startswith("# My project")
|
||||
|
||||
|
||||
def test_skips_malformed_json_instead_of_clobbering_it(home):
|
||||
settings = _write(home / ".claude" / "settings.json", '{"permissions": {oops')
|
||||
|
||||
report = context_tool_cleanup.purge_context_tool_artifacts()
|
||||
|
||||
assert settings.read_text() == '{"permissions": {oops'
|
||||
assert any("skipped" in line for line in report)
|
||||
|
||||
|
||||
def test_is_idempotent(home):
|
||||
_write(home / ".headroom" / "bin" / "rtk", "binary")
|
||||
_write(
|
||||
home / ".claude" / "settings.json",
|
||||
json.dumps({"hooks": {"PreToolUse": [{"hooks": [{"command": "rtk rewrite"}]}]}}),
|
||||
)
|
||||
|
||||
assert context_tool_cleanup.purge_context_tool_artifacts()
|
||||
# Steady state after the first run: nothing left to report.
|
||||
assert context_tool_cleanup.purge_context_tool_artifacts() == []
|
||||
|
||||
|
||||
def test_no_op_on_a_clean_machine(home):
|
||||
assert context_tool_cleanup.purge_context_tool_artifacts() == []
|
||||
|
||||
|
||||
def test_purge_reports_on_stderr_so_json_stdout_stays_parseable(home):
|
||||
"""`wrap openclaw --prepare-only` emits machine-readable JSON as its whole contract.
|
||||
|
||||
The purge runs from the `wrap` group callback, i.e. before that JSON is
|
||||
written. Reporting on stdout prepended a human line to it and broke every
|
||||
``json.loads(stdout)`` consumer — but only on the single run that actually
|
||||
had something to remove, so a clean CI machine never caught it.
|
||||
"""
|
||||
from click.testing import CliRunner
|
||||
|
||||
from headroom.cli.main import main
|
||||
|
||||
_write(home / ".headroom" / "bin" / "rtk", "binary")
|
||||
|
||||
result = CliRunner().invoke(
|
||||
main, ["wrap", "openclaw", "--prepare-only", "--gateway-provider-id", "codex"]
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
# Whole of stdout must still parse — no cleanup preamble.
|
||||
assert json.loads(result.stdout)["enabled"] is True
|
||||
assert "Retired CLI context tool cleanup" in result.stderr
|
||||
|
||||
|
||||
def test_help_does_not_purge(home, monkeypatch):
|
||||
"""`--help` must stay read-only — reading help should not delete files."""
|
||||
from click.testing import CliRunner
|
||||
|
||||
from headroom.cli.main import main
|
||||
|
||||
binary = _write(home / ".headroom" / "bin" / "rtk", "binary")
|
||||
monkeypatch.setattr("sys.argv", ["headroom", "wrap", "codex", "--help"])
|
||||
|
||||
result = CliRunner().invoke(main, ["wrap", "codex", "--help"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert binary.exists(), "--help performed filesystem cleanup"
|
||||
|
||||
|
||||
def test_selfheal_does_not_purge(home, monkeypatch):
|
||||
"""`wrap selfheal` runs from a SessionStart hook — no config surgery there.
|
||||
|
||||
It fires on every new conversation, where rewriting ~/.claude.json would race
|
||||
Claude Code's own writer.
|
||||
"""
|
||||
from click.testing import CliRunner
|
||||
|
||||
from headroom.cli.main import main
|
||||
|
||||
binary = _write(home / ".headroom" / "bin" / "rtk", "binary")
|
||||
monkeypatch.setattr("sys.argv", ["headroom", "wrap", "selfheal"])
|
||||
|
||||
CliRunner().invoke(main, ["wrap", "selfheal", "--marker", "headroom-wrap-selfheal"])
|
||||
|
||||
assert binary.exists(), "selfheal performed filesystem cleanup"
|
||||
|
|
@ -23,7 +23,6 @@ def _sample_stats() -> dict:
|
|||
"savings_usd": 12.34,
|
||||
"compression_savings_usd": 12.34,
|
||||
"cache_savings_usd": 5.25,
|
||||
"cli_tokens_avoided": 0,
|
||||
},
|
||||
"requests": {
|
||||
"total": 128,
|
||||
|
|
@ -37,7 +36,6 @@ def _sample_stats() -> dict:
|
|||
"input": 245_000,
|
||||
"output": 88_000,
|
||||
"saved": 143_000,
|
||||
"cli_tokens_avoided": 0,
|
||||
"total_before_compression": 388_000,
|
||||
"savings_percent": 36.86,
|
||||
},
|
||||
|
|
|
|||
|
|
@ -1,233 +0,0 @@
|
|||
"""Behavior-driven Playwright validation for context-tool (RTK) availability
|
||||
messaging on the dashboard (issue #1831).
|
||||
|
||||
Local runs only import/collect this file (Windows dev sandboxes don't run the
|
||||
real browser here); the "Dashboard Playwright" CI check executes it for real.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import json
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.dashboard import get_dashboard_html
|
||||
|
||||
playwright = pytest.importorskip("playwright.sync_api")
|
||||
Page = playwright.Page
|
||||
expect = playwright.expect
|
||||
sync_playwright = playwright.sync_playwright
|
||||
|
||||
|
||||
def _base_stats() -> dict:
|
||||
return {
|
||||
"cost": {
|
||||
"savings_usd": 12.34,
|
||||
"compression_savings_usd": 12.34,
|
||||
"cache_savings_usd": 5.25,
|
||||
"cli_tokens_avoided": 0,
|
||||
},
|
||||
"requests": {
|
||||
"total": 128,
|
||||
"cached": 96,
|
||||
"rate_limited": 0,
|
||||
"failed": 0,
|
||||
"by_provider": {"anthropic": 128},
|
||||
"by_model": {"claude-opus-4-6": 128},
|
||||
},
|
||||
"tokens": {
|
||||
"input": 245_000,
|
||||
"output": 88_000,
|
||||
"saved": 143_000,
|
||||
"cli_tokens_avoided": 0,
|
||||
"total_before_compression": 388_000,
|
||||
"savings_percent": 36.86,
|
||||
},
|
||||
"overhead": {"average_ms": 14.2, "min_ms": 4.5, "max_ms": 42.7},
|
||||
"ttfb": {"average_ms": 1320.0, "min_ms": 420.0, "max_ms": 2900.0},
|
||||
"latency": {"average_ms": 1510.0, "min_ms": 520.0, "max_ms": 3300.0},
|
||||
"waste_signals": {"json_bloat": 95_000, "repetition": 48_000},
|
||||
"savings_history": [
|
||||
["2026-04-01T00:00:00Z", 12_000],
|
||||
["2026-04-05T00:00:00Z", 143_000],
|
||||
],
|
||||
"persistent_savings": {
|
||||
"display_session": {},
|
||||
"lifetime": {"tokens_saved": 143_000, "compression_savings_usd": 12.34},
|
||||
},
|
||||
"pipeline_timing": {},
|
||||
"compression_cache": {"mode": "cache"},
|
||||
"prefix_cache": {"by_provider": {}, "totals": {}, "prefix_freeze": {}},
|
||||
}
|
||||
|
||||
|
||||
def _sample_stats(*, available: bool, tokens_saved: int = 0) -> dict:
|
||||
"""Build a /stats payload with context_tool/cli_filtering availability set.
|
||||
|
||||
`_base_stats()` has no `savings` key at all, so writing
|
||||
`savings.by_layer.cli_filtering` requires a `.setdefault(...)` chain
|
||||
rather than direct key assignment (would otherwise raise `KeyError`).
|
||||
"""
|
||||
stats = copy.deepcopy(_base_stats())
|
||||
stats["tokens"]["cli_tokens_avoided"] = tokens_saved
|
||||
stats["context_tool"] = {
|
||||
"configured": "rtk",
|
||||
"label": "RTK",
|
||||
"available": available,
|
||||
"stats": {"tool": "rtk", "label": "RTK", "installed": available},
|
||||
}
|
||||
cli_filtering = (
|
||||
stats.setdefault("savings", {}).setdefault("by_layer", {}).setdefault("cli_filtering", {})
|
||||
)
|
||||
cli_filtering.update(
|
||||
{
|
||||
"tool": "rtk",
|
||||
"label": "RTK",
|
||||
"available": available,
|
||||
"tokens": tokens_saved,
|
||||
"tokens_saved": tokens_saved,
|
||||
"session": {},
|
||||
"lifetime": {"tokens_saved": 0},
|
||||
"session_savings_pct": 0.0,
|
||||
}
|
||||
)
|
||||
return stats
|
||||
|
||||
|
||||
def _sample_history(*, available: bool, lifetime_tokens_saved: int = 456_700) -> dict:
|
||||
return {
|
||||
"history": [
|
||||
{
|
||||
"timestamp": "2026-04-05T00:00:00Z",
|
||||
"total_tokens_saved": 143_000,
|
||||
"compression_savings_usd": 12.34,
|
||||
},
|
||||
],
|
||||
"series": {"daily": [], "weekly": [], "monthly": []},
|
||||
"lifetime": {"tokens_saved": 143_000, "compression_savings_usd": 12.34},
|
||||
"cli_filtering": {
|
||||
"tool": "rtk",
|
||||
"label": "RTK",
|
||||
"available": available,
|
||||
"lifetime": {"tokens_saved": lifetime_tokens_saved},
|
||||
"session": {},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _install_dashboard_routes(page: Page, stats: dict, history: dict) -> None:
|
||||
health = {"status": "healthy", "version": "0.3.0"}
|
||||
dashboard_html = get_dashboard_html()
|
||||
|
||||
def handler(route) -> None: # type: ignore[no-untyped-def]
|
||||
# Match on the URL path only: the dashboard fetches /stats?cached=1,
|
||||
# so suffix checks against the full URL miss it and the request
|
||||
# escapes the harness to the real network.
|
||||
path = urlsplit(route.request.url).path
|
||||
if path in ("/dashboard", "/"):
|
||||
route.fulfill(status=200, content_type="text/html", body=dashboard_html)
|
||||
return
|
||||
if "/stats-history" in path:
|
||||
route.fulfill(
|
||||
status=200,
|
||||
content_type="application/json",
|
||||
body=json.dumps(history),
|
||||
)
|
||||
return
|
||||
if path.endswith("/stats"):
|
||||
route.fulfill(status=200, content_type="application/json", body=json.dumps(stats))
|
||||
return
|
||||
if path.endswith("/health"):
|
||||
route.fulfill(status=200, content_type="application/json", body=json.dumps(health))
|
||||
return
|
||||
route.continue_()
|
||||
|
||||
page.route("**/*", handler)
|
||||
|
||||
|
||||
def test_dashboard_session_view_shows_not_installed_message_when_unavailable() -> None:
|
||||
"""Session view shows a distinct "not installed" message, not `0`, when
|
||||
`context_tool.available` is False (the #1831 bug this fix addresses).
|
||||
"""
|
||||
stats = _sample_stats(available=False, tokens_saved=0)
|
||||
history = _sample_history(available=False)
|
||||
|
||||
with sync_playwright() as pw:
|
||||
browser = pw.chromium.launch()
|
||||
page = browser.new_page(viewport={"width": 1720, "height": 1400}, color_scheme="dark")
|
||||
_install_dashboard_routes(page, stats, history)
|
||||
page.goto("http://headroom.local/dashboard", wait_until="load")
|
||||
|
||||
expect(page.get_by_text("RTK not installed", exact=True)).to_be_visible()
|
||||
expect(page.get_by_text("not installed", exact=True)).to_be_visible()
|
||||
expect(page.get_by_text("RTK 0 this session (0.0%)", exact=True)).to_have_count(0)
|
||||
|
||||
browser.close()
|
||||
|
||||
|
||||
def test_dashboard_session_view_shows_real_zero_row_when_installed_but_zero() -> None:
|
||||
"""Boundary value 0.0: installed but genuinely zero savings still renders
|
||||
the real number, not the "not installed" message -- proves the new guard
|
||||
doesn't over-trigger on the exact case the #1831 reporter would hit again.
|
||||
"""
|
||||
stats = _sample_stats(available=True, tokens_saved=0)
|
||||
history = _sample_history(available=True, lifetime_tokens_saved=0)
|
||||
|
||||
with sync_playwright() as pw:
|
||||
browser = pw.chromium.launch()
|
||||
page = browser.new_page(viewport={"width": 1720, "height": 1400}, color_scheme="dark")
|
||||
_install_dashboard_routes(page, stats, history)
|
||||
page.goto("http://headroom.local/dashboard", wait_until="load")
|
||||
|
||||
expect(page.get_by_text("RTK 0 this session (0.0%)", exact=True)).to_be_visible()
|
||||
expect(page.get_by_text("RTK not installed", exact=True)).to_have_count(0)
|
||||
# The Token Usage panel's "not installed" row uses `x-show`, which
|
||||
# toggles CSS display and keeps the node in the DOM (unlike the
|
||||
# ternary-swapped "RTK not installed" text above, which is genuinely
|
||||
# absent). Assert hidden, not absent, matching the repo's existing
|
||||
# `x-show` convention in tests/test_dashboard_cache_lifetime_playwright.py.
|
||||
expect(page.get_by_text("not installed", exact=True)).to_be_hidden()
|
||||
|
||||
browser.close()
|
||||
|
||||
|
||||
def test_dashboard_historical_tab_hides_lifetime_card_when_unavailable() -> None:
|
||||
"""The Historical tab's lifetime card stays hidden (same as the existing
|
||||
hard-failure hide-card behavior) when `cli_filtering.available` is False.
|
||||
"""
|
||||
stats = _sample_stats(available=True, tokens_saved=5_000)
|
||||
history = _sample_history(available=False)
|
||||
|
||||
with sync_playwright() as pw:
|
||||
browser = pw.chromium.launch()
|
||||
page = browser.new_page(viewport={"width": 1720, "height": 1400}, color_scheme="dark")
|
||||
_install_dashboard_routes(page, stats, history)
|
||||
page.goto("http://headroom.local/dashboard", wait_until="load")
|
||||
|
||||
page.get_by_role("button", name="Historical").click()
|
||||
expect(page.get_by_text("Historical Summary")).to_be_visible()
|
||||
expect(page.get_by_text("RTK Lifetime Saved")).to_have_count(0)
|
||||
|
||||
browser.close()
|
||||
|
||||
|
||||
def test_dashboard_historical_tab_shows_lifetime_card_when_available() -> None:
|
||||
"""The Historical tab's lifetime card still renders the real number when
|
||||
`cli_filtering.available` is True (existing behavior, unchanged).
|
||||
"""
|
||||
stats = _sample_stats(available=True, tokens_saved=5_000)
|
||||
history = _sample_history(available=True, lifetime_tokens_saved=456_700)
|
||||
|
||||
with sync_playwright() as pw:
|
||||
browser = pw.chromium.launch()
|
||||
page = browser.new_page(viewport={"width": 1720, "height": 1400}, color_scheme="dark")
|
||||
_install_dashboard_routes(page, stats, history)
|
||||
page.goto("http://headroom.local/dashboard", wait_until="load")
|
||||
|
||||
page.get_by_role("button", name="Historical").click()
|
||||
expect(page.get_by_text("RTK Lifetime Saved")).to_be_visible()
|
||||
expect(page.get_by_text("456.7k")).to_be_visible()
|
||||
|
||||
browser.close()
|
||||
|
|
@ -2,6 +2,9 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import stat
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom import fsutil
|
||||
|
|
@ -39,6 +42,49 @@ def test_read_text_roundtrips_utf8_non_ascii(tmp_path):
|
|||
assert fsutil.read_text(p) == 'project = "比赛/机器人"\n'
|
||||
|
||||
|
||||
def test_write_text_leaves_original_intact_when_write_fails(tmp_path, monkeypatch):
|
||||
"""A failed write must leave the previous file whole, not truncated.
|
||||
|
||||
``open(path, "w")`` truncates to zero bytes before the new content lands, so
|
||||
a crash mid-write used to leave users with a half-written ~/.claude.json.
|
||||
"""
|
||||
p = tmp_path / "settings.json"
|
||||
fsutil.write_text(p, '{"permissions": {"allow": ["Bash"]}}\n')
|
||||
|
||||
def boom(src, dst):
|
||||
raise OSError("no space left on device")
|
||||
|
||||
monkeypatch.setattr(os, "replace", boom)
|
||||
with pytest.raises(OSError):
|
||||
fsutil.write_text(p, '{"hooks": {}}\n')
|
||||
|
||||
assert p.read_text() == '{"permissions": {"allow": ["Bash"]}}\n'
|
||||
# No .tmp litter left behind in the config directory.
|
||||
assert [f.name for f in tmp_path.iterdir()] == ["settings.json"]
|
||||
|
||||
|
||||
def test_write_text_follows_symlink_instead_of_replacing_it(tmp_path):
|
||||
"""Dotfile managers symlink these configs; os.replace would clobber the link."""
|
||||
real = tmp_path / "real.json"
|
||||
real.write_text("{}\n")
|
||||
link = tmp_path / "settings.json"
|
||||
link.symlink_to(real)
|
||||
|
||||
fsutil.write_text(link, '{"env": {}}\n')
|
||||
|
||||
assert link.is_symlink()
|
||||
assert real.read_text() == '{"env": {}}\n'
|
||||
|
||||
|
||||
def test_write_text_preserves_existing_mode(tmp_path):
|
||||
"""mkstemp creates 0600 — an existing 0644 config must not be silently tightened."""
|
||||
p = tmp_path / "settings.json"
|
||||
p.write_text("{}\n")
|
||||
p.chmod(0o644)
|
||||
fsutil.write_text(p, '{"a": 1}\n')
|
||||
assert stat.S_IMODE(p.stat().st_mode) == 0o644
|
||||
|
||||
|
||||
def test_read_text_falls_back_to_locale_encoding(tmp_path, monkeypatch):
|
||||
"""A file a tool wrote in the locale encoding (e.g. GBK) still decodes."""
|
||||
monkeypatch.setattr(fsutil.locale, "getpreferredencoding", lambda *_: "gbk")
|
||||
|
|
|
|||
|
|
@ -1,176 +0,0 @@
|
|||
"""Tests for managed lean-ctx installation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import tarfile
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.lean_ctx import get_lean_ctx_path, installer
|
||||
|
||||
|
||||
def test_get_lean_ctx_path_finds_windows_managed_binary(tmp_path: Path) -> None:
|
||||
managed_dir = tmp_path / ".headroom" / "bin"
|
||||
managed_dir.mkdir(parents=True)
|
||||
managed_path = managed_dir / "lean-ctx.exe"
|
||||
managed_path.write_bytes(b"binary")
|
||||
|
||||
with patch("headroom.lean_ctx.LEAN_CTX_BIN_DIR", managed_dir):
|
||||
with patch("headroom.lean_ctx.LEAN_CTX_BIN_PATH", managed_dir / "lean-ctx"):
|
||||
with patch("headroom.lean_ctx.shutil.which", return_value=None):
|
||||
assert get_lean_ctx_path() == managed_path
|
||||
|
||||
|
||||
def test_get_target_triple_uses_override(monkeypatch) -> None:
|
||||
monkeypatch.setenv("HEADROOM_LEAN_CTX_TARGET", "x86_64-pc-windows-msvc")
|
||||
assert installer._get_target_triple() == "x86_64-pc-windows-msvc"
|
||||
|
||||
|
||||
def test_detect_runtime_target_triple_handles_linux_gnu() -> None:
|
||||
with patch.object(installer.platform, "system", return_value="Linux"):
|
||||
with patch.object(installer.platform, "machine", return_value="x86_64"):
|
||||
with patch.object(installer, "_is_musl", return_value=False):
|
||||
assert installer._detect_runtime_target_triple() == "x86_64-unknown-linux-gnu"
|
||||
|
||||
|
||||
def test_detect_runtime_target_triple_handles_linux_musl_arm() -> None:
|
||||
with patch.object(installer.platform, "system", return_value="Linux"):
|
||||
with patch.object(installer.platform, "machine", return_value="aarch64"):
|
||||
with patch.object(installer, "_is_musl", return_value=True):
|
||||
assert installer._detect_runtime_target_triple() == "aarch64-unknown-linux-musl"
|
||||
|
||||
|
||||
def test_get_download_url_uses_windows_zip(monkeypatch) -> None:
|
||||
monkeypatch.delenv("HEADROOM_LEAN_CTX_TARGET", raising=False)
|
||||
monkeypatch.setenv("LEAN_CTX_TARGET", "x86_64-pc-windows-msvc")
|
||||
|
||||
url, ext = installer._get_download_url("v1.2.3")
|
||||
|
||||
assert url == f"{installer.GITHUB_RELEASE_URL}/v1.2.3/lean-ctx-x86_64-pc-windows-msvc.zip"
|
||||
assert ext == "zip"
|
||||
assert installer._binary_name_for_target("x86_64-pc-windows-msvc") == "lean-ctx.exe"
|
||||
|
||||
|
||||
def test_download_lean_ctx_skips_verify_for_non_native_target(monkeypatch, tmp_path: Path) -> None:
|
||||
archive = io.BytesIO()
|
||||
with tarfile.open(fileobj=archive, mode="w:gz") as tf:
|
||||
info = tarfile.TarInfo(name="lean-ctx")
|
||||
payload = b"fake-binary"
|
||||
info.size = len(payload)
|
||||
tf.addfile(info, io.BytesIO(payload))
|
||||
archive_bytes = archive.getvalue()
|
||||
|
||||
class _Response:
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def read(self) -> bytes:
|
||||
return archive_bytes
|
||||
|
||||
monkeypatch.setenv("HEADROOM_LEAN_CTX_TARGET", "x86_64-apple-darwin")
|
||||
|
||||
with patch.object(installer, "LEAN_CTX_BIN_DIR", tmp_path):
|
||||
with patch.object(installer, "urlopen", return_value=_Response()):
|
||||
with patch.object(installer.subprocess, "run") as subprocess_run:
|
||||
installed_path = installer.download_lean_ctx("v3.4.7")
|
||||
|
||||
assert installed_path == tmp_path / "lean-ctx"
|
||||
assert installed_path.exists()
|
||||
subprocess_run.assert_not_called()
|
||||
|
||||
|
||||
def test_download_lean_ctx_extracts_zip_binary(monkeypatch, tmp_path: Path) -> None:
|
||||
archive = io.BytesIO()
|
||||
with zipfile.ZipFile(archive, mode="w") as zf:
|
||||
zf.writestr("lean-ctx.exe", b"fake-windows-binary")
|
||||
|
||||
class _Response:
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def read(self) -> bytes:
|
||||
return archive.getvalue()
|
||||
|
||||
monkeypatch.setenv("HEADROOM_LEAN_CTX_TARGET", "x86_64-pc-windows-msvc")
|
||||
|
||||
with patch.object(installer, "LEAN_CTX_BIN_DIR", tmp_path):
|
||||
with patch.object(installer, "urlopen", return_value=_Response()):
|
||||
installed_path = installer.download_lean_ctx("v3.4.7")
|
||||
|
||||
assert installed_path == tmp_path / "lean-ctx.exe"
|
||||
assert installed_path.read_bytes() == b"fake-windows-binary"
|
||||
|
||||
|
||||
def test_download_lean_ctx_verifies_native_target(monkeypatch, tmp_path: Path) -> None:
|
||||
monkeypatch.delenv("HEADROOM_LEAN_CTX_TARGET", raising=False)
|
||||
monkeypatch.delenv("LEAN_CTX_TARGET", raising=False)
|
||||
|
||||
archive = io.BytesIO()
|
||||
with tarfile.open(fileobj=archive, mode="w:gz") as tf:
|
||||
info = tarfile.TarInfo(name="dist/lean-ctx")
|
||||
payload = b"fake-native-binary"
|
||||
info.size = len(payload)
|
||||
tf.addfile(info, io.BytesIO(payload))
|
||||
|
||||
class _Response:
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def read(self) -> bytes:
|
||||
return archive.getvalue()
|
||||
|
||||
run_result = SimpleNamespace(returncode=0, stdout="lean-ctx 3.4.7", stderr="")
|
||||
|
||||
with patch.object(installer, "LEAN_CTX_BIN_DIR", tmp_path):
|
||||
with patch.object(
|
||||
installer, "_detect_runtime_target_triple", return_value="x86_64-unknown-linux-gnu"
|
||||
):
|
||||
with patch.object(installer, "urlopen", return_value=_Response()):
|
||||
with patch.object(
|
||||
installer.subprocess, "run", return_value=run_result
|
||||
) as subprocess_run:
|
||||
installed_path = installer.download_lean_ctx("v3.4.7")
|
||||
|
||||
assert installed_path == tmp_path / "lean-ctx"
|
||||
subprocess_run.assert_called_once_with(
|
||||
[str(installed_path), "--version"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
encoding="utf-8",
|
||||
errors="replace",
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
|
||||
def test_download_lean_ctx_rejects_invalid_download_url(tmp_path: Path) -> None:
|
||||
with patch.object(installer, "LEAN_CTX_BIN_DIR", tmp_path):
|
||||
with patch.object(installer, "_get_target_triple", return_value="x86_64-unknown-linux-gnu"):
|
||||
with patch.object(
|
||||
installer,
|
||||
"_get_download_url",
|
||||
return_value=("file:///tmp/lean-ctx.tar.gz", "tar.gz"),
|
||||
):
|
||||
with pytest.raises(RuntimeError, match="Invalid URL scheme"):
|
||||
installer.download_lean_ctx("v3.4.7")
|
||||
|
||||
|
||||
def test_ensure_lean_ctx_returns_none_when_download_fails() -> None:
|
||||
with patch("headroom.lean_ctx.get_lean_ctx_path", return_value=None):
|
||||
with patch.object(
|
||||
installer, "download_lean_ctx", side_effect=RuntimeError("download failed")
|
||||
):
|
||||
assert installer.ensure_lean_ctx() is None
|
||||
|
|
@ -1,10 +1,10 @@
|
|||
"""Tests for loop detection and loop-weighting in Headroom Learn.
|
||||
|
||||
Covers the gap these changes close: RTK re-fetch loops (repeated, successful
|
||||
Covers the gap these changes close: re-fetch loops (repeated, successful
|
||||
but insufficient calls) were invisible to failure-only analysis and, even when
|
||||
surfaced, were ranked no higher than a one-off rule. These tests pin:
|
||||
|
||||
1. ``detect_loops`` finds RTK re-fetch loops and error loops, and ignores
|
||||
1. ``detect_loops`` finds re-fetch loops and error loops, and ignores
|
||||
one-offs — collapsing output-limit variants to one signature.
|
||||
2. The digest surfaces detected loops as a high-priority section.
|
||||
3. ``apply_loop_weighting`` lifts a loop guardrail above a one-off rule using
|
||||
|
|
@ -20,7 +20,7 @@ from headroom.learn.analyzer import SessionAnalyzer, _build_digest
|
|||
from headroom.learn.fixtures import (
|
||||
error_loop_session,
|
||||
one_off_error_session,
|
||||
rtk_refetch_loop_session,
|
||||
refetch_loop_session,
|
||||
)
|
||||
from headroom.learn.loops import (
|
||||
_canonical_signature,
|
||||
|
|
@ -48,19 +48,19 @@ def _project() -> ProjectInfo:
|
|||
|
||||
|
||||
class TestDetectLoops:
|
||||
def test_rtk_refetch_loop_detected_despite_no_errors(self):
|
||||
loops = detect_loops([rtk_refetch_loop_session(repetitions=5)])
|
||||
def test_refetch_loop_detected_despite_no_errors(self):
|
||||
loops = detect_loops([refetch_loop_session(repetitions=5)])
|
||||
assert len(loops) == 1
|
||||
lp = loops[0]
|
||||
assert lp.count == 5
|
||||
assert lp.is_error_loop is False
|
||||
assert lp.kind == "rtk-refetch-loop"
|
||||
assert lp.kind == "refetch-loop"
|
||||
# Waste counts the 4 redundant re-fetches (not the first legit call).
|
||||
assert lp.wasted_tokens > 0
|
||||
|
||||
def test_output_limit_variants_collapse_to_one_signature(self):
|
||||
# The five calls differ only by `head -50/-100/...`; same signature.
|
||||
session = rtk_refetch_loop_session(repetitions=5)
|
||||
session = refetch_loop_session(repetitions=5)
|
||||
sigs = {_canonical_signature(tc) for tc in session.tool_calls}
|
||||
assert len(sigs) == 1
|
||||
|
||||
|
|
@ -75,13 +75,13 @@ class TestDetectLoops:
|
|||
|
||||
def test_min_occurrences_threshold(self):
|
||||
# Two repetitions is a retry, not a loop, at the default threshold.
|
||||
assert detect_loops([rtk_refetch_loop_session(repetitions=2)]) == []
|
||||
assert detect_loops([rtk_refetch_loop_session(repetitions=3)])
|
||||
assert detect_loops([refetch_loop_session(repetitions=2)]) == []
|
||||
assert detect_loops([refetch_loop_session(repetitions=3)])
|
||||
|
||||
def test_error_loop_waste_exceeds_refetch_loop_first_call_credit(self):
|
||||
# Error loops waste every call; re-fetch loops credit the first call.
|
||||
err = detect_loops([error_loop_session(repetitions=4)])[0]
|
||||
ref = detect_loops([rtk_refetch_loop_session(repetitions=4)])[0]
|
||||
ref = detect_loops([refetch_loop_session(repetitions=4)])[0]
|
||||
assert err.count == ref.count
|
||||
# Same count, but error loop counts all N and re-fetch counts N-1.
|
||||
assert err.wasted_tokens >= 0 and ref.wasted_tokens >= 0
|
||||
|
|
@ -94,9 +94,9 @@ class TestDetectLoops:
|
|||
|
||||
class TestDigestSurfacesLoops:
|
||||
def test_digest_includes_detected_loops_section(self):
|
||||
digest = _build_digest(_project(), [rtk_refetch_loop_session()])
|
||||
digest = _build_digest(_project(), [refetch_loop_session()])
|
||||
assert "Detected Loops" in digest
|
||||
assert "rtk-refetch-loop" in digest
|
||||
assert "refetch-loop" in digest
|
||||
assert "tokens wasted" in digest
|
||||
|
||||
def test_digest_without_loops_has_no_loop_section(self):
|
||||
|
|
@ -128,7 +128,7 @@ class TestApplyLoopWeighting:
|
|||
)
|
||||
|
||||
def test_loop_rule_boosted_above_one_off(self):
|
||||
loops = detect_loops([rtk_refetch_loop_session(repetitions=5)])
|
||||
loops = detect_loops([refetch_loop_session(repetitions=5)])
|
||||
recs = [self._one_off_rec(), self._loop_rec()]
|
||||
apply_loop_weighting(recs, loops)
|
||||
|
||||
|
|
@ -148,7 +148,7 @@ class TestApplyLoopWeighting:
|
|||
assert recs[0].is_loop_guardrail is False
|
||||
|
||||
def test_unrelated_rule_not_credited(self):
|
||||
loops = detect_loops([rtk_refetch_loop_session(repetitions=5)])
|
||||
loops = detect_loops([refetch_loop_session(repetitions=5)])
|
||||
recs = [self._one_off_rec()] # about uv/python, not the grep loop
|
||||
apply_loop_weighting(recs, loops)
|
||||
assert recs[0].is_loop_guardrail is False
|
||||
|
|
@ -165,7 +165,7 @@ class TestAnalyzeEndToEnd:
|
|||
# Pure re-fetch loop: zero errors, no events. Must NOT early-return.
|
||||
mock_call_llm.return_value = {"context_file_rules": [], "memory_file_rules": []}
|
||||
analyzer = SessionAnalyzer(model="test-model")
|
||||
analyzer.analyze(_project(), [rtk_refetch_loop_session()])
|
||||
analyzer.analyze(_project(), [refetch_loop_session()])
|
||||
mock_call_llm.assert_called_once() # the guard let it through
|
||||
|
||||
@patch("headroom.learn.analyzer._call_llm")
|
||||
|
|
@ -190,7 +190,7 @@ class TestAnalyzeEndToEnd:
|
|||
"memory_file_rules": [],
|
||||
}
|
||||
analyzer = SessionAnalyzer(model="test-model")
|
||||
result = analyzer.analyze(_project(), [rtk_refetch_loop_session(repetitions=6)])
|
||||
result = analyzer.analyze(_project(), [refetch_loop_session(repetitions=6)])
|
||||
|
||||
# After weighting, the loop guardrail ranks first despite the LLM's order.
|
||||
assert result.recommendations[0].is_loop_guardrail is True
|
||||
|
|
|
|||
|
|
@ -1,27 +0,0 @@
|
|||
"""CI wrapper for the RTK-loop eval (benchmarks/rtk_loop_learn_eval.py).
|
||||
|
||||
The deterministic path runs everywhere and gates the loop-weighting behavior
|
||||
end-to-end. The real-LLM path is opt-in via the repo's ``real_llm`` marker and
|
||||
only runs when an API key is present.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from benchmarks.rtk_loop_learn_eval import run_eval
|
||||
|
||||
|
||||
def test_rtk_loop_eval_deterministic():
|
||||
card = run_eval(use_real_llm=False)
|
||||
assert card.passed, "RTK-loop eval failed:\n" + card.render()
|
||||
|
||||
|
||||
@pytest.mark.real_llm
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get("ANTHROPIC_API_KEY"),
|
||||
reason="real_llm eval needs ANTHROPIC_API_KEY",
|
||||
)
|
||||
def test_rtk_loop_eval_real_llm():
|
||||
card = run_eval(use_real_llm=True)
|
||||
assert card.passed, "RTK-loop eval (real LLM) failed:\n" + card.render()
|
||||
|
|
@ -2,7 +2,6 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
|
@ -136,7 +135,6 @@ def test_per_resource_getters_no_mkdir(fake_home: Path) -> None:
|
|||
paths.proxy_log_path()
|
||||
paths.debug_400_dir()
|
||||
paths.bin_dir()
|
||||
paths.rtk_path()
|
||||
paths.deploy_root()
|
||||
paths.beacon_lock_path(8787)
|
||||
paths.models_config_path()
|
||||
|
|
@ -320,18 +318,6 @@ def test_bin_dir_default(fake_home: Path) -> None:
|
|||
assert paths.bin_dir() == fake_home / ".headroom" / "bin"
|
||||
|
||||
|
||||
def test_rtk_path_suffix(fake_home: Path) -> None:
|
||||
expected_name = "rtk.exe" if os.name == "nt" else "rtk"
|
||||
assert paths.rtk_path().name == expected_name
|
||||
assert paths.rtk_path().parent == paths.bin_dir()
|
||||
|
||||
|
||||
def test_lean_ctx_path_suffix(fake_home: Path) -> None:
|
||||
expected_name = "lean-ctx.exe" if os.name == "nt" else "lean-ctx"
|
||||
assert paths.lean_ctx_path().name == expected_name
|
||||
assert paths.lean_ctx_path().parent == paths.bin_dir()
|
||||
|
||||
|
||||
def test_deploy_root_default(fake_home: Path) -> None:
|
||||
assert paths.deploy_root() == fake_home / ".headroom" / "deploy"
|
||||
|
||||
|
|
@ -381,24 +367,6 @@ def test_proxy_log_path_follows_workspace_env(
|
|||
assert paths.proxy_log_path() == ws / "logs" / "proxy.log"
|
||||
|
||||
|
||||
def test_rtk_path_follows_workspace_env(
|
||||
fake_home: Path, clean_env: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
ws = tmp_path / "state"
|
||||
clean_env.setenv(paths.HEADROOM_WORKSPACE_DIR_ENV, str(ws))
|
||||
expected_name = "rtk.exe" if os.name == "nt" else "rtk"
|
||||
assert paths.rtk_path() == ws / "bin" / expected_name
|
||||
|
||||
|
||||
def test_lean_ctx_path_follows_workspace_env(
|
||||
fake_home: Path, clean_env: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
ws = tmp_path / "state"
|
||||
clean_env.setenv(paths.HEADROOM_WORKSPACE_DIR_ENV, str(ws))
|
||||
expected_name = "lean-ctx.exe" if os.name == "nt" else "lean-ctx"
|
||||
assert paths.lean_ctx_path() == ws / "bin" / expected_name
|
||||
|
||||
|
||||
def test_beacon_lock_path_follows_workspace_env(
|
||||
fake_home: Path, clean_env: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -1,51 +0,0 @@
|
|||
"""``headroom perf`` must surface CLI context-tool (RTK) lifetime savings.
|
||||
|
||||
RTK keeps its savings in its own counter, which never lands in ``proxy.log``,
|
||||
so ``headroom perf`` used to omit them entirely — the report showed only
|
||||
proxy-compression savings. These tests pin that the report (text + JSON)
|
||||
includes the context-tool lifetime savings when available, and degrades
|
||||
cleanly to proxy-only when the tool is absent.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import headroom.proxy.helpers as helpers
|
||||
from headroom.perf import analyzer
|
||||
|
||||
_FAKE_RTK = {
|
||||
"installed": True,
|
||||
"tool": "rtk",
|
||||
"label": "RTK",
|
||||
"lifetime": {"tokens_saved": 26_853_652, "commands": 8000, "savings_pct": 68.9},
|
||||
}
|
||||
|
||||
|
||||
def test_build_perf_summary_includes_cli_filtering(monkeypatch):
|
||||
monkeypatch.setattr(helpers, "_get_context_tool_stats", lambda: _FAKE_RTK)
|
||||
report = analyzer.parse_log_files(last_n_hours=0.0)
|
||||
summary = analyzer.build_perf_summary(report)
|
||||
|
||||
assert summary["cli_filtering"] is not None
|
||||
assert summary["cli_filtering"]["tool"] == "rtk"
|
||||
assert summary["cli_filtering"]["tokens_saved"] == 26_853_652
|
||||
assert summary["cli_filtering"]["savings_pct"] == 68.9
|
||||
|
||||
|
||||
def test_format_report_shows_cli_filtering(monkeypatch):
|
||||
monkeypatch.setattr(helpers, "_get_context_tool_stats", lambda: _FAKE_RTK)
|
||||
report = analyzer.parse_log_files(last_n_hours=0.0)
|
||||
text = analyzer.format_report(report)
|
||||
|
||||
assert "RTK CLI Filtering" in text
|
||||
assert "26,853,652" in text
|
||||
|
||||
|
||||
def test_perf_omits_cli_filtering_when_tool_absent(monkeypatch):
|
||||
monkeypatch.setattr(helpers, "_get_context_tool_stats", lambda: None)
|
||||
report = analyzer.parse_log_files(last_n_hours=0.0)
|
||||
|
||||
summary = analyzer.build_perf_summary(report)
|
||||
assert summary["cli_filtering"] is None
|
||||
|
||||
text = analyzer.format_report(report)
|
||||
assert "CLI Filtering" not in text
|
||||
|
|
@ -1,16 +1,10 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import shutil
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.dashboard import get_dashboard_html
|
||||
from headroom.proxy import helpers as proxy_helpers
|
||||
|
||||
|
||||
class _StatsStub:
|
||||
|
|
@ -30,269 +24,8 @@ class _ToinStub:
|
|||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_rtk_stats_cache(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
monkeypatch.delenv("HEADROOM_RTK_GAIN_SCOPE", raising=False)
|
||||
def _stub_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("HEADROOM_REQUIRE_RUST_CORE", "false")
|
||||
proxy_helpers._rtk_stats_cache.update(
|
||||
{"expires_at": 0.0, "has_value": False, "tool": None, "value": None}
|
||||
)
|
||||
proxy_helpers._rtk_session_baseline.update(
|
||||
{
|
||||
"initialized": False,
|
||||
"tool": None,
|
||||
"total_commands": 0,
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
"tokens_saved": 0,
|
||||
"total_time_ms": 0,
|
||||
"captured_at": 0.0,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_get_rtk_stats_memoizes_subprocess_calls(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
|
||||
now = {"value": 100.0}
|
||||
calls = {"run": 0}
|
||||
totals = [
|
||||
{
|
||||
"total_commands": 7,
|
||||
"total_input": 2000,
|
||||
"total_output": 766,
|
||||
"total_saved": 1234,
|
||||
"avg_savings_pct": 61.7,
|
||||
"total_time_ms": 700,
|
||||
},
|
||||
{
|
||||
"total_commands": 9,
|
||||
"total_input": 2600,
|
||||
"total_output": 1100,
|
||||
"total_saved": 1500,
|
||||
"avg_savings_pct": 57.69,
|
||||
"total_time_ms": 1000,
|
||||
},
|
||||
]
|
||||
|
||||
def _fake_run(args, **kwargs):
|
||||
calls["run"] += 1
|
||||
assert [str(args[0]).replace("\\", "/")] + args[1:] == [
|
||||
"/usr/bin/rtk",
|
||||
"gain",
|
||||
"--format",
|
||||
"json",
|
||||
]
|
||||
summary = totals[min(calls["run"] - 1, len(totals) - 1)]
|
||||
return SimpleNamespace(
|
||||
returncode=0,
|
||||
stdout=json.dumps({"summary": summary}),
|
||||
)
|
||||
|
||||
monkeypatch.setattr(proxy_helpers.time, "monotonic", lambda: now["value"])
|
||||
monkeypatch.setattr(shutil, "which", lambda name: "/usr/bin/rtk")
|
||||
monkeypatch.setattr(subprocess, "run", _fake_run)
|
||||
|
||||
first = proxy_helpers._get_rtk_stats()
|
||||
second = proxy_helpers._get_rtk_stats()
|
||||
|
||||
assert first == second
|
||||
assert first["tool"] == "rtk"
|
||||
assert first["label"] == "RTK"
|
||||
assert first["installed"] is True
|
||||
assert first["scope"] == "global"
|
||||
assert first["total_commands"] == 0
|
||||
assert first["input_tokens"] == 0
|
||||
assert first["output_tokens"] == 0
|
||||
assert first["tokens_saved"] == 0
|
||||
assert first["session_savings_pct"] is None
|
||||
assert first["avg_savings_pct"] == 61.7
|
||||
assert first["avg_savings_pct_scope"] == "lifetime"
|
||||
assert first["lifetime_total_commands"] == 7
|
||||
assert first["lifetime_input_tokens"] == 2000
|
||||
assert first["lifetime_output_tokens"] == 766
|
||||
assert first["lifetime_tokens_saved"] == 1234
|
||||
assert first["session_baseline_total_commands"] == 7
|
||||
assert first["session_baseline_input_tokens"] == 2000
|
||||
assert first["session_baseline_output_tokens"] == 766
|
||||
assert first["session_baseline_tokens_saved"] == 1234
|
||||
assert first["session"]["tokens_saved"] == 0
|
||||
assert first["lifetime"]["savings_pct"] == 61.7
|
||||
assert first["sample_ttl_seconds"] == proxy_helpers.CONTEXT_TOOL_STATS_CACHE_TTL_SECONDS
|
||||
assert calls["run"] == 1
|
||||
|
||||
now["value"] += proxy_helpers.RTK_STATS_CACHE_TTL_SECONDS + 0.1
|
||||
third = proxy_helpers._get_rtk_stats()
|
||||
|
||||
assert third["tool"] == "rtk"
|
||||
assert third["label"] == "RTK"
|
||||
assert third["installed"] is True
|
||||
assert third["total_commands"] == 2
|
||||
assert third["input_tokens"] == 600
|
||||
assert third["output_tokens"] == 334
|
||||
assert third["tokens_saved"] == 266
|
||||
assert third["session_savings_pct"] == pytest.approx(44.3333)
|
||||
assert third["session_avg_time_ms"] == 150.0
|
||||
assert third["lifetime_total_commands"] == 9
|
||||
assert third["lifetime_input_tokens"] == 2600
|
||||
assert third["lifetime_output_tokens"] == 1100
|
||||
assert third["lifetime_tokens_saved"] == 1500
|
||||
assert third["session_baseline_total_commands"] == 7
|
||||
assert third["session_baseline_input_tokens"] == 2000
|
||||
assert third["session_baseline_output_tokens"] == 766
|
||||
assert third["session_baseline_tokens_saved"] == 1234
|
||||
assert third["session"] == {
|
||||
"commands": 2,
|
||||
"input_tokens": 600,
|
||||
"output_tokens": 334,
|
||||
"tokens_saved": 266,
|
||||
"savings_pct": pytest.approx(44.3333),
|
||||
"total_time_ms": 300,
|
||||
"avg_time_ms": 150.0,
|
||||
}
|
||||
assert calls["run"] == 2
|
||||
|
||||
|
||||
def test_get_rtk_stats_can_read_project_scoped_gain(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
calls = {"run": 0}
|
||||
|
||||
def _fake_run(args, **kwargs):
|
||||
calls["run"] += 1
|
||||
assert [str(args[0]).replace("\\", "/")] + args[1:] == [
|
||||
"/usr/bin/rtk",
|
||||
"gain",
|
||||
"--project",
|
||||
"--format",
|
||||
"json",
|
||||
]
|
||||
return SimpleNamespace(
|
||||
returncode=0,
|
||||
stdout=json.dumps(
|
||||
{
|
||||
"summary": {
|
||||
"total_commands": 1,
|
||||
"total_input": 100,
|
||||
"total_output": 75,
|
||||
"total_saved": 25,
|
||||
}
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
monkeypatch.setenv("HEADROOM_RTK_GAIN_SCOPE", "project")
|
||||
monkeypatch.setattr(shutil, "which", lambda name: "/usr/bin/rtk")
|
||||
monkeypatch.setattr(subprocess, "run", _fake_run)
|
||||
|
||||
payload = proxy_helpers._read_rtk_lifetime_stats()
|
||||
|
||||
assert payload is not None
|
||||
assert payload["scope"] == "project"
|
||||
assert payload["total_commands"] == 1
|
||||
assert payload["tokens_saved"] == 25
|
||||
assert calls["run"] == 1
|
||||
|
||||
|
||||
def test_get_rtk_stats_invalid_scope_defaults_to_global(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
calls = {"run": 0}
|
||||
|
||||
def _fake_run(args, **kwargs):
|
||||
calls["run"] += 1
|
||||
assert [str(args[0]).replace("\\", "/")] + args[1:] == [
|
||||
"/usr/bin/rtk",
|
||||
"gain",
|
||||
"--format",
|
||||
"json",
|
||||
]
|
||||
return SimpleNamespace(returncode=0, stdout=json.dumps({"summary": {}}))
|
||||
|
||||
mock_warning = MagicMock()
|
||||
monkeypatch.setenv("HEADROOM_RTK_GAIN_SCOPE", "workspace")
|
||||
monkeypatch.setattr(proxy_helpers.logger, "warning", mock_warning)
|
||||
monkeypatch.setattr(shutil, "which", lambda name: "/usr/bin/rtk")
|
||||
monkeypatch.setattr(subprocess, "run", _fake_run)
|
||||
|
||||
payload = proxy_helpers._read_rtk_lifetime_stats()
|
||||
|
||||
assert payload is not None
|
||||
assert payload["scope"] == "global"
|
||||
assert calls["run"] == 1
|
||||
warning_calls = " ".join(str(call) for call in mock_warning.call_args_list)
|
||||
assert "event=rtk_gain_scope_invalid" in warning_calls
|
||||
|
||||
|
||||
def test_get_context_tool_stats_reads_lean_ctx_gain(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("HEADROOM_CONTEXT_TOOL", "lean-ctx")
|
||||
now = {"value": 100.0}
|
||||
calls = {"run": 0}
|
||||
totals = [
|
||||
{
|
||||
"total_commands": 3,
|
||||
"total_input_tokens": 1000,
|
||||
"total_output_tokens": 600,
|
||||
"tokens_saved": 400,
|
||||
"avg_savings_pct": 40.0,
|
||||
},
|
||||
{
|
||||
"total_commands": 5,
|
||||
"total_input_tokens": 1250,
|
||||
"total_output_tokens": 775,
|
||||
"tokens_saved": 475,
|
||||
"avg_savings_pct": 38.0,
|
||||
},
|
||||
]
|
||||
|
||||
def _fake_run(args, **kwargs):
|
||||
calls["run"] += 1
|
||||
assert [str(args[0]).replace("\\", "/")] + args[1:] == [
|
||||
"/usr/bin/lean-ctx",
|
||||
"gain",
|
||||
"--json",
|
||||
]
|
||||
summary = totals[min(calls["run"] - 1, len(totals) - 1)]
|
||||
return SimpleNamespace(returncode=0, stdout=json.dumps({"summary": summary}))
|
||||
|
||||
monkeypatch.setattr(proxy_helpers.time, "monotonic", lambda: now["value"])
|
||||
monkeypatch.setattr(
|
||||
"headroom.lean_ctx.get_lean_ctx_path",
|
||||
lambda: Path("/usr/bin/lean-ctx"),
|
||||
)
|
||||
monkeypatch.setattr(subprocess, "run", _fake_run)
|
||||
|
||||
first = proxy_helpers._get_context_tool_stats()
|
||||
second = proxy_helpers._get_context_tool_stats()
|
||||
|
||||
assert first == second
|
||||
assert first["tool"] == "lean-ctx"
|
||||
assert first["label"] == "lean-ctx"
|
||||
assert first["installed"] is True
|
||||
assert first["total_commands"] == 0
|
||||
assert first["tokens_saved"] == 0
|
||||
assert first["avg_savings_pct"] == 40.0
|
||||
assert first["session_savings_pct"] is None
|
||||
assert first["lifetime_total_commands"] == 3
|
||||
assert first["lifetime_input_tokens"] == 1000
|
||||
assert first["lifetime_output_tokens"] == 600
|
||||
assert first["lifetime_tokens_saved"] == 400
|
||||
assert calls["run"] == 1
|
||||
|
||||
now["value"] += proxy_helpers.CONTEXT_TOOL_STATS_CACHE_TTL_SECONDS + 0.1
|
||||
third = proxy_helpers._get_context_tool_stats()
|
||||
|
||||
assert third["tool"] == "lean-ctx"
|
||||
assert third["label"] == "lean-ctx"
|
||||
assert third["installed"] is True
|
||||
assert third["total_commands"] == 2
|
||||
assert third["input_tokens"] == 250
|
||||
assert third["output_tokens"] == 175
|
||||
assert third["tokens_saved"] == 75
|
||||
assert third["avg_savings_pct"] == 38.0
|
||||
assert third["avg_savings_pct_scope"] == "lifetime"
|
||||
assert third["session_savings_pct"] == 30.0
|
||||
assert third["lifetime_total_commands"] == 5
|
||||
assert third["lifetime_tokens_saved"] == 475
|
||||
assert third["session"]["savings_pct"] == 30.0
|
||||
assert calls["run"] == 2
|
||||
|
||||
|
||||
def test_stats_cached_query_reuses_short_ttl_snapshot(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
|
|
@ -302,7 +35,7 @@ def test_stats_cached_query_reuses_short_ttl_snapshot(monkeypatch: pytest.Monkey
|
|||
import headroom.proxy.server as server
|
||||
from headroom.proxy.server import ProxyConfig, create_app
|
||||
|
||||
calls = {"store": 0, "telemetry": 0, "feedback": 0, "context_tool": 0}
|
||||
calls = {"store": 0, "telemetry": 0, "feedback": 0}
|
||||
now = {"value": 100.0}
|
||||
|
||||
monkeypatch.setattr(server.time, "monotonic", lambda: now["value"])
|
||||
|
|
@ -322,18 +55,6 @@ def test_stats_cached_query_reuses_short_ttl_snapshot(monkeypatch: pytest.Monkey
|
|||
lambda: _StatsStub(calls, "feedback", {}),
|
||||
)
|
||||
|
||||
def _fake_context_tool_stats() -> dict[str, int | bool | float | str]:
|
||||
calls["context_tool"] += 1
|
||||
return {
|
||||
"tool": "rtk",
|
||||
"label": "RTK",
|
||||
"installed": True,
|
||||
"total_commands": 1,
|
||||
"tokens_saved": 5,
|
||||
"avg_savings_pct": 10.0,
|
||||
}
|
||||
|
||||
monkeypatch.setattr(server, "_get_context_tool_stats", _fake_context_tool_stats)
|
||||
monkeypatch.setattr(server, "get_toin", lambda: _ToinStub())
|
||||
|
||||
app = create_app(
|
||||
|
|
@ -361,274 +82,13 @@ def test_stats_cached_query_reuses_short_ttl_snapshot(monkeypatch: pytest.Monkey
|
|||
assert third.status_code == 200
|
||||
assert uncached.status_code == 200
|
||||
|
||||
assert calls == {"store": 3, "telemetry": 3, "feedback": 3, "context_tool": 3}
|
||||
assert first.json()["context_tool"]["configured"] == "rtk"
|
||||
assert first.json()["context_tool"]["label"] == "RTK"
|
||||
assert first.json()["cli_filtering"]["tokens_saved"] == 5
|
||||
assert first.json()["tokens"]["saved"] == 5
|
||||
assert calls == {"store": 3, "telemetry": 3, "feedback": 3}
|
||||
assert first.json()["tokens"]["proxy_compression_saved"] == 0
|
||||
assert first.json()["tokens"]["cli_filtering_saved"] == 5
|
||||
assert first.json()["tokens"]["rtk_saved"] == 5
|
||||
assert first.json()["tokens"]["lean_ctx_saved"] == 0
|
||||
assert first.json()["tokens"]["all_layers_saved"] == 5
|
||||
assert (
|
||||
first.json()["tokens"]["savings_percent"]
|
||||
== first.json()["tokens"]["all_layers_savings_percent"]
|
||||
)
|
||||
assert first.json()["savings"]["by_layer"]["compression"]["tokens"] == 0
|
||||
assert first.json()["savings"]["by_layer"]["compression"]["cli_filtering_tokens"] == 5
|
||||
assert first.json()["savings"]["by_layer"]["compression"]["rtk_tokens"] == 5
|
||||
assert first.json()["savings"]["by_layer"]["compression"]["lean_ctx_tokens"] == 0
|
||||
assert first.json()["savings"]["by_layer"]["compression"]["all_layers_tokens"] == 5
|
||||
|
||||
|
||||
def test_stats_reports_lean_ctx_as_selected_cli_filter(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
pytest.importorskip("fastapi")
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import headroom.proxy.server as server
|
||||
from headroom.proxy.server import ProxyConfig, create_app
|
||||
|
||||
monkeypatch.setattr(
|
||||
server,
|
||||
"get_compression_store",
|
||||
lambda: _StatsStub({"store": 0}, "store", {}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
server,
|
||||
"get_telemetry_collector",
|
||||
lambda: _StatsStub({"telemetry": 0}, "telemetry", {}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
server,
|
||||
"get_compression_feedback",
|
||||
lambda: _StatsStub({"feedback": 0}, "feedback", {}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
server,
|
||||
"_get_context_tool_stats",
|
||||
lambda: {
|
||||
"tool": "lean-ctx",
|
||||
"label": "lean-ctx",
|
||||
"installed": True,
|
||||
"total_commands": 1,
|
||||
"tokens_saved": 9,
|
||||
"avg_savings_pct": 11.0,
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(server, "get_toin", lambda: _ToinStub())
|
||||
|
||||
app = create_app(
|
||||
ProxyConfig(
|
||||
optimize=False,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
log_requests=False,
|
||||
ccr_inject_tool=False,
|
||||
ccr_handle_responses=False,
|
||||
ccr_context_tracking=False,
|
||||
)
|
||||
)
|
||||
|
||||
with TestClient(app) as client:
|
||||
response = client.get("/stats")
|
||||
|
||||
payload = response.json()
|
||||
assert response.status_code == 200
|
||||
assert payload["context_tool"]["configured"] == "lean-ctx"
|
||||
assert payload["savings"]["by_layer"]["cli_filtering"]["label"] == "lean-ctx"
|
||||
assert payload["tokens"]["cli_filtering_saved"] == 9
|
||||
assert payload["tokens"]["rtk_saved"] == 0
|
||||
assert payload["tokens"]["lean_ctx_saved"] == 9
|
||||
assert payload["savings"]["by_layer"]["compression"]["rtk_tokens"] == 0
|
||||
assert payload["savings"]["by_layer"]["compression"]["lean_ctx_tokens"] == 9
|
||||
|
||||
|
||||
def test_stats_cli_filtering_available_false_when_not_installed(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Reproduction: savings.by_layer.cli_filtering.available reflects `installed`
|
||||
when the context tool isn't installed. On origin/main, `available` doesn't
|
||||
exist in this dict at all (`KeyError`); this asserts the fixed key/value.
|
||||
"""
|
||||
pytest.importorskip("fastapi")
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import headroom.proxy.server as server
|
||||
from headroom.proxy.server import ProxyConfig, create_app
|
||||
|
||||
monkeypatch.setattr(
|
||||
server,
|
||||
"get_compression_store",
|
||||
lambda: _StatsStub({"store": 0}, "store", {}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
server,
|
||||
"get_telemetry_collector",
|
||||
lambda: _StatsStub({"telemetry": 0}, "telemetry", {}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
server,
|
||||
"get_compression_feedback",
|
||||
lambda: _StatsStub({"feedback": 0}, "feedback", {}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
server,
|
||||
"_get_context_tool_stats",
|
||||
lambda: {
|
||||
"tool": "rtk",
|
||||
"label": "RTK",
|
||||
"installed": False,
|
||||
"total_commands": 0,
|
||||
"tokens_saved": 0,
|
||||
"avg_savings_pct": 0.0,
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(server, "get_toin", lambda: _ToinStub())
|
||||
|
||||
app = create_app(
|
||||
ProxyConfig(
|
||||
optimize=False,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
log_requests=False,
|
||||
ccr_inject_tool=False,
|
||||
ccr_handle_responses=False,
|
||||
ccr_context_tracking=False,
|
||||
)
|
||||
)
|
||||
|
||||
with TestClient(app) as client:
|
||||
response = client.get("/stats")
|
||||
|
||||
payload = response.json()
|
||||
assert response.status_code == 200
|
||||
assert payload["savings"]["by_layer"]["cli_filtering"]["available"] is False
|
||||
# Preservation: context_tool.available keeps matching the same `installed`
|
||||
# value it always did, now computed via the hoisted local.
|
||||
assert payload["context_tool"]["available"] is False
|
||||
|
||||
|
||||
def test_stats_cli_filtering_available_true_at_boundary_zero_tokens_saved(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Boundary value 0.0: installed but genuinely zero savings must still
|
||||
report `available: True` with a real `0`, never collapsing into the
|
||||
"not installed" state. This is the negative-space guard against the fix
|
||||
over-triggering on the #1831 reporter's original zero-figures symptom.
|
||||
"""
|
||||
pytest.importorskip("fastapi")
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import headroom.proxy.server as server
|
||||
from headroom.proxy.server import ProxyConfig, create_app
|
||||
|
||||
monkeypatch.setattr(
|
||||
server,
|
||||
"get_compression_store",
|
||||
lambda: _StatsStub({"store": 0}, "store", {}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
server,
|
||||
"get_telemetry_collector",
|
||||
lambda: _StatsStub({"telemetry": 0}, "telemetry", {}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
server,
|
||||
"get_compression_feedback",
|
||||
lambda: _StatsStub({"feedback": 0}, "feedback", {}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
server,
|
||||
"_get_context_tool_stats",
|
||||
lambda: {
|
||||
"tool": "rtk",
|
||||
"label": "RTK",
|
||||
"installed": True,
|
||||
"total_commands": 0,
|
||||
"tokens_saved": 0,
|
||||
"avg_savings_pct": 0.0,
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(server, "get_toin", lambda: _ToinStub())
|
||||
|
||||
app = create_app(
|
||||
ProxyConfig(
|
||||
optimize=False,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
log_requests=False,
|
||||
ccr_inject_tool=False,
|
||||
ccr_handle_responses=False,
|
||||
ccr_context_tracking=False,
|
||||
)
|
||||
)
|
||||
|
||||
with TestClient(app) as client:
|
||||
response = client.get("/stats")
|
||||
|
||||
payload = response.json()
|
||||
assert response.status_code == 200
|
||||
assert payload["savings"]["by_layer"]["cli_filtering"]["available"] is True
|
||||
assert payload["savings"]["by_layer"]["cli_filtering"]["tokens_saved"] == 0
|
||||
assert payload["context_tool"]["available"] is True
|
||||
|
||||
|
||||
def test_cost_merge_uses_generic_cli_filtering_name() -> None:
|
||||
from headroom.proxy.cost import merge_cost_stats
|
||||
|
||||
payload = merge_cost_stats(
|
||||
{"savings_usd": 1.23456, "other": "kept"},
|
||||
{"totals": {"net_savings_usd": 0.25}},
|
||||
cli_tokens_avoided=12,
|
||||
)
|
||||
|
||||
assert payload is not None
|
||||
assert payload["compression_savings_usd"] == 1.2346
|
||||
assert payload["cache_savings_usd"] == 0.25
|
||||
assert payload["cli_tokens_avoided"] == 12
|
||||
assert payload["cli_filtering_tokens_avoided"] == 12
|
||||
assert payload["cli_filtering_tokens_included_in_compression"] is True
|
||||
assert payload["cli_tokens_included_in_compression"] is True
|
||||
|
||||
|
||||
def test_session_summary_uses_generic_cli_filtering_keys() -> None:
|
||||
from headroom.proxy.cost import build_session_summary
|
||||
|
||||
proxy = SimpleNamespace(
|
||||
config=SimpleNamespace(mode="token"),
|
||||
logger=SimpleNamespace(_logs=[]),
|
||||
cost_tracker=SimpleNamespace(
|
||||
stats=lambda: {
|
||||
"cost_with_headroom_usd": 2.0,
|
||||
"savings_usd": 0.5,
|
||||
}
|
||||
),
|
||||
)
|
||||
metrics = SimpleNamespace(
|
||||
requests_by_model={"gpt-test": 1},
|
||||
tokens_saved_total=20,
|
||||
)
|
||||
|
||||
payload = build_session_summary(
|
||||
proxy,
|
||||
metrics,
|
||||
{"totals": {"net_savings_usd": 0.2}},
|
||||
cli_tokens_avoided=7,
|
||||
total_tokens_before=100,
|
||||
)
|
||||
|
||||
assert payload["compression"]["cli_filtering_tokens_avoided"] == 7
|
||||
assert payload["compression"]["total_tokens_saved_with_cli_filtering"] == 27
|
||||
assert payload["compression"]["total_tokens_before_with_cli_filtering"] == 100
|
||||
assert payload["compression"]["rtk_tokens_avoided"] == 7
|
||||
assert payload["cost"]["breakdown"]["cli_filtering_savings_usd"] is None
|
||||
assert payload["cost"]["breakdown"]["rtk_savings_usd"] is None
|
||||
# Metrics fixture has no codex_ws counters -> no codex_ws block.
|
||||
assert "codex_ws" not in payload
|
||||
# The retired CLI context tools must leave no trace in the payload.
|
||||
payload = first.json()
|
||||
assert "context_tool" not in payload
|
||||
assert "cli_filtering" not in payload
|
||||
assert not any("rtk" in key or "lean_ctx" in key for key in payload["tokens"])
|
||||
|
||||
|
||||
def test_session_summary_surfaces_codex_ws_counters() -> None:
|
||||
|
|
@ -647,13 +107,7 @@ def test_session_summary_surfaces_codex_ws_counters() -> None:
|
|||
codex_ws_unit_tokens_saved_sum=4321,
|
||||
)
|
||||
|
||||
payload = build_session_summary(
|
||||
proxy,
|
||||
metrics,
|
||||
{},
|
||||
cli_tokens_avoided=0,
|
||||
total_tokens_before=0,
|
||||
)
|
||||
payload = build_session_summary(proxy, metrics, {}, total_tokens_before=0)
|
||||
|
||||
assert payload["codex_ws"] == {
|
||||
"units_total": 12,
|
||||
|
|
@ -685,7 +139,6 @@ def test_stats_reset_clears_runtime_proxy_counters(monkeypatch: pytest.MonkeyPat
|
|||
"get_compression_feedback",
|
||||
lambda: _StatsStub({"feedback": 0}, "feedback", {}),
|
||||
)
|
||||
monkeypatch.setattr(server, "_get_context_tool_stats", lambda: None)
|
||||
monkeypatch.setattr(server, "get_toin", lambda: _ToinStub())
|
||||
|
||||
app = create_app(
|
||||
|
|
@ -733,16 +186,18 @@ def test_dashboard_uses_cached_stats_and_lazy_history_feed_polling() -> None:
|
|||
assert '@click="toggleFeed()"' in html
|
||||
assert "this.viewMode === 'history'" in html
|
||||
assert "this.feedOpen" in html
|
||||
assert "CLI Filtering (rtk)" not in html
|
||||
assert "RTK Filtered" not in html
|
||||
assert "|| 'RTK'" not in html
|
||||
assert "rtkShareOfTotal" not in html
|
||||
assert "Lean-ctx" in html
|
||||
assert "Context Tool" in html
|
||||
assert "cliFilteringLabel + ' Filtered (this session)'" in html
|
||||
# Lifetime CLI-filtering savings moved from the session card to the
|
||||
# history tab as "Lifetime Saved" in #2198 (persist lifetime metrics).
|
||||
assert "cliFilteringLabel) + ' Lifetime Saved'" in html
|
||||
# The retired CLI context tools left no panel, label or getter behind.
|
||||
for gone in (
|
||||
"CLI Filtering (rtk)",
|
||||
"RTK Filtered",
|
||||
"|| 'RTK'",
|
||||
"rtkShareOfTotal",
|
||||
"Lean-ctx",
|
||||
"Context Tool",
|
||||
"cliFiltering",
|
||||
"cli_filtering",
|
||||
):
|
||||
assert gone not in html, f"dashboard still references {gone!r}"
|
||||
|
||||
|
||||
def test_dashboard_session_metrics_do_not_repeat_proxy_tokens_without_new_context() -> None:
|
||||
|
|
|
|||
|
|
@ -297,8 +297,7 @@ class TestOpenAIResponsesCompression:
|
|||
assert response.status_code == 200
|
||||
|
||||
stats = openai_responses_client.get("/stats").json()
|
||||
# With bypass, proxy compression should not save tokens. The headline
|
||||
# saved count may include RTK CLI savings from the developer shell.
|
||||
# With bypass, proxy compression should not save tokens.
|
||||
assert stats["tokens"]["proxy_compression_saved"] == 0
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1352,125 +1352,6 @@ def test_dashboard_includes_history_toggle_and_endpoint(tmp_path, monkeypatch):
|
|||
assert "topModels[topModels.length - 1] = selected;" in html
|
||||
|
||||
|
||||
def test_stats_history_includes_cli_filtering(tmp_path, monkeypatch):
|
||||
"""The /stats-history response must include cli_filtering (RTK) lifetime stats.
|
||||
|
||||
Before this fix the endpoint returned only proxy compression data; after a
|
||||
restart the Historical tab showed no RTK savings at all.
|
||||
"""
|
||||
pytest.importorskip("fastapi")
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import headroom.proxy.server as server
|
||||
from headroom.proxy.server import ProxyConfig, create_app
|
||||
|
||||
savings_path = tmp_path / "proxy_savings.json"
|
||||
monkeypatch.setenv("HEADROOM_SAVINGS_PATH", str(savings_path))
|
||||
|
||||
_rtk_lifetime_payload = {
|
||||
"tool": "rtk",
|
||||
"label": "RTK",
|
||||
"tokens_saved": 999,
|
||||
"session": {"tokens_saved": 200, "commands": 5},
|
||||
"lifetime": {"tokens_saved": 999, "commands": 42},
|
||||
}
|
||||
monkeypatch.setattr(server, "_get_context_tool_stats", lambda: _rtk_lifetime_payload)
|
||||
|
||||
config = ProxyConfig(
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
log_requests=False,
|
||||
)
|
||||
|
||||
with TestClient(create_app(config)) as client:
|
||||
response = client.get("/stats-history")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
assert "cli_filtering" in data, "Historical /stats-history must include cli_filtering"
|
||||
assert data["cli_filtering"] is not None
|
||||
assert data["cli_filtering"]["tool"] == "rtk"
|
||||
assert data["cli_filtering"]["label"] == "RTK"
|
||||
assert data["cli_filtering"]["lifetime"]["tokens_saved"] == 999
|
||||
|
||||
|
||||
def test_stats_history_cli_filtering_available_false_when_not_installed(tmp_path, monkeypatch):
|
||||
"""Reproduction: /stats-history's curated cli_filtering block must carry
|
||||
`available` reflecting the backend `installed` flag. On origin/main this
|
||||
key doesn't exist in the curated dict at all (`KeyError`); this asserts
|
||||
the fixed key/value. The tool being merely absent must NOT collapse the
|
||||
block to `None` -- it stays populated with `available: False` and zeroed
|
||||
counters so the Historical tab can distinguish absence from a hard
|
||||
read failure.
|
||||
"""
|
||||
pytest.importorskip("fastapi")
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import headroom.proxy.server as server
|
||||
from headroom.proxy.server import ProxyConfig, create_app
|
||||
|
||||
savings_path = tmp_path / "proxy_savings.json"
|
||||
monkeypatch.setenv("HEADROOM_SAVINGS_PATH", str(savings_path))
|
||||
|
||||
_rtk_not_installed_payload = {
|
||||
"tool": "rtk",
|
||||
"label": "RTK",
|
||||
"installed": False,
|
||||
"tokens_saved": 0,
|
||||
"session": {"tokens_saved": 0, "commands": 0},
|
||||
"lifetime": {"tokens_saved": 0, "commands": 0},
|
||||
}
|
||||
monkeypatch.setattr(server, "_get_context_tool_stats", lambda: _rtk_not_installed_payload)
|
||||
|
||||
config = ProxyConfig(
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
log_requests=False,
|
||||
)
|
||||
|
||||
with TestClient(create_app(config)) as client:
|
||||
response = client.get("/stats-history")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
assert data["cli_filtering"] is not None
|
||||
assert data["cli_filtering"]["available"] is False
|
||||
|
||||
|
||||
def test_stats_history_cli_filtering_stays_none_on_hard_read_failure(tmp_path, monkeypatch):
|
||||
"""Preservation: /stats-history's cli_filtering key stays `None` only when
|
||||
the underlying stats read hard-fails (exception), not merely because the
|
||||
tool is absent -- the Historical tab keeps hiding the card in that case,
|
||||
unchanged from prior behavior.
|
||||
"""
|
||||
pytest.importorskip("fastapi")
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import headroom.proxy.server as server
|
||||
from headroom.proxy.server import ProxyConfig, create_app
|
||||
|
||||
savings_path = tmp_path / "proxy_savings.json"
|
||||
monkeypatch.setenv("HEADROOM_SAVINGS_PATH", str(savings_path))
|
||||
|
||||
def _raise() -> dict:
|
||||
raise RuntimeError("simulated hard stats-read failure")
|
||||
|
||||
monkeypatch.setattr(server, "_get_context_tool_stats", _raise)
|
||||
|
||||
config = ProxyConfig(
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
log_requests=False,
|
||||
)
|
||||
|
||||
with TestClient(create_app(config)) as client:
|
||||
response = client.get("/stats-history")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
assert data["cli_filtering"] is None
|
||||
|
||||
|
||||
def test_coercion_helpers_reject_non_finite_values():
|
||||
"""Non-finite inputs fail open to the default -- never raise, never leak NaN/inf.
|
||||
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ from __future__ import annotations
|
|||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from headroom.proxy import server
|
||||
from headroom.proxy.models import ProxyConfig
|
||||
from headroom.proxy.server import create_app
|
||||
|
||||
|
|
@ -166,17 +165,6 @@ def test_stats_recent_requests_includes_token_incomplete_requests() -> None:
|
|||
|
||||
def test_agent_usage_totals_use_proxy_only_savings(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("HEADROOM_REQUIRE_RUST_CORE", "false")
|
||||
monkeypatch.setattr(
|
||||
server,
|
||||
"_get_context_tool_stats",
|
||||
lambda: {
|
||||
"tool": "rtk",
|
||||
"label": "RTK",
|
||||
"tokens_saved": 500,
|
||||
"session": {},
|
||||
"lifetime": {},
|
||||
},
|
||||
)
|
||||
app = create_app(
|
||||
ProxyConfig(
|
||||
optimize=False,
|
||||
|
|
@ -220,7 +208,7 @@ def test_agent_usage_totals_use_proxy_only_savings(monkeypatch: pytest.MonkeyPat
|
|||
assert response.status_code == 200
|
||||
payload = response.json()
|
||||
|
||||
assert payload["tokens"]["saved"] == 600
|
||||
assert payload["tokens"]["saved"] == 100
|
||||
assert payload["agent_usage"]["totals"]["before_tokens"] == 1000
|
||||
assert payload["agent_usage"]["totals"]["tokens_saved"] == 100
|
||||
assert payload["agent_usage"]["totals"]["savings_percent"] == 10.0
|
||||
|
|
|
|||
|
|
@ -1,45 +0,0 @@
|
|||
"""Tests for RTK/context-tool availability detection in Docker environments."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from headroom.proxy.helpers import (
|
||||
_context_tool_zero_payload,
|
||||
_read_rtk_lifetime_stats,
|
||||
)
|
||||
|
||||
|
||||
class TestRtkNotInstalledPayload:
|
||||
"""When rtk binary is absent, the payload must report installed=False."""
|
||||
|
||||
def test_zero_payload_marks_not_installed(self) -> None:
|
||||
payload = _context_tool_zero_payload(tool="rtk", installed=False)
|
||||
assert payload["installed"] is False
|
||||
assert payload["total_commands"] == 0
|
||||
assert payload["tokens_saved"] == 0
|
||||
|
||||
def test_read_rtk_returns_not_installed_when_binary_missing(self) -> None:
|
||||
with patch("headroom.rtk.get_rtk_path", return_value=None):
|
||||
result = _read_rtk_lifetime_stats()
|
||||
assert result is not None
|
||||
assert result["installed"] is False
|
||||
assert result["tokens_saved"] == 0
|
||||
|
||||
def test_installed_payload_marks_installed(self) -> None:
|
||||
payload = _context_tool_zero_payload(tool="rtk", installed=True)
|
||||
assert payload["installed"] is True
|
||||
|
||||
|
||||
class TestDashboardAvailabilityFlag:
|
||||
"""The stats endpoint must surface context_tool.available for the dashboard."""
|
||||
|
||||
def test_available_false_when_tool_not_installed(self) -> None:
|
||||
stats = {"installed": False, "tokens_saved": 0}
|
||||
available = bool(stats.get("installed", False))
|
||||
assert available is False
|
||||
|
||||
def test_available_true_when_tool_installed(self) -> None:
|
||||
stats = {"installed": True, "tokens_saved": 42}
|
||||
available = bool(stats.get("installed", False))
|
||||
assert available is True
|
||||
|
|
@ -1,106 +0,0 @@
|
|||
"""Tests for host-target rtk installation overrides."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import stat
|
||||
import tarfile
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
from headroom.rtk import get_rtk_path, installer
|
||||
|
||||
|
||||
def test_get_rtk_path_finds_windows_managed_binary(tmp_path: Path) -> None:
|
||||
managed_dir = tmp_path / ".headroom" / "bin"
|
||||
managed_dir.mkdir(parents=True)
|
||||
managed_path = managed_dir / "rtk.exe"
|
||||
managed_path.write_bytes(b"binary")
|
||||
|
||||
with patch("headroom.rtk.RTK_BIN_DIR", managed_dir):
|
||||
with patch("headroom.rtk.RTK_BIN_PATH", managed_dir / "rtk"):
|
||||
with patch("headroom.rtk.shutil.which", return_value=None):
|
||||
assert get_rtk_path() == managed_path
|
||||
|
||||
|
||||
def test_get_target_triple_uses_override(monkeypatch) -> None:
|
||||
monkeypatch.setenv("HEADROOM_RTK_TARGET", "x86_64-pc-windows-msvc")
|
||||
assert installer._get_target_triple() == "x86_64-pc-windows-msvc"
|
||||
|
||||
|
||||
def test_download_rtk_skips_verify_for_non_native_target(monkeypatch, tmp_path: Path) -> None:
|
||||
archive = io.BytesIO()
|
||||
with tarfile.open(fileobj=archive, mode="w:gz") as tf:
|
||||
info = tarfile.TarInfo(name="rtk")
|
||||
payload = b"fake-binary"
|
||||
info.size = len(payload)
|
||||
tf.addfile(info, io.BytesIO(payload))
|
||||
archive_bytes = archive.getvalue()
|
||||
|
||||
class _Response:
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def read(self) -> bytes:
|
||||
return archive_bytes
|
||||
|
||||
monkeypatch.setenv("HEADROOM_RTK_TARGET", "x86_64-apple-darwin")
|
||||
|
||||
with patch.object(installer, "RTK_BIN_DIR", tmp_path):
|
||||
with patch.object(installer, "urlopen", return_value=_Response()):
|
||||
with patch.object(installer.subprocess, "run") as subprocess_run:
|
||||
installed_path = installer.download_rtk("v0.42.4")
|
||||
|
||||
assert installed_path == tmp_path / "rtk"
|
||||
assert installed_path.exists()
|
||||
subprocess_run.assert_not_called()
|
||||
|
||||
|
||||
def test_register_claude_hooks_survives_forked_daemon(tmp_path: Path) -> None:
|
||||
"""rtk init that exits fast but leaves a child holding stdout must not hang.
|
||||
|
||||
Regression: capturing through pipes made subprocess.run drain until EOF,
|
||||
which a lingering grandchild deferred past the 10s timeout even though the
|
||||
hooks were already registered. Output now goes to a temp file, so we wait
|
||||
only on the direct child.
|
||||
"""
|
||||
fake_rtk = tmp_path / "rtk"
|
||||
fake_rtk.write_text("#!/bin/bash\n( sleep 30 ) &\necho done\nexit 0\n")
|
||||
fake_rtk.chmod(fake_rtk.stat().st_mode | stat.S_IEXEC)
|
||||
|
||||
assert installer.register_claude_hooks(fake_rtk) is True
|
||||
|
||||
|
||||
def test_register_agent_hooks_passes_agent_flag_for_non_claude(tmp_path: Path, monkeypatch) -> None:
|
||||
calls: list[list[str]] = []
|
||||
|
||||
class FakeResult:
|
||||
returncode = 0
|
||||
|
||||
def fake_run(args, **kwargs):
|
||||
calls.append(args)
|
||||
return FakeResult()
|
||||
|
||||
monkeypatch.setattr(installer.subprocess, "run", fake_run)
|
||||
|
||||
assert installer.register_agent_hooks(Path("rtk"), agent="cursor") is True
|
||||
assert calls == [["rtk", "init", "--global", "--auto-patch", "--agent", "cursor"]]
|
||||
|
||||
|
||||
def test_register_agent_hooks_omits_agent_flag_for_claude(tmp_path: Path, monkeypatch) -> None:
|
||||
calls: list[list[str]] = []
|
||||
|
||||
class FakeResult:
|
||||
returncode = 0
|
||||
|
||||
def fake_run(args, **kwargs):
|
||||
calls.append(args)
|
||||
return FakeResult()
|
||||
|
||||
monkeypatch.setattr(installer.subprocess, "run", fake_run)
|
||||
|
||||
assert installer.register_agent_hooks(Path("rtk"), agent="claude") is True
|
||||
assert calls == [["rtk", "init", "--global", "--auto-patch"]]
|
||||
|
|
@ -1,309 +0,0 @@
|
|||
"""Session RTK savings must be the delta from the proxy-startup baseline.
|
||||
|
||||
Regression for the scope-mixing bug: the dashboard's *session* RTK number must
|
||||
be computed from token deltas since the baseline pinned at proxy startup — NOT
|
||||
from RTK's lifetime average (which dilutes a 62%-this-session rate down to an
|
||||
18.5% all-time number). This exercises the real ``_get_context_tool_stats()``
|
||||
plumbing rather than asserting the arithmetic in the abstract.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import headroom.proxy.helpers as helpers
|
||||
|
||||
|
||||
def _reset(monkeypatch):
|
||||
monkeypatch.delenv(helpers._RTK_GAIN_SCOPE_ENV, raising=False)
|
||||
monkeypatch.setenv("HEADROOM_CONTEXT_TOOL", "rtk")
|
||||
helpers._context_tool_stats_cache.update(
|
||||
{"expires_at": 0.0, "has_value": False, "tool": None, "value": None}
|
||||
)
|
||||
helpers._context_tool_session_baseline.update(
|
||||
{
|
||||
"initialized": False,
|
||||
"tool": None,
|
||||
"total_commands": 0,
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
"tokens_saved": 0,
|
||||
"total_time_ms": 0,
|
||||
"captured_at": 0.0,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _bust_cache():
|
||||
helpers._context_tool_stats_cache.update(
|
||||
{"expires_at": 0.0, "has_value": False, "tool": None, "value": None}
|
||||
)
|
||||
|
||||
|
||||
def test_session_savings_is_delta_not_lifetime_average(monkeypatch):
|
||||
_reset(monkeypatch)
|
||||
|
||||
state: dict = {"summary": None}
|
||||
|
||||
def fake_lifetime(tool):
|
||||
return helpers._context_tool_summary_payload(
|
||||
tool="rtk", installed=True, scope="global", summary=state["summary"]
|
||||
)
|
||||
|
||||
monkeypatch.setattr(helpers, "_read_context_tool_lifetime_stats", fake_lifetime)
|
||||
|
||||
# First poll pins the baseline to the current lifetime → session delta is 0,
|
||||
# but the lifetime number is preserved untouched.
|
||||
state["summary"] = {"total_input": 1000, "total_output": 400, "total_saved": 600}
|
||||
first = helpers._get_context_tool_stats()
|
||||
assert first is not None
|
||||
assert first["session"]["tokens_saved"] == 0
|
||||
assert first["lifetime"]["tokens_saved"] == 600
|
||||
|
||||
# Lifetime advances (more RTK commands run this session); the session number
|
||||
# is the DELTA, not the 800 lifetime total.
|
||||
_bust_cache()
|
||||
state["summary"] = {"total_input": 1300, "total_output": 500, "total_saved": 800}
|
||||
second = helpers._get_context_tool_stats()
|
||||
assert second["session"]["tokens_saved"] == 200 # 800 - 600
|
||||
assert second["lifetime"]["tokens_saved"] == 800
|
||||
# Session % is derived from the delta (200 saved / 300 input delta), not the
|
||||
# lifetime-diluted average.
|
||||
assert second["session"]["savings_pct"] == round(200 / 300 * 100, 4)
|
||||
|
||||
|
||||
# --- Failure semantics: a failed read is "no data", never a zero counter ---
|
||||
|
||||
|
||||
def _fake_run_raises(*args, **kwargs):
|
||||
import subprocess
|
||||
|
||||
raise subprocess.TimeoutExpired(cmd="rtk", timeout=5)
|
||||
|
||||
|
||||
def test_rtk_reader_returns_none_on_timeout(monkeypatch):
|
||||
import headroom.rtk as rtk_mod
|
||||
|
||||
monkeypatch.setattr(rtk_mod, "get_rtk_path", lambda: "/fake/rtk")
|
||||
monkeypatch.setattr(helpers, "run", _fake_run_raises)
|
||||
assert helpers._read_rtk_lifetime_stats() is None
|
||||
|
||||
|
||||
def test_rtk_reader_returns_none_on_nonzero_exit(monkeypatch, caplog):
|
||||
import logging
|
||||
from types import SimpleNamespace
|
||||
|
||||
import headroom.rtk as rtk_mod
|
||||
|
||||
monkeypatch.setattr(rtk_mod, "get_rtk_path", lambda: "/fake/rtk")
|
||||
monkeypatch.setattr(
|
||||
helpers,
|
||||
"run",
|
||||
lambda *a, **k: SimpleNamespace(returncode=1, stdout="", stderr="boom"),
|
||||
)
|
||||
with caplog.at_level(logging.WARNING):
|
||||
assert helpers._read_rtk_lifetime_stats() is None
|
||||
assert "rtk_stats_subprocess_failed" in caplog.text
|
||||
|
||||
|
||||
def test_rtk_reader_returns_none_on_bad_json(monkeypatch):
|
||||
from types import SimpleNamespace
|
||||
|
||||
import headroom.rtk as rtk_mod
|
||||
|
||||
monkeypatch.setattr(rtk_mod, "get_rtk_path", lambda: "/fake/rtk")
|
||||
monkeypatch.setattr(
|
||||
helpers,
|
||||
"run",
|
||||
lambda *a, **k: SimpleNamespace(returncode=0, stdout="not-json{", stderr=""),
|
||||
)
|
||||
assert helpers._read_rtk_lifetime_stats() is None
|
||||
|
||||
|
||||
def test_rtk_reader_not_installed_keeps_zero_payload(monkeypatch):
|
||||
import headroom.rtk as rtk_mod
|
||||
|
||||
monkeypatch.setattr(rtk_mod, "get_rtk_path", lambda: None)
|
||||
payload = helpers._read_rtk_lifetime_stats()
|
||||
assert payload is not None
|
||||
assert payload["installed"] is False
|
||||
assert payload["tokens_saved"] == 0
|
||||
|
||||
|
||||
def test_lean_ctx_reader_returns_none_on_failure_and_logs(monkeypatch, caplog):
|
||||
import logging
|
||||
|
||||
import headroom.lean_ctx as lean_mod
|
||||
|
||||
monkeypatch.setattr(lean_mod, "get_lean_ctx_path", lambda: "/fake/lean-ctx")
|
||||
monkeypatch.setattr(helpers, "run", _fake_run_raises)
|
||||
with caplog.at_level(logging.WARNING):
|
||||
assert helpers._read_lean_ctx_lifetime_stats() is None
|
||||
assert "stats_subprocess_failed" in caplog.text
|
||||
|
||||
|
||||
# --- Baseline guards: pin only from successful installed-tool reads ---
|
||||
|
||||
|
||||
def _payload(saved: int, *, input_tokens: int = 1000, installed: bool = True):
|
||||
return helpers._context_tool_summary_payload(
|
||||
tool="rtk",
|
||||
installed=installed,
|
||||
scope="global",
|
||||
summary={"total_input": input_tokens, "total_output": 400, "total_saved": saved},
|
||||
)
|
||||
|
||||
|
||||
def _stub_reads(monkeypatch, sequence):
|
||||
calls = {"n": 0}
|
||||
|
||||
def fake(tool):
|
||||
idx = min(calls["n"], len(sequence) - 1)
|
||||
calls["n"] += 1
|
||||
item = sequence[idx]
|
||||
return item() if callable(item) else item
|
||||
|
||||
monkeypatch.setattr(helpers, "_read_context_tool_lifetime_stats", fake)
|
||||
return calls
|
||||
|
||||
|
||||
def test_transient_failure_does_not_repin_baseline_or_inflate_session(monkeypatch):
|
||||
"""The headline regression, end to end through the real reader.
|
||||
|
||||
A transient rtk subprocess failure between two identical successful reads
|
||||
must not re-pin the session baseline; today the reader converts the
|
||||
failure into a synthetic zero payload and recovery reports the tool's
|
||||
entire lifetime as session savings.
|
||||
"""
|
||||
import json as json_mod
|
||||
from types import SimpleNamespace
|
||||
|
||||
import headroom.rtk as rtk_mod
|
||||
|
||||
_reset(monkeypatch)
|
||||
monkeypatch.setattr(rtk_mod, "get_rtk_path", lambda: "/fake/rtk")
|
||||
|
||||
good = json_mod.dumps(
|
||||
{"summary": {"total_input": 1000, "total_output": 400, "total_saved": 600}}
|
||||
)
|
||||
behaviors = [
|
||||
lambda: SimpleNamespace(returncode=0, stdout=good, stderr=""),
|
||||
_fake_run_raises,
|
||||
lambda: SimpleNamespace(returncode=0, stdout=good, stderr=""),
|
||||
]
|
||||
calls = {"n": 0}
|
||||
|
||||
def fake_run(*args, **kwargs):
|
||||
behavior = behaviors[min(calls["n"], len(behaviors) - 1)]
|
||||
calls["n"] += 1
|
||||
return behavior()
|
||||
|
||||
monkeypatch.setattr(helpers, "run", fake_run)
|
||||
|
||||
first = helpers._get_context_tool_stats()
|
||||
assert first["session"]["tokens_saved"] == 0
|
||||
assert first["lifetime"]["tokens_saved"] == 600
|
||||
|
||||
_bust_cache()
|
||||
helpers._get_context_tool_stats()
|
||||
# Baseline survives the failed poll untouched.
|
||||
assert helpers._context_tool_session_baseline["tokens_saved"] == 600
|
||||
|
||||
_bust_cache()
|
||||
recovered = helpers._get_context_tool_stats()
|
||||
# Recovery must NOT report the full lifetime as session savings.
|
||||
assert recovered["session"]["tokens_saved"] == 0
|
||||
assert recovered["counter_reset_detected"] is False
|
||||
|
||||
|
||||
def test_boot_fail_then_poll_fail_never_pins_zero_baseline(monkeypatch):
|
||||
import asyncio
|
||||
|
||||
_reset(monkeypatch)
|
||||
_stub_reads(monkeypatch, [None, None, _payload(600)])
|
||||
|
||||
asyncio.run(helpers.initialize_context_tool_session_baseline())
|
||||
assert helpers._context_tool_session_baseline["initialized"] is False
|
||||
|
||||
_bust_cache()
|
||||
assert helpers._get_context_tool_stats() is None
|
||||
# Lazy-init must not have pinned zeros from the failed poll.
|
||||
assert helpers._context_tool_session_baseline["initialized"] is False
|
||||
|
||||
_bust_cache()
|
||||
recovered = helpers._get_context_tool_stats()
|
||||
assert recovered["session"]["tokens_saved"] == 0
|
||||
assert recovered["lifetime"]["tokens_saved"] == 600
|
||||
|
||||
|
||||
def test_stats_reset_with_failing_read_defers_to_next_success(monkeypatch):
|
||||
import asyncio
|
||||
|
||||
_reset(monkeypatch)
|
||||
_stub_reads(monkeypatch, [_payload(600), None, _payload(650)])
|
||||
|
||||
first = helpers._get_context_tool_stats()
|
||||
assert first["session"]["tokens_saved"] == 0
|
||||
|
||||
# /stats/reset while rtk is down: old baseline dropped, pin deferred.
|
||||
asyncio.run(helpers.initialize_context_tool_session_baseline())
|
||||
assert helpers._context_tool_session_baseline["initialized"] is False
|
||||
|
||||
_bust_cache()
|
||||
after = helpers._get_context_tool_stats()
|
||||
# First successful read after the deferred reset pins fresh: delta 0.
|
||||
assert after["session"]["tokens_saved"] == 0
|
||||
assert after["lifetime"]["tokens_saved"] == 650
|
||||
|
||||
|
||||
def test_genuine_counter_reset_still_repins(monkeypatch):
|
||||
_reset(monkeypatch)
|
||||
_stub_reads(monkeypatch, [_payload(600), _payload(50)])
|
||||
|
||||
helpers._get_context_tool_stats()
|
||||
_bust_cache()
|
||||
second = helpers._get_context_tool_stats()
|
||||
assert second["counter_reset_detected"] is True
|
||||
assert second["session"]["tokens_saved"] == 0
|
||||
assert second["lifetime"]["tokens_saved"] == 50
|
||||
|
||||
|
||||
def test_not_installed_payload_does_not_repin_baseline(monkeypatch):
|
||||
_reset(monkeypatch)
|
||||
_stub_reads(
|
||||
monkeypatch,
|
||||
[_payload(600), _payload(0, input_tokens=0, installed=False), _payload(600)],
|
||||
)
|
||||
|
||||
helpers._get_context_tool_stats()
|
||||
_bust_cache()
|
||||
absent = helpers._get_context_tool_stats()
|
||||
# Tool vanished at resolution time: honest zeros display, baseline intact.
|
||||
assert absent["installed"] is False
|
||||
assert helpers._context_tool_session_baseline["tokens_saved"] == 600
|
||||
|
||||
_bust_cache()
|
||||
back = helpers._get_context_tool_stats()
|
||||
assert back["session"]["tokens_saved"] == 0
|
||||
assert back["counter_reset_detected"] is False
|
||||
|
||||
|
||||
def test_tool_switch_with_failing_first_read_does_not_zero_pin(monkeypatch):
|
||||
_reset(monkeypatch)
|
||||
_stub_reads(monkeypatch, [_payload(600), None])
|
||||
|
||||
helpers._get_context_tool_stats()
|
||||
monkeypatch.setenv("HEADROOM_CONTEXT_TOOL", "lean-ctx")
|
||||
_bust_cache()
|
||||
assert helpers._get_context_tool_stats() is None
|
||||
# Switching tools with a failing first read must not pin a zero baseline
|
||||
# for the new tool.
|
||||
assert helpers._context_tool_session_baseline.get("tool") != "lean-ctx"
|
||||
|
||||
|
||||
def test_failed_poll_caches_none_for_ttl(monkeypatch):
|
||||
_reset(monkeypatch)
|
||||
calls = _stub_reads(monkeypatch, [None])
|
||||
|
||||
assert helpers._get_context_tool_stats() is None
|
||||
assert helpers._get_context_tool_stats() is None
|
||||
# Second call inside the TTL is served from cache — no re-read storm.
|
||||
assert calls["n"] == 1
|
||||
|
|
@ -1,17 +1,14 @@
|
|||
"""Proxy-runtime subprocess calls must decode UTF-8 explicitly.
|
||||
|
||||
On Windows, ``subprocess.run(text=True)`` without an ``encoding`` decodes
|
||||
child output with the console code page (cp1252). rtk's emoji-laden output
|
||||
then kills the reader thread with ``UnicodeDecodeError: 'charmap' codec
|
||||
can't decode byte ...`` (seen in user proxy logs). These tests pin the
|
||||
child output with the console code page (cp1252). A child's emoji-laden
|
||||
output then kills the reader thread with ``UnicodeDecodeError: 'charmap'
|
||||
codec can't decode byte ...`` (seen in user proxy logs). These tests pin the
|
||||
``encoding="utf-8"`` kwarg on every proxy-runtime subprocess call.
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
|
||||
import headroom.lean_ctx
|
||||
import headroom.proxy.helpers as helpers
|
||||
import headroom.rtk
|
||||
from headroom.proxy.interceptors import astgrep
|
||||
|
||||
|
||||
|
|
@ -23,28 +20,6 @@ def _capture_run(captured, returncode=0, stdout='{"summary": {}}'):
|
|||
return fake_run
|
||||
|
||||
|
||||
def test_rtk_stats_subprocess_uses_utf8(monkeypatch):
|
||||
captured: dict = {}
|
||||
monkeypatch.setattr(helpers, "run", _capture_run(captured))
|
||||
monkeypatch.setattr(headroom.rtk, "get_rtk_path", lambda: "/fake/rtk")
|
||||
|
||||
helpers._read_rtk_lifetime_stats()
|
||||
|
||||
assert captured["encoding"] == "utf-8"
|
||||
assert captured["errors"] == "replace"
|
||||
|
||||
|
||||
def test_lean_ctx_stats_subprocess_uses_utf8(monkeypatch):
|
||||
captured: dict = {}
|
||||
monkeypatch.setattr(helpers, "run", _capture_run(captured))
|
||||
monkeypatch.setattr(headroom.lean_ctx, "get_lean_ctx_path", lambda: "/fake/lean-ctx")
|
||||
|
||||
helpers._read_lean_ctx_lifetime_stats()
|
||||
|
||||
assert captured["encoding"] == "utf-8"
|
||||
assert captured["errors"] == "replace"
|
||||
|
||||
|
||||
def test_ast_grep_subprocess_uses_utf8(monkeypatch):
|
||||
captured: dict = {}
|
||||
monkeypatch.setattr(astgrep.subprocess, "run", _capture_run(captured, returncode=1, stdout=""))
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sys
|
||||
import threading
|
||||
from datetime import datetime, timedelta
|
||||
|
|
@ -44,9 +45,6 @@ def _make_snapshot(
|
|||
|
||||
def test_tracker_notify_active_update_and_basic_state(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(SubscriptionTracker, "_load_persisted_state", lambda self: None)
|
||||
# PR-G2: keep the unit test deterministic — do not let
|
||||
# ``update_contribution`` call out to ``rtk gain`` via the proxy helper.
|
||||
monkeypatch.setattr(SubscriptionTracker, "_poll_rtk_delta", lambda self: 0)
|
||||
tracker = SubscriptionTracker(enabled=False)
|
||||
|
||||
assert tracker.is_available() is False
|
||||
|
|
@ -67,7 +65,6 @@ def test_tracker_notify_active_update_and_basic_state(monkeypatch: pytest.Monkey
|
|||
tracker.update_contribution(
|
||||
tokens_submitted=10,
|
||||
tokens_saved_compression=5,
|
||||
tokens_saved_cli_filtering=-1,
|
||||
tokens_saved_cache_reads=3,
|
||||
compression_savings_usd=1.25,
|
||||
cache_savings_usd=-2.0,
|
||||
|
|
@ -75,13 +72,9 @@ def test_tracker_notify_active_update_and_basic_state(monkeypatch: pytest.Monkey
|
|||
contribution = tracker._state.contribution
|
||||
assert contribution.tokens_submitted == 10
|
||||
assert contribution.tokens_saved_compression == 5
|
||||
assert contribution.tokens_saved_cli_filtering == 0
|
||||
assert contribution.tokens_saved_rtk == 0
|
||||
assert contribution.tokens_saved_cache_reads == 3
|
||||
assert contribution.to_dict()["tokens_saved"]["compression"] == 5
|
||||
assert contribution.to_dict()["tokens_saved"]["proxy_compression"] == 5
|
||||
assert contribution.to_dict()["tokens_saved"]["cli_filtering"] == 0
|
||||
assert contribution.to_dict()["tokens_saved"]["rtk"] == 0
|
||||
assert contribution.compression_savings_usd == 1.25
|
||||
assert contribution.cache_savings_usd == 0.0
|
||||
|
||||
|
|
@ -267,46 +260,24 @@ async def test_maybe_poll_runs_transcript_scan_off_event_loop(
|
|||
|
||||
|
||||
def test_persist_and_load_state_round_trip(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# PR-G2: ``update_contribution`` polls RTK by default; pin the helper to
|
||||
# 0 so the round-trip is deterministic.
|
||||
monkeypatch.setattr(SubscriptionTracker, "_poll_rtk_delta", lambda self: 0)
|
||||
|
||||
persist_path = tmp_path / "tracker-state.json"
|
||||
tracker = SubscriptionTracker(persist_path=persist_path)
|
||||
tracker.update_contribution(
|
||||
tokens_submitted=11,
|
||||
tokens_saved_compression=2,
|
||||
tokens_saved_cli_filtering=3,
|
||||
tokens_saved_cache_reads=4,
|
||||
compression_savings_usd=1.5,
|
||||
cache_savings_usd=2.5,
|
||||
)
|
||||
# PR-G2: also write a raw RTK delta directly to assert the persisted
|
||||
# ``rtk_raw`` field round-trips independently of cli_filtering.
|
||||
tracker.update_contribution(tokens_saved_rtk=9)
|
||||
tracker._state.poll_count = 7
|
||||
tracker._persist_state()
|
||||
|
||||
loader = SubscriptionTracker(persist_path=persist_path)
|
||||
assert loader._state.contribution.tokens_submitted == 11
|
||||
assert loader._state.contribution.tokens_saved_compression == 2
|
||||
# PR-G2: the raw counters now round-trip independently of the legacy
|
||||
# dashboard alias.
|
||||
assert loader._state.contribution.tokens_saved_cli_filtering == 3
|
||||
assert loader._state.contribution.tokens_saved_rtk == 9
|
||||
assert loader._state.contribution.tokens_saved_cache_reads == 4
|
||||
# ``compression`` is ``proxy_compression + cli_filtering_saved()`` =
|
||||
# ``2 + max(3, 9)`` = 11 after PR-G2 (was 5 when rtk mirrored
|
||||
# cli_filtering).
|
||||
assert loader._state.contribution.to_dict()["tokens_saved"]["compression"] == 11
|
||||
assert loader._state.contribution.to_dict()["tokens_saved"]["compression"] == 2
|
||||
assert loader._state.contribution.to_dict()["tokens_saved"]["proxy_compression"] == 2
|
||||
# Dashboard ``cli_filtering`` / ``rtk`` keys remain ``max(cli, rtk)``
|
||||
# for legacy display — 9 wins. Raw counters expose the un-aliased
|
||||
# values for the tracker's own round-trip.
|
||||
assert loader._state.contribution.to_dict()["tokens_saved"]["cli_filtering"] == 9
|
||||
assert loader._state.contribution.to_dict()["tokens_saved"]["rtk"] == 9
|
||||
assert loader._state.contribution.to_dict()["tokens_saved"]["cli_filtering_raw"] == 3
|
||||
assert loader._state.contribution.to_dict()["tokens_saved"]["rtk_raw"] == 9
|
||||
assert loader._state.contribution.compression_savings_usd == 1.5
|
||||
assert loader._state.contribution.cache_savings_usd == 2.5
|
||||
assert loader._state.poll_count == 7
|
||||
|
|
@ -317,3 +288,52 @@ def test_persist_and_load_state_round_trip(tmp_path: Path, monkeypatch: pytest.M
|
|||
|
||||
missing = SubscriptionTracker(persist_path=tmp_path / "missing.json")
|
||||
assert missing._state.poll_count == 0
|
||||
|
||||
|
||||
def test_load_state_written_before_cli_context_tools_were_removed(tmp_path: Path) -> None:
|
||||
"""A pre-removal state file must still load; the retired keys are ignored.
|
||||
|
||||
Older releases persisted ``cli_filtering`` / ``cli_filtering_raw`` / ``rtk``
|
||||
/ ``rtk_raw`` counters, and wrote the dashboard-facing ``compression`` as
|
||||
proxy-compression *plus* CLI filtering. Those files are still on users'
|
||||
disks, so the loader must neither raise on the extra keys nor double-count
|
||||
the inflated ``compression`` value — it prefers ``proxy_compression``.
|
||||
"""
|
||||
persist_path = tmp_path / "tracker-state.json"
|
||||
persist_path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"poll_count": 7,
|
||||
"contribution": {
|
||||
"tokens_submitted": 11,
|
||||
"tokens_saved": {
|
||||
# 2 (proxy) + 9 (retired CLI layer) as older code wrote it.
|
||||
"compression": 11,
|
||||
"proxy_compression": 2,
|
||||
"cli_filtering": 9,
|
||||
"cli_filtering_raw": 3,
|
||||
"rtk": 9,
|
||||
"rtk_raw": 9,
|
||||
"cache_reads": 4,
|
||||
"total": 15,
|
||||
},
|
||||
"savings_usd": {"compression": 1.5, "cache": 2.5},
|
||||
},
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
loader = SubscriptionTracker(persist_path=persist_path)
|
||||
|
||||
assert loader._state.poll_count == 7
|
||||
assert loader._state.contribution.tokens_submitted == 11
|
||||
# The raw proxy field wins, so the retired layer's tokens are not counted.
|
||||
assert loader._state.contribution.tokens_saved_compression == 2
|
||||
assert loader._state.contribution.tokens_saved_cache_reads == 4
|
||||
assert loader._state.contribution.compression_savings_usd == 1.5
|
||||
assert loader._state.contribution.cache_savings_usd == 2.5
|
||||
# The retired keys are gone from what the tracker now emits.
|
||||
emitted = loader._state.contribution.to_dict()["tokens_saved"]
|
||||
assert "cli_filtering" not in emitted
|
||||
assert "rtk" not in emitted
|
||||
|
|
|
|||
|
|
@ -1,668 +0,0 @@
|
|||
"""Tests for PR-G2 — RTK ``tokens_saved`` data-plane wiring.
|
||||
|
||||
Phase G of the Headroom realignment retires the dead ``tokens_saved_rtk``
|
||||
field by sourcing it from RTK's own stats endpoint (``rtk gain --format
|
||||
json`` via :func:`headroom.proxy.helpers._get_rtk_stats`) and writing the
|
||||
per-call delta into ``HeadroomContribution.tokens_saved_rtk``.
|
||||
|
||||
PR-G2 remediation (C1): the tracker reads the SESSION-incremental
|
||||
``session.tokens_saved`` field of the helper payload, NOT the raw
|
||||
``lifetime_tokens_saved`` counter. The helper de-baselines per proxy
|
||||
session at startup, so the first poll after process startup correctly
|
||||
reads 0 instead of the entire pre-Headroom RTK history.
|
||||
|
||||
These tests pin the wiring:
|
||||
|
||||
1. The delta is computed correctly across two consecutive
|
||||
:meth:`update_contribution` calls (monotonic session counter advances).
|
||||
2. ``tokens_saved_rtk`` is exactly zero when ``_get_rtk_stats()`` returns
|
||||
``None`` (RTK not installed / not selected).
|
||||
3. ``_last_rtk_tokens_saved`` advances monotonically; deltas are not
|
||||
replayed across calls when the session counter does not move.
|
||||
4. First poll reads 0 when the helper reports a fresh session baseline
|
||||
(the C1 regression fix — previously this poll emitted the entire RTK
|
||||
lifetime as a phantom delta).
|
||||
|
||||
Realignment build constraints honored:
|
||||
|
||||
- No silent fallback: a transient ``_get_rtk_stats()`` exception is
|
||||
structured-logged and yields ``tokens_saved_rtk = 0`` (test 4).
|
||||
- Configurable: ``HEADROOM_RTK_WIRING=disabled`` opts the polling out and
|
||||
produces a clean zero, exercised by ``test_disabled_env_returns_zero``.
|
||||
- Structured logs: each failure path emits a ``event=…`` line; the
|
||||
``caplog`` assertions below pin the log payload so the "no silent
|
||||
fallback" constraint is verified.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
import headroom.subscription.tracker as tracker_module
|
||||
from headroom.subscription.tracker import SubscriptionTracker
|
||||
|
||||
|
||||
def _build_tracker(monkeypatch: pytest.MonkeyPatch) -> SubscriptionTracker:
|
||||
"""Construct a tracker with persistence + multi-worker lock disabled.
|
||||
|
||||
Tests use ``_build_tracker`` to keep persistence side effects out of
|
||||
unit tests and to force the RTK poll lock to "owner" so polling runs.
|
||||
"""
|
||||
|
||||
monkeypatch.setattr(SubscriptionTracker, "_load_persisted_state", lambda self: None)
|
||||
monkeypatch.setattr(SubscriptionTracker, "_try_acquire_rtk_poll_lock", lambda self: True)
|
||||
return SubscriptionTracker(enabled=True)
|
||||
|
||||
|
||||
def _session_payload(tokens_saved: int, *, lifetime: int | None = None) -> dict[str, Any]:
|
||||
"""Build a stats payload mimicking ``_get_context_tool_stats``.
|
||||
|
||||
The tracker reads ``session.tokens_saved``. We always include the
|
||||
lifetime field so we can verify the tracker no longer reads it.
|
||||
"""
|
||||
|
||||
if lifetime is None:
|
||||
lifetime = tokens_saved + 50_000 # arbitrary pre-Headroom history
|
||||
return {
|
||||
"tokens_saved": tokens_saved, # session-incremental (canonical)
|
||||
"lifetime_tokens_saved": lifetime,
|
||||
"session": {"tokens_saved": tokens_saved},
|
||||
"lifetime": {"tokens_saved": lifetime},
|
||||
}
|
||||
|
||||
|
||||
def _stub_rtk_stats(
|
||||
monkeypatch: pytest.MonkeyPatch, payloads: list[dict[str, Any] | None]
|
||||
) -> list[int]:
|
||||
"""Stub ``_get_rtk_stats`` to return ``payloads`` in order.
|
||||
|
||||
Returns a counter list (mutated by the stub) so callers can assert the
|
||||
number of polls.
|
||||
"""
|
||||
|
||||
call_count: list[int] = [0]
|
||||
|
||||
def fake_get_rtk_stats() -> dict[str, Any] | None:
|
||||
idx = call_count[0]
|
||||
call_count[0] += 1
|
||||
if idx >= len(payloads):
|
||||
return payloads[-1]
|
||||
return payloads[idx]
|
||||
|
||||
monkeypatch.setattr(
|
||||
"headroom.proxy.helpers._get_rtk_stats",
|
||||
fake_get_rtk_stats,
|
||||
)
|
||||
return call_count
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 1 — delta computed correctly across two consecutive polls
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_tokens_saved_rtk_populated_from_session_field(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""First call seeds the baseline at the session counter, not lifetime.
|
||||
|
||||
PR-G2 remediation (C1): previously the tracker read
|
||||
``lifetime_tokens_saved`` and emitted the entire pre-Headroom RTK
|
||||
history as a phantom delta on the first poll. After the C1 fix the
|
||||
tracker reads ``session.tokens_saved`` which the helper has already
|
||||
de-baselined per proxy session.
|
||||
"""
|
||||
|
||||
tracker = _build_tracker(monkeypatch)
|
||||
monkeypatch.delenv(tracker_module._RTK_WIRING_ENV, raising=False)
|
||||
_stub_rtk_stats(
|
||||
monkeypatch,
|
||||
[
|
||||
_session_payload(tokens_saved=100, lifetime=50_100),
|
||||
_session_payload(tokens_saved=175, lifetime=50_175),
|
||||
],
|
||||
)
|
||||
|
||||
# First call — session counter is 100 (50 000 lifetime history was
|
||||
# rebaselined by the helper at proxy startup, so we DON'T see it).
|
||||
tracker.update_contribution()
|
||||
contribution_after_first = tracker._state.contribution.tokens_saved_rtk
|
||||
assert contribution_after_first == 100
|
||||
assert tracker._last_rtk_tokens_saved == 100
|
||||
|
||||
# Second call — delta is 175 - 100 = 75; cumulative contribution = 175.
|
||||
tracker.update_contribution()
|
||||
assert tracker._state.contribution.tokens_saved_rtk == 175
|
||||
assert tracker._last_rtk_tokens_saved == 175
|
||||
|
||||
|
||||
def test_first_poll_zero_when_session_baseline_fresh(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""C1 fix verification: a freshly-baselined session yields zero on first poll.
|
||||
|
||||
The helper's session baseline is captured at proxy startup. A brand
|
||||
new proxy with no RTK invocations since startup reports
|
||||
``session.tokens_saved == 0`` even though ``lifetime_tokens_saved``
|
||||
may be enormous (months of accumulated RTK history). The tracker must
|
||||
NOT emit the lifetime as a phantom delta.
|
||||
"""
|
||||
|
||||
tracker = _build_tracker(monkeypatch)
|
||||
monkeypatch.delenv(tracker_module._RTK_WIRING_ENV, raising=False)
|
||||
_stub_rtk_stats(
|
||||
monkeypatch,
|
||||
[
|
||||
# Pre-Headroom lifetime = 50 000 tokens. Helper rebaselines at
|
||||
# startup so session = 0.
|
||||
_session_payload(tokens_saved=0, lifetime=50_000),
|
||||
],
|
||||
)
|
||||
|
||||
tracker.update_contribution()
|
||||
|
||||
assert tracker._state.contribution.tokens_saved_rtk == 0, (
|
||||
"first poll must NOT emit pre-Headroom RTK history as a phantom delta"
|
||||
)
|
||||
assert tracker._last_rtk_tokens_saved == 0
|
||||
|
||||
|
||||
def test_delta_computed_correctly_across_polls(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Three consecutive polls — each adds only the new RTK delta."""
|
||||
|
||||
tracker = _build_tracker(monkeypatch)
|
||||
monkeypatch.delenv(tracker_module._RTK_WIRING_ENV, raising=False)
|
||||
_stub_rtk_stats(
|
||||
monkeypatch,
|
||||
[
|
||||
_session_payload(tokens_saved=0), # baseline at zero
|
||||
_session_payload(tokens_saved=50),
|
||||
_session_payload(tokens_saved=250),
|
||||
],
|
||||
)
|
||||
|
||||
tracker.update_contribution()
|
||||
assert tracker._state.contribution.tokens_saved_rtk == 0
|
||||
assert tracker._last_rtk_tokens_saved == 0
|
||||
|
||||
tracker.update_contribution()
|
||||
assert tracker._state.contribution.tokens_saved_rtk == 50
|
||||
assert tracker._last_rtk_tokens_saved == 50
|
||||
|
||||
tracker.update_contribution()
|
||||
# 50 + (250 - 50) = 250 cumulative; delta on the third call was 200.
|
||||
assert tracker._state.contribution.tokens_saved_rtk == 250
|
||||
assert tracker._last_rtk_tokens_saved == 250
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 2 — ``tokens_saved_rtk = 0`` when stats endpoint returns None
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_rtk_stats_none_yields_zero_delta(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""No RTK selected / installed — contribution stays at zero, no throw."""
|
||||
|
||||
tracker = _build_tracker(monkeypatch)
|
||||
monkeypatch.delenv(tracker_module._RTK_WIRING_ENV, raising=False)
|
||||
_stub_rtk_stats(monkeypatch, [None, None])
|
||||
|
||||
tracker.update_contribution()
|
||||
tracker.update_contribution()
|
||||
|
||||
assert tracker._state.contribution.tokens_saved_rtk == 0
|
||||
assert tracker._last_rtk_tokens_saved == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 3 — monotonic advancement; no replay on flat poll
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_last_rtk_advances_monotonically(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Two polls returning the same session total contribute exactly once."""
|
||||
|
||||
tracker = _build_tracker(monkeypatch)
|
||||
monkeypatch.delenv(tracker_module._RTK_WIRING_ENV, raising=False)
|
||||
_stub_rtk_stats(
|
||||
monkeypatch,
|
||||
[
|
||||
_session_payload(tokens_saved=42),
|
||||
_session_payload(tokens_saved=42), # no movement
|
||||
_session_payload(tokens_saved=42), # still no movement
|
||||
],
|
||||
)
|
||||
|
||||
tracker.update_contribution()
|
||||
assert tracker._state.contribution.tokens_saved_rtk == 42
|
||||
assert tracker._last_rtk_tokens_saved == 42
|
||||
|
||||
tracker.update_contribution()
|
||||
assert tracker._state.contribution.tokens_saved_rtk == 42 # unchanged
|
||||
assert tracker._last_rtk_tokens_saved == 42
|
||||
|
||||
tracker.update_contribution()
|
||||
assert tracker._state.contribution.tokens_saved_rtk == 42
|
||||
assert tracker._last_rtk_tokens_saved == 42
|
||||
|
||||
|
||||
def test_counter_regression_rebaselines_without_negative_delta(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Helper rebaselines the session counter — re-baseline, do not subtract."""
|
||||
|
||||
tracker = _build_tracker(monkeypatch)
|
||||
monkeypatch.delenv(tracker_module._RTK_WIRING_ENV, raising=False)
|
||||
_stub_rtk_stats(
|
||||
monkeypatch,
|
||||
[
|
||||
_session_payload(tokens_saved=500),
|
||||
_session_payload(tokens_saved=100), # regression!
|
||||
_session_payload(tokens_saved=150),
|
||||
],
|
||||
)
|
||||
|
||||
tracker.update_contribution()
|
||||
assert tracker._state.contribution.tokens_saved_rtk == 500
|
||||
assert tracker._last_rtk_tokens_saved == 500
|
||||
|
||||
tracker.update_contribution()
|
||||
# Regression: contribution stays at 500 (no negative subtraction).
|
||||
assert tracker._state.contribution.tokens_saved_rtk == 500
|
||||
# Baseline now points at the new (smaller) session total so subsequent
|
||||
# polls can compute a meaningful delta.
|
||||
assert tracker._last_rtk_tokens_saved == 100
|
||||
|
||||
tracker.update_contribution()
|
||||
# 150 - 100 = 50 new delta; contribution = 500 + 50 = 550.
|
||||
assert tracker._state.contribution.tokens_saved_rtk == 550
|
||||
assert tracker._last_rtk_tokens_saved == 150
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 4 — transient exception in the stats endpoint
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_rtk_stats_exception_zero_delta_no_throw_with_log(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
"""A raised ``_get_rtk_stats()`` is caught, structured-logged, yields 0.
|
||||
|
||||
PR-G2 remediation (H3): pins the loud-log requirement so the
|
||||
"no silent fallback" constraint is verified.
|
||||
"""
|
||||
|
||||
tracker = _build_tracker(monkeypatch)
|
||||
monkeypatch.delenv(tracker_module._RTK_WIRING_ENV, raising=False)
|
||||
|
||||
def boom() -> dict[str, Any] | None:
|
||||
raise RuntimeError("transient subprocess failure")
|
||||
|
||||
monkeypatch.setattr("headroom.proxy.helpers._get_rtk_stats", boom)
|
||||
|
||||
caplog.set_level(logging.WARNING, logger="headroom.subscription.tracker")
|
||||
|
||||
# Must not raise.
|
||||
tracker.update_contribution()
|
||||
|
||||
assert tracker._state.contribution.tokens_saved_rtk == 0
|
||||
assert tracker._last_rtk_tokens_saved == 0
|
||||
assert any(
|
||||
"event=subscription_rtk_stats_fetch_failed" in rec.getMessage() for rec in caplog.records
|
||||
), "expected structured log on RTK stats fetch failure"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 5 — explicit env-var opt-out
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_disabled_env_returns_zero(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""``HEADROOM_RTK_WIRING=disabled`` skips the poll entirely."""
|
||||
|
||||
tracker = _build_tracker(monkeypatch)
|
||||
monkeypatch.setenv(tracker_module._RTK_WIRING_ENV, "disabled")
|
||||
|
||||
polls = _stub_rtk_stats(
|
||||
monkeypatch,
|
||||
[_session_payload(tokens_saved=999)],
|
||||
)
|
||||
|
||||
tracker.update_contribution()
|
||||
|
||||
# Stats endpoint never called when wiring is disabled.
|
||||
assert polls[0] == 0
|
||||
assert tracker._state.contribution.tokens_saved_rtk == 0
|
||||
assert tracker._last_rtk_tokens_saved == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 6 — explicit override from caller (back-compat for callers that
|
||||
# already know the RTK delta out-of-band).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_explicit_rtk_override_skips_poll(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Caller-supplied ``tokens_saved_rtk`` short-circuits the poll."""
|
||||
|
||||
tracker = _build_tracker(monkeypatch)
|
||||
monkeypatch.delenv(tracker_module._RTK_WIRING_ENV, raising=False)
|
||||
|
||||
polls = _stub_rtk_stats(monkeypatch, [_session_payload(tokens_saved=999)])
|
||||
|
||||
tracker.update_contribution(tokens_saved_rtk=17)
|
||||
|
||||
# Stats endpoint not consulted when the caller passes an explicit value.
|
||||
assert polls[0] == 0
|
||||
assert tracker._state.contribution.tokens_saved_rtk == 17
|
||||
assert tracker._last_rtk_tokens_saved == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 7 — cli_filtering decoupled from rtk
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_cli_filtering_no_longer_mirrors_rtk(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Pre-PR-G2 bug: ``cli_filtering`` and ``rtk`` were always equal.
|
||||
|
||||
After PR-G2 they are independent counters fed by separate sources.
|
||||
"""
|
||||
|
||||
tracker = _build_tracker(monkeypatch)
|
||||
monkeypatch.delenv(tracker_module._RTK_WIRING_ENV, raising=False)
|
||||
_stub_rtk_stats(monkeypatch, [_session_payload(tokens_saved=25)])
|
||||
|
||||
tracker.update_contribution(tokens_saved_cli_filtering=8)
|
||||
|
||||
assert tracker._state.contribution.tokens_saved_cli_filtering == 8
|
||||
# rtk comes from the polled delta, not from cli_filtering.
|
||||
assert tracker._state.contribution.tokens_saved_rtk == 25
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 8 (H1) — invalid HEADROOM_RTK_WIRING fails loudly at startup
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_garbage_wiring_env_raises_at_startup(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""PR-G2 remediation (H1, M4): bad env value crashes startup loudly.
|
||||
|
||||
Previously the typo would be silently swallowed at every
|
||||
``update_contribution`` call. Now :func:`configure_subscription_tracker`
|
||||
validates eagerly and raises ``ValueError``.
|
||||
"""
|
||||
|
||||
monkeypatch.setenv(tracker_module._RTK_WIRING_ENV, "garbage")
|
||||
# Reset the singleton so configure() actually runs the validator.
|
||||
monkeypatch.setattr(tracker_module, "_tracker_instance", None)
|
||||
|
||||
with pytest.raises(ValueError, match="HEADROOM_RTK_WIRING"):
|
||||
tracker_module.configure_subscription_tracker(enabled=True)
|
||||
|
||||
|
||||
def test_garbage_wiring_env_logs_loudly_at_runtime(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
"""If env is flipped to garbage AFTER startup, runtime path emits ERROR.
|
||||
|
||||
This is the defence-in-depth tier — startup-validation is the primary
|
||||
barrier (test above) but a env-var rotation could still flip the value
|
||||
mid-run.
|
||||
"""
|
||||
|
||||
tracker = _build_tracker(monkeypatch)
|
||||
# Set garbage AFTER tracker construction so the constructor doesn't see it.
|
||||
monkeypatch.setenv(tracker_module._RTK_WIRING_ENV, "garbage")
|
||||
|
||||
caplog.set_level(logging.ERROR, logger="headroom.subscription.tracker")
|
||||
|
||||
tracker.update_contribution()
|
||||
|
||||
assert tracker._state.contribution.tokens_saved_rtk == 0
|
||||
assert any(
|
||||
rec.levelno >= logging.ERROR and "event=subscription_rtk_invalid_env" in rec.getMessage()
|
||||
for rec in caplog.records
|
||||
), "expected ERROR-level structured log on invalid HEADROOM_RTK_WIRING"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 9 (C2) — restart-seeding behavior: no phantom delta on second process
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_restart_does_not_emit_phantom_delta(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Any
|
||||
) -> None:
|
||||
"""PR-G2 remediation (C2): post-restart first poll must not phantom.
|
||||
|
||||
Scenario:
|
||||
1. Tracker A runs, accumulates ``c.tokens_saved_rtk = 100``, persists.
|
||||
2. Process restarts (tracker B loads from disk).
|
||||
3. First poll on tracker B: helper returns ``session.tokens_saved = 5``
|
||||
(small new value since startup). Delta = 5 - 0 = 5. Cumulative =
|
||||
100 + 5 = 105. NOT 100 + 50 000 (lifetime).
|
||||
|
||||
The C1 fix (read session, not lifetime) inherently dissolves this
|
||||
because the helper rebaselines session counters at every proxy
|
||||
startup. This test verifies that property.
|
||||
"""
|
||||
|
||||
monkeypatch.delenv(tracker_module._RTK_WIRING_ENV, raising=False)
|
||||
monkeypatch.setattr(SubscriptionTracker, "_try_acquire_rtk_poll_lock", lambda self: True)
|
||||
|
||||
persist_path = tmp_path / "state.json"
|
||||
|
||||
# Phase 1 — tracker A runs and persists state with non-zero counters.
|
||||
_stub_rtk_stats(
|
||||
monkeypatch,
|
||||
[_session_payload(tokens_saved=100, lifetime=50_100)],
|
||||
)
|
||||
tracker_a = SubscriptionTracker(persist_path=persist_path, enabled=True)
|
||||
tracker_a.update_contribution()
|
||||
assert tracker_a._state.contribution.tokens_saved_rtk == 100
|
||||
tracker_a._persist_state()
|
||||
|
||||
# Phase 2 — simulate process restart. New tracker loads state from
|
||||
# disk. Helper rebaselines (session counter starts fresh at 5 — only
|
||||
# one RTK invocation since restart).
|
||||
_stub_rtk_stats(
|
||||
monkeypatch,
|
||||
[_session_payload(tokens_saved=5, lifetime=50_105)],
|
||||
)
|
||||
tracker_b = SubscriptionTracker(persist_path=persist_path, enabled=True)
|
||||
# Loaded from disk.
|
||||
assert tracker_b._state.contribution.tokens_saved_rtk == 100
|
||||
# Tracker B's _last_rtk_tokens_saved starts at 0 (correct — the
|
||||
# session baseline was just re-pinned in the helper).
|
||||
assert tracker_b._last_rtk_tokens_saved == 0
|
||||
|
||||
tracker_b.update_contribution()
|
||||
# 100 (loaded) + 5 (new session delta) = 105. NOT 50 100 + anything.
|
||||
assert tracker_b._state.contribution.tokens_saved_rtk == 105
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 10 (M2 + M3) — legacy state file migration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_legacy_state_migrates_rtk_from_cli_filtering(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path: Any,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
"""PR-G2 remediation (M2): pre-G2 state has no ``rtk_raw`` key.
|
||||
|
||||
Pre-G2 the ``rtk`` field silently mirrored ``cli_filtering`` (the
|
||||
exact bug PR-G2 retires). When loading a legacy file we treat the
|
||||
aliased ``rtk`` value as the authoritative rtk_raw so historical
|
||||
accumulation isn't silently zeroed. A migration log line is emitted.
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
persist_path = tmp_path / "legacy.json"
|
||||
persist_path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"contribution": {
|
||||
"tokens_submitted": 50,
|
||||
"tokens_saved": {
|
||||
"proxy_compression": 10,
|
||||
"cli_filtering": 42,
|
||||
"rtk": 42, # pre-G2 alias
|
||||
"cache_reads": 3,
|
||||
# NO rtk_raw / cli_filtering_raw keys (legacy)
|
||||
},
|
||||
"savings_usd": {"compression": 0.0, "cache": 0.0},
|
||||
},
|
||||
"poll_count": 7,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
caplog.set_level(logging.INFO, logger="headroom.subscription.tracker")
|
||||
|
||||
tracker = SubscriptionTracker(persist_path=persist_path, enabled=True)
|
||||
|
||||
# Legacy ``rtk == cli_filtering`` got carried forward into rtk_raw.
|
||||
assert tracker._state.contribution.tokens_saved_rtk == 42
|
||||
assert tracker._state.contribution.tokens_saved_cli_filtering == 42
|
||||
# Migration log emitted.
|
||||
assert any(
|
||||
"event=subscription_state_legacy_load" in rec.getMessage() for rec in caplog.records
|
||||
), "expected legacy migration structured log"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 11 (H2) — helper logs structured warning on subprocess failure
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_rtk_subprocess_failure_logs_structured_warning(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""PR-G2 remediation (H2): synthetic-zero path must log loudly.
|
||||
|
||||
Without this, a broken RTK and a healthy "0 tokens saved" RTK are
|
||||
indistinguishable at the tracker layer.
|
||||
|
||||
Implementation note: earlier attempts used pytest's ``caplog`` fixture
|
||||
(both scoped to ``logger="headroom.proxy"`` and root-level capture).
|
||||
Both passed locally but failed in CI — likely a logger-propagation /
|
||||
handler-config difference in the CI test harness. The robust approach
|
||||
is to mock ``_helpers.logger.warning`` directly: when the production
|
||||
code calls ``logger.warning(...)`` the mock intercepts regardless of
|
||||
propagation, formatters, or handler order.
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import headroom.rtk as _rtk
|
||||
from headroom.proxy import helpers as _helpers
|
||||
|
||||
# Point get_rtk_path at a definitely-nonexistent absolute path so the
|
||||
# real ``subprocess.run`` raises FileNotFoundError → except branch
|
||||
# fires the structured warning.
|
||||
monkeypatch.setattr(_rtk, "get_rtk_path", lambda: "/nonexistent/headroom-test-rtk")
|
||||
|
||||
mock_warning = MagicMock()
|
||||
monkeypatch.setattr(_helpers.logger, "warning", mock_warning)
|
||||
|
||||
# Failed reads return None ("no data") rather than a synthetic zero
|
||||
# payload — the zero re-pinned the session baseline and inflated session
|
||||
# savings by the tool's whole lifetime on recovery.
|
||||
payload = _helpers._read_rtk_lifetime_stats()
|
||||
assert payload is None
|
||||
|
||||
# Concatenate all warning call args so the failure message shows what
|
||||
# the helper actually emitted (debug aid for CI flakes).
|
||||
all_warning_calls = " ".join(str(call) for call in mock_warning.call_args_list)
|
||||
assert "event=rtk_stats_subprocess_failed" in all_warning_calls, (
|
||||
f"expected structured warning; actual logger.warning calls: {mock_warning.call_args_list}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 12 (C3) — multi-worker poll deduplication via file lock
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_multi_worker_only_one_polls(monkeypatch: pytest.MonkeyPatch, tmp_path: Any) -> None:
|
||||
"""PR-G2 remediation (C3): two trackers sharing a state path elect one owner.
|
||||
|
||||
The owner polls; the non-owner returns 0 from ``_poll_rtk_delta``.
|
||||
Without this gate each worker would add the same RTK delta to its
|
||||
own ``c.tokens_saved_rtk``, inflating dashboard savings by N× workers.
|
||||
"""
|
||||
|
||||
monkeypatch.delenv(tracker_module._RTK_WIRING_ENV, raising=False)
|
||||
|
||||
# Both trackers share a state directory so they share the lock file.
|
||||
persist_path = tmp_path / "state.json"
|
||||
lock_path = tmp_path / ".rtk_poll_lock"
|
||||
monkeypatch.setenv(tracker_module._RTK_POLL_LOCK_ENV, str(lock_path))
|
||||
|
||||
_stub_rtk_stats(
|
||||
monkeypatch,
|
||||
[_session_payload(tokens_saved=100)],
|
||||
)
|
||||
|
||||
# Worker A — first to attempt acquisition wins.
|
||||
tracker_a = SubscriptionTracker(persist_path=persist_path, enabled=True)
|
||||
# Worker B — same lock path; flock will fail.
|
||||
tracker_b = SubscriptionTracker(persist_path=persist_path, enabled=True)
|
||||
|
||||
tracker_a.update_contribution()
|
||||
tracker_b.update_contribution()
|
||||
|
||||
# Owner polled and got 100; non-owner returned 0.
|
||||
a_rtk = tracker_a._state.contribution.tokens_saved_rtk
|
||||
b_rtk = tracker_b._state.contribution.tokens_saved_rtk
|
||||
# One worker saw the full 100; the other saw 0. (Order is OS-dependent
|
||||
# but exactly one owns the lock.)
|
||||
assert {a_rtk, b_rtk} == {0, 100}, (
|
||||
f"expected exactly one worker to poll; got a={a_rtk}, b={b_rtk}"
|
||||
)
|
||||
|
||||
# Cleanup so subsequent tests don't see a stale lock.
|
||||
tracker_a._release_rtk_poll_lock()
|
||||
tracker_b._release_rtk_poll_lock()
|
||||
|
||||
|
||||
def test_rtk_stats_mid_window_failure_preserves_high_water_mark(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A failed poll (None) mid-window must not reset the high-water mark.
|
||||
|
||||
Failed stat reads now arrive as None ("no data"); the recovery poll's
|
||||
delta is computed against the preserved mark, so no phantom contribution
|
||||
lands and nothing is lost.
|
||||
"""
|
||||
|
||||
tracker = _build_tracker(monkeypatch)
|
||||
monkeypatch.delenv(tracker_module._RTK_WIRING_ENV, raising=False)
|
||||
_stub_rtk_stats(monkeypatch, [_session_payload(100), None, _session_payload(150)])
|
||||
|
||||
tracker.update_contribution()
|
||||
assert tracker._state.contribution.tokens_saved_rtk == 100
|
||||
assert tracker._last_rtk_tokens_saved == 100
|
||||
|
||||
tracker.update_contribution()
|
||||
# Outage poll: zero contribution, mark preserved.
|
||||
assert tracker._state.contribution.tokens_saved_rtk == 100
|
||||
assert tracker._last_rtk_tokens_saved == 100
|
||||
|
||||
tracker.update_contribution()
|
||||
# Recovery: only the true delta lands.
|
||||
assert tracker._state.contribution.tokens_saved_rtk == 150
|
||||
assert tracker._last_rtk_tokens_saved == 150
|
||||
|
|
@ -165,17 +165,6 @@ class TestProxyCLITelemetryBanner:
|
|||
|
||||
assert "HEADROOM_TELEMETRY=off" in result.output or "--no-telemetry" in result.output
|
||||
|
||||
def test_banner_shows_context_tool(self, runner, monkeypatch):
|
||||
monkeypatch.setenv("HEADROOM_CONTEXT_TOOL", "lean-ctx")
|
||||
|
||||
from headroom.cli.main import main
|
||||
|
||||
with patch("headroom.proxy.server.run_server", side_effect=SystemExit(0)):
|
||||
result = runner.invoke(main, ["proxy"])
|
||||
|
||||
assert result.exit_code == 0
|
||||
assert "Context Tool: lean-ctx" in result.output
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# wrap CLI telemetry notice
|
||||
|
|
|
|||
|
|
@ -1,47 +0,0 @@
|
|||
"""RTK is opt-in (off by default): enabled only via --rtk / HEADROOM_RTK=1.
|
||||
|
||||
Regression for the RTK-default flip: the three RTK entry points must no-op
|
||||
unless explicitly opted in, and every wrap subcommand must expose --rtk.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from unittest.mock import patch
|
||||
|
||||
from click.testing import CliRunner
|
||||
|
||||
from headroom.cli import wrap
|
||||
|
||||
|
||||
def _no_rtk_env() -> dict[str, str]:
|
||||
env = dict(os.environ)
|
||||
env.pop("HEADROOM_RTK", None)
|
||||
return env
|
||||
|
||||
|
||||
def test_rtk_opt_in_off_by_default() -> None:
|
||||
with patch.dict(os.environ, _no_rtk_env(), clear=True):
|
||||
assert wrap._rtk_opt_in() is False
|
||||
|
||||
|
||||
def test_rtk_opt_in_on_via_env() -> None:
|
||||
for val in ("1", "true", "yes", "on"):
|
||||
with patch.dict(os.environ, {"HEADROOM_RTK": val}):
|
||||
assert wrap._rtk_opt_in() is True
|
||||
|
||||
|
||||
def test_rtk_entry_points_noop_when_not_opted_in(tmp_path) -> None:
|
||||
agents = tmp_path / "AGENTS.md"
|
||||
with patch.dict(os.environ, _no_rtk_env(), clear=True):
|
||||
assert wrap._setup_rtk() is None
|
||||
assert wrap._ensure_rtk_binary() is None
|
||||
assert wrap._inject_rtk_instructions(agents) is False
|
||||
assert not agents.exists() # nothing written when RTK is off
|
||||
|
||||
|
||||
def test_rtk_flag_present_on_subcommands() -> None:
|
||||
runner = CliRunner()
|
||||
for tool in ("claude", "codex", "copilot", "aider", "continue"):
|
||||
out = runner.invoke(wrap.wrap, [tool, "--help"]).output
|
||||
assert "--rtk" in out, f"--rtk missing from `wrap {tool} --help`"
|
||||
|
|
@ -741,7 +741,6 @@ headroom wrap claude --port 9999
|
|||
| Option / arg | Default | Meaning |
|
||||
|---|---|---|
|
||||
| `--port`, `-p` | `8787` | Proxy port |
|
||||
| `--no-rtk` | off | Skip `rtk` installation and hook registration |
|
||||
| `--no-proxy` | off | Reuse an existing proxy |
|
||||
| `--learn` | off | Enable live traffic learning |
|
||||
| `--verbose`, `-v` | off | Verbose output |
|
||||
|
|
@ -760,7 +759,6 @@ headroom wrap codex --backend anyllm --anyllm-provider groq
|
|||
| Option / arg | Default | Meaning |
|
||||
|---|---|---|
|
||||
| `--port`, `-p` | `8787` | Proxy port |
|
||||
| `--no-rtk` | off | Skip `rtk` installation and `AGENTS.md` injection |
|
||||
| `--no-proxy` | off | Reuse an existing proxy |
|
||||
| `--learn` | off | Enable live traffic learning |
|
||||
| `--backend` | unset | Proxy backend override |
|
||||
|
|
@ -781,7 +779,6 @@ headroom wrap copilot --backend anyllm --anyllm-provider groq -- --model gpt-4o
|
|||
| Option / arg | Default | Meaning |
|
||||
|---|---|---|
|
||||
| `--port`, `-p` | `8787` | Proxy port |
|
||||
| `--no-rtk` | off | Skip `rtk` installation and GitHub Copilot instructions injection |
|
||||
| `--no-proxy` | off | Reuse an existing proxy |
|
||||
| `--learn` | off | Enable live traffic learning |
|
||||
| `--backend` | unset | Proxy backend override |
|
||||
|
|
@ -805,7 +802,6 @@ headroom wrap aider --backend litellm-vertex --region us-central1
|
|||
| Option / arg | Default | Meaning |
|
||||
|---|---|---|
|
||||
| `--port`, `-p` | `8787` | Proxy port |
|
||||
| `--no-rtk` | off | Skip `rtk` installation and `CONVENTIONS.md` injection |
|
||||
| `--no-proxy` | off | Reuse an existing proxy |
|
||||
| `--learn` | off | Enable live traffic learning |
|
||||
| `--backend` | unset | Proxy backend override |
|
||||
|
|
@ -821,13 +817,11 @@ Requires the `aider` binary on the host.
|
|||
```bash
|
||||
headroom wrap cursor
|
||||
headroom wrap cursor --port 9999
|
||||
headroom wrap cursor --no-rtk
|
||||
```
|
||||
|
||||
| Option | Default | Meaning |
|
||||
|---|---|---|
|
||||
| `--port`, `-p` | `8787` | Proxy port |
|
||||
| `--no-rtk` | off | Skip `rtk` installation and `.cursorrules` injection |
|
||||
| `--no-proxy` | off | Reuse an existing proxy |
|
||||
| `--learn` | off | Enable live traffic learning |
|
||||
| `--verbose`, `-v` | off | Verbose output |
|
||||
|
|
|
|||
|
|
@ -72,7 +72,7 @@ docker run --rm -it \
|
|||
`wrap` is host-oriented in Docker-native mode:
|
||||
|
||||
- the wrapper starts the Headroom proxy in Docker
|
||||
- container-side prep writes Headroom config, memory, and `rtk` guidance into mounted host files
|
||||
- container-side prep writes Headroom config and memory into mounted host files
|
||||
- the target CLI itself is launched on the host by the wrapper
|
||||
|
||||
Supported host wrap flows:
|
||||
|
|
|
|||
|
|
@ -60,7 +60,6 @@ Examples:
|
|||
| Memory bridge state | `${WORKSPACE_DIR}/bridge_state.json` | — |
|
||||
| Proxy log directory | `${WORKSPACE_DIR}/logs/` | — |
|
||||
| HTTP 400 debug dumps | `${WORKSPACE_DIR}/logs/debug_400/` | — |
|
||||
| Vendored `rtk` binary | `${WORKSPACE_DIR}/bin/rtk[.exe]` | — |
|
||||
| Deployment profiles | `${WORKSPACE_DIR}/deploy/` | — |
|
||||
| Beacon lock file | `${WORKSPACE_DIR}/.beacon_lock_<port>` | — |
|
||||
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Add a link
Reference in a new issue