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

systemprompt_api/services/proxy/engine/
mod.rs
