PicoBot/src/bus/mod.rs

pub mod dispatcher;
pub mod message;

pub use dispatcher::OutboundDispatcher;
pub use message::{
    ChannelContext, ChatMessage, CommittedMessage, CommittedTurnDelta, CompletionStatus,
    ContentBlock, ControlMessage, InboundMessage, MediaItem, MediaRef, MessageSource,
    OutboundMessage, ProviderReasoningState, SourceKind,
};

use std::sync::Arc;
use tokio::sync::{Mutex, mpsc};

// ============================================================================
// MessageBus - Async message queue for Channel <-> Agent communication
// ============================================================================

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>>,
    // Control channel for session management operations
    control_tx: mpsc::Sender<ControlMessage>,
    control_rx: Mutex<mpsc::Receiver<ControlMessage>>,
}

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);
        let (control_tx, control_rx) = mpsc::channel(capacity);
        Arc::new(Self {
            inbound_tx,
            outbound_tx,
            inbound_rx: Mutex::new(inbound_rx),
            outbound_rx: Mutex::new(outbound_rx),
            control_tx,
            control_rx: Mutex::new(control_rx),
        })
    }

    /// Publish an inbound message (Channel -> Bus)
    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 an inbound message (Agent -> Bus)
    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 an outbound message (Agent -> Bus)
    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)
    }

    /// Publish an outbound message and wait for the dispatcher to report the
    /// actual channel delivery result.
    pub async fn deliver_outbound(&self, mut msg: OutboundMessage) -> Result<(), BusError> {
        let (delivery_tx, mut delivery_rx) = tokio::sync::watch::channel(None);
        msg.delivery = Some(delivery_tx);
        self.publish_outbound(msg).await?;

        tokio::time::timeout(std::time::Duration::from_secs(120), async {
            loop {
                delivery_rx.changed().await.map_err(|_| BusError::Closed)?;
                if let Some(result) = delivery_rx.borrow().clone() {
                    return result.map_err(BusError::DeliveryFailed);
                }
            }
        })
        .await
        .map_err(|_| BusError::DeliveryTimedOut)?
    }

    /// Consume an outbound message (Dispatcher -> Bus)
    pub async fn consume_outbound(&self) -> Option<OutboundMessage> {
        self.outbound_rx.lock().await.recv().await
    }

    /// Publish a control message (Channel -> Bus for session management)
    pub async fn publish_control(&self, msg: ControlMessage) -> Result<(), BusError> {
        tracing::debug!(op = ?msg.op, "Bus: publishing control message");
        self.control_tx
            .send(msg)
            .await
            .map_err(|_| BusError::Closed)
    }

    /// Consume a control message (ControlProcessor -> Bus)
    pub async fn consume_control(&self) -> Option<ControlMessage> {
        self.control_rx.lock().await.recv().await
    }

    /// Snapshot of the current depth and capacity of each bus queue.
    pub fn queue_depths(&self) -> QueueDepths {
        QueueDepths {
            inbound_depth: (self.inbound_tx.max_capacity() - self.inbound_tx.capacity()) as u64,
            inbound_cap: self.inbound_tx.max_capacity() as u64,
            outbound_depth: (self.outbound_tx.max_capacity() - self.outbound_tx.capacity()) as u64,
            outbound_cap: self.outbound_tx.max_capacity() as u64,
            control_depth: (self.control_tx.max_capacity() - self.control_tx.capacity()) as u64,
            control_cap: self.control_tx.max_capacity() as u64,
        }
    }
}

/// Read-only snapshot of MessageBus queue utilization.
#[derive(serde::Serialize)]
pub struct QueueDepths {
    pub inbound_depth: u64,
    pub inbound_cap: u64,
    pub outbound_depth: u64,
    pub outbound_cap: u64,
    pub control_depth: u64,
    pub control_cap: u64,
}

// ============================================================================
// BusError
// ============================================================================

#[derive(Debug)]
pub enum BusError {
    Closed,
    DeliveryFailed(String),
    DeliveryTimedOut,
}

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"),
            BusError::DeliveryFailed(error) => write!(f, "Outbound delivery failed: {error}"),
            BusError::DeliveryTimedOut => write!(f, "Outbound delivery confirmation timed out"),
        }
    }
}

impl std::error::Error for BusError {}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;

    #[tokio::test]
    async fn queue_depths_report_retained_sender_usage() {
        let bus = MessageBus::new(3);

        let empty = bus.queue_depths();
        assert_eq!(empty.inbound_depth, 0);
        assert_eq!(empty.inbound_cap, 3);
        assert_eq!(empty.outbound_depth, 0);
        assert_eq!(empty.outbound_cap, 3);
        assert_eq!(empty.control_depth, 0);
        assert_eq!(empty.control_cap, 3);

        bus.publish_outbound(OutboundMessage {
            channel: "test".to_string(),
            chat_id: "chat".to_string(),
            content: "queued".to_string(),
            reply_to: None,
            media: vec![],
            metadata: HashMap::new(),
            delivery: None,
        })
        .await
        .unwrap();

        assert_eq!(bus.queue_depths().outbound_depth, 1);
        bus.consume_outbound().await.unwrap();
        assert_eq!(bus.queue_depths().outbound_depth, 0);
    }
}