1use anyhow::Result;
2use rmcp::model::{
3 CallToolResult, CancelTaskParams, CancelTaskResult, CompleteRequestParams, CompleteResult,
4 GetPromptRequestParams, GetPromptResult, GetTaskInfoParams, GetTaskPayloadResult,
5 GetTaskResult, GetTaskResultParams, ListPromptsResult, ListResourceTemplatesResult,
6 ListResourcesResult, ListTasksResult, ListToolsResult, PaginatedRequestParams,
7 ReadResourceRequestParams, ReadResourceResult, ServerInfo, SetLevelRequestParams,
8 SubscribeRequestParams, UnsubscribeRequestParams,
9};
10use std::sync::Arc;
11
12use crate::{
13 context::PluginContext,
14 error::{PluginError, PluginResult},
15 helpers::{
16 CompletionRouter, PromptRouter, ResourceRouter, TaskRouter, ToolCallRequest, ToolRouter,
17 },
18 proto,
19 runtime::{
20 BulkHandler, CancelStreamHandler, ChannelHandler, CloseStreamHandler, HealthFuture,
21 HealthHandler, InitFuture, InitHandler, InitializeFuture, InitializeHandler,
22 MeshEventHandler, MeshVisibility, OpenStreamFuture, OpenStreamHandler, Plugin,
23 PluginInitializeRequest, PluginMetadata, PluginStartupPolicy, SetLogLevelFuture,
24 SetLogLevelHandler, StreamErrorHandler, SubscribeFuture, SubscribeHandler,
25 UnsubscribeHandler,
26 },
27};
28
29#[derive(Clone)]
30pub struct SimplePlugin {
31 metadata: PluginMetadata,
32 operation_router: Option<ToolRouter>,
33 prompt_router: Option<PromptRouter>,
34 resource_router: Option<ResourceRouter>,
35 completion_router: Option<CompletionRouter>,
36 task_router: Option<TaskRouter>,
37 initialize_handler: Option<InitializeHandler>,
38 on_initialized: Option<InitHandler>,
39 health_handler: Option<HealthHandler>,
40 subscribe_handler: Option<SubscribeHandler>,
41 unsubscribe_handler: Option<UnsubscribeHandler>,
42 set_log_level_handler: Option<SetLogLevelHandler>,
43 channel_handler: Option<ChannelHandler>,
44 bulk_handler: Option<BulkHandler>,
45 mesh_event_handler: Option<MeshEventHandler>,
46 open_stream_handler: Option<OpenStreamHandler>,
47 cancel_stream_handler: Option<CancelStreamHandler>,
48 close_stream_handler: Option<CloseStreamHandler>,
49 stream_error_handler: Option<StreamErrorHandler>,
50}
51
52impl SimplePlugin {
53 pub fn new(metadata: PluginMetadata) -> Self {
54 Self {
55 metadata,
56 operation_router: None,
57 prompt_router: None,
58 resource_router: None,
59 completion_router: None,
60 task_router: None,
61 initialize_handler: None,
62 on_initialized: None,
63 health_handler: None,
64 subscribe_handler: None,
65 unsubscribe_handler: None,
66 set_log_level_handler: None,
67 channel_handler: None,
68 bulk_handler: None,
69 mesh_event_handler: None,
70 open_stream_handler: None,
71 cancel_stream_handler: None,
72 close_stream_handler: None,
73 stream_error_handler: None,
74 }
75 }
76
77 pub fn with_capabilities(mut self, capabilities: Vec<String>) -> Self {
78 self.metadata = self.metadata.with_capabilities(capabilities);
79 self
80 }
81
82 pub fn with_manifest(mut self, manifest: proto::PluginManifest) -> Self {
83 self.metadata = self.metadata.with_manifest(manifest);
84 self
85 }
86
87 pub fn with_startup_policy(mut self, startup_policy: PluginStartupPolicy) -> Self {
88 self.metadata = self.metadata.with_startup_policy(startup_policy);
89 self
90 }
91
92 pub fn with_operation_router(mut self, router: ToolRouter) -> Self {
93 self.operation_router = Some(router);
94 self
95 }
96
97 pub fn extend_operation_router(mut self, router: ToolRouter) -> Self {
98 match &mut self.operation_router {
99 Some(existing) => existing.extend(router),
100 None => self.operation_router = Some(router),
101 }
102 self
103 }
104
105 pub fn with_prompt_router(mut self, router: PromptRouter) -> Self {
106 self.prompt_router = Some(router);
107 self
108 }
109
110 pub fn with_resource_router(mut self, router: ResourceRouter) -> Self {
111 self.resource_router = Some(router);
112 self
113 }
114
115 pub fn with_completion_router(mut self, router: CompletionRouter) -> Self {
116 self.completion_router = Some(router);
117 self
118 }
119
120 pub fn with_task_router(mut self, router: TaskRouter) -> Self {
121 self.task_router = Some(router);
122 self
123 }
124
125 pub fn on_initialize<F>(mut self, handler: F) -> Self
126 where
127 F: for<'a, 'ctx> Fn(
128 PluginInitializeRequest,
129 &'a mut PluginContext<'ctx>,
130 ) -> InitializeFuture<'a>
131 + Send
132 + Sync
133 + 'static,
134 {
135 self.initialize_handler = Some(Arc::new(handler));
136 self
137 }
138
139 pub fn on_initialized<F>(mut self, handler: F) -> Self
140 where
141 F: for<'a, 'ctx> Fn(&'a mut PluginContext<'ctx>) -> InitFuture<'a> + Send + Sync + 'static,
142 {
143 self.on_initialized = Some(Arc::new(handler));
144 self
145 }
146
147 pub fn with_health<F>(mut self, handler: F) -> Self
148 where
149 F: for<'a, 'ctx> Fn(&'a mut PluginContext<'ctx>) -> HealthFuture<'a>
150 + Send
151 + Sync
152 + 'static,
153 {
154 self.health_handler = Some(Arc::new(handler));
155 self
156 }
157
158 pub fn with_subscribe_resource<F>(mut self, handler: F) -> Self
159 where
160 F: for<'a, 'ctx> Fn(
161 SubscribeRequestParams,
162 &'a mut PluginContext<'ctx>,
163 ) -> SubscribeFuture<'a>
164 + Send
165 + Sync
166 + 'static,
167 {
168 self.subscribe_handler = Some(Arc::new(handler));
169 self
170 }
171
172 pub fn with_unsubscribe_resource<F>(mut self, handler: F) -> Self
173 where
174 F: for<'a, 'ctx> Fn(
175 UnsubscribeRequestParams,
176 &'a mut PluginContext<'ctx>,
177 ) -> SubscribeFuture<'a>
178 + Send
179 + Sync
180 + 'static,
181 {
182 self.unsubscribe_handler = Some(Arc::new(handler));
183 self
184 }
185
186 pub fn with_set_log_level<F>(mut self, handler: F) -> Self
187 where
188 F: for<'a, 'ctx> Fn(
189 SetLevelRequestParams,
190 &'a mut PluginContext<'ctx>,
191 ) -> SetLogLevelFuture<'a>
192 + Send
193 + Sync
194 + 'static,
195 {
196 self.set_log_level_handler = Some(Arc::new(handler));
197 self
198 }
199
200 pub fn on_channel_message<F>(mut self, handler: F) -> Self
201 where
202 F: for<'a, 'ctx> Fn(proto::ChannelMessage, &'a mut PluginContext<'ctx>) -> InitFuture<'a>
203 + Send
204 + Sync
205 + 'static,
206 {
207 self.channel_handler = Some(Arc::new(handler));
208 self
209 }
210
211 pub fn on_bulk_transfer_message<F>(mut self, handler: F) -> Self
212 where
213 F: for<'a, 'ctx> Fn(
214 proto::BulkTransferMessage,
215 &'a mut PluginContext<'ctx>,
216 ) -> InitFuture<'a>
217 + Send
218 + Sync
219 + 'static,
220 {
221 self.bulk_handler = Some(Arc::new(handler));
222 self
223 }
224
225 pub fn on_mesh_event<F>(mut self, handler: F) -> Self
226 where
227 F: for<'a, 'ctx> Fn(proto::MeshEvent, &'a mut PluginContext<'ctx>) -> InitFuture<'a>
228 + Send
229 + Sync
230 + 'static,
231 {
232 self.mesh_event_handler = Some(Arc::new(handler));
233 self
234 }
235
236 pub fn on_open_stream<F>(mut self, handler: F) -> Self
237 where
238 F: for<'a, 'ctx> Fn(
239 proto::OpenStreamRequest,
240 &'a mut PluginContext<'ctx>,
241 ) -> OpenStreamFuture<'a>
242 + Send
243 + Sync
244 + 'static,
245 {
246 self.open_stream_handler = Some(Arc::new(handler));
247 self
248 }
249
250 pub fn on_cancel_stream<F>(mut self, handler: F) -> Self
251 where
252 F: for<'a, 'ctx> Fn(
253 proto::CancelStreamNotification,
254 &'a mut PluginContext<'ctx>,
255 ) -> InitFuture<'a>
256 + Send
257 + Sync
258 + 'static,
259 {
260 self.cancel_stream_handler = Some(Arc::new(handler));
261 self
262 }
263
264 pub fn on_close_stream<F>(mut self, handler: F) -> Self
265 where
266 F: for<'a, 'ctx> Fn(
267 proto::CloseStreamNotification,
268 &'a mut PluginContext<'ctx>,
269 ) -> InitFuture<'a>
270 + Send
271 + Sync
272 + 'static,
273 {
274 self.close_stream_handler = Some(Arc::new(handler));
275 self
276 }
277
278 pub fn on_stream_error<F>(mut self, handler: F) -> Self
279 where
280 F: for<'a, 'ctx> Fn(proto::StreamError, &'a mut PluginContext<'ctx>) -> InitFuture<'a>
281 + Send
282 + Sync
283 + 'static,
284 {
285 self.stream_error_handler = Some(Arc::new(handler));
286 self
287 }
288}
289
290#[crate::async_trait]
291impl Plugin for SimplePlugin {
292 fn plugin_id(&self) -> &str {
293 &self.metadata.plugin_id
294 }
295
296 fn plugin_version(&self) -> String {
297 self.metadata.plugin_version.clone()
298 }
299
300 fn server_info(&self) -> ServerInfo {
301 self.metadata.server_info.clone()
302 }
303
304 fn capabilities(&self) -> Vec<String> {
305 self.metadata.capabilities.clone()
306 }
307
308 fn manifest(&self) -> Option<proto::PluginManifest> {
309 self.metadata.manifest.clone()
310 }
311
312 async fn initialize(
313 &mut self,
314 request: PluginInitializeRequest,
315 context: &mut PluginContext<'_>,
316 ) -> PluginResult<()> {
317 match self.metadata.startup_policy {
318 PluginStartupPolicy::Any => {}
319 PluginStartupPolicy::PrivateMeshOnly
320 if request.mesh_visibility != MeshVisibility::Private =>
321 {
322 return Err(PluginError::startup_disabled(format!(
323 "Plugin '{}' requires a private mesh",
324 self.metadata.plugin_id
325 )));
326 }
327 PluginStartupPolicy::PublicMeshOnly
328 if request.mesh_visibility != MeshVisibility::Public =>
329 {
330 return Err(PluginError::startup_disabled(format!(
331 "Plugin '{}' requires a public mesh",
332 self.metadata.plugin_id
333 )));
334 }
335 PluginStartupPolicy::PrivateMeshOnly | PluginStartupPolicy::PublicMeshOnly => {}
336 }
337 match &self.initialize_handler {
338 Some(handler) => handler(request, context).await,
339 None => Ok(()),
340 }
341 }
342
343 async fn on_initialized(&mut self, context: &mut PluginContext<'_>) -> Result<()> {
344 match &self.on_initialized {
345 Some(handler) => handler(context).await,
346 None => Ok(()),
347 }
348 }
349
350 async fn health(&mut self, context: &mut PluginContext<'_>) -> Result<String> {
351 match &self.health_handler {
352 Some(handler) => handler(context).await,
353 None => Ok("ok".into()),
354 }
355 }
356
357 async fn list_tools(
358 &mut self,
359 _context: &mut PluginContext<'_>,
360 ) -> PluginResult<Option<ListToolsResult>> {
361 Ok(self
362 .operation_router
363 .as_ref()
364 .map(|router| router.list_tools_result()))
365 }
366
367 async fn call_tool(
368 &mut self,
369 request: ToolCallRequest,
370 context: &mut PluginContext<'_>,
371 ) -> PluginResult<Option<CallToolResult>> {
372 match &self.operation_router {
373 Some(router) => Ok(Some(router.call(request, context).await?)),
374 None => Ok(None),
375 }
376 }
377
378 async fn list_prompts(
379 &mut self,
380 _request: Option<PaginatedRequestParams>,
381 _context: &mut PluginContext<'_>,
382 ) -> PluginResult<Option<ListPromptsResult>> {
383 Ok(self
384 .prompt_router
385 .as_ref()
386 .map(|router| router.list_prompts_result()))
387 }
388
389 async fn get_prompt(
390 &mut self,
391 request: GetPromptRequestParams,
392 context: &mut PluginContext<'_>,
393 ) -> PluginResult<Option<GetPromptResult>> {
394 match &self.prompt_router {
395 Some(router) => Ok(Some(router.get(request, context).await?)),
396 None => Ok(None),
397 }
398 }
399
400 async fn list_resources(
401 &mut self,
402 _request: Option<PaginatedRequestParams>,
403 _context: &mut PluginContext<'_>,
404 ) -> PluginResult<Option<ListResourcesResult>> {
405 Ok(self
406 .resource_router
407 .as_ref()
408 .map(|router| router.list_resources_result()))
409 }
410
411 async fn read_resource(
412 &mut self,
413 request: ReadResourceRequestParams,
414 context: &mut PluginContext<'_>,
415 ) -> PluginResult<Option<ReadResourceResult>> {
416 match &self.resource_router {
417 Some(router) => Ok(Some(router.read(request, context).await?)),
418 None => Ok(None),
419 }
420 }
421
422 async fn list_resource_templates(
423 &mut self,
424 _request: Option<PaginatedRequestParams>,
425 _context: &mut PluginContext<'_>,
426 ) -> PluginResult<Option<ListResourceTemplatesResult>> {
427 Ok(self
428 .resource_router
429 .as_ref()
430 .map(|router| router.list_resource_templates_result()))
431 }
432
433 async fn subscribe_resource(
434 &mut self,
435 request: SubscribeRequestParams,
436 context: &mut PluginContext<'_>,
437 ) -> PluginResult<Option<()>> {
438 match &self.subscribe_handler {
439 Some(handler) => Ok(Some(handler(request, context).await?)),
440 None => Ok(None),
441 }
442 }
443
444 async fn unsubscribe_resource(
445 &mut self,
446 request: UnsubscribeRequestParams,
447 context: &mut PluginContext<'_>,
448 ) -> PluginResult<Option<()>> {
449 match &self.unsubscribe_handler {
450 Some(handler) => Ok(Some(handler(request, context).await?)),
451 None => Ok(None),
452 }
453 }
454
455 async fn complete(
456 &mut self,
457 request: CompleteRequestParams,
458 context: &mut PluginContext<'_>,
459 ) -> PluginResult<Option<CompleteResult>> {
460 match &self.completion_router {
461 Some(router) => Ok(Some(router.complete(request, context).await?)),
462 None => Ok(None),
463 }
464 }
465
466 async fn set_log_level(
467 &mut self,
468 request: SetLevelRequestParams,
469 context: &mut PluginContext<'_>,
470 ) -> PluginResult<Option<()>> {
471 match &self.set_log_level_handler {
472 Some(handler) => Ok(Some(handler(request, context).await?)),
473 None => Ok(None),
474 }
475 }
476
477 async fn list_tasks(
478 &mut self,
479 request: Option<PaginatedRequestParams>,
480 context: &mut PluginContext<'_>,
481 ) -> PluginResult<Option<ListTasksResult>> {
482 match &self.task_router {
483 Some(router) => router.list_tasks(request, context).await,
484 None => Ok(None),
485 }
486 }
487
488 async fn get_task_info(
489 &mut self,
490 request: GetTaskInfoParams,
491 context: &mut PluginContext<'_>,
492 ) -> PluginResult<Option<GetTaskResult>> {
493 match &self.task_router {
494 Some(router) => router.get_task_info(request, context).await,
495 None => Ok(None),
496 }
497 }
498
499 async fn get_task_result(
500 &mut self,
501 request: GetTaskResultParams,
502 context: &mut PluginContext<'_>,
503 ) -> PluginResult<Option<GetTaskPayloadResult>> {
504 match &self.task_router {
505 Some(router) => router.get_task_result(request, context).await,
506 None => Ok(None),
507 }
508 }
509
510 async fn cancel_task(
511 &mut self,
512 request: CancelTaskParams,
513 context: &mut PluginContext<'_>,
514 ) -> PluginResult<Option<CancelTaskResult>> {
515 match &self.task_router {
516 Some(router) => router.cancel_task(request, context).await,
517 None => Ok(None),
518 }
519 }
520
521 async fn on_channel_message(
522 &mut self,
523 message: proto::ChannelMessage,
524 context: &mut PluginContext<'_>,
525 ) -> Result<()> {
526 match &self.channel_handler {
527 Some(handler) => handler(message, context).await,
528 None => Ok(()),
529 }
530 }
531
532 async fn on_bulk_transfer_message(
533 &mut self,
534 message: proto::BulkTransferMessage,
535 context: &mut PluginContext<'_>,
536 ) -> Result<()> {
537 match &self.bulk_handler {
538 Some(handler) => handler(message, context).await,
539 None => Ok(()),
540 }
541 }
542
543 async fn on_mesh_event(
544 &mut self,
545 event: proto::MeshEvent,
546 context: &mut PluginContext<'_>,
547 ) -> Result<()> {
548 match &self.mesh_event_handler {
549 Some(handler) => handler(event, context).await,
550 None => Ok(()),
551 }
552 }
553
554 async fn open_stream(
555 &mut self,
556 request: proto::OpenStreamRequest,
557 context: &mut PluginContext<'_>,
558 ) -> PluginResult<Option<proto::OpenStreamResponse>> {
559 match &self.open_stream_handler {
560 Some(handler) => handler(request, context).await,
561 None => Ok(None),
562 }
563 }
564
565 async fn on_cancel_stream(
566 &mut self,
567 notification: proto::CancelStreamNotification,
568 context: &mut PluginContext<'_>,
569 ) -> Result<()> {
570 match &self.cancel_stream_handler {
571 Some(handler) => handler(notification, context).await,
572 None => Ok(()),
573 }
574 }
575
576 async fn on_close_stream(
577 &mut self,
578 notification: proto::CloseStreamNotification,
579 context: &mut PluginContext<'_>,
580 ) -> Result<()> {
581 match &self.close_stream_handler {
582 Some(handler) => handler(notification, context).await,
583 None => Ok(()),
584 }
585 }
586
587 async fn on_stream_error(
588 &mut self,
589 error: proto::StreamError,
590 context: &mut PluginContext<'_>,
591 ) -> Result<()> {
592 match &self.stream_error_handler {
593 Some(handler) => handler(error, context).await,
594 None => Ok(()),
595 }
596 }
597}