From 0dc2e1cb3f7278332d450644831007316d6ac18c Mon Sep 17 00:00:00 2001 From: Yasser Sheikh <158480094+ysheikh2@users.noreply.github.com> Date: Tue, 16 Jun 2026 16:45:24 +0200 Subject: [PATCH] feat(bedrock): cross-region + Converse compression; bundle proxy binary in images (#999) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description The native Bedrock path (Phase D) compresses + signs Anthropic-on-Bedrock requests, but two real-world cases slipped through, and the native binary that powers it was never shipped. This PR closes those gaps as a focused set of give-backs. Aligns with the Rust migration plan (see below). ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [x] New feature (non-breaking change that adds functionality) - [ ] Breaking change - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - **Cross-region inference-profile detection** via a new `bedrock::vendor` module (`canonical_vendor()`), following the design proposed in #953: strip a known geo prefix (`eu.`/`us.`/`apac.`/`global.`) then match the canonical vendor. Geo-prefixed Anthropic profiles (`eu.anthropic.…`) now get live-zone compression instead of being silently skipped; geo-prefixed non-Anthropic vendors stay correctly excluded. - **Converse-body compression (two parts)**: 1. `run_anthropic_compression` no longer bails to passthrough when the body lacks an InvokeModel `anthropic_version` envelope; envelope re-emit stays gated on successful parse. 2. The **live-zone dispatcher now recognizes Bedrock Converse content blocks**. Converse blocks carry no `type` discriminator (the variant is the key: `{"text": …}` vs Anthropic's `{"type":"text","text":…}`), so real Converse user-message text was still passing through uncompressed. A typeless block whose `text` is a JSON string now routes through the same surgical text path. Anthropic blocks always carry `type`, so the Anthropic path is byte-for-byte unchanged; non-text Converse blocks (`{"image":…}`, `{"toolUse":…}`) stay unrecognized and no-op. - **Correct `/converse` upstream routing**: the non-streaming handler resolved the upstream action from a hard-coded `"invoke"`, so `/converse` requests were forwarded to Bedrock's `/invoke` endpoint. It now resolves the action from the inbound path (`extract_invoke_action`), mirroring the streaming handler's `extract_streaming_action`. SigV4 signs the same URL it forwards, so the signature stays consistent. - **`aws-config` `sso` feature**: SSO profiles now resolve through the default credential chain for SigV4 — the credential chain in `docs/bedrock.md` already promised SSO; this makes the code match. - **Ship the `headroom-proxy` binary in published images** (`Dockerfile`): built in the builder stage (`--locked`, with the cargo registry cache mounted at `CARGO_HOME`) and copied into both the debian and distroless runtime images. - **Docs** (`docs/bedrock.md`): document cross-region inference profiles and a "Running the proxy" section. AWS credentials mount at `/home/nonroot/.aws` (the default nonroot image home) where the SDK looks for `~/.aws`, with a note on the root-image alternative. ## Related issues - Closes #976 — ship the `headroom-proxy` binary in published images (this PR implements the exact fix proposed there). - Addresses the **cross-region inference-profile** half of #953 via its proposed `canonical_vendor()` design. Non-Anthropic vendor compression parity (Nova/GLM/MiniMax/ Kimi) is the natural follow-up — `bedrock::vendor` is the shared resolver it can build on. - Extends the native Bedrock InvokeModel compression requested in #734 (the Bedrock slice of #510) to cross-region profiles and Converse bodies. - Partially enables #181 (native, Python-free packaging): the native binary now ships in the images, though full Python-free distribution remains out of scope. ## Alignment with the Rust migration plan Per `docs/spec/022-rust-migration.md`, the migration is **proxy-first**: `headroom-proxy` is the deployable Rust artifact, native routes replace Python passthroughs one at a time (Stage 4 = provider expansion, Bedrock included), and the binary is meant to be "built, tested, and **released together with the Python package**." Two ways this PR advances that: - The binary-in-images change makes the codebase do what the spec already states (ship the artifact) — closing the gap that forced downstreams to build from source. - Hardening the native Bedrock route (cross-region, Converse routing + body compression) is exactly the Stage-4 provider-expansion work, keeping the native path at parity with real traffic so it can be the default rather than a passthrough. ## Testing - [x] Unit tests pass (`cargo test -p headroom-core -p headroom-proxy` — full suites, 0 failures) - [x] Linting passes (`cargo clippy -p headroom-core -p headroom-proxy --all-targets -- -D warnings`) - [x] Formatting passes (`cargo fmt -- --check`) - [x] New tests added — `bedrock::vendor` (foundation + inference-profile matching), `extract_invoke_action` + converse upstream URL, and live-zone Converse text-block routing (`block_has_string_text_field`, converse-vs-anthropic dispatch equivalence). - [x] Manual testing performed ### Test Output ```text $ cargo test -p headroom-core -p headroom-proxy # all suites: ok, 0 failed $ cargo clippy -p headroom-core -p headroom-proxy --all-targets -- -D warnings # Finished, no warnings $ cargo fmt -- --check # clean # image validation (local, proxy/code extras): $ docker build --target runtime ... # debian: /usr/local/bin/headroom-proxy, --help OK $ docker build --target runtime-slim ... # distroless: binary links + --help OK ``` ## Real Behavior Proof - Environment: native Bedrock proxy against `bedrock-runtime.eu-west-2`, SSO profile, model `eu.anthropic.claude-haiku-4-5-20251001-v1:0`. - Exact command / steps: POST a large multi-turn Converse body to `/model/eu.anthropic.claude-haiku-4-5-20251001-v1:0/converse`; separately build the `runtime` + `runtime-slim` targets and run `/usr/local/bin/headroom-proxy --help`. - Observed result: before — `bedrock_compression_skipped` (geo-prefixed id not recognized), forwarded uncompressed to the wrong `/invoke` upstream; after — geo-prefixed id recognized, `/converse` forwarded to the `/converse` upstream, live-zone dispatcher compresses the Converse user-message text, measurable token savings. Images contain a runnable `headroom-proxy` in both variants. - Not tested: non-Anthropic vendor compression parity (#953 follow-up); Converse `toolResult` nested-text compression (follow-up — only top-level Converse text blocks compress today). ## 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 - [x] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes - An earlier revision flipped the EventStream `Accept` default (`*/*`/absent → passthrough); **dropped** — `*/*` is what most clients (incl. reqwest and the proxy's own metrics tests) send while expecting SSE, so forcing passthrough breaks the standard SSE path. - The binary build adds the native-proxy compile to the image build; happy to gate it behind a build arg if maintainers prefer it opt-in. - Addressed a Copilot review round: corrected the `/converse` upstream routing, the stale `run_anthropic_compression` comment, the Dockerfile cargo cache mount + `--locked`, and the nonroot AWS-credentials docs example. --- CHANGELOG.md | 8 + Cargo.lock | 53 ++++++ Cargo.toml | 2 +- Dockerfile | 15 ++ .../headroom-core/src/transforms/live_zone.rs | 85 ++++++++- crates/headroom-proxy/src/bedrock/invoke.rs | 168 ++++++++++++++---- .../src/bedrock/invoke_streaming.rs | 40 +++-- crates/headroom-proxy/src/bedrock/mod.rs | 1 + crates/headroom-proxy/src/bedrock/vendor.rs | 82 +++++++++ docs/bedrock.md | 46 ++++- 10 files changed, 435 insertions(+), 65 deletions(-) create mode 100644 crates/headroom-proxy/src/bedrock/vendor.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index dba10dff1..08fcd9382 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,8 +16,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * **memory:** opt-in Apple-GPU (MPS) embedding offload via `HEADROOM_EMBEDDER_RUNTIME=pytorch_mps`. When set (and Apple MPS is available), the memory embedder runs on the torch sentence-transformers backend on the Apple GPU instead of the default ONNX CPU embedder, freeing the CPU under load. If MPS or the dependencies are unavailable, Headroom logs a warning and uses the existing default embedder selection path (ONNX when available, then the pre-existing local fallback). MPS encode calls are serialized internally (torch-MPS is not thread-safe). Adds the new `[pytorch-mps]` extra (`pip install 'headroom-ai[pytorch-mps]'`). Default behavior is unchanged. +### Features + +* **proxy:** cross-region Bedrock inference-profile detection — geo-prefixed model IDs (`eu.`/`us.`/`apac.`/`global.`) are now resolved to their canonical vendor, so Anthropic cross-region profiles (e.g. `eu.anthropic.claude-haiku-4-5-20251001-v1:0`) receive live-zone compression instead of being silently skipped ([#999](https://github.com/chopratejas/headroom/pull/999)). +* **proxy:** Converse-body compression on the native Bedrock route — the live-zone dispatcher now recognizes Bedrock Converse content blocks (typeless `{"text": …}`, not only Anthropic `{"type":"text", …}`), so Converse user-message text compresses; `run_anthropic_compression` no longer bails to passthrough when the body lacks an InvokeModel `anthropic_version` envelope, and envelope re-emit stays gated on successful parse ([#999](https://github.com/chopratejas/headroom/pull/999)). +* **docker:** bundle `headroom-proxy` binary in published `runtime` and `runtime-slim` images — closes [#976](https://github.com/chopratejas/headroom/issues/976) ([#999](https://github.com/chopratejas/headroom/pull/999)). + ### Bug Fixes +* **proxy:** enable SSO credential resolution in the native Bedrock route via the `aws-config` `sso` feature flag, making the credential chain match what `docs/bedrock.md` already documented ([#999](https://github.com/chopratejas/headroom/pull/999)). +* **proxy:** route native Bedrock `/model/{id}/converse` requests to the upstream Converse endpoint instead of the hard-coded `/invoke` action — the non-streaming handler now resolves the action from the inbound path, matching the streaming handler ([#999](https://github.com/chopratejas/headroom/pull/999)). * **ccr:** make retrieval store TTL configurable with `HEADROOM_CCR_TTL_SECONDS`, expose the effective TTL in `/v1/retrieve/stats`, and distinguish expired retrievals from missing hashes. * **proxy:** add native Bedrock `/model/{id}/converse-stream` route and forward it through the existing streaming EventStream/SSE pipeline. * **wrap (codex):** fix `headroom wrap codex` producing a `config.toml` with duplicate top-level `model_provider` / `openai_base_url` keys (TOML-spec error) when the user had already configured their own provider. The injector now rewrites pre-existing top-level `model_provider` and `openai_base_url` lines in place — the previous value is kept in a `# was: …` trailing comment — instead of unconditionally prepending a duplicate, so `codex` can start against the proxy. The pre-wrap snapshot mechanism continues to byte-for-byte restore the original file on `headroom unwrap codex`. diff --git a/Cargo.lock b/Cargo.lock index e02e2837b..abcac5c9c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -257,6 +257,8 @@ checksum = "50f156acdd2cf55f5aa53ee416c4ac851cf1222694506c0b1f78c85695e9ca9d" dependencies = [ "aws-credential-types", "aws-runtime", + "aws-sdk-sso", + "aws-sdk-ssooidc", "aws-sdk-sts", "aws-smithy-async", "aws-smithy-http", @@ -267,11 +269,14 @@ dependencies = [ "aws-types", "bytes", "fastrand", + "hex", "http 1.4.0", + "sha1", "time", "tokio", "tracing", "url", + "zeroize", ] [[package]] @@ -333,6 +338,54 @@ dependencies = [ "uuid", ] +[[package]] +name = "aws-sdk-sso" +version = "1.98.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d69c77aafa20460c68b6b3213c84f6423b6e76dbf89accd3e1789a686ffd9489" +dependencies = [ + "aws-credential-types", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-observability", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-types", + "bytes", + "fastrand", + "http 0.2.12", + "http 1.4.0", + "regex-lite", + "tracing", +] + +[[package]] +name = "aws-sdk-ssooidc" +version = "1.100.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c7e7b09346d5ca22a2a08267555843a6a0127fb20d8964cb6ecfb8fdb190225" +dependencies = [ + "aws-credential-types", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-observability", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-types", + "bytes", + "fastrand", + "http 0.2.12", + "http 1.4.0", + "regex-lite", + "tracing", +] + [[package]] name = "aws-sdk-sts" version = "1.103.0" diff --git a/Cargo.toml b/Cargo.toml index 6f0f59c1d..dbb26a7f6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -62,7 +62,7 @@ pyo3 = { version = "0.24", features = ["abi3-py310"] } # (env vars, profiles, IMDS, ECS task role, etc); `aws-credential-types` # exposes `Credentials` so the signer accepts whatever the chain returned. aws-sigv4 = { version = "1", default-features = false, features = ["sign-http", "http1"] } -aws-config = { version = "1", default-features = false, features = ["behavior-version-latest", "rustls", "rt-tokio"] } +aws-config = { version = "1", default-features = false, features = ["behavior-version-latest", "rustls", "rt-tokio", "sso"] } aws-credential-types = { version = "1", default-features = false } # `Identity` lives in aws-smithy-runtime-api; the SigV4 builder # accepts `&Identity`. Pinning the version explicitly avoids a diff --git a/Dockerfile b/Dockerfile index 59281b8a5..4e4628db5 100644 --- a/Dockerfile +++ b/Dockerfile @@ -60,6 +60,17 @@ RUN --mount=type=cache,target=/root/.cache/uv \ RUN cd /tmp && python -c "from headroom._core import DiffCompressor, SmartCrusher; \ print(f'build-stage rust core verify OK: {DiffCompressor.__name__}, {SmartCrusher.__name__}')" +# Build the native Rust reverse proxy binary and stage it for the runtime +# images (issue #976). These images already run "the proxy"; bundling the +# native `headroom-proxy` binary lets operators front the Python proxy with +# the Rust SigV4 / live-zone compression path from the same image. The +# binary is copied out of the cache-mounted target dir into a persistent +# path so the COPY in the runtime stages can pick it up. +RUN --mount=type=cache,target=/usr/local/cargo/registry \ + --mount=type=cache,target=/build/target \ + cargo build --release --locked --bin headroom-proxy && \ + cp target/release/headroom-proxy /usr/local/bin/headroom-proxy + # ---- Runtime stage (python-slim): supports root/nonroot via build arg ---- FROM python:${PYTHON_VERSION}-slim AS runtime-slim-base @@ -72,6 +83,8 @@ RUN apt-get update && \ COPY --from=builder ${PYTHON_SITE_PACKAGES} ${PYTHON_SITE_PACKAGES} COPY --from=builder /usr/local/bin/headroom /usr/local/bin/headroom +# Native Rust reverse proxy binary (issue #976). +COPY --from=builder /usr/local/bin/headroom-proxy /usr/local/bin/headroom-proxy RUN mkdir -p /home/nonroot /data && \ if [ "$RUNTIME_USER" = "nonroot" ]; then \ @@ -104,6 +117,8 @@ ARG RUNTIME_USER=nonroot ARG PYTHON_SITE_PACKAGES COPY --from=builder ${PYTHON_SITE_PACKAGES} ${PYTHON_SITE_PACKAGES} +# Native Rust reverse proxy binary (issue #976). +COPY --from=builder /usr/local/bin/headroom-proxy /usr/local/bin/headroom-proxy USER ${RUNTIME_USER} WORKDIR /app diff --git a/crates/headroom-core/src/transforms/live_zone.rs b/crates/headroom-core/src/transforms/live_zone.rs index d621d7165..6105dd346 100644 --- a/crates/headroom-core/src/transforms/live_zone.rs +++ b/crates/headroom-core/src/transforms/live_zone.rs @@ -1009,6 +1009,23 @@ enum SlotKind { /// latest user message. Errors out on shapes the dispatcher does not /// support (e.g. structured-array `content` inside a tool_result — /// rare; we degrade to NoChange in that case). +/// Whether a content block (no `type` key) carries a JSON-string `text` +/// field — the Bedrock Converse text-block shape (`{"text": "..."}`). +/// Used to route typeless Converse text through the Anthropic text path. +/// Blocks whose `text` is absent or non-string (e.g. `{"image": ...}`, +/// `{"toolUse": ...}`) return false and stay unrecognized → no-op. +fn block_has_string_text_field(block_json: &str) -> bool { + #[derive(Deserialize)] + struct Probe<'a> { + #[serde(borrow, default)] + text: Option<&'a RawValue>, + } + serde_json::from_str::>(block_json) + .ok() + .and_then(|p| p.text) + .is_some_and(|t| t.get().trim_start().starts_with('"')) +} + fn plan_block_replacements( body_raw: &[u8], target_msg_idx: usize, @@ -1072,7 +1089,18 @@ fn plan_block_replacements( let header: BlockHeader<'_> = serde_json::from_str(block_raw.get()).map_err(|_| PlanError::ParseFailed)?; - let block_type = header.r#type.unwrap_or("unknown").to_string(); + // Bedrock Converse content blocks carry no `type` discriminator — + // the variant is the key itself (`{"text": ...}`, `{"image": ...}`, + // `{"toolUse": ...}`). A typeless block whose `text` field is a + // JSON string is Converse text; route it through the same surgical + // path as an Anthropic `{"type":"text","text":...}` block so + // Converse user-message text compresses too. Anthropic blocks + // always carry `type`, so this never alters the Anthropic path. + let block_type = match header.r#type { + Some(t) => t.to_string(), + None if block_has_string_text_field(block_raw.get()) => "text".to_string(), + None => "unknown".to_string(), + }; if HOT_ZONE_BLOCK_TYPES.iter().any(|t| *t == block_type) { slots.push(PlanSlot { @@ -1611,6 +1639,61 @@ mod tests { assert!(matches!(out, LiveZoneOutcome::NoChange { .. })); } + #[test] + fn block_has_string_text_field_detects_converse_text_only() { + // Converse text block: typeless, string `text` → recognized. + assert!(block_has_string_text_field(r#"{"text":"hello"}"#)); + // Non-text Converse blocks must NOT be mistaken for text. + assert!(!block_has_string_text_field( + r#"{"image":{"format":"png"}}"# + )); + assert!(!block_has_string_text_field(r#"{"toolUse":{"name":"x"}}"#)); + // `text` present but not a JSON string → not Converse text. + assert!(!block_has_string_text_field(r#"{"text":["a"]}"#)); + assert!(!block_has_string_text_field(r#"{"text":{"v":1}}"#)); + } + + #[test] + fn converse_typeless_text_block_routes_like_anthropic_text() { + // Bedrock Converse content blocks omit the `type` discriminator — + // `{"text": "..."}` instead of `{"type":"text","text":"..."}`. The + // dispatcher must treat the two identically so Converse user-message + // text compresses like Anthropic text. + let payload = "{\"k\": \"v\", \"n\": 1}\n".repeat(200); + let converse = body(json!({ + "messages": [{"role": "user", "content": [{"text": payload}]}] + })); + let anthropic = body(json!({ + "messages": [{"role": "user", "content": [{"type": "text", "text": payload}]}] + })); + let c = compress_anthropic_live_zone(&converse, 0, AuthMode::Payg, DEFAULT_MODEL).unwrap(); + let a = compress_anthropic_live_zone(&anthropic, 0, AuthMode::Payg, DEFAULT_MODEL).unwrap(); + + // Identical dispatch outcome (both Modified or both NoChange). + assert_eq!( + std::mem::discriminant(&c), + std::mem::discriminant(&a), + "converse text block must dispatch like an anthropic text block" + ); + let cm = match &c { + LiveZoneOutcome::NoChange { manifest } => manifest, + LiveZoneOutcome::Modified { manifest, .. } => manifest, + }; + let am = match &a { + LiveZoneOutcome::NoChange { manifest } => manifest, + LiveZoneOutcome::Modified { manifest, .. } => manifest, + }; + // The Converse block is now classified the same as Anthropic text + // (before this change it was an unrecognized typeless block). + assert_eq!(cm.block_outcomes.len(), 1); + assert_eq!(am.block_outcomes.len(), 1); + assert_eq!( + cm.block_outcomes[0].block_type, + am.block_outcomes[0].block_type + ); + assert_eq!(cm.block_outcomes[0].block_type, "text"); + } + #[test] fn manifest_records_messages_below_floor() { let b = body(json!({ diff --git a/crates/headroom-proxy/src/bedrock/invoke.rs b/crates/headroom-proxy/src/bedrock/invoke.rs index 84c93a6cb..2571e8bcb 100644 --- a/crates/headroom-proxy/src/bedrock/invoke.rs +++ b/crates/headroom-proxy/src/bedrock/invoke.rs @@ -65,10 +65,7 @@ use crate::proxy::AppState; // would risk drift from the middleware's resolution + WARN log. use headroom_core::auth_mode::AuthMode; -/// Anthropic vendor prefix as encoded in Bedrock model ids -/// (`anthropic.claude-3-haiku-...`). Literal-match per project rule -/// "no regexes for parsing the model ID". -const ANTHROPIC_VENDOR_PREFIX: &str = "anthropic."; +use crate::bedrock::vendor::is_anthropic_model_id; /// RAII guard that observes the `bedrock_invoke_latency_seconds` /// histogram on drop. Created at handler entry; observed when the @@ -102,6 +99,27 @@ impl Drop for LatencyGuard { /// region. Only used when `Config::bedrock_endpoint` is `None`. const BEDROCK_RUNTIME_HOST_TEMPLATE: &str = "bedrock-runtime.{region}.amazonaws.com"; +/// Bedrock non-streaming action path segments. `invoke` is the legacy +/// InvokeModel surface; `converse` is the unified Converse surface. +/// Both mount the same handler (see `proxy.rs`), so the action is +/// resolved from the inbound path — otherwise `/converse` requests +/// would be forwarded to the upstream `/invoke` endpoint. +const INVOKE_ACTION: &str = "invoke"; +const CONVERSE_ACTION: &str = "converse"; + +/// Resolve the Bedrock action from the inbound request path. Mirrors +/// `invoke_streaming::extract_streaming_action` for the non-streaming +/// surfaces (`/invoke`, `/converse`). +fn extract_invoke_action(path: &str) -> Option<&'static str> { + if path.ends_with(&format!("/{INVOKE_ACTION}")) { + Some(INVOKE_ACTION) + } else if path.ends_with(&format!("/{CONVERSE_ACTION}")) { + Some(CONVERSE_ACTION) + } else { + None + } +} + /// Axum POST handler for `/model/{model_id}/invoke`. /// /// Buffers the body so the live-zone compressor + SigV4 signer can @@ -155,7 +173,7 @@ pub async fn handle_invoke( "bedrock invoke route received request" ); - let is_anthropic = model_id.starts_with(ANTHROPIC_VENDOR_PREFIX); + let is_anthropic = is_anthropic_model_id(&model_id); let outbound_body: Bytes = if is_anthropic { run_anthropic_compression(&body, &state, auth_mode, &request_id) } else { @@ -169,9 +187,30 @@ pub async fn handle_invoke( body.clone() }; + // Resolve the Bedrock action from the inbound path so `/converse` + // forwards to the upstream Converse endpoint instead of `/invoke`. + // Both paths mount this handler (see `proxy.rs`); the streaming + // sibling resolves its action the same way. + let action = match extract_invoke_action(uri.path()) { + Some(a) => a, + None => { + tracing::error!( + event = "bedrock_invoke_action_invalid", + request_id = %request_id, + path = %uri.path(), + "bedrock invoke: unrecognized action path" + ); + return error_response( + StatusCode::BAD_REQUEST, + "bedrock_invoke_action_invalid", + "Unsupported Bedrock action path", + ); + } + }; + // Build the upstream URL based on configured endpoint or // region-derived default. - let upstream_url = match build_bedrock_upstream(&state, &model_id, &uri, "invoke") { + let upstream_url = match build_bedrock_upstream(&state, &model_id, &uri, action) { Ok(u) => u, Err(msg) => { tracing::error!( @@ -363,24 +402,26 @@ fn run_anthropic_compression( _auth_mode: AuthMode, request_id: &str, ) -> Bytes { - // Validate envelope shape. If the body isn't a valid Bedrock - // envelope we still forward verbatim — the compressor would have - // refused too — but log loudly. - if let Err(e) = BedrockEnvelope::parse(body) { - tracing::warn!( - event = "bedrock_envelope_parse_error", + // Detect envelope shape. A parseable InvokeModel envelope takes the + // re-emit path below (anthropic_version pinned first); a non-envelope + // body (e.g. a Converse-shaped payload) still runs through the + // compressor but skips envelope re-emit. The body is NOT guaranteed + // unchanged on parse failure — we log which path we took. + let parsed_envelope = BedrockEnvelope::parse(body).is_ok(); + if parsed_envelope { + tracing::info!( + event = "bedrock_envelope_parsed", request_id = %request_id, - error = %e, - "bedrock invoke: envelope parse failed; passing body through unchanged" + body_bytes = body.len(), + "bedrock invoke: envelope validated; dispatching to live-zone compressor" + ); + } else { + tracing::info!( + event = "bedrock_envelope_parse_skipped", + request_id = %request_id, + "bedrock invoke: envelope parse skipped; attempting generic anthropic compression" ); - return body.clone(); } - tracing::info!( - event = "bedrock_envelope_parsed", - request_id = %request_id, - body_bytes = body.len(), - "bedrock invoke: envelope validated; dispatching to live-zone compressor" - ); // PR-E3: Bedrock uses IAM-signed AWS SigV4 downstream. Inbound // requests to the proxy may or may not carry their own auth, but @@ -410,20 +451,24 @@ fn run_anthropic_compression( body.clone() } AnthropicOutcome::Compressed { body: new_body, .. } => { - // Defence-in-depth: re-emit so anthropic_version is the - // first key. With preserve_order this is a no-op on the - // happy path. - match BedrockEnvelope::ensure_anthropic_version_first(&new_body) { - Ok(b) => b, - Err(e) => { - tracing::error!( - event = "bedrock_envelope_reemit_failed", - request_id = %request_id, - error = %e, - "bedrock invoke: failed to re-emit envelope; falling back to original body" - ); - body.clone() + if parsed_envelope { + // Defence-in-depth: re-emit so anthropic_version is the + // first key. With preserve_order this is a no-op on the + // happy path. + match BedrockEnvelope::ensure_anthropic_version_first(&new_body) { + Ok(b) => b, + Err(e) => { + tracing::error!( + event = "bedrock_envelope_reemit_failed", + request_id = %request_id, + error = %e, + "bedrock invoke: failed to re-emit envelope; falling back to original body" + ); + body.clone() + } } + } else { + new_body } } } @@ -538,12 +583,57 @@ fn error_response(status: StatusCode, event: &str, msg: &str) -> Response { mod tests { use super::*; + // Vendor/model-id classification is tested in `bedrock::vendor`. + #[test] - fn anthropic_vendor_prefix_match() { - assert!("anthropic.claude-3-haiku-20240307-v1:0".starts_with(ANTHROPIC_VENDOR_PREFIX)); - assert!("anthropic.claude-3-5-sonnet-20241022-v2:0".starts_with(ANTHROPIC_VENDOR_PREFIX)); - assert!(!"amazon.titan-text-express-v1".starts_with(ANTHROPIC_VENDOR_PREFIX)); - assert!(!"meta.llama3-70b-instruct-v1:0".starts_with(ANTHROPIC_VENDOR_PREFIX)); + fn extract_invoke_action_supports_both_bedrock_paths() { + assert_eq!( + extract_invoke_action("/model/anthropic.claude-3-haiku-20240307-v1:0/invoke"), + Some(INVOKE_ACTION) + ); + assert_eq!( + extract_invoke_action("/model/anthropic.claude-3-haiku-20240307-v1:0/converse"), + Some(CONVERSE_ACTION) + ); + // Streaming actions are handled by `invoke_streaming`, not here. + assert_eq!( + extract_invoke_action( + "/model/anthropic.claude-3-haiku-20240307-v1:0/invoke-with-response-stream" + ), + None + ); + assert_eq!(extract_invoke_action("/model/foo/unknown"), None); + } + + #[test] + fn build_upstream_routes_converse_to_converse_endpoint() { + use crate::config::Config; + let mut config = Config::for_test(Url::parse("http://up:8080").unwrap()); + config.bedrock_region = "us-west-2".to_string(); + let state = AppState { + config: std::sync::Arc::new(config), + client: reqwest::Client::new(), + bedrock_credentials: None, + drift_state: crate::cache_stabilization::drift_detector::DriftState::new(8), + vertex_token_source: std::sync::Arc::new(crate::vertex::StaticTokenSource::new( + "test".to_string(), + )), + }; + let uri: Uri = "/model/anthropic.claude-3-haiku-20240307-v1:0/converse" + .parse() + .unwrap(); + let action = extract_invoke_action(uri.path()).unwrap(); + let url = build_bedrock_upstream( + &state, + "anthropic.claude-3-haiku-20240307-v1:0", + &uri, + action, + ) + .unwrap(); + assert_eq!( + url.as_str(), + "https://bedrock-runtime.us-west-2.amazonaws.com/model/anthropic.claude-3-haiku-20240307-v1:0/converse" + ); } #[test] diff --git a/crates/headroom-proxy/src/bedrock/invoke_streaming.rs b/crates/headroom-proxy/src/bedrock/invoke_streaming.rs index 347195f6e..ce66429f9 100644 --- a/crates/headroom-proxy/src/bedrock/invoke_streaming.rs +++ b/crates/headroom-proxy/src/bedrock/invoke_streaming.rs @@ -74,8 +74,7 @@ use crate::proxy::AppState; // `Extension` extractor. use headroom_core::auth_mode::AuthMode; -/// Anthropic vendor prefix as encoded in Bedrock model ids. -const ANTHROPIC_VENDOR_PREFIX: &str = "anthropic."; +use crate::bedrock::vendor::is_anthropic_model_id; /// AWS Bedrock Runtime DNS template. const BEDROCK_RUNTIME_HOST_TEMPLATE: &str = "bedrock-runtime.{region}.amazonaws.com"; @@ -153,7 +152,7 @@ pub async fn handle_invoke_streaming( ); // 1. Live-zone compression for Anthropic-shape bodies (same as D1). - let is_anthropic = model_id.starts_with(ANTHROPIC_VENDOR_PREFIX); + let is_anthropic = is_anthropic_model_id(&model_id); let outbound_body: Bytes = if is_anthropic { run_anthropic_compression(&body, &state, auth_mode, &request_id) } else { @@ -855,14 +854,13 @@ fn run_anthropic_compression( ) -> Bytes { use crate::bedrock::envelope::BedrockEnvelope; - if let Err(e) = BedrockEnvelope::parse(body) { - tracing::warn!( - event = "bedrock_envelope_parse_error", + let parsed_envelope = BedrockEnvelope::parse(body).is_ok(); + if !parsed_envelope { + tracing::info!( + event = "bedrock_envelope_parse_skipped", request_id = %request_id, - error = %e, - "bedrock invoke-streaming: envelope parse failed; passing body through unchanged" + "bedrock invoke-streaming: envelope parse skipped; attempting generic anthropic compression" ); - return body.clone(); } // PR-E3: Bedrock channel hard-codes OAuth so cache_control @@ -887,17 +885,21 @@ fn run_anthropic_compression( body.clone() } AnthropicOutcome::Compressed { body: new_body, .. } => { - match BedrockEnvelope::ensure_anthropic_version_first(&new_body) { - Ok(b) => b, - Err(e) => { - tracing::error!( - event = "bedrock_envelope_reemit_failed", - request_id = %request_id, - error = %e, - "bedrock invoke-streaming: failed to re-emit envelope" - ); - body.clone() + if parsed_envelope { + match BedrockEnvelope::ensure_anthropic_version_first(&new_body) { + Ok(b) => b, + Err(e) => { + tracing::error!( + event = "bedrock_envelope_reemit_failed", + request_id = %request_id, + error = %e, + "bedrock invoke-streaming: failed to re-emit envelope" + ); + body.clone() + } } + } else { + new_body } } } diff --git a/crates/headroom-proxy/src/bedrock/mod.rs b/crates/headroom-proxy/src/bedrock/mod.rs index 780ed4d18..d80f87092 100644 --- a/crates/headroom-proxy/src/bedrock/mod.rs +++ b/crates/headroom-proxy/src/bedrock/mod.rs @@ -52,6 +52,7 @@ pub mod eventstream_to_sse; pub mod invoke; pub mod invoke_streaming; pub mod sigv4; +pub mod vendor; pub use auth_mode_layer::classify_and_attach_auth_mode; pub use envelope::{BedrockEnvelope, EnvelopeError}; diff --git a/crates/headroom-proxy/src/bedrock/vendor.rs b/crates/headroom-proxy/src/bedrock/vendor.rs new file mode 100644 index 000000000..bfa378be1 --- /dev/null +++ b/crates/headroom-proxy/src/bedrock/vendor.rs @@ -0,0 +1,82 @@ +//! Canonical Bedrock vendor resolution. +//! +//! Bedrock model IDs are `.--` (e.g. +//! `anthropic.claude-3-haiku-20240307-v1:0`). Cross-region inference +//! profiles prepend a geo routing token before the vendor +//! (`eu.anthropic.…`, `us.anthropic.…`, `apac.anthropic.…`, +//! `global.anthropic.…`). Matching the bare `anthropic.` prefix alone +//! silently skips compression for those profiles, so we strip a known +//! geo prefix first, then take the leading dot-segment as the vendor. +//! Literal matching only — no regexes (project rule). + +/// Cross-region inference-profile routing prefixes AWS prepends to the +/// vendor segment. Stripped (once) before vendor resolution. Kept to a +/// closed, known set so an unrelated `something.anthropic.x` id is not +/// mistaken for an Anthropic inference profile. +const GEO_PREFIXES: [&str; 4] = ["eu.", "us.", "apac.", "global."]; + +/// Resolve the canonical vendor of a Bedrock model id, stripping a +/// cross-region inference-profile geo prefix first. +/// +/// `eu.anthropic.claude-…` → `anthropic`; `amazon.titan-…` → `amazon`; +/// `global.amazon.nova-…` → `amazon`. +pub fn canonical_vendor(model_id: &str) -> &str { + let stripped = GEO_PREFIXES + .iter() + .find_map(|p| model_id.strip_prefix(p)) + .unwrap_or(model_id); + stripped.split('.').next().unwrap_or(stripped) +} + +/// Whether the model id is an Anthropic-shape model — a foundation +/// model (`anthropic.…`) or a cross-region inference profile +/// (`.anthropic.…`) — i.e. eligible for the live-zone Anthropic +/// compression + envelope pipeline. +pub fn is_anthropic_model_id(model_id: &str) -> bool { + canonical_vendor(model_id) == "anthropic" +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn canonical_vendor_strips_known_cross_region_prefix() { + assert_eq!( + canonical_vendor("anthropic.claude-3-haiku-20240307-v1:0"), + "anthropic" + ); + assert_eq!( + canonical_vendor("eu.anthropic.claude-haiku-4-5-20251001-v1:0"), + "anthropic" + ); + assert_eq!(canonical_vendor("amazon.titan-text-express-v1"), "amazon"); + assert_eq!(canonical_vendor("global.amazon.nova-lite-v1:0"), "amazon"); + // Unknown leading token is NOT stripped — stays as its own vendor. + assert_eq!(canonical_vendor("random.anthropic.x"), "random"); + } + + #[test] + fn anthropic_model_id_matches_foundation_and_inference_profiles() { + assert!(is_anthropic_model_id( + "anthropic.claude-3-haiku-20240307-v1:0" + )); + assert!(is_anthropic_model_id( + "eu.anthropic.claude-haiku-4-5-20251001-v1:0" + )); + assert!(is_anthropic_model_id( + "us.anthropic.claude-3-5-sonnet-20241022-v2:0" + )); + assert!(is_anthropic_model_id( + "apac.anthropic.claude-3-5-sonnet-20240620-v1:0" + )); + assert!(is_anthropic_model_id( + "global.anthropic.claude-haiku-4-5-20251001-v1:0" + )); + // Non-Anthropic vendors, including geo-prefixed, stay false. + assert!(!is_anthropic_model_id("amazon.titan-text-express-v1")); + assert!(!is_anthropic_model_id("meta.llama3-70b-instruct-v1:0")); + assert!(!is_anthropic_model_id("eu.amazon.nova-lite-v1:0")); + assert!(!is_anthropic_model_id("mistral.voxtral-mini-3b-2507")); + } +} diff --git a/docs/bedrock.md b/docs/bedrock.md index 31d502c44..c3a51bc45 100644 --- a/docs/bedrock.md +++ b/docs/bedrock.md @@ -16,6 +16,40 @@ This document covers how to deploy the Bedrock-native surface, how compression p | Per-model + per-region Prometheus metrics | PR-D3 — exposed at `GET /metrics` | | OAuth compression policy gates (no auto cache_control, lossless-only) | Phase F PR-F2/F3 (gates the marker D3 wires) | +## Running the proxy + +The native surface lives in the `headroom-proxy` binary, which ships in the published +container images (every `proxy`-extra tag) at `/usr/local/bin/headroom-proxy`. You can run +it directly from any published image — no separate build: + +```sh +docker run --rm -p 8787:8787 \ + -v "$HOME/.aws:/home/nonroot/.aws:ro" \ + -e HEADROOM_PROXY_AWS_PROFILE=my-profile \ + --entrypoint headroom-proxy \ + ghcr.io/chopratejas/headroom:latest \ + --listen 0.0.0.0:8787 \ + --upstream https://bedrock-runtime.us-east-1.amazonaws.com \ + --bedrock-region us-east-1 +``` + +The published images default to the `nonroot` user (home `/home/nonroot`), so AWS +credentials are mounted at `/home/nonroot/.aws` — that is where the SDK looks for +`~/.aws`. For a root-based image (`RUNTIME_USER=root` build), mount to `/root/.aws` +instead, or pass `--user root`. + +Then point the AWS SDK / CLI at the proxy: + +```sh +AWS_ENDPOINT_URL_BEDROCK_RUNTIME=http://localhost:8787 \ + aws bedrock-runtime invoke-model --model-id anthropic.claude-3-haiku-20240307-v1:0 ... +``` + +The proxy can also drop in front of the Python proxy (`--upstream http://127.0.0.1:8788`) +so non-Bedrock traffic is forwarded while Bedrock requests are signed + compressed +natively. The default `--enable-bedrock-native=true` mounts the Bedrock routes; everything +else is passed through to `--upstream`. + ## AWS credential configuration The proxy uses the [aws-config default credential chain](https://docs.aws.amazon.com/sdkref/latest/guide/standardized-credentials.html), resolved once at startup. @@ -75,17 +109,19 @@ headroom-proxy \ ## Supported model IDs -The proxy classifies model IDs by **literal vendor prefix** — no regexes. Any model ID starting with `anthropic.` is treated as Anthropic-shape: the live-zone compression dispatcher runs over the body, the envelope is re-emitted with `anthropic_version` preserved as the first key, and the request is signed with SigV4. +The proxy classifies model IDs by **literal vendor match** — no regexes. It strips a known cross-region inference-profile geo prefix (`eu.`, `us.`, `apac.`, `global.`) if present, then takes the leading dot-segment as the vendor. A model is treated as Anthropic-shape when that canonical vendor is `anthropic` — so both the bare `anthropic.…` foundation models and the geo-prefixed inference profiles (`eu.anthropic.…`, `us.anthropic.…`, `apac.anthropic.…`, `global.anthropic.…`) qualify. For those, the live-zone compression dispatcher runs over the body, the envelope is re-emitted with `anthropic_version` preserved as the first key, and the request is signed with SigV4. Examples that hit the Anthropic compression path: -- `anthropic.claude-3-haiku-20240307-v1:0` +- `anthropic.claude-3-haiku-20240307-v1:0` (foundation model) - `anthropic.claude-3-5-sonnet-20241022-v2:0` -- `anthropic.claude-3-opus-20240229-v1:0` +- `eu.anthropic.claude-haiku-4-5-20251001-v1:0` (EU cross-region inference profile) +- `us.anthropic.claude-3-5-sonnet-20241022-v2:0` (US inference profile) +- `global.anthropic.claude-haiku-4-5-20251001-v1:0` -Other Bedrock vendors (`amazon.titan-...`, `meta.llama3-...`, `cohere.command-...`, `ai21.j2-...`, `stability.stable-diffusion-...`) are signed and forwarded **without compression** — the proxy does not yet understand their body shapes and would risk corrupting them. These model IDs log `event=bedrock_compression_skipped, reason=non_anthropic_vendor` per request. Full Anthropic envelopes only. +Other Bedrock vendors (`amazon.titan-...`, `meta.llama3-...`, `cohere.command-...`, `ai21.j2-...`, `stability.stable-diffusion-...`, and their geo-prefixed inference profiles such as `eu.amazon.nova-...`) are signed and forwarded **without compression** — the proxy does not yet understand their body shapes and would risk corrupting them. These model IDs log `event=bedrock_compression_skipped, reason=non_anthropic_vendor` per request. Full Anthropic envelopes only. -The contract: **any new model ID that AWS adds under the `anthropic.` prefix automatically picks up the full compression + signing pipeline.** No code change in the proxy is needed for new versions of Claude on Bedrock. +The contract: **any new model ID that AWS adds under the `anthropic.` vendor (as a bare prefix or behind a cross-region geo prefix) automatically picks up the full compression + signing pipeline.** No code change in the proxy is needed for new versions of Claude on Bedrock. ## Compression behaviour