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

systemprompt_api/services/proxy/engine/
mod.rs
