From b08ec15b0d392b8b8cf93dbadaee4b7e6b465f1c Mon Sep 17 00:00:00 2001 From: Yasser Sheikh <158480094+ysheikh2@users.noreply.github.com> Date: Sat, 13 Jun 2026 00:18:43 +0200 Subject: [PATCH] fix(proxy): add native Bedrock converse-stream route (#917) ## Description Adds native Bedrock `POST /model/{model_id}/converse-stream` routing in `headroom-proxy` by reusing the existing streaming handler and preserving route-specific upstream action forwarding. This addresses a gap where native Bedrock streaming support existed for `invoke-with-response-stream` but not `converse-stream`, even though both share the same EventStream transport and SSE translation path in this proxy. Fixes #919 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring ## Changes Made - Add route mount in `crates/headroom-proxy/src/proxy.rs`: - `POST /model/:model_id/converse-stream` -> `bedrock::invoke_streaming::handle_invoke_streaming` - Update streaming handler URL construction in `crates/headroom-proxy/src/bedrock/invoke_streaming.rs`: - infer action from inbound path (`invoke-with-response-stream` or `converse-stream`) - build upstream URL with the resolved action - return structured `400` for unsupported streaming action paths - Add unit tests in `crates/headroom-proxy/src/bedrock/invoke_streaming.rs`: - action extraction coverage for both streaming paths - upstream URL construction coverage for `converse-stream` - Add integration coverage in `crates/headroom-proxy/tests/integration_bedrock_streaming.rs`: - `converse_stream_route_translates_to_sse` - Add changelog entry under `Unreleased` bug fixes in `CHANGELOG.md`. ## Testing - `cargo fmt --all` - `cargo test -p headroom-proxy --test integration_bedrock_streaming -- --nocapture` - `cargo test -p headroom-proxy --test integration_bedrock_metrics -- --nocapture` ## Real behavior proof - **Setup tested on** - macOS (darwin) - Rust workspace local dev build - `headroom-proxy` integration tests using wiremock upstream (no AWS dependency) - **Exact commands run after patch** - `cargo test -p headroom-proxy --test integration_bedrock_streaming -- --nocapture` - `cargo test -p headroom-proxy --test integration_bedrock_metrics -- --nocapture` - **After-fix evidence + observed result** - New integration test `converse_stream_route_translates_to_sse` passes. - Streaming suite result: `10 passed; 0 failed`. - Metrics suite result: `4 passed; 0 failed`. - Logs show requests reaching `/model/.../converse-stream` and flowing through Bedrock streaming path. - **What I did not test** - Live AWS Bedrock calls against real credentials/models. - End-to-end CLI/runtime behavior outside Rust integration test harness. --- CHANGELOG.md | 1 + .../src/bedrock/eventstream_to_sse.rs | 32 ++++++- .../src/bedrock/invoke_streaming.rs | 89 ++++++++++++++++++- crates/headroom-proxy/src/proxy.rs | 11 ++- .../tests/integration_bedrock_streaming.rs | 59 ++++++++++++ 5 files changed, 184 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a3e44c9d3..e67e3d517 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Bug Fixes * **ccr:** make retrieval store TTL configurable with `HEADROOM_CCR_TTL_SECONDS`, expose the effective TTL in `/v1/retrieve/stats`, and distinguish expired retrievals from missing hashes. +* **proxy:** add native Bedrock `/model/{id}/converse-stream` route and forward it through the existing streaming EventStream/SSE pipeline. ## [0.25.0](https://github.com/chopratejas/headroom/compare/v0.24.0...v0.25.0) (2026-06-12) diff --git a/crates/headroom-proxy/src/bedrock/eventstream_to_sse.rs b/crates/headroom-proxy/src/bedrock/eventstream_to_sse.rs index d4e5c8f17..f7e3107e6 100644 --- a/crates/headroom-proxy/src/bedrock/eventstream_to_sse.rs +++ b/crates/headroom-proxy/src/bedrock/eventstream_to_sse.rs @@ -154,12 +154,18 @@ pub fn translate_message( event_type: event_type.to_string(), }) } - (OutputMode::Sse, "chunk") => { + (OutputMode::Sse, "chunk") + | (OutputMode::Sse, "messageStart") + | (OutputMode::Sse, "contentBlockStart") + | (OutputMode::Sse, "contentBlockDelta") + | (OutputMode::Sse, "contentBlockStop") + | (OutputMode::Sse, "messageStop") + | (OutputMode::Sse, "metadata") => { tracing::info!( event = "bedrock_eventstream_translated_to_sse", event_type = event_type, payload_bytes = message.payload.len(), - "translated bedrock eventstream chunk to sse frame" + "translated bedrock eventstream message to sse frame" ); Ok(TranslateOutcome::Emit(payload_to_sse_frame( &message.payload, @@ -380,6 +386,28 @@ mod tests { } } + #[test] + fn translate_converse_event_to_sse_frame() { + let bytes = crate::bedrock::eventstream::MessageBuilder::new() + .header_string(":event-type", "contentBlockDelta") + .header_string(":message-type", "event") + .payload(Bytes::from_static( + br#"{"contentBlockIndex":0,"delta":{"text":"hi"}}"#, + )) + .build(); + let msg = parse(&bytes).unwrap(); + let outcome = translate_message(&msg, OutputMode::Sse).unwrap(); + match outcome { + TranslateOutcome::Emit(b) => { + let s = std::str::from_utf8(&b).unwrap(); + assert!(s.starts_with("data: ")); + assert!(s.ends_with("\n\n")); + assert!(s.contains("contentBlockIndex")); + } + other => panic!("expected Emit; got {other:?}"), + } + } + #[test] fn missing_event_type_is_loud() { // A message lacking :event-type must not silently translate. diff --git a/crates/headroom-proxy/src/bedrock/invoke_streaming.rs b/crates/headroom-proxy/src/bedrock/invoke_streaming.rs index ea52bee58..347195f6e 100644 --- a/crates/headroom-proxy/src/bedrock/invoke_streaming.rs +++ b/crates/headroom-proxy/src/bedrock/invoke_streaming.rs @@ -80,8 +80,9 @@ const ANTHROPIC_VENDOR_PREFIX: &str = "anthropic."; /// AWS Bedrock Runtime DNS template. const BEDROCK_RUNTIME_HOST_TEMPLATE: &str = "bedrock-runtime.{region}.amazonaws.com"; -/// Path action for the streaming route. +/// Path action for the streaming routes. const STREAMING_ACTION: &str = "invoke-with-response-stream"; +const CONVERSE_STREAM_ACTION: &str = "converse-stream"; /// RAII guard that observes the `bedrock_invoke_latency_seconds` /// histogram on drop. Mirrors the [`crate::bedrock::invoke`] guard @@ -166,8 +167,26 @@ pub async fn handle_invoke_streaming( body.clone() }; - // 2. Build upstream URL. - let upstream_url = match build_bedrock_streaming_upstream(&state, &model_id, &uri) { + // 2. Resolve the Bedrock streaming action from the inbound path and + // build the upstream URL. + let action = match extract_streaming_action(uri.path()) { + Some(a) => a, + None => { + tracing::error!( + event = "bedrock_streaming_action_invalid", + request_id = %request_id, + path = %uri.path(), + "bedrock invoke-streaming: unrecognized streaming action path" + ); + return error_response( + StatusCode::BAD_REQUEST, + "bedrock_streaming_action_invalid", + "Unsupported Bedrock streaming action path", + ); + } + }; + + let upstream_url = match build_bedrock_streaming_upstream(&state, &model_id, &uri, action) { Ok(u) => u, Err(msg) => { tracing::error!( @@ -888,6 +907,7 @@ fn build_bedrock_streaming_upstream( state: &AppState, model_id: &str, uri: &Uri, + action: &str, ) -> Result { let base = match state.config.bedrock_endpoint.as_ref() { Some(u) => u.clone(), @@ -901,7 +921,7 @@ fn build_bedrock_streaming_upstream( let path = format!( "/model/{model_id}/{action}", model_id = model_id, - action = STREAMING_ACTION, + action = action, ); let mut joined = base; joined.set_path(&path); @@ -911,6 +931,16 @@ fn build_bedrock_streaming_upstream( Ok(joined) } +fn extract_streaming_action(path: &str) -> Option<&'static str> { + if path.ends_with(&format!("/{STREAMING_ACTION}")) { + Some(STREAMING_ACTION) + } else if path.ends_with(&format!("/{CONVERSE_STREAM_ACTION}")) { + Some(CONVERSE_STREAM_ACTION) + } else { + None + } +} + 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() { @@ -996,6 +1026,7 @@ mod tests { &state, "anthropic.claude-3-haiku-20240307-v1:0", &uri, + STREAMING_ACTION, ) .unwrap(); assert_eq!( @@ -1012,4 +1043,54 @@ mod tests { assert!(s.ends_with("\n\n")); assert!(s.contains("bedrock_eventstream_crc_mismatch")); } + + #[test] + fn build_streaming_upstream_supports_converse_stream_action() { + use crate::config::Config; + let mut config = Config::for_test(Url::parse("http://up:8080").unwrap()); + config.bedrock_region = "eu-west-1".to_string(); + let state = AppState { + config: std::sync::Arc::new(config), + client: reqwest::Client::new(), + bedrock_credentials: None, + drift_state: crate::cache_stabilization::drift_detector::DriftState::new(8), + vertex_token_source: std::sync::Arc::new(crate::vertex::StaticTokenSource::new( + "test".to_string(), + )), + }; + let uri: Uri = "/model/anthropic.claude-3-haiku-20240307-v1:0/converse-stream" + .parse() + .unwrap(); + let url = build_bedrock_streaming_upstream( + &state, + "anthropic.claude-3-haiku-20240307-v1:0", + &uri, + CONVERSE_STREAM_ACTION, + ) + .unwrap(); + assert_eq!( + url.as_str(), + "https://bedrock-runtime.eu-west-1.amazonaws.com/model/anthropic.claude-3-haiku-20240307-v1:0/converse-stream" + ); + } + + #[test] + fn extract_streaming_action_supports_both_bedrock_paths() { + assert_eq!( + extract_streaming_action( + "/model/anthropic.claude-3-haiku-20240307-v1:0/invoke-with-response-stream" + ), + Some(STREAMING_ACTION) + ); + assert_eq!( + extract_streaming_action( + "/model/anthropic.claude-3-haiku-20240307-v1:0/converse-stream" + ), + Some(CONVERSE_STREAM_ACTION) + ); + assert_eq!( + extract_streaming_action("/model/anthropic.claude-3-haiku-20240307-v1:0/invoke"), + None + ); + } } diff --git a/crates/headroom-proxy/src/proxy.rs b/crates/headroom-proxy/src/proxy.rs index d3fe821cc..89820d6cc 100644 --- a/crates/headroom-proxy/src/proxy.rs +++ b/crates/headroom-proxy/src/proxy.rs @@ -209,15 +209,22 @@ pub fn build_app(state: AppState) -> Router { "/model/:model_id/converse", post(crate::bedrock::invoke::handle_invoke), ) - // PR-D2: streaming counterpart. Bedrock's protocol is + // PR-D2/PR-D5: streaming counterparts. Bedrock's protocol is // binary EventStream; the handler parses incrementally, // optionally translates each chunk to an SSE frame, and // tees translated frames into AnthropicStreamState for - // telemetry. See `bedrock::invoke_streaming`. + // telemetry. `invoke-with-response-stream` and + // `converse-stream` share the same wire framing and + // processing pipeline, so both route to the same handler. + // See `bedrock::invoke_streaming`. .route( "/model/:model_id/invoke-with-response-stream", post(crate::bedrock::invoke_streaming::handle_invoke_streaming), ) + .route( + "/model/:model_id/converse-stream", + post(crate::bedrock::invoke_streaming::handle_invoke_streaming), + ) .route_layer(axum::middleware::from_fn( crate::bedrock::classify_and_attach_auth_mode, )) diff --git a/crates/headroom-proxy/tests/integration_bedrock_streaming.rs b/crates/headroom-proxy/tests/integration_bedrock_streaming.rs index 41ade8b40..5bb4ba099 100644 --- a/crates/headroom-proxy/tests/integration_bedrock_streaming.rs +++ b/crates/headroom-proxy/tests/integration_bedrock_streaming.rs @@ -218,6 +218,18 @@ async fn mount_eventstream_upstream(upstream: &MockServer, body: Bytes) { .await; } +async fn mount_eventstream_upstream_for_action(upstream: &MockServer, body: Bytes, action: &str) { + Mock::given(method("POST")) + .and(path(format!("/model/{TEST_MODEL}/{action}"))) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "application/vnd.amazon.eventstream") + .set_body_bytes(body.to_vec()), + ) + .mount(upstream) + .await; +} + #[tokio::test] async fn eventstream_translated_to_sse() { let _ = tracing_subscriber::fmt() @@ -465,6 +477,53 @@ async fn client_can_choose_eventstream_or_sse() { proxy.shutdown().await; } +#[tokio::test] +async fn converse_stream_route_translates_to_sse() { + let upstream = MockServer::start().await; + let bedrock_bytes = synthesize_bedrock_stream(); + mount_eventstream_upstream_for_action(&upstream, bedrock_bytes, "converse-stream").await; + let proxy = bedrock_proxy(&upstream, |c| { + c.compression_mode = headroom_proxy::config::CompressionMode::Off; + }) + .await; + + let body = serde_json::to_vec(&json!({ + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 16, + "messages": [{"role":"user","content":"hi"}] + })) + .unwrap(); + + let resp = reqwest::Client::new() + .post(format!( + "{}/model/{TEST_MODEL}/converse-stream", + proxy.url() + )) + .header("content-type", "application/json") + .header("accept", "text/event-stream") + .body(body) + .send() + .await + .unwrap(); + + assert_eq!(resp.status(), 200); + let ct = resp + .headers() + .get("content-type") + .and_then(|v| v.to_str().ok()) + .unwrap_or(""); + assert!( + ct.starts_with("text/event-stream"), + "converse-stream should emit SSE when client asks for SSE; got {ct}" + ); + + let text = resp.text().await.unwrap(); + assert!(text.contains("event: content_block_delta")); + assert!(text.contains("\"text\":\"OK\"")); + + proxy.shutdown().await; +} + // ─── Test 5: Property test — never panic on adversarial bytes ───── proptest! {