monarch_hyperactor/

actor.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::collections::HashMap;
use std::collections::VecDeque;
use std::error::Error;
use std::fmt::Debug;
use std::future::pending;
use std::ops::Deref;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::Once;
use std::sync::OnceLock;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering as AtomicOrdering;
use std::time::SystemTime;

use async_trait::async_trait;
use hyperactor::Actor;
use hyperactor::ActorHandle;
use hyperactor::Context;
use hyperactor::Endpoint as _;
use hyperactor::Handler;
use hyperactor::Instance;
use hyperactor::OncePortHandle;
use hyperactor::Proc;
use hyperactor::RemoteSpawn;
use hyperactor::actor::ActorError;
use hyperactor::actor::ActorErrorKind;
use hyperactor::actor::ActorStatus;
use hyperactor::actor::Signal;
use hyperactor::context::Actor as ContextActor;
use hyperactor::mailbox::MessageEnvelope;
use hyperactor::mailbox::Undeliverable;
use hyperactor::mailbox::UndeliverableMessageError;
use hyperactor::mailbox::UndeliverableReason;
use hyperactor::supervision::ActorSupervisionEvent;
use hyperactor_config::Flattrs;
use hyperactor_mesh::ProcMeshRef;
use hyperactor_mesh::actor_mesh::ActorMeshRef;
use hyperactor_mesh::casting::update_undeliverable_envelope_for_casting;
use hyperactor_mesh::comm::multicast::CAST_POINT;
use hyperactor_mesh::comm::multicast::CastInfo;
use hyperactor_mesh::host_mesh::HostMeshRef;
use hyperactor_mesh::introspect::ActiveHandler;
use hyperactor_mesh::introspect::EXECUTION;
use hyperactor_mesh::introspect::Execution;
use hyperactor_mesh::supervision::MeshFailure;
use hyperactor_mesh::transport::default_bind_spec;
use hyperactor_mesh::value_mesh::ValueOverlay;
use monarch_types::PickledPyObject;
use monarch_types::SerializablePyErr;
use monarch_types::py_global;
use ndslice::Point;
use ndslice::extent;
use pyo3::IntoPyObjectExt;
use pyo3::exceptions::PyBaseException;
use pyo3::exceptions::PyRuntimeError;
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::PyDict;
use pyo3::types::PyList;
use pyo3::types::PyType;
use serde::Deserialize;
use serde::Serialize;
use serde_multipart::Part;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use typeuri::Named;

use crate::buffers::FrozenBuffer;
use crate::config::ACTOR_QUEUE_DISPATCH;
use crate::context::PyInstance;
use crate::local_state_broker::BrokerId;
use crate::local_state_broker::LocalStateBrokerMessage;
use crate::mailbox::EitherPortRef;
use crate::mailbox::PyMailbox;
use crate::mailbox::PythonUndeliverableMessageEnvelope;
use crate::metrics::ENDPOINT_ACTOR_COUNT;
use crate::metrics::ENDPOINT_ACTOR_ERROR;
use crate::metrics::ENDPOINT_ACTOR_LATENCY_US_HISTOGRAM;
use crate::metrics::ENDPOINT_ACTOR_PANIC;
use crate::pickle::PicklingState;
use crate::pickle::pickle_to_part;
use crate::proc::PyActorAddr;
use crate::pympsc;
use crate::pytokio::PyPythonTask;
use crate::pytokio::PythonTask;
use crate::runtime::GilSite;
use crate::runtime::get_tokio_runtime;
use crate::runtime::monarch_with_gil;
use crate::runtime::monarch_with_gil_blocking;
use crate::supervision::PyMeshFailure;

py_global!(
    unhandled_fault_hook_exception,
    "monarch._src.actor.supervision",
    "UnhandledFaultHookException"
);

#[pyclass(module = "monarch._rust_bindings.monarch_hyperactor.actor")]
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub enum UnflattenArg {
    Mailbox,
    PyObject,
}

#[pyclass(module = "monarch._rust_bindings.monarch_hyperactor.actor")]
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub enum MethodSpecifier {
    /// Call method 'name', send its return value to the response port.
    ReturnsResponse { name: String },
    /// Call method 'name', send the response port as the first argument.
    ExplicitPort { name: String },
    /// Construct the object
    Init {},
}

impl std::fmt::Display for MethodSpecifier {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.name())
    }
}

#[pymethods]
impl MethodSpecifier {
    #[getter(name)]
    fn py_name(&self) -> &str {
        self.name()
    }
}

impl MethodSpecifier {
    pub(crate) fn name(&self) -> &str {
        match self {
            MethodSpecifier::ReturnsResponse { name } => name,
            MethodSpecifier::ExplicitPort { name } => name,
            MethodSpecifier::Init {} => "__init__",
        }
    }
}

/// The payload of a single actor response, without rank information.
///
/// The rank is captured by the overlay's range key, so it is stripped
/// from the value to enable RLE dedup: two ranks returning the same
/// payload will have byte-identical values and can be coalesced into
/// a single run.
#[derive(Clone, Debug, Serialize, Deserialize, Named, PartialEq, Eq)]
pub enum PythonResponseMessage {
    Result {
        part: serde_multipart::Part,
        refs: Vec<MeshRef>,
    },
    Exception {
        part: serde_multipart::Part,
        refs: Vec<MeshRef>,
    },
}

wirevalue::register_type!(PythonResponseMessage);
wirevalue::register_type!(ValueOverlay<PythonResponseMessage>);

impl PythonResponseMessage {
    /// Decode this response's payload, reuniting its out-of-band `refs` table
    /// so mesh references reconstruct. Mirrors [`PythonMessage::decode`] for the
    /// accumulated (valuemesh / `.call()`) path.
    pub(crate) fn decode(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        let (part, refs) = match self {
            PythonResponseMessage::Result { part, refs }
            | PythonResponseMessage::Exception { part, refs } => (part, refs),
        };
        let mesh_references = refs.iter().cloned().map(Some).collect();
        let mut state = PicklingState::from_parts(part.clone(), VecDeque::new(), mesh_references);
        state.unpickle(py)
    }
}

/// Newtype wrapper around [`ValueOverlay<PythonResponseMessage>`] needed
/// because `PythonMessageKind` is a `#[pyclass]` enum, requiring all variant
/// fields to implement PyO3 traits. `ValueOverlay` is defined in another crate
/// and does not implement `PyClass`.
#[pyclass(frozen, module = "monarch._rust_bindings.monarch_hyperactor.actor")]
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct AccumulatedResponses(ValueOverlay<PythonResponseMessage>);

#[pyclass(module = "monarch._rust_bindings.monarch_hyperactor.actor")]
#[derive(Clone, Debug, Serialize, Deserialize, Named, PartialEq)]
pub enum PythonMessageKind {
    CallMethod {
        name: MethodSpecifier,
        response_port: Option<EitherPortRef>,
    },
    Result {
        rank: Option<usize>,
    },
    Exception {
        rank: Option<usize>,
    },
    Uninit {},
    CallMethodIndirect {
        name: MethodSpecifier,
        local_state_broker: (String, usize),
        id: usize,
        // specify whether the argument to unflatten the local mailbox,
        // or the next argument of the local state.
        unflatten_args: Vec<UnflattenArg>,
    },
    AccumulatedResponses(AccumulatedResponses),
}
wirevalue::register_type!(PythonMessageKind);

impl Default for PythonMessageKind {
    fn default() -> Self {
        PythonMessageKind::Uninit {}
    }
}

fn mailbox<'py, T: Actor>(py: Python<'py>, cx: &Context<'_, T>) -> Bound<'py, PyAny> {
    let mailbox: PyMailbox = cx.mailbox_for_py().clone().into();
    mailbox.into_bound_py_any(py).unwrap()
}

/// A serializable reference to a mesh (actor, proc, or host).
///
/// Serialized as a typed multipart part via [`MeshRefRepr`]: under the multipart
/// serializer each `MeshRef` becomes its own typed part, and inline bincode
/// elsewhere.
#[derive(Clone, Debug, Named, PartialEq, Eq)]
pub enum MeshRef {
    Actor(Box<ActorMeshRef<PythonActor>>),
    Proc(Box<ProcMeshRef>),
    Host(Box<HostMeshRef>),
}

/// Wire representation of [`MeshRef`] stored in a typed multipart part.
#[doc(hidden)]
#[derive(Clone, Debug, Serialize, Deserialize, Named)]
pub enum MeshRefRepr {
    Actor(Box<ActorMeshRef<PythonActor>>),
    Proc(Box<ProcMeshRef>),
    Host(Box<HostMeshRef>),
}

impl TryFrom<&MeshRef> for MeshRefRepr {
    type Error = serde_multipart::Error;
    fn try_from(m: &MeshRef) -> serde_multipart::Result<Self> {
        Ok(match m {
            MeshRef::Actor(r) => MeshRefRepr::Actor(r.clone()),
            MeshRef::Proc(r) => MeshRefRepr::Proc(r.clone()),
            MeshRef::Host(r) => MeshRefRepr::Host(r.clone()),
        })
    }
}

impl TryFrom<MeshRefRepr> for MeshRef {
    type Error = serde_multipart::Error;
    fn try_from(r: MeshRefRepr) -> serde_multipart::Result<Self> {
        Ok(match r {
            MeshRefRepr::Actor(r) => MeshRef::Actor(r),
            MeshRefRepr::Proc(r) => MeshRef::Proc(r),
            MeshRefRepr::Host(r) => MeshRef::Host(r),
        })
    }
}

serde_multipart::part_codec! {
    impl MeshRef
    {
        type Repr = MeshRefRepr;
    }
}

impl MeshRef {
    /// Reconstruct the Python mesh wrapper this reference points at.
    ///
    /// Mirrors the `py_*_from_bytes` reconstructors, but takes an
    /// already-deserialized [`MeshRef`] from the message's `refs` table.
    pub(crate) fn reconstruct(self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        match self {
            MeshRef::Proc(r) => {
                Ok(Py::new(py, crate::proc_mesh::PyProcMesh::new_ref(*r))?.into_any())
            }
            MeshRef::Host(r) => {
                Ok(Py::new(py, crate::host_mesh::PyHostMesh::new_ref(*r))?.into_any())
            }
            MeshRef::Actor(r) => {
                let inner = crate::actor_mesh::PythonActorMeshImpl::new_ref(*r);
                let async_mesh = crate::actor_mesh::AsyncActorMesh::from_impl(Arc::new(inner));
                let mesh = crate::actor_mesh::PythonActorMesh::from_impl(Arc::from(async_mesh));
                Ok(Py::new(py, mesh)?.into_any())
            }
        }
    }
}

/// An opaque carrier so a `MeshRef` can ride in `PythonMessage.refs` across
/// the Python boundary (the message getter out, the `PicklingState` ctor in).
#[pyclass(frozen, module = "monarch._rust_bindings.monarch_hyperactor.actor")]
pub struct PyMeshRef {
    pub(crate) inner: MeshRef,
}

impl<'py> IntoPyObject<'py> for MeshRef {
    type Target = PyMeshRef;
    type Output = Bound<'py, PyMeshRef>;
    type Error = PyErr;

    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
        Bound::new(py, PyMeshRef { inner: self })
    }
}

/// Extract the serializable [`MeshRef`] from a resolved mesh wrapper (the
/// inverse of [`MeshRef::reconstruct`]), for the sender-side pending fill.
pub(crate) fn mesh_ref_from_pyobject(value: &Bound<'_, PyAny>) -> PyResult<MeshRef> {
    if let Ok(m) = value.downcast::<crate::proc_mesh::PyProcMesh>() {
        return Ok(MeshRef::Proc(Box::new(m.borrow().mesh_ref()?)));
    }
    if let Ok(m) = value.downcast::<crate::host_mesh::PyHostMesh>() {
        return Ok(MeshRef::Host(Box::new(m.borrow().mesh_ref().map_err(
            |e| pyo3::exceptions::PyValueError::new_err(e.to_string()),
        )?)));
    }
    if let Ok(m) = value.downcast::<crate::actor_mesh::PythonActorMesh>() {
        return Ok(MeshRef::Actor(Box::new(m.borrow().get_inner().mesh_ref()?)));
    }
    Err(pyo3::exceptions::PyRuntimeError::new_err(
        "pending pickle did not resolve to a mesh reference",
    ))
}

#[pyclass(frozen, module = "monarch._rust_bindings.monarch_hyperactor.actor")]
#[derive(Clone, Serialize, Deserialize, Named, Default, PartialEq)]
pub struct PythonMessage {
    pub kind: PythonMessageKind,
    pub message: Part,
    /// Mesh references carried out-of-band from the pickled `message`.
    pub refs: Vec<MeshRef>,
}

/// Extract the endpoint method name from a [`PythonMessage`].
fn python_message_endpoint_name(msg: &PythonMessage) -> Option<String> {
    match &msg.kind {
        PythonMessageKind::CallMethod { name, .. }
        | PythonMessageKind::CallMethodIndirect { name, .. } => Some(name.name().to_string()),
        _ => None,
    }
}

// We use manual `submit!` instead of `register_type!` because PythonMessage is a
// struct, so the default `endpoint_name` (which delegates to `arm_unchecked`)
// always returns None. The custom implementation inspects `PythonMessageKind` to
// extract the method name. This registration handles direct (non-cast) dispatch.
wirevalue::submit! {
    wirevalue::TypeInfo {
        typename: <PythonMessage as wirevalue::Named>::typename,
        typehash: <PythonMessage as wirevalue::Named>::typehash,
        typeid: <PythonMessage as wirevalue::Named>::typeid,
        port: <PythonMessage as wirevalue::Named>::port,
        dump: Some(<PythonMessage as wirevalue::NamedDumpable>::dump),
        arm_unchecked: <PythonMessage as wirevalue::Named>::arm_unchecked,
        endpoint_name: |ptr| {
            // SAFETY: ptr points to a PythonMessage.
            let msg = unsafe { &*(ptr as *const PythonMessage) };
            python_message_endpoint_name(msg)
        },
    }
}

impl From<ValueOverlay<PythonResponseMessage>> for PythonMessage {
    fn from(overlay: ValueOverlay<PythonResponseMessage>) -> Self {
        PythonMessage {
            kind: PythonMessageKind::AccumulatedResponses(AccumulatedResponses(overlay)),
            message: Default::default(),
            refs: Vec::new(),
        }
    }
}

impl PythonMessage {
    /// Consume this message and extract a `ValueOverlay<PythonResponseMessage>`.
    ///
    /// Handles both already-collected responses and leaf `Result`/`Exception`
    /// messages by wrapping them in a single-run overlay.
    pub fn into_overlay(self) -> anyhow::Result<ValueOverlay<PythonResponseMessage>> {
        match self.kind {
            PythonMessageKind::AccumulatedResponses(overlay) => Ok(overlay.0),
            PythonMessageKind::Result { rank, .. } => {
                let rank = rank.expect("accumulated response should have a rank");
                let mut overlay = ValueOverlay::new();
                overlay.push_run(
                    rank..rank + 1,
                    PythonResponseMessage::Result {
                        part: self.message,
                        refs: self.refs,
                    },
                )?;
                Ok(overlay)
            }
            PythonMessageKind::Exception { rank, .. } => {
                let rank = rank.expect("accumulated exception should have a rank");
                let mut overlay = ValueOverlay::new();
                overlay.push_run(
                    rank..rank + 1,
                    PythonResponseMessage::Exception {
                        part: self.message,
                        refs: self.refs,
                    },
                )?;
                Ok(overlay)
            }
            other => {
                anyhow::bail!(
                    "unexpected message kind {:?} in collected responses reducer",
                    other
                );
            }
        }
    }
}

struct ResolvedCallMethod {
    method: MethodSpecifier,
    bytes: FrozenBuffer,
    local_state: Option<Py<PyAny>>,
    mesh_references: Vec<MeshRef>,
    /// Implements PortProtocol
    /// Concretely either a Port, DroppingPort, or LocalPort
    response_port: ResponsePort,
}

enum ResponsePort {
    Dropping,
    Port(Port),
    Local(LocalPort),
}

impl ResponsePort {
    fn into_py_any(self, py: Python<'_>) -> PyResult<Py<PyAny>> {
        match self {
            ResponsePort::Dropping => DroppingPort.into_py_any(py),
            ResponsePort::Port(port) => port.into_py_any(py),
            ResponsePort::Local(port) => port.into_py_any(py),
        }
    }
}

/// Message sent through the queue in queue-dispatch mode.
/// Contains pre-resolved components ready for Python consumption.
#[pyclass(frozen, module = "monarch._rust_bindings.monarch_hyperactor.actor")]
pub struct QueuedMessage {
    #[pyo3(get)]
    pub context: Py<crate::context::PyContext>,
    #[pyo3(get)]
    pub method: MethodSpecifier,
    #[pyo3(get)]
    pub bytes: FrozenBuffer,
    #[pyo3(get)]
    pub local_state: Py<PyAny>,
    #[pyo3(get)]
    pub refs: Py<PyAny>,
    #[pyo3(get)]
    pub response_port: Py<PyAny>,
}

impl PythonMessage {
    pub fn new_from_buf(kind: PythonMessageKind, message: impl Into<Part>) -> Self {
        Self::new_from_buf_with_refs(kind, message, Vec::new())
    }

    pub fn new_from_buf_with_refs(
        kind: PythonMessageKind,
        message: impl Into<Part>,
        refs: Vec<MeshRef>,
    ) -> Self {
        Self {
            kind,
            message: message.into(),
            refs,
        }
    }

    pub fn into_rank(self, rank: usize) -> Self {
        let rank = Some(rank);
        match self.kind {
            PythonMessageKind::Result { .. } => PythonMessage {
                kind: PythonMessageKind::Result { rank },
                message: self.message,
                refs: self.refs,
            },
            PythonMessageKind::Exception { .. } => PythonMessage {
                kind: PythonMessageKind::Exception { rank },
                message: self.message,
                refs: self.refs,
            },
            _ => panic!("PythonMessage is not a response but {:?}", self),
        }
    }
    async fn resolve_indirect_call(
        self,
        cx: &Context<'_, PythonActor>,
    ) -> anyhow::Result<ResolvedCallMethod> {
        match self.kind {
            PythonMessageKind::CallMethodIndirect {
                name,
                local_state_broker,
                id,
                unflatten_args,
            } => {
                let broker = BrokerId::new(local_state_broker).resolve(cx).await;
                let (send, recv) = cx.open_once_port();
                broker.post(cx, LocalStateBrokerMessage::Get(id, send));
                let state = recv.recv().await?;
                let mut state_it = state.state.into_iter();
                monarch_with_gil(GilSite::EndpointDispatch, |py| {
                    let mailbox = mailbox(py, cx);
                    let local_state = Some(
                        PyList::new(
                            py,
                            unflatten_args.into_iter().map(|x| -> Bound<'_, PyAny> {
                                match x {
                                    UnflattenArg::Mailbox => mailbox.clone(),
                                    UnflattenArg::PyObject => {
                                        state_it.next().unwrap().into_bound(py)
                                    }
                                }
                            }),
                        )
                        .unwrap()
                        .into(),
                    );
                    let response_port = ResponsePort::Local(LocalPort {
                        instance: cx.into(),
                        inner: Some(state.response_port),
                    });
                    Ok(ResolvedCallMethod {
                        method: name,
                        bytes: FrozenBuffer {
                            inner: self.message.into_bytes(),
                        },
                        local_state,
                        mesh_references: self.refs,
                        response_port,
                    })
                })
                .await
            }
            PythonMessageKind::CallMethod {
                name,
                response_port,
            } => {
                let method_name = name.name().to_string();
                let response_port = response_port.map_or(ResponsePort::Dropping, |port_ref| {
                    let point = cx.cast_point();
                    // Carry operation context onto the reply: copy
                    // OPERATION_*-marked keys from the inbound
                    // envelope, falling back to the method name when
                    // the caller didn't stamp.
                    let mut reply_headers = hyperactor_config::Flattrs::new();
                    hyperactor_config::attrs::copy_marked_flattrs(
                        &mut reply_headers,
                        cx.headers(),
                        hyperactor_config::attrs::OPERATION_CONTEXT_HEADER,
                    );
                    if reply_headers
                        .get(hyperactor::mailbox::headers::OPERATION_ENDPOINT)
                        .is_none()
                    {
                        reply_headers.set(
                            hyperactor::mailbox::headers::OPERATION_ENDPOINT,
                            format!("{}()", method_name),
                        );
                    }
                    ResponsePort::Port(Port::with_reply_headers(
                        port_ref,
                        cx.instance().clone_for_py(),
                        Some(point.rank()),
                        reply_headers,
                    ))
                });
                Ok(ResolvedCallMethod {
                    method: name,
                    bytes: FrozenBuffer {
                        inner: self.message.into_bytes(),
                    },
                    local_state: None,
                    mesh_references: self.refs,
                    response_port,
                })
            }
            _ => {
                panic!("unexpected message kind {:?}", self.kind)
            }
        }
    }
}

impl std::fmt::Debug for PythonMessage {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PythonMessage")
            .field("kind", &self.kind)
            .field(
                "message",
                &wirevalue::HexFmt(&(*self.message.to_bytes())[..]).to_string(),
            )
            .field("refs", &self.refs.len())
            .finish()
    }
}

#[pymethods]
impl PythonMessage {
    #[new]
    #[pyo3(signature = (kind, message, refs))]
    pub fn new(
        kind: PythonMessageKind,
        message: PyRef<'_, FrozenBuffer>,
        refs: &Bound<'_, PyList>,
    ) -> PyResult<Self> {
        let mesh_refs: Vec<MeshRef> = refs
            .iter()
            .map(|item| Ok(item.downcast::<PyMeshRef>()?.borrow().inner.clone()))
            .collect::<PyResult<_>>()?;
        Ok(PythonMessage::new_from_buf_with_refs(
            kind,
            message.inner.clone(),
            mesh_refs,
        ))
    }

    #[getter]
    fn kind(&self) -> PythonMessageKind {
        self.kind.clone()
    }

    /// Decode this message's payload, reuniting the out-of-band `refs` table so
    /// the `pop_mesh_reference` sentinels in the pickle stream resolve. The raw
    /// bytes are deliberately not exposed: a payload can only be read back
    /// through here, so a decode can never silently drop mesh references.
    #[pyo3(signature = (local_state=None))]
    fn decode(
        &self,
        py: Python<'_>,
        local_state: Option<&Bound<'_, PyList>>,
    ) -> PyResult<Py<PyAny>> {
        let tensor_engine_references: VecDeque<Py<PyAny>> = local_state
            .map(|list| list.iter().map(|item| item.unbind()).collect())
            .unwrap_or_default();
        let mesh_references: VecDeque<Option<MeshRef>> =
            self.refs.iter().cloned().map(Some).collect();
        let mut state = PicklingState::from_parts(
            self.message.clone(),
            tensor_engine_references,
            mesh_references,
        );
        state.unpickle(py)
    }

    #[getter]
    fn refs(&self) -> Vec<MeshRef> {
        self.refs.clone()
    }
}

#[pyclass(module = "monarch._rust_bindings.monarch_hyperactor.actor")]
pub(super) struct PythonActorHandle {
    pub(super) inner: ActorHandle<PythonActor>,
}

#[pymethods]
impl PythonActorHandle {
    // TODO: do the pickling in rust
    fn send(&self, instance: &PyInstance, message: &PythonMessage) -> PyResult<()> {
        self.inner.post(instance.deref(), message.clone());
        Ok(())
    }

    fn bind(&self) -> PyActorAddr {
        self.inner.bind::<PythonActor>().into_actor_addr().into()
    }
}

/// Dispatch mode for Python actors.
#[derive(Debug)]
pub enum PythonActorDispatchMode {
    /// Direct dispatch: Rust acquires the GIL and calls Python handlers directly.
    Direct,
    /// Queue dispatch: Rust enqueues messages to a channel; Python dequeues and dispatches.
    Queue {
        /// Channel sender for enqueuing messages to Python.
        sender: pympsc::Sender,
        /// Channel receiver, taken during Actor::init to start the message loop.
        receiver: Option<pympsc::PyReceiver>,
    },
}

// In-flight handler execution tracking for a Python actor -- the producer
// side of the mesh `execution` field. Per-actor Rust state, read GIL-free
// by the introspect seam. Producer invariants (PE-*, monarch_hyperactor-
// local; documented inline, no registry -- not our crate):
//   PE-2: only the real user-method invocation is bracketed (in
//         `_Actor.handle`), never Init/unpickling/plumbing.
//   PE-3: the snapshot reads `Arc` state only (atomic load + `try_lock`);
//         the `Mutex` is held only for an insert/remove, never across user
//         code, so a handler wedged holding the GIL never blocks the
//         snapshot and the read never touches the GIL.
//   PE-4: tokens are >= 1; `0` is a reserved no-op sentinel (returned by
//         the binding when an instance has no tracker), so the
//         unconditional Python `finally` cannot collide it with a real
//         token.

/// EX-4 cap: at most this many distinct in-flight handler names are
/// reported per snapshot; `truncated` is set when exceeded.
const MAX_ACTIVE_HANDLERS: usize = 64;

/// One in-flight handler invocation.
#[derive(Debug)]
struct ActiveEntry {
    name: String,
    started_at: SystemTime,
}

/// Per-actor in-flight handler tracker. Cheap to read concurrently: the
/// count is a lock-free atomic and the per-handler detail sits behind a
/// `try_lock` held only for an insert/remove (never across user code), so
/// a wedged actor stays introspectable (PE-3).
#[derive(Debug)]
pub(crate) struct ExecutionTracker {
    /// Lock-free count of in-flight invocations; always readable.
    active_count: AtomicU64,
    /// Monotonic token source, initialized to 1 so issued tokens are
    /// `>= 1` and `0` stays reserved as the no-op sentinel (PE-4).
    next_token: AtomicU64,
    /// token -> entry for the in-flight invocations.
    handlers: Mutex<HashMap<u64, ActiveEntry>>,
}

/// Aggregate raw in-flight entries into the reported per-handler view:
/// grouped by handler name, oldest-first with a stable tie-break on
/// `name`, capped at `max` (EX-4). Pure, so it can be unit-tested with
/// explicit timestamps.
fn aggregate_active(
    handlers: &HashMap<u64, ActiveEntry>,
    max: usize,
) -> (Vec<ActiveHandler>, bool) {
    let mut by_name: HashMap<&str, (u64, SystemTime)> = HashMap::new();
    for entry in handlers.values() {
        let slot = by_name
            .entry(entry.name.as_str())
            .or_insert((0, entry.started_at));
        slot.0 += 1;
        if entry.started_at < slot.1 {
            slot.1 = entry.started_at;
        }
    }
    let mut out: Vec<ActiveHandler> = by_name
        .into_iter()
        .map(|(name, (active_count, oldest_since))| ActiveHandler {
            name: name.to_string(),
            active_count,
            oldest_since,
        })
        .collect();
    // EX-4: oldest-first, stable tie-break on name.
    out.sort_by(|a, b| {
        a.oldest_since
            .cmp(&b.oldest_since)
            .then_with(|| a.name.cmp(&b.name))
    });
    let truncated = out.len() > max;
    if truncated {
        out.truncate(max);
    }
    (out, truncated)
}

impl ExecutionTracker {
    pub(crate) fn new() -> Self {
        Self {
            active_count: AtomicU64::new(0),
            next_token: AtomicU64::new(1),
            handlers: Mutex::new(HashMap::new()),
        }
    }

    /// Record the start of a handler invocation; returns its token.
    pub(crate) fn start(&self, name: String) -> u64 {
        let token = self.next_token.fetch_add(1, AtomicOrdering::Relaxed);
        self.handlers
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .insert(
                token,
                ActiveEntry {
                    name,
                    started_at: SystemTime::now(),
                },
            );
        self.active_count.fetch_add(1, AtomicOrdering::Relaxed);
        token
    }

    /// Record the end of a handler invocation. Idempotent (a token is
    /// removed at most once) and a no-op for the `0` sentinel (PE-4).
    pub(crate) fn finish(&self, token: u64) {
        if token == 0 {
            return;
        }
        let removed = self
            .handlers
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .remove(&token)
            .is_some();
        if removed {
            self.active_count.fetch_sub(1, AtomicOrdering::Relaxed);
        }
    }

    /// Point-in-time snapshot for the introspect seam. Never blocks: the
    /// count is read lock-free and the per-handler detail is best-effort
    /// behind `try_lock` (EX-2: a miss yields `complete: false`, it never
    /// drops the field).
    pub(crate) fn snapshot(&self) -> Execution {
        let active_count = self.active_count.load(AtomicOrdering::Relaxed);
        match self.handlers.try_lock() {
            Ok(guard) => {
                let (active_handlers, truncated) = aggregate_active(&guard, MAX_ACTIVE_HANDLERS);
                Execution {
                    active_count,
                    active_handlers,
                    complete: true,
                    truncated,
                }
            }
            Err(_) => Execution {
                active_count,
                active_handlers: Vec::new(),
                complete: false,
                truncated: false,
            },
        }
    }
}

#[cfg(test)]
mod execution_tracker_tests {
    use std::time::Duration;
    use std::time::UNIX_EPOCH;

    use super::*;

    fn at(secs: u64) -> SystemTime {
        UNIX_EPOCH + Duration::from_secs(secs)
    }

    #[test]
    fn aggregates_by_name_oldest_first() {
        let mut h = HashMap::new();
        h.insert(
            1,
            ActiveEntry {
                name: "b".to_string(),
                started_at: at(10),
            },
        );
        h.insert(
            2,
            ActiveEntry {
                name: "a".to_string(),
                started_at: at(20),
            },
        );
        h.insert(
            3,
            ActiveEntry {
                name: "a".to_string(),
                started_at: at(30),
            },
        );
        let (out, truncated) = aggregate_active(&h, MAX_ACTIVE_HANDLERS);
        assert!(!truncated);
        assert_eq!(out.len(), 2);
        // Oldest-first: b (10) before a (20).
        assert_eq!(out[0].name, "b");
        assert_eq!(out[0].active_count, 1);
        assert_eq!(out[0].oldest_since, at(10));
        // "a" aggregates two invocations; oldest_since is the min (20).
        assert_eq!(out[1].name, "a");
        assert_eq!(out[1].active_count, 2);
        assert_eq!(out[1].oldest_since, at(20));
    }

    #[test]
    fn tie_break_on_name_when_same_oldest() {
        let mut h = HashMap::new();
        h.insert(
            1,
            ActiveEntry {
                name: "zebra".to_string(),
                started_at: at(5),
            },
        );
        h.insert(
            2,
            ActiveEntry {
                name: "alpha".to_string(),
                started_at: at(5),
            },
        );
        let (out, _) = aggregate_active(&h, MAX_ACTIVE_HANDLERS);
        assert_eq!(out[0].name, "alpha");
        assert_eq!(out[1].name, "zebra");
    }

    #[test]
    fn truncates_to_n_oldest() {
        let mut h = HashMap::new();
        for i in 0..(MAX_ACTIVE_HANDLERS as u64 + 6) {
            h.insert(
                i,
                ActiveEntry {
                    name: format!("h{:03}", i),
                    started_at: at(i),
                },
            );
        }
        let (out, truncated) = aggregate_active(&h, MAX_ACTIVE_HANDLERS);
        assert!(truncated);
        assert_eq!(out.len(), MAX_ACTIVE_HANDLERS);
        // Prefix of the N oldest.
        assert_eq!(out[0].name, "h000");
        assert_eq!(
            out[MAX_ACTIVE_HANDLERS - 1].name,
            format!("h{:03}", MAX_ACTIVE_HANDLERS - 1)
        );
    }

    #[test]
    fn start_assigns_nonzero_distinct_tokens() {
        let t = ExecutionTracker::new();
        let a = t.start("a".to_string());
        let b = t.start("b".to_string());
        assert!(a >= 1);
        assert!(b >= 1);
        assert_ne!(a, b);
        let snap = t.snapshot();
        assert_eq!(snap.active_count, 2);
        assert!(snap.complete);
        assert_eq!(snap.active_handlers.len(), 2);
    }

    #[test]
    fn finish_is_idempotent_and_zero_is_noop() {
        let t = ExecutionTracker::new();
        let tok = t.start("a".to_string());
        t.finish(tok);
        assert_eq!(t.snapshot().active_count, 0);
        // Double-finish must not underflow the count.
        t.finish(tok);
        assert_eq!(t.snapshot().active_count, 0);
        // The 0 sentinel is a no-op.
        t.finish(0);
        assert_eq!(t.snapshot().active_count, 0);
    }
}

/// An actor for which message handlers are implemented in Python.
#[derive(Debug)]
#[hyperactor::export(
    handlers = [
        PythonMessage,
        MeshFailure,
    ],
)]
#[hyperactor::spawnable]
pub struct PythonActor {
    /// The Python object that we delegate message handling to.
    actor: Py<PyAny>,
    /// Stores a reference to the Python event loop to run Python coroutines on.
    task_locals: pyo3_async_runtimes::TaskLocals,
    /// Instance object that we keep across handle calls so that we can store
    /// information from the Init (spawn rank, controller) and provide it to other calls.
    instance: Option<Py<crate::context::PyInstance>>,
    /// Dispatch mode for this actor.
    dispatch_mode: PythonActorDispatchMode,
    /// The location in the actor mesh at which this actor was spawned.
    spawn_point: OnceLock<Option<Point>>,
    /// Initial message to process during PythonActor::init.
    init_message: Option<PythonMessage>,
    /// User-provided mesh base-name string plumbed from
    /// `PythonActorParams`. This is the base name the caller
    /// supplied when the mesh was spawned, narrowly used to populate
    /// `MeshFailure.actor_mesh_name` on the direct actor-handled
    /// supervision path without a lookup. It is not actor display
    /// text (`display_name` handles that) and it is not a general
    /// side channel; downstream code must not consume this field for
    /// any other purpose.
    mesh_base_name: Option<String>,

    /// Per-actor in-flight handler tracker (producer of the mesh
    /// `execution` field). Read GIL-free by the introspect seam; a clone
    /// of this `Arc` is injected into the actor's `PyInstance` so
    /// `_Actor.handle` can bracket each invocation.
    execution_tracker: Arc<ExecutionTracker>,
}

impl PythonActor {
    pub(crate) fn new(
        actor_type: PickledPyObject,
        init_message: Option<PythonMessage>,
        spawn_point: Option<Point>,
        mesh_base_name: Option<String>,
    ) -> Result<Self, anyhow::Error> {
        let use_queue_dispatch = hyperactor_config::global::get(ACTOR_QUEUE_DISPATCH);
        if !use_queue_dispatch {
            static WARNED: Once = Once::new();
            WARNED.call_once(|| {
                tracing::warn!(
                    "actor_queue_dispatch=false is deprecated and direct dispatch will be removed in a future release"
                );
            });
        }

        Ok(monarch_with_gil_blocking(
            GilSite::ActorConstruct,
            |py| -> Result<Self, SerializablePyErr> {
                let unpickled = actor_type.unpickle(py)?;
                let class_type: &Bound<'_, PyType> = unpickled.downcast()?;
                let actor: Py<PyAny> = class_type.call0()?.into_py_any(py)?;

                let task_locals = Python::detach(py, create_task_locals);

                let dispatch_mode = if use_queue_dispatch {
                    let (sender, receiver) = pympsc::channel().map_err(|e| {
                        let py_err = PyRuntimeError::new_err(e.to_string());
                        SerializablePyErr::from(py, &py_err)
                    })?;
                    PythonActorDispatchMode::Queue {
                        sender,
                        receiver: Some(receiver),
                    }
                } else {
                    PythonActorDispatchMode::Direct
                };

                Ok(Self {
                    actor,
                    task_locals,
                    instance: None,
                    dispatch_mode,
                    spawn_point: OnceLock::from(spawn_point),
                    init_message,
                    mesh_base_name,
                    execution_tracker: Arc::new(ExecutionTracker::new()),
                })
            },
        )?)
    }

    fn cancel_tasks_and_stop_python_loop(
        py: Python<'_>,
        task_locals: &pyo3_async_runtimes::TaskLocals,
    ) -> PyResult<()> {
        let asyncio = py.import("asyncio")?;
        let event_loop = task_locals.event_loop(py);
        let tasks = asyncio.call_method1("all_tasks", (&event_loop,))?;
        let mut has_tasks = false;
        for task in tasks.try_iter()? {
            let task = task?;
            let cancel = task.getattr("cancel")?;
            event_loop.call_method1("call_soon_threadsafe", (cancel,))?;
            has_tasks = true;
        }
        if has_tasks {
            asyncio
                .call_method1(
                    "run_coroutine_threadsafe",
                    (asyncio.call_method1("sleep", (0,))?, &event_loop),
                )?
                .call_method0("result")?;
        }
        let stop = event_loop.getattr("stop")?;
        event_loop.call_method1("call_soon_threadsafe", (stop,))?;
        Ok(())
    }

    fn cancel_pending_python_tasks_and_stop_loop(&self) -> anyhow::Result<()> {
        let task_locals = &self.task_locals;
        monarch_with_gil_blocking(GilSite::Stop, |py| -> anyhow::Result<()> {
            Self::cancel_tasks_and_stop_python_loop(py, task_locals)
                .map_err(|err| anyhow::Error::from(SerializablePyErr::from(py, &err)))?;
            Ok(())
        })
    }

    /// Get-or-create the actor's cached `PyInstance`, injecting a clone of
    /// the execution tracker (PE-1) so `_Actor.handle` can bracket each
    /// invocation. All four `self.instance` creation sites route through
    /// this so the tracker is never silently absent -- notably the
    /// supervision path, which can run before the first endpoint.
    fn ensure_py_instance(
        &mut self,
        py: Python<'_>,
        src: impl Into<crate::context::PyInstance>,
    ) -> Py<crate::context::PyInstance> {
        let tracker = self.execution_tracker.clone();
        self.instance
            .get_or_insert_with(|| {
                let mut inst: crate::context::PyInstance = src.into();
                inst.set_execution_tracker(tracker);
                inst.into_pyobject(py).unwrap().into()
            })
            .clone_ref(py)
    }

    /// Bootstrap the root client actor, creating a new proc for it.
    /// This is the legacy entry point that creates its own proc.
    pub(crate) fn bootstrap_client(py: Python<'_>) -> (&'static Instance<Self>, ActorHandle<Self>) {
        static ROOT_CLIENT_INSTANCE: OnceLock<Instance<PythonActor>> = OnceLock::new();

        let client_proc = Proc::direct(
            default_bind_spec().binding_addr(),
            "mesh_root_client_proc".into(),
        )
        .unwrap();

        Self::bootstrap_client_inner(py, client_proc, &ROOT_CLIENT_INSTANCE)
    }

    /// Bootstrap the client proc, storing the root client instance in given static.
    /// This is passed in because we require storage, as the instance is shared.
    /// This can be simplified when we remove v0.
    pub(crate) fn bootstrap_client_inner(
        py: Python<'_>,
        client_proc: Proc,
        root_client_instance: &'static OnceLock<Instance<PythonActor>>,
    ) -> (&'static Instance<Self>, ActorHandle<Self>) {
        let actor_mesh_mod = py
            .import("monarch._src.actor.actor_mesh")
            .expect("import actor_mesh");
        let root_client_class = actor_mesh_mod
            .getattr("RootClientActor")
            .expect("get RootClientActor");

        let actor_type =
            PickledPyObject::pickle(&actor_mesh_mod.getattr("_Actor").expect("get _Actor"))
                .expect("pickle _Actor");

        let init_frozen_buffer: FrozenBuffer = root_client_class
            .call_method0("_pickled_init_args")
            .expect("call RootClientActor._pickled_init_args")
            .extract()
            .expect("extract FrozenBuffer from _pickled_init_args");
        let init_message = PythonMessage::new_from_buf(
            PythonMessageKind::CallMethod {
                name: MethodSpecifier::Init {},
                response_port: None,
            },
            init_frozen_buffer,
        );

        let mut actor = PythonActor::new(
            actor_type,
            Some(init_message),
            Some(extent!().point_of_rank(0).unwrap()),
            None, // root client actor has no user-facing mesh name
        )
        .expect("create client PythonActor");

        let ai = client_proc
            .actor_instance(
                root_client_class
                    .getattr("name")
                    .expect("get RootClientActor.name")
                    .extract()
                    .expect("extract RootClientActor.name"),
            )
            .expect("root instance create");

        let handle = ai.handle;
        let signal_rx = ai.signal;
        let supervision_rx = ai.supervision;
        let work_rx = ai.work;

        root_client_instance
            .set(ai.instance)
            .map_err(|_| "already initialized root client instance")
            .unwrap();
        let instance = root_client_instance.get().unwrap();

        // The root client PythonActor uses a custom run loop that
        // bypasses Actor::init, so mark it as system explicitly
        // (matching GlobalClientActor::fresh_instance).
        instance.set_system();

        // Bind to ensure the Undeliverable<MessageEnvelope> port is bound.
        let _client_ref = handle.bind::<PythonActor>();

        get_tokio_runtime().spawn(async move {
            // This is gross. Sorry.
            actor.init(instance).await.unwrap();

            let mut signal_rx = signal_rx;
            let mut supervision_rx = supervision_rx;
            let mut work_rx = work_rx;
            let mut need_drain = false;
            let mut err = 'messages: loop {
                tokio::select! {
                    work = work_rx.recv() => {
                        let work = work.expect("inconsistent work queue state");
                        if let Err(err) = work.handle(&mut actor, instance).await {
                            // Check for UnhandledFaultHookException on the raw
                            // anyhow::Error before wrapping in ActorErrorKind.
                            // If __supervise__ already processed the supervision
                            // event and the hook raised, don't re-handle it via
                            // handle_supervision_event — that would call
                            // __supervise__ a second time.
                            let is_hook_exception = monarch_with_gil(GilSite::Supervise, |py| {
                                err.downcast_ref::<pyo3::PyErr>()
                                    .is_some_and(|pyerr| {
                                        pyerr.is_instance(
                                            py,
                                            &unhandled_fault_hook_exception(py),
                                        )
                                    })
                            }).await;

                            let kind = ActorErrorKind::processing(err);
                            let err = ActorError {
                                actor_id: Box::new(instance.self_addr().clone()),
                                kind: Box::new(kind),
                            };

                            if is_hook_exception {
                                break Some(err);
                            }

                            // Give the actor a chance to handle the error produced
                            // in its own message handler. This is important because
                            // we want Undeliverable<MessageEnvelope>, which returns
                            // an Err typically, to create a supervision event and
                            // call __supervise__.
                            let supervision_event = actor_error_to_event(instance, &actor, err);
                            // If the immediate supervision event isn't handled, continue with
                            // exiting the loop.
                            // Else, continue handling messages.
                            if let Err(err) = instance.handle_supervision_event(&mut actor, supervision_event).await {
                                while let Ok(supervision_event) = supervision_rx.try_recv() {
                                    if let Err(err) = instance.handle_supervision_event(&mut actor, supervision_event).await {
                                        break 'messages Some(err);
                                    }
                                }
                                break Some(err);
                            }
                        }
                    }
                    signal = signal_rx.recv() => {
                        tracing::info!(actor_id = %instance.self_addr(), "client received signal {signal:?}");
                        match signal {
                            Some(signal@(Signal::Stop(_) | Signal::DrainAndStop(_))) => {
                                need_drain = matches!(signal, Signal::DrainAndStop(_));
                                break None;
                            },
                            Some(Signal::ExitRequested(_)) => break None,
                            Some(Signal::ChildStopped(_)) => {},
                            Some(Signal::Kill(reason)) => {
                                break Some(ActorError { actor_id: Box::new(instance.self_addr().clone()), kind: Box::new(ActorErrorKind::Aborted(reason)) })
                            },
                            None => {
                                break Some(ActorError {
                                    actor_id: Box::new(instance.self_addr().clone()),
                                    kind: Box::new(ActorErrorKind::SignalChannelClosed),
                                })
                            },
                        }
                    }
                    Some(supervision_event) = supervision_rx.recv() => {
                        if let Err(err) = instance.handle_supervision_event(&mut actor, supervision_event).await {
                            break Some(err);
                        }
                    }
                };
            };
            if need_drain {
                let mut n = 0;
                while let Ok(work) = work_rx.try_recv() {
                    if let Err(e) = work.handle(&mut actor, instance).await {
                        err = Some(ActorError {
                            actor_id: Box::new(instance.self_addr().clone()),
                            kind: Box::new(ActorErrorKind::processing(e)),
                        });
                        break;
                    }
                    n += 1;
                }
                tracing::debug!(actor_id = %instance.self_addr(), "client drained {} messages before stopping", n);
            }
            if let Some(err) = err {
                let event = actor_error_to_event(instance, &actor, err);
                // The proc supervision handler will send to ProcAgent, which
                // just records it in v1. We want to crash instead, as nothing will
                // monitor the client ProcAgent for now.
                tracing::error!(
                    actor_id = %instance.self_addr(),
                    "could not propagate supervision event {} because it reached the global client: signaling KeyboardInterrupt to main thread",
                    event,
                );

                // This is running in a background thread, and thus cannot run
                // Py_FinalizeEx when it exits the process to properly shut down
                // all python objects.
                // We use _thread.interrupt_main to raise a KeyboardInterrupt
                // to the main thread at some point in the future.
                // There is no way to propagate the exception message, but it
                // will at least run proper shutdown code as long as BaseException
                // isn't caught.
                monarch_with_gil_blocking(GilSite::Stop, |py| {
                    // Use _thread.interrupt_main to force the client to exit if it has an
                    // unhandled supervision event.
                    let thread_mod = py.import("_thread").expect("import _thread");
                    let interrupt_main = thread_mod
                        .getattr("interrupt_main")
                        .expect("get interrupt_main");

                    // Ignore any exception from calling interrupt_main
                    if let Err(e) = interrupt_main.call0() {
                        tracing::error!("unable to interrupt main, exiting the process instead: {:?}", e);
                        eprintln!("unable to interrupt main, exiting the process with code 1 instead: {:?}", e);
                        std::process::exit(1);
                    }
                });
            } else {
                tracing::info!(actor_id = %instance.self_addr(), "client stopped");
                instance.change_status(hyperactor::actor::ActorStatus::Stopped("client stopped".into()));
            }
        });

        (root_client_instance.get().unwrap(), handle)
    }
}

fn actor_error_to_event(
    instance: &Instance<PythonActor>,
    actor: &PythonActor,
    err: ActorError,
) -> ActorSupervisionEvent {
    match *err.kind {
        ActorErrorKind::UnhandledSupervisionEvent(event) => *event,
        _ => {
            let status = ActorStatus::generic_failure(err.kind.to_string());
            ActorSupervisionEvent::new(
                instance.self_addr().clone(),
                actor.display_name(),
                status,
                None,
            )
        }
    }
}

pub(crate) fn root_client_actor(py: Python<'_>) -> &'static Instance<PythonActor> {
    static ROOT_CLIENT_ACTOR: OnceLock<&'static Instance<PythonActor>> = OnceLock::new();

    // Release the GIL before waiting on ROOT_CLIENT_ACTOR, because PythonActor::bootstrap_client
    // may release/reacquire the GIL; if thread 0 holds the GIL blocking on ROOT_CLIENT_ACTOR.get_or_init
    // while thread 1 blocks on acquiring the GIL inside PythonActor::bootstrap_client, we get
    // a deadlock.
    py.detach(|| {
        ROOT_CLIENT_ACTOR.get_or_init(|| {
            monarch_with_gil_blocking(GilSite::Bootstrap, |py| {
                let (client, _handle) = PythonActor::bootstrap_client(py);
                client
            })
        })
    })
}

#[async_trait]
impl Actor for PythonActor {
    async fn init(&mut self, this: &Instance<Self>) -> Result<(), anyhow::Error> {
        // PE-1: install the read side eagerly so the actor reports
        // `execution` from its first handled message. The callback runs on
        // the introspect task (off the actor loop) and only reads `Arc`
        // state (PE-3), so it is `Send + Sync`, non-blocking, and infallible.
        let tracker = self.execution_tracker.clone();
        this.set_attrs_snapshot(move || {
            let mut attrs = hyperactor_config::Attrs::new();
            attrs.set(EXECUTION, tracker.snapshot());
            attrs
        });

        if let PythonActorDispatchMode::Queue { receiver, .. } = &mut self.dispatch_mode {
            let receiver = receiver.take().unwrap();

            monarch_with_gil(GilSite::DispatchInit, |py| {
                let self_instance = self.ensure_py_instance(py, this);
                let actor_mesh_mod = py.import("monarch._src.actor.actor_mesh")?;

                let tl = &self.task_locals;
                let awaitable = actor_mesh_mod.call_method(
                    "_dispatch_loop",
                    (self.actor.clone_ref(py), receiver, self_instance),
                    None,
                )?;
                let future = pyo3_async_runtimes::into_future_with_locals(tl, awaitable)?;
                tokio::spawn(async move {
                    if let Err(e) = future.await {
                        tracing::error!("message loop error: {}", e);
                    }
                });
                Ok::<_, anyhow::Error>(())
            })
            .await?;
        }

        if let Some(init_message) = self.init_message.take() {
            let spawn_point = self.spawn_point.get().unwrap().as_ref().expect("PythonActor should never be spawned with init_message unless spawn_point also specified").clone();
            let mut headers = Flattrs::new();
            headers.set(CAST_POINT, spawn_point);
            let cx = Context::new(this, headers);
            <Self as Handler<PythonMessage>>::handle(self, &cx, init_message).await?;
        }

        Ok(())
    }

    async fn cleanup(
        &mut self,
        this: &Instance<Self>,
        err: Option<&ActorError>,
    ) -> anyhow::Result<()> {
        // Calls the "__cleanup__" method on the python instance to allow the actor
        // to control its own cleanup.
        // No headers because this isn't in the context of a message.
        let cx = Context::new(this, Flattrs::new());
        // Turn the ActorError into a representation of the error. We may not
        // have an original exception object or traceback, so we just pass in
        // the message.
        let err_as_str = err.map(|e| e.to_string());
        let future = monarch_with_gil(GilSite::EndpointCleanup, |py| {
            let py_cx = match &self.instance {
                Some(instance) => crate::context::PyContext::new(&cx, instance.clone_ref(py)),
                None => {
                    let py_instance: crate::context::PyInstance = this.into();
                    crate::context::PyContext::new(
                        &cx,
                        py_instance
                            .into_py_any(py)?
                            .downcast_bound(py)
                            .map_err(PyErr::from)?
                            .clone()
                            .unbind(),
                    )
                }
            }
            .into_bound_py_any(py)?;
            let actor = self.actor.bind(py);
            // Some tests don't use the Actor base class, so add this check
            // to be defensive.
            match actor.hasattr("__cleanup__") {
                Ok(false) | Err(_) => {
                    // No cleanup found, default to returning None
                    return Ok(None);
                }
                _ => {}
            }
            let awaitable = actor
                .call_method("__cleanup__", (&py_cx, err_as_str), None)
                .map_err(|err| anyhow::Error::from(SerializablePyErr::from(py, &err)))?;
            if awaitable.is_none() {
                Ok(None)
            } else {
                pyo3_async_runtimes::into_future_with_locals(&self.task_locals, awaitable)
                    .map(Some)
                    .map_err(anyhow::Error::from)
            }
        })
        .await;
        let cleanup_result = match future {
            Ok(Some(future)) => future.await.map(|_| ()).map_err(anyhow::Error::from),
            Ok(None) => Ok(()),
            Err(err) => Err(err),
        };
        let loop_shutdown_result = self.cancel_pending_python_tasks_and_stop_loop();
        cleanup_result?;
        loop_shutdown_result?;
        Ok(())
    }

    fn display_name(&self) -> Option<String> {
        self.instance.as_ref().and_then(|instance| {
            monarch_with_gil_blocking(GilSite::DisplayName, |py| {
                instance.bind(py).str().ok().map(|s| s.to_string())
            })
        })
    }

    async fn handle_undeliverable_message(
        &mut self,
        ins: &Instance<Self>,
        reason: UndeliverableReason,
        mut envelope: Undeliverable<MessageEnvelope>,
    ) -> Result<(), anyhow::Error> {
        if envelope
            .as_message()
            .is_some_and(|envelope| envelope.sender() != ins.self_addr())
        {
            // This can happen if the sender is comm. Update the envelope.
            envelope = update_undeliverable_envelope_for_casting(envelope);
        }
        let envelope = match envelope {
            Undeliverable::Returned(envelope) => envelope,
            Undeliverable::Report(report) => {
                return Err(UndeliverableMessageError::Report { report }.into());
            }
        };
        assert_eq!(
            envelope.sender(),
            ins.self_addr(),
            "undeliverable message was returned to the wrong actor. \
            Return address = {}, src actor = {}, dest handler port = {}, message type = {}, envelope headers = {}",
            envelope.sender(),
            ins.self_addr(),
            envelope.dest(),
            envelope.data().typename().unwrap_or("unknown"),
            envelope.headers()
        );

        let cx = Context::new(ins, envelope.headers().clone());

        let (envelope, handled) = monarch_with_gil(GilSite::EndpointDispatch, |py| {
            let py_cx = match &self.instance {
                Some(instance) => crate::context::PyContext::new(&cx, instance.clone_ref(py)),
                None => {
                    let py_instance: crate::context::PyInstance = ins.into();
                    crate::context::PyContext::new(
                        &cx,
                        py_instance
                            .into_py_any(py)?
                            .downcast_bound(py)
                            .map_err(PyErr::from)?
                            .clone()
                            .unbind(),
                    )
                }
            }
            .into_bound_py_any(py)?;
            let py_envelope = PythonUndeliverableMessageEnvelope {
                inner: Some(Undeliverable::Returned(envelope)),
            }
            .into_bound_py_any(py)?;
            let handled = self
                .actor
                .call_method(
                    py,
                    "_handle_undeliverable_message",
                    (&py_cx, &py_envelope),
                    None,
                )
                .map_err(|err| anyhow::Error::from(SerializablePyErr::from(py, &err)))?
                .extract::<bool>(py)?;
            Ok::<_, anyhow::Error>((
                py_envelope
                    .downcast::<PythonUndeliverableMessageEnvelope>()
                    .map_err(PyErr::from)?
                    .try_borrow_mut()
                    .map_err(PyErr::from)?
                    .take()?,
                handled,
            ))
        })
        .await?;

        if !handled {
            hyperactor::actor::handle_undeliverable_message(ins, reason, envelope)
        } else {
            Ok(())
        }
    }

    async fn handle_supervision_event(
        &mut self,
        this: &Instance<Self>,
        event: &ActorSupervisionEvent,
    ) -> Result<bool, anyhow::Error> {
        let cx = Context::new(this, Flattrs::new());
        self.handle(
            &cx,
            MeshFailure {
                // Populate the mesh name from the base-name string
                // plumbed through PythonActorParams at spawn time —
                // no lookup.
                actor_mesh_name: self.mesh_base_name.clone(),
                event: event.clone(),
                crashed_ranks: vec![],
            },
        )
        .await
        .map(|_| true)
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, Named)]
pub struct PythonActorParams {
    // The pickled actor class to instantiate.
    actor_type: PickledPyObject,
    // Python message to process as part of the actor initialization.
    init_message: Option<PythonMessage>,
    // User-provided mesh base-name string under which this actor
    // was spawned. The base name the caller passed when the mesh
    // was spawned, plumbed through `PythonActor` narrowly to
    // populate `MeshFailure.actor_mesh_name` on the direct
    // actor-handled supervision path without a lookup. It is not
    // actor display text (`display_name` handles that) and it is
    // not a general side channel; downstream code must not consume
    // this field for any other purpose. Kept separate from
    // `supervision_display_name`, which is a rendered supervision
    // display string passed through `spawn_with_name(...)`.
    mesh_base_name: Option<String>,
}

impl PythonActorParams {
    pub(crate) fn new(
        actor_type: PickledPyObject,
        init_message: Option<PythonMessage>,
        mesh_base_name: Option<String>,
    ) -> Self {
        Self {
            actor_type,
            init_message,
            mesh_base_name,
        }
    }
}

#[async_trait]
impl RemoteSpawn for PythonActor {
    type Params = PythonActorParams;

    async fn new(
        PythonActorParams {
            actor_type,
            init_message,
            mesh_base_name,
        }: PythonActorParams,
        environment: Flattrs,
    ) -> Result<Self, anyhow::Error> {
        let spawn_point = environment.get(CAST_POINT);
        Self::new(actor_type, init_message, spawn_point, mesh_base_name)
    }
}

/// Create a new TaskLocals with its own asyncio event loop in a dedicated thread.
fn create_task_locals() -> pyo3_async_runtimes::TaskLocals {
    monarch_with_gil_blocking(GilSite::TaskLocals, |py| {
        let asyncio = Python::import(py, "asyncio").unwrap();
        let event_loop = asyncio.call_method0("new_event_loop").unwrap();
        let task_locals = pyo3_async_runtimes::TaskLocals::new(event_loop.clone())
            .copy_context(py)
            .unwrap();

        let kwargs = PyDict::new(py);
        let target = event_loop.getattr("run_forever").unwrap();
        kwargs.set_item("target", target).unwrap();
        // Need to make this a daemon thread, otherwise shutdown will hang.
        kwargs.set_item("daemon", true).unwrap();
        let thread = py
            .import("threading")
            .unwrap()
            .call_method("Thread", (), Some(&kwargs))
            .unwrap();
        thread.call_method0("start").unwrap();
        task_locals
    })
}

// [Panics in async endpoints]
// This class exists to solve a deadlock when an async endpoint calls into some
// Rust code that panics.
//
// When an async endpoint is invoked and calls into Rust, the following sequence happens:
//
// hyperactor message -> PythonActor::handle() -> call _Actor.handle() in Python
//   -> convert the resulting coroutine into a Rust future, but scheduled on
//      the Python asyncio event loop (`into_future_with_locals`)
//   -> set a callback on Python asyncio loop to ping a channel that fulfills
//      the Rust future when the Python coroutine has finished. ('PyTaskCompleter`)
//
// This works fine for normal results and Python exceptions: we will take the
// result of the callback and send it through the channel, where it will be
// returned to the `await`er of the Rust future.
//
// This DOESN'T work for panics. The behavior of a panic in pyo3-bound code is
// that it will get caught by pyo3 and re-thrown to Python as a PanicException.
// And if that PanicException ever makes it back to Rust, it will get unwound
// instead of passed around as a normal PyErr type.
//
// So:
//   - Endpoint panics.
//   - This panic is captured as a PanicException in Python and
//     stored as the result of the Python asyncio task.
//   - When the callback in `PyTaskCompleter` queries the status of the task to
//     pass it back to the Rust awaiter, instead of getting a Result type, it
//     just starts resumes unwinding the PanicException
//   - This triggers a deadlock, because the whole task dies without ever
//     pinging the response channel, and the Rust awaiter will never complete.
//
// We work around this by passing a side-channel to our Python task so that it,
// in Python, can catch the PanicException and notify the Rust awaiter manually.
// In this way we can guarantee that the awaiter will complete even if the
// `PyTaskCompleter` callback explodes.
#[pyclass(module = "monarch._rust_bindings.monarch_hyperactor.actor")]
struct PanicFlag {
    sender: Option<tokio::sync::oneshot::Sender<Py<PyAny>>>,
}

#[pymethods]
impl PanicFlag {
    fn signal_panic(&mut self, ex: Py<PyAny>) {
        self.sender.take().unwrap().send(ex).unwrap();
    }
}

#[async_trait]
impl Handler<PythonMessage> for PythonActor {
    #[tracing::instrument(level = "debug", skip_all)]
    async fn handle(
        &mut self,
        cx: &Context<PythonActor>,
        message: PythonMessage,
    ) -> anyhow::Result<()> {
        match &self.dispatch_mode {
            PythonActorDispatchMode::Direct => self.handle_direct(cx, message).await,
            PythonActorDispatchMode::Queue { sender, .. } => {
                let sender = sender.clone();
                self.handle_queue(cx, sender, message).await
            }
        }
    }
}

impl PythonActor {
    /// Handle a message using direct dispatch (current behavior).
    async fn handle_direct(
        &mut self,
        cx: &Context<'_, PythonActor>,
        message: PythonMessage,
    ) -> anyhow::Result<()> {
        let resolved = message.resolve_indirect_call(cx).await?;
        let endpoint = resolved.method.to_string();

        // Create a channel for signaling panics in async endpoints.
        // See [Panics in async endpoints].
        let (sender, receiver) = oneshot::channel();

        let future = monarch_with_gil(
            GilSite::EndpointDispatch,
            |py| -> Result<_, SerializablePyErr> {
                let inst = self.ensure_py_instance(py, cx);

                let awaitable = self.actor.call_method(
                    py,
                    "handle",
                    (
                        crate::context::PyContext::new(cx, inst.clone_ref(py)),
                        resolved.method,
                        resolved.bytes,
                        PanicFlag {
                            sender: Some(sender),
                        },
                        resolved
                            .local_state
                            .unwrap_or_else(|| PyList::empty(py).unbind().into()),
                        resolved.mesh_references.into_py_any(py)?,
                        resolved.response_port.into_py_any(py)?,
                    ),
                    None,
                )?;

                pyo3_async_runtimes::into_future_with_locals(
                    &self.task_locals,
                    awaitable.into_bound(py),
                )
                .map_err(|err| err.into())
            },
        )
        .await?;

        // Spawn a child actor to await the Python handler method.
        tokio::spawn(handle_async_endpoint_panic(
            cx.signal_sender(),
            PythonTask::new(future)?,
            receiver,
            cx.self_addr().to_string(),
            endpoint,
        ));
        Ok(())
    }

    /// Handle a message using queue dispatch.
    /// Resolves the message on the Rust side and enqueues it for Python to process.
    async fn handle_queue(
        &mut self,
        cx: &Context<'_, PythonActor>,
        sender: pympsc::Sender,
        message: PythonMessage,
    ) -> anyhow::Result<()> {
        let resolved = message.resolve_indirect_call(cx).await?;

        let queued_msg = monarch_with_gil(
            GilSite::QueueDispatch,
            |py| -> anyhow::Result<QueuedMessage> {
                let inst = self.ensure_py_instance(py, cx);

                let py_context = crate::context::PyContext::new(cx, inst.clone_ref(py));
                let py_context_obj = Py::new(py, py_context)?;

                Ok(QueuedMessage {
                    context: py_context_obj,
                    method: resolved.method,
                    bytes: resolved.bytes,
                    local_state: resolved
                        .local_state
                        .unwrap_or_else(|| PyList::empty(py).unbind().into()),
                    refs: resolved.mesh_references.into_py_any(py)?,
                    response_port: resolved.response_port.into_py_any(py)?,
                })
            },
        )
        .await?;

        sender
            .send(queued_msg)
            .map_err(|_| anyhow::anyhow!("failed to send message to queue"))?;

        Ok(())
    }
}

#[async_trait]
impl Handler<MeshFailure> for PythonActor {
    async fn handle(&mut self, cx: &Context<Self>, message: MeshFailure) -> anyhow::Result<()> {
        // If the message is not about a failure, don't call __supervise__.
        // This includes messages like "stop", because those are not errors that
        // need to be propagated.
        if !message.event.actor_status.is_failed() {
            tracing::info!(
                "ignoring non-failure supervision event from child: {}",
                message
            );
            return Ok(());
        }
        // TODO: Consider routing supervision messages through the queue for Queue mode.
        // For now, supervision is always handled directly since it requires immediate response.

        // `_Actor.__supervise__` is `async def`, so calling it returns a
        // coroutine, which we schedule on the actor's asyncio event loop --
        // the same loop that runs endpoint coroutines. A sync user
        // `__supervise__` is dispatched under `fake_sync_state` inside
        // `_Actor.__supervise__`, mirroring `__cleanup__`.
        let (display_name, fut) = monarch_with_gil(GilSite::Supervise, |py| {
            let inst = self.ensure_py_instance(py, cx);
            // Compute display_name here since we can't call self.display_name() due to borrow.
            let display_name: Option<String> = inst.bind(py).str().ok().map(|s| s.to_string());
            let actor_bound = self.actor.bind(py);
            // The _Actor class always has a __supervise__ method, so this should
            // never happen.
            if !actor_bound.hasattr("__supervise__")? {
                return Err(anyhow::anyhow!(
                    "no __supervise__ method on {:?}",
                    actor_bound
                ));
            }
            let awaitable = actor_bound.call_method(
                "__supervise__",
                (
                    crate::context::PyContext::new(cx, inst.clone_ref(py)),
                    PyMeshFailure::from(message.clone()),
                ),
                None,
            )?;
            let fut = pyo3_async_runtimes::into_future_with_locals(&self.task_locals, awaitable)?;
            anyhow::Ok((display_name, fut))
        })
        .await?;

        let awaited = fut.await;

        monarch_with_gil(GilSite::Supervise, |py| match awaited {
            Ok(s) => {
                if s.bind(py).is_truthy()? {
                    // If the return value is truthy, then the exception was handled
                    // and doesn't need to be propagated.
                    // TODO: We also don't want to deliver multiple supervision
                    // events from the same mesh if an earlier one is handled.
                    tracing::info!(
                        name = "ActorMeshStatus",
                        status = "SupervisionError::Handled",
                        // only care about the event sender when the message is handled
                        actor_name = message.actor_mesh_name,
                        event = %message.event,
                        "__supervise__ on {} handled a supervision event, not reporting any further",
                        cx.self_addr(),
                    );
                    Ok(())
                } else {
                    // For a falsey return value, we propagate the supervision event
                    // to the next owning actor. We do this by returning a new
                    // error. This will not set the causal chain for ActorSupervisionEvent,
                    // so make sure to include the original event in the error message
                    // to provide context.

                    // False -- we propagate the event onward, but update it with the fact that
                    // this actor is now the event creator.
                    for (actor_name, status) in [
                        (
                            message
                                .actor_mesh_name
                                .as_deref()
                                .unwrap_or_else(|| message.event.actor_id.log_name()),
                            "SupervisionError::Unhandled",
                        ),
                        (cx.self_addr().log_name(), "UnhandledSupervisionEvent"),
                    ] {
                        tracing::info!(
                            name = "ActorMeshStatus",
                            status,
                            actor_name,
                            event = %message.event,
                            "__supervise__ on {} did not handle a supervision event, reporting to the next next owner",
                            cx.self_addr(),
                        );
                    }
                    let err = ActorErrorKind::UnhandledSupervisionEvent(Box::new(
                        ActorSupervisionEvent::new(
                            cx.self_addr().clone(),
                            display_name.clone(),
                            ActorStatus::Failed(ActorErrorKind::UnhandledSupervisionEvent(
                                Box::new(message.event.clone()),
                            )),
                            None,
                        ),
                    ));
                    Err(anyhow::Error::new(err))
                }
            }
            Err(err) => {
                // If __supervise__ raised UnhandledFaultHookException,
                // return the PyErr directly without wrapping in
                // ActorErrorKind. The custom run loop detects this by
                // downcasting the anyhow::Error to PyErr.
                if err.is_instance(py, &unhandled_fault_hook_exception(py)) {
                    return Err(err.into());
                }

                // Any other exception will supersede in the propagation chain,
                // and will become its own supervision failure.
                // Include the event it was handling in the error message.

                // Add to caused_by chain.
                for (actor_name, status) in [
                    (
                        message
                            .actor_mesh_name
                            .as_deref()
                            .unwrap_or_else(|| message.event.actor_id.log_name()),
                        "SupervisionError::__supervise__::exception",
                    ),
                    (cx.self_addr().log_name(), "UnhandledSupervisionEvent"),
                ] {
                    tracing::info!(
                        name = "ActorMeshStatus",
                        status,
                        actor_name,
                        event = %message.event,
                        "__supervise__ on {} threw an exception",
                        cx.self_addr(),
                    );
                }
                let err = ActorErrorKind::UnhandledSupervisionEvent(Box::new(
                    ActorSupervisionEvent::new(
                        cx.self_addr().clone(),
                        display_name,
                        ActorStatus::Failed(ActorErrorKind::ErrorDuringHandlingSupervision(
                            err.to_string(),
                            Box::new(message.event.clone()),
                        )),
                        None,
                    ),
                ));
                Err(anyhow::Error::new(err))
            }
        })
        .await
    }
}

async fn handle_async_endpoint_panic(
    panic_sender: mpsc::UnboundedSender<Signal>,
    task: PythonTask,
    side_channel: oneshot::Receiver<Py<PyAny>>,
    actor_id: String,
    endpoint: String,
) {
    // Create attributes for metrics with actor_id and endpoint
    let attributes =
        hyperactor_telemetry::kv_pairs!("actor_id" => actor_id, "endpoint" => endpoint);

    // Record the start time for latency measurement
    let start_time = std::time::Instant::now();

    // Increment throughput counter
    ENDPOINT_ACTOR_COUNT.add(1, attributes);

    let err_or_never = async {
        // The side channel will resolve with a value if a panic occured during
        // processing of the async endpoint, see [Panics in async endpoints].
        match side_channel.await {
            Ok(value) => {
                monarch_with_gil(GilSite::AwaitDrive, |py| -> Option<SerializablePyErr> {
                    let err: PyErr = value
                        .downcast_bound::<PyBaseException>(py)
                        .unwrap()
                        .clone()
                        .into();
                    ENDPOINT_ACTOR_PANIC.add(1, attributes);
                    Some(err.into())
                })
                .await
            }
            // An Err means that the sender has been dropped without sending.
            // That's okay, it just means that the Python task has completed.
            // In that case, just never resolve this future. We expect the other
            // branch of the select to finish eventually.
            Err(_) => pending().await,
        }
    };
    let future = task.take();
    if let Some(panic) = tokio::select! {
        result = future => {
            match result {
                Ok(_) => None,
                Err(e) => Some(e.into()),
            }
        },
        result = err_or_never => {
            result
        }
    } {
        // Record error and panic metrics
        ENDPOINT_ACTOR_ERROR.add(1, attributes);
        if panic_sender.send(Signal::Kill(panic.to_string())).is_err() {
            tracing::warn!("dropped panic signal: actor already stopped: {panic}");
        }
    }

    // Record latency in microseconds
    let elapsed_micros = start_time.elapsed().as_micros() as f64;
    ENDPOINT_ACTOR_LATENCY_US_HISTOGRAM.record(elapsed_micros, attributes);
}

#[pyclass(module = "monarch._rust_bindings.monarch_hyperactor.actor")]
struct LocalPort {
    instance: PyInstance,
    inner: Option<OncePortHandle<Result<Py<PyAny>, Py<PyAny>>>>,
}

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

pub(crate) fn to_py_error<T>(e: T) -> PyErr
where
    T: Error,
{
    PyErr::new::<PyValueError, _>(e.to_string())
}

#[pymethods]
impl LocalPort {
    fn send(&mut self, obj: Py<PyAny>) -> PyResult<()> {
        let port = self.inner.take().expect("use local port once");
        port.post(self.instance.deref(), Ok(obj));
        Ok(())
    }
    fn resolve_and_send(&mut self, obj: Py<PyAny>) -> PyResult<PyPythonTask> {
        self.send(obj)?;
        PyPythonTask::new(async { Ok(()) })
    }
    fn exception(&mut self, e: Py<PyAny>) -> PyResult<()> {
        let port = self.inner.take().expect("use local port once");
        port.post(self.instance.deref(), Err(e));
        Ok(())
    }
}

/// A port that drops all messages sent to it.
/// Used when there is no response port for a message.
/// Any exceptions sent to it are re-raised in the current actor.
#[pyclass(module = "monarch._rust_bindings.monarch_hyperactor.actor")]
#[derive(Debug)]
pub struct DroppingPort;

#[pymethods]
impl DroppingPort {
    #[new]
    fn new() -> Self {
        DroppingPort
    }

    fn send(&self, _obj: Py<PyAny>) -> PyResult<()> {
        Ok(())
    }

    fn resolve_and_send(&self, obj: Py<PyAny>) -> PyResult<PyPythonTask> {
        self.send(obj)?;
        PyPythonTask::new(async { Ok(()) })
    }

    fn send_message(&self, _message: PythonMessage) -> PyResult<()> {
        Ok(())
    }

    fn exception(&self, e: Bound<'_, PyAny>) -> PyResult<()> {
        // Unwrap ActorError to get the inner exception, matching Python behavior.
        let exc = if let Ok(inner) = e.getattr("exception") {
            inner
        } else {
            e
        };
        Err(PyErr::from_value(exc))
    }

    #[getter]
    fn get_return_undeliverable(&self) -> bool {
        true
    }

    #[setter]
    fn set_return_undeliverable(&self, _value: bool) {}
}

/// A port that sends messages to a remote receiver.
/// Wraps an EitherPortRef with the actor instance needed for sending.
#[pyclass(module = "monarch._src.actor.actor_mesh")]
pub struct Port {
    port_ref: EitherPortRef,
    instance: Instance<PythonActor>,
    rank: Option<usize>,
    /// Operation-context headers captured from the inbound request,
    /// re-emitted on every reply so failure surfaces can name the
    /// operation.
    reply_headers: hyperactor_config::Flattrs,
}

#[pymethods]
impl Port {
    #[new]
    fn new(
        port_ref: EitherPortRef,
        instance: &crate::context::PyInstance,
        rank: Option<usize>,
    ) -> Self {
        Self {
            port_ref,
            instance: instance.clone().into_instance(),
            rank,
            reply_headers: hyperactor_config::Flattrs::new(),
        }
    }

    #[getter("_port_ref")]
    fn port_ref_py(&self) -> EitherPortRef {
        self.port_ref.clone()
    }

    #[getter("_rank")]
    fn rank_py(&self) -> Option<usize> {
        self.rank
    }

    #[getter]
    fn get_return_undeliverable(&self) -> bool {
        self.port_ref.get_return_undeliverable()
    }

    #[setter]
    fn set_return_undeliverable(&mut self, value: bool) {
        self.port_ref.set_return_undeliverable(value);
    }

    #[tracing::instrument(level = "debug", skip_all)]
    fn send(&mut self, py: Python<'_>, obj: Py<PyAny>) -> PyResult<()> {
        let message = PythonMessage::new_from_buf(
            PythonMessageKind::Result { rank: self.rank },
            pickle_to_part(py, &obj)?,
        );

        self.port_ref
            .post_with_headers(&self.instance, self.reply_headers.clone(), message)
            .map_err(|e| PyRuntimeError::new_err(e.to_string()))
    }

    #[tracing::instrument(level = "debug", skip_all)]
    fn send_message(&mut self, message: PythonMessage) -> PyResult<()> {
        self.port_ref
            .post_with_headers(&self.instance, self.reply_headers.clone(), message)
            .map_err(|e| PyRuntimeError::new_err(e.to_string()))
    }

    fn exception(&mut self, py: Python<'_>, e: Py<PyAny>) -> PyResult<()> {
        let message = PythonMessage::new_from_buf(
            PythonMessageKind::Exception { rank: self.rank },
            pickle_to_part(py, &e)?,
        );

        self.port_ref
            .post_with_headers(&self.instance, self.reply_headers.clone(), message)
            .map_err(|e| PyRuntimeError::new_err(e.to_string()))
    }
}

impl Port {
    /// Constructor that attaches operation-context headers captured
    /// from the inbound request. The Python `#[new]` constructor
    /// defaults to empty headers.
    pub(crate) fn with_reply_headers(
        port_ref: EitherPortRef,
        instance: Instance<PythonActor>,
        rank: Option<usize>,
        reply_headers: hyperactor_config::Flattrs,
    ) -> Self {
        Self {
            port_ref,
            instance,
            rank,
            reply_headers,
        }
    }
}

pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
    hyperactor_mod.add_class::<PythonActorHandle>()?;
    hyperactor_mod.add_class::<PythonMessage>()?;
    hyperactor_mod.add_class::<PyMeshRef>()?;
    hyperactor_mod.add_class::<PythonMessageKind>()?;
    hyperactor_mod.add_class::<MethodSpecifier>()?;
    hyperactor_mod.add_class::<UnflattenArg>()?;
    hyperactor_mod.add_class::<PanicFlag>()?;
    hyperactor_mod.add_class::<QueuedMessage>()?;
    hyperactor_mod.add_class::<DroppingPort>()?;
    hyperactor_mod.add_class::<Port>()?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use hyperactor as reference;
    use hyperactor::accum::ReducerSpec;
    use hyperactor::accum::StreamingReducerOpts;
    use hyperactor::id::Label;
    use hyperactor::testing::ids::test_port_id;
    use hyperactor_mesh::Error as MeshError;
    use hyperactor_mesh::host_mesh::host_agent::ProcState;
    use hyperactor_mesh::mesh_id::ResourceId;
    use hyperactor_mesh::resource::Status;
    use hyperactor_mesh::resource::{self};
    use pyo3::PyTypeInfo;

    use super::*;
    use crate::actor::to_py_error;

    #[test]
    fn test_python_message_part_codec() {
        let reducer_spec = ReducerSpec {
            typehash: 123,
            builder_params: Some(wirevalue::Any::serialize(&"abcdefg12345".to_string()).unwrap()),
        };
        let port_ref = hyperactor::PortRef::<PythonMessage>::attest_reducible(
            test_port_id("world_0", "client", 123),
            Some(reducer_spec),
            StreamingReducerOpts::default(),
        );
        let message = PythonMessage {
            kind: PythonMessageKind::CallMethod {
                name: MethodSpecifier::ReturnsResponse {
                    name: "test".to_string(),
                },
                response_port: Some(EitherPortRef::Unbounded(port_ref.clone().into())),
            },
            message: Part::from(vec![1, 2, 3]),
            refs: Vec::new(),
        };
        {
            let mut multipart_message =
                wirevalue::Any::<wirevalue::encoding::Multipart>::serialize(&message).unwrap();
            let mut ports = vec![];
            multipart_message
                .visit_multipart_parts_mut::<reference::PortRefRepr, anyhow::Error>(|b| {
                    ports.push(b.clone());
                    Ok(())
                })
                .unwrap();
            assert_eq!(ports.len(), 1);
            assert_eq!(ports[0].port_addr(), port_ref.port_addr());
            assert_eq!(ports[0].reducer_spec(), port_ref.reducer_spec());
            assert_eq!(
                ports[0].get_return_undeliverable(),
                port_ref.get_return_undeliverable()
            );
            assert!(!ports[0].unsplit());
            assert_eq!(
                message,
                multipart_message
                    .deserialized_unchecked::<PythonMessage>()
                    .unwrap()
            );
        }

        let no_port_message = PythonMessage {
            kind: PythonMessageKind::CallMethod {
                name: MethodSpecifier::ReturnsResponse {
                    name: "test".to_string(),
                },
                response_port: None,
            },
            ..message
        };
        {
            let mut multipart_message =
                wirevalue::Any::<wirevalue::encoding::Multipart>::serialize(&no_port_message)
                    .unwrap();
            let mut ports = vec![];
            multipart_message
                .visit_multipart_parts_mut::<reference::PortRefRepr, anyhow::Error>(|b| {
                    ports.push(b.clone());
                    Ok(())
                })
                .unwrap();
            assert_eq!(ports.len(), 0);
            assert_eq!(
                no_port_message,
                multipart_message
                    .deserialized_unchecked::<PythonMessage>()
                    .unwrap()
            );
        }
    }

    #[test]
    fn test_python_message_refs_travel_as_parts() {
        // A non-live proc mesh ref, built in-memory from ids (no spawn).
        fn proc_mesh_ref(seed: u64, label: &str) -> MeshRef {
            let proc_id = hyperactor::ProcId::new(
                hyperactor::id::Uid::Instance(seed, None),
                Some(Label::new("local").unwrap()),
            );
            let proc_addr = hyperactor::ProcAddr::new(
                proc_id,
                hyperactor::channel::ChannelAddr::Local(seed).into(),
            );
            let agent: hyperactor::ActorRef<hyperactor_mesh::proc_agent::ProcAgent> =
                hyperactor::ActorRef::attest(
                    proc_addr.actor_addr(hyperactor_mesh::proc_agent::PROC_AGENT_ACTOR_NAME),
                );
            let proc_ref = hyperactor_mesh::proc_mesh::ProcRef::new(proc_addr, 0, agent);
            MeshRef::Proc(Box::new(
                hyperactor_mesh::proc_mesh::ProcMeshRef::new_singleton(
                    hyperactor_mesh::mesh_id::ProcMeshId::singleton(Label::new(label).unwrap()),
                    proc_ref,
                )
                .unwrap(),
            ))
        }

        let message = PythonMessage {
            kind: PythonMessageKind::CallMethod {
                name: MethodSpecifier::ReturnsResponse {
                    name: "test".to_string(),
                },
                response_port: None,
            },
            message: Part::from(vec![1, 2, 3]),
            refs: vec![proc_mesh_ref(1, "a"), proc_mesh_ref(2, "b")],
        };

        let mut multipart_message =
            wirevalue::Any::<wirevalue::encoding::Multipart>::serialize(&message).unwrap();
        let mut parts = vec![];
        multipart_message
            .visit_multipart_parts_mut::<MeshRefRepr, anyhow::Error>(|b| {
                parts.push(b.clone());
                Ok(())
            })
            .unwrap();
        // Each MeshRef rides as its own typed part on the multipart wire.
        assert_eq!(parts.len(), 2);
        // And the message round-trips, reuniting the refs from those parts.
        assert_eq!(
            message,
            multipart_message
                .deserialized_unchecked::<PythonMessage>()
                .unwrap()
        );
    }

    #[test]
    fn to_py_error_preserves_proc_creation_message() {
        // State<ProcState> w/ `state.is_none()`
        let state: resource::State<ProcState> = resource::State {
            id: ResourceId::instance(Label::new("my-proc").unwrap()),
            status: Status::Failed("boom".into()),
            state: None,
            generation: 0,
            timestamp: std::time::SystemTime::now(),
        };

        // A ProcCreationError
        let mesh_agent: hyperactor::ActorRef<hyperactor_mesh::host_mesh::HostAgent> =
            hyperactor::ActorRef::attest(test_port_id("hello_0", "actor", 0).actor_addr());
        let expected_prefix = format!(
            "error creating proc (host rank 0) on host mesh agent {}",
            mesh_agent
        );
        let err = MeshError::ProcCreationError {
            host_rank: 0,
            mesh_agent,
            state: Box::new(state),
        };

        let rust_msg = err.to_string();
        let pyerr = to_py_error(err);

        pyo3::Python::initialize();
        monarch_with_gil_blocking(GilSite::Test, |py| {
            assert!(pyerr.get_type(py).is(PyValueError::type_object(py)));
            let py_msg = pyerr.value(py).to_string();

            // 1) Bridge preserves the exact message
            assert_eq!(py_msg, rust_msg);
            // 2) Contains the structured state and failure status
            assert!(py_msg.contains(", state: "));
            assert!(py_msg.contains("\"status\":{\"Failed\":\"boom\"}"));
            // 3) Starts with the expected prefix
            assert!(py_msg.starts_with(&expected_prefix));
        });
    }
}