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();
crates/galaxy_agent_rig/src/stream.rs
+5
View File
@@ -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()
crates/galaxy_agent_rig/src/stream_unit_tests.rs
+16
View File
@@ -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);
}
}