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
This commit is contained in:
Dale Seo
GitHub
parent
1543a1afee
commit
24ba5265e2
8 changed files with 811 additions and 39 deletions
|
|
@ -153,7 +153,7 @@ server-side-http = [
|
|||
transport-worker = ["dep:tokio-stream"]
|
||||
|
||||
# SSE stream parsing utilities (used by streamable HTTP client for SSE-formatted responses)
|
||||
client-side-sse = ["dep:sse-stream", "dep:http", "base64"]
|
||||
client-side-sse = ["dep:sse-stream", "dep:http", "dep:bytes", "base64"]
|
||||
|
||||
# Streamable HTTP client
|
||||
transport-streamable-http-client = ["client-side-sse", "transport-worker"]
|
||||
|
|
|
|||
|
|
@ -47,6 +47,33 @@ where
|
|||
.await
|
||||
}
|
||||
|
||||
async fn get_stream_with_max_sse_event_size(
|
||||
&self,
|
||||
uri: std::sync::Arc<str>,
|
||||
session_id: std::sync::Arc<str>,
|
||||
last_event_id: Option<String>,
|
||||
mut auth_token: Option<String>,
|
||||
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
max_sse_event_size: usize,
|
||||
) -> Result<
|
||||
futures::stream::BoxStream<'static, Result<sse_stream::Sse, sse_stream::Error>>,
|
||||
crate::transport::streamable_http_client::StreamableHttpError<Self::Error>,
|
||||
> {
|
||||
if auth_token.is_none() {
|
||||
auth_token = Some(self.get_access_token().await?);
|
||||
}
|
||||
self.http_client
|
||||
.get_stream_with_max_sse_event_size(
|
||||
uri,
|
||||
session_id,
|
||||
last_event_id,
|
||||
auth_token,
|
||||
custom_headers,
|
||||
max_sse_event_size,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn post_message(
|
||||
&self,
|
||||
uri: std::sync::Arc<str>,
|
||||
|
|
@ -65,4 +92,31 @@ where
|
|||
.post_message(uri, message, session_id, auth_token, custom_headers)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn post_message_with_max_sse_event_size(
|
||||
&self,
|
||||
uri: std::sync::Arc<str>,
|
||||
message: crate::model::ClientJsonRpcMessage,
|
||||
session_id: Option<std::sync::Arc<str>>,
|
||||
mut auth_token: Option<String>,
|
||||
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
max_sse_event_size: usize,
|
||||
) -> Result<
|
||||
crate::transport::streamable_http_client::StreamableHttpPostResponse,
|
||||
StreamableHttpError<Self::Error>,
|
||||
> {
|
||||
if auth_token.is_none() {
|
||||
auth_token = Some(self.get_access_token().await?);
|
||||
}
|
||||
self.http_client
|
||||
.post_message_with_max_sse_event_size(
|
||||
uri,
|
||||
message,
|
||||
session_id,
|
||||
auth_token,
|
||||
custom_headers,
|
||||
max_sse_event_size,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -5,13 +5,166 @@ use std::{
|
|||
time::Duration,
|
||||
};
|
||||
|
||||
use futures::{Stream, stream::BoxStream};
|
||||
use sse_stream::{Error as SseError, Sse};
|
||||
use bytes::Bytes;
|
||||
use futures::{Stream, StreamExt, stream::BoxStream};
|
||||
use sse_stream::{Error as SseError, Sse, SseStream};
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::model::ServerJsonRpcMessage;
|
||||
|
||||
pub type BoxedSseResponse = BoxStream<'static, Result<Sse, SseError>>;
|
||||
|
||||
/// Maximum raw size of one SSE event accepted from a remote server.
|
||||
pub(crate) const DEFAULT_MAX_SSE_EVENT_SIZE: usize = 16 * 1024 * 1024;
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
enum BoundedSseStreamError {
|
||||
#[error(transparent)]
|
||||
Source(Box<dyn std::error::Error + Send + Sync>),
|
||||
#[error("SSE event exceeded the maximum size of {max_size} bytes")]
|
||||
EventTooLarge { max_size: usize },
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct SseEventSizeLimiter {
|
||||
max_size: usize,
|
||||
retained_size: usize,
|
||||
line_size: usize,
|
||||
line_is_comment: bool,
|
||||
previous_was_cr: bool,
|
||||
}
|
||||
|
||||
impl SseEventSizeLimiter {
|
||||
fn new(max_size: usize) -> Self {
|
||||
Self {
|
||||
max_size,
|
||||
retained_size: 0,
|
||||
line_size: 0,
|
||||
line_is_comment: false,
|
||||
previous_was_cr: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn observe(&mut self, chunk: &[u8]) -> Result<(), ()> {
|
||||
for &byte in chunk {
|
||||
if self.previous_was_cr {
|
||||
self.previous_was_cr = false;
|
||||
if byte == b'\n' {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
match byte {
|
||||
b'\r' => {
|
||||
self.finish_line()?;
|
||||
self.previous_was_cr = true;
|
||||
}
|
||||
b'\n' => self.finish_line()?,
|
||||
_ => {
|
||||
if self.line_size == 0 {
|
||||
self.line_is_comment = byte == b':';
|
||||
}
|
||||
self.line_size = self.line_size.saturating_add(1);
|
||||
self.check_limit()?;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn finish_line(&mut self) -> Result<(), ()> {
|
||||
if self.line_size == 0 {
|
||||
self.retained_size = 0;
|
||||
} else if !self.line_is_comment {
|
||||
// The SSE parser inserts a newline when joining multiple data fields.
|
||||
self.retained_size = self
|
||||
.retained_size
|
||||
.saturating_add(self.line_size)
|
||||
.saturating_add(1);
|
||||
}
|
||||
self.line_size = 0;
|
||||
self.line_is_comment = false;
|
||||
self.check_limit()
|
||||
}
|
||||
|
||||
fn check_limit(&self) -> Result<(), ()> {
|
||||
if self.retained_size.saturating_add(self.line_size) > self.max_size {
|
||||
Err(())
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pin_project_lite::pin_project! {
|
||||
struct BoundedSseByteStream<S> {
|
||||
#[pin]
|
||||
inner: S,
|
||||
limiter: SseEventSizeLimiter,
|
||||
failed: bool,
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, E> Stream for BoundedSseByteStream<S>
|
||||
where
|
||||
S: Stream<Item = Result<Bytes, E>>,
|
||||
E: std::error::Error + Send + Sync + 'static,
|
||||
{
|
||||
type Item = Result<Bytes, BoundedSseStreamError>;
|
||||
|
||||
fn poll_next(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut std::task::Context<'_>,
|
||||
) -> Poll<Option<Self::Item>> {
|
||||
let mut this = self.project();
|
||||
if *this.failed {
|
||||
return Poll::Ready(None);
|
||||
}
|
||||
|
||||
match ready!(this.inner.as_mut().poll_next(cx)) {
|
||||
Some(Ok(chunk)) => {
|
||||
if this.limiter.observe(&chunk).is_err() {
|
||||
*this.failed = true;
|
||||
Poll::Ready(Some(Err(BoundedSseStreamError::EventTooLarge {
|
||||
max_size: this.limiter.max_size,
|
||||
})))
|
||||
} else {
|
||||
Poll::Ready(Some(Ok(chunk)))
|
||||
}
|
||||
}
|
||||
Some(Err(error)) => {
|
||||
*this.failed = true;
|
||||
Poll::Ready(Some(Err(BoundedSseStreamError::Source(Box::new(error)))))
|
||||
}
|
||||
None => Poll::Ready(None),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn bounded_sse_stream<S, E>(stream: S, max_event_size: usize) -> BoxedSseResponse
|
||||
where
|
||||
S: Stream<Item = Result<Bytes, E>> + Send + 'static,
|
||||
E: std::error::Error + Send + Sync + 'static,
|
||||
{
|
||||
let stream = BoundedSseByteStream {
|
||||
inner: stream,
|
||||
limiter: SseEventSizeLimiter::new(max_event_size),
|
||||
failed: false,
|
||||
};
|
||||
SseStream::from_bytes_stream(stream).boxed()
|
||||
}
|
||||
|
||||
fn is_event_too_large_error(error: &SseError) -> bool {
|
||||
matches!(
|
||||
error,
|
||||
SseError::Body(error)
|
||||
if matches!(
|
||||
error.downcast_ref::<BoundedSseStreamError>(),
|
||||
Some(BoundedSseStreamError::EventTooLarge { .. })
|
||||
)
|
||||
)
|
||||
}
|
||||
|
||||
pub trait SseRetryPolicy: std::fmt::Debug + Send + Sync {
|
||||
fn retry(&self, current_times: usize) -> Option<Duration>;
|
||||
}
|
||||
|
|
@ -124,6 +277,10 @@ pub(crate) trait SseStreamReconnect {
|
|||
tracing::warn!("sse stream error: {error}");
|
||||
}
|
||||
}
|
||||
fn map_fatal_stream_error(&mut self, error: SseError) -> Option<Self::Error> {
|
||||
tracing::warn!("fatal sse stream error: {error}");
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
pin_project_lite::pin_project! {
|
||||
|
|
@ -248,6 +405,10 @@ where
|
|||
}
|
||||
}
|
||||
Some(Err(e)) => {
|
||||
if is_event_too_large_error(&e) {
|
||||
this.state.set(SseAutoReconnectStreamState::Terminated);
|
||||
return Poll::Ready(this.connector.map_fatal_stream_error(e).map(Err));
|
||||
}
|
||||
this.connector
|
||||
.handle_stream_error(&e, this.last_event_id.as_deref());
|
||||
let retrying = this
|
||||
|
|
@ -327,3 +488,202 @@ where
|
|||
self.poll_next(cx)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
use super::*;
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
enum TestReconnectError {
|
||||
#[error("SSE stream error: {0}")]
|
||||
Sse(SseError),
|
||||
#[error("unexpected reconnect")]
|
||||
Reconnect,
|
||||
}
|
||||
|
||||
struct CountingReconnect {
|
||||
attempts: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
impl SseStreamReconnect for CountingReconnect {
|
||||
type Error = TestReconnectError;
|
||||
type Future = futures::future::Ready<Result<BoxedSseResponse, Self::Error>>;
|
||||
|
||||
fn retry_connection(&mut self, _last_event_id: Option<&str>) -> Self::Future {
|
||||
self.attempts.fetch_add(1, Ordering::Relaxed);
|
||||
futures::future::ready(Err(TestReconnectError::Reconnect))
|
||||
}
|
||||
|
||||
fn map_fatal_stream_error(&mut self, error: SseError) -> Option<Self::Error> {
|
||||
Some(TestReconnectError::Sse(error))
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bounded_sse_stream_rejects_unterminated_event_over_limit() {
|
||||
let source = futures::stream::iter([
|
||||
Ok::<_, std::io::Error>(Bytes::from_static(b"data: aaaaa")),
|
||||
Ok(Bytes::from_static(b"aaaaaa")),
|
||||
]);
|
||||
let mut stream = bounded_sse_stream(source, 16);
|
||||
|
||||
let error = stream.next().await.unwrap().unwrap_err();
|
||||
|
||||
assert!(
|
||||
error.to_string().contains("maximum size of 16 bytes"),
|
||||
"unexpected error: {error}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bounded_sse_stream_resets_limit_after_event_terminator() {
|
||||
let source = futures::stream::iter([
|
||||
Ok::<_, std::io::Error>(Bytes::from_static(b"data: a\r")),
|
||||
Ok(Bytes::from_static(b"\n\r\ndata: b\n\n")),
|
||||
]);
|
||||
let mut stream = bounded_sse_stream(source, 8);
|
||||
|
||||
let first = stream.next().await.unwrap().unwrap();
|
||||
let second = stream.next().await.unwrap().unwrap();
|
||||
|
||||
assert_eq!(
|
||||
(first.data.as_deref(), second.data.as_deref()),
|
||||
(Some("a"), Some("b"))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bounded_sse_stream_passes_event_at_exact_limit() {
|
||||
let source =
|
||||
futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static(b"data: ab\n\n"))]);
|
||||
let mut stream = bounded_sse_stream(source, 9);
|
||||
|
||||
let event = stream.next().await.unwrap().unwrap();
|
||||
assert_eq!(event.data.as_deref(), Some("ab"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bounded_sse_stream_rejects_oversize_split_across_many_chunks() {
|
||||
let source = futures::stream::iter(
|
||||
std::iter::repeat_with(|| Ok::<_, std::io::Error>(Bytes::from_static(b"data: x")))
|
||||
.take(20),
|
||||
);
|
||||
let mut stream = bounded_sse_stream(source, 32);
|
||||
|
||||
let mut found_error = false;
|
||||
while let Some(item) = stream.next().await {
|
||||
if item.is_err() {
|
||||
found_error = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
assert!(found_error, "expected oversize error");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bounded_sse_stream_handles_crlf_split_across_chunks() {
|
||||
let source = futures::stream::iter([
|
||||
Ok::<_, std::io::Error>(Bytes::from_static(b"data: hello\r")),
|
||||
Ok(Bytes::from_static(b"\n\ndata: world\n\n")),
|
||||
]);
|
||||
let mut stream = bounded_sse_stream(source, 64);
|
||||
|
||||
let first = stream.next().await.unwrap().unwrap();
|
||||
let second = stream.next().await.unwrap().unwrap();
|
||||
|
||||
assert_eq!(first.data.as_deref(), Some("hello"));
|
||||
assert_eq!(second.data.as_deref(), Some("world"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bounded_sse_stream_discards_completed_comment_lines_from_limit() {
|
||||
let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static(
|
||||
b": ping\n: ping\n: ping\n",
|
||||
))]);
|
||||
let mut stream = bounded_sse_stream(source, 6);
|
||||
|
||||
assert!(stream.next().await.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bounded_sse_stream_comments_do_not_reset_accumulated_data() {
|
||||
let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static(
|
||||
b"data: a\n: ping\ndata: b\n",
|
||||
))]);
|
||||
let mut stream = bounded_sse_stream(source, 14);
|
||||
|
||||
assert!(stream.next().await.unwrap().is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bounded_sse_stream_counts_multiline_data_join_newlines() {
|
||||
let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static(
|
||||
b"data: aaa\ndata: bbb\n\n",
|
||||
))]);
|
||||
let mut stream = bounded_sse_stream(source, 18);
|
||||
|
||||
let error = stream.next().await.unwrap().unwrap_err();
|
||||
assert!(
|
||||
error.to_string().contains("maximum size"),
|
||||
"unexpected error: {error}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bounded_sse_stream_propagates_source_error() {
|
||||
let source = futures::stream::iter([Err::<Bytes, _>(std::io::Error::new(
|
||||
std::io::ErrorKind::ConnectionReset,
|
||||
"connection reset",
|
||||
))]);
|
||||
let mut stream = bounded_sse_stream(source, 1024);
|
||||
|
||||
let error = stream.next().await.unwrap().unwrap_err();
|
||||
assert!(
|
||||
error.to_string().contains("connection reset"),
|
||||
"unexpected error: {error}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn is_event_too_large_error_detects_oversize() {
|
||||
let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from(vec![b'A'; 100]))]);
|
||||
let mut stream = bounded_sse_stream(source, 8);
|
||||
|
||||
let error = stream.next().await.unwrap().unwrap_err();
|
||||
assert!(is_event_too_large_error(&error));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn is_event_too_large_error_rejects_other_errors() {
|
||||
let source =
|
||||
futures::stream::iter([Err::<Bytes, _>(std::io::Error::other("something else"))]);
|
||||
let mut stream = bounded_sse_stream(source, 1024);
|
||||
|
||||
let error = stream.next().await.unwrap().unwrap_err();
|
||||
assert!(!is_event_too_large_error(&error));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn oversized_event_returns_error_without_reconnecting() {
|
||||
let source = futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from(vec![b'A'; 100]))]);
|
||||
let attempts = Arc::new(AtomicUsize::new(0));
|
||||
let connector = CountingReconnect {
|
||||
attempts: attempts.clone(),
|
||||
};
|
||||
let stream = SseAutoReconnectStream::new(
|
||||
bounded_sse_stream(source, 8),
|
||||
connector,
|
||||
Arc::new(NeverRetry),
|
||||
);
|
||||
let mut stream = std::pin::pin!(stream);
|
||||
|
||||
let result = stream.next().await;
|
||||
|
||||
assert!(
|
||||
matches!(result, Some(Err(TestReconnectError::Sse(_))))
|
||||
&& attempts.load(Ordering::Relaxed) == 0
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,16 +1,19 @@
|
|||
use std::{borrow::Cow, collections::HashMap, sync::Arc};
|
||||
|
||||
use futures::{StreamExt, stream::BoxStream};
|
||||
use futures::stream::BoxStream;
|
||||
use http::{HeaderName, HeaderValue, header::WWW_AUTHENTICATE};
|
||||
use reqwest::header::ACCEPT;
|
||||
use sse_stream::{Sse, SseStream};
|
||||
use sse_stream::Sse;
|
||||
|
||||
use crate::{
|
||||
model::{ClientJsonRpcMessage, JsonRpcMessage, ServerJsonRpcMessage},
|
||||
transport::{
|
||||
common::http_header::{
|
||||
EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_SESSION_ID, JSON_MIME_TYPE,
|
||||
extract_scope_from_header, validate_custom_header,
|
||||
common::{
|
||||
client_side_sse::{DEFAULT_MAX_SSE_EVENT_SIZE, bounded_sse_stream},
|
||||
http_header::{
|
||||
EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_SESSION_ID, JSON_MIME_TYPE,
|
||||
extract_scope_from_header, validate_custom_header,
|
||||
},
|
||||
},
|
||||
streamable_http_client::*,
|
||||
},
|
||||
|
|
@ -53,6 +56,26 @@ impl StreamableHttpClient for reqwest::Client {
|
|||
last_event_id: Option<String>,
|
||||
auth_token: Option<String>,
|
||||
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
) -> Result<BoxStream<'static, Result<Sse, SseError>>, StreamableHttpError<Self::Error>> {
|
||||
self.get_stream_with_max_sse_event_size(
|
||||
uri,
|
||||
session_id,
|
||||
last_event_id,
|
||||
auth_token,
|
||||
custom_headers,
|
||||
DEFAULT_MAX_SSE_EVENT_SIZE,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn get_stream_with_max_sse_event_size(
|
||||
&self,
|
||||
uri: Arc<str>,
|
||||
session_id: Arc<str>,
|
||||
last_event_id: Option<String>,
|
||||
auth_token: Option<String>,
|
||||
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
max_sse_event_size: usize,
|
||||
) -> Result<BoxStream<'static, Result<Sse, SseError>>, StreamableHttpError<Self::Error>> {
|
||||
let mut request_builder = self
|
||||
.get(uri.as_ref())
|
||||
|
|
@ -84,7 +107,7 @@ impl StreamableHttpClient for reqwest::Client {
|
|||
return Err(StreamableHttpError::UnexpectedContentType(None));
|
||||
}
|
||||
}
|
||||
let event_stream = SseStream::from_bytes_stream(response.bytes_stream()).boxed();
|
||||
let event_stream = bounded_sse_stream(response.bytes_stream(), max_sse_event_size);
|
||||
Ok(event_stream)
|
||||
}
|
||||
|
||||
|
|
@ -119,6 +142,26 @@ impl StreamableHttpClient for reqwest::Client {
|
|||
session_id: Option<Arc<str>>,
|
||||
auth_token: Option<String>,
|
||||
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
|
||||
self.post_message_with_max_sse_event_size(
|
||||
uri,
|
||||
message,
|
||||
session_id,
|
||||
auth_token,
|
||||
custom_headers,
|
||||
DEFAULT_MAX_SSE_EVENT_SIZE,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn post_message_with_max_sse_event_size(
|
||||
&self,
|
||||
uri: Arc<str>,
|
||||
message: ClientJsonRpcMessage,
|
||||
session_id: Option<Arc<str>>,
|
||||
auth_token: Option<String>,
|
||||
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
max_sse_event_size: usize,
|
||||
) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
|
||||
let mut request = self
|
||||
.post(uri.as_ref())
|
||||
|
|
@ -223,7 +266,7 @@ impl StreamableHttpClient for reqwest::Client {
|
|||
}
|
||||
match content_type.as_deref() {
|
||||
Some(ct) if ct.as_bytes().starts_with(EVENT_STREAM_MIME_TYPE.as_bytes()) => {
|
||||
let event_stream = SseStream::from_bytes_stream(response.bytes_stream()).boxed();
|
||||
let event_stream = bounded_sse_stream(response.bytes_stream(), max_sse_event_size);
|
||||
Ok(StreamableHttpPostResponse::Sse(event_stream, session_id))
|
||||
}
|
||||
Some(ct) if ct.as_bytes().starts_with(JSON_MIME_TYPE.as_bytes()) => {
|
||||
|
|
@ -364,6 +407,57 @@ mod tests {
|
|||
assert!(parse_json_rpc_error(body).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn post_sse_response_honors_configured_event_limit() -> anyhow::Result<()> {
|
||||
use std::{collections::HashMap, net::SocketAddr, sync::Arc};
|
||||
|
||||
use axum::{Router, routing::post};
|
||||
use futures::StreamExt;
|
||||
|
||||
use crate::transport::streamable_http_client::{
|
||||
StreamableHttpClient, StreamableHttpPostResponse,
|
||||
};
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(SocketAddr::from(([127, 0, 0, 1], 0))).await?;
|
||||
let addr = listener.local_addr()?;
|
||||
let server = tokio::spawn(async move {
|
||||
let app = Router::new().route(
|
||||
"/mcp",
|
||||
post(|| async {
|
||||
(
|
||||
[(http::header::CONTENT_TYPE, "text/event-stream")],
|
||||
"data: this event is too large\n",
|
||||
)
|
||||
}),
|
||||
);
|
||||
axum::serve(listener, app).await
|
||||
});
|
||||
let message = ClientJsonRpcMessage::request(
|
||||
ClientRequest::PingRequest(PingRequest::default()),
|
||||
RequestId::Number(1),
|
||||
);
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
let response = client
|
||||
.post_message_with_max_sse_event_size(
|
||||
Arc::<str>::from(format!("http://{addr}/mcp")),
|
||||
message,
|
||||
None,
|
||||
None,
|
||||
HashMap::new(),
|
||||
16,
|
||||
)
|
||||
.await?;
|
||||
let StreamableHttpPostResponse::Sse(mut stream, _) = response else {
|
||||
anyhow::bail!("expected SSE response");
|
||||
};
|
||||
let error = stream.next().await.unwrap().unwrap_err();
|
||||
|
||||
server.abort();
|
||||
assert!(error.to_string().contains("maximum size of 16 bytes"));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn default_http_client_does_not_leak_custom_headers_to_redirect_target()
|
||||
-> anyhow::Result<()> {
|
||||
|
|
|
|||
|
|
@ -170,31 +170,61 @@ pub(crate) fn unexpected_message_response(expect: &str) -> Response<BoxBody<Byte
|
|||
|
||||
pub(crate) async fn expect_json<B>(
|
||||
body: B,
|
||||
max_bytes: usize,
|
||||
) -> Result<ClientJsonRpcMessage, Response<BoxBody<Bytes, Infallible>>>
|
||||
where
|
||||
B: Body + Send + 'static,
|
||||
B::Error: Display,
|
||||
{
|
||||
match body.collect().await {
|
||||
Ok(bytes) => {
|
||||
match serde_json::from_reader::<_, ClientJsonRpcMessage>(bytes.aggregate().reader()) {
|
||||
Ok(message) => Ok(message),
|
||||
Err(e) => {
|
||||
let mut collected = bytes::BytesMut::new();
|
||||
let mut body = std::pin::pin!(body);
|
||||
loop {
|
||||
let frame = futures::future::poll_fn(|cx| body.as_mut().poll_frame(cx)).await;
|
||||
match frame {
|
||||
None => break,
|
||||
Some(Ok(frame)) => {
|
||||
let Ok(mut data) = frame.into_data() else {
|
||||
continue;
|
||||
};
|
||||
if data.remaining() > max_bytes.saturating_sub(collected.len()) {
|
||||
let response = Response::builder()
|
||||
.status(http::StatusCode::UNSUPPORTED_MEDIA_TYPE)
|
||||
.status(http::StatusCode::PAYLOAD_TOO_LARGE)
|
||||
.body(
|
||||
Full::new(Bytes::from(format!("fail to deserialize request body {e}")))
|
||||
.boxed(),
|
||||
Full::new(Bytes::from(format!(
|
||||
"Payload Too Large: request body exceeds {max_bytes} bytes"
|
||||
)))
|
||||
.boxed(),
|
||||
)
|
||||
.expect("valid response");
|
||||
Err(response)
|
||||
return Err(response);
|
||||
}
|
||||
while data.has_remaining() {
|
||||
let chunk = data.chunk();
|
||||
let chunk_len = chunk.len();
|
||||
collected.extend_from_slice(chunk);
|
||||
data.advance(chunk_len);
|
||||
}
|
||||
}
|
||||
Some(Err(e)) => {
|
||||
let response = Response::builder()
|
||||
.status(http::StatusCode::INTERNAL_SERVER_ERROR)
|
||||
.body(
|
||||
Full::new(Bytes::from(format!("Failed to read request body: {e}"))).boxed(),
|
||||
)
|
||||
.expect("valid response");
|
||||
return Err(response);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
match serde_json::from_slice::<ClientJsonRpcMessage>(&collected) {
|
||||
Ok(message) => Ok(message),
|
||||
Err(e) => {
|
||||
let response = Response::builder()
|
||||
.status(http::StatusCode::INTERNAL_SERVER_ERROR)
|
||||
.body(Full::new(Bytes::from(format!("Failed to read request body: {e}"))).boxed())
|
||||
.status(http::StatusCode::UNSUPPORTED_MEDIA_TYPE)
|
||||
.body(
|
||||
Full::new(Bytes::from(format!("fail to deserialize request body {e}"))).boxed(),
|
||||
)
|
||||
.expect("valid response");
|
||||
Err(response)
|
||||
}
|
||||
|
|
@ -203,9 +233,17 @@ where
|
|||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::convert::Infallible;
|
||||
|
||||
use futures::stream;
|
||||
use http_body::Frame;
|
||||
use http_body_util::StreamBody;
|
||||
|
||||
use super::*;
|
||||
use crate::model::{EmptyResult, JsonRpcResponse, JsonRpcVersion2_0, RequestId, ServerResult};
|
||||
|
||||
const INITIALIZE_REQUEST: &str = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"test","version":"0.1"}}}"#;
|
||||
|
||||
fn dummy_message() -> ServerJsonRpcMessage {
|
||||
ServerJsonRpcMessage::Response(JsonRpcResponse {
|
||||
jsonrpc: JsonRpcVersion2_0,
|
||||
|
|
@ -245,4 +283,62 @@ mod tests {
|
|||
assert!(msg.message.is_none());
|
||||
assert_eq!(msg.retry, Some(Duration::from_secs(5)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn expect_json_accepts_body_under_limit() {
|
||||
let body = Full::new(Bytes::from_static(INITIALIZE_REQUEST.as_bytes()));
|
||||
let result = expect_json(body, 4 * 1024 * 1024).await;
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn expect_json_accepts_non_contiguous_body_data() {
|
||||
let split = INITIALIZE_REQUEST.len() / 2;
|
||||
let data = Bytes::copy_from_slice(&INITIALIZE_REQUEST.as_bytes()[..split]).chain(
|
||||
Bytes::copy_from_slice(&INITIALIZE_REQUEST.as_bytes()[split..]),
|
||||
);
|
||||
let result = expect_json(Full::new(data), 4 * 1024 * 1024).await;
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn expect_json_accepts_multi_frame_body() {
|
||||
let split = INITIALIZE_REQUEST.len() / 2;
|
||||
let body = StreamBody::new(stream::iter([
|
||||
Ok::<_, Infallible>(Frame::data(Bytes::copy_from_slice(
|
||||
&INITIALIZE_REQUEST.as_bytes()[..split],
|
||||
))),
|
||||
Ok(Frame::data(Bytes::copy_from_slice(
|
||||
&INITIALIZE_REQUEST.as_bytes()[split..],
|
||||
))),
|
||||
]));
|
||||
let result = expect_json(body, 4 * 1024 * 1024).await;
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn expect_json_rejects_oversized_body() {
|
||||
let big_body = Full::new(Bytes::from(vec![b'x'; 128]));
|
||||
let result = expect_json(big_body, 64).await;
|
||||
let response = result.unwrap_err();
|
||||
assert_eq!(response.status(), http::StatusCode::PAYLOAD_TOO_LARGE);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn expect_json_rejects_oversized_multi_frame_body() {
|
||||
let body = StreamBody::new(stream::iter([
|
||||
Ok::<_, Infallible>(Frame::data(Bytes::from_static(b"12345678"))),
|
||||
Ok(Frame::data(Bytes::from_static(b"9"))),
|
||||
]));
|
||||
let response = expect_json(body, 8).await.unwrap_err();
|
||||
assert_eq!(response.status(), http::StatusCode::PAYLOAD_TOO_LARGE);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn expect_json_returns_415_for_invalid_json_under_limit() {
|
||||
let body = Full::new(Bytes::from("not valid json"));
|
||||
let result = expect_json(body, 4 * 1024 * 1024).await;
|
||||
let response = result.unwrap_err();
|
||||
assert_eq!(response.status(), http::StatusCode::UNSUPPORTED_MEDIA_TYPE);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,20 +1,23 @@
|
|||
use std::{borrow::Cow, collections::HashMap, sync::Arc};
|
||||
|
||||
use bytes::Bytes;
|
||||
use futures::{StreamExt, stream::BoxStream};
|
||||
use futures::stream::BoxStream;
|
||||
use http::{HeaderName, HeaderValue, Method, Request, StatusCode, header::WWW_AUTHENTICATE};
|
||||
use http_body_util::{BodyExt, Full};
|
||||
use hyper::body::Incoming;
|
||||
use hyper_util::rt::TokioIo;
|
||||
use sse_stream::{Sse, SseStream};
|
||||
use sse_stream::Sse;
|
||||
use tokio::net::UnixStream;
|
||||
|
||||
use crate::{
|
||||
model::{ClientJsonRpcMessage, ServerJsonRpcMessage},
|
||||
transport::{
|
||||
common::http_header::{
|
||||
EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_SESSION_ID, JSON_MIME_TYPE,
|
||||
extract_scope_from_header, validate_custom_header,
|
||||
common::{
|
||||
client_side_sse::{DEFAULT_MAX_SSE_EVENT_SIZE, bounded_sse_stream},
|
||||
http_header::{
|
||||
EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_SESSION_ID, JSON_MIME_TYPE,
|
||||
extract_scope_from_header, validate_custom_header,
|
||||
},
|
||||
},
|
||||
streamable_http_client::*,
|
||||
},
|
||||
|
|
@ -168,6 +171,26 @@ impl StreamableHttpClient for UnixSocketHttpClient {
|
|||
session_id: Option<Arc<str>>,
|
||||
auth_token: Option<String>,
|
||||
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
|
||||
self.post_message_with_max_sse_event_size(
|
||||
uri,
|
||||
message,
|
||||
session_id,
|
||||
auth_token,
|
||||
custom_headers,
|
||||
DEFAULT_MAX_SSE_EVENT_SIZE,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn post_message_with_max_sse_event_size(
|
||||
&self,
|
||||
uri: Arc<str>,
|
||||
message: ClientJsonRpcMessage,
|
||||
session_id: Option<Arc<str>>,
|
||||
auth_token: Option<String>,
|
||||
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
max_sse_event_size: usize,
|
||||
) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
|
||||
let json_body = serde_json::to_string(&message)
|
||||
.map_err(|e| StreamableHttpError::Client(UnixSocketError::Json(e)))?;
|
||||
|
|
@ -282,7 +305,8 @@ impl StreamableHttpClient for UnixSocketHttpClient {
|
|||
|
||||
match content_type {
|
||||
Some(ref ct) if ct.as_bytes().starts_with(EVENT_STREAM_MIME_TYPE.as_bytes()) => {
|
||||
let sse_stream = SseStream::new(response.into_body()).boxed();
|
||||
let sse_stream =
|
||||
bounded_sse_stream(response.into_body().into_data_stream(), max_sse_event_size);
|
||||
Ok(StreamableHttpPostResponse::Sse(sse_stream, session_id))
|
||||
}
|
||||
Some(ref ct) if ct.as_bytes().starts_with(JSON_MIME_TYPE.as_bytes()) => {
|
||||
|
|
@ -357,6 +381,27 @@ impl StreamableHttpClient for UnixSocketHttpClient {
|
|||
auth_token: Option<String>,
|
||||
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
) -> Result<BoxStream<'static, Result<Sse, sse_stream::Error>>, StreamableHttpError<Self::Error>>
|
||||
{
|
||||
self.get_stream_with_max_sse_event_size(
|
||||
uri,
|
||||
session_id,
|
||||
last_event_id,
|
||||
auth_token,
|
||||
custom_headers,
|
||||
DEFAULT_MAX_SSE_EVENT_SIZE,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn get_stream_with_max_sse_event_size(
|
||||
&self,
|
||||
uri: Arc<str>,
|
||||
session_id: Arc<str>,
|
||||
last_event_id: Option<String>,
|
||||
auth_token: Option<String>,
|
||||
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
max_sse_event_size: usize,
|
||||
) -> Result<BoxStream<'static, Result<Sse, sse_stream::Error>>, StreamableHttpError<Self::Error>>
|
||||
{
|
||||
let mut builder = Request::builder()
|
||||
.method(Method::GET)
|
||||
|
|
@ -444,7 +489,10 @@ impl StreamableHttpClient for UnixSocketHttpClient {
|
|||
}
|
||||
}
|
||||
|
||||
Ok(SseStream::new(response.into_body()).boxed())
|
||||
Ok(bounded_sse_stream(
|
||||
response.into_body().into_data_stream(),
|
||||
max_sse_event_size,
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -13,7 +13,9 @@ use thiserror::Error;
|
|||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::debug;
|
||||
|
||||
use super::common::client_side_sse::{ExponentialBackoff, SseRetryPolicy, SseStreamReconnect};
|
||||
use super::common::client_side_sse::{
|
||||
DEFAULT_MAX_SSE_EVENT_SIZE, ExponentialBackoff, SseRetryPolicy, SseStreamReconnect,
|
||||
};
|
||||
use crate::{
|
||||
RoleClient,
|
||||
model::{
|
||||
|
|
@ -266,6 +268,12 @@ impl StreamableHttpPostResponse {
|
|||
}
|
||||
}
|
||||
|
||||
/// HTTP backend used by [`StreamableHttpClientTransport`].
|
||||
///
|
||||
/// Custom implementations that parse SSE responses must override
|
||||
/// [`Self::post_message_with_max_sse_event_size`] and
|
||||
/// [`Self::get_stream_with_max_sse_event_size`] to enforce the transport's
|
||||
/// configured event-size limit.
|
||||
pub trait StreamableHttpClient: Clone + Send + 'static {
|
||||
type Error: std::error::Error + Send + Sync + 'static;
|
||||
fn post_message(
|
||||
|
|
@ -278,6 +286,29 @@ pub trait StreamableHttpClient: Clone + Send + 'static {
|
|||
) -> impl Future<Output = Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>>>
|
||||
+ Send
|
||||
+ '_;
|
||||
/// Send a message while enforcing a maximum raw SSE event size.
|
||||
///
|
||||
/// `max_sse_event_size` is not a per-request option: it is the
|
||||
/// transport-wide [`StreamableHttpClientTransportConfig::max_sse_event_size`]
|
||||
/// value, passed identically on every call because the limit must be applied
|
||||
/// inside the client (at the raw byte layer, before SSE parsing) rather than
|
||||
/// by the caller.
|
||||
///
|
||||
/// Custom clients that parse SSE responses should override this method.
|
||||
/// The default implementation delegates to [`Self::post_message`].
|
||||
fn post_message_with_max_sse_event_size(
|
||||
&self,
|
||||
uri: Arc<str>,
|
||||
message: ClientJsonRpcMessage,
|
||||
session_id: Option<Arc<str>>,
|
||||
auth_header: Option<String>,
|
||||
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
_max_sse_event_size: usize,
|
||||
) -> impl Future<Output = Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>>>
|
||||
+ Send
|
||||
+ '_ {
|
||||
self.post_message(uri, message, session_id, auth_header, custom_headers)
|
||||
}
|
||||
fn delete_session(
|
||||
&self,
|
||||
uri: Arc<str>,
|
||||
|
|
@ -299,6 +330,33 @@ pub trait StreamableHttpClient: Clone + Send + 'static {
|
|||
>,
|
||||
> + Send
|
||||
+ '_;
|
||||
/// Open an SSE stream while enforcing a maximum raw event size.
|
||||
///
|
||||
/// `max_sse_event_size` is not a per-request option: it is the
|
||||
/// transport-wide [`StreamableHttpClientTransportConfig::max_sse_event_size`]
|
||||
/// value, passed identically on every call because the limit must be applied
|
||||
/// inside the client (at the raw byte layer, before SSE parsing) rather than
|
||||
/// by the caller.
|
||||
///
|
||||
/// Custom clients that parse SSE responses should override this method.
|
||||
/// The default implementation delegates to [`Self::get_stream`].
|
||||
fn get_stream_with_max_sse_event_size(
|
||||
&self,
|
||||
uri: Arc<str>,
|
||||
session_id: Arc<str>,
|
||||
last_event_id: Option<String>,
|
||||
auth_header: Option<String>,
|
||||
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
_max_sse_event_size: usize,
|
||||
) -> impl Future<
|
||||
Output = Result<
|
||||
BoxStream<'static, Result<Sse, SseError>>,
|
||||
StreamableHttpError<Self::Error>,
|
||||
>,
|
||||
> + Send
|
||||
+ '_ {
|
||||
self.get_stream(uri, session_id, last_event_id, auth_header, custom_headers)
|
||||
}
|
||||
}
|
||||
|
||||
#[non_exhaustive]
|
||||
|
|
@ -313,6 +371,7 @@ struct StreamableHttpClientReconnect<C> {
|
|||
pub uri: Arc<str>,
|
||||
pub auth_header: Option<String>,
|
||||
pub custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
pub max_sse_event_size: usize,
|
||||
}
|
||||
|
||||
impl<C: StreamableHttpClient> SseStreamReconnect for StreamableHttpClientReconnect<C> {
|
||||
|
|
@ -324,13 +383,25 @@ impl<C: StreamableHttpClient> SseStreamReconnect for StreamableHttpClientReconne
|
|||
let session_id = self.session_id.clone();
|
||||
let auth_header = self.auth_header.clone();
|
||||
let custom_headers = self.custom_headers.clone();
|
||||
let max_sse_event_size = self.max_sse_event_size;
|
||||
let last_event_id = last_event_id.map(|s| s.to_owned());
|
||||
Box::pin(async move {
|
||||
client
|
||||
.get_stream(uri, session_id, last_event_id, auth_header, custom_headers)
|
||||
.get_stream_with_max_sse_event_size(
|
||||
uri,
|
||||
session_id,
|
||||
last_event_id,
|
||||
auth_header,
|
||||
custom_headers,
|
||||
max_sse_event_size,
|
||||
)
|
||||
.await
|
||||
})
|
||||
}
|
||||
|
||||
fn map_fatal_stream_error(&mut self, error: SseError) -> Option<Self::Error> {
|
||||
Some(StreamableHttpError::Sse(error))
|
||||
}
|
||||
}
|
||||
|
||||
/// Info retained for cleaning up the session when the worker exits.
|
||||
|
|
@ -485,6 +556,7 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
|
|||
uri: Arc<str>,
|
||||
auth_header: Option<String>,
|
||||
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
max_sse_event_size: usize,
|
||||
retry_config: Arc<dyn SseRetryPolicy>,
|
||||
) -> impl Stream<Item = Result<ServerJsonRpcMessage, StreamableHttpError<C::Error>>> + Send + 'static
|
||||
{
|
||||
|
|
@ -496,6 +568,7 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
|
|||
uri,
|
||||
auth_header,
|
||||
custom_headers,
|
||||
max_sse_event_size,
|
||||
},
|
||||
retry_config,
|
||||
)
|
||||
|
|
@ -514,6 +587,7 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
|
|||
uri: Arc<str>,
|
||||
auth_header: Option<String>,
|
||||
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
max_sse_event_size: usize,
|
||||
retry_config: Arc<dyn SseRetryPolicy>,
|
||||
) -> BoxStream<'static, Result<ServerJsonRpcMessage, StreamableHttpError<C::Error>>> {
|
||||
match session_id {
|
||||
|
|
@ -524,6 +598,7 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
|
|||
uri,
|
||||
auth_header,
|
||||
custom_headers,
|
||||
max_sse_event_size,
|
||||
retry_config,
|
||||
)
|
||||
.boxed(),
|
||||
|
|
@ -591,15 +666,17 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
|
|||
uri: Arc<str>,
|
||||
auth_header: Option<String>,
|
||||
custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
max_sse_event_size: usize,
|
||||
) -> Result<(Option<Arc<str>>, HashMap<HeaderName, HeaderValue>), StreamableHttpError<C::Error>>
|
||||
{
|
||||
let (init_msg, new_session_id_str) = client
|
||||
.post_message(
|
||||
.post_message_with_max_sse_event_size(
|
||||
uri.clone(),
|
||||
saved_init_request,
|
||||
None,
|
||||
auth_header.clone(),
|
||||
custom_headers.clone(),
|
||||
max_sse_event_size,
|
||||
)
|
||||
.await?
|
||||
.expect_initialized::<C::Error>()
|
||||
|
|
@ -624,12 +701,13 @@ impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
|
|||
&negotiated_version,
|
||||
);
|
||||
client
|
||||
.post_message(
|
||||
.post_message_with_max_sse_event_size(
|
||||
uri,
|
||||
initialized_notification,
|
||||
new_session_id.clone(),
|
||||
auth_header,
|
||||
initialized_headers,
|
||||
max_sse_event_size,
|
||||
)
|
||||
.await?
|
||||
.expect_accepted_or_json::<C::Error>()?;
|
||||
|
|
@ -670,12 +748,13 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
let saved_init_request = initialize_request.clone();
|
||||
let (message, session_id) = match self
|
||||
.client
|
||||
.post_message(
|
||||
.post_message_with_max_sse_event_size(
|
||||
config.uri.clone(),
|
||||
initialize_request,
|
||||
None,
|
||||
config.auth_header.clone(),
|
||||
config.custom_headers.clone(),
|
||||
config.max_sse_event_size,
|
||||
)
|
||||
.await
|
||||
{
|
||||
|
|
@ -730,12 +809,13 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
&negotiated_version,
|
||||
);
|
||||
self.client
|
||||
.post_message(
|
||||
.post_message_with_max_sse_event_size(
|
||||
config.uri.clone(),
|
||||
initialized_notification.message,
|
||||
session_id.clone(),
|
||||
config.auth_header.clone(),
|
||||
initialized_headers,
|
||||
config.max_sse_event_size,
|
||||
)
|
||||
.await
|
||||
.map_err(WorkerQuitReason::fatal_context(
|
||||
|
|
@ -765,15 +845,17 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
let config_uri = config.uri.clone();
|
||||
let config_auth_header = config.auth_header.clone();
|
||||
let spawn_headers = protocol_headers.clone();
|
||||
let max_sse_event_size = config.max_sse_event_size;
|
||||
|
||||
streams.spawn(async move {
|
||||
match client
|
||||
.get_stream(
|
||||
.get_stream_with_max_sse_event_size(
|
||||
uri.clone(),
|
||||
session_id.clone(),
|
||||
None,
|
||||
auth_header.clone(),
|
||||
spawn_headers.clone(),
|
||||
max_sse_event_size,
|
||||
)
|
||||
.await
|
||||
{
|
||||
|
|
@ -786,6 +868,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
uri: config_uri,
|
||||
auth_header: config_auth_header,
|
||||
custom_headers: spawn_headers,
|
||||
max_sse_event_size,
|
||||
},
|
||||
retry_config,
|
||||
);
|
||||
|
|
@ -855,12 +938,13 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
);
|
||||
let response = self
|
||||
.client
|
||||
.post_message(
|
||||
.post_message_with_max_sse_event_size(
|
||||
config.uri.clone(),
|
||||
message.clone(),
|
||||
session_id.clone(),
|
||||
config.auth_header.clone(),
|
||||
request_headers,
|
||||
config.max_sse_event_size,
|
||||
)
|
||||
.await;
|
||||
let send_result = match response {
|
||||
|
|
@ -879,6 +963,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
config.uri.clone(),
|
||||
config.auth_header.clone(),
|
||||
config.custom_headers.clone(),
|
||||
config.max_sse_event_size,
|
||||
)
|
||||
.await
|
||||
{
|
||||
|
|
@ -926,14 +1011,16 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
let config_uri = config.uri.clone();
|
||||
let config_auth = config.auth_header.clone();
|
||||
let spawn_headers = protocol_headers.clone();
|
||||
let max_sse_event_size = config.max_sse_event_size;
|
||||
streams.spawn(async move {
|
||||
match client
|
||||
.get_stream(
|
||||
.get_stream_with_max_sse_event_size(
|
||||
uri,
|
||||
new_sid.clone(),
|
||||
None,
|
||||
auth_header.clone(),
|
||||
spawn_headers.clone(),
|
||||
max_sse_event_size,
|
||||
)
|
||||
.await
|
||||
{
|
||||
|
|
@ -946,6 +1033,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
uri: config_uri,
|
||||
auth_header: config_auth,
|
||||
custom_headers: spawn_headers,
|
||||
max_sse_event_size,
|
||||
},
|
||||
retry_config,
|
||||
);
|
||||
|
|
@ -981,12 +1069,13 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
);
|
||||
let retry_response = self
|
||||
.client
|
||||
.post_message(
|
||||
.post_message_with_max_sse_event_size(
|
||||
config.uri.clone(),
|
||||
message,
|
||||
session_id.clone(),
|
||||
config.auth_header.clone(),
|
||||
retry_headers,
|
||||
config.max_sse_event_size,
|
||||
)
|
||||
.await;
|
||||
match retry_response {
|
||||
|
|
@ -1021,6 +1110,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
config.uri.clone(),
|
||||
config.auth_header.clone(),
|
||||
protocol_headers.clone(),
|
||||
config.max_sse_event_size,
|
||||
self.config.retry_config.clone(),
|
||||
);
|
||||
streams.spawn(Self::execute_sse_stream(
|
||||
|
|
@ -1064,6 +1154,7 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
|
|||
config.uri.clone(),
|
||||
config.auth_header.clone(),
|
||||
protocol_headers.clone(),
|
||||
config.max_sse_event_size,
|
||||
self.config.retry_config.clone(),
|
||||
);
|
||||
streams.spawn(Self::execute_sse_stream(
|
||||
|
|
@ -1339,6 +1430,12 @@ pub struct StreamableHttpClientTransportConfig {
|
|||
pub auth_header: Option<String>,
|
||||
/// Custom HTTP headers to include with every request
|
||||
pub custom_headers: HashMap<HeaderName, HeaderValue>,
|
||||
/// Maximum raw size of one SSE event accepted from the server.
|
||||
///
|
||||
/// The built-in reqwest and Unix socket clients enforce this value. Custom
|
||||
/// [`StreamableHttpClient`] implementations must override the corresponding
|
||||
/// `*_with_max_sse_event_size` methods to enforce it.
|
||||
pub max_sse_event_size: usize,
|
||||
/// Enables transparent recovery when the server reports an expired session (`HTTP 404`).
|
||||
///
|
||||
/// When enabled, the transport performs one automatic recovery attempt:
|
||||
|
|
@ -1397,6 +1494,12 @@ impl StreamableHttpClientTransportConfig {
|
|||
self
|
||||
}
|
||||
|
||||
/// Set the maximum raw size of one SSE event accepted from the server.
|
||||
pub fn max_sse_event_size(mut self, bytes: usize) -> Self {
|
||||
self.max_sse_event_size = bytes;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set whether the transport should attempt transparent re-initialization on session expiration
|
||||
/// See [`Self::reinit_on_expired_session`] for details.
|
||||
/// # Example
|
||||
|
|
@ -1420,6 +1523,7 @@ impl Default for StreamableHttpClientTransportConfig {
|
|||
allow_stateless: true,
|
||||
auth_header: None,
|
||||
custom_headers: HashMap::new(),
|
||||
max_sse_event_size: DEFAULT_MAX_SSE_EVENT_SIZE,
|
||||
reinit_on_expired_session: true,
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -39,6 +39,9 @@ use crate::{
|
|||
},
|
||||
};
|
||||
|
||||
/// Default maximum POST request body size (4 MiB).
|
||||
pub(crate) const DEFAULT_MAX_REQUEST_BODY_BYTES: usize = 4 * 1024 * 1024;
|
||||
|
||||
#[non_exhaustive]
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct StreamableHttpServerConfig {
|
||||
|
|
@ -100,6 +103,12 @@ pub struct StreamableHttpServerConfig {
|
|||
/// };
|
||||
/// ```
|
||||
pub session_store: Option<Arc<dyn SessionStore>>,
|
||||
/// Maximum POST request body size in bytes.
|
||||
///
|
||||
/// Enforced while streaming the body, independent of `Content-Length`,
|
||||
/// chunked transfer encoding, or HTTP version. Oversized payloads receive
|
||||
/// a `413 Payload Too Large` response.
|
||||
pub max_request_body_bytes: usize,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for dyn SessionStore {
|
||||
|
|
@ -119,6 +128,7 @@ impl Default for StreamableHttpServerConfig {
|
|||
allowed_hosts: vec!["localhost".into(), "127.0.0.1".into(), "::1".into()],
|
||||
allowed_origins: vec![],
|
||||
session_store: None,
|
||||
max_request_body_bytes: DEFAULT_MAX_REQUEST_BODY_BYTES,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -172,6 +182,12 @@ impl StreamableHttpServerConfig {
|
|||
self.cancellation_token = token;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the maximum POST request body size in bytes.
|
||||
pub fn with_max_request_body_bytes(mut self, bytes: usize) -> Self {
|
||||
self.max_request_body_bytes = bytes;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[expect(
|
||||
|
|
@ -1141,7 +1157,7 @@ where
|
|||
|
||||
// json deserialize request body
|
||||
let (part, body) = request.into_parts();
|
||||
let mut message = match expect_json(body).await {
|
||||
let mut message = match expect_json(body, self.config.max_request_body_bytes).await {
|
||||
Ok(message) => message,
|
||||
Err(response) => return Ok(response),
|
||||
};
|
||||
|
|
|
|||
Loading…
Reference in a new issue