bedrock.rs

   1use std::pin::Pin;
   2use std::str::FromStr;
   3use std::sync::Arc;
   4
   5use crate::ui::InstructionListItem;
   6use anyhow::{Context as _, Result, anyhow};
   7use aws_config::stalled_stream_protection::StalledStreamProtectionConfig;
   8use aws_config::{BehaviorVersion, Region};
   9use aws_credential_types::Credentials;
  10use aws_http_client::AwsHttpClient;
  11use bedrock::bedrock_client::Client as BedrockClient;
  12use bedrock::bedrock_client::config::timeout::TimeoutConfig;
  13use bedrock::bedrock_client::types::{
  14    CachePointBlock, CachePointType, ContentBlockDelta, ContentBlockStart, ConverseStreamOutput,
  15    ReasoningContentBlockDelta, StopReason,
  16};
  17use bedrock::{
  18    BedrockAnyToolChoice, BedrockAutoToolChoice, BedrockBlob, BedrockError, BedrockInnerContent,
  19    BedrockMessage, BedrockModelMode, BedrockStreamingResponse, BedrockThinkingBlock,
  20    BedrockThinkingTextBlock, BedrockTool, BedrockToolChoice, BedrockToolConfig,
  21    BedrockToolInputSchema, BedrockToolResultBlock, BedrockToolResultContentBlock,
  22    BedrockToolResultStatus, BedrockToolSpec, BedrockToolUseBlock, Model, value_to_aws_document,
  23};
  24use collections::{BTreeMap, HashMap};
  25use credentials_provider::CredentialsProvider;
  26use editor::{Editor, EditorElement, EditorStyle};
  27use futures::{FutureExt, Stream, StreamExt, future::BoxFuture, stream::BoxStream};
  28use gpui::{
  29    AnyView, App, AsyncApp, Context, Entity, FontStyle, FontWeight, Subscription, Task, TextStyle,
  30    WhiteSpace,
  31};
  32use gpui_tokio::Tokio;
  33use http_client::HttpClient;
  34use language_model::{
  35    AuthenticateError, LanguageModel, LanguageModelCacheConfiguration,
  36    LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelId, LanguageModelName,
  37    LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName,
  38    LanguageModelProviderState, LanguageModelRequest, LanguageModelToolChoice,
  39    LanguageModelToolResultContent, LanguageModelToolUse, MessageContent, RateLimiter, Role,
  40    TokenUsage,
  41};
  42use schemars::JsonSchema;
  43use serde::{Deserialize, Serialize};
  44use serde_json::Value;
  45use settings::{Settings, SettingsStore};
  46use smol::lock::OnceCell;
  47use strum::{EnumIter, IntoEnumIterator, IntoStaticStr};
  48use theme::ThemeSettings;
  49use tokio::runtime::Handle;
  50use ui::{Icon, IconName, List, Tooltip, prelude::*};
  51use util::ResultExt;
  52
  53use crate::AllLanguageModelSettings;
  54
  55const PROVIDER_ID: &str = "amazon-bedrock";
  56const PROVIDER_NAME: &str = "Amazon Bedrock";
  57
  58#[derive(Default, Clone, Deserialize, Serialize, PartialEq, Debug)]
  59pub struct BedrockCredentials {
  60    pub access_key_id: String,
  61    pub secret_access_key: String,
  62    pub session_token: Option<String>,
  63    pub region: String,
  64}
  65
  66#[derive(Default, Clone, Debug, PartialEq)]
  67pub struct AmazonBedrockSettings {
  68    pub available_models: Vec<AvailableModel>,
  69    pub region: Option<String>,
  70    pub endpoint: Option<String>,
  71    pub profile_name: Option<String>,
  72    pub role_arn: Option<String>,
  73    pub authentication_method: Option<BedrockAuthMethod>,
  74}
  75
  76#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, EnumIter, IntoStaticStr, JsonSchema)]
  77pub enum BedrockAuthMethod {
  78    #[serde(rename = "named_profile")]
  79    NamedProfile,
  80    #[serde(rename = "sso")]
  81    SingleSignOn,
  82    /// IMDSv2, PodIdentity, env vars, etc.
  83    #[serde(rename = "default")]
  84    Automatic,
  85}
  86
  87#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
  88pub struct AvailableModel {
  89    pub name: String,
  90    pub display_name: Option<String>,
  91    pub max_tokens: u64,
  92    pub cache_configuration: Option<LanguageModelCacheConfiguration>,
  93    pub max_output_tokens: Option<u64>,
  94    pub default_temperature: Option<f32>,
  95    pub mode: Option<ModelMode>,
  96}
  97
  98#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize, JsonSchema)]
  99#[serde(tag = "type", rename_all = "lowercase")]
 100pub enum ModelMode {
 101    #[default]
 102    Default,
 103    Thinking {
 104        /// The maximum number of tokens to use for reasoning. Must be lower than the model's `max_output_tokens`.
 105        budget_tokens: Option<u64>,
 106    },
 107}
 108
 109impl From<ModelMode> for BedrockModelMode {
 110    fn from(value: ModelMode) -> Self {
 111        match value {
 112            ModelMode::Default => BedrockModelMode::Default,
 113            ModelMode::Thinking { budget_tokens } => BedrockModelMode::Thinking { budget_tokens },
 114        }
 115    }
 116}
 117
 118impl From<BedrockModelMode> for ModelMode {
 119    fn from(value: BedrockModelMode) -> Self {
 120        match value {
 121            BedrockModelMode::Default => ModelMode::Default,
 122            BedrockModelMode::Thinking { budget_tokens } => ModelMode::Thinking { budget_tokens },
 123        }
 124    }
 125}
 126
 127/// The URL of the base AWS service.
 128///
 129/// Right now we're just using this as the key to store the AWS credentials
 130/// under in the keychain.
 131const AMAZON_AWS_URL: &str = "https://amazonaws.com";
 132
 133// These environment variables all use a `ZED_` prefix because we don't want to overwrite the user's AWS credentials.
 134const ZED_BEDROCK_ACCESS_KEY_ID_VAR: &str = "ZED_ACCESS_KEY_ID";
 135const ZED_BEDROCK_SECRET_ACCESS_KEY_VAR: &str = "ZED_SECRET_ACCESS_KEY";
 136const ZED_BEDROCK_SESSION_TOKEN_VAR: &str = "ZED_SESSION_TOKEN";
 137const ZED_AWS_PROFILE_VAR: &str = "ZED_AWS_PROFILE";
 138const ZED_BEDROCK_REGION_VAR: &str = "ZED_AWS_REGION";
 139const ZED_AWS_CREDENTIALS_VAR: &str = "ZED_AWS_CREDENTIALS";
 140const ZED_AWS_ENDPOINT_VAR: &str = "ZED_AWS_ENDPOINT";
 141
 142pub struct State {
 143    credentials: Option<BedrockCredentials>,
 144    settings: Option<AmazonBedrockSettings>,
 145    credentials_from_env: bool,
 146    _subscription: Subscription,
 147}
 148
 149impl State {
 150    fn reset_credentials(&self, cx: &mut Context<Self>) -> Task<Result<()>> {
 151        let credentials_provider = <dyn CredentialsProvider>::global(cx);
 152        cx.spawn(async move |this, cx| {
 153            credentials_provider
 154                .delete_credentials(AMAZON_AWS_URL, &cx)
 155                .await
 156                .log_err();
 157            this.update(cx, |this, cx| {
 158                this.credentials = None;
 159                this.credentials_from_env = false;
 160                this.settings = None;
 161                cx.notify();
 162            })
 163        })
 164    }
 165
 166    fn set_credentials(
 167        &mut self,
 168        credentials: BedrockCredentials,
 169        cx: &mut Context<Self>,
 170    ) -> Task<Result<()>> {
 171        let credentials_provider = <dyn CredentialsProvider>::global(cx);
 172        cx.spawn(async move |this, cx| {
 173            credentials_provider
 174                .write_credentials(
 175                    AMAZON_AWS_URL,
 176                    "Bearer",
 177                    &serde_json::to_vec(&credentials)?,
 178                    &cx,
 179                )
 180                .await?;
 181            this.update(cx, |this, cx| {
 182                this.credentials = Some(credentials);
 183                cx.notify();
 184            })
 185        })
 186    }
 187
 188    fn is_authenticated(&self) -> bool {
 189        let derived = self
 190            .settings
 191            .as_ref()
 192            .and_then(|s| s.authentication_method.as_ref());
 193        let creds = self.credentials.as_ref();
 194
 195        derived.is_some() || creds.is_some()
 196    }
 197
 198    fn authenticate(&self, cx: &mut Context<Self>) -> Task<Result<(), AuthenticateError>> {
 199        if self.is_authenticated() {
 200            return Task::ready(Ok(()));
 201        }
 202
 203        let credentials_provider = <dyn CredentialsProvider>::global(cx);
 204        cx.spawn(async move |this, cx| {
 205            let (credentials, from_env) =
 206                if let Ok(credentials) = std::env::var(ZED_AWS_CREDENTIALS_VAR) {
 207                    (credentials, true)
 208                } else {
 209                    let (_, credentials) = credentials_provider
 210                        .read_credentials(AMAZON_AWS_URL, &cx)
 211                        .await?
 212                        .ok_or_else(|| AuthenticateError::CredentialsNotFound)?;
 213                    (
 214                        String::from_utf8(credentials)
 215                            .context("invalid {PROVIDER_NAME} credentials")?,
 216                        false,
 217                    )
 218                };
 219
 220            let credentials: BedrockCredentials =
 221                serde_json::from_str(&credentials).context("failed to parse credentials")?;
 222
 223            this.update(cx, |this, cx| {
 224                this.credentials = Some(credentials);
 225                this.credentials_from_env = from_env;
 226                cx.notify();
 227            })?;
 228
 229            Ok(())
 230        })
 231    }
 232
 233    fn get_region(&self) -> String {
 234        // Get region - from credentials or directly from settings
 235        let credentials_region = self.credentials.as_ref().map(|s| s.region.clone());
 236        let settings_region = self.settings.as_ref().and_then(|s| s.region.clone());
 237
 238        // Use credentials region if available, otherwise use settings region, finally fall back to default
 239        credentials_region
 240            .or(settings_region)
 241            .unwrap_or(String::from("us-east-1"))
 242    }
 243}
 244
 245pub struct BedrockLanguageModelProvider {
 246    http_client: AwsHttpClient,
 247    handler: tokio::runtime::Handle,
 248    state: gpui::Entity<State>,
 249}
 250
 251impl BedrockLanguageModelProvider {
 252    pub fn new(http_client: Arc<dyn HttpClient>, cx: &mut App) -> Self {
 253        let state = cx.new(|cx| State {
 254            credentials: None,
 255            settings: Some(AllLanguageModelSettings::get_global(cx).bedrock.clone()),
 256            credentials_from_env: false,
 257            _subscription: cx.observe_global::<SettingsStore>(|_, cx| {
 258                cx.notify();
 259            }),
 260        });
 261
 262        let tokio_handle = Tokio::handle(cx);
 263
 264        let coerced_client = AwsHttpClient::new(http_client.clone(), tokio_handle.clone());
 265
 266        Self {
 267            http_client: coerced_client,
 268            handler: tokio_handle.clone(),
 269            state,
 270        }
 271    }
 272
 273    fn create_language_model(&self, model: bedrock::Model) -> Arc<dyn LanguageModel> {
 274        Arc::new(BedrockModel {
 275            id: LanguageModelId::from(model.id().to_string()),
 276            model,
 277            http_client: self.http_client.clone(),
 278            handler: self.handler.clone(),
 279            state: self.state.clone(),
 280            client: OnceCell::new(),
 281            request_limiter: RateLimiter::new(4),
 282        })
 283    }
 284}
 285
 286impl LanguageModelProvider for BedrockLanguageModelProvider {
 287    fn id(&self) -> LanguageModelProviderId {
 288        LanguageModelProviderId(PROVIDER_ID.into())
 289    }
 290
 291    fn name(&self) -> LanguageModelProviderName {
 292        LanguageModelProviderName(PROVIDER_NAME.into())
 293    }
 294
 295    fn icon(&self) -> IconName {
 296        IconName::AiBedrock
 297    }
 298
 299    fn default_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {
 300        Some(self.create_language_model(bedrock::Model::default()))
 301    }
 302
 303    fn default_fast_model(&self, cx: &App) -> Option<Arc<dyn LanguageModel>> {
 304        let region = self.state.read(cx).get_region();
 305        Some(self.create_language_model(bedrock::Model::default_fast(region.as_str())))
 306    }
 307
 308    fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> {
 309        let mut models = BTreeMap::default();
 310
 311        for model in bedrock::Model::iter() {
 312            if !matches!(model, bedrock::Model::Custom { .. }) {
 313                // TODO: Sonnet 3.7 vs. 3.7 Thinking bug is here.
 314                models.insert(model.id().to_string(), model);
 315            }
 316        }
 317
 318        // Override with available models from settings
 319        for model in AllLanguageModelSettings::get_global(cx)
 320            .bedrock
 321            .available_models
 322            .iter()
 323        {
 324            models.insert(
 325                model.name.clone(),
 326                bedrock::Model::Custom {
 327                    name: model.name.clone(),
 328                    display_name: model.display_name.clone(),
 329                    max_tokens: model.max_tokens,
 330                    max_output_tokens: model.max_output_tokens,
 331                    default_temperature: model.default_temperature,
 332                    cache_configuration: model.cache_configuration.as_ref().map(|config| {
 333                        bedrock::BedrockModelCacheConfiguration {
 334                            max_cache_anchors: config.max_cache_anchors,
 335                            min_total_token: config.min_total_token,
 336                        }
 337                    }),
 338                },
 339            );
 340        }
 341
 342        models
 343            .into_values()
 344            .map(|model| self.create_language_model(model))
 345            .collect()
 346    }
 347
 348    fn is_authenticated(&self, cx: &App) -> bool {
 349        self.state.read(cx).is_authenticated()
 350    }
 351
 352    fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>> {
 353        self.state.update(cx, |state, cx| state.authenticate(cx))
 354    }
 355
 356    fn configuration_view(&self, window: &mut Window, cx: &mut App) -> AnyView {
 357        cx.new(|cx| ConfigurationView::new(self.state.clone(), window, cx))
 358            .into()
 359    }
 360
 361    fn reset_credentials(&self, cx: &mut App) -> Task<Result<()>> {
 362        self.state
 363            .update(cx, |state, cx| state.reset_credentials(cx))
 364    }
 365}
 366
 367impl LanguageModelProviderState for BedrockLanguageModelProvider {
 368    type ObservableEntity = State;
 369
 370    fn observable_entity(&self) -> Option<gpui::Entity<Self::ObservableEntity>> {
 371        Some(self.state.clone())
 372    }
 373}
 374
 375struct BedrockModel {
 376    id: LanguageModelId,
 377    model: Model,
 378    http_client: AwsHttpClient,
 379    handler: tokio::runtime::Handle,
 380    client: OnceCell<BedrockClient>,
 381    state: gpui::Entity<State>,
 382    request_limiter: RateLimiter,
 383}
 384
 385impl BedrockModel {
 386    fn get_or_init_client(&self, cx: &AsyncApp) -> anyhow::Result<&BedrockClient> {
 387        self.client
 388            .get_or_try_init_blocking(|| {
 389                let (auth_method, credentials, endpoint, region, settings) =
 390                    cx.read_entity(&self.state, |state, _cx| {
 391                        let auth_method = state
 392                            .settings
 393                            .as_ref()
 394                            .and_then(|s| s.authentication_method.clone());
 395
 396                        let endpoint = state.settings.as_ref().and_then(|s| s.endpoint.clone());
 397
 398                        let region = state.get_region();
 399
 400                        (
 401                            auth_method,
 402                            state.credentials.clone(),
 403                            endpoint,
 404                            region,
 405                            state.settings.clone(),
 406                        )
 407                    })?;
 408
 409                let mut config_builder = aws_config::defaults(BehaviorVersion::latest())
 410                    .stalled_stream_protection(StalledStreamProtectionConfig::disabled())
 411                    .http_client(self.http_client.clone())
 412                    .region(Region::new(region))
 413                    .timeout_config(TimeoutConfig::disabled());
 414
 415                if let Some(endpoint_url) = endpoint {
 416                    if !endpoint_url.is_empty() {
 417                        config_builder = config_builder.endpoint_url(endpoint_url);
 418                    }
 419                }
 420
 421                match auth_method {
 422                    None => {
 423                        if let Some(creds) = credentials {
 424                            let aws_creds = Credentials::new(
 425                                creds.access_key_id,
 426                                creds.secret_access_key,
 427                                creds.session_token,
 428                                None,
 429                                "zed-bedrock-provider",
 430                            );
 431                            config_builder = config_builder.credentials_provider(aws_creds);
 432                        }
 433                    }
 434                    Some(BedrockAuthMethod::NamedProfile)
 435                    | Some(BedrockAuthMethod::SingleSignOn) => {
 436                        // Currently NamedProfile and SSO behave the same way but only the instructions change
 437                        // Until we support BearerAuth through SSO, this will not change.
 438                        let profile_name = settings
 439                            .and_then(|s| s.profile_name)
 440                            .unwrap_or_else(|| "default".to_string());
 441
 442                        if !profile_name.is_empty() {
 443                            config_builder = config_builder.profile_name(profile_name);
 444                        }
 445                    }
 446                    Some(BedrockAuthMethod::Automatic) => {
 447                        // Use default credential provider chain
 448                    }
 449                }
 450
 451                let config = self.handler.block_on(config_builder.load());
 452                anyhow::Ok(BedrockClient::new(&config))
 453            })
 454            .context("initializing Bedrock client")?;
 455
 456        self.client.get().context("Bedrock client not initialized")
 457    }
 458
 459    fn stream_completion(
 460        &self,
 461        request: bedrock::Request,
 462        cx: &AsyncApp,
 463    ) -> Result<
 464        BoxFuture<'static, BoxStream<'static, Result<BedrockStreamingResponse, BedrockError>>>,
 465    > {
 466        let runtime_client = self
 467            .get_or_init_client(cx)
 468            .cloned()
 469            .context("Bedrock client not initialized")?;
 470        let owned_handle = self.handler.clone();
 471
 472        Ok(async move {
 473            let request = bedrock::stream_completion(runtime_client, request, owned_handle);
 474            request.await.unwrap_or_else(|e| {
 475                futures::stream::once(async move { Err(BedrockError::ClientError(e)) }).boxed()
 476            })
 477        }
 478        .boxed())
 479    }
 480}
 481
 482impl LanguageModel for BedrockModel {
 483    fn id(&self) -> LanguageModelId {
 484        self.id.clone()
 485    }
 486
 487    fn name(&self) -> LanguageModelName {
 488        LanguageModelName::from(self.model.display_name().to_string())
 489    }
 490
 491    fn provider_id(&self) -> LanguageModelProviderId {
 492        LanguageModelProviderId(PROVIDER_ID.into())
 493    }
 494
 495    fn provider_name(&self) -> LanguageModelProviderName {
 496        LanguageModelProviderName(PROVIDER_NAME.into())
 497    }
 498
 499    fn supports_tools(&self) -> bool {
 500        self.model.supports_tool_use()
 501    }
 502
 503    fn supports_images(&self) -> bool {
 504        false
 505    }
 506
 507    fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
 508        match choice {
 509            LanguageModelToolChoice::Auto | LanguageModelToolChoice::Any => {
 510                self.model.supports_tool_use()
 511            }
 512            // Add support for None - we'll filter tool calls at response
 513            LanguageModelToolChoice::None => self.model.supports_tool_use(),
 514        }
 515    }
 516
 517    fn telemetry_id(&self) -> String {
 518        format!("bedrock/{}", self.model.id())
 519    }
 520
 521    fn max_token_count(&self) -> u64 {
 522        self.model.max_token_count()
 523    }
 524
 525    fn max_output_tokens(&self) -> Option<u64> {
 526        Some(self.model.max_output_tokens())
 527    }
 528
 529    fn count_tokens(
 530        &self,
 531        request: LanguageModelRequest,
 532        cx: &App,
 533    ) -> BoxFuture<'static, Result<u64>> {
 534        get_bedrock_tokens(request, cx)
 535    }
 536
 537    fn stream_completion(
 538        &self,
 539        request: LanguageModelRequest,
 540        cx: &AsyncApp,
 541    ) -> BoxFuture<
 542        'static,
 543        Result<
 544            BoxStream<'static, Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>,
 545            LanguageModelCompletionError,
 546        >,
 547    > {
 548        let Ok(region) = cx.read_entity(&self.state, |state, _cx| state.get_region()) else {
 549            return async move { Err(anyhow::anyhow!("App State Dropped").into()) }.boxed();
 550        };
 551
 552        let model_id = match self.model.cross_region_inference_id(&region) {
 553            Ok(s) => s,
 554            Err(e) => {
 555                return async move { Err(e.into()) }.boxed();
 556            }
 557        };
 558
 559        let deny_tool_calls = request.tool_choice == Some(LanguageModelToolChoice::None);
 560
 561        let request = match into_bedrock(
 562            request,
 563            model_id,
 564            self.model.default_temperature(),
 565            self.model.max_output_tokens(),
 566            self.model.mode(),
 567            self.model.supports_caching(),
 568        ) {
 569            Ok(request) => request,
 570            Err(err) => return futures::future::ready(Err(err.into())).boxed(),
 571        };
 572
 573        let owned_handle = self.handler.clone();
 574
 575        let request = self.stream_completion(request, cx);
 576        let future = self.request_limiter.stream(async move {
 577            let response = request.map_err(|err| anyhow!(err))?.await;
 578            let events = map_to_language_model_completion_events(response, owned_handle);
 579
 580            if deny_tool_calls {
 581                Ok(deny_tool_use_events(events).boxed())
 582            } else {
 583                Ok(events.boxed())
 584            }
 585        });
 586
 587        async move { Ok(future.await?.boxed()) }.boxed()
 588    }
 589
 590    fn cache_configuration(&self) -> Option<LanguageModelCacheConfiguration> {
 591        self.model
 592            .cache_configuration()
 593            .map(|config| LanguageModelCacheConfiguration {
 594                max_cache_anchors: config.max_cache_anchors,
 595                should_speculate: false,
 596                min_total_token: config.min_total_token,
 597            })
 598    }
 599}
 600
 601fn deny_tool_use_events(
 602    events: impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>,
 603) -> impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>> {
 604    events.map(|event| {
 605        match event {
 606            Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) => {
 607                // Convert tool use to an error message if model decided to call it
 608                Ok(LanguageModelCompletionEvent::Text(format!(
 609                    "\n\n[Error: Tool calls are disabled in this context. Attempted to call '{}']",
 610                    tool_use.name
 611                )))
 612            }
 613            other => other,
 614        }
 615    })
 616}
 617
 618pub fn into_bedrock(
 619    request: LanguageModelRequest,
 620    model: String,
 621    default_temperature: f32,
 622    max_output_tokens: u64,
 623    mode: BedrockModelMode,
 624    supports_caching: bool,
 625) -> Result<bedrock::Request> {
 626    let mut new_messages: Vec<BedrockMessage> = Vec::new();
 627    let mut system_message = String::new();
 628
 629    for message in request.messages {
 630        if message.contents_empty() {
 631            continue;
 632        }
 633
 634        match message.role {
 635            Role::User | Role::Assistant => {
 636                let mut bedrock_message_content: Vec<BedrockInnerContent> = message
 637                    .content
 638                    .into_iter()
 639                    .filter_map(|content| match content {
 640                        MessageContent::Text(text) => {
 641                            if !text.is_empty() {
 642                                Some(BedrockInnerContent::Text(text))
 643                            } else {
 644                                None
 645                            }
 646                        }
 647                        MessageContent::Thinking { text, signature } => {
 648                            if model.contains(Model::DeepSeekR1.request_id()) {
 649                                // DeepSeekR1 doesn't support thinking blocks
 650                                // And the AWS API demands that you strip them
 651                                return None;
 652                            }
 653                            let thinking = BedrockThinkingTextBlock::builder()
 654                                .text(text)
 655                                .set_signature(signature)
 656                                .build()
 657                                .context("failed to build reasoning block")
 658                                .log_err()?;
 659
 660                            Some(BedrockInnerContent::ReasoningContent(
 661                                BedrockThinkingBlock::ReasoningText(thinking),
 662                            ))
 663                        }
 664                        MessageContent::RedactedThinking(blob) => {
 665                            if model.contains(Model::DeepSeekR1.request_id()) {
 666                                // DeepSeekR1 doesn't support thinking blocks
 667                                // And the AWS API demands that you strip them
 668                                return None;
 669                            }
 670                            let redacted =
 671                                BedrockThinkingBlock::RedactedContent(BedrockBlob::new(blob));
 672
 673                            Some(BedrockInnerContent::ReasoningContent(redacted))
 674                        }
 675                        MessageContent::ToolUse(tool_use) => {
 676                            let input = if tool_use.input.is_null() {
 677                                // Bedrock API requires valid JsonValue, not null, for tool use input
 678                                value_to_aws_document(&serde_json::json!({}))
 679                            } else {
 680                                value_to_aws_document(&tool_use.input)
 681                            };
 682                            BedrockToolUseBlock::builder()
 683                                .name(tool_use.name.to_string())
 684                                .tool_use_id(tool_use.id.to_string())
 685                                .input(input)
 686                                .build()
 687                                .context("failed to build Bedrock tool use block")
 688                                .log_err()
 689                                .map(BedrockInnerContent::ToolUse)
 690                        },
 691                        MessageContent::ToolResult(tool_result) => {
 692                            BedrockToolResultBlock::builder()
 693                                .tool_use_id(tool_result.tool_use_id.to_string())
 694                                .content(match tool_result.content {
 695                                    LanguageModelToolResultContent::Text(text) => {
 696                                        BedrockToolResultContentBlock::Text(text.to_string())
 697                                    }
 698                                    LanguageModelToolResultContent::Image(_) => {
 699                                        BedrockToolResultContentBlock::Text(
 700                                            // TODO: Bedrock image support
 701                                            "[Tool responded with an image, but Zed doesn't support these in Bedrock models yet]".to_string()
 702                                        )
 703                                    }
 704                                })
 705                                .status({
 706                                    if tool_result.is_error {
 707                                        BedrockToolResultStatus::Error
 708                                    } else {
 709                                        BedrockToolResultStatus::Success
 710                                    }
 711                                })
 712                                .build()
 713                                .context("failed to build Bedrock tool result block")
 714                                .log_err()
 715                                .map(BedrockInnerContent::ToolResult)
 716                        }
 717                        _ => None,
 718                    })
 719                    .collect();
 720                if message.cache && supports_caching {
 721                    bedrock_message_content.push(BedrockInnerContent::CachePoint(
 722                        CachePointBlock::builder()
 723                            .r#type(CachePointType::Default)
 724                            .build()
 725                            .context("failed to build cache point block")?,
 726                    ));
 727                }
 728                let bedrock_role = match message.role {
 729                    Role::User => bedrock::BedrockRole::User,
 730                    Role::Assistant => bedrock::BedrockRole::Assistant,
 731                    Role::System => unreachable!("System role should never occur here"),
 732                };
 733                if let Some(last_message) = new_messages.last_mut() {
 734                    if last_message.role == bedrock_role {
 735                        last_message.content.extend(bedrock_message_content);
 736                        continue;
 737                    }
 738                }
 739                new_messages.push(
 740                    BedrockMessage::builder()
 741                        .role(bedrock_role)
 742                        .set_content(Some(bedrock_message_content))
 743                        .build()
 744                        .context("failed to build Bedrock message")?,
 745                );
 746            }
 747            Role::System => {
 748                if !system_message.is_empty() {
 749                    system_message.push_str("\n\n");
 750                }
 751                system_message.push_str(&message.string_contents());
 752            }
 753        }
 754    }
 755
 756    let mut tool_spec: Vec<BedrockTool> = request
 757        .tools
 758        .iter()
 759        .filter_map(|tool| {
 760            Some(BedrockTool::ToolSpec(
 761                BedrockToolSpec::builder()
 762                    .name(tool.name.clone())
 763                    .description(tool.description.clone())
 764                    .input_schema(BedrockToolInputSchema::Json(value_to_aws_document(
 765                        &tool.input_schema,
 766                    )))
 767                    .build()
 768                    .log_err()?,
 769            ))
 770        })
 771        .collect();
 772
 773    if !tool_spec.is_empty() && supports_caching {
 774        tool_spec.push(BedrockTool::CachePoint(
 775            CachePointBlock::builder()
 776                .r#type(CachePointType::Default)
 777                .build()
 778                .context("failed to build cache point block")?,
 779        ));
 780    }
 781
 782    let tool_choice = match request.tool_choice {
 783        Some(LanguageModelToolChoice::Auto) | None => {
 784            BedrockToolChoice::Auto(BedrockAutoToolChoice::builder().build())
 785        }
 786        Some(LanguageModelToolChoice::Any) => {
 787            BedrockToolChoice::Any(BedrockAnyToolChoice::builder().build())
 788        }
 789        Some(LanguageModelToolChoice::None) => {
 790            // For None, we still use Auto but will filter out tool calls in the response
 791            BedrockToolChoice::Auto(BedrockAutoToolChoice::builder().build())
 792        }
 793    };
 794    let tool_config: BedrockToolConfig = BedrockToolConfig::builder()
 795        .set_tools(Some(tool_spec))
 796        .tool_choice(tool_choice)
 797        .build()?;
 798
 799    Ok(bedrock::Request {
 800        model,
 801        messages: new_messages,
 802        max_tokens: max_output_tokens,
 803        system: Some(system_message),
 804        tools: Some(tool_config),
 805        thinking: if let BedrockModelMode::Thinking { budget_tokens } = mode {
 806            Some(bedrock::Thinking::Enabled { budget_tokens })
 807        } else {
 808            None
 809        },
 810        metadata: None,
 811        stop_sequences: Vec::new(),
 812        temperature: request.temperature.or(Some(default_temperature)),
 813        top_k: None,
 814        top_p: None,
 815    })
 816}
 817
 818// TODO: just call the ConverseOutput.usage() method:
 819// https://docs.rs/aws-sdk-bedrockruntime/latest/aws_sdk_bedrockruntime/operation/converse/struct.ConverseOutput.html#method.output
 820pub fn get_bedrock_tokens(
 821    request: LanguageModelRequest,
 822    cx: &App,
 823) -> BoxFuture<'static, Result<u64>> {
 824    cx.background_executor()
 825        .spawn(async move {
 826            let messages = request.messages;
 827            let mut tokens_from_images = 0;
 828            let mut string_messages = Vec::with_capacity(messages.len());
 829
 830            for message in messages {
 831                use language_model::MessageContent;
 832
 833                let mut string_contents = String::new();
 834
 835                for content in message.content {
 836                    match content {
 837                        MessageContent::Text(text) | MessageContent::Thinking { text, .. } => {
 838                            string_contents.push_str(&text);
 839                        }
 840                        MessageContent::RedactedThinking(_) => {}
 841                        MessageContent::Image(image) => {
 842                            tokens_from_images += image.estimate_tokens();
 843                        }
 844                        MessageContent::ToolUse(_tool_use) => {
 845                            // TODO: Estimate token usage from tool uses.
 846                        }
 847                        MessageContent::ToolResult(tool_result) => match tool_result.content {
 848                            LanguageModelToolResultContent::Text(text) => {
 849                                string_contents.push_str(&text);
 850                            }
 851                            LanguageModelToolResultContent::Image(image) => {
 852                                tokens_from_images += image.estimate_tokens();
 853                            }
 854                        },
 855                    }
 856                }
 857
 858                if !string_contents.is_empty() {
 859                    string_messages.push(tiktoken_rs::ChatCompletionRequestMessage {
 860                        role: match message.role {
 861                            Role::User => "user".into(),
 862                            Role::Assistant => "assistant".into(),
 863                            Role::System => "system".into(),
 864                        },
 865                        content: Some(string_contents),
 866                        name: None,
 867                        function_call: None,
 868                    });
 869                }
 870            }
 871
 872            // Tiktoken doesn't yet support these models, so we manually use the
 873            // same tokenizer as GPT-4.
 874            tiktoken_rs::num_tokens_from_messages("gpt-4", &string_messages)
 875                .map(|tokens| (tokens + tokens_from_images) as u64)
 876        })
 877        .boxed()
 878}
 879
 880pub fn map_to_language_model_completion_events(
 881    events: Pin<Box<dyn Send + Stream<Item = Result<BedrockStreamingResponse, BedrockError>>>>,
 882    handle: Handle,
 883) -> impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>> {
 884    struct RawToolUse {
 885        id: String,
 886        name: String,
 887        input_json: String,
 888    }
 889
 890    struct State {
 891        events: Pin<Box<dyn Send + Stream<Item = Result<BedrockStreamingResponse, BedrockError>>>>,
 892        tool_uses_by_index: HashMap<i32, RawToolUse>,
 893    }
 894
 895    futures::stream::unfold(
 896        State {
 897            events,
 898            tool_uses_by_index: HashMap::default(),
 899        },
 900        move |mut state: State| {
 901            let inner_handle = handle.clone();
 902            async move {
 903                inner_handle
 904                    .spawn(async {
 905                        while let Some(event) = state.events.next().await {
 906                            match event {
 907                                Ok(event) => match event {
 908                                    ConverseStreamOutput::ContentBlockDelta(cb_delta) => {
 909                                        match cb_delta.delta {
 910                                            Some(ContentBlockDelta::Text(text_out)) => {
 911                                                let completion_event =
 912                                                    LanguageModelCompletionEvent::Text(text_out);
 913                                                return Some((Some(Ok(completion_event)), state));
 914                                            }
 915
 916                                            Some(ContentBlockDelta::ToolUse(text_out)) => {
 917                                                if let Some(tool_use) = state
 918                                                    .tool_uses_by_index
 919                                                    .get_mut(&cb_delta.content_block_index)
 920                                                {
 921                                                    tool_use.input_json.push_str(text_out.input());
 922                                                }
 923                                            }
 924
 925                                            Some(ContentBlockDelta::ReasoningContent(thinking)) => {
 926                                                match thinking {
 927                                                    ReasoningContentBlockDelta::RedactedContent(
 928                                                        redacted,
 929                                                    ) => {
 930                                                        let thinking_event =
 931                                                            LanguageModelCompletionEvent::Thinking {
 932                                                                text: String::from_utf8(
 933                                                                    redacted.into_inner(),
 934                                                                )
 935                                                                .unwrap_or("REDACTED".to_string()),
 936                                                                signature: None,
 937                                                            };
 938
 939                                                        return Some((
 940                                                            Some(Ok(thinking_event)),
 941                                                            state,
 942                                                        ));
 943                                                    }
 944                                                    ReasoningContentBlockDelta::Signature(
 945                                                        signature,
 946                                                    ) => {
 947                                                        return Some((
 948                                                            Some(Ok(LanguageModelCompletionEvent::Thinking {
 949                                                                text: "".to_string(),
 950                                                                signature: Some(signature)
 951                                                            })),
 952                                                            state,
 953                                                        ));
 954                                                    }
 955                                                    ReasoningContentBlockDelta::Text(thoughts) => {
 956                                                        let thinking_event =
 957                                                            LanguageModelCompletionEvent::Thinking {
 958                                                                text: thoughts.to_string(),
 959                                                                signature: None
 960                                                            };
 961
 962                                                        return Some((
 963                                                            Some(Ok(thinking_event)),
 964                                                            state,
 965                                                        ));
 966                                                    }
 967                                                    _ => {}
 968                                                }
 969                                            }
 970                                            _ => {}
 971                                        }
 972                                    }
 973                                    ConverseStreamOutput::ContentBlockStart(cb_start) => {
 974                                        if let Some(ContentBlockStart::ToolUse(text_out)) =
 975                                            cb_start.start
 976                                        {
 977                                            let tool_use = RawToolUse {
 978                                                id: text_out.tool_use_id,
 979                                                name: text_out.name,
 980                                                input_json: String::new(),
 981                                            };
 982
 983                                            state
 984                                                .tool_uses_by_index
 985                                                .insert(cb_start.content_block_index, tool_use);
 986                                        }
 987                                    }
 988                                    ConverseStreamOutput::ContentBlockStop(cb_stop) => {
 989                                        if let Some(tool_use) = state
 990                                            .tool_uses_by_index
 991                                            .remove(&cb_stop.content_block_index)
 992                                        {
 993                                            let tool_use_event = LanguageModelToolUse {
 994                                                id: tool_use.id.into(),
 995                                                name: tool_use.name.into(),
 996                                                is_input_complete: true,
 997                                                raw_input: tool_use.input_json.clone(),
 998                                                input: if tool_use.input_json.is_empty() {
 999                                                    Value::Null
1000                                                } else {
1001                                                    serde_json::Value::from_str(
1002                                                        &tool_use.input_json,
1003                                                    )
1004                                                    .map_err(|err| anyhow!(err))
1005                                                    .unwrap()
1006                                                },
1007                                            };
1008
1009                                            return Some((
1010                                                Some(Ok(LanguageModelCompletionEvent::ToolUse(
1011                                                    tool_use_event,
1012                                                ))),
1013                                                state,
1014                                            ));
1015                                        }
1016                                    }
1017
1018                                    ConverseStreamOutput::Metadata(cb_meta) => {
1019                                        if let Some(metadata) = cb_meta.usage {
1020                                            let completion_event =
1021                                                LanguageModelCompletionEvent::UsageUpdate(
1022                                                    TokenUsage {
1023                                                        input_tokens: metadata.input_tokens as u64,
1024                                                        output_tokens: metadata.output_tokens as u64,
1025                                                        cache_creation_input_tokens:
1026                                                            metadata.cache_write_input_tokens.unwrap_or_default() as u64,
1027                                                        cache_read_input_tokens:
1028                                                            metadata.cache_read_input_tokens.unwrap_or_default() as u64,
1029                                                    },
1030                                                );
1031                                            return Some((Some(Ok(completion_event)), state));
1032                                        }
1033                                    }
1034                                    ConverseStreamOutput::MessageStop(message_stop) => {
1035                                        let reason = match message_stop.stop_reason {
1036                                            StopReason::ContentFiltered => {
1037                                                LanguageModelCompletionEvent::Stop(
1038                                                    language_model::StopReason::EndTurn,
1039                                                )
1040                                            }
1041                                            StopReason::EndTurn => {
1042                                                LanguageModelCompletionEvent::Stop(
1043                                                    language_model::StopReason::EndTurn,
1044                                                )
1045                                            }
1046                                            StopReason::GuardrailIntervened => {
1047                                                LanguageModelCompletionEvent::Stop(
1048                                                    language_model::StopReason::EndTurn,
1049                                                )
1050                                            }
1051                                            StopReason::MaxTokens => {
1052                                                LanguageModelCompletionEvent::Stop(
1053                                                    language_model::StopReason::EndTurn,
1054                                                )
1055                                            }
1056                                            StopReason::StopSequence => {
1057                                                LanguageModelCompletionEvent::Stop(
1058                                                    language_model::StopReason::EndTurn,
1059                                                )
1060                                            }
1061                                            StopReason::ToolUse => {
1062                                                LanguageModelCompletionEvent::Stop(
1063                                                    language_model::StopReason::ToolUse,
1064                                                )
1065                                            }
1066                                            _ => LanguageModelCompletionEvent::Stop(
1067                                                language_model::StopReason::EndTurn,
1068                                            ),
1069                                        };
1070                                        return Some((Some(Ok(reason)), state));
1071                                    }
1072                                    _ => {}
1073                                },
1074
1075                                Err(err) => return Some((Some(Err(anyhow!(err).into())), state)),
1076                            }
1077                        }
1078                        None
1079                    })
1080                    .await
1081                    .log_err()
1082                    .flatten()
1083            }
1084        },
1085    )
1086    .filter_map(|event| async move { event })
1087}
1088
1089struct ConfigurationView {
1090    access_key_id_editor: Entity<Editor>,
1091    secret_access_key_editor: Entity<Editor>,
1092    session_token_editor: Entity<Editor>,
1093    region_editor: Entity<Editor>,
1094    state: gpui::Entity<State>,
1095    load_credentials_task: Option<Task<()>>,
1096}
1097
1098impl ConfigurationView {
1099    const PLACEHOLDER_ACCESS_KEY_ID_TEXT: &'static str = "XXXXXXXXXXXXXXXX";
1100    const PLACEHOLDER_SECRET_ACCESS_KEY_TEXT: &'static str =
1101        "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX";
1102    const PLACEHOLDER_SESSION_TOKEN_TEXT: &'static str = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX";
1103    const PLACEHOLDER_REGION: &'static str = "us-east-1";
1104
1105    fn new(state: gpui::Entity<State>, window: &mut Window, cx: &mut Context<Self>) -> Self {
1106        cx.observe(&state, |_, _, cx| {
1107            cx.notify();
1108        })
1109        .detach();
1110
1111        let load_credentials_task = Some(cx.spawn({
1112            let state = state.clone();
1113            async move |this, cx| {
1114                if let Some(task) = state
1115                    .update(cx, |state, cx| state.authenticate(cx))
1116                    .log_err()
1117                {
1118                    // We don't log an error, because "not signed in" is also an error.
1119                    let _ = task.await;
1120                }
1121                this.update(cx, |this, cx| {
1122                    this.load_credentials_task = None;
1123                    cx.notify();
1124                })
1125                .log_err();
1126            }
1127        }));
1128
1129        Self {
1130            access_key_id_editor: cx.new(|cx| {
1131                let mut editor = Editor::single_line(window, cx);
1132                editor.set_placeholder_text(Self::PLACEHOLDER_ACCESS_KEY_ID_TEXT, cx);
1133                editor
1134            }),
1135            secret_access_key_editor: cx.new(|cx| {
1136                let mut editor = Editor::single_line(window, cx);
1137                editor.set_placeholder_text(Self::PLACEHOLDER_SECRET_ACCESS_KEY_TEXT, cx);
1138                editor
1139            }),
1140            session_token_editor: cx.new(|cx| {
1141                let mut editor = Editor::single_line(window, cx);
1142                editor.set_placeholder_text(Self::PLACEHOLDER_SESSION_TOKEN_TEXT, cx);
1143                editor
1144            }),
1145            region_editor: cx.new(|cx| {
1146                let mut editor = Editor::single_line(window, cx);
1147                editor.set_placeholder_text(Self::PLACEHOLDER_REGION, cx);
1148                editor
1149            }),
1150            state,
1151            load_credentials_task,
1152        }
1153    }
1154
1155    fn save_credentials(
1156        &mut self,
1157        _: &menu::Confirm,
1158        _window: &mut Window,
1159        cx: &mut Context<Self>,
1160    ) {
1161        let access_key_id = self
1162            .access_key_id_editor
1163            .read(cx)
1164            .text(cx)
1165            .to_string()
1166            .trim()
1167            .to_string();
1168        let secret_access_key = self
1169            .secret_access_key_editor
1170            .read(cx)
1171            .text(cx)
1172            .to_string()
1173            .trim()
1174            .to_string();
1175        let session_token = self
1176            .session_token_editor
1177            .read(cx)
1178            .text(cx)
1179            .to_string()
1180            .trim()
1181            .to_string();
1182        let session_token = if session_token.is_empty() {
1183            None
1184        } else {
1185            Some(session_token)
1186        };
1187        let region = self
1188            .region_editor
1189            .read(cx)
1190            .text(cx)
1191            .to_string()
1192            .trim()
1193            .to_string();
1194        let region = if region.is_empty() {
1195            "us-east-1".to_string()
1196        } else {
1197            region
1198        };
1199
1200        let state = self.state.clone();
1201        cx.spawn(async move |_, cx| {
1202            state
1203                .update(cx, |state, cx| {
1204                    let credentials: BedrockCredentials = BedrockCredentials {
1205                        region: region.clone(),
1206                        access_key_id: access_key_id.clone(),
1207                        secret_access_key: secret_access_key.clone(),
1208                        session_token: session_token.clone(),
1209                    };
1210
1211                    state.set_credentials(credentials, cx)
1212                })?
1213                .await
1214        })
1215        .detach_and_log_err(cx);
1216    }
1217
1218    fn reset_credentials(&mut self, window: &mut Window, cx: &mut Context<Self>) {
1219        self.access_key_id_editor
1220            .update(cx, |editor, cx| editor.set_text("", window, cx));
1221        self.secret_access_key_editor
1222            .update(cx, |editor, cx| editor.set_text("", window, cx));
1223        self.session_token_editor
1224            .update(cx, |editor, cx| editor.set_text("", window, cx));
1225        self.region_editor
1226            .update(cx, |editor, cx| editor.set_text("", window, cx));
1227
1228        let state = self.state.clone();
1229        cx.spawn(async move |_, cx| {
1230            state
1231                .update(cx, |state, cx| state.reset_credentials(cx))?
1232                .await
1233        })
1234        .detach_and_log_err(cx);
1235    }
1236
1237    fn make_text_style(&self, cx: &Context<Self>) -> TextStyle {
1238        let settings = ThemeSettings::get_global(cx);
1239        TextStyle {
1240            color: cx.theme().colors().text,
1241            font_family: settings.ui_font.family.clone(),
1242            font_features: settings.ui_font.features.clone(),
1243            font_fallbacks: settings.ui_font.fallbacks.clone(),
1244            font_size: rems(0.875).into(),
1245            font_weight: settings.ui_font.weight,
1246            font_style: FontStyle::Normal,
1247            line_height: relative(1.3),
1248            background_color: None,
1249            underline: None,
1250            strikethrough: None,
1251            white_space: WhiteSpace::Normal,
1252            text_overflow: None,
1253            text_align: Default::default(),
1254            line_clamp: None,
1255        }
1256    }
1257
1258    fn make_input_styles(&self, cx: &Context<Self>) -> Div {
1259        let bg_color = cx.theme().colors().editor_background;
1260        let border_color = cx.theme().colors().border;
1261
1262        h_flex()
1263            .w_full()
1264            .px_2()
1265            .py_1()
1266            .bg(bg_color)
1267            .border_1()
1268            .border_color(border_color)
1269            .rounded_sm()
1270    }
1271
1272    fn should_render_editor(&self, cx: &Context<Self>) -> bool {
1273        self.state.read(cx).is_authenticated()
1274    }
1275}
1276
1277impl Render for ConfigurationView {
1278    fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
1279        let env_var_set = self.state.read(cx).credentials_from_env;
1280        let bedrock_settings = self.state.read(cx).settings.as_ref();
1281        let bedrock_method = bedrock_settings
1282            .as_ref()
1283            .and_then(|s| s.authentication_method.clone());
1284
1285        if self.load_credentials_task.is_some() {
1286            return div().child(Label::new("Loading credentials...")).into_any();
1287        }
1288
1289        if self.should_render_editor(cx) {
1290            return h_flex()
1291                .mt_1()
1292                .p_1()
1293                .justify_between()
1294                .rounded_md()
1295                .border_1()
1296                .border_color(cx.theme().colors().border)
1297                .bg(cx.theme().colors().background)
1298                .child(
1299                    h_flex()
1300                        .gap_1()
1301                        .child(Icon::new(IconName::Check).color(Color::Success))
1302                        .child(Label::new(if env_var_set {
1303                            format!("Access Key ID is set in {ZED_BEDROCK_ACCESS_KEY_ID_VAR}, Secret Key is set in {ZED_BEDROCK_SECRET_ACCESS_KEY_VAR}, Region is set in {ZED_BEDROCK_REGION_VAR} environment variables.")
1304                        } else {
1305                            match bedrock_method {
1306                                Some(BedrockAuthMethod::Automatic) => "You are using automatic credentials".into(),
1307                                Some(BedrockAuthMethod::NamedProfile) => {
1308                                    "You are using named profile".into()
1309                                },
1310                                Some(BedrockAuthMethod::SingleSignOn) => "You are using a single sign on profile".into(),
1311                                None => "You are using static credentials".into(),
1312                            }
1313                        })),
1314                )
1315                .child(
1316                    Button::new("reset-key", "Reset Key")
1317                        .icon(Some(IconName::Trash))
1318                        .icon_size(IconSize::Small)
1319                        .icon_position(IconPosition::Start)
1320                        .disabled(env_var_set || bedrock_method.is_some())
1321                        .when(env_var_set, |this| {
1322                            this.tooltip(Tooltip::text(format!("To reset your credentials, unset the {ZED_BEDROCK_ACCESS_KEY_ID_VAR}, {ZED_BEDROCK_SECRET_ACCESS_KEY_VAR}, and {ZED_BEDROCK_REGION_VAR} environment variables.")))
1323                        })
1324                        .when(bedrock_method.is_some(), |this| {
1325                            this.tooltip(Tooltip::text("You cannot reset credentials as they're being derived, check Zed settings to understand how"))
1326                        })
1327                        .on_click(cx.listener(|this, _, window, cx| this.reset_credentials(window, cx))),
1328                )
1329                .into_any();
1330        }
1331
1332        v_flex()
1333            .size_full()
1334            .on_action(cx.listener(ConfigurationView::save_credentials))
1335            .child(Label::new("To use Zed's assistant with Bedrock, you can set a custom authentication strategy through the settings.json, or use static credentials."))
1336            .child(Label::new("But, to access models on AWS, you need to:").mt_1())
1337            .child(
1338                List::new()
1339                    .child(
1340                        InstructionListItem::new(
1341                            "Grant permissions to the strategy you'll use according to the:",
1342                            Some("Prerequisites"),
1343                            Some("https://docs.aws.amazon.com/bedrock/latest/userguide/inference-prereq.html"),
1344                        )
1345                    )
1346                    .child(
1347                        InstructionListItem::new(
1348                            "Select the models you would like access to:",
1349                            Some("Bedrock Model Catalog"),
1350                            Some("https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/modelaccess"),
1351                        )
1352                    )
1353            )
1354            .child(self.render_static_credentials_ui(cx))
1355            .child(self.render_common_fields(cx))
1356            .child(
1357                Label::new(
1358                    format!("You can also assign the {ZED_BEDROCK_ACCESS_KEY_ID_VAR}, {ZED_BEDROCK_SECRET_ACCESS_KEY_VAR} AND {ZED_BEDROCK_REGION_VAR} environment variables and restart Zed."),
1359                )
1360                    .size(LabelSize::Small)
1361                    .color(Color::Muted)
1362                    .my_1(),
1363            )
1364            .child(
1365                Label::new(
1366                    format!("Optionally, if your environment uses AWS CLI profiles, you can set {ZED_AWS_PROFILE_VAR}; if it requires a custom endpoint, you can set {ZED_AWS_ENDPOINT_VAR}; and if it requires a Session Token, you can set {ZED_BEDROCK_SESSION_TOKEN_VAR}."),
1367                )
1368                    .size(LabelSize::Small)
1369                    .color(Color::Muted),
1370            )
1371            .into_any()
1372    }
1373}
1374
1375impl ConfigurationView {
1376    fn render_access_key_id_editor(&self, cx: &mut Context<Self>) -> impl IntoElement {
1377        let text_style = self.make_text_style(cx);
1378
1379        EditorElement::new(
1380            &self.access_key_id_editor,
1381            EditorStyle {
1382                background: cx.theme().colors().editor_background,
1383                local_player: cx.theme().players().local(),
1384                text: text_style,
1385                ..Default::default()
1386            },
1387        )
1388    }
1389
1390    fn render_secret_key_editor(&self, cx: &mut Context<Self>) -> impl IntoElement {
1391        let text_style = self.make_text_style(cx);
1392
1393        EditorElement::new(
1394            &self.secret_access_key_editor,
1395            EditorStyle {
1396                background: cx.theme().colors().editor_background,
1397                local_player: cx.theme().players().local(),
1398                text: text_style,
1399                ..Default::default()
1400            },
1401        )
1402    }
1403
1404    fn render_session_token_editor(&self, cx: &mut Context<Self>) -> impl IntoElement {
1405        let text_style = self.make_text_style(cx);
1406
1407        EditorElement::new(
1408            &self.session_token_editor,
1409            EditorStyle {
1410                background: cx.theme().colors().editor_background,
1411                local_player: cx.theme().players().local(),
1412                text: text_style,
1413                ..Default::default()
1414            },
1415        )
1416    }
1417
1418    fn render_region_editor(&self, cx: &mut Context<Self>) -> impl IntoElement {
1419        let text_style = self.make_text_style(cx);
1420
1421        EditorElement::new(
1422            &self.region_editor,
1423            EditorStyle {
1424                background: cx.theme().colors().editor_background,
1425                local_player: cx.theme().players().local(),
1426                text: text_style,
1427                ..Default::default()
1428            },
1429        )
1430    }
1431
1432    fn render_static_credentials_ui(&self, cx: &mut Context<Self>) -> AnyElement {
1433        v_flex()
1434            .my_2()
1435            .gap_1p5()
1436            .child(
1437                Label::new("Static Keys")
1438                    .size(LabelSize::Default)
1439                    .weight(FontWeight::BOLD),
1440            )
1441            .child(
1442                Label::new(
1443                    "This method uses your AWS access key ID and secret access key directly.",
1444                )
1445            )
1446            .child(
1447                List::new()
1448                    .child(InstructionListItem::new(
1449                        "Create an IAM user in the AWS console with programmatic access",
1450                        Some("IAM Console"),
1451                        Some("https://us-east-1.console.aws.amazon.com/iam/home?region=us-east-1#/users"),
1452                    ))
1453                    .child(InstructionListItem::new(
1454                        "Attach the necessary Bedrock permissions to this ",
1455                        Some("user"),
1456                        Some("https://docs.aws.amazon.com/bedrock/latest/userguide/inference-prereq.html"),
1457                    ))
1458                    .child(InstructionListItem::text_only(
1459                        "Copy the access key ID and secret access key when provided",
1460                    ))
1461                    .child(InstructionListItem::text_only(
1462                        "Enter these credentials below",
1463                    )),
1464            )
1465            .child(
1466                v_flex()
1467                    .gap_0p5()
1468                    .child(Label::new("Access Key ID").size(LabelSize::Small))
1469                    .child(
1470                        self.make_input_styles(cx)
1471                            .child(self.render_access_key_id_editor(cx)),
1472                    ),
1473            )
1474            .child(
1475                v_flex()
1476                    .gap_0p5()
1477                    .child(Label::new("Secret Access Key").size(LabelSize::Small))
1478                    .child(self.make_input_styles(cx).child(self.render_secret_key_editor(cx))),
1479            )
1480            .child(
1481                v_flex()
1482                    .gap_0p5()
1483                    .child(Label::new("Session Token (Optional)").size(LabelSize::Small))
1484                    .child(
1485                        self.make_input_styles(cx)
1486                            .child(self.render_session_token_editor(cx)),
1487                    ),
1488            )
1489            .into_any_element()
1490    }
1491
1492    fn render_common_fields(&self, cx: &mut Context<Self>) -> AnyElement {
1493        v_flex()
1494            .gap_0p5()
1495            .child(Label::new("Region").size(LabelSize::Small))
1496            .child(
1497                self.make_input_styles(cx)
1498                    .child(self.render_region_editor(cx)),
1499            )
1500            .into_any_element()
1501    }
1502}