mesh-llm-plugin 0.75.0

Plugin protocol and runtime primitives for mesh-llm
Documentation
use anyhow::Result;
use rmcp::model::{CallToolResult, ListToolsResult, ServerInfo};
use std::collections::BTreeMap;
use std::sync::Arc;

use crate::{
    context::PluginContext,
    error::{PluginResult, PluginRpcResult},
    helpers::{ToolCallRequest, ToolRouter},
    proto,
    runtime::{
        HealthFuture, Plugin, PluginInitializeRequest, PluginMetadata, PluginStartupPolicy,
        RpcMethodFuture, RpcMethodHandler,
    },
    simple_plugin::SimplePlugin,
};

pub struct InternalRpcPluginBuilder {
    plugin: SimplePlugin,
    rpc_handlers: BTreeMap<String, RpcMethodHandler>,
}

impl InternalRpcPluginBuilder {
    pub fn new(metadata: PluginMetadata) -> Self {
        Self {
            plugin: SimplePlugin::new(metadata),
            rpc_handlers: BTreeMap::new(),
        }
    }

    pub fn with_capabilities(mut self, capabilities: Vec<String>) -> Self {
        self.plugin = self.plugin.with_capabilities(capabilities);
        self
    }

    pub fn with_manifest(mut self, manifest: proto::PluginManifest) -> Self {
        self.plugin = self.plugin.with_manifest(manifest);
        self
    }

    pub fn with_startup_policy(mut self, startup_policy: PluginStartupPolicy) -> Self {
        self.plugin = self.plugin.with_startup_policy(startup_policy);
        self
    }

    pub fn with_operation_router(mut self, router: ToolRouter) -> Self {
        self.plugin = self.plugin.with_operation_router(router);
        self
    }

    pub fn with_health<F>(mut self, handler: F) -> Self
    where
        F: for<'a, 'ctx> Fn(&'a mut PluginContext<'ctx>) -> HealthFuture<'a>
            + Send
            + Sync
            + 'static,
    {
        self.plugin = self.plugin.with_health(handler);
        self
    }

    pub fn rpc_method<F>(mut self, method: impl Into<String>, handler: F) -> Self
    where
        F: for<'a, 'ctx> Fn(proto::RpcRequest, &'a mut PluginContext<'ctx>) -> RpcMethodFuture<'a>
            + Send
            + Sync
            + 'static,
    {
        self.rpc_handlers.insert(method.into(), Arc::new(handler));
        self
    }

    pub fn build(self) -> InternalRpcPlugin {
        InternalRpcPlugin {
            plugin: self.plugin,
            rpc_handlers: self.rpc_handlers,
        }
    }
}

#[derive(Clone)]
pub struct InternalRpcPlugin {
    plugin: SimplePlugin,
    rpc_handlers: BTreeMap<String, RpcMethodHandler>,
}

#[crate::async_trait]
impl Plugin for InternalRpcPlugin {
    fn plugin_id(&self) -> &str {
        self.plugin.plugin_id()
    }

    fn plugin_version(&self) -> String {
        self.plugin.plugin_version()
    }

    fn server_info(&self) -> ServerInfo {
        self.plugin.server_info()
    }

    fn capabilities(&self) -> Vec<String> {
        self.plugin.capabilities()
    }

    fn manifest(&self) -> Option<proto::PluginManifest> {
        self.plugin.manifest()
    }

    async fn initialize(
        &mut self,
        request: PluginInitializeRequest,
        context: &mut PluginContext<'_>,
    ) -> PluginResult<()> {
        self.plugin.initialize(request, context).await
    }

    async fn on_initialized(&mut self, context: &mut PluginContext<'_>) -> Result<()> {
        <SimplePlugin as Plugin>::on_initialized(&mut self.plugin, context).await
    }

    async fn health(&mut self, context: &mut PluginContext<'_>) -> Result<String> {
        self.plugin.health(context).await
    }

    async fn list_tools(
        &mut self,
        context: &mut PluginContext<'_>,
    ) -> PluginResult<Option<ListToolsResult>> {
        self.plugin.list_tools(context).await
    }

    async fn call_tool(
        &mut self,
        request: ToolCallRequest,
        context: &mut PluginContext<'_>,
    ) -> PluginResult<Option<CallToolResult>> {
        self.plugin.call_tool(request, context).await
    }

    async fn handle_rpc(
        &mut self,
        request: proto::RpcRequest,
        context: &mut PluginContext<'_>,
    ) -> PluginRpcResult {
        if let Some(handler) = self.rpc_handlers.get(&request.method).cloned() {
            return handler(request, context).await;
        }
        self.plugin.handle_rpc(request, context).await
    }
}