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 {
211 steps,
212 step_count,
213 ..workflow
214 })
215 }
216
217 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 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
274fn parse_workflow_row(row: &rusqlite::Row) -> rusqlite::Result<Workflow> {
278 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
303fn 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}