ai_crew_sync/tools/
events.rs

1use std::time::Duration;
2
3use rmcp::{
4    ErrorData, Json, handler::server::wrapper::Parameters, service::RequestContext, tool,
5    tool_router,
6};
7use schemars::JsonSchema;
8use serde::Deserialize;
9use sqlx::PgPool;
10use uuid::Uuid;
11
12use super::{Bus, auth_of};
13use crate::{
14    auth::AuthCtx,
15    events::BusEvent,
16    model::{WaitEvent, WaitResult},
17};
18
19const DEFAULT_TIMEOUT_SECS: i64 = 25;
20const MAX_TIMEOUT_SECS: i64 = 55;
21
22#[derive(Debug, Deserialize, JsonSchema)]
23pub struct WaitArgs {
24    /// How long to wait before giving up, in seconds (5-55, default 25).
25    /// Kept under a minute so HTTP intermediaries do not cut the call.
26    #[serde(default)]
27    pub timeout_seconds: Option<i64>,
28    /// Restrict to certain event kinds: any of "message", "task", "lock",
29    /// "note". Omit to wake on anything relevant to you.
30    #[serde(default)]
31    pub kinds: Option<Vec<String>>,
32}
33
34async fn unread_dms(pool: &PgPool, agent_id: Uuid) -> Result<i64, sqlx::Error> {
35    sqlx::query_scalar(
36        r#"
37        SELECT count(*)
38        FROM messages m
39        LEFT JOIN read_cursors c ON c.agent_id = $1 AND c.scope = 'inbox'
40        WHERE m.recipient_agent_id = $1
41          AND m.id > COALESCE(c.last_message_id, 0)
42        "#,
43    )
44    .bind(agent_id)
45    .fetch_one(pool)
46    .await
47}
48
49async fn unread_anything(pool: &PgPool, auth: &AuthCtx) -> Result<i64, sqlx::Error> {
50    sqlx::query_scalar(
51        r#"
52        SELECT count(*)
53        FROM messages m
54        LEFT JOIN read_cursors c ON c.agent_id = $1 AND c.scope = 'all'
55        WHERE m.team_id = $2
56          AND m.id > COALESCE(c.last_message_id, 0)
57          AND m.sender_agent_id <> $1
58          AND (m.channel_id IS NOT NULL OR m.recipient_agent_id = $1)
59        "#,
60    )
61    .bind(auth.agent_id)
62    .bind(auth.team_id)
63    .fetch_one(pool)
64    .await
65}
66
67/// Turn a raw bus event into a one-line summary, resolving ids to names.
68async fn describe(pool: &PgPool, event: &BusEvent) -> Option<WaitEvent> {
69    match event.kind() {
70        "message" => {
71            let id = event.message_id()?;
72            let row: (String, Option<String>, Option<String>, String) =
73                sqlx::query_as::<_, (String, Option<String>, Option<String>, String)>(
74                    r#"
75                    SELECT s.name, ch.name, r.name, left(m.body, 200)
76                    FROM messages m
77                    JOIN agents s ON s.id = m.sender_agent_id
78                    LEFT JOIN channels ch ON ch.id = m.channel_id
79                    LEFT JOIN agents r ON r.id = m.recipient_agent_id
80                    WHERE m.id = $1
81                    "#,
82                )
83                .bind(id)
84                .fetch_optional(pool)
85                .await
86                .ok()
87                .flatten()?;
88            let (sender, channel, recipient, body) = row;
89            let target = channel
90                .map(|c| format!("#{c}"))
91                .or(recipient.map(|r| format!("@{r}")))
92                .unwrap_or_default();
93            Some(WaitEvent {
94                kind: "message".into(),
95                summary: format!("{sender} → {target}: {body}"),
96            })
97        }
98        "task" => {
99            let key = event.0.get("key").and_then(|v| v.as_str())?;
100            let status = event
101                .0
102                .get("status")
103                .and_then(|v| v.as_str())
104                .unwrap_or("?");
105            let holder = match event.0.get("claimed_by").and_then(|v| v.as_str()) {
106                Some(uuid) => {
107                    sqlx::query_scalar::<_, String>("SELECT name FROM agents WHERE id = $1::uuid")
108                        .bind(uuid)
109                        .fetch_optional(pool)
110                        .await
111                        .ok()
112                        .flatten()
113                        .map(|n| format!(" by {n}"))
114                        .unwrap_or_default()
115                }
116                None => String::new(),
117            };
118            Some(WaitEvent {
119                kind: "task".into(),
120                summary: format!("task '{key}' is now {status}{holder}"),
121            })
122        }
123        "lock" => {
124            let name = event.0.get("name").and_then(|v| v.as_str())?;
125            let what = event
126                .0
127                .get("event")
128                .and_then(|v| v.as_str())
129                .unwrap_or("changed");
130            Some(WaitEvent {
131                kind: "lock".into(),
132                summary: format!("lock '{name}' {what}"),
133            })
134        }
135        "note" => {
136            let scope = event
137                .0
138                .get("scope")
139                .and_then(|v| v.as_str())
140                .unwrap_or("global");
141            let key = event.0.get("key").and_then(|v| v.as_str())?;
142            Some(WaitEvent {
143                kind: "note".into(),
144                summary: format!("note {scope}/{key} was updated"),
145            })
146        }
147        _ => None,
148    }
149}
150
151fn suggestion_for(events: &[WaitEvent], unread: i64) -> String {
152    if events.iter().any(|e| e.kind == "message") || unread > 0 {
153        "Call read_messages to fetch the new messages.".into()
154    } else if events.iter().any(|e| e.kind == "task") {
155        "Call list_tasks (or get_task) to see what changed.".into()
156    } else if events.iter().any(|e| e.kind == "lock") {
157        "Call list_locks (or retry acquire_lock) now.".into()
158    } else if events.iter().any(|e| e.kind == "note") {
159        "Call get_note to read the updated note.".into()
160    } else {
161        "Nothing happened; do other work or wait again.".into()
162    }
163}
164
165#[tool_router(router = events_router, vis = "pub")]
166impl Bus {
167    #[tool(
168        description = "Block until something happens on the bus that concerns you (a message \
169                       arrives, a task changes state, a lock is released, a note is updated) or \
170                       the timeout elapses. Use this instead of polling read_messages in a loop: \
171                       call it when you are waiting on teammates and idle. Returns immediately \
172                       if you already have unread messages."
173    )]
174    async fn wait_for_updates(
175        &self,
176        ctx: RequestContext<rmcp::RoleServer>,
177        Parameters(args): Parameters<WaitArgs>,
178    ) -> Result<Json<WaitResult>, ErrorData> {
179        let auth = auth_of(&ctx)?;
180        let timeout = args
181            .timeout_seconds
182            .unwrap_or(DEFAULT_TIMEOUT_SECS)
183            .clamp(5, MAX_TIMEOUT_SECS);
184        let kind_filter: Option<Vec<String>> = args
185            .kinds
186            .map(|ks| ks.into_iter().map(|k| k.trim().to_lowercase()).collect());
187        let wants = |kind: &str| {
188            kind_filter
189                .as_ref()
190                .map(|ks| ks.iter().any(|k| k == kind))
191                .unwrap_or(true)
192        };
193
194        // Subscribe before checking the database so nothing slips between the
195        // check and the wait.
196        let mut rx = self.hub.subscribe();
197
198        let pending = unread_anything(&self.db, &auth).await.map_err(|e| {
199            tracing::error!(error = %e, "unread check failed");
200            ErrorData::internal_error("database error", None)
201        })?;
202        if pending > 0 && wants("message") {
203            let dms = unread_dms(&self.db, auth.agent_id).await.unwrap_or(0);
204            return Ok(Json(WaitResult {
205                woke: true,
206                timed_out: false,
207                events: vec![WaitEvent {
208                    kind: "message".into(),
209                    summary: format!("{pending} unread message(s) already waiting"),
210                }],
211                unread_direct_messages: dms,
212                suggestion: "Call read_messages to fetch them.".into(),
213            }));
214        }
215
216        let deadline = tokio::time::Instant::now() + Duration::from_secs(timeout as u64);
217        let mut events: Vec<WaitEvent> = Vec::new();
218
219        while events.is_empty() {
220            let event = tokio::select! {
221                _ = tokio::time::sleep_until(deadline) => break,
222                recv = rx.recv() => match recv {
223                    Ok(ev) => ev,
224                    // Lagged: we missed events; report a generic wake so the
225                    // caller re-syncs from the database.
226                    Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
227                        events.push(WaitEvent {
228                            kind: "unknown".into(),
229                            summary: "event stream lagged; re-check the bus".into(),
230                        });
231                        break;
232                    }
233                    Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
234                },
235            };
236
237            if !event.visible_to(auth.team_id, auth.agent_id) || !wants(event.kind()) {
238                continue;
239            }
240            // Your own messages are not news to you.
241            if event.kind() == "message" && event.sender_agent_id() == Some(auth.agent_id) {
242                continue;
243            }
244            if let Some(described) = describe(&self.db, &event).await {
245                events.push(described);
246                // Grace window: batch events that arrive together.
247                let grace = tokio::time::Instant::now() + Duration::from_millis(150);
248                while let Ok(Ok(more)) = tokio::time::timeout_at(grace, rx.recv()).await {
249                    if more.visible_to(auth.team_id, auth.agent_id)
250                        && wants(more.kind())
251                        && !(more.kind() == "message"
252                            && more.sender_agent_id() == Some(auth.agent_id))
253                        && let Some(d) = describe(&self.db, &more).await
254                    {
255                        events.push(d);
256                    }
257                }
258            }
259        }
260
261        let unread = unread_dms(&self.db, auth.agent_id).await.unwrap_or(0);
262        let woke = !events.is_empty();
263        Ok(Json(WaitResult {
264            woke,
265            timed_out: !woke,
266            suggestion: suggestion_for(&events, unread),
267            events,
268            unread_direct_messages: unread,
269        }))
270    }
271}
ai_crew_sync/tools/events.rs

ai_crew_sync/tools/
events.rs
