huskies: merge 1163 story Replace perm_rx lock-as-presence-signal with a permission router
This commit is contained in:
Huskies Agent
@@ -0,0 +1,397 @@
|
||||
//! 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<PermissionForward>,
|
||||
}
|
||||
|
||||
/// 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<Vec<ResponderSlot>>,
|
||||
}
|
||||
|
||||
impl ResponderRegistry {
|
||||
/// Create an empty registry.
|
||||
pub fn new() -> Arc<Self> {
|
||||
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<Self>) -> (ResponderGuard, mpsc::Receiver<PermissionForward>) {
|
||||
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<mpsc::Sender<PermissionForward>>;
|
||||
}
|
||||
|
||||
impl CloneSenders for Vec<ResponderSlot> {
|
||||
fn clone_senders(&self) -> Vec<mpsc::Sender<PermissionForward>> {
|
||||
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<ResponderRegistry>,
|
||||
}
|
||||
|
||||
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<PermissionForward>,
|
||||
registry: Arc<ResponderRegistry>,
|
||||
) -> 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<PendingInner>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct PendingInner {
|
||||
by_request_id: HashMap<String, oneshot::Sender<PermissionDecision>>,
|
||||
by_location: HashMap<String, VecDeque<String>>,
|
||||
}
|
||||
|
||||
impl PendingPermReplies {
|
||||
/// Create an empty tracker.
|
||||
pub fn new() -> Arc<Self> {
|
||||
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<String>,
|
||||
request_id: impl Into<String>,
|
||||
tx: oneshot::Sender<PermissionDecision>,
|
||||
) {
|
||||
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<oneshot::Sender<PermissionDecision>> {
|
||||
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<oneshot::Sender<PermissionDecision>> {
|
||||
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<PermissionDecision>) {
|
||||
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());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user