systemprompt_api/services/middleware/session/
mod.rs

1//! Session-establishment middleware.
2//!
3//! [`SessionMiddleware`] resolves or mints the per-request session: it skips
4//! untracked paths, short-circuits detected bots into anonymous contexts,
5//! validates an existing JWT session, and refreshes or recreates the session
6//! when the token is stale, issuing a `Set-Cookie` for newly minted tokens.
7//!
8//! Copyright (c) systemprompt.io — Business Source License 1.1.
9//! See <https://systemprompt.io> for licensing details.
10
11mod lifecycle;
12mod skip;
13
14pub use skip::should_skip_session_tracking;
15
16use axum::extract::Request;
17use axum::http::header;
18use axum::middleware::Next;
19use axum::response::Response;
20use std::sync::Arc;
21use systemprompt_analytics::AnalyticsService;
22use systemprompt_identifiers::{AgentName, ContextId, SessionId, UserId};
23use systemprompt_models::api::ApiError;
24use systemprompt_models::auth::UserType;
25use systemprompt_models::execution::context::RequestContext;
26use systemprompt_oauth::services::SessionCreationService;
27use systemprompt_runtime::AppContext;
28use systemprompt_security::{
29    CookieExtractor, HeaderExtractor, TokenExtractor, extract_user_context,
30};
31use systemprompt_traits::AnalyticsProvider;
32use systemprompt_users::UserService;
33use uuid::Uuid;
34
35#[derive(Clone, Debug)]
36pub struct SessionMiddleware {
37    analytics_service: Arc<AnalyticsService>,
38    session_creation_service: Arc<SessionCreationService>,
39}
40
41impl SessionMiddleware {
42    pub fn new(ctx: &AppContext) -> anyhow::Result<Self> {
43        let user_service = UserService::new(ctx.db_pool())?;
44        let concrete = Arc::clone(ctx.analytics_service());
45        let analytics: Arc<dyn AnalyticsProvider> = concrete;
46        let session_creation_service = Arc::new(SessionCreationService::new(
47            analytics,
48            Arc::new(user_service),
49        ));
50
51        Ok(Self {
52            analytics_service: Arc::clone(ctx.analytics_service()),
53            session_creation_service,
54        })
55    }
56
57    pub async fn handle(&self, mut request: Request, next: Next) -> Result<Response, ApiError> {
58        let headers = request.headers();
59        let uri = request.uri().clone();
60        let method = request.method().clone();
61
62        let should_skip = should_skip_session_tracking(uri.path());
63
64        tracing::debug!(
65            path = %uri.path(),
66            should_skip = should_skip,
67            "Session middleware evaluating request"
68        );
69
70        let trace_id = HeaderExtractor::extract_trace_id(headers);
71
72        let (req_ctx, jwt_cookie) = if should_skip {
73            (
74                self.anonymous_context("untracked", trace_id, headers, &uri)
75                    .await?,
76                None,
77            )
78        } else {
79            self.tracked_context(trace_id, headers, &uri, &method)
80                .await?
81        };
82
83        tracing::debug!(
84            path = %uri.path(),
85            session_id = %req_ctx.session_id(),
86            "Session middleware setting context"
87        );
88
89        request.extensions_mut().insert(req_ctx);
90
91        let mut response = next.run(request).await;
92
93        if let Some(token) = jwt_cookie {
94            let cookie = format!(
95                "{}={token}; HttpOnly; SameSite=Strict; Path=/; Max-Age=604800",
96                CookieExtractor::DEFAULT_COOKIE_NAME
97            );
98            if let Ok(cookie_value) = cookie.parse() {
99                response
100                    .headers_mut()
101                    .insert(header::SET_COOKIE, cookie_value);
102            }
103        }
104
105        Ok(response)
106    }
107
108    /// Builds an untracked anonymous context. `session_prefix` distinguishes
109    /// the synthetic session id (`untracked_*` for skip-tracking paths,
110    /// `bot_*` for detected crawlers) so the two cases stay legible in logs
111    /// and analytics without two near-identical constructors.
112    async fn anonymous_context(
113        &self,
114        session_prefix: &str,
115        trace_id: systemprompt_identifiers::TraceId,
116        headers: &http::HeaderMap,
117        uri: &http::Uri,
118    ) -> Result<RequestContext, ApiError> {
119        let (user_id, fingerprint) = self
120            .session_creation_service
121            .ensure_anonymous_user(headers, Some(uri))
122            .await
123            .map_err(|e| {
124                tracing::error!(error = %e, session_prefix, "Failed to ensure anonymous user");
125                ApiError::internal_error("Service temporarily unavailable")
126            })?;
127
128        Ok(RequestContext::new(
129            SessionId::new(format!("{session_prefix}_{}", Uuid::new_v4())),
130            trace_id,
131            ContextId::generate(),
132            AgentName::system(),
133        )
134        .with_actor(systemprompt_identifiers::Actor::anonymous(user_id))
135        .with_user_type(UserType::Anon)
136        .with_tracked(false)
137        .with_fingerprint_hash(fingerprint))
138    }
139
140    async fn tracked_context(
141        &self,
142        trace_id: systemprompt_identifiers::TraceId,
143        headers: &http::HeaderMap,
144        uri: &http::Uri,
145        method: &http::Method,
146    ) -> Result<(RequestContext, Option<String>), ApiError> {
147        let analytics = self.analytics_service.extract_analytics(headers, Some(uri));
148        let is_bot = AnalyticsService::is_bot(&analytics);
149
150        tracing::debug!(
151            path = %uri.path(),
152            is_bot = is_bot,
153            user_agent = ?analytics.user_agent,
154            "Session middleware bot check"
155        );
156
157        if is_bot {
158            return Ok((
159                self.anonymous_context("bot", trace_id, headers, uri)
160                    .await?,
161                None,
162            ));
163        }
164
165        let token_result = TokenExtractor::browser_only().extract(headers).ok();
166
167        let (session_id, user_id, jwt_token, jwt_cookie, fingerprint_hash) = self
168            .resolve_session(token_result, headers, uri, method)
169            .await?;
170
171        let context_id =
172            HeaderExtractor::extract_context_id(headers).unwrap_or_else(ContextId::generate);
173
174        let mut ctx = RequestContext::new(session_id, trace_id, context_id, AgentName::system())
175            .with_actor(systemprompt_identifiers::Actor::user(user_id))
176            .with_auth_token(jwt_token)
177            .with_user_type(UserType::Anon)
178            .with_tracked(true);
179        if let Some(fp) = fingerprint_hash {
180            ctx = ctx.with_fingerprint_hash(fp);
181        }
182        Ok((ctx, jwt_cookie))
183    }
184
185    async fn resolve_session(
186        &self,
187        token_result: Option<String>,
188        headers: &http::HeaderMap,
189        uri: &http::Uri,
190        method: &http::Method,
191    ) -> Result<(SessionId, UserId, String, Option<String>, Option<String>), ApiError> {
192        let Some(token) = token_result else {
193            let (sid, uid, token, is_new, fp) =
194                lifecycle::create_new_session(&self.session_creation_service, headers, uri, method)
195                    .await?;
196            let jwt_cookie = if is_new { Some(token.clone()) } else { None };
197            return Ok((sid, uid, token, jwt_cookie, Some(fp)));
198        };
199
200        let Ok(jwt_context) = extract_user_context(&token) else {
201            let (sid, uid, token, is_new, fp) =
202                lifecycle::create_new_session(&self.session_creation_service, headers, uri, method)
203                    .await?;
204            let jwt_cookie = if is_new { Some(token.clone()) } else { None };
205            return Ok((sid, uid, token, jwt_cookie, Some(fp)));
206        };
207
208        let session_exists = self
209            .analytics_service
210            .find_active_session_by_id(&jwt_context.session_id)
211            .await
212            .map_err(|e| {
213                tracing::warn!(error = %e, "find_active_session_by_id failed");
214                e
215            })
216            .ok()
217            .flatten()
218            .is_some();
219
220        if session_exists {
221            return Ok((
222                jwt_context.session_id,
223                jwt_context.user_id,
224                token,
225                None,
226                None,
227            ));
228        }
229
230        tracing::info!(
231            old_session_id = %jwt_context.session_id,
232            user_id = %jwt_context.user_id,
233            "JWT valid but session missing, refreshing with new session"
234        );
235        match lifecycle::refresh_session_for_user(
236            &self.session_creation_service,
237            &self.analytics_service,
238            &jwt_context.user_id,
239            headers,
240            uri,
241        )
242        .await
243        {
244            Ok((sid, uid, new_token, _, fp)) => {
245                Ok((sid, uid, new_token.clone(), Some(new_token), Some(fp)))
246            },
247            Err(e) if e.error_key.as_deref() == Some("user_not_found") => {
248                tracing::warn!(
249                    user_id = %jwt_context.user_id,
250                    "JWT references non-existent user, creating new anonymous session"
251                );
252                let (sid, uid, token, _, fp) = lifecycle::create_new_session(
253                    &self.session_creation_service,
254                    headers,
255                    uri,
256                    method,
257                )
258                .await?;
259                Ok((sid, uid, token.clone(), Some(token), Some(fp)))
260            },
261            Err(e) => Err(e),
262        }
263    }
264}
systemprompt_api/services/middleware/session/mod.rs

systemprompt_api/services/middleware/session/
mod.rs
