diff --git a/app/src/ai/agent_sdk/driver/environment.rs b/app/src/ai/agent_sdk/driver/environment.rs index 45a6cb77a67..342f0701411 100644 --- a/app/src/ai/agent_sdk/driver/environment.rs +++ b/app/src/ai/agent_sdk/driver/environment.rs @@ -48,6 +48,7 @@ const CLONE_FAILURE_OUTPUT_TRUNCATION_MARKER: &str = "\n… clone output truncat const SETUP_COMMAND_OUTPUT_TRUNCATION_MARKER: &str = "\n… setup command output truncated …\n"; const SETUP_COMMAND_TIMEOUT: Duration = Duration::from_mins(30); const SETUP_COMMAND_CWD_RESET_TIMEOUT: Duration = Duration::from_secs(30); +const CHECKOUT_HELPER_TIMEOUT: Duration = Duration::from_mins(30); #[derive(Clone, Copy, Debug)] pub enum SetupCommandPhase { @@ -129,6 +130,8 @@ pub enum PrepareEnvironmentError { UnsupportedRepositoryForge { repo_name: String }, #[error("Environment checkout helper failed: {reason}")] CheckoutHelper { reason: &'static str }, + #[error("Environment checkout helper timed out after {timeout_seconds}s")] + CheckoutHelperTimedOut { timeout_seconds: u64 }, #[error("Terminal driver error while preparing environment: {source}")] TerminalDriver { source: AgentDriverError }, } @@ -1024,7 +1027,10 @@ async fn remove_repository_origins_from_repos( let batch = checkout_requests_batch(&selected, working_dir)?; let result = execute_checkout_helper(&batch, true, spawner) .await - .map_err(|_| PrepareEnvironmentError::RemoveRepositoryOrigins)?; + .map_err(|error| match error { + error @ PrepareEnvironmentError::CheckoutHelperTimedOut { .. } => error, + _ => PrepareEnvironmentError::RemoveRepositoryOrigins, + })?; if result.command_result.exit_code == 0.into() { Ok(()) } else { @@ -1204,6 +1210,18 @@ struct CheckoutHelperResult { report: Option, } +async fn await_checkout_helper( + operation: impl Future>, + deadline: impl Future, +) -> Result { + match select(Box::pin(operation), Box::pin(deadline)).await { + Either::Left((result, _)) => result, + Either::Right((_, _)) => Err(PrepareEnvironmentError::CheckoutHelperTimedOut { + timeout_seconds: CHECKOUT_HELPER_TIMEOUT.as_secs(), + }), + } +} + #[tracing::instrument(skip_all, err, fields( tags.cloud_agent = true, remove_origins_only = remove_origins_only, @@ -1246,7 +1264,11 @@ async fn execute_checkout_helper( remove_origins_only, active_shell_type(spawner).await, ); - let result = execute_command(command, spawner).await; + let result = await_checkout_helper( + execute_command(command, spawner), + Timer::after(CHECKOUT_HELPER_TIMEOUT), + ) + .await; let _ = std::fs::remove_file(&requests_file); let command_result = result?; let report = read_checkout_report(&report_file, batch.repositories.len()); diff --git a/app/src/ai/agent_sdk/driver/environment_tests.rs b/app/src/ai/agent_sdk/driver/environment_tests.rs index f6483a23f9c..96927766c45 100644 --- a/app/src/ai/agent_sdk/driver/environment_tests.rs +++ b/app/src/ai/agent_sdk/driver/environment_tests.rs @@ -18,8 +18,9 @@ use super::super::environment_checkout_protocol::{ use super::{ CloneFailureIdentityDiagnostics, PrepareEnvironmentError, RepositoryCloneRequest, ResolvedRepository, SETUP_COMMAND_OUTPUT_TRUNCATION_MARKER, SetupCommandPhase, - WorkspaceConfiguration, await_setup_phase, build_checkout_helper_command, environment_snapshot, - merge_repos_deduped, parse_resolved_head_sha, read_checkout_report, reported_resolved_heads, + WorkspaceConfiguration, await_checkout_helper, await_setup_phase, + build_checkout_helper_command, environment_snapshot, merge_repos_deduped, + parse_resolved_head_sha, read_checkout_report, reported_resolved_heads, repository_clone_requests, setup_command_failure, single_repo_name, validate_repository_preparation_overrides, }; @@ -118,6 +119,51 @@ fn duplicate_resolved_repositories_fail_before_clone() { assert!(WorkspaceConfiguration::from_resolved(repositories, Vec::new(), Vec::new()).is_err()); } +#[test] +fn checkout_helper_timeout_covers_command_start_and_exit() { + for command_started in [false, true] { + block_on(async { + let (start_tx, start_rx) = oneshot::channel::<()>(); + let (exit_tx, exit_rx) = oneshot::channel::(); + let (deadline_tx, deadline_rx) = oneshot::channel::<()>(); + let operation = async { + start_rx.await.unwrap(); + Ok(exit_rx.await.unwrap()) + }; + let mut wait = Box::pin(await_checkout_helper(operation, deadline_rx)); + assert!(poll!(wait.as_mut()).is_pending()); + if command_started { + start_tx.send(()).unwrap(); + assert!(poll!(wait.as_mut()).is_pending()); + } + deadline_tx.send(()).unwrap(); + + assert!(matches!( + wait.await, + Err(PrepareEnvironmentError::CheckoutHelperTimedOut { + timeout_seconds: 1800 + }) + )); + assert!(exit_tx.is_canceled()); + }); + } +} + +#[test] +fn checkout_helper_completion_preserves_command_result() { + assert_eq!( + block_on(await_checkout_helper(ready(Ok(17)), pending::<()>())).unwrap(), + 17 + ); + assert!(matches!( + block_on(await_checkout_helper( + ready(Err::<(), _>(PrepareEnvironmentError::InvalidRuntimeState)), + pending::<()>() + )), + Err(PrepareEnvironmentError::InvalidRuntimeState) + )); +} + #[test] fn setup_timeout_covers_command_start_and_exit_without_resetting_deadline() { for command_started in [false, true] { diff --git a/app/src/terminal/input.rs b/app/src/terminal/input.rs index 031320dad15..f8700d15f25 100644 --- a/app/src/terminal/input.rs +++ b/app/src/terminal/input.rs @@ -7711,6 +7711,16 @@ impl Input { return; } + // The final bootstrap precmd replaces the active block after the Bootstrapped hook. + if !self + .model + .lock() + .block_list() + .is_bootstrapping_precmd_done() + { + return; + } + let command = self.get_command(ctx); if self.can_execute_command(ctx).is_no() { return; diff --git a/app/src/terminal/input_tests.rs b/app/src/terminal/input_tests.rs index e4a4fc1b410..1a8263d40bb 100644 --- a/app/src/terminal/input_tests.rs +++ b/app/src/terminal/input_tests.rs @@ -91,7 +91,9 @@ use crate::terminal::input::slash_commands::SlashCommandsEvent; use crate::terminal::keys::TerminalKeybindings; use crate::terminal::local_shell::LocalShellState; use crate::terminal::local_tty::shell::ShellStarter; -use crate::terminal::model::ansi::{Handler, PromptMetadata}; +use crate::terminal::model::ansi::{ + CommandFinishedValue, CompletionMetadata, Handler, PrecmdValue, PromptMetadata, +}; use crate::terminal::model::block::{BlockId, SerializedBlock}; use crate::terminal::model::blocks::{BlockListPoint, insert_block}; use crate::terminal::model::grid::Dimensions as _; @@ -165,6 +167,47 @@ fn subscribe_to_execute_command_events( events } +#[test] +fn pending_command_waits_for_final_bootstrap_precmd() { + App::test((), |mut app| async move { + initialize_app(&mut app); + let terminal = add_window_with_bootstrapped_terminal(&mut app, None, None).await; + let events = subscribe_to_execute_command_events(&mut app, &terminal); + let input = terminal.read(&app, |terminal, _| terminal.input().clone()); + terminal.update(&mut app, |terminal, ctx| { + terminal.execute_command_or_set_pending("echo ready", ctx); + }); + assert!(input.read(&app, |input, _| input.has_pending_command())); + assert!(events.borrow().is_empty()); + + terminal.update(&mut app, |terminal, ctx| { + { + let mut model = terminal.model.lock(); + let session_id = model.block_list().active_block().session_id().unwrap(); + let completion_metadata = CompletionMetadata::default(); + model.command_finished(CommandFinishedValue { + completion_metadata: completion_metadata.clone(), + session_id: Some(session_id.as_u64()), + }); + model.precmd_with_completion_metadata(PrecmdValue { + completion_metadata, + prompt_metadata: PromptMetadata { + session_id: Some(session_id.as_u64()), + ..Default::default() + }, + }); + } + terminal.model_event_dispatcher().update(ctx, |_, ctx| { + ctx.emit(ModelEvent::BootstrapPrecmdDone); + }); + }); + + assert!(!input.read(&app, |input, _| input.has_pending_command())); + assert_eq!(events.borrow().len(), 1); + assert_eq!(events.borrow()[0].0, "echo ready"); + }); +} + #[test] fn unshared_terminal_executes_command_locally() { App::test((), |mut app| async move {