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 {
211            steps,
212            step_count,
213            ..workflow
214        })
215    }
216
217    /// List workflows, optionally filtered by status.
218    ///
219    /// Pass `None` to return all workflows. Results are ordered by `created_at`
220    /// descending (most recent first). The `step_count` field is populated.
221    /// Full step objects are only fetched by [`get_workflow`].
222    pub fn list_workflows(&self, status_filter: Option<&str>) -> Result<Vec<Workflow>> {
223        let conn = self.conn.lock().unwrap();
224        let workflows: Vec<Workflow> = match status_filter {
225            Some(s) => {
226                let mut stmt = conn.prepare(
227                    "SELECT w.id, w.name, w.status, w.input, w.output, w.error,
228                            w.metadata, w.created_at, w.updated_at,
229                            COUNT(s.id) AS step_count
230                     FROM _adb_workflows w
231                     LEFT JOIN _adb_workflow_steps s ON s.workflow_id = w.id
232                     WHERE w.status = ?1
233                     GROUP BY w.id
234                     ORDER BY w.created_at DESC",
235                )?;
236                let rows = stmt.query_map(params![s], parse_workflow_row_with_count)?;
237                rows.map(|r| r.map_err(AgentDbError::Sqlite))
238                    .collect::<Result<Vec<_>>>()?
239            }
240            None => {
241                let mut stmt = conn.prepare(
242                    "SELECT w.id, w.name, w.status, w.input, w.output, w.error,
243                            w.metadata, w.created_at, w.updated_at,
244                            COUNT(s.id) AS step_count
245                     FROM _adb_workflows w
246                     LEFT JOIN _adb_workflow_steps s ON s.workflow_id = w.id
247                     GROUP BY w.id
248                     ORDER BY w.created_at DESC",
249                )?;
250                let rows = stmt.query_map([], parse_workflow_row_with_count)?;
251                rows.map(|r| r.map_err(AgentDbError::Sqlite))
252                    .collect::<Result<Vec<_>>>()?
253            }
254        };
255        Ok(workflows)
256    }
257
258    // ── Internal helpers ─────────────────────────────────────────────────
259
260    fn steps_for_workflow(&self, workflow_id: &str) -> Result<Vec<WorkflowStep>> {
261        let conn = self.conn.lock().unwrap();
262        let mut stmt = conn.prepare(
263            "SELECT id, workflow_id, step_index, name, status,
264                    input, output, error, started_at, completed_at
265             FROM _adb_workflow_steps
266             WHERE workflow_id = ?1
267             ORDER BY step_index ASC",
268        )?;
269        let rows = stmt.query_map(params![workflow_id], parse_step_row)?;
270        rows.map(|r| r.map_err(AgentDbError::Sqlite)).collect()
271    }
272}
273
274// ── Row parsers ──────────────────────────────────────────────────────────────
275
276// Used by get_workflow (no step_count column in query)
277fn parse_workflow_row(row: &rusqlite::Row) -> rusqlite::Result<Workflow> {
278    // Column order: id(0) name(1) status(2) input(3) output(4) error(5) metadata(6) created_at(7) updated_at(8)
279    let input_str: Option<String> = row.get(3)?;
280    let output_str: Option<String> = row.get(4)?;
281    let meta_str: Option<String> = row.get(6)?;
282    Ok(Workflow {
283        id: row.get(0)?,
284        name: row.get(1)?,
285        status: row.get(2)?,
286        input: input_str
287            .as_deref()
288            .and_then(|s| serde_json::from_str(s).ok()),
289        output: output_str
290            .as_deref()
291            .and_then(|s| serde_json::from_str(s).ok()),
292        error: row.get(5)?,
293        metadata: meta_str
294            .as_deref()
295            .and_then(|s| serde_json::from_str(s).ok()),
296        created_at: row.get(7)?,
297        updated_at: row.get(8)?,
298        step_count: 0,
299        steps: vec![],
300    })
301}
302
303// Used by list_workflows (includes step_count at column 9)
304fn parse_workflow_row_with_count(row: &rusqlite::Row) -> rusqlite::Result<Workflow> {
305    let input_str: Option<String> = row.get(3)?;
306    let output_str: Option<String> = row.get(4)?;
307    let meta_str: Option<String> = row.get(6)?;
308    Ok(Workflow {
309        id: row.get(0)?,
310        name: row.get(1)?,
311        status: row.get(2)?,
312        input: input_str
313            .as_deref()
314            .and_then(|s| serde_json::from_str(s).ok()),
315        output: output_str
316            .as_deref()
317            .and_then(|s| serde_json::from_str(s).ok()),
318        error: row.get(5)?,
319        metadata: meta_str
320            .as_deref()
321            .and_then(|s| serde_json::from_str(s).ok()),
322        created_at: row.get(7)?,
323        updated_at: row.get(8)?,
324        step_count: row.get(9)?,
325        steps: vec![],
326    })
327}
328
329fn parse_step_row(row: &rusqlite::Row) -> rusqlite::Result<WorkflowStep> {
330    let input_str: Option<String> = row.get(5)?;
331    let output_str: Option<String> = row.get(6)?;
332    Ok(WorkflowStep {
333        id: row.get(0)?,
334        workflow_id: row.get(1)?,
335        step_index: row.get(2)?,
336        name: row.get(3)?,
337        status: row.get(4)?,
338        input: input_str
339            .as_deref()
340            .and_then(|s| serde_json::from_str(s).ok()),
341        output: output_str
342            .as_deref()
343            .and_then(|s| serde_json::from_str(s).ok()),
344        error: row.get(7)?,
345        started_at: row.get(8)?,
346        completed_at: row.get(9)?,
347    })
348}