Update to Rust MCP SDK version 2 (#175)

Author: Matthew Hawkins

Cargo.lock

@@ -50,9 +50,3 @@ expect_used = "deny"
 indexing_slicing = "deny"
 unwrap_used = "deny"
 panic = "deny"
-
-[patch.crates-io]
-# The Rust MCP SDK has not had a release in a long time, so we need to depend on it through GitHub to pick up things
-# like Streamable HTTP support. This particular commit is pinned because it is before a change to the way shutdown is
-# handled through Axum which broke the ability to shut down the server with Ctrl-C.
-rmcp = { git = 'https://github.com/modelcontextprotocol/rust-sdk.git', rev = '915bc3fb37fe259fba1a853a9dd03566593f3310' }

Cargo.toml

@@ -10,13 +10,14 @@ anyhow = "1.0.98"
 apollo-compiler.workspace = true
 apollo-federation.workspace = true
 apollo-mcp-registry = { path = "../apollo-mcp-registry" }
+axum = "0.8.4"
 bon = "3.6.3"
 clap = { version = "4.5.36", features = ["derive", "env"] }
 futures.workspace = true
 lz-str = "0.2.1"
 regex = "1.11.1"
 reqwest.workspace = true
-rmcp = { version = "0.1", features = [
+rmcp = { version = "0.2", features = [
   "server",
   "transport-io",
   "transport-sse-server",

crates/apollo-mcp-server/Cargo.toml

@@ -1,7 +1,7 @@
 use apollo_compiler::{Schema, validation::Valid};
 use apollo_federation::{ApiSchemaOptions, Supergraph};
 use apollo_mcp_registry::uplink::schema::{SchemaState, event::Event as SchemaEvent};
-use futures::{FutureExt as _, Stream, StreamExt as _, future, stream};
+use futures::{FutureExt as _, Stream, StreamExt as _, stream};
 use reqwest::header::HeaderMap;
 
 use crate::{
@@ -144,34 +144,36 @@ impl StateMachine {
         }
     }
 
-    #[allow(clippy::expect_used)]
     fn ctrl_c_stream() -> impl Stream<Item = ServerEvent> {
-        #[cfg(not(unix))]
-        {
-            async {
-                tokio::signal::ctrl_c()
-                    .await
-                    .expect("Failed to install CTRL+C signal handler");
-            }
+        shutdown_signal()
             .map(|_| ServerEvent::Shutdown)
             .into_stream()
             .boxed()
-        }
+    }
+}
 
-        #[cfg(unix)]
-        future::select(
-            tokio::signal::ctrl_c().map(|s| s.ok()).boxed(),
-            async {
-                tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
-                    .expect("Failed to install SIGTERM signal handler")
-                    .recv()
-                    .await
-            }
-            .boxed(),
-        )
-        .map(|_| ServerEvent::Shutdown)
-        .into_stream()
-        .boxed()
+#[allow(clippy::expect_used)]
+async fn shutdown_signal() {
+    let ctrl_c = async {
+        tokio::signal::ctrl_c()
+            .await
+            .expect("Failed to install CTRL+C signal handler");
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
+            .expect("Failed to install SIGTERM signal handler")
+            .recv()
+            .await;
+    };
+
+    #[cfg(not(unix))]
+    let terminate = std::future::pending::<()>();
+
+    tokio::select! {
+        _ = ctrl_c => {},
+        _ = terminate => {},
     }
 }
 

crates/apollo-mcp-server/src/server/states.rs

@@ -1,12 +1,11 @@
 use std::{net::SocketAddr, sync::Arc};
 
 use apollo_compiler::{Name, Schema, ast::OperationType, validation::Valid};
+use rmcp::transport::StreamableHttpService;
+use rmcp::transport::streamable_http_server::session::local::LocalSessionManager;
 use rmcp::{
     ServiceExt as _,
-    transport::{
-        SseServer, StreamableHttpServer, sse_server::SseServerConfig, stdio,
-        streamable_http_server::axum::StreamableHttpServerConfig,
-    },
+    transport::{SseServer, sse_server::SseServerConfig, stdio},
 };
 use tokio::sync::{Mutex, RwLock};
 use tokio_util::sync::CancellationToken;
@@ -20,7 +19,7 @@ use crate::{
     server::Transport,
 };
 
-use super::{Config, Running};
+use super::{Config, Running, shutdown_signal};
 
 pub(super) struct Starting {
     pub(super) config: Config,
@@ -118,14 +117,16 @@ impl Starting {
                 info!(port = ?port, address = ?address, "Starting MCP server in Streamable HTTP mode");
                 let running = running.clone();
                 let listen_address = SocketAddr::new(address, port);
-                StreamableHttpServer::serve_with_config(StreamableHttpServerConfig {
-                    bind: listen_address,
-                    path: "/mcp".to_string(),
-                    ct: cancellation_token,
-                    sse_keep_alive: None,
-                })
-                .await?
-                .with_service(move || running.clone());
+                let service = StreamableHttpService::new(
+                    move || Ok(running.clone()),
+                    LocalSessionManager::default().into(),
+                    Default::default(),
+                );
+                let router = axum::Router::new().nest_service("/mcp", service);
+                let tcp_listener = tokio::net::TcpListener::bind(listen_address).await?;
+                axum::serve(tcp_listener, router)
+                    .with_graceful_shutdown(shutdown_signal())
+                    .await?;
             }
             Transport::SSE { address, port } => {
                 info!(port = ?port, address = ?address, "Starting MCP server in SSE mode");