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            .filter(|s| !s.is_empty())
113            .and_then(|s| ContextId::try_new(s).ok())
114            .unwrap_or_else(ContextId::generate);
115
116        let (trace_id, task_id, auth_token, agent_name) =
117            extract_common_headers(&self.token_extractor, headers);
118
119        Ok(build_context(BuildContextParams {
120            jwt_context,
121            session_id,
122            user_id,
123            trace_id,
124            context_id,
125            agent_name,
126            task_id,
127            auth_token,
128        }))
129    }
130
131    pub async fn extract_mcp_a2a(
132        &self,
133        headers: &HeaderMap,
134    ) -> Result<RequestContext, ContextExtractionError> {
135        self.extract_standard(headers).await
136    }
137
138    /// Decode a JWT for the gateway authn path and confirm the user
139    /// still exists. Returns the parsed claim set; the caller assembles
140    /// whatever request-scoped struct it needs (the gateway builds an
141    /// `AuthedPrincipal` directly — there's no need to materialise a
142    /// full `RequestContext` at the auth/gate stage).
143    pub async fn decode_for_gateway(
144        &self,
145        jwt_token: &systemprompt_identifiers::JwtToken,
146    ) -> Result<JwtUserContext, ContextExtractionError> {
147        let jwt_context = self
148            .jwt_extractor
149            .extract_user_context(jwt_token.as_str())
150            .map_err(|e| ContextExtractionError::InvalidToken(e.to_string()))?;
151
152        if jwt_context.session_id.as_str().is_empty() {
153            return Err(ContextExtractionError::MissingSessionId);
154        }
155        if jwt_context.user_id.as_str().is_empty() {
156            return Err(ContextExtractionError::MissingUserId);
157        }
158
159        validate_user_exists(&self.db_pool, &jwt_context, "gateway").await?;
160
161        Ok(jwt_context)
162    }
163
164    async fn extract_from_request_impl(
165        &self,
166        request: Request<Body>,
167    ) -> Result<(RequestContext, Request<Body>), ContextExtractionError> {
168        use crate::services::middleware::context::sources::{ContextIdSource, PayloadSource};
169
170        let headers = request.headers().clone();
171        let has_auth = headers.get("authorization").is_some();
172
173        if headers.get("x-context-id").is_some() && !has_auth {
174            return Err(ContextExtractionError::ForbiddenHeader {
175                header: "X-Context-ID".to_string(),
176                reason: "Context ID must be in request body (A2A spec). Use contextId field in \
177                         message."
178                    .to_string(),
179            });
180        }
181
182        let jwt_context = self.extract_jwt_context(&headers)?;
183
184        if jwt_context.session_id.as_str().is_empty() {
185            return Err(ContextExtractionError::MissingSessionId);
186        }
187        if jwt_context.user_id.as_str().is_empty() {
188            return Err(ContextExtractionError::MissingUserId);
189        }
190
191        validate_user_exists(&self.db_pool, &jwt_context, " (A2A route)").await?;
192        validate_session_exists(
193            self.analytics_provider.as_ref(),
194            &jwt_context,
195            &headers,
196            " (A2A route)",
197        )
198        .await?;
199
200        let (body_bytes, reconstructed_request) =
201            PayloadSource::read_and_reconstruct(request).await?;
202
203        let context_source = PayloadSource::extract_context_source(&body_bytes)?;
204        let (context_id, task_id_from_payload) = match context_source {
205            ContextIdSource::Direct(id) => (
206                ContextId::try_new(id).map_err(|e| ContextExtractionError::InvalidHeaderValue {
207                    header: "contextId".to_string(),
208                    reason: e.to_string(),
209                })?,
210                None,
211            ),
212            ContextIdSource::FromTask { task_id } => (ContextId::generate(), Some(task_id)),
213        };
214
215        let (trace_id, task_id_from_header, auth_token, agent_name) =
216            extract_common_headers(&self.token_extractor, &headers);
217
218        let task_id = task_id_from_payload.or(task_id_from_header);
219
220        let session_id = jwt_context.session_id.clone();
221        let user_id = jwt_context.user_id.clone();
222        let ctx = build_context(BuildContextParams {
223            jwt_context,
224            session_id,
225            user_id,
226            trace_id,
227            context_id,
228            agent_name,
229            task_id,
230            auth_token,
231        });
232
233        Ok((ctx, reconstructed_request))
234    }
235}
236
237#[async_trait]
238impl ContextExtractor for JwtContextExtractor {
239    async fn extract_from_headers(
240        &self,
241        headers: &HeaderMap,
242    ) -> Result<RequestContext, ContextExtractionError> {
243        self.extract_standard(headers).await
244    }
245
246    async fn extract_from_request(
247        &self,
248        request: Request<Body>,
249    ) -> Result<(RequestContext, Request<Body>), ContextExtractionError> {
250        self.extract_from_request_impl(request).await
251    }
252
253    async fn extract_user_only(
254        &self,
255        headers: &HeaderMap,
256    ) -> Result<RequestContext, ContextExtractionError> {
257        self.extract_standard(headers).await
258    }
259}
systemprompt_api/services/middleware/jwt/context.rs

systemprompt_api/services/middleware/jwt/
context.rs
