systemprompt_agent/repository/task/
queries.rs

1//! Task read queries.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use crate::models::TaskRow;
7use sqlx::PgPool;
8use std::sync::Arc;
9use systemprompt_identifiers::{AgentName, ContextId, SessionId, TaskId, TraceId, UserId};
10use systemprompt_traits::RepositoryError;
11
12use super::constructor::TaskConstructor;
13use crate::models::a2a::Task;
14
15pub async fn get_task(
16    constructor: &TaskConstructor,
17    task_id: &TaskId,
18) -> Result<Option<Task>, RepositoryError> {
19    constructor.construct_task_from_task_id(task_id).await
20}
21
22pub async fn list_tasks_by_context(
23    pool: &Arc<PgPool>,
24    constructor: &TaskConstructor,
25    context_id: &ContextId,
26) -> Result<Vec<Task>, RepositoryError> {
27    let context_id_str = context_id.as_str();
28    let rows = sqlx::query_as!(
29        TaskRow,
30        r#"SELECT
31            task_id as "task_id!: TaskId",
32            context_id as "context_id!: ContextId",
33            status as "status!",
34            status_timestamp,
35            user_id as "user_id?: UserId",
36            session_id as "session_id?: SessionId",
37            trace_id as "trace_id?: TraceId",
38            agent_name as "agent_name?: AgentName",
39            started_at,
40            completed_at,
41            execution_time_ms,
42            error_message,
43            metadata,
44            created_at as "created_at!",
45            updated_at as "updated_at!"
46        FROM agent_tasks WHERE context_id = $1 ORDER BY created_at ASC"#,
47        context_id_str
48    )
49    .fetch_all(pool.as_ref())
50    .await
51    .map_err(RepositoryError::database)?;
52
53    let task_ids: Vec<TaskId> = rows.iter().map(|r| r.task_id.clone()).collect();
54    let tasks = constructor.construct_tasks_batch(&task_ids).await?;
55
56    Ok(tasks)
57}
58
59pub async fn get_tasks_by_user_id(
60    pool: &Arc<PgPool>,
61    constructor: &TaskConstructor,
62    user_id: &UserId,
63    limit: Option<i32>,
64    offset: Option<i32>,
65) -> Result<Vec<Task>, RepositoryError> {
66    let lim = limit.map_or(1000, i64::from);
67    let off = offset.map_or(0, i64::from);
68    let user_id_str = user_id.as_str();
69
70    let rows = sqlx::query_as!(
71        TaskRow,
72        r#"SELECT
73            task_id as "task_id!: TaskId",
74            context_id as "context_id!: ContextId",
75            status as "status!",
76            status_timestamp,
77            user_id as "user_id?: UserId",
78            session_id as "session_id?: SessionId",
79            trace_id as "trace_id?: TraceId",
80            agent_name as "agent_name?: AgentName",
81            started_at,
82            completed_at,
83            execution_time_ms,
84            error_message,
85            metadata,
86            created_at as "created_at!",
87            updated_at as "updated_at!"
88        FROM agent_tasks WHERE user_id = $1 ORDER BY created_at DESC LIMIT $2 OFFSET $3"#,
89        user_id_str,
90        lim,
91        off
92    )
93    .fetch_all(pool.as_ref())
94    .await
95    .map_err(RepositoryError::database)?;
96
97    let task_ids: Vec<TaskId> = rows.iter().map(|r| r.task_id.clone()).collect();
98    let tasks = constructor.construct_tasks_batch(&task_ids).await?;
99
100    Ok(tasks)
101}
102
103#[derive(Debug, Clone)]
104pub struct TaskContextInfo {
105    pub context_id: ContextId,
106    pub user_id: Option<UserId>,
107}
108
109pub async fn get_task_context_info(
110    pool: &Arc<PgPool>,
111    task_id: &TaskId,
112) -> Result<Option<TaskContextInfo>, RepositoryError> {
113    let task_id_str = task_id.as_str();
114    let row = sqlx::query!(
115        r#"SELECT
116            context_id as "context_id!: ContextId",
117            user_id as "user_id?: UserId"
118        FROM agent_tasks WHERE task_id = $1"#,
119        task_id_str
120    )
121    .fetch_optional(pool.as_ref())
122    .await
123    .map_err(RepositoryError::database)?;
124
125    Ok(row.map(|r| TaskContextInfo {
126        context_id: r.context_id,
127        user_id: r.user_id,
128    }))
129}
systemprompt_agent/repository/task/queries.rs

systemprompt_agent/repository/task/
queries.rs
