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