monarch_hyperactor/

endpoint.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 std::cell::Cell;
use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;

use hyperactor::ActorAddr;
use hyperactor::Instance;
use hyperactor::accum::Accumulator;
use hyperactor::accum::CommReducer;
use hyperactor::accum::ReducerFactory;
use hyperactor::accum::ReducerSpec;
use hyperactor::mailbox::OncePortReceiver;
use hyperactor::mailbox::PortReceiver;
use hyperactor_mesh::value_mesh::ValueOverlay;
use hyperactor_mesh::value_mesh::rle;
use monarch_types::py_global;
use ndslice::Extent;
use ndslice::Shape;
use pyo3::prelude::*;
use pyo3::types::PyDict;
use pyo3::types::PyTuple;
use serde_multipart::Part;
use typeuri::Named;

use crate::actor::MeshRef;
use crate::actor::MethodSpecifier;
use crate::actor::PythonActor;
use crate::actor::PythonMessage;
use crate::actor::PythonMessageKind;
use crate::actor::PythonResponseMessage;
use crate::actor_mesh::AllOrChoose;
use crate::actor_mesh::PythonActorMesh;
use crate::actor_mesh::SupervisableActorMesh;
use crate::actor_mesh::to_all_or_choose;
use crate::context::PyInstance;
use crate::mailbox::EitherPortRef;
use crate::mailbox::PythonOncePortRef;
use crate::mailbox::PythonPortRef;
use crate::metrics::ENDPOINT_BROADCAST_ERROR;
use crate::metrics::ENDPOINT_BROADCAST_THROUGHPUT;
use crate::metrics::ENDPOINT_CALL_ERROR;
use crate::metrics::ENDPOINT_CALL_LATENCY_US_HISTOGRAM;
use crate::metrics::ENDPOINT_CALL_ONE_ERROR;
use crate::metrics::ENDPOINT_CALL_ONE_LATENCY_US_HISTOGRAM;
use crate::metrics::ENDPOINT_CALL_ONE_THROUGHPUT;
use crate::metrics::ENDPOINT_CALL_THROUGHPUT;
use crate::metrics::ENDPOINT_CHOOSE_ERROR;
use crate::metrics::ENDPOINT_CHOOSE_LATENCY_US_HISTOGRAM;
use crate::metrics::ENDPOINT_CHOOSE_THROUGHPUT;
use crate::metrics::ENDPOINT_STREAM_ERROR;
use crate::metrics::ENDPOINT_STREAM_LATENCY_US_HISTOGRAM;
use crate::metrics::ENDPOINT_STREAM_THROUGHPUT;
use crate::pickle::PendingMessage;
use crate::pickle::PicklingState;
use crate::pytokio::PyPythonTask;
use crate::pytokio::PythonTask;
use crate::runtime::GilSite;
use crate::runtime::monarch_with_gil_blocking;
use crate::shape::PyExtent;
use crate::shape::PyShape;
use crate::supervision::Supervisable;
use crate::supervision::SupervisionError;
use crate::value_mesh::PyValueMesh;

py_global!(get_context, "monarch._src.actor.actor_mesh", "context");
py_global!(
    create_endpoint_message,
    "monarch._src.actor.actor_mesh",
    "_create_endpoint_message"
);
py_global!(
    dispatch_actor_rref,
    "monarch._src.actor.actor_mesh",
    "_dispatch_actor_rref"
);
py_global!(make_future, "monarch._src.actor.future", "Future");

/// The type of endpoint operation being performed.
///
/// Used to select the appropriate telemetry metrics for each operation type.
#[derive(Clone, Copy, Debug)]
pub(crate) enum EndpointAdverb {
    Call,
    CallOne,
    Choose,
    Stream,
}

impl EndpointAdverb {
    fn as_str(self) -> &'static str {
        match self {
            Self::Call => "call",
            Self::CallOne => "call_one",
            Self::Choose => "choose",
            Self::Stream => "stream",
        }
    }
}

/// RAII guard for recording endpoint call telemetry.
///
/// Records latency on drop, similar to Python's `@_with_telemetry` decorator.
/// Call `mark_error()` before dropping to also record an error.
pub struct RecordEndpointGuard {
    start: tokio::time::Instant,
    method_name: String,
    actor_count: usize,
    adverb: EndpointAdverb,
    error_occurred: Cell<bool>,
}

impl RecordEndpointGuard {
    fn new(
        start: tokio::time::Instant,
        method_name: String,
        actor_count: usize,
        adverb: EndpointAdverb,
    ) -> Self {
        let attributes = hyperactor_telemetry::kv_pairs!(
            "method" => method_name.clone()
        );
        match adverb {
            EndpointAdverb::Call => {
                ENDPOINT_CALL_THROUGHPUT.add(1, attributes);
            }
            EndpointAdverb::CallOne => {
                ENDPOINT_CALL_ONE_THROUGHPUT.add(1, attributes);
            }
            EndpointAdverb::Choose => {
                ENDPOINT_CHOOSE_THROUGHPUT.add(1, attributes);
            }
            EndpointAdverb::Stream => {
                // Throughput already recorded once at stream creation in py_stream_collector
            }
        }

        Self {
            start,
            method_name,
            actor_count,
            adverb,
            error_occurred: Cell::new(false),
        }
    }

    fn mark_error(&self) {
        self.error_occurred.set(true);
    }
}

impl Drop for RecordEndpointGuard {
    fn drop(&mut self) {
        let actor_count_str = self.actor_count.to_string();
        let attributes = hyperactor_telemetry::kv_pairs!(
            "method" => self.method_name.clone(),
            "actor_count" => actor_count_str
        );

        let duration_us = self.start.elapsed().as_micros();

        match self.adverb {
            EndpointAdverb::Call => {
                ENDPOINT_CALL_LATENCY_US_HISTOGRAM.record(duration_us as f64, attributes);
            }
            EndpointAdverb::CallOne => {
                ENDPOINT_CALL_ONE_LATENCY_US_HISTOGRAM.record(duration_us as f64, attributes);
            }
            EndpointAdverb::Choose => {
                ENDPOINT_CHOOSE_LATENCY_US_HISTOGRAM.record(duration_us as f64, attributes);
            }
            EndpointAdverb::Stream => {
                ENDPOINT_STREAM_LATENCY_US_HISTOGRAM.record(duration_us as f64, attributes);
            }
        }

        if self.error_occurred.get() {
            match self.adverb {
                EndpointAdverb::Call => {
                    ENDPOINT_CALL_ERROR.add(1, attributes);
                }
                EndpointAdverb::CallOne => {
                    ENDPOINT_CALL_ONE_ERROR.add(1, attributes);
                }
                EndpointAdverb::Choose => {
                    ENDPOINT_CHOOSE_ERROR.add(1, attributes);
                }
                EndpointAdverb::Stream => {
                    ENDPOINT_STREAM_ERROR.add(1, attributes);
                }
            }
        }
    }
}

/// Send-safe RAII guard for an OTEL-style endpoint span.
///
/// We only need endpoint spans for telemetry slices, not for `tracing` context
/// propagation. So this guard emits synthetic trace events directly into the
/// unified telemetry dispatcher instead of holding a real `tracing::Span`
/// across `.await` points.
pub(crate) struct SpanGuard {
    id: u64,
}

impl SpanGuard {
    fn actor_endpoint(name: &'static str, actor_id: &ActorAddr, mesh: &str, method: &str) -> Self {
        Self {
            id: hyperactor_telemetry::start_user_span(
                name,
                hyperactor_telemetry::sinks::perfetto::ENDPOINT_TELEMETRY_TARGET,
                [
                    (
                        "actor_id",
                        hyperactor_telemetry::trace_dispatcher::FieldValue::Str(
                            actor_id.to_string(),
                        ),
                    ),
                    (
                        "mesh",
                        hyperactor_telemetry::trace_dispatcher::FieldValue::Str(mesh.to_string()),
                    ),
                    (
                        "method",
                        hyperactor_telemetry::trace_dispatcher::FieldValue::Str(method.to_string()),
                    ),
                ],
            ),
        }
    }

    fn remote(name: &'static str, actor_id: &ActorAddr, call_name: &str) -> Self {
        Self {
            id: hyperactor_telemetry::start_user_span(
                name,
                hyperactor_telemetry::sinks::perfetto::ENDPOINT_TELEMETRY_TARGET,
                [
                    (
                        "actor_id",
                        hyperactor_telemetry::trace_dispatcher::FieldValue::Str(
                            actor_id.to_string(),
                        ),
                    ),
                    (
                        "call_name",
                        hyperactor_telemetry::trace_dispatcher::FieldValue::Str(
                            call_name.to_string(),
                        ),
                    ),
                ],
            ),
        }
    }
}

impl Drop for SpanGuard {
    fn drop(&mut self) {
        hyperactor_telemetry::end_user_span(self.id);
    }
}

fn supervision_error_to_pyerr(err: PyErr, qualified_endpoint_name: &Option<String>) -> PyErr {
    match qualified_endpoint_name {
        Some(endpoint) => monarch_with_gil_blocking(GilSite::Supervise, |py| {
            SupervisionError::set_endpoint_on_err(py, err, endpoint.clone())
        }),
        None => err,
    }
}

async fn collect_value(
    rx: &mut PortReceiver<PythonMessage>,
    supervision_monitor: &Option<Arc<dyn Supervisable>>,
    instance: &Instance<PythonActor>,
    qualified_endpoint_name: &Option<String>,
) -> PyResult<(Part, Vec<MeshRef>, Option<usize>)> {
    enum RaceResult {
        Message(Box<PythonMessage>),
        SupervisionError(PyErr),
        RecvError(String),
    }

    let race_result = match supervision_monitor {
        Some(sup) => {
            tokio::select! {
                biased;
                result = sup.supervision_event(instance) => {
                    match result {
                        Some(err) => RaceResult::SupervisionError(err),
                        None => {
                            match rx.recv().await {
                                Ok(msg) => RaceResult::Message(Box::new(msg)),
                                Err(e) => RaceResult::RecvError(e.to_string()),
                            }
                        }
                    }
                }
                msg = rx.recv() => {
                    match msg {
                        Ok(m) => RaceResult::Message(Box::new(m)),
                        Err(e) => RaceResult::RecvError(e.to_string()),
                    }
                }
            }
        }
        _ => match rx.recv().await {
            Ok(msg) => RaceResult::Message(Box::new(msg)),
            Err(e) => RaceResult::RecvError(e.to_string()),
        },
    };

    match race_result {
        RaceResult::Message(boxed) => {
            let PythonMessage {
                kind,
                message,
                refs,
            } = *boxed;
            match kind {
                PythonMessageKind::Result { rank, .. } => Ok((message, refs, rank)),
                PythonMessageKind::Exception { .. } => {
                    monarch_with_gil_blocking(GilSite::Traceback, |py| {
                        let mesh_references: VecDeque<Option<MeshRef>> =
                            refs.into_iter().map(Some).collect();
                        let mut state =
                            PicklingState::from_parts(message, VecDeque::new(), mesh_references);
                        Err(PyErr::from_value(state.unpickle(py)?.into_bound(py)))
                    })
                }
                other => Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "unexpected message kind {:?}",
                    other
                ))),
            }
        }
        RaceResult::RecvError(e) => Err(pyo3::exceptions::PyEOFError::new_err(format!(
            "Port closed: {}",
            e
        ))),
        RaceResult::SupervisionError(err) => {
            Err(supervision_error_to_pyerr(err, qualified_endpoint_name))
        }
    }
}

#[tracing::instrument(level = "debug", skip_all)]
async fn collect_valuemesh(
    extent: Extent,
    rx: OncePortReceiver<PythonMessage>,
    method_name: String,
    supervision_monitor: Option<Arc<dyn Supervisable>>,
    instance: &Instance<PythonActor>,
    qualified_endpoint_name: Option<String>,
) -> PyResult<Py<PyAny>> {
    let start = tokio::time::Instant::now();

    let expected_count = extent.num_ranks();

    let record_guard = RecordEndpointGuard::new(
        start,
        method_name.clone(),
        expected_count,
        EndpointAdverb::Call,
    );

    enum RaceResult {
        Collected(Box<PythonMessage>),
        SupervisionError(PyErr),
        RecvError(String),
    }

    let race_result = match &supervision_monitor {
        Some(sup) => {
            tokio::select! {
                biased;
                result = sup.supervision_event(instance) => {
                    match result {
                        Some(err) => RaceResult::SupervisionError(err),
                        None => RaceResult::RecvError(
                            "supervision monitor closed unexpectedly".to_string()
                        ),
                    }
                }
                batch = rx.recv() => {
                    match batch {
                        Ok(b) => RaceResult::Collected(Box::new(b)),
                        Err(e) => RaceResult::RecvError(e.to_string()),
                    }
                }
            }
        }
        None => match rx.recv().await {
            Ok(batch) => RaceResult::Collected(Box::new(batch)),
            Err(e) => RaceResult::RecvError(e.to_string()),
        },
    };

    match race_result {
        RaceResult::Collected(boxed) => {
            let msg = *boxed;
            let overlay = msg.into_overlay().map_err(|e| {
                pyo3::exceptions::PyRuntimeError::new_err(format!(
                    "failed to extract overlay from collected responses: {e}"
                ))
            })?;
            monarch_with_gil_blocking(GilSite::ReplyConvert, |py| {
                // Out-of-band mesh refs reunite only while the `PicklingState` is
                // live, i.e. during decode, so a ref-carrying response must be
                // decoded here (eagerly, at accumulation): a lazy decode on access
                // would have no state to reunite against, the REFS-1 failure. But
                // decoding *every* response here would revert D96180139, which made
                // valuemesh values unpickle lazily on access so a large
                // OnceBuffer-accumulated `.call()` does not pay one big unpickle
                // burst at the end of collection.
                //
                // Reconcile the two by gating on refs. A `.call()` fans one
                // endpoint over the mesh, so the batch is uniform: either every
                // response carries refs (the endpoint returns a mesh, the minority)
                // or none does (plain values, the common case). With no refs we
                // keep the raw parts and let them unpickle on access (D96180139,
                // preserved); only with refs present do we decode eagerly to
                // reunite them.
                let has_refs = overlay.runs().any(|(_, payload)| {
                    let (PythonResponseMessage::Result { refs, .. }
                    | PythonResponseMessage::Exception { refs, .. }) = payload;
                    !refs.is_empty()
                });

                if !has_refs {
                    let mut parts = Vec::with_capacity(expected_count);
                    for (range, payload) in overlay.runs() {
                        match payload {
                            PythonResponseMessage::Result { part, .. } => {
                                parts.extend(range.clone().map(|_| part.clone()));
                            }
                            PythonResponseMessage::Exception { .. } => {
                                record_guard.mark_error();
                                return Err(PyErr::from_value(payload.decode(py)?.into_bound(py)));
                            }
                        }
                    }
                    return Ok(PyValueMesh::build_from_parts(&extent, parts)?
                        .into_pyobject(py)?
                        .into_any()
                        .unbind());
                }

                let mut objects: Vec<Py<PyAny>> = Vec::with_capacity(expected_count);
                for (range, payload) in overlay.runs() {
                    match payload {
                        PythonResponseMessage::Result { .. } => {
                            let obj = payload.decode(py)?;
                            objects.extend(range.clone().map(|_| obj.clone_ref(py)));
                        }
                        PythonResponseMessage::Exception { .. } => {
                            record_guard.mark_error();
                            return Err(PyErr::from_value(payload.decode(py)?.into_bound(py)));
                        }
                    }
                }
                Ok(PyValueMesh::build_from_objects(&extent, objects)?
                    .into_pyobject(py)?
                    .into_any()
                    .unbind())
            })
        }
        RaceResult::RecvError(e) => {
            record_guard.mark_error();
            Err(pyo3::exceptions::PyEOFError::new_err(format!(
                "Port closed: {}",
                e
            )))
        }
        RaceResult::SupervisionError(err) => {
            record_guard.mark_error();
            Err(supervision_error_to_pyerr(err, &qualified_endpoint_name))
        }
    }
}

fn value_collector(
    mut receiver: PortReceiver<PythonMessage>,
    method_name: String,
    supervision_monitor: Option<Arc<dyn Supervisable>>,
    instance: Instance<PythonActor>,
    qualified_endpoint_name: Option<String>,
    adverb: EndpointAdverb,
    span_guard: SpanGuard,
) -> PyResult<PyPythonTask> {
    Ok(PythonTask::new(async move {
        let _span_guard = span_guard;
        let start = tokio::time::Instant::now();

        let record_guard = RecordEndpointGuard::new(start, method_name, 1, adverb);

        match collect_value(
            &mut receiver,
            &supervision_monitor,
            &instance,
            &qualified_endpoint_name,
        )
        .await
        {
            Ok((message, refs, _)) => monarch_with_gil_blocking(GilSite::ReplyConvert, |py| {
                let mesh_references: VecDeque<Option<MeshRef>> =
                    refs.into_iter().map(Some).collect();
                let mut state =
                    PicklingState::from_parts(message, VecDeque::new(), mesh_references);
                state.unpickle(py)
            }),
            Err(e) => {
                record_guard.mark_error();
                Err(e)
            }
        }
    })?
    .into())
}

/// A streaming iterator that yields futures for each response from actors.
///
/// Implements Python's iterator protocol (`__iter__`/`__next__`) to yield
/// `Future` objects that resolve to individual actor responses.
#[pyclass(
    name = "ValueStream",
    module = "monarch._rust_bindings.monarch_hyperactor.endpoint"
)]
pub struct PyValueStream {
    receiver: Arc<tokio::sync::Mutex<PortReceiver<PythonMessage>>>,
    /// Supervisor for monitoring actor health during streaming.
    supervision_monitor: Option<Arc<dyn Supervisable>>,
    instance: Instance<PythonActor>,
    remaining: AtomicUsize,
    method_name: String,
    qualified_endpoint_name: Option<String>,
    start: tokio::time::Instant,
    actor_count: usize,
    future_class: Py<PyAny>,
}

#[pymethods]
impl PyValueStream {
    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }

    fn __next__(&self, py: Python<'_>) -> PyResult<Option<Py<PyAny>>> {
        let remaining = self.remaining.load(Ordering::Relaxed);
        if remaining == 0 {
            return Ok(None);
        }
        self.remaining.store(remaining - 1, Ordering::Relaxed);

        let receiver = self.receiver.clone();
        let supervision_monitor = self.supervision_monitor.clone();
        let instance = self.instance.clone_for_py();
        let qualified_endpoint_name = self.qualified_endpoint_name.clone();
        let start = self.start;
        let method_name = self.method_name.clone();
        let actor_count = self.actor_count;

        let task: PyPythonTask = PythonTask::new(async move {
            let record_guard =
                RecordEndpointGuard::new(start, method_name, actor_count, EndpointAdverb::Stream);

            let mut rx_guard = receiver.lock().await;

            match collect_value(
                &mut rx_guard,
                &supervision_monitor,
                &instance,
                &qualified_endpoint_name,
            )
            .await
            {
                Ok((message, refs, _)) => monarch_with_gil_blocking(GilSite::ReplyConvert, |py| {
                    let mesh_references: VecDeque<Option<MeshRef>> =
                        refs.into_iter().map(Some).collect();
                    let mut state =
                        PicklingState::from_parts(message, VecDeque::new(), mesh_references);
                    state.unpickle(py)
                }),
                Err(e) => {
                    record_guard.mark_error();
                    Err(e)
                }
            }
        })?
        .into();

        let kwargs = PyDict::new(py);
        kwargs.set_item("coro", task)?;
        let future = self.future_class.call(py, (), Some(&kwargs))?;
        Ok(Some(future))
    }
}

fn wrap_in_future(py: Python<'_>, task: PyPythonTask) -> PyResult<Py<PyAny>> {
    let kwargs = PyDict::new(py);
    kwargs.set_item("coro", task)?;
    let future = make_future(py).call((), Some(&kwargs))?;
    Ok(future.unbind())
}

/// Trait that defines the core operations an endpoint must provide.
/// Both ActorEndpoint and RemoteEndpoint implement this trait.
pub(crate) trait Endpoint {
    /// Get the extent of the endpoint's targets.
    fn get_extent(&self, py: Python<'_>) -> PyResult<Extent>;

    /// Get the method name for this endpoint.
    fn get_method_name(&self) -> &str;

    /// Create and send a message with the given args/kwargs.
    fn send_message<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
        port_ref: Option<EitherPortRef>,
        selection: AllOrChoose,
        instance: &Instance<PythonActor>,
    ) -> PyResult<()>;

    /// Like `send_message` but stamps `caller_headers` onto the
    /// outgoing request envelope. Implementations that can carry
    /// headers override this; the default delegates to `send_message`
    /// and drops them.
    fn send_message_with_headers<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
        port_ref: Option<EitherPortRef>,
        selection: AllOrChoose,
        instance: &Instance<PythonActor>,
        _caller_headers: hyperactor_config::Flattrs,
    ) -> PyResult<()> {
        self.send_message(py, args, kwargs, port_ref, selection, instance)
    }

    /// Build the operation-context envelope headers to stamp on an
    /// outgoing request for this endpoint invocation. The result is
    /// empty when the endpoint cannot supply a qualified name
    /// (e.g. `RemoteEndpoint`'s `get_qualified_name` returns `None`),
    /// in which case callers see the unchanged dispatch surface.
    fn build_operation_context_headers(
        &self,
        adverb: EndpointAdverb,
    ) -> hyperactor_config::Flattrs {
        let adverb_str = match adverb {
            EndpointAdverb::Call => "call",
            EndpointAdverb::CallOne => "call_one",
            EndpointAdverb::Choose => "choose",
            EndpointAdverb::Stream => "stream",
        };
        let attrs = crate::operation_context::build_operation_context_attrs(
            self.get_qualified_name(),
            Some(adverb_str),
        );
        let mut headers = hyperactor_config::Flattrs::new();
        crate::operation_context::stamp_operation_context(&mut headers, &attrs);
        headers
    }

    /// Get the supervision_monitor for this endpoint (if any).
    fn get_supervision_monitor(&self) -> Option<Arc<dyn Supervisable>>;

    /// Get the qualified endpoint name for error messages (if any).
    fn get_qualified_name(&self) -> Option<String>;

    /// Open an OTEL-style span for an endpoint invocation. Each impl attaches
    /// the fields the perfetto sink needs to synthesize the display name
    /// (`{mesh}.{method}.{adverb}` for ActorEndpoint, `{call_name}.{adverb}`
    /// for Remote) and route the slice to an actor-specific track. The adverb is the span name,
    /// so no formatting happens at the call site and the sink formats only when
    /// it renders the slice.
    fn enter_endpoint_span(&self, adverb: EndpointAdverb, actor_id: &ActorAddr) -> SpanGuard;

    fn get_current_instance(&self, py: Python<'_>) -> PyResult<Instance<PythonActor>> {
        let context = get_context(py).call0()?;
        let py_instance: PyRef<PyInstance> = context.getattr("actor_instance")?.extract()?;
        Ok(py_instance.clone().into_instance())
    }

    fn open_response_port(
        &self,
        instance: &Instance<PythonActor>,
    ) -> (PythonPortRef, PortReceiver<PythonMessage>) {
        let (p, receiver) = instance.mailbox_for_py().open_port::<PythonMessage>();
        (PythonPortRef { inner: p.bind() }, receiver)
    }

    fn open_reduce_response_port(
        &self,
        instance: &Instance<PythonActor>,
    ) -> (PythonOncePortRef, OncePortReceiver<PythonMessage>) {
        let (p, receiver) = instance
            .mailbox_for_py()
            .open_reduce_port(PythonResponseMessageAccumulator);
        (PythonOncePortRef::from(p.bind()), receiver)
    }

    /// Call the endpoint on all actors and collect all responses into a ValueMesh.
    #[tracing::instrument(level = "debug", skip_all)]
    fn call<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
    ) -> PyResult<Py<PyAny>> {
        let instance = self.get_current_instance(py)?;
        let span_guard = self.enter_endpoint_span(EndpointAdverb::Call, instance.self_addr());

        let extent = self.get_extent(py)?;
        let method_name = self.get_method_name().to_string();
        let (port_ref, receiver) = self.open_reduce_response_port(&instance);

        let supervision_monitor = self.get_supervision_monitor();
        let qualified_endpoint_name = self.get_qualified_name();

        let caller_headers = self.build_operation_context_headers(EndpointAdverb::Call);
        self.send_message_with_headers(
            py,
            args,
            kwargs,
            Some(EitherPortRef::Once(port_ref)),
            AllOrChoose::All,
            &instance,
            caller_headers,
        )?;

        let instance_for_task = instance.clone_for_py();
        let task: PyPythonTask = PythonTask::new(async move {
            let _span_guard = span_guard;
            collect_valuemesh(
                extent,
                receiver,
                method_name,
                supervision_monitor,
                &instance_for_task,
                qualified_endpoint_name,
            )
            .await
        })?
        .into();

        wrap_in_future(py, task)
    }

    /// Load balanced sends a message to one chosen actor and awaits a result.
    fn choose<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
    ) -> PyResult<Py<PyAny>> {
        let instance = self.get_current_instance(py)?;
        let span_guard = self.enter_endpoint_span(EndpointAdverb::Choose, instance.self_addr());
        let (port_ref, receiver) = self.open_response_port(&instance);

        let caller_headers = self.build_operation_context_headers(EndpointAdverb::Choose);
        self.send_message_with_headers(
            py,
            args,
            kwargs,
            Some(EitherPortRef::Unbounded(port_ref)),
            AllOrChoose::Choose,
            &instance,
            caller_headers,
        )?;

        let task = value_collector(
            receiver,
            self.get_method_name().to_string(),
            self.get_supervision_monitor(),
            instance.clone_for_py(),
            self.get_qualified_name(),
            EndpointAdverb::Choose,
            span_guard,
        )?;

        wrap_in_future(py, task)
    }

    /// Call the endpoint on exactly one actor (the mesh must have exactly one actor).
    fn call_one<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
    ) -> PyResult<Py<PyAny>> {
        let extent = self.get_extent(py)?;

        if extent.num_ranks() != 1 {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "call_one requires exactly 1 actor, but mesh has {}",
                extent.num_ranks()
            )));
        }

        let instance = self.get_current_instance(py)?;
        let span_guard = self.enter_endpoint_span(EndpointAdverb::CallOne, instance.self_addr());
        let (port_ref, receiver) = self.open_response_port(&instance);

        let caller_headers = self.build_operation_context_headers(EndpointAdverb::CallOne);
        self.send_message_with_headers(
            py,
            args,
            kwargs,
            Some(EitherPortRef::Unbounded(port_ref)),
            AllOrChoose::All,
            &instance,
            caller_headers,
        )?;

        let task = value_collector(
            receiver,
            self.get_method_name().to_string(),
            self.get_supervision_monitor(),
            instance.clone_for_py(),
            self.get_qualified_name(),
            EndpointAdverb::CallOne,
            span_guard,
        )?;

        wrap_in_future(py, task)
    }

    /// Call the endpoint on all actors and return an iterator of Futures.
    fn stream<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
    ) -> PyResult<Py<PyAny>> {
        let extent = self.get_extent(py)?;
        let method_name = self.get_method_name().to_string();

        let instance = self.get_current_instance(py)?;
        let (port_ref, receiver) = self.open_response_port(&instance);

        let caller_headers = self.build_operation_context_headers(EndpointAdverb::Stream);
        self.send_message_with_headers(
            py,
            args,
            kwargs,
            Some(EitherPortRef::Unbounded(port_ref)),
            AllOrChoose::All,
            &instance,
            caller_headers,
        )?;

        let actor_count = extent.num_ranks();
        let start = tokio::time::Instant::now();
        let supervision_monitor = self.get_supervision_monitor();
        let qualified_endpoint_name = self.get_qualified_name();
        let future_class = make_future(py).unbind();

        let attributes = hyperactor_telemetry::kv_pairs!(
            "method" => method_name.clone()
        );
        ENDPOINT_STREAM_THROUGHPUT.add(1, attributes);

        let stream = PyValueStream {
            receiver: Arc::new(tokio::sync::Mutex::new(receiver)),
            supervision_monitor,
            instance: instance.clone_for_py(),
            remaining: AtomicUsize::new(actor_count),
            method_name,
            qualified_endpoint_name,
            start,
            actor_count,
            future_class,
        };

        Ok(stream.into_pyobject(py)?.unbind().into())
    }

    /// Send a message to all actors without waiting for responses (fire-and-forget).
    fn broadcast<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
    ) -> PyResult<()> {
        let instance = self.get_current_instance(py)?;
        let method_name = self.get_method_name();
        let attributes = hyperactor_telemetry::kv_pairs!(
            "method" => method_name.to_string()
        );

        match self.send_message(py, args, kwargs, None, AllOrChoose::All, &instance) {
            Ok(()) => {
                ENDPOINT_BROADCAST_THROUGHPUT.add(1, attributes);
                Ok(())
            }
            Err(e) => {
                ENDPOINT_BROADCAST_ERROR.add(1, attributes);
                Err(e)
            }
        }
    }
}

#[pyclass(
    name = "ActorEndpoint",
    module = "monarch._rust_bindings.monarch_hyperactor.endpoint"
)]
pub struct ActorEndpoint {
    inner: Arc<dyn SupervisableActorMesh>,
    shape: Shape,
    method: MethodSpecifier,
    mesh_name: String,
    signature: Option<Py<PyAny>>,
    proc_mesh: Option<Py<PyAny>>,
    propagator: Option<Py<PyAny>>,
}

impl ActorEndpoint {
    fn create_message<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
        port_ref: Option<EitherPortRef>,
    ) -> PyResult<PendingMessage> {
        let port_ref_py: Py<PyAny> = match port_ref {
            Some(pr) => pr.clone().into_pyobject(py)?.unbind(),
            None => py.None(),
        };

        let result = create_endpoint_message(py).call1((
            self.method.clone(),
            self.signature
                .as_ref()
                .map_or_else(|| py.None(), |s| s.clone_ref(py)),
            args,
            kwargs
                .map_or_else(|| PyDict::new(py), |d| d.clone())
                .into_any(),
            port_ref_py,
            self.proc_mesh
                .as_ref()
                .map_or_else(|| py.None(), |p| p.clone_ref(py)),
        ))?;
        let mut pending: PyRefMut<'_, PendingMessage> = result.extract()?;
        pending.take()
    }
}

impl Endpoint for ActorEndpoint {
    fn get_extent(&self, _py: Python<'_>) -> PyResult<Extent> {
        Ok(self.shape.extent())
    }

    fn get_method_name(&self) -> &str {
        self.method.name()
    }

    fn send_message<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
        port_ref: Option<EitherPortRef>,
        selection: AllOrChoose,
        instance: &Instance<PythonActor>,
    ) -> PyResult<()> {
        let message = self.create_message(py, args, kwargs, port_ref)?;
        self.inner.cast_unresolved(message, selection, instance)
    }

    fn send_message_with_headers<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
        port_ref: Option<EitherPortRef>,
        selection: AllOrChoose,
        instance: &Instance<PythonActor>,
        caller_headers: hyperactor_config::Flattrs,
    ) -> PyResult<()> {
        let message = self.create_message(py, args, kwargs, port_ref)?;
        self.inner
            .cast_unresolved_with_headers(message, selection, instance, caller_headers)
    }

    fn get_supervision_monitor(&self) -> Option<Arc<dyn Supervisable>> {
        Some(self.inner.clone())
    }

    fn get_qualified_name(&self) -> Option<String> {
        Some(format!("{}.{}()", self.mesh_name, self.method.name()))
    }

    fn enter_endpoint_span(&self, adverb: EndpointAdverb, actor_id: &ActorAddr) -> SpanGuard {
        let mesh = self.mesh_name.as_str();
        let method = self.method.name();
        SpanGuard::actor_endpoint(adverb.as_str(), actor_id, mesh, method)
    }
}

#[pymethods]
impl ActorEndpoint {
    /// Create a new ActorEndpoint.
    #[new]
    #[pyo3(signature = (actor_mesh, method, shape, mesh_name, signature=None, proc_mesh=None, propagator=None))]
    fn new(
        actor_mesh: PythonActorMesh,
        method: MethodSpecifier,
        shape: PyShape,
        mesh_name: String,
        signature: Option<Py<PyAny>>,
        proc_mesh: Option<Py<PyAny>>,
        propagator: Option<Py<PyAny>>,
    ) -> Self {
        Self {
            inner: actor_mesh.get_inner(),
            shape: shape.get_inner().clone(),
            method,
            mesh_name,
            signature,
            proc_mesh,
            propagator,
        }
    }

    /// Get the method specifier (used by actor_rref for tensor dispatch).
    #[getter]
    fn _name(&self) -> MethodSpecifier {
        self.method.clone()
    }

    /// Get the signature (used for argument checking in _dispatch_actor_rref).
    #[getter]
    fn _signature(&self, py: Python<'_>) -> Py<PyAny> {
        self.signature
            .clone()
            .unwrap_or_else(|| py.None().into_any())
    }

    /// Get the actor mesh (used by actor_rref for sending messages).
    #[getter]
    fn _actor_mesh(&self) -> PythonActorMesh {
        PythonActorMesh::from_impl(self.inner.clone())
    }

    /// Propagation method for tensor shape inference.
    /// Delegates to Python _do_propagate helper.
    fn _propagate<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyAny>,
        kwargs: &Bound<'py, PyAny>,
        fake_args: &Bound<'py, PyAny>,
        fake_kwargs: &Bound<'py, PyAny>,
    ) -> PyResult<Py<PyAny>> {
        let do_propagate = py
            .import("monarch._src.actor.endpoint")?
            .getattr("_do_propagate")?;
        let propagator = self
            .propagator
            .as_ref()
            .map(|p| p.clone_ref(py).into_bound(py))
            .unwrap_or_else(|| py.None().into_bound(py));
        let cache = PyDict::new(py);
        do_propagate
            .call1((&propagator, args, kwargs, fake_args, fake_kwargs, cache))?
            .extract()
    }

    /// Propagation for fetch operations.
    /// Returns None if no propagator is provided, otherwise calls _propagate.
    fn _fetch_propagate<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyAny>,
        kwargs: &Bound<'py, PyAny>,
        fake_args: &Bound<'py, PyAny>,
        fake_kwargs: &Bound<'py, PyAny>,
    ) -> PyResult<Py<PyAny>> {
        if self.propagator.is_none() {
            return Ok(py.None());
        }
        self._propagate(py, args, kwargs, fake_args, fake_kwargs)
    }

    /// Propagation for pipe operations.
    /// Requires an explicit callable propagator.
    fn _pipe_propagate<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyAny>,
        kwargs: &Bound<'py, PyAny>,
        fake_args: &Bound<'py, PyAny>,
        fake_kwargs: &Bound<'py, PyAny>,
    ) -> PyResult<Py<PyAny>> {
        // Check if propagator is callable
        let is_callable = self
            .propagator
            .as_ref()
            .map(|p| p.bind(py).is_callable())
            .unwrap_or(false);
        if !is_callable {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "Must specify explicit callable for pipe",
            ));
        }
        self._propagate(py, args, kwargs, fake_args, fake_kwargs)
    }

    /// Get the rref result by calling the Python dispatch helper.
    #[pyo3(signature = (*args, **kwargs))]
    fn rref<'py>(
        slf: PyRef<'py, Self>,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
    ) -> PyResult<Py<PyAny>> {
        let kwargs_dict = kwargs.map_or_else(|| PyDict::new(py), |d| d.clone());

        // Call _dispatch_actor_rref(endpoint, args, kwargs)
        let result = dispatch_actor_rref(py).call1((slf.into_pyobject(py)?, args, kwargs_dict))?;

        Ok(result.unbind())
    }

    /// Call the endpoint on all actors and collect all responses into a ValueMesh.
    #[pyo3(signature = (*args, **kwargs), name = "call")]
    fn py_call<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
    ) -> PyResult<Py<PyAny>> {
        self.call(py, args, kwargs)
    }

    /// Load balanced sends a message to one chosen actor and awaits a result.
    #[pyo3(signature = (*args, **kwargs), name = "choose")]
    fn py_choose<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
    ) -> PyResult<Py<PyAny>> {
        self.choose(py, args, kwargs)
    }

    /// Call the endpoint on exactly one actor (the mesh must have exactly one actor).
    #[pyo3(signature = (*args, **kwargs), name = "call_one")]
    fn py_call_one<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
    ) -> PyResult<Py<PyAny>> {
        self.call_one(py, args, kwargs)
    }

    /// Call the endpoint on all actors and return an iterator of Futures.
    #[pyo3(signature = (*args, **kwargs), name = "stream")]
    fn py_stream<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
    ) -> PyResult<Py<PyAny>> {
        self.stream(py, args, kwargs)
    }

    /// Send a message to all actors without waiting for responses (fire-and-forget).
    #[pyo3(signature = (*args, **kwargs), name = "broadcast")]
    fn py_broadcast<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
    ) -> PyResult<()> {
        self.broadcast(py, args, kwargs)
    }

    /// Send a message with optional port for response (used by actor_mesh.send).
    fn _send<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: &Bound<'py, PyDict>,
        port: Option<EitherPortRef>,
        selection: &str,
    ) -> PyResult<()> {
        let instance = self.get_current_instance(py)?;
        let sel = to_all_or_choose(selection)?;
        self.send_message(py, args, Some(kwargs), port, sel, &instance)
    }
}

/// A Rust wrapper for Python's RemoteImpl endpoint.
///
/// This allows us to implement the adverb methods (call, choose, call_one, stream, broadcast)
/// in Rust while delegating the actual send logic to the Python RemoteImpl._send() method.
#[pyclass(
    name = "Remote",
    module = "monarch._rust_bindings.monarch_hyperactor.endpoint"
)]
pub struct Remote {
    /// The wrapped Python RemoteImpl object
    inner: Py<PyAny>,
}

impl Endpoint for Remote {
    fn get_extent(&self, py: Python<'_>) -> PyResult<Extent> {
        let extent: PyExtent = self.inner.call_method0(py, "_get_extent")?.extract(py)?;
        Ok(extent.into())
    }

    fn get_method_name(&self) -> &str {
        "unknown"
    }

    fn send_message<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
        port_ref: Option<EitherPortRef>,
        selection: AllOrChoose,
        _instance: &Instance<PythonActor>,
    ) -> PyResult<()> {
        let send_kwargs = PyDict::new(py);
        match port_ref {
            Some(pr) => send_kwargs.set_item("port", pr.clone())?,
            None => send_kwargs.set_item("port", py.None())?,
        }

        send_kwargs.set_item("selection", selection.as_str())?;

        let kwargs_dict = kwargs.map_or_else(|| PyDict::new(py), |d| d.clone());
        self.inner
            .call_method(py, "_send", (args.clone(), kwargs_dict), Some(&send_kwargs))?;

        Ok(())
    }

    fn get_supervision_monitor(&self) -> Option<Arc<dyn Supervisable>> {
        None // Remote endpoints don't have supervision_monitors
    }

    fn get_qualified_name(&self) -> Option<String> {
        None // Remote endpoints don't have qualified names
    }

    fn enter_endpoint_span(&self, adverb: EndpointAdverb, actor_id: &ActorAddr) -> SpanGuard {
        let call_name = monarch_with_gil_blocking(GilSite::DisplayName, |py| {
            self.inner
                .call_method0(py, "_call_name")
                .ok()
                .and_then(|v| v.extract::<String>(py).ok())
        });
        let call_name = call_name.as_deref().unwrap_or("");
        SpanGuard::remote(adverb.as_str(), actor_id, call_name)
    }
}

#[pymethods]
impl Remote {
    /// Create a new Remote wrapping a Python RemoteImpl object.
    #[new]
    fn new(remote: Py<PyAny>) -> Self {
        Self { inner: remote }
    }

    /// Call the endpoint on all actors and collect all responses into a ValueMesh.
    #[pyo3(signature = (*args, **kwargs), name = "call")]
    fn py_call<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
    ) -> PyResult<Py<PyAny>> {
        self.call(py, args, kwargs)
    }

    /// Load balanced sends a message to one chosen actor and awaits a result.
    #[pyo3(signature = (*args, **kwargs), name = "choose")]
    fn py_choose<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
    ) -> PyResult<Py<PyAny>> {
        self.choose(py, args, kwargs)
    }

    /// Call the endpoint on exactly one actor (the mesh must have exactly one actor).
    #[pyo3(signature = (*args, **kwargs), name = "call_one")]
    fn py_call_one<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
    ) -> PyResult<Py<PyAny>> {
        self.call_one(py, args, kwargs)
    }

    /// Call the endpoint on all actors and return an iterator of Futures.
    #[pyo3(signature = (*args, **kwargs), name = "stream")]
    fn py_stream<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
    ) -> PyResult<Py<PyAny>> {
        self.stream(py, args, kwargs)
    }

    /// Send a message to all actors without waiting for responses (fire-and-forget).
    #[pyo3(signature = (*args, **kwargs), name = "broadcast")]
    fn py_broadcast<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
    ) -> PyResult<()> {
        self.broadcast(py, args, kwargs)
    }

    /// Get the rref result by calling the wrapped Remote's rref method.
    #[pyo3(signature = (*args, **kwargs))]
    fn rref<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
    ) -> PyResult<Py<PyAny>> {
        let kwargs_dict = kwargs.map_or_else(|| PyDict::new(py), |d| d.clone());
        self.inner.call_method(py, "rref", args, Some(&kwargs_dict))
    }

    /// Get the call name by delegating to the wrapped Remote's _call_name.
    fn _call_name(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        self.inner.call_method0(py, "_call_name")
    }

    /// Get the maybe_resolvable property from the wrapped RemoteImpl.
    #[getter]
    fn _maybe_resolvable(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        self.inner.getattr(py, "_maybe_resolvable")
    }

    /// Get the resolvable property from the wrapped RemoteImpl.
    #[getter]
    fn _resolvable(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        self.inner.getattr(py, "_resolvable")
    }

    /// Get the remote_impl from the wrapped RemoteImpl.
    /// This is needed for function_to_import_path() in function.py to work correctly.
    #[getter]
    fn _remote_impl(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        self.inner.getattr(py, "_remote_impl")
    }

    /// Propagation method for tensor shape inference.
    /// Delegates to the wrapped Remote's _propagate.
    fn _propagate<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyAny>,
        kwargs: &Bound<'py, PyAny>,
        fake_args: &Bound<'py, PyAny>,
        fake_kwargs: &Bound<'py, PyAny>,
    ) -> PyResult<Py<PyAny>> {
        self.inner
            .call_method1(py, "_propagate", (args, kwargs, fake_args, fake_kwargs))
    }

    /// Propagation for fetch operations.
    /// Delegates to the wrapped Remote's _fetch_propagate.
    fn _fetch_propagate<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyAny>,
        kwargs: &Bound<'py, PyAny>,
        fake_args: &Bound<'py, PyAny>,
        fake_kwargs: &Bound<'py, PyAny>,
    ) -> PyResult<Py<PyAny>> {
        self.inner.call_method1(
            py,
            "_fetch_propagate",
            (args, kwargs, fake_args, fake_kwargs),
        )
    }

    /// Propagation for pipe operations.
    /// Delegates to the wrapped Remote's _pipe_propagate.
    fn _pipe_propagate<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyAny>,
        kwargs: &Bound<'py, PyAny>,
        fake_args: &Bound<'py, PyAny>,
        fake_kwargs: &Bound<'py, PyAny>,
    ) -> PyResult<Py<PyAny>> {
        self.inner.call_method1(
            py,
            "_pipe_propagate",
            (args, kwargs, fake_args, fake_kwargs),
        )
    }

    /// Send a message with optional port for response.
    /// Delegates to the wrapped RemoteImpl's _send.
    fn _send<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: &Bound<'py, PyDict>,
        port: Option<Py<PyAny>>,
        selection: &str,
    ) -> PyResult<()> {
        self.inner.call_method(
            py,
            "_send",
            (args, kwargs),
            Some(&{
                let d = PyDict::new(py);
                d.set_item("port", port.unwrap_or_else(|| py.None()))?;
                d.set_item("selection", selection)?;
                d
            }),
        )?;
        Ok(())
    }

    /// Make RemoteEndpoint callable - delegates to rref() like Remote.__call__.
    #[pyo3(signature = (*args, **kwargs))]
    fn __call__<'py>(
        &self,
        py: Python<'py>,
        args: &Bound<'py, PyTuple>,
        kwargs: Option<&Bound<'py, PyDict>>,
    ) -> PyResult<Py<PyAny>> {
        self.rref(py, args, kwargs)
    }
}

pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> {
    module.add_class::<PyValueStream>()?;
    module.add_class::<ActorEndpoint>()?;
    module.add_class::<Remote>()?;

    Ok(())
}

#[derive(Named)]
struct PythonResponseMessageReducer;

impl CommReducer for PythonResponseMessageReducer {
    type Update = PythonMessage;

    fn reduce(&self, left: Self::Update, right: Self::Update) -> anyhow::Result<Self::Update> {
        Ok(ValueOverlay::try_from_runs(rle::merge_value_runs(
            left.into_overlay()?.into_runs(),
            right.into_overlay()?.into_runs(),
        ))?
        .into())
    }
}

inventory::submit! {
    ReducerFactory {
        typehash_f: <PythonResponseMessageReducer as Named>::typehash,
        builder_f: |_| Ok(Box::new(PythonResponseMessageReducer)),
    }
}

struct PythonResponseMessageAccumulator;

impl Accumulator for PythonResponseMessageAccumulator {
    type State = PythonMessage;
    type Update = PythonMessage;

    fn accumulate(&self, state: &mut Self::State, update: Self::Update) -> anyhow::Result<()> {
        *state = ValueOverlay::try_from_runs(rle::merge_value_runs(
            std::mem::take(state).into_overlay()?.into_runs(),
            update.into_overlay()?.into_runs(),
        ))?
        .into();

        Ok(())
    }

    fn reducer_spec(&self) -> Option<ReducerSpec> {
        Some(ReducerSpec {
            typehash: <PythonResponseMessageReducer as Named>::typehash(),
            builder_params: None,
        })
    }
}

#[cfg(test)]
mod tests {
    use hyperactor::ActorAddr;
    use hyperactor::mailbox::headers::OPERATION_ADVERB;
    use hyperactor::mailbox::headers::OPERATION_ENDPOINT;

    use super::*;

    /// Minimal `Endpoint` impl that only serves `get_qualified_name`.
    /// The default `build_operation_context_headers` consults no other
    /// method, so the rest are unreachable.
    struct TestEndpoint {
        qualified_name: Option<String>,
    }

    impl Endpoint for TestEndpoint {
        fn get_extent(&self, _py: Python<'_>) -> PyResult<Extent> {
            unreachable!()
        }
        fn get_method_name(&self) -> &str {
            unreachable!()
        }
        fn send_message<'py>(
            &self,
            _py: Python<'py>,
            _args: &Bound<'py, PyTuple>,
            _kwargs: Option<&Bound<'py, PyDict>>,
            _port_ref: Option<EitherPortRef>,
            _selection: AllOrChoose,
            _instance: &Instance<PythonActor>,
        ) -> PyResult<()> {
            unreachable!()
        }
        fn get_supervision_monitor(&self) -> Option<Arc<dyn Supervisable>> {
            None
        }
        fn get_qualified_name(&self) -> Option<String> {
            self.qualified_name.clone()
        }
        fn enter_endpoint_span(&self, _adverb: EndpointAdverb, _actor_id: &ActorAddr) -> SpanGuard {
            unreachable!()
        }
    }

    /// OC-1 request-side producer: each `EndpointAdverb` maps to the
    /// expected wire adverb string, and the qualified endpoint name
    /// from `get_qualified_name()` flows through to `OPERATION_ENDPOINT`.
    #[test]
    fn test_rc1_build_operation_context_headers_stamps_each_adverb() {
        let ep = TestEndpoint {
            qualified_name: Some("training.Philosopher.ping()".to_string()),
        };
        for (adverb, expected_adverb) in [
            (EndpointAdverb::Call, "call"),
            (EndpointAdverb::CallOne, "call_one"),
            (EndpointAdverb::Choose, "choose"),
            (EndpointAdverb::Stream, "stream"),
        ] {
            let headers = ep.build_operation_context_headers(adverb);
            assert_eq!(
                headers.get(OPERATION_ENDPOINT).as_deref(),
                Some("training.Philosopher.ping()"),
                "adverb {:?}: OPERATION_ENDPOINT",
                adverb,
            );
            assert_eq!(
                headers.get(OPERATION_ADVERB).as_deref(),
                Some(expected_adverb),
                "adverb {:?}: OPERATION_ADVERB",
                adverb,
            );
        }
    }

    /// OC-1 request-side producer: when the endpoint has no
    /// qualified name (e.g. `Remote::get_qualified_name` returns
    /// `None`), `OPERATION_ENDPOINT` is omitted while `OPERATION_ADVERB`
    /// is still stamped.
    #[test]
    fn test_rc1_build_operation_context_headers_omits_endpoint_when_no_qualified_name() {
        let ep = TestEndpoint {
            qualified_name: None,
        };
        let headers = ep.build_operation_context_headers(EndpointAdverb::CallOne);
        assert!(
            headers.get(OPERATION_ENDPOINT).is_none(),
            "OPERATION_ENDPOINT must be absent when endpoint has no qualified name",
        );
        assert_eq!(
            headers.get(OPERATION_ADVERB).as_deref(),
            Some("call_one"),
            "OPERATION_ADVERB should still be stamped",
        );
    }
}