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//!
7//! Copyright (c) systemprompt.io — Business Source License 1.1.
8//! See <https://systemprompt.io> for licensing details.
9
10mod 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}