headroom/tests/test_proxy_handler_helpers.py

Ignoring revisions in .git-blame-ignore-revs. Click here to bypass and see the normal blame view.

958 lines
32 KiB
Python
Raw Normal View History

from __future__ import annotations
import asyncio
import base64
import builtins
import json
from types import SimpleNamespace
from unittest.mock import patch
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
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
)
from headroom.proxy.helpers import _headroom_bypass_enabled
from headroom.proxy.server import HeadroomProxy
def _jwt(payload: object) -> str:
header = {"alg": "none", "typ": "JWT"}
def encode(part: object) -> str:
raw = json.dumps(part, separators=(",", ":")).encode("utf-8")
return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=")
return f"{encode(header)}.{encode(payload)}."
class _ImageCompressor:
def __init__(self, compressed_message):
self._compressed_message = compressed_message
def compress(self, messages, provider): # noqa: ANN001, ANN201
assert provider == "anthropic"
return [self._compressed_message]
class _FreshCompressor:
instances = 0
def __init__(self):
type(self).instances += 1
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""
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',
]
),
)
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
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
request = httpx.Request("POST", url, headers=headers, content=content)
return httpx.Response(200, request=request, content=b"{}")
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
)
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
)
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
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
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
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
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
fix(proxy): cancel retry backoff on shutdown (#1834) ## Description During proxy shutdown, an in-flight retrying request can currently stay asleep inside `_retry_request()` and keep the client socket hanging until the retry timer expires or an external supervisor kills the process. This wires retry backoff to a proxy-scoped shutdown event so shutdown interrupts those waits immediately and returns a clear `503` response instead of leaving the request stalled. Closes #1821. ## 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 - Added a proxy-scoped shutdown event in `headroom/proxy/server.py`. - Cleared that event at startup and set it at shutdown before teardown proceeds. - Replaced both retry-backoff sleeps with a helper that wakes on either timeout or shutdown. - Returned a shutdown `503` with `retry-after: 0` when shutdown interrupts retry backoff. - Stopped the shutdown interruption logs from falling back to the raw upstream URL when no safe path string is available. - Added focused regressions for retry-backoff interruption and shutdown event signaling. - Updated the existing Retry-After tests to observe the new shutdown-aware wait helper instead of the old raw sleep hook. ## Testing - [x] Unit tests pass (`uv run pytest tests/test_proxy_handler_helpers.py tests/test_proxy_pipeline_lifecycle.py -q`) - [x] Unit tests pass (`uv run pytest tests/test_proxy_retry_429.py -q`) - [x] Linting passes (`uv run ruff check headroom/proxy/server.py tests/test_proxy_handler_helpers.py tests/test_proxy_retry_429.py tests/test_proxy_pipeline_lifecycle.py`) - [ ] Type checking passes (`uv run mypy headroom`) - [x] New tests added for new functionality when applicable - [ ] Manual testing performed ### Test Output ```text uv run pytest tests/test_proxy_handler_helpers.py tests/test_proxy_pipeline_lifecycle.py -q 32 passed, 1 warning in 13.05s uv run pytest tests/test_proxy_retry_429.py -q 10 passed, 1 warning in 1.12s uv run ruff check headroom/proxy/server.py tests/test_proxy_handler_helpers.py tests/test_proxy_retry_429.py tests/test_proxy_pipeline_lifecycle.py All checks passed! ``` ## Real Behavior Proof - Environment: Windows, project `uv` environment, focused proxy retry and shutdown regressions. - Exact command / steps: copy the updated shutdown regression files into a detached `origin/main` worktree and run `tests/test_proxy_handler_helpers.py` plus `tests/test_proxy_pipeline_lifecycle.py`, then rerun those files on this branch and separately rerun `tests/test_proxy_retry_429.py` after updating the existing Retry-After tests to patch the shutdown-aware wait helper. - Observed result: base fails because retry backoff still returns the original `429` and `shutdown()` leaves the retry event unset; head passes the focused file, preserves the existing Retry-After assertions, and returns a shutdown `503` with `retry-after: 0` while signaling retry waiters during shutdown. - Not tested: live systemd-managed shutdown on Linux or a full VS Code / Claude Code 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 - [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 ## Additional Notes This is intentionally scoped to retry backoff during shutdown. It does not try to cancel unrelated in-flight request work or change the broader retry policy outside shutdown.
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
def test_anthropic_tool_sort_and_context_append_helpers() -> None:
tools = [
{"type": "function", "function": {"name": "beta"}},
{"name": "alpha"},
{"type": "tool"},
]
sorted_tools = AnthropicHandlerMixin._sort_tools_deterministically(tools)
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",
]
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).
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"}]}]
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"}}],
}
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]
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
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"}]
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
)
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"}])
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"}]}],
)
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
)
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
)