mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat: add lean-ctx context tool support
This commit is contained in:
parent
44231f68cd
commit
4b061792b2
24 changed files with 1379 additions and 138 deletions
|
|
@ -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 ...`.
|
||||
|
||||
---
|
||||
|
||||
|
|
|
|||
|
|
@ -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 <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`.
|
||||
|
||||
## Modes
|
||||
|
||||
| Mode | Behavior | Use Case |
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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_<port>` | — |
|
||||
|
||||
|
|
|
|||
|
|
@ -100,7 +100,7 @@ headroom wrap [OPTIONS] -- <command> [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
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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 <id> # 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.")
|
||||
|
|
|
|||
|
|
@ -109,7 +109,7 @@
|
|||
</div>
|
||||
<div class="mt-2 text-xs text-gray-500">
|
||||
<span x-show="stats.cost?.savings_usd > 0"
|
||||
x-text="formatNumber(stats.tokens?.proxy_compression_saved || 0) + ' proxy tokens only; RTK excluded from $'"></span>
|
||||
x-text="formatNumber(stats.tokens?.proxy_compression_saved || 0) + ' proxy tokens only; ' + cliFilteringLabel + ' excluded from $'"></span>
|
||||
<span x-show="!(stats.cost?.savings_usd > 0)"
|
||||
x-text="formatNumber(stats.requests?.total || 0) + ' requests processed'"></span>
|
||||
</div>
|
||||
|
|
@ -125,7 +125,7 @@
|
|||
<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="'RTK ' + formatNumber(stats.tokens?.rtk_saved || 0) + ' (' + rtkShareOfTotal.toFixed(1) + '%)'"></span>
|
||||
<span x-text="cliFilteringLabel + ' ' + formatNumber(cliFilteringSaved) + ' (' + cliFilteringShareOfTotal.toFixed(1) + '%)'"></span>
|
||||
</div>
|
||||
<div class="mt-2 h-8">
|
||||
<svg class="w-full h-full" viewBox="0 0 100 32" preserveAspectRatio="none">
|
||||
|
|
@ -698,8 +698,8 @@
|
|||
<span class="font-mono text-sm" x-text="formatNumber(stats.tokens?.total_before_compression || 0)"></span>
|
||||
</div>
|
||||
<div class="flex justify-between items-center">
|
||||
<span class="text-sm text-gray-400">RTK Filtered</span>
|
||||
<span class="font-mono text-sm text-emerald-400" x-text="formatNumber(stats.tokens?.rtk_saved || 0)"></span>
|
||||
<span class="text-sm text-gray-400" x-text="cliFilteringLabel + ' Filtered'"></span>
|
||||
<span class="font-mono text-sm text-emerald-400" x-text="formatNumber(cliFilteringSaved)"></span>
|
||||
</div>
|
||||
<div class="flex justify-between items-center">
|
||||
<span class="text-sm text-gray-400">Proxy Removed</span>
|
||||
|
|
@ -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 ---
|
||||
|
|
|
|||
46
headroom/lean_ctx/__init__.py
Normal file
46
headroom/lean_ctx/__init__.py
Normal file
|
|
@ -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
|
||||
179
headroom/lean_ctx/installer.py
Normal file
179
headroom/lean_ctx/installer.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
),
|
||||
},
|
||||
|
|
|
|||
|
|
@ -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"):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
},
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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[@]}"
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
176
tests/test_lean_ctx_installer.py
Normal file
176
tests/test_lean_ctx_installer.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue