mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
Add wrap commands for Codex/Cursor/Aider with rtk instructions, fix savings metrics
- headroom wrap codex: injects rtk instructions into AGENTS.md - headroom wrap cursor: injects into .cursorrules, prints config steps - headroom wrap aider: injects into CONVENTIONS.md, sets both env vars - All with --no-rtk flag to skip rtk setup Unified savings metrics across three layers: - cli_filtering: tokens avoided by rtk before reaching context - compression: tokens removed by proxy (SmartCrusher, etc.) - prefix_cache: provider cache discount with honest attribution Fix flaky test_process_stats_collected when psutil not installed.
This commit is contained in:
chopratejas
parent
0e09097735
commit
e2aac4863a
3 changed files with 527 additions and 69 deletions
|
|
@ -2,6 +2,9 @@
|
|||
|
||||
Usage:
|
||||
headroom wrap claude # Start proxy + rtk + claude
|
||||
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 claude --no-rtk # Without rtk hooks
|
||||
headroom wrap claude --port 9999 # Custom proxy port
|
||||
headroom wrap claude -- --model opus # Pass args to claude
|
||||
|
|
@ -112,6 +115,184 @@ def _setup_rtk(verbose: bool = False) -> Path | None:
|
|||
return rtk_path
|
||||
|
||||
|
||||
# rtk instructions for tools without hook support (Codex, Cursor, Aider).
|
||||
# These get injected into AGENTS.md / .cursorrules so the LLM voluntarily
|
||||
# uses rtk-prefixed commands. Kept concise to minimize instruction overhead.
|
||||
RTK_INSTRUCTIONS_BLOCK = """\
|
||||
<!-- headroom:rtk-instructions -->
|
||||
# RTK (Rust Token Killer) - Token-Optimized Commands
|
||||
|
||||
When running shell commands, **always prefix with `rtk`**. This reduces context
|
||||
usage by 60-90% with zero behavior change. If rtk has no filter for a command,
|
||||
it passes through unchanged — so it is always safe to use.
|
||||
|
||||
## Key Commands
|
||||
```bash
|
||||
# Git (59-80% savings)
|
||||
rtk git status rtk git diff rtk git log
|
||||
|
||||
# Files & Search (60-75% savings)
|
||||
rtk ls <path> rtk read <file> rtk grep <pattern>
|
||||
rtk find <pattern> rtk diff <file>
|
||||
|
||||
# Test (90-99% savings) — shows failures only
|
||||
rtk pytest tests/ rtk cargo test rtk test <cmd>
|
||||
|
||||
# Build & Lint (80-90% savings) — shows errors only
|
||||
rtk tsc rtk lint rtk cargo build
|
||||
rtk prettier --check rtk mypy rtk ruff check
|
||||
|
||||
# Analysis (70-90% savings)
|
||||
rtk err <cmd> rtk log <file> rtk json <file>
|
||||
rtk summary <cmd> rtk deps rtk env
|
||||
|
||||
# GitHub (26-87% savings)
|
||||
rtk gh pr view <n> rtk gh run list rtk gh issue list
|
||||
|
||||
# Infrastructure (85% savings)
|
||||
rtk docker ps rtk kubectl get rtk docker logs <c>
|
||||
|
||||
# Package managers (70-90% savings)
|
||||
rtk pip list rtk pnpm install rtk npm run <script>
|
||||
```
|
||||
|
||||
## Rules
|
||||
- In command chains, prefix each segment: `rtk git add . && rtk git commit -m "msg"`
|
||||
- For debugging, use raw command without rtk prefix
|
||||
- `rtk proxy <cmd>` runs command without filtering but tracks usage
|
||||
<!-- /headroom:rtk-instructions -->
|
||||
"""
|
||||
|
||||
# Marker used to detect if instructions are already injected
|
||||
_RTK_MARKER = "<!-- headroom:rtk-instructions -->"
|
||||
|
||||
|
||||
def _ensure_rtk_binary(verbose: bool = False) -> Path | None:
|
||||
"""Ensure rtk binary is installed (download if needed). No hook registration."""
|
||||
from headroom.rtk import get_rtk_path
|
||||
from headroom.rtk.installer import ensure_rtk
|
||||
|
||||
rtk_path = get_rtk_path()
|
||||
|
||||
if rtk_path:
|
||||
if verbose:
|
||||
click.echo(f" rtk found at {rtk_path}")
|
||||
return rtk_path
|
||||
|
||||
click.echo(" Downloading rtk (Rust Token Killer)...")
|
||||
rtk_path = ensure_rtk()
|
||||
if rtk_path:
|
||||
click.echo(f" rtk installed at {rtk_path}")
|
||||
return rtk_path
|
||||
|
||||
click.echo(" rtk download failed — continuing without it")
|
||||
return None
|
||||
|
||||
|
||||
def _inject_rtk_instructions(file_path: Path, verbose: bool = False) -> bool:
|
||||
"""Inject rtk instructions into a file (AGENTS.md, .cursorrules, etc.).
|
||||
|
||||
Idempotent — skips if marker already present. Appends to existing content.
|
||||
Returns True if instructions were written.
|
||||
"""
|
||||
if file_path.exists():
|
||||
existing = file_path.read_text()
|
||||
if _RTK_MARKER in existing:
|
||||
if verbose:
|
||||
click.echo(f" rtk instructions already in {file_path.name}")
|
||||
return True
|
||||
# Append to existing file
|
||||
with open(file_path, "a") as f:
|
||||
f.write("\n\n" + RTK_INSTRUCTIONS_BLOCK)
|
||||
else:
|
||||
file_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
file_path.write_text(RTK_INSTRUCTIONS_BLOCK)
|
||||
|
||||
click.echo(f" rtk instructions injected into {file_path}")
|
||||
return True
|
||||
|
||||
|
||||
def _ensure_proxy(port: int, no_proxy: bool) -> subprocess.Popen | None:
|
||||
"""Start or verify proxy. Returns process handle if we started it."""
|
||||
if not no_proxy:
|
||||
if _check_proxy(port):
|
||||
click.echo(f" Proxy already running on port {port}")
|
||||
return None
|
||||
else:
|
||||
click.echo(f" Starting Headroom proxy on port {port}...")
|
||||
try:
|
||||
proc = _start_proxy(port)
|
||||
click.echo(f" Proxy ready on http://127.0.0.1:{port}")
|
||||
return proc
|
||||
except RuntimeError as e:
|
||||
click.echo(f" Error: {e}")
|
||||
raise SystemExit(1) from e
|
||||
else:
|
||||
if not _check_proxy(port):
|
||||
click.echo(f" Warning: No proxy detected on port {port}")
|
||||
return None
|
||||
|
||||
|
||||
def _make_cleanup(proxy_proc_holder: list) -> Any:
|
||||
"""Create a cleanup function that terminates the proxy on exit."""
|
||||
|
||||
def cleanup(signum: int | None = None, frame: Any = None) -> None:
|
||||
proc = proxy_proc_holder[0] if proxy_proc_holder else None
|
||||
if proc and proc.poll() is None:
|
||||
proc.terminate()
|
||||
try:
|
||||
proc.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
proc.kill()
|
||||
|
||||
return cleanup
|
||||
|
||||
|
||||
def _launch_tool(
|
||||
binary: str,
|
||||
args: tuple,
|
||||
env: dict[str, str],
|
||||
port: int,
|
||||
no_proxy: bool,
|
||||
tool_label: str,
|
||||
env_vars_display: list[str],
|
||||
) -> None:
|
||||
"""Common logic: start proxy, launch tool, clean up."""
|
||||
proxy_holder: list[subprocess.Popen | None] = [None]
|
||||
cleanup = _make_cleanup(proxy_holder)
|
||||
signal.signal(signal.SIGINT, cleanup)
|
||||
signal.signal(signal.SIGTERM, cleanup)
|
||||
|
||||
try:
|
||||
click.echo()
|
||||
padded = f"HEADROOM WRAP: {tool_label}".center(47)
|
||||
click.echo(" ╔═══════════════════════════════════════════════╗")
|
||||
click.echo(f" ║{padded}║")
|
||||
click.echo(" ╚═══════════════════════════════════════════════╝")
|
||||
click.echo()
|
||||
|
||||
proxy_holder[0] = _ensure_proxy(port, no_proxy)
|
||||
|
||||
click.echo()
|
||||
click.echo(f" Launching {tool_label} (API routed through Headroom)...")
|
||||
for var in env_vars_display:
|
||||
click.echo(f" {var}")
|
||||
if args:
|
||||
click.echo(f" Extra args: {' '.join(args)}")
|
||||
click.echo()
|
||||
|
||||
result = subprocess.run([binary, *args], env=env)
|
||||
raise SystemExit(result.returncode)
|
||||
|
||||
except SystemExit:
|
||||
raise
|
||||
except Exception as e:
|
||||
click.echo(f" Error: {e}")
|
||||
raise SystemExit(1) from e
|
||||
finally:
|
||||
cleanup()
|
||||
|
||||
|
||||
@main.group()
|
||||
def wrap() -> None:
|
||||
"""Wrap CLI tools to run through Headroom.
|
||||
|
|
@ -121,13 +302,19 @@ def wrap() -> None:
|
|||
the target tool so all API calls route through Headroom automatically.
|
||||
|
||||
\b
|
||||
Example:
|
||||
headroom wrap claude # Proxy + rtk + Claude Code
|
||||
headroom wrap claude --no-rtk # Proxy only, no rtk hooks
|
||||
headroom wrap claude -- -p # Pass --print flag to claude
|
||||
Supported tools:
|
||||
headroom wrap claude # Claude Code (Anthropic)
|
||||
headroom wrap codex # OpenAI Codex CLI
|
||||
headroom wrap aider # Aider
|
||||
headroom wrap cursor # Cursor (prints config instructions)
|
||||
"""
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Claude Code
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@wrap.command()
|
||||
@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")
|
||||
|
|
@ -138,17 +325,7 @@ def claude(port: int, no_rtk: bool, no_proxy: bool, verbose: bool, claude_args:
|
|||
"""Launch Claude Code through Headroom proxy.
|
||||
|
||||
\b
|
||||
This command:
|
||||
1. Starts a Headroom optimization proxy on localhost
|
||||
2. Installs rtk and registers Claude Code hooks (compresses CLI output)
|
||||
3. Sets ANTHROPIC_API_URL to route through the proxy
|
||||
4. Launches 'claude' with your arguments
|
||||
|
||||
\b
|
||||
All API calls from Claude Code flow through Headroom, which:
|
||||
- Compresses tool outputs (SmartCrusher, Kompress, CodeCompressor)
|
||||
- Preserves prefix cache (frozen message optimization)
|
||||
- Tracks token savings and cache hit rates
|
||||
Sets ANTHROPIC_BASE_URL to route all Anthropic API calls through Headroom.
|
||||
|
||||
\b
|
||||
Examples:
|
||||
|
|
@ -157,24 +334,15 @@ def claude(port: int, no_rtk: bool, no_proxy: bool, verbose: bool, claude_args:
|
|||
headroom wrap claude --port 9999 # Custom proxy port
|
||||
headroom wrap claude --no-rtk # Skip rtk (proxy only)
|
||||
"""
|
||||
# Check that claude CLI is available
|
||||
claude_bin = shutil.which("claude")
|
||||
if not claude_bin:
|
||||
click.echo("Error: 'claude' not found in PATH.")
|
||||
click.echo("Install Claude Code: https://docs.anthropic.com/en/docs/claude-code")
|
||||
raise SystemExit(1)
|
||||
|
||||
proxy_proc: subprocess.Popen | None = None
|
||||
|
||||
def cleanup(signum: int | None = None, frame: Any = None) -> None:
|
||||
"""Clean up proxy on exit."""
|
||||
if proxy_proc and proxy_proc.poll() is None:
|
||||
proxy_proc.terminate()
|
||||
try:
|
||||
proxy_proc.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
proxy_proc.kill()
|
||||
|
||||
# Setup rtk before launching (Claude-specific)
|
||||
proxy_holder: list[subprocess.Popen | None] = [None]
|
||||
cleanup = _make_cleanup(proxy_holder)
|
||||
signal.signal(signal.SIGINT, cleanup)
|
||||
signal.signal(signal.SIGTERM, cleanup)
|
||||
|
||||
|
|
@ -185,31 +353,14 @@ def claude(port: int, no_rtk: bool, no_proxy: bool, verbose: bool, claude_args:
|
|||
click.echo(" ╚═══════════════════════════════════════════════╝")
|
||||
click.echo()
|
||||
|
||||
# Step 1: Start proxy
|
||||
if not no_proxy:
|
||||
if _check_proxy(port):
|
||||
click.echo(f" Proxy already running on port {port}")
|
||||
else:
|
||||
click.echo(f" Starting Headroom proxy on port {port}...")
|
||||
try:
|
||||
proxy_proc = _start_proxy(port)
|
||||
click.echo(f" Proxy ready on http://127.0.0.1:{port}")
|
||||
except RuntimeError as e:
|
||||
click.echo(f" Error: {e}")
|
||||
raise SystemExit(1) from e
|
||||
else:
|
||||
if not _check_proxy(port):
|
||||
click.echo(f" Warning: No proxy detected on port {port}")
|
||||
proxy_holder[0] = _ensure_proxy(port, no_proxy)
|
||||
|
||||
# Step 2: Setup rtk
|
||||
if not no_rtk:
|
||||
click.echo(" Setting up rtk...")
|
||||
_setup_rtk(verbose=verbose)
|
||||
else:
|
||||
if verbose:
|
||||
click.echo(" Skipping rtk (--no-rtk)")
|
||||
elif verbose:
|
||||
click.echo(" Skipping rtk (--no-rtk)")
|
||||
|
||||
# Step 3: Launch claude
|
||||
click.echo()
|
||||
click.echo(" Launching Claude Code (API routed through Headroom)...")
|
||||
click.echo(f" ANTHROPIC_BASE_URL=http://127.0.0.1:{port}")
|
||||
|
|
@ -220,12 +371,7 @@ def claude(port: int, no_rtk: bool, no_proxy: bool, verbose: bool, claude_args:
|
|||
env = os.environ.copy()
|
||||
env["ANTHROPIC_BASE_URL"] = f"http://127.0.0.1:{port}"
|
||||
|
||||
# Run claude — this blocks until claude exits
|
||||
result = subprocess.run(
|
||||
[claude_bin, *claude_args],
|
||||
env=env,
|
||||
)
|
||||
|
||||
result = subprocess.run([claude_bin, *claude_args], env=env)
|
||||
raise SystemExit(result.returncode)
|
||||
|
||||
except SystemExit:
|
||||
|
|
@ -235,3 +381,213 @@ def claude(port: int, no_rtk: bool, no_proxy: bool, verbose: bool, claude_args:
|
|||
raise SystemExit(1) from e
|
||||
finally:
|
||||
cleanup()
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# OpenAI Codex CLI
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@wrap.command()
|
||||
@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-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)")
|
||||
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
|
||||
@click.argument("codex_args", nargs=-1, type=click.UNPROCESSED)
|
||||
def codex(port: int, no_rtk: bool, no_proxy: bool, verbose: bool, codex_args: tuple) -> None:
|
||||
"""Launch OpenAI Codex CLI through Headroom proxy.
|
||||
|
||||
\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).
|
||||
|
||||
\b
|
||||
Examples:
|
||||
headroom wrap codex # Start proxy + rtk + codex
|
||||
headroom wrap codex -- "fix the bug" # Pass prompt to codex
|
||||
headroom wrap codex --no-rtk # Skip rtk setup
|
||||
headroom wrap codex --port 9999 # Custom proxy port
|
||||
"""
|
||||
codex_bin = shutil.which("codex")
|
||||
if not codex_bin:
|
||||
click.echo("Error: 'codex' not found in PATH.")
|
||||
click.echo("Install Codex CLI: npm install -g @openai/codex")
|
||||
raise SystemExit(1)
|
||||
|
||||
# Setup rtk for Codex (binary + AGENTS.md instructions, no hooks)
|
||||
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)
|
||||
|
||||
# Also inject into global ~/.codex/AGENTS.md
|
||||
global_agents = Path.home() / ".codex" / "AGENTS.md"
|
||||
_inject_rtk_instructions(global_agents, verbose=verbose)
|
||||
|
||||
env = os.environ.copy()
|
||||
env["OPENAI_BASE_URL"] = f"http://127.0.0.1:{port}/v1"
|
||||
|
||||
_launch_tool(
|
||||
binary=codex_bin,
|
||||
args=codex_args,
|
||||
env=env,
|
||||
port=port,
|
||||
no_proxy=no_proxy,
|
||||
tool_label="CODEX",
|
||||
env_vars_display=[f"OPENAI_BASE_URL=http://127.0.0.1:{port}/v1"],
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Aider
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@wrap.command()
|
||||
@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-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)")
|
||||
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
|
||||
@click.argument("aider_args", nargs=-1, type=click.UNPROCESSED)
|
||||
def aider(port: int, no_rtk: bool, no_proxy: bool, verbose: bool, aider_args: tuple) -> None:
|
||||
"""Launch aider through Headroom proxy.
|
||||
|
||||
\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.
|
||||
|
||||
\b
|
||||
Examples:
|
||||
headroom wrap aider # Start proxy + rtk + 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
|
||||
"""
|
||||
aider_bin = shutil.which("aider")
|
||||
if not aider_bin:
|
||||
click.echo("Error: 'aider' not found in PATH.")
|
||||
click.echo("Install aider: pip install aider-chat")
|
||||
raise SystemExit(1)
|
||||
|
||||
# Setup rtk for aider (binary + CONVENTIONS.md instructions)
|
||||
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)
|
||||
|
||||
env = os.environ.copy()
|
||||
env["OPENAI_API_BASE"] = f"http://127.0.0.1:{port}/v1"
|
||||
env["ANTHROPIC_BASE_URL"] = f"http://127.0.0.1:{port}"
|
||||
|
||||
_launch_tool(
|
||||
binary=aider_bin,
|
||||
args=aider_args,
|
||||
env=env,
|
||||
port=port,
|
||||
no_proxy=no_proxy,
|
||||
tool_label="AIDER",
|
||||
env_vars_display=[
|
||||
f"OPENAI_API_BASE=http://127.0.0.1:{port}/v1",
|
||||
f"ANTHROPIC_BASE_URL=http://127.0.0.1:{port}",
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Cursor
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@wrap.command()
|
||||
@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-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)")
|
||||
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
|
||||
def cursor(port: int, no_rtk: bool, no_proxy: bool, verbose: bool) -> None:
|
||||
"""Start Headroom proxy for use with 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.
|
||||
|
||||
\b
|
||||
After running this command, open Cursor and configure:
|
||||
Settings > Models > OpenAI API Key > Advanced > Override Base URL
|
||||
|
||||
\b
|
||||
Example:
|
||||
headroom wrap cursor # Start proxy + rtk + instructions
|
||||
headroom wrap cursor --no-rtk # Proxy only, no rtk
|
||||
headroom wrap cursor --port 9999 # Custom proxy port
|
||||
"""
|
||||
proxy_holder: list[subprocess.Popen | None] = [None]
|
||||
cleanup = _make_cleanup(proxy_holder)
|
||||
signal.signal(signal.SIGINT, cleanup)
|
||||
signal.signal(signal.SIGTERM, cleanup)
|
||||
|
||||
try:
|
||||
click.echo()
|
||||
click.echo(" ╔═══════════════════════════════════════════════╗")
|
||||
click.echo(" ║ HEADROOM WRAP: CURSOR ║")
|
||||
click.echo(" ╚═══════════════════════════════════════════════╝")
|
||||
click.echo()
|
||||
|
||||
proxy_holder[0] = _ensure_proxy(port, no_proxy)
|
||||
|
||||
# Setup rtk for Cursor (binary + .cursorrules instructions)
|
||||
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)
|
||||
|
||||
click.echo()
|
||||
click.echo(" Headroom proxy is running. Configure Cursor:")
|
||||
click.echo()
|
||||
click.echo(" For OpenAI models:")
|
||||
click.echo(f" Base URL: http://127.0.0.1:{port}/v1")
|
||||
click.echo(" API Key: your-openai-api-key")
|
||||
click.echo()
|
||||
click.echo(" For Anthropic models:")
|
||||
click.echo(f" Base URL: http://127.0.0.1:{port}")
|
||||
click.echo(" API Key: your-anthropic-api-key")
|
||||
click.echo()
|
||||
click.echo(" In Cursor:")
|
||||
click.echo(" Settings > Models > OpenAI API Key > Override OpenAI Base URL")
|
||||
click.echo(f" Set to: http://127.0.0.1:{port}/v1")
|
||||
if not no_rtk:
|
||||
click.echo()
|
||||
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.")
|
||||
click.echo()
|
||||
|
||||
# Block until Ctrl+C
|
||||
try:
|
||||
while True:
|
||||
time.sleep(1)
|
||||
proc = proxy_holder[0]
|
||||
if proc and proc.poll() is not None:
|
||||
click.echo(" Proxy process exited unexpectedly.")
|
||||
raise SystemExit(1)
|
||||
except KeyboardInterrupt:
|
||||
click.echo("\n Shutting down...")
|
||||
|
||||
except SystemExit:
|
||||
raise
|
||||
except Exception as e:
|
||||
click.echo(f" Error: {e}")
|
||||
raise SystemExit(1) from e
|
||||
finally:
|
||||
cleanup()
|
||||
|
|
|
|||
|
|
@ -217,6 +217,46 @@ _CACHE_ECONOMICS = {
|
|||
}
|
||||
|
||||
|
||||
def _get_rtk_stats() -> dict[str, Any] | None:
|
||||
"""Get rtk (Rust Token Killer) savings stats if rtk is installed.
|
||||
|
||||
Reads from rtk's tracking database via `rtk gain --format json`.
|
||||
Returns None if rtk is not installed.
|
||||
"""
|
||||
import shutil
|
||||
import subprocess as _sp
|
||||
|
||||
rtk_bin = shutil.which("rtk")
|
||||
if not rtk_bin:
|
||||
# Check headroom-managed install
|
||||
rtk_managed = Path.home() / ".headroom" / "bin" / "rtk"
|
||||
if rtk_managed.exists():
|
||||
rtk_bin = str(rtk_managed)
|
||||
else:
|
||||
return None
|
||||
|
||||
try:
|
||||
result = _sp.run(
|
||||
[rtk_bin, "gain", "--format", "json"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5,
|
||||
)
|
||||
if result.returncode == 0 and result.stdout.strip():
|
||||
data = json.loads(result.stdout)
|
||||
summary = data.get("summary", {})
|
||||
return {
|
||||
"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),
|
||||
}
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return {"installed": True, "total_commands": 0, "tokens_saved": 0, "avg_savings_pct": 0.0}
|
||||
|
||||
|
||||
def _build_prefix_cache_stats(
|
||||
metrics: PrometheusMetrics,
|
||||
cost_tracker: CostTracker | None,
|
||||
|
|
@ -322,24 +362,27 @@ def _build_prefix_cache_stats(
|
|||
metrics.prefix_freeze_tokens_preserved - metrics.prefix_freeze_compression_foregone
|
||||
),
|
||||
},
|
||||
"attribution": (
|
||||
"Prefix caching is performed by the LLM provider (Anthropic, OpenAI). "
|
||||
"Headroom reports cache stats as observed from API responses. "
|
||||
"CacheAligner and prefix freeze improve cache hit rates by stabilizing "
|
||||
"the message prefix, but baseline caching happens without Headroom."
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _merge_cost_stats(
|
||||
cost_stats: dict | None,
|
||||
cache_stats: dict,
|
||||
cli_tokens_avoided: int = 0,
|
||||
) -> dict | None:
|
||||
"""Add prefix cache savings to overall cost stats.
|
||||
"""Add prefix cache and CLI filtering savings to overall cost stats.
|
||||
|
||||
Compression savings and cache savings are computed on different token
|
||||
scopes (Headroom tracks user-message tokens; Anthropic's cache metrics
|
||||
cover the entire prompt including system/tools). We keep both as
|
||||
additive line items without cross-contaminating the counterfactual.
|
||||
|
||||
- compression_savings_usd: from removing tokens (CostTracker scope)
|
||||
- cache_savings_usd: from prefix cache discounts (full-prompt scope)
|
||||
- savings_usd: sum of both (hero metric)
|
||||
- cost_with/without_headroom_usd: unchanged (compression scope only)
|
||||
Three savings layers, each on a different scope:
|
||||
- compression_savings_usd: tokens removed by proxy (SmartCrusher, etc.)
|
||||
- cache_savings_usd: prefix cache discount from provider
|
||||
- cli_filtering_savings_usd: tokens avoided by rtk before reaching context
|
||||
- savings_usd: sum of all (hero metric)
|
||||
"""
|
||||
if cost_stats is None:
|
||||
return None
|
||||
|
|
@ -347,11 +390,24 @@ def _merge_cost_stats(
|
|||
cache_net = cache_stats.get("totals", {}).get("net_savings_usd", 0.0)
|
||||
compression_savings = cost_stats.get("savings_usd", 0.0)
|
||||
|
||||
# Estimate CLI filtering savings: tokens_avoided * avg input price
|
||||
# Use the average input price from the cost tracker if available
|
||||
cli_savings_usd = 0.0
|
||||
if cli_tokens_avoided > 0 and LITELLM_AVAILABLE:
|
||||
# Use a conservative estimate: average across models seen
|
||||
total_input_cost = cost_stats.get("total_input_cost_usd", 0.0)
|
||||
total_input_tokens = cost_stats.get("total_input_tokens", 0)
|
||||
if total_input_tokens > 0:
|
||||
avg_price_per_token = total_input_cost / total_input_tokens
|
||||
cli_savings_usd = cli_tokens_avoided * avg_price_per_token
|
||||
|
||||
return {
|
||||
**cost_stats,
|
||||
"savings_usd": round(compression_savings + cache_net, 4),
|
||||
"savings_usd": round(compression_savings + cache_net + cli_savings_usd, 4),
|
||||
"compression_savings_usd": round(compression_savings, 4),
|
||||
"cache_savings_usd": round(cache_net, 4),
|
||||
"cli_filtering_savings_usd": round(cli_savings_usd, 4),
|
||||
"cli_tokens_avoided": cli_tokens_avoided,
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -6240,10 +6296,42 @@ 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()
|
||||
cli_tokens_avoided = (
|
||||
cli_filtering_stats.get("tokens_saved", 0) if cli_filtering_stats else 0
|
||||
)
|
||||
|
||||
# Calculate total tokens before compression
|
||||
total_tokens_before = m.tokens_input_total + m.tokens_saved_total
|
||||
|
||||
# Build unified savings summary (all layers)
|
||||
compression_tokens = m.tokens_saved_total
|
||||
cache_net_usd = prefix_cache_stats.get("totals", {}).get("net_savings_usd", 0.0)
|
||||
total_tokens_all_layers = compression_tokens + cli_tokens_avoided
|
||||
|
||||
return {
|
||||
"savings": {
|
||||
"total_tokens": total_tokens_all_layers,
|
||||
"by_layer": {
|
||||
"cli_filtering": {
|
||||
"tokens": cli_tokens_avoided,
|
||||
"description": "Tokens avoided by CLI output filtering (rtk) before reaching context",
|
||||
},
|
||||
"compression": {
|
||||
"tokens": compression_tokens,
|
||||
"description": "Tokens removed by proxy compression (SmartCrusher, ContentRouter, etc.)",
|
||||
},
|
||||
"prefix_cache": {
|
||||
"discount_usd": round(cache_net_usd, 4),
|
||||
"description": (
|
||||
"Cost discount from provider prefix caching. "
|
||||
"Headroom's CacheAligner improves hit rates; "
|
||||
"baseline caching is provider-native."
|
||||
),
|
||||
},
|
||||
},
|
||||
},
|
||||
"requests": {
|
||||
"total": m.requests_total,
|
||||
"cached": m.requests_cached,
|
||||
|
|
@ -6256,6 +6344,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
"input": m.tokens_input_total,
|
||||
"output": m.tokens_output_total,
|
||||
"saved": m.tokens_saved_total,
|
||||
"cli_tokens_avoided": cli_tokens_avoided,
|
||||
"total_before_compression": total_tokens_before,
|
||||
"savings_percent": round(
|
||||
(m.tokens_saved_total / total_tokens_before * 100)
|
||||
|
|
@ -6298,6 +6387,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
"cost": _merge_cost_stats(
|
||||
proxy.cost_tracker.stats() if proxy.cost_tracker else None,
|
||||
prefix_cache_stats,
|
||||
cli_tokens_avoided=cli_tokens_avoided,
|
||||
),
|
||||
"compression": {
|
||||
"ccr_entries": compression_stats.get("entry_count", 0),
|
||||
|
|
@ -6327,6 +6417,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
),
|
||||
},
|
||||
"toin": get_toin().get_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,
|
||||
"recent_requests": proxy.logger.get_recent(10) if proxy.logger else [],
|
||||
|
|
|
|||
|
|
@ -445,9 +445,14 @@ class TestProcessStatsCollection:
|
|||
|
||||
stats = tracker.get_process_stats()
|
||||
|
||||
# Should have real values from the current process
|
||||
assert stats.rss_bytes > 0 # Process must use some memory
|
||||
assert stats.vms_bytes > 0
|
||||
# psutil is optional — without it, stats are all zeros (graceful degradation)
|
||||
try:
|
||||
import psutil # noqa: F401
|
||||
|
||||
assert stats.rss_bytes > 0 # Process must use some memory
|
||||
assert stats.vms_bytes > 0
|
||||
except ImportError:
|
||||
assert stats.rss_bytes == 0 # No psutil → zeros expected
|
||||
assert stats.percent >= 0 # Could be 0 on some systems
|
||||
|
||||
def test_process_stats_in_report(self):
|
||||
|
|
@ -456,5 +461,11 @@ class TestProcessStatsCollection:
|
|||
|
||||
report = tracker.get_report()
|
||||
|
||||
assert report.process.rss_bytes > 0
|
||||
assert report.process.rss_mb > 0
|
||||
try:
|
||||
import psutil # noqa: F401
|
||||
|
||||
assert report.process.rss_bytes > 0
|
||||
assert report.process.rss_mb > 0
|
||||
except ImportError:
|
||||
# Without psutil, process stats are zeros — that's expected
|
||||
assert report.process.rss_bytes == 0
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue