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::{fix_streamed_json, 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 .with_context(|| format!("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 self.model.supports_thinking()
646 }
647
648 fn supported_effort_levels(&self) -> Vec<language_model::LanguageModelEffortLevel> {
649 if self.model.supports_adaptive_thinking() {
650 vec![
651 language_model::LanguageModelEffortLevel {
652 name: "Low".into(),
653 value: "low".into(),
654 is_default: false,
655 },
656 language_model::LanguageModelEffortLevel {
657 name: "Medium".into(),
658 value: "medium".into(),
659 is_default: false,
660 },
661 language_model::LanguageModelEffortLevel {
662 name: "High".into(),
663 value: "high".into(),
664 is_default: true,
665 },
666 language_model::LanguageModelEffortLevel {
667 name: "Max".into(),
668 value: "max".into(),
669 is_default: false,
670 },
671 ]
672 } else {
673 Vec::new()
674 }
675 }
676
677 fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
678 match choice {
679 LanguageModelToolChoice::Auto | LanguageModelToolChoice::Any => {
680 self.model.supports_tool_use()
681 }
682 // Add support for None - we'll filter tool calls at response
683 LanguageModelToolChoice::None => self.model.supports_tool_use(),
684 }
685 }
686
687 fn supports_streaming_tools(&self) -> bool {
688 true
689 }
690
691 fn telemetry_id(&self) -> String {
692 format!("bedrock/{}", self.model.id())
693 }
694
695 fn max_token_count(&self) -> u64 {
696 self.model.max_token_count()
697 }
698
699 fn max_output_tokens(&self) -> Option<u64> {
700 Some(self.model.max_output_tokens())
701 }
702
703 fn count_tokens(
704 &self,
705 request: LanguageModelRequest,
706 cx: &App,
707 ) -> BoxFuture<'static, Result<u64>> {
708 get_bedrock_tokens(request, cx)
709 }
710
711 fn stream_completion(
712 &self,
713 request: LanguageModelRequest,
714 cx: &AsyncApp,
715 ) -> BoxFuture<
716 'static,
717 Result<
718 BoxStream<'static, Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>,
719 LanguageModelCompletionError,
720 >,
721 > {
722 let (region, allow_global, allow_extended_context) =
723 cx.read_entity(&self.state, |state, _cx| {
724 (
725 state.get_region(),
726 state.get_allow_global(),
727 state.get_allow_extended_context(),
728 )
729 });
730
731 let model_id = match self.model.cross_region_inference_id(®ion, allow_global) {
732 Ok(s) => s,
733 Err(e) => {
734 return async move { Err(e.into()) }.boxed();
735 }
736 };
737
738 let deny_tool_calls = request.tool_choice == Some(LanguageModelToolChoice::None);
739
740 let use_extended_context = allow_extended_context && self.model.supports_extended_context();
741
742 let request = match into_bedrock(
743 request,
744 model_id,
745 self.model.default_temperature(),
746 self.model.max_output_tokens(),
747 self.model.thinking_mode(),
748 self.model.supports_caching(),
749 self.model.supports_tool_use(),
750 use_extended_context,
751 ) {
752 Ok(request) => request,
753 Err(err) => return futures::future::ready(Err(err.into())).boxed(),
754 };
755
756 let request = self.stream_completion(request, cx);
757 let display_name = self.model.display_name().to_string();
758 let future = self.request_limiter.stream(async move {
759 let response = request.await.map_err(|err| match err {
760 BedrockError::Validation(ref msg) => {
761 if msg.contains("model identifier is invalid") {
762 LanguageModelCompletionError::Other(anyhow!(
763 "{display_name} is not available in {region}. \
764 Try switching to a region where this model is supported."
765 ))
766 } else {
767 LanguageModelCompletionError::BadRequestFormat {
768 provider: PROVIDER_NAME,
769 message: msg.clone(),
770 }
771 }
772 }
773 BedrockError::RateLimited => LanguageModelCompletionError::RateLimitExceeded {
774 provider: PROVIDER_NAME,
775 retry_after: None,
776 },
777 BedrockError::ServiceUnavailable => {
778 LanguageModelCompletionError::ServerOverloaded {
779 provider: PROVIDER_NAME,
780 retry_after: None,
781 }
782 }
783 BedrockError::AccessDenied(msg) => LanguageModelCompletionError::PermissionError {
784 provider: PROVIDER_NAME,
785 message: msg,
786 },
787 BedrockError::InternalServer(msg) => {
788 LanguageModelCompletionError::ApiInternalServerError {
789 provider: PROVIDER_NAME,
790 message: msg,
791 }
792 }
793 other => LanguageModelCompletionError::Other(anyhow!(other)),
794 })?;
795 let events = map_to_language_model_completion_events(response);
796
797 if deny_tool_calls {
798 Ok(deny_tool_use_events(events).boxed())
799 } else {
800 Ok(events.boxed())
801 }
802 });
803
804 async move { Ok(future.await?.boxed()) }.boxed()
805 }
806
807 fn cache_configuration(&self) -> Option<LanguageModelCacheConfiguration> {
808 self.model
809 .cache_configuration()
810 .map(|config| LanguageModelCacheConfiguration {
811 max_cache_anchors: config.max_cache_anchors,
812 should_speculate: false,
813 min_total_token: config.min_total_token,
814 })
815 }
816}
817
818fn deny_tool_use_events(
819 events: impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>,
820) -> impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>> {
821 events.map(|event| {
822 match event {
823 Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) => {
824 // Convert tool use to an error message if model decided to call it
825 Ok(LanguageModelCompletionEvent::Text(format!(
826 "\n\n[Error: Tool calls are disabled in this context. Attempted to call '{}']",
827 tool_use.name
828 )))
829 }
830 other => other,
831 }
832 })
833}
834
835pub fn into_bedrock(
836 request: LanguageModelRequest,
837 model: String,
838 default_temperature: f32,
839 max_output_tokens: u64,
840 thinking_mode: BedrockModelMode,
841 supports_caching: bool,
842 supports_tool_use: bool,
843 allow_extended_context: bool,
844) -> Result<bedrock::Request> {
845 let mut new_messages: Vec<BedrockMessage> = Vec::new();
846 let mut system_message = String::new();
847
848 // Track whether messages contain tool content - Bedrock requires toolConfig
849 // when tool blocks are present, so we may need to add a dummy tool
850 let mut messages_contain_tool_content = false;
851
852 for message in request.messages {
853 if message.contents_empty() {
854 continue;
855 }
856
857 match message.role {
858 Role::User | Role::Assistant => {
859 let mut bedrock_message_content: Vec<BedrockInnerContent> = message
860 .content
861 .into_iter()
862 .filter_map(|content| match content {
863 MessageContent::Text(text) => {
864 if !text.is_empty() {
865 Some(BedrockInnerContent::Text(text))
866 } else {
867 None
868 }
869 }
870 MessageContent::Thinking { text, signature } => {
871 if model.contains(Model::DeepSeekR1.request_id()) {
872 // DeepSeekR1 doesn't support thinking blocks
873 // And the AWS API demands that you strip them
874 return None;
875 }
876 if signature.is_none() {
877 // Thinking blocks without a signature are invalid
878 // (e.g. from cancellation mid-think) and must be
879 // stripped to avoid API errors.
880 return None;
881 }
882 let thinking = BedrockThinkingTextBlock::builder()
883 .text(text)
884 .set_signature(signature)
885 .build()
886 .context("failed to build reasoning block")
887 .log_err()?;
888
889 Some(BedrockInnerContent::ReasoningContent(
890 BedrockThinkingBlock::ReasoningText(thinking),
891 ))
892 }
893 MessageContent::RedactedThinking(blob) => {
894 if model.contains(Model::DeepSeekR1.request_id()) {
895 // DeepSeekR1 doesn't support thinking blocks
896 // And the AWS API demands that you strip them
897 return None;
898 }
899 let redacted =
900 BedrockThinkingBlock::RedactedContent(BedrockBlob::new(blob));
901
902 Some(BedrockInnerContent::ReasoningContent(redacted))
903 }
904 MessageContent::ToolUse(tool_use) => {
905 messages_contain_tool_content = true;
906 let input = if tool_use.input.is_null() {
907 // Bedrock API requires valid JsonValue, not null, for tool use input
908 value_to_aws_document(&serde_json::json!({}))
909 } else {
910 value_to_aws_document(&tool_use.input)
911 };
912 BedrockToolUseBlock::builder()
913 .name(tool_use.name.to_string())
914 .tool_use_id(tool_use.id.to_string())
915 .input(input)
916 .build()
917 .context("failed to build Bedrock tool use block")
918 .log_err()
919 .map(BedrockInnerContent::ToolUse)
920 }
921 MessageContent::ToolResult(tool_result) => {
922 messages_contain_tool_content = true;
923 BedrockToolResultBlock::builder()
924 .tool_use_id(tool_result.tool_use_id.to_string())
925 .content(match tool_result.content {
926 LanguageModelToolResultContent::Text(text) => {
927 BedrockToolResultContentBlock::Text(text.to_string())
928 }
929 LanguageModelToolResultContent::Image(image) => {
930 use base64::Engine;
931
932 match base64::engine::general_purpose::STANDARD
933 .decode(image.source.as_bytes())
934 {
935 Ok(image_bytes) => {
936 match BedrockImageBlock::builder()
937 .format(BedrockImageFormat::Png)
938 .source(BedrockImageSource::Bytes(
939 BedrockBlob::new(image_bytes),
940 ))
941 .build()
942 {
943 Ok(image_block) => {
944 BedrockToolResultContentBlock::Image(
945 image_block,
946 )
947 }
948 Err(err) => {
949 BedrockToolResultContentBlock::Text(
950 format!(
951 "[Failed to build image block: {}]",
952 err
953 ),
954 )
955 }
956 }
957 }
958 Err(err) => {
959 BedrockToolResultContentBlock::Text(format!(
960 "[Failed to decode tool result image: {}]",
961 err
962 ))
963 }
964 }
965 }
966 })
967 .status({
968 if tool_result.is_error {
969 BedrockToolResultStatus::Error
970 } else {
971 BedrockToolResultStatus::Success
972 }
973 })
974 .build()
975 .context("failed to build Bedrock tool result block")
976 .log_err()
977 .map(BedrockInnerContent::ToolResult)
978 }
979 MessageContent::Image(image) => {
980 use base64::Engine;
981
982 let image_bytes = base64::engine::general_purpose::STANDARD
983 .decode(image.source.as_bytes())
984 .context("failed to decode base64 image data")
985 .log_err()?;
986
987 BedrockImageBlock::builder()
988 .format(BedrockImageFormat::Png)
989 .source(BedrockImageSource::Bytes(BedrockBlob::new(image_bytes)))
990 .build()
991 .context("failed to build Bedrock image block")
992 .log_err()
993 .map(BedrockInnerContent::Image)
994 }
995 })
996 .collect();
997 if message.cache && supports_caching {
998 bedrock_message_content.push(BedrockInnerContent::CachePoint(
999 CachePointBlock::builder()
1000 .r#type(CachePointType::Default)
1001 .build()
1002 .context("failed to build cache point block")?,
1003 ));
1004 }
1005 let bedrock_role = match message.role {
1006 Role::User => bedrock::BedrockRole::User,
1007 Role::Assistant => bedrock::BedrockRole::Assistant,
1008 Role::System => unreachable!("System role should never occur here"),
1009 };
1010 if bedrock_message_content.is_empty() {
1011 continue;
1012 }
1013
1014 if let Some(last_message) = new_messages.last_mut()
1015 && last_message.role == bedrock_role
1016 {
1017 last_message.content.extend(bedrock_message_content);
1018 continue;
1019 }
1020 new_messages.push(
1021 BedrockMessage::builder()
1022 .role(bedrock_role)
1023 .set_content(Some(bedrock_message_content))
1024 .build()
1025 .context("failed to build Bedrock message")?,
1026 );
1027 }
1028 Role::System => {
1029 if !system_message.is_empty() {
1030 system_message.push_str("\n\n");
1031 }
1032 system_message.push_str(&message.string_contents());
1033 }
1034 }
1035 }
1036
1037 let mut tool_spec: Vec<BedrockTool> = if supports_tool_use {
1038 request
1039 .tools
1040 .iter()
1041 .filter_map(|tool| {
1042 Some(BedrockTool::ToolSpec(
1043 BedrockToolSpec::builder()
1044 .name(tool.name.clone())
1045 .description(tool.description.clone())
1046 .input_schema(BedrockToolInputSchema::Json(value_to_aws_document(
1047 &tool.input_schema,
1048 )))
1049 .build()
1050 .log_err()?,
1051 ))
1052 })
1053 .collect()
1054 } else {
1055 Vec::new()
1056 };
1057
1058 // Bedrock requires toolConfig when messages contain tool use/result blocks.
1059 // If no tools are defined but messages contain tool content (e.g., when
1060 // summarising a conversation that used tools), add a dummy tool to satisfy
1061 // the API requirement.
1062 if supports_tool_use && tool_spec.is_empty() && messages_contain_tool_content {
1063 tool_spec.push(BedrockTool::ToolSpec(
1064 BedrockToolSpec::builder()
1065 .name("_placeholder")
1066 .description("Placeholder tool to satisfy Bedrock API requirements when conversation history contains tool usage")
1067 .input_schema(BedrockToolInputSchema::Json(value_to_aws_document(
1068 &serde_json::json!({"type": "object", "properties": {}}),
1069 )))
1070 .build()
1071 .context("failed to build placeholder tool spec")?,
1072 ));
1073 }
1074
1075 if !tool_spec.is_empty() && supports_caching {
1076 tool_spec.push(BedrockTool::CachePoint(
1077 CachePointBlock::builder()
1078 .r#type(CachePointType::Default)
1079 .build()
1080 .context("failed to build cache point block")?,
1081 ));
1082 }
1083
1084 let tool_choice = match request.tool_choice {
1085 Some(LanguageModelToolChoice::Auto) | None => {
1086 BedrockToolChoice::Auto(BedrockAutoToolChoice::builder().build())
1087 }
1088 Some(LanguageModelToolChoice::Any) => {
1089 BedrockToolChoice::Any(BedrockAnyToolChoice::builder().build())
1090 }
1091 Some(LanguageModelToolChoice::None) => {
1092 // For None, we still use Auto but will filter out tool calls in the response
1093 BedrockToolChoice::Auto(BedrockAutoToolChoice::builder().build())
1094 }
1095 };
1096 let tool_config = if tool_spec.is_empty() {
1097 None
1098 } else {
1099 Some(
1100 BedrockToolConfig::builder()
1101 .set_tools(Some(tool_spec))
1102 .tool_choice(tool_choice)
1103 .build()?,
1104 )
1105 };
1106
1107 Ok(bedrock::Request {
1108 model,
1109 messages: new_messages,
1110 max_tokens: max_output_tokens,
1111 system: Some(system_message),
1112 tools: tool_config,
1113 thinking: if request.thinking_allowed {
1114 match thinking_mode {
1115 BedrockModelMode::Thinking { budget_tokens } => {
1116 Some(bedrock::Thinking::Enabled { budget_tokens })
1117 }
1118 BedrockModelMode::AdaptiveThinking {
1119 effort: default_effort,
1120 } => {
1121 let effort = request
1122 .thinking_effort
1123 .as_deref()
1124 .and_then(|e| match e {
1125 "low" => Some(bedrock::BedrockAdaptiveThinkingEffort::Low),
1126 "medium" => Some(bedrock::BedrockAdaptiveThinkingEffort::Medium),
1127 "high" => Some(bedrock::BedrockAdaptiveThinkingEffort::High),
1128 "max" => Some(bedrock::BedrockAdaptiveThinkingEffort::Max),
1129 _ => None,
1130 })
1131 .unwrap_or(default_effort);
1132 Some(bedrock::Thinking::Adaptive { effort })
1133 }
1134 BedrockModelMode::Default => None,
1135 }
1136 } else {
1137 None
1138 },
1139 metadata: None,
1140 stop_sequences: Vec::new(),
1141 temperature: request.temperature.or(Some(default_temperature)),
1142 top_k: None,
1143 top_p: None,
1144 allow_extended_context,
1145 })
1146}
1147
1148// TODO: just call the ConverseOutput.usage() method:
1149// https://docs.rs/aws-sdk-bedrockruntime/latest/aws_sdk_bedrockruntime/operation/converse/struct.ConverseOutput.html#method.output
1150pub fn get_bedrock_tokens(
1151 request: LanguageModelRequest,
1152 cx: &App,
1153) -> BoxFuture<'static, Result<u64>> {
1154 cx.background_executor()
1155 .spawn(async move {
1156 let messages = request.messages;
1157 let mut tokens_from_images = 0;
1158 let mut string_messages = Vec::with_capacity(messages.len());
1159
1160 for message in messages {
1161 use language_model::MessageContent;
1162
1163 let mut string_contents = String::new();
1164
1165 for content in message.content {
1166 match content {
1167 MessageContent::Text(text) | MessageContent::Thinking { text, .. } => {
1168 string_contents.push_str(&text);
1169 }
1170 MessageContent::RedactedThinking(_) => {}
1171 MessageContent::Image(image) => {
1172 tokens_from_images += image.estimate_tokens();
1173 }
1174 MessageContent::ToolUse(_tool_use) => {
1175 // TODO: Estimate token usage from tool uses.
1176 }
1177 MessageContent::ToolResult(tool_result) => match tool_result.content {
1178 LanguageModelToolResultContent::Text(text) => {
1179 string_contents.push_str(&text);
1180 }
1181 LanguageModelToolResultContent::Image(image) => {
1182 tokens_from_images += image.estimate_tokens();
1183 }
1184 },
1185 }
1186 }
1187
1188 if !string_contents.is_empty() {
1189 string_messages.push(tiktoken_rs::ChatCompletionRequestMessage {
1190 role: match message.role {
1191 Role::User => "user".into(),
1192 Role::Assistant => "assistant".into(),
1193 Role::System => "system".into(),
1194 },
1195 content: Some(string_contents),
1196 name: None,
1197 function_call: None,
1198 });
1199 }
1200 }
1201
1202 // Tiktoken doesn't yet support these models, so we manually use the
1203 // same tokenizer as GPT-4.
1204 tiktoken_rs::num_tokens_from_messages("gpt-4", &string_messages)
1205 .map(|tokens| (tokens + tokens_from_images) as u64)
1206 })
1207 .boxed()
1208}
1209
1210pub fn map_to_language_model_completion_events(
1211 events: Pin<Box<dyn Send + Stream<Item = Result<BedrockStreamingResponse, anyhow::Error>>>>,
1212) -> impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>> {
1213 struct RawToolUse {
1214 id: String,
1215 name: String,
1216 input_json: String,
1217 }
1218
1219 struct State {
1220 events: Pin<Box<dyn Send + Stream<Item = Result<BedrockStreamingResponse, anyhow::Error>>>>,
1221 tool_uses_by_index: HashMap<i32, RawToolUse>,
1222 emitted_tool_use: bool,
1223 }
1224
1225 let initial_state = State {
1226 events,
1227 tool_uses_by_index: HashMap::default(),
1228 emitted_tool_use: false,
1229 };
1230
1231 futures::stream::unfold(initial_state, |mut state| async move {
1232 match state.events.next().await {
1233 Some(event_result) => match event_result {
1234 Ok(event) => {
1235 let result = match event {
1236 ConverseStreamOutput::ContentBlockDelta(cb_delta) => match cb_delta.delta {
1237 Some(ContentBlockDelta::Text(text)) => {
1238 Some(Ok(LanguageModelCompletionEvent::Text(text)))
1239 }
1240 Some(ContentBlockDelta::ToolUse(tool_output)) => {
1241 if let Some(tool_use) = state
1242 .tool_uses_by_index
1243 .get_mut(&cb_delta.content_block_index)
1244 {
1245 tool_use.input_json.push_str(tool_output.input());
1246 if let Ok(input) = serde_json::from_str::<serde_json::Value>(
1247 &fix_streamed_json(&tool_use.input_json),
1248 ) {
1249 Some(Ok(LanguageModelCompletionEvent::ToolUse(
1250 LanguageModelToolUse {
1251 id: tool_use.id.clone().into(),
1252 name: tool_use.name.clone().into(),
1253 is_input_complete: false,
1254 raw_input: tool_use.input_json.clone(),
1255 input,
1256 thought_signature: None,
1257 },
1258 )))
1259 } else {
1260 None
1261 }
1262 } else {
1263 None
1264 }
1265 }
1266 Some(ContentBlockDelta::ReasoningContent(thinking)) => match thinking {
1267 ReasoningContentBlockDelta::Text(thoughts) => {
1268 Some(Ok(LanguageModelCompletionEvent::Thinking {
1269 text: thoughts,
1270 signature: None,
1271 }))
1272 }
1273 ReasoningContentBlockDelta::Signature(sig) => {
1274 Some(Ok(LanguageModelCompletionEvent::Thinking {
1275 text: "".into(),
1276 signature: Some(sig),
1277 }))
1278 }
1279 ReasoningContentBlockDelta::RedactedContent(redacted) => {
1280 let content = String::from_utf8(redacted.into_inner())
1281 .unwrap_or("REDACTED".to_string());
1282 Some(Ok(LanguageModelCompletionEvent::Thinking {
1283 text: content,
1284 signature: None,
1285 }))
1286 }
1287 _ => None,
1288 },
1289 _ => None,
1290 },
1291 ConverseStreamOutput::ContentBlockStart(cb_start) => {
1292 if let Some(ContentBlockStart::ToolUse(tool_start)) = cb_start.start {
1293 state.tool_uses_by_index.insert(
1294 cb_start.content_block_index,
1295 RawToolUse {
1296 id: tool_start.tool_use_id,
1297 name: tool_start.name,
1298 input_json: String::new(),
1299 },
1300 );
1301 }
1302 None
1303 }
1304 ConverseStreamOutput::MessageStart(_) => None,
1305 ConverseStreamOutput::ContentBlockStop(cb_stop) => state
1306 .tool_uses_by_index
1307 .remove(&cb_stop.content_block_index)
1308 .map(|tool_use| {
1309 state.emitted_tool_use = true;
1310
1311 let input = parse_tool_arguments(&tool_use.input_json)
1312 .unwrap_or_else(|_| Value::Object(Default::default()));
1313
1314 Ok(LanguageModelCompletionEvent::ToolUse(
1315 LanguageModelToolUse {
1316 id: tool_use.id.into(),
1317 name: tool_use.name.into(),
1318 is_input_complete: true,
1319 raw_input: tool_use.input_json,
1320 input,
1321 thought_signature: None,
1322 },
1323 ))
1324 }),
1325 ConverseStreamOutput::Metadata(cb_meta) => cb_meta.usage.map(|metadata| {
1326 Ok(LanguageModelCompletionEvent::UsageUpdate(TokenUsage {
1327 input_tokens: metadata.input_tokens as u64,
1328 output_tokens: metadata.output_tokens as u64,
1329 cache_creation_input_tokens: metadata
1330 .cache_write_input_tokens
1331 .unwrap_or_default()
1332 as u64,
1333 cache_read_input_tokens: metadata
1334 .cache_read_input_tokens
1335 .unwrap_or_default()
1336 as u64,
1337 }))
1338 }),
1339 ConverseStreamOutput::MessageStop(message_stop) => {
1340 let stop_reason = if state.emitted_tool_use {
1341 // Some models (e.g. Kimi) send EndTurn even when
1342 // they've made tool calls. Trust the content over
1343 // the stop reason.
1344 language_model::StopReason::ToolUse
1345 } else {
1346 match message_stop.stop_reason {
1347 StopReason::ToolUse => language_model::StopReason::ToolUse,
1348 _ => language_model::StopReason::EndTurn,
1349 }
1350 };
1351 Some(Ok(LanguageModelCompletionEvent::Stop(stop_reason)))
1352 }
1353 _ => None,
1354 };
1355
1356 Some((result, state))
1357 }
1358 Err(err) => Some((
1359 Some(Err(LanguageModelCompletionError::Other(anyhow!(err)))),
1360 state,
1361 )),
1362 },
1363 None => None,
1364 }
1365 })
1366 .filter_map(|result| async move { result })
1367}
1368
1369struct ConfigurationView {
1370 access_key_id_editor: Entity<InputField>,
1371 secret_access_key_editor: Entity<InputField>,
1372 session_token_editor: Entity<InputField>,
1373 bearer_token_editor: Entity<InputField>,
1374 state: Entity<State>,
1375 load_credentials_task: Option<Task<()>>,
1376 focus_handle: FocusHandle,
1377}
1378
1379impl ConfigurationView {
1380 const PLACEHOLDER_ACCESS_KEY_ID_TEXT: &'static str = "XXXXXXXXXXXXXXXX";
1381 const PLACEHOLDER_SECRET_ACCESS_KEY_TEXT: &'static str =
1382 "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX";
1383 const PLACEHOLDER_SESSION_TOKEN_TEXT: &'static str = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX";
1384 const PLACEHOLDER_BEARER_TOKEN_TEXT: &'static str = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX";
1385
1386 fn new(state: Entity<State>, window: &mut Window, cx: &mut Context<Self>) -> Self {
1387 let focus_handle = cx.focus_handle();
1388
1389 cx.observe(&state, |_, _, cx| {
1390 cx.notify();
1391 })
1392 .detach();
1393
1394 let access_key_id_editor = cx.new(|cx| {
1395 InputField::new(window, cx, Self::PLACEHOLDER_ACCESS_KEY_ID_TEXT)
1396 .label("Access Key ID")
1397 .tab_index(0)
1398 .tab_stop(true)
1399 });
1400
1401 let secret_access_key_editor = cx.new(|cx| {
1402 InputField::new(window, cx, Self::PLACEHOLDER_SECRET_ACCESS_KEY_TEXT)
1403 .label("Secret Access Key")
1404 .tab_index(1)
1405 .tab_stop(true)
1406 });
1407
1408 let session_token_editor = cx.new(|cx| {
1409 InputField::new(window, cx, Self::PLACEHOLDER_SESSION_TOKEN_TEXT)
1410 .label("Session Token (Optional)")
1411 .tab_index(2)
1412 .tab_stop(true)
1413 });
1414
1415 let bearer_token_editor = cx.new(|cx| {
1416 InputField::new(window, cx, Self::PLACEHOLDER_BEARER_TOKEN_TEXT)
1417 .label("Bedrock API Key")
1418 .tab_index(3)
1419 .tab_stop(true)
1420 });
1421
1422 let load_credentials_task = Some(cx.spawn({
1423 let state = state.clone();
1424 async move |this, cx| {
1425 if let Some(task) = Some(state.update(cx, |state, cx| state.authenticate(cx))) {
1426 // We don't log an error, because "not signed in" is also an error.
1427 let _ = task.await;
1428 }
1429 this.update(cx, |this, cx| {
1430 this.load_credentials_task = None;
1431 cx.notify();
1432 })
1433 .log_err();
1434 }
1435 }));
1436
1437 Self {
1438 access_key_id_editor,
1439 secret_access_key_editor,
1440 session_token_editor,
1441 bearer_token_editor,
1442 state,
1443 load_credentials_task,
1444 focus_handle,
1445 }
1446 }
1447
1448 fn save_credentials(
1449 &mut self,
1450 _: &menu::Confirm,
1451 _window: &mut Window,
1452 cx: &mut Context<Self>,
1453 ) {
1454 let access_key_id = self
1455 .access_key_id_editor
1456 .read(cx)
1457 .text(cx)
1458 .trim()
1459 .to_string();
1460 let secret_access_key = self
1461 .secret_access_key_editor
1462 .read(cx)
1463 .text(cx)
1464 .trim()
1465 .to_string();
1466 let session_token = self
1467 .session_token_editor
1468 .read(cx)
1469 .text(cx)
1470 .trim()
1471 .to_string();
1472 let session_token = if session_token.is_empty() {
1473 None
1474 } else {
1475 Some(session_token)
1476 };
1477 let bearer_token = self
1478 .bearer_token_editor
1479 .read(cx)
1480 .text(cx)
1481 .trim()
1482 .to_string();
1483 let bearer_token = if bearer_token.is_empty() {
1484 None
1485 } else {
1486 Some(bearer_token)
1487 };
1488
1489 let state = self.state.clone();
1490 cx.spawn(async move |_, cx| {
1491 state
1492 .update(cx, |state, cx| {
1493 let credentials = BedrockCredentials {
1494 access_key_id,
1495 secret_access_key,
1496 session_token,
1497 bearer_token,
1498 };
1499
1500 state.set_static_credentials(credentials, cx)
1501 })
1502 .await
1503 })
1504 .detach_and_log_err(cx);
1505 }
1506
1507 fn reset_credentials(&mut self, window: &mut Window, cx: &mut Context<Self>) {
1508 self.access_key_id_editor
1509 .update(cx, |editor, cx| editor.set_text("", window, cx));
1510 self.secret_access_key_editor
1511 .update(cx, |editor, cx| editor.set_text("", window, cx));
1512 self.session_token_editor
1513 .update(cx, |editor, cx| editor.set_text("", window, cx));
1514 self.bearer_token_editor
1515 .update(cx, |editor, cx| editor.set_text("", window, cx));
1516
1517 let state = self.state.clone();
1518 cx.spawn(async move |_, cx| state.update(cx, |state, cx| state.reset_auth(cx)).await)
1519 .detach_and_log_err(cx);
1520 }
1521
1522 fn should_render_editor(&self, cx: &Context<Self>) -> bool {
1523 self.state.read(cx).is_authenticated()
1524 }
1525
1526 fn on_tab(&mut self, _: &menu::SelectNext, window: &mut Window, cx: &mut Context<Self>) {
1527 window.focus_next(cx);
1528 }
1529
1530 fn on_tab_prev(
1531 &mut self,
1532 _: &menu::SelectPrevious,
1533 window: &mut Window,
1534 cx: &mut Context<Self>,
1535 ) {
1536 window.focus_prev(cx);
1537 }
1538}
1539
1540impl Render for ConfigurationView {
1541 fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
1542 let state = self.state.read(cx);
1543 let env_var_set = state.credentials_from_env;
1544 let auth = state.auth.clone();
1545 let settings_auth_method = state
1546 .settings
1547 .as_ref()
1548 .and_then(|s| s.authentication_method.clone());
1549
1550 if self.load_credentials_task.is_some() {
1551 return div().child(Label::new("Loading credentials...")).into_any();
1552 }
1553
1554 let configured_label = match &auth {
1555 Some(BedrockAuth::Automatic) => {
1556 "Using automatic credentials (AWS default chain)".into()
1557 }
1558 Some(BedrockAuth::NamedProfile { profile_name }) => {
1559 format!("Using AWS profile: {profile_name}")
1560 }
1561 Some(BedrockAuth::SingleSignOn { profile_name }) => {
1562 format!("Using AWS SSO profile: {profile_name}")
1563 }
1564 Some(BedrockAuth::IamCredentials { .. }) if env_var_set => {
1565 format!(
1566 "Using IAM credentials from {} and {} environment variables",
1567 ZED_BEDROCK_ACCESS_KEY_ID_VAR.name, ZED_BEDROCK_SECRET_ACCESS_KEY_VAR.name
1568 )
1569 }
1570 Some(BedrockAuth::IamCredentials { .. }) => "Using IAM credentials".into(),
1571 Some(BedrockAuth::ApiKey { .. }) if env_var_set => {
1572 format!(
1573 "Using Bedrock API Key from {} environment variable",
1574 ZED_BEDROCK_BEARER_TOKEN_VAR.name
1575 )
1576 }
1577 Some(BedrockAuth::ApiKey { .. }) => "Using Bedrock API Key".into(),
1578 None => "Not authenticated".into(),
1579 };
1580
1581 // Determine if credentials can be reset
1582 // Settings-derived auth (non-ApiKey) cannot be reset from UI
1583 let is_settings_derived = matches!(
1584 settings_auth_method,
1585 Some(BedrockAuthMethod::Automatic)
1586 | Some(BedrockAuthMethod::NamedProfile)
1587 | Some(BedrockAuthMethod::SingleSignOn)
1588 );
1589
1590 let tooltip_label = if env_var_set {
1591 Some(format!(
1592 "To reset your credentials, unset the {}, {}, and {} or {} environment variables.",
1593 ZED_BEDROCK_ACCESS_KEY_ID_VAR.name,
1594 ZED_BEDROCK_SECRET_ACCESS_KEY_VAR.name,
1595 ZED_BEDROCK_SESSION_TOKEN_VAR.name,
1596 ZED_BEDROCK_BEARER_TOKEN_VAR.name
1597 ))
1598 } else if is_settings_derived {
1599 Some(
1600 "Authentication method is configured in settings. Edit settings.json to change."
1601 .to_string(),
1602 )
1603 } else {
1604 None
1605 };
1606
1607 if self.should_render_editor(cx) {
1608 return ConfiguredApiCard::new(configured_label)
1609 .disabled(env_var_set || is_settings_derived)
1610 .on_click(cx.listener(|this, _, window, cx| this.reset_credentials(window, cx)))
1611 .when_some(tooltip_label, |this, label| this.tooltip_label(label))
1612 .into_any_element();
1613 }
1614
1615 v_flex()
1616 .min_w_0()
1617 .w_full()
1618 .track_focus(&self.focus_handle)
1619 .on_action(cx.listener(Self::on_tab))
1620 .on_action(cx.listener(Self::on_tab_prev))
1621 .on_action(cx.listener(ConfigurationView::save_credentials))
1622 .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."))
1623 .child(Label::new("But first, to access models on AWS, you need to:").mt_1())
1624 .child(
1625 List::new()
1626 .child(
1627 ListBulletItem::new("")
1628 .child(Label::new(
1629 "Grant permissions to the strategy you'll use according to the:",
1630 ))
1631 .child(ButtonLink::new(
1632 "Prerequisites",
1633 "https://docs.aws.amazon.com/bedrock/latest/userguide/inference-prereq.html",
1634 )),
1635 )
1636 .child(
1637 ListBulletItem::new("")
1638 .child(Label::new("Select the models you would like access to:"))
1639 .child(ButtonLink::new(
1640 "Bedrock Model Catalog",
1641 "https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/model-catalog",
1642 )),
1643 ),
1644 )
1645 .child(self.render_static_credentials_ui())
1646 .into_any()
1647 }
1648}
1649
1650impl ConfigurationView {
1651 fn render_static_credentials_ui(&self) -> impl IntoElement {
1652 let section_header = |title: SharedString| {
1653 h_flex()
1654 .gap_2()
1655 .child(Label::new(title).size(LabelSize::Default))
1656 .child(Divider::horizontal())
1657 };
1658
1659 let list_item = List::new()
1660 .child(
1661 ListBulletItem::new("")
1662 .child(Label::new(
1663 "For access keys: Create an IAM user in the AWS console with programmatic access",
1664 ))
1665 .child(ButtonLink::new(
1666 "IAM Console",
1667 "https://us-east-1.console.aws.amazon.com/iam/home?region=us-east-1#/users",
1668 )),
1669 )
1670 .child(
1671 ListBulletItem::new("")
1672 .child(Label::new("For Bedrock API Keys: Generate an API key from the"))
1673 .child(ButtonLink::new(
1674 "Bedrock Console",
1675 "https://docs.aws.amazon.com/bedrock/latest/userguide/api-keys-use.html",
1676 )),
1677 )
1678 .child(
1679 ListBulletItem::new("")
1680 .child(Label::new("Attach the necessary Bedrock permissions to"))
1681 .child(ButtonLink::new(
1682 "this user",
1683 "https://docs.aws.amazon.com/bedrock/latest/userguide/inference-prereq.html",
1684 )),
1685 )
1686 .child(ListBulletItem::new(
1687 "Enter either access keys OR a Bedrock API Key below (not both)",
1688 ));
1689
1690 v_flex()
1691 .my_2()
1692 .tab_group()
1693 .gap_1p5()
1694 .child(section_header("Static Credentials".into()))
1695 .child(Label::new(
1696 "This method uses your AWS access key ID and secret access key, or a Bedrock API Key.",
1697 ))
1698 .child(list_item)
1699 .child(self.access_key_id_editor.clone())
1700 .child(self.secret_access_key_editor.clone())
1701 .child(self.session_token_editor.clone())
1702 .child(
1703 Label::new(format!(
1704 "You can also set the {}, {} and {} environment variables (or {} for Bedrock API Key authentication) and restart Zed.",
1705 ZED_BEDROCK_ACCESS_KEY_ID_VAR.name,
1706 ZED_BEDROCK_SECRET_ACCESS_KEY_VAR.name,
1707 ZED_BEDROCK_REGION_VAR.name,
1708 ZED_BEDROCK_BEARER_TOKEN_VAR.name
1709 ))
1710 .size(LabelSize::Small)
1711 .color(Color::Muted),
1712 )
1713 .child(
1714 Label::new(format!(
1715 "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 {}.",
1716 ZED_AWS_PROFILE_VAR.name,
1717 ZED_AWS_ENDPOINT_VAR.name,
1718 ZED_BEDROCK_SESSION_TOKEN_VAR.name
1719 ))
1720 .size(LabelSize::Small)
1721 .color(Color::Muted)
1722 .mt_1()
1723 .mb_2p5(),
1724 )
1725 .child(section_header("Using the an API key".into()))
1726 .child(self.bearer_token_editor.clone())
1727 .child(
1728 Label::new(format!(
1729 "Region is configured via {} environment variable or settings.json (defaults to us-east-1).",
1730 ZED_BEDROCK_REGION_VAR.name
1731 ))
1732 .size(LabelSize::Small)
1733 .color(Color::Muted)
1734 )
1735 }
1736}