fix(transport): bound buffered Streamable HTTP responses

crates/rmcp/src/transport/common.rs

@@ -15,6 +15,12 @@ mod reqwest;
 #[cfg(feature = "client-side-sse")]
 pub mod client_side_sse;
 
+#[cfg(any(
+    feature = "transport-streamable-http-client-reqwest",
+    all(unix, feature = "transport-streamable-http-client-unix-socket")
+))]
+pub(crate) mod client_side_body;
+
 #[cfg(feature = "auth")]
 pub mod auth;
 

crates/rmcp/src/transport/common/auth/streamable_http_client.rs

@@ -5,7 +5,9 @@ use tracing::debug;
 
 use crate::transport::{
     auth::{AuthClient, AuthError},
-    streamable_http_client::{StreamableHttpClient, StreamableHttpError},
+    streamable_http_client::{
+        StreamableHttpClient, StreamableHttpError, StreamableHttpResponseLimits,
+    },
 };
 
 impl<C> AuthClient<C>
@@ -208,19 +210,217 @@ where
         })
         .await
     }
+
+    async fn post_message_with_response_limits(
+        &self,
+        uri: std::sync::Arc<str>,
+        message: crate::model::ClientJsonRpcMessage,
+        session_id: Option<std::sync::Arc<str>>,
+        auth_token: Option<String>,
+        custom_headers: HashMap<HeaderName, HeaderValue>,
+        limits: StreamableHttpResponseLimits,
+    ) -> Result<
+        crate::transport::streamable_http_client::StreamableHttpPostResponse,
+        StreamableHttpError<Self::Error>,
+    > {
+        self.call_reacting_to_challenges(auth_token, |token| {
+            let uri = uri.clone();
+            let message = message.clone();
+            let session_id = session_id.clone();
+            let custom_headers = custom_headers.clone();
+            async move {
+                self.http_client
+                    .post_message_with_response_limits(
+                        uri,
+                        message,
+                        session_id,
+                        token,
+                        custom_headers,
+                        limits,
+                    )
+                    .await
+            }
+        })
+        .await
+    }
 }
 
 #[cfg(all(test, feature = "transport-streamable-http-client-reqwest"))]
 mod tests {
+    use std::sync::{Arc, Mutex};
+
     use super::*;
     use crate::transport::{
         auth::{
             AuthorizationManager, AuthorizationMetadata, CredentialRefreshGuard, CredentialStore,
+            InMemoryCredentialStore, OAuthHttpClient, OAuthHttpClientFuture, OAuthHttpRequest,
             StoredCredentials,
         },
-        streamable_http_client::AuthRequiredError,
+        streamable_http_client::{AuthRequiredError, StreamableHttpPostResponse},
     };
 
+    struct RefreshClient;
+
+    impl OAuthHttpClient for RefreshClient {
+        fn execute(&self, request: OAuthHttpRequest) -> OAuthHttpClientFuture<'_> {
+            Box::pin(async move {
+                assert_eq!(request.request.uri(), "https://auth.example.com/token");
+                assert_eq!(request.request.method(), http::Method::POST);
+                Ok(oauth2::http::Response::builder()
+                    .status(200)
+                    .header("content-type", "application/json")
+                    .body(br#"{"access_token":"fresh-token","token_type":"Bearer"}"#.to_vec())
+                    .unwrap())
+            })
+        }
+    }
+
+    type RecordedPostCall = (Option<String>, StreamableHttpResponseLimits);
+
+    #[derive(Clone, Default)]
+    struct RecordingLimitsClient {
+        calls: Arc<Mutex<Vec<RecordedPostCall>>>,
+    }
+
+    impl StreamableHttpClient for RecordingLimitsClient {
+        type Error = std::io::Error;
+
+        async fn post_message(
+            &self,
+            _: Arc<str>,
+            _: crate::model::ClientJsonRpcMessage,
+            _: Option<Arc<str>>,
+            _: Option<String>,
+            _: HashMap<HeaderName, HeaderValue>,
+        ) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
+            panic!("response limits must not use the legacy POST method");
+        }
+
+        async fn delete_session(
+            &self,
+            _: Arc<str>,
+            _: Arc<str>,
+            _: Option<String>,
+            _: HashMap<HeaderName, HeaderValue>,
+        ) -> Result<(), StreamableHttpError<Self::Error>> {
+            unreachable!("this test only sends POST requests")
+        }
+
+        async fn get_stream(
+            &self,
+            _: Arc<str>,
+            _: Option<Arc<str>>,
+            _: Option<String>,
+            _: Option<String>,
+            _: HashMap<HeaderName, HeaderValue>,
+        ) -> Result<
+            futures::stream::BoxStream<'static, Result<sse_stream::Sse, sse_stream::Error>>,
+            StreamableHttpError<Self::Error>,
+        > {
+            unreachable!("this test only sends POST requests")
+        }
+
+        async fn post_message_with_response_limits(
+            &self,
+            uri: Arc<str>,
+            message: crate::model::ClientJsonRpcMessage,
+            session_id: Option<Arc<str>>,
+            auth_token: Option<String>,
+            custom_headers: HashMap<HeaderName, HeaderValue>,
+            limits: StreamableHttpResponseLimits,
+        ) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
+            assert_eq!(uri.as_ref(), "https://mcp.example.com/mcp");
+            assert_eq!(session_id.as_deref(), Some("test-session"));
+            assert_eq!(serde_json::to_value(message).unwrap()["method"], "ping");
+            assert_eq!(
+                custom_headers.get(&HeaderName::from_static("x-test-header")),
+                Some(&HeaderValue::from_static("preserved")),
+            );
+            let mut calls = self.calls.lock().unwrap();
+            calls.push((auth_token, limits));
+            if calls.len() == 1 {
+                Err(StreamableHttpError::AuthRequired(AuthRequiredError::new(
+                    "Bearer".into(),
+                )))
+            } else {
+                Err(StreamableHttpError::ResponseBodyTooLarge {
+                    limit: limits.max_json_response_size,
+                })
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn response_limits_survive_auth_refresh_and_size_errors_do_not_retry() {
+        let mut manager = AuthorizationManager::new_with_oauth_http_client(
+            "https://mcp.example.com/mcp",
+            Arc::new(RefreshClient),
+        )
+        .await
+        .unwrap();
+        manager.set_metadata(AuthorizationMetadata {
+            authorization_endpoint: "https://auth.example.com/authorize".into(),
+            token_endpoint: "https://auth.example.com/token".into(),
+            ..Default::default()
+        });
+        manager.configure_client_id("client").unwrap();
+        let store = InMemoryCredentialStore::new();
+        let credentials = serde_json::from_value(serde_json::json!({
+            "access_token": "old-token",
+            "refresh_token": "refresh-token",
+            "token_type": "Bearer",
+        }))
+        .unwrap();
+        store
+            .save(StoredCredentials::new(
+                "client".into(),
+                Some(credentials),
+                vec![],
+                None,
+            ))
+            .await
+            .unwrap();
+        manager.set_credential_store(store);
+        let recording = RecordingLimitsClient::default();
+        let client = AuthClient::new(recording.clone(), manager);
+        let mut headers = HashMap::new();
+        headers.insert(
+            HeaderName::from_static("x-test-header"),
+            HeaderValue::from_static("preserved"),
+        );
+        let limits = StreamableHttpResponseLimits {
+            max_sse_event_size: 17,
+            max_json_response_size: 31,
+            max_error_response_size: 13,
+        };
+        let result = client
+            .post_message_with_response_limits(
+                "https://mcp.example.com/mcp".into(),
+                serde_json::from_value(
+                    serde_json::json!({"jsonrpc":"2.0", "id":1, "method":"ping"}),
+                )
+                .unwrap(),
+                Some("test-session".into()),
+                None,
+                headers,
+                limits,
+            )
+            .await;
+        assert!(matches!(
+            result,
+            Err(StreamableHttpError::ResponseBodyTooLarge { limit: 31 })
+        ));
+        let calls = recording.calls.lock().unwrap();
+        assert_eq!(calls.len(), 2);
+        assert_eq!(calls[0].0.as_deref(), Some("old-token"));
+        assert_eq!(calls[1].0.as_deref(), Some("fresh-token"));
+        for (_, observed) in calls.iter() {
+            assert_eq!(observed.max_sse_event_size, 17);
+            assert_eq!(observed.max_json_response_size, 31);
+            assert_eq!(observed.max_error_response_size, 13);
+        }
+    }
+
     struct UnavailableStore;
 
     #[async_trait::async_trait]