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

auto refresh expired bedrock tokens

Browse Source
This commit is contained in:
Josh Woodcock
2026-05-27 13:29:11 -05:00
parent d55befec1c
commit 0e72791483
6 changed files with 139 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/bedrock/external_config.rs
+14 -1
View File
@@ -8,6 +8,7 @@ pub struct ExternalBedrockConfig {
pub profile: Option<String>,
pub region: Option<String>,
pub models: Vec<BedrockModelConfig>,
pub auth_refresh_command: Option<String>,
}
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,
}
}
app/src/ai/blocklist/controller.rs
+34 -1
View File
@@ -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<ResponseStreamId>,
}
/// 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,
app/src/ai/blocklist/controller/response_stream.rs
+5
View File
@@ -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
app/src/server/server_api.rs
+16
View File
@@ -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,
}
app/src/settings/ai.rs
+1 -1
View File
@@ -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,
app/src/terminal/view.rs
+69 -8
View File
@@ -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<Self>) {
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