mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
15 commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
Tejas Chopra
|
f9807fd69e
|
feat(proxy): let extensions report cost savings and their own latency (#3051)
## What
Two changes that let a proxy extension report **what it saved** and
**what it cost**, so both show up under `/stats`, the dashboard, and
Prometheus.
`record_scope_savings` already existed and already accepted `usd` — the
one channel in the proxy that can express savings *without* tokens. Two
things stopped it working end to end.
### 1. Savings were silently dropped on Gemini traffic (bug)
`bind_scope` shares one attribution ledger between ASGI middleware and
the request handler. Anthropic and OpenAI call it; **Gemini never did**,
so anything an extension recorded into the request scope was discarded
for Gemini traffic only — silently, because an empty ledger and an
unbound one are indistinguishable at the outcome funnel. Now bound at
all four Gemini tag sites.
### 2. An extension's own latency was invisible (gap)
`overhead_ms` is measured *inside* the handler, and an ASGI extension
**wraps** that handler — so every millisecond it spends reaches the
client while every timing surface stays flat. An extension that halves
the bill and adds 200 ms per request is a trade the operator has to see
both halves of, and only one half was reaching the dashboard.
`record_scope_timing(scope, stage, ms)` is the symmetric counterpart to
`record_scope_savings`, carried on the same bound ledger and merged into
`RequestOutcome.pipeline_timing` at the outcome funnel — one place, so
every provider picks it up at once.
## API surface
```python
from headroom.proxy.savings_attribution import record_scope_savings, record_scope_timing
record_scope_savings(scope, "my_extension", tokens=0, usd=0.004) # money without tokens
record_scope_timing(scope, "my_extension", elapsed_ms)
```
Both take the ASGI `scope`, because middleware has no other way in.
Documented in `extensions.py` — the module extension authors actually
read, and the stability contract for this interface.
- Savings → `/stats` `savings.by_source`, dashboard card,
`headroom_savings_attributed_usd_total{source=...}`
- Timing → `/stats` `pipeline_timing`, dashboard Performance panel,
`headroom_transform_timing_ms_*`
**Attribution only.** These rows explain the headline total; they are
never added to it.
## Changes to existing behavior
- `public_tags` now strips `_headroom_stage_timing` as well as
`_headroom_savings_attribution`. Both ride on `tags` because that is the
one dict reaching the outcome funnel from every handler, and a list and
a dict must not land in a string-keyed label store.
- `pipeline_timing` passed to `metrics.record_request` is merged rather
than passed through **only when an extension contributed timings**; with
no extension the handler's own dict is passed through unchanged
(asserted by identity in the tests).
- Stage names are extension-supplied, so they are capped at 16 and
namespaced `ext:` — `deep_copy` reported by a plugin must never
accumulate into the same series as `deep_copy` measured by the pipeline.
A handler's own timing wins a collision (unreachable while the prefix
stands; the safe way round if it ever goes).
## Failure modes
Both calls are bounded (32 sources, 16 stages), never raise, and never
change a response — telemetry from a plugin must not be able to break
the request it is describing. Non-positive and non-numeric durations are
ignored: a zero is a clock artifact, not an observation, and averaging
it in would drag the mean down exactly where the stage is cheapest to
skip. `timings_from_tags` tolerates junk on the tag.
## Test-double fix
Three Gemini test fakes (`FakeRequest`, `_FakeRequest`,
`_VertexGeminiImageRequest`) had no `.scope`, which every real Starlette
`Request` has. They now do. This is a double that had drifted from the
type it stands in for; the alternative was weakening the handler to
tolerate a request shape that cannot occur in production.
---
## Real behavior proof
**Setup:** macOS 15.4 (darwin 25.4.0), Python 3.12.13, this branch at
`
|
||
|
Ashish Patel
|
41dab2d099
|
fix(ccr): verify a scanned marker's hash before advertising it (#2908)
## Description
`CCRToolInjector.scan_for_markers()` decides whether a compression
marker is Headroom's own by *shape* alone — any bracket marker carrying
a 24-hex hash counts, per the generic fallback pattern
(`\[.*?compressed.*?hash=([a-f0-9]{24})\]`). Other context tools emit
exactly that shape. Once a foreign hash is scanned,
`has_compressed_content` flips true and the retrieve tool + "Available
hashes" system instruction get injected for a hash this proxy never
stored — the model calls `headroom_retrieve`, gets a guaranteed miss,
and re-does work it already had. Two wasted turns per adopted foreign
hash.
Closes #2836
## 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/ccr/tool_injection.py`: added
`CCRToolInjector.verify_ownership()` — filters `detected_hashes` down to
hashes the compression store actually recognizes, via the same
`store.exists()` check the retrieve endpoint itself performs. Added a
small `_HashOwnershipStore` Protocol (structural typing, not a hard
dependency on the concrete `CompressionStore` class) and a
`compression_store` constructor field for dependency injection/testing.
`scan_for_markers()` itself is untouched — kept store-independent (pure
regex) rather than baking the check into the scan loop, since that
approach broke 24 existing tests that correctly test "does this shape
match" in isolation.
- `headroom/proxy/handlers/anthropic.py`,
`headroom/proxy/handlers/openai.py`: call `injector.verify_ownership()`
right after `scan_for_markers()` — the two real per-request call sites.
- `verify_ownership()` is also called inside `process_request()` (the
convenience wrapper `batch.py`'s Google path uses), so that path is
covered without a separate call site edit.
- `tests/test_ccr_tool_injection.py`: new `TestVerifyOwnership` class (6
tests) — the exact issue repro, a real-hash-survives case, mixed
own/foreign hashes, explicit store override, store-exception safety
(must not raise), and no-op-on-empty-hashes.
- `tests/test_proxy_anthropic_cache_stability.py`: 3 pre-existing tests
needed updating for the new (correct) behavior — two `_FakeInjector`
test doubles needed a `verify_ownership()` stub added, and one real
end-to-end test needed a genuine store entry seeded (via
`explicit_hash`) for the hash its hand-typed marker references, instead
of asserting on an unverified shape-only match.
## Testing
- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [ ] Type checking passes (`mypy headroom`) — not run locally, will
confirm via CI
- [x] New tests added for new functionality
- [x] Manual testing performed (see Real Behavior Proof)
### Test Output
```text
$ .venv/Scripts/python -m pytest tests/test_ccr_marker_policy.py tests/test_ccr_tool_always_on.py tests/test_ccr_tool_injection.py tests/test_proxy/test_ccr_frozen_prefix_coupling.py tests/test_proxy_anthropic_cache_stability.py tests/test_proxy_handlers_batch.py tests/test_proxy_handler_helpers.py -q
151 passed in 15.87s
$ .venv/Scripts/python -m pytest tests/test_proxy_ccr.py tests/test_proxy_openai_responses_stream_ccr.py tests/test_anthropic_ccr_workspace_unbound.py tests/test_compression_store.py tests/test_no_ccr_lossy.py tests/test_proxy_handlers_batch.py -q
121 passed in 20.73s
$ .venv/Scripts/ruff check . && .venv/Scripts/ruff format --check . # touched files only
All checks passed / already formatted
```
## Real Behavior Proof
- Environment: Windows 11, Python 3.14.5, local venv
- Exact command / steps: ran the issue's exact 3-line repro
(`CCRToolInjector.scan_for_markers()` on the foreign marker text, then
`verify_ownership()`) before and after the fix; separately verified a
genuinely-Headroom-stored hash (via `store.store(...,
explicit_hash=...)`) still survives verification and still drives
injection
- Observed result: before the fix (scan only, no verify step exists yet)
`has_compressed_content` is `True` for the foreign marker — matches the
bug report exactly. After adding `verify_ownership()`: foreign marker →
`detected_hashes == []`, `has_compressed_content is False`; real stored
hash → `detected_hashes == [real_hash]`, `has_compressed_content is
True`.
- Not tested: have not driven this through a live two-context-tool proxy
session (e.g. Headroom alongside another CCR-shaped tool in the same
conversation) — verified at the unit/integration level (the exact repro
plus the real proxy handler call sites via
`test_proxy_anthropic_cache_stability.py`'s end-to-end `TestClient`
tests), not via a live multi-tool session.
## 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 (N/A —
internal CCR safety behavior, no user-facing docs reference the
marker-adoption mechanism)
- [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 (release-please
generates this automatically from commit messages)
## Additional Notes
Design note on why `verify_ownership()` is a separate step rather than
baked into `scan_for_markers()`: my first attempt did exactly that and
broke 24 tests across `test_ccr_tool_injection.py`,
`test_ccr_marker_policy.py`, `test_proxy_anthropic_cache_stability.py`,
and `test_proxy_handler_helpers.py` — all of them legitimately testing
"does the regex detect this marker shape" independent of any store
state. Keeping the scan pure and adding an explicit, separately-testable
verification step kept that test surface intact while still closing the
real gap at the three places that actually decide whether to advertise
the retrieve tool.
|
||
|
Abhay Singh
|
1f2c681c0b
|
fix(proxy/batch): don't crash an OpenAI batch on a valid-JSON non-object line (#2316)
## Description
A single malformed line can abort compression for an entire OpenAI batch
upload.
`_compress_batch_jsonl` parses each JSONL line and immediately reads the
request body:
```python
request_obj = json.loads(line)
body = request_obj.get("body", {})
messages = body.get("messages", [])
...
except json.JSONDecodeError as e:
...
compressed_lines.append(line) # keep original on error
```
`json.loads` returns a valid JSON *value*, which isn't necessarily an
object. A line like `[1, 2, 3]`, `"hello"`, or `null` parses fine, but
`request_obj.get(...)` on a list/str/None raises `AttributeError`.
Likewise a request object whose `body` is present but not a dict
(`{"body": "..."}`) makes `body.get("messages", ...)` raise. The
surrounding `except json.JSONDecodeError` doesn't catch
`AttributeError`, so the exception propagates out of
`_compress_batch_jsonl` and the whole batch-create request fails.
This is the OpenAI batch upload path; the file is user-supplied, so a
single stray non-object line takes the batch down instead of just being
passed through like the other non-compressible cases already are.
## Fix
Guard for non-object shapes and pass them through unchanged:
```python
request_obj = json.loads(line)
if not isinstance(request_obj, dict):
compressed_lines.append(line)
total_requests += 1
continue
body = request_obj.get("body", {})
if not isinstance(body, dict):
body = {}
```
Closes #
## 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/handlers/batch.py`: in `_compress_batch_jsonl`, pass
through a non-dict parsed line and coalesce a non-dict `body` to `{}`.
- `tests/test_proxy_handlers_batch.py`: new test that array / string /
null lines and a non-dict `body` pass through without crashing and are
preserved.
- `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/proxy/handlers/batch.py tests/test_proxy_handlers_batch.py
All checks passed!
$ uvx mypy@1.20.2 --ignore-missing-imports headroom/proxy/handlers/batch.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` OOM-kills this box (ML stack import), so I
reproduced the per-line handling with a dependency-free script and left
the full pytest to CI.
- Exact command / steps: ran lines `[1,2,3]`, `"hello"`, `null`, and
`{"body": "not-a-dict"}` through the OLD (bare `.get`) and NEW
(isinstance-guarded) logic, plus a normal request and a `not-json` line
as controls.
- Observed result: OLD raises `AttributeError` on each non-object line
and on the non-dict body; NEW passes the non-object lines through
unchanged, coalesces the non-dict body to `{}`, still processes a normal
request, and still passes `not-json` through as a JSON-decode error.
- Not tested: a live OpenAI batch upload end-to-end; full local `pytest`
deferred to CI (OOM).
## 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 the full suite
imports the ML stack, which I can't run here. The new test reuses the
existing `DummyBatchHandler` / `install_batch_support_modules` harness
in `tests/test_proxy_handlers_batch.py` (the same one the neighbouring
invalid-line test uses), so it runs under the normal CI pytest job;
behaviour is additionally verified by the standalone proof above.
---------
Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
|
||
|
Rod Boev
|
2483f57002
|
fix(gemini): resolve native CCR retrieval calls (#2253)
## Description Buffered native Gemini requests currently return `headroom_retrieve` function calls to the client because `GeminiHandlerMixin` never invokes the shared CCR response handler. This wires native Gemini request and response translation into the provider handler while reusing the existing Google CCR extraction, retrieval, round-limit, mixed-tool, and `functionResponse` machinery. Streaming native Gemini and Gemini's OpenAI-compatible `MALFORMED_FUNCTION_CALL` behavior remain separate surfaces. This follows the current support boundary documented in https://github.com/headroomlabs-ai/headroom/pull/2044. Closes #2041 ## 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) - [x] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - Invoke `CCRResponseHandler` for successful buffered native Gemini responses containing `headroom_retrieve`. - Build Gemini-native continuation requests with the model `functionCall` and matching user `functionResponse`. - Inject the existing Google CCR function declaration while preserving sibling Gemini tool configurations. - Preserve mixed client-tool responses, streaming requests, non-CCR responses, and upstream error bodies. - Preserve Google `functionCall.id` as `functionResponse.id` through the shared CCR identity contract. - Leave streaming requests outside buffered CCR injection. - Fail closed when an exclusive CCR call remains unresolved after continuation. - Update the CCR documentation to describe buffered native Gemini support and the mixed-tool boundary. ## Testing - [x] Unit tests pass (`uv run pytest tests/test_proxy_handlers_batch.py -k "gemini_native_ccr or gemini_stream" -q`) - [x] Linting passes (`uv run ruff check headroom/proxy/handlers/gemini.py tests/test_proxy_handlers_batch.py`) - [ ] Type checking passes (`uv run mypy headroom`) - [x] New tests added for new functionality when applicable - [ ] Manual testing performed ### Test Output ```text Focused tests: `14 passed, 22 deselected` for `uv run pytest tests/test_proxy_handlers_batch.py -k "gemini_native_ccr or gemini_stream" -q`; `43 passed` for `uv run pytest tests/test_ccr_response_handler.py tests/test_ccr_response_handler_extra.py -q`. Scoped Ruff check and format check passed for the changed Python files. Full repository format remains blocked by pre-existing formatting outside this target. ``` ## Real Behavior Proof - Environment: Windows, Python 3.12, commit `7b3a92a8`; local native-shape behavioral harness with no Gemini credentials. - Exact command / steps: run the focused native Gemini handler tests, then capture a live `generateContent` request and continuation after a Gemini credential is available. - Observed result: local tests prove the buffered `functionCall` to `functionResponse` continuation, mixed-tool preservation, declaration preservation, error forwarding, and retrieval-result shapes. - Not tested: owner-reaching live Gemini continuation and native Gemini streaming CCR continuation. ## 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] 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 ## Additional Notes The patch keeps Gemini wire translation in `GeminiHandlerMixin` and extends the provider-neutral CCR identity fields for Google call ids. Native streaming continuation and the OpenAI-compatible Gemini round-two failure are outside this PR. --------- Co-authored-by: JerrettDavis <mxjerrett@gmail.com> |
||
|
Abhay Singh
|
a5d7e12c90
|
fix(proxy/batch): preserve sibling tool configs on Google batch requests (#2177)
## Description
When Headroom optimizes a Google/Gemini batch request, it silently drops
every tool config that isn't `functionDeclarations`.
In `handle_google_batch_create` the per-item optimizer extracts the
function declarations:
```python
tools = req_content.get("tools")
existing_funcs = None
if tools:
for tool in tools:
if "functionDeclarations" in tool:
existing_funcs = tool["functionDeclarations"]
break
```
and then rebuilds the forwarded request's tools as a single entry:
```python
if existing_funcs is not None:
compressed_req_content["tools"] = [{"functionDeclarations": existing_funcs}]
```
Gemini's `tools` array is a list of heterogeneous entries —
`{"functionDeclarations": [...]}` can sit alongside `{"googleSearch":
{}}` and `{"codeExecution": {}}`. Collapsing the array to one
`functionDeclarations` entry discards those siblings, so a batch request
that combines function calling with Google Search or code execution
reaches Google with those features stripped out. The request still
succeeds, so the loss is silent — the model just never grounds against
Search / never runs code.
The branch fires whenever the item had any `functionDeclarations` (or
CCR injected a retrieval tool), i.e. exactly the requests most likely to
also declare Search/code-execution.
## Fix
Rebuild the tools list from the original, replacing only the
`functionDeclarations` entry with the (possibly CCR-injected) funcs and
appending a new entry when the original had none:
```python
rebuilt_tools = []
replaced = False
for tool in tools or []:
if "functionDeclarations" in tool:
rebuilt_tools.append({**tool, "functionDeclarations": existing_funcs})
replaced = True
else:
rebuilt_tools.append(tool)
if not replaced:
rebuilt_tools.append({"functionDeclarations": existing_funcs})
compressed_req_content["tools"] = rebuilt_tools
```
Sibling entries (`googleSearch`, `codeExecution`, ...) are preserved in
place; the search/no-search behavior of the request is unchanged.
Closes #
## 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/handlers/batch.py`: preserve
non-`functionDeclarations` tool entries when rebuilding the optimized
Gemini batch request's tools array.
- `tests/test_proxy_handlers_batch.py`: new regression test asserting
`googleSearch` / `codeExecution` survive alongside
`functionDeclarations` in the forwarded body (uses the existing
`RealConvHandler` harness with the real converters).
- `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/proxy/handlers/batch.py tests/test_proxy_handlers_batch.py
All checks passed!
$ uvx mypy@1.20.2 --ignore-missing-imports headroom/proxy/handlers/batch.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` OOM-kills this box (ML stack import), so I
reproduced the array rebuild with a dependency-free script and left the
full pytest to CI.
- Exact command / steps: fed a tools array of
`[{functionDeclarations:[get_weather]}, {googleSearch:{}},
{codeExecution:{}}]` (plus a CCR-injected retrieval function) through
the OLD single-entry rebuild and the NEW preserving rebuild; also the
search-only case where CCR injects the first `functionDeclarations`.
- Observed result: OLD → `[{functionDeclarations:[...]}]` only
(googleSearch and codeExecution gone); NEW → all three entries retained
with the injected retrieval function present in `functionDeclarations`;
the search-only case gains a `functionDeclarations` entry while keeping
`googleSearch`.
- Not tested: a live Gemini batch submission; full local `pytest`
deferred to CI (OOM).
## 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 the full suite
imports the ML stack, which I can't run here. The new test reuses the
in-file `RealConvHandler` harness (same one the existing
`..._preserves_functioncall_response_order` test uses) so it runs under
the normal CI pytest job; behaviour is additionally verified by the
standalone proof above.
Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
Co-authored-by: Tejas Chopra <chopratejas@gmail.com>
|
||
|
Abhay Singh
|
4056117d90
|
fix(proxy/gemini): preserve non-text content across the compression round-trip (#2079)
## Description Two related content-loss bugs in the Gemini `contents[]` <-> `messages[]` compression round-trip. Both drop or misplace real user content that entries with **non-text** parts should carry through untouched. They share the same theme (non-text preservation), so they're bundled here as two commits. ### 1. Google batch handler restores preserved entries by the wrong index (`handlers/batch.py`) The `batchGenerateContent` handler restored preserved (non-text) entries with the raw-index loop that commit #836 (`_rebuild_gemini_contents`) replaced in the three non-batch Gemini handlers: ```python for orig_idx, original_content in preserved_contents.items(): if orig_idx < len(optimized_contents): optimized_contents[orig_idx] = original_content ``` `preserved_indices` are indices into the **original** `contents[]`, but `optimized_contents` is a **shorter** list (text-less entries produce no message). Indexing `optimized_contents` by `orig_idx` overwrites the wrong entry and drops any preserved entry whose original index is past the optimized length. For: ```python [user text, model functionCall, user functionResponse, model text] ``` the batch was forwarded to Google as **two** entries: the model's answer overwritten by the functionCall, and the functionResponse dropped. Unlike `gemini.py` there is no `if optimized_messages != messages` gate, so it runs on every mixed batch item. **Fix:** use the shared `_rebuild_gemini_contents` interleaving helper. ### 2. Code-execution parts not detected as non-text (`handlers/gemini.py`) `_has_non_text_parts` only recognized `inlineData`/`fileData`/`functionCall`/`functionResponse`. Gemini's code-execution feature emits `executableCode` and `codeExecutionResult` parts, echoed back in `contents[]` on later turns. Because they weren't detected: - a mixed `text`+`executableCode` entry lost its code payload (only the text survived the round-trip); - a text-less `executableCode`+`codeExecutionResult` entry was treated as a phantom in `_rebuild_gemini_contents` — it consumed the next optimized message, dropping the whole code turn and shifting a following user turn into the model's role slot (corrupting role alternation). **Fix:** add both keys to the non-text detection so those entries are preserved verbatim. Closes: no issue filed — both found while auditing the Gemini contents<->messages round-trip. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made - `headroom/proxy/handlers/batch.py`: use `_rebuild_gemini_contents` instead of the raw-index restore loop. - `headroom/proxy/handlers/gemini.py`: recognize `executableCode` / `codeExecutionResult` in `_has_non_text_parts`. - `tests/test_proxy_handlers_batch.py`: add `test_handle_google_batch_create_preserves_functioncall_response_order`, driving the handler with the **real** Gemini converters (the existing batch tests stub them, which hid the bug); mix `GeminiHandlerMixin` into the shared `DummyBatchHandler` so `_rebuild_gemini_contents` is available. - `tests/test_google_multimodal.py`: extend the parametrized `test_each_non_text_key_detected` to the two new keys, and add `test_code_execution_entry_survives`. ## Testing - [x] New regression tests added (`tests/test_proxy_handlers_batch.py`, `tests/test_google_multimodal.py`) - [x] Linting/formatting clean — run with the CI-pinned `ruff==0.15.17` - [ ] Full `pytest` deferred to CI (local-OOM reason below). ```text $ uvx ruff@0.15.17 check headroom/proxy/handlers/batch.py headroom/proxy/handlers/gemini.py \ tests/test_proxy_handlers_batch.py tests/test_google_multimodal.py All checks passed! ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.10, headroom from this branch. Importing `headroom` pulls in the torch/transformers stack and a full `pytest` gets OOM-killed on this box, so I verified the interleaving/detection with dependency-free scripts (replicating the Gemini converters, the old loop, and `_rebuild_gemini_contents`) and left the full pytest to CI. - Exact command / steps: ran two standalone scripts. Script 1 rebuilds a Gemini batch request with `preserved_indices` holding a `functionCall`/`functionResponse` pair and compares the old raw-index loop against `_rebuild_gemini_contents`. Script 2 feeds a `codeExecutionResult` entry through `_has_non_text_parts` and the preserve path with and without the two new allowlist keys. Also ran `uvx ruff@0.15.17 check` on the changed files and tests. - Observed result: the old batch loop drops the `functionResponse` and overwrites the answer (4 parts collapse to 2); `_rebuild_gemini_contents` keeps all 4. Without the new keys the code-execution entry is dropped/shifted (2 parts, code absent); with them it survives intact (3 parts, code present). Lint clean. See the two blocks below. Batch fix (bug #1): ```text preserved_indices: [1, 2] OLD result parts: ['text', 'functionCall'] len 2 NEW result parts: ['text', 'functionCall', 'functionResponse', 'text'] len 4 GEMINI BATCH REBUILD FIX VERIFIED (old drops response + overwrites answer; new keeps all 4) ``` Code-execution fix (bug #2): ```text (b) OLD len=2 NEW len=3 (a) OLD has code=False NEW has code=True GEMINI CODE-EXECUTION PRESERVE FIX VERIFIED (old drops/shifts; new keeps intact) ``` - Not tested: a live Google/Gemini round-trip (handlers stubbed, as the existing tests do). Full local `pytest` deferred to CI (OOM, per above). ## 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 — ran lint + standalone logic checks; full pytest deferred to CI (local OOM, disclosed above) - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes - Two small behavioral changes (one loop -> shared helper, two keys added to an allowlist) plus regression tests; no new dependencies. Both complete/extend the non-text preservation the non-batch handlers already do (the #836 line). - @JerrettDavis tagging you since you reviewed the recent Gemini fixes. Both of these drop content (functionResponse/images on batch; code-execution on the normal round-trip), so they seemed worth surfacing together. Thanks. --------- Co-authored-by: Tejas Chopra <chopratejas@gmail.com> |
||
|
Parideboy
|
46d5d685d9
|
fix(proxy): bound HF tokenizer load and offload token counting off event loop (#1738)
## Description Fixes #1701. On Windows, `headroom proxy --anthropic-api-url https://api.deepseek.com/anthropic` froze: the first `/v1/messages` request took ~610s (`optimization_latency_ms=609972`) with only router/lifecycle markers, and afterwards the whole server was a zombie — `/livez`, `/readyz` and `/health` hung until the process was killed. `HEADROOM_DETECT_BACKEND=python` was already set, so this was not the #575/#845 native-detect deadlock. Root cause: DeepSeek model names route to the HuggingFace tokenizer backend (`MODEL_PATTERNS` in `headroom/tokenizers/registry.py`). `HuggingFaceTokenizer` loads lazily, so the registry's construction-time fallback never fires; the first `count_messages` calls `AutoTokenizer.from_pretrained(..., trust_remote_code=True)` — unbounded network downloads/retries — and this ran **synchronously inside the async Anthropic messages handler** (`get_tokenizer(model)` + `tokenizer.count_messages(messages)`), outside the 30s `_run_compression_in_executor` bound. huggingface_hub retry chains on a restricted network easily reach ~10 minutes, blocking the entire asyncio event loop; subsequent on-loop counting kept it pinned. tiktoken got a bounded eager load for the same bug class long ago (#956); the HF backend never did. ## Type of Change - [x] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Refactoring (no functional changes) ## Changes Made - `headroom/tokenizers/huggingface.py`: `_load_tokenizer` now tries the local HF cache first (`local_files_only=True`, no network), then bounds the network load with `HEADROOM_HF_TOKENIZER_LOAD_TIMEOUT_SECS` (default 10s; `0` disables network loads) on a daemon thread. Timeouts/failures return `None` (cached by `lru_cache`, so the hub is probed at most once per process per tokenizer) and `count_messages` fails open to char-based estimation via the existing `_use_fallback()` path. - `headroom/proxy/handlers/anthropic.py`: new `AnthropicHandlerMixin._count_tokens_offloaded(model, messages)` runs `get_tokenizer` + `count_messages` on the compression executor bounded by `COMPRESSION_TIMEOUT_SECONDS`, failing open to `EstimatingTokenCounter` (downgrade logged once per model). Used in `handle_anthropic_messages` (the issue's hot path, both count sites) and `handle_anthropic_batch_create`; the batch path's inline `anthropic_pipeline.apply()` is now offloaded via `_run_compression_in_executor` (mirrors the #1612 image-compression offload). - `headroom/proxy/handlers/batch.py`: the two remaining inline `openai_pipeline.apply()` calls (`handle_google_batch_create`, `_compress_batch_jsonl`) are offloaded the same way; existing `except` blocks keep the pass-through fail-open semantics. - Tests: `tests/test_huggingface_tokenizer_timeout.py` (cache-first, bounded timeout, failure caching, timeout=0, fail-open estimation), `tests/test_tokenizer_count_offload.py` (wiring guards, runs on `headroom-compress` worker, event loop stays responsive during slow tokenizer work, fail-open), plus `_run_compression_in_executor` stub on the batch test double. ## Testing - [x] All existing tests pass - [x] Added new tests for the changes - [ ] Manual testing performed ``` $ python -m pytest tests/test_huggingface_tokenizer_timeout.py tests/test_tokenizer_count_offload.py tests/test_image_compression_offload.py tests/test_gemini_compression_offload.py tests/test_tokenizers tests/test_proxy_handlers_batch.py -q 50 passed $ ruff check . # No issues found $ ruff format --check . # 1043 files already formatted $ mypy headroom --ignore-missing-imports # 0 errors ``` ## Real Behavior Proof - Environment: Windows 11 Pro (10.0.26200), Python 3.13, local checkout of this branch with the Rust core built. - Exact command / steps: `python -m pytest tests/test_tokenizer_count_offload.py -q` — includes `test_count_tokens_offloaded_keeps_loop_responsive`, which reproduces the issue's mechanism: a tokenizer whose `count_messages` blocks (stand-in for the unbounded `AutoTokenizer.from_pretrained` network load) while an asyncio ticker measures event-loop liveness. Also `python -m pytest tests/test_huggingface_tokenizer_timeout.py -q` with a `from_pretrained` stub that sleeps 60s and `HEADROOM_HF_TOKENIZER_LOAD_TIMEOUT_SECS=0.2`. - Observed result: with the fix, the slow count runs on a `headroom-compress` worker thread and the loop keeps ticking (`ticks >= 5`; inline it yields ~0 — the zombie). The 60s-hung HF load unblocks at the 0.2s timeout, falls back to estimation, and the second call returns instantly (failure cached, no re-probe). All 10 new tests pass. - Not tested: live reproduction against `api.deepseek.com` from a network where HF hub downloads stall (the reporter's exact environment); actual HF vocab download timing on a healthy network. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review Co-authored-by: Claude Fable 5 <noreply@anthropic.com> |
||
|
chopratejas
|
3ec549288a |
fix(proxy): thread tags into 13 outcome sites + synth /v1/models + free-fn _extract_tags
Three fixes bundled; all in admin / cache-hit paths where tests didn't catch the regression. ## (A) 13 RequestOutcome sites missing tags= An AST audit found that 13 of 21 ``RequestOutcome(...)`` construction sites across the four handler files emitted outcomes without threading ``tags=``. Affected paths: * ``handle_anthropic_messages`` — the ``from_response_cache=True`` early-return outcome (Claude Code cache-hit turns dashboard-blind) * ``handle_openai_chat`` — same cache-hit early-return (Codex + Cursor + Continue cache-hit turns dashboard-blind) * ``handle_openai_responses_ws`` — the per-turn outcome inside the Codex WS session. The stale comment that said "ws_session_tags is not yet bound" was wrong — ``ws_tags`` was already extracted at handler entry * ``handle_anthropic_batch_create / batch_passthrough / batch_results`` * ``handle_passthrough`` (OpenAI Models / Files / List-Batches) * ``handle_google_batch_create / batch_passthrough / batch_results`` * ``_google_batch_passthrough`` (internal helper) * ``handle_batch_create`` (OpenAI batch entry) * ``handle_gemini_count_tokens`` (also fixed in #479; identical) Pattern of the fix is uniform: pull tags from headers and thread them into the ``RequestOutcome`` construction. New contract test ``test_handler_outcome_tag_invariant.py`` walks each handler file's AST and asserts every ``RequestOutcome`` site inside any ``handle_*`` or ``*_passthrough`` method passes both ``tags=`` and ``client=``. Future handlers get a clear test failure with file + line + method name if they regress. ## (B) Issue #478 — /v1/models 403 under Codex ChatGPT auth Codex Desktop with ChatGPT-subscription OAuth polls ``/v1/models`` to populate its model picker. Forwarding to ``chatgpt.com/backend-api/ models`` returned 403 to OAuth tokens. Fix: synthesize an OpenAI- compatible payload locally from a known-supported model set (``gpt-5.5`` through ``gpt-5``). All other ChatGPT-auth paths still forward as before — only model-metadata gets the local response. ## (C) Move _extract_tags to free function (mixin-isolation test compat) Handlers called ``self._extract_tags(headers)``. That worked in production where ``HeadroomProxy`` composes every mixin and defines the method, but broke tests that instantiate a single mixin via ``object.__new__(OpenAIHandlerMixin)``. The free-function form removes that coupling — handlers import ``extract_tags`` from ``headroom.proxy.helpers`` and call directly. ``HeadroomProxy. _extract_tags`` is kept as a thin wrapper for any external caller still using the method form. 17 call sites migrated. ## Zero behavior change for existing users Claude Code, Codex, Cursor, Continue, Aider, Gemini-routed harnesses all hit handlers that already extracted tags. Their wire bytes to upstream LLMs are byte-identical. Only the dashboard view gains tags on previously-blind paths. Closes #478. |
||
|
chopratejas
|
299d995066 |
test: wire test dummies to RequestOutcome funnel for batch + copilot-auth paths
Full regression sweep found 7 failures in test dummies (out of 4242
tests) that didn't have the production handler interface my refactor
now requires. All same root cause: the dummies need
``_record_request_outcome`` to delegate to the funnel; the
copilot-auth passthrough dummy also needs ``_next_request_id``
because the migrated passthrough handler now allocates an ID at
record-time.
Failures:
test_proxy_handlers_batch.py (6 sites — all DummyBatchHandler)
test_proxy_copilot_auth_hooks.py (1 site — Dummy in passthrough test)
Fix is the same pattern used in the earlier dummy fixes
(test_anthropic_pre_upstream_backpressure, test_openai_codex_routing,
test_openai_codex_ws_lifecycle):
async def _record_request_outcome(self, outcome):
from headroom.proxy.outcome import emit_request_outcome
await emit_request_outcome(self, outcome)
After the fix: 22 / 22 in the previously-failing tests; full
regression sweep 4242 / 4242 with zero failures (179 skipped, all
opt-in real-API).
Unrelated env issues observed in the same sweep but skipped:
* tests/test_memory/* — huggingface-hub<2.0 / transformers version
drift in local venv. Pre-existing, not caused by this refactor.
* tests/integrations/* — same env class.
* tests/test_realignment_live_multi_turn.py — opt-in live tests
needing API keys.
|
||
|
chopratejas
|
f0dcc02775 |
fix: A3 — byte-faithful Python forwarders; serialize canonical only when mutated
Eliminates P0-2 universally. Every Python forwarder (server.py
`_retry_request`, handlers/streaming.py `_stream_response`,
handlers/openai.py `_ws_http_fallback`, handlers/batch.py `_batch_passthrough`
+ batch-create + Google batch passthrough, handlers/anthropic.py CCR
continuation + batch endpoint) now switches from `httpx ... json=body` to
`httpx ... content=raw_bytes`. The default httpx JSON encoder was
re-serializing every request with `, `/`: ` separators and `\\uXXXX` ASCII
escapes — collapsing Anthropic prompt-cache hit-rate.
Forwarder strategy:
- unmutated body → forward `await request.body()` verbatim;
- mutated body → re-serialize once via the new
`serialize_body_canonical(body) -> bytes` helper (compact separators,
`ensure_ascii=False`, dict insertion order preserved).
`HEADROOM_PROXY_PYTHON_FORWARDER_MODE` env var configures the mode:
- `byte_faithful` (default) — the new behavior;
- `legacy_json_kwarg` — explicit operator opt-in for emergency rollback.
Documented in `docs/content/docs/configuration.mdx`. NOT a fallback —
unknown values raise loudly per build constraint #4.
`BodyMutationTracker` accompanies each request through the handler so
transform sites mark the tracker (`memory_injection`,
`image_compression`, `compression_*`, `batch_compression`,
`ccr_continuation`, etc.). At forwarder dispatch we additionally compare
the final body dict against the parsed original bytes as a structural
safety net — any silent mutation we missed still triggers canonical
re-serialization.
A2 follow-up: `handlers/openai.py:534-540` (Chat Completions memory
injection) was prepending a system message; replaced with
`append_text_to_latest_user_chat_message`, the OpenAI Chat Completions
analog of `_append_context_to_latest_non_frozen_user_turn`. The cache
hot zone (system messages) is now sacrosanct on /v1/chat/completions
too. Honors `HEADROOM_MEMORY_INJECTION_MODE=disabled`.
Structured logging: every forwarder emits an `event=outbound_request`
log line with `forwarder`, `path`, `body_bytes`, `body_mutated`,
`mutation_reasons`, `source` (passthrough|canonical|legacy),
`request_id`. Never logs Authorization or full body.
`_read_request_json` factored to share `_read_request_body_bytes` with
new `read_request_json_with_bytes` so the anthropic handler can capture
both the parsed dict and the original (decompressed) bytes.
Tests:
- `tests/test_proxy_byte_faithful_forwarding.py` (28 tests):
SHA-256 byte-equality on /v1/messages and streaming, unicode
preservation, numeric precision, mutation-tracker invariants,
canonical-serializer properties, legacy-mode rollback, OpenAI
Chat memory routing.
- Existing test mocks updated to accept the new `**kwargs` on
`_retry_request` (no behavior change).
- `tests/test_proxy_handlers_batch.py` updated to read the captured
`content=` bytes (formerly `json=`).
- One A2 test corrected (`test_anthropic_tool_sort_and_context_append_helpers`)
to match the live-zone-tail semantics introduced by A2.
Constraints satisfied: configurable env var; no new regex / hardcodes;
no silent fallback (`legacy_json_kwarg` is operator opt-in);
performant (`prepare_outbound_body_bytes` is O(1) for passthrough);
elegant single-responsibility helpers; structured tracing logs.
|
||
|
JerrettDavis
|
a5a4486a30 |
test: apply linux ruff formatting
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> |
||
|
JerrettDavis
|
50e712ac24 |
test: normalize batch handler test formatting
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> |
||
|
JerrettDavis
|
65481d243a |
test: restore CCR module cleanup in batch tests
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> |
||
|
JerrettDavis
|
6e9ea54a04 |
test: fix PR formatting and cover log compressor
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> |
||
|
JerrettDavis
|
38bf3e639c |
test: expand coverage across helper slices
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> |