From f5654f592d3f94ec1ae0461b8e4d95df6ed118e8 Mon Sep 17 00:00:00 2001 From: seemeroland Date: Wed, 30 Sep 2026 15:56:08 -0700 Subject: [PATCH 1/3] Bound MCP tool calls and environment setup commands Add a 30-minute logical MCP deadline with bounded cancellation and no ambiguous transport replay. Bound setup commands, report redacted command-specific failures, and retain postmortem history without keeping stalled processes alive. Co-Authored-By: Warp --- Cargo.lock | 1 + app/src/ai/agent_sdk/driver.rs | 13 + app/src/ai/agent_sdk/driver/environment.rs | 112 ++++- .../ai/agent_sdk/driver/environment_tests.rs | 148 ++++++- .../agent_sdk/driver/error_classification.rs | 7 + .../driver/error_classification_tests.rs | 30 +- app/src/ai/agent_sdk/driver/terminal.rs | 40 +- app/src/ai/agent_sdk/driver_tests.rs | 32 ++ .../action_model/execute/call_mcp_tool.rs | 17 +- app/src/ai/mcp/reconnecting_peer.rs | 13 +- crates/mcp/Cargo.toml | 1 + crates/mcp/src/lib.rs | 2 + crates/mcp/src/tool_call.rs | 164 +++++++ crates/mcp/src/tool_call_tests.rs | 411 ++++++++++++++++++ 14 files changed, 975 insertions(+), 16 deletions(-) create mode 100644 crates/mcp/src/tool_call.rs create mode 100644 crates/mcp/src/tool_call_tests.rs diff --git a/Cargo.lock b/Cargo.lock index 7e2a9168d86..de3d7fa523a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8232,6 +8232,7 @@ dependencies = [ "futures", "handlebars", "http 1.4.0", + "instant", "log", "oauth2", "pin-project-lite", diff --git a/app/src/ai/agent_sdk/driver.rs b/app/src/ai/agent_sdk/driver.rs index 22afc267e30..891c0fe9819 100644 --- a/app/src/ai/agent_sdk/driver.rs +++ b/app/src/ai/agent_sdk/driver.rs @@ -887,6 +887,8 @@ pub enum AgentDriverError { command: String, output: Option, }, + #[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 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) => { diff --git a/app/src/ai/agent_sdk/driver/environment.rs b/app/src/ai/agent_sdk/driver/environment.rs index 1caa3021f80..c2622b6934b 100644 --- a/app/src/ai/agent_sdk/driver/environment.rs +++ b/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_secs(30 * 60); +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 startup or completion", + 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, }, + #[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) -> PrepareEnvi PrepareEnvironmentError::SetupCommand { command, output } } +async fn await_setup_phase( + command_index: usize, + command: &str, + phase: SetupCommandPhase, + operation: impl Future>, + deadline: impl Future, +) -> Result { + 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,14 +609,22 @@ 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; @@ -549,7 +632,20 @@ async fn prepare_environment_impl( } 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 diff --git a/app/src/ai/agent_sdk/driver/environment_tests.rs b/app/src/ai/agent_sdk/driver/environment_tests.rs index 301f03061be..67eefc66abf 100644 --- a/app/src/ai/agent_sdk/driver/environment_tests.rs +++ b/app/src/ai/agent_sdk/driver/environment_tests.rs @@ -3,22 +3,28 @@ use std::path::{Path, PathBuf}; use cloud_object_models::CodeForge; use command::blocking::Command; +use futures::channel::oneshot; +use futures::executor::block_on; +use futures::future::{pending, ready}; +use futures::poll; use tempfile::TempDir; use warp_cli::agent::{ RepositoryForge, RepositoryHeadRef, RepositoryIdentity, RepositoryPreparationOverride, }; use warp_completer::completer::{CommandExitStatus, CommandOutput}; +use warp_core::command::ExitCode; use super::{ CloneFailureCredentialIdentity, CloneFailureIdentityDiagnostics, PrepareEnvironmentError, - RepositoryCloneRequest, SETUP_COMMAND_OUTPUT_TRUNCATION_MARKER, - build_git_credential_query_command, build_parallel_clone_command, + RepositoryCloneRequest, SETUP_COMMAND_OUTPUT_TRUNCATION_MARKER, SetupCommandPhase, + await_setup_phase, build_git_credential_query_command, build_parallel_clone_command, build_remove_repository_origins_command, build_resolved_head_command, checkout_command_for, clone_failure_identity_diagnostics, environment_snapshot, is_valid_git_object_id, merge_repos_deduped, parse_resolved_head_sha, parse_resolved_head_shas, read_failed_repo_names, repository_clone_requests, setup_command_failure, single_repo_name, unique_clone_hosts, validate_repository_preparation_overrides, }; +use crate::ai::agent_sdk::driver::AgentDriverError; use crate::ai::cloud_environments::{AmbientAgentEnvironment, SourceRepo}; use crate::terminal::shell::ShellType; @@ -31,6 +37,144 @@ fn command_output(stdout: &str, stderr: &str, status: CommandExitStatus) -> Comm } } +#[test] +fn setup_timeout_covers_command_start_and_exit_without_resetting_deadline() { + 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_setup_phase( + 2, + "./setup.sh", + SetupCommandPhase::Execute, + operation, + deadline_rx, + )); + assert!(poll!(wait.as_mut()).is_pending()); + let start_tx = if command_started { + start_tx.send(()).unwrap(); + assert!(poll!(wait.as_mut()).is_pending()); + None + } else { + Some(start_tx) + }; + deadline_tx.send(()).unwrap(); + let error = wait.await.unwrap_err(); + assert_eq!( + error.to_string(), + "Setup command #2 timed out after 1800s while waiting for command startup or completion: `./setup.sh`" + ); + if let Some(start_tx) = start_tx { + assert!(start_tx.is_canceled()); + } + assert!(exit_tx.is_canceled()); + }); + } +} + +#[test] +fn setup_deadline_preserves_success_and_nonzero_exit_codes() { + for code in [0, 7] { + let result = block_on(await_setup_phase( + 1, + "./setup.sh", + SetupCommandPhase::Execute, + ready(Ok(ExitCode::from(code))), + pending::<()>(), + )) + .unwrap(); + assert_eq!(result, ExitCode::from(code)); + } +} + +#[test] +fn setup_deadline_preserves_shell_exit_error() { + let result = block_on(await_setup_phase::<()>( + 1, + "./setup.sh", + SetupCommandPhase::Execute, + ready(Err(PrepareEnvironmentError::TerminalDriver { + source: AgentDriverError::SetupCommandExitedShell { + command: "./setup.sh".to_string(), + }, + })), + pending::<()>(), + )); + assert!(matches!( + result, + Err(PrepareEnvironmentError::TerminalDriver { + source: AgentDriverError::SetupCommandExitedShell { .. } + }) + )); +} + +#[test] +fn setup_reset_timeout_is_distinct_and_names_the_configured_command() { + let error = block_on(await_setup_phase::( + 3, + "./setup.sh", + SetupCommandPhase::ResetWorkingDirectory, + pending(), + ready(()), + )) + .unwrap_err(); + assert_eq!( + error.to_string(), + "Setup command #3 timed out after 30s while resetting the working directory after the command: `./setup.sh`" + ); +} + +#[test] +fn setup_timeout_redacts_command_before_constructing_error() { + let secret = "AKIAIOSFODNN7EXAMPLE"; + let command = format!("./setup.sh {secret}"); + let error = block_on(await_setup_phase::<()>( + 1, + &command, + SetupCommandPhase::Execute, + pending(), + ready(()), + )) + .unwrap_err(); + assert!(!error.to_string().contains(secret)); + assert!(!format!("{error:?}").contains(secret)); + assert!(error.to_string().contains(&"*".repeat(secret.len()))); +} + +#[test] +fn setup_timeout_stops_following_work_and_drops_the_pending_command() { + block_on(async { + let (command_tx, command_rx) = oneshot::channel::<()>(); + let (deadline_tx, deadline_rx) = oneshot::channel::<()>(); + let mut reached_next_command = false; + let mut setup = Box::pin(async { + await_setup_phase( + 1, + "./setup.sh", + SetupCommandPhase::Execute, + async { + command_rx.await.unwrap(); + Ok(()) + }, + deadline_rx, + ) + .await?; + reached_next_command = true; + Ok::<(), PrepareEnvironmentError>(()) + }); + assert!(poll!(setup.as_mut()).is_pending()); + deadline_tx.send(()).unwrap(); + assert!(setup.await.is_err()); + assert!(!reached_next_command); + assert!(command_tx.is_canceled()); + }); +} + #[test] fn clone_error_includes_short_output() { let error = PrepareEnvironmentError::CloneRepo { diff --git a/app/src/ai/agent_sdk/driver/error_classification.rs b/app/src/ai/agent_sdk/driver/error_classification.rs index cb36b9ef591..7588410b2a1 100644 --- a/app/src/ai/agent_sdk/driver/error_classification.rs +++ b/app/src/ai/agent_sdk/driver/error_classification.rs @@ -193,6 +193,13 @@ pub fn classify_driver_error(error: &AgentDriverError) -> (AgentTaskState, TaskS AgentTaskState::Failed, setup_command_status_update(error, command, output.as_deref()), ), + AgentDriverError::SetupCommandTimedOut { .. } => ( + AgentTaskState::Failed, + TaskStatusUpdate::with_error_code( + error.to_string(), + PlatformErrorCode::EnvironmentSetupFailed, + ), + ), // The shell died while an environment setup command was running // (e.g. the command ran `exit`). This is a user-side environment // configuration problem, so classify as FAILED. diff --git a/app/src/ai/agent_sdk/driver/error_classification_tests.rs b/app/src/ai/agent_sdk/driver/error_classification_tests.rs index e7d88a2fe66..d0d31971291 100644 --- a/app/src/ai/agent_sdk/driver/error_classification_tests.rs +++ b/app/src/ai/agent_sdk/driver/error_classification_tests.rs @@ -6,7 +6,7 @@ use warp_graphql::platform_error::{PlatformErrorInfo, PlatformErrorMessageFormat use super::{classify_driver_error, markdown_code_span}; use crate::ai::agent::{RenderableAIError, TransientNetworkErrorKind}; use crate::ai::agent_sdk::driver::AgentDriverError; -use crate::ai::agent_sdk::driver::environment::PrepareEnvironmentError; +use crate::ai::agent_sdk::driver::environment::{PrepareEnvironmentError, SetupCommandPhase}; use crate::ai::agent_sdk::driver::terminal::{BootstrapError, ShareSessionError}; use crate::server::server_api::ai::TaskGitCredentialsError; @@ -23,6 +23,34 @@ fn assert_state_and_code( ); } +#[test] +fn setup_timeout_reports_failed_with_command_and_duration() { + for (phase, timeout_seconds) in [ + (SetupCommandPhase::Execute, 1800), + (SetupCommandPhase::ResetWorkingDirectory, 30), + ] { + let error = AgentDriverError::from(PrepareEnvironmentError::SetupCommandTimedOut { + command_index: 2, + command: "./setup.sh".to_string(), + phase, + timeout_seconds, + }); + assert!(matches!( + &error, + AgentDriverError::SetupCommandTimedOut { .. } + )); + let (state, update) = classify_driver_error(&error); + assert_eq!(state, AgentTaskState::Failed); + assert_eq!( + update.error_code, + Some(PlatformErrorCode::EnvironmentSetupFailed) + ); + assert!(update.message.contains("Setup command #2")); + assert!(update.message.contains("./setup.sh")); + assert!(update.message.contains(&format!("{timeout_seconds}s"))); + } +} + #[test] fn retryable_dependency_credentials_failure_is_error_with_structured_metadata() { let info = PlatformErrorInfo { diff --git a/app/src/ai/agent_sdk/driver/terminal.rs b/app/src/ai/agent_sdk/driver/terminal.rs index 6bf14ae4815..b403d6294a0 100644 --- a/app/src/ai/agent_sdk/driver/terminal.rs +++ b/app/src/ai/agent_sdk/driver/terminal.rs @@ -9,8 +9,10 @@ use std::time::Duration; use anyhow::Context as _; use futures::channel::oneshot; +use instant::Instant; use session_sharing_protocol::common::{Role, SessionId}; use session_sharing_protocol::sharer::SessionRetentionReason; +use uuid::Uuid; use warp_cli::share::{ShareAccessLevel, ShareRequest, ShareSubject}; use warp_completer::completer::CommandOutput; use warp_core::command::ExitCode; @@ -556,6 +558,12 @@ impl TerminalDriver { } let command_string = command.to_string(); + let command_id = Uuid::new_v4(); + let submitted_at = Instant::now(); + let session_id = self.shared_session_id; + log::info!( + "Terminal command lifecycle: event=submitted command_id={command_id} session_id={session_id:?}" + ); // Store a secret-redacted copy for shell-exit attribution: the text // flows into error reports (server task status, Sentry) if the shell // dies, so never retain the raw command here. @@ -572,9 +580,16 @@ impl TerminalDriver { let block_id = start_rx .await .map_err(|_| AgentDriverError::InvalidRuntimeState)??; + log::info!( + "Terminal command lifecycle: event=started command_id={command_id} session_id={session_id:?} block_id={block_id:?} elapsed_ms={}", + submitted_at.elapsed().as_millis() + ); Ok(CommandHandle { exit_status_rx: exit_rx, block_id, + command_id, + submitted_at, + session_id, }) }) } @@ -773,6 +788,9 @@ pub(crate) struct BlockOutputMatch { pub(crate) struct CommandHandle { exit_status_rx: oneshot::Receiver>, block_id: BlockId, + command_id: Uuid, + submitted_at: Instant, + session_id: Option, } impl CommandHandle { @@ -786,12 +804,24 @@ impl Future for CommandHandle { type Output = Result; fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - Pin::new(&mut self.exit_status_rx) + let result = Pin::new(&mut self.exit_status_rx) .poll(cx) .map(|result| match result { Ok(exit_status) => exit_status, Err(_) => Err(AgentDriverError::InvalidRuntimeState), - }) + }); + if let Poll::Ready(status) = &result { + log::info!( + "Terminal command lifecycle: event=finished command_id={} session_id={:?} block_id={:?} elapsed_ms={} exit_code={:?} exit_status_received={}", + self.command_id, + self.session_id, + self.block_id, + self.submitted_at.elapsed().as_millis(), + status.as_ref().ok().map(|code| code.value()), + status.is_ok() + ); + } + result } } @@ -823,6 +853,12 @@ impl TerminalDriver { } } crate::terminal::view::Event::Exited => { + log::warn!( + "Terminal command lifecycle: event=shell_exited session_id={:?} pending_start={} waiting_exit={}", + self.shared_session_id, + self.pending_command_start.is_some(), + self.waiting_command.is_some() + ); // The shell process exited before bootstrap completed — // cancel the wait immediately rather than sitting out the // full 60 s timeout. No specific reason is known at this diff --git a/app/src/ai/agent_sdk/driver_tests.rs b/app/src/ai/agent_sdk/driver_tests.rs index db277a6f9db..0f9c9284519 100644 --- a/app/src/ai/agent_sdk/driver_tests.rs +++ b/app/src/ai/agent_sdk/driver_tests.rs @@ -10,6 +10,7 @@ use chrono::Local; use cloud_object_models::CodeForge; use futures::channel::oneshot; use futures::executor::block_on; +use futures::poll; use repo_metadata::{DirectoryWatcher, RepoMetadataEvent, RepoMetadataModel, RepositoryIdentifier}; use serde_json::json; use session_sharing_protocol::common::{AgentAttachment, ParticipantId}; @@ -64,6 +65,37 @@ use crate::workspace::view::tests::initialize_app as initialize_workspace_test_a // ── IdleTimeoutSender tests ────────────────────────────────────────────────────── +#[test] +fn setup_timeout_does_not_retain_a_potentially_running_command() { + App::test((), |mut app| async move { + initialize_app_for_terminal_view(&mut app); + let terminal = add_window_with_terminal(&mut app, None); + let temp = TempDir::new().unwrap(); + let driver = app.add_model(|ctx| { + let terminal_driver = + super::terminal::TerminalDriver::create_from_existing_view(terminal, ctx); + let mut driver = + AgentDriver::new_for_test(temp.path().to_path_buf(), terminal_driver, ctx); + driver.idle_on_fail = Some(Duration::from_secs(30 * 60)); + driver + }); + let spawner = driver.update(&mut app, |_, ctx| ctx.spawner()); + let error = AgentDriverError::SetupCommandTimedOut { + message: "Setup command #1 timed out after 1800s: ./setup.sh".to_string(), + }; + let mut linger = Box::pin(AgentDriver::linger_after_failure( + &spawner, + "environment_setup", + &error, + )); + assert!(poll!(linger.as_mut()).is_ready()); + assert_eq!( + setup_failure_status_update(&error).error_code, + Some(PlatformErrorCode::EnvironmentSetupFailed) + ); + }); +} + #[test] fn driver_keeps_uninterpreted_factory_experiments() { App::test((), |mut app| async move { diff --git a/app/src/ai/blocklist/action_model/execute/call_mcp_tool.rs b/app/src/ai/blocklist/action_model/execute/call_mcp_tool.rs index a6b967dacf4..279d6b65541 100644 --- a/app/src/ai/blocklist/action_model/execute/call_mcp_tool.rs +++ b/app/src/ai/blocklist/action_model/execute/call_mcp_tool.rs @@ -3,6 +3,8 @@ use futures::future::BoxFuture; #[cfg(not(target_family = "wasm"))] use itertools::Itertools; #[cfg(not(target_family = "wasm"))] +use uuid::Uuid; +#[cfg(not(target_family = "wasm"))] use warpui::SingletonEntity; use warpui::{Entity, EntityId, ModelContext, ModelHandle}; @@ -87,6 +89,8 @@ impl CallMCPToolExecutor { #[cfg(not(target_family = "wasm"))] { let server_output_id = get_server_output_id(input.conversation_id, ctx); + let conversation_id = input.conversation_id; + let action_id = input.action.id.clone(); let AIAgentAction { action: AIAgentActionType::CallMCPTool { @@ -140,16 +144,27 @@ impl CallMCPToolExecutor { }; let name_owned_inner = name_owned.clone(); + let operation_id = Uuid::new_v4(); + log::info!( + "MCP tool action: event=started operation_id={operation_id} conversation_id={conversation_id:?} action_id={action_id:?} server_id={server_id:?} tool_name={name_owned:?}" + ); ActionExecution::new_async( async move { reconnecting_peer .call_tool( rmcp::model::CallToolRequestParams::new(name_owned_inner) .with_arguments(arguments), + operation_id, ) .await }, - move |res, ctx| handle_call_tool_result(res, server_output_id, name_clone, ctx), + move |res, ctx| { + let result = handle_call_tool_result(res, server_output_id, name_clone, ctx); + log::info!( + "MCP tool action: event=result_applied operation_id={operation_id} conversation_id={conversation_id:?} action_id={action_id:?}" + ); + result + }, ) } } diff --git a/app/src/ai/mcp/reconnecting_peer.rs b/app/src/ai/mcp/reconnecting_peer.rs index 992c032e400..fd3083d093d 100644 --- a/app/src/ai/mcp/reconnecting_peer.rs +++ b/app/src/ai/mcp/reconnecting_peer.rs @@ -2,6 +2,7 @@ use std::future::Future; +use mcp::tool_call::call_tool_with_deadline; use uuid::Uuid; use warpui::ModelSpawner; @@ -9,9 +10,8 @@ use super::TemplatableMCPServerManager; /// A wrapper around an MCP server connection that transparently handles reconnection. /// -/// When making requests (e.g., `call_tool` or `read_resource`), this type checks if the -/// underlying transport is closed and automatically triggers reconnection before retrying -/// the request. +/// Reconnects before dispatching requests. Only resource reads are retried after a +/// transport failure; a tool may have executed even when its response was lost. #[derive(Clone)] pub struct ReconnectingPeer { installation_uuid: Uuid, @@ -121,9 +121,12 @@ impl ReconnectingPeer { pub async fn call_tool( &self, params: rmcp::model::CallToolRequestParams, + operation_id: Uuid, ) -> Result { - self.with_reconnect_retry(params, |peer, p| async move { peer.call_tool(p).await }) - .await + call_tool_with_deadline(operation_id, self.installation_uuid, params, async { + self.get_connected_peer().await.map_err(Into::into) + }) + .await } /// Reads a resource from the MCP server. diff --git a/crates/mcp/Cargo.toml b/crates/mcp/Cargo.toml index ea5a1e9ccda..6cb3d109675 100644 --- a/crates/mcp/Cargo.toml +++ b/crates/mcp/Cargo.toml @@ -14,6 +14,7 @@ cloud_object_models.workspace = true futures.workspace = true handlebars.workspace = true http = "1" +instant.workspace = true log.workspace = true oauth2.workspace = true pin-project-lite = "0.2" diff --git a/crates/mcp/src/lib.rs b/crates/mcp/src/lib.rs index 90d345f5b76..8e2f6db71e8 100644 --- a/crates/mcp/src/lib.rs +++ b/crates/mcp/src/lib.rs @@ -4,6 +4,8 @@ pub mod oauth; pub mod runtime; #[cfg(not(target_family = "wasm"))] pub mod sse_transport; +#[cfg(not(target_family = "wasm"))] +pub mod tool_call; use uuid::Uuid; diff --git a/crates/mcp/src/tool_call.rs b/crates/mcp/src/tool_call.rs new file mode 100644 index 00000000000..85d45c1f0fe --- /dev/null +++ b/crates/mcp/src/tool_call.rs @@ -0,0 +1,164 @@ +use std::future::Future; +use std::time::Duration; + +use futures::FutureExt; +use futures::future::{Either, select}; +use instant::Instant; +use rmcp::model::{ + CallToolRequest, CallToolRequestParams, CallToolResult, CancelledNotificationParam, + ClientRequest, ErrorData, RequestId, ServerResult, +}; +use rmcp::service::PeerRequestOptions; +use rmcp::{Peer, RoleClient, ServiceError}; +use uuid::Uuid; +use warpui::r#async::Timer; + +/// Total connection, dispatch, and response budget for one tool call. +pub const TOOL_CALL_TIMEOUT: Duration = Duration::from_secs(30 * 60); +const CANCELLATION_TIMEOUT: Duration = Duration::from_secs(5); + +/// Calls a tool once, preserving transport resumption but bounding the logical request. +pub async fn call_tool_with_deadline( + operation_id: Uuid, + installation_id: Uuid, + params: CallToolRequestParams, + connect: impl Future, ServiceError>>, +) -> Result { + call_tool_with_deadline_inner( + operation_id, + installation_id, + params, + connect, + Timer::after(TOOL_CALL_TIMEOUT).map(|_| ()), + || Timer::after(CANCELLATION_TIMEOUT).map(|_| ()), + ) + .await +} + +async fn call_tool_with_deadline_inner( + operation_id: Uuid, + installation_id: Uuid, + params: CallToolRequestParams, + connect: Connect, + deadline: Deadline, + cancellation_deadline: CancellationDeadline, +) -> Result +where + Connect: Future, ServiceError>>, + Deadline: Future, + CancellationDeadline: FnOnce() -> CancellationTimer, + CancellationTimer: Future, +{ + let mut lifecycle = ToolCallLifecycle { + operation_id, + installation_id, + started_at: Instant::now(), + request_id: None, + phase: "connecting", + outcome: "pending", + }; + lifecycle.log("started"); + let mut cancellation = None; + + let completed = { + let operation = async { + let peer = connect.await?; + lifecycle.phase = "dispatching"; + lifecycle.log("peer_ready"); + let request = ClientRequest::CallToolRequest(CallToolRequest::new(params)); + // rmcp's own timeout awaits cancellation delivery before returning. Keep the + // logical deadline outside that path so a stalled send cannot defeat it. + let handle = peer + .send_cancellable_request(request, PeerRequestOptions::no_options()) + .await?; + lifecycle.request_id = Some(handle.id.clone()); + cancellation = Some((peer, handle.id.clone())); + lifecycle.phase = "awaiting_response"; + lifecycle.log("request_enqueued"); + match handle.await_response().await? { + ServerResult::CallToolResult(result) => Ok(result), + _ => Err(ServiceError::UnexpectedResponse), + } + }; + match select(Box::pin(operation), Box::pin(deadline)).await { + Either::Left((result, _)) => Some(result), + Either::Right(((), _)) => None, + } + }; + + if let Some(result) = completed { + lifecycle.outcome = match &result { + Ok(result) if result.is_error == Some(true) => "tool_error", + Ok(_) => "success", + Err(_) => "service_error", + }; + return result; + } + + lifecycle.outcome = "deadline_exceeded"; + lifecycle.log("deadline"); + if let Some((peer, request_id)) = cancellation { + let cancel = peer.notify_cancelled(CancelledNotificationParam::new( + Some(request_id), + Some("Warp MCP tool call deadline exceeded".to_owned()), + )); + let cancellation_outcome = + match select(Box::pin(cancel), Box::pin(cancellation_deadline())).await { + Either::Left((Ok(()), _)) => "cancellation_sent", + Either::Left((Err(_), _)) => "cancellation_failed", + Either::Right(((), _)) => "cancellation_deadline_exceeded", + }; + lifecycle.log(cancellation_outcome); + } + + let message = if lifecycle.phase == "connecting" { + format!( + "MCP tool call timed out after {} seconds while connecting. The tool was not dispatched.", + TOOL_CALL_TIMEOUT.as_secs() + ) + } else { + format!( + "MCP tool call deadline exceeded after {} seconds without a result. The tool may have executed; its outcome is unknown. Do not retry automatically.", + TOOL_CALL_TIMEOUT.as_secs() + ) + }; + Err(ServiceError::McpError(ErrorData::internal_error( + message, None, + ))) +} + +struct ToolCallLifecycle { + operation_id: Uuid, + installation_id: Uuid, + started_at: Instant, + request_id: Option, + phase: &'static str, + outcome: &'static str, +} + +impl ToolCallLifecycle { + fn log(&self, event: &str) { + log::info!( + "MCP tool lifecycle: event={event} operation_id={} installation_id={} request_id={:?} phase={} elapsed_ms={} outcome={}", + self.operation_id, + self.installation_id, + self.request_id, + self.phase, + self.started_at.elapsed().as_millis(), + self.outcome, + ); + } +} + +impl Drop for ToolCallLifecycle { + fn drop(&mut self) { + if self.outcome == "pending" { + self.outcome = "dropped"; + } + self.log("finished"); + } +} + +#[cfg(test)] +#[path = "tool_call_tests.rs"] +mod tests; diff --git a/crates/mcp/src/tool_call_tests.rs b/crates/mcp/src/tool_call_tests.rs new file mode 100644 index 00000000000..84cee7204dc --- /dev/null +++ b/crates/mcp/src/tool_call_tests.rs @@ -0,0 +1,411 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use axum::body::Body; +use axum::extract::State; +use axum::http::{HeaderMap, StatusCode}; +use axum::response::{IntoResponse, Response}; +use axum::routing::post; +use axum::{Json, Router}; +use futures::future::{pending, ready}; +use rmcp::model::{CallToolRequestParams, ErrorData, ProgressNotificationParam}; +use rmcp::service::NotificationContext; +use rmcp::transport::StreamableHttpClientTransport; +use rmcp::{ClientHandler, RoleClient, ServiceError, ServiceExt}; +use serde_json::{Value, json}; +use tokio::sync::{Notify, oneshot}; +use tokio::task::JoinHandle; +use uuid::Uuid; + +use super::call_tool_with_deadline_inner; + +#[derive(Clone, Copy)] +enum Reply { + Success, + ToolError, + ProtocolError, + HttpError, + WrongId, + IncompleteSse, + ResumableSse, + StuckSend, +} + +struct ServerState { + reply: Reply, + calls: AtomicUsize, + cancellations: AtomicUsize, + resumes: AtomicUsize, + request_seen: Notify, + request_id: Mutex>, +} + +struct TestServer { + uri: String, + state: Arc, + task: JoinHandle<()>, +} +struct ProgressObserver(Arc); + +impl ClientHandler for ProgressObserver { + async fn on_progress(&self, _: ProgressNotificationParam, _: NotificationContext) { + self.0.notify_one(); + } +} + +impl Drop for TestServer { + fn drop(&mut self) { + self.task.abort(); + } +} + +async fn start_server(reply: Reply) -> TestServer { + let state = Arc::new(ServerState { + reply, + calls: AtomicUsize::new(0), + cancellations: AtomicUsize::new(0), + resumes: AtomicUsize::new(0), + request_seen: Notify::new(), + request_id: Mutex::new(None), + }); + let router = Router::new() + .route( + "/mcp", + post(handle_post) + .get(handle_get) + .delete(|| async { StatusCode::OK }), + ) + .with_state(state.clone()); + let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)) + .await + .unwrap(); + let uri = format!("http://{}/mcp", listener.local_addr().unwrap()); + let task = tokio::spawn(async move { + axum::serve(listener, router).await.unwrap(); + }); + TestServer { uri, state, task } +} + +fn result(id: Value, is_error: bool) -> Value { + json!({ + "jsonrpc": "2.0", + "id": id, + "result": {"content": [], "isError": is_error} + }) +} + +fn sse(body: String) -> Response { + ([("content-type", "text/event-stream")], Body::from(body)).into_response() +} + +async fn handle_post( + State(state): State>, + Json(request): Json, +) -> Response { + match request["method"].as_str().unwrap() { + "initialize" => { + let mut response = Json(json!({ + "jsonrpc": "2.0", + "id": request["id"], + "result": { + "protocolVersion": "2025-03-26", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "deadline-test", "version": "1"} + } + })) + .into_response(); + if matches!(state.reply, Reply::ResumableSse) { + response + .headers_mut() + .insert("mcp-session-id", "test-session".parse().unwrap()); + } + response + } + "notifications/initialized" => StatusCode::ACCEPTED.into_response(), + "notifications/cancelled" => { + state.cancellations.fetch_add(1, Ordering::SeqCst); + StatusCode::ACCEPTED.into_response() + } + "tools/call" => { + state.calls.fetch_add(1, Ordering::SeqCst); + *state.request_id.lock().unwrap() = Some(request["id"].clone()); + state.request_seen.notify_one(); + if request["params"]["name"] == "succeed" { + return Json(result(request["id"].clone(), false)).into_response(); + } + match state.reply { + Reply::Success => Json(result(request["id"].clone(), false)).into_response(), + Reply::ToolError => Json(result(request["id"].clone(), true)).into_response(), + Reply::ProtocolError => Json(json!({ + "jsonrpc": "2.0", + "id": request["id"], + "error": {"code": -32602, "message": "test protocol error"} + })) + .into_response(), + Reply::HttpError => StatusCode::BAD_GATEWAY.into_response(), + Reply::WrongId => Json(result(json!("unrelated"), false)).into_response(), + Reply::IncompleteSse => sse(format!( + "data: {}\n\nevent: message\ndata: {{\"jsonrpc\":\"2.0\"", + progress(&request), + )), + Reply::ResumableSse => sse(format!( + "id: resume-1\nretry: 0\ndata: {}\n\n", + progress(&request), + )), + Reply::StuckSend => pending().await, + } + } + _ => StatusCode::BAD_REQUEST.into_response(), + } +} + +fn progress(request: &Value) -> Value { + json!({ + "jsonrpc": "2.0", + "method": "notifications/progress", + "params": { + "progressToken": request["params"]["_meta"]["progressToken"], + "progress": 1 + } + }) +} + +async fn handle_get(State(state): State>, headers: HeaderMap) -> Response { + if headers.get("last-event-id").is_none() { + return StatusCode::METHOD_NOT_ALLOWED.into_response(); + } + state.resumes.fetch_add(1, Ordering::SeqCst); + let id = state.request_id.lock().unwrap().clone().unwrap(); + sse(format!("data: {}\n\n", result(id, false))) +} + +#[tokio::test] +async fn deadline_includes_connecting_without_dispatch() { + let result = call_tool_with_deadline_inner( + Uuid::nil(), + Uuid::nil(), + CallToolRequestParams::new("test"), + pending(), + ready(()), + || ready(()), + ) + .await; + assert!(result.unwrap_err().to_string().contains("not dispatched")); +} + +#[tokio::test] +async fn connection_error_is_preserved() { + let result = call_tool_with_deadline_inner( + Uuid::nil(), + Uuid::nil(), + CallToolRequestParams::new("test"), + ready(Err(ServiceError::McpError(ErrorData::internal_error( + "test connection error", + None, + )))), + pending(), + || ready(()), + ) + .await; + assert!( + result + .unwrap_err() + .to_string() + .contains("test connection error") + ); +} + +#[tokio::test] +async fn successful_and_error_results_complete_without_retry() { + for reply in [ + Reply::Success, + Reply::ToolError, + Reply::ProtocolError, + Reply::HttpError, + ] { + let server = start_server(reply).await; + let service = + ().serve(StreamableHttpClientTransport::from_uri(server.uri.clone())) + .await + .unwrap(); + let result = tokio::time::timeout( + Duration::from_secs(5), + call_tool_with_deadline_inner( + Uuid::nil(), + Uuid::nil(), + CallToolRequestParams::new("test"), + ready(Ok(service.peer().clone())), + pending(), + || ready(()), + ), + ) + .await + .unwrap(); + match reply { + Reply::Success => assert_eq!(result.unwrap().is_error, Some(false)), + Reply::ToolError => assert_eq!(result.unwrap().is_error, Some(true)), + Reply::ProtocolError => assert!( + result + .unwrap_err() + .to_string() + .contains("test protocol error") + ), + Reply::HttpError => assert!(result.is_err()), + _ => unreachable!(), + } + assert_eq!(server.state.calls.load(Ordering::SeqCst), 1); + assert_eq!(server.state.cancellations.load(Ordering::SeqCst), 0); + service.cancel().await.unwrap(); + } +} + +#[tokio::test] +async fn incomplete_sse_and_wrong_response_id_return_outcome_unknown() { + for reply in [Reply::IncompleteSse, Reply::WrongId] { + let server = start_server(reply).await; + let progress_seen = Arc::new(Notify::new()); + let service = ProgressObserver(progress_seen.clone()) + .serve(StreamableHttpClientTransport::from_uri(server.uri.clone())) + .await + .unwrap(); + let peer = service.peer().clone(); + let (expire, deadline) = oneshot::channel(); + let task = tokio::spawn(call_tool_with_deadline_inner( + Uuid::nil(), + Uuid::nil(), + CallToolRequestParams::new("test"), + ready(Ok(peer)), + async { deadline.await.unwrap() }, + pending, + )); + server.state.request_seen.notified().await; + if matches!(reply, Reply::IncompleteSse) { + tokio::time::timeout(Duration::from_secs(5), progress_seen.notified()) + .await + .unwrap(); + } + assert!(!task.is_finished()); + expire.send(()).unwrap(); + let result = tokio::time::timeout(Duration::from_secs(5), task) + .await + .unwrap() + .unwrap(); + assert!( + result + .unwrap_err() + .to_string() + .contains("outcome is unknown") + ); + assert_eq!(server.state.calls.load(Ordering::SeqCst), 1); + assert_eq!(server.state.cancellations.load(Ordering::SeqCst), 1); + service.cancel().await.unwrap(); + } +} + +#[tokio::test] +async fn timed_out_request_does_not_close_the_connection_or_replay_the_tool() { + let server = start_server(Reply::IncompleteSse).await; + let service = + ().serve(StreamableHttpClientTransport::from_uri(server.uri.clone())) + .await + .unwrap(); + let peer = service.peer().clone(); + let (expire, deadline) = oneshot::channel(); + let task = tokio::spawn(call_tool_with_deadline_inner( + Uuid::nil(), + Uuid::nil(), + CallToolRequestParams::new("test"), + ready(Ok(peer.clone())), + async { deadline.await.unwrap() }, + pending, + )); + server.state.request_seen.notified().await; + let succeeding = call_tool_with_deadline_inner( + Uuid::nil(), + Uuid::nil(), + CallToolRequestParams::new("succeed"), + ready(Ok(peer)), + pending(), + pending, + ); + assert!( + tokio::time::timeout(Duration::from_secs(5), succeeding) + .await + .unwrap() + .is_ok() + ); + expire.send(()).unwrap(); + assert!( + tokio::time::timeout(Duration::from_secs(5), task) + .await + .unwrap() + .unwrap() + .is_err() + ); + assert_eq!(server.state.calls.load(Ordering::SeqCst), 2); + assert_eq!(server.state.cancellations.load(Ordering::SeqCst), 1); + assert!(!service.is_transport_closed()); + service.cancel().await.unwrap(); +} + +#[tokio::test] +async fn stalled_send_cannot_block_deadline_or_cancellation() { + let server = start_server(Reply::StuckSend).await; + let service = + ().serve(StreamableHttpClientTransport::from_uri(server.uri.clone())) + .await + .unwrap(); + let peer = service.peer().clone(); + let (expire, deadline) = oneshot::channel(); + let task = tokio::spawn(call_tool_with_deadline_inner( + Uuid::nil(), + Uuid::nil(), + CallToolRequestParams::new("test"), + ready(Ok(peer)), + async { deadline.await.unwrap() }, + || ready(()), + )); + server.state.request_seen.notified().await; + expire.send(()).unwrap(); + let result = tokio::time::timeout(Duration::from_secs(5), task) + .await + .unwrap() + .unwrap(); + assert!( + result + .unwrap_err() + .to_string() + .contains("outcome is unknown") + ); + assert_eq!(server.state.calls.load(Ordering::SeqCst), 1); + assert_eq!(server.state.cancellations.load(Ordering::SeqCst), 0); + service.cancellation_token().cancel(); +} + +#[tokio::test] +async fn resumable_sse_keeps_waiting_for_the_matching_result() { + let server = start_server(Reply::ResumableSse).await; + let service = + ().serve(StreamableHttpClientTransport::from_uri(server.uri.clone())) + .await + .unwrap(); + let result = tokio::time::timeout( + Duration::from_secs(5), + call_tool_with_deadline_inner( + Uuid::nil(), + Uuid::nil(), + CallToolRequestParams::new("test"), + ready(Ok(service.peer().clone())), + pending(), + || ready(()), + ), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(result.is_error, Some(false)); + assert_eq!(server.state.calls.load(Ordering::SeqCst), 1); + assert_eq!(server.state.resumes.load(Ordering::SeqCst), 1); + service.cancel().await.unwrap(); +} From ba70b80f3b45407ce5e392a9334aae40c0a04b4d Mon Sep 17 00:00:00 2001 From: seemeroland Date: Wed, 30 Sep 2026 17:04:44 -0700 Subject: [PATCH 2/3] bump rmcp version to pick up fix --- Cargo.lock | 66 ++++++++++-- Cargo.toml | 2 +- app/src/ai/agent_sdk/driver/environment.rs | 2 +- crates/mcp/src/oauth.rs | 8 +- crates/mcp/src/runtime.rs | 9 +- crates/mcp/src/tool_call_tests.rs | 112 ++++++++++++--------- 6 files changed, 134 insertions(+), 65 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index de3d7fa523a..793177a931c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4018,6 +4018,16 @@ dependencies = [ "darling_macro 0.23.0", ] +[[package]] +name = "darling" +version = "0.24.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed17f5901b6630b993ca003def43f2f8ef4014fc13b047b57aad617ff32bc2ec" +dependencies = [ + "darling_core 0.24.1", + "darling_macro 0.24.1", +] + [[package]] name = "darling_core" version = "0.20.8" @@ -4059,6 +4069,19 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "darling_core" +version = "0.24.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6837e2cf7485aaae18f86181d2f0e9a7ed297a025e220aeabf63fdebd3a2ddff" +dependencies = [ + "ident_case", + "proc-macro2", + "quote", + "strsim 0.11.1", + "syn 3.0.6", +] + [[package]] name = "darling_macro" version = "0.20.8" @@ -4092,6 +4115,17 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "darling_macro" +version = "0.24.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ac7135c3ef02b2f7833bbeb1be5ba7f966dcde8a87c6b87f65a778d71a02785" +dependencies = [ + "darling_core 0.24.1", + "quote", + "syn 3.0.6", +] + [[package]] name = "dary_heap" version = "0.3.7" @@ -10572,9 +10606,9 @@ dependencies = [ [[package]] name = "process-wrap" -version = "9.1.0" +version = "10.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e842efad9119158434d193c6682e2ebee4b44d6ad801d7b349623b3f57cdf55" +checksum = "1f21b97672d2dc848e7b25701ab4535618b92f4861c13cc3f7f7bed52ad3c8da" dependencies = [ "futures", "indexmap 2.12.0", @@ -11639,15 +11673,17 @@ dependencies = [ [[package]] name = "rmcp" -version = "2.2.0" +version = "3.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "14db48ee17a9ba61810ab1a9c1beb7d06d8136ae39ac25a1137f10d357af01af" +checksum = "fae7019994ae0fe4ada40b732f798f3ff26f0f04facb1477f1bf37eb4f18a2d3" dependencies = [ "async-trait", - "base64 0.22.1", + "base64 0.23.1", + "bytes", "chrono", "futures", "http 1.4.0", + "indexmap 2.12.0", "oauth2", "pastey 0.2.1", "pin-project-lite", @@ -11664,19 +11700,20 @@ dependencies = [ "tokio-util", "tracing", "url", + "uuid", ] [[package]] name = "rmcp-macros" -version = "2.2.0" +version = "3.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "783d787bf21813b285f13019adc49e11af501c658890c1e519f31f937c68b7e3" +checksum = "c7e66877136e969d5a2d561ee5841b17717a0619aebd64abb47a2a2f46e34683" dependencies = [ - "darling 0.23.0", + "darling 0.24.1", "proc-macro2", "quote", "serde_json", - "syn 2.0.117", + "syn 3.0.6", ] [[package]] @@ -13333,6 +13370,17 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "syn" +version = "3.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8593e8e72159ed2257d083c7a454a85cbf854f37a0966d8d483aff8c8a3ebcee" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "syn_derive" version = "0.1.8" diff --git a/Cargo.toml b/Cargo.toml index 1b59f01736a..efb42a437ec 100644 --- a/Cargo.toml +++ b/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 / diff --git a/app/src/ai/agent_sdk/driver/environment.rs b/app/src/ai/agent_sdk/driver/environment.rs index c2622b6934b..2b209a35f8b 100644 --- a/app/src/ai/agent_sdk/driver/environment.rs +++ b/app/src/ai/agent_sdk/driver/environment.rs @@ -41,7 +41,7 @@ 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_secs(30 * 60); +const SETUP_COMMAND_TIMEOUT: Duration = Duration::from_secs(10); const SETUP_COMMAND_CWD_RESET_TIMEOUT: Duration = Duration::from_secs(30); #[derive(Clone, Copy, Debug)] diff --git a/crates/mcp/src/oauth.rs b/crates/mcp/src/oauth.rs index a9566fde23e..99312d8d235 100644 --- a/crates/mcp/src/oauth.rs +++ b/crates/mcp/src/oauth.rs @@ -3,8 +3,8 @@ use std::collections::HashMap; use futures::future::BoxFuture; use oauth2::{RefreshToken, TokenResponse as _}; use rmcp::transport::auth::{ - AuthClient, AuthorizationManager, CredentialStore, InMemoryCredentialStore, OAuthClientConfig, - OAuthState, StoredCredentials, + AuthClient, AuthorizationManager, AuthorizationRequest, CredentialStore, + InMemoryCredentialStore, OAuthClientConfig, OAuthState, StoredCredentials, }; use rmcp::transport::{AuthError, AuthorizationSession}; use serde::de::DeserializeOwned; @@ -335,7 +335,7 @@ pub async fn make_authenticated_client( )); } - let metadata = auth_manager.discover_metadata().await?; + let metadata = auth_manager.resolve_metadata().await?.metadata; // Configure the auth manager's OAuth client using dynamic or static client registration. let mut oauth_state = if let Some(provider) = metadata @@ -387,7 +387,7 @@ pub async fn make_authenticated_client( // Try dynamic client registration. let mut oauth_state = OAuthState::Unauthorized(auth_manager); oauth_state - .start_authorization(&[], &redirect_uri, Some("Warp")) + .start_authorization(AuthorizationRequest::new(&redirect_uri).with_client_name("Warp")) .await?; oauth_state }; diff --git a/crates/mcp/src/runtime.rs b/crates/mcp/src/runtime.rs index abca3143447..2ea83a9bb63 100644 --- a/crates/mcp/src/runtime.rs +++ b/crates/mcp/src/runtime.rs @@ -12,6 +12,7 @@ use cfg_if::cfg_if; use cloud_object_models::{StaticEnvVar, TransportType}; use futures::FutureExt as _; use rmcp::ServiceExt as _; +use rmcp::model::{ClientConfig, Implementation}; use rmcp::transport::ConfigureCommandExt as _; use simple_logger::SimpleLogger; use tokio::io::AsyncBufReadExt as _; @@ -585,13 +586,13 @@ async fn send_initialize_request( }) } -/// Creates a [`ClientInfo`] for the MCP client. +/// Creates a [`ClientConfig`] for the MCP client. /// /// This tells the MCP server who we are and what capabilities we have. -fn make_client_info() -> rmcp::model::ClientInfo { - rmcp::model::ClientInfo::new( +fn make_client_info() -> ClientConfig { + ClientConfig::new( Default::default(), - rmcp::model::Implementation::new( + Implementation::new( warp_core::channel::ChannelState::app_id().to_string(), warp_core::channel::ChannelState::app_version() .map(|v| v.to_string()) diff --git a/crates/mcp/src/tool_call_tests.rs b/crates/mcp/src/tool_call_tests.rs index 84cee7204dc..c8abd65ef65 100644 --- a/crates/mcp/src/tool_call_tests.rs +++ b/crates/mcp/src/tool_call_tests.rs @@ -9,10 +9,9 @@ use axum::response::{IntoResponse, Response}; use axum::routing::post; use axum::{Json, Router}; use futures::future::{pending, ready}; -use rmcp::model::{CallToolRequestParams, ErrorData, ProgressNotificationParam}; -use rmcp::service::NotificationContext; +use rmcp::model::{CallToolRequestParams, ErrorData}; use rmcp::transport::StreamableHttpClientTransport; -use rmcp::{ClientHandler, RoleClient, ServiceError, ServiceExt}; +use rmcp::{ServiceError, ServiceExt}; use serde_json::{Value, json}; use tokio::sync::{Notify, oneshot}; use tokio::task::JoinHandle; @@ -23,10 +22,12 @@ use super::call_tool_with_deadline_inner; #[derive(Clone, Copy)] enum Reply { Success, + SseSuccess, ToolError, ProtocolError, HttpError, WrongId, + EmptySse, IncompleteSse, ResumableSse, StuckSend, @@ -46,13 +47,6 @@ struct TestServer { state: Arc, task: JoinHandle<()>, } -struct ProgressObserver(Arc); - -impl ClientHandler for ProgressObserver { - async fn on_progress(&self, _: ProgressNotificationParam, _: NotificationContext) { - self.0.notify_one(); - } -} impl Drop for TestServer { fn drop(&mut self) { @@ -136,6 +130,10 @@ async fn handle_post( } match state.reply { Reply::Success => Json(result(request["id"].clone(), false)).into_response(), + Reply::SseSuccess => sse(format!( + "data: {}\n\n", + result(request["id"].clone(), false) + )), Reply::ToolError => Json(result(request["id"].clone(), true)).into_response(), Reply::ProtocolError => Json(json!({ "jsonrpc": "2.0", @@ -145,6 +143,7 @@ async fn handle_post( .into_response(), Reply::HttpError => StatusCode::BAD_GATEWAY.into_response(), Reply::WrongId => Json(result(json!("unrelated"), false)).into_response(), + Reply::EmptySse => sse(String::new()), Reply::IncompleteSse => sse(format!( "data: {}\n\nevent: message\ndata: {{\"jsonrpc\":\"2.0\"", progress(&request), @@ -220,6 +219,7 @@ async fn connection_error_is_preserved() { async fn successful_and_error_results_complete_without_retry() { for reply in [ Reply::Success, + Reply::SseSuccess, Reply::ToolError, Reply::ProtocolError, Reply::HttpError, @@ -243,7 +243,9 @@ async fn successful_and_error_results_complete_without_retry() { .await .unwrap(); match reply { - Reply::Success => assert_eq!(result.unwrap().is_error, Some(false)), + Reply::Success | Reply::SseSuccess => { + assert_eq!(result.unwrap().is_error, Some(false)) + } Reply::ToolError => assert_eq!(result.unwrap().is_error, Some(true)), Reply::ProtocolError => assert!( result @@ -259,53 +261,71 @@ async fn successful_and_error_results_complete_without_retry() { service.cancel().await.unwrap(); } } - #[tokio::test] -async fn incomplete_sse_and_wrong_response_id_return_outcome_unknown() { - for reply in [Reply::IncompleteSse, Reply::WrongId] { - let server = start_server(reply).await; - let progress_seen = Arc::new(Notify::new()); - let service = ProgressObserver(progress_seen.clone()) - .serve(StreamableHttpClientTransport::from_uri(server.uri.clone())) +async fn wrong_response_id_returns_outcome_unknown_at_deadline() { + let server = start_server(Reply::WrongId).await; + let service = + ().serve(StreamableHttpClientTransport::from_uri(server.uri.clone())) .await .unwrap(); - let peer = service.peer().clone(); - let (expire, deadline) = oneshot::channel(); - let task = tokio::spawn(call_tool_with_deadline_inner( - Uuid::nil(), - Uuid::nil(), - CallToolRequestParams::new("test"), - ready(Ok(peer)), - async { deadline.await.unwrap() }, - pending, - )); - server.state.request_seen.notified().await; - if matches!(reply, Reply::IncompleteSse) { - tokio::time::timeout(Duration::from_secs(5), progress_seen.notified()) + let (expire, deadline) = oneshot::channel(); + let task = tokio::spawn(call_tool_with_deadline_inner( + Uuid::nil(), + Uuid::nil(), + CallToolRequestParams::new("test"), + ready(Ok(service.peer().clone())), + async { deadline.await.unwrap() }, + pending, + )); + server.state.request_seen.notified().await; + expire.send(()).unwrap(); + let result = tokio::time::timeout(Duration::from_secs(5), task) + .await + .unwrap() + .unwrap(); + assert!( + result + .unwrap_err() + .to_string() + .contains("outcome is unknown") + ); + assert_eq!(server.state.calls.load(Ordering::SeqCst), 1); + assert_eq!(server.state.cancellations.load(Ordering::SeqCst), 1); + service.cancel().await.unwrap(); +} + +#[tokio::test] +async fn ended_sse_without_result_fails_without_waiting_for_deadline() { + for reply in [Reply::EmptySse, Reply::IncompleteSse] { + let server = start_server(reply).await; + let service = + ().serve(StreamableHttpClientTransport::from_uri(server.uri.clone())) .await .unwrap(); - } - assert!(!task.is_finished()); - expire.send(()).unwrap(); - let result = tokio::time::timeout(Duration::from_secs(5), task) - .await - .unwrap() - .unwrap(); - assert!( - result - .unwrap_err() - .to_string() - .contains("outcome is unknown") - ); + let result = tokio::time::timeout( + Duration::from_secs(5), + call_tool_with_deadline_inner( + Uuid::nil(), + Uuid::nil(), + CallToolRequestParams::new("test"), + ready(Ok(service.peer().clone())), + pending(), + pending, + ), + ) + .await + .expect("a terminated response must not leave the tool call pending"); + assert!(matches!(result, Err(ServiceError::TransportClosed))); assert_eq!(server.state.calls.load(Ordering::SeqCst), 1); - assert_eq!(server.state.cancellations.load(Ordering::SeqCst), 1); + assert_eq!(server.state.cancellations.load(Ordering::SeqCst), 0); + assert!(!service.is_transport_closed()); service.cancel().await.unwrap(); } } #[tokio::test] async fn timed_out_request_does_not_close_the_connection_or_replay_the_tool() { - let server = start_server(Reply::IncompleteSse).await; + let server = start_server(Reply::WrongId).await; let service = ().serve(StreamableHttpClientTransport::from_uri(server.uri.clone())) .await From c48fd6db65a29fb2e4feb6419c70b0aeaac95e75 Mon Sep 17 00:00:00 2001 From: seemeroland Date: Wed, 30 Sep 2026 17:14:41 -0700 Subject: [PATCH 3/3] Undo testing change, and update message --- app/src/ai/agent_sdk/driver/environment.rs | 4 ++-- app/src/ai/agent_sdk/driver/environment_tests.rs | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/app/src/ai/agent_sdk/driver/environment.rs b/app/src/ai/agent_sdk/driver/environment.rs index 2b209a35f8b..5d7e6484dd5 100644 --- a/app/src/ai/agent_sdk/driver/environment.rs +++ b/app/src/ai/agent_sdk/driver/environment.rs @@ -41,7 +41,7 @@ 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_secs(10); +const SETUP_COMMAND_TIMEOUT: Duration = Duration::from_mins(30); const SETUP_COMMAND_CWD_RESET_TIMEOUT: Duration = Duration::from_secs(30); #[derive(Clone, Copy, Debug)] @@ -62,7 +62,7 @@ impl SetupCommandPhase { impl fmt::Display for SetupCommandPhase { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.write_str(match self { - Self::Execute => "waiting for command startup or completion", + Self::Execute => "waiting for command to complete", Self::ResetWorkingDirectory => "resetting the working directory after the command", }) } diff --git a/app/src/ai/agent_sdk/driver/environment_tests.rs b/app/src/ai/agent_sdk/driver/environment_tests.rs index 67eefc66abf..0595118dffd 100644 --- a/app/src/ai/agent_sdk/driver/environment_tests.rs +++ b/app/src/ai/agent_sdk/driver/environment_tests.rs @@ -67,7 +67,7 @@ fn setup_timeout_covers_command_start_and_exit_without_resetting_deadline() { let error = wait.await.unwrap_err(); assert_eq!( error.to_string(), - "Setup command #2 timed out after 1800s while waiting for command startup or completion: `./setup.sh`" + "Setup command #2 timed out after 1800s while waiting for command to complete: `./setup.sh`" ); if let Some(start_tx) = start_tx { assert!(start_tx.is_canceled());