diff --git a/codex-rs/rmcp-client/tests/streamable_http_remote.rs b/codex-rs/rmcp-client/tests/streamable_http_remote.rs index 5cfdf4be86..070b3c815d 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_remote.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_remote.rs @@ -9,10 +9,12 @@ mod streamable_http_test_support; use std::sync::Arc; use std::time::Duration; +use anyhow::Context; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; use codex_rmcp_client::ElicitationAction; use codex_rmcp_client::ElicitationResponse; +use codex_rmcp_client::McpProtocolMode; use codex_rmcp_client::RmcpClient; use futures::FutureExt as _; use pretty_assertions::assert_eq; @@ -20,9 +22,11 @@ use rmcp::model::ClientCapabilities; use rmcp::model::Implementation; use rmcp::model::InitializeRequestParams; use rmcp::model::ProtocolVersion; +use serde_json::json; use wiremock::Mock; use wiremock::MockServer; use wiremock::ResponseTemplate; +use wiremock::matchers::body_partial_json; use wiremock::matchers::method; use streamable_http_test_support::call_echo_tool; @@ -56,13 +60,28 @@ async fn streamable_http_remote_client_round_trips_through_exec_server() -> anyh /// A timed-out MCP handshake must release the remote executor for later requests. #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn streamable_http_handshake_timeout_unblocks_remote_executor() -> anyhow::Result<()> { + stalled_handshake_unblocks_remote_executor(McpProtocolMode::Legacy, "initialize").await +} + +/// MCP 2026 discovery must also release the remote executor when it times out. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn streamable_http_discovery_timeout_unblocks_remote_executor() -> anyhow::Result<()> { + stalled_handshake_unblocks_remote_executor(McpProtocolMode::V20260728, "server/discover").await +} + +async fn stalled_handshake_unblocks_remote_executor( + protocol_mode: McpProtocolMode, + expected_method: &str, +) -> anyhow::Result<()> { let http_server = MockServer::start().await; Mock::given(method("POST")) + .and(body_partial_json(json!({ "method": expected_method }))) .respond_with(ResponseTemplate::new(200).set_delay(Duration::from_secs(10))) + .expect(1) .mount(&http_server) .await; let exec_server = spawn_exec_server().await?; - let client = RmcpClient::new_streamable_http_client( + let client = RmcpClient::new_streamable_http_client_with_protocol_mode( "stalled-remote-http", &format!("{}/mcp", http_server.uri()), Some("test-bearer".to_string()), @@ -72,6 +91,7 @@ async fn streamable_http_handshake_timeout_unblocks_remote_executor() -> anyhow: AuthKeyringBackendKind::default(), Arc::new(exec_server.client.clone()), /*auth_provider*/ None, + protocol_mode, ) .await?; let params = InitializeRequestParams::new( @@ -95,7 +115,8 @@ async fn streamable_http_handshake_timeout_unblocks_remote_executor() -> anyhow: }), ) .await - .expect_err("the stalled MCP handshake must time out"); + .err() + .context("the stalled MCP handshake must time out")?; assert!( error.to_string().contains("timed out"), "unexpected MCP handshake error: {error}" @@ -106,6 +127,6 @@ async fn streamable_http_handshake_timeout_unblocks_remote_executor() -> anyhow: exec_server.client.environment_info(), ) .await - .expect("a timed-out handshake must not leave the serial executor blocked")?; + .context("a timed-out handshake must not leave the serial executor blocked")??; Ok(()) }