Bound MCP tool calls and environment setup commands

Cargo.toml

@@ -420,7 +420,7 @@ typed-path = "0.10.0"
 streaming-iterator = "0.1.0"
 derivative = "2.2.0"
 parquet = { version = "55.0.0", features = ["arrow"] }
-rmcp = { version = "2.2" }
+rmcp = { version = "3.5" }
 
 [profile.release]
 # Use line-tables-only (debug = 1) rather than full debuginfo (debug = 2 /

app/src/ai/agent_sdk/driver.rs

@@ -887,6 +887,8 @@ pub enum AgentDriverError {
         command: String,
         output: Option<String>,
     },
+    #[error("Environment setup failed: {message}")]
+    SetupCommandTimedOut { message: String },
     #[error("Cloud provider setup failed")]
     CloudProviderSetupFailed(#[from] cloud_provider::CloudProviderSetupError),
     #[error("Could not resolve working directory {}", path.display())]
@@ -1060,6 +1062,9 @@ impl From<PrepareEnvironmentError> for AgentDriverError {
                     output,
                 }
             }
+            PrepareEnvironmentError::SetupCommandTimedOut { .. } => {
+                AgentDriverError::SetupCommandTimedOut { message }
+            }
             _ => AgentDriverError::EnvironmentSetupFailed(message),
         }
     }
@@ -1660,6 +1665,7 @@ impl AgentDriver {
                     err,
                     AgentDriverError::EnvironmentSetupFailed(_)
                         | AgentDriverError::SetupCommandFailed { .. }
+                        | AgentDriverError::SetupCommandTimedOut { .. }
                         | AgentDriverError::SetupCommandExitedShell { .. }
                 ) {
                     let _ = foreground_for_error
@@ -2531,6 +2537,13 @@ impl AgentDriver {
         stage: &str,
         error: &AgentDriverError,
     ) {
+        // The terminal may still be executing setup; do not retain it or accept debug turns.
+        if matches!(error, AgentDriverError::SetupCommandTimedOut { .. }) {
+            log::warn!(
+                "Environment setup lifecycle: event=timeout_teardown stage={stage} retention_skipped=true"
+            );
+            return;
+        }
         let idle_on_fail = match foreground.spawn(|me, _| me.idle_on_fail).await {
             Ok(idle_on_fail) => idle_on_fail,
             Err(spawn_error) => {

app/src/ai/agent_sdk/driver/environment.rs

@@ -11,14 +11,15 @@ use ai::index::full_source_code_embedding::manager::{
 use chrono::Utc;
 use cloud_object_models::CodeForge;
 use futures::channel::oneshot;
-use futures::future::join_all;
+use futures::future::{Either, join_all, select};
+use instant::Instant;
 use repo_metadata::repositories::{DetectedRepositories, RepoDetectionSource};
 use uuid::Uuid;
 use warp_cli::agent::{RepositoryForge, RepositoryHeadRef, RepositoryPreparationOverride};
 use warp_completer::completer::{CommandExitStatus, CommandOutput};
 use warp_core::command::ExitCode;
 use warp_core::{safe_info, safe_warn};
-use warpui::r#async::FutureExt;
+use warpui::r#async::{FutureExt, Timer};
 use warpui::{ModelContext, ModelSpawner, SingletonEntity};
 
 #[cfg(feature = "local_fs")]
@@ -30,6 +31,7 @@ use crate::ai::agent_sdk::environment_snapshot::{
 };
 use crate::ai::agent_sdk::setup_observability::{SetupClientEventReporter, SetupStep};
 use crate::ai::cloud_environments::SourceRepo;
+use crate::server::telemetry::secret_redaction::redact_secrets_in_string;
 use crate::terminal::model::BlockId;
 use crate::terminal::model::session::command_executor::shell_escape_single_quotes;
 use crate::terminal::shell::ShellType;
@@ -39,6 +41,32 @@ const ENVIRONMENT_SNAPSHOT_CAPTURE_TIMEOUT: Duration = Duration::from_secs(2);
 const CLONE_FAILURE_IDENTITY_QUERY_TIMEOUT: Duration = Duration::from_secs(2);
 const CLONE_FAILURE_OUTPUT_TRUNCATION_MARKER: &str = "\n… clone output truncated …\n";
 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);
+
+#[derive(Clone, Copy, Debug)]
+pub enum SetupCommandPhase {
+    Execute,
+    ResetWorkingDirectory,
+}
+
+impl SetupCommandPhase {
+    fn timeout(self) -> Duration {
+        match self {
+            Self::Execute => SETUP_COMMAND_TIMEOUT,
+            Self::ResetWorkingDirectory => SETUP_COMMAND_CWD_RESET_TIMEOUT,
+        }
+    }
+}
+
+impl fmt::Display for SetupCommandPhase {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str(match self {
+            Self::Execute => "waiting for command to complete",
+            Self::ResetWorkingDirectory => "resetting the working directory after the command",
+        })
+    }
+}
 
 #[derive(Debug, thiserror::Error)]
 pub enum PrepareEnvironmentError {
@@ -71,6 +99,15 @@ pub enum PrepareEnvironmentError {
         command: String,
         output: Option<String>,
     },
+    #[error(
+        "Setup command #{command_index} timed out after {timeout_seconds}s while {phase}: `{command}`"
+    )]
+    SetupCommandTimedOut {
+        command_index: usize,
+        command: String,
+        phase: SetupCommandPhase,
+        timeout_seconds: u64,
+    },
     #[error("Failed to change directory into {repo_name}")]
     ChangeDirectory { repo_name: String },
     #[error(
@@ -105,6 +142,44 @@ fn setup_command_failure(command: String, output: Option<String>) -> PrepareEnvi
     PrepareEnvironmentError::SetupCommand { command, output }
 }
 
+async fn await_setup_phase<T>(
+    command_index: usize,
+    command: &str,
+    phase: SetupCommandPhase,
+    operation: impl Future<Output = Result<T, PrepareEnvironmentError>>,
+    deadline: impl Future,
+) -> Result<T, PrepareEnvironmentError> {
+    let started_at = Instant::now();
+    log::info!(
+        "Environment setup lifecycle: event=phase_started command_index={command_index} phase={phase:?} timeout_seconds={}",
+        phase.timeout().as_secs()
+    );
+    match select(Box::pin(operation), Box::pin(deadline)).await {
+        Either::Left((result, _)) => {
+            log::info!(
+                "Environment setup lifecycle: event=phase_finished command_index={command_index} phase={phase:?} elapsed_ms={} result_ok={}",
+                started_at.elapsed().as_millis(),
+                result.is_ok()
+            );
+            result
+        }
+        Either::Right((_, _)) => {
+            log::warn!(
+                "Environment setup lifecycle: event=phase_timed_out command_index={command_index} phase={phase:?} elapsed_ms={}",
+                started_at.elapsed().as_millis()
+            );
+            let mut command = command.to_owned();
+            redact_secrets_in_string(&mut command);
+            Err(PrepareEnvironmentError::SetupCommandTimedOut {
+                command_index,
+                command,
+                phase,
+                timeout_seconds: phase.timeout().as_secs(),
+            })
+        }
+    }
+}
+
 fn clone_failure_output_suffix(output: Option<&str>) -> String {
     output
         .filter(|output| !output.is_empty())
@@ -534,22 +609,43 @@ async fn prepare_environment_impl(
                 // environment variable.
                 execute_command("export CI=true".to_string(), spawner).await?;
 
-                for command in setup_commands {
+                for (index, command) in setup_commands.into_iter().enumerate() {
+                    let command_index = index + 1;
                     let command_for_error = command.clone();
                     safe_info!(
                         safe: ("Running setup command"),
                         full: ("Running setup command: {command}")
                     );
 
-                    let command_result = execute_command(command, spawner).await?;
+                    let command_result = await_setup_phase(
+                        command_index,
+                        &command_for_error,
+                        SetupCommandPhase::Execute,
+                        execute_command(command, spawner),
+                        Timer::after(SETUP_COMMAND_TIMEOUT),
+                    )
+                    .await?;
                     if command_result.exit_code != 0.into() {
                         let output =
                             fetch_block_output_plaintext(&command_result.block_id, spawner).await;
                         return Err(setup_command_failure(command_for_error, output));
                     }
 
                     let working_dir_string = working_dir.to_string_lossy().to_string();
-                    if let Err(error) = cd_in_terminal(working_dir_string, spawner).await {
+                    let reset_result = await_setup_phase(
+                        command_index,
+                        &command_for_error,
+                        SetupCommandPhase::ResetWorkingDirectory,
+                        cd_in_terminal(working_dir_string, spawner),
+                        Timer::after(SETUP_COMMAND_CWD_RESET_TIMEOUT),
+                    )
+                    .await;
+                    if matches!(
+                        &reset_result,
+                        Err(PrepareEnvironmentError::SetupCommandTimedOut { .. })
+                    ) {
+                        reset_result?;
+                    } else if let Err(error) = reset_result {
                         log::warn!(
                             "Failed to reset working directory after setup command: {error}"
                         );
@@ -578,6 +674,12 @@ async fn prepare_environment_impl(
         Ok(())
     };
 
+    // A timed-out command may still own the terminal; issuing cleanup commands can hang again.
+    let setup_result = match setup_result {
+        Err(error @ PrepareEnvironmentError::SetupCommandTimedOut { .. }) => return Err(error),
+        result => result,
+    };
+
     // Fill in a forge-appropriate identity for any repo whose effective
     // identity is still exactly what bootstrap set — i.e. nothing (a setup
     // command, or anything else run above) has claimed it yet. This runs