fix: accept stringified numeric response IDs (#1021)

crates/rmcp/src/model.rs

@@ -240,6 +240,22 @@ impl NumberOrString {
             NumberOrString::String(s) => Value::String(s.to_string()),
         }
     }
+
+    pub(crate) fn numeric_string_value(&self) -> Option<i64> {
+        match self {
+            Self::String(id) => id.parse().ok(),
+            Self::Number(_) => None,
+        }
+    }
+
+    pub(crate) fn matches_response_id(&self, response_id: &Self) -> bool {
+        self == response_id
+            || matches!(
+                self,
+                Self::Number(request_id)
+                    if response_id.numeric_string_value() == Some(*request_id)
+            )
+    }
 }
 
 impl std::fmt::Display for NumberOrString {

crates/rmcp/src/service.rs

@@ -305,6 +305,17 @@ pub trait ProgressTokenProvider: Send + Sync + 'static {
 pub type AtomicU32RequestIdProvider = AtomicU32Provider;
 pub type AtomicU32ProgressTokenProvider = AtomicU32Provider;
 
+pub(crate) fn remove_pending_request<T>(
+    pending_requests: &mut HashMap<RequestId, T>,
+    response_id: &RequestId,
+) -> Option<T> {
+    pending_requests.remove(response_id).or_else(|| {
+        response_id
+            .numeric_string_value()
+            .and_then(|id| pending_requests.remove(&RequestId::Number(id)))
+    })
+}
+
 #[derive(Debug, Default)]
 pub struct AtomicU32Provider {
     id: AtomicU64,
@@ -1481,7 +1492,9 @@ where
                     id,
                     ..
                 })) => {
-                    if let Some(responder) = local_responder_pool.remove(&id) {
+                    if let Some(responder) =
+                        remove_pending_request(&mut local_responder_pool, &id)
+                    {
                         let response_result = responder.send(Ok(result));
                         if let Err(_error) = response_result {
                             tracing::warn!(%id, "Error sending response");
@@ -1495,7 +1508,9 @@ where
                         tracing::debug!(?error, "received id-less peer error");
                         continue;
                     };
-                    if let Some(responder) = local_responder_pool.remove(&id) {
+                    if let Some(responder) =
+                        remove_pending_request(&mut local_responder_pool, &id)
+                    {
                         let service_error = if error.is_transport_closed() {
                             ServiceError::TransportClosed
                         } else {

crates/rmcp/src/service/client.rs

@@ -687,7 +687,7 @@ where
     let (response, response_id) =
         expect_response(transport, "initialize response", service, peer.clone()).await?;
 
-    if id != response_id {
+    if !id.matches_response_id(&response_id) {
         return Err(ClientInitializeError::ConflictInitResponseId(
             id,
             response_id,
@@ -753,7 +753,7 @@ where
 
         match expect_response(transport, "discover response", service, peer.clone()).await {
             Ok((ServerResult::DiscoverResult(result), response_id)) => {
-                if response_id != id {
+                if !id.matches_response_id(&response_id) {
                     return Err(ClientInitializeError::ConflictInitResponseId(
                         id,
                         response_id,

crates/rmcp/src/transport/streamable_http_client.rs

@@ -499,8 +499,14 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
         pending_stream_response_ids: &mut HashSet<RequestId>,
         message: &ServerJsonRpcMessage,
     ) {
-        if let Some(id) = Self::server_response_id(message) {
-            pending_stream_response_ids.remove(id);
+        let Some(response_id) = Self::server_response_id(message) else {
+            return;
+        };
+        if pending_stream_response_ids.remove(response_id) {
+            return;
+        }
+        if let Some(id) = response_id.numeric_string_value() {
+            pending_stream_response_ids.remove(&RequestId::Number(id));
         }
     }
 
@@ -1382,7 +1388,10 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
                 }
                 Event::ServerMessage(mut json_rpc_message) => {
                     if let Some(response_id) = Self::server_response_id(&json_rpc_message)
-                        && let Some(stream_ct) = request_stream_cancellations.remove(response_id)
+                        && let Some(stream_ct) = crate::service::remove_pending_request(
+                            &mut request_stream_cancellations,
+                            response_id,
+                        )
                     {
                         stream_ct.cancel();
                     }
@@ -1848,4 +1857,39 @@ mod tests {
             vec!["legacy"]
         );
     }
+
+    #[cfg(feature = "transport-streamable-http-client-reqwest")]
+    #[test]
+    fn clear_stream_response_pending_accepts_stringified_numeric_id() {
+        let mut pending = HashSet::from([NumberOrString::Number(1)]);
+        let response = ServerJsonRpcMessage::response(
+            ServerResult::ListToolsResult(ListToolsResult::default()),
+            NumberOrString::String("1".into()),
+        );
+
+        StreamableHttpClientWorker::<reqwest::Client>::clear_stream_response_pending(
+            &mut pending,
+            &response,
+        );
+
+        assert!(pending.is_empty());
+    }
+
+    #[cfg(feature = "transport-streamable-http-client-reqwest")]
+    #[test]
+    fn clear_stream_response_pending_prefers_exact_string_id() {
+        let string_id = NumberOrString::String("1".into());
+        let mut pending = HashSet::from([NumberOrString::Number(1), string_id.clone()]);
+        let response = ServerJsonRpcMessage::response(
+            ServerResult::ListToolsResult(ListToolsResult::default()),
+            string_id,
+        );
+
+        StreamableHttpClientWorker::<reqwest::Client>::clear_stream_response_pending(
+            &mut pending,
+            &response,
+        );
+
+        assert_eq!(pending, HashSet::from([NumberOrString::Number(1)]));
+    }
 }