diff --git a/Cargo.lock b/Cargo.lock index 20f19de832c567c6866731f128f049bd77d7be57..de9cb227c6cfb799099abf446c1bdee61ec85bff 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1441,9 +1441,9 @@ dependencies = [ [[package]] name = "aws-config" -version = "1.8.8" +version = "1.8.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "37cf2b6af2a95a20e266782b4f76f1a5e12bf412a9db2de9c1e9123b9d8c0ad8" +checksum = "1856b1b48b65f71a4dd940b1c0931f9a7b646d4a924b9828ffefc1454714668a" dependencies = [ "aws-credential-types", "aws-runtime", @@ -1507,9 +1507,9 @@ dependencies = [ [[package]] name = "aws-runtime" -version = "1.5.12" +version = "1.5.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bfa006bb32360ed90ac51203feafb9d02e3d21046e1fd3a450a404b90ea73e5d" +checksum = "9f2402da1a5e16868ba98725e5d73f26b8116eaa892e56f2cd0bf5eec7985f70" dependencies = [ "aws-credential-types", "aws-sigv4", @@ -1532,9 +1532,9 @@ dependencies = [ [[package]] name = "aws-sdk-bedrockruntime" -version = "1.109.0" +version = "1.112.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fbfdfd941dcb253c17bf70baddbf1e5b22f19e29d313d2e049bad4b1dadb2011" +checksum = "c06c037e6823696d752702ec2bad758d3cf95d1b92b712c8ac7e93824b5e2391" dependencies = [ "aws-credential-types", "aws-runtime", @@ -1614,9 +1614,9 @@ dependencies = [ [[package]] name = "aws-sdk-sso" -version = "1.86.0" +version = "1.88.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4a0abbfab841446cce6e87af853a3ba2cc1bc9afcd3f3550dd556c43d434c86d" +checksum = "d05b276777560aa9a196dbba2e3aada4d8006d3d7eeb3ba7fe0c317227d933c4" dependencies = [ "aws-credential-types", "aws-runtime", @@ -1636,9 +1636,9 @@ dependencies = [ [[package]] name = "aws-sdk-ssooidc" -version = "1.88.0" +version = "1.90.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a68d675582afea0e94d38b6ca9c5aaae4ca14f1d36faa6edb19b42e687e70d7" +checksum = "f9be14d6d9cd761fac3fd234a0f47f7ed6c0df62d83c0eeb7012750e4732879b" dependencies = [ "aws-credential-types", "aws-runtime", @@ -1658,9 +1658,9 @@ dependencies = [ [[package]] name = "aws-sdk-sts" -version = "1.88.0" +version = "1.90.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d30990923f4f675523c51eb1c0dec9b752fb267b36a61e83cbc219c9d86da715" +checksum = "98a862d704c817d865c8740b62d8bbeb5adcb30965e93b471df8a5bcefa20a80" dependencies = [ "aws-credential-types", "aws-runtime", @@ -1681,9 +1681,9 @@ dependencies = [ [[package]] name = "aws-sigv4" -version = "1.3.5" +version = "1.3.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bffc03068fbb9c8dd5ce1c6fb240678a5cffb86fb2b7b1985c999c4b83c8df68" +checksum = "c35452ec3f001e1f2f6db107b6373f1f48f05ec63ba2c5c9fa91f07dad32af11" dependencies = [ "aws-credential-types", "aws-smithy-eventstream", @@ -1740,9 +1740,9 @@ dependencies = [ [[package]] name = "aws-smithy-eventstream" -version = "0.60.12" +version = "0.60.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9656b85088f8d9dc7ad40f9a6c7228e1e8447cdf4b046c87e152e0805dea02fa" +checksum = "e29a304f8319781a39808847efb39561351b1bb76e933da7aa90232673638658" dependencies = [ "aws-smithy-types", "bytes 1.10.1", @@ -1751,9 +1751,9 @@ dependencies = [ [[package]] name = "aws-smithy-http" -version = "0.62.4" +version = "0.62.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3feafd437c763db26aa04e0cc7591185d0961e64c61885bece0fb9d50ceac671" +checksum = "445d5d720c99eed0b4aa674ed00d835d9b1427dd73e04adaf2f94c6b2d6f9fca" dependencies = [ "aws-smithy-eventstream", "aws-smithy-runtime-api", @@ -1761,6 +1761,7 @@ dependencies = [ "bytes 1.10.1", "bytes-utils", "futures-core", + "futures-util", "http 0.2.12", "http 1.3.1", "http-body 0.4.6", @@ -1772,9 +1773,9 @@ dependencies = [ [[package]] name = "aws-smithy-http-client" -version = "1.1.3" +version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1053b5e587e6fa40ce5a79ea27957b04ba660baa02b28b7436f64850152234f1" +checksum = "623254723e8dfd535f566ee7b2381645f8981da086b5c4aa26c0c41582bb1d2c" dependencies = [ "aws-smithy-async", "aws-smithy-runtime-api", @@ -1802,9 +1803,9 @@ dependencies = [ [[package]] name = "aws-smithy-json" -version = "0.61.6" +version = "0.61.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cff418fc8ec5cadf8173b10125f05c2e7e1d46771406187b2c878557d4503390" +checksum = "2db31f727935fc63c6eeae8b37b438847639ec330a9161ece694efba257e0c54" dependencies = [ "aws-smithy-types", ] @@ -1830,9 +1831,9 @@ dependencies = [ [[package]] name = "aws-smithy-runtime" -version = "1.9.3" +version = "1.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "40ab99739082da5347660c556689256438defae3bcefd66c52b095905730e404" +checksum = "0bbe9d018d646b96c7be063dd07987849862b0e6d07c778aad7d93d1be6c1ef0" dependencies = [ "aws-smithy-async", "aws-smithy-http", @@ -1854,9 +1855,9 @@ dependencies = [ [[package]] name = "aws-smithy-runtime-api" -version = "1.9.1" +version = "1.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3683c5b152d2ad753607179ed71988e8cfd52964443b4f74fd8e552d0bbfeb46" +checksum = "ec7204f9fd94749a7c53b26da1b961b4ac36bf070ef1e0b94bb09f79d4f6c193" dependencies = [ "aws-smithy-async", "aws-smithy-types", @@ -1871,9 +1872,9 @@ dependencies = [ [[package]] name = "aws-smithy-types" -version = "1.3.3" +version = "1.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f5b3a7486f6690ba25952cabf1e7d75e34d69eaff5081904a47bc79074d6457" +checksum = "25f535879a207fce0db74b679cfc3e91a3159c8144d717d55f5832aea9eef46e" dependencies = [ "base64-simd", "bytes 1.10.1", @@ -1897,18 +1898,18 @@ dependencies = [ [[package]] name = "aws-smithy-xml" -version = "0.60.11" +version = "0.60.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e9c34127e8c624bc2999f3b657e749c1393bedc9cd97b92a804db8ced4d2e163" +checksum = "eab77cdd036b11056d2a30a7af7b775789fb024bf216acc13884c6c97752ae56" dependencies = [ "xmlparser", ] [[package]] name = "aws-types" -version = "1.3.9" +version = "1.3.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e2fd329bf0e901ff3f60425691410c69094dc2a1f34b331f37bfc4e9ac1565a1" +checksum = "d79fb68e3d7fe5d4833ea34dc87d2e97d26d3086cb3da660bb6b1f76d98680b6" dependencies = [ "aws-credential-types", "aws-smithy-async", diff --git a/Cargo.toml b/Cargo.toml index f46ffa2583c022be8704e95684ddce65b19d3fab..a8002e207d7ba9d3699832ac76be530e1979ead4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -455,15 +455,15 @@ async-task = "4.7" async-trait = "0.1" async-tungstenite = "0.31.0" async_zip = { version = "0.0.18", features = ["deflate", "deflate64"] } -aws-config = { version = "1.6.1", features = ["behavior-version-latest"] } -aws-credential-types = { version = "1.2.2", features = [ +aws-config = { version = "1.8.10", features = ["behavior-version-latest"] } +aws-credential-types = { version = "1.2.8", features = [ "hardcoded-credentials", ] } -aws-sdk-bedrockruntime = { version = "1.80.0", features = [ +aws-sdk-bedrockruntime = { version = "1.112.0", features = [ "behavior-version-latest", ] } -aws-smithy-runtime-api = { version = "1.7.4", features = ["http-1x", "client"] } -aws-smithy-types = { version = "1.3.0", features = ["http-body-1-x"] } +aws-smithy-runtime-api = { version = "1.9.2", features = ["http-1x", "client"] } +aws-smithy-types = { version = "1.3.4", features = ["http-body-1-x"] } backtrace = "0.3" base64 = "0.22" bincode = "1.2.1" diff --git a/crates/bedrock/src/bedrock.rs b/crates/bedrock/src/bedrock.rs index ec0b4070906fdfd31195668312b3e7b425cd28ee..744dde38076a5a12c9bc957a75e2435b1b753d96 100644 --- a/crates/bedrock/src/bedrock.rs +++ b/crates/bedrock/src/bedrock.rs @@ -87,7 +87,7 @@ pub async fn stream_completion( Ok(None) => None, Err(err) => Some(( Err(BedrockError::ClientError(anyhow!( - "{:?}", + "{}", aws_sdk_bedrockruntime::error::DisplayErrorContext(err) ))), stream, diff --git a/crates/language_models/src/provider/bedrock.rs b/crates/language_models/src/provider/bedrock.rs index b85a038bb235d97bd9de8614f19764ecabf7bbfe..9273234161a8169abf68190ca8fe4627b8f769dc 100644 --- a/crates/language_models/src/provider/bedrock.rs +++ b/crates/language_models/src/provider/bedrock.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use anyhow::{Context as _, Result, anyhow}; use aws_config::stalled_stream_protection::StalledStreamProtectionConfig; use aws_config::{BehaviorVersion, Region}; -use aws_credential_types::Credentials; +use aws_credential_types::{Credentials, Token}; use aws_http_client::AwsHttpClient; use bedrock::bedrock_client::Client as BedrockClient; use bedrock::bedrock_client::config::timeout::TimeoutConfig; @@ -30,18 +30,19 @@ use gpui::{ use gpui_tokio::Tokio; use http_client::HttpClient; use language_model::{ - AuthenticateError, LanguageModel, LanguageModelCacheConfiguration, + AuthenticateError, EnvVar, LanguageModel, LanguageModelCacheConfiguration, LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelId, LanguageModelName, LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest, LanguageModelToolChoice, LanguageModelToolResultContent, LanguageModelToolUse, MessageContent, RateLimiter, Role, - TokenUsage, + TokenUsage, env_var, }; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use serde_json::Value; use settings::{BedrockAvailableModel as AvailableModel, Settings, SettingsStore}; use smol::lock::OnceCell; +use std::sync::LazyLock; use strum::{EnumIter, IntoEnumIterator, IntoStaticStr}; use ui::{ButtonLink, ConfiguredApiCard, List, ListBulletItem, prelude::*}; use ui_input::InputField; @@ -54,12 +55,52 @@ actions!(bedrock, [Tab, TabPrev]); const PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("amazon-bedrock"); const PROVIDER_NAME: LanguageModelProviderName = LanguageModelProviderName::new("Amazon Bedrock"); +/// Credentials stored in the keychain for static authentication. +/// Region is handled separately since it's orthogonal to auth method. #[derive(Default, Clone, Deserialize, Serialize, PartialEq, Debug)] pub struct BedrockCredentials { pub access_key_id: String, pub secret_access_key: String, pub session_token: Option, - pub region: String, + pub bearer_token: Option, +} + +/// Resolved authentication configuration for Bedrock. +/// Settings take priority over UX-provided credentials. +#[derive(Clone, Debug, PartialEq)] +pub enum BedrockAuth { + /// Use default AWS credential provider chain (IMDSv2, PodIdentity, env vars, etc.) + Automatic, + /// Use AWS named profile from ~/.aws/credentials or ~/.aws/config + NamedProfile { profile_name: String }, + /// Use AWS SSO profile + SingleSignOn { profile_name: String }, + /// Use IAM credentials (access key + secret + optional session token) + IamCredentials { + access_key_id: String, + secret_access_key: String, + session_token: Option, + }, + /// Use Bedrock API Key (bearer token authentication) + ApiKey { api_key: String }, +} + +impl BedrockCredentials { + /// Convert stored credentials to the appropriate auth variant. + /// Prefers API key if present, otherwise uses IAM credentials. + fn into_auth(self) -> Option { + if let Some(api_key) = self.bearer_token.filter(|t| !t.is_empty()) { + Some(BedrockAuth::ApiKey { api_key }) + } else if !self.access_key_id.is_empty() && !self.secret_access_key.is_empty() { + Some(BedrockAuth::IamCredentials { + access_key_id: self.access_key_id, + secret_access_key: self.secret_access_key, + session_token: self.session_token.filter(|t| !t.is_empty()), + }) + } else { + None + } + } } #[derive(Default, Clone, Debug, PartialEq)] @@ -79,6 +120,8 @@ pub enum BedrockAuthMethod { NamedProfile, #[serde(rename = "sso")] SingleSignOn, + #[serde(rename = "api_key")] + ApiKey, /// IMDSv2, PodIdentity, env vars, etc. #[serde(rename = "default")] Automatic, @@ -90,6 +133,7 @@ impl From for BedrockAuthMethod { settings::BedrockAuthMethodContent::SingleSignOn => BedrockAuthMethod::SingleSignOn, settings::BedrockAuthMethodContent::Automatic => BedrockAuthMethod::Automatic, settings::BedrockAuthMethodContent::NamedProfile => BedrockAuthMethod::NamedProfile, + settings::BedrockAuthMethodContent::ApiKey => BedrockAuthMethod::ApiKey, } } } @@ -130,23 +174,26 @@ impl From for ModelMode { const AMAZON_AWS_URL: &str = "https://amazonaws.com"; // These environment variables all use a `ZED_` prefix because we don't want to overwrite the user's AWS credentials. -const ZED_BEDROCK_ACCESS_KEY_ID_VAR: &str = "ZED_ACCESS_KEY_ID"; -const ZED_BEDROCK_SECRET_ACCESS_KEY_VAR: &str = "ZED_SECRET_ACCESS_KEY"; -const ZED_BEDROCK_SESSION_TOKEN_VAR: &str = "ZED_SESSION_TOKEN"; -const ZED_AWS_PROFILE_VAR: &str = "ZED_AWS_PROFILE"; -const ZED_BEDROCK_REGION_VAR: &str = "ZED_AWS_REGION"; -const ZED_AWS_CREDENTIALS_VAR: &str = "ZED_AWS_CREDENTIALS"; -const ZED_AWS_ENDPOINT_VAR: &str = "ZED_AWS_ENDPOINT"; +static ZED_BEDROCK_ACCESS_KEY_ID_VAR: LazyLock = env_var!("ZED_ACCESS_KEY_ID"); +static ZED_BEDROCK_SECRET_ACCESS_KEY_VAR: LazyLock = env_var!("ZED_SECRET_ACCESS_KEY"); +static ZED_BEDROCK_SESSION_TOKEN_VAR: LazyLock = env_var!("ZED_SESSION_TOKEN"); +static ZED_AWS_PROFILE_VAR: LazyLock = env_var!("ZED_AWS_PROFILE"); +static ZED_BEDROCK_REGION_VAR: LazyLock = env_var!("ZED_AWS_REGION"); +static ZED_AWS_ENDPOINT_VAR: LazyLock = env_var!("ZED_AWS_ENDPOINT"); +static ZED_BEDROCK_BEARER_TOKEN_VAR: LazyLock = env_var!("ZED_BEDROCK_BEARER_TOKEN"); pub struct State { - credentials: Option, + /// The resolved authentication method. Settings take priority over UX credentials. + auth: Option, + /// Raw settings from settings.json settings: Option, + /// Whether credentials came from environment variables (only relevant for static credentials) credentials_from_env: bool, _subscription: Subscription, } impl State { - fn reset_credentials(&self, cx: &mut Context) -> Task> { + fn reset_auth(&self, cx: &mut Context) -> Task> { let credentials_provider = ::global(cx); cx.spawn(async move |this, cx| { credentials_provider @@ -154,19 +201,19 @@ impl State { .await .log_err(); this.update(cx, |this, cx| { - this.credentials = None; + this.auth = None; this.credentials_from_env = false; - this.settings = None; cx.notify(); }) }) } - fn set_credentials( + fn set_static_credentials( &mut self, credentials: BedrockCredentials, cx: &mut Context, ) -> Task> { + let auth = credentials.clone().into_auth(); let credentials_provider = ::global(cx); cx.spawn(async move |this, cx| { credentials_provider @@ -178,50 +225,131 @@ impl State { ) .await?; this.update(cx, |this, cx| { - this.credentials = Some(credentials); + this.auth = auth; + this.credentials_from_env = false; cx.notify(); }) }) } fn is_authenticated(&self) -> bool { - let derived = self - .settings - .as_ref() - .and_then(|s| s.authentication_method.as_ref()); - let creds = self.credentials.as_ref(); - - derived.is_some() || creds.is_some() + self.auth.is_some() } + /// Resolve authentication. Settings take priority over UX-provided credentials. fn authenticate(&self, cx: &mut Context) -> Task> { if self.is_authenticated() { return Task::ready(Ok(())); } + // Step 1: Check if settings specify an auth method (enterprise control) + if let Some(settings) = &self.settings { + if let Some(method) = &settings.authentication_method { + let profile_name = settings + .profile_name + .clone() + .unwrap_or_else(|| "default".to_string()); + + let auth = match method { + BedrockAuthMethod::Automatic => BedrockAuth::Automatic, + BedrockAuthMethod::NamedProfile => BedrockAuth::NamedProfile { profile_name }, + BedrockAuthMethod::SingleSignOn => BedrockAuth::SingleSignOn { profile_name }, + BedrockAuthMethod::ApiKey => { + // ApiKey method means "use static credentials from keychain/env" + // Fall through to load them below + return self.load_static_credentials(cx); + } + }; + + return cx.spawn(async move |this, cx| { + this.update(cx, |this, cx| { + this.auth = Some(auth); + this.credentials_from_env = false; + cx.notify(); + })?; + Ok(()) + }); + } + } + + // Step 2: No settings auth method - try to load static credentials + self.load_static_credentials(cx) + } + + /// Load static credentials from environment variables or keychain. + fn load_static_credentials( + &self, + cx: &mut Context, + ) -> Task> { let credentials_provider = ::global(cx); cx.spawn(async move |this, cx| { - let (credentials, from_env) = - if let Ok(credentials) = std::env::var(ZED_AWS_CREDENTIALS_VAR) { - (credentials, true) - } else { - let (_, credentials) = credentials_provider - .read_credentials(AMAZON_AWS_URL, cx) - .await? - .ok_or_else(|| AuthenticateError::CredentialsNotFound)?; + // Try environment variables first + let (auth, from_env) = if let Some(bearer_token) = &ZED_BEDROCK_BEARER_TOKEN_VAR.value { + if !bearer_token.is_empty() { ( - String::from_utf8(credentials) - .context("invalid {PROVIDER_NAME} credentials")?, - false, + Some(BedrockAuth::ApiKey { + api_key: bearer_token.to_string(), + }), + true, ) - }; + } else { + (None, false) + } + } else if let Some(access_key_id) = &ZED_BEDROCK_ACCESS_KEY_ID_VAR.value { + if let Some(secret_access_key) = &ZED_BEDROCK_SECRET_ACCESS_KEY_VAR.value { + if !access_key_id.is_empty() && !secret_access_key.is_empty() { + let session_token = ZED_BEDROCK_SESSION_TOKEN_VAR + .value + .as_deref() + .filter(|s| !s.is_empty()) + .map(|s| s.to_string()); + ( + Some(BedrockAuth::IamCredentials { + access_key_id: access_key_id.to_string(), + secret_access_key: secret_access_key.to_string(), + session_token, + }), + true, + ) + } else { + (None, false) + } + } else { + (None, false) + } + } else { + (None, false) + }; + + // If we got auth from env vars, use it + if let Some(auth) = auth { + this.update(cx, |this, cx| { + this.auth = Some(auth); + this.credentials_from_env = from_env; + cx.notify(); + })?; + return Ok(()); + } + + // Try keychain + let (_, credentials_bytes) = credentials_provider + .read_credentials(AMAZON_AWS_URL, cx) + .await? + .ok_or(AuthenticateError::CredentialsNotFound)?; + + let credentials_str = String::from_utf8(credentials_bytes) + .context("invalid {PROVIDER_NAME} credentials")?; let credentials: BedrockCredentials = - serde_json::from_str(&credentials).context("failed to parse credentials")?; + serde_json::from_str(&credentials_str).context("failed to parse credentials")?; + + let auth = credentials + .into_auth() + .ok_or(AuthenticateError::CredentialsNotFound)?; this.update(cx, |this, cx| { - this.credentials = Some(credentials); - this.credentials_from_env = from_env; + this.auth = Some(auth); + this.credentials_from_env = false; cx.notify(); })?; @@ -229,15 +357,19 @@ impl State { }) } + /// Get the resolved region. Checks env var, then settings, then defaults to us-east-1. fn get_region(&self) -> String { - // Get region - from credentials or directly from settings - let credentials_region = self.credentials.as_ref().map(|s| s.region.clone()); - let settings_region = self.settings.as_ref().and_then(|s| s.region.clone()); - - // Use credentials region if available, otherwise use settings region, finally fall back to default - credentials_region - .or(settings_region) - .unwrap_or(String::from("us-east-1")) + // Priority: env var > settings > default + if let Some(region) = ZED_BEDROCK_REGION_VAR.value.as_deref() { + if !region.is_empty() { + return region.to_string(); + } + } + + self.settings + .as_ref() + .and_then(|s| s.region.clone()) + .unwrap_or_else(|| "us-east-1".to_string()) } fn get_allow_global(&self) -> bool { @@ -257,7 +389,7 @@ pub struct BedrockLanguageModelProvider { impl BedrockLanguageModelProvider { pub fn new(http_client: Arc, cx: &mut App) -> Self { let state = cx.new(|cx| State { - credentials: None, + auth: None, settings: Some(AllLanguageModelSettings::get_global(cx).bedrock.clone()), credentials_from_env: false, _subscription: cx.observe_global::(|_, cx| { @@ -266,7 +398,7 @@ impl BedrockLanguageModelProvider { }); Self { - http_client: AwsHttpClient::new(http_client.clone()), + http_client: AwsHttpClient::new(http_client), handle: Tokio::handle(cx), state, } @@ -312,7 +444,6 @@ impl LanguageModelProvider for BedrockLanguageModelProvider { for model in bedrock::Model::iter() { if !matches!(model, bedrock::Model::Custom { .. }) { - // TODO: Sonnet 3.7 vs. 3.7 Thinking bug is here. models.insert(model.id().to_string(), model); } } @@ -366,8 +497,7 @@ impl LanguageModelProvider for BedrockLanguageModelProvider { } fn reset_credentials(&self, cx: &mut App) -> Task> { - self.state - .update(cx, |state, cx| state.reset_credentials(cx)) + self.state.update(cx, |state, cx| state.reset_auth(cx)) } } @@ -393,25 +523,11 @@ impl BedrockModel { fn get_or_init_client(&self, cx: &AsyncApp) -> anyhow::Result<&BedrockClient> { self.client .get_or_try_init_blocking(|| { - let (auth_method, credentials, endpoint, region, settings) = - cx.read_entity(&self.state, |state, _cx| { - let auth_method = state - .settings - .as_ref() - .and_then(|s| s.authentication_method.clone()); - - let endpoint = state.settings.as_ref().and_then(|s| s.endpoint.clone()); - - let region = state.get_region(); - - ( - auth_method, - state.credentials.clone(), - endpoint, - region, - state.settings.clone(), - ) - })?; + let (auth, endpoint, region) = cx.read_entity(&self.state, |state, _cx| { + let endpoint = state.settings.as_ref().and_then(|s| s.endpoint.clone()); + let region = state.get_region(); + (state.auth.clone(), endpoint, region) + })?; let mut config_builder = aws_config::defaults(BehaviorVersion::latest()) .stalled_stream_protection(StalledStreamProtectionConfig::disabled()) @@ -425,37 +541,39 @@ impl BedrockModel { config_builder = config_builder.endpoint_url(endpoint_url); } - match auth_method { - None => { - if let Some(creds) = credentials { - let aws_creds = Credentials::new( - creds.access_key_id, - creds.secret_access_key, - creds.session_token, - None, - "zed-bedrock-provider", - ); - config_builder = config_builder.credentials_provider(aws_creds); - } + match auth { + Some(BedrockAuth::Automatic) | None => { + // Use default AWS credential provider chain } - Some(BedrockAuthMethod::NamedProfile) - | Some(BedrockAuthMethod::SingleSignOn) => { - // Currently NamedProfile and SSO behave the same way but only the instructions change - // Until we support BearerAuth through SSO, this will not change. - let profile_name = settings - .and_then(|s| s.profile_name) - .unwrap_or_else(|| "default".to_string()); - + Some(BedrockAuth::NamedProfile { profile_name }) + | Some(BedrockAuth::SingleSignOn { profile_name }) => { if !profile_name.is_empty() { config_builder = config_builder.profile_name(profile_name); } } - Some(BedrockAuthMethod::Automatic) => { - // Use default credential provider chain + Some(BedrockAuth::IamCredentials { + access_key_id, + secret_access_key, + session_token, + }) => { + let aws_creds = Credentials::new( + access_key_id, + secret_access_key, + session_token, + None, + "zed-bedrock-provider", + ); + config_builder = config_builder.credentials_provider(aws_creds); + } + Some(BedrockAuth::ApiKey { api_key }) => { + config_builder = config_builder + .auth_scheme_preference(["httpBearerAuth".into()]) // https://github.com/smithy-lang/smithy-rs/pull/4241 + .token_provider(Token::new(api_key, None)); } } let config = self.handle.block_on(config_builder.load()); + anyhow::Ok(BedrockClient::new(&config)) }) .context("initializing Bedrock client")?; @@ -1024,7 +1142,7 @@ struct ConfigurationView { access_key_id_editor: Entity, secret_access_key_editor: Entity, session_token_editor: Entity, - region_editor: Entity, + bearer_token_editor: Entity, state: Entity, load_credentials_task: Option>, focus_handle: FocusHandle, @@ -1035,7 +1153,7 @@ impl ConfigurationView { const PLACEHOLDER_SECRET_ACCESS_KEY_TEXT: &'static str = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"; const PLACEHOLDER_SESSION_TOKEN_TEXT: &'static str = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"; - const PLACEHOLDER_REGION: &'static str = "us-east-1"; + const PLACEHOLDER_BEARER_TOKEN_TEXT: &'static str = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"; fn new(state: Entity, window: &mut Window, cx: &mut Context) -> Self { let focus_handle = cx.focus_handle(); @@ -1066,9 +1184,9 @@ impl ConfigurationView { .tab_stop(true) }); - let region_editor = cx.new(|cx| { - InputField::new(window, cx, Self::PLACEHOLDER_REGION) - .label("Region") + let bearer_token_editor = cx.new(|cx| { + InputField::new(window, cx, Self::PLACEHOLDER_BEARER_TOKEN_TEXT) + .label("Bedrock API Key") .tab_index(3) .tab_stop(true) }); @@ -1095,7 +1213,7 @@ impl ConfigurationView { access_key_id_editor, secret_access_key_editor, session_token_editor, - region_editor, + bearer_token_editor, state, load_credentials_task, focus_handle, @@ -1131,25 +1249,30 @@ impl ConfigurationView { } else { Some(session_token) }; - let region = self.region_editor.read(cx).text(cx).trim().to_string(); - let region = if region.is_empty() { - "us-east-1".to_string() + let bearer_token = self + .bearer_token_editor + .read(cx) + .text(cx) + .trim() + .to_string(); + let bearer_token = if bearer_token.is_empty() { + None } else { - region + Some(bearer_token) }; let state = self.state.clone(); cx.spawn(async move |_, cx| { state .update(cx, |state, cx| { - let credentials: BedrockCredentials = BedrockCredentials { - region: region.clone(), - access_key_id: access_key_id.clone(), - secret_access_key: secret_access_key.clone(), - session_token: session_token.clone(), + let credentials = BedrockCredentials { + access_key_id, + secret_access_key, + session_token, + bearer_token, }; - state.set_credentials(credentials, cx) + state.set_static_credentials(credentials, cx) })? .await }) @@ -1163,16 +1286,12 @@ impl ConfigurationView { .update(cx, |editor, cx| editor.set_text("", window, cx)); self.session_token_editor .update(cx, |editor, cx| editor.set_text("", window, cx)); - self.region_editor + self.bearer_token_editor .update(cx, |editor, cx| editor.set_text("", window, cx)); let state = self.state.clone(); - cx.spawn(async move |_, cx| { - state - .update(cx, |state, cx| state.reset_credentials(cx))? - .await - }) - .detach_and_log_err(cx); + cx.spawn(async move |_, cx| state.update(cx, |state, cx| state.reset_auth(cx))?.await) + .detach_and_log_err(cx); } fn should_render_editor(&self, cx: &Context) -> bool { @@ -1195,9 +1314,11 @@ impl ConfigurationView { impl Render for ConfigurationView { fn render(&mut self, _window: &mut Window, cx: &mut Context) -> impl IntoElement { - let env_var_set = self.state.read(cx).credentials_from_env; - let bedrock_settings = self.state.read(cx).settings.as_ref(); - let bedrock_method = bedrock_settings + let state = self.state.read(cx); + let env_var_set = state.credentials_from_env; + let auth = state.auth.clone(); + let settings_auth_method = state + .settings .as_ref() .and_then(|s| s.authentication_method.clone()); @@ -1205,34 +1326,62 @@ impl Render for ConfigurationView { return div().child(Label::new("Loading credentials...")).into_any(); } - let configured_label = if env_var_set { - format!( - "Access Key ID is set in {ZED_BEDROCK_ACCESS_KEY_ID_VAR}, Secret Key is set in {ZED_BEDROCK_SECRET_ACCESS_KEY_VAR}, Region is set in {ZED_BEDROCK_REGION_VAR} environment variables." - ) - } else { - match bedrock_method { - Some(BedrockAuthMethod::Automatic) => "You are using automatic credentials.".into(), - Some(BedrockAuthMethod::NamedProfile) => "You are using named profile.".into(), - Some(BedrockAuthMethod::SingleSignOn) => { - "You are using a single sign on profile.".into() - } - None => "You are using static credentials.".into(), + let configured_label = match &auth { + Some(BedrockAuth::Automatic) => { + "Using automatic credentials (AWS default chain)".into() } + Some(BedrockAuth::NamedProfile { profile_name }) => { + format!("Using AWS profile: {profile_name}") + } + Some(BedrockAuth::SingleSignOn { profile_name }) => { + format!("Using AWS SSO profile: {profile_name}") + } + Some(BedrockAuth::IamCredentials { .. }) if env_var_set => { + format!( + "Using IAM credentials from {} and {} environment variables", + ZED_BEDROCK_ACCESS_KEY_ID_VAR.name, ZED_BEDROCK_SECRET_ACCESS_KEY_VAR.name + ) + } + Some(BedrockAuth::IamCredentials { .. }) => "Using IAM credentials".into(), + Some(BedrockAuth::ApiKey { .. }) if env_var_set => { + format!( + "Using Bedrock API Key from {} environment variable", + ZED_BEDROCK_BEARER_TOKEN_VAR.name + ) + } + Some(BedrockAuth::ApiKey { .. }) => "Using Bedrock API Key".into(), + None => "Not authenticated".into(), }; + // Determine if credentials can be reset + // Settings-derived auth (non-ApiKey) cannot be reset from UI + let is_settings_derived = matches!( + settings_auth_method, + Some(BedrockAuthMethod::Automatic) + | Some(BedrockAuthMethod::NamedProfile) + | Some(BedrockAuthMethod::SingleSignOn) + ); + let tooltip_label = if env_var_set { Some(format!( - "To reset your credentials, unset the {ZED_BEDROCK_ACCESS_KEY_ID_VAR}, {ZED_BEDROCK_SECRET_ACCESS_KEY_VAR}, and {ZED_BEDROCK_REGION_VAR} environment variables." + "To reset your credentials, unset the {}, {}, and {} or {} environment variables.", + ZED_BEDROCK_ACCESS_KEY_ID_VAR.name, + ZED_BEDROCK_SECRET_ACCESS_KEY_VAR.name, + ZED_BEDROCK_SESSION_TOKEN_VAR.name, + ZED_BEDROCK_BEARER_TOKEN_VAR.name )) - } else if bedrock_method.is_some() { - Some("You cannot reset credentials as they're being derived, check Zed settings to understand how.".to_string()) + } else if is_settings_derived { + Some( + "Authentication method is configured in settings. Edit settings.json to change." + .to_string(), + ) } else { None }; if self.should_render_editor(cx) { return ConfiguredApiCard::new(configured_label) - .disabled(env_var_set || bedrock_method.is_some()) + .disabled(env_var_set || is_settings_derived) .on_click(cx.listener(|this, _, window, cx| this.reset_credentials(window, cx))) .when_some(tooltip_label, |this, label| this.tooltip_label(label)) .into_any_element(); @@ -1262,7 +1411,7 @@ impl Render for ConfigurationView { .child(self.render_static_credentials_ui()) .child( Label::new( - format!("You can also assign the {ZED_BEDROCK_ACCESS_KEY_ID_VAR}, {ZED_BEDROCK_SECRET_ACCESS_KEY_VAR} AND {ZED_BEDROCK_REGION_VAR} environment variables and restart Zed."), + format!("You can also assign the {}, {} AND {} environment variables (or {} for Bedrock API Key authentication) and restart Zed.", ZED_BEDROCK_ACCESS_KEY_ID_VAR.name, ZED_BEDROCK_SECRET_ACCESS_KEY_VAR.name, ZED_BEDROCK_REGION_VAR.name, ZED_BEDROCK_BEARER_TOKEN_VAR.name), ) .size(LabelSize::Small) .color(Color::Muted) @@ -1270,7 +1419,7 @@ impl Render for ConfigurationView { ) .child( Label::new( - format!("Optionally, if your environment uses AWS CLI profiles, you can set {ZED_AWS_PROFILE_VAR}; if it requires a custom endpoint, you can set {ZED_AWS_ENDPOINT_VAR}; and if it requires a Session Token, you can set {ZED_BEDROCK_SESSION_TOKEN_VAR}."), + format!("Optionally, if your environment uses AWS CLI profiles, you can set {}; if it requires a custom endpoint, you can set {}; and if it requires a Session Token, you can set {}.", ZED_AWS_PROFILE_VAR.name, ZED_AWS_ENDPOINT_VAR.name, ZED_BEDROCK_SESSION_TOKEN_VAR.name), ) .size(LabelSize::Small) .color(Color::Muted), @@ -1292,31 +1441,47 @@ impl ConfigurationView { ) .child( Label::new( - "This method uses your AWS access key ID and secret access key directly.", + "This method uses your AWS access key ID and secret access key, or a Bedrock API Key.", ) ) .child( List::new() .child( ListBulletItem::new("") - .child(Label::new("Create an IAM user in the AWS console with programmatic access")) + .child(Label::new("For access keys: Create an IAM user in the AWS console with programmatic access")) .child(ButtonLink::new("IAM Console", "https://us-east-1.console.aws.amazon.com/iam/home?region=us-east-1#/users")) ) .child( ListBulletItem::new("") - .child(Label::new("Attach the necessary Bedrock permissions to this")) - .child(ButtonLink::new("user", "https://docs.aws.amazon.com/bedrock/latest/userguide/inference-prereq.html")) + .child(Label::new("For Bedrock API Keys: Generate an API key from the")) + .child(ButtonLink::new("Bedrock Console", "https://docs.aws.amazon.com/bedrock/latest/userguide/api-keys-use.html")) ) .child( - ListBulletItem::new("Copy the access key ID and secret access key when provided") + ListBulletItem::new("") + .child(Label::new("Attach the necessary Bedrock permissions to this")) + .child(ButtonLink::new("user", "https://docs.aws.amazon.com/bedrock/latest/userguide/inference-prereq.html")) ) .child( - ListBulletItem::new("Enter these credentials below") - ) + ListBulletItem::new("Enter either access keys OR a Bedrock API Key below (not both)") + ), ) .child(self.access_key_id_editor.clone()) .child(self.secret_access_key_editor.clone()) .child(self.session_token_editor.clone()) - .child(self.region_editor.clone()) + .child( + Label::new("OR") + .size(LabelSize::Default) + .weight(FontWeight::BOLD) + .my_1(), + ) + .child(self.bearer_token_editor.clone()) + .child( + Label::new( + format!("Region is configured via {} environment variable or settings.json (defaults to us-east-1).", ZED_BEDROCK_REGION_VAR.name), + ) + .size(LabelSize::Small) + .color(Color::Muted) + .mt_2(), + ) } } diff --git a/crates/settings/src/settings_content/language_model.rs b/crates/settings/src/settings_content/language_model.rs index b106f3d9925cb4afe058cff44649f998c8b73d8a..e523286e5f56af88110c2d4a7d874c22195ea2b1 100644 --- a/crates/settings/src/settings_content/language_model.rs +++ b/crates/settings/src/settings_content/language_model.rs @@ -83,6 +83,8 @@ pub enum BedrockAuthMethodContent { NamedProfile, #[serde(rename = "sso")] SingleSignOn, + #[serde(rename = "api_key")] + ApiKey, /// IMDSv2, PodIdentity, env vars, etc. #[serde(rename = "default")] Automatic,