x0x 0.19.47

//! Direct agent-to-agent messaging.
//!
//! This module provides point-to-point communication between agents,
//! bypassing the gossip layer for private, efficient, reliable delivery.
//!
//! ## Overview
//!
//! While gossip pub/sub is great for broadcast and eventually-consistent
//! data sharing, many use cases require direct communication:
//!
//! - Private messages between two agents
//! - Request/response patterns
//! - Large file transfers
//! - Real-time coordination
//!
//! ## Wire Format
//!
//! Direct messages use stream type byte `0x10` to distinguish from gossip:
//!
//! ```text
//! [0x10][sender_agent_id: 32 bytes][payload: N bytes]
//! ```
//!
//! The sender's AgentId is included in the message so the receiver can
//! identify who sent it, even if multiple agents share a machine.
//!
//! ## Security Model
//!
//! **Sender identity verification.** Each [`DirectMessage`] carries a `verified`
//! field that indicates whether the claimed `sender` AgentId was cross-referenced
//! against the identity discovery cache (which contains signed identity
//! announcements). When `verified` is `true`, the AgentId→MachineId binding
//! was confirmed. When `false`, the AgentId is self-asserted only.
//!
//! The underlying QUIC connection is always authenticated by the sender's
//! [`MachineId`](crate::identity::MachineId) via ML-DSA-65 signatures.
//!
//! **Trust annotations.** Each message also carries a `trust_decision` field
//! from [`TrustEvaluator`](crate::trust::TrustEvaluator), reflecting the
//! full trust evaluation including contact store trust level and machine
//! pinning. Messages are never dropped — applications inspect these fields
//! to decide how to handle each message.
//!
//! ## Example
//!
//! ```rust,ignore
//! use x0x::{Agent, DirectMessage};
//!
//! // Agent A sends to Agent B
//! let outcome = agent_a.connect_to_agent(&agent_b_id).await?;
//! agent_a.send_direct(&agent_b_id, b"hello".to_vec()).await?;
//!
//! // Agent B receives
//! let msg = agent_b.recv_direct().await?;
//! assert_eq!(msg.sender, agent_a.agent_id());
//! assert_eq!(msg.payload, b"hello");
//! ```

use crate::dm::DmPath;
use crate::error::{NetworkError, NetworkResult};
use crate::identity::{AgentId, MachineId};
use crate::trust::TrustDecision;
use serde::Serialize;
use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{SystemTime, UNIX_EPOCH};
use tokio::sync::{broadcast, mpsc, Notify, RwLock};

/// Stream type byte for direct messages (distinct from gossip: 0, 1, 2).
pub const DIRECT_MESSAGE_STREAM_TYPE: u8 = 0x10;

/// Maximum payload size for direct messages (16 MB).
pub const MAX_DIRECT_PAYLOAD_SIZE: usize = 16 * 1024 * 1024;

/// Per-subscriber direct-message buffer depth.
///
/// Each `subscribe_direct()` caller gets an independent queue of this size so
/// one slow SSE/WebSocket/file-transfer consumer cannot force drops for every
/// other consumer.
const DIRECT_SUBSCRIBER_BUFFER: usize = 8192;

/// Keep direct diagnostics for recently active disconnected peers for at most
/// this long. Connected peers are always retained.
const DIRECT_DIAGNOSTICS_IDLE_TTL_MS: u64 = 24 * 60 * 60 * 1000;

/// Minimum retained direct peer/lifecycle diagnostics entries before old idle
/// records are evicted. The effective cap also scales with connected peers.
const DIRECT_DIAGNOSTICS_MIN_RETAIN: usize = 1024;

/// X0X-0041: capacity of the prefer-newest-connection broadcast channel.
///
/// Sized for bursty supersede churn (multiple peers replacing in tight
/// succession). Slow subscribers may observe `RecvError::Lagged`; callers
/// reconcile against the lifecycle table by calling
/// [`DirectMessaging::current_generation`] after a lag.
const LIFECYCLE_REPLACED_BROADCAST_CAPACITY: usize = 256;

#[doc(hidden)]
pub struct RawQuicAckRaceTestHook {
    first_attempt_started: Notify,
    first_attempt_result_release: Notify,
    replaced_short_circuit: Notify,
}

impl std::fmt::Debug for RawQuicAckRaceTestHook {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("RawQuicAckRaceTestHook")
            .finish_non_exhaustive()
    }
}

impl Default for RawQuicAckRaceTestHook {
    fn default() -> Self {
        Self {
            first_attempt_started: Notify::new(),
            first_attempt_result_release: Notify::new(),
            replaced_short_circuit: Notify::new(),
        }
    }
}

impl RawQuicAckRaceTestHook {
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    pub async fn wait_first_attempt_started(&self) {
        self.first_attempt_started.notified().await;
    }

    pub fn release_first_attempt_result(&self) {
        self.first_attempt_result_release.notify_one();
    }

    pub async fn wait_replaced_short_circuit(&self) {
        self.replaced_short_circuit.notified().await;
    }

    pub(crate) fn notify_first_attempt_started(&self) {
        self.first_attempt_started.notify_one();
    }

    pub(crate) async fn hold_first_attempt_result(&self) {
        self.first_attempt_result_release.notified().await;
    }

    pub(crate) fn notify_replaced_short_circuit(&self) {
        self.replaced_short_circuit.notify_one();
    }
}

/// A direct message received from another agent.
///
/// # Security Note
///
/// The `sender` field is **self-asserted** by the sender and not cryptographically
/// verified. However, `machine_id` is authenticated via the QUIC connection's
/// ML-DSA-65 handshake, so you can trust which machine sent this message.
///
/// The claimed `sender` AgentId is only as trustworthy as the machine that sent it.
/// For sensitive operations, verify the AgentId→MachineId binding against a
/// trusted source (e.g., a signed identity announcement).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DirectMessage {
    /// The AgentId claimed by the sender.
    ///
    /// **Warning:** This is self-asserted and not cryptographically verified.
    /// Use `machine_id` for authenticated sender identity, or check the
    /// `verified` field which cross-references the identity discovery cache.
    pub sender: AgentId,
    /// The MachineId the message was sent from (authenticated via QUIC).
    ///
    /// This is derived from the QUIC connection's peer identity and is
    /// cryptographically verified via ML-DSA-65 signatures.
    pub machine_id: MachineId,
    /// The message payload.
    pub payload: Vec<u8>,
    /// Unix timestamp (milliseconds) when the message was received.
    pub received_at: u64,
    /// Whether the sender's AgentId was verified against the identity
    /// discovery cache.
    ///
    /// `true` if the cache contains an entry mapping this `sender` AgentId
    /// to this `machine_id`. `false` if the AgentId could not be verified
    /// (self-asserted only — the sender may still be legitimate but hasn't
    /// been seen via a signed identity announcement yet).
    pub verified: bool,
    /// Trust decision from [`TrustEvaluator`](crate::trust::TrustEvaluator)
    /// for the `(sender, machine_id)` pair.
    ///
    /// `None` if the trust system was unavailable at receive time.
    /// When present, reflects the full trust evaluation including contact
    /// store trust level and machine pinning.
    pub trust_decision: Option<TrustDecision>,
}

impl DirectMessage {
    /// Create a new `DirectMessage` with default verification fields.
    ///
    /// `verified` defaults to `false` and `trust_decision` to `None`.
    /// Use [`new_verified`](Self::new_verified) to set these fields.
    #[must_use]
    pub fn new(sender: AgentId, machine_id: MachineId, payload: Vec<u8>) -> Self {
        Self::new_verified(sender, machine_id, payload, false, None)
    }

    /// Create a new `DirectMessage` with explicit verification fields.
    #[must_use]
    pub fn new_verified(
        sender: AgentId,
        machine_id: MachineId,
        payload: Vec<u8>,
        verified: bool,
        trust_decision: Option<TrustDecision>,
    ) -> Self {
        let received_at = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_millis() as u64)
            .unwrap_or(0);

        Self {
            sender,
            machine_id,
            payload,
            received_at,
            verified,
            trust_decision,
        }
    }

    /// Returns the payload as a UTF-8 string if valid.
    #[must_use]
    pub fn payload_str(&self) -> Option<&str> {
        std::str::from_utf8(&self.payload).ok()
    }
}

fn now_unix_ms_lossy() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_millis() as u64)
        .unwrap_or(0)
}

fn direct_diagnostics_retain_limit(connected_len: usize) -> usize {
    DIRECT_DIAGNOSTICS_MIN_RETAIN.max(connected_len.saturating_mul(2))
}

fn dm_path_label(path: DmPath) -> &'static str {
    match path {
        DmPath::Loopback => "loopback",
        DmPath::GossipInbox => "gossip_inbox",
        DmPath::RawQuic => "raw_quic",
        DmPath::RawQuicAcked => "raw_quic_acked",
    }
}

/// Compute a stable, short content digest for `dm.trace` correlation.
///
/// Returns the first 16 hex characters (64 bits) of a BLAKE3 hash of the
/// supplied bytes. Used to correlate sender-side and receiver-side
/// `dm.trace` lines for a single message regardless of whether it travelled
/// the raw-QUIC or gossip-inbox path. The full payload is always available
/// to both ends — the gossip path decrypts before fan-out — so the same
/// input produces the same digest on both sides.
#[must_use]
pub fn dm_payload_digest_hex(bytes: &[u8]) -> String {
    let hash = blake3::hash(bytes);
    let hex = hex::encode(hash.as_bytes());
    hex[..16].to_string()
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum DirectSubscriberPush {
    Delivered,
    DeliveredWithEviction,
    Closed,
}

#[derive(Debug)]
struct DirectSubscriberQueue {
    queue: Mutex<VecDeque<DirectMessage>>,
    notify: Notify,
    closed: AtomicBool,
    capacity: usize,
}

impl DirectSubscriberQueue {
    fn new(capacity: usize) -> Self {
        Self {
            queue: Mutex::new(VecDeque::with_capacity(capacity)),
            notify: Notify::new(),
            closed: AtomicBool::new(false),
            capacity,
        }
    }

    fn push_drop_oldest(&self, msg: DirectMessage) -> DirectSubscriberPush {
        if self.closed.load(Ordering::Relaxed) {
            return DirectSubscriberPush::Closed;
        }
        let mut queue = match self.queue.lock() {
            Ok(queue) => queue,
            Err(e) => {
                tracing::error!("direct subscriber queue poisoned: {e}");
                return DirectSubscriberPush::Closed;
            }
        };
        if self.closed.load(Ordering::Relaxed) {
            return DirectSubscriberPush::Closed;
        }
        let evicted = if queue.len() >= self.capacity {
            queue.pop_front();
            true
        } else {
            false
        };
        queue.push_back(msg);
        drop(queue);
        self.notify.notify_one();
        if evicted {
            DirectSubscriberPush::DeliveredWithEviction
        } else {
            DirectSubscriberPush::Delivered
        }
    }

    fn pop_front(&self) -> Option<DirectMessage> {
        match self.queue.lock() {
            Ok(mut queue) => queue.pop_front(),
            Err(e) => {
                tracing::error!("direct subscriber queue poisoned: {e}");
                None
            }
        }
    }

    fn is_closed(&self) -> bool {
        self.closed.load(Ordering::Relaxed)
    }

    fn close(&self) {
        self.closed.store(true, Ordering::Relaxed);
        self.notify.notify_waiters();
    }
}

/// Receiver for direct messages.
///
/// Each receiver owns an independent bounded queue. Cloning a receiver
/// creates a fresh subscription rather than sharing cursor state, preserving
/// the old broadcast-style "every subscriber sees every future message"
/// semantics without tokio broadcast's global lag/drop behaviour. A slow
/// receiver keeps its stream open, but its oldest buffered direct events are
/// evicted once the queue reaches capacity.
#[derive(Debug)]
pub struct DirectMessageReceiver {
    id: Option<u64>,
    queue: Arc<DirectSubscriberQueue>,
    subscribers: Arc<Mutex<HashMap<u64, Arc<DirectSubscriberQueue>>>>,
    next_subscriber_id: Arc<AtomicU64>,
    capacity: usize,
}

impl DirectMessageReceiver {
    /// Create and register a new receiver in the shared subscriber registry.
    fn new(
        subscribers: Arc<Mutex<HashMap<u64, Arc<DirectSubscriberQueue>>>>,
        next_subscriber_id: Arc<AtomicU64>,
        capacity: usize,
    ) -> Self {
        let queue = Arc::new(DirectSubscriberQueue::new(capacity));
        let id = next_subscriber_id.fetch_add(1, Ordering::Relaxed);
        let registered = match subscribers.lock() {
            Ok(mut guard) => {
                guard.insert(id, Arc::clone(&queue));
                Some(id)
            }
            Err(e) => {
                tracing::error!("direct subscriber registry poisoned: {e}");
                queue.close();
                None
            }
        };

        Self {
            id: registered,
            queue,
            subscribers,
            next_subscriber_id,
            capacity,
        }
    }

    /// Receive the next direct message.
    ///
    /// Returns `None` if this subscriber was dropped because it fell behind,
    /// the daemon is shutting down, or the channel closed.
    pub async fn recv(&mut self) -> Option<DirectMessage> {
        loop {
            let notified = self.queue.notify.notified();
            if let Some(msg) = self.queue.pop_front() {
                return Some(msg);
            }
            if self.queue.is_closed() {
                return None;
            }
            notified.await;
        }
    }

    /// Try to receive a message without blocking.
    ///
    /// Returns `None` if no message is available or channel is closed.
    pub fn try_recv(&mut self) -> Option<DirectMessage> {
        self.queue.pop_front()
    }
}

impl Clone for DirectMessageReceiver {
    fn clone(&self) -> Self {
        Self::new(
            Arc::clone(&self.subscribers),
            Arc::clone(&self.next_subscriber_id),
            self.capacity,
        )
    }
}

impl Drop for DirectMessageReceiver {
    fn drop(&mut self) {
        let Some(id) = self.id.take() else {
            return;
        };
        match self.subscribers.lock() {
            Ok(mut guard) => {
                guard.remove(&id);
            }
            Err(e) => tracing::error!("direct subscriber registry poisoned on drop: {e}"),
        }
        self.queue.close();
    }
}

#[derive(Debug, Default)]
struct DirectDiagnosticsCounters {
    outgoing_send_total: AtomicU64,
    outgoing_send_succeeded: AtomicU64,
    outgoing_send_failed: AtomicU64,
    outgoing_path_loopback: AtomicU64,
    outgoing_path_raw_quic: AtomicU64,
    outgoing_path_gossip_inbox: AtomicU64,
    incoming_envelopes_total: AtomicU64,
    incoming_decode_failed: AtomicU64,
    incoming_signature_failed: AtomicU64,
    incoming_trust_rejected: AtomicU64,
    incoming_delivered_to_subscribe: AtomicU64,
    subscriber_channel_lagged: AtomicU64,
    subscriber_events_evicted: AtomicU64,
    subscriber_channel_closed: AtomicU64,
}

#[derive(Debug, Clone, Default)]
struct DirectPeerDiagnosticsState {
    avg_rtt_ms: Option<u32>,
    last_send_at_ms: Option<u64>,
    last_recv_at_ms: Option<u64>,
    send_succeeded: u64,
    send_failed: u64,
    recv_count: u64,
    preferred_path: Option<&'static str>,
}

impl DirectPeerDiagnosticsState {
    fn last_activity_ms(&self) -> Option<u64> {
        self.last_send_at_ms.max(self.last_recv_at_ms)
    }
}

#[derive(Debug, Clone, Default)]
struct DirectLifecycleState {
    generation: Option<u64>,
    blocked_reason: Option<String>,
    last_updated_at_ms: Option<u64>,
}

/// Global direct-message diagnostics exposed by `/diagnostics/dm`.
#[derive(Debug, Clone, Default, Serialize)]
pub struct DmDiagnosticsStats {
    pub outgoing_send_total: u64,
    pub outgoing_send_succeeded: u64,
    pub outgoing_send_failed: u64,
    pub outgoing_path_loopback: u64,
    pub outgoing_path_raw_quic: u64,
    pub outgoing_path_gossip_inbox: u64,
    pub incoming_envelopes_total: u64,
    pub incoming_decode_failed: u64,
    pub incoming_signature_failed: u64,
    pub incoming_trust_rejected: u64,
    pub incoming_delivered_to_subscribe: u64,
    /// Number of oldest buffered events evicted from slow subscriber queues.
    pub subscriber_events_evicted: u64,
    /// Backward-compatible alias for slow-subscriber pressure events.
    pub subscriber_channel_lagged: u64,
    pub subscriber_channel_closed: u64,
}

/// Per-peer direct-message diagnostics exposed by `/diagnostics/dm`.
#[derive(Debug, Clone, Default, Serialize)]
pub struct DmPeerDiagnostics {
    pub avg_rtt_ms: Option<u32>,
    pub last_send_ms_ago: Option<u64>,
    pub last_recv_ms_ago: Option<u64>,
    pub send_succeeded: u64,
    pub send_failed: u64,
    pub recv_count: u64,
    pub preferred_path: String,
}

/// Snapshot of the direct-message diagnostics surface.
#[derive(Debug, Clone, Default, Serialize)]
pub struct DmDiagnosticsSnapshot {
    pub stats: DmDiagnosticsStats,
    pub per_peer: BTreeMap<String, DmPeerDiagnostics>,
    pub subscriber_count: usize,
    pub subscriber_capacity: usize,
}

/// Tracks connections and mappings for direct messaging.
///
/// This maintains the MachineId → AgentId reverse lookup needed to
/// identify message senders, since ant-quic only knows about MachineIds.
#[derive(Debug)]
pub struct DirectMessaging {
    /// Reverse mapping: MachineId → AgentId.
    /// Built from discovered agents.
    machine_to_agent: Arc<RwLock<HashMap<MachineId, AgentId>>>,

    /// Currently connected agents (AgentId → MachineId).
    connected_agents: Arc<RwLock<HashMap<AgentId, MachineId>>>,

    /// Per-subscriber queues for received direct messages.
    subscribers: Arc<Mutex<HashMap<u64, Arc<DirectSubscriberQueue>>>>,

    /// Monotonic id source for subscriber queues.
    next_subscriber_id: Arc<AtomicU64>,

    /// Queue capacity assigned to each subscriber.
    subscriber_capacity: usize,

    /// Global direct-message diagnostics counters.
    diagnostics: Arc<DirectDiagnosticsCounters>,

    /// Per-peer direct-message diagnostics state.
    peer_diagnostics: Arc<Mutex<HashMap<AgentId, DirectPeerDiagnosticsState>>>,

    /// Hot peer lifecycle table keyed by MachineId.
    lifecycle: Arc<Mutex<HashMap<MachineId, DirectLifecycleState>>>,

    /// X0X-0041: prefer-newest-connection broadcast.
    ///
    /// Fires whenever ant-quic emits a `Replaced` lifecycle event so DM retry
    /// loops can short-circuit the current attempt and target the new
    /// generation immediately. The payload is `(machine_id, new_generation)`.
    /// Late or absent subscribers do not block the producer.
    lifecycle_replaced_tx: broadcast::Sender<(MachineId, u64)>,

    /// Test hook for deterministic raw-QUIC ACK/Replaced race coverage.
    raw_quic_ack_race_test_hook: Arc<Mutex<Option<Arc<RawQuicAckRaceTestHook>>>>,

    /// Internal sender for the receiver task.
    internal_tx: mpsc::Sender<DirectMessage>,

    /// Internal receiver (owned by the processing task).
    internal_rx: Arc<tokio::sync::Mutex<mpsc::Receiver<DirectMessage>>>,
}

impl DirectMessaging {
    /// Create a new DirectMessaging instance.
    #[must_use]
    pub fn new() -> Self {
        Self::with_subscriber_capacity(DIRECT_SUBSCRIBER_BUFFER)
    }

    fn with_subscriber_capacity(subscriber_capacity: usize) -> Self {
        let subscriber_capacity = subscriber_capacity.max(1);
        // Each subscribe_direct() caller gets an independent queue. This
        // preserves fan-out semantics without tokio::sync::broadcast's global
        // lag/drop behaviour. If a subscriber fills its own queue, the oldest
        // buffered event in that subscriber queue is evicted and counted.
        let (internal_tx, internal_rx) = mpsc::channel(subscriber_capacity);
        // X0X-0041: prefer-newest-connection broadcast. Capacity sized for
        // bursty supersede churn during reconnect storms; lagging subscribers
        // recover via lifecycle table reread.
        let (lifecycle_replaced_tx, _) = broadcast::channel(LIFECYCLE_REPLACED_BROADCAST_CAPACITY);

        Self {
            machine_to_agent: Arc::new(RwLock::new(HashMap::new())),
            connected_agents: Arc::new(RwLock::new(HashMap::new())),
            subscribers: Arc::new(Mutex::new(HashMap::new())),
            next_subscriber_id: Arc::new(AtomicU64::new(1)),
            subscriber_capacity,
            diagnostics: Arc::new(DirectDiagnosticsCounters::default()),
            peer_diagnostics: Arc::new(Mutex::new(HashMap::new())),
            lifecycle: Arc::new(Mutex::new(HashMap::new())),
            lifecycle_replaced_tx,
            raw_quic_ack_race_test_hook: Arc::new(Mutex::new(None)),
            internal_tx,
            internal_rx: Arc::new(tokio::sync::Mutex::new(internal_rx)),
        }
    }

    /// Register a mapping from MachineId to AgentId.
    ///
    /// Called when an agent is discovered or connected.
    pub async fn register_agent(&self, agent_id: AgentId, machine_id: MachineId) {
        let mut map = self.machine_to_agent.write().await;
        map.insert(machine_id, agent_id);
        tracing::debug!(
            "Registered agent mapping: {:?} -> {:?}",
            machine_id,
            agent_id
        );
    }

    /// Look up AgentId from MachineId.
    pub async fn lookup_agent(&self, machine_id: &MachineId) -> Option<AgentId> {
        let map = self.machine_to_agent.read().await;
        map.get(machine_id).copied()
    }

    /// Mark an agent as connected.
    pub async fn mark_connected(&self, agent_id: AgentId, machine_id: MachineId) {
        // Ensure we have the mapping
        self.register_agent(agent_id, machine_id).await;

        let mut connected = self.connected_agents.write().await;
        connected.insert(agent_id, machine_id);
        self.record_lifecycle_established(machine_id, None);
        tracing::info!("Agent connected: {:?}", agent_id);
    }

    /// Mark an agent as disconnected.
    pub async fn mark_disconnected(&self, agent_id: &AgentId) {
        let mut connected = self.connected_agents.write().await;
        connected.remove(agent_id);
        // NetworkEvent::PeerDisconnected carries no ant-quic lifecycle
        // generation. A delayed disconnect for a superseded old connection can
        // therefore arrive after a newer Established/Replaced event. Do not
        // write a lifecycle block here; generation-bearing Closed events are
        // the authoritative source for the send fast-fail table.
        tracing::info!("Agent disconnected: {:?}", agent_id);
    }

    /// Record an established lifecycle generation for a machine.
    pub fn record_lifecycle_established(&self, machine_id: MachineId, generation: Option<u64>) {
        self.update_lifecycle(machine_id, |state| {
            if let Some(generation) = generation {
                state.generation = Some(generation);
            }
            state.blocked_reason = None;
        });
    }

    /// Record that a newer generation replaced the old one.
    ///
    /// X0X-0041: this also broadcasts on the prefer-newest channel so DM retry
    /// loops mid-attempt can short-circuit and target the new generation.
    pub fn record_lifecycle_replaced(&self, machine_id: MachineId, new_generation: u64) {
        self.update_lifecycle(machine_id, |state| {
            state.generation = Some(new_generation);
            state.blocked_reason = None;
        });
        // Best-effort broadcast — slow / absent subscribers do not block the
        // producer. `send` returns Err only when there are zero receivers,
        // which is normal during steady state.
        let _ = self
            .lifecycle_replaced_tx
            .send((machine_id, new_generation));
    }

    /// X0X-0041: return the current active lifecycle generation for a peer,
    /// if known. The lifecycle table is updated from ant-quic
    /// `Established`/`Replaced` events; raw-DM uses this as a hint to detect
    /// connection supersede mid-send.
    #[must_use]
    pub fn current_generation(&self, machine_id: &MachineId) -> Option<u64> {
        match self.lifecycle.lock() {
            Ok(guard) => guard.get(machine_id).and_then(|state| state.generation),
            Err(e) => {
                tracing::error!("direct lifecycle registry poisoned: {e}");
                None
            }
        }
    }

    /// X0X-0041: subscribe to prefer-newest-connection events. The broadcast
    /// payload is `(machine_id, new_generation)` and fires synchronously with
    /// every [`Self::record_lifecycle_replaced`] call.
    #[must_use]
    pub fn subscribe_lifecycle_replaced(&self) -> broadcast::Receiver<(MachineId, u64)> {
        self.lifecycle_replaced_tx.subscribe()
    }

    #[doc(hidden)]
    pub fn set_raw_quic_ack_race_test_hook_for_testing(
        &self,
        hook: Option<Arc<RawQuicAckRaceTestHook>>,
    ) {
        match self.raw_quic_ack_race_test_hook.lock() {
            Ok(mut guard) => *guard = hook,
            Err(e) => tracing::error!("raw QUIC ACK race test hook poisoned: {e}"),
        }
    }

    pub(crate) fn raw_quic_ack_race_test_hook(&self) -> Option<Arc<RawQuicAckRaceTestHook>> {
        match self.raw_quic_ack_race_test_hook.lock() {
            Ok(guard) => guard.clone(),
            Err(e) => {
                tracing::error!("raw QUIC ACK race test hook poisoned: {e}");
                None
            }
        }
    }

    /// Record a closing/closed lifecycle state for a machine.
    pub fn record_lifecycle_blocked(
        &self,
        machine_id: MachineId,
        generation: Option<u64>,
        reason: impl Into<String>,
    ) {
        let reason = reason.into();
        self.update_lifecycle(machine_id, |state| {
            if let Some(generation) = generation {
                match state.generation {
                    Some(current) if current != generation => return,
                    Some(_) => {}
                    None => state.generation = Some(generation),
                }
            }
            state.blocked_reason = Some(reason);
        });
    }

    /// Returns the current lifecycle block reason for a machine, if any.
    #[must_use]
    pub fn lifecycle_block_reason(&self, machine_id: &MachineId) -> Option<String> {
        match self.lifecycle.lock() {
            Ok(guard) => guard
                .get(machine_id)
                .and_then(|state| state.blocked_reason.clone()),
            Err(e) => {
                tracing::error!("direct lifecycle registry poisoned: {e}");
                None
            }
        }
    }

    /// Check if an agent is currently connected.
    pub async fn is_connected(&self, agent_id: &AgentId) -> bool {
        let connected = self.connected_agents.read().await;
        connected.contains_key(agent_id)
    }

    /// Get the MachineId for a connected agent.
    pub async fn get_machine_id(&self, agent_id: &AgentId) -> Option<MachineId> {
        let connected = self.connected_agents.read().await;
        connected.get(agent_id).copied()
    }

    /// Get all currently connected agents.
    pub async fn connected_agents(&self) -> Vec<AgentId> {
        let connected = self.connected_agents.read().await;
        connected.keys().copied().collect()
    }

    /// Get a receiver for direct messages.
    pub fn subscribe(&self) -> DirectMessageReceiver {
        DirectMessageReceiver::new(
            Arc::clone(&self.subscribers),
            Arc::clone(&self.next_subscriber_id),
            self.subscriber_capacity,
        )
    }

    /// Current number of live direct-message subscribers.
    ///
    /// Used by diagnostics to distinguish "message dispatched to N SSE/WS
    /// consumers" from "no one is listening".
    pub fn subscriber_count(&self) -> usize {
        match self.subscribers.lock() {
            Ok(guard) => guard.len(),
            Err(e) => {
                tracing::error!("direct subscriber registry poisoned: {e}");
                0
            }
        }
    }

    /// Record that a logical outgoing DM was accepted for sending.
    pub(crate) fn record_outgoing_started(&self, agent_id: AgentId, avg_rtt_ms: Option<u32>) {
        self.diagnostics
            .outgoing_send_total
            .fetch_add(1, Ordering::Relaxed);
        let now_ms = now_unix_ms_lossy();
        self.with_peer_diagnostics(agent_id, |peer| {
            peer.last_send_at_ms = Some(now_ms);
            if let Some(rtt) = avg_rtt_ms.filter(|rtt| *rtt > 0) {
                peer.avg_rtt_ms = Some(rtt);
            }
        });
    }

    /// Record a successful logical outgoing DM.
    pub(crate) fn record_outgoing_succeeded(&self, agent_id: AgentId, path: DmPath) {
        self.diagnostics
            .outgoing_send_succeeded
            .fetch_add(1, Ordering::Relaxed);
        match path {
            DmPath::Loopback => {
                self.diagnostics
                    .outgoing_path_loopback
                    .fetch_add(1, Ordering::Relaxed);
            }
            DmPath::RawQuic | DmPath::RawQuicAcked => {
                self.diagnostics
                    .outgoing_path_raw_quic
                    .fetch_add(1, Ordering::Relaxed);
            }
            DmPath::GossipInbox => {
                self.diagnostics
                    .outgoing_path_gossip_inbox
                    .fetch_add(1, Ordering::Relaxed);
            }
        }
        let path_label = dm_path_label(path);
        self.with_peer_diagnostics(agent_id, |peer| {
            peer.send_succeeded = peer.send_succeeded.saturating_add(1);
            peer.preferred_path = Some(path_label);
        });
    }

    /// Record a failed logical outgoing DM.
    pub(crate) fn record_outgoing_failed(&self, agent_id: AgentId) {
        self.diagnostics
            .outgoing_send_failed
            .fetch_add(1, Ordering::Relaxed);
        self.with_peer_diagnostics(agent_id, |peer| {
            peer.send_failed = peer.send_failed.saturating_add(1);
        });
    }

    /// Record a DM inbox decode failure.
    pub(crate) fn record_incoming_decode_failed(&self) {
        self.diagnostics
            .incoming_decode_failed
            .fetch_add(1, Ordering::Relaxed);
    }

    /// Record a DM inbox signature failure.
    pub(crate) fn record_incoming_signature_failed(&self) {
        self.diagnostics
            .incoming_signature_failed
            .fetch_add(1, Ordering::Relaxed);
    }

    /// Record a DM trust-policy rejection.
    pub(crate) fn record_incoming_trust_rejected(&self, agent_id: AgentId) {
        self.diagnostics
            .incoming_trust_rejected
            .fetch_add(1, Ordering::Relaxed);
        self.with_peer_diagnostics(agent_id, |_| {});
    }

    /// Snapshot direct-message diagnostics for API surfaces.
    #[must_use]
    pub fn diagnostics_snapshot(&self) -> DmDiagnosticsSnapshot {
        let stats = DmDiagnosticsStats {
            outgoing_send_total: self.diagnostics.outgoing_send_total.load(Ordering::Relaxed),
            outgoing_send_succeeded: self
                .diagnostics
                .outgoing_send_succeeded
                .load(Ordering::Relaxed),
            outgoing_send_failed: self
                .diagnostics
                .outgoing_send_failed
                .load(Ordering::Relaxed),
            outgoing_path_loopback: self
                .diagnostics
                .outgoing_path_loopback
                .load(Ordering::Relaxed),
            outgoing_path_raw_quic: self
                .diagnostics
                .outgoing_path_raw_quic
                .load(Ordering::Relaxed),
            outgoing_path_gossip_inbox: self
                .diagnostics
                .outgoing_path_gossip_inbox
                .load(Ordering::Relaxed),
            incoming_envelopes_total: self
                .diagnostics
                .incoming_envelopes_total
                .load(Ordering::Relaxed),
            incoming_decode_failed: self
                .diagnostics
                .incoming_decode_failed
                .load(Ordering::Relaxed),
            incoming_signature_failed: self
                .diagnostics
                .incoming_signature_failed
                .load(Ordering::Relaxed),
            incoming_trust_rejected: self
                .diagnostics
                .incoming_trust_rejected
                .load(Ordering::Relaxed),
            incoming_delivered_to_subscribe: self
                .diagnostics
                .incoming_delivered_to_subscribe
                .load(Ordering::Relaxed),
            subscriber_channel_lagged: self
                .diagnostics
                .subscriber_channel_lagged
                .load(Ordering::Relaxed),
            subscriber_events_evicted: self
                .diagnostics
                .subscriber_events_evicted
                .load(Ordering::Relaxed),
            subscriber_channel_closed: self
                .diagnostics
                .subscriber_channel_closed
                .load(Ordering::Relaxed),
        };

        let now_ms = now_unix_ms_lossy();
        let per_peer = match self.peer_diagnostics.lock() {
            Ok(guard) => guard
                .iter()
                .map(|(agent_id, peer)| {
                    (
                        hex::encode(agent_id.as_bytes()),
                        DmPeerDiagnostics {
                            avg_rtt_ms: peer.avg_rtt_ms,
                            last_send_ms_ago: peer
                                .last_send_at_ms
                                .map(|ts| now_ms.saturating_sub(ts)),
                            last_recv_ms_ago: peer
                                .last_recv_at_ms
                                .map(|ts| now_ms.saturating_sub(ts)),
                            send_succeeded: peer.send_succeeded,
                            send_failed: peer.send_failed,
                            recv_count: peer.recv_count,
                            preferred_path: peer.preferred_path.unwrap_or("unknown").to_string(),
                        },
                    )
                })
                .collect(),
            Err(e) => {
                tracing::error!("direct peer diagnostics registry poisoned: {e}");
                BTreeMap::new()
            }
        };

        DmDiagnosticsSnapshot {
            stats,
            per_peer,
            subscriber_count: self.subscriber_count(),
            subscriber_capacity: self.subscriber_capacity,
        }
    }

    /// Enqueue a local loopback direct message through the normal delivery path.
    ///
    /// This is used when an agent sends a DM to its own [`AgentId`]. It avoids
    /// creating an impossible QUIC self-connection while preserving the same
    /// subscriber and pull-API fan-out semantics as a remote direct message.
    pub(crate) async fn handle_loopback(
        &self,
        machine_id: MachineId,
        agent_id: AgentId,
        payload: Vec<u8>,
    ) -> u64 {
        self.handle_incoming(
            machine_id,
            agent_id,
            payload,
            true,
            Some(TrustDecision::Accept),
        )
        .await
    }

    /// Process an incoming direct message from the network.
    ///
    /// Called by the network layer when a direct message is received.
    /// The `verified` and `trust_decision` fields are populated by the
    /// caller based on the identity discovery cache and contact store.
    ///
    /// Returns the number of subscribers that successfully received the
    /// message. Slow subscribers keep their streams open, but when a queue
    /// is full the oldest event in that subscriber's queue is evicted and
    /// counted in [`DmDiagnosticsStats`].
    pub async fn handle_incoming(
        &self,
        machine_id: MachineId,
        sender_agent_id: AgentId,
        payload: Vec<u8>,
        verified: bool,
        trust_decision: Option<TrustDecision>,
    ) -> u64 {
        self.diagnostics
            .incoming_envelopes_total
            .fetch_add(1, Ordering::Relaxed);
        let now_ms = now_unix_ms_lossy();
        self.with_peer_diagnostics(sender_agent_id, |peer| {
            peer.last_recv_at_ms = Some(now_ms);
            peer.recv_count = peer.recv_count.saturating_add(1);
        });

        let msg = DirectMessage::new_verified(
            sender_agent_id,
            machine_id,
            payload,
            verified,
            trust_decision,
        );

        let subscribers = self.subscriber_snapshot();
        let mut delivered = 0_u64;
        let mut remove_ids = Vec::new();
        for (id, queue) in subscribers {
            match queue.push_drop_oldest(msg.clone()) {
                DirectSubscriberPush::Delivered => {
                    delivered = delivered.saturating_add(1);
                }
                DirectSubscriberPush::DeliveredWithEviction => {
                    self.diagnostics
                        .subscriber_channel_lagged
                        .fetch_add(1, Ordering::Relaxed);
                    self.diagnostics
                        .subscriber_events_evicted
                        .fetch_add(1, Ordering::Relaxed);
                    tracing::warn!(
                        subscriber_id = id,
                        capacity = self.subscriber_capacity,
                        "direct subscriber queue full; evicted oldest buffered event"
                    );
                    delivered = delivered.saturating_add(1);
                }
                DirectSubscriberPush::Closed => {
                    self.diagnostics
                        .subscriber_channel_closed
                        .fetch_add(1, Ordering::Relaxed);
                    remove_ids.push(id);
                }
            }
        }
        if delivered > 0 {
            self.diagnostics
                .incoming_delivered_to_subscribe
                .fetch_add(1, Ordering::Relaxed);
        }
        if !remove_ids.is_empty() {
            self.remove_subscribers(&remove_ids);
        }

        // Also enqueue on the internal pull-API channel (consumed by
        // `recv_direct()`). This is a best-effort, non-blocking enqueue: the
        // mpsc receiver is typically idle in long-running daemons that only
        // use `subscribe_direct()` for SSE/WS fan-out. If we awaited a
        // bounded `send` here, a cold `internal_rx` would back-pressure this
        // task, which in turn stalls `start_direct_listener` →
        // `NetworkNode::spawn_receiver` → `Node::recv` and causes ant-quic
        // reader tasks to queue up on their forward channel. The per-subscriber
        // queues above are the authoritative delivery surface for daemons; the
        // internal channel is a convenience for library users that keep calling
        // `recv_direct()`.
        if self.internal_tx.try_send(msg).is_err() {
            tracing::trace!("direct internal_tx full or closed, skipping pull-API copy");
        }

        delivered
    }

    fn update_lifecycle(
        &self,
        machine_id: MachineId,
        update: impl FnOnce(&mut DirectLifecycleState),
    ) {
        match self.lifecycle.lock() {
            Ok(mut guard) => {
                let state = guard.entry(machine_id).or_default();
                state.last_updated_at_ms = Some(now_unix_ms_lossy());
                update(state);
                if guard.len() > DIRECT_DIAGNOSTICS_MIN_RETAIN {
                    if let Some(connected) = self.connected_machine_snapshot() {
                        Self::prune_lifecycle_locked(&mut guard, &connected);
                    }
                }
            }
            Err(e) => tracing::error!("direct lifecycle registry poisoned: {e}"),
        }
    }

    fn with_peer_diagnostics(
        &self,
        agent_id: AgentId,
        update: impl FnOnce(&mut DirectPeerDiagnosticsState),
    ) {
        match self.peer_diagnostics.lock() {
            Ok(mut guard) => {
                let peer = guard.entry(agent_id).or_default();
                update(peer);
                if guard.len() > DIRECT_DIAGNOSTICS_MIN_RETAIN {
                    if let Some(connected) = self.connected_agent_snapshot() {
                        Self::prune_peer_diagnostics_locked(&mut guard, &connected);
                    }
                }
            }
            Err(e) => tracing::error!("direct peer diagnostics registry poisoned: {e}"),
        }
    }

    fn connected_agent_snapshot(&self) -> Option<HashSet<AgentId>> {
        match self.connected_agents.try_read() {
            Ok(guard) => Some(guard.keys().copied().collect()),
            Err(_) => None,
        }
    }

    fn connected_machine_snapshot(&self) -> Option<HashSet<MachineId>> {
        match self.connected_agents.try_read() {
            Ok(guard) => Some(guard.values().copied().collect()),
            Err(_) => None,
        }
    }

    fn prune_peer_diagnostics_locked(
        guard: &mut HashMap<AgentId, DirectPeerDiagnosticsState>,
        connected: &HashSet<AgentId>,
    ) {
        let limit = direct_diagnostics_retain_limit(connected.len());
        if guard.len() <= limit {
            return;
        }

        let now = now_unix_ms_lossy();
        guard.retain(|agent_id, state| {
            connected.contains(agent_id)
                || state
                    .last_activity_ms()
                    .is_some_and(|last| now.saturating_sub(last) <= DIRECT_DIAGNOSTICS_IDLE_TTL_MS)
        });

        if guard.len() <= limit {
            return;
        }
        let mut idle: Vec<(AgentId, u64)> = guard
            .iter()
            .filter(|(agent_id, _)| !connected.contains(agent_id))
            .map(|(agent_id, state)| (*agent_id, state.last_activity_ms().unwrap_or(0)))
            .collect();
        idle.sort_by_key(|(_, last)| *last);
        let remove_count = guard.len().saturating_sub(limit).min(idle.len());
        for (agent_id, _) in idle.into_iter().take(remove_count) {
            guard.remove(&agent_id);
        }
    }

    fn prune_lifecycle_locked(
        guard: &mut HashMap<MachineId, DirectLifecycleState>,
        connected: &HashSet<MachineId>,
    ) {
        let limit = direct_diagnostics_retain_limit(connected.len());
        if guard.len() <= limit {
            return;
        }

        let now = now_unix_ms_lossy();
        guard.retain(|machine_id, state| {
            connected.contains(machine_id)
                || state
                    .last_updated_at_ms
                    .is_some_and(|last| now.saturating_sub(last) <= DIRECT_DIAGNOSTICS_IDLE_TTL_MS)
        });

        if guard.len() <= limit {
            return;
        }
        let mut idle: Vec<(MachineId, u64)> = guard
            .iter()
            .filter(|(machine_id, _)| !connected.contains(machine_id))
            .map(|(machine_id, state)| (*machine_id, state.last_updated_at_ms.unwrap_or(0)))
            .collect();
        idle.sort_by_key(|(_, last)| *last);
        let remove_count = guard.len().saturating_sub(limit).min(idle.len());
        for (machine_id, _) in idle.into_iter().take(remove_count) {
            guard.remove(&machine_id);
        }
    }

    fn subscriber_snapshot(&self) -> Vec<(u64, Arc<DirectSubscriberQueue>)> {
        match self.subscribers.lock() {
            Ok(guard) => guard.iter().map(|(id, tx)| (*id, tx.clone())).collect(),
            Err(e) => {
                tracing::error!("direct subscriber registry poisoned: {e}");
                Vec::new()
            }
        }
    }

    fn remove_subscribers(&self, ids: &[u64]) {
        match self.subscribers.lock() {
            Ok(mut guard) => {
                for id in ids {
                    guard.remove(id);
                }
            }
            Err(e) => tracing::error!("direct subscriber registry poisoned: {e}"),
        }
    }

    /// Receive the next direct message (blocking).
    pub async fn recv(&self) -> Option<DirectMessage> {
        let mut rx = self.internal_rx.lock().await;
        rx.recv().await
    }

    /// Encode a direct message for transmission.
    ///
    /// Format: `[0x10][sender_agent_id: 32 bytes][payload]`
    pub fn encode_message(sender_agent_id: &AgentId, payload: &[u8]) -> NetworkResult<Vec<u8>> {
        if payload.len() > MAX_DIRECT_PAYLOAD_SIZE {
            return Err(NetworkError::PayloadTooLarge {
                size: payload.len(),
                max: MAX_DIRECT_PAYLOAD_SIZE,
            });
        }

        let mut buf = Vec::with_capacity(1 + 32 + payload.len());
        buf.push(DIRECT_MESSAGE_STREAM_TYPE);
        buf.extend_from_slice(&sender_agent_id.0);
        buf.extend_from_slice(payload);
        Ok(buf)
    }

    /// Decode a direct message from the wire.
    ///
    /// Returns (sender_agent_id, payload) on success.
    pub fn decode_message(data: &[u8]) -> NetworkResult<(AgentId, Vec<u8>)> {
        // Minimum size: 1 (type) + 32 (agent_id) = 33 bytes
        if data.len() < 33 {
            return Err(NetworkError::InvalidMessage(
                "Direct message too short".to_string(),
            ));
        }

        if data[0] != DIRECT_MESSAGE_STREAM_TYPE {
            return Err(NetworkError::InvalidMessage(format!(
                "Invalid stream type byte: expected {}, got {}",
                DIRECT_MESSAGE_STREAM_TYPE, data[0]
            )));
        }

        let mut agent_id_bytes = [0u8; 32];
        agent_id_bytes.copy_from_slice(&data[1..33]);
        let sender = AgentId(agent_id_bytes);

        let payload = data[33..].to_vec();

        Ok((sender, payload))
    }
}

impl Default for DirectMessaging {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn dm_payload_digest_is_stable_and_short() {
        let payload = b"hello world".to_vec();
        let digest = dm_payload_digest_hex(&payload);
        assert_eq!(digest.len(), 16);
        assert!(digest.chars().all(|c| c.is_ascii_hexdigit()));
        assert_eq!(dm_payload_digest_hex(&payload), digest);

        let other = dm_payload_digest_hex(b"different");
        assert_ne!(other, digest);
    }

    #[test]
    fn test_encode_decode_roundtrip() {
        let agent_id = AgentId([42u8; 32]);
        let payload = b"hello world".to_vec();

        let encoded = DirectMessaging::encode_message(&agent_id, &payload).unwrap();

        assert_eq!(encoded[0], DIRECT_MESSAGE_STREAM_TYPE);
        assert_eq!(encoded.len(), 1 + 32 + payload.len());

        let (decoded_agent, decoded_payload) = DirectMessaging::decode_message(&encoded).unwrap();

        assert_eq!(decoded_agent, agent_id);
        assert_eq!(decoded_payload, payload);
    }

    #[test]
    fn test_decode_too_short() {
        let short_data = vec![DIRECT_MESSAGE_STREAM_TYPE; 10];
        let result = DirectMessaging::decode_message(&short_data);
        assert!(result.is_err());
    }

    #[test]
    fn test_decode_wrong_type() {
        let mut data = vec![0x00; 50]; // Wrong type byte
        data[0] = 0x01;
        let result = DirectMessaging::decode_message(&data);
        assert!(result.is_err());
    }

    #[test]
    fn test_encode_payload_too_large() {
        let agent_id = AgentId([1u8; 32]);
        let payload = vec![0u8; MAX_DIRECT_PAYLOAD_SIZE + 1];
        let result = DirectMessaging::encode_message(&agent_id, &payload);
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_register_and_lookup() {
        let dm = DirectMessaging::new();
        let agent_id = AgentId([1u8; 32]);
        let machine_id = MachineId([2u8; 32]);

        dm.register_agent(agent_id, machine_id).await;

        let lookup = dm.lookup_agent(&machine_id).await;
        assert_eq!(lookup, Some(agent_id));
    }

    #[tokio::test]
    async fn test_connection_tracking() {
        let dm = DirectMessaging::new();
        let agent_id = AgentId([1u8; 32]);
        let machine_id = MachineId([2u8; 32]);

        assert!(!dm.is_connected(&agent_id).await);

        dm.mark_connected(agent_id, machine_id).await;
        assert!(dm.is_connected(&agent_id).await);
        assert_eq!(dm.get_machine_id(&agent_id).await, Some(machine_id));

        let connected = dm.connected_agents().await;
        assert_eq!(connected, vec![agent_id]);

        dm.mark_disconnected(&agent_id).await;
        assert!(!dm.is_connected(&agent_id).await);
    }

    #[tokio::test]
    async fn test_message_subscription() {
        let dm = DirectMessaging::new();
        let mut rx = dm.subscribe();

        let sender = AgentId([1u8; 32]);
        let machine_id = MachineId([2u8; 32]);
        let payload = b"test message".to_vec();

        dm.handle_incoming(machine_id, sender, payload.clone(), true, None)
            .await;

        let msg = rx.recv().await.unwrap();
        assert_eq!(msg.sender, sender);
        assert_eq!(msg.machine_id, machine_id);
        assert_eq!(msg.payload, payload);
        assert!(msg.verified);
        assert!(msg.trust_decision.is_none());

        let snap = dm.diagnostics_snapshot();
        assert_eq!(snap.stats.incoming_envelopes_total, 1);
        assert_eq!(snap.stats.incoming_delivered_to_subscribe, 1);
        assert_eq!(snap.stats.subscriber_channel_lagged, 0);
    }

    #[tokio::test]
    async fn test_message_subscription_clone_gets_independent_queue() {
        let dm = DirectMessaging::new();
        let mut rx1 = dm.subscribe();
        let mut rx2 = rx1.clone();

        let sender = AgentId([3u8; 32]);
        let machine_id = MachineId([4u8; 32]);
        let payload = b"fanout".to_vec();

        dm.handle_incoming(machine_id, sender, payload.clone(), true, None)
            .await;

        assert_eq!(rx1.recv().await.unwrap().payload, payload);
        assert_eq!(rx2.recv().await.unwrap().payload, payload);
        assert_eq!(dm.subscriber_count(), 2);
    }

    #[tokio::test]
    async fn dm_subscriber_bounded_drop_oldest_keeps_stream_alive() {
        let dm = DirectMessaging::with_subscriber_capacity(2);
        let mut lagging_rx = dm.subscribe();
        let sender = AgentId([5u8; 32]);
        let machine_id = MachineId([6u8; 32]);

        for idx in 0_u64..=2 {
            dm.handle_incoming(machine_id, sender, idx.to_be_bytes().to_vec(), true, None)
                .await;
        }

        let snap = dm.diagnostics_snapshot();
        assert_eq!(snap.stats.subscriber_channel_lagged, 1);
        assert_eq!(snap.stats.subscriber_events_evicted, 1);
        assert_eq!(snap.subscriber_count, 1);

        let first = lagging_rx.recv().await.unwrap();
        assert_eq!(first.payload, 1_u64.to_be_bytes().to_vec());
        let second = lagging_rx.recv().await.unwrap();
        assert_eq!(second.payload, 2_u64.to_be_bytes().to_vec());
    }

    #[test]
    fn x0x_0041_current_generation_tracks_established_and_replaced() {
        let dm = DirectMessaging::new();
        let machine_id = MachineId([0xAB; 32]);
        assert_eq!(dm.current_generation(&machine_id), None);
        dm.record_lifecycle_established(machine_id, Some(7));
        assert_eq!(dm.current_generation(&machine_id), Some(7));
        dm.record_lifecycle_replaced(machine_id, 9);
        assert_eq!(dm.current_generation(&machine_id), Some(9));
    }

    #[tokio::test]
    async fn x0x_0041_subscribe_lifecycle_replaced_broadcasts_supersede() {
        let dm = DirectMessaging::new();
        let mut rx = dm.subscribe_lifecycle_replaced();
        let machine_a = MachineId([0xA1; 32]);
        let machine_b = MachineId([0xB2; 32]);

        // Established events do NOT fire on the prefer-newest channel.
        dm.record_lifecycle_established(machine_a, Some(1));
        // Use try_recv to confirm no event has been queued.
        match rx.try_recv() {
            Err(tokio::sync::broadcast::error::TryRecvError::Empty) => {}
            other => panic!("expected Empty for Established, got {other:?}"),
        }

        dm.record_lifecycle_replaced(machine_a, 2);
        let (m, gen) = rx.recv().await.expect("Replaced event");
        assert_eq!(m, machine_a);
        assert_eq!(gen, 2);

        // Supersede on a different peer also lands on the broadcast.
        dm.record_lifecycle_replaced(machine_b, 5);
        let (m, gen) = rx.recv().await.expect("Replaced event");
        assert_eq!(m, machine_b);
        assert_eq!(gen, 5);
    }

    #[test]
    fn test_lifecycle_blocks_only_current_generation() {
        let dm = DirectMessaging::new();
        let machine_id = MachineId([7u8; 32]);

        dm.record_lifecycle_established(machine_id, Some(1));
        assert!(dm.lifecycle_block_reason(&machine_id).is_none());

        dm.record_lifecycle_replaced(machine_id, 2);
        dm.record_lifecycle_blocked(machine_id, Some(1), "closed: superseded");
        assert!(dm.lifecycle_block_reason(&machine_id).is_none());

        dm.record_lifecycle_blocked(machine_id, Some(2), "closed: timed out");
        assert_eq!(
            dm.lifecycle_block_reason(&machine_id).as_deref(),
            Some("closed: timed out")
        );

        dm.record_lifecycle_established(machine_id, Some(3));
        assert!(dm.lifecycle_block_reason(&machine_id).is_none());
    }

    #[test]
    fn direct_diagnostics_prune_idle_entries_to_scaled_bound() {
        fn agent_id_from_u32(id: u32) -> AgentId {
            let mut bytes = [0u8; 32];
            bytes[..4].copy_from_slice(&id.to_be_bytes());
            AgentId(bytes)
        }

        let now = now_unix_ms_lossy();
        let connected = agent_id_from_u32(1);
        let mut connected_set = HashSet::new();
        connected_set.insert(connected);

        let mut guard = HashMap::new();
        guard.insert(
            connected,
            DirectPeerDiagnosticsState {
                last_recv_at_ms: Some(0),
                ..DirectPeerDiagnosticsState::default()
            },
        );
        for id in 2..1100 {
            guard.insert(
                agent_id_from_u32(id),
                DirectPeerDiagnosticsState {
                    last_recv_at_ms: Some(now),
                    ..DirectPeerDiagnosticsState::default()
                },
            );
        }

        DirectMessaging::prune_peer_diagnostics_locked(&mut guard, &connected_set);

        assert!(guard.len() <= DIRECT_DIAGNOSTICS_MIN_RETAIN);
        assert!(guard.contains_key(&connected));
    }

    #[test]
    fn test_direct_message_payload_str() {
        let msg = DirectMessage::new(AgentId([1u8; 32]), MachineId([2u8; 32]), b"hello".to_vec());
        assert_eq!(msg.payload_str(), Some("hello"));

        let binary_msg =
            DirectMessage::new(AgentId([1u8; 32]), MachineId([2u8; 32]), vec![0xff, 0xfe]);
        assert!(binary_msg.payload_str().is_none());
    }
}