akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions
2590 commits 410 branches 169 tags 142 MiB
Commit graph

3 commits

Author SHA1 Message Date
Abhay Singh
12f9f58cb3
fix(backends/litellm): None-guard core token counts in OpenAI usage block (#2324) 
## Description

`LiteLLMBackend.send_openai_message` builds the OpenAI-shape response
body. The core token counts are copied straight off LiteLLM's `Usage`
object with no guard, even though the cache fields immediately below
already use the defensive `int(getattr(..., 0) or 0)` form:

```python
usage_block: dict[str, Any] = {
    "prompt_tokens": response.usage.prompt_tokens,
    "completion_tokens": response.usage.completion_tokens,
    "total_tokens": response.usage.total_tokens,
}

# Defensive getattr right below:
cache_read = int(getattr(response.usage, "cache_read_input_tokens", 0) or 0)
cache_write = int(getattr(response.usage, "cache_creation_input_tokens", 0) or 0)
```

A provider can leave any of `prompt_tokens` / `completion_tokens` /
`total_tokens` as `None` on the `Usage` object. That `None` then lands
in `response.body["usage"]`, and the backend-routed OpenAI handler reads
it straight into arithmetic and the outcome ledger:

```python
usage = backend_response.body.get("usage", {})
output_tokens = usage.get("completion_tokens", 0)            # present key -> None, not the default
total_input_tokens = usage.get("prompt_tokens", optimized_tokens)  # present key -> None
...
uncached_input_tokens = max(0, total_input_tokens - cache_read_tokens - cache_write_tokens)  # None - int -> TypeError
...
RequestOutcome(..., output_tokens=output_tokens, ...)       # declared int; None crashes recording (e.g. prometheus += )
```

So an OpenAI-format request routed through a `--backend` (Bedrock /
Vertex / LiteLLM) whose provider returns a `None` count crashes on the
`max(0, None - ...)` subtraction, or later in outcome recording.
`.get(key, default)` does not help here because the key is present with
a `None` value, so the default never applies. This is the same class of
bug as the Anthropic-shape mapping and is fixed the same way.

## Fix

Coerce the three counts to `int` with the same defensive form already
used for the cache fields two lines down:

```python
usage_block: dict[str, Any] = {
    "prompt_tokens": int(getattr(response.usage, "prompt_tokens", 0) or 0),
    "completion_tokens": int(getattr(response.usage, "completion_tokens", 0) or 0),
    "total_tokens": int(getattr(response.usage, "total_tokens", 0) or 0),
}
```

No change for a normal integer usage; only a `None` (or absent) value
now becomes `0`.

## 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/backends/litellm.py`: `int`-coerce `prompt_tokens` /
`completion_tokens` / `total_tokens` in the `send_openai_message` usage
block.
- `tests/test_backends/test_litellm_cache_stats.py`: add
`test_none_core_counts_coerced_to_zero`, driving `send_openai_message`
with a `None`-count usage and asserting the block emits `int` `0`s.
- `CHANGELOG.md`: Bug Fixes entry.

## Testing

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

### Test Output

```text
$ uvx ruff@0.15.17 check headroom/backends/litellm.py tests/test_backends/test_litellm_cache_stats.py
All checks passed!
$ uvx ruff@0.15.17 format --check headroom/backends/litellm.py tests/test_backends/test_litellm_cache_stats.py
2 files already formatted
$ uvx mypy@1.20.2 --ignore-missing-imports headroom/backends/litellm.py
Success: no issues found in 1 source file
```

## Real Behavior Proof

- Environment: Windows 11, Python 3.12, `uvx ruff@0.15.17` / `uvx
mypy@1.20.2`. A full `pytest` OOMs this box (ML-stack import), so I
reproduced the field logic with a dependency-free script and left the
full pytest to CI.
- Exact command / steps: modeled the OLD (bare copy) and NEW
(`int(getattr(..., 0) or 0)`) field derivations for a `None` count, an
integer, and zero, then simulated the two downstream operations the
handler performs: `output_tokens += ...` and `max(0, prompt_tokens -
read - write)`.
- Observed result: OLD produced `None` and both downstream operations
raised `TypeError`; NEW produced `0` and both succeeded; an integer
count passed through unchanged. The added unit test drives
`send_openai_message` end to end (mocked `acompletion`) and asserts the
block emits `int` `0`s.
- Not tested: a live LiteLLM/Bedrock request that returns `None` counts;
the added test reuses the existing `_FakeUsage` / `_make_response` /
mocked-`acompletion` harness in `test_litellm_cache_stats.py`.

## 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
- [ ] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable

## Additional Notes

The "unit tests pass locally" box is unchecked because a local pytest
run imports the ML stack and OOMs this box; the added test reuses the
existing mocked-`acompletion` harness in `test_litellm_cache_stats.py`
and runs under the normal CI pytest job, and the behavior is
corroborated by the standalone proof above.

---------

Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
 2026-08-11 23:48:43 -05:00
Manmit Singh
54cfa361d3
fix(bedrock): fail fast when session-token auth lacks botocore (#1553) 
## Description

With `--backend bedrock` and **temporary** AWS credentials
(`AWS_SESSION_TOKEN`, as produced by SSO / STS assume-role /
`credential_process`), every request fails. litellm self-signs Bedrock
requests without botocore for *static* IAM keys, but as soon as a
session token is present it takes the `_auth_with_aws_session_token`
path in `litellm/llms/bedrock/base_aws_llm.py`, which imports
`botocore`. botocore is an optional dependency — it ships only with
headroom's `bedrock` extra, and the default Docker image is built with
`HEADROOM_EXTRAS=proxy,code`, so botocore is absent. The failure
surfaces only at request time as a misleading `authentication_error: No
module named 'botocore'` (and as a bare `Invalid API key` in Claude
Code).

This PR makes the Bedrock backend **fail fast at startup** with an
actionable message when a session token is set but botocore is missing —
directly addressing the "clearer error message" the reporter asked for.
It mirrors the existing optional-dependency guard pattern already used
for boto3 in `backends/litellm.py`.

Scope note: this does not change what the published image ships —
whether to add botocore/`bedrock` to the default image extras is a
separate sizing decision I left to maintainers. Static-credential
Bedrock users (who never hit the botocore path) are unaffected.

Refs #1551

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)

## Changes Made

- `headroom/backends/litellm.py`: when initializing the Bedrock backend
with `AWS_SESSION_TOKEN` set and `botocore` not importable, raise an
`ImportError` pointing at `pip install 'headroom-ai[bedrock]'` instead
of letting the request fail later with a misleading auth error.
- `tests/test_backends/test_bedrock_botocore_preflight.py`: regression
tests — the guard raises an actionable error for the
session-token-without-botocore case, and stays quiet for the
static-credential case.
- `CHANGELOG.md`: note under Unreleased → Fixed.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check`, `ruff format --check`)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

Regression test fails before the fix (no guard → no error raised),
passes after:

```text
# before fix (guard removed)
FAILED tests/test_backends/test_bedrock_botocore_preflight.py::test_bedrock_session_token_without_botocore_raises_actionable

# after fix
tests/test_backends/test_bedrock_botocore_preflight.py ..  [100%]
2 passed, 1 warning in 0.13s
```

`ruff check` / `ruff format --check` on the changed files: clean.

## Real Behavior Proof

- Environment: macOS (arm64), Python venv, editable install (`pip
install -e .`, no `bedrock` extra → botocore absent, matching the
reported slim-image condition), `pytest`.
- Exact command / steps: `python -m pytest
tests/test_backends/test_bedrock_botocore_preflight.py`. (1) Removed the
guard and ran the test → it failed because
`LiteLLMBackend(provider="bedrock")` with `AWS_SESSION_TOKEN` set and
botocore absent did NOT raise (reproducing the original "no early
signal" behavior). (2) Applied the guard. (3) Re-ran → both tests pass,
and the raised `ImportError` contains the `headroom-ai[bedrock]` install
hint.
- Observed result: with `AWS_SESSION_TOKEN` set and botocore not
importable, the backend now raises a clear, actionable `ImportError` at
construction time instead of deferring to litellm's later `No module
named 'botocore'` auth error. Without a session token the guard does not
fire, so static-credential users are unaffected.
- Not tested: I did not run a live Bedrock request against AWS with real
temporary credentials (no AWS account/STS access in this environment);
the reporter already confirmed that installing botocore makes the
identical request succeed, and this change surfaces that requirement at
startup.

## 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
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective
- [x] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md

Co-authored-by: JD Davis <mxjerrett@gmail.com>
 2026-07-01 21:02:15 -05:00
chopratejas
20dc1f28f3 fix(proxy): Strands MCP bundle + backend path fixes + Codex fail-closed protection 
Three logically-related sets of proxy changes ship in this branch:

1. Strands integration on the Bedrock path (HeadroomBundle + 4 OpenAI
   handler fixes + LiteLLM cache stats + dep pin)
2. /stats MCP aggregation (cross-process events log → proxy summary)
3. Codex compression-failure fail-closed (WS + HTTP /v1/responses)

== 1. Strands integration on the Bedrock path ==

* HeadroomBundle (headroom/integrations/strands/bundle.py): single-helper
  MCP wiring for a Strands Agent — Headroom MCP server (headroom_compress
  / headroom_retrieve / headroom_stats) plus optional Serena MCP and
  optional in-process compression hook. Constructor builds unstarted
  MCPClient instances per server; Strands' Agent owns the subprocess
  lifecycle. Default config: MCP enabled, Serena enabled, hook OFF
  (proxy is the single source of truth for compression). User-side
  integration is two lines in any Strands app.

* headroom/proxy/handlers/openai.py — backend path now:
  - calls PrefixCacheTracker.update_from_response (was direct-OpenAI only)
  - intercepts CCR headroom_retrieve tool_calls server-side, mirroring
    the Anthropic handler pattern; NO silent fallback, re-raises on
    CCR errors (per feedback_no_silent_fallbacks)
  - works for both non-streaming and streaming paths

* headroom/proxy/handlers/streaming.py: _stream_openai_via_backend now
  accepts prefix_tracker + optimized_messages, parses cache stats from
  the SSE final-usage frame (cache_creation_input_tokens added to the
  state machine), records CCR retrieve feedback via a new
  _record_ccr_feedback_from_openai_sse helper. Streaming CCR intercept
  is intentionally out of scope (mirrors Anthropic streaming behaviour).

* headroom/backends/litellm.py: send_openai_message response usage block
  now carries cache_read_input_tokens / cache_creation_input_tokens
  (Anthropic/Bedrock dialect) and prompt_tokens_details.cached_tokens
  (OpenAI dialect). Backwards-compatible — cold-start callers see the
  same 3-key shape; cache keys appear only when the underlying provider
  returns them. Pinned by test_no_cache_fields_means_no_cache_keys.

* headroom/proxy/auth_mode.py: ("strands-agents/", "strands") added to
  CLIENT_UA_MAP. Production callers should also set X-Client: strands
  since the default openai-python UA carries no Strands signal.

* pyproject.toml: huggingface-hub>=1.5.0,<2.0 pinned in [ml] so a sibling
  install (e.g. strands-agents) can't drag the version below the floor
  transformers 5.x requires (otherwise Kompress silently goes
  "unavailable").

== 2. /stats MCP aggregation ==

* headroom/proxy/cost.py: _aggregate_mcp_events() reads the cross-process
  shared events file the Headroom MCP server already writes to and
  surfaces summary.mcp with three new keys:
    - compressions       (count of headroom_compress invocations)
    - tokens_removed     (sum of input - output across those)
    - retrievals         (count of headroom_retrieve — the load-bearing
                          over-compression alarm; if it grows linearly
                          with turn count, lossy compressors are
                          dropping info the model actually needs)
  Defensive on every axis — missing MCP SDK, missing file, malformed
  events, read errors — never blocks /stats.

* examples/strands_bundle_demo.py: stats panel prints the new fields so
  the demo shows the full proxy-HTTP + MCP-tool story in one view.

== 3. Codex compression-failure fail-closed protection ==

Reported by Camille (2026-05-21): Codex threads were locking with
"ran out of room in the model's context window" after Headroom's
compression timed out on an oversized response.create frame and
forwarded the original ~1.7 MB frame to the upstream, which then
rejected it. Codex's auto-compact heuristic gates on the upstream-
reported total_usage_tokens (which Headroom had been shrinking on
earlier turns), so its compaction never fired and the thread locked.

Validated against open Codex issues (CLI + Desktop share codex-rs/core):
* #16068 — confirms compaction gates on total_usage_tokens,
  estimated_token_count is computed but only logged
* #19806 — confirms image token estimator unbounded, contributes to
  the same ContextManager.get_total_token_usage → auto-compaction chain

* headroom/proxy/helpers.py: decide_compression_failure_action() with a
  unit-tested decision matrix:
    - asyncio.TimeoutError                              → refuse, always
    - non-timeout failure + frame > 256 KiB (configurable) → refuse
    - non-timeout failure + small frame                 → forward (legacy)
  Operator escape hatches:
    - HEADROOM_WS_FAIL_OPEN_ON_COMPRESSION_FAILURE=1 restores legacy
    - HEADROOM_WS_COMPRESSION_FAIL_THRESHOLD_BYTES tunes the threshold

* headroom/proxy/handlers/openai.py (WS /v1/responses): consults the
  helper after compression failure. On refuse: close client websocket
  code 1009 with "headroom: compression <reason> — please compact
  context and retry" reason; set termination_cause for the outer
  lifecycle finally; return.

* headroom/proxy/handlers/openai.py (HTTP /v1/responses): same helper.
  On refuse: raise HTTPException(413) with a structured error body so
  FastAPI's HTTPException handler emits a clean 413. The existing
  `except HTTPException: raise` guard in this handler already ensures
  the 413 propagates without being swallowed by the 502 catch-all.

Anthropic /v1/messages NOT changed in this branch: no equivalent bug
report on Anthropic-protocol clients, Claude Code (Anthropic-owned)
handles context overflow via its own cache_control/ephemeral
primitives, and Cursor/Aider don't maintain the local-Y estimate the
Codex bug requires. Deferred until a real report lands; the patch is
a one-liner reusing the same helper.

== Tests + verification ==

* tests/test_backends/test_litellm_cache_stats.py — 3 tests pinning
  cache-stat surfacing across Anthropic/OpenAI dialects + backwards-
  compat for no-cache responses.
* tests/test_proxy/test_openai_backend_path.py — 5 tests (Bedrock cache
  fields, OpenAI fallback shape, CCR intercept with provider="openai",
  CCR re-raise on exception, streaming signature contract).
* tests/test_proxy/test_mcp_stats_aggregation.py — 5 tests pinning the
  aggregator across compress+retrieve mixes, empty events, unknown event
  types, missing token fields, and read failures.
* tests/test_proxy/test_compression_failure_action.py — 12 tests pinning
  the fail-closed decision matrix (timeout always refuses, small
  transient passes through, oversize refuses, env override variants,
  custom threshold, invalid threshold falls back, 0/negative ignored).

* examples/strands_bedrock_demo.py — model_id bumped from deprecated
  Claude 3 Haiku to Sonnet 4.5 (the deprecated model now errors on
  account access).
* examples/strands_via_proxy_demo.py — proxy + Bedrock cache + streaming
  smoke test.
* examples/strands_mcp_dispatch_test.py — pure MCP round-trip probe.
* examples/strands_bundle_demo.py — full Strands + HeadroomBundle E2E
  demo (this is the shape a real Strands user copies into their app).

Full pytest: 5327 passed, 178 skipped. The previously-failing
test_core_operations.py::TestAddBatch::test_add_batch_basic passes now
that the huggingface-hub pin in pyproject.toml unblocks transformers
imports.

E2E verified live against AWS Bedrock (Sonnet 4.5):
* cache_write=10,438 on turn A → cache_read=10,438 on turn B
* streaming SSE final usage frame carries cache_read_input_tokens
* 78.7% reduction on a 50 KB JSON tool_result via SmartCrusher (
  dispatched per-content-type by ContentRouter)
* Strands Agent + HeadroomBundle: model autonomously called
  headroom_compress + headroom_retrieve via MCP; CompressionStore
  round-trip succeeded; final answer correct.
 2026-05-21 11:00:14 -07:00