mcp_host/server/
session.rs

1//! Session management
2//!
3//! Tracks individual MCP client connections with state storage, lifecycle state machine,
4//! and per-session tool/resource/prompt customization.
5
6use std::collections::{HashMap, HashSet};
7use std::sync::{Arc, RwLock};
8
9use dashmap::DashMap;
10use serde_json::Value;
11use uuid::Uuid;
12
13use crate::protocol::capabilities::ClientCapabilities;
14use crate::protocol::types::Implementation;
15use crate::registry::prompts::Prompt;
16use crate::registry::resources::Resource;
17use crate::registry::tools::Tool;
18use crate::server::profile::SessionProfile;
19
20/// Session state storage (thread-safe)
21pub type SessionState = Arc<RwLock<HashMap<String, Value>>>;
22
23/// Session lifecycle states
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum SessionLifecycle {
26    /// Session created but not initialized
27    Created,
28    /// Session initialized and ready for requests
29    Ready,
30    /// Session encountered errors but still operational
31    Degraded,
32    /// Session closed
33    Closed,
34}
35
36impl SessionLifecycle {
37    /// Check if session can accept requests
38    pub fn can_accept_requests(&self) -> bool {
39        matches!(self, Self::Ready | Self::Degraded)
40    }
41
42    /// Check if session is healthy
43    pub fn is_healthy(&self) -> bool {
44        matches!(self, Self::Ready)
45    }
46}
47
48/// Session represents a single MCP client connection
49#[derive(Clone)]
50pub struct Session {
51    /// Unique session ID
52    pub id: String,
53
54    /// Client information (set during initialization)
55    pub client_info: Option<Implementation>,
56
57    /// Client capabilities (set during initialization)
58    pub capabilities: Option<ClientCapabilities>,
59
60    /// Negotiated protocol version (set during initialization)
61    pub protocol_version: Option<String>,
62
63    /// Session lifecycle state (Arc-wrapped for shared state across clones)
64    lifecycle: Arc<RwLock<SessionLifecycle>>,
65
66    /// Error count for degraded state transitions
67    error_count: Arc<RwLock<u32>>,
68
69    /// Per-session state storage
70    state: SessionState,
71
72    /// Notification channel for sending updates to client
73    notification_tx:
74        Option<tokio::sync::mpsc::UnboundedSender<crate::transport::traits::JsonRpcNotification>>,
75
76    // Tool customization
77    /// Tool overrides (same name, different implementation)
78    tool_overrides: Arc<DashMap<String, Arc<dyn Tool>>>,
79    /// Extra tools added to this session
80    tool_extras: Arc<DashMap<String, Arc<dyn Tool>>>,
81    /// Tools hidden from this session
82    tool_hidden: Arc<RwLock<HashSet<String>>>,
83    /// Tool aliases (alias -> actual name)
84    tool_aliases: Arc<RwLock<HashMap<String, String>>>,
85
86    // Resource customization
87    /// Resource overrides (same URI, different implementation)
88    resource_overrides: Arc<DashMap<String, Arc<dyn Resource>>>,
89    /// Extra resources added to this session
90    resource_extras: Arc<DashMap<String, Arc<dyn Resource>>>,
91    /// Resources hidden from this session
92    resource_hidden: Arc<RwLock<HashSet<String>>>,
93
94    // Prompt customization
95    /// Prompt overrides (same name, different implementation)
96    prompt_overrides: Arc<DashMap<String, Arc<dyn Prompt>>>,
97    /// Extra prompts added to this session
98    prompt_extras: Arc<DashMap<String, Arc<dyn Prompt>>>,
99    /// Prompts hidden from this session
100    prompt_hidden: Arc<RwLock<HashSet<String>>>,
101}
102
103impl Session {
104    /// Create new session with random UUID
105    pub fn new() -> Self {
106        Self {
107            id: Uuid::new_v4().to_string(),
108            client_info: None,
109            capabilities: None,
110            protocol_version: None,
111            lifecycle: Arc::new(RwLock::new(SessionLifecycle::Created)),
112            error_count: Arc::new(RwLock::new(0)),
113            state: Arc::new(RwLock::new(HashMap::new())),
114            notification_tx: None,
115            // Tool customization
116            tool_overrides: Arc::new(DashMap::new()),
117            tool_extras: Arc::new(DashMap::new()),
118            tool_hidden: Arc::new(RwLock::new(HashSet::new())),
119            tool_aliases: Arc::new(RwLock::new(HashMap::new())),
120            // Resource customization
121            resource_overrides: Arc::new(DashMap::new()),
122            resource_extras: Arc::new(DashMap::new()),
123            resource_hidden: Arc::new(RwLock::new(HashSet::new())),
124            // Prompt customization
125            prompt_overrides: Arc::new(DashMap::new()),
126            prompt_extras: Arc::new(DashMap::new()),
127            prompt_hidden: Arc::new(RwLock::new(HashSet::new())),
128        }
129    }
130
131    /// Create session with specific ID
132    pub fn with_id(id: impl Into<String>) -> Self {
133        Self {
134            id: id.into(),
135            client_info: None,
136            capabilities: None,
137            protocol_version: None,
138            lifecycle: Arc::new(RwLock::new(SessionLifecycle::Created)),
139            error_count: Arc::new(RwLock::new(0)),
140            state: Arc::new(RwLock::new(HashMap::new())),
141            notification_tx: None,
142            // Tool customization
143            tool_overrides: Arc::new(DashMap::new()),
144            tool_extras: Arc::new(DashMap::new()),
145            tool_hidden: Arc::new(RwLock::new(HashSet::new())),
146            tool_aliases: Arc::new(RwLock::new(HashMap::new())),
147            // Resource customization
148            resource_overrides: Arc::new(DashMap::new()),
149            resource_extras: Arc::new(DashMap::new()),
150            resource_hidden: Arc::new(RwLock::new(HashSet::new())),
151            // Prompt customization
152            prompt_overrides: Arc::new(DashMap::new()),
153            prompt_extras: Arc::new(DashMap::new()),
154            prompt_hidden: Arc::new(RwLock::new(HashSet::new())),
155        }
156    }
157
158    /// Initialize session with client info and capabilities
159    /// Transitions: Created -> Ready
160    pub fn initialize(
161        &mut self,
162        client_info: Implementation,
163        capabilities: ClientCapabilities,
164        protocol_version: String,
165    ) {
166        self.client_info = Some(client_info);
167        self.capabilities = Some(capabilities);
168        self.protocol_version = Some(protocol_version);
169        *self.lifecycle.write().unwrap() = SessionLifecycle::Ready;
170    }
171
172    /// Set notification channel for this session
173    pub fn set_notification_channel(
174        &mut self,
175        tx: tokio::sync::mpsc::UnboundedSender<crate::transport::traits::JsonRpcNotification>,
176    ) {
177        self.notification_tx = Some(tx);
178    }
179
180    /// Check if session is initialized (Ready or Degraded)
181    pub fn is_initialized(&self) -> bool {
182        self.lifecycle.read().unwrap().can_accept_requests()
183    }
184
185    /// Get the negotiated protocol version
186    pub fn protocol_version(&self) -> Option<&str> {
187        self.protocol_version.as_deref()
188    }
189
190    /// Record an error, potentially transitioning to Degraded state
191    /// Transitions: Ready -> Degraded (after threshold errors)
192    pub fn record_error(&mut self) {
193        if let Ok(mut count) = self.error_count.write() {
194            *count += 1;
195            // Transition to Degraded after 3 consecutive errors
196            if *count >= 3 && *self.lifecycle.read().unwrap() == SessionLifecycle::Ready {
197                *self.lifecycle.write().unwrap() = SessionLifecycle::Degraded;
198            }
199        }
200    }
201
202    /// Record success, potentially recovering from Degraded state
203    /// Transitions: Degraded -> Ready (resets error count)
204    pub fn record_success(&mut self) {
205        if let Ok(mut count) = self.error_count.write() {
206            *count = 0;
207            if *self.lifecycle.read().unwrap() == SessionLifecycle::Degraded {
208                *self.lifecycle.write().unwrap() = SessionLifecycle::Ready;
209            }
210        }
211    }
212
213    /// Close the session
214    /// Transitions: Any -> Closed
215    pub fn close(&mut self) {
216        *self.lifecycle.write().unwrap() = SessionLifecycle::Closed;
217    }
218
219    /// Get current lifecycle state
220    pub fn lifecycle(&self) -> SessionLifecycle {
221        *self.lifecycle.read().unwrap()
222    }
223
224    /// Get current error count
225    pub fn error_count(&self) -> u32 {
226        self.error_count.read().map(|c| *c).unwrap_or(0)
227    }
228
229    /// Get state value
230    pub fn get_state(&self, key: &str) -> Option<Value> {
231        self.state.read().ok()?.get(key).cloned()
232    }
233
234    /// Set state value
235    pub fn set_state(&self, key: impl Into<String>, value: Value) {
236        if let Ok(mut state) = self.state.write() {
237            state.insert(key.into(), value);
238        }
239    }
240
241    /// Remove state value
242    pub fn remove_state(&self, key: &str) -> Option<Value> {
243        self.state.write().ok()?.remove(key)
244    }
245
246    /// Clear all state
247    pub fn clear_state(&self) {
248        if let Ok(mut state) = self.state.write() {
249            state.clear();
250        }
251    }
252
253    /// Get all state keys
254    pub fn state_keys(&self) -> Vec<String> {
255        self.state
256            .read()
257            .ok()
258            .map(|state| state.keys().cloned().collect())
259            .unwrap_or_default()
260    }
261
262    // ==================== Tool Management ====================
263
264    /// Add an extra tool to this session
265    pub fn add_tool(&self, tool: Arc<dyn Tool>) {
266        let name = tool.name().to_string();
267        self.tool_extras.insert(name, tool);
268    }
269
270    /// Override a tool (same name, different implementation)
271    pub fn override_tool(&self, name: impl Into<String>, tool: Arc<dyn Tool>) {
272        self.tool_overrides.insert(name.into(), tool);
273    }
274
275    /// Hide a tool from this session
276    pub fn hide_tool(&self, name: impl Into<String>) {
277        if let Ok(mut hidden) = self.tool_hidden.write() {
278            hidden.insert(name.into());
279        }
280    }
281
282    /// Unhide a tool
283    pub fn unhide_tool(&self, name: &str) {
284        if let Ok(mut hidden) = self.tool_hidden.write() {
285            hidden.remove(name);
286        }
287    }
288
289    /// Add a tool alias (alias -> target)
290    pub fn alias_tool(&self, alias: impl Into<String>, target: impl Into<String>) {
291        if let Ok(mut aliases) = self.tool_aliases.write() {
292            aliases.insert(alias.into(), target.into());
293        }
294    }
295
296    /// Remove a tool alias
297    pub fn remove_tool_alias(&self, alias: &str) {
298        if let Ok(mut aliases) = self.tool_aliases.write() {
299            aliases.remove(alias);
300        }
301    }
302
303    /// Check if a tool is hidden
304    pub fn is_tool_hidden(&self, name: &str) -> bool {
305        self.tool_hidden
306            .read()
307            .map(|hidden| hidden.contains(name))
308            .unwrap_or(false)
309    }
310
311    /// Resolve a tool alias to its target name
312    pub fn resolve_tool_alias<'a>(&self, name: &'a str) -> std::borrow::Cow<'a, str> {
313        self.tool_aliases
314            .read()
315            .ok()
316            .and_then(|aliases| aliases.get(name).cloned())
317            .map(std::borrow::Cow::Owned)
318            .unwrap_or(std::borrow::Cow::Borrowed(name))
319    }
320
321    /// Get a tool override for this session
322    pub fn get_tool_override(&self, name: &str) -> Option<Arc<dyn Tool>> {
323        self.tool_overrides.get(name).map(|r| Arc::clone(&r))
324    }
325
326    /// Get an extra tool added to this session
327    pub fn get_tool_extra(&self, name: &str) -> Option<Arc<dyn Tool>> {
328        self.tool_extras.get(name).map(|r| Arc::clone(&r))
329    }
330
331    /// Get all tool overrides
332    pub fn tool_overrides(&self) -> &Arc<DashMap<String, Arc<dyn Tool>>> {
333        &self.tool_overrides
334    }
335
336    /// Get all extra tools
337    pub fn tool_extras(&self) -> &Arc<DashMap<String, Arc<dyn Tool>>> {
338        &self.tool_extras
339    }
340
341    /// Remove an extra tool from this session
342    pub fn remove_tool(&self, name: &str) -> Option<Arc<dyn Tool>> {
343        self.tool_extras.remove(name).map(|(_, tool)| tool)
344    }
345
346    /// Remove a tool override
347    pub fn remove_tool_override(&self, name: &str) -> Option<Arc<dyn Tool>> {
348        self.tool_overrides.remove(name).map(|(_, tool)| tool)
349    }
350
351    // ==================== Resource Management ====================
352
353    /// Add an extra resource to this session
354    pub fn add_resource(&self, resource: Arc<dyn Resource>) {
355        let uri = resource.uri().to_string();
356        self.resource_extras.insert(uri, resource);
357    }
358
359    /// Override a resource (same URI, different implementation)
360    pub fn override_resource(&self, uri: impl Into<String>, resource: Arc<dyn Resource>) {
361        self.resource_overrides.insert(uri.into(), resource);
362    }
363
364    /// Hide a resource from this session
365    pub fn hide_resource(&self, uri: impl Into<String>) {
366        if let Ok(mut hidden) = self.resource_hidden.write() {
367            hidden.insert(uri.into());
368        }
369    }
370
371    /// Unhide a resource
372    pub fn unhide_resource(&self, uri: &str) {
373        if let Ok(mut hidden) = self.resource_hidden.write() {
374            hidden.remove(uri);
375        }
376    }
377
378    /// Check if a resource is hidden
379    pub fn is_resource_hidden(&self, uri: &str) -> bool {
380        self.resource_hidden
381            .read()
382            .map(|hidden| hidden.contains(uri))
383            .unwrap_or(false)
384    }
385
386    /// Get a resource override for this session
387    pub fn get_resource_override(&self, uri: &str) -> Option<Arc<dyn Resource>> {
388        self.resource_overrides.get(uri).map(|r| Arc::clone(&r))
389    }
390
391    /// Get an extra resource added to this session
392    pub fn get_resource_extra(&self, uri: &str) -> Option<Arc<dyn Resource>> {
393        self.resource_extras.get(uri).map(|r| Arc::clone(&r))
394    }
395
396    /// Get all resource overrides
397    pub fn resource_overrides(&self) -> &Arc<DashMap<String, Arc<dyn Resource>>> {
398        &self.resource_overrides
399    }
400
401    /// Get all extra resources
402    pub fn resource_extras(&self) -> &Arc<DashMap<String, Arc<dyn Resource>>> {
403        &self.resource_extras
404    }
405
406    /// Remove an extra resource from this session
407    pub fn remove_resource(&self, uri: &str) -> Option<Arc<dyn Resource>> {
408        self.resource_extras
409            .remove(uri)
410            .map(|(_, resource)| resource)
411    }
412
413    /// Remove a resource override
414    pub fn remove_resource_override(&self, uri: &str) -> Option<Arc<dyn Resource>> {
415        self.resource_overrides
416            .remove(uri)
417            .map(|(_, resource)| resource)
418    }
419
420    // ==================== Prompt Management ====================
421
422    /// Add an extra prompt to this session
423    pub fn add_prompt(&self, prompt: Arc<dyn Prompt>) {
424        let name = prompt.name().to_string();
425        self.prompt_extras.insert(name, prompt);
426    }
427
428    /// Override a prompt (same name, different implementation)
429    pub fn override_prompt(&self, name: impl Into<String>, prompt: Arc<dyn Prompt>) {
430        self.prompt_overrides.insert(name.into(), prompt);
431    }
432
433    /// Hide a prompt from this session
434    pub fn hide_prompt(&self, name: impl Into<String>) {
435        if let Ok(mut hidden) = self.prompt_hidden.write() {
436            hidden.insert(name.into());
437        }
438    }
439
440    /// Unhide a prompt
441    pub fn unhide_prompt(&self, name: &str) {
442        if let Ok(mut hidden) = self.prompt_hidden.write() {
443            hidden.remove(name);
444        }
445    }
446
447    /// Check if a prompt is hidden
448    pub fn is_prompt_hidden(&self, name: &str) -> bool {
449        self.prompt_hidden
450            .read()
451            .map(|hidden| hidden.contains(name))
452            .unwrap_or(false)
453    }
454
455    /// Get a prompt override for this session
456    pub fn get_prompt_override(&self, name: &str) -> Option<Arc<dyn Prompt>> {
457        self.prompt_overrides.get(name).map(|r| Arc::clone(&r))
458    }
459
460    /// Get an extra prompt added to this session
461    pub fn get_prompt_extra(&self, name: &str) -> Option<Arc<dyn Prompt>> {
462        self.prompt_extras.get(name).map(|r| Arc::clone(&r))
463    }
464
465    /// Get all prompt overrides
466    pub fn prompt_overrides(&self) -> &Arc<DashMap<String, Arc<dyn Prompt>>> {
467        &self.prompt_overrides
468    }
469
470    /// Get all extra prompts
471    pub fn prompt_extras(&self) -> &Arc<DashMap<String, Arc<dyn Prompt>>> {
472        &self.prompt_extras
473    }
474
475    /// Remove an extra prompt from this session
476    pub fn remove_prompt(&self, name: &str) -> Option<Arc<dyn Prompt>> {
477        self.prompt_extras.remove(name).map(|(_, prompt)| prompt)
478    }
479
480    /// Remove a prompt override
481    pub fn remove_prompt_override(&self, name: &str) -> Option<Arc<dyn Prompt>> {
482        self.prompt_overrides.remove(name).map(|(_, prompt)| prompt)
483    }
484
485    // ==================== Batch Operations & Notifications ====================
486
487    /// Perform batch operations without sending notifications until complete
488    ///
489    /// This allows you to make multiple changes and send a single notification.
490    ///
491    /// # Example
492    /// ```rust,ignore
493    /// session.batch(|batch| {
494    ///     batch.add_tool(Arc::new(Tool1));
495    ///     batch.add_tool(Arc::new(Tool2));
496    ///     batch.remove_resource("old://resource");
497    ///     batch.hide_prompt("deprecated");
498    /// });
499    /// // Single notification sent after closure completes
500    /// ```
501    pub fn batch<F>(&self, f: F)
502    where
503        F: FnOnce(&mut SessionBatch<'_>),
504    {
505        let mut batch = SessionBatch {
506            session: self,
507            tools_changed: false,
508            resources_changed: false,
509            prompts_changed: false,
510        };
511
512        f(&mut batch);
513
514        // Send notifications for what changed
515        if batch.tools_changed {
516            self.notify_tools_changed();
517        }
518        if batch.resources_changed {
519            self.notify_resources_changed();
520        }
521        if batch.prompts_changed {
522            self.notify_prompts_changed();
523        }
524    }
525
526    /// Send tools/list_changed notification
527    fn notify_tools_changed(&self) {
528        if let Some(ref tx) = self.notification_tx {
529            let notification = crate::transport::traits::JsonRpcNotification::new(
530                "notifications/tools/list_changed",
531                None,
532            );
533            let _ = tx.send(notification);
534        }
535    }
536
537    /// Send resources/list_changed notification
538    fn notify_resources_changed(&self) {
539        if let Some(ref tx) = self.notification_tx {
540            let notification = crate::transport::traits::JsonRpcNotification::new(
541                "notifications/resources/list_changed",
542                None,
543            );
544            let _ = tx.send(notification);
545        }
546    }
547
548    /// Send prompts/list_changed notification
549    fn notify_prompts_changed(&self) {
550        if let Some(ref tx) = self.notification_tx {
551            let notification = crate::transport::traits::JsonRpcNotification::new(
552                "notifications/prompts/list_changed",
553                None,
554            );
555            let _ = tx.send(notification);
556        }
557    }
558
559    // ==================== Profile Management ====================
560
561    /// Apply a session profile
562    ///
563    /// This adds/overrides/hides tools, resources, and prompts according to the profile
564    pub fn apply_profile(&self, profile: &SessionProfile) {
565        // Apply tool configuration
566        for tool in &profile.tool_extras {
567            self.add_tool(Arc::clone(tool));
568        }
569        for (name, tool) in &profile.tool_overrides {
570            self.override_tool(name.clone(), Arc::clone(tool));
571        }
572        for name in &profile.tool_hidden {
573            self.hide_tool(name.clone());
574        }
575        for (alias, target) in &profile.tool_aliases {
576            self.alias_tool(alias.clone(), target.clone());
577        }
578
579        // Apply resource configuration
580        for resource in &profile.resource_extras {
581            self.add_resource(Arc::clone(resource));
582        }
583        for (uri, resource) in &profile.resource_overrides {
584            self.override_resource(uri.clone(), Arc::clone(resource));
585        }
586        for uri in &profile.resource_hidden {
587            self.hide_resource(uri.clone());
588        }
589
590        // Apply prompt configuration
591        for prompt in &profile.prompt_extras {
592            self.add_prompt(Arc::clone(prompt));
593        }
594        for (name, prompt) in &profile.prompt_overrides {
595            self.override_prompt(name.clone(), Arc::clone(prompt));
596        }
597        for name in &profile.prompt_hidden {
598            self.hide_prompt(name.clone());
599        }
600    }
601
602    /// Clear all session customizations
603    pub fn clear_customizations(&self) {
604        // Clear tools
605        self.tool_overrides.clear();
606        self.tool_extras.clear();
607        if let Ok(mut hidden) = self.tool_hidden.write() {
608            hidden.clear();
609        }
610        if let Ok(mut aliases) = self.tool_aliases.write() {
611            aliases.clear();
612        }
613
614        // Clear resources
615        self.resource_overrides.clear();
616        self.resource_extras.clear();
617        if let Ok(mut hidden) = self.resource_hidden.write() {
618            hidden.clear();
619        }
620
621        // Clear prompts
622        self.prompt_overrides.clear();
623        self.prompt_extras.clear();
624        if let Ok(mut hidden) = self.prompt_hidden.write() {
625            hidden.clear();
626        }
627    }
628}
629
630// Implement Debug manually since trait objects don't implement Debug
631impl std::fmt::Debug for Session {
632    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
633        f.debug_struct("Session")
634            .field("id", &self.id)
635            .field("client_info", &self.client_info)
636            .field("capabilities", &self.capabilities)
637            .field("lifecycle", &self.lifecycle)
638            .field("error_count", &self.error_count())
639            .field("tool_overrides_count", &self.tool_overrides.len())
640            .field("tool_extras_count", &self.tool_extras.len())
641            .field("resource_overrides_count", &self.resource_overrides.len())
642            .field("resource_extras_count", &self.resource_extras.len())
643            .field("prompt_overrides_count", &self.prompt_overrides.len())
644            .field("prompt_extras_count", &self.prompt_extras.len())
645            .finish()
646    }
647}
648
649impl Default for Session {
650    fn default() -> Self {
651        Self::new()
652    }
653}
654
655/// Batch operations helper for making multiple changes with a single notification
656pub struct SessionBatch<'a> {
657    session: &'a Session,
658    tools_changed: bool,
659    resources_changed: bool,
660    prompts_changed: bool,
661}
662
663impl<'a> SessionBatch<'a> {
664    // Tool operations
665    pub fn add_tool(&mut self, tool: Arc<dyn crate::registry::tools::Tool>) {
666        self.session.add_tool(tool);
667        self.tools_changed = true;
668    }
669
670    pub fn override_tool(
671        &mut self,
672        name: impl Into<String>,
673        tool: Arc<dyn crate::registry::tools::Tool>,
674    ) {
675        self.session.override_tool(name, tool);
676        self.tools_changed = true;
677    }
678
679    pub fn remove_tool(&mut self, name: &str) -> Option<Arc<dyn crate::registry::tools::Tool>> {
680        let result = self.session.remove_tool(name);
681        if result.is_some() {
682            self.tools_changed = true;
683        }
684        result
685    }
686
687    pub fn hide_tool(&mut self, name: impl Into<String>) {
688        self.session.hide_tool(name);
689        self.tools_changed = true;
690    }
691
692    pub fn unhide_tool(&mut self, name: &str) {
693        self.session.unhide_tool(name);
694        self.tools_changed = true;
695    }
696
697    // Resource operations
698    pub fn add_resource(&mut self, resource: Arc<dyn crate::registry::resources::Resource>) {
699        self.session.add_resource(resource);
700        self.resources_changed = true;
701    }
702
703    pub fn override_resource(
704        &mut self,
705        uri: impl Into<String>,
706        resource: Arc<dyn crate::registry::resources::Resource>,
707    ) {
708        self.session.override_resource(uri, resource);
709        self.resources_changed = true;
710    }
711
712    pub fn remove_resource(
713        &mut self,
714        uri: &str,
715    ) -> Option<Arc<dyn crate::registry::resources::Resource>> {
716        let result = self.session.remove_resource(uri);
717        if result.is_some() {
718            self.resources_changed = true;
719        }
720        result
721    }
722
723    pub fn hide_resource(&mut self, uri: impl Into<String>) {
724        self.session.hide_resource(uri);
725        self.resources_changed = true;
726    }
727
728    pub fn unhide_resource(&mut self, uri: &str) {
729        self.session.unhide_resource(uri);
730        self.resources_changed = true;
731    }
732
733    // Prompt operations
734    pub fn add_prompt(&mut self, prompt: Arc<dyn crate::registry::prompts::Prompt>) {
735        self.session.add_prompt(prompt);
736        self.prompts_changed = true;
737    }
738
739    pub fn override_prompt(
740        &mut self,
741        name: impl Into<String>,
742        prompt: Arc<dyn crate::registry::prompts::Prompt>,
743    ) {
744        self.session.override_prompt(name, prompt);
745        self.prompts_changed = true;
746    }
747
748    pub fn remove_prompt(
749        &mut self,
750        name: &str,
751    ) -> Option<Arc<dyn crate::registry::prompts::Prompt>> {
752        let result = self.session.remove_prompt(name);
753        if result.is_some() {
754            self.prompts_changed = true;
755        }
756        result
757    }
758
759    pub fn hide_prompt(&mut self, name: impl Into<String>) {
760        self.session.hide_prompt(name);
761        self.prompts_changed = true;
762    }
763
764    pub fn unhide_prompt(&mut self, name: &str) {
765        self.session.unhide_prompt(name);
766        self.prompts_changed = true;
767    }
768}
769
770#[cfg(test)]
771mod tests {
772    use super::*;
773
774    #[test]
775    fn test_session_creation() {
776        let session = Session::new();
777        assert!(!session.id.is_empty());
778        assert_eq!(session.lifecycle(), SessionLifecycle::Created);
779        assert!(!session.is_initialized());
780        assert!(session.client_info.is_none());
781        assert!(session.capabilities.is_none());
782    }
783
784    #[test]
785    fn test_session_with_id() {
786        let session = Session::with_id("test-session");
787        assert_eq!(session.id, "test-session");
788        assert_eq!(session.lifecycle(), SessionLifecycle::Created);
789    }
790
791    #[test]
792    fn test_session_initialization() {
793        let mut session = Session::new();
794        let client_info = Implementation {
795            name: "test-client".to_string(),
796            version: "1.0.0".to_string(),
797        };
798        let capabilities = ClientCapabilities::default();
799
800        session.initialize(client_info.clone(), capabilities, "2025-06-18".to_string());
801
802        assert!(session.is_initialized());
803        assert_eq!(session.lifecycle(), SessionLifecycle::Ready);
804        assert_eq!(session.client_info.unwrap().name, "test-client");
805    }
806
807    #[test]
808    fn test_session_lifecycle_transitions() {
809        let mut session = Session::new();
810        assert_eq!(session.lifecycle(), SessionLifecycle::Created);
811        assert!(!session.lifecycle().can_accept_requests());
812
813        // Initialize -> Ready
814        session.initialize(
815            Implementation {
816                name: "test".into(),
817                version: "1.0".into(),
818            },
819            ClientCapabilities::default(),
820            "2025-06-18".to_string(),
821        );
822        assert_eq!(session.lifecycle(), SessionLifecycle::Ready);
823        assert!(session.lifecycle().can_accept_requests());
824        assert!(session.lifecycle().is_healthy());
825
826        // Record errors -> Degraded (after 3)
827        session.record_error();
828        assert_eq!(session.lifecycle(), SessionLifecycle::Ready);
829        session.record_error();
830        assert_eq!(session.lifecycle(), SessionLifecycle::Ready);
831        session.record_error();
832        assert_eq!(session.lifecycle(), SessionLifecycle::Degraded);
833        assert!(session.lifecycle().can_accept_requests());
834        assert!(!session.lifecycle().is_healthy());
835
836        // Record success -> Ready (recovery)
837        session.record_success();
838        assert_eq!(session.lifecycle(), SessionLifecycle::Ready);
839        assert_eq!(session.error_count(), 0);
840
841        // Close -> Closed
842        session.close();
843        assert_eq!(session.lifecycle(), SessionLifecycle::Closed);
844        assert!(!session.lifecycle().can_accept_requests());
845    }
846
847    #[test]
848    fn test_session_state() {
849        let session = Session::new();
850
851        // Set state
852        session.set_state("key1", Value::String("value1".to_string()));
853        session.set_state("key2", Value::Number(42.into()));
854
855        // Get state
856        assert_eq!(
857            session.get_state("key1"),
858            Some(Value::String("value1".to_string()))
859        );
860        assert_eq!(session.get_state("key2"), Some(Value::Number(42.into())));
861        assert_eq!(session.get_state("nonexistent"), None);
862
863        // State keys
864        let keys = session.state_keys();
865        assert_eq!(keys.len(), 2);
866        assert!(keys.contains(&"key1".to_string()));
867        assert!(keys.contains(&"key2".to_string()));
868
869        // Remove state
870        let removed = session.remove_state("key1");
871        assert_eq!(removed, Some(Value::String("value1".to_string())));
872        assert_eq!(session.get_state("key1"), None);
873
874        // Clear state
875        session.clear_state();
876        assert_eq!(session.state_keys().len(), 0);
877    }
878
879    #[test]
880    fn test_session_batch() {
881        let mut session = Session::new();
882        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
883        session.set_notification_channel(tx);
884
885        // Test that batch operations track changes
886        session.batch(|batch| {
887            batch.hide_tool("unwanted");
888            batch.hide_resource("blocked://resource");
889            batch.hide_prompt("deprecated");
890        });
891
892        assert!(session.is_tool_hidden("unwanted"));
893        assert!(session.is_resource_hidden("blocked://resource"));
894        assert!(session.is_prompt_hidden("deprecated"));
895    }
896
897    #[test]
898    fn test_session_clone() {
899        let session1 = Session::with_id("test");
900        session1.set_state("shared", Value::Bool(true));
901
902        let session2 = session1.clone();
903
904        // Both sessions share the same state storage
905        assert_eq!(session1.id, session2.id);
906        assert_eq!(session2.get_state("shared"), Some(Value::Bool(true)));
907
908        // Modifying state in one affects the other
909        session2.set_state("shared", Value::Bool(false));
910        assert_eq!(session1.get_state("shared"), Some(Value::Bool(false)));
911    }
912}
mcp_host/server/session.rs

mcp_host/server/
session.rs
