zesdex/apps/interfaces/ws/src/lib.rs

//! WebSocket interface — real-time bidirectional communication.
//!
//! Enables web clients and other WS-capable consumers to connect
//! and participate in sessions. Built on Axum's WebSocket support.

use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::response::IntoResponse;
use axum::routing::get;
use axum::Router;
use futures_util::stream::StreamExt;
use futures_util::SinkExt;
use std::sync::Arc;
use tracing::info;

/// Shared application state for the WS server.
pub struct WsState {
    pub store_base_dir: std::path::PathBuf,
    pub session_id: Option<String>,
}

/// Build the WebSocket router.
pub fn build_router(state: Arc<WsState>) -> Router {
    Router::new()
        .route("/ws", get(ws_handler))
        .with_state(state)
}

/// WebSocket upgrade handler.
async fn ws_handler(
    ws: WebSocketUpgrade,
    axum::extract::State(state): axum::extract::State<Arc<WsState>>,
) -> impl IntoResponse {
    ws.on_upgrade(move |socket| handle_socket(socket, state))
}

/// Handle an established WebSocket connection.
async fn handle_socket(mut socket: WebSocket, state: Arc<WsState>) {
    // Channel for sending text messages to the WebSocket send task.
    // The receiver side runs in a spawned task that forwards each
    // string as a `Message::Text` to the client.
    let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<String>();

    info!("WebSocket client connected");

    // Send a welcome message
    let welcome = serde_json::json!({
        "type": "connected",
        "session": state.session_id,
        "message": "Connected to Zesdex WebSocket server"
    });
    // axum 0.8 Message::Text wraps Utf8Bytes; convert via .into()
    let _ = socket.send(Message::Text(welcome.to_string().into())).await;

    // Split the socket into sender and receiver halves
    let (mut sender, mut receiver) = socket.split();

    // Spawn task to forward messages from channel to WebSocket sender
    let send_task = tokio::spawn(async move {
        while let Some(msg) = rx.recv().await {
            if sender.send(Message::Text(msg.into())).await.is_err() {
                break;
            }
        }
    });

    // Receive messages from the client
    // receiver is SplitStream<WebSocket> — use StreamExt::next()
    while let Some(Ok(msg)) = receiver.next().await {
        match msg {
            Message::Text(text) => {
                let text_str = text.to_string();
                info!("Received WS message: {text_str}");
                
                if let Ok(val) = serde_json::from_str::<serde_json::Value>(&text_str) {
                    if val.get("type").and_then(|v| v.as_str()) == Some("prompt") {
                        if let Some(prompt) = val.get("message").and_then(|v| v.as_str()) {
                            let session_dir = std::env::current_dir().unwrap_or_default();
                            let workspace_roots = vec![session_dir.clone()];
                            let turn_events = Arc::new(std::sync::Mutex::new(std::collections::VecDeque::new()));
                            let in_flight = Arc::new(std::sync::atomic::AtomicBool::new(false));
                            let abort = Arc::new(std::sync::atomic::AtomicBool::new(false));
                            let api_key = std::env::var("OPENAI_API_KEY").unwrap_or_default();
                            let model = val.get("model").and_then(|v| v.as_str()).unwrap_or("gpt-4o").to_string();

                            let params = zesdex_domain::agent::AgentTurnParams {
                                messages: vec![zesdex_domain::core::ChatMessage::user(prompt)],
                                session_dir: session_dir.clone(),
                                workspace_roots: workspace_roots.clone(),
                                turn_events: turn_events.clone(),
                                in_flight,
                                abort,
                                api_key: api_key.clone(),
                                model: model.clone(),
                                api_base: None,
                            };

                            let client = std::sync::Arc::new(zesdex_infrastructure::llm::provider::LlmClient::new(
                                api_key,
                                model,
                                None,
                            ));

                            let tool_ctx = zesdex_infrastructure::tools::ToolCtx::builder()
                                .session_dir(session_dir)
                                .workspaces(workspace_roots)
                                .turn_events(turn_events.clone())
                                .build();

                            let tool_executor = std::sync::Arc::new(
                                zesdex_infrastructure::tools::executor::InfrastructureToolExecutor::new(tool_ctx),
                            );

                            let tools = zesdex_infrastructure::tools::all_tools();
                            let defs = zesdex_infrastructure::tools::tool_defs(&tools);

                            let turn_service = zesdex_application::agent::turn_service::AgentTurnServiceImpl::new(
                                client,
                                tool_executor,
                                defs,
                            );

                            use zesdex_application::agent::AgentTurnService;
                            tokio::spawn(async move {
                                let _ = turn_service.run_turn(params).await;
                            });

                            let tx_clone = tx.clone();
                            tokio::spawn(async move {
                                let mut done = false;
                                while !done {
                                    tokio::time::sleep(std::time::Duration::from_millis(50)).await;
                                    let events: Vec<_> = {
                                        if let Ok(mut q) = turn_events.lock() {
                                            q.drain(..).collect()
                                        } else {
                                            vec![]
                                        }
                                    };
                                    for ev in events {
                                        match ev {
                                            zesdex_infrastructure::TurnEvent::StreamToken(tok) => {
                                                let json = serde_json::json!({ "type": "token", "content": tok });
                                                let _ = tx_clone.send(json.to_string());
                                            }
                                            zesdex_infrastructure::TurnEvent::Done => {
                                                let json = serde_json::json!({ "type": "done" });
                                                let _ = tx_clone.send(json.to_string());
                                                done = true;
                                            }
                                            zesdex_infrastructure::TurnEvent::Error(err) => {
                                                let json = serde_json::json!({ "type": "error", "message": err });
                                                let _ = tx_clone.send(json.to_string());
                                                done = true;
                                            }
                                            _ => {}
                                        }
                                    }
                                }
                            });
                            continue;
                        }
                    }
                }

                // Fallback echo
                let response = serde_json::json!({
                    "type": "echo",
                    "data": text_str
                });
                let _ = tx.send(response.to_string());
            }
            Message::Close(_) => break,
            _ => {}
        }
    }

    send_task.abort();
    info!("WebSocket client disconnected");
}

/// Run the WebSocket server standalone.
pub async fn run_server(port: u16) -> anyhow::Result<()> {
    let state = Arc::new(WsState {
        store_base_dir: std::path::PathBuf::from("."),
        session_id: None,
    });
    let app = build_router(state);
    let addr = std::net::SocketAddr::from(([0, 0, 0, 0], port));
    info!("WebSocket server listening on ws://{addr}");
    let listener = tokio::net::TcpListener::bind(addr).await?;
    axum::serve(listener, app).await?;
    Ok(())
}