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