Skip to main content

rmcp/transport/streamable_http_server/session/
local.rs

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