Merge pull request #7 from tidewave-ai/sd-override-protocol-version

Add --override-protocol-version

Author: SteffenDE

Cargo.lock

@@ -4,7 +4,7 @@ version = "0.2.0"
 edition = "2024"
 
 [dependencies]
-rmcp = { git = "https://github.com/modelcontextprotocol/rust-sdk.git", rev = "076dc2c2cd8910bee56bae13f29bbcff8c279666", features = [
+rmcp = { git = "https://github.com/modelcontextprotocol/rust-sdk.git", rev = "4c34b64b7f8dcabf94d52a9c6518c6b49c1f0451", features = [
     "server",
     "client",
     "reqwest",
@@ -29,7 +29,7 @@ version = "0.9"
 features = ["vendored"]
 
 [dev-dependencies]
-rmcp = { git = "https://github.com/modelcontextprotocol/rust-sdk.git", rev = "076dc2c2cd8910bee56bae13f29bbcff8c279666", features = [
+rmcp = { git = "https://github.com/modelcontextprotocol/rust-sdk.git", rev = "4c34b64b7f8dcabf94d52a9c6518c6b49c1f0451", features = [
     "server",
     "client",
     "reqwest",

Cargo.toml

@@ -1,6 +1,8 @@
 use anyhow::Context;
 use clap::Parser;
-use rmcp::transport::StreamableHttpServer;
+use rmcp::transport::streamable_http_server::{
+    StreamableHttpService, session::local::LocalSessionManager,
+};
 use tracing_subscriber::FmtSubscriber;
 
 use rmcp::{
@@ -12,6 +14,10 @@ use rmcp::{
 pub struct Echo;
 #[tool(tool_box)]
 impl Echo {
+    pub fn new() -> Self {
+        Self {}
+    }
+
     #[tool(description = "Echo a message")]
     fn echo(&self, #[tool(param)] message: String) -> String {
         message
@@ -49,11 +55,17 @@ async fn main() -> anyhow::Result<()> {
 
     tracing::subscriber::set_global_default(subscriber).context("Failed to set up logging")?;
 
-    let ct = StreamableHttpServer::serve(args.address)
-        .await?
-        .with_service(Echo::default);
+    let service = StreamableHttpService::new(
+        || Ok(Echo::new()),
+        LocalSessionManager::default().into(),
+        Default::default(),
+    );
+
+    let router = axum::Router::new().nest_service("/mcp", service);
+    let tcp_listener = tokio::net::TcpListener::bind(args.address).await?;
+    let _ = axum::serve(tcp_listener, router)
+        .with_graceful_shutdown(async { tokio::signal::ctrl_c().await.unwrap() })
+        .await;
 
-    tokio::signal::ctrl_c().await?;
-    ct.cancel();
     Ok(())
 }

examples/echo_streamable.rs

@@ -18,4 +18,8 @@ pub struct Args {
     /// Initial retry interval in seconds. Default is 5 seconds
     #[arg(long, default_value = "5")]
     pub initial_retry_interval: u64,
+
+    #[arg(long)]
+    /// Override the protocol version returned to the client
+    pub override_protocol_version: Option<String>,
 }

src/cli.rs

@@ -64,13 +64,9 @@ pub(crate) async fn connect_with_streamable(app_state: &AppState) -> Result<SseC
         rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig {
             uri: app_state.url.clone().into(),
             // we don't want the sdk to perform any retries
-            retry_config: std::sync::Arc::new(
-                rmcp::transport::common::client_side_sse::FixedInterval {
-                    max_times: Some(0),
-                    duration: Duration::from_millis(0),
-                },
-            ),
+            retry_config: std::sync::Arc::new(rmcp::transport::common::client_side_sse::NeverRetry),
             channel_buffer_capacity: 16,
+            allow_stateless: true,
         },
     );
 

src/core.rs

@@ -2,7 +2,7 @@ use anyhow::{Context, Result, anyhow};
 use clap::Parser;
 use futures::StreamExt;
 use rmcp::{
-    model::{ClientJsonRpcMessage, ErrorCode, ServerJsonRpcMessage},
+    model::{ClientJsonRpcMessage, ErrorCode, ProtocolVersion, ServerJsonRpcMessage},
     transport::{StreamableHttpClientTransport, Transport, sse_client::SseClientTransport},
 };
 use std::env;
@@ -128,12 +128,33 @@ async fn main() -> Result<()> {
     debug!("Starting MCP proxy with URL: {}", sse_url);
     debug!("Max disconnected time: {:?}s", args.max_disconnected_time);
 
+    // Parse protocol version override if provided
+    let override_protocol_version = if let Some(version_str) = args.override_protocol_version {
+        let protocol_version = match version_str.as_str() {
+            "2024-11-05" => ProtocolVersion::V_2024_11_05,
+            "2025-03-26" => ProtocolVersion::V_2025_03_26,
+            _ => {
+                return Err(anyhow!(
+                    "Unsupported protocol version: {}. Supported versions are: 2024-11-05, 2025-03-26",
+                    version_str
+                ));
+            }
+        };
+        Some(protocol_version)
+    } else {
+        None
+    };
+
     // Set up communication channels
     let (reconnect_tx, mut reconnect_rx) = tokio::sync::mpsc::channel(10);
     let (timer_tx, mut timer_rx) = tokio::sync::mpsc::channel(10);
 
     // Initialize application state
-    let mut app_state = AppState::new(sse_url.clone(), args.max_disconnected_time);
+    let mut app_state = AppState::new(
+        sse_url.clone(),
+        args.max_disconnected_time,
+        override_protocol_version,
+    );
     // Pass channel senders to state
     app_state.reconnect_tx = Some(reconnect_tx.clone());
     app_state.timer_tx = Some(timer_tx.clone());

src/main.rs

@@ -8,7 +8,7 @@ use anyhow::Result;
 use futures::SinkExt;
 use rmcp::model::{
     ClientJsonRpcMessage, ClientNotification, ClientRequest, EmptyResult, InitializedNotification,
-    InitializedNotificationMethod, RequestId, ServerJsonRpcMessage,
+    InitializedNotificationMethod, ProtocolVersion, RequestId, ServerJsonRpcMessage, ServerResult,
 };
 use std::collections::HashMap;
 use std::time::{Duration, Instant};
@@ -49,6 +49,8 @@ pub struct AppState {
     pub url: String,
     /// Maximum time to try reconnecting in seconds (None = infinity)
     pub max_disconnected_time: Option<u64>,
+    /// Override protocol version
+    pub override_protocol_version: Option<ProtocolVersion>,
     /// When we were disconnected
     pub disconnected_since: Option<Instant>,
     /// Current state of the application
@@ -78,10 +80,15 @@ pub struct AppState {
 }
 
 impl AppState {
-    pub fn new(url: String, max_disconnected_time: Option<u64>) -> Self {
+    pub fn new(
+        url: String,
+        max_disconnected_time: Option<u64>,
+        override_protocol_version: Option<ProtocolVersion>,
+    ) -> Self {
         Self {
             url,
             max_disconnected_time,
+            override_protocol_version,
             disconnected_since: None,
             state: ProxyState::Connecting,
             connect_tries: 0,
@@ -286,6 +293,7 @@ impl AppState {
                                 "Initial connection successful, received init response. Waiting for client initialized."
                             );
                             self.state = ProxyState::WaitingForClientInitialized;
+                            message = self.maybe_overwrite_protocol_version(message);
                         }
                     }
                     // --- End Initialization Response Handling ---
@@ -537,4 +545,25 @@ impl AppState {
         // Not a response/error, return Some(original_message)
         Some(message)
     }
+
+    fn maybe_overwrite_protocol_version(
+        &mut self,
+        message: ServerJsonRpcMessage,
+    ) -> ServerJsonRpcMessage {
+        if let Some(protocol_version) = &self.override_protocol_version {
+            match message {
+                ServerJsonRpcMessage::Response(mut resp) => {
+                    if let ServerResult::InitializeResult(mut initialize_result) = resp.result {
+                        initialize_result.protocol_version = protocol_version.clone();
+                        resp.result = ServerResult::InitializeResult(initialize_result);
+                        return ServerJsonRpcMessage::Response(resp);
+                    }
+                    ServerJsonRpcMessage::Response(resp)
+                }
+                other => other,
+            }
+        } else {
+            message
+        }
+    }
 }

src/state.rs

@@ -110,7 +110,7 @@ async fn create_sse_server(
     address: SocketAddr,
 ) -> Result<(tokio::process::Child, String)> {
     let url = if server_name == "echo_streamable" {
-        format!("http://{}", address)
+        format!("http://{}/mcp", address)
     } else {
         format!("http://{}/sse", address)
     };
@@ -348,7 +348,7 @@ async fn initial_connection_retry(server_name: &str) -> Result<()> {
 
     const BIND_ADDRESS: &str = "127.0.0.1:8184";
     let server_url = if server_name == "echo_streamable" {
-        format!("http://{}", BIND_ADDRESS)
+        format!("http://{}/mcp", BIND_ADDRESS)
     } else {
         format!("http://{}/sse", BIND_ADDRESS)
     };
@@ -545,3 +545,98 @@ async fn test_ping_when_disconnected() -> Result<()> {
 
     Ok(())
 }
+
+async fn protocol_version_override(server_name: &str) -> Result<()> {
+    const BIND_ADDRESS: &str = "127.0.0.1:8186";
+
+    // Phase 1: Test normal behavior (no override)
+    {
+        let (server_handle, server_url) =
+            create_sse_server(server_name, BIND_ADDRESS.parse()?).await?;
+        let (child, mut reader, stderr_reader, mut stdin) =
+            spawn_proxy(&server_url, vec![]).await?;
+        let stderr_buffer = collect_stderr(stderr_reader);
+        let _guard = TestGuard::new(child, server_handle, stderr_buffer);
+
+        // Send initialization message with 2025-03-26
+        let init_message = r#"{"jsonrpc":"2.0","id":"init-normal","method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"test","version":"0.1.0"}}}"#;
+        stdin.write_all(init_message.as_bytes()).await?;
+        stdin.write_all(b"\n").await?;
+
+        // Read the initialization response
+        let mut response = String::new();
+        timeout(Duration::from_secs(10), reader.read_line(&mut response)).await??;
+
+        // Verify the response contains the original protocol version
+        assert!(
+            response.contains("\"protocolVersion\":\"2025-03-26\""),
+            "Expected server to respond with 2025-03-26 protocol version, got: {}",
+            response
+        );
+
+        // Send initialized notification
+        let initialized_message = r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#;
+        stdin.write_all(initialized_message.as_bytes()).await?;
+        stdin.write_all(b"\n").await?;
+
+        // Clean shutdown
+        drop(stdin);
+    }
+
+    // Give a moment for cleanup
+    sleep(Duration::from_millis(500)).await;
+
+    // Phase 2: Test with protocol version override
+    {
+        let (server_handle, server_url) =
+            create_sse_server(server_name, BIND_ADDRESS.parse()?).await?;
+        let (child, mut reader, stderr_reader, mut stdin) = spawn_proxy(
+            &server_url,
+            vec!["--override-protocol-version", "2024-11-05"],
+        )
+        .await?;
+        let stderr_buffer = collect_stderr(stderr_reader);
+        let _guard = TestGuard::new(child, server_handle, stderr_buffer);
+
+        // Send initialization message with 2025-03-26 (same as phase 1)
+        let init_message = r#"{"jsonrpc":"2.0","id":"init-override","method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"test","version":"0.1.0"}}}"#;
+        stdin.write_all(init_message.as_bytes()).await?;
+        stdin.write_all(b"\n").await?;
+
+        // Read the initialization response
+        let mut response = String::new();
+        timeout(Duration::from_secs(10), reader.read_line(&mut response)).await??;
+
+        // Verify the response contains the overridden protocol version
+        assert!(
+            response.contains("\"protocolVersion\":\"2024-11-05\""),
+            "Expected proxy to override protocol version to 2024-11-05, got: {}",
+            response
+        );
+
+        // Verify it does NOT contain the original version
+        assert!(
+            !response.contains("\"protocolVersion\":\"2025-03-26\""),
+            "Protocol version should have been overridden from 2025-03-26 to 2024-11-05, got: {}",
+            response
+        );
+
+        // Send initialized notification
+        let initialized_message = r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#;
+        stdin.write_all(initialized_message.as_bytes()).await?;
+        stdin.write_all(b"\n").await?;
+
+        // Clean shutdown
+        drop(stdin);
+    }
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn test_protocol_version_override() -> Result<()> {
+    protocol_version_override("echo").await?;
+    protocol_version_override("echo_streamable").await?;
+
+    Ok(())
+}