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