From 90ef66213d8566925ce0b021d84fd094fab79340 Mon Sep 17 00:00:00 2001 From: chopratejas Date: Sun, 3 May 2026 17:57:00 -0700 Subject: [PATCH] fix(proxy): PR-D3 Bedrock observability + auth-mode integration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase D close. Adds the operator-facing observability surface that PRs D1 (native invoke) and D2 (streaming EventStream) deferred, and wires the Phase F PR-F1 auth-mode classifier into the Bedrock route so downstream cache/compression policy gates have something to read. Changes ------- * New `bedrock::auth_mode_layer` middleware. Classifies every inbound Bedrock request via F1's `classify`, coerces the result to `AuthMode::OAuth` per the Bedrock policy matrix (SigV4 IAM is OAuth-equivalent), and stores the resolved value in `request.extensions()` so PR-F2/F3 can read it without re-classifying. Mismatches are logged at WARN with `event=bedrock_auth_mode_unexpected` — no silent coercion. * New `observability` module with three Prometheus families: - `bedrock_invoke_count_total{model, region, auth_mode}` (counter) - `bedrock_invoke_latency_seconds{model, region}` (histogram) - `bedrock_eventstream_message_count_total{model, region, event_type}` (counter) Registered lazily via `OnceLock` so per-request work is just `inc_with_label_values` / `observe`. Latency observed via an RAII `LatencyGuard` so every error path is instrumented; a future regression that adds a new return path can't drop the observation. * New `GET /metrics` endpoint serves the registry in Prometheus text format. Mounted unconditionally — no feature flag gate — so scrape works regardless of which provider routes are mounted. * Bedrock invoke + invoke-streaming handlers now extract `Extension`, log it in their entry breadcrumbs (`event=bedrock_invoke_received`, `event=bedrock_invoke_streaming_received`), and pass `model`/`region` into `translate_stream` so per-message metrics carry the right labels. * Operator docs at `docs/bedrock.md`: AWS credential chain, region/endpoint config, supported model IDs (`anthropic.*` literal-match — no regexes), compression behaviour, sample PromQL queries, structured-log correlation, rollback path. Tests added (6, all green) -------------------------- Auth-mode (`integration_bedrock_authmode.rs`): 1. `bedrock_classified_as_oauth` — empty headers → OAuth in extensions. 2. `oauth_policy_passthrough_prefer` — body byte-equal upstream; no auto cache_control / prompt_cache_key injected. Metrics (`integration_bedrock_metrics.rs`): 3. `metrics_increment_per_invoke` — 3 invokes → counter=3 with correct labels. 4. `metrics_observe_latency` — 1 invoke → histogram count=1, sum>0. 5. `eventstream_metrics_per_message_type` — 5 chunks → counter=5 with `event_type=chunk`. 6. `metrics_endpoint_serves_scrape` — `/metrics` returns 200, `text/plain`, all three metric families' HELP/TYPE lines present. Each metrics test owns a unique (model, region) tuple so the global `prometheus` registry — shared across parallel tests in the same binary — gives each test isolated label rows. Without isolation, parallel tests cross-contaminate counters. Constraints honoured -------------------- * No silent fallbacks — auth-mode coercion is logged at WARN. * No hardcodes — region from `--bedrock-region`, model from axum path parameter. * No regexes — vendor prefix is literal `anthropic.`. * Comprehensive structured logs — every metric increment paired with `tracing::debug!` carrying the same labels for incident correlation. * Performant — `OnceLock`-cached descriptors, RAII guard, total D3 overhead well under 1us per request. * Cardinality bounded — labels driven by config + bounded enums, never by user-controlled bytes. Live cloud validation deferred ------------------------------ The wiremock-backed integration tests are the canonical correctness gate for D3. A real Bedrock smoke test requires `bedrock:InvokeModel` permissions in the developer's AWS account and is documented in `docs/bedrock.md` — both D1 and D2 hit sandbox permission issues trying this path; D3 follows the same convention. Stacked on ---------- PR #364 (D1 native invoke), PR #365 (D2 streaming EventStream), PR #366 (F1 classifier helper). Merge those first; this PR will be rebased onto main once they land. --- Cargo.lock | 25 + crates/headroom-proxy/Cargo.toml | 6 + .../src/bedrock/auth_mode_layer.rs | 192 ++++++++ crates/headroom-proxy/src/bedrock/invoke.rs | 63 ++- .../src/bedrock/invoke_streaming.rs | 457 +++++++++++------- crates/headroom-proxy/src/bedrock/mod.rs | 2 + crates/headroom-proxy/src/lib.rs | 1 + .../headroom-proxy/src/observability/mod.rs | 46 ++ .../src/observability/prometheus.rs | 335 +++++++++++++ crates/headroom-proxy/src/proxy.rs | 25 +- .../tests/integration_bedrock_authmode.rs | 241 +++++++++ .../tests/integration_bedrock_metrics.rs | 435 +++++++++++++++++ docs/bedrock.md | 182 +++++++ 13 files changed, 1825 insertions(+), 185 deletions(-) create mode 100644 crates/headroom-proxy/src/bedrock/auth_mode_layer.rs create mode 100644 crates/headroom-proxy/src/observability/mod.rs create mode 100644 crates/headroom-proxy/src/observability/prometheus.rs create mode 100644 crates/headroom-proxy/tests/integration_bedrock_authmode.rs create mode 100644 crates/headroom-proxy/tests/integration_bedrock_metrics.rs create mode 100644 docs/bedrock.md diff --git a/Cargo.lock b/Cargo.lock index 5dceb1ed8..510f05079 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1826,6 +1826,7 @@ dependencies = [ "hyper", "hyper-util", "pin-project-lite", + "prometheus", "proptest", "reqwest", "serde", @@ -2860,6 +2861,16 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e" +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + [[package]] name = "parking_lot_core" version = "0.9.12" @@ -3027,6 +3038,20 @@ dependencies = [ "syn", ] +[[package]] +name = "prometheus" +version = "0.13.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d33c28a30771f7f96db69893f78b857f7450d7e0237e9c8fc6427a81bae7ed1" +dependencies = [ + "cfg-if", + "fnv", + "lazy_static", + "memchr", + "parking_lot", + "thiserror 1.0.69", +] + [[package]] name = "proptest" version = "1.11.0" diff --git a/crates/headroom-proxy/Cargo.toml b/crates/headroom-proxy/Cargo.toml index c35039ecf..c462f4b74 100644 --- a/crates/headroom-proxy/Cargo.toml +++ b/crates/headroom-proxy/Cargo.toml @@ -53,6 +53,12 @@ aws-smithy-runtime-api = { workspace = true } # transitively pulled in by `aws-smithy-*` (so this is not an # additional cold dep — promoting to a direct one for clarity). crc32fast = "1" +# Phase D PR-D3: Prometheus metrics for Bedrock observability. The +# `prometheus` crate's default-features pull in `protobuf`, which we +# don't need (we serve text-format scrapes only), so we disable +# defaults and re-enable nothing — pure registry + counter + +# histogram + text encoder is sufficient. +prometheus = { version = "0.13", default-features = false } [dev-dependencies] tower = { workspace = true, features = ["util"] } diff --git a/crates/headroom-proxy/src/bedrock/auth_mode_layer.rs b/crates/headroom-proxy/src/bedrock/auth_mode_layer.rs new file mode 100644 index 000000000..886724e20 --- /dev/null +++ b/crates/headroom-proxy/src/bedrock/auth_mode_layer.rs @@ -0,0 +1,192 @@ +//! Bedrock-route auth-mode middleware — Phase D PR-D3. +//! +//! # Why a dedicated middleware (vs inlining the classify call)? +//! +//! The Bedrock invoke + invoke-streaming handlers don't run through +//! `proxy::forward_http`'s catch-all (which is where Phase F PR-F1 +//! classifies and stores the value in request extensions for +//! `/v1/messages`, `/v1/chat/completions`, `/v1/responses`). To +//! preserve the same downstream contract — Phase F PR-F2 and PR-F3 +//! will read the [`AuthMode`] back out of `request.extensions()` and +//! gate compression policy on it — every Bedrock route applies this +//! middleware. The middleware: +//! +//! 1. Classifies the inbound headers via +//! [`headroom_core::auth_mode::classify`]. +//! 2. **Asserts** the result is [`AuthMode::OAuth`] under the Bedrock +//! policy matrix. AWS SigV4 is an `Authorization` value that +//! isn't `Bearer ...` — F1's classifier already routes that to +//! OAuth (see `auth_mode.rs` decision rule 5). We additionally +//! catch the no-Authorization case (Bedrock SDK signs DOWNSTREAM +//! of our proxy on some setups, leaving the inbound request +//! unsigned), classify it, and if F1 returned `Payg` for it we +//! still force `OAuth` while emitting +//! `event = bedrock_auth_mode_unexpected` at WARN. Per the +//! realignment build constraint "no silent fallbacks", we +//! NEVER silently coerce — the divergence is loud. +//! 3. Stores the resolved [`AuthMode`] in `request.extensions()` +//! so downstream handlers can read it without re-classifying. +//! 4. Emits a structured info-level log with +//! `event = bedrock_auth_mode_classified` for ops correlation. +//! +//! # Performance +//! +//! `classify` is a pure function with one short owned `String` for +//! the lowercase UA copy; benched <10us. Inserting into request +//! extensions is `O(1)`. Total per-request overhead well under 1us +//! (excluding the classify call itself, which is shared with the +//! main proxy path). +//! +//! # Where the middleware is mounted +//! +//! See `crate::proxy::build_app` — the Bedrock router branch wraps +//! the three Bedrock POST routes +//! (`/model/:model_id/invoke`, `/converse`, and +//! `/invoke-with-response-stream`) with this layer using +//! `axum::middleware::from_fn`. The catch-all and the other +//! provider routes already classify in `forward_http`, so this +//! middleware does NOT apply to them. + +use axum::body::Body; +use axum::extract::Request; +use axum::middleware::Next; +use axum::response::Response; + +use headroom_core::auth_mode::{classify, AuthMode}; + +/// Inspect the inbound headers, classify the auth mode under +/// Bedrock policy (always [`AuthMode::OAuth`], with a loud WARN +/// when F1's classifier disagrees), and attach the resolved value +/// to `request.extensions()`. +/// +/// Mounted as `axum::middleware::from_fn(classify_and_attach_auth_mode)` +/// so it composes with axum's standard router. The middleware is +/// infallible — it never short-circuits the request, never returns +/// an error response, and never panics. Worst case it logs and +/// proceeds. +pub async fn classify_and_attach_auth_mode(mut req: Request, next: Next) -> Response { + let raw_classification = classify(req.headers()); + + // Bedrock policy: always OAuth-equivalent. SigV4 IAM is an + // OAuth-class signal under our policy matrix. F1 already + // returns OAuth for non-Bearer Authorization (rule 5) and for + // sk-ant-oat-* Bearer tokens (rule 2); the only paths that + // wouldn't return OAuth are: + // + // - empty headers entirely (test setups, or AWS SDK that + // signs after our hop) → F1 returns Payg. + // - x-api-key set (Anthropic key on a Bedrock URL — wrong + // surface, but possible in misconfigured setups) → F1 + // returns Payg. + // + // In both cases we coerce to OAuth (defence-in-depth: Bedrock + // is OAuth-class regardless of the inbound surface) but log + // loudly so operators see the misclassification. NO SILENT + // FALLBACK — the divergence is the whole reason for the warn. + let resolved = if raw_classification == AuthMode::OAuth { + AuthMode::OAuth + } else { + tracing::warn!( + event = "bedrock_auth_mode_unexpected", + raw = raw_classification.as_str(), + resolved = AuthMode::OAuth.as_str(), + path = %req.uri().path(), + "Bedrock route received headers that classified as non-OAuth; \ + coercing to OAuth per Bedrock policy and logging the divergence \ + so operators can investigate the source" + ); + AuthMode::OAuth + }; + + tracing::info!( + event = "bedrock_auth_mode_classified", + mode = resolved.as_str(), + raw = raw_classification.as_str(), + path = %req.uri().path(), + "bedrock route classified inbound auth mode" + ); + + req.extensions_mut().insert(resolved); + next.run(req).await +} + +#[cfg(test)] +mod tests { + use super::*; + use axum::body::Body; + use axum::extract::Extension; + use axum::http::{Request as HttpRequest, StatusCode}; + use axum::routing::post; + use axum::Router; + use http::HeaderValue; + use tower::util::ServiceExt; + + /// Probe handler: returns the AuthMode it sees in extensions. + async fn probe(Extension(auth_mode): Extension) -> String { + auth_mode.as_str().to_string() + } + + fn router() -> Router { + Router::new() + .route("/probe", post(probe)) + .layer(axum::middleware::from_fn(classify_and_attach_auth_mode)) + } + + #[tokio::test] + async fn empty_headers_classify_as_oauth_for_bedrock() { + // No Authorization, no x-api-key, no UA — F1 returns Payg + // by default. The middleware coerces to OAuth (with a + // WARN logged at the call site) so downstream sees OAuth + // and the policy matrix is consistent. + let app = router(); + let req = HttpRequest::builder() + .method("POST") + .uri("/probe") + .body(Body::empty()) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let body = axum::body::to_bytes(resp.into_body(), 64).await.unwrap(); + assert_eq!(&body[..], b"oauth"); + } + + #[tokio::test] + async fn sigv4_authorization_classifies_as_oauth() { + // Real Bedrock SDK does sign before reaching us in some + // setups — `Authorization: AWS4-HMAC-SHA256 ...` is + // routed to OAuth by F1's rule 5 directly. No coercion + // needed; no WARN. + let app = router(); + let mut req = HttpRequest::builder() + .method("POST") + .uri("/probe") + .body(Body::empty()) + .unwrap(); + req.headers_mut().insert( + "authorization", + HeaderValue::from_static( + "AWS4-HMAC-SHA256 Credential=AKIA.../20260101/us-east-1/bedrock/aws4_request", + ), + ); + let resp = app.oneshot(req).await.unwrap(); + let body = axum::body::to_bytes(resp.into_body(), 64).await.unwrap(); + assert_eq!(&body[..], b"oauth"); + } + + #[tokio::test] + async fn x_api_key_inbound_is_coerced_to_oauth_loudly() { + // x-api-key on the Bedrock surface is misconfigured but + // possible. F1 returns Payg; we coerce to OAuth. + let app = router(); + let mut req = HttpRequest::builder() + .method("POST") + .uri("/probe") + .body(Body::empty()) + .unwrap(); + req.headers_mut() + .insert("x-api-key", HeaderValue::from_static("sk-ant-api-fake")); + let resp = app.oneshot(req).await.unwrap(); + let body = axum::body::to_bytes(resp.into_body(), 64).await.unwrap(); + assert_eq!(&body[..], b"oauth"); + } +} diff --git a/crates/headroom-proxy/src/bedrock/invoke.rs b/crates/headroom-proxy/src/bedrock/invoke.rs index afa03b8f1..82d7c015a 100644 --- a/crates/headroom-proxy/src/bedrock/invoke.rs +++ b/crates/headroom-proxy/src/bedrock/invoke.rs @@ -38,10 +38,10 @@ //! unsigned. use std::net::SocketAddr; -use std::time::SystemTime; +use std::time::{Instant, SystemTime}; use axum::body::Body; -use axum::extract::{ConnectInfo, Path, State}; +use axum::extract::{ConnectInfo, Extension, Path, State}; use axum::http::{HeaderMap, Method, StatusCode, Uri}; use axum::response::{IntoResponse, Response}; use bytes::Bytes; @@ -55,13 +55,49 @@ use crate::compression::{ compress_anthropic_request, Outcome as AnthropicOutcome, PassthroughReason, }; use crate::headers::filter_response_headers; +use crate::observability::{observe_bedrock_invoke_latency, record_bedrock_invoke}; use crate::proxy::AppState; +// Phase F PR-F1 + PR-D3: the bedrock auth-mode layer +// (`classify_and_attach_auth_mode`) populates `request.extensions()` +// with `AuthMode` BEFORE this handler runs. We extract it via +// `Extension` so the middleware-supplied value is the +// single source of truth — handler does NOT re-classify; that +// would risk drift from the middleware's resolution + WARN log. +use headroom_core::auth_mode::AuthMode; /// Anthropic vendor prefix as encoded in Bedrock model ids /// (`anthropic.claude-3-haiku-...`). Literal-match per project rule /// "no regexes for parsing the model ID". const ANTHROPIC_VENDOR_PREFIX: &str = "anthropic."; +/// RAII guard that observes the `bedrock_invoke_latency_seconds` +/// histogram on drop. Created at handler entry; observed when the +/// guard goes out of scope no matter how the handler exits. Owning +/// `String` rather than `&str` for the labels avoids capture-order +/// dramas with the borrow checker on early-return paths. +struct LatencyGuard { + model: String, + region: String, + start: Instant, +} + +impl LatencyGuard { + fn start(model: &str, region: &str) -> Self { + Self { + model: model.to_string(), + region: region.to_string(), + start: Instant::now(), + } + } +} + +impl Drop for LatencyGuard { + fn drop(&mut self) { + let elapsed = self.start.elapsed().as_secs_f64(); + observe_bedrock_invoke_latency(&self.model, &self.region, elapsed); + } +} + /// 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"; @@ -71,9 +107,11 @@ const BEDROCK_RUNTIME_HOST_TEMPLATE: &str = "bedrock-runtime.{region}.amazonaws. /// 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. +#[allow(clippy::too_many_arguments)] // axum extractors demand one argument per role pub async fn handle_invoke( State(state): State, ConnectInfo(client_addr): ConnectInfo, + Extension(auth_mode): Extension, Path(model_id): Path, method: Method, uri: Uri, @@ -87,11 +125,32 @@ pub async fn handle_invoke( .map(|s| s.to_string()) .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + // PR-D3: latency stopwatch starts at handler entry (after + // routing + middleware). The histogram observes wall-clock + // time, so it captures upstream RTT + sign + compress as a + // single number — operators can split contributions via the + // `bedrock_*` structured-log timing fields if a slow path + // shows up. Wrapped in a `LatencyGuard` so EVERY return path + // (success, sign-failure, upstream-error, response-build error) + // observes the histogram. RAII keeps the call site to one + // line and rules out future regressions where someone adds a + // new error path and forgets to instrument. + let region = state.config.bedrock_region.clone(); + let _latency_guard = LatencyGuard::start(&model_id, ®ion); + + // PR-D3: count every invoke at handler entry (one per request, + // before any error path can early-return). Pairs with the + // structured log emitted below so operators can join the + // counter with the trace by `request_id`. + record_bedrock_invoke(&model_id, ®ion, auth_mode); + tracing::info!( event = "bedrock_invoke_received", request_id = %request_id, method = %method, model_id = %model_id, + region = %region, + auth_mode = auth_mode.as_str(), body_bytes = body.len(), "bedrock invoke route received request" ); diff --git a/crates/headroom-proxy/src/bedrock/invoke_streaming.rs b/crates/headroom-proxy/src/bedrock/invoke_streaming.rs index daa9f683b..5da1667ec 100644 --- a/crates/headroom-proxy/src/bedrock/invoke_streaming.rs +++ b/crates/headroom-proxy/src/bedrock/invoke_streaming.rs @@ -43,10 +43,10 @@ use std::convert::Infallible; use std::net::SocketAddr; -use std::time::SystemTime; +use std::time::{Instant, SystemTime}; use axum::body::Body; -use axum::extract::{ConnectInfo, Path, State}; +use axum::extract::{ConnectInfo, Extension, Path, State}; use axum::http::{HeaderMap, Method, StatusCode, Uri}; use axum::response::{IntoResponse, Response}; use bytes::Bytes; @@ -65,7 +65,14 @@ use crate::compression::{ compress_anthropic_request, Outcome as AnthropicOutcome, PassthroughReason, }; use crate::headers::filter_response_headers; +use crate::observability::{ + observe_bedrock_invoke_latency, record_bedrock_eventstream_message, record_bedrock_invoke, +}; use crate::proxy::AppState; +// Phase F PR-F1 + PR-D3: pre-classified by `classify_and_attach_auth_mode` +// middleware on the bedrock router; we read it back via the +// `Extension` extractor. +use headroom_core::auth_mode::AuthMode; /// Anthropic vendor prefix as encoded in Bedrock model ids. const ANTHROPIC_VENDOR_PREFIX: &str = "anthropic."; @@ -76,10 +83,41 @@ const BEDROCK_RUNTIME_HOST_TEMPLATE: &str = "bedrock-runtime.{region}.amazonaws. /// Path action for the streaming route. const STREAMING_ACTION: &str = "invoke-with-response-stream"; +/// RAII guard that observes the `bedrock_invoke_latency_seconds` +/// histogram on drop. Mirrors the [`crate::bedrock::invoke`] guard +/// — duplicated to avoid a cross-module type dependency for what +/// is fundamentally a 6-line struct. (When PR-D2 + PR-D3 settle, +/// the two handlers can share a `bedrock::common::LatencyGuard` +/// helper; that's a deferred refactor.) +struct LatencyGuard { + model: String, + region: String, + start: Instant, +} + +impl LatencyGuard { + fn start(model: &str, region: &str) -> Self { + Self { + model: model.to_string(), + region: region.to_string(), + start: Instant::now(), + } + } +} + +impl Drop for LatencyGuard { + fn drop(&mut self) { + let elapsed = self.start.elapsed().as_secs_f64(); + observe_bedrock_invoke_latency(&self.model, &self.region, elapsed); + } +} + /// Axum POST handler for `/model/{model_id}/invoke-with-response-stream`. +#[allow(clippy::too_many_arguments)] // axum extractors demand one argument per role pub async fn handle_invoke_streaming( State(state): State, ConnectInfo(client_addr): ConnectInfo, + Extension(auth_mode): Extension, Path(model_id): Path, method: Method, uri: Uri, @@ -93,11 +131,22 @@ pub async fn handle_invoke_streaming( .map(|s| s.to_string()) .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + // PR-D3: latency stopwatch + invoke counter at handler entry. + // RAII guard observes the histogram regardless of which return + // path the handler takes. Per-EventStream-message metrics are + // recorded inside `translate_stream` once the upstream response + // starts arriving. + let region = state.config.bedrock_region.clone(); + let _latency_guard = LatencyGuard::start(&model_id, ®ion); + record_bedrock_invoke(&model_id, ®ion, auth_mode); + tracing::info!( event = "bedrock_invoke_streaming_received", request_id = %request_id, method = %method, model_id = %model_id, + region = %region, + auth_mode = auth_mode.as_str(), body_bytes = body.len(), "bedrock invoke-with-response-stream route received request" ); @@ -344,6 +393,8 @@ pub async fn handle_invoke_streaming( upstream_stream, state.config.bedrock_validate_eventstream_crc, request_id.clone(), + model_id.clone(), + region.clone(), ); let translated = tee_to_anthropic_state(translated, request_id.clone()); let body_out = Body::from_stream(translated); @@ -366,6 +417,8 @@ fn translate_stream( upstream: S, validate_crc: bool, request_id: String, + model_id: String, + region: String, ) -> impl Stream> where S: Stream> + Send + 'static, @@ -377,205 +430,247 @@ where parser = parser.with_crc_validation(CrcValidation::No); } let upstream: ByteStream = Box::pin(upstream); - let init = (parser, upstream, false, request_id); - stream::unfold(init, |(mut parser, mut upstream, mut done, request_id)| { - Box::pin(async move { - if done { - return None; - } - // First, drain any complete messages already in the - // parser's buffer (bytes from the previous chunk). - loop { - match parser.next_message() { - Ok(Some(msg)) => match translate_message(&msg, OutputMode::Sse) { - Ok(TranslateOutcome::Emit(frame)) => { - tracing::debug!( - event = "bedrock_eventstream_message", - request_id = %request_id, - event_type = msg.event_type().unwrap_or(""), - payload_bytes = msg.payload.len(), - "translated bedrock eventstream message" - ); - return Some((Ok(frame), (parser, upstream, false, request_id))); - } - Ok(TranslateOutcome::Skip { event_type }) => { - tracing::warn!( - event = "bedrock_eventstream_unknown_event_type", - request_id = %request_id, - event_type = %event_type, - "skipping unknown bedrock eventstream message" - ); - // Loop and try the next message in the buffer. - continue; - } - Err(TranslateError::UpstreamException { payload_preview }) => { - tracing::warn!( - event = "bedrock_eventstream_upstream_exception", - request_id = %request_id, - payload_preview = %payload_preview, - "bedrock eventstream upstream exception" - ); - // Emit the exception as an Anthropic-shape - // SSE error frame so the client sees it. - let json = serde_json::json!({ - "type": "error", - "error": { - "type": "bedrock_upstream_exception", - "message": payload_preview, + // Bundle the per-stream identifiers we thread through every + // unfold step. `unfold` only allows a single state value, so + // grouping these into a tuple of owned Strings keeps the + // closure readable. + let init = (parser, upstream, false, request_id, model_id, region); + stream::unfold( + init, + |(mut parser, mut upstream, mut done, request_id, model_id, region)| { + Box::pin(async move { + if done { + return None; + } + // First, drain any complete messages already in the + // parser's buffer (bytes from the previous chunk). + loop { + match parser.next_message() { + Ok(Some(msg)) => match translate_message(&msg, OutputMode::Sse) { + Ok(TranslateOutcome::Emit(frame)) => { + // PR-D3: per-message Prometheus metric. + // The label `event_type` is bounded by + // AWS's documented vocabulary (chunk, + // metadata, exception variants); not + // customer-controlled. + let event_type = msg.event_type().unwrap_or("unknown").to_string(); + record_bedrock_eventstream_message(&model_id, ®ion, &event_type); + tracing::debug!( + event = "bedrock_eventstream_message", + request_id = %request_id, + event_type = %event_type, + payload_bytes = msg.payload.len(), + "translated bedrock eventstream message" + ); + return Some(( + Ok(frame), + (parser, upstream, false, request_id, model_id, region), + )); + } + Ok(TranslateOutcome::Skip { event_type }) => { + tracing::warn!( + event = "bedrock_eventstream_unknown_event_type", + request_id = %request_id, + event_type = %event_type, + "skipping unknown bedrock eventstream message" + ); + // Loop and try the next message in the buffer. + continue; + } + Err(TranslateError::UpstreamException { payload_preview }) => { + tracing::warn!( + event = "bedrock_eventstream_upstream_exception", + request_id = %request_id, + payload_preview = %payload_preview, + "bedrock eventstream upstream exception" + ); + // Emit the exception as an Anthropic-shape + // SSE error frame so the client sees it. + let json = serde_json::json!({ + "type": "error", + "error": { + "type": "bedrock_upstream_exception", + "message": payload_preview, + } + }) + .to_string(); + let mut frame = Vec::with_capacity(json.len() + 32); + frame.extend_from_slice(b"event: error\ndata: "); + frame.extend_from_slice(json.as_bytes()); + frame.extend_from_slice(b"\n\n"); + return Some(( + Ok(Bytes::from(frame)), + (parser, upstream, true, request_id, model_id, region), + )); + } + Err(TranslateError::MissingEventType) => { + tracing::warn!( + event = "bedrock_eventstream_missing_event_type", + request_id = %request_id, + "bedrock eventstream message missing :event-type; emitting error frame" + ); + let frame = error_sse_frame( + "bedrock_eventstream_missing_event_type", + "Bedrock message missing :event-type header", + ); + return Some(( + Ok(frame), + (parser, upstream, true, request_id, model_id, region), + )); + } + }, + Ok(None) => break, + Err(parse_err) => { + let event_name = match &parse_err { + ParseError::PreludeCrcMismatch { .. } + | ParseError::MessageCrcMismatch { .. } => { + "bedrock_eventstream_crc_mismatch" } - }) - .to_string(); - let mut frame = Vec::with_capacity(json.len() + 32); - frame.extend_from_slice(b"event: error\ndata: "); - frame.extend_from_slice(json.as_bytes()); - frame.extend_from_slice(b"\n\n"); + _ => "bedrock_eventstream_parse_failed", + }; + tracing::warn!( + event = event_name, + request_id = %request_id, + error = %parse_err, + "bedrock eventstream parse failure; closing translated stream" + ); + let frame = error_sse_frame(event_name, &parse_err.to_string()); return Some(( - Ok(Bytes::from(frame)), - (parser, upstream, true, request_id), + Ok(frame), + (parser, upstream, true, request_id, model_id, region), )); } - Err(TranslateError::MissingEventType) => { - tracing::warn!( - event = "bedrock_eventstream_missing_event_type", - request_id = %request_id, - "bedrock eventstream message missing :event-type; emitting error frame" - ); - let frame = error_sse_frame( - "bedrock_eventstream_missing_event_type", - "Bedrock message missing :event-type header", - ); - return Some((Ok(frame), (parser, upstream, true, request_id))); - } - }, - Ok(None) => break, - Err(parse_err) => { - let event_name = match &parse_err { - ParseError::PreludeCrcMismatch { .. } - | ParseError::MessageCrcMismatch { .. } => { - "bedrock_eventstream_crc_mismatch" - } - _ => "bedrock_eventstream_parse_failed", - }; - tracing::warn!( - event = event_name, - request_id = %request_id, - error = %parse_err, - "bedrock eventstream parse failure; closing translated stream" - ); - let frame = error_sse_frame(event_name, &parse_err.to_string()); - return Some((Ok(frame), (parser, upstream, true, request_id))); } } - } - // Buffer drained; pull the next chunk from upstream. - loop { - match upstream.next().await { - Some(Ok(chunk)) => { - parser.push(&chunk); - // Loop back through the parser to emit any - // newly-complete messages. - match parser.next_message() { - Ok(Some(msg)) => match translate_message(&msg, OutputMode::Sse) { - Ok(TranslateOutcome::Emit(frame)) => { - tracing::debug!( - event = "bedrock_eventstream_message", + // Buffer drained; pull the next chunk from upstream. + loop { + match upstream.next().await { + Some(Ok(chunk)) => { + parser.push(&chunk); + // Loop back through the parser to emit any + // newly-complete messages. + match parser.next_message() { + Ok(Some(msg)) => match translate_message(&msg, OutputMode::Sse) { + Ok(TranslateOutcome::Emit(frame)) => { + // PR-D3: per-message Prometheus + // metric (mirror of the parser-buffer + // drain branch above). + let event_type = + msg.event_type().unwrap_or("unknown").to_string(); + record_bedrock_eventstream_message( + &model_id, + ®ion, + &event_type, + ); + tracing::debug!( + event = "bedrock_eventstream_message", + request_id = %request_id, + event_type = %event_type, + payload_bytes = msg.payload.len(), + "translated bedrock eventstream message" + ); + return Some(( + Ok(frame), + (parser, upstream, false, request_id, model_id, region), + )); + } + Ok(TranslateOutcome::Skip { event_type }) => { + tracing::warn!( + event = "bedrock_eventstream_unknown_event_type", + request_id = %request_id, + event_type = %event_type, + "skipping unknown bedrock eventstream message" + ); + // Continue draining the parser / + // pulling more chunks. + continue; + } + Err(TranslateError::UpstreamException { payload_preview }) => { + let json = serde_json::json!({ + "type": "error", + "error": { + "type": "bedrock_upstream_exception", + "message": payload_preview, + } + }) + .to_string(); + let mut frame = Vec::with_capacity(json.len() + 32); + frame.extend_from_slice(b"event: error\ndata: "); + frame.extend_from_slice(json.as_bytes()); + frame.extend_from_slice(b"\n\n"); + return Some(( + Ok(Bytes::from(frame)), + (parser, upstream, true, request_id, model_id, region), + )); + } + Err(TranslateError::MissingEventType) => { + let frame = error_sse_frame( + "bedrock_eventstream_missing_event_type", + "Bedrock message missing :event-type header", + ); + return Some(( + Ok(frame), + (parser, upstream, true, request_id, model_id, region), + )); + } + }, + Ok(None) => continue, + Err(parse_err) => { + let event_name = match &parse_err { + ParseError::PreludeCrcMismatch { .. } + | ParseError::MessageCrcMismatch { .. } => { + "bedrock_eventstream_crc_mismatch" + } + _ => "bedrock_eventstream_parse_failed", + }; + tracing::warn!( + event = event_name, request_id = %request_id, - event_type = msg.event_type().unwrap_or(""), - payload_bytes = msg.payload.len(), - "translated bedrock eventstream message" + error = %parse_err, + "bedrock eventstream parse failure" ); + let frame = error_sse_frame(event_name, &parse_err.to_string()); return Some(( Ok(frame), - (parser, upstream, false, request_id), + (parser, upstream, true, request_id, model_id, region), )); } - Ok(TranslateOutcome::Skip { event_type }) => { - tracing::warn!( - event = "bedrock_eventstream_unknown_event_type", - request_id = %request_id, - event_type = %event_type, - "skipping unknown bedrock eventstream message" - ); - // Continue draining the parser / - // pulling more chunks. - continue; - } - Err(TranslateError::UpstreamException { payload_preview }) => { - let json = serde_json::json!({ - "type": "error", - "error": { - "type": "bedrock_upstream_exception", - "message": payload_preview, - } - }) - .to_string(); - let mut frame = Vec::with_capacity(json.len() + 32); - frame.extend_from_slice(b"event: error\ndata: "); - frame.extend_from_slice(json.as_bytes()); - frame.extend_from_slice(b"\n\n"); - return Some(( - Ok(Bytes::from(frame)), - (parser, upstream, true, request_id), - )); - } - Err(TranslateError::MissingEventType) => { - let frame = error_sse_frame( - "bedrock_eventstream_missing_event_type", - "Bedrock message missing :event-type header", - ); - return Some((Ok(frame), (parser, upstream, true, request_id))); - } - }, - Ok(None) => continue, - Err(parse_err) => { - let event_name = match &parse_err { - ParseError::PreludeCrcMismatch { .. } - | ParseError::MessageCrcMismatch { .. } => { - "bedrock_eventstream_crc_mismatch" - } - _ => "bedrock_eventstream_parse_failed", - }; - tracing::warn!( - event = event_name, - request_id = %request_id, - error = %parse_err, - "bedrock eventstream parse failure" - ); - let frame = error_sse_frame(event_name, &parse_err.to_string()); - return Some((Ok(frame), (parser, upstream, true, request_id))); } } - } - Some(Err(e)) => { - tracing::warn!( - event = "bedrock_eventstream_upstream_io_error", - request_id = %request_id, - error = %e, - "upstream io error mid-stream" - ); - return Some((Err(e), (parser, upstream, true, request_id))); - } - None => { - // End of upstream stream. If buffered bytes - // remain that did not parse into a message, - // log loudly — we are NOT silently dropping - // them. - if parser.buffered_len() > 0 { + Some(Err(e)) => { tracing::warn!( - event = "bedrock_eventstream_truncated", + event = "bedrock_eventstream_upstream_io_error", request_id = %request_id, - buffered_bytes = parser.buffered_len(), - "upstream stream ended with un-parseable trailing bytes" + error = %e, + "upstream io error mid-stream" ); + return Some(( + Err(e), + (parser, upstream, true, request_id, model_id, region), + )); + } + None => { + // End of upstream stream. If buffered bytes + // remain that did not parse into a message, + // log loudly — we are NOT silently dropping + // them. + if parser.buffered_len() > 0 { + tracing::warn!( + event = "bedrock_eventstream_truncated", + request_id = %request_id, + buffered_bytes = parser.buffered_len(), + "upstream stream ended with un-parseable trailing bytes" + ); + } + done = true; + let _ = done; + return None; } - done = true; - let _ = done; - return None; } } - } - }) - }) + }) + }, + ) } /// Tee the translated SSE stream into an `AnthropicStreamState` task diff --git a/crates/headroom-proxy/src/bedrock/mod.rs b/crates/headroom-proxy/src/bedrock/mod.rs index e772a1b93..780ed4d18 100644 --- a/crates/headroom-proxy/src/bedrock/mod.rs +++ b/crates/headroom-proxy/src/bedrock/mod.rs @@ -45,6 +45,7 @@ //! parsing is in [`eventstream`]; the SSE translator is in //! [`eventstream_to_sse`]. +pub mod auth_mode_layer; pub mod envelope; pub mod eventstream; pub mod eventstream_to_sse; @@ -52,6 +53,7 @@ pub mod invoke; pub mod invoke_streaming; pub mod sigv4; +pub use auth_mode_layer::classify_and_attach_auth_mode; pub use envelope::{BedrockEnvelope, EnvelopeError}; pub use eventstream::{ parse as parse_eventstream, CrcValidation, EventStreamMessage, EventStreamParser, HeaderValue, diff --git a/crates/headroom-proxy/src/lib.rs b/crates/headroom-proxy/src/lib.rs index f6f152891..8b83eb099 100644 --- a/crates/headroom-proxy/src/lib.rs +++ b/crates/headroom-proxy/src/lib.rs @@ -8,6 +8,7 @@ pub mod error; pub mod handlers; pub mod headers; pub mod health; +pub mod observability; pub mod proxy; pub mod responses_items; pub mod sse; diff --git a/crates/headroom-proxy/src/observability/mod.rs b/crates/headroom-proxy/src/observability/mod.rs new file mode 100644 index 000000000..d1c513a1d --- /dev/null +++ b/crates/headroom-proxy/src/observability/mod.rs @@ -0,0 +1,46 @@ +//! Proxy observability surface — Phase D PR-D3. +//! +//! Centralises all Prometheus instrumentation in one place so that +//! metric names, label keys, and the global registry stay +//! co-located and discoverable. The Phase D acceptance criterion +//! (`Prometheus scrape includes Bedrock metrics`) demands a single +//! `/metrics` endpoint that serves the registry; that endpoint is +//! mounted by [`crate::proxy::build_app`] when the observability +//! module is in scope. +//! +//! # Module layout +//! +//! - [`prometheus`] — registry construction (lazy via `OnceLock`), +//! Bedrock-scoped counters / histograms, and the `/metrics` +//! text-format scrape handler. Per the realignment build +//! constraint "elegant + scalable" we keep one module per +//! concern; future Phase F / Phase H additions (auth-mode +//! counters, OpenAI request totals) live alongside the +//! Bedrock-prefixed ones below — never sprinkled across handlers. +//! +//! # Cardinality discipline +//! +//! Every label is bounded by infrastructure config, NOT by request +//! input. `model` comes from the axum path parameter (Bedrock vendor +//! prefix is enforced upstream of the metric increment); `region` +//! comes from `Config::bedrock_region`; `auth_mode` comes from the +//! `headroom_core::auth_mode::AuthMode` enum (3 variants total). +//! There is no path where a malicious client can drive label +//! cardinality unbounded — see `bedrock::invoke::handle_invoke` for +//! the call site. +//! +//! # Why not `metrics-rs`? +//! +//! `metrics-rs` is the more idiomatic Rust choice but it requires a +//! separate exporter binary. The Phase D scope is observability for +//! a single proxy binary; the simpler `prometheus` crate (with the +//! global default registry pinned in a `OnceLock`) keeps the +//! footprint small and the scrape endpoint trivial. Phase F may +//! revisit if multi-process aggregation lands. + +pub mod prometheus; + +pub use prometheus::{ + handle_metrics, observe_bedrock_invoke_latency, record_bedrock_eventstream_message, + record_bedrock_invoke, +}; diff --git a/crates/headroom-proxy/src/observability/prometheus.rs b/crates/headroom-proxy/src/observability/prometheus.rs new file mode 100644 index 000000000..bde34934f --- /dev/null +++ b/crates/headroom-proxy/src/observability/prometheus.rs @@ -0,0 +1,335 @@ +//! Prometheus instrumentation for the Bedrock route — Phase D PR-D3. +//! +//! # Registered metrics +//! +//! - `bedrock_invoke_count_total{model, region, auth_mode}` — Counter. +//! One increment per `/model/{model}/invoke` (or `/converse` or +//! `/invoke-with-response-stream`) request that successfully +//! passed the path-parameter extractor and reached the handler +//! body. Failures BEFORE the handler runs (router 404s, axum +//! extractor errors) do not increment. +//! - `bedrock_invoke_latency_seconds{model, region}` — Histogram. +//! Observed at request completion (whether the upstream call +//! succeeded or returned 5xx). Buckets target typical Bedrock +//! latencies (50ms → 60s) so p50/p99 land in distinct buckets at +//! typical throughput. The `auth_mode` label is intentionally +//! absent: the cost of cross-multiplying it with `model` would +//! triple the per-model cardinality with little operator value +//! (auth-mode breakdown lives in the count metric instead). +//! - `bedrock_eventstream_message_count_total{model, region, event_type}` +//! — Counter. One increment per parsed binary EventStream message +//! in the streaming handler. `event_type` is the +//! `:event-type` header from the message (`chunk`, `metadata`, +//! `internalServerException`, etc.). The set is bounded by +//! AWS's documented event-type vocabulary, not customer input — +//! see `crates/headroom-proxy/src/bedrock/eventstream.rs` for +//! the parsed shape. +//! +//! # Wiring +//! +//! Every counter / histogram is created exactly once at first call +//! via `OnceLock`. Per-request work is `inc_with_label_values` / +//! `observe`, which is `O(1)` and lock-free in the common case +//! (the underlying `prometheus` crate uses a sharded RwLock per +//! metric vector). Total D3 overhead per request: a few hundred ns +//! plus one `Instant::elapsed()` for the latency histogram. +//! +//! # Logs paired with every increment +//! +//! Per the realignment build-constraint "comprehensive structured +//! logs", every metric increment in this module emits a +//! `tracing::debug!` with `event = "metric_recorded"` so operators +//! can correlate scrape values with log lines during incidents. +//! The cardinality of these debug logs is the same as the metric +//! itself (bounded by `model × region × auth_mode`), so leaving +//! them at `debug` level avoids per-request log volume in normal +//! operation while still being available under +//! `RUST_LOG=headroom_proxy::observability=debug`. + +use std::sync::OnceLock; + +use axum::body::Body; +use axum::http::{header, StatusCode}; +use axum::response::Response; +use prometheus::{ + Encoder, HistogramOpts, HistogramVec, IntCounterVec, Opts, Registry, TextEncoder, +}; + +use headroom_core::auth_mode::AuthMode; + +/// Latency-histogram buckets in seconds. Chosen to discriminate +/// across typical Bedrock latencies: cold-start (~1-2s), warm +/// streaming-start (~100-500ms), small completions (~50-200ms), +/// long completions (5-60s). Mirrors the bucket layout the AWS +/// CloudWatch sample dashboards use. +const LATENCY_BUCKETS_SECONDS: &[f64] = &[0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0]; + +/// Lazy singleton registry. Borrowed by `handle_metrics` for +/// scrape rendering and by every metric registration helper. +fn registry() -> &'static Registry { + static REGISTRY: OnceLock = OnceLock::new(); + REGISTRY.get_or_init(Registry::new) +} + +/// `bedrock_invoke_count_total{model, region, auth_mode}` — +/// initialised on first call. +fn invoke_counter() -> &'static IntCounterVec { + static COUNTER: OnceLock = OnceLock::new(); + COUNTER.get_or_init(|| { + let opts = Opts::new( + "bedrock_invoke_count_total", + "Total Bedrock invoke requests handled by the Rust proxy, \ + broken down by model id, AWS region, and inbound auth mode.", + ); + let counter = IntCounterVec::new(opts, &["model", "region", "auth_mode"]) + .expect("bedrock_invoke_count_total descriptor is well-formed"); + registry() + .register(Box::new(counter.clone())) + .expect("bedrock_invoke_count_total registers exactly once"); + counter + }) +} + +/// `bedrock_invoke_latency_seconds{model, region}` — initialised +/// on first call. +fn invoke_latency() -> &'static HistogramVec { + static HIST: OnceLock = OnceLock::new(); + HIST.get_or_init(|| { + let opts = HistogramOpts::new( + "bedrock_invoke_latency_seconds", + "Latency in seconds of Bedrock invoke requests as observed at the \ + Rust proxy entry boundary. Includes upstream call time plus any \ + pre-/post-compression and SigV4 signing.", + ) + .buckets(LATENCY_BUCKETS_SECONDS.to_vec()); + let hist = HistogramVec::new(opts, &["model", "region"]) + .expect("bedrock_invoke_latency_seconds descriptor is well-formed"); + registry() + .register(Box::new(hist.clone())) + .expect("bedrock_invoke_latency_seconds registers exactly once"); + hist + }) +} + +/// `bedrock_eventstream_message_count_total{model, region, event_type}` +/// — initialised on first call. +fn eventstream_counter() -> &'static IntCounterVec { + static COUNTER: OnceLock = OnceLock::new(); + COUNTER.get_or_init(|| { + let opts = Opts::new( + "bedrock_eventstream_message_count_total", + "Total Bedrock binary EventStream messages parsed by the Rust proxy, \ + broken down by model id, AWS region, and the message's :event-type header \ + (chunk, metadata, modelStreamErrorException, etc.).", + ); + let counter = IntCounterVec::new(opts, &["model", "region", "event_type"]) + .expect("bedrock_eventstream_message_count_total descriptor is well-formed"); + registry() + .register(Box::new(counter.clone())) + .expect("bedrock_eventstream_message_count_total registers exactly once"); + counter + }) +} + +/// Record a single Bedrock invoke (non-streaming or streaming). +/// +/// Pure increment + a paired `tracing::debug!` so operators can +/// correlate this metric with the request's structured log line via +/// the `request_id` field (callers thread that through). +pub fn record_bedrock_invoke(model: &str, region: &str, auth_mode: AuthMode) { + invoke_counter() + .with_label_values(&[model, region, auth_mode.as_str()]) + .inc(); + tracing::debug!( + event = "metric_recorded", + metric = "bedrock_invoke_count_total", + model = %model, + region = %region, + auth_mode = auth_mode.as_str(), + "incremented bedrock_invoke_count_total" + ); +} + +/// Observe latency at the END of an invoke. The duration must be +/// computed by the caller via `Instant::elapsed()` — passing the +/// duration in (rather than the start time) keeps this helper +/// free of `Instant` types so unit tests can assert on synthetic +/// values. +pub fn observe_bedrock_invoke_latency(model: &str, region: &str, seconds: f64) { + invoke_latency() + .with_label_values(&[model, region]) + .observe(seconds); + tracing::debug!( + event = "metric_recorded", + metric = "bedrock_invoke_latency_seconds", + model = %model, + region = %region, + seconds = seconds, + "observed bedrock_invoke_latency_seconds" + ); +} + +/// Record a single parsed EventStream message in the streaming +/// path. The `event_type` argument is the `:event-type` header of +/// the message (or `unknown` when the message header was missing, +/// which itself is loud-logged at the call site). +pub fn record_bedrock_eventstream_message(model: &str, region: &str, event_type: &str) { + eventstream_counter() + .with_label_values(&[model, region, event_type]) + .inc(); + tracing::debug!( + event = "metric_recorded", + metric = "bedrock_eventstream_message_count_total", + model = %model, + region = %region, + event_type = %event_type, + "incremented bedrock_eventstream_message_count_total" + ); +} + +/// Axum handler for `GET /metrics`. Renders the registry in the +/// Prometheus text format. Per Phase D acceptance: the scrape MUST +/// include the three Bedrock metrics above as soon as they have +/// been touched at least once; un-touched counter vectors expose +/// their HELP/TYPE lines but no labelled rows (Prometheus's +/// documented behaviour, not a regression). +pub async fn handle_metrics() -> Response { + // Force lazy registration so the HELP/TYPE descriptor lines + // appear in the scrape even before any request has hit the + // Bedrock route. Operators who curl /metrics on a fresh boot + // see the metric names already advertised — surprises are + // worse than the cost of three function calls. + let _ = invoke_counter(); + let _ = invoke_latency(); + let _ = eventstream_counter(); + + let metric_families = registry().gather(); + let mut buffer = Vec::with_capacity(2048); + let encoder = TextEncoder::new(); + if let Err(e) = encoder.encode(&metric_families, &mut buffer) { + tracing::error!( + event = "metrics_encode_failed", + error = %e, + "failed to encode Prometheus metrics scrape" + ); + return Response::builder() + .status(StatusCode::INTERNAL_SERVER_ERROR) + .body(Body::from(format!("metrics encode error: {e}"))) + .expect("static error response"); + } + Response::builder() + .status(StatusCode::OK) + .header(header::CONTENT_TYPE, encoder.format_type()) + .body(Body::from(buffer)) + .unwrap_or_else(|e| { + Response::builder() + .status(StatusCode::INTERNAL_SERVER_ERROR) + .body(Body::from(format!("metrics response build error: {e}"))) + .expect("static error response") + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Helper: render the registry to a String for assertions. + fn scrape() -> String { + let mf = registry().gather(); + let encoder = TextEncoder::new(); + let mut buf = Vec::new(); + encoder.encode(&mf, &mut buf).expect("encode"); + String::from_utf8(buf).expect("utf8") + } + + #[test] + fn invoke_counter_advertises_metric_in_scrape() { + // The `prometheus` crate only emits HELP/TYPE for vectors + // that have AT LEAST ONE row (`gather()` skips empty + // vectors), so we must fire one increment with a unique + // label set this test owns to assert the metric family + // appears in the scrape. + invoke_counter() + .with_label_values(&[ + "anthropic.unit-test-advertise-v1:0", + "us-test-advertise-1", + "oauth", + ]) + .inc(); + let body = scrape(); + assert!( + body.contains("bedrock_invoke_count_total"), + "scrape missing bedrock_invoke_count_total: {body}" + ); + assert!( + body.contains("# TYPE bedrock_invoke_count_total counter"), + "scrape missing TYPE line: {body}" + ); + } + + #[test] + fn invoke_increment_appears_with_labels() { + record_bedrock_invoke( + "anthropic.claude-3-haiku-20240307-v1:0", + "us-east-1", + AuthMode::OAuth, + ); + let body = scrape(); + // The label-set rendering uses lexical ordering of label + // names — auth_mode, model, region — so we assert on the + // values without locking in the ordering of the columns. + assert!( + body.contains("auth_mode=\"oauth\""), + "scrape missing auth_mode label: {body}" + ); + assert!( + body.contains("model=\"anthropic.claude-3-haiku-20240307-v1:0\""), + "scrape missing model label: {body}" + ); + assert!( + body.contains("region=\"us-east-1\""), + "scrape missing region label: {body}" + ); + } + + #[test] + fn latency_histogram_records_observation() { + observe_bedrock_invoke_latency("anthropic.claude-3-haiku-20240307-v1:0", "us-east-1", 0.42); + let body = scrape(); + assert!( + body.contains("bedrock_invoke_latency_seconds_bucket"), + "histogram bucket lines missing: {body}" + ); + assert!( + body.contains("bedrock_invoke_latency_seconds_sum"), + "histogram sum line missing: {body}" + ); + assert!( + body.contains("bedrock_invoke_latency_seconds_count"), + "histogram count line missing: {body}" + ); + } + + #[test] + fn eventstream_counter_records_event_type_label() { + record_bedrock_eventstream_message( + "anthropic.claude-3-haiku-20240307-v1:0", + "us-east-1", + "chunk", + ); + record_bedrock_eventstream_message( + "anthropic.claude-3-haiku-20240307-v1:0", + "us-east-1", + "metadata", + ); + let body = scrape(); + assert!( + body.contains("event_type=\"chunk\""), + "scrape missing event_type=chunk: {body}" + ); + assert!( + body.contains("event_type=\"metadata\""), + "scrape missing event_type=metadata: {body}" + ); + } +} diff --git a/crates/headroom-proxy/src/proxy.rs b/crates/headroom-proxy/src/proxy.rs index dd794669a..dddf46964 100644 --- a/crates/headroom-proxy/src/proxy.rs +++ b/crates/headroom-proxy/src/proxy.rs @@ -88,6 +88,14 @@ pub fn build_app(state: AppState) -> Router { let mut router = Router::new() .route("/healthz", get(healthz)) .route("/healthz/upstream", get(healthz_upstream)) + // PR-D3: Prometheus scrape endpoint. Renders the global + // registry in text format. The handler is stateless — no + // `AppState` needed — and idempotent across concurrent + // scrapes (`prometheus`'s registry uses internal locking). + // Mounted unconditionally because it has no dependencies on + // any feature flag; an operator who doesn't want it scraped + // simply firewalls the path. + .route("/metrics", get(crate::observability::handle_metrics)) // PR-C2: explicit POST route for /v1/chat/completions. The // handler buffers the body and re-injects it into // `forward_http`, which runs the OpenAI live-zone gate @@ -116,7 +124,16 @@ pub fn build_app(state: AppState) -> Router { // identical for `anthropic.claude-*` model IDs (Bedrock just // accepts both legacy `invoke` and modern `converse` paths). if state.config.enable_bedrock_native { - router = router + // PR-D3: Bedrock-scoped auth-mode middleware. Build a + // sub-router with ONLY the Bedrock routes, attach the + // auth-mode layer (so it fires before the handler runs and + // is scoped to these routes alone — `/v1/messages`, + // `/healthz`, etc. do NOT run through this middleware), and + // merge it into the parent router. The merge composes + // routes without changing their layer stacks; the parent's + // `with_state` (applied at the end) hands `AppState` to the + // Bedrock handlers identically. + let bedrock_router: Router = Router::new() .route( "/model/:model_id/invoke", post(crate::bedrock::invoke::handle_invoke), @@ -133,7 +150,11 @@ pub fn build_app(state: AppState) -> Router { .route( "/model/:model_id/invoke-with-response-stream", post(crate::bedrock::invoke_streaming::handle_invoke_streaming), - ); + ) + .route_layer(axum::middleware::from_fn( + crate::bedrock::classify_and_attach_auth_mode, + )); + router = router.merge(bedrock_router); if !state.config.bedrock_validate_eventstream_crc { tracing::warn!( event = "bedrock_eventstream_crc_validation_disabled", diff --git a/crates/headroom-proxy/tests/integration_bedrock_authmode.rs b/crates/headroom-proxy/tests/integration_bedrock_authmode.rs new file mode 100644 index 000000000..81fe30085 --- /dev/null +++ b/crates/headroom-proxy/tests/integration_bedrock_authmode.rs @@ -0,0 +1,241 @@ +//! Integration tests for the Bedrock auth-mode middleware +//! (Phase D PR-D3). +//! +//! Coverage: +//! +//! 1. `bedrock_classified_as_oauth` — POST a Bedrock invoke request +//! with no Authorization header (the most common SDK pattern when +//! AWS credentials live downstream of the proxy). Assert the +//! middleware coerces the result to `AuthMode::OAuth` per the +//! Bedrock policy matrix and that the value lands in +//! `request.extensions()` where downstream Phase F handlers can +//! pick it up. +//! 2. `oauth_policy_passthrough_prefer` — fire a request with an +//! Anthropic body containing NO `cache_control` markers; assert +//! the upstream-bound body is byte-equal to the inbound body. +//! The OAuth policy matrix forbids auto-injecting `cache_control` +//! or `prompt_cache_key`; D3 wires the marker, F2 enforces the +//! policy. Until F2 lands, the proof is the byte-equality (no +//! mutation observed at the upstream boundary). + +mod common; + +use aws_credential_types::Credentials; +use axum::body::Body; +use axum::extract::{Extension, State}; +use axum::http::StatusCode; +use axum::routing::post; +use axum::Router; +use bytes::Bytes; +use common::start_proxy_with_state; +use headroom_core::auth_mode::AuthMode; +use headroom_proxy::AppState; +use serde_json::{json, Value}; +use sha2::{Digest, Sha256}; +use std::net::SocketAddr; +use std::sync::{Arc, Mutex}; +use tokio::sync::oneshot; +use url::Url; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +const TEST_MODEL: &str = "anthropic.claude-3-haiku-20240307-v1:0"; + +fn test_credentials() -> Credentials { + Credentials::new( + "AKIAEXAMPLEAKIDFORTEST", + "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + None, + None, + "test", + ) +} + +#[derive(Default, Clone, Debug)] +struct CapturedRequest { + body: Option>, +} + +type Capture = Arc>; + +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(); + Mock::given(method("POST")) + .and(path(format!("/model/{TEST_MODEL}/invoke"))) + .respond_with(move |req: &wiremock::Request| { + let mut c = captured_clone.lock().unwrap(); + c.body = Some(req.body.clone()); + ResponseTemplate::new(200).set_body_string(response_body.clone()) + }) + .mount(upstream) + .await; + captured +} + +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 +} + +/// Test 1: With no Authorization header, the bedrock auth-mode +/// middleware classifies as OAuth (Bedrock policy matrix). We boot +/// a separate axum app that mounts the same middleware in front of +/// a probe handler; the probe reads the AuthMode out of +/// `request.extensions()` and echoes it back. This is the canonical +/// "extension was set" assertion the spec asks for. +#[tokio::test] +async fn bedrock_classified_as_oauth() { + use headroom_proxy::bedrock::classify_and_attach_auth_mode; + + async fn probe(Extension(auth_mode): Extension) -> String { + auth_mode.as_str().to_string() + } + let app = Router::new() + .route("/model/:model_id/invoke", post(probe)) + .route_layer(axum::middleware::from_fn(classify_and_attach_auth_mode)); + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (tx, rx) = oneshot::channel::<()>(); + let task = tokio::spawn(async move { + let _ = axum::serve( + listener, + app.into_make_service_with_connect_info::(), + ) + .with_graceful_shutdown(async move { + let _ = rx.await; + }) + .await; + }); + + // Bedrock SDK style: no Authorization header in the inbound + // request to our proxy (the SDK signs at the egress side, or + // the customer is using IAM-instance-credential downstream of + // our hop). NO x-api-key. NO x-goog-api-key. F1 returns Payg by + // default; the bedrock middleware must coerce to OAuth. + let resp = reqwest::Client::new() + .post(format!( + "http://{addr}/model/{TEST_MODEL}/invoke", + addr = addr, + TEST_MODEL = TEST_MODEL, + )) + .header("content-type", "application/json") + .body(r#"{"anthropic_version":"bedrock-2023-05-31","max_tokens":8,"messages":[]}"#) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + let body_text = resp.text().await.unwrap(); + assert_eq!( + body_text, "oauth", + "bedrock route must classify as OAuth; saw {body_text}" + ); + let _ = tx.send(()); + let _ = task.await; +} + +/// Test 2: confirm the upstream-bound body is byte-equal to the +/// inbound body. The OAuth policy forbids auto-injecting +/// `cache_control`; D3's contribution is to MARK the request as +/// OAuth so PR-F2 can gate the cache-control walker. For now the +/// invariant is "no mutation visible at the upstream boundary" +/// when compression mode is `off`. +#[tokio::test] +async fn oauth_policy_passthrough_prefer() { + 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(); + let received = got.body.expect("upstream got body"); + // Byte-equality (sha256 hashes match). + let inbound_hash = sha256_hex(&body); + let received_hash = sha256_hex(&received); + assert_eq!( + inbound_hash, received_hash, + "upstream body must be byte-equal to inbound body under OAuth policy: \ + inbound={inbound_hash}, received={received_hash}" + ); + // Belt-and-braces: parse the upstream body and assert NO + // cache_control marker was added to any message. + let parsed: Value = serde_json::from_slice(&received).unwrap(); + let messages = parsed["messages"].as_array().expect("messages array"); + for (i, msg) in messages.iter().enumerate() { + // `cache_control` may live on either the message itself or + // on individual content blocks. Assert neither path got + // synthesised by us. + assert!( + msg.get("cache_control").is_none(), + "messages[{i}] gained a cache_control marker; OAuth policy forbids auto-injection" + ); + if let Some(content) = msg.get("content").and_then(|v| v.as_array()) { + for (j, block) in content.iter().enumerate() { + assert!( + block.get("cache_control").is_none(), + "messages[{i}].content[{j}] gained a cache_control marker" + ); + } + } + } + // And NO prompt_cache_key at the top level. + assert!( + parsed.get("prompt_cache_key").is_none(), + "top-level prompt_cache_key must NOT be auto-injected under OAuth" + ); + proxy.shutdown().await; +} + +/// Helper: SHA-256 hex of bytes. Mirrors `integration_bedrock_invoke.rs`. +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 + }) +} + +/// Pin the unused-import lint silencers — these symbols are +/// referenced by the assertions but the linter is paranoid about +/// `axum::body::Body` and `AppState` only being used in a single +/// type-position. +#[allow(dead_code)] +fn _pin(_: Body, _: State, _: Bytes, _: StatusCode) {} diff --git a/crates/headroom-proxy/tests/integration_bedrock_metrics.rs b/crates/headroom-proxy/tests/integration_bedrock_metrics.rs new file mode 100644 index 000000000..0bdb79180 --- /dev/null +++ b/crates/headroom-proxy/tests/integration_bedrock_metrics.rs @@ -0,0 +1,435 @@ +//! Integration tests for the Phase D PR-D3 Prometheus instrumentation. +//! +//! Coverage: +//! +//! 1. `metrics_increment_per_invoke` — fire 3 invoke calls; assert +//! `bedrock_invoke_count_total` registers 3 increments tagged +//! with the right `model` + `region` + `auth_mode=oauth`. +//! 2. `metrics_observe_latency` — fire one invoke; assert +//! `bedrock_invoke_latency_seconds` observed exactly one sample. +//! 3. `eventstream_metrics_per_message_type` — drive D2's streaming +//! path with a captured Bedrock binary stream that yields N +//! `chunk` messages and assert the counter registers +//! `event_type=chunk` with N. (The Anthropic-on-Bedrock vocabulary +//! in D2's translator only accepts `:event-type=chunk`; metadata +//! frames are not produced by Bedrock for the Anthropic shape. +//! We assert the chunk path; a future PR-H2 may add metadata +//! frame support and extend this test.) +//! 4. `metrics_endpoint_serves_scrape` — GET `/metrics` and assert +//! the three Bedrock metric families appear in the text-format +//! output. +//! +//! All tests use wiremock as the upstream — no live AWS dependency. + +mod common; + +use aws_credential_types::Credentials; +use bytes::{Bytes, BytesMut}; +use common::start_proxy_with_state; +use headroom_proxy::bedrock::MessageBuilder; +use serde_json::json; +use url::Url; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +// Each test in this file owns a UNIQUE (model, region) tuple so the +// global Prometheus registry — shared across all parallel tests in +// the same binary — gives each test isolated label rows. Without +// this isolation, parallel-running tests would cross-contaminate the +// counters they read back. Bumping a counter never tears down its +// row, so absolute counts are not assertable after-the-fact; +// per-tuple isolation gives each test a fresh row to assert deltas +// against. +const TEST_MODEL_INVOKE_COUNT: &str = "anthropic.claude-3-haiku-test-invoke-count-v1:0"; +const TEST_MODEL_LATENCY: &str = "anthropic.claude-3-haiku-test-latency-v1:0"; +const TEST_MODEL_EVENTSTREAM: &str = "anthropic.claude-3-haiku-test-eventstream-v1:0"; +const TEST_MODEL_SCRAPE: &str = "anthropic.claude-3-haiku-test-scrape-v1:0"; +const TEST_REGION_INVOKE_COUNT: &str = "us-test-invoke-count-1"; +const TEST_REGION_LATENCY: &str = "us-test-latency-1"; +const TEST_REGION_EVENTSTREAM: &str = "us-test-eventstream-1"; +const TEST_REGION_SCRAPE: &str = "us-test-scrape-1"; + +fn test_credentials() -> Credentials { + Credentials::new( + "AKIAEXAMPLEAKIDFORTEST", + "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + None, + None, + "test", + ) +} + +async fn bedrock_proxy_with_region( + upstream: &MockServer, + region: &str, + customize: impl FnOnce(&mut headroom_proxy::Config), +) -> common::ProxyHandle { + let endpoint: Url = upstream.uri().parse().unwrap(); + let region = region.to_string(); + start_proxy_with_state( + &upstream.uri(), + |c| { + c.bedrock_endpoint = Some(endpoint); + c.bedrock_region = region; + customize(c); + }, + |s| s.with_bedrock_credentials(test_credentials()), + ) + .await +} + +async fn mount_simple_invoke_for(upstream: &MockServer, model: &str) { + Mock::given(method("POST")) + .and(path(format!("/model/{model}/invoke"))) + .respond_with(ResponseTemplate::new(200).set_body_string(r#"{"id":"msg_x","content":[]}"#)) + .mount(upstream) + .await; +} + +/// Fetch the proxy's `/metrics` text-format scrape. +async fn scrape_metrics(proxy_url: &str) -> String { + let resp = reqwest::Client::new() + .get(format!("{proxy_url}/metrics")) + .send() + .await + .expect("metrics scrape"); + assert_eq!(resp.status(), 200, "metrics endpoint must return 200"); + let ct = resp + .headers() + .get("content-type") + .and_then(|v| v.to_str().ok()) + .unwrap_or("") + .to_string(); + assert!( + ct.starts_with("text/plain"), + "metrics content-type must be text/plain (Prometheus text format); got {ct}" + ); + resp.text().await.unwrap() +} + +/// Count the number of Prometheus text lines that contain the +/// metric name + every label key/value pair in `label_pairs`. The +/// label-set rendering uses lexical ordering of label names so we +/// MUST NOT compare exact substrings — instead, every pair must +/// appear in the same line, in any order. +fn count_lines_with_labels( + scrape: &str, + metric: &str, + label_pairs: &[(&str, &str)], +) -> Option { + for line in scrape.lines() { + if !line.starts_with(metric) { + continue; + } + if !label_pairs + .iter() + .all(|(k, v)| line.contains(&format!("{k}=\"{v}\""))) + { + continue; + } + // Counter / gauge: " " tail. We split on the last + // whitespace, parse as u64. + if let Some(value_str) = line.rsplit_once(' ').map(|(_, v)| v.trim()) { + if let Ok(value) = value_str.parse::() { + return Some(value); + } + if let Ok(f) = value_str.parse::() { + return Some(f as u64); + } + } + } + None +} + +/// Test 1: `bedrock_invoke_count_total` increments per request +/// with the right model / region / auth_mode labels. +#[tokio::test] +async fn metrics_increment_per_invoke() { + let upstream = MockServer::start().await; + mount_simple_invoke_for(&upstream, TEST_MODEL_INVOKE_COUNT).await; + let proxy = bedrock_proxy_with_region(&upstream, TEST_REGION_INVOKE_COUNT, |c| { + c.compression_mode = headroom_proxy::config::CompressionMode::Off; + }) + .await; + + // Per-tuple-isolated counter — start at 0 (no other test + // touches this label set), so absolute count == invocations. + let payload = json!({ + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 8, + "messages": [{"role":"user","content":"hi"}] + }); + let body = serde_json::to_vec(&payload).unwrap(); + for _ in 0..3 { + let resp = reqwest::Client::new() + .post(format!( + "{}/model/{TEST_MODEL_INVOKE_COUNT}/invoke", + proxy.url() + )) + .header("content-type", "application/json") + .body(body.clone()) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + } + + let after = scrape_metrics(&proxy.url()).await; + let after_count = count_lines_with_labels( + &after, + "bedrock_invoke_count_total", + &[ + ("model", TEST_MODEL_INVOKE_COUNT), + ("region", TEST_REGION_INVOKE_COUNT), + ("auth_mode", "oauth"), + ], + ) + .expect("counter row must appear after first request"); + assert_eq!( + after_count, 3, + "expected exactly 3 increments on isolated labels; got {after_count}" + ); + + proxy.shutdown().await; +} + +/// Test 2: `bedrock_invoke_latency_seconds` records exactly one +/// sample for one request. +#[tokio::test] +async fn metrics_observe_latency() { + let upstream = MockServer::start().await; + mount_simple_invoke_for(&upstream, TEST_MODEL_LATENCY).await; + let proxy = bedrock_proxy_with_region(&upstream, TEST_REGION_LATENCY, |c| { + c.compression_mode = headroom_proxy::config::CompressionMode::Off; + }) + .await; + + let payload = json!({ + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 8, + "messages": [{"role":"user","content":"hi"}] + }); + let body = serde_json::to_vec(&payload).unwrap(); + let resp = reqwest::Client::new() + .post(format!("{}/model/{TEST_MODEL_LATENCY}/invoke", proxy.url())) + .header("content-type", "application/json") + .body(body) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + + let after = scrape_metrics(&proxy.url()).await; + let after_count = count_lines_with_labels( + &after, + "bedrock_invoke_latency_seconds_count", + &[ + ("model", TEST_MODEL_LATENCY), + ("region", TEST_REGION_LATENCY), + ], + ) + .expect("histogram count row must appear after first request"); + assert_eq!( + after_count, 1, + "expected exactly 1 latency observation on isolated labels; got {after_count}" + ); + + // Sum line for the same labels must appear and be > 0. + let sum_line = after + .lines() + .find(|l| { + l.starts_with("bedrock_invoke_latency_seconds_sum") + && l.contains(&format!("model=\"{TEST_MODEL_LATENCY}\"")) + }) + .expect("histogram sum line must appear for our labels"); + let sum_value: f64 = sum_line + .rsplit_once(' ') + .map(|(_, v)| v.trim()) + .and_then(|s| s.parse().ok()) + .unwrap_or(0.0); + assert!( + sum_value > 0.0, + "histogram sum must reflect a real observation > 0s; saw {sum_value}" + ); + + proxy.shutdown().await; +} + +/// Synthesise N chunk EventStream messages. +fn synthesize_chunks(n: usize) -> Bytes { + let mut buf = BytesMut::new(); + for i in 0..n { + let payload = serde_json::to_string(&json!({ + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": format!("t{i}")} + })) + .unwrap(); + let bytes = MessageBuilder::new() + .header_string(":event-type", "chunk") + .header_string(":content-type", "application/json") + .header_string(":message-type", "event") + .payload(Bytes::from(payload)) + .build(); + buf.extend_from_slice(&bytes); + } + buf.freeze() +} + +/// Test 3: per-EventStream-message metrics increment with the +/// correct `event_type` label. +#[tokio::test] +async fn eventstream_metrics_per_message_type() { + let upstream = MockServer::start().await; + let chunks = synthesize_chunks(5); + Mock::given(method("POST")) + .and(path(format!( + "/model/{TEST_MODEL_EVENTSTREAM}/invoke-with-response-stream" + ))) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "application/vnd.amazon.eventstream") + .set_body_bytes(chunks.to_vec()), + ) + .mount(&upstream) + .await; + + let proxy = bedrock_proxy_with_region(&upstream, TEST_REGION_EVENTSTREAM, |c| { + c.compression_mode = headroom_proxy::config::CompressionMode::Off; + }) + .await; + + // Default Accept → SSE translation, which is the path that + // parses messages and increments the counter (passthrough mode + // forwards bytes verbatim and therefore can't categorize event + // types — the spec defers that to a future H2 PR). + let payload = json!({ + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 8, + "messages": [{"role":"user","content":"hi"}] + }); + let body = serde_json::to_vec(&payload).unwrap(); + let resp = reqwest::Client::new() + .post(format!( + "{}/model/{TEST_MODEL_EVENTSTREAM}/invoke-with-response-stream", + proxy.url() + )) + .header("content-type", "application/json") + .body(body) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + // Drain the response body so the translator runs to completion. + let _ = resp.bytes().await.unwrap(); + + let after = scrape_metrics(&proxy.url()).await; + let after_count = count_lines_with_labels( + &after, + "bedrock_eventstream_message_count_total", + &[ + ("model", TEST_MODEL_EVENTSTREAM), + ("region", TEST_REGION_EVENTSTREAM), + ("event_type", "chunk"), + ], + ) + .expect("eventstream chunk counter row must appear after first stream"); + assert_eq!( + after_count, 5, + "expected 5 chunk increments on isolated labels; got {after_count}" + ); + + proxy.shutdown().await; +} + +/// Test 4: `/metrics` endpoint serves a valid Prometheus text-format +/// scrape that includes the three Bedrock metric families. Every +/// metric family must be touched at least once for the +/// `prometheus` crate to render its HELP/TYPE lines (`gather()` +/// skips empty vectors), so this test explicitly drives both the +/// invoke and the streaming routes — each populates a different +/// family, and the latency histogram comes for free with the +/// invoke route. +#[tokio::test] +async fn metrics_endpoint_serves_scrape() { + let upstream = MockServer::start().await; + mount_simple_invoke_for(&upstream, TEST_MODEL_SCRAPE).await; + // Mount the streaming endpoint too, so the third metric family + // (eventstream message counter) gets at least one increment + // and its HELP/TYPE lines render in the scrape. + let chunks = synthesize_chunks(1); + Mock::given(method("POST")) + .and(path(format!( + "/model/{TEST_MODEL_SCRAPE}/invoke-with-response-stream" + ))) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "application/vnd.amazon.eventstream") + .set_body_bytes(chunks.to_vec()), + ) + .mount(&upstream) + .await; + let proxy = bedrock_proxy_with_region(&upstream, TEST_REGION_SCRAPE, |c| { + c.compression_mode = headroom_proxy::config::CompressionMode::Off; + }) + .await; + + // Fire one invoke (populates invoke_count + invoke_latency) + // and one streaming invoke (populates eventstream_count). + let payload = json!({ + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 8, + "messages": [{"role":"user","content":"hi"}] + }); + let body = serde_json::to_vec(&payload).unwrap(); + let resp = reqwest::Client::new() + .post(format!("{}/model/{TEST_MODEL_SCRAPE}/invoke", proxy.url())) + .header("content-type", "application/json") + .body(body.clone()) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + let stream_resp = reqwest::Client::new() + .post(format!( + "{}/model/{TEST_MODEL_SCRAPE}/invoke-with-response-stream", + proxy.url() + )) + .header("content-type", "application/json") + .body(body) + .send() + .await + .unwrap(); + assert_eq!(stream_resp.status(), 200); + let _ = stream_resp.bytes().await.unwrap(); + + let scrape = scrape_metrics(&proxy.url()).await; + // HELP + TYPE lines are advertised even before any increment; + // after the increment the labelled rows also appear. + assert!( + scrape.contains("# HELP bedrock_invoke_count_total"), + "scrape missing bedrock_invoke_count_total HELP: {scrape}" + ); + assert!( + scrape.contains("# TYPE bedrock_invoke_count_total counter"), + "scrape missing bedrock_invoke_count_total TYPE: {scrape}" + ); + assert!( + scrape.contains("# HELP bedrock_invoke_latency_seconds"), + "scrape missing bedrock_invoke_latency_seconds HELP" + ); + assert!( + scrape.contains("# TYPE bedrock_invoke_latency_seconds histogram"), + "scrape missing bedrock_invoke_latency_seconds TYPE" + ); + assert!( + scrape.contains("# HELP bedrock_eventstream_message_count_total"), + "scrape missing bedrock_eventstream_message_count_total HELP" + ); + assert!( + scrape.contains("# TYPE bedrock_eventstream_message_count_total counter"), + "scrape missing bedrock_eventstream_message_count_total TYPE" + ); + + proxy.shutdown().await; +} diff --git a/docs/bedrock.md b/docs/bedrock.md new file mode 100644 index 000000000..31d502c44 --- /dev/null +++ b/docs/bedrock.md @@ -0,0 +1,182 @@ +# AWS Bedrock — Operator Guide + +Headroom's Rust proxy ships a native AWS Bedrock InvokeModel surface. After Phase D (PRs D1–D3), Anthropic-on-Bedrock requests are signed, compressed, and observed by the proxy directly — no LiteLLM Python shim on the request path. + +This document covers how to deploy the Bedrock-native surface, how compression policy is applied, and how to read the Prometheus metrics the proxy exports. + +## What's in scope + +| Capability | Status | +|---|---| +| `POST /model/{model}/invoke` | PR-D1 — native Rust handler | +| `POST /model/{model}/converse` | PR-D1 — same handler (Bedrock accepts both paths for the Anthropic envelope) | +| `POST /model/{model}/invoke-with-response-stream` | PR-D2 — binary EventStream parsed and translated to SSE | +| AWS SigV4 signing (post-compression) | PR-D1 | +| `AuthMode::OAuth` classification | PR-D3 — Bedrock IAM is OAuth-equivalent under the policy matrix | +| Per-model + per-region Prometheus metrics | PR-D3 — exposed at `GET /metrics` | +| OAuth compression policy gates (no auto cache_control, lossless-only) | Phase F PR-F2/F3 (gates the marker D3 wires) | + +## AWS credential configuration + +The proxy uses the [aws-config default credential chain](https://docs.aws.amazon.com/sdkref/latest/guide/standardized-credentials.html), resolved once at startup. + +The chain searches in this order, stopping at the first source that yields valid credentials: + +1. **Environment variables** — `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, optional `AWS_SESSION_TOKEN`. Useful for ECS task roles that inject creds via env or for `aws sts assume-role` shells. +2. **Shared credentials file** — `~/.aws/credentials`. Profile selected by `--aws-profile` (or `AWS_PROFILE`). Falls back to `[default]`. +3. **IAM instance profile / IMDS** — when running on EC2. +4. **ECS task role / EKS pod identity** — when running on the AWS-managed compute platforms. +5. **AWS SSO** — `~/.aws/sso/cache/...` when `aws sso login` has been run. + +If the chain does NOT resolve any credentials at startup, the proxy logs `event=bedrock_credentials_unavailable` at WARN and continues to start. Bedrock invoke routes will then return `500` with `event=bedrock_credentials_missing` per request — there is **no silent fallback to unsigned requests**, by design. + +### Required IAM permissions + +The proxy needs: + +- `bedrock:InvokeModel` for non-streaming +- `bedrock:InvokeModelWithResponseStream` for streaming + +Scope these to the specific model ARNs you intend to use. Example IAM policy snippet: + +```json +{ + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": [ + "bedrock:InvokeModel", + "bedrock:InvokeModelWithResponseStream" + ], + "Resource": [ + "arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-3-haiku-20240307-v1:0", + "arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-3-5-sonnet-20241022-v2:0" + ] + } + ] +} +``` + +## Region configuration + +```sh +headroom-proxy \ + --upstream http://unused-when-bedrock-only \ + --bedrock-region us-east-1 +``` + +| Flag | Env var | Default | Notes | +|---|---|---|---| +| `--bedrock-region` | `HEADROOM_PROXY_BEDROCK_REGION` (or `AWS_REGION`) | `us-east-1` | Drives both the SigV4 region and the derived endpoint hostname. | +| `--bedrock-endpoint` | `HEADROOM_PROXY_BEDROCK_ENDPOINT` | derived from region | Override for FIPS endpoints (`bedrock-runtime-fips.{region}.amazonaws.com`), VPC endpoints, or local mock servers. | +| `--aws-profile` | `HEADROOM_PROXY_AWS_PROFILE` | unset | Selects the named profile from the shared credentials file. | +| `--enable-bedrock-native` | `HEADROOM_PROXY_ENABLE_BEDROCK_NATIVE` | `true` | Set to `false` to mount no Bedrock routes at all (Bedrock requests will then fall through to the catch-all and fail without SigV4). | + +## Supported model IDs + +The proxy classifies model IDs by **literal vendor prefix** — no regexes. Any model ID starting with `anthropic.` is treated as Anthropic-shape: the live-zone compression dispatcher runs over the body, the envelope is re-emitted with `anthropic_version` preserved as the first key, and the request is signed with SigV4. + +Examples that hit the Anthropic compression path: + +- `anthropic.claude-3-haiku-20240307-v1:0` +- `anthropic.claude-3-5-sonnet-20241022-v2:0` +- `anthropic.claude-3-opus-20240229-v1:0` + +Other Bedrock vendors (`amazon.titan-...`, `meta.llama3-...`, `cohere.command-...`, `ai21.j2-...`, `stability.stable-diffusion-...`) are signed and forwarded **without compression** — the proxy does not yet understand their body shapes and would risk corrupting them. These model IDs log `event=bedrock_compression_skipped, reason=non_anthropic_vendor` per request. Full Anthropic envelopes only. + +The contract: **any new model ID that AWS adds under the `anthropic.` prefix automatically picks up the full compression + signing pipeline.** No code change in the proxy is needed for new versions of Claude on Bedrock. + +## Compression behaviour + +Bedrock requests are subject to the **same** live-zone compression rules as direct Anthropic (`/v1/messages`): + +- Only the live-zone messages (latest user turn, latest tool/output blocks) are eligible for compression. +- The cache hot zone (older messages, system prompt, tools list) is byte-faithful passthrough. +- The dispatcher only mutates body bytes when at least one block compressed. The byte-equality invariant for unchanged blocks is enforced at `debug_assert!` granularity. + +### OAuth policy (PR-D3 → PR-F2/F3) + +The Bedrock auth-mode middleware classifies every Bedrock request as `AuthMode::OAuth`. Even when the inbound request has no Authorization header (the common case where the AWS SDK signs after our proxy), the middleware **coerces** to OAuth and emits `event=bedrock_auth_mode_unexpected` at WARN if F1's classifier disagreed — so the divergence is loud, not silent. + +Under the OAuth policy matrix (see `docs/auth-modes.md`): + +- **No auto-`cache_control` injection.** OAuth subscriptions pin the cache scope to `(account, model, session)`; auto-injecting markers can void cache hits. +- **No auto-`prompt_cache_key`.** Same reasoning. +- **Lossless-only compressors.** Lossy compressors (text rewriting, summarisation) are gated off for OAuth. + +PR-D3 lands the classification + the resulting `AuthMode` in `request.extensions()`. PR-F2 and PR-F3 wire the actual policy gates that read it. Until those PRs land, the Bedrock route uses the existing dispatcher (which is a no-op in `compression_mode=off`); the OAuth contract above is the documented forward direction. + +### Cache safety + +The bytes signed by SigV4 are exactly the bytes Bedrock receives — the signer hashes the post-compression body. There is no "sign before compress" shortcut that would produce a signature mismatched to the wire payload. Compression mutates the body once, then the signer runs once, then the bytes are forwarded once. + +## Prometheus metrics + +The proxy exposes a `GET /metrics` endpoint that serves the standard Prometheus text-format scrape. Three Bedrock-specific metric families are exported: + +| Metric | Type | Labels | Source | +|---|---|---|---| +| `bedrock_invoke_count_total` | Counter | `model`, `region`, `auth_mode` | One increment per `/model/.../invoke` (and `/converse` and `/invoke-with-response-stream`) request. | +| `bedrock_invoke_latency_seconds` | Histogram | `model`, `region` | Observed at request completion (success or failure). | +| `bedrock_eventstream_message_count_total` | Counter | `model`, `region`, `event_type` | One increment per parsed binary EventStream message in the streaming path. `event_type` is the `:event-type` header (`chunk`, `metadata`, `internalServerException`, etc.). | + +All labels are bounded by infrastructure config (`region` from `--bedrock-region`, `auth_mode` from the 3-variant enum) or by the path parameter (`model`, supplied by the axum extractor — never by user-controlled body bytes). Cardinality is bounded by deployment fan-out, not by traffic volume. + +### Sample PromQL queries + +**p99 latency by model:** +```promql +histogram_quantile( + 0.99, + sum by (model, le) (rate(bedrock_invoke_latency_seconds_bucket[5m])) +) +``` + +**Request rate by region (RPS):** +```promql +sum by (region) (rate(bedrock_invoke_count_total[1m])) +``` + +**EventStream message rate by event type (debugging the streaming path):** +```promql +sum by (event_type) (rate(bedrock_eventstream_message_count_total[1m])) +``` + +**Error breakdown by HTTP status (cross-references the structured logs `event=bedrock_upstream_error`):** +```promql +sum by (model) (rate(bedrock_invoke_count_total{auth_mode="oauth"}[5m])) + / sum by (model) (rate(bedrock_invoke_latency_seconds_count[5m])) +``` + +(The denominator is the total observed latency samples — useful for sanity-checking that every counted invoke also got a histogram observation. They should be equal.) + +### Structured-log correlation + +Every metric increment in the Bedrock path is paired with a `tracing::debug!` log line carrying: + +- `event = "metric_recorded"` +- `metric = "bedrock_invoke_count_total" | "bedrock_invoke_latency_seconds" | "bedrock_eventstream_message_count_total"` +- the same labels as the metric + +Enable with `RUST_LOG=headroom_proxy::observability=debug` for incident correlation. In normal operation keep this at the default `info` level — debug volume per request is bounded by the same cardinality the metric uses. + +## Live cloud validation + +The PR-D1, D2, D3 implementations are exercised end-to-end against a wiremock upstream (`crates/headroom-proxy/tests/integration_bedrock_*.rs`). The wiremock-based tests are the canonical correctness gate. + +A real Bedrock smoke test (`aws bedrock-runtime invoke-model ...` through the proxy) requires `bedrock:InvokeModel` permissions in the developer's AWS account. Set the proxy upstream to the proxy URL (`http://localhost:8787`) via the AWS SDK's `AWS_ENDPOINT_URL_BEDROCK_RUNTIME` env var: + +```sh +AWS_ENDPOINT_URL_BEDROCK_RUNTIME=http://localhost:8787 \ + aws bedrock-runtime invoke-model \ + --model-id anthropic.claude-3-haiku-20240307-v1:0 \ + --body '{"anthropic_version":"bedrock-2023-05-31","max_tokens":32,"messages":[{"role":"user","content":"hi"}]}' \ + /tmp/out.json +``` + +If the SDK signs the request before sending to the proxy, the proxy will see a SigV4 `Authorization` header and classify as OAuth via the standard rule. If the SDK is configured to sign downstream of the proxy (some IAM-instance-profile setups), the proxy still classifies as OAuth via the middleware's coerce-and-log fallback. + +## Rollback + +Set `--enable-bedrock-native=false` to unmount all Bedrock routes; the catch-all proxy then forwards Bedrock requests unchanged to `--upstream`. This is an emergency rollback only — without SigV4 re-signing, the catch-all path will fail closed unless the upstream is itself a Bedrock-aware proxy (e.g., the Python LiteLLM converter on a different port).