fix(transport): bound buffered Streamable HTTP responses

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>
@@ -213,11 +215,44 @@ 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;
+    use std::sync::{Arc, Mutex};
 
     use oauth2::{AccessToken, RefreshToken, basic::BasicTokenType};
 
@@ -228,9 +263,171 @@ mod tests {
             InMemoryCredentialStore, OAuthHttpClient, OAuthHttpClientFuture, OAuthHttpRequest,
             OAuthTokenResponse, StoredCredentials, VendorExtraTokenFields,
         },
-        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 {
+                    max_size: 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 { max_size: 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]

crates/rmcp/src/transport/common/client_side_http.rs

@@ -5,9 +5,6 @@ use futures::{Stream, TryStreamExt};
 
 use crate::transport::streamable_http_client::StreamableHttpError;
 
-pub(super) const MAX_JSON_RESPONSE_SIZE: usize = 16 * 1024 * 1024;
-pub(super) const MAX_ERROR_RESPONSE_SIZE: usize = 64 * 1024;
-
 /// Count bytes before retaining them; a length hint can reject a body early,
 /// but never replaces the streamed check or determines allocation size.
 pub(super) async fn read_bounded_body<S, E>(
@@ -38,9 +35,45 @@ where
     Ok(body)
 }
 
+/// Read at most `max_size` bytes for a diagnostic, dropping the rest of the stream.
+///
+/// Unlike a protocol message, a diagnostic need not be complete. Stop as soon as
+/// the prefix is full, without polling for another chunk or EOF. `Content-Length`
+/// is not consulted: an oversized declaration must not discard a useful prefix.
+/// The returned bytes must never be parsed as a complete JSON-RPC message.
+pub(super) async fn read_truncated_body<S, E>(
+    stream: S,
+    max_size: usize,
+) -> Result<Vec<u8>, StreamableHttpError<E>>
+where
+    S: Stream<Item = Result<Bytes, E>>,
+    E: std::error::Error + Send + Sync + 'static,
+{
+    let mut body = Vec::new();
+    futures::pin_mut!(stream);
+    while body.len() < max_size {
+        let Some(chunk) = stream
+            .try_next()
+            .await
+            .map_err(StreamableHttpError::Client)?
+        else {
+            break;
+        };
+        let count = chunk.len().min(max_size - body.len());
+        body.extend_from_slice(&chunk[..count]);
+    }
+    Ok(body)
+}
+
 #[cfg(test)]
 mod tests {
-    use std::io;
+    use std::{
+        io,
+        sync::{
+            Arc,
+            atomic::{AtomicUsize, Ordering},
+        },
+    };
 
     use futures::stream;
     use rstest::rstest;
@@ -119,5 +152,72 @@ mod tests {
                 .unwrap()
                 .is_empty()
         );
+        let stream = stream::iter([Ok::<_, io::Error>(Bytes::from_static(b"x"))]);
+        assert!(matches!(
+            read_bounded_body(stream, None, 0).await,
+            Err(StreamableHttpError::ResponseBodyTooLarge { max_size: 0 })
+        ));
+    }
+
+    #[tokio::test]
+    async fn truncated_body_keeps_short_bodies_complete() {
+        let stream = stream::iter([
+            Ok::<_, io::Error>(Bytes::from_static(b"ab")),
+            Ok(Bytes::from_static(b"c")),
+        ]);
+        assert_eq!(read_truncated_body(stream, 4).await.unwrap(), b"abc");
+    }
+
+    #[tokio::test]
+    async fn truncated_body_stops_at_exact_limit_without_polling_for_eof() {
+        let stream = futures::StreamExt::chain(
+            stream::iter([
+                Ok::<_, io::Error>(Bytes::from_static(b"ab")),
+                Ok(Bytes::from_static(b"cd")),
+            ]),
+            stream::poll_fn(|_| panic!("a full diagnostic prefix must not poll for EOF")),
+        );
+        assert_eq!(read_truncated_body(stream, 4).await.unwrap(), b"abcd");
+    }
+
+    #[tokio::test]
+    async fn truncated_body_stops_in_the_middle_of_an_oversized_chunk() {
+        let polls = Arc::new(AtomicUsize::new(0));
+        let count = polls.clone();
+        let stream = stream::poll_fn(move |_| {
+            assert!(
+                count.fetch_add(1, Ordering::SeqCst) < 2,
+                "must not drain the body"
+            );
+            std::task::Poll::Ready(Some(Ok::<_, io::Error>(Bytes::from_static(b"abc"))))
+        });
+        assert_eq!(read_truncated_body(stream, 4).await.unwrap(), b"abca");
+        assert_eq!(polls.load(Ordering::SeqCst), 2);
+    }
+
+    #[tokio::test]
+    async fn truncated_body_zero_limit_does_not_poll() {
+        let stream = stream::poll_fn(|_| -> std::task::Poll<Option<Result<Bytes, io::Error>>> {
+            panic!("zero diagnostic limit must not read the body")
+        });
+        assert!(read_truncated_body(stream, 0).await.unwrap().is_empty());
+    }
+
+    #[tokio::test]
+    async fn truncated_body_preserves_io_errors_before_the_limit() {
+        let stream = stream::iter([
+            Ok(Bytes::from_static(b"a")),
+            Err(io::Error::other("read failed")),
+        ]);
+        assert!(matches!(
+            read_truncated_body(stream, 4).await,
+            Err(StreamableHttpError::Client(_))
+        ));
+    }
+
+    #[tokio::test]
+    async fn truncated_body_counts_bytes_not_utf8_characters() {
+        let stream = stream::iter([Ok::<_, io::Error>(Bytes::from_static("é".as_bytes()))]);
+        assert_eq!(read_truncated_body(stream, 1).await.unwrap(), [0xc3]);
     }
 }