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

systemprompt_api/services/proxy/engine/
mod.rs
