mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat(bedrock): cross-region + Converse compression; bundle proxy binary in images (#999)
## 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.
This commit is contained in:
Yasser Sheikh
GitHub
parent
0d4571f72f
commit
0dc2e1cb3f
10 changed files with 435 additions and 65 deletions
|
|
@ -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`.
|
||||
|
|
|
|||
53
Cargo.lock
generated
53
Cargo.lock
generated
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
15
Dockerfile
15
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
|
||||
|
|
|
|||
|
|
@ -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::<Probe<'_>>(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!({
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -74,8 +74,7 @@ use crate::proxy::AppState;
|
|||
// `Extension<AuthMode>` 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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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};
|
||||
|
|
|
|||
82
crates/headroom-proxy/src/bedrock/vendor.rs
Normal file
82
crates/headroom-proxy/src/bedrock/vendor.rs
Normal file
|
|
@ -0,0 +1,82 @@
|
|||
//! Canonical Bedrock vendor resolution.
|
||||
//!
|
||||
//! Bedrock model IDs are `<vendor>.<model>-<date>-<rev>` (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
|
||||
/// (`<geo>.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"));
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue