systemprompt_agent/services/
execution_tracking.rs

1//! Recording agent execution steps (understanding, planning, skill usage, tool
2//! calls, completion) and transitioning them through their lifecycle.
3//!
4//! Copyright (c) systemprompt.io — Business Source License 1.1.
5//! See <https://systemprompt.io> for licensing details.
6
7use crate::services::shared::Result;
8use chrono::{DateTime, Utc};
9use std::sync::Arc;
10use systemprompt_identifiers::{SkillId, TaskId};
11use systemprompt_models::{ExecutionStep, PlannedTool, StepContent, StepId, TrackedStep};
12
13use crate::repository::execution::ExecutionStepRepository;
14
15#[derive(Debug, Clone)]
16pub struct ExecutionTrackingService {
17    repository: Arc<ExecutionStepRepository>,
18}
19
20impl ExecutionTrackingService {
21    pub const fn new(repository: Arc<ExecutionStepRepository>) -> Self {
22        Self { repository }
23    }
24
25    pub async fn track(&self, task_id: TaskId, content: StepContent) -> Result<ExecutionStep> {
26        let step = ExecutionStep::new(task_id, content);
27        self.repository.create(&step).await?;
28        Ok(step)
29    }
30
31    pub async fn track_async(
32        &self,
33        task_id: TaskId,
34        content: StepContent,
35    ) -> Result<(TrackedStep, ExecutionStep)> {
36        let step = ExecutionStep::new(task_id, content);
37        self.repository.create(&step).await?;
38
39        let tracked = TrackedStep {
40            step_id: step.step_id.clone(),
41            started_at: step.started_at,
42        };
43
44        Ok((tracked, step))
45    }
46
47    pub async fn complete(
48        &self,
49        tracked: TrackedStep,
50        result: Option<serde_json::Value>,
51    ) -> Result<()> {
52        self.repository
53            .complete_step(&tracked.step_id, tracked.started_at, result)
54            .await
55            .map_err(Into::into)
56    }
57
58    pub async fn complete_planning(
59        &self,
60        tracked: TrackedStep,
61        reasoning: Option<String>,
62        planned_tools: Option<Vec<PlannedTool>>,
63    ) -> Result<ExecutionStep> {
64        self.repository
65            .complete_planning_step(
66                &tracked.step_id,
67                tracked.started_at,
68                reasoning,
69                planned_tools,
70            )
71            .await
72            .map_err(Into::into)
73    }
74
75    pub async fn fail(&self, tracked: &TrackedStep, error: String) -> Result<()> {
76        self.repository
77            .fail_step(&tracked.step_id, tracked.started_at, &error)
78            .await
79            .map_err(Into::into)
80    }
81
82    pub async fn fail_step(
83        &self,
84        step_id: &StepId,
85        started_at: DateTime<Utc>,
86        error: String,
87    ) -> Result<()> {
88        self.repository
89            .fail_step(step_id, started_at, &error)
90            .await
91            .map_err(Into::into)
92    }
93
94    pub async fn list_steps_by_task(&self, task_id: &TaskId) -> Result<Vec<ExecutionStep>> {
95        self.repository
96            .list_by_task(task_id)
97            .await
98            .map_err(Into::into)
99    }
100
101    pub async fn find_step(&self, step_id: &StepId) -> Result<Option<ExecutionStep>> {
102        self.repository.get(step_id).await.map_err(Into::into)
103    }
104
105    pub async fn fail_in_progress_steps(&self, task_id: &TaskId, error: &str) -> Result<u64> {
106        self.repository
107            .fail_in_progress_steps_for_task(task_id, error)
108            .await
109            .map_err(Into::into)
110    }
111
112    pub async fn track_understanding(&self, task_id: TaskId) -> Result<ExecutionStep> {
113        self.track(task_id, StepContent::understanding()).await
114    }
115
116    pub async fn track_planning_async(
117        &self,
118        task_id: TaskId,
119        reasoning: Option<String>,
120        planned_tools: Option<Vec<PlannedTool>>,
121    ) -> Result<(TrackedStep, ExecutionStep)> {
122        self.track_async(task_id, StepContent::planning(reasoning, planned_tools))
123            .await
124    }
125
126    pub async fn track_skill_usage(
127        &self,
128        task_id: TaskId,
129        skill_id: SkillId,
130        skill_name: impl Into<String>,
131    ) -> Result<ExecutionStep> {
132        self.track(task_id, StepContent::skill_usage(skill_id, skill_name))
133            .await
134    }
135
136    pub async fn track_tool_execution(
137        &self,
138        task_id: TaskId,
139        tool_name: impl Into<String>,
140        tool_arguments: serde_json::Value,
141    ) -> Result<(TrackedStep, ExecutionStep)> {
142        self.track_async(
143            task_id,
144            StepContent::tool_execution(tool_name, tool_arguments),
145        )
146        .await
147    }
148
149    pub async fn track_completion(&self, task_id: TaskId) -> Result<ExecutionStep> {
150        self.track(task_id, StepContent::completion()).await
151    }
152}
systemprompt_agent/services/execution_tracking.rs

systemprompt_agent/services/
execution_tracking.rs
