mesh_llm_plugin/
internal_rpc.rs

1use anyhow::Result;
2use rmcp::model::{CallToolResult, ListToolsResult, ServerInfo};
3use std::collections::BTreeMap;
4use std::sync::Arc;
5
6use crate::{
7    context::PluginContext,
8    error::{PluginResult, PluginRpcResult},
9    helpers::{ToolCallRequest, ToolRouter},
10    proto,
11    runtime::{
12        HealthFuture, Plugin, PluginInitializeRequest, PluginMetadata, PluginStartupPolicy,
13        RpcMethodFuture, RpcMethodHandler,
14    },
15    simple_plugin::SimplePlugin,
16};
17
18pub struct InternalRpcPluginBuilder {
19    plugin: SimplePlugin,
20    rpc_handlers: BTreeMap<String, RpcMethodHandler>,
21}
22
23impl InternalRpcPluginBuilder {
24    pub fn new(metadata: PluginMetadata) -> Self {
25        Self {
26            plugin: SimplePlugin::new(metadata),
27            rpc_handlers: BTreeMap::new(),
28        }
29    }
30
31    pub fn with_capabilities(mut self, capabilities: Vec<String>) -> Self {
32        self.plugin = self.plugin.with_capabilities(capabilities);
33        self
34    }
35
36    pub fn with_manifest(mut self, manifest: proto::PluginManifest) -> Self {
37        self.plugin = self.plugin.with_manifest(manifest);
38        self
39    }
40
41    pub fn with_startup_policy(mut self, startup_policy: PluginStartupPolicy) -> Self {
42        self.plugin = self.plugin.with_startup_policy(startup_policy);
43        self
44    }
45
46    pub fn with_operation_router(mut self, router: ToolRouter) -> Self {
47        self.plugin = self.plugin.with_operation_router(router);
48        self
49    }
50
51    pub fn with_health<F>(mut self, handler: F) -> Self
52    where
53        F: for<'a, 'ctx> Fn(&'a mut PluginContext<'ctx>) -> HealthFuture<'a>
54            + Send
55            + Sync
56            + 'static,
57    {
58        self.plugin = self.plugin.with_health(handler);
59        self
60    }
61
62    pub fn rpc_method<F>(mut self, method: impl Into<String>, handler: F) -> Self
63    where
64        F: for<'a, 'ctx> Fn(proto::RpcRequest, &'a mut PluginContext<'ctx>) -> RpcMethodFuture<'a>
65            + Send
66            + Sync
67            + 'static,
68    {
69        self.rpc_handlers.insert(method.into(), Arc::new(handler));
70        self
71    }
72
73    pub fn build(self) -> InternalRpcPlugin {
74        InternalRpcPlugin {
75            plugin: self.plugin,
76            rpc_handlers: self.rpc_handlers,
77        }
78    }
79}
80
81#[derive(Clone)]
82pub struct InternalRpcPlugin {
83    plugin: SimplePlugin,
84    rpc_handlers: BTreeMap<String, RpcMethodHandler>,
85}
86
87#[crate::async_trait]
88impl Plugin for InternalRpcPlugin {
89    fn plugin_id(&self) -> &str {
90        self.plugin.plugin_id()
91    }
92
93    fn plugin_version(&self) -> String {
94        self.plugin.plugin_version()
95    }
96
97    fn server_info(&self) -> ServerInfo {
98        self.plugin.server_info()
99    }
100
101    fn capabilities(&self) -> Vec<String> {
102        self.plugin.capabilities()
103    }
104
105    fn manifest(&self) -> Option<proto::PluginManifest> {
106        self.plugin.manifest()
107    }
108
109    async fn initialize(
110        &mut self,
111        request: PluginInitializeRequest,
112        context: &mut PluginContext<'_>,
113    ) -> PluginResult<()> {
114        self.plugin.initialize(request, context).await
115    }
116
117    async fn on_initialized(&mut self, context: &mut PluginContext<'_>) -> Result<()> {
118        <SimplePlugin as Plugin>::on_initialized(&mut self.plugin, context).await
119    }
120
121    async fn health(&mut self, context: &mut PluginContext<'_>) -> Result<String> {
122        self.plugin.health(context).await
123    }
124
125    async fn list_tools(
126        &mut self,
127        context: &mut PluginContext<'_>,
128    ) -> PluginResult<Option<ListToolsResult>> {
129        self.plugin.list_tools(context).await
130    }
131
132    async fn call_tool(
133        &mut self,
134        request: ToolCallRequest,
135        context: &mut PluginContext<'_>,
136    ) -> PluginResult<Option<CallToolResult>> {
137        self.plugin.call_tool(request, context).await
138    }
139
140    async fn handle_rpc(
141        &mut self,
142        request: proto::RpcRequest,
143        context: &mut PluginContext<'_>,
144    ) -> PluginRpcResult {
145        if let Some(handler) = self.rpc_handlers.get(&request.method).cloned() {
146            return handler(request, context).await;
147        }
148        self.plugin.handle_rpc(request, context).await
149    }
150}
mesh_llm_plugin/internal_rpc.rs

mesh_llm_plugin/
internal_rpc.rs
