Skip to main content

starweaver_context/
message_bus.rs

1//! Subscriber/cursor message bus records for inter-agent communication.
2
3use std::collections::{BTreeMap, BTreeSet};
4
5use chrono::{DateTime, Utc};
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8use starweaver_core::Metadata;
9use uuid::Uuid;
10
11/// Default maximum number of messages retained by a [`MessageBus`].
12pub const DEFAULT_MESSAGE_BUS_MAXLEN: usize = 500;
13
14/// Inter-agent or user steering message.
15#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
16pub struct BusMessage {
17    /// Stable message id used for idempotent send and consume.
18    pub id: String,
19    /// Renderable text or JSON content.
20    #[serde(default)]
21    pub content: Value,
22    /// Sender identifier such as `user`, `main`, or a subagent id.
23    pub source: String,
24    /// Recipient agent id. `None` means broadcast to all subscribers.
25    #[serde(default, skip_serializing_if = "Option::is_none")]
26    pub target: Option<String>,
27    /// Optional template string. The Rust context stores the value; rendering is caller-owned.
28    #[serde(default, skip_serializing_if = "Option::is_none")]
29    pub template: Option<String>,
30    /// Message creation time.
31    #[serde(default = "Utc::now")]
32    pub timestamp: DateTime<Utc>,
33    /// Additional message metadata.
34    #[serde(default, skip_serializing_if = "Metadata::is_empty")]
35    pub metadata: Metadata,
36}
37
38impl BusMessage {
39    /// Create a topic-scoped message.
40    #[must_use]
41    pub fn new(topic: impl Into<String>, payload: Value) -> Self {
42        let topic = topic.into();
43        let mut metadata = Metadata::default();
44        metadata.insert("starweaver.topic".to_string(), Value::String(topic));
45        Self {
46            id: Uuid::new_v4().simple().to_string(),
47            content: payload,
48            source: "system".to_string(),
49            target: None,
50            template: None,
51            timestamp: Utc::now(),
52            metadata,
53        }
54    }
55
56    /// Create a text bus message.
57    #[must_use]
58    pub fn text(content: impl Into<String>, source: impl Into<String>) -> Self {
59        let content = content.into();
60        Self {
61            id: Uuid::new_v4().simple().to_string(),
62            content: Value::String(content),
63            source: source.into(),
64            target: None,
65            template: None,
66            timestamp: Utc::now(),
67            metadata: Metadata::default(),
68        }
69    }
70
71    /// Set an explicit message id.
72    #[must_use]
73    pub fn with_id(mut self, id: impl Into<String>) -> Self {
74        self.id = id.into();
75        self
76    }
77
78    /// Set target recipient.
79    #[must_use]
80    pub fn with_target(mut self, target: impl Into<String>) -> Self {
81        self.target = Some(target.into());
82        self
83    }
84
85    /// Set source sender.
86    #[must_use]
87    pub fn with_source(mut self, source: impl Into<String>) -> Self {
88        self.source = source.into();
89        self
90    }
91
92    /// Set template metadata.
93    #[must_use]
94    pub fn with_template(mut self, template: impl Into<String>) -> Self {
95        self.template = Some(template.into());
96        self
97    }
98
99    /// Return message text for logs, steering, and display.
100    #[must_use]
101    pub fn content_text(&self) -> String {
102        if let Some(text) = self.content.as_str() {
103            return text.to_string();
104        }
105        if let Some(text) = self.content.get("text").and_then(Value::as_str) {
106            return text.to_string();
107        }
108        if let Some(text) = self.content.get("message").and_then(Value::as_str) {
109            return text.to_string();
110        }
111        self.content.to_string()
112    }
113
114    /// Render message content. Minimal template support keeps raw content unless no JSON text exists.
115    #[must_use]
116    pub fn render_text(&self) -> String {
117        let content = self.content_text();
118        if let Some(template) = &self.template {
119            template.replace("{{ content }}", &content)
120        } else {
121            content
122        }
123    }
124}
125
126impl Eq for BusMessage {}
127
128/// Subscriber/cursor message bus with idempotent send and consume.
129#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
130pub struct MessageBus {
131    #[serde(default)]
132    messages: Vec<BusMessage>,
133    #[serde(default)]
134    message_ids: BTreeSet<String>,
135    #[serde(default)]
136    cursors: BTreeMap<String, usize>,
137    #[serde(default)]
138    consumed_ids: BTreeMap<String, BTreeSet<String>>,
139    #[serde(default = "default_maxlen")]
140    maxlen: usize,
141}
142
143impl Default for MessageBus {
144    fn default() -> Self {
145        Self::new()
146    }
147}
148
149impl MessageBus {
150    /// Create an empty message bus.
151    #[must_use]
152    pub const fn new() -> Self {
153        Self::with_maxlen(DEFAULT_MESSAGE_BUS_MAXLEN)
154    }
155
156    /// Create an empty message bus with a bounded retention length.
157    #[must_use]
158    pub const fn with_maxlen(maxlen: usize) -> Self {
159        Self {
160            messages: Vec::new(),
161            message_ids: BTreeSet::new(),
162            cursors: BTreeMap::new(),
163            consumed_ids: BTreeMap::new(),
164            maxlen,
165        }
166    }
167
168    /// Register a subscriber. New subscribers start at the current tail.
169    pub fn subscribe(&mut self, agent_id: impl Into<String>) {
170        let agent_id = agent_id.into();
171        if !self.cursors.contains_key(&agent_id) {
172            self.cursors.insert(agent_id.clone(), self.messages.len());
173            self.consumed_ids.insert(agent_id, BTreeSet::new());
174        }
175    }
176
177    /// Remove a subscriber and its cursor state.
178    pub fn unsubscribe(&mut self, agent_id: &str) {
179        self.cursors.remove(agent_id);
180        self.consumed_ids.remove(agent_id);
181    }
182
183    /// Return whether the given agent currently has a subscriber cursor.
184    #[must_use]
185    pub fn is_subscribed(&self, agent_id: &str) -> bool {
186        self.cursors.contains_key(agent_id)
187    }
188
189    /// Send a message idempotently.
190    pub fn send(&mut self, message: BusMessage) -> BusMessage {
191        if self.message_ids.contains(&message.id) {
192            return self
193                .messages
194                .iter()
195                .find(|existing| existing.id == message.id)
196                .cloned()
197                .unwrap_or(message);
198        }
199        self.message_ids.insert(message.id.clone());
200        self.messages.push(message.clone());
201        self.trim();
202        message
203    }
204
205    /// Consume unread messages for a subscriber.
206    pub fn consume(&mut self, agent_id: impl Into<String>) -> Vec<BusMessage> {
207        self.consume_matching(agent_id, |_| true)
208    }
209
210    /// Consume unread messages matching a predicate while preserving other pending messages.
211    pub fn consume_matching(
212        &mut self,
213        agent_id: impl Into<String>,
214        predicate: impl Fn(&BusMessage) -> bool,
215    ) -> Vec<BusMessage> {
216        let agent_id = agent_id.into();
217        if !self.cursors.contains_key(&agent_id) {
218            self.subscribe(agent_id.clone());
219        }
220        let cursor = self.cursors.get(&agent_id).copied().unwrap_or_default();
221        let consumed = self.consumed_ids.entry(agent_id.clone()).or_default();
222        let mut result = Vec::new();
223        for message in self.messages.iter().skip(cursor) {
224            if is_deliverable(message, &agent_id)
225                && !consumed.contains(&message.id)
226                && predicate(message)
227            {
228                consumed.insert(message.id.clone());
229                result.push(message.clone());
230            }
231        }
232        let next_cursor = self
233            .messages
234            .iter()
235            .enumerate()
236            .skip(cursor)
237            .find(|(_index, message)| {
238                is_deliverable(message, &agent_id) && !consumed.contains(&message.id)
239            })
240            .map_or(self.messages.len(), |(index, _message)| index);
241        self.cursors.insert(agent_id, next_cursor);
242        result
243    }
244
245    /// Return unread messages without advancing the cursor.
246    #[must_use]
247    pub fn peek(&self, agent_id: &str) -> Vec<BusMessage> {
248        let Some(cursor) = self.cursors.get(agent_id).copied() else {
249            return Vec::new();
250        };
251        let consumed = self.consumed_ids.get(agent_id);
252        self.messages
253            .iter()
254            .skip(cursor)
255            .filter(|message| {
256                is_deliverable(message, agent_id)
257                    && consumed.is_none_or(|ids| !ids.contains(&message.id))
258            })
259            .cloned()
260            .collect()
261    }
262
263    /// Mark specific unread messages as consumed for a subscriber.
264    ///
265    /// Returns the number of currently unread messages that were marked consumed.
266    pub fn mark_consumed(
267        &mut self,
268        agent_id: impl Into<String>,
269        message_ids: &BTreeSet<String>,
270    ) -> usize {
271        let agent_id = agent_id.into();
272        if !self.cursors.contains_key(&agent_id) || message_ids.is_empty() {
273            return 0;
274        }
275        let cursor = self.cursors.get(&agent_id).copied().unwrap_or_default();
276        let consumed = self.consumed_ids.entry(agent_id.clone()).or_default();
277        let mut marked = 0;
278        for message in self.messages.iter().skip(cursor) {
279            if message_ids.contains(&message.id)
280                && is_deliverable(message, &agent_id)
281                && !consumed.contains(&message.id)
282            {
283                consumed.insert(message.id.clone());
284                marked += 1;
285            }
286        }
287        marked
288    }
289
290    /// Return whether there are pending messages for one subscriber.
291    #[must_use]
292    pub fn has_pending(&self, agent_id: &str) -> bool {
293        !self.peek(agent_id).is_empty()
294    }
295
296    /// Return whether any queued message has the provided topic.
297    #[must_use]
298    pub fn has_topic(&self, topic: &str) -> bool {
299        self.messages.iter().any(|message| {
300            message
301                .metadata
302                .get("starweaver.topic")
303                .and_then(Value::as_str)
304                == Some(topic)
305        })
306    }
307
308    /// Clear all messages and subscriber state.
309    pub fn clear(&mut self) {
310        self.messages.clear();
311        self.message_ids.clear();
312        self.cursors.clear();
313        self.consumed_ids.clear();
314    }
315
316    /// Return number of retained messages.
317    #[must_use]
318    pub const fn len(&self) -> usize {
319        self.messages.len()
320    }
321
322    /// Return number of subscribers.
323    #[must_use]
324    pub fn subscriber_count(&self) -> usize {
325        self.cursors.len()
326    }
327
328    /// Return whether the bus has no messages.
329    #[must_use]
330    pub const fn is_empty(&self) -> bool {
331        self.messages.is_empty()
332    }
333
334    /// Return retained messages.
335    #[must_use]
336    pub fn messages(&self) -> &[BusMessage] {
337        &self.messages
338    }
339
340    fn trim(&mut self) {
341        if self.maxlen == 0 {
342            self.clear();
343            return;
344        }
345        if self.messages.len() <= self.maxlen {
346            return;
347        }
348        let trim_count = self.messages.len() - self.maxlen;
349        let trimmed_ids = self
350            .messages
351            .iter()
352            .take(trim_count)
353            .map(|message| message.id.clone())
354            .collect::<BTreeSet<_>>();
355        self.messages.drain(0..trim_count);
356        for id in &trimmed_ids {
357            self.message_ids.remove(id);
358        }
359        for cursor in self.cursors.values_mut() {
360            *cursor = cursor.saturating_sub(trim_count);
361        }
362        for consumed in self.consumed_ids.values_mut() {
363            for id in &trimmed_ids {
364                consumed.remove(id);
365            }
366        }
367    }
368}
369
370const fn default_maxlen() -> usize {
371    DEFAULT_MESSAGE_BUS_MAXLEN
372}
373
374fn is_deliverable(message: &BusMessage, agent_id: &str) -> bool {
375    message.target.is_none() || message.target.as_deref() == Some(agent_id)
376}