mod json_rpc;
#[cfg(test)]
mod tests;
use std::collections::HashMap;
use std::sync::mpsc::{self, SyncSender};
use std::sync::{Arc, Mutex, RwLock};
use std::time::Duration;
use crate::app::App;
use crate::application::Application;
use crate::core::New;
use crate::header::Header;
use crate::mime_type::MimeType;
use crate::range::Range;
use crate::request::Request;
use crate::response::{Response, STATUS_CODE_REASON_PHRASE};
use crate::server::ConnectionInfo;
const PROTOCOL_VERSION: &str = "2024-11-05";
#[derive(Clone, Debug)]
pub struct McpContent {
pub kind: &'static str,
pub text: String,
pub mime_type: Option<String>,
pub uri: Option<String>,
}
impl McpContent {
pub fn text(s: impl Into<String>) -> Self {
McpContent { kind: "text", text: s.into(), mime_type: None, uri: None }
}
pub fn json(s: impl Into<String>) -> Self {
McpContent { kind: "text", text: s.into(), mime_type: Some("application/json".to_string()), uri: None }
}
pub fn image(data: impl Into<String>, mime_type: impl Into<String>) -> Self {
McpContent { kind: "image", text: data.into(), mime_type: Some(mime_type.into()), uri: None }
}
pub fn embedded(uri: impl Into<String>, text: impl Into<String>, mime_type: impl Into<String>) -> Self {
McpContent { kind: "resource", text: text.into(), mime_type: Some(mime_type.into()), uri: Some(uri.into()) }
}
fn to_content_json(&self) -> String {
match self.kind {
"image" => format!(
r#"{{"type":"image","data":"{}","mimeType":"{}"}}"#,
json_escape(&self.text),
json_escape(self.mime_type.as_deref().unwrap_or("application/octet-stream")),
),
"resource" => format!(
r#"{{"type":"resource","resource":{{"uri":"{}","mimeType":"{}","text":"{}"}}}}"#,
json_escape(self.uri.as_deref().unwrap_or("")),
json_escape(self.mime_type.as_deref().unwrap_or("text/plain")),
json_escape(&self.text),
),
_ => format!(r#"{{"type":"text","text":"{}"}}"#, json_escape(&self.text)),
}
}
fn mime(&self) -> &str {
self.mime_type.as_deref().unwrap_or("text/plain")
}
}
#[derive(Clone, Debug)]
pub struct PromptMessage {
pub role: &'static str,
pub content: McpContent,
}
impl PromptMessage {
pub fn user(text: impl Into<String>) -> Self {
PromptMessage { role: "user", content: McpContent::text(text) }
}
pub fn assistant(text: impl Into<String>) -> Self {
PromptMessage { role: "assistant", content: McpContent::text(text) }
}
fn to_json(&self) -> String {
format!(
r#"{{"role":"{}","content":{}}}"#,
self.role,
self.content.to_content_json(),
)
}
}
#[derive(Clone)]
pub struct PromptArgDef {
pub name: String,
pub description: String,
pub required: bool,
}
impl PromptArgDef {
pub fn required(name: impl Into<String>, description: impl Into<String>) -> Self {
PromptArgDef { name: name.into(), description: description.into(), required: true }
}
pub fn optional(name: impl Into<String>, description: impl Into<String>) -> Self {
PromptArgDef { name: name.into(), description: description.into(), required: false }
}
}
#[derive(Debug, Clone, Default)]
pub struct McpContext {
pub client_name: Option<String>,
pub client_version: Option<String>,
pub session_id: Option<String>,
pub auth_claims: Option<String>,
pub progress_token: Option<String>,
sse_clients: Option<Arc<Mutex<Vec<SseSender>>>>,
}
impl McpContext {
pub fn report_progress(&self, progress: f64, total: Option<f64>, message: Option<&str>) {
let (Some(token), Some(sse_clients)) = (&self.progress_token, &self.sse_clients) else {
return;
};
let total_field = match total {
Some(t) => format!(r#","total":{t}"#),
None => String::new(),
};
let message_field = match message {
Some(m) => format!(r#","message":"{}""#, json_escape(m)),
None => String::new(),
};
let params = format!(
r#"{{"progressToken":{token},"progress":{progress}{total_field}{message_field}}}"#
);
let json = render_notification("notifications/progress", Some(¶ms));
broadcast_sse_to(sse_clients, &json);
}
}
#[derive(Clone, Default)]
struct StoredClientInfo {
name: Option<String>,
version: Option<String>,
}
#[derive(Debug, Clone, Copy, Default)]
pub struct ToolAnnotations {
pub read_only_hint: Option<bool>,
pub destructive_hint: Option<bool>,
pub idempotent_hint: Option<bool>,
pub open_world_hint: Option<bool>,
}
impl ToolAnnotations {
fn to_json(self) -> String {
let mut fields = Vec::with_capacity(4);
if let Some(v) = self.read_only_hint {
fields.push(format!(r#""readOnlyHint":{v}"#));
}
if let Some(v) = self.destructive_hint {
fields.push(format!(r#""destructiveHint":{v}"#));
}
if let Some(v) = self.idempotent_hint {
fields.push(format!(r#""idempotentHint":{v}"#));
}
if let Some(v) = self.open_world_hint {
fields.push(format!(r#""openWorldHint":{v}"#));
}
format!("{{{}}}", fields.join(","))
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum LogLevel {
Debug,
Info,
Notice,
Warning,
Error,
Critical,
Alert,
Emergency,
}
impl LogLevel {
pub fn parse(s: &str) -> Option<Self> {
match s {
"debug" => Some(LogLevel::Debug),
"info" => Some(LogLevel::Info),
"notice" => Some(LogLevel::Notice),
"warning" => Some(LogLevel::Warning),
"error" => Some(LogLevel::Error),
"critical" => Some(LogLevel::Critical),
"alert" => Some(LogLevel::Alert),
"emergency" => Some(LogLevel::Emergency),
_ => None,
}
}
pub fn as_str(self) -> &'static str {
match self {
LogLevel::Debug => "debug",
LogLevel::Info => "info",
LogLevel::Notice => "notice",
LogLevel::Warning => "warning",
LogLevel::Error => "error",
LogLevel::Critical => "critical",
LogLevel::Alert => "alert",
LogLevel::Emergency => "emergency",
}
}
}
type ToolFn = Arc<dyn Fn(McpContext, &str) -> Result<McpContent, String> + Send + Sync>;
type ResourceFn = Arc<dyn Fn(&str) -> Result<McpContent, String> + Send + Sync>;
type PromptFn = Arc<dyn Fn(&str) -> Result<Vec<PromptMessage>, String> + Send + Sync>;
#[derive(Clone)]
struct ToolDef {
name: String,
description: String,
input_schema: String,
annotations: Option<ToolAnnotations>,
handler: ToolFn,
}
#[derive(Clone)]
struct ResourceDef {
uri_template: String,
name: String,
description: String,
handler: ResourceFn,
}
#[derive(Clone)]
struct PromptDef {
name: String,
description: String,
arguments: Vec<PromptArgDef>,
handler: PromptFn,
}
#[derive(Clone)]
pub struct McpServer {
server_name: String,
server_version: String,
path: String,
tools: Arc<RwLock<Vec<ToolDef>>>,
resources: Arc<RwLock<Vec<ResourceDef>>>,
prompts: Arc<RwLock<Vec<PromptDef>>>,
fallback: Option<Arc<dyn Application + Send + Sync>>,
auth_token: Option<String>,
page_size: Option<usize>,
sessions: Arc<Mutex<HashMap<String, StoredClientInfo>>>,
sse_clients: Arc<Mutex<Vec<SseSender>>>,
logging_enabled: bool,
min_log_level: Arc<Mutex<LogLevel>>,
}
type SseSender = SyncSender<Vec<u8>>;
const SSE_CHANNEL_CAPACITY: usize = 32;
const SSE_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15);
impl McpServer {
pub fn new(name: impl Into<String>, version: impl Into<String>) -> Self {
McpServer {
server_name: name.into(),
server_version: version.into(),
path: "/mcp".to_string(),
tools: Arc::new(RwLock::new(Vec::new())),
resources: Arc::new(RwLock::new(Vec::new())),
prompts: Arc::new(RwLock::new(Vec::new())),
fallback: None,
auth_token: None,
page_size: None,
sessions: Arc::new(Mutex::new(HashMap::new())),
sse_clients: Arc::new(Mutex::new(Vec::new())),
logging_enabled: false,
min_log_level: Arc::new(Mutex::new(LogLevel::Debug)),
}
}
pub fn page_size(mut self, n: usize) -> Self {
self.page_size = Some(n.max(1));
self
}
pub fn notify(&self, method: &str, params_json: Option<&str>) {
let json = render_notification(method, params_json);
broadcast_sse_to(&self.sse_clients, &json);
}
fn start_sse_stream(&self) -> Response {
let (tx, rx) = mpsc::sync_channel::<Vec<u8>>(SSE_CHANNEL_CAPACITY);
self.sse_clients.lock().unwrap().push(tx);
let mut response = Response::new();
response.status_code = *STATUS_CODE_REASON_PHRASE.n200_ok.status_code;
response.reason_phrase = STATUS_CODE_REASON_PHRASE.n200_ok.reason_phrase.to_string();
response.headers.push(Header {
name: Header::_CONTENT_TYPE.to_string(),
value: "text/event-stream".to_string(),
});
response.headers.push(Header {
name: Header::_CACHE_CONTROL.to_string(),
value: "no-cache".to_string(),
});
response.headers.push(Header {
name: "X-Accel-Buffering".to_string(),
value: "no".to_string(),
});
response.stream_pipe = Some(Box::new(SseChannelReader::new(rx)));
response
}
pub fn logging_enabled(mut self) -> Self {
self.logging_enabled = true;
self
}
pub fn log(&self, level: LogLevel, logger: Option<&str>, data_json: &str) {
if level < *self.min_log_level.lock().unwrap() {
return;
}
let logger_field = match logger {
Some(l) => format!(r#","logger":"{}""#, json_escape(l)),
None => String::new(),
};
let params = format!(r#"{{"level":"{}"{logger_field},"data":{data_json}}}"#, level.as_str());
self.notify("notifications/message", Some(¶ms));
}
fn do_set_log_level(&self, body: &str) -> Result<String, (i32, String)> {
let params = json_rpc::extract_raw(body, "params")
.ok_or((json_rpc::INVALID_PARAMS, "Missing params".to_string()))?;
let level_str = json_rpc::extract_str(¶ms, "level")
.ok_or((json_rpc::INVALID_PARAMS, "Missing level".to_string()))?;
let level = LogLevel::parse(&level_str)
.ok_or_else(|| (json_rpc::INVALID_PARAMS, format!("Unknown log level: {level_str}")))?;
*self.min_log_level.lock().unwrap() = level;
Ok("{}".to_string())
}
pub fn register_tool<F>(&self, name: &str, description: &str, input_schema: &str, handler: F)
where
F: Fn(&str) -> Result<McpContent, String> + Send + Sync + 'static,
{
self.tools.write().unwrap().push(ToolDef {
name: name.to_string(),
description: description.to_string(),
input_schema: input_schema.to_string(),
annotations: None,
handler: Arc::new(move |_ctx: McpContext, args: &str| handler(args)),
});
self.notify("notifications/tools/list_changed", None);
}
pub fn remove_tool(&self, name: &str) -> bool {
let removed = {
let mut tools = self.tools.write().unwrap();
let before = tools.len();
tools.retain(|t| t.name != name);
tools.len() != before
};
if removed {
self.notify("notifications/tools/list_changed", None);
}
removed
}
pub fn register_resource<F>(&self, uri_template: &str, name: &str, description: &str, handler: F)
where
F: Fn(&str) -> Result<McpContent, String> + Send + Sync + 'static,
{
self.resources.write().unwrap().push(ResourceDef {
uri_template: uri_template.to_string(),
name: name.to_string(),
description: description.to_string(),
handler: Arc::new(handler),
});
self.notify("notifications/resources/list_changed", None);
}
pub fn remove_resource(&self, uri_template: &str) -> bool {
let removed = {
let mut resources = self.resources.write().unwrap();
let before = resources.len();
resources.retain(|r| r.uri_template != uri_template);
resources.len() != before
};
if removed {
self.notify("notifications/resources/list_changed", None);
}
removed
}
pub fn register_prompt<F>(&self, name: &str, description: &str, handler: F)
where
F: Fn(&str) -> Result<Vec<PromptMessage>, String> + Send + Sync + 'static,
{
self.prompts.write().unwrap().push(PromptDef {
name: name.to_string(),
description: description.to_string(),
arguments: vec![],
handler: Arc::new(handler),
});
self.notify("notifications/prompts/list_changed", None);
}
pub fn remove_prompt(&self, name: &str) -> bool {
let removed = {
let mut prompts = self.prompts.write().unwrap();
let before = prompts.len();
prompts.retain(|p| p.name != name);
prompts.len() != before
};
if removed {
self.notify("notifications/prompts/list_changed", None);
}
removed
}
pub fn require_bearer(mut self, token: impl Into<String>) -> Self {
self.auth_token = Some(token.into());
self
}
pub fn wrap(mut self, app: impl Application + Send + Sync + 'static) -> Self {
self.fallback = Some(Arc::new(app));
self
}
pub fn at(mut self, path: impl Into<String>) -> Self {
self.path = path.into();
self
}
pub fn tool<F>(self, name: &str, description: &str, input_schema: &str, handler: F) -> Self
where
F: Fn(&str) -> Result<McpContent, String> + Send + Sync + 'static,
{
self.tools.write().unwrap().push(ToolDef {
name: name.to_string(),
description: description.to_string(),
input_schema: input_schema.to_string(),
annotations: None,
handler: Arc::new(move |_ctx: McpContext, args: &str| handler(args)),
});
self
}
pub fn tool_annotated<F>(
self,
name: &str,
description: &str,
input_schema: &str,
annotations: ToolAnnotations,
handler: F,
) -> Self
where
F: Fn(&str) -> Result<McpContent, String> + Send + Sync + 'static,
{
self.tools.write().unwrap().push(ToolDef {
name: name.to_string(),
description: description.to_string(),
input_schema: input_schema.to_string(),
annotations: Some(annotations),
handler: Arc::new(move |_ctx: McpContext, args: &str| handler(args)),
});
self
}
pub fn tool_with_context<F>(self, name: &str, description: &str, input_schema: &str, handler: F) -> Self
where
F: Fn(McpContext, &str) -> Result<McpContent, String> + Send + Sync + 'static,
{
self.tools.write().unwrap().push(ToolDef {
name: name.to_string(),
description: description.to_string(),
input_schema: input_schema.to_string(),
annotations: None,
handler: Arc::new(handler),
});
self
}
pub fn resource<F>(self, uri_template: &str, name: &str, description: &str, handler: F) -> Self
where
F: Fn(&str) -> Result<McpContent, String> + Send + Sync + 'static,
{
self.resources.write().unwrap().push(ResourceDef {
uri_template: uri_template.to_string(),
name: name.to_string(),
description: description.to_string(),
handler: Arc::new(handler),
});
self
}
pub fn prompt<F>(self, name: &str, description: &str, handler: F) -> Self
where
F: Fn(&str) -> Result<Vec<PromptMessage>, String> + Send + Sync + 'static,
{
self.prompts.write().unwrap().push(PromptDef {
name: name.to_string(),
description: description.to_string(),
arguments: vec![],
handler: Arc::new(handler),
});
self
}
pub fn prompt_with_args<F>(
self,
name: &str,
description: &str,
args: Vec<PromptArgDef>,
handler: F,
) -> Self
where
F: Fn(&str) -> Result<Vec<PromptMessage>, String> + Send + Sync + 'static,
{
self.prompts.write().unwrap().push(PromptDef {
name: name.to_string(),
description: description.to_string(),
arguments: args,
handler: Arc::new(handler),
});
self
}
pub fn handle_request(&self, body: &str) -> Response {
self.handle_request_with_context(body, McpContext::default())
}
pub fn handle_request_with_context(&self, body: &str, ctx: McpContext) -> Response {
let trimmed = body.trim_start();
if trimmed.starts_with('[') {
return self.handle_batch(trimmed, ctx);
}
let method = match json_rpc::extract_str(body, "method") {
Some(m) => m,
None => return rpc_error(None, json_rpc::INVALID_REQUEST, "Missing method"),
};
let id = json_rpc::extract_id(body);
if method == "notifications/initialized" || (id.is_none() && method != "ping") {
return no_content();
}
let result = self.dispatch(&method, body, ctx);
let id_str = id.as_deref().unwrap_or("null");
let is_ok = result.is_ok();
let mut response = json_response(&Self::format_result(id_str, &result));
if method == "initialize" && is_ok {
self.start_session(body, &mut response);
}
response
}
fn handle_batch(&self, array_body: &str, ctx: McpContext) -> Response {
let elements = json_rpc::split_array_elements(array_body);
if elements.is_empty() {
return rpc_error(None, json_rpc::INVALID_REQUEST, "Invalid Request");
}
let mut parts: Vec<String> = Vec::new();
let mut session_init_body: Option<String> = None;
for elem in &elements {
let method = match json_rpc::extract_str(elem, "method") {
Some(m) => m,
None => {
parts.push(Self::format_result(
"null",
&Err((json_rpc::INVALID_REQUEST, "Missing method".to_string())),
));
continue;
}
};
let id = json_rpc::extract_id(elem);
if method == "notifications/initialized" || (id.is_none() && method != "ping") {
continue;
}
let result = self.dispatch(&method, elem, ctx.clone());
let id_str = id.as_deref().unwrap_or("null");
let is_ok = result.is_ok();
if method == "initialize" && is_ok && session_init_body.is_none() {
session_init_body = Some(elem.clone());
}
parts.push(Self::format_result(id_str, &result));
}
if parts.is_empty() {
return no_content();
}
let mut response = json_response(&format!("[{}]", parts.join(",")));
if let Some(init_body) = session_init_body {
self.start_session(&init_body, &mut response);
}
response
}
fn dispatch(&self, method: &str, body: &str, ctx: McpContext) -> Result<String, (i32, String)> {
match method {
"initialize" => self.do_initialize(body),
"ping" => Ok("{}".to_string()),
"tools/list" => self.do_tools_list(body),
"tools/call" => self.do_tools_call(body, ctx),
"resources/list" => self.do_resources_list(body),
"resources/read" => self.do_resources_read(body),
"prompts/list" => self.do_prompts_list(body),
"prompts/get" => self.do_prompts_get(body),
"logging/setLevel" => self.do_set_log_level(body),
_ => Err((json_rpc::METHOD_NOT_FOUND, format!("Unknown method: {method}"))),
}
}
fn format_result(id_str: &str, result: &Result<String, (i32, String)>) -> String {
match result {
Ok(result_json) => format!(
r#"{{"jsonrpc":"2.0","result":{result_json},"id":{id_str}}}"#
),
Err((code, msg)) => {
let escaped = json_escape(msg);
format!(
r#"{{"jsonrpc":"2.0","error":{{"code":{code},"message":"{escaped}"}},"id":{id_str}}}"#
)
}
}
}
fn start_session(&self, body: &str, response: &mut Response) {
let client_info = json_rpc::extract_raw(body, "params")
.and_then(|p| json_rpc::extract_raw(&p, "clientInfo"));
let (name, version) = match &client_info {
Some(info) => (
json_rpc::extract_str(info, "name"),
json_rpc::extract_str(info, "version"),
),
None => (None, None),
};
eprintln!(
"[mcp] initialize from client {} v{}",
name.as_deref().unwrap_or("unknown"),
version.as_deref().unwrap_or("unknown"),
);
let session_id = crate::request_id::generate_request_id();
self.sessions
.lock()
.unwrap()
.insert(session_id.clone(), StoredClientInfo { name, version });
response.headers.push(Header {
name: "Mcp-Session-Id".to_string(),
value: session_id,
});
}
fn do_initialize(&self, body: &str) -> Result<String, (i32, String)> {
let params = json_rpc::extract_raw(body, "params");
let client_version = params.as_deref().and_then(|p| json_rpc::extract_str(p, "protocolVersion"));
let negotiated_version: &str = match client_version.as_deref() {
Some(v) if v < PROTOCOL_VERSION => v,
_ => PROTOCOL_VERSION,
};
let logging_cap = if self.logging_enabled { r#","logging":{}"# } else { "" };
let caps = format!(
r#"{{"tools":{{"listChanged":true}},"resources":{{"subscribe":false,"listChanged":true}},"prompts":{{"listChanged":true}}{logging_cap}}}"#
);
Ok(format!(
r#"{{"protocolVersion":"{}","capabilities":{caps},"serverInfo":{{"name":"{}","version":"{}"}}}}"#,
json_escape(negotiated_version),
json_escape(&self.server_name),
json_escape(&self.server_version),
))
}
fn do_tools_list(&self, body: &str) -> Result<String, (i32, String)> {
let items: Vec<String> = self.tools.read().unwrap().iter().map(|t| {
let annotations = match t.annotations {
Some(a) => format!(r#","annotations":{}"#, a.to_json()),
None => String::new(),
};
format!(
r#"{{"name":"{}","description":"{}","inputSchema":{}{}}}"#,
json_escape(&t.name),
json_escape(&t.description),
t.input_schema,
annotations,
)
}).collect();
let (page, next_cursor) = self.paginate(&items, body)?;
Ok(format!(r#"{{"tools":[{}]{}}}"#, page.join(","), next_cursor_json(&next_cursor)))
}
fn do_tools_call(&self, body: &str, ctx: McpContext) -> Result<String, (i32, String)> {
let params = json_rpc::extract_raw(body, "params")
.ok_or((json_rpc::INVALID_PARAMS, "Missing params".to_string()))?;
let name = json_rpc::extract_str(¶ms, "name")
.ok_or((json_rpc::INVALID_PARAMS, "Missing tool name".to_string()))?;
let args = json_rpc::extract_raw(¶ms, "arguments")
.unwrap_or_else(|| "{}".to_string());
let progress_token = json_rpc::extract_raw(¶ms, "_meta")
.and_then(|meta| json_rpc::extract_raw(&meta, "progressToken"));
let ctx = McpContext { progress_token, ..ctx };
let handler = {
let tools = self.tools.read().unwrap();
tools.iter().find(|t| t.name == name).map(|t| t.handler.clone())
}.ok_or_else(|| (json_rpc::INVALID_PARAMS, format!("Unknown tool: {name}")))?;
match handler(ctx, &args) {
Ok(c) => Ok(format!(
r#"{{"content":[{}],"isError":false}}"#,
c.to_content_json(),
)),
Err(e) => {
let escaped = json_escape(&e);
Ok(format!(
r#"{{"content":[{{"type":"text","text":"{escaped}"}}],"isError":true}}"#
))
}
}
}
fn do_resources_list(&self, body: &str) -> Result<String, (i32, String)> {
let items: Vec<String> = self.resources.read().unwrap().iter().map(|r| {
format!(
r#"{{"uri":"{}","name":"{}","description":"{}","mimeType":"text/plain"}}"#,
json_escape(&r.uri_template),
json_escape(&r.name),
json_escape(&r.description),
)
}).collect();
let (page, next_cursor) = self.paginate(&items, body)?;
Ok(format!(r#"{{"resources":[{}]{}}}"#, page.join(","), next_cursor_json(&next_cursor)))
}
fn do_resources_read(&self, body: &str) -> Result<String, (i32, String)> {
let params = json_rpc::extract_raw(body, "params")
.ok_or((json_rpc::INVALID_PARAMS, "Missing params".to_string()))?;
let uri = json_rpc::extract_str(¶ms, "uri")
.ok_or((json_rpc::INVALID_PARAMS, "Missing uri".to_string()))?;
let handler = {
let resources = self.resources.read().unwrap();
resources.iter().find(|r| uri_matches(&r.uri_template, &uri)).map(|r| r.handler.clone())
}.ok_or_else(|| (json_rpc::INVALID_PARAMS, format!("Resource not found: {uri}")))?;
match handler(&uri) {
Ok(c) => {
let text_esc = json_escape(&c.text);
let uri_esc = json_escape(&uri);
Ok(format!(
r#"{{"contents":[{{"uri":"{uri_esc}","mimeType":"{}","text":"{text_esc}"}}]}}"#,
c.mime(),
))
}
Err(e) => Err((json_rpc::INVALID_PARAMS, e)),
}
}
fn do_prompts_list(&self, body: &str) -> Result<String, (i32, String)> {
let items: Vec<String> = self.prompts.read().unwrap().iter().map(|p| {
let arg_defs: Vec<String> = p.arguments.iter().map(|a| {
format!(
r#"{{"name":"{}","description":"{}","required":{}}}"#,
json_escape(&a.name),
json_escape(&a.description),
a.required,
)
}).collect();
format!(
r#"{{"name":"{}","description":"{}","arguments":[{}]}}"#,
json_escape(&p.name),
json_escape(&p.description),
arg_defs.join(","),
)
}).collect();
let (page, next_cursor) = self.paginate(&items, body)?;
Ok(format!(r#"{{"prompts":[{}]{}}}"#, page.join(","), next_cursor_json(&next_cursor)))
}
fn paginate<'a>(&self, items: &'a [String], body: &str) -> Result<(&'a [String], Option<String>), (i32, String)> {
let page_size = match self.page_size {
Some(n) => n,
None => return Ok((items, None)),
};
let cursor = json_rpc::extract_raw(body, "params")
.and_then(|p| json_rpc::extract_str(&p, "cursor"));
let offset = match cursor {
Some(c) => decode_cursor(&c)
.ok_or((json_rpc::INVALID_PARAMS, "Invalid cursor".to_string()))?,
None => 0,
};
if offset >= items.len() {
return Ok((&[], None));
}
let end = (offset + page_size).min(items.len());
let next_cursor = if end < items.len() { Some(encode_cursor(end)) } else { None };
Ok((&items[offset..end], next_cursor))
}
fn do_prompts_get(&self, body: &str) -> Result<String, (i32, String)> {
let params = json_rpc::extract_raw(body, "params")
.ok_or((json_rpc::INVALID_PARAMS, "Missing params".to_string()))?;
let name = json_rpc::extract_str(¶ms, "name")
.ok_or((json_rpc::INVALID_PARAMS, "Missing prompt name".to_string()))?;
let args = json_rpc::extract_raw(¶ms, "arguments")
.unwrap_or_else(|| "{}".to_string());
let (description, handler) = {
let prompts = self.prompts.read().unwrap();
prompts.iter().find(|p| p.name == name).map(|p| (p.description.clone(), p.handler.clone()))
}.ok_or_else(|| (json_rpc::INVALID_PARAMS, format!("Unknown prompt: {name}")))?;
match handler(&args) {
Ok(msgs) => {
let msg_jsons: Vec<String> = msgs.iter().map(|m| m.to_json()).collect();
Ok(format!(
r#"{{"description":"{}","messages":[{}]}}"#,
json_escape(&description),
msg_jsons.join(","),
))
}
Err(e) => Err((json_rpc::INVALID_PARAMS, e)),
}
}
fn context_for(&self, request: &Request) -> McpContext {
let session_id = request
.get_header("Mcp-Session-Id".to_string())
.map(|h| h.value.clone());
let (client_name, client_version) = match &session_id {
Some(sid) => match self.sessions.lock().unwrap().get(sid) {
Some(info) => (info.name.clone(), info.version.clone()),
None => (None, None),
},
None => (None, None),
};
McpContext {
client_name,
client_version,
session_id,
auth_claims: None,
progress_token: None,
sse_clients: Some(self.sse_clients.clone()),
}
}
}
fn render_notification(method: &str, params_json: Option<&str>) -> String {
let params_field = match params_json {
Some(p) => format!(r#","params":{p}"#),
None => String::new(),
};
format!(r#"{{"jsonrpc":"2.0","method":"{}"{}}}"#, json_escape(method), params_field)
}
fn broadcast_sse_to(clients: &Arc<Mutex<Vec<SseSender>>>, json: &str) {
let frame = format!("data: {json}\n\n").into_bytes();
let mut clients = clients.lock().unwrap();
clients.retain(|tx| tx.try_send(frame.clone()).is_ok());
}
struct SseChannelReader {
rx: mpsc::Receiver<Vec<u8>>,
leftover: Vec<u8>,
}
impl SseChannelReader {
fn new(rx: mpsc::Receiver<Vec<u8>>) -> Self {
SseChannelReader { rx, leftover: Vec::new() }
}
}
impl std::io::Read for SseChannelReader {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
if self.leftover.is_empty() {
loop {
match self.rx.recv_timeout(SSE_KEEPALIVE_INTERVAL) {
Ok(frame) => { self.leftover = frame; break; }
Err(mpsc::RecvTimeoutError::Timeout) => {
self.leftover = b": keep-alive\n\n".to_vec();
break;
}
Err(mpsc::RecvTimeoutError::Disconnected) => return Ok(0),
}
}
}
let n = self.leftover.len().min(buf.len());
buf[..n].copy_from_slice(&self.leftover[..n]);
self.leftover.drain(..n);
Ok(n)
}
}
impl Application for McpServer {
fn execute(&self, request: &Request, connection: &ConnectionInfo) -> Result<Response, String> {
if request.request_uri == self.path {
if let Some(expected) = &self.auth_token {
let provided = request.headers.iter()
.find(|h| h.name.eq_ignore_ascii_case("authorization"))
.map(|h| h.value.as_str())
.unwrap_or("");
let bearer = provided.strip_prefix("Bearer ").unwrap_or("");
if bearer != expected.as_str() {
return Ok(unauthorized());
}
}
return Ok(match request.method.as_str() {
"POST" => {
let body = std::str::from_utf8(&request.body).unwrap_or("");
let ctx = self.context_for(request);
self.handle_request_with_context(body, ctx)
}
"GET" => self.start_sse_stream(),
"OPTIONS" => {
let mut r = Response::new();
r.status_code = *STATUS_CODE_REASON_PHRASE.n200_ok.status_code;
r.reason_phrase = STATUS_CODE_REASON_PHRASE.n200_ok.reason_phrase.to_string();
r.headers.push(Header {
name: "Allow".to_string(),
value: "GET, POST, OPTIONS".to_string(),
});
r
}
_ => {
let mut r = Response::new();
r.status_code = *STATUS_CODE_REASON_PHRASE.n405_method_not_allowed.status_code;
r.reason_phrase = STATUS_CODE_REASON_PHRASE.n405_method_not_allowed.reason_phrase.to_string();
r.headers.push(Header {
name: "Allow".to_string(),
value: "GET, POST, OPTIONS".to_string(),
});
r.content_range_list = vec![Range::get_content_range(
b"MCP endpoint only accepts GET (SSE) or POST".to_vec(),
MimeType::TEXT_PLAIN.to_string(),
)];
r
}
});
}
match &self.fallback {
Some(app) => app.execute(request, connection),
None => App::new().execute(request, connection),
}
}
}
pub fn extract_arg(arguments: &str, name: &str) -> Option<String> {
json_rpc::extract_str(arguments, name)
}
fn json_response(body: &str) -> Response {
let mut r = Response::new();
r.status_code = *STATUS_CODE_REASON_PHRASE.n200_ok.status_code;
r.reason_phrase = STATUS_CODE_REASON_PHRASE.n200_ok.reason_phrase.to_string();
r.content_range_list = vec![Range::get_content_range(
body.as_bytes().to_vec(),
MimeType::APPLICATION_JSON.to_string(),
)];
r
}
fn no_content() -> Response {
let mut r = Response::new();
r.status_code = *STATUS_CODE_REASON_PHRASE.n202_accepted.status_code;
r.reason_phrase = STATUS_CODE_REASON_PHRASE.n202_accepted.reason_phrase.to_string();
r
}
fn unauthorized() -> Response {
let mut r = Response::new();
r.status_code = *STATUS_CODE_REASON_PHRASE.n401_unauthorized.status_code;
r.reason_phrase = STATUS_CODE_REASON_PHRASE.n401_unauthorized.reason_phrase.to_string();
r.headers.push(Header {
name: "WWW-Authenticate".to_string(),
value: "Bearer".to_string(),
});
r.content_range_list = vec![Range::get_content_range(
b"Unauthorized".to_vec(),
MimeType::TEXT_PLAIN.to_string(),
)];
r
}
fn rpc_error(id: Option<&str>, code: i32, message: &str) -> Response {
let id_str = id.unwrap_or("null");
let escaped = json_escape(message);
json_response(&format!(
r#"{{"jsonrpc":"2.0","error":{{"code":{code},"message":"{escaped}"}},"id":{id_str}}}"#
))
}
pub(crate) fn json_escape(s: &str) -> String {
let mut out = String::with_capacity(s.len() + 4);
for ch in s.chars() {
match ch {
'"' => out.push_str("\\\""),
'\\' => out.push_str("\\\\"),
'\n' => out.push_str("\\n"),
'\r' => out.push_str("\\r"),
'\t' => out.push_str("\\t"),
c if (c as u32) < 0x20 => { let _ = std::fmt::Write::write_fmt(&mut out, format_args!("\\u{:04x}", c as u32)); }
c => out.push(c),
}
}
out
}
fn next_cursor_json(next_cursor: &Option<String>) -> String {
match next_cursor {
Some(c) => format!(r#","nextCursor":"{}""#, json_escape(c)),
None => String::new(),
}
}
fn encode_cursor(offset: usize) -> String {
base64_encode(offset.to_string().as_bytes())
}
fn decode_cursor(cursor: &str) -> Option<usize> {
String::from_utf8(base64_decode(cursor)?).ok()?.parse().ok()
}
const BASE64_TABLE: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
fn base64_encode(data: &[u8]) -> String {
let mut out = String::with_capacity((data.len() + 2) / 3 * 4);
for chunk in data.chunks(3) {
let b0 = chunk[0] as u32;
let b1 = chunk.get(1).copied().unwrap_or(0) as u32;
let b2 = chunk.get(2).copied().unwrap_or(0) as u32;
let n = (b0 << 16) | (b1 << 8) | b2;
out.push(BASE64_TABLE[((n >> 18) & 0x3F) as usize] as char);
out.push(BASE64_TABLE[((n >> 12) & 0x3F) as usize] as char);
out.push(if chunk.len() > 1 { BASE64_TABLE[((n >> 6) & 0x3F) as usize] as char } else { '=' });
out.push(if chunk.len() > 2 { BASE64_TABLE[(n & 0x3F) as usize] as char } else { '=' });
}
out
}
fn base64_decode(s: &str) -> Option<Vec<u8>> {
fn sextet(c: u8) -> Option<u32> {
match c {
b'A'..=b'Z' => Some((c - b'A') as u32),
b'a'..=b'z' => Some((c - b'a' + 26) as u32),
b'0'..=b'9' => Some((c - b'0' + 52) as u32),
b'+' => Some(62),
b'/' => Some(63),
_ => None,
}
}
let trimmed = s.trim_end_matches('=');
let bytes = trimmed.as_bytes();
let mut out = Vec::with_capacity(bytes.len() * 3 / 4 + 3);
for chunk in bytes.chunks(4) {
if chunk.len() == 1 {
return None; }
let vals: Vec<u32> = chunk.iter().map(|&b| sextet(b)).collect::<Option<Vec<_>>>()?;
let n = vals.iter().enumerate().fold(0u32, |acc, (i, &v)| acc | (v << (18 - 6 * i)));
out.push(((n >> 16) & 0xFF) as u8);
if vals.len() > 2 { out.push(((n >> 8) & 0xFF) as u8); }
if vals.len() > 3 { out.push((n & 0xFF) as u8); }
}
Some(out)
}
fn uri_matches(template: &str, uri: &str) -> bool {
match template.find('{') {
Some(pos) => uri.starts_with(&template[..pos]),
None => template == uri,
}
}