huskies: merge 1240 story Matrix bot shows live progress by editing its placeholder message

2026-07-21 12:17:42 +00:00
parent afde373676
commit c906707a4f
2 changed files with 318 additions and 41 deletions
@@ -1,14 +1,16 @@
 //! Matrix handle_message — runs the LLM turn for a verified incoming message and
 //! streams the assistant reply back to the room.

-use crate::chat::ChatTransport;
 use crate::chat::util::drain_complete_paragraphs;
+use crate::chat::{ChatTransport, MessageId};
 use crate::llm::providers::claude_code::{CANCELLED, ClaudeCodeProvider, ClaudeCodeResult};
 use crate::slog;
-use matrix_sdk::ruma::OwnedRoomId;
+use matrix_sdk::ruma::{OwnedEventId, OwnedRoomId};
+use std::collections::HashSet;
 use std::sync::Arc;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::time::Duration;
+use tokio::sync::Mutex as TokioMutex;
 use tokio::sync::watch;

 use super::super::context::BotContext;
@@ -45,6 +47,85 @@ pub(in crate::chat::transport::matrix::bot) fn spawn_digging_in_watcher(
    })
 }

+/// One live-progress update to apply to the room's placeholder message
+/// while a turn runs (story 1240).
+enum ProgressUpdate {
+    /// A tool is about to run — shown to the user as e.g. "Using Read...".
+    Activity(String),
+    /// The model is thinking, with no tool call yet.
+    Thinking,
+    /// A committed chunk of assistant text — finalizes the current placeholder.
+    Text(String),
+}
+
+/// The provisional state currently displayed in the open placeholder, used
+/// to debounce repeated identical [`ProgressUpdate`]s into a single edit.
+#[derive(Clone, PartialEq, Eq)]
+enum DisplayState {
+    Thinking,
+    Tool(String),
+}
+
+/// Drain `updates`, applying each to `room_id` via `transport`.
+///
+/// `current_id` starts as the id of the "Working..." placeholder already
+/// posted by `post_working_notice` (`None` if that send failed). Provisional
+/// updates (`Activity`/`Thinking`) edit the open placeholder in place —
+/// skipped entirely when the state is unchanged from the last one shown, so
+/// a long run of identical signals (e.g. many thinking-token deltas, or the
+/// same tool_use signalled twice via both the real-time and fallback paths)
+/// produces at most one edit (AC2/AC5). `Text` always finalizes: it edits
+/// (or, if none is open, sends) the placeholder with real content and then
+/// clears `current_id`, so the *next* provisional update lazily opens a
+/// fresh placeholder below it (AC3) — and if no further update ever arrives,
+/// nothing extra is ever created, satisfying AC4 without needing a delete
+/// capability the transport doesn't have.
+async fn run_progress_updates(
+    transport: Arc<dyn ChatTransport>,
+    room_id: String,
+    bot_sent_event_ids: Arc<TokioMutex<HashSet<OwnedEventId>>>,
+    mut current_id: Option<MessageId>,
+    mut updates: tokio::sync::mpsc::UnboundedReceiver<ProgressUpdate>,
+) {
+    let mut display_state: Option<DisplayState> = None;
+    while let Some(update) = updates.recv().await {
+        let (text, new_state, finalizes) = match update {
+            ProgressUpdate::Activity(name) => {
+                let state = DisplayState::Tool(name.clone());
+                if display_state.as_ref() == Some(&state) {
+                    continue;
+                }
+                (format!("Using {name}..."), Some(state), false)
+            }
+            ProgressUpdate::Thinking => {
+                if display_state.as_ref() == Some(&DisplayState::Thinking) {
+                    continue;
+                }
+                (
+                    "Thinking...".to_string(),
+                    Some(DisplayState::Thinking),
+                    false,
+                )
+            }
+            ProgressUpdate::Text(text) => (text, None, true),
+        };
+
+        let html = markdown_to_html(&text);
+        if let Some(id) = &current_id {
+            let _ = transport.edit_message(&room_id, id, &text, &html).await;
+        } else if let Ok(msg_id) = transport.send_message(&room_id, &text, &html).await {
+            if let Ok(event_id) = msg_id.parse() {
+                bot_sent_event_ids.lock().await.insert(event_id);
+            }
+            current_id = Some(msg_id);
+        }
+        display_state = new_state;
+        if finalizes {
+            current_id = None;
+        }
+    }
+}
+
 pub(in crate::chat::transport::matrix::bot) async fn handle_message(
    room_id_str: String,
    room_id: OwnedRoomId,
@@ -52,6 +133,7 @@ pub(in crate::chat::transport::matrix::bot) async fn handle_message(
    sender: String,
    user_message: String,
    mut cancel_rx: watch::Receiver<bool>,
+    placeholder_id: Option<MessageId>,
 ) {
    // Look up the room's existing Claude Code session ID (if any) so we can
    // resume the conversation with structured API messages instead of
@@ -103,28 +185,24 @@ pub(in crate::chat::transport::matrix::bot) async fn handle_message(

    let provider = ClaudeCodeProvider::new();

-    // Channel for sending complete paragraphs to the Matrix posting task.
-    let (msg_tx, mut msg_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
-    let msg_tx_for_callback = msg_tx.clone();
+    // Channel for sending live-progress updates to the Matrix posting task.
+    let (progress_tx, progress_rx) = tokio::sync::mpsc::unbounded_channel::<ProgressUpdate>();
+    let progress_tx_for_callback = progress_tx.clone();
+    let progress_tx_for_thinking = progress_tx.clone();
+    let progress_tx_for_activity = progress_tx.clone();

-    // Spawn a task to post messages via the transport as they arrive so we
-    // don't block the LLM stream while waiting for send round-trips.
+    // Spawn a task to apply progress updates via the transport as they arrive
+    // so we don't block the LLM stream while waiting for send/edit round-trips.
    let post_transport = Arc::clone(&ctx.transport);
    let post_room_id = room_id_str.clone();
    let sent_ids = Arc::clone(&ctx.bot_sent_event_ids);
-    let sent_ids_for_post = Arc::clone(&sent_ids);
-    let post_task = tokio::spawn(async move {
-        while let Some(chunk) = msg_rx.recv().await {
-            let html = markdown_to_html(&chunk);
-            if let Ok(msg_id) = post_transport
-                .send_message(&post_room_id, &chunk, &html)
-                .await
-                && let Ok(event_id) = msg_id.parse()
-            {
-                sent_ids_for_post.lock().await.insert(event_id);
-            }
-        }
-    });
+    let post_task = tokio::spawn(run_progress_updates(
+        post_transport,
+        post_room_id,
+        sent_ids,
+        placeholder_id,
+        progress_rx,
+    ));

    // Shared state between the sync token callback and the async outer scope.
    let buffer = Arc::new(std::sync::Mutex::new(String::new()));
@@ -158,11 +236,15 @@ pub(in crate::chat::transport::matrix::bot) async fn handle_message(
            let paragraphs = drain_complete_paragraphs(&mut buf);
            for chunk in paragraphs {
                sent_any_chunk_for_callback.store(true, Ordering::Relaxed);
-                let _ = msg_tx_for_callback.send(chunk);
+                let _ = progress_tx_for_callback.send(ProgressUpdate::Text(chunk));
            }
        },
-        |_thinking| {}, // Discard thinking tokens
-        |_activity| {}, // Discard activity signals
+        move |_thinking| {
+            let _ = progress_tx_for_thinking.send(ProgressUpdate::Thinking);
+        },
+        move |activity| {
+            let _ = progress_tx_for_activity.send(ProgressUpdate::Activity(activity.to_string()));
+        },
    );
    tokio::pin!(chat_fut);

@@ -194,7 +276,7 @@ pub(in crate::chat::transport::matrix::bot) async fn handle_message(
            usage,
        }) => {
            let reply = if !remaining.is_empty() {
-                let _ = msg_tx.send(remaining.clone());
+                let _ = progress_tx.send(ProgressUpdate::Text(remaining.clone()));
                remaining
            } else if !did_send_any {
                // Nothing was streamed at all (e.g. only tool calls with no
@@ -207,7 +289,7 @@ pub(in crate::chat::transport::matrix::bot) async fn handle_message(
                    .map(|m| m.content.clone())
                    .unwrap_or_default();
                if !last_text.is_empty() {
-                    let _ = msg_tx.send(last_text.clone());
+                    let _ = progress_tx.send(ProgressUpdate::Text(last_text.clone()));
                }
                last_text
            } else {
@@ -232,14 +314,14 @@ pub(in crate::chat::transport::matrix::bot) async fn handle_message(
            } else {
                format!("Error processing your request: {e}")
            };
-            let _ = msg_tx.send(err_msg.clone());
+            let _ = progress_tx.send(ProgressUpdate::Text(err_msg.clone()));
            (err_msg, None, None)
        }
    };

-    // Drop the sender to signal the posting task that no more messages will
-    // arrive, then wait for all pending Matrix sends to complete.
-    drop(msg_tx);
+    // Drop the sender to signal the posting task that no more updates will
+    // arrive, then wait for all pending Matrix sends/edits to complete.
+    drop(progress_tx);
    let _ = post_task.await;

    if was_cancelled {
@@ -335,18 +417,24 @@ pub(in crate::chat::transport::matrix::bot) async fn handle_message(
 #[cfg(test)]
 mod tests {
    use super::*;
-    use crate::chat::MessageId;
    use async_trait::async_trait;

-    /// Captures every message sent through it, for assertion in tests.
+    /// Captures every message sent and edit applied through it, for
+    /// assertion in tests. Each `send_message` call returns a fresh,
+    /// distinguishable id (`msg-0`, `msg-1`, ...) so tests can tell which
+    /// placeholder a subsequent edit landed on.
    struct CapturingTransport {
        sent: std::sync::Mutex<Vec<(String, String)>>,
+        edits: std::sync::Mutex<Vec<(String, String, String)>>,
+        next_id: std::sync::atomic::AtomicUsize,
    }

    impl CapturingTransport {
        fn new() -> Self {
            Self {
                sent: std::sync::Mutex::new(Vec::new()),
+                edits: std::sync::Mutex::new(Vec::new()),
+                next_id: std::sync::atomic::AtomicUsize::new(0),
            }
        }

@@ -361,6 +449,20 @@ mod tests {
                .last()
                .map(|(_, plain)| plain.clone())
        }
+
+        fn edit_count(&self) -> usize {
+            self.edits.lock().unwrap().len()
+        }
+
+        /// `(message_id, plain_text)` for every edit, in call order.
+        fn edits(&self) -> Vec<(String, String)> {
+            self.edits
+                .lock()
+                .unwrap()
+                .iter()
+                .map(|(id, plain, _html)| (id.clone(), plain.clone()))
+                .collect()
+        }
    }

    #[async_trait]
@@ -375,16 +477,24 @@ mod tests {
                .lock()
                .unwrap()
                .push((room_id.to_string(), plain.to_string()));
-            Ok("msg-id".to_string())
+            let n = self
+                .next_id
+                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+            Ok(format!("msg-{n}"))
        }

        async fn edit_message(
            &self,
            _room_id: &str,
-            _original_message_id: &str,
-            _plain: &str,
-            _html: &str,
+            original_message_id: &str,
+            plain: &str,
+            html: &str,
        ) -> Result<(), String> {
+            self.edits.lock().unwrap().push((
+                original_message_id.to_string(),
+                plain.to_string(),
+                html.to_string(),
+            ));
            Ok(())
        }

@@ -493,4 +603,164 @@ mod tests {
        second_handle.await.unwrap();
        assert_eq!(transport.sent_count(), 1);
    }
+
+    // ── run_progress_updates (story 1240) ─────────────────────────────────
+
+    fn spawn_progress(
+        transport: Arc<CapturingTransport>,
+        placeholder_id: Option<MessageId>,
+    ) -> (
+        tokio::sync::mpsc::UnboundedSender<ProgressUpdate>,
+        tokio::task::JoinHandle<()>,
+    ) {
+        let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<ProgressUpdate>();
+        let bot_sent_event_ids = Arc::new(TokioMutex::new(HashSet::new()));
+        let handle = tokio::spawn(run_progress_updates(
+            transport as Arc<dyn ChatTransport>,
+            "!room:example.com".to_string(),
+            bot_sent_event_ids,
+            placeholder_id,
+            rx,
+        ));
+        (tx, handle)
+    }
+
+    /// AC 1: an activity signal edits the existing placeholder to name the
+    /// tool, and a thinking signal edits it to a "Thinking..." state.
+    #[tokio::test]
+    async fn activity_and_thinking_edit_the_open_placeholder() {
+        let transport = Arc::new(CapturingTransport::new());
+        let (tx, handle) = spawn_progress(transport.clone(), Some("placeholder-0".to_string()));
+
+        tx.send(ProgressUpdate::Activity("Read".to_string()))
+            .unwrap();
+        tx.send(ProgressUpdate::Thinking).unwrap();
+        drop(tx);
+        handle.await.unwrap();
+
+        assert_eq!(transport.sent_count(), 0, "no new message should be sent");
+        let edits = transport.edits();
+        assert_eq!(edits.len(), 2);
+        assert_eq!(
+            edits[0],
+            ("placeholder-0".to_string(), "Using Read...".to_string())
+        );
+        assert_eq!(
+            edits[1],
+            ("placeholder-0".to_string(), "Thinking...".to_string())
+        );
+    }
+
+    /// AC 2: repeating the same tool-call signal (e.g. the real-time and
+    /// fallback activity paths both firing for one tool_use block) produces
+    /// exactly one edit, not two.
+    #[tokio::test]
+    async fn repeated_identical_activity_produces_one_edit() {
+        let transport = Arc::new(CapturingTransport::new());
+        let (tx, handle) = spawn_progress(transport.clone(), Some("placeholder-0".to_string()));
+
+        tx.send(ProgressUpdate::Activity("Bash".to_string()))
+            .unwrap();
+        tx.send(ProgressUpdate::Activity("Bash".to_string()))
+            .unwrap();
+        tx.send(ProgressUpdate::Activity("Bash".to_string()))
+            .unwrap();
+        drop(tx);
+        handle.await.unwrap();
+
+        assert_eq!(transport.edit_count(), 1, "duplicate signals must debounce");
+    }
+
+    /// AC 3: an intermediate text block finalizes the current placeholder
+    /// (edits it with the real text) and the next provisional update opens a
+    /// fresh placeholder below it rather than reusing the finalized one.
+    #[tokio::test]
+    async fn intermediate_text_finalizes_then_next_update_opens_fresh_placeholder() {
+        let transport = Arc::new(CapturingTransport::new());
+        let (tx, handle) = spawn_progress(transport.clone(), Some("placeholder-0".to_string()));
+
+        tx.send(ProgressUpdate::Text("Here's part one.".to_string()))
+            .unwrap();
+        tx.send(ProgressUpdate::Activity("Edit".to_string()))
+            .unwrap();
+        drop(tx);
+        handle.await.unwrap();
+
+        let edits = transport.edits();
+        assert_eq!(
+            edits[0],
+            ("placeholder-0".to_string(), "Here's part one.".to_string())
+        );
+        assert_eq!(
+            transport.sent_count(),
+            1,
+            "a fresh placeholder must be sent"
+        );
+        assert_eq!(transport.last_message().unwrap(), "Using Edit...");
+        // The subsequent edit (from the AC5-style test below) would need to
+        // land on this newly sent placeholder, not the already-finalized one
+        // — confirmed here since this update was sent, not edited.
+        assert_eq!(edits.len(), 1, "the second update was sent, not edited");
+    }
+
+    /// AC 4: when a text update is the last thing to arrive, it finalizes the
+    /// open placeholder in place and nothing further is ever sent — no
+    /// trailing placeholder is left behind.
+    #[tokio::test]
+    async fn final_text_finalizes_with_no_trailing_placeholder() {
+        let transport = Arc::new(CapturingTransport::new());
+        let (tx, handle) = spawn_progress(transport.clone(), Some("placeholder-0".to_string()));
+
+        tx.send(ProgressUpdate::Activity("Read".to_string()))
+            .unwrap();
+        tx.send(ProgressUpdate::Text("All done.".to_string()))
+            .unwrap();
+        drop(tx);
+        handle.await.unwrap();
+
+        assert_eq!(transport.sent_count(), 0, "no trailing placeholder");
+        let edits = transport.edits();
+        assert_eq!(edits.last().unwrap().1, "All done.");
+    }
+
+    /// AC 5: a representative turn (tool call, duplicate tool signal,
+    /// thinking burst, another tool call, an intermediate text block, one
+    /// more tool call, final text) produces roughly five edits — not
+    /// hundreds — even though several of those updates arrive many times.
+    #[tokio::test]
+    async fn representative_turn_produces_roughly_five_edits() {
+        let transport = Arc::new(CapturingTransport::new());
+        let (tx, handle) = spawn_progress(transport.clone(), Some("placeholder-0".to_string()));
+
+        tx.send(ProgressUpdate::Activity("Read".to_string()))
+            .unwrap();
+        tx.send(ProgressUpdate::Activity("Read".to_string()))
+            .unwrap(); // fallback duplicate of the same tool_use — must debounce
+        for _ in 0..50 {
+            tx.send(ProgressUpdate::Thinking).unwrap(); // a burst of thinking-token deltas
+        }
+        tx.send(ProgressUpdate::Activity("Bash".to_string()))
+            .unwrap();
+        tx.send(ProgressUpdate::Text(
+            "Here's what I found so far.".to_string(),
+        ))
+        .unwrap();
+        tx.send(ProgressUpdate::Activity("Edit".to_string()))
+            .unwrap();
+        tx.send(ProgressUpdate::Text("Done.".to_string())).unwrap();
+        drop(tx);
+        handle.await.unwrap();
+
+        assert_eq!(
+            transport.edit_count(),
+            5,
+            "one edit per distinct state change, regardless of how many \
+             identical updates arrived in between"
+        );
+        assert_eq!(
+            transport.sent_count(),
+            1,
+            "one fresh placeholder for the second half"
+        );
+    }
 }
@@ -398,16 +398,20 @@ fn parse_question_reply(
 /// — never once per resulting agent turn — so a burst of messages the
 /// dispatcher later coalesces into a single turn still yields one notice per
 /// message the user actually sent (story 1239).
-async fn post_working_notice(ctx: &BotContext, room_id_str: &str) {
+async fn post_working_notice(
+    ctx: &BotContext,
+    room_id_str: &str,
+) -> Option<crate::chat::MessageId> {
    let html = markdown_to_html(handle_message::DIGGING_IN_MESSAGE);
-    if let Ok(msg_id) = ctx
+    let msg_id = ctx
        .transport
        .send_message(room_id_str, handle_message::DIGGING_IN_MESSAGE, &html)
        .await
-        && let Ok(event_id) = msg_id.parse()
-    {
+        .ok()?;
+    if let Ok(event_id) = msg_id.parse() {
        ctx.bot_sent_event_ids.lock().await.insert(event_id);
    }
+    Some(msg_id)
 }

 pub(in crate::chat::transport::matrix::bot) async fn on_room_message(
@@ -1586,8 +1590,9 @@ pub(in crate::chat::transport::matrix::bot) async fn on_room_message(
    }

    // Acknowledge receipt immediately, before the message is handed to the
-    // dispatcher below (story 1239).
-    post_working_notice(&ctx, &room_id_str).await;
+    // dispatcher below (story 1239). Its message id seeds the live-progress
+    // placeholder that handle_message edits in place as the turn runs (story 1240).
+    let placeholder_id = post_working_notice(&ctx, &room_id_str).await;

    // Hand the message to the protocol-agnostic dispatcher instead of spawning
    // directly.  The dispatcher applies a coalesce window and a per-session
@@ -1601,6 +1606,7 @@ pub(in crate::chat::transport::matrix::bot) async fn on_room_message(
                let incoming_room_id = incoming_room_id.clone();
                let ctx = ctx_for_factory.clone();
                let sender = sender.clone();
+                let placeholder_id = placeholder_id.clone();
                Box::pin(async move {
                    handle_message(
                        room_id_str,
@@ -1609,6 +1615,7 @@ pub(in crate::chat::transport::matrix::bot) async fn on_room_message(
                        sender,
                        coalesced,
                        cancel_rx,
+                        placeholder_id,
                    )
                    .await;
                })