systemprompt_agent/repository/execution/
mod.rs1mod mutations;
8mod parse;
9
10use sqlx::PgPool;
11use std::sync::Arc;
12use systemprompt_database::DbPool;
13use systemprompt_identifiers::TaskId;
14use systemprompt_models::{ExecutionStep, StepId};
15use systemprompt_traits::RepositoryError;
16
17use parse::{ParseStepParams, parse_step};
18
19#[derive(Debug, Clone)]
20pub struct ExecutionStepRepository {
21 pool: Arc<PgPool>,
22 write_pool: Arc<PgPool>,
23}
24
25impl ExecutionStepRepository {
26 pub fn new(db: &DbPool) -> Result<Self, crate::error::AgentError> {
27 let pool = db
28 .pool_arc()
29 .map_err(|e| crate::error::AgentError::Init(e.to_string()))?;
30 let write_pool = db
31 .write_pool_arc()
32 .map_err(|e| crate::error::AgentError::Init(e.to_string()))?;
33 Ok(Self { pool, write_pool })
34 }
35
36 pub async fn get(&self, step_id: &StepId) -> Result<Option<ExecutionStep>, RepositoryError> {
37 let step_id_str = step_id.as_str();
38 let row = sqlx::query!(
39 r#"SELECT step_id, task_id as "task_id!: TaskId", status, content,
40 started_at as "started_at!", completed_at, duration_ms, error_message
41 FROM task_execution_steps WHERE step_id = $1"#,
42 step_id_str
43 )
44 .fetch_optional(&*self.pool)
45 .await
46 .map_err(|e| {
47 RepositoryError::Internal(format!("Failed to get execution step: {step_id}: {e}"))
48 })?;
49 row.map(|r| {
50 parse_step(ParseStepParams {
51 step_id: r.step_id,
52 task_id: r.task_id,
53 status: r.status,
54 content: r.content,
55 started_at: r.started_at,
56 completed_at: r.completed_at,
57 duration_ms: r.duration_ms,
58 error_message: r.error_message,
59 })
60 })
61 .transpose()
62 }
63
64 pub async fn list_by_task(
65 &self,
66 task_id: &TaskId,
67 ) -> Result<Vec<ExecutionStep>, RepositoryError> {
68 let rows = sqlx::query!(
69 r#"SELECT step_id, task_id as "task_id!: TaskId", status, content,
70 started_at as "started_at!", completed_at, duration_ms, error_message
71 FROM task_execution_steps WHERE task_id = $1 ORDER BY started_at ASC"#,
72 task_id.as_str()
73 )
74 .fetch_all(&*self.pool)
75 .await
76 .map_err(|e| {
77 RepositoryError::Internal(format!(
78 "Failed to list execution steps for task: {task_id}: {e}"
79 ))
80 })?;
81 rows.into_iter()
82 .map(|r| {
83 parse_step(ParseStepParams {
84 step_id: r.step_id,
85 task_id: r.task_id,
86 status: r.status,
87 content: r.content,
88 started_at: r.started_at,
89 completed_at: r.completed_at,
90 duration_ms: r.duration_ms,
91 error_message: r.error_message,
92 })
93 })
94 .collect()
95 }
96
97 pub async fn mcp_execution_id_exists(
98 &self,
99 mcp_execution_id: &str,
100 ) -> Result<bool, RepositoryError> {
101 let exists = sqlx::query_scalar!(
102 r#"SELECT EXISTS(SELECT 1 FROM mcp_tool_executions WHERE mcp_execution_id = $1) as "exists!""#,
103 mcp_execution_id
104 )
105 .fetch_one(&*self.pool)
106 .await
107 .map_err(|e| RepositoryError::Internal(format!("Failed to check mcp_execution_id existence: {e}")))?;
108
109 Ok(exists)
110 }
111}