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 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 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 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 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#[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 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_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 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) = ¬ification.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 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 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 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 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 };
718 let (tx, rx) = tokio::sync::mpsc::channel(capacity);
719 if is_replacing_dead_primary {
720 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 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 Some(id) => {
760 let request_wise = self
761 .tx_router
762 .get_mut(&id)
763 .ok_or(SessionError::ChannelClosed(Some(id)))?;
764
765 if let Some(interval) = retry_interval {
767 request_wise.tx.send_priming(interval).await?;
768 }
769
770 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 None => {
782 if let Some(interval) = retry_interval {
784 self.common.send_priming(interval).await?;
785 }
786
787 let (tx, _rx) = tokio::sync::mpsc::channel(1);
789 self.common.tx = tx;
790
791 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 http_request_id: Option<HttpRequestId>,
827 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 event_tx: Sender<SessionEvent>,
851}
852
853impl LocalSessionHandle {
854 pub fn id(&self) -> &SessionId {
856 &self.id
857 }
858
859 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 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 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 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 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 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 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 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 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 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 }
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 }
1213 }
1214 }
1215 }
1216}
1217
1218#[derive(Debug, Clone)]
1219#[non_exhaustive]
1220pub struct SessionConfig {
1221 pub channel_capacity: usize,
1223 pub keep_alive: Option<Duration>,
1235 pub sse_retry: Option<Duration>,
1240 pub completed_cache_ttl: Duration,
1244 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
1270pub 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 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}