cascade_agent/orchestrator/
mod.rs1pub 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#[async_trait]
13pub trait OrchestratorConnection: Send + Sync {
14 async fn push(&self, message: OrchestratorMessage);
16
17 async fn recv(&mut self) -> Option<OrchestratorMessage>;
19
20 fn is_connected(&self) -> bool;
22}
23
24pub 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
40pub 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
82pub 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
98pub 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}