Skip to main content

systemprompt_users/repository/session/
mod.rs

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