cascade_agent/orchestrator/
server.rs

1use super::types::OrchestratorMessage;
2use futures_util::{SinkExt, StreamExt};
3use tokio::net::TcpListener;
4use tokio::sync::{broadcast, mpsc};
5use tokio_tungstenite::accept_async;
6
7/// A simple WebSocket server for the orchestrator to connect to.
8/// The agent runs this server; a future orchestrator UI connects as a client.
9pub struct OrchestratorServer {
10    bind_address: String,
11    // broadcast channel: server sends to all connected orchestrator clients
12    broadcast_tx: broadcast::Sender<OrchestratorMessage>,
13    // channel for server to receive messages from orchestrator clients
14    inbound_tx: mpsc::Sender<OrchestratorMessage>,
15}
16
17impl OrchestratorServer {
18    pub fn new(
19        bind_address: &str,
20    ) -> (
21        Self,
22        broadcast::Sender<OrchestratorMessage>,
23        mpsc::Receiver<OrchestratorMessage>,
24    ) {
25        let (broadcast_tx, _) = broadcast::channel(256);
26        let (inbound_tx, inbound_rx) = mpsc::channel(256);
27        let server = Self {
28            bind_address: bind_address.to_string(),
29            broadcast_tx: broadcast_tx.clone(),
30            inbound_tx,
31        };
32        (server, broadcast_tx, inbound_rx)
33    }
34
35    pub async fn run(self) -> crate::error::Result<()> {
36        let listener = TcpListener::bind(&self.bind_address).await.map_err(|e| {
37            crate::error::AgentError::OrchestratorError(format!("Failed to bind: {}", e))
38        })?;
39
40        tracing::info!(target: "orchestrator", "Server listening on {}", self.bind_address);
41
42        while let Ok((stream, addr)) = listener.accept().await {
43            let broadcast_tx = self.broadcast_tx.clone();
44            let inbound_tx = self.inbound_tx.clone();
45
46            tokio::spawn(async move {
47                let ws_stream = match accept_async(stream).await {
48                    Ok(s) => s,
49                    Err(e) => {
50                        tracing::error!(target: "orchestrator", "Accept failed from {}: {}", addr, e);
51                        return;
52                    }
53                };
54
55                tracing::info!(target: "orchestrator", "Orchestrator connected from {}", addr);
56                let (mut ws_sink, mut ws_stream) = ws_stream.split();
57
58                // Subscribe to broadcast
59                let mut rx = broadcast_tx.subscribe();
60
61                // Forward broadcast → WebSocket
62                let ws_to_client = async {
63                    while let Ok(msg) = rx.recv().await {
64                        let json = serde_json::to_string(&msg).unwrap_or_default();
65                        if ws_sink
66                            .send(tokio_tungstenite::tungstenite::protocol::Message::Text(
67                                json.into(),
68                            ))
69                            .await
70                            .is_err()
71                        {
72                            break;
73                        }
74                    }
75                };
76
77                // Forward WebSocket → inbound
78                let client_to_server = async {
79                    while let Some(Ok(msg)) = ws_stream.next().await {
80                        if msg.is_text() {
81                            let text = msg.to_text().unwrap_or_default();
82                            if let Ok(parsed) = serde_json::from_str::<OrchestratorMessage>(text) {
83                                if inbound_tx.send(parsed).await.is_err() {
84                                    break;
85                                }
86                            }
87                        }
88                    }
89                };
90
91                tokio::select! {
92                    _ = ws_to_client => {},
93                    _ = client_to_server => {},
94                }
95
96                tracing::info!(target: "orchestrator", "Orchestrator disconnected from {}", addr);
97            });
98        }
99
100        Ok(())
101    }
102}
cascade_agent/orchestrator/server.rs

cascade_agent/orchestrator/
server.rs
