rmcp_soddygo/transport/
streamable_http_client.rs

1use std::{
2    borrow::Cow,
3    collections::{HashMap, HashSet},
4    sync::Arc,
5    time::Duration,
6};
7
8use futures::{Stream, StreamExt, future::BoxFuture, stream::BoxStream};
9use http::{HeaderName, HeaderValue};
10pub use sse_stream::Error as SseError;
11use sse_stream::Sse;
12use thiserror::Error;
13use tokio_util::sync::CancellationToken;
14use tracing::debug;
15
16use super::common::client_side_sse::{
17    DEFAULT_MAX_SSE_EVENT_SIZE, ExponentialBackoff, SseRetryPolicy, SseStreamReconnect,
18};
19use crate::{
20    RoleClient,
21    model::{
22        ClientJsonRpcMessage, ClientNotification, ClientRequest, ErrorData, GetMeta,
23        InitializedNotification, JsonObject, ProtocolVersion, RequestId, ServerJsonRpcMessage,
24        ServerResult,
25    },
26    transport::{
27        common::{client_side_sse::SseAutoReconnectStream, mcp_headers},
28        worker::{Worker, WorkerQuitReason, WorkerSendRequest, WorkerTransport},
29    },
30};
31
32type BoxedSseStream = BoxStream<'static, Result<Sse, SseError>>;
33const SESSION_CLEANUP_TIMEOUT: Duration = Duration::from_secs(5);
34
35fn build_request_headers(
36    base: &HashMap<HeaderName, HeaderValue>,
37    message: &ClientJsonRpcMessage,
38    tool_cache: &HashMap<String, Arc<JsonObject>>,
39    version: &ProtocolVersion,
40) -> HashMap<HeaderName, HeaderValue> {
41    use serde_json::Value;
42
43    let mut headers = base.clone();
44    if *version >= ProtocolVersion::STANDARD_HEADERS {
45        if let Ok(value) = serde_json::to_value(message) {
46            let schema = value
47                .get("method")
48                .and_then(Value::as_str)
49                .filter(|method| *method == "tools/call")
50                .and_then(|_| value.get("params"))
51                .and_then(|params| params.get("name"))
52                .and_then(Value::as_str)
53                .and_then(|name| tool_cache.get(name))
54                .map(Arc::as_ref);
55            for (name, val) in mcp_headers::standard_request_headers(&value, schema) {
56                headers.insert(name, val);
57            }
58        }
59    }
60    headers
61}
62
63fn request_version_headers(
64    base: &HashMap<HeaderName, HeaderValue>,
65    message: &ClientJsonRpcMessage,
66    fallback: &ProtocolVersion,
67    tool_cache: &HashMap<String, Arc<JsonObject>>,
68) -> (ProtocolVersion, HashMap<HeaderName, HeaderValue>) {
69    let version = match message {
70        ClientJsonRpcMessage::Request(request) => request
71            .request
72            .get_meta()
73            .protocol_version()
74            .unwrap_or_else(|| fallback.clone()),
75        _ => fallback.clone(),
76    };
77    let mut headers = build_request_headers(base, message, tool_cache, &version);
78    if let Ok(value) = HeaderValue::from_str(version.as_str()) {
79        headers.insert(HeaderName::from_static("mcp-protocol-version"), value);
80    }
81    (version, headers)
82}
83
84fn cache_tools_from_response(
85    cache: &mut HashMap<String, Arc<JsonObject>>,
86    message: &ServerJsonRpcMessage,
87) {
88    if let ServerJsonRpcMessage::Response(response) = message {
89        if let ServerResult::ListToolsResult(list) = &response.result {
90            for tool in &list.tools {
91                if let Err(reason) =
92                    mcp_headers::validate_param_header_annotations(&tool.input_schema)
93                {
94                    tracing::warn!(tool = %tool.name, "ignoring x-mcp-header annotations: {reason}");
95                    continue;
96                }
97                cache.insert(tool.name.to_string(), tool.input_schema.clone());
98            }
99        }
100    }
101}
102
103fn negotiate_version_headers(
104    init_response: &ServerJsonRpcMessage,
105    base: HashMap<HeaderName, HeaderValue>,
106) -> (ProtocolVersion, HashMap<HeaderName, HeaderValue>) {
107    let mut version = ProtocolVersion::default();
108    let mut headers = base;
109    if let ServerJsonRpcMessage::Response(response) = init_response {
110        if let ServerResult::InitializeResult(init_result) = &response.result {
111            version = init_result.protocol_version.clone();
112            // HeaderName::from_static requires lowercase
113            if let Ok(hv) = HeaderValue::from_str(init_result.protocol_version.as_str()) {
114                headers.insert(HeaderName::from_static("mcp-protocol-version"), hv);
115            }
116        }
117    }
118    (version, headers)
119}
120
121#[derive(Debug)]
122#[non_exhaustive]
123pub struct AuthRequiredError {
124    pub www_authenticate_header: String,
125}
126
127impl AuthRequiredError {
128    /// Create a new `AuthRequiredError` instance.
129    pub fn new(www_authenticate_header: String) -> Self {
130        Self {
131            www_authenticate_header,
132        }
133    }
134}
135
136#[derive(Debug)]
137#[non_exhaustive]
138pub struct InsufficientScopeError {
139    pub www_authenticate_header: String,
140    pub required_scope: Option<String>,
141}
142
143impl InsufficientScopeError {
144    /// Create a new `InsufficientScopeError` instance.
145    pub fn new(www_authenticate_header: String, required_scope: Option<String>) -> Self {
146        Self {
147            www_authenticate_header,
148            required_scope,
149        }
150    }
151
152    /// check if scope upgrade is possible (i.e., we know what scope is required)
153    pub fn can_upgrade(&self) -> bool {
154        self.required_scope.is_some()
155    }
156
157    /// get the required scope for upgrade
158    pub fn get_required_scope(&self) -> Option<&str> {
159        self.required_scope.as_deref()
160    }
161}
162
163#[derive(Error, Debug)]
164#[non_exhaustive]
165pub enum StreamableHttpError<E: std::error::Error + Send + Sync + 'static> {
166    #[error("SSE error: {0}")]
167    Sse(#[from] SseError),
168    #[error("Io error: {0}")]
169    Io(#[from] std::io::Error),
170    #[error("Client error: {0}")]
171    Client(E),
172    #[error("unexpected end of stream")]
173    UnexpectedEndOfStream,
174    #[error("unexpected server response: {0}")]
175    UnexpectedServerResponse(Cow<'static, str>),
176    #[error("Unexpected content type: {0:?}")]
177    UnexpectedContentType(Option<String>),
178    #[error("Server does not support SSE")]
179    ServerDoesNotSupportSse,
180    #[error("Server does not support delete session")]
181    ServerDoesNotSupportDeleteSession,
182    #[error("Tokio join error: {0}")]
183    TokioJoinError(#[from] tokio::task::JoinError),
184    #[error("Deserialize error: {0}")]
185    Deserialize(#[from] serde_json::Error),
186    #[error("Transport channel closed")]
187    TransportChannelClosed,
188    #[error("Missing session id in HTTP response")]
189    MissingSessionIdInResponse,
190    #[cfg(feature = "auth")]
191    #[error("Auth error: {0}")]
192    Auth(#[from] crate::transport::auth::AuthError),
193    #[error("Auth required")]
194    AuthRequired(AuthRequiredError),
195    #[error("Insufficient scope")]
196    InsufficientScope(InsufficientScopeError),
197    #[error("Header name '{0}' is reserved and conflicts with default headers")]
198    ReservedHeaderConflict(String),
199    #[error("Session expired (HTTP 404)")]
200    SessionExpired,
201}
202
203#[derive(Debug, Clone, Error)]
204#[non_exhaustive]
205pub enum StreamableHttpProtocolError {
206    #[error("Missing session id in response")]
207    MissingSessionIdInResponse,
208}
209
210#[expect(
211    clippy::large_enum_variant,
212    reason = "boxing the streaming response would add an allocation to the common response path"
213)]
214#[non_exhaustive]
215pub enum StreamableHttpPostResponse {
216    Accepted,
217    Json(ServerJsonRpcMessage, Option<String>),
218    Sse(BoxedSseStream, Option<String>),
219}
220
221impl std::fmt::Debug for StreamableHttpPostResponse {
222    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
223        match self {
224            Self::Accepted => write!(f, "Accepted"),
225            Self::Json(arg0, arg1) => f.debug_tuple("Json").field(arg0).field(arg1).finish(),
226            Self::Sse(_, arg1) => f.debug_tuple("Sse").field(arg1).finish(),
227        }
228    }
229}
230
231impl StreamableHttpPostResponse {
232    pub async fn expect_initialized<E>(
233        self,
234    ) -> Result<(ServerJsonRpcMessage, Option<String>), StreamableHttpError<E>>
235    where
236        E: std::error::Error + Send + Sync + 'static,
237    {
238        match self {
239            Self::Json(message, session_id) => Ok((message, session_id)),
240            Self::Sse(mut stream, session_id) => {
241                while let Some(event) = stream.next().await {
242                    let event = event?;
243                    let payload = event.data.unwrap_or_default();
244                    if payload.trim().is_empty() {
245                        continue;
246                    }
247
248                    let message: ServerJsonRpcMessage = serde_json::from_str(&payload)?;
249
250                    if matches!(message, ServerJsonRpcMessage::Response(_)) {
251                        return Ok((message, session_id));
252                    }
253
254                    debug!(
255                        ?message,
256                        "received message before initialize response; continuing to drain stream"
257                    );
258                }
259
260                Err(StreamableHttpError::UnexpectedServerResponse(
261                    "empty sse stream".into(),
262                ))
263            }
264            _ => Err(StreamableHttpError::UnexpectedServerResponse(
265                "expect initialized, accepted".into(),
266            )),
267        }
268    }
269
270    pub fn expect_json<E>(self) -> Result<ServerJsonRpcMessage, StreamableHttpError<E>>
271    where
272        E: std::error::Error + Send + Sync + 'static,
273    {
274        match self {
275            Self::Json(message, ..) => Ok(message),
276            got => Err(StreamableHttpError::UnexpectedServerResponse(
277                format!("expect json, got {got:?}").into(),
278            )),
279        }
280    }
281
282    pub fn expect_accepted_or_json<E>(self) -> Result<(), StreamableHttpError<E>>
283    where
284        E: std::error::Error + Send + Sync + 'static,
285    {
286        match self {
287            Self::Accepted => Ok(()),
288            // Tolerate servers that return 200 with JSON for notifications
289            Self::Json(..) => Ok(()),
290            got => Err(StreamableHttpError::UnexpectedServerResponse(
291                format!("expect accepted or json, got {got:?}").into(),
292            )),
293        }
294    }
295}
296
297/// HTTP backend used by [`StreamableHttpClientTransport`].
298///
299/// Custom implementations that parse SSE responses must override
300/// [`Self::post_message_with_max_sse_event_size`] and
301/// [`Self::get_stream_with_max_sse_event_size`] to enforce the transport's
302/// configured event-size limit.
303pub trait StreamableHttpClient: Clone + Send + 'static {
304    type Error: std::error::Error + Send + Sync + 'static;
305    fn post_message(
306        &self,
307        uri: Arc<str>,
308        message: ClientJsonRpcMessage,
309        session_id: Option<Arc<str>>,
310        auth_header: Option<String>,
311        custom_headers: HashMap<HeaderName, HeaderValue>,
312    ) -> impl Future<Output = Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>>>
313    + Send
314    + '_;
315    /// Send a message while enforcing a maximum raw SSE event size.
316    ///
317    /// `max_sse_event_size` is not a per-request option: it is the
318    /// transport-wide [`StreamableHttpClientTransportConfig::max_sse_event_size`]
319    /// value, passed identically on every call because the limit must be applied
320    /// inside the client (at the raw byte layer, before SSE parsing) rather than
321    /// by the caller.
322    ///
323    /// Custom clients that parse SSE responses should override this method.
324    /// The default implementation delegates to [`Self::post_message`].
325    fn post_message_with_max_sse_event_size(
326        &self,
327        uri: Arc<str>,
328        message: ClientJsonRpcMessage,
329        session_id: Option<Arc<str>>,
330        auth_header: Option<String>,
331        custom_headers: HashMap<HeaderName, HeaderValue>,
332        _max_sse_event_size: usize,
333    ) -> impl Future<Output = Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>>>
334    + Send
335    + '_ {
336        self.post_message(uri, message, session_id, auth_header, custom_headers)
337    }
338    fn delete_session(
339        &self,
340        uri: Arc<str>,
341        session_id: Arc<str>,
342        auth_header: Option<String>,
343        custom_headers: HashMap<HeaderName, HeaderValue>,
344    ) -> impl Future<Output = Result<(), StreamableHttpError<Self::Error>>> + Send + '_;
345    fn get_stream(
346        &self,
347        uri: Arc<str>,
348        session_id: Arc<str>,
349        last_event_id: Option<String>,
350        auth_header: Option<String>,
351        custom_headers: HashMap<HeaderName, HeaderValue>,
352    ) -> impl Future<
353        Output = Result<
354            BoxStream<'static, Result<Sse, SseError>>,
355            StreamableHttpError<Self::Error>,
356        >,
357    > + Send
358    + '_;
359    /// Open an SSE stream while enforcing a maximum raw event size.
360    ///
361    /// `max_sse_event_size` is not a per-request option: it is the
362    /// transport-wide [`StreamableHttpClientTransportConfig::max_sse_event_size`]
363    /// value, passed identically on every call because the limit must be applied
364    /// inside the client (at the raw byte layer, before SSE parsing) rather than
365    /// by the caller.
366    ///
367    /// Custom clients that parse SSE responses should override this method.
368    /// The default implementation delegates to [`Self::get_stream`].
369    fn get_stream_with_max_sse_event_size(
370        &self,
371        uri: Arc<str>,
372        session_id: Arc<str>,
373        last_event_id: Option<String>,
374        auth_header: Option<String>,
375        custom_headers: HashMap<HeaderName, HeaderValue>,
376        _max_sse_event_size: usize,
377    ) -> impl Future<
378        Output = Result<
379            BoxStream<'static, Result<Sse, SseError>>,
380            StreamableHttpError<Self::Error>,
381        >,
382    > + Send
383    + '_ {
384        self.get_stream(uri, session_id, last_event_id, auth_header, custom_headers)
385    }
386}
387
388#[non_exhaustive]
389pub struct RetryConfig {
390    pub max_times: Option<usize>,
391    pub min_duration: Duration,
392}
393
394struct StreamableHttpClientReconnect<C> {
395    pub client: C,
396    pub session_id: Arc<str>,
397    pub uri: Arc<str>,
398    pub auth_header: Option<String>,
399    pub custom_headers: HashMap<HeaderName, HeaderValue>,
400    pub max_sse_event_size: usize,
401}
402
403impl<C: StreamableHttpClient> SseStreamReconnect for StreamableHttpClientReconnect<C> {
404    type Error = StreamableHttpError<C::Error>;
405    type Future = BoxFuture<'static, Result<BoxedSseStream, Self::Error>>;
406    fn retry_connection(&mut self, last_event_id: Option<&str>) -> Self::Future {
407        let client = self.client.clone();
408        let uri = self.uri.clone();
409        let session_id = self.session_id.clone();
410        let auth_header = self.auth_header.clone();
411        let custom_headers = self.custom_headers.clone();
412        let max_sse_event_size = self.max_sse_event_size;
413        let last_event_id = last_event_id.map(|s| s.to_owned());
414        Box::pin(async move {
415            client
416                .get_stream_with_max_sse_event_size(
417                    uri,
418                    session_id,
419                    last_event_id,
420                    auth_header,
421                    custom_headers,
422                    max_sse_event_size,
423                )
424                .await
425        })
426    }
427
428    fn map_fatal_stream_error(&mut self, error: SseError) -> Option<Self::Error> {
429        Some(StreamableHttpError::Sse(error))
430    }
431}
432
433/// Info retained for cleaning up the session when the worker exits.
434struct SessionCleanupInfo<C> {
435    client: C,
436    uri: Arc<str>,
437    session_id: Arc<str>,
438    auth_header: Option<String>,
439    protocol_headers: HashMap<HeaderName, HeaderValue>,
440}
441
442#[derive(Debug, Clone, Default)]
443#[non_exhaustive]
444pub struct StreamableHttpClientWorker<C: StreamableHttpClient> {
445    pub client: C,
446    pub config: StreamableHttpClientTransportConfig,
447}
448
449impl<C: StreamableHttpClient + Default> StreamableHttpClientWorker<C> {
450    pub fn new_simple(url: impl Into<Arc<str>>) -> Self {
451        Self {
452            client: C::default(),
453            config: StreamableHttpClientTransportConfig {
454                uri: url.into(),
455                ..Default::default()
456            },
457        }
458    }
459}
460
461impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
462    pub fn new(client: C, config: StreamableHttpClientTransportConfig) -> Self {
463        Self { client, config }
464    }
465}
466
467impl<C: StreamableHttpClient> StreamableHttpClientWorker<C> {
468    fn client_request_id(message: &ClientJsonRpcMessage) -> Option<RequestId> {
469        match message {
470            ClientJsonRpcMessage::Request(request) => Some(request.id.clone()),
471            _ => None,
472        }
473    }
474
475    fn server_response_id(message: &ServerJsonRpcMessage) -> Option<&RequestId> {
476        match message {
477            ServerJsonRpcMessage::Response(response) => Some(&response.id),
478            ServerJsonRpcMessage::Error(error) => error.id.as_ref(),
479            _ => None,
480        }
481    }
482
483    fn mark_stream_response_pending(
484        pending_stream_response_ids: &mut HashSet<RequestId>,
485        request_id: Option<RequestId>,
486    ) {
487        if let Some(request_id) = request_id {
488            pending_stream_response_ids.insert(request_id);
489        }
490    }
491
492    fn clear_stream_response_pending(
493        pending_stream_response_ids: &mut HashSet<RequestId>,
494        message: &ServerJsonRpcMessage,
495    ) {
496        if let Some(id) = Self::server_response_id(message) {
497            pending_stream_response_ids.remove(id);
498        }
499    }
500
501    async fn drain_queued_stream_messages(
502        sse_worker_rx: &mut tokio::sync::mpsc::Receiver<ServerJsonRpcMessage>,
503        context: &mut super::worker::WorkerContext<Self>,
504        pending_stream_response_ids: &mut HashSet<RequestId>,
505    ) -> Result<(), WorkerQuitReason<StreamableHttpError<C::Error>>> {
506        loop {
507            match sse_worker_rx.try_recv() {
508                Ok(message) => {
509                    Self::clear_stream_response_pending(pending_stream_response_ids, &message);
510                    context.send_to_handler(message).await?;
511                }
512                Err(tokio::sync::mpsc::error::TryRecvError::Empty) => return Ok(()),
513                Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => return Ok(()),
514            }
515        }
516    }
517
518    async fn fail_pending_stream_responses(
519        context: &mut super::worker::WorkerContext<Self>,
520        pending_stream_response_ids: &mut HashSet<RequestId>,
521    ) -> Result<(), WorkerQuitReason<StreamableHttpError<C::Error>>> {
522        if pending_stream_response_ids.is_empty() {
523            return Ok(());
524        }
525
526        let pending_ids = std::mem::take(pending_stream_response_ids);
527        for id in pending_ids {
528            context
529                .send_to_handler(ServerJsonRpcMessage::error(
530                    ErrorData::internal_error(
531                        "streamable HTTP session was re-initialized before the response arrived",
532                        None,
533                    ),
534                    Some(id),
535                ))
536                .await?;
537        }
538        Ok(())
539    }
540
541    /// Convert a raw SSE stream into a JSON-RPC message stream without
542    /// reconnection logic.
543    fn raw_sse_to_jsonrpc(
544        stream: BoxedSseStream,
545    ) -> impl Stream<Item = Result<ServerJsonRpcMessage, StreamableHttpError<C::Error>>> + Send + 'static
546    {
547        stream.filter_map(|event| async {
548            match event {
549                Err(e) => Some(Err(StreamableHttpError::Sse(e))),
550                Ok(sse) => {
551                    let is_message =
552                        matches!(sse.event.as_deref(), None | Some("") | Some("message"));
553                    if !is_message {
554                        return None;
555                    }
556                    let data = sse.data?;
557                    if data.trim().is_empty() {
558                        return None;
559                    }
560                    match serde_json::from_str::<ServerJsonRpcMessage>(&data) {
561                        Ok(msg) => Some(Ok(msg)),
562                        Err(e) => {
563                            tracing::debug!("failed to deserialize server message: {e}");
564                            None
565                        }
566                    }
567                }
568            }
569        })
570    }
571
572    /// Convert an SSE stream into JSON-RPC messages with reconnect semantics.
573    ///
574    /// This is used for request-scoped SSE responses as well as the standalone
575    /// GET stream. A request-scoped stream can close before its response arrives,
576    /// and SEP-1699 requires the client to honor `retry` and resume with
577    /// `Last-Event-ID` in that case.
578    fn reconnecting_sse_to_jsonrpc(
579        stream: BoxedSseStream,
580        client: C,
581        session_id: Arc<str>,
582        uri: Arc<str>,
583        auth_header: Option<String>,
584        custom_headers: HashMap<HeaderName, HeaderValue>,
585        max_sse_event_size: usize,
586        retry_config: Arc<dyn SseRetryPolicy>,
587    ) -> impl Stream<Item = Result<ServerJsonRpcMessage, StreamableHttpError<C::Error>>> + Send + 'static
588    {
589        SseAutoReconnectStream::new(
590            stream,
591            StreamableHttpClientReconnect {
592                client,
593                session_id,
594                uri,
595                auth_header,
596                custom_headers,
597                max_sse_event_size,
598            },
599            retry_config,
600        )
601    }
602
603    /// Convert a POST response SSE stream into JSON-RPC messages.
604    ///
605    /// Stateful sessions can resume via GET when the response stream closes
606    /// before the server sends the matching JSON-RPC response. Stateless
607    /// transports do not have enough state to resume, so they keep the raw
608    /// SSE-to-JSON-RPC mapping.
609    fn response_sse_to_jsonrpc(
610        stream: BoxedSseStream,
611        session_id: Option<Arc<str>>,
612        client: C,
613        uri: Arc<str>,
614        auth_header: Option<String>,
615        custom_headers: HashMap<HeaderName, HeaderValue>,
616        max_sse_event_size: usize,
617        retry_config: Arc<dyn SseRetryPolicy>,
618    ) -> BoxStream<'static, Result<ServerJsonRpcMessage, StreamableHttpError<C::Error>>> {
619        match session_id {
620            Some(session_id) => Self::reconnecting_sse_to_jsonrpc(
621                stream,
622                client,
623                session_id,
624                uri,
625                auth_header,
626                custom_headers,
627                max_sse_event_size,
628                retry_config,
629            )
630            .boxed(),
631            None => Self::raw_sse_to_jsonrpc(stream).boxed(),
632        }
633    }
634
635    async fn execute_sse_stream(
636        sse_stream: impl Stream<Item = Result<ServerJsonRpcMessage, StreamableHttpError<C::Error>>>
637        + Send
638        + 'static,
639        sse_worker_tx: tokio::sync::mpsc::Sender<ServerJsonRpcMessage>,
640        close_on_response: bool,
641        ct: CancellationToken,
642    ) -> Result<(), StreamableHttpError<C::Error>> {
643        let mut sse_stream = std::pin::pin!(sse_stream);
644        loop {
645            let message = tokio::select! {
646                event = sse_stream.next() => {
647                    event
648                }
649                _ = ct.cancelled() => {
650                    tracing::debug!("cancelled");
651                    break;
652                }
653            };
654            let Some(message) = message.transpose()? else {
655                break;
656            };
657            let is_response = matches!(
658                message,
659                ServerJsonRpcMessage::Response(_) | ServerJsonRpcMessage::Error(_)
660            );
661            let yield_result = sse_worker_tx.send(message).await;
662            if yield_result.is_err() {
663                tracing::trace!("streamable http transport worker dropped, exiting");
664                break;
665            }
666            if close_on_response && is_response {
667                tracing::debug!("got response, draining sse stream for connection reuse");
668                // Consume the remaining stream so the HTTP/1.1 connection
669                // returns to the pool cleanly.
670                let _ = tokio::time::timeout(std::time::Duration::from_millis(50), async {
671                    while sse_stream.next().await.is_some() {}
672                })
673                .await;
674                break;
675            }
676        }
677        Ok(())
678    }
679
680    fn spawn_common_stream(
681        streams: &mut tokio::task::JoinSet<Result<(), StreamableHttpError<C::Error>>>,
682        client: C,
683        session_id: Arc<str>,
684        config: &StreamableHttpClientTransportConfig,
685        protocol_headers: HashMap<HeaderName, HeaderValue>,
686        sse_worker_tx: tokio::sync::mpsc::Sender<ServerJsonRpcMessage>,
687        transport_task_ct: CancellationToken,
688    ) {
689        let uri = config.uri.clone();
690        let auth_header = config.auth_header.clone();
691        let retry_config = config.retry_config.clone();
692        let reconnect_uri = config.uri.clone();
693        let reconnect_auth_header = config.auth_header.clone();
694        let max_sse_event_size = config.max_sse_event_size;
695
696        streams.spawn(async move {
697            match client
698                .get_stream_with_max_sse_event_size(
699                    uri,
700                    session_id.clone(),
701                    None,
702                    auth_header,
703                    protocol_headers.clone(),
704                    max_sse_event_size,
705                )
706                .await
707            {
708                Ok(stream) => {
709                    let sse_stream = SseAutoReconnectStream::new(
710                        stream,
711                        StreamableHttpClientReconnect {
712                            client,
713                            session_id,
714                            uri: reconnect_uri,
715                            auth_header: reconnect_auth_header,
716                            custom_headers: protocol_headers,
717                            max_sse_event_size,
718                        },
719                        retry_config,
720                    );
721                    Self::execute_sse_stream(
722                        sse_stream,
723                        sse_worker_tx,
724                        false,
725                        transport_task_ct.child_token(),
726                    )
727                    .await
728                }
729                Err(StreamableHttpError::ServerDoesNotSupportSse) => {
730                    tracing::debug!("server doesn't support sse, skip common stream");
731                    Ok(())
732                }
733                Err(error) => {
734                    tracing::error!("fail to get common stream: {error}");
735                    Err(error)
736                }
737            }
738        });
739    }
740
741    /// Performs a transparent re-initialization handshake after a session-expired 404.
742    ///
743    /// Takes an owned clone of the client (avoiding `&self` across `.await` so the
744    /// future remains `Send` without requiring `C: Sync`).  POSTs the saved
745    /// initialize request without a session ID, extracts the new session ID and
746    /// protocol version, sends `notifications/initialized`, and returns the new
747    /// `(session_id, protocol_headers)` pair.  The init result message is **not**
748    /// forwarded to the handler because the handler already processed the original
749    /// initialization.
750    async fn perform_reinitialization(
751        client: C,
752        saved_init_request: ClientJsonRpcMessage,
753        uri: Arc<str>,
754        auth_header: Option<String>,
755        custom_headers: HashMap<HeaderName, HeaderValue>,
756        max_sse_event_size: usize,
757    ) -> Result<
758        (
759            Option<Arc<str>>,
760            ProtocolVersion,
761            HashMap<HeaderName, HeaderValue>,
762        ),
763        StreamableHttpError<C::Error>,
764    > {
765        let (init_msg, new_session_id_str) = client
766            .post_message_with_max_sse_event_size(
767                uri.clone(),
768                saved_init_request,
769                None,
770                auth_header.clone(),
771                custom_headers.clone(),
772                max_sse_event_size,
773            )
774            .await?
775            .expect_initialized::<C::Error>()
776            .await?;
777
778        let new_session_id: Option<Arc<str>> = new_session_id_str.map(|s| Arc::from(s.as_str()));
779
780        let (negotiated_version, new_protocol_headers) =
781            negotiate_version_headers(&init_msg, custom_headers);
782
783        let initialized_notification = ClientJsonRpcMessage::notification(
784            ClientNotification::InitializedNotification(InitializedNotification {
785                method: Default::default(),
786                extensions: Default::default(),
787            }),
788        );
789        // SEP-2243: notifications carry no Mcp-Param-*, so an empty tool cache suffices.
790        let initialized_headers = build_request_headers(
791            &new_protocol_headers,
792            &initialized_notification,
793            &HashMap::new(),
794            &negotiated_version,
795        );
796        client
797            .post_message_with_max_sse_event_size(
798                uri,
799                initialized_notification,
800                new_session_id.clone(),
801                auth_header,
802                initialized_headers,
803                max_sse_event_size,
804            )
805            .await?
806            .expect_accepted_or_json::<C::Error>()?;
807
808        Ok((new_session_id, negotiated_version, new_protocol_headers))
809    }
810}
811
812impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
813    type Role = RoleClient;
814    type Error = StreamableHttpError<C::Error>;
815    fn err_closed() -> Self::Error {
816        StreamableHttpError::TransportChannelClosed
817    }
818    fn err_join(e: tokio::task::JoinError) -> Self::Error {
819        StreamableHttpError::TokioJoinError(e)
820    }
821    fn config(&self) -> super::worker::WorkerConfig {
822        super::worker::WorkerConfig {
823            name: Some("StreamableHttpClientWorker".into()),
824            channel_buffer_capacity: self.config.channel_buffer_capacity,
825        }
826    }
827    async fn run(
828        self,
829        mut context: super::worker::WorkerContext<Self>,
830    ) -> Result<(), WorkerQuitReason<Self::Error>> {
831        let channel_buffer_capacity = self.config.channel_buffer_capacity;
832        let (sse_worker_tx, mut sse_worker_rx) =
833            tokio::sync::mpsc::channel::<ServerJsonRpcMessage>(channel_buffer_capacity);
834        let config = self.config.clone();
835        let transport_task_ct = context.cancellation_token.clone();
836        let _drop_guard = transport_task_ct.clone().drop_guard();
837        let WorkerSendRequest {
838            responder,
839            message: startup_request,
840        } = context.recv_from_handler().await?;
841        let is_legacy_startup = matches!(
842            &startup_request,
843            ClientJsonRpcMessage::Request(request)
844                if matches!(&request.request, ClientRequest::InitializeRequest(_))
845        );
846        let mut saved_init_request = is_legacy_startup.then(|| startup_request.clone());
847        let empty_tool_cache = HashMap::new();
848        let (bootstrap_version, bootstrap_headers) = if is_legacy_startup {
849            (ProtocolVersion::default(), config.custom_headers.clone())
850        } else {
851            request_version_headers(
852                &config.custom_headers,
853                &startup_request,
854                &ProtocolVersion::default(),
855                &empty_tool_cache,
856            )
857        };
858        let (message, session_id) = match self
859            .client
860            .post_message_with_max_sse_event_size(
861                config.uri.clone(),
862                startup_request,
863                None,
864                config.auth_header.clone(),
865                bootstrap_headers.clone(),
866                config.max_sse_event_size,
867            )
868            .await
869        {
870            Ok(res) => {
871                let _ = responder.send(Ok(()));
872                res.expect_initialized::<C::Error>().await.map_err(
873                    WorkerQuitReason::fatal_context("process initialize response"),
874                )?
875            }
876            Err(err) => {
877                let msg = format!("{:?}", err);
878                let _ = responder.send(Err(err));
879                return Err(WorkerQuitReason::fatal(
880                    StreamableHttpError::TransportChannelClosed,
881                    msg,
882                ));
883            }
884        };
885        let mut session_id: Option<Arc<str>> = if let Some(session_id) = session_id {
886            Some(session_id.into())
887        } else {
888            if !self.config.allow_stateless {
889                return Err(WorkerQuitReason::fatal(
890                    StreamableHttpError::<C::Error>::MissingSessionIdInResponse,
891                    "process initialize response",
892                ));
893            }
894            None
895        };
896        let (mut negotiated_version, mut protocol_headers) = if is_legacy_startup {
897            negotiate_version_headers(&message, config.custom_headers.clone())
898        } else {
899            (bootstrap_version, bootstrap_headers)
900        };
901        // SEP-2243: tool input schemas (name -> schema) cached from tools/list responses,
902        // used to promote annotated tools/call arguments to Mcp-Param-* headers.
903        let mut tool_header_cache: HashMap<String, Arc<JsonObject>> = HashMap::new();
904
905        // Store session info for cleanup when run() exits (not spawned, so cleanup completes before close() returns)
906        let mut session_cleanup_info = session_id.as_ref().map(|sid| SessionCleanupInfo {
907            client: self.client.clone(),
908            uri: config.uri.clone(),
909            session_id: sid.clone(),
910            auth_header: config.auth_header.clone(),
911            protocol_headers: protocol_headers.clone(),
912        });
913
914        context.send_to_handler(message).await?;
915        if is_legacy_startup {
916            let initialized_notification = context.recv_from_handler().await?;
917            let initialized_headers = build_request_headers(
918                &protocol_headers,
919                &initialized_notification.message,
920                &tool_header_cache,
921                &negotiated_version,
922            );
923            self.client
924                .post_message_with_max_sse_event_size(
925                    config.uri.clone(),
926                    initialized_notification.message,
927                    session_id.clone(),
928                    config.auth_header.clone(),
929                    initialized_headers,
930                    config.max_sse_event_size,
931                )
932                .await
933                .map_err(WorkerQuitReason::fatal_context(
934                    "send initialized notification",
935                ))?
936                .expect_accepted_or_json::<C::Error>()
937                .map_err(WorkerQuitReason::fatal_context(
938                    "process initialized notification response",
939                ))?;
940            let _ = initialized_notification.responder.send(Ok(()));
941        }
942        #[expect(
943            clippy::large_enum_variant,
944            reason = "the event is short-lived and boxing would add allocation in the event loop"
945        )]
946        enum Event<W: Worker, E: std::error::Error + Send + Sync + 'static> {
947            ClientMessage(WorkerSendRequest<W>),
948            ServerMessage(ServerJsonRpcMessage),
949            StreamResult(Result<(), StreamableHttpError<E>>),
950        }
951        let mut streams = tokio::task::JoinSet::new();
952        let mut pending_stream_response_ids = HashSet::new();
953        let mut awaiting_fallback_initialized = false;
954        if let Some(session_id) = &session_id {
955            Self::spawn_common_stream(
956                &mut streams,
957                self.client.clone(),
958                session_id.clone(),
959                &config,
960                protocol_headers.clone(),
961                sse_worker_tx.clone(),
962                transport_task_ct.clone(),
963            );
964        }
965        // Main event loop - capture exit reason so we can do cleanup before returning
966        let loop_result: Result<(), WorkerQuitReason<Self::Error>> = 'main_loop: loop {
967            let event = tokio::select! {
968                _ = transport_task_ct.cancelled() => {
969                    tracing::debug!("cancelled");
970                    break 'main_loop Err(WorkerQuitReason::Cancelled);
971                }
972                message = context.recv_from_handler() => {
973                    match message {
974                        Ok(msg) => Event::ClientMessage(msg),
975                        Err(e) => break 'main_loop Err(e),
976                    }
977                },
978                message = sse_worker_rx.recv() => {
979                    let Some(message) = message else {
980                        tracing::trace!("transport dropped, exiting");
981                        break 'main_loop Err(WorkerQuitReason::HandlerTerminated);
982                    };
983                    Event::ServerMessage(message)
984                },
985                terminated_stream = streams.join_next(), if !streams.is_empty() => {
986                    match terminated_stream {
987                        Some(result) => {
988                            Event::StreamResult(result.map_err(StreamableHttpError::TokioJoinError).and_then(std::convert::identity))
989                        }
990                        None => {
991                            continue
992                        }
993                    }
994                }
995            };
996            match event {
997                Event::ClientMessage(send_request) => {
998                    let WorkerSendRequest { message, responder } = send_request;
999                    let is_fallback_initialize = saved_init_request.is_none()
1000                        && matches!(
1001                            &message,
1002                            ClientJsonRpcMessage::Request(request)
1003                                if matches!(
1004                                    &request.request,
1005                                    ClientRequest::InitializeRequest(_)
1006                                )
1007                        );
1008                    if is_fallback_initialize {
1009                        saved_init_request = Some(message.clone());
1010                        // Servers do not assign sessions to `server/discover`, so a
1011                        // fallback initialize starts from a clean slate: no session
1012                        // ID, no cleanup state, and no streams to tear down.
1013                        debug_assert!(
1014                            session_id.is_none()
1015                                && session_cleanup_info.is_none()
1016                                && streams.is_empty(),
1017                            "discover bootstrap must not create session state"
1018                        );
1019
1020                        let response = self
1021                            .client
1022                            .post_message_with_max_sse_event_size(
1023                                config.uri.clone(),
1024                                message,
1025                                None,
1026                                config.auth_header.clone(),
1027                                config.custom_headers.clone(),
1028                                config.max_sse_event_size,
1029                            )
1030                            .await;
1031                        let response = match response {
1032                            Ok(response) => {
1033                                let _ = responder.send(Ok(()));
1034                                response
1035                            }
1036                            Err(error) => {
1037                                let _ = responder.send(Err(error));
1038                                continue;
1039                            }
1040                        };
1041                        let (initialize_response, new_session_id) = response
1042                            .expect_initialized::<C::Error>()
1043                            .await
1044                            .map_err(WorkerQuitReason::fatal_context(
1045                                "process fallback initialize response",
1046                            ))?;
1047                        session_id = new_session_id.map(Arc::from);
1048                        if session_id.is_none() && !config.allow_stateless {
1049                            return Err(WorkerQuitReason::fatal(
1050                                StreamableHttpError::<C::Error>::MissingSessionIdInResponse,
1051                                "process fallback initialize response",
1052                            ));
1053                        }
1054                        (negotiated_version, protocol_headers) = negotiate_version_headers(
1055                            &initialize_response,
1056                            config.custom_headers.clone(),
1057                        );
1058                        session_cleanup_info =
1059                            session_id.as_ref().map(|session_id| SessionCleanupInfo {
1060                                client: self.client.clone(),
1061                                uri: config.uri.clone(),
1062                                session_id: session_id.clone(),
1063                                auth_header: config.auth_header.clone(),
1064                                protocol_headers: protocol_headers.clone(),
1065                            });
1066                        context.send_to_handler(initialize_response).await?;
1067                        awaiting_fallback_initialized = true;
1068                        continue;
1069                    }
1070
1071                    let request_id = Self::client_request_id(&message);
1072                    let inline_version = match &message {
1073                        ClientJsonRpcMessage::Request(request) => {
1074                            request.request.get_meta().protocol_version()
1075                        }
1076                        _ => None,
1077                    };
1078                    let is_initialized_notification = matches!(
1079                        &message,
1080                        ClientJsonRpcMessage::Notification(notification)
1081                            if matches!(
1082                                &notification.notification,
1083                                ClientNotification::InitializedNotification(_)
1084                            )
1085                    );
1086                    // Pass a clone to the first attempt so `message` is retained for a
1087                    // potential re-init retry. `post_message` takes ownership and the
1088                    // trait cannot be changed, so the clone is unavoidable.
1089                    let (request_version, request_headers) = request_version_headers(
1090                        &protocol_headers,
1091                        &message,
1092                        &negotiated_version,
1093                        &tool_header_cache,
1094                    );
1095                    if inline_version.is_some() {
1096                        negotiated_version = request_version.clone();
1097                        if let Ok(value) = HeaderValue::from_str(request_version.as_str()) {
1098                            protocol_headers
1099                                .insert(HeaderName::from_static("mcp-protocol-version"), value);
1100                        }
1101                        if let Some(cleanup) = &mut session_cleanup_info {
1102                            cleanup.protocol_headers = protocol_headers.clone();
1103                        }
1104                    }
1105                    let response = self
1106                        .client
1107                        .post_message_with_max_sse_event_size(
1108                            config.uri.clone(),
1109                            message.clone(),
1110                            session_id.clone(),
1111                            config.auth_header.clone(),
1112                            request_headers,
1113                            config.max_sse_event_size,
1114                        )
1115                        .await;
1116                    let send_result = match response {
1117                        Err(StreamableHttpError::SessionExpired) => {
1118                            if let Some(saved_init_request) = saved_init_request
1119                                .as_ref()
1120                                .filter(|_| config.reinit_on_expired_session)
1121                            {
1122                                // The server discarded the session (HTTP 404). Perform a
1123                                // fresh handshake once and replay the original message.
1124                                tracing::info!(
1125                                    "session expired (HTTP 404), attempting transparent re-initialization"
1126                                );
1127                                match Self::perform_reinitialization(
1128                                    self.client.clone(),
1129                                    saved_init_request.clone(),
1130                                    config.uri.clone(),
1131                                    config.auth_header.clone(),
1132                                    config.custom_headers.clone(),
1133                                    config.max_sse_event_size,
1134                                )
1135                                .await
1136                                {
1137                                    Ok((
1138                                        new_session_id,
1139                                        new_negotiated_version,
1140                                        new_protocol_headers,
1141                                    )) => {
1142                                        // Old streams hold the stale session ID. Stop them first
1143                                        // so no late stale-session messages can arrive after the
1144                                        // pending requests below are completed.
1145                                        streams.abort_all();
1146                                        while streams.join_next().await.is_some() {}
1147
1148                                        // Forward any already queued response messages and fail
1149                                        // the remaining accepted requests so callers do not wait
1150                                        // forever for responses that can no longer arrive.
1151                                        Self::drain_queued_stream_messages(
1152                                            &mut sse_worker_rx,
1153                                            &mut context,
1154                                            &mut pending_stream_response_ids,
1155                                        )
1156                                        .await?;
1157                                        Self::fail_pending_stream_responses(
1158                                            &mut context,
1159                                            &mut pending_stream_response_ids,
1160                                        )
1161                                        .await?;
1162
1163                                        session_id = new_session_id;
1164                                        negotiated_version = new_negotiated_version;
1165                                        protocol_headers = new_protocol_headers;
1166                                        session_cleanup_info =
1167                                            session_id.as_ref().map(|sid| SessionCleanupInfo {
1168                                                client: self.client.clone(),
1169                                                uri: config.uri.clone(),
1170                                                session_id: sid.clone(),
1171                                                auth_header: config.auth_header.clone(),
1172                                                protocol_headers: protocol_headers.clone(),
1173                                            });
1174
1175                                        if let Some(new_sid) = &session_id {
1176                                            Self::spawn_common_stream(
1177                                                &mut streams,
1178                                                self.client.clone(),
1179                                                new_sid.clone(),
1180                                                &config,
1181                                                protocol_headers.clone(),
1182                                                sse_worker_tx.clone(),
1183                                                transport_task_ct.clone(),
1184                                            );
1185                                        }
1186
1187                                        let (_, retry_headers) = request_version_headers(
1188                                            &protocol_headers,
1189                                            &message,
1190                                            &negotiated_version,
1191                                            &tool_header_cache,
1192                                        );
1193                                        let retry_response = self
1194                                            .client
1195                                            .post_message_with_max_sse_event_size(
1196                                                config.uri.clone(),
1197                                                message,
1198                                                session_id.clone(),
1199                                                config.auth_header.clone(),
1200                                                retry_headers,
1201                                                config.max_sse_event_size,
1202                                            )
1203                                            .await;
1204                                        match retry_response {
1205                                            Err(e) => Err(e),
1206                                            Ok(StreamableHttpPostResponse::Accepted) => {
1207                                                Self::mark_stream_response_pending(
1208                                                    &mut pending_stream_response_ids,
1209                                                    request_id,
1210                                                );
1211                                                tracing::trace!(
1212                                                    "client message accepted after re-init"
1213                                                );
1214                                                Ok(())
1215                                            }
1216                                            Ok(StreamableHttpPostResponse::Json(msg, ..)) => {
1217                                                cache_tools_from_response(
1218                                                    &mut tool_header_cache,
1219                                                    &msg,
1220                                                );
1221                                                context.send_to_handler(msg).await?;
1222                                                Ok(())
1223                                            }
1224                                            Ok(StreamableHttpPostResponse::Sse(stream, ..)) => {
1225                                                Self::mark_stream_response_pending(
1226                                                    &mut pending_stream_response_ids,
1227                                                    request_id,
1228                                                );
1229                                                let sse_stream = Self::response_sse_to_jsonrpc(
1230                                                    stream,
1231                                                    session_id.clone(),
1232                                                    self.client.clone(),
1233                                                    config.uri.clone(),
1234                                                    config.auth_header.clone(),
1235                                                    protocol_headers.clone(),
1236                                                    config.max_sse_event_size,
1237                                                    self.config.retry_config.clone(),
1238                                                );
1239                                                streams.spawn(Self::execute_sse_stream(
1240                                                    sse_stream,
1241                                                    sse_worker_tx.clone(),
1242                                                    true,
1243                                                    transport_task_ct.child_token(),
1244                                                ));
1245                                                tracing::trace!("got new sse stream after re-init");
1246                                                Ok(())
1247                                            }
1248                                        }
1249                                    }
1250                                    Err(reinit_err) => Err(reinit_err),
1251                                }
1252                            } else {
1253                                Err(StreamableHttpError::SessionExpired)
1254                            }
1255                        }
1256                        Err(e) => Err(e),
1257                        Ok(StreamableHttpPostResponse::Accepted) => {
1258                            Self::mark_stream_response_pending(
1259                                &mut pending_stream_response_ids,
1260                                request_id,
1261                            );
1262                            tracing::trace!("client message accepted");
1263                            Ok(())
1264                        }
1265                        Ok(StreamableHttpPostResponse::Json(message, ..)) => {
1266                            cache_tools_from_response(&mut tool_header_cache, &message);
1267                            context.send_to_handler(message).await?;
1268                            Ok(())
1269                        }
1270                        Ok(StreamableHttpPostResponse::Sse(stream, ..)) => {
1271                            Self::mark_stream_response_pending(
1272                                &mut pending_stream_response_ids,
1273                                request_id,
1274                            );
1275                            let sse_stream = Self::response_sse_to_jsonrpc(
1276                                stream,
1277                                session_id.clone(),
1278                                self.client.clone(),
1279                                config.uri.clone(),
1280                                config.auth_header.clone(),
1281                                protocol_headers.clone(),
1282                                config.max_sse_event_size,
1283                                self.config.retry_config.clone(),
1284                            );
1285                            streams.spawn(Self::execute_sse_stream(
1286                                sse_stream,
1287                                sse_worker_tx.clone(),
1288                                true,
1289                                transport_task_ct.child_token(),
1290                            ));
1291                            tracing::trace!("got new sse stream");
1292                            Ok(())
1293                        }
1294                    };
1295                    if send_result.is_ok()
1296                        && awaiting_fallback_initialized
1297                        && is_initialized_notification
1298                    {
1299                        if let Some(session_id) = &session_id {
1300                            Self::spawn_common_stream(
1301                                &mut streams,
1302                                self.client.clone(),
1303                                session_id.clone(),
1304                                &config,
1305                                protocol_headers.clone(),
1306                                sse_worker_tx.clone(),
1307                                transport_task_ct.clone(),
1308                            );
1309                        }
1310                        awaiting_fallback_initialized = false;
1311                    }
1312                    let _ = responder.send(send_result);
1313                }
1314                Event::ServerMessage(json_rpc_message) => {
1315                    Self::clear_stream_response_pending(
1316                        &mut pending_stream_response_ids,
1317                        &json_rpc_message,
1318                    );
1319                    cache_tools_from_response(&mut tool_header_cache, &json_rpc_message);
1320                    // send the message to the handler
1321                    if let Err(e) = context.send_to_handler(json_rpc_message).await {
1322                        break 'main_loop Err(e);
1323                    }
1324                }
1325                Event::StreamResult(result) => {
1326                    if result.is_err() {
1327                        tracing::warn!(
1328                            "sse client event stream terminated with error: {:?}",
1329                            result
1330                        );
1331                    }
1332                }
1333            }
1334        };
1335
1336        // Cleanup session before returning (ensures close() waits for session deletion)
1337        // Use a timeout to prevent indefinite hangs if the server is unresponsive
1338        if let Some(cleanup) = session_cleanup_info {
1339            let cleanup_session_id = cleanup.session_id.clone();
1340            match tokio::time::timeout(
1341                SESSION_CLEANUP_TIMEOUT,
1342                cleanup.client.delete_session(
1343                    cleanup.uri,
1344                    cleanup.session_id,
1345                    cleanup.auth_header,
1346                    cleanup.protocol_headers,
1347                ),
1348            )
1349            .await
1350            {
1351                Ok(Ok(_)) => {
1352                    tracing::info!(
1353                        session_id = cleanup_session_id.as_ref(),
1354                        "delete session success"
1355                    )
1356                }
1357                Ok(Err(StreamableHttpError::ServerDoesNotSupportDeleteSession)) => {
1358                    tracing::info!(
1359                        session_id = cleanup_session_id.as_ref(),
1360                        "server doesn't support delete session"
1361                    )
1362                }
1363                Ok(Err(e)) => {
1364                    tracing::error!(
1365                        session_id = cleanup_session_id.as_ref(),
1366                        "fail to delete session: {e}"
1367                    );
1368                }
1369                Err(_elapsed) => {
1370                    tracing::warn!(
1371                        session_id = cleanup_session_id.as_ref(),
1372                        "session cleanup timed out after {:?}",
1373                        SESSION_CLEANUP_TIMEOUT
1374                    );
1375                }
1376            }
1377        }
1378
1379        loop_result
1380    }
1381}
1382
1383/// A client-agnostic HTTP transport for RMCP that supports streaming responses.
1384///
1385/// This transport allows you to choose your preferred HTTP client implementation
1386/// by implementing the [`StreamableHttpClient`] trait. The transport handles
1387/// session management, SSE streaming, and automatic reconnection.
1388///
1389/// # Usage
1390///
1391/// ## Using reqwest
1392///
1393/// ```rust,no_run
1394/// use rmcp::transport::StreamableHttpClientTransport;
1395///
1396/// // Enable the reqwest feature in Cargo.toml:
1397/// // rmcp = { version = "0.5", features = ["transport-streamable-http-client-reqwest"] }
1398///
1399/// let transport = StreamableHttpClientTransport::from_uri("http://localhost:8000/mcp");
1400/// ```
1401///
1402/// ## Using a custom HTTP client
1403///
1404/// ```rust,no_run
1405/// use rmcp::transport::streamable_http_client::{
1406///     StreamableHttpClient,
1407///     StreamableHttpClientTransport,
1408///     StreamableHttpClientTransportConfig
1409/// };
1410/// use std::sync::Arc;
1411/// use std::collections::HashMap;
1412/// use futures::stream::BoxStream;
1413/// use rmcp::model::ClientJsonRpcMessage;
1414/// use http::{HeaderName, HeaderValue};
1415/// use sse_stream::{Sse, Error as SseError};
1416///
1417/// #[derive(Clone)]
1418/// struct MyHttpClient;
1419///
1420/// #[derive(Debug, thiserror::Error)]
1421/// struct MyError;
1422///
1423/// impl std::fmt::Display for MyError {
1424///     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1425///         write!(f, "MyError")
1426///     }
1427/// }
1428///
1429/// impl StreamableHttpClient for MyHttpClient {
1430///     type Error = MyError;
1431///
1432///     async fn post_message(
1433///         &self,
1434///         _uri: Arc<str>,
1435///         _message: ClientJsonRpcMessage,
1436///         _session_id: Option<Arc<str>>,
1437///         _auth_header: Option<String>,
1438///         _custom_headers: HashMap<HeaderName, HeaderValue>,
1439///     ) -> Result<rmcp::transport::streamable_http_client::StreamableHttpPostResponse, rmcp::transport::streamable_http_client::StreamableHttpError<Self::Error>> {
1440///         todo!()
1441///     }
1442///
1443///     async fn delete_session(
1444///         &self,
1445///         _uri: Arc<str>,
1446///         _session_id: Arc<str>,
1447///         _auth_header: Option<String>,
1448///         _custom_headers: HashMap<HeaderName, HeaderValue>,
1449///     ) -> Result<(), rmcp::transport::streamable_http_client::StreamableHttpError<Self::Error>> {
1450///         todo!()
1451///     }
1452///
1453///     async fn get_stream(
1454///         &self,
1455///         _uri: Arc<str>,
1456///         _session_id: Arc<str>,
1457///         _last_event_id: Option<String>,
1458///         _auth_header: Option<String>,
1459///         _custom_headers: HashMap<HeaderName, HeaderValue>,
1460///     ) -> Result<BoxStream<'static, Result<Sse, SseError>>, rmcp::transport::streamable_http_client::StreamableHttpError<Self::Error>> {
1461///         todo!()
1462///     }
1463/// }
1464///
1465/// let transport = StreamableHttpClientTransport::with_client(
1466///     MyHttpClient,
1467///     StreamableHttpClientTransportConfig::with_uri("http://localhost:8000/mcp")
1468/// );
1469/// ```
1470///
1471/// # Feature Flags
1472///
1473/// - `transport-streamable-http-client`: Base feature providing the generic transport infrastructure
1474/// - `transport-streamable-http-client-reqwest`: Includes reqwest HTTP client support with convenience methods
1475pub type StreamableHttpClientTransport<C> = WorkerTransport<StreamableHttpClientWorker<C>>;
1476
1477impl<C: StreamableHttpClient> StreamableHttpClientTransport<C> {
1478    /// Creates a new transport with a custom HTTP client implementation.
1479    ///
1480    /// This method allows you to use any HTTP client that implements the [`StreamableHttpClient`] trait.
1481    /// Use this when you want to use a custom HTTP client or when the reqwest feature is not enabled.
1482    ///
1483    /// # Arguments
1484    ///
1485    /// * `client` - Your HTTP client implementation
1486    /// * `config` - Transport configuration including the server URI
1487    ///
1488    /// # Example
1489    ///
1490    /// ```rust,no_run
1491    /// use rmcp::transport::streamable_http_client::{
1492    ///     StreamableHttpClient,
1493    ///     StreamableHttpClientTransport,
1494    ///     StreamableHttpClientTransportConfig
1495    /// };
1496    /// use std::sync::Arc;
1497    /// use std::collections::HashMap;
1498    /// use futures::stream::BoxStream;
1499    /// use rmcp::model::ClientJsonRpcMessage;
1500    /// use http::{HeaderName, HeaderValue};
1501    /// use sse_stream::{Sse, Error as SseError};
1502    ///
1503    /// // Define your custom client
1504    /// #[derive(Clone)]
1505    /// struct MyHttpClient;
1506    ///
1507    /// #[derive(Debug, thiserror::Error)]
1508    /// struct MyError;
1509    ///
1510    /// impl std::fmt::Display for MyError {
1511    ///     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1512    ///         write!(f, "MyError")
1513    ///     }
1514    /// }
1515    ///
1516    /// impl StreamableHttpClient for MyHttpClient {
1517    ///     type Error = MyError;
1518    ///
1519    ///     async fn post_message(
1520    ///         &self,
1521    ///         _uri: Arc<str>,
1522    ///         _message: ClientJsonRpcMessage,
1523    ///         _session_id: Option<Arc<str>>,
1524    ///         _auth_header: Option<String>,
1525    ///         _custom_headers: HashMap<HeaderName, HeaderValue>,
1526    ///     ) -> Result<rmcp::transport::streamable_http_client::StreamableHttpPostResponse, rmcp::transport::streamable_http_client::StreamableHttpError<Self::Error>> {
1527    ///         todo!()
1528    ///     }
1529    ///
1530    ///     async fn delete_session(
1531    ///         &self,
1532    ///         _uri: Arc<str>,
1533    ///         _session_id: Arc<str>,
1534    ///         _auth_header: Option<String>,
1535    ///         _custom_headers: HashMap<HeaderName, HeaderValue>,
1536    ///     ) -> Result<(), rmcp::transport::streamable_http_client::StreamableHttpError<Self::Error>> {
1537    ///         todo!()
1538    ///     }
1539    ///
1540    ///     async fn get_stream(
1541    ///         &self,
1542    ///         _uri: Arc<str>,
1543    ///         _session_id: Arc<str>,
1544    ///         _last_event_id: Option<String>,
1545    ///         _auth_header: Option<String>,
1546    ///         _custom_headers: HashMap<HeaderName, HeaderValue>,
1547    ///     ) -> Result<BoxStream<'static, Result<Sse, SseError>>, rmcp::transport::streamable_http_client::StreamableHttpError<Self::Error>> {
1548    ///         todo!()
1549    ///     }
1550    /// }
1551    ///
1552    /// let transport = StreamableHttpClientTransport::with_client(
1553    ///     MyHttpClient,
1554    ///     StreamableHttpClientTransportConfig::with_uri("http://localhost:8000/mcp")
1555    /// );
1556    /// ```
1557    pub fn with_client(client: C, config: StreamableHttpClientTransportConfig) -> Self {
1558        let worker = StreamableHttpClientWorker::new(client, config);
1559        WorkerTransport::spawn(worker)
1560    }
1561}
1562#[derive(Debug, Clone)]
1563#[non_exhaustive]
1564pub struct StreamableHttpClientTransportConfig {
1565    pub uri: Arc<str>,
1566    pub retry_config: Arc<dyn SseRetryPolicy>,
1567    pub channel_buffer_capacity: usize,
1568    /// if true, the transport will not require a session to be established
1569    pub allow_stateless: bool,
1570    /// The value to send in the authorization header
1571    pub auth_header: Option<String>,
1572    /// Custom HTTP headers to include with every request
1573    pub custom_headers: HashMap<HeaderName, HeaderValue>,
1574    /// Maximum raw size of one SSE event accepted from the server.
1575    ///
1576    /// The built-in reqwest and Unix socket clients enforce this value. Custom
1577    /// [`StreamableHttpClient`] implementations must override the corresponding
1578    /// `*_with_max_sse_event_size` methods to enforce it.
1579    pub max_sse_event_size: usize,
1580    /// Enables transparent recovery when the server reports an expired session (`HTTP 404`).
1581    ///
1582    /// When enabled, the transport performs one automatic recovery attempt:
1583    /// 1. Replays the original `initialize` handshake to create a new session.
1584    /// 2. Re-establishes streaming state for that session.
1585    /// 3. Retries the in-flight request that failed with `SessionExpired`.
1586    ///
1587    /// This recovery is best-effort and bounded to a single attempt. If recovery fails,
1588    /// the original failure path is preserved and the error is returned to the caller.
1589    pub reinit_on_expired_session: bool,
1590}
1591
1592impl StreamableHttpClientTransportConfig {
1593    pub fn with_uri(uri: impl Into<Arc<str>>) -> Self {
1594        Self {
1595            uri: uri.into(),
1596            ..Default::default()
1597        }
1598    }
1599
1600    /// Set the authorization header to send with requests
1601    ///
1602    /// # Arguments
1603    ///
1604    /// * `value` - A bearer token without the `Bearer ` prefix
1605    pub fn auth_header<T: Into<String>>(mut self, value: T) -> Self {
1606        // set our authorization header
1607        self.auth_header = Some(value.into());
1608        self
1609    }
1610
1611    /// Set custom HTTP headers to include with every request
1612    ///
1613    /// # Arguments
1614    ///
1615    /// * `custom_headers` - A HashMap of header names to header values
1616    ///
1617    /// # Example
1618    ///
1619    /// ```rust,no_run
1620    /// use std::collections::HashMap;
1621    /// use http::{HeaderName, HeaderValue};
1622    /// use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig;
1623    ///
1624    /// let mut headers = HashMap::new();
1625    /// headers.insert(
1626    ///     HeaderName::from_static("x-custom-header"),
1627    ///     HeaderValue::from_static("custom-value")
1628    /// );
1629    ///
1630    /// let config = StreamableHttpClientTransportConfig::with_uri("http://localhost:8000")
1631    ///     .custom_headers(headers);
1632    /// ```
1633    pub fn custom_headers(mut self, custom_headers: HashMap<HeaderName, HeaderValue>) -> Self {
1634        self.custom_headers = custom_headers;
1635        self
1636    }
1637
1638    /// Set the maximum raw size of one SSE event accepted from the server.
1639    pub fn max_sse_event_size(mut self, bytes: usize) -> Self {
1640        self.max_sse_event_size = bytes;
1641        self
1642    }
1643
1644    /// Set whether the transport should attempt transparent re-initialization on session expiration
1645    /// See [`Self::reinit_on_expired_session`] for details.
1646    /// # Example
1647    /// ```rust,no_run
1648    /// use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig;
1649    /// let config = StreamableHttpClientTransportConfig::with_uri("http://localhost:8000")
1650    ///     .reinit_on_expired_session(true);
1651    /// ```
1652    pub fn reinit_on_expired_session(mut self, enable: bool) -> Self {
1653        self.reinit_on_expired_session = enable;
1654        self
1655    }
1656}
1657
1658impl Default for StreamableHttpClientTransportConfig {
1659    fn default() -> Self {
1660        Self {
1661            uri: "localhost".into(),
1662            retry_config: Arc::new(ExponentialBackoff::default()),
1663            channel_buffer_capacity: 16,
1664            allow_stateless: true,
1665            auth_header: None,
1666            custom_headers: HashMap::new(),
1667            max_sse_event_size: DEFAULT_MAX_SSE_EVENT_SIZE,
1668            reinit_on_expired_session: true,
1669        }
1670    }
1671}
rmcp_soddygo/transport/streamable_http_client.rs

rmcp_soddygo/transport/
streamable_http_client.rs
