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::{PersistMessagesTxParams, 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    constructor: TaskConstructor,
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            .finish_non_exhaustive()
64    }
65}
66
67impl TaskRepository {
68    pub fn new(
69        db: &DbPool,
70        sessions: DynSessionUsageCounters,
71    ) -> 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            constructor: TaskConstructor::new(db)?,
82            sessions,
83        })
84    }
85
86    pub async fn create_task(
87        &self,
88        params: RepoCreateTaskParams<'_>,
89    ) -> Result<String, RepositoryError> {
90        let result = create_task(CreateTaskParams {
91            pool: &self.write_pool,
92            task: params.task,
93            user_id: params.user_id,
94            session_id: params.session_id,
95            trace_id: params.trace_id,
96            agent_name: params.agent_name,
97        })
98        .await?;
99
100        if let Err(e) = self.sessions.increment_task_count(params.session_id).await {
101            tracing::warn!(error = %e, session_id = %params.session_id, "Failed to increment session task count");
102        }
103
104        Ok(result)
105    }
106
107    pub async fn get_task(
108        &self,
109        task_id: &systemprompt_identifiers::TaskId,
110    ) -> Result<Option<Task>, RepositoryError> {
111        get_task(&self.constructor, task_id).await
112    }
113
114    pub async fn list_tasks_by_context(
115        &self,
116        context_id: &systemprompt_identifiers::ContextId,
117    ) -> Result<Vec<Task>, RepositoryError> {
118        list_tasks_by_context(&self.pool, &self.constructor, context_id).await
119    }
120
121    pub async fn get_tasks_by_user_id(
122        &self,
123        user_id: &UserId,
124        limit: Option<i32>,
125        offset: Option<i32>,
126    ) -> Result<Vec<Task>, RepositoryError> {
127        get_tasks_by_user_id(&self.pool, &self.constructor, user_id, limit, offset).await
128    }
129
130    pub async fn track_agent_in_context(
131        &self,
132        context_id: &systemprompt_identifiers::ContextId,
133        agent_name: &str,
134    ) -> Result<(), RepositoryError> {
135        track_agent_in_context(&self.write_pool, context_id, agent_name).await
136    }
137
138    pub async fn update_task_state(
139        &self,
140        task_id: &systemprompt_identifiers::TaskId,
141        state: TaskState,
142        timestamp: &chrono::DateTime<chrono::Utc>,
143    ) -> Result<(), RepositoryError> {
144        update_task_state(&self.write_pool, task_id, state, timestamp).await
145    }
146
147    pub async fn apply_notification_status(
148        &self,
149        task_id: &systemprompt_identifiers::TaskId,
150        state: &str,
151        timestamp: &chrono::DateTime<chrono::Utc>,
152    ) -> Result<(), RepositoryError> {
153        apply_notification_status(&self.write_pool, task_id, state, timestamp).await
154    }
155
156    pub async fn update_task_failed_with_error(
157        &self,
158        task_id: &systemprompt_identifiers::TaskId,
159        error_message: &str,
160        timestamp: &chrono::DateTime<chrono::Utc>,
161    ) -> Result<(), RepositoryError> {
162        update_task_failed_with_error(&self.write_pool, task_id, error_message, timestamp).await
163    }
164
165    pub async fn get_task_context_info(
166        &self,
167        task_id: &systemprompt_identifiers::TaskId,
168    ) -> Result<Option<TaskContextInfo>, RepositoryError> {
169        get_task_context_info(&self.pool, task_id).await
170    }
171
172    pub async fn validate_task_ownership(
173        &self,
174        task_id: &systemprompt_identifiers::TaskId,
175        user_id: &UserId,
176    ) -> Result<(), RepositoryError> {
177        let result = sqlx::query_scalar!(
178            "SELECT t.task_id FROM agent_tasks t JOIN user_contexts c ON t.context_id = \
179             c.context_id WHERE t.task_id = $1 AND c.user_id = $2",
180            task_id.as_str(),
181            user_id.as_str()
182        )
183        .fetch_optional(self.pool.as_ref())
184        .await
185        .map_err(RepositoryError::database)?;
186
187        match result {
188            Some(_) => Ok(()),
189            None => Err(RepositoryError::NotFound(format!(
190                "Task {task_id} not found or user {user_id} does not have access"
191            ))),
192        }
193    }
194}
systemprompt_agent/repository/task/mod.rs

systemprompt_agent/repository/task/
mod.rs
