PicoBot/src/client/mod.rs

pub use crate::protocol::{WsInbound, WsOutbound, serialize_inbound, serialize_outbound};

use crate::command::adapter::InputAdapter;
use crate::command::adapters::cli::CliInputAdapter;
use crate::command::context::AdapterContext;
use futures_util::{SinkExt, StreamExt};
use tokio_tungstenite::{connect_async, tungstenite::Message};

use crate::cli::{InputCommand, InputEvent, InputHandler};

fn parse_message(raw: &str) -> Result<WsOutbound, serde_json::Error> {
    serde_json::from_str(raw)
}

fn format_json(value: &serde_json::Value) -> String {
    serde_json::to_string_pretty(value).unwrap_or_else(|_| value.to_string())
}

pub async fn run(gateway_url: &str) -> Result<(), Box<dyn std::error::Error>> {
    let (ws_stream, _) = connect_async(gateway_url).await?;
    tracing::info!(url = %gateway_url, "Connected to gateway");

    let (mut sender, mut receiver) = ws_stream.split();

    let mut input = InputHandler::new();
    let mut current_session_id: Option<String> = None;
    // Track message IDs that were already streamed so we can skip
    // the duplicate AssistantResponse that arrives afterwards.
    let mut streamed_message_ids: std::collections::HashSet<String> =
        std::collections::HashSet::new();
    input
        .write_output("picobot CLI - Commands: /new [title], /save [filepath], /quit\n")
        .await?;

    // Main loop: poll both stdin and WebSocket
    loop {
        tokio::select! {
            // Handle WebSocket messages
            msg = receiver.next() => {
                match msg {
                    Some(Ok(Message::Text(text))) => {
                        let text = text.to_string();
                        if let Ok(outbound) = parse_message(&text) {
                            match outbound {
                                WsOutbound::AssistantResponse { id, content, .. }
                                    // Skip if already fully streamed via StreamDelta
                                    if !streamed_message_ids.remove(&id) => {
                                        input.write_response(&content).await?;
                                    }
                                WsOutbound::ToolCall { tool_name, arguments, .. } => {
                                    input.write_output(&format!("Tool call: {}\n{}\n", tool_name, format_json(&arguments))).await?;
                                }
                                WsOutbound::ToolResult { tool_name, content, .. } => {
                                    input.write_output(&format!("Tool result: {}\n{}\n", tool_name, content)).await?;
                                }
                                WsOutbound::ToolPending { tool_name, content, resume_hint, .. } => {
                                    input.write_output(&format!("Tool pending: {}\n{}\n{}\n", tool_name, content, resume_hint)).await?;
                                }
                                WsOutbound::Error { message, .. } => {
                                    input.write_output(&format!("Error: {}", message)).await?;
                                }
                                WsOutbound::SessionEstablished { session_id } => {
                                    current_session_id = Some(session_id.clone());
                                    #[cfg(debug_assertions)]
                                    tracing::debug!(session_id = %session_id, "Session established");
                                    input.write_output(&format!("Session: {}\n", session_id)).await?;
                                }
                                WsOutbound::SessionCreated { session_id, title } => {
                                    current_session_id = Some(session_id.clone());
                                    input.write_output(&format!("Created session: {} ({})\n", session_id, title)).await?;
                                }
                                WsOutbound::SessionSaved { session_id, filepath } => {
                                    input.write_output(&format!("Saved session {} to: {}\n", session_id, filepath)).await?;
                                }
                                WsOutbound::StreamDelta { id, delta, .. } => {
                                    // Track that this message is being streamed
                                    streamed_message_ids.insert(id.clone());
                                    // 在终端直接输出流式增量文本
                                    if !delta.is_empty() {
                                        input.write_output(&delta).await?;
                                    }
                                }
                                WsOutbound::StreamEnd { .. } => {
                                    // 流式结束，输出换行
                                    input.write_output("\n").await?;
                                }
                                _ => {}
                            }
                        }
                    }
                    Some(Ok(Message::Close(_))) | None => {
                        tracing::info!("Gateway disconnected");
                        input.write_output("Gateway disconnected").await?;
                        break;
                    }
                    _ => {}
                }
            }
            // Handle stdin input
            result = input.read_input("> ") => {
                match result {
                    Ok(Some(event)) => {
                        match event {
                            InputEvent::Command(InputCommand::Exit) => {
                                input.write_output("Goodbye!").await?;
                                break;
                            }
                            InputEvent::Command(InputCommand::New(title)) => {
                                // 使用 CliInputAdapter 构建 Command
                                let adapter = CliInputAdapter::new();
                                let ctx = AdapterContext::new("cli")
                                    .with_session_id(current_session_id.as_deref().unwrap_or(""));

                                // 构建输入字符串
                                let input_str = match title {
                                    Some(t) => format!("/new {}", t),
                                    None => "/new".to_string(),
                                };

                                // 解析为 Command
                                match adapter.try_parse(&input_str, ctx) {
                                    Ok(Some(command)) => {
                                        // 序列化为 JSON
                                        let json = serde_json::to_string(&command).unwrap_or_default();
                                        // 通过 Command 消息发送
                                        let inbound = WsInbound::Command { payload: json };
                                        if let Ok(text) = serialize_inbound(&inbound) {
                                            let _ = sender.send(Message::Text(text.into())).await;
                                        }
                                    }
                                    Ok(None) => {
                                        tracing::warn!("Failed to parse /new command");
                                    }
                                    Err(e) => {
                                        tracing::error!(error = %e, "Error parsing /new command");
                                    }
                                }
                                continue;
                            }
                            InputEvent::Command(InputCommand::Save(filepath)) => {
                                // 使用 CliInputAdapter 构建 Command
                                let adapter = CliInputAdapter::new();
                                let ctx = AdapterContext::new("cli")
                                    .with_session_id(current_session_id.as_deref().unwrap_or(""));

                                // 构建输入字符串
                                let input_str = match filepath {
                                    Some(p) => format!("/save {}", p),
                                    None => "/save".to_string(),
                                };

                                // 解析为 Command
                                match adapter.try_parse(&input_str, ctx) {
                                    Ok(Some(command)) => {
                                        // 序列化为 JSON
                                        let json = serde_json::to_string(&command).unwrap_or_default();
                                        // 通过 Command 消息发送
                                        let inbound = WsInbound::Command { payload: json };
                                        if let Ok(text) = serialize_inbound(&inbound) {
                                            let _ = sender.send(Message::Text(text.into())).await;
                                        }
                                    }
                                    Ok(None) => {
                                        tracing::warn!("Failed to parse /save command");
                                    }
                                    Err(e) => {
                                        tracing::error!(error = %e, "Error parsing /save command");
                                    }
                                }
                                continue;
                            }
                            InputEvent::Command(InputCommand::Sessions) => {
                                // 使用 CliInputAdapter 构建 Command
                                let adapter = CliInputAdapter::new();
                                let ctx = AdapterContext::new("cli")
                                    .with_session_id(current_session_id.as_deref().unwrap_or(""));

                                // 解析为 Command
                                match adapter.try_parse("/list", ctx) {
                                    Ok(Some(command)) => {
                                        // 序列化为 JSON
                                        let json = serde_json::to_string(&command).unwrap_or_default();
                                        // 通过 Command 消息发送
                                        let inbound = WsInbound::Command { payload: json };
                                        if let Ok(text) = serialize_inbound(&inbound) {
                                            let _ = sender.send(Message::Text(text.into())).await;
                                        }
                                    }
                                    Ok(None) => {
                                        tracing::warn!("Failed to parse /list command");
                                    }
                                    Err(e) => {
                                        tracing::error!(error = %e, "Error parsing /list command");
                                    }
                                }
                                continue;
                            }
                            InputEvent::Command(InputCommand::Use(session_id)) => {
                                // 使用 CliInputAdapter 构建 Command
                                let adapter = CliInputAdapter::new();
                                let ctx = AdapterContext::new("cli")
                                    .with_session_id(current_session_id.as_deref().unwrap_or(""));

                                // 构建输入字符串
                                let input_str = format!("/use {}", session_id);

                                // 解析为 Command
                                match adapter.try_parse(&input_str, ctx) {
                                    Ok(Some(command)) => {
                                        // 序列化为 JSON
                                        let json = serde_json::to_string(&command).unwrap_or_default();
                                        // 通过 Command 消息发送
                                        let inbound = WsInbound::Command { payload: json };
                                        if let Ok(text) = serialize_inbound(&inbound) {
                                            let _ = sender.send(Message::Text(text.into())).await;
                                        }
                                        // 更新当前会话 ID
                                        current_session_id = Some(session_id.clone());
                                    }
                                    Ok(None) => {
                                        tracing::warn!("Failed to parse /use command");
                                    }
                                    Err(e) => {
                                        tracing::error!(error = %e, "Error parsing /use command");
                                    }
                                }
                                continue;
                            }
                            InputEvent::Message(msg) => {
                                let inbound = WsInbound::Message {
                                    content: msg.content,
                                    attachments: Vec::new(),
                                    channel: None,
                                    chat_id: current_session_id.clone(),
                                    sender_id: None,
                                };
                                if let Ok(text) = serialize_inbound(&inbound)
                                    && sender.send(Message::Text(text.into())).await.is_err() {
                                        tracing::error!("Failed to send message to gateway");
                                        break;
                                    }
                            }
                        }
                    }
                    Ok(None) => break,
                    Err(e) => {
                        tracing::error!(error = %e, "Input error");
                        break;
                    }
                }
            }
        }
    }

    Ok(())
}