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    ///
277    /// If `notify` is true, sends a `notifications/tools/list_changed` notification.
278    pub fn hide_tool(&self, name: impl Into<String>, notify: bool) {
279        if let Ok(mut hidden) = self.tool_hidden.write() {
280            hidden.insert(name.into());
281        }
282        if notify {
283            self.notify_tools_changed();
284        }
285    }
286
287    /// Unhide a tool
288    ///
289    /// If `notify` is true, sends a `notifications/tools/list_changed` notification.
290    pub fn unhide_tool(&self, name: &str, notify: bool) {
291        if let Ok(mut hidden) = self.tool_hidden.write() {
292            hidden.remove(name);
293        }
294        if notify {
295            self.notify_tools_changed();
296        }
297    }
298
299    /// Add a tool alias (alias -> target)
300    pub fn alias_tool(&self, alias: impl Into<String>, target: impl Into<String>) {
301        if let Ok(mut aliases) = self.tool_aliases.write() {
302            aliases.insert(alias.into(), target.into());
303        }
304    }
305
306    /// Remove a tool alias
307    pub fn remove_tool_alias(&self, alias: &str) {
308        if let Ok(mut aliases) = self.tool_aliases.write() {
309            aliases.remove(alias);
310        }
311    }
312
313    /// Check if a tool is hidden
314    pub fn is_tool_hidden(&self, name: &str) -> bool {
315        self.tool_hidden
316            .read()
317            .map(|hidden| hidden.contains(name))
318            .unwrap_or(false)
319    }
320
321    /// Resolve a tool alias to its target name
322    pub fn resolve_tool_alias<'a>(&self, name: &'a str) -> std::borrow::Cow<'a, str> {
323        self.tool_aliases
324            .read()
325            .ok()
326            .and_then(|aliases| aliases.get(name).cloned())
327            .map(std::borrow::Cow::Owned)
328            .unwrap_or(std::borrow::Cow::Borrowed(name))
329    }
330
331    /// Get a tool override for this session
332    pub fn get_tool_override(&self, name: &str) -> Option<Arc<dyn Tool>> {
333        self.tool_overrides.get(name).map(|r| Arc::clone(&r))
334    }
335
336    /// Get an extra tool added to this session
337    pub fn get_tool_extra(&self, name: &str) -> Option<Arc<dyn Tool>> {
338        self.tool_extras.get(name).map(|r| Arc::clone(&r))
339    }
340
341    /// Get all tool overrides
342    pub fn tool_overrides(&self) -> &Arc<DashMap<String, Arc<dyn Tool>>> {
343        &self.tool_overrides
344    }
345
346    /// Get all extra tools
347    pub fn tool_extras(&self) -> &Arc<DashMap<String, Arc<dyn Tool>>> {
348        &self.tool_extras
349    }
350
351    /// Remove an extra tool from this session
352    pub fn remove_tool(&self, name: &str) -> Option<Arc<dyn Tool>> {
353        self.tool_extras.remove(name).map(|(_, tool)| tool)
354    }
355
356    /// Remove a tool override
357    pub fn remove_tool_override(&self, name: &str) -> Option<Arc<dyn Tool>> {
358        self.tool_overrides.remove(name).map(|(_, tool)| tool)
359    }
360
361    // ==================== Resource Management ====================
362
363    /// Add an extra resource to this session
364    pub fn add_resource(&self, resource: Arc<dyn Resource>) {
365        let uri = resource.uri().to_string();
366        self.resource_extras.insert(uri, resource);
367    }
368
369    /// Override a resource (same URI, different implementation)
370    pub fn override_resource(&self, uri: impl Into<String>, resource: Arc<dyn Resource>) {
371        self.resource_overrides.insert(uri.into(), resource);
372    }
373
374    /// Hide a resource from this session
375    pub fn hide_resource(&self, uri: impl Into<String>) {
376        if let Ok(mut hidden) = self.resource_hidden.write() {
377            hidden.insert(uri.into());
378        }
379    }
380
381    /// Unhide a resource
382    pub fn unhide_resource(&self, uri: &str) {
383        if let Ok(mut hidden) = self.resource_hidden.write() {
384            hidden.remove(uri);
385        }
386    }
387
388    /// Check if a resource is hidden
389    pub fn is_resource_hidden(&self, uri: &str) -> bool {
390        self.resource_hidden
391            .read()
392            .map(|hidden| hidden.contains(uri))
393            .unwrap_or(false)
394    }
395
396    /// Get a resource override for this session
397    pub fn get_resource_override(&self, uri: &str) -> Option<Arc<dyn Resource>> {
398        self.resource_overrides.get(uri).map(|r| Arc::clone(&r))
399    }
400
401    /// Get an extra resource added to this session
402    pub fn get_resource_extra(&self, uri: &str) -> Option<Arc<dyn Resource>> {
403        self.resource_extras.get(uri).map(|r| Arc::clone(&r))
404    }
405
406    /// Get all resource overrides
407    pub fn resource_overrides(&self) -> &Arc<DashMap<String, Arc<dyn Resource>>> {
408        &self.resource_overrides
409    }
410
411    /// Get all extra resources
412    pub fn resource_extras(&self) -> &Arc<DashMap<String, Arc<dyn Resource>>> {
413        &self.resource_extras
414    }
415
416    /// Remove an extra resource from this session
417    pub fn remove_resource(&self, uri: &str) -> Option<Arc<dyn Resource>> {
418        self.resource_extras
419            .remove(uri)
420            .map(|(_, resource)| resource)
421    }
422
423    /// Remove a resource override
424    pub fn remove_resource_override(&self, uri: &str) -> Option<Arc<dyn Resource>> {
425        self.resource_overrides
426            .remove(uri)
427            .map(|(_, resource)| resource)
428    }
429
430    // ==================== Prompt Management ====================
431
432    /// Add an extra prompt to this session
433    pub fn add_prompt(&self, prompt: Arc<dyn Prompt>) {
434        let name = prompt.name().to_string();
435        self.prompt_extras.insert(name, prompt);
436    }
437
438    /// Override a prompt (same name, different implementation)
439    pub fn override_prompt(&self, name: impl Into<String>, prompt: Arc<dyn Prompt>) {
440        self.prompt_overrides.insert(name.into(), prompt);
441    }
442
443    /// Hide a prompt from this session
444    pub fn hide_prompt(&self, name: impl Into<String>) {
445        if let Ok(mut hidden) = self.prompt_hidden.write() {
446            hidden.insert(name.into());
447        }
448    }
449
450    /// Unhide a prompt
451    pub fn unhide_prompt(&self, name: &str) {
452        if let Ok(mut hidden) = self.prompt_hidden.write() {
453            hidden.remove(name);
454        }
455    }
456
457    /// Check if a prompt is hidden
458    pub fn is_prompt_hidden(&self, name: &str) -> bool {
459        self.prompt_hidden
460            .read()
461            .map(|hidden| hidden.contains(name))
462            .unwrap_or(false)
463    }
464
465    /// Get a prompt override for this session
466    pub fn get_prompt_override(&self, name: &str) -> Option<Arc<dyn Prompt>> {
467        self.prompt_overrides.get(name).map(|r| Arc::clone(&r))
468    }
469
470    /// Get an extra prompt added to this session
471    pub fn get_prompt_extra(&self, name: &str) -> Option<Arc<dyn Prompt>> {
472        self.prompt_extras.get(name).map(|r| Arc::clone(&r))
473    }
474
475    /// Get all prompt overrides
476    pub fn prompt_overrides(&self) -> &Arc<DashMap<String, Arc<dyn Prompt>>> {
477        &self.prompt_overrides
478    }
479
480    /// Get all extra prompts
481    pub fn prompt_extras(&self) -> &Arc<DashMap<String, Arc<dyn Prompt>>> {
482        &self.prompt_extras
483    }
484
485    /// Remove an extra prompt from this session
486    pub fn remove_prompt(&self, name: &str) -> Option<Arc<dyn Prompt>> {
487        self.prompt_extras.remove(name).map(|(_, prompt)| prompt)
488    }
489
490    /// Remove a prompt override
491    pub fn remove_prompt_override(&self, name: &str) -> Option<Arc<dyn Prompt>> {
492        self.prompt_overrides.remove(name).map(|(_, prompt)| prompt)
493    }
494
495    // ==================== Batch Operations & Notifications ====================
496
497    /// Perform batch operations without sending notifications until complete
498    ///
499    /// This allows you to make multiple changes and send a single notification.
500    ///
501    /// # Example
502    /// ```rust,ignore
503    /// session.batch(|batch| {
504    ///     batch.add_tool(Arc::new(Tool1));
505    ///     batch.add_tool(Arc::new(Tool2));
506    ///     batch.remove_resource("old://resource");
507    ///     batch.hide_prompt("deprecated");
508    /// });
509    /// // Single notification sent after closure completes
510    /// ```
511    pub fn batch<F>(&self, f: F)
512    where
513        F: FnOnce(&mut SessionBatch<'_>),
514    {
515        let mut batch = SessionBatch {
516            session: self,
517            tools_changed: false,
518            resources_changed: false,
519            prompts_changed: false,
520        };
521
522        f(&mut batch);
523
524        // Send notifications for what changed
525        if batch.tools_changed {
526            self.notify_tools_changed();
527        }
528        if batch.resources_changed {
529            self.notify_resources_changed();
530        }
531        if batch.prompts_changed {
532            self.notify_prompts_changed();
533        }
534    }
535
536    /// Send tools/list_changed notification
537    fn notify_tools_changed(&self) {
538        if let Some(ref tx) = self.notification_tx {
539            let notification = crate::transport::traits::JsonRpcNotification::new(
540                "notifications/tools/list_changed",
541                None,
542            );
543            let _ = tx.send(notification);
544        }
545    }
546
547    /// Send resources/list_changed notification
548    fn notify_resources_changed(&self) {
549        if let Some(ref tx) = self.notification_tx {
550            let notification = crate::transport::traits::JsonRpcNotification::new(
551                "notifications/resources/list_changed",
552                None,
553            );
554            let _ = tx.send(notification);
555        }
556    }
557
558    /// Send prompts/list_changed notification
559    fn notify_prompts_changed(&self) {
560        if let Some(ref tx) = self.notification_tx {
561            let notification = crate::transport::traits::JsonRpcNotification::new(
562                "notifications/prompts/list_changed",
563                None,
564            );
565            let _ = tx.send(notification);
566        }
567    }
568
569    // ==================== Profile Management ====================
570
571    /// Apply a session profile
572    ///
573    /// This adds/overrides/hides tools, resources, and prompts according to the profile
574    pub fn apply_profile(&self, profile: &SessionProfile) {
575        // Apply tool configuration
576        for tool in &profile.tool_extras {
577            self.add_tool(Arc::clone(tool));
578        }
579        for (name, tool) in &profile.tool_overrides {
580            self.override_tool(name.clone(), Arc::clone(tool));
581        }
582        for name in &profile.tool_hidden {
583            self.hide_tool(name.clone(), false); // no notify during profile apply
584        }
585        for (alias, target) in &profile.tool_aliases {
586            self.alias_tool(alias.clone(), target.clone());
587        }
588
589        // Apply resource configuration
590        for resource in &profile.resource_extras {
591            self.add_resource(Arc::clone(resource));
592        }
593        for (uri, resource) in &profile.resource_overrides {
594            self.override_resource(uri.clone(), Arc::clone(resource));
595        }
596        for uri in &profile.resource_hidden {
597            self.hide_resource(uri.clone());
598        }
599
600        // Apply prompt configuration
601        for prompt in &profile.prompt_extras {
602            self.add_prompt(Arc::clone(prompt));
603        }
604        for (name, prompt) in &profile.prompt_overrides {
605            self.override_prompt(name.clone(), Arc::clone(prompt));
606        }
607        for name in &profile.prompt_hidden {
608            self.hide_prompt(name.clone());
609        }
610    }
611
612    /// Clear all session customizations
613    pub fn clear_customizations(&self) {
614        // Clear tools
615        self.tool_overrides.clear();
616        self.tool_extras.clear();
617        if let Ok(mut hidden) = self.tool_hidden.write() {
618            hidden.clear();
619        }
620        if let Ok(mut aliases) = self.tool_aliases.write() {
621            aliases.clear();
622        }
623
624        // Clear resources
625        self.resource_overrides.clear();
626        self.resource_extras.clear();
627        if let Ok(mut hidden) = self.resource_hidden.write() {
628            hidden.clear();
629        }
630
631        // Clear prompts
632        self.prompt_overrides.clear();
633        self.prompt_extras.clear();
634        if let Ok(mut hidden) = self.prompt_hidden.write() {
635            hidden.clear();
636        }
637    }
638}
639
640// Implement Debug manually since trait objects don't implement Debug
641impl std::fmt::Debug for Session {
642    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
643        f.debug_struct("Session")
644            .field("id", &self.id)
645            .field("client_info", &self.client_info)
646            .field("capabilities", &self.capabilities)
647            .field("lifecycle", &self.lifecycle)
648            .field("error_count", &self.error_count())
649            .field("tool_overrides_count", &self.tool_overrides.len())
650            .field("tool_extras_count", &self.tool_extras.len())
651            .field("resource_overrides_count", &self.resource_overrides.len())
652            .field("resource_extras_count", &self.resource_extras.len())
653            .field("prompt_overrides_count", &self.prompt_overrides.len())
654            .field("prompt_extras_count", &self.prompt_extras.len())
655            .finish()
656    }
657}
658
659impl Default for Session {
660    fn default() -> Self {
661        Self::new()
662    }
663}
664
665/// Batch operations helper for making multiple changes with a single notification
666pub struct SessionBatch<'a> {
667    session: &'a Session,
668    tools_changed: bool,
669    resources_changed: bool,
670    prompts_changed: bool,
671}
672
673impl<'a> SessionBatch<'a> {
674    // Tool operations
675    pub fn add_tool(&mut self, tool: Arc<dyn crate::registry::tools::Tool>) {
676        self.session.add_tool(tool);
677        self.tools_changed = true;
678    }
679
680    pub fn override_tool(
681        &mut self,
682        name: impl Into<String>,
683        tool: Arc<dyn crate::registry::tools::Tool>,
684    ) {
685        self.session.override_tool(name, tool);
686        self.tools_changed = true;
687    }
688
689    pub fn remove_tool(&mut self, name: &str) -> Option<Arc<dyn crate::registry::tools::Tool>> {
690        let result = self.session.remove_tool(name);
691        if result.is_some() {
692            self.tools_changed = true;
693        }
694        result
695    }
696
697    pub fn hide_tool(&mut self, name: impl Into<String>) {
698        self.session.hide_tool(name, false); // batch sends notification at end
699        self.tools_changed = true;
700    }
701
702    pub fn unhide_tool(&mut self, name: &str) {
703        self.session.unhide_tool(name, false); // batch sends notification at end
704        self.tools_changed = true;
705    }
706
707    // Resource operations
708    pub fn add_resource(&mut self, resource: Arc<dyn crate::registry::resources::Resource>) {
709        self.session.add_resource(resource);
710        self.resources_changed = true;
711    }
712
713    pub fn override_resource(
714        &mut self,
715        uri: impl Into<String>,
716        resource: Arc<dyn crate::registry::resources::Resource>,
717    ) {
718        self.session.override_resource(uri, resource);
719        self.resources_changed = true;
720    }
721
722    pub fn remove_resource(
723        &mut self,
724        uri: &str,
725    ) -> Option<Arc<dyn crate::registry::resources::Resource>> {
726        let result = self.session.remove_resource(uri);
727        if result.is_some() {
728            self.resources_changed = true;
729        }
730        result
731    }
732
733    pub fn hide_resource(&mut self, uri: impl Into<String>) {
734        self.session.hide_resource(uri);
735        self.resources_changed = true;
736    }
737
738    pub fn unhide_resource(&mut self, uri: &str) {
739        self.session.unhide_resource(uri);
740        self.resources_changed = true;
741    }
742
743    // Prompt operations
744    pub fn add_prompt(&mut self, prompt: Arc<dyn crate::registry::prompts::Prompt>) {
745        self.session.add_prompt(prompt);
746        self.prompts_changed = true;
747    }
748
749    pub fn override_prompt(
750        &mut self,
751        name: impl Into<String>,
752        prompt: Arc<dyn crate::registry::prompts::Prompt>,
753    ) {
754        self.session.override_prompt(name, prompt);
755        self.prompts_changed = true;
756    }
757
758    pub fn remove_prompt(
759        &mut self,
760        name: &str,
761    ) -> Option<Arc<dyn crate::registry::prompts::Prompt>> {
762        let result = self.session.remove_prompt(name);
763        if result.is_some() {
764            self.prompts_changed = true;
765        }
766        result
767    }
768
769    pub fn hide_prompt(&mut self, name: impl Into<String>) {
770        self.session.hide_prompt(name);
771        self.prompts_changed = true;
772    }
773
774    pub fn unhide_prompt(&mut self, name: &str) {
775        self.session.unhide_prompt(name);
776        self.prompts_changed = true;
777    }
778}
779
780#[cfg(test)]
781mod tests {
782    use super::*;
783
784    #[test]
785    fn test_session_creation() {
786        let session = Session::new();
787        assert!(!session.id.is_empty());
788        assert_eq!(session.lifecycle(), SessionLifecycle::Created);
789        assert!(!session.is_initialized());
790        assert!(session.client_info.is_none());
791        assert!(session.capabilities.is_none());
792    }
793
794    #[test]
795    fn test_session_with_id() {
796        let session = Session::with_id("test-session");
797        assert_eq!(session.id, "test-session");
798        assert_eq!(session.lifecycle(), SessionLifecycle::Created);
799    }
800
801    #[test]
802    fn test_session_initialization() {
803        let mut session = Session::new();
804        let client_info = Implementation {
805            name: "test-client".to_string(),
806            version: "1.0.0".to_string(),
807        };
808        let capabilities = ClientCapabilities::default();
809
810        session.initialize(client_info.clone(), capabilities, "2025-06-18".to_string());
811
812        assert!(session.is_initialized());
813        assert_eq!(session.lifecycle(), SessionLifecycle::Ready);
814        assert_eq!(session.client_info.unwrap().name, "test-client");
815    }
816
817    #[test]
818    fn test_session_lifecycle_transitions() {
819        let mut session = Session::new();
820        assert_eq!(session.lifecycle(), SessionLifecycle::Created);
821        assert!(!session.lifecycle().can_accept_requests());
822
823        // Initialize -> Ready
824        session.initialize(
825            Implementation {
826                name: "test".into(),
827                version: "1.0".into(),
828            },
829            ClientCapabilities::default(),
830            "2025-06-18".to_string(),
831        );
832        assert_eq!(session.lifecycle(), SessionLifecycle::Ready);
833        assert!(session.lifecycle().can_accept_requests());
834        assert!(session.lifecycle().is_healthy());
835
836        // Record errors -> Degraded (after 3)
837        session.record_error();
838        assert_eq!(session.lifecycle(), SessionLifecycle::Ready);
839        session.record_error();
840        assert_eq!(session.lifecycle(), SessionLifecycle::Ready);
841        session.record_error();
842        assert_eq!(session.lifecycle(), SessionLifecycle::Degraded);
843        assert!(session.lifecycle().can_accept_requests());
844        assert!(!session.lifecycle().is_healthy());
845
846        // Record success -> Ready (recovery)
847        session.record_success();
848        assert_eq!(session.lifecycle(), SessionLifecycle::Ready);
849        assert_eq!(session.error_count(), 0);
850
851        // Close -> Closed
852        session.close();
853        assert_eq!(session.lifecycle(), SessionLifecycle::Closed);
854        assert!(!session.lifecycle().can_accept_requests());
855    }
856
857    #[test]
858    fn test_session_state() {
859        let session = Session::new();
860
861        // Set state
862        session.set_state("key1", Value::String("value1".to_string()));
863        session.set_state("key2", Value::Number(42.into()));
864
865        // Get state
866        assert_eq!(
867            session.get_state("key1"),
868            Some(Value::String("value1".to_string()))
869        );
870        assert_eq!(session.get_state("key2"), Some(Value::Number(42.into())));
871        assert_eq!(session.get_state("nonexistent"), None);
872
873        // State keys
874        let keys = session.state_keys();
875        assert_eq!(keys.len(), 2);
876        assert!(keys.contains(&"key1".to_string()));
877        assert!(keys.contains(&"key2".to_string()));
878
879        // Remove state
880        let removed = session.remove_state("key1");
881        assert_eq!(removed, Some(Value::String("value1".to_string())));
882        assert_eq!(session.get_state("key1"), None);
883
884        // Clear state
885        session.clear_state();
886        assert_eq!(session.state_keys().len(), 0);
887    }
888
889    #[test]
890    fn test_session_batch() {
891        let mut session = Session::new();
892        let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
893        session.set_notification_channel(tx);
894
895        // Test that batch operations track changes
896        session.batch(|batch| {
897            batch.hide_tool("unwanted");
898            batch.hide_resource("blocked://resource");
899            batch.hide_prompt("deprecated");
900        });
901
902        assert!(session.is_tool_hidden("unwanted"));
903        assert!(session.is_resource_hidden("blocked://resource"));
904        assert!(session.is_prompt_hidden("deprecated"));
905    }
906
907    #[test]
908    fn test_session_clone() {
909        let session1 = Session::with_id("test");
910        session1.set_state("shared", Value::Bool(true));
911
912        let session2 = session1.clone();
913
914        // Both sessions share the same state storage
915        assert_eq!(session1.id, session2.id);
916        assert_eq!(session2.get_state("shared"), Some(Value::Bool(true)));
917
918        // Modifying state in one affects the other
919        session2.set_state("shared", Value::Bool(false));
920        assert_eq!(session1.get_state("shared"), Some(Value::Bool(false)));
921    }
922}
mcp_host/server/session.rs

mcp_host/server/
session.rs
