#![cfg(not(feature = "local"))]
use std::{
collections::HashMap,
io,
sync::{
Arc, Mutex as StdMutex,
atomic::{AtomicBool, AtomicUsize, Ordering::SeqCst},
},
time::Duration,
};
use futures::{StreamExt, stream::BoxStream};
use http::{HeaderName, HeaderValue};
use rmcp::{
model::{
CallToolRequestParams, CancelledNotificationParam, ClientInfo, ClientJsonRpcMessage,
ClientRequest, DiscoverResult, ProtocolVersion, Request, RequestId, RequestMetaObject,
ServerJsonRpcMessage,
},
service::{
ClientLifecycleMode, PeerRequestOptions, RequestHandle, RoleClient, RunningService,
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, oneshot},
task::JoinHandle,
time::timeout,
};
use tokio_stream::wrappers::UnboundedReceiverStream;
use tokio_util::sync::CancellationToken;
const TEST_TIMEOUT: Duration = Duration::from_secs(5);
type PostResult = Result<StreamableHttpPostResponse, StreamableHttpError<io::Error>>;
type Call = JoinHandle<anyhow::Result<()>>;
type SseReceiver = mpsc::UnboundedReceiver<Result<Sse, SseError>>;
#[derive(Default)]
struct Counts {
initialized: AtomicUsize,
hold_reinitialization: AtomicBool,
manual_controls: AtomicBool,
reject_unmatched_cancellation: AtomicBool,
local_streams: StdMutex<HashMap<RequestId, CancellationToken>>,
local_cancellations: AtomicUsize,
deleted: AtomicUsize,
cancelled: AtomicUsize,
posted: AtomicUsize,
active: AtomicUsize,
peak: AtomicUsize,
}
struct ActivePost(Arc<Counts>);
impl Drop for ActivePost {
fn drop(&mut self) {
self.0.active.fetch_sub(1, SeqCst);
}
}
struct LocalStreamRegistration {
id: RequestId,
counts: Arc<Counts>,
dropped: Option<oneshot::Sender<()>>,
}
impl Drop for LocalStreamRegistration {
fn drop(&mut self) {
self.counts.local_streams.lock().unwrap().remove(&self.id);
let _ = self.dropped.take().unwrap().send(());
}
}
struct Posted {
id: Value,
name: String,
session: Option<Arc<str>>,
reply: oneshot::Sender<PostResult>,
returned: oneshot::Receiver<()>,
}
struct ControlPost {
message: Value,
session: Option<Arc<str>>,
reply: oneshot::Sender<PostResult>,
}
fn response(id: Value, result: Value) -> ServerJsonRpcMessage {
serde_json::from_value(json!({ "jsonrpc": "2.0", "id": id, "result": result }))
.expect("valid scripted response")
}
fn sse(message: Value) -> Result<Sse, SseError> {
Ok(Sse {
event: Some("message".into()),
data: Some(message.to_string()),
id: None,
retry: None,
})
}
async fn next_event<T>(receiver: &mut mpsc::UnboundedReceiver<T>) -> T {
timeout(TEST_TIMEOUT, receiver.recv())
.await
.expect("expected scripted event")
.expect("scripted client remains connected")
}
impl Posted {
fn result(&self) -> ServerJsonRpcMessage {
response(
self.id.clone(),
json!({ "content": [{ "type": "text", "text": self.name }] }),
)
}
fn finish(self, result: PostResult) {
self.reply.send(result).expect("POST is still waiting");
}
fn succeed(self) {
let result = StreamableHttpPostResponse::Json(self.result(), None);
self.finish(Ok(result));
}
fn expire(self) {
self.finish(Err(StreamableHttpError::SessionExpired));
}
async fn finish_and_wait(self, result: PostResult) -> anyhow::Result<()> {
let Self {
reply, returned, ..
} = self;
reply.send(result).expect("POST is still waiting");
timeout(TEST_TIMEOUT, returned).await??;
Ok(())
}
async fn expire_and_wait(self) -> anyhow::Result<()> {
self.finish_and_wait(Err(StreamableHttpError::SessionExpired))
.await
}
async fn start_sse(self) -> anyhow::Result<oneshot::Sender<()>> {
let message = serde_json::to_value(self.result()).unwrap();
let (release, released) = oneshot::channel();
let stream = futures::stream::once(async move {
released.await.expect("release the SSE response");
sse(message)
})
.boxed();
self.finish_and_wait(Ok(StreamableHttpPostResponse::Sse(stream, None)))
.await?;
Ok(release)
}
async fn start_local_sse(
self,
counts: Arc<Counts>,
) -> anyhow::Result<(CancellationToken, oneshot::Receiver<()>)> {
let id: RequestId = serde_json::from_value(self.id.clone())?;
let cancellation = CancellationToken::new();
let eof = cancellation.clone();
counts
.local_streams
.lock()
.unwrap()
.insert(id.clone(), cancellation.clone());
let (dropped, closed) = oneshot::channel();
let registration = LocalStreamRegistration {
id,
counts,
dropped: Some(dropped),
};
let (started, listening) = oneshot::channel();
let stream = futures::stream::once(async move {
let _registration = registration;
let _ = started.send(());
cancellation.cancelled().await;
None::<Result<Sse, SseError>>
})
.filter_map(futures::future::ready)
.boxed();
self.finish_and_wait(Ok(StreamableHttpPostResponse::Sse(stream, None)))
.await?;
timeout(TEST_TIMEOUT, listening).await??;
Ok((eof, closed))
}
}
#[derive(Clone)]
struct ScriptedClient {
started: mpsc::UnboundedSender<Posted>,
controls: mpsc::UnboundedSender<ControlPost>,
incoming: Arc<Mutex<Option<SseReceiver>>>,
reinitializing: mpsc::UnboundedSender<oneshot::Sender<()>>,
counts: Arc<Counts>,
}
impl ScriptedClient {
async fn control_post(&self, message: Value, session: Option<Arc<str>>) -> PostResult {
if !self.counts.manual_controls.load(SeqCst) {
return Ok(StreamableHttpPostResponse::Accepted);
}
let (reply, result) = oneshot::channel();
self.controls
.send(ControlPost {
message,
session,
reply,
})
.expect("test remains connected");
result.await.expect("test answers the control POST")
}
}
impl StreamableHttpClient for ScriptedClient {
type Error = io::Error;
async fn post_message(
&self,
_uri: Arc<str>,
message: ClientJsonRpcMessage,
session: Option<Arc<str>>,
_auth_header: Option<String>,
_custom_headers: HashMap<HeaderName, HeaderValue>,
) -> PostResult {
let value = serde_json::to_value(message).unwrap();
match value["method"].as_str() {
Some("server/discover") => Ok(StreamableHttpPostResponse::Json(
response(
value["id"].clone(),
serde_json::to_value(DiscoverResult::new(
vec![ProtocolVersion::V_2026_07_28],
serde_json::from_value(json!({ "tools": {} })).unwrap(),
))
.unwrap(),
),
None,
)),
Some("initialize") => {
let generation = self.counts.initialized.fetch_add(1, SeqCst) + 1;
if generation > 1 && self.counts.hold_reinitialization.load(SeqCst) {
let (release, released) = oneshot::channel();
self.reinitializing
.send(release)
.expect("test remains connected");
released.await.expect("test releases reinitialization");
}
Ok(StreamableHttpPostResponse::Json(
response(
value["id"].clone(),
json!({
"protocolVersion": "2025-11-25",
"capabilities": { "tools": {} },
"serverInfo": { "name": "scripted", "version": "1" },
}),
),
Some(format!("session-{generation}")),
))
}
Some("notifications/initialized") => Ok(StreamableHttpPostResponse::Accepted),
Some("notifications/cancelled") => {
let id: RequestId =
serde_json::from_value(value["params"]["requestId"].clone()).unwrap();
let local = self.counts.local_streams.lock().unwrap().get(&id).cloned();
if let Some(cancellation) = local {
self.counts.local_cancellations.fetch_add(1, SeqCst);
cancellation.cancel();
return Ok(StreamableHttpPostResponse::Accepted);
}
self.counts.cancelled.fetch_add(1, SeqCst);
if self.counts.reject_unmatched_cancellation.load(SeqCst) {
return Err(StreamableHttpError::Client(io::Error::other(
"local event cancellation reached the http POST path",
)));
}
self.control_post(value, session).await
}
Some("tools/call") => {
self.counts.posted.fetch_add(1, SeqCst);
let active = self.counts.active.fetch_add(1, SeqCst) + 1;
self.counts.peak.fetch_max(active, SeqCst);
let _active = ActivePost(self.counts.clone());
let (reply, response) = oneshot::channel();
let (finished, returned) = oneshot::channel();
self.started
.send(Posted {
id: value["id"].clone(),
name: value["params"]["name"].as_str().unwrap().to_owned(),
session,
reply,
returned,
})
.expect("test remains connected");
let response = response.await.expect("test answers each POST");
let _ = finished.send(());
response
}
None if value.get("result").is_some() || value.get("error").is_some() => {
self.control_post(value, session).await
}
method => panic!("unexpected scripted method: {method:?}"),
}
}
async fn delete_session(
&self,
_uri: Arc<str>,
_session: Arc<str>,
_auth_header: Option<String>,
_custom_headers: HashMap<HeaderName, HeaderValue>,
) -> Result<(), StreamableHttpError<Self::Error>> {
assert_eq!(
self.counts.active.load(SeqCst),
0,
"POSTs must stop before deleting the session"
);
self.counts.deleted.fetch_add(1, SeqCst);
Ok(())
}
async fn get_stream(
&self,
_uri: Arc<str>,
_session: Option<Arc<str>>,
_last_event_id: Option<String>,
_auth_header: Option<String>,
_custom_headers: HashMap<HeaderName, HeaderValue>,
) -> Result<BoxStream<'static, Result<Sse, SseError>>, StreamableHttpError<Self::Error>> {
Ok(match self.incoming.lock().await.take() {
Some(incoming) => UnboundedReceiverStream::new(incoming).boxed(),
None => futures::stream::pending().boxed(),
})
}
}
struct Harness {
client: RunningService<RoleClient, ClientInfo>,
started: mpsc::UnboundedReceiver<Posted>,
controls: mpsc::UnboundedReceiver<ControlPost>,
incoming: mpsc::UnboundedSender<Result<Sse, SseError>>,
reinitializations: mpsc::UnboundedReceiver<oneshot::Sender<()>>,
counts: Arc<Counts>,
}
fn config() -> StreamableHttpClientTransportConfig {
StreamableHttpClientTransportConfig::with_uri("http://scripted/mcp")
}
fn transport_error(error: &anyhow::Error) -> &StreamableHttpError<io::Error> {
let service_error = error
.downcast_ref::<rmcp::ServiceError>()
.expect("expected a service error");
let rmcp::ServiceError::TransportSend(transport_error) = service_error else {
panic!("expected a transport error, got {service_error:?}");
};
transport_error
.error
.downcast_ref::<StreamableHttpError<io::Error>>()
.expect("expected a streamable http error")
}
fn assert_recovery_timeout(error: anyhow::Error) {
assert!(matches!(
transport_error(&error),
StreamableHttpError::SessionRecoveryTimeout
));
}
impl Harness {
async fn start(config: StreamableHttpClientTransportConfig) -> anyhow::Result<Self> {
Self::with_lifecycle(config, ClientLifecycleMode::Initialize).await
}
async fn with_lifecycle(
config: StreamableHttpClientTransportConfig,
lifecycle: ClientLifecycleMode,
) -> anyhow::Result<Self> {
let (started, requests) = mpsc::unbounded_channel();
let (control_tx, controls) = mpsc::unbounded_channel();
let (incoming, incoming_rx) = mpsc::unbounded_channel();
let (reinitializing, reinitializations) = mpsc::unbounded_channel();
let counts = Arc::new(Counts::default());
let transport = StreamableHttpClientTransport::with_client(
ScriptedClient {
started,
controls: control_tx,
incoming: Arc::new(Mutex::new(Some(incoming_rx))),
reinitializing,
counts: counts.clone(),
},
config,
);
let client =
serve_client_with_lifecycle(ClientInfo::default(), transport, lifecycle).await?;
Ok(Self {
client,
started: requests,
controls,
incoming,
reinitializations,
counts,
})
}
fn call(&self, name: impl Into<String>) -> Call {
let name = name.into();
let peer = self.client.peer().clone();
tokio::spawn(async move {
let result = peer
.call_tool(CallToolRequestParams::new(name.clone()))
.await?;
anyhow::ensure!(serde_json::to_value(result)?["content"][0]["text"] == name);
Ok(())
})
}
async fn cancellable(&self, name: &'static str) -> anyhow::Result<RequestHandle<RoleClient>> {
self.cancellable_with_options(name, PeerRequestOptions::no_options())
.await
}
async fn cancellable_with_options(
&self,
name: &'static str,
options: PeerRequestOptions,
) -> anyhow::Result<RequestHandle<RoleClient>> {
Ok(self
.client
.peer()
.send_cancellable_request(
ClientRequest::CallToolRequest(Request::new(CallToolRequestParams::new(name))),
options,
)
.await?)
}
fn notify_cancellation(&self, id: RequestId) -> JoinHandle<Result<(), rmcp::ServiceError>> {
let peer = self.client.peer().clone();
tokio::spawn(async move {
peer.notify_cancelled(CancelledNotificationParam::new(Some(id), None))
.await
})
}
async fn next(&mut self) -> Posted {
next_event(&mut self.started).await
}
async fn next_control(&mut self) -> ControlPost {
next_event(&mut self.controls).await
}
async fn exchange_ping(&mut self, id: &str) {
self.incoming
.send(sse(json!({ "jsonrpc": "2.0", "id": id, "method": "ping" })))
.expect("common SSE stream remains open");
let control = self.next_control().await;
assert_eq!(control.message["id"], id);
assert!(control.message["result"].is_object());
assert_eq!(control.session.as_deref(), Some("session-1"));
control
.reply
.send(Ok(StreamableHttpPostResponse::Accepted))
.expect("reply POST is still waiting");
}
async fn finish(
self,
calls: Vec<Call>,
posted: usize,
initialized: usize,
) -> anyhow::Result<()> {
for call in calls {
timeout(TEST_TIMEOUT, call).await???;
}
assert_eq!(self.counts.posted.load(SeqCst), posted);
assert_eq!(self.counts.initialized.load(SeqCst), initialized);
assert_eq!(self.counts.active.load(SeqCst), 0);
self.client.cancel().await?;
Ok(())
}
}
#[tokio::test]
async fn json_limits_allow_overlap_and_preserve_response_ids() -> anyhow::Result<()> {
let mut zero = config();
zero.max_concurrent_requests = 0;
for (config, limit, total) in [
(config().max_concurrent_requests(2), 2, 5),
(config().max_concurrent_requests(1), 1, 3),
(zero, 1, 2),
(config(), 16, 17),
] {
let mut harness = Harness::start(config).await?;
let calls = (0..total)
.map(|index| harness.call(format!("request-{index}")))
.collect();
let mut pending = Vec::new();
for _ in 0..limit {
pending.push(harness.next().await);
}
assert_eq!(harness.counts.active.load(SeqCst), limit);
for _ in limit..total {
pending.pop().unwrap().succeed();
pending.push(harness.next().await);
}
for request in pending.into_iter().rev() {
request.succeed();
}
let counts = harness.counts.clone();
harness.finish(calls, total, 1).await?;
assert_eq!(counts.peak.load(SeqCst), limit);
}
Ok(())
}
#[tokio::test]
async fn early_sse_response_releases_the_post_slot() -> anyhow::Result<()> {
let mut harness = Harness::start(config().max_concurrent_requests(1)).await?;
let first = harness.call("first");
let release = harness.next().await.start_sse().await?;
let second = harness.call("second");
harness.next().await.succeed();
timeout(TEST_TIMEOUT, second).await???;
assert!(!first.is_finished(), "the SSE response is still blocked");
release.send(()).unwrap();
assert_eq!(harness.counts.peak.load(SeqCst), 1);
harness.finish(vec![first], 2, 1).await
}
#[tokio::test]
async fn a_common_stream_response_finishes_the_request_before_its_post_returns()
-> anyhow::Result<()> {
let (stream, incoming) = mpsc::unbounded_channel();
for (late, orphan) in [
(Ok(StreamableHttpPostResponse::Accepted), None),
(Err(StreamableHttpError::SessionExpired), None),
(
Ok(StreamableHttpPostResponse::Sse(
UnboundedReceiverStream::new(incoming).boxed(),
None,
)),
Some(stream),
),
] {
let mut harness = Harness::start(config().max_concurrent_requests(1)).await?;
let request = harness.cancellable("completed-on-common-stream").await?;
let post = harness.next().await;
harness
.incoming
.send(sse(serde_json::to_value(post.result())?))?;
timeout(TEST_TIMEOUT, request.await_response()).await??;
post.finish_and_wait(late).await?;
let next = harness.call("after-completed-request");
let post = harness.next().await;
assert_eq!(
post.name, "after-completed-request",
"do not retry completed work"
);
assert_eq!(post.session.as_deref(), Some("session-1"));
post.succeed();
timeout(TEST_TIMEOUT, next).await???;
if let Some(stream) = orphan {
timeout(Duration::from_secs(1), stream.closed())
.await
.expect("a late SSE response must not leave an orphan stream");
}
harness.finish(vec![], 2, 1).await?;
}
Ok(())
}
#[tokio::test]
async fn common_stream_response_keeps_numeric_and_string_ids_distinct() -> anyhow::Result<()> {
use rmcp::transport::Transport;
let (started, mut requests) = mpsc::unbounded_channel();
let (incoming, incoming_rx) = mpsc::unbounded_channel();
let counts = Arc::new(Counts::default());
let mut transport = StreamableHttpClientTransport::with_client(
ScriptedClient {
started,
controls: mpsc::unbounded_channel().0,
incoming: Arc::new(Mutex::new(Some(incoming_rx))),
reinitializing: mpsc::unbounded_channel().0,
counts: counts.clone(),
},
config(),
);
for message in [
json!({ "jsonrpc": "2.0", "id": 0, "method": "initialize", "params": ClientInfo::default() }),
json!({ "jsonrpc": "2.0", "method": "notifications/initialized" }),
] {
timeout(
TEST_TIMEOUT,
transport.send(serde_json::from_value(message)?),
)
.await??;
}
assert!(timeout(TEST_TIMEOUT, transport.receive()).await?.is_some());
let call = |id, name| {
ClientJsonRpcMessage::request(
ClientRequest::CallToolRequest(Request::new(CallToolRequestParams::new(name))),
id,
)
};
let mut numeric_send = Box::pin(transport.send(call(RequestId::Number(7), "numeric")));
assert!(futures::poll!(numeric_send.as_mut()).is_pending());
let numeric_post = next_event(&mut requests).await;
let numeric_result = numeric_post.result();
let release_numeric = numeric_post.start_sse().await?;
timeout(TEST_TIMEOUT, numeric_send).await??;
let mut string_send = Box::pin(transport.send(call(RequestId::String("7".into()), "string")));
assert!(futures::poll!(string_send.as_mut()).is_pending());
let string_post = next_event(&mut requests).await;
let string_result = serde_json::to_value(string_post.result())?;
string_post.finish(Ok(StreamableHttpPostResponse::Accepted));
timeout(TEST_TIMEOUT, string_send).await??;
incoming.send(sse(string_result.clone()))?;
let received = timeout(TEST_TIMEOUT, transport.receive())
.await?
.expect("string-id response");
assert_eq!(serde_json::to_value(received)?, string_result);
release_numeric
.send(())
.expect("the distinct numeric-id stream must remain open");
let received = timeout(TEST_TIMEOUT, transport.receive())
.await?
.expect("numeric-id response");
assert_eq!(
serde_json::to_value(received)?,
serde_json::to_value(numeric_result)?
);
assert_eq!(counts.posted.load(SeqCst), 2);
assert_eq!(counts.cancelled.load(SeqCst), 0);
transport.close().await?;
Ok(())
}
#[tokio::test]
async fn a_common_stream_response_prevents_replay_during_session_recovery() -> anyhow::Result<()> {
let mut harness = Harness::start(config().max_concurrent_requests(2)).await?;
let blocking = harness.call("blocking-recovery");
let blocked = harness.next().await;
let request = harness.cancellable("completed-during-recovery").await?;
let expired = harness.next().await;
let final_response = sse(serde_json::to_value(expired.result())?);
expired.expire_and_wait().await?;
harness.incoming.send(final_response)?;
timeout(TEST_TIMEOUT, request.await_response()).await??;
blocked.succeed();
timeout(TEST_TIMEOUT, blocking).await???;
let next = harness.call("after-completed-request");
let post = harness.next().await;
assert_eq!(
post.name, "after-completed-request",
"do not retry completed work"
);
let initializations = harness.counts.initialized.load(SeqCst);
assert!((1..=2).contains(&initializations));
post.succeed();
timeout(TEST_TIMEOUT, next).await???;
harness.finish(vec![], 3, initializations).await
}
#[tokio::test]
async fn concurrent_session_expiry_shares_one_reinitialization() -> anyhow::Result<()> {
let mut harness = Harness::start(config().max_concurrent_requests(2)).await?;
let calls = vec![harness.call("first"), harness.call("second")];
let first = harness.next().await;
let second = harness.next().await;
let mut originals = HashMap::new();
for request in [first, second] {
assert_eq!(request.session.as_deref(), Some("session-1"));
originals.insert(request.name.clone(), request.id.clone());
request.expire();
}
for _ in 0..2 {
let retry = harness.next().await;
assert_eq!(retry.session.as_deref(), Some("session-2"));
assert_eq!(originals.remove(&retry.name), Some(retry.id.clone()));
retry.succeed();
}
assert!(originals.is_empty());
harness.finish(calls, 4, 2).await
}
#[tokio::test]
async fn cancellation_still_runs_while_recovery_waits_for_old_posts() -> anyhow::Result<()> {
let mut harness = Harness::start(config().max_concurrent_requests(3)).await?;
let hanging = harness.cancellable("hanging").await?;
let mut blocked = harness.next().await;
let expired = harness.call("expired");
harness.next().await.expire_and_wait().await?;
assert_eq!(harness.counts.initialized.load(SeqCst), 1);
assert_eq!(harness.counts.active.load(SeqCst), 1);
timeout(Duration::from_secs(1), hanging.cancel(None))
.await
.expect("cancellation must bypass the session recovery wait")?;
timeout(TEST_TIMEOUT, blocked.reply.closed()).await?;
let retry = harness.next().await;
assert_eq!(retry.name, "expired");
assert_eq!(retry.session.as_deref(), Some("session-2"));
retry.succeed();
harness.finish(vec![expired], 3, 2).await
}
#[tokio::test]
async fn server_replies_still_run_while_recovery_waits_for_old_posts() -> anyhow::Result<()> {
let mut harness = Harness::start(config().max_concurrent_requests(2)).await?;
harness.counts.manual_controls.store(true, SeqCst);
let waiting = harness.call("waiting-for-client");
let blocked = harness.next().await;
let expired = harness.call("expired");
harness.next().await.expire_and_wait().await?;
harness.exchange_ping("recovery-ping").await;
blocked.succeed();
let retry = harness.next().await;
assert_eq!(retry.name, "expired");
assert_eq!(retry.session.as_deref(), Some("session-2"));
retry.succeed();
harness.finish(vec![waiting, expired], 3, 2).await
}
#[tokio::test]
async fn a_version_barrier_allows_the_server_reply_it_needs() -> anyhow::Result<()> {
let mut harness = Harness::start(config().max_concurrent_requests(2)).await?;
harness.counts.manual_controls.store(true, SeqCst);
let mut meta = RequestMetaObject::new();
meta.set_protocol_version(ProtocolVersion::V_2025_06_18);
let barrier = harness
.cancellable_with_options("barrier", PeerRequestOptions::no_options().with_meta(meta))
.await?;
let blocked = harness.next().await;
let queued = harness.cancellable("after-barrier").await?;
harness.exchange_ping("barrier-ping").await;
assert_eq!(harness.counts.posted.load(SeqCst), 1);
blocked.succeed();
timeout(TEST_TIMEOUT, barrier.await_response()).await??;
let next = harness.next().await;
assert_eq!(next.name, "after-barrier");
next.succeed();
timeout(TEST_TIMEOUT, queued.await_response()).await??;
harness.finish(vec![], 2, 1).await
}
#[tokio::test]
async fn a_cancelled_version_barrier_does_not_block_unrelated_posts() -> anyhow::Result<()> {
let mut harness = Harness::start(config().max_concurrent_requests(2)).await?;
let blocking = harness.call("blocking");
let blocked = harness.next().await;
let mut meta = RequestMetaObject::new();
meta.set_protocol_version(ProtocolVersion::V_2025_06_18);
let barrier = harness
.cancellable_with_options(
"cancelled-barrier",
PeerRequestOptions::no_options().with_meta(meta),
)
.await?;
timeout(TEST_TIMEOUT, barrier.cancel(None)).await??;
let next = harness.call("after-cancelled-barrier");
let post = harness.next().await;
assert_eq!(post.name, "after-cancelled-barrier");
post.succeed();
timeout(TEST_TIMEOUT, next).await???;
assert!(!blocked.reply.is_closed(), "the older POST is still held");
blocked.succeed();
harness.finish(vec![blocking], 2, 1).await
}
#[tokio::test]
async fn dropping_a_parked_barrier_send_unblocks_the_ordinary_queue() -> anyhow::Result<()> {
use std::task::{Context, Wake, Waker};
use rmcp::transport::Transport;
struct WakeSignal(mpsc::UnboundedSender<()>);
impl Wake for WakeSignal {
fn wake(self: Arc<Self>) {
let _ = self.0.send(());
}
}
let (started, mut requests) = mpsc::unbounded_channel();
let counts = Arc::new(Counts::default());
let mut config = config().max_concurrent_requests(2);
config.channel_buffer_capacity = 1;
let mut transport = StreamableHttpClientTransport::with_client(
ScriptedClient {
started,
controls: mpsc::unbounded_channel().0,
incoming: Arc::new(Mutex::new(None)),
reinitializing: mpsc::unbounded_channel().0,
counts: counts.clone(),
},
config,
);
for message in [
json!({ "jsonrpc": "2.0", "id": 0, "method": "initialize", "params": ClientInfo::default() }),
json!({ "jsonrpc": "2.0", "method": "notifications/initialized" }),
] {
timeout(
TEST_TIMEOUT,
transport.send(serde_json::from_value(message)?),
)
.await??;
}
assert!(timeout(TEST_TIMEOUT, transport.receive()).await?.is_some());
let call = |id, params| {
ClientJsonRpcMessage::request(
ClientRequest::CallToolRequest(Request::new(params)),
RequestId::Number(id),
)
};
let mut first = Box::pin(transport.send(call(1, CallToolRequestParams::new("held"))));
assert!(futures::poll!(first.as_mut()).is_pending());
let blocked = next_event(&mut requests).await;
let mut meta = RequestMetaObject::new();
meta.set_protocol_version(ProtocolVersion::V_2025_06_18);
let mut barrier_request = Request::new(CallToolRequestParams::new("abandoned-barrier"));
barrier_request.extensions.insert(meta);
let mut barrier = Box::pin(transport.send(ClientJsonRpcMessage::request(
ClientRequest::CallToolRequest(barrier_request),
RequestId::Number(2),
)));
assert!(futures::poll!(barrier.as_mut()).is_pending());
let mut next = Box::pin(transport.send(call(3, CallToolRequestParams::new("after-barrier"))));
let (wake, mut woke) = mpsc::unbounded_channel();
let waker = Waker::from(Arc::new(WakeSignal(wake)));
assert!(
next.as_mut()
.poll(&mut Context::from_waker(&waker))
.is_pending()
);
next_event(&mut woke).await;
assert_eq!(counts.posted.load(SeqCst), 1, "the barrier must be parked");
drop(barrier);
assert!(futures::poll!(next.as_mut()).is_pending());
let post = timeout(Duration::from_secs(1), requests.recv())
.await
.expect("dropping the parked send must wake the worker")
.unwrap();
assert_eq!(post.name, "after-barrier");
post.succeed();
timeout(TEST_TIMEOUT, next).await??;
assert!(timeout(TEST_TIMEOUT, transport.receive()).await?.is_some());
assert!(!blocked.reply.is_closed(), "the first POST remains held");
blocked.succeed();
timeout(TEST_TIMEOUT, first).await??;
assert_eq!(counts.posted.load(SeqCst), 2);
assert_eq!(counts.cancelled.load(SeqCst), 0);
transport.close().await?;
Ok(())
}
#[tokio::test]
async fn recovery_deadline_drops_ambiguous_posts_without_retrying_them() -> anyhow::Result<()> {
assert_eq!(config().session_recovery_timeout, Duration::from_secs(5));
let mut harness = Harness::start(
config()
.max_concurrent_requests(3)
.session_recovery_timeout(Duration::from_millis(50)),
)
.await?;
let ambiguous = harness.call("possibly-applied");
let mut blocked = harness.next().await;
let expired = harness.call("expired");
let rejected = harness.next().await;
let rejected_id = rejected.id.clone();
rejected.expire_and_wait().await?;
assert_recovery_timeout(timeout(TEST_TIMEOUT, ambiguous).await??.unwrap_err());
timeout(TEST_TIMEOUT, blocked.reply.closed()).await?;
let retry = harness.next().await;
assert_eq!(retry.name, "expired");
assert_eq!(retry.id, rejected_id);
assert_eq!(retry.session.as_deref(), Some("session-2"));
retry.succeed();
harness.finish(vec![expired], 3, 2).await
}
#[tokio::test]
async fn reinitialization_has_its_own_deadline() -> anyhow::Result<()> {
let mut harness =
Harness::start(config().session_recovery_timeout(Duration::from_millis(50))).await?;
harness.counts.hold_reinitialization.store(true, SeqCst);
let expired = harness.call("expired");
harness.next().await.expire_and_wait().await?;
let mut reinitialization = next_event(&mut harness.reinitializations).await;
assert_recovery_timeout(timeout(TEST_TIMEOUT, expired).await??.unwrap_err());
timeout(TEST_TIMEOUT, reinitialization.closed()).await?;
harness.finish(vec![], 1, 2).await
}
#[tokio::test]
async fn an_expired_retry_is_not_retried_again() -> anyhow::Result<()> {
let mut harness = Harness::start(config().max_concurrent_requests(2)).await?;
let call = harness.call("expires-twice");
let first = harness.next().await;
let id = first.id.clone();
first.expire();
let retry = harness.next().await;
assert_eq!(retry.id, id);
assert_eq!(retry.session.as_deref(), Some("session-2"));
retry.expire();
let error = timeout(TEST_TIMEOUT, call).await??.unwrap_err();
assert!(error.to_string().contains("Session expired"));
harness.finish(vec![], 2, 2).await
}
#[tokio::test]
async fn a_lost_post_response_is_not_retried() -> anyhow::Result<()> {
let mut harness = Harness::start(config().max_concurrent_requests(2)).await?;
let call = harness.call("possibly-applied");
harness
.next()
.await
.finish(Err(StreamableHttpError::Client(io::Error::new(
io::ErrorKind::ConnectionReset,
"scripted response lost",
))));
let error = timeout(TEST_TIMEOUT, call).await??.unwrap_err();
assert!(error.to_string().contains("scripted response lost"));
harness.finish(vec![], 1, 1).await
}
#[tokio::test]
async fn cancellation_bypasses_queued_posts_at_capacity() -> anyhow::Result<()> {
for (lifecycle, legacy_notifications, initializations) in [
(ClientLifecycleMode::Initialize, 2, 1),
(
ClientLifecycleMode::Discover {
preferred_versions: vec![ProtocolVersion::V_2026_07_28],
},
0,
0,
),
] {
let mut harness =
Harness::with_lifecycle(config().max_concurrent_requests(1), lifecycle).await?;
let request = harness.cancellable("cancel-me").await?;
let mut blocked = harness.next().await;
let cancelled = harness.cancellable("never-send").await?;
timeout(TEST_TIMEOUT, cancelled.cancel(None)).await??;
let queued = harness.cancellable("queued").await?;
timeout(TEST_TIMEOUT, request.cancel(None)).await??;
timeout(TEST_TIMEOUT, blocked.reply.closed()).await?;
let next = harness.next().await;
assert_eq!(next.name, "queued");
next.succeed();
timeout(TEST_TIMEOUT, queued.await_response()).await??;
assert_eq!(harness.counts.cancelled.load(SeqCst), legacy_notifications);
harness.finish(vec![], 2, initializations).await?;
}
Ok(())
}
#[tokio::test]
async fn legacy_cancellation_reaches_the_adapter_before_its_local_stream_drops()
-> anyhow::Result<()> {
let mut harness = Harness::start(config().max_concurrent_requests(1)).await?;
harness.counts.manual_controls.store(true, SeqCst);
let stale = harness.notify_cancellation(RequestId::Number(999));
let held = harness.next_control().await;
harness.counts.manual_controls.store(false, SeqCst);
harness
.counts
.reject_unmatched_cancellation
.store(true, SeqCst);
let event = harness.cancellable("local-events").await?;
let (eof, dropped) = harness
.next()
.await
.start_local_sse(harness.counts.clone())
.await?;
let peer = harness.client.peer().clone();
let mut cancel = Box::pin(peer.notify_cancelled(CancelledNotificationParam::new(
Some(event.id.clone()),
None,
)));
assert!(futures::poll!(cancel.as_mut()).is_pending());
let probe = harness.call("after-cancellation");
harness.next().await.succeed();
timeout(TEST_TIMEOUT, probe).await???;
eof.cancel();
let probe = harness.call("after-eof");
harness.next().await.succeed();
timeout(TEST_TIMEOUT, probe).await???;
assert!(
harness
.counts
.local_streams
.lock()
.unwrap()
.contains_key(&event.id),
"the adapter's stream registration must survive queued cancellation"
);
assert!(!held.reply.is_closed(), "the control slot is still held");
held.reply
.send(Ok(StreamableHttpPostResponse::Accepted))
.unwrap();
timeout(TEST_TIMEOUT, stale).await???;
timeout(TEST_TIMEOUT, cancel).await??;
timeout(TEST_TIMEOUT, dropped).await??;
assert!(matches!(
timeout(TEST_TIMEOUT, event.await_response()).await?,
Err(rmcp::ServiceError::Cancelled { .. })
));
assert_eq!(harness.counts.local_cancellations.load(SeqCst), 1);
assert_eq!(harness.counts.cancelled.load(SeqCst), 1);
harness.finish(vec![], 3, 1).await
}
#[tokio::test]
async fn control_timeout_is_configurable_and_releases_its_slot() -> anyhow::Result<()> {
assert_eq!(config().control_request_timeout, Duration::from_secs(5));
let mut harness =
Harness::start(config().control_request_timeout(Duration::from_millis(50))).await?;
harness.counts.manual_controls.store(true, SeqCst);
let cancellation = harness.notify_cancellation(RequestId::Number(999));
let mut cancelled_post = harness.next_control().await;
harness.incoming.send(sse(json!({
"jsonrpc": "2.0", "id": "timed-reply", "method": "ping"
})))?;
let error = timeout(Duration::from_secs(1), cancellation)
.await??
.unwrap_err();
assert!(matches!(
transport_error(&error.into()),
StreamableHttpError::ControlRequestTimeout
));
timeout(Duration::from_secs(1), cancelled_post.reply.closed()).await?;
let mut reply_post = harness.next_control().await;
assert_eq!(reply_post.message["id"], "timed-reply");
assert!(reply_post.message["result"].is_object());
harness.counts.manual_controls.store(false, SeqCst);
let next = harness.notify_cancellation(RequestId::Number(1000));
timeout(Duration::from_secs(1), reply_post.reply.closed()).await?;
timeout(Duration::from_secs(1), next).await???;
assert_eq!(harness.counts.cancelled.load(SeqCst), 2);
harness.finish(vec![], 0, 1).await
}
#[tokio::test]
async fn a_hanging_legacy_control_does_not_delay_local_cancellation() -> anyhow::Result<()> {
let mut harness = Harness::start(config().max_concurrent_requests(1)).await?;
harness.counts.manual_controls.store(true, SeqCst);
let stale = harness.notify_cancellation(RequestId::Number(999));
let mut held = harness.next_control().await;
assert_eq!(held.message["params"]["requestId"], 999);
harness.counts.manual_controls.store(false, SeqCst);
let live = harness.cancellable("live-post").await?;
let mut blocked = harness.next().await;
let cancel_post = tokio::spawn(async move { live.cancel(None).await });
timeout(Duration::from_secs(1), blocked.reply.closed()).await?;
let streaming = harness.cancellable("live-stream").await?;
let mut stream = harness.next().await.start_sse().await?;
let cancel_stream = harness.notify_cancellation(streaming.id.clone());
assert!(
!held.reply.is_closed(),
"the old control POST is still held"
);
assert_eq!(harness.counts.cancelled.load(SeqCst), 1);
let error = timeout(Duration::from_secs(10), stale).await??.unwrap_err();
assert!(matches!(
transport_error(&error.into()),
StreamableHttpError::ControlRequestTimeout
));
timeout(TEST_TIMEOUT, held.reply.closed()).await?;
timeout(TEST_TIMEOUT, cancel_post).await???;
timeout(TEST_TIMEOUT, cancel_stream).await???;
timeout(TEST_TIMEOUT, stream.closed()).await?;
assert!(matches!(
timeout(TEST_TIMEOUT, streaming.await_response()).await?,
Err(rmcp::ServiceError::Cancelled { .. })
));
assert_eq!(harness.counts.cancelled.load(SeqCst), 3);
harness.finish(vec![], 2, 1).await
}
#[tokio::test]
async fn close_drops_blocked_posts_before_deleting_the_session() -> anyhow::Result<()> {
let mut harness = Harness::start(config().max_concurrent_requests(2)).await?;
let calls = [harness.call("first"), harness.call("second")];
let posts = [harness.next().await, harness.next().await];
timeout(TEST_TIMEOUT, harness.client.cancel()).await??;
assert!(posts.iter().all(|post| post.reply.is_closed()));
assert_eq!(harness.counts.active.load(SeqCst), 0);
assert_eq!(harness.counts.deleted.load(SeqCst), 1);
for call in calls {
assert!(timeout(TEST_TIMEOUT, call).await??.is_err());
}
Ok(())
}