Cover remote MCP discovery timeout cleanup (#37248)

## What changed

Extend the executor-backed Streamable HTTP timeout test to cover both legacy
`initialize` and MCP 2026 `server/discover` requests. Verify that each stalled
handshake times out and releases the serial executor for subsequent requests.

GitOrigin-RevId: 3493904433c95265b9f645813723e890b7610abb
This commit is contained in:
jif
2026-08-06 09:46:04 +00:00
committed by copyberry
parent 7a0e974e08
commit 74b8f8db93
@@ -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(())
}