499 lines
20 KiB
Rust
499 lines
20 KiB
Rust
//! Matrix bot context — shared state for the Matrix bot (rooms, history, permissions).
|
||
use crate::chat::ChatTransport;
|
||
use crate::service::timer::TimerStore;
|
||
use crate::services::Services;
|
||
use matrix_sdk::ruma::{OwnedEventId, OwnedRoomId, OwnedUserId};
|
||
use std::collections::{BTreeMap, HashSet, VecDeque};
|
||
use std::sync::Arc;
|
||
use tokio::sync::Mutex as TokioMutex;
|
||
use tokio::sync::RwLock;
|
||
|
||
use super::history::ConversationHistory;
|
||
|
||
/// Maximum number of incoming event IDs retained for deduplication.
|
||
///
|
||
/// FIFO ring (VecDeque + HashSet): when full, the oldest entry is evicted.
|
||
/// Matrix sync replays are temporally clustered — any event replayed more than
|
||
/// 1 000 events later was long since processed, so FIFO eviction is correct.
|
||
/// 1 000 × ~80 B ≈ 80 KB, negligible memory cost.
|
||
pub const SEEN_EVENT_IDS_CAP: usize = 1_000;
|
||
|
||
/// Bounded FIFO set for deduplicating incoming Matrix event IDs.
|
||
pub struct SeenEventIds {
|
||
deque: VecDeque<OwnedEventId>,
|
||
set: HashSet<OwnedEventId>,
|
||
cap: usize,
|
||
}
|
||
|
||
impl SeenEventIds {
|
||
/// Create a new set with the given capacity cap.
|
||
pub fn new(cap: usize) -> Self {
|
||
Self {
|
||
deque: VecDeque::with_capacity(cap),
|
||
set: HashSet::with_capacity(cap),
|
||
cap,
|
||
}
|
||
}
|
||
|
||
/// Insert an event ID. Returns `true` if the ID was new (never seen),
|
||
/// `false` if already present (duplicate — caller should skip processing).
|
||
/// When the set is full, the oldest entry is evicted before inserting.
|
||
pub fn insert(&mut self, id: OwnedEventId) -> bool {
|
||
if self.set.contains(&id) {
|
||
return false;
|
||
}
|
||
if self.deque.len() == self.cap
|
||
&& let Some(old) = self.deque.pop_front()
|
||
{
|
||
self.set.remove(&old);
|
||
}
|
||
self.deque.push_back(id.clone());
|
||
self.set.insert(id);
|
||
true
|
||
}
|
||
}
|
||
|
||
/// Shared context injected into Matrix event handlers.
|
||
#[derive(Clone)]
|
||
pub struct BotContext {
|
||
/// Shared services bundle (project root, agent pool, bot identity, permissions).
|
||
pub services: Arc<Services>,
|
||
/// Matrix-specific parsed user ID (e.g. `@timmy:homeserver.local`).
|
||
/// Transport-specific — kept separate from `services.bot_user_id` (String)
|
||
/// because Matrix SDK APIs require `OwnedUserId` for comparisons and
|
||
/// `.localpart()` extraction.
|
||
pub matrix_user_id: OwnedUserId,
|
||
/// All room IDs the bot listens in.
|
||
pub target_room_ids: Vec<OwnedRoomId>,
|
||
pub allowed_users: Vec<String>,
|
||
/// Shared, per-room rolling conversation history.
|
||
pub history: ConversationHistory,
|
||
/// Maximum number of entries to keep per room before trimming the oldest.
|
||
pub history_size: usize,
|
||
/// Event IDs of messages the bot has sent. Used to detect replies to the
|
||
/// bot so it can continue a conversation thread without requiring an
|
||
/// explicit `@mention` on every follow-up.
|
||
pub bot_sent_event_ids: Arc<TokioMutex<HashSet<OwnedEventId>>>,
|
||
/// Per-room htop monitoring sessions. Keyed by room ID; each entry holds
|
||
/// a stop-signal sender that the background task watches.
|
||
pub htop_sessions: super::super::htop::HtopSessions,
|
||
/// Chat transport used for sending and editing messages.
|
||
///
|
||
/// All message I/O goes through this abstraction so the bot logic works
|
||
/// with any platform, not just Matrix.
|
||
pub transport: Arc<dyn ChatTransport>,
|
||
/// Persistent store for pending deferred-start timers.
|
||
pub timer_store: Arc<TimerStore>,
|
||
/// In gateway mode: the currently active project (shared with the gateway HTTP handler).
|
||
/// `None` in standalone single-project mode.
|
||
pub gateway_active_project: Option<Arc<RwLock<String>>>,
|
||
/// In gateway mode: valid project names accepted by the `switch` command.
|
||
/// Empty in standalone mode.
|
||
pub gateway_projects: Vec<String>,
|
||
/// In gateway mode: mapping of project name → base URL (e.g. `"http://localhost:3001"`).
|
||
/// Used to proxy bot commands to the active project over WebSocket (`/ws`).
|
||
/// Empty in standalone mode.
|
||
pub gateway_project_urls: BTreeMap<String, String>,
|
||
/// Pipeline transition events buffered since the last LLM turn.
|
||
///
|
||
/// A background task appends one compact audit line per real stage
|
||
/// transition. `handle_message` drains this buffer and injects it as a
|
||
/// `<system-reminder>` block at the head of the next user prompt so Timmy
|
||
/// sees pipeline activity without requiring a separate message.
|
||
pub pending_pipeline_events: Arc<TokioMutex<Vec<String>>>,
|
||
/// Gateway aggregate transition events buffered since the last LLM turn.
|
||
///
|
||
/// In gateway mode a background task appends one compact audit line per
|
||
/// `GatewayStatusEvent` received from the gateway broadcaster. Drained
|
||
/// alongside `pending_pipeline_events` on each user message. Always
|
||
/// empty in standalone (non-gateway) mode.
|
||
pub pending_gateway_events: Arc<TokioMutex<Vec<String>>>,
|
||
/// Bounded FIFO set of already-handled incoming event IDs.
|
||
///
|
||
/// The Matrix sync loop can replay events on reconnect. This set ensures
|
||
/// each event is processed at most once. Insert the event ID before any
|
||
/// side-effecting work; return early if the insert returns `false`.
|
||
pub handled_incoming_event_ids: Arc<TokioMutex<SeenEventIds>>,
|
||
}
|
||
|
||
impl BotContext {
|
||
/// Resolve the effective project root for command dispatch.
|
||
///
|
||
/// In gateway mode the bot's `project_root` is the gateway config directory.
|
||
/// Each project lives in a subdirectory named after the project, so the
|
||
/// effective root for commands is `project_root / active_project_name`.
|
||
/// In standalone (single-project) mode this returns `project_root` unchanged.
|
||
pub async fn effective_project_root(&self) -> std::path::PathBuf {
|
||
if let Some(ref ap) = self.gateway_active_project {
|
||
let name = ap.read().await.clone();
|
||
self.services.project_root.join(&name)
|
||
} else {
|
||
self.services.project_root.clone()
|
||
}
|
||
}
|
||
|
||
/// Returns `true` if the bot is running in gateway mode.
|
||
pub fn is_gateway(&self) -> bool {
|
||
self.gateway_active_project.is_some()
|
||
}
|
||
|
||
/// Return the base URL for the currently active project, if in gateway mode.
|
||
pub async fn active_project_url(&self) -> Option<String> {
|
||
let ap = self.gateway_active_project.as_ref()?;
|
||
let name = ap.read().await.clone();
|
||
self.gateway_project_urls.get(&name).cloned()
|
||
}
|
||
|
||
/// Proxy a bot command to the active project over a WebSocket RPC call.
|
||
///
|
||
/// Connects to `{base_url}/ws`, sends an `rpc_request` frame for the
|
||
/// `bot.command` method, and returns the Markdown response from the
|
||
/// `rpc_response` frame. Returns an error message string if the
|
||
/// connection or command fails.
|
||
pub async fn proxy_bot_command(&self, command: &str, args: &str) -> Option<String> {
|
||
use futures::{SinkExt, StreamExt};
|
||
use tokio_tungstenite::tungstenite::Message as WsMsg;
|
||
|
||
let base_url = self.active_project_url().await?;
|
||
|
||
// Convert http(s):// → ws(s)://
|
||
let ws_base = if let Some(rest) = base_url.strip_prefix("https://") {
|
||
format!("wss://{rest}")
|
||
} else if let Some(rest) = base_url.strip_prefix("http://") {
|
||
format!("ws://{rest}")
|
||
} else {
|
||
base_url.clone()
|
||
};
|
||
let ws_url = format!("{ws_base}/ws");
|
||
|
||
let correlation_id = uuid::Uuid::new_v4().to_string();
|
||
let request = serde_json::json!({
|
||
"kind": "rpc_request",
|
||
"version": 1,
|
||
"correlation_id": correlation_id,
|
||
"ttl_ms": 30_000u64,
|
||
"method": "bot.command",
|
||
"params": { "command": command, "args": args },
|
||
});
|
||
let request_text = match serde_json::to_string(&request) {
|
||
Ok(t) => t,
|
||
Err(e) => return Some(format!("Failed to serialize RPC request: {e}")),
|
||
};
|
||
|
||
let ws_stream = match tokio_tungstenite::connect_async(&ws_url).await {
|
||
Ok((stream, _)) => stream,
|
||
Err(e) => {
|
||
return Some(format!(
|
||
"Failed to connect to project server at {ws_url}: {e}"
|
||
));
|
||
}
|
||
};
|
||
|
||
let (mut sink, mut stream) = ws_stream.split();
|
||
|
||
if let Err(e) = sink.send(WsMsg::Text(request_text.into())).await {
|
||
return Some(format!("Failed to send RPC request: {e}"));
|
||
}
|
||
|
||
while let Some(msg) = stream.next().await {
|
||
match msg {
|
||
Ok(WsMsg::Text(text)) => {
|
||
let Ok(frame) = serde_json::from_str::<serde_json::Value>(&text) else {
|
||
continue;
|
||
};
|
||
if frame.get("kind").and_then(|v| v.as_str()) != Some("rpc_response") {
|
||
continue;
|
||
}
|
||
if frame
|
||
.get("correlation_id")
|
||
.and_then(|v| v.as_str())
|
||
.map(|id| id != correlation_id)
|
||
.unwrap_or(true)
|
||
{
|
||
continue;
|
||
}
|
||
let ok = frame.get("ok").and_then(|v| v.as_bool()).unwrap_or(false);
|
||
if ok {
|
||
return frame
|
||
.get("result")
|
||
.and_then(|r| r.get("response"))
|
||
.and_then(|v| v.as_str())
|
||
.map(String::from)
|
||
.or_else(|| {
|
||
Some("Command succeeded with no response text".to_string())
|
||
});
|
||
} else {
|
||
let err = frame
|
||
.get("error")
|
||
.and_then(|v| v.as_str())
|
||
.unwrap_or("unknown error");
|
||
return Some(format!("Project server command failed: {err}"));
|
||
}
|
||
}
|
||
Ok(WsMsg::Close(_)) => break,
|
||
Err(e) => return Some(format!("WebSocket error: {e}")),
|
||
_ => continue,
|
||
}
|
||
}
|
||
|
||
Some("Connection closed before receiving command response".to_string())
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Tests
|
||
// ---------------------------------------------------------------------------
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use std::collections::HashMap;
|
||
use std::path::PathBuf;
|
||
use tokio::sync::mpsc;
|
||
|
||
fn make_user_id(s: &str) -> OwnedUserId {
|
||
s.parse().unwrap()
|
||
}
|
||
|
||
/// Build a test `Services` bundle with the given project root.
|
||
fn test_services(project_root: PathBuf) -> Arc<Services> {
|
||
let (_perm_tx, perm_rx) = mpsc::unbounded_channel();
|
||
Arc::new(Services {
|
||
project_root,
|
||
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_timeout_secs: 120,
|
||
status: Arc::new(crate::service::status::StatusBroadcaster::new()),
|
||
chat_dispatcher: Arc::new(crate::chat::dispatcher::ChatDispatcher::new(1_500)),
|
||
})
|
||
}
|
||
|
||
/// Build a minimal `BotContext` for testing with the given Services and
|
||
/// optional gateway active project.
|
||
fn test_bot_context(
|
||
services: Arc<Services>,
|
||
gateway_active_project: Option<Arc<RwLock<String>>>,
|
||
gateway_projects: Vec<String>,
|
||
gateway_project_urls: BTreeMap<String, String>,
|
||
) -> BotContext {
|
||
BotContext {
|
||
services,
|
||
matrix_user_id: make_user_id("@bot:example.com"),
|
||
target_room_ids: vec![],
|
||
allowed_users: vec![],
|
||
history: Arc::new(TokioMutex::new(HashMap::new())),
|
||
history_size: 20,
|
||
bot_sent_event_ids: Arc::new(TokioMutex::new(std::collections::HashSet::new())),
|
||
htop_sessions: Arc::new(TokioMutex::new(HashMap::new())),
|
||
transport: Arc::new(crate::chat::transport::whatsapp::WhatsAppTransport::new(
|
||
"test-phone".to_string(),
|
||
"test-token".to_string(),
|
||
"pipeline_notification".to_string(),
|
||
)),
|
||
timer_store: Arc::new(crate::service::timer::TimerStore::load(
|
||
std::path::PathBuf::from("/tmp/timers.json"),
|
||
)),
|
||
gateway_active_project,
|
||
gateway_projects,
|
||
gateway_project_urls,
|
||
pending_pipeline_events: Arc::new(TokioMutex::new(Vec::new())),
|
||
pending_gateway_events: Arc::new(TokioMutex::new(Vec::new())),
|
||
handled_incoming_event_ids: Arc::new(TokioMutex::new(SeenEventIds::new(
|
||
SEEN_EVENT_IDS_CAP,
|
||
))),
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn bot_context_is_clone() {
|
||
// BotContext must be Clone for the Matrix event handler injection.
|
||
fn assert_clone<T: Clone>() {}
|
||
assert_clone::<BotContext>();
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn effective_project_root_standalone_returns_project_root() {
|
||
let services = test_services(PathBuf::from("/projects/myapp"));
|
||
let ctx = test_bot_context(services, None, vec![], BTreeMap::new());
|
||
assert_eq!(
|
||
ctx.effective_project_root().await,
|
||
PathBuf::from("/projects/myapp")
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn effective_project_root_gateway_uses_active_project_subdir() {
|
||
let services = test_services(PathBuf::from("/gateway"));
|
||
let active = Arc::new(RwLock::new("huskies".to_string()));
|
||
let ctx = test_bot_context(
|
||
services,
|
||
Some(Arc::clone(&active)),
|
||
vec!["huskies".into(), "robot-studio".into()],
|
||
BTreeMap::from([
|
||
("huskies".into(), "http://localhost:3001".into()),
|
||
("robot-studio".into(), "http://localhost:3002".into()),
|
||
]),
|
||
);
|
||
assert_eq!(
|
||
ctx.effective_project_root().await,
|
||
PathBuf::from("/gateway/huskies")
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn effective_project_root_gateway_reflects_project_switch() {
|
||
let services = test_services(PathBuf::from("/gateway"));
|
||
let active = Arc::new(RwLock::new("huskies".to_string()));
|
||
let ctx = test_bot_context(
|
||
services,
|
||
Some(Arc::clone(&active)),
|
||
vec!["huskies".into(), "robot-studio".into()],
|
||
BTreeMap::from([
|
||
("huskies".into(), "http://localhost:3001".into()),
|
||
("robot-studio".into(), "http://localhost:3002".into()),
|
||
]),
|
||
);
|
||
|
||
assert_eq!(
|
||
ctx.effective_project_root().await,
|
||
PathBuf::from("/gateway/huskies")
|
||
);
|
||
|
||
*active.write().await = "robot-studio".to_string();
|
||
|
||
assert_eq!(
|
||
ctx.effective_project_root().await,
|
||
PathBuf::from("/gateway/robot-studio")
|
||
);
|
||
}
|
||
|
||
// -- SeenEventIds deduplication ----------------------------------------
|
||
|
||
fn make_event_id(s: &str) -> OwnedEventId {
|
||
s.parse().unwrap()
|
||
}
|
||
|
||
/// AC3: the same event_id presented twice is deduplicated — the second
|
||
/// insert returns false, so any downstream handler would execute only once.
|
||
#[test]
|
||
fn insert_same_event_id_twice_returns_false_on_second() {
|
||
let mut seen = SeenEventIds::new(10);
|
||
let id = make_event_id("$event1:example.com");
|
||
assert!(seen.insert(id.clone()), "first insert must be new");
|
||
assert!(!seen.insert(id), "second insert must be a duplicate");
|
||
}
|
||
|
||
/// AC4: two different event_ids with the same body content both return
|
||
/// true — dedupe is keyed strictly on event_id, never on content.
|
||
#[test]
|
||
fn insert_different_event_ids_same_body_both_new() {
|
||
let mut seen = SeenEventIds::new(10);
|
||
let id1 = make_event_id("$event1:example.com");
|
||
let id2 = make_event_id("$event2:example.com");
|
||
assert!(seen.insert(id1), "first event_id must be new");
|
||
assert!(
|
||
seen.insert(id2),
|
||
"second event_id with identical body must also be new"
|
||
);
|
||
}
|
||
|
||
/// The set evicts the oldest entry (FIFO) when the cap is reached so that
|
||
/// subsequent inserts still work and memory stays bounded.
|
||
#[test]
|
||
fn seen_event_ids_evicts_oldest_at_cap() {
|
||
let cap = 3;
|
||
let mut seen = SeenEventIds::new(cap);
|
||
let id0 = make_event_id("$ev0:example.com");
|
||
let id1 = make_event_id("$ev1:example.com");
|
||
let id2 = make_event_id("$ev2:example.com");
|
||
let id3 = make_event_id("$ev3:example.com");
|
||
|
||
assert!(seen.insert(id0.clone())); // deque: [id0]
|
||
assert!(seen.insert(id1.clone())); // deque: [id0, id1]
|
||
assert!(seen.insert(id2.clone())); // deque: [id0, id1, id2] — full
|
||
// Cap reached — inserting id3 evicts id0 (oldest).
|
||
assert!(seen.insert(id3.clone())); // deque: [id1, id2, id3]
|
||
// id0 was evicted, so re-inserting it returns true (treated as new).
|
||
assert!(seen.insert(id0), "evicted entry should be re-insertable");
|
||
// Re-inserting id0 evicts id1 (new oldest). id2 and id3 are still present.
|
||
assert!(!seen.insert(id2.clone()), "id2 still in set");
|
||
assert!(!seen.insert(id3), "id3 still in set");
|
||
// id1 was evicted — treated as new again.
|
||
assert!(
|
||
seen.insert(id1),
|
||
"id1 was evicted and should be re-insertable"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn bot_context_has_no_require_verified_devices_field() {
|
||
let services = test_services(PathBuf::from("/tmp"));
|
||
let ctx = test_bot_context(services, None, vec![], BTreeMap::new());
|
||
let _cloned = ctx.clone();
|
||
}
|
||
|
||
/// A bot command issued in gateway mode must round-trip over WebSocket
|
||
/// (using the `bot.command` RPC method) and must NOT use HTTP transport.
|
||
#[tokio::test]
|
||
async fn proxy_bot_command_uses_websocket_not_http() {
|
||
use futures::{SinkExt, StreamExt};
|
||
use tokio::net::TcpListener;
|
||
use tokio_tungstenite::tungstenite::Message as WsMsg;
|
||
|
||
// Bind an ephemeral port for our mock WebSocket server.
|
||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||
let port = listener.local_addr().unwrap().port();
|
||
|
||
// Spawn a minimal WS server: accept one connection, verify the
|
||
// request uses the `bot.command` RPC method (not HTTP), and reply.
|
||
let server = tokio::spawn(async move {
|
||
let (tcp, _addr) = listener.accept().await.unwrap();
|
||
let mut ws = tokio_tungstenite::accept_async(tcp).await.unwrap();
|
||
while let Some(Ok(msg)) = ws.next().await {
|
||
if let WsMsg::Text(text) = msg {
|
||
let req: serde_json::Value =
|
||
serde_json::from_str(&text).expect("valid JSON from proxy");
|
||
assert_eq!(
|
||
req["kind"], "rpc_request",
|
||
"transport must use rpc_request, not HTTP"
|
||
);
|
||
assert_eq!(req["method"], "bot.command");
|
||
assert_eq!(req["params"]["command"], "status");
|
||
let correlation_id = req["correlation_id"].clone();
|
||
let resp = serde_json::json!({
|
||
"kind": "rpc_response",
|
||
"correlation_id": correlation_id,
|
||
"ok": true,
|
||
"result": { "response": "all systems go" },
|
||
});
|
||
ws.send(WsMsg::Text(resp.to_string().into())).await.unwrap();
|
||
break;
|
||
}
|
||
}
|
||
});
|
||
|
||
let base_url = format!("http://127.0.0.1:{port}");
|
||
let services = test_services(PathBuf::from("/gateway"));
|
||
let active = Arc::new(RwLock::new("huskies".to_string()));
|
||
let ctx = test_bot_context(
|
||
services,
|
||
Some(Arc::clone(&active)),
|
||
vec!["huskies".into()],
|
||
BTreeMap::from([("huskies".into(), base_url)]),
|
||
);
|
||
|
||
let result = ctx.proxy_bot_command("status", "").await;
|
||
assert_eq!(
|
||
result.as_deref(),
|
||
Some("all systems go"),
|
||
"proxy must return the response text from the rpc_response frame"
|
||
);
|
||
|
||
server.await.unwrap();
|
||
}
|
||
}
|