systemprompt_agent/repository/context/
queries.rs1use chrono::{DateTime, Utc};
7
8use super::ContextRepository;
9use crate::models::context::{ContextKind, ContextStateEvent, UserContext, UserContextWithStats};
10use crate::repository::task::constructor::TaskConstructor;
11use systemprompt_identifiers::{ContextId, SessionId, TaskId, UserId};
12use systemprompt_traits::RepositoryError;
13
14impl ContextRepository {
15 pub async fn find_user_id_for_context(
16 &self,
17 context_id: &ContextId,
18 ) -> Result<Option<UserId>, RepositoryError> {
19 let row = sqlx::query_scalar!(
20 r#"SELECT user_id FROM user_contexts WHERE context_id = $1"#,
21 context_id.as_str(),
22 )
23 .fetch_optional(&*self.pool)
24 .await
25 .map_err(RepositoryError::database)?;
26 Ok(row.map(UserId::new))
27 }
28
29 pub async fn get_context(
30 &self,
31 context_id: &ContextId,
32 user_id: &UserId,
33 ) -> Result<UserContext, RepositoryError> {
34 let row = sqlx::query!(
35 r#"SELECT
36 context_id as "context_id!",
37 user_id as "user_id!",
38 name as "name!",
39 kind as "kind!: ContextKind",
40 created_at as "created_at!",
41 updated_at as "updated_at!"
42 FROM user_contexts WHERE context_id = $1 AND user_id = $2"#,
43 context_id.as_str(),
44 user_id.as_str()
45 )
46 .fetch_one(&*self.pool)
47 .await
48 .map_err(|e| match e {
49 sqlx::Error::RowNotFound => RepositoryError::NotFound(format!(
50 "Context {} not found for user {}",
51 context_id, user_id
52 )),
53 _ => RepositoryError::database(e),
54 })?;
55
56 Ok(UserContext {
57 context_id: ContextId::new(row.context_id),
58 user_id: UserId::new(row.user_id),
59 name: row.name,
60 kind: row.kind,
61 created_at: row.created_at,
62 updated_at: row.updated_at,
63 })
64 }
65
66 pub async fn list_contexts_basic(
67 &self,
68 user_id: &UserId,
69 ) -> Result<Vec<UserContext>, RepositoryError> {
70 let rows = sqlx::query!(
71 r#"SELECT
72 context_id as "context_id!",
73 user_id as "user_id!",
74 name as "name!",
75 kind as "kind!: ContextKind",
76 created_at as "created_at!",
77 updated_at as "updated_at!"
78 FROM user_contexts WHERE user_id = $1 ORDER BY updated_at DESC"#,
79 user_id.as_str()
80 )
81 .fetch_all(&*self.pool)
82 .await
83 .map_err(RepositoryError::database)?;
84
85 Ok(rows
86 .into_iter()
87 .map(|r| UserContext {
88 context_id: ContextId::new(r.context_id),
89 user_id: UserId::new(r.user_id),
90 name: r.name,
91 kind: r.kind,
92 created_at: r.created_at,
93 updated_at: r.updated_at,
94 })
95 .collect())
96 }
97
98 pub async fn list_contexts_with_stats(
99 &self,
100 user_id: &UserId,
101 ) -> Result<Vec<UserContextWithStats>, RepositoryError> {
102 let rows = sqlx::query!(
103 r#"SELECT
104 c.context_id as "context_id!",
105 c.user_id as "user_id!",
106 c.name as "name!",
107 c.kind as "kind!: ContextKind",
108 c.created_at as "created_at!",
109 c.updated_at as "updated_at!",
110 COALESCE(COUNT(DISTINCT t.task_id), 0)::bigint as "task_count!",
111 COALESCE(COUNT(DISTINCT m.id), 0)::bigint as "message_count!",
112 MAX(m.created_at) as last_message_at
113 FROM user_contexts c
114 LEFT JOIN agent_tasks t ON t.context_id = c.context_id
115 LEFT JOIN task_messages m ON m.task_id = t.task_id
116 WHERE c.user_id = $1
117 GROUP BY c.context_id
118 ORDER BY c.updated_at DESC"#,
119 user_id.as_str()
120 )
121 .fetch_all(&*self.pool)
122 .await
123 .map_err(RepositoryError::database)?;
124
125 Ok(rows
126 .into_iter()
127 .map(|r| UserContextWithStats {
128 context_id: ContextId::new(r.context_id),
129 user_id: UserId::new(r.user_id),
130 name: r.name,
131 kind: r.kind,
132 created_at: r.created_at,
133 updated_at: r.updated_at,
134 task_count: r.task_count,
135 message_count: r.message_count,
136 last_message_at: r.last_message_at,
137 })
138 .collect())
139 }
140
141 pub async fn find_by_session_id(
142 &self,
143 session_id: &SessionId,
144 ) -> Result<Option<UserContext>, RepositoryError> {
145 let row = sqlx::query!(
146 r#"SELECT
147 context_id as "context_id!",
148 user_id as "user_id!",
149 name as "name!",
150 kind as "kind!: ContextKind",
151 created_at as "created_at!",
152 updated_at as "updated_at!"
153 FROM user_contexts WHERE session_id = $1
154 ORDER BY created_at DESC LIMIT 1"#,
155 session_id.as_str()
156 )
157 .fetch_optional(&*self.pool)
158 .await
159 .map_err(RepositoryError::database)?;
160
161 Ok(row.map(|r| UserContext {
162 context_id: ContextId::new(r.context_id),
163 user_id: UserId::new(r.user_id),
164 name: r.name,
165 kind: r.kind,
166 created_at: r.created_at,
167 updated_at: r.updated_at,
168 }))
169 }
170
171 pub async fn get_context_events_since(
172 &self,
173 context_id: &ContextId,
174 last_seen: DateTime<Utc>,
175 ) -> Result<Vec<ContextStateEvent>, RepositoryError> {
176 let mut events = Vec::new();
177
178 let task_ids: Vec<String> = sqlx::query_scalar!(
179 r#"SELECT t.task_id as "task_id!" FROM agent_tasks t
180 WHERE t.context_id = $1 AND t.updated_at > $2
181 ORDER BY t.updated_at ASC"#,
182 context_id.as_str(),
183 last_seen
184 )
185 .fetch_all(&*self.pool)
186 .await
187 .map_err(RepositoryError::database)?;
188
189 if !task_ids.is_empty() {
190 let constructor = TaskConstructor::new(&self.db_pool)?;
191 let task_ids_typed: Vec<TaskId> = task_ids.iter().map(TaskId::new).collect();
192 let tasks = constructor.construct_tasks_batch(&task_ids_typed).await?;
193
194 for task in tasks {
195 events.push(ContextStateEvent::TaskStatusChanged {
196 task,
197 context_id: context_id.clone(),
198 timestamp: Utc::now(),
199 });
200 }
201 }
202
203 let context_updates = sqlx::query!(
204 r#"SELECT
205 context_id as "context_id!",
206 name as "name!",
207 updated_at as "updated_at!"
208 FROM user_contexts
209 WHERE context_id = $1 AND updated_at > $2
210 ORDER BY updated_at ASC"#,
211 context_id.as_str(),
212 last_seen
213 )
214 .fetch_all(&*self.pool)
215 .await
216 .map_err(RepositoryError::database)?;
217
218 for row in context_updates {
219 events.push(ContextStateEvent::ContextUpdated {
220 context_id: ContextId::new(row.context_id),
221 name: row.name,
222 timestamp: row.updated_at,
223 });
224 }
225
226 events.sort_by_key(ContextStateEvent::timestamp);
227
228 Ok(events)
229 }
230}