//! Permission router — replaces holding `perm_rx`'s mutex as a presence //! signal for whether an interactive session is listening for MCP //! `prompt_permission` requests. //! //! Before this module, `Services` held the MCP-side receiver behind a //! `tokio::sync::Mutex`. Whichever transport task (Matrix bot, sled uplink, //! WS chat session) locked it became "the" active listener for the process, //! and `tool_prompt_permission` used `try_lock()` as a cheap presence check. //! That conflated two unrelated concerns — "is anyone listening" and "who //! drains the queue" — behind a single lock, and made it impossible for more //! than one responder (e.g. a Matrix bot and several open browser tabs) to be //! registered at once without one of them blocking forever on `.lock().await`. //! //! The replacement: [`spawn_permission_router`] is the sole, permanent owner //! of the MCP-side receiver. Responders call [`ResponderRegistry::register`] //! to obtain a private bounded channel and an RAII [`ResponderGuard`] that //! unregisters them on drop. The router only ever `try_send`s into a //! responder's channel — it never awaits foreign I/O, so a wedged responder //! (registered but not draining) can stall neither the router nor any other //! part of the server. use crate::http::context::{PermissionDecision, PermissionForward}; use std::collections::{HashMap, VecDeque}; use std::sync::Arc; use std::sync::Mutex as StdMutex; use std::sync::atomic::{AtomicU64, Ordering}; use tokio::sync::{Mutex as TokioMutex, mpsc, oneshot}; /// Capacity of a single responder's private inbox channel. /// /// Bounded so a wedged responder's queue fills up (and the router falls back /// to the next responder, or fail-closes) instead of growing without limit. pub const RESPONDER_CHANNEL_CAPACITY: usize = 16; struct ResponderSlot { id: u64, tx: mpsc::Sender, } /// Tracks which tasks are currently registered to receive forwarded /// permission requests. /// /// Registering does not grant exclusive ownership: multiple responders (the /// Matrix bot, a sled uplink, several WS browser tabs) may be registered at /// once. Each incoming request is dispatched to the first registered /// responder whose private channel accepts it via `try_send`. pub struct ResponderRegistry { next_id: AtomicU64, slots: StdMutex>, } impl ResponderRegistry { /// Create an empty registry. pub fn new() -> Arc { Arc::new(Self { next_id: AtomicU64::new(0), slots: StdMutex::new(Vec::new()), }) } /// `true` when no responder is currently registered. /// /// `tool_prompt_permission` uses this in place of the old /// `perm_rx.try_lock().is_ok()` check to auto-deny immediately instead of /// forwarding a request nobody is listening for. pub fn is_empty(&self) -> bool { self.slots.lock().unwrap().is_empty() } /// Register a new responder. Returns a private receiver for forwarded /// requests and an RAII guard — dropping the guard unregisters the /// responder, so a disconnected transport stops being counted as present. pub fn register(self: &Arc) -> (ResponderGuard, mpsc::Receiver) { let (tx, rx) = mpsc::channel(RESPONDER_CHANNEL_CAPACITY); let id = self.next_id.fetch_add(1, Ordering::Relaxed); self.slots.lock().unwrap().push(ResponderSlot { id, tx }); ( ResponderGuard { id, registry: Arc::clone(self), }, rx, ) } /// Dispatch one forwarded request to the first responder that accepts it. /// /// Never blocks: uses `try_send` against each registered responder in /// turn. If every responder's channel is full (or none are registered), /// the request is fail-closed with [`PermissionDecision::Deny`] rather /// than dropped silently. pub fn dispatch(&self, forward: PermissionForward) { let slots = self.slots.lock().unwrap().clone_senders(); let mut remaining = forward; for tx in &slots { match tx.try_send(remaining) { Ok(()) => return, Err(mpsc::error::TrySendError::Full(fwd)) | Err(mpsc::error::TrySendError::Closed(fwd)) => { remaining = fwd; } } } let _ = remaining.response_tx.send(PermissionDecision::Deny); } fn unregister(&self, id: u64) { self.slots.lock().unwrap().retain(|s| s.id != id); } } /// Helper so `dispatch` can iterate over a snapshot of senders without /// holding the `std::sync::Mutex` across the (non-blocking, but still /// non-trivial) `try_send` calls. trait CloneSenders { fn clone_senders(&self) -> Vec>; } impl CloneSenders for Vec { fn clone_senders(&self) -> Vec> { self.iter().map(|s| s.tx.clone()).collect() } } /// RAII guard returned by [`ResponderRegistry::register`]. /// /// Unregisters the responder on drop so a disconnected Matrix bot, sled /// uplink, or WS session stops being counted as an active listener. pub struct ResponderGuard { id: u64, registry: Arc, } impl Drop for ResponderGuard { fn drop(&mut self) { self.registry.unregister(self.id); } } /// Spawn the router task: the sole, permanent owner of the MCP-side receiver. /// /// Never awaits responder I/O — only [`ResponderRegistry::dispatch`], which /// is itself non-blocking. Exits when `perm_rx` closes (server shutdown). pub fn spawn_permission_router( mut perm_rx: mpsc::UnboundedReceiver, registry: Arc, ) -> tokio::task::JoinHandle<()> { tokio::spawn(async move { while let Some(forward) = perm_rx.recv().await { registry.dispatch(forward); } }) } // ── Pending replies ────────────────────────────────────────────────────── /// Tracks permission requests awaiting a plain-language (yes/no) reply from a /// chat transport, keyed by `request_id` rather than by conversation location /// (room / sender / channel). /// /// Keying by location alone meant a second concurrent request for the same /// room silently dropped the first request's oneshot sender when it /// overwrote the map entry. Keying by `request_id` fixes that; a per-location /// FIFO queue lets [`PendingPermReplies::resolve_oldest`] map a bare "yes"/"no" /// reply (which carries no request_id) back to the oldest still-pending /// request for that location. pub struct PendingPermReplies { inner: TokioMutex, } #[derive(Default)] struct PendingInner { by_request_id: HashMap>, by_location: HashMap>, } impl PendingPermReplies { /// Create an empty tracker. pub fn new() -> Arc { Arc::new(Self { inner: TokioMutex::new(PendingInner::default()), }) } /// Register a pending reply for `request_id`, queued under `location`. pub async fn insert( &self, location: impl Into, request_id: impl Into, tx: oneshot::Sender, ) { let request_id = request_id.into(); let mut inner = self.inner.lock().await; inner .by_location .entry(location.into()) .or_default() .push_back(request_id.clone()); inner.by_request_id.insert(request_id, tx); } /// Resolve the oldest pending request queued for `location`, removing it /// from both the location queue and the request_id map. Used when a chat /// transport receives a bare yes/no reply with no request_id attached. pub async fn resolve_oldest( &self, location: &str, ) -> Option> { let mut inner = self.inner.lock().await; loop { let queue = inner.by_location.get_mut(location)?; let request_id = queue.pop_front()?; if queue.is_empty() { inner.by_location.remove(location); } if let Some(tx) = inner.by_request_id.remove(&request_id) { return Some(tx); } // request_id was already removed (e.g. by a timeout) — try the // next queued entry for this location instead of returning None. } } /// Remove a specific pending reply by `request_id`, also dequeuing it from /// `location`'s FIFO so a later bare reply doesn't try to resolve an /// already-decided (e.g. timed-out) request. Used by responder timeout /// tasks. pub async fn remove_by_request_id( &self, location: &str, request_id: &str, ) -> Option> { let mut inner = self.inner.lock().await; if let Some(queue) = inner.by_location.get_mut(location) { queue.retain(|id| id != request_id); if queue.is_empty() { inner.by_location.remove(location); } } inner.by_request_id.remove(request_id) } } // ── Tests ───────────────────────────────────────────────────────────────── #[cfg(test)] mod tests { use super::*; use serde_json::json; fn make_forward( request_id: &str, ) -> (PermissionForward, oneshot::Receiver) { let (tx, rx) = oneshot::channel(); ( PermissionForward { request_id: request_id.to_string(), tool_name: "Bash".to_string(), tool_input: json!({"command": "echo hi"}), response_tx: tx, }, rx, ) } #[tokio::test] async fn registry_starts_empty() { let registry = ResponderRegistry::new(); assert!(registry.is_empty()); } #[tokio::test] async fn register_makes_registry_non_empty_until_guard_dropped() { let registry = ResponderRegistry::new(); let (guard, _rx) = registry.register(); assert!( !registry.is_empty(), "registry must be non-empty while registered" ); drop(guard); assert!( registry.is_empty(), "dropping the guard must unregister the responder" ); } #[tokio::test] async fn dispatch_delivers_to_registered_responder() { let registry = ResponderRegistry::new(); let (_guard, mut rx) = registry.register(); let (fwd, _response_rx) = make_forward("req-1"); registry.dispatch(fwd); let received = rx.recv().await.expect("responder must receive forward"); assert_eq!(received.request_id, "req-1"); } #[tokio::test] async fn dispatch_denies_when_no_responder_registered() { let registry = ResponderRegistry::new(); let (fwd, response_rx) = make_forward("req-2"); registry.dispatch(fwd); let decision = response_rx.await.expect("must receive a decision"); assert_eq!( decision, PermissionDecision::Deny, "no registered responder must fail-closed" ); } #[tokio::test] async fn wedged_responder_falls_through_to_deny_without_blocking() { let registry = ResponderRegistry::new(); let (_guard, _rx) = registry.register(); // Fill the responder's bounded channel without draining it, so it // looks "wedged" from the router's point of view. for i in 0..RESPONDER_CHANNEL_CAPACITY { let (fwd, _rx) = make_forward(&format!("filler-{i}")); registry.dispatch(fwd); } let (fwd, response_rx) = make_forward("req-overflow"); // Must return immediately — no await on responder I/O. registry.dispatch(fwd); let decision = tokio::time::timeout(std::time::Duration::from_secs(2), response_rx) .await .expect("dispatch must not stall when the responder is wedged") .expect("must receive a decision"); assert_eq!(decision, PermissionDecision::Deny); } #[tokio::test] async fn pending_replies_two_concurrent_requests_same_location_both_resolve() { let pending = PendingPermReplies::new(); let (tx1, rx1) = oneshot::channel(); let (tx2, rx2) = oneshot::channel(); pending.insert("room-1", "req-a", tx1).await; pending.insert("room-1", "req-b", tx2).await; // Regression: previously keying by location alone meant the second // insert overwrote (and dropped) the first request's sender. let first = pending .resolve_oldest("room-1") .await .expect("first pending reply must still be present"); let _ = first.send(PermissionDecision::Approve); assert_eq!(rx1.await.unwrap(), PermissionDecision::Approve); let second = pending .resolve_oldest("room-1") .await .expect("second pending reply must still be present"); let _ = second.send(PermissionDecision::Deny); assert_eq!(rx2.await.unwrap(), PermissionDecision::Deny); } #[tokio::test] async fn remove_by_request_id_prevents_later_resolution() { let pending = PendingPermReplies::new(); let (tx, _rx) = oneshot::channel(); pending.insert("room-1", "req-timeout", tx).await; let removed = pending.remove_by_request_id("room-1", "req-timeout").await; assert!(removed.is_some(), "must return the removed sender"); let resolved = pending.resolve_oldest("room-1").await; assert!( resolved.is_none(), "a request removed by timeout must not be resolvable afterwards" ); } #[tokio::test] async fn resolve_oldest_skips_already_removed_entries() { let pending = PendingPermReplies::new(); let (tx1, _rx1) = oneshot::channel(); let (tx2, rx2) = oneshot::channel(); pending.insert("room-1", "req-a", tx1).await; pending.insert("room-1", "req-b", tx2).await; // req-a times out (removed) before the user replies. pending.remove_by_request_id("room-1", "req-a").await; // A bare "yes" reply must resolve req-b, not silently return None. let tx = pending .resolve_oldest("room-1") .await .expect("must skip the timed-out entry and resolve the next one"); let _ = tx.send(PermissionDecision::Approve); assert_eq!(rx2.await.unwrap(), PermissionDecision::Approve); } #[tokio::test] async fn resolve_oldest_returns_none_for_unknown_location() { let pending = PendingPermReplies::new(); assert!(pending.resolve_oldest("no-such-room").await.is_none()); } }