From 7edb27ab2496b070cbe835b31eb2f828798ddfaa Mon Sep 17 00:00:00 2001 From: rongabbay <110092603+rongabbay@users.noreply.github.com> Date: Tue, 16 Jun 2026 21:56:16 +0200 Subject: [PATCH] feat(proxy): compress AWS Bedrock InvokeModel requests via configurable upstream (#720) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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:`, 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.
Proxy /stats output + content-length bug note ```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.
## 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 --- docs/content/docs/proxy.mdx | 22 ++ headroom/cli/proxy.py | 14 + headroom/providers/proxy_routes.py | 17 + headroom/proxy/handlers/__init__.py | 2 + headroom/proxy/handlers/bedrock.py | 300 +++++++++++++++ headroom/proxy/models.py | 8 + headroom/proxy/server.py | 12 + tests/test_proxy/test_bedrock_passthrough.py | 363 +++++++++++++++++++ 8 files changed, 738 insertions(+) create mode 100644 headroom/proxy/handlers/bedrock.py create mode 100644 tests/test_proxy/test_bedrock_passthrough.py diff --git a/docs/content/docs/proxy.mdx b/docs/content/docs/proxy.mdx index 625d6983e..6567d514f 100644 --- a/docs/content/docs/proxy.mdx +++ b/docs/content/docs/proxy.mdx @@ -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-` | +| `--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. + + +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. + + ## 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 ``` diff --git a/headroom/cli/proxy.py b/headroom/cli/proxy.py index 5c9b6e87b..f411d74d9 100644 --- a/headroom/cli/proxy.py +++ b/headroom/cli/proxy.py @@ -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, diff --git a/headroom/providers/proxy_routes.py b/headroom/providers/proxy_routes.py index 998ac99f9..7d948bac8 100644 --- a/headroom/providers/proxy_routes.py +++ b/headroom/providers/proxy_routes.py @@ -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( diff --git a/headroom/proxy/handlers/__init__.py b/headroom/proxy/handlers/__init__.py index d0d1c8ce7..0acbdd357 100644 --- a/headroom/proxy/handlers/__init__.py +++ b/headroom/proxy/handlers/__init__.py @@ -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", diff --git a/headroom/proxy/handlers/bedrock.py b/headroom/proxy/handlers/bedrock.py new file mode 100644 index 000000000..e3e391f23 --- /dev/null +++ b/headroom/proxy/handlers/bedrock.py @@ -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), + ) diff --git a/headroom/proxy/models.py b/headroom/proxy/models.py index ee4352071..83b578bca 100644 --- a/headroom/proxy/models.py +++ b/headroom/proxy/models.py @@ -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 diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 5aa80b8b6..f64a4c87d 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -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), diff --git a/tests/test_proxy/test_bedrock_passthrough.py b/tests/test_proxy/test_bedrock_passthrough.py new file mode 100644 index 000000000..d867fd4bf --- /dev/null +++ b/tests/test_proxy/test_bedrock_passthrough.py @@ -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