1use std::pin::Pin;
2use std::sync::Arc;
3
4use anyhow::{Context as _, Result, anyhow};
5use aws_config::stalled_stream_protection::StalledStreamProtectionConfig;
6use aws_config::{BehaviorVersion, Region};
7use aws_credential_types::{Credentials, Token};
8use aws_http_client::AwsHttpClient;
9use bedrock::bedrock_client::Client as BedrockClient;
10use bedrock::bedrock_client::config::timeout::TimeoutConfig;
11use bedrock::bedrock_client::types::{
12 CachePointBlock, CachePointType, ContentBlockDelta, ContentBlockStart, ConverseStreamOutput,
13 ReasoningContentBlockDelta, StopReason,
14};
15use bedrock::{
16 BedrockAnyToolChoice, BedrockAutoToolChoice, BedrockBlob, BedrockError, BedrockImageBlock,
17 BedrockImageFormat, BedrockImageSource, BedrockInnerContent, BedrockMessage, BedrockModelMode,
18 BedrockStreamingResponse, BedrockThinkingBlock, BedrockThinkingTextBlock, BedrockTool,
19 BedrockToolChoice, BedrockToolConfig, BedrockToolInputSchema, BedrockToolResultBlock,
20 BedrockToolResultContentBlock, BedrockToolResultStatus, BedrockToolSpec, BedrockToolUseBlock,
21 Model, value_to_aws_document,
22};
23use collections::{BTreeMap, HashMap};
24use credentials_provider::CredentialsProvider;
25use futures::{FutureExt, Stream, StreamExt, future::BoxFuture, stream::BoxStream};
26use gpui::{
27 AnyView, App, AsyncApp, Context, Entity, FocusHandle, Subscription, Task, Window, actions,
28};
29use gpui_tokio::Tokio;
30use http_client::HttpClient;
31use language_model::{
32 AuthenticateError, EnvVar, IconOrSvg, LanguageModel, LanguageModelCacheConfiguration,
33 LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelId, LanguageModelName,
34 LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName,
35 LanguageModelProviderState, LanguageModelRequest, LanguageModelToolChoice,
36 LanguageModelToolResultContent, LanguageModelToolUse, MessageContent, RateLimiter, Role,
37 TokenUsage, env_var,
38};
39use schemars::JsonSchema;
40use serde::{Deserialize, Serialize};
41use serde_json::Value;
42use settings::{BedrockAvailableModel as AvailableModel, Settings, SettingsStore};
43use smol::lock::OnceCell;
44use std::sync::LazyLock;
45use strum::{EnumIter, IntoEnumIterator, IntoStaticStr};
46use ui::{ButtonLink, ConfiguredApiCard, Divider, List, ListBulletItem, prelude::*};
47use ui_input::InputField;
48use util::ResultExt;
49
50use crate::AllLanguageModelSettings;
51use crate::provider::util::parse_tool_arguments;
52
53actions!(bedrock, [Tab, TabPrev]);
54
55const PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("amazon-bedrock");
56const PROVIDER_NAME: LanguageModelProviderName = LanguageModelProviderName::new("Amazon Bedrock");
57
58/// Credentials stored in the keychain for static authentication.
59/// Region is handled separately since it's orthogonal to auth method.
60#[derive(Default, Clone, Deserialize, Serialize, PartialEq, Debug)]
61pub struct BedrockCredentials {
62 pub access_key_id: String,
63 pub secret_access_key: String,
64 pub session_token: Option<String>,
65 pub bearer_token: Option<String>,
66}
67
68/// Resolved authentication configuration for Bedrock.
69/// Settings take priority over UX-provided credentials.
70#[derive(Clone, Debug, PartialEq)]
71pub enum BedrockAuth {
72 /// Use default AWS credential provider chain (IMDSv2, PodIdentity, env vars, etc.)
73 Automatic,
74 /// Use AWS named profile from ~/.aws/credentials or ~/.aws/config
75 NamedProfile { profile_name: String },
76 /// Use AWS SSO profile
77 SingleSignOn { profile_name: String },
78 /// Use IAM credentials (access key + secret + optional session token)
79 IamCredentials {
80 access_key_id: String,
81 secret_access_key: String,
82 session_token: Option<String>,
83 },
84 /// Use Bedrock API Key (bearer token authentication)
85 ApiKey { api_key: String },
86}
87
88impl BedrockCredentials {
89 /// Convert stored credentials to the appropriate auth variant.
90 /// Prefers API key if present, otherwise uses IAM credentials.
91 fn into_auth(self) -> Option<BedrockAuth> {
92 if let Some(api_key) = self.bearer_token.filter(|t| !t.is_empty()) {
93 Some(BedrockAuth::ApiKey { api_key })
94 } else if !self.access_key_id.is_empty() && !self.secret_access_key.is_empty() {
95 Some(BedrockAuth::IamCredentials {
96 access_key_id: self.access_key_id,
97 secret_access_key: self.secret_access_key,
98 session_token: self.session_token.filter(|t| !t.is_empty()),
99 })
100 } else {
101 None
102 }
103 }
104}
105
106#[derive(Default, Clone, Debug, PartialEq)]
107pub struct AmazonBedrockSettings {
108 pub available_models: Vec<AvailableModel>,
109 pub region: Option<String>,
110 pub endpoint: Option<String>,
111 pub profile_name: Option<String>,
112 pub role_arn: Option<String>,
113 pub authentication_method: Option<BedrockAuthMethod>,
114 pub allow_global: Option<bool>,
115 pub allow_extended_context: Option<bool>,
116}
117
118#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, EnumIter, IntoStaticStr, JsonSchema)]
119pub enum BedrockAuthMethod {
120 #[serde(rename = "named_profile")]
121 NamedProfile,
122 #[serde(rename = "sso")]
123 SingleSignOn,
124 #[serde(rename = "api_key")]
125 ApiKey,
126 /// IMDSv2, PodIdentity, env vars, etc.
127 #[serde(rename = "default")]
128 Automatic,
129}
130
131impl From<settings::BedrockAuthMethodContent> for BedrockAuthMethod {
132 fn from(value: settings::BedrockAuthMethodContent) -> Self {
133 match value {
134 settings::BedrockAuthMethodContent::SingleSignOn => BedrockAuthMethod::SingleSignOn,
135 settings::BedrockAuthMethodContent::Automatic => BedrockAuthMethod::Automatic,
136 settings::BedrockAuthMethodContent::NamedProfile => BedrockAuthMethod::NamedProfile,
137 settings::BedrockAuthMethodContent::ApiKey => BedrockAuthMethod::ApiKey,
138 }
139 }
140}
141
142#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize, JsonSchema)]
143#[serde(tag = "type", rename_all = "lowercase")]
144pub enum ModelMode {
145 #[default]
146 Default,
147 Thinking {
148 /// The maximum number of tokens to use for reasoning. Must be lower than the model's `max_output_tokens`.
149 budget_tokens: Option<u64>,
150 },
151 AdaptiveThinking {
152 effort: bedrock::BedrockAdaptiveThinkingEffort,
153 },
154}
155
156impl From<ModelMode> for BedrockModelMode {
157 fn from(value: ModelMode) -> Self {
158 match value {
159 ModelMode::Default => BedrockModelMode::Default,
160 ModelMode::Thinking { budget_tokens } => BedrockModelMode::Thinking { budget_tokens },
161 ModelMode::AdaptiveThinking { effort } => BedrockModelMode::AdaptiveThinking { effort },
162 }
163 }
164}
165
166impl From<BedrockModelMode> for ModelMode {
167 fn from(value: BedrockModelMode) -> Self {
168 match value {
169 BedrockModelMode::Default => ModelMode::Default,
170 BedrockModelMode::Thinking { budget_tokens } => ModelMode::Thinking { budget_tokens },
171 BedrockModelMode::AdaptiveThinking { effort } => ModelMode::AdaptiveThinking { effort },
172 }
173 }
174}
175
176/// The URL of the base AWS service.
177///
178/// Right now we're just using this as the key to store the AWS credentials
179/// under in the keychain.
180const AMAZON_AWS_URL: &str = "https://amazonaws.com";
181
182// These environment variables all use a `ZED_` prefix because we don't want to overwrite the user's AWS credentials.
183static ZED_BEDROCK_ACCESS_KEY_ID_VAR: LazyLock<EnvVar> = env_var!("ZED_ACCESS_KEY_ID");
184static ZED_BEDROCK_SECRET_ACCESS_KEY_VAR: LazyLock<EnvVar> = env_var!("ZED_SECRET_ACCESS_KEY");
185static ZED_BEDROCK_SESSION_TOKEN_VAR: LazyLock<EnvVar> = env_var!("ZED_SESSION_TOKEN");
186static ZED_AWS_PROFILE_VAR: LazyLock<EnvVar> = env_var!("ZED_AWS_PROFILE");
187static ZED_BEDROCK_REGION_VAR: LazyLock<EnvVar> = env_var!("ZED_AWS_REGION");
188static ZED_AWS_ENDPOINT_VAR: LazyLock<EnvVar> = env_var!("ZED_AWS_ENDPOINT");
189static ZED_BEDROCK_BEARER_TOKEN_VAR: LazyLock<EnvVar> = env_var!("ZED_BEDROCK_BEARER_TOKEN");
190
191pub struct State {
192 /// The resolved authentication method. Settings take priority over UX credentials.
193 auth: Option<BedrockAuth>,
194 /// Raw settings from settings.json
195 settings: Option<AmazonBedrockSettings>,
196 /// Whether credentials came from environment variables (only relevant for static credentials)
197 credentials_from_env: bool,
198 _subscription: Subscription,
199}
200
201impl State {
202 fn reset_auth(&self, cx: &mut Context<Self>) -> Task<Result<()>> {
203 let credentials_provider = <dyn CredentialsProvider>::global(cx);
204 cx.spawn(async move |this, cx| {
205 credentials_provider
206 .delete_credentials(AMAZON_AWS_URL, cx)
207 .await
208 .log_err();
209 this.update(cx, |this, cx| {
210 this.auth = None;
211 this.credentials_from_env = false;
212 cx.notify();
213 })
214 })
215 }
216
217 fn set_static_credentials(
218 &mut self,
219 credentials: BedrockCredentials,
220 cx: &mut Context<Self>,
221 ) -> Task<Result<()>> {
222 let auth = credentials.clone().into_auth();
223 let credentials_provider = <dyn CredentialsProvider>::global(cx);
224 cx.spawn(async move |this, cx| {
225 credentials_provider
226 .write_credentials(
227 AMAZON_AWS_URL,
228 "Bearer",
229 &serde_json::to_vec(&credentials)?,
230 cx,
231 )
232 .await?;
233 this.update(cx, |this, cx| {
234 this.auth = auth;
235 this.credentials_from_env = false;
236 cx.notify();
237 })
238 })
239 }
240
241 fn is_authenticated(&self) -> bool {
242 self.auth.is_some()
243 }
244
245 /// Resolve authentication. Settings take priority over UX-provided credentials.
246 fn authenticate(&self, cx: &mut Context<Self>) -> Task<Result<(), AuthenticateError>> {
247 if self.is_authenticated() {
248 return Task::ready(Ok(()));
249 }
250
251 // Step 1: Check if settings specify an auth method (enterprise control)
252 if let Some(settings) = &self.settings {
253 if let Some(method) = &settings.authentication_method {
254 let profile_name = settings
255 .profile_name
256 .clone()
257 .unwrap_or_else(|| "default".to_string());
258
259 let auth = match method {
260 BedrockAuthMethod::Automatic => BedrockAuth::Automatic,
261 BedrockAuthMethod::NamedProfile => BedrockAuth::NamedProfile { profile_name },
262 BedrockAuthMethod::SingleSignOn => BedrockAuth::SingleSignOn { profile_name },
263 BedrockAuthMethod::ApiKey => {
264 // ApiKey method means "use static credentials from keychain/env"
265 // Fall through to load them below
266 return self.load_static_credentials(cx);
267 }
268 };
269
270 return cx.spawn(async move |this, cx| {
271 this.update(cx, |this, cx| {
272 this.auth = Some(auth);
273 this.credentials_from_env = false;
274 cx.notify();
275 })?;
276 Ok(())
277 });
278 }
279 }
280
281 // Step 2: No settings auth method - try to load static credentials
282 self.load_static_credentials(cx)
283 }
284
285 /// Load static credentials from environment variables or keychain.
286 fn load_static_credentials(
287 &self,
288 cx: &mut Context<Self>,
289 ) -> Task<Result<(), AuthenticateError>> {
290 let credentials_provider = <dyn CredentialsProvider>::global(cx);
291 cx.spawn(async move |this, cx| {
292 // Try environment variables first
293 let (auth, from_env) = if let Some(bearer_token) = &ZED_BEDROCK_BEARER_TOKEN_VAR.value {
294 if !bearer_token.is_empty() {
295 (
296 Some(BedrockAuth::ApiKey {
297 api_key: bearer_token.to_string(),
298 }),
299 true,
300 )
301 } else {
302 (None, false)
303 }
304 } else if let Some(access_key_id) = &ZED_BEDROCK_ACCESS_KEY_ID_VAR.value {
305 if let Some(secret_access_key) = &ZED_BEDROCK_SECRET_ACCESS_KEY_VAR.value {
306 if !access_key_id.is_empty() && !secret_access_key.is_empty() {
307 let session_token = ZED_BEDROCK_SESSION_TOKEN_VAR
308 .value
309 .as_deref()
310 .filter(|s| !s.is_empty())
311 .map(|s| s.to_string());
312 (
313 Some(BedrockAuth::IamCredentials {
314 access_key_id: access_key_id.to_string(),
315 secret_access_key: secret_access_key.to_string(),
316 session_token,
317 }),
318 true,
319 )
320 } else {
321 (None, false)
322 }
323 } else {
324 (None, false)
325 }
326 } else {
327 (None, false)
328 };
329
330 // If we got auth from env vars, use it
331 if let Some(auth) = auth {
332 this.update(cx, |this, cx| {
333 this.auth = Some(auth);
334 this.credentials_from_env = from_env;
335 cx.notify();
336 })?;
337 return Ok(());
338 }
339
340 // Try keychain
341 let (_, credentials_bytes) = credentials_provider
342 .read_credentials(AMAZON_AWS_URL, cx)
343 .await?
344 .ok_or(AuthenticateError::CredentialsNotFound)?;
345
346 let credentials_str = String::from_utf8(credentials_bytes)
347 .context("invalid {PROVIDER_NAME} credentials")?;
348
349 let credentials: BedrockCredentials =
350 serde_json::from_str(&credentials_str).context("failed to parse credentials")?;
351
352 let auth = credentials
353 .into_auth()
354 .ok_or(AuthenticateError::CredentialsNotFound)?;
355
356 this.update(cx, |this, cx| {
357 this.auth = Some(auth);
358 this.credentials_from_env = false;
359 cx.notify();
360 })?;
361
362 Ok(())
363 })
364 }
365
366 /// Get the resolved region. Checks env var, then settings, then defaults to us-east-1.
367 fn get_region(&self) -> String {
368 // Priority: env var > settings > default
369 if let Some(region) = ZED_BEDROCK_REGION_VAR.value.as_deref() {
370 if !region.is_empty() {
371 return region.to_string();
372 }
373 }
374
375 self.settings
376 .as_ref()
377 .and_then(|s| s.region.clone())
378 .unwrap_or_else(|| "us-east-1".to_string())
379 }
380
381 fn get_allow_global(&self) -> bool {
382 self.settings
383 .as_ref()
384 .and_then(|s| s.allow_global)
385 .unwrap_or(false)
386 }
387
388 fn get_allow_extended_context(&self) -> bool {
389 self.settings
390 .as_ref()
391 .and_then(|s| s.allow_extended_context)
392 .unwrap_or(false)
393 }
394}
395
396pub struct BedrockLanguageModelProvider {
397 http_client: AwsHttpClient,
398 handle: tokio::runtime::Handle,
399 state: Entity<State>,
400}
401
402impl BedrockLanguageModelProvider {
403 pub fn new(http_client: Arc<dyn HttpClient>, cx: &mut App) -> Self {
404 let state = cx.new(|cx| State {
405 auth: None,
406 settings: Some(AllLanguageModelSettings::get_global(cx).bedrock.clone()),
407 credentials_from_env: false,
408 _subscription: cx.observe_global::<SettingsStore>(|_, cx| {
409 cx.notify();
410 }),
411 });
412
413 Self {
414 http_client: AwsHttpClient::new(http_client),
415 handle: Tokio::handle(cx),
416 state,
417 }
418 }
419
420 fn create_language_model(&self, model: bedrock::Model) -> Arc<dyn LanguageModel> {
421 Arc::new(BedrockModel {
422 id: LanguageModelId::from(model.id().to_string()),
423 model,
424 http_client: self.http_client.clone(),
425 handle: self.handle.clone(),
426 state: self.state.clone(),
427 client: OnceCell::new(),
428 request_limiter: RateLimiter::new(4),
429 })
430 }
431}
432
433impl LanguageModelProvider for BedrockLanguageModelProvider {
434 fn id(&self) -> LanguageModelProviderId {
435 PROVIDER_ID
436 }
437
438 fn name(&self) -> LanguageModelProviderName {
439 PROVIDER_NAME
440 }
441
442 fn icon(&self) -> IconOrSvg {
443 IconOrSvg::Icon(IconName::AiBedrock)
444 }
445
446 fn default_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {
447 Some(self.create_language_model(bedrock::Model::default()))
448 }
449
450 fn default_fast_model(&self, cx: &App) -> Option<Arc<dyn LanguageModel>> {
451 let region = self.state.read(cx).get_region();
452 Some(self.create_language_model(bedrock::Model::default_fast(region.as_str())))
453 }
454
455 fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> {
456 let mut models = BTreeMap::default();
457
458 for model in bedrock::Model::iter() {
459 if !matches!(model, bedrock::Model::Custom { .. }) {
460 models.insert(model.id().to_string(), model);
461 }
462 }
463
464 // Override with available models from settings
465 for model in AllLanguageModelSettings::get_global(cx)
466 .bedrock
467 .available_models
468 .iter()
469 {
470 models.insert(
471 model.name.clone(),
472 bedrock::Model::Custom {
473 name: model.name.clone(),
474 display_name: model.display_name.clone(),
475 max_tokens: model.max_tokens,
476 max_output_tokens: model.max_output_tokens,
477 default_temperature: model.default_temperature,
478 cache_configuration: model.cache_configuration.as_ref().map(|config| {
479 bedrock::BedrockModelCacheConfiguration {
480 max_cache_anchors: config.max_cache_anchors,
481 min_total_token: config.min_total_token,
482 }
483 }),
484 },
485 );
486 }
487
488 models
489 .into_values()
490 .map(|model| self.create_language_model(model))
491 .collect()
492 }
493
494 fn is_authenticated(&self, cx: &App) -> bool {
495 self.state.read(cx).is_authenticated()
496 }
497
498 fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>> {
499 self.state.update(cx, |state, cx| state.authenticate(cx))
500 }
501
502 fn configuration_view(
503 &self,
504 _target_agent: language_model::ConfigurationViewTargetAgent,
505 window: &mut Window,
506 cx: &mut App,
507 ) -> AnyView {
508 cx.new(|cx| ConfigurationView::new(self.state.clone(), window, cx))
509 .into()
510 }
511
512 fn reset_credentials(&self, cx: &mut App) -> Task<Result<()>> {
513 self.state.update(cx, |state, cx| state.reset_auth(cx))
514 }
515}
516
517impl LanguageModelProviderState for BedrockLanguageModelProvider {
518 type ObservableEntity = State;
519
520 fn observable_entity(&self) -> Option<Entity<Self::ObservableEntity>> {
521 Some(self.state.clone())
522 }
523}
524
525struct BedrockModel {
526 id: LanguageModelId,
527 model: Model,
528 http_client: AwsHttpClient,
529 handle: tokio::runtime::Handle,
530 client: OnceCell<BedrockClient>,
531 state: Entity<State>,
532 request_limiter: RateLimiter,
533}
534
535impl BedrockModel {
536 fn get_or_init_client(&self, cx: &AsyncApp) -> anyhow::Result<&BedrockClient> {
537 self.client
538 .get_or_try_init_blocking(|| {
539 let (auth, endpoint, region) = cx.read_entity(&self.state, |state, _cx| {
540 let endpoint = state.settings.as_ref().and_then(|s| s.endpoint.clone());
541 let region = state.get_region();
542 (state.auth.clone(), endpoint, region)
543 });
544
545 let mut config_builder = aws_config::defaults(BehaviorVersion::latest())
546 .stalled_stream_protection(StalledStreamProtectionConfig::disabled())
547 .http_client(self.http_client.clone())
548 .region(Region::new(region))
549 .timeout_config(TimeoutConfig::disabled());
550
551 if let Some(endpoint_url) = endpoint
552 && !endpoint_url.is_empty()
553 {
554 config_builder = config_builder.endpoint_url(endpoint_url);
555 }
556
557 match auth {
558 Some(BedrockAuth::Automatic) | None => {
559 // Use default AWS credential provider chain
560 }
561 Some(BedrockAuth::NamedProfile { profile_name })
562 | Some(BedrockAuth::SingleSignOn { profile_name }) => {
563 if !profile_name.is_empty() {
564 config_builder = config_builder.profile_name(profile_name);
565 }
566 }
567 Some(BedrockAuth::IamCredentials {
568 access_key_id,
569 secret_access_key,
570 session_token,
571 }) => {
572 let aws_creds = Credentials::new(
573 access_key_id,
574 secret_access_key,
575 session_token,
576 None,
577 "zed-bedrock-provider",
578 );
579 config_builder = config_builder.credentials_provider(aws_creds);
580 }
581 Some(BedrockAuth::ApiKey { api_key }) => {
582 config_builder = config_builder
583 .auth_scheme_preference(["httpBearerAuth".into()]) // https://github.com/smithy-lang/smithy-rs/pull/4241
584 .token_provider(Token::new(api_key, None));
585 }
586 }
587
588 let config = self.handle.block_on(config_builder.load());
589
590 anyhow::Ok(BedrockClient::new(&config))
591 })
592 .context("initializing Bedrock client")?;
593
594 self.client.get().context("Bedrock client not initialized")
595 }
596
597 fn stream_completion(
598 &self,
599 request: bedrock::Request,
600 cx: &AsyncApp,
601 ) -> BoxFuture<
602 'static,
603 Result<BoxStream<'static, Result<BedrockStreamingResponse, anyhow::Error>>, BedrockError>,
604 > {
605 let Ok(runtime_client) = self
606 .get_or_init_client(cx)
607 .cloned()
608 .context("Bedrock client not initialized")
609 else {
610 return futures::future::ready(Err(BedrockError::Other(anyhow!("App state dropped"))))
611 .boxed();
612 };
613
614 let task = Tokio::spawn(cx, bedrock::stream_completion(runtime_client, request));
615 async move { task.await.map_err(|e| BedrockError::Other(e.into()))? }.boxed()
616 }
617}
618
619impl LanguageModel for BedrockModel {
620 fn id(&self) -> LanguageModelId {
621 self.id.clone()
622 }
623
624 fn name(&self) -> LanguageModelName {
625 LanguageModelName::from(self.model.display_name().to_string())
626 }
627
628 fn provider_id(&self) -> LanguageModelProviderId {
629 PROVIDER_ID
630 }
631
632 fn provider_name(&self) -> LanguageModelProviderName {
633 PROVIDER_NAME
634 }
635
636 fn supports_tools(&self) -> bool {
637 self.model.supports_tool_use()
638 }
639
640 fn supports_images(&self) -> bool {
641 self.model.supports_images()
642 }
643
644 fn supports_thinking(&self) -> bool {
645 matches!(
646 self.model.mode(),
647 BedrockModelMode::Thinking { .. } | BedrockModelMode::AdaptiveThinking { .. }
648 )
649 }
650
651 fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
652 match choice {
653 LanguageModelToolChoice::Auto | LanguageModelToolChoice::Any => {
654 self.model.supports_tool_use()
655 }
656 // Add support for None - we'll filter tool calls at response
657 LanguageModelToolChoice::None => self.model.supports_tool_use(),
658 }
659 }
660
661 fn supports_streaming_tools(&self) -> bool {
662 true
663 }
664
665 fn telemetry_id(&self) -> String {
666 format!("bedrock/{}", self.model.id())
667 }
668
669 fn max_token_count(&self) -> u64 {
670 self.model.max_token_count()
671 }
672
673 fn max_output_tokens(&self) -> Option<u64> {
674 Some(self.model.max_output_tokens())
675 }
676
677 fn count_tokens(
678 &self,
679 request: LanguageModelRequest,
680 cx: &App,
681 ) -> BoxFuture<'static, Result<u64>> {
682 get_bedrock_tokens(request, cx)
683 }
684
685 fn stream_completion(
686 &self,
687 request: LanguageModelRequest,
688 cx: &AsyncApp,
689 ) -> BoxFuture<
690 'static,
691 Result<
692 BoxStream<'static, Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>,
693 LanguageModelCompletionError,
694 >,
695 > {
696 let (region, allow_global, allow_extended_context) =
697 cx.read_entity(&self.state, |state, _cx| {
698 (
699 state.get_region(),
700 state.get_allow_global(),
701 state.get_allow_extended_context(),
702 )
703 });
704
705 let model_id = match self.model.cross_region_inference_id(®ion, allow_global) {
706 Ok(s) => s,
707 Err(e) => {
708 return async move { Err(e.into()) }.boxed();
709 }
710 };
711
712 let deny_tool_calls = request.tool_choice == Some(LanguageModelToolChoice::None);
713
714 let use_extended_context = allow_extended_context && self.model.supports_extended_context();
715
716 let request = match into_bedrock(
717 request,
718 model_id,
719 self.model.default_temperature(),
720 self.model.max_output_tokens(),
721 self.model.mode(),
722 self.model.supports_caching(),
723 self.model.supports_tool_use(),
724 use_extended_context,
725 ) {
726 Ok(request) => request,
727 Err(err) => return futures::future::ready(Err(err.into())).boxed(),
728 };
729
730 let request = self.stream_completion(request, cx);
731 let display_name = self.model.display_name().to_string();
732 let future = self.request_limiter.stream(async move {
733 let response = request.await.map_err(|err| match err {
734 BedrockError::Validation(ref msg) => {
735 if msg.contains("model identifier is invalid") {
736 LanguageModelCompletionError::Other(anyhow!(
737 "{display_name} is not available in {region}. \
738 Try switching to a region where this model is supported."
739 ))
740 } else {
741 LanguageModelCompletionError::BadRequestFormat {
742 provider: PROVIDER_NAME,
743 message: msg.clone(),
744 }
745 }
746 }
747 BedrockError::RateLimited => LanguageModelCompletionError::RateLimitExceeded {
748 provider: PROVIDER_NAME,
749 retry_after: None,
750 },
751 BedrockError::ServiceUnavailable => {
752 LanguageModelCompletionError::ServerOverloaded {
753 provider: PROVIDER_NAME,
754 retry_after: None,
755 }
756 }
757 BedrockError::AccessDenied(msg) => LanguageModelCompletionError::PermissionError {
758 provider: PROVIDER_NAME,
759 message: msg,
760 },
761 BedrockError::InternalServer(msg) => {
762 LanguageModelCompletionError::ApiInternalServerError {
763 provider: PROVIDER_NAME,
764 message: msg,
765 }
766 }
767 other => LanguageModelCompletionError::Other(anyhow!(other)),
768 })?;
769 let events = map_to_language_model_completion_events(response);
770
771 if deny_tool_calls {
772 Ok(deny_tool_use_events(events).boxed())
773 } else {
774 Ok(events.boxed())
775 }
776 });
777
778 async move { Ok(future.await?.boxed()) }.boxed()
779 }
780
781 fn cache_configuration(&self) -> Option<LanguageModelCacheConfiguration> {
782 self.model
783 .cache_configuration()
784 .map(|config| LanguageModelCacheConfiguration {
785 max_cache_anchors: config.max_cache_anchors,
786 should_speculate: false,
787 min_total_token: config.min_total_token,
788 })
789 }
790}
791
792fn deny_tool_use_events(
793 events: impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>,
794) -> impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>> {
795 events.map(|event| {
796 match event {
797 Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) => {
798 // Convert tool use to an error message if model decided to call it
799 Ok(LanguageModelCompletionEvent::Text(format!(
800 "\n\n[Error: Tool calls are disabled in this context. Attempted to call '{}']",
801 tool_use.name
802 )))
803 }
804 other => other,
805 }
806 })
807}
808
809pub fn into_bedrock(
810 request: LanguageModelRequest,
811 model: String,
812 default_temperature: f32,
813 max_output_tokens: u64,
814 mode: BedrockModelMode,
815 supports_caching: bool,
816 supports_tool_use: bool,
817 allow_extended_context: bool,
818) -> Result<bedrock::Request> {
819 let mut new_messages: Vec<BedrockMessage> = Vec::new();
820 let mut system_message = String::new();
821
822 // Track whether messages contain tool content - Bedrock requires toolConfig
823 // when tool blocks are present, so we may need to add a dummy tool
824 let mut messages_contain_tool_content = false;
825
826 for message in request.messages {
827 if message.contents_empty() {
828 continue;
829 }
830
831 match message.role {
832 Role::User | Role::Assistant => {
833 let mut bedrock_message_content: Vec<BedrockInnerContent> = message
834 .content
835 .into_iter()
836 .filter_map(|content| match content {
837 MessageContent::Text(text) => {
838 if !text.is_empty() {
839 Some(BedrockInnerContent::Text(text))
840 } else {
841 None
842 }
843 }
844 MessageContent::Thinking { text, signature } => {
845 if model.contains(Model::DeepSeekR1.request_id()) {
846 // DeepSeekR1 doesn't support thinking blocks
847 // And the AWS API demands that you strip them
848 return None;
849 }
850 if signature.is_none() {
851 // Thinking blocks without a signature are invalid
852 // (e.g. from cancellation mid-think) and must be
853 // stripped to avoid API errors.
854 return None;
855 }
856 let thinking = BedrockThinkingTextBlock::builder()
857 .text(text)
858 .set_signature(signature)
859 .build()
860 .context("failed to build reasoning block")
861 .log_err()?;
862
863 Some(BedrockInnerContent::ReasoningContent(
864 BedrockThinkingBlock::ReasoningText(thinking),
865 ))
866 }
867 MessageContent::RedactedThinking(blob) => {
868 if model.contains(Model::DeepSeekR1.request_id()) {
869 // DeepSeekR1 doesn't support thinking blocks
870 // And the AWS API demands that you strip them
871 return None;
872 }
873 let redacted =
874 BedrockThinkingBlock::RedactedContent(BedrockBlob::new(blob));
875
876 Some(BedrockInnerContent::ReasoningContent(redacted))
877 }
878 MessageContent::ToolUse(tool_use) => {
879 messages_contain_tool_content = true;
880 let input = if tool_use.input.is_null() {
881 // Bedrock API requires valid JsonValue, not null, for tool use input
882 value_to_aws_document(&serde_json::json!({}))
883 } else {
884 value_to_aws_document(&tool_use.input)
885 };
886 BedrockToolUseBlock::builder()
887 .name(tool_use.name.to_string())
888 .tool_use_id(tool_use.id.to_string())
889 .input(input)
890 .build()
891 .context("failed to build Bedrock tool use block")
892 .log_err()
893 .map(BedrockInnerContent::ToolUse)
894 }
895 MessageContent::ToolResult(tool_result) => {
896 messages_contain_tool_content = true;
897 BedrockToolResultBlock::builder()
898 .tool_use_id(tool_result.tool_use_id.to_string())
899 .content(match tool_result.content {
900 LanguageModelToolResultContent::Text(text) => {
901 BedrockToolResultContentBlock::Text(text.to_string())
902 }
903 LanguageModelToolResultContent::Image(image) => {
904 use base64::Engine;
905
906 match base64::engine::general_purpose::STANDARD
907 .decode(image.source.as_bytes())
908 {
909 Ok(image_bytes) => {
910 match BedrockImageBlock::builder()
911 .format(BedrockImageFormat::Png)
912 .source(BedrockImageSource::Bytes(
913 BedrockBlob::new(image_bytes),
914 ))
915 .build()
916 {
917 Ok(image_block) => {
918 BedrockToolResultContentBlock::Image(
919 image_block,
920 )
921 }
922 Err(err) => {
923 BedrockToolResultContentBlock::Text(
924 format!(
925 "[Failed to build image block: {}]",
926 err
927 ),
928 )
929 }
930 }
931 }
932 Err(err) => {
933 BedrockToolResultContentBlock::Text(format!(
934 "[Failed to decode tool result image: {}]",
935 err
936 ))
937 }
938 }
939 }
940 })
941 .status({
942 if tool_result.is_error {
943 BedrockToolResultStatus::Error
944 } else {
945 BedrockToolResultStatus::Success
946 }
947 })
948 .build()
949 .context("failed to build Bedrock tool result block")
950 .log_err()
951 .map(BedrockInnerContent::ToolResult)
952 }
953 MessageContent::Image(image) => {
954 use base64::Engine;
955
956 let image_bytes = base64::engine::general_purpose::STANDARD
957 .decode(image.source.as_bytes())
958 .context("failed to decode base64 image data")
959 .log_err()?;
960
961 BedrockImageBlock::builder()
962 .format(BedrockImageFormat::Png)
963 .source(BedrockImageSource::Bytes(BedrockBlob::new(image_bytes)))
964 .build()
965 .context("failed to build Bedrock image block")
966 .log_err()
967 .map(BedrockInnerContent::Image)
968 }
969 })
970 .collect();
971 if message.cache && supports_caching {
972 bedrock_message_content.push(BedrockInnerContent::CachePoint(
973 CachePointBlock::builder()
974 .r#type(CachePointType::Default)
975 .build()
976 .context("failed to build cache point block")?,
977 ));
978 }
979 let bedrock_role = match message.role {
980 Role::User => bedrock::BedrockRole::User,
981 Role::Assistant => bedrock::BedrockRole::Assistant,
982 Role::System => unreachable!("System role should never occur here"),
983 };
984 if bedrock_message_content.is_empty() {
985 continue;
986 }
987
988 if let Some(last_message) = new_messages.last_mut()
989 && last_message.role == bedrock_role
990 {
991 last_message.content.extend(bedrock_message_content);
992 continue;
993 }
994 new_messages.push(
995 BedrockMessage::builder()
996 .role(bedrock_role)
997 .set_content(Some(bedrock_message_content))
998 .build()
999 .context("failed to build Bedrock message")?,
1000 );
1001 }
1002 Role::System => {
1003 if !system_message.is_empty() {
1004 system_message.push_str("\n\n");
1005 }
1006 system_message.push_str(&message.string_contents());
1007 }
1008 }
1009 }
1010
1011 let mut tool_spec: Vec<BedrockTool> = if supports_tool_use {
1012 request
1013 .tools
1014 .iter()
1015 .filter_map(|tool| {
1016 Some(BedrockTool::ToolSpec(
1017 BedrockToolSpec::builder()
1018 .name(tool.name.clone())
1019 .description(tool.description.clone())
1020 .input_schema(BedrockToolInputSchema::Json(value_to_aws_document(
1021 &tool.input_schema,
1022 )))
1023 .build()
1024 .log_err()?,
1025 ))
1026 })
1027 .collect()
1028 } else {
1029 Vec::new()
1030 };
1031
1032 // Bedrock requires toolConfig when messages contain tool use/result blocks.
1033 // If no tools are defined but messages contain tool content (e.g., when
1034 // summarising a conversation that used tools), add a dummy tool to satisfy
1035 // the API requirement.
1036 if supports_tool_use && tool_spec.is_empty() && messages_contain_tool_content {
1037 tool_spec.push(BedrockTool::ToolSpec(
1038 BedrockToolSpec::builder()
1039 .name("_placeholder")
1040 .description("Placeholder tool to satisfy Bedrock API requirements when conversation history contains tool usage")
1041 .input_schema(BedrockToolInputSchema::Json(value_to_aws_document(
1042 &serde_json::json!({"type": "object", "properties": {}}),
1043 )))
1044 .build()
1045 .context("failed to build placeholder tool spec")?,
1046 ));
1047 }
1048
1049 if !tool_spec.is_empty() && supports_caching {
1050 tool_spec.push(BedrockTool::CachePoint(
1051 CachePointBlock::builder()
1052 .r#type(CachePointType::Default)
1053 .build()
1054 .context("failed to build cache point block")?,
1055 ));
1056 }
1057
1058 let tool_choice = match request.tool_choice {
1059 Some(LanguageModelToolChoice::Auto) | None => {
1060 BedrockToolChoice::Auto(BedrockAutoToolChoice::builder().build())
1061 }
1062 Some(LanguageModelToolChoice::Any) => {
1063 BedrockToolChoice::Any(BedrockAnyToolChoice::builder().build())
1064 }
1065 Some(LanguageModelToolChoice::None) => {
1066 // For None, we still use Auto but will filter out tool calls in the response
1067 BedrockToolChoice::Auto(BedrockAutoToolChoice::builder().build())
1068 }
1069 };
1070 let tool_config = if tool_spec.is_empty() {
1071 None
1072 } else {
1073 Some(
1074 BedrockToolConfig::builder()
1075 .set_tools(Some(tool_spec))
1076 .tool_choice(tool_choice)
1077 .build()?,
1078 )
1079 };
1080
1081 Ok(bedrock::Request {
1082 model,
1083 messages: new_messages,
1084 max_tokens: max_output_tokens,
1085 system: Some(system_message),
1086 tools: tool_config,
1087 thinking: if request.thinking_allowed {
1088 match mode {
1089 BedrockModelMode::Thinking { budget_tokens } => {
1090 Some(bedrock::Thinking::Enabled { budget_tokens })
1091 }
1092 BedrockModelMode::AdaptiveThinking { effort } => {
1093 Some(bedrock::Thinking::Adaptive { effort })
1094 }
1095 BedrockModelMode::Default => None,
1096 }
1097 } else {
1098 None
1099 },
1100 metadata: None,
1101 stop_sequences: Vec::new(),
1102 temperature: request.temperature.or(Some(default_temperature)),
1103 top_k: None,
1104 top_p: None,
1105 allow_extended_context,
1106 })
1107}
1108
1109// TODO: just call the ConverseOutput.usage() method:
1110// https://docs.rs/aws-sdk-bedrockruntime/latest/aws_sdk_bedrockruntime/operation/converse/struct.ConverseOutput.html#method.output
1111pub fn get_bedrock_tokens(
1112 request: LanguageModelRequest,
1113 cx: &App,
1114) -> BoxFuture<'static, Result<u64>> {
1115 cx.background_executor()
1116 .spawn(async move {
1117 let messages = request.messages;
1118 let mut tokens_from_images = 0;
1119 let mut string_messages = Vec::with_capacity(messages.len());
1120
1121 for message in messages {
1122 use language_model::MessageContent;
1123
1124 let mut string_contents = String::new();
1125
1126 for content in message.content {
1127 match content {
1128 MessageContent::Text(text) | MessageContent::Thinking { text, .. } => {
1129 string_contents.push_str(&text);
1130 }
1131 MessageContent::RedactedThinking(_) => {}
1132 MessageContent::Image(image) => {
1133 tokens_from_images += image.estimate_tokens();
1134 }
1135 MessageContent::ToolUse(_tool_use) => {
1136 // TODO: Estimate token usage from tool uses.
1137 }
1138 MessageContent::ToolResult(tool_result) => match tool_result.content {
1139 LanguageModelToolResultContent::Text(text) => {
1140 string_contents.push_str(&text);
1141 }
1142 LanguageModelToolResultContent::Image(image) => {
1143 tokens_from_images += image.estimate_tokens();
1144 }
1145 },
1146 }
1147 }
1148
1149 if !string_contents.is_empty() {
1150 string_messages.push(tiktoken_rs::ChatCompletionRequestMessage {
1151 role: match message.role {
1152 Role::User => "user".into(),
1153 Role::Assistant => "assistant".into(),
1154 Role::System => "system".into(),
1155 },
1156 content: Some(string_contents),
1157 name: None,
1158 function_call: None,
1159 });
1160 }
1161 }
1162
1163 // Tiktoken doesn't yet support these models, so we manually use the
1164 // same tokenizer as GPT-4.
1165 tiktoken_rs::num_tokens_from_messages("gpt-4", &string_messages)
1166 .map(|tokens| (tokens + tokens_from_images) as u64)
1167 })
1168 .boxed()
1169}
1170
1171pub fn map_to_language_model_completion_events(
1172 events: Pin<Box<dyn Send + Stream<Item = Result<BedrockStreamingResponse, anyhow::Error>>>>,
1173) -> impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>> {
1174 struct RawToolUse {
1175 id: String,
1176 name: String,
1177 input_json: String,
1178 }
1179
1180 struct State {
1181 events: Pin<Box<dyn Send + Stream<Item = Result<BedrockStreamingResponse, anyhow::Error>>>>,
1182 tool_uses_by_index: HashMap<i32, RawToolUse>,
1183 emitted_tool_use: bool,
1184 }
1185
1186 let initial_state = State {
1187 events,
1188 tool_uses_by_index: HashMap::default(),
1189 emitted_tool_use: false,
1190 };
1191
1192 futures::stream::unfold(initial_state, |mut state| async move {
1193 match state.events.next().await {
1194 Some(event_result) => match event_result {
1195 Ok(event) => {
1196 let result = match event {
1197 ConverseStreamOutput::ContentBlockDelta(cb_delta) => match cb_delta.delta {
1198 Some(ContentBlockDelta::Text(text)) => {
1199 Some(Ok(LanguageModelCompletionEvent::Text(text)))
1200 }
1201 Some(ContentBlockDelta::ToolUse(tool_output)) => {
1202 if let Some(tool_use) = state
1203 .tool_uses_by_index
1204 .get_mut(&cb_delta.content_block_index)
1205 {
1206 tool_use.input_json.push_str(tool_output.input());
1207 if let Ok(input) = serde_json::from_str::<serde_json::Value>(
1208 &partial_json_fixer::fix_json(&tool_use.input_json),
1209 ) {
1210 Some(Ok(LanguageModelCompletionEvent::ToolUse(
1211 LanguageModelToolUse {
1212 id: tool_use.id.clone().into(),
1213 name: tool_use.name.clone().into(),
1214 is_input_complete: false,
1215 raw_input: tool_use.input_json.clone(),
1216 input,
1217 thought_signature: None,
1218 },
1219 )))
1220 } else {
1221 None
1222 }
1223 } else {
1224 None
1225 }
1226 }
1227 Some(ContentBlockDelta::ReasoningContent(thinking)) => match thinking {
1228 ReasoningContentBlockDelta::Text(thoughts) => {
1229 Some(Ok(LanguageModelCompletionEvent::Thinking {
1230 text: thoughts,
1231 signature: None,
1232 }))
1233 }
1234 ReasoningContentBlockDelta::Signature(sig) => {
1235 Some(Ok(LanguageModelCompletionEvent::Thinking {
1236 text: "".into(),
1237 signature: Some(sig),
1238 }))
1239 }
1240 ReasoningContentBlockDelta::RedactedContent(redacted) => {
1241 let content = String::from_utf8(redacted.into_inner())
1242 .unwrap_or("REDACTED".to_string());
1243 Some(Ok(LanguageModelCompletionEvent::Thinking {
1244 text: content,
1245 signature: None,
1246 }))
1247 }
1248 _ => None,
1249 },
1250 _ => None,
1251 },
1252 ConverseStreamOutput::ContentBlockStart(cb_start) => {
1253 if let Some(ContentBlockStart::ToolUse(tool_start)) = cb_start.start {
1254 state.tool_uses_by_index.insert(
1255 cb_start.content_block_index,
1256 RawToolUse {
1257 id: tool_start.tool_use_id,
1258 name: tool_start.name,
1259 input_json: String::new(),
1260 },
1261 );
1262 }
1263 None
1264 }
1265 ConverseStreamOutput::MessageStart(_) => None,
1266 ConverseStreamOutput::ContentBlockStop(cb_stop) => state
1267 .tool_uses_by_index
1268 .remove(&cb_stop.content_block_index)
1269 .map(|tool_use| {
1270 state.emitted_tool_use = true;
1271
1272 let input = parse_tool_arguments(&tool_use.input_json)
1273 .unwrap_or_else(|_| Value::Object(Default::default()));
1274
1275 Ok(LanguageModelCompletionEvent::ToolUse(
1276 LanguageModelToolUse {
1277 id: tool_use.id.into(),
1278 name: tool_use.name.into(),
1279 is_input_complete: true,
1280 raw_input: tool_use.input_json,
1281 input,
1282 thought_signature: None,
1283 },
1284 ))
1285 }),
1286 ConverseStreamOutput::Metadata(cb_meta) => cb_meta.usage.map(|metadata| {
1287 Ok(LanguageModelCompletionEvent::UsageUpdate(TokenUsage {
1288 input_tokens: metadata.input_tokens as u64,
1289 output_tokens: metadata.output_tokens as u64,
1290 cache_creation_input_tokens: metadata
1291 .cache_write_input_tokens
1292 .unwrap_or_default()
1293 as u64,
1294 cache_read_input_tokens: metadata
1295 .cache_read_input_tokens
1296 .unwrap_or_default()
1297 as u64,
1298 }))
1299 }),
1300 ConverseStreamOutput::MessageStop(message_stop) => {
1301 let stop_reason = if state.emitted_tool_use {
1302 // Some models (e.g. Kimi) send EndTurn even when
1303 // they've made tool calls. Trust the content over
1304 // the stop reason.
1305 language_model::StopReason::ToolUse
1306 } else {
1307 match message_stop.stop_reason {
1308 StopReason::ToolUse => language_model::StopReason::ToolUse,
1309 _ => language_model::StopReason::EndTurn,
1310 }
1311 };
1312 Some(Ok(LanguageModelCompletionEvent::Stop(stop_reason)))
1313 }
1314 _ => None,
1315 };
1316
1317 Some((result, state))
1318 }
1319 Err(err) => Some((
1320 Some(Err(LanguageModelCompletionError::Other(anyhow!(err)))),
1321 state,
1322 )),
1323 },
1324 None => None,
1325 }
1326 })
1327 .filter_map(|result| async move { result })
1328}
1329
1330struct ConfigurationView {
1331 access_key_id_editor: Entity<InputField>,
1332 secret_access_key_editor: Entity<InputField>,
1333 session_token_editor: Entity<InputField>,
1334 bearer_token_editor: Entity<InputField>,
1335 state: Entity<State>,
1336 load_credentials_task: Option<Task<()>>,
1337 focus_handle: FocusHandle,
1338}
1339
1340impl ConfigurationView {
1341 const PLACEHOLDER_ACCESS_KEY_ID_TEXT: &'static str = "XXXXXXXXXXXXXXXX";
1342 const PLACEHOLDER_SECRET_ACCESS_KEY_TEXT: &'static str =
1343 "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX";
1344 const PLACEHOLDER_SESSION_TOKEN_TEXT: &'static str = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX";
1345 const PLACEHOLDER_BEARER_TOKEN_TEXT: &'static str = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX";
1346
1347 fn new(state: Entity<State>, window: &mut Window, cx: &mut Context<Self>) -> Self {
1348 let focus_handle = cx.focus_handle();
1349
1350 cx.observe(&state, |_, _, cx| {
1351 cx.notify();
1352 })
1353 .detach();
1354
1355 let access_key_id_editor = cx.new(|cx| {
1356 InputField::new(window, cx, Self::PLACEHOLDER_ACCESS_KEY_ID_TEXT)
1357 .label("Access Key ID")
1358 .tab_index(0)
1359 .tab_stop(true)
1360 });
1361
1362 let secret_access_key_editor = cx.new(|cx| {
1363 InputField::new(window, cx, Self::PLACEHOLDER_SECRET_ACCESS_KEY_TEXT)
1364 .label("Secret Access Key")
1365 .tab_index(1)
1366 .tab_stop(true)
1367 });
1368
1369 let session_token_editor = cx.new(|cx| {
1370 InputField::new(window, cx, Self::PLACEHOLDER_SESSION_TOKEN_TEXT)
1371 .label("Session Token (Optional)")
1372 .tab_index(2)
1373 .tab_stop(true)
1374 });
1375
1376 let bearer_token_editor = cx.new(|cx| {
1377 InputField::new(window, cx, Self::PLACEHOLDER_BEARER_TOKEN_TEXT)
1378 .label("Bedrock API Key")
1379 .tab_index(3)
1380 .tab_stop(true)
1381 });
1382
1383 let load_credentials_task = Some(cx.spawn({
1384 let state = state.clone();
1385 async move |this, cx| {
1386 if let Some(task) = Some(state.update(cx, |state, cx| state.authenticate(cx))) {
1387 // We don't log an error, because "not signed in" is also an error.
1388 let _ = task.await;
1389 }
1390 this.update(cx, |this, cx| {
1391 this.load_credentials_task = None;
1392 cx.notify();
1393 })
1394 .log_err();
1395 }
1396 }));
1397
1398 Self {
1399 access_key_id_editor,
1400 secret_access_key_editor,
1401 session_token_editor,
1402 bearer_token_editor,
1403 state,
1404 load_credentials_task,
1405 focus_handle,
1406 }
1407 }
1408
1409 fn save_credentials(
1410 &mut self,
1411 _: &menu::Confirm,
1412 _window: &mut Window,
1413 cx: &mut Context<Self>,
1414 ) {
1415 let access_key_id = self
1416 .access_key_id_editor
1417 .read(cx)
1418 .text(cx)
1419 .trim()
1420 .to_string();
1421 let secret_access_key = self
1422 .secret_access_key_editor
1423 .read(cx)
1424 .text(cx)
1425 .trim()
1426 .to_string();
1427 let session_token = self
1428 .session_token_editor
1429 .read(cx)
1430 .text(cx)
1431 .trim()
1432 .to_string();
1433 let session_token = if session_token.is_empty() {
1434 None
1435 } else {
1436 Some(session_token)
1437 };
1438 let bearer_token = self
1439 .bearer_token_editor
1440 .read(cx)
1441 .text(cx)
1442 .trim()
1443 .to_string();
1444 let bearer_token = if bearer_token.is_empty() {
1445 None
1446 } else {
1447 Some(bearer_token)
1448 };
1449
1450 let state = self.state.clone();
1451 cx.spawn(async move |_, cx| {
1452 state
1453 .update(cx, |state, cx| {
1454 let credentials = BedrockCredentials {
1455 access_key_id,
1456 secret_access_key,
1457 session_token,
1458 bearer_token,
1459 };
1460
1461 state.set_static_credentials(credentials, cx)
1462 })
1463 .await
1464 })
1465 .detach_and_log_err(cx);
1466 }
1467
1468 fn reset_credentials(&mut self, window: &mut Window, cx: &mut Context<Self>) {
1469 self.access_key_id_editor
1470 .update(cx, |editor, cx| editor.set_text("", window, cx));
1471 self.secret_access_key_editor
1472 .update(cx, |editor, cx| editor.set_text("", window, cx));
1473 self.session_token_editor
1474 .update(cx, |editor, cx| editor.set_text("", window, cx));
1475 self.bearer_token_editor
1476 .update(cx, |editor, cx| editor.set_text("", window, cx));
1477
1478 let state = self.state.clone();
1479 cx.spawn(async move |_, cx| state.update(cx, |state, cx| state.reset_auth(cx)).await)
1480 .detach_and_log_err(cx);
1481 }
1482
1483 fn should_render_editor(&self, cx: &Context<Self>) -> bool {
1484 self.state.read(cx).is_authenticated()
1485 }
1486
1487 fn on_tab(&mut self, _: &menu::SelectNext, window: &mut Window, cx: &mut Context<Self>) {
1488 window.focus_next(cx);
1489 }
1490
1491 fn on_tab_prev(
1492 &mut self,
1493 _: &menu::SelectPrevious,
1494 window: &mut Window,
1495 cx: &mut Context<Self>,
1496 ) {
1497 window.focus_prev(cx);
1498 }
1499}
1500
1501impl Render for ConfigurationView {
1502 fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
1503 let state = self.state.read(cx);
1504 let env_var_set = state.credentials_from_env;
1505 let auth = state.auth.clone();
1506 let settings_auth_method = state
1507 .settings
1508 .as_ref()
1509 .and_then(|s| s.authentication_method.clone());
1510
1511 if self.load_credentials_task.is_some() {
1512 return div().child(Label::new("Loading credentials...")).into_any();
1513 }
1514
1515 let configured_label = match &auth {
1516 Some(BedrockAuth::Automatic) => {
1517 "Using automatic credentials (AWS default chain)".into()
1518 }
1519 Some(BedrockAuth::NamedProfile { profile_name }) => {
1520 format!("Using AWS profile: {profile_name}")
1521 }
1522 Some(BedrockAuth::SingleSignOn { profile_name }) => {
1523 format!("Using AWS SSO profile: {profile_name}")
1524 }
1525 Some(BedrockAuth::IamCredentials { .. }) if env_var_set => {
1526 format!(
1527 "Using IAM credentials from {} and {} environment variables",
1528 ZED_BEDROCK_ACCESS_KEY_ID_VAR.name, ZED_BEDROCK_SECRET_ACCESS_KEY_VAR.name
1529 )
1530 }
1531 Some(BedrockAuth::IamCredentials { .. }) => "Using IAM credentials".into(),
1532 Some(BedrockAuth::ApiKey { .. }) if env_var_set => {
1533 format!(
1534 "Using Bedrock API Key from {} environment variable",
1535 ZED_BEDROCK_BEARER_TOKEN_VAR.name
1536 )
1537 }
1538 Some(BedrockAuth::ApiKey { .. }) => "Using Bedrock API Key".into(),
1539 None => "Not authenticated".into(),
1540 };
1541
1542 // Determine if credentials can be reset
1543 // Settings-derived auth (non-ApiKey) cannot be reset from UI
1544 let is_settings_derived = matches!(
1545 settings_auth_method,
1546 Some(BedrockAuthMethod::Automatic)
1547 | Some(BedrockAuthMethod::NamedProfile)
1548 | Some(BedrockAuthMethod::SingleSignOn)
1549 );
1550
1551 let tooltip_label = if env_var_set {
1552 Some(format!(
1553 "To reset your credentials, unset the {}, {}, and {} or {} environment variables.",
1554 ZED_BEDROCK_ACCESS_KEY_ID_VAR.name,
1555 ZED_BEDROCK_SECRET_ACCESS_KEY_VAR.name,
1556 ZED_BEDROCK_SESSION_TOKEN_VAR.name,
1557 ZED_BEDROCK_BEARER_TOKEN_VAR.name
1558 ))
1559 } else if is_settings_derived {
1560 Some(
1561 "Authentication method is configured in settings. Edit settings.json to change."
1562 .to_string(),
1563 )
1564 } else {
1565 None
1566 };
1567
1568 if self.should_render_editor(cx) {
1569 return ConfiguredApiCard::new(configured_label)
1570 .disabled(env_var_set || is_settings_derived)
1571 .on_click(cx.listener(|this, _, window, cx| this.reset_credentials(window, cx)))
1572 .when_some(tooltip_label, |this, label| this.tooltip_label(label))
1573 .into_any_element();
1574 }
1575
1576 v_flex()
1577 .min_w_0()
1578 .w_full()
1579 .track_focus(&self.focus_handle)
1580 .on_action(cx.listener(Self::on_tab))
1581 .on_action(cx.listener(Self::on_tab_prev))
1582 .on_action(cx.listener(ConfigurationView::save_credentials))
1583 .child(Label::new("To use Zed's agent with Bedrock, you can set a custom authentication strategy through your settings file or use static credentials."))
1584 .child(Label::new("But first, to access models on AWS, you need to:").mt_1())
1585 .child(
1586 List::new()
1587 .child(
1588 ListBulletItem::new("")
1589 .child(Label::new(
1590 "Grant permissions to the strategy you'll use according to the:",
1591 ))
1592 .child(ButtonLink::new(
1593 "Prerequisites",
1594 "https://docs.aws.amazon.com/bedrock/latest/userguide/inference-prereq.html",
1595 )),
1596 )
1597 .child(
1598 ListBulletItem::new("")
1599 .child(Label::new("Select the models you would like access to:"))
1600 .child(ButtonLink::new(
1601 "Bedrock Model Catalog",
1602 "https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/model-catalog",
1603 )),
1604 ),
1605 )
1606 .child(self.render_static_credentials_ui())
1607 .into_any()
1608 }
1609}
1610
1611impl ConfigurationView {
1612 fn render_static_credentials_ui(&self) -> impl IntoElement {
1613 let section_header = |title: SharedString| {
1614 h_flex()
1615 .gap_2()
1616 .child(Label::new(title).size(LabelSize::Default))
1617 .child(Divider::horizontal())
1618 };
1619
1620 let list_item = List::new()
1621 .child(
1622 ListBulletItem::new("")
1623 .child(Label::new(
1624 "For access keys: Create an IAM user in the AWS console with programmatic access",
1625 ))
1626 .child(ButtonLink::new(
1627 "IAM Console",
1628 "https://us-east-1.console.aws.amazon.com/iam/home?region=us-east-1#/users",
1629 )),
1630 )
1631 .child(
1632 ListBulletItem::new("")
1633 .child(Label::new("For Bedrock API Keys: Generate an API key from the"))
1634 .child(ButtonLink::new(
1635 "Bedrock Console",
1636 "https://docs.aws.amazon.com/bedrock/latest/userguide/api-keys-use.html",
1637 )),
1638 )
1639 .child(
1640 ListBulletItem::new("")
1641 .child(Label::new("Attach the necessary Bedrock permissions to"))
1642 .child(ButtonLink::new(
1643 "this user",
1644 "https://docs.aws.amazon.com/bedrock/latest/userguide/inference-prereq.html",
1645 )),
1646 )
1647 .child(ListBulletItem::new(
1648 "Enter either access keys OR a Bedrock API Key below (not both)",
1649 ));
1650
1651 v_flex()
1652 .my_2()
1653 .tab_group()
1654 .gap_1p5()
1655 .child(section_header("Static Credentials".into()))
1656 .child(Label::new(
1657 "This method uses your AWS access key ID and secret access key, or a Bedrock API Key.",
1658 ))
1659 .child(list_item)
1660 .child(self.access_key_id_editor.clone())
1661 .child(self.secret_access_key_editor.clone())
1662 .child(self.session_token_editor.clone())
1663 .child(
1664 Label::new(format!(
1665 "You can also set the {}, {} and {} environment variables (or {} for Bedrock API Key authentication) and restart Zed.",
1666 ZED_BEDROCK_ACCESS_KEY_ID_VAR.name,
1667 ZED_BEDROCK_SECRET_ACCESS_KEY_VAR.name,
1668 ZED_BEDROCK_REGION_VAR.name,
1669 ZED_BEDROCK_BEARER_TOKEN_VAR.name
1670 ))
1671 .size(LabelSize::Small)
1672 .color(Color::Muted),
1673 )
1674 .child(
1675 Label::new(format!(
1676 "Optionally, if your environment uses AWS CLI profiles, you can set {}; if it requires a custom endpoint, you can set {}; and if it requires a Session Token, you can set {}.",
1677 ZED_AWS_PROFILE_VAR.name,
1678 ZED_AWS_ENDPOINT_VAR.name,
1679 ZED_BEDROCK_SESSION_TOKEN_VAR.name
1680 ))
1681 .size(LabelSize::Small)
1682 .color(Color::Muted)
1683 .mt_1()
1684 .mb_2p5(),
1685 )
1686 .child(section_header("Using the an API key".into()))
1687 .child(self.bearer_token_editor.clone())
1688 .child(
1689 Label::new(format!(
1690 "Region is configured via {} environment variable or settings.json (defaults to us-east-1).",
1691 ZED_BEDROCK_REGION_VAR.name
1692 ))
1693 .size(LabelSize::Small)
1694 .color(Color::Muted)
1695 )
1696 }
1697}