2026-04-21 23:09:32 -05:00
|
|
|
from __future__ import annotations
|
2026-04-24 15:33:30 +02:00
|
|
|
|
2026-05-10 12:51:28 -04:00
|
|
|
import asyncio
|
2026-04-21 23:09:32 -05:00
|
|
|
import base64
|
2026-04-23 20:41:59 -04:00
|
|
|
import builtins
|
2026-04-21 23:09:32 -05:00
|
|
|
import json
|
2026-05-10 12:51:28 -04:00
|
|
|
from types import SimpleNamespace
|
2026-04-23 20:41:59 -04:00
|
|
|
from unittest.mock import patch
|
2026-04-24 15:33:30 +02:00
|
|
|
|
2026-05-10 12:51:28 -04:00
|
|
|
import httpx
|
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
2026-06-10 01:05:30 -05:00
|
|
|
from fastapi.responses import StreamingResponse
|
2026-05-10 12:51:28 -04:00
|
|
|
|
2026-04-21 23:09:32 -05:00
|
|
|
from headroom.proxy.handlers.anthropic import AnthropicHandlerMixin
|
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
2026-06-10 01:05:30 -05:00
|
|
|
from headroom.proxy.handlers.openai import (
|
|
|
|
|
OpenAIHandlerMixin,
|
|
|
|
|
_decode_openai_bearer_payload,
|
|
|
|
|
_passthrough_usage_from_json,
|
fix(codex): poll /wham/usage for subscription limits (handshake no longer sends x-codex-* headers) (#924)
## Description
Codex's subscription usage window (primary/secondary rate-limit gauges)
stopped populating for ChatGPT-OAuth sessions. This PR restores it by
polling Codex's dedicated usage endpoint instead of relying on response
headers that are no longer sent.
### Why the previous approach no longer works
The existing code populates `CodexRateLimitState` from `x-codex-*`
rate-limit headers captured on the `/v1/responses` WebSocket handshake
(`update_from_headers` at WS accept). That worked when OpenAI returned
`x-codex-primary-used-percent`, `x-codex-primary-window-minutes`, etc.
on the handshake response.
OpenAI has since stopped sending those headers on the ChatGPT WebSocket
handshake. I confirmed this by faithfully replaying a real Plus-account
handshake (both `prewarm` and regular `request_kind`): no `x-codex-*`
headers come back on either. This matches OpenAI's own move to a
dedicated usage endpoint (`GET /backend-api/codex/usage` in CodexApi
mode) and reports such as openai/codex#14728. So `update_from_headers`
now runs on every accept but finds nothing to parse, and the window
silently stays empty.
The headers aren't coming back, so there is nothing to fix in the
parsing path. The data now lives behind a request we have to make
ourselves.
## Type of Change
- [x] Bug fix (non-breaking change that fixes an issue)
## Changes Made
- `subscription/codex_rate_limits.py`:
- `parse_codex_usage_payload()` / `update_from_usage_payload()` — map
the `GET /backend-api/wham/usage` JSON (`rate_limit.primary_window` /
`secondary_window` with `used_percent`, `limit_window_seconds`,
`reset_at`; `credits`; `rate_limit_reached_type`) into the existing
`CodexRateLimitState`. `limit_window_seconds` is converted to
window-minutes with the same round-up codex-rs uses (`(secs + 59) //
60`).
- `maybe_schedule_usage_poll()` — fire-and-forget, throttled to one
request per 60s, scoped to ChatGPT sessions (requires both a Bearer
token and `ChatGPT-Account-Id`; API-key traffic is skipped). Uses an
in-flight guard so concurrent accepts don't stack polls. Endpoint is
overridable via `HEADROOM_CODEX_USAGE_URL`.
- `proxy/handlers/openai.py`:
- At the Codex WS accept site, after the now-usually-empty
`update_from_headers` block, schedule the usage poll. Wrapped in
`contextlib.suppress` and fully non-blocking so it can never delay or
fail the WebSocket accept.
The old header-capture path is intentionally left in place as a no-cost
fallback in case OpenAI restores the headers.
## 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 (live `/wham/usage` replay against a Plus
account returned HTTP 200 with the expected schema; payload fixture in
tests mirrors that real shape)
## Test Output
```
$ uv run pytest tests/test_codex_rate_limits.py -q
........................................ [100%]
41 passed in 0.16s
$ uv run ruff check headroom/subscription/codex_rate_limits.py headroom/proxy/handlers/openai.py tests/test_codex_rate_limits.py
All checks passed!
$ uv run mypy headroom/subscription/codex_rate_limits.py
Success: no issues found in 1 source file
```
## Additional Notes
- New tests cover: full-payload mapping, window-minutes round-up,
credits balance kept only when `has_credits`, promo object vs string,
empty payload returns `None`, missing `used_percent` skipped,
header-gating (requires Bearer + account-id), poll throttling, and
no-event-loop safety.
- Scoping to `ChatGPT-Account-Id` keeps the poll off API-key traffic,
and the 60s throttle plus in-flight guard bound it to at most one
lightweight GET per minute per running proxy.
---------
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-13 00:03:14 +02:00
|
|
|
_prefers_http1_passthrough,
|
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
2026-06-10 01:05:30 -05:00
|
|
|
)
|
2026-05-09 13:47:53 -07:00
|
|
|
from headroom.proxy.helpers import _headroom_bypass_enabled
|
2026-05-10 12:51:28 -04:00
|
|
|
from headroom.proxy.server import HeadroomProxy
|
2026-04-24 15:33:30 +02:00
|
|
|
|
|
|
|
|
|
2026-04-21 23:09:32 -05:00
|
|
|
def _jwt(payload: object) -> str:
|
|
|
|
|
header = {"alg": "none", "typ": "JWT"}
|
2026-04-24 15:33:30 +02:00
|
|
|
|
2026-04-21 23:09:32 -05:00
|
|
|
def encode(part: object) -> str:
|
|
|
|
|
raw = json.dumps(part, separators=(",", ":")).encode("utf-8")
|
|
|
|
|
return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=")
|
2026-04-24 15:33:30 +02:00
|
|
|
|
2026-04-21 23:09:32 -05:00
|
|
|
return f"{encode(header)}.{encode(payload)}."
|
2026-04-24 15:33:30 +02:00
|
|
|
|
|
|
|
|
|
2026-04-21 23:09:32 -05:00
|
|
|
class _ImageCompressor:
|
|
|
|
|
def __init__(self, compressed_message):
|
|
|
|
|
self._compressed_message = compressed_message
|
2026-04-24 15:33:30 +02:00
|
|
|
|
2026-04-21 23:09:32 -05:00
|
|
|
def compress(self, messages, provider): # noqa: ANN001, ANN201
|
|
|
|
|
assert provider == "anthropic"
|
|
|
|
|
return [self._compressed_message]
|
2026-04-24 15:33:30 +02:00
|
|
|
|
|
|
|
|
|
2026-04-23 20:41:59 -04:00
|
|
|
class _FreshCompressor:
|
|
|
|
|
instances = 0
|
|
|
|
|
|
|
|
|
|
def __init__(self):
|
|
|
|
|
type(self).instances += 1
|
|
|
|
|
|
|
|
|
|
|
2026-05-10 12:51:28 -04:00
|
|
|
class _TimeoutHttpClient:
|
|
|
|
|
async def request(self, **kwargs): # noqa: ANN001, ANN201
|
|
|
|
|
raise httpx.ConnectTimeout("connect timed out")
|
|
|
|
|
|
|
|
|
|
|
fix(codex): poll /wham/usage for subscription limits (handshake no longer sends x-codex-* headers) (#924)
## Description
Codex's subscription usage window (primary/secondary rate-limit gauges)
stopped populating for ChatGPT-OAuth sessions. This PR restores it by
polling Codex's dedicated usage endpoint instead of relying on response
headers that are no longer sent.
### Why the previous approach no longer works
The existing code populates `CodexRateLimitState` from `x-codex-*`
rate-limit headers captured on the `/v1/responses` WebSocket handshake
(`update_from_headers` at WS accept). That worked when OpenAI returned
`x-codex-primary-used-percent`, `x-codex-primary-window-minutes`, etc.
on the handshake response.
OpenAI has since stopped sending those headers on the ChatGPT WebSocket
handshake. I confirmed this by faithfully replaying a real Plus-account
handshake (both `prewarm` and regular `request_kind`): no `x-codex-*`
headers come back on either. This matches OpenAI's own move to a
dedicated usage endpoint (`GET /backend-api/codex/usage` in CodexApi
mode) and reports such as openai/codex#14728. So `update_from_headers`
now runs on every accept but finds nothing to parse, and the window
silently stays empty.
The headers aren't coming back, so there is nothing to fix in the
parsing path. The data now lives behind a request we have to make
ourselves.
## Type of Change
- [x] Bug fix (non-breaking change that fixes an issue)
## Changes Made
- `subscription/codex_rate_limits.py`:
- `parse_codex_usage_payload()` / `update_from_usage_payload()` — map
the `GET /backend-api/wham/usage` JSON (`rate_limit.primary_window` /
`secondary_window` with `used_percent`, `limit_window_seconds`,
`reset_at`; `credits`; `rate_limit_reached_type`) into the existing
`CodexRateLimitState`. `limit_window_seconds` is converted to
window-minutes with the same round-up codex-rs uses (`(secs + 59) //
60`).
- `maybe_schedule_usage_poll()` — fire-and-forget, throttled to one
request per 60s, scoped to ChatGPT sessions (requires both a Bearer
token and `ChatGPT-Account-Id`; API-key traffic is skipped). Uses an
in-flight guard so concurrent accepts don't stack polls. Endpoint is
overridable via `HEADROOM_CODEX_USAGE_URL`.
- `proxy/handlers/openai.py`:
- At the Codex WS accept site, after the now-usually-empty
`update_from_headers` block, schedule the usage poll. Wrapped in
`contextlib.suppress` and fully non-blocking so it can never delay or
fail the WebSocket accept.
The old header-capture path is intentionally left in place as a no-cost
fallback in case OpenAI restores the headers.
## 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 (live `/wham/usage` replay against a Plus
account returned HTTP 200 with the expected schema; payload fixture in
tests mirrors that real shape)
## Test Output
```
$ uv run pytest tests/test_codex_rate_limits.py -q
........................................ [100%]
41 passed in 0.16s
$ uv run ruff check headroom/subscription/codex_rate_limits.py headroom/proxy/handlers/openai.py tests/test_codex_rate_limits.py
All checks passed!
$ uv run mypy headroom/subscription/codex_rate_limits.py
Success: no issues found in 1 source file
```
## Additional Notes
- New tests cover: full-payload mapping, window-minutes round-up,
credits balance kept only when `has_credits`, promo object vs string,
empty payload returns `None`, missing `used_percent` skipped,
header-gating (requires Bearer + account-id), poll throttling, and
no-event-loop safety.
- Scoping to `ChatGPT-Account-Id` keeps the poll off API-key traffic,
and the 60s throttle plus in-flight guard bound it to at most one
lightweight GET per minute per running proxy.
---------
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-13 00:03:14 +02:00
|
|
|
class _RecordingHttpClient:
|
|
|
|
|
def __init__(self, label: str) -> None:
|
|
|
|
|
self.label = label
|
|
|
|
|
self.calls = 0
|
|
|
|
|
|
|
|
|
|
async def request(self, **kwargs): # noqa: ANN001, ANN201
|
|
|
|
|
self.calls += 1
|
|
|
|
|
request = httpx.Request(kwargs["method"], kwargs["url"])
|
|
|
|
|
return httpx.Response(
|
|
|
|
|
200,
|
|
|
|
|
request=request,
|
|
|
|
|
headers={"content-type": "application/json"},
|
|
|
|
|
json={"client": self.label},
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class _ChatGPTAccountRequest:
|
|
|
|
|
method = "GET"
|
|
|
|
|
headers = {}
|
|
|
|
|
url = SimpleNamespace(path="/backend-api/me", query="")
|
|
|
|
|
|
|
|
|
|
async def body(self) -> bytes:
|
|
|
|
|
return b""
|
|
|
|
|
|
|
|
|
|
|
2026-05-10 12:51:28 -04:00
|
|
|
class _PassthroughRequest:
|
|
|
|
|
method = "GET"
|
|
|
|
|
headers = {}
|
|
|
|
|
url = SimpleNamespace(path="/favicon.ico", query="")
|
|
|
|
|
|
|
|
|
|
async def body(self) -> bytes:
|
|
|
|
|
return b""
|
|
|
|
|
|
|
|
|
|
|
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
2026-06-10 01:05:30 -05:00
|
|
|
class _VertexPassthroughRequest:
|
|
|
|
|
method = "POST"
|
|
|
|
|
headers = {}
|
|
|
|
|
url = SimpleNamespace(
|
|
|
|
|
path="/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:generateContent",
|
|
|
|
|
query="",
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
async def body(self) -> bytes:
|
|
|
|
|
return b'{"contents":[]}'
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class _VertexStreamPassthroughRequest:
|
|
|
|
|
method = "POST"
|
|
|
|
|
headers = {}
|
|
|
|
|
url = SimpleNamespace(
|
|
|
|
|
path="/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:streamGenerateContent",
|
|
|
|
|
query="alt=sse",
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
async def body(self) -> bytes:
|
|
|
|
|
return b'{"contents":[]}'
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class _VertexGeminiImageRequest:
|
|
|
|
|
method = "POST"
|
|
|
|
|
headers = {}
|
|
|
|
|
query_params = {}
|
|
|
|
|
url = SimpleNamespace(
|
|
|
|
|
path="/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:generateContent",
|
|
|
|
|
query="",
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
async def body(self) -> bytes:
|
|
|
|
|
return json.dumps(
|
|
|
|
|
{
|
|
|
|
|
"contents": [
|
|
|
|
|
{
|
|
|
|
|
"role": "user",
|
|
|
|
|
"parts": [
|
|
|
|
|
{
|
|
|
|
|
"inlineData": {
|
|
|
|
|
"mimeType": "image/png",
|
|
|
|
|
"data": "aW1hZ2U=",
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
],
|
|
|
|
|
}
|
|
|
|
|
]
|
|
|
|
|
}
|
|
|
|
|
).encode("utf-8")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class _VertexUsageClient:
|
|
|
|
|
async def request(self, **kwargs): # noqa: ANN001, ANN201
|
|
|
|
|
request = httpx.Request(kwargs["method"], kwargs["url"], content=kwargs["content"])
|
|
|
|
|
return httpx.Response(
|
|
|
|
|
200,
|
|
|
|
|
request=request,
|
|
|
|
|
headers={"content-type": "application/json"},
|
|
|
|
|
json={
|
|
|
|
|
"candidates": [{"content": {"parts": [{"text": "ok"}]}}],
|
|
|
|
|
"usageMetadata": {
|
|
|
|
|
"promptTokenCount": 11,
|
|
|
|
|
"candidatesTokenCount": 7,
|
|
|
|
|
"cachedContentTokenCount": 3,
|
|
|
|
|
},
|
|
|
|
|
},
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class _AsyncChunks(httpx.AsyncByteStream):
|
|
|
|
|
def __init__(self, chunks: list[bytes]) -> None:
|
|
|
|
|
self._chunks = chunks
|
|
|
|
|
|
|
|
|
|
async def __aiter__(self): # noqa: ANN204
|
|
|
|
|
for chunk in self._chunks:
|
|
|
|
|
yield chunk
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class _VertexStreamClient:
|
|
|
|
|
def __init__(self) -> None:
|
|
|
|
|
self.sent_url = ""
|
|
|
|
|
|
|
|
|
|
def build_request(self, method, url, headers, content): # noqa: ANN001, ANN201
|
|
|
|
|
self.sent_url = str(url)
|
|
|
|
|
return httpx.Request(method, url, headers=headers, content=content)
|
|
|
|
|
|
|
|
|
|
async def send(self, request, stream=False): # noqa: ANN001, ANN201
|
|
|
|
|
assert stream is True
|
|
|
|
|
return httpx.Response(
|
|
|
|
|
200,
|
|
|
|
|
request=request,
|
|
|
|
|
headers={"content-type": "text/event-stream"},
|
|
|
|
|
stream=_AsyncChunks(
|
|
|
|
|
[
|
|
|
|
|
b'data: {"candidates":[{"content":{"parts":[{"text":"hello"}]}}]}\n\n',
|
|
|
|
|
b'data: {"usageMetadata":{"promptTokenCount":13,'
|
|
|
|
|
b'"candidatesTokenCount":5,"cachedContentTokenCount":2}}\n\n',
|
|
|
|
|
]
|
|
|
|
|
),
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
2026-05-10 12:51:28 -04:00
|
|
|
class _RetryThenSuccessClient:
|
|
|
|
|
def __init__(self) -> None:
|
|
|
|
|
self.attempts = 0
|
|
|
|
|
|
fix(proxy): add an Anthropic buffered read-timeout override (#1331)
## Description
Buffered Anthropic `/v1/messages` requests still use Headroom's generic
300-second read timeout, which can produce proxy-generated `502
ReadTimeout` errors on long turns. This adds a dedicated buffered
Anthropic timeout, keeps it applied across CCR and memory continuations
plus batch paths, and makes the direct server entrypoint enforce the
same positive-integer contract as the Click CLI. Closes #1261.
## Type of Change
- [x] Bug fix (non-breaking change that fixes an issue)
## Changes Made
- Added `anthropic_buffered_request_timeout_seconds` for buffered
Anthropic reads.
- Routed `/v1/messages`, CCR continuation, memory continuation, batch
create, batch passthrough, and batch results through that timeout.
- Enforced the same positive-integer validation for
`HEADROOM_ANTHROPIC_BUFFERED_REQUEST_TIMEOUT_SECONDS` and
`--anthropic-buffered-request-timeout-seconds` in both startup paths.
- Added focused regressions and updated `CHANGELOG.md`.
## Testing
- [x] `uv run pytest tests/test_proxy/test_anthropic_buffered_timeout.py
tests/test_cli_proxy_improvements.py::TestNewEnvVarWiring`
- [x] `uv run ruff check .`
- [x] `uv run ruff format . --check`
### Test Output
```text
$ uv run pytest tests/test_proxy/test_anthropic_buffered_timeout.py tests/test_cli_proxy_improvements.py::TestNewEnvVarWiring
17 passed in 3.42s
$ uv run ruff check .
All checks passed!
$ uv run ruff format . --check
966 files already formatted
```
## Real Behavior Proof
- Environment: local FastAPI `TestClient` with stubbed retry and HTTP
client seams
- Exact command / steps: run `uv run pytest
tests/test_proxy/test_anthropic_buffered_timeout.py
tests/test_cli_proxy_improvements.py::TestNewEnvVarWiring`; the tests
build `ProxyConfig(request_timeout_seconds=7, connect_timeout_seconds=3,
anthropic_buffered_request_timeout_seconds=19)`, drive `/v1/messages`,
`/v1/messages/batches`, `/v1/messages/batches/{batch_id}/results`, a CCR
continuation, and a memory continuation through `TestClient`, then
verify `HEADROOM_ANTHROPIC_BUFFERED_REQUEST_TIMEOUT_SECONDS=0` falls
back to `600`, `--anthropic-buffered-request-timeout-seconds 0` is
rejected, and default proxy timeouts stay `read=300` and `write=300`
- Observed result: buffered Anthropic paths use
`httpx.Timeout(connect=3, read=19, write=7, pool=3)`, continuation
requests stay on that same budget, invalid zero-valued startup config is
rejected or ignored back to the default, and unrelated proxy timeout
defaults stay unchanged
- Not tested: live upstream Anthropic latency beyond the focused
stubbed-timeout regression
## 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 added tests that prove the fix
- [x] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable
2026-06-23 23:46:31 -04:00
|
|
|
async def post(self, url, content, headers, timeout=None): # noqa: ANN001, ANN201
|
2026-05-10 12:51:28 -04:00
|
|
|
self.attempts += 1
|
|
|
|
|
if self.attempts == 1:
|
|
|
|
|
raise httpx.ConnectTimeout("connect timed out")
|
fix(proxy): add an Anthropic buffered read-timeout override (#1331)
## Description
Buffered Anthropic `/v1/messages` requests still use Headroom's generic
300-second read timeout, which can produce proxy-generated `502
ReadTimeout` errors on long turns. This adds a dedicated buffered
Anthropic timeout, keeps it applied across CCR and memory continuations
plus batch paths, and makes the direct server entrypoint enforce the
same positive-integer contract as the Click CLI. Closes #1261.
## Type of Change
- [x] Bug fix (non-breaking change that fixes an issue)
## Changes Made
- Added `anthropic_buffered_request_timeout_seconds` for buffered
Anthropic reads.
- Routed `/v1/messages`, CCR continuation, memory continuation, batch
create, batch passthrough, and batch results through that timeout.
- Enforced the same positive-integer validation for
`HEADROOM_ANTHROPIC_BUFFERED_REQUEST_TIMEOUT_SECONDS` and
`--anthropic-buffered-request-timeout-seconds` in both startup paths.
- Added focused regressions and updated `CHANGELOG.md`.
## Testing
- [x] `uv run pytest tests/test_proxy/test_anthropic_buffered_timeout.py
tests/test_cli_proxy_improvements.py::TestNewEnvVarWiring`
- [x] `uv run ruff check .`
- [x] `uv run ruff format . --check`
### Test Output
```text
$ uv run pytest tests/test_proxy/test_anthropic_buffered_timeout.py tests/test_cli_proxy_improvements.py::TestNewEnvVarWiring
17 passed in 3.42s
$ uv run ruff check .
All checks passed!
$ uv run ruff format . --check
966 files already formatted
```
## Real Behavior Proof
- Environment: local FastAPI `TestClient` with stubbed retry and HTTP
client seams
- Exact command / steps: run `uv run pytest
tests/test_proxy/test_anthropic_buffered_timeout.py
tests/test_cli_proxy_improvements.py::TestNewEnvVarWiring`; the tests
build `ProxyConfig(request_timeout_seconds=7, connect_timeout_seconds=3,
anthropic_buffered_request_timeout_seconds=19)`, drive `/v1/messages`,
`/v1/messages/batches`, `/v1/messages/batches/{batch_id}/results`, a CCR
continuation, and a memory continuation through `TestClient`, then
verify `HEADROOM_ANTHROPIC_BUFFERED_REQUEST_TIMEOUT_SECONDS=0` falls
back to `600`, `--anthropic-buffered-request-timeout-seconds 0` is
rejected, and default proxy timeouts stay `read=300` and `write=300`
- Observed result: buffered Anthropic paths use
`httpx.Timeout(connect=3, read=19, write=7, pool=3)`, continuation
requests stay on that same budget, invalid zero-valued startup config is
rejected or ignored back to the default, and unrelated proxy timeout
defaults stay unchanged
- Not tested: live upstream Anthropic latency beyond the focused
stubbed-timeout regression
## 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 added tests that prove the fix
- [x] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable
2026-06-23 23:46:31 -04:00
|
|
|
del timeout
|
2026-05-10 12:51:28 -04:00
|
|
|
request = httpx.Request("POST", url, headers=headers, content=content)
|
|
|
|
|
return httpx.Response(200, request=request, content=b"{}")
|
|
|
|
|
|
|
|
|
|
|
2026-04-21 23:09:32 -05:00
|
|
|
def test_decode_openai_bearer_payload_handles_missing_and_non_mapping_payloads() -> None:
|
|
|
|
|
assert _decode_openai_bearer_payload({}) is None
|
|
|
|
|
assert _decode_openai_bearer_payload({"authorization": "Basic abc"}) is None
|
|
|
|
|
assert (
|
|
|
|
|
_decode_openai_bearer_payload({"authorization": f"Bearer {_jwt(['not', 'a', 'dict'])}"})
|
|
|
|
|
is None
|
|
|
|
|
)
|
2026-04-24 15:33:30 +02:00
|
|
|
|
|
|
|
|
|
2026-04-21 23:09:32 -05:00
|
|
|
def test_openai_handler_prefix_helpers_cover_edge_cases() -> None:
|
|
|
|
|
assert OpenAIHandlerMixin._strict_previous_turn_frozen_count([], 2) == 2
|
|
|
|
|
assert (
|
|
|
|
|
OpenAIHandlerMixin._strict_previous_turn_frozen_count(
|
|
|
|
|
[{"role": "assistant"}, {"role": "user"}],
|
|
|
|
|
0,
|
|
|
|
|
)
|
|
|
|
|
== 1
|
|
|
|
|
)
|
|
|
|
|
assert (
|
|
|
|
|
OpenAIHandlerMixin._strict_previous_turn_frozen_count(
|
|
|
|
|
[{"role": "user"}, {"role": "assistant"}],
|
|
|
|
|
0,
|
|
|
|
|
)
|
|
|
|
|
== 2
|
|
|
|
|
)
|
2026-04-24 15:33:30 +02:00
|
|
|
|
2026-04-21 23:09:32 -05:00
|
|
|
original = [{"role": "system", "content": "keep"}, {"role": "user", "content": "hello"}]
|
|
|
|
|
restored, changed = OpenAIHandlerMixin._restore_frozen_prefix(
|
|
|
|
|
original,
|
|
|
|
|
[],
|
|
|
|
|
frozen_message_count=1,
|
|
|
|
|
)
|
|
|
|
|
assert restored == [{"role": "system", "content": "keep"}]
|
|
|
|
|
assert changed == 1
|
2026-04-24 15:33:30 +02:00
|
|
|
|
2026-04-21 23:09:32 -05:00
|
|
|
restored, changed = OpenAIHandlerMixin._restore_frozen_prefix(
|
|
|
|
|
original,
|
|
|
|
|
[{"role": "system", "content": "changed"}, {"role": "user", "content": "hello"}],
|
|
|
|
|
frozen_message_count=1,
|
|
|
|
|
)
|
|
|
|
|
assert restored == original
|
|
|
|
|
assert changed == 1
|
2026-04-24 15:33:30 +02:00
|
|
|
|
|
|
|
|
|
2026-05-09 13:47:53 -07:00
|
|
|
def test_headroom_bypass_helper_is_transport_neutral() -> None:
|
|
|
|
|
assert _headroom_bypass_enabled({"x-headroom-bypass": "true"}) is True
|
|
|
|
|
assert _headroom_bypass_enabled({"x-headroom-bypass": " TRUE "}) is True
|
|
|
|
|
assert _headroom_bypass_enabled({"x-headroom-mode": "passthrough"}) is True
|
|
|
|
|
assert _headroom_bypass_enabled({"x-headroom-mode": " PASSTHROUGH "}) is True
|
|
|
|
|
assert _headroom_bypass_enabled({"x-headroom-bypass": "false"}) is False
|
|
|
|
|
assert _headroom_bypass_enabled({}) is False
|
|
|
|
|
assert _headroom_bypass_enabled(None) is False
|
|
|
|
|
assert OpenAIHandlerMixin._headroom_bypass_enabled({"x-headroom-bypass": "true"}) is True
|
|
|
|
|
|
|
|
|
|
|
2026-05-10 12:51:28 -04:00
|
|
|
def test_openai_passthrough_connect_timeout_returns_502() -> None:
|
|
|
|
|
handler = object.__new__(OpenAIHandlerMixin)
|
|
|
|
|
handler.http_client = _TimeoutHttpClient()
|
|
|
|
|
|
|
|
|
|
async def run():
|
|
|
|
|
return await handler.handle_passthrough(
|
|
|
|
|
_PassthroughRequest(),
|
|
|
|
|
"https://api.openai.com",
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
response = asyncio.run(run())
|
|
|
|
|
|
|
|
|
|
assert response.status_code == 502
|
|
|
|
|
payload = json.loads(response.body)
|
|
|
|
|
assert payload["error"]["type"] == "connection_error"
|
|
|
|
|
assert "Failed to connect to upstream API" in payload["error"]["message"]
|
|
|
|
|
|
|
|
|
|
|
fix(codex): poll /wham/usage for subscription limits (handshake no longer sends x-codex-* headers) (#924)
## Description
Codex's subscription usage window (primary/secondary rate-limit gauges)
stopped populating for ChatGPT-OAuth sessions. This PR restores it by
polling Codex's dedicated usage endpoint instead of relying on response
headers that are no longer sent.
### Why the previous approach no longer works
The existing code populates `CodexRateLimitState` from `x-codex-*`
rate-limit headers captured on the `/v1/responses` WebSocket handshake
(`update_from_headers` at WS accept). That worked when OpenAI returned
`x-codex-primary-used-percent`, `x-codex-primary-window-minutes`, etc.
on the handshake response.
OpenAI has since stopped sending those headers on the ChatGPT WebSocket
handshake. I confirmed this by faithfully replaying a real Plus-account
handshake (both `prewarm` and regular `request_kind`): no `x-codex-*`
headers come back on either. This matches OpenAI's own move to a
dedicated usage endpoint (`GET /backend-api/codex/usage` in CodexApi
mode) and reports such as openai/codex#14728. So `update_from_headers`
now runs on every accept but finds nothing to parse, and the window
silently stays empty.
The headers aren't coming back, so there is nothing to fix in the
parsing path. The data now lives behind a request we have to make
ourselves.
## Type of Change
- [x] Bug fix (non-breaking change that fixes an issue)
## Changes Made
- `subscription/codex_rate_limits.py`:
- `parse_codex_usage_payload()` / `update_from_usage_payload()` — map
the `GET /backend-api/wham/usage` JSON (`rate_limit.primary_window` /
`secondary_window` with `used_percent`, `limit_window_seconds`,
`reset_at`; `credits`; `rate_limit_reached_type`) into the existing
`CodexRateLimitState`. `limit_window_seconds` is converted to
window-minutes with the same round-up codex-rs uses (`(secs + 59) //
60`).
- `maybe_schedule_usage_poll()` — fire-and-forget, throttled to one
request per 60s, scoped to ChatGPT sessions (requires both a Bearer
token and `ChatGPT-Account-Id`; API-key traffic is skipped). Uses an
in-flight guard so concurrent accepts don't stack polls. Endpoint is
overridable via `HEADROOM_CODEX_USAGE_URL`.
- `proxy/handlers/openai.py`:
- At the Codex WS accept site, after the now-usually-empty
`update_from_headers` block, schedule the usage poll. Wrapped in
`contextlib.suppress` and fully non-blocking so it can never delay or
fail the WebSocket accept.
The old header-capture path is intentionally left in place as a no-cost
fallback in case OpenAI restores the headers.
## 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 (live `/wham/usage` replay against a Plus
account returned HTTP 200 with the expected schema; payload fixture in
tests mirrors that real shape)
## Test Output
```
$ uv run pytest tests/test_codex_rate_limits.py -q
........................................ [100%]
41 passed in 0.16s
$ uv run ruff check headroom/subscription/codex_rate_limits.py headroom/proxy/handlers/openai.py tests/test_codex_rate_limits.py
All checks passed!
$ uv run mypy headroom/subscription/codex_rate_limits.py
Success: no issues found in 1 source file
```
## Additional Notes
- New tests cover: full-payload mapping, window-minutes round-up,
credits balance kept only when `has_credits`, promo object vs string,
empty payload returns `None`, missing `used_percent` skipped,
header-gating (requires Bearer + account-id), poll throttling, and
no-event-loop safety.
- Scoping to `ChatGPT-Account-Id` keeps the poll off API-key traffic,
and the 60s throttle plus in-flight guard bound it to at most one
lightweight GET per minute per running proxy.
---------
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-13 00:03:14 +02:00
|
|
|
def test_prefers_http1_passthrough_matches_chatgpt_hosts_only() -> None:
|
|
|
|
|
assert _prefers_http1_passthrough("https://chatgpt.com") is True
|
|
|
|
|
assert _prefers_http1_passthrough("https://chatgpt.com/backend-api/me") is True
|
|
|
|
|
assert _prefers_http1_passthrough("https://api.chatgpt.com") is True
|
|
|
|
|
assert _prefers_http1_passthrough("https://CHATGPT.COM/backend-api/me") is True
|
|
|
|
|
assert _prefers_http1_passthrough("https://api.openai.com") is False
|
|
|
|
|
assert _prefers_http1_passthrough("https://notchatgpt.com") is False
|
|
|
|
|
assert _prefers_http1_passthrough("https://chatgpt.com.evil.com") is False
|
|
|
|
|
assert _prefers_http1_passthrough("") is False
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_chatgpt_passthrough_uses_http1_client() -> None:
|
|
|
|
|
handler = object.__new__(OpenAIHandlerMixin)
|
|
|
|
|
handler.http_client = _RecordingHttpClient("h2")
|
|
|
|
|
handler.http_client_h1 = _RecordingHttpClient("h1")
|
|
|
|
|
|
|
|
|
|
response = asyncio.run(
|
|
|
|
|
handler.handle_passthrough(_ChatGPTAccountRequest(), "https://chatgpt.com")
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
assert response.status_code == 200
|
|
|
|
|
assert json.loads(response.body)["client"] == "h1"
|
|
|
|
|
assert handler.http_client.calls == 0
|
|
|
|
|
assert handler.http_client_h1.calls == 1
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_non_chatgpt_passthrough_uses_default_client() -> None:
|
|
|
|
|
handler = object.__new__(OpenAIHandlerMixin)
|
|
|
|
|
handler.http_client = _RecordingHttpClient("h2")
|
|
|
|
|
handler.http_client_h1 = _RecordingHttpClient("h1")
|
|
|
|
|
|
|
|
|
|
response = asyncio.run(
|
|
|
|
|
handler.handle_passthrough(_PassthroughRequest(), "https://api.openai.com")
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
assert response.status_code == 200
|
|
|
|
|
assert json.loads(response.body)["client"] == "h2"
|
|
|
|
|
assert handler.http_client.calls == 1
|
|
|
|
|
assert handler.http_client_h1.calls == 0
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_chatgpt_passthrough_falls_back_when_h1_client_missing() -> None:
|
|
|
|
|
handler = object.__new__(OpenAIHandlerMixin)
|
|
|
|
|
handler.http_client = _RecordingHttpClient("h2")
|
|
|
|
|
handler.http_client_h1 = None
|
|
|
|
|
|
|
|
|
|
response = asyncio.run(
|
|
|
|
|
handler.handle_passthrough(_ChatGPTAccountRequest(), "https://chatgpt.com")
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
assert response.status_code == 200
|
|
|
|
|
assert json.loads(response.body)["client"] == "h2"
|
|
|
|
|
assert handler.http_client.calls == 1
|
|
|
|
|
|
|
|
|
|
|
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
2026-06-10 01:05:30 -05:00
|
|
|
def test_passthrough_usage_normalizes_vertex_usage_metadata() -> None:
|
|
|
|
|
usage = _passthrough_usage_from_json(
|
|
|
|
|
{
|
|
|
|
|
"usageMetadata": {
|
|
|
|
|
"promptTokenCount": 11,
|
|
|
|
|
"candidatesTokenCount": 7,
|
|
|
|
|
"cachedContentTokenCount": 3,
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
assert usage == {
|
|
|
|
|
"input_tokens": 11,
|
|
|
|
|
"output_tokens": 7,
|
|
|
|
|
"cache_read_input_tokens": 3,
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_vertex_passthrough_records_usage_metadata_for_dashboard() -> None:
|
|
|
|
|
handler = object.__new__(HeadroomProxy)
|
|
|
|
|
handler.http_client = _VertexUsageClient()
|
|
|
|
|
outcomes = []
|
|
|
|
|
|
|
|
|
|
async def next_request_id(): # noqa: ANN202
|
|
|
|
|
return "req_vertex"
|
|
|
|
|
|
|
|
|
|
async def record(outcome): # noqa: ANN001, ANN202
|
|
|
|
|
outcomes.append(outcome)
|
|
|
|
|
|
|
|
|
|
handler._next_request_id = next_request_id
|
|
|
|
|
handler._record_request_outcome = record
|
|
|
|
|
|
|
|
|
|
response = asyncio.run(
|
|
|
|
|
handler.handle_passthrough(
|
|
|
|
|
_VertexPassthroughRequest(),
|
|
|
|
|
"https://vertex.test",
|
|
|
|
|
"generateContent",
|
|
|
|
|
"vertex:google",
|
|
|
|
|
)
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
assert response.status_code == 200
|
|
|
|
|
assert len(outcomes) == 1
|
|
|
|
|
outcome = outcomes[0]
|
|
|
|
|
assert outcome.provider == "vertex:google"
|
|
|
|
|
assert outcome.model == "gemini-2.0-flash"
|
|
|
|
|
assert outcome.optimized_tokens == 11
|
|
|
|
|
assert outcome.output_tokens == 7
|
|
|
|
|
assert outcome.cache_read_tokens == 3
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_vertex_stream_passthrough_preserves_chunks_and_records_usage() -> None:
|
|
|
|
|
handler = object.__new__(HeadroomProxy)
|
|
|
|
|
handler.http_client = _VertexStreamClient()
|
|
|
|
|
outcomes = []
|
|
|
|
|
|
|
|
|
|
async def next_request_id(): # noqa: ANN202
|
|
|
|
|
return "req_vertex_stream"
|
|
|
|
|
|
|
|
|
|
async def record(outcome): # noqa: ANN001, ANN202
|
|
|
|
|
outcomes.append(outcome)
|
|
|
|
|
|
|
|
|
|
handler._next_request_id = next_request_id
|
|
|
|
|
handler._record_request_outcome = record
|
|
|
|
|
|
|
|
|
|
response = asyncio.run(
|
|
|
|
|
handler.handle_passthrough(
|
|
|
|
|
_VertexStreamPassthroughRequest(),
|
|
|
|
|
"https://vertex.test",
|
|
|
|
|
"streamGenerateContent",
|
|
|
|
|
"vertex:google",
|
|
|
|
|
)
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
assert isinstance(response, StreamingResponse)
|
|
|
|
|
|
|
|
|
|
async def collect(): # noqa: ANN202
|
|
|
|
|
return [chunk async for chunk in response.body_iterator]
|
|
|
|
|
|
|
|
|
|
chunks = asyncio.run(collect())
|
|
|
|
|
|
|
|
|
|
assert len(chunks) == 2
|
|
|
|
|
assert chunks[0].startswith(b'data: {"candidates"')
|
|
|
|
|
assert b'"usageMetadata"' in chunks[1]
|
|
|
|
|
assert len(outcomes) == 1
|
|
|
|
|
outcome = outcomes[0]
|
|
|
|
|
assert outcome.provider == "vertex:google"
|
|
|
|
|
assert outcome.model == "gemini-2.0-flash"
|
|
|
|
|
assert outcome.optimized_tokens == 13
|
|
|
|
|
assert outcome.output_tokens == 5
|
|
|
|
|
assert outcome.cache_read_tokens == 2
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_stream_finalizer_records_vertex_provider_for_dashboard() -> None:
|
|
|
|
|
handler = object.__new__(HeadroomProxy)
|
|
|
|
|
handler.config = SimpleNamespace(log_full_messages=False)
|
|
|
|
|
outcomes = []
|
|
|
|
|
|
|
|
|
|
async def record(outcome): # noqa: ANN001, ANN202
|
|
|
|
|
outcomes.append(outcome)
|
|
|
|
|
|
|
|
|
|
handler._record_request_outcome = record
|
|
|
|
|
|
|
|
|
|
asyncio.run(
|
|
|
|
|
handler._finalize_stream_response(
|
|
|
|
|
body={"contents": [{"role": "user", "parts": [{"text": "hello"}]}]},
|
|
|
|
|
provider="gemini",
|
|
|
|
|
outcome_provider="vertex:google",
|
|
|
|
|
model="gemini-2.0-flash",
|
|
|
|
|
request_id="req_vertex_stream_final",
|
|
|
|
|
original_tokens=20,
|
|
|
|
|
optimized_tokens=12,
|
|
|
|
|
tokens_saved=8,
|
|
|
|
|
transforms_applied=["test-transform"],
|
|
|
|
|
optimization_latency=3.0,
|
|
|
|
|
stream_state={
|
|
|
|
|
"input_tokens": 12,
|
|
|
|
|
"output_tokens": 5,
|
|
|
|
|
"cache_read_input_tokens": 2,
|
|
|
|
|
"cache_creation_input_tokens": 0,
|
|
|
|
|
"cache_creation_ephemeral_5m_input_tokens": 0,
|
|
|
|
|
"cache_creation_ephemeral_1h_input_tokens": 0,
|
|
|
|
|
"total_bytes": 100,
|
|
|
|
|
"sse_buffer": bytearray(),
|
|
|
|
|
"ttfb_ms": 4.0,
|
|
|
|
|
},
|
|
|
|
|
start_time=0.0,
|
|
|
|
|
tags={"route": "vertex"},
|
|
|
|
|
)
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
assert len(outcomes) == 1
|
|
|
|
|
outcome = outcomes[0]
|
|
|
|
|
assert outcome.provider == "vertex:google"
|
|
|
|
|
assert outcome.model == "gemini-2.0-flash"
|
|
|
|
|
assert outcome.optimized_tokens == 12
|
|
|
|
|
assert outcome.output_tokens == 5
|
|
|
|
|
assert outcome.tokens_saved == 8
|
|
|
|
|
assert outcome.cache_read_tokens == 2
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_vertex_gemini_non_text_generate_records_dashboard_outcome() -> None:
|
|
|
|
|
handler = object.__new__(HeadroomProxy)
|
|
|
|
|
handler.memory_handler = None
|
|
|
|
|
handler.rate_limiter = None
|
|
|
|
|
outcomes = []
|
|
|
|
|
upstream_urls = []
|
|
|
|
|
|
|
|
|
|
async def next_request_id(): # noqa: ANN202
|
|
|
|
|
return "req_vertex_image"
|
|
|
|
|
|
|
|
|
|
async def record(outcome): # noqa: ANN001, ANN202
|
|
|
|
|
outcomes.append(outcome)
|
|
|
|
|
|
|
|
|
|
async def retry_request(method, url, headers, body): # noqa: ANN001, ANN202
|
|
|
|
|
upstream_urls.append(url)
|
|
|
|
|
request = httpx.Request(method, url, headers=headers)
|
|
|
|
|
return httpx.Response(
|
|
|
|
|
200,
|
|
|
|
|
request=request,
|
|
|
|
|
headers={"content-type": "application/json"},
|
|
|
|
|
json={
|
|
|
|
|
"usageMetadata": {
|
|
|
|
|
"promptTokenCount": 31,
|
|
|
|
|
"candidatesTokenCount": 4,
|
|
|
|
|
"cachedContentTokenCount": 6,
|
|
|
|
|
}
|
|
|
|
|
},
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
handler._next_request_id = next_request_id
|
|
|
|
|
handler._record_request_outcome = record
|
|
|
|
|
handler._retry_request = retry_request
|
|
|
|
|
|
|
|
|
|
response = asyncio.run(
|
|
|
|
|
handler.handle_gemini_generate_content(
|
|
|
|
|
_VertexGeminiImageRequest(),
|
|
|
|
|
"gemini-2.0-flash",
|
|
|
|
|
"https://vertex.test",
|
|
|
|
|
"vertex:google",
|
|
|
|
|
)
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
assert response.status_code == 200
|
|
|
|
|
assert upstream_urls == [
|
|
|
|
|
"https://vertex.test/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:generateContent"
|
|
|
|
|
]
|
|
|
|
|
assert response.headers["x-headroom-tokens-before"] == "31"
|
|
|
|
|
assert response.headers["x-headroom-tokens-after"] == "31"
|
|
|
|
|
assert response.headers["x-headroom-tokens-saved"] == "0"
|
|
|
|
|
assert len(outcomes) == 1
|
|
|
|
|
outcome = outcomes[0]
|
|
|
|
|
assert outcome.provider == "vertex:google"
|
|
|
|
|
assert outcome.model == "gemini-2.0-flash"
|
|
|
|
|
assert outcome.original_tokens == 31
|
|
|
|
|
assert outcome.optimized_tokens == 31
|
|
|
|
|
assert outcome.output_tokens == 4
|
|
|
|
|
assert outcome.cache_read_tokens == 6
|
|
|
|
|
assert outcome.num_messages == 1
|
|
|
|
|
|
|
|
|
|
|
2026-05-10 12:51:28 -04:00
|
|
|
def test_retry_request_retries_connect_timeout() -> None:
|
|
|
|
|
proxy = object.__new__(HeadroomProxy)
|
|
|
|
|
proxy.http_client = _RetryThenSuccessClient()
|
|
|
|
|
proxy.config = SimpleNamespace(
|
|
|
|
|
retry_enabled=True,
|
|
|
|
|
retry_max_attempts=2,
|
|
|
|
|
retry_base_delay_ms=0,
|
|
|
|
|
retry_max_delay_ms=0,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
response = asyncio.run(
|
|
|
|
|
proxy._retry_request(
|
|
|
|
|
"POST",
|
|
|
|
|
"https://api.openai.com/v1/responses",
|
|
|
|
|
{},
|
|
|
|
|
{"model": "gpt-5"},
|
|
|
|
|
)
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
assert response.status_code == 200
|
|
|
|
|
assert proxy.http_client.attempts == 2
|
|
|
|
|
|
|
|
|
|
|
2026-07-06 09:24:47 -04:00
|
|
|
def test_retry_request_returns_503_when_shutdown_interrupts_retry_sleep() -> None:
|
|
|
|
|
class _Always429Client:
|
|
|
|
|
def __init__(self) -> None:
|
|
|
|
|
self.attempts = 0
|
|
|
|
|
|
|
|
|
|
async def post(self, url, **kwargs): # type: ignore[no-untyped-def]
|
|
|
|
|
self.attempts += 1
|
|
|
|
|
return httpx.Response(
|
|
|
|
|
429,
|
|
|
|
|
request=httpx.Request("POST", url),
|
|
|
|
|
json={"error": {"message": "slow down"}},
|
|
|
|
|
headers={"retry-after": "30"},
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
proxy = object.__new__(HeadroomProxy)
|
|
|
|
|
proxy.http_client = _Always429Client()
|
|
|
|
|
proxy.config = SimpleNamespace(
|
|
|
|
|
retry_enabled=True,
|
|
|
|
|
retry_max_attempts=3,
|
|
|
|
|
retry_base_delay_ms=30000,
|
|
|
|
|
retry_max_delay_ms=30000,
|
|
|
|
|
)
|
|
|
|
|
proxy._shutdown_event = asyncio.Event()
|
|
|
|
|
proxy._shutdown_event.set()
|
|
|
|
|
|
|
|
|
|
response = asyncio.run(
|
|
|
|
|
proxy._retry_request(
|
|
|
|
|
"POST",
|
|
|
|
|
"https://api.anthropic.test/v1/messages",
|
|
|
|
|
{},
|
|
|
|
|
{"model": "claude-3-5-sonnet"},
|
|
|
|
|
)
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
assert response.status_code == 503
|
|
|
|
|
assert response.json() == {
|
|
|
|
|
"error": {
|
|
|
|
|
"type": "shutdown",
|
|
|
|
|
"message": "Proxy is shutting down; retry backoff cancelled.",
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
assert response.headers["retry-after"] == "0"
|
|
|
|
|
assert proxy.http_client.attempts == 1
|
|
|
|
|
|
|
|
|
|
|
2026-04-21 23:09:32 -05:00
|
|
|
def test_anthropic_tool_sort_and_context_append_helpers() -> None:
|
|
|
|
|
tools = [
|
|
|
|
|
{"type": "function", "function": {"name": "beta"}},
|
|
|
|
|
{"name": "alpha"},
|
|
|
|
|
{"type": "tool"},
|
|
|
|
|
]
|
2026-04-24 15:33:30 +02:00
|
|
|
|
2026-04-21 23:09:32 -05:00
|
|
|
sorted_tools = AnthropicHandlerMixin._sort_tools_deterministically(tools)
|
2026-04-24 15:33:30 +02:00
|
|
|
|
2026-04-21 23:09:32 -05:00
|
|
|
assert [AnthropicHandlerMixin._tool_sort_key(tool)[0] for tool in sorted_tools] == [
|
|
|
|
|
"alpha",
|
|
|
|
|
"beta",
|
|
|
|
|
"tool",
|
|
|
|
|
]
|
|
|
|
|
assert AnthropicHandlerMixin._sort_tools_deterministically(None) is None
|
fix: preserve anthropic passthrough tool order (#1427)
## Description
Preserves Anthropic `tools` order when Headroom is forwarding a
passthrough/no-optimize request. This fixes a Claude Code style
`tool_result` continuation failure against stricter Anthropic-compatible
upstreams that treat the client's original tool ordering as part of the
conversation state.
Closes #1417
## 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
- Preserve client-provided Anthropic `tools` order when `optimize=False`
or the request is explicitly in Headroom passthrough/bypass mode.
- Keep deterministic tool sorting for optimized requests where Headroom
may rewrite the body for cache stability.
- Avoid sorting batch-request tools before the no-optimize passthrough
branch.
- Add regression coverage for the Anthropic HTTP path to prove
no-optimize forwarding keeps `Read`, then `Bash` tool order.
- Update existing cache-stability and byte-faithful forwarding tests so
no-optimize/passthrough expects preserved client order while optimized
mode still proves deterministic sorting.
## Testing
- [x] Focused unit tests pass (`pytest` on touched proxy test files)
- [x] Linting passes (`ruff check` and `ruff format --check` on touched
files)
- [x] Type checking passes (`mypy headroom`)
- [x] New regression tests added
- [x] Manual testing performed
### Test Output
```text
$ uv run --no-project --python /opt/homebrew/opt/python@3.13/bin/python3.13 --with pytest --with pytest-asyncio --with anyio --with 'httpx[http2]' --with fastapi --with pydantic --with tiktoken --with click --with rich --with opentelemetry-api --with opentelemetry-sdk --with zstandard --with openai --with mcp --with uvicorn pytest tests/test_proxy_handler_helpers.py tests/test_anthropic_stage_timings.py tests/test_proxy_anthropic_cache_stability.py tests/test_proxy_byte_faithful_forwarding.py -q
============================= test session starts ==============================
platform darwin -- Python 3.13.11, pytest-9.1.1, pluggy-1.6.0
rootdir: /Users/vinaygupta/Desktop/git/headroom-fix-anthropic-tool-order
configfile: pyproject.toml
plugins: anyio-4.14.1, asyncio-1.4.0
asyncio: mode=Mode.AUTO, debug=False, asyncio_default_fixture_loop_scope=None, asyncio_default_test_loop_scope=function
collected 87 items
tests/test_proxy_handler_helpers.py .......................... [ 29%]
tests/test_anthropic_stage_timings.py .... [ 34%]
tests/test_proxy_anthropic_cache_stability.py ......................... [ 63%]
tests/test_proxy_byte_faithful_forwarding.py ........................... [ 94%]
..... [100%]
=============================== warnings summary ===============================
.../fastapi/testclient.py:1: StarletteDeprecationWarning: Using `httpx` with `starlette.testclient` is deprecated; install `httpx2` instead.
======================== 87 passed, 1 warning in 5.13s =========================
$ uv run --no-project --python /opt/homebrew/opt/python@3.13/bin/python3.13 --with ruff==0.15.17 ruff format --check headroom/proxy/handlers/anthropic.py tests/test_proxy_handler_helpers.py tests/test_anthropic_stage_timings.py tests/test_proxy_anthropic_cache_stability.py tests/test_proxy_byte_faithful_forwarding.py
5 files already formatted
$ uv run --no-project --python /opt/homebrew/opt/python@3.13/bin/python3.13 --with ruff==0.15.17 ruff check headroom/proxy/handlers/anthropic.py tests/test_proxy_handler_helpers.py tests/test_anthropic_stage_timings.py tests/test_proxy_anthropic_cache_stability.py tests/test_proxy_byte_faithful_forwarding.py
All checks passed!
```
## Real Behavior Proof
- Environment: macOS, Python 3.13.11, local fake Anthropic-compatible
upstream, local Headroom proxy launched with `--no-optimize --no-cache
--no-rate-limit --stateless`.
- Exact command / steps: ran a local reproduction harness that starts a
fake `/v1/messages` upstream and Headroom proxy, then sends a Claude
Code style two-turn flow: first assistant `Bash` `tool_use`, then user
`tool_result`.
- Observed result: after this patch, both direct and proxied flows
returned `200` for `first_tool_use` and `second_tool_result`. The fake
upstream log showed the proxied `tools` array remained `["Read",
"Bash"]` on both turns.
```text
DIRECT
first_tool_use: 200
second_tool_result: 200
PROXIED
first_tool_use: 200
second_tool_result: 200
UPSTREAM REQUEST LOG
proxied first turn tools: ["Read", "Bash"]
proxied tool_result turn tools: ["Read", "Bash"]
```
- Not tested: full `pytest`, full-repo `ruff check .`, `mypy headroom`,
or a live third-party Anthropic-compatible provider.
## 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
- [x] I have updated the CHANGELOG.md if applicable
## Screenshots (if applicable)
N/A
## Additional Notes
- This PR intentionally does not add documentation because it fixes
passthrough behavior rather than introducing a new user-facing option.
- The code-comment checklist item is left unchecked because the change
is covered by a small helper docstring and regression tests; no extra
inline comments seemed necessary.
- `CHANGELOG.md` is left unchanged because this is a narrowly scoped bug
fix.
- Local pytest collection for these proxy tests required a local
`headroom._core` extension symlink, which was removed before committing.
2026-06-30 08:38:51 -05:00
|
|
|
assert AnthropicHandlerMixin._tools_for_forwarding(tools, preserve_order=True) == tools
|
|
|
|
|
assert [
|
|
|
|
|
AnthropicHandlerMixin._tool_sort_key(tool)[0]
|
|
|
|
|
for tool in AnthropicHandlerMixin._tools_for_forwarding(tools, preserve_order=False) or []
|
|
|
|
|
] == [
|
|
|
|
|
"alpha",
|
|
|
|
|
"beta",
|
|
|
|
|
"tool",
|
|
|
|
|
]
|
2026-04-21 23:09:32 -05:00
|
|
|
assert (
|
|
|
|
|
AnthropicHandlerMixin._append_context_to_latest_non_frozen_user_turn(
|
|
|
|
|
[], "ctx", frozen_message_count=0
|
|
|
|
|
)
|
|
|
|
|
== []
|
|
|
|
|
)
|
|
|
|
|
assert AnthropicHandlerMixin._append_context_to_latest_non_frozen_user_turn(
|
|
|
|
|
[{"role": "user", "content": "hello"}],
|
|
|
|
|
"ctx",
|
|
|
|
|
frozen_message_count=0,
|
|
|
|
|
) == [{"role": "user", "content": "hello\n\nctx"}]
|
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.
2026-05-02 09:02:10 -07:00
|
|
|
# PR-A2 semantics: list-content user messages get the context appended
|
|
|
|
|
# to the first text block (live-zone-tail injection).
|
2026-04-21 23:09:32 -05:00
|
|
|
assert AnthropicHandlerMixin._append_context_to_latest_non_frozen_user_turn(
|
|
|
|
|
[{"role": "user", "content": [{"type": "text", "text": "hello"}]}],
|
|
|
|
|
"ctx",
|
|
|
|
|
frozen_message_count=0,
|
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.
2026-05-02 09:02:10 -07:00
|
|
|
) == [{"role": "user", "content": [{"type": "text", "text": "hello\n\nctx"}]}]
|
2026-04-24 15:33:30 +02:00
|
|
|
|
|
|
|
|
|
2026-04-21 23:09:32 -05:00
|
|
|
def test_anthropic_image_compression_helper_only_rewrites_latest_eligible_turn() -> None:
|
|
|
|
|
image_message = {
|
|
|
|
|
"role": "user",
|
|
|
|
|
"content": [{"type": "image", "source": {"type": "base64", "data": "abc"}}],
|
|
|
|
|
}
|
|
|
|
|
compressed = {
|
|
|
|
|
"role": "user",
|
|
|
|
|
"content": [{"type": "image", "source": {"type": "base64", "data": "xyz"}}],
|
|
|
|
|
}
|
2026-04-24 15:33:30 +02:00
|
|
|
|
2026-04-21 23:09:32 -05:00
|
|
|
assert (
|
|
|
|
|
AnthropicHandlerMixin._compress_latest_user_turn_images_cache_safe(
|
|
|
|
|
[],
|
|
|
|
|
frozen_message_count=0,
|
|
|
|
|
compressor=_ImageCompressor(compressed),
|
|
|
|
|
)
|
|
|
|
|
== []
|
|
|
|
|
)
|
|
|
|
|
assert AnthropicHandlerMixin._compress_latest_user_turn_images_cache_safe(
|
|
|
|
|
[image_message],
|
|
|
|
|
frozen_message_count=1,
|
|
|
|
|
compressor=_ImageCompressor(compressed),
|
|
|
|
|
) == [image_message]
|
|
|
|
|
assert AnthropicHandlerMixin._compress_latest_user_turn_images_cache_safe(
|
|
|
|
|
[{"role": "assistant", "content": image_message["content"]}],
|
|
|
|
|
frozen_message_count=0,
|
|
|
|
|
compressor=_ImageCompressor(compressed),
|
|
|
|
|
) == [{"role": "assistant", "content": image_message["content"]}]
|
|
|
|
|
assert AnthropicHandlerMixin._compress_latest_user_turn_images_cache_safe(
|
|
|
|
|
[{"role": "user", "content": "no-image"}],
|
|
|
|
|
frozen_message_count=0,
|
|
|
|
|
compressor=_ImageCompressor(compressed),
|
|
|
|
|
) == [{"role": "user", "content": "no-image"}]
|
|
|
|
|
assert AnthropicHandlerMixin._compress_latest_user_turn_images_cache_safe(
|
|
|
|
|
[image_message],
|
|
|
|
|
frozen_message_count=0,
|
|
|
|
|
compressor=_ImageCompressor(image_message),
|
|
|
|
|
) == [image_message]
|
|
|
|
|
assert AnthropicHandlerMixin._compress_latest_user_turn_images_cache_safe(
|
|
|
|
|
[image_message],
|
|
|
|
|
frozen_message_count=0,
|
|
|
|
|
compressor=_ImageCompressor(compressed),
|
|
|
|
|
) == [compressed]
|
2026-04-24 15:33:30 +02:00
|
|
|
|
|
|
|
|
|
2026-04-23 20:41:59 -04:00
|
|
|
def test_proxy_helper_creates_fresh_image_compressors(monkeypatch) -> None:
|
|
|
|
|
from headroom.proxy import helpers
|
|
|
|
|
|
|
|
|
|
monkeypatch.setattr(helpers, "_image_compressor_available", None)
|
|
|
|
|
_FreshCompressor.instances = 0
|
|
|
|
|
|
|
|
|
|
with patch("headroom.image.ImageCompressor", _FreshCompressor):
|
|
|
|
|
first = helpers._get_image_compressor()
|
|
|
|
|
second = helpers._get_image_compressor()
|
|
|
|
|
|
|
|
|
|
assert isinstance(first, _FreshCompressor)
|
|
|
|
|
assert isinstance(second, _FreshCompressor)
|
|
|
|
|
assert first is not second
|
|
|
|
|
assert _FreshCompressor.instances == 2
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_proxy_helper_caches_image_stack_import_failure(monkeypatch) -> None:
|
|
|
|
|
from headroom.proxy import helpers
|
|
|
|
|
|
|
|
|
|
real_import = builtins.__import__
|
|
|
|
|
calls = 0
|
|
|
|
|
|
|
|
|
|
def fake_import(name, *args, **kwargs): # noqa: ANN001, ANN202
|
|
|
|
|
nonlocal calls
|
|
|
|
|
if name == "headroom.image":
|
|
|
|
|
calls += 1
|
|
|
|
|
raise ImportError("image extras unavailable")
|
|
|
|
|
return real_import(name, *args, **kwargs)
|
|
|
|
|
|
|
|
|
|
monkeypatch.setattr(helpers, "_image_compressor_available", None)
|
|
|
|
|
monkeypatch.setattr(builtins, "__import__", fake_import)
|
|
|
|
|
|
|
|
|
|
assert helpers._get_image_compressor() is None
|
|
|
|
|
assert helpers._get_image_compressor() is None
|
|
|
|
|
assert calls == 1
|
|
|
|
|
assert helpers._image_compressor_available is False
|
|
|
|
|
|
|
|
|
|
|
2026-04-21 23:09:32 -05:00
|
|
|
def test_anthropic_cache_delta_helpers_cover_string_list_and_role_mismatch() -> None:
|
|
|
|
|
previous_original = [{"role": "user", "content": "hello"}]
|
|
|
|
|
previous_forwarded = [{"role": "user", "content": "HELLO"}]
|
2026-04-24 15:33:30 +02:00
|
|
|
|
2026-04-21 23:09:32 -05:00
|
|
|
assert AnthropicHandlerMixin._extract_cache_stable_delta(
|
|
|
|
|
[{"role": "user", "content": "hello"}, {"role": "assistant", "content": "next"}],
|
|
|
|
|
previous_original,
|
|
|
|
|
previous_forwarded,
|
|
|
|
|
) == (previous_forwarded, [{"role": "assistant", "content": "next"}])
|
|
|
|
|
assert (
|
|
|
|
|
AnthropicHandlerMixin._extract_cache_stable_delta(
|
|
|
|
|
[{"role": "assistant", "content": "hello"}],
|
|
|
|
|
previous_original,
|
|
|
|
|
previous_forwarded,
|
|
|
|
|
)
|
|
|
|
|
is None
|
|
|
|
|
)
|
2026-04-24 15:33:30 +02:00
|
|
|
|
2026-04-21 23:09:32 -05:00
|
|
|
string_suffix = AnthropicHandlerMixin._extract_cache_stable_last_message_suffix(
|
|
|
|
|
[{"role": "user", "content": "hello world"}],
|
|
|
|
|
previous_original,
|
|
|
|
|
previous_forwarded,
|
|
|
|
|
)
|
|
|
|
|
assert string_suffix == ([], previous_forwarded[0], [{"role": "user", "content": " world"}])
|
2026-04-24 15:33:30 +02:00
|
|
|
|
2026-04-21 23:09:32 -05:00
|
|
|
list_suffix = AnthropicHandlerMixin._extract_cache_stable_last_message_suffix(
|
|
|
|
|
[
|
|
|
|
|
{
|
|
|
|
|
"role": "user",
|
|
|
|
|
"content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}],
|
|
|
|
|
}
|
|
|
|
|
],
|
|
|
|
|
[{"role": "user", "content": [{"type": "text", "text": "a"}]}],
|
|
|
|
|
[{"role": "user", "content": [{"type": "text", "text": "A"}]}],
|
|
|
|
|
)
|
|
|
|
|
assert list_suffix == (
|
|
|
|
|
[],
|
|
|
|
|
{"role": "user", "content": [{"type": "text", "text": "A"}]},
|
|
|
|
|
[{"role": "user", "content": [{"type": "text", "text": "b"}]}],
|
|
|
|
|
)
|
2026-04-24 15:33:30 +02:00
|
|
|
|
2026-04-21 23:09:32 -05:00
|
|
|
assert AnthropicHandlerMixin._merge_appended_message_delta(
|
|
|
|
|
{"role": "user", "content": "HELLO"},
|
|
|
|
|
{"role": "user", "content": " world"},
|
|
|
|
|
) == {"role": "user", "content": "HELLO world"}
|
|
|
|
|
assert AnthropicHandlerMixin._merge_appended_message_delta(
|
|
|
|
|
{"role": "user", "content": [{"type": "text", "text": "A"}]},
|
|
|
|
|
{"role": "user", "content": [{"type": "text", "text": "b"}]},
|
|
|
|
|
) == {"role": "user", "content": [{"type": "text", "text": "A"}, {"type": "text", "text": "b"}]}
|
|
|
|
|
assert (
|
|
|
|
|
AnthropicHandlerMixin._merge_appended_message_delta(
|
|
|
|
|
{"role": "user", "content": "A"},
|
|
|
|
|
{"role": "assistant", "content": "B"},
|
|
|
|
|
)
|
|
|
|
|
is None
|
|
|
|
|
)
|
2026-04-24 15:33:30 +02:00
|
|
|
|
|
|
|
|
|
2026-04-21 23:09:32 -05:00
|
|
|
def test_anthropic_assistant_message_helper_requires_assistant_role() -> None:
|
|
|
|
|
assert AnthropicHandlerMixin._assistant_message_from_response_json(None) is None
|
|
|
|
|
assert AnthropicHandlerMixin._assistant_message_from_response_json({"role": "user"}) is None
|
|
|
|
|
assert AnthropicHandlerMixin._assistant_message_from_response_json(
|
|
|
|
|
{"role": "assistant", "content": [{"type": "text", "text": "ok"}]}
|
|
|
|
|
) == {"role": "assistant", "content": [{"type": "text", "text": "ok"}]}
|
fix(ccr): scope proactive expansion by workspace (cross-project leak)
Closes the cross-project context leak Jocelyn reported 2026-05-26:
working on a Ruby/Rails project (daphni-rails), an unrelated Python
file (an Ollama inference provider from project `tamag0`) was being
injected into context as "Proactive Context Expansion - relevant to
your query". Two completely different projects, two different
languages, two different working directories — but the same proxy
process was serving both, and the in-memory ContextTracker had no
workspace identity to filter on.
Root cause
----------
`self.ccr_context_tracker` is one instance per proxy process. Every
session, every project, every user shared the same `_contexts` dict.
`track_compression()` stored sample content with no provenance key;
`analyze_query()` ran lexical keyword overlap across the full dict
without filtering. Within the 5-minute age window, surface-level
token matches ("provider", "session", "oauth", generic code/test
structure) scored above the 0.3 relevance threshold, recommendations
came back, and execute_expansions() injected the full original
content into a foreign session.
Refuted: this is NOT a race condition (joce's hypothesis). It
reproduces single-threaded, one-request-at-a-time. Plain shared
mutable state.
Fix
---
Add a required `workspace_key` to the tracker API and filter on it
inside `analyze_query`:
1. `CompressedContext` gets a `workspace_key: str` field.
2. `track_compression(..., workspace_key=...)` is now keyword-only,
no default — fail-loud on missing.
3. `analyze_query(..., workspace_key=...)` is also keyword-only; an
empty workspace_key short-circuits to `[]` (fail-closed per
`feedback_no_silent_fallbacks`).
4. The loop at `analyze_query` skips any entry whose workspace_key
differs from the request's.
In the Anthropic proxy handler:
5. New `_resolve_ccr_workspace(request, body)` static helper uses the
memory subsystem's `ProjectResolver` so CCR and memory agree on
project identity. Tier order: x-headroom-project-id →
x-headroom-cwd → CLI override → cwd: line in system prompt.
6. Both track and analyze sites gate on `ccr_workspace_key` being
non-empty — turning off proactive expansion entirely when project
identity can't be resolved is the safest default (it's an
optimization, not correctness).
7. `format_expansions_for_context(expansions, workspace_label=...)`
was already wired (GH #462 Fix C); the call site now passes the
label so the injected block declares its provenance, symmetric
with the memory injection header.
Affected population
-------------------
- Default mode (no `--cache`): bug fixed.
- Cache mode: was never affected — proactive expansion short-
circuits in cache mode to preserve prefix stability.
Tests
-----
- 6 new workspace-scoping tests in `test_ccr_context_tracker.py`:
same-workspace match still works, cross-workspace silently
filtered, empty workspace_key fail-closes, two workspaces each
see only their own, workspace_label propagates to formatter, LRU
cross-workspace doesn't leak even with full tracker.
- 6 new `_resolve_ccr_workspace` resolver tests in
`test_proxy_handler_helpers.py`: explicit project-id wins, cwd
header → key+label, two cwds get distinct keys, no-signal
fail-closed, system-prompt cwd: fallback, malformed request
fail-closed.
- 32 existing tracker tests updated to pass `workspace_key="ws-test"`.
- 55/55 tests pass; ci-precheck green.
Defense-in-depth follow-up
--------------------------
The compression_store itself (`headroom/cache/compression_store.py`)
also lacks workspace scoping — a CCR `headroom_retrieve` call from
Project B for a hash created by Project A would succeed. The
practical attack surface is closed by this PR (hashes only reach
Project B's model via proactive expansion, now gated), but
defense-in-depth hardening of the store is worth a separate PR.
Filed as task #44.
2026-05-26 13:23:51 -07:00
|
|
|
|
|
|
|
|
|
|
|
|
|
# ============================================================================
|
|
|
|
|
# CCR workspace resolution (cross-project leak fix, 2026-05-26).
|
|
|
|
|
#
|
|
|
|
|
# These tests pin the `_resolve_ccr_workspace` static helper that the
|
|
|
|
|
# anthropic handler uses to scope the proactive-expansion cache by
|
|
|
|
|
# project identity. The resolver shares its tier order with the memory
|
|
|
|
|
# subsystem's ProjectResolver: x-headroom-project-id → x-headroom-cwd →
|
|
|
|
|
# system-prompt `cwd:` line. Returns `("", None)` on no signal — the
|
|
|
|
|
# fail-closed signal that callers gate on.
|
|
|
|
|
# ============================================================================
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _fake_request(headers: dict[str, str]) -> SimpleNamespace:
|
|
|
|
|
"""Minimal Starlette/FastAPI-shaped request object for resolver tests."""
|
|
|
|
|
return SimpleNamespace(headers=headers)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_resolve_ccr_workspace_explicit_project_id_wins() -> None:
|
|
|
|
|
"""x-headroom-project-id is the highest-priority signal."""
|
|
|
|
|
request = _fake_request({"x-headroom-project-id": "my-cool-project"})
|
|
|
|
|
body = {}
|
|
|
|
|
key, label = AnthropicHandlerMixin._resolve_ccr_workspace(request, body)
|
|
|
|
|
assert key == "my-cool-project"
|
|
|
|
|
assert label == "my-cool-project"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_resolve_ccr_workspace_cwd_header() -> None:
|
|
|
|
|
"""x-headroom-cwd produces a stable per-cwd key + basename label."""
|
|
|
|
|
request = _fake_request({"x-headroom-cwd": "/home/user/code/daphni-rails"})
|
|
|
|
|
body = {}
|
|
|
|
|
key, label = AnthropicHandlerMixin._resolve_ccr_workspace(request, body)
|
|
|
|
|
# Key format: "{basename}-{sha256[:16]}" — stable per absolute cwd.
|
|
|
|
|
assert key.startswith("daphni-rails-")
|
|
|
|
|
assert len(key) >= len("daphni-rails-") + 16
|
|
|
|
|
assert label == "daphni-rails"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_resolve_ccr_workspace_two_cwds_get_distinct_keys() -> None:
|
|
|
|
|
"""Two different cwds produce different workspace keys (cross-leak prevention)."""
|
|
|
|
|
key_a, _ = AnthropicHandlerMixin._resolve_ccr_workspace(
|
|
|
|
|
_fake_request({"x-headroom-cwd": "/home/user/code/daphni-rails"}), {}
|
|
|
|
|
)
|
|
|
|
|
key_b, _ = AnthropicHandlerMixin._resolve_ccr_workspace(
|
|
|
|
|
_fake_request({"x-headroom-cwd": "/home/user/code/tamag0"}), {}
|
|
|
|
|
)
|
|
|
|
|
assert key_a != key_b, "different cwds must yield different workspace keys"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_resolve_ccr_workspace_no_signal_returns_empty() -> None:
|
|
|
|
|
"""No project-id, no cwd header, no system prompt → fail-closed signal."""
|
|
|
|
|
request = _fake_request({})
|
|
|
|
|
body = {}
|
|
|
|
|
key, label = AnthropicHandlerMixin._resolve_ccr_workspace(request, body)
|
|
|
|
|
assert key == ""
|
|
|
|
|
assert label is None
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_resolve_ccr_workspace_system_prompt_cwd_fallback() -> None:
|
|
|
|
|
"""System prompt with `cwd:` line is the lowest-tier fallback."""
|
|
|
|
|
request = _fake_request({})
|
|
|
|
|
body = {
|
|
|
|
|
"system": [{"type": "text", "text": "You are helpful.\ncwd: /home/u/code/my-project\nGo."}]
|
|
|
|
|
}
|
|
|
|
|
key, label = AnthropicHandlerMixin._resolve_ccr_workspace(request, body)
|
|
|
|
|
# The label is the basename of the cwd extracted from the prompt.
|
|
|
|
|
assert label == "my-project"
|
|
|
|
|
assert key.startswith("my-project-")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_resolve_ccr_workspace_malformed_request_returns_empty() -> None:
|
|
|
|
|
"""A request whose headers attribute can't be dict()-ed fails closed, not crashes."""
|
|
|
|
|
|
|
|
|
|
class _BrokenHeaders:
|
|
|
|
|
def __iter__(self):
|
|
|
|
|
raise RuntimeError("boom")
|
|
|
|
|
|
|
|
|
|
request = SimpleNamespace(headers=_BrokenHeaders())
|
|
|
|
|
body = {}
|
|
|
|
|
# The helper catches the exception, logs it, and returns the fail-
|
|
|
|
|
# closed sentinel ("", None). Critically, it does NOT raise — the
|
|
|
|
|
# proxy must continue serving the request even if CCR scoping fails.
|
|
|
|
|
key, label = AnthropicHandlerMixin._resolve_ccr_workspace(request, body)
|
|
|
|
|
assert key == ""
|
|
|
|
|
assert label is None
|
fix(proxy): freeze must forward cached (compressed) prefix byte-identical — stop token-mode cache busting (#1850)
The freeze path (both providers) emits the agent's ORIGINAL bytes for a
frozen message, but the provider cached whatever we FORWARDED last turn
(the compressed form). Forwarding original then mismatches the cached
prefix and busts it from that point — re-creating the whole suffix.
Measured on a real SWE-bench run: 100% of attributed misses were
prefix_change, ~56% of ALL cache-writes were bust-induced (2.8M tokens),
driving cache_create +150% and cost +41% vs baseline.
Cache mode already avoided this via _extract_cache_stable_delta (replay
the previously-forwarded prefix, compress only the delta). Token mode
called apply(frozen_count) directly, which forwards original for the
frozen region.
Fix: add a shared, provider-agnostic overlay_cached_prefix() that
replays the previously-forwarded (cached, compressed) prefix
byte-identical, append-only guarded and idempotent, and apply it in BOTH
the Anthropic and OpenAI handlers right before forwarding. This makes
freezing byte-identical in every mode, so the only remaining difference
between "token" and "cache" mode is how large a mutable
(still-compressible) tail each leaves — not whether the frozen prefix
busts the cache.
Tests:
- test_cache_prefix_overlay.py: the helper (replay, append-only guard,
idempotence).
- test_cross_turn_cache_safety.py: the invariant that was missing —
drive the REAL tracker + freeze + overlay over multiple append-only
turns against a simulated provider prefix cache and assert the forwarded
prefix stays byte-identical turn-over-turn. Load-bearing: it fails
(detects the bust) without the overlay.
## Description
<!-- Briefly explain the change and why it is needed. -->
Closes #
## Type of Change
- [ ] 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
-
## Testing
<!-- Check what you actually ran, then paste the real command output
below. -->
- [ ] Unit tests pass (`pytest`)
- [ ] Linting passes (`ruff check .`)
- [ ] Type checking passes (`mypy headroom`)
- [ ] New tests added for new functionality
- [ ] Manual testing performed
### Test Output
```text
# Paste relevant command output or artifact links here
```
## Real Behavior Proof
- Environment:
- Exact command / steps:
- Observed result:
- Not tested:
## Review Readiness
- [ ] I have performed a self-review
- [ ] This PR is ready for human review
## Checklist
- [ ] My code follows the project's style guidelines
- [ ] I have performed a self-review of my code
- [ ] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [ ] My changes generate no new warnings
- [ ] 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
- [ ] I have updated the CHANGELOG.md if applicable
## Screenshots (if applicable)
Add screenshots to help explain your changes.
## Additional Notes
<!-- Mention any N/A checklist items, tradeoffs, follow-ups, or
maintainer context. -->
2026-07-06 14:54:39 -07:00
|
|
|
|
|
|
|
|
|
|
|
|
|
class TestHasNewCcrMarkers:
|
|
|
|
|
"""#1850: replayed (overlay) markers must not count as new-this-turn.
|
|
|
|
|
|
|
|
|
|
``overlay_cached_prefix`` replays the previously-forwarded compressed prefix
|
|
|
|
|
byte-identical to keep the messages cache warm — which reintroduces its old
|
|
|
|
|
``hash=…`` markers. If those replayed markers counted as "new", the handler
|
|
|
|
|
would re-inject the retrieve tool every frozen turn and bust the *tools*
|
|
|
|
|
cache. ``has_new_ccr_markers`` filters them out.
|
|
|
|
|
"""
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
|
def _hashes(*contents: str) -> list[str]:
|
|
|
|
|
from headroom.ccr.tool_injection import CCRToolInjector
|
|
|
|
|
|
|
|
|
|
inj = CCRToolInjector(
|
|
|
|
|
provider="anthropic", inject_tool=False, inject_system_instructions=False
|
|
|
|
|
)
|
|
|
|
|
inj.scan_for_markers([{"role": "user", "content": c} for c in contents])
|
|
|
|
|
return inj.detected_hashes
|
|
|
|
|
|
|
|
|
|
def test_replayed_markers_are_not_new(self):
|
|
|
|
|
from headroom.proxy.helpers import has_new_ccr_markers
|
|
|
|
|
|
|
|
|
|
marker = "[100 items compressed to 10. Retrieve more: hash=abc123def456abc123def456]"
|
|
|
|
|
current = self._hashes(marker)
|
|
|
|
|
assert current, "sanity: the marker must be detected"
|
|
|
|
|
# Every marker was already in what we forwarded last turn → nothing new.
|
|
|
|
|
assert (
|
|
|
|
|
has_new_ccr_markers(
|
|
|
|
|
current_detected_hashes=current,
|
|
|
|
|
previous_forwarded_messages=[{"role": "user", "content": marker}],
|
|
|
|
|
provider="anthropic",
|
|
|
|
|
)
|
|
|
|
|
is False
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
def test_genuinely_new_marker_is_detected(self):
|
|
|
|
|
from headroom.proxy.helpers import has_new_ccr_markers
|
|
|
|
|
|
|
|
|
|
old = "[100 items compressed to 10. Retrieve more: hash=abc123def456abc123def456]"
|
|
|
|
|
new = "[50 items compressed to 5. Retrieve more: hash=deadbeefdeadbeefdeadbeef]"
|
|
|
|
|
current = self._hashes(old, new)
|
|
|
|
|
# Only `old` was forwarded before; `new` is fresh → override must fire.
|
|
|
|
|
assert (
|
|
|
|
|
has_new_ccr_markers(
|
|
|
|
|
current_detected_hashes=current,
|
|
|
|
|
previous_forwarded_messages=[{"role": "user", "content": old}],
|
|
|
|
|
provider="anthropic",
|
|
|
|
|
)
|
|
|
|
|
is True
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
def test_no_previous_forward_means_all_new(self):
|
|
|
|
|
from headroom.proxy.helpers import has_new_ccr_markers
|
|
|
|
|
|
|
|
|
|
marker = "[100 items compressed to 10. Retrieve more: hash=abc123def456abc123def456]"
|
|
|
|
|
assert (
|
|
|
|
|
has_new_ccr_markers(
|
|
|
|
|
current_detected_hashes=self._hashes(marker),
|
|
|
|
|
previous_forwarded_messages=None,
|
|
|
|
|
provider="anthropic",
|
|
|
|
|
)
|
|
|
|
|
is True
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
def test_no_markers_means_nothing_new(self):
|
|
|
|
|
from headroom.proxy.helpers import has_new_ccr_markers
|
|
|
|
|
|
|
|
|
|
assert (
|
|
|
|
|
has_new_ccr_markers(
|
|
|
|
|
current_detected_hashes=[],
|
|
|
|
|
previous_forwarded_messages=None,
|
|
|
|
|
provider="anthropic",
|
|
|
|
|
)
|
|
|
|
|
is False
|
|
|
|
|
)
|