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