Compare commits

..
5 Commits
34 changed files with 2101 additions and 596 deletions
+5 -3
View File
@@ -5,6 +5,7 @@ This file provides guidance when working with code in this repository.
## Development Commands
### Build and Run
- `cargo build` - Build Galaxy and verify feature and bug-fix changes
- `cargo run` - Build and run Warp locally
- `cargo bundle --bin warp` - Bundle the main app
- `make run` - Build and run `galaxy-oss`
@@ -28,8 +29,9 @@ Environment variables:
- `WS_SERVER_URL` - WebSocket endpoint (default: `ws://localhost:8080/graphql/v2`)
### Testing
- Before launching `galaxy-oss`, check for existing app instances and close the prior development instance so only one app instance runs at a time. Distinguish the app process from its `terminal-server` helper; do not count the helper as a duplicate app. Leave the newly launched instance running for the user.
- During interactive feature and bug-fix verification, make the requested code changes first and run only `cargo run --bin galaxy-oss` for the user to verify. Leave the app running for the user; do not stop it or treat the command timeout as a failure. Do not run `cargo fmt`, `cargo check`, `cargo test`, `cargo clippy`, presubmit, or other validation commands until the user confirms the behavior. Run formatting, tests, and linting only as the final cleanup step after interactive approval.
- This development machine is CLI-only and has no GUI. The user performs GUI and most behavioral testing on a remote machine.
- Verify feature and bug-fix changes with `cargo build` only. Do not use `cargo run`, `make run`, or launch the app to verify results, and do not wait for interactive approval before completing the work.
- After `cargo build` succeeds and all requested work and final review are complete, commit the changes and push the current branch to both configured remotes, `origin` and `samsung`, using ordinary non-force pushes. If the build fails, report the blocker and do not push.
- `cargo nextest run --no-fail-fast --workspace --exclude command-signatures-v2` - Run tests with nextest
- `cargo nextest run -p galaxy_completer --features v2` - Run completer tests with v2 features
- `cargo test --doc` - Run doc tests
@@ -143,7 +145,7 @@ ACP: controller.rs → ResponseStream → acp_output_stream (session-owned lifec
- `convert.rs` — Bedrock request construction and prompt-caching behavior
- `client.rs` — AWS SDK client construction, runtime creation, and independent background streaming calls
- `models.rs` — Model registry and cross-region inference prefix logic
- `discovery.rs` — Checks bundled `galaxy_bedrock_model_catalog` candidates with AWS `GetFoundationModelAvailability`; agreement, authorization, entitlement, and regional availability must all pass. Cached access is never restored as authorization; per-model enablement survives successful refreshes. Sweeps run once at startup when enabled and once per explicit connection Test; settings saves, reloads, and model toggles must never trigger another sweep. Live verification belongs to `LLMPreferences` and is scoped to the exact connection configuration.
- `discovery.rs` — Discovers the live inventory with AWS `ListFoundationModels` and paginated `ListInferenceProfiles`, resolving profile-only model details with `GetFoundationModel`. The bundled `galaxy_bedrock_model_catalog` is optional capability/token-budget enrichment, never an inventory allowlist. Invocation uses AWS-returned model IDs/profile ARNs rather than guessed prefixes. `GetFoundationModelAvailability` still requires agreement, authorization, entitlement, and regional availability; cached access is never restored as authorization. Saved model enablement and names survive refreshes, and failed discovery must not overwrite preferences with an empty catalog. Sweeps run once at startup when enabled and once per explicit connection Test; settings saves, credential reloads, login completion, and model toggles must never trigger another sweep. Live verification belongs to `LLMPreferences` and is scoped to the exact connection configuration. Prompt-cache preferences must never hide models; unknown cache capabilities use no caching.
- `diagnostic.rs` — Debug logging (enabled via `GALAXY_BEDROCK_DIAGNOSTICS=1`)
- `external_config.rs` — Fallback config from Claude Code/OpenCode settings
+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 {
+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
+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,
@@ -187,30 +187,14 @@ 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_model_route(model_id, settings.bedrock_models.value())
.with_managed_credentials(
&api_key_manager.aws_credentials_refresh_strategy(),
api_key_manager.aws_credentials_state(),
);
return ProviderConfig::Bedrock(config);
}
ProviderConfig::None
@@ -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.,
)
+356 -224
View File
@@ -2,12 +2,17 @@
use std::io::{BufRead, BufReader, Write};
use std::net::{TcpListener, TcpStream};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::Duration;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use base64::Engine;
use futures::channel::oneshot;
use futures::StreamExt;
use galaxy_agent_rig::ChatGPTCredentialResolver;
use galaxyui::{Entity, ModelContext, SingletonEntity};
use galaxyui_extras::secure_storage::AppContextExt as _;
use instant::Instant;
use rand::RngCore;
use sha2::{Digest, Sha256};
@@ -21,6 +26,10 @@ const CALLBACK_HOST: &str = "localhost";
const CALLBACK_PORTS: [u16; 2] = [1455, 1457];
const CALLBACK_PATH: &str = "/auth/callback";
const CALLBACK_TIMEOUT: Duration = Duration::from_secs(2 * 60);
const AUTH_STORAGE_KEY: &str = "galaxy.chatgpt.oauth";
const SIGN_IN_REQUIRED: &str = "ChatGPT sign-in required. Connect or reconnect ChatGPT in Providers.";
type Reply = oneshot::Sender<Result<ChatGPTAuthCredentials, String>>;
static RESOLVER: OnceLock<Mutex<Option<async_channel::Sender<Reply>>>> = OnceLock::new();
/// Current state of the local ChatGPT subscription connection.
#[derive(Clone, Debug, PartialEq, Eq)]
@@ -45,26 +54,42 @@ struct OAuthCallback {
/// Coordinates browser-based OAuth authorization for ChatGPT subscriptions.
pub(crate) struct ChatGPTAuthModel {
state: ChatGPTAuthState,
record: Option<AuthRecord>,
refresh: RefreshCoordinator,
sign_in_generation: u64,
sign_in_cancel: Option<Arc<AtomicBool>>,
}
impl ChatGPTAuthModel {
pub(crate) fn new() -> Self {
let state = match has_or_import_auth_credentials() {
Ok(()) => ChatGPTAuthState::Connected,
Err(error) => {
log::debug!(
"[chatgpt/auth] No usable persisted ChatGPT credentials at startup: {error}"
);
ChatGPTAuthState::NotConnected
}
pub(crate) fn new(ctx: &mut ModelContext<Self>) -> Self {
let (tx, rx) = async_channel::bounded(64);
*RESOLVER.get_or_init(|| Mutex::new(None)).lock().unwrap() = Some(tx);
ctx.spawn_stream_local(rx, |model, reply, ctx| model.resolve(reply, ctx), |_, _| {});
let record = ctx.secure_storage().read_value(AUTH_STORAGE_KEY).ok()
.and_then(|json| serde_json::from_str::<AuthRecord>(&json).ok())
.filter(|record| record.credentials(now()).is_ok() || record.can_refresh());
let state = if record.is_some() {
ChatGPTAuthState::Connected
} else {
ChatGPTAuthState::NotConnected
};
Self { state }
Self {
state,
record,
refresh: RefreshCoordinator::default(),
sign_in_generation: 0,
sign_in_cancel: None,
}
}
#[cfg(test)]
pub(crate) fn new_for_test() -> Self {
Self {
state: ChatGPTAuthState::NotConnected,
record: None,
refresh: RefreshCoordinator::default(),
sign_in_generation: 0,
sign_in_cancel: None,
}
}
@@ -72,22 +97,26 @@ impl ChatGPTAuthModel {
&self.state
}
/// Reuses persisted credentials when possible, otherwise starts browser OAuth with PKCE.
pub(crate) fn connect(&mut self, ctx: &mut ModelContext<Self>) {
fn begin_sign_in(&mut self) -> Option<u64> {
if matches!(
self.state,
ChatGPTAuthState::AwaitingBrowser | ChatGPTAuthState::ExchangingToken
) {
return;
return None;
}
if has_or_import_auth_credentials().is_ok() {
self.state = ChatGPTAuthState::Connected;
ctx.emit(ChatGPTAuthModelEvent::StateChanged);
return;
}
self.refresh.invalidate();
self.sign_in_generation += 1;
Some(self.sign_in_generation)
}
let (redirect_uri, callback_rx) = match start_oauth_callback_server() {
/// Explicit connection always starts a fresh browser OAuth flow with PKCE.
pub(crate) fn connect(&mut self, ctx: &mut ModelContext<Self>) {
let Some(generation) = self.begin_sign_in() else {
return;
};
let cancelled = Arc::new(AtomicBool::new(false));
let (redirect_uri, callback_rx) = match start_oauth_callback_server(cancelled.clone()) {
Ok(server) => server,
Err(error) => {
self.fail(error, ctx);
@@ -106,8 +135,8 @@ impl ChatGPTAuthModel {
}
};
self.state = ChatGPTAuthState::AwaitingBrowser;
ctx.emit(ChatGPTAuthModelEvent::StateChanged);
self.sign_in_cancel = Some(cancelled);
self.set_state(ChatGPTAuthState::AwaitingBrowser, ctx);
ctx.open_url(authorize_url.as_str());
let _ = ctx.spawn(
@@ -117,6 +146,9 @@ impl ChatGPTAuthModel {
})?
},
move |model, callback, ctx| {
if model.sign_in_generation != generation {
return;
}
let callback = match callback {
Ok(callback) => callback,
Err(error) => {
@@ -129,19 +161,27 @@ impl ChatGPTAuthModel {
return;
}
model.state = ChatGPTAuthState::ExchangingToken;
ctx.emit(ChatGPTAuthModelEvent::StateChanged);
model.set_state(ChatGPTAuthState::ExchangingToken, ctx);
let _ = ctx.spawn(
async move {
exchange_code_for_tokens(&callback.code, &code_verifier, &redirect_uri)
.await
},
|model, result, ctx| match result {
Ok(()) => {
model.state = ChatGPTAuthState::Connected;
ctx.emit(ChatGPTAuthModelEvent::StateChanged);
move |model, result, ctx| {
if model.sign_in_generation != generation {
return;
}
match result.and_then(|record| {
save_auth_record(&record, ctx)?;
Ok(record)
}) {
Ok(record) => {
model.record = Some(record);
model.sign_in_cancel = None;
model.set_state(ChatGPTAuthState::Connected, ctx);
}
Err(error) => model.fail(error, ctx),
}
Err(error) => model.fail(error, ctx),
},
);
},
@@ -149,8 +189,77 @@ impl ChatGPTAuthModel {
}
fn fail(&mut self, error: String, ctx: &mut ModelContext<Self>) {
self.state = ChatGPTAuthState::Failed(error);
ctx.emit(ChatGPTAuthModelEvent::StateChanged);
if let Some(cancelled) = self.sign_in_cancel.take() {
cancelled.store(true, Ordering::SeqCst);
}
self.set_state(ChatGPTAuthState::Failed(error), ctx);
}
fn set_state(&mut self, state: ChatGPTAuthState, ctx: &mut ModelContext<Self>) {
if self.state != state {
self.state = state;
ctx.emit(ChatGPTAuthModelEvent::StateChanged);
}
}
pub(crate) fn cancel_sign_in(&mut self, ctx: &mut ModelContext<Self>) {
if let Some(cancelled) = self.sign_in_cancel.take() {
cancelled.store(true, Ordering::SeqCst);
self.sign_in_generation += 1;
self.refresh.invalidate();
self.set_state(ChatGPTAuthState::NotConnected, ctx);
}
}
fn resolve(&mut self, reply: Reply, ctx: &mut ModelContext<Self>) {
if reply.is_canceled() {
return;
}
if matches!(self.state, ChatGPTAuthState::AwaitingBrowser | ChatGPTAuthState::ExchangingToken) {
let _ = reply.send(Err("ChatGPT sign-in is in progress.".to_owned()));
return;
}
let Some(record) = self.record.as_ref() else {
let _ = reply.send(Err(SIGN_IN_REQUIRED.to_owned()));
return;
};
if let Ok(credentials) = record.credentials(now()) {
let _ = reply.send(Ok(credentials));
return;
}
if !record.can_refresh() {
let _ = reply.send(Err(SIGN_IN_REQUIRED.to_owned()));
self.fail(SIGN_IN_REQUIRED.to_owned(), ctx);
return;
}
let record = record.clone();
let Some(generation) = self.refresh.begin(reply) else {
return;
};
let _ = ctx.spawn(
async move { refresh_tokens(&record).await },
move |model, result, ctx| {
let completed = model.refresh.complete(generation, result, |record| {
save_auth_record(record, ctx)
});
match completed {
None => {}
Some(Ok(record)) => {
model.record = Some(record);
model.set_state(ChatGPTAuthState::Connected, ctx);
}
Some(Err(error)) => {
if error.reauthenticate {
model.record = None;
if save_auth_record(&AuthRecord::default(), ctx).is_err() {
log::warn!("[chatgpt/auth] Could not persist reconnect-required state");
}
}
model.fail(error.message, ctx);
}
}
},
);
}
}
@@ -194,6 +303,7 @@ fn random_hex_string(byte_count: usize) -> String {
}
fn start_oauth_callback_server(
cancelled: Arc<AtomicBool>,
) -> Result<(String, oneshot::Receiver<Result<OAuthCallback, String>>), String> {
let mut bind_errors = Vec::new();
let mut bound = None;
@@ -222,6 +332,10 @@ fn start_oauth_callback_server(
.spawn(move || {
let deadline = Instant::now() + CALLBACK_TIMEOUT;
loop {
if cancelled.load(Ordering::SeqCst) {
let _ = callback_tx.send(Err("ChatGPT sign-in was cancelled.".to_owned()));
return;
}
match listener.accept() {
Ok((stream, _)) => {
let _ = callback_tx.send(handle_callback_connection(stream));
@@ -323,26 +437,19 @@ fn parse_callback_request_line(request_line: &str) -> Result<OAuthCallback, Stri
let mut code = None;
let mut state = None;
let mut oauth_error = None;
let mut error_description = None;
let mut oauth_error = false;
for (key, value) in url.query_pairs() {
match key.as_ref() {
"code" if !value.is_empty() => code = Some(value.into_owned()),
"state" if !value.is_empty() => state = Some(value.into_owned()),
"error" if !value.is_empty() => oauth_error = Some(value.into_owned()),
"error_description" if !value.is_empty() => {
error_description = Some(value.into_owned())
}
"error" if !value.is_empty() => oauth_error = true,
"code" | "state" | "error" | "error_description" => {}
_ => {}
}
}
if let Some(error) = oauth_error {
return Err(format!(
"ChatGPT authorization failed: {error} ({})",
error_description.as_deref().unwrap_or("no description")
));
if oauth_error {
return Err("ChatGPT browser authorization failed. Please try signing in again.".to_owned());
}
Ok(OAuthCallback {
code: code.ok_or("ChatGPT OAuth callback did not include an authorization code.")?,
@@ -358,157 +465,101 @@ fn callback_response_page(title: &str, message: &str, is_error: bool) -> String
.replace("__ACCENT__", accent)
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub(crate) struct ChatGPTAuthCredentials {
pub(crate) access_token: String,
pub(crate) account_id: Option<String>,
}
pub(crate) use galaxy_agent_rig::ChatGPTSubscriptionCredentials as ChatGPTAuthCredentials;
/// Attempts to read tokens from `~/.codex/auth.json` and write them to Rig's auth file.
/// Returns `Ok(())` if credentials were found and successfully imported.
pub(crate) fn import_codex_credentials() -> Result<(), String> {
let codex_path = codex_auth_file_path().ok_or("Cannot determine codex auth path")?;
let bytes = std::fs::read(&codex_path).map_err(|e| format!("{e}"))?;
let doc: serde_json::Value = serde_json::from_slice(&bytes).map_err(|e| format!("{e}"))?;
let tokens = doc.get("tokens").ok_or("No tokens object")?;
let access_token = tokens
.get("access_token")
.and_then(|value| value.as_str())
.filter(|token| !token.is_empty());
let refresh_token = tokens
.get("refresh_token")
.and_then(|value| value.as_str())
.filter(|token| !token.is_empty())
.map(ToOwned::to_owned);
if access_token.is_none() && refresh_token.is_none() {
return Err("Codex auth file does not contain an access or refresh token".to_string());
}
let expires_at = access_token.and_then(extract_expiration_timestamp);
let id_token = tokens
.get("id_token")
.and_then(|v| v.as_str())
.map(ToOwned::to_owned);
let account_id = tokens
.get("account_id")
.and_then(|v| v.as_str())
.map(ToOwned::to_owned)
.or_else(|| extract_account_id(id_token.as_deref()))
.or_else(|| extract_account_id(access_token));
let record = AuthRecord {
access_token: access_token.map(ToOwned::to_owned),
refresh_token,
id_token,
expires_at,
account_id,
};
write_auth_file(&record)
pub(crate) fn credential_resolver() -> ChatGPTCredentialResolver {
Arc::new(|| Box::pin(load_or_refresh_auth_credentials()))
}
pub(crate) async fn load_or_refresh_auth_credentials() -> Result<ChatGPTAuthCredentials, String> {
if let Ok(credentials) = load_auth_credentials() {
return Ok(credentials);
}
let _ = import_codex_credentials();
if let Ok(credentials) = load_auth_credentials() {
return Ok(credentials);
}
galaxy_agent_rig::refresh_chatgpt_subscription_credentials().await?;
load_auth_credentials()
}
fn has_or_import_auth_credentials() -> Result<(), String> {
if has_auth_credentials() {
return Ok(());
}
import_codex_credentials()?;
if has_auth_credentials() {
Ok(())
} else {
Err("ChatGPT credentials are missing an access or refresh token.".to_string())
let (tx, rx) = oneshot::channel();
RESOLVER.get_or_init(|| Mutex::new(None)).lock().unwrap()
.as_ref().ok_or("ChatGPT authorization service is unavailable.")?
.try_send(tx).map_err(|_| "ChatGPT authorization service is busy or unavailable.")?;
match futures::future::select(
Box::pin(rx),
Box::pin(async_io::Timer::after(Duration::from_secs(35))),
).await {
futures::future::Either::Left((result, _)) => {
result.map_err(|_| "ChatGPT authorization service stopped.".to_owned())?
}
futures::future::Either::Right(_) => {
Err("ChatGPT credential refresh timed out. Retry or sign in again.".to_owned())
}
}
}
fn has_auth_credentials() -> bool {
let Some(path) = auth_file_path() else {
return false;
};
let Ok(bytes) = std::fs::read(path) else {
return false;
};
let Ok(record) = serde_json::from_slice::<AuthRecord>(&bytes) else {
return false;
};
record
.refresh_token
.as_deref()
.is_some_and(|token| !token.trim().is_empty())
|| record.access_token.as_deref().is_some_and(|token| {
!token.trim().is_empty()
&& !token_is_expired(
record
.expires_at
.or_else(|| extract_expiration_timestamp(token)),
)
})
#[derive(Clone, Debug)]
struct AuthFailure {
message: String,
reauthenticate: bool,
}
fn load_auth_credentials() -> Result<ChatGPTAuthCredentials, String> {
let path = auth_file_path().ok_or("Cannot determine ChatGPT auth file path")?;
let bytes = std::fs::read(&path)
.map_err(|error| format!("Failed to read {}: {error}", path.display()))?;
let record: AuthRecord = serde_json::from_slice(&bytes)
.map_err(|error| format!("Failed to parse {}: {error}", path.display()))?;
let access_token = record
.access_token
.as_deref()
.filter(|token| !token.trim().is_empty())
.ok_or("ChatGPT auth file does not contain an access token")?;
#[derive(Default)]
struct RefreshCoordinator {
generation: u64,
active: Option<(u64, Vec<Reply>)>,
}
let expires_at = record
.expires_at
.or_else(|| extract_expiration_timestamp(access_token));
if token_is_expired(expires_at) {
return Err("ChatGPT access token is expired".to_string());
impl RefreshCoordinator {
fn begin(&mut self, reply: Reply) -> Option<u64> {
if let Some((_, replies)) = &mut self.active {
replies.push(reply);
return None;
}
self.generation += 1;
self.active = Some((self.generation, vec![reply]));
Some(self.generation)
}
let account_id = record
.account_id
.clone()
.filter(|account_id| !account_id.trim().is_empty())
.or_else(|| extract_account_id(record.id_token.as_deref()))
.or_else(|| extract_account_id(Some(access_token)));
fn invalidate(&mut self) {
self.generation += 1;
if let Some((_, replies)) = self.active.take() {
for reply in replies {
let _ = reply.send(Err("ChatGPT session was replaced. Please sign in again.".to_owned()));
}
}
}
Ok(ChatGPTAuthCredentials {
access_token: access_token.to_string(),
account_id,
})
}
fn codex_auth_file_path() -> Option<std::path::PathBuf> {
if let Some(codex_home) = std::env::var_os("CODEX_HOME") {
return Some(std::path::PathBuf::from(codex_home).join("auth.json"));
fn complete(
&mut self,
generation: u64,
result: Result<AuthRecord, AuthFailure>,
persist: impl FnOnce(&AuthRecord) -> Result<(), String>,
) -> Option<Result<AuthRecord, AuthFailure>> {
if self.active.as_ref().is_none_or(|(active, _)| *active != generation) {
return None;
}
let (_, replies) = self.active.take()?;
let result = result.and_then(|record| {
let credentials = record.access_context().map_err(|message| AuthFailure {
message, reauthenticate: true,
})?;
persist(&record).map_err(|_| AuthFailure {
message: "Could not securely save refreshed ChatGPT credentials. Please sign in again.".to_owned(),
reauthenticate: true,
})?;
Ok((record, credentials))
});
for reply in replies {
let response = result.as_ref().map(|(_, credentials)| credentials.clone())
.map_err(|error| error.message.clone());
let _ = reply.send(response);
}
Some(result.map(|(record, _)| record))
}
std::env::var_os("HOME").map(|h| std::path::PathBuf::from(h).join(".codex").join("auth.json"))
}
fn token_is_expired(expires_at: Option<i64>) -> bool {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|duration| duration.as_secs() as i64)
.unwrap_or_default();
expires_at.is_none_or(|expires_at| now >= expires_at - 60)
fn now() -> i64 {
std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH)
.map(|duration| duration.as_secs() as i64).unwrap_or_default()
}
async fn exchange_code_for_tokens(
code: &str,
code_verifier: &str,
redirect_uri: &str,
) -> Result<(), String> {
) -> Result<AuthRecord, String> {
let body = url::form_urlencoded::Serializer::new(String::new())
.append_pair("grant_type", "authorization_code")
.append_pair("client_id", CHATGPT_CLIENT_ID)
@@ -516,7 +567,34 @@ async fn exchange_code_for_tokens(
.append_pair("redirect_uri", redirect_uri)
.append_pair("code_verifier", code_verifier)
.finish();
let response = reqwest::Client::new()
let tokens = token_request(body).await.map_err(|error| error.message)?;
record_from_tokens(tokens, None, now()).map_err(|error| error.message)
}
async fn refresh_tokens(previous: &AuthRecord) -> Result<AuthRecord, AuthFailure> {
let refresh_token = previous.refresh_token.as_deref().filter(|token| !token.trim().is_empty())
.ok_or_else(|| AuthFailure {
message: SIGN_IN_REQUIRED.to_owned(),
reauthenticate: true,
})?;
let body = url::form_urlencoded::Serializer::new(String::new())
.append_pair("grant_type", "refresh_token")
.append_pair("client_id", CHATGPT_CLIENT_ID)
.append_pair("refresh_token", refresh_token)
.finish();
let tokens = token_request(body).await?;
record_from_tokens(tokens, Some(previous), now())
}
async fn token_request(body: String) -> Result<TokenResponse, AuthFailure> {
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(30))
.redirect(reqwest::redirect::Policy::none())
.build().map_err(|_| AuthFailure {
message: "Could not create the ChatGPT authorization client.".to_owned(),
reauthenticate: false,
})?;
let response = client
.post(CHATGPT_TOKEN_URL)
.header(
reqwest::header::CONTENT_TYPE,
@@ -525,67 +603,83 @@ async fn exchange_code_for_tokens(
.body(body)
.send()
.await
.map_err(|error| format!("ChatGPT token exchange request failed: {error}"))?;
let status = response.status();
if !status.is_success() {
let body = response.text().await.unwrap_or_default();
return Err(format!("ChatGPT token exchange failed: {status} {body}"));
.map_err(|_| AuthFailure {
message: "ChatGPT authorization request failed. Retry the connection.".to_owned(),
reauthenticate: false,
})?;
let status = response.status().as_u16();
let mut bytes = Vec::new();
let mut stream = response.bytes_stream();
while let Some(chunk) = stream.next().await {
let chunk = chunk.map_err(|_| AuthFailure {
message: "Could not read the ChatGPT authorization response.".to_owned(),
reauthenticate: (200..300).contains(&status),
})?;
if bytes.len().saturating_add(chunk.len()) > 64 * 1024 {
return Err(AuthFailure {
message: "ChatGPT authorization returned an oversized response.".to_owned(),
reauthenticate: (200..300).contains(&status),
});
}
bytes.extend_from_slice(&chunk);
}
let tokens: TokenResponse = response
.json()
.await
.map_err(|error| format!("Could not parse the ChatGPT token response: {error}"))?;
parse_token_response(status, &bytes)
}
fn parse_token_response(status: u16, bytes: &[u8]) -> Result<TokenResponse, AuthFailure> {
if (200..300).contains(&status) {
return serde_json::from_slice(bytes).map_err(|_| AuthFailure {
message: "ChatGPT returned invalid credentials. Please sign in again.".to_owned(),
reauthenticate: true,
});
}
let body = serde_json::from_slice::<serde_json::Value>(bytes).unwrap_or_default();
let code = body.get("error").and_then(|error| {
error.as_str().or_else(|| error.get("code").and_then(|code| code.as_str()))
});
let reauthenticate = matches!(status, 401 | 403)
|| (status == 400 && matches!(code,
Some("invalid_grant" | "refresh_token_reused" | "refresh_token_expired" | "refresh_token_invalid")
));
let message = if reauthenticate {
format!("ChatGPT authorization expired or was rejected (HTTP {status}). Please sign in again.")
} else {
format!("ChatGPT authorization failed (HTTP {status}). Retry the connection.")
};
Err(AuthFailure { message, reauthenticate })
}
fn record_from_tokens(
tokens: TokenResponse,
previous: Option<&AuthRecord>,
now: i64,
) -> Result<AuthRecord, AuthFailure> {
let expires_at = extract_expiration_timestamp(&tokens.access_token).or_else(|| {
tokens.expires_in.map(|expires_in| {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|duration| duration.as_secs() as i64)
.unwrap_or_default()
+ expires_in as i64
})
tokens.expires_in.and_then(|expires| i64::try_from(expires).ok())
.map(|expires| now.saturating_add(expires))
});
let account_id = extract_account_id(tokens.id_token.as_deref())
.or_else(|| extract_account_id(Some(&tokens.access_token)));
write_auth_file(&AuthRecord {
.or_else(|| extract_account_id(Some(&tokens.access_token)))
.or_else(|| previous.and_then(|record| record.account_id.clone()));
let record = AuthRecord {
access_token: Some(tokens.access_token),
refresh_token: tokens.refresh_token,
id_token: tokens.id_token,
refresh_token: tokens.refresh_token.or_else(|| previous.and_then(|record| record.refresh_token.clone())),
id_token: tokens.id_token.or_else(|| previous.and_then(|record| record.id_token.clone())),
expires_at,
account_id,
})
};
record.credentials(now).map_err(|_| AuthFailure {
message: "ChatGPT returned invalid or expired credentials. Please sign in again.".to_owned(),
reauthenticate: true,
})?;
Ok(record)
}
fn write_auth_file(record: &AuthRecord) -> Result<(), String> {
let path = auth_file_path().ok_or("Cannot determine auth file path")?;
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)
.map_err(|e| format!("Failed to create auth directory: {e}"))?;
}
let json =
serde_json::to_vec_pretty(record).map_err(|e| format!("Failed to serialize auth: {e}"))?;
std::fs::write(&path, json).map_err(|e| format!("Failed to write auth file: {e}"))?;
Ok(())
}
fn auth_file_path() -> Option<std::path::PathBuf> {
#[cfg(target_os = "windows")]
{
std::env::var_os("APPDATA").map(|d| {
std::path::PathBuf::from(d)
.join("chatgpt")
.join("auth.json")
})
}
#[cfg(not(target_os = "windows"))]
{
std::env::var_os("XDG_CONFIG_HOME")
.map(std::path::PathBuf::from)
.or_else(|| {
std::env::var_os("HOME").map(|h| std::path::PathBuf::from(h).join(".config"))
})
.map(|d| d.join("chatgpt").join("auth.json"))
}
fn save_auth_record(record: &AuthRecord, ctx: &galaxyui::AppContext) -> Result<(), String> {
let json = serde_json::to_string(record)
.map_err(|_| "Could not encode ChatGPT credentials.".to_owned())?;
ctx.secure_storage().write_value_with_owner_only_fallback(AUTH_STORAGE_KEY, &json)
.map_err(|_| "Could not securely save ChatGPT credentials. Please sign in again.".to_owned())
}
fn extract_expiration_timestamp(token: &str) -> Option<i64> {
@@ -633,7 +727,7 @@ struct TokenResponse {
expires_in: Option<u64>,
}
#[derive(serde::Deserialize, serde::Serialize)]
#[derive(Clone, Default, serde::Deserialize, serde::Serialize)]
struct AuthRecord {
access_token: Option<String>,
refresh_token: Option<String>,
@@ -641,3 +735,41 @@ struct AuthRecord {
expires_at: Option<i64>,
account_id: Option<String>,
}
impl std::fmt::Debug for AuthRecord {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ChatGPTAuthRecord")
.field("credentials", &"<redacted>")
.field("expires_at", &self.expires_at)
.finish()
}
}
impl AuthRecord {
fn can_refresh(&self) -> bool {
self.refresh_token.as_deref().is_some_and(|token| !token.trim().is_empty())
}
fn access_context(&self) -> Result<ChatGPTAuthCredentials, String> {
let access_token = self.access_token.as_deref().filter(|token| !token.trim().is_empty())
.ok_or_else(|| SIGN_IN_REQUIRED.to_owned())?;
let account_id = self.account_id.clone().filter(|id| !id.trim().is_empty())
.or_else(|| extract_account_id(self.id_token.as_deref()))
.or_else(|| extract_account_id(Some(access_token)));
Ok(ChatGPTAuthCredentials { access_token: access_token.to_owned(), account_id })
}
fn credentials(&self, now: i64) -> Result<ChatGPTAuthCredentials, String> {
let expiry = self.expires_at.or_else(|| {
self.access_token.as_deref().and_then(extract_expiration_timestamp)
});
if expiry.is_none_or(|expiry| now >= expiry.saturating_sub(60)) {
return Err(SIGN_IN_REQUIRED.to_owned());
}
self.access_context()
}
}
#[cfg(test)]
#[path = "chatgpt_auth_tests.rs"]
mod tests;
+171
View File
@@ -0,0 +1,171 @@
use std::cell::Cell;
use futures::FutureExt;
use super::*;
fn record(access: &str, refresh: &str) -> AuthRecord {
AuthRecord {
access_token: Some(access.to_owned()),
refresh_token: Some(refresh.to_owned()),
id_token: None,
expires_at: Some(10_000),
account_id: Some("test-account".to_owned()),
}
}
#[test]
fn concurrent_requests_share_one_refresh_and_wait_for_persistence() {
let mut coordinator = RefreshCoordinator::default();
let (first_tx, mut first_rx) = oneshot::channel();
let (second_tx, mut second_rx) = oneshot::channel();
let generation = coordinator.begin(first_tx).unwrap();
assert_eq!(coordinator.begin(second_tx), None);
let saved = Cell::new(false);
let result = coordinator.complete(generation, Ok(record("rotated-access", "rotated-refresh")), |record| {
assert_eq!(first_rx.try_recv().unwrap(), None);
assert_eq!(second_rx.try_recv().unwrap(), None);
assert_eq!(record.refresh_token.as_deref(), Some("rotated-refresh"));
saved.set(true);
Ok(())
});
assert!(result.unwrap().is_ok());
assert!(saved.get());
let first = first_rx.now_or_never().unwrap().unwrap().unwrap();
let second = second_rx.now_or_never().unwrap().unwrap().unwrap();
assert_eq!(first, second);
assert_eq!(first.access_token, "rotated-access");
}
#[test]
fn failed_persistence_never_publishes_rotated_credentials() {
let mut coordinator = RefreshCoordinator::default();
let (tx, rx) = oneshot::channel();
let generation = coordinator.begin(tx).unwrap();
let result = coordinator.complete(generation, Ok(record("rotated", "new-refresh")), |_| {
Err("storage unavailable".to_owned())
}).unwrap().unwrap_err();
assert!(result.reauthenticate);
assert!(rx.now_or_never().unwrap().unwrap().is_err());
}
#[test]
fn replaced_sign_in_rejects_late_refresh_without_overwriting_new_session() {
let mut coordinator = RefreshCoordinator::default();
let (old_tx, old_rx) = oneshot::channel();
let old = coordinator.begin(old_tx).unwrap();
coordinator.invalidate();
assert!(old_rx.now_or_never().unwrap().unwrap().is_err());
let (new_tx, new_rx) = oneshot::channel();
let current = coordinator.begin(new_tx).unwrap();
assert_ne!(old, current);
assert!(coordinator.complete(old, Ok(record("stale", "stale-refresh")), |_| {
panic!("A stale refresh must not write credentials")
}).is_none());
assert!(coordinator.complete(current, Ok(record("current", "current-refresh")), |_| Ok(()))
.unwrap().is_ok());
assert_eq!(new_rx.now_or_never().unwrap().unwrap().unwrap().access_token, "current");
}
#[test]
fn reconnect_starts_a_new_generation_even_with_a_cached_refresh_token() {
let mut model = ChatGPTAuthModel::new_for_test();
model.record = Some(record("old-access", "already-used-refresh"));
model.state = ChatGPTAuthState::Connected;
let first = model.begin_sign_in().unwrap();
model.state = ChatGPTAuthState::AwaitingBrowser;
assert_eq!(model.begin_sign_in(), None);
model.state = ChatGPTAuthState::Failed("refresh_token_reused".to_owned());
let replacement = model.begin_sign_in().unwrap();
assert!(replacement > first);
}
#[test]
fn rotated_tokens_replace_previous_tokens_and_omitted_refresh_tokens_are_preserved() {
let previous = record("old-access", "old-refresh");
for refresh in [None, Some("new-refresh".to_owned())] {
let result = record_from_tokens(TokenResponse {
access_token: "new-access".to_owned(),
refresh_token: refresh.clone(),
id_token: None,
expires_in: Some(3600),
}, Some(&previous), 100).unwrap();
assert_eq!(result.access_token.as_deref(), Some("new-access"));
assert_eq!(result.refresh_token, refresh.or_else(|| Some("old-refresh".to_owned())));
assert_eq!(result.expires_at, Some(3700));
assert_eq!(result.account_id.as_deref(), Some("test-account"));
}
}
#[test]
fn rejected_refresh_tokens_require_reauthentication_without_echoing_tokens() {
for code in ["refresh_token_reused", "invalid_grant", "refresh_token_expired"] {
let body = serde_json::to_vec(&serde_json::json!({
"error": {"code": code, "message": "sensitive-refresh-token"}
})).unwrap();
let error = parse_token_response(401, &body).err().unwrap();
assert!(error.reauthenticate);
assert!(error.message.contains("sign in"));
assert!(!error.message.contains("sensitive-refresh-token"));
}
let error = parse_token_response(400, br#"{"error":"invalid_grant"}"#).err().unwrap();
assert!(error.reauthenticate);
}
#[test]
fn transient_refresh_failure_preserves_the_session_and_does_not_save_rejected_data() {
let error = parse_token_response(503, br#"{"message":"sensitive-data"}"#).err().unwrap();
assert!(!error.reauthenticate);
assert!(!error.message.contains("sensitive-data"));
let mut coordinator = RefreshCoordinator::default();
let (tx, rx) = oneshot::channel();
let generation = coordinator.begin(tx).unwrap();
let result = coordinator.complete(generation, Err(error), |_| {
panic!("Transient failures must not overwrite stored credentials")
}).unwrap().unwrap_err();
assert!(!result.reauthenticate);
assert!(rx.now_or_never().unwrap().unwrap().is_err());
}
#[test]
fn credential_debug_output_is_redacted_and_expiry_has_a_refresh_margin() {
let credentials = ChatGPTAuthCredentials {
access_token: "sensitive-access-token".to_owned(),
account_id: Some("account".to_owned()),
};
assert!(!format!("{credentials:?}").contains("sensitive-access-token"));
let credentials = record("sensitive-access-token", "sensitive-refresh-token");
assert!(!format!("{credentials:?}").contains("sensitive-access-token"));
assert!(!format!("{credentials:?}").contains("sensitive-refresh-token"));
assert!(credentials.credentials(100).is_ok());
assert!(credentials.credentials(9940).is_err());
assert!(AuthRecord::default().credentials(100).is_err());
}
#[test]
fn invalid_token_payloads_are_not_saved_or_exposed() {
assert!(parse_token_response(200, b"invalid-sensitive-data").is_err());
assert!(record_from_tokens(TokenResponse {
access_token: String::new(),
refresh_token: Some("refresh".to_owned()),
id_token: None,
expires_in: Some(3600),
}, None, 100).is_err());
assert!(record_from_tokens(TokenResponse {
access_token: "opaque-token-without-expiry".to_owned(),
refresh_token: None,
id_token: None,
expires_in: None,
}, None, 100).is_err());
}
#[test]
fn browser_authorization_keeps_pkce_and_state_validation() {
let url = build_authorize_url("http://localhost:1455/auth/callback", "challenge", "state").unwrap();
let params = url.query_pairs().collect::<std::collections::HashMap<_, _>>();
assert_eq!(params.get("code_challenge_method").map(|v| v.as_ref()), Some("S256"));
assert_eq!(params.get("state").map(|v| v.as_ref()), Some("state"));
let callback = parse_callback_request_line("GET /auth/callback?code=code&state=state HTTP/1.1").unwrap();
assert_eq!(callback.state, "state");
assert!(parse_callback_request_line("GET /auth/callback?code=code HTTP/1.1").is_err());
}
+8 -23
View File
@@ -181,28 +181,14 @@ 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_model_route(model_id, settings.bedrock_models.value())
.with_managed_credentials(
&api_key_manager.aws_credentials_refresh_strategy(),
api_key_manager.aws_credentials_state(),
);
return ProviderConfig::Bedrock(config);
}
ProviderConfig::None
@@ -229,9 +215,8 @@ impl CrosscheckReviewer {
model: config.model.unwrap_or(model_id),
reasoning_effort: config.reasoning_effort,
max_output_tokens: Some(u64::from(REVIEWER_MAX_OUTPUT_TOKENS)),
auth_file: None,
},
);
).with_credential_resolver(crate::ai::chatgpt_auth::credential_resolver());
runtime
.complete_text(
prompt::CROSSCHECK_REVIEWER_SYSTEM_PROMPT.to_string(),
+42 -53
View File
@@ -27,6 +27,8 @@ use warp_multi_agent_api as api;
use super::custom_model_routers::{self, CustomModelRouter, ModelConfigError};
use super::execution_profiles::profiles::AIExecutionProfilesModel;
use crate::ai::acp::{acp_launch_fingerprint, acp_provider_selection_identity};
#[cfg(not(target_family = "wasm"))]
use crate::ai::provider::external_config::ExternalBedrockConfig;
use crate::auth::auth_manager::{AuthManager, AuthManagerEvent};
use crate::auth::AuthStateProvider;
use crate::network::{NetworkStatus, NetworkStatusEvent, NetworkStatusKind};
@@ -616,6 +618,7 @@ pub struct LLMPreferences {
#[cfg(not(target_family = "wasm"))]
chatgpt_subscription_models_refresh_in_flight: bool,
bedrock_models_refresh_in_flight: bool,
bedrock_verification_generation: u64,
#[cfg(not(target_family = "wasm"))]
bedrock_verified_config: Option<super::provider::client::BedrockClientConfig>,
#[cfg(not(target_family = "wasm"))]
@@ -796,6 +799,7 @@ impl LLMPreferences {
#[cfg(not(target_family = "wasm"))]
chatgpt_subscription_models_refresh_in_flight: false,
bedrock_models_refresh_in_flight: false,
bedrock_verification_generation: 0,
#[cfg(not(target_family = "wasm"))]
bedrock_verified_config: None,
#[cfg(not(target_family = "wasm"))]
@@ -844,7 +848,9 @@ impl LLMPreferences {
return;
}
self.bedrock_models_refresh_in_flight = true;
let generation = self.begin_bedrock_verification();
let config = bedrock_discovery_config(settings);
let revision = super::provider::authentication::connection_revision(ctx);
let requested_config = config.clone();
self.bedrock_verified_config = None;
self.bedrock_verified_model_ids.clear();
@@ -860,6 +866,11 @@ impl LLMPreferences {
discovery.models
});
me.bedrock_models_refresh_in_flight = false;
if me.bedrock_verification_generation != generation
|| super::provider::authentication::connection_revision(ctx) != revision
{
return;
}
if !*AISettings::as_ref(ctx).bedrock_enabled.value() {
return;
}
@@ -888,26 +899,37 @@ impl LLMPreferences {
ctx.emit(LLMPreferencesEvent::UpdatedAvailableLLMs);
}
Err(error) => {
log::debug!("[bedrock] Startup model discovery unavailable: {error}");
log::warn!("[bedrock] Startup model discovery unavailable: {error}");
}
}
},
);
}
#[cfg(not(target_family = "wasm"))]
pub(crate) fn begin_bedrock_verification(&mut self) -> u64 {
self.bedrock_verification_generation += 1;
self.bedrock_verification_generation
}
#[cfg(not(target_family = "wasm"))]
pub(crate) fn record_bedrock_verification(
&mut self,
generation: u64,
config: super::provider::client::BedrockClientConfig,
models: &[BedrockModelConfig],
ctx: &mut ModelContext<Self>,
) {
) -> bool {
if self.bedrock_verification_generation != generation {
return false;
}
self.bedrock_verified_model_ids =
models.iter().map(|model| model.model_id.clone()).collect();
self.bedrock_verified_config = Some(config);
self.inject_bedrock_models(ctx);
self.ensure_default_model_present();
ctx.emit(LLMPreferencesEvent::UpdatedAvailableLLMs);
true
}
#[cfg(not(target_family = "wasm"))]
@@ -945,48 +967,14 @@ impl LLMPreferences {
let region = settings.bedrock_region.value().clone();
let cross_region = *settings.bedrock_cross_region_inference.value();
// Check if user wants only 1-hour cache models
use crate::ai::provider::external_config::ExternalBedrockConfig;
let external_config = ExternalBedrockConfig::load();
let require_1h_cache = external_config.enable_prompt_caching_1h;
let mut effective = discovered_models;
// Filter out models that don't support 1-hour caching if required
if require_1h_cache {
effective.retain(|model| {
// 1-hour caching is supported by Claude 4.5+ models
// Opus 4.5+, Sonnet 4.5+, Haiku 4.5+
let supports_1h = model.model_id.contains("4-5")
|| model.model_id.contains("4.5")
|| model.model_id.contains("-4-6") // Opus/Sonnet 4.6+ also support 1h
|| model.model_id.contains("4.6")
|| model.model_id.contains("-4-7")
|| model.model_id.contains("4.7")
|| model.model_id.contains("-4-8")
|| model.model_id.contains("4.8");
if !supports_1h {
log::info!(
"[bedrock] Filtering out model {} - does not support 1-hour cache (ENABLE_PROMPT_CACHING_1H=1)",
model.model_id
);
}
supports_1h
});
if effective.is_empty() {
log::warn!("[bedrock] No models left after filtering for 1-hour cache support!");
}
}
let effective = effective;
for model in effective {
for model in discovered_models {
if !model.enabled || !self.bedrock_model_is_verified(&model.model_id, ctx) {
continue;
}
let metadata = model_metadata(&model.model_id);
if metadata.is_some_and(|metadata| metadata.supports_agent_runtime() == Some(false)) {
if model.runtime_supported == Some(false)
|| metadata.is_some_and(|metadata| metadata.supports_agent_runtime() == Some(false))
{
log::debug!(
"[bedrock] Excluding {} from Agent Mode because its catalog metadata marks it as incompatible with streaming Converse",
model.model_id
@@ -1046,7 +1034,6 @@ impl LLMPreferences {
// Default agent mode to ANTHROPIC_MODEL from external config if set,
// otherwise Claude Opus 4.6, falling back to the first available model.
// Note: external_config already loaded above for filtering
let external_config_for_default = ExternalBedrockConfig::load();
if let Some(id) = external_config_for_default
@@ -1082,7 +1069,11 @@ impl LLMPreferences {
.map(|m| m.id.clone())
})
{
self.models_by_feature.agent_mode.default_id = id;
if !self.models_by_feature.agent_mode.choices.iter().any(|model| {
model.id == self.models_by_feature.agent_mode.default_id
}) {
self.models_by_feature.agent_mode.default_id = id;
}
}
// Default coding to Claude Sonnet 4.6, falling back to the first available model.
@@ -1095,12 +1086,18 @@ impl LLMPreferences {
.or_else(|| self.models_by_feature.coding.choices.first())
.map(|m| m.id.clone())
{
self.models_by_feature.coding.default_id = id;
if !self.models_by_feature.coding.choices.iter().any(|model| {
model.id == self.models_by_feature.coding.default_id
}) {
self.models_by_feature.coding.default_id = id;
}
}
// Default CLI agent to the first available model.
if let Some(ref mut cli) = self.models_by_feature.cli_agent {
if let Some(id) = cli.choices.first().map(|m| m.id.clone()) {
if let Some(id) = cli.choices.first().filter(|_| {
!cli.choices.iter().any(|model| model.id == cli.default_id)
}).map(|m| m.id.clone()) {
cli.default_id = id;
}
}
@@ -3471,13 +3468,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)
}
+48
View File
@@ -159,6 +159,51 @@ fn openai_model(model_id: &str) -> OpenAIModelConfig {
}
}
#[test]
fn live_bedrock_models_are_not_hidden_by_missing_catalog_metadata() {
App::test((), |mut app| async move {
initialize_settings_for_tests(&mut app);
let id = "future.new-live-model";
AISettings::handle(&app).update(&mut app, |settings, ctx| {
settings.bedrock_enabled.set_value(true, ctx).unwrap();
settings.bedrock_models.set_value(
vec![bedrock_model("future.other-model"), bedrock_model(id)], ctx,
).unwrap();
});
let mut preferences = empty_preferences();
preferences.models_by_feature.agent_mode.default_id = id.into();
preferences.models_by_feature.coding.default_id = id.into();
app.read(|ctx| {
preferences.bedrock_verified_config =
Some(super::bedrock_discovery_config(AISettings::as_ref(ctx)));
preferences.bedrock_verified_model_ids.insert(id.to_owned());
preferences.bedrock_verified_model_ids.insert("future.other-model".to_owned());
preferences.inject_bedrock_models(ctx);
});
assert!(preferences.models_by_feature.agent_mode.choices.iter().any(|model| model.id.as_str() == id));
assert_eq!(preferences.models_by_feature.agent_mode.default_id.as_str(), id);
assert_eq!(preferences.models_by_feature.coding.default_id.as_str(), id);
});
}
#[test]
fn newer_bedrock_discovery_invalidates_older_results() {
App::test((), |mut app| async move {
initialize_settings_for_tests(&mut app);
let preferences = app.add_model(|_| empty_preferences());
let models = [bedrock_model("future.new-live-model")];
preferences.update(&mut app, |preferences, ctx| {
let first = preferences.begin_bedrock_verification();
let latest = preferences.begin_bedrock_verification();
let config = super::bedrock_discovery_config(AISettings::as_ref(ctx));
assert!(!preferences.record_bedrock_verification(first, config.clone(), &models, ctx));
assert!(preferences.bedrock_verified_model_ids.is_empty());
assert!(preferences.record_bedrock_verification(latest, config, &models, ctx));
assert!(preferences.bedrock_verified_model_ids.contains("future.new-live-model"));
});
});
}
fn bedrock_model(model_id: &str) -> BedrockModelConfig {
BedrockModelConfig {
enabled: true,
@@ -166,6 +211,7 @@ fn bedrock_model(model_id: &str) -> BedrockModelConfig {
model_id: model_id.to_string(),
display_name: model_id.to_string(),
vision_supported: false,
..Default::default()
}
}
@@ -235,6 +281,7 @@ fn empty_preferences() -> LLMPreferences {
catalog_refresh_generations: HashMap::new(),
chatgpt_subscription_models_refresh_in_flight: false,
bedrock_models_refresh_in_flight: false,
bedrock_verification_generation: 0,
bedrock_verified_config: None,
bedrock_verified_model_ids: HashSet::new(),
acp_selections: HashMap::new(),
@@ -451,6 +498,7 @@ fn acp_models_are_injected_only_while_acp_is_enabled() {
catalog_refresh_generations: HashMap::new(),
chatgpt_subscription_models_refresh_in_flight: false,
bedrock_models_refresh_in_flight: false,
bedrock_verification_generation: 0,
bedrock_verified_config: None,
bedrock_verified_model_ids: HashSet::new(),
acp_selections: HashMap::new(),
+3 -1
View File
@@ -85,10 +85,12 @@ 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.add_singleton_model(|_| chatgpt_auth::ChatGPTAuthModel::new());
app.add_singleton_model(chatgpt_auth::ChatGPTAuthModel::new);
#[cfg(not(target_family = "wasm"))]
app.add_singleton_model(claude_oauth::ClaudeOAuthModel::new);
blocklist::keyboard_navigable_buttons::init(app);
+333
View File
@@ -0,0 +1,333 @@
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,
invocation_model_id: None,
cross_region_inference: *settings.bedrock_cross_region_inference.value(),
}
}
pub(crate) fn with_model_route(
mut self,
model_id: &str,
models: &[crate::settings::BedrockModelConfig],
) -> Self {
self.invocation_model_id = models.iter().find(|model| {
model.model_id == model_id
|| model.invocation_model_id == model_id
|| super::models::foundation_model_id(&model.model_id)
== super::models::foundation_model_id(model_id)
}).and_then(|model| {
(!model.invocation_model_id.is_empty()).then(|| model.invocation_model_id.clone())
});
self
}
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;
+294
View File
@@ -0,0 +1,294 @@
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,
invocation_model_id: 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));
}
#[test]
fn model_routes_use_discovered_ids_and_match_legacy_cross_region_selections() {
let model = crate::settings::BedrockModelConfig {
model_id: "future.new-model".to_owned(),
invocation_model_id: "arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.future.new-model".to_owned(),
..Default::default()
};
for selected in ["future.new-model", "us.future.new-model"] {
let resolved = config(BedrockAuthMethod::Profile)
.with_model_route(selected, &[model.clone()]);
assert_eq!(resolved.invocation_model_id.as_deref(), Some(model.invocation_model_id.as_str()));
}
let resolved = config(BedrockAuthMethod::Profile).with_model_route("unrelated.model", &[model]);
assert_eq!(resolved.invocation_model_id, None);
}
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()));
}
}
+26 -31
View File
@@ -1,4 +1,5 @@
use std::sync::{Arc, Mutex};
use std::time::SystemTime;
use anyhow::Result;
use aws_config::BehaviorVersion;
@@ -27,6 +28,7 @@ fn strip_context_marker(model_id: &str) -> String {
pub struct BedrockClient {
runtime_client: BedrockRuntimeClient,
region: String,
invocation_model_id: Option<String>,
}
#[derive(Debug, Clone, PartialEq)]
@@ -37,40 +39,18 @@ pub struct BedrockClientConfig {
pub access_key_id: String,
pub secret_access_key: String,
pub session_token: Option<String>,
pub credentials_expiry: Option<SystemTime>,
pub invocation_model_id: Option<String>,
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 +95,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 +136,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
@@ -165,6 +157,7 @@ impl BedrockClient {
Ok(Self {
runtime_client,
region,
invocation_model_id: config.invocation_model_id,
})
}
@@ -195,14 +188,16 @@ impl BedrockClient {
tool_result_archive: Vec<ConversationMessage>,
) -> Result<LegacyResponseStream, BedrockError> {
let base_model_id = strip_context_marker(model_id);
let effective_model_id = if cross_region_inference {
let effective_model_id = if let Some(model) = &self.invocation_model_id {
model.clone()
} else if cross_region_inference {
apply_cross_region_prefix(&base_model_id, &self.region)
} else {
base_model_id
};
let external_config = ExternalBedrockConfig::load();
let caching_config = CachingConfig::from_external_config(&external_config);
let caching_config = CachingConfig::from_external_config(&external_config).for_model(model_id);
if !caching_config.enabled {
log::info!("[bedrock] Prompt caching disabled (DISABLE_PROMPT_CACHING=1)");
+16
View File
@@ -38,6 +38,22 @@ impl CachingConfig {
extended_ttl_requested: external.enable_prompt_caching_1h,
}
}
pub fn for_model(mut self, model_id: &str) -> Self {
let model = galaxy_bedrock_model_catalog::model_metadata(model_id);
let known_claude_4 = model.is_some_and(|metadata| {
metadata.provider_name == "Anthropic"
&& metadata.supports_agent_runtime() == Some(true)
&& (metadata.model_id.starts_with("anthropic.claude-opus-4")
|| metadata.model_id.starts_with("anthropic.claude-sonnet-4")
|| metadata.model_id.starts_with("anthropic.claude-haiku-4"))
});
self.enabled &= known_claude_4;
self.extended_ttl_requested &= self.enabled && model.is_some_and(|metadata| {
["4-5", "4-6", "4-7", "4-8"].iter().any(|version| metadata.model_id.contains(version))
});
self
}
}
pub struct ConvertedRequest {
+12
View File
@@ -1,3 +1,15 @@
#[test]
fn cache_preferences_do_not_enable_unsupported_or_unknown_model_caching() {
let external = super::external_config::ExternalBedrockConfig {
enable_prompt_caching_1h: true,
..Default::default()
};
for id in ["future.new-model-v100", "openai.gpt-5.6-sol"] {
let config = super::convert::CachingConfig::from_external_config(&external).for_model(id);
assert!(!config.enabled);
assert!(!config.extended_ttl_requested);
}
}
use aws_sdk_bedrockruntime::types::{ContentBlock, ConversationRole, SystemContentBlock, Tool};
use serde_json::json;
+169 -26
View File
@@ -1,17 +1,25 @@
//! AWS Bedrock control-plane discovery.
//!
//! The foundation-model catalog is only a candidate list. Every candidate is
//! checked with `GetFoundationModelAvailability` before it is offered to the
//! user or persisted in Galaxy settings.
//! AWS supplies the model inventory and invocation routes. Bundled metadata
//! enriches known capabilities but never determines inventory membership.
use std::collections::{BTreeMap, HashSet};
use std::future::Future;
use aws_config::BehaviorVersion;
use aws_sdk_bedrock::operation::list_inference_profiles::ListInferenceProfilesOutput;
use aws_sdk_bedrock::types::{
FoundationModelSummary, InferenceProfileStatus, InferenceProfileSummary,
InferenceProfileType,
};
use aws_sdk_bedrock::Client;
use aws_sdk_bedrockruntime::config::Region;
use futures::{stream, StreamExt};
use galaxy_bedrock_model_catalog::models as catalog_models;
use galaxy_bedrock_model_catalog::model_metadata;
use serde_json::json;
use super::client::{BedrockClientConfig, BedrockError};
use super::models::foundation_model_id;
use crate::ai::remote_logging::{RemoteLogLevel, RemoteLogRecord};
use crate::settings::ai::BedrockModelConfig;
@@ -29,9 +37,66 @@ pub(crate) async fn discover_available_models(
.await
.map_err(|error| error.to_string())?;
let client = Client::new(&aws_config);
// The bundled catalog defines candidates; AWS defines account access.
// Never add arbitrary entries from ListFoundationModels to this catalog.
let candidates = catalog_models().to_vec();
let inventory = client.list_foundation_models().send().await
.map_err(|error| format!("Unable to list Bedrock models: {}", sanitized_error(&error)))?;
let mut foundation_models = inventory.model_summaries.unwrap_or_default();
let profiles = if config.cross_region_inference {
collect_profile_pages(|token| {
let client = client.clone();
async move {
client.list_inference_profiles()
.set_next_token(token)
.send().await
.map_err(|error| format!("Unable to list Bedrock inference profiles: {}", sanitized_error(&error)))
}
}).await?
} else {
Vec::new()
};
let mut known_ids: HashSet<String> = foundation_models.iter()
.map(|model| model.model_id().to_owned()).collect();
// A profile can introduce a model before it appears in the regional listing.
for profile in profiles.iter().filter(|profile| profile.status() == &InferenceProfileStatus::Active) {
for model in profile.models() {
let Some(arn) = model.model_arn() else { continue };
if !arn.contains(":foundation-model/") {
continue;
}
let id = foundation_model_id(arn);
if !known_ids.insert(id.to_owned()) {
continue;
}
let details = match client.get_foundation_model().model_identifier(id).send().await {
Ok(output) => match output.model_details {
Some(details) => details,
None => {
log::warn!("[bedrock] AWS returned no metadata for profile model {id}");
continue;
}
},
Err(error) => {
let error = sanitized_error(&error);
log::warn!("[bedrock] Unable to describe profile model {id}: {error}");
continue;
}
};
let summary = FoundationModelSummary::builder()
.model_id(details.model_id)
.model_arn(details.model_arn)
.set_model_name(details.model_name)
.set_provider_name(details.provider_name)
.set_input_modalities(details.input_modalities)
.set_output_modalities(details.output_modalities)
.set_response_streaming_supported(details.response_streaming_supported)
.set_inference_types_supported(details.inference_types_supported)
.set_model_lifecycle(details.model_lifecycle)
.build().map_err(|error| format!("Invalid Bedrock model metadata: {error}"))?;
foundation_models.push(summary);
}
}
let inventory_count = foundation_models.len();
let candidates = models_from_inventory(foundation_models, &profiles, config.cross_region_inference);
let candidate_count = candidates.len();
let sweep_id = uuid::Uuid::new_v4().to_string();
let region = aws_config
.region()
@@ -39,30 +104,30 @@ pub(crate) async fn discover_available_models(
.unwrap_or("UNKNOWN")
.to_owned();
log::info!(
"[bedrock] Starting availability sweep {sweep_id}: {} catalog models, region={region}",
"[bedrock] Starting availability sweep {sweep_id}: {} live candidates, region={region}",
candidates.len()
);
let checks = candidates.into_iter().map(|metadata| {
let checks = candidates.into_iter().map(|mut model| {
let client = client.clone();
let sweep_id = sweep_id.clone();
let region = region.clone();
async move {
let model_id = metadata.model_id.as_str();
let model_id = model.model_id.clone();
let availability = match client
.get_foundation_model_availability()
.model_id(model_id)
.model_id(&model_id)
.send()
.await
{
Ok(availability) => availability,
Err(error) => {
let error = crate::ai::remote_logging::sanitize_error(format!("{error:?}"));
let error = sanitized_error(&error);
log::warn!("[bedrock] Availability sweep {sweep_id}: model={model_id}, region={region}, decision=excluded, error={error}");
return (None, RemoteLogRecord {
level: RemoteLogLevel::Warn,
message: "Bedrock model availability check failed".to_owned(),
context: json!({"event": "bedrock_model_availability", "sweep_id": sweep_id, "model_id": model_id, "aws_region": region, "decision": "excluded", "error": error}),
});
}, false);
}
};
let agreement = availability.agreement_availability().map(|value| value.status().as_str());
@@ -70,7 +135,7 @@ pub(crate) async fn discover_available_models(
let entitlement = availability.entitlement_availability().as_str();
let region_status = availability.region_availability().as_str();
let access_allowed = model_availability_is_usable(agreement, authorization, entitlement, region_status);
let runtime_supported = metadata.supports_agent_runtime() != Some(false);
let runtime_supported = model.runtime_supported != Some(false);
let allowed = access_allowed && runtime_supported;
let decision = if allowed { "allowed" } else { "excluded" };
log::info!("[bedrock] Availability sweep {sweep_id}: model={model_id}, agreement={}, authorization={authorization}, entitlement={entitlement}, region_status={region_status}, runtime_supported={runtime_supported}, decision={decision}", agreement.unwrap_or("MISSING"));
@@ -79,23 +144,20 @@ pub(crate) async fn discover_available_models(
message: "Bedrock model availability checked".to_owned(),
context: json!({"event": "bedrock_model_availability", "sweep_id": sweep_id, "model_id": model_id, "aws_region": region, "agreement": agreement, "authorization": authorization, "entitlement": entitlement, "region_availability": region_status, "agreement_error": availability.agreement_availability().and_then(|value| value.error_message()).map(crate::ai::remote_logging::sanitize_error), "access_allowed": access_allowed, "runtime_supported": runtime_supported, "decision": decision}),
};
if !allowed { return (None, record); }
(Some(BedrockModelConfig {
model_id: model_id.to_owned(),
display_name: metadata.model_name.clone(),
vision_supported: metadata.supports_vision_input(),
enabled: true,
access_verified: true,
}), record)
model.access_verified = allowed;
(allowed.then_some(model), record, true)
}
});
let results = stream::iter(checks)
.buffer_unordered(AVAILABILITY_CHECK_CONCURRENCY)
.collect::<Vec<_>>()
.await;
if candidate_count > 0 && !results.iter().any(|(_, _, checked)| *checked) {
return Err("All Bedrock availability checks failed; saved model preferences are unchanged".to_owned());
}
let (models, mut records): (Vec<_>, Vec<_>) = results.into_iter().unzip();
let (models, mut records): (Vec<_>, Vec<_>) = results.into_iter()
.map(|(model, record, _)| (model, record)).unzip();
let mut models = models.into_iter().flatten().collect::<Vec<_>>();
models.sort_by(|left, right| left.display_name.cmp(&right.display_name));
log::info!(
@@ -106,11 +168,84 @@ pub(crate) async fn discover_available_models(
records.push(RemoteLogRecord {
level: RemoteLogLevel::Info,
message: "Bedrock availability sweep completed".to_owned(),
context: json!({"event": "bedrock_availability_sweep_completed", "sweep_id": sweep_id, "aws_region": region, "catalog_count": catalog_models().len(), "checked_count": records.len(), "allowed_count": models.len()}),
context: json!({"event": "bedrock_availability_sweep_completed", "sweep_id": sweep_id, "aws_region": region, "inventory_count": inventory_count, "candidate_count": candidate_count, "checked_count": records.len(), "allowed_count": models.len()}),
});
Ok(BedrockDiscovery { models, records })
}
fn sanitized_error(error: &impl std::fmt::Debug) -> String {
crate::ai::remote_logging::sanitize_error(format!("{error:?}"))
}
async fn collect_profile_pages<F, Fut>(mut fetch: F) -> Result<Vec<InferenceProfileSummary>, String>
where
F: FnMut(Option<String>) -> Fut,
Fut: Future<Output = Result<ListInferenceProfilesOutput, String>>,
{
let mut profiles = Vec::new();
let mut token = None;
let mut seen = HashSet::new();
loop {
let page = fetch(token).await?;
profiles.extend(page.inference_profile_summaries.unwrap_or_default());
token = page.next_token.filter(|token| !token.is_empty());
let Some(next) = &token else { return Ok(profiles) };
if !seen.insert(next.clone()) {
return Err("AWS repeated an inference-profile pagination token".to_owned());
}
}
}
fn models_from_inventory(
foundation_models: Vec<FoundationModelSummary>,
profiles: &[InferenceProfileSummary],
cross_region: bool,
) -> Vec<BedrockModelConfig> {
let mut routes = profiles.iter()
.filter(|profile| profile.status() == &InferenceProfileStatus::Active)
.collect::<Vec<_>>();
routes.sort_by_key(|profile| (
profile.r#type() != &InferenceProfileType::SystemDefined,
profile.inference_profile_id().starts_with("global."),
profile.inference_profile_id().to_owned(),
));
let mut models = BTreeMap::new();
for summary in foundation_models {
let id = summary.model_id();
let metadata = model_metadata(id);
let text_output = summary.output_modalities().iter().any(|modality| modality.as_str() == "TEXT");
let incompatible = (!summary.output_modalities().is_empty() && !text_output)
|| summary.response_streaming_supported() == Some(false)
|| metadata.is_some_and(|metadata| metadata.supports_agent_runtime() == Some(false));
if incompatible {
continue;
}
let profile = cross_region.then(|| routes.iter().find(|profile| {
!profile.models().is_empty() && profile.models().iter().all(|model| {
model.model_arn().is_some_and(|arn| foundation_model_id(arn) == id)
})
})).flatten();
let invocation = if let Some(profile) = profile {
profile.inference_profile_arn().to_owned()
} else if summary.inference_types_supported().iter().any(|kind| kind.as_str() == "ON_DEMAND") {
id.to_owned()
} else {
continue;
};
models.entry(id.to_owned()).or_insert_with(|| BedrockModelConfig {
model_id: id.to_owned(),
invocation_model_id: invocation,
display_name: summary.model_name().unwrap_or(id).to_owned(),
provider_name: summary.provider_name().unwrap_or_default().to_owned(),
vision_supported: summary.input_modalities().iter().any(|modality| modality.as_str() == "IMAGE"),
response_streaming_supported: summary.response_streaming_supported(),
runtime_supported: (text_output && summary.response_streaming_supported() == Some(true)).then_some(true),
..Default::default()
});
}
models.into_values().collect()
}
fn model_availability_is_usable(
agreement_status: Option<&str>,
authorization_status: &str,
@@ -133,6 +268,14 @@ pub(crate) fn preserve_model_toggles(
.find(|previous| previous.model_id == model.model_id)
{
model.enabled = previous.enabled;
model.display_name.clone_from(&previous.display_name);
}
}
for previous in existing {
if !discovered.iter().any(|model| model.model_id == previous.model_id) {
let mut preserved = previous.clone();
preserved.access_verified = false;
discovered.push(preserved);
}
}
discovered
@@ -161,7 +304,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 =
+140
View File
@@ -1,4 +1,143 @@
use super::model_availability_is_usable;
use aws_sdk_bedrock::operation::list_inference_profiles::ListInferenceProfilesOutput;
use aws_sdk_bedrock::types::{
FoundationModelSummary, InferenceProfileModel, InferenceProfileStatus,
InferenceProfileSummary, InferenceProfileType, InferenceType, ModelModality,
};
fn foundation(id: &str, on_demand: bool) -> FoundationModelSummary {
FoundationModelSummary::builder()
.model_id(id)
.model_arn(format!("arn:aws:bedrock:us-east-1::foundation-model/{id}"))
.model_name("New live model")
.provider_name("New provider")
.input_modalities(ModelModality::Text)
.output_modalities(ModelModality::Text)
.response_streaming_supported(true)
.set_inference_types_supported(on_demand.then(|| vec![InferenceType::OnDemand]))
.build()
.unwrap()
}
fn inference_profile(id: &str, model: &str) -> InferenceProfileSummary {
InferenceProfileSummary::builder()
.inference_profile_id(id)
.inference_profile_name(id)
.inference_profile_arn(format!("arn:aws:bedrock:us-east-1:123456789012:inference-profile/{id}"))
.status(InferenceProfileStatus::Active)
.r#type(InferenceProfileType::SystemDefined)
.models(InferenceProfileModel::builder()
.model_arn(format!("arn:aws:bedrock:us-east-1::foundation-model/{model}"))
.build())
.build()
.unwrap()
}
#[test]
fn new_live_model_does_not_need_a_bundled_catalog_entry() {
let id = "future.new-model-v100";
assert!(galaxy_bedrock_model_catalog::model_metadata(id).is_none());
let profile = inference_profile("us.future.new-model-v100", id);
let models = super::models_from_inventory(
vec![foundation(id, false)], &[profile.clone()], true,
);
assert_eq!(models.len(), 1);
assert_eq!(models[0].model_id, id);
assert_eq!(models[0].invocation_model_id, profile.inference_profile_arn());
assert_eq!(models[0].response_streaming_supported, Some(true));
assert_eq!(models[0].provider_name, "New provider");
}
#[test]
fn cross_region_disabled_requires_an_actual_on_demand_route() {
let id = "future.profile-only";
let profile = inference_profile("us.future.profile-only", id);
assert!(super::models_from_inventory(
vec![foundation(id, false)], &[profile.clone()], false,
).is_empty());
let models = super::models_from_inventory(
vec![foundation(id, true)], &[profile], false,
);
assert_eq!(models.len(), 1);
assert_eq!(models[0].invocation_model_id, id);
}
#[test]
fn inventory_deduplicates_routes_and_rejects_known_incompatible_models() {
let id = "future.new-model-v100";
let profile = inference_profile("us.future.new-model-v100", id);
let models = super::models_from_inventory(
vec![foundation(id, false), foundation(id, false)],
&[profile.clone(), profile],
true,
);
assert_eq!(models.len(), 1);
assert!(super::models_from_inventory(
vec![foundation("openai.gpt-5.5", true)], &[], true,
).is_empty());
let mut not_streaming = foundation(id, true);
not_streaming.response_streaming_supported = Some(false);
assert!(super::models_from_inventory(vec![not_streaming], &[], true).is_empty());
let mut image_model = foundation(id, true);
image_model.output_modalities = Some(vec![ModelModality::Image]);
assert!(super::models_from_inventory(vec![image_model], &[], true).is_empty());
}
#[tokio::test]
async fn inference_profile_discovery_follows_every_page() {
let mut calls = Vec::new();
let profiles = super::collect_profile_pages(|token| {
calls.push(token.clone());
let page = match token.as_deref() {
None => ListInferenceProfilesOutput::builder()
.inference_profile_summaries(inference_profile("us.future.first", "future.first"))
.next_token("page-2").build(),
Some("page-2") => ListInferenceProfilesOutput::builder()
.inference_profile_summaries(inference_profile("us.future.second", "future.second"))
.build(),
Some(token) => panic!("Unexpected token: {token}"),
};
std::future::ready(Ok(page))
}).await.unwrap();
assert_eq!(profiles.len(), 2);
assert_eq!(calls, [None, Some("page-2".to_owned())]);
}
#[tokio::test]
async fn failed_or_repeating_profile_pages_do_not_return_partial_success() {
let result = super::collect_profile_pages(|_| {
std::future::ready(Err("network unavailable".to_owned()))
}).await;
assert_eq!(result.unwrap_err(), "network unavailable");
let result = super::collect_profile_pages(|_| {
std::future::ready(Ok(ListInferenceProfilesOutput::builder().next_token("repeated").build()))
}).await;
assert!(result.is_err());
}
#[test]
fn transiently_missing_models_keep_preferences_without_restoring_authorization() {
let mut saved = crate::settings::BedrockModelConfig {
enabled: false,
access_verified: true,
model_id: "future.existing".to_owned(),
display_name: "My model name".to_owned(),
..Default::default()
};
let preserved = super::preserve_model_toggles(&[saved.clone()], Vec::new());
assert_eq!(preserved.len(), 1);
assert!(!preserved[0].enabled);
assert!(!preserved[0].access_verified);
saved.invocation_model_id = "old-route".to_owned();
let mut discovered = saved.clone();
discovered.enabled = true;
discovered.display_name = "AWS name".to_owned();
discovered.invocation_model_id = "new-route".to_owned();
let preserved = super::preserve_model_toggles(&[saved], vec![discovered]);
assert!(!preserved[0].enabled);
assert_eq!(preserved[0].display_name, "My model name");
assert_eq!(preserved[0].invocation_model_id, "new-route");
}
#[test]
fn cached_access_is_not_authorization_and_refresh_preserves_disabled_models() {
@@ -8,6 +147,7 @@ fn cached_access_is_not_authorization_and_refresh_preserves_disabled_models() {
vision_supported: false,
enabled: false,
access_verified: true,
..Default::default()
};
let saved = serde_json::to_value(&model).unwrap();
assert!(saved.get("access_verified").is_none());
+2
View File
@@ -268,6 +268,8 @@ fn get_test_config() -> Option<BedrockClientConfig> {
access_key_id: String::new(),
secret_access_key: String::new(),
session_token: None,
credentials_expiry: None,
invocation_model_id: None,
cross_region_inference: false,
})
}
+1
View File
@@ -146,6 +146,7 @@ fn parse_claude_code_model_map(
model_id: arn,
display_name,
vision_supported: true,
..Default::default()
}
})
.collect()
+2
View File
@@ -23,6 +23,8 @@ fn get_test_config() -> Option<BedrockClientConfig> {
access_key_id: String::new(),
secret_access_key: String::new(),
session_token: None,
credentials_expiry: None,
invocation_model_id: None,
cross_region_inference: false,
})
}
+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;
+8
View File
@@ -1,3 +1,11 @@
pub(crate) fn foundation_model_id(model_id: &str) -> &str {
let id = model_id.rsplit('/').next().unwrap_or(model_id);
let id = id.strip_suffix("[1m]").or_else(|| id.strip_suffix("[1M]")).unwrap_or(id);
["global.", "us.", "eu.", "apac.", "jp.", "au."]
.iter()
.find_map(|prefix| id.strip_prefix(prefix))
.unwrap_or(id)
}
pub fn apply_cross_region_prefix(model_id: &str, region: &str) -> String {
if model_id.starts_with("arn:") {
return model_id.to_string();
+5 -5
View File
@@ -314,9 +314,8 @@ pub(crate) async fn provider_runtime_for_request(
model,
reasoning_effort: config.reasoning_effort,
max_output_tokens: config.max_output_tokens.map(u64::from),
auth_file: None,
},
)),
).with_credential_resolver(crate::ai::chatgpt_auth::credential_resolver())),
OpenAIProviderKind::Anthropic => {
let credential = config.api_key.unwrap_or_default();
let stored = crate::ai::claude_subscription::is_stored_credential(&credential);
@@ -346,14 +345,15 @@ pub(crate) async fn provider_runtime_for_request(
},
crate::ai::provider::ProviderConfig::Bedrock(config) => {
let max_output_tokens = Some(bedrock_max_output_tokens(&model));
let cross_region_inference = config.cross_region_inference;
let cross_region_inference = config.invocation_model_id.is_none() && config.cross_region_inference;
let caching_config =
CachingConfig::from_external_config(&ExternalBedrockConfig::load());
CachingConfig::from_external_config(&ExternalBedrockConfig::load()).for_model(&model);
let invocation_model = config.invocation_model_id.clone().unwrap_or(model);
let client = BedrockClient::from_config(config).await?;
Arc::new(BedrockRuntime::from_aws_client(
client.runtime_client(),
BedrockRigConfig {
model,
model: invocation_model,
region: client.region().to_string(),
cross_region_inference,
prompt_caching: caching_config.enabled,
+24
View File
@@ -834,10 +834,34 @@ pub struct BedrockModelConfig {
#[serde(default)]
#[schemars(description = "Whether the model supports image/vision input.")]
pub vision_supported: bool,
#[serde(default)]
pub invocation_model_id: String,
#[serde(default)]
pub provider_name: String,
#[serde(default)]
pub response_streaming_supported: Option<bool>,
#[serde(default)]
pub runtime_supported: Option<bool>,
}
impl settings_value::SettingsValue for BedrockModelConfig {}
impl Default for BedrockModelConfig {
fn default() -> Self {
Self {
enabled: true,
access_verified: false,
model_id: String::new(),
display_name: String::new(),
vision_supported: false,
invocation_model_id: String::new(),
provider_name: String::new(),
response_streaming_supported: None,
runtime_supported: None,
}
}
}
fn default_context_size() -> u32 {
200_000
}
+17
View File
@@ -32,6 +32,22 @@ fn create_test_request_limit_info(
}
}
#[test]
fn legacy_bedrock_model_settings_default_live_capabilities_without_authorization() {
let model: BedrockModelConfig = serde_json::from_value(serde_json::json!({
"model_id": "future.new-model",
"display_name": "Saved model",
"enabled": false,
"vision_supported": true
})).unwrap();
assert_eq!(model.invocation_model_id, "");
assert_eq!(model.provider_name, "");
assert_eq!(model.response_streaming_supported, None);
assert_eq!(model.runtime_supported, None);
assert!(!model.access_verified);
assert!(!model.enabled);
}
fn add_ai_enablement_dependencies_for_test(app: &mut App) {
app.add_singleton_model(|_| AuthStateProvider::new_for_test());
app.add_singleton_model(UserWorkspaces::default_mock);
@@ -446,6 +462,7 @@ fn orchestration_is_enabled_when_ai_is_enabled() {
model_id: "test-model".to_string(),
display_name: "Test model".to_string(),
vision_supported: false,
..Default::default()
}],
ctx,
)
+16 -1
View File
@@ -43,7 +43,7 @@ use strum::IntoEnumIterator;
use super::execution_profile_view::{ExecutionProfileView, ExecutionProfileViewEvent};
use super::provider_setup_view::{
model_catalog_table, AcpProviderDraft, BedrockProviderDraft, ProviderSetupProviderType,
ProviderSetupView, ProviderSetupViewEvent,
ProviderSetupView, ProviderSetupViewAction, ProviderSetupViewEvent,
};
use super::settings_page::{
build_sub_header, build_toggle_element, render_body_item_label,
@@ -2720,6 +2720,7 @@ pub enum AISettingsPageAction {
},
RemoveOpenAIProvider(usize),
EditBedrockProvider,
RefreshBedrockModels,
RemoveBedrockProvider,
RemoveAcpProvider(String),
ToggleFileBasedMcp,
@@ -3457,6 +3458,12 @@ impl TypedActionView for AISettingsPageView {
AISettingsPageAction::EditBedrockProvider => {
self.begin_inline_bedrock_setup(ctx);
}
AISettingsPageAction::RefreshBedrockModels => {
self.begin_inline_bedrock_setup(ctx);
self.provider_setup_body.update(ctx, |body, ctx| {
body.handle_action(&ProviderSetupViewAction::Next, ctx);
});
}
AISettingsPageAction::RemoveBedrockProvider => {
AISettings::handle(ctx).update(ctx, |settings, ctx| {
report_if_error!(settings.bedrock_enabled.set_value(false, ctx));
@@ -6641,6 +6648,7 @@ struct ProviderSettingsWidget {
bedrock_add_button: ViewHandle<ActionButton>,
acp_add_button: ViewHandle<ActionButton>,
bedrock_edit_button: ViewHandle<ActionButton>,
bedrock_refresh_button: ViewHandle<ActionButton>,
bedrock_remove_button: ViewHandle<ActionButton>,
provider_cards: Vec<OpenAIProviderCardState>,
acp_provider_cards: Vec<AcpProviderCardState>,
@@ -6749,6 +6757,11 @@ impl ProviderSettingsWidget {
ctx.dispatch_typed_action(AISettingsPageAction::EditBedrockProvider);
})
});
let bedrock_refresh_button = ctx.add_typed_action_view(|_| {
ActionButton::new("Refresh models", SecondaryTheme).on_click(|ctx| {
ctx.dispatch_typed_action(AISettingsPageAction::RefreshBedrockModels);
})
});
let bedrock_remove_button = ctx.add_typed_action_view(|_| {
ActionButton::new("Delete", DangerSecondaryTheme).on_click(|ctx| {
ctx.dispatch_typed_action(AISettingsPageAction::RemoveBedrockProvider);
@@ -6766,6 +6779,7 @@ impl ProviderSettingsWidget {
bedrock_add_button,
acp_add_button,
bedrock_edit_button,
bedrock_refresh_button,
bedrock_remove_button,
provider_cards,
acp_provider_cards,
@@ -7275,6 +7289,7 @@ impl ProviderSettingsWidget {
.with_child(
Flex::row()
.with_spacing(8.)
.with_child(ChildView::new(&self.bedrock_refresh_button).finish())
.with_child(ChildView::new(&self.bedrock_edit_button).finish())
.with_child(ChildView::new(&self.bedrock_remove_button).finish())
.finish(),
+82 -27
View File
@@ -220,6 +220,7 @@ pub enum ProviderSetupViewAction {
ToggleModel(usize),
CycleModelCapability(usize, CapabilityKey),
ConnectChatGPT,
CancelChatGPTSignIn,
ConnectClaude,
SignOutClaude,
CancelClaudeSignIn,
@@ -687,10 +688,8 @@ impl ProviderSetupView {
self.draft_bedrock = draft;
self.draft_models.clear();
self.discovery_state = DiscoveryState::Idle;
self.bedrock_region_request += 1;
self.sync_editors(ctx);
if self.draft_bedrock.region.is_empty() {
self.load_bedrock_profile_region(ctx);
}
self.sync_provider_type_buttons(ctx);
self.sync_bedrock_auth_buttons(ctx);
self.sync_model_switches(ctx);
@@ -1130,8 +1129,23 @@ impl ProviderSetupView {
}
}
fn bedrock_client_config(&self) -> crate::ai::provider::client::BedrockClientConfig {
crate::ai::provider::client::BedrockClientConfig {
auth_method: self.draft_bedrock.auth_method,
profile: self.draft_bedrock.profile.clone(),
region: self.draft_bedrock.region.clone(),
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,
invocation_model_id: None,
cross_region_inference: self.draft_bedrock.cross_region_inference,
}
}
fn begin_discovery(&mut self, ctx: &mut ViewContext<Self>) {
self.cancel_acp_setup();
self.bedrock_region_request += 1;
self.step = ProviderSetupStep::Discover;
self.discovery_state = DiscoveryState::Loading;
self.acp_setup_status.clear();
@@ -1145,33 +1159,45 @@ impl ProviderSetupView {
return;
}
ProviderSetupProviderType::Bedrock => {
let config = crate::ai::provider::client::BedrockClientConfig {
auth_method: self.draft_bedrock.auth_method,
profile: self.draft_bedrock.profile.clone(),
region: self.draft_bedrock.region.clone(),
access_key_id: self.draft_bedrock.access_key_id.clone(),
secret_access_key: self.draft_bedrock.secret_access_key.clone(),
session_token: None,
cross_region_inference: self.draft_bedrock.cross_region_inference,
};
let config = self.bedrock_client_config();
let requested_config = config.clone();
let generation = self.acp_discovery_generation;
let verification_generation = LLMPreferences::handle(ctx)
.update(ctx, |preferences, _| preferences.begin_bedrock_verification());
ctx.spawn(
async move {
crate::ai::provider::discovery::discover_available_models(config).await
},
move |me, result, ctx| match result {
move |me, result, ctx| {
if generation != me.acp_discovery_generation
|| me.provider_type != ProviderSetupProviderType::Bedrock
|| me.step != ProviderSetupStep::Discover
|| me.bedrock_client_config() != requested_config
{
return;
}
match result {
Ok(discovery) => {
for record in discovery.records {
crate::ai::remote_logging::log_view_event(ctx, record);
}
let models = discovery.models;
LLMPreferences::handle(ctx).update(ctx, |preferences, ctx| {
let accepted = LLMPreferences::handle(ctx).update(ctx, |preferences, ctx| {
preferences.record_bedrock_verification(
verification_generation,
requested_config,
&models,
ctx,
);
)
});
if !accepted {
me.discovery_state = DiscoveryState::Failed(
"A newer Bedrock connection test superseded this result. Retry the test.".to_owned(),
);
me.update_next_button(ctx);
ctx.notify();
return;
}
me.draft_bedrock.models =
crate::ai::provider::discovery::preserve_model_toggles(
&me.draft_bedrock.models,
@@ -1189,6 +1215,7 @@ impl ProviderSetupView {
me.update_next_button(ctx);
ctx.notify();
}
}
},
);
return;
@@ -1564,10 +1591,10 @@ impl ProviderSetupView {
);
}
if matches!(
state,
ChatGPTAuthState::NotConnected | ChatGPTAuthState::Failed(_)
) {
let pending = matches!(state,
ChatGPTAuthState::AwaitingBrowser | ChatGPTAuthState::ExchangingToken
);
{
children.push(
appearance
.ui_builder()
@@ -1575,10 +1602,21 @@ impl ProviderSetupView {
ButtonVariant::Secondary,
self.chatgpt_connect_mouse_state.clone(),
)
.with_text_label("Connect ChatGPT".to_owned())
.with_text_label(if pending {
"Cancel sign-in"
} else if matches!(state, ChatGPTAuthState::Connected) {
"Reconnect ChatGPT"
} else {
"Connect ChatGPT"
}.to_owned())
.build()
.on_click(|ctx, _, _| {
ctx.dispatch_typed_action(ProviderSetupViewAction::ConnectChatGPT);
.on_click(move |ctx, _, _| {
let action = if pending {
ProviderSetupViewAction::CancelChatGPTSignIn
} else {
ProviderSetupViewAction::ConnectChatGPT
};
ctx.dispatch_typed_action(action);
})
.finish(),
);
@@ -1835,7 +1873,7 @@ impl ProviderSetupView {
.finish()
}
fn render_discovery(&self, appearance: &Appearance) -> Box<dyn galaxyui::Element> {
fn render_discovery(&self, appearance: &Appearance, app: &AppContext) -> Box<dyn galaxyui::Element> {
let (message, color) = match &self.discovery_state {
DiscoveryState::Idle | DiscoveryState::Loading => (
if self.provider_type == ProviderSetupProviderType::Acp
@@ -1849,7 +1887,7 @@ impl ProviderSetupView {
),
DiscoveryState::Failed(error) => (error.clone(), appearance.theme().ui_error_color()),
};
Flex::column()
let mut column = Flex::column()
.with_spacing(12.)
.with_child(
Text::new(
@@ -1887,8 +1925,13 @@ impl ProviderSetupView {
.with_color(color)
.soft_wrap(true)
.finish()
})
.finish()
});
if self.provider_type == ProviderSetupProviderType::ChatGPTSubscription
&& matches!(self.discovery_state, DiscoveryState::Failed(_))
{
column = column.with_child(self.render_chatgpt_auth(appearance, app));
}
column.finish()
}
fn render_model_table(
@@ -2442,7 +2485,7 @@ impl View for ProviderSetupView {
let content = match self.step {
ProviderSetupStep::ProviderType => self.render_provider_type(appearance),
ProviderSetupStep::Configure => self.render_configure(appearance, app),
ProviderSetupStep::Discover => self.render_discovery(appearance),
ProviderSetupStep::Discover => self.render_discovery(appearance, app),
ProviderSetupStep::Models => self.render_models(appearance),
};
Flex::column()
@@ -2462,6 +2505,10 @@ impl TypedActionView for ProviderSetupView {
match action {
ProviderSetupViewAction::SelectProvider(kind) => {
if self.provider_type != *kind {
#[cfg(not(target_family = "wasm"))]
if self.provider_type == ProviderSetupProviderType::ChatGPTSubscription {
ChatGPTAuthModel::handle(ctx).update(ctx, |model, ctx| model.cancel_sign_in(ctx));
}
self.cancel_claude_signin();
self.draft_api_key = None;
self.claude_credentials_available = false;
@@ -2562,6 +2609,10 @@ impl TypedActionView for ProviderSetupView {
ProviderSetupViewAction::Cancel => {
self.cancel_acp_setup();
self.cancel_claude_signin();
#[cfg(not(target_family = "wasm"))]
if self.provider_type == ProviderSetupProviderType::ChatGPTSubscription {
ChatGPTAuthModel::handle(ctx).update(ctx, |model, ctx| model.cancel_sign_in(ctx));
}
self.discovery_state = DiscoveryState::Idle;
ctx.emit(ProviderSetupViewEvent::Close);
}
@@ -2607,6 +2658,10 @@ impl TypedActionView for ProviderSetupView {
#[cfg(not(target_family = "wasm"))]
ChatGPTAuthModel::handle(ctx).update(ctx, |model, ctx| model.connect(ctx));
}
ProviderSetupViewAction::CancelChatGPTSignIn => {
#[cfg(not(target_family = "wasm"))]
ChatGPTAuthModel::handle(ctx).update(ctx, |model, ctx| model.cancel_sign_in(ctx));
}
ProviderSetupViewAction::ConnectClaude => {
#[cfg(not(target_family = "wasm"))]
self.connect_claude(ctx);
+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();
+52 -22
View File
@@ -1,6 +1,7 @@
use std::path::PathBuf;
use std::sync::Arc;
use async_trait::async_trait;
use futures::future::BoxFuture;
use futures::StreamExt;
use galaxy_agent_core::{
AgentError, AgentErrorKind, AgentEvent, AgentEventStream, AgentRuntime, ConversationMessage,
@@ -8,35 +9,50 @@ use galaxy_agent_core::{
TurnRequest, turn_control,
};
use rig_core::client::CompletionClient;
use rig_core::providers::chatgpt;
use rig_core::providers::chatgpt::{self, ChatGPTAuth};
use crate::request::build_completion_request;
use crate::stream::start_model_turn;
/// Refreshes Galaxy's cached ChatGPT subscription credentials without allowing
/// Rig to fall back to its interactive device-code flow.
pub async fn refresh_chatgpt_subscription_credentials() -> Result<(), String> {
let client = chatgpt::Client::builder()
.oauth()
.allow_device_flow(false)
.originator("galaxy")
.build()
.map_err(|error| error.to_string())?;
client.authorize().await.map_err(|error| error.to_string())
#[derive(Clone, PartialEq, Eq)]
pub struct ChatGPTSubscriptionCredentials {
pub access_token: String,
pub account_id: Option<String>,
}
impl std::fmt::Debug for ChatGPTSubscriptionCredentials {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ChatGPTSubscriptionCredentials")
.field("access_token", &"<redacted>")
.field("account_id", &self.account_id.as_ref().map(|_| "<redacted>"))
.finish()
}
}
pub type ChatGPTCredentialResolver = Arc<
dyn Fn() -> BoxFuture<'static, Result<ChatGPTSubscriptionCredentials, String>> + Send + Sync,
>;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ChatGPTSubscriptionRuntimeConfig {
pub model: String,
pub reasoning_effort: Option<String>,
pub max_output_tokens: Option<u64>,
pub auth_file: Option<PathBuf>,
}
#[derive(Clone, Debug)]
#[derive(Clone)]
pub struct ChatGPTSubscriptionRuntime {
config: ChatGPTSubscriptionRuntimeConfig,
descriptor: RuntimeDescriptor,
credential_resolver: Option<ChatGPTCredentialResolver>,
}
impl std::fmt::Debug for ChatGPTSubscriptionRuntime {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ChatGPTSubscriptionRuntime")
.field("config", &self.config)
.field("has_credential_resolver", &self.credential_resolver.is_some())
.finish()
}
}
fn reasoning_additional_params(reasoning_effort: Option<&str>) -> Option<serde_json::Value> {
@@ -56,7 +72,24 @@ impl ChatGPTSubscriptionRuntime {
kind: RuntimeKind::Provider,
capabilities: RuntimeCapabilities::provider(),
};
Self { config, descriptor }
Self { config, descriptor, credential_resolver: None }
}
pub fn with_credential_resolver(mut self, resolver: ChatGPTCredentialResolver) -> Self {
self.credential_resolver = Some(resolver);
self
}
async fn resolve_auth(&self) -> Result<ChatGPTAuth, AgentError> {
let resolver = self.credential_resolver.as_ref().ok_or_else(|| {
AgentError::new(AgentErrorKind::Authentication, "ChatGPT sign-in required. Reconnect ChatGPT in Providers.")
})?;
let credentials = resolver().await
.map_err(|message| AgentError::new(AgentErrorKind::Authentication, message))?;
Ok(ChatGPTAuth::AccessToken {
access_token: credentials.access_token,
account_id: credentials.account_id,
})
}
/// Completes a text-only request for callers such as crosscheck reviewers.
@@ -117,14 +150,11 @@ impl AgentRuntime for ChatGPTSubscriptionRuntime {
request: TurnRequest,
control: TurnControl,
) -> Result<AgentEventStream, AgentError> {
let mut builder = chatgpt::Client::builder()
.oauth()
let auth = self.resolve_auth().await?;
let client = chatgpt::Client::builder()
.api_key(auth)
.allow_device_flow(false)
.originator("galaxy");
if let Some(auth_file) = &self.config.auth_file {
builder = builder.auth_file(auth_file);
}
let client = builder
.originator("galaxy")
.build()
.map_err(|error| AgentError::new(AgentErrorKind::Configuration, error.to_string()))?;
let model = client.completion_model(&self.config.model);
+108 -3
View File
@@ -1,13 +1,14 @@
use galaxy_agent_core::StopReason;
use rig_core::client::CompletionClient;
use rig_core::completion::{AssistantContent, CompletionModel, Message};
use rig_core::message::{ToolResultContent, UserContent};
use rig_core::providers::chatgpt::ChatGPTAuth;
use rig_core::test_utils::RecordingHttpClient;
use rig_core::test_utils::{RecordingHttpClient, SequencedStreamingHttpClient};
use super::*;
const COMPLETED_RESPONSE: &str = r#"data: {"type":"response.output_text.delta","delta":"ok"}
data: {"type":"response.completed","response":{"id":"resp_test","object":"response","created_at":1,"status":"completed","error":null,"incomplete_details":null,"instructions":null,"max_output_tokens":null,"model":"gpt-5.3-codex","usage":{"input_tokens":1,"input_tokens_details":{"cached_tokens":0},"output_tokens":1,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":2},"output":[{"type":"message","id":"msg_test","status":"completed","role":"assistant","content":[{"type":"output_text","annotations":[],"text":"ok"}]}],"tools":[]}}
data: {"type":"response.completed","sequence_number":54,"response":{"id":"resp_test","object":"response","created_at":1,"status":"completed","error":null,"incomplete_details":null,"instructions":null,"max_output_tokens":null,"model":"gpt-5.3-codex","usage":{"input_tokens":1,"input_tokens_details":{"cached_tokens":0},"output_tokens":1,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":2},"output":[{"type":"message","id":"msg_test","status":"completed","role":"assistant","content":[{"type":"output_text","annotations":[],"text":"ok"}]}],"tools":[]}}
data: [DONE]"#;
#[test]
@@ -16,7 +17,6 @@ fn runtime_descriptor_identifies_chatgpt_subscription() {
model: "gpt-5.3-codex".to_string(),
reasoning_effort: Some("high".to_string()),
max_output_tokens: None,
auth_file: None,
});
assert_eq!(
@@ -25,6 +25,43 @@ fn runtime_descriptor_identifies_chatgpt_subscription() {
);
}
#[tokio::test]
async fn runtime_resolves_current_credentials_without_rig_oauth() {
let runtime = ChatGPTSubscriptionRuntime::new(ChatGPTSubscriptionRuntimeConfig {
model: "test-model".to_owned(),
reasoning_effort: None,
max_output_tokens: None,
}).with_credential_resolver(std::sync::Arc::new(|| {
Box::pin(async {
Ok(ChatGPTSubscriptionCredentials {
access_token: "resolved-token".to_owned(),
account_id: Some("resolved-account".to_owned()),
})
})
}));
match runtime.resolve_auth().await.unwrap() {
ChatGPTAuth::AccessToken { access_token, account_id } => {
assert_eq!(access_token, "resolved-token");
assert_eq!(account_id.as_deref(), Some("resolved-account"));
}
ChatGPTAuth::OAuth => panic!("The runtime must not own OAuth refresh"),
}
}
#[tokio::test]
async fn missing_or_failed_resolver_requires_sign_in_without_device_flow() {
let runtime = ChatGPTSubscriptionRuntime::new(ChatGPTSubscriptionRuntimeConfig {
model: "test-model".to_owned(),
reasoning_effort: None,
max_output_tokens: None,
});
assert_eq!(runtime.resolve_auth().await.unwrap_err().kind, AgentErrorKind::Authentication);
let runtime = runtime.with_credential_resolver(std::sync::Arc::new(|| {
Box::pin(async { Err("Sign in again".to_owned()) })
}));
assert_eq!(runtime.resolve_auth().await.unwrap_err().kind, AgentErrorKind::Authentication);
}
#[test]
fn reasoning_effort_is_encoded_for_the_responses_request() {
let request = TurnRequest::new(
@@ -115,3 +152,71 @@ async fn chatgpt_follow_up_request_preserves_responses_call_ids() {
assert_eq!(input[0]["call_id"], "call_native_1");
assert_eq!(input[1]["call_id"], "call_native_1");
}
#[tokio::test]
async fn chatgpt_metadata_does_not_interrupt_text_or_completion() {
let metadata = r#"data: {"type":"response.metadata","response_id":"resp_test","sequence_number":51,"metadata":{"openai_chatgpt_moderation_metadata":{"prompt":{},"generation":{},"tool_call":{},"tool_response":{}}}}"#;
let delta = serde_json::json!({
"type": "response.output_text.delta",
"item_id": "msg_test",
"output_index": 0,
"content_index": 0,
"sequence_number": 52,
"delta": "ok",
});
let completed = COMPLETED_RESPONSE
.lines()
.find(|line| line.starts_with("data: {\"type\":\"response.completed\""))
.expect("fixture should contain a completed response");
let response =
format!("{metadata}\n\ndata: {delta}\n\n{metadata}\n\n{completed}\n\ndata: [DONE]\n\n");
let client = chatgpt::Client::builder()
.api_key(ChatGPTAuth::AccessToken {
access_token: "test-token".to_string(),
account_id: None,
})
.http_client(SequencedStreamingHttpClient::new(vec![Ok(response.into())]))
.build()
.expect("client should build");
let request = build_completion_request(
TurnRequest::new(
"gpt-5.3-codex".to_string(),
vec![ConversationMessage {
role: MessageRole::User,
content: MessageContent::Text("hello".to_string()),
}],
),
None,
true,
true,
None,
)
.expect("request should convert");
let (_sender, control) = turn_control();
let mut stream = start_model_turn(
client.completion_model("gpt-5.3-codex"),
request,
control,
None,
)
.await
.expect("request should reach the mocked provider");
let mut text = String::new();
let mut usage = None;
let mut stop_reason = None;
while let Some(event) = stream.next().await {
match event.expect("metadata should not fail the turn") {
AgentEvent::TextDelta { text: delta } => text.push_str(&delta),
AgentEvent::UsageUpdated { usage: updated } => usage = Some(updated),
AgentEvent::TurnStopped { reason } => stop_reason = Some(reason),
_ => {}
}
}
assert_eq!(text, "ok");
assert_eq!(
usage.expect("completion should report usage").output_tokens,
1
);
assert_eq!(stop_reason, Some(StopReason::Completed));
}
+12
View File
@@ -160,6 +160,13 @@ where
});
}
Ok(StreamedAssistantContent::Unknown(value)) => {
// ChatGPT subscription streams include informational moderation
// metadata that Rig does not map to assistant content.
if value.get("type").and_then(serde_json::Value::as_str)
== Some("response.metadata")
{
continue;
}
if is_keepalive_event(&value) {
yield Ok(AgentEvent::KeepAlive);
} else {
@@ -512,6 +519,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()
@@ -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);
}
}