From 04cff915095c385000f58149a1fa7c16dd6046ee Mon Sep 17 00:00:00 2001 From: Josh Woodcock Date: Wed, 14 Oct 2026 12:43:32 -0500 Subject: [PATCH] fix bedrock authentication refresh --- app/src/ai/aws_credentials.rs | 56 +++- app/src/ai/blocklist/block.rs | 4 +- app/src/ai/blocklist/controller.rs | 29 +- .../blocklist/controller/response_stream.rs | 28 +- .../aws_bedrock_credentials_error.rs | 39 +-- app/src/ai/crosscheck/reviewer.rs | 27 +- app/src/ai/llms.rs | 10 +- app/src/ai/mod.rs | 2 + app/src/ai/provider/authentication.rs | 316 ++++++++++++++++++ app/src/ai/provider/authentication_tests.rs | 277 +++++++++++++++ app/src/ai/provider/client.rs | 48 ++- app/src/ai/provider/discovery.rs | 2 +- app/src/ai/provider/e2e_tests.rs | 1 + app/src/ai/provider/integration_tests.rs | 1 + app/src/ai/provider/mod.rs | 2 + app/src/settings_view/provider_setup_view.rs | 1 + app/src/terminal/view.rs | 150 +++------ crates/galaxy_agent_rig/src/stream.rs | 5 + .../galaxy_agent_rig/src/stream_unit_tests.rs | 16 + 19 files changed, 777 insertions(+), 237 deletions(-) create mode 100644 app/src/ai/provider/authentication.rs create mode 100644 app/src/ai/provider/authentication_tests.rs diff --git a/app/src/ai/aws_credentials.rs b/app/src/ai/aws_credentials.rs index 1829721d..12b45fe6 100644 --- a/app/src/ai/aws_credentials.rs +++ b/app/src/ai/aws_credentials.rs @@ -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 { - 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 { diff --git a/app/src/ai/blocklist/block.rs b/app/src/ai/blocklist/block.rs index faef21f8..f60bfc74 100644 --- a/app/src/ai/blocklist/block.rs +++ b/app/src/ai/blocklist/block.rs @@ -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 diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 5d810dc2..7a23ffd3 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -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.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, ) { 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, diff --git a/app/src/ai/blocklist/controller/response_stream.rs b/app/src/ai/blocklist/controller/response_stream.rs index a80f5347..7e7340c1 100644 --- a/app/src/ai/blocklist/controller/response_stream.rs +++ b/app/src/ai/blocklist/controller/response_stream.rs @@ -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 diff --git a/app/src/ai/blocklist/inline_action/aws_bedrock_credentials_error.rs b/app/src/ai/blocklist/inline_action/aws_bedrock_credentials_error.rs index 96e6800c..9aa89ea2 100644 --- a/app/src/ai/blocklist/inline_action/aws_bedrock_credentials_error.rs +++ b/app/src/ai/blocklist/inline_action/aws_bedrock_credentials_error.rs @@ -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, @@ -52,7 +50,6 @@ impl AwsBedrockCredentialsErrorView { pub fn new( model_name: String, login_command: String, - auto_login_triggered: bool, ctx: &mut ViewContext, ) -> 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., ) diff --git a/app/src/ai/crosscheck/reviewer.rs b/app/src/ai/crosscheck/reviewer.rs index 2e6bbcca..73ebb556 100644 --- a/app/src/ai/crosscheck/reviewer.rs +++ b/app/src/ai/crosscheck/reviewer.rs @@ -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 diff --git a/app/src/ai/llms.rs b/app/src/ai/llms.rs index 6c96743e..8779a857 100644 --- a/app/src/ai/llms.rs +++ b/app/src/ai/llms.rs @@ -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) } diff --git a/app/src/ai/mod.rs b/app/src/ai/mod.rs index 0d63273b..c956cd27 100644 --- a/app/src/ai/mod.rs +++ b/app/src/ai/mod.rs @@ -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"))] diff --git a/app/src/ai/provider/authentication.rs b/app/src/ai/provider/authentication.rs new file mode 100644 index 00000000..6431588e --- /dev/null +++ b/app/src/ai/provider/authentication.rs @@ -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, 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::() { + 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, 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, +} + +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; +type CredentialFuture = Shared>; + +#[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 { + 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, + ) -> 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::() + .map_or(0, |authentication| authentication.revision()) +} + +#[cfg(test)] +#[path = "authentication_tests.rs"] +mod tests; diff --git a/app/src/ai/provider/authentication_tests.rs b/app/src/ai/provider/authentication_tests.rs new file mode 100644 index 00000000..3b7d2ab1 --- /dev/null +++ b/app/src/ai/provider/authentication_tests.rs @@ -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())); + } +} diff --git a/app/src/ai/provider/client.rs b/app/src/ai/provider/client.rs index f179937e..4976e593 100644 --- a/app/src/ai/provider/client.rs +++ b/app/src/ai/provider/client.rs @@ -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, + pub credentials_expiry: Option, 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 diff --git a/app/src/ai/provider/discovery.rs b/app/src/ai/provider/discovery.rs index 3dd6263f..2fac4e1b 100644 --- a/app/src/ai/provider/discovery.rs +++ b/app/src/ai/provider/discovery.rs @@ -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 = diff --git a/app/src/ai/provider/e2e_tests.rs b/app/src/ai/provider/e2e_tests.rs index aedb3563..434739dc 100644 --- a/app/src/ai/provider/e2e_tests.rs +++ b/app/src/ai/provider/e2e_tests.rs @@ -268,6 +268,7 @@ fn get_test_config() -> Option { access_key_id: String::new(), secret_access_key: String::new(), session_token: None, + credentials_expiry: None, cross_region_inference: false, }) } diff --git a/app/src/ai/provider/integration_tests.rs b/app/src/ai/provider/integration_tests.rs index 155b3424..5a52394a 100644 --- a/app/src/ai/provider/integration_tests.rs +++ b/app/src/ai/provider/integration_tests.rs @@ -23,6 +23,7 @@ fn get_test_config() -> Option { access_key_id: String::new(), secret_access_key: String::new(), session_token: None, + credentials_expiry: None, cross_region_inference: false, }) } diff --git a/app/src/ai/provider/mod.rs b/app/src/ai/provider/mod.rs index 253c09db..e7394c83 100644 --- a/app/src/ai/provider/mod.rs +++ b/app/src/ai/provider/mod.rs @@ -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; diff --git a/app/src/settings_view/provider_setup_view.rs b/app/src/settings_view/provider_setup_view.rs index 4e8278b4..7d089565 100644 --- a/app/src/settings_view/provider_setup_view.rs +++ b/app/src/settings_view/provider_setup_view.rs @@ -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(); diff --git a/app/src/terminal/view.rs b/app/src/terminal/view.rs index a513b994..4610cb8e 100644 --- a/app/src/terminal/view.rs +++ b/app/src/terminal/view.rs @@ -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) { - 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, + manual: bool, + ctx: &mut ViewContext, + ) { + 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(); diff --git a/crates/galaxy_agent_rig/src/stream.rs b/crates/galaxy_agent_rig/src/stream.rs index 3379040c..bae1df20 100644 --- a/crates/galaxy_agent_rig/src/stream.rs +++ b/crates/galaxy_agent_rig/src/stream.rs @@ -512,6 +512,11 @@ fn text_indicates_authentication_failure(text: &str) -> bool { "unrecognized client", "security token included in the request is invalid", "unable to locate credentials", + "unauthorizedexception", + "session token not found or invalid", + "sso session has expired", + "failed to load aws credentials", + "bedrock credentials not configured", "credential should be scoped to a valid region", ] .iter() diff --git a/crates/galaxy_agent_rig/src/stream_unit_tests.rs b/crates/galaxy_agent_rig/src/stream_unit_tests.rs index ac737561..f103a275 100644 --- a/crates/galaxy_agent_rig/src/stream_unit_tests.rs +++ b/crates/galaxy_agent_rig/src/stream_unit_tests.rs @@ -105,6 +105,10 @@ fn flattened_bedrock_credential_errors_are_authentication_failures() { r#"{"__type":"ExpiredTokenException","message":"The security token included in the request is expired"}"#, "UnrecognizedClientException: The security token included in the request is invalid", "AccessDeniedException: not authorized to perform bedrock:ConverseStream", + "UnauthorizedException: Session token not found or invalid", + "SSO session has expired or is otherwise invalid", + "Failed to load AWS credentials: credentials were not loaded", + "Bedrock credentials not configured", ] { let mapped = map_completion_error(CompletionError::ProviderError(message.to_string())); assert_eq!(mapped.kind, AgentErrorKind::Authentication); @@ -112,3 +116,15 @@ fn flattened_bedrock_credential_errors_are_authentication_failures() { assert!(mapped.user_message.is_some()); } } + +#[test] +fn unrelated_provider_failures_do_not_become_authentication_failures() { + for message in [ + "HTTP client error: connection reset", + "ProviderError: temporarily unavailable", + "ValidationException: invalid model", + ] { + let mapped = map_completion_error(CompletionError::ProviderError(message.to_owned())); + assert_ne!(mapped.kind, AgentErrorKind::Authentication); + } +}