diff --git a/signal-gateway/src/assistant/mod.rs b/signal-gateway/src/assistant/mod.rs index 482ea2f..52b7128 100644 --- a/signal-gateway/src/assistant/mod.rs +++ b/signal-gateway/src/assistant/mod.rs @@ -54,10 +54,9 @@ impl AssistantAgent { /// Create a new assistant agent with the given assistant implementation. /// /// Spawns a background worker task that processes requests serially. - pub fn new(assistant: Box) -> Self { + pub fn new(assistant: Box, cancellation_token: CancellationToken) -> 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, cancellation_token.clone()); diff --git a/signal-gateway/src/gateway/mod.rs b/signal-gateway/src/gateway/mod.rs index b83386d..36f6351 100644 --- a/signal-gateway/src/gateway/mod.rs +++ b/signal-gateway/src/gateway/mod.rs @@ -343,6 +343,7 @@ impl Gateway { extra_tool_executors: Vec>, assistant_factory: Option, ) -> Arc { + let child_token = token.child_token(); let (signal_alert_mq_tx, signal_alert_mq_rx) = unbounded_channel(); let prometheus = config @@ -369,7 +370,7 @@ impl Gateway { // Initialize the assistant agent using the factory if one was provided if let Some(factory) = assistant_factory { let assistant = factory(Arc::downgrade(&gateway) as Weak); - let agent = AssistantAgent::new(assistant); + let agent = AssistantAgent::new(assistant, child_token); gateway .assistant .set(agent)