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

systemprompt_agent/services/
execution_tracking.rs
