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 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 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 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
230fn 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}