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