mixa/zed
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

bedrock: Add Bedrock API key authentication support (#41393)

Browse Source
This commit is contained in:
Shardul Vaidya
2025-12-17 06:54:57 -05:00
committed by GitHub
parent c0b3422941
commit edf21a38c1
5 changed files with 353 additions and 185 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
Cargo.lock generated
View File

@@ -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",

10
Cargo.toml
View File

@@ -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"

2
crates/bedrock/src/bedrock.rs
View File

@@ -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,

459
crates/language_models/src/provider/bedrock.rs
View File

@@ -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<String>,
pub region: String,
pub bearer_token: Option<String>,
}
/// 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<String>,
},
/// 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<BedrockAuth> {
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<settings::BedrockAuthMethodContent> 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<BedrockModelMode> 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<EnvVar> = env_var!("ZED_ACCESS_KEY_ID");
static ZED_BEDROCK_SECRET_ACCESS_KEY_VAR: LazyLock<EnvVar> = env_var!("ZED_SECRET_ACCESS_KEY");
static ZED_BEDROCK_SESSION_TOKEN_VAR: LazyLock<EnvVar> = env_var!("ZED_SESSION_TOKEN");
static ZED_AWS_PROFILE_VAR: LazyLock<EnvVar> = env_var!("ZED_AWS_PROFILE");
static ZED_BEDROCK_REGION_VAR: LazyLock<EnvVar> = env_var!("ZED_AWS_REGION");
static ZED_AWS_ENDPOINT_VAR: LazyLock<EnvVar> = env_var!("ZED_AWS_ENDPOINT");
static ZED_BEDROCK_BEARER_TOKEN_VAR: LazyLock<EnvVar> = env_var!("ZED_BEDROCK_BEARER_TOKEN");
pub struct State {
credentials: Option<BedrockCredentials>,
/// The resolved authentication method. Settings take priority over UX credentials.
auth: Option<BedrockAuth>,
/// Raw settings from settings.json
settings: Option<AmazonBedrockSettings>,
/// 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<Self>) -> Task<Result<()>> {
fn reset_auth(&self, cx: &mut Context<Self>) -> Task<Result<()>> {
let credentials_provider = <dyn CredentialsProvider>::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<Self>,
) -> Task<Result<()>> {
let auth = credentials.clone().into_auth();
let credentials_provider = <dyn CredentialsProvider>::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<Self>) -> Task<Result<(), AuthenticateError>> {
if self.is_authenticated() {
return Task::ready(Ok(()));
}
let credentials_provider = <dyn CredentialsProvider>::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)?;
(
String::from_utf8(credentials)
.context("invalid {PROVIDER_NAME} credentials")?,
false,
)
// 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<Self>,
) -> Task<Result<(), AuthenticateError>> {
let credentials_provider = <dyn CredentialsProvider>::global(cx);
cx.spawn(async move |this, cx| {
// Try environment variables first
let (auth, from_env) = if let Some(bearer_token) = &ZED_BEDROCK_BEARER_TOKEN_VAR.value {
if !bearer_token.is_empty() {
(
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());
// Priority: env var > settings > default
if let Some(region) = ZED_BEDROCK_REGION_VAR.value.as_deref() {
if !region.is_empty() {
return region.to_string();
}
}
// 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"))
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<dyn HttpClient>, 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::<SettingsStore>(|_, 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<Result<()>> {
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<InputField>,
secret_access_key_editor: Entity<InputField>,
session_token_editor: Entity<InputField>,
region_editor: Entity<InputField>,
bearer_token_editor: Entity<InputField>,
state: Entity<State>,
load_credentials_task: Option<Task<()>>,
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<State>, window: &mut Window, cx: &mut Context<Self>) -> 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<Self>) -> bool {
@@ -1195,9 +1314,11 @@ impl ConfigurationView {
impl Render for ConfigurationView {
fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> 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("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("")
.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("Copy the access key ID and secret access key when provided")
)
.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(),
)
}
}

2
crates/settings/src/settings_content/language_model.rs
View File

@@ -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,