agentmux_srv/backend/
messagebus.rs

1// Copyright 2025-2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Local MessageBus for inter-agent communication.
5//!
6//! Provides agent registration, point-to-point messaging, terminal injection,
7//! and broadcast — all over localhost with no cloud dependency.
8
9use std::collections::{HashMap, VecDeque};
10
11use parking_lot::Mutex;
12use serde::{Deserialize, Serialize};
13use tokio::sync::mpsc;
14
15/// Maximum messages queued per offline agent.
16const MAX_OFFLINE_QUEUE: usize = 1000;
17
18/// Maximum messages returned per read_messages call.
19const MAX_READ_LIMIT: usize = 500;
20
21/// Message time-to-live in seconds (1 hour).
22const MESSAGE_TTL_SECS: u64 = 3600;
23
24// ---- Types ----
25
26#[derive(Debug, Clone, Serialize, Deserialize)]
27#[serde(rename_all = "snake_case")]
28pub enum MessageType {
29    Send,
30    Inject,
31    Broadcast,
32}
33
34#[derive(Debug, Clone, Serialize, Deserialize)]
35#[serde(rename_all = "snake_case")]
36pub enum Priority {
37    Normal,
38    High,
39    Urgent,
40}
41
42impl Default for Priority {
43    fn default() -> Self {
44        Priority::Normal
45    }
46}
47
48#[derive(Debug, Clone, Serialize, Deserialize)]
49pub struct BusMessage {
50    pub id: String,
51    pub from: String,
52    pub to: String,
53    #[serde(rename = "type")]
54    pub msg_type: MessageType,
55    pub payload: String,
56    #[serde(default)]
57    pub priority: Priority,
58    pub timestamp: u64,
59    /// Identity M4c-2d (SPEC_AGENT_IDENTITY_CARRIED_NOT_DERIVED_2026_09_23.md
60    /// §6.5.9): the sender's UID when its HTTP request was attributed (its
61    /// `X-Agent-Token`), beside the `from` name it claimed. Empty =
62    /// Unattributed (the WebSocket bus has no per-request caller). Set by
63    /// this srv only; a bus message never leaves it for another srv.
64    #[serde(default, skip_serializing_if = "String::is_empty")]
65    pub from_uid: String,
66}
67
68impl BusMessage {
69    pub fn new(from: &str, to: &str, msg_type: MessageType, payload: &str, priority: Priority) -> Self {
70        let now = agentmux_common::time::now_secs_u64();
71        Self {
72            id: uuid::Uuid::new_v4().to_string(),
73            from: from.to_string(),
74            to: to.to_string(),
75            msg_type,
76            payload: payload.to_string(),
77            priority,
78            timestamp: now,
79            from_uid: String::new(),
80        }
81    }
82
83    /// The sender's attributed UID (see [`Self::from_uid`]).
84    pub fn with_from_uid(mut self, from_uid: &str) -> Self {
85        self.from_uid = from_uid.to_string();
86        self
87    }
88
89    fn is_expired(&self) -> bool {
90        let now = agentmux_common::time::now_secs_u64();
91        now.saturating_sub(self.timestamp) > MESSAGE_TTL_SECS
92    }
93}
94
95/// Info about a connected agent.
96#[derive(Debug, Clone, Serialize)]
97pub struct AgentInfo {
98    pub id: String,
99    pub registered_at: u64,
100    pub last_seen: u64,
101    pub connection_type: String, // "websocket" or "http"
102}
103
104/// Internal agent connection state.
105struct AgentConnection {
106    info: AgentInfo,
107    /// Channel sender for pushing messages to this agent's WebSocket.
108    /// None for HTTP-polling agents (they read from the offline queue instead).
109    ws_sender: Option<mpsc::UnboundedSender<BusMessage>>,
110}
111
112// ---- MessageBus ----
113
114pub struct MessageBus {
115    agents: Mutex<HashMap<String, AgentConnection>>,
116    offline_queues: Mutex<HashMap<String, VecDeque<BusMessage>>>,
117}
118
119impl MessageBus {
120    pub fn new() -> Self {
121        Self {
122            agents: Mutex::new(HashMap::new()),
123            offline_queues: Mutex::new(HashMap::new()),
124        }
125    }
126
127    /// Register an agent on the bus with a WebSocket push channel.
128    /// Returns a receiver for messages pushed to this agent.
129    pub fn register(&self, agent_id: &str, connection_type: &str) -> mpsc::UnboundedReceiver<BusMessage> {
130        let (tx, rx) = mpsc::unbounded_channel();
131        let now = agentmux_common::time::now_secs_u64();
132
133        let conn = AgentConnection {
134            info: AgentInfo {
135                id: agent_id.to_string(),
136                registered_at: now,
137                last_seen: now,
138                connection_type: connection_type.to_string(),
139            },
140            ws_sender: Some(tx),
141        };
142
143        self.agents.lock().insert(agent_id.to_string(), conn);
144
145        tracing::info!("messagebus: agent '{}' registered ({})", agent_id, connection_type);
146
147        // Drain any offline queued messages
148        self.drain_offline_queue(agent_id);
149
150        rx
151    }
152
153    /// Register an HTTP-polling agent (no WebSocket push channel).
154    /// Messages sent to this agent are queued in the offline queue and
155    /// retrieved via `read_messages`.
156    pub fn register_http(&self, agent_id: &str) {
157        let now = agentmux_common::time::now_secs_u64();
158
159        let conn = AgentConnection {
160            info: AgentInfo {
161                id: agent_id.to_string(),
162                registered_at: now,
163                last_seen: now,
164                connection_type: "http".to_string(),
165            },
166            ws_sender: None,
167        };
168
169        self.agents.lock().insert(agent_id.to_string(), conn);
170
171        tracing::info!("messagebus: agent '{}' registered (http-polling)", agent_id);
172    }
173
174    /// Unregister an agent from the bus.
175    pub fn unregister(&self, agent_id: &str) {
176        if self.agents.lock().remove(agent_id).is_some() {
177            tracing::info!("messagebus: agent '{}' unregistered", agent_id);
178        }
179    }
180
181    /// Update last_seen timestamp for an agent (called on HTTP polling).
182    pub fn touch(&self, agent_id: &str) {
183        if let Some(conn) = self.agents.lock().get_mut(agent_id) {
184            conn.info.last_seen = agentmux_common::time::now_secs_u64();
185        }
186    }
187
188    /// Send a message to a specific agent.
189    pub fn send(&self, msg: BusMessage) -> Result<(), String> {
190        let target = msg.to.clone();
191
192        // Try to push via WebSocket channel (short lock scope)
193        let sent = {
194            let agents = self.agents.lock();
195            if let Some(conn) = agents.get(&target) {
196                if let Some(ref tx) = conn.ws_sender {
197                    tx.send(msg.clone()).is_ok()
198                } else {
199                    false
200                }
201            } else {
202                false
203            }
204        };
205
206        if sent {
207            return Ok(());
208        }
209
210        // Agent not connected or send failed — queue for later
211        self.queue_offline(msg);
212        Ok(())
213    }
214
215    /// Inject a message into an agent's terminal (jekt).
216    /// This is the same as send but with MessageType::Inject.
217    pub fn inject(
218        &self,
219        from: &str,
220        from_uid: &str,
221        target: &str,
222        message: &str,
223        priority: Priority,
224    ) -> Result<String, String> {
225        let msg = BusMessage::new(from, target, MessageType::Inject, message, priority).with_from_uid(from_uid);
226        let msg_id = msg.id.clone();
227        self.send(msg)?;
228        Ok(msg_id)
229    }
230
231    /// Broadcast a message to all connected agents (except sender).
232    pub fn broadcast(&self, from: &str, from_uid: &str, payload: &str, priority: Priority) -> Result<usize, String> {
233        // Collect senders under short lock
234        let targets: Vec<(String, mpsc::UnboundedSender<BusMessage>)> = {
235            let agents = self.agents.lock();
236            agents
237                .iter()
238                .filter(|(id, _)| id.as_str() != from)
239                .filter_map(|(id, conn)| {
240                    conn.ws_sender.as_ref().map(|tx| (id.clone(), tx.clone()))
241                })
242                .collect()
243        };
244
245        let mut delivered = 0;
246        for (agent_id, tx) in targets {
247            let msg = BusMessage::new(from, &agent_id, MessageType::Broadcast, payload, priority.clone())
248                .with_from_uid(from_uid);
249            if tx.send(msg).is_ok() {
250                delivered += 1;
251            }
252        }
253
254        Ok(delivered)
255    }
256
257    /// List all registered agents.
258    pub fn list_agents(&self) -> Vec<AgentInfo> {
259        self.agents.lock().values().map(|c| c.info.clone()).collect()
260    }
261
262    /// Read (and drain) queued offline messages for an agent.
263    /// Used by HTTP-polling agents that don't have a WebSocket connection.
264    pub fn read_messages(&self, agent_id: &str, limit: usize) -> Vec<BusMessage> {
265        self.touch(agent_id);
266        let limit = limit.min(MAX_READ_LIMIT);
267        let mut queues = self.offline_queues.lock();
268        let queue = match queues.get_mut(agent_id) {
269            Some(q) => q,
270            None => return Vec::new(),
271        };
272
273        // Purge expired messages
274        queue.retain(|m| !m.is_expired());
275
276        let count = limit.min(queue.len());
277        queue.drain(..count).collect()
278    }
279
280    /// Delete specific messages by ID from an agent's offline queue.
281    pub fn delete_messages(&self, agent_id: &str, message_ids: &[String]) -> usize {
282        let mut queues = self.offline_queues.lock();
283        let queue = match queues.get_mut(agent_id) {
284            Some(q) => q,
285            None => return 0,
286        };
287
288        let before = queue.len();
289        queue.retain(|m| !message_ids.contains(&m.id));
290        before - queue.len()
291    }
292
293    // ---- Internal ----
294
295    fn queue_offline(&self, msg: BusMessage) {
296        let target = msg.to.clone();
297        let mut queues = self.offline_queues.lock();
298        let queue = queues.entry(target).or_insert_with(VecDeque::new);
299
300        // Evict oldest if at capacity
301        if queue.len() >= MAX_OFFLINE_QUEUE {
302            queue.pop_front();
303        }
304        queue.push_back(msg);
305    }
306
307    fn drain_offline_queue(&self, agent_id: &str) {
308        // Collect messages and the sender under separate lock scopes
309        let messages: Vec<BusMessage> = {
310            let mut queues = self.offline_queues.lock();
311            match queues.get_mut(agent_id) {
312                Some(queue) => {
313                    queue.retain(|m| !m.is_expired());
314                    queue.drain(..).collect()
315                }
316                None => return,
317            }
318        };
319
320        if messages.is_empty() {
321            return;
322        }
323
324        // Get the sender clone, then drop the lock before sending
325        let tx = {
326            let agents = self.agents.lock();
327            agents
328                .get(agent_id)
329                .and_then(|conn| conn.ws_sender.clone())
330        };
331
332        if let Some(tx) = tx {
333            let count = messages.len();
334            for msg in messages {
335                let _ = tx.send(msg);
336            }
337            tracing::info!("messagebus: drained {} offline messages to '{}'", count, agent_id);
338        }
339    }
340}
341
342#[cfg(test)]
343mod tests {
344    use super::*;
345
346    /// Identity M4c-2d: an attributed sender's UID rides the message beside
347    /// its claimed name; an Unattributed one's is omitted.
348    #[test]
349    fn a_bus_message_carries_an_attributed_from_uid() {
350        let m = BusMessage::new("agenty", "b", MessageType::Send, "p", Priority::Normal);
351        assert!(serde_json::to_value(&m).unwrap().get("from_uid").is_none());
352        let m = m.with_from_uid("uid-y");
353        assert_eq!(serde_json::to_value(&m).unwrap()["from_uid"], "uid-y");
354    }
355}
agentmux_srv/backend/messagebus.rs

agentmux_srv/backend/
messagebus.rs
