fix(maestro): remove pending-RPC entries on caller cancellation (NAP-259)

`BrokerHandle::resolve_attach` (daemon) and `Registry::create_session` /
`kill_session` (registry) insert a pending oneshot into a map, then await
the reply. If the awaiting future is cancelled — e.g. an axum handler
dropped because the browser closed the tab or aborted the XHR — the entry
was never removed, leaking one map slot per cancelled request and growing
unboundedly under UI-driven retry traffic.

Introduce a small RAII `PendingGuard` that removes the entry on drop, so
cleanup happens on every exit path including cancellation. Because a sync
`Drop` can't lock a tokio mutex, move the pending maps to
`std::sync::Mutex` (only ever locked for a tiny insert/remove, never held
across an await): broker `pending`, and the registry's `pending_create` /
`pending_kill` lifted out of the async-locked `State`. Removal is
idempotent, so the happy path (reply already taken by `complete_*`) and
the guard's drop don't conflict; `await_rpc` no longer needs its manual
timeout cleanup.

Tested: cancelling an in-flight `resolve_attach` (broker) and
`kill_session` (registry) leaves the pending map empty.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

crates/maestro-daemon/src/broker.rs

@@ -6,13 +6,33 @@
 //! send and the reply correlation live here so the handler stays thin.
 
 use std::collections::HashMap;
-use std::sync::Arc;
+use std::hash::Hash;
+use std::sync::{Arc, Mutex};
 use std::time::Duration;
 
 use maestro_protocol::{AttachCandidates, DaemonOutbound, ReplyTo, RpcResult};
-use tokio::sync::{mpsc, oneshot, Mutex};
+use tokio::sync::{mpsc, oneshot};
 use uuid::Uuid;
 
+/// RAII cleanup for a pending-RPC entry. Removing the entry on drop guarantees
+/// it is gone on *every* exit path of the awaiting call — success, error,
+/// timeout, and crucially future cancellation (the axum handler is dropped when
+/// the client hangs up) — so the map can't grow without bound (NAP-259).
+///
+/// The map is a `std::sync::Mutex` (never held across an await) precisely so the
+/// sync `Drop` can lock it. Removal is idempotent: on the happy path the reply
+/// handler has already taken the entry and this is a harmless no-op.
+struct PendingGuard<'a, K: Eq + Hash, V> {
+    map: &'a Mutex<HashMap<K, V>>,
+    key: K,
+}
+
+impl<K: Eq + Hash, V> Drop for PendingGuard<'_, K, V> {
+    fn drop(&mut self) {
+        self.map.lock().unwrap().remove(&self.key);
+    }
+}
+
 /// How long to wait for a `ResolveAttachResult` after the request goes out.
 /// Matches the registry's RPC timeout — a healthy registry replies in ms.
 pub const RESOLVE_TIMEOUT: Duration = Duration::from_secs(10);
@@ -55,7 +75,13 @@ impl BrokerHandle {
     pub async fn resolve_attach(&self, session_id: Uuid) -> Result<AttachCandidates, BrokerError> {
         let reply_to = Uuid::new_v4();
         let (tx, rx) = oneshot::channel();
-        self.pending.lock().await.insert(reply_to, tx);
+        self.pending.lock().unwrap().insert(reply_to, tx);
+        // From here on, every exit path (including cancellation of this future)
+        // removes the pending entry when `_guard` drops (NAP-259).
+        let _guard = PendingGuard {
+            map: &*self.pending,
+            key: reply_to,
+        };
 
         if self
             .out_tx
@@ -66,7 +92,6 @@ impl BrokerHandle {
             .await
             .is_err()
         {
-            self.pending.lock().await.remove(&reply_to);
             return Err(BrokerError::NoControlLink);
         }
 
@@ -76,10 +101,7 @@ impl BrokerHandle {
                 Err(BrokerError::Registry(error.code, error.message))
             }
             Ok(Err(_)) => Err(BrokerError::ReplyClosed),
-            Err(_) => {
-                self.pending.lock().await.remove(&reply_to);
-                Err(BrokerError::Timeout(RESOLVE_TIMEOUT))
-            }
+            Err(_) => Err(BrokerError::Timeout(RESOLVE_TIMEOUT)),
         }
     }
 
@@ -91,8 +113,8 @@ impl BrokerHandle {
         reply_to: ReplyTo,
         result: RpcResult<AttachCandidates>,
     ) -> bool {
-        let mut pending = self.pending.lock().await;
-        match pending.remove(&reply_to) {
+        let entry = self.pending.lock().unwrap().remove(&reply_to);
+        match entry {
             Some(tx) => tx.send(result).is_ok(),
             None => {
                 tracing::debug!(?reply_to, "no waiter for ResolveAttachResult");
@@ -101,3 +123,39 @@ impl BrokerHandle {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Cancelling (dropping) an in-flight `resolve_attach` must not leak its
+    /// pending-reply entry — otherwise a browser hanging up mid-attach grows
+    /// the map without bound (NAP-259).
+    #[tokio::test]
+    async fn cancelled_resolve_attach_removes_pending_entry() {
+        // Keep the receiver alive so `out_tx.send()` succeeds and the call
+        // parks on the reply oneshot rather than failing fast.
+        let (handle, _out_rx) = BrokerHandle::new();
+
+        {
+            let fut = handle.resolve_attach(Uuid::new_v4());
+            tokio::pin!(fut);
+            tokio::select! {
+                biased;
+                _ = &mut fut => panic!("resolve must not complete without a reply"),
+                _ = tokio::time::sleep(Duration::from_millis(20)) => {}
+            }
+            assert_eq!(
+                handle.pending.lock().unwrap().len(),
+                1,
+                "entry should be in-flight before cancellation"
+            );
+        }
+
+        assert_eq!(
+            handle.pending.lock().unwrap().len(),
+            0,
+            "pending entry must be removed on cancellation (NAP-259)"
+        );
+    }
+}

crates/maestro-registry/src/state.rs

@@ -10,7 +10,8 @@
 //! handful of daemons and these critical sections are tiny.
 
 use std::collections::{BTreeMap, HashMap};
-use std::sync::Arc;
+use std::hash::Hash;
+use std::sync::{Arc, Mutex as StdMutex};
 use std::time::Duration;
 
 use maestro_protocol::{
@@ -22,6 +23,26 @@ use uuid::Uuid;
 
 use crate::jwt::issue_pair_attach_jwt;
 
+/// RAII cleanup for a pending-RPC entry. Removing the entry on drop guarantees
+/// it is gone on *every* exit path of the awaiting call — success, error,
+/// timeout, and crucially future cancellation (an axum handler is dropped when
+/// the HTTP client hangs up) — so the map can't grow without bound (NAP-259).
+///
+/// The pending maps live behind a `std::sync::Mutex` (never held across an
+/// await) precisely so this sync `Drop` can lock them. Removal is idempotent:
+/// on the happy path the reply handler has already taken the entry, making the
+/// drop a harmless no-op.
+struct PendingGuard<'a, K: Eq + Hash, V> {
+    map: &'a StdMutex<HashMap<K, V>>,
+    key: K,
+}
+
+impl<K: Eq + Hash, V> Drop for PendingGuard<'_, K, V> {
+    fn drop(&mut self) {
+        self.map.lock().unwrap().remove(&self.key);
+    }
+}
+
 /// Capacity of the registry-wide broadcast channel that feeds
 /// `WS /api/events` subscribers. Subscribers that fall this far behind get a
 /// `Lagged` event and are expected to resync via `GET /api/sessions`.
@@ -80,6 +101,11 @@ impl From<DaemonView> for DaemonInfo {
 pub struct Registry {
     inner: Arc<Mutex<State>>,
     events: broadcast::Sender<RegistryEvent>,
+    /// Outstanding R->D RPC waiters. Held behind a `std::sync::Mutex` (separate
+    /// from `State`'s async lock) so a [`PendingGuard`] can clean them up on
+    /// drop — including when the awaiting HTTP handler is cancelled (NAP-259).
+    pending_create: Arc<StdMutex<HashMap<ReplyTo, oneshot::Sender<RpcResult<Session>>>>>,
+    pending_kill: Arc<StdMutex<HashMap<ReplyTo, oneshot::Sender<RpcResult<()>>>>>,
 }
 
 #[derive(Debug, Default)]
@@ -87,8 +113,6 @@ struct State {
     daemons: HashMap<String, DaemonEntry>,
     /// `(daemon_id, session_id) -> Session`.
     sessions: BTreeMap<(String, Uuid), Session>,
-    pending_create: HashMap<ReplyTo, oneshot::Sender<RpcResult<Session>>>,
-    pending_kill: HashMap<ReplyTo, oneshot::Sender<RpcResult<()>>>,
 }
 
 #[derive(Debug)]
@@ -126,6 +150,8 @@ impl Registry {
         Self {
             inner: Arc::new(Mutex::new(State::default())),
             events,
+            pending_create: Arc::new(StdMutex::new(HashMap::new())),
+            pending_kill: Arc::new(StdMutex::new(HashMap::new())),
         }
     }
 
@@ -420,8 +446,8 @@ impl Registry {
 
     /// Resolve a CreateSession RPC reply. Returns true if a waiter was found.
     pub async fn complete_create(&self, reply_to: ReplyTo, result: RpcResult<Session>) -> bool {
-        let mut state = self.inner.lock().await;
-        match state.pending_create.remove(&reply_to) {
+        let entry = self.pending_create.lock().unwrap().remove(&reply_to);
+        match entry {
             Some(tx) => tx.send(result).is_ok(),
             None => false,
         }
@@ -429,8 +455,8 @@ impl Registry {
 
     /// Resolve a KillSession RPC reply. Returns true if a waiter was found.
     pub async fn complete_kill(&self, reply_to: ReplyTo, result: RpcResult<()>) -> bool {
-        let mut state = self.inner.lock().await;
-        match state.pending_kill.remove(&reply_to) {
+        let entry = self.pending_kill.lock().unwrap().remove(&reply_to);
+        match entry {
             Some(tx) => tx.send(result).is_ok(),
             None => false,
         }
@@ -450,8 +476,7 @@ impl Registry {
         let reply_to = Uuid::new_v4();
         let (tx, rx) = oneshot::channel();
         let outbound = {
-            let mut state = self.inner.lock().await;
-            state.pending_create.insert(reply_to, tx);
+            let state = self.inner.lock().await;
             state
                 .daemons
                 .get(daemon_id)
@@ -459,16 +484,22 @@ impl Registry {
                 .outbound
                 .clone()
         };
-        let send_result = outbound
+        self.pending_create.lock().unwrap().insert(reply_to, tx);
+        // Every exit path past here — send failure, timeout, a successful reply,
+        // or cancellation of this future — removes the pending entry when
+        // `_guard` drops (NAP-259).
+        let _guard = PendingGuard {
+            map: &*self.pending_create,
+            key: reply_to,
+        };
+        if outbound
             .send(DaemonInbound::CreateSession { reply_to, request })
-            .await;
-        if send_result.is_err() {
-            self.inner.lock().await.pending_create.remove(&reply_to);
+            .await
+            .is_err()
+        {
             return Err(RegistryError::DaemonGone);
         }
-        let rpc = await_rpc(rx, reply_to, |s| &mut s.pending_create, &self.inner)
-            .await
-            .map(rpc_into_result)?;
+        let rpc = await_rpc(rx).await.map(rpc_into_result)?;
         // Eagerly mirror the session into the catalog on success — the daemon
         // also emits `SessionAdded` asynchronously, but the HTTP caller would
         // otherwise see an empty `GET /api/sessions` immediately after a
@@ -491,8 +522,7 @@ impl Registry {
         let reply_to = Uuid::new_v4();
         let (tx, rx) = oneshot::channel();
         let outbound = {
-            let mut state = self.inner.lock().await;
-            state.pending_kill.insert(reply_to, tx);
+            let state = self.inner.lock().await;
             state
                 .daemons
                 .get(daemon_id)
@@ -500,16 +530,21 @@ impl Registry {
                 .outbound
                 .clone()
         };
-        let send_result = outbound
+        self.pending_kill.lock().unwrap().insert(reply_to, tx);
+        // See `create_session`: `_guard` removes the pending entry on every exit
+        // path, including cancellation of this future (NAP-259).
+        let _guard = PendingGuard {
+            map: &*self.pending_kill,
+            key: reply_to,
+        };
+        if outbound
             .send(DaemonInbound::KillSession { reply_to, id })
-            .await;
-        if send_result.is_err() {
-            self.inner.lock().await.pending_kill.remove(&reply_to);
+            .await
+            .is_err()
+        {
             return Err(RegistryError::DaemonGone);
         }
-        let rpc = await_rpc(rx, reply_to, |s| &mut s.pending_kill, &self.inner)
-            .await
-            .map(rpc_into_result)?;
+        let rpc = await_rpc(rx).await.map(rpc_into_result)?;
         // Mirror of the create-path eager update: drop the session locally so
         // subsequent reads see consistent state. The daemon also emits
         // `SessionRemoved`; `remove_session` is idempotent.
@@ -567,22 +602,14 @@ impl State {
     }
 }
 
-/// Await `rx` with [`RPC_TIMEOUT`]; on timeout / channel close, also clean up
-/// the pending map so a late reply doesn't leak.
-async fn await_rpc<T>(
-    rx: oneshot::Receiver<RpcResult<T>>,
-    reply_to: ReplyTo,
-    pick: impl FnOnce(&mut State) -> &mut HashMap<ReplyTo, oneshot::Sender<RpcResult<T>>>,
-    inner: &Arc<Mutex<State>>,
-) -> Result<RpcResult<T>, RegistryError> {
+/// Await `rx` with [`RPC_TIMEOUT`]. Cleanup of the pending map on every exit
+/// path (timeout, channel close, cancellation) is handled by the caller's
+/// [`PendingGuard`], so this only has to classify the outcome.
+async fn await_rpc<T>(rx: oneshot::Receiver<RpcResult<T>>) -> Result<RpcResult<T>, RegistryError> {
     match tokio::time::timeout(RPC_TIMEOUT, rx).await {
         Ok(Ok(result)) => Ok(result),
         Ok(Err(_)) => Err(RegistryError::ReplyChannelClosed),
-        Err(_) => {
-            let mut state = inner.lock().await;
-            pick(&mut state).remove(&reply_to);
-            Err(RegistryError::Timeout(RPC_TIMEOUT))
-        }
+        Err(_) => Err(RegistryError::Timeout(RPC_TIMEOUT)),
     }
 }
 
@@ -675,3 +702,49 @@ mod tests {
         assert!(reg.daemons().await.is_empty());
     }
 }
+
+#[cfg(test)]
+mod nap259_tests {
+    use super::*;
+
+    /// Cancelling an in-flight `kill_session` must not leak its pending-reply
+    /// entry — a client that hangs up mid-request would otherwise grow the map
+    /// without bound (NAP-259).
+    #[tokio::test]
+    async fn cancelled_kill_session_removes_pending_entry() {
+        let registry = Registry::new();
+        // Register a daemon with a live outbound receiver so the R->D send
+        // succeeds and the call parks on the reply oneshot.
+        let (tx, _rx) = mpsc::channel::<DaemonInbound>(8);
+        registry
+            .register(
+                "d1".into(),
+                "host".into(),
+                "0.1.0".into(),
+                "ws://127.0.0.1:1".into(),
+                tx,
+            )
+            .await;
+
+        {
+            let fut = registry.kill_session("d1", Uuid::new_v4());
+            tokio::pin!(fut);
+            tokio::select! {
+                biased;
+                _ = &mut fut => panic!("kill must not complete without a reply"),
+                _ = tokio::time::sleep(Duration::from_millis(20)) => {}
+            }
+            assert_eq!(
+                registry.pending_kill.lock().unwrap().len(),
+                1,
+                "entry should be in-flight before cancellation"
+            );
+        }
+
+        assert_eq!(
+            registry.pending_kill.lock().unwrap().len(),
+            0,
+            "pending entry must be removed on cancellation (NAP-259)"
+        );
+    }
+}