mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy): PR-D1 native Bedrock InvokeModel route + SigV4
Phase D PR-D1 lands the first native Rust path for AWS Bedrock,
replacing the lossy LiteLLM Python converter for Anthropic-on-Bedrock
non-streaming requests. Eliminates part of P4-37 and P4-39.
What landed
-----------
- New crates/headroom-proxy/src/bedrock/ module:
- envelope.rs: parses the {"anthropic_version": "...", ...}
Bedrock body shape; re-emits with anthropic_version preserved
as the first key (relies on serde_json preserve_order).
- sigv4.rs: AWS SigV4 signer wrapping the aws-sigv4 crate.
Forces PayloadChecksumKind::XAmzSha256 so x-amz-content-sha256
is in the canonical request, hashed over the post-compression
body bytes (the bytes that actually hit Bedrock). No silent
fallback: signing failures return 5xx with
event=bedrock_sigv4_failed.
- invoke.rs: POST handler for /model/{model_id}/invoke
(and /converse - same wire shape for anthropic.claude-*).
Detects Anthropic vendor via literal starts_with("anthropic.")
(no regex per project rule), routes Anthropic-shape bodies
through the existing compress_anthropic_request live-zone
dispatcher, then signs and forwards to the configured Bedrock
endpoint.
- Modified:
- proxy.rs: routes /model/:model_id/invoke and
/model/:model_id/converse when enable_bedrock_native is on
(default). Adds bedrock_credentials: Option<Arc<Credentials>>
to AppState.
- config.rs: new flags --bedrock-region (default us-east-1,
env HEADROOM_PROXY_BEDROCK_REGION), --bedrock-endpoint
(operator override for FIPS/VPC/test setups),
--enable-bedrock-native (default true), --aws-profile.
- main.rs: resolves AWS credentials at startup via
aws_config::defaults(BehaviorVersion::latest()). Failure logs
event=bedrock_credentials_unavailable at WARN; the handler
refuses to forward unsigned (event=bedrock_credentials_missing).
- Cargo.toml: workspace deps aws-sigv4, aws-config,
aws-credential-types, aws-smithy-runtime-api.
Tests
-----
8 integration tests under
crates/headroom-proxy/tests/integration_bedrock_invoke.rs:
1. native_envelope_round_trip_byte_equal
2. sigv4_signed_correctly_after_compression - confirms
authorization is SigV4-shape and x-amz-content-sha256
matches sha256(body received by upstream).
3. thinking_block_preserved_through_bedrock
4. redacted_thinking_preserved
5. document_block_preserved
6. tool_result_array_with_image_preserved
7. stop_sequence_null_only_when_present - pins that the proxy
does NOT inject stop_sequence: null (P4-37 hardcode).
8. tool_use_input_byte_equal_preserves_key_order
All eight pass. Full workspace test run is green; clippy + fmt
clean. make ci-precheck (rust + python + commitlint) passes
locally.
Build constraints honoured
--------------------------
- No silent fallbacks: missing creds / signing failures return
5xx with structured event=... log; no path ever forwards
unsigned.
- No hardcodes: region, endpoint, profile, enable-flag all
configurable via CLI + env.
- No regexes: vendor detection is str::starts_with.
- Comprehensive structured logs: event=bedrock_invoke_received,
bedrock_envelope_parsed, bedrock_compression_skipped,
bedrock_credentials_missing, sigv4_signed,
bedrock_invoke_forwarded, etc.
- Performant: body buffered once, passed by &[u8] to signer
(zero-copy), Bytes::clone only for ownership transfer to
reqwest. Sign exactly once per request.
- Elegant: 4 small focused modules mirror handlers/ + sse/.
- Tests use realistic Anthropic block content (real thinking,
redacted_thinking, document, base64 image fixtures).
Streaming (PR-D2) and observability (PR-D3) follow.
This commit is contained in:
chopratejas
parent
1352f621fc
commit
f2d4fe39cb
13 changed files with 2509 additions and 31 deletions
588
Cargo.lock
generated
588
Cargo.lock
generated
|
|
@ -240,6 +240,321 @@ dependencies = [
|
|||
"arrayvec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-config"
|
||||
version = "1.8.16"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "50f156acdd2cf55f5aa53ee416c4ac851cf1222694506c0b1f78c85695e9ca9d"
|
||||
dependencies = [
|
||||
"aws-credential-types",
|
||||
"aws-runtime",
|
||||
"aws-sdk-sts",
|
||||
"aws-smithy-async",
|
||||
"aws-smithy-http",
|
||||
"aws-smithy-json",
|
||||
"aws-smithy-runtime",
|
||||
"aws-smithy-runtime-api",
|
||||
"aws-smithy-types",
|
||||
"aws-types",
|
||||
"bytes",
|
||||
"fastrand",
|
||||
"http 1.4.0",
|
||||
"time",
|
||||
"tokio",
|
||||
"tracing",
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-credential-types"
|
||||
version = "1.2.14"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8f20799b373a1be121fe3005fba0c2090af9411573878f224df44b42727fcaf7"
|
||||
dependencies = [
|
||||
"aws-smithy-async",
|
||||
"aws-smithy-runtime-api",
|
||||
"aws-smithy-types",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-lc-rs"
|
||||
version = "1.16.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0ec6fb3fe69024a75fa7e1bfb48aa6cf59706a101658ea01bfd33b2b248a038f"
|
||||
dependencies = [
|
||||
"aws-lc-sys",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-lc-sys"
|
||||
version = "0.40.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f50037ee5e1e41e7b8f9d161680a725bd1626cb6f8c7e901f91f942850852fe7"
|
||||
dependencies = [
|
||||
"cc",
|
||||
"cmake",
|
||||
"dunce",
|
||||
"fs_extra",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-runtime"
|
||||
version = "1.7.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5dcd93c82209ac7413532388067dce79be5a8780c1786e5fae3df22e4dee2864"
|
||||
dependencies = [
|
||||
"aws-credential-types",
|
||||
"aws-sigv4",
|
||||
"aws-smithy-async",
|
||||
"aws-smithy-http",
|
||||
"aws-smithy-runtime",
|
||||
"aws-smithy-runtime-api",
|
||||
"aws-smithy-types",
|
||||
"aws-types",
|
||||
"bytes",
|
||||
"bytes-utils",
|
||||
"fastrand",
|
||||
"http 1.4.0",
|
||||
"http-body 1.0.1",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"tracing",
|
||||
"uuid",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-sdk-sts"
|
||||
version = "1.103.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c2249b81a2e73a8027c41c378463a81ec39b8510f184f2caab87de912af0f49b"
|
||||
dependencies = [
|
||||
"aws-credential-types",
|
||||
"aws-runtime",
|
||||
"aws-smithy-async",
|
||||
"aws-smithy-http",
|
||||
"aws-smithy-json",
|
||||
"aws-smithy-observability",
|
||||
"aws-smithy-query",
|
||||
"aws-smithy-runtime",
|
||||
"aws-smithy-runtime-api",
|
||||
"aws-smithy-types",
|
||||
"aws-smithy-xml",
|
||||
"aws-types",
|
||||
"fastrand",
|
||||
"http 0.2.12",
|
||||
"http 1.4.0",
|
||||
"regex-lite",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-sigv4"
|
||||
version = "1.4.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "68dc0b907359b120170613b5c09ccc61304eac3998ff6274b97d93ee6490115a"
|
||||
dependencies = [
|
||||
"aws-credential-types",
|
||||
"aws-smithy-http",
|
||||
"aws-smithy-runtime-api",
|
||||
"aws-smithy-types",
|
||||
"bytes",
|
||||
"form_urlencoded",
|
||||
"hex",
|
||||
"hmac",
|
||||
"http 0.2.12",
|
||||
"http 1.4.0",
|
||||
"percent-encoding",
|
||||
"sha2 0.11.0",
|
||||
"time",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-smithy-async"
|
||||
version = "1.2.14"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2ffcaf626bdda484571968400c326a244598634dc75fd451325a54ad1a59acfc"
|
||||
dependencies = [
|
||||
"futures-util",
|
||||
"pin-project-lite",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-smithy-http"
|
||||
version = "0.63.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ba1ab2dc1c2c3749ead27180d333c42f11be8b0e934058fb4b2258ee8dbe5231"
|
||||
dependencies = [
|
||||
"aws-smithy-runtime-api",
|
||||
"aws-smithy-types",
|
||||
"bytes",
|
||||
"bytes-utils",
|
||||
"futures-core",
|
||||
"futures-util",
|
||||
"http 1.4.0",
|
||||
"http-body 1.0.1",
|
||||
"http-body-util",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"pin-utils",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-smithy-http-client"
|
||||
version = "1.1.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6a2f165a7feee6f263028b899d0a181987f4fa7179a6411a32a439fba7c5f769"
|
||||
dependencies = [
|
||||
"aws-smithy-async",
|
||||
"aws-smithy-runtime-api",
|
||||
"aws-smithy-types",
|
||||
"h2",
|
||||
"http 1.4.0",
|
||||
"hyper",
|
||||
"hyper-rustls",
|
||||
"hyper-util",
|
||||
"pin-project-lite",
|
||||
"rustls",
|
||||
"rustls-native-certs",
|
||||
"rustls-pki-types",
|
||||
"tokio",
|
||||
"tokio-rustls",
|
||||
"tower",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-smithy-json"
|
||||
version = "0.62.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9648b0bb82a2eedd844052c6ad2a1a822d1f8e3adee5fbf668366717e428856a"
|
||||
dependencies = [
|
||||
"aws-smithy-types",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-smithy-observability"
|
||||
version = "0.2.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a06c2315d173edbf1920da8ba3a7189695827002e4c0fc961973ab1c54abca9c"
|
||||
dependencies = [
|
||||
"aws-smithy-runtime-api",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-smithy-query"
|
||||
version = "0.60.15"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1a56d79744fb3edb5d722ef79d86081e121d3b9422cb209eb03aea6aa4f21ebd"
|
||||
dependencies = [
|
||||
"aws-smithy-types",
|
||||
"urlencoding",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-smithy-runtime"
|
||||
version = "1.11.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0504b1ab12debb5959e5165ee5fe97dd387e7aa7ea6a477bfd7635dfe769a4f5"
|
||||
dependencies = [
|
||||
"aws-smithy-async",
|
||||
"aws-smithy-http",
|
||||
"aws-smithy-http-client",
|
||||
"aws-smithy-observability",
|
||||
"aws-smithy-runtime-api",
|
||||
"aws-smithy-types",
|
||||
"bytes",
|
||||
"fastrand",
|
||||
"http 0.2.12",
|
||||
"http 1.4.0",
|
||||
"http-body 0.4.6",
|
||||
"http-body 1.0.1",
|
||||
"http-body-util",
|
||||
"pin-project-lite",
|
||||
"pin-utils",
|
||||
"tokio",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-smithy-runtime-api"
|
||||
version = "1.12.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b71a13df6ada0aafbf21a73bdfcdf9324cfa9df77d96b8446045be3cde61b42e"
|
||||
dependencies = [
|
||||
"aws-smithy-async",
|
||||
"aws-smithy-runtime-api-macros",
|
||||
"aws-smithy-types",
|
||||
"bytes",
|
||||
"http 0.2.12",
|
||||
"http 1.4.0",
|
||||
"pin-project-lite",
|
||||
"tokio",
|
||||
"tracing",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-smithy-runtime-api-macros"
|
||||
version = "1.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8d7396fd9500589e62e460e987ecb671bad374934e55ec3b5f498cc7a8a8a7b7"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-smithy-types"
|
||||
version = "1.4.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9d73dbfbaa8e4bc57b9045137680b958d274823509a360abfd8e1d514d40c95c"
|
||||
dependencies = [
|
||||
"base64-simd",
|
||||
"bytes",
|
||||
"bytes-utils",
|
||||
"http 0.2.12",
|
||||
"http 1.4.0",
|
||||
"http-body 0.4.6",
|
||||
"http-body 1.0.1",
|
||||
"http-body-util",
|
||||
"itoa",
|
||||
"num-integer",
|
||||
"pin-project-lite",
|
||||
"pin-utils",
|
||||
"ryu",
|
||||
"serde",
|
||||
"time",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-smithy-xml"
|
||||
version = "0.60.15"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0ce02add1aa3677d022f8adf81dcbe3046a95f17a1b1e8979c145cd21d3d22b3"
|
||||
dependencies = [
|
||||
"xmlparser",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-types"
|
||||
version = "1.3.15"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2f4bbcaa9304ea40902d3d5f42a0428d1bd895a2b0f6999436fb279ffddc58ac"
|
||||
dependencies = [
|
||||
"aws-credential-types",
|
||||
"aws-smithy-async",
|
||||
"aws-smithy-runtime-api",
|
||||
"aws-smithy-types",
|
||||
"rustc_version",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "axum"
|
||||
version = "0.7.9"
|
||||
|
|
@ -252,8 +567,8 @@ dependencies = [
|
|||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"http",
|
||||
"http-body",
|
||||
"http 1.4.0",
|
||||
"http-body 1.0.1",
|
||||
"http-body-util",
|
||||
"hyper",
|
||||
"hyper-util",
|
||||
|
|
@ -287,8 +602,8 @@ dependencies = [
|
|||
"async-trait",
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"http",
|
||||
"http-body",
|
||||
"http 1.4.0",
|
||||
"http-body 1.0.1",
|
||||
"http-body-util",
|
||||
"mime",
|
||||
"pin-project-lite",
|
||||
|
|
@ -322,6 +637,16 @@ version = "0.22.1"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6"
|
||||
|
||||
[[package]]
|
||||
name = "base64-simd"
|
||||
version = "0.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "339abbe78e73178762e23bea9dfd08e697eb3f3301cd4be981c0f78ba5859195"
|
||||
dependencies = [
|
||||
"outref",
|
||||
"vsimd",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "base64ct"
|
||||
version = "1.8.3"
|
||||
|
|
@ -387,6 +712,15 @@ dependencies = [
|
|||
"generic-array",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "block-buffer"
|
||||
version = "0.12.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cdd35008169921d80bc60d3d0ab416eecb028c4cd653352907921d95084790be"
|
||||
dependencies = [
|
||||
"hybrid-array",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bstr"
|
||||
version = "1.12.1"
|
||||
|
|
@ -434,6 +768,16 @@ version = "1.11.1"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33"
|
||||
|
||||
[[package]]
|
||||
name = "bytes-utils"
|
||||
version = "0.1.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7dafe3a8757b027e2be6e4e5601ed563c55989fcf1546e933c66c8eb3a058d35"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"either",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bytesize"
|
||||
version = "1.3.3"
|
||||
|
|
@ -546,6 +890,21 @@ version = "1.1.0"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9"
|
||||
|
||||
[[package]]
|
||||
name = "cmake"
|
||||
version = "0.1.58"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c0f78a02292a74a88ac736019ab962ece0bc380e3f977bf72e376c5d78ff0678"
|
||||
dependencies = [
|
||||
"cc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cmov"
|
||||
version = "0.5.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3f88a43d011fc4a6876cb7344703e297c71dda42494fee094d5f7c76bf13f746"
|
||||
|
||||
[[package]]
|
||||
name = "color_quant"
|
||||
version = "1.1.0"
|
||||
|
|
@ -608,6 +967,12 @@ dependencies = [
|
|||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "const-oid"
|
||||
version = "0.10.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c"
|
||||
|
||||
[[package]]
|
||||
name = "constant_time_eq"
|
||||
version = "0.4.2"
|
||||
|
|
@ -773,6 +1138,24 @@ dependencies = [
|
|||
"typenum",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crypto-common"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "77727bb15fa921304124b128af125e7e3b968275d1b108b379190264f4423710"
|
||||
dependencies = [
|
||||
"hybrid-array",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ctutils"
|
||||
version = "0.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7d5515a3834141de9eafb9717ad39eea8247b5674e6066c404e8c4b365d2a29e"
|
||||
dependencies = [
|
||||
"cmov",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "darling"
|
||||
version = "0.20.11"
|
||||
|
|
@ -911,8 +1294,20 @@ version = "0.10.7"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292"
|
||||
dependencies = [
|
||||
"block-buffer",
|
||||
"crypto-common",
|
||||
"block-buffer 0.10.4",
|
||||
"crypto-common 0.1.7",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "digest"
|
||||
version = "0.11.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2"
|
||||
dependencies = [
|
||||
"block-buffer 0.12.0",
|
||||
"const-oid",
|
||||
"crypto-common 0.2.1",
|
||||
"ctutils",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -956,6 +1351,12 @@ dependencies = [
|
|||
"litrs",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "dunce"
|
||||
version = "1.0.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813"
|
||||
|
||||
[[package]]
|
||||
name = "either"
|
||||
version = "1.15.0"
|
||||
|
|
@ -1170,6 +1571,12 @@ dependencies = [
|
|||
"percent-encoding",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fs_extra"
|
||||
version = "1.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
|
||||
|
||||
[[package]]
|
||||
name = "futures"
|
||||
version = "0.3.32"
|
||||
|
|
@ -1329,7 +1736,7 @@ dependencies = [
|
|||
"fnv",
|
||||
"futures-core",
|
||||
"futures-sink",
|
||||
"http",
|
||||
"http 1.4.0",
|
||||
"indexmap",
|
||||
"slab",
|
||||
"tokio",
|
||||
|
|
@ -1415,7 +1822,7 @@ dependencies = [
|
|||
"rusqlite",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2",
|
||||
"sha2 0.10.9",
|
||||
"tempfile",
|
||||
"thiserror 1.0.69",
|
||||
"tiktoken-rs",
|
||||
|
|
@ -1441,6 +1848,10 @@ dependencies = [
|
|||
name = "headroom-proxy"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"aws-config",
|
||||
"aws-credential-types",
|
||||
"aws-sigv4",
|
||||
"aws-smithy-runtime-api",
|
||||
"axum",
|
||||
"bytes",
|
||||
"bytesize",
|
||||
|
|
@ -1448,7 +1859,7 @@ dependencies = [
|
|||
"futures",
|
||||
"futures-util",
|
||||
"headroom-core",
|
||||
"http",
|
||||
"http 1.4.0",
|
||||
"http-body-util",
|
||||
"humantime",
|
||||
"hyper",
|
||||
|
|
@ -1458,7 +1869,7 @@ dependencies = [
|
|||
"reqwest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2",
|
||||
"sha2 0.10.9",
|
||||
"thiserror 1.0.69",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
|
|
@ -1494,6 +1905,12 @@ version = "0.5.2"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c"
|
||||
|
||||
[[package]]
|
||||
name = "hex"
|
||||
version = "0.4.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70"
|
||||
|
||||
[[package]]
|
||||
name = "hf-hub"
|
||||
version = "0.4.3"
|
||||
|
|
@ -1501,7 +1918,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "629d8f3bbeda9d148036d6b0de0a3ab947abd08ce90626327fc3547a49d59d97"
|
||||
dependencies = [
|
||||
"dirs",
|
||||
"http",
|
||||
"http 1.4.0",
|
||||
"indicatif 0.17.11",
|
||||
"libc",
|
||||
"log",
|
||||
|
|
@ -1521,7 +1938,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "aef3982638978efa195ff11b305f51f1f22f4f0a6cabee7af79b383ebee6a213"
|
||||
dependencies = [
|
||||
"dirs",
|
||||
"http",
|
||||
"http 1.4.0",
|
||||
"indicatif 0.18.4",
|
||||
"libc",
|
||||
"log",
|
||||
|
|
@ -1535,12 +1952,32 @@ dependencies = [
|
|||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hmac"
|
||||
version = "0.13.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6303bc9732ae41b04cb554b844a762b4115a61bfaa81e3e83050991eeb56863f"
|
||||
dependencies = [
|
||||
"digest 0.11.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hmac-sha256"
|
||||
version = "1.1.14"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ec9d92d097f4749b64e8cc33d924d9f40a2d4eb91402b458014b781f5733d60f"
|
||||
|
||||
[[package]]
|
||||
name = "http"
|
||||
version = "0.2.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "601cbb57e577e2f5ef5be8e7b83f0f63994f25aa94d673e54a92d5c516d101f1"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"fnv",
|
||||
"itoa",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "http"
|
||||
version = "1.4.0"
|
||||
|
|
@ -1551,6 +1988,17 @@ dependencies = [
|
|||
"itoa",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "http-body"
|
||||
version = "0.4.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7ceab25649e9960c0311ea418d17bee82c0dcec1bd053b5f9a66e265a693bed2"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"http 0.2.12",
|
||||
"pin-project-lite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "http-body"
|
||||
version = "1.0.1"
|
||||
|
|
@ -1558,7 +2006,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"http",
|
||||
"http 1.4.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -1569,8 +2017,8 @@ checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a"
|
|||
dependencies = [
|
||||
"bytes",
|
||||
"futures-core",
|
||||
"http",
|
||||
"http-body",
|
||||
"http 1.4.0",
|
||||
"http-body 1.0.1",
|
||||
"pin-project-lite",
|
||||
]
|
||||
|
||||
|
|
@ -1592,6 +2040,15 @@ version = "2.3.0"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "135b12329e5e3ce057a9f972339ea52bc954fe1e9358ef27f95e89716fbc5424"
|
||||
|
||||
[[package]]
|
||||
name = "hybrid-array"
|
||||
version = "0.4.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "08d46837a0ed51fe95bd3b05de33cd64a1ee88fc797477ca48446872504507c5"
|
||||
dependencies = [
|
||||
"typenum",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hyper"
|
||||
version = "1.9.0"
|
||||
|
|
@ -1603,8 +2060,8 @@ dependencies = [
|
|||
"futures-channel",
|
||||
"futures-core",
|
||||
"h2",
|
||||
"http",
|
||||
"http-body",
|
||||
"http 1.4.0",
|
||||
"http-body 1.0.1",
|
||||
"httparse",
|
||||
"httpdate",
|
||||
"itoa",
|
||||
|
|
@ -1620,10 +2077,11 @@ version = "0.27.9"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f"
|
||||
dependencies = [
|
||||
"http",
|
||||
"http 1.4.0",
|
||||
"hyper",
|
||||
"hyper-util",
|
||||
"rustls",
|
||||
"rustls-native-certs",
|
||||
"tokio",
|
||||
"tokio-rustls",
|
||||
"tower-service",
|
||||
|
|
@ -1656,8 +2114,8 @@ dependencies = [
|
|||
"bytes",
|
||||
"futures-channel",
|
||||
"futures-util",
|
||||
"http",
|
||||
"http-body",
|
||||
"http 1.4.0",
|
||||
"http-body 1.0.1",
|
||||
"hyper",
|
||||
"ipnet",
|
||||
"libc",
|
||||
|
|
@ -2150,7 +2608,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "d89e7ee0cfbedfc4da3340218492196241d89eefb6dab27de5df917a6d2e78cf"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"digest",
|
||||
"digest 0.10.7",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -2509,6 +2967,12 @@ dependencies = [
|
|||
"ureq 3.3.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "outref"
|
||||
version = "0.5.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e"
|
||||
|
||||
[[package]]
|
||||
name = "parking_lot_core"
|
||||
version = "0.9.12"
|
||||
|
|
@ -2555,6 +3019,12 @@ version = "0.2.17"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd"
|
||||
|
||||
[[package]]
|
||||
name = "pin-utils"
|
||||
version = "0.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184"
|
||||
|
||||
[[package]]
|
||||
name = "pkg-config"
|
||||
version = "0.3.33"
|
||||
|
|
@ -3078,6 +3548,12 @@ dependencies = [
|
|||
"regex-syntax",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "regex-lite"
|
||||
version = "0.1.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973"
|
||||
|
||||
[[package]]
|
||||
name = "regex-syntax"
|
||||
version = "0.8.10"
|
||||
|
|
@ -3096,8 +3572,8 @@ dependencies = [
|
|||
"futures-core",
|
||||
"futures-util",
|
||||
"h2",
|
||||
"http",
|
||||
"http-body",
|
||||
"http 1.4.0",
|
||||
"http-body 1.0.1",
|
||||
"http-body-util",
|
||||
"hyper",
|
||||
"hyper-rustls",
|
||||
|
|
@ -3177,6 +3653,15 @@ version = "2.1.2"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe"
|
||||
|
||||
[[package]]
|
||||
name = "rustc_version"
|
||||
version = "0.4.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92"
|
||||
dependencies = [
|
||||
"semver",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustix"
|
||||
version = "1.1.4"
|
||||
|
|
@ -3196,6 +3681,7 @@ version = "0.23.39"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7c2c118cb077cca2822033836dfb1b975355dfb784b5e8da48f7b6c5db74e60e"
|
||||
dependencies = [
|
||||
"aws-lc-rs",
|
||||
"log",
|
||||
"once_cell",
|
||||
"ring",
|
||||
|
|
@ -3205,6 +3691,18 @@ dependencies = [
|
|||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustls-native-certs"
|
||||
version = "0.8.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "612460d5f7bea540c490b2b6395d8e34a953e52b491accd6c86c8164c5932a63"
|
||||
dependencies = [
|
||||
"openssl-probe",
|
||||
"rustls-pki-types",
|
||||
"schannel",
|
||||
"security-framework",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustls-pki-types"
|
||||
version = "1.14.0"
|
||||
|
|
@ -3221,6 +3719,7 @@ version = "0.103.13"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e"
|
||||
dependencies = [
|
||||
"aws-lc-rs",
|
||||
"ring",
|
||||
"rustls-pki-types",
|
||||
"untrusted",
|
||||
|
|
@ -3398,7 +3897,7 @@ checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba"
|
|||
dependencies = [
|
||||
"cfg-if",
|
||||
"cpufeatures 0.2.17",
|
||||
"digest",
|
||||
"digest 0.10.7",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -3409,7 +3908,18 @@ checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283"
|
|||
dependencies = [
|
||||
"cfg-if",
|
||||
"cpufeatures 0.2.17",
|
||||
"digest",
|
||||
"digest 0.10.7",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sha2"
|
||||
version = "0.11.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"cpufeatures 0.3.0",
|
||||
"digest 0.11.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -3923,8 +4433,8 @@ dependencies = [
|
|||
"bitflags",
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"http",
|
||||
"http-body",
|
||||
"http 1.4.0",
|
||||
"http-body 1.0.1",
|
||||
"iri-string",
|
||||
"pin-project-lite",
|
||||
"tower",
|
||||
|
|
@ -4036,7 +4546,7 @@ dependencies = [
|
|||
"byteorder",
|
||||
"bytes",
|
||||
"data-encoding",
|
||||
"http",
|
||||
"http 1.4.0",
|
||||
"httparse",
|
||||
"log",
|
||||
"rand 0.8.6",
|
||||
|
|
@ -4176,7 +4686,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "e994ba84b0bd1b1b0cf92878b7ef898a5c1760108fe7b6010327e274917a808c"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"http",
|
||||
"http 1.4.0",
|
||||
"httparse",
|
||||
"log",
|
||||
]
|
||||
|
|
@ -4193,6 +4703,12 @@ dependencies = [
|
|||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "urlencoding"
|
||||
version = "2.1.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "daf8dba3b7eb870caf1ddeed7bc9d2a049f3cfdfae7cb521b087cc33ae4c49da"
|
||||
|
||||
[[package]]
|
||||
name = "utf-8"
|
||||
version = "0.7.6"
|
||||
|
|
@ -4257,6 +4773,12 @@ version = "0.9.5"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
|
||||
|
||||
[[package]]
|
||||
name = "vsimd"
|
||||
version = "0.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64"
|
||||
|
||||
[[package]]
|
||||
name = "wait-timeout"
|
||||
version = "0.2.1"
|
||||
|
|
@ -4714,7 +5236,7 @@ dependencies = [
|
|||
"base64 0.22.1",
|
||||
"deadpool",
|
||||
"futures",
|
||||
"http",
|
||||
"http 1.4.0",
|
||||
"http-body-util",
|
||||
"hyper",
|
||||
"hyper-util",
|
||||
|
|
@ -4827,6 +5349,12 @@ version = "0.6.3"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4"
|
||||
|
||||
[[package]]
|
||||
name = "xmlparser"
|
||||
version = "0.13.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4"
|
||||
|
||||
[[package]]
|
||||
name = "y4m"
|
||||
version = "0.8.0"
|
||||
|
|
|
|||
12
Cargo.toml
12
Cargo.toml
|
|
@ -56,3 +56,15 @@ axum = "0.7"
|
|||
tower = "0.5"
|
||||
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] }
|
||||
pyo3 = "0.22"
|
||||
# Phase D PR-D1: AWS SigV4 signing for native Bedrock InvokeModel route.
|
||||
# `aws-sigv4` provides the canonical-request + signing-key implementation;
|
||||
# `aws-config` resolves credentials from the standard provider chain
|
||||
# (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-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
|
||||
# silent semver bump from the transitive dep tree.
|
||||
aws-smithy-runtime-api = { version = "1", default-features = false, features = ["client"] }
|
||||
|
|
|
|||
|
|
@ -41,6 +41,12 @@ humantime = "2"
|
|||
bytesize = "1"
|
||||
tokio-util = { version = "0.7" }
|
||||
headroom-core = { path = "../headroom-core" }
|
||||
# Phase D PR-D1: native Bedrock InvokeModel route. SigV4 + AWS
|
||||
# default credential chain.
|
||||
aws-sigv4 = { workspace = true }
|
||||
aws-config = { workspace = true }
|
||||
aws-credential-types = { workspace = true }
|
||||
aws-smithy-runtime-api = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
tower = { workspace = true, features = ["util"] }
|
||||
|
|
|
|||
239
crates/headroom-proxy/src/bedrock/envelope.rs
Normal file
239
crates/headroom-proxy/src/bedrock/envelope.rs
Normal file
|
|
@ -0,0 +1,239 @@
|
|||
//! Bedrock envelope: parse the body shape Bedrock expects, hand the
|
||||
//! Anthropic-shape sub-body to the compressor, and re-emit with
|
||||
//! `anthropic_version` preserved as the FIRST key.
|
||||
//!
|
||||
//! # Wire shape
|
||||
//!
|
||||
//! Bedrock InvokeModel for any `anthropic.claude-*` model expects:
|
||||
//!
|
||||
//! ```json
|
||||
//! {
|
||||
//! "anthropic_version": "bedrock-2023-05-31",
|
||||
//! "messages": [...],
|
||||
//! "max_tokens": 1024,
|
||||
//! ...rest_of_anthropic_body
|
||||
//! }
|
||||
//! ```
|
||||
//!
|
||||
//! `model` is in the URL path (`/model/{model}/invoke`), NOT the body
|
||||
//! — the opposite of direct-Anthropic `/v1/messages`. `anthropic_version`
|
||||
//! is REQUIRED and must be a literal Bedrock-recognized version
|
||||
//! string (e.g. `"bedrock-2023-05-31"`).
|
||||
//!
|
||||
//! # Why key order matters
|
||||
//!
|
||||
//! Bedrock's request validator does NOT depend on key order, but our
|
||||
//! cache-safety contract (`I1` in REALIGNMENT/02-architecture.md) is
|
||||
//! that *unmodified* bytes round-trip byte-equal. If the compressor
|
||||
//! returns `NoChange` we forward the original buffered bytes
|
||||
//! verbatim — we never re-serialize. If the compressor returns
|
||||
//! `Modified`, the byte slice it produced already preserves
|
||||
//! key order from the input via the `preserve_order` feature on
|
||||
//! `serde_json` (the workspace turns this on by default).
|
||||
//!
|
||||
//! This module's [`BedrockEnvelope`] is used ONLY for re-emitting
|
||||
//! the envelope shape after compression. It calls the compressor
|
||||
//! against the body bytes directly — there is no decode-then-encode
|
||||
//! round trip on the no-change path.
|
||||
|
||||
use bytes::Bytes;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use thiserror::Error;
|
||||
|
||||
/// The literal `anthropic_version` field name. Bedrock-strict.
|
||||
const ANTHROPIC_VERSION_KEY: &str = "anthropic_version";
|
||||
|
||||
/// Parsed Bedrock InvokeModel envelope. Holds the literal
|
||||
/// `anthropic_version` string plus the rest of the body as a parsed
|
||||
/// `serde_json::Value` (so callers can inspect / route it). The
|
||||
/// original byte-equal body is also stashed for the no-change
|
||||
/// passthrough path.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BedrockEnvelope {
|
||||
/// The literal `anthropic_version` string, e.g.
|
||||
/// `"bedrock-2023-05-31"`. Required.
|
||||
pub anthropic_version: String,
|
||||
/// The full parsed body, including `anthropic_version`. Useful
|
||||
/// for tests + logs; the compressor takes raw bytes.
|
||||
pub body: Value,
|
||||
}
|
||||
|
||||
/// Errors surfaced when parsing a Bedrock envelope.
|
||||
#[derive(Debug, Error)]
|
||||
pub enum EnvelopeError {
|
||||
/// Body was not valid JSON.
|
||||
#[error("body is not valid JSON: {0}")]
|
||||
NotJson(serde_json::Error),
|
||||
/// Body was JSON but not a top-level object.
|
||||
#[error("body is not a JSON object")]
|
||||
NotObject,
|
||||
/// `anthropic_version` field missing.
|
||||
#[error("missing required `anthropic_version` field")]
|
||||
MissingAnthropicVersion,
|
||||
/// `anthropic_version` was present but not a string.
|
||||
#[error("`anthropic_version` is not a string")]
|
||||
AnthropicVersionNotString,
|
||||
}
|
||||
|
||||
impl BedrockEnvelope {
|
||||
/// Parse a Bedrock envelope from raw JSON bytes.
|
||||
///
|
||||
/// This does NOT mutate the bytes. Compression dispatch happens
|
||||
/// elsewhere (the handler hands the same byte slice to
|
||||
/// `compress_anthropic_request`).
|
||||
pub fn parse(body: &[u8]) -> Result<Self, EnvelopeError> {
|
||||
let value: Value = serde_json::from_slice(body).map_err(EnvelopeError::NotJson)?;
|
||||
let obj = value.as_object().ok_or(EnvelopeError::NotObject)?;
|
||||
let av = obj
|
||||
.get(ANTHROPIC_VERSION_KEY)
|
||||
.ok_or(EnvelopeError::MissingAnthropicVersion)?;
|
||||
let av_str = av
|
||||
.as_str()
|
||||
.ok_or(EnvelopeError::AnthropicVersionNotString)?
|
||||
.to_string();
|
||||
Ok(Self {
|
||||
anthropic_version: av_str,
|
||||
body: value,
|
||||
})
|
||||
}
|
||||
|
||||
/// Re-emit the envelope as JSON bytes with `anthropic_version`
|
||||
/// preserved as the FIRST key.
|
||||
///
|
||||
/// Used only after the compressor returns `Modified` and we need
|
||||
/// to reassemble the body. With `serde_json`'s `preserve_order`
|
||||
/// feature (workspace default), the parsed object already preserves
|
||||
/// insertion order; we just need to make sure `anthropic_version`
|
||||
/// is the first key. If the compressor's output already has it
|
||||
/// first (which it will — the compressor only mutates the
|
||||
/// `messages` content slot, not key ordering), we hand the bytes
|
||||
/// back unchanged.
|
||||
///
|
||||
/// `body` is the (possibly-compressed) JSON bytes returned by the
|
||||
/// compression dispatcher.
|
||||
///
|
||||
/// Returns `Ok(bytes)` on success. On any structural error, returns
|
||||
/// the input bytes unchanged so the byte-fidelity contract is
|
||||
/// preserved (the caller will surface the error).
|
||||
pub fn ensure_anthropic_version_first(body: &[u8]) -> Result<Bytes, EnvelopeError> {
|
||||
// Parse with preserve_order (workspace serde_json default).
|
||||
let mut value: Value = serde_json::from_slice(body).map_err(EnvelopeError::NotJson)?;
|
||||
let map = value.as_object_mut().ok_or(EnvelopeError::NotObject)?;
|
||||
|
||||
// If `anthropic_version` is missing, that's a logical error
|
||||
// for the Bedrock surface — surface it loudly.
|
||||
if !map.contains_key(ANTHROPIC_VERSION_KEY) {
|
||||
return Err(EnvelopeError::MissingAnthropicVersion);
|
||||
}
|
||||
|
||||
// Already first? Then `body` round-trips byte-equal — no work.
|
||||
// Note: serde_json::Map iteration order, with preserve_order,
|
||||
// is insertion order. We check the first key directly.
|
||||
if map
|
||||
.keys()
|
||||
.next()
|
||||
.map(|k| k.as_str() == ANTHROPIC_VERSION_KEY)
|
||||
.unwrap_or(false)
|
||||
{
|
||||
return Ok(Bytes::copy_from_slice(body));
|
||||
}
|
||||
|
||||
// Reorder: pull anthropic_version, drain rest, rebuild with
|
||||
// anthropic_version first. This only runs when the compressor
|
||||
// moved the field (it shouldn't, but if it ever does we
|
||||
// reassert the invariant rather than ship a Bedrock-rejecting
|
||||
// body).
|
||||
let av = map
|
||||
.remove(ANTHROPIC_VERSION_KEY)
|
||||
.expect("contains_key guard above");
|
||||
let mut new_map = serde_json::Map::with_capacity(map.len() + 1);
|
||||
new_map.insert(ANTHROPIC_VERSION_KEY.to_string(), av);
|
||||
for (k, v) in std::mem::take(map) {
|
||||
new_map.insert(k, v);
|
||||
}
|
||||
let rebuilt = Value::Object(new_map);
|
||||
let bytes = serde_json::to_vec(&rebuilt).map_err(EnvelopeError::NotJson)?;
|
||||
Ok(Bytes::from(bytes))
|
||||
}
|
||||
}
|
||||
|
||||
/// Helper for reading just the model id from a Bedrock URL path.
|
||||
///
|
||||
/// Bedrock paths look like `/model/{model_id}/invoke`. The `{model_id}`
|
||||
/// segment can contain dots (`anthropic.claude-3-haiku-20240307-v1:0`),
|
||||
/// hyphens, colons, and digits — all of which are allowed in URL path
|
||||
/// segments. We use Axum's `:model_id` capture; this helper exists for
|
||||
/// callers who only have the raw path.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ModelPath {
|
||||
pub model_id: String,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn parses_minimal_envelope() {
|
||||
let body = json!({
|
||||
"anthropic_version": "bedrock-2023-05-31",
|
||||
"max_tokens": 16,
|
||||
"messages": [{"role": "user", "content": "hi"}]
|
||||
});
|
||||
let bytes = serde_json::to_vec(&body).unwrap();
|
||||
let env = BedrockEnvelope::parse(&bytes).expect("parses");
|
||||
assert_eq!(env.anthropic_version, "bedrock-2023-05-31");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_anthropic_version_errors() {
|
||||
let body = json!({"max_tokens": 16, "messages": []});
|
||||
let bytes = serde_json::to_vec(&body).unwrap();
|
||||
let err = BedrockEnvelope::parse(&bytes).expect_err("must error");
|
||||
assert!(matches!(err, EnvelopeError::MissingAnthropicVersion));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn anthropic_version_not_string_errors() {
|
||||
let body = json!({"anthropic_version": 123, "max_tokens": 16});
|
||||
let bytes = serde_json::to_vec(&body).unwrap();
|
||||
let err = BedrockEnvelope::parse(&bytes).expect_err("must error");
|
||||
assert!(matches!(err, EnvelopeError::AnthropicVersionNotString));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn not_an_object_errors() {
|
||||
let bytes = b"[1,2,3]";
|
||||
let err = BedrockEnvelope::parse(bytes).expect_err("must error");
|
||||
assert!(matches!(err, EnvelopeError::NotObject));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_json_errors() {
|
||||
let bytes = b"not json";
|
||||
let err = BedrockEnvelope::parse(bytes).expect_err("must error");
|
||||
assert!(matches!(err, EnvelopeError::NotJson(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ensure_first_no_op_when_already_first() {
|
||||
let body = br#"{"anthropic_version":"bedrock-2023-05-31","max_tokens":16}"#;
|
||||
let out = BedrockEnvelope::ensure_anthropic_version_first(body).unwrap();
|
||||
assert_eq!(&out[..], &body[..]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ensure_first_reorders_when_not_first() {
|
||||
// anthropic_version comes second.
|
||||
let body = br#"{"max_tokens":16,"anthropic_version":"bedrock-2023-05-31"}"#;
|
||||
let out = BedrockEnvelope::ensure_anthropic_version_first(body).unwrap();
|
||||
let out_str = std::str::from_utf8(&out).unwrap();
|
||||
// First key after `{"` must be `anthropic_version`.
|
||||
assert!(
|
||||
out_str.starts_with(r#"{"anthropic_version":"bedrock-2023-05-31""#),
|
||||
"got {out_str}"
|
||||
);
|
||||
}
|
||||
}
|
||||
555
crates/headroom-proxy/src/bedrock/invoke.rs
Normal file
555
crates/headroom-proxy/src/bedrock/invoke.rs
Normal file
|
|
@ -0,0 +1,555 @@
|
|||
//! POST `/model/{model_id}/invoke` handler — Phase D PR-D1.
|
||||
//!
|
||||
//! # Pipeline
|
||||
//!
|
||||
//! 1. Extract `{model_id}` from the path. The Bedrock convention is
|
||||
//! `anthropic.claude-3-haiku-20240307-v1:0` — dot-separated
|
||||
//! `<vendor>.<model>-<date>-<rev>`.
|
||||
//! 2. If the vendor is `anthropic`, parse the body as a Bedrock
|
||||
//! envelope (`{"anthropic_version": "...", ...rest}`) and run the
|
||||
//! live-zone Anthropic compression dispatcher over the body bytes.
|
||||
//! The dispatcher is the SAME one `/v1/messages` uses; Bedrock's
|
||||
//! body shape is just Anthropic-without-the-`model`-field.
|
||||
//! 3. Re-emit the (possibly compressed) body with `anthropic_version`
|
||||
//! preserved as the first key.
|
||||
//! 4. Build the upstream URL (`https://bedrock-runtime.{region}.amazonaws.com/model/{model}/invoke`
|
||||
//! or operator override).
|
||||
//! 5. Sign the (post-compression) body bytes with AWS SigV4. Sign
|
||||
//! over `host`, `x-amz-date`, `x-amz-content-sha256` plus any
|
||||
//! extra headers (`content-type`, `accept`).
|
||||
//! 6. Forward to Bedrock; stream the response back to the client.
|
||||
//!
|
||||
//! # Failure modes
|
||||
//!
|
||||
//! - **Missing credentials**: log
|
||||
//! `event=bedrock_credentials_missing` at WARN, return `500` with
|
||||
//! a JSON error body. NEVER forwards an unsigned request.
|
||||
//! - **Envelope parse failure**: log `event=bedrock_envelope_parse_error`,
|
||||
//! pass the bytes through unchanged. Bedrock will reject anyway,
|
||||
//! but the failure is the customer's, not ours — we just route.
|
||||
//! This matches the Anthropic compression path's
|
||||
//! `Outcome::Passthrough` behaviour.
|
||||
//! - **Non-anthropic model**: skip compression, but still sign +
|
||||
//! forward. Other vendors (Amazon Titan, Cohere, AI21, Meta) have
|
||||
//! different body shapes that the proxy doesn't yet understand;
|
||||
//! we pass them through opaquely.
|
||||
//! - **SigV4 signing failure**: log
|
||||
//! `event=bedrock_sigv4_failed`, return `500`. NEVER forwards
|
||||
//! unsigned.
|
||||
|
||||
use std::net::SocketAddr;
|
||||
use std::time::SystemTime;
|
||||
|
||||
use axum::body::Body;
|
||||
use axum::extract::{ConnectInfo, Path, State};
|
||||
use axum::http::{HeaderMap, Method, StatusCode, Uri};
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use bytes::Bytes;
|
||||
use futures_util::StreamExt as _;
|
||||
use http::HeaderName;
|
||||
use url::Url;
|
||||
|
||||
use crate::bedrock::envelope::BedrockEnvelope;
|
||||
use crate::bedrock::sigv4::{sign_request, SigningInputs};
|
||||
use crate::compression::{
|
||||
compress_anthropic_request, Outcome as AnthropicOutcome, PassthroughReason,
|
||||
};
|
||||
use crate::headers::filter_response_headers;
|
||||
use crate::proxy::AppState;
|
||||
|
||||
/// 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.";
|
||||
|
||||
/// AWS Bedrock Runtime DNS template. The `{}` placeholder is the
|
||||
/// region. Only used when `Config::bedrock_endpoint` is `None`.
|
||||
const BEDROCK_RUNTIME_HOST_TEMPLATE: &str = "bedrock-runtime.{region}.amazonaws.com";
|
||||
|
||||
/// Axum POST handler for `/model/{model_id}/invoke`.
|
||||
///
|
||||
/// Buffers the body so the live-zone compressor + SigV4 signer can
|
||||
/// inspect it. Both are required to be applied to the SAME byte slice
|
||||
/// — the signer hashes whatever the forwarder will actually send.
|
||||
pub async fn handle_invoke(
|
||||
State(state): State<AppState>,
|
||||
ConnectInfo(client_addr): ConnectInfo<SocketAddr>,
|
||||
Path(model_id): Path<String>,
|
||||
method: Method,
|
||||
uri: Uri,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
) -> Response {
|
||||
let _ = client_addr; // accepted for ConnectInfo extractor; not used directly today
|
||||
let request_id = headers
|
||||
.get("x-request-id")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(|s| s.to_string())
|
||||
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
|
||||
|
||||
tracing::info!(
|
||||
event = "bedrock_invoke_received",
|
||||
request_id = %request_id,
|
||||
method = %method,
|
||||
model_id = %model_id,
|
||||
body_bytes = body.len(),
|
||||
"bedrock invoke route received request"
|
||||
);
|
||||
|
||||
let is_anthropic = model_id.starts_with(ANTHROPIC_VENDOR_PREFIX);
|
||||
let outbound_body: Bytes = if is_anthropic {
|
||||
run_anthropic_compression(&body, &state, &request_id)
|
||||
} else {
|
||||
tracing::info!(
|
||||
event = "bedrock_compression_skipped",
|
||||
request_id = %request_id,
|
||||
model_id = %model_id,
|
||||
reason = "non_anthropic_vendor",
|
||||
"bedrock invoke: skipping live-zone compression for non-anthropic vendor"
|
||||
);
|
||||
body.clone()
|
||||
};
|
||||
|
||||
// Build the upstream URL based on configured endpoint or
|
||||
// region-derived default.
|
||||
let upstream_url = match build_bedrock_upstream(&state, &model_id, &uri, "invoke") {
|
||||
Ok(u) => u,
|
||||
Err(msg) => {
|
||||
tracing::error!(
|
||||
event = "bedrock_endpoint_invalid",
|
||||
request_id = %request_id,
|
||||
error = %msg,
|
||||
"bedrock invoke: failed to construct upstream URL"
|
||||
);
|
||||
return error_response(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"bedrock_endpoint_invalid",
|
||||
&msg,
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
// Resolve credentials. No silent fallback: missing creds → 5xx.
|
||||
let creds = match state.bedrock_credentials.as_ref() {
|
||||
Some(c) => c.clone(),
|
||||
None => {
|
||||
tracing::warn!(
|
||||
event = "bedrock_credentials_missing",
|
||||
request_id = %request_id,
|
||||
model_id = %model_id,
|
||||
"bedrock invoke: refusing to forward without AWS credentials"
|
||||
);
|
||||
return error_response(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"bedrock_credentials_missing",
|
||||
"AWS credentials not configured; refusing to forward unsigned",
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
// Build the headers we sign + forward. Start from the inbound
|
||||
// headers, drop the ones the upstream client manages, then sign.
|
||||
let extra_signed: Vec<(String, String)> = collect_signed_headers(&headers, &upstream_url);
|
||||
let extra_signed_refs: Vec<(&str, &str)> = extra_signed
|
||||
.iter()
|
||||
.map(|(k, v)| (k.as_str(), v.as_str()))
|
||||
.collect();
|
||||
|
||||
let sign_inputs = SigningInputs {
|
||||
method: method.as_str(),
|
||||
url: &upstream_url,
|
||||
region: &state.config.bedrock_region,
|
||||
credentials: creds.as_ref(),
|
||||
body: &outbound_body,
|
||||
extra_signed_headers: &extra_signed_refs,
|
||||
time: SystemTime::now(),
|
||||
};
|
||||
let signed = match sign_request(&sign_inputs) {
|
||||
Ok(s) => s,
|
||||
Err(e) => {
|
||||
tracing::error!(
|
||||
event = "bedrock_sigv4_failed",
|
||||
request_id = %request_id,
|
||||
model_id = %model_id,
|
||||
error = %e,
|
||||
"bedrock invoke: SigV4 signing failed; refusing to forward"
|
||||
);
|
||||
return error_response(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"bedrock_sigv4_failed",
|
||||
&e.to_string(),
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
// Compose the outgoing header map. Start with the headers we'll
|
||||
// forward (filter out hop-by-hop / Host / Content-Length;
|
||||
// reqwest sets those itself), then layer the SigV4 outputs on
|
||||
// top — they replace any pre-existing copies of the same name.
|
||||
let mut outbound_headers = HeaderMap::new();
|
||||
for (name, value) in extra_signed.iter() {
|
||||
if let (Ok(n), Ok(v)) = (
|
||||
HeaderName::from_bytes(name.as_bytes()),
|
||||
http::HeaderValue::from_str(value),
|
||||
) {
|
||||
outbound_headers.insert(n, v);
|
||||
}
|
||||
}
|
||||
for (name, value) in signed.entries.iter() {
|
||||
if let (Ok(n), Ok(v)) = (
|
||||
HeaderName::from_bytes(name.as_bytes()),
|
||||
http::HeaderValue::from_str(value),
|
||||
) {
|
||||
outbound_headers.insert(n, v);
|
||||
}
|
||||
}
|
||||
|
||||
// Forward. We surface upstream errors as 502; the byte path
|
||||
// streams the response back to the client.
|
||||
let reqwest_method = match reqwest::Method::from_bytes(method.as_str().as_bytes()) {
|
||||
Ok(m) => m,
|
||||
Err(e) => {
|
||||
tracing::error!(
|
||||
event = "bedrock_invalid_method",
|
||||
request_id = %request_id,
|
||||
error = %e,
|
||||
"bedrock invoke: invalid HTTP method"
|
||||
);
|
||||
return error_response(
|
||||
StatusCode::BAD_REQUEST,
|
||||
"bedrock_invalid_method",
|
||||
&e.to_string(),
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
let upstream_resp = state
|
||||
.client
|
||||
.request(reqwest_method, upstream_url.clone())
|
||||
.headers(outbound_headers)
|
||||
.body(outbound_body.clone())
|
||||
.send()
|
||||
.await;
|
||||
|
||||
let upstream_resp = match upstream_resp {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
event = "bedrock_upstream_error",
|
||||
request_id = %request_id,
|
||||
error = %e,
|
||||
"bedrock invoke: upstream request failed"
|
||||
);
|
||||
let status = if e.is_timeout() {
|
||||
StatusCode::GATEWAY_TIMEOUT
|
||||
} else {
|
||||
StatusCode::BAD_GATEWAY
|
||||
};
|
||||
return error_response(status, "bedrock_upstream_error", &e.to_string());
|
||||
}
|
||||
};
|
||||
|
||||
let status =
|
||||
StatusCode::from_u16(upstream_resp.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY);
|
||||
let resp_headers = filter_response_headers(upstream_resp.headers());
|
||||
|
||||
tracing::info!(
|
||||
event = "bedrock_invoke_forwarded",
|
||||
request_id = %request_id,
|
||||
model_id = %model_id,
|
||||
upstream_status = status.as_u16(),
|
||||
upstream_url = %upstream_url,
|
||||
"bedrock invoke: response forwarded"
|
||||
);
|
||||
|
||||
// Stream the response body back without buffering.
|
||||
let stream = upstream_resp
|
||||
.bytes_stream()
|
||||
.map(|r| r.map_err(std::io::Error::other));
|
||||
let body_out = Body::from_stream(stream);
|
||||
|
||||
let mut builder = Response::builder().status(status);
|
||||
if let Some(h) = builder.headers_mut() {
|
||||
h.extend(resp_headers);
|
||||
if let Ok(v) = http::HeaderValue::from_str(&request_id) {
|
||||
h.insert(HeaderName::from_static("x-request-id"), v);
|
||||
}
|
||||
}
|
||||
builder.body(body_out).unwrap_or_else(|e| {
|
||||
tracing::error!(
|
||||
event = "bedrock_response_build_failed",
|
||||
request_id = %request_id,
|
||||
error = %e,
|
||||
"bedrock invoke: failed to build response"
|
||||
);
|
||||
Response::builder()
|
||||
.status(StatusCode::INTERNAL_SERVER_ERROR)
|
||||
.body(Body::from("internal handler error"))
|
||||
.expect("static response")
|
||||
})
|
||||
}
|
||||
|
||||
/// Run the live-zone Anthropic compressor over a Bedrock-shape body.
|
||||
///
|
||||
/// The compressor only inspects `messages` — it doesn't care that the
|
||||
/// Bedrock body has `anthropic_version` instead of `model`. The
|
||||
/// `Outcome::Compressed` body bytes still preserve key order via
|
||||
/// `serde_json`'s `preserve_order` feature, so the caller's
|
||||
/// re-emission step (`ensure_anthropic_version_first`) almost always
|
||||
/// no-ops. We still call it as a defence-in-depth assertion that
|
||||
/// the byte order is correct before signing.
|
||||
fn run_anthropic_compression(body: &Bytes, state: &AppState, 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",
|
||||
request_id = %request_id,
|
||||
error = %e,
|
||||
"bedrock invoke: envelope parse failed; passing body through unchanged"
|
||||
);
|
||||
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"
|
||||
);
|
||||
|
||||
let outcome = compress_anthropic_request(
|
||||
body,
|
||||
state.config.compression_mode,
|
||||
state.config.cache_control_auto_frozen,
|
||||
request_id,
|
||||
);
|
||||
match outcome {
|
||||
AnthropicOutcome::NoCompression => body.clone(),
|
||||
AnthropicOutcome::Passthrough { reason } => {
|
||||
tracing::info!(
|
||||
event = "bedrock_compression_passthrough",
|
||||
request_id = %request_id,
|
||||
reason = ?reason,
|
||||
"bedrock invoke: live-zone dispatcher fell through to passthrough"
|
||||
);
|
||||
// The compressor's passthrough variants all leave bytes
|
||||
// unchanged. Forward the original.
|
||||
let _ = (PassthroughReason::ModeOff, PassthroughReason::NoMessages); // pin types
|
||||
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()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the upstream URL for the Bedrock route. Honours the
|
||||
/// operator-supplied `bedrock_endpoint` first, falling back to the
|
||||
/// region-derived default. The path/query portion is taken from the
|
||||
/// original URI verbatim — Bedrock's path schema (`/model/{id}/{action}`)
|
||||
/// is identical to the proxy's external path.
|
||||
fn build_bedrock_upstream(
|
||||
state: &AppState,
|
||||
model_id: &str,
|
||||
uri: &Uri,
|
||||
action: &str,
|
||||
) -> Result<Url, String> {
|
||||
let base = match state.config.bedrock_endpoint.as_ref() {
|
||||
Some(u) => u.clone(),
|
||||
None => {
|
||||
let host =
|
||||
BEDROCK_RUNTIME_HOST_TEMPLATE.replace("{region}", &state.config.bedrock_region);
|
||||
Url::parse(&format!("https://{host}/"))
|
||||
.map_err(|e| format!("bedrock derived base URL parse error: {e}"))?
|
||||
}
|
||||
};
|
||||
// Compose the path. We trust the captured `model_id` (Axum
|
||||
// already URL-decoded it) and append `/{action}`.
|
||||
let path = format!(
|
||||
"/model/{model_id}/{action}",
|
||||
model_id = model_id,
|
||||
action = action,
|
||||
);
|
||||
let mut joined = base;
|
||||
joined.set_path(&path);
|
||||
if let Some(q) = uri.query() {
|
||||
joined.set_query(Some(q));
|
||||
}
|
||||
Ok(joined)
|
||||
}
|
||||
|
||||
/// Build the list of headers to sign + forward. Drops hop-by-hop,
|
||||
/// `host`, `content-length` (reqwest manages those), `authorization`
|
||||
/// (we replace it with the SigV4 output). Lower-cases names for
|
||||
/// canonical-request consistency.
|
||||
fn collect_signed_headers(headers: &HeaderMap, upstream_url: &Url) -> Vec<(String, String)> {
|
||||
let mut out: Vec<(String, String)> = Vec::with_capacity(headers.len() + 1);
|
||||
for (name, value) in headers.iter() {
|
||||
let n = name.as_str().to_ascii_lowercase();
|
||||
if matches!(
|
||||
n.as_str(),
|
||||
"host"
|
||||
| "content-length"
|
||||
| "connection"
|
||||
| "keep-alive"
|
||||
| "proxy-authenticate"
|
||||
| "proxy-authorization"
|
||||
| "te"
|
||||
| "trailers"
|
||||
| "transfer-encoding"
|
||||
| "upgrade"
|
||||
| "authorization"
|
||||
| "x-amz-date"
|
||||
| "x-amz-content-sha256"
|
||||
) {
|
||||
// Drop client-managed + signer-managed headers.
|
||||
continue;
|
||||
}
|
||||
if n.starts_with("x-headroom-") {
|
||||
// Internal headers are stripped from upstream traffic
|
||||
// (PR-A5). The Bedrock route inherits the same default.
|
||||
continue;
|
||||
}
|
||||
if let Ok(v) = value.to_str() {
|
||||
out.push((n, v.to_string()));
|
||||
}
|
||||
}
|
||||
// Signer requires `host` in the canonical request. Add it
|
||||
// explicitly from the upstream URL — the inbound `host` header
|
||||
// (the proxy's listening hostname) is wrong for the canonical
|
||||
// request.
|
||||
if let Some(host) = upstream_url.host_str() {
|
||||
let host_value = match upstream_url.port() {
|
||||
Some(p) => format!("{host}:{p}"),
|
||||
None => host.to_string(),
|
||||
};
|
||||
out.push(("host".to_string(), host_value));
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
fn error_response(status: StatusCode, event: &str, msg: &str) -> Response {
|
||||
let body = serde_json::json!({
|
||||
"error": {
|
||||
"type": event,
|
||||
"message": msg,
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
let _ = event;
|
||||
let mut resp = Response::builder()
|
||||
.status(status)
|
||||
.body(Body::from(body))
|
||||
.expect("static error response");
|
||||
resp.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
http::HeaderValue::from_static("application/json"),
|
||||
);
|
||||
resp.into_response()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[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));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_upstream_uses_region_default() {
|
||||
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,
|
||||
};
|
||||
let uri: Uri = "/model/anthropic.claude-3-haiku-20240307-v1:0/invoke"
|
||||
.parse()
|
||||
.unwrap();
|
||||
let url = build_bedrock_upstream(
|
||||
&state,
|
||||
"anthropic.claude-3-haiku-20240307-v1:0",
|
||||
&uri,
|
||||
"invoke",
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
url.as_str(),
|
||||
"https://bedrock-runtime.us-west-2.amazonaws.com/model/anthropic.claude-3-haiku-20240307-v1:0/invoke"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_upstream_honors_explicit_endpoint() {
|
||||
use crate::config::Config;
|
||||
let mut config = Config::for_test(Url::parse("http://up:8080").unwrap());
|
||||
config.bedrock_endpoint = Some(Url::parse("http://127.0.0.1:9999").unwrap());
|
||||
let state = AppState {
|
||||
config: std::sync::Arc::new(config),
|
||||
client: reqwest::Client::new(),
|
||||
bedrock_credentials: None,
|
||||
};
|
||||
let uri: Uri = "/model/anthropic.claude-3-haiku-20240307-v1:0/invoke"
|
||||
.parse()
|
||||
.unwrap();
|
||||
let url = build_bedrock_upstream(
|
||||
&state,
|
||||
"anthropic.claude-3-haiku-20240307-v1:0",
|
||||
&uri,
|
||||
"invoke",
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
url.as_str(),
|
||||
"http://127.0.0.1:9999/model/anthropic.claude-3-haiku-20240307-v1:0/invoke"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn collect_signed_headers_strips_client_managed() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert("content-type", "application/json".parse().unwrap());
|
||||
headers.insert("host", "proxy.example".parse().unwrap());
|
||||
headers.insert("authorization", "Bearer x".parse().unwrap());
|
||||
headers.insert("x-headroom-mode", "live".parse().unwrap());
|
||||
headers.insert("accept", "application/json".parse().unwrap());
|
||||
let upstream =
|
||||
Url::parse("https://bedrock-runtime.us-east-1.amazonaws.com/model/x/invoke").unwrap();
|
||||
let out = collect_signed_headers(&headers, &upstream);
|
||||
let names: Vec<&str> = out.iter().map(|(k, _)| k.as_str()).collect();
|
||||
assert!(names.contains(&"content-type"));
|
||||
assert!(names.contains(&"accept"));
|
||||
assert!(names.contains(&"host"));
|
||||
assert!(!names.contains(&"authorization"));
|
||||
assert!(!names.contains(&"x-headroom-mode"));
|
||||
// host must be the upstream host, not the proxy host.
|
||||
let host = out
|
||||
.iter()
|
||||
.find(|(k, _)| k == "host")
|
||||
.map(|(_, v)| v.as_str())
|
||||
.unwrap();
|
||||
assert_eq!(host, "bedrock-runtime.us-east-1.amazonaws.com");
|
||||
}
|
||||
}
|
||||
52
crates/headroom-proxy/src/bedrock/mod.rs
Normal file
52
crates/headroom-proxy/src/bedrock/mod.rs
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
//! Native AWS Bedrock InvokeModel route — Phase D PR-D1.
|
||||
//!
|
||||
//! # Why a separate module?
|
||||
//!
|
||||
//! The Python proxy currently routes Anthropic-on-Bedrock through the
|
||||
//! `litellm` shim (`headroom/backends/litellm.py`). That shim
|
||||
//! lossy-converts every request and response between Anthropic and
|
||||
//! OpenAI shapes, dropping `thinking`, `redacted_thinking`,
|
||||
//! `document`, `search_result`, `image`, `server_tool_use`, and
|
||||
//! `mcp_tool_use` blocks (P4-37). It also hardcodes
|
||||
//! `stop_sequence: null` (§11.1 violation) and re-wraps
|
||||
//! `function_call.arguments` as a parsed JSON object (§4.4 — P4-43).
|
||||
//!
|
||||
//! Phase D rebuilds the Bedrock surface natively in Rust. PR-D1
|
||||
//! handles the **non-streaming** `POST /model/{model}/invoke` route:
|
||||
//!
|
||||
//! 1. Parse the Bedrock envelope (`{"anthropic_version": "...",
|
||||
//! ...rest_of_anthropic_body}`).
|
||||
//! 2. Route Anthropic-shape bodies through the live-zone compression
|
||||
//! path (the same one `/v1/messages` uses).
|
||||
//! 3. Re-emit the envelope with `anthropic_version` preserved as the
|
||||
//! first key — Bedrock is strict about schema validation.
|
||||
//! 4. Sign the **outgoing** body bytes with AWS SigV4 (after
|
||||
//! compression) and forward to the configured Bedrock endpoint.
|
||||
//!
|
||||
//! # Cache safety
|
||||
//!
|
||||
//! The signed bytes are exactly the bytes Bedrock receives. If the
|
||||
//! compressor mutated the body, the SigV4 signature is computed
|
||||
//! against the post-compression bytes; the upstream verifier will
|
||||
//! accept them. There is no "sign before compress" path — that would
|
||||
//! produce a signature that doesn't match the wire payload.
|
||||
//!
|
||||
//! # Module layout
|
||||
//!
|
||||
//! - [`envelope`] — `BedrockEnvelope` parse + emit (preserves
|
||||
//! `anthropic_version` ordering byte-equal).
|
||||
//! - [`sigv4`] — AWS SigV4 signing helper. Wraps the `aws-sigv4`
|
||||
//! crate with the project's no-fallback / structured-logging
|
||||
//! policy.
|
||||
//! - [`invoke`] — POST handler for `/model/{model}/invoke`.
|
||||
//!
|
||||
//! Streaming (`/model/{model}/invoke-with-response-stream`) is
|
||||
//! Phase D PR-D2.
|
||||
|
||||
pub mod envelope;
|
||||
pub mod invoke;
|
||||
pub mod sigv4;
|
||||
|
||||
pub use envelope::{BedrockEnvelope, EnvelopeError};
|
||||
pub use invoke::handle_invoke;
|
||||
pub use sigv4::{sign_request, SigV4Error, SigningInputs};
|
||||
324
crates/headroom-proxy/src/bedrock/sigv4.rs
Normal file
324
crates/headroom-proxy/src/bedrock/sigv4.rs
Normal file
|
|
@ -0,0 +1,324 @@
|
|||
//! AWS SigV4 request signing for the Bedrock InvokeModel route.
|
||||
//!
|
||||
//! # What this module does
|
||||
//!
|
||||
//! Wraps the `aws-sigv4` crate with the project's policies:
|
||||
//!
|
||||
//! - **No silent fallbacks.** A failed sign-attempt returns a
|
||||
//! structured error; the handler surfaces 5xx and logs an event.
|
||||
//! We NEVER forward an unsigned request to Bedrock, because the
|
||||
//! AWS endpoint will reject anyway and the user would see an
|
||||
//! opaque 403.
|
||||
//! - **Body bytes are hashed AFTER compression.** The signing
|
||||
//! inputs include a `&[u8]` that is the EXACT byte slice the
|
||||
//! forwarder will send upstream. There is no separate "hash the
|
||||
//! pre-compression body" code path.
|
||||
//! - **Structured logs at every decision point.** Every code path
|
||||
//! emits `tracing::info!`/`warn!` with an `event = ...` field so
|
||||
//! operators can confirm signing happened.
|
||||
//!
|
||||
//! # What this module deliberately does NOT do
|
||||
//!
|
||||
//! - It does not resolve credentials — that's `aws-config`'s job and
|
||||
//! happens at app-startup time so per-request signing is cheap.
|
||||
//! The handler holds the resolved [`aws_credential_types::Credentials`]
|
||||
//! in `AppState` and passes them in.
|
||||
//! - It does not buffer the body — callers buffer the body in the
|
||||
//! compression gate (it's the same byte slice).
|
||||
//! - It does not handle SigV4a (the cross-region variant). Bedrock
|
||||
//! uses standard SigV4 per region.
|
||||
|
||||
use std::time::SystemTime;
|
||||
|
||||
use aws_credential_types::Credentials;
|
||||
use aws_sigv4::http_request::{
|
||||
sign, PayloadChecksumKind, SignableBody, SignableRequest, SigningSettings,
|
||||
};
|
||||
use aws_sigv4::sign::v4;
|
||||
use aws_smithy_runtime_api::client::identity::Identity;
|
||||
use thiserror::Error;
|
||||
use url::Url;
|
||||
|
||||
/// AWS service name used in the SigV4 string-to-sign for Bedrock.
|
||||
/// Documented at
|
||||
/// <https://docs.aws.amazon.com/bedrock/latest/userguide/security-iam.html>
|
||||
pub const BEDROCK_SERVICE_NAME: &str = "bedrock";
|
||||
|
||||
/// Inputs needed to sign a Bedrock request. Borrows the body so we
|
||||
/// avoid a copy on the hot path.
|
||||
#[derive(Debug)]
|
||||
pub struct SigningInputs<'a> {
|
||||
/// HTTP method (always `"POST"` for InvokeModel today, but we
|
||||
/// keep this explicit so a future GET-shaped surface doesn't
|
||||
/// require redoing the call site).
|
||||
pub method: &'a str,
|
||||
/// Fully-qualified upstream URL (scheme + host + path + query).
|
||||
/// Used for canonical-request URI normalization.
|
||||
pub url: &'a Url,
|
||||
/// AWS region the upstream endpoint lives in.
|
||||
pub region: &'a str,
|
||||
/// AWS credentials (resolved from the `aws-config` default chain
|
||||
/// at app-startup time).
|
||||
pub credentials: &'a Credentials,
|
||||
/// Body bytes to sign. MUST be the exact bytes the forwarder
|
||||
/// will send to Bedrock (post-compression).
|
||||
pub body: &'a [u8],
|
||||
/// Extra headers the canonical request must include in the
|
||||
/// signed-headers list. The signer always includes `host`,
|
||||
/// `x-amz-date`, and `x-amz-content-sha256`; callers add
|
||||
/// anything else they want covered (e.g. `accept-encoding`,
|
||||
/// `content-type`).
|
||||
pub extra_signed_headers: &'a [(&'a str, &'a str)],
|
||||
/// Time to use in the signature. Production uses
|
||||
/// `SystemTime::now()`; tests pin a known time to make the
|
||||
/// canonical request deterministic.
|
||||
pub time: SystemTime,
|
||||
}
|
||||
|
||||
/// Headers that the signer will write into the outbound request.
|
||||
/// The handler must add every entry to the upstream-bound HeaderMap
|
||||
/// before sending — Bedrock validates each header against the
|
||||
/// canonical request.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SignedHeaders {
|
||||
pub entries: Vec<(String, String)>,
|
||||
/// Lowercase hex SHA-256 of the body. Surfaced for tests + logs.
|
||||
pub signature: String,
|
||||
}
|
||||
|
||||
/// Errors surfaced by the signing path.
|
||||
#[derive(Debug, Error)]
|
||||
pub enum SigV4Error {
|
||||
/// `aws-sigv4` rejected the request (URL parse, malformed header,
|
||||
/// etc).
|
||||
#[error("sigv4 signing failed: {0}")]
|
||||
Sign(String),
|
||||
/// The signing-params builder rejected the inputs (e.g. missing
|
||||
/// region — should never happen because we validate at startup).
|
||||
#[error("sigv4 builder error: {0}")]
|
||||
Builder(String),
|
||||
}
|
||||
|
||||
/// Sign a Bedrock request and return the headers the handler must add
|
||||
/// to the outbound request.
|
||||
///
|
||||
/// # Cache safety
|
||||
///
|
||||
/// The body bytes passed in MUST be the bytes the proxy is about to
|
||||
/// send upstream. If the compressor mutated the body, those mutated
|
||||
/// bytes are what get signed — Bedrock will accept because the
|
||||
/// signature covers the wire payload, not the original.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns [`SigV4Error::Sign`] when `aws-sigv4` rejects the request
|
||||
/// (malformed URL, etc). Returns [`SigV4Error::Builder`] when the
|
||||
/// signing-params builder rejects the inputs.
|
||||
pub fn sign_request(inputs: &SigningInputs<'_>) -> Result<SignedHeaders, SigV4Error> {
|
||||
// Build the identity wrapper around the credentials. Identity is
|
||||
// the type the SigV4 signer accepts; it can in principle hold
|
||||
// alternative auth schemes (Bearer, etc) but we only ever use it
|
||||
// for AWS creds.
|
||||
let identity: Identity = Identity::new(inputs.credentials.clone(), None);
|
||||
|
||||
// Default settings + force `x-amz-content-sha256` into the
|
||||
// canonical request. Bedrock validates the content hash to
|
||||
// catch any in-flight body mutation; with `NoHeader` (the
|
||||
// crate-level default) the signer would skip the header,
|
||||
// which means a downstream gateway that DOES check it would
|
||||
// 403.
|
||||
let mut settings = SigningSettings::default();
|
||||
settings.payload_checksum_kind = PayloadChecksumKind::XAmzSha256;
|
||||
|
||||
let signing_params = v4::SigningParams::builder()
|
||||
.identity(&identity)
|
||||
.region(inputs.region)
|
||||
.name(BEDROCK_SERVICE_NAME)
|
||||
.time(inputs.time)
|
||||
.settings(settings)
|
||||
.build()
|
||||
.map_err(|e| SigV4Error::Builder(e.to_string()))?
|
||||
.into();
|
||||
|
||||
// The signer needs the URL as a string; Url's Display impl is
|
||||
// canonical RFC 3986 form which is what aws-sigv4 expects.
|
||||
let url_string = inputs.url.to_string();
|
||||
|
||||
let signable = SignableRequest::new(
|
||||
inputs.method,
|
||||
url_string,
|
||||
inputs.extra_signed_headers.iter().copied(),
|
||||
SignableBody::Bytes(inputs.body),
|
||||
)
|
||||
.map_err(|e| SigV4Error::Sign(e.to_string()))?;
|
||||
|
||||
let signing_output =
|
||||
sign(signable, &signing_params).map_err(|e| SigV4Error::Sign(e.to_string()))?;
|
||||
|
||||
let signature = signing_output.signature().to_string();
|
||||
let (instructions, _signature) = signing_output.into_parts();
|
||||
let (header_entries, _query_params) = instructions.into_parts();
|
||||
let entries = header_entries
|
||||
.into_iter()
|
||||
.map(|h| (h.name().to_string(), h.value().to_string()))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
tracing::info!(
|
||||
event = "sigv4_signed",
|
||||
forwarder = "rust_proxy",
|
||||
region = inputs.region,
|
||||
service = BEDROCK_SERVICE_NAME,
|
||||
method = inputs.method,
|
||||
host = inputs.url.host_str().unwrap_or(""),
|
||||
body_bytes = inputs.body.len(),
|
||||
headers_added = entries.len(),
|
||||
"bedrock request signed with sigv4"
|
||||
);
|
||||
|
||||
Ok(SignedHeaders { entries, signature })
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::time::Duration;
|
||||
|
||||
fn fixed_time() -> SystemTime {
|
||||
// 2026-05-03T12:00:00Z — pinned so canonical request is
|
||||
// deterministic across runs.
|
||||
SystemTime::UNIX_EPOCH + Duration::from_secs(1_777_910_400)
|
||||
}
|
||||
|
||||
fn fixture_credentials() -> Credentials {
|
||||
Credentials::new("AKIA_TEST", "secret_test", None, None, "test")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn signs_minimal_request_produces_expected_headers() {
|
||||
let url = Url::parse(
|
||||
"https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-3-haiku-20240307-v1:0/invoke",
|
||||
)
|
||||
.unwrap();
|
||||
let creds = fixture_credentials();
|
||||
let body = br#"{"anthropic_version":"bedrock-2023-05-31","max_tokens":16,"messages":[]}"#;
|
||||
let inputs = SigningInputs {
|
||||
method: "POST",
|
||||
url: &url,
|
||||
region: "us-east-1",
|
||||
credentials: &creds,
|
||||
body,
|
||||
extra_signed_headers: &[("content-type", "application/json")],
|
||||
time: fixed_time(),
|
||||
};
|
||||
let signed = sign_request(&inputs).expect("sigv4 sign");
|
||||
// The signer always emits `authorization`, `x-amz-date`, and
|
||||
// `x-amz-content-sha256`. Confirm all three are present.
|
||||
let names: Vec<String> = signed
|
||||
.entries
|
||||
.iter()
|
||||
.map(|(k, _)| k.to_ascii_lowercase())
|
||||
.collect();
|
||||
assert!(
|
||||
names.iter().any(|n| n == "authorization"),
|
||||
"must add authorization; got {names:?}"
|
||||
);
|
||||
assert!(
|
||||
names.iter().any(|n| n == "x-amz-date"),
|
||||
"must add x-amz-date"
|
||||
);
|
||||
assert!(
|
||||
names.iter().any(|n| n == "x-amz-content-sha256"),
|
||||
"must add x-amz-content-sha256"
|
||||
);
|
||||
assert_eq!(signed.signature.len(), 64, "sigv4 signature is 32-byte hex");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn signature_is_deterministic_for_fixed_inputs() {
|
||||
let url = Url::parse(
|
||||
"https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-3-haiku-20240307-v1:0/invoke",
|
||||
)
|
||||
.unwrap();
|
||||
let creds = fixture_credentials();
|
||||
let body = br#"{"anthropic_version":"bedrock-2023-05-31","max_tokens":16}"#;
|
||||
let mk = || SigningInputs {
|
||||
method: "POST",
|
||||
url: &url,
|
||||
region: "us-east-1",
|
||||
credentials: &creds,
|
||||
body,
|
||||
extra_signed_headers: &[("content-type", "application/json")],
|
||||
time: fixed_time(),
|
||||
};
|
||||
let a = sign_request(&mk()).unwrap();
|
||||
let b = sign_request(&mk()).unwrap();
|
||||
assert_eq!(a.signature, b.signature);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn changing_body_changes_signature() {
|
||||
let url = Url::parse(
|
||||
"https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-3-haiku-20240307-v1:0/invoke",
|
||||
)
|
||||
.unwrap();
|
||||
let creds = fixture_credentials();
|
||||
let inputs_a = SigningInputs {
|
||||
method: "POST",
|
||||
url: &url,
|
||||
region: "us-east-1",
|
||||
credentials: &creds,
|
||||
body: br#"{"anthropic_version":"bedrock-2023-05-31","max_tokens":16}"#,
|
||||
extra_signed_headers: &[("content-type", "application/json")],
|
||||
time: fixed_time(),
|
||||
};
|
||||
let inputs_b = SigningInputs {
|
||||
body: br#"{"anthropic_version":"bedrock-2023-05-31","max_tokens":32}"#,
|
||||
..SigningInputs {
|
||||
method: "POST",
|
||||
url: &url,
|
||||
region: "us-east-1",
|
||||
credentials: &creds,
|
||||
body: &[],
|
||||
extra_signed_headers: &[("content-type", "application/json")],
|
||||
time: fixed_time(),
|
||||
}
|
||||
};
|
||||
let a = sign_request(&inputs_a).unwrap();
|
||||
let b = sign_request(&inputs_b).unwrap();
|
||||
assert_ne!(
|
||||
a.signature, b.signature,
|
||||
"different body bytes must yield different signatures"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn changing_region_changes_signature() {
|
||||
let url = Url::parse(
|
||||
"https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-3-haiku-20240307-v1:0/invoke",
|
||||
)
|
||||
.unwrap();
|
||||
let creds = fixture_credentials();
|
||||
let body = br#"{"anthropic_version":"bedrock-2023-05-31"}"#;
|
||||
let east = SigningInputs {
|
||||
method: "POST",
|
||||
url: &url,
|
||||
region: "us-east-1",
|
||||
credentials: &creds,
|
||||
body,
|
||||
extra_signed_headers: &[],
|
||||
time: fixed_time(),
|
||||
};
|
||||
let west = SigningInputs {
|
||||
method: "POST",
|
||||
url: &url,
|
||||
region: "us-west-2",
|
||||
credentials: &creds,
|
||||
body,
|
||||
extra_signed_headers: &[],
|
||||
time: fixed_time(),
|
||||
};
|
||||
let a = sign_request(&east).unwrap();
|
||||
let b = sign_request(&west).unwrap();
|
||||
assert_ne!(a.signature, b.signature);
|
||||
}
|
||||
}
|
||||
|
|
@ -296,6 +296,61 @@ pub struct CliArgs {
|
|||
action = clap::ArgAction::Set,
|
||||
)]
|
||||
pub enable_conversations_passthrough: bool,
|
||||
|
||||
/// Phase D PR-D1: enable the native Bedrock InvokeModel route.
|
||||
/// When `true` (default), `POST /model/{model_id}/invoke` is
|
||||
/// handled by the Rust `bedrock::invoke` handler — Anthropic-shape
|
||||
/// bodies run through the live-zone compression path and the
|
||||
/// proxy re-signs the request with SigV4 before forwarding to
|
||||
/// the configured Bedrock endpoint. When `false`, the routes are
|
||||
/// not mounted and requests fall through to the catch-all
|
||||
/// (which forwards to `--upstream` byte-equal but does NOT
|
||||
/// re-sign — operators MUST run an unsigned upstream that
|
||||
/// happens to know what to do, otherwise this fails closed).
|
||||
///
|
||||
/// Source priority: CLI flag → `HEADROOM_PROXY_ENABLE_BEDROCK_NATIVE`
|
||||
/// env var → default (`true`).
|
||||
#[arg(
|
||||
long = "enable-bedrock-native",
|
||||
env = "HEADROOM_PROXY_ENABLE_BEDROCK_NATIVE",
|
||||
default_value_t = true,
|
||||
action = clap::ArgAction::Set,
|
||||
)]
|
||||
pub enable_bedrock_native: bool,
|
||||
|
||||
/// AWS region to use when signing Bedrock requests. Default
|
||||
/// `us-east-1`. The Bedrock endpoint URL derived from this
|
||||
/// region is `https://bedrock-runtime.{region}.amazonaws.com`
|
||||
/// (override via `--bedrock-endpoint` for FIPS or VPC endpoints).
|
||||
///
|
||||
/// Source priority: CLI flag → `HEADROOM_PROXY_BEDROCK_REGION`
|
||||
/// env var → `AWS_REGION` env var → default (`us-east-1`).
|
||||
#[arg(
|
||||
long = "bedrock-region",
|
||||
env = "HEADROOM_PROXY_BEDROCK_REGION",
|
||||
default_value = "us-east-1"
|
||||
)]
|
||||
pub bedrock_region: String,
|
||||
|
||||
/// Bedrock endpoint base URL. When unset (the common case), the
|
||||
/// proxy derives `https://bedrock-runtime.{bedrock_region}.amazonaws.com`
|
||||
/// from the configured region. Override for FIPS endpoints
|
||||
/// (`bedrock-runtime-fips.{region}.amazonaws.com`), VPC endpoints,
|
||||
/// or local-mock test setups.
|
||||
///
|
||||
/// Source priority: CLI flag → `HEADROOM_PROXY_BEDROCK_ENDPOINT`
|
||||
/// env var → derived-from-region.
|
||||
#[arg(long = "bedrock-endpoint", env = "HEADROOM_PROXY_BEDROCK_ENDPOINT")]
|
||||
pub bedrock_endpoint: Option<Url>,
|
||||
|
||||
/// AWS profile name passed to the `aws-config` default credential
|
||||
/// chain. When unset, the chain uses the default behaviour
|
||||
/// (env vars → `[default]` profile → IMDS / ECS task role).
|
||||
///
|
||||
/// Source priority: CLI flag → `HEADROOM_PROXY_AWS_PROFILE`
|
||||
/// env var → `AWS_PROFILE` env var → default chain.
|
||||
#[arg(long = "aws-profile", env = "HEADROOM_PROXY_AWS_PROFILE")]
|
||||
pub aws_profile: Option<String>,
|
||||
}
|
||||
|
||||
fn parse_duration(s: &str) -> Result<Duration, String> {
|
||||
|
|
@ -350,6 +405,20 @@ pub struct Config {
|
|||
/// NOT gate compression of conversation items (that's
|
||||
/// C5+/B-phase territory).
|
||||
pub enable_conversations_passthrough: bool,
|
||||
/// PR-D1: enable the native Bedrock InvokeModel route. Default
|
||||
/// `true`. When disabled, the explicit Rust handlers are not
|
||||
/// mounted; operators relying on the Python LiteLLM converter
|
||||
/// keep their existing path.
|
||||
pub enable_bedrock_native: bool,
|
||||
/// PR-D1: AWS region used to sign Bedrock requests + (when no
|
||||
/// explicit endpoint is set) derive the Bedrock endpoint URL.
|
||||
pub bedrock_region: String,
|
||||
/// PR-D1: Bedrock endpoint base URL. `None` means
|
||||
/// "derive from region" (`https://bedrock-runtime.{region}.amazonaws.com`).
|
||||
pub bedrock_endpoint: Option<Url>,
|
||||
/// PR-D1: optional AWS profile name. When `None`, the default
|
||||
/// credential chain (env → `[default]` profile → IMDS) is used.
|
||||
pub aws_profile: Option<String>,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
|
|
@ -378,6 +447,10 @@ impl Config {
|
|||
strip_internal_headers: args.strip_internal_headers,
|
||||
enable_responses_streaming: args.enable_responses_streaming,
|
||||
enable_conversations_passthrough: args.enable_conversations_passthrough,
|
||||
enable_bedrock_native: args.enable_bedrock_native,
|
||||
bedrock_region: args.bedrock_region,
|
||||
bedrock_endpoint: args.bedrock_endpoint,
|
||||
aws_profile: args.aws_profile,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -408,6 +481,14 @@ impl Config {
|
|||
// production traffic will hit.
|
||||
enable_responses_streaming: true,
|
||||
enable_conversations_passthrough: true,
|
||||
// PR-D1: bedrock route default-on so tests exercise
|
||||
// it without per-test opt-in. Tests that set
|
||||
// `bedrock_endpoint` to a wiremock URL get the full
|
||||
// sign-and-forward path.
|
||||
enable_bedrock_native: true,
|
||||
bedrock_region: "us-east-1".to_string(),
|
||||
bedrock_endpoint: None,
|
||||
aws_profile: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
//! headroom-proxy library: transparent reverse proxy in front of the Python
|
||||
//! Headroom proxy. Used by both `main.rs` and the integration tests.
|
||||
|
||||
pub mod bedrock;
|
||||
pub mod compression;
|
||||
pub mod config;
|
||||
pub mod error;
|
||||
|
|
|
|||
|
|
@ -31,7 +31,37 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
|||
"headroom-proxy starting"
|
||||
);
|
||||
|
||||
let state = AppState::new(config.clone())?;
|
||||
let mut state = AppState::new(config.clone())?;
|
||||
|
||||
// PR-D1: resolve AWS credentials at startup via the `aws-config`
|
||||
// default chain. Loaded once so per-request signing is cheap.
|
||||
// Failure is NOT fatal — the proxy may run in front of a non-AWS
|
||||
// upstream — but the Bedrock invoke handler refuses to forward
|
||||
// unsigned requests when `bedrock_credentials` is `None`
|
||||
// (see `bedrock::invoke::handle_invoke`).
|
||||
if config.enable_bedrock_native {
|
||||
match load_bedrock_credentials(&config).await {
|
||||
Ok(creds) => {
|
||||
state = state.with_bedrock_credentials(creds);
|
||||
tracing::info!(
|
||||
event = "bedrock_credentials_loaded",
|
||||
region = %config.bedrock_region,
|
||||
profile = ?config.aws_profile,
|
||||
"AWS credentials resolved for Bedrock SigV4 signing"
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
event = "bedrock_credentials_unavailable",
|
||||
region = %config.bedrock_region,
|
||||
profile = ?config.aws_profile,
|
||||
error = %e,
|
||||
"AWS credentials not available at startup; Bedrock invoke will 5xx until creds are configured"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let app = build_app(state).into_make_service_with_connect_info::<SocketAddr>();
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(config.listen).await?;
|
||||
|
|
@ -64,6 +94,31 @@ fn init_tracing(level: &str) {
|
|||
.try_init();
|
||||
}
|
||||
|
||||
/// PR-D1: resolve AWS credentials for Bedrock SigV4 signing.
|
||||
///
|
||||
/// Uses the `aws-config` default chain (env vars → shared profile
|
||||
/// file → IMDS / ECS task role). Honours `Config::aws_profile` when
|
||||
/// set; otherwise the chain picks up `AWS_PROFILE` from the
|
||||
/// environment automatically.
|
||||
async fn load_bedrock_credentials(
|
||||
config: &Config,
|
||||
) -> Result<aws_credential_types::Credentials, Box<dyn std::error::Error + Send + Sync>> {
|
||||
use aws_config::BehaviorVersion;
|
||||
use aws_credential_types::provider::ProvideCredentials;
|
||||
|
||||
let mut loader = aws_config::defaults(BehaviorVersion::latest())
|
||||
.region(aws_config::Region::new(config.bedrock_region.clone()));
|
||||
if let Some(profile) = config.aws_profile.as_deref() {
|
||||
loader = loader.profile_name(profile);
|
||||
}
|
||||
let aws_config = loader.load().await;
|
||||
let creds_provider = aws_config
|
||||
.credentials_provider()
|
||||
.ok_or("no credentials provider configured")?;
|
||||
let creds = creds_provider.provide_credentials().await?;
|
||||
Ok(creds)
|
||||
}
|
||||
|
||||
async fn shutdown_signal() {
|
||||
let ctrl_c = async {
|
||||
let _ = tokio::signal::ctrl_c().await;
|
||||
|
|
|
|||
|
|
@ -34,6 +34,14 @@ use crate::websocket::ws_handler;
|
|||
pub struct AppState {
|
||||
pub config: Arc<Config>,
|
||||
pub client: reqwest::Client,
|
||||
/// PR-D1: AWS credentials resolved at startup via the
|
||||
/// `aws-config` default chain. `None` when the proxy boots
|
||||
/// without AWS creds available (operator running locally
|
||||
/// against a non-Bedrock upstream); the Bedrock invoke handler
|
||||
/// returns 5xx with a structured `event=bedrock_credentials_missing`
|
||||
/// log so failures are LOUD — no silent fallback to unsigned
|
||||
/// requests.
|
||||
pub bedrock_credentials: Option<Arc<aws_credential_types::Credentials>>,
|
||||
}
|
||||
|
||||
impl AppState {
|
||||
|
|
@ -52,8 +60,19 @@ impl AppState {
|
|||
Ok(Self {
|
||||
config: Arc::new(config),
|
||||
client,
|
||||
bedrock_credentials: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// PR-D1: attach AWS credentials resolved out-of-band (via
|
||||
/// `aws-config`'s default chain at startup). Returns the
|
||||
/// modified state; intended to be chained off `AppState::new`.
|
||||
/// Tests that don't exercise the Bedrock route can leave
|
||||
/// credentials unset (the catch-all paths never read them).
|
||||
pub fn with_bedrock_credentials(mut self, creds: aws_credential_types::Credentials) -> Self {
|
||||
self.bedrock_credentials = Some(Arc::new(creds));
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the axum app. `/healthz` and `/healthz/upstream` are intercepted;
|
||||
|
|
@ -84,6 +103,32 @@ pub fn build_app(state: AppState) -> Router {
|
|||
post(crate::handlers::responses::handle_responses),
|
||||
);
|
||||
|
||||
// PR-D1: native AWS Bedrock InvokeModel route. Mounts only when
|
||||
// `enable_bedrock_native` is on (default). The handler runs the
|
||||
// live-zone compressor over Anthropic-shape bodies, signs with
|
||||
// SigV4, and forwards to the configured Bedrock endpoint. The
|
||||
// `/converse` route mounts the same handler — the wire shape is
|
||||
// identical for `anthropic.claude-*` model IDs (Bedrock just
|
||||
// accepts both legacy `invoke` and modern `converse` paths).
|
||||
if state.config.enable_bedrock_native {
|
||||
router = router
|
||||
.route(
|
||||
"/model/:model_id/invoke",
|
||||
post(crate::bedrock::invoke::handle_invoke),
|
||||
)
|
||||
.route(
|
||||
"/model/:model_id/converse",
|
||||
post(crate::bedrock::invoke::handle_invoke),
|
||||
);
|
||||
} else {
|
||||
tracing::warn!(
|
||||
event = "bedrock_native_disabled",
|
||||
"Bedrock native InvokeModel route disabled by \
|
||||
--enable-bedrock-native=false; Bedrock requests will fall \
|
||||
through to the catch-all (no SigV4 re-signing — fails closed)"
|
||||
);
|
||||
}
|
||||
|
||||
// PR-C4: Conversations API (passthrough-with-instrumentation).
|
||||
// The flag is read once at app-build time so router shape
|
||||
// matches the configured policy. When disabled, requests still
|
||||
|
|
|
|||
|
|
@ -44,11 +44,28 @@ pub async fn start_proxy(upstream: &str) -> ProxyHandle {
|
|||
pub async fn start_proxy_with<F>(upstream: &str, customize: F) -> ProxyHandle
|
||||
where
|
||||
F: FnOnce(&mut Config),
|
||||
{
|
||||
start_proxy_with_state(upstream, customize, |s| s).await
|
||||
}
|
||||
|
||||
/// Start a proxy with both a Config customizer and an AppState
|
||||
/// post-processor. PR-D1: tests that exercise the Bedrock route
|
||||
/// inject credentials via `with_bedrock_credentials` here.
|
||||
#[allow(dead_code)]
|
||||
pub async fn start_proxy_with_state<F, G>(
|
||||
upstream: &str,
|
||||
customize: F,
|
||||
customize_state: G,
|
||||
) -> ProxyHandle
|
||||
where
|
||||
F: FnOnce(&mut Config),
|
||||
G: FnOnce(AppState) -> AppState,
|
||||
{
|
||||
let upstream_url: Url = upstream.parse().expect("valid upstream url");
|
||||
let mut config = Config::for_test(upstream_url);
|
||||
customize(&mut config);
|
||||
let state = AppState::new(config.clone()).expect("app state");
|
||||
let state = customize_state(state);
|
||||
let app = build_app(state).into_make_service_with_connect_info::<SocketAddr>();
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
|
|
|
|||
563
crates/headroom-proxy/tests/integration_bedrock_invoke.rs
Normal file
563
crates/headroom-proxy/tests/integration_bedrock_invoke.rs
Normal file
|
|
@ -0,0 +1,563 @@
|
|||
//! Integration tests for the native Bedrock InvokeModel route
|
||||
//! (Phase D PR-D1).
|
||||
//!
|
||||
//! These tests boot the real Rust proxy in front of a wiremock
|
||||
//! upstream that pretends to be the Bedrock runtime endpoint
|
||||
//! (`https://bedrock-runtime.{region}.amazonaws.com`). The proxy is
|
||||
//! configured with `bedrock_endpoint = wiremock_url` so SigV4-signed
|
||||
//! requests are routed to the mock instead of real AWS — no live
|
||||
//! AWS dependency.
|
||||
//!
|
||||
//! Coverage matrix (per PR-D1 spec, REALIGNMENT/06-phase-D-bedrock-vertex.md):
|
||||
//!
|
||||
//! 1. `native_envelope_round_trip_byte_equal` — small body,
|
||||
//! compression-mode off; bytes round-trip byte-equal upstream.
|
||||
//! 2. `sigv4_signed_correctly_after_compression` — confirms the
|
||||
//! `authorization` header arrives at upstream and the
|
||||
//! `x-amz-content-sha256` matches the (post-compression) body.
|
||||
//! 3. `thinking_block_preserved_through_bedrock` — Anthropic
|
||||
//! `thinking` block round-trips byte-equal with compression off.
|
||||
//! Validates the live-zone dispatcher doesn't strip the block.
|
||||
//! 4. `redacted_thinking_preserved` — `redacted_thinking` block
|
||||
//! round-trips byte-equal.
|
||||
//! 5. `document_block_preserved` — `document` block round-trips.
|
||||
//! 6. `tool_result_array_with_image_preserved` — `tool_result` content
|
||||
//! array containing a base64 `image` block round-trips byte-equal
|
||||
//! when compression is off.
|
||||
//! 7. `stop_sequence_null_only_when_present` — mock the upstream to
|
||||
//! return a Bedrock-shape response that does NOT include
|
||||
//! `stop_sequence`; the proxy must not inject a `null` value for
|
||||
//! it. (This is an end-to-end check that Phase D doesn't regress
|
||||
//! P4-37's hardcoded null.)
|
||||
//! 8. `tool_use_input_byte_equal_preserves_key_order` — `tool_use.input`
|
||||
//! object keys must arrive in the same order they were sent
|
||||
//! (the `serde_json::preserve_order` feature backs this).
|
||||
|
||||
mod common;
|
||||
|
||||
use aws_credential_types::Credentials;
|
||||
use common::start_proxy_with_state;
|
||||
use serde_json::{json, Value};
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use url::Url;
|
||||
use wiremock::matchers::{method, path};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
|
||||
/// What we capture from each upstream request — body + the
|
||||
/// authorization-shaped headers we care about.
|
||||
#[derive(Default, Clone, Debug)]
|
||||
struct CapturedRequest {
|
||||
body: Option<Vec<u8>>,
|
||||
authorization: Option<String>,
|
||||
x_amz_date: Option<String>,
|
||||
x_amz_content_sha256: Option<String>,
|
||||
host: Option<String>,
|
||||
content_type: Option<String>,
|
||||
}
|
||||
|
||||
type Capture = Arc<Mutex<CapturedRequest>>;
|
||||
|
||||
const TEST_MODEL: &str = "anthropic.claude-3-haiku-20240307-v1:0";
|
||||
|
||||
async fn mount_capture_invoke(upstream: &MockServer, response_body: &str) -> Capture {
|
||||
let captured: Capture = Arc::new(Mutex::new(CapturedRequest::default()));
|
||||
let captured_clone = captured.clone();
|
||||
let response_body = response_body.to_string();
|
||||
let model = TEST_MODEL.to_string();
|
||||
Mock::given(method("POST"))
|
||||
.and(path(format!("/model/{model}/invoke")))
|
||||
.respond_with(move |req: &wiremock::Request| {
|
||||
let mut c = captured_clone.lock().unwrap();
|
||||
c.body = Some(req.body.clone());
|
||||
c.authorization = req
|
||||
.headers
|
||||
.get("authorization")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(str::to_string);
|
||||
c.x_amz_date = req
|
||||
.headers
|
||||
.get("x-amz-date")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(str::to_string);
|
||||
c.x_amz_content_sha256 = req
|
||||
.headers
|
||||
.get("x-amz-content-sha256")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(str::to_string);
|
||||
c.host = req
|
||||
.headers
|
||||
.get("host")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(str::to_string);
|
||||
c.content_type = req
|
||||
.headers
|
||||
.get("content-type")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(str::to_string);
|
||||
ResponseTemplate::new(200).set_body_string(response_body.clone())
|
||||
})
|
||||
.mount(upstream)
|
||||
.await;
|
||||
captured
|
||||
}
|
||||
|
||||
fn sha256_hex(bytes: &[u8]) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(bytes);
|
||||
hasher
|
||||
.finalize()
|
||||
.iter()
|
||||
.fold(String::with_capacity(64), |mut acc, b| {
|
||||
use std::fmt::Write as _;
|
||||
let _ = write!(acc, "{b:02x}");
|
||||
acc
|
||||
})
|
||||
}
|
||||
|
||||
#[track_caller]
|
||||
fn assert_byte_equal_sha256(inbound: &[u8], received: &[u8]) {
|
||||
let inbound_hash = sha256_hex(inbound);
|
||||
let received_hash = sha256_hex(received);
|
||||
assert_eq!(
|
||||
inbound.len(),
|
||||
received.len(),
|
||||
"byte length mismatch: inbound={}, upstream-received={}",
|
||||
inbound.len(),
|
||||
received.len(),
|
||||
);
|
||||
assert_eq!(
|
||||
inbound_hash, received_hash,
|
||||
"SHA-256 mismatch: inbound={inbound_hash}, upstream-received={received_hash}",
|
||||
);
|
||||
}
|
||||
|
||||
fn test_credentials() -> Credentials {
|
||||
Credentials::new(
|
||||
"AKIAEXAMPLEAKIDFORTEST",
|
||||
"wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
|
||||
None,
|
||||
None,
|
||||
"test",
|
||||
)
|
||||
}
|
||||
|
||||
/// Boot a proxy pointed at the wiremock upstream as the Bedrock
|
||||
/// endpoint. The fake-upstream URL goes into `bedrock_endpoint`;
|
||||
/// the regular `upstream` field is set to a sentinel because the
|
||||
/// Bedrock route bypasses `forward_http`.
|
||||
async fn bedrock_proxy(
|
||||
upstream: &MockServer,
|
||||
customize: impl FnOnce(&mut headroom_proxy::Config),
|
||||
) -> common::ProxyHandle {
|
||||
let endpoint: Url = upstream.uri().parse().unwrap();
|
||||
start_proxy_with_state(
|
||||
&upstream.uri(),
|
||||
|c| {
|
||||
c.bedrock_endpoint = Some(endpoint);
|
||||
customize(c);
|
||||
},
|
||||
|s| s.with_bedrock_credentials(test_credentials()),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn native_envelope_round_trip_byte_equal() {
|
||||
// Compression off → body must arrive byte-equal at upstream.
|
||||
let upstream = MockServer::start().await;
|
||||
let captured = mount_capture_invoke(&upstream, r#"{"id":"msg_x","content":[]}"#).await;
|
||||
let proxy = bedrock_proxy(&upstream, |c| {
|
||||
c.compression = true;
|
||||
c.compression_mode = headroom_proxy::config::CompressionMode::Off;
|
||||
})
|
||||
.await;
|
||||
|
||||
let payload = json!({
|
||||
"anthropic_version": "bedrock-2023-05-31",
|
||||
"max_tokens": 64,
|
||||
"messages": [
|
||||
{"role": "user", "content": "hi"}
|
||||
]
|
||||
});
|
||||
let body = serde_json::to_vec(&payload).unwrap();
|
||||
let resp = reqwest::Client::new()
|
||||
.post(format!("{}/model/{TEST_MODEL}/invoke", proxy.url()))
|
||||
.header("content-type", "application/json")
|
||||
.body(body.clone())
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 200);
|
||||
|
||||
let got = captured.lock().unwrap().clone();
|
||||
assert_byte_equal_sha256(&body, got.body.as_deref().unwrap());
|
||||
proxy.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sigv4_signed_correctly_after_compression() {
|
||||
// The signature must cover the bytes that actually hit upstream.
|
||||
// We confirm: (a) `authorization` is present, (b)
|
||||
// `x-amz-content-sha256` matches sha256(body received by upstream).
|
||||
let upstream = MockServer::start().await;
|
||||
let captured = mount_capture_invoke(&upstream, r#"{"id":"msg_x","content":[]}"#).await;
|
||||
let proxy = bedrock_proxy(&upstream, |c| {
|
||||
c.compression = true;
|
||||
c.compression_mode = headroom_proxy::config::CompressionMode::LiveZone;
|
||||
})
|
||||
.await;
|
||||
|
||||
let payload = json!({
|
||||
"anthropic_version": "bedrock-2023-05-31",
|
||||
"max_tokens": 64,
|
||||
"messages": [{"role": "user", "content": "hi"}]
|
||||
});
|
||||
let body = serde_json::to_vec(&payload).unwrap();
|
||||
let resp = reqwest::Client::new()
|
||||
.post(format!("{}/model/{TEST_MODEL}/invoke", proxy.url()))
|
||||
.header("content-type", "application/json")
|
||||
.body(body)
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 200);
|
||||
|
||||
let got = captured.lock().unwrap().clone();
|
||||
let auth = got.authorization.expect("authorization header present");
|
||||
assert!(
|
||||
auth.starts_with("AWS4-HMAC-SHA256 "),
|
||||
"authorization must be SigV4-shape; got {auth}"
|
||||
);
|
||||
assert!(
|
||||
auth.contains("Credential=AKIAEXAMPLEAKIDFORTEST/"),
|
||||
"authorization must reference the test access key id; got {auth}"
|
||||
);
|
||||
assert!(
|
||||
auth.contains("/bedrock/aws4_request"),
|
||||
"authorization scope must reference the bedrock service; got {auth}"
|
||||
);
|
||||
assert!(got.x_amz_date.is_some(), "x-amz-date must be present");
|
||||
let body_received = got.body.expect("upstream got body");
|
||||
let expected_sha = sha256_hex(&body_received);
|
||||
assert_eq!(
|
||||
got.x_amz_content_sha256.as_deref(),
|
||||
Some(expected_sha.as_str()),
|
||||
"x-amz-content-sha256 must match sha256 of bytes the upstream received"
|
||||
);
|
||||
proxy.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn thinking_block_preserved_through_bedrock() {
|
||||
// Anthropic `thinking` block: cache hot zone item. With
|
||||
// compression OFF, body round-trips byte-equal upstream. (The
|
||||
// dispatcher only mutates the live zone — the latest user
|
||||
// message — so even with compression on, an assistant `thinking`
|
||||
// block is left alone. We pin the byte-equal contract to the
|
||||
// off-mode path because that's what the litellm Python shim
|
||||
// would have lost — P4-37 evidence.)
|
||||
let upstream = MockServer::start().await;
|
||||
let captured = mount_capture_invoke(&upstream, r#"{"id":"msg_x","content":[]}"#).await;
|
||||
let proxy = bedrock_proxy(&upstream, |c| {
|
||||
c.compression_mode = headroom_proxy::config::CompressionMode::Off;
|
||||
})
|
||||
.await;
|
||||
|
||||
let payload = json!({
|
||||
"anthropic_version": "bedrock-2023-05-31",
|
||||
"max_tokens": 1024,
|
||||
"messages": [
|
||||
{"role": "user", "content": "What's 2+2?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "thinking",
|
||||
"thinking": "The user is asking a basic arithmetic question. 2+2=4.",
|
||||
"signature": "EpYBCkYIBRgCKkAhello_world_signature_payload="
|
||||
},
|
||||
{"type": "text", "text": "4"}
|
||||
]
|
||||
}
|
||||
]
|
||||
});
|
||||
let body = serde_json::to_vec(&payload).unwrap();
|
||||
let resp = reqwest::Client::new()
|
||||
.post(format!("{}/model/{TEST_MODEL}/invoke", proxy.url()))
|
||||
.header("content-type", "application/json")
|
||||
.body(body.clone())
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 200);
|
||||
|
||||
let got = captured.lock().unwrap().clone();
|
||||
assert_byte_equal_sha256(&body, got.body.as_deref().unwrap());
|
||||
let parsed: Value = serde_json::from_slice(got.body.as_deref().unwrap()).unwrap();
|
||||
assert_eq!(parsed["messages"][1]["content"][0]["type"], "thinking");
|
||||
assert_eq!(
|
||||
parsed["messages"][1]["content"][0]["signature"],
|
||||
"EpYBCkYIBRgCKkAhello_world_signature_payload="
|
||||
);
|
||||
proxy.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn redacted_thinking_preserved() {
|
||||
// `redacted_thinking` blocks: opaque encrypted payloads from
|
||||
// the model. Must round-trip BYTE-EQUAL — the proxy never
|
||||
// inspects them.
|
||||
let upstream = MockServer::start().await;
|
||||
let captured = mount_capture_invoke(&upstream, r#"{"id":"msg_x","content":[]}"#).await;
|
||||
let proxy = bedrock_proxy(&upstream, |c| {
|
||||
c.compression_mode = headroom_proxy::config::CompressionMode::Off;
|
||||
})
|
||||
.await;
|
||||
|
||||
let payload = json!({
|
||||
"anthropic_version": "bedrock-2023-05-31",
|
||||
"max_tokens": 1024,
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hi"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "redacted_thinking",
|
||||
"data": "EuYBCogBAaR_o9XJEnEx_3Q9d5z9_redacted_payload"
|
||||
},
|
||||
{"type": "text", "text": "Hello!"}
|
||||
]
|
||||
}
|
||||
]
|
||||
});
|
||||
let body = serde_json::to_vec(&payload).unwrap();
|
||||
let resp = reqwest::Client::new()
|
||||
.post(format!("{}/model/{TEST_MODEL}/invoke", proxy.url()))
|
||||
.header("content-type", "application/json")
|
||||
.body(body.clone())
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 200);
|
||||
|
||||
let got = captured.lock().unwrap().clone();
|
||||
assert_byte_equal_sha256(&body, got.body.as_deref().unwrap());
|
||||
let parsed: Value = serde_json::from_slice(got.body.as_deref().unwrap()).unwrap();
|
||||
assert_eq!(
|
||||
parsed["messages"][1]["content"][0]["type"],
|
||||
"redacted_thinking"
|
||||
);
|
||||
assert_eq!(
|
||||
parsed["messages"][1]["content"][0]["data"],
|
||||
"EuYBCogBAaR_o9XJEnEx_3Q9d5z9_redacted_payload"
|
||||
);
|
||||
proxy.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn document_block_preserved() {
|
||||
// `document` block (PDF or text-document attachment). Has nested
|
||||
// `source.media_type` + `source.data` (base64). The litellm
|
||||
// Python shim drops these silently; the Rust route must NOT.
|
||||
let upstream = MockServer::start().await;
|
||||
let captured = mount_capture_invoke(&upstream, r#"{"id":"msg_x","content":[]}"#).await;
|
||||
let proxy = bedrock_proxy(&upstream, |c| {
|
||||
c.compression_mode = headroom_proxy::config::CompressionMode::Off;
|
||||
})
|
||||
.await;
|
||||
|
||||
let payload = json!({
|
||||
"anthropic_version": "bedrock-2023-05-31",
|
||||
"max_tokens": 256,
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "document",
|
||||
"source": {
|
||||
"type": "base64",
|
||||
"media_type": "application/pdf",
|
||||
"data": "JVBERi0xLjQKJfbk/N8KMSAwIG9iago8PAovVHlwZSAvQ2F0YWxvZwo+PgplbmRvYmoK"
|
||||
},
|
||||
"title": "Quarterly Report",
|
||||
"context": "Q4 2025 financial summary"
|
||||
},
|
||||
{"type": "text", "text": "Summarize."}
|
||||
]
|
||||
}]
|
||||
});
|
||||
let body = serde_json::to_vec(&payload).unwrap();
|
||||
let resp = reqwest::Client::new()
|
||||
.post(format!("{}/model/{TEST_MODEL}/invoke", proxy.url()))
|
||||
.header("content-type", "application/json")
|
||||
.body(body.clone())
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 200);
|
||||
|
||||
let got = captured.lock().unwrap().clone();
|
||||
assert_byte_equal_sha256(&body, got.body.as_deref().unwrap());
|
||||
let parsed: Value = serde_json::from_slice(got.body.as_deref().unwrap()).unwrap();
|
||||
let doc = &parsed["messages"][0]["content"][0];
|
||||
assert_eq!(doc["type"], "document");
|
||||
assert_eq!(doc["source"]["media_type"], "application/pdf");
|
||||
assert_eq!(doc["title"], "Quarterly Report");
|
||||
proxy.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_result_array_with_image_preserved() {
|
||||
// `tool_result.content` is an array; one element is an `image`
|
||||
// block with base64 source. The litellm shim flattened these
|
||||
// to a single string (P4-37). The Rust route must keep the
|
||||
// array shape verbatim.
|
||||
let upstream = MockServer::start().await;
|
||||
let captured = mount_capture_invoke(&upstream, r#"{"id":"msg_x","content":[]}"#).await;
|
||||
let proxy = bedrock_proxy(&upstream, |c| {
|
||||
c.compression_mode = headroom_proxy::config::CompressionMode::Off;
|
||||
})
|
||||
.await;
|
||||
|
||||
let payload = json!({
|
||||
"anthropic_version": "bedrock-2023-05-31",
|
||||
"max_tokens": 256,
|
||||
"messages": [
|
||||
{"role": "user", "content": "Take a screenshot"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{
|
||||
"type": "tool_use",
|
||||
"id": "toolu_xyz",
|
||||
"name": "screenshot",
|
||||
"input": {"region": "full"}
|
||||
}]
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_xyz",
|
||||
"content": [
|
||||
{"type": "text", "text": "Captured at 2026-05-03T12:00:00Z"},
|
||||
{
|
||||
"type": "image",
|
||||
"source": {
|
||||
"type": "base64",
|
||||
"media_type": "image/png",
|
||||
"data": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP8z8DwHwAFAQH/9zJEHwAAAABJRU5ErkJggg=="
|
||||
}
|
||||
}
|
||||
]
|
||||
}]
|
||||
}
|
||||
]
|
||||
});
|
||||
let body = serde_json::to_vec(&payload).unwrap();
|
||||
let resp = reqwest::Client::new()
|
||||
.post(format!("{}/model/{TEST_MODEL}/invoke", proxy.url()))
|
||||
.header("content-type", "application/json")
|
||||
.body(body.clone())
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 200);
|
||||
|
||||
let got = captured.lock().unwrap().clone();
|
||||
assert_byte_equal_sha256(&body, got.body.as_deref().unwrap());
|
||||
let parsed: Value = serde_json::from_slice(got.body.as_deref().unwrap()).unwrap();
|
||||
let tool_result = &parsed["messages"][2]["content"][0];
|
||||
assert_eq!(tool_result["type"], "tool_result");
|
||||
let inner = &tool_result["content"];
|
||||
assert!(inner.is_array(), "tool_result.content must remain an array");
|
||||
assert_eq!(inner[1]["type"], "image");
|
||||
assert_eq!(inner[1]["source"]["media_type"], "image/png");
|
||||
proxy.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stop_sequence_null_only_when_present() {
|
||||
// Synthesise a Bedrock-shape response that does NOT include
|
||||
// `stop_sequence` and confirm the proxy doesn't add `null` for
|
||||
// it. P4-37: the litellm Python shim hardcoded `stop_sequence:
|
||||
// null` in the converted response shape; the Rust path must not
|
||||
// do that — Bedrock's response is forwarded verbatim.
|
||||
let upstream = MockServer::start().await;
|
||||
let response_no_stop_sequence = r#"{"id":"msg_a","type":"message","role":"assistant","model":"claude-3-haiku-20240307","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","usage":{"input_tokens":3,"output_tokens":1}}"#;
|
||||
let _captured = mount_capture_invoke(&upstream, response_no_stop_sequence).await;
|
||||
let proxy = bedrock_proxy(&upstream, |c| {
|
||||
c.compression_mode = headroom_proxy::config::CompressionMode::Off;
|
||||
})
|
||||
.await;
|
||||
|
||||
let payload = json!({
|
||||
"anthropic_version": "bedrock-2023-05-31",
|
||||
"max_tokens": 16,
|
||||
"messages": [{"role": "user", "content": "hi"}]
|
||||
});
|
||||
let body = serde_json::to_vec(&payload).unwrap();
|
||||
let resp = reqwest::Client::new()
|
||||
.post(format!("{}/model/{TEST_MODEL}/invoke", proxy.url()))
|
||||
.header("content-type", "application/json")
|
||||
.body(body)
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 200);
|
||||
let resp_text = resp.text().await.unwrap();
|
||||
let resp_parsed: Value = serde_json::from_str(&resp_text).unwrap();
|
||||
assert!(
|
||||
resp_parsed.get("stop_sequence").is_none(),
|
||||
"stop_sequence must NOT be present on the response when upstream omitted it; got {resp_text}"
|
||||
);
|
||||
assert_eq!(resp_parsed["stop_reason"], "end_turn");
|
||||
proxy.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_use_input_byte_equal_preserves_key_order() {
|
||||
// `tool_use.input` is a JSON object whose key order must
|
||||
// round-trip exactly. P4-43: the litellm shim parsed
|
||||
// `function.arguments` into a dict and re-stringified, breaking
|
||||
// key order. The Rust path uses serde_json's preserve_order
|
||||
// feature throughout — confirm the bytes arrive byte-equal.
|
||||
let upstream = MockServer::start().await;
|
||||
let captured = mount_capture_invoke(&upstream, r#"{"id":"msg_x","content":[]}"#).await;
|
||||
let proxy = bedrock_proxy(&upstream, |c| {
|
||||
c.compression_mode = headroom_proxy::config::CompressionMode::Off;
|
||||
})
|
||||
.await;
|
||||
|
||||
// Hand-craft the body bytes so we can pin exact key order.
|
||||
// BTreeMap-default serialization would alphabetize the keys
|
||||
// (city < country < units); we send the opposite so any
|
||||
// accidental re-encode shows up as a byte mismatch.
|
||||
let body = br#"{"anthropic_version":"bedrock-2023-05-31","max_tokens":64,"messages":[{"role":"assistant","content":[{"type":"tool_use","id":"toolu_zoom","name":"get_weather","input":{"units":"metric","country":"FR","city":"Paris"}}]}]}"#;
|
||||
|
||||
let resp = reqwest::Client::new()
|
||||
.post(format!("{}/model/{TEST_MODEL}/invoke", proxy.url()))
|
||||
.header("content-type", "application/json")
|
||||
.body(body.to_vec())
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 200);
|
||||
|
||||
let got = captured.lock().unwrap().clone();
|
||||
let received = got.body.expect("upstream got body");
|
||||
assert_byte_equal_sha256(body, &received);
|
||||
|
||||
// Defensive: sanity-check the input key order by looking at the
|
||||
// raw substring.
|
||||
let received_str = std::str::from_utf8(&received).unwrap();
|
||||
let units_pos = received_str.find("\"units\"").expect("units present");
|
||||
let country_pos = received_str.find("\"country\"").expect("country present");
|
||||
let city_pos = received_str.find("\"city\"").expect("city present");
|
||||
assert!(
|
||||
units_pos < country_pos && country_pos < city_pos,
|
||||
"tool_use.input key order must be units→country→city; got: {received_str}"
|
||||
);
|
||||
proxy.shutdown().await;
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue