rmcp_soddygo/transport/streamable_http_server/session/
local.rs

1use std::{
2    collections::{HashMap, HashSet, VecDeque},
3    num::ParseIntError,
4    time::{Duration, Instant},
5};
6
7use futures::{Stream, StreamExt};
8use thiserror::Error;
9use tokio::sync::{
10    mpsc::{Receiver, Sender},
11    oneshot,
12};
13use tokio_stream::wrappers::ReceiverStream;
14use tracing::instrument;
15
16use crate::{
17    RoleServer,
18    model::{
19        CancelledNotificationParam, ClientJsonRpcMessage, ClientNotification, ClientRequest,
20        JsonRpcNotification, JsonRpcRequest, Notification, ProgressNotificationParam,
21        ProgressToken, RequestId, ServerJsonRpcMessage, ServerNotification,
22    },
23    transport::{
24        WorkerTransport,
25        common::server_side_http::{SessionId, session_id},
26        worker::{Worker, WorkerContext, WorkerQuitReason, WorkerSendRequest},
27    },
28};
29
30#[derive(Debug, Default)]
31#[non_exhaustive]
32pub struct LocalSessionManager {
33    pub sessions: tokio::sync::RwLock<HashMap<SessionId, LocalSessionHandle>>,
34    pub session_config: SessionConfig,
35}
36
37#[derive(Debug, Error)]
38#[non_exhaustive]
39pub enum LocalSessionManagerError {
40    #[error("Session not found: {0}")]
41    SessionNotFound(SessionId),
42    #[error("Session error: {0}")]
43    SessionError(#[from] SessionError),
44    #[error("Invalid event id: {0}")]
45    InvalidEventId(#[from] EventIdParseError),
46}
47impl SessionManager for LocalSessionManager {
48    type Error = LocalSessionManagerError;
49    type Transport = WorkerTransport<LocalSessionWorker>;
50    async fn create_session(&self) -> Result<(SessionId, Self::Transport), Self::Error> {
51        let id = session_id();
52        let (handle, worker) = create_local_session(id.clone(), self.session_config.clone());
53        self.sessions.write().await.insert(id.clone(), handle);
54        Ok((id, WorkerTransport::spawn(worker)))
55    }
56    async fn initialize_session(
57        &self,
58        id: &SessionId,
59        message: ClientJsonRpcMessage,
60    ) -> Result<ServerJsonRpcMessage, Self::Error> {
61        let sessions = self.sessions.read().await;
62        let handle = sessions
63            .get(id)
64            .ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
65        let response = handle.initialize(message).await?;
66        Ok(response)
67    }
68    async fn close_session(&self, id: &SessionId) -> Result<(), Self::Error> {
69        let handle = {
70            let mut sessions = self.sessions.write().await;
71            sessions.remove(id)
72        };
73        if let Some(handle) = handle {
74            match handle.close().await {
75                // Worker already exited — nothing left to clean up.
76                Ok(()) | Err(SessionError::SessionServiceTerminated) => {}
77                Err(e) => return Err(e.into()),
78            }
79        }
80        Ok(())
81    }
82    async fn has_session(&self, id: &SessionId) -> Result<bool, Self::Error> {
83        let sessions = self.sessions.read().await;
84        Ok(sessions.contains_key(id))
85    }
86    async fn create_stream(
87        &self,
88        id: &SessionId,
89        message: ClientJsonRpcMessage,
90    ) -> Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error> {
91        let sessions = self.sessions.read().await;
92        let handle = sessions
93            .get(id)
94            .ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
95        let receiver = handle.establish_request_wise_channel().await?;
96        let http_request_id = receiver.http_request_id;
97        handle.push_message(message, http_request_id).await?;
98
99        let priming = self.session_config.sse_retry.map(|retry| {
100            let event_id = match http_request_id {
101                Some(id) => format!("0/{id}"),
102                None => "0".into(),
103            };
104            ServerSseMessage::priming(event_id, retry)
105        });
106        Ok(futures::stream::iter(priming).chain(ReceiverStream::new(receiver.inner)))
107    }
108
109    async fn create_standalone_stream(
110        &self,
111        id: &SessionId,
112    ) -> Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error> {
113        let sessions = self.sessions.read().await;
114        let handle = sessions
115            .get(id)
116            .ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
117        let receiver = handle.establish_common_channel().await?;
118        Ok(ReceiverStream::new(receiver.inner))
119    }
120
121    async fn resume(
122        &self,
123        id: &SessionId,
124        last_event_id: String,
125    ) -> Result<impl Stream<Item = ServerSseMessage> + Send + 'static, Self::Error> {
126        let sessions = self.sessions.read().await;
127        let handle = sessions
128            .get(id)
129            .ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
130        let receiver = handle.resume(last_event_id.parse()?).await?;
131        Ok(ReceiverStream::new(receiver.inner))
132    }
133
134    async fn accept_message(
135        &self,
136        id: &SessionId,
137        message: ClientJsonRpcMessage,
138    ) -> Result<(), Self::Error> {
139        let sessions = self.sessions.read().await;
140        let handle = sessions
141            .get(id)
142            .ok_or(LocalSessionManagerError::SessionNotFound(id.clone()))?;
143        handle.push_message(message, None).await?;
144        Ok(())
145    }
146
147    async fn restore_session(
148        &self,
149        id: SessionId,
150    ) -> Result<RestoreOutcome<Self::Transport>, Self::Error> {
151        let mut sessions = self.sessions.write().await;
152        if sessions.contains_key(&id) {
153            // A concurrent request already restored this session.
154            return Ok(RestoreOutcome::AlreadyPresent);
155        }
156        let (handle, worker) = create_local_session(id.clone(), self.session_config.clone());
157        sessions.insert(id, handle);
158        Ok(RestoreOutcome::Restored(WorkerTransport::spawn(worker)))
159    }
160}
161
162/// `<index>/request_id>`
163#[derive(Debug, Clone, PartialEq, Eq, Hash)]
164pub struct EventId {
165    http_request_id: Option<HttpRequestId>,
166    index: usize,
167}
168
169impl std::fmt::Display for EventId {
170    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
171        write!(f, "{}", self.index)?;
172        match &self.http_request_id {
173            Some(http_request_id) => write!(f, "/{http_request_id}"),
174            None => write!(f, ""),
175        }
176    }
177}
178
179#[derive(Debug, Clone, Error)]
180#[non_exhaustive]
181pub enum EventIdParseError {
182    #[error("Invalid index: {0}")]
183    InvalidIndex(ParseIntError),
184    #[error("Invalid numeric request id: {0}")]
185    InvalidNumericRequestId(ParseIntError),
186    #[error("Missing request id type")]
187    InvalidRequestIdType,
188    #[error("Missing request id")]
189    MissingRequestId,
190}
191
192impl std::str::FromStr for EventId {
193    type Err = EventIdParseError;
194    fn from_str(s: &str) -> Result<Self, Self::Err> {
195        if let Some((index, request_id)) = s.split_once("/") {
196            let index = usize::from_str(index).map_err(EventIdParseError::InvalidIndex)?;
197            let request_id = u64::from_str(request_id).map_err(EventIdParseError::InvalidIndex)?;
198            Ok(EventId {
199                http_request_id: Some(request_id),
200                index,
201            })
202        } else {
203            let index = usize::from_str(s).map_err(EventIdParseError::InvalidIndex)?;
204            Ok(EventId {
205                http_request_id: None,
206                index,
207            })
208        }
209    }
210}
211
212use super::{RestoreOutcome, ServerSseMessage, SessionManager};
213
214struct CachedTx {
215    tx: Sender<ServerSseMessage>,
216    cache: VecDeque<ServerSseMessage>,
217    http_request_id: Option<HttpRequestId>,
218    capacity: usize,
219    starting_index: usize,
220}
221
222impl CachedTx {
223    fn new(
224        tx: Sender<ServerSseMessage>,
225        http_request_id: Option<HttpRequestId>,
226        starting_index: usize,
227    ) -> Self {
228        Self {
229            cache: VecDeque::with_capacity(tx.capacity()),
230            capacity: tx.capacity(),
231            tx,
232            http_request_id,
233            starting_index,
234        }
235    }
236    fn new_common(tx: Sender<ServerSseMessage>) -> Self {
237        Self::new(tx, None, 0)
238    }
239
240    fn next_event_id(&self) -> EventId {
241        let index = self.cache.back().map_or(self.starting_index, |m| {
242            m.event_id
243                .as_deref()
244                .unwrap_or_default()
245                .parse::<EventId>()
246                .expect("valid event id")
247                .index
248                + 1
249        });
250        EventId {
251            http_request_id: self.http_request_id,
252            index,
253        }
254    }
255
256    async fn send(&mut self, message: ServerJsonRpcMessage) {
257        let event_id = self.next_event_id();
258        let message = ServerSseMessage::new(event_id.to_string(), message);
259        self.cache_and_send(message).await;
260    }
261
262    async fn send_priming(&mut self, retry: Duration) {
263        let event_id = self.next_event_id();
264        let message = ServerSseMessage::priming(event_id.to_string(), retry);
265        self.cache_and_send(message).await;
266    }
267
268    async fn cache_and_send(&mut self, message: ServerSseMessage) {
269        if self.cache.len() >= self.capacity {
270            self.cache.pop_front();
271            self.cache.push_back(message.clone());
272        } else {
273            self.cache.push_back(message.clone());
274        }
275        let _ = self.tx.send(message).await.inspect_err(|e| {
276            let event_id = &e.0.event_id;
277            tracing::trace!(?event_id, "trying to send message in a closed session")
278        });
279    }
280
281    async fn sync(&mut self, index: usize) -> Result<(), SessionError> {
282        let Some(front) = self.cache.front() else {
283            return Ok(());
284        };
285        let front_event_id = front
286            .event_id
287            .as_deref()
288            .unwrap_or_default()
289            .parse::<EventId>()?;
290        let sync_index = index.saturating_sub(front_event_id.index);
291        if sync_index > self.cache.len() {
292            // invalid index
293            return Err(SessionError::InvalidEventId);
294        }
295        for message in self.cache.iter().skip(sync_index) {
296            let send_result = self.tx.send(message.clone()).await;
297            if send_result.is_err() {
298                let event_id: EventId = message.event_id.as_deref().unwrap_or_default().parse()?;
299                return Err(SessionError::ChannelClosed(Some(event_id.index as u64)));
300            }
301        }
302        Ok(())
303    }
304}
305
306struct HttpRequestWise {
307    resources: HashSet<ResourceKey>,
308    tx: CachedTx,
309    completed_at: Option<Instant>,
310}
311
312type HttpRequestId = u64;
313#[derive(Debug, Clone, Hash, PartialEq, Eq)]
314enum ResourceKey {
315    McpRequestId(RequestId),
316    ProgressToken(ProgressToken),
317}
318
319pub struct LocalSessionWorker {
320    id: SessionId,
321    next_http_request_id: HttpRequestId,
322    tx_router: HashMap<HttpRequestId, HttpRequestWise>,
323    resource_router: HashMap<ResourceKey, HttpRequestId>,
324    common: CachedTx,
325    /// Shadow senders for secondary SSE streams (e.g. from POST EventSource
326    /// reconnections). These keep the HTTP connections alive via SSE keep-alive
327    /// without receiving notifications, preventing MCP clients from entering
328    /// infinite reconnect loops when multiple EventSource connections compete
329    /// to replace the common channel.
330    shadow_txs: Vec<Sender<ServerSseMessage>>,
331    event_rx: Receiver<SessionEvent>,
332    session_config: SessionConfig,
333}
334
335impl LocalSessionWorker {
336    pub fn id(&self) -> &SessionId {
337        &self.id
338    }
339}
340
341#[derive(Debug, Error)]
342#[non_exhaustive]
343pub enum SessionError {
344    #[error("Invalid request id: {0}")]
345    DuplicatedRequestId(HttpRequestId),
346    #[error("Channel closed: {0:?}")]
347    ChannelClosed(Option<HttpRequestId>),
348    #[error("Cannot parse event id: {0}")]
349    EventIdParseError(#[from] EventIdParseError),
350    #[error("Session service terminated")]
351    SessionServiceTerminated,
352    #[error("Invalid event id")]
353    InvalidEventId,
354    #[error("IO error: {0}")]
355    Io(#[from] std::io::Error),
356}
357
358impl From<SessionError> for std::io::Error {
359    fn from(value: SessionError) -> Self {
360        match value {
361            SessionError::Io(io) => io,
362            _ => std::io::Error::other(format!("Session error: {value}")),
363        }
364    }
365}
366
367enum OutboundChannel {
368    RequestWise { id: HttpRequestId, close: bool },
369    Common,
370}
371#[derive(Debug)]
372#[non_exhaustive]
373pub struct StreamableHttpMessageReceiver {
374    pub http_request_id: Option<HttpRequestId>,
375    pub inner: Receiver<ServerSseMessage>,
376}
377
378impl LocalSessionWorker {
379    fn unregister_resource(&mut self, resource: &ResourceKey) {
380        let Some(http_request_id) = self.resource_router.remove(resource) else {
381            return;
382        };
383        tracing::trace!(?resource, http_request_id, "unregister resource");
384        let Some(channel) = self.tx_router.get_mut(&http_request_id) else {
385            tracing::warn!(http_request_id, "http request wise channel not found");
386            return;
387        };
388        if !channel.resources.is_empty() && !matches!(resource, ResourceKey::McpRequestId(_)) {
389            return;
390        }
391        tracing::debug!(http_request_id, "close http request wise channel");
392        let resources: Vec<_> = channel.resources.drain().collect();
393        channel.completed_at = Some(Instant::now());
394        // Close the sender so the client's SSE stream ends,
395        // but keep the entry so the cache is available for
396        // late resume requests.
397        let (closed_tx, _) = tokio::sync::mpsc::channel(1);
398        channel.tx.tx = closed_tx;
399        for resource in resources {
400            self.resource_router.remove(&resource);
401        }
402    }
403    fn register_resource(&mut self, resource: ResourceKey, http_request_id: HttpRequestId) {
404        tracing::trace!(?resource, http_request_id, "register resource");
405        if let Some(channel) = self.tx_router.get_mut(&http_request_id) {
406            channel.resources.insert(resource.clone());
407            self.resource_router.insert(resource, http_request_id);
408        }
409    }
410    fn register_request(
411        &mut self,
412        request: &JsonRpcRequest<ClientRequest>,
413        http_request_id: HttpRequestId,
414    ) {
415        use crate::model::GetMeta;
416        self.register_resource(
417            ResourceKey::McpRequestId(request.id.clone()),
418            http_request_id,
419        );
420        if let Some(progress_token) = request.request.get_meta().get_progress_token() {
421            self.register_resource(
422                ResourceKey::ProgressToken(progress_token.clone()),
423                http_request_id,
424            );
425        }
426    }
427    fn catch_cancellation_notification(
428        &mut self,
429        notification: &JsonRpcNotification<ClientNotification>,
430    ) {
431        if let ClientNotification::CancelledNotification(n) = &notification.notification {
432            if let Some(request_id) = n.params.request_id.clone() {
433                let resource = ResourceKey::McpRequestId(request_id);
434                self.unregister_resource(&resource);
435            }
436        }
437    }
438    fn evict_expired_channels(&mut self) {
439        let ttl = self.session_config.completed_cache_ttl;
440        self.tx_router
441            .retain(|_, rw| rw.completed_at.is_none_or(|at| at.elapsed() < ttl));
442    }
443    fn next_http_request_id(&mut self) -> HttpRequestId {
444        let id = self.next_http_request_id;
445        self.next_http_request_id = self.next_http_request_id.wrapping_add(1);
446        id
447    }
448    async fn establish_request_wise_channel(
449        &mut self,
450    ) -> Result<StreamableHttpMessageReceiver, SessionError> {
451        let http_request_id = self.next_http_request_id();
452        let (tx, rx) = tokio::sync::mpsc::channel(self.session_config.channel_capacity);
453        let starting_index = usize::from(self.session_config.sse_retry.is_some());
454        self.tx_router.insert(
455            http_request_id,
456            HttpRequestWise {
457                resources: Default::default(),
458                tx: CachedTx::new(tx, Some(http_request_id), starting_index),
459                completed_at: None,
460            },
461        );
462        tracing::debug!(http_request_id, "establish new request wise channel");
463        Ok(StreamableHttpMessageReceiver {
464            http_request_id: Some(http_request_id),
465            inner: rx,
466        })
467    }
468    fn resolve_outbound_channel(&self, message: &ServerJsonRpcMessage) -> OutboundChannel {
469        match &message {
470            ServerJsonRpcMessage::Request(_) => OutboundChannel::Common,
471            ServerJsonRpcMessage::Notification(JsonRpcNotification {
472                notification:
473                    ServerNotification::ProgressNotification(Notification {
474                        params: ProgressNotificationParam { progress_token, .. },
475                        ..
476                    }),
477                ..
478            }) => {
479                let id = self
480                    .resource_router
481                    .get(&ResourceKey::ProgressToken(progress_token.clone()));
482
483                if let Some(id) = id {
484                    OutboundChannel::RequestWise {
485                        id: *id,
486                        close: false,
487                    }
488                } else {
489                    OutboundChannel::Common
490                }
491            }
492            ServerJsonRpcMessage::Notification(JsonRpcNotification {
493                notification:
494                    ServerNotification::CancelledNotification(Notification {
495                        params: CancelledNotificationParam { request_id, .. },
496                        ..
497                    }),
498                ..
499            }) => {
500                if let Some(req_id) = request_id {
501                    if let Some(id) = self
502                        .resource_router
503                        .get(&ResourceKey::McpRequestId(req_id.clone()))
504                    {
505                        OutboundChannel::RequestWise {
506                            id: *id,
507                            close: false,
508                        }
509                    } else {
510                        OutboundChannel::Common
511                    }
512                } else {
513                    OutboundChannel::Common
514                }
515            }
516            ServerJsonRpcMessage::Notification(_) => OutboundChannel::Common,
517            ServerJsonRpcMessage::Response(json_rpc_response) => {
518                if let Some(id) = self
519                    .resource_router
520                    .get(&ResourceKey::McpRequestId(json_rpc_response.id.clone()))
521                {
522                    OutboundChannel::RequestWise {
523                        id: *id,
524                        close: true,
525                    }
526                } else {
527                    OutboundChannel::Common
528                }
529            }
530            ServerJsonRpcMessage::Error(json_rpc_error) => {
531                if let Some(id) = json_rpc_error.id.clone().and_then(|rid| {
532                    self.resource_router
533                        .get(&ResourceKey::McpRequestId(rid))
534                        .copied()
535                }) {
536                    OutboundChannel::RequestWise { id, close: true }
537                } else {
538                    OutboundChannel::Common
539                }
540            }
541        }
542    }
543    async fn handle_server_message(
544        &mut self,
545        message: ServerJsonRpcMessage,
546    ) -> Result<(), SessionError> {
547        let outbound_channel = self.resolve_outbound_channel(&message);
548        match outbound_channel {
549            OutboundChannel::RequestWise { id, close } => {
550                if let Some(request_wise) = self.tx_router.get_mut(&id) {
551                    request_wise.tx.send(message).await;
552                    if close {
553                        if let Some(channel) = self.tx_router.remove(&id) {
554                            for resource in channel.resources {
555                                self.resource_router.remove(&resource);
556                            }
557                        }
558                    }
559                } else {
560                    return Err(SessionError::ChannelClosed(Some(id)));
561                }
562            }
563            OutboundChannel::Common => self.common.send(message).await,
564        }
565        Ok(())
566    }
567    async fn resume(
568        &mut self,
569        last_event_id: EventId,
570    ) -> Result<StreamableHttpMessageReceiver, SessionError> {
571        // Clean up closed shadow senders before processing
572        self.shadow_txs.retain(|tx| !tx.is_closed());
573
574        match last_event_id.http_request_id {
575            Some(http_request_id) => {
576                let request_wise = self
577                    .tx_router
578                    .get_mut(&http_request_id)
579                    .ok_or(SessionError::ChannelClosed(Some(http_request_id)))?;
580                let is_completed = request_wise.completed_at.is_some();
581                let (tx, rx) = tokio::sync::mpsc::channel(self.session_config.channel_capacity);
582                request_wise.tx.tx = tx;
583                let index = last_event_id.index;
584                request_wise.tx.sync(index).await?;
585                if is_completed {
586                    // Drop the sender after replaying so the stream ends
587                    // instead of hanging indefinitely.
588                    let (closed_tx, _) = tokio::sync::mpsc::channel(1);
589                    request_wise.tx.tx = closed_tx;
590                }
591                Ok(StreamableHttpMessageReceiver {
592                    http_request_id: Some(http_request_id),
593                    inner: rx,
594                })
595            }
596            None => self.resume_or_shadow_common(last_event_id.index).await,
597        }
598    }
599
600    /// Resume the common channel, or create a shadow stream if the primary is
601    /// still active.
602    ///
603    /// When the primary common channel is dead (receiver dropped), replace it
604    /// so this stream becomes the new primary notification channel. Cached
605    /// messages are replayed from `last_event_index` so the client receives
606    /// any events it missed (including server-initiated requests).
607    ///
608    /// When the primary is still active, create a "shadow" stream — an idle SSE
609    /// connection kept alive by keep-alive pings. This prevents multiple
610    /// EventSource connections (e.g. from POST response reconnections) from
611    /// killing each other by repeatedly replacing the common channel sender.
612    async fn resume_or_shadow_common(
613        &mut self,
614        last_event_index: usize,
615    ) -> Result<StreamableHttpMessageReceiver, SessionError> {
616        let is_replacing_dead_primary = self.common.tx.is_closed();
617        let capacity = if is_replacing_dead_primary {
618            self.session_config.channel_capacity
619        } else {
620            1 // Shadow streams only need keep-alive pings
621        };
622        let (tx, rx) = tokio::sync::mpsc::channel(capacity);
623        if is_replacing_dead_primary {
624            // Primary common channel is dead — replace it.
625            tracing::debug!("Replacing dead common channel with new primary");
626            self.common.tx = tx;
627            // Replay cached messages from where the client left off so
628            // server-initiated requests and notifications are not lost.
629            self.common.sync(last_event_index).await?;
630        } else {
631            // Primary common channel is still active. Create a shadow stream
632            // that stays alive via SSE keep-alive but doesn't receive
633            // notifications. This prevents competing EventSource connections
634            // from killing each other's channels.
635            const MAX_SHADOW_STREAMS: usize = 32;
636
637            if self.shadow_txs.len() >= MAX_SHADOW_STREAMS {
638                tracing::warn!(
639                    shadow_count = self.shadow_txs.len(),
640                    "Shadow stream limit reached, dropping oldest"
641                );
642                self.shadow_txs.remove(0);
643            }
644            tracing::debug!(
645                shadow_count = self.shadow_txs.len(),
646                "Common channel active, creating shadow stream"
647            );
648            self.shadow_txs.push(tx);
649        }
650        Ok(StreamableHttpMessageReceiver {
651            http_request_id: None,
652            inner: rx,
653        })
654    }
655
656    async fn close_sse_stream(
657        &mut self,
658        http_request_id: Option<HttpRequestId>,
659        retry_interval: Option<Duration>,
660    ) -> Result<(), SessionError> {
661        match http_request_id {
662            // Close a request-wise stream
663            Some(id) => {
664                let request_wise = self
665                    .tx_router
666                    .get_mut(&id)
667                    .ok_or(SessionError::ChannelClosed(Some(id)))?;
668
669                // Send priming event if retry interval is specified
670                if let Some(interval) = retry_interval {
671                    request_wise.tx.send_priming(interval).await;
672                }
673
674                // Close the stream by dropping the sender
675                let (tx, _rx) = tokio::sync::mpsc::channel(1);
676                request_wise.tx.tx = tx;
677
678                tracing::debug!(
679                    http_request_id = id,
680                    "closed SSE stream for server-initiated disconnection"
681                );
682                Ok(())
683            }
684            // Close the standalone (common) stream
685            None => {
686                // Send priming event if retry interval is specified
687                if let Some(interval) = retry_interval {
688                    self.common.send_priming(interval).await;
689                }
690
691                // Close the stream by dropping the sender
692                let (tx, _rx) = tokio::sync::mpsc::channel(1);
693                self.common.tx = tx;
694
695                // Also close all shadow streams
696                self.shadow_txs.clear();
697
698                tracing::debug!("closed standalone SSE stream for server-initiated disconnection");
699                Ok(())
700            }
701        }
702    }
703}
704
705#[derive(Debug)]
706#[non_exhaustive]
707pub enum SessionEvent {
708    ClientMessage {
709        message: ClientJsonRpcMessage,
710        http_request_id: Option<HttpRequestId>,
711    },
712    EstablishRequestWiseChannel {
713        responder: oneshot::Sender<Result<StreamableHttpMessageReceiver, SessionError>>,
714    },
715    CloseRequestWiseChannel {
716        id: HttpRequestId,
717        responder: oneshot::Sender<Result<(), SessionError>>,
718    },
719    Resume {
720        last_event_id: EventId,
721        responder: oneshot::Sender<Result<StreamableHttpMessageReceiver, SessionError>>,
722    },
723    InitializeRequest {
724        request: ClientJsonRpcMessage,
725        responder: oneshot::Sender<Result<ServerJsonRpcMessage, SessionError>>,
726    },
727    Close,
728    CloseSseStream {
729        /// The HTTP request ID to close. If `None`, closes the standalone (common) stream.
730        http_request_id: Option<HttpRequestId>,
731        /// Optional retry interval. If provided, a priming event is sent before closing.
732        retry_interval: Option<Duration>,
733        responder: oneshot::Sender<Result<(), SessionError>>,
734    },
735}
736
737#[derive(Debug, Clone)]
738#[non_exhaustive]
739pub enum SessionQuitReason {
740    ServiceTerminated,
741    ClientTerminated,
742    ExpectInitializeRequest,
743    ExpectInitializeResponse,
744    Cancelled,
745}
746
747#[derive(Debug, Clone)]
748pub struct LocalSessionHandle {
749    id: SessionId,
750    // after all event_tx drop, inner task will be terminated
751    event_tx: Sender<SessionEvent>,
752}
753
754impl LocalSessionHandle {
755    /// Get the session id
756    pub fn id(&self) -> &SessionId {
757        &self.id
758    }
759
760    /// Close the session
761    pub async fn close(&self) -> Result<(), SessionError> {
762        self.event_tx
763            .send(SessionEvent::Close)
764            .await
765            .map_err(|_| SessionError::SessionServiceTerminated)?;
766        Ok(())
767    }
768
769    /// Send a message to the session
770    pub async fn push_message(
771        &self,
772        message: ClientJsonRpcMessage,
773        http_request_id: Option<HttpRequestId>,
774    ) -> Result<(), SessionError> {
775        self.event_tx
776            .send(SessionEvent::ClientMessage {
777                message,
778                http_request_id,
779            })
780            .await
781            .map_err(|_| SessionError::SessionServiceTerminated)?;
782        Ok(())
783    }
784
785    /// establish a channel for a http-request, the corresponded message from server will be
786    /// sent through this channel. The channel will be closed when the request is completed,
787    /// or you can close it manually by calling [`LocalSessionHandle::close_request_wise_channel`].
788    pub async fn establish_request_wise_channel(
789        &self,
790    ) -> Result<StreamableHttpMessageReceiver, SessionError> {
791        let (tx, rx) = tokio::sync::oneshot::channel();
792        self.event_tx
793            .send(SessionEvent::EstablishRequestWiseChannel { responder: tx })
794            .await
795            .map_err(|_| SessionError::SessionServiceTerminated)?;
796        rx.await
797            .map_err(|_| SessionError::SessionServiceTerminated)?
798    }
799
800    /// close the http-request wise channel.
801    pub async fn close_request_wise_channel(
802        &self,
803        request_id: HttpRequestId,
804    ) -> Result<(), SessionError> {
805        let (tx, rx) = tokio::sync::oneshot::channel();
806        self.event_tx
807            .send(SessionEvent::CloseRequestWiseChannel {
808                id: request_id,
809                responder: tx,
810            })
811            .await
812            .map_err(|_| SessionError::SessionServiceTerminated)?;
813        rx.await
814            .map_err(|_| SessionError::SessionServiceTerminated)?
815    }
816
817    /// Establish a common channel for general purpose messages.
818    pub async fn establish_common_channel(
819        &self,
820    ) -> Result<StreamableHttpMessageReceiver, SessionError> {
821        let (tx, rx) = tokio::sync::oneshot::channel();
822        self.event_tx
823            .send(SessionEvent::Resume {
824                last_event_id: EventId {
825                    http_request_id: None,
826                    index: 0,
827                },
828                responder: tx,
829            })
830            .await
831            .map_err(|_| SessionError::SessionServiceTerminated)?;
832        rx.await
833            .map_err(|_| SessionError::SessionServiceTerminated)?
834    }
835
836    /// Resume streaming response by the last event id. This is suitable for both request wise and common channel.
837    pub async fn resume(
838        &self,
839        last_event_id: EventId,
840    ) -> Result<StreamableHttpMessageReceiver, SessionError> {
841        let (tx, rx) = tokio::sync::oneshot::channel();
842        self.event_tx
843            .send(SessionEvent::Resume {
844                last_event_id,
845                responder: tx,
846            })
847            .await
848            .map_err(|_| SessionError::SessionServiceTerminated)?;
849        rx.await
850            .map_err(|_| SessionError::SessionServiceTerminated)?
851    }
852
853    /// Send an initialize request to the session. And wait for the initialized response.
854    ///
855    /// This is used to establish a session with the server.
856    pub async fn initialize(
857        &self,
858        request: ClientJsonRpcMessage,
859    ) -> Result<ServerJsonRpcMessage, SessionError> {
860        let (tx, rx) = tokio::sync::oneshot::channel();
861        self.event_tx
862            .send(SessionEvent::InitializeRequest {
863                request,
864                responder: tx,
865            })
866            .await
867            .map_err(|_| SessionError::SessionServiceTerminated)?;
868        rx.await
869            .map_err(|_| SessionError::SessionServiceTerminated)?
870    }
871
872    /// Close an SSE stream for a specific request.
873    ///
874    /// This closes the SSE connection for a POST request stream, but keeps the session
875    /// and message cache active. Clients can reconnect using the `Last-Event-ID` header
876    /// via a GET request to resume receiving messages.
877    ///
878    /// # Arguments
879    ///
880    /// * `http_request_id` - The HTTP request ID of the stream to close
881    /// * `retry_interval` - Optional retry interval. If provided, a priming event is sent
882    pub async fn close_sse_stream(
883        &self,
884        http_request_id: HttpRequestId,
885        retry_interval: Option<Duration>,
886    ) -> Result<(), SessionError> {
887        let (tx, rx) = tokio::sync::oneshot::channel();
888        self.event_tx
889            .send(SessionEvent::CloseSseStream {
890                http_request_id: Some(http_request_id),
891                retry_interval,
892                responder: tx,
893            })
894            .await
895            .map_err(|_| SessionError::SessionServiceTerminated)?;
896        rx.await
897            .map_err(|_| SessionError::SessionServiceTerminated)?
898    }
899
900    /// Close the standalone SSE stream.
901    ///
902    /// This closes the standalone SSE connection (established via GET request),
903    /// but keeps the session and message cache active. Clients can reconnect using
904    /// the `Last-Event-ID` header via a GET request to resume receiving messages.
905    ///
906    /// # Arguments
907    ///
908    /// * `retry_interval` - Optional retry interval. If provided, a priming event is sent
909    pub async fn close_standalone_sse_stream(
910        &self,
911        retry_interval: Option<Duration>,
912    ) -> Result<(), SessionError> {
913        let (tx, rx) = tokio::sync::oneshot::channel();
914        self.event_tx
915            .send(SessionEvent::CloseSseStream {
916                http_request_id: None,
917                retry_interval,
918                responder: tx,
919            })
920            .await
921            .map_err(|_| SessionError::SessionServiceTerminated)?;
922        rx.await
923            .map_err(|_| SessionError::SessionServiceTerminated)?
924    }
925}
926
927pub type SessionTransport = WorkerTransport<LocalSessionWorker>;
928
929#[allow(clippy::large_enum_variant)]
930#[derive(Debug, Error)]
931#[non_exhaustive]
932pub enum LocalSessionWorkerError {
933    #[error("transport terminated")]
934    TransportTerminated,
935    #[error("unexpected message: {0:?}")]
936    UnexpectedEvent(SessionEvent),
937    #[error("fail to send initialize request {0}")]
938    FailToSendInitializeRequest(SessionError),
939    #[error("fail to handle message: {0}")]
940    FailToHandleMessage(SessionError),
941    #[deprecated(note = "idle timeout now surfaces as WorkerQuitReason::IdleTimeout")]
942    #[error("keep alive timeout after {}ms", _0.as_millis())]
943    KeepAliveTimeout(Duration),
944    #[error("init timeout after {}ms", _0.as_millis())]
945    InitTimeout(Duration),
946    #[error("Transport closed")]
947    TransportClosed,
948    #[error("Tokio join error {0}")]
949    TokioJoinError(#[from] tokio::task::JoinError),
950}
951impl Worker for LocalSessionWorker {
952    type Error = LocalSessionWorkerError;
953    type Role = RoleServer;
954    fn err_closed() -> Self::Error {
955        LocalSessionWorkerError::TransportClosed
956    }
957    fn err_join(e: tokio::task::JoinError) -> Self::Error {
958        LocalSessionWorkerError::TokioJoinError(e)
959    }
960    fn config(&self) -> crate::transport::worker::WorkerConfig {
961        crate::transport::worker::WorkerConfig {
962            name: Some(format!("streamable-http-session-{}", self.id)),
963            channel_buffer_capacity: self.session_config.channel_capacity,
964        }
965    }
966    #[instrument(name = "streamable_http_session", skip_all, fields(id = self.id.as_ref()))]
967    async fn run(
968        mut self,
969        mut context: WorkerContext<Self>,
970    ) -> Result<(), WorkerQuitReason<Self::Error>> {
971        enum InnerEvent {
972            FromHttpService(SessionEvent),
973            FromHandler(WorkerSendRequest<LocalSessionWorker>),
974        }
975        let init_timeout = self.session_config.init_timeout.unwrap_or(Duration::MAX);
976        let evt = tokio::select! {
977            evt = self.event_rx.recv() => evt.ok_or_else(|| {
978                WorkerQuitReason::fatal(
979                    LocalSessionWorkerError::TransportTerminated,
980                    "get initialize request",
981                )
982            })?,
983            _ = context.cancellation_token.cancelled() => {
984                return Err(WorkerQuitReason::Cancelled);
985            }
986            _ = tokio::time::sleep(init_timeout) => {
987                return Err(WorkerQuitReason::fatal(
988                    LocalSessionWorkerError::InitTimeout(init_timeout),
989                    "waiting for initialize request",
990                ));
991            }
992        };
993        let SessionEvent::InitializeRequest { request, responder } = evt else {
994            return Err(WorkerQuitReason::fatal(
995                LocalSessionWorkerError::UnexpectedEvent(evt),
996                "get initialize request",
997            ));
998        };
999        context.send_to_handler(request).await?;
1000        let send_initialize_response = context.recv_from_handler().await?;
1001        responder
1002            .send(Ok(send_initialize_response.message))
1003            .map_err(|_| {
1004                WorkerQuitReason::fatal(
1005                    LocalSessionWorkerError::FailToSendInitializeRequest(
1006                        SessionError::SessionServiceTerminated,
1007                    ),
1008                    "send initialize response",
1009                )
1010            })?;
1011        send_initialize_response
1012            .responder
1013            .send(Ok(()))
1014            .map_err(|_| WorkerQuitReason::HandlerTerminated)?;
1015        let ct = context.cancellation_token.clone();
1016        let keep_alive = self.session_config.keep_alive.unwrap_or(Duration::MAX);
1017        loop {
1018            self.evict_expired_channels();
1019            let keep_alive_timeout = tokio::time::sleep(keep_alive);
1020            let event = tokio::select! {
1021                event = self.event_rx.recv() => {
1022                    if let Some(event) = event {
1023                        InnerEvent::FromHttpService(event)
1024                    } else {
1025                        return Err(WorkerQuitReason::fatal(LocalSessionWorkerError::TransportTerminated, "waiting next session event"))
1026                    }
1027                },
1028                from_handler = context.recv_from_handler() => {
1029                    InnerEvent::FromHandler(from_handler?)
1030                }
1031                _ = ct.cancelled() => {
1032                    return Err(WorkerQuitReason::Cancelled)
1033                }
1034                _ = keep_alive_timeout => {
1035                    return Err(WorkerQuitReason::IdleTimeout(keep_alive))
1036                }
1037            };
1038            match event {
1039                InnerEvent::FromHandler(WorkerSendRequest { message, responder }) => {
1040                    // catch response
1041                    let to_unregister = match &message {
1042                        crate::model::JsonRpcMessage::Response(json_rpc_response) => {
1043                            let request_id = json_rpc_response.id.clone();
1044                            Some(ResourceKey::McpRequestId(request_id))
1045                        }
1046                        crate::model::JsonRpcMessage::Error(json_rpc_error) => {
1047                            json_rpc_error.id.clone().map(ResourceKey::McpRequestId)
1048                        }
1049                        _ => {
1050                            None
1051                            // no need to unregister resource
1052                        }
1053                    };
1054                    let handle_result = self
1055                        .handle_server_message(message)
1056                        .await
1057                        .map_err(LocalSessionWorkerError::FailToHandleMessage);
1058                    let _ = responder.send(handle_result).inspect_err(|error| {
1059                        tracing::warn!(?error, "failed to send message to http service handler");
1060                    });
1061                    if let Some(to_unregister) = to_unregister {
1062                        self.unregister_resource(&to_unregister);
1063                    }
1064                }
1065                InnerEvent::FromHttpService(SessionEvent::ClientMessage {
1066                    message: json_rpc_message,
1067                    http_request_id,
1068                }) => {
1069                    match &json_rpc_message {
1070                        crate::model::JsonRpcMessage::Request(request) => {
1071                            if let Some(http_request_id) = http_request_id {
1072                                self.register_request(request, http_request_id)
1073                            }
1074                        }
1075                        crate::model::JsonRpcMessage::Notification(notification) => {
1076                            self.catch_cancellation_notification(notification)
1077                        }
1078                        _ => {}
1079                    }
1080                    context.send_to_handler(json_rpc_message).await?;
1081                }
1082                InnerEvent::FromHttpService(SessionEvent::EstablishRequestWiseChannel {
1083                    responder,
1084                }) => {
1085                    let handle_result = self.establish_request_wise_channel().await;
1086                    let _ = responder.send(handle_result);
1087                }
1088                InnerEvent::FromHttpService(SessionEvent::CloseRequestWiseChannel {
1089                    id,
1090                    responder,
1091                }) => {
1092                    let _handle_result = self.tx_router.remove(&id);
1093                    let _ = responder.send(Ok(()));
1094                }
1095                InnerEvent::FromHttpService(SessionEvent::Resume {
1096                    last_event_id,
1097                    responder,
1098                }) => {
1099                    let handle_result = self.resume(last_event_id).await;
1100                    let _ = responder.send(handle_result);
1101                }
1102                InnerEvent::FromHttpService(SessionEvent::Close) => {
1103                    return Err(WorkerQuitReason::TransportClosed);
1104                }
1105                InnerEvent::FromHttpService(SessionEvent::CloseSseStream {
1106                    http_request_id,
1107                    retry_interval,
1108                    responder,
1109                }) => {
1110                    let handle_result =
1111                        self.close_sse_stream(http_request_id, retry_interval).await;
1112                    let _ = responder.send(handle_result);
1113                }
1114                _ => {
1115                    // ignore
1116                }
1117            }
1118        }
1119    }
1120}
1121
1122#[derive(Debug, Clone)]
1123#[non_exhaustive]
1124pub struct SessionConfig {
1125    /// the capacity of the channel for the session. Default is 16.
1126    pub channel_capacity: usize,
1127    /// The session will be closed after this duration of inactivity.
1128    ///
1129    /// This serves as a safety net for cleaning up sessions whose HTTP
1130    /// connections have silently dropped (e.g., due to an HTTP/2
1131    /// `RST_STREAM`). Without a timeout, such sessions become zombies:
1132    /// the session worker keeps running indefinitely because the session
1133    /// handle's sender is still held in the session manager, preventing
1134    /// the worker's event channel from closing.
1135    ///
1136    /// Defaults to 5 minutes. Set to `None` to disable (not recommended
1137    /// for long-running servers behind proxies).
1138    pub keep_alive: Option<Duration>,
1139    /// SSE retry interval for priming events on request-wise streams.
1140    /// When set, the session layer prepends a priming event with the correct
1141    /// stream-identifying event ID to each request-wise SSE stream.
1142    /// Default is 3 seconds, matching `StreamableHttpServerConfig::default()`.
1143    pub sse_retry: Option<Duration>,
1144    /// How long to retain completed request-wise channel caches for late
1145    /// resume requests. After this duration, completed entries are evicted
1146    /// and resume will return an error. Default is 60 seconds.
1147    pub completed_cache_ttl: Duration,
1148    /// Maximum duration to wait for the `initialize` request after session
1149    /// creation. If not received within this window, the session is
1150    /// terminated. Default is 60 seconds. Set to `None` to disable.
1151    pub init_timeout: Option<Duration>,
1152}
1153
1154impl SessionConfig {
1155    pub const DEFAULT_CHANNEL_CAPACITY: usize = 16;
1156    pub const DEFAULT_KEEP_ALIVE: Duration = Duration::from_secs(300);
1157    pub const DEFAULT_SSE_RETRY: Duration = Duration::from_secs(3);
1158    pub const DEFAULT_COMPLETED_CACHE_TTL: Duration = Duration::from_secs(60);
1159    pub const DEFAULT_INIT_TIMEOUT: Duration = Duration::from_secs(60);
1160}
1161
1162impl Default for SessionConfig {
1163    fn default() -> Self {
1164        Self {
1165            channel_capacity: Self::DEFAULT_CHANNEL_CAPACITY,
1166            keep_alive: Some(Self::DEFAULT_KEEP_ALIVE),
1167            sse_retry: Some(Self::DEFAULT_SSE_RETRY),
1168            completed_cache_ttl: Self::DEFAULT_COMPLETED_CACHE_TTL,
1169            init_timeout: Some(Self::DEFAULT_INIT_TIMEOUT),
1170        }
1171    }
1172}
1173
1174/// Create a new session with the given id and configuration.
1175///
1176/// This function will return a pair of [`LocalSessionHandle`] and [`LocalSessionWorker`].
1177///
1178/// You can run the [`LocalSessionWorker`] as a transport for mcp server. And use the [`LocalSessionHandle`] operate the session.
1179pub fn create_local_session(
1180    id: impl Into<SessionId>,
1181    config: SessionConfig,
1182) -> (LocalSessionHandle, LocalSessionWorker) {
1183    let id = id.into();
1184    let (event_tx, event_rx) = tokio::sync::mpsc::channel(config.channel_capacity);
1185    let (common_tx, _) = tokio::sync::mpsc::channel(config.channel_capacity);
1186    let common = CachedTx::new_common(common_tx);
1187    tracing::info!(session_id = ?id, "create new session");
1188    let handle = LocalSessionHandle {
1189        event_tx,
1190        id: id.clone(),
1191    };
1192    let session_worker = LocalSessionWorker {
1193        next_http_request_id: 0,
1194        id,
1195        tx_router: HashMap::new(),
1196        resource_router: HashMap::new(),
1197        common,
1198        shadow_txs: Vec::new(),
1199        event_rx,
1200        session_config: config.clone(),
1201    };
1202    (handle, session_worker)
1203}
rmcp_soddygo/transport/streamable_http_server/session/local.rs

rmcp_soddygo/transport/streamable_http_server/session/
local.rs
