bedrock.rs

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