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_mcp::McpServerRegistry;
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: String,
21    pub tool_name: String,
22    pub server_name: String,
23    pub server_endpoint: String,
24    pub input: serde_json::Value,
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.to_string(),
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(Path(service_name): Path<String>) -> impl IntoResponse {
112    let base_url = match Config::get() {
113        Ok(c) => c.api_external_url.clone(),
114        Err(e) => {
115            tracing::error!(error = %e, "Failed to get config");
116            return (
117                StatusCode::INTERNAL_SERVER_ERROR,
118                Json(serde_json::json!({"error": "Configuration unavailable"})),
119            )
120                .into_response();
121        },
122    };
123
124    let scopes = match get_mcp_server_scopes(&service_name).await {
125        Some(s) => s,
126        None => vec!["user".to_string()],
127    };
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
142pub async fn handle_mcp_authorization_server(
143    Path(_service_name): Path<String>,
144) -> impl IntoResponse {
145    let base_url = match Config::get() {
146        Ok(c) => c.api_external_url.clone(),
147        Err(e) => {
148            tracing::error!(error = %e, "Failed to get config");
149            return (
150                StatusCode::INTERNAL_SERVER_ERROR,
151                Json(serde_json::json!({"error": "Configuration unavailable"})),
152            )
153                .into_response();
154        },
155    };
156
157    let metadata = McpAuthorizationServerMetadata {
158        issuer: base_url.clone(),
159        authorization_endpoint: format!("{}/api/v1/core/oauth/authorize", base_url),
160        token_endpoint: format!("{}/api/v1/core/oauth/token", base_url),
161        registration_endpoint: Some(format!("{}/api/v1/core/oauth/register", base_url)),
162        scopes_supported: vec!["user".to_string(), "admin".to_string()],
163        response_types_supported: vec![ResponseType::Code.to_string()],
164        grant_types_supported: vec![
165            GrantType::AuthorizationCode.to_string(),
166            GrantType::RefreshToken.to_string(),
167        ],
168        code_challenge_methods_supported: vec![PkceMethod::S256.to_string()],
169        token_endpoint_auth_methods_supported: vec![
170            TokenAuthMethod::None.to_string(),
171            TokenAuthMethod::ClientSecretPost.to_string(),
172            TokenAuthMethod::ClientSecretBasic.to_string(),
173        ],
174    };
175
176    (StatusCode::OK, Json(metadata)).into_response()
177}
178
179pub(crate) async fn get_mcp_server_scopes(service_name: &str) -> Option<Vec<String>> {
180    if McpServerRegistry::validate().is_err() {
181        return None;
182    }
183    let registry = systemprompt_mcp::services::registry::RegistryManager;
184    match McpRegistryProvider::get_server(&registry, service_name).await {
185        Ok(server_info) if server_info.oauth.required => {
186            let scopes: Vec<String> = server_info
187                .oauth
188                .scopes
189                .iter()
190                .map(ToString::to_string)
191                .collect();
192            if scopes.is_empty() {
193                None
194            } else {
195                Some(scopes)
196            }
197        },
198        _ => None,
199    }
200}
201
202pub fn router(ctx: &AppContext) -> Router {
203    let engine = ProxyEngine::new();
204
205    let repo = match ToolUsageRepository::new(ctx.db_pool()) {
206        Ok(r) => Arc::new(r),
207        Err(e) => {
208            tracing::error!(error = %e, "Failed to initialize MCP tool usage repository");
209            return Router::new();
210        },
211    };
212
213    let state = McpState {
214        ctx: ctx.clone(),
215        repo,
216    };
217
218    Router::new()
219        .route("/executions/{id}", get(handle_get_execution))
220        .route(
221            "/{service_name}/mcp/.well-known/oauth-protected-resource",
222            get(handle_mcp_protected_resource),
223        )
224        .route(
225            "/{service_name}/mcp/.well-known/oauth-authorization-server",
226            get(handle_mcp_authorization_server),
227        )
228        .route(
229            "/{service_name}/{*path}",
230            any({
231                let ctx_clone = ctx.clone();
232                move |Path((service_name, path)): Path<(String, String)>, request| {
233                    let engine = engine.clone();
234                    let ctx = ctx_clone.clone();
235                    async move {
236                        engine
237                            .handle_mcp_request_with_path(
238                                Path((service_name, path)),
239                                State(ctx),
240                                request,
241                            )
242                            .await
243                    }
244                }
245            }),
246        )
247        .with_state(state)
248}
249
250systemprompt_runtime::register_module_api!(
251    "mcp",
252    ServiceCategory::Mcp,
253    router,
254    true,
255    systemprompt_runtime::ModuleType::Proxy
256);
systemprompt_api/routes/proxy/mcp.rs

systemprompt_api/routes/proxy/
mcp.rs
