agentkit_task_manager/
lib.rs

1use std::collections::{BTreeMap, VecDeque};
2use std::sync::Arc;
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::time::Duration;
5
6use agentkit_core::{
7    Item, MetadataMap, TaskId, ToolCallId, ToolResultPart, TurnCancellation, TurnId,
8};
9use agentkit_tools_core::{
10    ApprovalRequest, OwnedToolContext, ToolError, ToolExecutionOutcome, ToolExecutor, ToolRequest,
11};
12use async_trait::async_trait;
13use thiserror::Error;
14use tokio::sync::{Mutex, Notify, mpsc};
15use tokio::task::JoinHandle;
16
17pub const TOOL_RESULT_FAILURE_KIND_METADATA_KEY: &str = "agentkit.tool.failure_kind";
18pub const TOOL_RESULT_FAILURE_KIND_PERMISSION_DENIED: &str = "permission_denied";
19
20#[derive(Clone, Copy, Debug, PartialEq, Eq)]
21pub enum TaskKind {
22    Foreground,
23    Background,
24}
25
26#[derive(Clone, Copy, Debug, PartialEq, Eq)]
27pub enum ContinuePolicy {
28    NotifyOnly,
29    RequestContinue,
30}
31
32#[derive(Clone, Copy, Debug, PartialEq, Eq)]
33pub enum DeliveryMode {
34    ToLoop,
35    Manual,
36}
37
38#[derive(Clone, Debug, PartialEq, Eq)]
39pub struct TaskSnapshot {
40    pub id: TaskId,
41    pub turn_id: TurnId,
42    pub call_id: ToolCallId,
43    pub tool_name: String,
44    pub kind: TaskKind,
45    pub metadata: MetadataMap,
46}
47
48#[derive(Clone, Debug, PartialEq)]
49pub enum TaskEvent {
50    Started(TaskSnapshot),
51    Detached(TaskSnapshot),
52    Completed(TaskSnapshot, ToolResultPart),
53    Cancelled(TaskSnapshot),
54    Failed(TaskSnapshot, ToolError),
55    ContinueRequested,
56}
57
58#[derive(Clone, Debug, PartialEq)]
59pub struct TaskApproval {
60    pub task_id: TaskId,
61    pub tool_request: ToolRequest,
62    pub approval: ApprovalRequest,
63}
64
65#[derive(Clone, Debug, PartialEq)]
66pub enum TaskResolution {
67    Item(Item),
68    Approval(TaskApproval),
69}
70
71#[derive(Clone, Debug, PartialEq)]
72pub enum TaskStartOutcome {
73    Ready(Box<TaskResolution>),
74    Pending { task_id: TaskId, kind: TaskKind },
75}
76
77#[derive(Clone, Debug, PartialEq)]
78pub enum TurnTaskUpdate {
79    Resolution(Box<TaskResolution>),
80    Detached(TaskSnapshot),
81}
82
83#[derive(Clone, Debug, Default, PartialEq)]
84pub struct PendingLoopUpdates {
85    pub resolutions: VecDeque<TaskResolution>,
86}
87
88/// How a task should be invoked. Mutually exclusive between plain execution
89/// and resuming after approval.
90#[derive(Clone, Debug, Default)]
91pub enum TaskLaunchKind {
92    /// Execute the tool with the active permission policy.
93    #[default]
94    Plain,
95    /// Re-execute a previously-interrupted call after the user approved it.
96    Approved(ApprovalRequest),
97}
98
99#[derive(Clone, Debug)]
100pub struct TaskLaunchRequest {
101    pub task_id: Option<TaskId>,
102    pub request: ToolRequest,
103    pub kind: TaskLaunchKind,
104}
105
106impl TaskLaunchRequest {
107    /// Plain launch (no prior approval / auth).
108    pub fn plain(task_id: Option<TaskId>, request: ToolRequest) -> Self {
109        Self {
110            task_id,
111            request,
112            kind: TaskLaunchKind::Plain,
113        }
114    }
115
116    /// Resume after the user approved the call.
117    pub fn approved(
118        task_id: Option<TaskId>,
119        request: ToolRequest,
120        approval: ApprovalRequest,
121    ) -> Self {
122        Self {
123            task_id,
124            request,
125            kind: TaskLaunchKind::Approved(approval),
126        }
127    }
128}
129
130#[derive(Clone)]
131pub struct TaskStartContext {
132    pub executor: Arc<dyn ToolExecutor>,
133    pub tool_context: OwnedToolContext,
134}
135
136#[derive(Debug, Error, Clone, PartialEq, Eq)]
137pub enum TaskManagerError {
138    #[error("task not found: {0}")]
139    NotFound(TaskId),
140    #[error("task manager internal error: {0}")]
141    Internal(String),
142}
143
144pub trait TaskRoutingPolicy: Send + Sync {
145    fn route(&self, request: &ToolRequest) -> RoutingDecision;
146}
147
148impl<F> TaskRoutingPolicy for F
149where
150    F: Fn(&ToolRequest) -> RoutingDecision + Send + Sync,
151{
152    fn route(&self, request: &ToolRequest) -> RoutingDecision {
153        self(request)
154    }
155}
156
157#[derive(Clone, Copy, Debug, PartialEq, Eq)]
158pub enum RoutingDecision {
159    Foreground,
160    Background,
161    ForegroundThenDetachAfter(Duration),
162}
163
164struct DefaultRoutingPolicy;
165
166impl TaskRoutingPolicy for DefaultRoutingPolicy {
167    fn route(&self, _request: &ToolRequest) -> RoutingDecision {
168        RoutingDecision::Foreground
169    }
170}
171
172#[async_trait]
173pub trait TaskManager: Send + Sync {
174    async fn start_task(
175        &self,
176        request: TaskLaunchRequest,
177        ctx: TaskStartContext,
178    ) -> Result<TaskStartOutcome, TaskManagerError>;
179
180    async fn wait_for_turn(
181        &self,
182        turn_id: &TurnId,
183        cancellation: Option<TurnCancellation>,
184    ) -> Result<Option<TurnTaskUpdate>, TaskManagerError>;
185
186    async fn take_pending_loop_updates(&self) -> Result<PendingLoopUpdates, TaskManagerError>;
187
188    async fn on_turn_interrupted(&self, turn_id: &TurnId) -> Result<(), TaskManagerError>;
189
190    fn handle(&self) -> TaskManagerHandle;
191}
192
193#[async_trait]
194trait TaskManagerControl: Send + Sync {
195    async fn next_event(&self) -> Option<TaskEvent>;
196    async fn cancel(&self, task_id: TaskId) -> Result<(), TaskManagerError>;
197    async fn list_running(&self) -> Vec<TaskSnapshot>;
198    async fn list_completed(&self) -> Vec<TaskSnapshot>;
199    async fn drain_ready_items(&self) -> Vec<Item>;
200    async fn set_continue_policy(
201        &self,
202        task_id: TaskId,
203        policy: ContinuePolicy,
204    ) -> Result<(), TaskManagerError>;
205    async fn set_delivery_mode(
206        &self,
207        task_id: TaskId,
208        mode: DeliveryMode,
209    ) -> Result<(), TaskManagerError>;
210    async fn wait_for_idle(&self);
211}
212
213#[derive(Clone)]
214pub struct TaskManagerHandle {
215    inner: Arc<dyn TaskManagerControl>,
216}
217
218impl TaskManagerHandle {
219    pub async fn next_event(&self) -> Option<TaskEvent> {
220        self.inner.next_event().await
221    }
222
223    pub async fn cancel(&self, task_id: TaskId) -> Result<(), TaskManagerError> {
224        self.inner.cancel(task_id).await
225    }
226
227    pub async fn list_running(&self) -> Vec<TaskSnapshot> {
228        self.inner.list_running().await
229    }
230
231    pub async fn list_completed(&self) -> Vec<TaskSnapshot> {
232        self.inner.list_completed().await
233    }
234
235    pub async fn drain_ready_items(&self) -> Vec<Item> {
236        self.inner.drain_ready_items().await
237    }
238
239    pub async fn set_continue_policy(
240        &self,
241        task_id: TaskId,
242        policy: ContinuePolicy,
243    ) -> Result<(), TaskManagerError> {
244        self.inner.set_continue_policy(task_id, policy).await
245    }
246
247    pub async fn set_delivery_mode(
248        &self,
249        task_id: TaskId,
250        mode: DeliveryMode,
251    ) -> Result<(), TaskManagerError> {
252        self.inner.set_delivery_mode(task_id, mode).await
253    }
254
255    /// Wait until all running tasks have completed.
256    pub async fn wait_for_idle(&self) {
257        self.inner.wait_for_idle().await
258    }
259}
260
261pub struct SimpleTaskManager {
262    state: Arc<HandleState>,
263}
264
265impl SimpleTaskManager {
266    pub fn new() -> Self {
267        Self {
268            state: Arc::new(HandleState::default()),
269        }
270    }
271}
272
273impl Default for SimpleTaskManager {
274    fn default() -> Self {
275        Self::new()
276    }
277}
278
279#[async_trait]
280impl TaskManager for SimpleTaskManager {
281    async fn start_task(
282        &self,
283        request: TaskLaunchRequest,
284        ctx: TaskStartContext,
285    ) -> Result<TaskStartOutcome, TaskManagerError> {
286        let task_id = request
287            .task_id
288            .clone()
289            .unwrap_or_else(|| self.state.next_task_id());
290        let outcome = match &request.kind {
291            TaskLaunchKind::Approved(approved) => {
292                ctx.executor
293                    .execute_approved_owned(request.request.clone(), approved, ctx.tool_context)
294                    .await
295            }
296            TaskLaunchKind::Plain => {
297                ctx.executor
298                    .execute_owned(request.request.clone(), ctx.tool_context)
299                    .await
300            }
301        };
302        Ok(TaskStartOutcome::Ready(Box::new(
303            map_outcome_to_resolution(Some(task_id), request.request, outcome),
304        )))
305    }
306
307    async fn wait_for_turn(
308        &self,
309        _turn_id: &TurnId,
310        _cancellation: Option<TurnCancellation>,
311    ) -> Result<Option<TurnTaskUpdate>, TaskManagerError> {
312        Ok(None)
313    }
314
315    async fn take_pending_loop_updates(&self) -> Result<PendingLoopUpdates, TaskManagerError> {
316        Ok(PendingLoopUpdates::default())
317    }
318
319    async fn on_turn_interrupted(&self, _turn_id: &TurnId) -> Result<(), TaskManagerError> {
320        Ok(())
321    }
322
323    fn handle(&self) -> TaskManagerHandle {
324        TaskManagerHandle {
325            inner: self.state.clone(),
326        }
327    }
328}
329
330#[derive(Default)]
331struct HandleState {
332    next_task_index: AtomicU64,
333    events_rx: Mutex<Option<mpsc::UnboundedReceiver<TaskEvent>>>,
334}
335
336impl HandleState {
337    fn next_task_id(&self) -> TaskId {
338        let next = self.next_task_index.fetch_add(1, Ordering::SeqCst) + 1;
339        TaskId::new(format!("task-{}", next))
340    }
341}
342
343#[async_trait]
344impl TaskManagerControl for HandleState {
345    async fn next_event(&self) -> Option<TaskEvent> {
346        let mut rx = self.events_rx.lock().await;
347        match rx.as_mut() {
348            Some(inner) => inner.recv().await,
349            None => None,
350        }
351    }
352
353    async fn cancel(&self, task_id: TaskId) -> Result<(), TaskManagerError> {
354        Err(TaskManagerError::NotFound(task_id))
355    }
356
357    async fn list_running(&self) -> Vec<TaskSnapshot> {
358        Vec::new()
359    }
360
361    async fn list_completed(&self) -> Vec<TaskSnapshot> {
362        Vec::new()
363    }
364
365    async fn drain_ready_items(&self) -> Vec<Item> {
366        Vec::new()
367    }
368
369    async fn set_continue_policy(
370        &self,
371        task_id: TaskId,
372        _policy: ContinuePolicy,
373    ) -> Result<(), TaskManagerError> {
374        Err(TaskManagerError::NotFound(task_id))
375    }
376
377    async fn set_delivery_mode(
378        &self,
379        task_id: TaskId,
380        _mode: DeliveryMode,
381    ) -> Result<(), TaskManagerError> {
382        Err(TaskManagerError::NotFound(task_id))
383    }
384
385    async fn wait_for_idle(&self) {}
386}
387
388pub struct AsyncTaskManager {
389    inner: Arc<AsyncInner>,
390    routing: Arc<dyn TaskRoutingPolicy>,
391}
392
393impl AsyncTaskManager {
394    pub fn new() -> Self {
395        let (event_tx, event_rx) = mpsc::unbounded_channel();
396        Self {
397            inner: Arc::new(AsyncInner {
398                state: Mutex::new(AsyncState::default()),
399                host_event_tx: event_tx,
400                host_event_rx: Mutex::new(event_rx),
401                notify: Notify::new(),
402            }),
403            routing: Arc::new(DefaultRoutingPolicy),
404        }
405    }
406
407    pub fn routing(mut self, policy: impl TaskRoutingPolicy + 'static) -> Self {
408        self.routing = Arc::new(policy);
409        self
410    }
411}
412
413impl Default for AsyncTaskManager {
414    fn default() -> Self {
415        Self::new()
416    }
417}
418
419#[derive(Default)]
420struct AsyncState {
421    next_task_index: u64,
422    tasks: BTreeMap<TaskId, TaskRecord>,
423    per_turn_running: BTreeMap<TurnId, usize>,
424    per_turn_updates: BTreeMap<TurnId, VecDeque<TurnTaskUpdate>>,
425    pending_loop_updates: VecDeque<TaskResolution>,
426    manual_ready_items: Vec<Item>,
427}
428
429struct TaskRecord {
430    snapshot: TaskSnapshot,
431    continue_policy: ContinuePolicy,
432    delivery_mode: DeliveryMode,
433    running: bool,
434    completed: bool,
435    join: Option<JoinHandle<()>>,
436}
437
438struct AsyncInner {
439    state: Mutex<AsyncState>,
440    host_event_tx: mpsc::UnboundedSender<TaskEvent>,
441    host_event_rx: Mutex<mpsc::UnboundedReceiver<TaskEvent>>,
442    notify: Notify,
443}
444
445impl AsyncInner {
446    async fn next_task_id(&self) -> TaskId {
447        let mut state = self.state.lock().await;
448        state.next_task_index += 1;
449        TaskId::new(format!("task-{}", state.next_task_index))
450    }
451}
452
453#[async_trait]
454impl TaskManager for AsyncTaskManager {
455    async fn start_task(
456        &self,
457        request: TaskLaunchRequest,
458        ctx: TaskStartContext,
459    ) -> Result<TaskStartOutcome, TaskManagerError> {
460        let route = self.routing.route(&request.request);
461        let task_id = match request.task_id.clone() {
462            Some(existing) => existing,
463            None => self.inner.next_task_id().await,
464        };
465        let initial_kind = match route {
466            RoutingDecision::Background => TaskKind::Background,
467            _ => TaskKind::Foreground,
468        };
469        let snapshot = TaskSnapshot {
470            id: task_id.clone(),
471            turn_id: request.request.turn_id.clone(),
472            call_id: request.request.call_id.clone(),
473            tool_name: request.request.tool_name.to_string(),
474            kind: initial_kind,
475            metadata: request.request.metadata.clone(),
476        };
477        let _ = self
478            .inner
479            .host_event_tx
480            .send(TaskEvent::Started(snapshot.clone()));
481
482        let mut state = self.inner.state.lock().await;
483        state.tasks.insert(
484            task_id.clone(),
485            TaskRecord {
486                snapshot: snapshot.clone(),
487                continue_policy: ContinuePolicy::NotifyOnly,
488                delivery_mode: DeliveryMode::ToLoop,
489                running: true,
490                completed: false,
491                join: None,
492            },
493        );
494        if initial_kind == TaskKind::Foreground {
495            *state
496                .per_turn_running
497                .entry(snapshot.turn_id.clone())
498                .or_default() += 1;
499        }
500        drop(state);
501
502        let event_tx = self.inner.host_event_tx.clone();
503        let inner = self.inner.clone();
504        let task_id_for_future = task_id.clone();
505        let turn_id = snapshot.turn_id.clone();
506        let kind = request.kind.clone();
507        let exec_request = request.request.clone();
508        let owned_ctx = ctx.tool_context.clone();
509        let executor = ctx.executor.clone();
510        let route_copy = route;
511        let join = tokio::spawn(async move {
512            if let RoutingDecision::ForegroundThenDetachAfter(duration) = route_copy {
513                let event_tx = event_tx.clone();
514                let inner = inner.clone();
515                let task_id = task_id_for_future.clone();
516                let turn_id = turn_id.clone();
517                tokio::spawn(async move {
518                    tokio::time::sleep(duration).await;
519                    let mut state = inner.state.lock().await;
520                    let snapshot = if let Some(record) = state.tasks.get_mut(&task_id)
521                        && record.running
522                        && record.snapshot.kind == TaskKind::Foreground
523                    {
524                        record.snapshot.kind = TaskKind::Background;
525                        Some(record.snapshot.clone())
526                    } else {
527                        None
528                    };
529                    if let Some(snapshot) = snapshot {
530                        if let Some(count) = state.per_turn_running.get_mut(&turn_id) {
531                            *count = count.saturating_sub(1);
532                            if *count == 0 {
533                                state.per_turn_running.remove(&turn_id);
534                            }
535                        }
536                        state
537                            .per_turn_updates
538                            .entry(turn_id.clone())
539                            .or_default()
540                            .push_back(TurnTaskUpdate::Detached(snapshot.clone()));
541                        let _ = event_tx.send(TaskEvent::Detached(snapshot));
542                        inner.notify.notify_waiters();
543                    }
544                });
545            }
546
547            let outcome = match &kind {
548                TaskLaunchKind::Approved(approval) => {
549                    executor
550                        .execute_approved_owned(exec_request.clone(), approval, owned_ctx)
551                        .await
552                }
553                TaskLaunchKind::Plain => {
554                    executor
555                        .execute_owned(exec_request.clone(), owned_ctx)
556                        .await
557                }
558            };
559
560            let resolution =
561                map_outcome_to_resolution(Some(task_id_for_future.clone()), exec_request, outcome);
562            let completed_result = match &resolution {
563                TaskResolution::Item(item) => item.parts.iter().find_map(|part| match part {
564                    agentkit_core::Part::ToolResult(result) => Some(result.clone()),
565                    _ => None,
566                }),
567                TaskResolution::Approval(_) => None,
568            };
569
570            let (snapshot, should_request_continue) = {
571                let mut state = inner.state.lock().await;
572                let Some(record) = state.tasks.get_mut(&task_id_for_future) else {
573                    return;
574                };
575                record.running = false;
576                record.completed = true;
577                let snapshot = record.snapshot.clone();
578                let continue_policy = record.continue_policy;
579                let delivery_mode = record.delivery_mode;
580                let current_kind = snapshot.kind;
581
582                if current_kind == TaskKind::Foreground {
583                    if let Some(count) = state.per_turn_running.get_mut(&turn_id) {
584                        *count = count.saturating_sub(1);
585                        if *count == 0 {
586                            state.per_turn_running.remove(&turn_id);
587                        }
588                    }
589                    state
590                        .per_turn_updates
591                        .entry(turn_id.clone())
592                        .or_default()
593                        .push_back(TurnTaskUpdate::Resolution(Box::new(resolution.clone())));
594                } else {
595                    match &resolution {
596                        TaskResolution::Item(_) if delivery_mode == DeliveryMode::ToLoop => {
597                            state.pending_loop_updates.push_back(resolution.clone());
598                        }
599                        TaskResolution::Approval(_) if delivery_mode == DeliveryMode::ToLoop => {
600                            state.pending_loop_updates.push_back(resolution.clone());
601                        }
602                        TaskResolution::Item(item) => {
603                            state.manual_ready_items.push(item.clone());
604                        }
605                        TaskResolution::Approval(_) => {}
606                    }
607                }
608
609                (
610                    snapshot,
611                    current_kind == TaskKind::Background
612                        && delivery_mode == DeliveryMode::ToLoop
613                        && continue_policy == ContinuePolicy::RequestContinue,
614                )
615            };
616
617            if let Some(result) = completed_result {
618                let _ = event_tx.send(TaskEvent::Completed(snapshot.clone(), result));
619            }
620            if should_request_continue {
621                let _ = event_tx.send(TaskEvent::ContinueRequested);
622            }
623            inner.notify.notify_waiters();
624        });
625
626        let mut state = self.inner.state.lock().await;
627        if let Some(record) = state.tasks.get_mut(&task_id) {
628            record.join = Some(join);
629        }
630        Ok(TaskStartOutcome::Pending {
631            task_id,
632            kind: initial_kind,
633        })
634    }
635
636    async fn wait_for_turn(
637        &self,
638        turn_id: &TurnId,
639        cancellation: Option<TurnCancellation>,
640    ) -> Result<Option<TurnTaskUpdate>, TaskManagerError> {
641        loop {
642            {
643                let mut state = self.inner.state.lock().await;
644                if let Some(queue) = state.per_turn_updates.get_mut(turn_id)
645                    && let Some(update) = queue.pop_front()
646                {
647                    return Ok(Some(update));
648                }
649                if state
650                    .per_turn_running
651                    .get(turn_id)
652                    .copied()
653                    .unwrap_or_default()
654                    == 0
655                {
656                    return Ok(None);
657                }
658            }
659            if cancellation
660                .as_ref()
661                .is_some_and(TurnCancellation::is_cancelled)
662            {
663                return Ok(None);
664            }
665            if let Some(cancellation) = cancellation.as_ref() {
666                tokio::select! {
667                    _ = self.inner.notify.notified() => {}
668                    _ = cancellation.cancelled() => return Ok(None),
669                }
670            } else {
671                self.inner.notify.notified().await;
672            }
673        }
674    }
675
676    async fn take_pending_loop_updates(&self) -> Result<PendingLoopUpdates, TaskManagerError> {
677        let mut state = self.inner.state.lock().await;
678        Ok(PendingLoopUpdates {
679            resolutions: std::mem::take(&mut state.pending_loop_updates),
680        })
681    }
682
683    async fn on_turn_interrupted(&self, turn_id: &TurnId) -> Result<(), TaskManagerError> {
684        let mut state = self.inner.state.lock().await;
685        let interrupted: Vec<TaskId> = state
686            .tasks
687            .iter()
688            .filter_map(|(id, record)| {
689                (record.snapshot.turn_id == *turn_id
690                    && record.snapshot.kind == TaskKind::Foreground
691                    && record.running)
692                    .then_some(id.clone())
693            })
694            .collect();
695        for task_id in interrupted {
696            if let Some(record) = state.tasks.get_mut(&task_id) {
697                record.running = false;
698                if let Some(join) = record.join.take() {
699                    join.abort();
700                }
701                let snapshot = record.snapshot.clone();
702                let _ = self
703                    .inner
704                    .host_event_tx
705                    .send(TaskEvent::Cancelled(snapshot));
706            }
707        }
708        state.per_turn_running.remove(turn_id);
709        self.inner.notify.notify_waiters();
710        Ok(())
711    }
712
713    fn handle(&self) -> TaskManagerHandle {
714        TaskManagerHandle {
715            inner: self.inner.clone(),
716        }
717    }
718}
719
720#[async_trait]
721impl TaskManagerControl for AsyncInner {
722    async fn next_event(&self) -> Option<TaskEvent> {
723        self.host_event_rx.lock().await.recv().await
724    }
725
726    async fn cancel(&self, task_id: TaskId) -> Result<(), TaskManagerError> {
727        let mut state = self.state.lock().await;
728        let record = state
729            .tasks
730            .get_mut(&task_id)
731            .ok_or_else(|| TaskManagerError::NotFound(task_id.clone()))?;
732        if let Some(join) = record.join.take() {
733            join.abort();
734        }
735        record.running = false;
736        let snapshot = record.snapshot.clone();
737        if record.snapshot.kind == TaskKind::Foreground
738            && let Some(count) = state.per_turn_running.get_mut(&snapshot.turn_id)
739        {
740            *count = count.saturating_sub(1);
741            if *count == 0 {
742                state.per_turn_running.remove(&snapshot.turn_id);
743            }
744        }
745        let _ = self.host_event_tx.send(TaskEvent::Cancelled(snapshot));
746        self.notify.notify_waiters();
747        Ok(())
748    }
749
750    async fn list_running(&self) -> Vec<TaskSnapshot> {
751        let state = self.state.lock().await;
752        state
753            .tasks
754            .values()
755            .filter(|record| record.running)
756            .map(|record| record.snapshot.clone())
757            .collect()
758    }
759
760    async fn list_completed(&self) -> Vec<TaskSnapshot> {
761        let state = self.state.lock().await;
762        state
763            .tasks
764            .values()
765            .filter(|record| record.completed)
766            .map(|record| record.snapshot.clone())
767            .collect()
768    }
769
770    async fn drain_ready_items(&self) -> Vec<Item> {
771        let mut state = self.state.lock().await;
772        std::mem::take(&mut state.manual_ready_items)
773    }
774
775    async fn set_continue_policy(
776        &self,
777        task_id: TaskId,
778        policy: ContinuePolicy,
779    ) -> Result<(), TaskManagerError> {
780        let mut state = self.state.lock().await;
781        let record = state
782            .tasks
783            .get_mut(&task_id)
784            .ok_or_else(|| TaskManagerError::NotFound(task_id.clone()))?;
785        record.continue_policy = policy;
786        Ok(())
787    }
788
789    async fn set_delivery_mode(
790        &self,
791        task_id: TaskId,
792        mode: DeliveryMode,
793    ) -> Result<(), TaskManagerError> {
794        let mut state = self.state.lock().await;
795        let record = state
796            .tasks
797            .get_mut(&task_id)
798            .ok_or_else(|| TaskManagerError::NotFound(task_id.clone()))?;
799        record.delivery_mode = mode;
800        Ok(())
801    }
802
803    async fn wait_for_idle(&self) {
804        loop {
805            {
806                let state = self.state.lock().await;
807                if !state.tasks.values().any(|r| r.running) {
808                    return;
809                }
810            }
811            self.notify.notified().await;
812        }
813    }
814}
815
816fn map_outcome_to_resolution(
817    task_id: Option<TaskId>,
818    request: ToolRequest,
819    outcome: ToolExecutionOutcome,
820) -> TaskResolution {
821    match outcome {
822        ToolExecutionOutcome::Completed(result) => TaskResolution::Item(Item {
823            id: None,
824            kind: agentkit_core::ItemKind::Tool,
825            parts: vec![agentkit_core::Part::ToolResult(result.result)],
826            metadata: result.metadata,
827            usage: None,
828            finish_reason: None,
829            created_at: None,
830        }),
831        ToolExecutionOutcome::Interrupted(
832            agentkit_tools_core::ToolInterruption::ApprovalRequired(mut approval),
833        ) => {
834            let task_id = task_id.unwrap_or_default();
835            approval.task_id = Some(task_id.clone());
836            TaskResolution::Approval(TaskApproval {
837                task_id,
838                tool_request: request,
839                approval,
840            })
841        }
842        ToolExecutionOutcome::FailedBeforeInvocation(error) => {
843            let mut metadata = request.metadata;
844            if matches!(error, ToolError::PermissionDenied(_)) {
845                metadata.insert(
846                    TOOL_RESULT_FAILURE_KIND_METADATA_KEY.into(),
847                    TOOL_RESULT_FAILURE_KIND_PERMISSION_DENIED.into(),
848                );
849            }
850            TaskResolution::Item(Item {
851                id: None,
852                kind: agentkit_core::ItemKind::Tool,
853                parts: vec![agentkit_core::Part::ToolResult(ToolResultPart {
854                    call_id: request.call_id,
855                    output: agentkit_core::ToolOutput::Text(error.to_string()),
856                    is_error: true,
857                    metadata,
858                })],
859                metadata: MetadataMap::new(),
860                usage: None,
861                finish_reason: None,
862                created_at: None,
863            })
864        }
865        ToolExecutionOutcome::Failed(error) => TaskResolution::Item(Item {
866            id: None,
867            kind: agentkit_core::ItemKind::Tool,
868            parts: vec![agentkit_core::Part::ToolResult(ToolResultPart {
869                call_id: request.call_id,
870                output: agentkit_core::ToolOutput::Text(error.to_string()),
871                is_error: true,
872                metadata: request.metadata,
873            })],
874            metadata: MetadataMap::new(),
875            usage: None,
876            finish_reason: None,
877            created_at: None,
878        }),
879    }
880}
881
882#[cfg(test)]
883mod tests {
884    use std::collections::BTreeMap;
885    use std::sync::Arc as StdArc;
886    use std::sync::atomic::{AtomicBool, Ordering as AtomicOrdering};
887
888    use agentkit_core::{
889        CancellationController, ItemKind, Part, SessionId, ToolOutput, TurnCancellation,
890    };
891    use agentkit_tools_core::{
892        ApprovalReason, PermissionChecker, PermissionDecision, ToolAnnotations, ToolInterruption,
893        ToolName, ToolResult, ToolSpec,
894    };
895    use serde_json::json;
896    use tokio::sync::Notify;
897    use tokio::time::{Duration, timeout};
898
899    use super::*;
900
901    struct AllowAllPermissions;
902
903    impl PermissionChecker for AllowAllPermissions {
904        fn evaluate(
905            &self,
906            _request: &dyn agentkit_tools_core::PermissionRequest,
907        ) -> PermissionDecision {
908            PermissionDecision::Allow
909        }
910    }
911
912    #[derive(Clone)]
913    enum TestBehavior {
914        Block {
915            entered: StdArc<AtomicBool>,
916            release: StdArc<Notify>,
917            output: &'static str,
918        },
919        Approval,
920    }
921
922    #[derive(Clone)]
923    struct TestExecutor {
924        behaviors: BTreeMap<String, TestBehavior>,
925    }
926
927    impl TestExecutor {
928        fn new(behaviors: impl IntoIterator<Item = (impl Into<String>, TestBehavior)>) -> Self {
929            Self {
930                behaviors: behaviors
931                    .into_iter()
932                    .map(|(name, behavior)| (name.into(), behavior))
933                    .collect(),
934            }
935        }
936    }
937
938    #[async_trait]
939    impl ToolExecutor for TestExecutor {
940        fn specs(&self) -> Vec<ToolSpec> {
941            self.behaviors
942                .keys()
943                .map(|name| ToolSpec {
944                    name: ToolName::new(name),
945                    description: format!("test tool {name}"),
946                    input_schema: json!({
947                        "type": "object",
948                        "properties": {},
949                        "additionalProperties": false
950                    }),
951                    output_schema: None,
952                    annotations: ToolAnnotations::default(),
953                    metadata: MetadataMap::new(),
954                })
955                .collect()
956        }
957
958        async fn execute(
959            &self,
960            request: ToolRequest,
961            _ctx: &mut agentkit_tools_core::ToolContext<'_>,
962        ) -> ToolExecutionOutcome {
963            match self.behaviors.get(request.tool_name.0.as_str()) {
964                Some(TestBehavior::Block {
965                    entered,
966                    release,
967                    output,
968                }) => {
969                    entered.store(true, AtomicOrdering::SeqCst);
970                    release.notified().await;
971                    ToolExecutionOutcome::Completed(ToolResult {
972                        result: ToolResultPart {
973                            call_id: request.call_id,
974                            output: ToolOutput::Text((*output).into()),
975                            is_error: false,
976                            metadata: request.metadata,
977                        },
978                        duration: None,
979                        metadata: MetadataMap::new(),
980                    })
981                }
982                Some(TestBehavior::Approval) => ToolExecutionOutcome::Interrupted(
983                    ToolInterruption::ApprovalRequired(ApprovalRequest {
984                        task_id: None,
985                        call_id: Some(request.call_id.clone()),
986                        id: "approval:test".into(),
987                        request_kind: "tool.test".into(),
988                        reason: ApprovalReason::SensitivePath,
989                        summary: "requires approval".into(),
990                        metadata: MetadataMap::new(),
991                    }),
992                ),
993                None => ToolExecutionOutcome::Failed(ToolError::Unavailable(
994                    request.tool_name.0.clone(),
995                )),
996            }
997        }
998    }
999
1000    struct NameRoutingPolicy {
1001        routes: BTreeMap<String, RoutingDecision>,
1002    }
1003
1004    impl NameRoutingPolicy {
1005        fn new(routes: impl IntoIterator<Item = (impl Into<String>, RoutingDecision)>) -> Self {
1006            Self {
1007                routes: routes
1008                    .into_iter()
1009                    .map(|(name, decision)| (name.into(), decision))
1010                    .collect(),
1011            }
1012        }
1013    }
1014
1015    impl TaskRoutingPolicy for NameRoutingPolicy {
1016        fn route(&self, request: &ToolRequest) -> RoutingDecision {
1017            self.routes
1018                .get(request.tool_name.0.as_str())
1019                .copied()
1020                .unwrap_or(RoutingDecision::Foreground)
1021        }
1022    }
1023
1024    fn make_request(tool_name: &str, turn_id: &str, call_id: &str) -> ToolRequest {
1025        ToolRequest {
1026            call_id: ToolCallId::new(call_id),
1027            tool_name: ToolName::new(tool_name),
1028            input: json!({}),
1029            session_id: SessionId::new("session-1"),
1030            turn_id: TurnId::new(turn_id),
1031            metadata: MetadataMap::new(),
1032        }
1033    }
1034
1035    fn make_context(
1036        executor: Arc<dyn ToolExecutor>,
1037        turn_id: &TurnId,
1038        cancellation: Option<TurnCancellation>,
1039    ) -> TaskStartContext {
1040        TaskStartContext {
1041            executor,
1042            tool_context: OwnedToolContext {
1043                session_id: SessionId::new("session-1"),
1044                turn_id: turn_id.clone(),
1045                metadata: MetadataMap::new(),
1046                permissions: Arc::new(AllowAllPermissions),
1047                resources: Arc::new(()),
1048                cancellation,
1049                execution_scope: None,
1050                approved_request: None,
1051            },
1052        }
1053    }
1054
1055    async fn next_event(handle: &TaskManagerHandle) -> TaskEvent {
1056        timeout(Duration::from_secs(1), handle.next_event())
1057            .await
1058            .expect("timed out waiting for task event")
1059            .expect("task event stream ended unexpectedly")
1060    }
1061
1062    async fn wait_until_entered(entered: &AtomicBool) {
1063        timeout(Duration::from_secs(1), async {
1064            while !entered.load(AtomicOrdering::SeqCst) {
1065                tokio::task::yield_now().await;
1066            }
1067        })
1068        .await
1069        .expect("task never entered execution");
1070    }
1071
1072    #[tokio::test]
1073    async fn simple_task_manager_executes_inline_and_assigns_task_ids() {
1074        let manager = SimpleTaskManager::new();
1075        let executor: Arc<dyn ToolExecutor> = Arc::new(TestExecutor::new([(
1076            "needs-approval",
1077            TestBehavior::Approval,
1078        )]));
1079        let request = make_request("needs-approval", "turn-1", "call-1");
1080
1081        let outcome = manager
1082            .start_task(
1083                TaskLaunchRequest {
1084                    task_id: None,
1085                    request: request.clone(),
1086                    kind: TaskLaunchKind::Plain,
1087                },
1088                make_context(executor, &request.turn_id, None),
1089            )
1090            .await
1091            .unwrap();
1092
1093        match outcome {
1094            TaskStartOutcome::Ready(resolution) => match *resolution {
1095                TaskResolution::Approval(task) => {
1096                    assert!(!task.task_id.0.is_empty());
1097                    assert_eq!(task.approval.task_id.as_ref(), Some(&task.task_id));
1098                    assert_eq!(task.tool_request.call_id, request.call_id);
1099                }
1100                other => panic!("unexpected task resolution: {other:?}"),
1101            },
1102            other => panic!("unexpected start outcome: {other:?}"),
1103        }
1104
1105        assert!(manager.handle().list_running().await.is_empty());
1106    }
1107
1108    #[tokio::test]
1109    async fn async_manager_interrupt_cancels_foreground_only() {
1110        let fg_release = StdArc::new(Notify::new());
1111        let fg_entered = StdArc::new(AtomicBool::new(false));
1112        let bg_release = StdArc::new(Notify::new());
1113        let bg_entered = StdArc::new(AtomicBool::new(false));
1114        let executor: Arc<dyn ToolExecutor> = Arc::new(TestExecutor::new([
1115            (
1116                "foreground",
1117                TestBehavior::Block {
1118                    entered: fg_entered.clone(),
1119                    release: fg_release.clone(),
1120                    output: "foreground-done",
1121                },
1122            ),
1123            (
1124                "background",
1125                TestBehavior::Block {
1126                    entered: bg_entered.clone(),
1127                    release: bg_release.clone(),
1128                    output: "background-done",
1129                },
1130            ),
1131        ]));
1132        let manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([
1133            ("foreground", RoutingDecision::Foreground),
1134            ("background", RoutingDecision::Background),
1135        ]));
1136        let handle = manager.handle();
1137        let turn_id = TurnId::new("turn-1");
1138
1139        let foreground = manager
1140            .start_task(
1141                TaskLaunchRequest {
1142                    task_id: None,
1143                    request: make_request("foreground", "turn-1", "call-fg"),
1144                    kind: TaskLaunchKind::Plain,
1145                },
1146                make_context(executor.clone(), &turn_id, None),
1147            )
1148            .await
1149            .unwrap();
1150        let background = manager
1151            .start_task(
1152                TaskLaunchRequest {
1153                    task_id: None,
1154                    request: make_request("background", "turn-1", "call-bg"),
1155                    kind: TaskLaunchKind::Plain,
1156                },
1157                make_context(executor.clone(), &turn_id, None),
1158            )
1159            .await
1160            .unwrap();
1161
1162        assert!(matches!(
1163            foreground,
1164            TaskStartOutcome::Pending {
1165                kind: TaskKind::Foreground,
1166                ..
1167            }
1168        ));
1169        let background_id = match background {
1170            TaskStartOutcome::Pending {
1171                task_id,
1172                kind: TaskKind::Background,
1173            } => task_id,
1174            other => panic!("unexpected background outcome: {other:?}"),
1175        };
1176
1177        let _ = next_event(&handle).await;
1178        let _ = next_event(&handle).await;
1179        wait_until_entered(fg_entered.as_ref()).await;
1180        wait_until_entered(bg_entered.as_ref()).await;
1181
1182        manager.on_turn_interrupted(&turn_id).await.unwrap();
1183
1184        match next_event(&handle).await {
1185            TaskEvent::Cancelled(snapshot) => assert_eq!(snapshot.tool_name, "foreground"),
1186            other => panic!("unexpected event after interrupt: {other:?}"),
1187        }
1188
1189        let running = handle.list_running().await;
1190        assert_eq!(running.len(), 1);
1191        assert_eq!(running[0].id, background_id);
1192        assert_eq!(running[0].tool_name, "background");
1193
1194        bg_release.notify_waiters();
1195        match next_event(&handle).await {
1196            TaskEvent::Completed(snapshot, result) => {
1197                assert_eq!(snapshot.id, background_id);
1198                assert_eq!(result.output, ToolOutput::Text("background-done".into()));
1199            }
1200            other => panic!("unexpected completion event: {other:?}"),
1201        }
1202    }
1203
1204    #[tokio::test]
1205    async fn async_manager_can_cancel_background_tasks_by_id() {
1206        let release = StdArc::new(Notify::new());
1207        let entered = StdArc::new(AtomicBool::new(false));
1208        let executor: Arc<dyn ToolExecutor> = Arc::new(TestExecutor::new([(
1209            "background",
1210            TestBehavior::Block {
1211                entered: entered.clone(),
1212                release,
1213                output: "done",
1214            },
1215        )]));
1216        let manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
1217            "background",
1218            RoutingDecision::Background,
1219        )]));
1220        let handle = manager.handle();
1221        let request = make_request("background", "turn-1", "call-1");
1222
1223        let task_id = match manager
1224            .start_task(
1225                TaskLaunchRequest {
1226                    task_id: None,
1227                    request: request.clone(),
1228                    kind: TaskLaunchKind::Plain,
1229                },
1230                make_context(executor, &request.turn_id, None),
1231            )
1232            .await
1233            .unwrap()
1234        {
1235            TaskStartOutcome::Pending { task_id, .. } => task_id,
1236            other => panic!("unexpected start outcome: {other:?}"),
1237        };
1238
1239        let _ = next_event(&handle).await;
1240        wait_until_entered(entered.as_ref()).await;
1241        handle.cancel(task_id.clone()).await.unwrap();
1242
1243        match next_event(&handle).await {
1244            TaskEvent::Cancelled(snapshot) => assert_eq!(snapshot.id, task_id),
1245            other => panic!("unexpected event after cancel: {other:?}"),
1246        }
1247
1248        assert!(handle.list_running().await.is_empty());
1249    }
1250
1251    #[tokio::test]
1252    async fn async_manager_manual_delivery_keeps_results_out_of_loop_updates() {
1253        let release = StdArc::new(Notify::new());
1254        let entered = StdArc::new(AtomicBool::new(false));
1255        let executor: Arc<dyn ToolExecutor> = Arc::new(TestExecutor::new([(
1256            "background",
1257            TestBehavior::Block {
1258                entered: entered.clone(),
1259                release: release.clone(),
1260                output: "manual-done",
1261            },
1262        )]));
1263        let manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
1264            "background",
1265            RoutingDecision::Background,
1266        )]));
1267        let handle = manager.handle();
1268        let request = make_request("background", "turn-1", "call-1");
1269
1270        let task_id = match manager
1271            .start_task(
1272                TaskLaunchRequest {
1273                    task_id: None,
1274                    request: request.clone(),
1275                    kind: TaskLaunchKind::Plain,
1276                },
1277                make_context(executor, &request.turn_id, None),
1278            )
1279            .await
1280            .unwrap()
1281        {
1282            TaskStartOutcome::Pending { task_id, .. } => task_id,
1283            other => panic!("unexpected start outcome: {other:?}"),
1284        };
1285
1286        let _ = next_event(&handle).await;
1287        wait_until_entered(entered.as_ref()).await;
1288        handle
1289            .set_continue_policy(task_id.clone(), ContinuePolicy::RequestContinue)
1290            .await
1291            .unwrap();
1292        handle
1293            .set_delivery_mode(task_id, DeliveryMode::Manual)
1294            .await
1295            .unwrap();
1296
1297        release.notify_waiters();
1298        match next_event(&handle).await {
1299            TaskEvent::Completed(_, result) => {
1300                assert_eq!(result.output, ToolOutput::Text("manual-done".into()))
1301            }
1302            other => panic!("unexpected event: {other:?}"),
1303        }
1304
1305        assert!(
1306            timeout(Duration::from_millis(50), handle.next_event())
1307                .await
1308                .is_err()
1309        );
1310        assert!(
1311            manager
1312                .take_pending_loop_updates()
1313                .await
1314                .unwrap()
1315                .resolutions
1316                .is_empty()
1317        );
1318
1319        let ready_items = handle.drain_ready_items().await;
1320        assert_eq!(ready_items.len(), 1);
1321        assert_eq!(ready_items[0].kind, ItemKind::Tool);
1322        match &ready_items[0].parts[0] {
1323            Part::ToolResult(result) => {
1324                assert_eq!(result.output, ToolOutput::Text("manual-done".into()))
1325            }
1326            other => panic!("unexpected ready item: {other:?}"),
1327        }
1328    }
1329
1330    #[tokio::test]
1331    async fn async_manager_to_loop_delivery_can_request_continue() {
1332        let release = StdArc::new(Notify::new());
1333        let entered = StdArc::new(AtomicBool::new(false));
1334        let executor: Arc<dyn ToolExecutor> = Arc::new(TestExecutor::new([(
1335            "background",
1336            TestBehavior::Block {
1337                entered: entered.clone(),
1338                release: release.clone(),
1339                output: "loop-done",
1340            },
1341        )]));
1342        let manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
1343            "background",
1344            RoutingDecision::Background,
1345        )]));
1346        let handle = manager.handle();
1347        let request = make_request("background", "turn-1", "call-1");
1348
1349        let task_id = match manager
1350            .start_task(
1351                TaskLaunchRequest {
1352                    task_id: None,
1353                    request: request.clone(),
1354                    kind: TaskLaunchKind::Plain,
1355                },
1356                make_context(
1357                    executor,
1358                    &request.turn_id,
1359                    Some(TurnCancellation::new(
1360                        CancellationController::new().handle(),
1361                    )),
1362                ),
1363            )
1364            .await
1365            .unwrap()
1366        {
1367            TaskStartOutcome::Pending { task_id, .. } => task_id,
1368            other => panic!("unexpected start outcome: {other:?}"),
1369        };
1370
1371        let _ = next_event(&handle).await;
1372        wait_until_entered(entered.as_ref()).await;
1373        handle
1374            .set_continue_policy(task_id, ContinuePolicy::RequestContinue)
1375            .await
1376            .unwrap();
1377
1378        release.notify_waiters();
1379        match next_event(&handle).await {
1380            TaskEvent::Completed(_, result) => {
1381                assert_eq!(result.output, ToolOutput::Text("loop-done".into()))
1382            }
1383            other => panic!("unexpected completion event: {other:?}"),
1384        }
1385        match next_event(&handle).await {
1386            TaskEvent::ContinueRequested => {}
1387            other => panic!("unexpected follow-up event: {other:?}"),
1388        }
1389
1390        let updates = manager.take_pending_loop_updates().await.unwrap();
1391        assert_eq!(updates.resolutions.len(), 1);
1392        assert!(handle.drain_ready_items().await.is_empty());
1393    }
1394
1395    #[tokio::test]
1396    async fn wait_for_idle_returns_after_loop_updates_are_queued() {
1397        let release = StdArc::new(Notify::new());
1398        let entered = StdArc::new(AtomicBool::new(false));
1399        let executor: Arc<dyn ToolExecutor> = Arc::new(TestExecutor::new([(
1400            "background",
1401            TestBehavior::Block {
1402                entered: entered.clone(),
1403                release: release.clone(),
1404                output: "idle-done",
1405            },
1406        )]));
1407        let manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([(
1408            "background",
1409            RoutingDecision::Background,
1410        )]));
1411        let handle = manager.handle();
1412        let request = make_request("background", "turn-1", "call-1");
1413
1414        let outcome = manager
1415            .start_task(
1416                TaskLaunchRequest {
1417                    task_id: None,
1418                    request: request.clone(),
1419                    kind: TaskLaunchKind::Plain,
1420                },
1421                make_context(executor, &request.turn_id, None),
1422            )
1423            .await
1424            .unwrap();
1425        assert!(matches!(outcome, TaskStartOutcome::Pending { .. }));
1426
1427        let _ = next_event(&handle).await;
1428        wait_until_entered(entered.as_ref()).await;
1429        release.notify_waiters();
1430
1431        timeout(Duration::from_secs(1), handle.wait_for_idle())
1432            .await
1433            .expect("wait_for_idle timed out");
1434
1435        let updates = manager.take_pending_loop_updates().await.unwrap();
1436        assert_eq!(updates.resolutions.len(), 1);
1437        match &updates.resolutions[0] {
1438            TaskResolution::Item(item) => match &item.parts[0] {
1439                Part::ToolResult(result) => {
1440                    assert_eq!(result.call_id, request.call_id);
1441                    assert_eq!(result.output, ToolOutput::Text("idle-done".into()));
1442                }
1443                other => panic!("unexpected tool item: {other:?}"),
1444            },
1445            other => panic!("unexpected pending update: {other:?}"),
1446        }
1447    }
1448}
agentkit_task_manager/lib.rs

agentkit_task_manager/
lib.rs
