llm/providers/bedrock/
provider.rs

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

llm/providers/bedrock/
provider.rs
