akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(proxy): compress AWS Bedrock InvokeModel requests via configurable upstream (#720)

Browse source 
## Description

Clients that speak **Bedrock to a local gateway** can't get proxy-level
compression. Claude Code launched with `CLAUDE_CODE_USE_BEDROCK=1` (and
any AWS SDK pointed at a custom endpoint) POSTs
`/model/{id}/invoke[-with-response-stream]` to
`AWS_ENDPOINT_URL_BEDROCK_RUNTIME`, never `/v1/messages`. Those requests
fell through the catch-all and were forwarded **verbatim — no
compression**.

`--backend bedrock` is the opposite direction: it accepts Anthropic
input and re-signs to AWS. It can't accept Bedrock-format input or
forward to a custom upstream. So the "client speaks Bedrock → local
re-signing gateway → AWS" topology (internal gateways, LiteLLM,
LocalStack; see #510) got nothing.

This adds a Bedrock InvokeModel passthrough that compresses the request
body with the **same** `anthropic_pipeline` used for `/v1/messages` —
the Bedrock InvokeModel body for Anthropic models *is* the Anthropic
Messages shape (`{anthropic_version, system, messages, max_tokens, …}`,
model in the URL), so there's no translation and no new compression
logic. The routes register **only** when `--bedrock-api-url` is set, so
default behavior is completely unchanged.

**Limitation (important):** rewriting the body invalidates the caller's
**SigV4** signature (it covers a hash of the body). Point
`--bedrock-api-url` at a gateway that re-signs or doesn't verify the
inbound signature (an internal gateway, LiteLLM, LocalStack, a corporate
Bedrock proxy) — **never raw AWS**, which would 403. For direct-to-AWS
compression, use `--backend bedrock` (which re-signs). The two are
complementary. This is documented in the flag help, the handler
docstring, the proxy docs, and the CHANGELOG.

Closes #734. Refs #510 (the Bedrock slice of the provider-agnostic
umbrella).

## Type of Change

- [ ] Bug fix (non-breaking change that fixes an issue)
- [x] 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

- New `--bedrock-api-url` flag (env: `BEDROCK_TARGET_API_URL`). When
set, registers `POST /model/{id}/invoke` and `POST
/model/{id}/invoke-with-response-stream`.
- `BedrockHandlerMixin` compresses the request body via the existing
`anthropic_pipeline`, then forwards to the configured upstream,
preserving path/query.
- Responses forwarded byte-faithfully (non-streaming JSON and the
streaming AWS event-stream alike — neither is parsed or mutated, since
all compression is request-side).
- `{model_id:path}` captures inference-profile ids with
dots/colons/slashes (e.g.
`us.anthropic.claude-sonnet-4-5-20250929-v1:0`).
- Fail-open: a malformed body or compression error forwards verbatim
rather than erroring.
- Routes register only when the flag is set — default behavior
unchanged.
- Files: `headroom/proxy/handlers/bedrock.py` (new —
`BedrockHandlerMixin`); `headroom/providers/proxy_routes.py` (gated
route registration); `headroom/cli/proxy.py`,
`headroom/proxy/server.py`, `headroom/proxy/models.py` (flag + config
wiring); `docs/content/docs/proxy.mdx`, `CHANGELOG.md` (docs).

## Testing

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

### Test Output

```text
$ pytest tests/test_proxy/test_bedrock_passthrough.py -q
.............. [100%]
14 passed in 12.46s
```

`tests/test_proxy/test_bedrock_passthrough.py` (14 tests) covers: route
gating (absent unless configured), body compression, non-message fields
preserved, inference-profile id capture + re-encoding, byte-faithful
streaming, fail-open on malformed body and on pipeline exceptions,
bypass when `optimize=False` and via the `x-headroom-bypass` header,
upstream connect failure surfacing as a 502, the content-length
regression, outcome recorded with `provider="bedrock"`, and
`BEDROCK_TARGET_API_URL` env wiring. `ruff check`/`format` clean.

## Real Behavior Proof

- Environment: macOS, Python 3.12; forked proxy on `:8788` with
`--bedrock-api-url` pointed at a local re-signing Bedrock gateway;
provider Anthropic Claude on Bedrock.
- Exact command / steps: `headroom proxy --port 8788 --bedrock-api-url
http://127.0.0.1:<gateway>`, then `curl -X POST
http://127.0.0.1:8788/model/claude-haiku-4-5/invoke --data
@bedrock_invoke.json` (a ~52k-token conversation with a large assistant
turn).
- Observed result: valid Claude response returned and the gateway
received the compressed body — proxy `/stats` reports `52,095 → 3,979
tokens` (92.4%, 48,116 removed), and the gateway's reported
`input_tokens: 3709` confirms the compressed body reached the model.
- Not tested: raw direct-to-AWS (out of scope by design — SigV4; use
`--backend bedrock`); non-Anthropic Bedrock model bodies (e.g.
Titan/Llama) — only the Anthropic Messages-shaped invoke body is
handled.

<details><summary>Proxy <code>/stats</code> output + content-length bug
note</summary>

```json
"compression": {
  "requests_compressed": 1,
  "avg_compression_pct": 92.4,
  "best_detail": "52,095 → 3,979 tokens",
  "total_tokens_removed": 48116
}
```

The first iteration of this proof surfaced a real bug — a shrunk body
still carried the inbound `Content-Length`, so httpx raised `Too little
data for declared Content-Length`. Fixed by dropping
`content-length`/`content-encoding` on the rewritten path so httpx
recomputes them; covered by a regression test.

</details>

## 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

`mypy headroom` is left unchecked above — type checking runs in the CI
matrix rather than locally on my side; the new code carries type hints
on all public functions. Design spec / feature request: #734.

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
rongabbay 2026-06-16 21:56:16 +02:00 • committed by GitHub
parent e36fccd8cf
commit 7edb27ab24
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 738 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

22
docs/content/docs/proxy.mdx
View file

@ -45,6 +45,7 @@ Telemetry is enabled by default. Opt out with `HEADROOM_TELEMETRY=off` or `--no-
| `--anthropic-api-url` | Anthropic default | Custom Anthropic API URL |
| `--gemini-api-url` | Gemini default | Custom Gemini API URL |
| `--backend` | `anthropic` | Backend: `anthropic`, `bedrock`, `openrouter`, `anyllm`, or `litellm-<provider>` |
| `--bedrock-api-url` | None | Bedrock InvokeModel upstream for the `/model/{id}/invoke` passthrough routes (see [Bedrock via a local gateway](#bedrock-via-a-local-gateway)) |
| `--no-telemetry` | `false` | Disable anonymous telemetry |
| `--stateless` | `false` | Disable filesystem writes and keep runtime state in memory |
@ -243,6 +244,24 @@ headroom proxy --backend azure
OPENROUTER_API_KEY=sk-or-... headroom proxy --backend openrouter
```
### Bedrock via a local gateway
`--backend bedrock` accepts **Anthropic** input (`/v1/messages`) and re-signs to AWS. Some setups are the other way around: the client already speaks **Bedrock** (e.g. Claude Code with `CLAUDE_CODE_USE_BEDROCK=1`, or any AWS SDK pointed at a custom endpoint), sending `POST /model/{id}/invoke` to a local gateway that re-signs and forwards to AWS (LiteLLM, LocalStack, a corporate Bedrock proxy).
`--bedrock-api-url` lets Headroom sit in that chain. It registers passthrough routes for `/model/{id}/invoke` and `/model/{id}/invoke-with-response-stream`, compresses the request body with the same pipeline as `/v1/messages`, and forwards to the gateway:
```bash
headroom proxy --bedrock-api-url http://127.0.0.1:4000
# then point the client's Bedrock endpoint at Headroom:
AWS_ENDPOINT_URL_BEDROCK_RUNTIME=http://127.0.0.1:8787 your-bedrock-client
```
The routes are registered **only** when `--bedrock-api-url` (or `BEDROCK_TARGET_API_URL`) is set — otherwise Bedrock requests fall through unchanged.
<Callout type="warn">
Rewriting the request body invalidates the caller's **SigV4** signature (it covers a hash of the body). Point `--bedrock-api-url` at a gateway that re-signs or does not verify the inbound signature — **never raw AWS**, which would reject the request with 403. For direct-to-AWS compression, use `--backend bedrock` (which re-signs). The two are complementary.
</Callout>
## Environment variables
```bash
@ -256,6 +275,9 @@ export OPENAI_TARGET_API_URL=https://custom.openai.endpoint.com
# Route Anthropic passthrough requests to a custom endpoint
export ANTHROPIC_TARGET_API_URL=https://litellm.company.internal
# Compress Bedrock InvokeModel traffic, forwarding to a re-signing gateway
export BEDROCK_TARGET_API_URL=http://127.0.0.1:4000
headroom proxy
```

14
headroom/cli/proxy.py
View file

@ -575,6 +575,16 @@ def _selected_context_tool() -> str:
default=None,
help="AWS profile name for Bedrock (default: use default credentials)",
)
@click.option(
"--bedrock-api-url",
default=None,
help=(
"Custom Bedrock InvokeModel upstream for the /model/{id}/invoke "
"passthrough routes. Point at a re-signing gateway (LiteLLM, "
"LocalStack), NOT raw AWS — rewriting the body breaks SigV4. "
"(env: BEDROCK_TARGET_API_URL)"
),
)
@click.option(
"--no-telemetry",
is_flag=True,
@ -661,6 +671,7 @@ def proxy(
region: str,
bedrock_region: str | None,
bedrock_profile: str | None,
bedrock_api_url: str | None,
no_telemetry: bool,
stateless: bool,
embedding_server: bool,
@ -876,6 +887,9 @@ def proxy(
backend=backend,
bedrock_region=bedrock_region or region,
bedrock_profile=bedrock_profile,
# CLI flag > env > unset. Matches the BEDROCK_TARGET_API_URL naming of
# the sibling *_TARGET_API_URL passthrough overrides.
bedrock_api_url=bedrock_api_url or os.environ.get("BEDROCK_TARGET_API_URL"),
anyllm_provider=effective_anyllm_provider,
# License / Usage Reporting (managed/enterprise)
license_key=license_key,

17
headroom/providers/proxy_routes.py
View file

@ -405,6 +405,23 @@ def register_provider_routes(app: FastAPI, proxy: Any) -> None:
async def anthropic_messages(request: Request):
return await proxy.handle_anthropic_messages(request)
# AWS Bedrock InvokeModel passthrough. Registered ONLY when an upstream is
# configured (`--bedrock-api-url` / BEDROCK_TARGET_API_URL): without it,
# `/model/{id}/invoke` keeps falling through to the catch-all (verbatim,
# signature-intact) so existing behavior is unchanged. The `{model_id:path}`
# converter captures inference-profile ids that contain dots, colons and
# slashes (e.g. `us.anthropic.claude-sonnet-4-5-20250929-v1:0`). See
# headroom/proxy/handlers/bedrock.py for the SigV4 caveat.
if getattr(proxy.config, "bedrock_api_url", None):
@app.post("/model/{model_id:path}/invoke")
async def bedrock_invoke(request: Request, model_id: str):
return await proxy.handle_bedrock_invoke(request, model_id, stream=False)
@app.post("/model/{model_id:path}/invoke-with-response-stream")
async def bedrock_invoke_stream(request: Request, model_id: str):
return await proxy.handle_bedrock_invoke(request, model_id, stream=True)
@app.post("/v1/messages/count_tokens")
async def anthropic_count_tokens(request: Request):
return await proxy.handle_passthrough(

2
headroom/proxy/handlers/__init__.py
View file

@ -7,6 +7,7 @@ __init__ for all self.* attributes (duck typing).
from headroom.proxy.handlers.anthropic import AnthropicHandlerMixin
from headroom.proxy.handlers.batch import BatchHandlerMixin
from headroom.proxy.handlers.bedrock import BedrockHandlerMixin
from headroom.proxy.handlers.gemini import GeminiHandlerMixin
from headroom.proxy.handlers.openai import OpenAIHandlerMixin
from headroom.proxy.handlers.streaming import StreamingMixin
@ -14,6 +15,7 @@ from headroom.proxy.handlers.streaming import StreamingMixin
__all__ = [
"AnthropicHandlerMixin",
"BatchHandlerMixin",
"BedrockHandlerMixin",
"GeminiHandlerMixin",
"OpenAIHandlerMixin",
"StreamingMixin",

300
headroom/proxy/handlers/bedrock.py Normal file
View file

@ -0,0 +1,300 @@
"""AWS Bedrock ``InvokeModel`` passthrough handler for HeadroomProxy.
Claude Code (and other clients) launched with ``CLAUDE_CODE_USE_BEDROCK=1``
talk to a Bedrock *runtime endpoint* over plain HTTP, POSTing to
``/model/{modelId}/invoke`` and ``/model/{modelId}/invoke-with-response-stream``
instead of the Anthropic ``/v1/messages`` route. Those requests previously fell
through Headroom's catch-all and were forwarded verbatim — no compression.
This mixin intercepts that Bedrock REST shape, compresses the request body with
the **same** ``anthropic_pipeline`` used for ``/v1/messages``, and forwards to a
configurable upstream (``config.bedrock_api_url``). The InvokeModel body for
Anthropic models *is* the Anthropic Messages shape
(``{anthropic_version, system, messages, max_tokens, …}``; the model travels in
the URL), so the existing pipeline applies with no translation.
LIMITATION — SigV4. Rewriting the body invalidates the caller's SigV4 signature
(the signature covers a hash of the body). These routes therefore register
**only** when ``--bedrock-api-url`` / ``BEDROCK_TARGET_API_URL`` is set, and the
target must be a gateway that re-signs or does not verify the inbound signature
(LiteLLM, LocalStack, a corporate Bedrock proxy) — never raw AWS. For
direct-to-AWS compression use ``--backend bedrock`` (which re-signs).
The response is forwarded byte-faithfully: the non-streaming reply is Anthropic
JSON and the streaming reply uses AWS event-stream binary framing — neither is
parsed or mutated, since all compression happens request-side.
"""
from __future__ import annotations
import json
import logging
import time
from typing import TYPE_CHECKING
from urllib.parse import quote
if TYPE_CHECKING:
from fastapi import Request
from fastapi.responses import Response, StreamingResponse
logger = logging.getLogger("headroom.proxy")
LOG_TAG = "bedrock_invoke"
class BedrockHandlerMixin:
"""Mixin providing the Bedrock InvokeModel passthrough handler."""
def _bedrock_upstream_base(self) -> str | None:
"""Resolved Bedrock upstream, or ``None`` when unconfigured.
Returns the normalized ``config.bedrock_api_url`` (trailing slash
stripped). ``None`` means the feature is off — the routes are not even
registered in that case, so a ``None`` here is a defensive guard only.
"""
base = getattr(self.config, "bedrock_api_url", None) # type: ignore[attr-defined]
return base.rstrip("/") if base else None
async def handle_bedrock_invoke(
self,
request: Request,
model_id: str,
*,
stream: bool,
) -> Response | StreamingResponse:
"""Compress and forward a Bedrock ``InvokeModel`` request.
Args:
request: The inbound FastAPI request.
model_id: The Bedrock model / inference-profile id captured from the
URL path (may contain ``.``, ``:`` and ``/``).
stream: ``True`` for ``invoke-with-response-stream``.
"""
from fastapi.responses import JSONResponse
from headroom.proxy.auth_mode import classify_client
from headroom.proxy.helpers import (
COMPRESSION_TIMEOUT_SECONDS,
MAX_MESSAGE_ARRAY_LENGTH,
_headroom_bypass_enabled,
_strip_internal_headers,
extract_tags,
read_request_json_with_bytes,
)
from headroom.proxy.modes import is_cache_mode
from headroom.utils import extract_user_query
start_time = time.time()
request_id = await self._next_request_id() # type: ignore[attr-defined]
base = self._bedrock_upstream_base()
if base is None:
# Routes only register when configured, so this is unreachable in
# practice; fail loud rather than silently forwarding nowhere.
return JSONResponse(
status_code=503,
content={
"error": {
"type": "configuration_error",
"message": "Bedrock passthrough requested but --bedrock-api-url is unset.",
}
},
)
suffix = "invoke-with-response-stream" if stream else "invoke"
url = f"{base}/model/{quote(model_id, safe='')}/{suffix}"
if request.url.query:
url = f"{url}?{request.url.query}"
# Outbound headers (case-insensitive drops). Two header sets:
# - verbatim: forwards the original bytes, so the inbound
# content-length / content-encoding still describe the body.
# - rewritten: the body we forward is decompressed JSON (possibly
# compressed by the pipeline), so content-length must be recomputed
# by httpx and the stale content-encoding dropped. Keeping the
# inbound content-length here is the classic "Too little data for
# declared Content-Length" footgun once the body shrinks.
# We never touch the auth headers — the upstream gateway owns
# (re-)signing.
in_headers = _strip_internal_headers(dict(request.headers.items()))
client = classify_client(dict(request.headers.items()))
tags = extract_tags(dict(request.headers.items()))
verbatim_drop = {"host", "accept-encoding"}
rewritten_drop = verbatim_drop | {"content-length", "content-encoding"}
verbatim_headers = {k: v for k, v in in_headers.items() if k.lower() not in verbatim_drop}
out_headers = {k: v for k, v in in_headers.items() if k.lower() not in rewritten_drop}
# Read the body up front so we can fail open to a verbatim forward on any
# parse error (a malformed body is the gateway's problem, not ours).
try:
body, raw = await read_request_json_with_bytes(request)
except Exception as err:
logger.warning(
"[%s] %s could not parse body; forwarding verbatim: %s",
request_id,
LOG_TAG,
err,
)
raw_only = await request.body()
return await self._forward_bedrock(
url=url,
headers=verbatim_headers,
content=raw_only,
stream=stream,
request_id=request_id,
)
messages = body.get("messages")
bypass = (
_headroom_bypass_enabled(request.headers)
or not getattr(self.config, "optimize", True) # type: ignore[attr-defined]
or is_cache_mode(getattr(self.config, "mode", "token")) # type: ignore[attr-defined]
or not isinstance(messages, list)
or not messages
or len(messages) > MAX_MESSAGE_ARRAY_LENGTH
)
outbound = raw
original_tokens = 0
optimized_tokens = 0
tokens_saved = 0
transforms_applied: tuple[str, ...] = ()
pipeline_timing: dict[str, float] | None = None
if not bypass:
try:
context_limit = self.anthropic_provider.get_context_limit(model_id) # type: ignore[attr-defined]
result = await self._run_compression_in_executor( # type: ignore[attr-defined]
lambda: self.anthropic_pipeline.apply( # type: ignore[attr-defined]
messages=messages,
model=model_id,
model_limit=context_limit,
context=extract_user_query(messages),
request_id=request_id,
),
timeout=COMPRESSION_TIMEOUT_SECONDS,
)
if result.messages != messages:
body["messages"] = result.messages
outbound = json.dumps(body).encode("utf-8")
original_tokens = result.tokens_before
optimized_tokens = result.tokens_after
tokens_saved = max(0, result.tokens_before - result.tokens_after)
transforms_applied = tuple(result.transforms_applied)
pipeline_timing = result.timing
logger.info(
"[%s] %s compressed %d→%d tokens (%d saved) model=%s",
request_id,
LOG_TAG,
result.tokens_before,
result.tokens_after,
tokens_saved,
model_id,
)
except Exception as err:
# Fail open: never break a request because compression failed.
logger.warning(
"[%s] %s compression failed; forwarding verbatim: %s",
request_id,
LOG_TAG,
err,
)
outbound = raw
out_headers["content-type"] = "application/json"
response = await self._forward_bedrock(
url=url,
headers=out_headers,
content=outbound,
stream=stream,
request_id=request_id,
)
# Best-effort metrics. Output tokens are left at 0 (the RequestOutcome
# contract treats 0 as "not measured") — Bedrock responses are forwarded
# byte-faithfully and never parsed. The valuable figure, request-side
# compression, is recorded in full.
try:
from headroom.proxy.outcome import RequestOutcome
await self._record_request_outcome( # type: ignore[attr-defined]
RequestOutcome(
request_id=request_id,
provider="bedrock",
model=model_id,
original_tokens=original_tokens,
optimized_tokens=optimized_tokens,
output_tokens=0,
tokens_saved=tokens_saved,
attempted_input_tokens=original_tokens,
total_latency_ms=(time.time() - start_time) * 1000,
transforms_applied=transforms_applied,
pipeline_timing=pipeline_timing,
tags=tags,
client=client,
)
)
except Exception:
logger.debug("[%s] %s outcome recording failed", request_id, LOG_TAG, exc_info=True)
return response
async def _forward_bedrock(
self,
*,
url: str,
headers: dict[str, str],
content: bytes,
stream: bool,
request_id: str,
) -> Response | StreamingResponse:
"""Stream a request to the Bedrock upstream, byte-faithfully.
Uses the canonical httpx-as-reverse-proxy pattern: open the upstream
with ``stream=True`` so status + headers are available immediately, then
hand the raw byte iterator to ``StreamingResponse`` and close the
upstream connection via a background task. Works for both the JSON
``invoke`` reply and the event-stream ``invoke-with-response-stream``
reply — neither is buffered or mutated.
"""
import httpx
from fastapi.responses import JSONResponse, StreamingResponse
from starlette.background import BackgroundTask
assert self.http_client is not None # type: ignore[attr-defined]
upstream_request = self.http_client.build_request( # type: ignore[attr-defined]
"POST",
url,
headers=headers,
content=content,
)
try:
upstream = await self.http_client.send(upstream_request, stream=True) # type: ignore[attr-defined]
except (httpx.ConnectError, httpx.TimeoutException) as err:
logger.warning("[%s] %s upstream connect failed: %s", request_id, LOG_TAG, err)
return JSONResponse(
status_code=502,
content={
"error": {
"type": "connection_error",
"message": f"Failed to connect to Bedrock upstream: {err}",
}
},
)
# Forward raw (still-encoded) bytes, so strip hop-by-hop headers that
# would conflict with StreamingResponse's own framing. content-encoding
# and content-type are preserved.
resp_headers = {
k: v
for k, v in upstream.headers.items()
if k.lower() not in ("content-length", "transfer-encoding", "connection")
}
media_type = upstream.headers.get("content-type")
return StreamingResponse(
upstream.aiter_raw(),
status_code=upstream.status_code,
headers=resp_headers,
media_type=media_type,
background=BackgroundTask(upstream.aclose),
)

8
headroom/proxy/models.py
View file

@ -107,6 +107,14 @@ class ProxyConfig:
backend: str = "anthropic"
bedrock_region: str = "us-west-2"
bedrock_profile: str | None = None
# Custom upstream for the Bedrock InvokeModel passthrough routes
# (`/model/{id}/invoke[-with-response-stream]`). When set, those routes are
# registered and compress the request body before forwarding here. Point it
# at a re-signing gateway (LiteLLM, LocalStack, a corporate Bedrock
# proxy) — NOT raw AWS, since rewriting the body invalidates the caller's
# SigV4 signature. Leave unset (default) to keep `--backend bedrock`'s
# direct-to-AWS, re-signing behavior unchanged.
bedrock_api_url: str | None = None
anyllm_provider: str = "openai"
# Optimization mode: "token" (rewrite for max compression) or

12
headroom/proxy/server.py
View file

@ -525,6 +525,7 @@ def _check_rust_core() -> tuple[str, str | None]:
from headroom.proxy.handlers import ( # noqa: E402
AnthropicHandlerMixin,
BatchHandlerMixin,
BedrockHandlerMixin,
GeminiHandlerMixin,
OpenAIHandlerMixin,
StreamingMixin,
@ -537,6 +538,7 @@ class HeadroomProxy(
OpenAIHandlerMixin,
GeminiHandlerMixin,
BatchHandlerMixin,
BedrockHandlerMixin,
):
"""Production-ready Headroom optimization proxy."""
@ -3483,6 +3485,7 @@ def _proxy_config_from_env() -> ProxyConfig:
backend=_get_env_str("HEADROOM_BACKEND", "anthropic"),
bedrock_region=_get_env_str("HEADROOM_BEDROCK_REGION", "us-west-2"),
bedrock_profile=os.environ.get("AWS_PROFILE"),
bedrock_api_url=os.environ.get("BEDROCK_TARGET_API_URL"),
anyllm_provider=_get_env_str("HEADROOM_ANYLLM_PROVIDER", "openai"),
disable_kompress=_get_env_bool("HEADROOM_DISABLE_KOMPRESS", False),
max_connections=_get_env_int("HEADROOM_MAX_CONNECTIONS", 500),
@ -3784,6 +3787,14 @@ if __name__ == "__main__":
"--bedrock-profile",
help="AWS profile for Bedrock backend (default: use default credentials)",
)
parser.add_argument(
"--bedrock-api-url",
help=(
"Custom Bedrock InvokeModel upstream for the /model/{id}/invoke "
"passthrough routes — point at a re-signing gateway, not raw AWS "
"(env: BEDROCK_TARGET_API_URL)"
),
)
parser.add_argument(
"--openrouter-api-key",
help="OpenRouter API key (or set OPENROUTER_API_KEY env var)",
@ -3930,6 +3941,7 @@ if __name__ == "__main__":
backend=_get_env_str("HEADROOM_BACKEND", args.backend), # type: ignore[arg-type]
bedrock_region=_get_env_str("HEADROOM_BEDROCK_REGION", args.bedrock_region),
bedrock_profile=args.bedrock_profile or os.environ.get("AWS_PROFILE"),
bedrock_api_url=_get_env_str("BEDROCK_TARGET_API_URL", args.bedrock_api_url),
anyllm_provider=_get_env_str("HEADROOM_ANYLLM_PROVIDER", args.anyllm_provider),
optimize=optimize,
min_tokens_to_crush=_get_env_int("HEADROOM_MIN_TOKENS", args.min_tokens),

363
tests/test_proxy/test_bedrock_passthrough.py Normal file
View file

@ -0,0 +1,363 @@
"""Tests for the AWS Bedrock InvokeModel passthrough handler.
Covers the routes registered by ``register_provider_routes`` when
``--bedrock-api-url`` is set, and the ``handle_bedrock_invoke`` behavior:
1. Routes register ONLY when ``bedrock_api_url`` is configured.
2. A large request body is compressed via ``anthropic_pipeline.apply`` and the
compressed messages are what gets forwarded upstream.
3. Inference-profile model ids (dots/colons) are captured whole and re-encoded
into the upstream URL.
4. The streaming route forwards upstream bytes byte-faithfully and closes the
upstream connection.
5. Fail-open: a malformed JSON body is forwarded verbatim, never a 500.
6. Bypass (``optimize=False``) forwards verbatim — no compression.
7. The request outcome is recorded with ``provider="bedrock"``.
8. ``BEDROCK_TARGET_API_URL`` feeds the env config path.
All forwarding is mocked at ``proxy.http_client`` so no real upstream is needed.
"""
from __future__ import annotations
import json
from unittest.mock import AsyncMock, MagicMock
import pytest
fastapi = pytest.importorskip("fastapi")
httpx = pytest.importorskip("httpx")
from fastapi.testclient import TestClient # noqa: E402
from headroom.proxy.server import ProxyConfig, create_app # noqa: E402
UPSTREAM = "http://127.0.0.1:4000"
SONNET_BEDROCK = "anthropic.claude-3-5-sonnet-20241022-v2:0"
INVOKE = f"/model/{SONNET_BEDROCK}/invoke"
INVOKE_STREAM = f"/model/{SONNET_BEDROCK}/invoke-with-response-stream"
class _FakeUpstream:
"""Minimal stand-in for an httpx streaming response."""
def __init__(
self,
status_code: int = 200,
headers: dict | None = None,
chunks: tuple[bytes, ...] = (b'{"ok":true}',),
) -> None:
self.status_code = status_code
self.headers = httpx.Headers(headers or {"content-type": "application/json"})
self._chunks = list(chunks)
self.closed = False
async def aiter_raw(self):
for chunk in self._chunks:
yield chunk
async def aclose(self):
self.closed = True
class _FakeResult:
"""Stand-in for TransformPipeline.apply's TransformResult."""
def __init__(
self,
messages: list[dict],
tokens_before: int,
tokens_after: int,
transforms: tuple[str, ...] = ("smartcrush",),
) -> None:
self.messages = messages
self.tokens_before = tokens_before
self.tokens_after = tokens_after
self.transforms_applied = list(transforms)
self.timing = {"total": 1.0}
def _make_config(**overrides) -> ProxyConfig:
base = {
"bedrock_api_url": UPSTREAM,
"optimize": True,
"cache_enabled": False,
"rate_limit_enabled": False,
"mode": "token",
}
base.update(overrides)
return ProxyConfig(**base)
def _install_fake_client(proxy, upstream: _FakeUpstream) -> MagicMock:
"""Replace proxy.http_client so forwarding never touches the network."""
client = MagicMock()
client.build_request = MagicMock(return_value=MagicMock(name="upstream_request"))
client.send = AsyncMock(return_value=upstream)
client.aclose = AsyncMock() # awaited by proxy.shutdown() on lifespan exit
proxy.http_client = client
return client
def _forwarded(client: MagicMock) -> tuple[str, dict]:
"""Return (url, parsed_json_body_or_raw) handed to build_request."""
call = client.build_request.call_args
url = call.args[1] if len(call.args) > 1 else call.kwargs["url"]
content = call.kwargs["content"]
try:
parsed = json.loads(content)
except (ValueError, TypeError):
parsed = content
return url, parsed
# ── route gating ──────────────────────────────────────────────────────
def _paths(cfg: ProxyConfig) -> set[str]:
app = create_app(cfg)
return {r.path for r in app.routes if hasattr(r, "path")}
def test_routes_absent_when_bedrock_api_url_unset():
paths = _paths(ProxyConfig())
assert "/model/{model_id:path}/invoke" not in paths
def test_routes_present_when_bedrock_api_url_set():
paths = _paths(_make_config())
assert "/model/{model_id:path}/invoke" in paths
assert "/model/{model_id:path}/invoke-with-response-stream" in paths
# ── compression ───────────────────────────────────────────────────────
def test_invoke_forwards_compressed_messages():
app = create_app(_make_config())
with TestClient(app) as client:
proxy = client.app.state.proxy
http = _install_fake_client(proxy, _FakeUpstream())
compressed = [{"role": "user", "content": "short"}]
proxy.anthropic_pipeline.apply = MagicMock(
return_value=_FakeResult(compressed, tokens_before=5000, tokens_after=200)
)
body = {
"anthropic_version": "bedrock-2023-05-31",
"messages": [{"role": "user", "content": "x" * 5000}],
"max_tokens": 100,
}
resp = client.post(INVOKE, json=body)
assert resp.status_code == 200
_, forwarded = _forwarded(http)
assert forwarded["messages"] == compressed
def test_compressed_body_drops_stale_content_length():
"""Regression: a shrunk body must not carry the inbound content-length, or
httpx raises 'Too little data for declared Content-Length' (caught in
live testing). content-encoding is dropped on the same path so a stale
gzip claim can't mislabel the re-serialized JSON."""
app = create_app(_make_config())
with TestClient(app) as client:
proxy = client.app.state.proxy
http = _install_fake_client(proxy, _FakeUpstream())
proxy.anthropic_pipeline.apply = MagicMock(
return_value=_FakeResult([{"role": "user", "content": "tiny"}], 5000, 100)
)
resp = client.post(
INVOKE,
json={"messages": [{"role": "user", "content": "x" * 8000}], "max_tokens": 8},
)
assert resp.status_code == 200
sent_headers = http.build_request.call_args.kwargs["headers"]
lower = {k.lower() for k in sent_headers}
assert "content-length" not in lower
assert "content-encoding" not in lower
def test_invoke_preserves_non_message_body_fields():
app = create_app(_make_config())
with TestClient(app) as client:
proxy = client.app.state.proxy
http = _install_fake_client(proxy, _FakeUpstream())
proxy.anthropic_pipeline.apply = MagicMock(
return_value=_FakeResult(
[{"role": "user", "content": "c"}], tokens_before=900, tokens_after=100
)
)
body = {
"anthropic_version": "bedrock-2023-05-31",
"messages": [{"role": "user", "content": "y" * 3000}],
"max_tokens": 256,
}
client.post(INVOKE, json=body)
_, forwarded = _forwarded(http)
assert forwarded["anthropic_version"] == "bedrock-2023-05-31"
assert forwarded["max_tokens"] == 256
# ── model id encoding ─────────────────────────────────────────────────
def test_inference_profile_model_id_is_captured_and_reencoded():
app = create_app(_make_config())
with TestClient(app) as client:
proxy = client.app.state.proxy
http = _install_fake_client(proxy, _FakeUpstream())
proxy.anthropic_pipeline.apply = MagicMock(
return_value=_FakeResult([{"role": "user", "content": "c"}], 900, 100)
)
profile = "us.anthropic.claude-sonnet-4-5-20250929-v1:0"
client.post(
f"/model/{profile}/invoke",
json={"messages": [{"role": "user", "content": "z" * 3000}], "max_tokens": 8},
)
url, _ = _forwarded(http)
# Colon is percent-encoded; the whole profile id survives in the path.
assert url == f"{UPSTREAM}/model/us.anthropic.claude-sonnet-4-5-20250929-v1%3A0/invoke"
# ── streaming ─────────────────────────────────────────────────────────
def test_invoke_with_response_stream_is_byte_faithful():
app = create_app(_make_config())
upstream = _FakeUpstream(chunks=(b"event-stream-chunk-1", b"event-stream-chunk-2"))
with TestClient(app) as client:
proxy = client.app.state.proxy
_install_fake_client(proxy, upstream)
proxy.anthropic_pipeline.apply = MagicMock(
return_value=_FakeResult([{"role": "user", "content": "c"}], 900, 100)
)
resp = client.post(
INVOKE_STREAM,
json={"messages": [{"role": "user", "content": "w" * 3000}], "max_tokens": 8},
)
assert resp.status_code == 200
assert resp.content == b"event-stream-chunk-1event-stream-chunk-2"
assert upstream.closed is True
# ── fail-open + bypass ────────────────────────────────────────────────
def test_malformed_body_is_forwarded_verbatim():
app = create_app(_make_config())
with TestClient(app) as client:
proxy = client.app.state.proxy
http = _install_fake_client(proxy, _FakeUpstream())
# Pipeline must NOT be invoked on an unparseable body.
proxy.anthropic_pipeline.apply = MagicMock(side_effect=AssertionError("should not run"))
resp = client.post(
INVOKE,
content=b"not-json-at-all",
headers={"content-type": "application/json"},
)
assert resp.status_code == 200
_, forwarded = _forwarded(http)
assert forwarded == b"not-json-at-all"
def test_optimize_disabled_forwards_verbatim():
app = create_app(_make_config(optimize=False))
with TestClient(app) as client:
proxy = client.app.state.proxy
http = _install_fake_client(proxy, _FakeUpstream())
proxy.anthropic_pipeline.apply = MagicMock(side_effect=AssertionError("should not run"))
body = {"messages": [{"role": "user", "content": "x" * 5000}], "max_tokens": 8}
resp = client.post(INVOKE, json=body)
assert resp.status_code == 200
_, forwarded = _forwarded(http)
assert forwarded["messages"] == body["messages"]
def test_compression_failure_forwards_verbatim():
"""Fail-open: if the pipeline raises, forward the ORIGINAL body untouched
rather than 500ing the request."""
app = create_app(_make_config())
with TestClient(app) as client:
proxy = client.app.state.proxy
http = _install_fake_client(proxy, _FakeUpstream())
proxy.anthropic_pipeline.apply = MagicMock(side_effect=RuntimeError("boom"))
body = {"messages": [{"role": "user", "content": "x" * 5000}], "max_tokens": 8}
resp = client.post(INVOKE, json=body)
assert resp.status_code == 200
_, forwarded = _forwarded(http)
assert forwarded["messages"] == body["messages"]
def test_bypass_header_skips_compression():
"""`x-headroom-bypass: true` forwards verbatim — the pipeline never runs."""
app = create_app(_make_config())
with TestClient(app) as client:
proxy = client.app.state.proxy
http = _install_fake_client(proxy, _FakeUpstream())
proxy.anthropic_pipeline.apply = MagicMock(side_effect=AssertionError("should not run"))
body = {"messages": [{"role": "user", "content": "x" * 5000}], "max_tokens": 8}
resp = client.post(INVOKE, json=body, headers={"x-headroom-bypass": "true"})
assert resp.status_code == 200
_, forwarded = _forwarded(http)
assert forwarded["messages"] == body["messages"]
def test_upstream_connect_error_returns_502():
"""A transport failure to the gateway surfaces as a clean 502, not a crash."""
app = create_app(_make_config())
with TestClient(app) as client:
proxy = client.app.state.proxy
client_mock = _install_fake_client(proxy, _FakeUpstream())
client_mock.send = AsyncMock(side_effect=httpx.ConnectError("no route"))
proxy.anthropic_pipeline.apply = MagicMock(
return_value=_FakeResult([{"role": "user", "content": "c"}], 900, 100)
)
resp = client.post(
INVOKE,
json={"messages": [{"role": "user", "content": "z" * 3000}], "max_tokens": 8},
)
assert resp.status_code == 502
# ── metrics ───────────────────────────────────────────────────────────
def test_outcome_recorded_with_bedrock_provider():
app = create_app(_make_config())
with TestClient(app) as client:
proxy = client.app.state.proxy
_install_fake_client(proxy, _FakeUpstream())
proxy._record_request_outcome = AsyncMock()
proxy.anthropic_pipeline.apply = MagicMock(
return_value=_FakeResult([{"role": "user", "content": "c"}], 1000, 250)
)
client.post(
INVOKE,
json={"messages": [{"role": "user", "content": "q" * 3000}], "max_tokens": 8},
)
assert proxy._record_request_outcome.await_count == 1
outcome = proxy._record_request_outcome.await_args.args[0]
assert outcome.provider == "bedrock"
assert outcome.tokens_saved == 750
# ── env config path ───────────────────────────────────────────────────
def test_env_var_feeds_config(monkeypatch):
from headroom.proxy.server import _proxy_config_from_env
monkeypatch.delenv("HEADROOM_PROXY_CONFIG_JSON", raising=False)
monkeypatch.setenv("BEDROCK_TARGET_API_URL", UPSTREAM)
cfg = _proxy_config_from_env()
assert cfg.bedrock_api_url == UPSTREAM