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