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    /// Retrieve a workflow and all of its steps.
165    pub fn get_workflow(&self, id: &str) -> Result<Workflow> {
166        let workflow = {
167            let conn = self.conn.lock().unwrap();
168            conn.query_row(
169                "SELECT id, name, status, input, output, metadata, created_at, updated_at
170                 FROM _adb_workflows
171                 WHERE id = ?1",
172                params![id],
173                parse_workflow_row,
174            )
175            .map_err(|_| AgentDbError::InvalidArgument(format!("workflow not found: {id}")))?
176        };
177        let steps = self.steps_for_workflow(id)?;
178        Ok(Workflow { steps, ..workflow })
179    }
180
181    /// List workflows, optionally filtered by status.
182    ///
183    /// Pass `None` to return all workflows. Results are ordered by `created_at`
184    /// descending (most recent first). Steps are **not** populated on list results
185    /// to keep the query lightweight; call [`get_workflow`] for step details.
186    pub fn list_workflows(&self, status_filter: Option<&str>) -> Result<Vec<Workflow>> {
187        let conn = self.conn.lock().unwrap();
188        let workflows: Vec<Workflow> = match status_filter {
189            Some(s) => {
190                let mut stmt = conn.prepare(
191                    "SELECT id, name, status, input, output, metadata, created_at, updated_at
192                     FROM _adb_workflows
193                     WHERE status = ?1
194                     ORDER BY created_at DESC",
195                )?;
196                let rows = stmt.query_map(params![s], parse_workflow_row)?;
197                rows.map(|r| r.map_err(AgentDbError::Sqlite))
198                    .collect::<Result<Vec<_>>>()?
199            }
200            None => {
201                let mut stmt = conn.prepare(
202                    "SELECT id, name, status, input, output, metadata, created_at, updated_at
203                     FROM _adb_workflows
204                     ORDER BY created_at DESC",
205                )?;
206                let rows = stmt.query_map([], parse_workflow_row)?;
207                rows.map(|r| r.map_err(AgentDbError::Sqlite))
208                    .collect::<Result<Vec<_>>>()?
209            }
210        };
211        Ok(workflows)
212    }
213
214    // ── Internal helpers ─────────────────────────────────────────────────
215
216    fn steps_for_workflow(&self, workflow_id: &str) -> Result<Vec<WorkflowStep>> {
217        let conn = self.conn.lock().unwrap();
218        let mut stmt = conn.prepare(
219            "SELECT id, workflow_id, step_index, name, status,
220                    input, output, error, started_at, completed_at
221             FROM _adb_workflow_steps
222             WHERE workflow_id = ?1
223             ORDER BY step_index ASC",
224        )?;
225        let rows = stmt.query_map(params![workflow_id], parse_step_row)?;
226        rows.map(|r| r.map_err(AgentDbError::Sqlite)).collect()
227    }
228}
229
230// ── Row parsers ──────────────────────────────────────────────────────────────
231
232fn parse_workflow_row(row: &rusqlite::Row) -> rusqlite::Result<Workflow> {
233    let input_str: Option<String> = row.get(3)?;
234    let output_str: Option<String> = row.get(4)?;
235    let meta_str: Option<String> = row.get(5)?;
236    Ok(Workflow {
237        id: row.get(0)?,
238        name: row.get(1)?,
239        status: row.get(2)?,
240        input: input_str
241            .as_deref()
242            .and_then(|s| serde_json::from_str(s).ok()),
243        output: output_str
244            .as_deref()
245            .and_then(|s| serde_json::from_str(s).ok()),
246        metadata: meta_str
247            .as_deref()
248            .and_then(|s| serde_json::from_str(s).ok()),
249        created_at: row.get(6)?,
250        updated_at: row.get(7)?,
251        steps: vec![],
252    })
253}
254
255fn parse_step_row(row: &rusqlite::Row) -> rusqlite::Result<WorkflowStep> {
256    let input_str: Option<String> = row.get(5)?;
257    let output_str: Option<String> = row.get(6)?;
258    Ok(WorkflowStep {
259        id: row.get(0)?,
260        workflow_id: row.get(1)?,
261        step_index: row.get(2)?,
262        name: row.get(3)?,
263        status: row.get(4)?,
264        input: input_str
265            .as_deref()
266            .and_then(|s| serde_json::from_str(s).ok()),
267        output: output_str
268            .as_deref()
269            .and_then(|s| serde_json::from_str(s).ok()),
270        error: row.get(7)?,
271        started_at: row.get(8)?,
272        completed_at: row.get(9)?,
273    })
274}