Skip to main content

agentdb/
workflows.rs

1use crate::error::{AgentDbError, Result};
2use crate::schema::now_ms;
3use rusqlite::params;
4use rusqlite::Connection;
5use serde_json::Value;
6use std::sync::{Arc, Mutex};
7use uuid::Uuid;
8
9/// A workflow and its current state.
10#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
11pub struct Workflow {
12    /// Unique identifier for this workflow.
13    pub id: String,
14    /// Human-readable name for the workflow.
15    pub name: String,
16    /// Current status (`"pending"`, `"running"`, `"completed"`, `"failed"`).
17    pub status: String,
18    /// JSON input payload passed to the workflow.
19    pub input: Option<Value>,
20    /// JSON output payload produced by the workflow.
21    pub output: Option<Value>,
22    /// Error message if the workflow failed.
23    pub error: Option<String>,
24    /// Arbitrary JSON metadata.
25    pub metadata: Option<Value>,
26    /// Unix-millisecond timestamp when the workflow was created.
27    pub created_at: i64,
28    /// Unix-millisecond timestamp of the most recent status change.
29    pub updated_at: i64,
30    /// Number of steps in this workflow (populated by `list_workflows`; full
31    /// step objects are only fetched by `get_workflow`).
32    pub step_count: i64,
33    /// Steps belonging to this workflow, ordered by `step_index`.
34    pub steps: Vec<WorkflowStep>,
35}
36
37/// A single step within a workflow.
38#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
39pub struct WorkflowStep {
40    /// Unique identifier for this step.
41    pub id: String,
42    /// ID of the parent workflow.
43    pub workflow_id: String,
44    /// Zero-based position of the step within the workflow.
45    pub step_index: i64,
46    /// Human-readable name for the step.
47    pub name: String,
48    /// Current status (`"pending"`, `"running"`, `"completed"`, `"failed"`).
49    pub status: String,
50    /// JSON input payload for this step.
51    pub input: Option<Value>,
52    /// JSON output payload produced by this step.
53    pub output: Option<Value>,
54    /// Error message if the step failed.
55    pub error: Option<String>,
56    /// Unix-millisecond timestamp when execution of this step began.
57    pub started_at: Option<i64>,
58    /// Unix-millisecond timestamp when execution of this step finished.
59    pub completed_at: Option<i64>,
60}
61
62/// Manages workflow lifecycle and step tracking.
63pub struct WorkflowStore {
64    conn: Arc<Mutex<Connection>>,
65}
66
67impl WorkflowStore {
68    pub(crate) fn new(conn: Arc<Mutex<Connection>>) -> Self {
69        Self { conn }
70    }
71
72    /// Create a new workflow in `pending` status.
73    pub fn create_workflow(
74        &self,
75        id: &str,
76        name: &str,
77        input: Option<Value>,
78        metadata: Option<Value>,
79    ) -> Result<()> {
80        let conn = self.conn.lock().unwrap();
81        let input_str = input.as_ref().map(|v| v.to_string());
82        let meta_str = metadata.as_ref().map(|v| v.to_string());
83        let now = now_ms();
84        conn.execute(
85            "INSERT INTO _adb_workflows
86                 (id, name, status, input, metadata, created_at, updated_at)
87             VALUES (?1, ?2, 'pending', ?3, ?4, ?5, ?6)",
88            params![id, name, input_str, meta_str, now, now],
89        )?;
90        Ok(())
91    }
92
93    /// Append a step to an existing workflow.
94    ///
95    /// The `step_index` is assigned automatically as one more than the current
96    /// maximum index in the workflow (0-based). Returns the new step ID.
97    pub fn add_step(&self, workflow_id: &str, name: &str, input: Option<Value>) -> Result<String> {
98        let step_id = Uuid::new_v4().to_string();
99        let conn = self.conn.lock().unwrap();
100        let input_str = input.as_ref().map(|v| v.to_string());
101        let step_index: i64 = conn
102            .query_row(
103                "SELECT COALESCE(MAX(step_index) + 1, 0)
104                 FROM _adb_workflow_steps
105                 WHERE workflow_id = ?1",
106                params![workflow_id],
107                |r| r.get(0),
108            )
109            .unwrap_or(0);
110        conn.execute(
111            "INSERT INTO _adb_workflow_steps
112                 (id, workflow_id, step_index, name, status, input)
113             VALUES (?1, ?2, ?3, ?4, 'pending', ?5)",
114            params![step_id, workflow_id, step_index, name, input_str],
115        )?;
116        Ok(step_id)
117    }
118
119    /// Update the status, output, and/or error of a workflow step.
120    ///
121    /// When `status` is `"running"` the `started_at` timestamp is recorded.
122    /// When `status` is `"completed"` or `"failed"` the `completed_at`
123    /// timestamp is recorded.
124    pub fn update_step(
125        &self,
126        step_id: &str,
127        status: &str,
128        output: Option<Value>,
129        error: Option<&str>,
130    ) -> Result<()> {
131        let conn = self.conn.lock().unwrap();
132        let output_str = output.as_ref().map(|v| v.to_string());
133        let now = now_ms();
134        let started_at: Option<i64> = if status == "running" { Some(now) } else { None };
135        let completed_at: Option<i64> = if status == "completed" || status == "failed" {
136            Some(now)
137        } else {
138            None
139        };
140        let changed = conn.execute(
141            "UPDATE _adb_workflow_steps
142             SET status       = ?2,
143                 output       = COALESCE(?3, output),
144                 error        = COALESCE(?4, error),
145                 started_at   = COALESCE(?5, started_at),
146                 completed_at = COALESCE(?6, completed_at)
147             WHERE id = ?1",
148            params![step_id, status, output_str, error, started_at, completed_at],
149        )?;
150        if changed == 0 {
151            return Err(AgentDbError::InvalidArgument(format!(
152                "step not found: {step_id}"
153            )));
154        }
155        Ok(())
156    }
157
158    /// Mark a workflow as `completed` and record its output.
159    pub fn complete_workflow(&self, id: &str, output: Option<Value>) -> Result<()> {
160        let conn = self.conn.lock().unwrap();
161        let output_str = output.as_ref().map(|v| v.to_string());
162        let now = now_ms();
163        let changed = conn.execute(
164            "UPDATE _adb_workflows
165             SET status = 'completed', output = ?2, updated_at = ?3
166             WHERE id = ?1",
167            params![id, output_str, now],
168        )?;
169        if changed == 0 {
170            return Err(AgentDbError::InvalidArgument(format!(
171                "workflow not found: {id}"
172            )));
173        }
174        Ok(())
175    }
176
177    /// Mark a workflow as `failed` and record an optional error message.
178    pub fn fail_workflow(&self, id: &str, error: Option<&str>) -> Result<()> {
179        let conn = self.conn.lock().unwrap();
180        let now = now_ms();
181        let changed = conn.execute(
182            "UPDATE _adb_workflows
183             SET status = 'failed', error = COALESCE(?2, error), updated_at = ?3
184             WHERE id = ?1",
185            params![id, error, now],
186        )?;
187        if changed == 0 {
188            return Err(AgentDbError::InvalidArgument(format!(
189                "workflow not found: {id}"
190            )));
191        }
192        Ok(())
193    }
194
195    /// Retrieve a workflow and all of its steps.
196    pub fn get_workflow(&self, id: &str) -> Result<Workflow> {
197        let workflow = {
198            let conn = self.conn.lock().unwrap();
199            conn.query_row(
200                "SELECT id, name, status, input, output, error, metadata, created_at, updated_at
201                 FROM _adb_workflows
202                 WHERE id = ?1",
203                params![id],
204                parse_workflow_row,
205            )
206            .map_err(|_| AgentDbError::InvalidArgument(format!("workflow not found: {id}")))?
207        };
208        let steps = self.steps_for_workflow(id)?;
209        let step_count = steps.len() as i64;
210        Ok(Workflow { steps, step_count, ..workflow })
211    }
212
213    /// List workflows, optionally filtered by status.
214    ///
215    /// Pass `None` to return all workflows. Results are ordered by `created_at`
216    /// descending (most recent first). The `step_count` field is populated.
217    /// Full step objects are only fetched by [`get_workflow`].
218    pub fn list_workflows(&self, status_filter: Option<&str>) -> Result<Vec<Workflow>> {
219        let conn = self.conn.lock().unwrap();
220        let workflows: Vec<Workflow> = match status_filter {
221            Some(s) => {
222                let mut stmt = conn.prepare(
223                    "SELECT w.id, w.name, w.status, w.input, w.output, w.error,
224                            w.metadata, w.created_at, w.updated_at,
225                            COUNT(s.id) AS step_count
226                     FROM _adb_workflows w
227                     LEFT JOIN _adb_workflow_steps s ON s.workflow_id = w.id
228                     WHERE w.status = ?1
229                     GROUP BY w.id
230                     ORDER BY w.created_at DESC",
231                )?;
232                let rows = stmt.query_map(params![s], parse_workflow_row_with_count)?;
233                rows.map(|r| r.map_err(AgentDbError::Sqlite))
234                    .collect::<Result<Vec<_>>>()?
235            }
236            None => {
237                let mut stmt = conn.prepare(
238                    "SELECT w.id, w.name, w.status, w.input, w.output, w.error,
239                            w.metadata, w.created_at, w.updated_at,
240                            COUNT(s.id) AS step_count
241                     FROM _adb_workflows w
242                     LEFT JOIN _adb_workflow_steps s ON s.workflow_id = w.id
243                     GROUP BY w.id
244                     ORDER BY w.created_at DESC",
245                )?;
246                let rows = stmt.query_map([], parse_workflow_row_with_count)?;
247                rows.map(|r| r.map_err(AgentDbError::Sqlite))
248                    .collect::<Result<Vec<_>>>()?
249            }
250        };
251        Ok(workflows)
252    }
253
254    // ── Internal helpers ─────────────────────────────────────────────────
255
256    fn steps_for_workflow(&self, workflow_id: &str) -> Result<Vec<WorkflowStep>> {
257        let conn = self.conn.lock().unwrap();
258        let mut stmt = conn.prepare(
259            "SELECT id, workflow_id, step_index, name, status,
260                    input, output, error, started_at, completed_at
261             FROM _adb_workflow_steps
262             WHERE workflow_id = ?1
263             ORDER BY step_index ASC",
264        )?;
265        let rows = stmt.query_map(params![workflow_id], parse_step_row)?;
266        rows.map(|r| r.map_err(AgentDbError::Sqlite)).collect()
267    }
268}
269
270// ── Row parsers ──────────────────────────────────────────────────────────────
271
272// Used by get_workflow (no step_count column in query)
273fn parse_workflow_row(row: &rusqlite::Row) -> rusqlite::Result<Workflow> {
274    // Column order: id(0) name(1) status(2) input(3) output(4) error(5) metadata(6) created_at(7) updated_at(8)
275    let input_str: Option<String> = row.get(3)?;
276    let output_str: Option<String> = row.get(4)?;
277    let meta_str: Option<String> = row.get(6)?;
278    Ok(Workflow {
279        id: row.get(0)?,
280        name: row.get(1)?,
281        status: row.get(2)?,
282        input: input_str
283            .as_deref()
284            .and_then(|s| serde_json::from_str(s).ok()),
285        output: output_str
286            .as_deref()
287            .and_then(|s| serde_json::from_str(s).ok()),
288        error: row.get(5)?,
289        metadata: meta_str
290            .as_deref()
291            .and_then(|s| serde_json::from_str(s).ok()),
292        created_at: row.get(7)?,
293        updated_at: row.get(8)?,
294        step_count: 0,
295        steps: vec![],
296    })
297}
298
299// Used by list_workflows (includes step_count at column 9)
300fn parse_workflow_row_with_count(row: &rusqlite::Row) -> rusqlite::Result<Workflow> {
301    let input_str: Option<String> = row.get(3)?;
302    let output_str: Option<String> = row.get(4)?;
303    let meta_str: Option<String> = row.get(6)?;
304    Ok(Workflow {
305        id: row.get(0)?,
306        name: row.get(1)?,
307        status: row.get(2)?,
308        input: input_str
309            .as_deref()
310            .and_then(|s| serde_json::from_str(s).ok()),
311        output: output_str
312            .as_deref()
313            .and_then(|s| serde_json::from_str(s).ok()),
314        error: row.get(5)?,
315        metadata: meta_str
316            .as_deref()
317            .and_then(|s| serde_json::from_str(s).ok()),
318        created_at: row.get(7)?,
319        updated_at: row.get(8)?,
320        step_count: row.get(9)?,
321        steps: vec![],
322    })
323}
324
325fn parse_step_row(row: &rusqlite::Row) -> rusqlite::Result<WorkflowStep> {
326    let input_str: Option<String> = row.get(5)?;
327    let output_str: Option<String> = row.get(6)?;
328    Ok(WorkflowStep {
329        id: row.get(0)?,
330        workflow_id: row.get(1)?,
331        step_index: row.get(2)?,
332        name: row.get(3)?,
333        status: row.get(4)?,
334        input: input_str
335            .as_deref()
336            .and_then(|s| serde_json::from_str(s).ok()),
337        output: output_str
338            .as_deref()
339            .and_then(|s| serde_json::from_str(s).ok()),
340        error: row.get(7)?,
341        started_at: row.get(8)?,
342        completed_at: row.get(9)?,
343    })
344}