systemprompt_agent/services/
execution_tracking.rs

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

systemprompt_agent/services/
execution_tracking.rs
