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::{ContextId, SessionId, UserId};
10use systemprompt_models::execution::context::{ContextExtractionError, RequestContext};
11use systemprompt_security::TokenExtractor;
12use systemprompt_traits::AnalyticsProvider;
13
14use super::params::{BuildContextParams, build_context, extract_common_headers};
15use super::token::{JwtExtractor, JwtUserContext};
16use super::validation::{validate_session_exists, validate_user_exists};
17
18#[derive(Clone)]
19pub struct JwtContextExtractor {
20    jwt_extractor: Arc<JwtExtractor>,
21    token_extractor: TokenExtractor,
22    db_pool: DbPool,
23    analytics_provider: Option<Arc<dyn AnalyticsProvider>>,
24}
25
26impl std::fmt::Debug for JwtContextExtractor {
27    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28        f.debug_struct("JwtContextExtractor")
29            .field("jwt_extractor", &self.jwt_extractor)
30            .field("token_extractor", &self.token_extractor)
31            .field("db_pool", &"DbPool")
32            .field("analytics_provider", &self.analytics_provider.is_some())
33            .finish()
34    }
35}
36
37impl JwtContextExtractor {
38    pub fn new(jwt_secret: &str, db_pool: &DbPool) -> Self {
39        Self {
40            jwt_extractor: Arc::new(JwtExtractor::new(jwt_secret)),
41            token_extractor: TokenExtractor::browser_only(),
42            db_pool: Arc::clone(db_pool),
43            analytics_provider: None,
44        }
45    }
46
47    pub fn with_analytics_provider(mut self, provider: Arc<dyn AnalyticsProvider>) -> Self {
48        self.analytics_provider = Some(provider);
49        self
50    }
51
52    fn extract_jwt_context(
53        &self,
54        headers: &HeaderMap,
55    ) -> Result<JwtUserContext, ContextExtractionError> {
56        let token = self
57            .token_extractor
58            .extract(headers)
59            .map_err(|_| ContextExtractionError::MissingAuthHeader)?;
60        self.jwt_extractor
61            .extract_user_context(&token)
62            .map_err(|e| ContextExtractionError::InvalidToken(e.to_string()))
63    }
64
65    pub async fn extract_standard(
66        &self,
67        headers: &HeaderMap,
68    ) -> Result<RequestContext, ContextExtractionError> {
69        let has_auth = headers.get("authorization").is_some();
70        let has_context_headers =
71            headers.get("x-user-id").is_some() && headers.get("x-session-id").is_some();
72
73        if has_context_headers && !has_auth {
74            return Err(ContextExtractionError::ForbiddenHeader {
75                header: "X-User-ID/X-Session-ID".to_string(),
76                reason: "Context headers require valid JWT for authentication".to_string(),
77            });
78        }
79
80        let jwt_context = self.extract_jwt_context(headers)?;
81
82        if jwt_context.session_id.as_str().is_empty() {
83            return Err(ContextExtractionError::MissingSessionId);
84        }
85        if jwt_context.user_id.as_str().is_empty() {
86            return Err(ContextExtractionError::MissingUserId);
87        }
88
89        validate_user_exists(&self.db_pool, &jwt_context, "").await?;
90        validate_session_exists(self.analytics_provider.as_ref(), &jwt_context, headers, "")
91            .await?;
92
93        let session_id = headers
94            .get("x-session-id")
95            .and_then(|h| h.to_str().ok())
96            .map_or_else(
97                || jwt_context.session_id.clone(),
98                |s| SessionId::new(s.to_string()),
99            );
100
101        let user_id = headers
102            .get("x-user-id")
103            .and_then(|h| h.to_str().ok())
104            .map_or_else(
105                || jwt_context.user_id.clone(),
106                |s| UserId::new(s.to_string()),
107            );
108
109        let context_id = headers
110            .get("x-context-id")
111            .and_then(|h| h.to_str().ok())
112            .map_or_else(
113                || ContextId::new(String::new()),
114                |s| ContextId::new(s.to_string()),
115            );
116
117        let (trace_id, task_id, auth_token, agent_name) =
118            extract_common_headers(&self.token_extractor, headers);
119
120        Ok(build_context(BuildContextParams {
121            jwt_context,
122            session_id,
123            user_id,
124            trace_id,
125            context_id,
126            agent_name,
127            task_id,
128            auth_token,
129        }))
130    }
131
132    pub async fn extract_mcp_a2a(
133        &self,
134        headers: &HeaderMap,
135    ) -> Result<RequestContext, ContextExtractionError> {
136        self.extract_standard(headers).await
137    }
138
139    /// Decode a JWT for the gateway authn path and confirm the user
140    /// still exists. Returns the parsed claim set; the caller assembles
141    /// whatever request-scoped struct it needs (the gateway builds an
142    /// `AuthedPrincipal` directly — there's no need to materialise a
143    /// full `RequestContext` at the auth/gate stage).
144    pub async fn decode_for_gateway(
145        &self,
146        jwt_token: &systemprompt_identifiers::JwtToken,
147    ) -> Result<JwtUserContext, ContextExtractionError> {
148        let jwt_context = self
149            .jwt_extractor
150            .extract_user_context(jwt_token.as_str())
151            .map_err(|e| ContextExtractionError::InvalidToken(e.to_string()))?;
152
153        if jwt_context.session_id.as_str().is_empty() {
154            return Err(ContextExtractionError::MissingSessionId);
155        }
156        if jwt_context.user_id.as_str().is_empty() {
157            return Err(ContextExtractionError::MissingUserId);
158        }
159
160        validate_user_exists(&self.db_pool, &jwt_context, "gateway").await?;
161
162        Ok(jwt_context)
163    }
164
165    async fn extract_from_request_impl(
166        &self,
167        request: Request<Body>,
168    ) -> Result<(RequestContext, Request<Body>), ContextExtractionError> {
169        use crate::services::middleware::context::sources::{
170            ContextIdSource, PayloadSource, TASK_BASED_CONTEXT_MARKER,
171        };
172
173        let headers = request.headers().clone();
174        let has_auth = headers.get("authorization").is_some();
175
176        if headers.get("x-context-id").is_some() && !has_auth {
177            return Err(ContextExtractionError::ForbiddenHeader {
178                header: "X-Context-ID".to_string(),
179                reason: "Context ID must be in request body (A2A spec). Use contextId field in \
180                         message."
181                    .to_string(),
182            });
183        }
184
185        let jwt_context = self.extract_jwt_context(&headers)?;
186
187        if jwt_context.session_id.as_str().is_empty() {
188            return Err(ContextExtractionError::MissingSessionId);
189        }
190        if jwt_context.user_id.as_str().is_empty() {
191            return Err(ContextExtractionError::MissingUserId);
192        }
193
194        validate_user_exists(&self.db_pool, &jwt_context, " (A2A route)").await?;
195        validate_session_exists(
196            self.analytics_provider.as_ref(),
197            &jwt_context,
198            &headers,
199            " (A2A route)",
200        )
201        .await?;
202
203        let (body_bytes, reconstructed_request) =
204            PayloadSource::read_and_reconstruct(request).await?;
205
206        let context_source = PayloadSource::extract_context_source(&body_bytes)?;
207        let (context_id, task_id_from_payload) = match context_source {
208            ContextIdSource::Direct(id) => (ContextId::new(id), None),
209            ContextIdSource::FromTask { task_id } => {
210                (ContextId::new(TASK_BASED_CONTEXT_MARKER), Some(task_id))
211            },
212        };
213
214        let (trace_id, task_id_from_header, auth_token, agent_name) =
215            extract_common_headers(&self.token_extractor, &headers);
216
217        let task_id = task_id_from_payload.or(task_id_from_header);
218
219        let session_id = jwt_context.session_id.clone();
220        let user_id = jwt_context.user_id.clone();
221        let ctx = build_context(BuildContextParams {
222            jwt_context,
223            session_id,
224            user_id,
225            trace_id,
226            context_id,
227            agent_name,
228            task_id,
229            auth_token,
230        });
231
232        Ok((ctx, reconstructed_request))
233    }
234}
235
236#[async_trait]
237impl ContextExtractor for JwtContextExtractor {
238    async fn extract_from_headers(
239        &self,
240        headers: &HeaderMap,
241    ) -> Result<RequestContext, ContextExtractionError> {
242        self.extract_standard(headers).await
243    }
244
245    async fn extract_from_request(
246        &self,
247        request: Request<Body>,
248    ) -> Result<(RequestContext, Request<Body>), ContextExtractionError> {
249        self.extract_from_request_impl(request).await
250    }
251
252    async fn extract_user_only(
253        &self,
254        headers: &HeaderMap,
255    ) -> Result<RequestContext, ContextExtractionError> {
256        self.extract_standard(headers).await
257    }
258}
systemprompt_api/services/middleware/jwt/context.rs

systemprompt_api/services/middleware/jwt/
context.rs
