systemprompt_api/services/proxy/engine/
mod.rs

1//! Reverse-proxy engine for MCP and agent backends.
2//!
3//! [`ProxyEngine`] resolves a service by name, enforces access via the proxy
4//! auth boundary, forwards the request to the local backend port, and streams
5//! the response back (with SSE keep-alive). For MCP it also maintains the
6//! session-identity cache so a session-only follow-up request can be enriched
7//! with the identity established on the authenticated initialize call.
8
9mod external;
10mod handlers;
11mod mcp_session;
12#[cfg(feature = "test-api")]
13pub mod test_api;
14
15use axum::body::Body;
16use axum::extract::Request;
17use axum::http::HeaderMap;
18use axum::response::Response;
19use std::collections::HashMap;
20use std::sync::Arc;
21use systemprompt_database::ServiceConfig;
22use systemprompt_identifiers::AgentName;
23use systemprompt_models::RequestContext;
24use systemprompt_runtime::AppContext;
25use tokio::sync::RwLock;
26
27use super::auth::{AccessValidator, build_mcp_unknown_service_challenge};
28use super::backend::{HeaderInjector, ProxyError, RequestBuilder, ResponseHandler, UrlResolver};
29use super::client::ClientPool;
30use super::resolver::ServiceResolver;
31use mcp_session::SessionCache;
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum ProxyKind {
35    Mcp,
36    Agent,
37}
38
39#[derive(Debug)]
40pub struct ProxyTarget<'a> {
41    pub service_name: &'a str,
42    pub path: &'a str,
43    pub kind: ProxyKind,
44}
45
46#[derive(Debug, Clone)]
47pub struct ProxyEngine {
48    client_pool: ClientPool,
49    session_cache: SessionCache,
50}
51
52impl Default for ProxyEngine {
53    fn default() -> Self {
54        Self::new()
55    }
56}
57
58impl ProxyEngine {
59    pub fn new() -> Self {
60        Self {
61            client_pool: ClientPool::new(),
62            session_cache: Arc::new(RwLock::new(HashMap::new())),
63        }
64    }
65
66    pub async fn proxy_request(
67        &self,
68        target: ProxyTarget<'_>,
69        request: Request<Body>,
70        ctx: AppContext,
71    ) -> Result<Response<Body>, ProxyError> {
72        let ProxyTarget {
73            service_name,
74            path,
75            kind: proxy_kind,
76        } = target;
77        if request.extensions().get::<RequestContext>().is_none() {
78            tracing::warn!("RequestContext missing from request extensions");
79        }
80
81        if matches!(proxy_kind, ProxyKind::Mcp)
82            && let Ok(Some(server_config)) = ctx.mcp_registry().find_server(service_name)
83            && server_config.is_external()
84        {
85            return self
86                .proxy_external_mcp(service_name, request, ctx, server_config)
87                .await;
88        }
89
90        let service = match ServiceResolver::resolve(service_name, &ctx).await {
91            Ok(svc) => svc,
92            Err(err) => {
93                return Err(unknown_service_error(
94                    service_name,
95                    proxy_kind,
96                    &request,
97                    &ctx,
98                    err,
99                ));
100            },
101        };
102
103        let req_ctx = request.extensions().get::<RequestContext>().cloned();
104        let authenticated_user = AccessValidator::validate(
105            request.headers(),
106            service_name,
107            &service,
108            &ctx,
109            req_ctx.as_ref(),
110        )
111        .await?;
112
113        let backend_url = UrlResolver::build_backend_url("http", "127.0.0.1", service.port, path);
114
115        let method_str = request.method().to_string();
116        let request_headers = request.headers().clone();
117        let mut headers = request_headers.clone();
118        let query = request.uri().query();
119        let full_url = UrlResolver::append_query_params(backend_url, query);
120
121        let req_context = self
122            .build_forward_context(req_ctx, &service, service_name, &request_headers)
123            .await?;
124
125        inject_forward_headers(&mut headers, &req_context, service_name);
126
127        let response = self
128            .send_to_backend(request, &full_url, &headers, service_name)
129            .await?;
130
131        if service.module_name == "mcp" {
132            mcp_session::handle_mcp_response(mcp_session::McpResponseCtx {
133                cache: &self.session_cache,
134                response: &response,
135                request_headers: &request_headers,
136                req_context: &req_context,
137                authenticated_user: authenticated_user.as_ref(),
138                service_name,
139                method_str: &method_str,
140            })
141            .await;
142        }
143
144        match ResponseHandler::build_response(response) {
145            Ok(resp) => Ok(resp),
146            Err(e) => {
147                tracing::error!(service = %service_name, error = %e, "Failed to build response");
148                Err(ProxyError::InvalidResponse {
149                    service: service_name.to_owned(),
150                    reason: format!("Failed to build response: {e}"),
151                })
152            },
153        }
154    }
155
156    async fn build_forward_context(
157        &self,
158        req_ctx: Option<RequestContext>,
159        service: &ServiceConfig,
160        service_name: &str,
161        request_headers: &HeaderMap,
162    ) -> Result<RequestContext, ProxyError> {
163        let mut req_context = req_ctx.ok_or_else(|| ProxyError::MissingContext {
164            message: "Request context required - proxy cannot operate without authentication"
165                .to_owned(),
166        })?;
167
168        if service.module_name == "agent" || service.module_name == "mcp" {
169            req_context = req_context.with_agent_name(AgentName::new(service_name.to_owned()));
170        }
171
172        if service.module_name == "mcp" && req_context.auth_token().as_str().is_empty() {
173            req_context = mcp_session::enrich_with_cached_identity(
174                &self.session_cache,
175                request_headers,
176                req_context,
177                service_name,
178            )
179            .await;
180        }
181
182        Ok(req_context)
183    }
184
185    async fn send_to_backend(
186        &self,
187        request: Request<Body>,
188        full_url: &str,
189        headers: &HeaderMap,
190        service_name: &str,
191    ) -> Result<reqwest::Response, ProxyError> {
192        let method_str = request.method().to_string();
193
194        let body = RequestBuilder::extract_body(request.into_body())
195            .await
196            .map_err(|e| ProxyError::BodyExtractionFailed { source: e })?;
197
198        let reqwest_method = RequestBuilder::parse_method(&method_str)
199            .map_err(|reason| ProxyError::InvalidMethod { reason })?;
200
201        let client = self.client_pool.get_default_client();
202
203        let req_builder =
204            RequestBuilder::build_request(&client, reqwest_method, full_url, headers, body);
205
206        req_builder.send().await.map_err(|e| {
207            tracing::error!(service = %service_name, url = %full_url, error = %e, "Connection failed");
208            ProxyError::ConnectionFailed {
209                service: service_name.to_owned(),
210                url: full_url.to_owned(),
211                source: e,
212            }
213        })
214    }
215}
216
217fn unknown_service_error(
218    service_name: &str,
219    proxy_kind: ProxyKind,
220    request: &Request<Body>,
221    ctx: &AppContext,
222    err: ProxyError,
223) -> ProxyError {
224    if proxy_kind == ProxyKind::Mcp && matches!(err, ProxyError::ServiceNotFound { .. }) {
225        let req_ctx = request.extensions().get::<RequestContext>().cloned();
226        if let Some(challenge) = build_mcp_unknown_service_challenge(
227            service_name,
228            request.headers(),
229            ctx,
230            req_ctx.as_ref(),
231        ) {
232            return challenge;
233        }
234    }
235    err
236}
237
238fn inject_forward_headers(
239    headers: &mut HeaderMap,
240    req_context: &RequestContext,
241    service_name: &str,
242) {
243    let has_auth_before = headers.get("authorization").is_some();
244    let ctx_has_token = !req_context.auth_token().as_str().is_empty();
245
246    HeaderInjector::inject_context(headers, req_context);
247
248    let has_auth_after = headers.get("authorization").is_some();
249    tracing::debug!(
250        service = %service_name,
251        has_auth_before = has_auth_before,
252        ctx_has_token = ctx_has_token,
253        has_auth_after = has_auth_after,
254        "Proxy forwarding request"
255    );
256}
systemprompt_api/services/proxy/engine/mod.rs

systemprompt_api/services/proxy/engine/
mod.rs
