Skip to main content

systemprompt_agent/repository/execution/
mod.rs

1//! Repository for `task_execution_steps` — per-task tool calls and intermediate
2//! state.
3//!
4//! Read paths live here; write paths (create, complete, fail) live in the
5//! `mutations` submodule.
6
7mod 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}