fix(bus): synchronous handler registration to de-flake round-trip tests

agent-mesh-bus/src/bus.rs

@@ -314,14 +314,15 @@ impl Bus {
         F: Fn(Vec<u8>) -> Fut + Send + Sync + 'static,
         Fut: Future<Output = Result<Vec<u8>>> + Send + 'static,
     {
-        let inbox = self.inbox.clone();
-        // Inbox::register_handler is async only because it takes a
-        // write lock; spawn it so the caller doesn't have to .await.
-        // The handler is in place by the time the next message is
-        // dispatched.
-        tokio::spawn(async move {
-            inbox.register_handler(topic, handler).await;
-        });
+        // Register synchronously so the handler is live the instant this
+        // returns. Previously this spawned the registration onto the
+        // runtime, which left a window — before the spawned task was
+        // polled — where a request dispatched to this topic found no
+        // handler and was silently dropped (`Inbox::dispatch_request`
+        // returns `Ok(None)`), timing out the asker. Direct-dial
+        // round-trip tests papered over that window with a fixed
+        // `sleep`; synchronous registration removes the race outright.
+        self.inbox.register_handler(topic, handler);
     }
 
     /// Publish a body to `peer_fp` on `topic`. Fire-and-forget — the
@@ -872,12 +873,9 @@ mod tests {
         bob_bus.handle_requests(topic.clone(), |body| async move {
             Ok(format!("echo: {}", String::from_utf8_lossy(&body)).into_bytes())
         });
-        // `handle_requests` registers on a spawned task; on the single-threaded
-        // runtime a couple of yields run it to completion (it only takes an
-        // uncontended lock) before we send — deterministic, no wall-clock sleep.
-        for _ in 0..4 {
-            tokio::task::yield_now().await;
-        }
+        // No yield/sleep needed: `handle_requests` registers the handler
+        // synchronously before it returns (see
+        // `handle_requests_registers_synchronously_no_spawn_race`).
 
         let reply = alice_bus
             .request(bob_fp, &topic, b"hi".to_vec(), Duration::from_secs(5))
@@ -889,6 +887,52 @@ mod tests {
         bob_bus.close().await.unwrap();
     }
 
+    /// Regression (#52 de-flake): `handle_requests` must register the
+    /// handler *before it returns*, with no spawn and no intervening
+    /// yield. It used to spawn the registration onto the runtime, so on a
+    /// `current_thread` runtime — where a spawned task is not polled until
+    /// the current task next yields — the handler was still absent the
+    /// instant `handle_requests` returned. A request dispatched into that
+    /// window found no handler and was silently dropped
+    /// (`Inbox::dispatch_request` -> `Ok(None)`), timing out the asker.
+    /// That is the exact flake the direct-dial round-trip tests used to
+    /// mask with a fixed `sleep(200ms)`.
+    ///
+    /// This test asserts the count with no yield between registration and
+    /// the check: it deterministically FAILS on the old spawn-based
+    /// implementation (count still 0) and PASSES on synchronous
+    /// registration (count 1). Run on the default single-threaded test
+    /// runtime so the "spawned task hasn't been polled yet" invariant
+    /// holds.
+    #[tokio::test]
+    async fn handle_requests_registers_synchronously_no_spawn_race() {
+        let user = UserKey::generate();
+        let bob = Arc::new(agent(&user, "bob"));
+        let bob_fp = bob.fingerprint();
+        let net = MeshNet::new();
+        let bob_bus =
+            Bus::bind_with_transport(bob, user.fingerprint(), Arc::new(net.transport_for(bob_fp)));
+
+        let topic = Topic::new(user.fingerprint(), "echo");
+        assert_eq!(
+            bob_bus.inbox.handler_count(),
+            0,
+            "no handler registered before handle_requests"
+        );
+
+        bob_bus.handle_requests(topic, |body| async move { Ok(body) });
+
+        // No sleep, no yield: on the old spawn-based code the spawned
+        // registration task has not run yet, so this would still read 0.
+        assert_eq!(
+            bob_bus.inbox.handler_count(),
+            1,
+            "handle_requests must register the handler before returning"
+        );
+
+        bob_bus.close().await.unwrap();
+    }
+
     /// A send to a fingerprint that isn't on the switchboard is `Unreachable`
     /// (mirrors the real transport's "peer not announced").
     #[tokio::test]

agent-mesh-bus/src/inbox.rs

@@ -101,7 +101,15 @@ pub struct Inbox {
     sequence: SequenceTracker,
     waiters: ReplyWaiter,
     subscriptions: RwLock<HashMap<String, broadcast::Sender<Vec<u8>>>>,
-    handlers: RwLock<HashMap<String, RequestHandler>>,
+    // Request handlers are guarded by a *synchronous* lock, not a
+    // `tokio::sync::RwLock`, so a handler can be registered without an
+    // `.await`. That lets `Bus::handle_requests` install the handler
+    // before it returns instead of on a spawned task — closing the
+    // registration race a directly-dialed request could otherwise lose
+    // (see `register_handler`). The guard is only ever held for a
+    // `get().cloned()` / `insert()` and never across an `.await`, so it
+    // cannot block the async runtime.
+    handlers: std::sync::RwLock<HashMap<String, RequestHandler>>,
 }
 
 impl Inbox {
@@ -116,7 +124,7 @@ impl Inbox {
             sequence: SequenceTracker::new(),
             waiters: ReplyWaiter::new(),
             subscriptions: RwLock::new(HashMap::new()),
-            handlers: RwLock::new(HashMap::new()),
+            handlers: std::sync::RwLock::new(HashMap::new()),
         }
     }
 
@@ -137,15 +145,31 @@ impl Inbox {
 
     /// Register a request handler for the given topic.
     ///
+    /// Synchronous by design: registration takes only the in-memory
+    /// `handlers` lock (never held across an `.await`), so the handler
+    /// is live the instant this returns. `Bus::handle_requests` relies
+    /// on that to avoid a spawn-and-race window where a freshly-dialed
+    /// request could arrive before the handler existed and be silently
+    /// dropped.
+    ///
     /// Re-registering replaces the previous handler for that topic.
-    pub async fn register_handler<F, Fut>(&self, topic: Topic, handler: F)
+    pub fn register_handler<F, Fut>(&self, topic: Topic, handler: F)
     where
         F: Fn(Vec<u8>) -> Fut + Send + Sync + 'static,
         Fut: Future<Output = Result<Vec<u8>>> + Send + 'static,
     {
         let key = topic.wire();
         let boxed: RequestHandler = Arc::new(move |body| Box::pin(handler(body)));
-        self.handlers.write().await.insert(key, boxed);
+        self.handlers
+            .write()
+            .expect("handlers lock poisoned")
+            .insert(key, boxed);
+    }
+
+    /// Number of registered request handlers (tests + diagnostics).
+    #[must_use]
+    pub fn handler_count(&self) -> usize {
+        self.handlers.read().expect("handlers lock poisoned").len()
     }
 
     /// Register an in-flight request waiter; returns the receiver
@@ -244,7 +268,7 @@ impl Inbox {
         body: Vec<u8>,
     ) -> Result<Option<OutgoingReply>> {
         let handler = {
-            let map = self.handlers.read().await;
+            let map = self.handlers.read().expect("handlers lock poisoned");
             map.get(&topic).cloned()
         };
         let Some(handler) = handler else {
@@ -396,11 +420,9 @@ mod tests {
         let topic = Topic::new(user.fingerprint(), "echo");
 
         let inbox = Inbox::new();
-        inbox
-            .register_handler(topic.clone(), |body| async move {
-                Ok([b"echo:".to_vec(), body].concat())
-            })
-            .await;
+        inbox.register_handler(topic.clone(), |body| async move {
+            Ok([b"echo:".to_vec(), body].concat())
+        });
 
         let req = BusMessage::Request {
             topic: topic.wire(),

agent-mesh-bus/tests/bus_roundtrip.rs

@@ -150,9 +150,9 @@ async fn request_reply_roundtrip_via_direct_dial_no_mdns() {
         Ok(format!("echo: {}", String::from_utf8_lossy(&body)).into_bytes())
     });
 
-    // Brief pause only for handler registration to settle — there is
-    // no discovery to wait on.
-    tokio::time::sleep(Duration::from_millis(200)).await;
+    // No pause needed: `handle_requests` registers the handler
+    // synchronously before returning, and there is no discovery to wait
+    // on (both buses bound quiet). The round-trip is fully deterministic.
 
     // The explicit dial route: bob's agent pubkey + his loopback addr.
     let bob_endpoint = PeerEndpoint::new(

agent-mesh-cli/tests/mcp_stdio.rs

@@ -152,8 +152,9 @@ async fn mcp_server_direct_addr_round_trips_to_live_responder() {
         let req: Value = serde_json::from_slice(&body).unwrap_or(Value::Null);
         Ok(serde_json::to_vec(&json!({ "echo": req["msg"] })).unwrap())
     });
-    // Only the handler registration needs a beat — there is no discovery.
-    tokio::time::sleep(Duration::from_millis(200)).await;
+    // No beat needed: `handle_requests` registers the handler
+    // synchronously before returning, and there is no discovery to wait
+    // on (the responder is dialed by explicit addr+pubkey).
 
     // 4. Round-trip via explicit addr+pubkey — no mesh_peers, no mDNS.
     let reply = client