From 4b061792b29f776f36831530ba6fdbfd6d71ad0b Mon Sep 17 00:00:00 2001 From: Gili Tzabari Date: Mon, 11 May 2026 16:30:02 -0400 Subject: [PATCH] feat: add lean-ctx context tool support --- README.md | 5 +- docs/content/docs/configuration.mdx | 14 ++ docs/content/docs/docker-install.mdx | 2 +- docs/content/docs/filesystem-contract.mdx | 1 + docs/spec/015-interfaces.md | 2 +- headroom/cli/proxy.py | 20 ++ headroom/cli/wrap.py | 230 ++++++++++++++----- headroom/dashboard/templates/dashboard.html | 28 ++- headroom/lean_ctx/__init__.py | 46 ++++ headroom/lean_ctx/installer.py | 179 +++++++++++++++ headroom/paths.py | 10 + headroom/proxy/cost.py | 28 ++- headroom/proxy/helpers.py | 237 +++++++++++++++++--- headroom/proxy/server.py | 48 +++- headroom/subscription/models.py | 25 ++- headroom/subscription/tracker.py | 17 +- scripts/install.ps1 | 44 +++- scripts/install.sh | 44 +++- tests/test_cli/test_wrap_bridge.py | 118 ++++++++++ tests/test_lean_ctx_installer.py | 176 +++++++++++++++ tests/test_paths.py | 15 ++ tests/test_proxy_dashboard_stats_cache.py | 207 ++++++++++++++++- tests/test_subscription_tracker.py | 10 +- tests/test_telemetry_warning.py | 11 + 24 files changed, 1379 insertions(+), 138 deletions(-) create mode 100644 headroom/lean_ctx/__init__.py create mode 100644 headroom/lean_ctx/installer.py create mode 100644 tests/test_lean_ctx_installer.py diff --git a/README.md b/README.md index 8ce942b85..a00bafb23 100644 --- a/README.md +++ b/README.md @@ -83,7 +83,7 @@ OPENAI_BASE_URL=http://localhost:8787/v1 your-app - **Cross-agent memory and learning.** Claude Code saves a fact, Codex reads it back. `headroom learn` mines failed sessions and writes corrections straight to `CLAUDE.md` / `AGENTS.md` / `GEMINI.md` — reliability compounds over time. - **Reversible (CCR).** Compression is not deletion. The model can always call `headroom_retrieve` to pull the original bytes. Nothing is thrown away. -Bundles the [RTK](https://github.com/rtk-ai/rtk) binary for shell-output rewriting — full [attribution below](#compared-to). +Bundles managed [RTK](https://github.com/rtk-ai/rtk) and [lean-ctx](https://github.com/yvgude/lean-ctx) binaries for local CLI context filtering — full [attribution below](#compared-to). --- @@ -264,10 +264,11 @@ Headroom runs **locally**, covers **every** content type (not just CLI or text), |----------------------------------|-------------------------------------------------|-------------------------------------|:-----:|:----------:| | **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) | CLI commands, MCP tools, editor rules | CLI wrapper · MCP | Yes | No | | [Compresr](https://compresr.ai), [Token Co.](https://thetokencompany.ai) | Text sent to their API | Hosted API call | No | No | | OpenAI Compaction | Conversation history | Provider-native | No | No | -> **Attribution.** Headroom ships with the excellent [RTK](https://github.com/rtk-ai/rtk) binary for shell-output rewriting — `git show` → `git show --short`, noisy `ls` → scoped, chatty installers → summarized. Huge thanks to the RTK team; their tool is a first-class part of our stack, and Headroom compresses everything downstream of it. +> **Attribution.** Headroom ships with the excellent [RTK](https://github.com/rtk-ai/rtk) binary for shell-output rewriting — `git show` → `git show --short`, noisy `ls` → scoped, chatty installers → summarized. Huge thanks to the RTK team; their tool is a first-class part of our stack, and Headroom compresses everything downstream of it. Headroom can also use [lean-ctx](https://github.com/yvgude/lean-ctx) as the selected CLI context tool; set `HEADROOM_CONTEXT_TOOL=lean-ctx` before running `headroom wrap ...`. --- diff --git a/docs/content/docs/configuration.mdx b/docs/content/docs/configuration.mdx index 2ee202482..8027c283c 100644 --- a/docs/content/docs/configuration.mdx +++ b/docs/content/docs/configuration.mdx @@ -5,6 +5,20 @@ 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 ` 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`. + ## Modes | Mode | Behavior | Use Case | diff --git a/docs/content/docs/docker-install.mdx b/docs/content/docs/docker-install.mdx index f61ca7761..b6f351263 100644 --- a/docs/content/docs/docker-install.mdx +++ b/docs/content/docs/docker-install.mdx @@ -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 `rtk` guidance into mounted host files +- container-side prep writes Headroom config, memory, and selected CLI context-tool setup into mounted host files - the target CLI itself is launched on the host by the wrapper Supported host wrap flows: diff --git a/docs/content/docs/filesystem-contract.mdx b/docs/content/docs/filesystem-contract.mdx index 64e90e11e..6aeda5451 100644 --- a/docs/content/docs/filesystem-contract.mdx +++ b/docs/content/docs/filesystem-contract.mdx @@ -55,6 +55,7 @@ Examples: | 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_` | — | diff --git a/docs/spec/015-interfaces.md b/docs/spec/015-interfaces.md index 73b6a619c..20b7069ee 100644 --- a/docs/spec/015-interfaces.md +++ b/docs/spec/015-interfaces.md @@ -100,7 +100,7 @@ headroom wrap [OPTIONS] -- [args...] | Flag | Default | Description | |------|---------|-------------| | `--port` | `8787` | Proxy port | -| `--no-rtk` | `false` | Skip RTK hooks | +| `--no-context-tool` / `--no-rtk` | `false` | Skip CLI context-tool setup | **Supported Commands:** - `claude` — Wrap Claude Code diff --git a/headroom/cli/proxy.py b/headroom/cli/proxy.py index afc473a58..8220aa048 100644 --- a/headroom/cli/proxy.py +++ b/headroom/cli/proxy.py @@ -12,6 +12,11 @@ from headroom.proxy.modes import PROXY_MODE_TOKEN, normalize_proxy_mode from .main import main +_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) @@ -20,6 +25,19 @@ def _get_env_bool(name: str, default: bool) -> bool: return val.lower() in ("true", "1", "yes", "on") +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( "--host", @@ -735,6 +753,7 @@ 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()}" click.echo(f""" ╔═══════════════════════════════════════════════════════════════════════╗ @@ -752,6 +771,7 @@ Starting proxy server... Memory: {memory_status} License: {license_status} {code_aware_line} +{context_tool_line} {extensions_line} {stateless_line}{telemetry_line} {backend_section} diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index 4be5c2976..9639ded25 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -1,13 +1,13 @@ """Wrap CLI commands to run through Headroom proxy. Usage: - headroom wrap claude # Start proxy + rtk + claude + headroom wrap claude # Start proxy + context tool + claude headroom wrap copilot -- --model ... # Start proxy + launch GitHub Copilot CLI headroom wrap codex # Start proxy + OpenAI Codex CLI headroom wrap aider # Start proxy + aider headroom wrap cursor # Start proxy + print Cursor config instructions headroom wrap openclaw # Install + configure OpenClaw plugin - headroom wrap claude --no-rtk # Without rtk hooks + headroom wrap claude --no-context-tool # Without CLI context-tool setup headroom wrap claude --port 9999 # Custom proxy port headroom wrap claude -- --model opus # Pass args to claude """ @@ -22,6 +22,7 @@ import signal import socket import subprocess import sys +import tempfile import time from pathlib import Path from typing import Any, cast @@ -77,12 +78,37 @@ from headroom.providers.openclaw import ( from .main import main +_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 _live_wrap_module() -> Any: """Return the current live wrap module instance.""" return cast(Any, sys.modules[__name__]) +def _selected_context_tool() -> str: + """Return the configured CLI context tool. + + RTK remains the default for backward compatibility. Set + ``HEADROOM_CONTEXT_TOOL=lean-ctx`` to let lean-ctx configure the supported + coding agent instead. + """ + + 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 + + def _print_telemetry_notice() -> None: """Print a telemetry notice when anonymous telemetry is enabled. @@ -244,6 +270,53 @@ def _setup_rtk(verbose: bool = False) -> Path | None: return rtk_path +def _setup_lean_ctx_agent(agent: str, verbose: bool = False) -> Path | None: + """Run lean-ctx agent setup for the requested coding tool.""" + + from headroom.lean_ctx import get_lean_ctx_path + from headroom.lean_ctx.installer import ensure_lean_ctx + + lean_ctx = get_lean_ctx_path() + if not lean_ctx: + click.echo(" Downloading lean-ctx...") + lean_ctx = ensure_lean_ctx() + if not lean_ctx: + click.echo(" lean-ctx download failed — continuing without it") + return None + + try: + with tempfile.TemporaryDirectory(prefix="headroom-lean-ctx-") as setup_cwd: + # lean-ctx writes project-local files when initialized from a git + # checkout. Run from a non-project directory so setup is limited to + # home-scoped agent config such as ~/.codex or ~/.claude. + result = subprocess.run( + [str(lean_ctx), "init", "--agent", agent], + capture_output=True, + text=True, + encoding="utf-8", + errors="replace", + timeout=30, + cwd=setup_cwd, + ) + except Exception as e: + click.echo(f" lean-ctx setup failed — continuing without it: {e}") + return None + + if result.returncode != 0: + detail = (result.stderr or result.stdout).strip() + suffix = f": {detail}" if detail else "" + click.echo(f" lean-ctx setup failed — continuing without it{suffix}") + return None + + if verbose: + detail = result.stdout.strip() + if detail: + click.echo(f" lean-ctx configured for {agent}: {detail}") + else: + click.echo(f" lean-ctx configured for {agent}") + return lean_ctx + + def _remove_claude_rtk_hooks(settings_path: Path | None = None) -> bool: """Remove Headroom/rtk-managed Claude hook entries from settings.json. @@ -1710,7 +1783,13 @@ def unwrap() -> None: @wrap.command(context_settings={"ignore_unknown_options": True}) @click.option("--port", "-p", default=8787, type=int, help="Proxy port (default: 8787)") -@click.option("--no-rtk", is_flag=True, help="Skip rtk installation and hook registration") +@click.option( + "--no-context-tool", + "--no-rtk", + "no_rtk", + is_flag=True, + help="Skip CLI context-tool setup", +) @click.option( "--no-mcp", is_flag=True, @@ -1756,13 +1835,16 @@ def claude( headroom wrap claude --resume # Resume a session headroom wrap claude -- -p # Claude in print mode headroom wrap claude --code-graph # With code graph intelligence - headroom wrap claude --no-rtk # Skip rtk (proxy only) + headroom wrap claude --no-context-tool # Skip CLI context-tool setup headroom wrap claude --no-mcp # Skip MCP retrieve tool registration headroom wrap claude --no-serena # Skip Serena MCP registration """ if prepare_only: if not no_rtk: - _prepare_wrap_rtk(verbose=verbose, label="Claude") + if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX: + _setup_lean_ctx_agent("claude", verbose=verbose) + else: + _prepare_wrap_rtk(verbose=verbose, label="Claude") return claude_bin = shutil.which("claude") @@ -1831,10 +1913,14 @@ def claude( ) if not no_rtk: - click.echo(" Setting up rtk...") - _setup_rtk(verbose=verbose) + if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX: + click.echo(" Setting up lean-ctx...") + _setup_lean_ctx_agent("claude", verbose=verbose) + else: + click.echo(" Setting up rtk...") + _setup_rtk(verbose=verbose) elif verbose: - click.echo(" Skipping rtk (--no-rtk)") + click.echo(" Skipping CLI context tool (--no-context-tool)") if not no_mcp: from headroom.mcp_registry import ClaudeRegistrar @@ -1945,9 +2031,11 @@ def unwrap_claude( @wrap.command(context_settings={"ignore_unknown_options": True}) @click.option("--port", "-p", default=8787, type=int, help="Proxy port (default: 8787)") @click.option( + "--no-context-tool", "--no-rtk", + "no_rtk", is_flag=True, - help="Skip rtk installation and Copilot instructions injection", + help="Skip CLI context-tool setup", ) @click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)") @click.option( @@ -2005,7 +2093,7 @@ def copilot( headroom wrap copilot -- --model claude-sonnet-4-20250514 headroom wrap copilot --backend anyllm --anyllm-provider groq -- --model gpt-4o headroom wrap copilot --provider-type openai --wire-api responses -- --model gpt-5.4 - headroom wrap copilot --no-rtk -- --prompt "explain this file" + headroom wrap copilot --no-context-tool -- --prompt "explain this file" """ copilot_bin = shutil.which("copilot") if not copilot_bin: @@ -2034,11 +2122,15 @@ def copilot( ) if not no_rtk: - click.echo(" Setting up rtk for Copilot...") - rtk_path = _ensure_rtk_binary(verbose=verbose) - if rtk_path: - copilot_instructions = Path.cwd() / ".github" / "copilot-instructions.md" - _inject_rtk_instructions(copilot_instructions, verbose=verbose) + if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX: + click.echo(" Setting up lean-ctx for Copilot...") + _setup_lean_ctx_agent("copilot", verbose=verbose) + else: + click.echo(" Setting up rtk for Copilot...") + rtk_path = _ensure_rtk_binary(verbose=verbose) + if rtk_path: + copilot_instructions = Path.cwd() / ".github" / "copilot-instructions.md" + _inject_rtk_instructions(copilot_instructions, verbose=verbose) env = os.environ.copy() openai_api_url: str | None = None @@ -2118,7 +2210,13 @@ def copilot( @wrap.command(context_settings={"ignore_unknown_options": True}) @click.option("--port", "-p", default=8787, type=int, help="Proxy port (default: 8787)") -@click.option("--no-rtk", is_flag=True, help="Skip rtk installation and AGENTS.md injection") +@click.option( + "--no-context-tool", + "--no-rtk", + "no_rtk", + is_flag=True, + help="Skip CLI context-tool setup", +) @click.option( "--no-mcp", is_flag=True, @@ -2171,16 +2269,16 @@ def codex( \b Sets OPENAI_BASE_URL to route all OpenAI API calls through Headroom. - Installs rtk and injects instructions into AGENTS.md so Codex uses - token-optimized commands (60-90% savings on shell output). Also + Sets up the selected CLI context tool so Codex uses token-optimized + commands (60-90% savings on shell output). Also registers the headroom MCP server in ~/.codex/config.toml so Codex can call ``headroom_retrieve`` on compression markers. \b Examples: - headroom wrap codex # Start proxy + rtk + mcp + codex + headroom wrap codex # Start proxy + context tool + mcp + codex headroom wrap codex -- "fix the bug" # Pass prompt to codex - headroom wrap codex --no-rtk # Skip rtk setup + headroom wrap codex --no-context-tool # Skip CLI context-tool setup headroom wrap codex --no-mcp # Skip MCP retrieve tool registration headroom wrap codex --no-serena # Skip Serena MCP registration headroom wrap codex --port 9999 # Custom proxy port @@ -2195,18 +2293,22 @@ def codex( _codex_config_file, _codex_backup_file = _codex_config_paths() _snapshot_codex_config_if_unwrapped(_codex_config_file, _codex_backup_file) - # Setup rtk for Codex (binary + AGENTS.md instructions, no hooks) + # Setup CLI context tool for Codex. if not no_rtk: - click.echo(" Setting up rtk for Codex...") - rtk_path = _ensure_rtk_binary(verbose=verbose) - if rtk_path: - # Inject into project AGENTS.md (Codex reads this automatically) - agents_md = Path.cwd() / "AGENTS.md" - _inject_rtk_instructions(agents_md, verbose=verbose) + if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX: + click.echo(" Setting up lean-ctx for Codex...") + _setup_lean_ctx_agent("codex", verbose=verbose) + else: + click.echo(" Setting up rtk for Codex...") + rtk_path = _ensure_rtk_binary(verbose=verbose) + if rtk_path: + # Inject into project AGENTS.md (Codex reads this automatically) + agents_md = Path.cwd() / "AGENTS.md" + _inject_rtk_instructions(agents_md, verbose=verbose) - # Also inject into global ~/.codex/AGENTS.md - global_agents = Path.home() / ".codex" / "AGENTS.md" - _inject_rtk_instructions(global_agents, verbose=verbose) + # Also inject into global ~/.codex/AGENTS.md + global_agents = Path.home() / ".codex" / "AGENTS.md" + _inject_rtk_instructions(global_agents, verbose=verbose) # Register headroom MCP server in ~/.codex/config.toml so Codex can # call headroom_retrieve on compression markers from the proxy. @@ -2320,7 +2422,13 @@ def codex( @wrap.command(context_settings={"ignore_unknown_options": True}) @click.option("--port", "-p", default=8787, type=int, help="Proxy port (default: 8787)") -@click.option("--no-rtk", is_flag=True, help="Skip rtk installation and conventions injection") +@click.option( + "--no-context-tool", + "--no-rtk", + "no_rtk", + is_flag=True, + help="Skip CLI context-tool setup", +) @click.option( "--code-graph", is_flag=True, @@ -2355,25 +2463,28 @@ def aider( \b Sets OPENAI_API_BASE to route all API calls through Headroom. - Installs rtk and injects instructions into .aider.conf.yml conventions - so aider uses token-optimized commands. + Sets up the selected CLI context tool so aider uses token-optimized commands. \b Examples: - headroom wrap aider # Start proxy + rtk + aider + headroom wrap aider # Start proxy + context tool + aider headroom wrap aider -- --model gpt-4o # Use GPT-4o headroom wrap aider -- --model claude-sonnet-4 # Use Claude - headroom wrap aider --no-rtk # Skip rtk setup + headroom wrap aider --no-context-tool # Skip CLI context-tool setup headroom wrap aider --backend litellm-vertex --region us-central1 """ - # Setup rtk for aider (binary + CONVENTIONS.md instructions) + # Setup CLI context tool for aider. if not no_rtk: - click.echo(" Setting up rtk for aider...") - rtk_path = _ensure_rtk_binary(verbose=verbose) - if rtk_path: - # aider reads CONVENTIONS.md from project root - conventions = Path.cwd() / "CONVENTIONS.md" - _inject_rtk_instructions(conventions, verbose=verbose) + if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX: + click.echo(" Setting up lean-ctx for aider...") + _setup_lean_ctx_agent("aider", verbose=verbose) + else: + click.echo(" Setting up rtk for aider...") + rtk_path = _ensure_rtk_binary(verbose=verbose) + if rtk_path: + # aider reads CONVENTIONS.md from project root + conventions = Path.cwd() / "CONVENTIONS.md" + _inject_rtk_instructions(conventions, verbose=verbose) if prepare_only: return @@ -2411,7 +2522,13 @@ def aider( @wrap.command(context_settings={"ignore_unknown_options": True}) @click.option("--port", "-p", default=8787, type=int, help="Proxy port (default: 8787)") -@click.option("--no-rtk", is_flag=True, help="Skip rtk installation and .cursorrules injection") +@click.option( + "--no-context-tool", + "--no-rtk", + "no_rtk", + is_flag=True, + help="Skip CLI context-tool setup", +) @click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)") @click.option( "--learn", is_flag=True, help="Enable live traffic learning (patterns saved to .cursor/rules/)" @@ -2432,8 +2549,8 @@ def cursor( \b Cursor reads its API configuration from its settings UI, not from - environment variables. This command starts the proxy, installs rtk - with .cursorrules instructions, and prints the Cursor settings. + environment variables. This command starts the proxy, sets up the selected + CLI context tool, and prints the Cursor settings. \b After running this command, open Cursor and configure: @@ -2441,16 +2558,20 @@ def cursor( \b Example: - headroom wrap cursor # Start proxy + rtk + instructions - headroom wrap cursor --no-rtk # Proxy only, no rtk + headroom wrap cursor # Start proxy + context-tool instructions + headroom wrap cursor --no-context-tool # Proxy only, no CLI context tool headroom wrap cursor --port 9999 # Custom proxy port """ if not no_rtk: - click.echo(" Setting up rtk for Cursor...") - rtk_path = _ensure_rtk_binary(verbose=verbose) - if rtk_path: - cursorrules = Path.cwd() / ".cursorrules" - _inject_rtk_instructions(cursorrules, verbose=verbose) + if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX: + click.echo(" Setting up lean-ctx for Cursor...") + _setup_lean_ctx_agent("cursor", verbose=verbose) + else: + click.echo(" Setting up rtk for Cursor...") + rtk_path = _ensure_rtk_binary(verbose=verbose) + if rtk_path: + cursorrules = Path.cwd() / ".cursorrules" + _inject_rtk_instructions(cursorrules, verbose=verbose) if prepare_only: return @@ -2476,7 +2597,10 @@ def cursor( click.echo(line) if not no_rtk: click.echo() - click.echo(" rtk instructions injected into .cursorrules") + if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX: + click.echo(" lean-ctx configured for Cursor") + else: + click.echo(" rtk instructions injected into .cursorrules") click.echo(" Cursor will use token-optimized commands automatically.") click.echo() click.echo(" Press Ctrl+C to stop the proxy.") diff --git a/headroom/dashboard/templates/dashboard.html b/headroom/dashboard/templates/dashboard.html index 7d44b366a..8d017b1cb 100644 --- a/headroom/dashboard/templates/dashboard.html +++ b/headroom/dashboard/templates/dashboard.html @@ -109,7 +109,7 @@
+ x-text="formatNumber(stats.tokens?.proxy_compression_saved || 0) + ' proxy tokens only; ' + cliFilteringLabel + ' excluded from $'">
@@ -125,7 +125,7 @@
/ - +
@@ -698,8 +698,8 @@
- RTK Filtered - + +
Proxy Removed @@ -1778,10 +1778,26 @@ return (this.stats.tokens?.proxy_compression_saved || 0) / total * 100; }, - get rtkShareOfTotal() { + 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.stats.tokens?.rtk_saved || 0) / total * 100; + return this.cliFilteringSaved / total * 100; }, // --- Compression Confidence --- diff --git a/headroom/lean_ctx/__init__.py b/headroom/lean_ctx/__init__.py new file mode 100644 index 000000000..fcad548c3 --- /dev/null +++ b/headroom/lean_ctx/__init__.py @@ -0,0 +1,46 @@ +"""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 diff --git a/headroom/lean_ctx/installer.py b/headroom/lean_ctx/installer.py new file mode 100644 index 000000000..045e9f828 --- /dev/null +++ b/headroom/lean_ctx/installer.py @@ -0,0 +1,179 @@ +"""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 . 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 = subprocess.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 = subprocess.run( + [str(target_path), "--version"], + capture_output=True, + text=True, + encoding="utf-8", + errors="replace", + 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 diff --git a/headroom/paths.py b/headroom/paths.py index dd24ab261..eb1633f8d 100644 --- a/headroom/paths.py +++ b/headroom/paths.py @@ -75,6 +75,8 @@ _CODEX_WIRE_DEBUG_DIR = "codex_wire" _BIN_DIR = "bin" _RTK_UNIX = "rtk" _RTK_WIN = "rtk.exe" +_LEAN_CTX_UNIX = "lean-ctx" +_LEAN_CTX_WIN = "lean-ctx.exe" _DEPLOY_DIR = "deploy" _PLUGINS_DIR = "plugins" @@ -274,6 +276,13 @@ def rtk_path() -> Path: 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.""" @@ -349,6 +358,7 @@ __all__ = [ "codex_wire_debug_dir", "bin_dir", "rtk_path", + "lean_ctx_path", "deploy_root", "beacon_lock_path", "models_config_path", diff --git a/headroom/proxy/cost.py b/headroom/proxy/cost.py index 0eed221ef..8cc51378a 100644 --- a/headroom/proxy/cost.py +++ b/headroom/proxy/cost.py @@ -300,12 +300,14 @@ def merge_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 rtk (token count only, no $ estimate) + - 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. RTK 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. 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. Prefix cache savings stay separate because they are a provider discount, not token removal. This avoids the non-monotonic moving-average repricing bug (#83). @@ -322,7 +324,9 @@ def merge_cost_stats( "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, } @@ -388,8 +392,8 @@ def build_session_summary( best_detail = f"{best['original']:,} → {best['optimized']:,} tokens" # Cost summary — dollar savings are proxy-compression only at model list - # price. rtk tokens are counted in token savings but have no model-specific - # price because they never reached the proxy request. + # price. CLI filtering tokens are counted in token savings but have no + # model-specific price because they never reached the proxy request. 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) @@ -414,6 +418,11 @@ 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, @@ -427,9 +436,14 @@ 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": ( - "rtk tokens are included in token savings only; dollar savings " + "CLI filtering tokens are included in token savings only; dollar savings " "use proxy compression tokens at model list price." ), }, diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index f346459a2..80c96c5aa 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -13,6 +13,7 @@ import json import logging import os import random +import subprocess import threading import time from collections import OrderedDict @@ -566,18 +567,28 @@ 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_STATS_CACHE_TTL_SECONDS = 5.0 -_rtk_stats_cache_lock = threading.Lock() -_rtk_stats_cache: dict[str, Any] = { +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, } -_rtk_session_baseline: dict[str, Any] = { +_context_tool_session_baseline: dict[str, Any] = { "initialized": False, + "tool": None, "total_commands": 0, "tokens_saved": 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 @@ -803,19 +814,59 @@ 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 _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 _read_rtk_lifetime_stats() -> dict[str, Any] | None: """Read rtk's current project-level lifetime stats.""" - import subprocess as _sp - from headroom.rtk import get_rtk_path rtk_path = get_rtk_path() if not rtk_path: - return None + return { + "tool": _CONTEXT_TOOL_RTK, + "label": _context_tool_label(_CONTEXT_TOOL_RTK), + "installed": False, + "total_commands": 0, + "tokens_saved": 0, + "avg_savings_pct": 0.0, + } try: - result = _sp.run( + result = subprocess.run( [str(rtk_path), "gain", "--project", "--format", "json"], capture_output=True, text=True, @@ -825,13 +876,17 @@ def _read_rtk_lifetime_stats() -> dict[str, Any] | None: data = json.loads(result.stdout) summary = data.get("summary", {}) payload = { + "tool": _CONTEXT_TOOL_RTK, + "label": _context_tool_label(_CONTEXT_TOOL_RTK), "installed": True, - "total_commands": summary.get("total_commands", 0), - "tokens_saved": summary.get("total_saved", 0), - "avg_savings_pct": summary.get("avg_savings_pct", 0.0), + "total_commands": _coerce_int(summary.get("total_commands", 0)), + "tokens_saved": _coerce_int(summary.get("total_saved", 0)), + "avg_savings_pct": _coerce_float(summary.get("avg_savings_pct", 0.0)), } else: return { + "tool": _CONTEXT_TOOL_RTK, + "label": _context_tool_label(_CONTEXT_TOOL_RTK), "installed": True, "total_commands": 0, "tokens_saved": 0, @@ -839,6 +894,8 @@ def _read_rtk_lifetime_stats() -> dict[str, Any] | None: } except Exception: return { + "tool": _CONTEXT_TOOL_RTK, + "label": _context_tool_label(_CONTEXT_TOOL_RTK), "installed": True, "total_commands": 0, "tokens_saved": 0, @@ -848,46 +905,153 @@ def _read_rtk_lifetime_stats() -> dict[str, Any] | None: return payload -def initialize_rtk_session_baseline() -> None: - """Pin the current rtk counters as the proxy-session baseline.""" +def _read_lean_ctx_lifetime_stats() -> dict[str, Any] | None: + """Read lean-ctx's current project-level lifetime stats.""" - payload = _read_rtk_lifetime_stats() - with _rtk_stats_cache_lock: - _rtk_session_baseline.update( + from headroom.lean_ctx import get_lean_ctx_path + + lean_ctx_path = get_lean_ctx_path() + if not lean_ctx_path: + return { + "tool": _CONTEXT_TOOL_LEAN_CTX, + "label": _context_tool_label(_CONTEXT_TOOL_LEAN_CTX), + "installed": False, + "total_commands": 0, + "tokens_saved": 0, + "avg_savings_pct": 0.0, + } + + base_payload = { + "tool": _CONTEXT_TOOL_LEAN_CTX, + "label": _context_tool_label(_CONTEXT_TOOL_LEAN_CTX), + "installed": True, + "total_commands": 0, + "tokens_saved": 0, + "avg_savings_pct": 0.0, + } + + try: + result = subprocess.run( + [str(lean_ctx_path), "gain", "--json"], + capture_output=True, + text=True, + timeout=5, + ) + if result.returncode != 0 or not result.stdout.strip(): + return dict(base_payload) + + data = json.loads(result.stdout) + summary = data.get("summary", data) if isinstance(data, dict) else {} + if not isinstance(summary, dict): + return dict(base_payload) + + return { + **base_payload, + "total_commands": _coerce_int( + _first_value( + summary, + ( + "total_commands", + "commands", + "command_count", + "totalCommandCount", + ), + ) + ), + "tokens_saved": _coerce_int( + _first_value( + summary, + ( + "total_saved", + "tokens_saved", + "total_tokens_saved", + "saved_tokens", + "totalSaved", + ), + ) + ), + "avg_savings_pct": _coerce_float( + _first_value( + summary, + ( + "avg_savings_pct", + "average_savings_pct", + "avgSavingsPct", + "savings_percent", + ), + 0.0, + ) + ), + } + except Exception: + return dict(base_payload) + + +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() + + +def initialize_context_tool_session_baseline() -> None: + """Pin the current context-tool counters as the proxy-session baseline.""" + + tool = _selected_context_tool() + payload = _read_context_tool_lifetime_stats(tool) + with _context_tool_stats_cache_lock: + _context_tool_session_baseline.update( { "initialized": True, + "tool": tool, "total_commands": int((payload or {}).get("total_commands", 0) or 0), "tokens_saved": int((payload or {}).get("tokens_saved", 0) or 0), } ) - _rtk_stats_cache.update( + _context_tool_stats_cache.update( { "expires_at": 0.0, "has_value": False, + "tool": None, "value": None, } ) -def _get_rtk_stats() -> dict[str, Any] | None: - """Get rtk savings for the current Headroom proxy session. +def initialize_rtk_session_baseline() -> None: + """Pin the current context-tool counters as the proxy-session baseline.""" - rtk persists project-level lifetime counters. Dashboard stats should be - session-local, so we subtract the counter snapshot captured at proxy - startup instead of resetting rtk's own history. + 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 _rtk_stats_cache_lock: - if _rtk_stats_cache["has_value"] and now < float(_rtk_stats_cache["expires_at"]): - return cast(dict[str, Any] | None, _rtk_stats_cache["value"]) + 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_rtk_lifetime_stats() - with _rtk_stats_cache_lock: - if not _rtk_session_baseline["initialized"]: - _rtk_session_baseline.update( + payload = _read_context_tool_lifetime_stats(tool) + with _context_tool_stats_cache_lock: + if ( + 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 or {}).get("total_commands", 0) or 0), "tokens_saved": int((payload or {}).get("tokens_saved", 0) or 0), } @@ -896,28 +1060,37 @@ def _get_rtk_stats() -> dict[str, Any] | None: if payload is not None: payload = { **payload, + "tool": tool, + "label": _context_tool_label(tool), "total_commands": max( int(payload.get("total_commands", 0) or 0) - - int(_rtk_session_baseline["total_commands"]), + - int(_context_tool_session_baseline["total_commands"]), 0, ), "tokens_saved": max( int(payload.get("tokens_saved", 0) or 0) - - int(_rtk_session_baseline["tokens_saved"]), + - int(_context_tool_session_baseline["tokens_saved"]), 0, ), } - _rtk_stats_cache.update( + _context_tool_stats_cache.update( { - "expires_at": time.monotonic() + RTK_STATS_CACHE_TTL_SECONDS, + "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"): diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 78fcdd958..30f66d1b0 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -116,11 +116,12 @@ from headroom.proxy.helpers import ( MAX_MESSAGE_ARRAY_LENGTH, # noqa: F401 MAX_REQUEST_BODY_SIZE, # noqa: F401 MAX_SSE_BUFFER_SIZE, # noqa: F401 + _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_rtk_session_baseline, + initialize_context_tool_session_baseline, is_anthropic_auth, # noqa: F401 jitter_delay_ms, ) @@ -1339,7 +1340,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: app.state.started_at = time.time() app.state.ready = False app.state.startup_error = None - initialize_rtk_session_baseline() + initialize_context_tool_session_baseline() try: try: @@ -1700,15 +1701,24 @@ 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 (rtk — tokens avoided before reaching context) - cli_filtering_stats = _get_rtk_stats() + # Fetch CLI filtering savings from the selected context tool. These + # tokens are avoided before they reach model context. + cli_filtering_stats = _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 ) + 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 # Calculate total tokens before Headroom-side reduction. Proxy - # compression and rtk both remove tokens before they reach model - # context, so dashboard-facing compression savings combines them. + # compression and the configured context tool both remove tokens before + # they reach model context, so dashboard-facing savings combines them. proxy_compression_tokens = m.tokens_saved_total all_layers_tokens_saved = proxy_compression_tokens + cli_tokens_avoided total_tokens_before = m.tokens_input_total + all_layers_tokens_saved @@ -1767,21 +1777,27 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: "total_tokens": total_tokens_all_layers, "by_layer": { "cli_filtering": { + "tool": cli_filtering_tool, + "label": cli_filtering_label, "tokens": cli_tokens_avoided, + "tokens_saved": cli_tokens_avoided, "included_in": "tokens.saved", "description": ( - "Tokens avoided by CLI output filtering (rtk) before reaching context. " + 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, - "rtk_tokens": cli_tokens_avoided, + "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 rtk CLI filtering." + "Dashboard token savings also includes CLI context-tool filtering." ), }, "prefix_cache": { @@ -1808,7 +1824,9 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: "output": m.tokens_output_total, "saved": all_layers_tokens_saved, "proxy_compression_saved": proxy_compression_tokens, - "rtk_saved": cli_tokens_avoided, + "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, @@ -1904,6 +1922,14 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: ), }, "toin": get_toin().get_stats(), + "context_tool": { + "configured": cli_filtering_tool, + "label": cli_filtering_label, + "available": bool( + cli_filtering_stats and cli_filtering_stats.get("installed", False) + ), + "stats": cli_filtering_stats, + }, "cli_filtering": cli_filtering_stats, "cache": await proxy.cache.stats() if proxy.cache else None, "rate_limiter": await proxy.rate_limiter.stats() if proxy.rate_limiter else None, @@ -1956,7 +1982,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: await proxy.metrics.reset_runtime() if proxy.cost_tracker: proxy.cost_tracker.reset_runtime() - initialize_rtk_session_baseline() + initialize_context_tool_session_baseline() async with _stats_snapshot_lock: _stats_snapshot["value"] = None _stats_snapshot["expires_at"] = 0.0 diff --git a/headroom/subscription/models.py b/headroom/subscription/models.py index 45b17620b..df93996b7 100644 --- a/headroom/subscription/models.py +++ b/headroom/subscription/models.py @@ -4,7 +4,7 @@ Mirrors the Anthropic OAuth usage API response exactly, including: - five_hour / seven_day rolling windows (utilization + reset times) - seven_day_opus / seven_day_sonnet per-model 7-day windows - extra_usage overage block (credits stored in cents by Anthropic) - - Headroom contribution: tokens conserved by compression, rtk, cache + - Headroom contribution: tokens conserved by compression, CLI filtering, cache - Window discrepancy detection (surge pricing, cache-miss anomalies) """ @@ -383,8 +383,11 @@ 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 - """Tokens avoided by CLI filtering (rtk) before reaching context.""" + """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).""" @@ -392,19 +395,26 @@ class HeadroomContribution: 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.tokens_saved_rtk + self.tokens_saved_cache_reads + return ( + self.tokens_saved_compression + + self.cli_filtering_saved() + + self.tokens_saved_cache_reads + ) def compression_saved(self) -> int: - """Tokens removed before model context by compression plus rtk.""" + """Tokens removed before model context by compression plus CLI filtering.""" - return self.tokens_saved_compression + self.tokens_saved_rtk + return self.tokens_saved_compression + self.cli_filtering_saved() 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.tokens_saved_rtk + return self.tokens_submitted + self.tokens_saved_compression + self.cli_filtering_saved() def efficiency_pct(self) -> float: raw = self.raw_without_headroom() @@ -418,7 +428,8 @@ class HeadroomContribution: "tokens_saved": { "compression": self.compression_saved(), "proxy_compression": self.tokens_saved_compression, - "rtk": self.tokens_saved_rtk, + "cli_filtering": self.cli_filtering_saved(), + "rtk": self.cli_filtering_saved(), "cache_reads": self.tokens_saved_cache_reads, "total": self.total_saved(), }, diff --git a/headroom/subscription/tracker.py b/headroom/subscription/tracker.py index 95260646b..9b28d0a62 100644 --- a/headroom/subscription/tracker.py +++ b/headroom/subscription/tracker.py @@ -193,6 +193,7 @@ class SubscriptionTracker(QuotaTracker): *, tokens_submitted: int = 0, tokens_saved_compression: int = 0, + tokens_saved_cli_filtering: int | None = None, tokens_saved_rtk: int = 0, tokens_saved_cache_reads: int = 0, compression_savings_usd: float = 0.0, @@ -204,9 +205,15 @@ class SubscriptionTracker(QuotaTracker): """ with self._lock: c = self._state.contribution + cli_filtering = ( + tokens_saved_rtk + if tokens_saved_cli_filtering is None + else tokens_saved_cli_filtering + ) c.tokens_submitted += max(tokens_submitted, 0) c.tokens_saved_compression += max(tokens_saved_compression, 0) - c.tokens_saved_rtk += max(tokens_saved_rtk, 0) + c.tokens_saved_cli_filtering += max(cli_filtering, 0) + c.tokens_saved_rtk += max(cli_filtering, 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) @@ -479,12 +486,14 @@ class SubscriptionTracker(QuotaTracker): c.tokens_submitted = int(contrib.get("tokens_submitted", 0)) saved = contrib.get("tokens_saved", {}) # Newer state writes dashboard-facing ``compression`` as - # proxy-compression + rtk. Prefer the raw proxy field when present - # so loading does not double-count rtk into the internal counter. + # proxy-compression + CLI filtering. Prefer the raw proxy field when + # present so loading does not double-count CLI filtering. c.tokens_saved_compression = int( saved.get("proxy_compression", saved.get("compression", 0)) ) - c.tokens_saved_rtk = int(saved.get("rtk", 0)) + cli_filtering = int(saved.get("cli_filtering", saved.get("rtk", 0))) + c.tokens_saved_cli_filtering = cli_filtering + c.tokens_saved_rtk = cli_filtering 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)) diff --git a/scripts/install.ps1 b/scripts/install.ps1 index 5c6c223ed..abe01875a 100644 --- a/scripts/install.ps1 +++ b/scripts/install.ps1 @@ -839,6 +839,38 @@ function Invoke-ClaudeRtkInit { } } +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, @@ -1614,6 +1646,7 @@ 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)) } @@ -1633,11 +1666,20 @@ 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) { Invoke-ClaudeRtkInit } + 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 } diff --git a/scripts/install.sh b/scripts/install.sh index 3296d8fb6..64e0f33dd 100644 --- a/scripts/install.sh +++ b/scripts/install.sh @@ -797,6 +797,38 @@ run_claude_rtk_init() { 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 @@ -1470,8 +1502,9 @@ main() { return fi - local known_args host_args port no_rtk no_proxy learn backend anyllm region + 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 proxy_args=() if [[ "${learn}" -eq 1 ]]; then @@ -1497,11 +1530,18 @@ 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 ]]; then + 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[@]}" diff --git a/tests/test_cli/test_wrap_bridge.py b/tests/test_cli/test_wrap_bridge.py index 48c475f38..2c7b6c674 100644 --- a/tests/test_cli/test_wrap_bridge.py +++ b/tests/test_cli/test_wrap_bridge.py @@ -3,12 +3,22 @@ from __future__ import annotations import json +import subprocess from pathlib import Path from unittest.mock import patch +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: + 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: @@ -29,6 +39,46 @@ def test_wrap_claude_prepare_only_skips_host_binary_lookup() -> None: 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"]) + + 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() @@ -43,6 +93,53 @@ def test_wrap_codex_prepare_only_updates_config(monkeypatch, tmp_path: Path) -> assert 'base_url = "http://127.0.0.1:8787/v1"' in config_file.read_text() +def test_wrap_codex_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", "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) runner = CliRunner() @@ -71,6 +168,27 @@ def test_wrap_cursor_prepare_only_injects_cursorrules(monkeypatch, tmp_path: Pat assert "headroom:rtk-instructions" in cursorrules.read_text() +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() + + def test_wrap_openclaw_prepare_only_emits_config_without_python_default() -> None: runner = CliRunner() diff --git a/tests/test_lean_ctx_installer.py b/tests/test_lean_ctx_installer.py new file mode 100644 index 000000000..1ba7734b6 --- /dev/null +++ b/tests/test_lean_ctx_installer.py @@ -0,0 +1,176 @@ +"""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 diff --git a/tests/test_paths.py b/tests/test_paths.py index 33013860a..6af26ef9c 100644 --- a/tests/test_paths.py +++ b/tests/test_paths.py @@ -326,6 +326,12 @@ def test_rtk_path_suffix(fake_home: Path) -> None: 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" @@ -384,6 +390,15 @@ def test_rtk_path_follows_workspace_env( 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: diff --git a/tests/test_proxy_dashboard_stats_cache.py b/tests/test_proxy_dashboard_stats_cache.py index eea9521a3..d9d970615 100644 --- a/tests/test_proxy_dashboard_stats_cache.py +++ b/tests/test_proxy_dashboard_stats_cache.py @@ -3,6 +3,7 @@ from __future__ import annotations import json import shutil import subprocess +from pathlib import Path from types import SimpleNamespace import pytest @@ -28,14 +29,19 @@ class _ToinStub: @pytest.fixture(autouse=True) -def _reset_rtk_stats_cache() -> None: - proxy_helpers._rtk_stats_cache.update({"expires_at": 0.0, "has_value": False, "value": None}) +def _reset_rtk_stats_cache(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False) + 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, "total_commands": 0, "tokens_saved": 0} + {"initialized": False, "tool": None, "total_commands": 0, "tokens_saved": 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 = [ @@ -60,6 +66,8 @@ def test_get_rtk_stats_memoizes_subprocess_calls(monkeypatch: pytest.MonkeyPatch assert first == second assert first == { + "tool": "rtk", + "label": "RTK", "installed": True, "total_commands": 0, "tokens_saved": 0, @@ -71,6 +79,8 @@ def test_get_rtk_stats_memoizes_subprocess_calls(monkeypatch: pytest.MonkeyPatch third = proxy_helpers._get_rtk_stats() assert third == { + "tool": "rtk", + "label": "RTK", "installed": True, "total_commands": 2, "tokens_saved": 266, @@ -79,6 +89,56 @@ def test_get_rtk_stats_memoizes_subprocess_calls(monkeypatch: pytest.MonkeyPatch assert calls["run"] == 2 +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, "tokens_saved": 400, "avg_savings_pct": 12.5}, + {"total_commands": 5, "tokens_saved": 475, "avg_savings_pct": 15.0}, + ] + + def _fake_run(args, **kwargs): + calls["run"] += 1 + assert args == ["/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", + "label": "lean-ctx", + "installed": True, + "total_commands": 0, + "tokens_saved": 0, + "avg_savings_pct": 12.5, + } + 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", + "label": "lean-ctx", + "installed": True, + "total_commands": 2, + "tokens_saved": 75, + "avg_savings_pct": 15.0, + } + assert calls["run"] == 2 + + def test_stats_cached_query_reuses_short_ttl_snapshot(monkeypatch: pytest.MonkeyPatch) -> None: pytest.importorskip("fastapi") from fastapi.testclient import TestClient @@ -86,7 +146,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, "rtk": 0} + calls = {"store": 0, "telemetry": 0, "feedback": 0, "context_tool": 0} now = {"value": 100.0} monkeypatch.setattr(server.time, "monotonic", lambda: now["value"]) @@ -106,16 +166,18 @@ def test_stats_cached_query_reuses_short_ttl_snapshot(monkeypatch: pytest.Monkey lambda: _StatsStub(calls, "feedback", {}), ) - def _fake_rtk_stats() -> dict[str, int | bool | float]: - calls["rtk"] += 1 + 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_rtk_stats", _fake_rtk_stats) + monkeypatch.setattr(server, "_get_context_tool_stats", _fake_context_tool_stats) monkeypatch.setattr(server, "get_toin", lambda: _ToinStub()) app = create_app( @@ -143,21 +205,142 @@ 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, "rtk": 3} + 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 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_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 + + def test_stats_reset_clears_runtime_proxy_counters(monkeypatch: pytest.MonkeyPatch) -> None: pytest.importorskip("fastapi") from fastapi.testclient import TestClient @@ -181,7 +364,7 @@ def test_stats_reset_clears_runtime_proxy_counters(monkeypatch: pytest.MonkeyPat "get_compression_feedback", lambda: _StatsStub({"feedback": 0}, "feedback", {}), ) - monkeypatch.setattr(server, "_get_rtk_stats", lambda: None) + monkeypatch.setattr(server, "_get_context_tool_stats", lambda: None) monkeypatch.setattr(server, "get_toin", lambda: _ToinStub()) app = create_app( @@ -224,3 +407,9 @@ def test_dashboard_uses_cached_stats_and_lazy_history_feed_polling() -> None: 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'" in html diff --git a/tests/test_subscription_tracker.py b/tests/test_subscription_tracker.py index b591283f0..00c069730 100644 --- a/tests/test_subscription_tracker.py +++ b/tests/test_subscription_tracker.py @@ -56,7 +56,7 @@ def test_tracker_notify_active_update_and_basic_state(monkeypatch: pytest.Monkey tracker.update_contribution( tokens_submitted=10, tokens_saved_compression=5, - tokens_saved_rtk=-1, + tokens_saved_cli_filtering=-1, tokens_saved_cache_reads=3, compression_savings_usd=1.25, cache_savings_usd=-2.0, @@ -64,10 +64,13 @@ 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 @@ -178,7 +181,7 @@ def test_persist_and_load_state_round_trip(tmp_path: Path) -> None: tracker.update_contribution( tokens_submitted=11, tokens_saved_compression=2, - tokens_saved_rtk=3, + tokens_saved_cli_filtering=3, tokens_saved_cache_reads=4, compression_savings_usd=1.5, cache_savings_usd=2.5, @@ -189,10 +192,13 @@ def test_persist_and_load_state_round_trip(tmp_path: Path) -> None: loader = SubscriptionTracker(persist_path=persist_path) assert loader._state.contribution.tokens_submitted == 11 assert loader._state.contribution.tokens_saved_compression == 2 + assert loader._state.contribution.tokens_saved_cli_filtering == 3 assert loader._state.contribution.tokens_saved_rtk == 3 assert loader._state.contribution.tokens_saved_cache_reads == 4 assert loader._state.contribution.to_dict()["tokens_saved"]["compression"] == 5 assert loader._state.contribution.to_dict()["tokens_saved"]["proxy_compression"] == 2 + assert loader._state.contribution.to_dict()["tokens_saved"]["cli_filtering"] == 3 + assert loader._state.contribution.to_dict()["tokens_saved"]["rtk"] == 3 assert loader._state.contribution.compression_savings_usd == 1.5 assert loader._state.contribution.cache_savings_usd == 2.5 assert loader._state.poll_count == 7 diff --git a/tests/test_telemetry_warning.py b/tests/test_telemetry_warning.py index 8c9755955..43186c4a4 100644 --- a/tests/test_telemetry_warning.py +++ b/tests/test_telemetry_warning.py @@ -139,6 +139,17 @@ 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