From f08256411d0be70ffb9f4c0905f8a238737cd000 Mon Sep 17 00:00:00 2001 From: Ming Wen Date: Mon, 25 May 2026 08:07:16 +0800 Subject: [PATCH 1/2] feat(bedrock): wire Converse API for all publishers + Anthropic streaming (#302 Phase G Step 3 / D7.2.b / D7.3 / D7.4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces per-publisher /invoke dispatch with the unified Converse API path for everything except Anthropic non-stream (which keeps its existing /invoke path for backward compat with operator deployments + the e2e test fixtures pinned in #320). All other dispatch — Meta / Mistral / Cohere / Amazon Titan / Amazon Nova / AI21 — and ALL chat_stream() requests — including Anthropic — now flow through the SDK's `.converse()` / `.converse_stream()`. This closes D7.2.b (Anthropic streaming), D7.3 (Meta dispatch), and D7.4 (Mistral / Amazon / Cohere / AI21 dispatch) per #302 Phase G's audit-corrected roadmap. ## Why Converse (vs. extending the per-publisher /invoke path) AWS introduced the Converse API exactly to solve the per-publisher body-shaping problem: one JSON envelope (`messages[].content[].text` + `system[]` + `inferenceConfig`), one URL pattern (`/model//converse[-stream]`), one response shape, works for every Bedrock publisher. The legacy /invoke path requires a per-publisher body translator on the gateway side (Anthropic's Messages, Meta's Llama prompt template, Cohere's command shape, etc.) — multiplying maintenance burden as AWS adds publishers. ## Eventstream framing `/converse-stream` emits `application/vnd.amazon.eventstream` binary frames. The aws-sdk-bedrockruntime SDK owns the frame decoder — we don't need our own (compare with the mock-bedrock fixture I wrote in api7/AISIX-Cloud#480 which hand-rolled the encoder for the test mock). The bridge layer just walks the typed `ConverseStreamOutput` event sequence and maps each event to zero, one, or two `ChatChunk`s: - `MessageStart` → role chunk (first emission only) - `ContentBlockDelta` (text) → content chunk - `MessageStop` → finish_reason chunk - `Metadata` → usage chunk - `ContentBlockStart` / `ContentBlockStop` / tool-use deltas → no chunk ## Error classification The new `map_aws_sdk_error_generic` helper preserves the existing 4xx-vs-5xx classification: ServiceError → `UpstreamStatus` carrying the HTTP code; TimeoutError or deadline-elapsed → `Timeout`; everything else → `Transport`. Mirrors the pattern in the legacy `map_sdk_error` + the audit-corrected Vertex/Azure token-mint classifications (#387, #388). ## Test changes Tests deleted (obsolete after Converse wiring): - chat_rejects_non_anthropic_publishers_with_publisher_named - chat_stream_returns_clear_not_implemented_error - chat_stream_anthropic_returns_d7_2_b_specific_error - chat_stream_non_anthropic_publisher_returns_d7_3_specific_error - chat_publisher_not_implemented_error_includes_model_id_and_publisher_name Tests updated: - chat_ignores_req_model_and_uses_ctx_model_name — re-anchored to the publisher-unknown error path (only publisher-resolution failures still produce Config errors before dispatch) Tests added (5 new for the Converse path): - chat_meta_publisher_dispatches_via_converse_url — `/model/meta.*/converse` path regex + Bedrock-shape response decode - chat_amazon_nova_publisher_dispatches_via_converse_url — same dispatch contract, different publisher prefix (regression guard against a future change that hard-codes per-publisher routing inside chat_converse) - chat_stream_for_anthropic_dispatches_via_converse_stream_url — `/converse-stream` URL pattern pin for Anthropic streaming (legacy /invoke had no stream variant) - chat_stream_for_meta_dispatches_via_converse_stream_url — same pattern for non-Anthropic streaming - chat_converse_request_body_uses_text_content_block_shape — pins outbound body shape: system messages lifted to top-level `system[]`, user messages emit typed `content: [{text: "..."}]` blocks (NOT Anthropic-style flat string), no top-level `model` field `cargo test -p aisix-provider-bedrock` → 41/41 PASS (was 36 with 5 deletions + 1 update + 5 additions). `cargo clippy -p aisix-provider-bedrock --all-targets -- -D warnings` clean. `cargo fmt --all` applied. ## References (CLAUDE.md §7) - Converse API spec: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html - Converse stream events: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ConverseStream.html - aws-sdk-bedrockruntime 1.130.0 operation modules: https://docs.rs/aws-sdk-bedrockruntime/1.130.0/aws_sdk_bedrockruntime/operation/ - mock-bedrock fixture (api7/AISIX-Cloud#480) uses the same `vnd.amazon.eventstream` frame format for its encoder; the bridge here consumes the SDK-decoded events rather than re-implementing the decoder ## Out of scope (deferred follow-ups) - D7.2.a Anthropic /invoke → Converse migration. Keeping legacy path until cp-api / dashboard test fixtures are updated to assert the Converse outbound body shape instead of the Anthropic Messages shape. The customer-visible ChatResponse is identical on both paths. - Tool-use / image / document `ContentBlock` variants — chat scenarios only carry `Text` blocks today; the `Role::Tool` arm in `build_converse_inputs` is a deliberate no-op pending a structured tool-use surface on `ChatFormat`. ## Unblocks AC.10 in api7/AISIX-Cloud#302: "Bedrock 上 Claude / Llama 都能 stream". Both halves are now done — Claude streaming via Converse, Llama via Converse (chat + stream). Live e2e against mock-bedrock (#480) is the next sub-step, separate PR. --- Cargo.lock | 1 + crates/aisix-provider-bedrock/Cargo.toml | 11 +- crates/aisix-provider-bedrock/src/bridge.rs | 899 +++++++++++++++----- crates/aisix-provider-bedrock/src/lib.rs | 26 +- 4 files changed, 721 insertions(+), 216 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b9076357..225c560d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -241,6 +241,7 @@ dependencies = [ "aisix-core", "aisix-gateway", "aisix-provider-anthropic", + "async-stream", "async-trait", "aws-config", "aws-credential-types", diff --git a/crates/aisix-provider-bedrock/Cargo.toml b/crates/aisix-provider-bedrock/Cargo.toml index c1bbe217..5c6d0fd4 100644 --- a/crates/aisix-provider-bedrock/Cargo.toml +++ b/crates/aisix-provider-bedrock/Cargo.toml @@ -12,11 +12,14 @@ description = "aisix: AWS Bedrock runtime provider bridge (skeleton — multi-pu aisix-core = { path = "../aisix-core" } aisix-gateway = { path = "../aisix-gateway" } # Reuse Anthropic Messages request/response wire types for the -# `anthropic.*` Bedrock publisher (Claude on Bedrock). The body shape -# is the Anthropic Messages API minus the API-key header plus -# `anthropic_version: "bedrock-2023-05-31"` — built on top of the -# anthropic crate's typed serializers/deserializers. +# `anthropic.*` Bedrock publisher (Claude on Bedrock) non-stream +# /invoke path. The body shape is the Anthropic Messages API minus +# the API-key header plus `anthropic_version: "bedrock-2023-05-31"`. +# Other publishers (Meta / Mistral / Cohere / Amazon Titan/Nova / +# AI21) and all streaming go through the unified Converse API +# instead — see `chat_converse` / `chat_converse_stream`. aisix-provider-anthropic = { path = "../aisix-provider-anthropic" } +async-stream = "0.3" async-trait.workspace = true thiserror.workspace = true tracing.workspace = true diff --git a/crates/aisix-provider-bedrock/src/bridge.rs b/crates/aisix-provider-bedrock/src/bridge.rs index eb13fdbe..ee218ce7 100644 --- a/crates/aisix-provider-bedrock/src/bridge.rs +++ b/crates/aisix-provider-bedrock/src/bridge.rs @@ -17,15 +17,23 @@ //! point at a private deployment / VPC endpoint. use aisix_gateway::{ - Bridge, BridgeContext, BridgeError, ChatChunkStream, ChatFormat, ChatResponse, + Bridge, BridgeContext, BridgeError, ChatChunk, ChatChunkStream, ChatDelta, ChatFormat, + ChatMessage, ChatResponse, FinishReason, Role, UsageStats, }; use async_trait::async_trait; use aws_credential_types::provider::SharedCredentialsProvider; use aws_credential_types::Credentials; use aws_sdk_bedrockruntime::config::{BehaviorVersion, Region}; use aws_sdk_bedrockruntime::error::SdkError; +use aws_sdk_bedrockruntime::operation::converse::ConverseError; +use aws_sdk_bedrockruntime::operation::converse_stream::ConverseStreamError; use aws_sdk_bedrockruntime::operation::invoke_model::InvokeModelError; use aws_sdk_bedrockruntime::primitives::Blob; +use aws_sdk_bedrockruntime::types::{ + ContentBlock, ContentBlockDelta, ConversationRole, ConverseStreamOutput, + InferenceConfiguration, Message as BedrockMessage, StopReason as SdkStopReason, + SystemContentBlock, +}; use aws_sdk_bedrockruntime::Client as BedrockClient; use aws_smithy_runtime_api::client::result::ServiceError; use serde::Deserialize; @@ -39,9 +47,9 @@ use crate::wire; /// Anthropic-on-Bedrock body-shape version pin per /// . -/// Goes in the request body as the `anthropic_version` field; the -/// `model` field is stripped because Bedrock keys dispatch off the -/// URL path, not the body. +/// Used by the legacy /invoke Anthropic non-stream path +/// ([`BedrockBridge::chat_anthropic`]); the Converse path doesn't +/// need it (the SDK shapes the body). const BEDROCK_ANTHROPIC_VERSION: &str = "bedrock-2023-05-31"; /// Family Bridge for AWS Bedrock Runtime. @@ -200,7 +208,11 @@ impl BedrockPublisher { }) } - /// Human-readable name used in publisher-not-implemented errors. + /// Human-readable name. Was used by the old "publisher-not- + /// implemented" error path; kept after Phase G Step 3's Converse + /// dispatch wiring made the legacy path obsolete because the same + /// label is still useful for telemetry tagging follow-ups. + #[allow(dead_code)] fn name(&self) -> &'static str { match self { Self::Anthropic => "anthropic", @@ -478,25 +490,28 @@ impl Bridge for BedrockBridge { // for any operator default_headers override. let _ = wire::reserved_sigv4_headers(); + // Phase G productionization (#302 Step 3): unified Converse + // wire for non-Anthropic publishers (Meta / Mistral / Cohere / + // Amazon Titan / Amazon Nova / AI21). Anthropic stays on the + // legacy /invoke path (chat_anthropic, wired via #320) for + // backward compat with existing operator deployments — the + // Anthropic Messages JSON envelope is what cp-api / dashboard + // tests already pin. Moving Anthropic to Converse is a clean + // follow-up since the SDK-decoded Converse response shape is + // equivalent for the customer-visible ChatResponse. + // + // chat_stream (below) dispatches ALL publishers through + // Converse because the legacy /invoke path never had a + // streaming variant (D7.2.b was the gap). match publisher { BedrockPublisher::Anthropic => self.chat_anthropic(req, ctx, upstream_id).await, - // Audit H2: surface the operator's actual model id rather - // than the enum's Debug taxonomy (`Other` / `` - // are internal labels that don't help the customer or the - // operator diagnose). The `publisher.name()` is the - // catalog-level identifier the operator pinned. - other => Err(BridgeError::Config(format!( - "bedrock dispatch for model id {upstream_id:?} (publisher={}) \ - not yet implemented — tracked under api7/AISIX-Cloud#302 \ - Phase G (D7.3+)", - other.name() - ))), + _ => self.chat_converse(req, ctx, upstream_id).await, } } async fn chat_stream( &self, - _req: &ChatFormat, + req: &ChatFormat, ctx: &BridgeContext, ) -> Result { let upstream_id = upstream_model(ctx)?; @@ -509,62 +524,34 @@ impl Bridge for BedrockBridge { (optionally prefixed with a cross-region inference profile like us. / eu. / apac.)" )) })?; - // Audit M4: distinguish "anthropic streaming not yet wired" - // (D7.2.b — same publisher as chat, just streaming) from - // "publisher X not yet wired at all" (D7.3+). Mixing them - // would mis-route the operator to the wrong follow-up - // tracking task. - match publisher { - BedrockPublisher::Anthropic => Err(BridgeError::Config( - "bedrock anthropic streaming is not yet implemented — \ - tracked under api7/AISIX-Cloud#302 Phase G (D7.2.b)" - .into(), - )), - other => Err(BridgeError::Config(format!( - "bedrock dispatch (chat_stream) for model id {upstream_id:?} \ - (publisher={}) not yet implemented — tracked under \ - api7/AISIX-Cloud#302 Phase G (D7.3+)", - other.name() - ))), - } + // Phase G productionization (#302 Step 3): unified Converse + // stream path for all publishers — same SDK call (.converse_stream) + // owns the AWS event-stream binary frame decoding internally, + // so the bridge layer doesn't ship its own decoder. + let _ = publisher; + self.chat_converse_stream(req, ctx, upstream_id).await } } impl BedrockBridge { - /// Dispatch Anthropic-on-Bedrock chat. Body shape per + /// Dispatch Anthropic-on-Bedrock chat via the legacy /invoke + /// path. Body shape per /// : /// the Anthropic Messages JSON minus the `model` field (Bedrock /// keys dispatch off the URL) plus `anthropic_version: - /// "bedrock-2023-05-31"`. + /// "bedrock-2023-05-31"`. Kept alongside `chat_converse` because + /// the existing customer deployments + e2e tests pin the + /// Anthropic-shape outbound body; future PR can fold this into + /// the Converse path once cp-api / dashboard tests have been + /// updated to assert the Converse envelope instead. async fn chat_anthropic( &self, req: &ChatFormat, ctx: &BridgeContext, upstream_id: &str, ) -> Result { - // Parse credentials. Per-request to keep the bridge stateless - // — credential rotation lands as soon as the PK snapshot - // refreshes, no client cache invalidation needed. - let creds = BedrockSecret::parse(&ctx.provider_key.secret)?; - let endpoint_url = { - #[cfg(test)] - { - self.endpoint_url_override - .as_deref() - .or(ctx.provider_key.api_base.as_deref()) - } - #[cfg(not(test))] - { - ctx.provider_key.api_base.as_deref() - } - }; - let client = build_client(&creds, endpoint_url)?; + let client = self.build_client_from_ctx(ctx)?; - // Build the Anthropic Messages body via the shared - // serializers, then shape it for Bedrock: - // 1. Strip `model` (Bedrock takes it via URL path) - // 2. Strip `stream` (Bedrock decides via Invoke vs InvokeWithResponseStream) - // 3. Add `anthropic_version` (Bedrock-specific pin) let (system, messages) = split_system(req).map_err(|e| BridgeError::Config(format!("{e}")))?; let anthropic_req = build_request(req, upstream_id, system, messages, false); @@ -582,9 +569,6 @@ impl BedrockBridge { BridgeError::Config(format!("serialize Anthropic request body bytes: {e}")) })?; - // Dispatch via the SDK. SigV4 + retries + content-type - // headers are handled by the SDK; we pass model id + - // accept/content-type + body bytes. let started = Instant::now(); let deadline = ctx.deadline; let resp = client @@ -601,6 +585,407 @@ impl BedrockBridge { .map_err(|e| BridgeError::UpstreamDecode(e.to_string()))?; Ok(response_into_chat_response(parsed)) } + + /// Resolve credentials + build an SDK client. Pulled out of + /// `chat_converse` / `chat_converse_stream` because both methods + /// need it and the body is identical. + fn build_client_from_ctx(&self, ctx: &BridgeContext) -> Result { + // Parse credentials per-request to keep the bridge stateless — + // credential rotation lands as soon as the PK snapshot + // refreshes, no client cache invalidation needed. + let creds = BedrockSecret::parse(&ctx.provider_key.secret)?; + let endpoint_url = { + #[cfg(test)] + { + self.endpoint_url_override + .as_deref() + .or(ctx.provider_key.api_base.as_deref()) + } + #[cfg(not(test))] + { + ctx.provider_key.api_base.as_deref() + } + }; + build_client(&creds, endpoint_url) + } + + /// Dispatch Bedrock chat via the unified Converse API. + /// + /// Request: `POST /model//converse`. The SDK builds the + /// Converse JSON envelope (`messages[]`, `system[]`, + /// `inferenceConfig`) from typed builders, signs SigV4, and POSTs + /// to the resolved endpoint. Response: typed `ConverseOutput` with + /// `output.message.content[]` text blocks, `stop_reason`, `usage` + /// (input/output/total tokens), and `metrics.latency_ms`. + /// + /// Unlike the per-publisher `/invoke` legacy path, Converse uses + /// ONE wire shape for every Bedrock publisher (Anthropic, Meta, + /// Mistral, Cohere, Amazon Titan/Nova, AI21). Adding a new + /// publisher to AWS's supported list requires zero gateway-side + /// code change — the SDK upgrade picks it up. + /// + /// Reference: + async fn chat_converse( + &self, + req: &ChatFormat, + ctx: &BridgeContext, + upstream_id: &str, + ) -> Result { + let client = self.build_client_from_ctx(ctx)?; + let (system_blocks, message_blocks) = build_converse_inputs(req)?; + if message_blocks.is_empty() { + return Err(BridgeError::Config( + "bedrock converse: messages must include at least one user / \ + assistant turn (system-only requests are not supported by Converse)" + .into(), + )); + } + + let started = Instant::now(); + let deadline = ctx.deadline; + let mut call = client.converse().model_id(upstream_id); + for sb in system_blocks { + call = call.system(sb); + } + for mb in message_blocks { + call = call.messages(mb); + } + // Bedrock's InferenceConfiguration is optional — we leave it + // unset by default so the upstream model's own defaults apply. + // A follow-up override-pipeline PR can wire temperature / + // max_tokens from RequestOverrides. + let _ = InferenceConfiguration::builder(); + + let resp = call + .send() + .await + .map_err(|e| map_converse_sdk_error(e, started, deadline))?; + Ok(converse_output_into_chat_response(resp, upstream_id)) + } + + /// Dispatch Bedrock chat via the unified Converse stream API. + /// + /// Request: `POST /model//converse-stream`. The SDK + /// reads the binary `vnd.amazon.eventstream` body internally and + /// yields typed [`ConverseStreamOutput`] events through the + /// returned stream — no per-byte frame decoder lives in the + /// bridge layer. + /// + /// Event sequence per AWS docs: + /// + /// ```text + /// MessageStart → {role: "assistant"} + /// ContentBlockStart → {start: {text: ""}, contentBlockIndex: 0} + /// ContentBlockDelta → {delta: {text: "..."}, contentBlockIndex: 0} + /// ContentBlockStop → {contentBlockIndex: 0} + /// MessageStop → {stopReason: "end_turn"} + /// Metadata → {usage: {...}, metrics: {latencyMs: ...}} + /// ``` + /// + /// Each event maps to zero, one, or two [`ChatChunk`]s — see + /// [`emit_converse_chunk`]. Stream closes cleanly when AWS sends + /// the final Metadata event and the underlying eventstream + /// connection terminates. + /// + /// Reference: + async fn chat_converse_stream( + &self, + req: &ChatFormat, + ctx: &BridgeContext, + upstream_id: &str, + ) -> Result { + let client = self.build_client_from_ctx(ctx)?; + let (system_blocks, message_blocks) = build_converse_inputs(req)?; + if message_blocks.is_empty() { + return Err(BridgeError::Config( + "bedrock converse: messages must include at least one user / \ + assistant turn (system-only requests are not supported by Converse)" + .into(), + )); + } + + let started = Instant::now(); + let deadline = ctx.deadline; + let mut call = client.converse_stream().model_id(upstream_id); + for sb in system_blocks { + call = call.system(sb); + } + for mb in message_blocks { + call = call.messages(mb); + } + + let mut resp = call + .send() + .await + .map_err(|e| map_converse_stream_sdk_error(e, started, deadline))?; + + let upstream_id_owned = upstream_id.to_string(); + let stream = async_stream::try_stream! { + let mut emitted_role = false; + loop { + match resp.stream.recv().await { + Ok(Some(event)) => { + for chunk in emit_converse_chunk(event, &upstream_id_owned, &mut emitted_role) { + yield chunk; + } + } + Ok(None) => break, + Err(e) => { + // SDK already maps eventstream-decode and + // transport failures into its error type. + // Surface as Transport to flow through the + // existing customer-visible error envelope. + Err(BridgeError::Transport(format!( + "bedrock converse-stream: {e}" + )))?; + } + } + } + }; + Ok(Box::pin(stream)) + } +} + +/// Translate the gateway [`ChatFormat`] into the SDK's typed +/// `(Vec, Vec)` Converse inputs. +/// +/// System messages → top-level `system[]` blocks (Converse splits +/// them out of `messages[]` per AWS spec). +/// User / assistant messages → `messages[]` with role + text content +/// block. Bedrock Converse rejects empty `messages[]` so the caller +/// must check the returned `Vec` length. +fn build_converse_inputs( + req: &ChatFormat, +) -> Result<(Vec, Vec), BridgeError> { + let mut systems: Vec = Vec::new(); + let mut messages: Vec = Vec::new(); + for msg in &req.messages { + match msg.role { + Role::System => { + if !msg.content.is_empty() { + systems.push(SystemContentBlock::Text(msg.content.clone())); + } + } + Role::User | Role::Assistant => { + let role = if matches!(msg.role, Role::User) { + ConversationRole::User + } else { + ConversationRole::Assistant + }; + let m = BedrockMessage::builder() + .role(role) + .content(ContentBlock::Text(msg.content.clone())) + .build() + .map_err(|e| { + BridgeError::Config(format!("bedrock converse: build message: {e}")) + })?; + messages.push(m); + } + Role::Tool => { + // Tool-result messages are part of Anthropic's tool-use + // protocol; Bedrock Converse supports them via + // `ContentBlock::ToolResult` but the gateway's ChatMessage + // surface doesn't carry the structured tool_use_id + + // content shape needed to round-trip cleanly. Skip + // silently for now; a tool-use-aware follow-up PR can + // wire this when the upstream ChatFormat extends to + // carry the structured payload. + } + } + } + Ok((systems, messages)) +} + +/// Translate a Converse response into the gateway [`ChatResponse`]. +/// Text content blocks are concatenated; tool-use / image / document +/// blocks (rare in chat scenarios) are skipped. +fn converse_output_into_chat_response( + resp: aws_sdk_bedrockruntime::operation::converse::ConverseOutput, + upstream_id: &str, +) -> ChatResponse { + let (text, finish) = match resp.output() { + Some(aws_sdk_bedrockruntime::types::ConverseOutput::Message(msg)) => { + let text: String = msg + .content() + .iter() + .filter_map(|cb| match cb { + ContentBlock::Text(t) => Some(t.as_str()), + _ => None, + }) + .collect::>() + .join(""); + (text, map_stop_reason(resp.stop_reason())) + } + _ => (String::new(), FinishReason::Stop), + }; + let usage = resp + .usage() + .map(|u| UsageStats { + prompt_tokens: u.input_tokens().max(0) as u32, + completion_tokens: u.output_tokens().max(0) as u32, + total_tokens: u.total_tokens().max(0) as u32, + ..Default::default() + }) + .unwrap_or_default(); + ChatResponse { + id: String::new(), + model: upstream_id.to_string(), + message: ChatMessage::assistant(text), + finish_reason: finish, + usage, + } +} + +/// Map AWS Converse `StopReason` to the gateway's [`FinishReason`] +/// taxonomy. Reference: +/// +fn map_stop_reason(stop: &SdkStopReason) -> FinishReason { + match stop { + SdkStopReason::EndTurn => FinishReason::Stop, + SdkStopReason::MaxTokens => FinishReason::Length, + SdkStopReason::ContentFiltered | SdkStopReason::GuardrailIntervened => { + FinishReason::ContentFilter + } + SdkStopReason::ToolUse => FinishReason::ToolCalls, + SdkStopReason::StopSequence => FinishReason::Stop, + _ => FinishReason::Stop, + } +} + +/// Translate one Converse stream event into zero, one, or two +/// [`ChatChunk`]s. +/// +/// - `MessageStart` → role chunk (role=assistant) +/// - `ContentBlockDelta` (text) → content chunk +/// - `MessageStop` → finish-reason chunk +/// - `Metadata` → usage chunk +/// - Other variants (`ContentBlockStart`, `ContentBlockStop`, +/// tool-use deltas) → no chunk emitted (covered by the AWS spec +/// but not surfaced through the OpenAI-style ChatChunk shape). +fn emit_converse_chunk( + event: ConverseStreamOutput, + upstream_id: &str, + emitted_role: &mut bool, +) -> Vec { + let mut out = Vec::new(); + match event { + ConverseStreamOutput::MessageStart(_m) => { + if !*emitted_role { + *emitted_role = true; + out.push(ChatChunk { + id: String::new(), + model: upstream_id.to_string(), + delta: ChatDelta { + role: Some(Role::Assistant), + ..Default::default() + }, + finish_reason: None, + usage: None, + }); + } + } + ConverseStreamOutput::ContentBlockDelta(d) => { + if let Some(ContentBlockDelta::Text(text)) = d.delta { + out.push(ChatChunk { + id: String::new(), + model: upstream_id.to_string(), + delta: ChatDelta { + content: Some(text), + ..Default::default() + }, + finish_reason: None, + usage: None, + }); + } + } + ConverseStreamOutput::MessageStop(m) => { + out.push(ChatChunk { + id: String::new(), + model: upstream_id.to_string(), + delta: ChatDelta::default(), + finish_reason: Some(map_stop_reason(&m.stop_reason)), + usage: None, + }); + } + ConverseStreamOutput::Metadata(m) => { + if let Some(u) = m.usage { + out.push(ChatChunk { + id: String::new(), + model: upstream_id.to_string(), + delta: ChatDelta::default(), + finish_reason: None, + usage: Some(UsageStats { + prompt_tokens: u.input_tokens.max(0) as u32, + completion_tokens: u.output_tokens.max(0) as u32, + total_tokens: u.total_tokens.max(0) as u32, + ..Default::default() + }), + }); + } + } + // ContentBlockStart / ContentBlockStop carry no customer-visible + // payload in chat scenarios; tool-use / image deltas are + // deferred. The catch-all is non-exhaustive because the SDK + // enum is non_exhaustive — future AWS event types fall here + // safely. + _ => {} + } + out +} + +/// Map the SDK's `ConverseError` SdkError variant to BridgeError. +/// Same classification rules as `map_sdk_error` (UpstreamStatus for +/// retryable upstream failures, Timeout on dispatch timeout). +fn map_converse_sdk_error( + e: SdkError, + started: Instant, + deadline: Option, +) -> BridgeError { + map_aws_sdk_error_generic(e.to_string(), &e, started, deadline) +} + +fn map_converse_stream_sdk_error( + e: SdkError, + started: Instant, + deadline: Option, +) -> BridgeError { + map_aws_sdk_error_generic(e.to_string(), &e, started, deadline) +} + +/// Common AWS SDK error classifier. Tries to recover a Bedrock HTTP +/// status (4xx vs 5xx) and routes through the same UpstreamStatus / +/// Transport / Timeout taxonomy as the legacy invoke path. The +/// `service_error()` accessor is generic across operation types so +/// we delegate to a trait object via the boxed-error path. +fn map_aws_sdk_error_generic( + msg: String, + err: &SdkError, + started: Instant, + deadline: Option, +) -> BridgeError +where + E: std::fmt::Debug, +{ + // Deadline-elapsed wins over upstream classification — surfacing + // an UpstreamStatus on a timed-out request would hide the + // operator's configured deadline. + if let Some(d) = deadline { + if started.elapsed() >= d { + return BridgeError::Timeout { + elapsed_ms: started.elapsed().as_millis() as u64, + }; + } + } + match err { + SdkError::ServiceError(svc) => { + let status = svc.raw().status().as_u16(); + BridgeError::upstream_status(status, msg) + } + SdkError::TimeoutError(_) => BridgeError::Timeout { + elapsed_ms: started.elapsed().as_millis() as u64, + }, + _ => BridgeError::Transport(msg), + } } #[cfg(test)] @@ -886,31 +1271,10 @@ mod tests { } } - #[tokio::test] - async fn chat_rejects_non_anthropic_publishers_with_publisher_named() { - // Other Bedrock publishers are recognized but not yet wired - // for dispatch — the error must call out which publisher - // got rejected so the operator can pin the follow-up task. - let bridge = BedrockBridge::new(); - let ctx = BridgeContext::new( - "req-1", - sample_model_with("meta.llama3-3-70b-instruct-v1:0"), - sample_pk_with_secret(valid_secret_json()), - ); - let req = ChatFormat::new("customer-facing-name", vec![ChatMessage::user("hi")]); - let err = bridge.chat(&req, &ctx).await.unwrap_err(); - match err { - BridgeError::Config(msg) => { - assert!(msg.contains("not yet implemented")); - assert!( - msg.contains("meta") || msg.contains("Meta"), - "publisher name must appear in error; got {msg}" - ); - assert!(msg.contains("D7.3+") || msg.contains("Phase G")); - } - other => panic!("expected Config error, got {other:?}"), - } - } + // Removed `chat_rejects_non_anthropic_publishers_with_publisher_named`: + // Phase G Step 3 wires non-Anthropic publishers via Converse, so the + // "not yet implemented" error path it pinned no longer exists. + // Coverage replaced by the new Converse-path tests in this module. #[tokio::test] async fn chat_with_invalid_secret_errors_before_dispatch() { @@ -975,50 +1339,47 @@ mod tests { #[tokio::test] async fn chat_ignores_req_model_and_uses_ctx_model_name() { - // D6 audit HIGH-1 regression: dispatch must read upstream - // id from ctx.model.model_name, NOT from req.model. We use - // a non-anthropic publisher on the upstream id so the chat - // call hits the publisher-not-implemented branch (proving - // dispatch read model_name), not the publisher-unknown branch. + // D6 audit HIGH-1 regression: dispatch must read upstream id + // from ctx.model.model_name, NOT from req.model. Phase G + // Step 3 wires Converse for all known publishers, so the old + // "publisher-not-implemented" proof path is gone — we now use + // an UNKNOWN publisher prefix on ctx.model.model_name so the + // chat call hits the publisher-resolution error AND the error + // message echoes the unknown id (proving model_name was the + // source). req.model is set to a known Anthropic id; if + // dispatch had used req.model the call would have routed + // successfully or hit a different error class. let bridge = BedrockBridge::new(); let ctx = BridgeContext::new( "req-1", - sample_model_with("meta.llama3-3-70b-instruct-v1:0"), + sample_model_with("unknown-publisher.foo-bar-v1:0"), sample_pk_with_secret(valid_secret_json()), ); - // req.model deliberately set to something the publisher - // resolver would also reject if used as source of truth. - let req = ChatFormat::new("totally-bogus-model-id", vec![ChatMessage::user("hi")]); + let req = ChatFormat::new( + "anthropic.claude-3-5-sonnet-20241022-v2:0", + vec![ChatMessage::user("hi")], + ); let err = bridge.chat(&req, &ctx).await.unwrap_err(); match err { BridgeError::Config(msg) => { assert!( - msg.contains("not yet implemented"), - "must hit publisher-not-implemented (proving model_name was used); got {msg}" + msg.contains("publisher unknown"), + "must hit publisher-resolution error (proving model_name was used); got {msg}" + ); + assert!( + msg.contains("unknown-publisher.foo-bar-v1:0"), + "must echo the ctx.model.model_name value (NOT req.model); got {msg}" ); } other => panic!("expected Config error, got {other:?}"), } } - #[tokio::test] - async fn chat_stream_returns_clear_not_implemented_error() { - let bridge = BedrockBridge::new(); - let ctx = BridgeContext::new( - "req-1", - sample_model_with("anthropic.claude-3-5-sonnet-20241022-v2:0"), - sample_pk_with_secret(valid_secret_json()), - ); - let req = ChatFormat::new("customer-facing", vec![ChatMessage::user("hi")]); - let err = bridge.chat_stream(&req, &ctx).await.err().unwrap(); - match err { - BridgeError::Config(msg) => { - assert!(msg.contains("streaming is not yet implemented")); - assert!(msg.contains("D7.2.b")); - } - other => panic!("expected Config error, got {other:?}"), - } - } + // Removed `chat_stream_returns_clear_not_implemented_error`: + // Phase G Step 3 wires Converse streaming for all publishers; the + // "streaming not yet implemented" error path it pinned no longer + // exists. The new chat_converse_stream-side tests in this module + // exercise the actual streaming path. // ─── Dispatch end-to-end against wiremock via endpoint_url override ── @@ -1472,40 +1833,12 @@ mod tests { assert_eq!(chat.message.content, "global ok"); } - /// Audit H2 regression: rejection error for a not-yet-wired - /// publisher must include the operator's model id (so they can - /// open the right follow-up tracking issue) and the publisher - /// name (so dashboards can group). The earlier message echoed - /// `BedrockPublisher::Other` Debug output (`Other` / - /// ``) which is internal taxonomy that doesn't help. - #[tokio::test] - async fn chat_publisher_not_implemented_error_includes_model_id_and_publisher_name() { - let bridge = BedrockBridge::new(); - let ctx = BridgeContext::new( - "req-1", - sample_model_with("meta.llama3-3-70b-instruct-v1:0"), - sample_pk_with_secret(valid_secret_json()), - ); - let req = ChatFormat::new("customer-facing", vec![ChatMessage::user("hi")]); - let err = bridge.chat(&req, &ctx).await.unwrap_err(); - match err { - BridgeError::Config(msg) => { - assert!( - msg.contains("meta.llama3-3-70b-instruct-v1:0"), - "must include operator's model id; got {msg}" - ); - assert!( - msg.contains("publisher=meta"), - "must name the publisher catalog identifier; got {msg}" - ); - assert!( - !msg.contains("Other") && !msg.contains(""), - "must not leak internal enum taxonomy; got {msg}" - ); - } - other => panic!("expected Config error, got {other:?}"), - } - } + // Removed `chat_publisher_not_implemented_error_includes_model_id_and_publisher_name` + // (and its preceding doc comment about the H2 regression): + // Phase G Step 3 wires Converse for all non-Anthropic publishers, + // so the "not yet implemented" error path the test pinned no + // longer exists. New Converse-path tests below cover the + // non-Anthropic dispatch contract. /// Audit M2 regression: defense-in-depth model-id char check. /// Even though the AWS SDK URL-encodes reserved chars, the gateway @@ -1536,57 +1869,12 @@ mod tests { } } - /// Audit M4 regression: `chat_stream` must distinguish "anthropic - /// streaming not wired yet" (D7.2.b — same publisher as chat - /// just streaming) from "publisher X not wired at all" (D7.3+). - /// Mixing them mis-routes operators to the wrong tracking task. - #[tokio::test] - async fn chat_stream_anthropic_returns_d7_2_b_specific_error() { - let bridge = BedrockBridge::new(); - let ctx = BridgeContext::new( - "req-1", - sample_model_with("anthropic.claude-3-5-sonnet-20241022-v2:0"), - sample_pk_with_secret(valid_secret_json()), - ); - let req = ChatFormat::new("customer-facing", vec![ChatMessage::user("hi")]); - let err = bridge.chat_stream(&req, &ctx).await.err().unwrap(); - match err { - BridgeError::Config(msg) => { - assert!( - msg.contains("anthropic streaming"), - "must call out anthropic streaming specifically; got {msg}" - ); - assert!(msg.contains("D7.2.b"), "must point at D7.2.b; got {msg}"); - } - other => panic!("expected Config error, got {other:?}"), - } - } - - #[tokio::test] - async fn chat_stream_non_anthropic_publisher_returns_d7_3_specific_error() { - let bridge = BedrockBridge::new(); - let ctx = BridgeContext::new( - "req-1", - sample_model_with("meta.llama3-3-70b-instruct-v1:0"), - sample_pk_with_secret(valid_secret_json()), - ); - let req = ChatFormat::new("customer-facing", vec![ChatMessage::user("hi")]); - let err = bridge.chat_stream(&req, &ctx).await.err().unwrap(); - match err { - BridgeError::Config(msg) => { - assert!( - msg.contains("publisher=meta"), - "must call out the publisher; got {msg}" - ); - assert!(msg.contains("D7.3+"), "must point at D7.3+; got {msg}"); - assert!( - !msg.contains("D7.2.b"), - "must NOT point at the anthropic-streaming task; got {msg}" - ); - } - other => panic!("expected Config error, got {other:?}"), - } - } + // Removed `chat_stream_anthropic_returns_d7_2_b_specific_error` + // and `chat_stream_non_anthropic_publisher_returns_d7_3_specific_error`: + // Phase G Step 3 wires Converse streaming for all publishers, so + // the "not yet implemented" error paths these tests pinned no + // longer exist. New chat_converse_stream-path tests below cover + // the actual streaming dispatch. #[tokio::test] async fn chat_anthropic_translates_system_messages_to_system_field() { @@ -1636,4 +1924,213 @@ mod tests { Some("user") ); } + + // ─── Phase G Step 3 — Converse path tests ────────────────────── + + /// Canned Converse non-stream response body. Shape per AWS docs: + /// . + fn default_converse_response_template() -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "output": { + "message": { + "role": "assistant", + "content": [{"text": "hello from converse"}] + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 7, + "outputTokens": 3, + "totalTokens": 10 + }, + "metrics": {"latencyMs": 1} + })) + } + + #[tokio::test] + async fn chat_meta_publisher_dispatches_via_converse_url() { + // Non-Anthropic publisher: Phase G Step 3 wires Meta (and the + // other 4 publishers) through the unified Converse API. + // Asserts URL path is `/model//converse` (NOT /invoke) and + // the SDK's typed response decodes into the gateway's + // ChatResponse correctly. + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex( + r"^/model/meta\.llama3-3-70b-instruct-v1(:0|%3A0)/converse$", + )) + .respond_with(default_converse_response_template()) + .expect(1) + .mount(&server) + .await; + + let bridge = BedrockBridge::new().with_endpoint_override(server.uri()); + let ctx = BridgeContext::new( + "req-1", + sample_model_with("meta.llama3-3-70b-instruct-v1:0"), + sample_pk_with_secret(valid_secret_json()), + ); + let req = ChatFormat::new("my-llama", vec![ChatMessage::user("hi")]); + let chat = bridge.chat(&req, &ctx).await.unwrap(); + assert_eq!(chat.message.content, "hello from converse"); + assert_eq!(chat.usage.prompt_tokens, 7); + assert_eq!(chat.usage.completion_tokens, 3); + assert_eq!(chat.usage.total_tokens, 10); + assert_eq!(chat.finish_reason, FinishReason::Stop); + } + + #[tokio::test] + async fn chat_amazon_nova_publisher_dispatches_via_converse_url() { + // Second non-Anthropic publisher — Amazon Nova. Same Converse + // dispatch as Meta, different model id prefix; pins that the + // dispatch is genuinely uniform across publishers (regression + // guard against a future change that hard-codes per-publisher + // routing inside chat_converse itself). + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex( + r"^/model/amazon\.nova-pro-v1(:0|%3A0)/converse$", + )) + .respond_with(default_converse_response_template()) + .expect(1) + .mount(&server) + .await; + + let bridge = BedrockBridge::new().with_endpoint_override(server.uri()); + let ctx = BridgeContext::new( + "req-1", + sample_model_with("amazon.nova-pro-v1:0"), + sample_pk_with_secret(valid_secret_json()), + ); + let req = ChatFormat::new("my-nova", vec![ChatMessage::user("hi")]); + let chat = bridge.chat(&req, &ctx).await.unwrap(); + assert_eq!(chat.message.content, "hello from converse"); + } + + #[tokio::test] + async fn chat_stream_for_anthropic_dispatches_via_converse_stream_url() { + // Stream dispatch goes through Converse for ALL publishers + // including Anthropic (the legacy /invoke path never had a + // stream variant). Wiremock here only needs to acknowledge + // the request — full event-stream binary frame decode is + // exercised at the SDK level, beyond the scope of unit + // tests; this guard just verifies the URL routing. + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex( + r"^/model/anthropic\.claude-3-5-sonnet-20241022-v2(:0|%3A0)/converse-stream$", + )) + // Returning a 200 with an empty body causes the SDK to + // emit a transport / decode error which we catch below. + // The test's job is to pin the URL pattern, not to + // exercise the eventstream decoder. + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + + let bridge = BedrockBridge::new().with_endpoint_override(server.uri()); + let ctx = BridgeContext::new( + "req-1", + sample_model_with("anthropic.claude-3-5-sonnet-20241022-v2:0"), + sample_pk_with_secret(valid_secret_json()), + ); + let req = ChatFormat::new("my-claude", vec![ChatMessage::user("hi")]); + // Either the chat_stream() call returns Ok (and the stream + // errors on first .next() because the empty body isn't a + // valid eventstream) OR returns Err directly from the SDK. + // Both prove the dispatch reached wiremock — wiremock's + // `.expect(1)` enforces this on drop. + let _ = bridge.chat_stream(&req, &ctx).await; + } + + #[tokio::test] + async fn chat_stream_for_meta_dispatches_via_converse_stream_url() { + // Same as the Anthropic-stream test but for a non-Anthropic + // publisher — pins that chat_stream wiring is uniform across + // publishers. + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex( + r"^/model/meta\.llama3-3-70b-instruct-v1(:0|%3A0)/converse-stream$", + )) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + + let bridge = BedrockBridge::new().with_endpoint_override(server.uri()); + let ctx = BridgeContext::new( + "req-1", + sample_model_with("meta.llama3-3-70b-instruct-v1:0"), + sample_pk_with_secret(valid_secret_json()), + ); + let req = ChatFormat::new("my-llama", vec![ChatMessage::user("hi")]); + let _ = bridge.chat_stream(&req, &ctx).await; + } + + #[tokio::test] + async fn chat_converse_request_body_uses_text_content_block_shape() { + // Pin the outbound request body shape: Converse expects + // `messages: [{role, content: [{text: "..."}]}]`. A + // regression that emitted Anthropic-shape `messages: [{role, + // content: "..."}]` (string instead of typed block array) + // would 400 upstream. We use Meta here so the test routes + // through chat_converse (the Anthropic publisher would route + // through legacy /invoke and skip the Converse body shape). + let server = MockServer::start().await; + let responder = CapturingResponder::default(); + Mock::given(method("POST")) + .and(path_regex(r"^/model/.+/converse$")) + .respond_with(responder.clone()) + .expect(1) + .mount(&server) + .await; + + let bridge = BedrockBridge::new().with_endpoint_override(server.uri()); + let ctx = BridgeContext::new( + "req-1", + sample_model_with("meta.llama3-3-70b-instruct-v1:0"), + sample_pk_with_secret(valid_secret_json()), + ); + let req = ChatFormat::new( + "my-llama", + vec![ + ChatMessage::system("you are concise"), + ChatMessage::user("hi"), + ], + ); + let _ = bridge.chat(&req, &ctx).await; + + let body = responder.captured_body.lock().unwrap().clone().unwrap(); + // System message must be lifted to top-level `system` array + // (Converse splits systems out of `messages` per AWS spec). + let system = body.get("system").and_then(|v| v.as_array()).unwrap(); + assert_eq!(system.len(), 1, "system block must be present; body={body}"); + assert_eq!( + system[0].get("text").and_then(|v| v.as_str()), + Some("you are concise") + ); + // messages[] must use typed content-block shape, NOT Anthropic- + // style flat string. + let messages = body.get("messages").and_then(|v| v.as_array()).unwrap(); + assert_eq!(messages.len(), 1); + let content = messages[0] + .get("content") + .and_then(|v| v.as_array()) + .expect("content must be an array of typed blocks"); + assert_eq!(content.len(), 1); + assert_eq!( + content[0].get("text").and_then(|v| v.as_str()), + Some("hi"), + "content block must carry `text` field; body={body}" + ); + // Per Bedrock Converse spec, the `model` field is NOT part of + // the body — model id flows via URL path. The SDK should not + // emit it. + assert!( + body.get("model").is_none(), + "Converse body must NOT carry top-level `model` field; body={body}" + ); + } } diff --git a/crates/aisix-provider-bedrock/src/lib.rs b/crates/aisix-provider-bedrock/src/lib.rs index 5adb1d23..7549ea98 100644 --- a/crates/aisix-provider-bedrock/src/lib.rs +++ b/crates/aisix-provider-bedrock/src/lib.rs @@ -10,18 +10,22 @@ //! "bedrock-2023-05-31"` in body not header) //! - [x] D7.6 — Cross-region inference profiles (`us.`/`eu.`/`apac.`/ //! `global.`/`us-gov.` prefixes stripped by [`bridge::BedrockPublisher::from_model_id`]) -//! - [ ] D7.2.b — Anthropic-on-Bedrock streaming via -//! `invoke_model_with_response_stream` (AWS event-stream framed, -//! NOT canonical SSE; reuses the Anthropic typed-event stream state -//! machine from `aisix-provider-anthropic`) -//! - [ ] D7.3 — Meta-on-Bedrock dispatch (Llama 3 / 3.1 / 3.2 / 3.3) -//! - [ ] D7.4 — Mistral / Amazon Titan / Amazon Nova / Cohere / AI21 -//! per-publisher request bodies +//! - [x] D7.2.b — All-publisher streaming via the unified Converse API +//! (`/model//converse-stream`). The SDK owns the +//! `vnd.amazon.eventstream` binary frame decoding; the bridge layer +//! maps typed events to ChatChunk via `emit_converse_chunk`. +//! - [x] D7.3 — Meta-on-Bedrock dispatch via Converse +//! (`/model/meta.*/converse`). SDK shapes the Converse JSON body +//! from typed builders. +//! - [x] D7.4 — Mistral / Amazon Titan / Amazon Nova / Cohere / AI21 +//! dispatch via Converse (same unified path; AWS handles +//! publisher-specific body shaping inside Bedrock). //! -//! Until D7.2.b lands, `chat_stream()` returns a clear -//! `BridgeError::Config(...)` referencing the streaming follow-up. -//! Publishers other than Anthropic return a publisher-specific -//! "not yet implemented" error from `chat()` / `chat_stream()`. +//! Anthropic non-streaming keeps the legacy `/invoke` path +//! (`chat_anthropic`) for backward compat with existing operator +//! deployments + e2e test fixtures. Anthropic streaming, and all +//! non-Anthropic dispatch (stream or non-stream), flow through +//! Converse. //! //! # Multi-publisher single-entry model //! From b97829c8641402ffb7d4121406185de43babf4db Mon Sep 17 00:00:00 2001 From: Ming Wen Date: Mon, 25 May 2026 08:18:29 +0800 Subject: [PATCH 2/2] fix(bedrock): preserve Bedrock wire + Retry-After + InferenceConfig (audit HIGH/MEDIUM on #389) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Independent audit-aigw-389-bedrock-converse flagged 2 HIGH + 2 MEDIUM bugs in the original Converse wiring. All fixed in this commit; 45/45 tests pass. ## HIGH-1 + HIGH-2 — `map_aws_sdk_error_generic` silently regressed the legacy /invoke path's error envelope hardening The original helper collapsed all ServiceError responses to `BridgeError::upstream_status(status, msg)`, a convenience constructor that sets `wire: Unknown` + `parsed: None` + `retry_after: None`. The legacy `chat_anthropic` /invoke path uses `map_service_error` which builds a fully-shaped `BridgeError::UpstreamStatus` with: - wire: Bedrock — required by error_translate to render Bedrock- shape errors back to OpenAI/Anthropic-shape clients - retry_after: parsed from upstream Retry-After header — required by the cooldown layer to honour AWS throttle hints - parsed.kind: extracted via .meta().code() — distinguishes a ThrottlingException 429 from a different throttle source Without this fix every non-Anthropic publisher + all streaming silently shipped degraded errors vs. the legacy path — broke the PR #323 audit (MEDIUM-2) hardening it had just restored. Fix: introduce `bedrock_service_error_to_upstream_status` (generic over `E: ProvideErrorMetadata`) that mirrors `map_service_error`'s field-by-field construction. Both Converse + ConverseStream errors now flow through it. ## MEDIUM-1 — `e.to_string()` produced opaque "service error" strings as the customer-visible message Same root cause; resolved by the HIGH fix above. The customer- visible `message` now uses the legacy path's canned status-keyed phrase ("upstream rate limited", "upstream authentication failed", etc.) — preserving the operator-ARN redaction the legacy path also enforces. ## MEDIUM-2 — `chat_converse` silently dropped ChatFormat's temperature/max_tokens/top_p The original `chat_converse` constructed an `InferenceConfiguration::builder()` and immediately discarded it with `let _ =`. Every non-Anthropic Bedrock customer using these knobs would see them vanish — real behavioural regression vs. the legacy chat_anthropic path which forwards them via build_request → AnthropicRequest. Fix: new `build_inference_config(req)` helper extracts temperature / max_tokens / top_p when set, returns None when all absent (so we don't emit an empty `inferenceConfig: {}` which 400s on some Bedrock publishers). Wired into both chat_converse and chat_converse_stream. ## Tests added (5 new for the audit fixes) - chat_converse_maps_upstream_429_with_retry_after_and_bedrock_wire pins the HIGH-1+2 fix: 429 + Retry-After: 42 → UpstreamStatus with status=429, message="upstream rate limited", wire=Bedrock, retry_after=Some(42s), parsed.kind="ThrottlingException" - chat_converse_maps_upstream_4xx_to_canned_message_with_bedrock_wire pins HIGH-1 for the non-throttle path: 403 with leaky ARN body → canned "upstream authentication failed" (NO ARN leak), wire=Bedrock, parsed.kind="AccessDeniedException" - chat_converse_wires_temperature_max_tokens_top_p_into_inference_config pins MEDIUM-2: ChatFormat knobs reach Bedrock's body as camelCase {temperature, maxTokens, topP} - chat_converse_omits_inference_config_when_no_knobs_set companion: empty knobs → no inferenceConfig field on the wire `cargo test -p aisix-provider-bedrock` → 45/45 PASS (was 41; +4 new audit-regression tests). `cargo clippy -p aisix-provider-bedrock --all-targets -- -D warnings` clean. ## LOWs (non-blocking, deferred per CLAUDE.md §8) - LOW-1: Role::Tool messages silently dropped on Converse path — deliberate no-op pending a structured tool-use surface on ChatFormat. Documented in `build_converse_inputs`. - LOW-2: Per-publisher dispatch test matrix coverage (Mistral / Cohere / AmazonTitan / AI21 not individually pinned). The match arm in chat() is structurally uniform (`_ => chat_converse`) so routing is correct by construction; the two existing Meta + Amazon Nova tests already cover the non-Anthropic dispatch contract. --- crates/aisix-provider-bedrock/src/bridge.rs | 345 ++++++++++++++++++-- 1 file changed, 323 insertions(+), 22 deletions(-) diff --git a/crates/aisix-provider-bedrock/src/bridge.rs b/crates/aisix-provider-bedrock/src/bridge.rs index ee218ce7..4625dcfa 100644 --- a/crates/aisix-provider-bedrock/src/bridge.rs +++ b/crates/aisix-provider-bedrock/src/bridge.rs @@ -36,6 +36,7 @@ use aws_sdk_bedrockruntime::types::{ }; use aws_sdk_bedrockruntime::Client as BedrockClient; use aws_smithy_runtime_api::client::result::ServiceError; +use aws_smithy_types::error::metadata::ProvideErrorMetadata; use serde::Deserialize; use std::time::{Duration, Instant}; @@ -650,11 +651,15 @@ impl BedrockBridge { for mb in message_blocks { call = call.messages(mb); } - // Bedrock's InferenceConfiguration is optional — we leave it - // unset by default so the upstream model's own defaults apply. - // A follow-up override-pipeline PR can wire temperature / - // max_tokens from RequestOverrides. - let _ = InferenceConfiguration::builder(); + // Audit MEDIUM-2 (PR #389): wire ChatFormat's + // temperature/max_tokens/top_p through to Converse's + // InferenceConfiguration. The legacy chat_anthropic path + // forwards these via build_request → AnthropicRequest; + // dropping them on the Converse path is a silent + // behavioural regression for every non-Anthropic customer. + if let Some(cfg) = build_inference_config(req) { + call = call.inference_config(cfg); + } let resp = call .send() @@ -713,6 +718,10 @@ impl BedrockBridge { for mb in message_blocks { call = call.messages(mb); } + // Mirror chat_converse on the stream path — same audit fix. + if let Some(cfg) = build_inference_config(req) { + call = call.inference_config(cfg); + } let mut resp = call .send() @@ -934,14 +943,14 @@ fn emit_converse_chunk( } /// Map the SDK's `ConverseError` SdkError variant to BridgeError. -/// Same classification rules as `map_sdk_error` (UpstreamStatus for -/// retryable upstream failures, Timeout on dispatch timeout). +/// Delegates to the generic helper so both Converse and ConverseStream +/// share the wire-shape + Retry-After preservation logic. fn map_converse_sdk_error( e: SdkError, started: Instant, deadline: Option, ) -> BridgeError { - map_aws_sdk_error_generic(e.to_string(), &e, started, deadline) + map_aws_sdk_error_generic(e, started, deadline) } fn map_converse_stream_sdk_error( @@ -949,22 +958,37 @@ fn map_converse_stream_sdk_error( started: Instant, deadline: Option, ) -> BridgeError { - map_aws_sdk_error_generic(e.to_string(), &e, started, deadline) + map_aws_sdk_error_generic(e, started, deadline) } -/// Common AWS SDK error classifier. Tries to recover a Bedrock HTTP -/// status (4xx vs 5xx) and routes through the same UpstreamStatus / -/// Transport / Timeout taxonomy as the legacy invoke path. The -/// `service_error()` accessor is generic across operation types so -/// we delegate to a trait object via the boxed-error path. +/// Common AWS SDK error classifier. Routes through the same +/// UpstreamStatus / Transport / Timeout taxonomy the legacy +/// `map_service_error` uses for `/invoke`, preserving: +/// +/// - `wire: UpstreamWire::Bedrock` — so `error_translate` renders +/// Bedrock-shape errors back to OpenAI/Anthropic-shape clients +/// - `retry_after: Some(...)` — parsed from the upstream's +/// `Retry-After` header so the cooldown layer honours the AWS +/// throttle hint instead of falling back to its default +/// - `parsed.kind: Some()` — so a 429 from +/// `ThrottlingException` is distinguishable from a 429 from a +/// different throttle source +/// +/// The `E: ProvideErrorMetadata` bound lets the helper extract +/// `.meta().code()` regardless of which operation the error came +/// from (Converse vs. ConverseStream variants of the same shape). +/// +/// Audit MEDIUM / HIGH on PR #389 — the prior implementation +/// collapsed to `BridgeError::upstream_status(status, msg)` which +/// sets `wire: Unknown` and `retry_after: None`, silently regressing +/// the legacy /invoke path's hardening (PR #323 MEDIUM-2). fn map_aws_sdk_error_generic( - msg: String, - err: &SdkError, + err: SdkError, started: Instant, deadline: Option, ) -> BridgeError where - E: std::fmt::Debug, + E: std::fmt::Debug + ProvideErrorMetadata, { // Deadline-elapsed wins over upstream classification — surfacing // an UpstreamStatus on a timed-out request would hide the @@ -977,15 +1001,92 @@ where } } match err { - SdkError::ServiceError(svc) => { - let status = svc.raw().status().as_u16(); - BridgeError::upstream_status(status, msg) - } + SdkError::ServiceError(svc) => bedrock_service_error_to_upstream_status(svc), SdkError::TimeoutError(_) => BridgeError::Timeout { elapsed_ms: started.elapsed().as_millis() as u64, }, - _ => BridgeError::Transport(msg), + other => BridgeError::Transport(format!("{other}")), + } +} + +/// Translate an AWS SDK `ServiceError` into a fully-shaped +/// `BridgeError::UpstreamStatus` — same fields as the legacy +/// `map_service_error` would emit for `/invoke`. Generic over the +/// operation error type so the same logic covers Converse + +/// ConverseStream + any future Bedrock op. +fn bedrock_service_error_to_upstream_status( + svc: ServiceError, +) -> BridgeError +where + E: ProvideErrorMetadata, +{ + let kind = svc.err().meta().code().map(str::to_string); + let raw = svc.raw(); + let status = raw.status().as_u16(); + // Convert smithy HeaderMap → http::HeaderMap so we can reuse the + // gateway-level `parse_retry_after` helper. Headers with invalid + // bytes are dropped (defensive — SDK should not produce them). + let mut hdrs = http::HeaderMap::new(); + for (k, v) in raw.headers() { + if let (Ok(name), Ok(val)) = ( + http::HeaderName::from_bytes(k.as_bytes()), + http::HeaderValue::from_str(v), + ) { + hdrs.insert(name, val); + } + } + let retry_after = aisix_gateway::parse_retry_after(&hdrs); + let message = match status { + 401 | 403 => "upstream authentication failed".to_string(), + 404 => "upstream model not found".to_string(), + 408 => "upstream request timeout".to_string(), + 429 => "upstream rate limited".to_string(), + _ => format!("upstream returned {status}"), + }; + // SECURITY: the AWS error message embeds operator-internal + // taxonomy (ARNs, region, account id, IAM role names). We surface + // only the AWS error CODE (e.g. "ThrottlingException") for the + // error_translate layer; the customer-visible `message` is the + // canned status-keyed phrase above. + let parsed = kind.as_ref().map(|k| { + Box::new(aisix_gateway::UpstreamErrorView { + kind: Some(k.clone()), + message: None, + code: None, + param: None, + }) + }); + BridgeError::UpstreamStatus { + status, + message, + parsed, + wire: aisix_gateway::UpstreamWire::Bedrock, + retry_after, + } +} + +/// Build an `InferenceConfiguration` from the gateway's +/// [`ChatFormat`] knobs, or `None` if no knobs are set. Caller +/// only attaches when `Some` so we don't ship an empty +/// `inferenceConfig: {}` to Bedrock (some publishers tolerate it, +/// others 400). +fn build_inference_config(req: &ChatFormat) -> Option { + if req.temperature.is_none() && req.max_tokens.is_none() && req.top_p.is_none() { + return None; } + let mut b = InferenceConfiguration::builder(); + if let Some(t) = req.temperature { + b = b.temperature(t); + } + if let Some(m) = req.max_tokens { + // Bedrock's API uses i32; ChatFormat's u32 is always non-negative. + let clamped = i32::try_from(m).unwrap_or(i32::MAX); + b = b.max_tokens(clamped); + } + if let Some(p) = req.top_p { + b = b.top_p(p); + } + Some(b.build()) } #[cfg(test)] @@ -2069,6 +2170,206 @@ mod tests { let _ = bridge.chat_stream(&req, &ctx).await; } + /// Audit HIGH-1+HIGH-2 (PR #389): Converse 4xx must preserve + /// `wire: Bedrock`, parse `Retry-After` into `retry_after`, and + /// extract the AWS error code into `parsed.kind` — matching the + /// legacy /invoke path's `map_service_error`. Without this fix, + /// non-Anthropic publishers and all streaming silently regress + /// `wire` to `Unknown` (breaks error_translate) and drop the + /// Retry-After hint (breaks cooldown auto-tuning). + #[tokio::test] + async fn chat_converse_maps_upstream_429_with_retry_after_and_bedrock_wire() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex(r"^/model/meta\..+/converse$")) + .respond_with( + ResponseTemplate::new(429) + .insert_header("Retry-After", "42") + .insert_header("x-amzn-ErrorType", "ThrottlingException") + .set_body_json(serde_json::json!({ + "__type": "ThrottlingException", + "message": "Rate exceeded" + })), + ) + .mount(&server) + .await; + + let bridge = BedrockBridge::new().with_endpoint_override(server.uri()); + let ctx = BridgeContext::new( + "req-1", + sample_model_with("meta.llama3-3-70b-instruct-v1:0"), + sample_pk_with_secret(valid_secret_json()), + ); + let req = ChatFormat::new("my-llama", vec![ChatMessage::user("hi")]); + let err = bridge.chat(&req, &ctx).await.unwrap_err(); + match err { + BridgeError::UpstreamStatus { + status, + message, + wire, + retry_after, + parsed, + } => { + assert_eq!(status, 429); + assert_eq!(message, "upstream rate limited"); + assert_eq!( + wire, + aisix_gateway::UpstreamWire::Bedrock, + "must preserve Bedrock wire for error_translate" + ); + assert_eq!( + retry_after, + Some(std::time::Duration::from_secs(42)), + "must propagate upstream Retry-After to cooldown layer" + ); + let parsed = parsed.expect("must expose AWS error kind to error_translate"); + assert_eq!( + parsed.kind.as_deref(), + Some("ThrottlingException"), + "must extract AWS error code into parsed.kind" + ); + } + other => panic!("expected UpstreamStatus, got {other:?}"), + } + } + + /// Audit HIGH-1 (PR #389): Converse 4xx that isn't 429 must + /// still get the canned status-keyed message + Bedrock wire, + /// matching the legacy /invoke path. + #[tokio::test] + async fn chat_converse_maps_upstream_4xx_to_canned_message_with_bedrock_wire() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex(r"^/model/meta\..+/converse$")) + .respond_with( + ResponseTemplate::new(403) + .insert_header("x-amzn-ErrorType", "AccessDeniedException") + .set_body_json(serde_json::json!({ + "__type": "AccessDeniedException", + "message": "User: arn:aws:iam::123:user/x is not authorized" + })), + ) + .mount(&server) + .await; + + let bridge = BedrockBridge::new().with_endpoint_override(server.uri()); + let ctx = BridgeContext::new( + "req-1", + sample_model_with("meta.llama3-3-70b-instruct-v1:0"), + sample_pk_with_secret(valid_secret_json()), + ); + let req = ChatFormat::new("my-llama", vec![ChatMessage::user("hi")]); + let err = bridge.chat(&req, &ctx).await.unwrap_err(); + match err { + BridgeError::UpstreamStatus { + status, + message, + wire, + parsed, + .. + } => { + assert_eq!(status, 403); + assert_eq!( + message, "upstream authentication failed", + "must use canned phrase (NOT echo upstream ARN)" + ); + assert_eq!(wire, aisix_gateway::UpstreamWire::Bedrock); + assert!( + !message.contains("arn:aws:iam"), + "canned phrase must not leak the operator-internal ARN; got {message}" + ); + let parsed = parsed.expect("must expose AWS error kind"); + assert_eq!(parsed.kind.as_deref(), Some("AccessDeniedException")); + } + other => panic!("expected UpstreamStatus, got {other:?}"), + } + } + + /// Audit MEDIUM-2 (PR #389): ChatFormat.temperature / + /// .max_tokens / .top_p must flow through to Converse's + /// InferenceConfiguration. Dropping them silently is a + /// behavioural regression for every non-Anthropic Bedrock + /// customer who was setting these knobs via the gateway. + #[tokio::test] + async fn chat_converse_wires_temperature_max_tokens_top_p_into_inference_config() { + let server = MockServer::start().await; + let responder = CapturingResponder::default(); + Mock::given(method("POST")) + .and(path_regex(r"^/model/.+/converse$")) + .respond_with(responder.clone()) + .expect(1) + .mount(&server) + .await; + + let bridge = BedrockBridge::new().with_endpoint_override(server.uri()); + let mut req = ChatFormat::new("my-llama", vec![ChatMessage::user("hi")]); + // Pick values that survive the f32 → JSON → f64 round-trip + // cleanly (powers of 1/2 / 1/4 / 1/8 ...). Avoids brittle + // `0.7 ≠ 0.6999999881` precision-comparison failures. + req.temperature = Some(0.5); + req.max_tokens = Some(128); + req.top_p = Some(0.75); + let ctx = BridgeContext::new( + "req-1", + sample_model_with("meta.llama3-3-70b-instruct-v1:0"), + sample_pk_with_secret(valid_secret_json()), + ); + let _ = bridge.chat(&req, &ctx).await; + + let body = responder.captured_body.lock().unwrap().clone().unwrap(); + let cfg = body + .get("inferenceConfig") + .and_then(|v| v.as_object()) + .expect("inferenceConfig must be set when ChatFormat carries knobs"); + // Bedrock uses camelCase per AWS spec. + assert_eq!( + cfg.get("temperature").and_then(|v| v.as_f64()), + Some(0.5), + "temperature must propagate; body={body}" + ); + assert_eq!( + cfg.get("maxTokens").and_then(|v| v.as_i64()), + Some(128), + "maxTokens (camelCase) must propagate; body={body}" + ); + assert_eq!( + cfg.get("topP").and_then(|v| v.as_f64()), + Some(0.75), + "topP (camelCase) must propagate; body={body}" + ); + } + + /// Companion to the wires-temperature test: when no knobs are + /// set, inferenceConfig must be OMITTED from the body — sending + /// an empty `{}` causes some Bedrock publishers to 400. + #[tokio::test] + async fn chat_converse_omits_inference_config_when_no_knobs_set() { + let server = MockServer::start().await; + let responder = CapturingResponder::default(); + Mock::given(method("POST")) + .and(path_regex(r"^/model/.+/converse$")) + .respond_with(responder.clone()) + .expect(1) + .mount(&server) + .await; + + let bridge = BedrockBridge::new().with_endpoint_override(server.uri()); + // ChatFormat::new leaves temperature/max_tokens/top_p None. + let req = ChatFormat::new("my-llama", vec![ChatMessage::user("hi")]); + let ctx = BridgeContext::new( + "req-1", + sample_model_with("meta.llama3-3-70b-instruct-v1:0"), + sample_pk_with_secret(valid_secret_json()), + ); + let _ = bridge.chat(&req, &ctx).await; + + let body = responder.captured_body.lock().unwrap().clone().unwrap(); + assert!( + body.get("inferenceConfig").is_none(), + "inferenceConfig must be absent when ChatFormat has no knobs; body={body}" + ); + } + #[tokio::test] async fn chat_converse_request_body_uses_text_content_block_shape() { // Pin the outbound request body shape: Converse expects