Fix handler lifecycle race condition for shutdown receiver

The message-server handler's register() function needed a ShutdownReceiver,
but it was only stored during the async setup() method. If an actor called
register() during its synchronous init function, setup() hadn't run yet.

Additionally, setup() had a tokio::select! that immediately took the
receiver from the mutex, racing with register().

Changes:
- Add ShutdownController to HandlerContext for early subscription
- Make ShutdownController Clone-able with Arc<Mutex<>>
- Create ShutdownController early in build_actor_resources()
- Message-server handler now subscribes during setup_host_functions_composite()
- Simplified setup() to just wait for registered_notify signal

Author: colinrozzi

crates/theater-handler-message-server/src/lib.rs

@@ -569,40 +569,30 @@ impl Handler for MessageServerHandler
     ) -> Pin<Box<dyn Future<Output = Result<()>> + Send>> {
         info!("Message server handler setup (passive mode)");
 
-        // Store the actor_handle and shutdown_receiver for use by register()
+        // Store the actor_handle for use by register()
         {
             let mut handle_guard = self.actor_handle.lock().unwrap();
             *handle_guard = Some(actor_handle);
         }
+        // Only store shutdown_receiver if not already set during setup_host_functions_composite()
+        // The early receiver from setup_host_functions_composite() is what register() needs
         {
             let mut shutdown_guard = self.shutdown_receiver.lock().unwrap();
-            *shutdown_guard = Some(shutdown_receiver);
+            if shutdown_guard.is_none() {
+                *shutdown_guard = Some(shutdown_receiver);
+            }
+            // If already set, drop the one passed here - the controller will notify both
         }
 
-        let shutdown_receiver_arc = self.shutdown_receiver.clone();
         let registered_notify = self.registered_notify.clone();
 
-        // Wait for either:
-        // 1. register() takes the receiver (notified via registered_notify)
-        // 2. Shutdown signal received (if register() was never called)
+        // Wait for register() to be called (or shutdown if never called).
+        // The shutdown receiver is left in the mutex for register() to take.
+        // If the actor never calls register(), the receiver stays in the mutex
+        // and will be dropped when the handler is dropped.
         Box::pin(async move {
-            tokio::select! {
-                // Wait for register() to take the receiver
-                _ = registered_notify.notified() => {
-                    info!("Message server handler: registered, setup complete");
-                }
-                // Wait for shutdown if receiver wasn't taken
-                _ = async {
-                    let receiver = {
-                        let mut guard = shutdown_receiver_arc.lock().unwrap();
-                        guard.take()
-                    };
-                    if let Some(rx) = receiver {
-                        rx.wait_for_shutdown().await;
-                        info!("Message server handler received shutdown signal");
-                    }
-                } => {}
-            }
+            registered_notify.notified().await;
+            info!("Message server handler: registered, setup complete");
             Ok(())
         })
     }
@@ -623,6 +613,16 @@ impl Handler for MessageServerHandler
             self.actor_id = Some(actor_id.clone());
         }
 
+        // Get a shutdown receiver early - this is needed by register() which may be
+        // called before the async setup() runs. This fixes the race condition where
+        // an actor calls register() immediately in init before setup() has stored
+        // the shutdown_receiver.
+        if let Some(shutdown_receiver) = ctx.subscribe_shutdown() {
+            let mut guard = self.shutdown_receiver.lock().unwrap();
+            *guard = Some(shutdown_receiver);
+            info!("Message server handler: got shutdown receiver during host function setup");
+        }
+
         // Check if already satisfied
         if ctx.is_satisfied("theater:simple/message-server-host") {
             info!("theater:simple/message-server-host already satisfied, skipping");

crates/theater/src/actor/runtime.rs

@@ -290,8 +290,11 @@ impl ActorRuntime {
             Some(Arc::new(RecordingInterceptor::new(chain.clone())))
         };
 
-        // Create a closure that sets up all handler host functions
-        let mut handler_ctx = HandlerContext::new();
+        // Create shutdown controller early so handlers can subscribe during setup
+        let mut shutdown_controller = ShutdownController::new();
+
+        // Create handler context with shutdown controller access
+        let mut handler_ctx = HandlerContext::with_shutdown_controller(shutdown_controller.clone());
         handler_ctx.actor_id = Some(id.clone());
         let handlers_for_setup = &mut handlers;
 
@@ -480,9 +483,8 @@ impl ActorRuntime {
             *instance_guard = Some(actor_instance);
         }
 
-        // Set up handlers
+        // Set up handlers - use the shutdown_controller created earlier
         let mut handler_tasks: Vec<JoinHandle<()>> = vec![];
-        let mut shutdown_controller = ShutdownController::new();
         let handler_actor_handle = actor_handle.clone();
         debug!("Setting up {} handlers", handlers.len());
         for mut handler in handlers {

crates/theater/src/handler/mod.rs

@@ -4,7 +4,7 @@ use crate::chain::ChainEvent;
 use crate::id::TheaterId;
 use crate::pack_bridge::{HostLinkerBuilder, LinkerError, PackInstance, TypeHash};
 use crate::config::actor_manifest::HandlerConfig;
-use crate::shutdown::ShutdownReceiver;
+use crate::shutdown::{ShutdownController, ShutdownReceiver};
 use anyhow::Result;
 use std::collections::HashSet;
 use std::future::Future;
@@ -17,22 +17,46 @@ use tokio::sync::RwLock;
 pub type SharedActorInstance = Arc<RwLock<Option<PackInstance>>>;
 
 /// Context passed to handlers during setup, tracking which imports are already satisfied
-#[derive(Debug, Clone, Default)]
+/// and providing access to the shutdown controller for handlers that need it.
+#[derive(Debug, Clone)]
 pub struct HandlerContext {
     /// Set of imports that have already been registered by other handlers
     pub satisfied_imports: HashSet<String>,
     /// The actor ID for the actor being set up
     pub actor_id: Option<TheaterId>,
+    /// Shutdown controller - handlers can subscribe to get shutdown signals
+    pub shutdown_controller: Option<ShutdownController>,
+}
+
+impl Default for HandlerContext {
+    fn default() -> Self {
+        Self::new()
+    }
 }
 
 impl HandlerContext {
     pub fn new() -> Self {
         Self {
             satisfied_imports: HashSet::new(),
             actor_id: None,
+            shutdown_controller: None,
         }
     }
 
+    /// Create a new context with a shutdown controller
+    pub fn with_shutdown_controller(shutdown_controller: ShutdownController) -> Self {
+        Self {
+            satisfied_imports: HashSet::new(),
+            actor_id: None,
+            shutdown_controller: Some(shutdown_controller),
+        }
+    }
+
+    /// Get a shutdown receiver from the controller, if available
+    pub fn subscribe_shutdown(&mut self) -> Option<ShutdownReceiver> {
+        self.shutdown_controller.as_mut().map(|c| c.subscribe())
+    }
+
     /// Check if an import is already satisfied
     pub fn is_satisfied(&self, import: &str) -> bool {
         self.satisfied_imports.contains(import)
@@ -155,6 +179,14 @@ impl HandlerRegistry {
 /// External handler crates can implement this trait and register their handlers
 /// with the Theater runtime without depending on the concrete `Handler` enum.
 ///
+/// ## Handler Lifecycle
+///
+/// 1. `create_instance()` - Clone/create handler instance with optional config
+/// 2. `setup_host_functions_composite()` - Register host functions (sync, during instantiation)
+///    - HandlerContext provides shutdown_controller for handlers that need early access
+/// 3. `init()` - Synchronous critical initialization (called before actor can receive calls)
+/// 4. `run()` - Async runtime loop (spawned as background task)
+///
 /// ## Composite Migration
 ///
 /// For handlers migrating to Composite's Graph ABI runtime, implement:
@@ -169,21 +201,61 @@ pub trait Handler: Send + Sync + 'static {
     /// with that config. Otherwise, clones the current instance.
     fn create_instance(&self, config: Option<&HandlerConfig>) -> Box<dyn Handler>;
 
-    /// Initialize the handler with actor references.
+    /// Synchronous initialization called BEFORE the actor can receive any calls.
     ///
-    /// This is called once when the actor starts. Handlers should store the provided
-    /// handles for later use but NOT start event loops here. Event loops should be
-    /// triggered by explicit actor calls (e.g., `enable-input()` for terminal).
+    /// This is the place for critical setup that must complete before host functions
+    /// can be used. For example, storing actor handles that host functions depend on.
+    ///
+    /// Note: If a handler needs a ShutdownReceiver for host functions, it should
+    /// subscribe via `ctx.subscribe_shutdown()` during `setup_host_functions_composite()`.
+    ///
+    /// This runs synchronously - do NOT do any async work here.
+    /// Default implementation does nothing.
+    fn init(
+        &mut self,
+        _actor_handle: ActorHandle,
+        _actor_instance: SharedActorInstance,
+    ) {
+        // Default: no-op
+    }
+
+    /// Async runtime loop that runs for the handler's lifetime.
+    ///
+    /// This is spawned as a background task AFTER init() completes and the actor
+    /// is ready to receive calls. Use this for event loops, message consumption,
+    /// or any long-running async operations.
     ///
     /// The `event_rx` parameter receives chain events as they're recorded. Most handlers
     /// can ignore this, but ReplayHandler uses it for streaming hash verification.
+    ///
+    /// Default implementation just waits for shutdown.
+    fn run(
+        &mut self,
+        shutdown_receiver: ShutdownReceiver,
+        _event_rx: broadcast::Receiver<ChainEvent>,
+    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send>> {
+        Box::pin(async move {
+            shutdown_receiver.wait_for_shutdown().await;
+            Ok(())
+        })
+    }
+
+    /// Initialize and run the handler.
+    ///
+    /// The runtime calls init() synchronously, then spawns run() as a background task.
+    /// Most handlers should override init() and/or run() rather than this method.
     fn setup(
         &mut self,
         actor_handle: ActorHandle,
         actor_instance: SharedActorInstance,
         shutdown_receiver: ShutdownReceiver,
         event_rx: broadcast::Receiver<ChainEvent>,
-    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send>>;
+    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send>> {
+        // Call init synchronously first
+        self.init(actor_handle, actor_instance);
+        // Then return the run future
+        self.run(shutdown_receiver, event_rx)
+    }
 
     /// Set up host functions for this handler (Composite Graph ABI runtime).
     ///

crates/theater/src/shutdown.rs

@@ -1,3 +1,4 @@
+use std::sync::{Arc, Mutex};
 use std::time::Duration;
 use tokio::sync::oneshot::{Receiver, Sender};
 use tracing::debug;
@@ -34,31 +35,45 @@ pub enum ShutdownType {
     Force,
 }
 
-/// Controller that can broadcast shutdown signals to multiple receivers
+/// Controller that can broadcast shutdown signals to multiple receivers.
+/// This type is Clone-able and all clones share the same subscriber list.
+#[derive(Clone)]
 pub struct ShutdownController {
-    subscribers: Vec<Sender<ShutdownSignal>>,
+    subscribers: Arc<Mutex<Vec<Sender<ShutdownSignal>>>>,
+}
+
+impl std::fmt::Debug for ShutdownController {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ShutdownController")
+            .field("subscriber_count", &self.subscribers.lock().unwrap().len())
+            .finish()
+    }
 }
 
 impl ShutdownController {
     /// Create a new ShutdownController and a ShutdownReceiver
     pub fn new() -> Self {
         Self {
-            subscribers: Vec::new(),
+            subscribers: Arc::new(Mutex::new(Vec::new())),
         }
     }
 
     /// Get a new receiver for this controller
     pub fn subscribe(&mut self) -> ShutdownReceiver {
         let (sender, receiver) = tokio::sync::oneshot::channel();
-        self.subscribers.push(sender);
+        self.subscribers.lock().unwrap().push(sender);
         ShutdownReceiver { receiver }
     }
 
     /// Signal all receivers to shutdown
     pub async fn signal_shutdown(self, shutdown_type: ShutdownType) {
         debug!("Signaling shutdown to all subscribers");
+        let subscribers = {
+            let mut guard = self.subscribers.lock().unwrap();
+            std::mem::take(&mut *guard)
+        };
         let mut receivers = Vec::new();
-        for sender in self.subscribers {
+        for sender in subscribers {
             let (responder, receiver) = tokio::sync::oneshot::channel();
             receivers.push(receiver);
             match sender.send(ShutdownSignal {