agent-mesh-bus 0.8.3

//! Request-reply correlation.
//!
//! Each outbound [`crate::inbox::BusMessage::Request`] carries a
//! [`CorrelationId`] generated by the requester. The matching
//! [`crate::inbox::BusMessage::Reply`] echoes that id back so the bus
//! can route the response to the right waiter. [`ReplyWaiter`] holds
//! the in-flight oneshot senders, keyed by correlation and bound to the
//! expected responder's authenticated agent fingerprint.

use agent_mesh_protocol::Fingerprint;
use rand::RngCore;
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Mutex;
use std::task::{Context, Poll};
use tokio::sync::oneshot;

/// A 16-byte correlation identifier.
///
/// 128 bits is enough collision space that the requester can draw a
/// fresh one at random without coordination, even at very high request
/// rates. The bytes are echoed verbatim in the reply.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct CorrelationId(pub [u8; 16]);

impl CorrelationId {
    /// Draw a fresh random correlation id from the OS RNG.
    #[must_use]
    pub fn new_random() -> Self {
        let mut id = [0u8; 16];
        rand::thread_rng().fill_bytes(&mut id);
        Self(id)
    }

    /// Full 32-character hex representation. For log lines and the
    /// `CorrelationId` round-trip in `BusMessage`.
    #[must_use]
    pub fn hex(&self) -> String {
        hex::encode(self.0)
    }
}

/// The result of attempting to deliver a reply to an in-flight waiter.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReplyDelivery {
    /// The expected peer's reply completed the waiter.
    Delivered,
    /// No waiter exists for the correlation id.
    Unknown,
    /// The supplied reply signer was not the request's target.
    PeerMismatch {
        /// Agent fingerprint the request was sent to.
        expected_peer_fp: Fingerprint,
        /// Verified agent fingerprint that signed the reply.
        actual_peer_fp: Fingerprint,
    },
    /// The request was cancelled while this reply was being delivered.
    ReceiverDropped,
}

struct PendingReply {
    /// Which registration this is, so only its own receiver can remove it.
    token: u64,
    expected_peer_fp: Fingerprint,
    sender: oneshot::Sender<Vec<u8>>,
}

/// Non-consuming responder check used before an inbound reply mutates replay
/// state. Delivery repeats this comparison while removing the waiter, so a
/// concurrent cancellation or replacement cannot bypass the binding.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ReplyPeerCheck {
    Expected,
    Unknown,
    PeerMismatch {
        expected_peer_fp: Fingerprint,
        actual_peer_fp: Fingerprint,
    },
}

/// Holds oneshot senders for in-flight requests, keyed by correlation and
/// bound to the expected responder fingerprint.
///
/// `register` returns a [`PendingReplyRx`]; the bus's inbox later calls
/// `deliver` when the matching reply arrives. However the request ends —
/// answered, timed out, failed to send, or cancelled by dropping its future —
/// dropping the receiver removes the waiter, and a later reply is reported as
/// [`ReplyDelivery::Unknown`].
#[derive(Default)]
pub struct ReplyWaiter {
    waiters: Mutex<HashMap<CorrelationId, PendingReply>>,
    next_token: AtomicU64,
}

impl ReplyWaiter {
    /// Build a fresh waiter with no in-flight requests.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Register a fresh correlation id, atomically binding it to
    /// `expected_peer_fp`, and return the receiver the caller will await.
    /// The waiter lives exactly as long as that receiver.
    pub fn register(&self, id: CorrelationId, expected_peer_fp: Fingerprint) -> PendingReplyRx<'_> {
        let (tx, rx) = oneshot::channel();
        let token = self.next_token.fetch_add(1, Ordering::Relaxed);
        self.waiters.lock().expect("reply waiter poisoned").insert(
            id,
            PendingReply {
                token,
                expected_peer_fp,
                sender: tx,
            },
        );
        PendingReplyRx {
            waiters: self,
            id,
            token,
            rx,
        }
    }

    /// Deliver a reply payload signed by `actual_peer_fp` to the waiter for
    /// `id`.
    ///
    /// The fingerprint comparison and matching removal happen while holding
    /// the same lock. A reply from any other verified signer returns
    /// [`ReplyDelivery::PeerMismatch`] without consuming the waiter, leaving
    /// the expected peer free to complete it later.
    pub fn deliver(
        &self,
        id: CorrelationId,
        actual_peer_fp: Fingerprint,
        payload: Vec<u8>,
    ) -> ReplyDelivery {
        let pending = {
            let mut waiters = self.waiters.lock().expect("reply waiter poisoned");
            let Some(pending) = waiters.get(&id) else {
                return ReplyDelivery::Unknown;
            };
            if pending.expected_peer_fp != actual_peer_fp {
                return ReplyDelivery::PeerMismatch {
                    expected_peer_fp: pending.expected_peer_fp,
                    actual_peer_fp,
                };
            }
            waiters
                .remove(&id)
                .expect("reply waiter disappeared while lock was held")
        };

        if pending.sender.send(payload).is_ok() {
            ReplyDelivery::Delivered
        } else {
            ReplyDelivery::ReceiverDropped
        }
    }

    /// Check whether `actual_peer_fp` is the expected signer for `id` without
    /// consuming or otherwise changing the waiter.
    pub(crate) fn check_peer(
        &self,
        id: CorrelationId,
        actual_peer_fp: Fingerprint,
    ) -> ReplyPeerCheck {
        let waiters = self.waiters.lock().expect("reply waiter poisoned");
        match waiters.get(&id) {
            None => ReplyPeerCheck::Unknown,
            Some(pending) if pending.expected_peer_fp == actual_peer_fp => ReplyPeerCheck::Expected,
            Some(pending) => ReplyPeerCheck::PeerMismatch {
                expected_peer_fp: pending.expected_peer_fp,
                actual_peer_fp,
            },
        }
    }

    /// Remove `id`'s waiter only if it is still registration `token`, so a
    /// replaced receiver cannot remove the registration that replaced it.
    fn cancel(&self, id: &CorrelationId, token: u64) {
        let mut waiters = self.waiters.lock().expect("reply waiter poisoned");
        if waiters
            .get(id)
            .is_some_and(|pending| pending.token == token)
        {
            waiters.remove(id);
        }
    }

    /// Number of in-flight waiters. For tests + diagnostics.
    #[must_use]
    pub fn pending(&self) -> usize {
        self.waiters.lock().expect("reply waiter poisoned").len()
    }
}

/// The receiver for one registered reply: awaits the payload, and removes its
/// waiter when dropped, so no way a request ends can leave one behind (#97).
pub struct PendingReplyRx<'a> {
    waiters: &'a ReplyWaiter,
    id: CorrelationId,
    token: u64,
    rx: oneshot::Receiver<Vec<u8>>,
}

impl Future for PendingReplyRx<'_> {
    type Output = Result<Vec<u8>, oneshot::error::RecvError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        Pin::new(&mut self.rx).poll(cx)
    }
}

impl Drop for PendingReplyRx<'_> {
    fn drop(&mut self) {
        self.waiters.cancel(&self.id, self.token);
    }
}

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

    fn peer(seed: u8) -> Fingerprint {
        Fingerprint([seed; 32])
    }

    #[test]
    fn correlation_ids_are_distinct() {
        let a = CorrelationId::new_random();
        let b = CorrelationId::new_random();
        assert_ne!(a, b, "two random ids must not collide");
        assert_eq!(a.hex().len(), 32);
    }

    #[test]
    fn correlation_id_hex_is_deterministic() {
        let id = CorrelationId([0xab; 16]);
        assert_eq!(id.hex(), "abababababababababababababababab");
    }

    #[tokio::test]
    async fn register_returns_receiver_completed_on_deliver() {
        let w = ReplyWaiter::new();
        let id = CorrelationId::new_random();
        let expected = peer(1);
        let rx = w.register(id, expected);
        assert_eq!(w.pending(), 1);
        assert_eq!(
            w.deliver(id, expected, b"reply".to_vec()),
            ReplyDelivery::Delivered
        );
        assert_eq!(rx.await.unwrap(), b"reply");
        assert_eq!(w.pending(), 0, "deliver removes the waiter");
    }

    #[test]
    fn deliver_unknown_id_is_reported() {
        let w = ReplyWaiter::new();
        assert_eq!(
            w.deliver(CorrelationId::new_random(), peer(1), b"orphan".to_vec()),
            ReplyDelivery::Unknown
        );
    }

    #[tokio::test]
    async fn mismatched_peer_does_not_consume_waiter() {
        let w = ReplyWaiter::new();
        let id = CorrelationId::new_random();
        let expected = peer(1);
        let attacker = peer(2);
        let rx = w.register(id, expected);

        assert_eq!(
            w.deliver(id, attacker, b"forged".to_vec()),
            ReplyDelivery::PeerMismatch {
                expected_peer_fp: expected,
                actual_peer_fp: attacker,
            }
        );
        assert_eq!(w.pending(), 1, "mismatch must leave the waiter pending");

        assert_eq!(
            w.deliver(id, expected, b"legitimate".to_vec()),
            ReplyDelivery::Delivered
        );
        assert_eq!(rx.await.unwrap(), b"legitimate");
        assert_eq!(w.pending(), 0);
    }

    #[tokio::test]
    async fn deliver_twice_only_first_wins() {
        let w = ReplyWaiter::new();
        let id = CorrelationId::new_random();
        let expected = peer(1);
        let rx = w.register(id, expected);
        assert_eq!(
            w.deliver(id, expected, b"first".to_vec()),
            ReplyDelivery::Delivered
        );
        // Second delivery → no waiter left, reports Unknown.
        assert_eq!(
            w.deliver(id, expected, b"second".to_vec()),
            ReplyDelivery::Unknown
        );
        assert_eq!(rx.await.unwrap(), b"first");
    }

    #[tokio::test]
    async fn dropping_a_replaced_receiver_leaves_the_newer_waiter() {
        let w = ReplyWaiter::new();
        let id = CorrelationId::new_random();
        let expected = peer(1);
        let first = w.register(id, expected);
        let second = w.register(id, expected);
        drop(first);
        assert_eq!(w.pending(), 1, "the newer registration survives");
        assert_eq!(
            w.deliver(id, expected, b"reply".to_vec()),
            ReplyDelivery::Delivered
        );
        assert_eq!(second.await.unwrap(), b"reply");
    }

    #[test]
    fn dropping_the_receiver_removes_the_waiter() {
        let w = ReplyWaiter::new();
        let id = CorrelationId::new_random();
        let expected = peer(1);
        drop(w.register(id, expected));
        assert_eq!(w.pending(), 0);
        assert_eq!(
            w.deliver(id, expected, b"late".to_vec()),
            ReplyDelivery::Unknown,
            "a late reply finds no waiter"
        );
    }
}