Skip to content

Commit 24ba526

Browse files
authored
fix: bound streamable HTTP memory usage (#970)
* fix: bound streamable HTTP memory usage * fix: handle SSE comment lines safely * refactor: narrow SSE limit public API surface
1 parent 1543a1a commit 24ba526

8 files changed

Lines changed: 811 additions & 39 deletions

File tree

crates/rmcp/Cargo.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,7 @@ server-side-http = [
153153
transport-worker = ["dep:tokio-stream"]
154154

155155
# SSE stream parsing utilities (used by streamable HTTP client for SSE-formatted responses)
156-
client-side-sse = ["dep:sse-stream", "dep:http", "base64"]
156+
client-side-sse = ["dep:sse-stream", "dep:http", "dep:bytes", "base64"]
157157

158158
# Streamable HTTP client
159159
transport-streamable-http-client = ["client-side-sse", "transport-worker"]

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

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,33 @@ where
4747
.await
4848
}
4949

50+
async fn get_stream_with_max_sse_event_size(
51+
&self,
52+
uri: std::sync::Arc<str>,
53+
session_id: std::sync::Arc<str>,
54+
last_event_id: Option<String>,
55+
mut auth_token: Option<String>,
56+
custom_headers: HashMap<HeaderName, HeaderValue>,
57+
max_sse_event_size: usize,
58+
) -> Result<
59+
futures::stream::BoxStream<'static, Result<sse_stream::Sse, sse_stream::Error>>,
60+
crate::transport::streamable_http_client::StreamableHttpError<Self::Error>,
61+
> {
62+
if auth_token.is_none() {
63+
auth_token = Some(self.get_access_token().await?);
64+
}
65+
self.http_client
66+
.get_stream_with_max_sse_event_size(
67+
uri,
68+
session_id,
69+
last_event_id,
70+
auth_token,
71+
custom_headers,
72+
max_sse_event_size,
73+
)
74+
.await
75+
}
76+
5077
async fn post_message(
5178
&self,
5279
uri: std::sync::Arc<str>,
@@ -65,4 +92,31 @@ where
6592
.post_message(uri, message, session_id, auth_token, custom_headers)
6693
.await
6794
}
95+
96+
async fn post_message_with_max_sse_event_size(
97+
&self,
98+
uri: std::sync::Arc<str>,
99+
message: crate::model::ClientJsonRpcMessage,
100+
session_id: Option<std::sync::Arc<str>>,
101+
mut auth_token: Option<String>,
102+
custom_headers: HashMap<HeaderName, HeaderValue>,
103+
max_sse_event_size: usize,
104+
) -> Result<
105+
crate::transport::streamable_http_client::StreamableHttpPostResponse,
106+
StreamableHttpError<Self::Error>,
107+
> {
108+
if auth_token.is_none() {
109+
auth_token = Some(self.get_access_token().await?);
110+
}
111+
self.http_client
112+
.post_message_with_max_sse_event_size(
113+
uri,
114+
message,
115+
session_id,
116+
auth_token,
117+
custom_headers,
118+
max_sse_event_size,
119+
)
120+
.await
121+
}
68122
}

0 commit comments

Comments
 (0)