402 lines
14 KiB
Rust
402 lines
14 KiB
Rust
|
//! Question router — mirrors `permission_router.rs`'s `ResponderRegistry` /
|
||
|
|
//! pending-reply pattern for the MCP `ask_question` tool (story 1228).
|
||
|
|
//!
|
||
|
|
//! Kept as a fully separate registry and pending-reply store from
|
||
|
|
//! `permission_router.rs` rather than reusing those types with a flag: AC4 of
|
||
|
|
//! story 1228 requires that a chat reply answering a permission prompt is
|
||
|
|
//! never treated as answering a pending question (and vice versa). Two
|
||
|
|
//! independent stores make that conflation structurally impossible instead of
|
||
|
|
//! relying on careful conditionals over a shared one.
|
||
|
|
|
||
|
|
use crate::http::context::{QuestionAnswer, QuestionForward};
|
||
|
|
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.
|
||
|
|
pub const QUESTION_RESPONDER_CHANNEL_CAPACITY: usize = 16;
|
||
|
|
|
||
|
|
struct ResponderSlot {
|
||
|
|
id: u64,
|
||
|
|
tx: mpsc::Sender<QuestionForward>,
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Tracks which tasks are currently registered to receive forwarded question
|
||
|
|
/// requests. Mirrors `permission_router::ResponderRegistry` exactly, but for
|
||
|
|
/// `ask_question` forwards instead of `prompt_permission` ones.
|
||
|
|
pub struct QuestionResponderRegistry {
|
||
|
|
next_id: AtomicU64,
|
||
|
|
slots: StdMutex<Vec<ResponderSlot>>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl QuestionResponderRegistry {
|
||
|
|
/// 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_ask_question` uses this to fail closed immediately (returning a
|
||
|
|
/// "no interactive session" result to the agent) instead of forwarding a
|
||
|
|
/// question 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.
|
||
|
|
pub fn register(self: &Arc<Self>) -> (QuestionResponderGuard, mpsc::Receiver<QuestionForward>) {
|
||
|
|
let (tx, rx) = mpsc::channel(QUESTION_RESPONDER_CHANNEL_CAPACITY);
|
||
|
|
let id = self.next_id.fetch_add(1, Ordering::Relaxed);
|
||
|
|
self.slots.lock().unwrap().push(ResponderSlot { id, tx });
|
||
|
|
(
|
||
|
|
QuestionResponderGuard {
|
||
|
|
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 an error result rather than dropped
|
||
|
|
/// silently.
|
||
|
|
pub fn dispatch(&self, forward: QuestionForward) {
|
||
|
|
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(Err(
|
||
|
|
"No interactive session is available to answer this question.".to_string(),
|
||
|
|
));
|
||
|
|
}
|
||
|
|
|
||
|
|
fn unregister(&self, id: u64) {
|
||
|
|
self.slots.lock().unwrap().retain(|s| s.id != id);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
trait CloneSenders {
|
||
|
|
fn clone_senders(&self) -> Vec<mpsc::Sender<QuestionForward>>;
|
||
|
|
}
|
||
|
|
|
||
|
|
impl CloneSenders for Vec<ResponderSlot> {
|
||
|
|
fn clone_senders(&self) -> Vec<mpsc::Sender<QuestionForward>> {
|
||
|
|
self.iter().map(|s| s.tx.clone()).collect()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// RAII guard returned by [`QuestionResponderRegistry::register`].
|
||
|
|
pub struct QuestionResponderGuard {
|
||
|
|
id: u64,
|
||
|
|
registry: Arc<QuestionResponderRegistry>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl Drop for QuestionResponderGuard {
|
||
|
|
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 [`QuestionResponderRegistry::dispatch`],
|
||
|
|
/// which is itself non-blocking. Exits when `question_rx` closes (server
|
||
|
|
/// shutdown).
|
||
|
|
pub fn spawn_question_router(
|
||
|
|
mut question_rx: mpsc::UnboundedReceiver<QuestionForward>,
|
||
|
|
registry: Arc<QuestionResponderRegistry>,
|
||
|
|
) -> tokio::task::JoinHandle<()> {
|
||
|
|
tokio::spawn(async move {
|
||
|
|
while let Some(forward) = question_rx.recv().await {
|
||
|
|
registry.dispatch(forward);
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
// ── Pending replies ──────────────────────────────────────────────────────
|
||
|
|
|
||
|
|
/// A pending question awaiting a chat reply, together with the metadata a
|
||
|
|
/// transport needs to parse that reply (option count, single/multi-select).
|
||
|
|
struct PendingQuestion {
|
||
|
|
tx: oneshot::Sender<Result<QuestionAnswer, String>>,
|
||
|
|
num_options: usize,
|
||
|
|
multi_select: bool,
|
||
|
|
labels: Vec<String>,
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Tracks questions awaiting a chat reply, keyed by `request_id` with a
|
||
|
|
/// per-location FIFO index — mirrors `permission_router::PendingPermReplies`.
|
||
|
|
pub struct PendingQuestionReplies {
|
||
|
|
inner: TokioMutex<PendingInner>,
|
||
|
|
}
|
||
|
|
|
||
|
|
#[derive(Default)]
|
||
|
|
struct PendingInner {
|
||
|
|
by_request_id: HashMap<String, PendingQuestion>,
|
||
|
|
by_location: HashMap<String, VecDeque<String>>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl PendingQuestionReplies {
|
||
|
|
/// 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>,
|
||
|
|
num_options: usize,
|
||
|
|
multi_select: bool,
|
||
|
|
labels: Vec<String>,
|
||
|
|
tx: oneshot::Sender<Result<QuestionAnswer, String>>,
|
||
|
|
) {
|
||
|
|
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,
|
||
|
|
PendingQuestion {
|
||
|
|
tx,
|
||
|
|
num_options,
|
||
|
|
multi_select,
|
||
|
|
labels,
|
||
|
|
},
|
||
|
|
);
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Peek the metadata (option count, multi-select, option labels) of the
|
||
|
|
/// oldest pending question queued for `location`, without consuming it.
|
||
|
|
/// Used so an invalid reply can trigger a re-prompt without losing the
|
||
|
|
/// pending request (AC4), and so a plain-text reply can be matched
|
||
|
|
/// against option labels (AC2).
|
||
|
|
pub async fn peek_oldest_meta(&self, location: &str) -> Option<(usize, bool, Vec<String>)> {
|
||
|
|
let inner = self.inner.lock().await;
|
||
|
|
let request_id = inner.by_location.get(location)?.front()?;
|
||
|
|
inner
|
||
|
|
.by_request_id
|
||
|
|
.get(request_id)
|
||
|
|
.map(|p| (p.num_options, p.multi_select, p.labels.clone()))
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Resolve the oldest pending question queued for `location`, removing it
|
||
|
|
/// from both the location queue and the request_id map.
|
||
|
|
pub async fn resolve_oldest(
|
||
|
|
&self,
|
||
|
|
location: &str,
|
||
|
|
) -> Option<oneshot::Sender<Result<QuestionAnswer, String>>> {
|
||
|
|
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(pending) = inner.by_request_id.remove(&request_id) {
|
||
|
|
return Some(pending.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. Used by responder timeout tasks.
|
||
|
|
pub async fn remove_by_request_id(
|
||
|
|
&self,
|
||
|
|
location: &str,
|
||
|
|
request_id: &str,
|
||
|
|
) -> Option<oneshot::Sender<Result<QuestionAnswer, String>>> {
|
||
|
|
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).map(|p| p.tx)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// ── Tests ─────────────────────────────────────────────────────────────────
|
||
|
|
|
||
|
|
#[cfg(test)]
|
||
|
|
mod tests {
|
||
|
|
use super::*;
|
||
|
|
use crate::http::context::QuestionSpec;
|
||
|
|
|
||
|
|
fn make_forward(
|
||
|
|
request_id: &str,
|
||
|
|
) -> (
|
||
|
|
QuestionForward,
|
||
|
|
oneshot::Receiver<Result<QuestionAnswer, String>>,
|
||
|
|
) {
|
||
|
|
let (tx, rx) = oneshot::channel();
|
||
|
|
(
|
||
|
|
QuestionForward {
|
||
|
|
request_id: request_id.to_string(),
|
||
|
|
question: QuestionSpec {
|
||
|
|
header: "Test".to_string(),
|
||
|
|
question: "Pick one?".to_string(),
|
||
|
|
options: vec![],
|
||
|
|
multi_select: false,
|
||
|
|
},
|
||
|
|
response_tx: tx,
|
||
|
|
},
|
||
|
|
rx,
|
||
|
|
)
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn registry_starts_empty() {
|
||
|
|
let registry = QuestionResponderRegistry::new();
|
||
|
|
assert!(registry.is_empty());
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn register_makes_registry_non_empty_until_guard_dropped() {
|
||
|
|
let registry = QuestionResponderRegistry::new();
|
||
|
|
let (guard, _rx) = registry.register();
|
||
|
|
assert!(!registry.is_empty());
|
||
|
|
drop(guard);
|
||
|
|
assert!(registry.is_empty());
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn dispatch_delivers_to_registered_responder() {
|
||
|
|
let registry = QuestionResponderRegistry::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_fails_closed_when_no_responder_registered() {
|
||
|
|
let registry = QuestionResponderRegistry::new();
|
||
|
|
let (fwd, response_rx) = make_forward("req-2");
|
||
|
|
registry.dispatch(fwd);
|
||
|
|
let result = response_rx.await.expect("must receive a result");
|
||
|
|
assert!(
|
||
|
|
result.is_err(),
|
||
|
|
"no registered responder must fail-closed with an error"
|
||
|
|
);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn pending_replies_two_concurrent_requests_same_location_both_resolve() {
|
||
|
|
let pending = PendingQuestionReplies::new();
|
||
|
|
let (tx1, rx1) = oneshot::channel();
|
||
|
|
let (tx2, rx2) = oneshot::channel();
|
||
|
|
pending
|
||
|
|
.insert(
|
||
|
|
"room-1",
|
||
|
|
"req-a",
|
||
|
|
3,
|
||
|
|
false,
|
||
|
|
vec!["A".to_string(), "B".to_string(), "C".to_string()],
|
||
|
|
tx1,
|
||
|
|
)
|
||
|
|
.await;
|
||
|
|
pending
|
||
|
|
.insert("room-1", "req-b", 2, true, vec![], tx2)
|
||
|
|
.await;
|
||
|
|
|
||
|
|
let first = pending
|
||
|
|
.resolve_oldest("room-1")
|
||
|
|
.await
|
||
|
|
.expect("first pending reply must still be present");
|
||
|
|
let _ = first.send(Ok(QuestionAnswer::Selected(vec![0])));
|
||
|
|
assert_eq!(
|
||
|
|
rx1.await.unwrap().unwrap(),
|
||
|
|
QuestionAnswer::Selected(vec![0])
|
||
|
|
);
|
||
|
|
|
||
|
|
let second = pending
|
||
|
|
.resolve_oldest("room-1")
|
||
|
|
.await
|
||
|
|
.expect("second pending reply must still be present");
|
||
|
|
let _ = second.send(Ok(QuestionAnswer::FreeText("other".to_string())));
|
||
|
|
assert_eq!(
|
||
|
|
rx2.await.unwrap().unwrap(),
|
||
|
|
QuestionAnswer::FreeText("other".to_string())
|
||
|
|
);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn peek_oldest_meta_does_not_consume() {
|
||
|
|
let pending = PendingQuestionReplies::new();
|
||
|
|
let (tx, _rx) = oneshot::channel();
|
||
|
|
pending
|
||
|
|
.insert(
|
||
|
|
"room-1",
|
||
|
|
"req-a",
|
||
|
|
4,
|
||
|
|
true,
|
||
|
|
vec!["X".to_string(), "Y".to_string()],
|
||
|
|
tx,
|
||
|
|
)
|
||
|
|
.await;
|
||
|
|
|
||
|
|
let meta = pending.peek_oldest_meta("room-1").await;
|
||
|
|
assert_eq!(
|
||
|
|
meta,
|
||
|
|
Some((4, true, vec!["X".to_string(), "Y".to_string()]))
|
||
|
|
);
|
||
|
|
|
||
|
|
// Peeking again must return the same entry — it was not consumed.
|
||
|
|
let meta_again = pending.peek_oldest_meta("room-1").await;
|
||
|
|
assert_eq!(
|
||
|
|
meta_again,
|
||
|
|
Some((4, true, vec!["X".to_string(), "Y".to_string()]))
|
||
|
|
);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn remove_by_request_id_prevents_later_resolution() {
|
||
|
|
let pending = PendingQuestionReplies::new();
|
||
|
|
let (tx, _rx) = oneshot::channel();
|
||
|
|
pending
|
||
|
|
.insert("room-1", "req-timeout", 2, false, vec![], tx)
|
||
|
|
.await;
|
||
|
|
|
||
|
|
let removed = pending.remove_by_request_id("room-1", "req-timeout").await;
|
||
|
|
assert!(removed.is_some());
|
||
|
|
|
||
|
|
let resolved = pending.resolve_oldest("room-1").await;
|
||
|
|
assert!(resolved.is_none());
|
||
|
|
assert!(pending.peek_oldest_meta("room-1").await.is_none());
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn resolve_oldest_returns_none_for_unknown_location() {
|
||
|
|
let pending = PendingQuestionReplies::new();
|
||
|
|
assert!(pending.resolve_oldest("no-such-room").await.is_none());
|
||
|
|
}
|
||
|
|
}
|