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