agentmux_srv\muxbus/
cloud_subscriber.rs

1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Cloud push subscription — replaces per-agent polling with a single
5//! sidecar-level WebSocket to muxbus.agentmux.ai.
6//!
7//! When MUXBUS_TOKEN is present, the subscriber:
8//!   1. Opens wss://muxbus.agentmux.ai/ws with the bearer token.
9//!   2. Listens for { type: "inject_available" } broadcast wake signals (zero metadata).
10//!   3. On each signal, polls REST GET /reactive/pending/:id for every locally-registered agent.
11//!   4. Delivers via ReactiveHandler; ACKs successful deliveries via REST POST /reactive/ack.
12//!   5. Reconnects with exponential back-off on any disconnect.
13//!
14//! The server broadcasts to ALL connected sidecars on every injection, so a subscriber
15//! cannot correlate the signal to any particular agent, account, or timing.
16//!
17//! Agents register/unregister via add_agent()/remove_agent().
18//! No MUXBUS_TOKEN → subscriber stays idle.
19
20use std::collections::HashSet;
21use std::sync::{Arc, Mutex, OnceLock};
22use std::time::Duration;
23
24use futures_util::{SinkExt, StreamExt};
25use serde::{Deserialize, Serialize};
26use tokio::sync::mpsc;
27use tokio_tungstenite::{connect_async_tls_with_config, tungstenite::Message};
28
29use crate::backend::reactive::handler::get_global_handler;
30use crate::backend::reactive::types::InjectionRequest;
31use crate::backend::storage::store::Store;
32
33const MUXBUS_WS_URL: &str = "wss://muxbus.agentmux.ai/ws";
34const MUXBUS_REST_URL: &str = "https://muxbus.agentmux.ai";
35const RECONNECT_DELAY_SECS: u64 = 5;
36const MAX_RECONNECT_DELAY_SECS: u64 = 60;
37
38// ── Wire protocol ─────────────────────────────────────────────────────────────
39
40#[derive(Serialize)]
41#[serde(tag = "type", rename_all = "snake_case")]
42enum ClientMsg {
43    Subscribe { agents: Vec<String> },
44    #[serde(rename = "subscribe:add")]
45    SubscribeAdd { agents: Vec<String> },
46    #[serde(rename = "subscribe:remove")]
47    SubscribeRemove { agents: Vec<String> },
48    // Ack removed — ACK is sent via REST POST /reactive/ack, not WS
49}
50
51#[derive(Deserialize)]
52#[serde(tag = "type", rename_all = "snake_case")]
53enum ServerMsg {
54    // Zero-metadata wake signal. The sidecar polls all its registered agents on receipt.
55    // No agent_id, no injection_id — subscribers gain no information from this frame.
56    InjectAvailable,
57    Subscribed {
58        #[allow(dead_code)]
59        agents: Vec<String>,
60    },
61    Pong,
62    Error {
63        message: String,
64    },
65    Evicted {
66        #[allow(dead_code)]
67        agents: Vec<String>,
68    },
69    #[serde(other)]
70    Unknown,
71}
72
73// ── Control channel ───────────────────────────────────────────────────────────
74
75enum CtrlMsg {
76    AddAgent(String),
77    RemoveAgent(String),
78    ReloadToken,
79}
80
81// ── Public struct ─────────────────────────────────────────────────────────────
82
83pub struct CloudSubscriber {
84    ctrl_tx: mpsc::UnboundedSender<CtrlMsg>,
85    /// In-memory copy of subscribed agents, kept in sync with the WS loop
86    /// so add_agent/remove_agent can be called before the WS connects.
87    agents: Arc<Mutex<HashSet<String>>>,
88}
89
90static GLOBAL_SUBSCRIBER: OnceLock<CloudSubscriber> = OnceLock::new();
91
92pub fn get_global_subscriber() -> Option<&'static CloudSubscriber> {
93    GLOBAL_SUBSCRIBER.get()
94}
95
96impl CloudSubscriber {
97    /// Initialize the global subscriber and start the background WS loop.
98    /// Should be called once at startup. No-op if already initialized.
99    pub fn init_global(wstore: Arc<Store>) {
100        let (ctrl_tx, ctrl_rx) = mpsc::unbounded_channel::<CtrlMsg>();
101        let agents = Arc::new(Mutex::new(HashSet::<String>::new()));
102        let subscriber = CloudSubscriber {
103            ctrl_tx,
104            agents: agents.clone(),
105        };
106        if GLOBAL_SUBSCRIBER.set(subscriber).is_err() {
107            return; // already initialized
108        }
109        tokio::spawn(run_loop(wstore, agents, ctrl_rx));
110    }
111
112    /// Notify the WS loop that a new agent is registered locally.
113    /// Normalizes to lowercase and skips the WS send if already subscribed,
114    /// so calling this on every COMMAND_AGENT_INPUT turn is safe.
115    pub fn add_agent(&self, agent_id: &str) {
116        let key = agent_id.to_lowercase();
117        let mut agents = self.agents.lock().unwrap();
118        if agents.contains(&key) {
119            return;
120        }
121        agents.insert(key.clone());
122        drop(agents);
123        let _ = self.ctrl_tx.send(CtrlMsg::AddAgent(key));
124    }
125
126    /// Notify the WS loop that an agent has been unregistered.
127    pub fn remove_agent(&self, agent_id: &str) {
128        let key = agent_id.to_lowercase();
129        self.agents.lock().unwrap().remove(&key);
130        let _ = self.ctrl_tx.send(CtrlMsg::RemoveAgent(key));
131    }
132
133    /// Called after muxbus.login completes — trigger a fresh WS connection
134    /// with the newly stored token.
135    pub fn reload_token(&self) {
136        let _ = self.ctrl_tx.send(CtrlMsg::ReloadToken);
137    }
138
139    /// Snapshot of the agents this sidecar has subscribed to the cloud relay
140    /// (Tier-4), sorted for stable output. Empty when no MUXBUS_TOKEN is set
141    /// (cloud disabled). Read-only; used by the discovery endpoint.
142    pub fn subscribed_agents(&self) -> Vec<String> {
143        let mut v: Vec<String> = self.agents.lock().unwrap().iter().cloned().collect();
144        v.sort();
145        v
146    }
147}
148
149// ── Background loop ───────────────────────────────────────────────────────────
150
151async fn run_loop(
152    wstore: Arc<Store>,
153    agents: Arc<Mutex<HashSet<String>>>,
154    mut ctrl_rx: mpsc::UnboundedReceiver<CtrlMsg>,
155) {
156    let http = reqwest::Client::new();
157    let mut delay_secs = RECONNECT_DELAY_SECS;
158
159    loop {
160        // Load token — refresh if expired.
161        // Two distinct None cases:
162        //   a) No credentials in DB at all → wait for muxbus.login ReloadToken signal
163        //   b) Credentials exist but refresh failed transiently → back off and retry
164        // has_stored_creds = true only when a retry is meaningful:
165        //   - access_token present AND still valid (use it now), OR
166        //   - access_token expired but refresh_token present (can refresh).
167        // An expired token with no refresh_token is a permanent failure → park like no-creds.
168        let has_stored_creds = wstore
169            .muxbus_load()
170            .ok()
171            .flatten()
172            .map(|c| !c.access_token.is_empty() && (c.is_valid() || !c.refresh_token.is_empty()))
173            .unwrap_or(false);
174        let token = match load_valid_token(&wstore, &http).await {
175            Some(t) => t,
176            None if !has_stored_creds => {
177                // No credentials at all — wait for muxbus.login to signal us
178                while let Some(msg) = ctrl_rx.recv().await {
179                    if matches!(msg, CtrlMsg::ReloadToken) { break; }
180                    // AddAgent / RemoveAgent already updated `agents` mutex
181                }
182                delay_secs = RECONNECT_DELAY_SECS; // reset back-off after fresh login
183                continue;
184            }
185            None => {
186                // Credentials present but refresh failed transiently — back off and retry.
187                // Loop inside the select! so AddAgent/RemoveAgent messages drain without
188                // cancelling the sleep; only ReloadToken (or sleep expiry) breaks out.
189                tracing::warn!(
190                    "cloud_subscriber: token refresh failed, retrying in {}s",
191                    delay_secs
192                );
193                let sleep = tokio::time::sleep(Duration::from_secs(delay_secs));
194                tokio::pin!(sleep);
195                'backoff: loop {
196                    tokio::select! {
197                        _ = &mut sleep => break 'backoff,
198                        Some(msg) = ctrl_rx.recv() => match msg {
199                            CtrlMsg::ReloadToken => {
200                                delay_secs = RECONNECT_DELAY_SECS;
201                                break 'backoff;
202                            }
203                            CtrlMsg::AddAgent(_) | CtrlMsg::RemoveAgent(_) => {
204                                // agents mutex already updated; continue sleeping
205                            }
206                        }
207                    }
208                }
209                delay_secs = (delay_secs * 2).min(MAX_RECONNECT_DELAY_SECS);
210                continue;
211            }
212        };
213
214        tracing::info!("cloud_subscriber: connecting to {}", MUXBUS_WS_URL);
215        let session_start = std::time::Instant::now();
216
217        match connect_and_run(&token, agents.clone(), &mut ctrl_rx, &wstore, &http).await {
218            Ok(()) => {
219                // Apply the same >30s healthy-session guard as the error branch: a server
220                // that immediately closes after handshake must not suppress back-off.
221                if session_start.elapsed().as_secs() > 30 {
222                    delay_secs = RECONNECT_DELAY_SECS;
223                }
224                tracing::info!("cloud_subscriber: disconnected cleanly");
225            }
226            Err(e) => {
227                // If the session was healthy for >30s before the error, treat it as
228                // network instability rather than a persistent failure and reset back-off.
229                if session_start.elapsed().as_secs() > 30 {
230                    delay_secs = RECONNECT_DELAY_SECS;
231                }
232                tracing::warn!("cloud_subscriber: error: {e}, reconnecting in {}s", delay_secs);
233            }
234        }
235
236        // Post-disconnect back-off. Loop inside the select! so AddAgent/RemoveAgent
237        // messages drain without cancelling the sleep; only ReloadToken breaks out early.
238        let sleep = tokio::time::sleep(Duration::from_secs(delay_secs));
239        tokio::pin!(sleep);
240        'reconnect: loop {
241            tokio::select! {
242                _ = &mut sleep => break 'reconnect,
243                Some(msg) = ctrl_rx.recv() => match msg {
244                    CtrlMsg::ReloadToken => {
245                        delay_secs = RECONNECT_DELAY_SECS;
246                        break 'reconnect;
247                    }
248                    CtrlMsg::AddAgent(_) | CtrlMsg::RemoveAgent(_) => {
249                        // agents mutex already updated; continue sleeping
250                    }
251                }
252            }
253        }
254        delay_secs = (delay_secs * 2).min(MAX_RECONNECT_DELAY_SECS);
255    }
256}
257
258async fn connect_and_run(
259    token: &str,
260    agents: Arc<Mutex<HashSet<String>>>,
261    ctrl_rx: &mut mpsc::UnboundedReceiver<CtrlMsg>,
262    _wstore: &Arc<Store>,
263    http: &reqwest::Client,
264) -> Result<(), String> {
265    use tokio_tungstenite::tungstenite::http::Request;
266
267    // Only set Authorization — tungstenite adds Connection/Upgrade/Sec-WebSocket-*
268    // automatically. Setting them manually causes duplicate headers in the handshake.
269    let request = Request::builder()
270        .uri(MUXBUS_WS_URL)
271        .header("Authorization", format!("Bearer {}", token))
272        .body(())
273        .map_err(|e| format!("build request: {e}"))?;
274
275    let (ws_stream, _) = connect_async_tls_with_config(request, None, false, None)
276        .await
277        .map_err(|e| format!("connect: {e}"))?;
278
279    tracing::info!("cloud_subscriber: WebSocket connected");
280    let (mut write, mut read) = ws_stream.split();
281
282    // Initial subscription for all currently-registered agents
283    let initial_agents: Vec<String> = agents.lock().unwrap().iter().cloned().collect();
284    let sub_msg = serde_json::to_string(&ClientMsg::Subscribe { agents: initial_agents })
285        .map_err(|e| format!("serialize subscribe: {e}"))?;
286    write.send(Message::Text(sub_msg.into()))
287        .await
288        .map_err(|e| format!("send subscribe: {e}"))?;
289
290    loop {
291        tokio::select! {
292            // Incoming WebSocket message from cloud
293            msg = read.next() => {
294                match msg {
295                    None => return Ok(()), // server closed
296                    Some(Err(e)) => return Err(format!("ws recv: {e}")),
297                    Some(Ok(Message::Text(text))) => {
298                        if let Ok(server_msg) = serde_json::from_str::<ServerMsg>(&text) {
299                            match handle_server_msg(server_msg, token, http, &agents).await {
300                                Ok(()) => {}
301                                // Eviction or expired token — close stream to trigger reconnect.
302                                Err(ref e) if e.starts_with("reconnect:") => {
303                                    tracing::info!("cloud_subscriber: {e}, reconnecting");
304                                    return Ok(());
305                                }
306                                Err(e) => {
307                                    tracing::warn!("cloud_subscriber: handle msg error: {e}");
308                                }
309                            }
310                        }
311                    }
312                    Some(Ok(Message::Ping(data))) => {
313                        let _ = write.send(Message::Pong(data)).await;
314                    }
315                    Some(Ok(Message::Close(_))) => return Ok(()),
316                    Some(Ok(_)) => {} // binary frames etc — ignore
317                }
318            }
319
320            // Control messages from register/unregister calls
321            ctrl = ctrl_rx.recv() => {
322                match ctrl {
323                    None => return Ok(()), // channel closed
324                    Some(CtrlMsg::AddAgent(id)) => {
325                        let msg = serde_json::to_string(&ClientMsg::SubscribeAdd { agents: vec![id] })
326                            .unwrap_or_default();
327                        let _ = write.send(Message::Text(msg.into())).await;
328                    }
329                    Some(CtrlMsg::RemoveAgent(id)) => {
330                        let msg = serde_json::to_string(&ClientMsg::SubscribeRemove { agents: vec![id] })
331                            .unwrap_or_default();
332                        let _ = write.send(Message::Text(msg.into())).await;
333                    }
334                    Some(CtrlMsg::ReloadToken) => {
335                        // Close and let the outer loop reconnect with the new token
336                        let _ = write.send(Message::Close(None)).await;
337                        return Ok(());
338                    }
339                }
340            }
341        }
342    }
343}
344
345/// Handle a message from the cloud server.
346/// The WS wake signal carries zero metadata — the sidecar polls all its registered
347/// agents via REST to find pending injections, so a compromised subscriber gains nothing.
348async fn handle_server_msg(
349    msg: ServerMsg,
350    token: &str,
351    http: &reqwest::Client,
352    agents: &Arc<Mutex<HashSet<String>>>,
353) -> Result<(), String> {
354    match msg {
355        ServerMsg::InjectAvailable => {
356            // Collect currently-registered agents without holding the lock across awaits.
357            let registered: Vec<String> = agents.lock().unwrap().iter().cloned().collect();
358
359            #[derive(Deserialize)]
360            struct PendingResp { injections: Vec<PendingInj> }
361            #[derive(Deserialize)]
362            struct PendingInj {
363                id: String,
364                source_agent: Option<String>,
365                message: String,
366                priority: Option<String>,
367            }
368
369            let handler = get_global_handler();
370            for agent_id in &registered {
371                let url = format!("{}/reactive/pending/{}", MUXBUS_REST_URL, agent_id);
372                let resp = match http
373                    .get(&url)
374                    .header("Authorization", format!("Bearer {}", token))
375                    .header("X-Agent-ID", agent_id)
376                    .send()
377                    .await
378                {
379                    Ok(r) => r,
380                    Err(e) => {
381                        tracing::warn!(agent_id = %agent_id, error = %e, "cloud_subscriber: fetch pending failed");
382                        continue;
383                    }
384                };
385
386                if resp.status() == reqwest::StatusCode::UNAUTHORIZED {
387                    // Token expired during session — reconnect to refresh.
388                    return Err("reconnect:token_expired".to_string());
389                }
390                if !resp.status().is_success() {
391                    tracing::warn!(
392                        status = %resp.status(),
393                        agent_id = %agent_id,
394                        "cloud_subscriber: fetch pending non-2xx"
395                    );
396                    continue;
397                }
398
399                let body: PendingResp = match resp.json().await {
400                    Ok(b) => b,
401                    Err(e) => {
402                        tracing::warn!(error = %e, "cloud_subscriber: parse pending failed");
403                        continue;
404                    }
405                };
406
407                let mut ack_ids: Vec<String> = Vec::new();
408                for inj in &body.injections {
409                    let req = InjectionRequest {
410                        target_agent: agent_id.clone(),
411                        message: inj.message.clone(),
412                        source_agent: inj.source_agent.clone(),
413                        request_id: Some(inj.id.clone()),
414                        priority: inj.priority.clone(),
415                        wait_for_idle: false,
416                    };
417                    let delivery = handler.inject_message(req);
418                    tracing::debug!(
419                        injection_id = %inj.id,
420                        agent_id = %agent_id,
421                        success = delivery.success,
422                        "cloud_subscriber: delivered injection"
423                    );
424                    // Only ACK successfully delivered injections — failed deliveries stay
425                    // in the queue so the cloud can retry (e.g. rate-limited or agent not ready).
426                    if delivery.success {
427                        ack_ids.push(inj.id.clone());
428                    }
429                }
430
431                if !ack_ids.is_empty() {
432                    let ack_url = format!("{}/reactive/ack", MUXBUS_REST_URL);
433                    let _ = http
434                        .post(&ack_url)
435                        .header("Authorization", format!("Bearer {}", token))
436                        .header("X-Agent-ID", agent_id)
437                        .json(&serde_json::json!({ "injection_ids": ack_ids }))
438                        .send()
439                        .await;
440                }
441            }
442        }
443        ServerMsg::Error { message } => {
444            tracing::warn!("cloud_subscriber: server error: {message}");
445        }
446        ServerMsg::Evicted { .. } => {
447            // Another sidecar evicted us. Reconnect to re-subscribe (reclaim).
448            return Err("reconnect:evicted".to_string());
449        }
450        ServerMsg::Pong | ServerMsg::Subscribed { .. } | ServerMsg::Unknown => {}
451    }
452    Ok(())
453}
454
455/// Load a valid (non-expired) access token, refreshing via refresh_token if needed.
456/// Saves refreshed credentials back to the store. Returns None if no credentials
457/// are stored, the token is expired and refresh fails, or the refresh_token is absent.
458async fn load_valid_token(wstore: &Store, http: &reqwest::Client) -> Option<String> {
459    let creds = match wstore.muxbus_load() {
460        Ok(Some(c)) if !c.access_token.is_empty() => c,
461        _ => return None,
462    };
463    if creds.is_valid() && !creds.nearly_expired() {
464        return Some(creds.access_token);
465    }
466    // Token expired or expiring within 300s — proactively refresh
467    tracing::info!("cloud_subscriber: access token expired, attempting refresh");
468    match crate::muxbus::pkce::refresh_token(&creds, http).await {
469        Ok(refreshed) => {
470            if let Err(e) = wstore.muxbus_save(&refreshed) {
471                tracing::warn!(error = %e, "cloud_subscriber: failed to save refreshed credentials");
472            }
473            Some(refreshed.access_token)
474        }
475        Err(e) => {
476            tracing::warn!(error = %e, "cloud_subscriber: proactive refresh failed");
477            // Fall back to the existing token if it's still within its validity window.
478            // A transient refresh failure should not prevent connecting with a working token.
479            if creds.is_valid() {
480                Some(creds.access_token)
481            } else {
482                None
483            }
484        }
485    }
486}
487
agentmux_srv\muxbus/cloud_subscriber.rs

agentmux_srv\muxbus/
cloud_subscriber.rs
