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

systemprompt_agent/repository/task/
mod.rs
