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#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
11pub struct Workflow {
12 pub id: String,
14 pub name: String,
16 pub status: String,
18 pub input: Option<Value>,
20 pub output: Option<Value>,
22 pub error: Option<String>,
24 pub metadata: Option<Value>,
26 pub created_at: i64,
28 pub updated_at: i64,
30 pub step_count: i64,
33 pub steps: Vec<WorkflowStep>,
35}
36
37#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
39pub struct WorkflowStep {
40 pub id: String,
42 pub workflow_id: String,
44 pub step_index: i64,
46 pub name: String,
48 pub status: String,
50 pub input: Option<Value>,
52 pub output: Option<Value>,
54 pub error: Option<String>,
56 pub started_at: Option<i64>,
58 pub completed_at: Option<i64>,
60}
61
62pub 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 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 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 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 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 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 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 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 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
270fn parse_workflow_row(row: &rusqlite::Row) -> rusqlite::Result<Workflow> {
274 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
299fn 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}