pass child token into assistant agent

This commit is contained in:
Chris Beck
2025-12-15 12:47:45 -07:00
parent 2670104b49
commit 51d3ee1a10
2 changed files with 3 additions and 3 deletions
+1 -2
View File
@@ -54,10 +54,9 @@ impl AssistantAgent {
/// Create a new assistant agent with the given assistant implementation. /// Create a new assistant agent with the given assistant implementation.
/// ///
/// Spawns a background worker task that processes requests serially. /// Spawns a background worker task that processes requests serially.
pub fn new(assistant: Box<dyn Assistant>) -> Self { pub fn new(assistant: Box<dyn Assistant>, cancellation_token: CancellationToken) -> Self {
let (input_tx, input_rx) = mpsc::channel(REQUEST_QUEUE_SIZE); let (input_tx, input_rx) = mpsc::channel(REQUEST_QUEUE_SIZE);
let (stop_tx, stop_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()); let worker = AssistantWorker::new(assistant, input_rx, stop_rx, cancellation_token.clone());
+2 -1
View File
@@ -343,6 +343,7 @@ impl Gateway {
extra_tool_executors: Vec<Arc<dyn ToolExecutor>>, extra_tool_executors: Vec<Arc<dyn ToolExecutor>>,
assistant_factory: Option<AssistantFactory>, assistant_factory: Option<AssistantFactory>,
) -> Arc<Self> { ) -> Arc<Self> {
let child_token = token.child_token();
let (signal_alert_mq_tx, signal_alert_mq_rx) = unbounded_channel(); let (signal_alert_mq_tx, signal_alert_mq_rx) = unbounded_channel();
let prometheus = config let prometheus = config
@@ -369,7 +370,7 @@ impl Gateway {
// Initialize the assistant agent using the factory if one was provided // Initialize the assistant agent using the factory if one was provided
if let Some(factory) = assistant_factory { if let Some(factory) = assistant_factory {
let assistant = factory(Arc::downgrade(&gateway) as Weak<dyn ToolExecutor>); let assistant = factory(Arc::downgrade(&gateway) as Weak<dyn ToolExecutor>);
let agent = AssistantAgent::new(assistant); let agent = AssistantAgent::new(assistant, child_token);
gateway gateway
.assistant .assistant
.set(agent) .set(agent)