Skip to main content

systemprompt_agent/services/
context_provider.rs

1//! Adapter exposing [`ContextRepository`] through the
2//! [`systemprompt_traits::ContextProvider`] trait so other crates can consume
3//! it without a direct database dependency.
4//!
5//! Copyright (c) systemprompt.io — Business Source License 1.1.
6//! See <https://systemprompt.io> for licensing details.
7
8use async_trait::async_trait;
9use systemprompt_identifiers::{ContextId, SessionId, UserId};
10use systemprompt_traits::{
11    ContextMaterializer, ContextProvider, ContextProviderError, ContextWithStats,
12    EnsureContextParams,
13};
14
15use crate::models::context::ContextKind;
16use crate::repository::ContextRepository;
17
18#[derive(Debug, Clone)]
19pub struct ContextProviderService {
20    repo: ContextRepository,
21}
22
23impl ContextProviderService {
24    #[must_use]
25    pub const fn new(repo: ContextRepository) -> Self {
26        Self { repo }
27    }
28}
29
30#[async_trait]
31impl ContextProvider for ContextProviderService {
32    async fn list_contexts_with_stats(
33        &self,
34        user_id: &UserId,
35    ) -> Result<Vec<ContextWithStats>, ContextProviderError> {
36        let contexts = self
37            .repo
38            .list_contexts_with_stats(user_id)
39            .await
40            .map_err(|e| ContextProviderError::Database(e.to_string()))?;
41
42        Ok(contexts
43            .into_iter()
44            .map(|c| ContextWithStats {
45                context_id: c.context_id,
46                user_id: c.user_id,
47                name: c.name,
48                created_at: c.created_at,
49                updated_at: c.updated_at,
50                task_count: c.task_count,
51                message_count: c.message_count,
52                last_message_at: c.last_message_at,
53            })
54            .collect())
55    }
56
57    async fn get_context(
58        &self,
59        context_id: &ContextId,
60        user_id: &UserId,
61    ) -> Result<ContextWithStats, ContextProviderError> {
62        let context = self
63            .repo
64            .get_context(context_id, user_id)
65            .await
66            .map_err(|e| match e {
67                systemprompt_traits::RepositoryError::NotFound(msg) => {
68                    ContextProviderError::NotFound(msg)
69                },
70                other => ContextProviderError::Database(other.to_string()),
71            })?;
72
73        let all_contexts = self
74            .repo
75            .list_contexts_with_stats(user_id)
76            .await
77            .map_err(|e| ContextProviderError::Database(e.to_string()))?;
78
79        let context_with_stats = all_contexts
80            .into_iter()
81            .find(|c| c.context_id == context.context_id)
82            .ok_or_else(|| {
83                ContextProviderError::NotFound(format!("Context {} not found", context_id))
84            })?;
85
86        Ok(ContextWithStats {
87            context_id: context_with_stats.context_id,
88            user_id: context_with_stats.user_id,
89            name: context_with_stats.name,
90            created_at: context_with_stats.created_at,
91            updated_at: context_with_stats.updated_at,
92            task_count: context_with_stats.task_count,
93            message_count: context_with_stats.message_count,
94            last_message_at: context_with_stats.last_message_at,
95        })
96    }
97
98    async fn create_context(
99        &self,
100        user_id: &UserId,
101        session_id: Option<&SessionId>,
102        name: &str,
103    ) -> Result<ContextId, ContextProviderError> {
104        self.repo
105            .create_context(user_id, session_id, name, ContextKind::User)
106            .await
107            .map_err(|e| ContextProviderError::Database(e.to_string()))
108    }
109
110    async fn update_context_name(
111        &self,
112        context_id: &ContextId,
113        user_id: &UserId,
114        name: &str,
115    ) -> Result<(), ContextProviderError> {
116        self.repo
117            .update_context_name(context_id, user_id, name)
118            .await
119            .map_err(|e| match e {
120                systemprompt_traits::RepositoryError::NotFound(msg) => {
121                    ContextProviderError::NotFound(msg)
122                },
123                other => ContextProviderError::Database(other.to_string()),
124            })
125    }
126
127    async fn delete_context(
128        &self,
129        context_id: &ContextId,
130        user_id: &UserId,
131    ) -> Result<(), ContextProviderError> {
132        self.repo
133            .delete_context(context_id, user_id)
134            .await
135            .map_err(|e| match e {
136                systemprompt_traits::RepositoryError::NotFound(msg) => {
137                    ContextProviderError::NotFound(msg)
138                },
139                other => ContextProviderError::Database(other.to_string()),
140            })
141    }
142}
143
144#[async_trait]
145impl ContextMaterializer for ContextProviderService {
146    async fn ensure_context(
147        &self,
148        params: EnsureContextParams<'_>,
149    ) -> Result<(), ContextProviderError> {
150        let kind = params
151            .kind
152            .parse::<ContextKind>()
153            .map_err(|e| ContextProviderError::Internal(e.to_string()))?;
154
155        self.repo
156            .ensure_context(&params, kind)
157            .await
158            .map_err(|e| ContextProviderError::Database(e.to_string()))
159    }
160}