public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix bedrock authentication refresh

Browse Source
This commit is contained in:
Josh Woodcock
2026-10-14 12:43:32 -05:00
parent 55eb87a983
commit 04cff91509
19 changed files with 777 additions and 237 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/aws_credentials.rs
+40 -16
View File
@@ -14,6 +14,10 @@ use tokio::sync::Mutex;
use vec1::vec1;
use crate::settings::{AISettings, AISettingsChangedEvent};
use crate::ai::provider::authentication::{
BedrockAuthConnection, connection_revision, credentials_are_current,
};
use crate::ai::provider::client::BedrockClientConfig;
use crate::terminal::event::{AfterBlockCompletedEvent, BlockType, UserBlockCompleted};
use crate::terminal::model_events::{ModelEvent, ModelEventDispatcher};
use crate::workspaces::user_workspaces::{UserWorkspaces, UserWorkspacesEvent};
@@ -56,7 +60,7 @@ fn aws_profile_reference_for_message(profile: &str, capitalize_first_word: bool)
}
}
fn user_facing_aws_credentials_error_message(err: &CredentialsError, profile: &str) -> String {
pub(crate) fn user_facing_aws_credentials_error_message(err: &CredentialsError, profile: &str) -> String {
match err {
CredentialsError::CredentialsNotLoaded(_) => format!(
"AWS credentials were not found for {}. Log in with the AWS CLI or update your AWS credentials configuration, then refresh.",
@@ -129,15 +133,17 @@ fn aws_credentials_state_for_error(err: LoadAwsCredentialsError) -> AwsCredentia
/// Loads AWS credentials from the AWS SDK.
///
/// # Arguments
/// * `profile` - AWS profile name. If empty, uses the default AWS SDK behavior
/// * `connection` - Effective AWS connection. If its profile is empty, uses the default AWS SDK behavior
/// (checks AWS_PROFILE env var, then uses "default").
pub async fn load_aws_credentials_from_sdk(
profile: &str,
connection: &BedrockClientConfig,
) -> Result<AwsCredentials, LoadAwsCredentialsError> {
let region_provider = aws_config::meta::region::RegionProviderChain::default_provider();
let loader =
aws_config::defaults(aws_config::BehaviorVersion::latest()).region(region_provider);
let loader = if profile.trim().is_empty() || profile.trim() == "default" {
let profile = &connection.profile;
let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
if !connection.region.is_empty() {
loader = loader.region(aws_config::Region::new(connection.region.clone()));
}
let loader = if profile.trim().is_empty() {
loader // Let AWS SDK use its default behavior
} else {
loader.profile_name(profile)
@@ -150,14 +156,15 @@ pub async fn load_aws_credentials_from_sdk(
let creds = provider.provide_credentials().await.map_err(|e| {
let message = user_facing_aws_credentials_error_message(&e, profile);
log::warn!("{e}");
// TODO(isaiah): turn this full SDK dump back down to debug once we've resolved
// the current customer-facing AWS credential issue and no longer need prod-visible
// provider internals for support debugging.
log::warn!("{e:#?}");
log::warn!("AWS credential reload failed: {message}");
log::info!("AWS credential load failure message shown to user: {message}");
LoadAwsCredentialsError::CredentialsLoadFailed(message)
})?;
if !credentials_are_current(creds.expiry(), SystemTime::now()) {
return Err(LoadAwsCredentialsError::CredentialsLoadFailed(
"AWS credentials have expired. Refresh your AWS login.".to_owned(),
));
}
Ok(AwsCredentials::new(
creds.access_key_id().to_string(),
@@ -198,8 +205,10 @@ impl AwsCredentialRefresher for ApiKeyManager {
..
}) = event
{
let auth_command = &AISettings::as_ref(ctx).bedrock_auth_refresh_command;
if command.trim().starts_with(auth_command.trim()) {
let auth_command = BedrockAuthConnection::current(ctx).command;
if !auth_command.trim().is_empty()
&& command.trim().starts_with(auth_command.trim())
{
log::debug!("Detected AWS auth command completion, refreshing credentials");
drop(refresh_aws_credentials(manager, ctx));
}
@@ -225,6 +234,8 @@ impl AwsCredentialRefresher for ApiKeyManager {
if matches!(
event,
AISettingsChangedEvent::BedrockProfile { .. }
| AISettingsChangedEvent::BedrockAuthMethod { .. }
| AISettingsChangedEvent::BedrockRegion { .. }
| AISettingsChangedEvent::BedrockAuthRefreshCommand { .. }
| AISettingsChangedEvent::BedrockEnabled { .. }
) {
@@ -266,15 +277,28 @@ fn refresh_aws_credentials_local_chain(
return Box::pin(async { Ok(()) });
}
let profile = (*AISettings::as_ref(ctx).bedrock_profile).clone();
let connection = BedrockAuthConnection::current(ctx);
if connection.config.auth_method == crate::settings::BedrockAuthMethod::StaticKeys {
manager.set_aws_credentials_state(AwsCredentialsState::Missing, ctx);
return Box::pin(async { Ok(()) });
}
let config = connection.config.clone();
let revision = connection_revision(ctx);
manager.set_aws_credentials_state(AwsCredentialsState::Refreshing, ctx);
let (tx, rx) = channel();
// credential fetch from aws cli's disk cache
let _ = ctx.spawn(
async move { load_aws_credentials_from_sdk(&profile).await },
async move { load_aws_credentials_from_sdk(&config).await },
move |manager, result, ctx| {
if connection_revision(ctx) != revision
|| BedrockAuthConnection::current(ctx) != connection
|| manager.aws_credentials_refresh_strategy() != AwsCredentialsRefreshStrategy::LocalChain
{
let _ = tx.send(Err("AWS connection changed during credential reload".to_owned()));
return;
}
let (new_state, tx_result) = match result {
Ok(credentials) => (
AwsCredentialsState::Loaded {
app/src/ai/blocklist/block.rs
+2 -2
View File
@@ -4439,7 +4439,7 @@ impl AIBlock {
}
let ai_settings = AISettings::as_ref(ctx);
let login_command = ai_settings.bedrock_auth_refresh_command.value().clone();
let login_command = crate::ai::provider::authentication::BedrockAuthConnection::current(ctx).command;
// Recreating a historical error view must not restart authentication.
let auto_login_enabled = attempt_auto_login && *ai_settings.bedrock_auto_login.value();
@@ -4452,7 +4452,7 @@ impl AIBlock {
let model_name = model_name.clone();
let view = ctx.add_typed_action_view(|ctx| {
AwsBedrockCredentialsErrorView::new(model_name, login_command, auto_login_enabled, ctx)
AwsBedrockCredentialsErrorView::new(model_name, login_command, ctx)
});
// Subscribe to events from the view and emit AIBlockEvents directly
app/src/ai/blocklist/controller.rs
+22 -7
View File
@@ -6610,7 +6610,11 @@ impl BlocklistAIController {
let prepared = match result {
Ok(prepared) => prepared,
Err(error) => {
self.fail_provider_startup(conversation_id, stream_id, error.to_string(), ctx);
let renderable = crate::ai::provider::authentication::provider_startup_error(
&error,
slot.response_stream.as_ref(ctx).model_id(),
);
self.fail_provider_startup_with_error(conversation_id, stream_id, renderable, ctx);
return;
}
};
@@ -8916,6 +8920,21 @@ impl BlocklistAIController {
stream_id: ResponseStreamId,
message: String,
ctx: &mut ModelContext<Self>,
) {
self.fail_provider_startup_with_error(
conversation_id,
stream_id,
RenderableAIError::other(message, false),
ctx,
);
}
fn fail_provider_startup_with_error(
&mut self,
conversation_id: AIConversationId,
stream_id: ResponseStreamId,
error: RenderableAIError,
ctx: &mut ModelContext<Self>,
) {
let response_stream = self
.active_provider_runs
@@ -8925,14 +8944,10 @@ impl BlocklistAIController {
let Some(response_stream) = response_stream else {
return;
};
let message = error.to_string();
BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| {
history_model.mark_response_stream_completed_with_error(
RenderableAIError::Other {
error_message: message.clone(),
will_attempt_resume: false,
waiting_for_network: false,
is_user_error: false,
},
error,
false,
&stream_id,
conversation_id,
app/src/ai/blocklist/controller/response_stream.rs
+5 -23
View File
@@ -187,30 +187,12 @@ impl ResponseStream {
// Fall back to Bedrock
if *settings.bedrock_enabled.value() {
let auth_method = *settings.bedrock_auth_method.value();
let region = settings.bedrock_region.value().clone();
let cross_region_inference = *settings.bedrock_cross_region_inference.value();
let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx);
let mut config = BedrockClientConfig {
auth_method,
profile: settings.bedrock_profile.value().clone(),
region,
access_key_id: settings.bedrock_access_key_id.value().clone(),
secret_access_key: settings.bedrock_secret_access_key.value().clone(),
session_token: None,
cross_region_inference,
};
if let ::ai::api_keys::AwsCredentialsState::Loaded { credentials, .. } =
api_key_manager.aws_credentials_state()
{
config.auth_method = crate::settings::BedrockAuthMethod::StaticKeys;
config.access_key_id = credentials.access_key().to_string();
config.secret_access_key = credentials.secret_key().to_string();
config.session_token = credentials.session_token().map(|s| s.to_string());
}
return ProviderConfig::Bedrock(config.with_external_fallbacks());
let config = BedrockClientConfig::from_settings(settings).with_managed_credentials(
&api_key_manager.aws_credentials_refresh_strategy(),
api_key_manager.aws_credentials_state(),
);
return ProviderConfig::Bedrock(config);
}
ProviderConfig::None
app/src/ai/blocklist/inline_action/aws_bedrock_credentials_error.rs
+13 -26
View File
@@ -35,8 +35,6 @@ pub struct AwsBedrockCredentialsErrorView {
model_name: String,
login_command: String,
/// Whether auto-login was triggered, which shows a simpler "running command" message.
auto_login_triggered: bool,
// Run button
run_button: ViewHandle<ActionButton>,
@@ -52,7 +50,6 @@ impl AwsBedrockCredentialsErrorView {
pub fn new(
model_name: String,
login_command: String,
auto_login_triggered: bool,
ctx: &mut ViewContext<Self>,
) -> Self {
// Run button
@@ -83,7 +80,6 @@ impl AwsBedrockCredentialsErrorView {
Self {
model_name,
login_command,
auto_login_triggered,
run_button,
configure_button,
auto_login_checkbox_handle: MouseStateHandle::default(),
@@ -104,23 +100,6 @@ impl View for AwsBedrockCredentialsErrorView {
let appearance = Appearance::as_ref(app);
let theme = appearance.theme();
// If auto-login was triggered, show a simple "running command" message
if self.auto_login_triggered {
return Flex::row()
.with_spacing(8.)
.with_cross_axis_alignment(CrossAxisAlignment::Center)
.with_child(
Text::new(
format!("Running `{}`...", self.login_command),
appearance.ui_font_family(),
14.,
)
.with_color(blended_colors::text_sub(theme, theme.surface_1()))
.with_selectable(false)
.finish(),
)
.finish();
}
let auto_login_enabled = *AISettings::as_ref(app).bedrock_auto_login.value();
@@ -149,11 +128,19 @@ impl View for AwsBedrockCredentialsErrorView {
let make_detail_text = || {
Text::new(
format!(
"Failed to authenticate with AWS Bedrock when using {}. \
Run `{}` to refresh credentials.",
self.model_name, self.login_command
),
if self.login_command.is_empty() {
format!(
"Failed to authenticate with AWS Bedrock when using {}. \
Refresh credentials for the AWS profile selected in Providers, then retry.",
self.model_name
)
} else {
format!(
"Failed to authenticate with AWS Bedrock when using {}. \
Run `{}` to refresh credentials.",
self.model_name, self.login_command
)
},
appearance.ui_font_family(),
14.,
)
app/src/ai/crosscheck/reviewer.rs
+6 -21
View File
@@ -181,28 +181,13 @@ impl CrosscheckReviewer {
// Fall back to Bedrock via external config
if *settings.bedrock_enabled.value() {
let auth_method = *settings.bedrock_auth_method.value();
let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx);
let mut config = crate::ai::provider::client::BedrockClientConfig {
auth_method,
profile: settings.bedrock_profile.value().clone(),
region: settings.bedrock_region.value().clone(),
access_key_id: settings.bedrock_access_key_id.value().clone(),
secret_access_key: settings.bedrock_secret_access_key.value().clone(),
session_token: None,
cross_region_inference: *settings.bedrock_cross_region_inference.value(),
};
if let ::ai::api_keys::AwsCredentialsState::Loaded { credentials, .. } =
api_key_manager.aws_credentials_state()
{
config.auth_method = crate::settings::BedrockAuthMethod::StaticKeys;
config.access_key_id = credentials.access_key().to_string();
config.secret_access_key = credentials.secret_key().to_string();
config.session_token = credentials.session_token().map(|s| s.to_string());
}
return ProviderConfig::Bedrock(config.with_external_fallbacks());
let config = crate::ai::provider::client::BedrockClientConfig::from_settings(settings)
.with_managed_credentials(
&api_key_manager.aws_credentials_refresh_strategy(),
api_key_manager.aws_credentials_state(),
);
return ProviderConfig::Bedrock(config);
}
ProviderConfig::None
app/src/ai/llms.rs
+1 -9
View File
@@ -3471,13 +3471,5 @@ mod tests;
fn bedrock_discovery_config(
settings: &AISettings,
) -> crate::ai::provider::client::BedrockClientConfig {
crate::ai::provider::client::BedrockClientConfig {
auth_method: *settings.bedrock_auth_method.value(),
profile: settings.bedrock_profile.value().clone(),
region: settings.bedrock_region.value().clone(),
access_key_id: settings.bedrock_access_key_id.value().clone(),
secret_access_key: settings.bedrock_secret_access_key.value().clone(),
session_token: None,
cross_region_inference: *settings.bedrock_cross_region_inference.value(),
}
crate::ai::provider::client::BedrockClientConfig::from_settings(settings)
}
app/src/ai/mod.rs
+2
View File
@@ -85,6 +85,8 @@ pub mod outline;
pub(crate) use ai::paths;
pub fn init(app: &mut AppContext) {
#[cfg(not(target_family = "wasm"))]
app.add_singleton_model(provider::authentication::BedrockAuthentication::new);
#[cfg(not(target_family = "wasm"))]
app.add_singleton_model(acp::AcpRuntimeModel::new);
#[cfg(not(target_family = "wasm"))]
app/src/ai/provider/authentication.rs
+316
View File
@@ -0,0 +1,316 @@
use std::collections::HashMap;
use std::time::SystemTime;
use ai::api_keys::{ApiKeyManager, AwsCredentials, AwsCredentialsRefreshStrategy, AwsCredentialsState};
use futures::channel::oneshot;
use futures::future::{BoxFuture, Shared};
use futures::FutureExt;
use galaxyui::{AppContext, Entity, ModelContext, SingletonEntity};
use settings::Setting;
use super::client::{BedrockClientConfig, BedrockError};
use crate::ai::agent::RenderableAIError;
use crate::settings::{AISettings, AISettingsChangedEvent, BedrockAuthMethod};
impl BedrockClientConfig {
pub(crate) fn from_settings(settings: &AISettings) -> Self {
Self {
auth_method: *settings.bedrock_auth_method.value(),
profile: settings.bedrock_profile.value().clone(),
region: settings.bedrock_region.value().clone(),
access_key_id: settings.bedrock_access_key_id.value().clone(),
secret_access_key: settings.bedrock_secret_access_key.value().clone(),
session_token: None,
credentials_expiry: None,
cross_region_inference: *settings.bedrock_cross_region_inference.value(),
}
}
pub(crate) fn with_managed_credentials(
mut self,
strategy: &AwsCredentialsRefreshStrategy,
state: &AwsCredentialsState,
) -> Self {
match strategy {
AwsCredentialsRefreshStrategy::LocalChain => self,
AwsCredentialsRefreshStrategy::OidcManaged { region, .. } => {
self.auth_method = BedrockAuthMethod::StaticKeys;
self.region.clone_from(region);
self.access_key_id.clear();
self.secret_access_key.clear();
self.session_token = None;
self.credentials_expiry = None;
if let AwsCredentialsState::Loaded { credentials, .. } = state {
self.access_key_id = credentials.access_key().to_owned();
self.secret_access_key = credentials.secret_key().to_owned();
self.session_token = credentials.session_token().map(str::to_owned);
self.credentials_expiry = Some(
credentials.expires_at().unwrap_or(SystemTime::UNIX_EPOCH),
);
}
self
}
}
}
}
pub(crate) fn credentials_are_current(expiry: Option<SystemTime>, now: SystemTime) -> bool {
expiry.is_none_or(|expiry| expiry > now)
}
pub(crate) fn provider_startup_error(error: &anyhow::Error, model: &str) -> RenderableAIError {
match error.downcast_ref::<BedrockError>() {
Some(BedrockError::CredentialsNotConfigured | BedrockError::CredentialsLoadFailed(_)) => {
RenderableAIError::AwsBedrockCredentialsExpiredOrInvalid {
model_name: model.to_owned(),
}
}
Some(
BedrockError::RegionNotConfigured
| BedrockError::ApiError(_)
| BedrockError::ModelNotFound(_)
| BedrockError::AccessDenied(_)
| BedrockError::Throttling(_)
| BedrockError::ValidationError(_),
) | None => RenderableAIError::other(error.to_string(), false),
}
}
pub(crate) fn login_command(
configured: &str,
profile: &str,
) -> String {
if configured != "aws sso login" {
return configured.to_owned();
}
if profile.is_empty() {
configured.to_owned()
} else {
format!("aws sso login --profile {}", shell_words::quote(profile))
}
}
pub(crate) fn login_command_parts(command: &str) -> Result<Vec<String>, String> {
let parts = shell_words::split(command)
.map_err(|_| "AWS login command contains invalid quoting".to_owned())?;
if parts.is_empty() {
return Err("AWS login command is not configured".to_owned());
}
Ok(parts)
}
#[derive(Clone, PartialEq)]
pub(crate) struct BedrockAuthConnection {
pub(crate) config: BedrockClientConfig,
pub(crate) enabled: bool,
pub(crate) auto_login: bool,
pub(crate) command: String,
}
impl BedrockAuthConnection {
pub(crate) fn current(ctx: &AppContext) -> Self {
let settings = AISettings::as_ref(ctx);
let config = BedrockClientConfig::from_settings(settings);
Self {
command: match config.auth_method {
BedrockAuthMethod::Sso => login_command(
settings.bedrock_auth_refresh_command.value(),
&config.profile,
),
BedrockAuthMethod::Profile | BedrockAuthMethod::StaticKeys => String::new(),
},
config,
enabled: *settings.bedrock_enabled.value(),
auto_login: *settings.bedrock_auto_login.value(),
}
}
fn identity(&self) -> String {
format!(
"{:?}:{}:{}:{}",
self.config.auth_method, self.config.profile, self.config.region, self.command
)
}
}
#[derive(Default)]
struct RecoveryAttempts {
exchanges: HashMap<String, (String, String)>,
}
impl RecoveryAttempts {
fn begin(&mut self, conversation: &str, exchange: &str, connection: &str, manual: bool) -> bool {
let attempt = (exchange.to_owned(), connection.to_owned());
if !manual && self.exchanges.get(conversation) == Some(&attempt) {
return false;
}
self.exchanges.insert(conversation.to_owned(), attempt);
true
}
}
type CredentialResult = Result<AwsCredentials, String>;
type CredentialFuture = Shared<BoxFuture<'static, CredentialResult>>;
#[derive(Default)]
pub(crate) struct BedrockAuthentication {
attempts: RecoveryAttempts,
active: Vec<(u64, BedrockAuthConnection, CredentialFuture)>,
generation: u64,
revision: u64,
}
impl BedrockAuthentication {
pub(crate) fn new(ctx: &mut ModelContext<Self>) -> Self {
ctx.subscribe_to_model(&AISettings::handle(ctx), |me, _, event, _| {
if matches!(
event,
AISettingsChangedEvent::BedrockEnabled { .. }
| AISettingsChangedEvent::BedrockAuthMethod { .. }
| AISettingsChangedEvent::BedrockProfile { .. }
| AISettingsChangedEvent::BedrockRegion { .. }
| AISettingsChangedEvent::BedrockAuthRefreshCommand { .. }
| AISettingsChangedEvent::BedrockAccessKeyId { .. }
| AISettingsChangedEvent::BedrockSecretAccessKey { .. }
) {
me.revision += 1;
me.active.clear();
}
});
Self::default()
}
pub(crate) fn revision(&self) -> u64 {
self.revision
}
pub(crate) fn recover(
&mut self,
conversation: &str,
exchange: &str,
manual: bool,
ctx: &mut ModelContext<Self>,
) -> BoxFuture<'static, Result<(), String>> {
let connection = BedrockAuthConnection::current(ctx);
if !connection.enabled || (!manual && !connection.auto_login) {
return Box::pin(async { Err("Automatic AWS login is disabled".to_owned()) });
}
if !self.attempts.begin(conversation, exchange, &connection.identity(), manual) {
return Box::pin(async {
Err("AWS authentication still failed after recovery. Use Refresh AWS Credentials to retry.".to_owned())
});
}
let strategy = ApiKeyManager::as_ref(ctx).aws_credentials_refresh_strategy();
if matches!(strategy, AwsCredentialsRefreshStrategy::OidcManaged { .. }) {
return ApiKeyManager::handle(ctx).update(ctx, |manager, ctx| {
crate::ai::aws_credentials::refresh_aws_credentials(manager, ctx)
});
}
if connection.config.auth_method == BedrockAuthMethod::StaticKeys {
return Box::pin(async {
Err("Update the configured AWS static keys to authenticate.".to_owned())
});
}
let (generation, future) = match self.active.iter().find(|(_, active, _)| active == &connection) {
Some((generation, _, future)) => (*generation, future.clone()),
None => {
self.generation += 1;
let config = connection.config.clone();
let command = connection.command.clone();
let future = recover_credentials(
connection.config.auth_method,
manual,
move || {
let config = config.clone();
async move {
crate::ai::aws_credentials::load_aws_credentials_from_sdk(&config)
.await
.map_err(|error| error.to_string())
}.boxed()
},
move || {
let command = command.clone();
async move {
let parts = login_command_parts(&command)?;
let status = tokio::process::Command::new(&parts[0])
.args(&parts[1..])
.status()
.await
.map_err(|error| format!("Failed to start AWS login: {error}"))?;
if !status.success() {
return Err(format!("AWS login command exited with status: {status}"));
}
Ok(())
}.boxed()
},
).boxed().shared();
self.active.push((self.generation, connection.clone(), future.clone()));
(self.generation, future)
}
};
let revision = self.revision;
let (tx, rx) = oneshot::channel();
let _ = ctx.spawn(future, move |me, result, ctx| {
me.active.retain(|(active, _, _)| *active != generation);
if me.revision != revision
|| BedrockAuthConnection::current(ctx) != connection
|| ApiKeyManager::as_ref(ctx).aws_credentials_refresh_strategy() != strategy
{
let _ = tx.send(Err("AWS connection changed during authentication".to_owned()));
return;
}
let outcome = result.as_ref().map(|_| ()).map_err(Clone::clone);
ApiKeyManager::handle(ctx).update(ctx, |manager, ctx| {
let state = match result {
Ok(credentials) => AwsCredentialsState::Loaded {
credentials,
loaded_at: SystemTime::now(),
},
Err(message) => AwsCredentialsState::Failed { message },
};
manager.set_aws_credentials_state(state, ctx);
});
let _ = tx.send(outcome);
});
async move {
rx.await.unwrap_or_else(|_| Err("AWS authentication was interrupted".to_owned()))
}.boxed()
}
}
async fn recover_credentials(
auth_method: BedrockAuthMethod,
login_first: bool,
load: impl Fn() -> BoxFuture<'static, CredentialResult>,
login: impl Fn() -> BoxFuture<'static, Result<(), String>>,
) -> CredentialResult {
match auth_method {
BedrockAuthMethod::Profile => return load().await,
BedrockAuthMethod::StaticKeys => {
return Err("Update the configured AWS static keys to authenticate.".to_owned());
}
BedrockAuthMethod::Sso => {}
}
if !login_first {
if let Ok(credentials) = load().await {
return Ok(credentials);
}
}
login().await?;
load().await
}
impl Entity for BedrockAuthentication {
type Event = ();
}
impl SingletonEntity for BedrockAuthentication {}
pub(crate) fn connection_revision(ctx: &AppContext) -> u64 {
ctx.try_get_singleton_model_as_ref::<BedrockAuthentication>()
.map_or(0, |authentication| authentication.revision())
}
#[cfg(test)]
#[path = "authentication_tests.rs"]
mod tests;
app/src/ai/provider/authentication_tests.rs
+277
View File
@@ -0,0 +1,277 @@
use std::time::{Duration, SystemTime};
use ai::api_keys::{AwsCredentials, AwsCredentialsRefreshStrategy, AwsCredentialsState};
use super::*;
fn config(auth_method: BedrockAuthMethod) -> BedrockClientConfig {
BedrockClientConfig {
auth_method,
profile: "default".to_owned(),
region: String::new(),
access_key_id: "configured-key".to_owned(),
secret_access_key: "configured-secret".to_owned(),
session_token: None,
credentials_expiry: None,
cross_region_inference: true,
}
}
#[test]
fn provider_profile_and_region_are_authoritative_including_defaults() {
galaxyui::App::test((), |mut app| async move {
crate::test_util::settings::initialize_settings_for_tests(&mut app);
for (profile, region) in [("default", ""), ("selected-profile", "eu-west-1")] {
AISettings::handle(&app).update(&mut app, |settings, ctx| {
settings.bedrock_profile.set_value(profile.to_owned(), ctx).unwrap();
settings.bedrock_region.set_value(region.to_owned(), ctx).unwrap();
});
app.read(|ctx| {
let resolved = BedrockClientConfig::from_settings(AISettings::as_ref(ctx));
assert_eq!(resolved.profile, profile);
assert_eq!(resolved.region, region);
});
}
});
}
#[test]
fn custom_login_commands_are_not_rewritten() {
for command in ["", "company-login --interactive", "aws sso login --profile custom"] {
assert_eq!(login_command(command, "selected"), command);
}
}
#[test]
fn default_login_targets_and_quotes_the_effective_profile() {
assert_eq!(
login_command("aws sso login", "default"),
"aws sso login --profile default"
);
let command = login_command(
"aws sso login",
"team profile",
);
assert_eq!(
login_command_parts(&command).unwrap(),
["aws", "sso", "login", "--profile", "team profile"]
);
assert!(login_command_parts("").is_err());
assert!(login_command_parts("aws 'unterminated").is_err());
}
#[test]
fn local_credential_snapshots_never_override_configured_authentication() {
let now = SystemTime::now();
for auth in [
BedrockAuthMethod::Profile,
BedrockAuthMethod::Sso,
BedrockAuthMethod::StaticKeys,
] {
for expiry in [None, Some(now - Duration::from_secs(1)), Some(now + Duration::from_secs(60))] {
let configured = config(auth);
let loaded = AwsCredentialsState::Loaded {
credentials: AwsCredentials::new(
"snapshot-key".to_owned(),
"snapshot-secret".to_owned(),
Some("snapshot-token".to_owned()),
expiry,
),
loaded_at: now,
};
assert_eq!(
configured.clone().with_managed_credentials(
&AwsCredentialsRefreshStrategy::LocalChain,
&loaded,
),
configured
);
}
}
}
#[test]
fn managed_oidc_credentials_preserve_expiry_and_do_not_fall_back_to_local_keys() {
let now = SystemTime::now();
let strategy = AwsCredentialsRefreshStrategy::OidcManaged {
task_id: Some("task".to_owned()),
role_arn: "test-role".to_owned(),
region: "us-east-1".to_owned(),
};
let loaded = AwsCredentialsState::Loaded {
credentials: AwsCredentials::new(
"managed-key".to_owned(),
"managed-secret".to_owned(),
Some("managed-token".to_owned()),
Some(now + Duration::from_secs(60)),
),
loaded_at: now,
};
let resolved = config(BedrockAuthMethod::Profile).with_managed_credentials(&strategy, &loaded);
assert_eq!(resolved.auth_method, BedrockAuthMethod::StaticKeys);
assert_eq!(resolved.region, "us-east-1");
assert_eq!(resolved.access_key_id, "managed-key");
assert_eq!(resolved.credentials_expiry, Some(now + Duration::from_secs(60)));
for state in [
AwsCredentialsState::Missing,
AwsCredentialsState::Disabled,
AwsCredentialsState::Refreshing,
AwsCredentialsState::Failed { message: "failed".to_owned() },
] {
let resolved = config(BedrockAuthMethod::Profile).with_managed_credentials(&strategy, &state);
assert!(resolved.access_key_id.is_empty());
assert!(resolved.secret_access_key.is_empty());
}
}
#[test]
fn automatic_recovery_is_bounded_per_user_exchange_but_manual_retry_is_available() {
let mut attempts = RecoveryAttempts::default();
assert!(attempts.begin("conversation", "query-1", "connection", false));
assert!(!attempts.begin("conversation", "query-1", "connection", false));
assert!(attempts.begin("conversation", "query-1", "connection", true));
assert!(attempts.begin("conversation", "query-2", "connection", false));
assert!(attempts.begin("conversation", "query-2", "new-connection", false));
assert!(attempts.begin("another-conversation", "query-2", "connection", false));
}
#[test]
fn credential_expiry_is_checked_at_the_boundary() {
let now = SystemTime::now();
assert!(credentials_are_current(None, now));
assert!(credentials_are_current(Some(now + Duration::from_secs(1)), now));
assert!(!credentials_are_current(Some(now), now));
assert!(!credentials_are_current(Some(now - Duration::from_secs(1)), now));
}
fn test_credentials() -> AwsCredentials {
AwsCredentials::new("test-key".to_owned(), "test-secret".to_owned(), None, None)
}
#[tokio::test]
async fn external_refresh_recovers_without_running_login() {
let credentials = recover_credentials(
BedrockAuthMethod::Sso,
false,
|| futures::future::ready(Ok(test_credentials())).boxed(),
|| futures::future::ready(Err("Login must not run".to_owned())).boxed(),
).await.unwrap();
assert_eq!(credentials, test_credentials());
}
#[tokio::test]
async fn expired_session_runs_login_then_awaits_a_fresh_reload() {
use std::sync::{Arc, atomic::{AtomicUsize, Ordering}};
let calls = Arc::new(AtomicUsize::new(0));
let load_calls = calls.clone();
let login_calls = calls.clone();
let result = recover_credentials(
BedrockAuthMethod::Sso,
false,
move || {
let index = load_calls.fetch_add(1, Ordering::SeqCst);
let result = match index {
0 => Err("expired".to_owned()),
2 => Ok(test_credentials()),
index => panic!("Unexpected reload order: {index}"),
};
futures::future::ready(result).boxed()
},
move || {
assert_eq!(login_calls.fetch_add(1, Ordering::SeqCst), 1);
futures::future::ready(Ok(())).boxed()
},
).await;
assert_eq!(result, Ok(test_credentials()));
assert_eq!(calls.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn manual_login_failure_does_not_reload_or_resume() {
let result = recover_credentials(
BedrockAuthMethod::Sso,
true,
|| panic!("Credentials must not reload after failed login"),
|| futures::future::ready(Err("login failed".to_owned())).boxed(),
).await;
assert_eq!(result, Err("login failed".to_owned()));
}
#[tokio::test]
async fn login_success_with_failed_reload_remains_an_error() {
let result = recover_credentials(
BedrockAuthMethod::Sso,
true,
|| futures::future::ready(Err("reload failed".to_owned())).boxed(),
|| futures::future::ready(Ok(())).boxed(),
).await;
assert_eq!(result, Err("reload failed".to_owned()));
}
#[tokio::test]
async fn shared_recovery_executes_login_once_for_overlapping_waiters() {
use std::sync::{Arc, atomic::{AtomicUsize, Ordering}};
let calls = Arc::new(AtomicUsize::new(0));
let login_calls = calls.clone();
let recovery = recover_credentials(
BedrockAuthMethod::Sso,
true,
|| futures::future::ready(Ok(test_credentials())).boxed(),
move || {
login_calls.fetch_add(1, Ordering::SeqCst);
futures::future::ready(Ok(())).boxed()
},
).boxed().shared();
let (first, second) = futures::join!(recovery.clone(), recovery);
assert_eq!(first, Ok(test_credentials()));
assert_eq!(second, Ok(test_credentials()));
assert_eq!(calls.load(Ordering::SeqCst), 1);
}
#[test]
fn credential_preparation_failures_surface_authentication_recovery() {
use super::super::client::BedrockError;
for error in [
BedrockError::CredentialsNotConfigured,
BedrockError::CredentialsLoadFailed("SSO session invalid".to_owned()),
] {
let error = provider_startup_error(&anyhow::Error::new(error), "selected-model");
assert!(error.is_aws_bedrock_credentials_error());
}
for error in [
BedrockError::RegionNotConfigured,
BedrockError::ApiError("transport".to_owned()),
BedrockError::ModelNotFound("model".to_owned()),
BedrockError::AccessDenied("policy".to_owned()),
BedrockError::Throttling("busy".to_owned()),
BedrockError::ValidationError("request".to_owned()),
] {
let error = provider_startup_error(&anyhow::Error::new(error), "selected-model");
assert!(!error.is_aws_bedrock_credentials_error());
}
assert!(!provider_startup_error(&anyhow::anyhow!("unrelated"), "selected-model")
.is_aws_bedrock_credentials_error());
}
#[tokio::test]
async fn profile_authentication_never_starts_a_login_flow() {
for manual in [false, true] {
let result = recover_credentials(
BedrockAuthMethod::Profile,
manual,
|| futures::future::ready(Ok(test_credentials())).boxed(),
|| panic!("Profile authentication must not execute login commands"),
).await;
assert_eq!(result, Ok(test_credentials()));
let result = recover_credentials(
BedrockAuthMethod::Profile,
manual,
|| futures::future::ready(Err("profile needs fresh credentials".to_owned())).boxed(),
|| panic!("A failed profile reload must not execute login commands"),
).await;
assert_eq!(result, Err("profile needs fresh credentials".to_owned()));
}
}
app/src/ai/provider/client.rs
+19 -29
View File
@@ -1,4 +1,5 @@
use std::sync::{Arc, Mutex};
use std::time::SystemTime;
use anyhow::Result;
use aws_config::BehaviorVersion;
@@ -37,40 +38,17 @@ pub struct BedrockClientConfig {
pub access_key_id: String,
pub secret_access_key: String,
pub session_token: Option<String>,
pub credentials_expiry: Option<SystemTime>,
pub cross_region_inference: bool,
}
impl BedrockClientConfig {
/// Applies external config (from Claude Code / OpenCode) as fallback values
/// when Galaxy's own settings are at their defaults.
pub fn with_external_fallbacks(mut self) -> Self {
let external = ExternalBedrockConfig::load();
if external.is_empty() {
return self;
}
if self.profile == "default" {
if let Some(profile) = external.profile {
log::info!("[bedrock] Using profile from external config: {profile}");
self.profile = profile;
}
}
if self.region.is_empty() {
if let Some(region) = external.region {
log::info!("[bedrock] Using region from external config: {region}");
self.region = region;
}
}
self
}
}
#[derive(Debug, thiserror::Error)]
pub enum BedrockError {
#[error("Bedrock credentials not configured")]
CredentialsNotConfigured,
#[error("Failed to load AWS credentials: {0}")]
CredentialsLoadFailed(String),
#[error("Bedrock region not configured and could not be auto-detected")]
RegionNotConfigured,
#[error("Bedrock API error: {0}")]
@@ -115,12 +93,20 @@ impl BedrockClient {
if config.access_key_id.is_empty() || config.secret_access_key.is_empty() {
return Err(BedrockError::CredentialsNotConfigured);
}
if !super::authentication::credentials_are_current(
config.credentials_expiry,
SystemTime::now(),
) {
return Err(BedrockError::CredentialsLoadFailed(
"AWS credentials have expired".to_owned(),
));
}
let creds = aws_credential_types::Credentials::new(
&config.access_key_id,
&config.secret_access_key,
config.session_token,
None,
config.credentials_expiry,
"warp-bedrock-static",
);
@@ -148,11 +134,15 @@ impl BedrockClient {
);
}
Err(e) => {
log::warn!("[bedrock] Failed to resolve AWS credentials from provider: {e:?}");
let message = crate::ai::aws_credentials::user_facing_aws_credentials_error_message(
&e,
&config.profile,
);
return Err(BedrockError::CredentialsLoadFailed(message));
}
}
} else {
log::warn!("[bedrock] No credentials provider found in resolved AWS config");
return Err(BedrockError::CredentialsNotConfigured);
}
let region = aws_config
app/src/ai/provider/discovery.rs
+1 -1
View File
@@ -161,7 +161,7 @@ async fn load_aws_config(
&config.access_key_id,
&config.secret_access_key,
config.session_token.clone(),
None,
config.credentials_expiry,
"galaxy-bedrock-discovery",
);
let mut loader =
app/src/ai/provider/e2e_tests.rs
+1
View File
@@ -268,6 +268,7 @@ fn get_test_config() -> Option<BedrockClientConfig> {
access_key_id: String::new(),
secret_access_key: String::new(),
session_token: None,
credentials_expiry: None,
cross_region_inference: false,
})
}
app/src/ai/provider/integration_tests.rs
+1
View File
@@ -23,6 +23,7 @@ fn get_test_config() -> Option<BedrockClientConfig> {
access_key_id: String::new(),
secret_access_key: String::new(),
session_token: None,
credentials_expiry: None,
cross_region_inference: false,
})
}
app/src/ai/provider/mod.rs
+2
View File
@@ -16,6 +16,8 @@ pub enum ProviderConfig {
// Provider-neutral module boundary; Bedrock-specific transport helpers live here without a
// provider-named directory so the provider layer can be reorganized independently of callers.
#[cfg(not(target_family = "wasm"))]
pub(crate) mod authentication;
#[cfg(not(target_family = "wasm"))]
pub mod client;
#[cfg(not(target_family = "wasm"))]
pub mod convert;
app/src/settings_view/provider_setup_view.rs
+1
View File
@@ -1152,6 +1152,7 @@ impl ProviderSetupView {
access_key_id: self.draft_bedrock.access_key_id.clone(),
secret_access_key: self.draft_bedrock.secret_access_key.clone(),
session_token: None,
credentials_expiry: None,
cross_region_inference: self.draft_bedrock.cross_region_inference,
};
let requested_config = config.clone();
app/src/terminal/view.rs
+47 -103
View File
@@ -29,6 +29,8 @@ use warp_util::remote_path::RemotePath;
use warp_util::standardized_path::StandardizedPath;
use crate::ai::block_context::BlockContext;
#[cfg(not(target_family = "wasm"))]
use crate::ai::provider::authentication::{BedrockAuthConnection, BedrockAuthentication};
use crate::global_resource_handles::GlobalResourceHandlesProvider;
pub(crate) mod docker_sandbox;
mod link_detection;
@@ -10910,89 +10912,55 @@ impl TerminalView {
self.remove_aws_bedrock_login_banner(ctx);
}
/// Runs the AWS login command configured in settings to refresh Bedrock credentials.
/// Doing this in PTY vs just a subprocess allows the user to see any output/errors
/// from the command directly in the terminal. Also, `aws login` commands may require
/// user interaction (e.g. "do you want to override X profile? y/n" is common)
/// Refreshes the configured AWS connection and resumes only its originating conversation.
fn run_aws_login_command(&mut self, ctx: &mut ViewContext<Self>) {
let settings_command = AISettings::as_ref(ctx)
.bedrock_auth_refresh_command
.value()
.clone();
// Use the configured command, but if it's the bare default ("aws sso login")
// prefer the external config's auth_refresh_command which includes --profile.
let login_command = if settings_command == "aws sso login" {
use crate::ai::provider::external_config::ExternalBedrockConfig;
let external = ExternalBedrockConfig::load();
external.auth_refresh_command.unwrap_or(settings_command)
let conversation_id = if FeatureFlag::AgentView.is_enabled() {
self.agent_view_controller
.as_ref(ctx)
.agent_view_state()
.active_conversation_id()
} else {
settings_command
BlocklistAIHistoryModel::as_ref(ctx).last_conversation_id(self.id())
};
self.recover_aws_authentication(conversation_id, true, ctx);
}
if login_command.is_empty() {
log::warn!("AWS login command is not configured");
return;
}
log::info!("[bedrock] Running AWS login command: {login_command}");
// Spawn the login command as an async subprocess so it can open the browser
// for SSO authentication. Once it completes, refresh credentials and resume
// the conversation.
fn recover_aws_authentication(
&mut self,
conversation_id: Option<AIConversationId>,
manual: bool,
ctx: &mut ViewContext<Self>,
) {
let exchange_id = conversation_id
.and_then(|id| Self::thread_start_exchange_id(&id, ctx));
let connection = BedrockAuthConnection::current(ctx);
let conversation_key = conversation_id
.map(|id| id.to_string())
.unwrap_or_else(|| format!("login-only-{}", self.view_id));
let exchange_key = exchange_id.map(|id| id.to_string()).unwrap_or_default();
let recovery = BedrockAuthentication::handle(ctx).update(ctx, |authentication, ctx| {
authentication.recover(&conversation_key, &exchange_key, manual, ctx)
});
let _ = ctx.spawn(
async move {
let parts: Vec<&str> = login_command.split_whitespace().collect();
let Some((cmd, args)) = parts.split_first() else {
return Err("Empty login command".to_string());
};
let result = tokio::process::Command::new(cmd)
.args(args)
.status()
.await;
match result {
Ok(status) if status.success() => Ok(()),
Ok(status) => Err(format!("AWS login command exited with status: {status}")),
Err(e) => Err(format!("Failed to spawn AWS login command: {e}")),
}
},
|_me, result, ctx| {
recovery,
move |me, result, ctx| {
match result {
Ok(()) => {
log::info!("[bedrock] AWS login completed successfully; reloading credentials before resuming");
// Refresh through ApiKeyManager and wait for its completion. A fixed
// delay can resume the request while the SDK still has stale SSO data.
let refresh = ApiKeyManager::handle(ctx).update(ctx, |manager, ctx| {
crate::ai::aws_credentials::refresh_aws_credentials(manager, ctx)
});
let _ = ctx.spawn(refresh, |me, result, ctx| {
match result {
Ok(()) => {
log::info!("[bedrock] AWS credentials refreshed; resuming conversation");
let conversation_id = if FeatureFlag::AgentView.is_enabled() {
me.agent_view_controller
.as_ref(ctx)
.agent_view_state()
.active_conversation_id()
} else {
BlocklistAIHistoryModel::as_ref(ctx)
.last_conversation_id(me.id())
};
if let Some(conversation_id) = conversation_id {
me.handle_resume_conversation(&conversation_id, ctx);
}
}
Err(error) => {
log::error!(
"[bedrock] AWS login succeeded but credential reload failed; not resuming: {error}"
);
}
}
});
let Some(conversation_id) = conversation_id else {
return;
};
if BedrockAuthConnection::current(ctx) != connection
|| Self::thread_start_exchange_id(&conversation_id, ctx) != exchange_id
|| BlocklistAIHistoryModel::as_ref(ctx)
.conversation_status(&conversation_id) != Some(&ConversationStatus::Error)
{
return;
}
log::info!("[bedrock] AWS credentials refreshed; resuming originating conversation");
me.handle_resume_conversation(&conversation_id, ctx);
}
Err(e) => {
log::error!("[bedrock] AWS login failed: {e}");
Err(error) => {
log::warn!("[bedrock] AWS authentication recovery stopped: {error}");
}
}
},
@@ -20741,36 +20709,12 @@ impl TerminalView {
ctx.emit(Event::OpenThemeChooser);
}
AIBlockEvent::RunAwsLoginCommand => {
self.run_aws_login_command(ctx);
let conversation_id = block.as_ref(ctx).conversation_id();
self.recover_aws_authentication(Some(conversation_id), true, ctx);
}
AIBlockEvent::RefreshAwsCredentials => {
let refresh = ApiKeyManager::handle(ctx).update(ctx, |manager, ctx| {
crate::ai::aws_credentials::refresh_aws_credentials(manager, ctx)
});
let _ = ctx.spawn(refresh, |me, result, ctx| {
match result {
Ok(()) => {
log::info!("[bedrock] AWS credentials reloaded; resuming conversation");
let conversation_id = if FeatureFlag::AgentView.is_enabled() {
me.agent_view_controller
.as_ref(ctx)
.agent_view_state()
.active_conversation_id()
} else {
BlocklistAIHistoryModel::as_ref(ctx).last_conversation_id(me.id())
};
if let Some(conversation_id) = conversation_id {
me.handle_resume_conversation(&conversation_id, ctx);
}
}
Err(error) => {
log::warn!(
"[bedrock] AWS credential reload failed; starting configured login command: {error}"
);
me.run_aws_login_command(ctx);
}
}
});
let conversation_id = block.as_ref(ctx).conversation_id();
self.recover_aws_authentication(Some(conversation_id), false, ctx);
}
}
ctx.notify();