cascade_agent/orchestrator/
websocket.rs

1use super::{types::OrchestratorMessage, OrchestratorConnection};
2use async_trait::async_trait;
3use futures_util::{SinkExt, StreamExt};
4use tokio::sync::mpsc;
5use tokio_tungstenite::{connect_async, tungstenite::protocol::Message};
6
7pub struct WebSocketOrchestrator {
8    outbound_tx: mpsc::Sender<OrchestratorMessage>,
9    inbound_rx: mpsc::Receiver<OrchestratorMessage>,
10}
11
12impl WebSocketOrchestrator {
13    pub async fn connect(url: &str) -> crate::error::Result<Self> {
14        let (ws_stream, _response) = connect_async(url).await.map_err(|e| {
15            crate::error::AgentError::OrchestratorError(format!("WebSocket connect failed: {}", e))
16        })?;
17
18        let (mut ws_sink, mut ws_stream) = ws_stream.split();
19
20        let (outbound_tx, mut outbound_rx) = mpsc::channel::<OrchestratorMessage>(256);
21        let (inbound_tx, inbound_rx) = mpsc::channel::<OrchestratorMessage>(256);
22
23        // Task 1: Forward outbound messages to WebSocket
24        tokio::spawn(async move {
25            while let Some(msg) = outbound_rx.recv().await {
26                let json = serde_json::to_string(&msg).unwrap_or_default();
27                if ws_sink.send(Message::Text(json.into())).await.is_err() {
28                    tracing::error!(target: "orchestrator", "Failed to send message to WebSocket");
29                    break;
30                }
31            }
32        });
33
34        // Task 2: Forward WebSocket messages to inbound channel
35        tokio::spawn(async move {
36            while let Some(Ok(msg)) = ws_stream.next().await {
37                if msg.is_text() {
38                    let text = msg.to_text().unwrap_or_default();
39                    if let Ok(parsed) = serde_json::from_str::<OrchestratorMessage>(text) {
40                        if inbound_tx.send(parsed).await.is_err() {
41                            break;
42                        }
43                    }
44                }
45            }
46        });
47
48        Ok(Self {
49            outbound_tx,
50            inbound_rx,
51        })
52    }
53}
54
55#[async_trait]
56impl OrchestratorConnection for WebSocketOrchestrator {
57    async fn push(&self, message: OrchestratorMessage) {
58        let _ = self.outbound_tx.send(message).await;
59    }
60
61    async fn recv(&mut self) -> Option<OrchestratorMessage> {
62        self.inbound_rx.recv().await
63    }
64
65    fn is_connected(&self) -> bool {
66        !self.outbound_tx.is_closed()
67    }
68}
cascade_agent/orchestrator/websocket.rs

cascade_agent/orchestrator/
websocket.rs
