huskies/server/src/service/permission_router.rs

//! 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, HashSet, 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)
    }
}

// ── Remembered permissions ───────────────────────────────────────────────

/// In-memory "don't ask again for this tool + target-pattern" store, scoped
/// per requesting agent session (story 1218).
///
/// Keyed by session id (see `http::mcp::session`) — for a per-story coding
/// agent this is the story id, stable across a `--resume` of the same story
/// since the worktree (and its `.mcp.json` header) doesn't change. Rules
/// remembered here are visible only within the session that recorded them,
/// so one agent's "don't ask again" never widens what another agent or
/// story is auto-approved for.
///
/// Deliberately **not** persisted to disk: this store lives only as long as
/// the server process does. Restarting the server clears every remembered
/// rule (agents will be prompted again once), which is an explicit tradeoff
/// — durability across restarts would require writing into each worktree's
/// own `.claude/settings.json`, which adds a second persistence path for a
/// case (server restarts mid-story) rare enough not to justify the added
/// complexity.
pub struct RememberedPermissions {
    inner: StdMutex<HashMap<String, HashSet<String>>>,
}

impl RememberedPermissions {
    /// Create an empty store.
    pub fn new() -> Arc<Self> {
        Arc::new(Self {
            inner: StdMutex::new(HashMap::new()),
        })
    }

    /// Record that `rule` is approved for `session_id` going forward.
    pub fn remember(&self, session_id: &str, rule: &str) {
        self.inner
            .lock()
            .unwrap()
            .entry(session_id.to_string())
            .or_default()
            .insert(rule.to_string());
    }

    /// `true` if `rule` was previously remembered for `session_id`.
    pub fn is_remembered(&self, session_id: &str, rule: &str) -> bool {
        self.inner
            .lock()
            .unwrap()
            .get(session_id)
            .is_some_and(|rules| rules.contains(rule))
    }
}

// ── 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());
    }

    // ── RememberedPermissions ───────────────────────────────────────

    #[test]
    fn remembered_permissions_starts_empty() {
        let store = RememberedPermissions::new();
        assert!(!store.is_remembered("1218", "Bash(git *)"));
    }

    #[test]
    fn remembered_permissions_recalls_after_remember() {
        let store = RememberedPermissions::new();
        store.remember("1218", "Bash(git *)");
        assert!(store.is_remembered("1218", "Bash(git *)"));
    }

    #[test]
    fn remembered_permissions_scoped_per_session() {
        let store = RememberedPermissions::new();
        store.remember("1218", "Bash(git *)");
        assert!(
            !store.is_remembered("1216", "Bash(git *)"),
            "a rule remembered for one story must not apply to another"
        );
    }

    #[test]
    fn remembered_permissions_scoped_per_rule() {
        let store = RememberedPermissions::new();
        store.remember("1218", "Bash(git *)");
        assert!(!store.is_remembered("1218", "Write"));
    }
}