agent_core/controller/session/
manager.rs

1// This implements a session manager that handles multiple sessions
2
3use std::collections::HashMap;
4use std::sync::Arc;
5
6use tokio::sync::{mpsc, RwLock};
7use tokio_util::sync::CancellationToken;
8
9use super::config::LLMSessionConfig;
10use super::LLMSession;
11use crate::client::error::LlmError;
12use crate::controller::types::FromLLMPayload;
13
14/// Manages multiple LLM sessions
15pub struct LLMSessionManager {
16    /// Map of session ID to session instance
17    sessions: RwLock<HashMap<i64, Arc<LLMSession>>>,
18}
19
20impl LLMSessionManager {
21    /// Creates a new session manager
22    pub fn new() -> Self {
23        Self {
24            sessions: RwLock::new(HashMap::new()),
25        }
26    }
27
28    /// Creates a new LLM session and starts it.
29    ///
30    /// # Arguments
31    /// * `config` - Session configuration (includes model, API key, etc.)
32    /// * `from_llm` - Channel sender for responses from the LLM
33    ///
34    /// # Returns
35    /// The session ID of the newly created session
36    ///
37    /// # Errors
38    /// Returns an error if the session fails to initialize (e.g., TLS setup failure)
39    pub async fn create_session(
40        &self,
41        config: LLMSessionConfig,
42        from_llm: mpsc::Sender<FromLLMPayload>,
43    ) -> Result<i64, LlmError> {
44        let cancel_token = CancellationToken::new();
45        let session = Arc::new(LLMSession::new(config, from_llm, cancel_token)?);
46        let session_id = session.id();
47
48        // Store the session
49        {
50            let mut sessions = self.sessions.write().await;
51            sessions.insert(session_id, Arc::clone(&session));
52        }
53
54        // Spawn the session's processing loop
55        let session_clone = Arc::clone(&session);
56        tokio::spawn(async move {
57            session_clone.start().await;
58        });
59
60        tracing::info!(session_id, "Session created and started");
61        Ok(session_id)
62    }
63
64    /// Retrieves a session by its ID.
65    ///
66    /// # Returns
67    /// The session if found, None otherwise
68    pub async fn get_session_by_id(&self, session_id: i64) -> Option<Arc<LLMSession>> {
69        let sessions = self.sessions.read().await;
70        sessions.get(&session_id).cloned()
71    }
72
73    /// Removes and shuts down a specific session.
74    ///
75    /// # Arguments
76    /// * `session_id` - The ID of the session to remove
77    ///
78    /// # Returns
79    /// true if the session was found and removed, false otherwise
80    pub async fn remove_session(&self, session_id: i64) -> bool {
81        let session = {
82            let mut sessions = self.sessions.write().await;
83            sessions.remove(&session_id)
84        };
85
86        if let Some(session) = session {
87            session.shutdown();
88            tracing::info!(session_id, "Session removed");
89            true
90        } else {
91            false
92        }
93    }
94
95    /// Shuts down all sessions managed by this manager.
96    /// This is idempotent and safe to call multiple times.
97    pub async fn shutdown(&self) {
98        let sessions: Vec<Arc<LLMSession>> = {
99            let mut sessions = self.sessions.write().await;
100            sessions.drain().map(|(_, s)| s).collect()
101        };
102
103        for session in sessions {
104            session.shutdown();
105        }
106
107        tracing::info!("Session manager shutdown complete");
108    }
109
110    /// Returns the number of active sessions
111    pub async fn session_count(&self) -> usize {
112        let sessions = self.sessions.read().await;
113        sessions.len()
114    }
115}
116
117impl Default for LLMSessionManager {
118    fn default() -> Self {
119        Self::new()
120    }
121}
agent_core/controller/session/manager.rs

agent_core/controller/session/
manager.rs
