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

agentkit_task_manager/
lib.rs
