diff --git a/crates/rmcp/src/service.rs b/crates/rmcp/src/service.rs index b1b2c2b7c..0f47fa9a2 100644 --- a/crates/rmcp/src/service.rs +++ b/crates/rmcp/src/service.rs @@ -168,12 +168,31 @@ pub trait ServiceRole: std::fmt::Debug + Send + Sync + 'static + Copy + Clone { fn enforce_peer_request_association( _peer_request: &Self::PeerReq, _peer_info: Option<&Self::PeerInfo>, - _has_pending_outbound_request: bool, + _association: PeerRequestAssociation, ) -> Result<(), McpError> { Ok(()) } } +/// How an inbound peer request relates to this side's in-flight outbound +/// requests (SEP-2260). +/// +/// SEP-2260 defines no wire field for association, so only stream-separating +/// transports (streamable HTTP) can observe it; other transports yield +/// [`Self::Unknown`]. +#[derive(Debug, Clone, PartialEq, Eq)] +#[expect(clippy::exhaustive_enums, reason = "intentionally exhaustive")] +pub enum PeerRequestAssociation { + /// Arrived on the response stream of an in-flight outbound request. + Associated, + /// Arrived on a stream tied to no in-flight outbound request (e.g. the + /// streamable HTTP standalone GET stream). + Unassociated, + /// The transport cannot distinguish streams; only the coarse in-flight + /// signal is available. + Unknown { has_pending_outbound_request: bool }, +} + pub(crate) fn uses_legacy_lifecycle( protocol_version: Option<&ProtocolVersion>, uses_discover_lifecycle: bool, @@ -182,6 +201,25 @@ pub(crate) fn uses_legacy_lifecycle( && protocol_version.is_none_or(|version| version < &ProtocolVersion::V_2026_07_28) } +pub(crate) fn peer_request_association( + request: &Req, + local_responder_pool: &std::collections::HashMap, +) -> PeerRequestAssociation { + match request.extensions().get::() { + None => PeerRequestAssociation::Unknown { + has_pending_outbound_request: !local_responder_pool.is_empty(), + }, + Some(InboundStreamOrigin::Unassociated) => PeerRequestAssociation::Unassociated, + Some(InboundStreamOrigin::OutboundRequest(id)) => { + if local_responder_pool.contains_key(id) { + PeerRequestAssociation::Associated + } else { + PeerRequestAssociation::Unassociated + } + } + } +} + tokio::task_local! { pub(crate) static ORIGINATING_REQUEST: RequestId; } @@ -205,10 +243,26 @@ pub(crate) fn in_request_handler_scope() -> bool { /// outside a handler they return an `invalid_request` error. The association /// is task-local and does not cross `tokio::spawn`, so use the task manager /// for long-running work. +/// +/// The client receive-side mirror is [`InboundStreamOrigin`]. #[derive(Debug, Clone, PartialEq, Eq)] #[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")] pub struct OriginatingRequestId(pub RequestId); +/// Marker in an inbound request's non-serialized [`Extensions`] recording +/// which HTTP response stream it arrived on: the receive-side mirror of +/// [`OriginatingRequestId`]. Never on the wire (SEP-2260 defines no wire +/// field); when absent, the coarse in-flight check applies. +#[derive(Debug, Clone, PartialEq, Eq)] +#[expect(clippy::exhaustive_enums, reason = "intentionally exhaustive")] +pub enum InboundStreamOrigin { + /// The standalone GET stream, or a POST response stream not tied to an + /// outbound request. + Unassociated, + /// The SSE response stream of the POST that carried this outbound request. + OutboundRequest(RequestId), +} + pub type TxJsonRpcMessage = JsonRpcMessage<::Req, ::Resp, ::Not>; pub type RxJsonRpcMessage = JsonRpcMessage< @@ -1448,7 +1502,7 @@ where if let Err(error) = R::enforce_peer_request_association( &request, peer.peer_info().as_deref(), - !local_responder_pool.is_empty(), + peer_request_association(&request, &local_responder_pool), ) { tracing::warn!(%id, message = %error.message, "rejected peer request"); // send directly: the sink proxy path would drop the @@ -1729,4 +1783,73 @@ mod sep2260_marker_tests { let request = send_and_capture(None).await; assert!(request.extensions().get::().is_none()); } + + #[test] + #[expect( + deprecated, + reason = "Sampling is deprecated by SEP-2577 but remains the canonical restricted request" + )] + fn peer_request_association_maps_stream_origin() { + use std::collections::HashMap; + + use crate::model::{ + CreateMessageRequest, CreateMessageRequestParams, SamplingMessage, ServerRequest, + }; + + fn sampling(origin: Option) -> ServerRequest { + let mut request = CreateMessageRequest::new(CreateMessageRequestParams::new( + vec![SamplingMessage::user_text("hi")], + 16, + )); + if let Some(origin) = origin { + request.extensions.insert(origin); + } + ServerRequest::CreateMessageRequest(request) + } + + let empty: HashMap = HashMap::new(); + let in_flight: HashMap = HashMap::from([(RequestId::Number(7), ())]); + + // No marker (stdio): coarse signal. + assert_eq!( + peer_request_association(&sampling(None), &in_flight), + PeerRequestAssociation::Unknown { + has_pending_outbound_request: true + } + ); + assert_eq!( + peer_request_association(&sampling(None), &empty), + PeerRequestAssociation::Unknown { + has_pending_outbound_request: false + } + ); + // Standalone GET stream: unassociated even with requests in flight. + assert_eq!( + peer_request_association( + &sampling(Some(InboundStreamOrigin::Unassociated)), + &in_flight + ), + PeerRequestAssociation::Unassociated + ); + // Originating POST stream of an in-flight request: associated. + assert_eq!( + peer_request_association( + &sampling(Some(InboundStreamOrigin::OutboundRequest( + RequestId::Number(7) + ))), + &in_flight + ), + PeerRequestAssociation::Associated + ); + // Stream of a request that is no longer in flight: unassociated. + assert_eq!( + peer_request_association( + &sampling(Some(InboundStreamOrigin::OutboundRequest( + RequestId::Number(8) + ))), + &in_flight + ), + PeerRequestAssociation::Unassociated + ); + } } diff --git a/crates/rmcp/src/service/client.rs b/crates/rmcp/src/service/client.rs index b23a1496d..5c8473147 100644 --- a/crates/rmcp/src/service/client.rs +++ b/crates/rmcp/src/service/client.rs @@ -238,14 +238,13 @@ impl ServiceRole for RoleClient { } } - // SEP-2260: with no outbound request in flight there is nothing the - // server request could be associated with, so reject it. With one in - // flight we cannot tell which request it belongs to (no wire field), so - // we accept — an under-approximation of the spec's SHOULD. + // SEP-2260: reject restricted server requests that arrived unassociated + // with any in-flight outbound request. Without stream separation + // (`Unknown`) the coarse in-flight check under-approximates the SHOULD. fn enforce_peer_request_association( peer_request: &Self::PeerReq, peer_info: Option<&Self::PeerInfo>, - has_pending_outbound_request: bool, + association: PeerRequestAssociation, ) -> Result<(), ErrorData> { let restricted = matches!( peer_request, @@ -258,13 +257,24 @@ impl ServiceRole for RoleClient { } let strict = peer_info.is_some_and(|info| info.protocol_version >= ProtocolVersion::V_2026_07_28); - if strict && !has_pending_outbound_request { - return Err(ErrorData::invalid_params( + if !strict { + return Ok(()); + } + let associated = match association { + PeerRequestAssociation::Associated => true, + PeerRequestAssociation::Unassociated => false, + PeerRequestAssociation::Unknown { + has_pending_outbound_request, + } => has_pending_outbound_request, + }; + if associated { + Ok(()) + } else { + Err(ErrorData::invalid_params( "SEP-2260: server-to-client requests must be associated with an in-flight client request", None, - )); + )) } - Ok(()) } async fn invalidate_response_cache(peer: &Peer, notification: &Self::PeerNot) { @@ -2215,3 +2225,72 @@ mod tests { assert_eq!(peer.discover(meta).await.unwrap(), expected); } } + +#[cfg(test)] +mod sep2260_association_tests { + use super::*; + use crate::{ + model::{ + CreateMessageRequest, CreateMessageRequestParams, SamplingMessage, ServerCapabilities, + }, + service::PeerRequestAssociation, + }; + + fn sampling_request() -> ServerRequest { + ServerRequest::CreateMessageRequest(CreateMessageRequest::new( + CreateMessageRequestParams::new(vec![SamplingMessage::user_text("hi")], 16), + )) + } + + fn server_info(version: ProtocolVersion) -> ServerInfo { + let mut info = ServerInfo::new(ServerCapabilities::default()); + info.protocol_version = version; + info + } + + fn enforce(info: &ServerInfo, association: PeerRequestAssociation) -> Result<(), ErrorData> { + RoleClient::enforce_peer_request_association(&sampling_request(), Some(info), association) + } + + #[test] + fn strict_rejects_unassociated() { + let info = server_info(ProtocolVersion::V_2026_07_28); + let err = enforce(&info, PeerRequestAssociation::Unassociated).unwrap_err(); + assert_eq!(err.code, crate::model::ErrorCode::INVALID_PARAMS); + } + + #[test] + fn strict_accepts_associated() { + let info = server_info(ProtocolVersion::V_2026_07_28); + assert!(enforce(&info, PeerRequestAssociation::Associated).is_ok()); + } + + #[test] + fn strict_unknown_falls_back_to_coarse_check() { + let info = server_info(ProtocolVersion::V_2026_07_28); + assert!( + enforce( + &info, + PeerRequestAssociation::Unknown { + has_pending_outbound_request: true + } + ) + .is_ok() + ); + assert!( + enforce( + &info, + PeerRequestAssociation::Unknown { + has_pending_outbound_request: false + } + ) + .is_err() + ); + } + + #[test] + fn legacy_protocol_accepts_even_unassociated() { + let info = server_info(ProtocolVersion::V_2025_11_25); + assert!(enforce(&info, PeerRequestAssociation::Unassociated).is_ok()); + } +} diff --git a/crates/rmcp/src/transport/streamable_http_client.rs b/crates/rmcp/src/transport/streamable_http_client.rs index c43fb7dd5..fe563ee3f 100644 --- a/crates/rmcp/src/transport/streamable_http_client.rs +++ b/crates/rmcp/src/transport/streamable_http_client.rs @@ -19,10 +19,11 @@ use super::common::client_side_sse::{ use crate::{ RoleClient, model::{ - ClientJsonRpcMessage, ClientNotification, ClientRequest, ErrorData, GetMeta, + ClientJsonRpcMessage, ClientNotification, ClientRequest, ErrorData, GetExtensions, GetMeta, InitializedNotification, JsonObject, ProtocolVersion, RequestId, ServerJsonRpcMessage, ServerResult, }, + service::InboundStreamOrigin, transport::{ common::{client_side_sse::SseAutoReconnectStream, mcp_headers}, worker::{Worker, WorkerQuitReason, WorkerSendRequest, WorkerTransport}, @@ -635,6 +636,7 @@ impl StreamableHttpClientWorker { + Send + 'static, sse_worker_tx: tokio::sync::mpsc::Sender, + origin: InboundStreamOrigin, close_on_response: bool, ct: CancellationToken, ) -> Result<(), StreamableHttpError> { @@ -649,9 +651,14 @@ impl StreamableHttpClientWorker { break; } }; - let Some(message) = message.transpose()? else { + let Some(mut message) = message.transpose()? else { break; }; + // SEP-2260: mark inbound requests with the stream they arrived on + // for the client receive-side association check. + if let ServerJsonRpcMessage::Request(request) = &mut message { + request.request.extensions_mut().insert(origin.clone()); + } let is_response = matches!( message, ServerJsonRpcMessage::Response(_) | ServerJsonRpcMessage::Error(_) @@ -719,6 +726,7 @@ impl StreamableHttpClientWorker { Self::execute_sse_stream( sse_stream, sse_worker_tx, + InboundStreamOrigin::Unassociated, false, transport_task_ct.child_token(), ) @@ -1283,9 +1291,18 @@ impl Worker for StreamableHttpClientWorker { ); } let stream_tx = sse_worker_tx.clone(); + let origin = match &stream_request_id { + Some(id) => { + InboundStreamOrigin::OutboundRequest( + id.clone(), + ) + } + None => InboundStreamOrigin::Unassociated, + }; streams.spawn(async move { let result = Self::execute_sse_stream( - sse_stream, stream_tx, true, stream_ct, + sse_stream, stream_tx, origin, true, + stream_ct, ) .await; (stream_request_id, result) @@ -1342,9 +1359,13 @@ impl Worker for StreamableHttpClientWorker { .insert(request_id.clone(), stream_ct.clone()); } let stream_tx = sse_worker_tx.clone(); + let origin = match &stream_request_id { + Some(id) => InboundStreamOrigin::OutboundRequest(id.clone()), + None => InboundStreamOrigin::Unassociated, + }; streams.spawn(async move { let result = Self::execute_sse_stream( - sse_stream, stream_tx, true, stream_ct, + sse_stream, stream_tx, origin, true, stream_ct, ) .await; (stream_request_id, result) @@ -1769,7 +1790,66 @@ mod tests { use serde_json::json; use super::*; - use crate::model::{ListToolsResult, NumberOrString, ServerResult, Tool}; + use crate::{ + model::{ + GetExtensions, ListToolsResult, NumberOrString, ServerRequest, ServerResult, Tool, + }, + service::InboundStreamOrigin, + }; + + #[expect( + deprecated, + reason = "Sampling is deprecated by SEP-2577 but remains the canonical restricted request" + )] + fn sampling_request_message(id: i64) -> ServerJsonRpcMessage { + use crate::model::{CreateMessageRequest, CreateMessageRequestParams, SamplingMessage}; + ServerJsonRpcMessage::request( + ServerRequest::CreateMessageRequest(CreateMessageRequest::new( + CreateMessageRequestParams::new(vec![SamplingMessage::user_text("hi")], 16), + )), + NumberOrString::Number(id), + ) + } + + #[tokio::test] + async fn execute_sse_stream_marks_inbound_requests_with_origin() { + for origin in [ + InboundStreamOrigin::Unassociated, + InboundStreamOrigin::OutboundRequest(RequestId::Number(3)), + ] { + let response = ServerJsonRpcMessage::response( + ServerResult::ListToolsResult(ListToolsResult::default()), + NumberOrString::Number(1), + ); + let stream = futures::stream::iter([Ok(sampling_request_message(9)), Ok(response)]); + let (tx, mut rx) = tokio::sync::mpsc::channel(4); + StreamableHttpClientWorker::::execute_sse_stream( + stream, + tx, + origin.clone(), + false, + CancellationToken::new(), + ) + .await + .expect("stream completes"); + + let ServerJsonRpcMessage::Request(request) = + rx.recv().await.expect("request forwarded") + else { + panic!("expected request first"); + }; + assert_eq!( + request.request.extensions().get::(), + Some(&origin), + "inbound requests must carry their stream origin" + ); + // Responses are correlated by JSON-RPC id; no marker needed or added. + assert!(matches!( + rx.recv().await.expect("response forwarded"), + ServerJsonRpcMessage::Response(_) + )); + } + } type ReconnectAttempt = (Option, Option); @@ -1868,6 +1948,132 @@ mod tests { ); } + #[derive(Clone, Default)] + struct ResumedRequestClient { + reconnects: Arc>>, + } + + impl StreamableHttpClient for ResumedRequestClient { + type Error = std::io::Error; + + async fn post_message( + &self, + _uri: Arc, + _message: ClientJsonRpcMessage, + _session_id: Option>, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result> { + Err(StreamableHttpError::UnexpectedServerResponse( + "unexpected POST".into(), + )) + } + + async fn delete_session( + &self, + _uri: Arc, + _session_id: Arc, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result<(), StreamableHttpError> { + Ok(()) + } + + async fn get_stream( + &self, + _uri: Arc, + session_id: Option>, + last_event_id: Option, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result> { + self.reconnects + .lock() + .expect("lock reconnects") + .push((session_id.map(|id| id.to_string()), last_event_id)); + let request = sampling_request_message(9); + let response = ServerJsonRpcMessage::response( + ServerResult::ListToolsResult(ListToolsResult::default()), + NumberOrString::Number(1), + ); + // Stay open after the response, like a live connection, so the + // post-response drain in `execute_sse_stream` doesn't trigger + // further reconnects. + Ok(futures::stream::iter([request, response].map(|message| { + Ok(Sse { + event: None, + data: Some(serde_json::to_string(&message).expect("serialize message")), + id: None, + retry: None, + }) + })) + .chain(futures::stream::pending()) + .boxed()) + } + } + + /// SEP-1699 resumes a broken POST SSE stream via GET + Last-Event-ID + /// beneath `execute_sse_stream`, so the SEP-2260 origin marker must span + /// resumes; if reconnection were hoisted above the marker attach point, + /// replayed associated requests would be wrongly rejected with -32602. + #[tokio::test] + async fn resumed_post_stream_requests_keep_outbound_origin() { + let initial = futures::stream::iter([Ok(Sse { + event: None, + data: None, + id: Some("e1".into()), + retry: Some(0), + })]) + .boxed(); + let client = ResumedRequestClient::default(); + let reconnects = client.reconnects.clone(); + let sse_stream = + StreamableHttpClientWorker::::response_sse_to_jsonrpc( + initial, + None, + client, + Arc::from("http://localhost/mcp"), + None, + HashMap::new(), + DEFAULT_MAX_SSE_EVENT_SIZE, + Arc::new(ExponentialBackoff { + max_times: Some(1), + base_duration: Duration::ZERO, + }), + ); + + let origin = InboundStreamOrigin::OutboundRequest(RequestId::Number(3)); + let (tx, mut rx) = tokio::sync::mpsc::channel(4); + StreamableHttpClientWorker::::execute_sse_stream( + sse_stream, + tx, + origin.clone(), + true, + CancellationToken::new(), + ) + .await + .expect("stream completes"); + + assert_eq!( + reconnects.lock().expect("lock reconnects").as_slice(), + &[(None, Some("e1".into()))], + "the request must arrive on the resumed connection" + ); + let ServerJsonRpcMessage::Request(request) = rx.recv().await.expect("request forwarded") + else { + panic!("expected request first"); + }; + assert_eq!( + request.request.extensions().get::(), + Some(&origin), + "origin marker must survive SSE resumption" + ); + assert!(matches!( + rx.recv().await.expect("response forwarded"), + ServerJsonRpcMessage::Response(_) + )); + } + fn tool(name: &'static str, annotation: serde_json::Value) -> Tool { let schema = json!({ "type": "object", diff --git a/crates/rmcp/tests/test_sep_2260_stream_enforcement.rs b/crates/rmcp/tests/test_sep_2260_stream_enforcement.rs new file mode 100644 index 000000000..4246b7d68 --- /dev/null +++ b/crates/rmcp/tests/test_sep_2260_stream_enforcement.rs @@ -0,0 +1,276 @@ +//! SEP-2260 follow-up (#1033): stream-based receive-side enforcement. +//! +//! Scripted streamable HTTP "server": answers a legacy initialize with +//! protocol 2026-07-28 and a session id (spec-legal; rmcp's own server is +//! stateless at that version, but the client must be correct against any +//! server), so the client has BOTH a standalone GET stream and strict +//! SEP-2260 enforcement. +#![cfg(all( + feature = "client", + feature = "transport-streamable-http-client", + not(feature = "local") +))] +#![allow(deprecated)] + +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use futures::{StreamExt, stream::BoxStream}; +use http::{HeaderName, HeaderValue}; +use rmcp::{ + ClientHandler, + model::{ + ClientInfo, ClientJsonRpcMessage, CreateMessageRequestParams, CreateMessageResult, + ProtocolVersion, SamplingMessage, ServerCapabilities, ServerInfo, ServerJsonRpcMessage, + }, + service::{ClientLifecycleMode, RequestContext, RoleClient, serve_client_with_lifecycle}, + transport::streamable_http_client::{ + StreamableHttpClient, StreamableHttpClientTransport, StreamableHttpClientTransportConfig, + StreamableHttpError, StreamableHttpPostResponse, + }, +}; +use serde_json::{Value, json}; +use sse_stream::{Error as SseError, Sse}; +use tokio::sync::{Mutex, mpsc}; + +fn to_sse(message: Value) -> Result { + Ok(Sse { + event: None, + data: Some(message.to_string()), + id: None, + retry: None, + }) +} + +fn message_stream(rx: mpsc::Receiver) -> BoxStream<'static, Result> { + tokio_stream::wrappers::ReceiverStream::new(rx) + .map(to_sse) + .boxed() +} + +/// Scripted server: initialize -> JSON init result (2026-07-28 + session); +/// first non-initialize request POST -> SSE stream fed by `post_stream`; +/// everything else -> Accepted. Every message the client POSTs is forwarded +/// to `posted`. +#[derive(Clone)] +struct ScriptedServer { + get_stream: Arc>>>, + post_stream: Arc>>>, + posted: mpsc::UnboundedSender, +} + +impl StreamableHttpClient for ScriptedServer { + type Error = std::io::Error; + + async fn post_message( + &self, + _uri: Arc, + message: ClientJsonRpcMessage, + _session_id: Option>, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result> { + let value = serde_json::to_value(&message).expect("serialize client message"); + self.posted.send(value.clone()).expect("test alive"); + if value["method"] == "initialize" { + let mut info = ServerInfo::new(ServerCapabilities::default()); + info.protocol_version = ProtocolVersion::V_2026_07_28; + let response = ServerJsonRpcMessage::response( + rmcp::model::ServerResult::InitializeResult(info), + serde_json::from_value(value["id"].clone()).expect("request id"), + ); + return Ok(StreamableHttpPostResponse::Json( + response, + Some("scripted-session".into()), + )); + } + if matches!(message, ClientJsonRpcMessage::Request(_)) { + let rx = self + .post_stream + .lock() + .await + .take() + .expect("exactly one non-initialize request POST in this test"); + return Ok(StreamableHttpPostResponse::Sse(message_stream(rx), None)); + } + Ok(StreamableHttpPostResponse::Accepted) + } + + async fn delete_session( + &self, + _uri: Arc, + _session_id: Arc, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result<(), StreamableHttpError> { + Ok(()) + } + + async fn get_stream( + &self, + _uri: Arc, + _session_id: Option>, + _last_event_id: Option, + _auth_header: Option, + _custom_headers: HashMap, + ) -> Result>, StreamableHttpError> { + match self.get_stream.lock().await.take() { + Some(rx) => Ok(message_stream(rx)), + // Reconnect after the scripted stream ends: stay silent. + None => Ok(futures::stream::pending().boxed()), + } + } +} + +#[derive(Clone)] +struct SamplingClient { + // Some(tx): forward sampling params; None: sampling must never reach the handler. + on_sampling: Option>, +} + +impl ClientHandler for SamplingClient { + async fn create_message( + &self, + params: CreateMessageRequestParams, + _context: RequestContext, + ) -> Result { + let Some(tx) = &self.on_sampling else { + panic!("sampling request must not reach the handler in this test"); + }; + tx.send(params).expect("test alive"); + Ok(CreateMessageResult::new( + SamplingMessage::assistant_text("pong"), + "test-model".to_string(), + )) + } + + fn get_info(&self) -> ClientInfo { + ClientInfo::default() + } +} + +fn sampling_request(id: u32) -> Value { + json!({ + "jsonrpc": "2.0", + "id": id, + "method": "sampling/createMessage", + "params": { + "messages": [{ "role": "user", "content": { "type": "text", "text": "hi" } }], + "maxTokens": 16 + } + }) +} + +fn tools_list_response(id: &Value) -> Value { + json!({ "jsonrpc": "2.0", "id": id, "result": { "tools": [] } }) +} + +async fn next_posted(posted: &mut mpsc::UnboundedReceiver) -> Value { + tokio::time::timeout(Duration::from_secs(5), posted.recv()) + .await + .expect("posted message within 5s") + .expect("channel open") +} + +/// Drive startup + one in-flight tools/list; return (client, posted rx, +/// get stream tx, post stream tx, in-flight call handle, tools/list id). +async fn setup( + handler: SamplingClient, +) -> ( + rmcp::service::RunningService, + mpsc::UnboundedReceiver, + mpsc::Sender, + mpsc::Sender, + tokio::task::JoinHandle>, + Value, +) { + let (get_tx, get_rx) = mpsc::channel(8); + let (post_tx, post_rx) = mpsc::channel(8); + let (posted_tx, mut posted_rx) = mpsc::unbounded_channel(); + let server = ScriptedServer { + get_stream: Arc::new(Mutex::new(Some(get_rx))), + post_stream: Arc::new(Mutex::new(Some(post_rx))), + posted: posted_tx, + }; + let transport = StreamableHttpClientTransport::with_client( + server, + StreamableHttpClientTransportConfig::with_uri("http://scripted/mcp"), + ); + let client = serve_client_with_lifecycle(handler, transport, ClientLifecycleMode::Initialize) + .await + .expect("initialize against scripted server"); + + // initialize + notifications/initialized already posted during startup. + assert_eq!(next_posted(&mut posted_rx).await["method"], "initialize"); + assert_eq!( + next_posted(&mut posted_rx).await["method"], + "notifications/initialized" + ); + + // Unrelated outbound request, kept in flight (response withheld). + let peer = client.peer().clone(); + let call = tokio::spawn(async move { peer.list_tools(None).await }); + let tools_list = next_posted(&mut posted_rx).await; + assert_eq!(tools_list["method"], "tools/list"); + let tools_list_id = tools_list["id"].clone(); + + (client, posted_rx, get_tx, post_tx, call, tools_list_id) +} + +/// #1033 scenario 1: a restricted request on the standalone GET stream while +/// an unrelated outbound request is in flight must be rejected with -32602. +/// (The coarse check from #1029 incorrectly accepted this.) +#[tokio::test] +async fn restricted_request_on_get_stream_rejected_while_unrelated_request_in_flight() +-> anyhow::Result<()> { + let (client, mut posted_rx, get_tx, post_tx, call, tools_list_id) = + setup(SamplingClient { on_sampling: None }).await; + + get_tx.send(sampling_request(100)).await?; + + let rejection = next_posted(&mut posted_rx).await; + assert_eq!( + rejection["id"], 100, + "reply to the sampling request: {rejection}" + ); + assert_eq!( + rejection["error"]["code"], -32602, + "SEP-2260: GET-stream request must be rejected even with an unrelated \ + request in flight, got {rejection}" + ); + + post_tx.send(tools_list_response(&tools_list_id)).await?; + call.await??; + client.cancel().await?; + Ok(()) +} + +/// Positive twin: the same restricted request arriving on the SSE stream of +/// the originating POST is dispatched to the handler and answered. +#[tokio::test] +async fn restricted_request_on_originating_post_stream_is_dispatched() -> anyhow::Result<()> { + let (sampled_tx, mut sampled_rx) = mpsc::unbounded_channel(); + let (client, mut posted_rx, _get_tx, post_tx, call, tools_list_id) = setup(SamplingClient { + on_sampling: Some(sampled_tx), + }) + .await; + + post_tx.send(sampling_request(200)).await?; + + let response = next_posted(&mut posted_rx).await; + assert_eq!( + response["id"], 200, + "reply to the sampling request: {response}" + ); + assert_eq!( + response["result"]["model"], "test-model", + "request on the originating POST stream must reach the handler, got {response}" + ); + tokio::time::timeout(Duration::from_secs(5), sampled_rx.recv()) + .await? + .expect("handler invoked"); + + post_tx.send(tools_list_response(&tools_list_id)).await?; + call.await??; + client.cancel().await?; + Ok(()) +}