fix(proxy): make budget enforcement actually work (#885)

## Description

`CostTracker._costs` was initialized but never written to, so
`get_period_cost()` always returned `0` and `check_budget()` always
returned "allowed" — the `--budget` flag was a silent no-op.
`_prune_old_costs()` was dead code with zero callers. This makes budget
enforcement actually work: requests are rejected once the configured
limit is reached.

Closes # <!-- no tracked issue; discovered during a proxy-pipeline audit
-->

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- **`headroom/proxy/cost.py`** — `record_tokens()` now computes the
request cost via `estimate_cost()` and appends it to `_costs`,
activating `_prune_old_costs()`. When a call site has no API usage
breakdown (cache/uncached all zero), `tokens_sent` is used as the input
count so input cost is not silently dropped. `COST_RETENTION_HOURS` 24 →
744 so retention covers the longest budget period (monthly sums from the
1st; 24h retention would have under-enforced monthly budgets).
- **`headroom/proxy/outcome.py`** — the request funnel passes
`output_tokens` through to `record_tokens()` so costs include output,
for all providers.
- **`headroom/cli/proxy.py`** — added `--budget-period
[hourly|daily|monthly]` (env `HEADROOM_BUDGET_PERIOD`); it existed in
`ProxyConfig` and the server entry point but was unreachable from the
main CLI. Fixed the `--budget` help text that wrongly said "resets at
midnight UTC".
- **`headroom/cli/main.py`** — minor registration/version plumbing.
- Tests: regression coverage for the full `record_tokens →
get_period_cost → check_budget` chain, the `tokens_sent` fallback, and
the `--budget-period` flag/env wiring.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
$ pytest tests/test_cost_tracker_counterfactual.py tests/test_request_outcome.py -q
40 passed

$ ruff check headroom/proxy/cost.py headroom/proxy/outcome.py headroom/cli/proxy.py
All checks passed!

$ mypy headroom/proxy/cost.py headroom/proxy/outcome.py headroom/cli/proxy.py --ignore-missing-imports
Success: no issues found
```

## Real Behavior Proof

- Environment: local macOS, Python 3.11, branch `fix/budget-enforcement`
at the PR head commit.
- Exact command / steps: `pytest
tests/test_cost_tracker_counterfactual.py::test_budget_enforced_after_recording_costs
-v` — sets `CostTracker(budget_limit_usd=0.0001)`, records ~$1.50 of
Sonnet input, then asserts `check_budget()` returns not-allowed with
`remaining == 0`.
- Observed result: budget is now enforced — `get_period_cost()` reflects
real spend and `check_budget()` rejects once the limit is exceeded (the
proxy returns HTTP 429 on that path). On `main` the same test fails
because `_costs` is never populated and `check_budget()` always returns
allowed.
- Not tested: live end-to-end rejection against a running proxy with
real upstream traffic; the running proxy needs a restart on this version
to pick up the fix.

```text
$ pytest tests/test_cost_tracker_counterfactual.py::test_budget_enforced_after_recording_costs \
         tests/test_cost_tracker_counterfactual.py::test_budget_input_cost_counted_without_usage_breakdown -v
2 passed
```

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [ ] I have updated the CHANGELOG.md if applicable

## Screenshots (if applicable)

N/A — CLI/backend change with no UI surface. See **Test Output** and
**Real Behavior Proof** above for terminal evidence.

## Additional Notes

- The `ci.yml` coverage-upload change originally added here (commit
`120696e5`) was superseded by an equivalent block the maintainer added
to `main`; the merge from main resolved to main's version. Codecov now
reports all modified lines covered.
- N/A checklist items: no docs or CHANGELOG entry — this is an internal
correctness fix to an existing flag.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

.github/workflows/ci.yml

@@ -191,6 +191,11 @@ jobs:
           cp "${SITE}/headroom/"_core*.so headroom/
           python -c "from headroom._core import DiffCompressor; print('headroom._core OK')"
 
+      # Coverage upload: without this, codecov only receives reports from
+      # the two native-e2e workflows (3 CLI test files total), so head
+      # coverage reads ~6% and codecov/patch fails for ANY diff not
+      # exercised by those files — a false negative on every PR. The main
+      # suite runs here; its coverage must be what codecov sees.
       - name: Run test shard ${{ matrix.shard }}/4
         run: |
           pytest tests scripts/tests \

headroom/cli/proxy.py

@@ -329,8 +329,19 @@ def _selected_context_tool() -> str:
     default=None,
     envvar="HEADROOM_BUDGET",
     help=(
-        "Daily budget limit in USD. Requests are rejected with 429 once the limit is reached. "
-        "Resets at midnight UTC. Env: HEADROOM_BUDGET."
+        "Budget limit in USD per --budget-period. Requests are rejected with 429 "
+        "once the limit is reached. Env: HEADROOM_BUDGET."
+    ),
+)
+@click.option(
+    "--budget-period",
+    type=click.Choice(["hourly", "daily", "monthly"]),
+    default="daily",
+    envvar="HEADROOM_BUDGET_PERIOD",
+    help=(
+        "Period the --budget limit applies to. Hourly resets on a rolling hour, "
+        "daily at local midnight, monthly on the 1st. Default: daily. "
+        "Env: HEADROOM_BUDGET_PERIOD."
     ),
 )
 # Code-aware compression (AST-based, requires `pip install headroom-ai[code]`).
@@ -621,6 +632,7 @@ def proxy(
     codex_wire_debug: bool,
     codex_wire_debug_dir: str | None,
     budget: float | None,
+    budget_period: str,
     code_aware_flag: bool | None,
     disable_kompress: bool,
     code_graph: bool,
@@ -827,6 +839,7 @@ def proxy(
         log_full_messages=log_messages
         or os.environ.get("HEADROOM_LOG_MESSAGES", "").lower() in ("true", "1", "yes", "on"),
         budget_limit_usd=budget,
+        budget_period=cast(Literal["hourly", "daily", "monthly"], budget_period),
         # Code-aware compression resolution:
         # 1. Explicit --code-aware / --no-code-aware always wins.
         # 2. Otherwise read HEADROOM_CODE_AWARE_ENABLED (truthy = on).

headroom/proxy/cost.py

@@ -556,7 +556,10 @@ class CostTracker:
     """
 
     MAX_COST_ENTRIES = 100_000
-    COST_RETENTION_HOURS = 24
+    # Used by _prune_old_costs(), called from record_tokens() on every request.
+    # Must be >= the longest budget_period (monthly = up to 31 days), otherwise
+    # get_period_cost() undercounts and check_budget() silently under-enforces.
+    COST_RETENTION_HOURS = 744  # 31 days
 
     def __init__(self, budget_limit_usd: float | None = None, budget_period: str = "daily"):
         self.budget_limit_usd = budget_limit_usd
@@ -667,8 +670,9 @@ class CostTracker:
         cache_write_5m_tokens: int = 0,
         cache_write_1h_tokens: int = 0,
         uncached_tokens: int = 0,
+        output_tokens: int = 0,
     ):
-        """Record token counts per model.
+        """Record token counts per model and accumulate request cost for budget enforcement.
 
         Args:
             model: Model name.
@@ -677,6 +681,7 @@ class CostTracker:
             cache_read_tokens: Cache read tokens from API response usage.
             cache_write_tokens: Cache write tokens from API response usage.
             uncached_tokens: Non-cached input tokens from API response usage.
+            output_tokens: Output tokens from API response usage.
         """
         self._tokens_saved_by_model[model] = (
             self._tokens_saved_by_model.get(model, 0) + tokens_saved
@@ -699,6 +704,24 @@ class CostTracker:
             self._api_uncached_by_model.get(model, 0) + uncached_tokens
         )
 
+        # Populate _costs so check_budget() has real data to enforce against.
+        # When the call site had no API usage breakdown (all cache/uncached
+        # fields are 0), fall back to tokens_sent so input cost isn't
+        # silently dropped from the budget.
+        input_tokens = uncached_tokens
+        if not (uncached_tokens or cache_read_tokens or cache_write_tokens):
+            input_tokens = tokens_sent
+        cost = self.estimate_cost(
+            model=model,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            cache_read_tokens=cache_read_tokens,
+            cache_write_tokens=cache_write_tokens,
+        )
+        if cost is not None:
+            self._costs.append((datetime.now(), cost))
+            self._prune_old_costs()
+
     def get_period_cost(self) -> float:
         """Get cost for current budget period."""
         now = datetime.now()

headroom/proxy/outcome.py

@@ -371,6 +371,7 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None:
             cache_write_5m_tokens=outcome.cache_write_5m_tokens,
             cache_write_1h_tokens=outcome.cache_write_1h_tokens,
             uncached_tokens=outcome.uncached_input_tokens,
+            output_tokens=outcome.output_tokens,
         )
 
     # 3. Per-request log (optional). The ``client`` outcome field is

tests/test_cli_proxy_env.py

@@ -222,6 +222,34 @@ class TestCLIProxyEnvVars:
         assert result.exit_code == 0, result.output
         assert captured_config["config"].budget_limit_usd == 100.5
 
+    def test_budget_period_flag_and_env(self, runner):
+        """--budget-period and HEADROOM_BUDGET_PERIOD should reach ProxyConfig."""
+        captured_config = {}
+
+        def mock_run_server(config, **kwargs):
+            captured_config["config"] = config
+
+        with patch("headroom.proxy.server.run_server", mock_run_server):
+            result = runner.invoke(
+                main,
+                ["proxy", "--budget", "50", "--budget-period", "monthly"],
+                catch_exceptions=False,
+            )
+
+        assert result.exit_code == 0, result.output
+        assert captured_config["config"].budget_period == "monthly"
+
+        with patch("headroom.proxy.server.run_server", mock_run_server):
+            result = runner.invoke(
+                main,
+                ["proxy"],
+                env={"HEADROOM_BUDGET_PERIOD": "hourly"},
+                catch_exceptions=False,
+            )
+
+        assert result.exit_code == 0, result.output
+        assert captured_config["config"].budget_period == "hourly"
+
     def test_code_aware_enabled_from_env(self, runner):
         """HEADROOM_CODE_AWARE_ENABLED env var should be passed to ProxyConfig."""
         captured_config = {}

tests/test_cost_tracker_counterfactual.py

@@ -108,3 +108,47 @@ def test_no_cost_without_headroom_field():
     stats = ct.stats()
 
     assert "cost_without_headroom_usd" not in stats
+
+
+def test_budget_enforced_after_recording_costs():
+    """record_tokens must populate cost history so check_budget enforces the limit.
+
+    Regression test: _costs was never written, so check_budget always
+    returned (True, budget_limit) and budgets were silently unenforced.
+    """
+    from headroom.proxy.server import CostTracker
+
+    ct = CostTracker(budget_limit_usd=0.0001, budget_period="daily")
+    allowed, remaining = ct.check_budget()
+    assert allowed  # nothing spent yet
+
+    # ~$1.50+ of Sonnet input at list price — far over the budget
+    ct.record_tokens(
+        "claude-sonnet-4-20250514",
+        tokens_saved=0,
+        tokens_sent=500_000,
+        uncached_tokens=500_000,
+        output_tokens=10_000,
+    )
+
+    assert ct.get_period_cost() > 0
+    allowed, remaining = ct.check_budget()
+    assert not allowed
+    assert remaining == 0
+
+
+def test_budget_input_cost_counted_without_usage_breakdown():
+    """When the call site has no API usage breakdown (cache/uncached all 0),
+    tokens_sent must be used as the input count — input cost must not be
+    silently dropped from the budget."""
+    from headroom.proxy.server import CostTracker
+
+    ct = CostTracker(budget_limit_usd=100.0)
+    ct.record_tokens(
+        "claude-sonnet-4-20250514",
+        tokens_saved=0,
+        tokens_sent=500_000,
+    )
+
+    # 500k input tokens at Sonnet list price is ~$1.50 — must be > output-only
+    assert ct.get_period_cost() > 0.5

tests/test_request_outcome.py

@@ -293,6 +293,7 @@ async def test_funnel_passes_canonical_record_tokens_shape() -> None:
         "cache_write_5m_tokens": 80,
         "cache_write_1h_tokens": 20,
         "uncached_tokens": 0,
+        "output_tokens": 50,
     }