From 0e727914832a0e137986aa2c606eba5b50834421 Mon Sep 17 00:00:00 2001 From: Josh Woodcock Date: Wed, 27 May 2026 13:29:11 -0500 Subject: [PATCH] auto refresh expired bedrock tokens --- app/src/ai/bedrock/external_config.rs | 15 +++- app/src/ai/blocklist/controller.rs | 35 ++++++++- .../blocklist/controller/response_stream.rs | 5 ++ app/src/server/server_api.rs | 16 ++++ app/src/settings/ai.rs | 2 +- app/src/terminal/view.rs | 77 +++++++++++++++++-- 6 files changed, 139 insertions(+), 11 deletions(-) diff --git a/app/src/ai/bedrock/external_config.rs b/app/src/ai/bedrock/external_config.rs index 12534784..abfd0158 100644 --- a/app/src/ai/bedrock/external_config.rs +++ b/app/src/ai/bedrock/external_config.rs @@ -8,6 +8,7 @@ pub struct ExternalBedrockConfig { pub profile: Option, pub region: Option, pub models: Vec, + pub auth_refresh_command: Option, } impl ExternalBedrockConfig { @@ -28,6 +29,7 @@ impl ExternalBedrockConfig { } else { claude_config.models }, + auth_refresh_command: claude_config.auth_refresh_command, } } @@ -55,9 +57,18 @@ fn parse_claude_code_config(path: PathBuf) -> ExternalBedrockConfig { Err(_) => return ExternalBedrockConfig::default(), }; + // Read the top-level awsAuthRefresh command (used by Claude Code for SSO login) + let auth_refresh_command = json + .get("awsAuthRefresh") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + let env = match json.get("env").and_then(|v| v.as_object()) { Some(e) => e, - None => return ExternalBedrockConfig::default(), + None => return ExternalBedrockConfig { + auth_refresh_command, + ..Default::default() + }, }; let profile = env @@ -76,6 +87,7 @@ fn parse_claude_code_config(path: PathBuf) -> ExternalBedrockConfig { profile, region, models, + auth_refresh_command, } } @@ -221,6 +233,7 @@ fn parse_opencode_config(path: PathBuf) -> ExternalBedrockConfig { profile, region, models: Vec::new(), + auth_refresh_command: None, } } diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 7274a7fa..655c7835 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -2508,7 +2508,22 @@ impl BlocklistAIController { }); } - let mut renderable_error: RenderableAIError = e.as_ref().into(); + let mut renderable_error: RenderableAIError = + if let AIApiError::Stream { stream_type, source } = e.as_ref() { + if *stream_type == "bedrock_converse" + && is_bedrock_credentials_error(&source.to_string()) + { + let model_name = + response_stream.as_ref(ctx).model_id().to_string(); + RenderableAIError::AwsBedrockCredentialsExpiredOrInvalid { + model_name, + } + } else { + e.as_ref().into() + } + } else { + e.as_ref().into() + }; if let RenderableAIError::Other { will_attempt_resume, waiting_for_network, @@ -2994,6 +3009,24 @@ pub struct ClientIdentifiers { pub response_stream_id: Option, } +/// Returns `true` if the given error message from a Bedrock stream indicates an +/// AWS credentials issue (expired, invalid, or missing session token). +fn is_bedrock_credentials_error(msg: &str) -> bool { + let lower = msg.to_lowercase(); + // "Session token not found or invalid" is the most common SSO expiry message. + // AccessDenied / ExpiredToken / UnrecognizedClient cover other credential failures. + // SSO cache file not found means the token file was deleted or never created. + lower.contains("session token not found") + || lower.contains("expiredtoken") + || lower.contains("expired token") + || lower.contains("unrecognizedclientexception") + || lower.contains("unauthorizedexception") + || (lower.contains("sso/cache") && lower.contains("notfound")) + || (lower.contains("sso/cache") && lower.contains("no such file")) + || (lower.contains("accessdenied") + && (lower.contains("token") || lower.contains("credential") || lower.contains("security"))) +} + #[allow(clippy::too_many_arguments)] fn input_for_query( query: String, diff --git a/app/src/ai/blocklist/controller/response_stream.rs b/app/src/ai/blocklist/controller/response_stream.rs index 6f16989d..a163ffb2 100644 --- a/app/src/ai/blocklist/controller/response_stream.rs +++ b/app/src/ai/blocklist/controller/response_stream.rs @@ -149,6 +149,11 @@ impl ResponseStream { &self.params.bedrock_messages_sent } + /// Returns the model ID associated with this response stream's request. + pub fn model_id(&self) -> &str { + self.params.model.as_str() + } + /// Returns true if we should attempt to resume the conversation after the stream finishes. pub fn should_resume_conversation_after_stream_finished(&self) -> bool { self.should_resume_conversation_after_stream_finished diff --git a/app/src/server/server_api.rs b/app/src/server/server_api.rs index 8617ecfb..53f2b24a 100644 --- a/app/src/server/server_api.rs +++ b/app/src/server/server_api.rs @@ -310,6 +310,22 @@ impl AIApiError { { false } + // Don't retry Bedrock credential/auth errors — they require user + // action (e.g. SSO re-login) and will always fail with the same token. + AIApiError::Stream { source, .. } + if { + let msg = source.to_string().to_lowercase(); + msg.contains("session token not found") + || msg.contains("expiredtoken") + || msg.contains("expired token") + || msg.contains("unrecognizedclientexception") + || msg.contains("unauthorizedexception") + || (msg.contains("sso/cache") && msg.contains("no such file")) + || (msg.contains("sso/cache") && msg.contains("notfound")) + } => + { + false + } // By default, retry on error. _ => true, } diff --git a/app/src/settings/ai.rs b/app/src/settings/ai.rs index 7dcb5f4e..efe029dd 100644 --- a/app/src/settings/ai.rs +++ b/app/src/settings/ai.rs @@ -1123,7 +1123,7 @@ define_settings_group!(AISettings, settings: [ // Whether to automatically run the login command when Bedrock credentials expire. bedrock_auto_login: BedrockAutoLogin { type: bool, - default: false, + default: true, supported_platforms: SupportedPlatforms::DESKTOP, sync_to_cloud: SyncToCloud::Globally(RespectUserSyncSetting::Yes), private: false, diff --git a/app/src/terminal/view.rs b/app/src/terminal/view.rs index 41ce37d7..19da5ad9 100644 --- a/app/src/terminal/view.rs +++ b/app/src/terminal/view.rs @@ -9421,24 +9421,85 @@ impl TerminalView { /// 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) fn run_aws_login_command(&mut self, ctx: &mut ViewContext) { - let login_command = AISettings::as_ref(ctx) + 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::bedrock::external_config::ExternalBedrockConfig; + let external = ExternalBedrockConfig::load(); + external.auth_refresh_command.unwrap_or(settings_command) + } else { + settings_command + }; + if login_command.is_empty() { log::warn!("AWS login command is not configured"); return; } - // Track that we're running an AWS login command so we can detect - // "command not found" if AWS CLI isn't installed - self.is_pending_aws_login = true; + log::info!("[bedrock] Running AWS login command: {login_command}"); - // Write the command to the PTY and execute it - let command_bytes = login_command.into_bytes(); - self.clear_line_editor_and_write_to_pty(command_bytes, ctx); - self.write_to_pty(vec![escape_sequences::C0::CR], ctx); + // 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. + 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| { + match result { + Ok(()) => { + log::info!("[bedrock] AWS login completed successfully, refreshing credentials and resuming"); + // Refresh credentials from the updated SSO cache + ApiKeyManager::handle(ctx).update( + ctx, + |manager, ctx| { + drop(crate::ai::aws_credentials::refresh_aws_credentials(manager, ctx)); + }, + ); + // Resume the conversation after a short delay to let credentials load + let _ = ctx.spawn( + async move { + tokio::time::sleep(std::time::Duration::from_millis(500)).await; + }, + |me, _, ctx| { + // Find the active conversation and resume it directly + 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(e) => { + log::error!("[bedrock] AWS login failed: {e}"); + } + } + }, + ); } /// Checks if the current model request could be served via AWS Bedrock and the user