systemprompt_agent/repository/execution/
mutations.rs

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

systemprompt_agent/repository/execution/
mutations.rs
