bedrock.rs

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