use std::num::NonZeroUsize;
use std::sync::Arc;
use crate::adapters::secondary::idempotency::InMemoryIdempotencyStore;
use crate::application::idempotency::IdempotencyStore;
use crate::application::mcp::input_required::StateKey;
use crate::application::registry::store::RegistryStore;
use crate::application::server::context::RequestContext;
use crate::application::server::live::{Handlers, Live, ServerHandle, ToolEntry};
use crate::domain::idempotency::IdempotencyError;
use crate::domain::ir::{IrErrorCode, JsonObject};
use crate::domain::registry::{
Primitive, PrimitiveId, PrimitiveKind, Prompt, Registered, Resource, ResourceTemplate, Tool,
};
pub type ToolHandler = Arc<dyn Fn(RequestContext, JsonObject) -> BoxFuture + Send + Sync>;
pub type BoxFuture = std::pin::Pin<Box<dyn Future<Output = JsonObject> + Send>>;
pub type BoxReadFuture =
std::pin::Pin<Box<dyn Future<Output = Result<JsonObject, DispatchError>> + Send>>;
pub type ResourceHandler = Arc<dyn Fn(RequestContext, String) -> BoxReadFuture + Send + Sync>;
pub type PromptHandler = Arc<dyn Fn(RequestContext, JsonObject) -> BoxReadFuture + Send + Sync>;
pub struct Server {
name: String,
version: String,
live: Arc<Live>,
idempotency: Arc<dyn IdempotencyStore<Result = JsonObject>>,
state: Arc<crate::application::server::context::StateMap>,
state_key: Option<StateKey>,
}
impl core::fmt::Debug for Server {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("Server")
.field("name", &self.name)
.field("version", &self.version)
.field("tool_count", &self.tool_count())
.finish_non_exhaustive()
}
}
impl Server {
#[cfg(feature = "stdio")]
pub async fn serve_stdio(
self,
) -> Result<
crate::adapters::primary::stdio::ServeReport,
crate::adapters::primary::stdio::ServeError,
> {
let mut loop_ = crate::application::server::serve::StdioBuilder::new(self).build();
loop_.serve().await
}
#[cfg(feature = "http")]
pub async fn serve_http(
self,
) -> Result<
crate::adapters::primary::transport::http::ServeReport,
crate::adapters::primary::transport::http::ServeError,
> {
crate::adapters::primary::transport::http::HttpBuilder::new(self)
.serve()
.await
}
#[must_use]
pub fn builder(name: impl Into<String>) -> ServerBuilder {
ServerBuilder {
name: name.into(),
version: env!("CARGO_PKG_VERSION").to_string(),
registry: RegistryStore::new(),
tools: Vec::new(),
resource_reads: Vec::new(),
prompt_gets: Vec::new(),
idempotency: None,
state: Vec::new(),
state_key: None,
}
}
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
#[must_use]
pub fn version(&self) -> &str {
&self.version
}
#[must_use]
pub fn tool_count(&self) -> usize {
self.live.tool_count()
}
#[must_use]
pub const fn request_state_key(&self) -> Option<&StateKey> {
self.state_key.as_ref()
}
pub fn ensure_request_state_key(&mut self, key: impl FnOnce() -> Option<StateKey>) {
if self.state_key.is_none() {
self.state_key = key();
}
}
#[must_use]
pub fn handle(&self) -> ServerHandle {
ServerHandle::new(Arc::clone(&self.live))
}
#[must_use]
pub const fn idempotency(&self) -> &Arc<dyn IdempotencyStore<Result = JsonObject>> {
&self.idempotency
}
pub fn with_state<T: Send + Sync + 'static>(&mut self, value: T) {
Arc::make_mut(&mut self.state).insert(
std::any::TypeId::of::<T>(),
Arc::new(value) as Arc<dyn std::any::Any + Send + Sync>,
);
}
pub fn add_tool(&mut self, tool: Tool, handler: ToolHandler) {
let name = tool.name.clone();
let _published = self
.live
.add(Registered::from(Primitive::from(tool)), |table| {
table.tools.retain(|entry| entry.name != name);
table.tools.push(ToolEntry { name, handler });
});
}
#[must_use]
pub fn state(&self) -> &Arc<crate::application::server::context::StateMap> {
&self.state
}
#[must_use]
pub fn registry(&self) -> &RegistryStore {
self.live.registry()
}
#[must_use]
pub fn handler(&self, name: &str) -> Option<ToolHandler> {
self.live.tool(name)
}
#[must_use]
pub fn reader(&self, uri: &str) -> Option<ResourceHandler> {
self.live.reader(uri)
}
#[must_use]
pub fn getter(&self, name: &str) -> Option<PromptHandler> {
self.live.getter(name)
}
pub async fn call(
&self,
ctx: &RequestContext,
name: &str,
args: JsonObject,
) -> Result<JsonObject, DispatchError> {
let handler = self.handler(name).ok_or(DispatchError::UnknownTool)?;
if self
.registry()
.get_visible(
ctx.peer().auth(),
&ctx.peer().profile(),
PrimitiveKind::Tool,
name,
)
.is_none()
{
return Err(DispatchError::UnknownTool);
}
Ok(handler(ctx.clone(), args).await)
}
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum DispatchError {
#[error("no such tool")]
UnknownTool,
#[error("the session did not negotiate what this tool requires")]
NotCallable,
#[error("no such method")]
UnknownMethod,
#[error("no such resource")]
UnknownResource,
#[error("no such prompt")]
UnknownPrompt,
#[error("the handler refused: {0}")]
HandlerFailed(String),
#[error("the call parameters are not a tool call")]
MalformedParams,
#[error("the initialize parameters are not a handshake")]
MalformedInitialize,
#[error("the session is on a revision this build cannot serve")]
UnnegotiableVersion,
#[error("the request `_meta` is missing a required key")]
MalformedRequestMeta,
#[error("this server does not implement the requested protocol version")]
UnsupportedVersion {
requested: String,
supported: Vec<String>,
},
#[error("request {} is already in flight on this session", .0)]
DuplicateInFlight(String),
#[error("the request needs client capabilities it did not declare")]
MissingClientCapability {
required: Vec<&'static str>,
},
#[error("requestState failed its integrity check")]
TamperedRequestState,
#[error("recorded results are unavailable: {0}")]
IdempotencyUnavailable(String),
}
impl DispatchError {
#[must_use]
pub const fn jsonrpc_code(&self) -> IrErrorCode {
match self {
Self::UnknownMethod => IrErrorCode::MethodNotFound,
Self::MalformedParams
| Self::MalformedInitialize
| Self::UnknownTool
| Self::UnknownResource
| Self::UnknownPrompt
| Self::MalformedRequestMeta
| Self::TamperedRequestState => IrErrorCode::InvalidParams,
Self::MissingClientCapability { .. } => IrErrorCode::MISSING_REQUIRED_CLIENT_CAPABILITY,
Self::DuplicateInFlight(_) => IrErrorCode::InvalidRequest,
Self::NotCallable | Self::UnnegotiableVersion | Self::IdempotencyUnavailable(_) => {
IrErrorCode::InternalError
}
Self::UnsupportedVersion { .. } => IrErrorCode::UNSUPPORTED_PROTOCOL_VERSION,
Self::HandlerFailed(_) => IrErrorCode::InternalError,
}
}
#[must_use]
pub fn wire_data(&self) -> Option<JsonObject> {
let data = match self {
Self::UnsupportedVersion {
requested,
supported,
} => serde_json::json!({ "requested": requested, "supported": supported }),
Self::MissingClientCapability { required } => {
let capabilities: serde_json::Map<String, serde_json::Value> = required
.iter()
.map(|name| {
(
(*name).to_owned(),
serde_json::Value::Object(serde_json::Map::new()),
)
})
.collect();
serde_json::json!({ "requiredCapabilities": capabilities })
}
_ => return None,
};
JsonObject::from_value(data).ok()
}
#[must_use]
pub fn as_wire(&self) -> (IrErrorCode, String) {
(self.jsonrpc_code(), self.to_string())
}
}
impl From<crate::application::server::result::OutcomeError> for DispatchError {
fn from(error: crate::application::server::result::OutcomeError) -> Self {
Self::HandlerFailed(error.to_string())
}
}
impl From<IdempotencyError> for DispatchError {
fn from(error: IdempotencyError) -> Self {
Self::IdempotencyUnavailable(error.to_string())
}
}
pub struct ServerBuilder {
name: String,
version: String,
registry: RegistryStore,
tools: Vec<ToolEntry>,
resource_reads: Vec<(PrimitiveId, ResourceHandler)>,
prompt_gets: Vec<(PrimitiveId, PromptHandler)>,
idempotency: Option<Arc<dyn IdempotencyStore<Result = JsonObject>>>,
state: Vec<(std::any::TypeId, Arc<dyn std::any::Any + Send + Sync>)>,
state_key: Option<StateKey>,
}
impl core::fmt::Debug for ServerBuilder {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("ServerBuilder")
.field("name", &self.name)
.field("version", &self.version)
.field("tool_count", &self.tools.len())
.field("idempotency_supplied", &self.idempotency.is_some())
.finish_non_exhaustive()
}
}
impl ServerBuilder {
#[must_use]
pub const fn request_state_key(mut self, key: StateKey) -> Self {
self.state_key = Some(key);
self
}
#[must_use]
pub fn version(mut self, version: impl Into<String>) -> Self {
self.version = version.into();
self
}
#[must_use]
pub fn idempotency(mut self, store: Arc<dyn IdempotencyStore<Result = JsonObject>>) -> Self {
self.idempotency = Some(store);
self
}
#[must_use]
pub fn with_state<T: Send + Sync + 'static>(mut self, value: T) -> Self {
self.state.push((
std::any::TypeId::of::<T>(),
Arc::new(value) as Arc<dyn std::any::Any + Send + Sync>,
));
self
}
#[must_use]
pub fn tool(self, tool: Tool, handler: ToolHandler) -> Self {
self.restricted_tool(Registered::from(Primitive::from(tool)), handler)
}
#[must_use]
pub fn restricted_tool(mut self, registered: Registered, handler: ToolHandler) -> Self {
if let Primitive::Tool(tool) = registered.primitive() {
self.tools.push(ToolEntry {
name: tool.name.clone(),
handler,
});
}
self.registry.register(registered);
self
}
pub fn simple_tool(
self,
name: impl Into<PrimitiveId>,
input_schema: &str,
handler: ToolHandler,
) -> Result<Self, crate::domain::ir::JsonObjectError> {
Ok(self.tool(Tool::new(name, input_schema)?, handler))
}
#[must_use]
pub fn resource_with(self, resource: Resource, reader: ResourceHandler) -> Self {
let mut builder = self;
builder.resource_reads.push((resource.uri.clone(), reader));
builder
.registry
.register(Registered::from(Primitive::from(resource)));
builder
}
#[must_use]
pub fn resource_template(self, template: ResourceTemplate) -> Self {
self.registry
.register(Registered::from(Primitive::from(template)));
self
}
#[must_use]
pub fn prompt_with(self, prompt: Prompt, getter: PromptHandler) -> Self {
let mut builder = self;
builder.prompt_gets.push((prompt.name.clone(), getter));
builder
.registry
.register(Registered::from(Primitive::from(prompt)));
builder
}
#[must_use]
pub fn build(self) -> Server {
let idempotency = self.idempotency.unwrap_or_else(|| {
const DEFAULT_KEYS: NonZeroUsize = match NonZeroUsize::new(4096) {
Some(value) => value,
None => NonZeroUsize::MIN,
};
Arc::new(InMemoryIdempotencyStore::with_capacity(DEFAULT_KEYS))
});
Server {
name: self.name,
version: self.version,
live: Arc::new(Live::new(
self.registry,
Handlers {
tools: self.tools,
resources: self.resource_reads,
prompts: self.prompt_gets,
},
)),
idempotency,
state: Arc::new(
self.state
.into_iter()
.collect::<std::collections::HashMap<_, _>>(),
),
state_key: self.state_key,
}
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use super::{DispatchError, Server};
use crate::application::server::context::{PeerSession, RequestContext};
use crate::application::session::SinkError;
use crate::domain::feature::FeatureProfile;
use crate::domain::ir::{JsonObject, JsonObjectError};
use crate::domain::registry::PrimitiveKind;
use crate::domain::session::{CancellationToken, SessionId};
use crate::domain::{ProtocolVersion, RequestId, TransportKind};
#[derive(Debug)]
enum Failure {
Message(String),
Fixture(JsonObjectError),
}
impl core::fmt::Display for Failure {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::Message(text) => f.write_str(text),
Self::Fixture(error) => write!(f, "bad fixture: {error}"),
}
}
}
impl std::error::Error for Failure {}
impl From<String> for Failure {
fn from(text: String) -> Self {
Self::Message(text)
}
}
impl From<&str> for Failure {
fn from(text: &str) -> Self {
Self::Message(text.to_string())
}
}
impl From<JsonObjectError> for Failure {
fn from(error: JsonObjectError) -> Self {
Self::Fixture(error)
}
}
type TestError = Failure;
fn answering(answer: &str) -> Result<super::ToolHandler, TestError> {
let answer = JsonObject::new(answer)?;
Ok(Arc::new(move |_ctx: RequestContext, _args: JsonObject| {
let answer = answer.clone();
Box::pin(async move { answer })
}))
}
struct Inert;
impl crate::application::session::FrameSink for Inert {
fn send(&self, _frame: &str) -> impl Future<Output = Result<(), SinkError>> + Send {
std::future::ready(Ok(()))
}
}
impl crate::application::server::context::CallbackExchange for Inert {
fn exchange(
&self,
_method: &str,
_params: &JsonObject,
) -> impl Future<Output = Result<JsonObject, crate::application::server::ExchangeError>> + Send
{
std::future::ready(Err(crate::application::server::ExchangeError::new(
"inert in these tests",
)))
}
}
fn peer_at(profile: FeatureProfile) -> PeerSession {
PeerSession::new(
ProtocolVersion::V2025_06_18,
profile,
SessionId::new(1),
TransportKind::Stdio,
Arc::new(Inert),
Arc::new(Inert),
)
}
fn ctx_for(profile: FeatureProfile) -> RequestContext {
RequestContext::for_request(
&peer_at(profile),
RequestId::number(1),
CancellationToken::new(),
None,
)
}
#[test]
fn a_registered_tool_is_in_the_registry_immediately() -> Result<(), TestError> {
let server = Server::builder("demo")
.simple_tool("echo", r#"{"type":"object"}"#, answering("{}")?)?
.build();
assert_eq!(server.tool_count(), 1);
assert_eq!(server.registry().snapshot().len(PrimitiveKind::Tool), 1);
Ok(())
}
#[test]
fn a_tool_added_after_the_build_lists_and_runs() -> Result<(), TestError> {
use crate::domain::registry::{Primitive, Tool};
let mut server = Server::builder("demo")
.simple_tool("echo", r#"{"type":"object"}"#, answering("{}")?)?
.build();
let raw = Tool::new(
"json_schema_2020_12",
r#"{"$schema":"https://json-schema.org/draft/2020-12/schema","type":"object","$defs":{"address":{"$anchor":"addressDef","type":"string"}},"if":{"properties":{"a":{"const":1}}},"then":{"required":["b"]}}"#,
)?;
server.add_tool(raw, answering(r#"{"ok":true}"#)?);
assert_eq!(server.tool_count(), 2);
assert_eq!(server.registry().snapshot().len(PrimitiveKind::Tool), 2);
assert!(server.handler("json_schema_2020_12").is_some());
let snapshot = server.registry().snapshot();
let advertised = snapshot
.list(PrimitiveKind::Tool)
.iter()
.find(|tool| tool.id().as_str() == "json_schema_2020_12")
.ok_or("the added tool is not in the registry")?;
let Primitive::Tool(advertised) = advertised.primitive() else {
return Err("a tool listing held something that is not a tool".into());
};
let text = advertised.input_schema.as_str();
for keyword in ["$schema", "$defs", "$anchor", "if", "then"] {
if !text.contains(keyword) {
return Err(format!("{keyword} did not reach the listing").into());
}
}
Ok(())
}
#[test]
fn a_malformed_input_schema_is_rejected_at_build_time() -> Result<(), TestError> {
let handler = answering("{}")?;
let built = Server::builder("demo").simple_tool("bad", "{not json", handler);
match built {
Ok(_) => Err(Failure::from(
"a malformed input schema was accepted at build time",
)),
Err(_) => Ok(()),
}
}
#[test]
fn a_second_tool_of_the_same_name_replaces_the_first() -> Result<(), TestError> {
let server = Server::builder("demo")
.simple_tool(
"echo",
r#"{"type":"object"}"#,
answering(r#"{"first":true}"#)?,
)?
.simple_tool(
"echo",
r#"{"type":"object"}"#,
answering(r#"{"second":true}"#)?,
)?
.build();
assert_eq!(server.tool_count(), 2, "both handlers are kept");
assert_eq!(
server.registry().snapshot().len(PrimitiveKind::Tool),
1,
"but the registry holds one name"
);
Ok(())
}
#[tokio::test]
async fn an_unknown_tool_is_refused() -> Result<(), TestError> {
let server = Server::builder("demo")
.simple_tool("echo", r#"{"type":"object"}"#, answering("{}")?)?
.build();
let ctx = ctx_for(FeatureProfile::all());
match server.call(&ctx, "nope", JsonObject::new("{}")?).await {
Ok(answer) => Err(Failure::from(format!("expected a refusal, got {answer}"))),
Err(DispatchError::UnknownTool) => Ok(()),
Err(other) => Err(Failure::from(format!(
"refused for the wrong reason: {other}"
))),
}
}
#[tokio::test]
async fn a_tool_needing_an_absent_feature_is_not_run() -> Result<(), TestError> {
use crate::domain::registry::Tool;
let structured = Tool::new("echo", r#"{"type":"object"}"#)?
.with_output_schema(r#"{"type":"object"}"#)?;
let server = Server::builder("demo")
.tool(structured, answering(r#"{"ran":true}"#)?)
.build();
let ctx = ctx_for(FeatureProfile::empty());
match server.call(&ctx, "echo", JsonObject::new("{}")?).await {
Ok(answer) => Err(Failure::from(format!(
"a tool the session cannot call returned {answer}"
))),
Err(DispatchError::UnknownTool) => Ok(()),
Err(other) => Err(Failure::from(format!(
"refused for the wrong reason: {other}"
))),
}
}
}