systemprompt_agent/services/a2a_server/processing/
persistence_service.rs

1//! Task persistence for the processing pipeline.
2//!
3//! [`PersistenceService`] wraps the A2A repositories to create tasks, update
4//! their state, and persist a completed task together with its messages —
5//! publishing any attached artifacts unless they were already published
6//! upstream.
7//!
8//! Copyright (c) systemprompt.io — Business Source License 1.1.
9//! See <https://systemprompt.io> for licensing details.
10
11use crate::services::shared::{AgentServiceError, Result};
12use std::sync::Arc;
13use systemprompt_identifiers::{SessionId, TaskId, TraceId, UserId};
14use systemprompt_models::{RequestContext, TaskMetadata};
15
16use crate::models::{Message, Task, TaskState, TaskStatus};
17use crate::repository::A2ARepositories;
18use crate::repository::task::UpdateTaskAndSaveMessagesParams;
19use crate::services::ArtifactPublishingService;
20
21#[derive(Debug)]
22pub struct PersistCompletedTaskServiceParams<'a> {
23    pub task: &'a Task,
24    pub user_message: &'a Message,
25    pub agent_message: &'a Message,
26    pub context: &'a RequestContext,
27    pub artifacts_already_published: bool,
28}
29
30pub struct PersistenceService {
31    repositories: Arc<A2ARepositories>,
32}
33
34impl std::fmt::Debug for PersistenceService {
35    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36        f.debug_struct("PersistenceService").finish_non_exhaustive()
37    }
38}
39
40impl PersistenceService {
41    pub const fn new(repositories: Arc<A2ARepositories>) -> Self {
42        Self { repositories }
43    }
44
45    pub async fn create_task(
46        &self,
47        task: &Task,
48        context: &RequestContext,
49        agent_name: &str,
50    ) -> Result<()> {
51        self.repositories
52            .tasks
53            .create_task(crate::repository::task::RepoCreateTaskParams {
54                task,
55                user_id: &UserId::new(context.user_id().as_str()),
56                session_id: &SessionId::new(context.session_id().as_str()),
57                trace_id: &TraceId::new(context.trace_id().as_str()),
58                agent_name,
59            })
60            .await
61            .map_err(|e| {
62                AgentServiceError::Internal(format!("Failed to persist task at start: {}", e))
63            })?;
64
65        tracing::info!(task_id = %task.id, "Task persisted to database");
66
67        Ok(())
68    }
69
70    pub async fn update_task_state(
71        &self,
72        task_id: &TaskId,
73        state: TaskState,
74        timestamp: &chrono::DateTime<chrono::Utc>,
75    ) -> Result<()> {
76        self.repositories
77            .tasks
78            .update_task_state(task_id, state, timestamp)
79            .await
80            .map_err(|e| AgentServiceError::Internal(format!("Failed to update task state: {e}")))
81    }
82
83    pub async fn persist_completed_task(
84        &self,
85        params: PersistCompletedTaskServiceParams<'_>,
86    ) -> Result<Task> {
87        let PersistCompletedTaskServiceParams {
88            task,
89            user_message,
90            agent_message,
91            context,
92            artifacts_already_published,
93        } = params;
94        let updated_task = self
95            .repositories
96            .tasks
97            .update_task_and_save_messages(UpdateTaskAndSaveMessagesParams {
98                task,
99                user_message,
100                agent_message,
101                user_id: Some(context.user_id()),
102                session_id: context.session_id(),
103                trace_id: context.trace_id(),
104            })
105            .await
106            .map_err(|e| {
107                AgentServiceError::Internal(format!(
108                    "Failed to update task and save messages: {}",
109                    e
110                ))
111            })?;
112
113        if !artifacts_already_published && let Some(artifacts) = &task.artifacts {
114            let context_id = &task.context_id;
115            let publishing_service = ArtifactPublishingService::new(
116                self.repositories.artifacts.clone(),
117                self.repositories.execution_steps.clone(),
118                self.repositories.tasks.clone(),
119            )?;
120            for artifact in artifacts {
121                publishing_service
122                    .publish_from_a2a(artifact, &task.id, context_id)
123                    .await
124                    .map_err(|e| {
125                        AgentServiceError::Internal(format!(
126                            "Failed to publish artifact {}: {}",
127                            artifact.id, e
128                        ))
129                    })?;
130            }
131
132            tracing::info!(
133                task_id = %task.id,
134                artifact_count = artifacts.len(),
135                "Published artifacts for task"
136            );
137        }
138
139        tracing::info!(
140            task_id = %task.id,
141            context_id = ?task.context_id,
142            user_id = %context.user_id(),
143            "Persisted task"
144        );
145
146        Ok(updated_task)
147    }
148
149    pub fn build_initial_task(
150        task_id: TaskId,
151        context_id: systemprompt_identifiers::ContextId,
152        agent_name: &str,
153    ) -> Task {
154        let metadata = TaskMetadata::new_agent_message(agent_name.to_owned());
155
156        Task {
157            id: task_id,
158            context_id,
159            status: TaskStatus {
160                state: TaskState::Submitted,
161                message: None,
162                timestamp: Some(chrono::Utc::now()),
163            },
164            history: None,
165            artifacts: None,
166            metadata: Some(metadata),
167            created_at: Some(chrono::Utc::now()),
168            last_modified: Some(chrono::Utc::now()),
169        }
170    }
171}
systemprompt_agent/services/a2a_server/processing/persistence_service.rs

systemprompt_agent/services/a2a_server/processing/
persistence_service.rs
