Skip to content

Commit c715051

Browse files
5herlockedona-agent
5herlocked
and
ona-agent
committed
Add AWS Bedrock API key (bearer token) authentication support
Fixes #41312 This adds support for AWS Bedrock API keys (bearer tokens) as an authentication method. Users can now authenticate using: - Traditional access keys (existing) - Bearer tokens via ZED_AWS_BEARER_TOKEN_BEDROCK environment variable (new) - Bearer tokens via UI input field (new) Changes: - Added bearer_token field to BedrockCredentials struct - Added ZED_AWS_BEARER_TOKEN_BEDROCK environment variable constant - Updated authentication logic to check for bearer tokens (read-only from ENV) - Implemented BearerTokenProvider using AWS SDK's ResolveIdentity trait - Modified client initialization to use token_provider when bearer token is present - Enhanced UI with bearer token input field and updated instructions - Updated all help text and tooltips to mention bearer token option Bearer tokens are: - Read from environment variables (never written) - Stored in Zed's credential store when provided via UI - Passed to AWS SDK using the proper token_provider configuration Also includes fix for duplicate Region field (issue #41313). References: - AWS Bedrock API Key Documentation: https://docs.aws.amazon.com/bedrock/latest/userguide/api-keys-use.html Co-authored-by: Ona <[email protected]>
1 parent db0f7a8 commit c715051

File tree

1 file changed

+92
-22
lines changed

1 file changed

+92
-22
lines changed

‎crates/language_models/src/provider/bedrock.rs‎

Lines changed: 92 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ use aws_config::stalled_stream_protection::StalledStreamProtectionConfig;
88
use aws_config::{BehaviorVersion, Region};
99
use aws_credential_types::Credentials;
1010
use aws_http_client::AwsHttpClient;
11+
use aws_smithy_runtime_api::client::identity::{Identity, IdentityFuture, ResolveIdentity};
12+
use aws_smithy_runtime_api::client::runtime_components::RuntimeComponents;
1113
use bedrock::bedrock_client::Client as BedrockClient;
1214
use bedrock::bedrock_client::config::timeout::TimeoutConfig;
1315
use bedrock::bedrock_client::types::{
@@ -50,12 +52,33 @@ use crate::AllLanguageModelSettings;
5052
const PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("amazon-bedrock");
5153
const PROVIDER_NAME: LanguageModelProviderName = LanguageModelProviderName::new("Amazon Bedrock");
5254

55+
#[derive(Clone)]
56+
struct BearerTokenProvider {
57+
token: String,
58+
}
59+
60+
impl BearerTokenProvider {
61+
fn new(token: String) -> Self {
62+
Self { token }
63+
}
64+
}
65+
66+
impl ResolveIdentity for BearerTokenProvider {
67+
fn resolve_identity<'a>(
68+
&'a self,
69+
_runtime_components: &'a RuntimeComponents,
70+
) -> IdentityFuture<'a> {
71+
IdentityFuture::ready(Ok(Identity::new(self.token.clone(), None)))
72+
}
73+
}
74+
5375
#[derive(Default, Clone, Deserialize, Serialize, PartialEq, Debug)]
5476
pub struct BedrockCredentials {
5577
pub access_key_id: String,
5678
pub secret_access_key: String,
5779
pub session_token: Option<String>,
5880
pub region: String,
81+
pub bearer_token: Option<String>,
5982
}
6083

6184
#[derive(Default, Clone, Debug, PartialEq)]
@@ -132,6 +155,7 @@ const ZED_AWS_PROFILE_VAR: &str = "ZED_AWS_PROFILE";
132155
const ZED_BEDROCK_REGION_VAR: &str = "ZED_AWS_REGION";
133156
const ZED_AWS_CREDENTIALS_VAR: &str = "ZED_AWS_CREDENTIALS";
134157
const ZED_AWS_ENDPOINT_VAR: &str = "ZED_AWS_ENDPOINT";
158+
const ZED_AWS_BEARER_TOKEN_BEDROCK: &str = "ZED_AWS_BEARER_TOKEN_BEDROCK";
135159

136160
pub struct State {
137161
credentials: Option<BedrockCredentials>,
@@ -199,6 +223,16 @@ impl State {
199223
let (credentials, from_env) =
200224
if let Ok(credentials) = std::env::var(ZED_AWS_CREDENTIALS_VAR) {
201225
(credentials, true)
226+
} else if let Ok(bearer_token) = std::env::var(ZED_AWS_BEARER_TOKEN_BEDROCK) {
227+
let region = std::env::var(ZED_BEDROCK_REGION_VAR).unwrap_or_else(|_| "us-east-1".to_string());
228+
let creds = BedrockCredentials {
229+
access_key_id: String::new(),
230+
secret_access_key: String::new(),
231+
session_token: None,
232+
region,
233+
bearer_token: Some(bearer_token),
234+
};
235+
(serde_json::to_string(&creds).unwrap(), true)
202236
} else {
203237
let (_, credentials) = credentials_provider
204238
.read_credentials(AMAZON_AWS_URL, cx)
@@ -254,7 +288,7 @@ impl BedrockLanguageModelProvider {
254288
});
255289

256290
Self {
257-
http_client: AwsHttpClient::new(http_client.clone()),
291+
http_client: AwsHttpClient::new(http_client),
258292
handle: Tokio::handle(cx),
259293
state,
260294
}
@@ -407,6 +441,15 @@ impl BedrockModel {
407441
.region(Region::new(region))
408442
.timeout_config(TimeoutConfig::disabled());
409443

444+
let bearer_token = credentials.as_ref().and_then(|c| c.bearer_token.as_ref());
445+
if let Some(token) = bearer_token {
446+
if !token.is_empty() {
447+
config_builder = config_builder.identity_cache(
448+
aws_config::identity::IdentityCache::no_cache()
449+
).token_provider(BearerTokenProvider::new(token.clone()));
450+
}
451+
}
452+
410453
if let Some(endpoint_url) = endpoint
411454
&& !endpoint_url.is_empty()
412455
{
@@ -416,20 +459,20 @@ impl BedrockModel {
416459
match auth_method {
417460
None => {
418461
if let Some(creds) = credentials {
419-
let aws_creds = Credentials::new(
420-
creds.access_key_id,
421-
creds.secret_access_key,
422-
creds.session_token,
423-
None,
424-
"zed-bedrock-provider",
425-
);
426-
config_builder = config_builder.credentials_provider(aws_creds);
462+
if creds.bearer_token.is_none() || creds.bearer_token.as_ref().map_or(true, |t| t.is_empty()) {
463+
let aws_creds = Credentials::new(
464+
creds.access_key_id,
465+
creds.secret_access_key,
466+
creds.session_token,
467+
None,
468+
"zed-bedrock-provider",
469+
);
470+
config_builder = config_builder.credentials_provider(aws_creds);
471+
}
427472
}
428473
}
429474
Some(BedrockAuthMethod::NamedProfile)
430475
| Some(BedrockAuthMethod::SingleSignOn) => {
431-
// Currently NamedProfile and SSO behave the same way but only the instructions change
432-
// Until we support BearerAuth through SSO, this will not change.
433476
let profile_name = settings
434477
.and_then(|s| s.profile_name)
435478
.unwrap_or_else(|| "default".to_string());
@@ -439,7 +482,6 @@ impl BedrockModel {
439482
}
440483
}
441484
Some(BedrockAuthMethod::Automatic) => {
442-
// Use default credential provider chain
443485
}
444486
}
445487

@@ -1009,6 +1051,7 @@ struct ConfigurationView {
10091051
access_key_id_editor: Entity<InputField>,
10101052
secret_access_key_editor: Entity<InputField>,
10111053
session_token_editor: Entity<InputField>,
1054+
bearer_token_editor: Entity<InputField>,
10121055
region_editor: Entity<InputField>,
10131056
state: Entity<State>,
10141057
load_credentials_task: Option<Task<()>>,
@@ -1019,6 +1062,7 @@ impl ConfigurationView {
10191062
const PLACEHOLDER_SECRET_ACCESS_KEY_TEXT: &'static str =
10201063
"XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX";
10211064
const PLACEHOLDER_SESSION_TOKEN_TEXT: &'static str = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX";
1065+
const PLACEHOLDER_BEARER_TOKEN_TEXT: &'static str = "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX";
10221066
const PLACEHOLDER_REGION: &'static str = "us-east-1";
10231067

10241068
fn new(state: Entity<State>, window: &mut Window, cx: &mut Context<Self>) -> Self {
@@ -1058,6 +1102,10 @@ impl ConfigurationView {
10581102
InputField::new(window, cx, Self::PLACEHOLDER_SESSION_TOKEN_TEXT)
10591103
.label("Session Token (Optional)")
10601104
}),
1105+
bearer_token_editor: cx.new(|cx| {
1106+
InputField::new(window, cx, Self::PLACEHOLDER_BEARER_TOKEN_TEXT)
1107+
.label("Bearer Token (Optional)")
1108+
}),
10611109
region_editor: cx
10621110
.new(|cx| InputField::new(window, cx, Self::PLACEHOLDER_REGION).label("Region")),
10631111
state,
@@ -1094,6 +1142,17 @@ impl ConfigurationView {
10941142
} else {
10951143
Some(session_token)
10961144
};
1145+
let bearer_token = self
1146+
.bearer_token_editor
1147+
.read(cx)
1148+
.text(cx)
1149+
.trim()
1150+
.to_string();
1151+
let bearer_token = if bearer_token.is_empty() {
1152+
None
1153+
} else {
1154+
Some(bearer_token)
1155+
};
10971156
let region = self.region_editor.read(cx).text(cx).trim().to_string();
10981157
let region = if region.is_empty() {
10991158
"us-east-1".to_string()
@@ -1110,6 +1169,7 @@ impl ConfigurationView {
11101169
access_key_id: access_key_id.clone(),
11111170
secret_access_key: secret_access_key.clone(),
11121171
session_token: session_token.clone(),
1172+
bearer_token: bearer_token.clone(),
11131173
};
11141174

11151175
state.set_credentials(credentials, cx)
@@ -1126,6 +1186,8 @@ impl ConfigurationView {
11261186
.update(cx, |editor, cx| editor.set_text("", window, cx));
11271187
self.session_token_editor
11281188
.update(cx, |editor, cx| editor.set_text("", window, cx));
1189+
self.bearer_token_editor
1190+
.update(cx, |editor, cx| editor.set_text("", window, cx));
11291191
self.region_editor
11301192
.update(cx, |editor, cx| editor.set_text("", window, cx));
11311193

@@ -1169,7 +1231,7 @@ impl Render for ConfigurationView {
11691231
.gap_1()
11701232
.child(Icon::new(IconName::Check).color(Color::Success))
11711233
.child(Label::new(if env_var_set {
1172-
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.")
1234+
format!("Credentials are set via environment variables ({ZED_BEDROCK_ACCESS_KEY_ID_VAR}/{ZED_AWS_BEARER_TOKEN_BEDROCK} and {ZED_BEDROCK_REGION_VAR}).")
11731235
} else {
11741236
match bedrock_method {
11751237
Some(BedrockAuthMethod::Automatic) => "You are using automatic credentials".into(),
@@ -1188,7 +1250,7 @@ impl Render for ConfigurationView {
11881250
.icon_position(IconPosition::Start)
11891251
.disabled(env_var_set || bedrock_method.is_some())
11901252
.when(env_var_set, |this| {
1191-
this.tooltip(Tooltip::text(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.")))
1253+
this.tooltip(Tooltip::text(format!("To reset your credentials, unset the {ZED_BEDROCK_ACCESS_KEY_ID_VAR}/{ZED_AWS_BEARER_TOKEN_BEDROCK} and {ZED_BEDROCK_REGION_VAR} environment variables.")))
11921254
})
11931255
.when(bedrock_method.is_some(), |this| {
11941256
this.tooltip(Tooltip::text("You cannot reset credentials as they're being derived, check Zed settings to understand how"))
@@ -1221,10 +1283,9 @@ impl Render for ConfigurationView {
12211283
)
12221284
)
12231285
.child(self.render_static_credentials_ui())
1224-
.child(self.region_editor.clone())
12251286
.child(
12261287
Label::new(
1227-
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."),
1288+
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 (or {ZED_AWS_BEARER_TOKEN_BEDROCK} for bearer token authentication) and restart Zed."),
12281289
)
12291290
.size(LabelSize::Small)
12301291
.color(Color::Muted)
@@ -1253,31 +1314,40 @@ impl ConfigurationView {
12531314
)
12541315
.child(
12551316
Label::new(
1256-
"This method uses your AWS access key ID and secret access key directly.",
1317+
"This method uses your AWS access key ID and secret access key, or a bearer token (API key).",
12571318
)
12581319
)
12591320
.child(
12601321
List::new()
12611322
.child(InstructionListItem::new(
1262-
"Create an IAM user in the AWS console with programmatic access",
1323+
"For access keys: Create an IAM user in the AWS console with programmatic access",
12631324
Some("IAM Console"),
12641325
Some("https://us-east-1.console.aws.amazon.com/iam/home?region=us-east-1#/users"),
12651326
))
1327+
.child(InstructionListItem::new(
1328+
"For bearer tokens: Generate an API key from the ",
1329+
Some("Bedrock Console"),
1330+
Some("https://docs.aws.amazon.com/bedrock/latest/userguide/api-keys-use.html"),
1331+
))
12661332
.child(InstructionListItem::new(
12671333
"Attach the necessary Bedrock permissions to this ",
12681334
Some("user"),
12691335
Some("https://docs.aws.amazon.com/bedrock/latest/userguide/inference-prereq.html"),
12701336
))
12711337
.child(InstructionListItem::text_only(
1272-
"Copy the access key ID and secret access key when provided",
1273-
))
1274-
.child(InstructionListItem::text_only(
1275-
"Enter these credentials below",
1338+
"Enter either access keys OR a bearer token below (not both)",
12761339
)),
12771340
)
12781341
.child(self.access_key_id_editor.clone())
12791342
.child(self.secret_access_key_editor.clone())
12801343
.child(self.session_token_editor.clone())
1344+
.child(
1345+
Label::new("OR")
1346+
.size(LabelSize::Default)
1347+
.weight(FontWeight::BOLD)
1348+
.my_1(),
1349+
)
1350+
.child(self.bearer_token_editor.clone())
12811351
.child(self.region_editor.clone())
12821352
.into_any_element()
12831353
}

0 commit comments

Comments
 (0)