diff --git a/crates/rmcp/Cargo.toml b/crates/rmcp/Cargo.toml index 7f172597d..032aa795c 100644 --- a/crates/rmcp/Cargo.toml +++ b/crates/rmcp/Cargo.toml @@ -365,6 +365,14 @@ required-features = [ ] path = "tests/test_streamable_http_malformed_json_request.rs" +[[test]] +name = "test_streamable_http_initialize_error" +required-features = [ + "client", + "transport-streamable-http-client", + "transport-streamable-http-client-reqwest", +] +path = "tests/test_streamable_http_initialize_error.rs" [[test]] name = "test_custom_request" diff --git a/crates/rmcp/src/transport/streamable_http_client.rs b/crates/rmcp/src/transport/streamable_http_client.rs index a5811429a..9f1a4bcdb 100644 --- a/crates/rmcp/src/transport/streamable_http_client.rs +++ b/crates/rmcp/src/transport/streamable_http_client.rs @@ -217,6 +217,9 @@ pub enum StreamableHttpError { /// Session recovery timed out. The server may have processed an interrupted POST. #[error("Session recovery timed out; the server may have processed the POST")] SessionRecoveryTimeout, + /// The server answered the session re-initialization with a JSON-RPC error. + #[error("Session re-initialization failed: {0}")] + ReinitializationFailed(ErrorData), /// A cancellation or reply POST did not finish in time. #[error("Control POST timed out")] ControlRequestTimeout, @@ -314,7 +317,10 @@ impl StreamableHttpPostResponse { let message: ServerJsonRpcMessage = serde_json::from_str(&payload)?; - if matches!(message, ServerJsonRpcMessage::Response(_)) { + if matches!( + message, + ServerJsonRpcMessage::Response(_) | ServerJsonRpcMessage::Error(_) + ) { return Ok((message, session_id)); } @@ -988,7 +994,8 @@ impl StreamableHttpClientWorker { /// protocol version, sends `notifications/initialized`, and returns the new /// `(session_id, protocol_headers)` pair. The init result message is **not** /// forwarded to the handler because the handler already processed the original - /// initialization. + /// initialization. If the server answers the initialize with a JSON-RPC error, + /// returns an error without sending `notifications/initialized`. async fn perform_reinitialization( client: C, saved_init_request: ClientJsonRpcMessage, @@ -1016,6 +1023,9 @@ impl StreamableHttpClientWorker { .await? .expect_initialized::() .await?; + if let ServerJsonRpcMessage::Error(error) = init_msg { + return Err(StreamableHttpError::ReinitializationFailed(error.error)); + } let new_session_id: Option> = new_session_id_str.map(|s| Arc::from(s.as_str())); diff --git a/crates/rmcp/tests/test_streamable_http_initialize_error.rs b/crates/rmcp/tests/test_streamable_http_initialize_error.rs new file mode 100644 index 000000000..63d139e35 --- /dev/null +++ b/crates/rmcp/tests/test_streamable_http_initialize_error.rs @@ -0,0 +1,226 @@ +#![cfg(all( + feature = "client", + feature = "transport-streamable-http-client", + feature = "transport-streamable-http-client-reqwest", + not(feature = "local") +))] + +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use axum::{ + Router, + body::{Body, Bytes}, + http::{HeaderMap, Response}, + routing::post, +}; +use futures::{StreamExt, stream}; +use rmcp::{ + ServiceError, ServiceExt, + model::{ClientConfig, ProtocolVersion}, + service::ClientInitializeError, + transport::{ + StreamableHttpClientTransport, + streamable_http_client::{StreamableHttpClientTransportConfig, StreamableHttpError}, + }, +}; +use serde_json::{Value, json}; + +fn initialize_error(id: Value) -> Value { + json!({ + "jsonrpc": "2.0", + "id": id, + "error": { + "code": -32602, + "message": "Unsupported protocol version", + "data": { "supported": ["2024-11-05"], "requested": "1.0.0" } + } + }) +} + +/// A single SSE event, optionally followed by a stream that never ends. +fn sse_response(message: Value, hold_open: bool) -> Response { + let event = format!("data: {message}\n\n"); + let body = if hold_open { + let first = stream::once(async move { Ok::<_, std::io::Error>(event) }); + Body::from_stream(first.chain(stream::pending())) + } else { + Body::from(event) + }; + Response::builder() + .status(200) + .header("content-type", "text/event-stream") + .body(body) + .unwrap() +} + +fn json_response(message: Value, session_id: Option<&str>) -> Response { + let mut builder = Response::builder() + .status(200) + .header("content-type", "application/json"); + if let Some(session_id) = session_id { + builder = builder.header("mcp-session-id", session_id); + } + builder.body(Body::from(message.to_string())).unwrap() +} + +async fn spawn(router: Router) -> String { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(listener, router).await.unwrap() }); + format!("http://{addr}/mcp") +} + +async fn spawn_sse_initialize_error_server(hold_open: bool) -> String { + let router = Router::new().route( + "/mcp", + post(move |body: Bytes| async move { + let request: Value = serde_json::from_slice(&body).unwrap(); + sse_response(initialize_error(request["id"].clone()), hold_open) + }), + ); + spawn(router).await +} + +async fn assert_initialize_jsonrpc_error(uri: String) { + let result = tokio::time::timeout( + Duration::from_secs(3), + ().serve(StreamableHttpClientTransport::from_uri(uri)), + ) + .await + .expect("serve() should not hang on an initialize error"); + match result { + Err(ClientInitializeError::JsonRpcError(error)) => { + assert_eq!(error.code.0, -32602); + assert_eq!(error.message, "Unsupported protocol version"); + assert_eq!( + error.data, + Some(json!({ "supported": ["2024-11-05"], "requested": "1.0.0" })) + ); + } + other => panic!("expected ClientInitializeError::JsonRpcError, got: {other:?}"), + } +} + +#[tokio::test] +async fn initialize_sse_error_then_stream_end_returns_jsonrpc_error() { + let uri = spawn_sse_initialize_error_server(false).await; + assert_initialize_jsonrpc_error(uri).await; +} + +#[tokio::test] +async fn initialize_sse_error_with_stream_held_open_returns_jsonrpc_error() { + let uri = spawn_sse_initialize_error_server(true).await; + assert_initialize_jsonrpc_error(uri).await; +} + +type PostLog = Arc)>>>; + +/// Accepts the first initialize as session `s1`, expires `s1` on the next +/// request, and answers the re-initialize with an error. +async fn spawn_reinit_error_server(error_over_sse: bool) -> (String, PostLog) { + let log = PostLog::default(); + let initialized = Arc::new(Mutex::new(false)); + let router = Router::new().route( + "/mcp", + post({ + let log = log.clone(); + move |headers: HeaderMap, body: Bytes| { + let log = log.clone(); + let initialized = initialized.clone(); + async move { + let request: Value = serde_json::from_slice(&body).unwrap(); + let method = request["method"].as_str().unwrap_or_default().to_owned(); + let session_id = headers + .get("mcp-session-id") + .map(|value| value.to_str().unwrap().to_owned()); + log.lock().unwrap().push((method.clone(), session_id)); + let id = request["id"].clone(); + let first_initialize = method == "initialize" + && !std::mem::replace(&mut *initialized.lock().unwrap(), true); + match method.as_str() { + "initialize" if first_initialize => json_response( + json!({ + "jsonrpc": "2.0", + "id": id, + "result": { + "protocolVersion": "2025-11-25", + "capabilities": {}, + "serverInfo": { "name": "mock", "version": "1" } + } + }), + Some("s1"), + ), + "initialize" if error_over_sse => sse_response(initialize_error(id), true), + "initialize" => json_response(initialize_error(id), None), + "notifications/initialized" => { + Response::builder().status(202).body(Body::empty()).unwrap() + } + _ => Response::builder().status(404).body(Body::empty()).unwrap(), + } + } + } + }), + ); + (spawn(router).await, log) +} + +async fn assert_reinit_error_fails_pending_request(error_over_sse: bool) { + let (uri, log) = spawn_reinit_error_server(error_over_sse).await; + let transport = StreamableHttpClientTransport::from_config( + StreamableHttpClientTransportConfig::with_uri(uri).reinit_on_expired_session(true), + ); + let client = ClientConfig::default() + .with_protocol_version(ProtocolVersion::LATEST_WITH_INITIALIZE) + .serve(transport) + .await + .expect("initial handshake"); + + let result = tokio::time::timeout(Duration::from_secs(3), client.list_tools(None)) + .await + .expect("pending request should fail promptly"); + + let posts = log.lock().unwrap().clone(); + let s1 = Some("s1".to_owned()); + assert_eq!( + posts, + vec![ + ("initialize".to_owned(), None), + ("notifications/initialized".to_owned(), s1.clone()), + ("tools/list".to_owned(), s1), + ("initialize".to_owned(), None), + ], + "no notifications/initialized or retry may follow a failed re-initialize" + ); + + let Err(ServiceError::TransportSend(error)) = result else { + panic!("expected TransportSend error, got: {result:?}"); + }; + match error + .error + .downcast_ref::>() + { + Some(StreamableHttpError::ReinitializationFailed(error)) => { + assert_eq!(error.code.0, -32602); + assert_eq!(error.message, "Unsupported protocol version"); + assert_eq!( + error.data, + Some(json!({ "supported": ["2024-11-05"], "requested": "1.0.0" })) + ); + } + other => panic!("expected ReinitializationFailed, got: {other:?}"), + } + let _ = client.cancel().await; +} + +#[tokio::test] +async fn reinit_json_error_fails_pending_request_without_retry() { + assert_reinit_error_fails_pending_request(false).await; +} + +#[tokio::test] +async fn reinit_sse_error_fails_pending_request_without_retry() { + assert_reinit_error_fails_pending_request(true).await; +}