1use std::collections::HashSet;
4use std::fmt;
5use std::sync::atomic::{AtomicU64, Ordering};
6use std::sync::{Arc, Mutex};
7
8use fastmcp_core::logging::{debug, targets, warn};
9use fastmcp_core::{McpContext, SessionState, Sha256Digest, sha256_bounded};
10use fastmcp_protocol::{
11 ClientCapabilities, ClientInfo, JsonRpcRequest, LogLevel, ResourceUpdatedNotificationParams,
12 ServerCapabilities, ServerInfo,
13};
14
15use crate::NotificationSender;
16
17static NEXT_SESSION_ID: AtomicU64 = AtomicU64::new(1);
18
19const fn checked_session_id_successor(current: u64) -> Option<u64> {
20 current.checked_add(1)
21}
22
23fn next_session_id() -> u64 {
24 NEXT_SESSION_ID
25 .try_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
26 checked_session_id_successor(current)
27 })
28 .unwrap_or_else(|_| {
29 panic!("process-local MCP session identity space exhausted")
34 })
35}
36
37pub(crate) const MAX_RESOURCE_SUBSCRIPTIONS_PER_SESSION: usize = 32;
39pub(crate) const MAX_RESOURCE_SUBSCRIPTION_BYTES_PER_SESSION: usize = 4 * 1024 * 1024;
41
42const RESOURCE_LOG_HASH_INPUT_LIMIT: usize = 4 * 1024;
44const RESOURCE_LOG_DIGEST_PREFIX_BYTES: usize = 8;
45
46#[derive(Clone, Copy)]
47struct SafeResourceLogLabel {
48 byte_len: usize,
49 hashed_bytes: usize,
50 digest_prefix: [u8; RESOURCE_LOG_DIGEST_PREFIX_BYTES],
51}
52
53impl fmt::Display for SafeResourceLogLabel {
54 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
55 write!(f, "bytes={},sha256_prefix=", self.byte_len)?;
56 for byte in self.digest_prefix {
57 write!(f, "{byte:02x}")?;
58 }
59 if self.hashed_bytes < self.byte_len {
60 write!(f, ",hashed_prefix_bytes={}", self.hashed_bytes)?;
61 }
62 Ok(())
63 }
64}
65
66fn safe_resource_log_label(uri: &str) -> SafeResourceLogLabel {
67 let bytes = uri.as_bytes();
68 let hashed_bytes = bytes.len().min(RESOURCE_LOG_HASH_INPUT_LIMIT);
69 let bounded_prefix = &bytes[..hashed_bytes];
70 let mut digest_prefix = [0_u8; RESOURCE_LOG_DIGEST_PREFIX_BYTES];
71 if let Ok(digest) = sha256_bounded(bounded_prefix, RESOURCE_LOG_HASH_INPUT_LIMIT) {
72 digest_prefix.copy_from_slice(&digest.as_bytes()[..RESOURCE_LOG_DIGEST_PREFIX_BYTES]);
73 }
74
75 SafeResourceLogLabel {
76 byte_len: bytes.len(),
77 hashed_bytes,
78 digest_prefix,
79 }
80}
81
82#[derive(Debug, Clone, Copy, PartialEq, Eq)]
84pub(crate) enum SubscriptionAdmission {
85 Accepted,
87 Duplicate,
89}
90
91#[derive(Debug, Clone, Copy, PartialEq, Eq)]
93pub(crate) enum SubscriptionAdmissionError {
94 CapacityExceeded,
96 RequestNotLive,
98}
99
100#[derive(Debug, Clone, Copy, PartialEq, Eq)]
102pub(crate) enum SubscriptionRemoval {
103 Removed,
105 NotSubscribed,
107}
108
109#[derive(Debug, Clone, Copy, PartialEq, Eq)]
111pub(crate) enum SubscriptionRemovalError {
112 RequestNotLive,
114}
115
116#[derive(Clone)]
123pub(crate) struct InitializationSnapshot {
124 initialized: bool,
125 client_info: Option<ClientInfo>,
126 client_capabilities: Option<ClientCapabilities>,
127 protocol_version: Option<String>,
128}
129
130#[derive(Clone, Default)]
133pub(crate) struct SessionPrincipalBinding {
134 fingerprint: Arc<Mutex<Option<Sha256Digest>>>,
135}
136
137impl fmt::Debug for SessionPrincipalBinding {
138 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
139 formatter
140 .debug_struct("SessionPrincipalBinding")
141 .field("bound", &self.is_bound_for_debug())
142 .finish()
143 }
144}
145
146impl SessionPrincipalBinding {
147 pub(crate) fn bind_or_verify(&self, fingerprint: Sha256Digest) -> bool {
150 let mut bound = self
151 .fingerprint
152 .lock()
153 .unwrap_or_else(std::sync::PoisonError::into_inner);
154 match *bound {
155 Some(existing) => existing == fingerprint,
156 None => {
157 *bound = Some(fingerprint);
158 true
159 }
160 }
161 }
162
163 pub(crate) fn verify_existing(&self, fingerprint: Sha256Digest) -> bool {
167 self.fingerprint
168 .lock()
169 .unwrap_or_else(std::sync::PoisonError::into_inner)
170 .is_some_and(|existing| existing == fingerprint)
171 }
172
173 fn is_bound_for_debug(&self) -> bool {
174 self.fingerprint
175 .lock()
176 .unwrap_or_else(std::sync::PoisonError::into_inner)
177 .is_some()
178 }
179}
180
181pub struct Session {
188 id: u64,
190 initialized: bool,
192 client_info: Option<ClientInfo>,
194 client_capabilities: Option<ClientCapabilities>,
196 server_info: ServerInfo,
198 server_capabilities: ServerCapabilities,
200 protocol_version: Option<String>,
202 resource_subscriptions: HashSet<String>,
204 resource_subscription_bytes: usize,
206 log_level: Option<LogLevel>,
208 state: SessionState,
210 principal_binding: SessionPrincipalBinding,
212}
213
214impl fmt::Debug for Session {
215 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
216 f.debug_struct("Session")
217 .field("initialized", &self.initialized)
218 .field("has_client_info", &self.client_info.is_some())
219 .field(
220 "has_client_capabilities",
221 &self.client_capabilities.is_some(),
222 )
223 .field("has_protocol_version", &self.protocol_version.is_some())
224 .field(
225 "resource_subscription_count",
226 &self.resource_subscriptions.len(),
227 )
228 .field(
229 "resource_subscription_bytes",
230 &self.resource_subscription_bytes,
231 )
232 .field("has_log_level", &self.log_level.is_some())
233 .field("state_entry_count", &self.state.len())
234 .field(
235 "principal_bound",
236 &self.principal_binding.is_bound_for_debug(),
237 )
238 .finish()
239 }
240}
241
242impl Session {
243 #[must_use]
245 pub fn new(server_info: ServerInfo, server_capabilities: ServerCapabilities) -> Self {
246 Self {
247 id: next_session_id(),
248 initialized: false,
249 client_info: None,
250 client_capabilities: None,
251 server_info,
252 server_capabilities,
253 protocol_version: None,
254 resource_subscriptions: HashSet::new(),
255 resource_subscription_bytes: 0,
256 log_level: None,
257 state: SessionState::new(),
258 principal_binding: SessionPrincipalBinding::default(),
259 }
260 }
261
262 #[must_use]
264 pub(crate) fn id(&self) -> u64 {
265 self.id
266 }
267
268 #[must_use]
273 pub fn state(&self) -> &SessionState {
274 &self.state
275 }
276
277 pub(crate) fn principal_binding(&self) -> SessionPrincipalBinding {
279 self.principal_binding.clone()
280 }
281
282 #[must_use]
284 pub fn is_initialized(&self) -> bool {
285 self.initialized
286 }
287
288 pub fn initialize(
290 &mut self,
291 client_info: ClientInfo,
292 client_capabilities: ClientCapabilities,
293 protocol_version: String,
294 ) {
295 self.client_info = Some(client_info);
296 self.client_capabilities = Some(client_capabilities);
297 self.protocol_version = Some(protocol_version);
298 self.initialized = true;
299 }
300
301 pub(crate) fn initialization_snapshot(&self) -> InitializationSnapshot {
303 InitializationSnapshot {
304 initialized: self.initialized,
305 client_info: self.client_info.clone(),
306 client_capabilities: self.client_capabilities.clone(),
307 protocol_version: self.protocol_version.clone(),
308 }
309 }
310
311 pub(crate) fn restore_initialization(&mut self, snapshot: InitializationSnapshot) {
313 self.initialized = snapshot.initialized;
314 self.client_info = snapshot.client_info;
315 self.client_capabilities = snapshot.client_capabilities;
316 self.protocol_version = snapshot.protocol_version;
317 }
318
319 #[must_use]
321 pub fn client_info(&self) -> Option<&ClientInfo> {
322 self.client_info.as_ref()
323 }
324
325 #[must_use]
327 pub fn client_capabilities(&self) -> Option<&ClientCapabilities> {
328 self.client_capabilities.as_ref()
329 }
330
331 #[must_use]
333 pub fn server_info(&self) -> &ServerInfo {
334 &self.server_info
335 }
336
337 #[must_use]
339 pub fn server_capabilities(&self) -> &ServerCapabilities {
340 &self.server_capabilities
341 }
342
343 #[must_use]
345 pub fn protocol_version(&self) -> Option<&str> {
346 self.protocol_version.as_deref()
347 }
348
349 pub(crate) fn subscribe_resource(
356 &mut self,
357 ctx: &McpContext,
358 uri: String,
359 ) -> Result<SubscriptionAdmission, SubscriptionAdmissionError> {
360 if ctx.ensure_live().is_err() {
361 return Err(SubscriptionAdmissionError::RequestNotLive);
362 }
363
364 if uri.len() > MAX_RESOURCE_SUBSCRIPTION_BYTES_PER_SESSION {
368 warn!(
369 target: targets::SESSION,
370 "resource subscription rejected by session limits; subscription_count={}; subscription_count_limit={}; retained_uri_bytes={}; retained_uri_bytes_limit={}; candidate_uri={}",
371 self.resource_subscriptions.len(),
372 MAX_RESOURCE_SUBSCRIPTIONS_PER_SESSION,
373 self.resource_subscription_bytes,
374 MAX_RESOURCE_SUBSCRIPTION_BYTES_PER_SESSION,
375 safe_resource_log_label(&uri),
376 );
377 return Err(SubscriptionAdmissionError::CapacityExceeded);
378 }
379
380 if self.resource_subscriptions.contains(&uri) {
381 return Ok(SubscriptionAdmission::Duplicate);
382 }
383
384 let prospective_bytes = self.resource_subscription_bytes.checked_add(uri.len());
385 if self.resource_subscriptions.len() >= MAX_RESOURCE_SUBSCRIPTIONS_PER_SESSION
386 || prospective_bytes
387 .is_none_or(|bytes| bytes > MAX_RESOURCE_SUBSCRIPTION_BYTES_PER_SESSION)
388 {
389 warn!(
390 target: targets::SESSION,
391 "resource subscription rejected by session limits; subscription_count={}; subscription_count_limit={}; retained_uri_bytes={}; retained_uri_bytes_limit={}; candidate_uri={}",
392 self.resource_subscriptions.len(),
393 MAX_RESOURCE_SUBSCRIPTIONS_PER_SESSION,
394 self.resource_subscription_bytes,
395 MAX_RESOURCE_SUBSCRIPTION_BYTES_PER_SESSION,
396 safe_resource_log_label(&uri),
397 );
398 return Err(SubscriptionAdmissionError::CapacityExceeded);
399 }
400
401 let Some(prospective_bytes) = prospective_bytes else {
402 return Err(SubscriptionAdmissionError::CapacityExceeded);
403 };
404
405 if self.resource_subscriptions.try_reserve(1).is_err() {
409 warn!(
410 target: targets::SESSION,
411 "resource subscription rejected because bounded storage could not be reserved; subscription_count={}; subscription_count_limit={}; retained_uri_bytes={}; retained_uri_bytes_limit={}; candidate_uri={}",
412 self.resource_subscriptions.len(),
413 MAX_RESOURCE_SUBSCRIPTIONS_PER_SESSION,
414 self.resource_subscription_bytes,
415 MAX_RESOURCE_SUBSCRIPTION_BYTES_PER_SESSION,
416 safe_resource_log_label(&uri),
417 );
418 return Err(SubscriptionAdmissionError::CapacityExceeded);
419 }
420
421 if ctx.ensure_live().is_err() {
424 return Err(SubscriptionAdmissionError::RequestNotLive);
425 }
426
427 if self.resource_subscriptions.insert(uri) {
428 self.resource_subscription_bytes = prospective_bytes;
429 Ok(SubscriptionAdmission::Accepted)
430 } else {
431 Ok(SubscriptionAdmission::Duplicate)
432 }
433 }
434
435 pub(crate) fn unsubscribe_resource(
437 &mut self,
438 ctx: &McpContext,
439 uri: &str,
440 ) -> Result<SubscriptionRemoval, SubscriptionRemovalError> {
441 self.unsubscribe_resource_with_precommit(ctx, uri, || {})
442 }
443
444 fn unsubscribe_resource_with_precommit(
445 &mut self,
446 ctx: &McpContext,
447 uri: &str,
448 before_commit: impl FnOnce(),
449 ) -> Result<SubscriptionRemoval, SubscriptionRemovalError> {
450 if ctx.ensure_live().is_err() {
451 return Err(SubscriptionRemovalError::RequestNotLive);
452 }
453
454 if uri.len() > MAX_RESOURCE_SUBSCRIPTION_BYTES_PER_SESSION {
457 return Ok(SubscriptionRemoval::NotSubscribed);
458 }
459
460 let subscribed = self.resource_subscriptions.contains(uri);
461 before_commit();
462
463 if ctx.ensure_live().is_err() {
467 return Err(SubscriptionRemovalError::RequestNotLive);
468 }
469
470 if !subscribed {
471 return Ok(SubscriptionRemoval::NotSubscribed);
472 }
473
474 if let Some(removed) = self.resource_subscriptions.take(uri) {
475 self.resource_subscription_bytes = self
476 .resource_subscription_bytes
477 .saturating_sub(removed.len());
478 Ok(SubscriptionRemoval::Removed)
479 } else {
480 Ok(SubscriptionRemoval::NotSubscribed)
481 }
482 }
483
484 pub(crate) fn rollback_resource_subscription(&mut self, uri: &str) {
486 if let Some(removed) = self.resource_subscriptions.take(uri) {
487 self.resource_subscription_bytes = self
488 .resource_subscription_bytes
489 .saturating_sub(removed.len());
490 }
491 }
492
493 pub(crate) fn restore_resource_subscription(&mut self, uri: String) {
495 let uri_len = uri.len();
496 if self.resource_subscriptions.insert(uri) {
497 self.resource_subscription_bytes = self
498 .resource_subscription_bytes
499 .checked_add(uri_len)
500 .expect("restoring a previously retained subscription cannot overflow");
501 }
502 debug_assert!(self.resource_subscriptions.len() <= MAX_RESOURCE_SUBSCRIPTIONS_PER_SESSION);
503 debug_assert!(
504 self.resource_subscription_bytes <= MAX_RESOURCE_SUBSCRIPTION_BYTES_PER_SESSION
505 );
506 }
507
508 #[must_use]
510 pub fn is_resource_subscribed(&self, uri: &str) -> bool {
511 self.resource_subscriptions.contains(uri)
512 }
513
514 #[must_use]
516 pub fn subscribed_resource_uris(&self) -> impl Iterator<Item = &str> {
517 self.resource_subscriptions.iter().map(String::as_str)
518 }
519
520 pub fn set_log_level(&mut self, level: LogLevel) {
522 self.log_level = Some(level);
523 }
524
525 pub(crate) fn restore_log_level(&mut self, level: Option<LogLevel>) {
527 self.log_level = level;
528 }
529
530 #[must_use]
532 pub fn log_level(&self) -> Option<LogLevel> {
533 self.log_level
534 }
535
536 #[must_use]
538 pub fn supports_sampling(&self) -> bool {
539 self.client_capabilities
540 .as_ref()
541 .is_some_and(|caps| caps.sampling.is_some())
542 }
543
544 #[must_use]
546 pub fn supports_elicitation(&self) -> bool {
547 self.client_capabilities
548 .as_ref()
549 .is_some_and(|caps| caps.elicitation.is_some())
550 }
551
552 #[must_use]
554 pub fn supports_roots(&self) -> bool {
555 self.client_capabilities
556 .as_ref()
557 .is_some_and(|caps| caps.roots.is_some())
558 }
559
560 pub fn notify_resource_updated(&self, uri: &str, sender: &NotificationSender) -> bool {
564 if !self.is_resource_subscribed(uri) {
565 return false;
566 }
567
568 let safe_uri = safe_resource_log_label(uri);
569 let params = ResourceUpdatedNotificationParams {
570 uri: uri.to_string(),
571 };
572 let payload = match serde_json::to_value(params) {
573 Ok(value) => value,
574 Err(err) => {
575 warn!(
576 target: targets::SESSION,
577 "failed to serialize resource update; uri={}: {}",
578 safe_uri,
579 err
580 );
581 return false;
582 }
583 };
584
585 debug!(
586 target: targets::SESSION,
587 "sending resource update notification; uri={}",
588 safe_uri
589 );
590 let notification =
591 JsonRpcRequest::notification("notifications/resources/updated", Some(payload));
592 if crate::catch_extension_unwind(|| sender(notification)).is_err() {
593 warn!(
594 target: targets::SESSION,
595 "resource update notification sender terminated unexpectedly; uri={}; detail=panic_payload_redacted",
596 safe_uri
597 );
598 return false;
599 }
600 true
601 }
602}
603
604#[cfg(test)]
605mod tests {
606 use super::*;
607 use asupersync::Cx;
608 use fastmcp_protocol::{ElicitationCapability, RootsCapability, SamplingCapability};
609 use std::sync::{Arc, Mutex};
610
611 fn make_server_info() -> ServerInfo {
612 ServerInfo {
613 name: "test".to_string(),
614 version: "1.0".to_string(),
615 }
616 }
617
618 fn make_client_info() -> ClientInfo {
619 ClientInfo {
620 name: "test-client".to_string(),
621 version: "1.0".to_string(),
622 }
623 }
624
625 fn make_session() -> Session {
626 Session::new(make_server_info(), ServerCapabilities::default())
627 }
628
629 fn live_context() -> McpContext {
630 McpContext::new(Cx::for_testing(), 1)
631 }
632
633 fn admit_resource(session: &mut Session, uri: impl Into<String>) -> SubscriptionAdmission {
634 session
635 .subscribe_resource(&live_context(), uri.into())
636 .expect("test subscription should be admitted")
637 }
638
639 fn remove_resource(session: &mut Session, uri: &str) -> SubscriptionRemoval {
640 session
641 .unsubscribe_resource(&live_context(), uri)
642 .expect("live test unsubscription should succeed")
643 }
644
645 #[test]
648 fn new_session_is_not_initialized() {
649 let session = make_session();
650 assert!(!session.is_initialized());
651 }
652
653 #[test]
654 fn session_id_successor_rejects_wraparound() {
655 assert_eq!(checked_session_id_successor(1), Some(2));
656 assert_eq!(checked_session_id_successor(u64::MAX - 1), Some(u64::MAX));
657 assert_eq!(checked_session_id_successor(u64::MAX), None);
658 }
659
660 #[test]
661 fn new_session_has_no_client_info() {
662 let session = make_session();
663 assert!(session.client_info().is_none());
664 }
665
666 #[test]
667 fn new_session_has_no_client_capabilities() {
668 let session = make_session();
669 assert!(session.client_capabilities().is_none());
670 }
671
672 #[test]
673 fn new_session_has_no_protocol_version() {
674 let session = make_session();
675 assert!(session.protocol_version().is_none());
676 }
677
678 #[test]
679 fn new_session_has_no_log_level() {
680 let session = make_session();
681 assert!(session.log_level().is_none());
682 }
683
684 #[test]
685 fn new_session_returns_server_info() {
686 let session = make_session();
687 assert_eq!(session.server_info().name, "test");
688 assert_eq!(session.server_info().version, "1.0");
689 }
690
691 #[test]
692 fn new_session_returns_server_capabilities() {
693 let caps = ServerCapabilities::default();
694 let session = Session::new(make_server_info(), caps);
695 let _ = session.server_capabilities();
697 }
698
699 #[test]
702 fn initialize_sets_initialized_flag() {
703 let mut session = make_session();
704 session.initialize(
705 make_client_info(),
706 ClientCapabilities::default(),
707 "2024-11-05".to_string(),
708 );
709 assert!(session.is_initialized());
710 }
711
712 #[test]
713 fn initialize_stores_client_info() {
714 let mut session = make_session();
715 session.initialize(
716 make_client_info(),
717 ClientCapabilities::default(),
718 "2024-11-05".to_string(),
719 );
720 let info = session.client_info().expect("client_info set");
721 assert_eq!(info.name, "test-client");
722 assert_eq!(info.version, "1.0");
723 }
724
725 #[test]
726 fn initialize_stores_client_capabilities() {
727 let mut session = make_session();
728 let caps = ClientCapabilities {
729 sampling: Some(SamplingCapability {}),
730 elicitation: None,
731 roots: None,
732 };
733 session.initialize(make_client_info(), caps, "2024-11-05".to_string());
734 let stored = session.client_capabilities().expect("caps set");
735 assert!(stored.sampling.is_some());
736 }
737
738 #[test]
739 fn initialize_stores_protocol_version() {
740 let mut session = make_session();
741 session.initialize(
742 make_client_info(),
743 ClientCapabilities::default(),
744 "2025-03-26".to_string(),
745 );
746 assert_eq!(session.protocol_version(), Some("2025-03-26"));
747 }
748
749 #[test]
752 fn subscribe_and_check_resource() {
753 let mut session = make_session();
754 assert!(!session.is_resource_subscribed("file:///a.txt"));
755 assert_eq!(
756 admit_resource(&mut session, "file:///a.txt"),
757 SubscriptionAdmission::Accepted
758 );
759 assert!(session.is_resource_subscribed("file:///a.txt"));
760 }
761
762 #[test]
763 fn unsubscribe_resource_removes_it() {
764 let mut session = make_session();
765 admit_resource(&mut session, "file:///a.txt");
766 assert_eq!(
767 remove_resource(&mut session, "file:///a.txt"),
768 SubscriptionRemoval::Removed
769 );
770 assert!(!session.is_resource_subscribed("file:///a.txt"));
771 }
772
773 #[test]
774 fn unsubscribe_nonexistent_resource_is_noop() {
775 let mut session = make_session();
776 assert_eq!(
777 remove_resource(&mut session, "file:///does-not-exist"),
778 SubscriptionRemoval::NotSubscribed
779 );
780 assert!(!session.is_resource_subscribed("file:///does-not-exist"));
781 }
782
783 #[test]
784 fn cancellation_during_unsubscribe_lookup_cannot_commit_removal() {
785 let mut session = make_session();
786 let uri = "resource://retained";
787 admit_resource(&mut session, uri);
788 let raw_cx = Cx::for_testing();
789 let ctx = McpContext::new(raw_cx.clone(), 44);
790
791 let result = session.unsubscribe_resource_with_precommit(&ctx, uri, || {
792 raw_cx.set_cancel_requested(true);
793 });
794
795 assert_eq!(result, Err(SubscriptionRemovalError::RequestNotLive));
796 assert!(session.is_resource_subscribed(uri));
797 assert_eq!(session.resource_subscription_bytes, uri.len());
798 }
799
800 #[test]
801 fn multiple_subscriptions_are_independent() {
802 let mut session = make_session();
803 admit_resource(&mut session, "a://1");
804 admit_resource(&mut session, "b://2");
805 assert!(session.is_resource_subscribed("a://1"));
806 assert!(session.is_resource_subscribed("b://2"));
807 remove_resource(&mut session, "a://1");
808 assert!(!session.is_resource_subscribed("a://1"));
809 assert!(session.is_resource_subscribed("b://2"));
810 }
811
812 #[test]
813 fn duplicate_subscribe_is_idempotent() {
814 let mut session = make_session();
815 assert_eq!(
816 admit_resource(&mut session, "r://x"),
817 SubscriptionAdmission::Accepted
818 );
819 let retained_bytes = session.resource_subscription_bytes;
820 assert_eq!(
821 admit_resource(&mut session, "r://x"),
822 SubscriptionAdmission::Duplicate
823 );
824 assert!(session.is_resource_subscribed("r://x"));
825 assert_eq!(session.resource_subscriptions.len(), 1);
826 assert_eq!(session.resource_subscription_bytes, retained_bytes);
827 remove_resource(&mut session, "r://x");
828 assert!(!session.is_resource_subscribed("r://x"));
829 assert_eq!(session.resource_subscription_bytes, 0);
830 }
831
832 #[test]
833 fn resource_subscription_count_admits_exact_cap_and_rejects_cap_plus_one() {
834 let mut session = make_session();
835 for index in 0..MAX_RESOURCE_SUBSCRIPTIONS_PER_SESSION {
836 assert_eq!(
837 admit_resource(&mut session, format!("test://subscription/{index}")),
838 SubscriptionAdmission::Accepted
839 );
840 }
841
842 assert_eq!(
843 session.resource_subscriptions.len(),
844 MAX_RESOURCE_SUBSCRIPTIONS_PER_SESSION
845 );
846 let retained_bytes = session.resource_subscription_bytes;
847 let over_limit_uri = "test://subscription/over-limit";
848 let rejection = session.subscribe_resource(&live_context(), over_limit_uri.to_string());
849 assert_eq!(rejection, Err(SubscriptionAdmissionError::CapacityExceeded));
850 assert!(!session.is_resource_subscribed(over_limit_uri));
851 assert_eq!(session.resource_subscription_bytes, retained_bytes);
852
853 remove_resource(&mut session, "test://subscription/0");
854 assert_eq!(
855 admit_resource(&mut session, over_limit_uri),
856 SubscriptionAdmission::Accepted
857 );
858 assert!(session.is_resource_subscribed(over_limit_uri));
859 assert_eq!(
860 session.resource_subscriptions.len(),
861 MAX_RESOURCE_SUBSCRIPTIONS_PER_SESSION
862 );
863 }
864
865 #[test]
866 fn resource_subscription_bytes_admit_exact_cap_reject_cap_plus_one_and_release() {
867 let mut session = make_session();
868 let exact_limit_uri = "é".repeat(MAX_RESOURCE_SUBSCRIPTION_BYTES_PER_SESSION / 2);
869 assert_eq!(
870 exact_limit_uri.len(),
871 MAX_RESOURCE_SUBSCRIPTION_BYTES_PER_SESSION
872 );
873
874 assert_eq!(
875 admit_resource(&mut session, exact_limit_uri.clone()),
876 SubscriptionAdmission::Accepted
877 );
878 assert!(session.is_resource_subscribed(&exact_limit_uri));
879 assert_eq!(
880 session.resource_subscription_bytes,
881 MAX_RESOURCE_SUBSCRIPTION_BYTES_PER_SESSION
882 );
883 assert_eq!(
884 admit_resource(&mut session, exact_limit_uri.clone()),
885 SubscriptionAdmission::Duplicate
886 );
887 assert_eq!(
888 session.resource_subscription_bytes,
889 MAX_RESOURCE_SUBSCRIPTION_BYTES_PER_SESSION
890 );
891
892 remove_resource(&mut session, &exact_limit_uri);
893 assert_eq!(session.resource_subscription_bytes, 0);
894 assert!(session.resource_subscriptions.is_empty());
895
896 let over_limit_uri = format!("{exact_limit_uri}x");
897 assert_eq!(
898 over_limit_uri.len(),
899 MAX_RESOURCE_SUBSCRIPTION_BYTES_PER_SESSION + 1
900 );
901 let rejection = session.subscribe_resource(&live_context(), over_limit_uri.clone());
902 assert_eq!(rejection, Err(SubscriptionAdmissionError::CapacityExceeded));
903 assert!(!session.is_resource_subscribed(&over_limit_uri));
904 assert_eq!(session.resource_subscription_bytes, 0);
905 assert_eq!(
906 remove_resource(&mut session, &over_limit_uri),
907 SubscriptionRemoval::NotSubscribed
908 );
909 }
910
911 #[test]
912 fn non_live_request_cannot_admit_or_duplicate_a_subscription() {
913 let mut session = make_session();
914 admit_resource(&mut session, "resource://retained");
915
916 let cancelled_cx = Cx::for_testing();
917 cancelled_cx.set_cancel_requested(true);
918 let cancelled = McpContext::new(cancelled_cx, 2);
919
920 assert_eq!(
921 session.subscribe_resource(&cancelled, "resource://new".to_string()),
922 Err(SubscriptionAdmissionError::RequestNotLive)
923 );
924 assert_eq!(
925 session.subscribe_resource(&cancelled, "resource://retained".to_string()),
926 Err(SubscriptionAdmissionError::RequestNotLive)
927 );
928 assert_eq!(
929 session.unsubscribe_resource(&cancelled, "resource://retained"),
930 Err(SubscriptionRemovalError::RequestNotLive)
931 );
932 assert!(!session.is_resource_subscribed("resource://new"));
933 assert!(session.is_resource_subscribed("resource://retained"));
934
935 let (expired, lease) = McpContext::new(Cx::for_testing(), 3)
936 .begin_request_scope()
937 .expect("fresh context should acquire its request lease");
938 drop(lease);
939 assert_eq!(
940 session.subscribe_resource(&expired, "resource://expired".to_string()),
941 Err(SubscriptionAdmissionError::RequestNotLive)
942 );
943 assert_eq!(
944 session.unsubscribe_resource(&expired, "resource://retained"),
945 Err(SubscriptionRemovalError::RequestNotLive)
946 );
947 assert!(!session.is_resource_subscribed("resource://expired"));
948 assert!(session.is_resource_subscribed("resource://retained"));
949 }
950
951 #[test]
954 fn set_log_level_and_read_back() {
955 let mut session = make_session();
956 session.set_log_level(LogLevel::Warning);
957 assert_eq!(session.log_level(), Some(LogLevel::Warning));
958 }
959
960 #[test]
961 fn set_log_level_overwrites_previous() {
962 let mut session = make_session();
963 session.set_log_level(LogLevel::Debug);
964 session.set_log_level(LogLevel::Error);
965 assert_eq!(session.log_level(), Some(LogLevel::Error));
966 }
967
968 #[test]
971 fn state_is_accessible() {
972 let session = make_session();
973 let state = session.state();
974 let val: Option<String> = state.get("key");
976 assert!(val.is_none());
977 }
978
979 #[test]
982 fn notify_resource_updated_returns_false_when_not_subscribed() {
983 let session = make_session();
984 let sender: NotificationSender = Arc::new(|_| {});
985 assert!(!session.notify_resource_updated("file:///a.txt", &sender));
986 }
987
988 #[test]
989 fn notify_resource_updated_sends_when_subscribed() {
990 let mut session = make_session();
991 admit_resource(&mut session, "file:///a.txt");
992
993 let sent = Arc::new(Mutex::new(Vec::new()));
994 let sent_clone = Arc::clone(&sent);
995 let sender: NotificationSender = Arc::new(move |req| {
996 sent_clone.lock().unwrap().push(req);
997 });
998
999 let result = session.notify_resource_updated("file:///a.txt", &sender);
1000 assert!(result);
1001
1002 let messages = sent.lock().unwrap();
1003 assert_eq!(messages.len(), 1);
1004 assert_eq!(messages[0].method, "notifications/resources/updated");
1005 }
1006
1007 #[test]
1008 fn notify_resource_updated_includes_uri_in_params() {
1009 let mut session = make_session();
1010 admit_resource(&mut session, "test://res");
1011
1012 let sent = Arc::new(Mutex::new(Vec::new()));
1013 let sent_clone = Arc::clone(&sent);
1014 let sender: NotificationSender = Arc::new(move |req| {
1015 sent_clone.lock().unwrap().push(req);
1016 });
1017
1018 session.notify_resource_updated("test://res", &sender);
1019
1020 let messages = sent.lock().unwrap();
1021 let params = messages[0].params.as_ref().expect("params present");
1022 let uri = params
1023 .get("uri")
1024 .and_then(|v| v.as_str())
1025 .expect("uri field");
1026 assert_eq!(uri, "test://res");
1027 }
1028
1029 #[test]
1030 fn notify_resource_updated_does_not_fire_for_other_uri() {
1031 let mut session = make_session();
1032 admit_resource(&mut session, "file:///a.txt");
1033
1034 let sent = Arc::new(Mutex::new(Vec::new()));
1035 let sent_clone = Arc::clone(&sent);
1036 let sender: NotificationSender = Arc::new(move |req| {
1037 sent_clone.lock().unwrap().push(req);
1038 });
1039
1040 let result = session.notify_resource_updated("file:///b.txt", &sender);
1041 assert!(!result);
1042 assert!(sent.lock().unwrap().is_empty());
1043 }
1044
1045 #[test]
1046 fn resource_update_log_label_is_bounded_and_redacts_uri_text() {
1047 const URI_CANARY: &str = "private://customer-pii-canary@example.invalid/secret-token";
1048 let metadata = format!("{}", safe_resource_log_label(URI_CANARY));
1049
1050 assert!(metadata.contains(&format!("bytes={}", URI_CANARY.len())));
1051 assert!(metadata.contains("sha256_prefix="));
1052 assert!(!metadata.contains(URI_CANARY));
1053
1054 let long_uri = "x".repeat(RESOURCE_LOG_HASH_INPUT_LIMIT + 1);
1055 let bounded_metadata = format!("{}", safe_resource_log_label(&long_uri));
1056 assert!(bounded_metadata.contains(&format!(
1057 "hashed_prefix_bytes={RESOURCE_LOG_HASH_INPUT_LIMIT}"
1058 )));
1059 assert!(!bounded_metadata.contains(&long_uri));
1060 }
1061
1062 #[test]
1063 fn notify_resource_updated_catches_sender_panic_and_returns_false() {
1064 let mut session = make_session();
1065 admit_resource(&mut session, "test://subscribed");
1066 let sender: NotificationSender = Arc::new(|_| {
1067 panic!("notification-sender-panic-payload-canary-2d941");
1068 });
1069
1070 assert!(!session.notify_resource_updated("test://subscribed", &sender));
1071 assert!(session.is_resource_subscribed("test://subscribed"));
1072 }
1073
1074 #[test]
1077 fn session_debug_format_includes_fields() {
1078 let session = make_session();
1079 let debug = format!("{:?}", session);
1080 assert!(debug.contains("Session"));
1081 assert!(debug.contains("initialized: false"));
1082 }
1083
1084 #[test]
1085 fn session_debug_reports_only_safe_counts_and_booleans() {
1086 const SERVER_NAME_CANARY: &str = "server-tenant-canary-85fd";
1087 const SERVER_VERSION_CANARY: &str = "server-version-canary-12aa";
1088 const CLIENT_NAME_CANARY: &str = "customer-pii-canary@example.invalid";
1089 const CLIENT_VERSION_CANARY: &str = "client-version-canary-53bc";
1090 const PROTOCOL_CANARY: &str = "protocol-canary-secret-71f2";
1091 const URI_CANARY: &str = "private://subscription-uri-canary-2bd4";
1092 const STATE_KEY_CANARY: &str = "fastmcp.auth.canary-key-59cf";
1093 const STATE_VALUE_CANARY: &str = "Bearer auth-token-canary-b300";
1094
1095 let mut session = Session::new(
1096 ServerInfo {
1097 name: SERVER_NAME_CANARY.to_string(),
1098 version: SERVER_VERSION_CANARY.to_string(),
1099 },
1100 ServerCapabilities::default(),
1101 );
1102 session.initialize(
1103 ClientInfo {
1104 name: CLIENT_NAME_CANARY.to_string(),
1105 version: CLIENT_VERSION_CANARY.to_string(),
1106 },
1107 ClientCapabilities::default(),
1108 PROTOCOL_CANARY.to_string(),
1109 );
1110 admit_resource(&mut session, URI_CANARY);
1111 assert!(session.state().set(STATE_KEY_CANARY, STATE_VALUE_CANARY));
1112 session.set_log_level(LogLevel::Debug);
1113
1114 let debug = format!("{session:?}");
1115 assert!(debug.contains("Session"));
1116 assert!(debug.contains("initialized: true"));
1117 assert!(debug.contains("resource_subscription_count: 1"));
1118 assert!(debug.contains("state_entry_count: 1"));
1119 for canary in [
1120 SERVER_NAME_CANARY,
1121 SERVER_VERSION_CANARY,
1122 CLIENT_NAME_CANARY,
1123 CLIENT_VERSION_CANARY,
1124 PROTOCOL_CANARY,
1125 URI_CANARY,
1126 STATE_KEY_CANARY,
1127 STATE_VALUE_CANARY,
1128 ] {
1129 assert!(
1130 !debug.contains(canary),
1131 "Session Debug leaked canary: {canary}"
1132 );
1133 }
1134 }
1135
1136 #[test]
1139 fn test_session_supports_sampling() {
1140 let mut session = Session::new(
1141 ServerInfo {
1142 name: "test".to_string(),
1143 version: "1.0".to_string(),
1144 },
1145 ServerCapabilities::default(),
1146 );
1147
1148 assert!(!session.supports_sampling());
1150
1151 session.initialize(
1153 ClientInfo {
1154 name: "test-client".to_string(),
1155 version: "1.0".to_string(),
1156 },
1157 ClientCapabilities {
1158 sampling: Some(SamplingCapability {}),
1159 elicitation: None,
1160 roots: None,
1161 },
1162 "2024-11-05".to_string(),
1163 );
1164
1165 assert!(session.supports_sampling());
1166 assert!(!session.supports_elicitation());
1167 assert!(!session.supports_roots());
1168 }
1169
1170 #[test]
1171 fn test_session_supports_elicitation() {
1172 let mut session = Session::new(
1173 ServerInfo {
1174 name: "test".to_string(),
1175 version: "1.0".to_string(),
1176 },
1177 ServerCapabilities::default(),
1178 );
1179
1180 session.initialize(
1181 ClientInfo {
1182 name: "test-client".to_string(),
1183 version: "1.0".to_string(),
1184 },
1185 ClientCapabilities {
1186 sampling: None,
1187 elicitation: Some(ElicitationCapability::form()),
1188 roots: None,
1189 },
1190 "2024-11-05".to_string(),
1191 );
1192
1193 assert!(!session.supports_sampling());
1194 assert!(session.supports_elicitation());
1195 assert!(!session.supports_roots());
1196 }
1197
1198 #[test]
1199 fn test_session_supports_roots() {
1200 let mut session = Session::new(
1201 ServerInfo {
1202 name: "test".to_string(),
1203 version: "1.0".to_string(),
1204 },
1205 ServerCapabilities::default(),
1206 );
1207
1208 session.initialize(
1209 ClientInfo {
1210 name: "test-client".to_string(),
1211 version: "1.0".to_string(),
1212 },
1213 ClientCapabilities {
1214 sampling: None,
1215 elicitation: None,
1216 roots: Some(RootsCapability { list_changed: true }),
1217 },
1218 "2024-11-05".to_string(),
1219 );
1220
1221 assert!(!session.supports_sampling());
1222 assert!(!session.supports_elicitation());
1223 assert!(session.supports_roots());
1224 }
1225
1226 #[test]
1227 fn test_session_supports_all_capabilities() {
1228 let mut session = Session::new(
1229 ServerInfo {
1230 name: "test".to_string(),
1231 version: "1.0".to_string(),
1232 },
1233 ServerCapabilities::default(),
1234 );
1235
1236 session.initialize(
1237 ClientInfo {
1238 name: "test-client".to_string(),
1239 version: "1.0".to_string(),
1240 },
1241 ClientCapabilities {
1242 sampling: Some(SamplingCapability {}),
1243 elicitation: Some(ElicitationCapability::both()),
1244 roots: Some(RootsCapability {
1245 list_changed: false,
1246 }),
1247 },
1248 "2024-11-05".to_string(),
1249 );
1250
1251 assert!(session.supports_sampling());
1252 assert!(session.supports_elicitation());
1253 assert!(session.supports_roots());
1254 }
1255
1256 #[test]
1257 fn test_session_no_capabilities() {
1258 let mut session = Session::new(
1259 ServerInfo {
1260 name: "test".to_string(),
1261 version: "1.0".to_string(),
1262 },
1263 ServerCapabilities::default(),
1264 );
1265
1266 session.initialize(
1267 ClientInfo {
1268 name: "test-client".to_string(),
1269 version: "1.0".to_string(),
1270 },
1271 ClientCapabilities::default(),
1272 "2024-11-05".to_string(),
1273 );
1274
1275 assert!(!session.supports_sampling());
1276 assert!(!session.supports_elicitation());
1277 assert!(!session.supports_roots());
1278 }
1279
1280 #[test]
1283 fn reinitialize_overwrites_client_info() {
1284 let mut session = make_session();
1285 session.initialize(
1286 make_client_info(),
1287 ClientCapabilities::default(),
1288 "2024-11-05".to_string(),
1289 );
1290 session.initialize(
1291 ClientInfo {
1292 name: "new-client".to_string(),
1293 version: "2.0".to_string(),
1294 },
1295 ClientCapabilities {
1296 sampling: Some(SamplingCapability {}),
1297 elicitation: None,
1298 roots: None,
1299 },
1300 "2025-03-26".to_string(),
1301 );
1302 assert!(session.is_initialized());
1303 let info = session.client_info().unwrap();
1304 assert_eq!(info.name, "new-client");
1305 assert_eq!(info.version, "2.0");
1306 assert_eq!(session.protocol_version(), Some("2025-03-26"));
1307 assert!(session.supports_sampling());
1308 }
1309
1310 #[test]
1313 fn state_persists_after_initialization() {
1314 let mut session = make_session();
1315 session.state().set("key", "before_init");
1316 session.initialize(
1317 make_client_info(),
1318 ClientCapabilities::default(),
1319 "2024-11-05".to_string(),
1320 );
1321 let val: Option<String> = session.state().get("key");
1322 assert_eq!(val.as_deref(), Some("before_init"));
1323 }
1324
1325 #[test]
1328 fn notify_resource_updated_after_unsubscribe_returns_false() {
1329 let mut session = make_session();
1330 admit_resource(&mut session, "r://x");
1331
1332 let sent = Arc::new(Mutex::new(Vec::new()));
1333 let sent_clone = Arc::clone(&sent);
1334 let sender: NotificationSender = Arc::new(move |req| {
1335 sent_clone.lock().unwrap().push(req);
1336 });
1337
1338 assert!(session.notify_resource_updated("r://x", &sender));
1340 assert_eq!(sent.lock().unwrap().len(), 1);
1341
1342 remove_resource(&mut session, "r://x");
1344 assert!(!session.notify_resource_updated("r://x", &sender));
1345 assert_eq!(sent.lock().unwrap().len(), 1);
1347 }
1348
1349 #[test]
1352 fn resubscribe_after_unsubscribe_works() {
1353 let mut session = make_session();
1354 admit_resource(&mut session, "r://x");
1355 remove_resource(&mut session, "r://x");
1356 assert!(!session.is_resource_subscribed("r://x"));
1357 admit_resource(&mut session, "r://x");
1358 assert!(session.is_resource_subscribed("r://x"));
1359 }
1360
1361 #[test]
1364 fn session_debug_after_init_shows_initialized_true() {
1365 let mut session = make_session();
1366 session.initialize(
1367 make_client_info(),
1368 ClientCapabilities::default(),
1369 "2024-11-05".to_string(),
1370 );
1371 let debug = format!("{:?}", session);
1372 assert!(debug.contains("initialized: true"));
1373 }
1374
1375 #[test]
1378 fn session_with_custom_server_capabilities() {
1379 use fastmcp_protocol::{LoggingCapability, TasksCapability, ToolsCapability};
1380 let caps = ServerCapabilities {
1381 tools: Some(ToolsCapability { list_changed: true }),
1382 logging: Some(LoggingCapability {}),
1383 tasks: Some(TasksCapability {
1384 list_changed: false,
1385 }),
1386 ..ServerCapabilities::default()
1387 };
1388 let session = Session::new(make_server_info(), caps);
1389 assert!(session.server_capabilities().tools.is_some());
1390 assert!(session.server_capabilities().logging.is_some());
1391 assert!(session.server_capabilities().tasks.is_some());
1392 }
1393
1394 #[test]
1397 fn set_log_level_all_variants() {
1398 let mut session = make_session();
1399 for level in [
1400 LogLevel::Debug,
1401 LogLevel::Info,
1402 LogLevel::Warning,
1403 LogLevel::Error,
1404 ] {
1405 session.set_log_level(level);
1406 assert_eq!(session.log_level(), Some(level));
1407 }
1408 }
1409
1410 #[test]
1413 fn log_level_persists_across_reinitialization() {
1414 let mut session = make_session();
1415 session.set_log_level(LogLevel::Warning);
1416 session.initialize(
1417 make_client_info(),
1418 ClientCapabilities::default(),
1419 "2024-11-05".to_string(),
1420 );
1421 assert_eq!(session.log_level(), Some(LogLevel::Warning));
1422 session.initialize(
1424 ClientInfo {
1425 name: "other".to_string(),
1426 version: "2.0".to_string(),
1427 },
1428 ClientCapabilities::default(),
1429 "2025-03-26".to_string(),
1430 );
1431 assert_eq!(session.log_level(), Some(LogLevel::Warning));
1432 }
1433
1434 #[test]
1435 fn resource_subscriptions_persist_across_reinitialization() {
1436 let mut session = make_session();
1437 admit_resource(&mut session, "file:///keep.txt");
1438 session.initialize(
1439 make_client_info(),
1440 ClientCapabilities::default(),
1441 "2024-11-05".to_string(),
1442 );
1443 assert!(session.is_resource_subscribed("file:///keep.txt"));
1444 }
1445
1446 #[test]
1447 fn state_set_after_init_persists_through_reinit() {
1448 let mut session = make_session();
1449 session.initialize(
1450 make_client_info(),
1451 ClientCapabilities::default(),
1452 "2024-11-05".to_string(),
1453 );
1454 session.state().set("counter", 42);
1455 session.initialize(
1457 ClientInfo {
1458 name: "new".to_string(),
1459 version: "3.0".to_string(),
1460 },
1461 ClientCapabilities::default(),
1462 "2025-03-26".to_string(),
1463 );
1464 let val: Option<i32> = session.state().get("counter");
1465 assert_eq!(val, Some(42));
1466 }
1467
1468 #[test]
1469 fn notify_resource_updated_fires_independently_per_subscription() {
1470 let mut session = make_session();
1471 admit_resource(&mut session, "a://1");
1472 admit_resource(&mut session, "b://2");
1473
1474 let sent = Arc::new(Mutex::new(Vec::new()));
1475 let sent_clone = Arc::clone(&sent);
1476 let sender: NotificationSender = Arc::new(move |req| {
1477 sent_clone.lock().unwrap().push(req);
1478 });
1479
1480 assert!(session.notify_resource_updated("a://1", &sender));
1482 assert_eq!(sent.lock().unwrap().len(), 1);
1483 let uri = sent.lock().unwrap()[0]
1484 .params
1485 .as_ref()
1486 .unwrap()
1487 .get("uri")
1488 .unwrap()
1489 .as_str()
1490 .unwrap()
1491 .to_string();
1492 assert_eq!(uri, "a://1");
1493
1494 assert!(session.notify_resource_updated("b://2", &sender));
1496 assert_eq!(sent.lock().unwrap().len(), 2);
1497 let uri2 = sent.lock().unwrap()[1]
1498 .params
1499 .as_ref()
1500 .unwrap()
1501 .get("uri")
1502 .unwrap()
1503 .as_str()
1504 .unwrap()
1505 .to_string();
1506 assert_eq!(uri2, "b://2");
1507 }
1508
1509 #[test]
1510 fn supports_elicitation_and_roots_false_before_init() {
1511 let session = make_session();
1512 assert!(!session.supports_elicitation());
1513 assert!(!session.supports_roots());
1514 }
1515
1516 #[test]
1517 fn principal_binding_is_write_once_and_shared_across_clones() {
1518 let session = make_session();
1519 let binding = session.principal_binding();
1520 let clone = binding.clone();
1521 let alice = Sha256Digest::from_bytes([0x11; 32]);
1522 let bob = Sha256Digest::from_bytes([0x22; 32]);
1523
1524 assert!(!binding.is_bound_for_debug());
1525 assert!(!binding.verify_existing(alice));
1526 assert!(binding.bind_or_verify(alice));
1527 assert!(binding.verify_existing(alice));
1528 assert!(!binding.verify_existing(bob));
1529 assert!(clone.bind_or_verify(alice));
1530 assert!(!clone.bind_or_verify(bob));
1531 assert!(binding.bind_or_verify(alice));
1532 }
1533
1534 #[test]
1535 fn session_debug_reports_only_principal_presence() {
1536 let session = make_session();
1537 let canary = Sha256Digest::from_bytes([0xAB; 32]);
1538 assert!(session.principal_binding().bind_or_verify(canary));
1539
1540 let debug = format!("{session:?}");
1541 assert!(debug.contains("principal_bound: true"));
1542 assert!(!debug.contains("abab"));
1543 }
1544}