PicoBot/src/bus/mod.rs

90 lines
3.2 KiB
Rust

pub mod message;
pub use crate::domain::messages::ContentBlock;
pub use message::{
ChatMessage, InboundMessage, MediaItem, OutboundMessage, SYSTEM_CONTEXT_AGENT_PROMPT,
SYSTEM_CONTEXT_HISTORY_COMPACTION, SYSTEM_CONTEXT_SCHEDULED_PROMPT,
};
use std::sync::Arc;
use tokio::sync::{Mutex, mpsc};
// ============================================================================
// MessageBus - async inbound/outbound queues
// ============================================================================
pub struct MessageBus {
inbound_tx: mpsc::Sender<InboundMessage>,
outbound_tx: mpsc::Sender<OutboundMessage>,
inbound_rx: Mutex<mpsc::Receiver<InboundMessage>>,
outbound_rx: Mutex<mpsc::Receiver<OutboundMessage>>,
}
impl MessageBus {
/// Create a new MessageBus with the given channel capacity
pub fn new(capacity: usize) -> Arc<Self> {
let (inbound_tx, inbound_rx) = mpsc::channel(capacity);
let (outbound_tx, outbound_rx) = mpsc::channel(capacity);
Arc::new(Self {
inbound_tx,
outbound_tx,
inbound_rx: Mutex::new(inbound_rx),
outbound_rx: Mutex::new(outbound_rx),
})
}
/// Publish a message to the inbound queue
pub async fn publish_inbound(&self, msg: InboundMessage) -> Result<(), BusError> {
#[cfg(debug_assertions)]
tracing::debug!(channel = %msg.channel, sender = %msg.sender_id, chat = %msg.chat_id, content_len = %msg.content.len(), media_count = %msg.media.len(), "Bus: publishing inbound message");
self.inbound_tx
.send(msg)
.await
.map_err(|_| BusError::Closed)
}
/// Consume a message from the inbound queue.
/// Returns `None` when the channel is closed (all senders dropped).
pub async fn consume_inbound(&self) -> Option<InboundMessage> {
let msg = self.inbound_rx.lock().await.recv().await?;
#[cfg(debug_assertions)]
tracing::debug!(channel = %msg.channel, sender = %msg.sender_id, chat = %msg.chat_id, "Bus: consuming inbound message");
Some(msg)
}
/// Publish a message to the outbound queue
pub async fn publish_outbound(&self, msg: OutboundMessage) -> Result<(), BusError> {
#[cfg(debug_assertions)]
tracing::debug!(channel = %msg.channel, chat_id = %msg.chat_id, content_len = %msg.content.len(), "Bus: publishing outbound message");
self.outbound_tx
.send(msg)
.await
.map_err(|_| BusError::Closed)
}
/// Consume an outbound message from the outbound queue.
/// Returns `None` when the channel is closed (all senders dropped).
pub async fn consume_outbound(&self) -> Option<OutboundMessage> {
self.outbound_rx.lock().await.recv().await
}
}
// ============================================================================
// BusError
// ============================================================================
#[derive(Debug)]
pub enum BusError {
Closed,
}
impl std::fmt::Display for BusError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
BusError::Closed => write!(f, "Bus channel closed"),
}
}
}
impl std::error::Error for BusError {}