IdeaSDK/crates/infrastructure/src/mailbox/mod.rs

//! [`InMemoryMailbox`] — the [`AgentMailbox`] adapter (Option 1, lot B-1).
//!
//! The driven side of the inter-agent rendezvous: a per-agent FIFO of
//! [`Ticket`]s, each paired with a one-shot reply channel. `enqueue` appends a
//! ticket and hands the caller a [`PendingReply`] over the receiver; `resolve`
//! feeds the target's `idea_reply` result into the **head** ticket's sender.
//!
//! All the concrete machinery the domain refuses to name lives here: the
//! [`std::collections::VecDeque`] queue, its [`std::sync::Mutex`], and the
//! [`tokio::sync::oneshot`] channel that bridges the resolving call to the awaiting
//! one. The domain port ([`domain::mailbox`]) stays I/O-free.
//!
//! ## Concurrency
//!
//! The map is guarded by a **synchronous** [`Mutex`] held only for the O(1)
//! enqueue/resolve/cancel mutations — **never across an `.await`** (the await is the
//! caller's, on the returned [`PendingReply`], outside the lock). Two `ask`s for the
//! **same** target serialise positionally in that target's `VecDeque`; two `ask`s
//! for **different** targets touch different queues and never contend on the data,
//! only briefly on the map mutex.

use std::collections::{HashMap, VecDeque};
use std::sync::Mutex;

use tokio::sync::oneshot;

use domain::ids::AgentId;
use domain::mailbox::{
    AgentMailbox, AgentQueueSnapshot, MailboxError, PendingReply, QueuedTicketSnapshot, Ticket,
    TicketId,
};

/// One queued request plus the sender that resolves its awaiting [`PendingReply`].
struct Slot {
    ticket: Ticket,
    reply: oneshot::Sender<String>,
}

/// In-memory, per-agent FIFO mailbox (the production [`AgentMailbox`]).
#[derive(Default)]
pub struct InMemoryMailbox {
    queues: Mutex<HashMap<AgentId, VecDeque<Slot>>>,
}

impl InMemoryMailbox {
    /// Creates an empty mailbox.
    #[must_use]
    pub fn new() -> Self {
        Self {
            queues: Mutex::new(HashMap::new()),
        }
    }

    /// Number of tickets currently queued for `agent` (test/inspection helper).
    #[must_use]
    pub fn pending(&self, agent: &AgentId) -> usize {
        self.queues
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .get(agent)
            .map_or(0, VecDeque::len)
    }

    /// The id of the ticket currently at the head of `agent`'s queue, if any
    /// (test/inspection helper).
    #[must_use]
    pub fn head_ticket(&self, agent: &AgentId) -> Option<TicketId> {
        self.queues
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .get(agent)
            .and_then(|q| q.front())
            .map(|slot| slot.ticket.id)
    }
}

impl AgentMailbox for InMemoryMailbox {
    fn enqueue(&self, agent: AgentId, ticket: Ticket) -> PendingReply {
        let (tx, rx) = oneshot::channel::<String>();
        {
            let mut queues = self
                .queues
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            queues
                .entry(agent)
                .or_default()
                .push_back(Slot { ticket, reply: tx });
        }
        // The await happens in the application layer, outside the map lock. A closed
        // channel (sender dropped without a value, e.g. the head was cancelled or the
        // session ended) maps to a typed `Cancelled` rather than a raw recv error.
        PendingReply::new(Box::pin(async move {
            rx.await.map_err(|_| MailboxError::Cancelled)
        }))
    }

    fn resolve(&self, agent: AgentId, result: String) -> Result<(), MailboxError> {
        // Pop the head slot under the lock, then send **outside** any await (the send
        // is non-blocking). If the receiver already went away (caller timed out), the
        // ticket is still correctly retired from the head — the queue advances.
        let slot = {
            let mut queues = self
                .queues
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            let queue = queues
                .get_mut(&agent)
                .filter(|q| !q.is_empty())
                .ok_or(MailboxError::NoPendingRequest(agent))?;
            queue.pop_front().expect("non-empty queue has a head")
        };
        // A dropped receiver (the awaiting caller timed out and went away) is fine:
        // the reply is simply discarded, the head ticket is already retired.
        let _ = slot.reply.send(result);
        Ok(())
    }

    fn resolve_ticket(
        &self,
        agent: AgentId,
        ticket_id: TicketId,
        result: String,
    ) -> Result<(), MailboxError> {
        // Remove the slot whose ticket id matches, anywhere in the queue (multi-thread
        // correlation), then send outside any await. A missing match is a typed
        // NoPendingRequest — never a panic.
        let slot = {
            let mut queues = self
                .queues
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            let queue = queues
                .get_mut(&agent)
                .filter(|q| !q.is_empty())
                .ok_or(MailboxError::NoPendingRequest(agent))?;
            let pos = queue
                .iter()
                .position(|s| s.ticket.id == ticket_id)
                .ok_or(MailboxError::NoPendingRequest(agent))?;
            queue.remove(pos).expect("position just found is in range")
        };
        let _ = slot.reply.send(result);
        Ok(())
    }

    fn cancel_head(&self, agent: AgentId, ticket_id: TicketId) {
        let mut queues = self
            .queues
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if let Some(queue) = queues.get_mut(&agent) {
            // Only retire the head, and only if it is *this* ticket: a head that has
            // since changed (already resolved, or someone else's ticket now in front)
            // must not be dropped. Idempotent and positional.
            if queue.front().map(|s| s.ticket.id) == Some(ticket_id) {
                queue.pop_front();
            }
            if queue.is_empty() {
                queues.remove(&agent);
            }
        }
    }
}

impl AgentQueueSnapshot for InMemoryMailbox {
    fn queue_for(&self, agent: AgentId) -> Vec<QueuedTicketSnapshot> {
        // Pure read: clone each ticket's *data* (never the `oneshot::Sender`) under
        // the lock, recomputing the FIFO position from the current order (0 = head).
        // The queue is observed, not mutated.
        let queues = self
            .queues
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        queues
            .get(&agent)
            .map(|queue| {
                queue
                    .iter()
                    .enumerate()
                    .map(|(idx, slot)| QueuedTicketSnapshot {
                        id: slot.ticket.id,
                        source: slot.ticket.source,
                        conversation: slot.ticket.conversation,
                        requester: slot.ticket.requester.clone(),
                        task: slot.ticket.task.clone(),
                        position: u32::try_from(idx).unwrap_or(u32::MAX),
                    })
                    .collect()
            })
            .unwrap_or_default()
    }
}

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

    fn agent(n: u128) -> AgentId {
        AgentId::from_uuid(uuid::Uuid::from_u128(n))
    }

    fn ticket(n: u128, task: &str) -> Ticket {
        Ticket::new(TicketId::from_uuid(uuid::Uuid::from_u128(n)), "Main", task)
    }

    #[tokio::test]
    async fn enqueue_then_resolve_wakes_the_pending_reply() {
        let mb = InMemoryMailbox::new();
        let a = agent(1);
        let pending = mb.enqueue(a, ticket(10, "do X"));

        mb.resolve(a, "done X".to_owned()).expect("resolve ok");

        let reply = pending.await.expect("reply received");
        assert_eq!(reply, "done X");
        assert_eq!(mb.pending(&a), 0, "queue drained after resolve");
    }

    #[tokio::test]
    async fn two_asks_same_target_resolve_fifo_head_first() {
        let mb = InMemoryMailbox::new();
        let a = agent(1);
        let p1 = mb.enqueue(a, ticket(10, "first"));
        let p2 = mb.enqueue(a, ticket(11, "second"));
        assert_eq!(mb.pending(&a), 2);
        assert_eq!(
            mb.head_ticket(&a),
            Some(TicketId::from_uuid(uuid::Uuid::from_u128(10)))
        );

        // First resolve goes to the FIRST (head) ticket.
        mb.resolve(a, "r1".to_owned()).unwrap();
        assert_eq!(p1.await.unwrap(), "r1");
        // Now the second is at the head.
        assert_eq!(
            mb.head_ticket(&a),
            Some(TicketId::from_uuid(uuid::Uuid::from_u128(11)))
        );
        mb.resolve(a, "r2".to_owned()).unwrap();
        assert_eq!(p2.await.unwrap(), "r2");
    }

    #[tokio::test]
    async fn different_targets_do_not_block_each_other() {
        let mb = InMemoryMailbox::new();
        let a = agent(1);
        let b = agent(2);
        let pa = mb.enqueue(a, ticket(10, "task a"));
        let _pb = mb.enqueue(b, ticket(20, "task b"));

        // Resolving B leaves A untouched; resolving A then completes pa.
        mb.resolve(b, "rb".to_owned()).unwrap();
        assert_eq!(mb.pending(&a), 1, "A's queue is independent of B's");
        mb.resolve(a, "ra".to_owned()).unwrap();
        assert_eq!(pa.await.unwrap(), "ra");
    }

    #[test]
    fn resolve_without_pending_is_a_typed_error() {
        let mb = InMemoryMailbox::new();
        let a = agent(1);
        assert_eq!(
            mb.resolve(a, "orphan".to_owned()),
            Err(MailboxError::NoPendingRequest(a))
        );
    }

    #[tokio::test]
    async fn cancel_head_retires_the_head_and_advances_the_queue() {
        let mb = InMemoryMailbox::new();
        let a = agent(1);
        let p1 = mb.enqueue(a, ticket(10, "stuck"));
        let p2 = mb.enqueue(a, ticket(11, "next"));

        // Caller of ticket 10 timed out: retire exactly its head ticket.
        mb.cancel_head(a, TicketId::from_uuid(uuid::Uuid::from_u128(10)));
        assert_eq!(mb.pending(&a), 1);
        assert_eq!(
            mb.head_ticket(&a),
            Some(TicketId::from_uuid(uuid::Uuid::from_u128(11)))
        );

        // The cancelled pending resolves to Cancelled (its sender was dropped).
        assert_eq!(p1.await, Err(MailboxError::Cancelled));

        // The next ticket is now resolvable normally.
        mb.resolve(a, "r2".to_owned()).unwrap();
        assert_eq!(p2.await.unwrap(), "r2");
    }

    #[test]
    fn cancel_head_is_a_noop_when_head_is_a_different_ticket() {
        let mb = InMemoryMailbox::new();
        let a = agent(1);
        let _p1 = mb.enqueue(a, ticket(10, "head"));
        // Try to cancel a ticket that is NOT the head ⇒ nothing retired.
        mb.cancel_head(a, TicketId::from_uuid(uuid::Uuid::from_u128(99)));
        assert_eq!(mb.pending(&a), 1);
        assert_eq!(
            mb.head_ticket(&a),
            Some(TicketId::from_uuid(uuid::Uuid::from_u128(10)))
        );
    }

    fn tid(n: u128) -> TicketId {
        TicketId::from_uuid(uuid::Uuid::from_u128(n))
    }

    #[test]
    fn snapshot_of_empty_queue_is_empty() {
        let mb = InMemoryMailbox::new();
        assert!(mb.queue_for(agent(1)).is_empty());
    }

    #[test]
    fn snapshot_preserves_fifo_order_and_positions() {
        let mb = InMemoryMailbox::new();
        let a = agent(1);
        let _p1 = mb.enqueue(a, ticket(10, "first"));
        let _p2 = mb.enqueue(a, ticket(11, "second"));

        let snap = mb.queue_for(a);
        assert_eq!(snap.len(), 2);
        assert_eq!(snap[0].id, tid(10));
        assert_eq!(snap[0].position, 0);
        assert_eq!(snap[0].task, "first");
        assert_eq!(snap[1].id, tid(11));
        assert_eq!(snap[1].position, 1);
        assert_eq!(snap[1].task, "second");
    }

    #[test]
    fn snapshot_carries_ticket_metadata() {
        use domain::conversation::ConversationId;
        use domain::input::InputSource;

        let mb = InMemoryMailbox::new();
        let a = agent(1);
        let from = agent(2);
        let conv = ConversationId::from_uuid(uuid::Uuid::from_u128(42));
        let _p = mb.enqueue(
            a,
            Ticket::from_agent(tid(10), from, conv, "Main", "delegate task"),
        );

        let snap = mb.queue_for(a);
        assert_eq!(snap.len(), 1);
        assert_eq!(snap[0].source, InputSource::agent(from));
        assert_eq!(snap[0].conversation, conv);
        assert_eq!(snap[0].requester, "Main");
        assert_eq!(snap[0].task, "delegate task");
    }

    #[test]
    fn snapshot_is_read_only() {
        let mb = InMemoryMailbox::new();
        let a = agent(1);
        let _p1 = mb.enqueue(a, ticket(10, "first"));
        let _p2 = mb.enqueue(a, ticket(11, "second"));

        let _ = mb.queue_for(a);
        let _ = mb.queue_for(a);

        // Observing the queue did not mutate it.
        assert_eq!(mb.pending(&a), 2);
        assert_eq!(mb.head_ticket(&a), Some(tid(10)));
    }

    #[test]
    fn snapshot_updates_after_resolve_ticket() {
        let mb = InMemoryMailbox::new();
        let a = agent(1);
        let _p1 = mb.enqueue(a, ticket(10, "first"));
        let _p2 = mb.enqueue(a, ticket(11, "second"));

        mb.resolve_ticket(a, tid(10), "done".to_owned()).unwrap();

        let snap = mb.queue_for(a);
        assert_eq!(snap.len(), 1);
        assert_eq!(snap[0].id, tid(11));
        assert_eq!(snap[0].position, 0, "remaining ticket becomes head");
    }

    #[test]
    fn snapshot_updates_after_cancel_head() {
        let mb = InMemoryMailbox::new();
        let a = agent(1);
        let _p1 = mb.enqueue(a, ticket(10, "head"));
        let _p2 = mb.enqueue(a, ticket(11, "next"));

        mb.cancel_head(a, tid(10));

        let snap = mb.queue_for(a);
        assert_eq!(snap.len(), 1);
        assert_eq!(snap[0].id, tid(11));
        assert_eq!(snap[0].position, 0, "next ticket recomputed to position 0");
    }
}