Skip to main content

systemprompt_api/routes/proxy/
mcp.rs

1use crate::services::proxy::ProxyEngine;
2use axum::extract::{Path, State};
3use axum::http::StatusCode;
4use axum::response::IntoResponse;
5use axum::routing::{any, get};
6use axum::{Json, Router};
7use serde::Serialize;
8use std::sync::Arc;
9use systemprompt_identifiers::McpExecutionId;
10use systemprompt_mcp::repository::ToolUsageRepository;
11use systemprompt_models::modules::ApiPaths;
12use systemprompt_models::{ApiError, Config};
13use systemprompt_oauth::{GrantType, PkceMethod, ResponseType, TokenAuthMethod};
14use systemprompt_runtime::{AppContext, ServiceCategory};
15use systemprompt_traits::McpRegistryProvider;
16
17#[derive(Debug, Serialize)]
18pub struct ToolExecutionResponse {
19    pub id: McpExecutionId,
20    pub tool_name: String,
21    pub server_name: String,
22    pub server_endpoint: String,
23    pub input: serde_json::Value,
24    #[serde(skip_serializing_if = "Option::is_none")]
25    pub output: Option<serde_json::Value>,
26    pub status: String,
27}
28
29#[derive(Clone, Debug)]
30pub struct McpState {
31    pub ctx: AppContext,
32    pub repo: Arc<ToolUsageRepository>,
33}
34
35pub async fn handle_get_execution(
36    Path(execution_id): Path<String>,
37    State(state): State<McpState>,
38) -> impl IntoResponse {
39    tracing::info!(execution_id = %execution_id, "Fetching execution");
40
41    let execution_id_typed = McpExecutionId::new(&execution_id);
42    match state.repo.find_by_id(&execution_id_typed).await {
43        Ok(Some(execution)) => {
44            let server_endpoint = ApiPaths::mcp_server_endpoint(&execution.server_name);
45
46            let input = match serde_json::from_str(&execution.input) {
47                Ok(v) => v,
48                Err(e) => {
49                    tracing::error!(execution_id = %execution_id, error = %e, "Invalid input JSON");
50                    return ApiError::internal_error(format!("Invalid input JSON: {e}"))
51                        .into_response();
52                },
53            };
54
55            let response = ToolExecutionResponse {
56                id: execution.mcp_execution_id,
57                tool_name: execution.tool_name,
58                server_name: execution.server_name.clone(),
59                server_endpoint,
60                input,
61                output: execution.output.as_deref().and_then(|s| {
62                    serde_json::from_str(s)
63                        .map_err(|e| {
64                            tracing::warn!(
65                                execution_id = %execution_id,
66                                error = %e,
67                                "Failed to parse execution output JSON"
68                            );
69                            e
70                        })
71                        .ok()
72                }),
73                status: execution.status,
74            };
75
76            tracing::info!(execution_id = %execution_id, "Execution found");
77            Json(response).into_response()
78        },
79        Ok(None) => {
80            ApiError::not_found(format!("Execution not found: {execution_id}")).into_response()
81        },
82        Err(e) => {
83            tracing::error!(execution_id = %execution_id, error = %e, "Failed to get execution");
84            ApiError::internal_error(format!("Failed to get execution: {e}")).into_response()
85        },
86    }
87}
88
89#[derive(Debug, Serialize)]
90pub struct McpProtectedResourceMetadata {
91    pub resource: String,
92    pub authorization_servers: Vec<String>,
93    pub scopes_supported: Vec<String>,
94    pub bearer_methods_supported: Vec<String>,
95    pub resource_documentation: Option<String>,
96}
97
98#[derive(Debug, Serialize)]
99pub struct McpAuthorizationServerMetadata {
100    pub issuer: String,
101    pub authorization_endpoint: String,
102    pub token_endpoint: String,
103    pub registration_endpoint: Option<String>,
104    pub scopes_supported: Vec<String>,
105    pub response_types_supported: Vec<String>,
106    pub grant_types_supported: Vec<String>,
107    pub code_challenge_methods_supported: Vec<String>,
108    pub token_endpoint_auth_methods_supported: Vec<String>,
109}
110
111pub async fn handle_mcp_protected_resource(
112    State(state): State<McpState>,
113    Path(service_name): Path<String>,
114) -> impl IntoResponse {
115    let base_url = match Config::get() {
116        Ok(c) => c.api_external_url.clone(),
117        Err(e) => {
118            tracing::error!(error = %e, "Failed to get config");
119            return (
120                StatusCode::INTERNAL_SERVER_ERROR,
121                Json(serde_json::json!({"error": "Configuration unavailable"})),
122            )
123                .into_response();
124        },
125    };
126
127    let scopes = get_mcp_server_scopes(state.ctx.mcp_registry(), &service_name)
128        .await
129        .unwrap_or_else(|| vec!["user".to_owned()]);
130
131    let resource_url = format!("{}/api/v1/mcp/{}/mcp", base_url, service_name);
132
133    let metadata = McpProtectedResourceMetadata {
134        resource: resource_url,
135        authorization_servers: vec![base_url.clone()],
136        scopes_supported: scopes,
137        bearer_methods_supported: vec!["header".to_owned()],
138        resource_documentation: Some(base_url.clone()),
139    };
140
141    (StatusCode::OK, Json(metadata)).into_response()
142}
143
144pub async fn handle_mcp_authorization_server(
145    Path(_service_name): Path<String>,
146) -> impl IntoResponse {
147    let base_url = match Config::get() {
148        Ok(c) => c.api_external_url.clone(),
149        Err(e) => {
150            tracing::error!(error = %e, "Failed to get config");
151            return (
152                StatusCode::INTERNAL_SERVER_ERROR,
153                Json(serde_json::json!({"error": "Configuration unavailable"})),
154            )
155                .into_response();
156        },
157    };
158
159    let metadata = McpAuthorizationServerMetadata {
160        issuer: base_url.clone(),
161        authorization_endpoint: format!("{}/api/v1/core/oauth/authorize", base_url),
162        token_endpoint: format!("{}/api/v1/core/oauth/token", base_url),
163        registration_endpoint: Some(format!("{}/api/v1/core/oauth/register", base_url)),
164        scopes_supported: vec!["user".to_owned(), "admin".to_owned()],
165        response_types_supported: vec![ResponseType::Code.to_string()],
166        grant_types_supported: vec![
167            GrantType::AuthorizationCode.to_string(),
168            GrantType::RefreshToken.to_string(),
169        ],
170        code_challenge_methods_supported: vec![PkceMethod::S256.to_string()],
171        token_endpoint_auth_methods_supported: vec![
172            TokenAuthMethod::None.to_string(),
173            TokenAuthMethod::ClientSecretPost.to_string(),
174            TokenAuthMethod::ClientSecretBasic.to_string(),
175        ],
176    };
177
178    (StatusCode::OK, Json(metadata)).into_response()
179}
180
181pub(in crate::routes) async fn get_mcp_server_scopes(
182    registry: &dyn McpRegistryProvider,
183    service_name: &str,
184) -> Option<Vec<String>> {
185    match registry.get_server(service_name).await {
186        Ok(server_info) if server_info.oauth.required => {
187            let scopes: Vec<String> = server_info
188                .oauth
189                .scopes
190                .iter()
191                .map(ToString::to_string)
192                .collect();
193            if scopes.is_empty() {
194                None
195            } else {
196                Some(scopes)
197            }
198        },
199        _ => None,
200    }
201}
202
203pub(in crate::routes) async fn get_mcp_server_scopes_from_resource(
204    registry: &dyn McpRegistryProvider,
205    resource_uri: &str,
206) -> Option<Vec<String>> {
207    let url = reqwest::Url::parse(resource_uri).ok()?;
208    let path = url.path();
209    let parts: Vec<&str> = path.split('/').collect();
210    if parts.len() < 6 || parts[1] != "api" || parts[3] != "mcp" || parts[5] != "mcp" {
211        return None;
212    }
213    let server_name = parts[4];
214    get_mcp_server_scopes(registry, server_name).await
215}
216
217pub fn router(ctx: &AppContext) -> Router {
218    let engine = ProxyEngine::new();
219
220    let repo = match ToolUsageRepository::new(ctx.db_pool()) {
221        Ok(r) => Arc::new(r),
222        Err(e) => {
223            tracing::error!(error = %e, "Failed to initialize MCP tool usage repository");
224            return Router::new();
225        },
226    };
227
228    let state = McpState {
229        ctx: ctx.clone(),
230        repo,
231    };
232
233    Router::new()
234        .route("/executions/{id}", get(handle_get_execution))
235        .route(
236            "/{service_name}/mcp/.well-known/oauth-protected-resource",
237            get(handle_mcp_protected_resource),
238        )
239        .route(
240            "/{service_name}/mcp/.well-known/oauth-authorization-server",
241            get(handle_mcp_authorization_server),
242        )
243        .route(
244            "/{service_name}/{*path}",
245            any({
246                let ctx_clone = ctx.clone();
247                move |Path((service_name, path)): Path<(String, String)>, request| {
248                    let engine = engine.clone();
249                    let ctx = ctx_clone.clone();
250                    async move {
251                        engine
252                            .handle_mcp_request_with_path(
253                                Path((service_name, path)),
254                                State(ctx),
255                                request,
256                            )
257                            .await
258                    }
259                }
260            }),
261        )
262        .with_state(state)
263}
264
265systemprompt_runtime::register_module_api!(
266    "mcp",
267    ServiceCategory::Mcp,
268    router,
269    true,
270    systemprompt_runtime::ModuleType::Proxy
271);