| 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::http_client::McpHttpClient; |
| 9 | use super::wire::{MAX_MCP_RESPONSE_BYTES, parse_sse_message_data}; |
| 10 | use super::{ERROR_BODY_PREVIEW_BYTES, McpHttpAuth, bounded_body_excerpt, mask_url_secrets}; |
| 11 | |
| 12 | pub(super) struct StreamableHttpTransport { |
| 13 | pub(super) client: McpHttpClient, |
| 14 | pub(super) url: String, |
| 15 | /// Request-time auth and custom header resolver for outbound POSTs. |
| 16 | pub(super) auth: McpHttpAuth, |
| 17 | pending_messages: VecDeque<Vec<u8>>, |
| 18 | /// Per-spec MCP session identifier returned by the server in the |
| 19 | /// first response (typically the `initialize` response). Attached |
| 20 | /// as the `Mcp-Session-Id` header on every subsequent outbound |
| 21 | /// request so the server can correlate messages within the same |
| 22 | /// session. |
| 23 | pub(super) session_id: Option<String>, |
| 24 | /// Protocol revision negotiated at `initialize`. Attached as the |
| 25 | /// `MCP-Protocol-Version` header on every subsequent outbound request |
| 26 | /// per the Streamable HTTP spec (absent means the server assumes |
| 27 | /// the 2025-03-26 default, so the negotiated value is always sent). |
| 28 | protocol_version: Option<String>, |
| 29 | } |
| 30 | |
| 31 | #[derive(Debug)] |
| 32 | pub(super) enum StreamableSendError { |
| 33 | Incompatible(String), |
| 34 | StaleSession(String), |
| 35 | Other(anyhow::Error), |
| 36 | } |
| 37 | |
| 38 | impl StreamableHttpTransport { |
| 39 | pub(super) fn new(client: McpHttpClient, url: String, auth: McpHttpAuth) -> Self { |
| 40 | Self { |
| 41 | client, |
| 42 | url, |
| 43 | auth, |
| 44 | pending_messages: VecDeque::new(), |
| 45 | session_id: None, |
| 46 | protocol_version: None, |
| 47 | } |
| 48 | } |
| 49 | |
| 50 | pub(super) fn set_protocol_version(&mut self, version: &str) { |
| 51 | self.protocol_version = Some(version.to_string()); |
| 52 | } |
| 53 | |
| 54 | pub(super) async fn send( |
| 55 | &mut self, |
| 56 | msg: Vec<u8>, |
| 57 | ) -> std::result::Result<(), StreamableSendError> { |
| 58 | // Reactive OAuth recovery (T4): a 401/403 may mean the server no |
| 59 | // longer accepts a token that the local expiry clock still trusts. |
| 60 | // Retry once after a forced refresh; then surface a login hint instead |
| 61 | // of a raw rejection that reads like a broken server. |
| 62 | let mut retried = false; |
| 63 | loop { |
| 64 | // Apply user-configured custom headers after protocol framing so |
| 65 | // reserved Accept / Content-Type overrides can be filtered out. |
| 66 | let headers = self |
| 67 | .auth |
| 68 | .resolved_headers() |
| 69 | .await |
| 70 | .map_err(StreamableSendError::Other)?; |
| 71 | let mut request = apply_safe_custom_headers( |
| 72 | with_default_mcp_http_headers(self.client.post(&self.url), true), |
| 73 | &headers, |
| 74 | ); |
| 75 | // Attach any previously captured session ID per the Streamable |
| 76 | // HTTP spec so the server can correlate this request to the |
| 77 | // existing session. |
| 78 | if let Some(ref sid) = self.session_id { |
| 79 | request = request.header("Mcp-Session-Id", sid.as_str()); |
| 80 | } |
| 81 | // Per the Streamable HTTP spec, subsequent requests carry the |
| 82 | // negotiated revision; absent means the server assumes 2025-03-26. |
| 83 | if let Some(ref version) = self.protocol_version { |
| 84 | request = request.header("MCP-Protocol-Version", version.as_str()); |
| 85 | } |
| 86 | let response = self |
| 87 | .client |
| 88 | .send(request.body(msg.clone())) |
| 89 | .await |
| 90 | .map_err(StreamableSendError::Other)?; |
| 91 | |
| 92 | let status = response.status(); |
| 93 | |
| 94 | // Capture session ID from any response (2xx, 202, 4xx, ...). The |
| 95 | // server may return it on the `initialize` response or on a |
| 96 | // best-effort GET preflight below. |
| 97 | if let Some(sid) = response |
| 98 | .headers() |
| 99 | .get("Mcp-Session-Id") |
| 100 | .and_then(|v| v.to_str().ok()) |
| 101 | && self.session_id.as_deref() != Some(sid) |
| 102 | { |
| 103 | let session_ref = crate::utils::redacted_identifier_for_log(sid); |
| 104 | tracing::debug!(target: "mcp", session = %session_ref, "captured MCP session ID"); |
| 105 | self.session_id = Some(sid.to_string()); |
| 106 | } |
| 107 | if status == StatusCode::ACCEPTED || status == StatusCode::NO_CONTENT { |
| 108 | return Ok(()); |
| 109 | } |
| 110 | |
| 111 | if status == StatusCode::UNAUTHORIZED || status == StatusCode::FORBIDDEN { |
| 112 | if !retried && let Some(oauth) = self.auth.oauth.as_ref() { |
| 113 | match oauth.force_refresh().await { |
| 114 | Ok(()) => { |
| 115 | retried = true; |
| 116 | continue; |
| 117 | } |
| 118 | Err(refresh_error) => { |
| 119 | return Err(StreamableSendError::Other(anyhow::anyhow!( |
| 120 | "MCP server {} rejected the request with {status} and refreshing the OAuth session failed: {refresh_error:#}. {hint}", |
| 121 | mask_url_secrets(&self.url), |
| 122 | hint = oauth_refresh_failed_hint(), |
| 123 | ))); |
| 124 | } |
| 125 | } |
| 126 | } |
| 127 | let hint = unauthorized_session_hint(self.auth.oauth_configured); |
| 128 | return Err(StreamableSendError::Other(anyhow::anyhow!( |
| 129 | "MCP server {} rejected the request with {status}; the session is no longer accepted. {hint}", |
| 130 | mask_url_secrets(&self.url), |
| 131 | ))); |
| 132 | } |
| 133 | |
| 134 | if !status.is_success() { |
| 135 | let body_excerpt = bounded_body_excerpt(response, ERROR_BODY_PREVIEW_BYTES).await; |
| 136 | let stale_session = self.session_id.is_some() |
| 137 | && is_streamable_http_stale_session_status(status, &body_excerpt); |
| 138 | let body_excerpt = self.auth.server_error_preview(&body_excerpt); |
| 139 | if stale_session { |
| 140 | return Err(StreamableSendError::StaleSession(format!( |
| 141 | "status={status} body={body_excerpt}" |
| 142 | ))); |
| 143 | } |
| 144 | if is_streamable_http_incompatible_status(status) { |
| 145 | return Err(StreamableSendError::Incompatible(format!( |
| 146 | "status={status} body={body_excerpt}" |
| 147 | ))); |
| 148 | } |
| 149 | return Err(StreamableSendError::Other(anyhow::anyhow!( |
| 150 | "MCP Streamable HTTP rejected (transport=http url={} status={}): {}", |
| 151 | mask_url_secrets(&self.url), |
| 152 | status, |
| 153 | body_excerpt, |
| 154 | ))); |
| 155 | } |
| 156 | |
| 157 | let content_type = response |
| 158 | .headers() |
| 159 | .get(CONTENT_TYPE) |
| 160 | .and_then(|value| value.to_str().ok()) |
| 161 | .map(str::to_string); |
| 162 | // Reject an over-large declared body before reading anything (fast |
| 163 | // path), then bound the read itself so chunked / length-less |
| 164 | // responses cannot OOM us either — Content-Length alone does not |
| 165 | // protect against a server that streams without declaring a length. |
| 166 | if let Some(len) = response.content_length() |
| 167 | && len > MAX_MCP_RESPONSE_BYTES as u64 |
| 168 | { |
| 169 | return Err(StreamableSendError::Other(anyhow::anyhow!( |
| 170 | "MCP response Content-Length {len} exceeds {} bytes — aborting", |
| 171 | MAX_MCP_RESPONSE_BYTES |
| 172 | ))); |
| 173 | } |
| 174 | let body = read_body_capped(response, MAX_MCP_RESPONSE_BYTES) |
| 175 | .await |
| 176 | .map_err(StreamableSendError::Other)?; |
| 177 | return self |
| 178 | .store_response_body(content_type.as_deref(), &body) |
| 179 | .map_err(StreamableSendError::Other); |
| 180 | } |
| 181 | } |
| 182 | |
| 183 | pub(super) async fn recv(&mut self) -> Result<Vec<u8>> { |
| 184 | self.pending_messages |
| 185 | .pop_front() |
| 186 | .context("MCP Streamable HTTP response queue is empty") |
| 187 | } |
| 188 | |
| 189 | fn store_response_body(&mut self, content_type: Option<&str>, body: &str) -> Result<()> { |
| 190 | if body.trim().is_empty() { |
| 191 | return Ok(()); |
| 192 | } |
| 193 | |
| 194 | let is_event_stream = content_type |
| 195 | .map(|value| value.to_ascii_lowercase().contains("text/event-stream")) |
| 196 | .unwrap_or(false) |
| 197 | || body.trim_start().starts_with("event:") |
| 198 | || body.trim_start().starts_with("data:"); |
| 199 | |
| 200 | if is_event_stream { |
| 201 | for msg in parse_sse_message_data(body) { |
| 202 | self.pending_messages.push_back(msg); |
| 203 | } |
| 204 | return Ok(()); |
| 205 | } |
| 206 | |
| 207 | self.pending_messages.push_back(body.as_bytes().to_vec()); |
| 208 | Ok(()) |
| 209 | } |
| 210 | } |
| 211 | |
| 212 | /// Read a response body through the byte stream, failing as soon as it |
| 213 | /// exceeds `max_bytes`. This bounds chunked and missing-Content-Length |
| 214 | /// responses exactly like declared ones (the declared-length fast path in |
| 215 | /// `send` only covers servers honest enough to announce their size). |
| 216 | /// MCP bodies are JSON or SSE, so lossy UTF-8 matches `.text()` behavior. |
| 217 | pub(super) async fn read_body_capped( |
| 218 | response: reqwest::Response, |
| 219 | max_bytes: usize, |
| 220 | ) -> Result<String> { |
| 221 | use futures_util::StreamExt; |
| 222 | |
| 223 | let mut stream = response.bytes_stream(); |
| 224 | let mut buf: Vec<u8> = Vec::new(); |
| 225 | while let Some(chunk) = stream.next().await { |
| 226 | let chunk = chunk.context("failed to read MCP response body")?; |
| 227 | if buf.len().saturating_add(chunk.len()) > max_bytes { |
| 228 | anyhow::bail!("MCP response body exceeds {max_bytes} bytes — aborting"); |
| 229 | } |
| 230 | buf.extend_from_slice(&chunk); |
| 231 | } |
| 232 | Ok(String::from_utf8_lossy(&buf).into_owned()) |
| 233 | } |
| 234 | |
| 235 | /// TUI recovery for a rejected OAuth session. Settings recovery and the |
| 236 | /// Streamable HTTP path share the oauth helpers so `/mcp login <name>` stays |
| 237 | /// the only advertised command; `/mcp auth` is not a command. |
| 238 | fn oauth_refresh_failed_hint() -> &'static str { |
| 239 | super::oauth::tui_reauth_refresh_failed_hint() |
| 240 | } |
| 241 | |
| 242 | /// TUI recovery for a rejected OAuth session. `oauth_configured` is the |
| 243 | /// server's configured auth path ([`McpHttpAuth::oauth_configured`]), not the |
| 244 | /// presence of a cached token, so a first-run OAuth server — a 401 with |
| 245 | /// nothing stored yet — is still pointed at `/mcp login <name>` rather than at |
| 246 | /// a bearer token it never had (#6030). Servers where a bearer credential is |
| 247 | /// genuinely configured (or that are plugin-contributed, where OAuth login is |
| 248 | /// disabled) keep the bearer-token copy. |
| 249 | fn unauthorized_session_hint(oauth_configured: bool) -> &'static str { |
| 250 | if oauth_configured { |
| 251 | super::oauth::tui_reauth_hint() |
| 252 | } else { |
| 253 | "Check the configured bearer token (or its environment variable)." |
| 254 | } |
| 255 | } |
| 256 | |
| 257 | fn is_streamable_http_incompatible_status(status: StatusCode) -> bool { |
| 258 | matches!( |
| 259 | status, |
| 260 | StatusCode::NOT_FOUND |
| 261 | | StatusCode::METHOD_NOT_ALLOWED |
| 262 | | StatusCode::NOT_ACCEPTABLE |
| 263 | | StatusCode::UNSUPPORTED_MEDIA_TYPE |
| 264 | | StatusCode::NOT_IMPLEMENTED |
| 265 | ) |
| 266 | } |
| 267 | |
| 268 | fn is_streamable_http_stale_session_status(status: StatusCode, body_excerpt: &str) -> bool { |
| 269 | if status == StatusCode::NOT_FOUND { |
| 270 | return true; |
| 271 | } |
| 272 | if status != StatusCode::BAD_REQUEST && status != StatusCode::UNAUTHORIZED { |
| 273 | return false; |
| 274 | } |
| 275 | let body = body_excerpt.to_ascii_lowercase(); |
| 276 | body.contains("session") && (body.contains("expired") || body.contains("invalid")) |
| 277 | } |
| 278 | |
| 279 | #[cfg(test)] |
| 280 | mod tests { |
| 281 | use super::{McpHttpAuth, oauth_refresh_failed_hint, unauthorized_session_hint}; |
| 282 | use crate::mcp::McpServerConfig; |
| 283 | |
| 284 | fn server_config(json: serde_json::Value) -> McpServerConfig { |
| 285 | serde_json::from_value(json).expect("MCP server config fixture") |
| 286 | } |
| 287 | |
| 288 | #[test] |
| 289 | fn oauth_configured_server_without_a_cached_token_names_login() { |
| 290 | // The OAuth fields are optional in MCP config, so a URL-based server |
| 291 | // with no manual bearer configuration is OAuth's to claim — including |
| 292 | // before the first login, when there is no runtime to observe. |
| 293 | let auth = McpHttpAuth::from_config( |
| 294 | "remote", |
| 295 | &server_config(serde_json::json!({ "url": "https://example.invalid/mcp" })), |
| 296 | None, |
| 297 | ); |
| 298 | assert!(auth.oauth.is_none(), "precondition: no cached credential"); |
| 299 | assert!(auth.oauth_configured, "a URL server is OAuth-servable"); |
| 300 | assert!( |
| 301 | unauthorized_session_hint(auth.oauth_configured).contains("/mcp login <name>"), |
| 302 | "a first-run OAuth 401 must name the login command, not a bearer token" |
| 303 | ); |
| 304 | |
| 305 | // A server whose bearer token is genuinely expected keeps that copy. |
| 306 | let bearer = McpHttpAuth::from_config( |
| 307 | "remote", |
| 308 | &server_config(serde_json::json!({ |
| 309 | "url": "https://example.invalid/mcp", |
| 310 | "bearer_token_env_var": "EXAMPLE_MCP_TOKEN", |
| 311 | })), |
| 312 | None, |
| 313 | ); |
| 314 | assert!(!bearer.oauth_configured); |
| 315 | assert!(unauthorized_session_hint(bearer.oauth_configured).contains("bearer token")); |
| 316 | } |
| 317 | |
| 318 | #[test] |
| 319 | fn unauthorized_oauth_hints_name_the_login_command() { |
| 320 | for hint in [oauth_refresh_failed_hint(), unauthorized_session_hint(true)] { |
| 321 | assert!( |
| 322 | hint.contains("/mcp login <name>"), |
| 323 | "OAuth recovery must name the implemented command" |
| 324 | ); |
| 325 | assert!( |
| 326 | !hint.contains("/mcp auth"), |
| 327 | "OAuth recovery must not advertise a missing /mcp auth command" |
| 328 | ); |
| 329 | } |
| 330 | assert!( |
| 331 | !unauthorized_session_hint(false).contains("/mcp"), |
| 332 | "bearer-token recovery should not send the user to OAuth login" |
| 333 | ); |
| 334 | } |
| 335 | } |
| 336 |