systemprompt_api/services/middleware/jwt/
context.rs

1//! JWT-backed request-context extractor.
2//!
3//! [`JwtContextExtractor`] implements [`ContextExtractor`] by validating the
4//! bearer token (signature, session existence, user existence, and JTI
5//! revocation) and building a `RequestContext`. It resolves the context id from
6//! the `x-context-id` header on standard routes and from the JSON-RPC body on
7//! A2A routes, and exposes a gateway decode path for pre-authenticated tokens.
8//!
9//! Copyright (c) systemprompt.io — Business Source License 1.1.
10//! See <https://systemprompt.io> for licensing details.
11
12use async_trait::async_trait;
13use axum::body::Body;
14use axum::extract::Request;
15use axum::http::HeaderMap;
16use std::sync::Arc;
17
18use crate::services::middleware::context::ContextExtractor;
19use systemprompt_identifiers::ContextId;
20use systemprompt_models::execution::context::{ContextExtractionError, RequestContext};
21use systemprompt_security::{JwtUserContext, TokenExtractor, extract_user_context};
22use systemprompt_traits::{AnalyticsProvider, UserProvider};
23
24use super::params::{BuildContextParams, build_context, extract_common_headers};
25use super::revocation::JtiRevocationChecker;
26use super::validation::{UserCache, user_is_admin, validate_session_exists, validate_user_exists};
27
28#[derive(Clone)]
29pub struct JwtContextExtractor {
30    token_extractor: TokenExtractor,
31    analytics_provider: Arc<dyn AnalyticsProvider>,
32    user_provider: Arc<dyn UserProvider>,
33    user_cache: Arc<UserCache>,
34    jti_revocation: JtiRevocationChecker,
35}
36
37impl std::fmt::Debug for JwtContextExtractor {
38    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
39        f.debug_struct("JwtContextExtractor")
40            .field("token_extractor", &self.token_extractor)
41            .finish_non_exhaustive()
42    }
43}
44
45impl JwtContextExtractor {
46    pub fn new(
47        analytics_provider: Arc<dyn AnalyticsProvider>,
48        user_provider: Arc<dyn UserProvider>,
49        jti_revocation: JtiRevocationChecker,
50    ) -> Self {
51        Self {
52            token_extractor: TokenExtractor::browser_only(),
53            analytics_provider,
54            user_provider,
55            user_cache: UserCache::new(),
56            jti_revocation,
57        }
58    }
59
60    fn extract_jwt_context(
61        &self,
62        headers: &HeaderMap,
63    ) -> Result<JwtUserContext, ContextExtractionError> {
64        let token = self
65            .token_extractor
66            .extract(headers)
67            .map_err(|_e| ContextExtractionError::MissingAuthHeader)?;
68        extract_user_context(&token)
69            .map_err(|e| ContextExtractionError::InvalidToken(e.to_string()))
70    }
71
72    async fn validate(
73        &self,
74        jwt_context: &JwtUserContext,
75        route_context: &str,
76    ) -> Result<systemprompt_traits::AuthUser, ContextExtractionError> {
77        if jwt_context.session_id.as_str().is_empty() {
78            return Err(ContextExtractionError::MissingSessionId);
79        }
80        if jwt_context.user_id.as_str().is_empty() {
81            return Err(ContextExtractionError::MissingUserId);
82        }
83        let validated = validate_user_exists(
84            &self.user_provider,
85            &self.user_cache,
86            jwt_context,
87            route_context,
88        )
89        .await?;
90        validate_session_exists(&self.analytics_provider, jwt_context, route_context).await?;
91        self.jti_revocation
92            .ensure_not_revoked(&jwt_context.jti)
93            .await?;
94        Ok(validated.user)
95    }
96
97    pub async fn extract_standard(
98        &self,
99        headers: &HeaderMap,
100    ) -> Result<RequestContext, ContextExtractionError> {
101        let jwt_context = self.extract_jwt_context(headers)?;
102        let user = self.validate(&jwt_context, "").await?;
103
104        let context_id = headers
105            .get("x-context-id")
106            .and_then(|h| h.to_str().ok())
107            .filter(|s| !s.is_empty())
108            .and_then(|s| ContextId::try_new(s).ok())
109            .unwrap_or_else(ContextId::generate);
110
111        let (trace_id, task_id, auth_token, agent_name) =
112            extract_common_headers(&self.token_extractor, headers);
113
114        let user_type = jwt_context.user_type.reconcile_with(user_is_admin(&user));
115        let session_id = jwt_context.session_id.clone();
116        let user_id = jwt_context.user_id.clone();
117
118        Ok(build_context(BuildContextParams {
119            jwt_context,
120            session_id,
121            user_id,
122            trace_id,
123            context_id,
124            agent_name,
125            task_id,
126            auth_token,
127            user_type,
128        }))
129    }
130
131    pub async fn decode_for_gateway(
132        &self,
133        jwt_token: &systemprompt_identifiers::JwtToken,
134    ) -> Result<(JwtUserContext, systemprompt_traits::AuthUser), ContextExtractionError> {
135        let jwt_context = extract_user_context(jwt_token.as_str())
136            .map_err(|e| ContextExtractionError::InvalidToken(e.to_string()))?;
137
138        let user = self.validate(&jwt_context, "gateway").await?;
139        Ok((jwt_context, user))
140    }
141
142    async fn extract_from_request_impl(
143        &self,
144        request: Request<Body>,
145    ) -> Result<(RequestContext, Request<Body>), ContextExtractionError> {
146        use crate::services::middleware::context::sources::{ContextIdSource, PayloadSource};
147
148        let headers = request.headers().clone();
149        let has_auth = headers.get("authorization").is_some();
150
151        if headers.get("x-context-id").is_some() && !has_auth {
152            return Err(ContextExtractionError::ForbiddenHeader {
153                header: "X-Context-ID".to_owned(),
154                reason: "Context ID must be in request body (A2A spec). Use contextId field in \
155                         message."
156                    .to_owned(),
157            });
158        }
159
160        let jwt_context = self.extract_jwt_context(&headers)?;
161        let user = self.validate(&jwt_context, " (A2A route)").await?;
162
163        let (body_bytes, reconstructed_request) =
164            PayloadSource::read_and_reconstruct(request).await?;
165
166        let context_source = PayloadSource::extract_context_source(&body_bytes)?;
167        let (context_id, task_id_from_payload) = match context_source {
168            ContextIdSource::Direct(id) => (
169                ContextId::try_new(id).map_err(|e| ContextExtractionError::InvalidHeaderValue {
170                    header: "contextId".to_owned(),
171                    reason: e.to_string(),
172                })?,
173                None,
174            ),
175            ContextIdSource::FromTask { task_id } => (ContextId::generate(), Some(task_id)),
176        };
177
178        let (trace_id, task_id_from_header, auth_token, agent_name) =
179            extract_common_headers(&self.token_extractor, &headers);
180
181        let task_id = task_id_from_payload.or(task_id_from_header);
182        let user_type = jwt_context.user_type.reconcile_with(user_is_admin(&user));
183
184        let session_id = jwt_context.session_id.clone();
185        let user_id = jwt_context.user_id.clone();
186        let ctx = build_context(BuildContextParams {
187            jwt_context,
188            session_id,
189            user_id,
190            trace_id,
191            context_id,
192            agent_name,
193            task_id,
194            auth_token,
195            user_type,
196        });
197
198        Ok((ctx, reconstructed_request))
199    }
200}
201
202#[async_trait]
203impl ContextExtractor for JwtContextExtractor {
204    async fn extract_from_headers(
205        &self,
206        headers: &HeaderMap,
207    ) -> Result<RequestContext, ContextExtractionError> {
208        self.extract_standard(headers).await
209    }
210
211    async fn extract_from_request(
212        &self,
213        request: Request<Body>,
214    ) -> Result<(RequestContext, Request<Body>), ContextExtractionError> {
215        self.extract_from_request_impl(request).await
216    }
217}
systemprompt_api/services/middleware/jwt/context.rs

systemprompt_api/services/middleware/jwt/
context.rs
