use async_trait::async_trait;
use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
use pmcp::types::capabilities::ElicitationCapabilities;
use pmcp::types::*;
use pmcp::{Server, ToolHandler};
use serde_json::{json, Value};
use std::hint::black_box;
struct BenchmarkTool;
#[async_trait]
impl ToolHandler for BenchmarkTool {
async fn handle(&self, args: Value, _extra: pmcp::RequestHandlerExtra) -> pmcp::Result<Value> {
let input = args
.get("input")
.and_then(|v| v.as_str())
.unwrap_or("default");
Ok(json!({
"result": format!("Processed: {}", input),
"length": input.len(),
"timestamp": chrono::Utc::now().timestamp()
}))
}
}
struct ComplexBenchmarkTool;
#[async_trait]
impl ToolHandler for ComplexBenchmarkTool {
async fn handle(&self, args: Value, _extra: pmcp::RequestHandlerExtra) -> pmcp::Result<Value> {
let empty_vec = Vec::new();
let data = args
.get("data")
.and_then(|v| v.as_array())
.unwrap_or(&empty_vec);
let operation = args
.get("operation")
.and_then(|v| v.as_str())
.unwrap_or("sum");
let result = match operation {
"sum" => data.iter().filter_map(|v| v.as_f64()).sum::<f64>(),
"count" => data.len() as f64,
"average" => {
let sum: f64 = data.iter().filter_map(|v| v.as_f64()).sum();
let count = data.iter().filter(|v| v.is_number()).count();
if count > 0 {
sum / count as f64
} else {
0.0
}
},
_ => 0.0,
};
Ok(json!({
"operation": operation,
"result": result,
"input_size": data.len(),
"processed_items": data.iter().filter(|v| v.is_number()).count()
}))
}
}
fn bench_server_creation(c: &mut Criterion) {
let mut group = c.benchmark_group("server_creation");
group.bench_function("basic_server_build", |b| {
b.iter(|| {
let server = black_box(
Server::builder()
.name("benchmark-server")
.version("1.0.0")
.tool("simple_tool", BenchmarkTool)
.build()
.unwrap(),
);
black_box(server)
})
});
group.bench_function("complex_server_build", |b| {
b.iter(|| {
let server = black_box(
Server::builder()
.name("complex-benchmark-server")
.version("2.0.0")
.tool("simple_tool", BenchmarkTool)
.tool("complex_tool", ComplexBenchmarkTool)
.tool("another_tool", BenchmarkTool)
.capabilities({
let mut caps = ServerCapabilities::default();
caps.tools = Some(ToolCapabilities {
list_changed: Some(true),
});
caps
})
.build()
.unwrap(),
);
black_box(server)
})
});
group.finish();
}
fn bench_request_processing(c: &mut Criterion) {
let mut group = c.benchmark_group("request_processing");
let requests = vec![
("ping", ClientRequest::Ping),
(
"list_tools",
ClientRequest::ListTools(ListToolsRequest::default()),
),
(
"simple_call_tool",
ClientRequest::CallTool(CallToolRequest::new(
"simple_tool",
json!({"input": "test"}),
)),
),
(
"complex_call_tool",
ClientRequest::CallTool(CallToolRequest::new(
"complex_tool",
json!({
"data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
"operation": "average"
}),
)),
),
(
"initialize",
ClientRequest::Initialize(InitializeRequest::new(
Implementation::new("benchmark-client", "1.0.0"),
ClientCapabilities::default(),
)),
),
];
for (name, request) in requests {
group.bench_with_input(
BenchmarkId::new("serialize_request", name),
&request,
|b, request| {
b.iter(|| {
let serialized = serde_json::to_string(black_box(request)).unwrap();
black_box(serialized)
})
},
);
group.bench_with_input(
BenchmarkId::new("roundtrip_request", name),
&request,
|b, request| {
b.iter(|| {
let serialized = serde_json::to_string(black_box(request)).unwrap();
let deserialized: ClientRequest = serde_json::from_str(&serialized).unwrap();
black_box(deserialized)
})
},
);
}
group.finish();
}
fn bench_response_generation(c: &mut Criterion) {
let mut group = c.benchmark_group("response_generation");
let responses = vec![
(
"simple_success",
JSONRPCResponse {
jsonrpc: "2.0".to_string(),
id: RequestId::Number(1),
payload: pmcp::types::jsonrpc::ResponsePayload::Result(json!("pong")),
},
),
(
"tool_list",
JSONRPCResponse {
jsonrpc: "2.0".to_string(),
id: RequestId::Number(2),
payload: pmcp::types::jsonrpc::ResponsePayload::Result(
serde_json::to_value(ListToolsResult::new(vec![
ToolInfo::new(
"tool1",
Some("First tool".to_string()),
json!({"type": "object"}),
),
ToolInfo::new(
"tool2",
Some("Second tool".to_string()),
json!({"type": "object"}),
),
]))
.unwrap(),
),
},
),
(
"tool_call_result",
JSONRPCResponse {
jsonrpc: "2.0".to_string(),
id: RequestId::Number(3),
payload: pmcp::types::jsonrpc::ResponsePayload::Result(
serde_json::to_value(CallToolResult::new(vec![Content::text(
"This is the result of a tool call with substantial output data.",
)]))
.unwrap(),
),
},
),
(
"error_response",
JSONRPCResponse {
jsonrpc: "2.0".to_string(),
id: RequestId::Number(4),
payload: pmcp::types::jsonrpc::ResponsePayload::Error(JSONRPCError {
code: -32602,
message: "Invalid params".to_string(),
data: Some(json!({"details": "Missing required parameter"})),
}),
},
),
];
for (name, response) in responses {
group.bench_with_input(
BenchmarkId::new("serialize_response", name),
&response,
|b, response| {
b.iter(|| {
let serialized = serde_json::to_string(black_box(response)).unwrap();
black_box(serialized)
})
},
);
}
group.finish();
}
fn bench_capabilities(c: &mut Criterion) {
let mut group = c.benchmark_group("capabilities");
let mut full_client_capabilities = ClientCapabilities::default();
full_client_capabilities.sampling = Some(SamplingCapabilities::default());
full_client_capabilities.elicitation = Some(ElicitationCapabilities::default());
full_client_capabilities.roots = Some(RootsCapabilities { list_changed: true });
group.bench_function("serialize_full_client_capabilities", |b| {
b.iter(|| {
let serialized = serde_json::to_string(black_box(&full_client_capabilities)).unwrap();
black_box(serialized)
})
});
let mut full_server_capabilities = ServerCapabilities::default();
full_server_capabilities.tools = Some(ToolCapabilities {
list_changed: Some(true),
});
full_server_capabilities.prompts = Some(PromptCapabilities {
list_changed: Some(true),
});
full_server_capabilities.resources = Some(ResourceCapabilities {
subscribe: Some(true),
list_changed: Some(true),
});
full_server_capabilities.logging = Some(LoggingCapabilities { levels: None });
group.bench_function("serialize_full_server_capabilities", |b| {
b.iter(|| {
let serialized = serde_json::to_string(black_box(&full_server_capabilities)).unwrap();
black_box(serialized)
})
});
group.bench_function("roundtrip_capabilities", |b| {
b.iter(|| {
let serialized = serde_json::to_string(black_box(&full_client_capabilities)).unwrap();
let deserialized: ClientCapabilities = serde_json::from_str(&serialized).unwrap();
black_box(deserialized)
})
});
group.finish();
}
fn bench_notifications(c: &mut Criterion) {
let mut group = c.benchmark_group("notifications");
let notifications = vec![
(
"progress",
Notification::Progress(ProgressNotification::new(
ProgressToken::String("task_123".to_string()),
75.0,
Some("Processing data...".to_string()),
)),
),
(
"cancelled",
Notification::Cancelled(
CancelledNotification::new(RequestId::Number(42))
.with_reason("User requested cancellation"),
),
),
(
"client_initialized",
Notification::Client(ClientNotification::Initialized),
),
(
"server_tools_changed",
Notification::Server(ServerNotification::ToolsChanged),
),
];
for (name, notification) in notifications {
group.bench_with_input(
BenchmarkId::new("serialize_notification", name),
¬ification,
|b, notification| {
b.iter(|| {
let serialized = serde_json::to_string(black_box(notification)).unwrap();
black_box(serialized)
})
},
);
}
group.finish();
}
fn bench_batch_operations(c: &mut Criterion) {
let mut group = c.benchmark_group("batch_operations");
let tool_calls: Vec<ClientRequest> = (0..50)
.map(|i| {
ClientRequest::CallTool(CallToolRequest::new(
"batch_tool",
json!({
"id": i,
"data": format!("Batch item {}", i),
"index": i
}),
))
})
.collect();
group.bench_function("serialize_batch_tool_calls", |b| {
b.iter(|| {
let mut results = Vec::new();
for request in &tool_calls {
let serialized = serde_json::to_string(black_box(request)).unwrap();
results.push(serialized);
}
black_box(results)
})
});
let responses: Vec<CallToolResult> = (0..50)
.map(|i| {
CallToolResult::new(vec![Content::text(format!(
"Result for batch item {}: processed successfully",
i
))])
})
.collect();
group.bench_function("serialize_batch_responses", |b| {
b.iter(|| {
let mut results = Vec::new();
for response in &responses {
let serialized = serde_json::to_string(black_box(response)).unwrap();
results.push(serialized);
}
black_box(results)
})
});
group.finish();
}
criterion_group!(
client_server_benches,
bench_server_creation,
bench_request_processing,
bench_response_generation,
bench_capabilities,
bench_notifications,
bench_batch_operations
);
criterion_main!(client_server_benches);