Skip to main content

fastmcp_server/
session.rs

1//! MCP session management.
2
3use 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            // `Session::new` is infallible, so there is no safe API-level way
30            // to report exhaustion. Reusing an ID would alias cancellation and
31            // request-ownership domains across live sessions; fail closed
32            // instead.
33            panic!("process-local MCP session identity space exhausted")
34        })
35}
36
37/// Default maximum number of retained resource subscriptions per session.
38pub(crate) const MAX_RESOURCE_SUBSCRIPTIONS_PER_SESSION: usize = 32;
39/// Default maximum aggregate UTF-8 bytes retained for resource subscription URIs.
40pub(crate) const MAX_RESOURCE_SUBSCRIPTION_BYTES_PER_SESSION: usize = 4 * 1024 * 1024;
41
42/// Maximum peer-controlled URI bytes hashed for one log correlation label.
43const 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/// Successful result of admitting a resource subscription.
83#[derive(Debug, Clone, Copy, PartialEq, Eq)]
84pub(crate) enum SubscriptionAdmission {
85    /// A new URI was retained by the session.
86    Accepted,
87    /// The URI was already retained; subscribe is idempotent.
88    Duplicate,
89}
90
91/// Fail-closed rejection of a resource-subscription admission.
92#[derive(Debug, Clone, Copy, PartialEq, Eq)]
93pub(crate) enum SubscriptionAdmissionError {
94    /// Retaining the URI would exceed a per-session bound.
95    CapacityExceeded,
96    /// The request context is cancelled, expired, or no longer owns a live lease.
97    RequestNotLive,
98}
99
100/// Successful result of an idempotent resource unsubscription.
101#[derive(Debug, Clone, Copy, PartialEq, Eq)]
102pub(crate) enum SubscriptionRemoval {
103    /// A retained URI was removed.
104    Removed,
105    /// The URI was not retained, so no mutation was necessary.
106    NotSubscribed,
107}
108
109/// Fail-closed rejection of a resource-unsubscription mutation.
110#[derive(Debug, Clone, Copy, PartialEq, Eq)]
111pub(crate) enum SubscriptionRemovalError {
112    /// The request context is cancelled, expired, or no longer owns a live lease.
113    RequestNotLive,
114}
115
116/// Exact initialization fields restored when an initialize response does not
117/// reach its commit boundary.
118///
119/// Session dispatch is serialized, so retaining one private snapshot while a
120/// response is finalized provides transaction-like rollback without exposing
121/// peer identity or capability contents outside the session module.
122#[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/// Shared, credential-free binding between a connection session and the
131/// verified principal admitted on its first authenticated frame.
132#[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    /// Binds an unclaimed session or verifies that a later request carries the
148    /// same principal. A mismatch fails closed without replacing the owner.
149    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    /// Verifies a control frame against an owner established by an ordinary
164    /// dispatched request. Control frames must never win an admission-order
165    /// race and claim an otherwise unbound session.
166    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
181/// An MCP session between client and server.
182///
183/// Tracks the state of an initialized MCP connection.
184/// Debug output is deliberately redacted to safe counts and presence booleans;
185/// it does not render client/server identity, protocol text, subscription URIs,
186/// or session-state contents.
187pub struct Session {
188    /// Process-local identity used to bind request ownership and cancellation.
189    id: u64,
190    /// Whether the session has been initialized.
191    initialized: bool,
192    /// Client info from initialization.
193    client_info: Option<ClientInfo>,
194    /// Client capabilities from initialization.
195    client_capabilities: Option<ClientCapabilities>,
196    /// Server info.
197    server_info: ServerInfo,
198    /// Server capabilities.
199    server_capabilities: ServerCapabilities,
200    /// Negotiated protocol version.
201    protocol_version: Option<String>,
202    /// Resource subscriptions for this session.
203    resource_subscriptions: HashSet<String>,
204    /// Aggregate UTF-8 bytes retained by unique resource subscription URIs.
205    resource_subscription_bytes: usize,
206    /// Session-scoped log level for log notifications.
207    log_level: Option<LogLevel>,
208    /// Per-session state storage.
209    state: SessionState,
210    /// Stable verified owner of session-scoped state and control operations.
211    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    /// Creates a new uninitialized session.
244    #[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    /// Returns the process-local identity of this connection session.
263    #[must_use]
264    pub(crate) fn id(&self) -> u64 {
265        self.id
266    }
267
268    /// Returns a reference to the session state.
269    ///
270    /// Session state persists across requests within this session and can be
271    /// used to store handler-specific data.
272    #[must_use]
273    pub fn state(&self) -> &SessionState {
274        &self.state
275    }
276
277    /// Returns the credential-free shared principal binding for this session.
278    pub(crate) fn principal_binding(&self) -> SessionPrincipalBinding {
279        self.principal_binding.clone()
280    }
281
282    /// Returns whether the session has been initialized.
283    #[must_use]
284    pub fn is_initialized(&self) -> bool {
285        self.initialized
286    }
287
288    /// Initializes the session with client info.
289    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    /// Captures every field mutated by [`Self::initialize`].
302    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    /// Restores an initialization snapshot after response finalization fails.
312    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    /// Returns the client info if initialized.
320    #[must_use]
321    pub fn client_info(&self) -> Option<&ClientInfo> {
322        self.client_info.as_ref()
323    }
324
325    /// Returns the client capabilities if initialized.
326    #[must_use]
327    pub fn client_capabilities(&self) -> Option<&ClientCapabilities> {
328        self.client_capabilities.as_ref()
329    }
330
331    /// Returns the server info.
332    #[must_use]
333    pub fn server_info(&self) -> &ServerInfo {
334        &self.server_info
335    }
336
337    /// Returns the server capabilities.
338    #[must_use]
339    pub fn server_capabilities(&self) -> &ServerCapabilities {
340        &self.server_capabilities
341    }
342
343    /// Returns the negotiated protocol version.
344    #[must_use]
345    pub fn protocol_version(&self) -> Option<&str> {
346        self.protocol_version.as_deref()
347    }
348
349    /// Subscribes to a resource URI for this session.
350    ///
351    /// Duplicate subscriptions are idempotent. New subscriptions that would
352    /// exceed the per-session count or aggregate UTF-8 byte limit are rejected
353    /// fail-closed and are not retained. A non-live request cannot mutate the
354    /// session.
355    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        // Reject an individually impossible candidate before hashing it for
365        // duplicate lookup. A URI larger than the aggregate byte cap can
366        // never be retained, even by an otherwise empty session.
367        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        // Reserve before the commit boundary so peer-driven growth reports a
406        // bounded admission failure instead of reaching HashSet's infallible
407        // allocation path and potentially panicking the session dispatcher.
408        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        // Recheck immediately before the synchronous commit so cancellation
422        // observed during admission cannot retain a new subscription.
423        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    /// Idempotently unsubscribes from a resource URI for this session.
436    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        // Such a URI cannot have passed subscription admission, so preserve
455        // idempotent success without hashing an unretainable peer string.
456        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        // Hashing a retained peer URI can be substantial at the configured
464        // bound. Recheck immediately before the synchronous removal so
465        // cancellation observed during that work cannot commit a mutation.
466        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    /// Rolls back a newly admitted subscription when request finalization fails.
485    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    /// Restores a removed subscription when request finalization fails.
494    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    /// Returns true if this session is subscribed to the given resource URI.
509    #[must_use]
510    pub fn is_resource_subscribed(&self, uri: &str) -> bool {
511        self.resource_subscriptions.contains(uri)
512    }
513
514    /// Returns the resource URIs this session currently subscribes to.
515    #[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    /// Sets the session log level for log notifications.
521    pub fn set_log_level(&mut self, level: LogLevel) {
522        self.log_level = Some(level);
523    }
524
525    /// Restores the exact pre-request log-level state, including `None`.
526    pub(crate) fn restore_log_level(&mut self, level: Option<LogLevel>) {
527        self.log_level = level;
528    }
529
530    /// Returns the current session log level for log notifications.
531    #[must_use]
532    pub fn log_level(&self) -> Option<LogLevel> {
533        self.log_level
534    }
535
536    /// Returns whether the client supports sampling (LLM completions).
537    #[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    /// Returns whether the client supports elicitation (user input requests).
545    #[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    /// Returns whether the client supports roots listing.
553    #[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    /// Sends a resource updated notification if the session is subscribed.
561    ///
562    /// Returns true if a notification was sent.
563    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    // ── New session initial state ────────────────────────────────────
646
647    #[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        // default caps should be accessible
696        let _ = session.server_capabilities();
697    }
698
699    // ── Initialization lifecycle ─────────────────────────────────────
700
701    #[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    // ── Resource subscriptions ───────────────────────────────────────
750
751    #[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    // ── Log level ────────────────────────────────────────────────────
952
953    #[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    // ── Session state ────────────────────────────────────────────────
969
970    #[test]
971    fn state_is_accessible() {
972        let session = make_session();
973        let state = session.state();
974        // fresh state should have no stored values
975        let val: Option<String> = state.get("key");
976        assert!(val.is_none());
977    }
978
979    // ── notify_resource_updated ──────────────────────────────────────
980
981    #[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    // ── Debug impl ───────────────────────────────────────────────────
1075
1076    #[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    // ── Existing capability tests ────────────────────────────────────
1137
1138    #[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        // Before initialization, no capabilities
1149        assert!(!session.supports_sampling());
1150
1151        // Initialize with sampling capability
1152        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    // ── Re-initialization ───────────────────────────────────────────
1281
1282    #[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    // ── State persistence through lifecycle ──────────────────────────
1311
1312    #[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    // ── Notification after unsubscribe ───────────────────────────────
1326
1327    #[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        // First notification should fire
1339        assert!(session.notify_resource_updated("r://x", &sender));
1340        assert_eq!(sent.lock().unwrap().len(), 1);
1341
1342        // Unsubscribe and try again
1343        remove_resource(&mut session, "r://x");
1344        assert!(!session.notify_resource_updated("r://x", &sender));
1345        // No new notifications sent
1346        assert_eq!(sent.lock().unwrap().len(), 1);
1347    }
1348
1349    // ── Subscribe → unsubscribe → re-subscribe ─────────────────────
1350
1351    #[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    // ── Debug format after initialization ───────────────────────────
1362
1363    #[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    // ── Non-default server capabilities ─────────────────────────────
1376
1377    #[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    // ── Log level with all variants ─────────────────────────────────
1395
1396    #[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    // ── Persistence across re-initialization ────────────────────────
1411
1412    #[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        // Re-initialize with different client info
1423        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        // Re-initialize
1456        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        // Notify first URI
1481        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        // Notify second URI
1495        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}