fix: allow SSE server router to be nested (#240)

* fix: allow SSE server router to be nested

Author: swallez

crates/rmcp/Cargo.toml

@@ -151,6 +151,7 @@ path = "tests/test_tool_macros.rs"
 [[test]]
 name = "test_with_python"
 required-features = [
+    "reqwest",
     "server",
     "client",
     "transport-sse-server",

crates/rmcp/src/transport/sse_server.rs

@@ -1,8 +1,8 @@
 use std::{collections::HashMap, io, net::SocketAddr, sync::Arc, time::Duration};
 
 use axum::{
-    Json, Router,
-    extract::{Query, State},
+    Extension, Json, Router,
+    extract::{NestedPath, Query, State},
     http::{StatusCode, request::Parts},
     response::{
         Response,
@@ -84,6 +84,7 @@ async fn post_event_handler(
 
 async fn sse_handler(
     State(app): State<App>,
+    nested_path: Option<Extension<NestedPath>>,
     parts: Parts,
 ) -> Result<Sse<impl Stream<Item = Result<Event, io::Error>>>, Response<String>> {
     let session = session_id();
@@ -115,12 +116,13 @@ async fn sse_handler(
         *response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
         return Err(response);
     }
+    let nested_path = nested_path.as_deref().map(NestedPath::as_str).unwrap_or("");
     let post_path = app.post_path.as_ref();
     let ping_interval = app.sse_ping_interval;
     let stream = futures::stream::once(futures::future::ok(
         Event::default()
             .event("endpoint")
-            .data(format!("{post_path}?sessionId={session}")),
+            .data(format!("{nested_path}{post_path}?sessionId={session}")),
     ))
     .chain(ReceiverStream::new(to_client_rx).map(|message| {
         match serde_json::to_string(&message) {
@@ -257,8 +259,6 @@ impl SseServer {
         Ok(sse_server)
     }
 
-    /// Warning: This function creates a new SseServer instance with the provided configuration.
-    /// `App.post_path` may be incorrect if using `Router` as an embedded router.
     pub fn new(config: SseServerConfig) -> (SseServer, Router) {
         let (app, transport_rx) = App::new(
             config.post_path.clone(),

crates/rmcp/tests/test_with_python.rs

@@ -1,15 +1,15 @@
+use axum::Router;
 use rmcp::{
     ServiceExt,
-    transport::{ConfigureCommandExt, SseServer, TokioChildProcess},
+    transport::{ConfigureCommandExt, SseServer, TokioChildProcess, sse_server::SseServerConfig},
 };
+use tokio::time::timeout;
+use tokio_util::sync::CancellationToken;
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
 mod common;
 use common::calculator::Calculator;
 
-const BIND_ADDRESS: &str = "127.0.0.1:8000";
-
-#[tokio::test]
-async fn test_with_python_client() -> anyhow::Result<()> {
+async fn init() -> anyhow::Result<()> {
     let _ = tracing_subscriber::registry()
         .with(
             tracing_subscriber::EnvFilter::try_from_default_env()
@@ -23,6 +23,14 @@ async fn test_with_python_client() -> anyhow::Result<()> {
         .spawn()?
         .wait()
         .await?;
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_with_python_client() -> anyhow::Result<()> {
+    init().await?;
+
+    const BIND_ADDRESS: &str = "127.0.0.1:8000";
 
     let ct = SseServer::serve(BIND_ADDRESS.parse()?)
         .await?
@@ -31,6 +39,7 @@ async fn test_with_python_client() -> anyhow::Result<()> {
     let status = tokio::process::Command::new("uv")
         .arg("run")
         .arg("tests/test_with_python/client.py")
+        .arg(format!("http://{BIND_ADDRESS}/sse"))
         .spawn()?
         .wait()
         .await?;
@@ -39,21 +48,61 @@ async fn test_with_python_client() -> anyhow::Result<()> {
     Ok(())
 }
 
+/// Test the SSE server in a nested Axum router.
+#[tokio::test]
+async fn test_nested_with_python_client() -> anyhow::Result<()> {
+    init().await?;
+
+    const BIND_ADDRESS: &str = "127.0.0.1:8001";
+
+    // Create an SSE router
+    let sse_config = SseServerConfig {
+        bind: BIND_ADDRESS.parse()?,
+        sse_path: "/sse".to_string(),
+        post_path: "/message".to_string(),
+        ct: CancellationToken::new(),
+        sse_keep_alive: None,
+    };
+
+    let listener = tokio::net::TcpListener::bind(&sse_config.bind).await?;
+
+    let (sse_server, sse_router) = SseServer::new(sse_config);
+    let ct = sse_server.with_service(Calculator::default);
+
+    let main_router = Router::new().nest("/nested", sse_router);
+
+    let server_ct = ct.clone();
+    let server = axum::serve(listener, main_router).with_graceful_shutdown(async move {
+        server_ct.cancelled().await;
+        tracing::info!("sse server cancelled");
+    });
+
+    tokio::spawn(async move {
+        let _ = server.await;
+        tracing::info!("sse server shutting down");
+    });
+
+    // Spawn the process with timeout, as failure to access the '/message' URL
+    // causes the client to never exit.
+    let status = timeout(
+        tokio::time::Duration::from_secs(5),
+        tokio::process::Command::new("uv")
+            .arg("run")
+            .arg("tests/test_with_python/client.py")
+            .arg(format!("http://{BIND_ADDRESS}/nested/sse"))
+            .spawn()?
+            .wait(),
+    )
+    .await?;
+    assert!(status?.success());
+    ct.cancel();
+    Ok(())
+}
+
 #[tokio::test]
 async fn test_with_python_server() -> anyhow::Result<()> {
-    let _ = tracing_subscriber::registry()
-        .with(
-            tracing_subscriber::EnvFilter::try_from_default_env()
-                .unwrap_or_else(|_| "debug".to_string().into()),
-        )
-        .with(tracing_subscriber::fmt::layer())
-        .try_init();
-    tokio::process::Command::new("uv")
-        .args(["pip", "install", "-r", "pyproject.toml"])
-        .current_dir("tests/test_with_python")
-        .spawn()?
-        .wait()
-        .await?;
+    init().await?;
+
     let transport = TokioChildProcess::new(tokio::process::Command::new("uv").configure(|cmd| {
         cmd.arg("run").arg("tests/test_with_python/server.py");
     }))?;

crates/rmcp/tests/test_with_python/client.py

@@ -1,10 +1,10 @@
 from mcp import ClientSession, StdioServerParameters, types
 from mcp.client.sse import sse_client
-
-
+import sys
 
 async def run():
-    async with sse_client("http://localhost:8000/sse") as (read, write):
+    url = sys.argv[1]
+    async with sse_client(url) as (read, write):
         async with ClientSession(
             read, write
         ) as session:
@@ -25,4 +25,4 @@ async def run():
 if __name__ == "__main__":
     import asyncio
 
-    asyncio.run(run())
+    asyncio.run(run())