Skip to main content

systemprompt_users/sessions/
mod.rs

1//! Authoritative session persistence contracts.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6mod behavioral;
7mod behavioral_queries;
8mod geo;
9mod mutations;
10mod providers;
11mod queries;
12use crate::Result;
13use chrono::{DateTime, Utc};
14use sqlx::PgPool;
15use std::sync::Arc;
16use systemprompt_database::DbPool;
17use systemprompt_identifiers::{SessionId, UserId};
18use systemprompt_traits::session_store::{
19    ActiveSessionLookup, CreateSessionParams, SessionBehavioralData, SessionRecord,
20    SessionSnapshot as AnalyticsSession,
21};
22#[derive(Clone, Debug)]
23pub struct SessionRepository {
24    pool: Arc<PgPool>,
25    write_pool: Arc<PgPool>,
26}
27impl SessionRepository {
28    pub fn new(db: &DbPool) -> Result<Self> {
29        Ok(Self {
30            pool: db.pool_arc()?,
31            write_pool: db.write_pool_arc()?,
32        })
33    }
34    pub async fn find_by_id(&self, session_id: &SessionId) -> Result<Option<AnalyticsSession>> {
35        queries::find_by_id(&self.write_pool, session_id).await
36    }
37    pub async fn find_active_by_id(
38        &self,
39        session_id: &SessionId,
40    ) -> Result<Option<ActiveSessionLookup>> {
41        queries::find_active_by_id(&self.write_pool, session_id).await
42    }
43    pub async fn revoke_session(&self, session_id: &SessionId) -> Result<()> {
44        mutations::revoke_session(&self.write_pool, session_id).await
45    }
46    pub async fn revoke_all_for_user(&self, user_id: &UserId) -> Result<u64> {
47        mutations::revoke_all_for_user(&self.write_pool, user_id).await
48    }
49    pub async fn find_by_fingerprint(
50        &self,
51        fingerprint_hash: &str,
52        user_id: &UserId,
53    ) -> Result<Option<AnalyticsSession>> {
54        queries::find_by_fingerprint(&self.pool, fingerprint_hash, user_id).await
55    }
56    pub async fn list_active_by_user(&self, user_id: &UserId) -> Result<Vec<AnalyticsSession>> {
57        queries::list_active_by_user(&self.pool, user_id).await
58    }
59    pub async fn increment_request_count(&self, session_id: &SessionId) -> Result<()> {
60        mutations::increment_request_count(&self.write_pool, session_id).await
61    }
62    pub async fn increment_task_count(&self, session_id: &SessionId) -> Result<()> {
63        mutations::increment_task_count(&self.write_pool, session_id).await
64    }
65    pub async fn increment_message_count(&self, session_id: &SessionId) -> Result<()> {
66        mutations::increment_message_count(&self.write_pool, session_id).await
67    }
68    pub async fn end_session(&self, session_id: &SessionId) -> Result<()> {
69        mutations::end_session(&self.write_pool, session_id).await
70    }
71    pub async fn mark_as_scanner(&self, session_id: &SessionId) -> Result<()> {
72        mutations::mark_as_scanner(&self.write_pool, session_id).await
73    }
74    pub async fn mark_converted(&self, session_id: &SessionId) -> Result<()> {
75        mutations::mark_converted(&self.write_pool, session_id).await
76    }
77    pub async fn mark_as_behavioral_bot(&self, session_id: &SessionId, reason: &str) -> Result<()> {
78        behavioral::mark_as_behavioral_bot(&self.write_pool, session_id, reason).await
79    }
80    pub async fn check_and_mark_behavioral_bot(
81        &self,
82        session_id: &SessionId,
83        request_count_threshold: i32,
84    ) -> Result<bool> {
85        behavioral::check_and_mark_behavioral_bot(
86            &self.write_pool,
87            session_id,
88            request_count_threshold,
89        )
90        .await
91    }
92    pub async fn cleanup_inactive(&self, inactive_hours: i32) -> Result<u64> {
93        mutations::cleanup_inactive(&self.write_pool, inactive_hours).await
94    }
95    pub async fn count_inactive(&self, inactive_hours: i32) -> Result<i64> {
96        queries::count_inactive(&self.pool, inactive_hours).await
97    }
98    pub async fn count_sessions_missing_geo(&self) -> Result<i64> {
99        mutations::count_sessions_missing_geo(&self.pool).await
100    }
101    pub async fn migrate_user_sessions(
102        &self,
103        old_user_id: &UserId,
104        new_user_id: &UserId,
105    ) -> Result<u64> {
106        mutations::migrate_user_sessions(&self.write_pool, old_user_id, new_user_id).await
107    }
108    pub async fn create_session(&self, params: &CreateSessionParams<'_>) -> Result<()> {
109        mutations::create_session(&self.write_pool, params).await
110    }
111    pub async fn find_recent_by_fingerprint(
112        &self,
113        fingerprint_hash: &str,
114        max_age_seconds: i64,
115    ) -> Result<Option<SessionRecord>> {
116        queries::find_recent_by_fingerprint(&self.write_pool, fingerprint_hash, max_age_seconds)
117            .await
118    }
119    pub async fn increment_ai_usage(
120        &self,
121        session_id: &SessionId,
122        tokens: i32,
123        cost_microdollars: i64,
124    ) -> Result<()> {
125        mutations::increment_ai_usage(&self.write_pool, session_id, tokens, cost_microdollars).await
126    }
127    pub async fn update_behavioral_detection(
128        &self,
129        session_id: &SessionId,
130        score: i32,
131        is_behavioral_bot: bool,
132        reason: Option<&str>,
133    ) -> Result<()> {
134        behavioral::update_behavioral_detection(
135            &self.write_pool,
136            session_id,
137            score,
138            is_behavioral_bot,
139            reason,
140        )
141        .await
142    }
143    pub async fn count_sessions_by_fingerprint(
144        &self,
145        fingerprint_hash: &str,
146        window_hours: i64,
147    ) -> Result<i64> {
148        behavioral_queries::count_sessions_by_fingerprint(
149            &self.write_pool,
150            fingerprint_hash,
151            window_hours,
152        )
153        .await
154    }
155    pub async fn get_session_for_behavioral_analysis(
156        &self,
157        session_id: &SessionId,
158    ) -> Result<Option<SessionBehavioralData>> {
159        behavioral_queries::get_session_for_behavioral_analysis(&self.write_pool, session_id).await
160    }
161    pub async fn count_unique_ips_by_fingerprint(
162        &self,
163        fingerprint_hash: &str,
164        window_days: i64,
165    ) -> Result<i64> {
166        behavioral_queries::count_unique_ips_by_fingerprint(
167            &self.write_pool,
168            fingerprint_hash,
169            window_days,
170        )
171        .await
172    }
173    pub async fn get_session_starts_by_fingerprint(
174        &self,
175        fingerprint_hash: &str,
176        window_days: i64,
177    ) -> Result<Vec<DateTime<Utc>>> {
178        behavioral_queries::get_session_starts_by_fingerprint(
179            &self.write_pool,
180            fingerprint_hash,
181            window_days,
182        )
183        .await
184    }
185    pub async fn get_session_velocity(
186        &self,
187        session_id: &SessionId,
188    ) -> Result<(Option<i64>, Option<i64>)> {
189        behavioral_queries::get_session_velocity(&self.write_pool, session_id).await
190    }
191}
192
193mod ai_provider;
194pub use ai_provider::UsersAiSessionProvider;
195
196mod fingerprint;
197
198mod store;