systemprompt_agent/services/a2a_server/processing/
persistence_service.rs

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

systemprompt_agent/services/a2a_server/processing/
persistence_service.rs
