use std::sync::Arc; use async_trait::async_trait; use tokio::sync::{mpsc, Mutex}; use uuid::Uuid; use crate::bus::{ControlInbound, ControlMessage, ControlOutbound, InboundMessage, MessageBus, OutboundMessage}; use crate::protocol::{parse_inbound, WsInbound, WsOutbound}; use super::base::{Channel, ChannelError}; // ============================================================================ // Client - Connected CLI client // ============================================================================ pub(crate) struct Client { sender: mpsc::Sender, current_session_id: Mutex>, } // ============================================================================ // CliChatChannel - Channel implementation for CLI chat // ============================================================================ pub struct CliChatChannel { bus: std::sync::Mutex>>, clients: Mutex>>, } impl CliChatChannel { pub fn new() -> Self { Self { bus: std::sync::Mutex::new(None), clients: Mutex::new(Vec::new()), } } /// Register a new client connection, returns (session_id, client) pub(crate) async fn register_client(&self, sender: mpsc::Sender) -> (String, Arc) { let client = Arc::new(Client { sender, current_session_id: Mutex::new(None), }); self.clients.lock().await.push(client.clone()); // Create initial session via control message let session_id = match self.create_session_via_control(None).await { Ok(id) => id, Err(e) => { tracing::error!(error = %e, "Failed to create initial session"); Uuid::new_v4().to_string() } }; // Set current session id in client { let mut current = client.current_session_id.lock().await; *current = Some(session_id.clone()); } (session_id, client) } /// Handle an inbound message from a client pub(crate) async fn handle_inbound(&self, client: Arc, raw_msg: &str) { match parse_inbound(raw_msg) { Ok(inbound) => { match self.handle_ws_inbound(client.clone(), inbound).await { Ok(()) => {} Err(e) => { tracing::warn!(error = %e, "Failed to handle inbound message"); let _ = client .sender .send(WsOutbound::Error { code: "INTERNAL_ERROR".to_string(), message: e.to_string(), }) .await; } } } Err(e) => { tracing::warn!(error = %e, "Failed to parse inbound message"); let _ = client .sender .send(WsOutbound::Error { code: "PARSE_ERROR".to_string(), message: e.to_string(), }) .await; } } } async fn handle_ws_inbound(&self, client: Arc, inbound: WsInbound) -> Result<(), ChannelError> { let bus = { let guard = self.bus.lock().unwrap(); guard.clone().ok_or_else(|| ChannelError::Other("Channel not started".to_string()))? }; let mut current_session_guard = client.current_session_id.lock().await; match inbound { WsInbound::UserInput { content, chat_id, .. } => { let chat_id = chat_id.or(current_session_guard.clone()).unwrap_or_else(|| Uuid::new_v4().to_string()); // If no session, create one first if current_session_guard.is_none() { let new_id = self.create_session_via_control(None).await?; *current_session_guard = Some(new_id); } let session_id = current_session_guard.clone().unwrap(); // Publish to bus for AI processing let msg = InboundMessage { channel: self.name().to_string(), sender_id: "cli".to_string(), chat_id: session_id.clone(), content, timestamp: current_timestamp(), media: Vec::new(), metadata: Default::default(), forwarded_metadata: Default::default(), }; bus.publish_inbound(msg).await?; } WsInbound::ClearHistory { chat_id, session_id } => { let target = session_id .or(chat_id) .or(current_session_guard.clone()) .ok_or_else(|| ChannelError::Other("No active session".to_string()))?; let (reply_tx, mut reply_rx) = mpsc::channel(1); bus.publish_control(ControlMessage { op: ControlInbound::ClearHistory { session_id: target.clone() }, reply_tx, }).await?; match reply_rx.recv().await { Some(Ok(ControlOutbound::HistoryCleared { .. })) => { let _ = client .sender .send(WsOutbound::HistoryCleared { session_id: target }) .await; } Some(Ok(_)) => { // Unexpected response type, ignore } Some(Err(e)) => { return Err(e); } None => { return Err(ChannelError::Other("Control channel closed".to_string())); } } } WsInbound::CreateSession { title } => { let new_id = self.create_session_via_control(title.as_deref()).await?; *current_session_guard = Some(new_id.clone()); let _ = client .sender .send(WsOutbound::SessionCreated { session_id: new_id, title: title.unwrap_or_default(), }) .await; } WsInbound::ListSessions { include_archived } => { let (reply_tx, mut reply_rx) = mpsc::channel(1); bus.publish_control(ControlMessage { op: ControlInbound::ListSessions { include_archived }, reply_tx, }).await?; match reply_rx.recv().await { Some(Ok(ControlOutbound::SessionList { sessions })) => { let _ = client .sender .send(WsOutbound::SessionList { sessions, current_session_id: current_session_guard.clone(), }) .await; } Some(Ok(_)) => { // Unexpected response type } Some(Err(e)) => { return Err(e); } None => { return Err(ChannelError::Other("Control channel closed".to_string())); } } } WsInbound::LoadSession { session_id } => { let (reply_tx, mut reply_rx) = mpsc::channel(1); bus.publish_control(ControlMessage { op: ControlInbound::LoadSession { session_id: session_id.clone() }, reply_tx, }).await?; match reply_rx.recv().await { Some(Ok(ControlOutbound::SessionLoaded { session_id, title, message_count })) => { *current_session_guard = Some(session_id.clone()); let _ = client .sender .send(WsOutbound::SessionLoaded { session_id, title, message_count, }) .await; } Some(Ok(_)) => { // Unexpected response type } Some(Err(e)) => { let _ = client .sender .send(WsOutbound::Error { code: "SESSION_NOT_FOUND".to_string(), message: format!("Session not found: {}", session_id), }) .await; } None => { return Err(ChannelError::Other("Control channel closed".to_string())); } } } WsInbound::RenameSession { session_id, title } => { let target = session_id.or(current_session_guard.clone()).ok_or_else(|| { ChannelError::Other("No active session".to_string()) })?; let (reply_tx, mut reply_rx) = mpsc::channel(1); bus.publish_control(ControlMessage { op: ControlInbound::RenameSession { session_id: target.clone(), title: title.clone() }, reply_tx, }).await?; match reply_rx.recv().await { Some(Ok(ControlOutbound::SessionRenamed { session_id, title })) => { let _ = client .sender .send(WsOutbound::SessionRenamed { session_id, title }) .await; } Some(Ok(_)) => { // Unexpected response type } Some(Err(e)) => { return Err(e); } None => { return Err(ChannelError::Other("Control channel closed".to_string())); } } } WsInbound::ArchiveSession { session_id } => { let target = session_id.or(current_session_guard.clone()).ok_or_else(|| { ChannelError::Other("No active session".to_string()) })?; let (reply_tx, mut reply_rx) = mpsc::channel(1); bus.publish_control(ControlMessage { op: ControlInbound::ArchiveSession { session_id: target.clone() }, reply_tx, }).await?; match reply_rx.recv().await { Some(Ok(ControlOutbound::SessionArchived { session_id })) => { let _ = client .sender .send(WsOutbound::SessionArchived { session_id }) .await; } Some(Ok(_)) => { // Unexpected response type } Some(Err(e)) => { return Err(e); } None => { return Err(ChannelError::Other("Control channel closed".to_string())); } } } WsInbound::DeleteSession { session_id } => { let target = session_id.or(current_session_guard.clone()).ok_or_else(|| { ChannelError::Other("No active session".to_string()) })?; let (reply_tx, mut reply_rx) = mpsc::channel(1); bus.publish_control(ControlMessage { op: ControlInbound::DeleteSession { session_id: target.clone() }, reply_tx, }).await?; match reply_rx.recv().await { Some(Ok(ControlOutbound::SessionDeleted { session_id })) => { let _ = client .sender .send(WsOutbound::SessionDeleted { session_id: session_id.clone() }) .await; // If deleting current session, create a new one if current_session_guard.as_deref() == Some(&target) { drop(reply_rx); if let Ok(new_id) = self.create_session_via_control(None).await { *current_session_guard = Some(new_id.clone()); let _ = client .sender .send(WsOutbound::SessionCreated { session_id: new_id, title: String::new(), }) .await; } } } Some(Ok(_)) => { // Unexpected response type } Some(Err(e)) => { return Err(e); } None => { return Err(ChannelError::Other("Control channel closed".to_string())); } } } WsInbound::Ping => { let _ = client.sender.send(WsOutbound::Pong).await; } } Ok(()) } /// Create a session via control message and return the session_id async fn create_session_via_control(&self, title: Option<&str>) -> Result { let bus = { let guard = self.bus.lock().unwrap(); guard.clone().ok_or_else(|| ChannelError::Other("Channel not started".to_string()))? }; let (reply_tx, mut reply_rx) = mpsc::channel(1); bus.publish_control(ControlMessage { op: ControlInbound::CreateSession { title: title.map(String::from) }, reply_tx, }).await?; match reply_rx.recv().await { Some(Ok(ControlOutbound::SessionCreated { session_id, .. })) => { Ok(session_id) } Some(Ok(_)) => { Err(ChannelError::Other("Unexpected response type".to_string())) } Some(Err(e)) => Err(e), None => Err(ChannelError::Other("Control channel closed".to_string())), } } } #[async_trait] impl Channel for CliChatChannel { fn name(&self) -> &str { "cli_chat" } fn is_running(&self) -> bool { self.bus.lock().unwrap().is_some() } async fn start(&self, bus: Arc) -> Result<(), ChannelError> { *self.bus.lock().unwrap() = Some(bus); Ok(()) } async fn stop(&self) -> Result<(), ChannelError> { *self.bus.lock().unwrap() = None; Ok(()) } async fn send(&self, msg: OutboundMessage) -> Result<(), ChannelError> { let clients = self.clients.lock().await.clone(); for client in clients { let outbound = WsOutbound::AssistantResponse { id: Uuid::new_v4().to_string(), content: msg.content.clone(), role: "assistant".to_string(), }; let _ = client.sender.send(outbound).await; } Ok(()) } } fn current_timestamp() -> i64 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap() .as_millis() as i64 }