starweaver_context/
message_bus.rs1use 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
11pub const DEFAULT_MESSAGE_BUS_MAXLEN: usize = 500;
13
14#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
16pub struct BusMessage {
17 pub id: String,
19 #[serde(default)]
21 pub content: Value,
22 pub source: String,
24 #[serde(default, skip_serializing_if = "Option::is_none")]
26 pub target: Option<String>,
27 #[serde(default, skip_serializing_if = "Option::is_none")]
29 pub template: Option<String>,
30 #[serde(default = "Utc::now")]
32 pub timestamp: DateTime<Utc>,
33 #[serde(default, skip_serializing_if = "Metadata::is_empty")]
35 pub metadata: Metadata,
36}
37
38impl BusMessage {
39 #[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 #[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 #[must_use]
73 pub fn with_id(mut self, id: impl Into<String>) -> Self {
74 self.id = id.into();
75 self
76 }
77
78 #[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 #[must_use]
87 pub fn with_source(mut self, source: impl Into<String>) -> Self {
88 self.source = source.into();
89 self
90 }
91
92 #[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 #[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 #[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#[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 #[must_use]
152 pub const fn new() -> Self {
153 Self::with_maxlen(DEFAULT_MESSAGE_BUS_MAXLEN)
154 }
155
156 #[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 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 pub fn unsubscribe(&mut self, agent_id: &str) {
179 self.cursors.remove(agent_id);
180 self.consumed_ids.remove(agent_id);
181 }
182
183 #[must_use]
185 pub fn is_subscribed(&self, agent_id: &str) -> bool {
186 self.cursors.contains_key(agent_id)
187 }
188
189 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 pub fn consume(&mut self, agent_id: impl Into<String>) -> Vec<BusMessage> {
207 self.consume_matching(agent_id, |_| true)
208 }
209
210 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 #[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 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 #[must_use]
292 pub fn has_pending(&self, agent_id: &str) -> bool {
293 !self.peek(agent_id).is_empty()
294 }
295
296 #[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 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 #[must_use]
318 pub const fn len(&self) -> usize {
319 self.messages.len()
320 }
321
322 #[must_use]
324 pub fn subscriber_count(&self) -> usize {
325 self.cursors.len()
326 }
327
328 #[must_use]
330 pub const fn is_empty(&self) -> bool {
331 self.messages.is_empty()
332 }
333
334 #[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}