use std::sync::Arc;
use rmcp::handler::server::router::tool::ToolRouter;
use rmcp::handler::server::wrapper::Parameters;
use rmcp::model::{
CallToolResult, ContentBlock, Implementation, ProtocolVersion, ServerCapabilities, ServerInfo,
};
use rmcp::service::RequestContext;
use rmcp::transport::stdio;
use rmcp::{
ErrorData as McpError, RoleServer, ServerHandler, ServiceExt, schemars, tool, tool_handler,
tool_router,
};
use serde::Deserialize;
use tokio::sync::Mutex;
use crate::auth::auth_manager::{AuthManager, header_location_for};
use crate::core::config_schema::Config;
use crate::core::errors::McpifyError;
use crate::data::store::{cached_store_connection, get_endpoint};
use crate::http::auth_extractor::extract_request_credentials;
use crate::tools::call_tool::call_operation;
use crate::tools::get_tool::get_operation;
use crate::tools::search_tool::search_operations;
fn default_search_limit() -> usize {
5
}
#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub struct SearchArgs {
pub query: String,
#[serde(default = "default_search_limit")]
pub limit: usize,
}
#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub struct GetArgs {
pub operation_id: String,
}
#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub struct CallArgs {
pub operation_id: String,
#[serde(default)]
pub arguments: serde_json::Value,
}
#[derive(Clone)]
pub struct McpifyServer {
api_version: String,
config: Config,
auth_manager: Arc<Mutex<AuthManager>>,
tool_router: ToolRouter<McpifyServer>,
}
#[tool_router]
impl McpifyServer {
pub fn new(api_version: String, config: Config, auth_manager: Arc<Mutex<AuthManager>>) -> Self {
Self {
api_version,
config,
auth_manager,
tool_router: Self::tool_router(),
}
}
#[tool(
description = "Semantic search for GitHub v3 REST API operations using a natural-language query."
)]
async fn search(
&self,
Parameters(args): Parameters<SearchArgs>,
) -> Result<CallToolResult, McpError> {
let api_version = self.api_version.clone();
self.run_tool("search", async move {
let conn = cached_store_connection(&api_version)?.lock().unwrap();
search_operations(&conn, &args.query, args.limit)
})
.await
}
#[tool(
description = "Return the schema, path, method, and documentation for a specific GitHub v3 REST API operationId."
)]
async fn get(&self, Parameters(args): Parameters<GetArgs>) -> Result<CallToolResult, McpError> {
let api_version = self.api_version.clone();
self.run_tool("get", async move {
let conn = cached_store_connection(&api_version)?.lock().unwrap();
get_operation(&conn, &args.operation_id)
})
.await
}
#[tool(
description = "Validate arguments, invoke a live GitHub v3 REST API API operation, and validate the response."
)]
async fn call(
&self,
Parameters(args): Parameters<CallArgs>,
context: RequestContext<RoleServer>,
) -> Result<CallToolResult, McpError> {
let api_version = self.api_version.clone();
let config = self.config.clone();
let auth_manager = self.auth_manager.clone();
let request_credentials = context
.extensions
.get::<axum::http::request::Parts>()
.and_then(|parts| {
let (header_location, header_name) = header_location_for(config.auth_method);
extract_request_credentials(&parts.headers, header_location, header_name).ok()
});
self.run_tool("call", async move {
let endpoint = {
let conn = cached_store_connection(&api_version)?.lock().unwrap();
get_endpoint(&conn, &args.operation_id)?.ok_or_else(|| {
McpifyError::NotFound(format!("unknown operationId '{}'", args.operation_id))
})?
};
let mut auth_manager = auth_manager.lock().await;
call_operation(
&endpoint,
&config,
&mut auth_manager,
&args.operation_id,
args.arguments,
request_credentials.as_ref(),
)
.await
})
.await
}
}
impl McpifyServer {
async fn run_tool<F>(&self, tool_name: &str, fut: F) -> Result<CallToolResult, McpError>
where
F: std::future::Future<Output = anyhow::Result<serde_json::Value>>,
{
match fut.await {
Ok(value) => {
let text =
serde_json::to_string_pretty(&value).unwrap_or_else(|_| value.to_string());
Ok(CallToolResult::success(vec![ContentBlock::text(text)]))
}
Err(err) => {
tracing::error!(tool = tool_name, error = %err, "tool execution failed");
Ok(CallToolResult::error(vec![ContentBlock::text(
err.to_string(),
)]))
}
}
}
}
#[tool_handler(router = self.tool_router.clone())]
impl ServerHandler for McpifyServer {
fn get_info(&self) -> ServerInfo {
ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
.with_server_info(Implementation::from_build_env())
.with_protocol_version(ProtocolVersion::V_2024_11_05)
.with_instructions(
"Exposes exactly 3 tools -- search, get, call -- backed by an embedded \
semantic database, so you never need the full API surface in context."
.to_string(),
)
}
}
pub async fn connect_stdio<S>(server: S) -> anyhow::Result<()>
where
S: rmcp::ServerHandler,
{
let running = server.serve(stdio()).await?;
tracing::info!("MCP server connected over stdio");
running.waiting().await?;
Ok(())
}