<h1 align="center">Context Budget Engine</h1>

<p align="center">The context layer that teaches AI coding agents to read less, and keeps improving itself.</p>

<p align="center">
  <a href="https://github.com/sivram04/context-budget-engine/actions/workflows/ci.yml"><img src="https://github.com/sivram04/context-budget-engine/actions/workflows/ci.yml/badge.svg" alt="CI"></a>
  <img src="https://img.shields.io/badge/python-3.11%2B-3776ab" alt="Python 3.11+">
  <img src="https://img.shields.io/badge/MCP-server-1a7f64" alt="MCP server">
  <a href="LICENSE"><img src="https://img.shields.io/badge/license-MIT-blue" alt="MIT license"></a>
</p>

<p align="center">
<strong>Most of an agent's tokens are not spent thinking.</strong><br>
They are spent re-reading whole files, long logs and search dumps, on every&nbsp;step.
</p>

<p align="center">
<strong>40&nbsp;tasks. 3&nbsp;runs each. −40%&nbsp;tokens. No&nbsp;drop in&nbsp;success.</strong><br>
Then I plugged it into Claude Code: same 10/10 answers with 22%&nbsp;fewer input&nbsp;tokens.
</p>

<p align="center">
Ask a question about a file. It returns the functions that matter, not the whole&nbsp;file.<br>
A gated loop rewrites its own reading policy, and <strong>only keeps changes that are proven&nbsp;cheaper without losing&nbsp;tasks.</strong>
</p>

<p align="center">
  <picture>
    <source media="(prefers-color-scheme: dark)" srcset="docs/results-dark.svg">
    <img src="docs/results-light.svg" alt="Input tokens per task: benchmark agent 26,394 with whole files vs 15,801 with the engine (−40%); Claude Code 49,838 with built-in Read vs 38,820 with read_smart (−22%); success equal or better in both" width="800">
  </picture>
</p>

<p align="center"><a href="#your-turn">Try it in one minute&nbsp;↓</a> · <a href="#results">See every number, including the ones that went wrong&nbsp;→</a></p>

## Your turn

Connect it to the AI assistant you already use. No API key needed.

```bash
git clone https://github.com/sivram04/context-budget-engine && cd context-budget-engine
python -m venv .venv && source .venv/bin/activate   # Windows: .venv\Scripts\activate
pip install -e .
claude mcp add cbe -- cbe-mcp
```

Then ask Claude Code something about your codebase and tell it to use `read_smart`. Run
`savings` at any time to see how many tokens it did not have to read.

<p align="center"><sub>Open source. Local. Works with Claude Code, Claude Desktop, Cursor and any MCP client.</sub></p>

## The principles

> Read the part, not the file. Measure, don't assume. A change ships only if it is cheaper **and** solves as many tasks. The loop edits policies, never the judge. Report the noise, not just the best run.

## What it does for you

| Feature | Description |
|---|---|
| **`read_smart`** | Reads a file the way the question needs: the matching functions (AST-based for Python), a skeleton of signatures, a grep window, or the whole file when it is small. Never returns more than the file itself |
| **`log_smart`** | The errors, warnings and last lines of a log, not 10,000 lines of INFO |
| **`search_smart`** | Regex search that returns short `file:line: text` hits, skips caches, binaries and huge files, and caps the result count |
| **`savings`** | Running total of tokens returned vs. what full reads would have cost, and the % saved |
| **Policy file** | One small YAML file (`policies/v1.yaml`) decides which strategy to use by file kind, size and whether there is a question. Strictly validated |
| **Self-improvement loop** | `cbe improve` runs the agent, finds token hot spots and failures, proposes a new policy (LLM optimizer or model-free local search), and passes it through a train + held-out gate |
| **Noise-aware gate** | Every candidate runs several times; it is accepted only if it beats the mean by a minimum effect size on training tasks and does not lose held-out tasks. Duplicate proposals are skipped for free |
| **Benchmark + judge** | 40 tasks on a real codebase: 20 exact-answer questions and 20 planted bugs judged by the repo's own test suite. Verified by `scripts/verify_tasks.py` |
| **Works without a key** | The optional `summarize` strategy uses a cheap model; with no key or a failed call it falls back to a model-free view |
| **Hard budget** | Paid runs stop at `CBE_MAX_USD`, save progress, and resume where they left off |

## What it will not do

- **It never edits your code.** It only decides what the assistant reads.
- **The loop cannot grade itself.** It may only write policy files: never the tasks, the judge, the gate or the agent.
- **It never claims a win the gate did not measure.** Rejected rounds and noisy runs are published below, next to the good ones.
- **It never sends your code anywhere.** The MCP tools run locally; only the optional `summarize` strategy calls a model, and only if you configure one.

## How it works

```
            ┌──────── cbe improve (self-improvement loop) ────────┐
            │  run → measure → diagnose → propose policy → gate   │
            └───────────────┬─────────────────────────────────────┘
                            │ edits only the policy (policies/vN.yaml)
 test agent / Claude Code   ▼
        └──► read_smart(path, query) ──► policy picks a strategy ──► reduced view
                                          full · skeleton · extract · grep_window
                                          log_tail · head · summarize (cheap model)
 benchmark + judge (read-only) ──► results/*.jsonl ──► cbe report
```

**The gate** (`src/cbe/improve.py`): a candidate must be cheaper on the 27 training tasks by a
minimum margin with no meaningful loss of success, *and* must not lose success on the 13
held-out tasks. **Proposals** come from an LLM that reads the failures and token hot spots,
and from a model-free local search so the loop still works with weak models.

## Results

All raw per-task results are in `results/`. Agent model: Gemini 3.1 Flash-Lite.

### 1. Baseline vs engine, 3 runs each, all 40 tasks

| Policy | Success, mean ± sd (runs) | Tokens / task, mean ± sd |
|---|---|---|
| v0: read whole files | 95.0% ± 0.0 (95, 95, 95) | 26,394 ± 1,673 |
| **v1: Context Budget Engine** | **96.7% ± 3.8 (100, 92, 98)** | **15,801 ± 1,037 (−40%)** |

Paired by task, v1 is cheaper on 27/40 tasks (median ratio 0.75). The −40% is far outside the
~6% run-to-run spread; the success difference is within noise. On bug fixes alone, the first
run used 49,476 → 24,490 tokens per task (−50%).

### 2. Claude Code, with vs without the MCP tools

Same model (Sonnet), same repo copy, isolated config; only the reading tools differ
(`scripts/claude_code_check.py`).

| | Correct | Input tokens / task | Output tokens / task | Turns / task |
|---|---|---|---|---|
| Built-in Read | 10/10 | 49,838 | 538 | 4.3 |
| **`read_smart`** | **10/10** | **38,820 (−22%)** | **376 (−30%)** | **3.2** |

Fewer input tokens on 8 of 10 questions. Caveats: one run per arm; the files are short; and the
list-price cost came out about equal ($0.031 vs $0.033 per question) because the MCP tool calls
caused extra cache writes, which are priced higher than cache reads.

### 3. The self-improvement loop, honestly

**First attempt (1 run per evaluation): 6 rounds, all rejected.** Three times the optimizer
proposed a policy identical to v1, and those re-runs of the *same* policy scored 88.9%, 96.3% and
92.6%. One run cannot tell a real gain from noise, so the gate was rebuilt.

**Noise-aware loop (3 runs per evaluation, ≥5% cheaper, ≤2 pts success drop):**

| Round | Outcome |
|---|---|
| 1 | **Accepted as v2**: training 19,201 → 18,195 tokens/task (−5.2%) at 97.5%; held-out 100% → 100% |
| 2, 3, 5 | Rejected: training success fell to 90.1%, 91.4%, 91.4% |
| 4 | Rejected: success fell to 95.1% |
| 6 | Rejected: only −2.3% tokens, below the bar |

One small verified improvement, and five candidates refused, four of which would have traded
success for tokens. All of them came from the model-free search: the Flash-Lite optimizer only
re-proposed existing policies. Whole study (noise check + loop): **$3.49**.

## Usage

### MCP tools

| Tool | What it does |
|---|---|
| `read_smart(path, query, full=false)` | Read the relevant part of a file for a question |
| `log_smart(path, lines)` | Errors, warnings and the tail of a log file |
| `search_smart(root, pattern)` | Compact repository search |
| `savings()` | Tokens saved so far |

Claude Desktop / Cursor: add to `claude_desktop_config.json` / `.cursor/mcp.json`:

```json
{ "mcpServers": { "cbe": { "command": "/full/path/to/.venv/bin/cbe-mcp" } } }
```

### Commands

| Command | What it does |
|---|---|
| `cbe benchmark --policy policies/v1.yaml` | Run a policy on the 40 tasks (resumable) |
| `cbe improve --repeats 3 --min-gain 0.05` | Run the gated self-improvement loop |
| `cbe report --held-only` | Summarise result files |
| `python scripts/noise_report.py <dir>` | Mean ± sd across repeated runs, paired by task |
| `python scripts/claude_code_check.py --n 10` | The Claude Code with/without comparison |

Any OpenAI-compatible endpoint works for the benchmark and the loop (Gemini, Groq, local Ollama).
Copy `.env.example` to `.env`:

```bash
CBE_PROVIDER=gemini GEMINI_API_KEY=... CBE_MAX_USD=3.0
```

## Benchmark

40 tasks on a fixed snapshot of a real codebase (`bench_repos/gateway`, an LLM gateway with 87
tests):

- **20 code-understanding questions** with exact answers, e.g. *"What is the default
  cosine-similarity threshold for a semantic cache hit?"*
- **20 bug fixes**: one bug is planted per task from a plain-language report; it passes only if
  the repo's own tests pass afterwards.

`scripts/verify_tasks.py` proves every planted bug breaks its tests and every expected answer
exists in the source. The 27 train / 13 held-out split is a stable hash of the task id.

## Project layout

```
src/cbe/strategies.py   ways to read less (AST-based for Python, generic for other languages)
src/cbe/policy.py       policy schema + strict validation (what the loop may change)
src/cbe/engine.py       picks a strategy per read, accounts tokens, keyless fallback
src/cbe/agent.py        small JSON-protocol coding agent used for measurement
src/cbe/bench/          tasks, workspaces, judge (read-only), runner
src/cbe/improve.py      diagnose / propose / noise-aware gate loop
src/cbe/mcp_server.py   MCP tools for real assistants
policies/               v0-full (baseline), v1 (hand-written)
results/                raw results, loop logs and accepted policies (v2)
```

## Contributing

```bash
pip install -e ".[dev]"
pytest
ruff check src tests scripts
python scripts/verify_tasks.py
```

Ideas that would help most: strategies for more languages (tree-sitter), a stronger optimizer
model for the loop, and benchmarks on larger codebases.

## License

MIT. See [`LICENSE`](LICENSE).