Skip to main content

llm/providers/bedrock/
provider.rs

1use super::mantle::{MantleAuth, MantleClient};
2use super::mappers::{default_cache_point, map_messages, map_tools};
3use super::streaming::{converse_events, decode_converse_event};
4use crate::catalog::Provider;
5use crate::catalog::transport::ModelTransport;
6use crate::provider::{LlmResponseStream, StreamingModelProvider, validate_reasoning};
7use crate::providers::openai::responses::streaming::decode_responses;
8use crate::providers::response_stream::{OpenedStream, error_stream, response_stream};
9use crate::{Context, LlmError, ProviderAuthMode, ProviderConnectionConfig, ProviderError, Result};
10use aws_config::Region;
11use aws_credential_types::provider::SharedCredentialsProvider;
12use aws_sdk_bedrockruntime::config::{BehaviorVersion, Credentials};
13use aws_sdk_bedrockruntime::error::SdkError;
14use aws_sdk_bedrockruntime::operation::converse_stream::ConverseStreamError;
15use aws_sdk_bedrockruntime::primitives::event_stream::EventReceiver;
16use aws_sdk_bedrockruntime::types::error::ConverseStreamOutputError;
17use aws_sdk_bedrockruntime::types::{ConverseStreamOutput, InferenceConfiguration};
18use aws_sdk_bedrockruntime::{Client, Config};
19use std::time::Duration;
20use tracing::{error, info, warn};
21
22const DEFAULT_MODEL: &str = "anthropic.claude-sonnet-4-5-20250929-v1:0";
23const DEFAULT_MAX_TOKENS: i32 = 16_384;
24const DEFAULT_REGION: &str = "us-east-1";
25
26/// AWS credentials for explicit authentication with Bedrock.
27#[derive(Clone)]
28pub struct AwsCredentials {
29    pub access_key_id: String,
30    pub secret_access_key: String,
31    pub session_token: Option<String>,
32}
33
34#[derive(Clone)]
35pub struct BedrockProvider {
36    client: Client,
37    mantle: MantleClient,
38    model: String,
39    inference_profile_arn: Option<String>,
40    idle_timeout: Duration,
41}
42
43impl BedrockProvider {
44    /// Create a provider using the default AWS credential chain
45    /// (env vars, `~/.aws/credentials`, IAM roles, SSO).
46    pub async fn new(connection: ProviderConnectionConfig) -> Self {
47        if connection.auth_mode == ProviderAuthMode::None {
48            return Self::from_config(None, region_from_env().as_deref(), connection);
49        }
50
51        let mut loader = aws_config::defaults(BehaviorVersion::latest());
52        if let Some(url) = &connection.base_url {
53            loader = loader.endpoint_url(url.clone());
54        }
55
56        let config = loader.load().await;
57        let region = config
58            .region()
59            .map(ToString::to_string)
60            .or_else(region_from_env)
61            .unwrap_or_else(|| DEFAULT_REGION.to_string());
62        let auth = mantle_auth(config.credentials_provider(), &region);
63        Self::from_parts(Client::new(&config), region, auth, connection)
64    }
65
66    /// Create a provider from explicit configuration without async credential discovery.
67    pub fn from_config(
68        credentials: Option<AwsCredentials>,
69        region: Option<&str>,
70        connection: ProviderConnectionConfig,
71    ) -> Self {
72        let region = region.unwrap_or(DEFAULT_REGION).to_string();
73        let auth = match connection.auth_mode {
74            ProviderAuthMode::Default => mantle_auth(credentials.clone().map(shared_credentials), &region),
75            ProviderAuthMode::None => MantleAuth::None,
76        };
77        let client = build_client(credentials, &region, connection.base_url.as_deref(), connection.auth_mode);
78        Self::from_parts(client, region, auth, connection)
79    }
80
81    fn from_parts(client: Client, region: String, auth: MantleAuth, connection: ProviderConnectionConfig) -> Self {
82        let mantle = MantleClient::new(region, auth, connection.base_url.clone());
83        Self {
84            client,
85            mantle,
86            model: DEFAULT_MODEL.to_string(),
87            inference_profile_arn: connection.inference_profile_arn,
88            idle_timeout: connection.idle_timeout,
89        }
90    }
91
92    pub fn with_model(mut self, model: &str) -> Self {
93        self.model = model.to_string();
94        self
95    }
96
97    pub fn with_inference_profile_arn(mut self, arn: impl Into<String>) -> Self {
98        self.inference_profile_arn = Some(arn.into());
99        self
100    }
101
102    /// Authenticate Responses-transport requests with a Bedrock API key instead
103    /// of `SigV4`. Equivalent to setting `AWS_BEARER_TOKEN_BEDROCK`.
104    pub fn with_bearer_token(mut self, token: impl Into<String>) -> Self {
105        self.mantle = self.mantle.with_auth(MantleAuth::BearerToken(token.into()));
106        self
107    }
108
109    fn request_model_id(&self) -> &str {
110        self.inference_profile_arn.as_deref().unwrap_or(&self.model)
111    }
112
113    /// The Responses-API transport for the current model, when the catalog says
114    /// it is not served by the Converse API.
115    fn mantle_transport(&self) -> Option<ModelTransport> {
116        self.model().and_then(|model| model.transport())
117    }
118
119    async fn send_converse_stream(
120        &self,
121        context: &Context,
122    ) -> Result<EventReceiver<ConverseStreamOutput, ConverseStreamOutputError>> {
123        let cache_point =
124            self.model().is_some_and(|m| m.supports_prompt_caching()).then(default_cache_point).transpose()?;
125        let (system_blocks, messages) = map_messages(context.messages(), cache_point.as_ref())?;
126        let settings = context.model_settings();
127        let max_tokens = settings.max_tokens.and_then(|m| i32::try_from(m).ok()).unwrap_or(DEFAULT_MAX_TOKENS);
128        let inference_config = InferenceConfiguration::builder()
129            .max_tokens(max_tokens)
130            .set_temperature(settings.temperature)
131            .set_top_p(settings.top_p)
132            .build();
133        let tool_config =
134            (!context.tools().is_empty()).then(|| map_tools(context.tools(), cache_point.as_ref())).transpose()?;
135
136        info!(
137            model = %self.model,
138            inference_profile_arn = ?self.inference_profile_arn,
139            "Sending Bedrock converse_stream request"
140        );
141
142        let request = self
143            .client
144            .converse_stream()
145            .model_id(self.request_model_id())
146            .set_messages(Some(messages))
147            .set_system((!system_blocks.is_empty()).then_some(system_blocks))
148            .inference_config(inference_config)
149            .set_tool_config(tool_config);
150
151        let response = request.send().await.map_err(|e| {
152            error!(model = %self.model, error = ?e, "Bedrock API error");
153            LlmError::from(e)
154        })?;
155
156        Ok(response.stream)
157    }
158}
159
160impl StreamingModelProvider for BedrockProvider {
161    fn model(&self) -> Option<crate::LlmModel> {
162        crate::LlmModel::parse_for(Provider::Bedrock, &self.model).ok()
163    }
164
165    fn stream_response(&self, context: &Context) -> LlmResponseStream {
166        let identity = self.model();
167        if let Err(error) = validate_reasoning(context, identity.as_ref()) {
168            return error_stream(error);
169        }
170        let provider = self.clone();
171        let context = context.clone();
172
173        let Some(transport) = self.mantle_transport() else {
174            return response_stream(
175                async move { Ok(OpenedStream::new(converse_events(provider.send_converse_stream(&context).await?))) },
176                |event, turn| Ok(decode_converse_event(event, turn)),
177                self.idle_timeout,
178            );
179        };
180
181        if let Some(arn) = self.inference_profile_arn.as_deref() {
182            warn!(
183                model = %self.model,
184                inference_profile_arn = %arn,
185                "Ignoring inferenceProfileArn: Responses requests use the selected model ID"
186            );
187        }
188
189        response_stream(
190            async move { provider.mantle.stream(identity.as_ref(), &provider.model, &transport, &context).await },
191            decode_responses(),
192            self.idle_timeout,
193        )
194    }
195}
196
197impl From<SdkError<ConverseStreamError>> for LlmError {
198    fn from(e: SdkError<ConverseStreamError>) -> Self {
199        let message = format!("Bedrock API error: {e}");
200        let status = e.raw_response().map(|r| r.status().as_u16());
201        let request_id = e.raw_response().and_then(|r| r.headers().get("x-amzn-requestid")).map(str::to_string);
202        let mut provider = match &e {
203            SdkError::TimeoutError(_) => ProviderError::timeout(message),
204            SdkError::DispatchFailure(_) => ProviderError::network(message),
205            SdkError::ResponseError(_) => ProviderError::server(message),
206            SdkError::ServiceError(svc) => {
207                let inner = svc.err();
208                if inner.is_throttling_exception() {
209                    ProviderError::rate_limit(message)
210                } else if inner.is_service_unavailable_exception()
211                    || inner.is_internal_server_exception()
212                    || inner.is_model_stream_error_exception()
213                {
214                    ProviderError::server(message)
215                } else {
216                    ProviderError::api(message)
217                }
218            }
219            _ => ProviderError::api(message),
220        };
221        provider = provider.with_http_metadata(status, request_id);
222        Self::from(provider)
223    }
224}
225
226fn build_client(
227    credentials: Option<AwsCredentials>,
228    region: &str,
229    base_url: Option<&str>,
230    auth_mode: ProviderAuthMode,
231) -> Client {
232    let mut config =
233        Config::builder().behavior_version(BehaviorVersion::latest()).region(Region::new(region.to_string()));
234
235    if auth_mode == ProviderAuthMode::None {
236        config = config.allow_no_auth();
237    } else if let Some(credentials) = credentials {
238        config = config.credentials_provider(shared_credentials(credentials));
239    }
240    if let Some(url) = base_url {
241        config = config.endpoint_url(url);
242    }
243
244    Client::from_conf(config.build())
245}
246
247fn shared_credentials(credentials: AwsCredentials) -> SharedCredentialsProvider {
248    SharedCredentialsProvider::new(Credentials::new(
249        credentials.access_key_id,
250        credentials.secret_access_key,
251        credentials.session_token,
252        None,
253        "aether-bedrock-provider",
254    ))
255}
256
257/// Resolve the credential scheme for the Responses transport.
258///
259/// A Bedrock API key takes precedence over the credential chain because it is
260/// the scheme the model catalog advertises for these endpoints; `SigV4` keeps
261/// SSO and IAM-role users working without extra configuration.
262fn mantle_auth(credentials: Option<SharedCredentialsProvider>, region: &str) -> MantleAuth {
263    if let Some(token) = MantleAuth::bearer_token_from_env() {
264        return MantleAuth::BearerToken(token);
265    }
266    match credentials {
267        Some(credentials) => MantleAuth::sigv4(credentials, region),
268        None => MantleAuth::None,
269    }
270}
271
272fn region_from_env() -> Option<String> {
273    ["AWS_REGION", "AWS_DEFAULT_REGION"].into_iter().find_map(|name| match std::env::var(name) {
274        Ok(value) if !value.is_empty() => Some(value),
275        _ => None,
276    })
277}
278
279#[cfg(test)]
280mod tests {
281    use super::*;
282    use crate::providers::test_capture_server::{CaptureServer, ResponseSpec, hello_context};
283    use crate::types::IsoString;
284    use crate::{AssistantReasoning, ChatMessage, EncryptedReasoningContent, LlmModel, MessageId};
285    use futures::StreamExt;
286    use utils::ReasoningEffort;
287
288    #[test]
289    fn test_display_name() {
290        assert_eq!(test_provider().display_name(), "AWS Bedrock (anthropic.claude-sonnet-4-5-20250929-v1:0)");
291    }
292
293    #[test]
294    fn test_with_model() {
295        let provider = test_provider().with_model("anthropic.claude-opus-4-20250514-v1:0");
296        assert_eq!(provider.display_name(), "AWS Bedrock (anthropic.claude-opus-4-20250514-v1:0)");
297    }
298
299    #[test]
300    fn test_default_values() {
301        let provider = test_provider();
302        assert_eq!(provider.model, "anthropic.claude-sonnet-4-5-20250929-v1:0");
303    }
304
305    #[tokio::test]
306    async fn auth_none_sends_unsigned_request_to_custom_endpoint() {
307        let mut server = CaptureServer::start_with_response("{}").await;
308        let provider = create_provider(&server.base_url, DEFAULT_MODEL).await;
309
310        let _ = provider.stream_response(&hello_context()).collect::<Vec<_>>().await;
311        let request = server.captured().await;
312
313        assert!(request.path.starts_with("/model/"), "{}", request.path);
314        assert!(!request.headers.contains_key("authorization"), "request was signed: {:?}", request.headers);
315        assert!(
316            !request.headers.contains_key("x-amz-security-token"),
317            "request included session token: {:?}",
318            request.headers
319        );
320    }
321
322    fn static_credentials() -> AwsCredentials {
323        AwsCredentials {
324            access_key_id: "AKIAIOSFODNN7EXAMPLE".to_string(),
325            secret_access_key: "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY".to_string(),
326            session_token: None,
327        }
328    }
329
330    #[tokio::test]
331    async fn mantle_disabled_uses_none_without_summary() {
332        let model = LlmModel::all()
333            .iter()
334            .find(|model| {
335                model.provider_enum() == Provider::Bedrock
336                    && model.transport().is_some()
337                    && model.supports_reasoning_off()
338            })
339            .unwrap();
340        let mut server = CaptureServer::start_responses().await;
341        let provider = create_provider(&server.base_url, &model.model_id()).await;
342        let mut context = hello_context();
343        context.set_reasoning_effort(ReasoningEffort::Disabled);
344        let responses = provider.stream_response(&context).collect::<Vec<_>>().await;
345        assert!(responses.iter().all(Result::is_ok), "{responses:?}");
346        let body = server.captured().await.body;
347        assert_eq!(body["reasoning"]["effort"], "none");
348        assert!(body["reasoning"]["summary"].is_null());
349    }
350
351    #[tokio::test]
352    async fn converse_disabled_is_a_non_retryable_transport_error() {
353        let model = LlmModel::all()
354            .iter()
355            .find(|model| {
356                model.provider_enum() == Provider::Bedrock
357                    && model.transport().is_none()
358                    && model.supports_reasoning_off()
359            })
360            .unwrap();
361        let provider =
362            BedrockProvider::new(ProviderConnectionConfig { auth_mode: ProviderAuthMode::None, ..Default::default() })
363                .await
364                .with_model(&model.model_id());
365        let mut context = hello_context();
366        context.set_reasoning_effort(crate::ReasoningEffort::Disabled);
367        let responses = provider.stream_response(&context).collect::<Vec<_>>().await;
368        assert_eq!(responses.len(), 1);
369        assert!(matches!(&responses[0], Err(LlmError::UnsupportedDisableTransport { .. })));
370        assert!(!responses[0].as_ref().unwrap_err().is_retryable());
371    }
372
373    #[tokio::test]
374    async fn responses_shape_models_are_sent_to_the_responses_endpoint() {
375        let mut server = CaptureServer::start_responses().await;
376        let provider = mantle_provider(&server).await;
377        let mut context = hello_context();
378        context.set_reasoning_effort(crate::ReasoningEffort::Xhigh);
379
380        let responses = provider.stream_response(&context).collect::<Vec<_>>().await;
381        let captured = server.captured().await;
382
383        assert!(!responses.is_empty());
384        assert_eq!(captured.path, "/responses");
385        assert_eq!(captured.body["model"], mantle_model());
386        assert_eq!(captured.body["stream"], true);
387        assert_eq!(captured.body["store"], false);
388        assert_eq!(captured.body["reasoning"]["effort"], "xhigh");
389        assert_eq!(captured.body["input"][0]["role"], "user");
390        assert_eq!(captured.body["input"][0]["content"][0]["type"], "input_text");
391        assert_eq!(captured.body["input"][0]["content"][0]["text"], "Hello");
392        assert!(captured.headers.get("authorization").is_none(), "{:?}", captured.headers);
393    }
394
395    #[tokio::test]
396    async fn http_200_failed_server_error_is_retryable_with_diagnostics() {
397        let spec = ResponseSpec::sse(include_str!("../../../tests/fixtures/openai_responses/04_failed_server.sse"))
398            .with_header("x-amzn-requestid", "amzn-req-123");
399        let mut server = CaptureServer::start_with_spec(spec).await;
400        let provider = mantle_provider(&server).await;
401
402        let responses = provider.stream_response(&hello_context()).collect::<Vec<_>>().await;
403        let _ = server.captured().await;
404
405        assert!(!responses.iter().any(|r| matches!(r, Ok(crate::LlmResponse::Done { .. }))));
406        let err = responses.iter().find_map(|r| r.as_ref().err()).expect("expected a failure");
407        assert!(err.is_retryable(), "server_error must be retryable: {err:?}");
408        let provider_error = err.provider().expect("expected provider error");
409        assert_eq!(provider_error.kind, crate::ProviderErrorKind::Server);
410        assert_eq!(provider_error.code.as_deref(), Some("server_error"));
411        assert_eq!(provider_error.http_status, Some(200));
412        assert_eq!(provider_error.request_id.as_deref(), Some("amzn-req-123"));
413    }
414
415    #[tokio::test]
416    async fn failed_event_with_unknown_code_is_terminal_without_done() {
417        let body = "event: response.created\ndata: {\"type\":\"response.created\",\"response\":{\"id\":\"r1\"}}\n\nevent: response.failed\ndata: {\"type\":\"response.failed\",\"response\":{\"error\":{\"code\":\"invalid_prompt\",\"message\":\"bad prompt\"}}}\n\n";
418        let mut server = CaptureServer::start_with_response(body).await;
419        let provider = mantle_provider(&server).await;
420
421        let responses = provider.stream_response(&hello_context()).collect::<Vec<_>>().await;
422        let _ = server.captured().await;
423
424        assert!(!responses.iter().any(|r| matches!(r, Ok(crate::LlmResponse::Done { .. }))));
425        let err = responses.iter().find_map(|r| r.as_ref().err()).expect("expected a failure");
426        assert!(!err.is_retryable(), "invalid_prompt must be terminal: {err:?}");
427    }
428
429    #[tokio::test]
430    async fn responses_transport_rejects_malformed_and_truncated_streams() {
431        for response in [
432            "data: {not-json}\n\n",
433            "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\"}}\n\ndata: [DONE]\n\n",
434        ] {
435            let mut server = CaptureServer::start_with_response(response).await;
436            let provider = mantle_provider(&server).await;
437
438            let responses = provider.stream_response(&hello_context()).collect::<Vec<_>>().await;
439            let _ = server.captured().await;
440
441            assert!(responses.iter().any(|response| {
442                response.as_ref().err().and_then(LlmError::provider).map(|provider| provider.kind)
443                    == Some(crate::ProviderErrorKind::StreamInterrupted)
444            }));
445            assert!(!responses.iter().any(|response| matches!(response, Ok(crate::LlmResponse::Done { .. }))));
446        }
447    }
448
449    #[tokio::test]
450    async fn responses_transport_drops_encrypted_reasoning_from_another_model() {
451        let mut server = CaptureServer::start_responses().await;
452        let provider = mantle_provider(&server).await;
453        let context = Context::new(
454            vec![ChatMessage::Assistant {
455                message_id: MessageId::new(),
456                content: "previous answer".to_string(),
457                reasoning: AssistantReasoning {
458                    summary_text: None,
459                    encrypted_content: Some(EncryptedReasoningContent {
460                        id: "reasoning-id".to_string(),
461                        model: "bedrock:openai.gpt-5.5".parse().unwrap(),
462                        content: "opaque-for-another-model".to_string(),
463                    }),
464                },
465                timestamp: IsoString::now(),
466                tool_calls: vec![],
467            }],
468            vec![],
469        );
470
471        let _ = provider.stream_response(&context).collect::<Vec<_>>().await;
472        let captured = server.captured().await;
473
474        assert!(
475            captured.body["input"].as_array().unwrap().iter().all(|item| item["type"] != "reasoning"),
476            "{}",
477            captured.body
478        );
479    }
480
481    #[tokio::test]
482    async fn converse_shape_models_do_not_use_the_responses_endpoint() {
483        let mut server = CaptureServer::start_with_response("{}").await;
484        let provider = create_provider(&server.base_url, DEFAULT_MODEL).await;
485
486        let _ = provider.stream_response(&hello_context()).collect::<Vec<_>>().await;
487        let request = server.captured().await;
488
489        assert!(request.path.starts_with("/model/"), "{}", request.path);
490    }
491
492    #[tokio::test]
493    async fn bearer_token_authenticates_responses_requests() {
494        let mut server = CaptureServer::start_responses().await;
495        let provider = BedrockProvider::from_config(
496            None,
497            Some("us-west-2"),
498            ProviderConnectionConfig { base_url: Some(server.base_url.clone()), ..Default::default() },
499        )
500        .with_bearer_token("test-token")
501        .with_model(&mantle_model());
502
503        let _ = provider.stream_response(&hello_context()).collect::<Vec<_>>().await;
504        let captured = server.captured().await;
505
506        assert_eq!(captured.headers.get("authorization").unwrap(), "Bearer test-token");
507    }
508
509    #[tokio::test]
510    async fn credential_chain_sigv4_signs_responses_requests() {
511        let mut server = CaptureServer::start_responses().await;
512        let provider = BedrockProvider::from_config(
513            Some(static_credentials()),
514            Some("us-west-2"),
515            ProviderConnectionConfig { base_url: Some(server.base_url.clone()), ..Default::default() },
516        )
517        .with_model(&mantle_model());
518
519        let _ = provider.stream_response(&hello_context()).collect::<Vec<_>>().await;
520        let captured = server.captured().await;
521
522        let authorization = captured.headers.get("authorization").expect("request was not signed").to_str().unwrap();
523        assert!(
524            authorization.starts_with("AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/"),
525            "unexpected authorization header: {authorization}"
526        );
527        assert!(authorization.contains("/us-west-2/bedrock/aws4_request"), "{authorization}");
528        assert!(captured.headers.contains_key("x-amz-date"), "{:?}", captured.headers);
529    }
530
531    #[test]
532    fn only_responses_shape_models_route_to_the_mantle_transport() {
533        let mantle = test_provider().with_model(&mantle_model());
534        let converse = test_provider().with_model(DEFAULT_MODEL);
535        let profile = test_provider().with_model("us.anthropic.claude-future-model-v99:0");
536
537        assert!(matches!(mantle.mantle_transport(), Some(ModelTransport::OpenAiResponses { .. })));
538        assert_eq!(converse.mantle_transport(), None);
539        assert_eq!(profile.mantle_transport(), None);
540    }
541
542    #[test]
543    fn explicit_connection_preserves_inference_profile() {
544        let provider = BedrockProvider::from_config(
545            None,
546            Some("us-west-2"),
547            ProviderConnectionConfig { inference_profile_arn: Some("arn:test".to_string()), ..Default::default() },
548        );
549
550        assert_eq!(provider.inference_profile_arn.as_deref(), Some("arn:test"));
551    }
552
553    #[test]
554    fn test_from_config_with_credentials() {
555        let provider =
556            BedrockProvider::from_config(Some(static_credentials()), None, ProviderConnectionConfig::default());
557        assert_eq!(provider.model, DEFAULT_MODEL);
558    }
559
560    #[test]
561    fn test_from_config_with_credentials_and_region() {
562        let credentials =
563            AwsCredentials { session_token: Some("FwoGZXIvYXdzEBYaD...".to_string()), ..static_credentials() };
564
565        let provider =
566            BedrockProvider::from_config(Some(credentials), Some("us-west-2"), ProviderConnectionConfig::default())
567                .with_model("anthropic.claude-opus-4-20250514-v1:0");
568
569        assert_eq!(provider.model, "anthropic.claude-opus-4-20250514-v1:0");
570    }
571
572    #[test]
573    fn test_from_config_with_region_only() {
574        let provider = BedrockProvider::from_config(None, Some("eu-west-1"), ProviderConnectionConfig::default());
575        assert_eq!(provider.model, DEFAULT_MODEL);
576    }
577
578    #[test]
579    fn catalog_foundation_id_resolves_context_window() {
580        let provider = test_provider().with_model("anthropic.claude-sonnet-4-5-20250929-v1:0");
581        assert!(provider.context_window().is_some());
582        assert_eq!(provider.model().unwrap().to_string(), "bedrock:anthropic.claude-sonnet-4-5-20250929-v1:0");
583    }
584
585    #[test]
586    fn cross_region_profile_id_in_catalog_resolves() {
587        let provider = test_provider().with_model("us.anthropic.claude-opus-4-6-v1");
588        assert!(provider.context_window().is_some());
589    }
590
591    #[test]
592    fn unknown_cross_region_profile_id_falls_through_to_profile() {
593        let id = "us.anthropic.claude-future-model-v99:0";
594        let provider = test_provider().with_model(id);
595        assert_eq!(provider.context_window(), None);
596        assert_eq!(provider.model().unwrap().to_string(), format!("bedrock:{id}"));
597        assert_eq!(provider.display_name(), format!("AWS Bedrock ({id})"));
598    }
599
600    #[tokio::test]
601    async fn separate_inference_profile_arn_is_used_as_request_model_id() {
602        let mut server = CaptureServer::start_with_response("{}").await;
603        let provider = BedrockProvider::new(ProviderConnectionConfig {
604            base_url: Some(server.base_url.clone()),
605            auth_mode: ProviderAuthMode::None,
606            inference_profile_arn: Some(application_inference_profile_arn().to_string()),
607            ..Default::default()
608        })
609        .await
610        .with_model(DEFAULT_MODEL);
611
612        let _ = provider.stream_response(&hello_context()).collect::<Vec<_>>().await;
613        let request = server.captured().await;
614
615        assert!(
616            request.path.contains("arn%3Aaws%3Abedrock%3Aus-west-2%3A000000000000%3Aapplication-inference-profile"),
617            "{}",
618            request.path
619        );
620        assert_eq!(provider.context_window(), Some(200_000));
621        assert_eq!(provider.model().unwrap().to_string(), "bedrock:anthropic.claude-sonnet-4-5-20250929-v1:0");
622    }
623
624    #[test]
625    fn with_inference_profile_arn_keeps_canonical_model_identity() {
626        let arn = inference_profile_arn("us.anthropic.claude-sonnet-4-5-20250929-v1:0");
627        let provider =
628            test_provider().with_model("anthropic.claude-sonnet-4-5-20250929-v1:0").with_inference_profile_arn(&arn);
629
630        assert_eq!(provider.request_model_id(), arn);
631        assert_eq!(provider.context_window(), Some(200_000));
632        assert_eq!(provider.model().unwrap().to_string(), "bedrock:anthropic.claude-sonnet-4-5-20250929-v1:0");
633    }
634
635    #[test]
636    fn prompt_caching_support_comes_from_canonical_model() {
637        let cached = test_provider().with_model("anthropic.claude-sonnet-4-5-20250929-v1:0");
638        assert!(cached.model().unwrap().supports_prompt_caching());
639
640        let unknown_profile = test_provider().with_model("us.anthropic.claude-future-model-v99:0");
641        assert!(!unknown_profile.model().unwrap().supports_prompt_caching());
642    }
643
644    fn inference_profile_arn(model: &str) -> String {
645        format!("arn:aws:bedrock:us-west-2:000000000000:inference-profile/{model}")
646    }
647
648    fn application_inference_profile_arn() -> &'static str {
649        "arn:aws:bedrock:us-west-2:000000000000:application-inference-profile/000000000000"
650    }
651
652    fn test_provider() -> BedrockProvider {
653        BedrockProvider::from_config(None, None, ProviderConnectionConfig::default())
654    }
655
656    fn mantle_model() -> String {
657        LlmModel::all()
658            .iter()
659            .find(|model| model.provider_enum() == Provider::Bedrock && model.transport().is_some())
660            .expect("catalog must expose at least one Responses-transport Bedrock model")
661            .model_id()
662            .to_string()
663    }
664
665    async fn mantle_provider(server: &CaptureServer) -> BedrockProvider {
666        create_provider(&server.base_url, &mantle_model()).await
667    }
668
669    async fn create_provider(base_url: &str, model: &str) -> BedrockProvider {
670        BedrockProvider::new(ProviderConnectionConfig {
671            base_url: Some(base_url.to_string()),
672            auth_mode: ProviderAuthMode::None,
673            ..Default::default()
674        })
675        .await
676        .with_model(model)
677    }
678}