systemprompt_api/services/middleware/context/middleware/
flavours.rs

1//! Route-specific context middleware flavours.
2//!
3//! Each flavour decides how a `RequestContext` is established for a class of
4//! route: [`PublicContextMiddleware`] admits anonymous traffic,
5//! [`UserOnlyContextMiddleware`] requires a real user from headers,
6//! [`A2AContextMiddleware`] recovers the context id from the JSON-RPC body, and
7//! [`McpContextMiddleware`] falls back to the session context so the MCP proxy
8//! can issue an OAuth challenge.
9//!
10//! Copyright (c) systemprompt.io — Business Source License 1.1.
11//! See <https://systemprompt.io> for licensing details.
12
13use std::sync::Arc;
14
15use axum::extract::Request;
16use axum::middleware::Next;
17use axum::response::Response;
18use systemprompt_identifiers::{AgentName, ContextId};
19use systemprompt_models::execution::context::RequestContext;
20use systemprompt_security::HeaderExtractor;
21use tracing::Instrument;
22
23use super::super::extractors::ContextExtractor;
24use super::error::log_error_response;
25use super::support::{DynExtractor, create_request_span, session_context_required_error};
26
27/// Public route flavour: admits `UserType::Anon`.
28///
29/// Forwards the session-derived [`RequestContext`] minted by
30/// `POST /oauth/session`, merging optional `x-context-id` / `x-agent-name`
31/// headers on top. Never touches the request body, and never invokes the
32/// extractor — the public gate has nothing to extract from anonymous traffic.
33#[derive(Clone, Copy, Debug, Default)]
34pub struct PublicContextMiddleware;
35
36impl PublicContextMiddleware {
37    #[must_use]
38    pub const fn new() -> Self {
39        Self
40    }
41
42    pub async fn handle(&self, mut request: Request, next: Next) -> Response {
43        let Some(mut req_ctx) = request.extensions().get::<RequestContext>().cloned() else {
44            let trace_id = HeaderExtractor::extract_trace_id(request.headers());
45            let path = request.uri().path().to_owned();
46            let method = request.method().to_string();
47            return session_context_required_error(&trace_id, &path, &method);
48        };
49
50        let headers = request.headers();
51        if let Some(context_id) = headers.get("x-context-id")
52            && let Ok(id) = context_id.to_str()
53        {
54            match ContextId::try_new(id.to_owned()) {
55                Ok(parsed) => req_ctx.execution.context_id = parsed,
56                Err(e) => {
57                    tracing::warn!(error = %e, "ignoring malformed x-context-id header");
58                },
59            }
60        }
61
62        if let Some(agent_name) = headers.get("x-agent-name")
63            && let Ok(name) = agent_name.to_str()
64        {
65            match AgentName::try_new(name.to_owned()) {
66                Ok(parsed) => req_ctx.execution.agent_name = parsed,
67                Err(e) => {
68                    tracing::warn!(error = %e, "ignoring malformed x-agent-name header");
69                },
70            }
71        }
72
73        let span = create_request_span(&req_ctx);
74        request.extensions_mut().insert(req_ctx);
75        next.run(request).instrument(span).await
76    }
77}
78
79#[derive(Clone)]
80pub struct UserOnlyContextMiddleware {
81    extractor: DynExtractor,
82}
83
84impl std::fmt::Debug for UserOnlyContextMiddleware {
85    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
86        f.debug_struct("UserOnlyContextMiddleware").finish()
87    }
88}
89
90impl UserOnlyContextMiddleware {
91    pub fn new<E>(extractor: E) -> Self
92    where
93        E: ContextExtractor + Send + Sync + 'static,
94    {
95        Self {
96            extractor: Arc::new(extractor),
97        }
98    }
99
100    pub async fn handle(&self, mut request: Request, next: Next) -> Response {
101        let trace_id = HeaderExtractor::extract_trace_id(request.headers());
102        let path = request.uri().path().to_owned();
103        let method = request.method().to_string();
104
105        match self.extractor.extract_from_headers(request.headers()).await {
106            Ok(context) => {
107                let span = create_request_span(&context);
108                request.extensions_mut().insert(context);
109                next.run(request).instrument(span).await
110            },
111            Err(e) => log_error_response(&e, &trace_id, &path, &method),
112        }
113    }
114}
115
116/// A2A flavour: requires a real user.
117///
118/// Parses the JSON-RPC body to recover `contextId` (the A2A wire spec carries
119/// it in the body, not headers). The body is read and rebuilt so downstream
120/// handlers can deserialise it again.
121#[derive(Clone)]
122pub struct A2AContextMiddleware {
123    extractor: DynExtractor,
124}
125
126impl std::fmt::Debug for A2AContextMiddleware {
127    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
128        f.debug_struct("A2AContextMiddleware").finish()
129    }
130}
131
132impl A2AContextMiddleware {
133    pub fn new<E>(extractor: E) -> Self
134    where
135        E: ContextExtractor + Send + Sync + 'static,
136    {
137        Self {
138            extractor: Arc::new(extractor),
139        }
140    }
141
142    pub async fn handle(&self, request: Request, next: Next) -> Response {
143        let trace_id = HeaderExtractor::extract_trace_id(request.headers());
144        let path = request.uri().path().to_owned();
145        let method = request.method().to_string();
146
147        match self.extractor.extract_from_request(request).await {
148            Ok((context, reconstructed_request)) => {
149                let span = create_request_span(&context);
150                let mut req = reconstructed_request;
151                req.extensions_mut().insert(context);
152                next.run(req).instrument(span).await
153            },
154            Err(e) => log_error_response(&e, &trace_id, &path, &method),
155        }
156    }
157}
158
159/// MCP flavour: headers-only extraction with session fallback.
160///
161/// Extracts a real user from headers when an `Authorization` header is present;
162/// otherwise forwards the session-derived [`RequestContext`] (Anon) so the
163/// downstream MCP proxy handler can emit an RFC 9728 `WWW-Authenticate` 401
164/// challenge to start the OAuth dance.
165///
166/// The session-context fallback is load-bearing: MCP clients (Cowork,
167/// Claude Code, etc.) only begin OAuth discovery on a 401 carrying the
168/// challenge — collapsing this to a 4xx-without-challenge breaks them. See
169/// `crates/tests/integration/api/routes_mcp_unauth_challenge.rs`.
170#[derive(Clone)]
171pub struct McpContextMiddleware {
172    extractor: DynExtractor,
173}
174
175impl std::fmt::Debug for McpContextMiddleware {
176    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
177        f.debug_struct("McpContextMiddleware").finish()
178    }
179}
180
181impl McpContextMiddleware {
182    pub fn new<E>(extractor: E) -> Self
183    where
184        E: ContextExtractor + Send + Sync + 'static,
185    {
186        Self {
187            extractor: Arc::new(extractor),
188        }
189    }
190
191    pub async fn handle(&self, request: Request, next: Next) -> Response {
192        let trace_id = HeaderExtractor::extract_trace_id(request.headers());
193        let path = request.uri().path().to_owned();
194        let method = request.method().to_string();
195
196        match self.extractor.extract_from_headers(request.headers()).await {
197            Ok(context) => {
198                let span = create_request_span(&context);
199                let mut req = request;
200                req.extensions_mut().insert(context);
201                next.run(req).instrument(span).await
202            },
203            Err(e) => {
204                if let Some(ctx) = request.extensions().get::<RequestContext>().cloned() {
205                    tracing::debug!(
206                        error = %e,
207                        trace_id = %trace_id,
208                        "MCP header extraction failed, using session context"
209                    );
210                    let span = create_request_span(&ctx);
211                    next.run(request).instrument(span).await
212                } else {
213                    session_context_required_error(&trace_id, &path, &method)
214                }
215            },
216        }
217    }
218}
systemprompt_api/services/middleware/context/middleware/flavours.rs

systemprompt_api/services/middleware/context/middleware/
flavours.rs
