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
5use async_trait::async_trait;
6use systemprompt_database::DbPool;
7use systemprompt_identifiers::{ContextId, SessionId, UserId};
8use systemprompt_traits::{ContextProvider, ContextProviderError, ContextWithStats};
9
10use crate::error::AgentError;
11use crate::repository::ContextRepository;
12
13#[derive(Debug, Clone)]
14pub struct ContextProviderService {
15    repo: ContextRepository,
16}
17
18impl ContextProviderService {
19    pub fn new(db_pool: &DbPool) -> Result<Self, AgentError> {
20        Ok(Self {
21            repo: ContextRepository::new(db_pool)?,
22        })
23    }
24}
25
26#[async_trait]
27impl ContextProvider for ContextProviderService {
28    async fn list_contexts_with_stats(
29        &self,
30        user_id: &UserId,
31    ) -> Result<Vec<ContextWithStats>, ContextProviderError> {
32        let contexts = self
33            .repo
34            .list_contexts_with_stats(user_id)
35            .await
36            .map_err(|e| ContextProviderError::Database(e.to_string()))?;
37
38        Ok(contexts
39            .into_iter()
40            .map(|c| ContextWithStats {
41                context_id: c.context_id,
42                user_id: c.user_id,
43                name: c.name,
44                created_at: c.created_at,
45                updated_at: c.updated_at,
46                task_count: c.task_count,
47                message_count: c.message_count,
48                last_message_at: c.last_message_at,
49            })
50            .collect())
51    }
52
53    async fn get_context(
54        &self,
55        context_id: &ContextId,
56        user_id: &UserId,
57    ) -> Result<ContextWithStats, ContextProviderError> {
58        let context = self
59            .repo
60            .get_context(context_id, user_id)
61            .await
62            .map_err(|e| match e {
63                systemprompt_traits::RepositoryError::NotFound(msg) => {
64                    ContextProviderError::NotFound(msg)
65                },
66                other => ContextProviderError::Database(other.to_string()),
67            })?;
68
69        let all_contexts = self
70            .repo
71            .list_contexts_with_stats(user_id)
72            .await
73            .map_err(|e| ContextProviderError::Database(e.to_string()))?;
74
75        let context_with_stats = all_contexts
76            .into_iter()
77            .find(|c| c.context_id == context.context_id)
78            .ok_or_else(|| {
79                ContextProviderError::NotFound(format!("Context {} not found", context_id))
80            })?;
81
82        Ok(ContextWithStats {
83            context_id: context_with_stats.context_id,
84            user_id: context_with_stats.user_id,
85            name: context_with_stats.name,
86            created_at: context_with_stats.created_at,
87            updated_at: context_with_stats.updated_at,
88            task_count: context_with_stats.task_count,
89            message_count: context_with_stats.message_count,
90            last_message_at: context_with_stats.last_message_at,
91        })
92    }
93
94    async fn create_context(
95        &self,
96        user_id: &UserId,
97        session_id: Option<&SessionId>,
98        name: &str,
99    ) -> Result<ContextId, ContextProviderError> {
100        self.repo
101            .create_context(user_id, session_id, name)
102            .await
103            .map_err(|e| ContextProviderError::Database(e.to_string()))
104    }
105
106    async fn update_context_name(
107        &self,
108        context_id: &ContextId,
109        user_id: &UserId,
110        name: &str,
111    ) -> Result<(), ContextProviderError> {
112        self.repo
113            .update_context_name(context_id, user_id, name)
114            .await
115            .map_err(|e| match e {
116                systemprompt_traits::RepositoryError::NotFound(msg) => {
117                    ContextProviderError::NotFound(msg)
118                },
119                other => ContextProviderError::Database(other.to_string()),
120            })
121    }
122
123    async fn delete_context(
124        &self,
125        context_id: &ContextId,
126        user_id: &UserId,
127    ) -> Result<(), ContextProviderError> {
128        self.repo
129            .delete_context(context_id, user_id)
130            .await
131            .map_err(|e| match e {
132                systemprompt_traits::RepositoryError::NotFound(msg) => {
133                    ContextProviderError::NotFound(msg)
134                },
135                other => ContextProviderError::Database(other.to_string()),
136            })
137    }
138}