398 lines
15 KiB
Rust
398 lines
15 KiB
Rust
//! 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());
|
|
}
|
|
}
|