hyperactor_mesh/

comm.rs

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 * All rights reserved.
 *
 * This source code is licensed under the BSD-style license found in the
 * LICENSE file in the root directory of this source tree.
 */

use crate::casting::CAST_ACTOR_MESH_ID;
use crate::comm::multicast::CAST_ORIGINATING_SENDER;
use crate::comm::multicast::CastEnvelope;
use crate::comm::multicast::CastMessageV1;
use crate::comm::multicast::ForwardMessageV1;
use crate::mesh_id::ActorMeshId;
use crate::resource;
pub mod multicast;

use std::cmp::Ordering;
use std::collections::HashMap;
use std::fmt::Debug;

use anyhow::Result;
use async_trait::async_trait;
use hyperactor::Actor;
use hyperactor::ActorAddr;
use hyperactor::ActorRef;
use hyperactor::Context;
use hyperactor::Endpoint as _;
use hyperactor::Handler;
use hyperactor::Instance;
use hyperactor::OncePortRefRepr;
use hyperactor::PortAddr;
use hyperactor::PortRef;
use hyperactor::PortRefRepr;
use hyperactor::RemoteEndpoint as _;
use hyperactor::RemoteMessage;
use hyperactor::accum::ReducerMode;
use hyperactor::mailbox::MailboxSender;
use hyperactor::mailbox::Undeliverable;
use hyperactor::mailbox::UndeliverableMailboxSender;
use hyperactor::mailbox::UndeliverableMessageError;
use hyperactor::mailbox::monitored_return_handle;
use hyperactor::ordering::SEQ_INFO;
use hyperactor::ordering::SeqInfo;
use hyperactor_config::CONFIG;
use hyperactor_config::ConfigAttr;
use hyperactor_config::Flattrs;
use hyperactor_config::attrs::declare_attrs;
use hyperactor_mesh_macros::sel;
use ndslice::Point;
use ndslice::Selection;
use ndslice::View;
use ndslice::selection::routing::RoutingFrame;
use serde::Deserialize;
use serde::Serialize;
use typeuri::Named;

use crate::comm::multicast::CastMessage;
use crate::comm::multicast::CastMessageEnvelope;
use crate::comm::multicast::ForwardMessage;
use crate::comm::multicast::set_cast_info_on_headers;

declare_attrs! {
    /// Whether to use native v1 casting in v1 ActorMesh.
    @meta(CONFIG = ConfigAttr::new(
        Some("HYPERACTOR_MESH_ENABLE_NATIVE_V1_CASTING".to_string()),
        Some("enable_native_v1_casting".to_string()),
    ))
    pub attr ENABLE_NATIVE_V1_CASTING: bool = true;

    /// The multicast phase that attached context to a delivery failure.
    pub attr MULTICAST_FAILURE_PHASE: String;

    /// The comm actor that attached multicast context to a delivery failure.
    pub attr MULTICAST_FAILURE_COMM_ACTOR: ActorAddr;

    /// The originating cast sender.
    pub attr MULTICAST_FAILURE_ORIGIN: ActorAddr;

    /// The return port used to send the undeliverable message to the origin.
    pub attr MULTICAST_FAILURE_RETURN_PORT: String;
}

fn annotate_multicast_failure(
    envelope: &mut hyperactor::mailbox::MessageEnvelope,
    comm_actor: &ActorAddr,
    phase: &str,
    origin: &ActorAddr,
    return_port: &PortAddr,
) {
    let actor_mesh_id = envelope.headers().get(CAST_ACTOR_MESH_ID);
    if let Some(failure) = envelope.root_delivery_failure_mut() {
        failure
            .attrs
            .set(MULTICAST_FAILURE_PHASE, phase.to_string());
        failure
            .attrs
            .set(MULTICAST_FAILURE_COMM_ACTOR, comm_actor.clone());
        failure.attrs.set(MULTICAST_FAILURE_ORIGIN, origin.clone());
        failure
            .attrs
            .set(MULTICAST_FAILURE_RETURN_PORT, return_port.to_string());
        if let Some(actor_mesh_id) = actor_mesh_id {
            failure.attrs.set(CAST_ACTOR_MESH_ID, actor_mesh_id);
        }
    }
}

/// Parameters to initialize the CommActor
#[derive(Debug, Clone, Serialize, Deserialize, Named, Default)]
pub struct CommActorParams {}
wirevalue::register_type!(CommActorParams);

/// A message buffered due to out-of-order delivery.
#[derive(Debug)]
struct Buffered {
    /// Sequence number of this message.
    seq: usize,
    /// Whether to deliver this message to this comm-actors actors.
    deliver_here: bool,
    /// Peer comm actors to forward message to.
    next_steps: HashMap<usize, Vec<RoutingFrame>>,
    /// The message to deliver.
    message: CastMessageEnvelope,
}

/// Bookkeeping to handle sequence numbers and in-order delivery for messages
/// sent to and through this comm actor.
#[derive(Debug, Default)]
struct ReceiveState {
    /// The sequence of the last received message.
    seq: usize,
    /// A buffer storing messages we received out-of-order, indexed by the seq
    /// that should precede it.
    buffer: HashMap<usize, Buffered>,
    /// A map of the last sequence number we sent to next steps, indexed by rank.
    last_seqs: HashMap<usize, usize>,
}

/// This is the comm actor used for efficient and scalable message multicasting
/// and result accumulation.
#[derive(Debug, Default)]
#[hyperactor::export(
    CommMeshConfig,
    CastMessage,
    ForwardMessage,
    CastMessageV1,
    ForwardMessageV1
)]
#[hyperactor::spawnable]
pub struct CommActor {
    /// Sequence numbers are maintained for each (actor mesh id, sender).
    send_seq: HashMap<(ActorMeshId, ActorAddr), usize>,
    /// Each sender is a unique stream.
    recv_state: HashMap<(ActorMeshId, ActorAddr), ReceiveState>,

    /// The comm actor's mesh configuration, or buffered messages if not yet configured.
    mesh_config: MeshConfigState,
}

#[derive(Debug)]
enum PendingMessage {
    Cast(CastMessage),
    Forward(ForwardMessage),
    ForwardV1(ForwardMessageV1),
}

#[derive(Debug)]
enum MeshConfigState {
    /// Config not yet received; buffer incoming messages until it arrives.
    NotConfigured(Vec<PendingMessage>),
    /// Config received; ready to route messages.
    Configured(CommMeshConfig),
}

impl Default for MeshConfigState {
    fn default() -> Self {
        MeshConfigState::NotConfigured(Vec::new())
    }
}

/// Configuration for how a `CommActor` determines its own rank and locates peers.
#[derive(Debug, Clone, Serialize, Deserialize, Named)]
pub struct CommMeshConfig {
    /// The rank of this comm actor on the root mesh.
    rank: usize,
    /// Key is the rank of the peer on the root mesh. Value is the peer's comm actor.
    peers: HashMap<usize, ActorRef<CommActor>>,
}
wirevalue::register_type!(CommMeshConfig);

impl CommMeshConfig {
    /// Create a new mesh configuration with the given rank and peer mapping.
    pub fn new(rank: usize, peers: HashMap<usize, ActorRef<CommActor>>) -> Self {
        Self { rank, peers }
    }

    /// Return the peer comm actor for the given rank.
    fn peer_for_rank(&self, rank: usize) -> Result<ActorRef<CommActor>> {
        self.peers
            .get(&rank)
            .cloned()
            .ok_or_else(|| anyhow::anyhow!("no peer for rank {}", rank))
    }

    /// Return the rank of the comm actor.
    fn self_rank(&self) -> usize {
        self.rank
    }
}

#[async_trait]
impl Actor for CommActor {
    async fn init(&mut self, this: &Instance<Self>) -> Result<(), anyhow::Error> {
        this.set_system();
        Ok(())
    }

    // This is an override of the default actor behavior.
    async fn handle_undeliverable_message(
        &mut self,
        cx: &Instance<Self>,
        _reason: hyperactor::mailbox::UndeliverableReason,
        undelivered: hyperactor::mailbox::Undeliverable<hyperactor::mailbox::MessageEnvelope>,
    ) -> Result<(), anyhow::Error> {
        self.return_delivery_failure_to_origin(cx, undelivered)
            .await
    }

    async fn handle_invalid_reference(
        &mut self,
        cx: &Instance<Self>,
        _invalid: hyperactor::mailbox::InvalidReference,
        undelivered: hyperactor::mailbox::Undeliverable<hyperactor::mailbox::MessageEnvelope>,
    ) -> Result<(), anyhow::Error> {
        self.return_delivery_failure_to_origin(cx, undelivered)
            .await
    }
}

impl CommActor {
    async fn return_delivery_failure_to_origin(
        &mut self,
        cx: &Instance<Self>,
        undelivered: hyperactor::mailbox::Undeliverable<hyperactor::mailbox::MessageEnvelope>,
    ) -> Result<(), anyhow::Error> {
        let mut message_envelope = match undelivered {
            Undeliverable::Returned(message_envelope) => message_envelope,
            Undeliverable::Report(report) => {
                anyhow::bail!(UndeliverableMessageError::Report { report });
            }
        };

        // 1. Case delivery failure at a "forwarding" step.
        if let Ok(ForwardMessage { message, .. }) =
            message_envelope.deserialized::<ForwardMessage>()
        {
            let sender = message.sender();
            let return_port = PortRef::attest_handler_port(sender);
            annotate_multicast_failure(
                &mut message_envelope,
                cx.self_addr(),
                "forward",
                sender,
                return_port.port_addr(),
            );

            // Needed so that the receiver of the undeliverable message can easily find the
            // original sender of the cast message.
            message_envelope.set_header(CAST_ORIGINATING_SENDER, sender.clone());

            return_port.post(cx, Undeliverable::Returned(message_envelope.clone()));
            return Ok(());
        }

        // 2. Case delivery failure at a "deliver here" step.
        if let Some(sender) = message_envelope.headers().get(CAST_ORIGINATING_SENDER) {
            let return_port = PortRef::attest_handler_port(&sender);
            annotate_multicast_failure(
                &mut message_envelope,
                cx.self_addr(),
                "deliver_here",
                &sender,
                return_port.port_addr(),
            );
            return_port.post(cx, Undeliverable::Returned(message_envelope.clone()));
            return Ok(());
        }

        // 3. A return of an undeliverable message was itself returned.
        UndeliverableMailboxSender
            .post(message_envelope, /*unused */ monitored_return_handle());
        Ok(())
    }

    /// Forward the message to the comm actor on the given peer rank.
    fn forward<M: RemoteMessage>(
        cx: &Context<Self>,
        config: &CommMeshConfig,
        rank: usize,
        message: M,
    ) -> Result<()>
    where
        CommActor: hyperactor::RemoteHandles<M>,
    {
        let child = config.peer_for_rank(rank)?;
        // TEMPORARY: until dropping v0 support
        if let Some(cast_actor_mesh_id) = cx.headers().get(CAST_ACTOR_MESH_ID) {
            let mut headers = Flattrs::new();
            headers.set(CAST_ACTOR_MESH_ID, cast_actor_mesh_id);
            child.post_with_headers(cx, headers, message);
        } else {
            child.post(cx, message);
        }
        Ok(())
    }

    fn handle_message(
        cx: &Context<Self>,
        config: &CommMeshConfig,
        deliver_here: bool,
        next_steps: HashMap<usize, Vec<RoutingFrame>>,
        sender: ActorAddr,
        mut message: CastMessageEnvelope,
        seq: usize,
        last_seqs: &mut HashMap<usize, usize>,
    ) -> Result<()> {
        split_ports(cx, message.data_mut(), deliver_here, &next_steps)?;

        // Deliver message here, if necessary.
        if deliver_here {
            // We should not copy cx.headers() because it contains auto-generated
            // headers from mailbox. We want fresh headers only containing
            // user-provided headers.
            let headers = message.headers().clone();
            Self::deliver_to_dest(cx, headers, &mut message, config)?;
        }

        // Forward to peers.
        next_steps
            .into_iter()
            .map(|(peer, dests)| {
                let last_seq = last_seqs.entry(peer).or_default();
                Self::forward(
                    cx,
                    config,
                    peer,
                    ForwardMessage {
                        dests,
                        sender: sender.clone(),
                        message: message.clone(),
                        seq,
                        last_seq: *last_seq,
                    },
                )?;
                *last_seq = seq;
                Ok(())
            })
            .collect::<Result<Vec<_>>>()?;

        Ok(())
    }

    fn deliver_to_dest<M: CastEnvelope>(
        cx: &Context<Self>,
        mut headers: Flattrs,
        message: &mut M,
        config: &CommMeshConfig,
    ) -> anyhow::Result<()> {
        let cast_point = message.cast_point(config)?;
        // Replace ranks with self ranks.
        replace_with_self_ranks(&cast_point, message.data_mut())?;

        set_cast_info_on_headers(&mut headers, cast_point, message.sender().clone());

        // Bind dest ONCE so we can pass to both the stamp helper and the
        // post call.
        let dest = cx
            .self_addr()
            .proc_addr()
            .actor_addr_uid(message.dest_port().actor_uid().clone())
            .port_addr(hyperactor::Port::handler_id(
                message.dest_port().port(),
                None,
            ));

        // Stamp SENDER_ACTOR_ID when headers already carry SEQ_INFO (V1
        // path). V0 path leaves SEQ_INFO absent here; MailboxExt::post will
        // assign it and stamp via its own helper call later. Flattrs::get
        // returns an owned typed value, so seq_info isn't borrowed from
        // headers and we can pass &mut headers to the helper without
        // a borrow-checker conflict.
        if let Some(seq_info) = headers.get(SEQ_INFO) {
            hyperactor::mailbox::headers::stamp_sender_actor_id(
                &mut headers,
                &seq_info,
                &dest,
                message.sender(),
            );
        }

        cx.post_with_external_seq_info(dest, headers, message.data().clone().erase_encoding());

        Ok(())
    }
}

// Split ports, if any, and update message with new ports. In this
// way, children actors will reply to this comm actor's ports, instead
// of to the original ports provided by parent.
fn split_ports(
    cx: &Context<CommActor>,
    data: &mut wirevalue::Any<wirevalue::encoding::Multipart>,
    deliver_here: bool,
    next_steps: &HashMap<usize, Vec<RoutingFrame>>,
) -> anyhow::Result<()> {
    // Split ports, if any, and update message with new ports. In this
    // way, children actors will reply to this comm actor's ports, instead
    // of to the original ports provided by parent.
    data.visit_multipart_parts_mut::<PortRefRepr, anyhow::Error>(|port| {
        if port.unsplit() {
            return Ok(());
        }

        let split = port.port_addr().split(
            cx,
            port.reducer_spec().clone(),
            ReducerMode::Streaming(port.streaming_opts().clone()),
            port.get_return_undeliverable(),
        )?;

        #[cfg(test)]
        tests::collect_split_port(port.port_addr(), &split, deliver_here);

        port.update_port_addr(split);
        Ok(())
    })?;

    data.visit_multipart_parts_mut::<OncePortRefRepr, anyhow::Error>(|port| {
        if port.unsplit() || port.reducer_spec().is_none() {
            // We can only split OncePorts that have reducers. Pass this
            // through; if it is used multiple times, it will cause a delivery
            // error downstream.
            return Ok(());
        }

        let peer_count = next_steps.len() + if deliver_here { 1 } else { 0 };
        let split = port.port_addr().split(
            cx,
            port.reducer_spec().clone(),
            ReducerMode::Once(peer_count),
            true,
        )?;

        #[cfg(test)]
        tests::collect_split_port(port.port_addr(), &split, deliver_here);

        port.update_port_addr(split);
        Ok(())
    })?;

    Ok(())
}

fn replace_with_self_ranks(
    cast_point: &Point,
    data: &mut wirevalue::Any<wirevalue::encoding::Multipart>,
) -> anyhow::Result<()> {
    data.visit_multipart_parts_mut::<resource::RankRepr, anyhow::Error>(
        |resource::RankRepr(rank)| {
            *rank = Some(cast_point.rank());
            Ok(())
        },
    )
}

#[async_trait]
impl Handler<CommMeshConfig> for CommActor {
    async fn handle(&mut self, cx: &Context<Self>, config: CommMeshConfig) -> Result<()> {
        let pending =
            match std::mem::replace(&mut self.mesh_config, MeshConfigState::Configured(config)) {
                MeshConfigState::NotConfigured(pending) => pending,
                MeshConfigState::Configured(_) => Vec::new(),
            };
        if !pending.is_empty() {
            tracing::info!(
                count = pending.len(),
                "replaying buffered pre-config messages"
            );
        }
        for msg in pending {
            match msg {
                PendingMessage::Cast(m) => self.handle(cx, m).await?,
                PendingMessage::Forward(m) => self.handle(cx, m).await?,
                PendingMessage::ForwardV1(m) => self.handle(cx, m).await?,
            }
        }
        Ok(())
    }
}

// TODO(T218630526): reliable casting for mutable topology
#[async_trait]
impl Handler<CastMessage> for CommActor {
    #[tracing::instrument(level = "debug", skip_all)]
    async fn handle(&mut self, cx: &Context<Self>, cast_message: CastMessage) -> Result<()> {
        let config = match &mut self.mesh_config {
            MeshConfigState::NotConfigured(pending) => {
                pending.push(PendingMessage::Cast(cast_message));
                return Ok(());
            }
            MeshConfigState::Configured(config) => config,
        };
        // Always forward the message to the root rank of the slice, casting starts from there.
        let slice = cast_message.dest.slice.clone();
        let selection = cast_message.dest.selection.clone();
        let frame = RoutingFrame::root(selection, slice);
        let rank = frame.slice.location(&frame.here)?;
        let seq = self
            .send_seq
            .entry(cast_message.message.stream_key())
            .or_default();
        let last_seq = *seq;
        *seq += 1;

        let fwd_message = ForwardMessage {
            dests: vec![frame],
            sender: cx.self_addr().clone(),
            message: cast_message.message,
            seq: *seq,
            last_seq,
        };

        // Optimization: if forwarding to ourselves, handle inline instead of
        // going through the message queue
        if config.self_rank() == rank {
            Handler::<ForwardMessage>::handle(self, cx, fwd_message).await?;
        } else {
            Self::forward(cx, config, rank, fwd_message)?;
        }
        Ok(())
    }
}

#[async_trait]
impl Handler<ForwardMessage> for CommActor {
    #[tracing::instrument(level = "debug", skip_all)]
    async fn handle(&mut self, cx: &Context<Self>, fwd_message: ForwardMessage) -> Result<()> {
        let config = match &mut self.mesh_config {
            MeshConfigState::NotConfigured(pending) => {
                pending.push(PendingMessage::Forward(fwd_message));
                return Ok(());
            }
            MeshConfigState::Configured(config) => config,
        };

        let ForwardMessage {
            sender,
            dests,
            message,
            seq,
            last_seq,
        } = fwd_message;

        // Resolve/dedup routing frames.
        let rank = config.self_rank();
        let (deliver_here, next_steps) =
            ndslice::selection::routing::resolve_routing(rank, dests, &mut |_| {
                panic!("Choice encountered in CommActor routing")
            })?;

        let recv_state = self.recv_state.entry(message.stream_key()).or_default();
        match recv_state.seq.cmp(&last_seq) {
            // We got the expected next message to deliver to this host.
            Ordering::Equal => {
                // We got an in-order operation, so handle it now.
                Self::handle_message(
                    cx,
                    config,
                    deliver_here,
                    next_steps,
                    sender.clone(),
                    message,
                    seq,
                    &mut recv_state.last_seqs,
                )?;
                recv_state.seq = seq;

                // Also deliver any pending operations from the recv buffer that
                // were received out-of-order that are now unblocked.
                while let Some(Buffered {
                    seq,
                    deliver_here,
                    next_steps,
                    message,
                }) = recv_state.buffer.remove(&recv_state.seq)
                {
                    Self::handle_message(
                        cx,
                        config,
                        deliver_here,
                        next_steps,
                        sender.clone(),
                        message,
                        seq,
                        &mut recv_state.last_seqs,
                    )?;
                    recv_state.seq = seq;
                }
            }
            // We got an out-of-order operation, so buffer it for now, until we
            // recieved the onces sequenced before it.
            Ordering::Less => {
                tracing::warn!(
                    "buffering out-of-order message with seq {} (last {}), expected {}: {:?}",
                    seq,
                    last_seq,
                    recv_state.seq,
                    message
                );
                recv_state.buffer.insert(
                    last_seq,
                    Buffered {
                        seq,
                        deliver_here,
                        next_steps,
                        message,
                    },
                );
            }
            // We already got this message -- just drop it.
            Ordering::Greater => {
                tracing::warn!("received duplicate message with seq {}: {:?}", seq, message);
            }
        }

        Ok(())
    }
}

#[async_trait]
impl Handler<CastMessageV1> for CommActor {
    async fn handle(&mut self, cx: &Context<Self>, cast_message: CastMessageV1) -> Result<()> {
        let slice = cast_message.dest_region.slice().clone();
        let frame = RoutingFrame::root(sel!(*), slice);
        let forward_message = ForwardMessageV1 {
            dests: vec![frame],
            message: cast_message,
        };
        self.handle(cx, forward_message).await
    }
}

#[async_trait]
impl Handler<ForwardMessageV1> for CommActor {
    async fn handle(&mut self, cx: &Context<Self>, fwd_message: ForwardMessageV1) -> Result<()> {
        let config = match &mut self.mesh_config {
            MeshConfigState::NotConfigured(pending) => {
                pending.push(PendingMessage::ForwardV1(fwd_message));
                return Ok(());
            }
            MeshConfigState::Configured(config) => config,
        };

        let ForwardMessageV1 { dests, mut message } = fwd_message;
        // Resolve/dedup routing frames.
        let rank_on_root_mesh = config.self_rank();
        let (deliver_here, next_steps) =
            ndslice::selection::routing::resolve_routing(rank_on_root_mesh, dests, &mut |_| {
                panic!("choice encountered in CommActor routing")
            })?;

        split_ports(cx, &mut message.data, deliver_here, &next_steps)?;

        // Deliver message here, if necessary.
        if deliver_here {
            let mut headers = message.headers().clone();
            let seq = message
                .seqs
                .get(message.cast_point(config)?.rank())
                .expect("mismatched seqs and dest_region");
            headers.set(
                SEQ_INFO,
                SeqInfo::Session {
                    session_id: message.session_id,
                    seq,
                },
            );
            Self::deliver_to_dest(cx, headers, &mut message, config)?;
        }

        // Forward to peers.
        for (peer_rank_on_root_mesh, dests) in next_steps {
            let forward_message = ForwardMessageV1 {
                dests,
                message: message.clone(),
            };
            Self::forward(cx, config, peer_rank_on_root_mesh, forward_message)?;
        }

        Ok(())
    }
}

pub mod test_utils {
    use anyhow::Result;
    use async_trait::async_trait;
    use hyperactor::Actor;
    use hyperactor::ActorAddr;
    use hyperactor::Context;
    use hyperactor::Handler;
    use hyperactor::PortRef;
    use serde::Deserialize;
    use serde::Serialize;
    use typeuri::Named;

    use super::*;

    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Named)]
    pub struct MyReply {
        pub sender: ActorAddr,
        pub value: u64,
    }

    #[derive(Debug, Named, Serialize, Deserialize, PartialEq, Clone)]
    #[expect(
        clippy::large_enum_variant,
        reason = "test fixture; CastAndReply carries PortRefs and boxing fields ripples into handler assertions"
    )]
    pub enum TestMessage {
        Forward(String),
        CastAndReply {
            arg: String,
            // Intentionally unsplit so this port should pass through unchanged.
            reply_to0: PortRef<String>,
            reply_to1: PortRef<u64>,
            reply_to2: PortRef<MyReply>,
        },
        CastAndReplyOnce {
            arg: String,
            reply_to: hyperactor::OncePortRef<u64>,
        },
        CastWithUnsplitPort {
            reply_to: PortRef<u64>,
        },
    }

    #[derive(Debug)]
    #[hyperactor::export(TestMessage)]
    #[hyperactor::spawnable]
    pub struct TestActor {
        // Forward the received message to this port, so it can be inspected by
        // the unit test.
        forward_port: PortRef<TestMessage>,
    }

    #[derive(Debug, Clone, Named, Serialize, Deserialize)]
    pub struct TestActorParams {
        pub forward_port: PortRef<TestMessage>,
    }

    #[async_trait]
    impl Actor for TestActor {}

    #[async_trait]
    impl hyperactor::RemoteSpawn for TestActor {
        type Params = TestActorParams;

        async fn new(params: Self::Params, _environment: Flattrs) -> Result<Self> {
            let Self::Params { forward_port } = params;
            Ok(Self { forward_port })
        }
    }

    #[async_trait]
    impl Handler<TestMessage> for TestActor {
        async fn handle(&mut self, cx: &Context<Self>, msg: TestMessage) -> anyhow::Result<()> {
            // For CastWithUnsplitPort, send a reply so the test can
            // verify that the unsplit port is still directly reachable.
            if let TestMessage::CastWithUnsplitPort { ref reply_to } = msg {
                reply_to.post(cx, 42);
            }
            self.forward_port.post(cx, msg);
            Ok(())
        }
    }

    // SENDER_ACTOR_ID capture fixture for the cast-correctness matrix tests.
    // Kept in test_utils so #[hyperactor::spawnable]'s inventory::submit!
    // registration is reliably picked up by Remote::collect() at runtime;
    // #[doc(hidden)] because this is a test fixture, not API.

    #[doc(hidden)]
    #[derive(Debug, Clone, Serialize, Deserialize, Named)]
    pub struct SenderCaptureMsg();

    #[doc(hidden)]
    #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Named)]
    pub struct CapturedSender(pub Option<ActorAddr>);

    #[doc(hidden)]
    #[derive(Debug)]
    #[hyperactor::export(SenderCaptureMsg)]
    #[hyperactor::spawnable]
    pub struct SenderCapturingActor {
        forward_port: PortRef<CapturedSender>,
    }

    #[doc(hidden)]
    #[derive(Debug, Clone, Named, Serialize, Deserialize)]
    pub struct SenderCapturingActorParams {
        pub forward_port: PortRef<CapturedSender>,
    }

    #[async_trait]
    impl Actor for SenderCapturingActor {}

    #[async_trait]
    impl hyperactor::RemoteSpawn for SenderCapturingActor {
        type Params = SenderCapturingActorParams;

        async fn new(params: Self::Params, _environment: Flattrs) -> Result<Self> {
            let Self::Params { forward_port } = params;
            Ok(Self { forward_port })
        }
    }

    #[async_trait]
    impl Handler<SenderCaptureMsg> for SenderCapturingActor {
        async fn handle(&mut self, cx: &Context<Self>, _: SenderCaptureMsg) -> anyhow::Result<()> {
            let sender = cx
                .headers()
                .get(hyperactor::mailbox::headers::SENDER_ACTOR_ID);
            self.forward_port.post(cx, CapturedSender(sender));
            Ok(())
        }
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;
    use std::collections::HashSet;
    use std::fmt::Display;
    use std::hash::Hash;
    use std::ops::Deref;
    use std::ops::DerefMut;
    use std::sync::Mutex;
    use std::sync::OnceLock;

    use hyperactor::accum;

    /// Common setup for pre-config buffering tests: a single proc with a
    /// TestActor (for observing delivery) and an unconfigured CommActor.
    /// Returns (client, rx, comm_handle, actor_mesh_id) plus handles
    /// that must be kept alive.
    async fn buffering_fixture(
        proc_name: &str,
    ) -> (
        hyperactor::Client,
        hyperactor::mailbox::PortReceiver<TestMessage>,
        hyperactor::ActorHandle<CommActor>,
        crate::mesh_id::ActorMeshId,
        // Drop guards: test actor handle, test actor ref.
        (hyperactor::ActorHandle<TestActor>, ActorRef<TestActor>),
    ) {
        use hyperactor::Proc;
        use hyperactor::RemoteSpawn;
        use hyperactor::channel::ChannelTransport;
        use hyperactor::id::Label;

        let proc = Proc::direct(ChannelTransport::Unix.any(), proc_name.to_string()).unwrap();
        let client = proc.client("client");

        let actor_mesh_id = crate::mesh_id::ActorMeshId::instance(Label::new("test").unwrap());

        let (tx, rx) = open_port(&client);
        let forward_port = tx.bind();
        let test_actor = TestActor::new(TestActorParams { forward_port }, Default::default())
            .await
            .unwrap();
        let test_handle = proc
            .spawn_with_uid(actor_mesh_id.uid().clone(), test_actor)
            .unwrap();
        let test_ref: ActorRef<TestActor> = test_handle.bind::<TestActor>();

        let comm_handle = proc.spawn(CommActor::default());

        (
            client,
            rx,
            comm_handle,
            actor_mesh_id,
            (test_handle, test_ref),
        )
    }

    /// Send CommMeshConfig (single-rank mesh pointing at self).
    fn send_config(client: &hyperactor::Client, comm_handle: &hyperactor::ActorHandle<CommActor>) {
        let comm_ref = comm_handle.bind::<CommActor>();
        let mut peers = HashMap::new();
        peers.insert(0, comm_ref);
        comm_handle.post(client, CommMeshConfig::new(0, peers));
    }

    /// Send a message before config, send config, send another after config,
    /// and verify both are delivered in order.
    async fn assert_buffered_and_replayed<M: hyperactor::Message>(
        proc_name: &str,
        mut make_msg: impl FnMut(&hyperactor::Client, &crate::mesh_id::ActorMeshId, &str) -> M,
    ) where
        CommActor: hyperactor::Handler<M>,
    {
        let (client, mut rx, comm_handle, actor_mesh_id, _guards) =
            buffering_fixture(proc_name).await;

        comm_handle.post(&client, make_msg(&client, &actor_mesh_id, "buffered"));
        send_config(&client, &comm_handle);
        comm_handle.post(&client, make_msg(&client, &actor_mesh_id, "direct"));

        assert_eq!(
            rx.recv().await.unwrap(),
            TestMessage::Forward("buffered".to_string()),
        );
        assert_eq!(
            rx.recv().await.unwrap(),
            TestMessage::Forward("direct".to_string()),
        );
        comm_handle.drain_and_stop("test done").ok();
    }

    #[async_timed_test(timeout_secs = 1)]
    async fn cast_before_config_is_buffered_and_replayed() {
        use ndslice::Slice;

        assert_buffered_and_replayed("test_cast", |client, actor_mesh_id, payload| {
            let actor_mesh_id = actor_mesh_id.clone();
            let slice = Slice::new_row_major(vec![1]);
            let shape = ndslice::Shape::new(vec!["rank".to_string()], slice.clone()).unwrap();
            let envelope = multicast::CastMessageEnvelope::new::<TestActor, TestMessage>(
                actor_mesh_id,
                client.self_addr().clone(),
                shape,
                hyperactor_config::Flattrs::new(),
                TestMessage::Forward(payload.to_string()),
            )
            .unwrap();
            multicast::CastMessage {
                dest: multicast::Uslice {
                    slice,
                    selection: sel!(*),
                },
                message: envelope,
            }
        })
        .await;
    }

    #[async_timed_test(timeout_secs = 1)]
    async fn forward_before_config_is_buffered_and_replayed() {
        use ndslice::Slice;
        use ndslice::selection::routing::RoutingFrame;

        let mut next_seq: usize = 0;
        assert_buffered_and_replayed("test_fwd", move |client, actor_mesh_id, payload| {
            let actor_mesh_id = actor_mesh_id.clone();
            let slice = Slice::new_row_major(vec![1]);
            let shape = ndslice::Shape::new(vec!["rank".to_string()], slice.clone()).unwrap();
            let envelope = multicast::CastMessageEnvelope::new::<TestActor, TestMessage>(
                actor_mesh_id,
                client.self_addr().clone(),
                shape,
                hyperactor_config::Flattrs::new(),
                TestMessage::Forward(payload.to_string()),
            )
            .unwrap();
            let frame = RoutingFrame::root(sel!(*), slice);
            let last_seq = next_seq;
            next_seq += 1;
            multicast::ForwardMessage {
                sender: client.self_addr().clone(),
                dests: vec![frame],
                seq: next_seq,
                last_seq,
                message: envelope,
            }
        })
        .await;
    }

    #[async_timed_test(timeout_secs = 1)]
    async fn forward_v1_before_config_is_buffered_and_replayed() {
        use ndslice::Region;
        use ndslice::Slice;
        use ndslice::selection::routing::RoutingFrame;

        assert_buffered_and_replayed("test_fwd_v1", |client, actor_mesh_id, payload| {
            let slice = Slice::new_row_major(vec![1]);
            let region = Region::new(vec!["rank".to_string()], slice.clone());
            let cast_msg = multicast::CastMessageV1::new::<TestActor, TestMessage>(
                client.self_addr().clone(),
                actor_mesh_id,
                region.clone(),
                hyperactor_config::Flattrs::new(),
                TestMessage::Forward(payload.to_string()),
                uuid::Uuid::new_v4(),
                crate::ValueMesh::from_single(region, 1u64),
            )
            .unwrap();
            let frame = RoutingFrame::root(sel!(*), slice);
            multicast::ForwardMessageV1 {
                dests: vec![frame],
                message: cast_msg,
            }
        })
        .await;
    }

    use hyperactor::ActorAddr;
    use hyperactor::ActorRef;
    use hyperactor::Endpoint as _;
    use hyperactor::Index;
    use hyperactor::OncePortRef;
    use hyperactor::PortAddr;
    use hyperactor::PortRef;
    use hyperactor::ProcAddr;
    use hyperactor::accum::Accumulator;
    use hyperactor::accum::ReducerSpec;
    use hyperactor::channel::ChannelAddr;
    use hyperactor::context;
    use hyperactor::context::Mailbox;
    use hyperactor::mailbox::DeliveryFailure;
    use hyperactor::mailbox::MessageEnvelope;
    use hyperactor::mailbox::PortReceiver;
    use hyperactor::mailbox::TransportFailure;
    use hyperactor::mailbox::TransportFailureReason;
    use hyperactor::mailbox::UndeliverableReason;
    use hyperactor::mailbox::open_port;
    use hyperactor::port::Port;
    use hyperactor_config;
    use hyperactor_mesh_macros::sel;
    use maplit::btreemap;
    use maplit::hashmap;
    use ndslice::Extent;
    use ndslice::Selection;
    use ndslice::ViewExt as _;
    use ndslice::extent;
    use ndslice::selection::test_utils::collect_commactor_routing_tree;
    use test_utils::*;
    use timed_test::async_timed_test;
    use tokio::time::Duration;

    use super::*;
    use crate::ActorMesh;
    use crate::ProcMesh;
    use crate::host_mesh::HostMesh;
    use crate::test_utils::local_host_mesh;
    use crate::testing;

    #[test]
    fn annotate_multicast_failure_adds_attrs_to_root_failure() {
        let proc_addr = ProcAddr::singleton(ChannelAddr::Local(1), "test");
        let origin = proc_addr.actor_addr("origin");
        let comm_actor = proc_addr.actor_addr("comm");
        let dest = proc_addr.actor_addr("dest").port_addr(Port::from(42));
        let return_port = origin.port_addr(Port::from(7));
        let mut envelope =
            MessageEnvelope::serialize(origin.clone(), dest.clone(), &(), Flattrs::new()).unwrap();
        envelope.push_delivery_failure(DeliveryFailure::new(UndeliverableReason::Transport(
            TransportFailure::new(dest, TransportFailureReason::NoRoute),
        )));

        annotate_multicast_failure(
            &mut envelope,
            &comm_actor,
            "deliver_here",
            &origin,
            &return_port,
        );

        let root_failure = envelope
            .root_delivery_failure()
            .expect("expected root delivery failure");
        assert_eq!(
            root_failure.attrs.get(MULTICAST_FAILURE_PHASE).as_deref(),
            Some("deliver_here")
        );
        assert_eq!(
            root_failure.attrs.get(MULTICAST_FAILURE_COMM_ACTOR),
            Some(comm_actor)
        );
        assert_eq!(
            root_failure.attrs.get(MULTICAST_FAILURE_ORIGIN),
            Some(origin)
        );
        assert_eq!(
            root_failure.attrs.get(MULTICAST_FAILURE_RETURN_PORT),
            Some(return_port.to_string())
        );
    }

    #[test]
    fn annotate_multicast_failure_records_forward_phase() {
        let proc_addr = ProcAddr::singleton(ChannelAddr::Local(1), "test");
        let origin = proc_addr.actor_addr("origin");
        let comm_actor = proc_addr.actor_addr("comm");
        let dest = proc_addr.actor_addr("dest").port_addr(Port::from(42));
        let return_port = origin.port_addr(Port::from(7));
        let mut envelope =
            MessageEnvelope::serialize(origin.clone(), dest.clone(), &(), Flattrs::new()).unwrap();
        envelope.push_delivery_failure(DeliveryFailure::new(UndeliverableReason::Transport(
            TransportFailure::new(dest, TransportFailureReason::NoRoute),
        )));

        annotate_multicast_failure(&mut envelope, &comm_actor, "forward", &origin, &return_port);

        let root_failure = envelope
            .root_delivery_failure()
            .expect("expected root delivery failure");
        assert_eq!(
            root_failure.attrs.get(MULTICAST_FAILURE_PHASE).as_deref(),
            Some("forward")
        );
    }

    #[test]
    fn annotate_multicast_failure_copies_actor_mesh_id_attr() {
        let proc_addr = ProcAddr::singleton(ChannelAddr::Local(1), "test");
        let origin = proc_addr.actor_addr("origin");
        let comm_actor = proc_addr.actor_addr("comm");
        let dest = proc_addr.actor_addr("dest").port_addr(Port::from(42));
        let return_port = origin.port_addr(Port::from(7));
        let actor_mesh_id =
            crate::mesh_id::ActorMeshId::instance(hyperactor::id::Label::new("mesh").unwrap());
        let mut headers = Flattrs::new();
        headers.set(CAST_ACTOR_MESH_ID, actor_mesh_id.clone());
        let mut envelope =
            MessageEnvelope::serialize(origin.clone(), dest.clone(), &(), headers).unwrap();
        envelope.push_delivery_failure(DeliveryFailure::new(UndeliverableReason::Transport(
            TransportFailure::new(dest, TransportFailureReason::NoRoute),
        )));

        annotate_multicast_failure(
            &mut envelope,
            &comm_actor,
            "deliver_here",
            &origin,
            &return_port,
        );

        let root_failure = envelope
            .root_delivery_failure()
            .expect("expected root delivery failure");
        assert_eq!(
            root_failure.attrs.get(CAST_ACTOR_MESH_ID),
            Some(actor_mesh_id)
        );
    }

    #[test]
    fn annotate_multicast_failure_noops_without_root_failure() {
        let proc_addr = ProcAddr::singleton(ChannelAddr::Local(1), "test");
        let origin = proc_addr.actor_addr("origin");
        let comm_actor = proc_addr.actor_addr("comm");
        let dest = proc_addr.actor_addr("dest").port_addr(Port::from(42));
        let return_port = origin.port_addr(Port::from(7));
        let mut envelope =
            MessageEnvelope::serialize(origin.clone(), dest, &(), Flattrs::new()).unwrap();

        annotate_multicast_failure(
            &mut envelope,
            &comm_actor,
            "deliver_here",
            &origin,
            &return_port,
        );

        assert!(envelope.root_delivery_failure().is_none());
    }

    struct Edge<T> {
        from: T,
        to: T,
        is_leaf: bool,
    }

    impl<T> From<(T, T, bool)> for Edge<T> {
        fn from((from, to, is_leaf): (T, T, bool)) -> Self {
            Self { from, to, is_leaf }
        }
    }

    // The relationship between original ports and split ports. The elements in
    // the tuple are (original port, split port, deliver_here).
    static SPLIT_PORT_TREE: OnceLock<Mutex<Vec<Edge<PortAddr>>>> = OnceLock::new();

    // Collect the relationships between original ports and split ports into
    // SPLIT_PORT_TREE. This is used by tests to verify that ports are split as expected.
    pub(crate) fn collect_split_port(original: &PortAddr, split: &PortAddr, deliver_here: bool) {
        let mutex = SPLIT_PORT_TREE.get_or_init(|| Mutex::new(vec![]));
        let mut tree = mutex.lock().unwrap();

        tree.deref_mut().push(Edge {
            from: original.clone(),
            to: split.clone(),
            is_leaf: deliver_here,
        });
    }

    // There could be other cast calls before the one we want to check, e.g. from
    // allocating the proc mesh, or spawning the actor mesh. Clear the collected
    // tree so it will only contain the cast we want to check.
    fn clear_collected_tree() {
        if let Some(tree) = SPLIT_PORT_TREE.get() {
            let mut tree: std::sync::MutexGuard<'_, Vec<Edge<PortAddr>>> = tree.lock().unwrap();
            tree.clear();
        }
    }

    // A representation of a tree.
    //   * Map's keys are the tree's leafs;
    //   * Map's values are the path from the root to that leaf.
    #[derive(PartialEq)]
    struct PathToLeaves<T>(BTreeMap<T, Vec<T>>);

    // Add a custom Debug trait impl so the result from assert_eq! is readable.
    impl<T: Display> Debug for PathToLeaves<T> {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            fn vec_to_string<T: Display>(v: &[T]) -> String {
                v.iter()
                    .map(ToString::to_string)
                    .collect::<Vec<String>>()
                    .join(", ")
            }

            for (src, path) in &self.0 {
                writeln!(f, "{} -> {}", src, vec_to_string(path))?;
            }
            Ok(())
        }
    }

    fn build_paths<T: Clone + Eq + Hash + Ord>(edges: &[Edge<T>]) -> PathToLeaves<T> {
        let mut child_parent_map = HashMap::new();
        let mut all_nodes = HashSet::new();
        let mut parents = HashSet::new();
        let mut children = HashSet::new();
        let mut dests = HashSet::new();

        // Build parent map and track all nodes and children
        for Edge { from, to, is_leaf } in edges {
            child_parent_map.insert(to.clone(), from.clone());
            all_nodes.insert(from.clone());
            all_nodes.insert(to.clone());
            parents.insert(from.clone());
            children.insert(to.clone());
            if *is_leaf {
                dests.insert(to.clone());
            }
        }

        // For each leaf, reconstruct path back to root
        let mut result = BTreeMap::new();
        for dest in dests {
            let mut path = vec![dest.clone()];
            let mut current = dest.clone();
            while let Some(parent) = child_parent_map.get(&current) {
                path.push(parent.clone());
                current = parent.clone();
            }
            path.reverse();
            result.insert(dest, path);
        }

        PathToLeaves(result)
    }

    #[test]
    fn test_build_paths() {
        // Given the tree:
        //     0
        //    / \
        //   1   4
        //  / \   \
        // 2   3   5
        let edges: Vec<_> = [
            (0, 1, false),
            (1, 2, true),
            (1, 3, true),
            (0, 4, true),
            (4, 5, true),
        ]
        .into_iter()
        .map(|(from, to, is_leaf)| Edge { from, to, is_leaf })
        .collect();

        let paths = build_paths(&edges);

        let expected = btreemap! {
            2 => vec![0, 1, 2],
            3 => vec![0, 1, 3],
            4 => vec![0, 4],
            5 => vec![0, 4, 5],
        };

        assert_eq!(paths.0, expected);
    }

    struct NoneAccumulator;

    impl Accumulator for NoneAccumulator {
        type State = u64;
        type Update = u64;

        fn accumulate(
            &self,
            _state: &mut Self::State,
            _update: Self::Update,
        ) -> anyhow::Result<()> {
            unimplemented!()
        }

        fn reducer_spec(&self) -> Option<ReducerSpec> {
            unimplemented!()
        }
    }

    async fn execute_cast_and_reply(
        ranks: Vec<ActorRef<TestActor>>,
        instance: &impl context::Actor,
        mut reply1_rx: PortReceiver<u64>,
        mut reply2_rx: PortReceiver<MyReply>,
        reply_tos: Vec<(PortRef<u64>, PortRef<MyReply>)>,
    ) {
        // Reply from each dest actor. The replies should be received by client.
        {
            for (rank, (dest_actor, (reply_to1, reply_to2))) in
                ranks.iter().zip(reply_tos.iter()).enumerate()
            {
                let rank_u64 = rank as u64;
                reply_to1.post(instance, rank_u64);
                let my_reply = MyReply {
                    sender: dest_actor.actor_addr().clone(),
                    value: rank_u64,
                };
                reply_to2.post(instance, my_reply.clone());

                assert_eq!(reply1_rx.recv().await.unwrap(), rank_u64);
                assert_eq!(reply2_rx.recv().await.unwrap(), my_reply);
            }
        }

        tracing::info!("the 1st updates from all dest actors were receivered by client");

        // Now send multiple replies from the dest actors. They should all be
        // received by client. Replies sent from the same dest actor should
        // be received in the same order as they were sent out.
        {
            let n = 100;
            let mut expected2: HashMap<ActorAddr, Vec<MyReply>> = hashmap! {};
            for (i, (dest_actor, (_reply_to1, reply_to2))) in
                ranks.iter().zip(reply_tos.iter()).enumerate()
            {
                let mut sent2 = vec![];
                for j in 0..n {
                    let value = (i * 100 + j) as u64;
                    let my_reply = MyReply {
                        sender: dest_actor.actor_addr().clone(),
                        value,
                    };
                    reply_to2.post(instance, my_reply.clone());
                    sent2.push(my_reply);
                }
                assert!(
                    expected2
                        .insert(dest_actor.actor_addr().clone(), sent2)
                        .is_none(),
                    "duplicate actor_id {} in map",
                    dest_actor.actor_addr()
                );
            }

            let mut received2: HashMap<ActorAddr, Vec<MyReply>> = hashmap! {};

            for _ in 0..(n * ranks.len()) {
                let my_reply = reply2_rx.recv().await.unwrap();
                received2
                    .entry(my_reply.sender.clone())
                    .or_default()
                    .push(my_reply);
            }
            assert_eq!(received2, expected2);
        }
    }

    async fn wait_for_with_timeout(
        receiver: &mut PortReceiver<u64>,
        expected: u64,
        dur: Duration,
    ) -> anyhow::Result<()> {
        // timeout wraps the entire async block
        tokio::time::timeout(dur, async {
            loop {
                let msg = receiver.recv().await.unwrap();
                if msg == expected {
                    break;
                }
            }
        })
        .await?;
        Ok(())
    }

    async fn execute_cast_and_accum(
        ranks: Vec<ActorRef<TestActor>>,
        instance: &impl context::Actor,
        mut reply1_rx: PortReceiver<u64>,
        reply_tos: Vec<(PortRef<u64>, PortRef<MyReply>)>,
    ) {
        // Now send multiple replies from the dest actors. They should all be
        // received by client. Replies sent from the same dest actor should
        // be received in the same order as they were sent out.
        let mut sum = 0;
        let n = 100;
        for (i, (_dest_actor, (reply_to1, _reply_to2))) in
            ranks.iter().zip(reply_tos.iter()).enumerate()
        {
            for j in 0..n {
                let value = (i + j) as u64;
                reply_to1.post(instance, value);
                sum += value;
            }
        }
        wait_for_with_timeout(&mut reply1_rx, sum, Duration::from_secs(8))
            .await
            .unwrap();
        // no more messages
        tokio::time::sleep(Duration::from_secs(2)).await;
        let msg = reply1_rx.try_recv().unwrap();
        assert_eq!(msg, None);
    }

    struct MeshSetupV1 {
        instance: &'static Instance<testing::TestRootClient>,
        actor_mesh_ref: crate::ActorMeshRef<TestActor>,
        reply1_rx: PortReceiver<u64>,
        reply2_rx: PortReceiver<MyReply>,
        reply_tos: Vec<(PortRef<u64>, PortRef<MyReply>)>,
        // Keep the host mesh alive so comm actors aren't shut down.
        host_mesh: HostMesh,
    }

    async fn setup_mesh_v1<A>(accum: Option<A>) -> MeshSetupV1
    where
        A: Accumulator<Update = u64, State = u64> + Send + Sync + 'static,
    {
        let instance = crate::testing::instance();
        // We have to use a in process host mesh, because SPLIT_PORT_TREE only
        // can collect paths from the same process.
        let host_mesh = local_host_mesh(8).await;
        let proc_mesh = host_mesh
            .spawn(instance, "test", extent!(gpu = 8), None, None)
            .await
            .unwrap();

        let (tx, mut rx) = hyperactor::mailbox::open_port(instance);
        let params = TestActorParams {
            forward_port: tx.bind(),
        };
        let actor_name =
            crate::mesh_id::ActorMeshId::instance(hyperactor::id::Label::new("test").unwrap());
        // Make this actor a "system" actor to avoid spawning a controller actor.
        // This test is verifying the whole comm tree, so we want fewer actors
        // involved.
        let actor_mesh = proc_mesh
            .spawn_with_name(&instance, actor_name, &params, None, true)
            .await
            .unwrap();
        let actor_mesh_ref: crate::ActorMeshRef<TestActor> = actor_mesh.deref().clone();

        let (reply_port_handle0, _) = open_port::<String>(instance);
        let reply_port_ref0 = reply_port_handle0.bind();
        let (reply_port_handle1, reply1_rx) = match accum {
            Some(a) => instance.mailbox().open_accum_port(a),
            None => open_port(instance),
        };
        let reply_port_ref1 = reply_port_handle1.bind();
        let (reply_port_handle2, reply2_rx) = open_port::<MyReply>(instance);
        let reply_port_ref2 = reply_port_handle2.bind();
        let message = TestMessage::CastAndReply {
            arg: "abc".to_string(),
            reply_to0: reply_port_ref0.clone().unsplit(),
            reply_to1: reply_port_ref1.clone(),
            reply_to2: reply_port_ref2.clone(),
        };

        clear_collected_tree();
        actor_mesh_ref.cast(instance, message).unwrap();

        let mut reply_tos = vec![];
        for _ in proc_mesh.extent().points() {
            let msg = rx.recv().await.expect("missing");
            match msg {
                TestMessage::CastAndReply {
                    arg,
                    reply_to0,
                    reply_to1,
                    reply_to2,
                } => {
                    assert_eq!(arg, "abc");
                    // port 0 is still the same as the original one because it is unsplit.
                    assert_eq!(reply_to0, reply_port_ref0.clone().unsplit());
                    // ports have been replaced by split ports.
                    assert_ne!(reply_to1, reply_port_ref1);
                    assert_ne!(reply_to2, reply_port_ref2);
                    reply_tos.push((reply_to1, reply_to2));
                }
                _ => {
                    panic!("unexpected message: {:?}", msg);
                }
            }
        }

        MeshSetupV1 {
            instance,
            actor_mesh_ref,
            reply1_rx,
            reply2_rx,
            reply_tos,
            host_mesh,
        }
    }

    async fn execute_cast_and_reply_v1() {
        let mut setup = setup_mesh_v1::<NoneAccumulator>(None).await;

        let ranks = setup.actor_mesh_ref.values().collect::<Vec<_>>();
        execute_cast_and_reply(
            ranks,
            setup.instance,
            setup.reply1_rx,
            setup.reply2_rx,
            setup.reply_tos,
        )
        .await;

        let _ = setup.host_mesh.shutdown(setup.instance).await;
    }

    #[async_timed_test(timeout_secs = 60)]
    async fn test_cast_and_reply_v1_retrofit() {
        let config = hyperactor_config::global::lock();
        let _guard = config.override_key(ENABLE_NATIVE_V1_CASTING, false);
        let _guard2 = config.override_key(
            hyperactor::config::ENABLE_DEST_ACTOR_REORDERING_BUFFER,
            false,
        );
        execute_cast_and_reply_v1().await
    }

    #[async_timed_test(timeout_secs = 60)]
    async fn test_cast_and_reply_v1_native() {
        let config = hyperactor_config::global::lock();
        let _guard = config.override_key(ENABLE_NATIVE_V1_CASTING, true);
        let _guard2 = config.override_key(
            hyperactor::config::ENABLE_DEST_ACTOR_REORDERING_BUFFER,
            true,
        );
        execute_cast_and_reply_v1().await
    }

    #[async_timed_test(timeout_secs = 60)]
    async fn test_cast_and_reply_v1_native_p2p() {
        let config = hyperactor_config::global::lock();
        let _guard = config.override_key(ENABLE_NATIVE_V1_CASTING, true);
        let _guard2 = config.override_key(
            hyperactor::config::ENABLE_DEST_ACTOR_REORDERING_BUFFER,
            true,
        );
        let _guard3 = config.override_key(crate::config::V1_CAST_POINT_TO_POINT_THRESHOLD, 1024);
        execute_cast_and_reply_v1().await
    }

    async fn execute_cast_and_accum_v1(config: &hyperactor_config::global::ConfigLock) {
        // Use temporary config for this test
        let _guard1 = config.override_key(hyperactor::config::SPLIT_MAX_BUFFER_SIZE, 1);

        let mut setup = setup_mesh_v1(Some(accum::sum::<u64>())).await;

        let ranks = setup.actor_mesh_ref.values().collect::<Vec<_>>();
        execute_cast_and_accum(ranks, setup.instance, setup.reply1_rx, setup.reply_tos).await;

        let _ = setup.host_mesh.shutdown(setup.instance).await;
    }

    #[async_timed_test(timeout_secs = 60)]
    async fn test_cast_and_accum_v1_retrofit() {
        let config = hyperactor_config::global::lock();
        let _guard = config.override_key(ENABLE_NATIVE_V1_CASTING, false);
        let _guard2 = config.override_key(
            hyperactor::config::ENABLE_DEST_ACTOR_REORDERING_BUFFER,
            false,
        );
        execute_cast_and_accum_v1(&config).await
    }

    #[async_timed_test(timeout_secs = 60)]
    async fn test_cast_and_accum_v1_native() {
        let config = hyperactor_config::global::lock();
        let _guard = config.override_key(ENABLE_NATIVE_V1_CASTING, true);
        let _guard2 = config.override_key(
            hyperactor::config::ENABLE_DEST_ACTOR_REORDERING_BUFFER,
            true,
        );
        execute_cast_and_accum_v1(&config).await
    }

    #[async_timed_test(timeout_secs = 60)]
    async fn test_cast_and_accum_v1_native_p2p() {
        let config = hyperactor_config::global::lock();
        let _guard = config.override_key(ENABLE_NATIVE_V1_CASTING, true);
        let _guard2 = config.override_key(
            hyperactor::config::ENABLE_DEST_ACTOR_REORDERING_BUFFER,
            true,
        );
        let _guard3 = config.override_key(crate::config::V1_CAST_POINT_TO_POINT_THRESHOLD, 1024);
        execute_cast_and_accum_v1(&config).await
    }

    struct OncePortMeshSetupV1 {
        instance: &'static Instance<testing::TestRootClient>,
        reply_rx: hyperactor::mailbox::OncePortReceiver<u64>,
        reply_tos: Vec<OncePortRef<u64>>,
        _reply_port_ref: OncePortRef<u64>,
        host_mesh: HostMesh,
    }

    async fn setup_once_port_mesh<A>(reducer: Option<A>) -> OncePortMeshSetupV1
    where
        A: Accumulator<State = u64, Update = u64> + Send + Sync + 'static,
    {
        let instance = crate::testing::instance();
        // We have to use a in process host mesh, because SPLIT_PORT_TREE only
        // can collect paths from the same process.
        let host_mesh = local_host_mesh(8).await;
        let proc_mesh = host_mesh
            .spawn(instance, "test", extent!(gpu = 8), None, None)
            .await
            .unwrap();

        let (tx, mut rx) = hyperactor::mailbox::open_port(instance);
        let params = TestActorParams {
            forward_port: tx.bind(),
        };
        let actor_name =
            crate::mesh_id::ActorMeshId::instance(hyperactor::id::Label::new("test").unwrap());
        // Make this actor a "system" actor to avoid spawning a controller actor.
        let actor_mesh: crate::ActorMesh<TestActor> = proc_mesh
            .spawn_with_name(&instance, actor_name, &params, None, true)
            .await
            .unwrap();
        let actor_mesh_ref = actor_mesh.deref().clone();

        let has_reducer = reducer.is_some();
        let (reply_port_handle, reply_rx) = match reducer {
            Some(reducer) => instance.mailbox().open_reduce_port(reducer),
            None => instance.mailbox().open_once_port::<u64>(),
        };
        let reply_port_ref = reply_port_handle.bind();

        let message = TestMessage::CastAndReplyOnce {
            arg: "abc".to_string(),
            reply_to: reply_port_ref.clone(),
        };

        clear_collected_tree();
        actor_mesh_ref.cast(instance, message).unwrap();

        let mut reply_tos = vec![];
        for _ in proc_mesh.extent().points() {
            let msg = rx.recv().await.expect("missing");
            match msg {
                TestMessage::CastAndReplyOnce { arg, reply_to } => {
                    assert_eq!(arg, "abc");
                    if has_reducer {
                        // With reducer: port is split.
                        assert_ne!(reply_to, reply_port_ref);
                    } else {
                        // Without reducer: port is passed through unchanged.
                        assert_eq!(reply_to, reply_port_ref);
                    }
                    reply_tos.push(reply_to);
                }
                _ => {
                    panic!("unexpected message: {:?}", msg);
                }
            }
        }

        OncePortMeshSetupV1 {
            instance,
            reply_rx,
            reply_tos,
            _reply_port_ref: reply_port_ref,
            host_mesh,
        }
    }

    async fn execute_cast_and_reply_once_v1() {
        // Test OncePort without accumulator - port is NOT split.
        // All destinations receive the same original port.
        // First reply is delivered, others fail at receiver (port closed).
        let mut setup = setup_once_port_mesh::<NoneAccumulator>(None).await;

        // All reply_tos point to the same port (not split).
        // Only the first message will be delivered successfully.
        let num_replies = setup.reply_tos.len();
        for (i, reply_to) in setup.reply_tos.into_iter().enumerate() {
            reply_to.post(setup.instance, i as u64);
        }

        // OncePort receives exactly one value (the first to arrive)
        let result = setup.reply_rx.recv().await.unwrap();
        // The result should be one of the values sent
        assert!(result < num_replies as u64);

        let _ = setup.host_mesh.shutdown(setup.instance).await;
    }

    #[async_timed_test(timeout_secs = 60)]
    async fn test_cast_and_reply_once_v1() {
        execute_cast_and_reply_once_v1().await
    }

    #[async_timed_test(timeout_secs = 60)]
    async fn test_cast_and_reply_once_v1_p2p() {
        let config = hyperactor_config::global::lock();
        let _guard = config.override_key(ENABLE_NATIVE_V1_CASTING, true);
        let _guard2 = config.override_key(
            hyperactor::config::ENABLE_DEST_ACTOR_REORDERING_BUFFER,
            true,
        );
        let _guard3 = config.override_key(crate::config::V1_CAST_POINT_TO_POINT_THRESHOLD, 1024);
        execute_cast_and_reply_once_v1().await
    }

    async fn execute_cast_and_accum_once_v1() {
        // Test OncePort splitting with sum accumulator.
        // Each destination actor replies with its rank.
        // The sum of all ranks should be received at the original port.
        let mut setup = setup_once_port_mesh(Some(accum::sum::<u64>())).await;

        // Each actor replies with its index
        let mut expected_sum = 0u64;
        for (i, reply_to) in setup.reply_tos.into_iter().enumerate() {
            reply_to.post(setup.instance, i as u64);
            expected_sum += i as u64;
        }

        // OncePort should receive the sum of all responses
        let result = setup.reply_rx.recv().await.unwrap();
        assert_eq!(result, expected_sum);

        let _ = setup.host_mesh.shutdown(setup.instance).await;
    }

    #[async_timed_test(timeout_secs = 60)]
    async fn test_cast_and_accum_once_v1() {
        execute_cast_and_accum_once_v1().await
    }

    #[async_timed_test(timeout_secs = 60)]
    async fn test_cast_and_accum_once_v1_p2p() {
        let config = hyperactor_config::global::lock();
        let _guard = config.override_key(ENABLE_NATIVE_V1_CASTING, true);
        let _guard2 = config.override_key(
            hyperactor::config::ENABLE_DEST_ACTOR_REORDERING_BUFFER,
            true,
        );
        let _guard3 = config.override_key(crate::config::V1_CAST_POINT_TO_POINT_THRESHOLD, 1024);
        execute_cast_and_accum_once_v1().await
    }

    #[async_timed_test(timeout_secs = 60)]
    async fn test_unsplit_port_not_split() {
        let instance = crate::testing::instance();
        let mut host_mesh = local_host_mesh(8).await;
        let proc_mesh = host_mesh
            .spawn(instance, "test", extent!(gpu = 8), None, None)
            .await
            .unwrap();

        let (tx, mut rx) = hyperactor::mailbox::open_port(instance);
        let params = TestActorParams {
            forward_port: tx.bind(),
        };
        let actor_name =
            crate::mesh_id::ActorMeshId::instance(hyperactor::id::Label::new("test").unwrap());
        let actor_mesh: ActorMesh<TestActor> = proc_mesh
            .spawn_with_name(&instance, actor_name, &params, None, true)
            .await
            .unwrap();
        let (reply_port_handle, mut reply_rx) = open_port::<u64>(instance);
        let reply_port_ref = reply_port_handle.bind().unsplit();

        let message = TestMessage::CastWithUnsplitPort {
            reply_to: reply_port_ref.clone(),
        };

        clear_collected_tree();
        actor_mesh.cast(instance, message).unwrap();

        // Verify that all destinations received the original port (not split).
        let num_points = proc_mesh.extent().points().count();
        for _ in 0..num_points {
            let msg = rx.recv().await.expect("missing");
            match msg {
                TestMessage::CastWithUnsplitPort { reply_to } => {
                    assert_eq!(
                        reply_to.port_addr(),
                        reply_port_ref.port_addr(),
                        "unsplit port should not be replaced by a comm actor split port"
                    );
                }
                _ => panic!("unexpected message: {:?}", msg),
            }
        }

        // All 8 actors sent replies directly to the same port.
        // Verify we receive all 8 replies.
        for _ in 0..8 {
            let val = reply_rx.recv().await.unwrap();
            assert_eq!(val, 42);
        }
        let _ = host_mesh.shutdown(instance).await;
    }

    // ────────────────────────────────────────────────────────────────────────
    // Cast correctness matrix: SENDER_ACTOR_ID stamping invariant verified at
    // the receiver across direct, V1 native p2p, V1 native tree, and V0 legacy
    // paths. Invariant: SENDER_ACTOR_ID == session owner (the actor whose
    // Sequencer assigned the SEQ_INFO for this session).

    struct SenderCaptureSetup {
        instance: &'static Instance<testing::TestRootClient>,
        host_mesh: HostMesh,
        proc_mesh: ProcMesh,
        actor_mesh: ActorMesh<SenderCapturingActor>,
        rx: hyperactor::mailbox::PortReceiver<CapturedSender>,
    }

    async fn setup_sender_capture_mesh(num_ranks: usize) -> SenderCaptureSetup {
        let instance = crate::testing::instance();
        let host_mesh = local_host_mesh(num_ranks).await;
        let proc_mesh = host_mesh
            .spawn(
                instance,
                "sender_capture",
                extent!(gpu = num_ranks),
                None,
                None,
            )
            .await
            .unwrap();
        let (tx, rx) = hyperactor::mailbox::open_port(instance);
        let params = SenderCapturingActorParams {
            forward_port: tx.bind(),
        };
        let actor_name = crate::mesh_id::ActorMeshId::instance(
            hyperactor::id::Label::new("sender_capture").unwrap(),
        );
        let actor_mesh: ActorMesh<SenderCapturingActor> = proc_mesh
            .spawn_with_name(&instance, actor_name, &params, None, true)
            .await
            .unwrap();
        SenderCaptureSetup {
            instance,
            host_mesh,
            proc_mesh,
            actor_mesh,
            rx,
        }
    }

    /// Direct send: SENDER_ACTOR_ID at the receiver equals the caller's
    /// `actor_addr` (MailboxExt::post stamps via the caller's Sequencer).
    #[async_timed_test(timeout_secs = 60)]
    async fn test_direct_send_sender_actor_id() {
        let mut setup = setup_sender_capture_mesh(1).await;

        let actor_ref = setup.actor_mesh.values().next().unwrap();
        actor_ref.post(setup.instance, SenderCaptureMsg());

        let captured = setup.rx.recv().await.unwrap();
        assert_eq!(
            captured,
            CapturedSender(Some(setup.instance.self_addr().clone())),
        );

        let _ = setup.host_mesh.shutdown(setup.instance).await;
    }

    /// V1 native p2p cast: SENDER_ACTOR_ID at the receiver equals the
    /// caster's `actor_addr`. SEQ_INFO is assigned at the cast origin's
    /// Sequencer; CommActor::deliver_to_dest stamps from message.sender()
    /// (which is the origin) on the leaf.
    #[async_timed_test(timeout_secs = 60)]
    async fn test_v1_native_p2p_sender_actor_id() {
        let config = hyperactor_config::global::lock();
        let _g1 = config.override_key(ENABLE_NATIVE_V1_CASTING, true);
        let _g2 = config.override_key(
            hyperactor::config::ENABLE_DEST_ACTOR_REORDERING_BUFFER,
            true,
        );
        let _g3 = config.override_key(crate::config::V1_CAST_POINT_TO_POINT_THRESHOLD, 1024);

        let mut setup = setup_sender_capture_mesh(1).await;
        setup
            .actor_mesh
            .cast(setup.instance, SenderCaptureMsg())
            .unwrap();

        let captured = setup.rx.recv().await.unwrap();
        assert_eq!(
            captured,
            CapturedSender(Some(setup.instance.self_addr().clone())),
        );

        let _ = setup.host_mesh.shutdown(setup.instance).await;
    }

    /// V1 native tree cast: SENDER_ACTOR_ID at every receiver equals the
    /// caster's `actor_addr`. Origin propagates through the comm-tree
    /// because each leaf's deliver_to_dest stamps from message.sender()
    /// when SEQ_INFO is present on headers.
    #[async_timed_test(timeout_secs = 60)]
    async fn test_v1_native_tree_sender_actor_id() {
        let config = hyperactor_config::global::lock();
        let _g1 = config.override_key(ENABLE_NATIVE_V1_CASTING, true);
        let _g2 = config.override_key(
            hyperactor::config::ENABLE_DEST_ACTOR_REORDERING_BUFFER,
            true,
        );
        let _g3 = config.override_key(crate::config::V1_CAST_POINT_TO_POINT_THRESHOLD, 0);

        let mut setup = setup_sender_capture_mesh(8).await;
        setup
            .actor_mesh
            .cast(setup.instance, SenderCaptureMsg())
            .unwrap();

        let num_points = setup.proc_mesh.extent().points().count();
        let expected = CapturedSender(Some(setup.instance.self_addr().clone()));
        for _ in 0..num_points {
            let captured = setup.rx.recv().await.unwrap();
            assert_eq!(captured, expected);
        }

        let _ = setup.host_mesh.shutdown(setup.instance).await;
    }
}