velo 0.3.0

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Control-plane handler constructors for the anchor lifecycle.
//!
//! This module provides four [`crate::messenger::Handler`] constructors:
//! - [`create_anchor_attach_handler`]: validates anchor existence, calls
//!   `transport.bind().await` (outside shard lock), then atomically stores
//!   the [`flume::Receiver`] in the anchor entry.
//! - [`create_anchor_detach_handler`]: clears attachment, cancels CancellationToken,
//!   injects [`crate::streaming::frame::StreamFrame::Detached`] sentinel; anchor stays in registry.
//! - [`create_anchor_finalize_handler`]: injects [`crate::streaming::frame::StreamFrame::Finalized`]
//!   sentinel, then removes anchor from registry.
//! - [`create_anchor_cancel_handler`]: removes anchor from registry with no sentinel injection.

use crate::observability::{HandlerOutcome, StreamingOp};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::{Duration, Instant};

use crate::streaming::anchor::AnchorManager;
use crate::streaming::handle::StreamAnchorHandle;

/// Number of consecutive missed heartbeat windows that trigger `Dropped` injection.
///
/// The reader pump tolerates `DETECTION_MULTIPLIER * heartbeat_interval` of total silence
/// (each window of length `heartbeat_interval`) before declaring the sender dead.
/// Both the producer (`StreamSender`) heartbeat cadence and the consumer (`reader_pump`)
/// per-window deadline are negotiated via `AnchorAttachResponse::heartbeat_interval_ms`,
/// but the multiplier itself is a protocol constant agreed by both sides.
pub const DETECTION_MULTIPLIER: u8 = 3;

/// Default heartbeat interval (milliseconds) used when `AnchorAttachResponse::Ok` is
/// deserialized from a wire payload that predates the `heartbeat_interval_ms` field.
/// Matches the historical hardcoded 5s constant.
fn default_heartbeat_interval_ms() -> u64 {
    5_000
}

// ---------------------------------------------------------------------------
// StreamCancelHandle
// ---------------------------------------------------------------------------

/// Compact wire handle encoding the sender's [`velo_ext::WorkerId`] (upper 64 bits)
/// and the sender's local stream ID (lower 64 bits) into a single `u128`.
///
/// Serializes via rmp-serde as a two-field struct `{hi: u64, lo: u64}` — not as raw
/// binary bytes — to guarantee correct round-tripping across msgpack boundaries.
/// Identical encoding to [`StreamAnchorHandle`] but scoped to the sender side.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct StreamCancelHandle(u128);

/// Private wire representation for rmp-serde serialization.
///
/// rmp-serde encodes a raw `u128` as a MessagePack binary blob (`bin8`), which
/// cannot be decoded back to a struct. By delegating to this two-field struct we
/// encode as a fixmap with named fields that round-trip correctly.
#[derive(Serialize, Deserialize)]
struct StreamCancelHandleWire {
    hi: u64,
    lo: u64,
}

impl StreamCancelHandle {
    /// Encode a sender [`velo_ext::WorkerId`] and stream ID into a [`StreamCancelHandle`].
    pub fn pack(worker_id: velo_ext::WorkerId, stream_id: u64) -> Self {
        Self(((worker_id.as_u64() as u128) << 64) | (stream_id as u128))
    }

    /// Decode the sender [`velo_ext::WorkerId`] and stream ID from this handle.
    pub fn unpack(self) -> (velo_ext::WorkerId, u64) {
        let hi = (self.0 >> 64) as u64;
        let lo = self.0 as u64;
        (velo_ext::WorkerId::from_u64(hi), lo)
    }
}

impl Serialize for StreamCancelHandle {
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        StreamCancelHandleWire {
            hi: (self.0 >> 64) as u64,
            lo: self.0 as u64,
        }
        .serialize(serializer)
    }
}

impl<'de> Deserialize<'de> for StreamCancelHandle {
    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        let wire = StreamCancelHandleWire::deserialize(deserializer)?;
        Ok(Self(((wire.hi as u128) << 64) | (wire.lo as u128)))
    }
}

// ---------------------------------------------------------------------------
// StreamCancelRequest
// ---------------------------------------------------------------------------

/// Payload for the `_stream_cancel` active message.
///
/// The receiver (sender-side worker) looks up `sender_stream_id` in the
/// [`SenderRegistry`] to find and cancel the corresponding [`SenderEntry`].
#[derive(Debug, Serialize, Deserialize)]
pub struct StreamCancelRequest {
    pub sender_stream_id: u64,
}

// ---------------------------------------------------------------------------
// SenderEntry + SenderRegistry
// ---------------------------------------------------------------------------

/// A single slot in the sender-side registry, representing an active [`crate::streaming::sender::StreamSender`].
///
/// Stored per active stream. The `_stream_cancel` handler retrieves and removes
/// the entry then triggers both the user-facing cancellation token and the
/// poison-drop mechanism via `rx_closer`.
pub struct SenderEntry {
    /// Fires when `_stream_cancel` is received — user-facing via `cancellation_token()`.
    pub cancel_token: tokio_util::sync::CancellationToken,

    /// Drop this to signal cancellation to `StreamSender::send()` via
    /// `poison_tx.is_disconnected()`. Wrapped in `Mutex<Option<...>>` so the
    /// cancel handler can take it exactly once.
    pub rx_closer: std::sync::Mutex<Option<flume::Receiver<()>>>,
}

/// Sender-side registry of active [`SenderEntry`] slots.
///
/// Keyed by the sender's local stream ID (`u64`). Mirrored in structure to the
/// anchor registry (`DashMap<u64, AnchorEntry>`) on the receiver side.
///
/// `pub` so that [`create_stream_cancel_handler`] can accept `Arc<SenderRegistry>`
/// at its public function signature. Callers outside this crate hold it via `Arc`.
#[derive(Default)]
pub struct SenderRegistry {
    pub senders: dashmap::DashMap<u64, SenderEntry>,
}

// ---------------------------------------------------------------------------
// create_stream_cancel_handler
// ---------------------------------------------------------------------------

/// Build the `_stream_cancel` handler.
///
/// When the consumer-side anchor receives a cancel request, it sends a
/// `_stream_cancel` active message to the sender's worker. This handler:
/// 1. Looks up the [`SenderEntry`] by `sender_stream_id`.
/// 2. Drops the `rx_closer` to poison the sender channel.
/// 3. Cancels the user-facing `cancel_token`.
///
/// Idempotent: if the entry is absent the handler returns `Ok(())` silently.
pub fn create_stream_cancel_handler(
    sender_registry: Arc<SenderRegistry>,
) -> crate::messenger::Handler {
    crate::messenger::Handler::am_handler(
        "_stream_cancel",
        move |ctx: crate::messenger::Context| {
            let req = serde_json::from_slice::<StreamCancelRequest>(&ctx.payload)?;
            if let Some((_, entry)) = sender_registry.senders.remove(&req.sender_stream_id) {
                // Poison the tx channel: drop the receiver end so
                // poison_tx.is_disconnected() is true in StreamSender::send()
                drop(entry.rx_closer.lock().unwrap().take());
                // Signal the token so user code can react proactively
                entry.cancel_token.cancel();
            }
            Ok(())
        },
    )
    .build()
}

// ---------------------------------------------------------------------------
// Request / Response types
// ---------------------------------------------------------------------------

/// Request to attach a transport sender to an anchor.
///
/// `session_id` is an opaque caller-assigned identifier that may be forwarded
/// to the transport layer for logging and routing purposes.
///
/// `stream_cancel_handle` encodes the sender's worker ID and local stream ID so that
/// the anchor can route `_stream_cancel` active messages back to the correct sender.
#[derive(Debug, Serialize, Deserialize)]
pub struct AnchorAttachRequest {
    pub handle: StreamAnchorHandle,
    pub session_id: u64,
    /// Encodes the sender's WorkerId + sender_stream_id. Stored in the anchor entry
    /// on successful attach so the anchor knows where to route upstream cancel AMs.
    pub stream_cancel_handle: StreamCancelHandle,
}

/// Response from the attach handler.
#[derive(Debug, Serialize, Deserialize)]
pub enum AnchorAttachResponse {
    /// Attach succeeded; caller can connect to `stream_endpoint`.
    ///
    /// `heartbeat_interval_ms` tells the sender how often it must emit a
    /// [`crate::streaming::frame::StreamFrame::Heartbeat`] when no data frames are flowing.
    /// The consumer's reader pump will tolerate `DETECTION_MULTIPLIER * heartbeat_interval_ms`
    /// of total silence before injecting `Dropped`. The field is carried as `u64` ms
    /// (rather than `Duration`) for stable msgpack encoding, and defaults to 5000ms
    /// when absent so older clients continue to deserialize new responses unchanged.
    Ok {
        stream_endpoint: String,
        #[serde(default = "default_heartbeat_interval_ms")]
        heartbeat_interval_ms: u64,
    },
    /// Attach failed; `reason` describes why.
    Err { reason: String },
}

/// Request to detach the current sender from an anchor without closing it.
///
/// After detach the anchor remains in the registry so a new sender may attach.
#[derive(Debug, Serialize, Deserialize)]
pub struct AnchorDetachRequest {
    pub handle: StreamAnchorHandle,
}

/// Request to finalize (permanently close) an anchor.
///
/// After finalize the anchor is removed from the registry.
#[derive(Debug, Serialize, Deserialize)]
pub struct AnchorFinalizeRequest {
    pub handle: StreamAnchorHandle,
}

/// Request to cancel an anchor with no sentinel injection.
///
/// Used when a sender exits before attaching or when an explicit abort is needed.
/// After cancel the anchor is removed from the registry.
#[derive(Debug, Serialize, Deserialize)]
pub struct AnchorCancelRequest {
    pub handle: StreamAnchorHandle,
}

// ---------------------------------------------------------------------------
// Reader pump
// ---------------------------------------------------------------------------

/// Reader pump: bridges transport frames to the anchor's delivery channel.
///
/// Spawned as a tokio task after successful attach. Reads from the transport
/// receiver, forwards to the anchor's frame_tx. Monitors for heartbeat
/// timeouts: `DETECTION_MULTIPLIER` consecutive `heartbeat_deadline` windows
/// with no frames trigger Dropped sentinel injection, registry removal
/// (LIVE-02), and cleanup. The deadline is negotiated at attach time via
/// `AnchorAttachResponse::heartbeat_interval_ms`.
pub(crate) async fn reader_pump(
    transport_rx: flume::Receiver<Vec<u8>>,
    frame_tx: flume::Sender<Vec<u8>>,
    cancel_token: tokio_util::sync::CancellationToken,
    ctx: crate::streaming::anchor::AnchorContext,
    local_id: u64,
    heartbeat_deadline: Duration,
) {
    let crate::streaming::anchor::AnchorContext {
        registry,
        mpsc_registry,
        metrics,
    } = ctx;
    let mut missed_heartbeats: u8 = 0;

    loop {
        tokio::select! {
            _ = cancel_token.cancelled() => break,
            result = tokio::time::timeout(heartbeat_deadline, transport_rx.recv_async()) => {
                match result {
                    Ok(Ok(bytes)) => {
                        // Any frame (data or heartbeat) proves liveness
                        missed_heartbeats = 0;
                        // Forward to anchor's frame channel
                        if frame_tx.send_async(bytes).await.is_err() {
                            break; // consumer dropped
                        }
                    }
                    Ok(Err(_)) => break, // transport channel closed
                    Err(_timeout) => {
                        missed_heartbeats += 1;
                        if missed_heartbeats >= DETECTION_MULTIPLIER {
                            // Inject Dropped sentinel -- sender is dead
                            let dropped_bytes = crate::streaming::sender::cached_dropped().clone();
                            let _ = frame_tx.send_async(dropped_bytes).await;
                            // LIVE-02: Full anchor cleanup -- remove from registry
                            // so no stale entry remains (ANCR-04)
                            if let Some((_, entry)) = registry.remove(&local_id) {
                                entry.cancel_token.cancel();
                                crate::streaming::anchor::set_active_anchor_gauge(
                                    metrics.as_ref(),
                                    &registry,
                                    &mpsc_registry,
                                );
                            }
                            break;
                        }
                    }
                }
            }
        }
    }
    // Cleanup: cancel token so other paths know the pump exited
    cancel_token.cancel();
}

// ---------------------------------------------------------------------------
// Handler constructors
// ---------------------------------------------------------------------------

/// Build the `_anchor_attach` handler.
///
/// Uses the bind-then-lock pattern: calls `transport.bind().await` OUTSIDE the
/// DashMap shard lock, then atomically checks and sets the attachment under the lock.
/// This avoids holding the shard lock across an async `.await` point.
///
/// Returns [`AnchorAttachResponse::Ok`] on success or [`AnchorAttachResponse::Err`] on
/// any failure (not found, already attached, transport error).
pub fn create_anchor_attach_handler(manager: Arc<AnchorManager>) -> crate::messenger::Handler {
    crate::messenger::Handler::typed_unary_async(
        "_anchor_attach",
        move |ctx: crate::messenger::TypedContext<AnchorAttachRequest>| {
            let manager = manager.clone();
            async move {
                let started = Instant::now();
                let req = ctx.input;

                // Defence-in-depth: reject MPSC handles at the SPSC attach
                // endpoint. The client-side `attach_stream_anchor` already
                // rejects these before the AM, but misbehaving or older
                // clients may still hit the wire.
                if req.handle.is_mpsc_stream() {
                    manager.record_streaming_operation(
                        StreamingOp::Attach,
                        HandlerOutcome::Error,
                        "unknown",
                        started,
                    );
                    return Ok(AnchorAttachResponse::Err {
                        reason: format!("anchor {} is mpsc; use _mpsc_anchor_attach", req.handle),
                    });
                }

                let (_, local_id) = req.handle.unpack();

                // Step 1: Quick check -- anchor exists and is unattached (drop lock)
                {
                    let entry = manager.registry.get(&local_id);
                    match entry {
                        None => {
                            manager.record_streaming_operation(
                                StreamingOp::Attach,
                                HandlerOutcome::Error,
                                "unknown",
                                started,
                            );
                            return Ok(AnchorAttachResponse::Err {
                                reason: format!("anchor {} not found", req.handle),
                            });
                        }
                        Some(e) if e.attachment => {
                            manager.record_streaming_operation(
                                StreamingOp::Attach,
                                HandlerOutcome::Error,
                                "unknown",
                                started,
                            );
                            return Ok(AnchorAttachResponse::Err {
                                reason: format!("anchor {} already attached", req.handle),
                            });
                        }
                        _ => {} // looks good, proceed
                    }
                } // DashMap ref dropped here

                // Step 2: Async bind OUTSIDE shard lock
                let (endpoint, receiver) =
                    match manager.transport.bind(local_id, req.session_id).await {
                        Ok(pair) => pair,
                        Err(e) => {
                            manager.record_streaming_operation(
                                StreamingOp::Attach,
                                HandlerOutcome::Error,
                                "unknown",
                                started,
                            );
                            return Ok(AnchorAttachResponse::Err {
                                reason: format!("transport error: {}", e),
                            });
                        }
                    };

                // Step 3: Atomically set attachment under shard lock
                use dashmap::mapref::entry::Entry;
                match manager.registry.entry(local_id) {
                    Entry::Vacant(_) => {
                        manager.record_streaming_operation(
                            StreamingOp::Attach,
                            HandlerOutcome::Error,
                            "unknown",
                            started,
                        );
                        Ok(AnchorAttachResponse::Err {
                            reason: format!("anchor {} removed during bind", req.handle),
                        })
                    }
                    Entry::Occupied(mut occ) => {
                        let entry = occ.get_mut();
                        if entry.attachment {
                            manager.record_streaming_operation(
                                StreamingOp::Attach,
                                HandlerOutcome::Error,
                                "unknown",
                                started,
                            );
                            Ok(AnchorAttachResponse::Err {
                                reason: format!("anchor {} already attached", req.handle),
                            })
                        } else {
                            // Derive a child token for this pump so detach can cancel it
                            // without poisoning the parent (which lives for the anchor's lifetime).
                            let pump_cancel = entry.cancel_token.child_token();
                            entry.active_pump_token = Some(pump_cancel.clone());
                            let pump_frame_tx = entry.frame_tx.clone();
                            // Snapshot the negotiated heartbeat interval before dropping the lock.
                            let heartbeat_interval = entry.heartbeat_interval;

                            // Mark as attached and store cancel handle for upstream cancel routing
                            entry.attachment = true;
                            entry.stream_cancel_handle = Some(req.stream_cancel_handle);

                            // Drop shard lock before spawning
                            drop(occ);

                            // Spawn reader pump as background task
                            let (_, local_id) = req.handle.unpack();
                            tokio::spawn(reader_pump(
                                receiver,      // transport receiver from bind
                                pump_frame_tx, // cloned from entry
                                pump_cancel,   // cloned from entry
                                manager.anchor_context(),
                                local_id, // anchor's local_id
                                heartbeat_interval,
                            ));

                            let transport_scheme =
                                endpoint.split("://").next().unwrap_or("unknown");
                            manager.record_streaming_operation(
                                StreamingOp::Attach,
                                HandlerOutcome::Success,
                                transport_scheme,
                                started,
                            );

                            Ok(AnchorAttachResponse::Ok {
                                stream_endpoint: endpoint,
                                heartbeat_interval_ms: heartbeat_interval.as_millis() as u64,
                            })
                        }
                    }
                }
            }
        },
    )
    .spawn()
    .build()
}

/// Build the `_anchor_detach` handler.
///
/// Atomically clears `attachment` via `DashMap::entry()`, then -- after dropping the
/// shard lock -- cancels the `CancellationToken` and injects a
/// [`crate::streaming::frame::StreamFrame::Detached`] sentinel into the frame channel.
/// The anchor remains in the registry so a new sender may re-attach.
///
/// Idempotent: if the anchor is not found, returns `Ok(())`.
pub fn create_anchor_detach_handler(manager: Arc<AnchorManager>) -> crate::messenger::Handler {
    crate::messenger::Handler::typed_unary_async(
        "_anchor_detach",
        move |ctx: crate::messenger::TypedContext<AnchorDetachRequest>| {
            let manager = manager.clone();
            async move {
                let started = Instant::now();
                let req = ctx.input;
                let (_, local_id) = req.handle.unpack();

                use dashmap::mapref::entry::Entry;
                // Atomically clear attachment and clone cancel_token + frame_tx
                // before dropping the shard lock (never hold DashMap ref across channel ops).
                let maybe_entry_info = match manager.registry.entry(local_id) {
                    Entry::Vacant(_) => None,
                    Entry::Occupied(mut occ) => {
                        let entry = occ.get_mut();
                        // Clear the attachment flag
                        entry.attachment = false;
                        // Take the child token (leaves None) so the next attach creates a fresh one
                        Some((entry.active_pump_token.take(), entry.frame_tx.clone()))
                    }
                };
                // shard lock is now dropped

                if let Some((maybe_pump_token, frame_tx)) = maybe_entry_info {
                    if let Some(pump_token) = maybe_pump_token {
                        pump_token.cancel();
                    }
                    let sentinel_bytes = crate::streaming::sender::cached_detached().clone();
                    let _ = frame_tx.try_send(sentinel_bytes);
                    manager.record_streaming_operation(
                        StreamingOp::Detach,
                        HandlerOutcome::Success,
                        "velo",
                        started,
                    );
                } else {
                    manager.record_streaming_operation(
                        StreamingOp::Detach,
                        HandlerOutcome::Error,
                        "velo",
                        started,
                    );
                }

                Ok(())
            }
        },
    )
    .spawn()
    .build()
}

/// Build the `_anchor_finalize` handler.
///
/// Atomically removes the anchor from the registry via `remove_anchor()`, injects a
/// [`crate::streaming::frame::StreamFrame::Finalized`] sentinel, and cancels the `CancellationToken`.
///
/// Idempotent: if the anchor is already absent, returns `Ok(())`.
pub fn create_anchor_finalize_handler(manager: Arc<AnchorManager>) -> crate::messenger::Handler {
    crate::messenger::Handler::typed_unary_async(
        "_anchor_finalize",
        move |ctx: crate::messenger::TypedContext<AnchorFinalizeRequest>| {
            let manager = manager.clone();
            async move {
                let started = Instant::now();
                let req = ctx.input;
                let (_, local_id) = req.handle.unpack();

                // remove_anchor cancels the token and returns the entry
                if let Some(entry) = manager.remove_anchor(local_id) {
                    let sentinel_bytes = crate::streaming::sender::cached_finalized().clone();
                    let _ = entry.frame_tx.try_send(sentinel_bytes);
                    manager.record_streaming_operation(
                        StreamingOp::Finalize,
                        HandlerOutcome::Success,
                        "velo",
                        started,
                    );
                } else {
                    manager.record_streaming_operation(
                        StreamingOp::Finalize,
                        HandlerOutcome::Error,
                        "velo",
                        started,
                    );
                }

                Ok(())
            }
        },
    )
    .spawn()
    .build()
}

/// Build the `_anchor_cancel` handler.
///
/// Removes the anchor from the registry with no sentinel injection.
/// Used when a sender aborts before or during attachment.
///
/// Idempotent: calling cancel on an already-absent anchor does not panic.
pub fn create_anchor_cancel_handler(manager: Arc<AnchorManager>) -> crate::messenger::Handler {
    crate::messenger::Handler::typed_unary_async(
        "_anchor_cancel",
        move |ctx: crate::messenger::TypedContext<AnchorCancelRequest>| {
            let manager = manager.clone();
            async move {
                let started = Instant::now();
                let req = ctx.input;
                let (_, local_id) = req.handle.unpack();

                // remove_anchor is a no-op (returns None) if anchor absent -- idempotent
                if let Some(entry) = manager.remove_anchor(local_id) {
                    entry.cancel_token.cancel();
                    manager.record_streaming_operation(
                        StreamingOp::Cancel,
                        HandlerOutcome::Success,
                        "velo",
                        started,
                    );
                } else {
                    manager.record_streaming_operation(
                        StreamingOp::Cancel,
                        HandlerOutcome::Error,
                        "velo",
                        started,
                    );
                }

                Ok(())
            }
        },
    )
    .spawn()
    .build()
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use anyhow::Result as AnyhowResult;
    use futures::StreamExt;
    use futures::future::BoxFuture;
    use std::sync::Arc;

    // -----------------------------------------------------------------------
    // MockFrameTransport (test-only)
    // -----------------------------------------------------------------------

    struct MockFrameTransport;

    impl crate::streaming::transport::FrameTransport for MockFrameTransport {
        fn bind(
            &self,
            _anchor_id: u64,
            _session_id: u64,
        ) -> BoxFuture<'_, AnyhowResult<(String, flume::Receiver<Vec<u8>>)>> {
            Box::pin(async {
                Ok((
                    "mock://test-endpoint".to_string(),
                    flume::bounded::<Vec<u8>>(256).1,
                ))
            })
        }

        fn connect(
            &self,
            _endpoint: &str,
            _anchor_id: u64,
            _session_id: u64,
        ) -> BoxFuture<'_, AnyhowResult<flume::Sender<Vec<u8>>>> {
            Box::pin(async { Ok(flume::bounded::<Vec<u8>>(256).0) })
        }
    }

    // -----------------------------------------------------------------------
    // Helper: make a test AnchorManager
    // -----------------------------------------------------------------------

    fn make_test_manager() -> Arc<AnchorManager> {
        let worker_id = velo_ext::WorkerId::from_u64(1);
        let transport = Arc::new(MockFrameTransport);
        Arc::new(AnchorManager::new(worker_id, transport))
    }

    // -----------------------------------------------------------------------
    // Test helpers for calling handler logic directly
    // -----------------------------------------------------------------------

    // We call the handler constructor only to verify it compiles and returns Handler.
    // For behavioral tests, we call the underlying AnchorManager APIs + simulate
    // the same logic the handler performs to verify correctness without needing
    // a running velo_messenger runtime.

    // -----------------------------------------------------------------------
    // Type serialization tests (Task 1 scope)
    // -----------------------------------------------------------------------

    #[test]
    fn test_anchor_attach_response_serde_ok() {
        let resp = AnchorAttachResponse::Ok {
            stream_endpoint: "mock://test-endpoint".to_string(),
            heartbeat_interval_ms: 5000,
        };
        let json = serde_json::to_string(&resp).expect("serialize Ok");
        let decoded: AnchorAttachResponse = serde_json::from_str(&json).expect("deserialize Ok");
        match decoded {
            AnchorAttachResponse::Ok {
                stream_endpoint,
                heartbeat_interval_ms,
            } => {
                assert_eq!(stream_endpoint, "mock://test-endpoint");
                assert_eq!(heartbeat_interval_ms, 5000);
            }
            other => panic!("expected Ok, got {:?}", other),
        }
    }

    #[test]
    fn test_anchor_attach_response_rmp_round_trip_non_default_heartbeat() {
        // msgpack must carry the negotiated interval losslessly so the sender
        // gets the cadence the consumer dictated.
        let resp = AnchorAttachResponse::Ok {
            stream_endpoint: "tcp://10.0.0.1:9000".to_string(),
            heartbeat_interval_ms: 1234,
        };
        let bytes = rmp_serde::to_vec(&resp).expect("rmp serialize Ok");
        let decoded: AnchorAttachResponse =
            rmp_serde::from_slice(&bytes).expect("rmp deserialize Ok");
        match decoded {
            AnchorAttachResponse::Ok {
                stream_endpoint,
                heartbeat_interval_ms,
            } => {
                assert_eq!(stream_endpoint, "tcp://10.0.0.1:9000");
                assert_eq!(heartbeat_interval_ms, 1234);
            }
            other => panic!("expected Ok, got {:?}", other),
        }
    }

    #[test]
    fn test_anchor_attach_response_serde_ok_backcompat_default_heartbeat() {
        // An older client/server that has no `heartbeat_interval_ms` field
        // must still deserialize cleanly into the 5000ms default. Construct
        // the legacy shape directly via JSON so we don't need a frozen wire
        // sample.
        let legacy_json = r#"{"Ok":{"stream_endpoint":"mock://legacy"}}"#;
        let decoded: AnchorAttachResponse =
            serde_json::from_str(legacy_json).expect("legacy Ok response must deserialize");
        match decoded {
            AnchorAttachResponse::Ok {
                stream_endpoint,
                heartbeat_interval_ms,
            } => {
                assert_eq!(stream_endpoint, "mock://legacy");
                assert_eq!(
                    heartbeat_interval_ms, 5000,
                    "missing field must default to 5000ms"
                );
            }
            other => panic!("expected Ok, got {:?}", other),
        }
    }

    #[test]
    fn test_anchor_attach_response_serde_err() {
        let resp = AnchorAttachResponse::Err {
            reason: "already attached".to_string(),
        };
        let json = serde_json::to_string(&resp).expect("serialize Err");
        let decoded: AnchorAttachResponse = serde_json::from_str(&json).expect("deserialize Err");
        match decoded {
            AnchorAttachResponse::Err { reason } => {
                assert!(reason.contains("already attached"));
            }
            other => panic!("expected Err, got {:?}", other),
        }
    }

    // -----------------------------------------------------------------------
    // Attach handler tests
    // -----------------------------------------------------------------------

    #[tokio::test]
    async fn test_anchor_attach_handler() {
        let manager = make_test_manager();
        let anchor = manager.create_anchor::<u8>();
        let handle = anchor.handle();
        let (_, local_id) = handle.unpack();

        // Simulate bind-then-lock attach handler logic:
        // Step 1: async bind outside shard lock
        let (endpoint, _receiver) = manager.transport.bind(local_id, 0).await.unwrap();

        // Step 2: atomically set attachment under shard lock
        use dashmap::mapref::entry::Entry;
        let result = match manager.registry.entry(local_id) {
            Entry::Vacant(_) => AnchorAttachResponse::Err {
                reason: format!("anchor {} not found", handle),
            },
            Entry::Occupied(mut occ) => {
                let entry = occ.get_mut();
                if entry.attachment {
                    AnchorAttachResponse::Err {
                        reason: format!("anchor {} already attached", handle),
                    }
                } else {
                    entry.attachment = true;
                    AnchorAttachResponse::Ok {
                        stream_endpoint: endpoint,
                        heartbeat_interval_ms: 5000,
                    }
                }
            }
        };

        match result {
            AnchorAttachResponse::Ok {
                stream_endpoint, ..
            } => {
                assert_eq!(stream_endpoint, "mock://test-endpoint");
            }
            other => panic!("expected Ok, got {:?}", other),
        }

        // Verify attachment is set
        assert!(
            manager
                .registry
                .get(&local_id)
                .map(|e| e.attachment)
                .unwrap_or(false),
            "attachment must be true after attach"
        );

        // Verify handler constructor compiles and returns Handler
        let _handler = create_anchor_attach_handler(manager.clone());
    }

    #[tokio::test]
    async fn test_anchor_attach_already_attached() {
        let manager = make_test_manager();
        let anchor = manager.create_anchor::<u8>();
        let handle = anchor.handle();
        let (_, local_id) = handle.unpack();

        // First attach: set attachment flag directly
        {
            use dashmap::mapref::entry::Entry;
            if let Entry::Occupied(mut occ) = manager.registry.entry(local_id) {
                let entry = occ.get_mut();
                entry.attachment = true;
            }
        }

        // Second attach via handler logic -- should fail
        use dashmap::mapref::entry::Entry;
        let result = match manager.registry.entry(local_id) {
            Entry::Vacant(_) => AnchorAttachResponse::Err {
                reason: format!("anchor {} not found", handle),
            },
            Entry::Occupied(mut occ) => {
                let entry = occ.get_mut();
                if entry.attachment {
                    AnchorAttachResponse::Err {
                        reason: format!("anchor {} already attached", handle),
                    }
                } else {
                    AnchorAttachResponse::Ok {
                        stream_endpoint: "unreachable".to_string(),
                        heartbeat_interval_ms: 5000,
                    }
                }
            }
        };

        match result {
            AnchorAttachResponse::Err { reason } => {
                assert!(
                    reason.contains("already attached"),
                    "reason must mention 'already attached', got: {reason}"
                );
            }
            other => panic!("expected Err, got {:?}", other),
        }
    }

    #[tokio::test]
    async fn test_anchor_attach_not_found() {
        let manager = make_test_manager();
        // Create a handle that is NOT in the registry
        let fake_handle = StreamAnchorHandle::pack(velo_ext::WorkerId::from_u64(1), 9999);

        // Simulate handler logic
        use dashmap::mapref::entry::Entry;
        let local_id = 9999u64;
        let result = match manager.registry.entry(local_id) {
            Entry::Vacant(_) => AnchorAttachResponse::Err {
                reason: format!("anchor {} not found", fake_handle),
            },
            Entry::Occupied(_) => panic!("should not be occupied"),
        };

        match result {
            AnchorAttachResponse::Err { reason } => {
                assert!(
                    reason.contains("not found"),
                    "reason must mention 'not found', got: {reason}"
                );
            }
            other => panic!("expected Err, got {:?}", other),
        }
    }

    // -----------------------------------------------------------------------
    // Detach handler tests
    // -----------------------------------------------------------------------

    #[tokio::test]
    async fn test_anchor_detach_handler() {
        let manager = make_test_manager();
        let mut stream = manager.create_anchor::<Vec<u8>>();
        let handle = stream.handle();
        let (_, local_id) = handle.unpack();

        // Simulate attach: set flag directly
        {
            use dashmap::mapref::entry::Entry;
            if let Entry::Occupied(mut occ) = manager.registry.entry(local_id) {
                let entry = occ.get_mut();
                entry.attachment = true;
            }
        }

        // Simulate detach handler logic (cancels child token, not parent)
        use dashmap::mapref::entry::Entry;
        let maybe_entry_info = match manager.registry.entry(local_id) {
            Entry::Vacant(_) => None,
            Entry::Occupied(mut occ) => {
                let entry = occ.get_mut();
                entry.attachment = false;
                Some((entry.active_pump_token.take(), entry.frame_tx.clone()))
            }
        };

        if let Some((maybe_pump_token, frame_tx)) = maybe_entry_info {
            if let Some(pump_token) = maybe_pump_token {
                pump_token.cancel();
            }
            let sentinel_bytes =
                rmp_serde::to_vec(&crate::streaming::frame::StreamFrame::<Vec<u8>>::Detached)
                    .expect("serialize Detached sentinel");
            let _ = frame_tx.try_send(sentinel_bytes);
        }

        // Verify: attachment is cleared
        assert!(
            manager
                .registry
                .get(&local_id)
                .map(|e| !e.attachment)
                .unwrap_or(false),
            "attachment must be false after detach"
        );

        // Verify: anchor still in registry
        assert!(
            manager.registry.contains_key(&local_id),
            "anchor must remain in registry after detach"
        );

        // Verify: Detached sentinel received via Stream interface
        let result = stream.next().await;
        assert!(
            matches!(
                result,
                Some(Ok(crate::streaming::frame::StreamFrame::Detached))
            ),
            "sentinel must be Detached, got {:?}",
            result
        );

        // Verify handler constructor compiles
        let _handler = create_anchor_detach_handler(manager.clone());
    }

    // -----------------------------------------------------------------------
    // Finalize handler tests
    // -----------------------------------------------------------------------

    #[tokio::test]
    async fn test_anchor_finalize_handler() {
        let manager = make_test_manager();
        let mut stream = manager.create_anchor::<Vec<u8>>();
        let handle = stream.handle();
        let (_, local_id) = handle.unpack();

        // Simulate attach: set flag directly
        {
            use dashmap::mapref::entry::Entry;
            if let Entry::Occupied(mut occ) = manager.registry.entry(local_id) {
                let entry = occ.get_mut();
                entry.attachment = true;
            }
        }

        // Simulate finalize handler logic
        if let Some(entry) = manager.remove_anchor(local_id) {
            let sentinel_bytes =
                rmp_serde::to_vec(&crate::streaming::frame::StreamFrame::<Vec<u8>>::Finalized)
                    .expect("serialize Finalized sentinel");
            let _ = entry.frame_tx.try_send(sentinel_bytes);
        }

        // Verify: anchor removed from registry
        assert!(
            !manager.registry.contains_key(&local_id),
            "anchor must be absent from registry after finalize"
        );

        // Verify: Finalized sentinel received via Stream interface
        let result = stream.next().await;
        assert!(
            matches!(
                result,
                Some(Ok(crate::streaming::frame::StreamFrame::Finalized))
            ),
            "sentinel must be Finalized, got {:?}",
            result
        );

        // Verify handler constructor compiles
        let _handler = create_anchor_finalize_handler(manager.clone());
    }

    // -----------------------------------------------------------------------
    // Cancel handler tests
    // -----------------------------------------------------------------------

    #[tokio::test]
    async fn test_anchor_cancel_handler() {
        let manager = make_test_manager();
        let anchor = manager.create_anchor::<u8>();
        let (_, local_id) = anchor.handle().unpack();

        // Simulate cancel handler logic
        if let Some(entry) = manager.remove_anchor(local_id) {
            entry.cancel_token.cancel();
        }

        // Verify: anchor removed
        assert!(
            !manager.registry.contains_key(&local_id),
            "anchor must be absent after cancel"
        );

        // Idempotent: cancel again -- must not panic
        if let Some(entry) = manager.remove_anchor(local_id) {
            entry.cancel_token.cancel();
        }
        // No panic -- test passes

        // Verify handler constructor compiles
        let _handler = create_anchor_cancel_handler(manager.clone());
    }

    // -----------------------------------------------------------------------
    // reader_pump tests (Plan 08-03, Task 2)
    // -----------------------------------------------------------------------

    /// Helper: set up infrastructure for reader_pump tests.
    /// Returns (transport_tx, frame_rx, cancel_token, registry, local_id).
    #[allow(clippy::type_complexity)]
    fn make_pump_test_infra() -> (
        flume::Sender<Vec<u8>>,   // transport_tx: simulates transport frames
        flume::Receiver<Vec<u8>>, // frame_rx: where pump writes to (consumer side)
        tokio_util::sync::CancellationToken,
        std::sync::Arc<dashmap::DashMap<u64, crate::streaming::anchor::AnchorEntry>>,
        u64, // local_id
    ) {
        let (transport_tx, transport_rx) = flume::bounded::<Vec<u8>>(256);
        let (frame_tx, frame_rx) = flume::bounded::<Vec<u8>>(256);
        let cancel_token = tokio_util::sync::CancellationToken::new();
        let registry = std::sync::Arc::new(dashmap::DashMap::new());
        let local_id = 1u64;

        // Insert an entry in the registry so the pump can remove it
        registry.insert(
            local_id,
            crate::streaming::anchor::AnchorEntry {
                frame_tx: frame_tx.clone(),
                cancel_token: cancel_token.clone(),
                active_pump_token: None,
                attachment: true,
                timeout_cancel: None,
                unattached_timeout: None,
                heartbeat_interval: Duration::from_secs(5),
                stream_cancel_handle: None,
            },
        );

        // Spawn the reader pump
        let pump_cancel = cancel_token.clone();
        let ctx = crate::streaming::anchor::AnchorContext {
            registry: registry.clone(),
            mpsc_registry: std::sync::Arc::new(dashmap::DashMap::new()),
            metrics: None,
        };
        tokio::spawn(reader_pump(
            transport_rx,
            frame_tx,
            pump_cancel,
            ctx,
            local_id,
            Duration::from_secs(5),
        ));

        (transport_tx, frame_rx, cancel_token, registry, local_id)
    }

    #[tokio::test]
    async fn test_pump_forwards_data_frames() {
        let (transport_tx, frame_rx, _cancel, _registry, _id) = make_pump_test_infra();

        // Send a data frame through the transport side
        let data_bytes =
            rmp_serde::to_vec(&crate::streaming::frame::StreamFrame::Item(42u32)).unwrap();
        transport_tx.send_async(data_bytes.clone()).await.unwrap();

        // Should arrive on the frame_rx side
        let received =
            tokio::time::timeout(std::time::Duration::from_millis(500), frame_rx.recv_async())
                .await
                .expect("timeout waiting for frame")
                .expect("frame_rx closed");

        assert_eq!(received, data_bytes, "pump must forward bytes unchanged");
    }

    #[tokio::test]
    async fn test_pump_resets_heartbeat_counter_on_frame() {
        tokio::time::pause();

        let (transport_tx, frame_rx, _cancel, registry, local_id) = make_pump_test_infra();

        // Wait 4.5 seconds (almost one heartbeat window)
        tokio::time::sleep(std::time::Duration::from_millis(4500)).await;

        // Send a frame to reset the counter
        let hb_bytes =
            rmp_serde::to_vec(&crate::streaming::frame::StreamFrame::<()>::Heartbeat).unwrap();
        transport_tx.send_async(hb_bytes).await.unwrap();

        // Wait another 4.5 seconds
        tokio::time::sleep(std::time::Duration::from_millis(4500)).await;

        // Send another frame
        transport_tx
            .send_async(
                rmp_serde::to_vec(&crate::streaming::frame::StreamFrame::<()>::Heartbeat).unwrap(),
            )
            .await
            .unwrap();

        // Drain forwarded frames
        while frame_rx.try_recv().is_ok() {}

        // The anchor should still be in the registry (counter resets each time)
        assert!(
            registry.contains_key(&local_id),
            "anchor must still be in registry -- heartbeat counter was reset"
        );
    }

    #[tokio::test]
    async fn test_pump_injects_dropped_after_3_missed_heartbeats() {
        tokio::time::pause();

        let (transport_tx, frame_rx, _cancel, _registry, _id) = make_pump_test_infra();

        // Keep transport_tx alive but don't send anything -- pump will timeout
        // 3 consecutive 5s windows with no frames trigger Dropped
        tokio::time::sleep(std::time::Duration::from_secs(16)).await;

        // Collect all frames from frame_rx
        let mut frames = Vec::new();
        while let Ok(bytes) = frame_rx.try_recv() {
            frames.push(bytes);
        }

        // The last frame should be a Dropped sentinel
        assert!(
            !frames.is_empty(),
            "must have received at least one frame (Dropped sentinel)"
        );
        let last = frames.last().unwrap();
        let decoded: crate::streaming::frame::StreamFrame<()> =
            rmp_serde::from_slice(last).expect("deserialize");
        assert!(
            matches!(decoded, crate::streaming::frame::StreamFrame::Dropped),
            "last frame must be Dropped, got {:?}",
            decoded
        );

        // Keep transport_tx alive for the duration of the test
        drop(transport_tx);
    }

    #[tokio::test]
    async fn test_pump_removes_registry_entry_after_3_missed_heartbeats() {
        tokio::time::pause();

        let (transport_tx, _frame_rx, _cancel, registry, local_id) = make_pump_test_infra();

        // Keep transport_tx alive but don't send -- pump will timeout
        tokio::time::sleep(std::time::Duration::from_secs(16)).await;

        // LIVE-02: anchor entry must be removed from registry
        assert!(
            !registry.contains_key(&local_id),
            "anchor must be removed from registry after 3 missed heartbeats (LIVE-02)"
        );

        // Keep transport_tx alive for the duration of the test
        drop(transport_tx);
    }

    #[tokio::test]
    async fn test_pump_exits_when_cancel_token_cancelled() {
        let (transport_tx, frame_rx, cancel_token, registry, local_id) = make_pump_test_infra();

        // Cancel the token
        cancel_token.cancel();

        // Give the pump a moment to exit
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;

        // Pump should have exited -- sending on transport_tx should not be forwarded
        let data = rmp_serde::to_vec(&crate::streaming::frame::StreamFrame::Item(99u32)).unwrap();
        let _ = transport_tx.try_send(data);

        // Allow propagation
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;

        // frame_rx should be empty (pump exited, nothing forwarded)
        assert!(
            frame_rx.try_recv().is_err(),
            "no frames should be forwarded after cancel"
        );

        // Pump calls cancel on exit, so token should be cancelled
        assert!(cancel_token.is_cancelled());

        // Registry entry may or may not be removed (cancel != heartbeat death)
        let _ = (registry, local_id);
    }

    #[tokio::test]
    async fn test_pump_exits_when_transport_closes() {
        let (transport_tx, _frame_rx, cancel_token, _registry, _id) = make_pump_test_infra();

        // Drop the transport sender -- transport channel closes
        drop(transport_tx);

        // Give the pump a moment to exit
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;

        // Pump should have exited and cancelled the token
        assert!(
            cancel_token.is_cancelled(),
            "cancel_token must be cancelled after pump exits due to transport close"
        );
    }

    #[tokio::test]
    async fn test_child_token_reattach_pump_survives() {
        let parent = tokio_util::sync::CancellationToken::new();
        let (frame_tx, frame_rx) = flume::bounded::<Vec<u8>>(256);
        let registry = std::sync::Arc::new(dashmap::DashMap::new());
        let local_id = 1u64;

        // --- First attach: spawn pump with child token ---
        let (tx1, rx1) = flume::bounded::<Vec<u8>>(256);
        let child1 = parent.child_token();

        registry.insert(
            local_id,
            crate::streaming::anchor::AnchorEntry {
                frame_tx: frame_tx.clone(),
                cancel_token: parent.clone(),
                active_pump_token: Some(child1.clone()),
                attachment: true,
                timeout_cancel: None,
                unattached_timeout: None,
                heartbeat_interval: Duration::from_secs(5),
                stream_cancel_handle: None,
            },
        );

        let mpsc_reg: std::sync::Arc<
            dashmap::DashMap<u64, crate::streaming::mpsc::anchor::MpscAnchorEntry>,
        > = std::sync::Arc::new(dashmap::DashMap::new());
        let ctx1 = crate::streaming::anchor::AnchorContext {
            registry: registry.clone(),
            mpsc_registry: mpsc_reg.clone(),
            metrics: None,
        };
        tokio::spawn(reader_pump(
            rx1,
            frame_tx.clone(),
            child1.clone(),
            ctx1,
            local_id,
            Duration::from_secs(5),
        ));

        // Send a frame -- pump should forward it
        let data1 = rmp_serde::to_vec(&crate::streaming::frame::StreamFrame::Item(1u32)).unwrap();
        tx1.send_async(data1.clone()).await.unwrap();
        let received =
            tokio::time::timeout(std::time::Duration::from_millis(500), frame_rx.recv_async())
                .await
                .expect("timeout")
                .expect("closed");
        assert_eq!(received, data1, "first pump must forward data");

        // --- Detach: cancel child, NOT parent ---
        child1.cancel();
        assert!(
            !parent.is_cancelled(),
            "parent must NOT be cancelled by child cancel"
        );

        // Give pump time to exit
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;

        // --- Reattach: new child from the same parent ---
        let (tx2, rx2) = flume::bounded::<Vec<u8>>(256);
        let child2 = parent.child_token();

        // Update the entry (simulates what _anchor_attach does)
        if let Some(mut entry) = registry.get_mut(&local_id) {
            entry.active_pump_token = Some(child2.clone());
            entry.attachment = true;
        }

        let ctx2 = crate::streaming::anchor::AnchorContext {
            registry: registry.clone(),
            mpsc_registry: mpsc_reg.clone(),
            metrics: None,
        };
        tokio::spawn(reader_pump(
            rx2,
            frame_tx.clone(),
            child2.clone(),
            ctx2,
            local_id,
            Duration::from_secs(5),
        ));

        // Send a frame through the new transport -- pump should forward it
        let data2 = rmp_serde::to_vec(&crate::streaming::frame::StreamFrame::Item(2u32)).unwrap();
        tx2.send_async(data2.clone()).await.unwrap();
        let received2 =
            tokio::time::timeout(std::time::Duration::from_millis(500), frame_rx.recv_async())
                .await
                .expect("timeout on reattach")
                .expect("closed on reattach");
        assert_eq!(
            received2, data2,
            "second pump must forward data after reattach"
        );

        // --- Finalize: cancel parent cascades to child ---
        parent.cancel();
        assert!(
            child2.is_cancelled(),
            "child must be cancelled when parent is cancelled"
        );
    }

    // -----------------------------------------------------------------------
    // StreamCancelHandle + SenderRegistry + create_stream_cancel_handler tests (Task 1)
    // -----------------------------------------------------------------------

    #[test]
    fn test_stream_cancel_handle_pack_unpack() {
        let worker_id = velo_ext::WorkerId::from_u64(0xDEAD_BEEF_1234_5678);
        let stream_id: u64 = 0xABCD_EF01_2345_6789;

        let handle = crate::streaming::control::StreamCancelHandle::pack(worker_id, stream_id);
        let (recovered_worker, recovered_stream) = handle.unpack();

        assert_eq!(
            recovered_worker, worker_id,
            "worker_id must round-trip through pack/unpack"
        );
        assert_eq!(
            recovered_stream, stream_id,
            "stream_id must round-trip through pack/unpack"
        );
    }

    #[test]
    fn test_stream_cancel_handle_serde() {
        let worker_id = velo_ext::WorkerId::from_u64(0xCAFE_BABE_0000_0001);
        let stream_id: u64 = 42;

        let handle = crate::streaming::control::StreamCancelHandle::pack(worker_id, stream_id);
        let encoded = rmp_serde::to_vec(&handle).expect("rmp_serde serialize must succeed");
        let decoded: crate::streaming::control::StreamCancelHandle =
            rmp_serde::from_slice(&encoded).expect("rmp_serde deserialize must succeed");

        assert_eq!(
            handle, decoded,
            "StreamCancelHandle must survive rmp_serde round-trip"
        );
        let (w, s) = decoded.unpack();
        assert_eq!(w, worker_id);
        assert_eq!(s, stream_id);
    }

    #[test]
    fn test_stream_cancel_handler_compiles() {
        let registry = std::sync::Arc::new(crate::streaming::control::SenderRegistry::default());
        let _handler = crate::streaming::control::create_stream_cancel_handler(registry);
        // Returns without panic — confirms the handler constructor compiles and runs.
    }

    #[tokio::test]
    async fn test_pump_exits_when_consumer_drops() {
        let (transport_tx, frame_rx, cancel_token, _registry, _id) = make_pump_test_infra();

        // Drop the frame_rx consumer side -- pump's send will fail
        drop(frame_rx);

        // Send data so the pump tries to forward and fails
        let data = rmp_serde::to_vec(&crate::streaming::frame::StreamFrame::Item(1u32)).unwrap();
        let _ = transport_tx.send_async(data).await;

        // Give the pump time to process and exit
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;

        // Pump should have exited and cancelled the token
        assert!(
            cancel_token.is_cancelled(),
            "cancel_token must be cancelled after pump exits due to consumer drop"
        );
    }
}