|
|
|
@@ -1,11 +1,11 @@
|
|
|
|
|
//! Background task that holds `perm_rx` for the bot's lifetime and forwards
|
|
|
|
|
//! permission requests to the configured Matrix room.
|
|
|
|
|
//! Background task that registers as a permission responder for the bot's
|
|
|
|
|
//! lifetime and forwards permission requests to the configured Matrix room.
|
|
|
|
|
//!
|
|
|
|
|
//! Before story 884, each chat message handler acquired `perm_rx` for the
|
|
|
|
|
//! duration of one chat_fut and dropped it afterwards. That meant whenever
|
|
|
|
|
//! the bot wasn't actively responding, `prompt_permission` auto-denied any
|
|
|
|
|
//! spawned coder bash call as "no interactive session" — making unattended
|
|
|
|
|
//! coder work impossible. This task holds the lock continuously while the
|
|
|
|
|
//! coder work impossible. This task stays registered continuously while the
|
|
|
|
|
//! bot is connected, so requests can flow at any time.
|
|
|
|
|
|
|
|
|
|
use crate::chat::ChatTransport;
|
|
|
|
@@ -20,13 +20,13 @@ use tokio::sync::Mutex as TokioMutex;
|
|
|
|
|
|
|
|
|
|
use super::format::markdown_to_html;
|
|
|
|
|
|
|
|
|
|
/// Spawn a background task that holds `services.perm_rx` for the bot's
|
|
|
|
|
/// lifetime and forwards each incoming permission request to `target_room`
|
|
|
|
|
/// as a chat message. Replies (yes/no) are resolved by the existing
|
|
|
|
|
/// `on_room_message` handler via `pending_perm_replies`.
|
|
|
|
|
/// Spawn a background task that registers with `services.permission_registry`
|
|
|
|
|
/// for the bot's lifetime and forwards each incoming permission request to
|
|
|
|
|
/// `target_room` as a chat message. Replies (yes/no) are resolved by the
|
|
|
|
|
/// existing `on_room_message` handler via `pending_perm_replies`.
|
|
|
|
|
///
|
|
|
|
|
/// Returns the JoinHandle so the caller can keep ownership; the task exits
|
|
|
|
|
/// only when the `perm_rx` channel is closed (bot shutdown).
|
|
|
|
|
/// only when its private channel closes (bot shutdown).
|
|
|
|
|
pub fn spawn_permission_listener(
|
|
|
|
|
services: Arc<Services>,
|
|
|
|
|
transport: Arc<dyn ChatTransport>,
|
|
|
|
@@ -34,7 +34,7 @@ pub fn spawn_permission_listener(
|
|
|
|
|
bot_sent_event_ids: Arc<TokioMutex<HashSet<OwnedEventId>>>,
|
|
|
|
|
) -> tokio::task::JoinHandle<()> {
|
|
|
|
|
tokio::spawn(async move {
|
|
|
|
|
let mut perm_rx = services.perm_rx.lock().await;
|
|
|
|
|
let (_responder_guard, mut perm_rx) = services.permission_registry.register();
|
|
|
|
|
let target_room_str = target_room.as_str().to_string();
|
|
|
|
|
slog!("[matrix-bot] permission listener started; forwarding requests to {target_room_str}");
|
|
|
|
|
|
|
|
|
@@ -57,24 +57,33 @@ pub fn spawn_permission_listener(
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Store the MCP oneshot sender so on_room_message can resolve it
|
|
|
|
|
// when the user replies yes/no in the target room.
|
|
|
|
|
// when the user replies yes/no in the target room. Keyed by
|
|
|
|
|
// request_id (not just room) so a second concurrent request for
|
|
|
|
|
// the same room doesn't drop the first's sender.
|
|
|
|
|
services
|
|
|
|
|
.pending_perm_replies
|
|
|
|
|
.lock()
|
|
|
|
|
.await
|
|
|
|
|
.insert(target_room.to_string(), perm_fwd.response_tx);
|
|
|
|
|
.insert(
|
|
|
|
|
target_room.to_string(),
|
|
|
|
|
perm_fwd.request_id.clone(),
|
|
|
|
|
perm_fwd.response_tx,
|
|
|
|
|
)
|
|
|
|
|
.await;
|
|
|
|
|
|
|
|
|
|
// Spawn a per-request timeout: auto-deny if the user does not
|
|
|
|
|
// respond within `permission_timeout_secs`.
|
|
|
|
|
let pending = Arc::clone(&services.pending_perm_replies);
|
|
|
|
|
let timeout_room_key = target_room.to_string();
|
|
|
|
|
let timeout_request_id = perm_fwd.request_id.clone();
|
|
|
|
|
let timeout_transport = Arc::clone(&transport);
|
|
|
|
|
let timeout_room_str = target_room_str.clone();
|
|
|
|
|
let timeout_sent_ids = Arc::clone(&bot_sent_event_ids);
|
|
|
|
|
let timeout_secs = services.permission_timeout_secs;
|
|
|
|
|
tokio::spawn(async move {
|
|
|
|
|
tokio::time::sleep(Duration::from_secs(timeout_secs)).await;
|
|
|
|
|
if let Some(tx) = pending.lock().await.remove(&timeout_room_key) {
|
|
|
|
|
if let Some(tx) = pending
|
|
|
|
|
.remove_by_request_id(&timeout_room_key, &timeout_request_id)
|
|
|
|
|
.await
|
|
|
|
|
{
|
|
|
|
|
let _ = tx.send(PermissionDecision::Deny);
|
|
|
|
|
let msg = "Permission request timed out — denied (fail-closed).";
|
|
|
|
|
let html = markdown_to_html(msg);
|
|
|
|
@@ -89,7 +98,7 @@ pub fn spawn_permission_listener(
|
|
|
|
|
});
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
slog!("[matrix-bot] permission listener exiting (perm_rx channel closed)");
|
|
|
|
|
slog!("[matrix-bot] permission listener exiting (channel closed)");
|
|
|
|
|
})
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
@@ -102,11 +111,11 @@ mod tests {
|
|
|
|
|
|
|
|
|
|
use super::*;
|
|
|
|
|
use crate::http::context::PermissionForward;
|
|
|
|
|
use crate::service::permission_router::PendingPermReplies;
|
|
|
|
|
use crate::services::Services;
|
|
|
|
|
use async_trait::async_trait;
|
|
|
|
|
use serde_json::json;
|
|
|
|
|
use std::collections::HashMap;
|
|
|
|
|
use tokio::sync::{mpsc, oneshot};
|
|
|
|
|
use tokio::sync::oneshot;
|
|
|
|
|
|
|
|
|
|
struct RecordingTransport {
|
|
|
|
|
sent: Arc<std::sync::Mutex<Vec<(String, String)>>>,
|
|
|
|
@@ -138,26 +147,40 @@ mod tests {
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
fn test_services_with_tx() -> (Arc<Services>, mpsc::UnboundedSender<PermissionForward>) {
|
|
|
|
|
let (perm_tx, perm_rx) = mpsc::unbounded_channel();
|
|
|
|
|
let services = Arc::new(Services {
|
|
|
|
|
/// Poll `cond` until it returns `true` or `timeout` elapses, sleeping
|
|
|
|
|
/// briefly between checks. Used instead of a single fixed sleep so tests
|
|
|
|
|
/// don't flake under slow/loaded CI where a spawned task hasn't yet run.
|
|
|
|
|
async fn wait_until(mut cond: impl FnMut() -> bool, timeout: std::time::Duration) -> bool {
|
|
|
|
|
let start = tokio::time::Instant::now();
|
|
|
|
|
loop {
|
|
|
|
|
if cond() {
|
|
|
|
|
return true;
|
|
|
|
|
}
|
|
|
|
|
if start.elapsed() > timeout {
|
|
|
|
|
return false;
|
|
|
|
|
}
|
|
|
|
|
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
fn test_services() -> Arc<Services> {
|
|
|
|
|
Arc::new(Services {
|
|
|
|
|
project_root: std::path::PathBuf::from("/tmp/test"),
|
|
|
|
|
agents: Arc::new(crate::agents::AgentPool::new_test(3000)),
|
|
|
|
|
bot_name: "Assistant".to_string(),
|
|
|
|
|
bot_user_id: "@bot:example.com".to_string(),
|
|
|
|
|
ambient_rooms: Arc::new(std::sync::Mutex::new(std::collections::HashSet::new())),
|
|
|
|
|
perm_rx: Arc::new(TokioMutex::new(perm_rx)),
|
|
|
|
|
pending_perm_replies: Arc::new(TokioMutex::new(HashMap::new())),
|
|
|
|
|
permission_registry: crate::service::permission_router::ResponderRegistry::new(),
|
|
|
|
|
pending_perm_replies: PendingPermReplies::new(),
|
|
|
|
|
permission_timeout_secs: 120,
|
|
|
|
|
status: Arc::new(crate::service::status::StatusBroadcaster::new()),
|
|
|
|
|
chat_dispatcher: Arc::new(crate::chat::dispatcher::ChatDispatcher::new(1_500)),
|
|
|
|
|
});
|
|
|
|
|
(services, perm_tx)
|
|
|
|
|
})
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[tokio::test]
|
|
|
|
|
async fn listener_forwards_request_to_target_room_when_no_chat_in_flight() {
|
|
|
|
|
let (services, perm_tx) = test_services_with_tx();
|
|
|
|
|
let services = test_services();
|
|
|
|
|
let sent: Arc<std::sync::Mutex<Vec<(String, String)>>> =
|
|
|
|
|
Arc::new(std::sync::Mutex::new(Vec::new()));
|
|
|
|
|
let transport: Arc<dyn crate::chat::ChatTransport> = Arc::new(RecordingTransport {
|
|
|
|
@@ -173,21 +196,33 @@ mod tests {
|
|
|
|
|
Arc::clone(&bot_sent_event_ids),
|
|
|
|
|
);
|
|
|
|
|
|
|
|
|
|
// Yield so the listener task acquires perm_rx and starts recv'ing.
|
|
|
|
|
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
|
|
|
|
|
// Wait until the listener task registers and starts recv'ing.
|
|
|
|
|
assert!(
|
|
|
|
|
wait_until(
|
|
|
|
|
|| !services.permission_registry.is_empty(),
|
|
|
|
|
std::time::Duration::from_secs(2)
|
|
|
|
|
)
|
|
|
|
|
.await,
|
|
|
|
|
"listener never registered as a responder"
|
|
|
|
|
);
|
|
|
|
|
|
|
|
|
|
let (response_tx, _response_rx) = oneshot::channel();
|
|
|
|
|
perm_tx
|
|
|
|
|
.send(PermissionForward {
|
|
|
|
|
request_id: "req-1".to_string(),
|
|
|
|
|
tool_name: "Bash".to_string(),
|
|
|
|
|
tool_input: json!({"command": "cargo test"}),
|
|
|
|
|
response_tx,
|
|
|
|
|
})
|
|
|
|
|
.expect("send PermissionForward");
|
|
|
|
|
services.permission_registry.dispatch(PermissionForward {
|
|
|
|
|
request_id: "req-1".to_string(),
|
|
|
|
|
tool_name: "Bash".to_string(),
|
|
|
|
|
tool_input: json!({"command": "cargo test"}),
|
|
|
|
|
response_tx,
|
|
|
|
|
});
|
|
|
|
|
|
|
|
|
|
// Give the listener a moment to process.
|
|
|
|
|
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
|
|
|
|
|
// Wait for the listener to process the forward.
|
|
|
|
|
assert!(
|
|
|
|
|
wait_until(
|
|
|
|
|
|| !sent.lock().unwrap().is_empty(),
|
|
|
|
|
std::time::Duration::from_secs(2)
|
|
|
|
|
)
|
|
|
|
|
.await,
|
|
|
|
|
"listener never sent the permission prompt"
|
|
|
|
|
);
|
|
|
|
|
|
|
|
|
|
// The transport must have received exactly one send_message to the
|
|
|
|
|
// target room with the prompt content.
|
|
|
|
@@ -205,17 +240,126 @@ mod tests {
|
|
|
|
|
recorded[0].1
|
|
|
|
|
);
|
|
|
|
|
|
|
|
|
|
// pending_perm_replies must contain an entry keyed by the target room
|
|
|
|
|
// (so the user-reply handler can resolve the request when they reply).
|
|
|
|
|
let pending = services.pending_perm_replies.lock().await;
|
|
|
|
|
// pending_perm_replies must contain an entry resolvable for the target
|
|
|
|
|
// room (so the user-reply handler can resolve the request when they
|
|
|
|
|
// reply). The insert happens just after send_message returns, so poll
|
|
|
|
|
// briefly rather than assuming it's already visible.
|
|
|
|
|
let mut resolved = None;
|
|
|
|
|
for _ in 0..50 {
|
|
|
|
|
resolved = services
|
|
|
|
|
.pending_perm_replies
|
|
|
|
|
.resolve_oldest(target_room.as_str())
|
|
|
|
|
.await;
|
|
|
|
|
if resolved.is_some() {
|
|
|
|
|
break;
|
|
|
|
|
}
|
|
|
|
|
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
|
|
|
|
|
}
|
|
|
|
|
assert!(
|
|
|
|
|
pending.contains_key(target_room.as_str()),
|
|
|
|
|
resolved.is_some(),
|
|
|
|
|
"pending_perm_replies missing entry for target room"
|
|
|
|
|
);
|
|
|
|
|
|
|
|
|
|
// bot_sent_event_ids must have recorded the prompt's event ID so the
|
|
|
|
|
// bot does not echo its own prompt back as user input.
|
|
|
|
|
let sent_ids = bot_sent_event_ids.lock().await;
|
|
|
|
|
assert_eq!(sent_ids.len(), 1, "expected one sent event ID recorded");
|
|
|
|
|
assert!(
|
|
|
|
|
wait_until(
|
|
|
|
|
|| bot_sent_event_ids
|
|
|
|
|
.try_lock()
|
|
|
|
|
.map(|s| s.len() == 1)
|
|
|
|
|
.unwrap_or(false),
|
|
|
|
|
std::time::Duration::from_secs(2)
|
|
|
|
|
)
|
|
|
|
|
.await,
|
|
|
|
|
"expected one sent event ID recorded"
|
|
|
|
|
);
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Regression: two concurrent permission requests for the same target
|
|
|
|
|
/// room must both be resolvable — keying `pending_perm_replies` by
|
|
|
|
|
/// request_id (rather than overwriting a single room-keyed entry) means
|
|
|
|
|
/// the first request's oneshot sender is never dropped by the second.
|
|
|
|
|
#[tokio::test]
|
|
|
|
|
async fn two_concurrent_requests_for_same_room_both_resolve() {
|
|
|
|
|
let services = test_services();
|
|
|
|
|
let sent: Arc<std::sync::Mutex<Vec<(String, String)>>> =
|
|
|
|
|
Arc::new(std::sync::Mutex::new(Vec::new()));
|
|
|
|
|
let transport: Arc<dyn crate::chat::ChatTransport> = Arc::new(RecordingTransport {
|
|
|
|
|
sent: Arc::clone(&sent),
|
|
|
|
|
});
|
|
|
|
|
let target_room: OwnedRoomId = "!perm:example.com".parse().unwrap();
|
|
|
|
|
let bot_sent_event_ids = Arc::new(TokioMutex::new(HashSet::new()));
|
|
|
|
|
|
|
|
|
|
spawn_permission_listener(
|
|
|
|
|
Arc::clone(&services),
|
|
|
|
|
Arc::clone(&transport),
|
|
|
|
|
target_room.clone(),
|
|
|
|
|
Arc::clone(&bot_sent_event_ids),
|
|
|
|
|
);
|
|
|
|
|
assert!(
|
|
|
|
|
wait_until(
|
|
|
|
|
|| !services.permission_registry.is_empty(),
|
|
|
|
|
std::time::Duration::from_secs(2)
|
|
|
|
|
)
|
|
|
|
|
.await,
|
|
|
|
|
"listener never registered as a responder"
|
|
|
|
|
);
|
|
|
|
|
|
|
|
|
|
let (tx1, rx1) = oneshot::channel();
|
|
|
|
|
services.permission_registry.dispatch(PermissionForward {
|
|
|
|
|
request_id: "req-a".to_string(),
|
|
|
|
|
tool_name: "Bash".to_string(),
|
|
|
|
|
tool_input: json!({}),
|
|
|
|
|
response_tx: tx1,
|
|
|
|
|
});
|
|
|
|
|
let (tx2, rx2) = oneshot::channel();
|
|
|
|
|
services.permission_registry.dispatch(PermissionForward {
|
|
|
|
|
request_id: "req-b".to_string(),
|
|
|
|
|
tool_name: "Write".to_string(),
|
|
|
|
|
tool_input: json!({}),
|
|
|
|
|
response_tx: tx2,
|
|
|
|
|
});
|
|
|
|
|
|
|
|
|
|
// Wait until the listener has processed both forwards.
|
|
|
|
|
assert!(
|
|
|
|
|
wait_until(
|
|
|
|
|
|| sent.lock().unwrap().len() >= 2,
|
|
|
|
|
std::time::Duration::from_secs(2)
|
|
|
|
|
)
|
|
|
|
|
.await,
|
|
|
|
|
"listener never sent both permission prompts"
|
|
|
|
|
);
|
|
|
|
|
|
|
|
|
|
// Poll rather than assuming the pending_perm_replies insert (which
|
|
|
|
|
// happens just after send_message returns) is already visible.
|
|
|
|
|
let mut first = None;
|
|
|
|
|
for _ in 0..50 {
|
|
|
|
|
first = services
|
|
|
|
|
.pending_perm_replies
|
|
|
|
|
.resolve_oldest(target_room.as_str())
|
|
|
|
|
.await;
|
|
|
|
|
if first.is_some() {
|
|
|
|
|
break;
|
|
|
|
|
}
|
|
|
|
|
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
|
|
|
|
|
}
|
|
|
|
|
let first = first.expect("first request's sender must not have been dropped");
|
|
|
|
|
let _ = first.send(PermissionDecision::Approve);
|
|
|
|
|
assert_eq!(rx1.await.unwrap(), PermissionDecision::Approve);
|
|
|
|
|
|
|
|
|
|
let mut second = None;
|
|
|
|
|
for _ in 0..50 {
|
|
|
|
|
second = services
|
|
|
|
|
.pending_perm_replies
|
|
|
|
|
.resolve_oldest(target_room.as_str())
|
|
|
|
|
.await;
|
|
|
|
|
if second.is_some() {
|
|
|
|
|
break;
|
|
|
|
|
}
|
|
|
|
|
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
|
|
|
|
|
}
|
|
|
|
|
let second = second.expect("second request's sender must still be present");
|
|
|
|
|
let _ = second.send(PermissionDecision::Deny);
|
|
|
|
|
assert_eq!(rx2.await.unwrap(), PermissionDecision::Deny);
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
Huskies Agent