auto refresh expired bedrock tokens
This commit is contained in:
@@ -8,6 +8,7 @@ pub struct ExternalBedrockConfig {
|
|||||||
pub profile: Option<String>,
|
pub profile: Option<String>,
|
||||||
pub region: Option<String>,
|
pub region: Option<String>,
|
||||||
pub models: Vec<BedrockModelConfig>,
|
pub models: Vec<BedrockModelConfig>,
|
||||||
|
pub auth_refresh_command: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl ExternalBedrockConfig {
|
impl ExternalBedrockConfig {
|
||||||
@@ -28,6 +29,7 @@ impl ExternalBedrockConfig {
|
|||||||
} else {
|
} else {
|
||||||
claude_config.models
|
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(),
|
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()) {
|
let env = match json.get("env").and_then(|v| v.as_object()) {
|
||||||
Some(e) => e,
|
Some(e) => e,
|
||||||
None => return ExternalBedrockConfig::default(),
|
None => return ExternalBedrockConfig {
|
||||||
|
auth_refresh_command,
|
||||||
|
..Default::default()
|
||||||
|
},
|
||||||
};
|
};
|
||||||
|
|
||||||
let profile = env
|
let profile = env
|
||||||
@@ -76,6 +87,7 @@ fn parse_claude_code_config(path: PathBuf) -> ExternalBedrockConfig {
|
|||||||
profile,
|
profile,
|
||||||
region,
|
region,
|
||||||
models,
|
models,
|
||||||
|
auth_refresh_command,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -221,6 +233,7 @@ fn parse_opencode_config(path: PathBuf) -> ExternalBedrockConfig {
|
|||||||
profile,
|
profile,
|
||||||
region,
|
region,
|
||||||
models: Vec::new(),
|
models: Vec::new(),
|
||||||
|
auth_refresh_command: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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 {
|
if let RenderableAIError::Other {
|
||||||
will_attempt_resume,
|
will_attempt_resume,
|
||||||
waiting_for_network,
|
waiting_for_network,
|
||||||
@@ -2994,6 +3009,24 @@ pub struct ClientIdentifiers {
|
|||||||
pub response_stream_id: Option<ResponseStreamId>,
|
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)]
|
#[allow(clippy::too_many_arguments)]
|
||||||
fn input_for_query(
|
fn input_for_query(
|
||||||
query: String,
|
query: String,
|
||||||
|
|||||||
@@ -149,6 +149,11 @@ impl ResponseStream {
|
|||||||
&self.params.bedrock_messages_sent
|
&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.
|
/// Returns true if we should attempt to resume the conversation after the stream finishes.
|
||||||
pub fn should_resume_conversation_after_stream_finished(&self) -> bool {
|
pub fn should_resume_conversation_after_stream_finished(&self) -> bool {
|
||||||
self.should_resume_conversation_after_stream_finished
|
self.should_resume_conversation_after_stream_finished
|
||||||
|
|||||||
@@ -310,6 +310,22 @@ impl AIApiError {
|
|||||||
{
|
{
|
||||||
false
|
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.
|
// By default, retry on error.
|
||||||
_ => true,
|
_ => true,
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1123,7 +1123,7 @@ define_settings_group!(AISettings, settings: [
|
|||||||
// Whether to automatically run the login command when Bedrock credentials expire.
|
// Whether to automatically run the login command when Bedrock credentials expire.
|
||||||
bedrock_auto_login: BedrockAutoLogin {
|
bedrock_auto_login: BedrockAutoLogin {
|
||||||
type: bool,
|
type: bool,
|
||||||
default: false,
|
default: true,
|
||||||
supported_platforms: SupportedPlatforms::DESKTOP,
|
supported_platforms: SupportedPlatforms::DESKTOP,
|
||||||
sync_to_cloud: SyncToCloud::Globally(RespectUserSyncSetting::Yes),
|
sync_to_cloud: SyncToCloud::Globally(RespectUserSyncSetting::Yes),
|
||||||
private: false,
|
private: false,
|
||||||
|
|||||||
@@ -9421,24 +9421,85 @@ impl TerminalView {
|
|||||||
/// from the command directly in the terminal. Also, `aws login` commands may require
|
/// 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)
|
/// 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>) {
|
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
|
.bedrock_auth_refresh_command
|
||||||
.value()
|
.value()
|
||||||
.clone();
|
.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() {
|
if login_command.is_empty() {
|
||||||
log::warn!("AWS login command is not configured");
|
log::warn!("AWS login command is not configured");
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Track that we're running an AWS login command so we can detect
|
log::info!("[bedrock] Running AWS login command: {login_command}");
|
||||||
// "command not found" if AWS CLI isn't installed
|
|
||||||
self.is_pending_aws_login = true;
|
|
||||||
|
|
||||||
// Write the command to the PTY and execute it
|
// Spawn the login command as an async subprocess so it can open the browser
|
||||||
let command_bytes = login_command.into_bytes();
|
// for SSO authentication. Once it completes, refresh credentials and resume
|
||||||
self.clear_line_editor_and_write_to_pty(command_bytes, ctx);
|
// the conversation.
|
||||||
self.write_to_pty(vec![escape_sequences::C0::CR], 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| {
|
||||||
|
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
|
/// Checks if the current model request could be served via AWS Bedrock and the user
|
||||||
|
|||||||
Reference in New Issue
Block a user