systemprompt_api/services/middleware/jwt/
context.rs

1use async_trait::async_trait;
2use axum::body::Body;
3use axum::extract::Request;
4use axum::http::HeaderMap;
5use std::sync::Arc;
6
7use crate::services::middleware::context::ContextExtractor;
8use systemprompt_database::DbPool;
9use systemprompt_identifiers::{
10    AgentName, ContextId, SessionId, SessionSource, TaskId, TraceId, UserId,
11};
12use systemprompt_models::execution::context::{ContextExtractionError, RequestContext};
13use systemprompt_security::{HeaderExtractor, TokenExtractor};
14use systemprompt_traits::{AnalyticsProvider, CreateSessionInput};
15use systemprompt_users::UserService;
16
17use super::token::{JwtExtractor, JwtUserContext};
18
19struct BuildContextParams {
20    jwt_context: JwtUserContext,
21    session_id: SessionId,
22    user_id: UserId,
23    trace_id: TraceId,
24    context_id: ContextId,
25    agent_name: AgentName,
26    task_id: Option<TaskId>,
27    auth_token: Option<String>,
28}
29
30#[derive(Clone)]
31pub struct JwtContextExtractor {
32    jwt_extractor: Arc<JwtExtractor>,
33    token_extractor: TokenExtractor,
34    db_pool: DbPool,
35    analytics_provider: Option<Arc<dyn AnalyticsProvider>>,
36}
37
38impl std::fmt::Debug for JwtContextExtractor {
39    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40        f.debug_struct("JwtContextExtractor")
41            .field("jwt_extractor", &self.jwt_extractor)
42            .field("token_extractor", &self.token_extractor)
43            .field("db_pool", &"DbPool")
44            .field("analytics_provider", &self.analytics_provider.is_some())
45            .finish()
46    }
47}
48
49impl JwtContextExtractor {
50    pub fn new(jwt_secret: &str, db_pool: &DbPool) -> Self {
51        Self {
52            jwt_extractor: Arc::new(JwtExtractor::new(jwt_secret)),
53            token_extractor: TokenExtractor::browser_only(),
54            db_pool: Arc::clone(db_pool),
55            analytics_provider: None,
56        }
57    }
58
59    pub fn with_analytics_provider(mut self, provider: Arc<dyn AnalyticsProvider>) -> Self {
60        self.analytics_provider = Some(provider);
61        self
62    }
63
64    fn extract_jwt_context(
65        &self,
66        headers: &HeaderMap,
67    ) -> Result<JwtUserContext, ContextExtractionError> {
68        let token = self
69            .token_extractor
70            .extract(headers)
71            .map_err(|_| ContextExtractionError::MissingAuthHeader)?;
72        self.jwt_extractor
73            .extract_user_context(&token)
74            .map_err(|e| ContextExtractionError::InvalidToken(e.to_string()))
75    }
76
77    async fn validate_user_exists(
78        &self,
79        jwt_context: &JwtUserContext,
80        route_context: &str,
81    ) -> Result<(), ContextExtractionError> {
82        let user_service = UserService::new(&self.db_pool).map_err(|e| {
83            ContextExtractionError::DatabaseError(format!("Failed to create user service: {e}"))
84        })?;
85        let user_exists = user_service
86            .find_by_id(&jwt_context.user_id)
87            .await
88            .map_err(|e| {
89                ContextExtractionError::DatabaseError(format!(
90                    "Failed to check user existence: {e}"
91                ))
92            })?;
93
94        if user_exists.is_none() {
95            tracing::info!(
96                session_id = %jwt_context.session_id.as_str(),
97                user_id = %jwt_context.user_id.as_str(),
98                route = %route_context,
99                "JWT validation failed: User no longer exists in database"
100            );
101
102            return Err(ContextExtractionError::UserNotFound(format!(
103                "User {} no longer exists",
104                jwt_context.user_id.as_str()
105            )));
106        }
107        Ok(())
108    }
109
110    async fn validate_session_exists(
111        &self,
112        jwt_context: &JwtUserContext,
113        headers: &HeaderMap,
114        route_context: &str,
115    ) -> Result<(), ContextExtractionError> {
116        let Some(analytics_provider) = &self.analytics_provider else {
117            return Ok(());
118        };
119
120        let session_exists = analytics_provider
121            .find_session_by_id(&jwt_context.session_id)
122            .await
123            .map_err(|e| {
124                ContextExtractionError::DatabaseError(format!("Failed to check session: {e}"))
125            })?
126            .is_some();
127
128        if session_exists {
129            return Ok(());
130        }
131
132        tracing::info!(
133            session_id = %jwt_context.session_id.as_str(),
134            user_id = %jwt_context.user_id.as_str(),
135            route = %route_context,
136            "Creating missing session for legacy token"
137        );
138
139        let config = systemprompt_models::Config::get().map_err(|e| {
140            ContextExtractionError::DatabaseError(format!("Failed to get config: {e}"))
141        })?;
142        let expires_at =
143            chrono::Utc::now() + chrono::Duration::seconds(config.jwt_access_token_expiration);
144        let analytics = analytics_provider.extract_analytics(headers, None);
145        let session_source = jwt_context
146            .client_id
147            .as_ref()
148            .map_or(SessionSource::Api, |c| {
149                SessionSource::from_client_id(c.as_str())
150            });
151
152        analytics_provider
153            .create_session(CreateSessionInput {
154                session_id: &jwt_context.session_id,
155                user_id: Some(&jwt_context.user_id),
156                analytics: &analytics,
157                session_source,
158                is_bot: false,
159                expires_at,
160            })
161            .await
162            .map_err(|e| {
163                ContextExtractionError::DatabaseError(format!("Failed to create session: {e}"))
164            })?;
165
166        Ok(())
167    }
168
169    fn extract_common_headers(
170        &self,
171        headers: &HeaderMap,
172    ) -> (TraceId, Option<TaskId>, Option<String>, AgentName) {
173        (
174            HeaderExtractor::extract_trace_id(headers),
175            HeaderExtractor::extract_task_id(headers),
176            self.token_extractor.extract(headers).ok(),
177            HeaderExtractor::extract_agent_name(headers),
178        )
179    }
180
181    fn build_context(params: BuildContextParams) -> RequestContext {
182        let BuildContextParams {
183            jwt_context,
184            session_id,
185            user_id,
186            trace_id,
187            context_id,
188            agent_name,
189            task_id,
190            auth_token,
191        } = params;
192        let mut ctx = RequestContext::new(session_id, trace_id, context_id, agent_name)
193            .with_user_id(user_id)
194            .with_user_type(jwt_context.user_type);
195
196        if let Some(client_id) = jwt_context.client_id {
197            ctx = ctx.with_client_id(client_id);
198        }
199        if let Some(t_id) = task_id {
200            ctx = ctx.with_task_id(t_id);
201        }
202        if let Some(token) = auth_token {
203            ctx = ctx.with_auth_token(token);
204        }
205        ctx
206    }
207
208    pub async fn extract_standard(
209        &self,
210        headers: &HeaderMap,
211    ) -> Result<RequestContext, ContextExtractionError> {
212        let has_auth = headers.get("authorization").is_some();
213        let has_context_headers =
214            headers.get("x-user-id").is_some() && headers.get("x-session-id").is_some();
215
216        if has_context_headers && !has_auth {
217            return Err(ContextExtractionError::ForbiddenHeader {
218                header: "X-User-ID/X-Session-ID".to_string(),
219                reason: "Context headers require valid JWT for authentication".to_string(),
220            });
221        }
222
223        let jwt_context = self.extract_jwt_context(headers)?;
224
225        if jwt_context.session_id.as_str().is_empty() {
226            return Err(ContextExtractionError::MissingSessionId);
227        }
228        if jwt_context.user_id.as_str().is_empty() {
229            return Err(ContextExtractionError::MissingUserId);
230        }
231
232        self.validate_user_exists(&jwt_context, "").await?;
233        self.validate_session_exists(&jwt_context, headers, "")
234            .await?;
235
236        let session_id = headers
237            .get("x-session-id")
238            .and_then(|h| h.to_str().ok())
239            .map_or_else(
240                || jwt_context.session_id.clone(),
241                |s| SessionId::new(s.to_string()),
242            );
243
244        let user_id = headers
245            .get("x-user-id")
246            .and_then(|h| h.to_str().ok())
247            .map_or_else(
248                || jwt_context.user_id.clone(),
249                |s| UserId::new(s.to_string()),
250            );
251
252        let context_id = headers
253            .get("x-context-id")
254            .and_then(|h| h.to_str().ok())
255            .map_or_else(
256                || ContextId::new(String::new()),
257                |s| ContextId::new(s.to_string()),
258            );
259
260        let (trace_id, task_id, auth_token, agent_name) = self.extract_common_headers(headers);
261
262        Ok(Self::build_context(BuildContextParams {
263            jwt_context,
264            session_id,
265            user_id,
266            trace_id,
267            context_id,
268            agent_name,
269            task_id,
270            auth_token,
271        }))
272    }
273
274    pub async fn extract_mcp_a2a(
275        &self,
276        headers: &HeaderMap,
277    ) -> Result<RequestContext, ContextExtractionError> {
278        self.extract_standard(headers).await
279    }
280
281    async fn extract_from_request_impl(
282        &self,
283        request: Request<Body>,
284    ) -> Result<(RequestContext, Request<Body>), ContextExtractionError> {
285        use crate::services::middleware::context::sources::{
286            ContextIdSource, PayloadSource, TASK_BASED_CONTEXT_MARKER,
287        };
288
289        let headers = request.headers().clone();
290        let has_auth = headers.get("authorization").is_some();
291
292        if headers.get("x-context-id").is_some() && !has_auth {
293            return Err(ContextExtractionError::ForbiddenHeader {
294                header: "X-Context-ID".to_string(),
295                reason: "Context ID must be in request body (A2A spec). Use contextId field in \
296                         message."
297                    .to_string(),
298            });
299        }
300
301        let jwt_context = self.extract_jwt_context(&headers)?;
302
303        if jwt_context.session_id.as_str().is_empty() {
304            return Err(ContextExtractionError::MissingSessionId);
305        }
306        if jwt_context.user_id.as_str().is_empty() {
307            return Err(ContextExtractionError::MissingUserId);
308        }
309
310        self.validate_user_exists(&jwt_context, " (A2A route)")
311            .await?;
312        self.validate_session_exists(&jwt_context, &headers, " (A2A route)")
313            .await?;
314
315        let (body_bytes, reconstructed_request) =
316            PayloadSource::read_and_reconstruct(request).await?;
317
318        let context_source = PayloadSource::extract_context_source(&body_bytes)?;
319        let (context_id, task_id_from_payload) = match context_source {
320            ContextIdSource::Direct(id) => (ContextId::new(id), None),
321            ContextIdSource::FromTask { task_id } => (
322                ContextId::new(TASK_BASED_CONTEXT_MARKER),
323                Some(TaskId::new(task_id)),
324            ),
325        };
326
327        let (trace_id, task_id_from_header, auth_token, agent_name) =
328            self.extract_common_headers(&headers);
329
330        let task_id = task_id_from_payload.or(task_id_from_header);
331
332        let session_id = jwt_context.session_id.clone();
333        let user_id = jwt_context.user_id.clone();
334        let ctx = Self::build_context(BuildContextParams {
335            jwt_context,
336            session_id,
337            user_id,
338            trace_id,
339            context_id,
340            agent_name,
341            task_id,
342            auth_token,
343        });
344
345        Ok((ctx, reconstructed_request))
346    }
347}
348
349#[async_trait]
350impl ContextExtractor for JwtContextExtractor {
351    async fn extract_from_headers(
352        &self,
353        headers: &HeaderMap,
354    ) -> Result<RequestContext, ContextExtractionError> {
355        self.extract_standard(headers).await
356    }
357
358    async fn extract_from_request(
359        &self,
360        request: Request<Body>,
361    ) -> Result<(RequestContext, Request<Body>), ContextExtractionError> {
362        self.extract_from_request_impl(request).await
363    }
364
365    async fn extract_user_only(
366        &self,
367        headers: &HeaderMap,
368    ) -> Result<RequestContext, ContextExtractionError> {
369        self.extract_standard(headers).await
370    }
371}
systemprompt_api/services/middleware/jwt/context.rs

systemprompt_api/services/middleware/jwt/
context.rs
