diff --git a/signal-gateway/src/assistant/mod.rs b/signal-gateway/src/assistant/mod.rs
index 30deab3..ee7f8ce 100644
--- a/signal-gateway/src/assistant/mod.rs
+++ b/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,
stop_tx: mpsc::Sender<()>,
+ cancellation_token: CancellationToken,
#[allow(dead_code)]
- worker_handle: tokio::task::JoinHandle<()>,
+ worker_handle: Option>,
}
impl AssistantAgent {
@@ -50,18 +52,20 @@ impl AssistantAgent {
pub fn new(assistant: Box) -> 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();
+ }
+}
diff --git a/signal-gateway/src/assistant/worker.rs b/signal-gateway/src/assistant/worker.rs
index 2276328..1cb237a 100644
--- a/signal-gateway/src/assistant/worker.rs
+++ b/signal-gateway/src/assistant/worker.rs
@@ -26,7 +26,10 @@ pub enum Input {
pub struct AssistantWorker {
assistant: Box,
input_rx: mpsc::Receiver,
+ /// 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,
input_rx: mpsc::Receiver,
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
- let cancel_token = CancellationToken::new();
+ // Make a child cancel token for each new request
+ 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
}
}
}