systemprompt_agent/repository/context/
mutations.rs

1//! Context insert/update/delete mutations.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use chrono::Utc;
7
8use super::ContextRepository;
9use crate::models::context::ContextKind;
10use systemprompt_identifiers::{ContextId, SessionId, UserId};
11use systemprompt_traits::RepositoryError;
12
13impl ContextRepository {
14    pub async fn create_context(
15        &self,
16        user_id: &UserId,
17        session_id: Option<&SessionId>,
18        name: &str,
19        kind: ContextKind,
20    ) -> Result<ContextId, RepositoryError> {
21        let context_id = ContextId::generate();
22        let now = Utc::now();
23        let session_id_str = session_id.map(SessionId::as_str);
24
25        sqlx::query!(
26            "INSERT INTO user_contexts (context_id, user_id, session_id, name, kind, created_at, \
27             updated_at)
28             VALUES ($1, $2, $3, $4, $5, $6, $6)",
29            context_id.as_str(),
30            user_id.as_str(),
31            session_id_str,
32            name,
33            kind.as_str(),
34            now
35        )
36        .execute(&*self.write_pool)
37        .await
38        .map_err(RepositoryError::database)?;
39
40        Ok(context_id)
41    }
42
43    pub async fn ensure_context(
44        &self,
45        params: &systemprompt_traits::EnsureContextParams<'_>,
46        kind: ContextKind,
47    ) -> Result<(), RepositoryError> {
48        let context_id = params.context_id;
49        let user_id = params.user_id;
50        let session_id = params.session_id;
51        let name = params.name;
52        let now = Utc::now();
53
54        sqlx::query!(
55            "INSERT INTO user_contexts (context_id, user_id, session_id, name, kind, created_at, \
56             updated_at)
57             VALUES ($1, $2, $3, $4, $5, $6, $6)
58             ON CONFLICT (context_id) DO NOTHING",
59            context_id.as_str(),
60            user_id.as_str(),
61            session_id.map(SessionId::as_str),
62            name,
63            kind.as_str(),
64            now
65        )
66        .execute(&*self.write_pool)
67        .await
68        .map_err(RepositoryError::database)?;
69
70        Ok(())
71    }
72
73    pub async fn get_or_create_cli_context(
74        &self,
75        user_id: &UserId,
76        session_id: &SessionId,
77        name: &str,
78    ) -> Result<ContextId, RepositoryError> {
79        let now = Utc::now();
80
81        let adopted = sqlx::query_scalar!(
82            r#"UPDATE user_contexts SET session_id = $1, updated_at = $2
83             WHERE context_id = (
84                 SELECT context_id FROM user_contexts
85                 WHERE user_id = $3 AND kind = $4 AND name = $5
86                 ORDER BY updated_at DESC LIMIT 1
87             )
88             RETURNING context_id"#,
89            session_id.as_str(),
90            now,
91            user_id.as_str(),
92            ContextKind::CliSession.as_str(),
93            name
94        )
95        .fetch_optional(&*self.write_pool)
96        .await
97        .map_err(RepositoryError::database)?;
98
99        match adopted {
100            Some(context_id) => Ok(ContextId::new(context_id)),
101            None => {
102                self.create_context(user_id, Some(session_id), name, ContextKind::CliSession)
103                    .await
104            },
105        }
106    }
107
108    pub async fn validate_context_ownership(
109        &self,
110        context_id: &ContextId,
111        user_id: &UserId,
112    ) -> Result<(), RepositoryError> {
113        let result = sqlx::query_scalar!(
114            "SELECT context_id FROM user_contexts WHERE context_id = $1 AND user_id = $2",
115            context_id.as_str(),
116            user_id.as_str()
117        )
118        .fetch_optional(&*self.pool)
119        .await
120        .map_err(RepositoryError::database)?;
121
122        match result {
123            Some(_) => Ok(()),
124            None => Err(RepositoryError::NotFound(format!(
125                "Context {} not found or user {} does not have access",
126                context_id, user_id
127            ))),
128        }
129    }
130
131    pub async fn update_context_name(
132        &self,
133        context_id: &ContextId,
134        user_id: &UserId,
135        name: &str,
136    ) -> Result<(), RepositoryError> {
137        let now = Utc::now();
138
139        let result = sqlx::query!(
140            "UPDATE user_contexts SET name = $1, updated_at = $2
141             WHERE context_id = $3 AND user_id = $4",
142            name,
143            now,
144            context_id.as_str(),
145            user_id.as_str()
146        )
147        .execute(&*self.write_pool)
148        .await
149        .map_err(RepositoryError::database)?;
150
151        if result.rows_affected() == 0 {
152            return Err(RepositoryError::NotFound(format!(
153                "Context {} not found for user {}",
154                context_id, user_id
155            )));
156        }
157
158        Ok(())
159    }
160
161    pub async fn delete_context(
162        &self,
163        context_id: &ContextId,
164        user_id: &UserId,
165    ) -> Result<(), RepositoryError> {
166        let result = sqlx::query!(
167            "DELETE FROM user_contexts WHERE context_id = $1 AND user_id = $2",
168            context_id.as_str(),
169            user_id.as_str()
170        )
171        .execute(&*self.write_pool)
172        .await
173        .map_err(RepositoryError::database)?;
174
175        if result.rows_affected() == 0 {
176            return Err(RepositoryError::NotFound(format!(
177                "Context {} not found for user {}",
178                context_id, user_id
179            )));
180        }
181
182        Ok(())
183    }
184}
systemprompt_agent/repository/context/mutations.rs

systemprompt_agent/repository/context/
mutations.rs
