systemprompt_users/sessions/
mod.rs1mod 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;