use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use tokio::sync::mpsc;
use uuid::Uuid;
use crate::{AgentId, Error, Result};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct MessageId(Uuid);
impl MessageId {
pub fn new() -> Self {
Self(Uuid::new_v4())
}
}
impl Default for MessageId {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum MessagePayload {
Start,
Data(serde_json::Value),
Complete,
Error(String),
StatusRequest,
StatusResponse(String),
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
pub id: MessageId,
pub from: AgentId,
pub to: AgentId,
pub payload: MessagePayload,
pub timestamp: i64,
}
impl Message {
pub fn new(from: AgentId, to: AgentId, payload: MessagePayload) -> Self {
Self {
id: MessageId::new(),
from,
to,
payload,
timestamp: chrono::Utc::now().timestamp(),
}
}
pub fn data(from: AgentId, to: AgentId, data: serde_json::Value) -> Self {
Self::new(from, to, MessagePayload::Data(data))
}
pub fn start(from: AgentId, to: AgentId) -> Self {
Self::new(from, to, MessagePayload::Start)
}
pub fn complete(from: AgentId, to: AgentId) -> Self {
Self::new(from, to, MessagePayload::Complete)
}
pub fn error(from: AgentId, to: AgentId, error: impl Into<String>) -> Self {
Self::new(from, to, MessagePayload::Error(error.into()))
}
}
pub struct MessageBus {
channels: HashMap<AgentId, mpsc::UnboundedSender<Message>>,
}
impl MessageBus {
pub fn new() -> Self {
Self {
channels: HashMap::new(),
}
}
pub fn register(&mut self, agent_id: AgentId) -> mpsc::UnboundedReceiver<Message> {
let (tx, rx) = mpsc::unbounded_channel();
self.channels.insert(agent_id, tx);
rx
}
pub fn unregister(&mut self, agent_id: &AgentId) {
self.channels.remove(agent_id);
}
pub async fn send(&self, message: Message) -> Result<()> {
let channel = self.channels
.get(&message.to)
.ok_or_else(|| Error::MessageBus(format!("Agent {} not registered", message.to)))?;
channel
.send(message)
.map_err(|e| Error::MessageBus(format!("Failed to send message: {}", e)))?;
Ok(())
}
pub async fn broadcast(&self, message: Message) -> Result<()> {
for channel in self.channels.values() {
channel
.send(message.clone())
.map_err(|e| Error::MessageBus(format!("Failed to broadcast: {}", e)))?;
}
Ok(())
}
pub fn agent_count(&self) -> usize {
self.channels.len()
}
}
impl Default for MessageBus {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_message_creation() {
let from = AgentId::new();
let to = AgentId::new();
let msg = Message::start(from, to);
assert_eq!(msg.from, from);
assert_eq!(msg.to, to);
assert!(matches!(msg.payload, MessagePayload::Start));
}
#[test]
fn test_message_bus_registration() {
let mut bus = MessageBus::new();
let agent_id = AgentId::new();
let _rx = bus.register(agent_id);
assert_eq!(bus.agent_count(), 1);
bus.unregister(&agent_id);
assert_eq!(bus.agent_count(), 0);
}
#[tokio::test]
async fn test_message_bus_send() {
let mut bus = MessageBus::new();
let agent1 = AgentId::new();
let agent2 = AgentId::new();
let mut rx = bus.register(agent2);
let _tx = bus.register(agent1);
let msg = Message::start(agent1, agent2);
bus.send(msg.clone()).await.unwrap();
let received = rx.recv().await.unwrap();
assert_eq!(received.id, msg.id);
assert_eq!(received.from, agent1);
}
}