soma_mcp_client/upstream/
http_body_cap.rs

1use std::{borrow::Cow, collections::HashMap, sync::Arc};
2
3use futures::{StreamExt, stream::BoxStream};
4use http::{HeaderName, HeaderValue};
5use reqwest::header::{ACCEPT, WWW_AUTHENTICATE};
6use rmcp::{
7    model::{ClientJsonRpcMessage, JsonRpcMessage, ServerJsonRpcMessage},
8    transport::{
9        common::http_header::{
10            EVENT_STREAM_MIME_TYPE, HEADER_LAST_EVENT_ID, HEADER_MCP_PROTOCOL_VERSION,
11            HEADER_SESSION_ID, JSON_MIME_TYPE,
12        },
13        streamable_http_client::{
14            AuthRequiredError, InsufficientScopeError, SseError, StreamableHttpClient,
15            StreamableHttpError, StreamableHttpPostResponse,
16        },
17    },
18};
19use sse_stream::{Sse, SseStream};
20
21#[derive(Clone)]
22pub struct BodyCappedHttpClient {
23    inner: reqwest::Client,
24    json_max_bytes: usize,
25    sse_event_max_bytes: usize,
26}
27
28impl BodyCappedHttpClient {
29    #[must_use]
30    pub fn new(inner: reqwest::Client, json_max_bytes: usize, sse_event_max_bytes: usize) -> Self {
31        Self {
32            inner,
33            json_max_bytes,
34            sse_event_max_bytes,
35        }
36    }
37
38    #[must_use]
39    pub fn default_with_caps(json_max_bytes: usize, sse_event_max_bytes: usize) -> Self {
40        let inner = reqwest::Client::builder()
41            .pool_max_idle_per_host(0)
42            .redirect(reqwest::redirect::Policy::none())
43            .build()
44            .expect("failed to build gateway HTTP client");
45        Self::new(inner, json_max_bytes, sse_event_max_bytes)
46    }
47}
48
49impl StreamableHttpClient for BodyCappedHttpClient {
50    type Error = reqwest::Error;
51
52    async fn post_message(
53        &self,
54        uri: Arc<str>,
55        message: ClientJsonRpcMessage,
56        session_id: Option<Arc<str>>,
57        auth_token: Option<String>,
58        custom_headers: HashMap<HeaderName, HeaderValue>,
59    ) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
60        let session_was_attached = session_id.is_some();
61        let mut request = self
62            .inner
63            .post(uri.as_ref())
64            .header(ACCEPT, [EVENT_STREAM_MIME_TYPE, JSON_MIME_TYPE].join(", "));
65        if let Some(token) = auth_token {
66            request = request.bearer_auth(token);
67        }
68        if let Some(session_id) = session_id {
69            request = request.header(HEADER_SESSION_ID, session_id.as_ref());
70        }
71        let response = apply_custom_headers(request, custom_headers)?
72            .json(&message)
73            .send()
74            .await
75            .map_err(StreamableHttpError::Client)?;
76        response_to_post_result(response, message, session_was_attached, self.json_max_bytes).await
77    }
78
79    async fn delete_session(
80        &self,
81        uri: Arc<str>,
82        session_id: Arc<str>,
83        auth_token: Option<String>,
84        custom_headers: HashMap<HeaderName, HeaderValue>,
85    ) -> Result<(), StreamableHttpError<Self::Error>> {
86        let mut request = self
87            .inner
88            .delete(uri.as_ref())
89            .header(HEADER_SESSION_ID, session_id.as_ref());
90        if let Some(token) = auth_token {
91            request = request.bearer_auth(token);
92        }
93        let response = apply_custom_headers(request, custom_headers)?
94            .send()
95            .await
96            .map_err(StreamableHttpError::Client)?;
97        if response.status() == reqwest::StatusCode::METHOD_NOT_ALLOWED {
98            return Err(StreamableHttpError::ServerDoesNotSupportDeleteSession);
99        }
100        response
101            .error_for_status()
102            .map(|_| ())
103            .map_err(StreamableHttpError::Client)
104    }
105
106    async fn get_stream(
107        &self,
108        uri: Arc<str>,
109        session_id: Option<Arc<str>>,
110        last_event_id: Option<String>,
111        auth_token: Option<String>,
112        custom_headers: HashMap<HeaderName, HeaderValue>,
113    ) -> Result<BoxStream<'static, Result<Sse, SseError>>, StreamableHttpError<Self::Error>> {
114        let mut request = self
115            .inner
116            .get(uri.as_ref())
117            .header(ACCEPT, [EVENT_STREAM_MIME_TYPE, JSON_MIME_TYPE].join(", "));
118        if let Some(session_id) = session_id {
119            request = request.header(HEADER_SESSION_ID, session_id.as_ref());
120        }
121        if let Some(last_event_id) = last_event_id {
122            request = request.header(HEADER_LAST_EVENT_ID, last_event_id);
123        }
124        if let Some(token) = auth_token {
125            request = request.bearer_auth(token);
126        }
127        let response = apply_custom_headers(request, custom_headers)?
128            .send()
129            .await
130            .map_err(StreamableHttpError::Client)?;
131        if response.status() == reqwest::StatusCode::METHOD_NOT_ALLOWED {
132            return Err(StreamableHttpError::ServerDoesNotSupportSse);
133        }
134        let response = response
135            .error_for_status()
136            .map_err(StreamableHttpError::Client)?;
137        ensure_stream_content_type(&response)?;
138        let capped = per_event_capped_stream(response.bytes_stream(), self.sse_event_max_bytes);
139        Ok(SseStream::from_bytes_stream(capped).boxed())
140    }
141}
142
143fn apply_custom_headers(
144    mut request: reqwest::RequestBuilder,
145    custom_headers: HashMap<HeaderName, HeaderValue>,
146) -> Result<reqwest::RequestBuilder, StreamableHttpError<reqwest::Error>> {
147    for (name, value) in custom_headers {
148        validate_custom_header(&name).map_err(StreamableHttpError::ReservedHeaderConflict)?;
149        request = request.header(name, value);
150    }
151    Ok(request)
152}
153
154fn validate_custom_header(name: &HeaderName) -> Result<(), String> {
155    let reserved = [
156        "accept",
157        HEADER_SESSION_ID,
158        HEADER_LAST_EVENT_ID,
159        HEADER_MCP_PROTOCOL_VERSION,
160    ];
161    if reserved
162        .iter()
163        .any(|reserved| name.as_str().eq_ignore_ascii_case(reserved))
164        && !name
165            .as_str()
166            .eq_ignore_ascii_case(HEADER_MCP_PROTOCOL_VERSION)
167    {
168        return Err(name.to_string());
169    }
170    Ok(())
171}
172
173async fn response_to_post_result(
174    response: reqwest::Response,
175    message: ClientJsonRpcMessage,
176    session_was_attached: bool,
177    max_bytes: usize,
178) -> Result<StreamableHttpPostResponse, StreamableHttpError<reqwest::Error>> {
179    if let Some(error) = auth_error(&response) {
180        return Err(error);
181    }
182    let status = response.status();
183    if matches!(
184        status,
185        reqwest::StatusCode::ACCEPTED | reqwest::StatusCode::NO_CONTENT
186    ) {
187        return Ok(StreamableHttpPostResponse::Accepted);
188    }
189    if status == reqwest::StatusCode::NOT_FOUND && session_was_attached {
190        return Err(StreamableHttpError::SessionExpired);
191    }
192    let content_type = content_type(&response);
193    let session_id = response
194        .headers()
195        .get(HEADER_SESSION_ID)
196        .and_then(|value| value.to_str().ok())
197        .map(ToOwned::to_owned);
198    if status.is_success() && response.content_length() == Some(0) && is_empty_ok(&message) {
199        return Ok(StreamableHttpPostResponse::Accepted);
200    }
201    if !status.is_success() {
202        return non_success_response(status, content_type, response, max_bytes).await;
203    }
204    match content_type.as_deref() {
205        Some(ct) if ct.as_bytes().starts_with(EVENT_STREAM_MIME_TYPE.as_bytes()) => {
206            let capped = per_event_capped_stream(response.bytes_stream(), max_bytes);
207            Ok(StreamableHttpPostResponse::Sse(
208                SseStream::from_bytes_stream(capped).boxed(),
209                session_id,
210            ))
211        }
212        Some(ct) if ct.as_bytes().starts_with(JSON_MIME_TYPE.as_bytes()) => {
213            let bytes = read_body_capped(response, max_bytes).await?;
214            match serde_json::from_slice::<ServerJsonRpcMessage>(&bytes) {
215                Ok(message) => Ok(StreamableHttpPostResponse::Json(message, session_id)),
216                Err(_) => Ok(StreamableHttpPostResponse::Accepted),
217            }
218        }
219        _ => Err(StreamableHttpError::UnexpectedContentType(content_type)),
220    }
221}
222
223fn auth_error(response: &reqwest::Response) -> Option<StreamableHttpError<reqwest::Error>> {
224    let header = response.headers().get(WWW_AUTHENTICATE)?.to_str().ok()?;
225    match response.status() {
226        reqwest::StatusCode::UNAUTHORIZED => Some(StreamableHttpError::AuthRequired(
227            AuthRequiredError::new(header.to_owned()),
228        )),
229        reqwest::StatusCode::FORBIDDEN => Some(StreamableHttpError::InsufficientScope(
230            InsufficientScopeError::new(header.to_owned(), extract_scope(header)),
231        )),
232        _ => None,
233    }
234}
235
236async fn non_success_response(
237    status: reqwest::StatusCode,
238    content_type: Option<String>,
239    response: reqwest::Response,
240    max_bytes: usize,
241) -> Result<StreamableHttpPostResponse, StreamableHttpError<reqwest::Error>> {
242    let bytes = read_body_capped(response, max_bytes).await?;
243    let body = String::from_utf8_lossy(&bytes);
244    if content_type
245        .as_deref()
246        .is_some_and(|ct| ct.as_bytes().starts_with(JSON_MIME_TYPE.as_bytes()))
247        && let Some(message) = parse_json_rpc_error(&body)
248    {
249        return Ok(StreamableHttpPostResponse::Json(message, None));
250    }
251    Err(StreamableHttpError::UnexpectedServerResponse(Cow::Owned(
252        format!("HTTP {status}: {body}"),
253    )))
254}
255
256async fn read_body_capped(
257    response: reqwest::Response,
258    max_bytes: usize,
259) -> Result<Vec<u8>, StreamableHttpError<reqwest::Error>> {
260    if let Some(length) = response.content_length()
261        && length > max_bytes as u64
262    {
263        return Err(too_large(format!(
264            "response_too_large: declared {length} bytes, max {max_bytes}"
265        )));
266    }
267    let mut stream = response.bytes_stream();
268    let mut bytes = Vec::new();
269    while let Some(chunk) = stream.next().await {
270        let chunk = chunk.map_err(StreamableHttpError::Client)?;
271        if bytes.len().saturating_add(chunk.len()) > max_bytes {
272            return Err(too_large(format!(
273                "response_too_large: streamed {} bytes, max {max_bytes}",
274                bytes.len() + chunk.len()
275            )));
276        }
277        bytes.extend_from_slice(&chunk);
278    }
279    Ok(bytes)
280}
281
282fn per_event_capped_stream(
283    inner: impl futures::Stream<Item = reqwest::Result<bytes::Bytes>> + Send + 'static,
284    max_bytes: usize,
285) -> BoxStream<'static, Result<bytes::Bytes, CappedStreamError>> {
286    inner
287        .scan((0usize, false), move |state, item| {
288            let result = match item {
289                Ok(chunk) => account_event_bytes(&chunk, state, max_bytes).map(|_| chunk),
290                Err(error) => Err(CappedStreamError::Reqwest(error)),
291            };
292            futures::future::ready(Some(result))
293        })
294        .boxed()
295}
296
297fn account_event_bytes(
298    chunk: &[u8],
299    state: &mut (usize, bool),
300    max_bytes: usize,
301) -> Result<(), CappedStreamError> {
302    for byte in chunk {
303        if state.1 && *byte == b'\n' {
304            state.0 = 0;
305            state.1 = false;
306            continue;
307        }
308        state.0 = state.0.saturating_add(1);
309        if state.0 > max_bytes {
310            return Err(CappedStreamError::TooLarge {
311                event_bytes: state.0,
312                max_bytes,
313            });
314        }
315        state.1 = *byte == b'\n';
316    }
317    Ok(())
318}
319
320#[derive(Debug)]
321enum CappedStreamError {
322    Reqwest(reqwest::Error),
323    TooLarge {
324        event_bytes: usize,
325        max_bytes: usize,
326    },
327}
328
329impl std::fmt::Display for CappedStreamError {
330    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
331        match self {
332            Self::Reqwest(error) => write!(f, "upstream stream error: {error}"),
333            Self::TooLarge {
334                event_bytes,
335                max_bytes,
336            } => write!(
337                f,
338                "response_too_large: single SSE event reached {event_bytes} bytes, max {max_bytes}"
339            ),
340        }
341    }
342}
343
344impl std::error::Error for CappedStreamError {}
345
346fn ensure_stream_content_type(
347    response: &reqwest::Response,
348) -> Result<(), StreamableHttpError<reqwest::Error>> {
349    match response.headers().get(reqwest::header::CONTENT_TYPE) {
350        Some(value) => {
351            let raw = value.as_bytes();
352            if raw.starts_with(EVENT_STREAM_MIME_TYPE.as_bytes())
353                || raw.starts_with(JSON_MIME_TYPE.as_bytes())
354            {
355                Ok(())
356            } else {
357                Err(StreamableHttpError::UnexpectedContentType(Some(
358                    String::from_utf8_lossy(raw).to_string(),
359                )))
360            }
361        }
362        None => Err(StreamableHttpError::UnexpectedContentType(None)),
363    }
364}
365
366fn content_type(response: &reqwest::Response) -> Option<String> {
367    response
368        .headers()
369        .get(reqwest::header::CONTENT_TYPE)
370        .map(|value| String::from_utf8_lossy(value.as_bytes()).to_string())
371}
372
373fn is_empty_ok(message: &ClientJsonRpcMessage) -> bool {
374    matches!(
375        message,
376        ClientJsonRpcMessage::Notification(_)
377            | ClientJsonRpcMessage::Response(_)
378            | ClientJsonRpcMessage::Error(_)
379    )
380}
381
382fn parse_json_rpc_error(body: &str) -> Option<ServerJsonRpcMessage> {
383    match serde_json::from_str::<ServerJsonRpcMessage>(body) {
384        Ok(message @ JsonRpcMessage::Error(_)) => Some(message),
385        _ => None,
386    }
387}
388
389fn extract_scope(header: &str) -> Option<String> {
390    let lower = header.to_ascii_lowercase();
391    let start = lower.find("scope=")? + "scope=".len();
392    let value = &header[start..];
393    if let Some(quoted) = value.strip_prefix('"') {
394        return quoted.split('"').next().map(ToOwned::to_owned);
395    }
396    let end = value
397        .find(|ch: char| ch == ',' || ch == ';' || ch.is_whitespace())
398        .unwrap_or(value.len());
399    (end > 0).then(|| value[..end].to_owned())
400}
401
402fn too_large(message: String) -> StreamableHttpError<reqwest::Error> {
403    StreamableHttpError::UnexpectedServerResponse(Cow::Owned(message))
404}
405
406#[cfg(test)]
407#[path = "http_body_cap_tests.rs"]
408mod tests;
soma_mcp_client/upstream/http_body_cap.rs

soma_mcp_client/upstream/
http_body_cap.rs
