huskies: merge 1213 story Chat "stop" command that immediately aborts the in-flight LLM turn

server/src/chat/dispatcher.rs

@@ -19,6 +19,7 @@
 use crate::slog;
 use std::collections::HashMap;
 use std::pin::Pin;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
 use tokio::sync::{mpsc, watch};
@@ -47,6 +48,20 @@ enum SessionMsg {
 
 struct SessionHandle {
     tx: mpsc::UnboundedSender<SessionMsg>,
+    /// `true` whenever the session task is coalescing, running, or draining a
+    /// batch — i.e. there is something for [`ChatDispatcher::stop`] to cancel.
+    /// `false` while the task is idle in Phase 1, waiting for a message.
+    active: Arc<AtomicBool>,
+}
+
+/// Result of a [`ChatDispatcher::stop`] call.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum StopOutcome {
+    /// A turn was actively coalescing, running, or draining pending messages
+    /// for the session, and a cancellation was sent.
+    Cancelled,
+    /// The session doesn't exist, or exists but has no turn in flight.
+    NothingRunning,
 }
 
 /// Coalescing, serialising dispatcher for chat-to-LLM message routing.
@@ -79,8 +94,14 @@ impl ChatDispatcher {
         let coalesce_ms = self.coalesce_ms;
         let handle = guard.entry(session_key.clone()).or_insert_with(|| {
             let (tx, rx) = mpsc::unbounded_channel();
-            tokio::spawn(session_task(session_key.clone(), rx, coalesce_ms));
-            SessionHandle { tx }
+            let active = Arc::new(AtomicBool::new(false));
+            tokio::spawn(session_task(
+                session_key.clone(),
+                rx,
+                coalesce_ms,
+                Arc::clone(&active),
+            ));
+            SessionHandle { tx, active }
         });
         let _ = handle.tx.send(SessionMsg::UserMessage {
             text: message,
@@ -90,15 +111,19 @@ impl ChatDispatcher {
 
     /// Stop the active LLM run for `session_key` and clear its pending queue.
     ///
-    /// Returns `true` if the session existed (whether or not anything was
-    /// actually running), `false` if no session for that key has been created.
-    pub fn stop(&self, session_key: &str) -> bool {
+    /// Returns [`StopOutcome::Cancelled`] and sends the cancellation only when
+    /// a turn is actually coalescing, running, or draining pending messages;
+    /// otherwise returns [`StopOutcome::NothingRunning`] without sending
+    /// anything (covers both "no session was ever created" and "session
+    /// exists but is idle").
+    pub fn stop(&self, session_key: &str) -> StopOutcome {
         let guard = self.sessions.lock().unwrap();
-        if let Some(handle) = guard.get(session_key) {
-            let _ = handle.tx.send(SessionMsg::Stop);
-            true
-        } else {
-            false
+        match guard.get(session_key) {
+            Some(handle) if handle.active.load(Ordering::SeqCst) => {
+                let _ = handle.tx.send(SessionMsg::Stop);
+                StopOutcome::Cancelled
+            }
+            _ => StopOutcome::NothingRunning,
         }
     }
 }
@@ -119,6 +144,7 @@ async fn session_task(
     session_key: String,
     mut rx: mpsc::UnboundedReceiver<SessionMsg>,
     coalesce_ms: u64,
+    active: Arc<AtomicBool>,
 ) {
     let coalesce_dur = Duration::from_millis(coalesce_ms);
 
@@ -131,6 +157,9 @@ async fn session_task(
                 Some(SessionMsg::UserMessage { text, factory }) => break (text, factory),
             }
         };
+        // From here until we loop back to Phase 1, there is something in
+        // flight (coalescing, running, or draining) for `stop()` to cancel.
+        active.store(true, Ordering::SeqCst);
 
         // ── Phase 2: coalesce window (debounce) ──────────────────────────────
         let mut batch: Vec<String> = vec![first_text];
@@ -160,6 +189,7 @@ async fn session_task(
         }
 
         if batch.is_empty() {
+            active.store(false, Ordering::SeqCst);
             continue; // Stop received during coalesce — restart
         }
 
@@ -209,6 +239,7 @@ async fn session_task(
             }
 
             if stopped || pending_texts.is_empty() {
+                active.store(false, Ordering::SeqCst);
                 break; // back to Phase 1
             }
 
@@ -351,7 +382,12 @@ mod tests {
         );
 
         // Stop immediately.
-        dispatcher.stop(&session);
+        let outcome = dispatcher.stop(&session);
+        assert_eq!(
+            outcome,
+            StopOutcome::Cancelled,
+            "a run was active, so stop must report Cancelled"
+        );
 
         // Wait longer than the run would have taken if not stopped.
         tokio::time::sleep(Duration::from_millis(700)).await;
@@ -364,4 +400,35 @@ mod tests {
             "stop should discard pending; got {count} spawns"
         );
     }
+
+    /// Stopping a session key that was never submitted reports NothingRunning
+    /// rather than silently no-opping.
+    #[tokio::test]
+    async fn stop_on_unknown_session_reports_nothing_running() {
+        let dispatcher = ChatDispatcher::new(30);
+        assert_eq!(
+            dispatcher.stop("never-seen-room"),
+            StopOutcome::NothingRunning
+        );
+    }
+
+    /// Stopping a session that exists but has no active turn (its one run
+    /// already completed) reports NothingRunning, not Cancelled.
+    #[tokio::test]
+    async fn stop_on_idle_session_reports_nothing_running() {
+        let spawn_count = Arc::new(AtomicUsize::new(0));
+        let dispatcher = ChatDispatcher::new(20);
+        let session = "room4".to_string();
+
+        dispatcher.submit(
+            session.clone(),
+            "hello".to_string(),
+            make_factory(Arc::clone(&spawn_count), 20),
+        );
+
+        // Wait long enough for the coalesce window and the run to finish.
+        tokio::time::sleep(Duration::from_millis(200)).await;
+
+        assert_eq!(dispatcher.stop(&session), StopOutcome::NothingRunning);
+    }
 }