1use crate::prompt::Prompt;
15use crate::protocol::{PROTOCOL_VERSION, SUPPORTED_VERSIONS, Implementation, method};
16use crate::resource::Resource;
17use crate::rpc::{self, Frame, Id, Message, RpcError};
18use crate::tool::Tool;
19use rustlavel_core::events::Event;
20use rustlavel_core::Json;
21use std::collections::BTreeMap;
22use std::time::Instant;
23
24pub struct Server {
29 info: Implementation,
30 instructions: Option<String>,
31 tools: BTreeMap<String, Tool>,
32 resources: BTreeMap<String, Resource>,
33 prompts: BTreeMap<String, Prompt>,
34}
35
36impl Server {
37 pub fn new(name: impl Into<String>, version: impl Into<String>) -> Self {
38 Server {
39 info: Implementation::new(name, version),
40 instructions: None,
41 tools: BTreeMap::new(),
42 resources: BTreeMap::new(),
43 prompts: BTreeMap::new(),
44 }
45 }
46
47 pub fn instructions(mut self, instructions: impl Into<String>) -> Self {
49 self.instructions = Some(instructions.into());
50 self
51 }
52
53 pub fn tool(mut self, tool: Tool) -> Self {
54 self.tools.insert(tool.name.clone(), tool);
55 self
56 }
57
58 pub fn resource(mut self, resource: Resource) -> Self {
59 self.resources.insert(resource.uri.clone(), resource);
60 self
61 }
62
63 pub fn prompt(mut self, prompt: Prompt) -> Self {
64 self.prompts.insert(prompt.name.clone(), prompt);
65 self
66 }
67
68 pub fn name(&self) -> &str {
69 &self.info.name
70 }
71
72 pub fn tool_names(&self) -> Vec<&str> {
73 self.tools.keys().map(String::as_str).collect()
74 }
75
76 pub fn capabilities(&self) -> Json {
82 let mut pairs: Vec<(&str, Json)> = Vec::new();
83 if !self.tools.is_empty() {
84 pairs.push(("tools", Json::object([("listChanged", Json::from(false))])));
85 }
86 if !self.resources.is_empty() {
87 pairs.push((
88 "resources",
89 Json::object([
90 ("listChanged", Json::from(false)),
91 ("subscribe", Json::from(false)),
92 ]),
93 ));
94 }
95 if !self.prompts.is_empty() {
96 pairs.push(("prompts", Json::object([("listChanged", Json::from(false))])));
97 }
98 Json::object(pairs)
99 }
100
101 pub async fn handle_text(&self, text: &str) -> Option<String> {
106 let frame = match Frame::parse(text) {
107 Ok(frame) => frame,
108 Err(error) => {
110 return Some(rpc::Response::failure(None, error).to_json().to_string());
111 }
112 };
113
114 let batched = frame.is_batch();
115 let mut answers = Vec::new();
116 for item in frame.items() {
117 if let Some(answer) = self.handle_value(item).await {
118 answers.push(answer);
119 }
120 }
121 rpc::encode(answers, batched)
122 }
123
124 pub async fn handle_value(&self, value: &Json) -> Option<Json> {
126 match Message::from_json(value) {
127 Ok(Message::Request(request)) => {
128 let id = request.id.clone();
129 Some(match self.dispatch(&request.method, request.params).await {
130 Ok(result) => rpc::Response::success(id, result).to_json(),
131 Err(error) => rpc::Response::failure(Some(id), error).to_json(),
132 })
133 }
134 Ok(Message::Notification(_)) => None,
137 Ok(Message::Response(_)) => None,
140 Err(error) => {
141 let id = value.get("id").and_then(Id::from_json);
144 Some(rpc::Response::failure(id, error).to_json())
145 }
146 }
147 }
148
149 pub async fn dispatch(&self, method: &str, params: Json) -> Result<Json, RpcError> {
151 match method {
152 method::INITIALIZE => Ok(self.initialize(¶ms)),
153 method::PING => Ok(Json::Object(Default::default())),
154 method::TOOLS_LIST => Ok(Json::object([(
155 "tools",
156 Json::Array(self.tools.values().map(|tool| tool.info().to_json()).collect()),
157 )])),
158 method::TOOLS_CALL => self.call_tool(¶ms).await,
159 method::RESOURCES_LIST => Ok(Json::object([(
160 "resources",
161 Json::Array(
162 self.resources.values().map(|resource| resource.info().to_json()).collect(),
163 ),
164 )])),
165 method::RESOURCES_READ => self.read_resource(¶ms).await,
166 method::PROMPTS_LIST => Ok(Json::object([(
167 "prompts",
168 Json::Array(self.prompts.values().map(|prompt| prompt.info().to_json()).collect()),
169 )])),
170 method::PROMPTS_GET => self.get_prompt(¶ms).await,
171 unknown => Err(RpcError::method_not_found(unknown)),
172 }
173 }
174
175 fn initialize(&self, params: &Json) -> Json {
181 let asked = params.get("protocolVersion").and_then(Json::as_str).unwrap_or("");
182 let version = if SUPPORTED_VERSIONS.contains(&asked) { asked } else { PROTOCOL_VERSION };
183
184 let mut pairs = vec![
185 ("protocolVersion", Json::from(version)),
186 ("capabilities", self.capabilities()),
187 ("serverInfo", self.info.to_json()),
188 ];
189 if let Some(instructions) = &self.instructions {
190 pairs.push(("instructions", Json::from(instructions.clone())));
191 }
192 Json::object(pairs)
193 }
194
195 async fn call_tool(&self, params: &Json) -> Result<Json, RpcError> {
196 let Some(name) = params.get("name").and_then(Json::as_str) else {
197 return Err(RpcError::invalid_params("`name` is required"));
198 };
199 let arguments = params
200 .as_object()
201 .and_then(|map| map.get("arguments"))
202 .cloned()
203 .unwrap_or(Json::Null);
204
205 let started = Instant::now();
206 let outcome = self.run_tool(name, arguments).await;
207
208 record_call(name, started, &outcome);
210
211 outcome
212 }
213
214 async fn run_tool(&self, name: &str, arguments: Json) -> Result<Json, RpcError> {
215 let Some(tool) = self.tools.get(name) else {
218 return Err(RpcError::invalid_params(format!("unknown tool `{name}`"))
219 .with_data(Json::object([(
220 "available",
221 Json::Array(self.tools.keys().cloned().map(Json::from).collect()),
222 )])));
223 };
224
225 if let Err(problems) = tool.schema.validate(&arguments) {
228 return Err(RpcError::invalid_params(problems.join("; ")).with_data(Json::object([(
229 "errors",
230 Json::Array(problems.into_iter().map(Json::from).collect()),
231 )])));
232 }
233
234 Ok(tool.call(arguments).await.to_json())
235 }
236
237 async fn read_resource(&self, params: &Json) -> Result<Json, RpcError> {
238 let Some(uri) = params.get("uri").and_then(Json::as_str) else {
239 return Err(RpcError::invalid_params("`uri` is required"));
240 };
241 let Some(resource) = self.resources.get(uri) else {
242 return Err(RpcError::resource_not_found(uri));
243 };
244
245 resource.contents().await.map_err(RpcError::internal_error)
246 }
247
248 async fn get_prompt(&self, params: &Json) -> Result<Json, RpcError> {
249 let Some(name) = params.get("name").and_then(Json::as_str) else {
250 return Err(RpcError::invalid_params("`name` is required"));
251 };
252 let Some(prompt) = self.prompts.get(name) else {
253 return Err(RpcError::invalid_params(format!("unknown prompt `{name}`")));
254 };
255
256 let arguments = params
257 .as_object()
258 .and_then(|map| map.get("arguments"))
259 .cloned()
260 .unwrap_or(Json::Null);
261
262 let missing = prompt.missing_arguments(&arguments);
263 if !missing.is_empty() {
264 return Err(RpcError::invalid_params(missing.join("; ")));
265 }
266
267 prompt.render(arguments).await.map_err(RpcError::internal_error)
268 }
269}
270
271fn record_call(name: &str, started: Instant, outcome: &Result<Json, RpcError>) {
277 if !rustlavel_core::events::has_subscribers() {
278 return;
279 }
280
281 let (ok, error) = match outcome {
282 Err(rpc) => (false, Some(rpc.message.clone())),
283 Ok(result) => match result.get("isError").and_then(Json::as_bool) {
284 Some(true) => (
285 false,
286 result.get("content.0.text").and_then(Json::as_str).map(str::to_string),
287 ),
288 _ => (true, None),
289 },
290 };
291
292 let mut event = Event::new("mcp.call").with("tool", name).with("ok", ok).took(started.elapsed());
293 if let Some(error) = error {
294 event = event.with("error", error);
295 }
296 event.dispatch();
297}
298
299#[cfg(test)]
300mod tests {
301 use super::*;
302 use crate::protocol::{PromptArgument, PromptMessage, ToolResult};
303 use crate::schema::Schema;
304 use rustlavel_core::Error;
305
306 fn server() -> Server {
307 Server::new("test-app", "0.1.0")
308 .instructions("Call `greet` to say hello.")
309 .tool(Tool::new(
310 "greet",
311 "Greet somebody by name",
312 Schema::object().string("name", "Who to greet"),
313 |args: Json| async move {
314 let name = args.get("name").and_then(Json::as_str).unwrap_or("world");
315 Ok(Json::from(format!("Hello, {name}!")))
316 },
317 ))
318 .tool(Tool::new("explodes", "Always panics", Schema::object(), |_: Json| async {
319 panic!("held wrong");
320 }))
321 .tool(Tool::new("fails", "Always errors", Schema::object(), |_: Json| async {
322 Err(Error::msg("upstream is down"))
323 }))
324 .resource(Resource::json("app://routes", "Routes", || async {
325 Ok(r#"["GET /"]"#.to_string())
326 }))
327 .prompt(
328 Prompt::new("review", "Review an order", |args: Json| async move {
329 let id = args.get("order_id").and_then(Json::as_str).unwrap_or("?");
330 Ok(vec![PromptMessage::user(format!("Review order {id}"))])
331 })
332 .argument(PromptArgument::new("order_id", "Which order")),
333 )
334 }
335
336 async fn call(server: &Server, method: &str, params: Json) -> Json {
339 server.dispatch(method, params).await.unwrap_or_else(|error| {
340 panic!("{method} failed: {error}");
341 })
342 }
343
344 #[tokio::test]
345 async fn the_handshake_reports_the_protocol_version_and_capabilities() {
346 let result = call(
347 &server(),
348 method::INITIALIZE,
349 Json::object([("protocolVersion", Json::from(PROTOCOL_VERSION))]),
350 )
351 .await;
352
353 assert_eq!(result.get("protocolVersion").unwrap().as_str(), Some("2025-06-18"));
354 assert_eq!(result.get("serverInfo.name").unwrap().as_str(), Some("test-app"));
355 assert_eq!(result.get("serverInfo.version").unwrap().as_str(), Some("0.1.0"));
356 assert!(result.get("capabilities.tools").is_some());
357 assert!(result.get("capabilities.resources").is_some());
358 assert_eq!(
359 result.get("instructions").unwrap().as_str(),
360 Some("Call `greet` to say hello.")
361 );
362 }
363
364 #[tokio::test]
365 async fn an_older_protocol_version_is_answered_in_kind() {
366 let older = call(
367 &server(),
368 method::INITIALIZE,
369 Json::object([("protocolVersion", Json::from("2024-11-05"))]),
370 )
371 .await;
372 assert_eq!(older.get("protocolVersion").unwrap().as_str(), Some("2024-11-05"));
373
374 let unknown = call(
376 &server(),
377 method::INITIALIZE,
378 Json::object([("protocolVersion", Json::from("1999-01-01"))]),
379 )
380 .await;
381 assert_eq!(unknown.get("protocolVersion").unwrap().as_str(), Some(PROTOCOL_VERSION));
382 }
383
384 #[tokio::test]
385 async fn a_server_with_nothing_registered_advertises_nothing() {
386 let bare = Server::new("bare", "0.1.0");
387 assert!(bare.capabilities().as_object().unwrap().is_empty());
388 }
389
390 #[tokio::test]
391 async fn tools_list_returns_every_tool_with_its_schema() {
392 let result = call(&server(), method::TOOLS_LIST, Json::Null).await;
393 let tools = result.get("tools").unwrap().as_array().unwrap();
394
395 let names: Vec<&str> =
397 tools.iter().map(|t| t.get("name").unwrap().as_str().unwrap()).collect();
398 assert_eq!(names, ["explodes", "fails", "greet"]);
399
400 let greet = tools.iter().find(|t| t.get("name").unwrap().as_str() == Some("greet")).unwrap();
401 assert_eq!(greet.get("description").unwrap().as_str(), Some("Greet somebody by name"));
402 assert_eq!(greet.get("inputSchema.properties.name.type").unwrap().as_str(), Some("string"));
403 }
404
405 #[tokio::test]
406 async fn tools_call_runs_the_handler_and_returns_content() {
407 let result = call(
408 &server(),
409 method::TOOLS_CALL,
410 Json::object([
411 ("name", Json::from("greet")),
412 ("arguments", Json::object([("name", Json::from("Ada"))])),
413 ]),
414 )
415 .await;
416
417 assert_eq!(result.get("isError").unwrap().as_bool(), Some(false));
418 assert_eq!(result.get("content.0.type").unwrap().as_str(), Some("text"));
419 assert_eq!(result.get("content.0.text").unwrap().as_str(), Some("Hello, Ada!"));
420 }
421
422 #[tokio::test]
423 async fn an_unknown_tool_is_a_protocol_error_listing_what_exists() {
424 let error = server()
425 .dispatch(method::TOOLS_CALL, Json::object([("name", Json::from("nope"))]))
426 .await
427 .unwrap_err();
428
429 assert_eq!(error.code, rpc::codes::INVALID_PARAMS);
430 assert!(error.message.contains("unknown tool `nope`"));
431 assert!(error.data.unwrap().to_string().contains("greet"));
432 }
433
434 #[tokio::test]
435 async fn an_argument_of_the_wrong_type_never_reaches_the_handler() {
436 let error = server()
437 .dispatch(
438 method::TOOLS_CALL,
439 Json::object([
440 ("name", Json::from("greet")),
441 ("arguments", Json::object([("name", Json::from(42))])),
442 ]),
443 )
444 .await
445 .unwrap_err();
446
447 assert_eq!(error.code, rpc::codes::INVALID_PARAMS);
448 assert!(error.message.contains("`name` must be a string, found a number"));
449 assert_eq!(
450 error.data.unwrap().get("errors.0").unwrap().as_str(),
451 Some("`name` must be a string, found a number")
452 );
453 }
454
455 #[tokio::test]
456 async fn a_missing_required_argument_is_rejected() {
457 let error = server()
458 .dispatch(method::TOOLS_CALL, Json::object([("name", Json::from("greet"))]))
459 .await
460 .unwrap_err();
461
462 assert!(error.message.contains("`name` is required"));
463 }
464
465 #[tokio::test]
466 async fn tools_call_without_a_name_is_invalid_params() {
467 let error = server().dispatch(method::TOOLS_CALL, Json::Null).await.unwrap_err();
468 assert_eq!(error.code, rpc::codes::INVALID_PARAMS);
469 }
470
471 #[tokio::test]
472 async fn a_panicking_tool_becomes_an_error_result_rather_than_a_crash() {
473 let result =
474 call(&server(), method::TOOLS_CALL, Json::object([("name", Json::from("explodes"))]))
475 .await;
476
477 assert_eq!(result.get("isError").unwrap().as_bool(), Some(true));
478 assert!(result.get("content.0.text").unwrap().as_str().unwrap().contains("held wrong"));
479
480 let after = call(&server(), method::TOOLS_LIST, Json::Null).await;
482 assert_eq!(after.get("tools").unwrap().as_array().unwrap().len(), 3);
483 }
484
485 #[tokio::test]
486 async fn a_failing_tool_reports_through_is_error_not_a_protocol_error() {
487 let result =
488 call(&server(), method::TOOLS_CALL, Json::object([("name", Json::from("fails"))])).await;
489
490 let parsed = ToolResult::from_json(&result).unwrap();
491 assert!(parsed.is_error);
492 assert!(parsed.text_content().contains("upstream is down"));
493 }
494
495 #[tokio::test]
496 async fn resources_list_and_read_round_trip() {
497 let listed = call(&server(), method::RESOURCES_LIST, Json::Null).await;
498 assert_eq!(listed.get("resources.0.uri").unwrap().as_str(), Some("app://routes"));
499
500 let read = call(
501 &server(),
502 method::RESOURCES_READ,
503 Json::object([("uri", Json::from("app://routes"))]),
504 )
505 .await;
506 assert_eq!(read.get("contents.0.text").unwrap().as_str(), Some(r#"["GET /"]"#));
507 }
508
509 #[tokio::test]
510 async fn an_unknown_resource_uses_the_reserved_code() {
511 let error = server()
512 .dispatch(method::RESOURCES_READ, Json::object([("uri", Json::from("app://ghost"))]))
513 .await
514 .unwrap_err();
515
516 assert_eq!(error.code, rpc::codes::RESOURCE_NOT_FOUND);
517 }
518
519 #[tokio::test]
520 async fn prompts_list_and_get_round_trip() {
521 let listed = call(&server(), method::PROMPTS_LIST, Json::Null).await;
522 assert_eq!(listed.get("prompts.0.name").unwrap().as_str(), Some("review"));
523
524 let got = call(
525 &server(),
526 method::PROMPTS_GET,
527 Json::object([
528 ("name", Json::from("review")),
529 ("arguments", Json::object([("order_id", Json::from("7"))])),
530 ]),
531 )
532 .await;
533 assert_eq!(got.get("messages.0.content.text").unwrap().as_str(), Some("Review order 7"));
534 }
535
536 #[tokio::test]
537 async fn a_prompt_missing_a_required_argument_is_rejected() {
538 let error = server()
539 .dispatch(method::PROMPTS_GET, Json::object([("name", Json::from("review"))]))
540 .await
541 .unwrap_err();
542
543 assert_eq!(error.code, rpc::codes::INVALID_PARAMS);
544 assert!(error.message.contains("`order_id` is required"));
545 }
546
547 #[tokio::test]
548 async fn an_unknown_method_is_method_not_found() {
549 let error = server().dispatch("tools/teleport", Json::Null).await.unwrap_err();
550 assert_eq!(error.code, rpc::codes::METHOD_NOT_FOUND);
551 }
552
553 #[tokio::test]
554 async fn ping_answers_with_an_empty_result() {
555 assert_eq!(call(&server(), method::PING, Json::Null).await.to_string(), "{}");
556 }
557
558 #[tokio::test]
561 async fn a_request_frame_comes_back_correlated_by_id() {
562 let answer = server()
563 .handle_text(r#"{"jsonrpc":"2.0","id":"abc","method":"tools/list"}"#)
564 .await
565 .unwrap();
566 let value = Json::parse(&answer).unwrap();
567
568 assert_eq!(value.get("jsonrpc").unwrap().as_str(), Some("2.0"));
569 assert_eq!(value.get("id").unwrap().as_str(), Some("abc"));
570 assert!(value.get("result.tools").is_some());
571 }
572
573 #[tokio::test]
574 async fn a_notification_frame_is_answered_with_silence() {
575 let answer =
576 server().handle_text(r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#).await;
577 assert_eq!(answer, None);
578 }
579
580 #[tokio::test]
581 async fn a_malformed_frame_becomes_a_parse_error_with_a_null_id() {
582 let answer = server().handle_text("{ this is not json").await.unwrap();
583 let value = Json::parse(&answer).unwrap();
584
585 assert_eq!(value.get("error.code").unwrap().as_i64(), Some(rpc::codes::PARSE_ERROR));
586 assert!(value.get("id").unwrap().is_null());
587 }
588
589 #[tokio::test]
590 async fn an_invalid_message_keeps_the_id_it_arrived_with() {
591 let answer = server().handle_text(r#"{"id":5,"method":"ping"}"#).await.unwrap();
592 let value = Json::parse(&answer).unwrap();
593
594 assert_eq!(value.get("error.code").unwrap().as_i64(), Some(rpc::codes::INVALID_REQUEST));
595 assert_eq!(value.get("id").unwrap().as_i64(), Some(5));
596 }
597
598 #[tokio::test]
599 async fn a_batch_is_answered_by_a_batch_of_the_same_requests() {
600 let answer = server()
601 .handle_text(
602 r#"[{"jsonrpc":"2.0","id":1,"method":"ping"},
603 {"jsonrpc":"2.0","method":"notifications/initialized"},
604 {"jsonrpc":"2.0","id":2,"method":"tools/list"},
605 "garbage"]"#,
606 )
607 .await
608 .unwrap();
609
610 let answers = Json::parse(&answer).unwrap();
611 let items = answers.as_array().unwrap();
612
613 assert_eq!(items.len(), 3);
615 assert_eq!(items[0].get("id").unwrap().as_i64(), Some(1));
616 assert_eq!(items[1].get("id").unwrap().as_i64(), Some(2));
617 assert_eq!(items[2].get("error.code").unwrap().as_i64(), Some(rpc::codes::INVALID_REQUEST));
618 }
619
620 #[tokio::test]
621 async fn a_batch_of_only_notifications_is_answered_with_silence() {
622 let answer = server()
623 .handle_text(r#"[{"jsonrpc":"2.0","method":"notifications/initialized"}]"#)
624 .await;
625 assert_eq!(answer, None);
626 }
627
628 #[tokio::test]
629 async fn an_empty_batch_is_an_invalid_request() {
630 let answer = server().handle_text("[]").await.unwrap();
631 let value = Json::parse(&answer).unwrap();
632 assert_eq!(value.get("error.code").unwrap().as_i64(), Some(rpc::codes::INVALID_REQUEST));
633 }
634
635 #[tokio::test]
636 async fn a_tool_call_is_reported_on_the_instrumentation_bus() {
637 use rustlavel_core::events::{Event, subscribe};
638 use std::sync::{Arc, Mutex};
639
640 type Recorded = Vec<(bool, Option<f64>)>;
643
644 let seen: Arc<Mutex<Recorded>> = Arc::new(Mutex::new(Vec::new()));
647 let sink = Arc::clone(&seen);
648 subscribe(move |event: &Event| {
649 if event.kind != "mcp.call" {
650 return;
651 }
652 if event.field("tool").and_then(Json::as_str) != Some("watched") {
653 return;
654 }
655 let ok = event.field("ok").and_then(Json::as_bool).unwrap_or(false);
656 sink.lock().unwrap().push((ok, event.duration_ms()));
657 });
658
659 let server = Server::new("events", "0.1.0").tool(Tool::new(
660 "watched",
661 "Observed by the event test",
662 Schema::object(),
663 |_: Json| async { Ok(Json::from("done")) },
664 ));
665
666 server
667 .dispatch(method::TOOLS_CALL, Json::object([("name", Json::from("watched"))]))
668 .await
669 .unwrap();
670 let _ = server
671 .dispatch(method::TOOLS_CALL, Json::object([("name", Json::from("watched-typo"))]))
672 .await;
673
674 let recorded = seen.lock().unwrap().clone();
675 assert_eq!(recorded.len(), 1, "only the real call names the watched tool");
676 assert!(recorded[0].0, "the call succeeded");
677 assert!(recorded[0].1.is_some(), "the event carries a duration");
678 }
679
680 #[tokio::test]
681 async fn a_failing_call_is_reported_as_not_ok() {
682 use rustlavel_core::events::{Event, subscribe};
683 use std::sync::{Arc, Mutex};
684
685 let seen: Arc<Mutex<Vec<bool>>> = Arc::new(Mutex::new(Vec::new()));
686 let sink = Arc::clone(&seen);
687 subscribe(move |event: &Event| {
688 if event.kind == "mcp.call"
689 && event.field("tool").and_then(Json::as_str) == Some("watched-failure")
690 {
691 sink.lock().unwrap().push(event.field("ok").and_then(Json::as_bool).unwrap_or(true));
692 }
693 });
694
695 let server = Server::new("events", "0.1.0").tool(Tool::new(
696 "watched-failure",
697 "Observed by the event test",
698 Schema::object(),
699 |_: Json| async { Err(Error::msg("nope")) },
700 ));
701
702 server
703 .dispatch(method::TOOLS_CALL, Json::object([("name", Json::from("watched-failure"))]))
704 .await
705 .unwrap();
706
707 assert_eq!(seen.lock().unwrap().clone(), [false]);
708 }
709}