hyperactor_mesh/

proc_agent.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.
 */

//! The mesh agent actor manages procs in ProcMeshes.

// EnumAsInner generates code that triggers a false positive
// unused_assignments lint on struct variant fields. #[allow] on the
// enum itself doesn't propagate into derive-macro-generated code, so
// the suppression must be at module scope.
#![allow(unused_assignments)]

use std::collections::HashMap;
use std::mem::take;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::RwLock;
use std::sync::RwLockReadGuard;
use std::time::Duration;

use async_trait::async_trait;
use enum_as_inner::EnumAsInner;
use hyperactor::Actor;
use hyperactor::ActorHandle;
use hyperactor::Bind;
use hyperactor::Context;
use hyperactor::Data;
use hyperactor::HandleClient;
use hyperactor::Handler;
use hyperactor::Instance;
use hyperactor::PortHandle;
use hyperactor::RefClient;
use hyperactor::Unbind;
use hyperactor::actor::ActorStatus;
use hyperactor::actor::remote::Remote;
use hyperactor::channel;
use hyperactor::channel::ChannelAddr;
use hyperactor::mailbox::BoxedMailboxSender;
use hyperactor::mailbox::DialMailboxRouter;
use hyperactor::mailbox::IntoBoxedMailboxSender;
use hyperactor::mailbox::MailboxClient;
use hyperactor::mailbox::MailboxSender;
use hyperactor::mailbox::MessageEnvelope;
use hyperactor::mailbox::Undeliverable;
use hyperactor::proc::Proc;
use hyperactor::reference as hyperactor_reference;
use hyperactor::supervision::ActorSupervisionEvent;
use hyperactor_config::CONFIG;
use hyperactor_config::ConfigAttr;
use hyperactor_config::attrs::declare_attrs;
use serde::Deserialize;
use serde::Serialize;
use typeuri::Named;

use crate::Name;
use crate::resource;

/// Actor name used when spawning the proc agent on user procs.
pub const PROC_AGENT_ACTOR_NAME: &str = "proc_agent";

declare_attrs! {
    /// Whether to self kill actors, procs, and hosts whose owner is not reachable.
    @meta(CONFIG = ConfigAttr::new(
        Some("HYPERACTOR_MESH_ORPHAN_TIMEOUT".to_string()),
        Some("mesh_orphan_timeout".to_string()),
    ))
    pub attr MESH_ORPHAN_TIMEOUT: Duration = Duration::from_secs(60);
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Named)]
pub enum GspawnResult {
    Success {
        rank: usize,
        actor_id: hyperactor_reference::ActorId,
    },
    Error(String),
}
wirevalue::register_type!(GspawnResult);

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Named)]
pub enum StopActorResult {
    Success,
    Timeout,
    NotFound,
}
wirevalue::register_type!(StopActorResult);

/// Deferred republish of introspect properties.
///
/// Sent as a zero-delay self-message from the supervision event
/// handler so it returns immediately without blocking the ProcAgent
/// message loop. Multiple rapid supervision events (e.g., 4 actors
/// failing simultaneously via broadcast) coalesce into a single
/// republish via the `introspect_dirty` flag.
///
/// Without this, calling `publish_introspect_properties` inline in
/// the supervision handler starves `GetRankStatus` polls from the
/// `ActorMeshController`, preventing `__supervise__` from firing
/// within the test timeout. See D94960791 for the root cause
/// analysis.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Named, Bind, Unbind)]
struct RepublishIntrospect;
wirevalue::register_type!(RepublishIntrospect);

/// Collect live actor children and system actor children from the
/// proc's instance DashMap using `all_instance_keys()` with point
/// lookups. This avoids the convoy starvation from `all_actor_ids()`
/// which holds shard read locks while doing heavy per-entry work.
/// See S12 in `introspect` module doc.
fn collect_live_children(proc: &hyperactor::Proc) -> (Vec<String>, Vec<String>) {
    let all_keys = proc.all_instance_keys();
    let mut children = Vec::with_capacity(all_keys.len());
    let mut system_children = Vec::new();
    for id in all_keys {
        if let Some(cell) = proc.get_instance(&id) {
            let ref_str = id.to_string();
            if cell.is_system() {
                system_children.push(ref_str.clone());
            }
            children.push(ref_str);
        }
    }
    (children, system_children)
}

#[derive(
    Debug,
    Clone,
    PartialEq,
    Serialize,
    Deserialize,
    Handler,
    HandleClient,
    RefClient,
    Named
)]
pub(crate) enum MeshAgentMessage {
    /// Configure the proc in the mesh.
    Configure {
        /// The rank of this proc in the mesh.
        rank: usize,
        /// The forwarder to send messages to unknown destinations.
        forwarder: ChannelAddr,
        /// The supervisor port to which the agent should report supervision events.
        supervisor: Option<hyperactor_reference::PortRef<ActorSupervisionEvent>>,
        /// An address book to use for direct dialing.
        address_book: HashMap<hyperactor_reference::ProcId, ChannelAddr>,
        /// The agent should write its rank to this port when it successfully
        /// configured.
        configured: hyperactor_reference::PortRef<usize>,
        /// If true, and supervisor is None, record supervision events to be reported
        record_supervision_events: bool,
    },

    Status {
        /// The status of the proc.
        /// To be replaced with fine-grained lifecycle status,
        /// and to use aggregation.
        status: hyperactor_reference::PortRef<(usize, bool)>,
    },

    /// Spawn an actor on the proc to the provided name.
    Gspawn {
        /// registered actor type
        actor_type: String,
        /// spawned actor name
        actor_name: String,
        /// serialized parameters
        params_data: Data,
        /// reply port; the proc should send its rank to indicated a spawned actor
        status_port: hyperactor_reference::PortRef<GspawnResult>,
    },

    /// Stop actors of a specific mesh name
    StopActor {
        /// The actor to stop
        actor_id: hyperactor_reference::ActorId,
        /// The timeout for waiting for the actor to stop
        timeout_ms: u64,
        /// The reason for stopping the actor
        reason: String,
        /// The result when trying to stop the actor
        #[reply]
        stopped: hyperactor_reference::OncePortRef<StopActorResult>,
    },
}

/// Internal configuration state of the mesh agent.
#[derive(Debug, EnumAsInner, Default)]
enum State {
    UnconfiguredV0 {
        sender: ReconfigurableMailboxSender,
    },

    ConfiguredV0 {
        sender: ReconfigurableMailboxSender,
        rank: usize,
        supervisor: Option<hyperactor_reference::PortRef<ActorSupervisionEvent>>,
    },

    V1,

    #[default]
    Invalid,
}

impl State {
    fn rank(&self) -> Option<usize> {
        match self {
            State::ConfiguredV0 { rank, .. } => Some(*rank),
            _ => None,
        }
    }

    fn supervisor(&self) -> Option<hyperactor_reference::PortRef<ActorSupervisionEvent>> {
        match self {
            State::ConfiguredV0 { supervisor, .. } => supervisor.clone(),
            _ => None,
        }
    }
}

/// Actor state used for v1 API.
#[derive(Debug)]
struct ActorInstanceState {
    create_rank: usize,
    spawn: Result<hyperactor_reference::ActorId, anyhow::Error>,
    /// If true, the actor has been stopped. There is no way to restart it, a new
    /// actor must be spawned.
    stopped: bool,
    /// The time at which the actor should be considered expired if no further
    /// keepalive is received. `None` meaning it will never expire.
    expiry_time: Option<std::time::SystemTime>,
}

#[derive(
    Clone,
    Debug,
    Default,
    PartialEq,
    Serialize,
    Deserialize,
    Named,
    Bind,
    Unbind
)]
struct SelfCheck {}

/// A mesh agent is responsible for managing procs in a [`ProcMesh`].
///
/// ## Supervision event ingestion (remote)
///
/// `ProcAgent` is the *process/rank-local* sink for
/// `ActorSupervisionEvent`s produced by the runtime (actor failures,
/// routing failures, undeliverables, etc.).
///
/// We **export** `ActorSupervisionEvent` as a handler so that other
/// procs—most importantly the process-global root client created by
/// `context()`—can forward undeliverables as supervision
/// events to the *currently active* mesh.
///
/// Without exporting this handler, `ActorSupervisionEvent` cannot be
/// addressed via `ActorRef`/`PortRef` across processes, and the
/// global-root-client undeliverable → supervision pipeline would
/// degrade to log-only behavior (events become undeliverable again or
/// are dropped).
///
/// See GC-1 in `global_context` module doc.
#[hyperactor::export(
    handlers=[
        MeshAgentMessage,
        ActorSupervisionEvent,
        resource::CreateOrUpdate<ActorSpec> { cast = true },
        resource::Stop { cast = true },
        resource::StopAll { cast = true },
        resource::GetState<ActorState> { cast = true },
        resource::KeepaliveGetState<ActorState> { cast = true },
        resource::GetRankStatus { cast = true },
        RepublishIntrospect { cast = true },
    ]
)]
pub struct ProcAgent {
    proc: Proc,
    remote: Remote,
    state: State,
    /// Actors created and tracked through the resource behavior.
    actor_states: HashMap<Name, ActorInstanceState>,
    /// If true, and supervisor is None, record supervision events to be reported
    /// to owning actors later.
    record_supervision_events: bool,
    /// If record_supervision_events is true, then this will contain the list
    /// of all events that were received.
    supervision_events: HashMap<hyperactor_reference::ActorId, Vec<ActorSupervisionEvent>>,
    /// True when supervision events have arrived but introspect
    /// properties haven't been republished yet.
    introspect_dirty: bool,
    /// If set, the StopAll handler will send the exit code through this
    /// channel instead of calling process::exit directly, allowing the
    /// caller to perform graceful shutdown (e.g. draining the mailbox server).
    shutdown_tx: Option<tokio::sync::oneshot::Sender<i32>>,
    /// If set, check for expired actors whose keepalive has lapsed.
    mesh_orphan_timeout: Option<Duration>,
}

impl ProcAgent {
    #[hyperactor::observe_result("MeshAgent")]
    pub(crate) async fn bootstrap(
        proc_id: hyperactor_reference::ProcId,
    ) -> Result<(Proc, ActorHandle<Self>), anyhow::Error> {
        let sender = ReconfigurableMailboxSender::new();
        let proc = Proc::configured(proc_id.clone(), BoxedMailboxSender::new(sender.clone()));

        let agent = ProcAgent {
            proc: proc.clone(),
            remote: Remote::collect(),
            state: State::UnconfiguredV0 { sender },
            actor_states: HashMap::new(),
            record_supervision_events: false,
            supervision_events: HashMap::new(),
            introspect_dirty: false,
            shutdown_tx: None,
            // v0 procs don't have an owner they can check for, so they should
            // never try to kill the children.
            mesh_orphan_timeout: None,
        };
        let handle = proc.spawn::<Self>("mesh", agent)?;
        Ok((proc, handle))
    }

    pub(crate) fn boot_v1(
        proc: Proc,
        shutdown_tx: Option<tokio::sync::oneshot::Sender<i32>>,
    ) -> Result<ActorHandle<Self>, anyhow::Error> {
        // We can't use Option<Duration> directly in config attrs because AttrValue
        // is not implemented for Option<Duration>. So we use a zero timeout to
        // indicate no timeout.
        let orphan_timeout = hyperactor_config::global::get(MESH_ORPHAN_TIMEOUT);
        let orphan_timeout = if orphan_timeout.is_zero() {
            None
        } else {
            Some(orphan_timeout)
        };
        let agent = ProcAgent {
            proc: proc.clone(),
            remote: Remote::collect(),
            state: State::V1,
            actor_states: HashMap::new(),
            record_supervision_events: true,
            supervision_events: HashMap::new(),
            introspect_dirty: false,
            shutdown_tx,
            mesh_orphan_timeout: orphan_timeout,
        };
        proc.spawn::<Self>(PROC_AGENT_ACTOR_NAME, agent)
    }

    async fn destroy_and_wait_except_current<'a>(
        &mut self,
        cx: &Context<'a, Self>,
        timeout: tokio::time::Duration,
        reason: &str,
    ) -> Result<
        (
            Vec<hyperactor_reference::ActorId>,
            Vec<hyperactor_reference::ActorId>,
        ),
        anyhow::Error,
    > {
        self.proc
            .destroy_and_wait_except_current::<Self>(timeout, Some(cx), true, reason)
            .await
    }

    /// Publish the current proc properties and children list for
    /// introspection. See S12 in `introspect` module doc.
    fn publish_introspect_properties(&self, cx: &impl hyperactor::context::Actor) {
        let (mut children, mut system_children) = collect_live_children(&self.proc);

        // Terminated actors appear as children but don't inflate
        // the actor count. Track them in stopped_children so the
        // TUI can filter/gray without per-child fetches.
        let mut stopped_children: Vec<String> = Vec::new();
        for id in self.proc.all_terminated_actor_ids() {
            let ref_str = id.to_string();
            stopped_children.push(ref_str.clone());
            // Terminated system actors must also appear in
            // system_children for correct filtering.
            if let Some(snapshot) = self.proc.terminated_snapshot(&id) {
                let snapshot_attrs: hyperactor_config::Attrs =
                    serde_json::from_str(&snapshot.attrs).unwrap_or_default();
                if snapshot_attrs
                    .get(hyperactor::introspect::IS_SYSTEM)
                    .copied()
                    .unwrap_or(false)
                {
                    system_children.push(ref_str.clone());
                }
            }
            if !children.contains(&ref_str) {
                children.push(ref_str);
            }
        }

        let stopped_retention_cap =
            hyperactor_config::global::get(hyperactor::config::TERMINATED_SNAPSHOT_RETENTION);

        // FI-5: is_poisoned iff failed_actor_count > 0.
        let failed_actor_count = self.supervision_events.len();

        // Attrs-based introspection.
        let num_live = children.len();
        let mut attrs = hyperactor_config::Attrs::new();
        attrs.set(crate::introspect::NODE_TYPE, "proc".to_string());
        attrs.set(
            crate::introspect::PROC_NAME,
            self.proc.proc_id().to_string(),
        );
        attrs.set(crate::introspect::NUM_ACTORS, num_live);
        attrs.set(hyperactor::introspect::CHILDREN, children);
        attrs.set(crate::introspect::SYSTEM_CHILDREN, system_children);
        attrs.set(crate::introspect::STOPPED_CHILDREN, stopped_children);
        attrs.set(
            crate::introspect::STOPPED_RETENTION_CAP,
            stopped_retention_cap,
        );
        attrs.set(crate::introspect::IS_POISONED, failed_actor_count > 0);
        attrs.set(crate::introspect::FAILED_ACTOR_COUNT, failed_actor_count);
        cx.instance().publish_attrs(attrs);
    }
}

#[async_trait]
impl Actor for ProcAgent {
    async fn init(&mut self, this: &Instance<Self>) -> Result<(), anyhow::Error> {
        this.set_system();
        self.proc.set_supervision_coordinator(this.port())?;
        self.publish_introspect_properties(this);

        // Resolve terminated actor snapshots via QueryChild so that
        // dead actors remain directly queryable by reference.
        let proc = self.proc.clone();
        let self_id = this.self_id().clone();
        this.set_query_child_handler(move |child_ref| {
            use hyperactor::introspect::IntrospectResult;

            if let hyperactor::reference::Reference::Actor(id) = child_ref {
                if let Some(snapshot) = proc.terminated_snapshot(id) {
                    return snapshot;
                }
            }

            // PA-1 (ProcAgent path): proc-node children used by
            // admin/TUI must be computed from live proc state at query
            // time, not solely from cached published_properties.
            // Therefore a direct proc.spawn() actor must appear on the
            // next QueryChild(Reference::Proc) response without an
            // extra publish event. See
            // test_query_child_proc_returns_live_children.
            if let hyperactor::reference::Reference::Proc(proc_id) = child_ref {
                if proc_id == proc.proc_id() {
                    let (mut children, mut system_children) = collect_live_children(&proc);

                    let mut stopped_children: Vec<String> = Vec::new();
                    for id in proc.all_terminated_actor_ids() {
                        let ref_str = id.to_string();
                        stopped_children.push(ref_str.clone());
                        if let Some(snapshot) = proc.terminated_snapshot(&id) {
                            let snapshot_attrs: hyperactor_config::Attrs =
                                serde_json::from_str(&snapshot.attrs).unwrap_or_default();
                            if snapshot_attrs
                                .get(hyperactor::introspect::IS_SYSTEM)
                                .copied()
                                .unwrap_or(false)
                            {
                                system_children.push(ref_str.clone());
                            }
                        }
                        if !children.contains(&ref_str) {
                            children.push(ref_str);
                        }
                    }

                    let stopped_retention_cap = hyperactor_config::global::get(
                        hyperactor::config::TERMINATED_SNAPSHOT_RETENTION,
                    );

                    let (is_poisoned, failed_actor_count) = proc
                        .get_instance(&self_id)
                        .and_then(|cell| cell.published_attrs())
                        .map(|attrs| {
                            let is_poisoned = attrs
                                .get(crate::introspect::IS_POISONED)
                                .copied()
                                .unwrap_or(false);
                            let failed_actor_count = attrs
                                .get(crate::introspect::FAILED_ACTOR_COUNT)
                                .copied()
                                .unwrap_or(0);
                            (is_poisoned, failed_actor_count)
                        })
                        .unwrap_or((false, 0));

                    // Build attrs for this proc node.
                    let num_live = children.len();
                    let mut attrs = hyperactor_config::Attrs::new();
                    attrs.set(crate::introspect::NODE_TYPE, "proc".to_string());
                    attrs.set(crate::introspect::PROC_NAME, proc_id.to_string());
                    attrs.set(crate::introspect::NUM_ACTORS, num_live);
                    attrs.set(crate::introspect::SYSTEM_CHILDREN, system_children);
                    attrs.set(crate::introspect::STOPPED_CHILDREN, stopped_children);
                    attrs.set(
                        crate::introspect::STOPPED_RETENTION_CAP,
                        stopped_retention_cap,
                    );
                    attrs.set(crate::introspect::IS_POISONED, is_poisoned);
                    attrs.set(crate::introspect::FAILED_ACTOR_COUNT, failed_actor_count);
                    let attrs_json =
                        serde_json::to_string(&attrs).unwrap_or_else(|_| "{}".to_string());

                    return IntrospectResult {
                        identity: proc_id.to_string(),
                        attrs: attrs_json,
                        children,
                        parent: None,
                        as_of: humantime::format_rfc3339_millis(std::time::SystemTime::now())
                            .to_string(),
                    };
                }
            }

            {
                let mut error_attrs = hyperactor_config::Attrs::new();
                error_attrs.set(hyperactor::introspect::ERROR_CODE, "not_found".to_string());
                error_attrs.set(
                    hyperactor::introspect::ERROR_MESSAGE,
                    format!("child {} not found", child_ref),
                );
                IntrospectResult {
                    identity: String::new(),
                    attrs: serde_json::to_string(&error_attrs).unwrap_or_else(|_| "{}".to_string()),
                    children: Vec::new(),
                    parent: None,
                    as_of: humantime::format_rfc3339_millis(std::time::SystemTime::now())
                        .to_string(),
                }
            }
        });

        if let Some(delay) = &self.mesh_orphan_timeout {
            this.self_message_with_delay(SelfCheck::default(), *delay)?;
        }
        Ok(())
    }
}

#[async_trait]
#[hyperactor::handle(MeshAgentMessage)]
impl MeshAgentMessageHandler for ProcAgent {
    async fn configure(
        &mut self,
        cx: &Context<Self>,
        rank: usize,
        forwarder: ChannelAddr,
        supervisor: Option<hyperactor_reference::PortRef<ActorSupervisionEvent>>,
        address_book: HashMap<hyperactor_reference::ProcId, ChannelAddr>,
        configured: hyperactor_reference::PortRef<usize>,
        record_supervision_events: bool,
    ) -> Result<(), anyhow::Error> {
        anyhow::ensure!(
            self.state.is_unconfigured_v0(),
            "mesh agent cannot be (re-)configured"
        );
        self.record_supervision_events = record_supervision_events;

        let client = MailboxClient::new(channel::dial(forwarder)?);
        let router =
            DialMailboxRouter::new_with_default_direct_addressed_remote_only(client.into_boxed());

        for (proc_id, addr) in address_book {
            router.bind(proc_id.into(), addr);
        }

        let sender = take(&mut self.state).into_unconfigured_v0().unwrap();
        assert!(sender.configure(router.into_boxed()));

        // This is a bit suboptimal: ideally we'd set the supervisor first, to correctly report
        // any errors that occur during configuration. However, these should anyway be correctly
        // caught on process exit.
        self.state = State::ConfiguredV0 {
            sender,
            rank,
            supervisor,
        };
        configured.send(cx, rank)?;

        Ok(())
    }

    async fn gspawn(
        &mut self,
        cx: &Context<Self>,
        actor_type: String,
        actor_name: String,
        params_data: Data,
        status_port: hyperactor_reference::PortRef<GspawnResult>,
    ) -> Result<(), anyhow::Error> {
        anyhow::ensure!(
            self.state.is_configured_v0(),
            "mesh agent is not v0 configured"
        );
        let actor_id = match self
            .remote
            .gspawn(
                &self.proc,
                &actor_type,
                &actor_name,
                params_data,
                cx.headers().clone(),
            )
            .await
        {
            Ok(id) => id,
            Err(err) => {
                status_port.send(cx, GspawnResult::Error(format!("gspawn failed: {}", err)))?;
                return Err(anyhow::anyhow!("gspawn failed"));
            }
        };
        status_port.send(
            cx,
            GspawnResult::Success {
                rank: self.state.rank().unwrap(),
                actor_id,
            },
        )?;
        self.publish_introspect_properties(cx);
        Ok(())
    }

    async fn stop_actor(
        &mut self,
        cx: &Context<Self>,
        actor_id: hyperactor_reference::ActorId,
        timeout_ms: u64,
        reason: String,
    ) -> Result<StopActorResult, anyhow::Error> {
        tracing::info!(
            name = "StopActor",
            %actor_id,
            actor_name = actor_id.name(),
            %reason,
        );

        let result = if let Some(mut status) = self.proc.stop_actor(&actor_id, reason) {
            match tokio::time::timeout(
                tokio::time::Duration::from_millis(timeout_ms),
                status.wait_for(|state: &ActorStatus| state.is_terminal()),
            )
            .await
            {
                Ok(_) => Ok(StopActorResult::Success),
                Err(_) => Ok(StopActorResult::Timeout),
            }
        } else {
            Ok(StopActorResult::NotFound)
        };
        self.publish_introspect_properties(cx);
        result
    }

    async fn status(
        &mut self,
        cx: &Context<Self>,
        status_port: hyperactor_reference::PortRef<(usize, bool)>,
    ) -> Result<(), anyhow::Error> {
        match &self.state {
            State::ConfiguredV0 { rank, .. } => {
                // v0 path: configured with a concrete rank
                status_port.send(cx, (*rank, true))?;
                Ok(())
            }
            State::UnconfiguredV0 { .. } => {
                // v0 path but not configured yet
                Err(anyhow::anyhow!(
                    "status unavailable: v0 agent not configured (waiting for Configure)"
                ))
            }
            State::V1 => {
                // v1/owned path does not support status (no rank semantics)
                Err(anyhow::anyhow!(
                    "status unsupported in v1/owned path (no rank)"
                ))
            }
            State::Invalid => Err(anyhow::anyhow!(
                "status unavailable: agent in invalid state"
            )),
        }
    }
}

#[async_trait]
impl Handler<ActorSupervisionEvent> for ProcAgent {
    async fn handle(
        &mut self,
        cx: &Context<Self>,
        event: ActorSupervisionEvent,
    ) -> anyhow::Result<()> {
        if self.record_supervision_events {
            if event.is_error() {
                tracing::warn!(
                    name = "SupervisionEvent",
                    proc_id = %self.proc.proc_id(),
                    %event,
                    "recording supervision error",
                );
            } else {
                tracing::debug!(
                    name = "SupervisionEvent",
                    proc_id = %self.proc.proc_id(),
                    %event,
                    "recording non-error supervision event",
                );
            }
            self.supervision_events
                .entry(event.actor_id.clone())
                .or_default()
                .push(event.clone());
            // Defer republish so introspection picks up is_poisoned /
            // failed_actor_count without blocking the message loop.
            // Multiple rapid events coalesce into one republish.
            if !self.introspect_dirty {
                self.introspect_dirty = true;
                let _ = cx.self_message_with_delay(
                    RepublishIntrospect,
                    std::time::Duration::from_millis(100),
                );
            }
        }
        if let Some(supervisor) = self.state.supervisor() {
            supervisor.send(cx, event)?;
        } else if !self.record_supervision_events {
            // If there is no supervisor, and nothing is recording these, crash
            // the whole process.
            tracing::error!(
                name = "supervision_event_transmit_failed",
                proc_id = %cx.self_id().proc_id(),
                %event,
                "could not propagate supervision event, crashing",
            );

            // We should have a custom "crash" function here, so that this works
            // in testing of the LocalAllocator, etc.
            std::process::exit(1);
        }
        Ok(())
    }
}

#[async_trait]
impl Handler<RepublishIntrospect> for ProcAgent {
    async fn handle(&mut self, cx: &Context<Self>, _: RepublishIntrospect) -> anyhow::Result<()> {
        if self.introspect_dirty {
            self.introspect_dirty = false;
            self.publish_introspect_properties(cx);
        }
        Ok(())
    }
}

// Implement the resource behavior for managing actors:

/// Actor spec.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Named)]
pub struct ActorSpec {
    /// registered actor type
    pub actor_type: String,
    /// serialized parameters
    pub params_data: Data,
}
wirevalue::register_type!(ActorSpec);

/// Actor state.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Named, Bind, Unbind)]
pub struct ActorState {
    /// The actor's ID.
    pub actor_id: hyperactor_reference::ActorId,
    /// The rank of the proc that created the actor. This is before any slicing.
    pub create_rank: usize,
    // TODO status: ActorStatus,
    pub supervision_events: Vec<ActorSupervisionEvent>,
}
wirevalue::register_type!(ActorState);

#[async_trait]
impl Handler<resource::CreateOrUpdate<ActorSpec>> for ProcAgent {
    async fn handle(
        &mut self,
        cx: &Context<Self>,
        create_or_update: resource::CreateOrUpdate<ActorSpec>,
    ) -> anyhow::Result<()> {
        if self.actor_states.contains_key(&create_or_update.name) {
            // There is no update.
            return Ok(());
        }
        let create_rank = create_or_update.rank.unwrap();
        // If there have been supervision events for any actors on this proc,
        // we disallow spawning new actors on it, as this proc may be in an
        // invalid state.
        if !self.supervision_events.is_empty() {
            self.actor_states.insert(
                create_or_update.name.clone(),
                ActorInstanceState {
                    spawn: Err(anyhow::anyhow!(
                        "Cannot spawn new actors on mesh with supervision events"
                    )),
                    create_rank,
                    stopped: false,
                    expiry_time: None,
                },
            );
            return Ok(());
        }

        let ActorSpec {
            actor_type,
            params_data,
        } = create_or_update.spec;
        self.actor_states.insert(
            create_or_update.name.clone(),
            ActorInstanceState {
                create_rank,
                spawn: self
                    .remote
                    .gspawn(
                        &self.proc,
                        &actor_type,
                        &create_or_update.name.to_string(),
                        params_data,
                        cx.headers().clone(),
                    )
                    .await,
                stopped: false,
                expiry_time: None,
            },
        );

        self.publish_introspect_properties(cx);
        Ok(())
    }
}

#[async_trait]
impl Handler<resource::Stop> for ProcAgent {
    async fn handle(&mut self, cx: &Context<Self>, message: resource::Stop) -> anyhow::Result<()> {
        // We don't remove the actor from the state map, instead we just store
        // its state as Stopped.
        let actor = self.actor_states.get_mut(&message.name);
        // Have to separate stop_actor from setting "stopped" because it borrows
        // as mutable and cannot have self borrowed mutably twice.
        let actor_id = match actor {
            Some(actor_state) => {
                match &actor_state.spawn {
                    Ok(actor_id) => {
                        if actor_state.stopped {
                            None
                        } else {
                            actor_state.stopped = true;
                            Some(actor_id.clone())
                        }
                    }
                    // If the original spawn had failed, the actor is still considered
                    // successfully stopped.
                    Err(_) => None,
                }
            }
            // TODO: represent unknown rank
            None => None,
        };
        let timeout = hyperactor_config::global::get(hyperactor::config::STOP_ACTOR_TIMEOUT);
        if let Some(actor_id) = actor_id {
            // The orphaned actor SelfCheck will consult the stopped field before
            // trying to stop any actor again.
            // While this function returns a Result, it never returns an Err
            // value so we can simply expect without any failure handling.
            self.stop_actor(cx, actor_id, timeout.as_millis() as u64, message.reason)
                .await
                .expect("stop_actor cannot fail");
        }

        Ok(())
    }
}

/// Handles `StopAll` by coordinating an orderly stop of child actors and then
/// exiting the process. This handler never returns to the caller: it calls
/// `std::process::exit(0/1)` after shutdown. Any sender must *not* expect a
/// reply or send any further message, and should watch `ProcStatus` instead.
#[async_trait]
impl Handler<resource::StopAll> for ProcAgent {
    async fn handle(
        &mut self,
        cx: &Context<Self>,
        message: resource::StopAll,
    ) -> anyhow::Result<()> {
        let timeout = hyperactor_config::global::get(hyperactor::config::STOP_ACTOR_TIMEOUT);
        // By passing in the self context, destroy_and_wait will stop this agent
        // last, after all others are stopped.
        let stop_result = self
            .destroy_and_wait_except_current(cx, timeout, &message.reason)
            .await;
        // Exit here to cleanup all remaining resources held by the process.
        // This means ProcAgent will never run cleanup or any other code
        // from exiting its root actor. Senders of this message should never
        // send any further messages or expect a reply.
        match stop_result {
            Ok((stopped_actors, aborted_actors)) => {
                // No need to clean up any state, the process is exiting.
                tracing::info!(
                    actor = %cx.self_id(),
                    "exiting process after receiving StopAll message on ProcAgent. \
                    stopped actors = {:?}, aborted actors = {:?}",
                    stopped_actors.into_iter().map(|a| a.to_string()).collect::<Vec<_>>(),
                    aborted_actors.into_iter().map(|a| a.to_string()).collect::<Vec<_>>(),
                );
                if let Some(tx) = self.shutdown_tx.take() {
                    let _ = tx.send(0);
                    return Ok(());
                }
                std::process::exit(0);
            }
            Err(e) => {
                tracing::error!(actor = %cx.self_id(), "failed to stop all actors on ProcAgent: {:?}", e);
                if let Some(tx) = self.shutdown_tx.take() {
                    let _ = tx.send(1);
                    return Ok(());
                }
                std::process::exit(1);
            }
        }
    }
}

#[async_trait]
impl Handler<resource::GetRankStatus> for ProcAgent {
    async fn handle(
        &mut self,
        cx: &Context<Self>,
        get_rank_status: resource::GetRankStatus,
    ) -> anyhow::Result<()> {
        use crate::StatusOverlay;
        use crate::resource::Status;

        let (rank, status) = match self.actor_states.get(&get_rank_status.name) {
            Some(ActorInstanceState {
                spawn: Ok(actor_id),
                create_rank,
                stopped,
                ..
            }) => {
                if *stopped {
                    (*create_rank, resource::Status::Stopped)
                } else {
                    let supervision_events = self
                        .supervision_events
                        .get(actor_id)
                        .map_or_else(Vec::new, |a| a.clone());
                    (
                        *create_rank,
                        if supervision_events.is_empty() {
                            resource::Status::Running
                        } else {
                            resource::Status::Failed(format!(
                                "because of supervision events: {:?}",
                                supervision_events
                            ))
                        },
                    )
                }
            }
            Some(ActorInstanceState {
                spawn: Err(e),
                create_rank,
                ..
            }) => (*create_rank, Status::Failed(e.to_string())),
            // TODO: represent unknown rank
            None => (usize::MAX, Status::NotExist),
        };

        // Send a sparse overlay update. If rank is unknown, emit an
        // empty overlay.
        let overlay = if rank == usize::MAX {
            StatusOverlay::new()
        } else {
            StatusOverlay::try_from_runs(vec![(rank..(rank + 1), status)])
                .expect("valid single-run overlay")
        };
        let result = get_rank_status.reply.send(cx, overlay);
        // Ignore errors, because returning Err from here would cause the ProcAgent
        // to be stopped, which would prevent querying and spawning other actors.
        // This only means some actor that requested the state of an actor failed to receive it.
        if let Err(e) = result {
            tracing::warn!(
                actor = %cx.self_id(),
                "failed to send GetRankStatus reply to {} due to error: {}",
                get_rank_status.reply.port_id().actor_id(),
                e
            );
        }
        Ok(())
    }
}

#[async_trait]
impl Handler<resource::GetState<ActorState>> for ProcAgent {
    async fn handle(
        &mut self,
        cx: &Context<Self>,
        get_state: resource::GetState<ActorState>,
    ) -> anyhow::Result<()> {
        let state = match self.actor_states.get(&get_state.name) {
            Some(ActorInstanceState {
                create_rank,
                spawn: Ok(actor_id),
                stopped,
                ..
            }) => {
                let supervision_events = self
                    .supervision_events
                    .get(actor_id)
                    .map_or_else(Vec::new, |a| a.clone());
                let status = if *stopped {
                    resource::Status::Stopped
                } else if supervision_events.is_empty() {
                    resource::Status::Running
                } else {
                    resource::Status::Failed(format!(
                        "because of supervision events: {:?}",
                        supervision_events
                    ))
                };
                resource::State {
                    name: get_state.name.clone(),
                    status,
                    state: Some(ActorState {
                        actor_id: actor_id.clone(),
                        create_rank: *create_rank,
                        supervision_events,
                    }),
                }
            }
            Some(ActorInstanceState { spawn: Err(e), .. }) => resource::State {
                name: get_state.name.clone(),
                status: resource::Status::Failed(e.to_string()),
                state: None,
            },
            None => resource::State {
                name: get_state.name.clone(),
                status: resource::Status::NotExist,
                state: None,
            },
        };

        let result = get_state.reply.send(cx, state);
        // Ignore errors, because returning Err from here would cause the ProcAgent
        // to be stopped, which would prevent querying and spawning other actors.
        // This only means some actor that requested the state of an actor failed to receive it.
        if let Err(e) = result {
            tracing::warn!(
                actor = %cx.self_id(),
                "failed to send GetState reply to {} due to error: {}",
                get_state.reply.port_id().actor_id(),
                e
            );
        }
        Ok(())
    }
}

#[async_trait]
impl Handler<resource::KeepaliveGetState<ActorState>> for ProcAgent {
    async fn handle(
        &mut self,
        cx: &Context<Self>,
        message: resource::KeepaliveGetState<ActorState>,
    ) -> anyhow::Result<()> {
        // Same impl as GetState, but additionally update the expiry time on the actor.
        if let Ok(instance_state) = self
            .actor_states
            .get_mut(&message.get_state.name)
            .ok_or_else(|| {
                anyhow::anyhow!(
                    "attempting to register a keepalive for an actor that doesn't exist: {}",
                    message.get_state.name
                )
            })
        {
            instance_state.expiry_time = Some(message.expires_after);
        }

        // Forward the rest of the impl to GetState.
        <Self as Handler<resource::GetState<ActorState>>>::handle(self, cx, message.get_state).await
    }
}

/// A local handler to get a new client instance on the proc.
/// This is used to create root client instances.
#[derive(Debug, hyperactor::Handler, hyperactor::HandleClient)]
pub struct NewClientInstance {
    #[reply]
    pub client_instance: PortHandle<Instance<()>>,
}

#[async_trait]
impl Handler<NewClientInstance> for ProcAgent {
    async fn handle(
        &mut self,
        cx: &Context<Self>,
        NewClientInstance { client_instance }: NewClientInstance,
    ) -> anyhow::Result<()> {
        let (instance, _handle) = self.proc.instance("client")?;
        client_instance.send(cx, instance)?;
        Ok(())
    }
}

/// A handler to get a clone of the proc managed by this agent.
/// This is used to obtain the local proc from a host mesh.
#[derive(Debug, hyperactor::Handler, hyperactor::HandleClient)]
pub struct GetProc {
    #[reply]
    pub proc: PortHandle<Proc>,
}

#[async_trait]
impl Handler<GetProc> for ProcAgent {
    async fn handle(
        &mut self,
        cx: &Context<Self>,
        GetProc { proc }: GetProc,
    ) -> anyhow::Result<()> {
        proc.send(cx, self.proc.clone())?;
        Ok(())
    }
}

#[async_trait]
impl Handler<SelfCheck> for ProcAgent {
    async fn handle(&mut self, cx: &Context<Self>, _: SelfCheck) -> anyhow::Result<()> {
        // Check each actor's expiry time. If the current time is past the expiry,
        // stop the actor. This allows automatic cleanup when a controller disappears
        // but owned resources remain. It is important that this check runs on the
        // same proc as the child actor itself, since the controller could be dead or
        // disconnected.
        let Some(duration) = &self.mesh_orphan_timeout else {
            return Ok(());
        };
        let duration = duration.clone();
        let now = std::time::SystemTime::now();

        // Collect expired actors before mutating, since stop_actor borrows &mut self.
        let expired: Vec<(Name, hyperactor_reference::ActorId)> = self
            .actor_states
            .iter()
            .filter_map(|(name, state)| {
                let expiry = state.expiry_time?;
                // If the actor was already stopped we don't need to stop it again.
                if now > expiry && !state.stopped {
                    if let Ok(actor_id) = &state.spawn {
                        return Some((name.clone(), actor_id.clone()));
                    }
                }
                None
            })
            .collect();

        if !expired.is_empty() {
            tracing::info!(
                "stopping {} orphaned actors past their keepalive expiry",
                expired.len(),
            );
        }

        let timeout = hyperactor_config::global::get(hyperactor::config::STOP_ACTOR_TIMEOUT);
        for (name, actor_id) in expired {
            if let Some(state) = self.actor_states.get_mut(&name) {
                state.stopped = true;
            }
            self.stop_actor(
                cx,
                actor_id,
                timeout.as_millis() as u64,
                "orphaned".to_string(),
            )
            .await
            .expect("stop_actor cannot fail");
        }

        // Reschedule.
        cx.self_message_with_delay(SelfCheck::default(), duration)?;
        Ok(())
    }
}

/// A mailbox sender that initially queues messages, and then relays them to
/// an underlying sender once configured.
#[derive(Clone)]
pub(crate) struct ReconfigurableMailboxSender {
    state: Arc<RwLock<ReconfigurableMailboxSenderState>>,
}

impl std::fmt::Debug for ReconfigurableMailboxSender {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        // Not super helpful, but we definitely don't wan to acquire any locks
        // in a Debug formatter.
        f.debug_struct("ReconfigurableMailboxSender").finish()
    }
}

/// A capability wrapper granting access to the configured mailbox
/// sender.
///
/// This type exists to tie the lifetime of any `&BoxedMailboxSender`
/// reference to a lock guard, so the underlying state cannot be
/// reconfigured while the reference is in use.
///
/// A **read** guard is sufficient because we only need to *observe*
/// and borrow the configured sender, not mutate state. While a
/// `RwLockReadGuard` is held, `configure()` cannot acquire the write
/// lock, so the state cannot transition from `Configured(..)` to any
/// other variant during the guard’s lifetime.
pub(crate) struct ReconfigurableMailboxSenderInner<'a> {
    guard: RwLockReadGuard<'a, ReconfigurableMailboxSenderState>,
}

impl<'a> ReconfigurableMailboxSenderInner<'a> {
    pub(crate) fn as_configured(&self) -> Option<&BoxedMailboxSender> {
        self.guard.as_configured()
    }
}

type Post = (MessageEnvelope, PortHandle<Undeliverable<MessageEnvelope>>);

#[derive(EnumAsInner, Debug)]
enum ReconfigurableMailboxSenderState {
    Queueing(Mutex<Vec<Post>>),
    Configured(BoxedMailboxSender),
}

impl ReconfigurableMailboxSender {
    pub(crate) fn new() -> Self {
        Self {
            state: Arc::new(RwLock::new(ReconfigurableMailboxSenderState::Queueing(
                Mutex::new(Vec::new()),
            ))),
        }
    }

    /// Configure this mailbox with the provided sender. This will first
    /// enqueue any pending messages onto the sender; future messages are
    /// posted directly to the configured sender.
    pub(crate) fn configure(&self, sender: BoxedMailboxSender) -> bool {
        // Hold the write lock until all queued messages are flushed.
        let mut state = self.state.write().unwrap();
        if state.is_configured() {
            return false;
        }

        // Install the configured sender exactly once.
        let queued = std::mem::replace(
            &mut *state,
            ReconfigurableMailboxSenderState::Configured(sender),
        );

        // Borrow the configured sender from the state (stable while
        // we hold the lock).
        let configured_sender = state.as_configured().expect("just configured");

        // Flush the old queue while still holding the write lock.
        for (envelope, return_handle) in queued.into_queueing().unwrap().into_inner().unwrap() {
            configured_sender.post(envelope, return_handle);
        }

        true
    }

    pub(crate) fn as_inner<'a>(
        &'a self,
    ) -> Result<ReconfigurableMailboxSenderInner<'a>, anyhow::Error> {
        let state = self.state.read().unwrap();
        if state.is_configured() {
            Ok(ReconfigurableMailboxSenderInner { guard: state })
        } else {
            Err(anyhow::anyhow!("cannot get inner sender: not configured"))
        }
    }
}

impl MailboxSender for ReconfigurableMailboxSender {
    fn post(
        &self,
        envelope: MessageEnvelope,
        return_handle: PortHandle<Undeliverable<MessageEnvelope>>,
    ) {
        match &*self.state.read().unwrap() {
            ReconfigurableMailboxSenderState::Queueing(queue) => {
                queue.lock().unwrap().push((envelope, return_handle));
            }
            ReconfigurableMailboxSenderState::Configured(sender) => {
                sender.post(envelope, return_handle);
            }
        }
    }

    fn post_unchecked(
        &self,
        envelope: MessageEnvelope,
        return_handle: PortHandle<Undeliverable<MessageEnvelope>>,
    ) {
        match &*self.state.read().unwrap() {
            ReconfigurableMailboxSenderState::Queueing(queue) => {
                queue.lock().unwrap().push((envelope, return_handle));
            }
            ReconfigurableMailboxSenderState::Configured(sender) => {
                sender.post_unchecked(envelope, return_handle);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use std::sync::Mutex;

    use hyperactor::mailbox::BoxedMailboxSender;
    use hyperactor::mailbox::Mailbox;
    use hyperactor::mailbox::MailboxSender;
    use hyperactor::mailbox::MessageEnvelope;
    use hyperactor::mailbox::PortHandle;
    use hyperactor::mailbox::Undeliverable;
    use hyperactor::testing::ids::test_actor_id;
    use hyperactor::testing::ids::test_port_id;
    use hyperactor_config::Flattrs;

    use super::*;

    #[derive(Debug, Clone)]
    struct QueueingMailboxSender {
        messages: Arc<Mutex<Vec<MessageEnvelope>>>,
    }

    impl QueueingMailboxSender {
        fn new() -> Self {
            Self {
                messages: Arc::new(Mutex::new(Vec::new())),
            }
        }

        fn get_messages(&self) -> Vec<MessageEnvelope> {
            self.messages.lock().unwrap().clone()
        }
    }

    impl MailboxSender for QueueingMailboxSender {
        fn post_unchecked(
            &self,
            envelope: MessageEnvelope,
            _return_handle: PortHandle<Undeliverable<MessageEnvelope>>,
        ) {
            self.messages.lock().unwrap().push(envelope);
        }
    }

    // Helper function to create a test message envelope
    fn envelope(data: u64) -> MessageEnvelope {
        MessageEnvelope::serialize(
            test_actor_id("world_0", "sender"),
            test_port_id("world_0", "receiver", 1),
            &data,
            Flattrs::new(),
        )
        .unwrap()
    }

    fn return_handle() -> PortHandle<Undeliverable<MessageEnvelope>> {
        let mbox = Mailbox::new_detached(test_actor_id("0", "test"));
        let (port, _receiver) = mbox.open_port::<Undeliverable<MessageEnvelope>>();
        port
    }

    #[test]
    fn test_queueing_before_configure() {
        let sender = ReconfigurableMailboxSender::new();

        let test_sender = QueueingMailboxSender::new();
        let boxed_sender = BoxedMailboxSender::new(test_sender.clone());

        let return_handle = return_handle();
        sender.post(envelope(1), return_handle.clone());
        sender.post(envelope(2), return_handle.clone());

        assert_eq!(test_sender.get_messages().len(), 0);

        sender.configure(boxed_sender);

        let messages = test_sender.get_messages();
        assert_eq!(messages.len(), 2);

        assert_eq!(messages[0].deserialized::<u64>().unwrap(), 1);
        assert_eq!(messages[1].deserialized::<u64>().unwrap(), 2);
    }

    #[test]
    fn test_direct_delivery_after_configure() {
        // Create a ReconfigurableMailboxSender
        let sender = ReconfigurableMailboxSender::new();

        let test_sender = QueueingMailboxSender::new();
        let boxed_sender = BoxedMailboxSender::new(test_sender.clone());
        sender.configure(boxed_sender);

        let return_handle = return_handle();
        sender.post(envelope(3), return_handle.clone());
        sender.post(envelope(4), return_handle.clone());

        let messages = test_sender.get_messages();
        assert_eq!(messages.len(), 2);

        assert_eq!(messages[0].deserialized::<u64>().unwrap(), 3);
        assert_eq!(messages[1].deserialized::<u64>().unwrap(), 4);
    }

    #[test]
    fn test_multiple_configurations() {
        let sender = ReconfigurableMailboxSender::new();
        let boxed_sender = BoxedMailboxSender::new(QueueingMailboxSender::new());

        assert!(sender.configure(boxed_sender.clone()));
        assert!(!sender.configure(boxed_sender));
    }

    #[test]
    fn test_mixed_queueing_and_direct_delivery() {
        let sender = ReconfigurableMailboxSender::new();

        let test_sender = QueueingMailboxSender::new();
        let boxed_sender = BoxedMailboxSender::new(test_sender.clone());

        let return_handle = return_handle();
        sender.post(envelope(5), return_handle.clone());
        sender.post(envelope(6), return_handle.clone());

        sender.configure(boxed_sender);

        sender.post(envelope(7), return_handle.clone());
        sender.post(envelope(8), return_handle.clone());

        let messages = test_sender.get_messages();
        assert_eq!(messages.len(), 4);

        assert_eq!(messages[0].deserialized::<u64>().unwrap(), 5);
        assert_eq!(messages[1].deserialized::<u64>().unwrap(), 6);
        assert_eq!(messages[2].deserialized::<u64>().unwrap(), 7);
        assert_eq!(messages[3].deserialized::<u64>().unwrap(), 8);
    }

    // A no-op actor used to test direct proc-level spawning.
    #[derive(Debug, Default)]
    #[hyperactor::export(handlers = [])]
    struct ExtraActor;
    impl hyperactor::Actor for ExtraActor {}

    // Verifies that QueryChild(Reference::Proc) on a ProcAgent returns
    // a live IntrospectResult whose children reflect actors spawned
    // directly on the proc — i.e. via proc.spawn(), which bypasses the
    // gspawn message handler and therefore never triggers
    // publish_introspect_properties.
    //
    // Exercises PA-1 (see mesh_admin module doc).
    //
    // Regression guard for the bug introduced in 9a08d559: removing
    // handle_introspect left publish_introspect_properties as the only
    // update path, which missed supervision-spawned actors (e.g. every
    // sieve actor after sieve[0]). See also
    // mesh_admin::tests::test_proc_children_reflect_directly_spawned_actors.
    #[tokio::test]
    async fn test_query_child_proc_returns_live_children() {
        use hyperactor::Proc;
        use hyperactor::actor::ActorStatus;
        use hyperactor::channel::ChannelTransport;
        use hyperactor::introspect::IntrospectMessage;
        use hyperactor::introspect::IntrospectResult;
        use hyperactor::reference as hyperactor_reference;

        let proc = Proc::direct(ChannelTransport::Unix.any(), "test_proc".to_string()).unwrap();
        let agent_handle = ProcAgent::boot_v1(proc.clone(), None).unwrap();

        // Wait for ProcAgent to finish init.
        agent_handle
            .status()
            .wait_for(|s| matches!(s, ActorStatus::Idle))
            .await
            .unwrap();

        // Client instance for opening reply ports.
        let client_proc = Proc::direct(ChannelTransport::Unix.any(), "client".to_string()).unwrap();
        let (client, _client_handle) = client_proc.instance("client").unwrap();

        let agent_id = proc.proc_id().actor_id(PROC_AGENT_ACTOR_NAME, 0);
        let port =
            hyperactor_reference::PortRef::<IntrospectMessage>::attest_message_port(&agent_id);

        // Helper: send QueryChild(Proc) and return the payload with a
        // timeout so a misrouted reply fails fast rather than hanging.
        let query = |client: &hyperactor::Instance<()>| {
            let (reply_port, reply_rx) = client.open_once_port::<IntrospectResult>();
            port.send(
                client,
                IntrospectMessage::QueryChild {
                    child_ref: hyperactor_reference::Reference::Proc(proc.proc_id().clone()),
                    reply: reply_port.bind(),
                },
            )
            .unwrap();
            reply_rx
        };
        let recv = |rx: hyperactor::mailbox::OncePortReceiver<IntrospectResult>| async move {
            tokio::time::timeout(std::time::Duration::from_secs(5), rx.recv())
                .await
                .expect("QueryChild(Proc) timed out — reply never delivered")
                .expect("reply channel closed")
        };

        // Initial query: ProcAgent itself should appear in children.
        let payload = recv(query(&client)).await;
        // Verify this is a proc node by checking attrs contain node_type=proc.
        let attrs: hyperactor_config::Attrs =
            serde_json::from_str(&payload.attrs).expect("valid attrs JSON");
        assert_eq!(
            attrs.get(crate::introspect::NODE_TYPE).map(String::as_str),
            Some("proc"),
            "expected node_type=proc in attrs, got {:?}",
            payload.attrs
        );
        assert!(
            payload
                .children
                .iter()
                .any(|c| c.contains(PROC_AGENT_ACTOR_NAME)),
            "initial children {:?} should contain proc_agent",
            payload.children
        );
        let initial_count = payload.children.len();

        // Spawn an actor directly on the proc, bypassing ProcAgent's
        // gspawn message handler. This is how supervision-spawned
        // actors (e.g. sieve children) are created.
        proc.spawn("extra_actor", ExtraActor).unwrap();

        // Second query: extra_actor must appear without any republish.
        let payload2 = recv(query(&client)).await;
        let attrs2: hyperactor_config::Attrs =
            serde_json::from_str(&payload2.attrs).expect("valid attrs JSON");
        assert_eq!(
            attrs2.get(crate::introspect::NODE_TYPE).map(String::as_str),
            Some("proc"),
            "expected node_type=proc in attrs, got {:?}",
            payload2.attrs
        );
        assert!(
            payload2.children.iter().any(|c| c.contains("extra_actor")),
            "after direct spawn, children {:?} should contain extra_actor",
            payload2.children
        );
        assert!(
            payload2.children.len() > initial_count,
            "expected at least {} children after direct spawn, got {:?}",
            initial_count + 1,
            payload2.children
        );
    }

    // Exercises S12 (see introspect module doc): introspection must
    // not impair actor liveness. Rapidly spawns and stops
    // actors while concurrently querying QueryChild(Reference::Proc).
    // The spawn/stop loop must complete within the timeout and the
    // iteration count must match -- if DashMap convoy starvation
    // blocks the proc, the timeout fires and the test fails.
    #[tokio::test]
    async fn test_rapid_spawn_stop_does_not_stall_proc_agent() {
        use std::sync::Arc;
        use std::sync::atomic::AtomicUsize;
        use std::sync::atomic::Ordering;

        use hyperactor::Proc;
        use hyperactor::actor::ActorStatus;
        use hyperactor::channel::ChannelTransport;
        use hyperactor::introspect::IntrospectMessage;
        use hyperactor::introspect::IntrospectResult;
        use hyperactor::reference as hyperactor_reference;

        let proc = Proc::direct(ChannelTransport::Unix.any(), "test_proc".to_string()).unwrap();
        let agent_handle = ProcAgent::boot_v1(proc.clone(), None).unwrap();

        agent_handle
            .status()
            .wait_for(|s| matches!(s, ActorStatus::Idle))
            .await
            .unwrap();

        let client_proc = Proc::direct(ChannelTransport::Unix.any(), "client".to_string()).unwrap();
        let (client, _client_handle) = client_proc.instance("client").unwrap();

        let agent_id = proc.proc_id().actor_id(PROC_AGENT_ACTOR_NAME, 0);
        let port =
            hyperactor_reference::PortRef::<IntrospectMessage>::attest_message_port(&agent_id);

        // Concurrent query task: send QueryChild(Proc) every 10ms.
        let query_client_proc =
            Proc::direct(ChannelTransport::Unix.any(), "query_client".to_string()).unwrap();
        let (query_client, _qc_handle) = query_client_proc.instance("qc").unwrap();
        let query_port = port.clone();
        let query_proc_id = proc.proc_id().clone();
        let query_count = Arc::new(AtomicUsize::new(0));
        let query_count_clone = query_count.clone();
        let query_task = tokio::spawn(async move {
            loop {
                let (reply_port, reply_rx) = query_client.open_once_port::<IntrospectResult>();
                if query_port
                    .send(
                        &query_client,
                        IntrospectMessage::QueryChild {
                            child_ref: hyperactor_reference::Reference::Proc(query_proc_id.clone()),
                            reply: reply_port.bind(),
                        },
                    )
                    .is_err()
                {
                    break;
                }
                match tokio::time::timeout(std::time::Duration::from_secs(2), reply_rx.recv()).await
                {
                    Ok(Ok(_)) => {
                        query_count_clone.fetch_add(1, Ordering::Relaxed);
                    }
                    _ => {} // Transient failures expected during churn
                }
                tokio::time::sleep(std::time::Duration::from_millis(10)).await;
            }
        });

        // Rapid spawn/stop loop with liveness timeout.
        const ITERATIONS: usize = 200;
        let mut completed = 0usize;
        let result = tokio::time::timeout(std::time::Duration::from_secs(30), async {
            for i in 0..ITERATIONS {
                let name = format!("churn_{}", i);
                let handle = proc.spawn(&name, ExtraActor).unwrap();
                let actor_id = handle.actor_id().clone();
                if let Some(mut status) = proc.stop_actor(&actor_id, "churn".to_string()) {
                    let _ = tokio::time::timeout(
                        std::time::Duration::from_secs(5),
                        status.wait_for(ActorStatus::is_terminal),
                    )
                    .await;
                }
                completed += 1;
            }
        })
        .await;

        query_task.abort();
        let _ = query_task.await; // Join to suppress noisy panic on drop.

        assert!(
            result.is_ok(),
            "spawn/stop loop stalled after {completed}/{ITERATIONS} iterations — \
             DashMap convoy starvation likely"
        );
        assert_eq!(
            completed, ITERATIONS,
            "expected {ITERATIONS} completed iterations, got {completed}"
        );
        assert!(
            query_count.load(Ordering::Relaxed) > 0,
            "concurrent QueryChild queries never succeeded — query task may not have run"
        );

        // Final consistency check: QueryChild should still work.
        let (reply_port, reply_rx) = client.open_once_port::<IntrospectResult>();
        port.send(
            &client,
            IntrospectMessage::QueryChild {
                child_ref: hyperactor_reference::Reference::Proc(proc.proc_id().clone()),
                reply: reply_port.bind(),
            },
        )
        .unwrap();
        let final_payload =
            tokio::time::timeout(std::time::Duration::from_secs(5), reply_rx.recv())
                .await
                .expect("final QueryChild timed out")
                .expect("final QueryChild channel closed");
        let attrs: hyperactor_config::Attrs =
            serde_json::from_str(&final_payload.attrs).expect("valid attrs JSON");
        assert_eq!(
            attrs.get(crate::introspect::NODE_TYPE).map(String::as_str),
            Some("proc"),
        );
    }
}