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