add --override-protocol-version

Author: SteffenDE

src/cli.rs

@@ -18,4 +18,8 @@ pub struct Args {
     /// Initial retry interal 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/core.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/main.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/state.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
+        }
+    }
 }