fix: include http_request_id in request-wise priming event IDs

Author: DaleSeo

crates/rmcp/src/transport/streamable_http_server/session/local.rs

@@ -4,7 +4,7 @@ use std::{
     time::Duration,
 };
 
-use futures::Stream;
+use futures::{Stream, StreamExt};
 use thiserror::Error;
 use tokio::sync::{
     mpsc::{Receiver, Sender},
@@ -86,10 +86,20 @@ impl SessionManager for LocalSessionManager {
             .get(id)
             .ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
         let receiver = handle.establish_request_wise_channel().await?;
-        handle
-            .push_message(message, receiver.http_request_id)
-            .await?;
-        Ok(ReceiverStream::new(receiver.inner))
+        let http_request_id = receiver.http_request_id;
+        handle.push_message(message, http_request_id).await?;
+
+        let priming_events: Vec<ServerSseMessage> = match self.session_config.sse_retry {
+            Some(retry) => {
+                let event_id = match http_request_id {
+                    Some(id) => format!("0/{id}"),
+                    None => "0".into(),
+                };
+                vec![ServerSseMessage::priming(event_id, retry)]
+            }
+            None => vec![],
+        };
+        Ok(futures::stream::iter(priming_events).chain(ReceiverStream::new(receiver.inner)))
     }
 
     async fn create_standalone_stream(
@@ -188,23 +198,29 @@ struct CachedTx {
     cache: VecDeque<ServerSseMessage>,
     http_request_id: Option<HttpRequestId>,
     capacity: usize,
+    starting_index: usize,
 }
 
 impl CachedTx {
-    fn new(tx: Sender<ServerSseMessage>, http_request_id: Option<HttpRequestId>) -> Self {
+    fn new(
+        tx: Sender<ServerSseMessage>,
+        http_request_id: Option<HttpRequestId>,
+        starting_index: usize,
+    ) -> Self {
         Self {
             cache: VecDeque::with_capacity(tx.capacity()),
             capacity: tx.capacity(),
             tx,
             http_request_id,
+            starting_index,
         }
     }
     fn new_common(tx: Sender<ServerSseMessage>) -> Self {
-        Self::new(tx, None)
+        Self::new(tx, None, 0)
     }
 
     fn next_event_id(&self) -> EventId {
-        let index = self.cache.back().map_or(0, |m| {
+        let index = self.cache.back().map_or(self.starting_index, |m| {
             m.event_id
                 .as_deref()
                 .unwrap_or_default()
@@ -405,11 +421,16 @@ impl LocalSessionWorker {
     ) -> Result<StreamableHttpMessageReceiver, SessionError> {
         let http_request_id = self.next_http_request_id();
         let (tx, rx) = tokio::sync::mpsc::channel(self.session_config.channel_capacity);
+        let starting_index = if self.session_config.sse_retry.is_some() {
+            1
+        } else {
+            0
+        };
         self.tx_router.insert(
             http_request_id,
             HttpRequestWise {
                 resources: Default::default(),
-                tx: CachedTx::new(tx, Some(http_request_id)),
+                tx: CachedTx::new(tx, Some(http_request_id), starting_index),
             },
         );
         tracing::debug!(http_request_id, "establish new request wise channel");
@@ -1072,18 +1093,25 @@ pub struct SessionConfig {
     /// Defaults to 5 minutes. Set to `None` to disable (not recommended
     /// for long-running servers behind proxies).
     pub keep_alive: Option<Duration>,
+    /// SSE retry interval for priming events on request-wise streams.
+    /// When set, the session layer prepends a priming event with the correct
+    /// stream-identifying event ID to each request-wise SSE stream.
+    /// Default is 3 seconds, matching [`StreamableHttpServerConfig::default()`].
+    pub sse_retry: Option<Duration>,
 }
 
 impl SessionConfig {
     pub const DEFAULT_CHANNEL_CAPACITY: usize = 16;
     pub const DEFAULT_KEEP_ALIVE: Duration = Duration::from_secs(300);
+    pub const DEFAULT_SSE_RETRY: Duration = Duration::from_secs(3);
 }
 
 impl Default for SessionConfig {
     fn default() -> Self {
         Self {
             channel_capacity: Self::DEFAULT_CHANNEL_CAPACITY,
             keep_alive: Some(Self::DEFAULT_KEEP_ALIVE),
+            sse_retry: Some(Self::DEFAULT_SSE_RETRY),
         }
     }
 }

crates/rmcp/src/transport/streamable_http_server/tower.rs

@@ -598,20 +598,14 @@ where
 
                 match message {
                     ClientJsonRpcMessage::Request(_) => {
+                        // Priming for request-wise streams is handled by the
+                        // session layer (SessionManager::create_stream) which
+                        // has access to the http_request_id for correct event IDs.
                         let stream = self
                             .session_manager
                             .create_stream(&session_id, message)
                             .await
                             .map_err(internal_error_response("get session"))?;
-                        // Prepend priming event if sse_retry configured
-                        let stream = if let Some(retry) = self.config.sse_retry {
-                            let priming = ServerSseMessage::priming("0", retry);
-                            futures::stream::once(async move { priming })
-                                .chain(stream)
-                                .left_stream()
-                        } else {
-                            stream.right_stream()
-                        };
                         Ok(sse_stream_response(
                             stream,
                             self.config.sse_keep_alive,

crates/rmcp/tests/test_streamable_http_priming.rs

@@ -2,7 +2,8 @@
 use std::time::Duration;
 
 use rmcp::transport::streamable_http_server::{
-    StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
+    StreamableHttpServerConfig, StreamableHttpService,
+    session::{SessionId, local::LocalSessionManager},
 };
 use tokio_util::sync::CancellationToken;
 
@@ -54,7 +55,7 @@ async fn test_priming_on_stream_start() -> anyhow::Result<()> {
     let events: Vec<&str> = body.split("\n\n").filter(|e| !e.is_empty()).collect();
     assert!(events.len() >= 2);
 
-    // Verify priming event (first event)
+    // Verify priming event (first event) — initialize uses "0" (no http_request_id)
     let priming_event = events[0];
     assert!(priming_event.contains("id: 0"));
     assert!(priming_event.contains("retry: 3000"));
@@ -71,6 +72,131 @@ async fn test_priming_on_stream_start() -> anyhow::Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn test_request_wise_priming_includes_http_request_id() -> anyhow::Result<()> {
+    let ct = CancellationToken::new();
+
+    let service: StreamableHttpService<Calculator, LocalSessionManager> =
+        StreamableHttpService::new(
+            || Ok(Calculator::new()),
+            Default::default(),
+            StreamableHttpServerConfig::default()
+                .with_sse_keep_alive(None)
+                .with_cancellation_token(ct.child_token()),
+        );
+
+    let router = axum::Router::new().nest_service("/mcp", service);
+    let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
+    let addr = tcp_listener.local_addr()?;
+
+    let handle = tokio::spawn({
+        let ct = ct.clone();
+        async move {
+            let _ = axum::serve(tcp_listener, router)
+                .with_graceful_shutdown(async move { ct.cancelled_owned().await })
+                .await;
+        }
+    });
+
+    let client = reqwest::Client::new();
+
+    // Initialize the session
+    let response = client
+        .post(format!("http://{addr}/mcp"))
+        .header("Content-Type", "application/json")
+        .header("Accept", "application/json, text/event-stream")
+        .body(r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-11-25","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}"#)
+        .send()
+        .await?;
+    assert_eq!(response.status(), 200);
+    let session_id: SessionId = response.headers()["mcp-session-id"].to_str()?.into();
+
+    // Send notifications/initialized
+    let status = client
+        .post(format!("http://{addr}/mcp"))
+        .header("Content-Type", "application/json")
+        .header("Accept", "application/json, text/event-stream")
+        .header("mcp-session-id", session_id.to_string())
+        .header("Mcp-Protocol-Version", "2025-06-18")
+        .body(r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#)
+        .send()
+        .await?
+        .status();
+    assert_eq!(status, 202);
+
+    // First tool call — should get http_request_id 0
+    let body = client
+        .post(format!("http://{addr}/mcp"))
+        .header("Content-Type", "application/json")
+        .header("Accept", "application/json, text/event-stream")
+        .header("mcp-session-id", session_id.to_string())
+        .header("Mcp-Protocol-Version", "2025-06-18")
+        .body(r#"{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"sum","arguments":{"a":1,"b":2}}}"#)
+        .send()
+        .await?
+        .text()
+        .await?;
+
+    let events: Vec<&str> = body.split("\n\n").filter(|e| !e.is_empty()).collect();
+    assert!(
+        events.len() >= 2,
+        "expected priming + response, got: {body}"
+    );
+
+    // Priming event should encode the http_request_id (0)
+    let priming = events[0];
+    assert!(
+        priming.contains("id: 0/0"),
+        "first request priming should be 0/0, got: {priming}"
+    );
+    assert!(priming.contains("retry: 3000"));
+
+    // Response event should use index 1 (since priming occupies index 0)
+    let response_event = events[1];
+    assert!(
+        response_event.contains("id: 1/0"),
+        "first response event id should be 1/0, got: {response_event}"
+    );
+    assert!(response_event.contains(r#""id":2"#));
+
+    // Second tool call — should get http_request_id 1
+    let body = client
+        .post(format!("http://{addr}/mcp"))
+        .header("Content-Type", "application/json")
+        .header("Accept", "application/json, text/event-stream")
+        .header("mcp-session-id", session_id.to_string())
+        .header("Mcp-Protocol-Version", "2025-06-18")
+        .body(r#"{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"sum","arguments":{"a":3,"b":4}}}"#)
+        .send()
+        .await?
+        .text()
+        .await?;
+
+    let events: Vec<&str> = body.split("\n\n").filter(|e| !e.is_empty()).collect();
+    assert!(
+        events.len() >= 2,
+        "expected priming + response, got: {body}"
+    );
+
+    let priming = events[0];
+    assert!(
+        priming.contains("id: 0/1"),
+        "second request priming should be 0/1, got: {priming}"
+    );
+
+    let response_event = events[1];
+    assert!(
+        response_event.contains("id: 1/1"),
+        "second response event id should be 1/1, got: {response_event}"
+    );
+    assert!(response_event.contains(r#""id":3"#));
+
+    ct.cancel();
+    handle.await?;
+
+    Ok(())
+}
+
 #[tokio::test]
 async fn test_priming_on_stream_close() -> anyhow::Result<()> {
     use std::sync::Arc;