Fix some issues with UI getting out of sync. (#7)

Reviewed-on: #7
Co-authored-by: Drew Galbraith <drew@tiramisu.one>
Co-committed-by: Drew Galbraith <drew@tiramisu.one>

src/tui/events.rs

@@ -4,7 +4,7 @@ use tokio::sync::mpsc;
 use tracing::debug;
 
 use super::AppState;
-use crate::core::types::{Role, UIEvent};
+use crate::core::types::{Role, StampedEvent, UIEvent};
 
 /// Drain all pending [`UIEvent`]s from `event_rx` and apply them to `state`.
 ///
@@ -19,15 +19,25 @@ use crate::core::types::{Role, UIEvent};
 /// | `ToolResult`           | Display tool result                                        |
 /// | `TurnComplete`         | No structural change; logged at debug level                |
 /// | `Error(msg)`           | Push `(Assistant, "[error] {msg}")`                        |
-pub(super) fn drain_ui_events(event_rx: &mut mpsc::Receiver<UIEvent>, state: &mut AppState) {
-    while let Ok(event) = event_rx.try_recv() {
-        match event {
+pub(super) fn drain_ui_events(event_rx: &mut mpsc::Receiver<StampedEvent>, state: &mut AppState) {
+    while let Ok(stamped) = event_rx.try_recv() {
+        // Discard events from before the most recent :clear.
+        if stamped.epoch < state.epoch {
+            debug!(
+                event_epoch = stamped.epoch,
+                state_epoch = state.epoch,
+                "dropping stale event"
+            );
+            continue;
+        }
+        match stamped.event {
             UIEvent::StreamDelta(chunk) => {
                 if let Some((Role::Assistant, content)) = state.messages.last_mut() {
                     content.push_str(&chunk);
                 } else {
                     state.messages.push((Role::Assistant, chunk));
                 }
+                state.content_changed = true;
             }
             UIEvent::ToolApprovalRequest {
                 tool_use_id,
@@ -47,6 +57,7 @@ pub(super) fn drain_ui_events(event_rx: &mut mpsc::Receiver<UIEvent>, state: &mu
                 state
                     .messages
                     .push((Role::Assistant, format!("[{tool_name}] {input_summary}")));
+                state.content_changed = true;
             }
             UIEvent::ToolResult {
                 tool_name,
@@ -58,6 +69,7 @@ pub(super) fn drain_ui_events(event_rx: &mut mpsc::Receiver<UIEvent>, state: &mu
                     Role::Assistant,
                     format!("[{tool_name} {prefix}] {output_summary}"),
                 ));
+                state.content_changed = true;
             }
             UIEvent::TurnComplete => {
                 debug!("turn complete");
@@ -69,6 +81,7 @@ pub(super) fn drain_ui_events(event_rx: &mut mpsc::Receiver<UIEvent>, state: &mu
                 state
                     .messages
                     .push((Role::Assistant, format!("[error] {msg}")));
+                state.content_changed = true;
             }
         }
     }
@@ -86,12 +99,17 @@ pub struct PendingApproval {
 mod tests {
     use super::*;
 
+    /// Wrap a [`UIEvent`] in a [`StampedEvent`] at epoch 0 for tests.
+    fn stamp(event: UIEvent) -> StampedEvent {
+        StampedEvent { epoch: 0, event }
+    }
+
     #[tokio::test]
     async fn drain_appends_to_existing_assistant() {
         let (tx, mut rx) = tokio::sync::mpsc::channel(8);
         let mut state = AppState::new();
         state.messages.push((Role::Assistant, "hello".to_string()));
-        tx.send(UIEvent::StreamDelta(" world".to_string()))
+        tx.send(stamp(UIEvent::StreamDelta(" world".to_string())))
             .await
             .unwrap();
         drop(tx);
@@ -104,7 +122,7 @@ mod tests {
         let (tx, mut rx) = tokio::sync::mpsc::channel(8);
         let mut state = AppState::new();
         state.messages.push((Role::User, "hi".to_string()));
-        tx.send(UIEvent::StreamDelta("hello".to_string()))
+        tx.send(stamp(UIEvent::StreamDelta("hello".to_string())))
             .await
             .unwrap();
         drop(tx);
@@ -118,11 +136,11 @@ mod tests {
     async fn drain_tool_approval_sets_pending() {
         let (tx, mut rx) = tokio::sync::mpsc::channel(8);
         let mut state = AppState::new();
-        tx.send(UIEvent::ToolApprovalRequest {
+        tx.send(stamp(UIEvent::ToolApprovalRequest {
             tool_use_id: "t1".to_string(),
             tool_name: "write_file".to_string(),
             input_summary: "path: foo.txt".to_string(),
-        })
+        }))
         .await
         .unwrap();
         drop(tx);
@@ -136,11 +154,11 @@ mod tests {
     async fn drain_tool_result_adds_message() {
         let (tx, mut rx) = tokio::sync::mpsc::channel(8);
         let mut state = AppState::new();
-        tx.send(UIEvent::ToolResult {
+        tx.send(stamp(UIEvent::ToolResult {
             tool_name: "read_file".to_string(),
             output_summary: "file contents...".to_string(),
             is_error: false,
-        })
+        }))
         .await
         .unwrap();
         drop(tx);
@@ -148,4 +166,42 @@ mod tests {
         assert_eq!(state.messages.len(), 1);
         assert!(state.messages[0].1.contains("read_file result"));
     }
+
+    #[tokio::test]
+    async fn drain_discards_stale_epoch_events() {
+        let (tx, mut rx) = tokio::sync::mpsc::channel(8);
+        let mut state = AppState::new();
+        state.epoch = 2;
+        // Event from epoch 1 should be discarded.
+        tx.send(StampedEvent {
+            epoch: 1,
+            event: UIEvent::StreamDelta("ghost".to_string()),
+        })
+        .await
+        .unwrap();
+        // Event from epoch 2 should be accepted.
+        tx.send(StampedEvent {
+            epoch: 2,
+            event: UIEvent::StreamDelta("real".to_string()),
+        })
+        .await
+        .unwrap();
+        drop(tx);
+        drain_ui_events(&mut rx, &mut state);
+        assert_eq!(state.messages.len(), 1);
+        assert_eq!(state.messages[0].1, "real");
+    }
+
+    #[tokio::test]
+    async fn drain_sets_content_changed() {
+        let (tx, mut rx) = tokio::sync::mpsc::channel(8);
+        let mut state = AppState::new();
+        assert!(!state.content_changed);
+        tx.send(stamp(UIEvent::StreamDelta("hi".to_string())))
+            .await
+            .unwrap();
+        drop(tx);
+        drain_ui_events(&mut rx, &mut state);
+        assert!(state.content_changed);
+    }
 }