bedrock.rs

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