systemprompt_agent/repository/task/
mod.rs

1//! Persistence for A2A tasks, their messages, parts, and execution state.
2//!
3//! [`TaskRepository`] is the repository facade over the `agent_tasks` table and
4//! its satellites (`task_messages`, `message_parts`, `task_execution_steps`).
5//! It splits reads and writes across separate pools, optionally drives session
6//! analytics and file-upload providers, and delegates aggregate reassembly to
7//! [`TaskConstructor`]. Query, mutation, and state-transition helpers live in
8//! the sibling submodules and are re-exported here.
9
10pub mod constructor;
11mod mutations;
12mod queries;
13mod state;
14mod task_messages;
15mod task_updates;
16
17pub use constructor::TaskConstructor;
18pub use mutations::{
19    CreateTaskParams, create_task, task_state_to_db_string, track_agent_in_context,
20};
21pub use queries::{
22    TaskContextInfo, get_task, get_task_context_info, get_tasks_by_user_id, list_tasks_by_context,
23};
24pub use state::{apply_notification_status, update_task_failed_with_error, update_task_state};
25pub use task_updates::UpdateTaskAndSaveMessagesParams;
26
27use crate::models::a2a::{Task, TaskState};
28use sqlx::PgPool;
29use std::sync::Arc;
30use systemprompt_database::DbPool;
31use systemprompt_identifiers::{SessionId, TraceId, UserId};
32use systemprompt_traits::{DynFileUploadProvider, DynSessionAnalyticsProvider, RepositoryError};
33
34#[expect(
35    missing_debug_implementations,
36    reason = "params struct holds non-Debug references"
37)]
38pub struct RepoCreateTaskParams<'a> {
39    pub task: &'a Task,
40    pub user_id: &'a UserId,
41    pub session_id: &'a SessionId,
42    pub trace_id: &'a TraceId,
43    pub agent_name: &'a str,
44}
45
46#[derive(Clone)]
47pub struct TaskRepository {
48    pool: Arc<PgPool>,
49    write_pool: Arc<PgPool>,
50    db_pool: DbPool,
51    pub(crate) session_analytics_provider: Option<DynSessionAnalyticsProvider>,
52    pub(crate) file_upload_provider: Option<DynFileUploadProvider>,
53}
54
55impl std::fmt::Debug for TaskRepository {
56    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
57        f.debug_struct("TaskRepository")
58            .field("pool", &"<PgPool>")
59            .field("write_pool", &"<PgPool>")
60            .field("db_pool", &"<DbPool>")
61            .field(
62                "session_analytics_provider",
63                &self.session_analytics_provider.is_some(),
64            )
65            .field("file_upload_provider", &self.file_upload_provider.is_some())
66            .finish()
67    }
68}
69
70impl TaskRepository {
71    pub fn new(db: &DbPool) -> Result<Self, crate::error::AgentError> {
72        let pool = db
73            .pool_arc()
74            .map_err(|e| crate::error::AgentError::Init(e.to_string()))?;
75        let write_pool = db
76            .write_pool_arc()
77            .map_err(|e| crate::error::AgentError::Init(e.to_string()))?;
78        Ok(Self {
79            pool,
80            write_pool,
81            db_pool: Arc::clone(db),
82            session_analytics_provider: None,
83            file_upload_provider: None,
84        })
85    }
86
87    #[must_use]
88    pub fn with_session_analytics_provider(
89        mut self,
90        provider: DynSessionAnalyticsProvider,
91    ) -> Self {
92        self.session_analytics_provider = Some(provider);
93        self
94    }
95
96    #[must_use]
97    pub fn with_file_upload_provider(mut self, provider: DynFileUploadProvider) -> Self {
98        self.file_upload_provider = Some(provider);
99        self
100    }
101
102    pub(crate) const fn db_pool(&self) -> &DbPool {
103        &self.db_pool
104    }
105
106    pub async fn create_task(
107        &self,
108        params: RepoCreateTaskParams<'_>,
109    ) -> Result<String, RepositoryError> {
110        let result = create_task(CreateTaskParams {
111            pool: &self.write_pool,
112            task: params.task,
113            user_id: params.user_id,
114            session_id: params.session_id,
115            trace_id: params.trace_id,
116            agent_name: params.agent_name,
117        })
118        .await?;
119
120        if let Some(ref provider) = self.session_analytics_provider {
121            if let Err(e) = provider.increment_task_count(params.session_id).await {
122                tracing::warn!(error = %e, "Failed to increment analytics task count");
123            }
124        }
125
126        Ok(result)
127    }
128
129    pub async fn get_task(
130        &self,
131        task_id: &systemprompt_identifiers::TaskId,
132    ) -> Result<Option<Task>, RepositoryError> {
133        get_task(&self.pool, &self.db_pool, task_id).await
134    }
135
136    pub async fn list_tasks_by_context(
137        &self,
138        context_id: &systemprompt_identifiers::ContextId,
139    ) -> Result<Vec<Task>, RepositoryError> {
140        list_tasks_by_context(&self.pool, &self.db_pool, context_id).await
141    }
142
143    pub async fn get_tasks_by_user_id(
144        &self,
145        user_id: &UserId,
146        limit: Option<i32>,
147        offset: Option<i32>,
148    ) -> Result<Vec<Task>, RepositoryError> {
149        get_tasks_by_user_id(&self.pool, &self.db_pool, user_id, limit, offset).await
150    }
151
152    pub async fn track_agent_in_context(
153        &self,
154        context_id: &systemprompt_identifiers::ContextId,
155        agent_name: &str,
156    ) -> Result<(), RepositoryError> {
157        track_agent_in_context(&self.write_pool, context_id, agent_name).await
158    }
159
160    pub async fn update_task_state(
161        &self,
162        task_id: &systemprompt_identifiers::TaskId,
163        state: TaskState,
164        timestamp: &chrono::DateTime<chrono::Utc>,
165    ) -> Result<(), RepositoryError> {
166        update_task_state(&self.write_pool, task_id, state, timestamp).await
167    }
168
169    pub async fn apply_notification_status(
170        &self,
171        task_id: &systemprompt_identifiers::TaskId,
172        state: &str,
173        timestamp: &chrono::DateTime<chrono::Utc>,
174    ) -> Result<(), RepositoryError> {
175        apply_notification_status(&self.write_pool, task_id, state, timestamp).await
176    }
177
178    pub async fn update_task_failed_with_error(
179        &self,
180        task_id: &systemprompt_identifiers::TaskId,
181        error_message: &str,
182        timestamp: &chrono::DateTime<chrono::Utc>,
183    ) -> Result<(), RepositoryError> {
184        update_task_failed_with_error(&self.write_pool, task_id, error_message, timestamp).await
185    }
186
187    pub async fn get_task_context_info(
188        &self,
189        task_id: &systemprompt_identifiers::TaskId,
190    ) -> Result<Option<TaskContextInfo>, RepositoryError> {
191        get_task_context_info(&self.pool, task_id).await
192    }
193
194    pub async fn validate_task_ownership(
195        &self,
196        task_id: &systemprompt_identifiers::TaskId,
197        user_id: &UserId,
198    ) -> Result<(), RepositoryError> {
199        let result = sqlx::query_scalar!(
200            "SELECT t.task_id FROM agent_tasks t JOIN user_contexts c ON t.context_id = \
201             c.context_id WHERE t.task_id = $1 AND c.user_id = $2",
202            task_id.as_str(),
203            user_id.as_str()
204        )
205        .fetch_optional(self.pool.as_ref())
206        .await
207        .map_err(RepositoryError::database)?;
208
209        match result {
210            Some(_) => Ok(()),
211            None => Err(RepositoryError::NotFound(format!(
212                "Task {task_id} not found or user {user_id} does not have access"
213            ))),
214        }
215    }
216}
systemprompt_agent/repository/task/mod.rs

systemprompt_agent/repository/task/
mod.rs
