Skip to main content

cascade_agent/orchestrator/
mod.rs

1pub mod server;
2pub mod types;
3pub mod websocket;
4
5use crate::config::OrchestratorSettings;
6use crate::error::Result;
7use async_trait::async_trait;
8use tokio::sync::broadcast;
9use types::OrchestratorMessage;
10
11/// Trait for bidirectional orchestrator communication.
12#[async_trait]
13pub trait OrchestratorConnection: Send + Sync {
14    /// Push a message from agent to orchestrator (fire-and-forget)
15    async fn push(&self, message: OrchestratorMessage);
16
17    /// Receive the next message from orchestrator (blocking await)
18    async fn recv(&mut self) -> Option<OrchestratorMessage>;
19
20    /// Check if connected
21    fn is_connected(&self) -> bool;
22}
23
24/// No-op implementation for when orchestrator is disabled.
25pub struct NoopOrchestrator;
26
27#[async_trait]
28impl OrchestratorConnection for NoopOrchestrator {
29    async fn push(&self, message: OrchestratorMessage) {
30        tracing::debug!(target: "orchestrator", "Noop push: {:?}", message);
31    }
32    async fn recv(&mut self) -> Option<OrchestratorMessage> {
33        std::future::pending().await
34    }
35    fn is_connected(&self) -> bool {
36        false
37    }
38}
39
40/// Server-mode orchestrator: the agent hosts a WebSocket server and communicates
41/// via broadcast (outbound) and mpsc (inbound) channels.
42pub struct ServerOrchestrator {
43    broadcast_tx: broadcast::Sender<OrchestratorMessage>,
44    inbound_rx: tokio::sync::mpsc::Receiver<OrchestratorMessage>,
45    #[allow(dead_code)]
46    server_handle: Option<tokio::task::JoinHandle<crate::error::Result<()>>>,
47}
48
49impl std::fmt::Debug for ServerOrchestrator {
50    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
51        f.debug_struct("ServerOrchestrator").finish()
52    }
53}
54
55impl ServerOrchestrator {
56    pub async fn new(bind_address: &str) -> Result<Self> {
57        let (server, broadcast_tx, inbound_rx) = server::OrchestratorServer::new(bind_address);
58        let handle = tokio::spawn(server.run());
59        Ok(Self {
60            broadcast_tx,
61            inbound_rx,
62            server_handle: Some(handle),
63        })
64    }
65}
66
67#[async_trait]
68impl OrchestratorConnection for ServerOrchestrator {
69    async fn push(&self, message: OrchestratorMessage) {
70        let _ = self.broadcast_tx.send(message);
71    }
72
73    async fn recv(&mut self) -> Option<OrchestratorMessage> {
74        self.inbound_rx.recv().await
75    }
76
77    fn is_connected(&self) -> bool {
78        self.broadcast_tx.receiver_count() > 0
79    }
80}
81
82/// Factory function to create the right orchestrator transport.
83pub fn create_orchestrator(
84    config: &OrchestratorSettings,
85) -> Result<Box<dyn OrchestratorConnection>> {
86    if !config.enabled {
87        return Ok(Box::new(NoopOrchestrator));
88    }
89    match config.transport.as_str() {
90        "websocket" => Ok(Box::new(NoopOrchestrator)),
91        other => Err(crate::error::AgentError::OrchestratorError(format!(
92            "Unknown transport: {}",
93            other
94        ))),
95    }
96}
97
98/// Async factory for transports that need initialization (like WebSocket).
99pub async fn create_orchestrator_async(
100    config: &OrchestratorSettings,
101) -> Result<Box<dyn OrchestratorConnection>> {
102    if !config.enabled {
103        return Ok(Box::new(NoopOrchestrator));
104    }
105    match config.transport.as_str() {
106        "websocket" => {
107            if let Some(url) = &config.connect_url {
108                let ws = websocket::WebSocketOrchestrator::connect(url).await?;
109                Ok(Box::new(ws))
110            } else {
111                let server = ServerOrchestrator::new(&config.bind_address).await?;
112                Ok(Box::new(server))
113            }
114        }
115        other => Err(crate::error::AgentError::OrchestratorError(format!(
116            "Unknown transport: {}",
117            other
118        ))),
119    }
120}