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    pub async fn decode_for_gateway(
139        &self,
140        jwt_token: &systemprompt_identifiers::JwtToken,
141    ) -> Result<JwtUserContext, ContextExtractionError> {
142        let jwt_context = self
143            .jwt_extractor
144            .extract_user_context(jwt_token.as_str())
145            .map_err(|e| ContextExtractionError::InvalidToken(e.to_string()))?;
146
147        if jwt_context.session_id.as_str().is_empty() {
148            return Err(ContextExtractionError::MissingSessionId);
149        }
150        if jwt_context.user_id.as_str().is_empty() {
151            return Err(ContextExtractionError::MissingUserId);
152        }
153
154        validate_user_exists(&self.db_pool, &jwt_context, "gateway").await?;
155
156        Ok(jwt_context)
157    }
158
159    async fn extract_from_request_impl(
160        &self,
161        request: Request<Body>,
162    ) -> Result<(RequestContext, Request<Body>), ContextExtractionError> {
163        use crate::services::middleware::context::sources::{ContextIdSource, PayloadSource};
164
165        let headers = request.headers().clone();
166        let has_auth = headers.get("authorization").is_some();
167
168        if headers.get("x-context-id").is_some() && !has_auth {
169            return Err(ContextExtractionError::ForbiddenHeader {
170                header: "X-Context-ID".to_string(),
171                reason: "Context ID must be in request body (A2A spec). Use contextId field in \
172                         message."
173                    .to_string(),
174            });
175        }
176
177        let jwt_context = self.extract_jwt_context(&headers)?;
178
179        if jwt_context.session_id.as_str().is_empty() {
180            return Err(ContextExtractionError::MissingSessionId);
181        }
182        if jwt_context.user_id.as_str().is_empty() {
183            return Err(ContextExtractionError::MissingUserId);
184        }
185
186        validate_user_exists(&self.db_pool, &jwt_context, " (A2A route)").await?;
187        validate_session_exists(
188            self.analytics_provider.as_ref(),
189            &jwt_context,
190            &headers,
191            " (A2A route)",
192        )
193        .await?;
194
195        let (body_bytes, reconstructed_request) =
196            PayloadSource::read_and_reconstruct(request).await?;
197
198        let context_source = PayloadSource::extract_context_source(&body_bytes)?;
199        let (context_id, task_id_from_payload) = match context_source {
200            ContextIdSource::Direct(id) => (
201                ContextId::try_new(id).map_err(|e| ContextExtractionError::InvalidHeaderValue {
202                    header: "contextId".to_string(),
203                    reason: e.to_string(),
204                })?,
205                None,
206            ),
207            ContextIdSource::FromTask { task_id } => (ContextId::generate(), Some(task_id)),
208        };
209
210        let (trace_id, task_id_from_header, auth_token, agent_name) =
211            extract_common_headers(&self.token_extractor, &headers);
212
213        let task_id = task_id_from_payload.or(task_id_from_header);
214
215        let session_id = jwt_context.session_id.clone();
216        let user_id = jwt_context.user_id.clone();
217        let ctx = build_context(BuildContextParams {
218            jwt_context,
219            session_id,
220            user_id,
221            trace_id,
222            context_id,
223            agent_name,
224            task_id,
225            auth_token,
226        });
227
228        Ok((ctx, reconstructed_request))
229    }
230}
231
232#[async_trait]
233impl ContextExtractor for JwtContextExtractor {
234    async fn extract_from_headers(
235        &self,
236        headers: &HeaderMap,
237    ) -> Result<RequestContext, ContextExtractionError> {
238        self.extract_standard(headers).await
239    }
240
241    async fn extract_from_request(
242        &self,
243        request: Request<Body>,
244    ) -> Result<(RequestContext, Request<Body>), ContextExtractionError> {
245        self.extract_from_request_impl(request).await
246    }
247
248    async fn extract_user_only(
249        &self,
250        headers: &HeaderMap,
251    ) -> Result<RequestContext, ContextExtractionError> {
252        self.extract_standard(headers).await
253    }
254}
systemprompt_api/services/middleware/jwt/context.rs

systemprompt_api/services/middleware/jwt/
context.rs
