use crate::jsonrpc::{ErrorObject, Outgoing, RequestId};
use crate::protocol;
use crate::protocol::{
CacheScope, CallToolResult, DiscoverResult, GetPromptResult, Prompt, PromptMessage,
PromptsListResult, ReadResourceResult, Resource, ResourceContents, ResourcesListResult,
ResultType, ServerInfo, Tool, ToolAnnotations, ToolsListResult,
};
use crate::zod::{Issues, Schema};
use serde::de::DeserializeOwned;
use serde_json::{Value, json};
use std::collections::BTreeMap;
use std::panic::AssertUnwindSafe;
use std::sync::Arc;
use std::time::Duration;
pub type ToolResult = Result<Value, String>;
type ToolFn = dyn Fn(&Value) -> Result<Value, ToolFailure> + Send + Sync;
type ResourceFn = dyn Fn() -> Result<ResourceContents, String> + Send + Sync;
type PromptFn = dyn Fn(&Value) -> Result<Vec<PromptMessage>, String> + Send + Sync;
static EMPTY_OBJECT: std::sync::OnceLock<Value> = std::sync::OnceLock::new();
fn empty_object() -> &'static Value {
EMPTY_OBJECT.get_or_init(|| json!({}))
}
pub const DEFAULT_LIST_TTL_MS: u64 = 60_000;
const RESOURCE_FAILURE: i64 = -32000;
enum Outcome<T> {
Done(T),
Panicked,
TimedOut,
}
fn run_with_timeout<T: Send + 'static>(
timeout: Duration,
f: impl FnOnce() -> T + Send + 'static,
) -> Outcome<T> {
let (tx, rx) = std::sync::mpsc::channel();
let spawned = std::thread::Builder::new()
.name("litemcp-handler".into())
.spawn(move || {
let _ = tx.send(std::panic::catch_unwind(AssertUnwindSafe(f)));
});
if spawned.is_err() {
return Outcome::Panicked;
}
match rx.recv_timeout(timeout) {
Ok(Ok(value)) => Outcome::Done(value),
Ok(Err(_)) => Outcome::Panicked,
Err(std::sync::mpsc::RecvTimeoutError::Timeout) => Outcome::TimedOut,
Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => Outcome::Panicked,
}
}
#[derive(Debug)]
enum ToolFailure {
Tool(String),
Protocol(i64, String),
}
#[derive(Clone)]
struct ToolDef {
description: Option<String>,
schema: Option<Arc<dyn Schema>>,
output_schema: Option<Arc<dyn Schema>>,
annotations: Option<ToolAnnotations>,
f: Arc<ToolFn>,
}
#[derive(Clone)]
struct ResourceDef {
name: String,
description: Option<String>,
f: Arc<ResourceFn>,
}
#[derive(Clone)]
struct PromptDef {
description: Option<String>,
f: Arc<PromptFn>,
}
#[derive(Debug, Clone, Default)]
pub struct McpHeaders {
pub method: Option<String>,
pub name: Option<String>,
}
#[derive(Clone)]
pub struct ServerCore {
name: String,
version: String,
tools: Arc<BTreeMap<String, ToolDef>>,
resources: Arc<BTreeMap<String, ResourceDef>>,
prompts: Arc<BTreeMap<String, PromptDef>>,
tool_timeout: Option<Duration>,
}
impl ServerCore {
pub fn dispatch(&self, body: &[u8]) -> Option<Vec<u8>> {
self.dispatch_with_headers(body, None)
}
pub fn dispatch_with_headers(
&self,
body: &[u8],
headers: Option<&McpHeaders>,
) -> Option<Vec<u8>> {
let incoming = match crate::jsonrpc::parse_incoming(body) {
Ok(incoming) => incoming,
Err(crate::jsonrpc::ParseFail::Syntax) => {
return Some(Outgoing::parse_error().to_vec());
}
Err(fail) => {
let id = fail.id().cloned();
let message = match fail {
crate::jsonrpc::ParseFail::Syntax => String::new(),
crate::jsonrpc::ParseFail::Invalid { message, .. } => message,
};
return Some(
Outgoing::invalid_request(id, format!("requête invalide : {message}")).to_vec(),
);
}
};
if incoming.is_notification() {
return None;
}
let id = incoming
.id
.as_ref()
.map(RequestId::to_value)
.unwrap_or(Value::Null);
let params = incoming.params.as_ref().unwrap_or_else(|| empty_object());
if let Some(headers) = headers {
if let Some(header_method) = &headers.method {
if !header_method.eq_ignore_ascii_case(&incoming.method) {
return Some(
Outgoing::failure(
id,
ErrorObject::with_data(
ErrorObject::HEADER_MISMATCH,
"en-tête Mcp-Method incohérent avec le corps de la requête",
json!({
"header": header_method,
"body": incoming.method,
}),
),
)
.to_vec(),
);
}
}
if let Some(header_name) = &headers.name {
let body_name = params.get("name").and_then(Value::as_str).unwrap_or("");
if !body_name.is_empty() && !header_name.eq_ignore_ascii_case(body_name) {
return Some(
Outgoing::failure(
id,
ErrorObject::with_data(
ErrorObject::HEADER_MISMATCH,
"en-tête Mcp-Name incohérent avec le corps de la requête",
json!({
"header": header_name,
"body": body_name,
}),
),
)
.to_vec(),
);
}
}
}
if let Some(err) = self.check_protocol_version(params) {
return Some(Outgoing::failure(id, err).to_vec());
}
let outgoing = match incoming.method.as_str() {
protocol::method::DISCOVER => Outgoing::success(
id,
serde_json::to_value(self.discover()).unwrap_or(Value::Null),
),
protocol::method::INITIALIZE => Outgoing::success(
id,
json!({
"resultType": "complete",
"protocolVersion": protocol::PROTOCOL_VERSION,
"capabilities": self.capabilities(),
"serverInfo": ServerInfo {
name: self.name.clone(),
version: self.version.clone(),
},
"_meta": {
protocol::META_SERVER_INFO: ServerInfo {
name: self.name.clone(),
version: self.version.clone(),
}
}
}),
),
protocol::method::PING => Outgoing::success(
id,
json!({
"resultType": "complete",
"_meta": self.meta(),
}),
),
protocol::method::TOOLS_LIST => Outgoing::success(
id,
serde_json::to_value(self.list_tools()).unwrap_or_else(|_| json!({ "tools": [] })),
),
protocol::method::TOOLS_CALL => self.call_tool(&id, params),
protocol::method::RESOURCES_LIST => Outgoing::success(
id,
serde_json::to_value(self.list_resources())
.unwrap_or_else(|_| json!({ "resources": [] })),
),
protocol::method::RESOURCES_READ => self.read_resource(&id, params),
protocol::method::PROMPTS_LIST => Outgoing::success(
id,
serde_json::to_value(self.list_prompts())
.unwrap_or_else(|_| json!({ "prompts": [] })),
),
protocol::method::PROMPTS_GET => self.get_prompt(&id, params),
other => Outgoing::failure(
id,
ErrorObject::new(
ErrorObject::METHOD_NOT_FOUND,
format!("Méthode inconnue : {other}"),
),
),
};
Some(outgoing.to_vec())
}
fn check_protocol_version(&self, params: &Value) -> Option<ErrorObject> {
let requested = params
.get("_meta")
.and_then(|meta| meta.get(protocol::META_PROTOCOL_VERSION))
.and_then(Value::as_str)?;
if protocol::SUPPORTED_VERSIONS.contains(&requested) {
return None;
}
Some(ErrorObject::with_data(
ErrorObject::UNSUPPORTED_PROTOCOL_VERSION,
format!("Version de protocole non supportée : {requested}"),
json!({
"supported": protocol::SUPPORTED_VERSIONS,
"requested": requested,
}),
))
}
fn meta(&self) -> Value {
json!({
protocol::META_SERVER_INFO: ServerInfo {
name: self.name.clone(),
version: self.version.clone(),
}
})
}
fn capabilities(&self) -> Value {
let mut capabilities = json!({ "tools": {} });
if !self.resources.is_empty() {
capabilities["resources"] = json!({});
}
if !self.prompts.is_empty() {
capabilities["prompts"] = json!({});
}
capabilities
}
fn discover(&self) -> DiscoverResult {
DiscoverResult {
result_type: ResultType::Complete,
supported_versions: protocol::SUPPORTED_VERSIONS.to_vec(),
capabilities: self.capabilities(),
meta: json!({
protocol::META_SERVER_INFO: ServerInfo {
name: self.name.clone(),
version: self.version.clone(),
}
}),
}
}
fn list_tools(&self) -> ToolsListResult {
let tools = self
.tools
.iter()
.map(|(name, def)| Tool {
name: name.clone(),
description: def.description.clone(),
annotations: def.annotations.clone(),
input_schema: Some(
def.schema
.as_ref()
.map(|schema| schema.json())
.unwrap_or_else(|| json!({ "type": "object" })),
),
output_schema: def.output_schema.as_ref().map(|schema| schema.json()),
})
.collect();
ToolsListResult {
result_type: ResultType::Complete,
ttl_ms: DEFAULT_LIST_TTL_MS,
cache_scope: CacheScope::Private,
tools,
meta: self.meta(),
}
}
fn list_resources(&self) -> ResourcesListResult {
let resources = self
.resources
.iter()
.map(|(uri, def)| Resource {
uri: uri.clone(),
name: def.name.clone(),
description: def.description.clone(),
mime_type: None,
})
.collect();
ResourcesListResult {
result_type: ResultType::Complete,
ttl_ms: DEFAULT_LIST_TTL_MS,
cache_scope: CacheScope::Private,
resources,
meta: Some(self.meta()),
}
}
fn list_prompts(&self) -> PromptsListResult {
let prompts = self
.prompts
.iter()
.map(|(name, def)| Prompt {
name: name.clone(),
description: def.description.clone(),
})
.collect();
PromptsListResult {
result_type: ResultType::Complete,
ttl_ms: DEFAULT_LIST_TTL_MS,
cache_scope: CacheScope::Private,
prompts,
meta: Some(self.meta()),
}
}
fn read_resource(&self, id: &Value, params: &Value) -> Outgoing {
let Some(uri) = params.get("uri").and_then(Value::as_str) else {
return Outgoing::failure(
id.clone(),
ErrorObject::new(ErrorObject::INVALID_PARAMS, "params.uri manquant"),
);
};
let Some(def) = self.resources.get(uri) else {
return Outgoing::failure(
id.clone(),
ErrorObject::new(
ErrorObject::INVALID_PARAMS,
format!("Ressource inconnue : {uri}"),
),
);
};
let handler = Arc::clone(&def.f);
let contents = match self.run_bounded(move || handler()) {
Ok(contents) => contents,
Err(failure) => return Outgoing::failure(id.clone(), failure),
};
match contents {
Ok(contents) => Outgoing::success(
id.clone(),
serde_json::to_value(ReadResourceResult {
result_type: ResultType::Complete,
ttl_ms: DEFAULT_LIST_TTL_MS,
cache_scope: CacheScope::Private,
contents: vec![contents],
meta: Some(self.meta()),
})
.unwrap_or(Value::Null),
),
Err(message) => Outgoing::failure(
id.clone(),
ErrorObject::new(RESOURCE_FAILURE, format!("Lecture impossible : {message}")),
),
}
}
fn get_prompt(&self, id: &Value, params: &Value) -> Outgoing {
let Some(name) = params.get("name").and_then(Value::as_str) else {
return Outgoing::failure(
id.clone(),
ErrorObject::new(ErrorObject::INVALID_PARAMS, "params.name manquant"),
);
};
let Some(def) = self.prompts.get(name) else {
return Outgoing::failure(
id.clone(),
ErrorObject::new(
ErrorObject::INVALID_PARAMS,
format!("Prompt inconnu : {name}"),
),
);
};
let arguments = match params.get("arguments") {
None => empty_object(),
Some(value) if value.is_object() => value,
Some(_) => {
return Outgoing::failure(
id.clone(),
ErrorObject::new(
ErrorObject::INVALID_PARAMS,
"params.arguments doit être un objet",
),
);
}
};
let handler = Arc::clone(&def.f);
let arguments = arguments.clone();
let messages = match self.run_bounded(move || handler(&arguments)) {
Ok(messages) => messages,
Err(failure) => return Outgoing::failure(id.clone(), failure),
};
match messages {
Ok(messages) => Outgoing::success(
id.clone(),
serde_json::to_value(GetPromptResult {
result_type: ResultType::Complete,
description: def.description.clone(),
messages,
meta: Some(self.meta()),
})
.unwrap_or(Value::Null),
),
Err(message) => Outgoing::failure(
id.clone(),
ErrorObject::new(RESOURCE_FAILURE, format!("Prompt impossible : {message}")),
),
}
}
fn run_bounded<T: Send + 'static>(
&self,
f: impl FnOnce() -> Result<T, String> + Send + 'static,
) -> Result<Result<T, String>, ErrorObject> {
match self.tool_timeout {
None => Ok(f()),
Some(timeout) => match run_with_timeout(timeout, f) {
Outcome::Done(result) => Ok(result),
Outcome::Panicked => Err(ErrorObject::new(
ErrorObject::INTERNAL_ERROR,
"erreur interne (panique du handler)",
)),
Outcome::TimedOut => Err(ErrorObject::new(
RESOURCE_FAILURE,
format!("délai d'exécution dépassé ({} ms)", timeout.as_millis()),
)),
},
}
}
fn call_tool(&self, id: &Value, params: &Value) -> Outgoing {
let Some(name) = params.get("name").and_then(Value::as_str) else {
return Outgoing::failure(
id.clone(),
ErrorObject::new(ErrorObject::INVALID_PARAMS, "params.name manquant"),
);
};
let Some(tool) = self.tools.get(name) else {
return Outgoing::failure(
id.clone(),
ErrorObject::new(
ErrorObject::INVALID_PARAMS,
format!("Outil inconnu : {name}"),
),
);
};
let arguments = match params.get("arguments") {
None => empty_object(),
Some(value) if value.is_object() => value,
Some(_) => {
return Outgoing::failure(
id.clone(),
ErrorObject::new(
ErrorObject::INVALID_PARAMS,
"params.arguments doit être un objet",
),
);
}
};
if let Some(schema) = &tool.schema {
let mut issues = Issues::new();
schema.validate("arguments", arguments, &mut issues);
if !issues.is_empty() {
let rendered = issues.render();
return Outgoing::failure(
id.clone(),
ErrorObject::with_data(
ErrorObject::INVALID_PARAMS,
format!("Arguments invalides : {}", rendered.join(" ; ")),
json!({ "issues": rendered }),
),
);
}
}
let execution = match self.tool_timeout {
None => (tool.f)(arguments),
Some(timeout) => {
let owned = arguments.clone();
let handler = Arc::clone(&tool.f);
match run_with_timeout(timeout, move || handler(&owned)) {
Outcome::Done(result) => result,
Outcome::Panicked => Err(ToolFailure::Protocol(
ErrorObject::INTERNAL_ERROR,
"erreur interne (panique du handler)".to_string(),
)),
Outcome::TimedOut => Err(ToolFailure::Tool(format!(
"l'exécution de l'outil a dépassé le délai imparti ({} ms)",
timeout.as_millis()
))),
}
}
};
let mut result = match execution {
Ok(value) => {
let mut call = CallToolResult::from(Ok::<Value, String>(value));
if let (Some(output), Some(structured)) =
(&tool.output_schema, &call.structured_content)
{
let mut issues = Issues::new();
output.validate("output", structured, &mut issues);
if !issues.is_empty() {
let rendered = issues.render().join(" ; ");
call = CallToolResult::from(Err::<Value, String>(format!(
"sortie non conforme au outputSchema : {rendered}"
)));
}
}
call
}
Err(ToolFailure::Tool(message)) => CallToolResult::from(Err::<Value, String>(message)),
Err(ToolFailure::Protocol(code, message)) => {
return Outgoing::failure(id.clone(), ErrorObject::new(code, message));
}
};
result.meta = Some(self.meta());
Outgoing::success(
id.clone(),
serde_json::to_value(result).unwrap_or(Value::Null),
)
}
}
#[derive(Clone)]
pub struct ServerBuilder {
name: String,
version: String,
route: String,
allow_origins: Vec<String>,
read_timeout: Duration,
max_connections: usize,
tools: BTreeMap<String, ToolDef>,
connection_ttl: Duration,
tool_timeout: Option<Duration>,
annotations: BTreeMap<String, ToolAnnotations>,
output_schemas: BTreeMap<String, Arc<dyn Schema>>,
resources: BTreeMap<String, ResourceDef>,
prompts: BTreeMap<String, PromptDef>,
}
impl ServerBuilder {
#[must_use]
pub fn new(name: &str, version: &str) -> Self {
Self {
name: name.to_string(),
version: version.to_string(),
route: "/mcp".to_string(),
allow_origins: Vec::new(),
read_timeout: crate::http::READ_TIMEOUT,
max_connections: crate::http::MAX_CONNECTIONS,
tools: BTreeMap::new(),
annotations: BTreeMap::new(),
output_schemas: BTreeMap::new(),
resources: BTreeMap::new(),
prompts: BTreeMap::new(),
connection_ttl: crate::http::CONNECTION_TTL,
tool_timeout: None,
}
}
#[must_use]
pub fn connection_ttl(mut self, ttl: Duration) -> Self {
self.connection_ttl = ttl;
self
}
#[must_use]
pub fn read_timeout(mut self, timeout: Duration) -> Self {
self.read_timeout = timeout;
self
}
#[must_use]
pub fn max_connections(mut self, max: usize) -> Self {
self.max_connections = max;
self
}
#[must_use]
pub fn route(mut self, route: &str) -> Self {
self.route = route.to_string();
self
}
#[must_use]
pub fn allow_origins(mut self, origins: &[&str]) -> Self {
self.allow_origins = origins.iter().map(|o| (*o).to_string()).collect();
self
}
#[must_use]
pub fn tool_timeout(mut self, timeout: Duration) -> Self {
self.tool_timeout = Some(timeout);
self
}
#[must_use]
pub fn output_schema<S: Schema + 'static>(mut self, tool: &str, schema: S) -> Self {
if let Some(def) = self.tools.get_mut(tool) {
def.output_schema = Some(Arc::new(schema));
} else {
self.output_schemas
.insert(tool.to_string(), Arc::new(schema));
}
self
}
#[must_use]
pub fn resource<F>(mut self, uri: &str, name: &str, description: &str, f: F) -> Self
where
F: Fn() -> Result<ResourceContents, String> + Send + Sync + 'static,
{
self.resources.insert(
uri.to_string(),
ResourceDef {
name: name.to_string(),
description: Some(description.to_string()),
f: Arc::new(f),
},
);
self
}
#[must_use]
pub fn prompt<F>(mut self, name: &str, description: &str, f: F) -> Self
where
F: Fn(&Value) -> Result<Vec<PromptMessage>, String> + Send + Sync + 'static,
{
self.prompts.insert(
name.to_string(),
PromptDef {
description: Some(description.to_string()),
f: Arc::new(f),
},
);
self
}
#[must_use]
pub fn with_tool(self, register: impl FnOnce(Self) -> Self) -> Self {
register(self)
}
pub fn tool<F>(mut self, name: &str, description: &str, f: F) -> Self
where
F: Fn(&Value) -> ToolResult + Send + Sync + 'static,
{
let wrapped = move |args: &Value| f(args).map_err(ToolFailure::Tool);
self.tools.insert(
name.to_string(),
ToolDef {
description: Some(description.to_string()),
schema: None,
output_schema: self.output_schemas.remove(name),
annotations: self.annotations.remove(name),
f: Arc::new(wrapped),
},
);
self
}
pub fn tool_typed<S, F>(mut self, name: &str, description: &str, schema: S, f: F) -> Self
where
S: Schema + 'static,
F: Fn(&Value) -> ToolResult + Send + Sync + 'static,
{
let wrapped = move |args: &Value| f(args).map_err(ToolFailure::Tool);
self.tools.insert(
name.to_string(),
ToolDef {
description: Some(description.to_string()),
schema: Some(Arc::new(schema)),
output_schema: self.output_schemas.remove(name),
annotations: self.annotations.remove(name),
f: Arc::new(wrapped),
},
);
self
}
pub fn tool_parse<S, T, F>(mut self, name: &str, description: &str, schema: S, f: F) -> Self
where
S: Schema + 'static,
T: DeserializeOwned,
F: Fn(T) -> ToolResult + Send + Sync + 'static,
{
let handler = Arc::new(move |value: &Value| match serde_json::to_vec(value) {
Err(error) => Err(ToolFailure::Protocol(
ErrorObject::INVALID_PARAMS,
format!("arguments illisibles : {error}"),
)),
Ok(buffer) => match serde_json::from_slice::<T>(&buffer) {
Ok(parsed) => f(parsed).map_err(ToolFailure::Tool),
Err(error) => Err(ToolFailure::Protocol(
ErrorObject::INVALID_PARAMS,
format!("arguments invalides pour le type attendu : {error}"),
)),
},
});
self.tools.insert(
name.to_string(),
ToolDef {
description: Some(description.to_string()),
schema: Some(Arc::new(schema)),
output_schema: self.output_schemas.remove(name),
annotations: self.annotations.remove(name),
f: handler,
},
);
self
}
#[must_use]
pub fn annotate(mut self, tool: &str, annotations: ToolAnnotations) -> Self {
if let Some(def) = self.tools.get_mut(tool) {
def.annotations = Some(annotations);
} else {
self.annotations.insert(tool.to_string(), annotations);
}
self
}
#[must_use]
pub fn core(&self) -> ServerCore {
let mut tools = self.tools.clone();
for (name, def) in tools.iter_mut() {
if let Some(annotations) = self.annotations.get(name) {
def.annotations = Some(annotations.clone());
}
if let Some(output) = self.output_schemas.get(name) {
def.output_schema = Some(Arc::clone(output));
}
}
ServerCore {
name: self.name.clone(),
version: self.version.clone(),
tools: Arc::new(tools),
resources: Arc::new(self.resources.clone()),
prompts: Arc::new(self.prompts.clone()),
tool_timeout: self.tool_timeout,
}
}
pub fn bind(self, addr: &str) -> std::io::Result<crate::http::BoundServer> {
if !self.route.starts_with('/') || self.route.contains('?') {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!(
"route invalide {:?} : elle doit commencer par '/' et ne pas contenir de '?'",
self.route
),
));
}
if self.route == crate::http::HEALTH_ROUTE {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
format!(
"route invalide {:?} : elle masque la sonde de vie {}",
self.route,
crate::http::HEALTH_ROUTE
),
));
}
let listener = std::net::TcpListener::bind(addr)?;
let mut tools = self.tools;
for (name, def) in tools.iter_mut() {
if let Some(annotations) = self.annotations.get(name) {
def.annotations = Some(annotations.clone());
}
if let Some(output) = self.output_schemas.get(name) {
def.output_schema = Some(Arc::clone(output));
}
}
let core = Arc::new(ServerCore {
name: self.name,
version: self.version,
tools: Arc::new(tools),
resources: Arc::new(self.resources),
prompts: Arc::new(self.prompts),
tool_timeout: self.tool_timeout,
});
Ok(crate::http::BoundServer {
listener,
core,
route: Arc::new(self.route),
allow_origins: Arc::new(self.allow_origins),
read_timeout: self.read_timeout,
max_connections: self.max_connections,
connection_ttl: self.connection_ttl,
live_connections: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
refusals_in_flight: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
shutdown: Arc::new(std::sync::atomic::AtomicBool::new(false)),
})
}
}