Skip to main content

agent_sdk/task/
watcher.rs

1use std::sync::Arc;
2use std::time::Duration;
3
4use tokio::time;
5use tracing::debug;
6
7use crate::error::SdkResult;
8use crate::task::store::{TaskStore, TaskSummary};
9
10pub struct TaskWatcher {
11    task_store: Arc<TaskStore>,
12    poll_interval: Duration,
13}
14
15impl TaskWatcher {
16    pub fn new(task_store: Arc<TaskStore>, poll_interval_ms: u64) -> Self {
17        Self {
18            task_store,
19            poll_interval: Duration::from_millis(poll_interval_ms),
20        }
21    }
22
23    pub async fn wait_for_completion(&self) -> SdkResult<TaskSummary> {
24        loop {
25            let summary = self.task_store.summary()?;
26
27            debug!(
28                pending = summary.pending,
29                in_progress = summary.in_progress,
30                completed = summary.completed,
31                failed = summary.failed,
32                "Task status"
33            );
34
35            if summary.is_done() {
36                return Ok(summary);
37            }
38
39            time::sleep(self.poll_interval).await;
40        }
41    }
42
43    pub fn current_summary(&self) -> SdkResult<TaskSummary> {
44        self.task_store.summary()
45    }
46}