systemprompt_agent/repository/execution/
mutations.rs1use chrono::{DateTime, Utc};
8use systemprompt_identifiers::TaskId;
9use systemprompt_models::{ExecutionStep, PlannedTool, StepContent, StepId, StepStatus};
10use systemprompt_traits::RepositoryError;
11
12use super::ExecutionStepRepository;
13use super::parse::{ParseStepParams, parse_step};
14
15impl ExecutionStepRepository {
16 pub async fn create(&self, step: &ExecutionStep) -> Result<(), RepositoryError> {
17 let step_id_str = step.step_id.as_str();
18 let task_id = &step.task_id;
19 let status_str = step.status.to_string();
20 let step_type_str = step.content.step_type().to_string();
21 let title = step.content.title();
22 let content_json = serde_json::to_value(&step.content).map_err(|e| {
23 RepositoryError::Internal(format!("Failed to serialize step content: {e}"))
24 })?;
25 sqlx::query!(
26 r#"INSERT INTO task_execution_steps (
27 step_id, task_id, step_type, title, status, content, started_at, completed_at, duration_ms, error_message
28 ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)"#,
29 step_id_str,
30 task_id.as_str(),
31 step_type_str,
32 title,
33 status_str,
34 content_json,
35 step.started_at,
36 step.completed_at,
37 step.duration_ms,
38 step.error_message
39 )
40 .execute(&*self.write_pool)
41 .await
42 .map_err(|e| RepositoryError::Internal(format!("Failed to create execution step: {e}")))?;
43 Ok(())
44 }
45
46 pub async fn complete_step(
47 &self,
48 step_id: &StepId,
49 started_at: DateTime<Utc>,
50 tool_result: Option<serde_json::Value>,
51 ) -> Result<(), RepositoryError> {
52 let completed_at = Utc::now();
53 let duration_ms = (completed_at - started_at).num_milliseconds() as i32;
54 let step_id_str = step_id.as_str();
55 let status_str = StepStatus::Completed.to_string();
56
57 if let Some(result) = tool_result {
58 sqlx::query!(
59 r#"UPDATE task_execution_steps SET
60 status = $2,
61 completed_at = $3,
62 duration_ms = $4,
63 content = jsonb_set(content, '{tool_result}', $5)
64 WHERE step_id = $1"#,
65 step_id_str,
66 status_str,
67 completed_at,
68 duration_ms,
69 result
70 )
71 .execute(&*self.write_pool)
72 .await
73 .map_err(|e| {
74 RepositoryError::Internal(format!(
75 "Failed to complete execution step: {step_id}: {e}"
76 ))
77 })?;
78 } else {
79 sqlx::query!(
80 r#"UPDATE task_execution_steps SET
81 status = $2,
82 completed_at = $3,
83 duration_ms = $4
84 WHERE step_id = $1"#,
85 step_id_str,
86 status_str,
87 completed_at,
88 duration_ms
89 )
90 .execute(&*self.write_pool)
91 .await
92 .map_err(|e| {
93 RepositoryError::Internal(format!(
94 "Failed to complete execution step: {step_id}: {e}"
95 ))
96 })?;
97 }
98
99 Ok(())
100 }
101
102 pub async fn fail_step(
103 &self,
104 step_id: &StepId,
105 started_at: DateTime<Utc>,
106 error_message: &str,
107 ) -> Result<(), RepositoryError> {
108 let completed_at = Utc::now();
109 let duration_ms = (completed_at - started_at).num_milliseconds() as i32;
110 let step_id_str = step_id.as_str();
111 let status_str = StepStatus::Failed.to_string();
112
113 sqlx::query!(
114 r#"UPDATE task_execution_steps SET
115 status = $2,
116 completed_at = $3,
117 duration_ms = $4,
118 error_message = $5
119 WHERE step_id = $1"#,
120 step_id_str,
121 status_str,
122 completed_at,
123 duration_ms,
124 error_message
125 )
126 .execute(&*self.write_pool)
127 .await
128 .map_err(|e| {
129 RepositoryError::Internal(format!("Failed to fail execution step: {step_id}: {e}"))
130 })?;
131
132 Ok(())
133 }
134
135 pub async fn fail_in_progress_steps_for_task(
136 &self,
137 task_id: &TaskId,
138 error_message: &str,
139 ) -> Result<u64, RepositoryError> {
140 let completed_at = Utc::now();
141 let in_progress_str = StepStatus::InProgress.to_string();
142 let failed_str = StepStatus::Failed.to_string();
143 let task_id_str = task_id.as_str();
144
145 let result = sqlx::query!(
146 r#"UPDATE task_execution_steps SET
147 status = $3,
148 completed_at = $4,
149 error_message = $5
150 WHERE task_id = $1 AND status = $2"#,
151 task_id_str,
152 in_progress_str,
153 failed_str,
154 completed_at,
155 error_message
156 )
157 .execute(&*self.write_pool)
158 .await
159 .map_err(|e| {
160 RepositoryError::Internal(format!(
161 "Failed to fail in-progress steps for task: {task_id}: {e}"
162 ))
163 })?;
164
165 Ok(result.rows_affected())
166 }
167
168 pub async fn complete_planning_step(
169 &self,
170 step_id: &StepId,
171 started_at: DateTime<Utc>,
172 reasoning: Option<String>,
173 planned_tools: Option<Vec<PlannedTool>>,
174 ) -> Result<ExecutionStep, RepositoryError> {
175 let completed_at = Utc::now();
176 let duration_ms = (completed_at - started_at).num_milliseconds() as i32;
177 let step_id_str = step_id.as_str();
178 let status_str = StepStatus::Completed.to_string();
179
180 let content = StepContent::planning(reasoning, planned_tools);
181 let content_json = serde_json::to_value(&content).map_err(|e| {
182 RepositoryError::Internal(format!("Failed to serialize planning content: {e}"))
183 })?;
184
185 let row = sqlx::query!(
186 r#"UPDATE task_execution_steps SET
187 status = $2,
188 completed_at = $3,
189 duration_ms = $4,
190 content = $5
191 WHERE step_id = $1
192 RETURNING step_id, task_id as "task_id!: TaskId", status, content,
193 started_at as "started_at!", completed_at, duration_ms, error_message"#,
194 step_id_str,
195 status_str,
196 completed_at,
197 duration_ms,
198 content_json
199 )
200 .fetch_one(&*self.write_pool)
201 .await
202 .map_err(|e| {
203 RepositoryError::Internal(format!("Failed to complete planning step: {step_id}: {e}"))
204 })?;
205
206 parse_step(ParseStepParams {
207 step_id: row.step_id,
208 task_id: row.task_id,
209 status: row.status,
210 content: row.content,
211 started_at: row.started_at,
212 completed_at: row.completed_at,
213 duration_ms: row.duration_ms,
214 error_message: row.error_message,
215 })
216 }
217}