1use std::future::Future;
5use std::pin::Pin;
6use std::sync::Arc;
7
8use std::borrow::Cow;
9use std::cell::RefCell;
10use std::collections::{BTreeMap, HashMap};
11
12use rmcp::handler::server::common::schema_for_type;
13use rmcp::model::{ElicitResult, ElicitationAction};
14use schemars::JsonSchema;
15use serde::de::DeserializeOwned;
16use serde::{Deserialize, Serialize};
17use serde_json::{Map, Value};
18use tokio_util::sync::CancellationToken;
19
20use crate::seams::{Feature, Space};
21
22pub type JsonObject = Map<String, Value>;
24
25#[derive(Clone)]
27pub struct Ctx {
28 space: Arc<Space>,
29 cancel: Option<CancellationToken>,
32 progress: Option<Arc<dyn Progress>>,
33}
34
35#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
38pub struct Ask {
39 pub key: String,
41 pub message: String,
42 pub form: Form,
43}
44
45#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
47#[serde(rename_all = "snake_case")]
48pub enum Form {
49 Fields(JsonObject),
52 Url(String),
55}
56
57#[derive(Clone, Copy, Debug, PartialEq, Eq)]
61pub struct NeedsInput;
62
63impl std::fmt::Display for NeedsInput {
64 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
65 f.write_str("the operation needs an answer from the user")
66 }
67}
68
69impl std::error::Error for NeedsInput {}
70
71impl From<NeedsInput> for String {
72 fn from(needs: NeedsInput) -> Self {
73 needs.to_string()
74 }
75}
76
77pub(crate) struct Input {
81 answers: Option<BTreeMap<String, Value>>,
82 effects: Effects,
83}
84
85#[derive(Default)]
87pub(crate) struct Effects {
88 pub asked: Vec<Ask>,
90 pub updated: Vec<String>,
92 pub attached: Vec<Content>,
94}
95
96tokio::task_local! {
97 static INPUT: RefCell<Input>;
98}
99
100pub(crate) async fn with_input<F: Future>(
103 answers: Option<BTreeMap<String, Value>>,
104 work: F,
105) -> (F::Output, Effects) {
106 let input = RefCell::new(Input {
107 answers,
108 effects: Effects::default(),
109 });
110 INPUT
111 .scope(input, async {
112 let output = work.await;
113 let effects = INPUT.with(|input| std::mem::take(&mut input.borrow_mut().effects));
114 (output, effects)
115 })
116 .await
117}
118
119pub(crate) trait Progress: Send + Sync {
122 fn report(
123 &self,
124 done: f64,
125 total: Option<f64>,
126 message: Option<String>,
127 ) -> Pin<Box<dyn Future<Output = ()> + Send + '_>>;
128}
129
130impl Ctx {
131 pub(crate) fn new(
132 space: Arc<Space>,
133 cancel: Option<CancellationToken>,
134 progress: Option<Arc<dyn Progress>>,
135 ) -> Self {
136 Self {
137 space,
138 cancel,
139 progress,
140 }
141 }
142
143 pub fn space(&self) -> &Space {
146 &self.space
147 }
148
149 pub(crate) fn shared_space(&self) -> Arc<Space> {
151 self.space.clone()
152 }
153
154 pub fn is_cancelled(&self) -> bool {
158 self.cancel
159 .as_ref()
160 .is_some_and(CancellationToken::is_cancelled)
161 }
162
163 pub async fn cancelled(&self) {
166 match &self.cancel {
167 Some(cancel) => cancel.cancelled().await,
168 None => std::future::pending().await,
169 }
170 }
171
172 pub fn ask<T: DeserializeOwned + JsonSchema + 'static>(
179 &self,
180 key: &str,
181 message: &str,
182 ) -> Result<Option<T>, NeedsInput> {
183 let fields = || Form::Fields(JsonObject::clone(&schema_for_type::<T>()));
184 match self.answer(key, message, fields)? {
185 Some(content) => match serde_json::from_value(content) {
186 Ok(answer) => Ok(Some(answer)),
187 Err(_) => {
189 self.record(key, message, fields());
190 Err(NeedsInput)
191 }
192 },
193 None => Ok(None),
194 }
195 }
196
197 pub fn ask_url(&self, key: &str, message: &str, url: &str) -> Result<bool, NeedsInput> {
203 let answered = self.answer(key, message, || Form::Url(url.to_owned()))?;
204 Ok(answered.is_some())
205 }
206
207 fn answer(
211 &self,
212 key: &str,
213 message: &str,
214 form: impl FnOnce() -> Form,
215 ) -> Result<Option<Value>, NeedsInput> {
216 let given = INPUT
217 .try_with(|input| {
218 let input = input.borrow();
219 let given = input.answers.as_ref()?.get(key)?;
220 serde_json::from_value::<ElicitResult>(given.clone()).ok()
221 })
222 .map_err(|_| NeedsInput)?;
224 match given {
225 Some(result) => Ok(match result.action {
226 ElicitationAction::Accept => Some(result.content.unwrap_or(Value::Null)),
227 _ => None,
228 }),
229 None => {
230 self.record(key, message, form());
231 Err(NeedsInput)
232 }
233 }
234 }
235
236 fn record(&self, key: &str, message: &str, form: Form) {
238 let _ = INPUT.try_with(|input| {
239 input.borrow_mut().effects.asked.push(Ask {
240 key: key.to_owned(),
241 message: message.to_owned(),
242 form,
243 });
244 });
245 }
246
247 pub fn attach(&self, content: Content) {
251 let _ = INPUT.try_with(|input| input.borrow_mut().effects.attached.push(content));
253 }
254
255 pub fn updated(&self, uri: impl Into<String>) {
258 let _ = INPUT.try_with(|input| input.borrow_mut().effects.updated.push(uri.into()));
260 }
261
262 pub async fn progress(&self, done: f64, total: Option<f64>, message: Option<&str>) {
265 if let Some(progress) = &self.progress {
266 progress
267 .report(done, total, message.map(str::to_owned))
268 .await;
269 }
270 }
271}
272
273pub(crate) type Handler<M, T> = Box<
277 dyn Fn(Arc<M>, Ctx, Value) -> Pin<Box<dyn Future<Output = Result<T, String>> + Send>>
278 + Send
279 + Sync,
280>;
281
282fn erase<M, A, R, E, F, Fut, T>(handler: F, finish: fn(R) -> Result<T, String>) -> Handler<M, T>
285where
286 M: Send + Sync + 'static,
287 A: DeserializeOwned,
288 R: 'static,
289 E: std::fmt::Display,
290 F: Fn(Arc<M>, Ctx, A) -> Fut + Send + Sync + 'static,
291 Fut: Future<Output = Result<R, E>> + Send + 'static,
292 T: 'static,
293{
294 Box::new(move |module, ctx, arguments| {
295 let started = serde_json::from_value(arguments)
299 .map(|arguments| handler(module, ctx, arguments))
300 .map_err(|error| format!("invalid arguments: {error}"));
301 Box::pin(async move {
302 let result = started?.await.map_err(|error| error.to_string())?;
303 finish(result)
304 })
305 })
306}
307
308#[derive(Clone, Debug)]
310pub struct Tool {
311 pub name: &'static str,
312 pub feature: Feature,
313 pub description: Option<String>,
315 pub input: Arc<JsonObject>,
316 pub output: Option<Arc<JsonObject>>,
319 pub long: bool,
321 pub ui: Option<&'static str>,
323 pub app_only: bool,
326 pub hints: Option<Hints>,
328 pub looks: Looks,
329}
330
331#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize)]
334pub struct Looks {
335 pub title: Option<&'static str>,
337 pub icons: &'static [Icon],
338 pub mime: Option<&'static str>,
340 pub audience: &'static [Role],
342 pub priority: Option<f32>,
344 pub about: Option<&'static str>,
347 pub website: Option<&'static str>,
349}
350
351#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
353pub struct Icon {
354 pub src: &'static str,
355 pub mime: Option<&'static str>,
356 pub sizes: &'static [&'static str],
358}
359
360#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
365pub struct Hints {
366 pub read_only: bool,
368 pub destructive: bool,
370 pub idempotent: bool,
372 pub open_world: bool,
374}
375
376impl Default for Hints {
377 fn default() -> Self {
378 Self {
379 read_only: false,
380 destructive: true,
381 idempotent: false,
382 open_world: true,
383 }
384 }
385}
386
387#[derive(Clone, Debug)]
391pub struct Resource {
392 pub uri: &'static str,
393 pub name: &'static str,
394 pub feature: Feature,
395 pub description: Option<String>,
397 pub looks: Looks,
398}
399
400impl Resource {
401 pub fn is_template(&self) -> bool {
402 self.uri.contains('{')
403 }
404}
405
406#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
408#[serde(rename_all = "snake_case")]
409pub enum Body {
410 Text {
411 mime: Cow<'static, str>,
412 text: String,
413 },
414 Bytes {
415 mime: Cow<'static, str>,
416 #[serde(with = "base64_bytes")]
417 bytes: Vec<u8>,
418 },
419}
420
421impl Body {
422 pub fn text(mime: impl Into<Cow<'static, str>>, text: impl Into<String>) -> Self {
423 Self::Text {
424 mime: mime.into(),
425 text: text.into(),
426 }
427 }
428
429 pub fn bytes(mime: impl Into<Cow<'static, str>>, bytes: impl Into<Vec<u8>>) -> Self {
430 Self::Bytes {
431 mime: mime.into(),
432 bytes: bytes.into(),
433 }
434 }
435
436 pub fn json<T: Serialize>(value: &T) -> Result<Self, String> {
438 let text = serde_json::to_string(value).map_err(|error| error.to_string())?;
439 Ok(Self::text("application/json", text))
440 }
441}
442
443#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
446#[serde(rename_all = "snake_case")]
447pub enum Content {
448 Text(String),
449 Image {
450 mime: Cow<'static, str>,
451 #[serde(with = "base64_bytes")]
452 data: Vec<u8>,
453 },
454 Audio {
455 mime: Cow<'static, str>,
456 #[serde(with = "base64_bytes")]
457 data: Vec<u8>,
458 },
459 Link {
461 uri: String,
462 name: String,
463 mime: Option<Cow<'static, str>>,
464 },
465 Resource {
467 uri: String,
468 body: Body,
469 },
470}
471
472mod base64_bytes {
475 use base64::Engine;
476 use base64::engine::general_purpose::STANDARD;
477 use serde::{Deserialize, Deserializer, Serializer};
478
479 pub fn serialize<S: Serializer>(bytes: &[u8], serializer: S) -> Result<S::Ok, S::Error> {
480 serializer.serialize_str(&STANDARD.encode(bytes))
481 }
482
483 pub fn deserialize<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Vec<u8>, D::Error> {
484 let text = String::deserialize(deserializer)?;
485 STANDARD.decode(text).map_err(serde::de::Error::custom)
486 }
487}
488
489pub(crate) struct Readable<M> {
491 pub resource: Resource,
492 pub read: Handler<M, Body>,
493}
494
495#[derive(Clone, Debug)]
498pub struct Prompt {
499 pub name: &'static str,
500 pub feature: Feature,
501 pub description: Option<String>,
503 pub arguments: Vec<PromptArgument>,
504 pub schema: Arc<JsonObject>,
507 pub looks: Looks,
508}
509
510#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
512pub struct PromptArgument {
513 pub name: String,
514 pub description: Option<String>,
515 pub required: bool,
516}
517
518#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
520#[serde(rename_all = "snake_case")]
521pub enum Role {
522 User,
523 Assistant,
524}
525
526#[derive(Clone, Debug, PartialEq, Eq)]
528pub struct Message {
529 pub role: Role,
530 pub content: Content,
531}
532
533impl Message {
534 pub fn text(role: Role, text: impl Into<String>) -> Self {
536 Self {
537 role,
538 content: Content::Text(text.into()),
539 }
540 }
541}
542
543pub(crate) struct Promptable<M> {
545 pub prompt: Prompt,
546 pub get: Handler<M, Vec<Message>>,
547}
548
549#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
552pub struct Completing {
553 pub typed: String,
555 #[serde(default)]
558 pub given: HashMap<String, String>,
559}
560
561pub(crate) struct Completer<M> {
562 pub target: &'static str,
564 pub argument: &'static str,
565 pub complete: Handler<M, Vec<String>>,
566}
567
568pub(crate) const APP_MIME: &str = "text/html;profile=mcp-app";
570
571#[derive(Clone, Copy, Debug, PartialEq, Eq)]
573pub struct App {
574 pub uri: &'static str,
575 pub html: &'static str,
576}
577
578pub(crate) struct Op<M> {
580 pub tool: Tool,
581 pub call: Handler<M, Value>,
582}
583
584pub struct Ops<M> {
586 pub(crate) list: Vec<Op<M>>,
587 pub(crate) resources: Vec<Readable<M>>,
588 pub(crate) prompts: Vec<Promptable<M>>,
589 pub(crate) completers: Vec<Completer<M>>,
590 pub(crate) apps: Vec<App>,
591 pub(crate) instructions: Option<&'static str>,
592 pub(crate) server: Looks,
593}
594
595impl<M: Send + Sync + 'static> Ops<M> {
596 pub(crate) fn new() -> Self {
597 Self {
598 list: Vec::new(),
599 resources: Vec::new(),
600 prompts: Vec::new(),
601 completers: Vec::new(),
602 apps: Vec::new(),
603 instructions: None,
604 server: Looks::default(),
605 }
606 }
607
608 pub fn tool<A, R, E, F, Fut>(
616 &mut self,
617 name: &'static str,
618 feature: Feature,
619 handler: F,
620 ) -> &mut Self
621 where
622 A: DeserializeOwned + JsonSchema + 'static,
623 R: Serialize + JsonSchema + 'static,
624 E: std::fmt::Display,
625 F: Fn(Arc<M>, Ctx, A) -> Fut + Send + Sync + 'static,
626 Fut: Future<Output = Result<R, E>> + Send + 'static,
627 {
628 self.add(name, feature, false, handler)
629 }
630
631 pub fn task<A, R, E, F, Fut>(
635 &mut self,
636 name: &'static str,
637 feature: Feature,
638 handler: F,
639 ) -> &mut Self
640 where
641 A: DeserializeOwned + JsonSchema + 'static,
642 R: Serialize + JsonSchema + 'static,
643 E: std::fmt::Display,
644 F: Fn(Arc<M>, Ctx, A) -> Fut + Send + Sync + 'static,
645 Fut: Future<Output = Result<R, E>> + Send + 'static,
646 {
647 self.add(name, feature, true, handler)
648 }
649
650 pub fn app(&mut self, tool: &str, uri: &'static str, html: &'static str) -> &mut Self {
660 assert!(
661 uri.starts_with("ui://"),
662 "the app {uri} must be a ui:// URI"
663 );
664 self.registered(tool).ui = Some(uri);
665 match self.apps.iter().find(|app| app.uri == uri) {
666 Some(app) => assert_eq!(app.html, html, "{uri} already serves another page"),
667 None => self.apps.push(App { uri, html }),
668 }
669 self
670 }
671
672 pub fn resource<A, E, F, Fut>(
682 &mut self,
683 uri: &'static str,
684 name: &'static str,
685 feature: Feature,
686 handler: F,
687 ) -> &mut Self
688 where
689 A: DeserializeOwned + JsonSchema + 'static,
690 E: std::fmt::Display,
691 F: Fn(Arc<M>, Ctx, A) -> Fut + Send + Sync + 'static,
692 Fut: Future<Output = Result<Body, E>> + Send + 'static,
693 {
694 assert!(
695 self.resources.iter().all(|r| r.resource.uri != uri),
696 "resource {uri} is registered twice"
697 );
698 let description = take_text(&mut JsonObject::clone(&schema_for_type::<A>()));
699 self.resources.push(Readable {
700 resource: Resource {
701 uri,
702 name,
703 feature,
704 description,
705 looks: Looks::default(),
706 },
707 read: erase(handler, Ok),
708 });
709 self
710 }
711
712 pub fn prompt<A, E, F, Fut>(
720 &mut self,
721 name: &'static str,
722 feature: Feature,
723 handler: F,
724 ) -> &mut Self
725 where
726 A: DeserializeOwned + JsonSchema + 'static,
727 E: std::fmt::Display,
728 F: Fn(Arc<M>, Ctx, A) -> Fut + Send + Sync + 'static,
729 Fut: Future<Output = Result<Vec<Message>, E>> + Send + 'static,
730 {
731 assert!(
732 self.prompts.iter().all(|p| p.prompt.name != name),
733 "prompt {name} is registered twice"
734 );
735 let mut schema = JsonObject::clone(&schema_for_type::<A>());
736 assert!(
737 is_object(&schema),
738 "the arguments of {name} must be a struct"
739 );
740 let description = take_text(&mut schema);
741 let arguments = arguments(&schema);
742 self.prompts.push(Promptable {
743 prompt: Prompt {
744 name,
745 feature,
746 description,
747 arguments,
748 schema: Arc::new(schema),
749 looks: Looks::default(),
750 },
751 get: erase(handler, Ok),
752 });
753 self
754 }
755
756 pub fn complete<E, F, Fut>(
764 &mut self,
765 target: &'static str,
766 argument: &'static str,
767 handler: F,
768 ) -> &mut Self
769 where
770 E: std::fmt::Display,
771 F: Fn(Arc<M>, Ctx, Completing) -> Fut + Send + Sync + 'static,
772 Fut: Future<Output = Result<Vec<String>, E>> + Send + 'static,
773 {
774 let known = self.prompts.iter().any(|p| p.prompt.name == target)
775 || self.resources.iter().any(|r| r.resource.uri == target);
776 assert!(known, "register {target} before completing it");
777 self.completers.push(Completer {
778 target,
779 argument,
780 complete: erase(handler, Ok),
781 });
782 self
783 }
784
785 pub fn app_only(&mut self, tool: &str) -> &mut Self {
793 self.registered(tool).app_only = true;
794 self
795 }
796
797 pub fn hints(&mut self, tool: &str, hints: Hints) -> &mut Self {
803 self.registered(tool).hints = Some(hints);
804 self
805 }
806
807 pub fn looks(&mut self, target: &str, looks: Looks) -> &mut Self {
814 if let Some(op) = self.list.iter_mut().find(|op| op.tool.name == target) {
815 op.tool.looks = looks;
816 } else if let Some(p) = self.prompts.iter_mut().find(|p| p.prompt.name == target) {
817 p.prompt.looks = looks;
818 } else if let Some(r) = self.resources.iter_mut().find(|r| r.resource.uri == target) {
819 r.resource.looks = looks;
820 } else {
821 panic!("register {target} before describing it");
822 }
823 self
824 }
825
826 pub fn server(&mut self, looks: Looks) -> &mut Self {
828 self.server = looks;
829 self
830 }
831
832 pub fn instructions(&mut self, text: &'static str) -> &mut Self {
835 self.instructions = Some(text);
836 self
837 }
838
839 fn registered(&mut self, name: &str) -> &mut Tool {
841 self.list
842 .iter_mut()
843 .map(|op| &mut op.tool)
844 .find(|tool| tool.name == name)
845 .unwrap_or_else(|| panic!("register {name} before describing it"))
846 }
847
848 fn add<A, R, E, F, Fut>(
849 &mut self,
850 name: &'static str,
851 feature: Feature,
852 long: bool,
853 handler: F,
854 ) -> &mut Self
855 where
856 A: DeserializeOwned + JsonSchema + 'static,
857 R: Serialize + JsonSchema + 'static,
858 E: std::fmt::Display,
859 F: Fn(Arc<M>, Ctx, A) -> Fut + Send + Sync + 'static,
860 Fut: Future<Output = Result<R, E>> + Send + 'static,
861 {
862 assert!(
863 self.list.iter().all(|op| op.tool.name != name),
864 "operation {name} is registered twice"
865 );
866 let mut input = JsonObject::clone(&schema_for_type::<A>());
868 assert!(
869 is_object(&input),
870 "the arguments of {name} must be a struct"
871 );
872 let description = take_text(&mut input);
873 let output = Some(JsonObject::clone(&schema_for_type::<R>()))
874 .filter(is_object)
875 .map(|mut output| {
876 take_text(&mut output);
877 Arc::new(output)
878 });
879 let call = erase(handler, |result: R| {
880 serde_json::to_value(result).map_err(|error| error.to_string())
881 });
882 self.list.push(Op {
883 tool: Tool {
884 name,
885 feature,
886 description,
887 input: Arc::new(input),
888 output,
889 long,
890 ui: None,
891 app_only: false,
892 hints: None,
893 looks: Looks::default(),
894 },
895 call,
896 });
897 self
898 }
899}
900
901fn arguments(schema: &JsonObject) -> Vec<PromptArgument> {
903 let required = |name: &str| {
904 schema
905 .get("required")
906 .and_then(Value::as_array)
907 .is_some_and(|required| required.iter().any(|r| r == name))
908 };
909 let Some(Value::Object(properties)) = schema.get("properties") else {
910 return Vec::new();
911 };
912 properties
913 .iter()
914 .map(|(name, property)| PromptArgument {
915 name: name.clone(),
916 description: property
917 .get("description")
918 .and_then(Value::as_str)
919 .map(str::to_owned),
920 required: required(name),
921 })
922 .collect()
923}
924
925fn is_object(schema: &JsonObject) -> bool {
926 schema.get("type").and_then(Value::as_str) == Some("object")
927}
928
929fn take_text(schema: &mut JsonObject) -> Option<String> {
933 schema.remove("title");
934 match schema.remove("description") {
935 Some(Value::String(text)) => Some(text),
936 _ => None,
937 }
938}