From 04cff915095c385000f58149a1fa7c16dd6046ee Mon Sep 17 00:00:00 2001 From: Josh Woodcock Date: Wed, 14 Oct 2026 12:43:32 -0500 Subject: [PATCH 01/12] fix bedrock authentication refresh --- app/src/ai/aws_credentials.rs | 56 +++- app/src/ai/blocklist/block.rs | 4 +- app/src/ai/blocklist/controller.rs | 29 +- .../blocklist/controller/response_stream.rs | 28 +- .../aws_bedrock_credentials_error.rs | 39 +-- app/src/ai/crosscheck/reviewer.rs | 27 +- app/src/ai/llms.rs | 10 +- app/src/ai/mod.rs | 2 + app/src/ai/provider/authentication.rs | 316 ++++++++++++++++++ app/src/ai/provider/authentication_tests.rs | 277 +++++++++++++++ app/src/ai/provider/client.rs | 48 ++- app/src/ai/provider/discovery.rs | 2 +- app/src/ai/provider/e2e_tests.rs | 1 + app/src/ai/provider/integration_tests.rs | 1 + app/src/ai/provider/mod.rs | 2 + app/src/settings_view/provider_setup_view.rs | 1 + app/src/terminal/view.rs | 150 +++------ crates/galaxy_agent_rig/src/stream.rs | 5 + .../galaxy_agent_rig/src/stream_unit_tests.rs | 16 + 19 files changed, 777 insertions(+), 237 deletions(-) create mode 100644 app/src/ai/provider/authentication.rs create mode 100644 app/src/ai/provider/authentication_tests.rs diff --git a/app/src/ai/aws_credentials.rs b/app/src/ai/aws_credentials.rs index 1829721d..12b45fe6 100644 --- a/app/src/ai/aws_credentials.rs +++ b/app/src/ai/aws_credentials.rs @@ -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 { - 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 { diff --git a/app/src/ai/blocklist/block.rs b/app/src/ai/blocklist/block.rs index faef21f8..f60bfc74 100644 --- a/app/src/ai/blocklist/block.rs +++ b/app/src/ai/blocklist/block.rs @@ -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 diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 5d810dc2..7a23ffd3 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -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.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, ) { 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, diff --git a/app/src/ai/blocklist/controller/response_stream.rs b/app/src/ai/blocklist/controller/response_stream.rs index a80f5347..7e7340c1 100644 --- a/app/src/ai/blocklist/controller/response_stream.rs +++ b/app/src/ai/blocklist/controller/response_stream.rs @@ -187,30 +187,12 @@ impl ResponseStream { // Fall back to Bedrock if *settings.bedrock_enabled.value() { - let auth_method = *settings.bedrock_auth_method.value(); - let region = settings.bedrock_region.value().clone(); - let cross_region_inference = *settings.bedrock_cross_region_inference.value(); let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx); - let mut config = BedrockClientConfig { - auth_method, - profile: settings.bedrock_profile.value().clone(), - region, - access_key_id: settings.bedrock_access_key_id.value().clone(), - secret_access_key: settings.bedrock_secret_access_key.value().clone(), - session_token: None, - cross_region_inference, - }; - - if let ::ai::api_keys::AwsCredentialsState::Loaded { credentials, .. } = - api_key_manager.aws_credentials_state() - { - config.auth_method = crate::settings::BedrockAuthMethod::StaticKeys; - config.access_key_id = credentials.access_key().to_string(); - config.secret_access_key = credentials.secret_key().to_string(); - config.session_token = credentials.session_token().map(|s| s.to_string()); - } - - return ProviderConfig::Bedrock(config.with_external_fallbacks()); + let config = BedrockClientConfig::from_settings(settings).with_managed_credentials( + &api_key_manager.aws_credentials_refresh_strategy(), + api_key_manager.aws_credentials_state(), + ); + return ProviderConfig::Bedrock(config); } ProviderConfig::None diff --git a/app/src/ai/blocklist/inline_action/aws_bedrock_credentials_error.rs b/app/src/ai/blocklist/inline_action/aws_bedrock_credentials_error.rs index 96e6800c..9aa89ea2 100644 --- a/app/src/ai/blocklist/inline_action/aws_bedrock_credentials_error.rs +++ b/app/src/ai/blocklist/inline_action/aws_bedrock_credentials_error.rs @@ -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, @@ -52,7 +50,6 @@ impl AwsBedrockCredentialsErrorView { pub fn new( model_name: String, login_command: String, - auto_login_triggered: bool, ctx: &mut ViewContext, ) -> 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., ) diff --git a/app/src/ai/crosscheck/reviewer.rs b/app/src/ai/crosscheck/reviewer.rs index 2e6bbcca..73ebb556 100644 --- a/app/src/ai/crosscheck/reviewer.rs +++ b/app/src/ai/crosscheck/reviewer.rs @@ -181,28 +181,13 @@ impl CrosscheckReviewer { // Fall back to Bedrock via external config if *settings.bedrock_enabled.value() { - let auth_method = *settings.bedrock_auth_method.value(); let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx); - let mut config = crate::ai::provider::client::BedrockClientConfig { - auth_method, - profile: settings.bedrock_profile.value().clone(), - region: settings.bedrock_region.value().clone(), - access_key_id: settings.bedrock_access_key_id.value().clone(), - secret_access_key: settings.bedrock_secret_access_key.value().clone(), - session_token: None, - cross_region_inference: *settings.bedrock_cross_region_inference.value(), - }; - - if let ::ai::api_keys::AwsCredentialsState::Loaded { credentials, .. } = - api_key_manager.aws_credentials_state() - { - config.auth_method = crate::settings::BedrockAuthMethod::StaticKeys; - config.access_key_id = credentials.access_key().to_string(); - config.secret_access_key = credentials.secret_key().to_string(); - config.session_token = credentials.session_token().map(|s| s.to_string()); - } - - return ProviderConfig::Bedrock(config.with_external_fallbacks()); + let config = crate::ai::provider::client::BedrockClientConfig::from_settings(settings) + .with_managed_credentials( + &api_key_manager.aws_credentials_refresh_strategy(), + api_key_manager.aws_credentials_state(), + ); + return ProviderConfig::Bedrock(config); } ProviderConfig::None diff --git a/app/src/ai/llms.rs b/app/src/ai/llms.rs index 6c96743e..8779a857 100644 --- a/app/src/ai/llms.rs +++ b/app/src/ai/llms.rs @@ -3471,13 +3471,5 @@ mod tests; fn bedrock_discovery_config( settings: &AISettings, ) -> crate::ai::provider::client::BedrockClientConfig { - crate::ai::provider::client::BedrockClientConfig { - auth_method: *settings.bedrock_auth_method.value(), - profile: settings.bedrock_profile.value().clone(), - region: settings.bedrock_region.value().clone(), - access_key_id: settings.bedrock_access_key_id.value().clone(), - secret_access_key: settings.bedrock_secret_access_key.value().clone(), - session_token: None, - cross_region_inference: *settings.bedrock_cross_region_inference.value(), - } + crate::ai::provider::client::BedrockClientConfig::from_settings(settings) } diff --git a/app/src/ai/mod.rs b/app/src/ai/mod.rs index 0d63273b..c956cd27 100644 --- a/app/src/ai/mod.rs +++ b/app/src/ai/mod.rs @@ -85,6 +85,8 @@ pub mod outline; pub(crate) use ai::paths; pub fn init(app: &mut AppContext) { + #[cfg(not(target_family = "wasm"))] + app.add_singleton_model(provider::authentication::BedrockAuthentication::new); #[cfg(not(target_family = "wasm"))] app.add_singleton_model(acp::AcpRuntimeModel::new); #[cfg(not(target_family = "wasm"))] diff --git a/app/src/ai/provider/authentication.rs b/app/src/ai/provider/authentication.rs new file mode 100644 index 00000000..6431588e --- /dev/null +++ b/app/src/ai/provider/authentication.rs @@ -0,0 +1,316 @@ +use std::collections::HashMap; +use std::time::SystemTime; + +use ai::api_keys::{ApiKeyManager, AwsCredentials, AwsCredentialsRefreshStrategy, AwsCredentialsState}; +use futures::channel::oneshot; +use futures::future::{BoxFuture, Shared}; +use futures::FutureExt; +use galaxyui::{AppContext, Entity, ModelContext, SingletonEntity}; +use settings::Setting; + +use super::client::{BedrockClientConfig, BedrockError}; +use crate::ai::agent::RenderableAIError; +use crate::settings::{AISettings, AISettingsChangedEvent, BedrockAuthMethod}; + +impl BedrockClientConfig { + pub(crate) fn from_settings(settings: &AISettings) -> Self { + Self { + auth_method: *settings.bedrock_auth_method.value(), + profile: settings.bedrock_profile.value().clone(), + region: settings.bedrock_region.value().clone(), + access_key_id: settings.bedrock_access_key_id.value().clone(), + secret_access_key: settings.bedrock_secret_access_key.value().clone(), + session_token: None, + credentials_expiry: None, + cross_region_inference: *settings.bedrock_cross_region_inference.value(), + } + } + + pub(crate) fn with_managed_credentials( + mut self, + strategy: &AwsCredentialsRefreshStrategy, + state: &AwsCredentialsState, + ) -> Self { + match strategy { + AwsCredentialsRefreshStrategy::LocalChain => self, + AwsCredentialsRefreshStrategy::OidcManaged { region, .. } => { + self.auth_method = BedrockAuthMethod::StaticKeys; + self.region.clone_from(region); + self.access_key_id.clear(); + self.secret_access_key.clear(); + self.session_token = None; + self.credentials_expiry = None; + if let AwsCredentialsState::Loaded { credentials, .. } = state { + self.access_key_id = credentials.access_key().to_owned(); + self.secret_access_key = credentials.secret_key().to_owned(); + self.session_token = credentials.session_token().map(str::to_owned); + self.credentials_expiry = Some( + credentials.expires_at().unwrap_or(SystemTime::UNIX_EPOCH), + ); + } + self + } + } + } +} + +pub(crate) fn credentials_are_current(expiry: Option, 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::() { + 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, 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, +} + +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; +type CredentialFuture = Shared>; + +#[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 { + 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, + ) -> 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::() + .map_or(0, |authentication| authentication.revision()) +} + +#[cfg(test)] +#[path = "authentication_tests.rs"] +mod tests; diff --git a/app/src/ai/provider/authentication_tests.rs b/app/src/ai/provider/authentication_tests.rs new file mode 100644 index 00000000..3b7d2ab1 --- /dev/null +++ b/app/src/ai/provider/authentication_tests.rs @@ -0,0 +1,277 @@ +use std::time::{Duration, SystemTime}; + +use ai::api_keys::{AwsCredentials, AwsCredentialsRefreshStrategy, AwsCredentialsState}; + +use super::*; + +fn config(auth_method: BedrockAuthMethod) -> BedrockClientConfig { + BedrockClientConfig { + auth_method, + profile: "default".to_owned(), + region: String::new(), + access_key_id: "configured-key".to_owned(), + secret_access_key: "configured-secret".to_owned(), + session_token: None, + credentials_expiry: None, + cross_region_inference: true, + } +} + +#[test] +fn provider_profile_and_region_are_authoritative_including_defaults() { + galaxyui::App::test((), |mut app| async move { + crate::test_util::settings::initialize_settings_for_tests(&mut app); + for (profile, region) in [("default", ""), ("selected-profile", "eu-west-1")] { + AISettings::handle(&app).update(&mut app, |settings, ctx| { + settings.bedrock_profile.set_value(profile.to_owned(), ctx).unwrap(); + settings.bedrock_region.set_value(region.to_owned(), ctx).unwrap(); + }); + app.read(|ctx| { + let resolved = BedrockClientConfig::from_settings(AISettings::as_ref(ctx)); + assert_eq!(resolved.profile, profile); + assert_eq!(resolved.region, region); + }); + } + }); +} + +#[test] +fn custom_login_commands_are_not_rewritten() { + for command in ["", "company-login --interactive", "aws sso login --profile custom"] { + assert_eq!(login_command(command, "selected"), command); + } +} + +#[test] +fn default_login_targets_and_quotes_the_effective_profile() { + assert_eq!( + login_command("aws sso login", "default"), + "aws sso login --profile default" + ); + let command = login_command( + "aws sso login", + "team profile", + ); + assert_eq!( + login_command_parts(&command).unwrap(), + ["aws", "sso", "login", "--profile", "team profile"] + ); + assert!(login_command_parts("").is_err()); + assert!(login_command_parts("aws 'unterminated").is_err()); +} + +#[test] +fn local_credential_snapshots_never_override_configured_authentication() { + let now = SystemTime::now(); + for auth in [ + BedrockAuthMethod::Profile, + BedrockAuthMethod::Sso, + BedrockAuthMethod::StaticKeys, + ] { + for expiry in [None, Some(now - Duration::from_secs(1)), Some(now + Duration::from_secs(60))] { + let configured = config(auth); + let loaded = AwsCredentialsState::Loaded { + credentials: AwsCredentials::new( + "snapshot-key".to_owned(), + "snapshot-secret".to_owned(), + Some("snapshot-token".to_owned()), + expiry, + ), + loaded_at: now, + }; + assert_eq!( + configured.clone().with_managed_credentials( + &AwsCredentialsRefreshStrategy::LocalChain, + &loaded, + ), + configured + ); + } + } +} + +#[test] +fn managed_oidc_credentials_preserve_expiry_and_do_not_fall_back_to_local_keys() { + let now = SystemTime::now(); + let strategy = AwsCredentialsRefreshStrategy::OidcManaged { + task_id: Some("task".to_owned()), + role_arn: "test-role".to_owned(), + region: "us-east-1".to_owned(), + }; + let loaded = AwsCredentialsState::Loaded { + credentials: AwsCredentials::new( + "managed-key".to_owned(), + "managed-secret".to_owned(), + Some("managed-token".to_owned()), + Some(now + Duration::from_secs(60)), + ), + loaded_at: now, + }; + let resolved = config(BedrockAuthMethod::Profile).with_managed_credentials(&strategy, &loaded); + assert_eq!(resolved.auth_method, BedrockAuthMethod::StaticKeys); + assert_eq!(resolved.region, "us-east-1"); + assert_eq!(resolved.access_key_id, "managed-key"); + assert_eq!(resolved.credentials_expiry, Some(now + Duration::from_secs(60))); + for state in [ + AwsCredentialsState::Missing, + AwsCredentialsState::Disabled, + AwsCredentialsState::Refreshing, + AwsCredentialsState::Failed { message: "failed".to_owned() }, + ] { + let resolved = config(BedrockAuthMethod::Profile).with_managed_credentials(&strategy, &state); + assert!(resolved.access_key_id.is_empty()); + assert!(resolved.secret_access_key.is_empty()); + } +} + +#[test] +fn automatic_recovery_is_bounded_per_user_exchange_but_manual_retry_is_available() { + let mut attempts = RecoveryAttempts::default(); + assert!(attempts.begin("conversation", "query-1", "connection", false)); + assert!(!attempts.begin("conversation", "query-1", "connection", false)); + assert!(attempts.begin("conversation", "query-1", "connection", true)); + assert!(attempts.begin("conversation", "query-2", "connection", false)); + assert!(attempts.begin("conversation", "query-2", "new-connection", false)); + assert!(attempts.begin("another-conversation", "query-2", "connection", false)); +} + +#[test] +fn credential_expiry_is_checked_at_the_boundary() { + let now = SystemTime::now(); + assert!(credentials_are_current(None, now)); + assert!(credentials_are_current(Some(now + Duration::from_secs(1)), now)); + assert!(!credentials_are_current(Some(now), now)); + assert!(!credentials_are_current(Some(now - Duration::from_secs(1)), now)); +} + +fn test_credentials() -> AwsCredentials { + AwsCredentials::new("test-key".to_owned(), "test-secret".to_owned(), None, None) +} + +#[tokio::test] +async fn external_refresh_recovers_without_running_login() { + let credentials = recover_credentials( + BedrockAuthMethod::Sso, + false, + || futures::future::ready(Ok(test_credentials())).boxed(), + || futures::future::ready(Err("Login must not run".to_owned())).boxed(), + ).await.unwrap(); + assert_eq!(credentials, test_credentials()); +} + +#[tokio::test] +async fn expired_session_runs_login_then_awaits_a_fresh_reload() { + use std::sync::{Arc, atomic::{AtomicUsize, Ordering}}; + + let calls = Arc::new(AtomicUsize::new(0)); + let load_calls = calls.clone(); + let login_calls = calls.clone(); + let result = recover_credentials( + BedrockAuthMethod::Sso, + false, + move || { + let index = load_calls.fetch_add(1, Ordering::SeqCst); + let result = match index { + 0 => Err("expired".to_owned()), + 2 => Ok(test_credentials()), + index => panic!("Unexpected reload order: {index}"), + }; + futures::future::ready(result).boxed() + }, + move || { + assert_eq!(login_calls.fetch_add(1, Ordering::SeqCst), 1); + futures::future::ready(Ok(())).boxed() + }, + ).await; + assert_eq!(result, Ok(test_credentials())); + assert_eq!(calls.load(Ordering::SeqCst), 3); +} + +#[tokio::test] +async fn manual_login_failure_does_not_reload_or_resume() { + let result = recover_credentials( + BedrockAuthMethod::Sso, + true, + || panic!("Credentials must not reload after failed login"), + || futures::future::ready(Err("login failed".to_owned())).boxed(), + ).await; + assert_eq!(result, Err("login failed".to_owned())); +} + +#[tokio::test] +async fn login_success_with_failed_reload_remains_an_error() { + let result = recover_credentials( + BedrockAuthMethod::Sso, + true, + || futures::future::ready(Err("reload failed".to_owned())).boxed(), + || futures::future::ready(Ok(())).boxed(), + ).await; + assert_eq!(result, Err("reload failed".to_owned())); +} + +#[tokio::test] +async fn shared_recovery_executes_login_once_for_overlapping_waiters() { + use std::sync::{Arc, atomic::{AtomicUsize, Ordering}}; + + let calls = Arc::new(AtomicUsize::new(0)); + let login_calls = calls.clone(); + let recovery = recover_credentials( + BedrockAuthMethod::Sso, + true, + || futures::future::ready(Ok(test_credentials())).boxed(), + move || { + login_calls.fetch_add(1, Ordering::SeqCst); + futures::future::ready(Ok(())).boxed() + }, + ).boxed().shared(); + let (first, second) = futures::join!(recovery.clone(), recovery); + assert_eq!(first, Ok(test_credentials())); + assert_eq!(second, Ok(test_credentials())); + assert_eq!(calls.load(Ordering::SeqCst), 1); +} + +#[test] +fn credential_preparation_failures_surface_authentication_recovery() { + use super::super::client::BedrockError; + + for error in [ + BedrockError::CredentialsNotConfigured, + BedrockError::CredentialsLoadFailed("SSO session invalid".to_owned()), + ] { + let error = provider_startup_error(&anyhow::Error::new(error), "selected-model"); + assert!(error.is_aws_bedrock_credentials_error()); + } + for error in [ + BedrockError::RegionNotConfigured, + BedrockError::ApiError("transport".to_owned()), + BedrockError::ModelNotFound("model".to_owned()), + BedrockError::AccessDenied("policy".to_owned()), + BedrockError::Throttling("busy".to_owned()), + BedrockError::ValidationError("request".to_owned()), + ] { + let error = provider_startup_error(&anyhow::Error::new(error), "selected-model"); + assert!(!error.is_aws_bedrock_credentials_error()); + } + assert!(!provider_startup_error(&anyhow::anyhow!("unrelated"), "selected-model") + .is_aws_bedrock_credentials_error()); +} + +#[tokio::test] +async fn profile_authentication_never_starts_a_login_flow() { + for manual in [false, true] { + let result = recover_credentials( + BedrockAuthMethod::Profile, + manual, + || futures::future::ready(Ok(test_credentials())).boxed(), + || panic!("Profile authentication must not execute login commands"), + ).await; + assert_eq!(result, Ok(test_credentials())); + let result = recover_credentials( + BedrockAuthMethod::Profile, + manual, + || futures::future::ready(Err("profile needs fresh credentials".to_owned())).boxed(), + || panic!("A failed profile reload must not execute login commands"), + ).await; + assert_eq!(result, Err("profile needs fresh credentials".to_owned())); + } +} diff --git a/app/src/ai/provider/client.rs b/app/src/ai/provider/client.rs index f179937e..4976e593 100644 --- a/app/src/ai/provider/client.rs +++ b/app/src/ai/provider/client.rs @@ -1,4 +1,5 @@ use std::sync::{Arc, Mutex}; +use std::time::SystemTime; use anyhow::Result; use aws_config::BehaviorVersion; @@ -37,40 +38,17 @@ pub struct BedrockClientConfig { pub access_key_id: String, pub secret_access_key: String, pub session_token: Option, + pub credentials_expiry: Option, pub cross_region_inference: bool, } -impl BedrockClientConfig { - /// Applies external config (from Claude Code / OpenCode) as fallback values - /// when Galaxy's own settings are at their defaults. - pub fn with_external_fallbacks(mut self) -> Self { - let external = ExternalBedrockConfig::load(); - if external.is_empty() { - return self; - } - - if self.profile == "default" { - if let Some(profile) = external.profile { - log::info!("[bedrock] Using profile from external config: {profile}"); - self.profile = profile; - } - } - - if self.region.is_empty() { - if let Some(region) = external.region { - log::info!("[bedrock] Using region from external config: {region}"); - self.region = region; - } - } - - self - } -} #[derive(Debug, thiserror::Error)] pub enum BedrockError { #[error("Bedrock credentials not configured")] CredentialsNotConfigured, + #[error("Failed to load AWS credentials: {0}")] + CredentialsLoadFailed(String), #[error("Bedrock region not configured and could not be auto-detected")] RegionNotConfigured, #[error("Bedrock API error: {0}")] @@ -115,12 +93,20 @@ impl BedrockClient { if config.access_key_id.is_empty() || config.secret_access_key.is_empty() { return Err(BedrockError::CredentialsNotConfigured); } + if !super::authentication::credentials_are_current( + config.credentials_expiry, + SystemTime::now(), + ) { + return Err(BedrockError::CredentialsLoadFailed( + "AWS credentials have expired".to_owned(), + )); + } let creds = aws_credential_types::Credentials::new( &config.access_key_id, &config.secret_access_key, config.session_token, - None, + config.credentials_expiry, "warp-bedrock-static", ); @@ -148,11 +134,15 @@ impl BedrockClient { ); } Err(e) => { - log::warn!("[bedrock] Failed to resolve AWS credentials from provider: {e:?}"); + let message = crate::ai::aws_credentials::user_facing_aws_credentials_error_message( + &e, + &config.profile, + ); + return Err(BedrockError::CredentialsLoadFailed(message)); } } } else { - log::warn!("[bedrock] No credentials provider found in resolved AWS config"); + return Err(BedrockError::CredentialsNotConfigured); } let region = aws_config diff --git a/app/src/ai/provider/discovery.rs b/app/src/ai/provider/discovery.rs index 3dd6263f..2fac4e1b 100644 --- a/app/src/ai/provider/discovery.rs +++ b/app/src/ai/provider/discovery.rs @@ -161,7 +161,7 @@ async fn load_aws_config( &config.access_key_id, &config.secret_access_key, config.session_token.clone(), - None, + config.credentials_expiry, "galaxy-bedrock-discovery", ); let mut loader = diff --git a/app/src/ai/provider/e2e_tests.rs b/app/src/ai/provider/e2e_tests.rs index aedb3563..434739dc 100644 --- a/app/src/ai/provider/e2e_tests.rs +++ b/app/src/ai/provider/e2e_tests.rs @@ -268,6 +268,7 @@ fn get_test_config() -> Option { access_key_id: String::new(), secret_access_key: String::new(), session_token: None, + credentials_expiry: None, cross_region_inference: false, }) } diff --git a/app/src/ai/provider/integration_tests.rs b/app/src/ai/provider/integration_tests.rs index 155b3424..5a52394a 100644 --- a/app/src/ai/provider/integration_tests.rs +++ b/app/src/ai/provider/integration_tests.rs @@ -23,6 +23,7 @@ fn get_test_config() -> Option { access_key_id: String::new(), secret_access_key: String::new(), session_token: None, + credentials_expiry: None, cross_region_inference: false, }) } diff --git a/app/src/ai/provider/mod.rs b/app/src/ai/provider/mod.rs index 253c09db..e7394c83 100644 --- a/app/src/ai/provider/mod.rs +++ b/app/src/ai/provider/mod.rs @@ -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; diff --git a/app/src/settings_view/provider_setup_view.rs b/app/src/settings_view/provider_setup_view.rs index 4e8278b4..7d089565 100644 --- a/app/src/settings_view/provider_setup_view.rs +++ b/app/src/settings_view/provider_setup_view.rs @@ -1152,6 +1152,7 @@ impl ProviderSetupView { access_key_id: self.draft_bedrock.access_key_id.clone(), secret_access_key: self.draft_bedrock.secret_access_key.clone(), session_token: None, + credentials_expiry: None, cross_region_inference: self.draft_bedrock.cross_region_inference, }; let requested_config = config.clone(); diff --git a/app/src/terminal/view.rs b/app/src/terminal/view.rs index a513b994..4610cb8e 100644 --- a/app/src/terminal/view.rs +++ b/app/src/terminal/view.rs @@ -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) { - 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, + manual: bool, + ctx: &mut ViewContext, + ) { + 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(); diff --git a/crates/galaxy_agent_rig/src/stream.rs b/crates/galaxy_agent_rig/src/stream.rs index 3379040c..bae1df20 100644 --- a/crates/galaxy_agent_rig/src/stream.rs +++ b/crates/galaxy_agent_rig/src/stream.rs @@ -512,6 +512,11 @@ fn text_indicates_authentication_failure(text: &str) -> bool { "unrecognized client", "security token included in the request is invalid", "unable to locate credentials", + "unauthorizedexception", + "session token not found or invalid", + "sso session has expired", + "failed to load aws credentials", + "bedrock credentials not configured", "credential should be scoped to a valid region", ] .iter() diff --git a/crates/galaxy_agent_rig/src/stream_unit_tests.rs b/crates/galaxy_agent_rig/src/stream_unit_tests.rs index ac737561..f103a275 100644 --- a/crates/galaxy_agent_rig/src/stream_unit_tests.rs +++ b/crates/galaxy_agent_rig/src/stream_unit_tests.rs @@ -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); + } +} From 1820561f2afd0694bd4c4104d4ae332b5e160fae Mon Sep 17 00:00:00 2001 From: Josh Woodcock Date: Wed, 14 Oct 2026 13:51:35 -0500 Subject: [PATCH 02/12] remove old okta login authentication refresh --- AGENTS.md | 2 +- .../blocklist/controller/response_stream.rs | 4 +- app/src/ai/crosscheck/reviewer.rs | 1 + app/src/ai/llms.rs | 85 ++++---- app/src/ai/llms_tests.rs | 48 +++++ app/src/ai/provider/authentication.rs | 17 ++ app/src/ai/provider/authentication_tests.rs | 17 ++ app/src/ai/provider/client.rs | 9 +- app/src/ai/provider/convert.rs | 16 ++ app/src/ai/provider/convert_tests.rs | 12 ++ app/src/ai/provider/discovery.rs | 193 +++++++++++++++--- app/src/ai/provider/discovery_tests.rs | 140 +++++++++++++ app/src/ai/provider/e2e_tests.rs | 1 + app/src/ai/provider/external_config.rs | 1 + app/src/ai/provider/integration_tests.rs | 1 + app/src/ai/provider/models.rs | 8 + app/src/ai/runtime/rig.rs | 7 +- app/src/settings/ai.rs | 24 +++ app/src/settings/ai_tests.rs | 17 ++ app/src/settings_view/provider_setup_view.rs | 52 +++-- 20 files changed, 566 insertions(+), 89 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index a86543e9..ee8875b3 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -143,7 +143,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 diff --git a/app/src/ai/blocklist/controller/response_stream.rs b/app/src/ai/blocklist/controller/response_stream.rs index 7e7340c1..ea77f7d9 100644 --- a/app/src/ai/blocklist/controller/response_stream.rs +++ b/app/src/ai/blocklist/controller/response_stream.rs @@ -188,7 +188,9 @@ impl ResponseStream { // Fall back to Bedrock if *settings.bedrock_enabled.value() { let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx); - let config = BedrockClientConfig::from_settings(settings).with_managed_credentials( + 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(), ); diff --git a/app/src/ai/crosscheck/reviewer.rs b/app/src/ai/crosscheck/reviewer.rs index 73ebb556..c1aa4d82 100644 --- a/app/src/ai/crosscheck/reviewer.rs +++ b/app/src/ai/crosscheck/reviewer.rs @@ -183,6 +183,7 @@ impl CrosscheckReviewer { if *settings.bedrock_enabled.value() { let api_key_manager = ::ai::api_keys::ApiKeyManager::as_ref(ctx); 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(), diff --git a/app/src/ai/llms.rs b/app/src/ai/llms.rs index 8779a857..5fb6c1a6 100644 --- a/app/src/ai/llms.rs +++ b/app/src/ai/llms.rs @@ -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, #[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, - ) { + ) -> 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; } } diff --git a/app/src/ai/llms_tests.rs b/app/src/ai/llms_tests.rs index e8546ff8..c0a83ce8 100644 --- a/app/src/ai/llms_tests.rs +++ b/app/src/ai/llms_tests.rs @@ -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(), diff --git a/app/src/ai/provider/authentication.rs b/app/src/ai/provider/authentication.rs index 6431588e..6f555ef3 100644 --- a/app/src/ai/provider/authentication.rs +++ b/app/src/ai/provider/authentication.rs @@ -22,10 +22,27 @@ impl BedrockClientConfig { 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, diff --git a/app/src/ai/provider/authentication_tests.rs b/app/src/ai/provider/authentication_tests.rs index 3b7d2ab1..7b463707 100644 --- a/app/src/ai/provider/authentication_tests.rs +++ b/app/src/ai/provider/authentication_tests.rs @@ -13,6 +13,7 @@ fn config(auth_method: BedrockAuthMethod) -> BedrockClientConfig { secret_access_key: "configured-secret".to_owned(), session_token: None, credentials_expiry: None, + invocation_model_id: None, cross_region_inference: true, } } @@ -144,6 +145,22 @@ fn credential_expiry_is_checked_at_the_boundary() { 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) } diff --git a/app/src/ai/provider/client.rs b/app/src/ai/provider/client.rs index 4976e593..82fe4266 100644 --- a/app/src/ai/provider/client.rs +++ b/app/src/ai/provider/client.rs @@ -28,6 +28,7 @@ fn strip_context_marker(model_id: &str) -> String { pub struct BedrockClient { runtime_client: BedrockRuntimeClient, region: String, + invocation_model_id: Option, } #[derive(Debug, Clone, PartialEq)] @@ -39,6 +40,7 @@ pub struct BedrockClientConfig { pub secret_access_key: String, pub session_token: Option, pub credentials_expiry: Option, + pub invocation_model_id: Option, pub cross_region_inference: bool, } @@ -155,6 +157,7 @@ impl BedrockClient { Ok(Self { runtime_client, region, + invocation_model_id: config.invocation_model_id, }) } @@ -185,14 +188,16 @@ impl BedrockClient { tool_result_archive: Vec, ) -> Result { 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)"); diff --git a/app/src/ai/provider/convert.rs b/app/src/ai/provider/convert.rs index dd3b86b8..b57e3af3 100644 --- a/app/src/ai/provider/convert.rs +++ b/app/src/ai/provider/convert.rs @@ -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 { diff --git a/app/src/ai/provider/convert_tests.rs b/app/src/ai/provider/convert_tests.rs index 25011584..b44b8cb7 100644 --- a/app/src/ai/provider/convert_tests.rs +++ b/app/src/ai/provider/convert_tests.rs @@ -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; diff --git a/app/src/ai/provider/discovery.rs b/app/src/ai/provider/discovery.rs index 2fac4e1b..86cae58f 100644 --- a/app/src/ai/provider/discovery.rs +++ b/app/src/ai/provider/discovery.rs @@ -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 = 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::>() .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::>(); 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(mut fetch: F) -> Result, String> +where + F: FnMut(Option) -> Fut, + Fut: Future>, +{ + 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, + profiles: &[InferenceProfileSummary], + cross_region: bool, +) -> Vec { + let mut routes = profiles.iter() + .filter(|profile| profile.status() == &InferenceProfileStatus::Active) + .collect::>(); + 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 diff --git a/app/src/ai/provider/discovery_tests.rs b/app/src/ai/provider/discovery_tests.rs index 560f5f45..1c9ed080 100644 --- a/app/src/ai/provider/discovery_tests.rs +++ b/app/src/ai/provider/discovery_tests.rs @@ -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()); diff --git a/app/src/ai/provider/e2e_tests.rs b/app/src/ai/provider/e2e_tests.rs index 434739dc..51019c74 100644 --- a/app/src/ai/provider/e2e_tests.rs +++ b/app/src/ai/provider/e2e_tests.rs @@ -269,6 +269,7 @@ fn get_test_config() -> Option { secret_access_key: String::new(), session_token: None, credentials_expiry: None, + invocation_model_id: None, cross_region_inference: false, }) } diff --git a/app/src/ai/provider/external_config.rs b/app/src/ai/provider/external_config.rs index 71bc8145..d9b5b2fa 100644 --- a/app/src/ai/provider/external_config.rs +++ b/app/src/ai/provider/external_config.rs @@ -146,6 +146,7 @@ fn parse_claude_code_model_map( model_id: arn, display_name, vision_supported: true, + ..Default::default() } }) .collect() diff --git a/app/src/ai/provider/integration_tests.rs b/app/src/ai/provider/integration_tests.rs index 5a52394a..82ff8c69 100644 --- a/app/src/ai/provider/integration_tests.rs +++ b/app/src/ai/provider/integration_tests.rs @@ -24,6 +24,7 @@ fn get_test_config() -> Option { secret_access_key: String::new(), session_token: None, credentials_expiry: None, + invocation_model_id: None, cross_region_inference: false, }) } diff --git a/app/src/ai/provider/models.rs b/app/src/ai/provider/models.rs index 9c9c69be..df3d777c 100644 --- a/app/src/ai/provider/models.rs +++ b/app/src/ai/provider/models.rs @@ -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(); diff --git a/app/src/ai/runtime/rig.rs b/app/src/ai/runtime/rig.rs index 85140b4a..61172595 100644 --- a/app/src/ai/runtime/rig.rs +++ b/app/src/ai/runtime/rig.rs @@ -346,14 +346,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, diff --git a/app/src/settings/ai.rs b/app/src/settings/ai.rs index be2b2c1e..04474c65 100644 --- a/app/src/settings/ai.rs +++ b/app/src/settings/ai.rs @@ -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, + #[serde(default)] + pub runtime_supported: Option, } 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 } diff --git a/app/src/settings/ai_tests.rs b/app/src/settings/ai_tests.rs index f98f61fc..97cd1775 100644 --- a/app/src/settings/ai_tests.rs +++ b/app/src/settings/ai_tests.rs @@ -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, ) diff --git a/app/src/settings_view/provider_setup_view.rs b/app/src/settings_view/provider_setup_view.rs index 7d089565..9ae55646 100644 --- a/app/src/settings_view/provider_setup_view.rs +++ b/app/src/settings_view/provider_setup_view.rs @@ -1130,6 +1130,20 @@ 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.cancel_acp_setup(); self.step = ProviderSetupStep::Discover; @@ -1145,34 +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, - credentials_expiry: 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, @@ -1190,6 +1215,7 @@ impl ProviderSetupView { me.update_next_button(ctx); ctx.notify(); } + } }, ); return; From 3ddd8f25aca986900656105039da3fffdc51382f Mon Sep 17 00:00:00 2001 From: Josh Woodcock Date: Mon, 19 Oct 2026 12:51:53 -0500 Subject: [PATCH 03/12] Add Bedrock model refresh and fix auto-region discovery race Co-Authored-By: Warp --- app/src/settings_view/ai_page.rs | 17 ++++++++++++++++- app/src/settings_view/provider_setup_view.rs | 5 ++--- 2 files changed, 18 insertions(+), 4 deletions(-) diff --git a/app/src/settings_view/ai_page.rs b/app/src/settings_view/ai_page.rs index 07f647f2..2233f8c5 100644 --- a/app/src/settings_view/ai_page.rs +++ b/app/src/settings_view/ai_page.rs @@ -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, acp_add_button: ViewHandle, bedrock_edit_button: ViewHandle, + bedrock_refresh_button: ViewHandle, bedrock_remove_button: ViewHandle, provider_cards: Vec, acp_provider_cards: Vec, @@ -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(), diff --git a/app/src/settings_view/provider_setup_view.rs b/app/src/settings_view/provider_setup_view.rs index 9ae55646..179ed0c7 100644 --- a/app/src/settings_view/provider_setup_view.rs +++ b/app/src/settings_view/provider_setup_view.rs @@ -687,10 +687,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); @@ -1146,6 +1144,7 @@ impl ProviderSetupView { fn begin_discovery(&mut self, ctx: &mut ViewContext) { self.cancel_acp_setup(); + self.bedrock_region_request += 1; self.step = ProviderSetupStep::Discover; self.discovery_state = DiscoveryState::Loading; self.acp_setup_status.clear(); From c894a11589f76a7e42a6e5dbad60e8e79d935f51 Mon Sep 17 00:00:00 2001 From: Josh Woodcock Date: Mon, 19 Oct 2026 12:53:11 -0500 Subject: [PATCH 04/12] Update ChatGPT authentication and sign-in recovery Co-Authored-By: Warp --- app/src/ai/chatgpt_auth.rs | 580 ++++++++++++------- app/src/ai/chatgpt_auth_tests.rs | 171 ++++++ app/src/ai/crosscheck/reviewer.rs | 3 +- app/src/ai/mod.rs | 2 +- app/src/ai/runtime/rig.rs | 3 +- app/src/settings_view/provider_setup_view.rs | 53 +- crates/galaxy_agent_rig/src/chatgpt.rs | 74 ++- crates/galaxy_agent_rig/src/chatgpt_tests.rs | 38 +- 8 files changed, 660 insertions(+), 264 deletions(-) create mode 100644 app/src/ai/chatgpt_auth_tests.rs diff --git a/app/src/ai/chatgpt_auth.rs b/app/src/ai/chatgpt_auth.rs index ce5a746b..f82300d0 100644 --- a/app/src/ai/chatgpt_auth.rs +++ b/app/src/ai/chatgpt_auth.rs @@ -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>; +static RESOLVER: OnceLock>>> = 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, + refresh: RefreshCoordinator, + sign_in_generation: u64, + sign_in_cancel: Option>, } 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 { + 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::(&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) { + fn begin_sign_in(&mut self) -> Option { 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) { + 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.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) { + if self.state != state { + self.state = state; + ctx.emit(ChatGPTAuthModelEvent::StateChanged); + } + } + + pub(crate) fn cancel_sign_in(&mut self, ctx: &mut ModelContext) { + 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) { + 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, ) -> Result<(String, oneshot::Receiver>), 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 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, -} +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 { - 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::(&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 { - 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)>, +} - 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 { + 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 { - 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, + persist: impl FnOnce(&AuthRecord) -> Result<(), String>, + ) -> Option> { + 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) -> 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 { 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 { + 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 { + 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 { + 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::(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 { 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 { - #[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 { @@ -633,7 +727,7 @@ struct TokenResponse { expires_in: Option, } -#[derive(serde::Deserialize, serde::Serialize)] +#[derive(Clone, Default, serde::Deserialize, serde::Serialize)] struct AuthRecord { access_token: Option, refresh_token: Option, @@ -641,3 +735,41 @@ struct AuthRecord { expires_at: Option, account_id: Option, } + +impl std::fmt::Debug for AuthRecord { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ChatGPTAuthRecord") + .field("credentials", &"") + .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 { + 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 { + 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; diff --git a/app/src/ai/chatgpt_auth_tests.rs b/app/src/ai/chatgpt_auth_tests.rs new file mode 100644 index 00000000..280e69be --- /dev/null +++ b/app/src/ai/chatgpt_auth_tests.rs @@ -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::>(); + 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()); +} diff --git a/app/src/ai/crosscheck/reviewer.rs b/app/src/ai/crosscheck/reviewer.rs index c1aa4d82..59f6ded1 100644 --- a/app/src/ai/crosscheck/reviewer.rs +++ b/app/src/ai/crosscheck/reviewer.rs @@ -215,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(), diff --git a/app/src/ai/mod.rs b/app/src/ai/mod.rs index c956cd27..23d8236a 100644 --- a/app/src/ai/mod.rs +++ b/app/src/ai/mod.rs @@ -90,7 +90,7 @@ pub fn init(app: &mut AppContext) { #[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); diff --git a/app/src/ai/runtime/rig.rs b/app/src/ai/runtime/rig.rs index 61172595..c4593a09 100644 --- a/app/src/ai/runtime/rig.rs +++ b/app/src/ai/runtime/rig.rs @@ -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); diff --git a/app/src/settings_view/provider_setup_view.rs b/app/src/settings_view/provider_setup_view.rs index 179ed0c7..ef1cb90e 100644 --- a/app/src/settings_view/provider_setup_view.rs +++ b/app/src/settings_view/provider_setup_view.rs @@ -220,6 +220,7 @@ pub enum ProviderSetupViewAction { ToggleModel(usize), CycleModelCapability(usize, CapabilityKey), ConnectChatGPT, + CancelChatGPTSignIn, ConnectClaude, SignOutClaude, CancelClaudeSignIn, @@ -1590,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() @@ -1601,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(), ); @@ -1861,7 +1873,7 @@ impl ProviderSetupView { .finish() } - fn render_discovery(&self, appearance: &Appearance) -> Box { + fn render_discovery(&self, appearance: &Appearance, app: &AppContext) -> Box { let (message, color) = match &self.discovery_state { DiscoveryState::Idle | DiscoveryState::Loading => ( if self.provider_type == ProviderSetupProviderType::Acp @@ -1875,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( @@ -1913,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( @@ -2468,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() @@ -2488,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; @@ -2588,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); } @@ -2633,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); diff --git a/crates/galaxy_agent_rig/src/chatgpt.rs b/crates/galaxy_agent_rig/src/chatgpt.rs index f4af1f7d..6036a266 100644 --- a/crates/galaxy_agent_rig/src/chatgpt.rs +++ b/crates/galaxy_agent_rig/src/chatgpt.rs @@ -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, } +impl std::fmt::Debug for ChatGPTSubscriptionCredentials { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ChatGPTSubscriptionCredentials") + .field("access_token", &"") + .field("account_id", &self.account_id.as_ref().map(|_| "")) + .finish() + } +} +pub type ChatGPTCredentialResolver = Arc< + dyn Fn() -> BoxFuture<'static, Result> + Send + Sync, +>; + #[derive(Clone, Debug, PartialEq, Eq)] pub struct ChatGPTSubscriptionRuntimeConfig { pub model: String, pub reasoning_effort: Option, pub max_output_tokens: Option, - pub auth_file: Option, } -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct ChatGPTSubscriptionRuntime { config: ChatGPTSubscriptionRuntimeConfig, descriptor: RuntimeDescriptor, + credential_resolver: Option, +} + +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 { @@ -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 { + 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 { - 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); diff --git a/crates/galaxy_agent_rig/src/chatgpt_tests.rs b/crates/galaxy_agent_rig/src/chatgpt_tests.rs index c176c402..192dce22 100644 --- a/crates/galaxy_agent_rig/src/chatgpt_tests.rs +++ b/crates/galaxy_agent_rig/src/chatgpt_tests.rs @@ -16,7 +16,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 +24,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( From 75b9334a0b4bb0d9c313976fbbcfe9d92a7e05a2 Mon Sep 17 00:00:00 2001 From: Ryan Ward Date: Mon, 19 Oct 2026 19:41:37 +0000 Subject: [PATCH 05/12] Handle ChatGPT subscription metadata without aborting streams --- AGENTS.md | 6 +- crates/galaxy_agent_rig/src/chatgpt_tests.rs | 73 +++++++++++++++++++- crates/galaxy_agent_rig/src/stream.rs | 7 ++ 3 files changed, 82 insertions(+), 4 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index ee8875b3..c7bd8c01 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -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 diff --git a/crates/galaxy_agent_rig/src/chatgpt_tests.rs b/crates/galaxy_agent_rig/src/chatgpt_tests.rs index 192dce22..f9371a2a 100644 --- a/crates/galaxy_agent_rig/src/chatgpt_tests.rs +++ b/crates/galaxy_agent_rig/src/chatgpt_tests.rs @@ -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] @@ -151,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)); +} diff --git a/crates/galaxy_agent_rig/src/stream.rs b/crates/galaxy_agent_rig/src/stream.rs index bae1df20..c67cb89e 100644 --- a/crates/galaxy_agent_rig/src/stream.rs +++ b/crates/galaxy_agent_rig/src/stream.rs @@ -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 { From 0383166f4ecd42b4d76f298c173080c14ea1159b Mon Sep 17 00:00:00 2001 From: Ryan Ward Date: Tue, 20 Oct 2026 18:27:34 +0000 Subject: [PATCH 06/12] Use the active model context limit for provider compaction --- app/src/ai/blocklist/controller.rs | 5 +- .../ai/runtime/provider_run_coordinator.rs | 9 ++++ .../runtime/provider_run_coordinator_tests.rs | 50 +++++++++++++++++++ 3 files changed, 62 insertions(+), 2 deletions(-) diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 7a23ffd3..bfafeb5a 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -6935,8 +6935,9 @@ impl BlocklistAIController { } let active_context_limit = active_run - .response_config - .max_context_tokens + .coordinator + .current_context_limit() + .or(active_run.response_config.max_context_tokens) .unwrap_or_else(|| { let preferences = LLMPreferences::as_ref(ctx); configured_llm_context_limit( diff --git a/app/src/ai/runtime/provider_run_coordinator.rs b/app/src/ai/runtime/provider_run_coordinator.rs index 435b4ac4..391e2138 100644 --- a/app/src/ai/runtime/provider_run_coordinator.rs +++ b/app/src/ai/runtime/provider_run_coordinator.rs @@ -251,6 +251,15 @@ impl ProviderRunCoordinator { self.profiles.get(profile).map(|profile| &profile.request) } + /// The pending model profile owns the context budget, including after a + /// preflight overflow has returned the run to its compaction boundary. + pub(crate) fn current_context_limit(&self) -> Option { + self.profiles + .get(self.run.profile().as_str()) + .and_then(|profile| profile.max_context_tokens) + .or(self.max_context_tokens) + } + /// Includes the current transcript, fixed request context, tools, and output reservation. /// Read at the model boundary so newly committed tool output is included. pub(crate) fn estimated_context_tokens(&self) -> Option<(u32, u32)> { diff --git a/app/src/ai/runtime/provider_run_coordinator_tests.rs b/app/src/ai/runtime/provider_run_coordinator_tests.rs index 88003166..fa685f81 100644 --- a/app/src/ai/runtime/provider_run_coordinator_tests.rs +++ b/app/src/ai/runtime/provider_run_coordinator_tests.rs @@ -1558,6 +1558,56 @@ fn context_estimate_tracks_live_transcript_and_reserves_fixed_context() { assert_eq!(fixed_before, fixed_after); } +#[tokio::test] +async fn compaction_boundary_keeps_the_monitor_context_limit() { + let base_runtime = Arc::new(ScriptedRuntime::new(vec![])); + let cli_runtime = Arc::new(ScriptedRuntime::new(vec![])); + let run = ProviderRun::new( + "monitor-compaction-budget", + vec![ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text("history".repeat(5_000)), + }], + CLI_MONITOR_PROVIDER_PROFILE, + ProviderRunLimits::default(), + ); + let mut profiles = BTreeMap::new(); + profiles.insert( + BASE_PROVIDER_PROFILE.to_string(), + ProviderRunProfile::new(base_runtime.clone(), request()), + ); + profiles.insert( + CLI_MONITOR_PROVIDER_PROFILE.to_string(), + ProviderRunProfile::new(cli_runtime.clone(), request()) + .with_max_context_tokens(Some(8_000)), + ); + let mut coordinator = ProviderRunCoordinator::new(run, profiles).unwrap(); + coordinator.set_max_context_tokens(Some(100_000)); + let (_sender, control) = turn_control(); + + let block = coordinator + .drive_until_blocked(control, |_| Ok(())) + .await + .unwrap(); + + assert_eq!(block, ProviderRunBlock::ContextWindowExceeded); + assert_eq!(coordinator.run().state(), &ProviderRunState::ReadyToCallModel); + assert_eq!(coordinator.current_context_limit(), Some(8_000)); + assert!(base_runtime.requests().is_empty()); + assert!(cli_runtime.requests().is_empty()); + + let work_id = coordinator.run().ready_work_id().unwrap(); + coordinator + .run_mut() + .continue_ready_with_observation( + &work_id, + MessageContent::Text("Command finished".to_string()), + BASE_PROVIDER_PROFILE, + ) + .unwrap(); + assert_eq!(coordinator.current_context_limit(), Some(100_000)); +} + #[tokio::test] async fn monitor_request_uses_its_profile_context_limit() { let base_runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()])); From cb830e1ba76b8da16e07ce5fdda376dbb62f9dde Mon Sep 17 00:00:00 2001 From: Ryan Ward Date: Tue, 20 Oct 2026 21:05:57 +0000 Subject: [PATCH 07/12] Run progressive compaction alongside foreground provider work --- AGENTS.md | 1 + app/src/ai/blocklist/controller.rs | 307 ++++++++++++-- app/src/ai/blocklist/controller_tests.rs | 387 +++++++++++++++++- .../ai/runtime/provider_run_coordinator.rs | 18 + .../runtime/provider_run_coordinator_tests.rs | 112 +++++ 5 files changed, 794 insertions(+), 31 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index c7bd8c01..df151ec7 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -219,6 +219,7 @@ Key invariants: - Bedrock prompt caching uses three cache points: system prompt, second-to-last history message, and tool configuration - `ensure_tool_results_paired()` enforces Bedrock's invariant that every `tool_use` has a matching `tool_result` - Progressive summaries are prepended to provider requests as a user/assistant pair; direct-provider runs compact their live transcript at a checkpointed `ReadyToCallModel` boundary before another model call, while background summarization remains independent for session-owned runtimes +- Direct-provider progressive compaction runs in a separate request alongside foreground work. Start early against the smaller usable history budget of the current and terminal-use profiles, including each profile's prompt/tools/output reservation. Apply a result only at a safe model boundary after validating its exact snapshotted prefix, preserving all newer history and recomputing live context accounting. Wait for pending compaction only after context exhaustion; cancellation and terminal cleanup must remain immediate. Cleanly completed conversations may accept a late summary while idle; stale, cancelled, or superseded generations must never replace a newer run's history. - Loop prevention in `controller.rs` detects repeated tool failures (3+ identical) and injects a corrective instruction - Direct-provider long-running shell follow-ups retain CLI tasks as UI projections while command monitoring and completion stay owned by the same provider run - A direct-provider command completion is only queued when the terminal reports it; the CLI task remains active until the provider run applies that completion at a safe boundary and deactivates it diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index bfafeb5a..b1a72b83 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -158,14 +158,24 @@ struct ActiveProviderRunIdentity { enum ActiveProviderProgressiveSummaryState { InFlight { identity: ActiveProviderRunIdentity, + waiting_for_context: bool, }, Ready { identity: ActiveProviderRunIdentity, + waiting_for_context: bool, plan: ActiveProgressiveSummaryPlan, result: anyhow::Result<(String, Usage, String)>, }, } +impl ActiveProviderProgressiveSummaryState { + fn identity(&self) -> &ActiveProviderRunIdentity { + match self { + Self::InFlight { identity, .. } | Self::Ready { identity, .. } => identity, + } + } +} + fn configured_llm_context_limit(info: &LLMInfo) -> u32 { [ info.context_window.default_max, @@ -6955,16 +6965,21 @@ impl BlocklistAIController { conversation.has_pending_progressive_summary() }; let transcript = active_run.coordinator.run().transcript(); - let Some((current_context_tokens, persistent_context_tokens)) = + let Some((current_context_tokens, current_persistent_context_tokens)) = self.active_provider_context_estimate(conversation_id, ctx) else { return false; }; + let (summary_context_limit, summary_persistent_context_tokens) = active_run + .coordinator + .progressive_summary_context_budget() + .unwrap_or((active_context_limit, current_persistent_context_tokens)); // Trigger and retain against the space history can actually use. Fixed // prompts and a large output reservation cannot be reduced by summarizing. - let history_capacity = active_context_limit.saturating_sub(persistent_context_tokens); + let history_capacity = + summary_context_limit.saturating_sub(summary_persistent_context_tokens); let current_history_tokens = - current_context_tokens.saturating_sub(persistent_context_tokens); + current_context_tokens.saturating_sub(current_persistent_context_tokens); if summary_pending || last_failure_context_tokens == Some(current_context_tokens) || (!force @@ -7030,7 +7045,7 @@ impl BlocklistAIController { retained_messages_count: transcript.len().saturating_sub(split_point), trigger_context_tokens: current_context_tokens, active_context_limit, - persistent_context_tokens, + persistent_context_tokens: current_persistent_context_tokens, candidates, summarized_prefix: transcript[..split_point].to_vec(), summarize_messages: vec![ConversationMessage { @@ -7051,6 +7066,7 @@ impl BlocklistAIController { conversation_id, ActiveProviderProgressiveSummaryState::InFlight { identity: identity.clone(), + waiting_for_context: force, }, ); BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, _| { @@ -7121,27 +7137,23 @@ impl BlocklistAIController { ctx: &mut ModelContext, ) { let conversation_id = identity.conversation_id; - let owns_summary = self + let waiting_for_context = match self .active_provider_progressive_summaries .get(&conversation_id) - .is_some_and(|state| { - matches!( - state, - ActiveProviderProgressiveSummaryState::InFlight { - identity: active_identity, - } if active_identity == &identity - ) - }); - if !owns_summary { - return; - } + { + Some(ActiveProviderProgressiveSummaryState::InFlight { + identity: active_identity, + waiting_for_context, + }) if active_identity == &identity => *waiting_for_context, + _ => return, + }; let identity_is_current = self.active_provider_runs .get(&conversation_id) .is_some_and(|slot| { slot.stream_id == identity.stream_id && slot.run_id == identity.run_id }); - if !identity_is_current { + if !identity_is_current && self.active_provider_runs.contains_key(&conversation_id) { self.active_provider_progressive_summaries .remove(&conversation_id); BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, _| { @@ -7156,11 +7168,29 @@ impl BlocklistAIController { conversation_id, ActiveProviderProgressiveSummaryState::Ready { identity, + waiting_for_context, plan, result, }, ); - self.drive_active_provider_run(conversation_id, ctx); + if !self.active_provider_runs.contains_key(&conversation_id) { + self.apply_ready_idle_provider_progressive_summary(conversation_id, ctx); + return; + } + // A background result must not re-enter an active model call or dispatch + // an awaiting tool batch again. Normal work consumes it at its next boundary. + if waiting_for_context + && self.active_provider_runs.get(&conversation_id).is_some_and(|slot| { + slot.run.as_ref().is_some_and(|run| { + matches!( + run.coordinator.run().state(), + ProviderRunState::ReadyToCallModel + ) + }) + }) + { + self.drive_active_provider_run(conversation_id, ctx); + } } fn apply_ready_active_provider_progressive_summary( @@ -7193,6 +7223,7 @@ impl BlocklistAIController { let Some(ActiveProviderProgressiveSummaryState::Ready { identity: _, + waiting_for_context, plan, result, }) = self @@ -7209,6 +7240,7 @@ impl BlocklistAIController { let mut summary_applied = false; let mut recorded_usage = None; + let mut failure = None; match result { Ok((summary, usage, summarizer_model_id)) if estimated_history_tokens(&progressive_summary_messages(summary.clone())) @@ -7218,6 +7250,7 @@ impl BlocklistAIController { "[progressive-summary] Summary did not reduce context for conversation {:?}", conversation_id, ); + failure = Some("summary did not reduce context".to_string()); recorded_usage = Some((usage, summarizer_model_id)); } Ok((summary, usage, summarizer_model_id)) => { @@ -7234,6 +7267,9 @@ impl BlocklistAIController { "[progressive-summary] Discarding stale background summary for active conversation {:?}; the snapshotted prefix no longer matches", conversation_id, ); + failure = Some( + "summary prefix no longer matches conversation history".to_string(), + ); } else { let drained = transcript[..plan.split_point].to_vec(); let retained = transcript[plan.split_point..].to_vec(); @@ -7241,6 +7277,14 @@ impl BlocklistAIController { let summary_prefix = progressive_summary_messages(summary.clone()); let summary_prefix_len = summary_prefix.len(); let summary_prefix_tokens = estimated_history_tokens(&summary_prefix); + let persistent_context_tokens = run + .coordinator + .estimated_context_tokens() + .map_or(plan.persistent_context_tokens, |(_, fixed)| fixed); + let active_context_limit = run + .coordinator + .current_context_limit() + .unwrap_or(plan.active_context_limit); if let Err(error) = run .coordinator .run_mut() @@ -7255,12 +7299,11 @@ impl BlocklistAIController { if let Ok(mut messages_sent) = run.messages_sent.lock() { *messages_sent = retained.clone(); } - let new_context_tokens = plan - .persistent_context_tokens + let new_context_tokens = persistent_context_tokens .saturating_add(summary_prefix_tokens) .saturating_add(estimated_history_tokens(&retained)); let new_usage = - new_context_tokens as f32 / plan.active_context_limit.max(1) as f32; + new_context_tokens as f32 / active_context_limit.max(1) as f32; BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, _| { let Some(conversation) = history.conversation_mut(&conversation_id) else { @@ -7278,7 +7321,7 @@ impl BlocklistAIController { summarizer_model_id, new_context_tokens, new_usage * 100.0, - plan.active_context_limit, + active_context_limit, retained_count, ); recorded_usage = Some((usage, summarizer_model_id)); @@ -7291,9 +7334,13 @@ impl BlocklistAIController { "[progressive-summary] Failed for active conversation {:?}: {error:#}", conversation_id, ); + failure = Some(error.to_string()); } } + let failure_context_tokens = self + .active_provider_context_estimate(conversation_id, ctx) + .map_or(plan.trigger_context_tokens, |(tokens, _)| tokens); if let Some(run) = self .active_provider_runs .get_mut(&conversation_id) @@ -7303,12 +7350,23 @@ impl BlocklistAIController { None } else { // Fail open and suppress another request at this unchanged context boundary. - Some(plan.trigger_context_tokens) + Some(failure_context_tokens) }; } if let Some((usage, summarizer_model_id)) = recorded_usage { self.record_progressive_summary_usage(conversation_id, usage, summarizer_model_id, ctx); } + if waiting_for_context && !summary_applied { + self.fail_active_provider_run( + conversation_id, + format!( + "context compaction could not recover enough space: {}", + failure.unwrap_or_else(|| "could not replace the summarized history".to_string()) + ), + ctx, + ); + return true; + } if let Err(error) = self.persist_active_provider_run(conversation_id, ctx) { self.fail_active_provider_run( conversation_id, @@ -7464,18 +7522,55 @@ impl BlocklistAIController { conversation_id: AIConversationId, ctx: &mut ModelContext, ) { + // A follow-up owns a new generation. Its history must never be blocked + // by, or replaced with, a detached summary from the previous generation. + // Cancellation also discards summaries before they can fail recovery. + let stale_summary = self + .active_provider_progressive_summaries + .get(&conversation_id) + .zip(self.active_provider_runs.get(&conversation_id)) + .is_some_and(|(summary, slot)| { + let identity = summary.identity(); + slot.stream_id != identity.stream_id || slot.run_id != identity.run_id + }); + let cancelling = self + .active_provider_runs + .get(&conversation_id) + .is_some_and(|slot| slot.cancellation_reason.is_some()); + if (stale_summary || cancelling) + && self + .active_provider_progressive_summaries + .remove(&conversation_id) + .is_some() + { + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, _| { + if let Some(conversation) = history.conversation_mut(&conversation_id) { + conversation.set_has_pending_progressive_summary(false); + } + }); + } if self.apply_ready_active_provider_progressive_summary(conversation_id, ctx) { return; } self.schedule_provider_command_output_poll(conversation_id, ctx); - self.begin_active_provider_progressive_summary(conversation_id, false, ctx); + // Ordinary background work never blocks the agent. Once a model has + // exhausted its context, polling must not relaunch that same request + // while its recovery summary is still running. Cancellation still drives. if matches!( - self.active_provider_progressive_summaries - .get(&conversation_id), - Some(ActiveProviderProgressiveSummaryState::InFlight { .. }) - ) { + self.active_provider_progressive_summaries.get(&conversation_id), + Some(ActiveProviderProgressiveSummaryState::InFlight { + waiting_for_context: true, + .. + }) + ) && self.active_provider_runs.get(&conversation_id).is_some_and(|slot| { + slot.cancellation_reason.is_none() + && slot.run.as_ref().is_some_and(|run| { + matches!(run.coordinator.run().state(), ProviderRunState::ReadyToCallModel) + }) + }) { return; } + self.begin_active_provider_progressive_summary(conversation_id, false, ctx); let Some(slot) = self.active_provider_runs.get_mut(&conversation_id) else { return; }; @@ -7977,6 +8072,24 @@ impl BlocklistAIController { ); return; } + match self.active_provider_progressive_summaries.get_mut(&conversation_id) { + Some(ActiveProviderProgressiveSummaryState::InFlight { + waiting_for_context, + .. + }) => { + *waiting_for_context = true; + return; + } + Some(ActiveProviderProgressiveSummaryState::Ready { + waiting_for_context, + .. + }) => { + *waiting_for_context = true; + self.drive_active_provider_run(conversation_id, ctx); + return; + } + None => {} + } if !self.begin_active_provider_progressive_summary(conversation_id, true, ctx) { self.fail_active_provider_run( conversation_id, @@ -8826,6 +8939,11 @@ impl BlocklistAIController { &outcome, ), ); + let completed_cleanly = matches!( + &outcome, + ProviderRunOutcome::Completed(completion) + if completion.stop_reason == StopReason::Completed + ); match outcome { ProviderRunOutcome::Completed(completion) => match completion.stop_reason { StopReason::Completed => { @@ -8890,7 +9008,140 @@ impl BlocklistAIController { } } } + // A background checkpoint remains useful after the answer finishes. Detach + // it while cleanup releases this generation, then retain it only if no + // successor has started or claimed its own summary in the meantime. + let pending_summary = if completed_cleanly { + self.active_provider_progressive_summaries + .remove(&conversation_id) + } else { + None + }; self.cleanup_active_provider_run(conversation_id, &stream_id, &response_stream, ctx); + if let Some(summary) = pending_summary { + let conversation_succeeded = BlocklistAIHistoryModel::as_ref(ctx) + .conversation_status(&conversation_id) + == Some(&ConversationStatus::Success); + if conversation_succeeded + && !self.active_provider_runs.contains_key(&conversation_id) + && !self + .active_provider_progressive_summaries + .contains_key(&conversation_id) + { + self.active_provider_progressive_summaries + .insert(conversation_id, summary); + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, _| { + if let Some(conversation) = history.conversation_mut(&conversation_id) { + conversation.set_has_pending_progressive_summary(true); + } + }); + self.apply_ready_idle_provider_progressive_summary(conversation_id, ctx); + } + } + } + + fn apply_ready_idle_provider_progressive_summary( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + if self.active_provider_runs.contains_key(&conversation_id) + || !matches!( + self.active_provider_progressive_summaries + .get(&conversation_id), + Some(ActiveProviderProgressiveSummaryState::Ready { .. }) + ) + { + return; + } + let Some(ActiveProviderProgressiveSummaryState::Ready { plan, result, .. }) = self + .active_provider_progressive_summaries + .remove(&conversation_id) + else { + return; + }; + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, _| { + if let Some(conversation) = history.conversation_mut(&conversation_id) { + conversation.set_has_pending_progressive_summary(false); + } + }); + if BlocklistAIHistoryModel::as_ref(ctx).conversation_status(&conversation_id) + != Some(&ConversationStatus::Success) + { + return; + } + + match result { + Ok((summary, usage, summarizer_model_id)) => { + let summary_prefix_tokens = + estimated_history_tokens(&progressive_summary_messages(summary.clone())); + let summary_reduces_context = + summary_prefix_tokens < estimated_history_tokens(&plan.summarized_prefix); + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| { + let Some(conversation) = history.conversation_mut(&conversation_id) else { + return; + }; + let mut transcript = conversation + .progressive_summary() + .map(|existing| progressive_summary_messages(existing.to_owned())) + .unwrap_or_default(); + transcript.extend_from_slice(conversation.bedrock_message_history()); + if transcript.get(..plan.split_point) + != Some(plan.summarized_prefix.as_slice()) + { + log::warn!( + "[progressive-summary] Discarding stale idle summary for conversation {:?}; the snapshotted prefix no longer matches", + conversation_id, + ); + return; + } + if !summary_reduces_context { + log::warn!( + "[progressive-summary] Idle summary did not reduce context for conversation {:?}", + conversation_id, + ); + return; + } + + let retained = transcript.split_off(plan.split_point); + let retained_count = retained.len(); + let new_context_tokens = plan + .persistent_context_tokens + .saturating_add(summary_prefix_tokens) + .saturating_add(estimated_history_tokens(&retained)); + let new_usage = + new_context_tokens as f32 / plan.active_context_limit.max(1) as f32; + conversation.archive_tool_results(transcript); + *conversation.bedrock_message_history_mut() = retained; + conversation.set_progressive_summary(Some(summary), 0); + conversation.set_current_context_tokens(new_context_tokens); + conversation.set_context_window_usage(new_usage.clamp(0.0, 1.0)); + conversation.write_updated_conversation_state(ctx); + log::info!( + "[progressive-summary] Applied background summary after active run completion {:?}: ~{} tokens ({:.1}% of {} context), {} tail messages preserved", + conversation_id, + new_context_tokens, + new_usage * 100.0, + plan.active_context_limit, + retained_count, + ); + }); + // This Ready result is consumed once, even when its prefix became + // stale or its response could not reduce the retained context. + self.record_progressive_summary_usage( + conversation_id, + usage, + summarizer_model_id, + ctx, + ); + } + Err(error) => { + log::error!( + "[progressive-summary] Failed after active conversation {:?} completed: {error:#}", + conversation_id, + ); + } + } } fn finalize_failed_provider_conversation( diff --git a/app/src/ai/blocklist/controller_tests.rs b/app/src/ai/blocklist/controller_tests.rs index 097e5005..e4de9412 100644 --- a/app/src/ai/blocklist/controller_tests.rs +++ b/app/src/ai/blocklist/controller_tests.rs @@ -4,16 +4,16 @@ use std::sync::{Arc, Mutex}; use ai::agent::action::{AskUserQuestionItem, AskUserQuestionType}; use chrono::Local; use galaxy_agent_core::{ - AgentError, AgentErrorKind, CompletedModelTurn, ContentPart, ConversationMessage, + turn_control, AgentError, AgentErrorKind, CompletedModelTurn, ContentPart, ConversationMessage, ExternalWorkId, MessageContent, MessageRole, PermissionKind, PermissionRequest, ProviderRun, ProviderRunFailure, ProviderRunFailureKind, ProviderRunId, ProviderRunLimits, ProviderRunOutcome, ProviderRunState, ProviderRunStep, RunEpoch, RuntimeCapabilities, - StopReason, ToolCall, TurnRequest, Usage, + StopReason, ToolCall, TurnCommand, TurnRequest, Usage, }; use galaxy_core::command::ExitCode; use uuid::Uuid; use warp_multi_agent_api::response_event; -use warpui::{App, EntityId, SingletonEntity}; +use warpui::{App, EntityId, ModelHandle, SingletonEntity}; use crate::ai::agent::conversation::{AIConversationId, ConversationStatus}; use crate::ai::agent::task::TaskId; @@ -3915,6 +3915,387 @@ fn optimistic_cli_subagent_completion_with_in_flight_stream_reports_success() { }); } +#[test] +fn progressive_summary_completion_preserves_an_in_flight_model_turn() { + App::test((), |mut app| async move { + initialize_app_for_terminal_view(&mut app); + let terminal = add_window_with_terminal(&mut app, None); + + terminal.update(&mut app, |terminal, ctx| { + let conversation_id = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| { + let conversation_id = + history.start_new_conversation(terminal.id(), false, false, false, ctx); + history + .conversation_mut(&conversation_id) + .unwrap() + .set_has_pending_progressive_summary(true); + conversation_id + }); + let stream_id = ResponseStreamId::new_for_test(); + let response_stream = + ctx.add_model(|_| ResponseStream::new_for_test(stream_id.clone())); + let mut slot = progressive_summary_test_slot(conversation_id, stream_id, response_stream); + let identity = super::ActiveProviderRunIdentity { + conversation_id, + stream_id: slot.stream_id.clone(), + run_id: slot.run_id.clone(), + }; + let (sender, foreground_control) = turn_control(); + slot.turn_control = Some(sender); + let prefix = vec![ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text("old project details".repeat(1_000)), + }]; + + terminal.ai_controller().update(ctx, |controller, ctx| { + controller.active_provider_runs.insert(conversation_id, slot); + controller.active_provider_progressive_summaries.insert( + conversation_id, + super::ActiveProviderProgressiveSummaryState::InFlight { + identity: identity.clone(), + waiting_for_context: true, + }, + ); + + controller.finish_active_provider_progressive_summary( + identity.clone(), + progressive_summary_test_plan(prefix), + Ok(( + "checkpoint".to_owned(), + Usage::default(), + "summary-model".to_owned(), + )), + ctx, + ); + + let active = &controller.active_provider_runs[&conversation_id]; + assert_eq!(active.stream_id, identity.stream_id); + assert_eq!(active.run_id, identity.run_id); + assert!(active.run.is_none()); + assert!(active.checkpoint.is_none()); + assert!(active.cancellation_reason.is_none()); + let steer = TurnCommand::Steer { + display_text: "continue".to_owned(), + model_text: "continue".to_owned(), + }; + assert!(foreground_control.try_receive().is_err()); + active + .turn_control + .as_ref() + .unwrap() + .try_send(steer.clone()) + .unwrap(); + assert_eq!(foreground_control.try_receive().unwrap(), steer); + match &controller.active_provider_progressive_summaries[&conversation_id] { + super::ActiveProviderProgressiveSummaryState::Ready { + identity: ready_identity, + waiting_for_context, + result, + .. + } => { + assert_eq!(ready_identity.run_id, identity.run_id); + assert!(*waiting_for_context); + assert_eq!(result.as_ref().unwrap().0, "checkpoint"); + } + super::ActiveProviderProgressiveSummaryState::InFlight { .. } => { + panic!("completed summary must remain ready until the next model boundary"); + } + } + }); + let conversation = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .unwrap(); + assert!(conversation.has_pending_progressive_summary()); + assert_eq!(conversation.progressive_summary(), None); + }); + }); +} + +#[test] +fn stale_progressive_summary_completion_preserves_a_successor_generation() { + App::test((), |mut app| async move { + initialize_app_for_terminal_view(&mut app); + let terminal = add_window_with_terminal(&mut app, None); + + terminal.update(&mut app, |terminal, ctx| { + let conversation_id = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| { + let conversation_id = + history.start_new_conversation(terminal.id(), false, false, false, ctx); + history + .conversation_mut(&conversation_id) + .unwrap() + .set_has_pending_progressive_summary(true); + conversation_id + }); + let stream_id = ResponseStreamId::new_for_test(); + let response_stream = + ctx.add_model(|_| ResponseStream::new_for_test(stream_id.clone())); + let slot = progressive_summary_test_slot(conversation_id, stream_id, response_stream); + let identity = super::ActiveProviderRunIdentity { + conversation_id, + stream_id: slot.stream_id.clone(), + run_id: slot.run_id.clone(), + }; + let stale_identity = super::ActiveProviderRunIdentity { + conversation_id, + stream_id: ResponseStreamId::new_for_test(), + run_id: ProviderRunId::new("previous-generation"), + }; + + terminal.ai_controller().update(ctx, |controller, ctx| { + controller.active_provider_runs.insert(conversation_id, slot); + controller.active_provider_progressive_summaries.insert( + conversation_id, + super::ActiveProviderProgressiveSummaryState::InFlight { + identity: identity.clone(), + waiting_for_context: false, + }, + ); + + controller.finish_active_provider_progressive_summary( + stale_identity, + progressive_summary_test_plan(Vec::new()), + Ok(( + "stale checkpoint".to_owned(), + Usage::default(), + "summary-model".to_owned(), + )), + ctx, + ); + + let active = &controller.active_provider_runs[&conversation_id]; + assert_eq!(active.run_id, identity.run_id); + assert_eq!(active.stream_id, identity.stream_id); + match &controller.active_provider_progressive_summaries[&conversation_id] { + super::ActiveProviderProgressiveSummaryState::InFlight { + identity: pending_identity, + waiting_for_context, + } => { + assert_eq!(pending_identity.run_id, identity.run_id); + assert_eq!(pending_identity.stream_id, identity.stream_id); + assert!(!*waiting_for_context); + } + super::ActiveProviderProgressiveSummaryState::Ready { .. } => { + panic!("stale completion must not replace the successor summary"); + } + } + }); + let conversation = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .unwrap(); + assert!(conversation.has_pending_progressive_summary()); + assert_eq!(conversation.progressive_summary(), None); + }); + }); +} + +#[test] +fn idle_progressive_summary_preserves_appended_messages_and_archives_tool_results() { + App::test((), |mut app| async move { + initialize_app_for_terminal_view(&mut app); + let terminal = add_window_with_terminal(&mut app, None); + + terminal.update(&mut app, |terminal, ctx| { + let prefix = vec![ + ConversationMessage { + role: MessageRole::Assistant, + content: MessageContent::ToolUse { + tool_use_id: "read-original".to_owned(), + name: "read_file".to_owned(), + input: serde_json::json!({"path": "src/original.rs"}), + }, + }, + ConversationMessage { + role: MessageRole::User, + content: MessageContent::ToolResult { + tool_use_id: "read-original".to_owned(), + content: "original file contents".repeat(1_000), + is_error: false, + }, + }, + ]; + let tail = vec![ + ConversationMessage { + role: MessageRole::Assistant, + content: MessageContent::Text("The change is ready".to_owned()), + }, + ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text("Appended while the summary was running".to_owned()), + }, + ]; + let conversation_id = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| { + let conversation_id = + history.start_new_conversation(terminal.id(), false, false, false, ctx); + let conversation = history.conversation_mut(&conversation_id).unwrap(); + conversation.append_to_bedrock_history(prefix.clone()); + conversation.append_to_bedrock_history(tail.clone()); + conversation.set_has_pending_progressive_summary(true); + conversation.set_status_for_test(ConversationStatus::Success); + conversation_id + }); + let identity = super::ActiveProviderRunIdentity { + conversation_id, + stream_id: ResponseStreamId::new_for_test(), + run_id: ProviderRunId::new("finished-generation"), + }; + + terminal.ai_controller().update(ctx, |controller, ctx| { + controller.active_provider_progressive_summaries.insert( + conversation_id, + super::ActiveProviderProgressiveSummaryState::InFlight { + identity: identity.clone(), + waiting_for_context: false, + }, + ); + controller.finish_active_provider_progressive_summary( + identity, + progressive_summary_test_plan(prefix.clone()), + Ok(( + "checkpoint".to_owned(), + Usage::default(), + "summary-model".to_owned(), + )), + ctx, + ); + assert!(!controller.active_provider_runs.contains_key(&conversation_id)); + assert!(!controller + .active_provider_progressive_summaries + .contains_key(&conversation_id)); + }); + + let conversation = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .unwrap(); + assert_eq!(conversation.progressive_summary(), Some("checkpoint")); + assert_eq!(conversation.bedrock_message_history(), tail.as_slice()); + assert_eq!(conversation.tool_result_archive(), prefix.as_slice()); + assert!(!conversation.has_pending_progressive_summary()); + assert_eq!(conversation.status(), &ConversationStatus::Success); + assert_eq!( + conversation.current_context_tokens(), + 100 + super::estimated_history_tokens(&super::progressive_summary_messages( + "checkpoint".to_owned() + )) + super::estimated_history_tokens(&tail) + ); + }); + }); +} + +#[test] +fn idle_progressive_summary_rejects_a_changed_prefix() { + App::test((), |mut app| async move { + initialize_app_for_terminal_view(&mut app); + let terminal = add_window_with_terminal(&mut app, None); + + terminal.update(&mut app, |terminal, ctx| { + let prefix = vec![ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text("original project context".repeat(1_000)), + }]; + let replacement = vec![ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text("replacement context that must survive".repeat(1_000)), + }]; + let conversation_id = + BlocklistAIHistoryModel::handle(ctx).update(ctx, |history, ctx| { + let conversation_id = + history.start_new_conversation(terminal.id(), false, false, false, ctx); + let conversation = history.conversation_mut(&conversation_id).unwrap(); + conversation.append_to_bedrock_history(replacement.clone()); + conversation.set_has_pending_progressive_summary(true); + conversation.set_status_for_test(ConversationStatus::Success); + conversation_id + }); + let identity = super::ActiveProviderRunIdentity { + conversation_id, + stream_id: ResponseStreamId::new_for_test(), + run_id: ProviderRunId::new("finished-generation"), + }; + + terminal.ai_controller().update(ctx, |controller, ctx| { + controller.active_provider_progressive_summaries.insert( + conversation_id, + super::ActiveProviderProgressiveSummaryState::InFlight { + identity: identity.clone(), + waiting_for_context: false, + }, + ); + controller.finish_active_provider_progressive_summary( + identity, + progressive_summary_test_plan(prefix), + Ok(( + "stale checkpoint".to_owned(), + Usage::default(), + "summary-model".to_owned(), + )), + ctx, + ); + assert!(!controller + .active_provider_progressive_summaries + .contains_key(&conversation_id)); + }); + + let conversation = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .unwrap(); + assert_eq!(conversation.progressive_summary(), None); + assert_eq!(conversation.bedrock_message_history(), replacement.as_slice()); + assert!(conversation.tool_result_archive().is_empty()); + assert!(!conversation.has_pending_progressive_summary()); + }); + }); +} + +fn progressive_summary_test_plan( + summarized_prefix: Vec, +) -> super::ActiveProgressiveSummaryPlan { + super::ActiveProgressiveSummaryPlan { + split_point: summarized_prefix.len(), + retained_messages_count: 1, + trigger_context_tokens: 10_000, + active_context_limit: 128_000, + persistent_context_tokens: 100, + candidates: Vec::new(), + summarize_messages: summarized_prefix.clone(), + summarized_prefix, + } +} + +fn progressive_summary_test_slot( + conversation_id: AIConversationId, + stream_id: ResponseStreamId, + response_stream: ModelHandle, +) -> super::ActiveProviderRunSlot { + let snapshot = provider_snapshot(conversation_id); + super::ActiveProviderRunSlot { + stream_id, + response_stream, + did_input_contain_user_query: true, + run_id: snapshot.run.id().clone(), + root_task_id: snapshot.root_task_id, + projection_target: snapshot.projection_target, + run: None, + checkpoint: None, + turn_control: None, + cancellation_reason: None, + committed_provider_batch: None, + finished_provider_batch: None, + command_action_refs: HashMap::new(), + command_monitor: None, + pending_monitor_observation: None, + pending_command_completion: None, + monitor_prose_continuations: 0, + live_command_output_tokens: 0, + command_output_poll_scheduled: false, + retry_status: None, + tool_call_progress: Vec::new(), + } +} + #[test] fn progressive_summary_compacts_short_histories_to_the_token_budget() { let messages = vec![ diff --git a/app/src/ai/runtime/provider_run_coordinator.rs b/app/src/ai/runtime/provider_run_coordinator.rs index 391e2138..5eb9e291 100644 --- a/app/src/ai/runtime/provider_run_coordinator.rs +++ b/app/src/ai/runtime/provider_run_coordinator.rs @@ -260,6 +260,24 @@ impl ProviderRunCoordinator { .or(self.max_context_tokens) } + /// Leave enough history space for a command-monitor turn before switching + /// profiles, accounting for each profile's own fixed prompt and output budget. + pub(crate) fn progressive_summary_context_budget(&self) -> Option<(u32, u32)> { + [self.run.profile().as_str(), CLI_MONITOR_PROVIDER_PROFILE] + .into_iter() + .filter_map(|name| { + let profile = self.profiles.get(name)?; + let limit = profile.max_context_tokens.or(self.max_context_tokens)?; + let mut request = profile.request.clone(); + request.prompt = None; + request.messages.clear(); + let fixed = u32::try_from(estimate_turn_request_tokens(&request)) + .unwrap_or(u32::MAX); + Some((limit, fixed)) + }) + .min_by_key(|(limit, fixed)| limit.saturating_sub(*fixed)) + } + /// Includes the current transcript, fixed request context, tools, and output reservation. /// Read at the model boundary so newly committed tool output is included. pub(crate) fn estimated_context_tokens(&self) -> Option<(u32, u32)> { diff --git a/app/src/ai/runtime/provider_run_coordinator_tests.rs b/app/src/ai/runtime/provider_run_coordinator_tests.rs index fa685f81..a8eed5b7 100644 --- a/app/src/ai/runtime/provider_run_coordinator_tests.rs +++ b/app/src/ai/runtime/provider_run_coordinator_tests.rs @@ -1608,6 +1608,118 @@ async fn compaction_boundary_keeps_the_monitor_context_limit() { assert_eq!(coordinator.current_context_limit(), Some(100_000)); } +#[test] +fn progressive_summary_budget_anticipates_switching_to_a_smaller_monitor() { + let runtime = Arc::new(ScriptedRuntime::new(vec![])); + let mut coordinator = coordinator(runtime.clone()); + coordinator.set_max_context_tokens(Some(500_000)); + coordinator + .insert_profile( + CLI_MONITOR_PROVIDER_PROFILE, + ProviderRunProfile::new(runtime, request()).with_max_context_tokens(Some(275_000)), + ) + .unwrap(); + + assert_eq!(coordinator.current_context_limit(), Some(500_000)); + assert_eq!( + coordinator + .progressive_summary_context_budget() + .map(|(limit, _)| limit), + Some(275_000) + ); + + coordinator + .run_mut() + .set_ready_profile(CLI_MONITOR_PROVIDER_PROFILE) + .unwrap(); + assert_eq!(coordinator.current_context_limit(), Some(275_000)); + assert_eq!( + coordinator + .progressive_summary_context_budget() + .map(|(limit, _)| limit), + Some(275_000) + ); + + coordinator + .run_mut() + .set_ready_profile(BASE_PROVIDER_PROFILE) + .unwrap(); + assert_eq!(coordinator.current_context_limit(), Some(500_000)); + assert_eq!( + coordinator + .progressive_summary_context_budget() + .map(|(limit, _)| limit), + Some(275_000) + ); +} + +#[test] +fn progressive_summary_budget_respects_smaller_active_limits_and_profile_fallbacks() { + for (base_limit, monitor_limit, fallback_limit, expected) in [ + ( + Some(200_000), + Some(Some(275_000)), + Some(500_000), + Some(200_000), + ), + (Some(500_000), Some(None), Some(275_000), Some(275_000)), + (Some(500_000), None, Some(275_000), Some(500_000)), + (None, None, None, None), + ] { + let runtime = Arc::new(ScriptedRuntime::new(vec![])); + let mut coordinator = ProviderRunCoordinator::from_request( + "progressive-summary-budget", + ProviderRunProfile::new(runtime.clone(), request()).with_max_context_tokens(base_limit), + Vec::new(), + ProviderRunLimits::default(), + ) + .unwrap(); + coordinator.set_max_context_tokens(fallback_limit); + if let Some(monitor_limit) = monitor_limit { + coordinator + .insert_profile( + CLI_MONITOR_PROVIDER_PROFILE, + ProviderRunProfile::new(runtime, request()) + .with_max_context_tokens(monitor_limit), + ) + .unwrap(); + } + + assert_eq!( + coordinator.current_context_limit(), + base_limit.or(fallback_limit) + ); + assert_eq!( + coordinator + .progressive_summary_context_budget() + .map(|(limit, _)| limit), + expected + ); + } +} + +#[test] +fn progressive_summary_budget_accounts_for_output_reservations_when_selecting_a_profile() { + let runtime = Arc::new(ScriptedRuntime::new(vec![])); + let mut coordinator = coordinator(runtime.clone()); + coordinator.set_max_context_tokens(Some(275_000)); + let mut monitor_request = request(); + monitor_request.max_output_tokens = Some(300_000); + coordinator + .insert_profile( + CLI_MONITOR_PROVIDER_PROFILE, + ProviderRunProfile::new(runtime, monitor_request) + .with_max_context_tokens(Some(500_000)), + ) + .unwrap(); + + let (limit, persistent_tokens) = coordinator.progressive_summary_context_budget().unwrap(); + + assert_eq!(coordinator.current_context_limit(), Some(275_000)); + assert_eq!(limit, 500_000); + assert!(persistent_tokens > 300_000); +} + #[tokio::test] async fn monitor_request_uses_its_profile_context_limit() { let base_runtime = Arc::new(ScriptedRuntime::new(vec![answer_turn()])); From 7ccb0e9abc908b1258d8e307677ecc7faf078134 Mon Sep 17 00:00:00 2001 From: Ryan Ward Date: Tue, 20 Oct 2026 21:57:15 +0000 Subject: [PATCH 08/12] Track task progress and organize Galaxy Drive plans --- AGENTS.md | 2 + app/src/ai/agent/todos/mod.rs | 12 +- app/src/ai/blocklist/controller.rs | 27 +- app/src/ai/document/mod.rs | 1 + app/src/ai/document/plan_archival.rs | 500 ++++++++++++++++++ app/src/ai/document/plan_archival_tests.rs | 417 +++++++++++++++ app/src/ai/execution_profiles/profiles.rs | 1 + app/src/ai/provider/response_translator.rs | 1 + app/src/ai/runtime/event_translator.rs | 142 +++-- app/src/ai/runtime/event_translator_tests.rs | 58 ++ app/src/ai/runtime/mod.rs | 2 + .../ai/runtime/provider_run_coordinator.rs | 20 +- .../runtime/provider_run_coordinator_tests.rs | 44 ++ app/src/ai/runtime/rig_request.rs | 17 +- app/src/ai/runtime/rig_request_tests.rs | 19 +- app/src/ai/runtime/task_progress.rs | 292 ++++++++++ app/src/ai/runtime/task_progress_tests.rs | 116 ++++ app/src/drive/index.rs | 146 ++++- app/src/drive/index_tests.rs | 67 ++- app/src/drive/items/item.rs | 24 +- app/src/lib.rs | 1 + app/src/local_object_repository.rs | 104 +++- app/src/local_object_repository_tests.rs | 263 +++++++++ crates/galaxy_agent_core/src/lib.rs | 2 + crates/galaxy_agent_core/src/provider_run.rs | 60 ++- .../src/provider_run_tests.rs | 197 ++++++- crates/galaxy_agent_core/src/task_list.rs | 129 +++++ .../galaxy_agent_core/src/task_list_tests.rs | 157 ++++++ crates/galaxy_agent_core/src/tool_policy.rs | 19 +- .../src/tool_policy_tests.rs | 56 ++ 30 files changed, 2791 insertions(+), 105 deletions(-) create mode 100644 app/src/ai/document/plan_archival.rs create mode 100644 app/src/ai/document/plan_archival_tests.rs create mode 100644 app/src/ai/runtime/task_progress.rs create mode 100644 app/src/ai/runtime/task_progress_tests.rs create mode 100644 crates/galaxy_agent_core/src/task_list.rs create mode 100644 crates/galaxy_agent_core/src/task_list_tests.rs diff --git a/AGENTS.md b/AGENTS.md index df151ec7..fe68bd59 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -212,6 +212,8 @@ Key invariants: - Same-conversation direct-provider follow-ups queue behind the cancelling generation; the old run keeps the active slot until terminal projection and cleanup, queued intent is checkpointed without credentials for restart recovery, and queued-only restore validates provider ownership and terminalizes the abandoned unprepared exchange from its persisted projection/stream identity before starting the successor; the next generation rebuilds provider history after cleanup so it includes the old generation's final committed output, and stale callbacks are ignored by stream identity - Known tools are in `KNOWN_TOOLS` in `response_translator.rs`; definitions are built by `tool_definition_for_name()` in `convert_request.rs` - Direct-provider normal and plan turns must advertise `read_plan`, `create_plan`, and `edit_plan` when the matching document capabilities are enabled; plan-creation requests should call `create_plan` after research rather than only returning prose +- Direct-provider normal and plan turns expose `update_task_list` for explicit progress updates. Keep the validated current list in durable `ProviderRun` state, inject it into each request, preserve it across compaction and restore, and reconcile restored UI updates idempotently. Finishing a response must never mark unfinished work complete. CLI monitors and completed-command assessments retain their focused tool sets. +- Completed local plans are archived intact into `Plans/Archive` after seven days only when their conversation has explicitly completed all checklist items and has no active provider run or child. Draft, unfinished, or reopened plans stay visible; resumed work restores an automatically archived plan, and manual recovery is respected. Archival never deletes plan contents. - Unknown or invalid tool calls receive one correlated synthetic error result and a visible `AgentOutput` message; the durable run owns any continuation - `recall_tool_history` is an inline completed tool batch. `ProviderRun` commits its synthetic result and starts a bounded next turn without routing it through client action execution - `recall_tool_history` excludes earlier calls to itself; archived tool results remain searchable by query or exact `tool_use_id` diff --git a/app/src/ai/agent/todos/mod.rs b/app/src/ai/agent/todos/mod.rs index 5852a1f8..8a3d7d73 100644 --- a/app/src/ai/agent/todos/mod.rs +++ b/app/src/ai/agent/todos/mod.rs @@ -83,7 +83,17 @@ impl AIAgentTodoList { .position(|item| item.id == completed_todo_id.clone().into()) .map(|i| self.pending_items.remove(i)) { - self.completed_items.push(item); + // Restored provider batches can redeliver a pending update and + // its completion. Keep one completed item per stable task ID. + if let Some(previous) = self + .completed_items + .iter_mut() + .find(|old| old.id == item.id) + { + *previous = item; + } else { + self.completed_items.push(item); + } } } } diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index b1a72b83..985e0fe3 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -6212,6 +6212,24 @@ impl BlocklistAIController { }, ctx, ); + // Projection can be persisted before a mixed tool batch commits. Seed + // from the actual visible list so replaying its inline update is a no-op. + let projected_task_list = BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .and_then(|conversation| { + crate::ai::runtime::task_list_from_api_tasks( + &conversation.compute_active_tasks(), + &root_task_id.to_string(), + ) + }); + let projector = ProviderRunResponseProjector::restored( + response_config.clone(), + projection_was_initialized, + ) + .with_restored_task_lists( + projected_task_list, + crate::ai::runtime::restored_task_list(coordinator.run()), + ); self.active_provider_runs.insert( conversation_id, ActiveProviderRunSlot { @@ -6223,10 +6241,7 @@ impl BlocklistAIController { projection_target, run: Some(ActiveProviderRun { coordinator, - projector: ProviderRunResponseProjector::restored( - response_config.clone(), - projection_was_initialized, - ), + projector, response_config, action_context, messages_sent, @@ -6681,9 +6696,11 @@ impl BlocklistAIController { .active_provider_runs .get(&conversation_id) .and_then(|slot| slot.cancellation_reason); + let projector = ProviderRunResponseProjector::new(response_config.clone()) + .with_task_list(coordinator.run().task_list().cloned()); let mut run = ActiveProviderRun { coordinator, - projector: ProviderRunResponseProjector::new(response_config.clone()), + projector, response_config, action_context, messages_sent, diff --git a/app/src/ai/document/mod.rs b/app/src/ai/document/mod.rs index 355fa10f..1cb9c1c9 100644 --- a/app/src/ai/document/mod.rs +++ b/app/src/ai/document/mod.rs @@ -1,3 +1,4 @@ pub mod ai_document_model; pub mod orchestration_config_block; +pub(crate) mod plan_archival; pub(in crate::ai) mod plan_publication; diff --git a/app/src/ai/document/plan_archival.rs b/app/src/ai/document/plan_archival.rs new file mode 100644 index 00000000..1d90ff67 --- /dev/null +++ b/app/src/ai/document/plan_archival.rs @@ -0,0 +1,500 @@ +//! Retain completed plans in a recoverable Drive folder after one week. + +use std::collections::{HashMap, HashSet}; +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use galaxy_core::user_preferences::GetUserPreferences as _; +use galaxyui::r#async::Timer; +use galaxyui::{Entity, ModelContext, SingletonEntity}; +use serde::{Deserialize, Serialize}; + +use super::ai_document_model::{ + AIDocumentId, AIDocumentModel, AIDocumentModelEvent, AIDocumentUpdateSource, + AIDocumentUserEditStatus, +}; +use crate::ai::agent::conversation::{AIConversation, AIConversationId, ConversationStatus}; +use crate::ai::artifacts::Artifact; +use crate::ai::blocklist::{BlocklistAIHistoryEvent, BlocklistAIHistoryModel}; +use crate::cloud_object::model::persistence::{CloudModel, CloudModelEvent}; +use crate::cloud_object::{CloudObject, CloudObjectLocation, Space}; +use crate::drive::CloudObjectTypeAndId; +use crate::local_object_repository::{ + local_owner, LocalObjectRepository, LocalObjectRepositoryEvent, +}; +use crate::server::ids::{ClientId, SyncId}; + +const PLAN_ARCHIVAL_KEY: &str = "CompletedPlanArchival"; +const PLAN_FOLDER_NAME: &str = "Plans"; +const ARCHIVE_FOLDER_NAME: &str = "Archive"; +const SWEEP_INTERVAL: Duration = Duration::from_secs(5 * 60); +const ARCHIVE_AFTER_DAYS: i64 = 7; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +struct CompletedPlan { + conversation_id: AIConversationId, + completed_at: DateTime, + #[serde(default)] + archived_to: Option, + #[serde(default)] + original_folder: Option, +} + +impl CompletedPlan { + fn is_due(&self, now: DateTime) -> bool { + self.archived_to.is_none() + && now.signed_duration_since(self.completed_at) + >= chrono::Duration::days(ARCHIVE_AFTER_DAYS) + } +} + +/// Completion records are local preferences; notebook contents and folder +/// membership continue to use the normal SQLite-backed Drive repository. +#[derive(Debug, Default, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub(crate) struct PlanArchivalModel { + #[serde(default)] + completed_plans: HashMap, + #[serde(default)] + reopened_plans: HashMap, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +struct ReopenedPlan { + conversation_id: AIConversationId, + /// The edited plan must reach a new agent exchange before the old task + /// list can be used as evidence of completion again. + saw_new_exchange: bool, +} + +impl PlanArchivalModel { + pub(crate) fn new(ctx: &mut ModelContext) -> Self { + let model = ctx + .private_user_preferences() + .read_value(PLAN_ARCHIVAL_KEY) + .unwrap_or_default() + .and_then(|json| Self::from_persisted(&json)) + .unwrap_or_default(); + + ctx.subscribe_to_model( + &BlocklistAIHistoryModel::handle(ctx), + |model, _, event, ctx| match event { + BlocklistAIHistoryEvent::AppendedExchange { conversation_id, .. } => { + // Clear completion before a successor turn can start. A + // previously archived plan becomes visible again as well. + model.begin_exchange(*conversation_id, ctx); + } + BlocklistAIHistoryEvent::UpdatedConversationStatus { .. } + | BlocklistAIHistoryEvent::RestoredConversations { .. } => { + model.sweep(ctx); + } + _ => {} + }, + ); + ctx.subscribe_to_model(&AIDocumentModel::handle(ctx), |model, _, event, ctx| { + match event { + AIDocumentModelEvent::DocumentUpdated { + document_id, + source: AIDocumentUpdateSource::User, + .. + } + | AIDocumentModelEvent::DocumentUserEditStatusUpdated { + document_id, + status: AIDocumentUserEditStatus::Dirty, + } => model.reopen_document(*document_id, ctx), + AIDocumentModelEvent::DocumentUserEditStatusUpdated { + document_id, + status: AIDocumentUserEditStatus::UpToDate, + } => model.accept_document_edits(*document_id, ctx), + _ => {} + } + }); + ctx.subscribe_to_model(&CloudModel::handle(ctx), |model, _, event, ctx| { + if matches!(event, CloudModelEvent::InitialLoadCompleted) { + model.sweep(ctx); + } + }); + ctx.subscribe_to_model(&LocalObjectRepository::handle(ctx), |model, _, event, ctx| { + let LocalObjectRepositoryEvent::NotebookContentEdited { id } = event else { + return; + }; + let document_id = CloudModel::as_ref(ctx) + .get_notebook(id) + .and_then(|notebook| notebook.model().ai_document_id); + if let Some(document_id) = document_id { + model.reopen_document(document_id, ctx); + } + }); + + // The first callback runs after singleton registration and session + // restoration. Periodic checks also catch provider cleanup occurring + // just after a successful status event, without delaying that cleanup. + ctx.spawn(async {}, |model, _, ctx| { + model.sweep(ctx); + model.schedule_sweep(ctx); + }); + + model + } + + fn from_persisted(json: &str) -> Option { + serde_json::from_str(json).ok().or_else(|| { + // Preserve completion records written before reopened-plan state + // was added to the preference payload. + serde_json::from_str(json).ok().map(|completed_plans| Self { + completed_plans, + ..Self::default() + }) + }) + } + + fn schedule_sweep(&self, ctx: &mut ModelContext) { + ctx.spawn( + async { Timer::after(SWEEP_INTERVAL).await }, + |model, _, ctx| { + model.sweep(ctx); + model.schedule_sweep(ctx); + }, + ); + } + + fn sweep(&mut self, ctx: &mut ModelContext) { + let before = self.completed_plans.clone(); + let reopened_before = self.reopened_plans.clone(); + self.reopened_plans.retain(|document_id, record| { + !record.saw_new_exchange + || AIDocumentId::try_from(document_id.as_str()) + .ok() + .and_then(|id| AIDocumentModel::as_ref(ctx).get_current_document(&id)) + .is_none_or(|document| document.user_edit_status.is_dirty()) + }); + let observations = { + let history = BlocklistAIHistoryModel::as_ref(ctx); + let documents = AIDocumentModel::as_ref(ctx); + let mut seen = HashSet::new(); + history + .all_live_conversations() + .into_iter() + .chain(history.all_cleared_conversations()) + .filter(|(_, conversation)| seen.insert(conversation.id())) + .map(|(_, conversation)| { + let mut document_ids = documents + .get_all_documents_for_conversation(conversation.id()) + .into_iter() + .map(|(id, _)| id.to_string()) + .collect::>(); + document_ids.extend(conversation.artifacts().iter().filter_map( + |artifact| { + if let Artifact::Plan { document_uid, .. } = artifact { + Some(document_uid.clone()) + } else { + None + } + }, + )); + ( + conversation.id(), + document_ids, + completed_work_time(conversation), + ) + }) + .collect::>() + }; + + for (conversation_id, document_ids, completed_at) in observations { + match completed_at { + Some(completed_at) => { + for document_id in document_ids { + self.observe_completion(document_id, conversation_id, completed_at); + } + } + None => self.resume_conversation(conversation_id, ctx), + } + } + + // Persisted completion records also cover unloaded conversations. + // AppendedExchange clears them before later work starts. Do not use + // metadata.last_modified_at as work identity: checkpoint and title + // saves update that timestamp after the exchange has already finished. + + let now = Utc::now(); + let due = self + .completed_plans + .iter() + .filter(|(_, record)| record.is_due(now)) + .map(|(document_id, _)| document_id.clone()) + .collect::>(); + for document_id in due { + self.archive_plan(&document_id, ctx); + } + + if before != self.completed_plans || reopened_before != self.reopened_plans { + self.persist(ctx); + } + } + + fn observe_completion( + &mut self, + document_id: String, + conversation_id: AIConversationId, + completed_at: DateTime, + ) { + if self.reopened_plans.contains_key(&document_id) { + return; + } + if self + .completed_plans + .get(&document_id) + .is_some_and(|record| record.conversation_id == conversation_id) + { + return; + } + self.completed_plans.insert( + document_id, + CompletedPlan { + conversation_id, + completed_at, + archived_to: None, + original_folder: None, + }, + ); + } + + fn archive_plan(&mut self, document_id: &str, ctx: &mut ModelContext) { + if self.reopened_plans.contains_key(document_id) { + return; + } + if let Some(id) = AIDocumentId::try_from(document_id).ok().filter(|id| { + AIDocumentModel::as_ref(ctx) + .get_current_document(id) + .is_some_and(|document| document.user_edit_status.is_dirty()) + }) { + self.reopen_document(id, ctx); + return; + } + let notebook = CloudModel::as_ref(ctx) + .get_all_active_notebooks() + .find(|notebook| { + notebook.permissions.owner == local_owner() + && notebook + .model() + .ai_document_id + .is_some_and(|id| id.to_string() == document_id) + }) + .cloned(); + let Some(notebook) = notebook else { + return; + }; + + let plans_folder = get_or_create_folder(PLAN_FOLDER_NAME, None, ctx); + let archive_folder = get_or_create_folder(ARCHIVE_FOLDER_NAME, Some(plans_folder), ctx); + let original_folder = notebook.metadata.folder_id; + if original_folder != Some(archive_folder) + && !LocalObjectRepository::handle(ctx).update(ctx, |repository, ctx| { + repository.move_object_to_location( + CloudObjectTypeAndId::Notebook(notebook.id), + CloudObjectLocation::Folder(archive_folder), + ctx, + ) + }) + { + return; + } + if let Some(record) = self.completed_plans.get_mut(document_id) { + record.original_folder = original_folder.filter(|id| *id != archive_folder); + record.archived_to = Some(archive_folder); + } + } + + fn resume_conversation( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + let document_ids = self + .completed_plans + .iter() + .filter(|(_, record)| record.conversation_id == conversation_id) + .map(|(document_id, _)| document_id.clone()) + .collect::>(); + if document_ids.is_empty() { + return; + } + for document_id in document_ids { + let Some(record) = self.completed_plans.remove(&document_id) else { + continue; + }; + restore_archived_plan(&document_id, record, ctx); + } + self.persist(ctx); + } + + fn begin_exchange( + &mut self, + conversation_id: AIConversationId, + ctx: &mut ModelContext, + ) { + self.resume_conversation(conversation_id, ctx); + let mut changed = false; + self.reopened_plans.retain(|document_id, record| { + if record.conversation_id != conversation_id { + return true; + } + record.saw_new_exchange = true; + changed = true; + let document = AIDocumentId::try_from(document_id.as_str()) + .ok() + .and_then(|id| AIDocumentModel::as_ref(ctx).get_current_document(&id)); + document.is_none_or(|document| document.user_edit_status.is_dirty()) + }); + if changed { + self.persist(ctx); + } + } + + fn reopen_document(&mut self, document_id: AIDocumentId, ctx: &mut ModelContext) { + let key = document_id.to_string(); + let completed = self.completed_plans.remove(&key); + // Restored free-floating document views can have a synthetic owner. + // Prefer the real owner retained in the completion/reopening record. + let conversation_id = completed + .as_ref() + .map(|record| record.conversation_id) + .or_else(|| { + self.reopened_plans + .get(&key) + .map(|record| record.conversation_id) + }) + .or_else(|| { + AIDocumentModel::as_ref(ctx).get_conversation_id_for_document_id(&document_id) + }); + if let Some(conversation_id) = conversation_id { + self.reopened_plans.insert( + key.clone(), + ReopenedPlan { + conversation_id, + saw_new_exchange: false, + }, + ); + } + if let Some(record) = completed { + restore_archived_plan(&key, record, ctx); + } + self.persist(ctx); + } + + fn accept_document_edits(&mut self, document_id: AIDocumentId, ctx: &mut ModelContext) { + let key = document_id.to_string(); + if self + .reopened_plans + .get(&key) + .is_some_and(|record| record.saw_new_exchange) + { + self.reopened_plans.remove(&key); + self.persist(ctx); + } + } + + fn persist(&self, ctx: &mut ModelContext) { + match serde_json::to_string(self) { + Ok(json) => { + if let Err(error) = ctx + .private_user_preferences() + .write_value(PLAN_ARCHIVAL_KEY, json) + { + log::error!("Failed to persist completed plan retention: {error}"); + } + } + Err(error) => log::error!("Failed to serialize completed plan retention: {error}"), + } + } +} + +fn restore_archived_plan( + document_id: &str, + record: CompletedPlan, + ctx: &mut ModelContext, +) { + let Some(archive_folder) = record.archived_to else { + return; + }; + let notebook = CloudModel::as_ref(ctx) + .get_all_active_notebooks() + .find(|notebook| { + notebook.permissions.owner == local_owner() + && notebook.metadata.folder_id == Some(archive_folder) + && notebook + .model() + .ai_document_id + .is_some_and(|id| id.to_string() == document_id) + }) + .cloned(); + let Some(notebook) = notebook else { + return; + }; + let original_folder = record.original_folder.filter(|id| { + let cloud_model = CloudModel::as_ref(ctx); + cloud_model.get_folder(id).is_some_and(|folder| { + folder.permissions.owner == local_owner() && !folder.is_trashed(cloud_model) + }) + }); + let destination = original_folder + .map(CloudObjectLocation::Folder) + .unwrap_or(CloudObjectLocation::Space(Space::Personal)); + LocalObjectRepository::handle(ctx).update(ctx, |repository, ctx| { + repository.move_object_to_location( + CloudObjectTypeAndId::Notebook(notebook.id), + destination, + ctx, + ); + }); +} + +fn completed_work_time(conversation: &AIConversation) -> Option> { + if !matches!(conversation.status(), ConversationStatus::Success) + || conversation.active_provider_run_json().is_some() + || conversation.has_active_subagent() + || !conversation + .active_todo_list() + .is_some_and(|list| list.is_finished()) + { + return None; + } + conversation + .latest_exchange() + .and_then(|exchange| exchange.finish_time) + .map(|time| time.with_timezone(&Utc)) +} + +fn get_or_create_folder( + name: &str, + parent: Option, + ctx: &mut ModelContext, +) -> SyncId { + if let Some(id) = CloudModel::as_ref(ctx) + .get_all_active_folders() + .find(|folder| { + folder.permissions.owner == local_owner() + && folder.model().name == name + && folder.metadata.folder_id == parent + }) + .map(|folder| folder.id) + { + return id; + } + let id = SyncId::ClientId(ClientId::new()); + LocalObjectRepository::handle(ctx).update(ctx, |repository, ctx| { + repository.create_folder_with_id(id, name.to_string(), parent, ctx); + }); + id +} + +impl Entity for PlanArchivalModel { + type Event = (); +} + +impl SingletonEntity for PlanArchivalModel {} + +#[cfg(test)] +#[path = "plan_archival_tests.rs"] +mod tests; diff --git a/app/src/ai/document/plan_archival_tests.rs b/app/src/ai/document/plan_archival_tests.rs new file mode 100644 index 00000000..0b5cdee1 --- /dev/null +++ b/app/src/ai/document/plan_archival_tests.rs @@ -0,0 +1,417 @@ +use std::sync::mpsc::sync_channel; + +use ai::document::AIDocumentId; +use galaxyui::App; + +use super::*; +use crate::notebooks::CloudNotebookModel; +use crate::persistence::ModelEvent; +use crate::test_util::settings::initialize_settings_for_tests; + +fn completed_plan(completed_at: DateTime) -> CompletedPlan { + CompletedPlan { + conversation_id: AIConversationId::new(), + completed_at, + archived_to: None, + original_folder: None, + } +} + +#[test] +fn completed_plan_archives_after_seven_full_days_only() { + let now = Utc::now(); + let plan = completed_plan(now - chrono::Duration::days(7)); + assert!(!plan.is_due(now - chrono::Duration::milliseconds(1))); + assert!(plan.is_due(now)); + assert!(!completed_plan(now + chrono::Duration::days(1)).is_due(now)); +} + +#[test] +fn archived_plan_deadline_and_recovery_location_survive_restart() { + let mut plan = completed_plan(Utc::now() - chrono::Duration::days(8)); + plan.archived_to = Some(SyncId::ClientId(ClientId::new())); + plan.original_folder = Some(SyncId::ClientId(ClientId::new())); + let serialized = serde_json::to_string(&plan).unwrap(); + let restored: CompletedPlan = serde_json::from_str(&serialized).unwrap(); + assert_eq!(restored, plan); + assert!(!restored.is_due(Utc::now())); +} + +#[test] +fn revisiting_a_completed_plan_preserves_its_original_deadline() { + let mut model = PlanArchivalModel { + completed_plans: HashMap::new(), + ..Default::default() + }; + let conversation_id = AIConversationId::new(); + let completed_at = Utc::now() - chrono::Duration::days(6); + model.observe_completion("plan".to_string(), conversation_id, completed_at); + model.observe_completion("plan".to_string(), conversation_id, Utc::now()); + assert_eq!(model.completed_plans["plan"].completed_at, completed_at); +} + +#[test] +fn legacy_completed_plan_records_are_preserved() { + let record = completed_plan(Utc::now() - chrono::Duration::days(4)); + let records = HashMap::from([("plan".to_string(), record.clone())]); + let json = serde_json::to_string(&records).unwrap(); + let model = PlanArchivalModel::from_persisted(&json).unwrap(); + assert_eq!(model.completed_plans["plan"], record); + assert!(model.reopened_plans.is_empty()); +} + +#[test] +fn dirty_and_reverted_plans_stay_visible_until_new_work_accepts_the_edits() { + App::test((), |mut app| async move { + initialize_settings_for_tests(&mut app); + app.add_singleton_model(|_| CloudModel::new(None, Vec::new(), None)); + app.add_singleton_model(|_| BlocklistAIHistoryModel::new_for_test()); + let documents = app.add_singleton_model(|_| AIDocumentModel::new_for_test()); + let repository = app.add_singleton_model(|ctx| { + LocalObjectRepository::new(None, None, ctx) + }); + let document_id = AIDocumentId::new(); + let notebook_id = SyncId::ClientId(ClientId::new()); + repository.update(&mut app, |repository, ctx| { + repository.create_notebook_with_id( + notebook_id, + None, + CloudNotebookModel { + title: "Reopened plan".to_string(), + data: "New unfinished work".to_string(), + ai_document_id: Some(document_id), + conversation_id: None, + }, + ctx, + ); + }); + let archival = app.add_model(PlanArchivalModel::new); + let conversation_id = AIConversationId::new(); + let old_completion = Utc::now() - chrono::Duration::days(8); + archival.update(&mut app, |model, ctx| { + model.observe_completion(document_id.to_string(), conversation_id, old_completion); + model.archive_plan(&document_id.to_string(), ctx); + }); + assert!(CloudModel::as_ref(&app) + .get_notebook(¬ebook_id) + .unwrap() + .metadata + .folder_id + .is_some()); + + // Typing in the editor emits Dirty; reverting a version also emits a + // user-originated DocumentUpdated event. Either must reopen the plan. + for event in [ + AIDocumentModelEvent::DocumentUserEditStatusUpdated { + document_id, + status: AIDocumentUserEditStatus::Dirty, + }, + AIDocumentModelEvent::DocumentUpdated { + document_id, + version: Default::default(), + source: AIDocumentUpdateSource::User, + }, + ] { + documents.update(&mut app, |_, ctx| ctx.emit(event)); + archival.update(&mut app, |model, ctx| { + model.observe_completion(document_id.to_string(), conversation_id, old_completion); + model.sweep(ctx); + assert!(!model.completed_plans.contains_key(&document_id.to_string())); + assert_eq!( + model.reopened_plans[&document_id.to_string()].conversation_id, + conversation_id, + ); + let json = serde_json::to_string(model).unwrap(); + let mut restored = PlanArchivalModel::from_persisted(&json).unwrap(); + restored.observe_completion( + document_id.to_string(), + conversation_id, + old_completion, + ); + assert!(restored.completed_plans.is_empty()); + }); + assert!(CloudModel::as_ref(&app) + .get_notebook(¬ebook_id) + .unwrap() + .metadata + .folder_id + .is_none()); + } + + archival.update(&mut app, |model, ctx| { + // Merely restoring clean editor state cannot close a reopened plan. + model.accept_document_edits(document_id, ctx); + assert!(model.reopened_plans.contains_key(&document_id.to_string())); + model.begin_exchange(conversation_id, ctx); + assert!(model.reopened_plans.contains_key(&document_id.to_string())); + model.accept_document_edits(document_id, ctx); + assert!(model.reopened_plans.is_empty()); + let finished_again = Utc::now(); + model.observe_completion(document_id.to_string(), conversation_id, finished_again); + assert_eq!(model.completed_plans[&document_id.to_string()].completed_at, finished_again); + assert!(!model.completed_plans[&document_id.to_string()].is_due(finished_again)); + }); + }); +} + +#[test] +fn a_successful_draft_conversation_does_not_complete_a_plan() { + let mut conversation = AIConversation::new(false, false); + conversation.set_status_for_test(ConversationStatus::Success); + assert!(completed_work_time(&conversation).is_none()); + + for status in [ + ConversationStatus::InProgress, + ConversationStatus::WaitingForEvents, + ConversationStatus::TransientError, + ConversationStatus::Error, + ConversationStatus::Cancelled, + ] { + conversation.set_status_for_test(status); + assert!(completed_work_time(&conversation).is_none()); + } +} + +#[test] +fn archival_preserves_notebook_and_resuming_restores_original_folder() { + App::test((), |mut app| async move { + initialize_settings_for_tests(&mut app); + let (sender, receiver) = sync_channel(32); + app.add_singleton_model(|_| CloudModel::new(None, Vec::new(), None)); + app.add_singleton_model(|_| BlocklistAIHistoryModel::new_for_test()); + app.add_singleton_model(|_| AIDocumentModel::new_for_test()); + let repository = app.add_singleton_model(|ctx| { + LocalObjectRepository::new(Some(sender), None, ctx) + }); + let archival = app.add_model(|_| PlanArchivalModel { + completed_plans: HashMap::new(), + ..Default::default() + }); + let original_folder = SyncId::ClientId(ClientId::new()); + let notebook_id = SyncId::ClientId(ClientId::new()); + let document_id = AIDocumentId::new(); + let notebook_model = CloudNotebookModel { + title: "Retained plan".to_string(), + data: "# Plan\n- [x] Finished work".to_string(), + ai_document_id: Some(document_id), + conversation_id: None, + }; + repository.update(&mut app, |repository, ctx| { + repository.create_folder_with_id( + original_folder, + "Project".to_string(), + None, + ctx, + ); + repository.create_notebook_with_id( + notebook_id, + Some(original_folder), + notebook_model.clone(), + ctx, + ); + }); + let conversation_id = AIConversationId::new(); + archival.update(&mut app, |model, ctx| { + model.observe_completion( + document_id.to_string(), + conversation_id, + Utc::now() - chrono::Duration::days(8), + ); + model.persist(ctx); + }); + + // Simulate startup with persisted completion records and no loaded + // conversations/documents. The notebook is still retained in Drive. + let archival = app.add_model(PlanArchivalModel::new); + archival.update(&mut app, |model, ctx| model.sweep(ctx)); + + let archive_folder = archival.read(&app, |model, _| { + model.completed_plans[&document_id.to_string()] + .archived_to + .expect("plan should have been archived") + }); + let cloud_model = CloudModel::as_ref(&app); + let notebook = cloud_model.get_notebook(¬ebook_id).unwrap(); + assert_eq!(notebook.model(), ¬ebook_model); + assert_eq!(notebook.metadata.folder_id, Some(archive_folder)); + assert!(!notebook.is_trashed(cloud_model)); + let folder = cloud_model.get_folder(&archive_folder).unwrap(); + assert_eq!(folder.model().name, ARCHIVE_FOLDER_NAME); + let parent = cloud_model + .get_folder(&folder.metadata.folder_id.unwrap()) + .unwrap(); + assert_eq!(parent.model().name, PLAN_FOLDER_NAME); + + let events = receiver.try_iter().collect::>(); + assert!(events.iter().any(|event| matches!( + event, + ModelEvent::UpsertNotebook { notebook } + if notebook.id == notebook_id + && notebook.metadata.folder_id == Some(archive_folder) + && notebook.model() == ¬ebook_model + ))); + assert!(!events.iter().any(|event| matches!(event, ModelEvent::DeleteObjects { .. }))); + + archival.update(&mut app, |model, ctx| { + model.resume_conversation(conversation_id, ctx); + assert!(model.completed_plans.is_empty()); + }); + let cloud_model = CloudModel::as_ref(&app); + let notebook = cloud_model.get_notebook(¬ebook_id).unwrap(); + assert_eq!(notebook.metadata.folder_id, Some(original_folder)); + assert_eq!(notebook.model(), ¬ebook_model); + assert!(!notebook.is_trashed(cloud_model)); + }); +} + +#[test] +fn manual_restore_does_not_rearchive_a_completed_plan() { + App::test((), |mut app| async move { + initialize_settings_for_tests(&mut app); + app.add_singleton_model(|_| CloudModel::new(None, Vec::new(), None)); + app.add_singleton_model(|_| BlocklistAIHistoryModel::new_for_test()); + app.add_singleton_model(|_| AIDocumentModel::new_for_test()); + let repository = app.add_singleton_model(|ctx| { + LocalObjectRepository::new(None, None, ctx) + }); + let archive_folder = SyncId::ClientId(ClientId::new()); + let notebook_id = SyncId::ClientId(ClientId::new()); + let document_id = AIDocumentId::new(); + repository.update(&mut app, |repository, ctx| { + repository.create_folder_with_id(archive_folder, "Archive".to_string(), None, ctx); + repository.create_notebook_with_id( + notebook_id, + Some(archive_folder), + CloudNotebookModel { + title: "Recovered plan".to_string(), + data: "Keep this plan visible".to_string(), + ai_document_id: Some(document_id), + conversation_id: None, + }, + ctx, + ); + }); + let mut record = completed_plan(Utc::now() - chrono::Duration::days(8)); + record.archived_to = Some(archive_folder); + let archival = app.add_model(|_| PlanArchivalModel { + completed_plans: HashMap::from([(document_id.to_string(), record)]), + ..Default::default() + }); + + assert!(repository.update(&mut app, |repository, ctx| { + repository.move_object_to_location( + CloudObjectTypeAndId::Notebook(notebook_id), + CloudObjectLocation::Space(Space::Personal), + ctx, + ) + })); + archival.update(&mut app, |model, ctx| model.sweep(ctx)); + + let cloud_model = CloudModel::as_ref(&app); + let notebook = cloud_model.get_notebook(¬ebook_id).unwrap(); + assert!(notebook.metadata.folder_id.is_none()); + assert_eq!(notebook.model().data, "Keep this plan visible"); + archival.read(&app, |model, _| { + assert!(!model.completed_plans[&document_id.to_string()].is_due(Utc::now())); + }); + }); +} + +#[test] +fn ordinary_notebooks_are_never_archived() { + App::test((), |mut app| async move { + initialize_settings_for_tests(&mut app); + app.add_singleton_model(|_| CloudModel::new(None, Vec::new(), None)); + app.add_singleton_model(|_| AIDocumentModel::new_for_test()); + let repository = app.add_singleton_model(|ctx| { + LocalObjectRepository::new(None, None, ctx) + }); + let notebook_id = SyncId::ClientId(ClientId::new()); + repository.update(&mut app, |repository, ctx| { + repository.create_notebook_with_id( + notebook_id, + None, + CloudNotebookModel { + title: "Ordinary notes".to_string(), + data: "Keep these notes".to_string(), + ai_document_id: None, + conversation_id: None, + }, + ctx, + ); + }); + let archival = app.add_model(|_| PlanArchivalModel { + completed_plans: HashMap::new(), + ..Default::default() + }); + archival.update(&mut app, |model, ctx| { + model.observe_completion( + notebook_id.to_string(), + AIConversationId::new(), + Utc::now() - chrono::Duration::days(8), + ); + model.archive_plan(¬ebook_id.to_string(), ctx); + }); + let cloud_model = CloudModel::as_ref(&app); + let notebook = cloud_model.get_notebook(¬ebook_id).unwrap(); + assert!(notebook.metadata.folder_id.is_none()); + assert_eq!(notebook.model().data, "Keep these notes"); + assert_eq!(cloud_model.get_all_active_folders().count(), 0); + }); +} + +#[test] +fn editing_an_archived_plan_from_drive_reopens_it_without_loading_its_conversation() { + App::test((), |mut app| async move { + initialize_settings_for_tests(&mut app); + app.add_singleton_model(|_| CloudModel::new(None, Vec::new(), None)); + app.add_singleton_model(|_| BlocklistAIHistoryModel::new_for_test()); + app.add_singleton_model(|_| AIDocumentModel::new_for_test()); + let repository = app.add_singleton_model(|ctx| { + LocalObjectRepository::new(None, None, ctx) + }); + let document_id = AIDocumentId::new(); + let notebook_id = SyncId::ClientId(ClientId::new()); + repository.update(&mut app, |repository, ctx| { + repository.create_notebook_with_id( + notebook_id, + None, + CloudNotebookModel { + title: "Completed plan".to_string(), + data: "- [x] Done".to_string(), + ai_document_id: Some(document_id), + conversation_id: None, + }, + ctx, + ); + }); + let archival = app.add_model(PlanArchivalModel::new); + let conversation_id = AIConversationId::new(); + let completed_at = Utc::now() - chrono::Duration::days(8); + archival.update(&mut app, |model, ctx| { + model.observe_completion(document_id.to_string(), conversation_id, completed_at); + model.archive_plan(&document_id.to_string(), ctx); + }); + repository.update(&mut app, |repository, ctx| { + repository.update_notebook_data( + notebook_id, + "- [ ] New unfinished work".to_string(), + ctx, + ); + }); + archival.update(&mut app, |model, ctx| { + model.observe_completion(document_id.to_string(), conversation_id, completed_at); + model.sweep(ctx); + assert!(model.completed_plans.is_empty()); + assert_eq!( + model.reopened_plans[&document_id.to_string()].conversation_id, + conversation_id, + ); + }); + let cloud_model = CloudModel::as_ref(&app); + let notebook = cloud_model.get_notebook(¬ebook_id).unwrap(); + assert!(notebook.metadata.folder_id.is_none()); + assert_eq!(notebook.model().data, "- [ ] New unfinished work"); + assert!(!notebook.is_trashed(cloud_model)); + }); +} diff --git a/app/src/ai/execution_profiles/profiles.rs b/app/src/ai/execution_profiles/profiles.rs index e585a0f8..91c5f412 100644 --- a/app/src/ai/execution_profiles/profiles.rs +++ b/app/src/ai/execution_profiles/profiles.rs @@ -1377,6 +1377,7 @@ impl AIExecutionProfilesModel { } LocalObjectRepositoryEvent::Rules | LocalObjectRepositoryEvent::Notebooks { .. } + | LocalObjectRepositoryEvent::NotebookContentEdited { .. } | LocalObjectRepositoryEvent::Workflows { .. } => {} } } diff --git a/app/src/ai/provider/response_translator.rs b/app/src/ai/provider/response_translator.rs index 9c9d561d..fd470747 100644 --- a/app/src/ai/provider/response_translator.rs +++ b/app/src/ai/provider/response_translator.rs @@ -1625,6 +1625,7 @@ const KNOWN_TOOLS: &[&str] = &[ "read_skill", "fetch_conversation", "recall_tool_history", + "update_task_list", ]; pub(super) fn is_known_tool(name: &str) -> bool { diff --git a/app/src/ai/runtime/event_translator.rs b/app/src/ai/runtime/event_translator.rs index 739bed4f..0b00fca8 100644 --- a/app/src/ai/runtime/event_translator.rs +++ b/app/src/ai/runtime/event_translator.rs @@ -2,13 +2,14 @@ use std::collections::HashMap; use galaxy_agent_core::{ AgentEvent, ProviderRunOutcome, RuntimeActivity, RuntimeActivityStatus, RuntimeCapabilities, - StopReason, Usage, + StopReason, TaskListSnapshot, Usage, }; use uuid::Uuid; use warp_multi_agent_api::response_event::stream_finished; use warp_multi_agent_api::{self as api, ClientAction, ResponseEvent}; use super::provider_run_coordinator::ProviderRunProjection; +use super::task_progress::{explicit_task_updates, task_list_from_plan_batch, todo_operations}; use crate::ai::agent::runtime_activity; use crate::ai::openai::response_translator::{build_stream_finished, StreamUsage}; use crate::ai::provider::response_translator::{ @@ -54,6 +55,8 @@ pub(crate) struct ProviderRunResponseProjector { translator: RuntimeResponseTranslator, has_started_model_turn: bool, plan_tasks_projected: bool, + task_list: Option, + pending_task_list: Option, finished: bool, } @@ -63,6 +66,8 @@ impl ProviderRunResponseProjector { translator: RuntimeResponseTranslator::new(config), has_started_model_turn: false, plan_tasks_projected: false, + task_list: None, + pending_task_list: None, finished: false, } } @@ -75,10 +80,45 @@ impl ProviderRunResponseProjector { translator: RuntimeResponseTranslator::restored(config, projection_was_initialized), has_started_model_turn: false, plan_tasks_projected: true, + task_list: None, + pending_task_list: None, finished: false, } } + pub(crate) fn with_task_list(mut self, task_list: Option) -> Self { + self.plan_tasks_projected |= task_list.is_some(); + self.task_list = task_list; + self + } + + pub(crate) fn with_restored_task_lists( + mut self, + visible: Option, + desired: Option, + ) -> Self { + self.plan_tasks_projected = visible.is_some() || desired.is_some(); + self.task_list = visible; + self.pending_task_list = + desired.filter(|snapshot| self.task_list.as_ref() != Some(snapshot)); + self + } + + fn reconcile_restored_tasks(&mut self) -> Vec { + let Some(snapshot) = self.pending_task_list.take() else { + return Vec::new(); + }; + let mut events = Vec::new(); + self.translator.initialize(&mut events); + events.extend( + todo_operations(self.task_list.as_ref(), &snapshot) + .into_iter() + .map(|operation| build_todo_update(&self.translator.config.task_id, operation)), + ); + self.task_list = Some(snapshot); + events + } + pub(crate) fn project( &mut self, projection: ProviderRunProjection, @@ -86,7 +126,8 @@ impl ProviderRunResponseProjector { if self.finished { return Err("provider run projection is already finished".to_string()); } - match projection { + let mut events = self.reconcile_restored_tasks(); + let projected = match projection { ProviderRunProjection::ModelTurnStarted { .. } => { if self.has_started_model_turn { self.translator.begin_followup_turn(); @@ -111,23 +152,28 @@ impl ProviderRunResponseProjector { } ProviderRunProjection::ModelTurnFinished { .. } => Ok(Vec::new()), ProviderRunProjection::ToolBatchReady { batch } => { - if !self.plan_tasks_projected { - if let Some(todos) = todos_from_plan_batch(&batch) { - self.plan_tasks_projected = true; - return Ok(vec![build_todo_update( - &self.translator.config.task_id, - api::message::update_todos::Operation::CreateTodoList( - api::CreateTodoList { - initial_todos: todos, - }, - ), - )]); - } + let mut snapshots = explicit_task_updates(&batch); + if snapshots.is_empty() && !self.plan_tasks_projected { + snapshots.extend(task_list_from_plan_batch(&batch)); + } + let mut events = Vec::new(); + for snapshot in snapshots { + events.extend( + todo_operations(self.task_list.as_ref(), &snapshot) + .into_iter() + .map(|operation| { + build_todo_update(&self.translator.config.task_id, operation) + }), + ); + self.task_list = Some(snapshot); + self.plan_tasks_projected = true; } // Tool batches indicate runtime activity, not completion of planned work. - Ok(Vec::new()) + Ok(events) } - } + }?; + events.extend(projected); + Ok(events) } pub(crate) fn set_task_id(&mut self, task_id: impl Into) { @@ -143,19 +189,22 @@ impl ProviderRunResponseProjector { return Err("provider run projection is already finished".to_string()); } self.finished = true; - match outcome { - ProviderRunOutcome::Completed(completion) => Ok(self + let mut events = self.reconcile_restored_tasks(); + let terminal = match outcome { + ProviderRunOutcome::Completed(completion) => self .translator - .finish_provider_run(completion.stop_reason.clone(), aggregate_usage)), - ProviderRunOutcome::Failed(failure) => Ok(self.translator.provider_failure( + .finish_provider_run(completion.stop_reason.clone(), aggregate_usage), + ProviderRunOutcome::Failed(failure) => self.translator.provider_failure( &failure.message, failure.source.as_ref(), aggregate_usage, - )), - ProviderRunOutcome::Cancelled { .. } => Ok(self + ), + ProviderRunOutcome::Cancelled { .. } => self .translator - .finish_provider_run(StopReason::Cancelled, aggregate_usage)), - } + .finish_provider_run(StopReason::Cancelled, aggregate_usage), + }; + events.extend(terminal); + Ok(events) } } @@ -174,42 +223,21 @@ fn authentication_provider_for_runtime(runtime_id: &str) -> api::LlmProvider { // Only explicit multi-step checklists in a plan become tasks. Other document bullets // (risks, examples, requirements) are not executable tasks. +#[cfg(test)] fn todos_from_plan_batch( batch: &galaxy_agent_core::PendingToolBatch, ) -> Option> { - let plan_call = batch.calls.iter().find(|pending| { - matches!( - pending.call.name.as_str(), - "create_plan" | "create_documents" - ) - })?; - let content = plan_call - .call - .arguments - .get("documents")? - .as_array()? - .first()? - .get("content")? - .as_str()?; - let mut lines = content.lines(); - lines.find(|line| matches!(line.trim(), "## Tasks" | "## Implementation Tasks"))?; - let todos = lines - .take_while(|line| !line.trim().starts_with("## ")) - .filter_map(|line| line.trim().strip_prefix("- [ ]").map(str::trim)) - .filter(|item| !item.is_empty()) - .take(50) - .enumerate() - .map(|(index, item)| api::TodoItem { - id: format!("plan-{}-{index}", plan_call.call.id), - title: item - .split_once(" - ") - .map_or(item, |(title, _)| title) - .trim() - .to_owned(), - description: item.to_owned(), - }) - .collect::>(); - (todos.len() > 1).then_some(todos) + task_list_from_plan_batch(batch).map(|snapshot| { + snapshot + .tasks + .into_iter() + .map(|task| api::TodoItem { + id: task.id, + description: task.title.clone(), + title: task.title, + }) + .collect() + }) } fn build_todo_update( diff --git a/app/src/ai/runtime/event_translator_tests.rs b/app/src/ai/runtime/event_translator_tests.rs index 39180670..d36af971 100644 --- a/app/src/ai/runtime/event_translator_tests.rs +++ b/app/src/ai/runtime/event_translator_tests.rs @@ -82,6 +82,64 @@ fn restored_provider_projection_skips_stream_initialization() { )); } +#[test] +fn restored_task_progress_replay_does_not_duplicate_the_persisted_checklist() { + use galaxy_agent_core::{TaskListItem, TaskListSnapshot, TaskListStatus}; + let snapshot = TaskListSnapshot { tasks: vec![ + TaskListItem { id: "inspect".to_owned(), title: "Inspect".to_owned(), status: TaskListStatus::Completed }, + TaskListItem { id: "verify".to_owned(), title: "Verify".to_owned(), status: TaskListStatus::InProgress }, + ] }; + let mut projector = ProviderRunResponseProjector::restored(RuntimeResponseConfig { + task_id: "task".to_owned(), conversation_id: "conversation".to_owned(), needs_create_task: false, + user_query: None, model_id: "model".to_owned(), max_context_tokens: Some(1000), + capabilities: RuntimeCapabilities::provider(), empty_output_message: None, + }, true).with_task_list(Some(snapshot.clone())); + let batch = galaxy_agent_core::PendingToolBatch { + work_id: galaxy_agent_core::ExternalWorkId { + run_id: galaxy_agent_core::ProviderRunId::new("run"), epoch: galaxy_agent_core::RunEpoch::new(1), + }, + calls: vec![galaxy_agent_core::PendingToolCall { + call: galaxy_agent_core::ToolCall { + id: "progress".to_owned(), name: galaxy_agent_core::UPDATE_TASK_LIST_NAME.to_owned(), + arguments: serde_json::to_value(&snapshot).unwrap(), + }, + state: galaxy_agent_core::PendingToolCallState::Resolved { result: galaxy_agent_core::ToolResult { + call_id: "progress".to_owned(), status: galaxy_agent_core::ToolResultStatus::Success, + content: snapshot.to_tool_result_content(), + }}, + }], + }; + assert!(projector.project(ProviderRunProjection::ToolBatchReady { batch }).unwrap().is_empty()); +} + +#[test] +fn restored_committed_task_snapshot_repairs_missing_projection_even_on_terminal_cleanup() { + use galaxy_agent_core::{TaskListItem, TaskListSnapshot, TaskListStatus}; + let snapshot = TaskListSnapshot { tasks: vec![TaskListItem { + id: "verify".to_owned(), title: "Verify".to_owned(), status: TaskListStatus::Completed, + }] }; + let mut projector = ProviderRunResponseProjector::restored(RuntimeResponseConfig { + task_id: "task".to_owned(), conversation_id: "conversation".to_owned(), needs_create_task: false, + user_query: None, model_id: "model".to_owned(), max_context_tokens: Some(1000), + capabilities: RuntimeCapabilities::provider(), empty_output_message: None, + }, true).with_restored_task_lists(None, Some(snapshot)); + let events = projector.finish(&galaxy_agent_core::ProviderRunOutcome::Cancelled { + reason: "restored cancellation".to_owned(), + }, &Usage::default()).unwrap(); + let operations = events.into_iter().filter_map(|event| { + let Some(response_event::Type::ClientActions(actions)) = event.r#type else { return None }; + Some(actions.actions.into_iter().filter_map(|action| { + let Some(client_action::Action::AddMessagesToTask(add)) = action.action else { return None }; + Some(add.messages.into_iter().filter_map(|message| { + let Some(message::Message::UpdateTodos(update)) = message.message else { return None }; + update.operation + }).collect::>()) + }).flatten().collect::>()) + }).flatten().collect::>(); + assert!(matches!(&operations[0], message::update_todos::Operation::CreateTodoList(_))); + assert!(matches!(&operations[1], message::update_todos::Operation::MarkTodosCompleted(mark) if mark.todo_ids == ["verify"])); +} + #[test] fn provider_activity_does_not_create_or_complete_task_lists() { let mut projector = ProviderRunResponseProjector::new(RuntimeResponseConfig { diff --git a/app/src/ai/runtime/mod.rs b/app/src/ai/runtime/mod.rs index 42e0048a..39c64c12 100644 --- a/app/src/ai/runtime/mod.rs +++ b/app/src/ai/runtime/mod.rs @@ -4,6 +4,7 @@ mod provider_run_coordinator; mod rig; mod rig_request; mod rig_tool; +mod task_progress; pub(crate) use child_context::add_parent_context; pub(crate) use event_translator::{ @@ -19,3 +20,4 @@ pub(crate) use rig::{ PreparedProviderRun, ProviderActionContext, }; pub(crate) use rig_request::OrchestrationModelOption; +pub(crate) use task_progress::{restored_task_list, task_list_from_api_tasks}; diff --git a/app/src/ai/runtime/provider_run_coordinator.rs b/app/src/ai/runtime/provider_run_coordinator.rs index 5eb9e291..efc035ea 100644 --- a/app/src/ai/runtime/provider_run_coordinator.rs +++ b/app/src/ai/runtime/provider_run_coordinator.rs @@ -12,11 +12,12 @@ use galaxy_agent_core::{ ModelFailureDisposition, PendingToolBatch, ProviderModelCall, ProviderRequestProfile, ProviderRun, ProviderRunFailureKind, ProviderRunId, ProviderRunLimits, ProviderRunOutcome, ProviderRunPhase, ProviderRunProtocolError, ProviderRunState, ProviderRunStep, RunEpoch, - RuntimeKind, StopReason, ToolEvent, TurnCommand, TurnControl, TurnRequest, Usage, + RuntimeKind, StopReason, TaskListSnapshot, ToolEvent, TurnCommand, TurnControl, TurnRequest, Usage, }; use instant::Instant; use warpui::r#async::Timer; +use super::task_progress::{add_task_list_context, initial_task_list, task_list_from_plan_batch}; use crate::ai::agent::conversation::AIConversationId; #[cfg(not(target_family = "wasm"))] use crate::ai::model_output_logging; @@ -205,6 +206,7 @@ impl ProviderRunCoordinator { limits, ); run.replace_tool_result_archive(tool_result_archive); + run.set_task_list(initial_task_list(&base_profile.request))?; let mut profiles = BTreeMap::new(); profiles.insert(BASE_PROVIDER_PROFILE.to_string(), base_profile); Self::new(run, profiles) @@ -271,6 +273,7 @@ impl ProviderRunCoordinator { let mut request = profile.request.clone(); request.prompt = None; request.messages.clear(); + add_task_list_context(&mut request, self.run.task_list()); let fixed = u32::try_from(estimate_turn_request_tokens(&request)) .unwrap_or(u32::MAX); Some((limit, fixed)) @@ -285,6 +288,7 @@ impl ProviderRunCoordinator { let mut request = profile.request.clone(); request.prompt = None; request.messages.clear(); + add_task_list_context(&mut request, self.run.task_list()); let persistent_tokens = estimate_turn_request_tokens(&request); request.messages = self.run.transcript().to_vec(); Some(( @@ -432,6 +436,9 @@ impl ProviderRunCoordinator { } } Some(ProviderRunStep::DispatchTools(batch)) => { + if self.run.task_list().is_none() { + self.run.set_task_list(task_list_from_plan_batch(&batch))?; + } if !self.checkpoint_or_fail(&mut checkpoint).await? { return self.terminal_block(); } @@ -500,8 +507,12 @@ impl ProviderRunCoordinator { else { return false; }; - let request = - request_for_model_call(profile.request.clone(), call, Some(max_context_tokens)); + let request = request_for_model_call( + profile.request.clone(), + call, + Some(max_context_tokens), + self.run.task_list(), + ); // Provider tokenizers and request envelopes differ. Compact before the // estimate reaches the advertised limit, leaving room for that uncertainty. let safe_context_tokens = u64::from(max_context_tokens) * 9 / 10; @@ -585,6 +596,7 @@ impl ProviderRunCoordinator { profile.request.clone(), &call, profile.max_context_tokens.or(self.max_context_tokens), + self.run.task_list(), ); let started_at = Instant::now(); let (attempt_sender, attempt_control) = turn_control(); @@ -1176,6 +1188,7 @@ fn request_for_model_call( mut template: TurnRequest, call: &ProviderModelCall, max_context_tokens: Option, + task_list: Option<&TaskListSnapshot>, ) -> TurnRequest { template.prompt = None; template.messages = call.messages.clone(); @@ -1184,6 +1197,7 @@ fn request_for_model_call( message .truncate_tagged_terminal_output_for_provider_request_with_context(max_context_tokens); } + add_task_list_context(&mut template, task_list); add_exploration_reminder(&mut template); template } diff --git a/app/src/ai/runtime/provider_run_coordinator_tests.rs b/app/src/ai/runtime/provider_run_coordinator_tests.rs index a8eed5b7..a4059281 100644 --- a/app/src/ai/runtime/provider_run_coordinator_tests.rs +++ b/app/src/ai/runtime/provider_run_coordinator_tests.rs @@ -945,6 +945,50 @@ async fn inline_tool_batches_commit_before_the_next_model_boundary() { )); } +#[tokio::test] +async fn task_progress_reaches_the_next_model_after_compaction_and_restore() { + use galaxy_agent_core::{TaskListItem, TaskListSnapshot, TaskListStatus}; + use crate::ai::runtime::task_progress::{add_task_list_context, initial_task_list, task_list_tool}; + + let initial = TaskListSnapshot { tasks: vec![TaskListItem { + id: "verify".to_owned(), title: "Build the change".to_owned(), status: TaskListStatus::InProgress, + }] }; + let current = TaskListSnapshot { tasks: vec![TaskListItem { + status: TaskListStatus::Completed, ..initial.tasks[0].clone() + }] }; + let runtime = Arc::new(ScriptedRuntime::new(vec![ + tool_turn_with_calls(vec![ToolCall { + id: "progress".to_owned(), + name: galaxy_agent_core::UPDATE_TASK_LIST_NAME.to_owned(), + arguments: serde_json::to_value(¤t).unwrap(), + }]), + answer_turn(), + ])); + let mut template = request(); + template.tools.push(task_list_tool()); + add_task_list_context(&mut template, Some(&initial)); + let mut coordinator = ProviderRunCoordinator::from_request( + "run", ProviderRunProfile::new(runtime.clone(), template.clone()), Vec::new(), ProviderRunLimits::default(), + ).unwrap(); + let (_sender, control) = turn_control(); + assert!(matches!(coordinator.drive_until_blocked(control, |_| Ok(())).await.unwrap(), ProviderRunBlock::ReadyToCallModel)); + assert_eq!(coordinator.run().task_list(), Some(¤t)); + let length = coordinator.run().transcript().len(); + coordinator.run_mut().compact_transcript_at_model_boundary(length, vec![ConversationMessage { + role: MessageRole::User, content: MessageContent::Text("Prior work summarized.".to_owned()), + }]).unwrap(); + coordinator.run_mut().replace_tool_result_archive(Vec::new()); + let restored: ProviderRun = serde_json::from_str(&serde_json::to_string(coordinator.run()).unwrap()).unwrap(); + restored.validate_restored_state().unwrap(); + let mut coordinator = ProviderRunCoordinator::new(restored, BTreeMap::from([ + (BASE_PROVIDER_PROFILE.to_owned(), ProviderRunProfile::new(runtime.clone(), template)), + ])).unwrap(); + let (_sender, control) = turn_control(); + coordinator.drive_until_blocked(control, |_| Ok(())).await.unwrap(); + assert_eq!(initial_task_list(&runtime.requests()[1]), Some(current)); + assert_eq!(runtime.requests()[1].messages.len(), 1); +} + #[tokio::test] async fn recoverable_start_failure_retries_the_same_work_identity() { let mut recoverable = AgentError::new(AgentErrorKind::Transport, "temporary network error"); diff --git a/app/src/ai/runtime/rig_request.rs b/app/src/ai/runtime/rig_request.rs index 668a97c0..766e0e07 100644 --- a/app/src/ai/runtime/rig_request.rs +++ b/app/src/ai/runtime/rig_request.rs @@ -15,6 +15,9 @@ use uuid::Uuid; use warp_multi_agent_api as api; use warp_multi_agent_api::ToolType; +use super::task_progress::{ + add_task_list_context, task_list_from_api_tasks, task_list_tool, TASK_PROGRESS_GUIDANCE, +}; use crate::ai::agent::api::RequestParams; use crate::ai::agent::{AIAgentContext, AIAgentInput, MCPContext, UserQueryMode}; use crate::ai::openai::client::OpenAIClientConfig; @@ -277,7 +280,9 @@ fn prepare_rig_turn_for_provider( // History recall cannot advance a running command. Keeping it in the CLI tool list lets // the model spend its monitor turn recalling a prior snapshot instead of scheduling // `read_shell_command_output`, so make polling the only inspection path here. - tools.retain(|tool| tool.name != "recall_tool_history"); + tools.retain(|tool| { + !matches!(tool.name.as_str(), "recall_tool_history" | "update_task_list") + }); } RigRequestMode::CompletedCommandAssessment => { // The caller deliberately disables tools for the final assessment. The inline history @@ -348,6 +353,8 @@ fn prepare_rig_turn_for_provider( request.system_prompt = Some(system_prompt); request.tools = tools; request.max_output_tokens = max_output_tokens; + let task_list = task_list_from_api_tasks(&tasks, &task_id); + add_task_list_context(&mut request, task_list.as_ref()); PreparedRigTurn { max_context_tokens, @@ -692,6 +699,7 @@ fn tool_definitions( .into_iter() .filter(|tool| tool_name_is_supported(&tool.name, &supported)) .collect::>(); + tools.push(task_list_tool()); let mut mcp_tool_aliases = HashMap::new(); if !supported.contains(&ToolType::CallMcpTool) { @@ -782,7 +790,9 @@ fn add_tool_title_parameter(schema: &mut serde_json::Value) -> String { } fn tools_are_inline_only(tools: &[ToolDefinition]) -> bool { - tools.iter().all(|tool| tool.name == "recall_tool_history") + tools.iter().all(|tool| { + matches!(tool.name.as_str(), "recall_tool_history" | "update_task_list") + }) } const MAX_PROVIDER_TOOL_NAME_BYTES: usize = 64; @@ -850,6 +860,9 @@ fn build_system_prompt( ); prompt.push_str(SHELL_EXECUTION_GUIDANCE); prompt.push_str(PROGRESS_COMMUNICATION_GUIDANCE); + if tools.iter().any(|tool| tool.name == "update_task_list") { + prompt.push_str(TASK_PROGRESS_GUIDANCE); + } let contexts = inputs.iter().filter_map(AIAgentInput::context).flatten(); let mut environment = Vec::new(); let mut request_time = None; diff --git a/app/src/ai/runtime/rig_request_tests.rs b/app/src/ai/runtime/rig_request_tests.rs index fe29c588..cce253c0 100644 --- a/app/src/ai/runtime/rig_request_tests.rs +++ b/app/src/ai/runtime/rig_request_tests.rs @@ -270,8 +270,9 @@ fn no_tools_turn_flattens_historical_tool_protocol_messages() { let prepared = prepare_rig_turn(&config(), params, Vec::new(), Vec::new()); - assert_eq!(prepared.request.tools.len(), 1); - assert_eq!(prepared.request.tools[0].name, "recall_tool_history"); + assert_eq!(prepared.request.tools.len(), 2); + assert!(prepared.request.tools.iter().any(|tool| tool.name == "recall_tool_history")); + assert!(prepared.request.tools.iter().any(|tool| tool.name == "update_task_list")); assert!(prepared .request .messages @@ -499,6 +500,20 @@ fn forced_cli_profile_uses_monitor_prompt_and_tools_for_an_initial_query() { .tools .iter() .any(|tool| tool.name == "recall_tool_history")); + assert!(!prepared.request.tools.iter().any(|tool| tool.name == "update_task_list")); +} + +#[test] +fn task_progress_is_available_with_explicit_status_guidance_and_ui_title() { + let mut params = RequestParams::new_for_test(); + params.input = vec![user_query("Implement and verify the fix")]; + let prepared = prepare_rig_turn(&config(), params, vec![ToolType::ReadFiles], Vec::new()); + let tool = prepared.request.tools.iter().find(|tool| tool.name == "update_task_list").unwrap(); + assert!(tool.input_schema["required"].as_array().unwrap().iter().any(|value| value == "_galaxy_tool_title")); + let prompt = prepared.request.system_prompt.unwrap(); + assert!(prompt.contains("## Task List Progress")); + assert!(prompt.contains("reconcile the list before the final response")); + assert!(prompt.contains("Planning or returning a draft does not complete")); } #[test] diff --git a/app/src/ai/runtime/task_progress.rs b/app/src/ai/runtime/task_progress.rs new file mode 100644 index 00000000..0003720c --- /dev/null +++ b/app/src/ai/runtime/task_progress.rs @@ -0,0 +1,292 @@ +//! Task state is kept outside the compactable transcript and projected only from +//! explicit, validated model updates. Finishing a response never finishes tasks. + +use galaxy_agent_core::{ + PendingToolBatch, PendingToolCallState, ProviderRun, ProviderRunState, TaskListItem, + TaskListSnapshot, TaskListStatus, ToolDefinition, ToolResultStatus, TurnRequest, + UPDATE_TASK_LIST_NAME, +}; +use warp_multi_agent_api as api; + +const SNAPSHOT_START: &str = "\n\n"; +const SNAPSHOT_END: &str = "\n\n"; +const TASK_LIST_REMINDER: &str = "\nKeep the current task list synchronized with the work at meaningful milestones and before finishing. If a step is now done or work has moved to another step, call update_task_list with the evidence-based status. Pending work must remain visible.\n"; + +pub(super) const TASK_PROGRESS_GUIDANCE: &str = "\n## Task List Progress\nUse `update_task_list` to keep the visible checklist current throughout multi-step work. It accepts the complete list, with stable task IDs, titles, and pending/in_progress/completed statuses. Keep at most one task in progress and place it before pending tasks. After creating a plan, use its checklist and IDs from the current task list; do not create a competing list. Update when starting the next step, completing a meaningful step, changing scope, or discovering a blocker, and reconcile the list before the final response. Preserve unfinished tasks when blocked or waiting for a decision; explain the blocker in the task title if needed. Mark completed only when the work and required verification actually succeeded. Planning or returning a draft does not complete its implementation tasks. Do not recreate an unchanged list, update after every poll, or spend turns just restating progress. Tool panes show work; task updates do not require prose announcements. CLI command monitoring must continue using its advertised monitoring tools.\n"; + +pub(super) fn task_list_tool() -> ToolDefinition { + ToolDefinition { + name: UPDATE_TASK_LIST_NAME.to_owned(), + description: "Create or update the visible task list for this conversation. Send the complete list with stable IDs, retaining completed and unfinished work. Update at meaningful milestones, when the current step changes, and before finishing; mark completed only after the work and required checks succeed. This records task progress only and does not execute the tasks or edit the plan document.".to_owned(), + input_schema: serde_json::json!({ + "type": "object", + "properties": { + "tasks": { + "type": "array", "minItems": 1, "maxItems": 50, + "items": { + "type": "object", + "properties": { + "id": {"type": "string", "minLength": 1, "maxLength": 200}, + "title": {"type": "string", "minLength": 1, "maxLength": 1000}, + "status": {"type": "string", "enum": ["pending", "in_progress", "completed"]} + }, + "required": ["id", "title", "status"], + "additionalProperties": false + } + } + }, + "required": ["tasks"], + "additionalProperties": false + }), + } +} + +pub(crate) fn task_list_from_api_tasks( + tasks: &[api::Task], + root_task_id: &str, +) -> Option { + use api::message::update_todos::Operation; + + let task = tasks.iter().find(|task| task.id == root_task_id)?; + let mut pending = Vec::::new(); + let mut completed = Vec::::new(); + for message in &task.messages { + let Some(api::message::Message::UpdateTodos(update)) = &message.message else { + continue; + }; + let Some(operation) = update.operation.as_ref() else { + continue; + }; + match operation { + Operation::CreateTodoList(create) => { + pending = create.initial_todos.clone(); + completed.clear(); + } + Operation::UpdatePendingTodos(update) => { + pending = update.updated_pending_todos.clone(); + } + Operation::MarkTodosCompleted(mark) => { + for id in &mark.todo_ids { + if let Some(index) = pending.iter().position(|task| task.id == *id) { + let item = pending.remove(index); + if let Some(old) = completed.iter_mut().find(|task| task.id == *id) { + *old = item; + } else { + completed.push(item); + } + } + } + } + } + } + let snapshot = TaskListSnapshot { + tasks: completed + .into_iter() + .map(|task| task_item(task, TaskListStatus::Completed)) + .chain(pending.into_iter().enumerate().map(|(index, task)| { + task_item( + task, + if index == 0 { + TaskListStatus::InProgress + } else { + TaskListStatus::Pending + }, + ) + })) + .collect(), + }; + snapshot.validate().ok()?; + Some(snapshot) +} + +fn task_item(task: api::TodoItem, status: TaskListStatus) -> TaskListItem { + TaskListItem { + id: task.id, + title: task.title, + status, + } +} + +pub(super) fn initial_task_list(request: &TurnRequest) -> Option { + let prompt = request.system_prompt.as_deref()?; + let (_, snapshot) = prompt.rsplit_once(SNAPSHOT_START)?; + let (snapshot, _) = snapshot.split_once(SNAPSHOT_END)?; + let snapshot: TaskListSnapshot = serde_json::from_str(snapshot).ok()?; + snapshot.validate().ok()?; + Some(snapshot) +} + +pub(super) fn add_task_list_context(request: &mut TurnRequest, snapshot: Option<&TaskListSnapshot>) { + let Some(prompt) = request.system_prompt.as_mut() else { + return; + }; + if prompt.ends_with(TASK_LIST_REMINDER) { + prompt.truncate(prompt.len() - TASK_LIST_REMINDER.len()); + } + if let Some(start) = prompt.rfind(SNAPSHOT_START) { + if let Some(end) = prompt[start..].find(SNAPSHOT_END) { + prompt.replace_range(start..start + end + SNAPSHOT_END.len(), ""); + } + } + if let Some(snapshot) = snapshot { + prompt.push_str(SNAPSHOT_START); + prompt.push_str(&serde_json::to_string(snapshot).expect("serializable task list")); + prompt.push_str(SNAPSHOT_END); + } + if request + .tools + .iter() + .any(|tool| tool.name == UPDATE_TASK_LIST_NAME) + { + prompt.push_str(TASK_LIST_REMINDER); + } +} + +/// Preserve the existing plan checklist behavior, with stable IDs shared by the +/// UI and the next provider request. Rejected tools cannot create task lists. +pub(super) fn task_list_from_plan_batch(batch: &PendingToolBatch) -> Option { + let plan_call = batch.calls.iter().find(|pending| { + matches!(pending.call.name.as_str(), "create_plan" | "create_documents") + && !matches!( + &pending.state, + PendingToolCallState::Resolved { result } + if result.status != galaxy_agent_core::ToolResultStatus::Success + ) + })?; + let content = plan_call + .call + .arguments + .get("documents")? + .as_array()? + .first()? + .get("content")? + .as_str()?; + let mut lines = content.lines(); + lines.find(|line| matches!(line.trim(), "## Tasks" | "## Implementation Tasks"))?; + let tasks = lines + .take_while(|line| !line.trim().starts_with("## ")) + .filter_map(|line| line.trim().strip_prefix("- [ ]").map(str::trim)) + .filter(|item| !item.is_empty()) + .take(50) + .enumerate() + .map(|(index, item)| TaskListItem { + id: format!("plan-{}-{index}", plan_call.call.id), + title: item + .split_once(" - ") + .map_or(item, |(title, _)| title) + .trim() + .chars() + .take(1_000) + .collect(), + status: TaskListStatus::Pending, + }) + .collect::>(); + if tasks.len() < 2 { + return None; + } + let snapshot = TaskListSnapshot { tasks }; + snapshot.validate().ok()?; + Some(snapshot) +} + +fn todo_item(task: &TaskListItem) -> api::TodoItem { + api::TodoItem { + id: task.id.clone(), + title: task.title.clone(), + description: task.title.clone(), + } +} + +pub(super) fn explicit_task_updates(batch: &PendingToolBatch) -> Vec { + batch + .calls + .iter() + .filter_map(|pending| { + if pending.call.name != UPDATE_TASK_LIST_NAME { + return None; + } + let PendingToolCallState::Resolved { result } = &pending.state else { + return None; + }; + if result.status != ToolResultStatus::Success { + return None; + } + galaxy_agent_core::task_list_from_tool_result(&result.content) + }) + .collect() +} + +pub(crate) fn restored_task_list(run: &ProviderRun) -> Option { + // A mixed batch can have projected an inline update before the other tools + // finish. Conversely an inline-only batch can commit during normalization + // before its UI projection was ever persisted. Reconcile both crash windows. + if let ProviderRunState::AwaitingTools { batch } = run.state() { + if let Some(snapshot) = explicit_task_updates(batch).pop() { + return Some(snapshot); + } + } + run.task_list().cloned() +} + +pub(super) fn todo_operations( + previous: Option<&TaskListSnapshot>, + snapshot: &TaskListSnapshot, +) -> Vec { + use api::message::update_todos::Operation; + + if previous == Some(snapshot) { + return Vec::new(); + } + // Existing completed items are immutable in the transcript protocol. A new + // scope or reopening a completed item therefore starts a replacement list. + let replace = previous.is_none_or(|previous| { + previous.tasks.iter().any(|old| { + old.status == TaskListStatus::Completed + && !snapshot.tasks.iter().any(|task| { + task.id == old.id + && task.title == old.title + && task.status == TaskListStatus::Completed + }) + }) + }); + let already_completed = |id: &str| { + !replace + && previous.is_some_and(|previous| { + previous + .tasks + .iter() + .any(|task| task.id == id && task.status == TaskListStatus::Completed) + }) + }; + let todos = snapshot + .tasks + .iter() + .filter(|task| !already_completed(&task.id)) + .map(todo_item) + .collect::>(); + let mut operations = vec![if replace { + Operation::CreateTodoList(api::CreateTodoList { + initial_todos: todos, + }) + } else { + Operation::UpdatePendingTodos(api::UpdatePendingTodos { + updated_pending_todos: todos, + }) + }]; + let todo_ids = snapshot + .tasks + .iter() + .filter(|task| task.status == TaskListStatus::Completed && !already_completed(&task.id)) + .map(|task| task.id.clone()) + .collect::>(); + if !todo_ids.is_empty() { + operations.push(Operation::MarkTodosCompleted(api::MarkTodosCompleted { + todo_ids, + })); + } + operations +} + +#[cfg(test)] +#[path = "task_progress_tests.rs"] +mod tests; diff --git a/app/src/ai/runtime/task_progress_tests.rs b/app/src/ai/runtime/task_progress_tests.rs new file mode 100644 index 00000000..0c794d6b --- /dev/null +++ b/app/src/ai/runtime/task_progress_tests.rs @@ -0,0 +1,116 @@ +use super::*; + +fn snapshot(tasks: &[(&str, TaskListStatus)]) -> TaskListSnapshot { + TaskListSnapshot { + tasks: tasks.iter().map(|(id, status)| TaskListItem { + id: (*id).to_owned(), + title: format!("Task {id}"), + status: *status, + }).collect(), + } +} + +fn task_with_operations(operations: Vec) -> api::Task { + api::Task { + id: "root".to_owned(), + messages: operations.into_iter().map(|operation| api::Message { + message: Some(api::message::Message::UpdateTodos(api::message::UpdateTodos { + operation: Some(operation), + })), + ..Default::default() + }).collect(), + ..Default::default() + } +} + +#[test] +fn task_status_updates_reuse_the_list_and_survive_a_followup_request() { + use TaskListStatus::{Completed, InProgress, Pending}; + use api::message::update_todos::Operation; + + let initial = snapshot(&[("inspect", InProgress), ("implement", Pending), ("verify", Pending)]); + let next = snapshot(&[("inspect", Completed), ("implement", InProgress), ("verify", Pending)]); + let updates = todo_operations(Some(&initial), &next); + assert!(matches!(&updates[0], Operation::UpdatePendingTodos(_))); + assert!(matches!(&updates[1], Operation::MarkTodosCompleted(mark) if mark.todo_ids == ["inspect"])); + let mut operations = todo_operations(None, &initial); + operations.extend(updates); + let root = task_with_operations(operations); + let unrelated_child = api::Task { id: "child".to_owned(), ..Default::default() }; + assert_eq!(task_list_from_api_tasks(&[unrelated_child, root], "root"), Some(next.clone())); + assert!(todo_operations(Some(&next), &next).is_empty()); + + let finished = snapshot(&[("inspect", Completed), ("implement", Completed), ("verify", Completed)]); + let operations = todo_operations(Some(&next), &finished); + let Operation::UpdatePendingTodos(update) = &operations[0] else { + panic!("normal progress preserves the existing checklist"); + }; + assert_eq!(update.updated_pending_todos.len(), 2); + assert!(!update.updated_pending_todos.iter().any(|task| task.id == "inspect")); +} + +#[test] +fn reopening_completed_work_replaces_the_protocol_list_without_losing_ids() { + use TaskListStatus::{Completed, InProgress}; + let old = snapshot(&[("inspect", Completed), ("verify", Completed)]); + let reopened = snapshot(&[("inspect", Completed), ("verify", InProgress)]); + let operations = todo_operations(Some(&old), &reopened); + assert!(matches!(&operations[0], api::message::update_todos::Operation::CreateTodoList(_))); + let root = task_with_operations(operations); + assert_eq!(task_list_from_api_tasks(&[root], "root"), Some(reopened)); +} + +#[test] +fn redelivered_progress_cannot_duplicate_completed_items_in_followup_state() { + use TaskListStatus::{Completed, InProgress, Pending}; + let initial = snapshot(&[("inspect", InProgress), ("verify", Pending)]); + let next = snapshot(&[("inspect", Completed), ("verify", InProgress)]); + let update = todo_operations(Some(&initial), &next); + let mut operations = todo_operations(None, &initial); + operations.extend(update.clone()); + operations.extend(update); + assert_eq!(task_list_from_api_tasks(&[task_with_operations(operations)], "root"), Some(next)); +} + +#[test] +fn current_snapshot_replaces_stale_prompt_state_without_accumulating_reminders() { + let initial = snapshot(&[("inspect", TaskListStatus::InProgress)]); + let current = snapshot(&[("inspect", TaskListStatus::Completed)]); + let mut request = TurnRequest::new("model", Vec::new()); + request.system_prompt = Some("Follow user instructions.\n".to_owned()); + request.tools = vec![task_list_tool()]; + add_task_list_context(&mut request, Some(&initial)); + assert_eq!(initial_task_list(&request), Some(initial)); + add_task_list_context(&mut request, Some(¤t)); + let prompt = request.system_prompt.clone(); + add_task_list_context(&mut request, Some(¤t)); + assert_eq!(request.system_prompt, prompt); + assert_eq!(initial_task_list(&request), Some(current)); + let prompt = request.system_prompt.unwrap(); + assert_eq!(prompt.matches(SNAPSHOT_START).count(), 1); + assert_eq!(prompt.matches(TASK_LIST_REMINDER).count(), 1); + assert!(prompt.starts_with("Follow user instructions.")); +} + +#[test] +fn rejected_plan_creation_cannot_seed_visible_or_durable_tasks() { + let batch = PendingToolBatch { + work_id: galaxy_agent_core::ExternalWorkId { + run_id: galaxy_agent_core::ProviderRunId::new("run"), + epoch: galaxy_agent_core::RunEpoch::new(1), + }, + calls: vec![galaxy_agent_core::PendingToolCall { + call: galaxy_agent_core::ToolCall { + id: "plan".to_owned(), + name: "create_plan".to_owned(), + arguments: serde_json::json!({"documents": [{"content": "## Tasks\n- [ ] Inspect\n- [ ] Verify"}]}), + }, + state: PendingToolCallState::Resolved { result: galaxy_agent_core::ToolResult { + call_id: "plan".to_owned(), + content: "Not advertised".to_owned(), + status: galaxy_agent_core::ToolResultStatus::Error, + }}, + }], + }; + assert!(task_list_from_plan_batch(&batch).is_none()); +} diff --git a/app/src/drive/index.rs b/app/src/drive/index.rs index 21819dec..d8eca1ed 100644 --- a/app/src/drive/index.rs +++ b/app/src/drive/index.rs @@ -41,7 +41,7 @@ use super::drive_helpers::{ use super::empty_trash_confirmation_dialog::{ EmptyTrashConfirmationDialog, EmptyTrashConfirmationEvent, }; -use super::folders::CloudFolder; +use super::folders::{CloudFolder, CloudFolderModel, FolderId}; use super::items::ai_fact_collection::WarpDriveAIFactCollection; use super::items::item::{tools_panel_menu_direction, ItemStates, WarpDriveRow}; use super::items::mcp_server_collection::WarpDriveMCPServerCollection; @@ -69,7 +69,7 @@ use crate::editor::{EditorView, Event as EditorEvent, SingleLineEditorOptions}; use crate::env_vars::CloudEnvVarCollection; use crate::features::FeatureFlag; use crate::local_object_repository::{local_owner, LocalObjectRepository}; -use crate::menu::{Event, Menu, MenuItem, MenuItemFields}; +use crate::menu::{Event, Menu, MenuItem, MenuItemFields, SubMenu}; use crate::network::NetworkStatus; use crate::notebooks::CloudNotebookModel; use crate::server::cloud_objects::update_manager::{ @@ -1060,6 +1060,21 @@ impl DriveIndex { false } + fn move_operation_allowed( + &self, + cloud_object_type_and_id: &CloudObjectTypeAndId, + app: &AppContext, + ) -> bool { + let cloud_model = CloudModel::as_ref(app); + if let Some(object) = cloud_model.get_by_uid(&cloud_object_type_and_id.uid()) { + if object.permissions().owner == local_owner() { + return !object.is_trashed(cloud_model) + && !object.metadata().has_pending_online_only_change(); + } + } + self.online_only_operation_allowed(cloud_object_type_and_id, app) + } + fn is_online(&self, app: &AppContext) -> bool { NetworkStatus::as_ref(app).is_online() } @@ -2685,7 +2700,7 @@ impl DriveIndex { // If the share dialog is open, we don't want to open the menu for the same object. let menu_open = self.menu_object_id_if_open == Some(warp_drive_item_id) && !share_dialog_open; - let can_move = self.online_only_operation_allowed(&row_object_id, app) + let can_move = self.move_operation_allowed(&row_object_id, app) && matches!(self.index_variant, DriveIndexVariant::MainIndex) && access_level.can_move_drive(); let mut is_focused = false; @@ -3239,12 +3254,15 @@ impl DriveIndex { ) { // TODO: we need to tell the user why 'move' is not working. let is_drop_target_valid = match new_location { - CloudObjectLocation::Folder(folder_id) => self.online_only_operation_allowed( - &CloudObjectTypeAndId::from_id_and_type(folder_id, ObjectType::Folder), - ctx, - ), - CloudObjectLocation::Space(_) => self.is_online(ctx), - CloudObjectLocation::Trash => self.is_online(ctx), + CloudObjectLocation::Folder(folder_id) => { + Self::is_local_folder(&folder_id, ctx) + || self.online_only_operation_allowed( + &CloudObjectTypeAndId::from_id_and_type(folder_id, ObjectType::Folder), + ctx, + ) + } + CloudObjectLocation::Space(space) => space == Space::Personal || self.is_online(ctx), + CloudObjectLocation::Trash => true, }; if is_drop_target_valid { @@ -3275,23 +3293,37 @@ impl DriveIndex { self.current_drop_target = None; ctx.notify(); - if !self.online_only_operation_allowed(cloud_object_type_and_id, ctx) { + if !self.move_operation_allowed(cloud_object_type_and_id, ctx) { return; } let cloud_model = CloudModel::handle(ctx); + let is_local = cloud_model + .as_ref(ctx) + .get_by_uid(&cloud_object_type_and_id.uid()) + .is_some_and(|object| object.permissions().owner == local_owner()); + let can_move = if is_local { + LocalObjectRepository::can_move_object_to_location( + *cloud_object_type_and_id, + new_location, + ctx, + ) + } else { + cloud_model.as_ref(ctx).can_move_object_to_location( + &cloud_object_type_and_id.uid(), + new_location, + ctx, + ) + }; // Only proceed if we can move the object to this location AND the operation results in a change. // Even though an object technically can be moved to its current location (can_move_object_to_location returns true) // we do not need to do any of the update logic that follows. - if !cloud_model.as_ref(ctx).can_move_object_to_location( - &cloud_object_type_and_id.uid(), - new_location, - ctx, - ) || cloud_model - .as_ref(ctx) - .object_location(&cloud_object_type_and_id.uid(), ctx) - .is_none_or(|location| location == new_location) + if !can_move + || cloud_model + .as_ref(ctx) + .object_location(&cloud_object_type_and_id.uid(), ctx) + .is_none_or(|location| location == new_location) { return; } @@ -3329,9 +3361,15 @@ impl DriveIndex { } // Otherwise allow object move to go through. - UpdateManager::handle(ctx).update(ctx, move |update_manager, ctx| { - update_manager.move_object_to_location(*cloud_object_type_and_id, new_location, ctx); - }); + if is_local { + LocalObjectRepository::handle(ctx).update(ctx, |repository, ctx| { + repository.move_object_to_location(*cloud_object_type_and_id, new_location, ctx); + }); + } else { + UpdateManager::handle(ctx).update(ctx, move |update_manager, ctx| { + update_manager.move_object_to_location(*cloud_object_type_and_id, new_location, ctx); + }); + } match new_location { CloudObjectLocation::Space(space) => self.open_section_of_space(space), @@ -4873,6 +4911,13 @@ impl DriveIndex { } } + if self.move_operation_allowed(cloud_object_type_and_id, app) && access_level.can_move_drive() + { + if let Some(move_menu) = self.local_move_menu_item(*cloud_object_type_and_id, app) { + menu_items.push(move_menu); + } + } + if can_trash && (!FeatureFlag::SharedWithMe.is_enabled() || access_level.can_trash()) { menu_items.push( MenuItemFields::new("Trash") @@ -4887,6 +4932,65 @@ impl DriveIndex { menu_items } + fn local_move_menu_item( + &self, + object_id: CloudObjectTypeAndId, + app: &AppContext, + ) -> Option> { + let cloud_model = CloudModel::as_ref(app); + let object = cloud_model.get_by_uid(&object_id.uid())?; + if object.permissions().owner != local_owner() { + return None; + } + + let mut destinations = Vec::new(); + if object.metadata().folder_id.is_some() { + destinations.push(( + "Personal (top level)".to_string(), + CloudObjectLocation::Space(Space::Personal), + )); + } + destinations.extend( + cloud_model + .get_all_objects_of_type::() + .filter(|folder| object.metadata().folder_id != Some(folder.id)) + .filter(|folder| { + LocalObjectRepository::can_move_object_to_location( + object_id, + CloudObjectLocation::Folder(folder.id), + app, + ) + }) + .map(|folder| { + ( + format!("{} / {}", folder.breadcrumbs(app), folder.model().name), + CloudObjectLocation::Folder(folder.id), + ) + }) + .sorted_by(|(left, _), (right, _)| left.cmp(right)), + ); + if destinations.is_empty() { + return None; + } + + let items = destinations + .into_iter() + .map(|(label, location)| { + MenuItemFields::new(label) + .with_on_select_action(DriveIndexAction::DropIndexItem { + cloud_object_type_and_id: object_id, + drop_target_location: location, + }) + .with_icon(Icon::Folder) + .into_item() + }) + .collect(); + Some(MenuItem::Submenu { + fields: MenuItemFields::new_submenu("Move to").with_icon(Icon::Move), + menu: SubMenu::new(items), + }) + } + /// Builder for a menu item to open a Warp Drive object in a pane. The icon and label depend /// on whether the object is editable or not. /// diff --git a/app/src/drive/index_tests.rs b/app/src/drive/index_tests.rs index dc5d9110..99fd5446 100644 --- a/app/src/drive/index_tests.rs +++ b/app/src/drive/index_tests.rs @@ -11,10 +11,12 @@ use crate::cloud_object::model::actions::ObjectActions; use crate::cloud_object::model::persistence::CloudModel; use crate::cloud_object::model::view::CloudViewModel; use crate::cloud_object::{ - CloudObjectSyncStatus, ObjectIdType, ObjectType, Owner, ServerCreationInfo, Space, + CloudObjectLocation, CloudObjectSyncStatus, ObjectIdType, ObjectType, Owner, + ServerCreationInfo, Space, }; use crate::drive::items::WarpDriveItemId; use crate::drive::CloudObjectTypeAndId; +use crate::local_object_repository::LocalObjectRepository; use crate::menu::MenuItem; use crate::network::NetworkStatus; use crate::notebooks::{CloudNotebook, CloudNotebookModel}; @@ -125,6 +127,69 @@ fn label_for_menu_item(item: &MenuItem) -> &str { } } +#[test] +fn local_drive_objects_can_move_into_and_out_of_plans_offline() { + App::test(ASSETS, |mut app| async move { + initialize_app(&mut app); + let repository = LocalObjectRepository::handle(&app); + let plans_id = SyncId::ClientId(ClientId::new()); + let notebook_id = SyncId::ClientId(ClientId::new()); + repository.update(&mut app, |repository, ctx| { + repository.create_folder_with_id(plans_id, "Plans".to_string(), None, ctx); + repository.create_notebook_with_id( + notebook_id, + None, + CloudNotebookModel::default(), + ctx, + ); + }); + NetworkStatus::handle(&app).update(&mut app, |network_status, ctx| { + network_status.reachability_changed(false, ctx); + }); + let index = create_index(&mut app); + let object_id = CloudObjectTypeAndId::Notebook(notebook_id); + + index.update(&mut app, |index, ctx| { + assert!(index.move_operation_allowed(&object_id, ctx)); + index.update_drop_target_location(CloudObjectLocation::Folder(plans_id), ctx); + assert_eq!( + index.current_drop_target, + Some(CloudObjectLocation::Folder(plans_id)), + ); + let Some(MenuItem::Submenu { fields, menu }) = + index.local_move_menu_item(object_id, ctx) + else { + panic!("expected a local Move to submenu"); + }; + assert_eq!(fields.label(), "Move to"); + assert_eq!(label_for_menu_item(&menu.items()[0]), "Personal / Plans"); + + index.move_object(&object_id, CloudObjectLocation::Folder(plans_id), ctx); + assert_eq!( + CloudModel::as_ref(ctx) + .get_notebook(¬ebook_id) + .unwrap() + .metadata + .folder_id, + Some(plans_id), + ); + let Some(MenuItem::Submenu { menu, .. }) = index.local_move_menu_item(object_id, ctx) + else { + panic!("expected moving the notebook back to the top level to be offered"); + }; + assert_eq!(label_for_menu_item(&menu.items()[0]), "Personal (top level)"); + + index.move_object(&object_id, CloudObjectLocation::Space(Space::Personal), ctx); + assert!(CloudModel::as_ref(ctx) + .get_notebook(¬ebook_id) + .unwrap() + .metadata + .folder_id + .is_none()); + }); + }); +} + #[test] fn test_retry_menu_item_visibility() { App::test(ASSETS, |mut app| async move { diff --git a/app/src/drive/items/item.rs b/app/src/drive/items/item.rs index 9c1e5931..a65c8ddf 100644 --- a/app/src/drive/items/item.rs +++ b/app/src/drive/items/item.rs @@ -27,6 +27,7 @@ use crate::drive::index::{ use crate::drive::items::WarpDriveItem; use crate::drive::panel::WARP_DRIVE_POSITION_ID; use crate::drive::CloudObjectTypeAndId; +use crate::local_object_repository::{local_owner, LocalObjectRepository}; use crate::menu::Menu; use crate::ui_components::blended_colors; use crate::ui_components::icons::{Icon, ICON_DIMENSIONS}; @@ -826,11 +827,24 @@ impl UiComponent for WarpDriveRow<'_> { return AcceptedByDropTarget::No; }; let cloud_model = CloudModel::handle(app); - if cloud_model.as_ref(app).can_move_object_to_location( - &item.uid(), - *location, - app, - ) { + let is_local = cloud_model + .as_ref(app) + .get_by_uid(&item.uid()) + .is_some_and(|object| { + object.permissions().owner == local_owner() + }); + let can_move = if is_local { + LocalObjectRepository::can_move_object_to_location( + item, *location, app, + ) + } else { + cloud_model.as_ref(app).can_move_object_to_location( + &item.uid(), + *location, + app, + ) + }; + if can_move { AcceptedByDropTarget::Yes } else { AcceptedByDropTarget::No diff --git a/app/src/lib.rs b/app/src/lib.rs index d8ded2d4..4ce9613b 100644 --- a/app/src/lib.rs +++ b/app/src/lib.rs @@ -2079,6 +2079,7 @@ pub(crate) fn initialize_app( // AIDocumentModel subscribes to UpdateManager so that it can be notified when notebooks are created on the server. ctx.add_singleton_model(AIDocumentModel::new); + ctx.add_singleton_model(ai::document::plan_archival::PlanArchivalModel::new); // AgentConversationsModel subscribes to UpdateManager for RTC task updates. ctx.add_singleton_model(AgentConversationsModel::new); diff --git a/app/src/local_object_repository.rs b/app/src/local_object_repository.rs index 90fd0a52..0041bf23 100644 --- a/app/src/local_object_repository.rs +++ b/app/src/local_object_repository.rs @@ -1,3 +1,4 @@ +use std::collections::HashSet; use std::sync::mpsc::SyncSender; use chrono::Utc; @@ -13,7 +14,7 @@ use crate::ai::mcp::templatable::{ }; use crate::auth::UserUid; use crate::cloud_object::model::generic_string_model::GenericStringObjectId; -use crate::cloud_object::model::persistence::{CloudModel, CloudModelEvent}; +use crate::cloud_object::model::persistence::{CloudModel, CloudModelEvent, UpdateSource}; use crate::cloud_object::{ CloudObject, CloudObjectLocation, CloudObjectMetadata, CloudObjectPermissions, CloudObjectStatuses, CloudObjectSyncStatus, GenericCloudObject, GenericStringObjectFormat, @@ -50,6 +51,7 @@ pub enum LocalObjectRepositoryEvent { Rules, ExecutionProfiles { id: Option }, Notebooks { id: Option }, + NotebookContentEdited { id: SyncId }, Workflows { id: Option }, } @@ -219,6 +221,95 @@ impl LocalObjectRepository { CloudModel::as_ref(app).get_folder(id).cloned() } + /// Local moves never require a server ID, connectivity, or a cloud ownership transfer. + pub fn can_move_object_to_location( + object_id: CloudObjectTypeAndId, + new_location: CloudObjectLocation, + app: &AppContext, + ) -> bool { + let cloud_model = CloudModel::as_ref(app); + let Some(object) = cloud_model.get_by_uid(&object_id.uid()) else { + return false; + }; + if object.permissions().owner != local_owner() || object.is_trashed(cloud_model) { + return false; + } + + match new_location { + CloudObjectLocation::Space(space) => space == Space::Personal, + CloudObjectLocation::Trash => true, + CloudObjectLocation::Folder(folder_id) => { + let mut ancestor_id = Some(folder_id); + let mut visited = HashSet::new(); + while let Some(id) = ancestor_id { + if object_id == CloudObjectTypeAndId::Folder(id) || !visited.insert(id) { + return false; + } + let Some(folder) = cloud_model.get_folder(&id) else { + return false; + }; + if folder.permissions.owner != local_owner() + || folder.metadata.trashed_ts.is_some() + { + return false; + } + ancestor_id = folder.metadata.folder_id; + } + true + } + } + } + + /// Reparent a local object and persist its metadata without enqueuing server work. + pub fn move_object_to_location( + &mut self, + object_id: CloudObjectTypeAndId, + new_location: CloudObjectLocation, + ctx: &mut ModelContext, + ) -> bool { + if !Self::can_move_object_to_location(object_id, new_location, ctx) { + return false; + } + + let event = CloudModel::handle(ctx).update(ctx, |cloud_model, ctx| { + let object = cloud_model.get_mut_by_uid(&object_id.uid())?; + let from_folder = object.metadata().folder_id; + let to_folder = match new_location { + CloudObjectLocation::Space(_) => None, + CloudObjectLocation::Folder(folder_id) => Some(folder_id), + CloudObjectLocation::Trash => from_folder, + }; + let event = if new_location == CloudObjectLocation::Trash { + object.metadata_mut().trashed_ts = Some(ServerTimestamp::new(Utc::now())); + CloudModelEvent::ObjectTrashed { + type_and_id: object_id, + source: UpdateSource::Local, + } + } else { + if from_folder == to_folder { + return None; + } + object.metadata_mut().folder_id = to_folder; + CloudModelEvent::ObjectMoved { + type_and_id: object_id, + source: UpdateSource::Local, + from_folder, + to_folder, + } + }; + set_locally_saved_metadata(object.metadata_mut()); + let persistence_event = object.upsert_event(); + ctx.emit(event); + ctx.notify(); + Some(persistence_event) + }); + let Some(event) = event else { + return false; + }; + self.save(event); + true + } + pub fn create_folder_with_id( &mut self, id: SyncId, @@ -383,7 +474,10 @@ impl LocalObjectRepository { let Some(notebook) = self.notebook(&id, ctx) else { return false; }; - self.update_notebook( + if notebook.model().data == data { + return true; + } + let updated = self.update_notebook( id, CloudNotebookModel { title: notebook.model().title.clone(), @@ -392,7 +486,11 @@ impl LocalObjectRepository { conversation_id: notebook.model().conversation_id.clone(), }, ctx, - ) + ); + if updated { + ctx.emit(LocalObjectRepositoryEvent::NotebookContentEdited { id }); + } + updated } pub fn update_notebook_title( diff --git a/app/src/local_object_repository_tests.rs b/app/src/local_object_repository_tests.rs index 6e25e0b5..14668976 100644 --- a/app/src/local_object_repository_tests.rs +++ b/app/src/local_object_repository_tests.rs @@ -172,6 +172,269 @@ fn create_update_and_delete_notebook_are_local_and_persisted() { }); } +#[test] +fn notebook_content_edit_events_exclude_title_edits_and_unchanged_autosaves() { + App::test((), |mut app| async move { + let receiver = initialize_app(&mut app); + let repository = LocalObjectRepository::handle(&app); + let id = SyncId::ClientId(ClientId::new()); + repository.update(&mut app, |repository, ctx| { + repository.create_notebook_with_id( + id, + None, + CloudNotebookModel { + title: "Plan".to_string(), + data: "Initial content".to_string(), + ..Default::default() + }, + ctx, + ); + }); + receiver.recv().unwrap(); + let (edit_sender, edits) = sync_channel(10); + app.update(|ctx| { + ctx.subscribe_to_model(&repository, move |_, event, _| { + if let LocalObjectRepositoryEvent::NotebookContentEdited { id } = event { + edit_sender.send(*id).unwrap(); + } + }); + }); + + assert!(repository.update(&mut app, |repository, ctx| { + repository.update_notebook_title(id, "Renamed plan".to_string(), ctx) + })); + receiver.recv().unwrap(); + assert!(edits.try_recv().is_err()); + + assert!(repository.update(&mut app, |repository, ctx| { + repository.update_notebook_data(id, "Initial content".to_string(), ctx) + })); + assert!(receiver.try_recv().is_err()); + assert!(edits.try_recv().is_err()); + + assert!(repository.update(&mut app, |repository, ctx| { + repository.update_notebook_data(id, "Changed content".to_string(), ctx) + })); + assert!(matches!( + receiver.recv().unwrap(), + ModelEvent::UpsertNotebook { notebook } if notebook.model().data == "Changed content" + )); + assert_eq!(edits.try_recv().unwrap(), id); + assert!(edits.try_recv().is_err()); + + assert!(!repository.update(&mut app, |repository, ctx| { + repository.update_notebook_data( + SyncId::ClientId(ClientId::new()), + "Missing notebook".to_string(), + ctx, + ) + })); + assert!(receiver.try_recv().is_err()); + assert!(edits.try_recv().is_err()); + }); +} + +#[test] +fn move_local_plan_between_folders_and_root_is_persisted() { + App::test((), |mut app| async move { + let receiver = initialize_app(&mut app); + let repository = LocalObjectRepository::handle(&app); + let plans_id = SyncId::ClientId(ClientId::new()); + let projects_id = SyncId::ClientId(ClientId::new()); + let notebook_id = SyncId::ClientId(ClientId::new()); + let document_id = crate::ai::document::ai_document_model::AIDocumentId::new(); + repository.update(&mut app, |repository, ctx| { + repository.create_folder_with_id(plans_id, "Plans".to_string(), None, ctx); + repository.create_folder_with_id(projects_id, "Projects".to_string(), None, ctx); + repository.create_notebook_with_id( + notebook_id, + Some(plans_id), + CloudNotebookModel { + title: "Release plan".to_string(), + data: "- [ ] Ship the release".to_string(), + ai_document_id: Some(document_id), + conversation_id: None, + }, + ctx, + ); + }); + for _ in 0..3 { + receiver.recv().unwrap(); + } + + for (location, expected_parent) in [ + (CloudObjectLocation::Folder(projects_id), Some(projects_id)), + (CloudObjectLocation::Space(Space::Personal), None), + (CloudObjectLocation::Folder(plans_id), Some(plans_id)), + ] { + assert!(repository.update(&mut app, |repository, ctx| { + repository.move_object_to_location( + CloudObjectTypeAndId::Notebook(notebook_id), + location, + ctx, + ) + })); + let ModelEvent::UpsertNotebook { notebook } = receiver.recv().unwrap() else { + panic!("expected the moved plan to be persisted locally"); + }; + assert_eq!(notebook.id, notebook_id); + assert_eq!(notebook.metadata.folder_id, expected_parent); + assert_eq!(notebook.model().ai_document_id, Some(document_id)); + assert_eq!(notebook.model().data, "- [ ] Ship the release"); + assert!(!notebook.metadata.has_pending_content_changes()); + assert!(!notebook.metadata.has_pending_online_only_change()); + assert_eq!( + repository.read(&app, |repository, app| { + repository + .notebook(¬ebook_id, app) + .unwrap() + .metadata + .folder_id + }), + expected_parent, + ); + } + }); +} + +#[test] +fn local_folder_moves_reject_cycles_and_missing_or_trashed_destinations() { + App::test((), |mut app| async move { + let receiver = initialize_app(&mut app); + let repository = LocalObjectRepository::handle(&app); + let parent_id = SyncId::ClientId(ClientId::new()); + let child_id = SyncId::ClientId(ClientId::new()); + let trashed_id = SyncId::ClientId(ClientId::new()); + repository.update(&mut app, |repository, ctx| { + repository.create_folder_with_id(parent_id, "Plans".to_string(), None, ctx); + repository.create_folder_with_id( + child_id, + "Nested".to_string(), + Some(parent_id), + ctx, + ); + repository.create_folder_with_id(trashed_id, "Deleted".to_string(), None, ctx); + repository.set_folder_trashed(trashed_id, true, ctx); + }); + for _ in 0..4 { + receiver.recv().unwrap(); + } + + for location in [ + CloudObjectLocation::Folder(parent_id), + CloudObjectLocation::Folder(child_id), + CloudObjectLocation::Folder(trashed_id), + CloudObjectLocation::Folder(SyncId::ClientId(ClientId::new())), + CloudObjectLocation::Space(Space::Shared), + ] { + assert!(!repository.update(&mut app, |repository, ctx| { + repository.move_object_to_location( + CloudObjectTypeAndId::Folder(parent_id), + location, + ctx, + ) + })); + } + assert!(receiver.try_recv().is_err()); + + assert!(repository.update(&mut app, |repository, ctx| { + repository.move_object_to_location( + CloudObjectTypeAndId::Folder(child_id), + CloudObjectLocation::Space(Space::Personal), + ctx, + ) + })); + assert!(matches!( + receiver.recv().unwrap(), + ModelEvent::UpsertFolder { folder } + if folder.id == child_id && folder.metadata.folder_id.is_none() + )); + }); +} + +#[test] +fn moving_local_objects_to_trash_is_persisted() { + App::test((), |mut app| async move { + let receiver = initialize_app(&mut app); + let repository = LocalObjectRepository::handle(&app); + let id = SyncId::ClientId(ClientId::new()); + repository.update(&mut app, |repository, ctx| { + repository.create_workflow_with_id( + id, + None, + Workflow::new("Build", "cargo build"), + ctx, + ); + }); + receiver.recv().unwrap(); + + assert!(repository.update(&mut app, |repository, ctx| { + repository.move_object_to_location( + CloudObjectTypeAndId::Workflow(id), + CloudObjectLocation::Trash, + ctx, + ) + })); + assert!(matches!( + receiver.recv().unwrap(), + ModelEvent::UpsertWorkflow { workflow } + if workflow.id == id && workflow.metadata.trashed_ts.is_some() + )); + }); +} + +#[test] +fn local_moves_do_not_transfer_objects_across_ownership_boundaries() { + App::test((), |mut app| async move { + let receiver = initialize_app(&mut app); + let repository = LocalObjectRepository::handle(&app); + let local_id = SyncId::ClientId(ClientId::new()); + let cloud_id = ClientId::new(); + let cloud_folder_id = ClientId::new(); + repository.update(&mut app, |repository, ctx| { + repository.create_notebook_with_id(local_id, None, CloudNotebookModel::default(), ctx); + }); + receiver.recv().unwrap(); + CloudModel::handle(&app).update(&mut app, |cloud_model, ctx| { + cloud_model.create_object( + SyncId::ClientId(cloud_id), + CloudNotebook::new_local( + CloudNotebookModel::default(), + Owner::mock_current_user(), + None, + cloud_id, + ), + ctx, + ); + cloud_model.create_object( + SyncId::ClientId(cloud_folder_id), + CloudFolder::new_local( + CloudFolderModel::new("Shared", false), + Owner::mock_current_user(), + None, + cloud_folder_id, + ), + ctx, + ); + }); + assert!(!repository.update(&mut app, |repository, ctx| { + repository.move_object_to_location( + CloudObjectTypeAndId::Notebook(local_id), + CloudObjectLocation::Folder(SyncId::ClientId(cloud_folder_id)), + ctx, + ) + })); + assert!(!repository.update(&mut app, |repository, ctx| { + repository.move_object_to_location( + CloudObjectTypeAndId::Notebook(SyncId::ClientId(cloud_id)), + CloudObjectLocation::Space(Space::Personal), + ctx, + ) + })); + assert!(receiver.try_recv().is_err()); + }); +} + #[test] fn create_update_open_trash_and_delete_folder_are_local_and_persisted() { App::test((), |mut app| async move { diff --git a/crates/galaxy_agent_core/src/lib.rs b/crates/galaxy_agent_core/src/lib.rs index 231b6d69..7fcd6cb0 100644 --- a/crates/galaxy_agent_core/src/lib.rs +++ b/crates/galaxy_agent_core/src/lib.rs @@ -6,10 +6,12 @@ mod provider_run; mod runtime; +mod task_list; mod tool_policy; mod types; pub use provider_run::*; pub use runtime::*; +pub use task_list::*; pub use tool_policy::*; pub use types::*; diff --git a/crates/galaxy_agent_core/src/provider_run.rs b/crates/galaxy_agent_core/src/provider_run.rs index 7a8c2533..16a8a7dc 100644 --- a/crates/galaxy_agent_core/src/provider_run.rs +++ b/crates/galaxy_agent_core/src/provider_run.rs @@ -6,8 +6,8 @@ use serde::{Deserialize, Serialize}; use crate::{ AgentError, ContentPart, ConversationMessage, MessageContent, MessageRole, PermissionDecision, - PermissionRequest, StopReason, ToolCall, ToolCallDecision, ToolPolicy, ToolResult, - ToolResultStatus, Usage, + PermissionRequest, StopReason, TaskListSnapshot, ToolCall, ToolCallDecision, ToolPolicy, + ToolResult, ToolResultStatus, UPDATE_TASK_LIST_NAME, Usage, task_list_from_tool_result, }; #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] @@ -357,6 +357,9 @@ pub enum ProviderRunProtocolError { InvalidTranscriptCompaction { message: String, }, + InvalidTaskList { + message: String, + }, InvalidRestoredState { message: String, }, @@ -422,6 +425,7 @@ impl fmt::Display for ProviderRunProtocolError { Self::InvalidTranscriptCompaction { message } => { write!(f, "invalid transcript compaction: {message}") } + Self::InvalidTaskList { message } => write!(f, "invalid task list: {message}"), Self::InvalidRestoredState { message } => { write!(f, "invalid restored provider run: {message}") } @@ -439,6 +443,8 @@ pub struct ProviderRun { epoch: RunEpoch, transcript: Vec, tool_result_archive: Vec, + #[serde(default)] + task_list: Option, usage: Usage, model_turns: u32, model_retries: u32, @@ -460,6 +466,7 @@ impl ProviderRun { epoch: RunEpoch::default(), transcript, tool_result_archive: Vec::new(), + task_list: None, usage: Usage::default(), model_turns: 0, model_retries: 0, @@ -489,6 +496,23 @@ impl ProviderRun { &self.transcript } + pub fn task_list(&self) -> Option<&TaskListSnapshot> { + self.task_list.as_ref() + } + + pub fn set_task_list( + &mut self, + task_list: Option, + ) -> Result<(), ProviderRunProtocolError> { + if let Some(snapshot) = &task_list { + snapshot + .validate() + .map_err(|message| ProviderRunProtocolError::InvalidTaskList { message })?; + } + self.task_list = task_list; + Ok(()) + } + pub fn tool_result_archive(&self) -> &[ConversationMessage] { &self.tool_result_archive } @@ -564,6 +588,11 @@ impl ProviderRun { if self.profile.as_str().is_empty() { return Err(invalid("request profile must not be empty".to_string())); } + if let Some(snapshot) = &self.task_list { + snapshot + .validate() + .map_err(|error| invalid(format!("invalid task list: {error}")))?; + } if self.limits.max_model_turns == Some(0) { return Err(invalid("model-turn limit must be at least one".to_string())); } @@ -1273,6 +1302,20 @@ impl ProviderRun { }); }; let next_epoch = self.next_epoch()?; + let mut task_list = self.task_list.clone(); + for pending in &batch.calls { + if pending.call.name != UPDATE_TASK_LIST_NAME { + continue; + } + let result = pending.state.result().expect("validated complete tool batch"); + if result.status == ToolResultStatus::Success { + task_list = Some(task_list_from_tool_result(&result.content).ok_or_else(|| { + ProviderRunProtocolError::InvalidTaskList { + message: "successful task update has no valid snapshot result".to_string(), + } + })?); + } + } let yielded_to_children = batch.calls.iter().any(|pending| { is_local_agent_launch(&pending.call) && pending @@ -1281,6 +1324,7 @@ impl ProviderRun { .is_some_and(|result| result.status == ToolResultStatus::Success) }); self.transcript.push(tool_result_message(results)); + self.task_list = task_list; self.epoch = next_epoch; // Launch acknowledgement is a yield, not an invitation to issue more tools. // The client persists child membership and supplies a new input once the group finishes. @@ -1634,6 +1678,18 @@ fn validate_pending_tool_state(pending: &PendingToolCall) -> Result<(), String> | PendingToolCallState::RecoveryPending | PendingToolCallState::Resolved { .. } => {} } + if pending.call.name == UPDATE_TASK_LIST_NAME { + if let Some(result) = pending.state.result() { + if result.status == ToolResultStatus::Success { + let snapshot = task_list_from_tool_result(&result.content) + .ok_or_else(|| "successful task update has no valid snapshot result".to_string())?; + let expected = TaskListSnapshot::from_arguments(&pending.call.arguments)?; + if snapshot != expected { + return Err("task update result does not match its arguments".to_string()); + } + } + } + } Ok(()) } diff --git a/crates/galaxy_agent_core/src/provider_run_tests.rs b/crates/galaxy_agent_core/src/provider_run_tests.rs index 9267f3e9..da299550 100644 --- a/crates/galaxy_agent_core/src/provider_run_tests.rs +++ b/crates/galaxy_agent_core/src/provider_run_tests.rs @@ -1,7 +1,7 @@ use serde_json::json; use super::*; -use crate::{AgentErrorKind, PermissionKind}; +use crate::{AgentErrorKind, PermissionKind, TaskListItem, TaskListStatus}; fn initial_messages() -> Vec { vec![ConversationMessage { @@ -354,6 +354,201 @@ fn unknown_malformed_and_inline_tools_are_pre_resolved_in_the_same_batch() { assert_eq!(batch.unresolved_call_ids(), vec!["external"]); } +fn task_snapshot(status: TaskListStatus) -> TaskListSnapshot { + TaskListSnapshot { + tasks: vec![TaskListItem { + id: "fix".to_string(), + title: "Fix the failure".to_string(), + status, + }], + } +} + +fn update_task_call(id: &str, snapshot: &TaskListSnapshot) -> ToolCall { + ToolCall { + id: id.to_string(), + name: UPDATE_TASK_LIST_NAME.to_string(), + arguments: serde_json::to_value(snapshot).unwrap(), + } +} + +#[test] +fn inline_task_updates_commit_atomically_in_call_order_with_other_tools() { + let mut run = run(); + let initial = task_snapshot(TaskListStatus::Pending); + let working = task_snapshot(TaskListStatus::InProgress); + let completed = task_snapshot(TaskListStatus::Completed); + run.set_task_list(Some(initial.clone())).unwrap(); + let batch = accept_tool_turn( + &mut run, + tool_turn( + vec![ + update_task_call("tasks-start", &working), + tool_call("external", "read_files"), + update_task_call("tasks-finish", &completed), + ], + &[UPDATE_TASK_LIST_NAME, "read_files"], + ), + ); + + assert_eq!(batch.unresolved_call_ids(), vec!["external"]); + assert_eq!(run.task_list(), Some(&initial)); + assert!(matches!( + run.commit_tool_batch(&batch.work_id), + Err(ProviderRunProtocolError::IncompleteToolBatch { .. }), + )); + assert_eq!(run.task_list(), Some(&initial)); + + run.complete_tool(&batch.work_id, successful_result("external", "contents")) + .unwrap(); + run.commit_tool_batch(&batch.work_id).unwrap(); + + assert_eq!(run.task_list(), Some(&completed)); + assert_eq!(run.state().phase(), ProviderRunPhase::ReadyToCallModel); + assert_serialization_round_trip(&run); + let committed = run.clone(); + assert!(run.commit_tool_batch(&batch.work_id).is_err()); + assert_eq!(run, committed); +} + +#[test] +fn rejected_task_updates_do_not_replace_durable_progress() { + let mut run = run(); + let initial = task_snapshot(TaskListStatus::InProgress); + run.set_task_list(Some(initial.clone())).unwrap(); + let batch = accept_tool_turn( + &mut run, + tool_turn( + vec![ToolCall { + id: "invalid-task-update".to_string(), + name: UPDATE_TASK_LIST_NAME.to_string(), + arguments: json!({"tasks": []}), + }], + &[UPDATE_TASK_LIST_NAME], + ), + ); + assert_eq!( + batch.calls[0].state.result().unwrap().status, + ToolResultStatus::Error, + ); + run.commit_tool_batch(&batch.work_id).unwrap(); + assert_eq!(run.task_list(), Some(&initial)); + + let batch = accept_tool_turn( + &mut run, + tool_turn( + vec![update_task_call( + "unadvertised-update", + &task_snapshot(TaskListStatus::Completed), + )], + &[], + ), + ); + run.commit_tool_batch(&batch.work_id).unwrap(); + assert_eq!(run.task_list(), Some(&initial)); +} + +#[test] +fn task_progress_survives_compaction_archive_eviction_and_restore() { + let mut run = run(); + let snapshot = task_snapshot(TaskListStatus::InProgress); + let batch = accept_tool_turn( + &mut run, + tool_turn( + vec![update_task_call("tasks", &snapshot)], + &[UPDATE_TASK_LIST_NAME], + ), + ); + assert!(batch.is_complete()); + run.commit_tool_batch(&batch.work_id).unwrap(); + run.compact_transcript_at_model_boundary( + run.transcript().len(), + vec![ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text("Earlier conversation summary".to_string()), + }], + ) + .unwrap(); + + let mut archive = run.tool_result_archive().to_vec(); + for index in 0..450 { + archive.extend([ + ConversationMessage { + role: MessageRole::Assistant, + content: MessageContent::ToolUse { + tool_use_id: format!("read-{index}"), + name: "read_files".to_string(), + input: json!({}), + }, + }, + ConversationMessage { + role: MessageRole::User, + content: MessageContent::ToolResult { + tool_use_id: format!("read-{index}"), + content: "contents".to_string(), + is_error: false, + }, + }, + ]); + } + run.replace_tool_result_archive(archive); + assert!(!serde_json::to_string(run.tool_result_archive()) + .unwrap() + .contains("task_list_snapshot")); + assert_eq!(run.task_list(), Some(&snapshot)); + assert_serialization_round_trip(&run); + run.validate_restored_state().unwrap(); +} + +#[test] +fn task_list_setter_and_restored_snapshots_validate_before_mutation() { + let mut run = run(); + let initial = task_snapshot(TaskListStatus::InProgress); + run.set_task_list(Some(initial.clone())).unwrap(); + assert!(matches!( + run.set_task_list(Some(TaskListSnapshot { tasks: Vec::new() })), + Err(ProviderRunProtocolError::InvalidTaskList { .. }), + )); + assert_eq!(run.task_list(), Some(&initial)); + + let invalid = mutate_run_json(&run, |value| value["task_list"] = json!({"tasks": []})); + assert!(restored_state_error(&invalid).contains("invalid task list")); + + let legacy = mutate_run_json(&run, |value| { + value.as_object_mut().unwrap().remove("task_list"); + }); + assert_eq!(legacy.task_list(), None); + legacy.validate_restored_state().unwrap(); + run.set_task_list(None).unwrap(); + assert_eq!(run.task_list(), None); +} + +#[test] +fn restoring_a_completed_inline_task_batch_commits_once_without_external_work() { + let mut run = run(); + let snapshot = task_snapshot(TaskListStatus::InProgress); + accept_tool_turn( + &mut run, + tool_turn( + vec![update_task_call("tasks", &snapshot)], + &[UPDATE_TASK_LIST_NAME], + ), + ); + let mut restored: ProviderRun = serde_json::from_str(&serde_json::to_string(&run).unwrap()).unwrap(); + restored.validate_restored_state().unwrap(); + assert_eq!(restored.task_list(), None); + assert!(restored.normalize_after_restore().unwrap().committed_tool_batch); + assert_eq!(restored.task_list(), Some(&snapshot)); + assert!(!restored.normalize_after_restore().unwrap().committed_tool_batch); + assert_eq!(restored.task_list(), Some(&snapshot)); + + let corrupted = mutate_run_json(&run, |value| { + value["state"]["AwaitingTools"]["batch"]["calls"][0]["state"]["Resolved"]["result"]["content"] = + json!(task_snapshot(TaskListStatus::Completed).to_tool_result_content()); + }); + assert!(restored_state_error(&corrupted).contains("does not match its arguments")); +} + #[test] fn parallel_tool_results_commit_atomically_in_original_call_order() { let mut run = run(); diff --git a/crates/galaxy_agent_core/src/task_list.rs b/crates/galaxy_agent_core/src/task_list.rs new file mode 100644 index 00000000..1f86acb8 --- /dev/null +++ b/crates/galaxy_agent_core/src/task_list.rs @@ -0,0 +1,129 @@ +use std::collections::HashSet; + +use serde::{Deserialize, Serialize}; +use serde_json::Value as JsonValue; + +pub const UPDATE_TASK_LIST_NAME: &str = "update_task_list"; +pub const MAX_TASK_LIST_ITEMS: usize = 50; +pub const MAX_TASK_ID_CHARS: usize = 200; +pub const MAX_TASK_TITLE_CHARS: usize = 1_000; + +/// A complete replacement of the task list. Callers retain task IDs across updates. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct TaskListSnapshot { + pub tasks: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct TaskListItem { + pub id: String, + pub title: String, + pub status: TaskListStatus, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum TaskListStatus { + Pending, + InProgress, + Completed, +} + +impl TaskListSnapshot { + pub fn from_arguments(arguments: &JsonValue) -> Result { + let mut arguments = arguments.clone(); + if let Some(fields) = arguments.as_object_mut() { + if let Some(title) = fields.remove("_galaxy_tool_title") { + if !title.is_string() { + return Err("The reserved tool title must be a string.".to_string()); + } + } + } + let snapshot: Self = serde_json::from_value(arguments) + .map_err(|error| format!("Invalid task list arguments: {error}"))?; + snapshot.validate()?; + Ok(snapshot) + } + + pub fn validate(&self) -> Result<(), String> { + if self.tasks.is_empty() || self.tasks.len() > MAX_TASK_LIST_ITEMS { + return Err(format!( + "Task lists must contain between 1 and {MAX_TASK_LIST_ITEMS} tasks." + )); + } + + let mut ids = HashSet::new(); + let mut has_in_progress = false; + let mut has_pending = false; + for task in &self.tasks { + if task.id.trim().is_empty() || task.id.chars().count() > MAX_TASK_ID_CHARS { + return Err(format!( + "Task IDs must be nonempty and at most {MAX_TASK_ID_CHARS} characters." + )); + } + if !ids.insert(task.id.as_str()) { + return Err(format!("Duplicate task ID: '{}'.", task.id)); + } + if task.title.trim().is_empty() || task.title.chars().count() > MAX_TASK_TITLE_CHARS + { + return Err(format!( + "Task titles must be nonempty and at most {MAX_TASK_TITLE_CHARS} characters." + )); + } + match task.status { + TaskListStatus::Pending => has_pending = true, + TaskListStatus::InProgress => { + if has_in_progress { + return Err("Only one task may be in progress at a time.".to_string()); + } + if has_pending { + return Err( + "The task in progress must precede all pending tasks; reorder the task list to reflect the current work." + .to_string(), + ); + } + has_in_progress = true; + } + TaskListStatus::Completed => {} + } + } + Ok(()) + } + + pub fn to_tool_result_content(&self) -> String { + serde_json::to_string(&TaskListToolResult { + kind: TaskListResultKind::Snapshot, + tasks: self.tasks.clone(), + }) + .expect("task list snapshots contain only serializable strings and statuses") + } +} + +/// Recognizes a canonical snapshot result. The caller must first verify the owning +/// tool name and successful result status when reading conversation history. +pub fn task_list_from_tool_result(content: &str) -> Option { + let result: TaskListToolResult = serde_json::from_str(content).ok()?; + let snapshot = TaskListSnapshot { tasks: result.tasks }; + snapshot.validate().ok()?; + Some(snapshot) +} + +#[derive(Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +struct TaskListToolResult { + #[serde(rename = "type")] + kind: TaskListResultKind, + tasks: Vec, +} + +#[derive(Serialize, Deserialize)] +enum TaskListResultKind { + #[serde(rename = "task_list_snapshot")] + Snapshot, +} + +#[cfg(test)] +#[path = "task_list_tests.rs"] +mod tests; diff --git a/crates/galaxy_agent_core/src/task_list_tests.rs b/crates/galaxy_agent_core/src/task_list_tests.rs new file mode 100644 index 00000000..da65b259 --- /dev/null +++ b/crates/galaxy_agent_core/src/task_list_tests.rs @@ -0,0 +1,157 @@ +use serde_json::json; + +use super::*; + +fn task(id: &str, status: TaskListStatus) -> TaskListItem { + TaskListItem { + id: id.to_string(), + title: format!("Task {id}"), + status, + } +} + +#[test] +fn snapshot_round_trips_through_a_recognizable_canonical_result() { + let arguments = json!({"tasks": [ + {"id": "inspect", "title": "Inspect the failure", "status": "completed"}, + {"id": "fix", "title": "Fix the failure", "status": "in_progress"}, + {"id": "review", "title": "Review the fix", "status": "pending"}, + ]}); + let snapshot = TaskListSnapshot::from_arguments(&arguments).unwrap(); + let content = snapshot.to_tool_result_content(); + + assert_eq!(task_list_from_tool_result(&content), Some(snapshot.clone())); + assert_eq!(snapshot.to_tool_result_content(), content); + assert_eq!( + serde_json::from_str::(&content).unwrap(), + json!({"type": "task_list_snapshot", "tasks": arguments["tasks"]}), + ); + assert_eq!(snapshot.tasks[1].id, "fix"); + assert_eq!(snapshot.tasks[1].status, TaskListStatus::InProgress); +} + +#[test] +fn snapshot_rejects_invented_statuses_partial_arguments_and_unknown_fields() { + for arguments in [ + json!({"tasks": []}), + json!({"tasks": [{"id": "one", "title": "Inspect", "status": "started"}]}), + json!({"tasks": [{"id": "one", "title": "Inspect"}]}), + json!({"tasks": [{"id": "one", "title": "Inspect", "status": "pending", "done": false}]}), + json!({"tasks": [{"id": "one", "title": "Inspect", "status": "pending"}], "append": true}), + json!([{"id": "one", "title": "Inspect", "status": "pending"}]), + ] { + assert!(TaskListSnapshot::from_arguments(&arguments).is_err()); + } +} + +#[test] +fn argument_parser_accepts_only_the_reserved_tool_title_outside_the_snapshot() { + let tasks = json!([{"id": "one", "title": "Inspect", "status": "in_progress"}]); + let plain = TaskListSnapshot::from_arguments(&json!({"tasks": tasks})).unwrap(); + assert_eq!( + TaskListSnapshot::from_arguments(&json!({ + "tasks": tasks, + "_galaxy_tool_title": "Update progress", + })), + Ok(plain), + ); + assert!(TaskListSnapshot::from_arguments(&json!({ + "tasks": tasks, + "_galaxy_tool_title": false, + })).is_err()); + assert!(TaskListSnapshot::from_arguments(&json!({ + "tasks": tasks, + "_galaxy_tool_title": "Update progress", + "other": "unsupported", + })).is_err()); + assert!(serde_json::from_value::(json!({ + "tasks": tasks, + "_galaxy_tool_title": "Update progress", + })).is_err()); +} + +#[test] +fn snapshot_enforces_stable_unique_nonempty_ids_and_nonempty_titles() { + for tasks in [ + vec![task(" ", TaskListStatus::Pending)], + vec![task("", TaskListStatus::Pending)], + vec![ + task("one", TaskListStatus::Completed), + task("one", TaskListStatus::Pending), + ], + vec![TaskListItem { + id: "one".to_string(), + title: "\t\n".to_string(), + status: TaskListStatus::Pending, + }], + ] { + assert!(TaskListSnapshot { tasks }.validate().is_err()); + } +} + +#[test] +fn snapshot_bounds_count_and_unicode_character_lengths() { + let mut snapshot = TaskListSnapshot { + tasks: (0..MAX_TASK_LIST_ITEMS) + .map(|index| task(&index.to_string(), TaskListStatus::Pending)) + .collect(), + }; + assert!(snapshot.validate().is_ok()); + snapshot.tasks.push(task("extra", TaskListStatus::Pending)); + assert!(snapshot.validate().is_err()); + + let mut snapshot = TaskListSnapshot { + tasks: vec![TaskListItem { + id: "🦀".repeat(MAX_TASK_ID_CHARS), + title: "界".repeat(MAX_TASK_TITLE_CHARS), + status: TaskListStatus::Pending, + }], + }; + assert!(snapshot.validate().is_ok()); + snapshot.tasks[0].id.push('🦀'); + assert!(snapshot.validate().is_err()); + snapshot.tasks[0].id.pop(); + snapshot.tasks[0].title.push('界'); + assert!(snapshot.validate().is_err()); +} + +#[test] +fn current_work_precedes_pending_tasks_while_completed_tasks_may_stay_in_place() { + let snapshot = TaskListSnapshot { + tasks: vec![ + task("done", TaskListStatus::Completed), + task("working", TaskListStatus::InProgress), + task("also-done", TaskListStatus::Completed), + task("next", TaskListStatus::Pending), + ], + }; + assert!(snapshot.validate().is_ok()); + + for tasks in [ + vec![ + task("later", TaskListStatus::Pending), + task("done", TaskListStatus::Completed), + task("working", TaskListStatus::InProgress), + ], + vec![ + task("one", TaskListStatus::InProgress), + task("two", TaskListStatus::InProgress), + ], + ] { + assert!(TaskListSnapshot { tasks }.validate().is_err()); + } +} + +#[test] +fn result_parser_requires_marker_and_validates_snapshot_instead_of_trusting_content() { + let tasks = json!([{"id": "one", "title": "Inspect", "status": "pending"}]); + for content in [ + json!({"tasks": tasks}).to_string(), + json!({"type": "something_else", "tasks": tasks}).to_string(), + json!({"type": "task_list_snapshot", "tasks": []}).to_string(), + json!({"type": "task_list_snapshot", "tasks": tasks, "unrelated": true}).to_string(), + "not JSON".to_string(), + ] { + assert_eq!(task_list_from_tool_result(&content), None); + } +} diff --git a/crates/galaxy_agent_core/src/tool_policy.rs b/crates/galaxy_agent_core/src/tool_policy.rs index 61b520ad..acf9eec9 100644 --- a/crates/galaxy_agent_core/src/tool_policy.rs +++ b/crates/galaxy_agent_core/src/tool_policy.rs @@ -3,8 +3,8 @@ use std::collections::{BTreeSet, HashMap, VecDeque}; use serde_json::Value as JsonValue; use crate::{ - ContentPart, ConversationMessage, MessageContent, ToolCall, ToolDefinition, ToolResult, - ToolResultStatus, + ContentPart, ConversationMessage, MessageContent, TaskListSnapshot, ToolCall, ToolDefinition, + ToolResult, ToolResultStatus, UPDATE_TASK_LIST_NAME, }; pub const RECALL_TOOL_HISTORY_NAME: &str = "recall_tool_history"; @@ -145,6 +145,21 @@ impl ToolPolicy { }); } + if call.name == UPDATE_TASK_LIST_NAME { + return match TaskListSnapshot::from_arguments(&call.arguments) { + Ok(snapshot) => ToolCallDecision::Inline(ToolResult { + call_id: call.id.clone(), + content: snapshot.to_tool_result_content(), + status: ToolResultStatus::Success, + }), + Err(error) => ToolCallDecision::Reject(ToolResult { + call_id: call.id.clone(), + content: format!("Error: {error}"), + status: ToolResultStatus::Error, + }), + }; + } + ToolCallDecision::Execute } } diff --git a/crates/galaxy_agent_core/src/tool_policy_tests.rs b/crates/galaxy_agent_core/src/tool_policy_tests.rs index 4b4dd597..5e876110 100644 --- a/crates/galaxy_agent_core/src/tool_policy_tests.rs +++ b/crates/galaxy_agent_core/src/tool_policy_tests.rs @@ -58,6 +58,62 @@ fn executes_only_tools_advertised_for_this_turn() { ); } +#[test] +fn task_updates_are_inline_only_when_advertised_and_keep_the_call_identity() { + let call = ToolCall { + id: "task-update-1".to_string(), + name: UPDATE_TASK_LIST_NAME.to_string(), + arguments: serde_json::json!({"tasks": [ + {"id": "fix", "title": "Fix the failure", "status": "in_progress"}, + {"id": "review", "title": "Review the fix", "status": "pending"}, + ]}), + }; + let unadvertised = ToolPolicy::new(&[definition("read_files")]); + let ToolCallDecision::Reject(result) = unadvertised.decide(&call, &[], &[]) else { + panic!("task updates cannot bypass advertised capabilities"); + }; + assert_eq!(result.call_id, call.id); + assert_eq!(result.status, ToolResultStatus::Error); + + let policy = ToolPolicy::new(&[definition(UPDATE_TASK_LIST_NAME)]); + let ToolCallDecision::Inline(result) = policy.decide(&call, &[], &[]) else { + panic!("valid task updates must not dispatch external actions"); + }; + assert_eq!(result.call_id, call.id); + assert_eq!(result.status, ToolResultStatus::Success); + let snapshot = crate::task_list_from_tool_result(&result.content).unwrap(); + assert_eq!(snapshot.tasks[0].id, "fix"); + assert_eq!(snapshot.tasks[1].id, "review"); +} + +#[test] +fn malformed_task_updates_return_one_correlated_error_instead_of_executing() { + let policy = ToolPolicy::new(&[definition(UPDATE_TASK_LIST_NAME)]); + for arguments in [ + serde_json::json!({"tasks": []}), + serde_json::json!({"tasks": [ + {"id": "same", "title": "First", "status": "completed"}, + {"id": "same", "title": "Second", "status": "pending"}, + ]}), + serde_json::json!({"tasks": [ + {"id": "one", "title": "First", "status": "invented"}, + ]}), + ] { + let call = ToolCall { + id: "invalid-task-update".to_string(), + name: UPDATE_TASK_LIST_NAME.to_string(), + arguments, + }; + let ToolCallDecision::Reject(result) = policy.decide(&call, &[], &[]) else { + panic!("invalid task updates must not execute or succeed"); + }; + assert_eq!(result.call_id, call.id); + assert_eq!(result.status, ToolResultStatus::Error); + assert!(result.content.starts_with("Error: ")); + assert_eq!(crate::task_list_from_tool_result(&result.content), None); + } +} + #[test] fn recall_searches_archived_and_live_results_with_live_results_most_recent() { let policy = ToolPolicy::new(&[definition(RECALL_TOOL_HISTORY_NAME)]); From 7948555bb93194b03254474d491df93f07618757 Mon Sep 17 00:00:00 2001 From: Ryan Ward Date: Wed, 21 Oct 2026 02:06:19 +0000 Subject: [PATCH 09/12] Keep AI session history local and update Galaxy branding --- AGENTS.md | 2 + app/src/ai/agent/mod.rs | 6 +- app/src/ai/agent/mod_tests.rs | 2 +- app/src/ai/agent_conversations_model.rs | 104 +++++++++ app/src/ai/agent_conversations_model/entry.rs | 20 ++ app/src/ai/agent_conversations_model_tests.rs | 218 +++++++++++++++++- app/src/ai/agent_management/view.rs | 177 +++++--------- .../harness/claude_code/parent_bridge.rs | 2 +- app/src/ai/agent_sdk/driver/terminal.rs | 10 +- app/src/ai/agent_tips.rs | 2 +- .../agent_view/agent_input_footer/mod.rs | 26 +-- .../agent_input_footer/toolbar_item.rs | 2 +- .../ai/blocklist/block/view_impl/common.rs | 6 +- .../history_model/conversation_loader.rs | 16 ++ .../inline_action/requested_command.rs | 2 +- app/src/ai/conversation_details_panel.rs | 24 ++ app/src/ai/custom_model_router_editor.rs | 4 +- app/src/ai/document/ai_document_model.rs | 4 +- .../execution_profiles/editor/ui_helpers.rs | 2 +- app/src/code/footer.rs | 2 +- app/src/drive/workflows/modal.rs | 2 +- .../pane_group/pane/local_harness_launch.rs | 8 +- app/src/prompt/editor_modal.rs | 2 +- app/src/search/command_palette/view.rs | 1 + app/src/search/command_search/warp_ai.rs | 7 +- .../static_commands/commands.rs | 6 +- app/src/settings_view/appearance_page.rs | 12 +- app/src/settings_view/features_page.rs | 16 +- app/src/settings_view/settings_file_footer.rs | 6 +- .../plugin_manager/codex.rs | 18 +- .../plugin_manager/codex_tests.rs | 2 +- app/src/terminal/input/models/data_source.rs | 2 +- app/src/terminal/ssh/wormhole.rs | 2 +- app/src/terminal/view.rs | 14 +- app/src/user_config/mod.rs | 2 +- app/src/workspace/action.rs | 2 + app/src/workspace/cli_install.rs | 18 +- app/src/workspace/mod.rs | 6 +- app/src/workspace/view.rs | 115 ++++++--- .../workspace/view/conversation_list/view.rs | 10 +- .../view/conversation_list/view_model.rs | 6 +- app/src/workspace/view/left_panel.rs | 9 +- app/src/workspace/view/vertical_tabs.rs | 2 +- app/src/workspace/view_tests.rs | 8 +- crates/galaxy_cli/src/agent.rs | 2 +- 45 files changed, 638 insertions(+), 271 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index fe68bd59..e043e826 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -214,6 +214,8 @@ Key invariants: - Direct-provider normal and plan turns must advertise `read_plan`, `create_plan`, and `edit_plan` when the matching document capabilities are enabled; plan-creation requests should call `create_plan` after research rather than only returning prose - Direct-provider normal and plan turns expose `update_task_list` for explicit progress updates. Keep the validated current list in durable `ProviderRun` state, inject it into each request, preserve it across compaction and restore, and reconcile restored UI updates idempotently. Finishing a response must never mark unfinished work complete. CLI monitors and completed-command assessments retain their focused tool sets. - Completed local plans are archived intact into `Plans/Archive` after seven days only when their conversation has explicitly completed all checklist items and has no active provider run or child. Draft, unfinished, or reopened plans stay visible; resumed work restores an automatically archived plan, and manual recovery is respected. Archival never deletes plan contents. +- Session history lists conversations available in local panes or local persistence, with local status/date/search filters. Restore from memory or SQLite only; never fall back to hosted attachment or server recovery, and report unavailable local data without substituting a new conversation. Ignore stale cloud/team/source/harness/artifact filters on this screen. +- Active UI and built-in agent instructions use Galaxy branding. Preserve legacy serialized names, protocol identifiers, environment variables, actual third-party package names, and upstream service references where compatibility requires them; do not relabel an unsupported hosted service as a Galaxy feature. - Unknown or invalid tool calls receive one correlated synthetic error result and a visible `AgentOutput` message; the durable run owns any continuation - `recall_tool_history` is an inline completed tool batch. `ProviderRun` commits its synthetic result and starts a bounded next turn without routing it through client action execution - `recall_tool_history` excludes earlier calls to itself; archived tool results remain searchable by query or exact `tool_use_id` diff --git a/app/src/ai/agent/mod.rs b/app/src/ai/agent/mod.rs index 48b67130..bf8a3491 100644 --- a/app/src/ai/agent/mod.rs +++ b/app/src/ai/agent/mod.rs @@ -754,7 +754,7 @@ pub enum RenderableAIError { impl RenderableAIError { const TRANSIENT_NETWORK_ERROR_MESSAGE: &'static str = - "Warp lost connection while receiving the agent response. This is usually temporary."; + "Galaxy lost connection while receiving the agent response. This is usually temporary."; /// User-facing message shown when an agent-issued command exits the shell. pub const AGENT_EXITED_SHELL_MESSAGE: &'static str = "The shell exited while the agent was running a command, so the run could not continue. Ensure the agent is not asked to run commands or source scripts that can exit the shell."; @@ -905,9 +905,9 @@ impl Display for RenderableAIError { } } Self::ServerOverloaded => { - write!(f, "Warp is currently overloaded. Please try again later.") + write!(f, "The AI service is currently overloaded. Please try again later.") } - Self::InternalWarpError => write!(f, "Internal Warp error."), + Self::InternalWarpError => write!(f, "Internal Galaxy error."), Self::ContextWindowExceeded(message) => { write!(f, "Context window exceeded: {message}") } diff --git a/app/src/ai/agent/mod_tests.rs b/app/src/ai/agent/mod_tests.rs index 44e48206..686195c5 100644 --- a/app/src/ai/agent/mod_tests.rs +++ b/app/src/ai/agent/mod_tests.rs @@ -150,7 +150,7 @@ fn transient_network_error_includes_user_facing_message_and_debug_details() { let rendered = error.to_string(); assert!( rendered.starts_with( - "Warp lost connection while receiving the agent response. This is usually temporary.\n\nDebug info: " + "Galaxy lost connection while receiving the agent response. This is usually temporary.\n\nDebug info: " ), "unexpected rendering: {rendered}" ); diff --git a/app/src/ai/agent_conversations_model.rs b/app/src/ai/agent_conversations_model.rs index ee124116..15e614f4 100644 --- a/app/src/ai/agent_conversations_model.rs +++ b/app/src/ai/agent_conversations_model.rs @@ -284,6 +284,16 @@ pub struct AgentManagementFilters { } impl AgentManagementFilters { + /// Retain only filters supported by the local session history screen. + /// Older app state may still contain cloud-only filter selections. + pub fn for_local_history(&self) -> Self { + Self { + status: self.status, + created_on: self.created_on, + ..Default::default() + } + } + pub fn reset_all_but_owner(&mut self) { self.status = StatusFilter::default(); self.source = SourceFilter::default(); @@ -1204,6 +1214,98 @@ impl AgentConversationsModel { } } + /// Returns locally recoverable sessions without hosted run or server-only rows. + pub fn get_local_entries( + &self, + filters: &AgentManagementFilters, + app: &AppContext, + ) -> Vec { + let history_model = BlocklistAIHistoryModel::as_ref(app); + let filters = filters.for_local_history(); + let mut entries = HashMap::new(); + + for metadata in self.conversations.values() { + let entry = entry::entry_for_conversation(metadata, history_model, app); + // Metadata from a live local session can precede its first disk write. + if !entry.backing.has_ambient_run + && (entry.backing.has_loaded_conversation + || entry.backing.has_local_persisted_data + || !entry.backing.has_cloud_data) + { + entries.insert(entry.id, entry.for_local_history()); + } + } + + for metadata in history_model.get_local_conversations_metadata() { + if !metadata.has_local_data && history_model.conversation(&metadata.id).is_none() { + continue; + } + let id = AgentConversationEntryId::Conversation(metadata.id); + entries.entry(id).or_insert_with(|| { + let nav_data = + ConversationNavigationData::from_historical_conversation_metadata(metadata); + entry::entry_for_historical_metadata(metadata, nav_data, history_model, app) + .for_local_history() + }); + } + + entries + .into_values() + .filter(|entry| entry.matches_filters(&filters, app)) + .sorted_by(|a, b| b.display.last_updated.cmp(&a.display.last_updated)) + .collect() + } + + pub fn get_local_entry_by_id( + &self, + id: &AgentConversationEntryId, + app: &AppContext, + ) -> Option { + let AgentConversationEntryId::Conversation(conversation_id) = id else { + return None; + }; + let entry = self.get_entry_by_id(id, app)?; + (!entry.backing.has_ambient_run + && (entry.backing.has_loaded_conversation + || entry.backing.has_local_persisted_data + || (!entry.backing.has_cloud_data + && self.conversations.contains_key(conversation_id)))) + .then(|| entry.for_local_history()) + } + + /// Resolve history navigation exclusively through an existing local pane or disk data. + pub fn resolve_local_open_action( + id: AgentConversationEntryId, + restore_layout: Option, + app: &AppContext, + ) -> Option { + let model = Self::as_ref(app); + let entry = model.get_local_entry_by_id(&id, app)?; + let conversation_id = entry.identity.local_conversation_id?; + let active_views = ActiveAgentViewsModel::as_ref(app); + let nav_data = model + .conversations + .get(&conversation_id) + .map(|metadata| &metadata.nav_data); + + if active_views.is_conversation_open(conversation_id, app) && nav_data.is_none() { + return active_views + .get_terminal_view_id_for_conversation(conversation_id, app) + .map(|terminal_view_id| WorkspaceAction::FocusTerminalViewInWorkspace { + terminal_view_id, + }); + } + + Some(WorkspaceAction::RestoreOrNavigateToConversation { + conversation_id, + window_id: nav_data.and_then(|data| data.window_id), + pane_view_locator: nav_data.and_then(|data| data.pane_view_locator), + terminal_view_id: nav_data.and_then(|data| data.terminal_view_id), + restore_layout, + local_only: true, + }) + } + /// Returns normalized, owned entries for agent management/navigation surfaces. pub fn get_entries( &self, @@ -1380,6 +1482,7 @@ impl AgentConversationsModel { pane_view_locator: nav_data.pane_view_locator, terminal_view_id: nav_data.terminal_view_id, restore_layout, + local_only: false, }); } @@ -1409,6 +1512,7 @@ impl AgentConversationsModel { pane_view_locator: None, terminal_view_id: nav_data.and_then(|nav_data| nav_data.terminal_view_id), restore_layout, + local_only: false, }); } } diff --git a/app/src/ai/agent_conversations_model/entry.rs b/app/src/ai/agent_conversations_model/entry.rs index 89d69a19..c4db9c81 100644 --- a/app/src/ai/agent_conversations_model/entry.rs +++ b/app/src/ai/agent_conversations_model/entry.rs @@ -166,6 +166,26 @@ pub struct AgentConversationCapabilities { } impl AgentConversationEntry { + /// Local history never offers hosted attachment, sharing, or server recovery. + pub(super) fn for_local_history(mut self) -> Self { + self.provenance = AgentConversationProvenance::LocalInteractive; + self.identity.ambient_agent_task_id = None; + self.identity.server_conversation_token = None; + self.identity.session_id = None; + self.display.session_status = None; + self.display.source = None; + self.display.environment_id = None; + self.display.harness = None; + self.display.request_usage = None; + self.display.executor = None; + self.capabilities.can_open = true; + self.capabilities.can_copy_link = false; + self.capabilities.can_share = false; + self.capabilities.can_fork_locally = false; + self.capabilities.can_cancel = false; + self + } + pub(super) fn matches_filters( &self, filters: &AgentManagementFilters, diff --git a/app/src/ai/agent_conversations_model_tests.rs b/app/src/ai/agent_conversations_model_tests.rs index 0bf62cc1..67de2391 100644 --- a/app/src/ai/agent_conversations_model_tests.rs +++ b/app/src/ai/agent_conversations_model_tests.rs @@ -16,7 +16,8 @@ use super::entry::{ use super::{ record_earliest_rtc_task_refresh_timestamp, AgentConversationsModel, AgentConversationsModelEvent, AgentManagementFilters, AgentRunDisplayStatus, ArtifactFilter, - ConversationMetadata, ConversationUpdateKind, EnvironmentFilter, HarnessFilter, OwnerFilter, + ConversationMetadata, ConversationUpdateKind, CreatedOnFilter, CreatorFilter, + EnvironmentFilter, HarnessFilter, OwnerFilter, SourceFilter, RtcTaskRefreshThrottleState, StatusFilter, TaskFetchError, TaskFetchState, MAX_PERSONAL_TASKS, MAX_TEAM_TASKS, }; @@ -28,11 +29,12 @@ use crate::ai::agent::conversation::{ }; use crate::ai::ambient_agents::task::{HarnessConfig, TaskPrincipalInfo, TaskStatusMessage}; use crate::ai::ambient_agents::{ - AgentConfigSnapshot, AmbientAgentTask, AmbientAgentTaskId, AmbientAgentTaskState, + AgentConfigSnapshot, AgentSource, AmbientAgentTask, AmbientAgentTaskId, AmbientAgentTaskState, }; use crate::ai::artifacts::Artifact; use crate::ai::blocklist::history_model::{ - BlocklistAIHistoryEvent, BlocklistAIHistoryModel, ConversationStatusUpdate, + BlocklistAIHistoryEvent, BlocklistAIHistoryModel, CloudConversationData, + ConversationStatusUpdate, }; use crate::ai::conversation_navigation::ConversationNavigationData; use crate::auth::AuthStateProvider; @@ -814,6 +816,216 @@ fn create_server_conversation_metadata( } } +#[test] +fn local_history_discards_obsolete_filters_but_preserves_status_and_date() { + let filters = AgentManagementFilters { + owners: OwnerFilter::All, + status: StatusFilter::Failed, + source: SourceFilter::Specific(AgentSource::CloudMode), + created_on: CreatedOnFilter::LastWeek, + creator: CreatorFilter::Specific { + name: "Former team".to_string(), + uid: "former-user".to_string(), + }, + artifact: ArtifactFilter::PullRequest, + environment: EnvironmentFilter::Specific("old-environment".to_string()), + harness: HarnessFilter::Specific(Harness::Oz), + }; + assert_eq!( + filters.for_local_history(), + AgentManagementFilters { + status: StatusFilter::Failed, + created_on: CreatedOnFilter::LastWeek, + ..Default::default() + } + ); +} + +#[test] +fn local_history_keeps_saved_legacy_sessions_and_excludes_hosted_rows() { + App::test((), |mut app| async move { + app.add_singleton_model(|_| AuthStateProvider::new_for_test()); + app.add_singleton_model(|_| ActiveAgentViewsModel::new()); + let conversation_id = AIConversationId::new(); + let task_id = "saved-root-task"; + let saved = persistence::model::AgentConversation { + conversation: persistence::model::AgentConversationRecord { + id: 0, + conversation_id: conversation_id.to_string(), + conversation_data: serde_json::to_string(&AgentConversationData { + server_conversation_token: Some("legacy-saved-token".to_string()), + ..Default::default() + }) + .unwrap(), + last_modified_at: Utc::now().naive_utc(), + }, + tasks: vec![create_api_task( + task_id, + vec![create_message("saved-message", task_id)], + )], + }; + let history = + app.add_singleton_model(|_| BlocklistAIHistoryModel::new(vec![], &[saved])); + history.update(&mut app, |model, _| { + model.merge_cloud_conversation_metadata(vec![create_server_conversation_metadata( + "Remote-only conversation", + "remote-only-token", + None, + )]); + }); + app.add_singleton_model(|_| { + let mut model = create_test_model(); + let task = create_test_task(&make_uuid(8400), "user-a", Utc::now()); + model.tasks.insert(task.task_id, task); + model + }); + app.update(|ctx| { + let model = AgentConversationsModel::as_ref(ctx); + let entries = model.get_local_entries(&AgentManagementFilters::default(), ctx); + assert_eq!(entries.len(), 1); + let entry = &entries[0]; + assert_eq!(entry.identity.local_conversation_id, Some(conversation_id)); + assert!(entry.identity.server_conversation_token.is_none()); + assert!(entry.display.harness.is_none()); + assert!(!entry.capabilities.can_share); + assert!(!entry.capabilities.can_copy_link); + assert!(model.get_local_entry_by_id(&entry.id, ctx).is_some()); + assert!(matches!( + AgentConversationsModel::resolve_local_open_action(entry.id, None, ctx), + Some(WorkspaceAction::RestoreOrNavigateToConversation { + conversation_id: restored_id, + local_only: true, + .. + }) if restored_id == conversation_id + )); + let remote_id = history + .as_ref(ctx) + .find_conversation_id_by_server_token(&ServerConversationToken::new( + "remote-only-token".to_string(), + )) + .unwrap(); + assert!(AgentConversationsModel::resolve_local_open_action( + AgentConversationEntryId::Conversation(remote_id), None, ctx, + ).is_none()); + }); + }); +} + +#[test] +fn local_history_loader_does_not_fetch_missing_cloud_data() { + App::test((), |mut app| async move { + add_entry_projection_test_models(&mut app); + let history = BlocklistAIHistoryModel::handle(&app); + history.update(&mut app, |model, _| { + model.merge_cloud_conversation_metadata(vec![create_server_conversation_metadata( + "Remote-only conversation", + "remote-only-token", + None, + )]); + }); + let future = app.update(|ctx| { + let model = history.as_ref(ctx); + let id = model + .find_conversation_id_by_server_token(&ServerConversationToken::new( + "remote-only-token".to_string(), + )) + .unwrap(); + model.load_local_conversation_data(id) + }); + // No ServerApiProvider is installed: any accidental server fallback would fail. + assert!(future.await.is_none()); + }); +} + +#[test] +fn local_history_recovers_current_live_state_without_hosted_task_shadowing() { + App::test((), |mut app| async move { + add_entry_projection_test_models(&mut app); + let conversation_id = AIConversationId::new(); + let run_id = make_uuid(8401); + let conversation = create_restored_conversation( + conversation_id, + "live-root-task", + AgentConversationData { + run_id: Some(run_id.clone()), + ..Default::default() + }, + ); + let history = BlocklistAIHistoryModel::handle(&app); + history.update(&mut app, |model, ctx| { + model.restore_conversations(EntityId::new(), vec![conversation], ctx); + model + .conversation_mut(&conversation_id) + .unwrap() + .set_last_event_sequence(42); + }); + let mut model = create_test_model(); + let task = create_test_task(&run_id, "user-a", Utc::now()); + model.tasks.insert(task.task_id, task); + let future = app.update(|ctx| { + let entries = model.get_local_entries(&AgentManagementFilters::default(), ctx); + assert_eq!(entries.len(), 1); + assert_eq!( + entries[0].id, + AgentConversationEntryId::Conversation(conversation_id) + ); + assert!(entries[0].identity.ambient_agent_task_id.is_none()); + assert!(model.get_local_entry_by_id(&entries[0].id, ctx).is_some()); + history + .as_ref(ctx) + .load_local_conversation_data(conversation_id) + }); + let Some(CloudConversationData::Oz(restored)) = future.await else { + panic!("live local session should restore without disk or server access"); + }; + assert_eq!(restored.id(), conversation_id); + assert_eq!(restored.last_event_sequence(), Some(42)); + }); +} + +#[test] +fn local_history_excludes_cached_ambient_sessions_from_list_and_lookup() { + App::test((), |mut app| async move { + add_entry_projection_test_models(&mut app); + let conversation_id = AIConversationId::new(); + let task_id = make_uuid(8402).parse().unwrap(); + let mut conversation = create_restored_conversation( + conversation_id, + "ambient-root-task", + AgentConversationData::default(), + ); + conversation.set_server_metadata(create_server_conversation_metadata( + "Cached hosted run", + "cached-ambient-token", + Some(task_id), + )); + let history = BlocklistAIHistoryModel::handle(&app); + history.update(&mut app, |model, ctx| { + model.restore_conversations(EntityId::new(), vec![conversation], ctx); + }); + let mut model = create_test_model(); + model.conversations.insert( + conversation_id, + create_test_conversation_metadata(conversation_id, "Cached hosted run"), + ); + let future = app.update(|ctx| { + assert!(model + .get_local_entries(&AgentManagementFilters::default(), ctx) + .is_empty()); + assert!(model + .get_local_entry_by_id( + &AgentConversationEntryId::Conversation(conversation_id), + ctx, + ) + .is_none()); + history + .as_ref(ctx) + .load_local_conversation_data(conversation_id) + }); + assert!(future.await.is_none()); + }); +} + #[test] fn test_get_entries_includes_task_only_entry() { App::test((), |mut app| async move { diff --git a/app/src/ai/agent_management/view.rs b/app/src/ai/agent_management/view.rs index fbd27d84..7a3aba9d 100644 --- a/app/src/ai/agent_management/view.rs +++ b/app/src/ai/agent_management/view.rs @@ -36,7 +36,7 @@ use siphasher::sip::SipHasher; use crate::ai::agent::conversation::AIConversationId; use crate::ai::agent_conversations_model::{ - AgentConversationEntry, AgentConversationEntryId, AgentConversationNavigationSubject, + AgentConversationEntry, AgentConversationEntryId, AgentConversationsModel, AgentConversationsModelEvent, AgentManagementFilters, ArtifactFilter, ConversationUpdateKind, CreatedOnFilter, CreatorFilter, EnvironmentFilter, HarnessFilter, OwnerFilter, SessionStatus, SourceFilter, StatusFilter, @@ -55,7 +55,6 @@ use crate::ai::agent_management::telemetry::{ }; use crate::ai::ambient_agents::{cancel_task_with_toast, AgentSource}; use crate::ai::artifacts::{Artifact, ArtifactButtonsRow, ArtifactButtonsRowEvent}; -use crate::ai::blocklist::format_credits; use crate::ai::conversation_details_panel::{ ConversationDetailsData, ConversationDetailsPanel, ConversationDetailsPanelEvent, }; @@ -89,7 +88,6 @@ use crate::workflows::WorkflowType; use crate::workspace::{ ForkedConversationDestination, RestoreConversationLayout, ToastStack, WorkspaceAction, }; -use crate::workspaces::user_workspaces::UserWorkspaces; use crate::{send_telemetry_from_ctx, AgentModeEntrypoint}; lazy_static! { @@ -311,7 +309,7 @@ impl AgentManagementView { }, ctx, ); - editor.set_placeholder_text("Search", ctx); + editor.set_placeholder_text("Search local sessions", ctx); editor }); ctx.subscribe_to_view(&search_editor, |me, _handle, event, ctx| { @@ -319,7 +317,7 @@ impl AgentManagementView { }); let new_agent_button = CompactibleActionButton::new( - "New agent".to_string(), + "New session".to_string(), None, ButtonSize::Small, AgentManagementViewAction::ShowAgentTypeSelector, @@ -333,7 +331,9 @@ impl AgentManagementView { let has_dismissed_setup_guide = *AISettings::as_ref(ctx).did_dismiss_cloud_setup_guide; - let filters = persisted_filters.map(|p| p.filters).unwrap_or_default(); + let filters = persisted_filters + .map(|persisted| persisted.filters.for_local_history()) + .unwrap_or_default(); let details_panel: ViewHandle = ctx.add_typed_action_view(|ctx| { @@ -378,28 +378,17 @@ impl AgentManagementView { view.update_creator_dropdown(ctx); view.update_environment_dropdown(ctx); - // Trigger server fetch if persisted filters differ from defaults - // (team tasks are not loaded at startup, so we need to fetch them) - if view.filters != AgentManagementFilters::default() { - view.trigger_filter_fetch(ctx); - } - view.get_tasks_from_model(ctx); view } fn get_view_state(&self, app: &AppContext) -> ViewState { let model = AgentConversationsModel::as_ref(app); - let has_items = model.has_items(); + let has_items = !model + .get_local_entries(&AgentManagementFilters::default(), app) + .is_empty(); - // If loading with zero items, show skeleton cards - // If loading with items, show list of interactive conversations (with loading indicator in header) - if model.is_loading() && !has_items { - return ViewState::Loading; - } - - // Show setup guide if: no items (zero state) or user clicked button to toggle on the guide - if !has_items || self.is_viewing_setup_guide { + if !has_items { return ViewState::SetupGuide { has_items }; } @@ -866,7 +855,7 @@ impl AgentManagementView { /// Common handler for filter changes: fetch, refresh tasks, and save. fn on_filter_changed(&mut self, ctx: &mut ViewContext) { - self.trigger_filter_fetch(ctx); + self.filters = self.filters.for_local_history(); self.get_tasks_from_model(ctx); ctx.dispatch_global_action("workspace:save_app", ()); } @@ -888,11 +877,9 @@ impl AgentManagementView { /// Shows the setup guide from a deep-link/action without toggling it off on repeated calls. pub(crate) fn show_setup_guide_from_link(&mut self, ctx: &mut ViewContext) { - if !self.is_viewing_setup_guide { - send_telemetry_from_ctx!(AgentManagementTelemetryEvent::OpenSetupGuide, ctx); - } - self.is_viewing_setup_guide = true; - ctx.notify(); + // Older links may target cloud onboarding. Keep them within local history. + self.is_viewing_setup_guide = false; + self.get_tasks_from_model(ctx); } #[cfg(test)] @@ -905,10 +892,10 @@ impl AgentManagementView { environment_id: String, ctx: &mut ViewContext, ) { - // This navigation should show the team/global task runs list. - self.filters.owners = OwnerFilter::All; + // A legacy environment link must not invisibly filter local sessions. + let _ = environment_id; + self.filters.owners = OwnerFilter::PersonalOnly; self.filters.reset_all_but_owner(); - self.filters.environment = EnvironmentFilter::Specific(environment_id); self.update_filter_buttons(ctx); // Clear search query. @@ -952,7 +939,7 @@ impl AgentManagementView { let model = AgentConversationsModel::as_ref(ctx); let search_query = self.search_query.trim().to_lowercase(); let cards: Vec = model - .get_entries(&self.filters, ctx) + .get_local_entries(&self.filters, ctx) .into_iter() .filter(|entry| { if search_query.is_empty() { @@ -962,16 +949,7 @@ impl AgentManagementView { }) .map(|entry| { let item_id = entry.id; - let copy_link_url = entry - .capabilities - .can_copy_link - .then(|| { - AgentConversationsModel::resolve_copy_link( - AgentConversationNavigationSubject::Entry(entry.id), - ctx, - ) - }) - .flatten(); + let copy_link_url = None; let mut config = Self::action_buttons_config_for_entry(&entry, None, copy_link_url); if FeatureFlag::AgentManagementDetailsView.is_enabled() { config.view_details_item_id = Some(item_id); @@ -1100,6 +1078,9 @@ impl AgentManagementView { { config.delete_conversation_id = None; } + if !entry.capabilities.can_fork_locally { + config.fork_conversation_id = None; + } config } @@ -1184,7 +1165,7 @@ impl AgentManagementView { AgentDetailsButtonEvent::DeleteConversation { conversation_id } => { let model = AgentConversationsModel::as_ref(ctx); let conversation_title = model - .get_entry_by_id(item_id, ctx) + .get_local_entry_by_id(item_id, ctx) .map(|entry| entry.display.title) .unwrap_or_else(|| "Conversation".to_string()); ctx.emit(AgentManagementViewEvent::ShowDeleteConfirmationDialog { @@ -1341,18 +1322,15 @@ impl AgentManagementView { ctx: &mut ViewContext, ) { let model = AgentConversationsModel::as_ref(ctx); - let Some(entry) = model.get_entry_by_id(item_id, ctx) else { + let Some(entry) = model.get_local_entry_by_id(item_id, ctx) else { return; }; - let open_action = AgentConversationsModel::resolve_open_action( - AgentConversationNavigationSubject::Entry(*item_id), + let open_action = AgentConversationsModel::resolve_local_open_action( + *item_id, Some(RestoreConversationLayout::NewTab), ctx, ); - let copy_link_url = AgentConversationsModel::resolve_copy_link( - AgentConversationNavigationSubject::Entry(*item_id), - ctx, - ); + let copy_link_url = None; let task = entry .identity .ambient_agent_task_id @@ -1362,7 +1340,8 @@ impl AgentManagementView { task.as_ref(), open_action, copy_link_url, - ); + ) + .for_local_history(); self.details_panel.update(ctx, |p, ctx| { p.set_conversation_details(data, ctx); @@ -1377,7 +1356,7 @@ impl AgentManagementView { ) { let model = AgentConversationsModel::as_ref(ctx); let Some((index, entry)) = self.items.iter().enumerate().find_map(|(index, card)| { - let entry = model.get_entry_by_id(&card.item_id, ctx)?; + let entry = model.get_local_entry_by_id(&card.item_id, ctx)?; (entry.identity.local_conversation_id == Some(conversation_id)) .then_some((index, entry)) }) else { @@ -1638,7 +1617,7 @@ impl AgentManagementView { }; let model = AgentConversationsModel::as_ref(app); - let Some(entry) = model.get_entry_by_id(&card_state.item_id, app) else { + let Some(entry) = model.get_local_entry_by_id(&card_state.item_id, app) else { return Empty::new().finish(); }; @@ -1673,7 +1652,7 @@ impl AgentManagementView { let mut card_content = Flex::column() .with_spacing(CARD_ROW_SPACING) .with_child(Self::render_header_row(card_state, entry, appearance)) - .with_child(Self::render_metadata_row(entry, appearance, app)); + .with_child(Self::render_metadata_row(entry, appearance)); // Add artifacts row if there is a buttons view if let Some(buttons_element) = artifact_buttons_element { @@ -1816,25 +1795,14 @@ impl AgentManagementView { fn render_metadata_row( entry: &AgentConversationEntry, appearance: &Appearance, - app: &AppContext, ) -> Box { let theme = appearance.theme(); let font_family = appearance.ui_font_family(); let font_size = appearance.ui_font_size(); let mut metadata_parts = Vec::new(); - if let Some(source) = &entry.display.source { - metadata_parts.push(format!("Source: {}", source.display_name())); - } - - let availability = HarnessAvailabilityModel::as_ref(app); - if availability.should_show_harness_selector() { - if let Some(harness) = entry.display.harness { - metadata_parts.push(format!( - "Harness: {}", - availability.display_name_for(harness) - )); - } + if let Some(directory) = &entry.display.working_directory { + metadata_parts.push(directory.clone()); } if let Some(executor) = &entry.display.executor { @@ -1859,8 +1827,8 @@ impl AgentManagementView { metadata_parts.push(format!("Run time: {run_time}")); } - if let Some(usage) = entry.display.request_usage.map(format_credits) { - metadata_parts.push(format!("Credits used: {usage}")); + if metadata_parts.is_empty() { + metadata_parts.push("Saved on this device".to_string()); } Text::new(metadata_parts.join(" • "), font_family, font_size) @@ -1870,13 +1838,7 @@ impl AgentManagementView { // Render the main page header based on the current view state fn render_header(&self, app: &AppContext) -> Box { - match self.get_view_state(app) { - ViewState::Loading => self.render_loading_header(app), - ViewState::SetupGuide { - has_items: has_tasks, - } => self.render_setup_guide_header(has_tasks, app), - ViewState::NoFilterMatches | ViewState::HasTasks => self.render_task_list_header(app), - } + self.render_task_list_header(app) } fn render_setup_guide_header(&self, has_tasks: bool, app: &AppContext) -> Box { @@ -1920,17 +1882,11 @@ impl AgentManagementView { let appearance = Appearance::as_ref(app); let theme = appearance.theme(); - let is_loading = AgentConversationsModel::handle(app) - .as_ref(app) - .is_loading(); - - let is_on_team = UserWorkspaces::as_ref(app).current_team().is_some(); - let size_switch_threshold = MEDIUM_SIZE_SWITCH_THRESHOLD * appearance.monospace_ui_scalar(); let build_header = |use_expanded: bool| { let title = Text::new_inline( - "Runs", + "AI sessions", appearance.ui_font_family(), appearance.ui_font_size() + 4., ) @@ -1938,12 +1894,6 @@ impl AgentManagementView { .with_color(theme.active_ui_text_color().into()) .finish(); - let setup_guide_button = if use_expanded { - self.setup_guide_button.expanded_button() - } else { - self.setup_guide_button.compact_button() - }; - let new_agent_button = if use_expanded { self.new_agent_button.expanded_button() } else { @@ -1956,17 +1906,7 @@ impl AgentManagementView { .with_spacing(4.) .with_child(Container::new(title).with_margin_right(12.).finish()); - if is_on_team { - header_top.add_child(ChildView::new(&self.personal_filter_button).finish()); - header_top.add_child(ChildView::new(&self.all_filter_button).finish()); - } - - if is_loading { - header_top.add_child(self.render_cloud_loading_icon(appearance)); - } - header_top.add_child(Expanded::new(1., Empty::new().finish()).finish()); - header_top.add_child(ChildView::new(setup_guide_button).finish()); if !cfg!(target_family = "wasm") { header_top.add_child(ChildView::new(new_agent_button).finish()); @@ -1976,19 +1916,7 @@ impl AgentManagementView { .with_spacing(4.) .with_cross_axis_alignment(CrossAxisAlignment::Center) .with_child(ChildView::new(&self.status_dropdown).finish()) - .with_child(ChildView::new(&self.source_dropdown).finish()) - .with_child(ChildView::new(&self.created_on_dropdown).finish()) - .with_child(ChildView::new(&self.artifact_dropdown).finish()); - - if HarnessAvailabilityModel::as_ref(app).should_show_harness_selector() { - filters_wrap.add_child(ChildView::new(&self.harness_dropdown).finish()); - } - - filters_wrap.add_child(ChildView::new(&self.environment_dropdown).finish()); - - if self.filters.owners != OwnerFilter::PersonalOnly { - filters_wrap.add_child(ChildView::new(&self.creator_dropdown).finish()); - } + .with_child(ChildView::new(&self.created_on_dropdown).finish()); if self.filters.is_filtering() { filters_wrap.add_child(ChildView::new(&self.clear_all_filters_button).finish()); @@ -2198,7 +2126,19 @@ impl View for AgentManagementView { fn render(&self, app: &AppContext) -> Box { let main_content = match self.get_view_state(app) { ViewState::Loading => self.render_loading_state(app), - ViewState::SetupGuide { .. } => ChildView::new(&self.cloud_setup_guide_view).finish(), + ViewState::SetupGuide { .. } => { + let appearance = Appearance::as_ref(app); + Align::new( + Text::new( + "No saved AI sessions on this device. Start a session to keep its history here.", + appearance.ui_font_family(), + appearance.ui_font_size(), + ) + .with_color(appearance.theme().nonactive_ui_text_color().into()) + .finish(), + ) + .finish() + } ViewState::NoFilterMatches => self.render_no_results_view(app), ViewState::HasTasks => self.render_default_scroll_view(app), }; @@ -2388,16 +2328,17 @@ impl TypedActionView for AgentManagementView { } AgentManagementViewAction::ShowAgentTypeSelector => { send_telemetry_from_ctx!( - AgentManagementTelemetryEvent::AgentTypeSelectorOpened, + AgentManagementTelemetryEvent::SpawnNewLocalAgent, ctx ); - self.is_agent_type_selector_open = true; - ctx.focus(&self.agent_type_selector); - ctx.notify(); + ctx.dispatch_typed_action(&WorkspaceAction::NewTabInAgentMode { + entrypoint: AgentModeEntrypoint::AgentManagementView, + zero_state_prompt_suggestion_type: None, + }); } AgentManagementViewAction::OpenSession { item_id } => { - let Some(action) = AgentConversationsModel::resolve_open_action( - AgentConversationNavigationSubject::Entry(*item_id), + let Some(action) = AgentConversationsModel::resolve_local_open_action( + *item_id, Some(RestoreConversationLayout::NewTab), ctx, ) else { diff --git a/app/src/ai/agent_sdk/driver/harness/claude_code/parent_bridge.rs b/app/src/ai/agent_sdk/driver/harness/claude_code/parent_bridge.rs index 9dec5b7d..5fcf49fa 100644 --- a/app/src/ai/agent_sdk/driver/harness/claude_code/parent_bridge.rs +++ b/app/src/ai/agent_sdk/driver/harness/claude_code/parent_bridge.rs @@ -41,7 +41,7 @@ const PARENT_BRIDGE_HOOK_OUTPUT_ACK_FILE_NAME: &str = "pending-hook-output.ack"; const PARENT_BRIDGE_MAX_CONTEXT_CHARS_ENV: &str = "OZ_PARENT_MAX_CONTEXT_CHARS"; const PARENT_BRIDGE_DEFAULT_MAX_CONTEXT_CHARS: usize = 6000; pub(super) const MESSAGE_BRIDGE_CONTEXT_PREAMBLE: &str = - "Oz mailbox update for this child run.\nSource: lead agent\nContext type: user-level coordination messages\n"; + "Galaxy mailbox update for this child run.\nSource: lead agent\nContext type: user-level coordination messages\n"; const PARENT_BRIDGE_REMAINING_MESSAGES_NOTE: &str = "\n\nAdditional lead agent mailbox messages remain queued and may be surfaced later."; diff --git a/app/src/ai/agent_sdk/driver/terminal.rs b/app/src/ai/agent_sdk/driver/terminal.rs index 013d8a93..b96cfbf9 100644 --- a/app/src/ai/agent_sdk/driver/terminal.rs +++ b/app/src/ai/agent_sdk/driver/terminal.rs @@ -41,9 +41,9 @@ use crate::workspaces::user_workspaces::UserWorkspaces; pub(crate) enum BootstrapError { /// The PTY or shell process failed before the bootstrap script completed. /// When `reason` is `Some`, the message is - /// "Shell spawn failed: {reason}. Check the Warp logs for details." + /// "Shell spawn failed: {reason}. Check the Galaxy logs for details." /// When `reason` is `None`, it is - /// "Shell spawn failed. Check the Warp logs for details." + /// "Shell spawn failed. Check the Galaxy logs for details." PtySpawnFailed { reason: Option }, /// The bootstrap script did not complete within the expected time. TimedOut, @@ -58,16 +58,16 @@ impl std::fmt::Display for BootstrapError { BootstrapError::PtySpawnFailed { reason: Some(r) } => { write!( f, - "Shell spawn failed: {r}. Check the Warp logs for details." + "Shell spawn failed: {r}. Check the Galaxy logs for details." ) } BootstrapError::PtySpawnFailed { reason: None } => { - write!(f, "Shell spawn failed. Check the Warp logs for details.") + write!(f, "Shell spawn failed. Check the Galaxy logs for details.") } BootstrapError::TimedOut => write!( f, "Terminal session did not start within the expected time. \ - Check the Warp logs for details." + Check the Galaxy logs for details." ), BootstrapError::InternalError => { write!(f, "An unexpected internal error occurred during bootstrap.") diff --git a/app/src/ai/agent_tips.rs b/app/src/ai/agent_tips.rs index e1976fc4..668d85cc 100644 --- a/app/src/ai/agent_tips.rs +++ b/app/src/ai/agent_tips.rs @@ -112,7 +112,7 @@ static DEFAULT_TIPS: LazyLock> = LazyLock::new(|| { kind: AgentTipKind::SlashCommands, }, AgentTip { - description: " to open the Command Palette and access Warp actions and shortcuts.".to_string(), + description: " to open the Command Palette and access Galaxy actions and shortcuts.".to_string(), link: Some("https://docs.warp.dev/terminal/command-palette".to_string()), binding_name: Some(TOGGLE_COMMAND_PALETTE_KEYBINDING_NAME), action: Some(WorkspaceAction::OpenPalette { diff --git a/app/src/ai/blocklist/agent_view/agent_input_footer/mod.rs b/app/src/ai/blocklist/agent_view/agent_input_footer/mod.rs index 8287c361..99623c16 100644 --- a/app/src/ai/blocklist/agent_view/agent_input_footer/mod.rs +++ b/app/src/ai/blocklist/agent_view/agent_input_footer/mod.rs @@ -443,7 +443,7 @@ impl AgentInputFooter { ActionButton::new("Enable notifications", InstallPluginButtonTheme) .with_icon(Icon::Download) .with_tooltip( - "Install the Warp plugin to enable rich agent notifications within Warp", + "Install the Galaxy plugin to enable rich agent notifications within Galaxy", ) .with_size(cli_button_size) .with_tooltip_alignment(TooltipAlignment::Left) @@ -456,7 +456,7 @@ impl AgentInputFooter { let plugin_instructions_button = ctx.add_typed_action_view(|_ctx| { ActionButton::new("Notifications setup instructions", InstallPluginButtonTheme) .with_icon(Icon::Info) - .with_tooltip("View instructions to install the Warp plugin") + .with_tooltip("View instructions to install the Galaxy plugin") .with_size(cli_button_size) .with_tooltip_alignment(TooltipAlignment::Left) .with_adjoined_side(AdjoinedSide::Right) @@ -468,9 +468,9 @@ impl AgentInputFooter { }); let update_plugin_button = ctx.add_typed_action_view(|_ctx| { - ActionButton::new("Update Warp plugin", InstallPluginButtonTheme) + ActionButton::new("Update Galaxy plugin", InstallPluginButtonTheme) .with_icon(Icon::Download) - .with_tooltip("A new version of the Warp plugin is available") + .with_tooltip("A new version of the Galaxy plugin is available") .with_size(cli_button_size) .with_tooltip_alignment(TooltipAlignment::Left) .with_adjoined_side(AdjoinedSide::Right) @@ -482,7 +482,7 @@ impl AgentInputFooter { let update_instructions_button = ctx.add_typed_action_view(|_ctx| { ActionButton::new("Plugin update instructions", InstallPluginButtonTheme) .with_icon(Icon::Info) - .with_tooltip("View instructions to update the Warp plugin") + .with_tooltip("View instructions to update the Galaxy plugin") .with_size(cli_button_size) .with_tooltip_alignment(TooltipAlignment::Left) .with_adjoined_side(AdjoinedSide::Right) @@ -1396,10 +1396,10 @@ impl AgentInputFooter { .cli_agent(ctx) .and_then(plugin_manager_for) .map(|m| m.install_success_message()) - .unwrap_or("Warp plugin installed. Please restart the session to activate."); + .unwrap_or("Galaxy plugin installed. Please restart the session to activate."); self.handle_plugin_operation( - "Installing Warp plugin...", - "Failed to install Warp plugin", + "Installing Galaxy plugin...", + "Failed to install Galaxy plugin", success_msg, PluginChipTelemetryKind::Install, |manager| async move { manager.install().await }, @@ -1413,10 +1413,10 @@ impl AgentInputFooter { .cli_agent(ctx) .and_then(plugin_manager_for) .map(|m| m.update_success_message()) - .unwrap_or("Warp plugin updated. Please restart the session to activate."); + .unwrap_or("Galaxy plugin updated. Please restart the session to activate."); self.handle_plugin_operation( - "Updating Warp plugin...", - "Failed to update Warp plugin", + "Updating Galaxy plugin...", + "Failed to update Galaxy plugin", success_msg, PluginChipTelemetryKind::Update, |manager| async move { manager.update().await }, @@ -2823,7 +2823,7 @@ impl ActionButtonTheme for ActiveMicButtonTheme { } } -/// Green-accented theme for the "Install Warp plugin" chip. +/// Green-accented theme for the "Install Galaxy plugin" chip. struct InstallPluginButtonTheme; impl ActionButtonTheme for InstallPluginButtonTheme { @@ -2863,7 +2863,7 @@ async fn write_install_log(agent: CLIAgent, err: &PluginInstallError) -> Option< let log_path = env::temp_dir().join("warp-plugin-install.log"); let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S UTC"); let contents = format!( - "Warp plugin installation — {agent:?}\n\ + "Galaxy plugin installation — {agent:?}\n\ {now}\n\ \n\ {log}", diff --git a/app/src/ai/blocklist/agent_view/agent_input_footer/toolbar_item.rs b/app/src/ai/blocklist/agent_view/agent_input_footer/toolbar_item.rs index 063ca0cc..4414c7a3 100644 --- a/app/src/ai/blocklist/agent_view/agent_input_footer/toolbar_item.rs +++ b/app/src/ai/blocklist/agent_view/agent_input_footer/toolbar_item.rs @@ -129,7 +129,7 @@ impl AgentToolbarItemKind { Self::NLDToggle => "Autodetection", Self::VoiceInput => "Voice Input", Self::FileAttach => "Attach File", - Self::ContextWindowUsage => "Warp Context Indicator", + Self::ContextWindowUsage => "Galaxy Context Indicator", Self::LLMContextUsage => "LLM Context Usage", Self::LLMCacheDetails => "LLM Cache Details", Self::FileExplorer => "File Explorer", diff --git a/app/src/ai/blocklist/block/view_impl/common.rs b/app/src/ai/blocklist/block/view_impl/common.rs index ec1de633..fb880dcc 100644 --- a/app/src/ai/blocklist/block/view_impl/common.rs +++ b/app/src/ai/blocklist/block/view_impl/common.rs @@ -107,7 +107,7 @@ pub const WAITING_FOR_USER_INPUT_MESSAGE: &str = "Agent waiting for instructions const IMAGE_SOURCE_LINK_LINE_INDEX: usize = 1; const ERROR_APOLOGY_TEXT: &str = "I'm sorry, I couldn't complete that request."; -const INTERNAL_WARP_ERROR: &str = "Internal Warp error."; +const INTERNAL_WARP_ERROR: &str = "Internal Galaxy error."; const GALAXY_STATUS_MESSAGES: &[&str] = &[ "Warping through the Galaxy...", @@ -3133,7 +3133,7 @@ pub fn render_failed_output(props: FailedOutputProps, app: &AppContext) -> Box Box { - "Warp is currently overloaded. Please try again later.".to_string() + "The AI service is currently overloaded. Please try again later.".to_string() } RenderableAIError::InternalWarpError => { format!("{ERROR_APOLOGY_TEXT}\n\n{INTERNAL_WARP_ERROR}") diff --git a/app/src/ai/blocklist/history_model/conversation_loader.rs b/app/src/ai/blocklist/history_model/conversation_loader.rs index eb1c8ef7..701060c5 100644 --- a/app/src/ai/blocklist/history_model/conversation_loader.rs +++ b/app/src/ai/blocklist/history_model/conversation_loader.rs @@ -226,6 +226,22 @@ impl AIConversationMetadata { } impl BlocklistAIHistoryModel { + /// Restore saved local sessions without falling back to a hosted transcript. + pub fn load_local_conversation_data( + &self, + conversation_id: AIConversationId, + ) -> galaxyui::r#async::BoxFuture<'static, Option> { + let conversation = self + .load_conversation_from_db(&conversation_id) + .filter(|conversation| { + !conversation + .server_metadata() + .is_some_and(|metadata| metadata.ambient_agent_task_id.is_some()) + }) + .map(|conversation| CloudConversationData::Oz(Box::new(conversation))); + box_future(futures::future::ready(conversation)) + } + /// Loads conversation data from the appropriate source (DB or server). /// /// This method automatically determines whether to load from the local database or diff --git a/app/src/ai/blocklist/inline_action/requested_command.rs b/app/src/ai/blocklist/inline_action/requested_command.rs index 078b600e..88899a6e 100644 --- a/app/src/ai/blocklist/inline_action/requested_command.rs +++ b/app/src/ai/blocklist/inline_action/requested_command.rs @@ -877,7 +877,7 @@ impl RequestedCommandView { ) if show_for_action_id == &self.action_id => { *shown.lock() = true; Some(render_autonomy_checkbox_setting_speedbump_footer( - "Always allow Oz to execute read-only commands (relies on model)", + "Always allow Galaxy to execute read-only commands (relies on model)", *checked, AIBlockAction::ToggleAutoexecuteReadonlyCommandsSpeedbumpCheckbox, self.autoexecute_readonly_commands_speedbump_checkbox_handle diff --git a/app/src/ai/conversation_details_panel.rs b/app/src/ai/conversation_details_panel.rs index 67920f4b..8c6362c4 100644 --- a/app/src/ai/conversation_details_panel.rs +++ b/app/src/ai/conversation_details_panel.rs @@ -246,9 +246,20 @@ pub struct ConversationDetailsData { harness: Option, /// Error details displayed when the API call to fetch run data failed. fetch_error: Option, + /// Local history resumes the saved session instead of offering a hosted fork. + local_history: bool, } impl ConversationDetailsData { + pub fn for_local_history(mut self) -> Self { + self.local_history = true; + self + } + + fn is_local_history(&self) -> bool { + self.local_history + } + fn directory_for_task(task: &AmbientAgentTask, app: &AppContext) -> Option { let history_model = BlocklistAIHistoryModel::as_ref(app); let conversation_id = history_model @@ -358,6 +369,7 @@ impl ConversationDetailsData { skill_spec: None, harness, fetch_error: None, + local_history: false, } } @@ -422,6 +434,7 @@ impl ConversationDetailsData { skill_spec, harness, fetch_error: None, + local_history: false, } } @@ -494,6 +507,7 @@ impl ConversationDetailsData { skill_spec, harness, fetch_error: None, + local_history: false, }; } @@ -521,6 +535,7 @@ impl ConversationDetailsData { skill_spec: None, harness, fetch_error: None, + local_history: false, } } @@ -552,6 +567,7 @@ impl ConversationDetailsData { skill_spec: None, harness: None, fetch_error, + local_history: false, } } @@ -594,6 +610,7 @@ impl ConversationDetailsData { skill_spec: None, harness, fetch_error: None, + local_history: false, } } } @@ -752,6 +769,10 @@ impl ConversationDetailsPanel { &self, app: &AppContext, ) -> Option { + if self.data.is_local_history() { + // History's Open action restores the original local session directly. + return None; + } if !AISettings::as_ref(app).is_any_ai_enabled(app) { return None; } @@ -894,6 +915,9 @@ impl ConversationDetailsPanel { open_action, data.copy_link_url.clone(), ); + if data.is_local_history() { + config.fork_conversation_id = None; + } if !status.as_ref().is_some_and(ConversationStatus::is_done) { config.delete_conversation_id = None; } diff --git a/app/src/ai/custom_model_router_editor.rs b/app/src/ai/custom_model_router_editor.rs index 5cea0bce..b6a7a811 100644 --- a/app/src/ai/custom_model_router_editor.rs +++ b/app/src/ai/custom_model_router_editor.rs @@ -687,7 +687,7 @@ impl CustomRouterEditorView { ); let rules_copy = FormattedText::new([ FormattedTextLine::Line(vec![FormattedTextFragment::plain_text( - "Rules are custom prompts that describe when to use a specific model. Warp intelligently matches your tasks against these rules.", + "Rules are custom prompts that describe when to use a specific model. Galaxy matches your tasks against these rules.", )]), FormattedTextLine::Line(vec![FormattedTextFragment::plain_text( "Rules are matched top to bottom — rules higher in the list take precedence over those below.", @@ -750,7 +750,7 @@ impl CustomRouterEditorView { FormattedTextLine::Line(vec![ FormattedTextFragment::bold("Complexity-based"), FormattedTextFragment::plain_text( - " routing chooses a model based on Warp's classification of the task's difficulty.", + " routing chooses a model based on Galaxy's classification of the task's difficulty.", ), ]), FormattedTextLine::Line(vec![ diff --git a/app/src/ai/document/ai_document_model.rs b/app/src/ai/document/ai_document_model.rs index f4e6fb07..1137e1b4 100644 --- a/app/src/ai/document/ai_document_model.rs +++ b/app/src/ai/document/ai_document_model.rs @@ -618,12 +618,12 @@ impl AIDocumentModel { ) }) .ok_or_else(|| { - format!("Plan document {ai_document_id} was not found in Warp Drive.") + format!("Plan document {ai_document_id} was not found in Galaxy Drive.") })?; let (sync_id, title, content) = notebook; if sync_id.into_server().is_none() { return Err(format!( - "Plan document {ai_document_id} is not backed by a saved Warp Drive notebook." + "Plan document {ai_document_id} is not backed by a saved Galaxy Drive notebook." )); } diff --git a/app/src/ai/execution_profiles/editor/ui_helpers.rs b/app/src/ai/execution_profiles/editor/ui_helpers.rs index 9401496f..310dc2a8 100644 --- a/app/src/ai/execution_profiles/editor/ui_helpers.rs +++ b/app/src/ai/execution_profiles/editor/ui_helpers.rs @@ -258,7 +258,7 @@ pub fn render_models_section( .with_child(render_filterable_dropdown_row( appearance, "Base model", - "This model serves as the primary engine behind the agent. It powers most interactions and invokes other models for tasks like planning or code generation when necessary. Warp may automatically switch to alternate models based on model availability or for auxiliary tasks such as conversation summarization.", + "This model serves as the primary engine behind the agent. It powers most interactions and invokes other models for tasks like planning or code generation when necessary. Galaxy may automatically switch to alternate models based on model availability or for auxiliary tasks such as conversation summarization.", &view.base_model_dropdown, )); diff --git a/app/src/code/footer.rs b/app/src/code/footer.rs index cc67c42d..30e69ce8 100644 --- a/app/src/code/footer.rs +++ b/app/src/code/footer.rs @@ -1729,7 +1729,7 @@ impl View for CodeFooterView { Self::render_status_text( theme, appearance, - "Use Oz to update this config".to_string(), + "Use Galaxy Agent to update this config".to_string(), ), ) .finish(), diff --git a/app/src/drive/workflows/modal.rs b/app/src/drive/workflows/modal.rs index 5cf06555..2ae1865d 100644 --- a/app/src/drive/workflows/modal.rs +++ b/app/src/drive/workflows/modal.rs @@ -1704,7 +1704,7 @@ impl WorkflowModal { .finish(); let button_with_tool_tip = appearance.ui_builder().tool_tip_on_element( - "Generate a title, descriptions, or parameters with Warp AI".to_string(), + "Generate a title, descriptions, or parameters with Galaxy Agent".to_string(), self.button_mouse_states.ai_assist_tool_tip.clone(), rendered_button, ParentAnchor::BottomMiddle, diff --git a/app/src/pane_group/pane/local_harness_launch.rs b/app/src/pane_group/pane/local_harness_launch.rs index cec992eb..b7a9f796 100644 --- a/app/src/pane_group/pane/local_harness_launch.rs +++ b/app/src/pane_group/pane/local_harness_launch.rs @@ -84,15 +84,15 @@ pub(super) fn validate_local_harness_shell(shell_type: Option) -> Res } } -const LOCAL_CLAUDE_CHILD_ORCHESTRATION_INSTRUCTIONS: &str = r#"You are a local Claude Code child agent launched by a lead agent in Warp. +const LOCAL_CLAUDE_CHILD_ORCHESTRATION_INSTRUCTIONS: &str = r#"You are a local Claude Code child agent launched by a lead agent in Galaxy. -Coordinate with the lead agent through the Oz CLI messaging environment: +Coordinate with the lead agent through Galaxy's local CLI messaging environment: - Your run id is in OZ_RUN_ID. - The lead agent id is in OZ_PARENT_RUN_ID. -- The Oz CLI command is in OZ_CLI. +- The Galaxy CLI command is in OZ_CLI (the compatibility environment variable name). If OZ_CLI, OZ_RUN_ID, or OZ_PARENT_RUN_ID is missing, report that blocker in your final response. -Do not use Claude Code Agent or SendMessage tools to contact the lead agent; use the Oz CLI commands below. +Do not use Claude Code Agent or SendMessage tools to contact the lead agent; use the Galaxy CLI commands below. Do not ask to inspect help before messaging. The command shapes below are complete. Send a message to the lead agent at start, when blocked, and when complete: diff --git a/app/src/prompt/editor_modal.rs b/app/src/prompt/editor_modal.rs index 4f503e79..38d4dcc2 100644 --- a/app/src/prompt/editor_modal.rs +++ b/app/src/prompt/editor_modal.rs @@ -53,7 +53,7 @@ const MODAL_CONTENT_FONT_SIZE: f32 = 14.; const CHECKBOX_SIZE: f32 = 16.; const MODAL_TITLE: &str = "Edit prompt"; -const WARP_PROMPT_SECTION_HEADER: &str = "Warp terminal prompt"; +const WARP_PROMPT_SECTION_HEADER: &str = "Galaxy terminal prompt"; const SHELL_PROMPT_SECTION_HEADER: &str = "Shell prompt (PS1)"; const RESTORE_DEFAULT_BUTTON: &str = "Restore default"; diff --git a/app/src/search/command_palette/view.rs b/app/src/search/command_palette/view.rs index 8933ea84..94edfb0c 100644 --- a/app/src/search/command_palette/view.rs +++ b/app/src/search/command_palette/view.rs @@ -853,6 +853,7 @@ impl View { conversation_id, terminal_view_id, restore_layout: None, + local_only: false, }); send_telemetry_from_app_ctx!(TelemetryEvent::SelectNavigationPaletteItem, ctx); } diff --git a/app/src/search/command_search/warp_ai.rs b/app/src/search/command_search/warp_ai.rs index b3270195..a74a7748 100644 --- a/app/src/search/command_search/warp_ai.rs +++ b/app/src/search/command_search/warp_ai.rs @@ -29,8 +29,9 @@ use crate::ui_components::icons::Icon as UIIcon; use crate::util::color::{ContrastingColor, MinimumAllowedContrast}; use crate::workflows::{AIWorkflowOrigin, WorkflowSource, WorkflowType}; -const OPEN_WARP_AI_ITEM_BODY_TEXT: &str = "Ask Warp AI for command suggestions"; -const TRANSLATE_WITH_WARP_AI_ITEM_BODY_TEXT: &str = "Translate into shell command using Warp AI"; +const OPEN_WARP_AI_ITEM_BODY_TEXT: &str = "Ask Galaxy Agent for command suggestions"; +const TRANSLATE_WITH_WARP_AI_ITEM_BODY_TEXT: &str = + "Translate into a shell command using Galaxy Agent"; #[derive(Clone, Debug)] pub enum WarpAISearchItem { @@ -133,7 +134,7 @@ impl SearchItem for WarpAISearchItem { } fn accessibility_label(&self) -> String { - format!("Warp AI: {}", self.item_body_text()) + format!("Galaxy Agent: {}", self.item_body_text()) } } diff --git a/app/src/search/slash_command_menu/static_commands/commands.rs b/app/src/search/slash_command_menu/static_commands/commands.rs index 2f95e885..70e61f77 100644 --- a/app/src/search/slash_command_menu/static_commands/commands.rs +++ b/app/src/search/slash_command_menu/static_commands/commands.rs @@ -70,7 +70,7 @@ pub const CREATE_DOCKER_SANDBOX: StaticCommand = StaticCommand { pub static CREATE_NEW_PROJECT: LazyLock = LazyLock::new(|| StaticCommand { name: "/create-new-project", - description: "Have Oz walk you through creating a new coding project", + description: "Have Galaxy walk you through creating a new coding project", icon_path: "bundled/svg/plus.svg", availability: Availability::LOCAL | Availability::AI_ENABLED, auto_enter_ai_mode: true, @@ -79,7 +79,7 @@ pub static CREATE_NEW_PROJECT: LazyLock = LazyLock::new(|| Static pub static EDIT_SKILL: LazyLock = LazyLock::new(|| StaticCommand { name: "/open-skill", - description: "Open a skill's markdown file in Warp's built-in editor", + description: "Open a skill's markdown file in Galaxy's built-in editor", icon_path: "bundled/svg/file-code-02.svg", availability: Availability::AI_ENABLED, auto_enter_ai_mode: false, @@ -119,7 +119,7 @@ pub const ADD_RULE: StaticCommand = StaticCommand { pub static EDIT: LazyLock = LazyLock::new(|| StaticCommand { name: "/open-file", - description: "Open a file in Warp's code editor", + description: "Open a file in Galaxy's code editor", icon_path: "bundled/svg/file-code-02.svg", availability: Availability::LOCAL, auto_enter_ai_mode: false, diff --git a/app/src/settings_view/appearance_page.rs b/app/src/settings_view/appearance_page.rs index 88cb83c6..ed48f6ed 100644 --- a/app/src/settings_view/appearance_page.rs +++ b/app/src/settings_view/appearance_page.rs @@ -253,7 +253,7 @@ pub fn init_actions_from_parent_view( // Add command palette entry for toggling between Warp and Classic input modes app.register_fixed_bindings(vec![FixedBinding::empty( - "Toggle Input Mode (Warp/Classic)".to_string(), + "Toggle Input Mode (Galaxy/Classic)".to_string(), builder(SettingsAction::AppearancePageToggle( AppearancePageAction::ToggleInputMode, )), @@ -1643,7 +1643,7 @@ impl AppearanceSettingsPageView { fn input_mode_dropdown_item_label(val: InputMode) -> &'static str { match val { - InputMode::PinnedToBottom => "Pin to the bottom (Warp mode)", + InputMode::PinnedToBottom => "Pin to the bottom (Galaxy mode)", InputMode::PinnedToTop => "Pin to the top (Reverse mode)", InputMode::Waterfall => "Start at the top (Classic mode)", } @@ -3061,7 +3061,7 @@ impl SettingsWidget for CustomAppIconWidget { ); let show_dock_icon_toggle = render_body_item::( - "Show Warp in Dock".into(), + "Show Galaxy in Dock".into(), None, LocalOnlyIconState::for_setting( ShowDockIconState::storage_key(), @@ -3101,7 +3101,7 @@ impl SettingsWidget for CustomAppIconWidget { appearance .ui_builder() .wrappable_text( - "You may need to restart Warp for MacOS to apply the preferred icon style.", + "You may need to restart Galaxy for macOS to apply the preferred icon style.", true, ) .with_style(UiComponentStyles { @@ -3590,7 +3590,7 @@ impl SettingsWidget for InputTypeWidget { .radio_buttons( self.radio_buttons_states.clone(), vec![ - RadioButtonItem::text("Warp"), + RadioButtonItem::text("Galaxy"), RadioButtonItem::text("Shell (PS1)"), ], view.input_type_radio_state.clone(), @@ -5093,7 +5093,7 @@ impl SettingsWidget for UseLatestUserPromptAsConversationTitleInTabNamesWidget { }) .finish(), Some( - "Show the latest user prompt instead of the generated conversation title for Oz and third-party agent sessions in vertical tabs." + "Show the latest user prompt instead of the generated conversation title for Galaxy and third-party agent sessions in vertical tabs." .to_string(), ), ) diff --git a/app/src/settings_view/features_page.rs b/app/src/settings_view/features_page.rs index f4150ccc..b9800128 100644 --- a/app/src/settings_view/features_page.rs +++ b/app/src/settings_view/features_page.rs @@ -303,7 +303,7 @@ pub fn init_actions_from_parent_view( ]; toggle_binding_pairs.push(ToggleSettingActionPair::new( - "reuse existing SSH ControlMaster in the Warp SSH wrapper", + "reuse existing SSH ControlMaster in the Galaxy SSH wrapper", builder(SettingsAction::FeaturesPageToggle( FeaturesPageAction::ToggleSshReuseControlMaster, )), @@ -714,7 +714,7 @@ pub fn init_actions_from_parent_view( if DefaultTerminal::can_warp_become_default() { app.register_fixed_bindings([FixedBinding::empty( - "Make Warp the default terminal", + "Make Galaxy the default terminal", builder(SettingsAction::FeaturesPageToggle( FeaturesPageAction::MakeWarpDefaultTerminal, )), @@ -4822,9 +4822,9 @@ impl SettingsWidget for LoginItemWidget { let general_settings = GeneralSettings::as_ref(app); let ui_builder = appearance.ui_builder(); #[cfg(target_os = "macos")] - let label = "Start Warp at login (requires macOS 13+)"; + let label = "Start Galaxy at login (requires macOS 13+)"; #[cfg(not(target_os = "macos"))] - let label = "Start Warp at login"; + let label = "Start Galaxy at login"; render_body_item::( label.into(), None, @@ -5101,7 +5101,7 @@ impl SettingsWidget for DefaultTerminalWidget { let default_terminal = DefaultTerminal::as_ref(app); if default_terminal.is_warp_default() { ui_builder - .wrappable_text("Warp is the default terminal", true) + .wrappable_text("Galaxy is the default terminal", true) .with_style(UiComponentStyles { font_color: Some(appearance.theme().disabled_ui_text_color().into()), margin: Some(Coords::default().bottom(16.)), @@ -5112,7 +5112,7 @@ impl SettingsWidget for DefaultTerminalWidget { } else { ui_builder .link( - "Make Warp the default terminal".to_string(), + "Make Galaxy the default terminal".to_string(), None, Some(Box::new(|ctx| { ctx.dispatch_typed_action(FeaturesPageAction::MakeWarpDefaultTerminal); @@ -5207,7 +5207,7 @@ impl SettingsWidget for DesktopNotificationsWidget { let ui_builder = appearance.ui_builder(); let mut column = Flex::column(); column.add_child(render_body_item::( - "Receive desktop notifications from Warp".into(), + "Receive desktop notifications from Galaxy".into(), Some(AdditionalInfo { mouse_state: self.additional_info_link.clone(), on_click_action: Some(FeaturesPageAction::OpenUrl(NOTIFICATIONS_DOCS_URL.into())), @@ -7552,7 +7552,7 @@ impl SettingsWidget for WindowSystemWidget { may be blurry if your Wayland compositor is using fraction scaling (ex: 125%)." .to_string(); if view.force_x11_changed { - secondary_text.push_str("\n\nRestart Warp for changes to take effect."); + secondary_text.push_str("\n\nRestart Galaxy for changes to take effect."); } let warp_theme = appearance.theme(); children.add_child( diff --git a/app/src/settings_view/settings_file_footer.rs b/app/src/settings_view/settings_file_footer.rs index 09902bd4..1dc30058 100644 --- a/app/src/settings_view/settings_file_footer.rs +++ b/app/src/settings_view/settings_file_footer.rs @@ -35,10 +35,10 @@ const OPEN_BUTTON_HEIGHT: f32 = 32.; const ALERT_ACTION_BUTTON_HEIGHT: f32 = 24.; /// Size of the leading icons (search-sm, code-02, alert-circle, oz). const FOOTER_ICON_SIZE: f32 = 16.; -/// Size of the Oz brand mark inside the "Fix with Oz" button. Matches the +/// Size of the agent icon inside the "Fix with Galaxy Agent" button. Matches the /// Figma spec and the workspace banner's secondary-button icon sizing. const ALERT_OZ_ICON_SIZE: f32 = 14.; -/// Horizontal padding inside the "Open file" / "Fix with Oz" action buttons. +/// Horizontal padding inside the "Open file" / "Fix with Galaxy Agent" action buttons. /// Matches the workspace banner's secondary button pad. const ALERT_BUTTON_HORIZONTAL_PADDING: f32 = 8.; /// Spacing between the two action buttons when they fit on one row. @@ -250,7 +250,7 @@ pub fn render_settings_error_alert( ui_font_family, text_color, mouse_states.alert_fix_with_oz_button.clone(), - "Fix with Oz", + "Fix with Galaxy Agent", Some(Icon::Oz), /*bordered=*/ false, WorkspaceAction::FixSettingsWithOz { error_description }, diff --git a/app/src/terminal/cli_agent_sessions/plugin_manager/codex.rs b/app/src/terminal/cli_agent_sessions/plugin_manager/codex.rs index 11d60b1f..596bb0c1 100644 --- a/app/src/terminal/cli_agent_sessions/plugin_manager/codex.rs +++ b/app/src/terminal/cli_agent_sessions/plugin_manager/codex.rs @@ -193,11 +193,11 @@ impl CliAgentPluginManager for CodexPluginManager { } fn install_success_message(&self) -> &'static str { - "Warp plugin installed. Please restart Codex to activate." + "Galaxy plugin installed. Please restart Codex to activate." } fn update_success_message(&self) -> &'static str { - "Warp plugin updated. Please restart Codex to activate." + "Galaxy plugin updated. Please restart Codex to activate." } fn install_instructions(&self) -> &'static PluginInstructions { @@ -273,17 +273,17 @@ impl CliAgentPluginManager for CodexPluginManager { static PLUGIN_INSTALL_INSTRUCTIONS: LazyLock = LazyLock::new(|| PluginInstructions { - title: "Install Warp Plugin for Codex", + title: "Install Galaxy notification plugin for Codex", subtitle: "Run the following commands, then restart Codex.", steps: &[ PluginInstructionStep { - description: "Add the Warp plugin marketplace repository", + description: "Add the notification plugin marketplace repository", command: "codex plugin marketplace add warpdotdev/codex-warp", executable: true, link: None, }, PluginInstructionStep { - description: "Install the Warp plugin", + description: "Install the notification plugin", command: "codex plugin add warp@codex-warp", executable: true, link: None, @@ -294,8 +294,8 @@ static PLUGIN_INSTALL_INSTRUCTIONS: LazyLock = static NATIVE_INSTALL_INSTRUCTIONS: LazyLock = LazyLock::new(|| { PluginInstructions { - title: "Enable Warp Notifications for Codex", - subtitle: "Update Codex to the latest version, then enable in-focus notifications so Warp can display them while you work.", + title: "Enable Galaxy notifications for Codex", + subtitle: "Update Codex to the latest version, then enable in-focus notifications so Galaxy can display them while you work.", steps: &[ PluginInstructionStep { description: "Update Codex to the latest version.", @@ -323,7 +323,7 @@ static EMPTY_INSTRUCTIONS: LazyLock = LazyLock::new(|| Plugi static PLUGIN_UPDATE_INSTRUCTIONS: LazyLock = LazyLock::new(|| { PluginInstructions { - title: "Update Warp Plugin for Codex", + title: "Update Galaxy notification plugin for Codex", subtitle: "Run the following commands, then restart Codex.", steps: &[ PluginInstructionStep { @@ -333,7 +333,7 @@ static PLUGIN_UPDATE_INSTRUCTIONS: LazyLock = LazyLock::new( link: None, }, PluginInstructionStep { - description: "Reinstall the Warp plugin", + description: "Reinstall the notification plugin", command: "codex plugin add warp@codex-warp", executable: true, link: None, diff --git a/app/src/terminal/cli_agent_sessions/plugin_manager/codex_tests.rs b/app/src/terminal/cli_agent_sessions/plugin_manager/codex_tests.rs index 71ec515b..d65e2e84 100644 --- a/app/src/terminal/cli_agent_sessions/plugin_manager/codex_tests.rs +++ b/app/src/terminal/cli_agent_sessions/plugin_manager/codex_tests.rs @@ -21,7 +21,7 @@ fn can_auto_install_is_false_without_codex_plugin() { fn install_instructions_are_native_without_codex_plugin() { let _guard = FeatureFlag::CodexPlugin.override_enabled(false); let instructions = CodexPluginManager::new(None, None, None).install_instructions(); - assert_eq!(instructions.title, "Enable Warp Notifications for Codex"); + assert_eq!(instructions.title, "Enable Galaxy notifications for Codex"); assert_eq!( instructions.steps[1].command, "[tui]\nnotification_condition = \"always\"" diff --git a/app/src/terminal/input/models/data_source.rs b/app/src/terminal/input/models/data_source.rs index e3ff9bdc..8d7ede10 100644 --- a/app/src/terminal/input/models/data_source.rs +++ b/app/src/terminal/input/models/data_source.rs @@ -45,7 +45,7 @@ use crate::terminal::view::ambient_agent::AmbientAgentViewModel; use crate::workspace::WorkspaceAction; use crate::workspaces::user_workspaces::UserWorkspaces; -const AUTO_BEDROCK_TOOLTIP: &str = "Warp uses Bedrock when the model Auto selects supports it; otherwise it may use Warp-hosted inference."; +const AUTO_BEDROCK_TOOLTIP: &str = "Galaxy uses Bedrock when the model Auto selects supports it. Other models use their configured provider."; #[derive(Clone, Debug)] pub struct AcceptModel { diff --git a/app/src/terminal/ssh/wormhole.rs b/app/src/terminal/ssh/wormhole.rs index 9b500cce..2b19ea36 100644 --- a/app/src/terminal/ssh/wormhole.rs +++ b/app/src/terminal/ssh/wormhole.rs @@ -79,7 +79,7 @@ pub fn wormhole_description( let description = FormattedText::new(vec![FormattedTextLine::Line(vec![ FormattedTextFragment::plain_text( - "Bring Galaxy's features to your remote session: blocks, full text editing, completions, Oz, and more." + "Bring Galaxy's features to your remote session: blocks, full text editing, completions, Galaxy Agent, and more." ), ])]); wormhole::render::build_description_row(description, theme, appearance, hyperlink_index.clone()) diff --git a/app/src/terminal/view.rs b/app/src/terminal/view.rs index 4610cb8e..c67da273 100644 --- a/app/src/terminal/view.rs +++ b/app/src/terminal/view.rs @@ -7801,24 +7801,24 @@ impl TerminalView { .get_pending_action(app) .map(|action| match &action.action { AIAgentActionType::RequestCommandOutput { command, .. } => { - format!("Oz needs your permission to run `{command}`") + format!("Galaxy needs your permission to run `{command}`") } AIAgentActionType::ReadFiles(..) => { - "Oz needs your permission to read files".to_string() + "Galaxy needs your permission to read files".to_string() } AIAgentActionType::SearchCodebase(..) => { - "Oz needs your permission to search your codebase".to_string() + "Galaxy needs your permission to search your codebase".to_string() } AIAgentActionType::RequestFileEdits { .. } => { - "Oz needs your permission to edit a file".to_string() + "Galaxy needs your permission to edit a file".to_string() } AIAgentActionType::WriteToLongRunningShellCommand { .. } => { - "Oz needs your permission to interact with a running shell command" + "Galaxy needs your permission to interact with a running shell command" .to_string() } - _ => "Oz needs your confirmation to continue".to_string(), + _ => "Galaxy needs your confirmation to continue".to_string(), }) - .unwrap_or("Oz needs your confirmation to continue".to_string()); + .unwrap_or("Galaxy needs your confirmation to continue".to_string()); return Some(AIBlockNotificationSummary { success: false, title, diff --git a/app/src/user_config/mod.rs b/app/src/user_config/mod.rs index 2407f2d1..594b7233 100644 --- a/app/src/user_config/mod.rs +++ b/app/src/user_config/mod.rs @@ -24,7 +24,7 @@ use crate::workflows::workflow::Workflow; lazy_static! { pub static ref LAUNCH_CONFIG_COMMENT: String = format!( - "# Warp Launch Configuration + "# Galaxy Launch Configuration # # # Use this to start a certain configuration of windows, tabs, and panes. diff --git a/app/src/workspace/action.rs b/app/src/workspace/action.rs index 6246b560..efd3a2e6 100644 --- a/app/src/workspace/action.rs +++ b/app/src/workspace/action.rs @@ -592,6 +592,8 @@ pub enum WorkspaceAction { /// If provided, use this layout to restore the conversation. /// Otherwise, fall back to the user's setting. restore_layout: Option, + /// History recovery must use this device's saved data without server fallback. + local_only: bool, }, /// Fork an existing AI conversation. /// Optionally summarizes the conversation after forking and/or sends an initial prompt. diff --git a/app/src/workspace/cli_install.rs b/app/src/workspace/cli_install.rs index e62c3fb7..1e676821 100644 --- a/app/src/workspace/cli_install.rs +++ b/app/src/workspace/cli_install.rs @@ -7,7 +7,7 @@ use command::blocking::Command; use galaxy_core::channel::ChannelState; use galaxy_util::path::ShellFamily; -/// Compute the target path where the Oz CLI symlink should be installed, based on channel +/// Compute the target path where the Galaxy AI CLI symlink should be installed, based on channel fn oz_install_target_path() -> PathBuf { PathBuf::from("/usr/local/bin").join(ChannelState::channel().cli_command_name()) } @@ -19,8 +19,8 @@ fn galaxyctrl_install_target_path() -> PathBuf { /// Compute the source path of the galaxyctrl wrapper inside the current app bundle. /// -/// Oz commands are part of the shared executable's normal argument parser, so -/// Oz can symlink directly to the current executable. Galaxy Control has a +/// Galaxy AI commands are part of the shared executable's normal argument parser, so +/// Galaxy AI can symlink directly to the current executable. Galaxy Control has a /// separate parser selected by the hidden `--galaxyctrl` flag, so its installed /// symlink must target the bundled wrapper that injects that flag. Without it, /// Galaxy Control subcommands such as `tab` would reach the normal parser and be @@ -196,21 +196,21 @@ fn uninstall_symlink(target: &Path, command_name: &str) -> Result<()> { Ok(()) } -/// Install the Oz CLI by symlinking the shared Warp executable into /usr/local/bin. +/// Install the Galaxy AI CLI by symlinking the shared Galaxy executable into /usr/local/bin. /// -/// The normal argument parser dispatches Oz subcommands directly. It also uses -/// the `oz`-prefixed invocation name to print CLI help rather than launch the +/// The normal argument parser dispatches Galaxy AI subcommands directly. It also +/// recognizes the CLI invocation name to print CLI help rather than launch the /// GUI when no subcommand is provided. pub fn install_oz() -> Result<()> { let oz_path = oz_install_target_path(); let current_binary = std::env::current_exe().context("Failed to get current executable path")?; - install_symlink(¤t_binary, &oz_path, "Oz CLI") + install_symlink(¤t_binary, &oz_path, "Galaxy AI CLI") } -/// Uninstall the Oz CLI by removing the symlink from /usr/local/bin +/// Uninstall the Galaxy AI CLI by removing the symlink from /usr/local/bin pub fn uninstall_oz() -> Result<()> { - uninstall_symlink(&oz_install_target_path(), "Oz command") + uninstall_symlink(&oz_install_target_path(), "Galaxy AI command") } /// Install Galaxy Control by symlinking its bundled wrapper into /usr/local/bin. diff --git a/app/src/workspace/mod.rs b/app/src/workspace/mod.rs index e52983c2..a35ebb87 100644 --- a/app/src/workspace/mod.rs +++ b/app/src/workspace/mod.rs @@ -1234,20 +1234,20 @@ pub fn init(app: &mut AppContext) { .with_context_predicate(id!("Workspace") & id!(flags::ENABLE_WARP_DRIVE))]); } - // Oz and Galaxy Control CLI install/uninstall actions (macOS only) + // Galaxy AI and Galaxy Control CLI install/uninstall actions (macOS only) #[cfg(target_os = "macos")] { app.register_editable_bindings([ EditableBinding::new( "workspace:install_cli", - "Install Oz CLI globally for use outside of Warp", + "Install Galaxy AI CLI globally for use outside of Galaxy", WorkspaceAction::InstallOz, ) .with_group(bindings::BindingGroup::Settings.as_str()) .with_context_predicate(id!("Workspace")), EditableBinding::new( "workspace:uninstall_cli", - "Undo global Oz CLI installation (oz will still work within Warp)", + "Remove global Galaxy AI CLI installation", WorkspaceAction::UninstallOz, ) .with_group(bindings::BindingGroup::Settings.as_str()) diff --git a/app/src/workspace/view.rs b/app/src/workspace/view.rs index 89fc13b7..434b54b6 100644 --- a/app/src/workspace/view.rs +++ b/app/src/workspace/view.rs @@ -8950,34 +8950,35 @@ impl Workspace { } } - /// Install the Oz CLI by creating a symlink in /usr/local/bin + /// Install the Galaxy AI CLI by creating a symlink in /usr/local/bin #[cfg(target_os = "macos")] fn install_oz(&mut self, ctx: &mut ViewContext) { ctx.spawn(async { cli_install::install_oz() }, |view, result, ctx| { let command_name = ChannelState::channel().cli_command_name(); - let message = format!("Installed the Oz CLI globally. You can now run '{command_name}' from any terminal outside of Warp."); - let toast = DismissibleToast::success(message).with_link( - ToastLink::new("Learn more".to_string()) - .with_href("https://docs.warp.dev/reference/cli".to_string()), + let message = format!("Installed the Galaxy AI CLI globally. You can now run '{command_name}' from any terminal."); + let toast = DismissibleToast::success(message); + view.handle_cli_command_result( + result, + toast, + "Failed to install Galaxy AI command", + ctx, ); - view.handle_cli_command_result(result, toast, "Failed to install Oz command", ctx); }); } - /// Uninstall the Oz CLI by removing the symlink from /usr/local/bin + /// Uninstall the Galaxy AI CLI by removing the symlink from /usr/local/bin #[cfg(target_os = "macos")] fn uninstall_oz(&mut self, ctx: &mut ViewContext) { ctx.spawn( async { cli_install::uninstall_oz() }, |view, result, ctx| { let toast = DismissibleToast::success( - "Removed the global Oz CLI installation — it still works inside Warp." - .to_string(), + "Removed the global Galaxy AI CLI installation.".to_string(), ); view.handle_cli_command_result( result, toast, - "Failed to uninstall Oz command", + "Failed to uninstall Galaxy AI command", ctx, ); }, @@ -9212,10 +9213,10 @@ impl Workspace { let window_id = self.window_id; let view_id = self.agent_management_view.id(); AgentConversationsModel::handle(ctx).update(ctx, |model, ctx| { + // This history surface consumes local conversations, not hosted run updates. + model.register_view_closed(window_id, view_id, ctx); if is_open { - model.register_view_open(window_id, view_id, ctx); - } else { - model.register_view_closed(window_id, view_id, ctx); + model.sync_conversations(ctx); } }); @@ -12978,6 +12979,7 @@ impl Workspace { pane_view_locator: Option, terminal_view_id: Option, mut restore_layout: Option, + local_only: bool, ctx: &mut ViewContext, ) { // If we have all required navigation data, try to navigate to the existing pane @@ -13022,8 +13024,27 @@ impl Workspace { UndoCloseStack::handle(ctx).update(ctx, |stack, ctx| { stack.discard_pane_group_parent(pane_view_locator.pane_group_id, ctx); }); - } else if has_blocking_long_running_command { - // Don't open in this pane and use the existing restore layout. + } else if has_blocking_long_running_command + || (local_only + && (ctx + .view_with_id::(window_id, terminal_view_id) + .is_none() + || ctx + .view_with_id::(window_id, pane_view_locator.pane_group_id) + .and_then(|pane_group| { + pane_group.as_ref(ctx).terminal_view_from_pane_id( + pane_view_locator.pane_id, + ctx, + ) + }) + .is_none_or(|terminal| terminal.id() != terminal_view_id) + || ctx.root_view_id(window_id).is_none())) + || (local_only + && BlocklistAIHistoryModel::as_ref(ctx) + .conversation(&conversation_id) + .is_none()) + { + // Local history restores from disk if its previous pane is unavailable. } else { BlocklistAIHistoryModel::handle(ctx).update(ctx, |history_model, ctx| { history_model.set_active_conversation_id( @@ -13070,13 +13091,13 @@ impl Workspace { // Handle based on effective layout match effective_layout { RestoreConversationLayout::ActivePane => { - self.restore_conversation_in_active_pane(conversation_id, ctx); + self.restore_conversation_in_active_pane(conversation_id, local_only, ctx); } RestoreConversationLayout::SplitPane => { - self.restore_conversation_in_split_pane(conversation_id, ctx); + self.restore_conversation_in_split_pane(conversation_id, local_only, ctx); } RestoreConversationLayout::NewTab => { - self.restore_conversation_in_new_tab(conversation_id, ctx); + self.restore_conversation_in_new_tab(conversation_id, local_only, ctx); } } } @@ -13087,6 +13108,7 @@ impl Workspace { fn restore_conversation_in_active_pane( &mut self, conversation_id: AIConversationId, + local_only: bool, ctx: &mut ViewContext, ) { let terminal_view_for_active_pane = self.active_session_view(ctx).filter(|_| { @@ -13097,7 +13119,7 @@ impl Workspace { // If we can't restore in the active pane, fall back to restoring in new tab. let Some(terminal_view) = &terminal_view_for_active_pane else { - self.restore_conversation_in_new_tab(conversation_id, ctx); + self.restore_conversation_in_new_tab(conversation_id, local_only, ctx); return; }; @@ -13143,7 +13165,7 @@ impl Workspace { }); // If there's an active long-running command in the active terminal, fall back to restoring in new tab. if active_pane_has_long_running_command { - self.restore_conversation_in_new_tab(conversation_id, ctx); + self.restore_conversation_in_new_tab(conversation_id, local_only, ctx); return; } @@ -13152,9 +13174,11 @@ impl Workspace { .as_ref(ctx) .get_conversation_metadata(&conversation_id) .is_some_and(|m| m.has_local_data); - let future = history_model - .as_ref(ctx) - .load_conversation_data(conversation_id, ctx); + let future = if local_only { + history_model.as_ref(ctx).load_local_conversation_data(conversation_id) + } else { + history_model.as_ref(ctx).load_conversation_data(conversation_id, ctx) + }; let terminal_view_for_closure = terminal_view.clone(); let window_id = ctx.window_id(); ctx.spawn(future, move |_workspace, conversation, ctx| { @@ -13169,7 +13193,12 @@ impl Workspace { ctx.notify(); }); WorkspaceToastStack::handle(ctx).update(ctx, |toast_stack, ctx| { - let toast = DismissibleToast::error("Failed to load conversation.".to_owned()); + let message = if local_only { + "This AI session is not available on this device." + } else { + "Failed to load conversation." + }; + let toast = DismissibleToast::error(message.to_owned()); toast_stack.add_ephemeral_toast(toast, window_id, ctx); }); return; @@ -13208,6 +13237,7 @@ impl Workspace { fn restore_conversation_in_split_pane( &mut self, conversation_id: AIConversationId, + local_only: bool, ctx: &mut ViewContext, ) { let window_id = ctx.window_id(); @@ -13223,14 +13253,21 @@ impl Workspace { }); let history_model = BlocklistAIHistoryModel::handle(ctx); - let future = history_model - .as_ref(ctx) - .load_conversation_data(conversation_id, ctx); + let future = if local_only { + history_model.as_ref(ctx).load_local_conversation_data(conversation_id) + } else { + history_model.as_ref(ctx).load_conversation_data(conversation_id, ctx) + }; ctx.spawn(future, move |_workspace, conversation, ctx| { let Some(conversation) = conversation else { log::warn!("Failed to load conversation {conversation_id}"); WorkspaceToastStack::handle(ctx).update(ctx, |toast_stack, ctx| { - let toast = DismissibleToast::error("Failed to load conversation.".to_owned()); + let message = if local_only { + "This AI session is not available on this device." + } else { + "Failed to load conversation." + }; + let toast = DismissibleToast::error(message.to_owned()); toast_stack.add_ephemeral_toast(toast, window_id, ctx); }); // Close the loading pane @@ -13262,6 +13299,7 @@ impl Workspace { fn restore_conversation_in_new_tab( &mut self, conversation_id: AIConversationId, + local_only: bool, ctx: &mut ViewContext, ) { let window_id = ctx.window_id(); @@ -13291,15 +13329,22 @@ impl Workspace { let loading_pane_id = new_pane_group.as_ref(ctx).focused_pane_id(ctx); let history_model = BlocklistAIHistoryModel::handle(ctx); - let future = history_model - .as_ref(ctx) - .load_conversation_data(conversation_id, ctx); + let future = if local_only { + history_model.as_ref(ctx).load_local_conversation_data(conversation_id) + } else { + history_model.as_ref(ctx).load_conversation_data(conversation_id, ctx) + }; ctx.spawn(future, move |workspace, conversation, ctx| { let Some(conversation) = conversation else { log::warn!("Failed to load conversation {conversation_id}"); WorkspaceToastStack::handle(ctx).update(ctx, |toast_stack, ctx| { - let toast = DismissibleToast::error("Failed to load conversation.".to_owned()); + let message = if local_only { + "This AI session is not available on this device." + } else { + "Failed to load conversation." + }; + let toast = DismissibleToast::error(message.to_owned()); toast_stack.add_ephemeral_toast(toast, window_id, ctx); }); // Close the loading tab @@ -21892,7 +21937,7 @@ impl Workspace { AISettings::as_ref(app) .is_any_ai_enabled(app) .then(|| WorkspaceBannerButtonDetails { - text: "Fix with Oz".to_owned(), + text: "Fix with Galaxy Agent".to_owned(), action: WorkspaceAction::FixSettingsWithOz { error_description: error.to_string(), }, @@ -24893,7 +24938,7 @@ impl TypedActionView for Workspace { } RunAISuggestedCommand(code) => { let command = code.trim().to_string(); - let workflow = Workflow::new("Command from Oz", command); + let workflow = Workflow::new("Command from Galaxy Agent", command); self.run_workflow_in_active_input( &WorkflowType::AIGenerated { workflow, @@ -25148,6 +25193,7 @@ impl TypedActionView for Workspace { conversation_id, terminal_view_id, restore_layout, + local_only, } => { self.restore_or_navigate_to_conversation( *conversation_id, @@ -25155,6 +25201,7 @@ impl TypedActionView for Workspace { *pane_view_locator, *terminal_view_id, *restore_layout, + *local_only, ctx, ); } diff --git a/app/src/workspace/view/conversation_list/view.rs b/app/src/workspace/view/conversation_list/view.rs index 062a7cd7..009fee07 100644 --- a/app/src/workspace/view/conversation_list/view.rs +++ b/app/src/workspace/view/conversation_list/view.rs @@ -30,7 +30,7 @@ use super::view_model::{ConversationEntry, ConversationListViewModel}; use crate::ai::active_agent_views_model::{ActiveAgentViewsModel, ConversationOrTaskId}; use crate::ai::agent::conversation::AIConversationId; use crate::ai::agent_conversations_model::{ - AgentConversationEntryId, AgentConversationNavigationSubject, AgentConversationsModel, + AgentConversationEntryId, AgentConversationsModel, }; use crate::ai::agent_management::telemetry::{AgentManagementTelemetryEvent, OpenedFrom}; use crate::ai::blocklist::history_model::BlocklistAIHistoryModel; @@ -707,8 +707,8 @@ impl ConversationListView { ctx.emit(Event::NewConversationInNewTab); } ListItem::Conversation { entry, .. } => { - if let Some(action) = AgentConversationsModel::resolve_open_action( - AgentConversationNavigationSubject::Entry(entry.id), + if let Some(action) = AgentConversationsModel::resolve_local_open_action( + entry.id, None, ctx, ) { @@ -1404,8 +1404,8 @@ impl TypedActionView for ConversationListView { }); } ConversationListViewAction::OpenItem { id } => { - let Some(action) = AgentConversationsModel::resolve_open_action( - AgentConversationNavigationSubject::Entry(*id), + let Some(action) = AgentConversationsModel::resolve_local_open_action( + *id, None, ctx, ) else { diff --git a/app/src/workspace/view/conversation_list/view_model.rs b/app/src/workspace/view/conversation_list/view_model.rs index cb90bccf..bf58aa18 100644 --- a/app/src/workspace/view/conversation_list/view_model.rs +++ b/app/src/workspace/view/conversation_list/view_model.rs @@ -77,7 +77,7 @@ impl ConversationListViewModel { fn refresh_cached_items(&mut self, ctx: &mut ModelContext) { let model = self.conversations_model.as_ref(ctx); self.cached_entry_ids = model - .get_entries( + .get_local_entries( &AgentManagementFilters { owners: OwnerFilter::PersonalOnly, status: StatusFilter::All, @@ -127,7 +127,7 @@ impl ConversationListViewModel { .cached_entry_ids .iter() .filter_map(|id| { - let item = conversations_model.get_entry_by_id(id, ctx)?; + let item = conversations_model.get_local_entry_by_id(id, ctx)?; match_indices_case_insensitive(&item.display.title, &search_query).map( |result| { @@ -165,7 +165,7 @@ impl ConversationListViewModel { ctx: &AppContext, ) -> Option { let model = self.conversations_model.as_ref(ctx); - model.get_entry_by_id(id, ctx) + model.get_local_entry_by_id(id, ctx) } pub fn current_ids(&self) -> impl Iterator { diff --git a/app/src/workspace/view/left_panel.rs b/app/src/workspace/view/left_panel.rs index bd2d5b07..450a3637 100644 --- a/app/src/workspace/view/left_panel.rs +++ b/app/src/workspace/view/left_panel.rs @@ -1089,9 +1089,7 @@ impl LeftPanelView { } } - /// When the conversation list view's visibility changes, - /// we need to update the conversation and tasks model to reflect the new state - /// (this information is used to decide whether or not we should poll for new tasks). + /// Refresh local history without subscribing this surface to hosted task polling. fn on_conversation_list_view_visibility_changed( &self, is_now_open: bool, @@ -1100,10 +1098,9 @@ impl LeftPanelView { let window_id = ctx.window_id(); let view_id = self.conversation_list_view.id(); AgentConversationsModel::handle(ctx).update(ctx, |model, ctx| { + model.register_view_closed(window_id, view_id, ctx); if is_now_open { - model.register_view_open(window_id, view_id, ctx); - } else { - model.register_view_closed(window_id, view_id, ctx); + model.sync_conversations(ctx); } }); } diff --git a/app/src/workspace/view/vertical_tabs.rs b/app/src/workspace/view/vertical_tabs.rs index 962958a3..5a70b485 100644 --- a/app/src/workspace/view/vertical_tabs.rs +++ b/app/src/workspace/view/vertical_tabs.rs @@ -4182,7 +4182,7 @@ fn terminal_kind_badge_label(is_oz_agent: bool, cli_agent: Option) -> if let Some(cli_agent) = cli_agent { cli_agent.display_name().to_string() } else if is_oz_agent { - "Oz".to_string() + "Galaxy Agent".to_string() } else { "Terminal".to_string() } diff --git a/app/src/workspace/view_tests.rs b/app/src/workspace/view_tests.rs index 1e00d060..1c97a218 100644 --- a/app/src/workspace/view_tests.rs +++ b/app/src/workspace/view_tests.rs @@ -928,7 +928,7 @@ fn restore_conversation_in_active_pane_enters_existing_live_conversation_without workspace.update(&mut app, |workspace, ctx| { assert_eq!(workspace.tab_count(), 1); - workspace.restore_conversation_in_active_pane(conversation_id, ctx); + workspace.restore_conversation_in_active_pane(conversation_id, false, ctx); assert_eq!(workspace.tab_count(), 1); }); @@ -3577,7 +3577,7 @@ fn test_standard_tab_context_menu_shows_hover_only_tab_bar() { } #[test] -fn test_open_cloud_agent_setup_guide_action_opens_management_view_and_is_idempotent() { +fn test_legacy_cloud_setup_guide_action_opens_local_history_without_onboarding() { let _agent_management_guard = FeatureFlag::AgentManagementView.override_enabled(true); App::test((), |mut app| async move { @@ -3598,7 +3598,7 @@ fn test_open_cloud_agent_setup_guide_action_opens_management_view_and_is_idempot .current_workspace_state .is_agent_management_view_open ); - assert!(workspace + assert!(!workspace .agent_management_view .as_ref(ctx) .is_showing_setup_guide()); @@ -3609,7 +3609,7 @@ fn test_open_cloud_agent_setup_guide_action_opens_management_view_and_is_idempot .current_workspace_state .is_agent_management_view_open ); - assert!(workspace + assert!(!workspace .agent_management_view .as_ref(ctx) .is_showing_setup_guide()); diff --git a/crates/galaxy_cli/src/agent.rs b/crates/galaxy_cli/src/agent.rs index 135e171e..07a738ae 100644 --- a/crates/galaxy_cli/src/agent.rs +++ b/crates/galaxy_cli/src/agent.rs @@ -169,7 +169,7 @@ impl Harness { pub fn display_name(self) -> &'static str { match self { - Self::Oz => "Oz", + Self::Oz => "Galaxy", Self::Claude => "Claude Code", Self::OpenCode => "OpenCode", Self::Gemini => "Gemini CLI", From d5fa711e5c261436c0c87269974374f42572d67e Mon Sep 17 00:00:00 2001 From: Ryan Ward Date: Wed, 21 Oct 2026 13:42:28 +0000 Subject: [PATCH 10/12] Check protected-file access before macOS startup guidance Run a conservative, read-only preflight off the UI thread before offering Full Disk Access setup. Successful or inconclusive checks stay silent; only a protected-access denial offers guidance. Honor the native reminder opt-out locally and expose manual settings access in the macOS app menu. --- app/src/app_menus.rs | 6 ++ app/src/full_disk_access.rs | 138 ++++++++++++++++++++++++++++++++++++ app/src/lib.rs | 5 ++ 3 files changed, 149 insertions(+) create mode 100644 app/src/full_disk_access.rs diff --git a/app/src/app_menus.rs b/app/src/app_menus.rs index ec716def..c852e77a 100644 --- a/app/src/app_menus.rs +++ b/app/src/app_menus.rs @@ -158,6 +158,12 @@ fn make_new_app_menu(ctx: &AppContext) -> Menu { let preferences_menu_items = vec![ updateable_custom_item_without_checkmark(CustomAction::ShowSettings, ctx), + MenuItem::Custom(CustomMenuItem::new( + "Full Disk Access Settings...", + crate::full_disk_access::open_settings, + no_updates, + None, + )), MenuItem::Separator, updateable_custom_item_without_checkmark(CustomAction::ToggleKeybindingsPage, ctx), updateable_custom_item_without_checkmark(CustomAction::ConfigureKeybindings, ctx), diff --git a/app/src/full_disk_access.rs b/app/src/full_disk_access.rs new file mode 100644 index 00000000..6ba44704 --- /dev/null +++ b/app/src/full_disk_access.rs @@ -0,0 +1,138 @@ +//! Optional, device-local guidance for macOS Full Disk Access. +//! +//! macOS requires the user to grant this permission in System Settings. The +//! stored marker records dismissal of the guidance, never authorization. +//! A conservative filesystem preflight checks effective protected-file access; +//! macOS provides no supported API for querying the Full Disk Access toggle. + +use galaxy_core::user_preferences::GetUserPreferences; +use galaxyui::modals::{AlertDialogWithCallbacks, ModalButton}; +use galaxyui::AppContext; + +const PROMPT_DISMISSED_KEY: &str = "FullDiskAccessPromptDismissed"; +const SETTINGS_URL: &str = + "x-apple.systempreferences:com.apple.preference.security?Privacy_AllFiles"; + +#[derive(Clone, Copy, PartialEq, Eq)] +enum ProtectedFileAccess { + Accessible, + Denied, + Unknown, +} + +/// Checks access after the macOS GUI has opened, unless reminders were disabled. +pub(crate) fn prompt_on_startup(ctx: &mut AppContext) { + if !cfg!(target_os = "macos") || prompt_dismissed(ctx) { + return; + } + + let (sender, receiver) = futures::channel::oneshot::channel(); + ctx.background_executor() + .spawn(async move { + let _ = sender.send(protected_file_access()); + }) + .detach(); + + let weak_app = ctx.weak_app(); + ctx.foreground_executor() + .spawn(async move { + if receiver.await.ok() != Some(ProtectedFileAccess::Denied) { + return; + } + if let Some(mut app) = weak_app.upgrade() { + app.update(|ctx| { + if ctx.window_ids().next().is_some() && !prompt_dismissed(ctx) { + show_guidance(ctx); + } + }); + } + }) + .detach(); +} + +fn prompt_dismissed(ctx: &AppContext) -> bool { + match ctx + .private_user_preferences() + .read_value(PROMPT_DISMISSED_KEY) + { + Ok(Some(_)) => true, + Ok(None) => false, + Err(error) => { + log::warn!("Could not read Full Disk Access guidance preference: {error}"); + true + }, + } +} + +/// This is an effective-access heuristic, not a supported permission query. +/// Open only the current user's database; never inspect or parse its contents. +/// Missing files and ambiguous failures must not produce a permission prompt. +fn protected_file_access() -> ProtectedFileAccess { + #[cfg(unix)] + { + use std::fs::OpenOptions; + use std::os::unix::fs::OpenOptionsExt; + + let Some(home) = dirs::home_dir() else { + return ProtectedFileAccess::Unknown; + }; + let path = home.join("Library/Application Support/com.apple.TCC/TCC.db"); + match OpenOptions::new() + .read(true) + .custom_flags(libc::O_NOFOLLOW | libc::O_NONBLOCK) + .open(path) + { + Ok(file) => { + let is_regular_file = file.metadata().is_ok_and(|metadata| metadata.is_file()); + drop(file); + if is_regular_file { + ProtectedFileAccess::Accessible + } else { + ProtectedFileAccess::Unknown + } + }, + // EACCES may reflect ordinary POSIX permissions or ACLs, so only + // EPERM counts as protected-access denial. Neither is proof of the + // setting's value; all other errors remain unknown. + Err(error) if error.raw_os_error() == Some(libc::EPERM) => ProtectedFileAccess::Denied, + Err(_) => ProtectedFileAccess::Unknown, + } + } + #[cfg(not(unix))] + { + ProtectedFileAccess::Unknown + } +} + +fn show_guidance(ctx: &mut AppContext) { + ctx.show_native_platform_modal(AlertDialogWithCallbacks::for_app( + "Protected-file access is unavailable", + "Full Disk Access is optional. Enabling it gives Galaxy, including its agents and \ + commands you run, broad access to protected files on this Mac.\n\n\ + To enable it, open System Settings > Privacy & Security > Full Disk Access and \ + turn on Galaxy. If Galaxy is missing, click + and select the installed Galaxy app. \ + macOS may ask you to quit and reopen Galaxy.\n\n\ + You can continue without this permission and reopen these settings from Galaxy's menu later.", + vec![ + ModalButton::for_app("Open System Settings", open_settings), + ModalButton::for_app("Not Now", |_| {}), + ], + remember_dismissal, + )); +} + +/// Opens the macOS settings page without claiming or changing authorization. +pub(crate) fn open_settings(ctx: &mut AppContext) { + if cfg!(target_os = "macos") { + ctx.open_url(SETTINGS_URL); + } +} + +fn remember_dismissal(ctx: &mut AppContext) { + if let Err(error) = ctx + .private_user_preferences() + .write_value(PROMPT_DISMISSED_KEY, true.to_string()) + { + log::warn!("Could not save Full Disk Access guidance preference: {error}"); + } +} diff --git a/app/src/lib.rs b/app/src/lib.rs index 4ce9613b..cfb17b90 100644 --- a/app/src/lib.rs +++ b/app/src/lib.rs @@ -61,6 +61,7 @@ mod experiments; mod external_secrets; #[cfg(target_family = "wasm")] mod font_fallback; +mod full_disk_access; pub mod galaxy_managed_paths_watcher; pub use galaxy_managed_paths_watcher as warp_managed_paths_watcher; mod global_resource_handles; @@ -2826,6 +2827,10 @@ fn launch(ctx: &mut galaxyui::AppContext, app_state: Option, launch_mo timer.mark_interval_end("WINDOWS_CREATED"); }); + if cfg!(target_os = "macos") && matches!(launch_mode, LaunchMode::App { .. }) { + full_disk_access::prompt_on_startup(ctx); + } + #[cfg(feature = "local_fs")] if matches!(launch_mode, LaunchMode::App { .. }) && PersistenceWriter::as_ref(ctx).sender().is_none() From f6c6434e5e86de058dd685c743e679b63615caef Mon Sep 17 00:00:00 2001 From: Ryan Ward Date: Thu, 22 Oct 2026 14:25:21 +0000 Subject: [PATCH 11/12] Fix LiteLLM context budgets and direct computer use Keep efficient context windows separate from expanded maxima throughout discovery, display and provider budgets. Expose correlated computer-use requests and actions with existing permissions, preserve screenshots as images, and retain them across checkpoints, compaction and explicit history recall. Validation: cargo build on Linux passed. Added focused regression coverage; test execution and native GUI verification remain outside this CLI build-only workflow. --- AGENTS.md | 14 +- app/src/ai/agent/api/impl.rs | 8 + app/src/ai/agent/api/impl_tests.rs | 31 ++++ app/src/ai/agent/provider_history.rs | 75 +++++++++ app/src/ai/agent/provider_history_tests.rs | 23 +++ app/src/ai/blocklist/action_model.rs | 55 ++++++- .../action_model/execute/use_computer.rs | 6 +- app/src/ai/blocklist/action_model_tests.rs | 20 +++ app/src/ai/blocklist/controller.rs | 6 +- app/src/ai/llms.rs | 115 ++++++++++---- app/src/ai/llms_tests.rs | 150 ++++++++++++++++++ app/src/ai/openai/response_translator.rs | 2 + app/src/ai/provider/request_translator.rs | 71 +++++++++ app/src/ai/provider/response_translator.rs | 10 ++ app/src/ai/runtime/event_translator_tests.rs | 1 + app/src/ai/runtime/mod.rs | 1 + .../runtime/provider_run_coordinator_tests.rs | 2 + app/src/ai/runtime/rig.rs | 27 ++-- app/src/ai/runtime/rig_request.rs | 18 +-- app/src/ai/runtime/rig_request_tests.rs | 2 + app/src/ai/runtime/rig_tests.rs | 21 +++ app/src/ai/runtime/rig_tool.rs | 124 +++++++++++++++ app/src/ai/runtime/rig_tool_tests.rs | 49 ++++++ app/src/ai/runtime/task_progress_tests.rs | 1 + crates/galaxy_agent_core/src/provider_run.rs | 63 +++++--- .../src/provider_run_tests.rs | 125 ++++++++++++++- crates/galaxy_agent_core/src/tool_policy.rs | 96 +++++++++-- .../src/tool_policy_tests.rs | 114 +++++++++++++ crates/galaxy_agent_core/src/types.rs | 44 +++++ crates/galaxy_agent_core/src/types_tests.rs | 3 + 30 files changed, 1176 insertions(+), 101 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index e043e826..0ed85621 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -174,9 +174,10 @@ base_url = "http://localhost:4000/v1" api_key = "sk-..." [[ai.providers.models]] -model_id = "claude-sonnet-4-20250514[1m]" -display_name = "Claude Sonnet 4 (1M)" +model_id = "claude-sonnet-4-20250514" +display_name = "Claude Sonnet 4" context_size = 1000000 +default_context_size = 200000 [[ai.providers]] name = "Ollama (Local)" @@ -189,10 +190,9 @@ context_size = 128000 ``` **OpenAI/LiteLLM model discovery**: -- Models can be auto-fetched from the `/models` endpoint via the Settings > OpenAI / LiteLLM page -- For each model, the system probes `{model_id}[1m]` with a minimal chat completion request -- If the `[1m]` variant is accepted (HTTP 200 or 429), it's used with 1M context window -- Otherwise, the base model ID is used with its reported context size +- Models can be auto-fetched from `/models` and enriched with LiteLLM `/model/info` via the Settings > OpenAI / LiteLLM page. +- Preserve the efficient/default context window separately from the provider's maximum. Ordinary requests, compaction, and context displays use the lower positive default/input/maximum budget, including any explicit profile cap; availability of a larger maximum must never expand that budget automatically. +- Retain lower budgets when merging metadata from both endpoints, apply `effective_context_window_percent` once, and preserve the maximum as metadata. Discovery must not probe or select an extended `[1m]` variant automatically. - Models injected via `ai.providers[]` are routed to their specific endpoint (per-model routing map) - Provider name shown as the description label in the model picker; icon shows OpenAI logo for all OpenAI-compatible providers @@ -211,6 +211,8 @@ Key invariants: - A restored `AwaitingModel` checkpoint has an uncertain remote outcome and must terminate as an explicit restore failure rather than replaying the call; known recoverable failures observed in-process retain the bounded model-retry lifecycle - Same-conversation direct-provider follow-ups queue behind the cancelling generation; the old run keeps the active slot until terminal projection and cleanup, queued intent is checkpointed without credentials for restart recovery, and queued-only restore validates provider ownership and terminalizes the abandoned unprepared exchange from its persisted projection/stream identity before starting the successor; the next generation rebuilds provider history after cleanup so it includes the old generation's final committed output, and stale callbacks are ignored by stream identity - Known tools are in `KNOWN_TOOLS` in `response_translator.rs`; definitions are built by `tool_definition_for_name()` in `convert_request.rs` +- Direct-provider computer use starts with `request_computer_use`; the computer-use profile exposes `use_computer`, including when its model falls back to the main model. Tool availability and every action remain subject to session capabilities and the existing profile permission checks. +- Computer-use screenshots must reach providers as image parts, survive checkpoints and tool-result archival, and remain correlated with their tool call on explicit history recall. Emit every result in a tool batch before its screenshot image groups, and estimate image tokens without counting encoded image bytes as text. - Direct-provider normal and plan turns must advertise `read_plan`, `create_plan`, and `edit_plan` when the matching document capabilities are enabled; plan-creation requests should call `create_plan` after research rather than only returning prose - Direct-provider normal and plan turns expose `update_task_list` for explicit progress updates. Keep the validated current list in durable `ProviderRun` state, inject it into each request, preserve it across compaction and restore, and reconcile restored UI updates idempotently. Finishing a response must never mark unfinished work complete. CLI monitors and completed-command assessments retain their focused tool sets. - Completed local plans are archived intact into `Plans/Archive` after seven days only when their conversation has explicitly completed all checklist items and has no active provider run or child. Draft, unfinished, or reopened plans stay visible; resumed work restores an automatically archived plan, and manual recovery is respected. Archival never deletes plan contents. diff --git a/app/src/ai/agent/api/impl.rs b/app/src/ai/agent/api/impl.rs index 00b0d802..df4c1cb2 100644 --- a/app/src/ai/agent/api/impl.rs +++ b/app/src/ai/agent/api/impl.rs @@ -28,12 +28,17 @@ pub(crate) fn prepare_direct_provider_params( .clone() .unwrap_or_else(|| get_supported_tools(params)); remove_orchestration_tools_if_disabled(&mut supported_tools, params.orchestration_enabled); + // Computer use controls this machine, even when the terminal is connected elsewhere. + // Only expose the initial request here; the dedicated profile owns subsequent actions. + supported_tools.retain(|tool| *tool != api::ToolType::UseComputer + && (*tool != api::ToolType::RequestComputerUse || params.computer_use_enabled)); let mut supported_cli_agent_tools = supported_tools_override.unwrap_or_else(|| get_supported_cli_agent_tools(params)); remove_orchestration_tools_if_disabled( &mut supported_cli_agent_tools, params.orchestration_enabled, ); + supported_cli_agent_tools.retain(|tool| !matches!(tool, api::ToolType::UseComputer | api::ToolType::RequestComputerUse)); if params.permission_mode == crate::ai::execution_profiles::PermissionMode::ReadOnly { let permitted = |tool: &api::ToolType| is_read_only_tool(tool) || (!params.command_allowlist.is_empty() && matches!(tool, api::ToolType::RunShellCommand | api::ToolType::ReadShellCommandOutput | api::ToolType::WriteToLongRunningShellCommand)); @@ -73,6 +78,9 @@ fn get_supported_tools(params: &RequestParams) -> Vec { api::ToolType::CreateDocuments, api::ToolType::EditDocuments, ]; + if params.computer_use_enabled { + supported_tools.push(api::ToolType::RequestComputerUse); + } if params.orchestration_enabled { supported_tools.push(api::ToolType::Subagent); diff --git a/app/src/ai/agent/api/impl_tests.rs b/app/src/ai/agent/api/impl_tests.rs index 00b13989..2bfef573 100644 --- a/app/src/ai/agent/api/impl_tests.rs +++ b/app/src/ai/agent/api/impl_tests.rs @@ -91,6 +91,37 @@ fn supported_tools_include_plan_document_capabilities() { assert!(supported_tools.contains(&api::ToolType::EditDocuments)); } +#[test] +fn computer_use_is_explicitly_enabled_and_kept_out_of_command_monitor_tools() { + let mut params = request_params_with_ask_user_question_enabled(false); + params.computer_use_enabled = true; + let (normal, cli) = super::prepare_direct_provider_params(&mut params); + assert!(normal.contains(&api::ToolType::RequestComputerUse)); + assert!(!normal.contains(&api::ToolType::UseComputer)); + assert!(!cli.contains(&api::ToolType::RequestComputerUse)); + assert!(!cli.contains(&api::ToolType::UseComputer)); + + params.computer_use_enabled = false; + let (normal, _) = super::prepare_direct_provider_params(&mut params); + assert!(!normal.contains(&api::ToolType::RequestComputerUse)); +} + +#[test] +fn computer_use_never_expands_explicit_tool_overrides_or_read_only_permissions() { + let mut params = request_params_with_ask_user_question_enabled(false); + params.computer_use_enabled = true; + params.supported_tools_override = Some(Vec::new()); + let (normal, cli) = super::prepare_direct_provider_params(&mut params); + assert!(normal.is_empty()); + assert!(cli.is_empty()); + + params.permission_mode = crate::ai::execution_profiles::PermissionMode::ReadOnly; + params.supported_tools_override = Some(vec![api::ToolType::RequestComputerUse, api::ToolType::UseComputer]); + let (normal, cli) = super::prepare_direct_provider_params(&mut params); + assert!(normal.is_empty()); + assert!(cli.is_empty()); +} + #[test] fn supported_tools_omit_subagents_when_orchestration_is_disabled() { let params = request_params_with_ask_user_question_enabled(false); diff --git a/app/src/ai/agent/provider_history.rs b/app/src/ai/agent/provider_history.rs index 4371903e..3b91c41d 100644 --- a/app/src/ai/agent/provider_history.rs +++ b/app/src/ai/agent/provider_history.rs @@ -8,6 +8,81 @@ use crate::ai::runtime::ProviderActionContext; const TOOL_PREFIX: &str = "galaxy:provider-tool:v1:"; +pub(crate) fn computer_use_tool(action: &super::AIAgentActionType) -> Option { + use api::message::tool_call::{self as tool, use_computer::action as action_api}; + use computer_use::{Action, Key, MouseButton, ScrollDirection, ScrollDistance}; + let coordinates = |point: computer_use::Vector2I| Some(api::Coordinates { x: point.x(), y: point.y() }); + let button = |button: &MouseButton| match button { + MouseButton::Left => action_api::MouseButton::Left as i32, + MouseButton::Right => action_api::MouseButton::Right as i32, + MouseButton::Middle => action_api::MouseButton::Middle as i32, + MouseButton::Back => action_api::MouseButton::Back as i32, + MouseButton::Forward => action_api::MouseButton::Forward as i32, + }; + let key = |key: &Key| Some(action_api::Key { data: Some(match key { + Key::Char(ch) => action_api::key::Data::Char(ch.to_string()), + Key::Keycode(code) => action_api::key::Data::Keycode(*code), + }) }); + match action { + super::AIAgentActionType::RequestComputerUse(request) => Some(tool::Tool::RequestComputerUse(tool::RequestComputerUse { + task_summary: request.task_summary.clone(), + screenshot_params: request.screenshot_params.map(computer_screenshot), + })), + super::AIAgentActionType::UseComputer(request) => Some(tool::Tool::UseComputer(tool::UseComputer { + action_summary: request.action_summary.clone(), + post_actions_screenshot_params: request.screenshot_params.map(computer_screenshot), + actions: request.actions.iter().map(|targeted| tool::use_computer::Action { + target: computer_target(targeted.target), + r#type: Some(match &targeted.action { + Action::MouseMove { to } => action_api::Type::MouseMove(action_api::MouseMove { to: coordinates(*to) }), + Action::MouseDown { button: b, at } => action_api::Type::MouseDown(action_api::MouseDown { button: button(b), at: coordinates(*at) }), + Action::MouseUp { button: b } => action_api::Type::MouseUp(action_api::MouseUp { button: button(b) }), + Action::MouseWheel { at, direction, distance } => action_api::Type::MouseWheel(action_api::MouseWheel { + at: coordinates(*at), + direction: match direction { + ScrollDirection::Up => action_api::mouse_wheel::Direction::Up as i32, + ScrollDirection::Down => action_api::mouse_wheel::Direction::Down as i32, + ScrollDirection::Left => action_api::mouse_wheel::Direction::Left as i32, + ScrollDirection::Right => action_api::mouse_wheel::Direction::Right as i32, + }, + distance: Some(match distance { + ScrollDistance::Pixels(value) => action_api::mouse_wheel::Distance::Pixels(*value), + ScrollDistance::Clicks(value) => action_api::mouse_wheel::Distance::Clicks(*value), + }), + }), + Action::TypeText { text } => action_api::Type::TypeText(action_api::TypeText { text: text.clone() }), + Action::KeyDown { key: k } => action_api::Type::KeyDown(action_api::KeyDown { key: key(k) }), + Action::KeyUp { key: k } => action_api::Type::KeyUp(action_api::KeyUp { key: key(k) }), + Action::Wait(duration) => action_api::Type::Wait(action_api::Wait { duration: Some(prost_types::Duration { + seconds: duration.as_secs() as i64, nanos: duration.subsec_nanos() as i32, + }) }), + }), + }).collect(), + })), + _ => None, + } +} + +fn computer_target(target: computer_use::Target) -> Option { + use api::message::tool_call::{computer_use_target, ComputerUseTarget}; + match target { + computer_use::Target::Screen => None, + computer_use::Target::Window { window_id, pid } => Some(ComputerUseTarget { target: Some(computer_use_target::Target::Window(computer_use_target::Window { window_id: window_id.to_string(), pid })) }), + } +} + +fn computer_screenshot(params: computer_use::ScreenshotParams) -> api::message::tool_call::ScreenshotParams { + api::message::tool_call::ScreenshotParams { + max_long_edge_px: params.max_long_edge_px.unwrap_or_default() as i32, + max_total_px: params.max_total_px.unwrap_or_default() as i32, + target: computer_target(params.target), + region: params.region.map(|region| api::message::tool_call::screenshot_params::Region { + top_left: Some(api::Coordinates { x: region.top_left.x(), y: region.top_left.y() }), + bottom_right: Some(api::Coordinates { x: region.bottom_right.x(), y: region.bottom_right.y() }), + }), + } +} + #[derive(Serialize, Deserialize)] struct SavedTool { context: ProviderActionContext, diff --git a/app/src/ai/agent/provider_history_tests.rs b/app/src/ai/agent/provider_history_tests.rs index 0963bd74..bb146de2 100644 --- a/app/src/ai/agent/provider_history_tests.rs +++ b/app/src/ai/agent/provider_history_tests.rs @@ -24,6 +24,29 @@ fn query(id: &str) -> api::Message { } } +#[test] +fn computer_tool_display_and_native_restoration_preserve_the_same_actions() { + let context = ProviderActionContext::new_for_test("task"); + let call = ToolCall { + id: "desktop-call".into(), name: "use_computer".into(), arguments: json!({ + "action_summary": "Type and press Return", "actions": [ + { "type": "type_text", "text": "hello" }, + { "type": "key_down", "keycode": 36 }, + { "type": "key_up", "keycode": 36 }, + ], + }), + }; + let action = context.action_from_tool_call(&call).unwrap(); + let message = tool_message(&context, &call, &action).unwrap(); + let Some(api::message::Message::ToolCall(display)) = message.message else { panic!("expected tool display"); }; + let Some(api::message::tool_call::Tool::UseComputer(display)) = display.tool else { panic!("expected computer display"); }; + let display_action = crate::ai::agent::AIAgentActionType::try_from(display).unwrap(); + assert_eq!(display_action, action.action); + let restored = super::restore_tool(&message.server_message_data).unwrap().unwrap(); + assert_eq!(restored.action, action.action); + assert_eq!(restored.id, action.id); +} + #[test] fn saved_tools_and_results_restore_as_individual_panes_across_turns() { let context = ProviderActionContext::new_for_test("task"); diff --git a/app/src/ai/blocklist/action_model.rs b/app/src/ai/blocklist/action_model.rs index 55b0bf08..e79c9b5a 100644 --- a/app/src/ai/blocklist/action_model.rs +++ b/app/src/ai/blocklist/action_model.rs @@ -418,19 +418,72 @@ fn domain_tool_result(action_result: &AIAgentActionResult, permission_denied: bo } else { ToolResultStatus::Success }; - let content = action_result.result.model_content(); + let (content, images) = computer_use_model_content(&action_result.result) + .unwrap_or_else(|| (action_result.result.model_content(), Vec::new())); let content = if permission_denied { format!("Permission denied by the user. {content}") } else { content }; ToolResult { + images: if permission_denied { Vec::new() } else { images }, call_id: action_result.id.to_string(), content, status, } } +fn computer_use_model_content( + result: &AIAgentActionResultType, +) -> Option<(String, Vec)> { + let (screenshot, mut metadata) = match result { + AIAgentActionResultType::RequestComputerUse(RequestComputerUseResult::Approved { + screenshot, + platform, + windows, + }) => ( + Some(screenshot), + serde_json::json!({ + "status": "approved", + "platform": format!("{platform:?}"), + "windows": windows, + }), + ), + AIAgentActionResultType::UseComputer(UseComputerResult::Success(result)) => ( + result.screenshot.as_ref(), + serde_json::json!({ + "status": "success", + "cursor_position": result.cursor_position.map(|point| serde_json::json!({ + "x": point.x(), "y": point.y(), + })), + "windows": result.windows, + "captured_window": result.captured_window.map(|window| serde_json::json!({ + "window_id": window.window_id, + "width_px": window.width_px, + "height_px": window.height_px, + })), + }), + ), + _ => return None, + }; + let images = screenshot + .map(|screenshot| { + metadata["screenshot"] = serde_json::json!({ + "width": screenshot.width, + "height": screenshot.height, + "original_width": screenshot.original_width, + "original_height": screenshot.original_height, + "coordinate_space": "Use original physical pixels; scale image coordinates by original_width/width and original_height/height.", + }); + vec![galaxy_agent_core::ToolResultImage { + data: screenshot.data.clone(), + mime_type: screenshot.mime_type.to_string(), + }] + }) + .unwrap_or_default(); + Some((metadata.to_string(), images)) +} + #[cfg(not(target_family = "wasm"))] fn action_tool_name(action: &AIAgentAction) -> String { action diff --git a/app/src/ai/blocklist/action_model/execute/use_computer.rs b/app/src/ai/blocklist/action_model/execute/use_computer.rs index 66ed3223..ba5d085c 100644 --- a/app/src/ai/blocklist/action_model/execute/use_computer.rs +++ b/app/src/ai/blocklist/action_model/execute/use_computer.rs @@ -24,10 +24,8 @@ impl UseComputerExecutor { return false; }; - // We unconditionally return true here because this action is only executed by - // the computer use subagent, which cannot begin without the user approving it via - // a `RequestComputerUse` action, and the approval extends to all computer use - // actions within that computer use subagent. + // The shared action boundary applies the active permission mode to every action. + // Once that check passes, the executor has no additional approval requirement. true } diff --git a/app/src/ai/blocklist/action_model_tests.rs b/app/src/ai/blocklist/action_model_tests.rs index d8ce5253..38f29821 100644 --- a/app/src/ai/blocklist/action_model_tests.rs +++ b/app/src/ai/blocklist/action_model_tests.rs @@ -366,6 +366,26 @@ fn domain_read_result_contains_the_file_contents_for_the_next_model_turn() { assert!(result.content.contains("pub fn answer() -> u8 { 42 }")); } +#[test] +fn computer_results_keep_real_image_bytes_and_coordinate_metadata() { + let screenshot = computer_use::Screenshot { + width: 500, height: 300, original_width: 1000, original_height: 600, + data: vec![1, 2, 3, 4], mime_type: "image/png".into(), + }; + let action = action_result("desktop", AIAgentActionResultType::RequestComputerUse(RequestComputerUseResult::Approved { + screenshot, platform: computer_use::Platform::Mac, windows: Vec::new(), + })); + let result = domain_tool_result(&action, false); + assert_eq!(result.images.len(), 1); + assert_eq!(result.images[0].data, vec![1, 2, 3, 4]); + assert_eq!(result.images[0].mime_type, "image/png"); + let metadata: serde_json::Value = serde_json::from_str(&result.content).unwrap(); + assert_eq!(metadata["platform"], "Mac"); + assert_eq!(metadata["screenshot"]["width"], 500); + assert_eq!(metadata["screenshot"]["original_width"], 1000); + assert!(domain_tool_result(&action, true).images.is_empty()); +} + #[test] fn action_permission_kinds_match_the_safety_boundary() { assert_eq!( diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 985e0fe3..46d948b9 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -1039,6 +1039,7 @@ fn convert_provider_tool_batch( matches!(pending.state, PendingToolCallState::RecoveryPending), )), Err(message) => invalid_results.push(ToolResult { + images: Vec::new(), call_id: pending.call.id.clone(), content: format!("Invalid {} tool input: {message}", pending.call.name), status: ToolResultStatus::Error, @@ -8460,7 +8461,10 @@ impl BlocklistAIController { matches!( pending.call.name.as_str(), "request_computer_use" | "use_computer" - ) + ) && pending.state.result().or_else(|| match event { + galaxy_agent_core::ToolEvent::Completed { result } if result.call_id == pending.call.id => Some(result), + _ => None, + }).is_some_and(|result| result.status == galaxy_agent_core::ToolResultStatus::Success) }); used_computer.then_some(COMPUTER_USE_PROVIDER_PROFILE) } diff --git a/app/src/ai/llms.rs b/app/src/ai/llms.rs index 5fb6c1a6..45ecb63a 100644 --- a/app/src/ai/llms.rs +++ b/app/src/ai/llms.rs @@ -2819,9 +2819,12 @@ pub(crate) fn openai_model_context_size( .max_input_tokens .unwrap_or(model.context_size) .min(model.context_size); - if provider_kind == OpenAIProviderKind::ChatGPTSubscription { - // The catalog's default window is the supported session budget. The - // advertised maximum must not silently opt subscriptions into long context. + if matches!( + provider_kind, + OpenAIProviderKind::ChatGPTSubscription | OpenAIProviderKind::LiteLLM + ) { + // Use the efficient session window for both subscriptions and proxies. + // An advertised maximum does not opt a connection into long context. // Older saved catalogs kept the default in max_input_tokens instead. return model .default_context_size @@ -2959,7 +2962,7 @@ fn openai_model_context_window( is_configurable: max_context_size > default_context_size, min: default_context_size, max: max_context_size, - // Subscription limits are already capped to the catalog's default above. + // Subscription and LiteLLM limits are capped to the catalog default above. default_max: max_context_size, } } @@ -2971,6 +2974,53 @@ fn u32_from_json_any(value: &serde_json::Value, keys: &[&str]) -> Option { .and_then(|value| u32::try_from(value).ok()) } +/// Keep an endpoint's recommended window separate from its optional expanded +/// maximum. A bare model ID carries no token metadata, so it must not replace +/// richer limits when the two catalog endpoints are combined. +#[cfg(not(target_family = "wasm"))] +fn openai_context_limits(metadata: &serde_json::Value) -> (u32, Option, Option) { + let positive_limit = |keys: &[&str]| { + keys.iter().find_map(|key| { + metadata[*key] + .as_u64() + .and_then(|value| u32::try_from(value).ok()) + .filter(|value| *value > 0) + }) + }; + let max_input_tokens = + positive_limit(&["max_input_tokens", "input_token_limit", "max_prompt_tokens"]); + let default_context_size = positive_limit(&[ + "default_context_size", + "default_context_window", + "context_window", + ]); + let maximum_context_size = positive_limit(&[ + "max_context_window", + "max_model_len", + "context_size", + "token_size", + ]); + let effective_context_percent = metadata["effective_context_window_percent"] + .as_u64() + .filter(|value| (1..=100).contains(value)) + .unwrap_or(100); + // A proxy may already reduce max_input_tokens by this percentage. Apply it + // only to the original session window, never to that input limit again. + let default_context_size = default_context_size + .map(|value| { + u32::try_from(u64::from(value) * effective_context_percent / 100) + .unwrap_or(value) + .max(1) + }) + .or(max_input_tokens) + .or(maximum_context_size); + let context_size = maximum_context_size + .or(default_context_size) + .or(max_input_tokens) + .unwrap_or(DEFAULT_DISCOVERED_MODEL_CONTEXT_SIZE); + (context_size, max_input_tokens, default_context_size) +} + #[cfg(not(target_family = "wasm"))] fn normalize_codex_release_version(version: &str) -> Option { let version = version.trim(); @@ -3209,19 +3259,11 @@ async fn fetch_from_litellm_model_info( let model_name = entry["model_name"].as_str()?; let model_info = &entry["model_info"]; - let max_input_tokens = model_info["max_input_tokens"] - .as_u64() - .and_then(|v| u32::try_from(v).ok()); + let (context_size, max_input_tokens, default_context_size) = + openai_context_limits(model_info); let max_output_tokens = model_info["max_output_tokens"] .as_u64() .and_then(|v| u32::try_from(v).ok()); - let context_size = u32_from_json_any( - model_info, - &["max_context_window", "context_window", "max_model_len"], - ) - .or(max_input_tokens) - .unwrap_or(200_000); - let vision_supported = model_info["supports_vision"].as_bool().unwrap_or(true); let display_name = model_name.replace(['-', '_'], " "); @@ -3271,7 +3313,7 @@ async fn fetch_from_litellm_model_info( vision_supported, context_size, max_input_tokens, - default_context_size: None, + default_context_size, max_output_tokens, provider, supports_system_messages: if model_name.starts_with("codex-gpt-") { @@ -3318,7 +3360,30 @@ async fn fetch_litellm_catalog( metadata .iter() .find(|info| info.model_id == model.model_id) - .cloned() + .map(|info| { + let mut enriched = info.clone(); + // /model/info supplies capabilities, but must not erase a + // smaller session budget supplied by /models (or vice versa). + enriched.default_context_size = model + .default_context_size + .into_iter() + .chain(info.default_context_size) + .min(); + enriched.max_input_tokens = model + .max_input_tokens + .into_iter() + .chain(info.max_input_tokens) + .min(); + enriched.context_size = match ( + model.default_context_size, + info.default_context_size, + ) { + (Some(_), Some(_)) => model.context_size.max(info.context_size), + (Some(_), None) => model.context_size, + _ => info.context_size, + }; + enriched + }) .unwrap_or(model) }) .collect()) @@ -3369,21 +3434,7 @@ async fn fetch_from_openai_models( .iter() .filter_map(|m| { let id = m["id"].as_str()?; - let max_input_tokens = u32_from_any( - m, - &["max_input_tokens", "input_token_limit", "max_prompt_tokens"], - ); - let context_size = u32_from_any( - m, - &[ - "max_context_window", - "max_model_len", - "context_window", - "token_size", - ], - ) - .or(max_input_tokens) - .unwrap_or(200_000); + let (context_size, max_input_tokens, default_context_size) = openai_context_limits(m); let max_output_tokens = u32_from_any( m, &[ @@ -3430,7 +3481,7 @@ async fn fetch_from_openai_models( .unwrap_or(true), context_size, max_input_tokens, - default_context_size: None, + default_context_size, max_output_tokens, provider, supports_system_messages: if id.starts_with("codex-gpt-") { diff --git a/app/src/ai/llms_tests.rs b/app/src/ai/llms_tests.rs index c0a83ce8..e8488faf 100644 --- a/app/src/ai/llms_tests.rs +++ b/app/src/ai/llms_tests.rs @@ -1,3 +1,4 @@ +use serde_json::json; use warpui::App; use super::*; @@ -1189,6 +1190,155 @@ fn chatgpt_default_window_limits_routing_without_a_separate_input_limit() { ); } +#[test] +fn litellm_efficient_windows_limit_model_picker_and_provider_routing() { + App::test((), |mut app| async move { + initialize_settings_for_tests(&mut app); + let models = [("terminal-model", 272_000), ("main-model", 500_000)] + .into_iter() + .map(|(id, efficient_window)| { + let mut model = openai_model(id); + model.context_size = 872_000; + model.max_input_tokens = Some(872_000); + model.default_context_size = Some(efficient_window); + model + }) + .collect(); + AISettings::handle(&app).update(&mut app, |settings, ctx| { + settings.openai_enabled.set_value(true, ctx).unwrap(); + settings + .openai_providers + .set_value( + vec![OpenAIProviderConfig { + kind: OpenAIProviderKind::LiteLLM, + enabled: true, + name: "LiteLLM".to_owned(), + base_url: "https://proxy.test/v1".to_owned(), + api_key: None, + project_id: None, + location: None, + models, + }], + ctx, + ) + .unwrap(); + }); + let mut preferences = empty_preferences(); + app.read(|ctx| preferences.inject_openai_models(ctx)); + for (id, efficient_window) in [("terminal-model", 272_000), ("main-model", 500_000)] { + let route = preferences.openai_client_config_for_model(id).unwrap(); + assert_eq!(route.max_input_tokens, Some(efficient_window)); + let model = preferences + .models_by_feature + .agent_mode + .choices + .iter() + .find(|model| model.id.as_str() == id) + .unwrap(); + assert_eq!(model.context_window.max, efficient_window); + assert_eq!(model.context_window.default_max, efficient_window); + assert!(!model.context_window.is_configurable); + } + }); +} + +#[tokio::test] +async fn litellm_discovery_preserves_efficient_windows_across_catalog_endpoints() { + let mut server = mockito::Server::new_async().await; + for (model_metadata, rich_metadata, maximum, efficient_window, request_budget) in [ + ( + json!({"id": "model"}), + Some(json!({ + "context_window": 500000, + "max_context_window": 872000, + "max_input_tokens": 872000 + })), + 872_000, + 500_000, + 500_000, + ), + ( + json!({"id": "model", "context_window": 272000, "max_context_window": 872000}), + Some(json!({"max_context_window": 872000})), + 872_000, + 272_000, + 272_000, + ), + ( + json!({"id": "model", "context_window": 272000, "max_context_window": 872000}), + None, + 872_000, + 272_000, + 272_000, + ), + ( + json!({"id": "model", "context_window": 500000, "max_context_window": 872000}), + Some(json!({"supports_vision": true})), + 872_000, + 500_000, + 500_000, + ), + ( + json!({"id": "model"}), + Some(json!({ + "context_window": 272000, + "max_context_window": 872000, + "max_input_tokens": 258400, + "effective_context_window_percent": 95 + })), + 872_000, + 258_400, + 258_400, + ), + ( + json!({"id": "model"}), + Some(json!({ + "context_window": 272000, + "max_context_window": 872000, + "max_input_tokens": 128000 + })), + 872_000, + 272_000, + 128_000, + ), + ] { + let catalog = server + .mock("GET", "/models") + .with_status(200) + .with_header("content-type", "application/json") + .with_body(json!({"data": [model_metadata]}).to_string()) + .create_async() + .await; + let metadata = server + .mock("GET", "/model/info") + .with_status(if rich_metadata.is_some() { 200 } else { 404 }) + .with_header("content-type", "application/json") + .with_body( + json!({"data": [{"model_name": "model", "model_info": rich_metadata}]}).to_string(), + ) + .create_async() + .await; + let models = fetch_litellm_catalog(&server.url(), None, &reqwest::Client::new()) + .await + .unwrap(); + assert_eq!(models.len(), 1); + assert_eq!(models[0].context_size, maximum); + assert_eq!(models[0].default_context_size, Some(efficient_window)); + assert_eq!( + openai_model_context_size(&models[0], OpenAIProviderKind::LiteLLM), + request_budget + ); + let window = openai_model_context_window(&models[0], OpenAIProviderKind::LiteLLM); + assert_eq!(window.default_max, request_budget); + assert_eq!(window.max, request_budget); + assert!(!window.is_configurable); + catalog.assert_async().await; + metadata.assert_async().await; + catalog.remove_async().await; + metadata.remove_async().await; + } +} + #[test] fn litellm_refresh_replaces_membership_and_preserves_preferences() { let mut retained = openai_model("retained"); diff --git a/app/src/ai/openai/response_translator.rs b/app/src/ai/openai/response_translator.rs index 18f491d4..ab6b50ef 100644 --- a/app/src/ai/openai/response_translator.rs +++ b/app/src/ai/openai/response_translator.rs @@ -663,6 +663,8 @@ fn estimate_cost_cents( } const KNOWN_TOOLS: &[&str] = &[ + "request_computer_use", + "use_computer", "run_shell_command", "read_files", "apply_file_diffs", diff --git a/app/src/ai/provider/request_translator.rs b/app/src/ai/provider/request_translator.rs index cded1d37..6f830adc 100644 --- a/app/src/ai/provider/request_translator.rs +++ b/app/src/ai/provider/request_translator.rs @@ -1714,6 +1714,8 @@ pub(crate) fn tool_name_is_supported(name: &str, supported: &HashSet has(ToolType::AskUserQuestion), "read_skill" => has(ToolType::ReadSkill), "fetch_conversation" => has(ToolType::FetchConversation), + "request_computer_use" => has(ToolType::RequestComputerUse), + "use_computer" => has(ToolType::UseComputer), // This tool is implemented entirely inside the direct-provider response // translator, so it does not need a client ToolType capability bit. "recall_tool_history" => true, @@ -1730,8 +1732,77 @@ pub(crate) const IMPLEMENTATION_PROGRESS_GUIDANCE: &str = "## Incremental Implem pub(crate) const SHELL_EXECUTION_GUIDANCE: &str = "## Shell Execution\nCommands use a persistent terminal. Never exit, log out of, replace with exec, or kill that shell or Galaxy. Do not append `exit $status` to a command to report its result; let the command's exit status be returned naturally. If a script needs exit, run it in an explicit child shell such as `sh -c '...'`. Stop a running agent command with interrupt_shell_command, not broad kill/pkill commands.\nRun builds, tests, installs, and workspace checks as separate shell tool calls so each has its own output and exit status. Keep stdout and stderr live: do not hide them in subprocess capture_output/PIPE, redirect everything to a file, or buffer a whole validation suite and print only its final tail. If a log file is needed, stream through tee while preserving the original command's exit status. Use non-interactive, non-watch modes and disable pagers. Give intentionally quiet commands a descriptive tool title so the user knows what is running; silence alone is not evidence of a hang. If repeated snapshots show no progress, investigate whether the command is waiting for input, watching, or buffering; do not blindly poll forever or start duplicate commands. Interrupt and correct an unintended watcher or blocked command before retrying.\n\n"; +fn computer_target_schema() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "type": { "type": "string", "enum": ["screen", "window"], "default": "screen" }, + "window_id": { "type": "integer", "minimum": 1 }, + "pid": { "type": "integer", "minimum": 1 }, + }, + "required": ["type"], + "description": "Omit to target the main screen. Use window only with a window_id and pid from the latest returned windows list; requires background computer use support.", + }) +} + +fn computer_screenshot_schema() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "max_long_edge_px": { "type": "integer", "minimum": 1, "maximum": 4096, "default": 1568 }, + "max_total_px": { "type": "integer", "minimum": 1, "maximum": 4194304, "default": 1150000 }, + "target": computer_target_schema(), + }, + "description": "Screenshot size limits and target. Omit for a main-screen screenshot with bounded image size.", + }) +} + pub fn default_tool_definitions() -> Vec { vec![ + ToolDefinition { + name: "request_computer_use".to_string(), + description: "Request desktop computer use on the machine running Galaxy and capture its initial screenshot. This uses the user's local desktop even when the terminal is connected by SSH. Call before interacting with the desktop; subsequent turns use the configured computer-use model and use_computer tool. Existing profile permissions apply. Screen Recording permission is needed to capture the screen; Accessibility permission is needed to control input on macOS.".to_string(), + input_schema: serde_json::json!({ + "type": "object", + "properties": { + "task_summary": { "type": "string", "description": "Specific desktop task and application to operate." }, + "screenshot_params": computer_screenshot_schema(), + }, + "required": ["task_summary"], + }), + }, + ToolDefinition { + name: "use_computer".to_string(), + description: "Perform a bounded batch of desktop actions after request_computer_use succeeds, then receive a screenshot. Use an empty actions array to inspect the desktop. Coordinates use original physical screenshot pixels, so scale coordinates from the displayed image using the returned original_width/width and original_height/height. Screen targets use main-display coordinates; window targets use coordinates relative to the exact enumerated window. Prefer small action batches and inspect the result. Named keys use OS keycodes; key accepts a single character. Button/key down and up are separate actions.".to_string(), + input_schema: serde_json::json!({ + "type": "object", + "properties": { + "action_summary": { "type": "string" }, + "actions": { + "type": "array", "maxItems": 20, + "items": { + "type": "object", + "properties": { + "type": { "type": "string", "enum": ["wait", "mouse_move", "mouse_down", "mouse_up", "mouse_wheel", "type_text", "key_down", "key_up"] }, + "x": { "type": "integer" }, "y": { "type": "integer" }, + "button": { "type": "string", "enum": ["left", "right", "middle", "back", "forward"] }, + "text": { "type": "string", "maxLength": 100000 }, + "key": { "type": "string", "minLength": 1, "maxLength": 1 }, + "keycode": { "type": "integer", "description": "Platform keycode for modifier/navigation keys, such as macOS Return 36, Escape 53, Command 55, Shift 56, Option 58, Control 59." }, + "seconds": { "type": "number", "minimum": 0, "maximum": 10 }, + "direction": { "type": "string", "enum": ["up", "down", "left", "right"] }, + "pixels": { "type": "integer", "minimum": 0 }, + "target": computer_target_schema(), + }, + "required": ["type"], + }, + }, + "take_screenshot": { "type": "boolean", "default": true }, + "screenshot_params": computer_screenshot_schema(), + }, + "required": ["action_summary", "actions"], + }), + }, ToolDefinition { name: "run_shell_command".to_string(), description: "Execute a shell command in the user's terminal and return its output. Use for builds, tests, git, package managers, and other shell work. Commands run in this conversation's persistent shell and environment, starting in the supplied working directory. Successful cd commands affect subsequent calls in this shell; subshell directory changes do not. Avoid prefixing commands with cd to the directory you are already in, and use pwd only when the current location is uncertain. Never exit, log out of, replace, or kill this persistent shell. Run each build/test/install separately and keep its output live; avoid buffered subprocess wrappers and redirecting all output to files. Set wait_until_complete=false for commands that might run longer than a few seconds; Galaxy will return a command_id that can be monitored with read_shell_command_output.".to_string(), diff --git a/app/src/ai/provider/response_translator.rs b/app/src/ai/provider/response_translator.rs index fd470747..df6b11cd 100644 --- a/app/src/ai/provider/response_translator.rs +++ b/app/src/ai/provider/response_translator.rs @@ -839,6 +839,14 @@ pub fn build_tool_call_message( serde_json::from_str(tool_input_json).unwrap_or(serde_json::json!({})); let tool = match tool_name { + "request_computer_use" | "use_computer" => { + crate::ai::runtime::action_from_tool_call( + task_id, + &galaxy_agent_core::ToolCall { id: tool_use_id.to_string(), name: tool_name.to_string(), arguments: input.clone() }, + &ai::skills::SkillPathOrigin::Local, + &std::collections::HashMap::new(), + ).ok().and_then(|action| crate::ai::agent::provider_history::computer_use_tool(&action.action)) + } "run_shell_command" => { let command = input .get("command") @@ -1596,6 +1604,8 @@ fn run_agents_harness_from_str(harness_type: &str) -> Option { /// Built-in tools that Galaxy knows how to execute directly. const KNOWN_TOOLS: &[&str] = &[ + "request_computer_use", + "use_computer", "run_shell_command", "read_files", "apply_file_diffs", diff --git a/app/src/ai/runtime/event_translator_tests.rs b/app/src/ai/runtime/event_translator_tests.rs index d36af971..bba27bfd 100644 --- a/app/src/ai/runtime/event_translator_tests.rs +++ b/app/src/ai/runtime/event_translator_tests.rs @@ -104,6 +104,7 @@ fn restored_task_progress_replay_does_not_duplicate_the_persisted_checklist() { arguments: serde_json::to_value(&snapshot).unwrap(), }, state: galaxy_agent_core::PendingToolCallState::Resolved { result: galaxy_agent_core::ToolResult { + images: Vec::new(), call_id: "progress".to_owned(), status: galaxy_agent_core::ToolResultStatus::Success, content: snapshot.to_tool_result_content(), }}, diff --git a/app/src/ai/runtime/mod.rs b/app/src/ai/runtime/mod.rs index 39c64c12..7abefc26 100644 --- a/app/src/ai/runtime/mod.rs +++ b/app/src/ai/runtime/mod.rs @@ -4,6 +4,7 @@ mod provider_run_coordinator; mod rig; mod rig_request; mod rig_tool; +pub(crate) use rig_tool::action_from_tool_call; mod task_progress; pub(crate) use child_context::add_parent_context; diff --git a/app/src/ai/runtime/provider_run_coordinator_tests.rs b/app/src/ai/runtime/provider_run_coordinator_tests.rs index a4059281..b535d26e 100644 --- a/app/src/ai/runtime/provider_run_coordinator_tests.rs +++ b/app/src/ai/runtime/provider_run_coordinator_tests.rs @@ -461,6 +461,7 @@ async fn coordinator_awaiting_tools( fn tool_result(call_id: &str, content: &str, status: ToolResultStatus) -> ToolEvent { ToolEvent::Completed { result: ToolResult { + images: Vec::new(), call_id: call_id.to_string(), content: content.to_string(), status, @@ -514,6 +515,7 @@ async fn one_run_drives_model_tool_and_followup_turns_with_atomic_history() { .complete_tool( &batch.work_id, ToolResult { + images: Vec::new(), call_id: "read-1".to_string(), content: "manifest contents".to_string(), status: ToolResultStatus::Success, diff --git a/app/src/ai/runtime/rig.rs b/app/src/ai/runtime/rig.rs index c4593a09..94e8876f 100644 --- a/app/src/ai/runtime/rig.rs +++ b/app/src/ai/runtime/rig.rs @@ -133,20 +133,25 @@ pub(crate) async fn prepare_provider_run( let skill_path_origin = params.session_context.skill_path_origin(); let mut cli_params = params.clone(); let computer_use_model = params.computer_use_model.as_str(); - let computer_use_profile = if params.computer_use_enabled - && !computer_use_model.is_empty() - && computer_use_model != "computer-use-agent-auto" - && computer_use_model != params.model.as_str() - && !matches!( - &computer_use_provider_config, - crate::ai::provider::ProviderConfig::None - ) { + let computer_use_profile = if supported_tools.contains(&ToolType::RequestComputerUse) { let mut computer_params = params.clone(); - computer_params.model = params.computer_use_model.clone(); + let computer_provider = if computer_use_model.is_empty() + || computer_use_model == "computer-use-agent-auto" + || computer_use_model == params.model.as_str() + || matches!(&computer_use_provider_config, crate::ai::provider::ProviderConfig::None) + { + base_provider_config.clone() + } else { + computer_params.model = params.computer_use_model.clone(); + computer_use_provider_config + }; + let mut computer_tools = supported_tools.clone(); + computer_tools.retain(|tool| *tool != ToolType::RequestComputerUse); + computer_tools.push(ToolType::UseComputer); let (runtime, prepared) = prepare_provider_profile( - computer_use_provider_config, + computer_provider, computer_params, - supported_tools.clone(), + computer_tools, supported_cli_agent_tools.clone(), None, ) diff --git a/app/src/ai/runtime/rig_request.rs b/app/src/ai/runtime/rig_request.rs index 766e0e07..8650961a 100644 --- a/app/src/ai/runtime/rig_request.rs +++ b/app/src/ai/runtime/rig_request.rs @@ -389,17 +389,7 @@ fn input_messages( ) -> Vec { let mut messages = Vec::new(); if !tool_results.is_empty() { - let mut parts = tool_results - .into_iter() - .map(|result| { - let is_error = result.is_error(); - ContentPart::ToolResult { - tool_use_id: result.call_id, - content: result.content, - is_error, - } - }) - .collect::>(); + let mut parts = ToolResult::into_batch_content_parts(tool_results); let content = if parts.len() == 1 { let ContentPart::ToolResult { tool_use_id, @@ -1043,6 +1033,12 @@ fn build_system_prompt( { prompt.push_str(CODING_DELEGATION_GUIDANCE); } + if tools.iter().any(|tool| tool.name == "request_computer_use") { + prompt.push_str("## Desktop Computer Use\nUse request_computer_use for tasks that require seeing or interacting with the user's desktop. This controls the machine running Galaxy, including when the terminal is connected to another host. Request the specific authorized task and wait for the screenshot. A successful request switches to the configured computer-use model; rejected or failed requests do not grant control.\n\n"); + } + if tools.iter().any(|tool| tool.name == "use_computer") { + prompt.push_str("## Desktop Computer Use\nInspect the latest returned screenshot before acting. Use use_computer with small action batches, then inspect the new screenshot. Use original physical pixel coordinates: multiply displayed image x by original_width/width and y by original_height/height. Every button/key press requires both down and up actions. For window targeting, select an exact window_id/pid from the latest windows list; never invent a window. Stop on errors or denied permission and report the actual cause. Galaxy's Full Disk Access is separate from Screen Recording and Accessibility permissions. When desktop work is complete and terminal/file work remains, call the appropriate normal tool and Galaxy resumes the base model. Only finish the response when the whole user's task is complete.\n\n"); + } prompt.push_str("## Available Tools\n"); if tools.is_empty() { prompt.push_str("No tools are available. Do not invent tool calls.\n"); diff --git a/app/src/ai/runtime/rig_request_tests.rs b/app/src/ai/runtime/rig_request_tests.rs index cce253c0..f14f8dad 100644 --- a/app/src/ai/runtime/rig_request_tests.rs +++ b/app/src/ai/runtime/rig_request_tests.rs @@ -533,6 +533,7 @@ fn lrc_snapshot_follow_up_uses_the_cli_monitor_prompt_and_tools() { ), }; let snapshot_tool_result = ToolResult { + images: Vec::new(), call_id: "run-call".to_string(), content: result.result.model_content(), status: ToolResultStatus::Success, @@ -842,6 +843,7 @@ fn normalized_tool_outcomes_are_the_only_action_results_sent_to_rig() { let tool_results = statuses .iter() .map(|(call_id, status, _)| ToolResult { + images: Vec::new(), call_id: (*call_id).to_string(), content: format!("normalized-{call_id}"), status: *status, diff --git a/app/src/ai/runtime/rig_tests.rs b/app/src/ai/runtime/rig_tests.rs index e685ff1c..b157252a 100644 --- a/app/src/ai/runtime/rig_tests.rs +++ b/app/src/ai/runtime/rig_tests.rs @@ -87,6 +87,27 @@ async fn selected_computer_use_model_gets_its_own_profile() { ); } +#[tokio::test] +async fn automatic_and_same_model_computer_use_keep_a_dedicated_profile() { + for selected in ["computer-use-agent-auto", "base-selection", "unroutable-model"] { + let mut params = RequestParams::new_for_test(); + params.model = LLMId::from("base-selection"); + params.computer_use_model = LLMId::from(selected); + params.computer_use_enabled = true; + let prepared = prepare_provider_run( + ProviderConfig::OpenAI(openai_config("base-provider-model")), + ProviderConfig::None, ProviderConfig::None, params, Vec::new(), + ).await.unwrap(); + assert!(prepared.base_profile.request.tools.iter().any(|tool| tool.name == "request_computer_use")); + assert!(!prepared.base_profile.request.tools.iter().any(|tool| tool.name == "use_computer")); + let computer = prepared.computer_use_profile.unwrap(); + assert_eq!(computer.request.model.as_str(), "base-provider-model"); + assert!(computer.request.tools.iter().any(|tool| tool.name == "use_computer")); + assert!(!computer.request.tools.iter().any(|tool| tool.name == "request_computer_use")); + assert!(computer.request.system_prompt.unwrap().contains("Inspect the latest returned screenshot before acting")); + } +} + #[test] fn subscription_context_budget_uses_catalog_limit_and_clamps_stale_overrides() { let mut config = openai_config("subscription-model"); diff --git a/app/src/ai/runtime/rig_tool.rs b/app/src/ai/runtime/rig_tool.rs index 50d50be4..f87ba6dc 100644 --- a/app/src/ai/runtime/rig_tool.rs +++ b/app/src/ai/runtime/rig_tool.rs @@ -55,6 +55,23 @@ pub(crate) fn action_from_tool_call( rationale: optional_string(input, "rationale")?, citations: Vec::new(), }, + "request_computer_use" => AIAgentActionType::RequestComputerUse( + crate::ai::agent::RequestComputerUseRequest { + task_summary: required_nonempty_string(input, "task_summary")?, + screenshot_params: Some(computer_screenshot_params(input)?), + }, + ), + "use_computer" => AIAgentActionType::UseComputer( + crate::ai::agent::UseComputerRequest { + action_summary: required_nonempty_string(input, "action_summary")?, + actions: computer_actions(input)?, + screenshot_params: if optional_boolean(input, "take_screenshot")?.unwrap_or(true) { + Some(computer_screenshot_params(input)?) + } else { + None + }, + }, + ), "read_files" => AIAgentActionType::ReadFiles(ReadFilesRequest { locations: limited_required_array(input, "files", MAX_BATCH_READ_ITEMS)? .iter() @@ -472,6 +489,113 @@ fn optional_nonnegative_i32(input: &serde_json::Value, key: &str) -> Result Result { + input.get(key).and_then(serde_json::Value::as_i64) + .and_then(|value| i32::try_from(value).ok()) + .ok_or_else(|| format!("invalid field {key:?}: expected a 32-bit integer")) +} + +fn computer_target(input: &serde_json::Value) -> Result { + let Some(target) = input.get("target") else { return Ok(computer_use::Target::Screen); }; + require_object(target, "target")?; + match required_string(target, "type")?.as_str() { + "screen" => Ok(computer_use::Target::Screen), + "window" => { + let window_id = optional_bounded_u64(target, "window_id", u64::from(u32::MAX))? + .filter(|id| *id > 0).ok_or("window target requires a non-zero window_id")? as u32; + let pid = computer_i32(target, "pid")?; + if pid <= 0 { return Err("window target requires a positive pid".to_string()); } + if !crate::features::FeatureFlag::BackgroundComputerUse.is_enabled() + || !computer_use::background_supported() + { + return Err("Background window targeting is unavailable; use the screen target.".to_string()); + } + Ok(computer_use::Target::Window { window_id, pid }) + } + _ => Err("target.type must be screen or window".to_string()), + } +} + +fn computer_screenshot_params(input: &serde_json::Value) -> Result { + let empty = serde_json::json!({}); + let params = input.get("screenshot_params").unwrap_or(&empty); + require_object(params, "screenshot_params")?; + let positive = |key, maximum, default| -> Result { + let value = optional_bounded_u64(params, key, maximum)?.unwrap_or(default); + if value == 0 { return Err(format!("{key} must be positive")); } + Ok(value as usize) + }; + Ok(computer_use::ScreenshotParams { + max_long_edge_px: Some(positive("max_long_edge_px", 4096, 1568)?), + max_total_px: Some(positive("max_total_px", 4_194_304, 1_150_000)?), + region: None, + target: computer_target(params)?, + }) +} + +fn computer_actions(input: &serde_json::Value) -> Result, String> { + use computer_use::{Action, Key, MouseButton, ScrollDirection, ScrollDistance, TargetedAction, Vector2I}; + let mut wait_seconds = 0.0; + limited_required_array(input, "actions", 20)?.iter().map(|input| { + require_object(input, "action")?; + let position = || -> Result { + Ok(Vector2I::new(computer_i32(input, "x")?, computer_i32(input, "y")?)) + }; + let button = || -> Result { + match required_string(input, "button")?.as_str() { + "left" => Ok(MouseButton::Left), "right" => Ok(MouseButton::Right), + "middle" => Ok(MouseButton::Middle), "back" => Ok(MouseButton::Back), + "forward" => Ok(MouseButton::Forward), + _ => Err("button must be left, right, middle, back, or forward".to_string()), + } + }; + let key = || -> Result { + match (input.get("key"), input.get("keycode")) { + (Some(_), None) => { + let text = required_string(input, "key")?; + let mut chars = text.chars(); + let key = chars.next().ok_or("key must contain exactly one character")?; + if chars.next().is_some() { return Err("key must contain exactly one character".to_string()); } + Ok(Key::Char(key)) + } + (None, Some(_)) => Ok(Key::Keycode(computer_i32(input, "keycode")?)), + _ => Err("provide exactly one of key or keycode".to_string()), + } + }; + let action = match required_string(input, "type")?.as_str() { + "mouse_move" => Action::MouseMove { to: position()? }, + "mouse_down" => Action::MouseDown { button: button()?, at: position()? }, + "mouse_up" => Action::MouseUp { button: button()? }, + "mouse_wheel" => Action::MouseWheel { + at: position()?, + direction: match required_string(input, "direction")?.as_str() { + "up" => ScrollDirection::Up, "down" => ScrollDirection::Down, + "left" => ScrollDirection::Left, "right" => ScrollDirection::Right, + _ => return Err("direction must be up, down, left, or right".to_string()), + }, + distance: ScrollDistance::Pixels(optional_nonnegative_i32(input, "pixels")?.ok_or("mouse_wheel requires pixels")?), + }, + "type_text" => { + let text = required_string(input, "text")?; + if text.len() > 100_000 { return Err("text must be no longer than 100000 bytes".to_string()); } + Action::TypeText { text } + } + "key_down" => Action::KeyDown { key: key()? }, + "key_up" => Action::KeyUp { key: key()? }, + "wait" => { + let seconds = input.get("seconds").and_then(serde_json::Value::as_f64) + .filter(|seconds| seconds.is_finite() && *seconds >= 0.0 && *seconds <= 10.0) + .ok_or("seconds must be between 0 and 10")?; + wait_seconds += seconds; + if wait_seconds > 10.0 { return Err("a computer-use batch may wait at most 10 seconds".to_string()); } + Action::Wait(Duration::from_secs_f64(seconds)) + } + _ => return Err("unknown computer action type".to_string()), + }; + Ok(TargetedAction { action, target: computer_target(input)? }) + }).collect() +} + fn skill_references( input: &serde_json::Value, skill_path_origin: &SkillPathOrigin, diff --git a/app/src/ai/runtime/rig_tool_tests.rs b/app/src/ai/runtime/rig_tool_tests.rs index 71a7d1fb..7b66189a 100644 --- a/app/src/ai/runtime/rig_tool_tests.rs +++ b/app/src/ai/runtime/rig_tool_tests.rs @@ -16,6 +16,55 @@ fn call(name: &str, arguments: serde_json::Value) -> ToolCall { } } +#[test] +fn computer_use_request_always_captures_a_bounded_initial_screenshot() { + let action = action_from_tool_call("task", &call("request_computer_use", serde_json::json!({ + "task_summary": "Inspect the browser", + })), &SkillPathOrigin::Local, &HashMap::new()).unwrap(); + let AIAgentActionType::RequestComputerUse(request) = action.action else { panic!("expected computer-use request"); }; + let screenshot = request.screenshot_params.unwrap(); + assert_eq!(screenshot.max_long_edge_px, Some(1568)); + assert_eq!(screenshot.max_total_px, Some(1_150_000)); + assert_eq!(screenshot.target, computer_use::Target::Screen); +} + +#[test] +fn computer_actions_preserve_coordinates_keys_and_screenshot_choices() { + let action = action_from_tool_call("task", &call("use_computer", serde_json::json!({ + "action_summary": "Select the field and type", + "actions": [ + { "type": "mouse_down", "button": "left", "x": 420, "y": 310 }, + { "type": "mouse_up", "button": "left" }, + { "type": "type_text", "text": "hello" }, + { "type": "key_down", "keycode": 36 }, + { "type": "key_up", "keycode": 36 }, + ], + "take_screenshot": false, + })), &SkillPathOrigin::Local, &HashMap::new()).unwrap(); + let AIAgentActionType::UseComputer(request) = action.action else { panic!("expected computer actions"); }; + assert_eq!(request.actions.len(), 5); + assert_eq!(request.actions[0].action, computer_use::Action::MouseDown { + button: computer_use::MouseButton::Left, at: computer_use::Vector2I::new(420, 310), + }); + assert_eq!(request.actions[3].action, computer_use::Action::KeyDown { key: computer_use::Key::Keycode(36) }); + assert_eq!(request.screenshot_params, None); +} + +#[test] +fn computer_actions_reject_invalid_targets_keys_and_unbounded_waiting() { + for actions in [ + serde_json::json!([{ "type": "key_down", "key": "Return" }]), + serde_json::json!([{ "type": "key_up", "key": "a", "keycode": 0 }]), + serde_json::json!([{ "type": "mouse_move", "x": 1, "y": 2, "target": { "type": "window", "window_id": 0, "pid": 12 } }]), + serde_json::json!([{ "type": "wait", "seconds": 6 }, { "type": "wait", "seconds": 6 }]), + serde_json::json!(vec![serde_json::json!({ "type": "wait", "seconds": 0 }); 21]), + ] { + assert!(action_from_tool_call("task", &call("use_computer", serde_json::json!({ + "action_summary": "Invalid request", "actions": actions, + })), &SkillPathOrigin::Local, &HashMap::new()).is_err()); + } +} + #[test] fn shell_calls_become_domain_actions_without_a_proto_round_trip() { let action = action_from_tool_call( diff --git a/app/src/ai/runtime/task_progress_tests.rs b/app/src/ai/runtime/task_progress_tests.rs index 0c794d6b..9b69c172 100644 --- a/app/src/ai/runtime/task_progress_tests.rs +++ b/app/src/ai/runtime/task_progress_tests.rs @@ -106,6 +106,7 @@ fn rejected_plan_creation_cannot_seed_visible_or_durable_tasks() { arguments: serde_json::json!({"documents": [{"content": "## Tasks\n- [ ] Inspect\n- [ ] Verify"}]}), }, state: PendingToolCallState::Resolved { result: galaxy_agent_core::ToolResult { + images: Vec::new(), call_id: "plan".to_owned(), content: "Not advertised".to_owned(), status: galaxy_agent_core::ToolResultStatus::Error, diff --git a/crates/galaxy_agent_core/src/provider_run.rs b/crates/galaxy_agent_core/src/provider_run.rs index 16a8a7dc..9c99f7db 100644 --- a/crates/galaxy_agent_core/src/provider_run.rs +++ b/crates/galaxy_agent_core/src/provider_run.rs @@ -4,6 +4,7 @@ use std::fmt; use serde::{Deserialize, Serialize}; +use crate::types::tool_image_group_call_id; use crate::{ AgentError, ContentPart, ConversationMessage, MessageContent, MessageRole, PermissionDecision, PermissionRequest, StopReason, TaskListSnapshot, ToolCall, ToolCallDecision, ToolPolicy, @@ -860,6 +861,7 @@ impl ProviderRun { .push(pending.call.id.clone()); pending.state = PendingToolCallState::Resolved { result: ToolResult { + images: Vec::new(), call_id: pending.call.id.clone(), content: "Tool execution was interrupted by application restart and was not replayed." .to_string(), @@ -1161,6 +1163,7 @@ impl ProviderRun { let content = reason.unwrap_or_else(|| "Tool permission was denied.".to_string()); call.state = PendingToolCallState::Resolved { result: ToolResult { + images: Vec::new(), call_id: call_id.to_string(), content, status: ToolResultStatus::Denied, @@ -1244,6 +1247,7 @@ impl ProviderRun { | PendingToolCallState::RecoveryPending => { call.state = PendingToolCallState::Resolved { result: ToolResult { + images: Vec::new(), call_id: call_id.to_string(), content: reason.into(), status: ToolResultStatus::Cancelled, @@ -1468,6 +1472,7 @@ impl ProviderRun { let state = if has_launch && !is_local_agent_launch(&call) { PendingToolCallState::Resolved { result: ToolResult { + images: Vec::new(), call_id: call.id.clone(), content: "Not executed: launching a child group yields this turn. Wait for the group's completion before scheduling dependent work.".to_string(), status: ToolResultStatus::Error, @@ -1476,6 +1481,7 @@ impl ProviderRun { } else if !call.arguments.is_object() { PendingToolCallState::Resolved { result: ToolResult { + images: Vec::new(), call_id: call.id.clone(), content: format!( "Error: '{}' received malformed arguments; expected a JSON object.", @@ -1587,6 +1593,7 @@ impl ProviderRun { if call.state.result().is_none() { call.state = PendingToolCallState::Resolved { result: ToolResult { + images: Vec::new(), call_id: call.call.id.clone(), content: result_content.clone(), status, @@ -1782,16 +1789,7 @@ fn assistant_message(turn: &CompletedModelTurn) -> ConversationMessage { fn tool_result_message(results: Vec) -> ConversationMessage { ConversationMessage { role: MessageRole::User, - content: MessageContent::MultiPart( - results - .into_iter() - .map(|result| ContentPart::ToolResult { - tool_use_id: result.call_id, - content: result.content, - is_error: !matches!(result.status, ToolResultStatus::Success), - }) - .collect(), - ), + content: MessageContent::MultiPart(ToolResult::into_batch_content_parts(results)), } } @@ -1862,17 +1860,42 @@ pub fn archive_tool_results( seen_results.insert(tool_use_id.clone()) } MessageContent::MultiPart(parts) => { - parts.retain(|part| match part { - ContentPart::ToolUse { tool_use_id, .. } => { - seen_uses.insert(tool_use_id.clone()) + let mut retained_results = HashSet::new(); + let mut keep_parts: Vec<_> = parts + .iter() + .map(|part| match part { + ContentPart::ToolUse { tool_use_id, .. } => { + seen_uses.insert(tool_use_id.clone()) + } + ContentPart::ToolResult { tool_use_id, .. } => { + let keep = seen_results.insert(tool_use_id.clone()); + if keep { + retained_results.insert(tool_use_id.clone()); + } + keep + } + ContentPart::Text(_) + | ContentPart::Reasoning { .. } + | ContentPart::Image { .. } => false, + }) + .collect(); + let mut keep_image_group = false; + for (index, part) in parts.iter().enumerate() { + match part { + ContentPart::Text(_) => { + keep_image_group = tool_image_group_call_id(part) + .is_some_and(|call_id| retained_results.contains(call_id)) + && matches!(parts.get(index + 1), Some(ContentPart::Image { .. })); + keep_parts[index] = keep_image_group; + } + ContentPart::Image { .. } => keep_parts[index] = keep_image_group, + ContentPart::Reasoning { .. } + | ContentPart::ToolUse { .. } + | ContentPart::ToolResult { .. } => keep_image_group = false, } - ContentPart::ToolResult { tool_use_id, .. } => { - seen_results.insert(tool_use_id.clone()) - } - ContentPart::Text(_) - | ContentPart::Reasoning { .. } - | ContentPart::Image { .. } => false, - }); + } + let mut retained = keep_parts.into_iter(); + parts.retain(|_| retained.next().unwrap_or(false)); !parts.is_empty() } MessageContent::Text(_) => false, diff --git a/crates/galaxy_agent_core/src/provider_run_tests.rs b/crates/galaxy_agent_core/src/provider_run_tests.rs index da299550..b7b391ce 100644 --- a/crates/galaxy_agent_core/src/provider_run_tests.rs +++ b/crates/galaxy_agent_core/src/provider_run_tests.rs @@ -1,7 +1,7 @@ use serde_json::json; use super::*; -use crate::{AgentErrorKind, PermissionKind, TaskListItem, TaskListStatus}; +use crate::{AgentErrorKind, PermissionKind, TaskListItem, TaskListStatus, ToolResultImage}; fn initial_messages() -> Vec { vec![ConversationMessage { @@ -88,6 +88,7 @@ fn accept_tool_turn(run: &mut ProviderRun, turn: CompletedModelTurn) -> PendingT fn successful_result(call_id: &str, content: &str) -> ToolResult { ToolResult { + images: Vec::new(), call_id: call_id.to_string(), content: content.to_string(), status: ToolResultStatus::Success, @@ -1673,6 +1674,127 @@ fn multipart_archive_is_not_multiplied_and_legacy_duplicates_are_repaired() { assert_eq!(archive, expected); } +#[test] +fn screenshot_results_survive_checkpoint_compaction_and_archive_deduplication() { + let mut run = ProviderRun::new( + "screenshots", + vec![ConversationMessage { + role: MessageRole::User, + content: MessageContent::MultiPart(vec![ContentPart::Image { + data: vec![99], + mime_type: "image/png".to_string(), + }]), + }], + "base", + ProviderRunLimits::default(), + ); + let batch = accept_tool_turn( + &mut run, + tool_turn( + vec![ + tool_call("capture", "use_computer"), + tool_call("capture-more", "use_computer"), + ], + &["use_computer"], + ), + ); + for (call_id, bytes) in [("capture", vec![1, 2]), ("capture-more", vec![3, 4])] { + let mut result = successful_result(call_id, "screenshot captured"); + result.images.push(ToolResultImage { + data: bytes, + mime_type: "image/png".to_string(), + }); + run.complete_tool(&batch.work_id, result).unwrap(); + } + // Pending typed results and their image payloads must survive a checkpoint. + assert_serialization_round_trip(&run); + run = serde_json::from_str(&serde_json::to_string(&run).unwrap()).unwrap(); + run.commit_tool_batch(&batch.work_id).unwrap(); + let transcript = run.transcript().to_vec(); + let summary = vec![ConversationMessage { + role: MessageRole::User, + content: MessageContent::Text("Screenshots inspected.".to_string()), + }]; + run.compact_transcript_at_model_boundary(transcript.len(), summary.clone()) + .unwrap(); + assert_eq!(run.transcript(), summary); + assert_eq!(run.tool_result_archive().len(), 2); + let MessageContent::MultiPart(parts) = &run.tool_result_archive()[1].content else { + panic!("expected archived screenshot result batch"); + }; + assert_eq!(parts.len(), 6); + assert!(matches!( + &parts[2], ContentPart::Text(text) if text == "Images returned by tool call capture:" + )); + assert!(matches!( + &parts[3], ContentPart::Image { data, .. } if data == &[1, 2] + )); + assert!(matches!( + &parts[4], ContentPart::Text(text) if text == "Images returned by tool call capture-more:" + )); + assert!(matches!( + &parts[5], ContentPart::Image { data, .. } if data == &[3, 4] + )); + let expected = run.tool_result_archive().to_vec(); + let duplicated = expected.iter().cloned().chain(transcript).collect(); + run.replace_tool_result_archive(duplicated); + assert_eq!(run.tool_result_archive(), expected); + assert_serialization_round_trip(&run); +} + +#[test] +fn archives_only_images_labelled_for_new_results_in_the_same_message() { + let result = successful_result("capture", "captured"); + let mut parts = ToolResult::into_batch_content_parts(vec![result]); + parts.extend([ + ContentPart::Image { + data: vec![99], + mime_type: "image/png".to_string(), + }, + ContentPart::Text("Images returned by tool call missing:".to_string()), + ContentPart::Image { + data: vec![88], + mime_type: "image/png".to_string(), + }, + ContentPart::Text("Images returned by tool call capture:".to_string()), + ContentPart::Image { + data: vec![1], + mime_type: "image/png".to_string(), + }, + ContentPart::Text("Unrelated user upload".to_string()), + ContentPart::Image { + data: vec![77], + mime_type: "image/png".to_string(), + }, + ]); + let mut archive = Vec::new(); + let message = ConversationMessage { + role: MessageRole::User, + content: MessageContent::MultiPart(parts), + }; + archive_tool_results(&mut archive, vec![message.clone()]); + let MessageContent::MultiPart(parts) = &archive[0].content else { + panic!("expected correlated tool result with one image"); + }; + assert_eq!(parts.len(), 3); + assert!(matches!(&parts[2], ContentPart::Image { data, .. } if data == &[1])); + let expected = archive.clone(); + archive_tool_results(&mut archive, vec![message]); + assert_eq!(archive, expected); +} + +#[test] +fn legacy_tool_results_without_images_still_restore() { + let result: ToolResult = serde_json::from_value(json!({ + "call_id": "legacy", + "content": "old text result", + "status": "Success" + })) + .unwrap(); + assert!(result.images.is_empty()); + assert_eq!(result.content, "old text result"); +} + #[test] fn local_child_launch_yields_and_does_not_execute_batched_follow_up_tools() { let mut run = run(); @@ -1718,6 +1840,7 @@ fn failed_child_launch_allows_a_corrective_model_turn() { run.complete_tool( &batch.work_id, ToolResult { + images: Vec::new(), call_id: "children".to_string(), content: "Launch failed".to_string(), status: ToolResultStatus::Error, diff --git a/crates/galaxy_agent_core/src/tool_policy.rs b/crates/galaxy_agent_core/src/tool_policy.rs index acf9eec9..b97c9fc7 100644 --- a/crates/galaxy_agent_core/src/tool_policy.rs +++ b/crates/galaxy_agent_core/src/tool_policy.rs @@ -2,9 +2,10 @@ use std::collections::{BTreeSet, HashMap, VecDeque}; use serde_json::Value as JsonValue; +use crate::types::tool_image_group_call_id; use crate::{ ContentPart, ConversationMessage, MessageContent, TaskListSnapshot, ToolCall, ToolDefinition, - ToolResult, ToolResultStatus, UPDATE_TASK_LIST_NAME, + ToolResult, ToolResultImage, ToolResultStatus, UPDATE_TASK_LIST_NAME, }; pub const RECALL_TOOL_HISTORY_NAME: &str = "recall_tool_history"; @@ -124,6 +125,7 @@ impl ToolPolicy { ) }; return ToolCallDecision::Reject(ToolResult { + images: Vec::new(), call_id: call.id.clone(), content: format!( "Error: '{}' is not a valid tool for this request; {available}. Do not invent tool names.", @@ -134,13 +136,15 @@ impl ToolPolicy { } if call.name == RECALL_TOOL_HISTORY_NAME { + let (content, images) = recall_tool_history_with_images( + messages, + archive, + ToolHistoryQuery::from_arguments(&call.arguments), + ); return ToolCallDecision::Inline(ToolResult { + images, call_id: call.id.clone(), - content: recall_tool_history( - messages, - archive, - ToolHistoryQuery::from_arguments(&call.arguments), - ), + content, status: ToolResultStatus::Success, }); } @@ -148,11 +152,13 @@ impl ToolPolicy { if call.name == UPDATE_TASK_LIST_NAME { return match TaskListSnapshot::from_arguments(&call.arguments) { Ok(snapshot) => ToolCallDecision::Inline(ToolResult { + images: Vec::new(), call_id: call.id.clone(), content: snapshot.to_tool_result_content(), status: ToolResultStatus::Success, }), Err(error) => ToolCallDecision::Reject(ToolResult { + images: Vec::new(), call_id: call.id.clone(), content: format!("Error: {error}"), status: ToolResultStatus::Error, @@ -201,6 +207,14 @@ pub fn recall_tool_history( archive: &[ConversationMessage], query: ToolHistoryQuery<'_>, ) -> String { + recall_tool_history_with_images(messages, archive, query).0 +} + +fn recall_tool_history_with_images( + messages: &[ConversationMessage], + archive: &[ConversationMessage], + query: ToolHistoryQuery<'_>, +) -> (String, Vec) { let mut entries = Vec::new(); collect_tool_entries(archive, &mut entries); collect_tool_entries(messages, &mut entries); @@ -222,26 +236,43 @@ pub fn recall_tool_history( .len() .checked_sub(1usize.saturating_add(query.offset_from_end)) else { - return "No matching tool calls found in conversation history.".to_string(); + return ( + "No matching tool calls found in conversation history.".to_string(), + Vec::new(), + ); }; let entry = filtered[index]; let result = truncate_recalled_result(&entry.result); - format!( - "Tool: {}\nTool Use ID: {}\nInput: {}\nResult:\n{result}", - entry.name, entry.tool_use_id, entry.input + ( + format!( + "Tool: {}\nTool Use ID: {}\nInput: {}\nResult:\n{result}", + entry.name, entry.tool_use_id, entry.input + ), + entry + .images + .iter() + .map(|(data, mime_type)| ToolResultImage { + data: data.to_vec(), + mime_type: (*mime_type).to_string(), + }) + .collect(), ) } -#[derive(Clone, Debug, PartialEq, Eq)] -struct ToolHistoryEntry { +#[derive(Clone, Debug, PartialEq)] +struct ToolHistoryEntry<'a> { tool_use_id: String, name: String, input: String, result: String, + images: Vec<(&'a [u8], &'a str)>, } -fn collect_tool_entries(messages: &[ConversationMessage], entries: &mut Vec) { +fn collect_tool_entries<'a>( + messages: &'a [ConversationMessage], + entries: &mut Vec>, +) { let mut pending = Vec::new(); for message in messages { match &message.content { @@ -256,6 +287,7 @@ fn collect_tool_entries(messages: &[ConversationMessage], entries: &mut Vec pair_result(tool_use_id, content, &mut pending, entries), MessageContent::MultiPart(parts) => { + let entries_start = entries.len(); for part in parts { match part { ContentPart::ToolUse { @@ -275,6 +307,41 @@ fn collect_tool_entries(messages: &[ConversationMessage], entries: &mut Vec {} } } + // ToolResult::into_batch_content_parts places labelled image groups + // after every correlated result. Only attach images to results from + // this same message, never a nearby user upload or another call. + let result_ids: BTreeSet<_> = parts + .iter() + .filter_map(|part| match part { + ContentPart::ToolResult { tool_use_id, .. } => Some(tool_use_id.as_str()), + ContentPart::Text(_) + | ContentPart::Reasoning { .. } + | ContentPart::Image { .. } + | ContentPart::ToolUse { .. } => None, + }) + .collect(); + let mut image_call_id = None; + for part in parts { + match part { + ContentPart::Text(_) => { + image_call_id = tool_image_group_call_id(part) + .filter(|call_id| result_ids.contains(call_id)); + } + ContentPart::Image { data, mime_type } => { + if let Some(entry) = image_call_id.and_then(|call_id| { + entries[entries_start..] + .iter_mut() + .rev() + .find(|entry| entry.tool_use_id == call_id) + }) { + entry.images.push((data, mime_type)); + } + } + ContentPart::Reasoning { .. } + | ContentPart::ToolUse { .. } + | ContentPart::ToolResult { .. } => image_call_id = None, + } + } } MessageContent::Text(_) => {} } @@ -285,7 +352,7 @@ fn pair_result( tool_use_id: &str, content: &str, pending: &mut Vec<(String, String, String)>, - entries: &mut Vec, + entries: &mut Vec>, ) { let Some(index) = pending.iter().position(|(id, _, _)| id == tool_use_id) else { return; @@ -296,6 +363,7 @@ fn pair_result( name, input, result: content.to_string(), + images: Vec::new(), }); } diff --git a/crates/galaxy_agent_core/src/tool_policy_tests.rs b/crates/galaxy_agent_core/src/tool_policy_tests.rs index 5e876110..1758b90b 100644 --- a/crates/galaxy_agent_core/src/tool_policy_tests.rs +++ b/crates/galaxy_agent_core/src/tool_policy_tests.rs @@ -212,6 +212,120 @@ fn recall_never_returns_a_prior_recall_call() { ); } +fn screenshot_exchange(id: &str, bytes: Vec) -> Vec { + let mut messages = tool_exchange(id, "use_computer", serde_json::json!({}), "captured"); + messages[1].content = MessageContent::MultiPart(ToolResult::into_batch_content_parts(vec![ + ToolResult { + call_id: id.to_string(), + content: "captured".to_string(), + status: ToolResultStatus::Success, + images: vec![ToolResultImage { + data: bytes, + mime_type: "image/png".to_string(), + }], + }, + ])); + messages +} + +#[test] +fn recall_attaches_only_the_exact_selected_screenshots_from_archived_or_live_results() { + let archive = screenshot_exchange("capture", vec![1, 2]); + let mut messages = screenshot_exchange("capture-more", vec![3, 4]); + messages.push(ConversationMessage { + role: MessageRole::User, + content: MessageContent::MultiPart(vec![ + ContentPart::Text("Images returned by tool call capture:".to_string()), + ContentPart::Image { + data: vec![99], + mime_type: "image/png".to_string(), + }, + ]), + }); + let policy = ToolPolicy::new(&[definition(RECALL_TOOL_HISTORY_NAME)]); + for (arguments, original_id, bytes) in [ + ( + serde_json::json!({"tool_use_id": "capture"}), + "capture", + vec![1, 2], + ), + ( + serde_json::json!({"tool_name": "use_computer"}), + "capture-more", + vec![3, 4], + ), + ( + serde_json::json!({"tool_name": "use_computer", "offset_from_end": 1}), + "capture", + vec![1, 2], + ), + ] { + let call = ToolCall { + id: "recall-image".to_string(), + name: RECALL_TOOL_HISTORY_NAME.to_string(), + arguments, + }; + let ToolCallDecision::Inline(result) = policy.decide(&call, &messages, &archive) else { + panic!("recall should execute inline"); + }; + assert_eq!(result.call_id, call.id); + assert!(result + .content + .contains(&format!("Tool Use ID: {original_id}\n"))); + assert_eq!( + result.images, + vec![ToolResultImage { + data: bytes, + mime_type: "image/png".to_string(), + }] + ); + let parts = ToolResult::into_batch_content_parts(vec![result]); + assert!(matches!( + &parts[0], ContentPart::ToolResult { tool_use_id, .. } if tool_use_id == "recall-image" + )); + assert!(matches!( + &parts[1], ContentPart::Text(text) if text == "Images returned by tool call recall-image:" + )); + assert!(matches!(&parts[2], ContentPart::Image { .. })); + } + let call = ToolCall { + id: "missing".to_string(), + name: RECALL_TOOL_HISTORY_NAME.to_string(), + arguments: serde_json::json!({"tool_use_id": "capture-missing"}), + }; + let ToolCallDecision::Inline(result) = policy.decide(&call, &messages, &archive) else { + panic!("missing recall should execute inline"); + }; + assert!(result.images.is_empty()); + assert_eq!( + result.content, + "No matching tool calls found in conversation history." + ); +} + +#[test] +fn recall_does_not_reattach_images_from_an_earlier_recall_result() { + let mut messages = screenshot_exchange("recall-previous", vec![99]); + let MessageContent::ToolUse { name, .. } = &mut messages[0].content else { + panic!("expected tool call"); + }; + *name = RECALL_TOOL_HISTORY_NAME.to_string(); + let policy = ToolPolicy::new(&[definition(RECALL_TOOL_HISTORY_NAME)]); + let call = ToolCall { + id: "recall-now".to_string(), + name: RECALL_TOOL_HISTORY_NAME.to_string(), + arguments: serde_json::json!({}), + }; + let ToolCallDecision::Inline(result) = policy.decide(&call, &messages, &[]) else { + panic!("recall should execute inline"); + }; + assert!(result.images.is_empty()); + assert_eq!( + result.content, + "No matching tool calls found in conversation history." + ); +} + #[test] fn loop_guard_detects_repeated_failures_and_resets_after_detection() { let mut guard = ToolLoopGuard::new(5, 3); diff --git a/crates/galaxy_agent_core/src/types.rs b/crates/galaxy_agent_core/src/types.rs index 8f5c720d..e17f88dc 100644 --- a/crates/galaxy_agent_core/src/types.rs +++ b/crates/galaxy_agent_core/src/types.rs @@ -190,9 +190,53 @@ pub struct ToolResult { pub call_id: String, pub content: String, pub status: ToolResultStatus, + /// Images returned by a tool, kept outside the text so provider adapters receive real images. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub images: Vec, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct ToolResultImage { + pub data: Vec, + pub mime_type: String, +} + +pub(crate) fn tool_image_group_call_id(part: &ContentPart) -> Option<&str> { + let ContentPart::Text(text) = part else { + return None; + }; + text.strip_prefix("Images returned by tool call ")? + .strip_suffix(':') + .filter(|call_id| !call_id.is_empty()) } impl ToolResult { + pub fn into_batch_content_parts(results: Vec) -> Vec { + let mut parts = Vec::new(); + let mut image_parts = Vec::new(); + for result in results { + let is_error = !matches!(result.status, ToolResultStatus::Success); + if !result.images.is_empty() { + image_parts.push(ContentPart::Text(format!( + "Images returned by tool call {}:", result.call_id + ))); + image_parts.extend(result.images.into_iter().map(|image| ContentPart::Image { + data: image.data, + mime_type: image.mime_type, + })); + } + parts.push(ContentPart::ToolResult { + tool_use_id: result.call_id, + content: result.content, + is_error, + }); + } + // OpenAI requires every tool result before the next user message. Its adapter splits + // adjacent images into a user message, so never interleave them with a result batch. + parts.extend(image_parts); + parts + } + pub fn is_error(&self) -> bool { matches!( self.status, diff --git a/crates/galaxy_agent_core/src/types_tests.rs b/crates/galaxy_agent_core/src/types_tests.rs index 54e88198..62a65df0 100644 --- a/crates/galaxy_agent_core/src/types_tests.rs +++ b/crates/galaxy_agent_core/src/types_tests.rs @@ -169,6 +169,7 @@ fn tool_events_keep_one_call_id_across_permission_and_execution() { }, ToolEvent::Completed { result: ToolResult { + images: Vec::new(), call_id: "call-1".to_string(), content: "ok".to_string(), status: ToolResultStatus::Success, @@ -182,11 +183,13 @@ fn tool_events_keep_one_call_id_across_permission_and_execution() { #[test] fn denied_results_are_errors_but_cancelled_results_are_distinct() { let denied = ToolResult { + images: Vec::new(), call_id: "denied".to_string(), content: "permission denied".to_string(), status: ToolResultStatus::Denied, }; let cancelled = ToolResult { + images: Vec::new(), call_id: "cancelled".to_string(), content: "cancelled".to_string(), status: ToolResultStatus::Cancelled, From 7dceca51cfab1a7ff3f9699cf6fe8cfa43d33cf1 Mon Sep 17 00:00:00 2001 From: Ryan Ward Date: Thu, 22 Oct 2026 16:12:16 +0000 Subject: [PATCH 12/12] Allow clarifying questions in Allow All profiles Separate clarification questions from action approvals while preserving Galaxy Decides and Read Only behavior. Respect question capabilities in explicit tool overrides, keep CLI monitors focused, and render approved asynchronous questions while they wait for answers. Validation: cargo build passed on Linux. Added permission, request and question-state regressions; tests and GUI execution were not run under the build-only workflow. --- AGENTS.md | 1 + app/src/ai/agent/api/impl.rs | 11 ++++ app/src/ai/agent/api/impl_tests.rs | 45 +++++++++++++- app/src/ai/blocklist/block.rs | 34 +++++------ .../inline_action/ask_user_question_view.rs | 29 +++++---- .../ask_user_question_view_tests.rs | 59 ++++++++++++++++++- app/src/ai/runtime/rig_request_tests.rs | 52 ++++++++++++++++ .../src/ai_execution_profile.rs | 12 ++-- .../src/ai_execution_profile_tests.rs | 23 ++++++++ 9 files changed, 227 insertions(+), 39 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 0ed85621..0874bb3f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -250,6 +250,7 @@ Key invariants: - Safe inspection and output polling do not add risk-review model calls. A narrow shell fast path accepts validated `pwd`, `ls`, Git version, and hardened Git status/diff forms; shell expansion, chaining, redirects, unknown flags and unhardened Git commands still require review. Explicit restrictions are checked before fast approval. Galaxy Decides skips user-question pauses. External ACP/CLI agents retain their own protocol permission flows; unstructured ACP requests must not grant Galaxy terminal mutations. - Allow All bypasses independent model review and human approval for Galaxy-owned tools, while explicit command/MCP denylists, protected paths, workspace restrictions, command ownership, and child delegation boundaries remain enforced. +- Ask First and Allow All expose `ask_user_question` for clarifying missing information; Allow All still bypasses action-approval panes. Galaxy Decides and Read Only omit clarification pauses. CLI monitoring and tool-free assessments must not gain question tools through an explicit tool override. - Profiles persist `command_allowlist` regular expressions. A matching complete submitted command bypasses permission review in every mode, including Read Only, only after explicit restrictions are checked; denylist matching uses the shell command decomposition as well. Profile permission-mode changes preserve the explicit list but clear remembered command digests. - Profiles persist `orchestration_providers`: absent/null includes all providers (including future additions); an explicit empty list permits none. Entries use configured provider names. Orchestration tool descriptions and pickers filter models, and RunAgents validates per-child/default model selection again at launch and recovery. Restricted profiles cannot use custom routers or external harnesses to bypass provider selection. diff --git a/app/src/ai/agent/api/impl.rs b/app/src/ai/agent/api/impl.rs index df4c1cb2..4312bb43 100644 --- a/app/src/ai/agent/api/impl.rs +++ b/app/src/ai/agent/api/impl.rs @@ -28,6 +28,16 @@ pub(crate) fn prepare_direct_provider_params( .clone() .unwrap_or_else(|| get_supported_tools(params)); remove_orchestration_tools_if_disabled(&mut supported_tools, params.orchestration_enabled); + let ask_user_question_enabled = FeatureFlag::AskUserQuestion.is_enabled() + && params.ask_user_question_enabled + && matches!( + params.permission_mode, + crate::ai::execution_profiles::PermissionMode::AskFirst + | crate::ai::execution_profiles::PermissionMode::AllowAll + ); + supported_tools.retain(|tool| { + *tool != api::ToolType::AskUserQuestion || ask_user_question_enabled + }); // Computer use controls this machine, even when the terminal is connected elsewhere. // Only expose the initial request here; the dedicated profile owns subsequent actions. supported_tools.retain(|tool| *tool != api::ToolType::UseComputer @@ -39,6 +49,7 @@ pub(crate) fn prepare_direct_provider_params( params.orchestration_enabled, ); supported_cli_agent_tools.retain(|tool| !matches!(tool, api::ToolType::UseComputer | api::ToolType::RequestComputerUse)); + supported_cli_agent_tools.retain(|tool| *tool != api::ToolType::AskUserQuestion); if params.permission_mode == crate::ai::execution_profiles::PermissionMode::ReadOnly { let permitted = |tool: &api::ToolType| is_read_only_tool(tool) || (!params.command_allowlist.is_empty() && matches!(tool, api::ToolType::RunShellCommand | api::ToolType::ReadShellCommandOutput | api::ToolType::WriteToLongRunningShellCommand)); diff --git a/app/src/ai/agent/api/impl_tests.rs b/app/src/ai/agent/api/impl_tests.rs index 2bfef573..6882499d 100644 --- a/app/src/ai/agent/api/impl_tests.rs +++ b/app/src/ai/agent/api/impl_tests.rs @@ -191,9 +191,7 @@ fn supported_tools_omits_ask_user_question_when_disabled() { #[test] fn supported_tools_includes_ask_user_question_when_enabled_and_feature_flag_is_enabled() { - if !FeatureFlag::AskUserQuestion.is_enabled() { - return; - } + let _flag = FeatureFlag::AskUserQuestion.override_enabled(true); let params = request_params_with_ask_user_question_enabled(true); let supported_tools = get_supported_tools(¶ms); @@ -201,6 +199,47 @@ fn supported_tools_includes_ask_user_question_when_enabled_and_feature_flag_is_e assert!(supported_tools.contains(&api::ToolType::AskUserQuestion)); } +#[test] +fn question_overrides_respect_mode_capability_and_feature_without_enabling_cli_pauses() { + use crate::ai::execution_profiles::PermissionMode; + + for (mode, enabled, feature_enabled, expected) in [ + (PermissionMode::AskFirst, true, true, true), + (PermissionMode::AllowAll, true, true, true), + (PermissionMode::AllowAll, false, true, false), + (PermissionMode::AllowAll, true, false, false), + (PermissionMode::GalaxyDecides, true, true, false), + (PermissionMode::ReadOnly, true, true, false), + ] { + let _flag = FeatureFlag::AskUserQuestion.override_enabled(feature_enabled); + let mut params = request_params_with_ask_user_question_enabled(enabled); + params.permission_mode = mode; + params.supported_tools_override = Some(vec![api::ToolType::AskUserQuestion]); + + let (normal, cli) = super::prepare_direct_provider_params(&mut params); + + assert_eq!( + normal.contains(&api::ToolType::AskUserQuestion), + expected, + "{mode:?}" + ); + assert!(!cli.contains(&api::ToolType::AskUserQuestion), "{mode:?}"); + } +} + +#[test] +fn question_availability_does_not_expand_empty_tool_overrides() { + let _flag = FeatureFlag::AskUserQuestion.override_enabled(true); + let mut params = request_params_with_ask_user_question_enabled(true); + params.permission_mode = crate::ai::execution_profiles::PermissionMode::AllowAll; + params.supported_tools_override = Some(Vec::new()); + + let (normal, cli) = super::prepare_direct_provider_params(&mut params); + + assert!(normal.is_empty()); + assert!(cli.is_empty()); +} + #[test] fn remote_supported_tools_include_search_codebase_when_connected_and_feature_flag_is_enabled() { let _flag = FeatureFlag::RemoteCodebaseIndexing.override_enabled(true); diff --git a/app/src/ai/blocklist/block.rs b/app/src/ai/blocklist/block.rs index f60bfc74..8b4447c5 100644 --- a/app/src/ai/blocklist/block.rs +++ b/app/src/ai/blocklist/block.rs @@ -3998,12 +3998,7 @@ impl AIBlock { }); self.ask_user_question_view = Some(view.clone()); - if self - .action_model - .as_ref(ctx) - .get_action_status(self.client_ids.conversation_id, action_id) - .is_some_and(|status| status.is_blocked()) - { + if view.as_ref(ctx).is_waiting_on_user_answers(ctx) { ctx.focus(&view); } if self.sync_ask_user_question_speedbump_footer(ctx) @@ -4826,16 +4821,20 @@ impl AIBlock { self.model.status(app).is_complete() } - /// Returns `true` if the block contains any actions that are blocked on user confirmation. + /// Returns `true` when actions need confirmation or a question is waiting for an answer. pub fn is_blocked_on_user_confirmation(&self, app: &AppContext) -> bool { - self.requested_action_ids - .iter() - .filter_map(|id| { - self.action_model - .as_ref(app) - .get_action_status(self.client_ids.conversation_id, id) - }) - .any(|status| status.is_blocked()) + self.ask_user_question_view + .as_ref() + .is_some_and(|view| view.as_ref(app).is_waiting_on_user_answers(app)) + || self + .requested_action_ids + .iter() + .filter_map(|id| { + self.action_model + .as_ref(app) + .get_action_status(self.client_ids.conversation_id, id) + }) + .any(|status| status.is_blocked()) } /// Returns the server output ID for this AI block, if available. @@ -5369,10 +5368,7 @@ impl AIBlock { did_focus_subview = true; } else if let Some(ask_user_question_view) = self.ask_user_question_view.as_ref().filter(|view| { - pending_action_id.is_some_and(|id| { - let view = view.as_ref(ctx); - view.action_id() == id && view.is_editing() - }) + view.as_ref(ctx).is_waiting_on_user_answers(ctx) }) { ctx.focus(ask_user_question_view); diff --git a/app/src/ai/blocklist/inline_action/ask_user_question_view.rs b/app/src/ai/blocklist/inline_action/ask_user_question_view.rs index 4502bd7b..905f1f21 100644 --- a/app/src/ai/blocklist/inline_action/ask_user_question_view.rs +++ b/app/src/ai/blocklist/inline_action/ask_user_question_view.rs @@ -188,6 +188,17 @@ enum AskUserQuestionPhase<'a> { }, } +fn ask_user_question_is_waiting_on_user_answers( + phase: AskUserQuestionPhase<'_>, + status: Option<&AIActionStatus>, +) -> bool { + matches!(phase, AskUserQuestionPhase::Editing) + && matches!( + status, + Some(AIActionStatus::Blocked | AIActionStatus::RunningAsync) + ) +} + #[derive(Clone, Debug, Eq, PartialEq)] /// Editing-phase state for the questionnaire. /// Holds the active question index and one draft slot per question. @@ -806,10 +817,6 @@ impl AskUserQuestionView { &self.action_id } - pub fn is_editing(&self) -> bool { - self.session.is_editing() - } - /// Sets the settings-link mouse handle used in the speedbump footer. /// The footer is only actually drawn during the completed/finished render paths. pub fn set_speedbump_settings_link( @@ -902,10 +909,12 @@ impl AskUserQuestionView { }) } - /// True when the action model reports this inline action is blocked waiting on user answers. - fn is_waiting_on_user_answers(&self, app: &AppContext) -> bool { - self.action_status(app) - .is_some_and(|status| status.is_blocked()) + /// Questions can wait before execution or in an approved async executor, as in Allow All. + pub(crate) fn is_waiting_on_user_answers(&self, app: &AppContext) -> bool { + ask_user_question_is_waiting_on_user_answers( + self.session.phase(), + self.action_status(app).as_ref(), + ) } pub fn should_render_inline(&self, app: &AppContext) -> bool { @@ -1170,8 +1179,8 @@ impl AskUserQuestionView { self.abort_auto_advance(); let action_id = self.action_id().clone(); let conversation_id = self.conversation_id; - // Completing the executor only stages the result; execute_action drives the action model - // to publish final status/events for this inline block. + // A running executor receives the answer directly. A blocked question must first start + // execution; execute_action starts it and leaves an already-running receiver intact. self.action_model.update(ctx, |action_model, ctx| { action_model .ask_user_question_executor(ctx) diff --git a/app/src/ai/blocklist/inline_action/ask_user_question_view_tests.rs b/app/src/ai/blocklist/inline_action/ask_user_question_view_tests.rs index 212b57c5..cc990022 100644 --- a/app/src/ai/blocklist/inline_action/ask_user_question_view_tests.rs +++ b/app/src/ai/blocklist/inline_action/ask_user_question_view_tests.rs @@ -1,10 +1,16 @@ +use std::sync::Arc; + use ai::agent::action::{AskUserQuestionItem, AskUserQuestionOption, AskUserQuestionType}; -use ai::agent::action_result::AskUserQuestionAnswerItem; +use ai::agent::action_result::{AskUserQuestionAnswerItem, AskUserQuestionResult}; use super::{ - ask_user_question_view_state, AskUserQuestionAction, AskUserQuestionEffect, - AskUserQuestionPhase, AskUserQuestionSession, AskUserQuestionViewState, + ask_user_question_is_waiting_on_user_answers, ask_user_question_view_state, + AskUserQuestionAction, AskUserQuestionEffect, AskUserQuestionPhase, + AskUserQuestionSession, AskUserQuestionViewState, }; +use crate::ai::agent::task::TaskId; +use crate::ai::agent::{AIAgentActionId, AIAgentActionResult, AIAgentActionResultType}; +use crate::ai::blocklist::action_model::AIActionStatus; fn build_question( question_id: &str, @@ -42,6 +48,53 @@ fn current_draft(session: &AskUserQuestionSession) -> Option<&super::QuestionDra session.current().and_then(|current| current.draft) } +#[test] +fn blocked_and_running_questions_accept_answers_until_the_session_completes() { + let mut session = build_session(vec![build_question("q1", "Destination", false, true, &[])]); + let blocked = AIActionStatus::Blocked; + let running = AIActionStatus::RunningAsync; + + for status in [&blocked, &running] { + assert!(ask_user_question_is_waiting_on_user_answers( + session.phase(), + Some(status), + )); + } + + assert!(matches!( + session.apply(AskUserQuestionAction::SkipAll), + AskUserQuestionEffect::Submit(_) + )); + for status in [&blocked, &running] { + assert!(!ask_user_question_is_waiting_on_user_answers( + session.phase(), + Some(status), + )); + } +} + +#[test] +fn queued_preprocessing_missing_and_finished_questions_do_not_accept_answers() { + let session = build_session(vec![build_question("q1", "Destination", false, true, &[])]); + let finished = AIActionStatus::Finished(Arc::new(AIAgentActionResult { + result: AIAgentActionResultType::AskUserQuestion(AskUserQuestionResult::Cancelled), + task_id: TaskId::new("question-task".to_owned()), + id: AIAgentActionId::from("question-action".to_owned()), + })); + + for status in [ + None, + Some(AIActionStatus::Queued), + Some(AIActionStatus::Preprocessing), + Some(finished), + ] { + assert!(!ask_user_question_is_waiting_on_user_answers( + session.phase(), + status.as_ref(), + )); + } +} + #[test] fn enter_on_other_row_focuses_the_other_input() { let mut session = build_session(vec![build_question("q1", "Only", false, true, &["Stable"])]); diff --git a/app/src/ai/runtime/rig_request_tests.rs b/app/src/ai/runtime/rig_request_tests.rs index f14f8dad..830fb161 100644 --- a/app/src/ai/runtime/rig_request_tests.rs +++ b/app/src/ai/runtime/rig_request_tests.rs @@ -242,6 +242,58 @@ fn plan_turn_requires_a_plan_tool_result() { assert!(prompt.contains("Do not return the plan only as prose")); } +#[test] +fn normal_and_plan_turns_advertise_clarifications_only_for_interactive_permission_modes() { + use crate::ai::execution_profiles::{ + AIExecutionProfile, AskUserQuestionPermission, PermissionMode, + }; + + let _flag = galaxy_core::features::FeatureFlag::AskUserQuestion.override_enabled(true); + for (permission_mode, expected) in [ + (PermissionMode::AskFirst, true), + (PermissionMode::AllowAll, true), + (PermissionMode::GalaxyDecides, false), + (PermissionMode::ReadOnly, false), + ] { + for plan_mode in [false, true] { + let mut profile = AIExecutionProfile { + permission_mode, + ..Default::default() + }; + profile.project_permission_mode(); + let mut params = RequestParams::new_for_test(); + params.permission_mode = permission_mode; + params.ask_user_question_enabled = + profile.ask_user_question != AskUserQuestionPermission::Never; + let mut input = user_query("Clarify the requested destination before proceeding"); + if plan_mode { + let AIAgentInput::UserQuery { + user_query_mode, .. + } = &mut input + else { + unreachable!("user_query returns a user query input"); + }; + *user_query_mode = UserQueryMode::Plan; + } + params.input = vec![input]; + let (normal, cli) = + crate::ai::agent::api::prepare_direct_provider_params(&mut params); + assert!(!cli.contains(&ToolType::AskUserQuestion)); + + let prepared = prepare_rig_turn(&config(), params, normal, cli); + assert_eq!( + prepared + .request + .tools + .iter() + .any(|tool| tool.name == "ask_user_question"), + expected, + "{permission_mode:?}, plan_mode={plan_mode}" + ); + } + } +} + #[test] fn no_tools_turn_flattens_historical_tool_protocol_messages() { let mut params = RequestParams::new_for_test(); diff --git a/crates/cloud_object_models/src/ai_execution_profile.rs b/crates/cloud_object_models/src/ai_execution_profile.rs index 5ad265c3..ada17173 100644 --- a/crates/cloud_object_models/src/ai_execution_profile.rs +++ b/crates/cloud_object_models/src/ai_execution_profile.rs @@ -661,10 +661,14 @@ impl AIExecutionProfile { PermissionMode::GalaxyDecides => RunAgentsPermission::GalaxyDecides, PermissionMode::AllowAll => RunAgentsPermission::AlwaysAllow, }; - self.ask_user_question = if self.permission_mode == PermissionMode::AskFirst { - AskUserQuestionPermission::AlwaysAsk - } else { - AskUserQuestionPermission::Never + // Allow All bypasses tool approval, while clarification can still need user input. + self.ask_user_question = match self.permission_mode { + PermissionMode::AskFirst | PermissionMode::AllowAll => { + AskUserQuestionPermission::AlwaysAsk + } + PermissionMode::ReadOnly | PermissionMode::GalaxyDecides => { + AskUserQuestionPermission::Never + } }; self.directory_allowlist.clear(); self.mcp_allowlist.clear(); diff --git a/crates/cloud_object_models/src/ai_execution_profile_tests.rs b/crates/cloud_object_models/src/ai_execution_profile_tests.rs index fb53ee6d..470396c5 100644 --- a/crates/cloud_object_models/src/ai_execution_profile_tests.rs +++ b/crates/cloud_object_models/src/ai_execution_profile_tests.rs @@ -61,6 +61,7 @@ fn serialized_profiles_only_store_one_permission_mode_and_bounded_approvals() { "write_to_pty", "run_agents", "computer_use", + "ask_user_question", "directory_allowlist", "mcp_allowlist", ] { @@ -127,3 +128,25 @@ fn allow_all_projects_every_tool_to_always_allow() { assert_eq!(profile.computer_use, ComputerUsePermission::AlwaysAllow); assert_eq!(profile.run_agents, RunAgentsPermission::AlwaysAllow); } + +#[test] +fn clarification_availability_follows_permission_mode_after_projection_and_restore() { + for (mode, expected) in [ + (PermissionMode::AskFirst, AskUserQuestionPermission::AlwaysAsk), + (PermissionMode::AllowAll, AskUserQuestionPermission::AlwaysAsk), + (PermissionMode::GalaxyDecides, AskUserQuestionPermission::Never), + (PermissionMode::ReadOnly, AskUserQuestionPermission::Never), + ] { + let mut profile = AIExecutionProfile { + permission_mode: mode, + ..Default::default() + }; + profile.project_permission_mode(); + assert_eq!(profile.ask_user_question, expected, "{mode:?}"); + + let saved = serde_json::to_value(&profile).unwrap(); + assert!(saved.get("ask_user_question").is_none()); + let restored: AIExecutionProfile = serde_json::from_value(saved).unwrap(); + assert_eq!(restored.ask_user_question, expected, "{mode:?}"); + } +}