rust-sdk/crates/rmcp/tests/test_streamable_http_event_store.rs

#![cfg(all(
    feature = "client",
    feature = "server",
    feature = "transport-streamable-http-client-reqwest",
    feature = "transport-streamable-http-server",
    not(feature = "local")
))]

use std::{
    collections::HashMap,
    sync::{
        Arc,
        atomic::{AtomicU64, Ordering},
    },
};

use futures::StreamExt;
use rmcp::{
    ErrorData, ServerHandler,
    model::{
        CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, InitializeResult,
        ProgressNotificationParam, ServerCapabilities,
    },
    service::RequestContext,
    transport::streamable_http_server::{
        StreamableHttpServerConfig, StreamableHttpService,
        session::{
            EventStore, EventStoreError, EventStream, ServerSseMessage, SessionState, SessionStore,
            SessionStoreError, local::LocalSessionManager, never::NeverSessionManager,
        },
    },
};
use tokio::sync::RwLock;
use tokio_util::sync::CancellationToken;

#[derive(Clone, Default)]
struct InMemorySessionStore(Arc<RwLock<HashMap<String, SessionState>>>);

#[async_trait::async_trait]
impl SessionStore for InMemorySessionStore {
    async fn load(&self, session_id: &str) -> Result<Option<SessionState>, SessionStoreError> {
        Ok(self.0.read().await.get(session_id).cloned())
    }

    async fn store(&self, session_id: &str, state: &SessionState) -> Result<(), SessionStoreError> {
        self.0
            .write()
            .await
            .insert(session_id.to_owned(), state.clone());
        Ok(())
    }

    async fn delete(&self, session_id: &str) -> Result<(), SessionStoreError> {
        self.0.write().await.remove(session_id);
        Ok(())
    }
}

#[derive(Clone)]
struct StoredEvent {
    stream_id: String,
    event: ServerSseMessage,
}

#[derive(Clone, Default)]
struct InMemoryEventStore {
    events: Arc<RwLock<Vec<StoredEvent>>>,
    next_id: Arc<AtomicU64>,
}

#[async_trait::async_trait]
impl EventStore for InMemoryEventStore {
    async fn store_event(
        &self,
        stream_id: &str,
        event: &ServerSseMessage,
    ) -> Result<String, EventStoreError> {
        let event_id = format!("event-{}", self.next_id.fetch_add(1, Ordering::Relaxed));
        let mut event = event.clone();
        event.event_id = Some(event_id.clone());
        self.events.write().await.push(StoredEvent {
            stream_id: stream_id.to_owned(),
            event,
        });
        Ok(event_id)
    }

    async fn replay_events_after(
        &self,
        last_event_id: &str,
    ) -> Result<EventStream, EventStoreError> {
        let events = self.events.read().await;
        let last_index = events
            .iter()
            .position(|stored| stored.event.event_id.as_deref() == Some(last_event_id))
            .ok_or_else(|| std::io::Error::other("event not found"))?;
        let stream_id = events[last_index].stream_id.clone();
        let replay = events
            .iter()
            .skip(last_index + 1)
            .filter(|stored| stored.stream_id == stream_id)
            .map(|stored| stored.event.clone())
            .collect::<Vec<_>>();
        Ok(Box::pin(futures::stream::iter(replay)))
    }
}

#[derive(Clone)]
struct ProgressServer;

impl ServerHandler for ProgressServer {
    fn get_info(&self) -> InitializeResult {
        InitializeResult::new(ServerCapabilities::builder().enable_tools().build())
    }

    async fn call_tool(
        &self,
        request: CallToolRequestParams,
        context: RequestContext<rmcp::RoleServer>,
    ) -> Result<CallToolResponse, ErrorData> {
        if request.name == "progress" || request.name == "slow-progress" {
            let progress_token = context
                .meta
                .get_progress_token()
                .expect("request includes progressToken");
            context
                .peer
                .notify_progress(
                    ProgressNotificationParam::new(progress_token, 50.0)
                        .with_total(100.0)
                        .with_message("working"),
                )
                .await
                .expect("progress notification is delivered");
        }
        if request.name == "slow-progress" {
            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
        }
        Ok(CallToolResult::success(vec![ContentBlock::text("done")]).into())
    }
}

async fn spawn_server(
    session_store: Arc<dyn SessionStore>,
    event_store: Arc<dyn EventStore>,
    cancellation_token: &CancellationToken,
    legacy_session_mode: bool,
) -> anyhow::Result<(String, tokio::task::JoinHandle<()>)> {
    let config = {
        let mut config = StreamableHttpServerConfig::default();
        config.sse_keep_alive = None;
        config.legacy_session_mode = legacy_session_mode;
        config.cancellation_token = cancellation_token.child_token();
        config.session_store = Some(session_store);
        config
    };
    let router = if legacy_session_mode {
        let session_manager =
            Arc::new(LocalSessionManager::default().with_event_store(event_store));
        let service = StreamableHttpService::new(|| Ok(ProgressServer), session_manager, config);
        axum::Router::new().nest_service("/mcp", service)
    } else {
        let session_manager =
            Arc::new(NeverSessionManager::default().with_event_store(event_store));
        let service = StreamableHttpService::new(|| Ok(ProgressServer), session_manager, config);
        axum::Router::new().nest_service("/mcp", service)
    };
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
    let address = listener.local_addr()?;
    let handle = tokio::spawn({
        let cancellation_token = cancellation_token.clone();
        async move {
            let _ = axum::serve(listener, router)
                .with_graceful_shutdown(async move { cancellation_token.cancelled_owned().await })
                .await;
        }
    });
    Ok((format!("http://{address}/mcp"), handle))
}

fn event_id_containing<'a>(body: &'a str, needle: &str) -> Option<&'a str> {
    body.split("\n\n")
        .find(|event| event.contains(needle))?
        .lines()
        .find_map(|line| line.strip_prefix("id: "))
}

#[tokio::test]
async fn restored_instance_replays_events_from_shared_store() -> anyhow::Result<()> {
    let session_store: Arc<dyn SessionStore> = Arc::new(InMemorySessionStore::default());
    let event_store: Arc<dyn EventStore> = Arc::new(InMemoryEventStore::default());
    let http = reqwest::Client::new();

    let cancellation_a = CancellationToken::new();
    let (url_a, server_a) = spawn_server(
        session_store.clone(),
        event_store.clone(),
        &cancellation_a,
        true,
    )
    .await?;
    let initialize = http
        .post(&url_a)
        .header("Content-Type", "application/json")
        .header("Accept", "application/json, text/event-stream")
        .body(r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}"#)
        .send()
        .await?;
    let session_id = initialize
        .headers()
        .get("mcp-session-id")
        .expect("initialize returns a session ID")
        .to_str()?
        .to_owned();
    let _ = initialize.text().await?;

    let initialized_status = http
        .post(&url_a)
        .header("Content-Type", "application/json")
        .header("Accept", "application/json, text/event-stream")
        .header("Mcp-Session-Id", &session_id)
        .header("Mcp-Protocol-Version", "2025-06-18")
        .body(r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#)
        .send()
        .await?
        .status();
    assert_eq!(initialized_status, reqwest::StatusCode::ACCEPTED);

    let original_body = http
        .post(&url_a)
        .header("Content-Type", "application/json")
        .header("Accept", "application/json, text/event-stream")
        .header("Mcp-Session-Id", &session_id)
        .header("Mcp-Protocol-Version", "2025-06-18")
        .body(r#"{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"progress","arguments":{},"_meta":{"progressToken":"progress-1"}}}"#)
        .send()
        .await?
        .text()
        .await?;
    let progress_event_id = event_id_containing(&original_body, "notifications/progress")
        .expect("progress event has a persisted event ID")
        .to_owned();

    cancellation_a.cancel();
    server_a.await?;

    let cancellation_b = CancellationToken::new();
    let (url_b, server_b) = spawn_server(session_store, event_store, &cancellation_b, true).await?;
    let replay = http
        .get(&url_b)
        .header("Accept", "text/event-stream")
        .header("Mcp-Session-Id", &session_id)
        .header("Mcp-Protocol-Version", "2025-06-18")
        .header("Last-Event-ID", progress_event_id)
        .send()
        .await?;
    assert_eq!(replay.status(), reqwest::StatusCode::OK);
    let replay_body = replay.text().await?;
    assert!(
        replay_body.contains(r#""id":2"#),
        "instance B should replay the final response stored by instance A: {replay_body}"
    );

    cancellation_b.cancel();
    server_b.await?;
    Ok(())
}

#[tokio::test]
async fn stateless_instance_replays_events_from_shared_store() -> anyhow::Result<()> {
    let session_store: Arc<dyn SessionStore> = Arc::new(InMemorySessionStore::default());
    let event_store = Arc::new(InMemoryEventStore::default());
    let http = reqwest::Client::new();

    let cancellation_a = CancellationToken::new();
    let (url_a, server_a) = spawn_server(
        session_store.clone(),
        event_store.clone(),
        &cancellation_a,
        false,
    )
    .await?;
    let original = http
        .post(&url_a)
        .header("Content-Type", "application/json")
        .header("Accept", "application/json, text/event-stream")
        .header("MCP-Protocol-Version", "2026-07-28")
        .header("Mcp-Method", "tools/call")
        .header("Mcp-Name", "slow-progress")
        .body(
            r#"{
                "jsonrpc":"2.0",
                "id":2,
                "method":"tools/call",
                "params":{
                    "name":"slow-progress",
                    "arguments":{},
                    "_meta":{
                        "progressToken":"progress-1",
                        "io.modelcontextprotocol/protocolVersion":"2026-07-28",
                        "io.modelcontextprotocol/clientInfo":{"name":"test","version":"1.0"},
                        "io.modelcontextprotocol/clientCapabilities":{}
                    }
                }
            }"#,
        )
        .send()
        .await?;
    assert_eq!(original.status(), reqwest::StatusCode::OK);
    assert!(
        original.headers().get("mcp-session-id").is_none(),
        "stateless response must not create a session"
    );
    let mut body = original.bytes_stream();
    let mut received = String::new();
    let progress_event_id = tokio::time::timeout(std::time::Duration::from_secs(2), async {
        loop {
            let chunk = body
                .next()
                .await
                .expect("response remains open until progress arrives")?;
            received.push_str(&String::from_utf8_lossy(&chunk));
            if let Some(event_id) = event_id_containing(&received, "notifications/progress") {
                return Ok::<_, reqwest::Error>(event_id.to_owned());
            }
        }
    })
    .await??;
    drop(body);

    tokio::time::timeout(std::time::Duration::from_secs(2), async {
        loop {
            let stored_response = event_store.events.read().await.iter().any(|stored| {
                stored.event.message.as_ref().is_some_and(|message| {
                    matches!(
                        message.as_ref(),
                        rmcp::model::ServerJsonRpcMessage::Response(_)
                    )
                })
            });
            if stored_response {
                break;
            }
            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
        }
    })
    .await?;

    cancellation_a.cancel();
    server_a.await?;

    let cancellation_b = CancellationToken::new();
    let (url_b, server_b) =
        spawn_server(session_store, event_store, &cancellation_b, false).await?;
    let replay = http
        .get(&url_b)
        .header("Accept", "text/event-stream")
        .header("MCP-Protocol-Version", "2026-07-28")
        .header("Last-Event-ID", progress_event_id)
        .send()
        .await?;
    assert_eq!(replay.status(), reqwest::StatusCode::OK);
    let replay_body = replay.text().await?;
    assert!(
        replay_body.contains(r#""id":2"#),
        "instance B should replay the stateless response stored by instance A: {replay_body}"
    );

    cancellation_b.cancel();
    server_b.await?;
    Ok(())
}