PicoBot/src/channels/cli_chat.rs
2026-07-15 17:48:14 +08:00

869 lines
34 KiB
Rust

use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{Mutex, mpsc};
use crate::bus::{ControlMessage, InboundMessage, MessageBus, OutboundMessage};
use crate::protocol::{HistoryMessage, SlashCommandInfo, WsInbound, WsOutbound, parse_inbound};
use crate::session::{SessionCommand, SessionEvent, UnifiedSessionId};
use super::base::{Channel, ChannelError};
// ============================================================================
// Client - Connected CLI client
// ============================================================================
pub(crate) struct Client {
sender: mpsc::Sender<WsOutbound>,
chat_id: String,
current_session_id: Mutex<Option<String>>,
}
impl Client {
pub(crate) fn chat_id(&self) -> &str {
&self.chat_id
}
}
// ============================================================================
// CliChatChannel - Channel implementation for CLI chat
// ============================================================================
pub struct CliChatChannel {
bus: std::sync::Mutex<Option<Arc<MessageBus>>>,
clients: Mutex<HashMap<String, Arc<Client>>>,
}
impl Default for CliChatChannel {
fn default() -> Self {
Self::new()
}
}
impl CliChatChannel {
pub fn new() -> Self {
Self {
bus: std::sync::Mutex::new(None),
clients: Mutex::new(HashMap::new()),
}
}
/// Register a new client connection, returns (session_id, client)
pub(crate) async fn register_client(
&self,
sender: mpsc::Sender<WsOutbound>,
requested_chat_id: Option<String>,
) -> (String, Arc<Client>) {
// Each WebSocket connection gets a stable chat scope. All user input and
// dialog controls for this client stay inside that scope unless the
// protocol explicitly carries a full session id.
let chat_id = requested_chat_id.unwrap_or_else(crate::util::short_id);
let client = Arc::new(Client {
sender,
chat_id: chat_id.clone(),
current_session_id: Mutex::new(None),
});
self.clients
.lock()
.await
.insert(chat_id.clone(), client.clone());
// Resume the current/most-recent dialog for a stable TUI identity. Only
// create a dialog when this client has never connected before.
let session_id = match self.resume_session_via_control(&chat_id).await {
Ok(id) => id,
Err(e) => {
tracing::error!(error = %e, "Failed to resume initial session");
UnifiedSessionId::new("cli_chat", &chat_id, crate::util::short_id()).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)
}
pub(crate) async fn unregister_client(&self, client: &Arc<Client>) {
let mut clients = self.clients.lock().await;
if clients
.get(client.chat_id())
.is_some_and(|registered| Arc::ptr_eq(registered, client))
{
clients.remove(client.chat_id());
}
}
/// Push a structured task-plan update to the WebSocket client that owns
/// the session. Other channels remain unaffected.
pub async fn publish_plan_changed(&self, event: crate::work::PlanChanged) {
let Some(session_id) = UnifiedSessionId::parse(&event.session_id) else {
return;
};
if session_id.channel != "cli_chat" {
return;
}
let client = self.clients.lock().await.get(&session_id.chat_id).cloned();
if let Some(client) = client {
let _ = client
.sender
.send(WsOutbound::PlanUpdated {
session_id: event.session_id,
reason: event.reason,
changed_item_ids: event.changed_item_ids,
plan: event.plan,
})
.await;
}
}
/// Handle an inbound message from a client
pub(crate) async fn handle_inbound(&self, client: Arc<Client>, 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<Client>,
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, ..
} => {
// All messages (including slash commands) go through the normal inbound flow
// SessionManager handles session creation/reuse internally
let msg = InboundMessage {
channel: self.name().to_string(),
sender_id: "cli".to_string(),
chat_id: chat_id.unwrap_or_else(|| client.chat_id.clone()),
content,
timestamp: crate::bus::message::current_timestamp(),
media: Vec::new(),
metadata: Default::default(),
forwarded_metadata: Default::default(),
};
bus.publish_inbound(msg).await?;
}
WsInbound::ClearHistory {
chat_id,
session_id,
} => {
let (reply_tx, mut reply_rx) = mpsc::channel(1);
let session_id = if let Some(session_id) = session_id {
Self::parse_client_session(&client, &session_id)?
} else if let Some(chat_id) = chat_id {
if chat_id != client.chat_id {
return Err(ChannelError::Other(
"Chat does not belong to this client".to_string(),
));
}
let (current_tx, mut current_rx) = mpsc::channel(1);
bus.publish_control(ControlMessage {
op: SessionCommand::GetCurrentDialog {
channel: "cli_chat".to_string(),
chat_id,
},
reply_tx: current_tx,
})
.await?;
match current_rx.recv().await {
Some(Ok(SessionEvent::CurrentDialog {
session_id: Some(session_id),
})) => session_id,
Some(Ok(SessionEvent::CurrentDialog { session_id: None })) => {
return Err(ChannelError::Other("No active session".to_string()));
}
Some(Ok(_)) => {
return Err(ChannelError::Other(
"Unexpected response type".to_string(),
));
}
Some(Err(e)) => return Err(e),
None => {
return Err(ChannelError::Other("Control channel closed".to_string()));
}
}
} else {
let target = current_session_guard
.clone()
.ok_or_else(|| ChannelError::Other("No active session".to_string()))?;
Self::parse_client_session(&client, &target)?
};
let target = session_id.to_string();
bus.publish_control(ControlMessage {
op: SessionCommand::ClearHistory { session_id },
reply_tx,
})
.await?;
match reply_rx.recv().await {
Some(Ok(SessionEvent::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, created_title) = self
.create_session_via_control(&client.chat_id, title.as_deref())
.await?;
*current_session_guard = Some(new_id.clone());
let _ = client
.sender
.send(WsOutbound::SessionCreated {
session_id: new_id,
title: created_title,
})
.await;
}
WsInbound::ListSessions { include_archived } => {
// List dialogs for the current chat
let chat_id = client.chat_id.clone();
let (reply_tx, mut reply_rx) = mpsc::channel(1);
bus.publish_control(ControlMessage {
op: SessionCommand::ListDialogs {
channel: "cli_chat".to_string(),
chat_id,
include_archived,
},
reply_tx,
})
.await?;
match reply_rx.recv().await {
Some(Ok(SessionEvent::DialogList {
dialogs,
current_dialog_id,
})) => {
// Convert DialogInfo to SessionSummary for backward compatibility
let sessions: Vec<crate::protocol::SessionSummary> = dialogs
.into_iter()
.map(|d| crate::protocol::SessionSummary {
session_id: d.session_id.to_string(),
title: d.title,
channel_name: d.session_id.channel.clone(),
chat_id: d.session_id.chat_id,
message_count: d.message_count,
last_active_at: d.last_active_at,
archived_at: d.archived_at,
})
.collect();
let current_session_id = current_dialog_id.map(|did| {
UnifiedSessionId::new("cli_chat", &client.chat_id, &did).to_string()
});
if let Some(ref session_id) = current_session_id {
*current_session_guard = Some(session_id.clone());
}
let _ = client
.sender
.send(WsOutbound::SessionList {
sessions,
current_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::LoadSession { session_id } => {
let (reply_tx, mut reply_rx) = mpsc::channel(1);
let unified_id = UnifiedSessionId::parse(&session_id)
.ok_or_else(|| ChannelError::Other("Invalid session ID format".to_string()))?;
if unified_id.channel != "cli_chat" || unified_id.chat_id != client.chat_id {
return Err(ChannelError::Other(
"Session does not belong to this client".to_string(),
));
}
bus.publish_control(ControlMessage {
op: SessionCommand::SwitchDialog {
channel: unified_id.channel.clone(),
chat_id: unified_id.chat_id.clone(),
dialog_id: unified_id.dialog_id.clone(),
},
reply_tx,
})
.await?;
match reply_rx.recv().await {
Some(Ok(SessionEvent::DialogSwitched { session_id })) => {
*current_session_guard = Some(session_id.to_string());
let _ = client
.sender
.send(WsOutbound::SessionLoaded {
session_id: session_id.to_string(),
title: "Session".to_string(),
message_count: 0,
})
.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::GetSessionHistory { session_id, limit } => {
let unified_id = Self::parse_client_session(&client, &session_id)?;
let (reply_tx, mut reply_rx) = mpsc::channel(1);
bus.publish_control(ControlMessage {
op: SessionCommand::GetDialogHistory {
session_id: unified_id,
limit: limit.unwrap_or(1_000).clamp(1, 2_000),
},
reply_tx,
})
.await?;
match reply_rx.recv().await {
Some(Ok(SessionEvent::DialogHistory {
session_id,
messages,
})) => {
let messages = messages
.into_iter()
.filter(|message| !message.content.is_empty())
.map(|message| HistoryMessage {
id: message.id,
seq: message.seq,
role: message.role,
content: message.content,
created_at: message.created_at,
})
.collect();
let _ = client
.sender
.send(WsOutbound::SessionHistory {
session_id: session_id.to_string(),
messages,
})
.await;
}
Some(Ok(_)) => {}
Some(Err(e)) => return Err(e),
None => {
return Err(ChannelError::Other("Control channel closed".to_string()));
}
}
}
WsInbound::GetSessionPlan { session_id } => {
let unified_id = Self::parse_client_session(&client, &session_id)?;
let (reply_tx, mut reply_rx) = mpsc::channel(1);
bus.publish_control(ControlMessage {
op: SessionCommand::GetTaskPlan {
session_id: unified_id,
},
reply_tx,
})
.await?;
match reply_rx.recv().await {
Some(Ok(SessionEvent::TaskPlan { session_id, plan })) => {
let _ = client
.sender
.send(WsOutbound::SessionPlan {
session_id: session_id.to_string(),
plan,
})
.await;
}
Some(Ok(_)) => {}
Some(Err(error)) => return Err(error),
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);
let unified_id = Self::parse_client_session(&client, &target)?;
bus.publish_control(ControlMessage {
op: SessionCommand::RenameDialog {
session_id: unified_id,
title: title.clone(),
},
reply_tx,
})
.await?;
match reply_rx.recv().await {
Some(Ok(SessionEvent::DialogRenamed { session_id, title })) => {
let _ = client
.sender
.send(WsOutbound::SessionRenamed {
session_id: session_id.to_string(),
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 was_current = current_session_guard.as_deref() == Some(&target);
let (reply_tx, mut reply_rx) = mpsc::channel(1);
let unified_id = Self::parse_client_session(&client, &target)?;
bus.publish_control(ControlMessage {
op: SessionCommand::ArchiveDialog {
session_id: unified_id,
},
reply_tx,
})
.await?;
match reply_rx.recv().await {
Some(Ok(SessionEvent::DialogArchived { session_id })) => {
let _ = client
.sender
.send(WsOutbound::SessionArchived {
session_id: session_id.to_string(),
})
.await;
if was_current {
let (new_id, title) = self
.create_session_via_control(&client.chat_id, None)
.await?;
*current_session_guard = Some(new_id.clone());
let _ = client
.sender
.send(WsOutbound::SessionCreated {
session_id: new_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::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);
let unified_id = Self::parse_client_session(&client, &target)?;
bus.publish_control(ControlMessage {
op: SessionCommand::DeleteDialog {
session_id: unified_id,
},
reply_tx,
})
.await?;
match reply_rx.recv().await {
Some(Ok(SessionEvent::DialogDeleted { session_id })) => {
let _ = client
.sender
.send(WsOutbound::SessionDeleted {
session_id: session_id.to_string(),
})
.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, title)) =
self.create_session_via_control(&client.chat_id, None).await
{
*current_session_guard = Some(new_id.clone());
let _ = client
.sender
.send(WsOutbound::SessionCreated {
session_id: new_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::GetSlashCommands => {
// Get commands from session manager via control message
let (reply_tx, mut reply_rx) = mpsc::channel(1);
bus.publish_control(ControlMessage {
op: SessionCommand::GetSlashCommands {
channel: "cli_chat".to_string(),
chat_id: client.chat_id.clone(),
},
reply_tx,
})
.await?;
if let Some(result) = reply_rx.recv().await {
match result {
Ok(SessionEvent::SlashCommandsList { commands }) => {
// Convert to SlashCommand to SlashCommandInfo
let command_infos: Vec<SlashCommandInfo> = commands
.into_iter()
.map(|cmd| SlashCommandInfo {
name: cmd.name.to_string(),
description: cmd.description.to_string(),
aliases: cmd.aliases.iter().map(|&a| a.to_string()).collect(),
})
.collect();
let _ = client
.sender
.send(WsOutbound::SlashCommandsList {
commands: command_infos,
})
.await;
}
Ok(SessionEvent::Error { code, message }) => {
let _ = client
.sender
.send(WsOutbound::Error { code, message })
.await;
}
Err(e) => {
let _ = client
.sender
.send(WsOutbound::Error {
code: "GET_COMMANDS_ERROR".to_string(),
message: e.to_string(),
})
.await;
}
_ => {}
}
}
}
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,
chat_id: &str,
title: Option<&str>,
) -> Result<(String, String), ChannelError> {
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: SessionCommand::CreateDialog {
channel: "cli_chat".to_string(),
chat_id: chat_id.to_string(),
title: title.map(String::from),
},
reply_tx,
})
.await?;
match reply_rx.recv().await {
Some(Ok(SessionEvent::DialogCreated { session_id, title })) => {
Ok((session_id.to_string(), title))
}
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 fn resume_session_via_control(&self, chat_id: &str) -> Result<String, ChannelError> {
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: SessionCommand::GetCurrentDialog {
channel: "cli_chat".to_string(),
chat_id: chat_id.to_string(),
},
reply_tx,
})
.await?;
if let Some(Ok(SessionEvent::CurrentDialog {
session_id: Some(session_id),
})) = reply_rx.recv().await
{
return Ok(session_id.to_string());
}
let (reply_tx, mut reply_rx) = mpsc::channel(1);
bus.publish_control(ControlMessage {
op: SessionCommand::ListDialogs {
channel: "cli_chat".to_string(),
chat_id: chat_id.to_string(),
include_archived: false,
},
reply_tx,
})
.await?;
if let Some(Ok(SessionEvent::DialogList { dialogs, .. })) = reply_rx.recv().await
&& let Some(dialog) = dialogs.first()
{
let session_id = dialog.session_id.clone();
let (reply_tx, mut reply_rx) = mpsc::channel(1);
bus.publish_control(ControlMessage {
op: SessionCommand::SwitchDialog {
channel: session_id.channel.clone(),
chat_id: session_id.chat_id.clone(),
dialog_id: session_id.dialog_id.clone(),
},
reply_tx,
})
.await?;
if let Some(Ok(SessionEvent::DialogSwitched { session_id })) = reply_rx.recv().await {
return Ok(session_id.to_string());
}
}
self.create_session_via_control(chat_id, None)
.await
.map(|(session_id, _)| session_id)
}
fn parse_client_session(
client: &Client,
session_id: &str,
) -> Result<UnifiedSessionId, ChannelError> {
let unified_id = UnifiedSessionId::parse(session_id)
.ok_or_else(|| ChannelError::Other("Invalid session ID format".to_string()))?;
if unified_id.channel != "cli_chat" || unified_id.chat_id != client.chat_id {
return Err(ChannelError::Other(
"Session does not belong to this client".to_string(),
));
}
Ok(unified_id)
}
}
#[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<MessageBus>) -> Result<(), ChannelError> {
*self.bus.lock().unwrap() = Some(bus);
Ok(())
}
async fn stop(&self) -> Result<(), ChannelError> {
*self.bus.lock().unwrap() = None;
self.clients.lock().await.clear();
Ok(())
}
async fn send(&self, msg: OutboundMessage) -> Result<(), ChannelError> {
let client = self.clients.lock().await.get(&msg.chat_id).cloned();
let Some(client) = client else {
tracing::debug!(chat_id = %msg.chat_id, "No active CLI client for outbound message");
return Ok(());
};
let message_type = msg.metadata.get("_type").map(String::as_str);
let session_id = msg.metadata.get("_session_id").cloned();
let outbound = if message_type == Some("notification") {
WsOutbound::SystemNotification {
content: msg.content,
session_id,
}
} else if message_type == Some("command") {
WsOutbound::CommandExecuted {
message: msg.content,
}
} else {
WsOutbound::AssistantResponse {
id: crate::util::short_id(),
content: msg.content,
role: "assistant".to_string(),
session_id,
}
};
if client.sender.send(outbound).await.is_err() {
let mut clients = self.clients.lock().await;
if clients
.get(&msg.chat_id)
.is_some_and(|registered| Arc::ptr_eq(registered, &client))
{
clients.remove(&msg.chat_id);
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn failed_sender_is_pruned_from_client_registry() {
let channel = CliChatChannel::new();
let (sender, receiver) = mpsc::channel(1);
drop(receiver);
let client = Arc::new(Client {
sender,
chat_id: "dead-client".to_string(),
current_session_id: Mutex::new(None),
});
channel
.clients
.lock()
.await
.insert("dead-client".to_string(), client);
channel
.send(OutboundMessage {
channel: "cli_chat".to_string(),
chat_id: "dead-client".to_string(),
content: "message".to_string(),
reply_to: None,
media: Vec::new(),
metadata: Default::default(),
delivery: None,
})
.await
.unwrap();
assert!(channel.clients.lock().await.is_empty());
}
#[tokio::test]
async fn stale_connection_cannot_unregister_replacement() {
let channel = CliChatChannel::new();
let (old_sender, _old_receiver) = mpsc::channel(1);
let (new_sender, _new_receiver) = mpsc::channel(1);
let old = Arc::new(Client {
sender: old_sender,
chat_id: "stable-client".to_string(),
current_session_id: Mutex::new(None),
});
let replacement = Arc::new(Client {
sender: new_sender,
chat_id: "stable-client".to_string(),
current_session_id: Mutex::new(None),
});
channel
.clients
.lock()
.await
.insert("stable-client".to_string(), replacement.clone());
channel.unregister_client(&old).await;
let registered = channel.clients.lock().await;
assert!(
registered
.get("stable-client")
.is_some_and(|client| Arc::ptr_eq(client, &replacement))
);
}
}