Fix environment checkout hanging during shell bootstrap

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<CheckoutReport>,
 }
 
+async fn await_checkout_helper<T>(
+    operation: impl Future<Output = Result<T, PrepareEnvironmentError>>,
+    deadline: impl Future,
+) -> Result<T, PrepareEnvironmentError> {
+    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());

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::<ExitCode>();
+            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] {

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;

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 {