From b17c6d81cc39676f322779dc2e5695cba28b8e65 Mon Sep 17 00:00:00 2001 From: JerrettDavis Date: Tue, 21 Apr 2026 21:04:04 -0500 Subject: [PATCH] refactor: extract provider logic into slices Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- README.md | 571 ++- headroom/cli/proxy.py | 1104 ++--- headroom/cli/wrap.py | 4280 +++++++++-------- headroom/client.py | 2139 +++++---- headroom/providers/aider/__init__.py | 5 + headroom/providers/aider/runtime.py | 24 + headroom/providers/claude/__init__.py | 5 + headroom/providers/claude/runtime.py | 10 + headroom/providers/codex/__init__.py | 5 + headroom/providers/codex/runtime.py | 23 + headroom/providers/copilot/__init__.py | 21 + headroom/providers/copilot/wrap.py | 120 + headroom/providers/cursor/__init__.py | 5 + headroom/providers/cursor/runtime.py | 44 + headroom/providers/gemini/__init__.py | 5 + headroom/providers/gemini/runtime.py | 5 + headroom/providers/openclaw/__init__.py | 15 + headroom/providers/openclaw/wrap.py | 89 + headroom/providers/proxy_routes.py | 337 ++ headroom/providers/registry.py | 316 ++ headroom/proxy/models.py | 512 ++- headroom/proxy/server.py | 5584 +++++++++++------------ tests/test_cli/test_wrap_aider.py | 39 + tests/test_cli/test_wrap_copilot.py | 462 +- tests/test_provider_cursor.py | 18 + tests/test_provider_registry.py | 83 + 26 files changed, 8141 insertions(+), 7680 deletions(-) create mode 100644 headroom/providers/aider/__init__.py create mode 100644 headroom/providers/aider/runtime.py create mode 100644 headroom/providers/claude/__init__.py create mode 100644 headroom/providers/claude/runtime.py create mode 100644 headroom/providers/codex/__init__.py create mode 100644 headroom/providers/codex/runtime.py create mode 100644 headroom/providers/copilot/__init__.py create mode 100644 headroom/providers/copilot/wrap.py create mode 100644 headroom/providers/cursor/__init__.py create mode 100644 headroom/providers/cursor/runtime.py create mode 100644 headroom/providers/gemini/__init__.py create mode 100644 headroom/providers/gemini/runtime.py create mode 100644 headroom/providers/openclaw/__init__.py create mode 100644 headroom/providers/openclaw/wrap.py create mode 100644 headroom/providers/proxy_routes.py create mode 100644 headroom/providers/registry.py create mode 100644 tests/test_cli/test_wrap_aider.py create mode 100644 tests/test_provider_cursor.py create mode 100644 tests/test_provider_registry.py diff --git a/README.md b/README.md index 7090a50c8..09d9d3bae 100644 --- a/README.md +++ b/README.md @@ -1,286 +1,285 @@ -
- -# Headroom - -**Compress everything your AI agent reads. Same answers, fraction of the tokens.** - -[![CI](https://github.com/chopratejas/headroom/actions/workflows/ci.yml/badge.svg)](https://github.com/chopratejas/headroom/actions/workflows/ci.yml) -[![codecov](https://codecov.io/gh/chopratejas/headroom/graph/badge.svg)](https://app.codecov.io/gh/chopratejas/headroom) -[![PyPI](https://img.shields.io/pypi/v/headroom-ai.svg)](https://pypi.org/project/headroom-ai/) -[![npm](https://img.shields.io/npm/v/headroom-ai.svg)](https://www.npmjs.com/package/headroom-ai) -[![Model: Kompress-base](https://img.shields.io/badge/model-Kompress--base-yellow.svg)](https://huggingface.co/chopratejas/kompress-base) -[![Tokens saved: 60B+](https://img.shields.io/badge/tokens%20saved-60B%2B-2ea44f)](https://headroomlabs.ai/dashboard) -[![License: Apache 2.0](https://img.shields.io/badge/license-Apache%202.0-blue.svg)](LICENSE) -[![Docs](https://img.shields.io/badge/docs-online-blue.svg)](https://headroom-docs.vercel.app/docs) - -Headroom in action - -
- ---- - -Every tool call, log line, DB read, RAG chunk, and file your agent injects into a prompt is mostly boilerplate. Headroom strips the noise and keeps the signal — **losslessly, locally, and without touching accuracy.** - -> **100 logs. One FATAL error buried at position 67. Both runs found it.** -> Baseline **10,144 tokens** → Headroom **1,260 tokens** — **87% fewer, identical answer.** -> `python examples/needle_in_haystack_test.py` - ---- - -## Quick start - -Works with Anthropic, OpenAI, Google, Bedrock, Vertex, Azure, OpenRouter, and 100+ models via LiteLLM. - -**Wrap your coding agent — one command:** - -```bash -pip install "headroom-ai[all]" - -headroom wrap claude # Claude Code -headroom wrap codex # Codex -headroom wrap cursor # Cursor -headroom wrap aider # Aider -headroom wrap copilot # GitHub Copilot CLI -``` - -**Prefer a one-time durable install instead of wrapping every launch:** - -```bash -headroom init -g # Detect installed user-scoped agents and wire them to Headroom -headroom init claude # Install repo-local Claude hooks for just this project -headroom init copilot -g # Install user-scoped Copilot hooks and provider routing -``` - -**Drop it into your own code — Python or TypeScript:** - -```python -from headroom import compress - -result = compress(messages, model="claude-sonnet-4-5") -response = client.messages.create(model="claude-sonnet-4-5", messages=result.messages) -print(f"Saved {result.tokens_saved} tokens ({result.compression_ratio:.0%})") -``` - -```typescript -import { compress } from 'headroom-ai'; -const result = await compress(messages, { model: 'gpt-4o' }); -``` - -**Or run it as a proxy — zero code changes, any language:** - -```bash -headroom proxy --port 8787 -ANTHROPIC_BASE_URL=http://localhost:8787 your-app -OPENAI_BASE_URL=http://localhost:8787/v1 your-app -``` - ---- - -## Why Headroom - -- **Accuracy-preserving.** GSM8K **0.870 → 0.870** (±0.000). TruthfulQA **+0.030**. SQuAD v2 and BFCL both **97%** accuracy after compression. Validated on public OSS benchmarks you can rerun yourself. -- **Runs on your machine.** No cloud API, no data egress. Compression latency is milliseconds — faster end-to-end for Sonnet / Opus / GPT-4 class models than a hosted service round-trip. -- **[Kompress-base](https://huggingface.co/chopratejas/kompress-base) on HuggingFace.** Our open-source text compressor, fine-tuned on real agentic traces — tool outputs, logs, RAG chunks, code. Install with `pip install "headroom-ai[ml]"`. -- **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). - ---- - -## How it fits - -``` - Your agent / app - (Claude Code, Cursor, Codex, LangChain, Agno, Strands, your own code…) - │ prompts · tool outputs · logs · RAG results · files - ▼ - ┌────────────────────────────────────────────────────┐ - │ Headroom (runs locally — your data stays here) │ - │ ─────────────────────────────────────────────── │ - │ CacheAligner → ContentRouter → CCR │ - │ ├─ SmartCrusher (JSON) │ - │ ├─ CodeCompressor (AST) │ - │ └─ Kompress-base (text, HF) │ - │ │ - │ Cross-agent memory · headroom learn · MCP │ - └────────────────────────────────────────────────────┘ - │ compressed prompt + retrieval tool - ▼ - LLM provider (Anthropic · OpenAI · Bedrock · …) -``` - -→ [Architecture](https://headroom-docs.vercel.app/docs/architecture) · [CCR reversible compression](https://headroom-docs.vercel.app/docs/ccr) · [Kompress-base model card](https://huggingface.co/chopratejas/kompress-base) - -### Canonical pipeline lifecycle - -Headroom now exposes one stable request lifecycle across `compress()`, the SDK, and the proxy: - -`Setup` → `Pre-Start` → `Post-Start` → `Input Received` → `Input Cached` → `Input Routed` → `Input Compressed` → `Input Remembered` → `Pre-Send` → `Post-Send` → `Response Received` - -- **Transforms** still do the work: CacheAligner, ContentRouter, SmartCrusher, CodeCompressor, Kompress-base, IntelligentContext / RollingWindow. -- **Pipeline extensions** observe or customize those lifecycle stages via `on_pipeline_event(...)`. -- **Compression hooks** still work and now sit alongside the canonical lifecycle instead of being the only extension seam. -- **Proxy extensions** remain the server/app integration seam for ASGI middleware, routes, and startup policy. - ---- - -## Proof - -**Savings on real agent workloads:** - -| Workload | Before | After | Savings | -|-------------------------------|-------:|-------:|--------:| -| Code search (100 results) | 17,765 | 1,408 | **92%** | -| SRE incident debugging | 65,694 | 5,118 | **92%** | -| GitHub issue triage | 54,174 | 14,761 | **73%** | -| Codebase exploration | 78,502 | 41,254 | **47%** | - -**Accuracy preserved on standard benchmarks:** - -| Benchmark | Category | N | Baseline | Headroom | Delta | -|------------|----------|----:|---------:|---------:|----------:| -| GSM8K | Math | 100 | 0.870 | 0.870 | **±0.000**| -| TruthfulQA | Factual | 100 | 0.530 | 0.560 | **+0.030**| -| SQuAD v2 | QA | 100 | — | **97%** | 19% compression | -| BFCL | Tools | 100 | — | **97%** | 32% compression | - -Reproduce: - -```bash -python -m headroom.evals suite --tier 1 -``` - -**Community, live:** - -
- - 60B+ tokens saved — community leaderboard - -

60B+ tokens saved by the community in the last 20 days — live leaderboard →

-
- -→ [Full benchmarks & methodology](https://headroom-docs.vercel.app/docs/benchmarks) - ---- - -## Built for coding agents - -| Agent | Durable init / one-shot wrap | Notes | -|--------------------|------------------------------------|------------------------------------------------------------------| -| **Claude Code** | `headroom init claude -g` / `headroom wrap claude` | `init` installs user or repo-local hooks; `wrap` is still useful for ad hoc sessions | -| **Codex** | `headroom init codex -g` / `headroom wrap codex --memory` | `init` installs provider config plus lifecycle hooks where supported | -| **Cursor** | `headroom wrap cursor` | Prints Cursor config — durable init not available yet | -| **Aider** | `headroom wrap aider` | Starts proxy, launches Aider | -| **Copilot CLI** | `headroom init copilot -g` / `headroom wrap copilot` | `init` installs hooks and BYOK provider routing for the current user | -| **OpenClaw** | `headroom init openclaw -g` / `headroom wrap openclaw` | Installs Headroom as ContextEngine plugin | - -MCP-native too — `headroom mcp install` exposes `headroom_compress`, `headroom_retrieve`, and `headroom_stats` to any MCP client. - -
- headroom learn in action -
- ---- - -## Integrations - -
-Drop Headroom into any stack - -| Your setup | Hook in with | -|-------------------------|------------------------------------------------------------------| -| Any Python app | `compress(messages, model=…)` | -| Any TypeScript app | `await compress(messages, { model })` | -| Anthropic / OpenAI SDK | `withHeadroom(new Anthropic())` · `withHeadroom(new OpenAI())` | -| Vercel AI SDK | `wrapLanguageModel({ model, middleware: headroomMiddleware() })` | -| LiteLLM | `litellm.callbacks = [HeadroomCallback()]` | -| LangChain | `HeadroomChatModel(your_llm)` | -| Agno | `HeadroomAgnoModel(your_model)` | -| Strands | [Strands guide](https://headroom-docs.vercel.app/docs/strands) | -| ASGI apps | `app.add_middleware(CompressionMiddleware)` | -| Multi-agent | `SharedContext().put / .get` | -| MCP clients | `headroom mcp install` | - -
- -
-What's inside - -- **SmartCrusher** — universal JSON: arrays of dicts, nested objects, mixed types. -- **CodeCompressor** — AST-aware for Python, JS, Go, Rust, Java, C++. -- **Kompress-base** — our HuggingFace model, trained on agentic traces. -- **Image compression** — 40–90% reduction via trained ML router. -- **CacheAligner** — stabilizes prefixes so Anthropic/OpenAI KV caches actually hit. -- **IntelligentContext** — score-based context fitting with learned importance. -- **CCR** — reversible compression; LLM retrieves originals on demand. -- **Cross-agent memory** — shared store, agent provenance, auto-dedup. -- **SharedContext** — compressed context passing across multi-agent workflows. -- **`headroom learn`** — plugin-based failure mining for Claude, Codex, Gemini. - -
- ---- - -## Install - -```bash -pip install "headroom-ai[all]" # Python, everything -npm install headroom-ai # TypeScript / Node -docker pull ghcr.io/chopratejas/headroom:latest -``` - -Granular extras: `[proxy]`, `[mcp]`, `[ml]` (Kompress-base), `[agno]`, `[langchain]`, `[evals]`. Requires **Python 3.10+**. - -→ [Installation guide](https://headroom-docs.vercel.app/docs/installation) — Docker tags, persistent service, PowerShell, devcontainers. - ---- - -## Documentation - -| Start here | Go deeper | -|-------------------------------------------------------------------------|------------------------------------------------------------------------| -| [Quickstart](https://headroom-docs.vercel.app/docs/quickstart) | [Architecture](https://headroom-docs.vercel.app/docs/architecture) | -| [Proxy](https://headroom-docs.vercel.app/docs/proxy) | [How compression works](https://headroom-docs.vercel.app/docs/how-compression-works) | -| [MCP tools](https://headroom-docs.vercel.app/docs/mcp) | [CCR — reversible compression](https://headroom-docs.vercel.app/docs/ccr) | -| [Memory](https://headroom-docs.vercel.app/docs/memory) | [Cache optimization](https://headroom-docs.vercel.app/docs/cache-optimization) | -| [Failure learning](https://headroom-docs.vercel.app/docs/failure-learning) | [Benchmarks](https://headroom-docs.vercel.app/docs/benchmarks) | -| [Configuration](https://headroom-docs.vercel.app/docs/configuration) | [Limitations](https://headroom-docs.vercel.app/docs/limitations) | - ---- - -## Compared to - -Headroom runs **locally**, covers **every** content type (not just CLI or text), works with every major framework, and is **reversible**. - -| | Scope | Deploy | Local | Reversible | -|----------------------------------|-------------------------------------------------|-------------------------------------|:-----:|:----------:| -| **Headroom** | All context — tools, RAG, logs, files, history | Proxy · library · middleware · MCP | Yes | Yes | -| [RTK](https://github.com/rtk-ai/rtk) | CLI command outputs | CLI wrapper | Yes | No | -| [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. - ---- - -## Contributing - -```bash -git clone https://github.com/chopratejas/headroom.git && cd headroom -pip install -e ".[dev]" && pytest -``` - -Devcontainers in `.devcontainer/` (default + `memory-stack` with Qdrant & Neo4j). See [CONTRIBUTING.md](CONTRIBUTING.md). - ---- - -## Community - -- **[Live leaderboard](https://headroomlabs.ai/dashboard)** — 60B+ tokens saved and counting. -- **[Discord](https://discord.gg/yRmaUNpsPJ)** — questions, feedback, war stories. -- **[Kompress-base on HuggingFace](https://huggingface.co/chopratejas/kompress-base)** — the model behind our text compression. - -## License - -Apache 2.0 — see [LICENSE](LICENSE). +
+ +# Headroom + +**Compress everything your AI agent reads. Same answers, fraction of the tokens.** + +[![CI](https://github.com/chopratejas/headroom/actions/workflows/ci.yml/badge.svg)](https://github.com/chopratejas/headroom/actions/workflows/ci.yml) +[![PyPI](https://img.shields.io/pypi/v/headroom-ai.svg)](https://pypi.org/project/headroom-ai/) +[![npm](https://img.shields.io/npm/v/headroom-ai.svg)](https://www.npmjs.com/package/headroom-ai) +[![Model: Kompress-base](https://img.shields.io/badge/model-Kompress--base-yellow.svg)](https://huggingface.co/chopratejas/kompress-base) +[![Tokens saved: 60B+](https://img.shields.io/badge/tokens%20saved-60B%2B-2ea44f)](https://headroomlabs.ai/dashboard) +[![License: Apache 2.0](https://img.shields.io/badge/license-Apache%202.0-blue.svg)](LICENSE) +[![Docs](https://img.shields.io/badge/docs-online-blue.svg)](https://headroom-docs.vercel.app/docs) + +Headroom in action + +
+ +--- + +Every tool call, log line, DB read, RAG chunk, and file your agent injects into a prompt is mostly boilerplate. Headroom strips the noise and keeps the signal — **losslessly, locally, and without touching accuracy.** + +> **100 logs. One FATAL error buried at position 67. Both runs found it.** +> Baseline **10,144 tokens** → Headroom **1,260 tokens** — **87% fewer, identical answer.** +> `python examples/needle_in_haystack_test.py` + +--- + +## Quick start + +Works with Anthropic, OpenAI, Google, Bedrock, Vertex, Azure, OpenRouter, and 100+ models via LiteLLM. + +**Wrap your coding agent — one command:** + +```bash +pip install "headroom-ai[all]" + +headroom wrap claude # Claude Code +headroom wrap codex # Codex +headroom wrap cursor # Cursor +headroom wrap aider # Aider +headroom wrap copilot # GitHub Copilot CLI +``` + +**Drop it into your own code — Python or TypeScript:** + +```python +from headroom import compress + +result = compress(messages, model="claude-sonnet-4-5") +response = client.messages.create(model="claude-sonnet-4-5", messages=result.messages) +print(f"Saved {result.tokens_saved} tokens ({result.compression_ratio:.0%})") +``` + +```typescript +import { compress } from 'headroom-ai'; +const result = await compress(messages, { model: 'gpt-4o' }); +``` + +**Or run it as a proxy — zero code changes, any language:** + +```bash +headroom proxy --port 8787 +ANTHROPIC_BASE_URL=http://localhost:8787 your-app +OPENAI_BASE_URL=http://localhost:8787/v1 your-app +``` + +--- + +## Why Headroom + +- **Accuracy-preserving.** GSM8K **0.870 → 0.870** (±0.000). TruthfulQA **+0.030**. SQuAD v2 and BFCL both **97%** accuracy after compression. Validated on public OSS benchmarks you can rerun yourself. +- **Runs on your machine.** No cloud API, no data egress. Compression latency is milliseconds — faster end-to-end for Sonnet / Opus / GPT-4 class models than a hosted service round-trip. +- **[Kompress-base](https://huggingface.co/chopratejas/kompress-base) on HuggingFace.** Our open-source text compressor, fine-tuned on real agentic traces — tool outputs, logs, RAG chunks, code. Install with `pip install "headroom-ai[ml]"`. +- **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). + +--- + +## How it fits + +``` + Your agent / app + (Claude Code, Cursor, Codex, LangChain, Agno, Strands, your own code…) + │ prompts · tool outputs · logs · RAG results · files + ▼ + ┌────────────────────────────────────────────────────┐ + │ Headroom (runs locally — your data stays here) │ + │ ─────────────────────────────────────────────── │ + │ CacheAligner → ContentRouter → CCR │ + │ ├─ SmartCrusher (JSON) │ + │ ├─ CodeCompressor (AST) │ + │ └─ Kompress-base (text, HF) │ + │ │ + │ Cross-agent memory · headroom learn · MCP │ + └────────────────────────────────────────────────────┘ + │ compressed prompt + retrieval tool + ▼ + LLM provider (Anthropic · OpenAI · Bedrock · …) +``` + +→ [Architecture](https://headroom-docs.vercel.app/docs/architecture) · [CCR reversible compression](https://headroom-docs.vercel.app/docs/ccr) · [Kompress-base model card](https://huggingface.co/chopratejas/kompress-base) + +### Canonical pipeline lifecycle + +Headroom now exposes one stable request lifecycle across `compress()`, the SDK, and the proxy: + +`Setup` → `Pre-Start` → `Post-Start` → `Input Received` → `Input Cached` → `Input Routed` → `Input Compressed` → `Input Remembered` → `Pre-Send` → `Post-Send` → `Response Received` + +- **Transforms** still do the work: CacheAligner, ContentRouter, SmartCrusher, CodeCompressor, Kompress-base, IntelligentContext / RollingWindow. +- **Pipeline extensions** observe or customize those lifecycle stages via `on_pipeline_event(...)`. +- **Compression hooks** still work and now sit alongside the canonical lifecycle instead of being the only extension seam. +- **Proxy extensions** remain the server/app integration seam for ASGI middleware, routes, and startup policy. + +### Provider slices + +Provider and tool-specific behavior is being moved behind dedicated modules under `headroom/providers/` so core orchestration stays focused on lifecycle, sequencing, and policy. + +- **CLI/tool slices**: `headroom/providers/claude`, `copilot`, `codex`, `openclaw` +- **Provider runtime slices**: `headroom/providers/claude`, `gemini`, plus shared backend/runtime dispatch in `headroom/providers/registry.py` +- **Core files stay orchestration-first**: `wrap.py`, `client.py`, `cli/proxy.py`, and `proxy/server.py` now delegate provider-specific env shaping, API target normalization, backend selection, and transport dispatch instead of inlining those rules. + +--- + +## Proof + +**Savings on real agent workloads:** + +| Workload | Before | After | Savings | +|-------------------------------|-------:|-------:|--------:| +| Code search (100 results) | 17,765 | 1,408 | **92%** | +| SRE incident debugging | 65,694 | 5,118 | **92%** | +| GitHub issue triage | 54,174 | 14,761 | **73%** | +| Codebase exploration | 78,502 | 41,254 | **47%** | + +**Accuracy preserved on standard benchmarks:** + +| Benchmark | Category | N | Baseline | Headroom | Delta | +|------------|----------|----:|---------:|---------:|----------:| +| GSM8K | Math | 100 | 0.870 | 0.870 | **±0.000**| +| TruthfulQA | Factual | 100 | 0.530 | 0.560 | **+0.030**| +| SQuAD v2 | QA | 100 | — | **97%** | 19% compression | +| BFCL | Tools | 100 | — | **97%** | 32% compression | + +Reproduce: + +```bash +python -m headroom.evals suite --tier 1 +``` + +**Community, live:** + +
+ + 60B+ tokens saved — community leaderboard + +

60B+ tokens saved by the community in the last 20 days — live leaderboard →

+
+ +→ [Full benchmarks & methodology](https://headroom-docs.vercel.app/docs/benchmarks) + +--- + +## Built for coding agents + +| Agent | One-command wrap | Notes | +|--------------------|------------------------------------|------------------------------------------------------------------| +| **Claude Code** | `headroom wrap claude` | `--memory` for cross-agent memory, `--code-graph` for codebase intel | +| **Codex** | `headroom wrap codex --memory` | Shares the same memory store as Claude | +| **Cursor** | `headroom wrap cursor` | Prints Cursor config — paste once, done | +| **Aider** | `headroom wrap aider` | Starts proxy, launches Aider | +| **Copilot CLI** | `headroom wrap copilot` | Starts proxy, launches Copilot | +| **OpenClaw** | `headroom wrap openclaw` | Installs Headroom as ContextEngine plugin | + +MCP-native too — `headroom mcp install` exposes `headroom_compress`, `headroom_retrieve`, and `headroom_stats` to any MCP client. + +
+ headroom learn in action +
+ +--- + +## Integrations + +
+Drop Headroom into any stack + +| Your setup | Hook in with | +|-------------------------|------------------------------------------------------------------| +| Any Python app | `compress(messages, model=…)` | +| Any TypeScript app | `await compress(messages, { model })` | +| Anthropic / OpenAI SDK | `withHeadroom(new Anthropic())` · `withHeadroom(new OpenAI())` | +| Vercel AI SDK | `wrapLanguageModel({ model, middleware: headroomMiddleware() })` | +| LiteLLM | `litellm.callbacks = [HeadroomCallback()]` | +| LangChain | `HeadroomChatModel(your_llm)` | +| Agno | `HeadroomAgnoModel(your_model)` | +| Strands | [Strands guide](https://headroom-docs.vercel.app/docs/strands) | +| ASGI apps | `app.add_middleware(CompressionMiddleware)` | +| Multi-agent | `SharedContext().put / .get` | +| MCP clients | `headroom mcp install` | + +
+ +
+What's inside + +- **SmartCrusher** — universal JSON: arrays of dicts, nested objects, mixed types. +- **CodeCompressor** — AST-aware for Python, JS, Go, Rust, Java, C++. +- **Kompress-base** — our HuggingFace model, trained on agentic traces. +- **Image compression** — 40–90% reduction via trained ML router. +- **CacheAligner** — stabilizes prefixes so Anthropic/OpenAI KV caches actually hit. +- **IntelligentContext** — score-based context fitting with learned importance. +- **CCR** — reversible compression; LLM retrieves originals on demand. +- **Cross-agent memory** — shared store, agent provenance, auto-dedup. +- **SharedContext** — compressed context passing across multi-agent workflows. +- **`headroom learn`** — plugin-based failure mining for Claude, Codex, Gemini. + +
+ +--- + +## Install + +```bash +pip install "headroom-ai[all]" # Python, everything +npm install headroom-ai # TypeScript / Node +docker pull ghcr.io/chopratejas/headroom:latest +``` + +Granular extras: `[proxy]`, `[mcp]`, `[ml]` (Kompress-base), `[agno]`, `[langchain]`, `[evals]`. Requires **Python 3.10+**. + +→ [Installation guide](https://headroom-docs.vercel.app/docs/installation) — Docker tags, persistent service, PowerShell, devcontainers. + +--- + +## Documentation + +| Start here | Go deeper | +|-------------------------------------------------------------------------|------------------------------------------------------------------------| +| [Quickstart](https://headroom-docs.vercel.app/docs/quickstart) | [Architecture](https://headroom-docs.vercel.app/docs/architecture) | +| [Proxy](https://headroom-docs.vercel.app/docs/proxy) | [How compression works](https://headroom-docs.vercel.app/docs/how-compression-works) | +| [MCP tools](https://headroom-docs.vercel.app/docs/mcp) | [CCR — reversible compression](https://headroom-docs.vercel.app/docs/ccr) | +| [Memory](https://headroom-docs.vercel.app/docs/memory) | [Cache optimization](https://headroom-docs.vercel.app/docs/cache-optimization) | +| [Failure learning](https://headroom-docs.vercel.app/docs/failure-learning) | [Benchmarks](https://headroom-docs.vercel.app/docs/benchmarks) | +| [Configuration](https://headroom-docs.vercel.app/docs/configuration) | [Limitations](https://headroom-docs.vercel.app/docs/limitations) | + +--- + +## Compared to + +Headroom runs **locally**, covers **every** content type (not just CLI or text), works with every major framework, and is **reversible**. + +| | Scope | Deploy | Local | Reversible | +|----------------------------------|-------------------------------------------------|-------------------------------------|:-----:|:----------:| +| **Headroom** | All context — tools, RAG, logs, files, history | Proxy · library · middleware · MCP | Yes | Yes | +| [RTK](https://github.com/rtk-ai/rtk) | CLI command outputs | CLI wrapper | Yes | No | +| [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. + +--- + +## Contributing + +```bash +git clone https://github.com/chopratejas/headroom.git && cd headroom +pip install -e ".[dev]" && pytest +``` + +Devcontainers in `.devcontainer/` (default + `memory-stack` with Qdrant & Neo4j). See [CONTRIBUTING.md](CONTRIBUTING.md). + +--- + +## Community + +- **[Live leaderboard](https://headroomlabs.ai/dashboard)** — 60B+ tokens saved and counting. +- **[Discord](https://discord.gg/yRmaUNpsPJ)** — questions, feedback, war stories. +- **[Kompress-base on HuggingFace](https://huggingface.co/chopratejas/kompress-base)** — the model behind our text compression. + +## License + +Apache 2.0 — see [LICENSE](LICENSE). diff --git a/headroom/cli/proxy.py b/headroom/cli/proxy.py index 5564ff639..7204406ae 100644 --- a/headroom/cli/proxy.py +++ b/headroom/cli/proxy.py @@ -1,550 +1,554 @@ -"""Proxy server CLI commands.""" - -import os -import sys - -import click - -from headroom.proxy.modes import PROXY_MODE_TOKEN, normalize_proxy_mode - -from .main import main - - -@main.command() -@click.option( - "--host", - default="127.0.0.1", - envvar="HEADROOM_HOST", - help="Host to bind to (default: 127.0.0.1, env: HEADROOM_HOST)", -) -@click.option( - "--port", - "-p", - default=8787, - type=int, - envvar="HEADROOM_PORT", - help="Port to bind to (default: 8787, env: HEADROOM_PORT)", -) -@click.option( - "--mode", - default=None, - type=click.Choice( - [ - "token", - "cache", - "token_mode", - "cache_mode", - "token_savings", - "cost_savings", - "token_headroom", - ] - ), - help=( - "Optimization mode: token (prioritize compression) or cache " - "(freeze prior turns for prefix-cache stability). " - "Legacy aliases are accepted. Default: token. Env: HEADROOM_MODE" - ), -) -@click.option( - "--intercept-tool-results", - is_flag=True, - help=( - "Opt in to tool_result interceptors (ast-grep Read outliner, etc.). " - "Off by default while this feature ships." - ), -) -@click.option("--no-optimize", is_flag=True, help="Disable optimization (passthrough mode)") -@click.option("--no-cache", is_flag=True, help="Disable semantic caching") -@click.option("--no-rate-limit", is_flag=True, help="Disable rate limiting") -@click.option( - "--retry-max-attempts", - type=int, - default=None, - help="Maximum upstream retry attempts for connect/read/5xx failures (default: 3)", -) -@click.option( - "--connect-timeout-seconds", - type=int, - default=None, - help="Upstream connection timeout in seconds (default: 10)", -) -@click.option( - "--anthropic-pre-upstream-concurrency", - type=int, - default=None, - envvar="HEADROOM_ANTHROPIC_PRE_UPSTREAM_CONCURRENCY", - help=( - "Cap the number of Anthropic HTTP requests that may run pre-upstream work " - "(request parse / deep-copy / first compression stage / memory context / upstream connect) " - "concurrently. Prevents cold-start replay storms from starving /livez and new Codex WS opens. " - "Default: max(2, min(8, os.cpu_count() or 4)). " - "Set to 0 or negative to disable (unbounded). " - "Env: HEADROOM_ANTHROPIC_PRE_UPSTREAM_CONCURRENCY." - ), -) -@click.option( - "--anthropic-pre-upstream-acquire-timeout-seconds", - type=float, - default=None, - envvar="HEADROOM_ANTHROPIC_PRE_UPSTREAM_ACQUIRE_TIMEOUT_SECONDS", - help=( - "Fail-fast timeout for waiting on the Anthropic pre-upstream semaphore " - "before returning 503 + Retry-After. " - "Default: 15.0 seconds. " - "Env: HEADROOM_ANTHROPIC_PRE_UPSTREAM_ACQUIRE_TIMEOUT_SECONDS." - ), -) -@click.option( - "--anthropic-pre-upstream-memory-context-timeout-seconds", - type=float, - default=None, - envvar="HEADROOM_ANTHROPIC_PRE_UPSTREAM_MEMORY_CONTEXT_TIMEOUT_SECONDS", - help=( - "Fail-open timeout for Anthropic memory-context lookup while the request " - "still holds a pre-upstream slot. " - "Default: 2.0 seconds. " - "Env: HEADROOM_ANTHROPIC_PRE_UPSTREAM_MEMORY_CONTEXT_TIMEOUT_SECONDS." - ), -) -@click.option("--log-file", default=None, help="Path to JSONL log file") -@click.option( - "--log-messages", - is_flag=True, - help="Enable full message logging (request/response content stored for live feed)", -) -@click.option( - "--budget", - type=float, - default=None, - envvar="HEADROOM_BUDGET", - help="Daily budget limit in USD (env: HEADROOM_BUDGET)", -) -# Code graph: indexes project + watches files for live reindex via codebase-memory-mcp -@click.option( - "--code-graph", - is_flag=True, - help="Enable code graph intelligence (indexes project, watches files for live reindex via codebase-memory-mcp)", -) -# Read lifecycle (ON by default: compresses stale/superseded Read outputs) -@click.option( - "--no-read-lifecycle", - is_flag=True, - help="Disable Read lifecycle management (stale/superseded Read compression)", -) -# Intelligent Context Management (ON by default) -@click.option( - "--no-intelligent-context", - is_flag=True, - help="Disable IntelligentContextManager (fall back to RollingWindow)", -) -@click.option( - "--no-intelligent-scoring", - is_flag=True, - help="Disable multi-factor importance scoring (use position-based)", -) -@click.option( - "--no-compress-first", - is_flag=True, - help="Disable trying deeper compression before dropping messages", -) -# Memory System (Multi-Provider Support) -@click.option( - "--memory", - is_flag=True, - help="Enable persistent user memory. Auto-detects provider and uses appropriate tool format. " - "Set x-headroom-user-id header for per-user memory (defaults to 'default').", -) -@click.option( - "--memory-db-path", - default="", - help="Path to memory database file (default: {cwd}/.headroom/memory.db)", -) -@click.option("--no-memory-tools", is_flag=True, help="Disable automatic memory tool injection") -@click.option( - "--no-memory-context", is_flag=True, help="Disable automatic memory context injection" -) -@click.option( - "--memory-top-k", - type=int, - default=10, - help="Number of memories to inject as context (default: 10)", -) -# Traffic Learning (live pattern extraction from proxy traffic) -@click.option( - "--learn", - is_flag=True, - help="Enable live traffic learning: extract error→recovery patterns, environment facts, " - "and user preferences from proxy traffic. Implies --memory. " - "Learned patterns are saved to agent-native memory files (MEMORY.md, .cursor/rules, AGENTS.md).", -) -@click.option( - "--no-learn", - is_flag=True, - help="Explicitly disable traffic learning even when --memory is set.", -) -# Backend configuration -@click.option( - "--backend", - default="anthropic", - help=( - "API backend: 'anthropic' (direct), 'bedrock' (AWS), 'openrouter' (OpenRouter), " - "'anyllm' (any-llm), or 'litellm-' (e.g., litellm-vertex)" - ), -) -@click.option( - "--anyllm-provider", - default="openai", - help="Provider for any-llm backend: openai, mistral, groq, ollama, etc. (default: openai)", -) -@click.option( - "--anthropic-api-url", - default=None, - help="Custom Anthropic API URL for passthrough endpoints (env: ANTHROPIC_TARGET_API_URL)", -) -@click.option( - "--openai-api-url", - default=None, - help="Custom OpenAI API URL for passthrough endpoints (env: OPENAI_TARGET_API_URL)", -) -@click.option( - "--gemini-api-url", - default=None, - help="Custom Gemini API URL for passthrough endpoints (env: GEMINI_TARGET_API_URL)", -) -@click.option( - "--cloudcode-api-url", - default=None, - help="Custom Cloud Code Assist API URL for compatibility endpoints (env: CLOUDCODE_TARGET_API_URL)", -) -@click.option( - "--region", - default="us-west-2", - help="Cloud region for Bedrock/Vertex/etc (default: us-west-2)", -) -@click.option( - "--bedrock-region", - default=None, - help="(deprecated, use --region) AWS region for Bedrock", -) -@click.option( - "--bedrock-profile", - default=None, - help="AWS profile name for Bedrock (default: use default credentials)", -) -@click.option( - "--no-telemetry", - is_flag=True, - help="Disable anonymous usage telemetry (env: HEADROOM_TELEMETRY=off)", -) -@click.option( - "--stateless", - is_flag=True, - help="Disable all filesystem writes — run purely in-memory. " - "For containerized / read-only / load-balanced deployments. " - "(env: HEADROOM_STATELESS=true)", -) -@click.pass_context -def proxy( - ctx: click.Context, - mode: str | None, - host: str, - port: int, - intercept_tool_results: bool, - no_optimize: bool, - no_cache: bool, - no_rate_limit: bool, - retry_max_attempts: int | None, - connect_timeout_seconds: int | None, - anthropic_pre_upstream_concurrency: int | None, - anthropic_pre_upstream_acquire_timeout_seconds: float | None, - anthropic_pre_upstream_memory_context_timeout_seconds: float | None, - log_file: str | None, - log_messages: bool, - budget: float | None, - code_graph: bool, - no_read_lifecycle: bool, - no_intelligent_context: bool, - no_intelligent_scoring: bool, - no_compress_first: bool, - memory: bool, - memory_db_path: str, - no_memory_tools: bool, - no_memory_context: bool, - memory_top_k: int, - learn: bool, - no_learn: bool, - backend: str, - anyllm_provider: str, - anthropic_api_url: str | None, - openai_api_url: str | None, - gemini_api_url: str | None, - cloudcode_api_url: str | None, - region: str, - bedrock_region: str | None, - bedrock_profile: str | None, - no_telemetry: bool, - stateless: bool, -) -> None: - """Start the optimization proxy server. - - \b - Examples: - headroom proxy Start proxy on port 8787 - headroom proxy --port 8080 Start proxy on port 8080 - headroom proxy --no-optimize Passthrough mode (no optimization) - - \b - Usage with Claude Code: - ANTHROPIC_BASE_URL=http://localhost:8787 claude - - \b - Usage with OpenAI-compatible clients: - OPENAI_BASE_URL=http://localhost:8787/v1 your-app - """ - # Import here to avoid slow startup - try: - from headroom.proxy.server import ProxyConfig, run_server - except ImportError as e: - click.echo("Error: Proxy dependencies not installed. Run: pip install headroom[proxy]") - click.echo(f"Details: {e}") - raise SystemExit(1) from None - - # Opt-in: turn on tool_result interceptors (ast-grep Read outline, etc.). - # Only fetch the bundled CLI tool binaries when the feature is enabled — - # otherwise we'd pay a network round-trip and risk a readonly-FS failure - # for capabilities the user hasn't asked for. The TransformPipeline reads - # this env var at construction time. - if intercept_tool_results: - from headroom.binaries import ensure_tools - - resolved_tools = ensure_tools() - critical_tools = ["ast-grep"] - missing = [t for t in critical_tools if not resolved_tools.get(t)] - if missing: - # User explicitly opted in — fail fast rather than silently starting - # with non-functional interceptors. They can retry with the tool - # installed, or drop the flag if they want pass-through behavior. - click.secho( - f"error: --intercept-tool-results requires tool(s) that could not " - f"be installed: {missing}. Run `headroom tools doctor` to diagnose, " - "or omit the flag to start the proxy without interceptors.", - fg="red", - err=True, - ) - sys.exit(1) - os.environ["HEADROOM_INTERCEPT_ENABLED"] = "1" - - # Resolve API URL overrides: CLI flag > env var > None - effective_anthropic_api_url = anthropic_api_url or os.environ.get("ANTHROPIC_TARGET_API_URL") - effective_openai_api_url = openai_api_url or os.environ.get("OPENAI_TARGET_API_URL") - effective_gemini_api_url = gemini_api_url or os.environ.get("GEMINI_TARGET_API_URL") - effective_cloudcode_api_url = cloudcode_api_url or os.environ.get("CLOUDCODE_TARGET_API_URL") - - # Resolve anyllm provider: env var takes precedence over CLI default (matches argparse path) - effective_anyllm_provider = os.environ.get("HEADROOM_ANYLLM_PROVIDER") or anyllm_provider - - # Resolve mode: CLI flag > env var > default - effective_mode: str = normalize_proxy_mode( - mode or os.environ.get("HEADROOM_MODE") or PROXY_MODE_TOKEN - ) - - # Stateless mode: CLI flag or env var - is_stateless = stateless or os.environ.get("HEADROOM_STATELESS", "").lower() in ( - "true", - "1", - "yes", - "on", - ) - - # Telemetry opt-out: --no-telemetry flag sets the env var - if no_telemetry: - os.environ["HEADROOM_TELEMETRY"] = "off" - - # Stateless mode: suppress TOIN filesystem persistence - if is_stateless: - os.environ["HEADROOM_TOIN_BACKEND"] = "none" - - # License key for managed/enterprise deployments (optional) - license_key = os.environ.get("HEADROOM_LICENSE_KEY") - - config = ProxyConfig( - host=host, - port=port, - anthropic_api_url=effective_anthropic_api_url, - openai_api_url=effective_openai_api_url, - gemini_api_url=effective_gemini_api_url, - cloudcode_api_url=effective_cloudcode_api_url, - mode=effective_mode, - optimize=not no_optimize, - cache_enabled=not no_cache, - rate_limit_enabled=not no_rate_limit, - retry_max_attempts=retry_max_attempts if retry_max_attempts is not None else 3, - connect_timeout_seconds=connect_timeout_seconds - if connect_timeout_seconds is not None - else 10, - log_file=None if is_stateless else log_file, - log_full_messages=log_messages - or os.environ.get("HEADROOM_LOG_MESSAGES", "").lower() in ("true", "1", "yes", "on"), - budget_limit_usd=budget, - # Code graph: live file watcher for incremental reindexing - code_graph_watcher=code_graph, - # Read lifecycle: ON by default (use --no-read-lifecycle to disable) - read_lifecycle=not no_read_lifecycle, - # Intelligent Context: ON by default (use --no-intelligent-context to disable) - intelligent_context=not no_intelligent_context, - intelligent_context_scoring=not no_intelligent_scoring, - intelligent_context_compress_first=not no_compress_first, - # Memory System (Multi-Provider with auto-detection) - # --learn implies --memory (need backend for storing patterns) - # Stateless mode disables memory (requires SQLite on disk) - memory_enabled=False if is_stateless else (memory or (learn and not no_learn)), - memory_db_path=memory_db_path, - memory_inject_tools=not no_memory_tools, - memory_inject_context=not no_memory_context, - memory_top_k=memory_top_k, - # Traffic Learning: only with --learn, never with --no-learn - # Stateless mode disables learning (requires filesystem) - traffic_learning_enabled=False if is_stateless else (learn and not no_learn), - traffic_learning_agent_type=os.environ.get("HEADROOM_AGENT_TYPE", "unknown"), - # Backend (Anthropic direct, Bedrock, LiteLLM, or any-llm) - backend=backend, - bedrock_region=bedrock_region or region, - bedrock_profile=bedrock_profile, - anyllm_provider=effective_anyllm_provider, - # License / Usage Reporting (managed/enterprise) - license_key=license_key, - # Stateless mode: disable all filesystem writes - stateless=is_stateless, - # Unit 4: bounded pre-upstream concurrency on the Anthropic HTTP - # path. ``None`` -> HeadroomProxy computes ``max(2, min(8, - # os.cpu_count() or 4))``; ``<= 0`` -> disabled (unbounded). - # Precedence: CLI > env > auto-compute (click's ``envvar`` - # handles the env-var fallback). - anthropic_pre_upstream_concurrency=anthropic_pre_upstream_concurrency, - anthropic_pre_upstream_acquire_timeout_seconds=( - anthropic_pre_upstream_acquire_timeout_seconds - if anthropic_pre_upstream_acquire_timeout_seconds is not None - else 15.0 - ), - anthropic_pre_upstream_memory_context_timeout_seconds=( - anthropic_pre_upstream_memory_context_timeout_seconds - if anthropic_pre_upstream_memory_context_timeout_seconds is not None - else 2.0 - ), - ) - - memory_status = "DISABLED" - if config.memory_enabled: - memory_status = "ENABLED (multi-provider)" - - license_status = "OSS (no license key)" - if license_key: - license_status = f"MANAGED (key={license_key[:8]}...)" - - anthropic_url = config.anthropic_api_url or "https://api.anthropic.com" - openai_url = config.openai_api_url or "https://api.openai.com" - cloudcode_url = config.cloudcode_api_url or "https://cloudcode-pa.googleapis.com" - backend_section = "" - - if config.backend == "anyllm" or config.backend.startswith("anyllm-"): - # any-llm backend - backend_section = """ - Set credentials for your provider (e.g., OPENAI_API_KEY, MISTRAL_API_KEY) - Providers: https://mozilla-ai.github.io/any-llm/providers/ -""" - elif config.backend != "anthropic": - # LiteLLM backend - from headroom.backends.litellm import get_provider_config - - provider = config.backend.replace("litellm-", "") - provider_config = get_provider_config(provider) - - # Build usage instructions from provider config - env_vars_str = ( - ", ".join(provider_config.env_vars) if provider_config.env_vars else "See docs" - ) - backend_section = f""" -IMPORTANT for {provider_config.display_name} users: - 1. Set credentials: {env_vars_str} - 2. Set a dummy Anthropic key: ANTHROPIC_API_KEY="sk-ant-dummy" - (Headroom ignores this - it uses your {provider_config.display_name} credentials) - 3. Set base URL: ANTHROPIC_BASE_URL=http://{config.host}:{config.port}""" - if provider_config.model_format_hint: - backend_section += f"\n 4. Use model names: {provider_config.model_format_hint}" - backend_section += "\n" - - # Build memory section if enabled - memory_section = "" - if config.memory_enabled: - memory_section = f""" -Memory (Multi-Provider): - - Auto-detects provider from request (Anthropic, OpenAI, Gemini, etc.) - - Anthropic: Uses native memory tool (memory_20250818) - subscription safe - - OpenAI/Gemini/Others: Uses function calling format - - All providers share the same semantic vector store backend - - Set x-headroom-user-id header for per-user memory (defaults to 'default') - - Tools: {"ENABLED" if config.memory_inject_tools else "DISABLED"} - - Context injection: {"ENABLED" if config.memory_inject_context else "DISABLED"} - - Database: {config.memory_db_path} -""" - - # Stateless mode warning - stateless_line = "" - if is_stateless: - stateless_line = ( - " Stateless: YES (no filesystem writes — memory, logs, TOIN disabled)\n" - ) - - from headroom.telemetry.beacon import is_telemetry_enabled - - # Build telemetry section for the startup banner - if is_telemetry_enabled(): - telemetry_line = ( - " Telemetry: ENABLED (anonymous aggregate stats)\n" - " Disable: HEADROOM_TELEMETRY=off or headroom proxy --no-telemetry" - ) - else: - telemetry_line = " Telemetry: DISABLED" - - click.echo(f""" -╔═══════════════════════════════════════════════════════════════════════╗ -║ HEADROOM PROXY ║ -║ The Context Optimization Layer for LLM Applications ║ -╚═══════════════════════════════════════════════════════════════════════╝ - -Starting proxy server... - - URL: http://{config.host}:{config.port} - Mode: {config.mode} - Optimization: {"ENABLED" if config.optimize else "DISABLED"} - Caching: {"ENABLED" if config.cache_enabled else "DISABLED"} - Rate Limit: {"ENABLED" if config.rate_limit_enabled else "DISABLED"} - Memory: {memory_status} - License: {license_status} -{stateless_line}{telemetry_line} -{backend_section} -Routing: - /v1/messages → {anthropic_url} - /v1/chat/completions → {openai_url} - /v1/responses → {openai_url} (HTTP + WebSocket) - /v1internal:streamGenerateContent → {cloudcode_url} - -Usage: - Claude Code: ANTHROPIC_BASE_URL=http://{config.host}:{config.port} claude - Codex / OpenAI: OPENAI_BASE_URL=http://{config.host}:{config.port}/v1 your-app -{memory_section} -Endpoints: - GET /livez Process liveness - GET /readyz Traffic readiness - GET /health Aggregate health - GET /stats Detailed statistics - GET /stats-history Durable compression history + display session - GET /metrics Prometheus metrics - -Press Ctrl+C to stop. -""") - - try: - run_server(config) - except KeyboardInterrupt: - click.echo("\nShutting down...") +"""Proxy server CLI commands.""" + +import os +import sys + +import click + +from headroom.providers.registry import resolve_api_overrides, resolve_api_targets +from headroom.proxy.modes import PROXY_MODE_TOKEN, normalize_proxy_mode + +from .main import main + + +@main.command() +@click.option( + "--host", + default="127.0.0.1", + envvar="HEADROOM_HOST", + help="Host to bind to (default: 127.0.0.1, env: HEADROOM_HOST)", +) +@click.option( + "--port", + "-p", + default=8787, + type=int, + envvar="HEADROOM_PORT", + help="Port to bind to (default: 8787, env: HEADROOM_PORT)", +) +@click.option( + "--mode", + default=None, + type=click.Choice( + [ + "token", + "cache", + "token_mode", + "cache_mode", + "token_savings", + "cost_savings", + "token_headroom", + ] + ), + help=( + "Optimization mode: token (prioritize compression) or cache " + "(freeze prior turns for prefix-cache stability). " + "Legacy aliases are accepted. Default: token. Env: HEADROOM_MODE" + ), +) +@click.option( + "--intercept-tool-results", + is_flag=True, + help=( + "Opt in to tool_result interceptors (ast-grep Read outliner, etc.). " + "Off by default while this feature ships." + ), +) +@click.option("--no-optimize", is_flag=True, help="Disable optimization (passthrough mode)") +@click.option("--no-cache", is_flag=True, help="Disable semantic caching") +@click.option("--no-rate-limit", is_flag=True, help="Disable rate limiting") +@click.option( + "--retry-max-attempts", + type=int, + default=None, + help="Maximum upstream retry attempts for connect/read/5xx failures (default: 3)", +) +@click.option( + "--connect-timeout-seconds", + type=int, + default=None, + help="Upstream connection timeout in seconds (default: 10)", +) +@click.option( + "--anthropic-pre-upstream-concurrency", + type=int, + default=None, + envvar="HEADROOM_ANTHROPIC_PRE_UPSTREAM_CONCURRENCY", + help=( + "Cap the number of Anthropic HTTP requests that may run pre-upstream work " + "(request parse / deep-copy / first compression stage / memory context / upstream connect) " + "concurrently. Prevents cold-start replay storms from starving /livez and new Codex WS opens. " + "Default: max(2, min(8, os.cpu_count() or 4)). " + "Set to 0 or negative to disable (unbounded). " + "Env: HEADROOM_ANTHROPIC_PRE_UPSTREAM_CONCURRENCY." + ), +) +@click.option( + "--anthropic-pre-upstream-acquire-timeout-seconds", + type=float, + default=None, + envvar="HEADROOM_ANTHROPIC_PRE_UPSTREAM_ACQUIRE_TIMEOUT_SECONDS", + help=( + "Fail-fast timeout for waiting on the Anthropic pre-upstream semaphore " + "before returning 503 + Retry-After. " + "Default: 15.0 seconds. " + "Env: HEADROOM_ANTHROPIC_PRE_UPSTREAM_ACQUIRE_TIMEOUT_SECONDS." + ), +) +@click.option( + "--anthropic-pre-upstream-memory-context-timeout-seconds", + type=float, + default=None, + envvar="HEADROOM_ANTHROPIC_PRE_UPSTREAM_MEMORY_CONTEXT_TIMEOUT_SECONDS", + help=( + "Fail-open timeout for Anthropic memory-context lookup while the request " + "still holds a pre-upstream slot. " + "Default: 2.0 seconds. " + "Env: HEADROOM_ANTHROPIC_PRE_UPSTREAM_MEMORY_CONTEXT_TIMEOUT_SECONDS." + ), +) +@click.option("--log-file", default=None, help="Path to JSONL log file") +@click.option( + "--log-messages", + is_flag=True, + help="Enable full message logging (request/response content stored for live feed)", +) +@click.option( + "--budget", + type=float, + default=None, + envvar="HEADROOM_BUDGET", + help="Daily budget limit in USD (env: HEADROOM_BUDGET)", +) +# Code graph: indexes project + watches files for live reindex via codebase-memory-mcp +@click.option( + "--code-graph", + is_flag=True, + help="Enable code graph intelligence (indexes project, watches files for live reindex via codebase-memory-mcp)", +) +# Read lifecycle (ON by default: compresses stale/superseded Read outputs) +@click.option( + "--no-read-lifecycle", + is_flag=True, + help="Disable Read lifecycle management (stale/superseded Read compression)", +) +# Intelligent Context Management (ON by default) +@click.option( + "--no-intelligent-context", + is_flag=True, + help="Disable IntelligentContextManager (fall back to RollingWindow)", +) +@click.option( + "--no-intelligent-scoring", + is_flag=True, + help="Disable multi-factor importance scoring (use position-based)", +) +@click.option( + "--no-compress-first", + is_flag=True, + help="Disable trying deeper compression before dropping messages", +) +# Memory System (Multi-Provider Support) +@click.option( + "--memory", + is_flag=True, + help="Enable persistent user memory. Auto-detects provider and uses appropriate tool format. " + "Set x-headroom-user-id header for per-user memory (defaults to 'default').", +) +@click.option( + "--memory-db-path", + default="", + help="Path to memory database file (default: {cwd}/.headroom/memory.db)", +) +@click.option("--no-memory-tools", is_flag=True, help="Disable automatic memory tool injection") +@click.option( + "--no-memory-context", is_flag=True, help="Disable automatic memory context injection" +) +@click.option( + "--memory-top-k", + type=int, + default=10, + help="Number of memories to inject as context (default: 10)", +) +# Traffic Learning (live pattern extraction from proxy traffic) +@click.option( + "--learn", + is_flag=True, + help="Enable live traffic learning: extract error→recovery patterns, environment facts, " + "and user preferences from proxy traffic. Implies --memory. " + "Learned patterns are saved to agent-native memory files (MEMORY.md, .cursor/rules, AGENTS.md).", +) +@click.option( + "--no-learn", + is_flag=True, + help="Explicitly disable traffic learning even when --memory is set.", +) +# Backend configuration +@click.option( + "--backend", + default="anthropic", + help=( + "API backend: 'anthropic' (direct), 'bedrock' (AWS), 'openrouter' (OpenRouter), " + "'anyllm' (any-llm), or 'litellm-' (e.g., litellm-vertex)" + ), +) +@click.option( + "--anyllm-provider", + default="openai", + help="Provider for any-llm backend: openai, mistral, groq, ollama, etc. (default: openai)", +) +@click.option( + "--anthropic-api-url", + default=None, + help="Custom Anthropic API URL for passthrough endpoints (env: ANTHROPIC_TARGET_API_URL)", +) +@click.option( + "--openai-api-url", + default=None, + help="Custom OpenAI API URL for passthrough endpoints (env: OPENAI_TARGET_API_URL)", +) +@click.option( + "--gemini-api-url", + default=None, + help="Custom Gemini API URL for passthrough endpoints (env: GEMINI_TARGET_API_URL)", +) +@click.option( + "--cloudcode-api-url", + default=None, + help="Custom Cloud Code Assist API URL for compatibility endpoints (env: CLOUDCODE_TARGET_API_URL)", +) +@click.option( + "--region", + default="us-west-2", + help="Cloud region for Bedrock/Vertex/etc (default: us-west-2)", +) +@click.option( + "--bedrock-region", + default=None, + help="(deprecated, use --region) AWS region for Bedrock", +) +@click.option( + "--bedrock-profile", + default=None, + help="AWS profile name for Bedrock (default: use default credentials)", +) +@click.option( + "--no-telemetry", + is_flag=True, + help="Disable anonymous usage telemetry (env: HEADROOM_TELEMETRY=off)", +) +@click.option( + "--stateless", + is_flag=True, + help="Disable all filesystem writes — run purely in-memory. " + "For containerized / read-only / load-balanced deployments. " + "(env: HEADROOM_STATELESS=true)", +) +@click.pass_context +def proxy( + ctx: click.Context, + mode: str | None, + host: str, + port: int, + intercept_tool_results: bool, + no_optimize: bool, + no_cache: bool, + no_rate_limit: bool, + retry_max_attempts: int | None, + connect_timeout_seconds: int | None, + anthropic_pre_upstream_concurrency: int | None, + anthropic_pre_upstream_acquire_timeout_seconds: float | None, + anthropic_pre_upstream_memory_context_timeout_seconds: float | None, + log_file: str | None, + log_messages: bool, + budget: float | None, + code_graph: bool, + no_read_lifecycle: bool, + no_intelligent_context: bool, + no_intelligent_scoring: bool, + no_compress_first: bool, + memory: bool, + memory_db_path: str, + no_memory_tools: bool, + no_memory_context: bool, + memory_top_k: int, + learn: bool, + no_learn: bool, + backend: str, + anyllm_provider: str, + anthropic_api_url: str | None, + openai_api_url: str | None, + gemini_api_url: str | None, + cloudcode_api_url: str | None, + region: str, + bedrock_region: str | None, + bedrock_profile: str | None, + no_telemetry: bool, + stateless: bool, +) -> None: + """Start the optimization proxy server. + + \b + Examples: + headroom proxy Start proxy on port 8787 + headroom proxy --port 8080 Start proxy on port 8080 + headroom proxy --no-optimize Passthrough mode (no optimization) + + \b + Usage with Claude Code: + ANTHROPIC_BASE_URL=http://localhost:8787 claude + + \b + Usage with OpenAI-compatible clients: + OPENAI_BASE_URL=http://localhost:8787/v1 your-app + """ + # Import here to avoid slow startup + try: + from headroom.proxy.server import ProxyConfig, run_server + except ImportError as e: + click.echo("Error: Proxy dependencies not installed. Run: pip install headroom[proxy]") + click.echo(f"Details: {e}") + raise SystemExit(1) from None + + # Opt-in: turn on tool_result interceptors (ast-grep Read outline, etc.). + # Only fetch the bundled CLI tool binaries when the feature is enabled — + # otherwise we'd pay a network round-trip and risk a readonly-FS failure + # for capabilities the user hasn't asked for. The TransformPipeline reads + # this env var at construction time. + if intercept_tool_results: + from headroom.binaries import ensure_tools + + resolved_tools = ensure_tools() + critical_tools = ["ast-grep"] + missing = [t for t in critical_tools if not resolved_tools.get(t)] + if missing: + # User explicitly opted in — fail fast rather than silently starting + # with non-functional interceptors. They can retry with the tool + # installed, or drop the flag if they want pass-through behavior. + click.secho( + f"error: --intercept-tool-results requires tool(s) that could not " + f"be installed: {missing}. Run `headroom tools doctor` to diagnose, " + "or omit the flag to start the proxy without interceptors.", + fg="red", + err=True, + ) + sys.exit(1) + os.environ["HEADROOM_INTERCEPT_ENABLED"] = "1" + + provider_api_overrides = resolve_api_overrides( + anthropic_api_url=anthropic_api_url, + openai_api_url=openai_api_url, + gemini_api_url=gemini_api_url, + cloudcode_api_url=cloudcode_api_url, + environ=os.environ, + ) + + # Resolve anyllm provider: env var takes precedence over CLI default (matches argparse path) + effective_anyllm_provider = os.environ.get("HEADROOM_ANYLLM_PROVIDER") or anyllm_provider + + # Resolve mode: CLI flag > env var > default + effective_mode: str = normalize_proxy_mode( + mode or os.environ.get("HEADROOM_MODE") or PROXY_MODE_TOKEN + ) + + # Stateless mode: CLI flag or env var + is_stateless = stateless or os.environ.get("HEADROOM_STATELESS", "").lower() in ( + "true", + "1", + "yes", + "on", + ) + + # Telemetry opt-out: --no-telemetry flag sets the env var + if no_telemetry: + os.environ["HEADROOM_TELEMETRY"] = "off" + + # Stateless mode: suppress TOIN filesystem persistence + if is_stateless: + os.environ["HEADROOM_TOIN_BACKEND"] = "none" + + # License key for managed/enterprise deployments (optional) + license_key = os.environ.get("HEADROOM_LICENSE_KEY") + + config = ProxyConfig( + host=host, + port=port, + anthropic_api_url=provider_api_overrides.anthropic, + openai_api_url=provider_api_overrides.openai, + gemini_api_url=provider_api_overrides.gemini, + cloudcode_api_url=provider_api_overrides.cloudcode, + mode=effective_mode, + optimize=not no_optimize, + cache_enabled=not no_cache, + rate_limit_enabled=not no_rate_limit, + retry_max_attempts=retry_max_attempts if retry_max_attempts is not None else 3, + connect_timeout_seconds=connect_timeout_seconds + if connect_timeout_seconds is not None + else 10, + log_file=None if is_stateless else log_file, + log_full_messages=log_messages + or os.environ.get("HEADROOM_LOG_MESSAGES", "").lower() in ("true", "1", "yes", "on"), + budget_limit_usd=budget, + # Code graph: live file watcher for incremental reindexing + code_graph_watcher=code_graph, + # Read lifecycle: ON by default (use --no-read-lifecycle to disable) + read_lifecycle=not no_read_lifecycle, + # Intelligent Context: ON by default (use --no-intelligent-context to disable) + intelligent_context=not no_intelligent_context, + intelligent_context_scoring=not no_intelligent_scoring, + intelligent_context_compress_first=not no_compress_first, + # Memory System (Multi-Provider with auto-detection) + # --learn implies --memory (need backend for storing patterns) + # Stateless mode disables memory (requires SQLite on disk) + memory_enabled=False if is_stateless else (memory or (learn and not no_learn)), + memory_db_path=memory_db_path, + memory_inject_tools=not no_memory_tools, + memory_inject_context=not no_memory_context, + memory_top_k=memory_top_k, + # Traffic Learning: only with --learn, never with --no-learn + # Stateless mode disables learning (requires filesystem) + traffic_learning_enabled=False if is_stateless else (learn and not no_learn), + traffic_learning_agent_type=os.environ.get("HEADROOM_AGENT_TYPE", "unknown"), + # Backend (Anthropic direct, Bedrock, LiteLLM, or any-llm) + backend=backend, + bedrock_region=bedrock_region or region, + bedrock_profile=bedrock_profile, + anyllm_provider=effective_anyllm_provider, + # License / Usage Reporting (managed/enterprise) + license_key=license_key, + # Stateless mode: disable all filesystem writes + stateless=is_stateless, + # Unit 4: bounded pre-upstream concurrency on the Anthropic HTTP + # path. ``None`` -> HeadroomProxy computes ``max(2, min(8, + # os.cpu_count() or 4))``; ``<= 0`` -> disabled (unbounded). + # Precedence: CLI > env > auto-compute (click's ``envvar`` + # handles the env-var fallback). + anthropic_pre_upstream_concurrency=anthropic_pre_upstream_concurrency, + anthropic_pre_upstream_acquire_timeout_seconds=( + anthropic_pre_upstream_acquire_timeout_seconds + if anthropic_pre_upstream_acquire_timeout_seconds is not None + else 15.0 + ), + anthropic_pre_upstream_memory_context_timeout_seconds=( + anthropic_pre_upstream_memory_context_timeout_seconds + if anthropic_pre_upstream_memory_context_timeout_seconds is not None + else 2.0 + ), + ) + + memory_status = "DISABLED" + if config.memory_enabled: + memory_status = "ENABLED (multi-provider)" + + license_status = "OSS (no license key)" + if license_key: + license_status = f"MANAGED (key={license_key[:8]}...)" + + provider_api_targets = resolve_api_targets(config.provider_api_overrides) + anthropic_url = provider_api_targets.anthropic + openai_url = provider_api_targets.openai + cloudcode_url = provider_api_targets.cloudcode + backend_section = "" + + if config.backend == "anyllm" or config.backend.startswith("anyllm-"): + # any-llm backend + backend_section = """ + Set credentials for your provider (e.g., OPENAI_API_KEY, MISTRAL_API_KEY) + Providers: https://mozilla-ai.github.io/any-llm/providers/ +""" + elif config.backend != "anthropic": + # LiteLLM backend + from headroom.backends.litellm import get_provider_config + + provider = config.backend.replace("litellm-", "") + provider_config = get_provider_config(provider) + + # Build usage instructions from provider config + env_vars_str = ( + ", ".join(provider_config.env_vars) if provider_config.env_vars else "See docs" + ) + backend_section = f""" +IMPORTANT for {provider_config.display_name} users: + 1. Set credentials: {env_vars_str} + 2. Set a dummy Anthropic key: ANTHROPIC_API_KEY="sk-ant-dummy" + (Headroom ignores this - it uses your {provider_config.display_name} credentials) + 3. Set base URL: ANTHROPIC_BASE_URL=http://{config.host}:{config.port}""" + if provider_config.model_format_hint: + backend_section += f"\n 4. Use model names: {provider_config.model_format_hint}" + backend_section += "\n" + + # Build memory section if enabled + memory_section = "" + if config.memory_enabled: + memory_section = f""" +Memory (Multi-Provider): + - Auto-detects provider from request (Anthropic, OpenAI, Gemini, etc.) + - Anthropic: Uses native memory tool (memory_20250818) - subscription safe + - OpenAI/Gemini/Others: Uses function calling format + - All providers share the same semantic vector store backend + - Set x-headroom-user-id header for per-user memory (defaults to 'default') + - Tools: {"ENABLED" if config.memory_inject_tools else "DISABLED"} + - Context injection: {"ENABLED" if config.memory_inject_context else "DISABLED"} + - Database: {config.memory_db_path} +""" + + # Stateless mode warning + stateless_line = "" + if is_stateless: + stateless_line = ( + " Stateless: YES (no filesystem writes — memory, logs, TOIN disabled)\n" + ) + + from headroom.telemetry.beacon import is_telemetry_enabled + + # Build telemetry section for the startup banner + if is_telemetry_enabled(): + telemetry_line = ( + " Telemetry: ENABLED (anonymous aggregate stats)\n" + " Disable: HEADROOM_TELEMETRY=off or headroom proxy --no-telemetry" + ) + else: + telemetry_line = " Telemetry: DISABLED" + + click.echo(f""" +╔═══════════════════════════════════════════════════════════════════════╗ +║ HEADROOM PROXY ║ +║ The Context Optimization Layer for LLM Applications ║ +╚═══════════════════════════════════════════════════════════════════════╝ + +Starting proxy server... + + URL: http://{config.host}:{config.port} + Mode: {config.mode} + Optimization: {"ENABLED" if config.optimize else "DISABLED"} + Caching: {"ENABLED" if config.cache_enabled else "DISABLED"} + Rate Limit: {"ENABLED" if config.rate_limit_enabled else "DISABLED"} + Memory: {memory_status} + License: {license_status} +{stateless_line}{telemetry_line} +{backend_section} +Routing: + /v1/messages → {anthropic_url} + /v1/chat/completions → {openai_url} + /v1/responses → {openai_url} (HTTP + WebSocket) + /v1internal:streamGenerateContent → {cloudcode_url} + +Usage: + Claude Code: ANTHROPIC_BASE_URL=http://{config.host}:{config.port} claude + Codex / OpenAI: OPENAI_BASE_URL=http://{config.host}:{config.port}/v1 your-app +{memory_section} +Endpoints: + GET /livez Process liveness + GET /readyz Traffic readiness + GET /health Aggregate health + GET /stats Detailed statistics + GET /stats-history Durable compression history + display session + GET /metrics Prometheus metrics + +Press Ctrl+C to stop. +""") + + try: + run_server(config) + except KeyboardInterrupt: + click.echo("\nShutting down...") diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index 6fbec6551..ae1731537 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -1,2211 +1,2069 @@ -"""Wrap CLI commands to run through Headroom proxy. - -Usage: - headroom wrap claude # Start proxy + rtk + 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 --port 9999 # Custom proxy port - headroom wrap claude -- --model opus # Pass args to claude -""" - -from __future__ import annotations - -import io -import json -import os -import shutil -import signal -import socket -import subprocess -import sys -import time -import urllib.error -import urllib.request -from pathlib import Path -from typing import Any - -# Fix Windows cp1252 encoding — box-drawing characters require UTF-8 -if sys.platform == "win32" and hasattr(sys.stdout, "buffer"): - if sys.stdout.encoding and sys.stdout.encoding.lower().replace("-", "") != "utf8": - sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace") - sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding="utf-8", errors="replace") - -import click - -from headroom.copilot_auth import DEFAULT_API_URL as COPILOT_API_URL -from headroom.copilot_auth import has_oauth_auth, resolve_client_bearer_token - -from .main import main - - -def _print_telemetry_notice() -> None: - """Print a telemetry notice when anonymous telemetry is enabled. - - Respects the HEADROOM_TELEMETRY and HEADROOM_TELEMETRY_WARN feature flags. - Does nothing when telemetry or warnings are disabled. - """ - from headroom.telemetry.beacon import format_telemetry_notice - - notice = format_telemetry_notice(prefix=" ") - if notice: - click.echo(notice) - - -# Proxy health check (reused from evals/suite_runner.py pattern) - - -def _check_proxy(port: int) -> bool: - """Check if Headroom proxy is running on given port.""" - try: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.settimeout(1) - s.connect(("127.0.0.1", port)) - return True - except (TimeoutError, ConnectionRefusedError, OSError): - return False - - -def _get_log_path() -> Path: - """Get path for proxy log file.""" - from headroom import paths as _paths - - log_dir = _paths.log_dir() - log_dir.mkdir(parents=True, exist_ok=True) - return log_dir / "proxy.log" - - -def _start_proxy( - port: int, - *, - learn: bool = False, - memory: bool = False, - agent_type: str = "unknown", - code_graph: bool = False, - backend: str | None = None, - anyllm_provider: str | None = None, - region: str | None = None, - openai_api_url: str | None = None, -) -> subprocess.Popen: - """Start Headroom proxy as a background subprocess. - - Logs are written to ~/.headroom/logs/proxy.log to avoid pipe buffer - deadlocks (macOS pipe buffer is ~64KB — a busy proxy fills it quickly, - blocking the process). - """ - cmd = [sys.executable, "-m", "headroom.cli", "proxy", "--port", str(port)] - - # Forward HEADROOM_MODE env var so the proxy respects the user's mode choice - headroom_mode = os.environ.get("HEADROOM_MODE") - if headroom_mode: - cmd.extend(["--mode", headroom_mode]) - - # Forward --learn flag to proxy subprocess - if learn: - cmd.append("--learn") - - # Forward --memory flag to proxy subprocess - if memory: - cmd.append("--memory") - - # Forward --code-graph flag to proxy subprocess (live file watcher) - if code_graph: - cmd.append("--code-graph") - - # Forward backend configuration to proxy subprocess - _backend = backend or os.environ.get("HEADROOM_BACKEND") - if _backend: - cmd.extend(["--backend", _backend]) - - _anyllm = anyllm_provider or os.environ.get("HEADROOM_ANYLLM_PROVIDER") - if _anyllm: - cmd.extend(["--anyllm-provider", _anyllm]) - - _region = region or os.environ.get("HEADROOM_REGION") - if _region: - cmd.extend(["--region", _region]) - - if openai_api_url: - cmd.extend(["--openai-api-url", openai_api_url]) - - log_path = _get_log_path() - log_file = open(log_path, "a") # noqa: SIM115 - - # Ensure proxy subprocess uses UTF-8 (Windows defaults to cp1252) - proxy_env = os.environ.copy() - proxy_env["PYTHONIOENCODING"] = "utf-8" - - # Tell the proxy which agent is being wrapped (for traffic learning output) - if agent_type != "unknown": - proxy_env["HEADROOM_AGENT_TYPE"] = agent_type - proxy_env.setdefault("HEADROOM_STACK", f"wrap_{agent_type}") - - proc = subprocess.Popen( - cmd, - stdout=log_file, - stderr=log_file, - env=proxy_env, - ) - - # Wait for proxy to be ready (up to 45 seconds). - # ML components (Kompress, Magika, Tree-sitter) load synchronously before - # uvicorn binds the port. On slower machines this can take 20-30 seconds. - for _i in range(45): - time.sleep(1) - if _check_proxy(port): - click.echo(f" Logs: {log_path}") - return proc - # Check if process died - if proc.poll() is not None: - log_file.close() - # Read last few lines of log for error context - try: - tail = log_path.read_text()[-500:] - except Exception: - tail = "(no log output)" - raise RuntimeError(f"Proxy exited with code {proc.returncode}: {tail}") - - proc.kill() - log_file.close() - raise RuntimeError(f"Proxy failed to start on port {port} within 45 seconds") - - -def _setup_rtk(verbose: bool = False) -> Path | None: - """Ensure rtk is installed and hooks are registered.""" - from headroom.rtk import get_rtk_path - from headroom.rtk.installer import ensure_rtk, register_claude_hooks - - rtk_path = get_rtk_path() - - if rtk_path: - if verbose: - click.echo(f" rtk found at {rtk_path}") - else: - click.echo(" Downloading rtk (Rust Token Killer)...") - rtk_path = ensure_rtk() - if rtk_path: - click.echo(f" rtk installed at {rtk_path}") - else: - click.echo(" rtk download failed — continuing without it") - return None - - # Register hooks (idempotent) - if register_claude_hooks(rtk_path): - if verbose: - click.echo(" rtk hooks registered in Claude Code") - else: - click.echo(" rtk hook registration failed — continuing without it") - - return rtk_path - - -_CBM_MCP_SERVER_NAME = "codebase-memory-mcp" - - -def _register_cbm_mcp_server(cbm_bin: str) -> None: - """Register codebase-memory-mcp as an MCP server in Claude Code. - - Uses ``claude mcp add`` so the tools appear in ``/mcp`` automatically. - Idempotent — skips if already registered. - """ - claude_cli = shutil.which("claude") - if not claude_cli: - return - - # Check if already registered - check = subprocess.run( - [claude_cli, "mcp", "get", _CBM_MCP_SERVER_NAME], - capture_output=True, - text=True, - ) - if check.returncode == 0: - return # Already registered - - result = subprocess.run( - [claude_cli, "mcp", "add", _CBM_MCP_SERVER_NAME, "-s", "user", "--", cbm_bin], - capture_output=True, - text=True, - ) - if result.returncode == 0: - click.echo(f" Code graph: registered {_CBM_MCP_SERVER_NAME} MCP server") - else: - pass # Non-critical — tools won't appear in /mcp but graph still works - - -def _setup_code_graph(verbose: bool = False) -> bool: - """Ensure codebase-memory-mcp is installed, registered as MCP server, and project is indexed. - - codebase-memory-mcp builds a knowledge graph of the codebase using - tree-sitter, enabling the LLM to query code structure (call chains, - function definitions, impact analysis) instead of reading entire files. - - Steps: - 1. Download the binary if not already present. - 2. Register as an MCP server in Claude Code (``claude mcp add``). - 3. Index the current project (fast, idempotent). - - With Claude Code's MCP Tool Search, the 14 graph tools add ~200 tokens - overhead per request (not the full ~1,915) — they're lazy-loaded. - - Returns True if graph is ready, False if setup failed. - """ - from headroom.graph.installer import ensure_cbm, get_cbm_path - - cbm_path = get_cbm_path() - if not cbm_path: - click.echo(" Code graph: downloading codebase-memory-mcp...") - cbm_path = ensure_cbm() - if cbm_path: - click.echo(f" Code graph: installed at {cbm_path}") - else: - click.echo(" Code graph: download failed — skipping") - return False - - cbm_bin = str(cbm_path) - - # Register as MCP server so tools appear in /mcp - _register_cbm_mcp_server(cbm_bin) - - # Index current project (fast — ~1s for most repos, idempotent) - project_dir = str(Path.cwd()) - try: - result = subprocess.run( - [ - cbm_bin, - "cli", - "index_repository", - json.dumps({"repo_path": project_dir, "mode": "fast"}), - ], - capture_output=True, - text=True, - timeout=30, - ) - if result.returncode == 0: - # Parse node/edge counts from output - for line in result.stdout.splitlines(): - if '"nodes"' in line: - try: - # Parse the JSON response to extract node/edge counts - import re - - m_nodes = re.search(r'"nodes":(\d+)', line) - m_edges = re.search(r'"edges":(\d+)', line) - if m_nodes and m_edges: - nodes = int(m_nodes.group(1)) - edges = int(m_edges.group(1)) - click.echo( - f" Code graph: indexed ({nodes:,} symbols, " - f"{edges:,} relationships)" - ) - else: - click.echo(" Code graph: indexed") - except (ValueError, AttributeError): - click.echo(" Code graph: indexed") - return True - click.echo(" Code graph: indexed") - return True - else: - if verbose: - click.echo(f" Code graph: indexing failed ({result.stderr[:100]})") - return False - except subprocess.TimeoutExpired: - click.echo(" Code graph: indexing timed out (will complete in background)") - return False - except Exception as e: - if verbose: - click.echo(f" Code graph: setup failed ({e})") - return False - - -# 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 = """\ - -# 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 rtk read rtk grep -rtk find rtk diff - -# Test (90-99% savings) — shows failures only -rtk pytest tests/ rtk cargo test rtk test - -# 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 rtk log rtk json -rtk summary rtk deps rtk env - -# GitHub (26-87% savings) -rtk gh pr view rtk gh run list rtk gh issue list - -# Infrastructure (85% savings) -rtk docker ps rtk kubectl get rtk docker logs - -# Package managers (70-90% savings) -rtk pip list rtk pnpm install rtk npm run