Skip to main content

systemprompt_agent/
state.rs

1use std::sync::Arc;
2use systemprompt_database::DbPool;
3use systemprompt_models::Config;
4use systemprompt_traits::{
5    AnalyticsProvider, DynFileUploadProvider, DynJwtValidationProvider, DynMcpServiceProvider,
6    DynProcessCleanupProvider, DynSessionAnalyticsProvider, DynUserProvider,
7};
8
9#[derive(Clone)]
10pub struct AgentState {
11    db_pool: DbPool,
12    config: Arc<Config>,
13    jwt_provider: DynJwtValidationProvider,
14    user_provider: Option<DynUserProvider>,
15    analytics_provider: Option<Arc<dyn AnalyticsProvider>>,
16    session_analytics_provider: Option<DynSessionAnalyticsProvider>,
17    file_upload_provider: Option<DynFileUploadProvider>,
18    mcp_service_provider: Option<DynMcpServiceProvider>,
19    process_cleanup_provider: Option<DynProcessCleanupProvider>,
20}
21
22impl AgentState {
23    #[must_use]
24    pub fn new(
25        db_pool: DbPool,
26        config: Arc<Config>,
27        jwt_provider: DynJwtValidationProvider,
28    ) -> Self {
29        Self {
30            db_pool,
31            config,
32            jwt_provider,
33            user_provider: None,
34            analytics_provider: None,
35            session_analytics_provider: None,
36            file_upload_provider: None,
37            mcp_service_provider: None,
38            process_cleanup_provider: None,
39        }
40    }
41
42    #[must_use]
43    pub fn with_user_provider(mut self, provider: DynUserProvider) -> Self {
44        self.user_provider = Some(provider);
45        self
46    }
47
48    #[must_use]
49    pub fn with_analytics_provider(mut self, provider: Arc<dyn AnalyticsProvider>) -> Self {
50        self.analytics_provider = Some(provider);
51        self
52    }
53
54    #[must_use]
55    pub fn with_session_analytics_provider(
56        mut self,
57        provider: DynSessionAnalyticsProvider,
58    ) -> Self {
59        self.session_analytics_provider = Some(provider);
60        self
61    }
62
63    #[must_use]
64    pub fn with_file_upload_provider(mut self, provider: DynFileUploadProvider) -> Self {
65        self.file_upload_provider = Some(provider);
66        self
67    }
68
69    #[must_use]
70    pub fn with_mcp_service_provider(mut self, provider: DynMcpServiceProvider) -> Self {
71        self.mcp_service_provider = Some(provider);
72        self
73    }
74
75    #[must_use]
76    pub fn with_process_cleanup_provider(mut self, provider: DynProcessCleanupProvider) -> Self {
77        self.process_cleanup_provider = Some(provider);
78        self
79    }
80
81    #[must_use]
82    pub const fn db_pool(&self) -> &DbPool {
83        &self.db_pool
84    }
85
86    #[must_use]
87    pub fn config(&self) -> &Config {
88        &self.config
89    }
90
91    #[must_use]
92    pub fn jwt_provider(&self) -> &DynJwtValidationProvider {
93        &self.jwt_provider
94    }
95
96    #[must_use]
97    pub fn user_provider(&self) -> Option<&DynUserProvider> {
98        self.user_provider.as_ref()
99    }
100
101    #[must_use]
102    pub fn analytics_provider(&self) -> Option<&Arc<dyn AnalyticsProvider>> {
103        self.analytics_provider.as_ref()
104    }
105
106    #[must_use]
107    pub fn session_analytics_provider(&self) -> Option<&DynSessionAnalyticsProvider> {
108        self.session_analytics_provider.as_ref()
109    }
110
111    #[must_use]
112    pub fn file_upload_provider(&self) -> Option<&DynFileUploadProvider> {
113        self.file_upload_provider.as_ref()
114    }
115
116    #[must_use]
117    pub fn mcp_service_provider(&self) -> Option<&DynMcpServiceProvider> {
118        self.mcp_service_provider.as_ref()
119    }
120
121    #[must_use]
122    pub fn process_cleanup_provider(&self) -> Option<&DynProcessCleanupProvider> {
123        self.process_cleanup_provider.as_ref()
124    }
125}
126
127impl std::fmt::Debug for AgentState {
128    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
129        f.debug_struct("AgentState")
130            .field("db_pool", &"<DbPool>")
131            .field("config", &"<Arc<Config>>")
132            .field("jwt_provider", &"<DynJwtValidationProvider>")
133            .field("user_provider", &self.user_provider.is_some())
134            .field("analytics_provider", &self.analytics_provider.is_some())
135            .field(
136                "session_analytics_provider",
137                &self.session_analytics_provider.is_some(),
138            )
139            .field("file_upload_provider", &self.file_upload_provider.is_some())
140            .field("mcp_service_provider", &self.mcp_service_provider.is_some())
141            .field(
142                "process_cleanup_provider",
143                &self.process_cleanup_provider.is_some(),
144            )
145            .finish()
146    }
147}