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)]
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 metadata: Option<Value>,
24 pub created_at: i64,
26 pub updated_at: i64,
28 pub steps: Vec<WorkflowStep>,
30}
31
32#[derive(Debug, Clone)]
34pub struct WorkflowStep {
35 pub id: String,
37 pub workflow_id: String,
39 pub step_index: i64,
41 pub name: String,
43 pub status: String,
45 pub input: Option<Value>,
47 pub output: Option<Value>,
49 pub error: Option<String>,
51 pub started_at: Option<i64>,
53 pub completed_at: Option<i64>,
55}
56
57pub 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 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 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 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 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 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 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 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 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
248fn 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}