daimon_core/
distributed.rs

1//! Core types and trait for distributed agent execution.
2//!
3//! Provider crates implement [`TaskBroker`] for their cloud-native message
4//! service. The main `daimon` crate re-exports everything from here.
5
6use std::collections::HashMap;
7use std::future::Future;
8use std::pin::Pin;
9
10use serde::{Deserialize, Serialize};
11
12use crate::error::Result;
13
14/// A unit of work submitted to a [`TaskBroker`].
15///
16/// Each task carries a unique ID and the input text for an agent prompt.
17/// Optional metadata lets callers tag tasks with routing hints, priority,
18/// or any application-specific data.
19#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct AgentTask {
21    /// Unique identifier for this task (generated on creation).
22    pub task_id: String,
23    /// The user input to prompt the agent with.
24    pub input: String,
25    /// Optional run ID for resumable execution via checkpoints.
26    pub run_id: Option<String>,
27    /// Arbitrary key-value metadata (routing hints, priority, etc.).
28    pub metadata: HashMap<String, serde_json::Value>,
29}
30
31impl AgentTask {
32    /// Creates a new task with a timestamp-based ID.
33    pub fn new(input: impl Into<String>) -> Self {
34        Self {
35            task_id: Self::generate_id(),
36            input: input.into(),
37            run_id: None,
38            metadata: HashMap::new(),
39        }
40    }
41
42    /// Assigns a checkpoint run ID for resumable execution.
43    pub fn with_run_id(mut self, run_id: impl Into<String>) -> Self {
44        self.run_id = Some(run_id.into());
45        self
46    }
47
48    /// Adds a metadata key-value pair.
49    pub fn with_metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
50        self.metadata.insert(key.into(), value);
51        self
52    }
53
54    fn generate_id() -> String {
55        use std::time::{SystemTime, UNIX_EPOCH};
56        let ts = SystemTime::now()
57            .duration_since(UNIX_EPOCH)
58            .unwrap_or_default()
59            .as_nanos();
60        format!("task-{ts:x}")
61    }
62}
63
64/// The result of a completed agent task.
65#[derive(Debug, Clone, Serialize, Deserialize)]
66pub struct TaskResult {
67    /// The task ID this result corresponds to.
68    pub task_id: String,
69    /// The agent's final text output.
70    pub output: String,
71    /// Number of ReAct iterations the agent performed.
72    pub iterations: usize,
73    /// Estimated cost in USD (if a cost model was configured).
74    pub cost: f64,
75    /// Error message if the task failed.
76    pub error: Option<String>,
77}
78
79/// Current status of a distributed task.
80#[derive(Debug, Clone, Serialize, Deserialize)]
81pub enum TaskStatus {
82    /// Submitted but not yet picked up by a worker.
83    Pending,
84    /// Currently being executed by a worker.
85    Running,
86    /// Completed successfully.
87    Completed(TaskResult),
88    /// Failed with an error.
89    Failed(String),
90}
91
92/// Trait for distributing agent tasks across workers.
93///
94/// Implement this for your message broker (AWS SQS, Google Pub/Sub,
95/// Azure Service Bus, Redis, NATS, RabbitMQ, etc.) to enable
96/// multi-process agent execution.
97pub trait TaskBroker: Send + Sync {
98    /// Submits a task for execution. Returns the task ID.
99    fn submit(&self, task: AgentTask) -> impl Future<Output = Result<String>> + Send;
100
101    /// Queries the current status of a task.
102    fn status(&self, task_id: &str) -> impl Future<Output = Result<TaskStatus>> + Send;
103
104    /// Blocks until a task is available and returns it.
105    /// Returns `None` if the broker is closed.
106    fn receive(&self) -> impl Future<Output = Result<Option<AgentTask>>> + Send;
107
108    /// Marks a task as completed with the given result.
109    fn complete(
110        &self,
111        task_id: &str,
112        result: TaskResult,
113    ) -> impl Future<Output = Result<()>> + Send;
114
115    /// Marks a task as failed with an error message.
116    fn fail(&self, task_id: &str, error: String) -> impl Future<Output = Result<()>> + Send;
117}
118
119/// Object-safe wrapper for [`TaskBroker`], enabling `Arc<dyn ErasedTaskBroker>`.
120pub trait ErasedTaskBroker: Send + Sync {
121    fn submit_erased<'a>(
122        &'a self,
123        task: AgentTask,
124    ) -> Pin<Box<dyn Future<Output = Result<String>> + Send + 'a>>;
125
126    fn status_erased<'a>(
127        &'a self,
128        task_id: &'a str,
129    ) -> Pin<Box<dyn Future<Output = Result<TaskStatus>> + Send + 'a>>;
130
131    fn receive_erased(
132        &self,
133    ) -> Pin<Box<dyn Future<Output = Result<Option<AgentTask>>> + Send + '_>>;
134
135    fn complete_erased<'a>(
136        &'a self,
137        task_id: &'a str,
138        result: TaskResult,
139    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>;
140
141    fn fail_erased<'a>(
142        &'a self,
143        task_id: &'a str,
144        error: String,
145    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>;
146}
147
148impl<T: TaskBroker> ErasedTaskBroker for T {
149    fn submit_erased<'a>(
150        &'a self,
151        task: AgentTask,
152    ) -> Pin<Box<dyn Future<Output = Result<String>> + Send + 'a>> {
153        Box::pin(self.submit(task))
154    }
155
156    fn status_erased<'a>(
157        &'a self,
158        task_id: &'a str,
159    ) -> Pin<Box<dyn Future<Output = Result<TaskStatus>> + Send + 'a>> {
160        Box::pin(self.status(task_id))
161    }
162
163    fn receive_erased(
164        &self,
165    ) -> Pin<Box<dyn Future<Output = Result<Option<AgentTask>>> + Send + '_>> {
166        Box::pin(self.receive())
167    }
168
169    fn complete_erased<'a>(
170        &'a self,
171        task_id: &'a str,
172        result: TaskResult,
173    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>> {
174        Box::pin(self.complete(task_id, result))
175    }
176
177    fn fail_erased<'a>(
178        &'a self,
179        task_id: &'a str,
180        error: String,
181    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>> {
182        Box::pin(self.fail(task_id, error))
183    }
184}
daimon_core/distributed.rs

daimon_core/
distributed.rs
