cortexai_mcp/
session_resources.rs

1//! Session and Trajectory MCP Resources
2//!
3//! Exposes conversation sessions and execution traces as MCP resources,
4//! allowing external tools to browse agent history.
5//!
6//! # Resource URIs
7//!
8//! - `session://list` - List all sessions
9//! - `session://{session_id}` - Get a specific session
10//! - `session://{session_id}/messages` - Get session messages
11//! - `session://{session_id}/turns/{turn_number}` - Get a specific turn
12//!
13//! - `trace://list` - List all traces
14//! - `trace://{task_id}` - Get a specific trace
15//! - `trace://{task_id}/steps` - Get all steps in a trace
16//! - `trace://{task_id}/summary` - Get trace summary
17//!
18//! # Example
19//!
20//! ```rust,ignore
21//! use cortexai_mcp::{McpServer, SessionResourceHandler, TraceResourceHandler};
22//! use cortexai_agents::session::MemorySessionStore;
23//! use cortexai_agents::trajectory::TrajectoryStore;
24//!
25//! let session_store = Arc::new(MemorySessionStore::new());
26//! let trace_store = Arc::new(TrajectoryStore::new(1000));
27//!
28//! let server = McpServer::builder()
29//!     .name("agent-server")
30//!     .add_resource(SessionResourceHandler::new(session_store))
31//!     .add_resource(TraceResourceHandler::new(trace_store))
32//!     .build();
33//! ```
34
35use async_trait::async_trait;
36use serde::{Deserialize, Serialize};
37use std::sync::Arc;
38use tracing::debug;
39
40use crate::error::McpError;
41use crate::protocol::{McpResource, ResourceContent};
42use crate::server::ResourceHandler;
43
44// =============================================================================
45// Session Resource Handler
46// =============================================================================
47
48/// Resource handler for conversation sessions
49///
50/// Exposes sessions stored in a SessionStore as MCP resources.
51pub struct SessionResourceHandler<S: SessionStoreRead> {
52    store: Arc<S>,
53    /// Optional filter by user ID
54    user_filter: Option<String>,
55    /// Maximum sessions to list
56    max_list_size: usize,
57}
58
59/// Trait for reading sessions (subset of SessionStore)
60#[async_trait]
61pub trait SessionStoreRead: Send + Sync {
62    /// List all session IDs
63    async fn list_ids(&self) -> Result<Vec<String>, String>;
64
65    /// Load a session by ID and return as JSON
66    async fn load_json(&self, session_id: &str) -> Result<Option<String>, String>;
67
68    /// Get session metadata (id, user_id, created_at, updated_at, message_count)
69    async fn get_metadata(&self, session_id: &str) -> Result<Option<SessionMetadata>, String>;
70
71    /// Get messages for a session as JSON
72    async fn get_messages_json(&self, session_id: &str) -> Result<Option<String>, String>;
73
74    /// Get a specific turn as JSON
75    async fn get_turn_json(
76        &self,
77        session_id: &str,
78        turn_number: u32,
79    ) -> Result<Option<String>, String>;
80}
81
82/// Session metadata for listing
83#[derive(Debug, Clone, Serialize, Deserialize)]
84pub struct SessionMetadata {
85    pub id: String,
86    pub user_id: Option<String>,
87    pub created_at: u64,
88    pub updated_at: u64,
89    pub message_count: usize,
90    pub turn_count: usize,
91}
92
93impl<S: SessionStoreRead> SessionResourceHandler<S> {
94    /// Create a new session resource handler
95    pub fn new(store: Arc<S>) -> Self {
96        Self {
97            store,
98            user_filter: None,
99            max_list_size: 100,
100        }
101    }
102
103    /// Filter sessions by user ID
104    pub fn with_user_filter(mut self, user_id: impl Into<String>) -> Self {
105        self.user_filter = Some(user_id.into());
106        self
107    }
108
109    /// Set maximum sessions to list
110    pub fn with_max_list_size(mut self, size: usize) -> Self {
111        self.max_list_size = size;
112        self
113    }
114
115    /// Parse a session URI and extract components
116    fn parse_uri(uri: &str) -> Option<SessionUriParts> {
117        let uri = uri.strip_prefix("session://")?;
118
119        if uri == "list" {
120            return Some(SessionUriParts::List);
121        }
122
123        let parts: Vec<&str> = uri.split('/').collect();
124
125        match parts.as_slice() {
126            [session_id] => Some(SessionUriParts::Session(session_id.to_string())),
127            [session_id, "messages"] => Some(SessionUriParts::Messages(session_id.to_string())),
128            [session_id, "turns", turn_num] => {
129                let turn: u32 = turn_num.parse().ok()?;
130                Some(SessionUriParts::Turn(session_id.to_string(), turn))
131            }
132            _ => None,
133        }
134    }
135}
136
137#[derive(Debug)]
138enum SessionUriParts {
139    List,
140    Session(String),
141    Messages(String),
142    Turn(String, u32),
143}
144
145#[async_trait]
146impl<S: SessionStoreRead + 'static> ResourceHandler for SessionResourceHandler<S> {
147    fn list(&self) -> Vec<McpResource> {
148        vec![McpResource {
149            uri: "session://list".to_string(),
150            name: "Sessions".to_string(),
151            description: Some("List all conversation sessions".to_string()),
152            mime_type: Some("application/json".to_string()),
153        }]
154    }
155
156    async fn read(&self, uri: &str) -> Result<ResourceContent, McpError> {
157        debug!(uri = %uri, "Reading session resource");
158
159        let parts = Self::parse_uri(uri)
160            .ok_or_else(|| McpError::ResourceNotFound(format!("Invalid session URI: {}", uri)))?;
161
162        match parts {
163            SessionUriParts::List => {
164                let ids = self.store.list_ids().await.map_err(McpError::Internal)?;
165
166                let mut sessions = Vec::new();
167                for id in ids.into_iter().take(self.max_list_size) {
168                    if let Ok(Some(meta)) = self.store.get_metadata(&id).await {
169                        // Apply user filter if set
170                        if let Some(ref filter) = self.user_filter {
171                            if meta.user_id.as_ref() != Some(filter) {
172                                continue;
173                            }
174                        }
175                        sessions.push(meta);
176                    }
177                }
178
179                let json = serde_json::to_string_pretty(&sessions)
180                    .map_err(|e| McpError::Internal(e.to_string()))?;
181
182                Ok(ResourceContent {
183                    uri: uri.to_string(),
184                    mime_type: Some("application/json".to_string()),
185                    text: Some(json),
186                    blob: None,
187                })
188            }
189
190            SessionUriParts::Session(session_id) => {
191                let json = self
192                    .store
193                    .load_json(&session_id)
194                    .await
195                    .map_err(McpError::Internal)?
196                    .ok_or_else(|| {
197                        McpError::ResourceNotFound(format!("Session not found: {}", session_id))
198                    })?;
199
200                Ok(ResourceContent {
201                    uri: uri.to_string(),
202                    mime_type: Some("application/json".to_string()),
203                    text: Some(json),
204                    blob: None,
205                })
206            }
207
208            SessionUriParts::Messages(session_id) => {
209                let json = self
210                    .store
211                    .get_messages_json(&session_id)
212                    .await
213                    .map_err(McpError::Internal)?
214                    .ok_or_else(|| {
215                        McpError::ResourceNotFound(format!("Session not found: {}", session_id))
216                    })?;
217
218                Ok(ResourceContent {
219                    uri: uri.to_string(),
220                    mime_type: Some("application/json".to_string()),
221                    text: Some(json),
222                    blob: None,
223                })
224            }
225
226            SessionUriParts::Turn(session_id, turn_number) => {
227                let json = self
228                    .store
229                    .get_turn_json(&session_id, turn_number)
230                    .await
231                    .map_err(McpError::Internal)?
232                    .ok_or_else(|| {
233                        McpError::ResourceNotFound(format!(
234                            "Turn {} not found in session {}",
235                            turn_number, session_id
236                        ))
237                    })?;
238
239                Ok(ResourceContent {
240                    uri: uri.to_string(),
241                    mime_type: Some("application/json".to_string()),
242                    text: Some(json),
243                    blob: None,
244                })
245            }
246        }
247    }
248}
249
250// =============================================================================
251// Trace Resource Handler
252// =============================================================================
253
254/// Resource handler for execution traces
255///
256/// Exposes trajectories stored in a TrajectoryStore as MCP resources.
257pub struct TraceResourceHandler<T: TraceStoreRead> {
258    store: Arc<T>,
259    /// Optional filter by agent name
260    agent_filter: Option<String>,
261    /// Optional filter by success status
262    success_filter: Option<bool>,
263    /// Maximum traces to list
264    max_list_size: usize,
265}
266
267/// Trait for reading traces (subset of TrajectoryStore)
268#[async_trait]
269pub trait TraceStoreRead: Send + Sync {
270    /// List all task IDs
271    fn list_ids(&self) -> Vec<String>;
272
273    /// Get a trace by task ID as JSON
274    fn get_json(&self, task_id: &str) -> Option<String>;
275
276    /// Get trace metadata
277    fn get_metadata(&self, task_id: &str) -> Option<TraceMetadata>;
278
279    /// Get steps for a trace as JSON
280    fn get_steps_json(&self, task_id: &str) -> Option<String>;
281
282    /// Get trace summary as JSON
283    fn get_summary_json(&self, task_id: &str) -> Option<String>;
284
285    /// Filter by agent name
286    fn filter_by_agent(&self, agent_name: &str) -> Vec<String>;
287
288    /// Filter by success status
289    fn filter_by_success(&self, success: bool) -> Vec<String>;
290}
291
292/// Trace metadata for listing
293#[derive(Debug, Clone, Serialize, Deserialize)]
294pub struct TraceMetadata {
295    pub task_id: String,
296    pub agent_name: String,
297    pub success: bool,
298    pub total_duration_ms: u64,
299    pub step_count: usize,
300    pub llm_calls: usize,
301    pub tool_calls: usize,
302}
303
304impl<T: TraceStoreRead> TraceResourceHandler<T> {
305    /// Create a new trace resource handler
306    pub fn new(store: Arc<T>) -> Self {
307        Self {
308            store,
309            agent_filter: None,
310            success_filter: None,
311            max_list_size: 100,
312        }
313    }
314
315    /// Filter traces by agent name
316    pub fn with_agent_filter(mut self, agent_name: impl Into<String>) -> Self {
317        self.agent_filter = Some(agent_name.into());
318        self
319    }
320
321    /// Filter traces by success status
322    pub fn with_success_filter(mut self, success: bool) -> Self {
323        self.success_filter = Some(success);
324        self
325    }
326
327    /// Set maximum traces to list
328    pub fn with_max_list_size(mut self, size: usize) -> Self {
329        self.max_list_size = size;
330        self
331    }
332
333    /// Parse a trace URI and extract components
334    fn parse_uri(uri: &str) -> Option<TraceUriParts> {
335        let uri = uri.strip_prefix("trace://")?;
336
337        if uri == "list" {
338            return Some(TraceUriParts::List);
339        }
340
341        let parts: Vec<&str> = uri.split('/').collect();
342
343        match parts.as_slice() {
344            [task_id] => Some(TraceUriParts::Trace(task_id.to_string())),
345            [task_id, "steps"] => Some(TraceUriParts::Steps(task_id.to_string())),
346            [task_id, "summary"] => Some(TraceUriParts::Summary(task_id.to_string())),
347            _ => None,
348        }
349    }
350}
351
352#[derive(Debug)]
353enum TraceUriParts {
354    List,
355    Trace(String),
356    Steps(String),
357    Summary(String),
358}
359
360#[async_trait]
361impl<T: TraceStoreRead + 'static> ResourceHandler for TraceResourceHandler<T> {
362    fn list(&self) -> Vec<McpResource> {
363        vec![McpResource {
364            uri: "trace://list".to_string(),
365            name: "Traces".to_string(),
366            description: Some("List all execution traces".to_string()),
367            mime_type: Some("application/json".to_string()),
368        }]
369    }
370
371    async fn read(&self, uri: &str) -> Result<ResourceContent, McpError> {
372        debug!(uri = %uri, "Reading trace resource");
373
374        let parts = Self::parse_uri(uri)
375            .ok_or_else(|| McpError::ResourceNotFound(format!("Invalid trace URI: {}", uri)))?;
376
377        match parts {
378            TraceUriParts::List => {
379                // Get IDs with optional filters
380                let ids = if let Some(ref agent) = self.agent_filter {
381                    self.store.filter_by_agent(agent)
382                } else if let Some(success) = self.success_filter {
383                    self.store.filter_by_success(success)
384                } else {
385                    self.store.list_ids()
386                };
387
388                let mut traces = Vec::new();
389                for id in ids.into_iter().take(self.max_list_size) {
390                    if let Some(meta) = self.store.get_metadata(&id) {
391                        traces.push(meta);
392                    }
393                }
394
395                let json = serde_json::to_string_pretty(&traces)
396                    .map_err(|e| McpError::Internal(e.to_string()))?;
397
398                Ok(ResourceContent {
399                    uri: uri.to_string(),
400                    mime_type: Some("application/json".to_string()),
401                    text: Some(json),
402                    blob: None,
403                })
404            }
405
406            TraceUriParts::Trace(task_id) => {
407                let json = self.store.get_json(&task_id).ok_or_else(|| {
408                    McpError::ResourceNotFound(format!("Trace not found: {}", task_id))
409                })?;
410
411                Ok(ResourceContent {
412                    uri: uri.to_string(),
413                    mime_type: Some("application/json".to_string()),
414                    text: Some(json),
415                    blob: None,
416                })
417            }
418
419            TraceUriParts::Steps(task_id) => {
420                let json = self.store.get_steps_json(&task_id).ok_or_else(|| {
421                    McpError::ResourceNotFound(format!("Trace not found: {}", task_id))
422                })?;
423
424                Ok(ResourceContent {
425                    uri: uri.to_string(),
426                    mime_type: Some("application/json".to_string()),
427                    text: Some(json),
428                    blob: None,
429                })
430            }
431
432            TraceUriParts::Summary(task_id) => {
433                let json = self.store.get_summary_json(&task_id).ok_or_else(|| {
434                    McpError::ResourceNotFound(format!("Trace not found: {}", task_id))
435                })?;
436
437                Ok(ResourceContent {
438                    uri: uri.to_string(),
439                    mime_type: Some("application/json".to_string()),
440                    text: Some(json),
441                    blob: None,
442                })
443            }
444        }
445    }
446}
447
448// =============================================================================
449// In-Memory Implementations
450// =============================================================================
451
452/// In-memory session store adapter for MCP resources
453pub struct MemorySessionStoreAdapter {
454    sessions: parking_lot::RwLock<std::collections::HashMap<String, String>>,
455}
456
457impl MemorySessionStoreAdapter {
458    pub fn new() -> Self {
459        Self {
460            sessions: parking_lot::RwLock::new(std::collections::HashMap::new()),
461        }
462    }
463
464    /// Store a session (JSON serialized)
465    pub fn store(&self, session_id: &str, session_json: &str) {
466        self.sessions
467            .write()
468            .insert(session_id.to_string(), session_json.to_string());
469    }
470
471    /// Store a session object
472    pub fn store_session<T: Serialize>(&self, session_id: &str, session: &T) {
473        if let Ok(json) = serde_json::to_string(session) {
474            self.store(session_id, &json);
475        }
476    }
477}
478
479impl Default for MemorySessionStoreAdapter {
480    fn default() -> Self {
481        Self::new()
482    }
483}
484
485#[async_trait]
486impl SessionStoreRead for MemorySessionStoreAdapter {
487    async fn list_ids(&self) -> Result<Vec<String>, String> {
488        Ok(self.sessions.read().keys().cloned().collect())
489    }
490
491    async fn load_json(&self, session_id: &str) -> Result<Option<String>, String> {
492        Ok(self.sessions.read().get(session_id).cloned())
493    }
494
495    async fn get_metadata(&self, session_id: &str) -> Result<Option<SessionMetadata>, String> {
496        let sessions = self.sessions.read();
497        let json = match sessions.get(session_id) {
498            Some(j) => j,
499            None => return Ok(None),
500        };
501
502        // Parse just enough to get metadata
503        let value: serde_json::Value = serde_json::from_str(json).map_err(|e| e.to_string())?;
504
505        Ok(Some(SessionMetadata {
506            id: value["id"].as_str().unwrap_or(session_id).to_string(),
507            user_id: value["user_id"].as_str().map(|s| s.to_string()),
508            created_at: value["created_at"].as_u64().unwrap_or(0),
509            updated_at: value["updated_at"].as_u64().unwrap_or(0),
510            message_count: value["messages"].as_array().map(|a| a.len()).unwrap_or(0),
511            turn_count: value["turns"].as_array().map(|a| a.len()).unwrap_or(0),
512        }))
513    }
514
515    async fn get_messages_json(&self, session_id: &str) -> Result<Option<String>, String> {
516        let sessions = self.sessions.read();
517        let json = match sessions.get(session_id) {
518            Some(j) => j,
519            None => return Ok(None),
520        };
521
522        let value: serde_json::Value = serde_json::from_str(json).map_err(|e| e.to_string())?;
523
524        let messages = &value["messages"];
525        Ok(Some(
526            serde_json::to_string_pretty(messages).map_err(|e| e.to_string())?,
527        ))
528    }
529
530    async fn get_turn_json(
531        &self,
532        session_id: &str,
533        turn_number: u32,
534    ) -> Result<Option<String>, String> {
535        let sessions = self.sessions.read();
536        let json = match sessions.get(session_id) {
537            Some(j) => j,
538            None => return Ok(None),
539        };
540
541        let value: serde_json::Value = serde_json::from_str(json).map_err(|e| e.to_string())?;
542
543        let turns = value["turns"].as_array();
544        let turn = turns.and_then(|t| t.get(turn_number as usize - 1));
545
546        match turn {
547            Some(t) => Ok(Some(
548                serde_json::to_string_pretty(t).map_err(|e| e.to_string())?,
549            )),
550            None => Ok(None),
551        }
552    }
553}
554
555/// In-memory trace store adapter for MCP resources
556pub struct MemoryTraceStoreAdapter {
557    traces: parking_lot::RwLock<std::collections::HashMap<String, String>>,
558}
559
560impl MemoryTraceStoreAdapter {
561    pub fn new() -> Self {
562        Self {
563            traces: parking_lot::RwLock::new(std::collections::HashMap::new()),
564        }
565    }
566
567    /// Store a trace (JSON serialized)
568    pub fn store(&self, task_id: &str, trace_json: &str) {
569        self.traces
570            .write()
571            .insert(task_id.to_string(), trace_json.to_string());
572    }
573
574    /// Store a trace object
575    pub fn store_trace<T: Serialize>(&self, task_id: &str, trace: &T) {
576        if let Ok(json) = serde_json::to_string(trace) {
577            self.store(task_id, &json);
578        }
579    }
580}
581
582impl Default for MemoryTraceStoreAdapter {
583    fn default() -> Self {
584        Self::new()
585    }
586}
587
588impl TraceStoreRead for MemoryTraceStoreAdapter {
589    fn list_ids(&self) -> Vec<String> {
590        self.traces.read().keys().cloned().collect()
591    }
592
593    fn get_json(&self, task_id: &str) -> Option<String> {
594        self.traces.read().get(task_id).cloned()
595    }
596
597    fn get_metadata(&self, task_id: &str) -> Option<TraceMetadata> {
598        let traces = self.traces.read();
599        let json = traces.get(task_id)?;
600
601        let value: serde_json::Value = serde_json::from_str(json).ok()?;
602
603        Some(TraceMetadata {
604            task_id: value["task_id"].as_str().unwrap_or(task_id).to_string(),
605            agent_name: value["agent_name"]
606                .as_str()
607                .unwrap_or("unknown")
608                .to_string(),
609            success: value["success"].as_bool().unwrap_or(false),
610            total_duration_ms: value["total_duration_ms"].as_u64().unwrap_or(0),
611            step_count: value["steps"].as_array().map(|a| a.len()).unwrap_or(0),
612            llm_calls: value["steps"]
613                .as_array()
614                .map(|steps| {
615                    steps
616                        .iter()
617                        .filter(|s| s["step_type"] == "llm_call")
618                        .count()
619                })
620                .unwrap_or(0),
621            tool_calls: value["steps"]
622                .as_array()
623                .map(|steps| {
624                    steps
625                        .iter()
626                        .filter(|s| s["step_type"] == "tool_call")
627                        .count()
628                })
629                .unwrap_or(0),
630        })
631    }
632
633    fn get_steps_json(&self, task_id: &str) -> Option<String> {
634        let traces = self.traces.read();
635        let json = traces.get(task_id)?;
636
637        let value: serde_json::Value = serde_json::from_str(json).ok()?;
638        let steps = &value["steps"];
639
640        serde_json::to_string_pretty(steps).ok()
641    }
642
643    fn get_summary_json(&self, task_id: &str) -> Option<String> {
644        let meta = self.get_metadata(task_id)?;
645        serde_json::to_string_pretty(&meta).ok()
646    }
647
648    fn filter_by_agent(&self, agent_name: &str) -> Vec<String> {
649        self.traces
650            .read()
651            .iter()
652            .filter_map(|(id, json)| {
653                let value: serde_json::Value = serde_json::from_str(json).ok()?;
654                if value["agent_name"].as_str() == Some(agent_name) {
655                    Some(id.clone())
656                } else {
657                    None
658                }
659            })
660            .collect()
661    }
662
663    fn filter_by_success(&self, success: bool) -> Vec<String> {
664        self.traces
665            .read()
666            .iter()
667            .filter_map(|(id, json)| {
668                let value: serde_json::Value = serde_json::from_str(json).ok()?;
669                if value["success"].as_bool() == Some(success) {
670                    Some(id.clone())
671                } else {
672                    None
673                }
674            })
675            .collect()
676    }
677}
678
679// =============================================================================
680// Tests
681// =============================================================================
682
683#[cfg(test)]
684mod tests {
685    use super::*;
686
687    fn sample_session_json() -> String {
688        serde_json::json!({
689            "id": "session-1",
690            "user_id": "user-123",
691            "created_at": 1700000000,
692            "updated_at": 1700001000,
693            "messages": [
694                {"role": "user", "content": "Hello"},
695                {"role": "assistant", "content": "Hi there!"}
696            ],
697            "turns": [
698                {
699                    "number": 1,
700                    "user_message": {"role": "user", "content": "Hello"},
701                    "assistant_response": {"role": "assistant", "content": "Hi there!"}
702                }
703            ]
704        })
705        .to_string()
706    }
707
708    fn sample_trace_json() -> String {
709        serde_json::json!({
710            "task_id": "task-1",
711            "agent_name": "research-agent",
712            "success": true,
713            "total_duration_ms": 1500,
714            "steps": [
715                {"step_type": "llm_call", "duration_ms": 500, "success": true},
716                {"step_type": "tool_call", "duration_ms": 200, "success": true},
717                {"step_type": "llm_call", "duration_ms": 600, "success": true}
718            ],
719            "metadata": {}
720        })
721        .to_string()
722    }
723
724    #[tokio::test]
725    async fn test_session_resource_list() {
726        let store = Arc::new(MemorySessionStoreAdapter::new());
727        store.store("session-1", &sample_session_json());
728        store.store("session-2", &sample_session_json());
729
730        let handler = SessionResourceHandler::new(store);
731        let resources = handler.list();
732
733        assert_eq!(resources.len(), 1);
734        assert_eq!(resources[0].uri, "session://list");
735    }
736
737    #[tokio::test]
738    async fn test_session_resource_read_list() {
739        let store = Arc::new(MemorySessionStoreAdapter::new());
740        store.store("session-1", &sample_session_json());
741
742        let handler = SessionResourceHandler::new(store);
743        let content = handler.read("session://list").await.unwrap();
744
745        assert!(content.text.is_some());
746        let text = content.text.unwrap();
747        assert!(text.contains("session-1"));
748    }
749
750    #[tokio::test]
751    async fn test_session_resource_read_session() {
752        let store = Arc::new(MemorySessionStoreAdapter::new());
753        store.store("session-1", &sample_session_json());
754
755        let handler = SessionResourceHandler::new(store);
756        let content = handler.read("session://session-1").await.unwrap();
757
758        assert!(content.text.is_some());
759        let text = content.text.unwrap();
760        assert!(text.contains("user-123"));
761    }
762
763    #[tokio::test]
764    async fn test_session_resource_read_messages() {
765        let store = Arc::new(MemorySessionStoreAdapter::new());
766        store.store("session-1", &sample_session_json());
767
768        let handler = SessionResourceHandler::new(store);
769        let content = handler.read("session://session-1/messages").await.unwrap();
770
771        assert!(content.text.is_some());
772        let text = content.text.unwrap();
773        assert!(text.contains("Hello"));
774        assert!(text.contains("Hi there!"));
775    }
776
777    #[tokio::test]
778    async fn test_session_resource_read_turn() {
779        let store = Arc::new(MemorySessionStoreAdapter::new());
780        store.store("session-1", &sample_session_json());
781
782        let handler = SessionResourceHandler::new(store);
783        let content = handler.read("session://session-1/turns/1").await.unwrap();
784
785        assert!(content.text.is_some());
786        let text = content.text.unwrap();
787        assert!(text.contains("Hello"));
788    }
789
790    #[tokio::test]
791    async fn test_session_resource_not_found() {
792        let store = Arc::new(MemorySessionStoreAdapter::new());
793        let handler = SessionResourceHandler::new(store);
794
795        let result = handler.read("session://nonexistent").await;
796        assert!(result.is_err());
797    }
798
799    #[tokio::test]
800    async fn test_trace_resource_list() {
801        let store = Arc::new(MemoryTraceStoreAdapter::new());
802        store.store("task-1", &sample_trace_json());
803
804        let handler = TraceResourceHandler::new(store);
805        let resources = handler.list();
806
807        assert_eq!(resources.len(), 1);
808        assert_eq!(resources[0].uri, "trace://list");
809    }
810
811    #[tokio::test]
812    async fn test_trace_resource_read_list() {
813        let store = Arc::new(MemoryTraceStoreAdapter::new());
814        store.store("task-1", &sample_trace_json());
815
816        let handler = TraceResourceHandler::new(store);
817        let content = handler.read("trace://list").await.unwrap();
818
819        assert!(content.text.is_some());
820        let text = content.text.unwrap();
821        assert!(text.contains("task-1"));
822        assert!(text.contains("research-agent"));
823    }
824
825    #[tokio::test]
826    async fn test_trace_resource_read_trace() {
827        let store = Arc::new(MemoryTraceStoreAdapter::new());
828        store.store("task-1", &sample_trace_json());
829
830        let handler = TraceResourceHandler::new(store);
831        let content = handler.read("trace://task-1").await.unwrap();
832
833        assert!(content.text.is_some());
834        let text = content.text.unwrap();
835        assert!(text.contains("research-agent"));
836        assert!(text.contains("1500"));
837    }
838
839    #[tokio::test]
840    async fn test_trace_resource_read_steps() {
841        let store = Arc::new(MemoryTraceStoreAdapter::new());
842        store.store("task-1", &sample_trace_json());
843
844        let handler = TraceResourceHandler::new(store);
845        let content = handler.read("trace://task-1/steps").await.unwrap();
846
847        assert!(content.text.is_some());
848        let text = content.text.unwrap();
849        assert!(text.contains("llm_call"));
850        assert!(text.contains("tool_call"));
851    }
852
853    #[tokio::test]
854    async fn test_trace_resource_read_summary() {
855        let store = Arc::new(MemoryTraceStoreAdapter::new());
856        store.store("task-1", &sample_trace_json());
857
858        let handler = TraceResourceHandler::new(store);
859        let content = handler.read("trace://task-1/summary").await.unwrap();
860
861        assert!(content.text.is_some());
862        let text = content.text.unwrap();
863        assert!(text.contains("llm_calls"));
864        assert!(text.contains("tool_calls"));
865    }
866
867    #[tokio::test]
868    async fn test_trace_filter_by_agent() {
869        let store = Arc::new(MemoryTraceStoreAdapter::new());
870        store.store("task-1", &sample_trace_json());
871        store.store(
872            "task-2",
873            &serde_json::json!({
874                "task_id": "task-2",
875                "agent_name": "other-agent",
876                "success": true,
877                "total_duration_ms": 500,
878                "steps": []
879            })
880            .to_string(),
881        );
882
883        let handler = TraceResourceHandler::new(store).with_agent_filter("research-agent");
884        let content = handler.read("trace://list").await.unwrap();
885
886        let text = content.text.unwrap();
887        assert!(text.contains("task-1"));
888        assert!(!text.contains("task-2"));
889    }
890
891    #[tokio::test]
892    async fn test_trace_filter_by_success() {
893        let store = Arc::new(MemoryTraceStoreAdapter::new());
894        store.store("task-1", &sample_trace_json()); // success: true
895        store.store(
896            "task-2",
897            &serde_json::json!({
898                "task_id": "task-2",
899                "agent_name": "agent",
900                "success": false,
901                "total_duration_ms": 500,
902                "steps": []
903            })
904            .to_string(),
905        );
906
907        let handler = TraceResourceHandler::new(store).with_success_filter(false);
908        let content = handler.read("trace://list").await.unwrap();
909
910        let text = content.text.unwrap();
911        assert!(!text.contains("task-1"));
912        assert!(text.contains("task-2"));
913    }
914
915    #[test]
916    fn test_session_uri_parsing() {
917        assert!(matches!(
918            SessionResourceHandler::<MemorySessionStoreAdapter>::parse_uri("session://list"),
919            Some(SessionUriParts::List)
920        ));
921
922        assert!(matches!(
923            SessionResourceHandler::<MemorySessionStoreAdapter>::parse_uri("session://abc"),
924            Some(SessionUriParts::Session(id)) if id == "abc"
925        ));
926
927        assert!(matches!(
928            SessionResourceHandler::<MemorySessionStoreAdapter>::parse_uri("session://abc/messages"),
929            Some(SessionUriParts::Messages(id)) if id == "abc"
930        ));
931
932        assert!(matches!(
933            SessionResourceHandler::<MemorySessionStoreAdapter>::parse_uri("session://abc/turns/1"),
934            Some(SessionUriParts::Turn(id, 1)) if id == "abc"
935        ));
936
937        assert!(
938            SessionResourceHandler::<MemorySessionStoreAdapter>::parse_uri("invalid").is_none()
939        );
940    }
941
942    #[test]
943    fn test_trace_uri_parsing() {
944        assert!(matches!(
945            TraceResourceHandler::<MemoryTraceStoreAdapter>::parse_uri("trace://list"),
946            Some(TraceUriParts::List)
947        ));
948
949        assert!(matches!(
950            TraceResourceHandler::<MemoryTraceStoreAdapter>::parse_uri("trace://task-1"),
951            Some(TraceUriParts::Trace(id)) if id == "task-1"
952        ));
953
954        assert!(matches!(
955            TraceResourceHandler::<MemoryTraceStoreAdapter>::parse_uri("trace://task-1/steps"),
956            Some(TraceUriParts::Steps(id)) if id == "task-1"
957        ));
958
959        assert!(matches!(
960            TraceResourceHandler::<MemoryTraceStoreAdapter>::parse_uri("trace://task-1/summary"),
961            Some(TraceUriParts::Summary(id)) if id == "task-1"
962        ));
963    }
964}
cortexai_mcp/session_resources.rs

cortexai_mcp/
session_resources.rs
