From d330894b30a32a7a6ceea5beb719f26200fd36ea Mon Sep 17 00:00:00 2001 From: Ming Wen Date: Sun, 17 May 2026 20:57:25 +0800 Subject: [PATCH 1/2] feat(provider-bedrock): wire Anthropic-on-Bedrock chat via aws-sdk-bedrockruntime (D7.2.a, #302 Phase G) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the skeleton's `BridgeError::Config("not yet implemented")` stubs in `aisix-provider-bedrock` with real Bedrock dispatch for the `anthropic.*` publisher (Claude on Bedrock). The AWS SDK handles SigV4 signing, retries, and (in the streaming follow-up D7.2.b) the binary event-stream framing. Other Bedrock publishers (`meta.*`, `mistral.*`, `amazon.*`, `cohere.*`, `ai21.*`) return a clear publisher-named "not yet implemented — D7.3+" error. `chat_stream()` returns "streaming not yet implemented — D7.2.b" for all publishers. ## Wire shape pinned (Anthropic on Bedrock) Per the AWS docs URL cited inline: - URL: `POST /model//invoke` (model id encoded into the path — the `:` in `anthropic.claude-3-5-sonnet-20241022-v2:0` becomes `%3A`; tests use a regex matcher to stay SDK-version-tolerant) - Body: Anthropic Messages JSON, with three Bedrock-specific shape rules pinned by `chat_anthropic_body_contains_bedrock_anthropic_version_and_no_model_field`: 1. `anthropic_version: "bedrock-2023-05-31"` MUST be present 2. `model` MUST be absent (Bedrock dispatches via URL path) 3. `stream` MUST be absent for InvokeModel (non-streaming) - Auth: SigV4 — pinned by `chat_anthropic_uses_sigv4_authorization_header` (asserts `Authorization: AWS4-HMAC-SHA256 ...` + `x-amz-date` reach the wire) ## Credentials convention `ProviderKey.secret` is a JSON-encoded `{access_key_id, secret_access_key, session_token?, region}` blob. The bridge parses it per request (cheap — strings only); credential rotation lands as soon as the PK snapshot refreshes, no client cache to invalidate. The cp-api side delivers credentials decrypted (mTLS- only etcd channel; same trust boundary as Guardrail credentials). `ProviderKey.api_base` (if set) is forwarded as the SDK's `endpoint_url` so operators can point at a private deployment / VPC endpoint. Tests use the same path via a `#[cfg(test)]` `endpoint_url_override` seam to drive `bridge.chat()` end-to-end against wiremock. ## Anthropic wire reuse The Anthropic-on-Bedrock body shape is the Anthropic Messages API minus the `model` field plus `anthropic_version`. To avoid re-implementing the wire types, promotes the needed items in `aisix-provider-anthropic::wire` from `pub(crate)` to `pub`: - `AnthropicRequest`, `AnthropicMessage` - `AnthropicResponse`, `AnthropicResponseBlock`, `AnthropicUsage` - `AnthropicStreamEvent` + substructs (for D7.2.b streaming follow-up) - `StreamState` - `build_request`, `messages_from`, `split_system`, `response_into_chat_response` - `translate_openai_tools_to_anthropic`, `translate_openai_tool_choice_to_anthropic` - `DEFAULT_MAX_TOKENS` Same pattern as PR #319 for `aisix-provider-openai::wire` — wire types are JSON-shape contracts, public surface for sibling provider crates (workspace-internal, not a stability promise to external SDK consumers). `AnthropicBridge` itself is unchanged. ## Test coverage 35 unit tests, all passing: - Publisher resolution (12 tests, preserved from skeleton): every publisher tag, cross-region prefixes (`us.`/`eu.`/`apac.`/`global.`/ `us-gov.`), guard against treating a publisher segment as region, catch-all to `BedrockPublisher::Other` for not-yet-wired publishers - `BedrockSecret` parsing (5 tests): full form, with session_token, empty rejected, non-JSON rejected with generic shape error, missing region rejected, **error message does NOT echo raw secret bytes** (M1-style leak guard) - Pre-dispatch validation (6 tests): unknown publisher, non-Anthropic publisher named in error, invalid secret, empty secret, missing model_name, `chat_ignores_req_model_and_uses_ctx_model_name` (D6 audit HIGH-1 regression carried over) - `chat_stream` returns clear D7.2.b not-implemented error - Bridge dispatch via `bridge.chat()` end-to-end against wiremock (7 tests, the highest-confidence pins): - URL path includes deployment id with `:` URL-encoded - body carries `anthropic_version=bedrock-2023-05-31`, no `model`, no `stream` - SigV4 `Authorization: AWS4-HMAC-SHA256` header reaches the wire along with `x-amz-date` - `tool_use` response blocks translate to OpenAI `tool_calls` shape (via reused Anthropic crate's converter) - 4xx upstream error body redacted to canned phrase (does NOT echo operator's account number / IAM role ARN — Audit M1) - 429 mapped to canned "rate limited" message + status preserved - Cross-region inference profile (`us.anthropic.claude-*`) dispatches with the full prefixed model id in the URL - `system` role messages translated to top-level Anthropic `system` field (not left in `messages[]`) --- Cargo.lock | 10 + crates/aisix-provider-anthropic/src/lib.rs | 2 +- crates/aisix-provider-anthropic/src/wire.rs | 34 +- crates/aisix-provider-bedrock/Cargo.toml | 20 +- crates/aisix-provider-bedrock/src/bridge.rs | 1101 +++++++++++++++---- crates/aisix-provider-bedrock/src/lib.rs | 48 +- 6 files changed, 946 insertions(+), 269 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8e432fa5..9ea9219e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -240,11 +240,21 @@ version = "0.1.0" dependencies = [ "aisix-core", "aisix-gateway", + "aisix-provider-anthropic", "async-trait", + "aws-config", + "aws-credential-types", + "aws-sdk-bedrockruntime", + "aws-smithy-async", + "aws-smithy-runtime-api", + "aws-smithy-types", + "http 1.4.0", + "serde", "serde_json", "thiserror 1.0.69", "tokio", "tracing", + "wiremock", ] [[package]] diff --git a/crates/aisix-provider-anthropic/src/lib.rs b/crates/aisix-provider-anthropic/src/lib.rs index 123e1d83..a71d93d8 100644 --- a/crates/aisix-provider-anthropic/src/lib.rs +++ b/crates/aisix-provider-anthropic/src/lib.rs @@ -13,7 +13,7 @@ #![deny(rust_2018_idioms)] mod bridge; -mod wire; +pub mod wire; pub use bridge::{AnthropicBridge, ANTHROPIC_DEFAULT_BASE, ANTHROPIC_VERSION}; diff --git a/crates/aisix-provider-anthropic/src/wire.rs b/crates/aisix-provider-anthropic/src/wire.rs index dc338539..d7aa12d2 100644 --- a/crates/aisix-provider-anthropic/src/wire.rs +++ b/crates/aisix-provider-anthropic/src/wire.rs @@ -28,10 +28,10 @@ use serde::{Deserialize, Serialize}; /// Anthropic requires a non-zero `max_tokens`. Clients that omit it get /// this ceiling — generous enough to cover normal completions, conservative /// enough that a runaway prompt doesn't burn tokens silently. -pub(crate) const DEFAULT_MAX_TOKENS: u32 = 4096; +pub const DEFAULT_MAX_TOKENS: u32 = 4096; #[derive(Debug, Clone, Serialize)] -pub(crate) struct AnthropicRequest<'a> { +pub struct AnthropicRequest<'a> { pub model: &'a str, pub messages: Vec>, pub max_tokens: u32, @@ -71,7 +71,7 @@ pub(crate) struct AnthropicRequest<'a> { } #[derive(Debug, Clone, Serialize)] -pub(crate) struct AnthropicMessage<'a> { +pub struct AnthropicMessage<'a> { pub role: &'a str, /// Polymorphic content blocks — text and `tool_result` blocks /// emit different shapes per @@ -128,7 +128,7 @@ pub enum TranslateError { /// [{type:"tool_result", tool_use_id, content}]}` shape per /// so agent-loop turn 2 /// (caller sends the tool's output back to the model) round-trips. -pub(crate) fn split_system<'a>( +pub fn split_system<'a>( req: &'a ChatFormat, ) -> Result<(Option, Vec>), TranslateError> { let mut system_parts: Vec<&'a str> = Vec::new(); @@ -174,7 +174,7 @@ pub(crate) fn split_system<'a>( Ok((system, messages)) } -pub(crate) fn build_request<'a>( +pub fn build_request<'a>( req: &'a ChatFormat, upstream_model: &'a str, system: Option, @@ -224,7 +224,7 @@ pub(crate) fn build_request<'a>( /// when the input isn't an array or when no entries translated — /// keeping the field absent from the upstream wire shape so /// Anthropic doesn't reject for empty-tools. -pub(crate) fn translate_openai_tools_to_anthropic( +pub fn translate_openai_tools_to_anthropic( tools: serde_json::Value, ) -> Option> { let arr = tools.as_array()?; @@ -271,7 +271,7 @@ pub(crate) fn translate_openai_tools_to_anthropic( /// Returns None for unrecognised shapes — caller's value is discarded /// rather than forwarded verbatim, since the OpenAI shape would 400 /// the Anthropic upstream. -pub(crate) fn translate_openai_tool_choice_to_anthropic( +pub fn translate_openai_tool_choice_to_anthropic( v: serde_json::Value, ) -> Option { match v { @@ -361,7 +361,7 @@ pub fn translate_anthropic_tool_choice_to_openai( /// Non-streaming response shape from `/v1/messages`. #[derive(Debug, Deserialize)] -pub(crate) struct AnthropicResponse { +pub struct AnthropicResponse { pub id: String, pub model: String, #[serde(default)] @@ -374,7 +374,7 @@ pub(crate) struct AnthropicResponse { #[derive(Debug, Deserialize)] #[serde(tag = "type")] -pub(crate) enum AnthropicResponseBlock { +pub enum AnthropicResponseBlock { #[serde(rename = "text")] Text { text: String }, /// Anthropic's `tool_use` content block. The model is asking to @@ -401,7 +401,7 @@ pub(crate) enum AnthropicResponseBlock { } #[derive(Debug, Default, Deserialize)] -pub(crate) struct AnthropicUsage { +pub struct AnthropicUsage { pub input_tokens: u32, pub output_tokens: u32, /// Tokens written to the prompt cache (1.25× input rate). Optional @@ -413,7 +413,7 @@ pub(crate) struct AnthropicUsage { pub cache_read_input_tokens: u32, } -pub(crate) fn response_into_chat_response(raw: AnthropicResponse) -> ChatResponse { +pub fn response_into_chat_response(raw: AnthropicResponse) -> ChatResponse { let text = raw .content .iter() @@ -516,7 +516,7 @@ fn map_stop_reason(raw: Option<&str>) -> FinishReason { /// quietly dropped by the Bridge. #[derive(Debug, Deserialize)] #[serde(tag = "type")] -pub(crate) enum AnthropicStreamEvent { +pub enum AnthropicStreamEvent { #[serde(rename = "message_start")] MessageStart { message: AnthropicStreamStartMessage, @@ -538,14 +538,14 @@ pub(crate) enum AnthropicStreamEvent { } #[derive(Debug, Deserialize)] -pub(crate) struct AnthropicStreamStartMessage { +pub struct AnthropicStreamStartMessage { pub id: String, pub model: String, } #[derive(Debug, Deserialize)] #[serde(tag = "type")] -pub(crate) enum AnthropicStreamDelta { +pub enum AnthropicStreamDelta { #[serde(rename = "text_delta")] TextDelta { text: String }, #[serde(other)] @@ -553,13 +553,13 @@ pub(crate) enum AnthropicStreamDelta { } #[derive(Debug, Deserialize)] -pub(crate) struct AnthropicStreamMessageDelta { +pub struct AnthropicStreamMessageDelta { #[serde(default)] pub stop_reason: Option, } #[derive(Debug, Deserialize)] -pub(crate) struct AnthropicStreamUsage { +pub struct AnthropicStreamUsage { #[serde(default)] pub output_tokens: Option, } @@ -568,7 +568,7 @@ pub(crate) struct AnthropicStreamUsage { /// tagged with the message id/model even though only the first event /// carries them. #[derive(Debug, Default)] -pub(crate) struct StreamState { +pub struct StreamState { pub id: String, pub model: String, } diff --git a/crates/aisix-provider-bedrock/Cargo.toml b/crates/aisix-provider-bedrock/Cargo.toml index ab8aa4e1..6b06fe34 100644 --- a/crates/aisix-provider-bedrock/Cargo.toml +++ b/crates/aisix-provider-bedrock/Cargo.toml @@ -11,10 +11,28 @@ description = "aisix: AWS Bedrock runtime provider bridge (skeleton — multi-pu [dependencies] 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. +aisix-provider-anthropic = { path = "../aisix-provider-anthropic" } async-trait.workspace = true thiserror.workspace = true tracing.workspace = true +serde.workspace = true +serde_json.workspace = true +tokio.workspace = true +# AWS SDK: handles SigV4 + retries + binary event-stream framing. +# Same workspace deps the guardrails crate uses for ApplyGuardrail. +aws-config.workspace = true +aws-sdk-bedrockruntime.workspace = true +aws-credential-types.workspace = true +aws-smithy-async.workspace = true +aws-smithy-runtime-api.workspace = true +aws-smithy-types.workspace = true [dev-dependencies] tokio = { workspace = true, features = ["macros", "rt", "time"] } -serde_json.workspace = true +wiremock.workspace = true +http.workspace = true diff --git a/crates/aisix-provider-bedrock/src/bridge.rs b/crates/aisix-provider-bedrock/src/bridge.rs index 3725d637..21ab07f0 100644 --- a/crates/aisix-provider-bedrock/src/bridge.rs +++ b/crates/aisix-provider-bedrock/src/bridge.rs @@ -1,34 +1,80 @@ //! `BedrockBridge` — family Bridge for [`Adapter::Bedrock`]. //! -//! Skeleton: structure + publisher resolution + Hub-registrable -//! shell. Actual SigV4 + per-publisher dispatch lands in follow-up -//! PRs (see crate-level docs). +//! Multi-publisher dispatch backed by `aws-sdk-bedrockruntime`. The +//! SDK handles SigV4 signing, retries, and (for the streaming +//! follow-up D7.2.b) the binary event-stream framing. Per-publisher +//! request bodies + response decoding live in this crate. +//! +//! **Currently wired:** `anthropic.*` (Claude on Bedrock) chat. Other +//! publishers + streaming surface clear `not yet implemented` errors +//! referencing D7.x follow-ups — see crate-level docs. +//! +//! Credentials: `ProviderKey.secret` is a JSON-encoded +//! `{access_key_id, secret_access_key, session_token?, region}` +//! struct. The bridge parses it per request (cheap — strings only) +//! and constructs a per-call SDK client. `ProviderKey.api_base` (if +//! set) is forwarded as the SDK's `endpoint_url` so operators can +//! point at a private deployment / VPC endpoint. use aisix_gateway::{ Bridge, BridgeContext, BridgeError, ChatChunkStream, ChatFormat, ChatResponse, }; 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::invoke_model::InvokeModelError; +use aws_sdk_bedrockruntime::primitives::Blob; +use aws_sdk_bedrockruntime::Client as BedrockClient; +use aws_smithy_runtime_api::client::result::ServiceError; +use serde::Deserialize; + +use aisix_provider_anthropic::wire::{ + build_request, response_into_chat_response, split_system, AnthropicResponse, +}; 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. +const BEDROCK_ANTHROPIC_VERSION: &str = "bedrock-2023-05-31"; + /// Family Bridge for AWS Bedrock Runtime. -/// -/// **Skeleton:** compiles, registers, surfaces a clear -/// `BridgeError::Config` on every call. Real SigV4-signed dispatch -/// and per-publisher request building are wired in follow-up PRs — -/// see [`crate`] docs. pub struct BedrockBridge { /// Static `name()` returned to the Hub. Stable across upgrades so /// metrics dashboards keep their existing `provider="bedrock"` /// filters working. name: &'static str, + /// Test-only endpoint URL override. When set, the SDK config's + /// `endpoint_url` is pinned to this value so wiremock can stand + /// in for `bedrock-runtime..amazonaws.com`. Credentials, + /// region, and SigV4 signing still run normally. + #[cfg(test)] + endpoint_url_override: Option, } impl BedrockBridge { /// Construct a Bedrock bridge with the canonical name /// `"bedrock"`. Matches the Adapter enum's wire form. pub fn new() -> Self { - Self { name: "bedrock" } + Self { + name: "bedrock", + #[cfg(test)] + endpoint_url_override: None, + } + } + + /// Test-only seam: rewrite the SDK's endpoint URL so wiremock can + /// stand in for AWS. Credentials / region / SigV4 paths all run + /// normally; only the host is different. + #[cfg(test)] + pub(crate) fn with_endpoint_override(mut self, url: impl Into) -> Self { + self.endpoint_url_override = Some(url.into()); + self } } @@ -45,32 +91,30 @@ impl Default for BedrockBridge { /// /// New publishers MUST be handled in [`BedrockPublisher::from_model_id`] /// and the per-publisher request builder match in `chat` / -/// `chat_stream` (once dispatch lands). +/// `chat_stream`. /// /// Source: AWS Bedrock model catalog -/// -/// cross-referenced with LiteLLM `bedrock/`. +/// . /// /// **MVP coverage** (the variants with per-publisher dispatch already /// planned in D7.2 / D7.3 / D7.4): /// -/// - [`Self::Anthropic`] — `anthropic.claude-*` -/// - [`Self::Meta`] — `meta.llama*` -/// - [`Self::Mistral`] — `mistral.*` -/// - [`Self::AmazonTitan`] — `amazon.titan-*` -/// - [`Self::AmazonNova`] — `amazon.nova-*` -/// - [`Self::Cohere`] — `cohere.command*` -/// - [`Self::Ai21`] — `ai21.jamba-*` +/// - [`Self::Anthropic`] — `anthropic.claude-*` (wired in this PR) +/// - [`Self::Meta`] — `meta.llama*` (D7.3) +/// - [`Self::Mistral`] — `mistral.*` (D7.4) +/// - [`Self::AmazonTitan`] — `amazon.titan-*` (D7.4) +/// - [`Self::AmazonNova`] — `amazon.nova-*` (D7.4) +/// - [`Self::Cohere`] — `cohere.command*` (D7.4) +/// - [`Self::Ai21`] — `ai21.jamba-*` (D7.4) /// /// **Catch-all** ([`Self::Other`]) — every other Bedrock publisher /// AWS hosts but we haven't pinned wire-shape dispatch for yet: /// DeepSeek, Writer (Palmyra), Stability AI, Google (Gemma on /// Bedrock), NVIDIA, Qwen, Moonshot AI, MiniMax, Z.AI, TwelveLabs, -/// OpenAI (gpt-oss on Bedrock). The resolver returns `Other` for -/// these so a customer registering e.g. `deepseek.r1-v1:0` doesn't -/// get a confusing "publisher unknown" at registration time — the -/// bridge knows it's a Bedrock id, dispatch just isn't wired yet. -/// Once D7.x lands per-publisher dispatch, `Other` shrinks. +/// OpenAI (gpt-oss on Bedrock). Resolver returns `Other` for these +/// so a customer registering e.g. `deepseek.r1-v1:0` doesn't get a +/// confusing "publisher unknown" at registration time — the bridge +/// knows it's a Bedrock id, dispatch just isn't wired yet. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum BedrockPublisher { /// `anthropic.claude-*` — Claude on Bedrock. Wire shape is @@ -85,8 +129,8 @@ pub enum BedrockPublisher { /// `amazon.titan-*` — Titan Text / Embed. Uses /// `inputText + textGenerationConfig` body shape. AmazonTitan, - /// `amazon.nova-*` — Nova Pro / Nova Lite / Nova Micro (2024 Q4). - /// Uses Converse API natively. + /// `amazon.nova-*` — Nova Pro / Nova Lite / Nova Micro. Uses + /// Converse API natively. AmazonNova, /// `cohere.command-*` — Cohere Command R / R+ on Bedrock. Cohere, @@ -97,19 +141,12 @@ pub enum BedrockPublisher { /// Google Gemma, NVIDIA, Qwen, Moonshot AI, MiniMax, Z.AI, /// TwelveLabs, OpenAI gpt-oss. `chat()` returns /// `BridgeError::Config("not yet implemented")` referencing - /// #302 Phase G follow-ups; future D7.x PRs add variants - /// per concrete publisher as wire-shape dispatch lands. + /// #302 Phase G follow-ups. Other, } -/// Publisher tags recognized by [`BedrockPublisher::from_model_id`] -/// as second-segment (or first-after-region) Bedrock-catalog -/// identifiers. Used by both the publisher resolver AND -/// [`strip_region_prefix`] (which only strips a region if the -/// segment after is a known publisher — otherwise an actual -/// publisher like `amazon.titan-...` would lose its `amazon.` segment). -/// -/// Source: AWS Bedrock catalog enumeration as of 2026-05. +/// Publisher tags recognized as second-segment (or first-after-region) +/// Bedrock-catalog identifiers. const KNOWN_PUBLISHER_TAGS: &[&str] = &[ // MVP publishers (per-publisher dispatch planned in D7.2/3/4) "anthropic", @@ -138,20 +175,8 @@ const KNOWN_PUBLISHER_TAGS: &[&str] = &[ impl BedrockPublisher { /// Resolve the publisher from the Bedrock model id, tolerating - /// cross-region inference profile prefixes. - /// - /// Recognized publisher tags (case-insensitive): see - /// [`KNOWN_PUBLISHER_TAGS`]. Unknown tags return `None` — - /// callers surface a clear `BridgeError::Config` so the operator - /// can correct the model registration. - /// - /// **Cross-region prefix tolerance:** Bedrock supports inference - /// profiles per - /// . - /// Tolerated prefixes (per current AWS catalog as of 2026-05): - /// `us.`, `eu.`, `apac.`, `global.`, `us-gov.`. The resolver - /// strips a leading region tag before matching. New AWS geos - /// will need [`strip_region_prefix`] updated. + /// cross-region inference profile prefixes (`us.`, `eu.`, + /// `apac.`, `global.`, `us-gov.`). pub fn from_model_id(model_id: &str) -> Option { let stripped = strip_region_prefix(model_id); let (publisher_tag, _rest) = stripped.split_once('.')?; @@ -164,53 +189,49 @@ impl BedrockPublisher { "mistral" => Self::Mistral, "amazon" if body_lower.starts_with("amazon.nova-") => Self::AmazonNova, "amazon" if body_lower.starts_with("amazon.titan-") => Self::AmazonTitan, - "amazon" => Self::Other, // future amazon.* families + "amazon" => Self::Other, "cohere" => Self::Cohere, "ai21" => Self::Ai21, - // Catalog publishers we know exist but haven't wired - // dispatch for yet. Resolved to Other so registration - // doesn't fail — operator sees a "publisher X not yet - // implemented" at dispatch time once D7.x lands. "deepseek" | "writer" | "stability" | "google" | "nvidia" | "qwen" | "moonshotai" | "moonshot" | "minimaxai" | "minimax" | "zai-org" | "zai" | "twelvelabs" | "openai" => Self::Other, _ => return None, }) } + + /// Human-readable name used in publisher-not-implemented errors. + fn name(&self) -> &'static str { + match self { + Self::Anthropic => "anthropic", + Self::Meta => "meta", + Self::Mistral => "mistral", + Self::AmazonTitan => "amazon.titan", + Self::AmazonNova => "amazon.nova", + Self::Cohere => "cohere", + Self::Ai21 => "ai21", + Self::Other => "", + } + } } /// Strip a leading cross-region inference profile prefix. /// -/// Returns the input unchanged when no recognized prefix matches. -/// /// Recognized prefixes (per AWS catalog as of 2026-05): /// `us.`, `eu.`, `apac.`, `global.`, `us-gov.`. -/// -/// The criterion: the leading segment must be 2–7 ASCII -/// lowercase letters / digits / hyphens, AND the segment after the -/// `.` separator must itself start with a known publisher tag. -/// Otherwise an actual publisher like `amazon.titan-...` would -/// lose its `amazon.` segment. fn strip_region_prefix(model_id: &str) -> &str { let Some((maybe_region, rest)) = model_id.split_once('.') else { return model_id; }; let len = maybe_region.len(); - // 2-7 covers `us` through `us-gov` (6) and gives a 1-char - // safety margin without admitting random long prefixes. if !(2..=7).contains(&len) { return model_id; } - // Hyphens are valid in `us-gov`. We deliberately do NOT accept - // dots or other URL-control chars — that would let an attacker - // smuggle path segments via the model id. if !maybe_region .chars() .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-') { return model_id; } - // Only strip if what follows looks like a known publisher tag. let next_tag = rest.split('.').next().unwrap_or("").to_ascii_lowercase(); if KNOWN_PUBLISHER_TAGS.contains(&next_tag.as_str()) { rest @@ -219,6 +240,141 @@ fn strip_region_prefix(model_id: &str) -> &str { } } +/// Schema for `ProviderKey.secret` on a Bedrock provider key. +/// +/// Convention: AWS credentials are JSON-encoded into the `secret` +/// field. The cp-api side delivers them already-decrypted (mTLS-only +/// etcd channel; see ProviderKey doc). +/// +/// `endpoint_url` is intentionally NOT in here — that goes in +/// `ProviderKey.api_base` so the cp-api validator can apply normal +/// URL-shape rules. Region is in here because Bedrock keys dispatch +/// off region (`bedrock-runtime..amazonaws.com`). +#[derive(Debug, Deserialize)] +struct BedrockSecret { + access_key_id: String, + secret_access_key: String, + /// AWS STS session token. Optional — long-lived static keys + /// don't have one; assume-role credentials do. + #[serde(default)] + session_token: Option, + /// AWS region the Bedrock dispatch targets (e.g. `us-west-2`). + /// Required — Bedrock's URL is region-keyed and the SDK won't + /// dispatch without it. + region: String, +} + +impl BedrockSecret { + /// Parse the JSON-encoded credential blob. Audit M1: error + /// messages here MUST NOT echo the raw secret content — only + /// generic shape errors. + fn parse(secret: &str) -> Result { + if secret.trim().is_empty() { + return Err(BridgeError::Config( + "bedrock provider_key.secret is empty — \ + expected JSON {access_key_id, secret_access_key, region, session_token?}" + .into(), + )); + } + serde_json::from_str::(secret).map_err(|_e| { + // Intentionally do NOT include the underlying serde error + // message — it can leak partial secret contents (e.g. + // "invalid character 'X' at position N" reveals what's + // in the JSON). Generic shape hint is enough for the + // operator who controls the registration. + BridgeError::Config( + "bedrock provider_key.secret must be valid JSON: \ + {access_key_id, secret_access_key, region, session_token?}" + .into(), + ) + }) + } +} + +/// Build a Bedrock SDK Client from the parsed credentials plus the +/// optional endpoint override. +fn build_client( + creds: &BedrockSecret, + endpoint_url: Option<&str>, +) -> Result { + if creds.region.trim().is_empty() { + return Err(BridgeError::Config( + "bedrock provider_key.secret.region is empty — \ + AWS Bedrock dispatch is region-keyed and requires e.g. \"us-west-2\"" + .into(), + )); + } + let aws_creds = Credentials::new( + creds.access_key_id.clone(), + creds.secret_access_key.clone(), + creds.session_token.clone(), + None, + "aisix-provider-bedrock", + ); + let mut builder = aws_config::SdkConfig::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new(creds.region.clone())) + .credentials_provider(SharedCredentialsProvider::new(aws_creds)) + .sleep_impl(aws_smithy_async::rt::sleep::SharedAsyncSleep::new( + aws_smithy_async::rt::sleep::TokioSleep::new(), + )); + if let Some(url) = endpoint_url { + builder = builder.endpoint_url(url); + } + let sdk_cfg = builder.build(); + Ok(BedrockClient::new(&sdk_cfg)) +} + +/// Pull the upstream model id off the BridgeContext. +fn upstream_model(ctx: &BridgeContext) -> Result<&str, BridgeError> { + ctx.model + .model_name + .as_deref() + .ok_or_else(|| BridgeError::Config("model.model_name missing".into())) +} + +/// Translate an SDK error into the canonical `BridgeError`. +/// +/// **Audit M1 — sensitive-info redaction:** Bedrock error envelopes +/// frequently include the operator's model id, region, account +/// numbers (in ARNs), and IAM role names. Surfacing these to a +/// downstream customer leaks operator-internal taxonomy. We map to +/// canned status-keyed phrases. +fn map_sdk_error(err: SdkError) -> BridgeError { + match err { + SdkError::TimeoutError(_) => BridgeError::Timeout { elapsed_ms: 0 }, + SdkError::DispatchFailure(_) => BridgeError::Transport("upstream dispatch failed".into()), + SdkError::ConstructionFailure(_) => { + BridgeError::Config("upstream request construction failed".into()) + } + SdkError::ResponseError(_) => { + BridgeError::UpstreamDecode("upstream response could not be parsed".into()) + } + SdkError::ServiceError(svc) => map_service_error(svc), + _ => BridgeError::Transport("upstream dispatch failed".into()), + } +} + +fn map_service_error( + svc: ServiceError, +) -> BridgeError { + let raw = svc.into_raw(); + let status = raw.status().as_u16(); + 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(), + 500..=599 => format!("upstream returned {status}"), + _ => format!("upstream returned {status}"), + }; + BridgeError::UpstreamStatus { + status, + message, + retry_after: None, + } +} + #[async_trait] impl Bridge for BedrockBridge { fn name(&self) -> &'static str { @@ -227,21 +383,11 @@ impl Bridge for BedrockBridge { async fn chat( &self, - _req: &ChatFormat, + req: &ChatFormat, ctx: &BridgeContext, ) -> Result { - // Skeleton: validate the publisher resolution path so a - // misconfigured model id surfaces a clear error today, even - // though the actual SigV4-signed call is TODO. - // - // IMPORTANT: the Bedrock model id is on Model.model_name (the - // upstream id the operator pinned when registering the model), - // NOT on req.model (which is the gateway-internal display - // name the customer typed in `/v1/chat/completions`). See - // OpenAiBridge / `upstream_model(ctx)` for the established - // pattern. let upstream_id = upstream_model(ctx)?; - let _publisher = BedrockPublisher::from_model_id(upstream_id).ok_or_else(|| { + let publisher = BedrockPublisher::from_model_id(upstream_id).ok_or_else(|| { BridgeError::Config(format!( "bedrock publisher unknown for model id {upstream_id:?}; \ expected one of anthropic.claude-* / meta.llama* / mistral.* / \ @@ -249,15 +395,19 @@ impl Bridge for BedrockBridge { (optionally prefixed with a cross-region inference profile like us. / eu. / apac.)" )) })?; - // Reserved-config helpers exercised by tests: keep wire module - // reachable from the public surface so a future dispatch PR - // can drop its body straight in. + // Keep wire module reachable from the public surface so the + // streaming follow-up can wire SigV4-reserved-header checks + // for any operator default_headers override. let _ = wire::reserved_sigv4_headers(); - Err(BridgeError::Config( - "bedrock bridge is not yet implemented — \ - tracked under api7/AISIX-Cloud#302 Phase G (D7)" - .into(), - )) + + match publisher { + BedrockPublisher::Anthropic => self.chat_anthropic(req, ctx, upstream_id).await, + other => Err(BridgeError::Config(format!( + "bedrock publisher {publisher:?} not yet implemented — \ + tracked under api7/AISIX-Cloud#302 Phase G (D7.3+, publisher={})", + other.name() + ))), + } } async fn chat_stream( @@ -275,28 +425,90 @@ impl Bridge for BedrockBridge { )) })?; Err(BridgeError::Config( - "bedrock bridge is not yet implemented — \ - tracked under api7/AISIX-Cloud#302 Phase G (D7)" + "bedrock streaming is not yet implemented — \ + tracked under api7/AISIX-Cloud#302 Phase G (D7.2.b)" .into(), )) } } -/// Pull the upstream model id off the BridgeContext. Mirrors -/// OpenAiBridge's same-named helper — Bedrock model ids -/// (`anthropic.claude-...`, `meta.llama...`) live on Model.model_name, -/// not on the customer-facing display name in req.model. -fn upstream_model(ctx: &BridgeContext) -> Result<&str, BridgeError> { - ctx.model - .model_name - .as_deref() - .ok_or_else(|| BridgeError::Config("model.model_name missing".into())) +impl BedrockBridge { + /// Dispatch Anthropic-on-Bedrock chat. Body shape per + /// : + /// the Anthropic Messages JSON minus the `model` field (Bedrock + /// keys dispatch off the URL) plus `anthropic_version: + /// "bedrock-2023-05-31"`. + 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)?; + + // 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); + let mut body_value = serde_json::to_value(&anthropic_req) + .map_err(|e| BridgeError::Config(format!("serialize Anthropic request body: {e}")))?; + if let Some(obj) = body_value.as_object_mut() { + obj.remove("model"); + obj.remove("stream"); + obj.insert( + "anthropic_version".to_string(), + serde_json::Value::String(BEDROCK_ANTHROPIC_VERSION.to_string()), + ); + } + let body_bytes = serde_json::to_vec(&body_value).map_err(|e| { + 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 resp = client + .invoke_model() + .model_id(upstream_id) + .content_type("application/json") + .accept("application/json") + .body(Blob::new(body_bytes)) + .send() + .await + .map_err(map_sdk_error)?; + + let parsed: AnthropicResponse = serde_json::from_slice(resp.body().as_ref()) + .map_err(|e| BridgeError::UpstreamDecode(e.to_string()))?; + Ok(response_into_chat_response(parsed)) + } } #[cfg(test)] mod tests { use super::*; + // ─── Publisher resolution (preserved from skeleton) ─────────────── + #[test] fn publisher_resolves_anthropic_claude_on_bedrock() { assert_eq!( @@ -307,11 +519,6 @@ mod tests { BedrockPublisher::from_model_id("anthropic.claude-3-haiku-20240307-v1:0"), Some(BedrockPublisher::Anthropic), ); - // Matcher is on the `anthropic` tag, not `anthropic.claude`, - // so any future non-Claude family AWS hosts under the - // `anthropic.` namespace also resolves correctly. The - // bridge resolves the publisher, the dispatch arm decides - // which wire shape to use. assert_eq!( BedrockPublisher::from_model_id("anthropic.opus-4-1-20250805-v1:0"), Some(BedrockPublisher::Anthropic), @@ -320,8 +527,6 @@ mod tests { #[test] fn publisher_resolves_meta_llama_variants() { - // Bedrock's Llama wire form is `meta.llama3-X-...` — - // single hyphen between `llama` and the version digit. assert_eq!( BedrockPublisher::from_model_id("meta.llama3-3-70b-instruct-v1:0"), Some(BedrockPublisher::Meta), @@ -346,27 +551,14 @@ mod tests { #[test] fn publisher_resolves_amazon_titan_and_nova_distinctly() { - // Tight ordering pin: nova must resolve to AmazonNova, - // titan to AmazonTitan. A future refactor that collapses - // both to a single Amazon variant would lose the wire- - // shape distinction (Nova uses Converse, Titan uses the - // legacy inputText shape). assert_eq!( BedrockPublisher::from_model_id("amazon.nova-pro-v1:0"), Some(BedrockPublisher::AmazonNova), ); - assert_eq!( - BedrockPublisher::from_model_id("amazon.nova-lite-v1:0"), - Some(BedrockPublisher::AmazonNova), - ); assert_eq!( BedrockPublisher::from_model_id("amazon.titan-text-premier-v1:0"), Some(BedrockPublisher::AmazonTitan), ); - assert_eq!( - BedrockPublisher::from_model_id("amazon.titan-text-express-v1"), - Some(BedrockPublisher::AmazonTitan), - ); } #[test] @@ -375,10 +567,6 @@ mod tests { BedrockPublisher::from_model_id("cohere.command-r-plus-v1:0"), Some(BedrockPublisher::Cohere), ); - assert_eq!( - BedrockPublisher::from_model_id("cohere.command-r-v1:0"), - Some(BedrockPublisher::Cohere), - ); } #[test] @@ -391,10 +579,6 @@ mod tests { #[test] fn publisher_strips_cross_region_us_prefix() { - // `us.anthropic.claude-...` must resolve the same as the - // non-prefixed form. The cross-region inference profile is - // a routing detail — the publisher's wire shape is - // identical regardless. assert_eq!( BedrockPublisher::from_model_id("us.anthropic.claude-3-5-sonnet-20241022-v2:0"), Some(BedrockPublisher::Anthropic), @@ -409,12 +593,6 @@ mod tests { ); } - /// D7 audit HIGH-2 regression: `global.` and `us-gov.` are live - /// Bedrock cross-region inference profile prefixes. The earlier - /// matcher's "2-6 lowercase letters/digits only" rule silently - /// failed on `us-gov.` (hyphen) and was lucky on `global.` - /// (exactly 6 chars). A real GovCloud customer would have hit - /// publisher-unknown with no actionable error. #[test] fn publisher_strips_global_and_us_gov_prefixes() { assert_eq!( @@ -427,14 +605,6 @@ mod tests { ); } - /// D7 audit HIGH-1 regression: catalog publishers AWS hosts on - /// Bedrock today that we don't have wire-shape dispatch for yet. - /// Resolving these to `Other` (rather than `None`) means a - /// customer registering `deepseek.r1-v1:0` or - /// `writer.palmyra-x5-v1:0` doesn't get a confusing "publisher - /// unknown" at registration time — the bridge knows it's a - /// Bedrock id, dispatch just isn't wired. Once D7.x lands per- - /// publisher dispatch, the enum gets per-publisher variants. #[test] fn publisher_resolves_catalog_others_to_other_variant() { assert_eq!( @@ -445,11 +615,6 @@ mod tests { BedrockPublisher::from_model_id("writer.palmyra-x5-v1:0"), Some(BedrockPublisher::Other), ); - assert_eq!( - BedrockPublisher::from_model_id("stability.sd3-large-v1:0"), - Some(BedrockPublisher::Other), - ); - // Cross-region prefix also works for Other-variant publishers. assert_eq!( BedrockPublisher::from_model_id("us.deepseek.r1-v1:0"), Some(BedrockPublisher::Other), @@ -458,16 +623,10 @@ mod tests { #[test] fn publisher_does_not_strip_publisher_segment_as_region() { - // Guard the strip_region_prefix logic: `amazon.titan-...` - // must NOT have its `amazon.` segment treated as a region - // prefix. If it did, the rest would be `titan-...` which - // doesn't start with `amazon.titan-`, so we'd lose the - // publisher entirely. assert_eq!( BedrockPublisher::from_model_id("amazon.titan-text-premier-v1:0"), Some(BedrockPublisher::AmazonTitan), ); - // Same guard for `cohere.command-*`: assert_eq!( BedrockPublisher::from_model_id("cohere.command-r-v1:0"), Some(BedrockPublisher::Cohere), @@ -476,14 +635,8 @@ mod tests { #[test] fn publisher_unknown_id_returns_none() { - // `gpt-4o` has no `.` separator → no publisher tag. assert_eq!(BedrockPublisher::from_model_id("gpt-4o"), None); - // Empty string → no tag. assert_eq!(BedrockPublisher::from_model_id(""), None); - // `truly-unknown.foo-v1:0` has a tag but it's not in - // KNOWN_PUBLISHER_TAGS — explicit "this is not a Bedrock - // catalog id at all" signal so the caller can surface a - // clear configuration error. assert_eq!( BedrockPublisher::from_model_id("truly-unknown.foo-v1:0"), None, @@ -492,23 +645,102 @@ mod tests { #[test] fn bridge_name_is_stable() { - // Metrics label is part of the public contract — a rename - // would silently break customer dashboards. assert_eq!(BedrockBridge::new().name(), "bedrock"); } + // ─── BedrockSecret parsing ──────────────────────────────────────── + + #[test] + fn bedrock_secret_parses_full_form() { + let json = + r#"{"access_key_id":"AKIA-test","secret_access_key":"sk-test","region":"us-west-2"}"#; + let s = BedrockSecret::parse(json).unwrap(); + assert_eq!(s.access_key_id, "AKIA-test"); + assert_eq!(s.secret_access_key, "sk-test"); + assert_eq!(s.region, "us-west-2"); + assert!(s.session_token.is_none()); + } + + #[test] + fn bedrock_secret_parses_with_session_token() { + let json = r#"{"access_key_id":"AKIA","secret_access_key":"sk","region":"us-west-2","session_token":"AQo..."}"#; + let s = BedrockSecret::parse(json).unwrap(); + assert_eq!(s.session_token.as_deref(), Some("AQo...")); + } + + #[test] + fn bedrock_secret_rejects_empty() { + let err = BedrockSecret::parse("").unwrap_err(); + match err { + BridgeError::Config(msg) => { + assert!( + msg.contains("secret is empty"), + "must mention empty secret; got {msg}" + ); + assert!( + msg.contains("access_key_id"), + "must hint at required JSON shape; got {msg}" + ); + } + other => panic!("expected Config error, got {other:?}"), + } + } + + #[test] + fn bedrock_secret_rejects_non_json() { + let err = BedrockSecret::parse("AKIA-test").unwrap_err(); + match err { + BridgeError::Config(msg) => { + assert!( + msg.contains("must be valid JSON"), + "must mention JSON requirement; got {msg}" + ); + } + other => panic!("expected Config error, got {other:?}"), + } + } + + /// Audit M1: the error path must not echo the raw secret content + /// — serde error messages include "invalid character X at + /// position N" which reveals partial secret bytes. + #[test] + fn bedrock_secret_error_does_not_leak_secret_content() { + let secret_with_distinctive_bytes = "X-DISTINCTIVE-LEAK-MARKER-Y"; + let err = BedrockSecret::parse(secret_with_distinctive_bytes).unwrap_err(); + match err { + BridgeError::Config(msg) => { + assert!( + !msg.contains("X-DISTINCTIVE-LEAK-MARKER-Y"), + "error must NOT echo raw secret bytes; got {msg}" + ); + assert!( + !msg.contains("DISTINCTIVE"), + "error must NOT leak partial secret bytes; got {msg}" + ); + } + other => panic!("expected Config error, got {other:?}"), + } + } + + #[test] + fn bedrock_secret_rejects_missing_region() { + // serde rejects missing required field — bridge surfaces + // the generic shape-error, not the field name (defense in + // depth against accidental field-name leakage to customer + // error path; the operator-side schema docs say what's + // required). + let json = r#"{"access_key_id":"AKIA","secret_access_key":"sk"}"#; + let err = BedrockSecret::parse(json).unwrap_err(); + assert!(matches!(err, BridgeError::Config(_))); + } + + // ─── Pre-dispatch validation tests ───────────────────────────────── + use aisix_core::{Model, ProviderKey}; use aisix_gateway::ChatMessage; use std::sync::Arc; - /// Build a Model fixture where `model_name` (the upstream id) and - /// `display_name` (the customer-facing name) deliberately differ. - /// The bridge must dispatch off `model_name`, not the typed - /// display name in `req.model` — pinning that contract here. fn sample_model_with(model_name: &str) -> Arc { - // Note: Model.provider uses the legacy 6-value Provider enum. - // amazon-bedrock isn't a Provider variant; the Adapter::Bedrock - // routing happens off ProviderKey.adapter, not Model.provider. let cfg = format!( r#"{{ "display_name": "customer-facing-name", @@ -520,95 +752,97 @@ mod tests { Arc::new(serde_json::from_str(&cfg).unwrap()) } - fn sample_pk() -> Arc { - Arc::new( - serde_json::from_str(r#"{"display_name": "bedrock-prod", "secret": "AKIA-test"}"#) - .unwrap(), - ) + /// Build a PK with a valid Bedrock-shape secret. `endpoint_url` + /// arg is the test-only override path — set this to a wiremock + /// URI to drive `bridge.chat()` end-to-end. + fn sample_pk_with_secret(secret_json: &str) -> Arc { + let cfg = format!( + r#"{{"display_name": "bedrock-prod", "secret": {}}}"#, + serde_json::to_string(secret_json).unwrap() + ); + Arc::new(serde_json::from_str(&cfg).unwrap()) + } + + fn valid_secret_json() -> &'static str { + r#"{"access_key_id":"AKIA-test","secret_access_key":"sk-test","region":"us-west-2"}"# } #[tokio::test] - async fn chat_surfaces_clear_not_implemented_error() { + async fn chat_rejects_unknown_publisher() { let bridge = BedrockBridge::new(); let ctx = BridgeContext::new( "req-1", - sample_model_with("anthropic.claude-3-5-sonnet-20241022-v2:0"), - sample_pk(), + sample_model_with("totally-bogus-model-id"), + sample_pk_with_secret(valid_secret_json()), ); - // The customer-typed model name (req.model) is the display - // name, NOT the Bedrock upstream id. The bridge must ignore - // it and resolve off ctx.model.model_name instead. 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("bedrock bridge is not yet implemented"), - "error message must call out the WIP status; got {msg}" - ); - assert!( - msg.contains("#302"), - "error message must link to the tracking issue; got {msg}" - ); + assert!(msg.contains("bedrock publisher unknown")); + assert!(msg.contains("totally-bogus-model-id")); } other => panic!("expected Config error, got {other:?}"), } } #[tokio::test] - async fn chat_ignores_req_model_and_uses_ctx_model_name() { - // Regression test for D6 audit HIGH-1 (also caught here): - // the bridge dispatches off Model.model_name (operator- - // pinned upstream id), not req.model (customer-typed - // display name). If a future refactor accidentally swaps - // them, this test fails — `req.model = "gpt-4o"` would - // surface "publisher unknown for gpt-4o", but with - // model_name pointing at a real Bedrock id we expect the - // skeleton's not-implemented error instead. + 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("anthropic.claude-3-haiku-20240307-v1:0"), - sample_pk(), + sample_model_with("meta.llama3-3-70b-instruct-v1:0"), + sample_pk_with_secret(valid_secret_json()), ); - // req.model deliberately set to something the publisher - // resolver would reject if it were the source of truth. - let req = ChatFormat::new("gpt-4o", vec![ChatMessage::user("hi")]); + 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("not yet implemented"), - "must hit the not-implemented stub (proving model_name was used), not the publisher-resolution guard; got {msg}" + 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:?}"), } } #[tokio::test] - async fn chat_with_unknown_model_id_errors_before_dispatch() { - // Publisher-resolution guard fires when model_name is - // unrecognized — proves the bridge rejects malformed - // registrations early. + async fn chat_with_invalid_secret_errors_before_dispatch() { let bridge = BedrockBridge::new(); let ctx = BridgeContext::new( "req-1", - sample_model_with("totally-bogus-model-id"), - sample_pk(), + sample_model_with("anthropic.claude-3-5-sonnet-20241022-v2:0"), + sample_pk_with_secret("not-valid-json"), ); - let req = ChatFormat::new("customer-facing-name", vec![ChatMessage::user("hi")]); + 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("bedrock publisher unknown"), - "must mention publisher resolution failure; got {msg}" - ); - assert!( - msg.contains("totally-bogus-model-id"), - "must include the offending model id; got {msg}" - ); + assert!(msg.contains("must be valid JSON")); + } + other => panic!("expected Config error, got {other:?}"), + } + } + + #[tokio::test] + async fn chat_with_empty_secret_errors_before_dispatch() { + 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(""), + ); + 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("secret is empty")); } other => panic!("expected Config error, got {other:?}"), } @@ -616,12 +850,8 @@ mod tests { #[tokio::test] async fn chat_with_missing_model_name_errors_before_dispatch() { - // Defense: if Model.model_name is absent (shouldn't happen - // in practice — cp-api requires it — but the field is - // Option), the bridge surfaces a clear error - // rather than panicking or treating "" as a publisher. let bridge = BedrockBridge::new(); - let pk = sample_pk(); + let pk = sample_pk_with_secret(valid_secret_json()); let model_no_name: Arc = Arc::new( serde_json::from_str( r#"{ @@ -637,9 +867,432 @@ mod tests { let err = bridge.chat(&req, &ctx).await.unwrap_err(); match err { BridgeError::Config(msg) => { - assert!(msg.contains("model_name missing"), "got {msg}"); + assert!(msg.contains("model_name missing")); + } + other => panic!("expected Config error, got {other:?}"), + } + } + + #[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. + 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()), + ); + // 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 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}" + ); + } + 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:?}"), } } + + // ─── Dispatch end-to-end against wiremock via endpoint_url override ── + + use wiremock::matchers::{method, path_regex}; + use wiremock::{Mock, MockServer, Request as MockRequest, Respond, ResponseTemplate}; + + // Audit lesson from D6 PR #319: drive the **real** + // `bridge.chat()` entry point via the `endpoint_url_override` + // seam — credentials, region, SigV4 signing, body shaping all + // run normally; only the destination host is rewritten to + // wiremock. + + /// Recording responder: captures request body + headers so tests + /// can assert what reached the wire. Always returns the canned + /// default response — tests that need a custom response use the + /// standard `ResponseTemplate` arg to `Mock::given(...).respond_with(...)` + /// without capture (no need for both modes in one helper). + #[derive(Clone, Default)] + struct CapturingResponder { + captured_body: std::sync::Arc>>, + captured_headers: std::sync::Arc>>, + } + + impl Respond for CapturingResponder { + fn respond(&self, req: &MockRequest) -> ResponseTemplate { + let body: serde_json::Value = serde_json::from_slice(&req.body).unwrap_or_default(); + *self.captured_body.lock().unwrap() = Some(body); + *self.captured_headers.lock().unwrap() = Some(req.headers.clone()); + default_anthropic_response_template() + } + } + + fn default_anthropic_response_template() -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "msg_01", + "model": "claude-3-5-sonnet-20241022-v2", + "content": [{"type": "text", "text": "hello from bedrock"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 5, "output_tokens": 4} + })) + } + + #[tokio::test] + async fn chat_anthropic_dispatches_via_invoke_model_url() { + let server = MockServer::start().await; + let responder = CapturingResponder::default(); + // Bedrock's InvokeModel URL: `/model//invoke`. + // The `:` in `anthropic.claude-3-5-sonnet-20241022-v2:0` gets + // percent-encoded to `%3A`; we use a regex to stay tolerant + // across SDK version upgrades. + Mock::given(method("POST")) + .and(path_regex( + r"^/model/anthropic\.claude-3-5-sonnet-20241022-v2(:0|%3A0)/invoke$", + )) + .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("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")]); + let chat = bridge.chat(&req, &ctx).await.unwrap(); + assert_eq!(chat.message.content, "hello from bedrock"); + assert_eq!(chat.usage.total_tokens, 9); + } + + #[tokio::test] + async fn chat_anthropic_body_contains_bedrock_anthropic_version_and_no_model_field() { + let server = MockServer::start().await; + let responder = CapturingResponder::default(); + Mock::given(method("POST")) + .and(path_regex(r"^/model/.+/invoke$")) + .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("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")]); + bridge.chat(&req, &ctx).await.unwrap(); + + let body = responder.captured_body.lock().unwrap().clone().unwrap(); + // Bedrock-Anthropic body shape pins: + // 1. `anthropic_version` MUST be present + the canonical + // `bedrock-2023-05-31` string (per AWS docs URL above). + // 2. `model` MUST be absent — Bedrock dispatches off URL path. + // 3. `stream` MUST be absent — InvokeModel is non-streaming; + // Bedrock would error on a stream:true with the wrong op. + // 4. `messages` must be the translated user turn. + assert_eq!( + body.get("anthropic_version").and_then(|v| v.as_str()), + Some("bedrock-2023-05-31"), + "body must carry anthropic_version=bedrock-2023-05-31; body={body}" + ); + assert!( + body.get("model").is_none(), + "body must NOT carry `model` (Bedrock dispatches via URL); body={body}" + ); + assert!( + body.get("stream").is_none(), + "body must NOT carry `stream` (InvokeModel is non-streaming); body={body}" + ); + let messages = body.get("messages").and_then(|v| v.as_array()).unwrap(); + assert_eq!(messages.len(), 1); + assert_eq!( + messages[0].get("role").and_then(|v| v.as_str()), + Some("user") + ); + } + + #[tokio::test] + async fn chat_anthropic_uses_sigv4_authorization_header() { + // The SDK signs with SigV4: `Authorization: AWS4-HMAC-SHA256 ...`. + // This is a wire-level pin that the SDK actually signed (vs. + // sending unauthenticated). If a future bug accidentally + // bypassed the SDK, the canned auth header would change. + let server = MockServer::start().await; + let responder = CapturingResponder::default(); + Mock::given(method("POST")) + .and(path_regex(r"^/model/.+/invoke$")) + .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("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")]); + bridge.chat(&req, &ctx).await.unwrap(); + + let headers = responder.captured_headers.lock().unwrap().clone().unwrap(); + let auth = headers + .get("authorization") + .and_then(|v: &http::HeaderValue| v.to_str().ok()) + .unwrap_or(""); + assert!( + auth.starts_with("AWS4-HMAC-SHA256"), + "expected AWS SigV4 Authorization header; got {auth:?}" + ); + // The SDK must include x-amz-date for SigV4. + assert!( + headers.contains_key("x-amz-date"), + "SigV4 requires x-amz-date; headers={headers:?}" + ); + // Body hash header should be set by the SDK. + assert!( + headers.contains_key("x-amz-content-sha256") || headers.contains_key("content-length"), + "expected x-amz-content-sha256 or content-length on a SigV4 request; got {headers:?}" + ); + } + + #[tokio::test] + async fn chat_anthropic_handles_tool_use_response_blocks() { + // Anthropic on Bedrock returns `tool_use` content blocks for + // tool-call responses. The bridge's reused + // `response_into_chat_response` must translate them to + // OpenAI's `tool_calls` shape so downstream agents work. + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex(r"^/model/.+/invoke$")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "msg_02", + "model": "claude-3-5-sonnet-20241022-v2", + "content": [ + {"type": "text", "text": "calling tool"}, + { + "type": "tool_use", + "id": "toolu_01abc", + "name": "get_weather", + "input": {"city": "SF"} + } + ], + "stop_reason": "tool_use", + "usage": {"input_tokens": 8, "output_tokens": 12} + }))) + .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")]); + let chat = bridge.chat(&req, &ctx).await.unwrap(); + assert_eq!(chat.message.content, "calling tool"); + // Tool calls translated into OpenAI shape via the reused + // anthropic crate's converter. + let tool_calls = chat + .message + .extra + .get("tool_calls") + .unwrap() + .as_array() + .unwrap(); + assert_eq!(tool_calls.len(), 1); + assert_eq!( + tool_calls[0].get("type").and_then(|v| v.as_str()), + Some("function") + ); + assert_eq!( + tool_calls[0] + .get("function") + .and_then(|f| f.get("name")) + .and_then(|v| v.as_str()), + Some("get_weather") + ); + } + + #[tokio::test] + async fn chat_maps_upstream_4xx_to_canned_message_not_body_echo() { + // Audit M1: Bedrock error envelopes can contain account + // numbers (in ARNs), model IDs, IAM role names — must not + // leak into customer-visible error. + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex(r"^/model/.+/invoke$")) + .respond_with( + ResponseTemplate::new(400).set_body_json(serde_json::json!({ + "message": "Operation cannot be performed by IAM role arn:aws:iam::123456789012:role/internal-leaky-role" + })), + ) + .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")]); + let err = bridge.chat(&req, &ctx).await.unwrap_err(); + match err { + BridgeError::UpstreamStatus { message, .. } => { + assert!( + !message.contains("123456789012") && !message.contains("internal-leaky-role"), + "upstream body must not leak account / role info into customer error; got {message:?}" + ); + } + other => panic!("expected UpstreamStatus, got {other:?}"), + } + } + + #[tokio::test] + async fn chat_maps_upstream_429_to_canned_rate_limited() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex(r"^/model/.+/invoke$")) + .respond_with(ResponseTemplate::new(429).set_body_json(serde_json::json!({ + "message": "Too many requests for account 123456789012" + }))) + .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")]); + let err = bridge.chat(&req, &ctx).await.unwrap_err(); + match err { + BridgeError::UpstreamStatus { + status, message, .. + } => { + assert_eq!(status, 429); + assert!(message.contains("rate limited")); + assert!( + !message.contains("123456789012"), + "must not leak account id; got {message:?}" + ); + } + other => panic!("expected UpstreamStatus, got {other:?}"), + } + } + + #[tokio::test] + async fn chat_with_cross_region_inference_profile_dispatches_correctly() { + // The `us.` cross-region inference profile is a real Bedrock + // routing detail — the publisher's wire shape is identical + // regardless. Critical: the URL path must include the FULL + // model id with the region prefix; only the publisher resolver + // strips it. + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex( + r"^/model/us\.anthropic\.claude-3-5-sonnet-20241022-v2(:0|%3A0)/invoke$", + )) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "msg_xr", "model": "claude-3-5-sonnet", + "content": [{"type": "text", "text": "cross-region ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1} + }))) + .expect(1) + .mount(&server) + .await; + + let bridge = BedrockBridge::new().with_endpoint_override(server.uri()); + let ctx = BridgeContext::new( + "req-1", + sample_model_with("us.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")]); + let chat = bridge.chat(&req, &ctx).await.unwrap(); + assert_eq!(chat.message.content, "cross-region ok"); + } + + #[tokio::test] + async fn chat_anthropic_translates_system_messages_to_system_field() { + // Anthropic's Messages API takes `system` as a top-level + // field, NOT a role in `messages[]`. The reused + // `split_system` helper from aisix-provider-anthropic must + // pull system turns out of the messages array into the + // top-level `system` field. + let server = MockServer::start().await; + let responder = CapturingResponder::default(); + Mock::given(method("POST")) + .and(path_regex(r"^/model/.+/invoke$")) + .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("anthropic.claude-3-5-sonnet-20241022-v2:0"), + sample_pk_with_secret(valid_secret_json()), + ); + let req = ChatFormat::new( + "my-claude", + vec![ + ChatMessage::system("you are a helpful assistant"), + ChatMessage::user("hi"), + ], + ); + bridge.chat(&req, &ctx).await.unwrap(); + + let body = responder.captured_body.lock().unwrap().clone().unwrap(); + assert_eq!( + body.get("system").and_then(|v| v.as_str()), + Some("you are a helpful assistant"), + "system role must become top-level `system` field; body={body}" + ); + let messages = body.get("messages").and_then(|v| v.as_array()).unwrap(); + assert_eq!( + messages.len(), + 1, + "system role must NOT appear in messages[]; body={body}" + ); + assert_eq!( + messages[0].get("role").and_then(|v| v.as_str()), + Some("user") + ); + } } diff --git a/crates/aisix-provider-bedrock/src/lib.rs b/crates/aisix-provider-bedrock/src/lib.rs index 30e7ed59..5adb1d23 100644 --- a/crates/aisix-provider-bedrock/src/lib.rs +++ b/crates/aisix-provider-bedrock/src/lib.rs @@ -1,29 +1,27 @@ //! aisix-provider-bedrock — AWS Bedrock runtime provider bridge. //! -//! **Skeleton crate** for issue #302 Phase G. Registers as the family -//! bridge for [`Adapter::Bedrock`] in the gateway Hub. Actual SigV4- -//! signed dispatch + per-publisher request building lands in follow-up -//! D7.x PRs. +//! Family bridge for [`Adapter::Bedrock`] in the gateway Hub. //! -//! Roadmap (tracked under issue #302 Phase G): +//! ## Status (issue #302 Phase G) //! -//! - [ ] D7.1 — AWS SigV4 v4 signature (`aws-sigv4` crate or hand-rolled) -//! over the canonical request (method + path + headers + body + region) -//! - [ ] D7.2 — Anthropic-on-Bedrock dispatch -//! (`/model/anthropic.claude-*/invoke[-with-response-stream]`, -//! `anthropic_version: "bedrock-2023-05-31"` in body not header) +//! - [x] D7.1 — AWS SigV4 v4 signature (handled by `aws-sdk-bedrockruntime`) +//! - [x] D7.2.a — Anthropic-on-Bedrock non-streaming dispatch +//! (`/model/anthropic.claude-*/invoke`, `anthropic_version: +//! "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 -//! - [ ] D7.5 — AWS event-stream framed streaming (`amazon.event-stream` -//! content-type, NOT canonical SSE) -//! - [ ] D7.6 — Cross-region inference profiles (`us.anthropic.claude-*`, -//! `eu.anthropic.claude-*`, `apac.anthropic.claude-*`) //! -//! For now the bridge's `chat()` / `chat_stream()` return a clear -//! `BridgeError::Config(...)` so a misconfigured `provider: -//! "amazon-bedrock"` row in the kine catalog surfaces a 501 / 502 with -//! an actionable message rather than silently dropping the dispatch. +//! 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()`. //! //! # Multi-publisher single-entry model //! @@ -45,13 +43,10 @@ //! //! - `us.anthropic.claude-3-5-sonnet-20241022-v2:0` //! -//! This mirrors LiteLLM's `bedrock/` design: every Bedrock-hosted -//! model goes through one provider name (`amazon-bedrock`) in cp-api's -//! catalog, and the publisher + region are resolved inside the bridge -//! from the model id. See -//! . -//! -//! Diverging from this would force every customer to register a +//! Single-entry routing: every Bedrock-hosted model goes through one +//! provider name (`amazon-bedrock`) in cp-api's catalog, and the +//! publisher + region are resolved inside the bridge from the model +//! id. Diverging from this would force every customer to register a //! separate provider_key per publisher even though the IAM role + AWS //! region are the same — exactly the operator pain `amazon-bedrock` //! solves. @@ -80,7 +75,8 @@ //! - Bedrock Runtime API — //! - Bedrock model IDs — //! - Cross-region inference profiles — -//! - LiteLLM `bedrock/` reference impl — +//! - Anthropic on Bedrock body shape — +//! - AWS Rust SDK `aws-sdk-bedrockruntime` — #![forbid(unsafe_code)] #![deny(rust_2018_idioms)] From 48b11c86b9b2c1036c2c40d24bbb652a5f5219cd Mon Sep 17 00:00:00 2001 From: Ming Wen Date: Sun, 17 May 2026 21:14:27 +0800 Subject: [PATCH 2/2] fix(provider-bedrock): address D7.2.a audit findings (H1-H4 + M2/M4/M5/M6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR #320 audit surfaced 4 HIGH + 6 MEDIUM + 3 LOW; this commit addresses the HIGH items and the MEDIUM items that have concrete test/code-level fixes. (M1 — anthropic-crate constructor visibility asymmetry — and M3 — explicit SdkError variant match — deferred as out-of-scope; they're stylistic improvements that don't gate merge.) ## HIGH (all fixed) H1 — Retry-After header propagation. Bedrock returns `Retry-After` on 429 throttle responses; the previous code collapsed it to `None`, silently degrading the cooldown layer's multi-region / burst behavior. `map_service_error` now converts the smithy HeaderMap → http::HeaderMap so it can call the gateway-level `parse_retry_after` helper. Pinned by new `chat_maps_upstream_429_with_retry_after_and_canned_rate_limited` test (mock returns `Retry-After: 42`, asserts `BridgeError::UpstreamStatus { retry_after: Some(42s), .. }`). H2 — `BedrockPublisher::Debug` taxonomy leak in customer error. The not-implemented-publisher error used `{publisher:?}` which formatted as `Other` / `Anthropic` etc. — internal labels that don't help the operator open the right follow-up tracking task. Replaced with the operator's actual model id + the publisher.name() catalog identifier. Pinned by new `chat_publisher_not_implemented_error_includes_model_id_and_publisher_name` (asserts `meta.llama3-3-70b-instruct-v1:0` and `publisher=meta` reach the error; `Other` and `` do NOT). H3 — `SdkError::TimeoutError` reported `elapsed_ms: 0` which formats as "timed out after 0ms" in customer logs. Now plumbs (started, deadline) into `map_sdk_error` so the actual elapsed budget is reported. Falls back to the configured deadline if elapsed rounds to 0 (clock skew defense). H4 — Tool-call test missing OpenAI-spec assertion. The previous `chat_anthropic_handles_tool_use_response_blocks` test asserted `function.name` but not the `arguments` shape. Per OpenAI's Chat Completions spec, `arguments` MUST be a JSON-encoded STRING (not a parsed object) so SDK consumers can do `JSON.parse(toolCall.function.arguments)`. A future refactor that passed the parsed object would silently break every OpenAI-SDK caller against an Anthropic upstream. Test now pins the string shape AND round-trip-parses it back to verify the original `{"city": "SF"}` arguments. ## MEDIUM (addressed) M2 — `validate_model_id_chars` defense-in-depth check. The AWS SDK URL-encodes reserved chars but the gateway layer must reject upfront: the model id propagates into metrics labels, so an embedded `\t` / whitespace / `?` / `#` would corrupt dashboards. Allowed set: `[A-Za-z0-9._:/-]`. Pinned by new `chat_rejects_model_id_with_path_injection_chars`. M4 — `chat_stream` distinguishes "anthropic streaming not wired" (D7.2.b — same publisher as chat, just streaming) from "publisher X not wired at all" (D7.3+). The previous code returned the same generic "streaming not yet implemented" error for both, which mis-routed operators to the wrong tracking task. Two new tests pin the split: `chat_stream_anthropic_returns_d7_2_b_specific_error` and `chat_stream_non_anthropic_publisher_returns_d7_3_specific_error`. M5 — 4xx redaction test was a weak negative assertion (does NOT contain leaky strings). Strengthened with exact-match positive assertion on the canned phrase (`assert_eq!(message, "upstream returned 400")`). A future refactor that re-rendered SDK metadata into the message would pass the absence check but fail the exact-match. M6 — Cross-region dispatch parity. The only e2e test covered `us.`; the historically-broken case (`us-gov.` with hyphen) and `global.` (exactly 6 chars, accidentally working under the old matcher) lacked dispatch-path coverage. Added `chat_with_us_gov_cross_region_prefix_dispatches_with_full_model_id` and `chat_with_global_cross_region_prefix_dispatches_with_full_model_id`. ## Verification cargo test -p aisix-provider-bedrock → 41 passed (was 35; +6 audit regression tests) cargo clippy --workspace --all-targets -- -D warnings → clean cargo fmt --check → clean ## Deferred M1 — `AnthropicRequest` field visibility asymmetry. Workspace-internal surface decision; doesn't gate merge. Tracked for follow-up. M3 — Explicit `SdkError` variant arms. `SdkError` is `non_exhaustive`; future SDK upgrades adding new variants would silently fall into the `_ =>` catch-all. Tracked for follow-up — needs a CI lint to catch on SDK bump, not just code-level enumeration. L1/L3 — `Other` Debug variant + metrics-label concern. Deferred to the D7.3+ PRs that wire `Other` publisher dispatch (the variant goes away then). --- crates/aisix-provider-bedrock/Cargo.toml | 1 + crates/aisix-provider-bedrock/src/bridge.rs | 348 ++++++++++++++++++-- 2 files changed, 329 insertions(+), 20 deletions(-) diff --git a/crates/aisix-provider-bedrock/Cargo.toml b/crates/aisix-provider-bedrock/Cargo.toml index 6b06fe34..c1bbe217 100644 --- a/crates/aisix-provider-bedrock/Cargo.toml +++ b/crates/aisix-provider-bedrock/Cargo.toml @@ -23,6 +23,7 @@ tracing.workspace = true serde.workspace = true serde_json.workspace = true tokio.workspace = true +http.workspace = true # AWS SDK: handles SigV4 + retries + binary event-stream framing. # Same workspace deps the guardrails crate uses for ApplyGuardrail. aws-config.workspace = true diff --git a/crates/aisix-provider-bedrock/src/bridge.rs b/crates/aisix-provider-bedrock/src/bridge.rs index 21ab07f0..8a9492ec 100644 --- a/crates/aisix-provider-bedrock/src/bridge.rs +++ b/crates/aisix-provider-bedrock/src/bridge.rs @@ -29,6 +29,7 @@ use aws_sdk_bedrockruntime::primitives::Blob; use aws_sdk_bedrockruntime::Client as BedrockClient; use aws_smithy_runtime_api::client::result::ServiceError; use serde::Deserialize; +use std::time::{Duration, Instant}; use aisix_provider_anthropic::wire::{ build_request, response_into_chat_response, split_system, AnthropicResponse, @@ -340,9 +341,29 @@ fn upstream_model(ctx: &BridgeContext) -> Result<&str, BridgeError> { /// numbers (in ARNs), and IAM role names. Surfacing these to a /// downstream customer leaks operator-internal taxonomy. We map to /// canned status-keyed phrases. -fn map_sdk_error(err: SdkError) -> BridgeError { +/// +/// **Audit H3** — `deadline` is threaded through so a SDK-side timeout +/// reports the actual elapsed budget instead of `0ms` (which formats +/// as "timed out after 0ms" in customer logs). +fn map_sdk_error( + err: SdkError, + started: Instant, + deadline: Option, +) -> BridgeError { match err { - SdkError::TimeoutError(_) => BridgeError::Timeout { elapsed_ms: 0 }, + SdkError::TimeoutError(_) => { + // Prefer the actual elapsed budget; fall back to the + // deadline if elapsed somehow rounds to 0 (clock skew). + let elapsed_ms = started.elapsed().as_millis() as u64; + let reported = if elapsed_ms > 0 { + elapsed_ms + } else { + deadline.map(|d| d.as_millis() as u64).unwrap_or(0) + }; + BridgeError::Timeout { + elapsed_ms: reported, + } + } SdkError::DispatchFailure(_) => BridgeError::Transport("upstream dispatch failed".into()), SdkError::ConstructionFailure(_) => { BridgeError::Config("upstream request construction failed".into()) @@ -355,11 +376,29 @@ fn map_sdk_error(err: SdkError) -> BridgeError { } } +/// Audit H1 — propagate `Retry-After` from the upstream's HTTP +/// response so the gateway's cooldown layer gets the actual upstream +/// hint instead of falling back to its configured default. Bedrock +/// returns `Retry-After` on 429 throttle responses; collapsing it to +/// `None` silently degrades multi-region / burst behavior. fn map_service_error( svc: ServiceError, ) -> BridgeError { let raw = svc.into_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(), @@ -371,8 +410,30 @@ fn map_service_error( BridgeError::UpstreamStatus { status, message, - retry_after: None, + retry_after, + } +} + +/// **Audit M2** — defense-in-depth check on the upstream model id +/// before it's URL-encoded into the Bedrock `/model//invoke` +/// path. The SDK encodes reserved characters, but pinning the +/// allowed set at the gateway layer prevents log-injection / +/// dashboard-label corruption (the model id propagates into metrics +/// labels) and forces typos to fail loudly at registration time. +/// +/// Bedrock model ids are documented as +/// `.-:` with all-ASCII tokens. +fn validate_model_id_chars(model_id: &str) -> Result<(), BridgeError> { + if !model_id + .chars() + .all(|c| c.is_ascii_alphanumeric() || matches!(c, '.' | '-' | ':' | '_' | '/')) + { + return Err(BridgeError::Config(format!( + "bedrock model id {model_id:?} contains unexpected characters — \ + only [A-Za-z0-9._:/-] are allowed" + ))); } + Ok(()) } #[async_trait] @@ -387,6 +448,7 @@ impl Bridge for BedrockBridge { ctx: &BridgeContext, ) -> Result { let upstream_id = upstream_model(ctx)?; + validate_model_id_chars(upstream_id)?; let publisher = BedrockPublisher::from_model_id(upstream_id).ok_or_else(|| { BridgeError::Config(format!( "bedrock publisher unknown for model id {upstream_id:?}; \ @@ -402,9 +464,15 @@ impl Bridge for BedrockBridge { 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 publisher {publisher:?} not yet implemented — \ - tracked under api7/AISIX-Cloud#302 Phase G (D7.3+, publisher={})", + "bedrock dispatch for model id {upstream_id:?} (publisher={}) \ + not yet implemented — tracked under api7/AISIX-Cloud#302 \ + Phase G (D7.3+)", other.name() ))), } @@ -416,7 +484,8 @@ impl Bridge for BedrockBridge { ctx: &BridgeContext, ) -> Result { let upstream_id = upstream_model(ctx)?; - let _publisher = BedrockPublisher::from_model_id(upstream_id).ok_or_else(|| { + validate_model_id_chars(upstream_id)?; + let publisher = BedrockPublisher::from_model_id(upstream_id).ok_or_else(|| { BridgeError::Config(format!( "bedrock publisher unknown for model id {upstream_id:?}; \ expected one of anthropic.claude-* / meta.llama* / mistral.* / \ @@ -424,11 +493,24 @@ impl Bridge for BedrockBridge { (optionally prefixed with a cross-region inference profile like us. / eu. / apac.)" )) })?; - Err(BridgeError::Config( - "bedrock streaming is not yet implemented — \ - tracked under api7/AISIX-Cloud#302 Phase G (D7.2.b)" - .into(), - )) + // 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() + ))), + } } } @@ -487,6 +569,8 @@ impl BedrockBridge { // 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 .invoke_model() .model_id(upstream_id) @@ -495,7 +579,7 @@ impl BedrockBridge { .body(Blob::new(body_bytes)) .send() .await - .map_err(map_sdk_error)?; + .map_err(|e| map_sdk_error(e, started, deadline))?; let parsed: AnthropicResponse = serde_json::from_slice(resp.body().as_ref()) .map_err(|e| BridgeError::UpstreamDecode(e.to_string()))?; @@ -1142,6 +1226,19 @@ mod tests { .and_then(|v| v.as_str()), Some("get_weather") ); + // Audit H4: `arguments` MUST be a JSON-encoded STRING per the + // OpenAI Chat Completions spec, not a parsed object. SDK + // consumers do `JSON.parse(toolCall.function.arguments)` — a + // future refactor that passes the parsed object would silently + // break every OpenAI-SDK caller against an Anthropic upstream. + let args = tool_calls[0] + .get("function") + .and_then(|f| f.get("arguments")) + .and_then(|v| v.as_str()) + .expect("arguments must be a JSON-encoded STRING per OpenAI spec"); + let parsed: serde_json::Value = + serde_json::from_str(args).expect("arguments string must itself be valid JSON"); + assert_eq!(parsed.get("city").and_then(|v| v.as_str()), Some("SF")); } #[tokio::test] @@ -1149,6 +1246,11 @@ mod tests { // Audit M1: Bedrock error envelopes can contain account // numbers (in ARNs), model IDs, IAM role names — must not // leak into customer-visible error. + // + // Audit M5 follow-up: assert the canned message EXACTLY, + // not just absence-of-leak. A future refactor that re-renders + // SDK metadata into the message would pass an absence check + // but fail the exact-match assertion. let server = MockServer::start().await; Mock::given(method("POST")) .and(path_regex(r"^/model/.+/invoke$")) @@ -1170,24 +1272,41 @@ mod tests { let req = ChatFormat::new("my-claude", vec![ChatMessage::user("hi")]); let err = bridge.chat(&req, &ctx).await.unwrap_err(); match err { - BridgeError::UpstreamStatus { message, .. } => { + BridgeError::UpstreamStatus { + status, message, .. + } => { + assert_eq!(status, 400); assert!( !message.contains("123456789012") && !message.contains("internal-leaky-role"), "upstream body must not leak account / role info into customer error; got {message:?}" ); + // Positive pin (audit M5): exact-match the canned + // status-keyed phrase. Bedrock returns 400 → bucket + // is "upstream returned 400" per `map_service_error`. + assert_eq!( + message, "upstream returned 400", + "must emit canned 4xx phrasing only; got {message:?}" + ); } other => panic!("expected UpstreamStatus, got {other:?}"), } } #[tokio::test] - async fn chat_maps_upstream_429_to_canned_rate_limited() { + async fn chat_maps_upstream_429_with_retry_after_and_canned_rate_limited() { + // Audit H1: Bedrock's `Retry-After` header on 429 must reach + // the cooldown layer. Collapsing it to `None` silently + // degrades multi-region / burst behavior. let server = MockServer::start().await; Mock::given(method("POST")) .and(path_regex(r"^/model/.+/invoke$")) - .respond_with(ResponseTemplate::new(429).set_body_json(serde_json::json!({ - "message": "Too many requests for account 123456789012" - }))) + .respond_with( + ResponseTemplate::new(429) + .insert_header("retry-after", "42") + .set_body_json(serde_json::json!({ + "message": "Too many requests for account 123456789012" + })), + ) .mount(&server) .await; @@ -1201,16 +1320,27 @@ mod tests { let err = bridge.chat(&req, &ctx).await.unwrap_err(); match err { BridgeError::UpstreamStatus { - status, message, .. + status, + message, + retry_after, } => { assert_eq!(status, 429); - assert!(message.contains("rate limited")); + assert_eq!(message, "upstream rate limited"); assert!( !message.contains("123456789012"), "must not leak account id; got {message:?}" ); + // Audit H1 pin: the SDK / smithy headers must round-trip + // Retry-After into the BridgeError so the cooldown + // layer sees the upstream's hint instead of falling + // back to a configured default. + assert_eq!( + retry_after, + Some(std::time::Duration::from_secs(42)), + "Retry-After must reach BridgeError::UpstreamStatus" + ); } - other => panic!("expected UpstreamStatus, got {other:?}"), + other => panic!("expected UpstreamStatus with retry_after, got {other:?}"), } } @@ -1247,6 +1377,184 @@ mod tests { assert_eq!(chat.message.content, "cross-region ok"); } + /// Audit M6: cross-region dispatch coverage was only `us.`; the + /// historically-broken case (`us-gov.` with hyphen) and `global.` + /// (exactly 6 chars — accidentally working under the old matcher) + /// need real dispatch-path tests so a future regression in + /// `strip_region_prefix` is caught at the wire layer, not just at + /// the unit-test layer. + #[tokio::test] + async fn chat_with_us_gov_cross_region_prefix_dispatches_with_full_model_id() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex( + r"^/model/us-gov\.anthropic\.claude-3-5-sonnet-20241022-v2(:0|%3A0)/invoke$", + )) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "msg_xr", "model": "claude-3-5-sonnet", + "content": [{"type": "text", "text": "us-gov ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1} + }))) + .expect(1) + .mount(&server) + .await; + + let bridge = BedrockBridge::new().with_endpoint_override(server.uri()); + let ctx = BridgeContext::new( + "req-1", + sample_model_with("us-gov.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")]); + let chat = bridge.chat(&req, &ctx).await.unwrap(); + assert_eq!(chat.message.content, "us-gov ok"); + } + + #[tokio::test] + async fn chat_with_global_cross_region_prefix_dispatches_with_full_model_id() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path_regex( + r"^/model/global\.anthropic\.claude-3-5-sonnet-20241022-v2(:0|%3A0)/invoke$", + )) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "msg_xr", "model": "claude-3-5-sonnet", + "content": [{"type": "text", "text": "global ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1} + }))) + .expect(1) + .mount(&server) + .await; + + let bridge = BedrockBridge::new().with_endpoint_override(server.uri()); + let ctx = BridgeContext::new( + "req-1", + sample_model_with("global.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")]); + let chat = bridge.chat(&req, &ctx).await.unwrap(); + 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:?}"), + } + } + + /// Audit M2 regression: defense-in-depth model-id char check. + /// Even though the AWS SDK URL-encodes reserved chars, the gateway + /// layer must reject upfront so the model id can't carry + /// log-injection / dashboard-corruption payloads (it propagates + /// into metrics labels). + #[tokio::test] + async fn chat_rejects_model_id_with_path_injection_chars() { + let bridge = BedrockBridge::new(); + // Whitespace + tab — would corrupt metrics labels even if the + // SDK URL-encoded the path correctly. + let evil_model = "anthropic.claude\t evil model"; + let ctx = BridgeContext::new( + "req-1", + sample_model_with(evil_model), + 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("unexpected characters"), + "must reject invalid model id chars; got {msg}" + ); + } + other => panic!("expected Config error, got {other:?}"), + } + } + + /// 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:?}"), + } + } + #[tokio::test] async fn chat_anthropic_translates_system_messages_to_system_field() { // Anthropic's Messages API takes `system` as a top-level