| 1 | use std::collections::VecDeque; |
| 2 | |
| 3 | use anyhow::{Context, Result}; |
| 4 | use reqwest::StatusCode; |
| 5 | use reqwest::header::CONTENT_TYPE; |
| 6 | |
| 7 | use super::headers::{apply_safe_custom_headers, with_default_mcp_http_headers}; |
| 8 | use super::{ |
| 9 | ERROR_BODY_PREVIEW_BYTES, McpHttpAuth, bounded_body_excerpt, mask_url_secrets, |
| 10 | parse_sse_message_data, |
| 11 | }; |
| 12 | |
| 13 | pub(super) struct StreamableHttpTransport { |
| 14 | pub(super) client: reqwest::Client, |
| 15 | pub(super) url: String, |
| 16 | /// Request-time auth and custom header resolver for outbound POSTs. |
| 17 | pub(super) auth: McpHttpAuth, |
| 18 | pending_messages: VecDeque<Vec<u8>>, |
| 19 | /// Per-spec MCP session identifier returned by the server in the |
| 20 | /// first response (typically the `initialize` response). Attached |
| 21 | /// as the `Mcp-Session-Id` header on every subsequent outbound |
| 22 | /// request so the server can correlate messages within the same |
| 23 | /// session. |
| 24 | pub(super) session_id: Option<String>, |
| 25 | } |
| 26 | |
| 27 | #[derive(Debug)] |
| 28 | pub(super) enum StreamableSendError { |
| 29 | Incompatible(String), |
| 30 | StaleSession(String), |
| 31 | Other(anyhow::Error), |
| 32 | } |
| 33 | |
| 34 | impl StreamableHttpTransport { |
| 35 | pub(super) fn new(client: reqwest::Client, url: String, auth: McpHttpAuth) -> Self { |
| 36 | Self { |
| 37 | client, |
| 38 | url, |
| 39 | auth, |
| 40 | pending_messages: VecDeque::new(), |
| 41 | session_id: None, |
| 42 | } |
| 43 | } |
| 44 | |
| 45 | pub(super) async fn send( |
| 46 | &mut self, |
| 47 | msg: Vec<u8>, |
| 48 | ) -> std::result::Result<(), StreamableSendError> { |
| 49 | // Apply user-configured custom headers after protocol framing so |
| 50 | // reserved Accept / Content-Type overrides can be filtered out. |
| 51 | let headers = self |
| 52 | .auth |
| 53 | .resolved_headers() |
| 54 | .await |
| 55 | .map_err(StreamableSendError::Other)?; |
| 56 | let mut request = apply_safe_custom_headers( |
| 57 | with_default_mcp_http_headers(self.client.post(&self.url), true), |
| 58 | &headers, |
| 59 | ); |
| 60 | // Attach any previously captured session ID per the Streamable |
| 61 | // HTTP spec so the server can correlate this request to the |
| 62 | // existing session. |
| 63 | if let Some(ref sid) = self.session_id { |
| 64 | request = request.header("Mcp-Session-Id", sid.as_str()); |
| 65 | } |
| 66 | let response = request |
| 67 | .body(msg) |
| 68 | .send() |
| 69 | .await |
| 70 | .map_err(|err| StreamableSendError::Other(err.into()))?; |
| 71 | |
| 72 | let status = response.status(); |
| 73 | |
| 74 | // Capture session ID from any response (2xx, 202, 4xx, ...). The |
| 75 | // server may return it on the `initialize` response or on a |
| 76 | // best-effort GET preflight below. |
| 77 | if let Some(sid) = response |
| 78 | .headers() |
| 79 | .get("Mcp-Session-Id") |
| 80 | .and_then(|v| v.to_str().ok()) |
| 81 | && self.session_id.as_deref() != Some(sid) |
| 82 | { |
| 83 | let session_ref = crate::utils::redacted_identifier_for_log(sid); |
| 84 | tracing::debug!(target: "mcp", session = %session_ref, "captured MCP session ID"); |
| 85 | self.session_id = Some(sid.to_string()); |
| 86 | } |
| 87 | if status == StatusCode::ACCEPTED || status == StatusCode::NO_CONTENT { |
| 88 | return Ok(()); |
| 89 | } |
| 90 | |
| 91 | if !status.is_success() { |
| 92 | let body_excerpt = bounded_body_excerpt(response, ERROR_BODY_PREVIEW_BYTES).await; |
| 93 | let stale_session = self.session_id.is_some() |
| 94 | && is_streamable_http_stale_session_status(status, &body_excerpt); |
| 95 | let body_excerpt = self.auth.server_error_preview(&body_excerpt); |
| 96 | if stale_session { |
| 97 | return Err(StreamableSendError::StaleSession(format!( |
| 98 | "status={status} body={body_excerpt}" |
| 99 | ))); |
| 100 | } |
| 101 | if is_streamable_http_incompatible_status(status) { |
| 102 | return Err(StreamableSendError::Incompatible(format!( |
| 103 | "status={status} body={body_excerpt}" |
| 104 | ))); |
| 105 | } |
| 106 | return Err(StreamableSendError::Other(anyhow::anyhow!( |
| 107 | "MCP Streamable HTTP rejected (transport=http url={} status={}): {}", |
| 108 | mask_url_secrets(&self.url), |
| 109 | status, |
| 110 | body_excerpt, |
| 111 | ))); |
| 112 | } |
| 113 | |
| 114 | let content_type = response |
| 115 | .headers() |
| 116 | .get(CONTENT_TYPE) |
| 117 | .and_then(|value| value.to_str().ok()) |
| 118 | .map(str::to_string); |
| 119 | // Reject an over-large declared body before reading anything (fast |
| 120 | // path), then bound the read itself so chunked / length-less |
| 121 | // responses cannot OOM us either — Content-Length alone does not |
| 122 | // protect against a server that streams without declaring a length. |
| 123 | if let Some(len) = response.content_length() |
| 124 | && len > super::MAX_MCP_RESPONSE_BYTES as u64 |
| 125 | { |
| 126 | return Err(StreamableSendError::Other(anyhow::anyhow!( |
| 127 | "MCP response Content-Length {len} exceeds {} bytes — aborting", |
| 128 | super::MAX_MCP_RESPONSE_BYTES |
| 129 | ))); |
| 130 | } |
| 131 | let body = read_body_capped(response, super::MAX_MCP_RESPONSE_BYTES) |
| 132 | .await |
| 133 | .map_err(StreamableSendError::Other)?; |
| 134 | self.store_response_body(content_type.as_deref(), &body) |
| 135 | .map_err(StreamableSendError::Other) |
| 136 | } |
| 137 | |
| 138 | pub(super) async fn recv(&mut self) -> Result<Vec<u8>> { |
| 139 | self.pending_messages |
| 140 | .pop_front() |
| 141 | .context("MCP Streamable HTTP response queue is empty") |
| 142 | } |
| 143 | |
| 144 | fn store_response_body(&mut self, content_type: Option<&str>, body: &str) -> Result<()> { |
| 145 | if body.trim().is_empty() { |
| 146 | return Ok(()); |
| 147 | } |
| 148 | |
| 149 | let is_event_stream = content_type |
| 150 | .map(|value| value.to_ascii_lowercase().contains("text/event-stream")) |
| 151 | .unwrap_or(false) |
| 152 | || body.trim_start().starts_with("event:") |
| 153 | || body.trim_start().starts_with("data:"); |
| 154 | |
| 155 | if is_event_stream { |
| 156 | for msg in parse_sse_message_data(body) { |
| 157 | self.pending_messages.push_back(msg); |
| 158 | } |
| 159 | return Ok(()); |
| 160 | } |
| 161 | |
| 162 | self.pending_messages.push_back(body.as_bytes().to_vec()); |
| 163 | Ok(()) |
| 164 | } |
| 165 | } |
| 166 | |
| 167 | /// Read a response body through the byte stream, failing as soon as it |
| 168 | /// exceeds `max_bytes`. This bounds chunked and missing-Content-Length |
| 169 | /// responses exactly like declared ones (the declared-length fast path in |
| 170 | /// `send` only covers servers honest enough to announce their size). |
| 171 | /// MCP bodies are JSON or SSE, so lossy UTF-8 matches `.text()` behavior. |
| 172 | pub(super) async fn read_body_capped( |
| 173 | response: reqwest::Response, |
| 174 | max_bytes: usize, |
| 175 | ) -> Result<String> { |
| 176 | use futures_util::StreamExt; |
| 177 | |
| 178 | let mut stream = response.bytes_stream(); |
| 179 | let mut buf: Vec<u8> = Vec::new(); |
| 180 | while let Some(chunk) = stream.next().await { |
| 181 | let chunk = chunk.context("failed to read MCP response body")?; |
| 182 | if buf.len().saturating_add(chunk.len()) > max_bytes { |
| 183 | anyhow::bail!("MCP response body exceeds {max_bytes} bytes — aborting"); |
| 184 | } |
| 185 | buf.extend_from_slice(&chunk); |
| 186 | } |
| 187 | Ok(String::from_utf8_lossy(&buf).into_owned()) |
| 188 | } |
| 189 | |
| 190 | fn is_streamable_http_incompatible_status(status: StatusCode) -> bool { |
| 191 | matches!( |
| 192 | status, |
| 193 | StatusCode::NOT_FOUND |
| 194 | | StatusCode::METHOD_NOT_ALLOWED |
| 195 | | StatusCode::NOT_ACCEPTABLE |
| 196 | | StatusCode::UNSUPPORTED_MEDIA_TYPE |
| 197 | | StatusCode::NOT_IMPLEMENTED |
| 198 | ) |
| 199 | } |
| 200 | |
| 201 | fn is_streamable_http_stale_session_status(status: StatusCode, body_excerpt: &str) -> bool { |
| 202 | if status == StatusCode::NOT_FOUND { |
| 203 | return true; |
| 204 | } |
| 205 | if status != StatusCode::BAD_REQUEST && status != StatusCode::UNAUTHORIZED { |
| 206 | return false; |
| 207 | } |
| 208 | let body = body_excerpt.to_ascii_lowercase(); |
| 209 | body.contains("session") && (body.contains("expired") || body.contains("invalid")) |
| 210 | } |
| 211 |