mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
5 commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
Ashish
|
a14ab45cf0
|
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
`
|
||
|
JD Davis
|
3c77e52ce4
|
feat: add Vertex AI proxy routing (#793)
## Description Adds first-class GCP Vertex AI proxy routing for publisher REST endpoints so Vertex requests are forwarded to a configurable regional Vertex host instead of falling through to the generic OpenAI/Anthropic/Gemini passthrough selection. Fixes #792 ## Type of Change - [x] New feature (non-breaking change that adds functionality) - [x] Documentation update ## Changes Made - Added a `vertex` provider target with `VERTEX_TARGET_API_URL` and `--vertex-api-url` support. - Registered explicit Vertex publisher routes for Google `generateContent`, `streamGenerateContent`, `countTokens` and Anthropic publisher `rawPredict`, `streamRawPredict` passthrough. - Added startup banner/routing output for Vertex AI. - Added focused tests for provider target resolution, CLI/env config, banner output, and route delegation. - Added `wiki/vertex.md` with usage examples and Google Cloud source links. ## Sources - Vertex AI Gemini inference reference: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/model-reference/inference - Google Cloud REST authentication: https://docs.cloud.google.com/docs/authentication/rest - Google Application Default Credentials: https://docs.cloud.google.com/docs/authentication/application-default-credentials ## Testing - [x] Linting passes (`python -m ruff check .`) - [x] New tests added for new functionality - [x] Focused unit tests pass - [ ] Full unit suite completed locally - [ ] Rust tests completed locally - [ ] Type checking passes locally ## Test Output ```text $ python -m ruff check . All checks passed! $ python -m pytest tests/test_provider_registry.py tests/test_provider_proxy_routes.py tests/test_cli_proxy_env.py tests/test_banner_upstream_targets.py -q 57 passed, 1 warning in 13.75s ``` Local limitations: - `python -m pytest tests scripts/tests -q` timed out after 1 hour on this Windows machine before completing. - `cargo test -p headroom-proxy --test integration_vertex_raw_predict` could not run because `cargo` is not installed on PATH in this environment. - The commit hook's `mypy` step fails locally on an existing Windows `fcntl` typing issue in `headroom/subscription/tracker.py`; `ruff`, `ruff-format`, and plugin-version hooks passed, and the commit was made with only `mypy` skipped. ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have made corresponding changes to the documentation - [x] I have added tests that prove the feature works - [ ] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable |
||
|
chopratejas
|
694589fec4 |
refactor(proxy): collapse 3 stream finalizers onto RequestOutcome.from_stream
Three streaming finalizers — ``_finalize_stream_response``,
``_stream_response_bedrock``, ``_stream_openai_via_backend`` — each
duplicated the same set of body- and config-derived fields when
constructing a ``RequestOutcome``:
* ``attempted_input_tokens = optimized_tokens + tokens_saved``
* ``num_messages = len(body.get("messages", []))``
* ``request_messages`` conditional on ``config.log_full_messages``
* ``transforms_applied`` list → tuple (frozen-dataclass contract)
* ``tags or {}`` normalization
* ``turn_id`` via ``compute_turn_id``
The last one was a real bug. Only the Bedrock site computed
``turn_id`` — sites 1 and 3 silently dropped it, breaking the
dashboard's multi-turn-session grouping for every Anthropic-SSE and
OpenAI-via-backend request. The new ``RequestOutcome.from_stream``
classmethod computes it uniformly so the three finalizers cannot
drift apart on derivation logic again.
Each call site now hands ``from_stream`` the body + provider-specific
cache/timing fields and gets a fully-constructed outcome back. The
funnel call after it stays identical (``await
self._record_request_outcome(outcome)``).
|
||
|
chopratejas
|
96a674f1b9 |
feat(proxy): add client (harness) identification — per-harness analytics across every handler
Cashes in the RequestOutcome refactor with a typed-field surface that
gives EVERY handler per-harness visibility — Codex / Claude Code /
aider / Cursor / Zed / opencode / DROID / antigravity / etc. — for one
field-add across the contract.
The one-field-add proof
Headroom went from "what fraction of OUR requests come from which
harness?" being unanswerable (handlers logged ad-hoc User-Agent strings
in heterogeneous tag dicts at 18 sites, with 9 of 18 not even
populating them) to one structured ``client: str | None`` value on
every observation flowing through the funnel. No new bookkeeping at
call sites; every handler picks it up via a single
``classify_client(headers)`` call at request entry.
Implementation
* New ``CLIENT_UA_MAP`` + ``classify_client()`` in
``headroom/proxy/auth_mode.py``. Substring match against
User-Agent; ``X-Client`` header overrides UA. Returns ``str | None``
so ``None`` is the loud "unidentified" signal rather than a silent
empty bucket.
* New ``RequestOutcome.client: str | None = None`` field.
* Funnel updates (in ``outcome.py``):
- Appends ``client=X`` to the PERF log line ONLY when set, so
``headroom perf --client X`` parsing stays clean for
unidentified traffic (no bogus ``client=`` token).
- Copies ``client`` into ``RequestLog.tags["client"]`` so the
dashboard's existing tag-based filtering surfaces per-harness
slicing with zero new columns.
* Every handler that constructs a RequestOutcome now passes
``client=client`` — wired across streaming.py (3 finalizers,
with ``_finalize_stream_response`` gaining a new optional kwarg
since it doesn't have direct access to headers), anthropic.py
(6 sites), openai.py (8 sites including Codex WS), gemini.py
(2 emitting sites), batch.py (5 sites).
Harnesses recognised
Anthropic ecosystem: claude-code, claude-cli, claude-vscode,
anthropic-cli
OpenAI ecosystem: codex-cli
Editors: cursor, zed
AI coding harnesses: aider, droid, opencode, github-copilot
Other: antigravity (Google experimental)
Adding a new client is a one-line edit to ``CLIENT_UA_MAP``.
Tests
* 8 new tests in ``test_request_outcome.py`` covering:
- ``client`` field round-trips on the value type
- ``classify_client`` against every recognised UA prefix
- ``X-Client`` header override beats UA match
- ``None`` for unknown traffic (the loud signal)
- Funnel appends ``client=X`` to PERF when set
- Funnel OMITS ``client=`` from PERF when None (no bogus empty)
- Funnel stamps ``client`` into ``RequestLog.tags``
* All 228 existing tests still pass (full sweep across streaming,
cache, Codex, Anthropic, OpenAI, Gemini, batch, auth-mode).
* ruff + ruff-format + mypy clean.
What's now true that wasn't before
Once this lands, the dashboard can answer:
* "Show me cache hit rate by harness"
→ ``GROUP BY tags.client FROM request_log``
* "Which harness contributes the most cache writes?"
→ same
* "Per-harness savings ratio"
→ same
* ``headroom perf --client codex`` / ``--client claude-code``
→ analyzer filters PERF log lines on ``client=X`` token
Zero new bookkeeping in handlers. Zero changes to Prometheus label
cardinality (kept the client dimension out of Prometheus on purpose —
the tags route is the right surface). The "what's our traffic split
by harness?" question is now answerable in three places (PERF log,
RequestLog tags, dashboard widgets that already filter on tags)
without any per-provider work.
|
||
|
chopratejas
|
e898f68b89 |
refactor(proxy): introduce RequestOutcome funnel; collapse 3 streaming finalizers
P0 audit (docs/superpowers/specs/P0-proxy-pipeline-audit.md) catalogued **18 metrics.record_request call sites** across 4 handler files with **4 distinct argument shapes**: 9 of 18 omitted `cached=`, 7 of 18 omitted `attempted_input_tokens=` (= bug #454/#455's "headline 0%"), only 4 sites emitted a `PERF` log line (= bug #327's "msgs=0" sibling — Codex traffic invisible to `headroom perf`), and `cache_hit` was hardcoded `False` at 9 of 18 RequestLog sites. The cause was structural, not tactical: every site was independently deciding what "record this completed request" meant. This PR puts a single value type + a single function between the handlers and the metrics layer. Two new files: * `headroom/proxy/outcome.py` — `RequestOutcome` frozen dataclass. Captures everything we ever need to record about one completed request: identity, tokens, cache stats (per-TTL splits + inferred flag for OpenAI), timing, transforms, diagnostics. Provider-specific fields default to neutral values so non-Anthropic handlers don't have to know about 5m/1h splits, non-OpenAI handlers don't have to know about inferred writes, etc. Computed properties (`cache_hit`, `cache_hit_pct`, `savings_pct`) make "forgot to compute it" mistakes structurally impossible. * `HeadroomProxy._record_request_outcome` in `server.py` — the single funnel. Owns the four downstream effects in canonical order: 1. `metrics.record_request(...)` with the FULL kwarg set 2. `cost_tracker.record_tokens(...)` with `(model, tokens_saved, optimized_tokens)` positional + all cache kwargs 3. `logger.log(RequestLog(...))` with `cache_hit` correctly derived 4. structured `PERF` log line in the canonical key=value shape Migrated three streaming finalizers in this PR: * `_finalize_stream_response` (Anthropic native + OpenAI HTTP streaming) * `_stream_response_bedrock` (Bedrock-native Anthropic streaming) * `_stream_openai_via_backend` (OpenAI/Azure backend via LiteLLM/AnyLLM) All three previously had inline, drifted versions of the four-call sequence. Each is now ~70 fewer lines: build a `RequestOutcome` from local context, call `self._record_request_outcome(outcome)`. The prefix-tracker mutation (Anthropic-specific) stays outside the funnel — different concern. Six more migrations queued for follow-up PRs (handle_anthropic_messages 6 sites, handle_openai_chat, handle_openai_responses, handle_openai_ responses_ws 2 sites, handle_gemini_*, handle_databricks_invocations). Each is mechanical now. Tests * New: `tests/test_request_outcome.py` — 14 tests covering value-type contract (frozen, derived properties, neutral defaults) + funnel contract (full record_request kwargs, canonical record_tokens shape, derived cache_hit in RequestLog, PERF log key=value format, optional cost_tracker/logger). Bind the real production method via descriptor binding so the test exercises the real implementation, not a fork. * All 135 existing streaming/cache/Codex tests pass with zero regressions (`tests/test_backend_streaming_cache_metrics.py`, `test_proxy_streaming_request_logger.py`, `test_proxy_streaming_resilience.py`, `test_proxy_anthropic_cache_stability.py`, `test_openai_codex_*`, `test_responses_ws_pyo3_compression.py`, `test_anthropic_pre_upstream_backpressure.py`). * `mypy headroom/proxy/{outcome,server,handlers/streaming}.py` clean. * `ruff check` clean. Surface impact * −238 lines from `handlers/streaming.py` (deduplication). * +92 lines in `server.py` (the funnel — counted ONCE, not 18×). * +130 lines in new `outcome.py` (frozen dataclass + docstrings). * Net production code: ~−16 lines today, ~−500 lines after the remaining six migrations land. Forward design constraints (per docs/superpowers/specs/P0-proxy-pipeline-audit.md §7) * KISS: one value type, one function, no factory hierarchies. * No regex in routing — handlers stay provider-specific in their upstream contract. Output unification only. * No silent fallbacks — `cache_hit` is computed, not defaulted. `cache_inferred=True` is the loud signal when OpenAI write count came from `_infer_openai_cache_write_tokens`. * PERF format frozen so `headroom/perf/analyzer.py` keeps parsing cleanly; P3 follow-up replaces the free-text shape with a structured event. |