1use rmcp::model::{AnnotateAble, RawResource, RawResourceTemplate};
2use rmcp::model::{
3 CallToolResult, CancelTaskParams, CancelTaskResult, CompleteRequestParams, CompleteResult,
4 CompletionInfo, Content, GetPromptRequestParams, GetPromptResult, GetTaskInfoParams,
5 GetTaskPayloadResult, GetTaskResult, GetTaskResultParams, Implementation, ListPromptsResult,
6 ListResourceTemplatesResult, ListResourcesResult, ListTasksResult, ListToolsResult,
7 PaginatedRequestParams, Prompt, PromptArgument, ReadResourceRequestParams, ReadResourceResult,
8 Resource, ResourceContents, ResourceTemplate, ServerCapabilities, ServerInfo, Task, TaskStatus,
9 Tool,
10};
11use schemars::JsonSchema;
12use serde::de::DeserializeOwned;
13use serde::{Deserialize, Serialize};
14use serde_json::json;
15use std::collections::{BTreeMap, BTreeSet, HashMap};
16use std::future::Future;
17use std::pin::Pin;
18use std::sync::Arc;
19
20use crate::{
21 context::PluginContext,
22 error::{PluginError, PluginResult, PluginRpcResult},
23 proto,
24};
25
26fn default_arguments() -> serde_json::Value {
27 json!({})
28}
29
30#[derive(Debug, Clone, Deserialize, Serialize)]
31pub struct ToolCallRequest {
32 pub name: String,
33 #[serde(default = "default_arguments")]
34 pub arguments: serde_json::Value,
35}
36
37impl ToolCallRequest {
38 pub fn arguments<T: DeserializeOwned>(&self) -> PluginResult<T> {
39 serde_json::from_value(self.arguments.clone()).map_err(|err| {
40 PluginError::invalid_params(format!(
41 "Invalid arguments for tool '{}': {err}",
42 self.name
43 ))
44 })
45 }
46
47 pub fn arguments_or_default<T>(&self) -> PluginResult<T>
48 where
49 T: DeserializeOwned + Default,
50 {
51 self.arguments()
52 }
53}
54
55pub type OperationRequest = ToolCallRequest;
56
57pub fn json_string<T: Serialize>(value: &T) -> PluginResult<String> {
58 serde_json::to_string(value).map_err(|err| PluginError::internal(err.to_string()))
59}
60
61pub fn json_bytes<T: Serialize>(value: &T) -> PluginResult<Vec<u8>> {
62 serde_json::to_vec(value).map_err(|err| PluginError::internal(err.to_string()))
63}
64
65pub fn structured_tool_result<T: Serialize>(value: T) -> PluginResult<CallToolResult> {
66 let value =
67 serde_json::to_value(value).map_err(|err| PluginError::internal(err.to_string()))?;
68 Ok(CallToolResult::structured(value))
69}
70
71pub fn tool_error(message: impl Into<String>) -> CallToolResult {
72 CallToolResult::error(vec![Content::text(message.into())])
73}
74
75pub fn operation_error(message: impl Into<String>) -> CallToolResult {
76 tool_error(message)
77}
78
79pub fn list_tools(tools: Vec<Tool>) -> ListToolsResult {
80 ListToolsResult {
81 tools,
82 meta: None,
83 next_cursor: None,
84 }
85}
86
87pub fn list_prompts(prompts: Vec<Prompt>) -> ListPromptsResult {
88 ListPromptsResult {
89 prompts,
90 meta: None,
91 next_cursor: None,
92 }
93}
94
95pub fn list_resources(resources: Vec<Resource>) -> ListResourcesResult {
96 ListResourcesResult {
97 resources,
98 meta: None,
99 next_cursor: None,
100 }
101}
102
103pub fn list_resource_templates(
104 resource_templates: Vec<ResourceTemplate>,
105) -> ListResourceTemplatesResult {
106 ListResourceTemplatesResult {
107 resource_templates,
108 meta: None,
109 next_cursor: None,
110 }
111}
112
113pub fn list_tasks(tasks: Vec<Task>) -> ListTasksResult {
114 ListTasksResult::new(tasks)
115}
116
117pub fn read_resource_result(contents: Vec<ResourceContents>) -> ReadResourceResult {
118 ReadResourceResult::new(contents)
119}
120
121pub fn get_prompt_result(messages: Vec<rmcp::model::PromptMessage>) -> GetPromptResult {
122 GetPromptResult::new(messages)
123}
124
125pub fn complete_result(values: Vec<String>) -> PluginResult<CompleteResult> {
126 let completion =
127 CompletionInfo::with_all_values(values).map_err(PluginError::invalid_params)?;
128 Ok(CompleteResult::new(completion))
129}
130
131pub fn prompt(
132 name: impl Into<String>,
133 description: impl Into<String>,
134 arguments: Option<Vec<PromptArgument>>,
135) -> Prompt {
136 Prompt::new(name, Some(description.into()), arguments)
137}
138
139pub fn prompt_argument(
140 name: impl Into<String>,
141 description: impl Into<String>,
142 required: bool,
143) -> PromptArgument {
144 PromptArgument::new(name)
145 .with_description(description)
146 .with_required(required)
147}
148
149pub fn text_resource(uri: impl Into<String>, name: impl Into<String>) -> Resource {
150 RawResource::new(uri, name).no_annotation()
151}
152
153pub fn resource_template(
154 uri_template: impl Into<String>,
155 name: impl Into<String>,
156) -> ResourceTemplate {
157 RawResourceTemplate::new(uri_template, name).no_annotation()
158}
159
160pub fn task(
161 task_id: impl Into<String>,
162 status: TaskStatus,
163 created_at: impl Into<String>,
164 last_updated_at: impl Into<String>,
165) -> Task {
166 Task::new(
167 task_id.into(),
168 status,
169 created_at.into(),
170 last_updated_at.into(),
171 )
172}
173
174pub fn get_task_result(task: Task) -> GetTaskResult {
175 GetTaskResult { meta: None, task }
176}
177
178pub fn get_task_payload_result<T: Serialize>(value: T) -> PluginResult<GetTaskPayloadResult> {
179 let value =
180 serde_json::to_value(value).map_err(|err| PluginError::internal(err.to_string()))?;
181 Ok(GetTaskPayloadResult::new(value))
182}
183
184pub fn cancel_task_result(task: Task) -> CancelTaskResult {
185 CancelTaskResult { meta: None, task }
186}
187
188#[allow(deprecated)]
189pub fn plugin_server_info_full(
190 implementation_name: impl Into<String>,
191 implementation_version: impl Into<String>,
192 title: impl Into<String>,
193 description: impl Into<String>,
194 instructions: Option<impl Into<String>>,
195) -> ServerInfo {
196 let info = ServerInfo::new(
197 ServerCapabilities::builder()
198 .enable_tools()
199 .enable_tool_list_changed()
200 .enable_prompts()
201 .enable_prompts_list_changed()
202 .enable_resources()
203 .enable_resources_list_changed()
204 .enable_resources_subscribe()
205 .enable_completions()
206 .enable_tasks()
207 .build(),
208 )
209 .with_server_info(
210 Implementation::new(implementation_name, implementation_version)
211 .with_title(title)
212 .with_description(description),
213 );
214 match instructions {
215 Some(instructions) => info.with_instructions(instructions.into()),
216 None => info,
217 }
218}
219
220pub fn plugin_server_info(
221 implementation_name: impl Into<String>,
222 implementation_version: impl Into<String>,
223 title: impl Into<String>,
224 description: impl Into<String>,
225 instructions: Option<impl Into<String>>,
226) -> ServerInfo {
227 let info = ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
228 .with_server_info(
229 Implementation::new(implementation_name, implementation_version)
230 .with_title(title)
231 .with_description(description),
232 );
233 match instructions {
234 Some(instructions) => info.with_instructions(instructions.into()),
235 None => info,
236 }
237}
238
239pub fn empty_object_schema() -> serde_json::Map<String, serde_json::Value> {
240 serde_json::json!({
241 "type": "object",
242 "additionalProperties": false
243 })
244 .as_object()
245 .cloned()
246 .unwrap()
247}
248
249pub fn json_schema_for<T: JsonSchema>() -> serde_json::Map<String, serde_json::Value> {
250 serde_json::to_value(schemars::schema_for!(T))
251 .ok()
252 .and_then(|value| value.as_object().cloned())
253 .unwrap_or_else(|| {
254 serde_json::json!({
255 "type": "object",
256 "additionalProperties": true
257 })
258 .as_object()
259 .cloned()
260 .unwrap()
261 })
262}
263
264pub fn tool_with_schema(
265 name: impl Into<String>,
266 description: impl Into<String>,
267 schema: serde_json::Map<String, serde_json::Value>,
268) -> Tool {
269 Tool::new(name.into(), description.into(), Arc::new(schema))
270}
271
272pub fn operation_with_schema(
273 name: impl Into<String>,
274 description: impl Into<String>,
275 schema: serde_json::Map<String, serde_json::Value>,
276) -> Tool {
277 tool_with_schema(name, description, schema)
278}
279
280pub fn json_schema_tool<T: JsonSchema>(
281 name: impl Into<String>,
282 description: impl Into<String>,
283) -> Tool {
284 tool_with_schema(name, description, json_schema_for::<T>())
285}
286
287pub fn json_schema_operation<T: JsonSchema>(
288 name: impl Into<String>,
289 description: impl Into<String>,
290) -> Tool {
291 json_schema_tool::<T>(name, description)
292}
293
294pub fn channel_message(
295 channel: impl Into<String>,
296 target_peer_id: impl Into<String>,
297 content_type: impl Into<String>,
298 body: Vec<u8>,
299 message_kind: impl Into<String>,
300) -> proto::ChannelMessage {
301 proto::ChannelMessage {
302 channel: channel.into(),
303 source_peer_id: String::new(),
304 target_peer_id: target_peer_id.into(),
305 content_type: content_type.into(),
306 body,
307 message_kind: message_kind.into(),
308 correlation_id: String::new(),
309 metadata_json: String::new(),
310 }
311}
312
313pub fn json_channel_message<T: Serialize>(
314 channel: impl Into<String>,
315 target_peer_id: impl Into<String>,
316 message_kind: impl Into<String>,
317 payload: &T,
318) -> PluginResult<proto::ChannelMessage> {
319 Ok(channel_message(
320 channel,
321 target_peer_id,
322 "application/json",
323 json_bytes(payload)?,
324 message_kind,
325 ))
326}
327
328pub fn json_reply_channel_message<T: Serialize>(
329 message: &proto::ChannelMessage,
330 message_kind: impl Into<String>,
331 payload: &T,
332) -> PluginResult<proto::ChannelMessage> {
333 let mut reply = json_channel_message(
334 message.channel.clone(),
335 message.source_peer_id.clone(),
336 message_kind,
337 payload,
338 )?;
339 reply.correlation_id = message.correlation_id.clone();
340 Ok(reply)
341}
342
343#[allow(clippy::too_many_arguments)]
344pub fn bulk_transfer_message(
345 kind: i32,
346 channel: impl Into<String>,
347 target_peer_id: impl Into<String>,
348 content_type: impl Into<String>,
349 total_bytes: u64,
350 offset: u64,
351 body: Vec<u8>,
352 final_chunk: bool,
353) -> proto::BulkTransferMessage {
354 proto::BulkTransferMessage {
355 kind,
356 transfer_id: String::new(),
357 channel: channel.into(),
358 source_peer_id: String::new(),
359 target_peer_id: target_peer_id.into(),
360 content_type: content_type.into(),
361 correlation_id: String::new(),
362 metadata_json: String::new(),
363 total_bytes,
364 offset,
365 body,
366 final_chunk,
367 }
368}
369
370pub fn accept_bulk_transfer_message(
371 message: &proto::BulkTransferMessage,
372) -> proto::BulkTransferMessage {
373 let mut response = bulk_transfer_message(
374 proto::bulk_transfer_message::Kind::Accept as i32,
375 message.channel.clone(),
376 message.source_peer_id.clone(),
377 message.content_type.clone(),
378 message.total_bytes,
379 0,
380 Vec::new(),
381 false,
382 );
383 response.transfer_id = message.transfer_id.clone();
384 response.correlation_id = message.correlation_id.clone();
385 response
386}
387
388pub struct BulkTransferSequence {
389 pub transfer_id: String,
390 pub correlation_id: String,
391 pub messages: Vec<proto::BulkTransferMessage>,
392}
393
394#[allow(clippy::too_many_arguments)]
395pub fn bulk_transfer_sequence(
396 channel: impl Into<String>,
397 target_peer_id: impl Into<String>,
398 content_type: impl Into<String>,
399 bytes: Vec<u8>,
400 chunk_size: usize,
401 correlation_id: impl Into<String>,
402 transfer_id: impl Into<String>,
403 metadata_json: impl Into<String>,
404) -> BulkTransferSequence {
405 let channel = channel.into();
406 let target_peer_id = target_peer_id.into();
407 let content_type = content_type.into();
408 let correlation_id = correlation_id.into();
409 let transfer_id = transfer_id.into();
410 let metadata_json = metadata_json.into();
411 let total_bytes = bytes.len() as u64;
412 let chunk_size = chunk_size.max(1);
413
414 let mut messages = Vec::new();
415
416 let mut offer = bulk_transfer_message(
417 proto::bulk_transfer_message::Kind::Offer as i32,
418 channel.clone(),
419 target_peer_id.clone(),
420 content_type.clone(),
421 total_bytes,
422 0,
423 Vec::new(),
424 false,
425 );
426 offer.transfer_id = transfer_id.clone();
427 offer.correlation_id = correlation_id.clone();
428 offer.metadata_json = metadata_json.clone();
429 messages.push(offer);
430
431 let mut offset = 0usize;
432 for chunk in bytes.chunks(chunk_size) {
433 let mut message = bulk_transfer_message(
434 proto::bulk_transfer_message::Kind::Chunk as i32,
435 channel.clone(),
436 target_peer_id.clone(),
437 content_type.clone(),
438 total_bytes,
439 offset as u64,
440 chunk.to_vec(),
441 false,
442 );
443 message.transfer_id = transfer_id.clone();
444 message.correlation_id = correlation_id.clone();
445 message.metadata_json = metadata_json.clone();
446 messages.push(message);
447 offset += chunk.len();
448 }
449
450 let mut complete = bulk_transfer_message(
451 proto::bulk_transfer_message::Kind::Complete as i32,
452 channel,
453 target_peer_id,
454 content_type,
455 total_bytes,
456 total_bytes,
457 Vec::new(),
458 true,
459 );
460 complete.transfer_id = transfer_id.clone();
461 complete.correlation_id = correlation_id.clone();
462 complete.metadata_json = metadata_json;
463 messages.push(complete);
464
465 BulkTransferSequence {
466 transfer_id,
467 correlation_id,
468 messages,
469 }
470}
471
472pub fn json_response<T: Serialize>(value: &T) -> PluginRpcResult {
473 Ok(proto::envelope::Payload::RpcResponse(proto::RpcResponse {
474 result_json: serde_json::to_string(value)
475 .map_err(|err| PluginError::internal(err.to_string()))?,
476 }))
477}
478
479pub fn parse_rpc_params<T: DeserializeOwned>(
480 request: &proto::RpcRequest,
481) -> Result<T, PluginError> {
482 serde_json::from_str(&request.params_json).map_err(|err| {
483 PluginError::invalid_params(format!("Invalid params for '{}': {err}", request.method))
484 })
485}
486
487pub fn parse_tool_call_request(request: &proto::RpcRequest) -> PluginResult<ToolCallRequest> {
488 parse_rpc_params(request)
489}
490
491pub fn parse_optional_json(raw: &str) -> Option<serde_json::Value> {
492 if raw.trim().is_empty() {
493 None
494 } else {
495 serde_json::from_str(raw).ok()
496 }
497}
498
499pub fn parse_get_prompt_request(
500 request: &proto::RpcRequest,
501) -> PluginResult<GetPromptRequestParams> {
502 parse_rpc_params(request)
503}
504
505pub fn parse_read_resource_request(
506 request: &proto::RpcRequest,
507) -> PluginResult<ReadResourceRequestParams> {
508 parse_rpc_params(request)
509}
510
511pub type ToolFuture<'a> = Pin<Box<dyn Future<Output = PluginResult<CallToolResult>> + Send + 'a>>;
512pub type JsonToolFuture<'a, T> = Pin<Box<dyn Future<Output = PluginResult<T>> + Send + 'a>>;
513pub type OperationFuture<'a> = ToolFuture<'a>;
514pub type JsonOperationFuture<'a, T> = JsonToolFuture<'a, T>;
515pub type PromptFuture<'a> =
516 Pin<Box<dyn Future<Output = PluginResult<GetPromptResult>> + Send + 'a>>;
517pub type ResourceFuture<'a> =
518 Pin<Box<dyn Future<Output = PluginResult<ReadResourceResult>> + Send + 'a>>;
519pub type CompletionFuture<'a> =
520 Pin<Box<dyn Future<Output = PluginResult<CompleteResult>> + Send + 'a>>;
521pub type TaskListFuture<'a> =
522 Pin<Box<dyn Future<Output = PluginResult<ListTasksResult>> + Send + 'a>>;
523pub type TaskInfoFuture<'a> =
524 Pin<Box<dyn Future<Output = PluginResult<GetTaskResult>> + Send + 'a>>;
525pub type TaskResultFuture<'a> =
526 Pin<Box<dyn Future<Output = PluginResult<GetTaskPayloadResult>> + Send + 'a>>;
527pub type TaskCancelFuture<'a> =
528 Pin<Box<dyn Future<Output = PluginResult<CancelTaskResult>> + Send + 'a>>;
529
530type ToolHandler = Arc<
531 dyn for<'a, 'ctx> Fn(ToolCallRequest, &'a mut PluginContext<'ctx>) -> ToolFuture<'a>
532 + Send
533 + Sync,
534>;
535
536#[derive(Clone)]
537pub struct ToolRouter {
538 tools: Vec<Tool>,
539 handlers: HashMap<String, ToolHandler>,
540}
541
542pub type OperationRouter = ToolRouter;
543
544impl ToolRouter {
545 pub fn new() -> Self {
546 Self {
547 tools: Vec::new(),
548 handlers: HashMap::new(),
549 }
550 }
551
552 pub fn add_raw<F>(&mut self, tool: Tool, handler: F)
553 where
554 F: for<'a, 'ctx> Fn(ToolCallRequest, &'a mut PluginContext<'ctx>) -> ToolFuture<'a>
555 + Send
556 + Sync
557 + 'static,
558 {
559 let name = tool.name.to_string();
560 self.tools.push(tool);
561 self.handlers.insert(name, Arc::new(handler));
562 }
563
564 pub fn extend(&mut self, mut other: Self) {
565 self.tools.append(&mut other.tools);
566 self.handlers.extend(other.handlers.drain());
567 }
568
569 pub fn add_json<TArgs, TResult, F>(&mut self, tool: Tool, handler: F)
570 where
571 TArgs: DeserializeOwned + Send + 'static,
572 TResult: Serialize + Send + 'static,
573 F: for<'a, 'ctx> Fn(TArgs, &'a mut PluginContext<'ctx>) -> JsonToolFuture<'a, TResult>
574 + Send
575 + Sync
576 + 'static,
577 {
578 let handler = Arc::new(handler);
579 self.add_raw(tool, move |request, context| {
580 let handler = Arc::clone(&handler);
581 Box::pin(async move {
582 let args: TArgs = request.arguments()?;
583 let value = handler(args, context).await?;
584 structured_tool_result(value)
585 })
586 });
587 }
588
589 pub fn add_json_default<TArgs, TResult, F>(&mut self, tool: Tool, handler: F)
590 where
591 TArgs: DeserializeOwned + Default + Send + 'static,
592 TResult: Serialize + Send + 'static,
593 F: for<'a, 'ctx> Fn(TArgs, &'a mut PluginContext<'ctx>) -> JsonToolFuture<'a, TResult>
594 + Send
595 + Sync
596 + 'static,
597 {
598 let handler = Arc::new(handler);
599 self.add_raw(tool, move |request, context| {
600 let handler = Arc::clone(&handler);
601 Box::pin(async move {
602 let args: TArgs = request.arguments_or_default()?;
603 let value = handler(args, context).await?;
604 structured_tool_result(value)
605 })
606 });
607 }
608
609 pub fn list_tools_result(&self) -> ListToolsResult {
610 list_tools(self.tools.clone())
611 }
612
613 pub async fn call(
614 &self,
615 request: ToolCallRequest,
616 context: &mut PluginContext<'_>,
617 ) -> PluginResult<CallToolResult> {
618 let Some(handler) = self.handlers.get(&request.name).cloned() else {
619 return Err(PluginError::method_not_found(format!(
620 "Unknown tool '{}'",
621 request.name
622 )));
623 };
624 handler(request, context).await
625 }
626}
627
628impl Default for ToolRouter {
629 fn default() -> Self {
630 Self::new()
631 }
632}
633
634type PromptHandler = Arc<
635 dyn for<'a, 'ctx> Fn(GetPromptRequestParams, &'a mut PluginContext<'ctx>) -> PromptFuture<'a>
636 + Send
637 + Sync,
638>;
639
640#[derive(Clone)]
641pub struct PromptRouter {
642 prompts: Vec<Prompt>,
643 handlers: HashMap<String, PromptHandler>,
644}
645
646impl PromptRouter {
647 pub fn new() -> Self {
648 Self {
649 prompts: Vec::new(),
650 handlers: HashMap::new(),
651 }
652 }
653
654 pub fn add<F>(&mut self, prompt: Prompt, handler: F)
655 where
656 F: for<'a, 'ctx> Fn(
657 GetPromptRequestParams,
658 &'a mut PluginContext<'ctx>,
659 ) -> PromptFuture<'a>
660 + Send
661 + Sync
662 + 'static,
663 {
664 let name = prompt.name.to_string();
665 self.prompts.push(prompt);
666 self.handlers.insert(name, Arc::new(handler));
667 }
668
669 pub fn list_prompts_result(&self) -> ListPromptsResult {
670 list_prompts(self.prompts.clone())
671 }
672
673 pub async fn get(
674 &self,
675 request: GetPromptRequestParams,
676 context: &mut PluginContext<'_>,
677 ) -> PluginResult<GetPromptResult> {
678 let Some(handler) = self.handlers.get(&request.name).cloned() else {
679 return Err(PluginError::invalid_params(format!(
680 "Unknown prompt '{}'",
681 request.name
682 )));
683 };
684 handler(request, context).await
685 }
686}
687
688impl Default for PromptRouter {
689 fn default() -> Self {
690 Self::new()
691 }
692}
693
694#[derive(Clone)]
695enum ResourceReadMatcher {
696 Exact(String),
697 Prefix(String),
698}
699
700impl ResourceReadMatcher {
701 fn matches(&self, uri: &str) -> bool {
702 match self {
703 Self::Exact(expected) => uri == expected,
704 Self::Prefix(prefix) => uri.starts_with(prefix),
705 }
706 }
707}
708
709type ResourceHandler = Arc<
710 dyn for<'a, 'ctx> Fn(
711 ReadResourceRequestParams,
712 &'a mut PluginContext<'ctx>,
713 ) -> ResourceFuture<'a>
714 + Send
715 + Sync,
716>;
717
718#[derive(Clone)]
719pub struct ResourceRouter {
720 resources: Vec<Resource>,
721 resource_templates: Vec<ResourceTemplate>,
722 handlers: Vec<(ResourceReadMatcher, ResourceHandler)>,
723}
724
725impl ResourceRouter {
726 pub fn new() -> Self {
727 Self {
728 resources: Vec::new(),
729 resource_templates: Vec::new(),
730 handlers: Vec::new(),
731 }
732 }
733
734 pub fn add_exact<F>(&mut self, resource: Resource, handler: F)
735 where
736 F: for<'a, 'ctx> Fn(
737 ReadResourceRequestParams,
738 &'a mut PluginContext<'ctx>,
739 ) -> ResourceFuture<'a>
740 + Send
741 + Sync
742 + 'static,
743 {
744 let uri = resource.raw.uri.to_string();
745 self.resources.push(resource);
746 self.handlers
747 .push((ResourceReadMatcher::Exact(uri), Arc::new(handler)));
748 }
749
750 pub fn add_prefix_template<F>(
751 &mut self,
752 resource_template: ResourceTemplate,
753 prefix: impl Into<String>,
754 handler: F,
755 ) where
756 F: for<'a, 'ctx> Fn(
757 ReadResourceRequestParams,
758 &'a mut PluginContext<'ctx>,
759 ) -> ResourceFuture<'a>
760 + Send
761 + Sync
762 + 'static,
763 {
764 self.resource_templates.push(resource_template);
765 self.handlers.push((
766 ResourceReadMatcher::Prefix(prefix.into()),
767 Arc::new(handler),
768 ));
769 }
770
771 pub fn list_resources_result(&self) -> ListResourcesResult {
772 list_resources(self.resources.clone())
773 }
774
775 pub fn list_resource_templates_result(&self) -> ListResourceTemplatesResult {
776 list_resource_templates(self.resource_templates.clone())
777 }
778
779 pub async fn read(
780 &self,
781 request: ReadResourceRequestParams,
782 context: &mut PluginContext<'_>,
783 ) -> PluginResult<ReadResourceResult> {
784 let Some((_, handler)) = self
785 .handlers
786 .iter()
787 .find(|(matcher, _)| matcher.matches(&request.uri))
788 else {
789 return Err(PluginError::invalid_params(format!(
790 "Unknown resource '{}'",
791 request.uri
792 )));
793 };
794 handler(request, context).await
795 }
796}
797
798impl Default for ResourceRouter {
799 fn default() -> Self {
800 Self::new()
801 }
802}
803
804#[derive(Clone)]
805enum CompletionMatcher {
806 PromptArgument {
807 prompt_name: String,
808 argument_name: Option<String>,
809 },
810 ResourceArgument {
811 resource_uri: String,
812 argument_name: Option<String>,
813 },
814}
815
816impl CompletionMatcher {
817 fn matches(&self, request: &CompleteRequestParams) -> bool {
818 match self {
819 Self::PromptArgument {
820 prompt_name,
821 argument_name,
822 } => {
823 request.r#ref.as_prompt_name() == Some(prompt_name.as_str())
824 && argument_name
825 .as_ref()
826 .map(|name| request.argument.name == *name)
827 .unwrap_or(true)
828 }
829 Self::ResourceArgument {
830 resource_uri,
831 argument_name,
832 } => {
833 request.r#ref.as_resource_uri() == Some(resource_uri.as_str())
834 && argument_name
835 .as_ref()
836 .map(|name| request.argument.name == *name)
837 .unwrap_or(true)
838 }
839 }
840 }
841}
842
843type CompletionHandler = Arc<
844 dyn for<'a, 'ctx> Fn(CompleteRequestParams, &'a mut PluginContext<'ctx>) -> CompletionFuture<'a>
845 + Send
846 + Sync,
847>;
848
849#[derive(Clone)]
850pub struct CompletionRouter {
851 handlers: Vec<(CompletionMatcher, CompletionHandler)>,
852}
853
854impl CompletionRouter {
855 pub fn new() -> Self {
856 Self {
857 handlers: Vec::new(),
858 }
859 }
860
861 pub fn add_prompt_argument_values(
862 &mut self,
863 prompt_name: impl Into<String>,
864 argument_name: impl Into<String>,
865 values: Vec<String>,
866 ) {
867 let values = Arc::new(values);
868 self.add_prompt_argument(prompt_name, argument_name, move |_request, _context| {
869 let values = values.clone();
870 Box::pin(async move { complete_result(values.as_ref().clone()) })
871 });
872 }
873
874 pub fn add_resource_argument_values(
875 &mut self,
876 resource_uri: impl Into<String>,
877 argument_name: impl Into<String>,
878 values: Vec<String>,
879 ) {
880 let values = Arc::new(values);
881 self.add_resource_argument(resource_uri, argument_name, move |_request, _context| {
882 let values = values.clone();
883 Box::pin(async move { complete_result(values.as_ref().clone()) })
884 });
885 }
886
887 pub fn add_prompt_argument<F>(
888 &mut self,
889 prompt_name: impl Into<String>,
890 argument_name: impl Into<String>,
891 handler: F,
892 ) where
893 F: for<'a, 'ctx> Fn(
894 CompleteRequestParams,
895 &'a mut PluginContext<'ctx>,
896 ) -> CompletionFuture<'a>
897 + Send
898 + Sync
899 + 'static,
900 {
901 self.handlers.push((
902 CompletionMatcher::PromptArgument {
903 prompt_name: prompt_name.into(),
904 argument_name: Some(argument_name.into()),
905 },
906 Arc::new(handler),
907 ));
908 }
909
910 pub fn add_prompt<F>(&mut self, prompt_name: impl Into<String>, handler: F)
911 where
912 F: for<'a, 'ctx> Fn(
913 CompleteRequestParams,
914 &'a mut PluginContext<'ctx>,
915 ) -> CompletionFuture<'a>
916 + Send
917 + Sync
918 + 'static,
919 {
920 self.handlers.push((
921 CompletionMatcher::PromptArgument {
922 prompt_name: prompt_name.into(),
923 argument_name: None,
924 },
925 Arc::new(handler),
926 ));
927 }
928
929 pub fn add_resource_argument<F>(
930 &mut self,
931 resource_uri: impl Into<String>,
932 argument_name: impl Into<String>,
933 handler: F,
934 ) where
935 F: for<'a, 'ctx> Fn(
936 CompleteRequestParams,
937 &'a mut PluginContext<'ctx>,
938 ) -> CompletionFuture<'a>
939 + Send
940 + Sync
941 + 'static,
942 {
943 self.handlers.push((
944 CompletionMatcher::ResourceArgument {
945 resource_uri: resource_uri.into(),
946 argument_name: Some(argument_name.into()),
947 },
948 Arc::new(handler),
949 ));
950 }
951
952 pub fn add_resource<F>(&mut self, resource_uri: impl Into<String>, handler: F)
953 where
954 F: for<'a, 'ctx> Fn(
955 CompleteRequestParams,
956 &'a mut PluginContext<'ctx>,
957 ) -> CompletionFuture<'a>
958 + Send
959 + Sync
960 + 'static,
961 {
962 self.handlers.push((
963 CompletionMatcher::ResourceArgument {
964 resource_uri: resource_uri.into(),
965 argument_name: None,
966 },
967 Arc::new(handler),
968 ));
969 }
970
971 pub async fn complete(
972 &self,
973 request: CompleteRequestParams,
974 context: &mut PluginContext<'_>,
975 ) -> PluginResult<CompleteResult> {
976 let Some((_, handler)) = self
977 .handlers
978 .iter()
979 .find(|(matcher, _)| matcher.matches(&request))
980 else {
981 return complete_result(vec![request.argument.value]);
982 };
983 handler(request, context).await
984 }
985}
986
987impl Default for CompletionRouter {
988 fn default() -> Self {
989 Self::new()
990 }
991}
992
993type TaskListHandler = Arc<
994 dyn for<'a, 'ctx> Fn(
995 Option<PaginatedRequestParams>,
996 &'a mut PluginContext<'ctx>,
997 ) -> TaskListFuture<'a>
998 + Send
999 + Sync,
1000>;
1001type TaskInfoHandler = Arc<
1002 dyn for<'a, 'ctx> Fn(GetTaskInfoParams, &'a mut PluginContext<'ctx>) -> TaskInfoFuture<'a>
1003 + Send
1004 + Sync,
1005>;
1006type TaskResultHandler = Arc<
1007 dyn for<'a, 'ctx> Fn(GetTaskResultParams, &'a mut PluginContext<'ctx>) -> TaskResultFuture<'a>
1008 + Send
1009 + Sync,
1010>;
1011type TaskCancelHandler = Arc<
1012 dyn for<'a, 'ctx> Fn(CancelTaskParams, &'a mut PluginContext<'ctx>) -> TaskCancelFuture<'a>
1013 + Send
1014 + Sync,
1015>;
1016
1017#[derive(Clone)]
1018pub struct TaskRouter {
1019 list_handler: Option<TaskListHandler>,
1020 info_handler: Option<TaskInfoHandler>,
1021 result_handler: Option<TaskResultHandler>,
1022 cancel_handler: Option<TaskCancelHandler>,
1023}
1024
1025impl TaskRouter {
1026 pub fn new() -> Self {
1027 Self {
1028 list_handler: None,
1029 info_handler: None,
1030 result_handler: None,
1031 cancel_handler: None,
1032 }
1033 }
1034
1035 pub fn with_list<F>(mut self, handler: F) -> Self
1036 where
1037 F: for<'a, 'ctx> Fn(
1038 Option<PaginatedRequestParams>,
1039 &'a mut PluginContext<'ctx>,
1040 ) -> TaskListFuture<'a>
1041 + Send
1042 + Sync
1043 + 'static,
1044 {
1045 self.list_handler = Some(Arc::new(handler));
1046 self
1047 }
1048
1049 pub fn with_get_info<F>(mut self, handler: F) -> Self
1050 where
1051 F: for<'a, 'ctx> Fn(GetTaskInfoParams, &'a mut PluginContext<'ctx>) -> TaskInfoFuture<'a>
1052 + Send
1053 + Sync
1054 + 'static,
1055 {
1056 self.info_handler = Some(Arc::new(handler));
1057 self
1058 }
1059
1060 pub fn with_get_result<F>(mut self, handler: F) -> Self
1061 where
1062 F: for<'a, 'ctx> Fn(
1063 GetTaskResultParams,
1064 &'a mut PluginContext<'ctx>,
1065 ) -> TaskResultFuture<'a>
1066 + Send
1067 + Sync
1068 + 'static,
1069 {
1070 self.result_handler = Some(Arc::new(handler));
1071 self
1072 }
1073
1074 pub fn with_cancel<F>(mut self, handler: F) -> Self
1075 where
1076 F: for<'a, 'ctx> Fn(CancelTaskParams, &'a mut PluginContext<'ctx>) -> TaskCancelFuture<'a>
1077 + Send
1078 + Sync
1079 + 'static,
1080 {
1081 self.cancel_handler = Some(Arc::new(handler));
1082 self
1083 }
1084
1085 pub async fn list_tasks(
1086 &self,
1087 request: Option<PaginatedRequestParams>,
1088 context: &mut PluginContext<'_>,
1089 ) -> PluginResult<Option<ListTasksResult>> {
1090 match &self.list_handler {
1091 Some(handler) => Ok(Some(handler(request, context).await?)),
1092 None => Ok(None),
1093 }
1094 }
1095
1096 pub async fn get_task_info(
1097 &self,
1098 request: GetTaskInfoParams,
1099 context: &mut PluginContext<'_>,
1100 ) -> PluginResult<Option<GetTaskResult>> {
1101 match &self.info_handler {
1102 Some(handler) => Ok(Some(handler(request, context).await?)),
1103 None => Ok(None),
1104 }
1105 }
1106
1107 pub async fn get_task_result(
1108 &self,
1109 request: GetTaskResultParams,
1110 context: &mut PluginContext<'_>,
1111 ) -> PluginResult<Option<GetTaskPayloadResult>> {
1112 match &self.result_handler {
1113 Some(handler) => Ok(Some(handler(request, context).await?)),
1114 None => Ok(None),
1115 }
1116 }
1117
1118 pub async fn cancel_task(
1119 &self,
1120 request: CancelTaskParams,
1121 context: &mut PluginContext<'_>,
1122 ) -> PluginResult<Option<CancelTaskResult>> {
1123 match &self.cancel_handler {
1124 Some(handler) => Ok(Some(handler(request, context).await?)),
1125 None => Ok(None),
1126 }
1127 }
1128}
1129
1130impl Default for TaskRouter {
1131 fn default() -> Self {
1132 Self::new()
1133 }
1134}
1135
1136#[derive(Clone, Debug, Default)]
1137pub struct SubscriptionSet {
1138 uris: BTreeSet<String>,
1139}
1140
1141impl SubscriptionSet {
1142 pub fn subscribe(&mut self, uri: impl Into<String>) {
1143 self.uris.insert(uri.into());
1144 }
1145
1146 pub fn unsubscribe(&mut self, uri: &str) {
1147 self.uris.remove(uri);
1148 }
1149
1150 pub fn list(&self) -> Vec<String> {
1151 self.uris.iter().cloned().collect()
1152 }
1153}
1154
1155#[derive(Clone, Debug)]
1156pub struct TaskRecord<T> {
1157 pub task: Task,
1158 pub payload: T,
1159}
1160
1161#[derive(Clone, Debug, Default)]
1162pub struct TaskStore<T> {
1163 tasks: BTreeMap<String, TaskRecord<T>>,
1164}
1165
1166impl<T> TaskStore<T> {
1167 pub fn insert(&mut self, task: Task, payload: T) {
1168 self.tasks
1169 .insert(task.task_id.clone(), TaskRecord { task, payload });
1170 }
1171
1172 pub fn list(&self) -> Vec<Task> {
1173 self.tasks.values().map(|task| task.task.clone()).collect()
1174 }
1175
1176 pub fn get(&self, task_id: &str) -> PluginResult<&TaskRecord<T>> {
1177 self.tasks
1178 .get(task_id)
1179 .ok_or_else(|| PluginError::invalid_params(format!("Unknown task '{task_id}'")))
1180 }
1181
1182 pub fn get_mut(&mut self, task_id: &str) -> PluginResult<&mut TaskRecord<T>> {
1183 self.tasks
1184 .get_mut(task_id)
1185 .ok_or_else(|| PluginError::invalid_params(format!("Unknown task '{task_id}'")))
1186 }
1187
1188 pub fn values(&self) -> impl Iterator<Item = &TaskRecord<T>> {
1189 self.tasks.values()
1190 }
1191}