improve assistant behavior when the service is shutting down

signal-gateway/src/assistant/mod.rs

@@ -13,6 +13,7 @@ pub use signal_gateway_assistant::{
 use crate::message_handler::AdminMessageResponse;
 use chrono::{DateTime, Utc};
 use tokio::sync::{mpsc, oneshot};
+use tokio_util::sync::CancellationToken;
 use worker::{AssistantWorker, Input};
 
 /// Size of the request queue for the assistant worker.
@@ -39,8 +40,9 @@ pub enum AssistantError {
 pub struct AssistantAgent {
     input_tx: mpsc::Sender<Input>,
     stop_tx: mpsc::Sender<()>,
+    cancellation_token: CancellationToken,
     #[allow(dead_code)]
-    worker_handle: tokio::task::JoinHandle<()>,
+    worker_handle: Option<tokio::task::JoinHandle<()>>,
 }
 
 impl AssistantAgent {
@@ -50,18 +52,20 @@ impl AssistantAgent {
     pub fn new(assistant: Box<dyn Assistant>) -> Self {
         let (input_tx, input_rx) = mpsc::channel(REQUEST_QUEUE_SIZE);
         let (stop_tx, stop_rx) = mpsc::channel(REQUEST_QUEUE_SIZE);
+        let cancellation_token = CancellationToken::new();
 
-        let worker = AssistantWorker::new(assistant, input_rx, stop_rx);
+        let worker = AssistantWorker::new(assistant, input_rx, stop_rx, cancellation_token.clone());
 
-        let worker_handle = tokio::spawn(async move {
+        let worker_handle = Some(tokio::spawn(async move {
             worker.run().await;
             tracing::info!("Assistant worker task exited");
-        });
+        }));
 
         Self {
             input_tx,
             stop_tx,
             worker_handle,
+            cancellation_token,
         }
     }
 
@@ -111,8 +115,14 @@ impl AssistantAgent {
         let _ = self.input_tx.try_send(Input::Debug);
     }
 
-    /// Request the worker to stop processing and cancel pending requests.
+    /// Request the worker to stop what it's doing and cancel pending requests (but not exit).
     pub fn request_stop(&self) {
         let _ = self.stop_tx.try_send(());
     }
 }
+
+impl Drop for AssistantAgent {
+    fn drop(&mut self) {
+        self.cancellation_token.cancel();
+    }
+}

signal-gateway/src/assistant/worker.rs

@@ -26,7 +26,10 @@ pub enum Input {
 pub struct AssistantWorker {
     assistant: Box<dyn Assistant>,
     input_rx: mpsc::Receiver<Input>,
+    /// Used when the user wants to cancel the current requests, but not shutdown the service
     stop_rx: mpsc::Receiver<()>,
+    /// Used when the user wants to shut down the service
+    cancellation_token: CancellationToken,
 }
 
 impl AssistantWorker {
@@ -35,11 +38,13 @@ impl AssistantWorker {
         assistant: Box<dyn Assistant>,
         input_rx: mpsc::Receiver<Input>,
         stop_rx: mpsc::Receiver<()>,
+        cancellation_token: CancellationToken,
     ) -> Self {
         Self {
             assistant,
             input_rx,
             stop_rx,
+            cancellation_token,
         }
     }
 
@@ -51,22 +56,34 @@ impl AssistantWorker {
                     let Some(input) = input else {
                         break; // Channel closed
                     };
-                    self.handle_input(input).await;
+                    let was_canceled = self.handle_input(input).await;
+                    if was_canceled {
+                        // If the overall cancellation token was canceled,
+                        // then don't even bother calling handle_stop, just exit
+                        if self.cancellation_token.is_cancelled() {
+                            return;
+                        }
+                        self.handle_stop().await;
+                    }
                 }
                 _ = self.stop_rx.recv() => {
                     self.handle_stop().await;
                 }
+                _ = self.cancellation_token.cancelled() => {
+                    return;
+                }
             }
         }
     }
 
-    async fn handle_input(&mut self, input: Input) {
+    async fn handle_input(&mut self, input: Input) -> bool {
         match input {
             Input::Prompt(msg, sender) => {
-                // Make a cancel token for each new request
+                // Make a child cancel token for each new request
-                let cancel_token = CancellationToken::new();
+                let cancel_token = self.cancellation_token.child_token();
 
                 let mut assistant_fut = self.assistant.prompt(msg, cancel_token.clone());
+                let mut was_canceled = false;
 
                 // If the assistant finishes normally, return its result.
                 // If we get a stop request, cancel the token, then wait for assistant to finish.
@@ -76,6 +93,7 @@ impl AssistantWorker {
                     },
                     _ = self.stop_rx.recv() => {
                         cancel_token.cancel();
+                        was_canceled = true;
                         assistant_fut.await
                     }
                 };
@@ -91,15 +109,20 @@ impl AssistantWorker {
                 };
 
                 let _ = sender.send(response);
+
+                was_canceled
             }
             Input::Record(msg) => {
                 self.assistant.record_message(msg).await;
+                false
             }
             Input::Compact => {
                 self.assistant.compact().await;
+                false
             }
             Input::Debug => {
                 self.assistant.debug_log();
+                false
             }
         }
     }