PicoBot/src/gateway/session.rs

68 lines
1.5 KiB
Rust

use std::sync::Arc;
use tokio::sync::{Mutex, mpsc};
use uuid::Uuid;
use crate::config::LLMProviderConfig;
use crate::agent::AgentLoop;
use crate::protocol::WsOutbound;
pub struct Session {
pub id: Uuid,
pub agent_loop: Arc<Mutex<AgentLoop>>,
pub user_tx: mpsc::Sender<WsOutbound>,
}
impl Session {
pub async fn new(
provider_config: LLMProviderConfig,
user_tx: mpsc::Sender<WsOutbound>,
) -> Result<Self, crate::agent::AgentError> {
let agent_loop = AgentLoop::new(provider_config)?;
Ok(Self {
id: Uuid::new_v4(),
agent_loop: Arc::new(Mutex::new(agent_loop)),
user_tx,
})
}
pub async fn send(&self, msg: WsOutbound) {
let _ = self.user_tx.send(msg).await;
}
}
use std::collections::HashMap;
use std::sync::RwLock;
pub struct SessionManager {
sessions: RwLock<HashMap<Uuid, Arc<Session>>>,
}
impl SessionManager {
pub fn new() -> Self {
Self {
sessions: RwLock::new(HashMap::new()),
}
}
pub fn add(&self, session: Arc<Session>) {
self.sessions.write().unwrap().insert(session.id, session);
}
pub fn remove(&self, id: &Uuid) {
self.sessions.write().unwrap().remove(id);
}
pub fn get(&self, id: &Uuid) -> Option<Arc<Session>> {
self.sessions.read().unwrap().get(id).cloned()
}
pub fn len(&self) -> usize {
self.sessions.read().unwrap().len()
}
}
impl Default for SessionManager {
fn default() -> Self {
Self::new()
}
}