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