llm/providers/bedrock/
provider.rs

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

llm/providers/bedrock/
provider.rs
