mesh_llm_plugin/
runtime.rs

1use anyhow::{Result, bail};
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 serde::de::DeserializeOwned;
11use std::collections::HashMap;
12use std::future::Future;
13use std::pin::Pin;
14use std::sync::{Arc, Mutex};
15
16use tokio::sync::{RwLock, mpsc, watch};
17
18use crate::{
19    PROTOCOL_VERSION,
20    context::{
21        PendingHostResponses, PluginContext, drain_pending_host_responses,
22        remove_pending_host_response,
23    },
24    error::{PluginError, PluginResult, PluginRpcResult},
25    helpers::{
26        ToolCallRequest, json_response, parse_get_prompt_request, parse_read_resource_request,
27        parse_rpc_params, parse_tool_call_request,
28    },
29    io::{
30        LocalReadHalf, LocalStream, LocalWriteHalf, connect_from_env, read_envelope_from,
31        write_envelope_to,
32    },
33    proto,
34};
35use serde::{Deserialize, Serialize};
36
37pub use crate::internal_rpc::{InternalRpcPlugin, InternalRpcPluginBuilder};
38pub use crate::simple_plugin::SimplePlugin;
39
40#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
41#[serde(rename_all = "snake_case")]
42pub enum MeshVisibility {
43    #[default]
44    Private,
45    Public,
46}
47
48impl MeshVisibility {
49    fn from_proto(value: i32) -> Self {
50        match proto::MeshVisibility::try_from(value).unwrap_or(proto::MeshVisibility::Unspecified) {
51            proto::MeshVisibility::Public => Self::Public,
52            _ => Self::Private,
53        }
54    }
55}
56
57#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
58#[serde(rename_all = "snake_case")]
59pub enum PluginStartupPolicy {
60    #[default]
61    Any,
62    PrivateMeshOnly,
63    PublicMeshOnly,
64}
65
66#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
67pub struct PluginInitializeRequest {
68    pub host_protocol_version: u32,
69    pub host_version: String,
70    pub host_info_json: String,
71    pub mesh_visibility: MeshVisibility,
72}
73
74impl From<proto::InitializeRequest> for PluginInitializeRequest {
75    fn from(value: proto::InitializeRequest) -> Self {
76        Self {
77            host_protocol_version: value.host_protocol_version,
78            host_version: value.host_version,
79            host_info_json: value.host_info_json,
80            mesh_visibility: MeshVisibility::from_proto(value.mesh_visibility),
81        }
82    }
83}
84
85#[derive(Clone)]
86pub struct PluginMetadata {
87    pub(super) plugin_id: String,
88    pub(super) plugin_version: String,
89    pub(super) server_info: ServerInfo,
90    pub(super) capabilities: Vec<String>,
91    pub(super) manifest: Option<proto::PluginManifest>,
92    pub(super) startup_policy: PluginStartupPolicy,
93}
94
95impl PluginMetadata {
96    pub fn new(
97        plugin_id: impl Into<String>,
98        plugin_version: impl Into<String>,
99        server_info: ServerInfo,
100    ) -> Self {
101        Self {
102            plugin_id: plugin_id.into(),
103            plugin_version: plugin_version.into(),
104            server_info,
105            capabilities: Vec::new(),
106            manifest: None,
107            startup_policy: PluginStartupPolicy::Any,
108        }
109    }
110
111    pub fn with_capabilities(mut self, capabilities: Vec<String>) -> Self {
112        self.capabilities = capabilities;
113        self
114    }
115
116    pub fn with_manifest(mut self, manifest: proto::PluginManifest) -> Self {
117        self.manifest = Some(manifest);
118        self
119    }
120
121    pub fn with_startup_policy(mut self, startup_policy: PluginStartupPolicy) -> Self {
122        self.startup_policy = startup_policy;
123        self
124    }
125}
126
127pub(super) type InitializeFuture<'a> = Pin<Box<dyn Future<Output = PluginResult<()>> + Send + 'a>>;
128pub(super) type InitFuture<'a> = Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>;
129pub(super) type HealthFuture<'a> = Pin<Box<dyn Future<Output = Result<String>> + Send + 'a>>;
130pub(super) type OpenStreamFuture<'a> =
131    Pin<Box<dyn Future<Output = PluginResult<Option<proto::OpenStreamResponse>>> + Send + 'a>>;
132pub(super) type RpcMethodFuture<'a> = Pin<Box<dyn Future<Output = PluginRpcResult> + Send + 'a>>;
133pub(super) type SubscribeFuture<'a> = Pin<Box<dyn Future<Output = PluginResult<()>> + Send + 'a>>;
134pub(super) type SetLogLevelFuture<'a> = Pin<Box<dyn Future<Output = PluginResult<()>> + Send + 'a>>;
135
136pub(super) type InitializeHandler = Arc<
137    dyn for<'a, 'ctx> Fn(
138            PluginInitializeRequest,
139            &'a mut PluginContext<'ctx>,
140        ) -> InitializeFuture<'a>
141        + Send
142        + Sync,
143>;
144pub(super) type InitHandler =
145    Arc<dyn for<'a, 'ctx> Fn(&'a mut PluginContext<'ctx>) -> InitFuture<'a> + Send + Sync>;
146pub(super) type HealthHandler =
147    Arc<dyn for<'a, 'ctx> Fn(&'a mut PluginContext<'ctx>) -> HealthFuture<'a> + Send + Sync>;
148pub(super) type SubscribeHandler = Arc<
149    dyn for<'a, 'ctx> Fn(SubscribeRequestParams, &'a mut PluginContext<'ctx>) -> SubscribeFuture<'a>
150        + Send
151        + Sync,
152>;
153pub(super) type UnsubscribeHandler = Arc<
154    dyn for<'a, 'ctx> Fn(
155            UnsubscribeRequestParams,
156            &'a mut PluginContext<'ctx>,
157        ) -> SubscribeFuture<'a>
158        + Send
159        + Sync,
160>;
161pub(super) type SetLogLevelHandler = Arc<
162    dyn for<'a, 'ctx> Fn(
163            SetLevelRequestParams,
164            &'a mut PluginContext<'ctx>,
165        ) -> SetLogLevelFuture<'a>
166        + Send
167        + Sync,
168>;
169pub(super) type ChannelHandler = Arc<
170    dyn for<'a, 'ctx> Fn(proto::ChannelMessage, &'a mut PluginContext<'ctx>) -> InitFuture<'a>
171        + Send
172        + Sync,
173>;
174pub(super) type BulkHandler = Arc<
175    dyn for<'a, 'ctx> Fn(proto::BulkTransferMessage, &'a mut PluginContext<'ctx>) -> InitFuture<'a>
176        + Send
177        + Sync,
178>;
179pub(super) type MeshEventHandler = Arc<
180    dyn for<'a, 'ctx> Fn(proto::MeshEvent, &'a mut PluginContext<'ctx>) -> InitFuture<'a>
181        + Send
182        + Sync,
183>;
184pub(super) type RpcMethodHandler = Arc<
185    dyn for<'a, 'ctx> Fn(proto::RpcRequest, &'a mut PluginContext<'ctx>) -> RpcMethodFuture<'a>
186        + Send
187        + Sync,
188>;
189pub(super) type OpenStreamHandler = Arc<
190    dyn for<'a, 'ctx> Fn(
191            proto::OpenStreamRequest,
192            &'a mut PluginContext<'ctx>,
193        ) -> OpenStreamFuture<'a>
194        + Send
195        + Sync,
196>;
197pub(super) type CancelStreamHandler = Arc<
198    dyn for<'a, 'ctx> Fn(
199            proto::CancelStreamNotification,
200            &'a mut PluginContext<'ctx>,
201        ) -> InitFuture<'a>
202        + Send
203        + Sync,
204>;
205pub(super) type CloseStreamHandler = Arc<
206    dyn for<'a, 'ctx> Fn(
207            proto::CloseStreamNotification,
208            &'a mut PluginContext<'ctx>,
209        ) -> InitFuture<'a>
210        + Send
211        + Sync,
212>;
213pub(super) type StreamErrorHandler = Arc<
214    dyn for<'a, 'ctx> Fn(proto::StreamError, &'a mut PluginContext<'ctx>) -> InitFuture<'a>
215        + Send
216        + Sync,
217>;
218
219#[crate::async_trait]
220pub trait Plugin: Send {
221    fn plugin_id(&self) -> &str;
222    fn plugin_version(&self) -> String;
223    fn server_info(&self) -> ServerInfo;
224
225    fn capabilities(&self) -> Vec<String> {
226        Vec::new()
227    }
228
229    fn manifest(&self) -> Option<proto::PluginManifest> {
230        None
231    }
232
233    async fn initialize(
234        &mut self,
235        _request: PluginInitializeRequest,
236        _context: &mut PluginContext<'_>,
237    ) -> PluginResult<()> {
238        Ok(())
239    }
240
241    async fn on_initialized(&mut self, _context: &mut PluginContext<'_>) -> Result<()> {
242        Ok(())
243    }
244
245    async fn health(&mut self, _context: &mut PluginContext<'_>) -> Result<String> {
246        Ok("ok".into())
247    }
248
249    async fn list_tools(
250        &mut self,
251        _context: &mut PluginContext<'_>,
252    ) -> PluginResult<Option<ListToolsResult>> {
253        Ok(None)
254    }
255
256    async fn call_tool(
257        &mut self,
258        _request: ToolCallRequest,
259        _context: &mut PluginContext<'_>,
260    ) -> PluginResult<Option<CallToolResult>> {
261        Ok(None)
262    }
263
264    async fn list_prompts(
265        &mut self,
266        _request: Option<PaginatedRequestParams>,
267        _context: &mut PluginContext<'_>,
268    ) -> PluginResult<Option<ListPromptsResult>> {
269        Ok(None)
270    }
271
272    async fn get_prompt(
273        &mut self,
274        _request: GetPromptRequestParams,
275        _context: &mut PluginContext<'_>,
276    ) -> PluginResult<Option<GetPromptResult>> {
277        Ok(None)
278    }
279
280    async fn list_resources(
281        &mut self,
282        _request: Option<PaginatedRequestParams>,
283        _context: &mut PluginContext<'_>,
284    ) -> PluginResult<Option<ListResourcesResult>> {
285        Ok(None)
286    }
287
288    async fn read_resource(
289        &mut self,
290        _request: ReadResourceRequestParams,
291        _context: &mut PluginContext<'_>,
292    ) -> PluginResult<Option<ReadResourceResult>> {
293        Ok(None)
294    }
295
296    async fn list_resource_templates(
297        &mut self,
298        _request: Option<PaginatedRequestParams>,
299        _context: &mut PluginContext<'_>,
300    ) -> PluginResult<Option<ListResourceTemplatesResult>> {
301        Ok(None)
302    }
303
304    async fn subscribe_resource(
305        &mut self,
306        _request: SubscribeRequestParams,
307        _context: &mut PluginContext<'_>,
308    ) -> PluginResult<Option<()>> {
309        Ok(None)
310    }
311
312    async fn unsubscribe_resource(
313        &mut self,
314        _request: UnsubscribeRequestParams,
315        _context: &mut PluginContext<'_>,
316    ) -> PluginResult<Option<()>> {
317        Ok(None)
318    }
319
320    async fn complete(
321        &mut self,
322        _request: CompleteRequestParams,
323        _context: &mut PluginContext<'_>,
324    ) -> PluginResult<Option<CompleteResult>> {
325        Ok(None)
326    }
327
328    async fn set_log_level(
329        &mut self,
330        _request: SetLevelRequestParams,
331        _context: &mut PluginContext<'_>,
332    ) -> PluginResult<Option<()>> {
333        Ok(None)
334    }
335
336    async fn list_tasks(
337        &mut self,
338        _request: Option<PaginatedRequestParams>,
339        _context: &mut PluginContext<'_>,
340    ) -> PluginResult<Option<ListTasksResult>> {
341        Ok(None)
342    }
343
344    async fn get_task_info(
345        &mut self,
346        _request: GetTaskInfoParams,
347        _context: &mut PluginContext<'_>,
348    ) -> PluginResult<Option<GetTaskResult>> {
349        Ok(None)
350    }
351
352    async fn get_task_result(
353        &mut self,
354        _request: GetTaskResultParams,
355        _context: &mut PluginContext<'_>,
356    ) -> PluginResult<Option<GetTaskPayloadResult>> {
357        Ok(None)
358    }
359
360    async fn cancel_task(
361        &mut self,
362        _request: CancelTaskParams,
363        _context: &mut PluginContext<'_>,
364    ) -> PluginResult<Option<CancelTaskResult>> {
365        Ok(None)
366    }
367
368    async fn invoke_service(
369        &mut self,
370        request: proto::InvokeServiceRequest,
371        context: &mut PluginContext<'_>,
372    ) -> PluginResult<Option<proto::InvokeServiceResponse>> {
373        match proto::ServiceKind::try_from(request.kind).unwrap_or(proto::ServiceKind::Unspecified)
374        {
375            proto::ServiceKind::Operation => {
376                let arguments = parse_service_input::<serde_json::Value>(&request.input_json)?;
377                let tool_request = ToolCallRequest {
378                    name: request.service_name,
379                    arguments,
380                };
381                match self.call_tool(tool_request, context).await? {
382                    Some(result) => Ok(Some(proto::InvokeServiceResponse {
383                        output_json: normalize_call_tool_output(&result)?,
384                        is_error: result.is_error.unwrap_or(false),
385                    })),
386                    None => Ok(None),
387                }
388            }
389            proto::ServiceKind::Prompt => {
390                let params = parse_service_input::<GetPromptRequestParams>(&request.input_json)?;
391                match self.get_prompt(params, context).await? {
392                    Some(result) => Ok(Some(proto::InvokeServiceResponse {
393                        output_json: serialize_service_output(&result)?,
394                        is_error: false,
395                    })),
396                    None => Ok(None),
397                }
398            }
399            proto::ServiceKind::Resource => {
400                let params = parse_service_input::<ReadResourceRequestParams>(&request.input_json)?;
401                match self.read_resource(params, context).await? {
402                    Some(result) => Ok(Some(proto::InvokeServiceResponse {
403                        output_json: serialize_service_output(&result)?,
404                        is_error: false,
405                    })),
406                    None => Ok(None),
407                }
408            }
409            proto::ServiceKind::Completion => {
410                let params = parse_service_input::<CompleteRequestParams>(&request.input_json)?;
411                match self.complete(params, context).await? {
412                    Some(result) => Ok(Some(proto::InvokeServiceResponse {
413                        output_json: serialize_service_output(&result)?,
414                        is_error: false,
415                    })),
416                    None => Ok(None),
417                }
418            }
419            proto::ServiceKind::Unspecified => Err(PluginError::invalid_request(
420                "Service invocation kind is required",
421            )),
422        }
423    }
424
425    async fn handle_rpc(
426        &mut self,
427        request: proto::RpcRequest,
428        context: &mut PluginContext<'_>,
429    ) -> PluginRpcResult {
430        match request.method.as_str() {
431            "tools/list" => match self.list_tools(context).await? {
432                Some(result) => json_response(&result),
433                None => Err(PluginError::method_not_found(
434                    "Unsupported MCP method 'tools/list'",
435                )),
436            },
437            "tools/call" => {
438                let tool_call = parse_tool_call_request(&request)?;
439                match self.call_tool(tool_call, context).await? {
440                    Some(result) => json_response(&result),
441                    None => Err(PluginError::method_not_found(
442                        "Unsupported MCP method 'tools/call'",
443                    )),
444                }
445            }
446            "prompts/list" => {
447                let params: Option<PaginatedRequestParams> = parse_rpc_params(&request)?;
448                match self.list_prompts(params, context).await? {
449                    Some(result) => json_response(&result),
450                    None => Err(PluginError::method_not_found(
451                        "Unsupported MCP method 'prompts/list'",
452                    )),
453                }
454            }
455            "prompts/get" => {
456                let params = parse_get_prompt_request(&request)?;
457                match self.get_prompt(params, context).await? {
458                    Some(result) => json_response(&result),
459                    None => Err(PluginError::method_not_found(
460                        "Unsupported MCP method 'prompts/get'",
461                    )),
462                }
463            }
464            "resources/list" => {
465                let params: Option<PaginatedRequestParams> = parse_rpc_params(&request)?;
466                match self.list_resources(params, context).await? {
467                    Some(result) => json_response(&result),
468                    None => Err(PluginError::method_not_found(
469                        "Unsupported MCP method 'resources/list'",
470                    )),
471                }
472            }
473            "resources/read" => {
474                let params = parse_read_resource_request(&request)?;
475                match self.read_resource(params, context).await? {
476                    Some(result) => json_response(&result),
477                    None => Err(PluginError::method_not_found(
478                        "Unsupported MCP method 'resources/read'",
479                    )),
480                }
481            }
482            "resources/templates/list" => {
483                let params: Option<PaginatedRequestParams> = parse_rpc_params(&request)?;
484                match self.list_resource_templates(params, context).await? {
485                    Some(result) => json_response(&result),
486                    None => Err(PluginError::method_not_found(
487                        "Unsupported MCP method 'resources/templates/list'",
488                    )),
489                }
490            }
491            "resources/subscribe" => {
492                let params: SubscribeRequestParams = parse_rpc_params(&request)?;
493                match self.subscribe_resource(params, context).await? {
494                    Some(()) => json_response(&serde_json::json!({})),
495                    None => Err(PluginError::method_not_found(
496                        "Unsupported MCP method 'resources/subscribe'",
497                    )),
498                }
499            }
500            "resources/unsubscribe" => {
501                let params: UnsubscribeRequestParams = parse_rpc_params(&request)?;
502                match self.unsubscribe_resource(params, context).await? {
503                    Some(()) => json_response(&serde_json::json!({})),
504                    None => Err(PluginError::method_not_found(
505                        "Unsupported MCP method 'resources/unsubscribe'",
506                    )),
507                }
508            }
509            "completion/complete" => {
510                let params: CompleteRequestParams = parse_rpc_params(&request)?;
511                match self.complete(params, context).await? {
512                    Some(result) => json_response(&result),
513                    None => Err(PluginError::method_not_found(
514                        "Unsupported MCP method 'completion/complete'",
515                    )),
516                }
517            }
518            "logging/setLevel" => {
519                let params: SetLevelRequestParams = parse_rpc_params(&request)?;
520                match self.set_log_level(params, context).await? {
521                    Some(()) => json_response(&serde_json::json!({})),
522                    None => Err(PluginError::method_not_found(
523                        "Unsupported MCP method 'logging/setLevel'",
524                    )),
525                }
526            }
527            "tasks/list" => {
528                let params: Option<PaginatedRequestParams> = parse_rpc_params(&request)?;
529                match self.list_tasks(params, context).await? {
530                    Some(result) => json_response(&result),
531                    None => Err(PluginError::method_not_found(
532                        "Unsupported MCP method 'tasks/list'",
533                    )),
534                }
535            }
536            "tasks/get" => {
537                let params: GetTaskInfoParams = parse_rpc_params(&request)?;
538                match self.get_task_info(params, context).await? {
539                    Some(result) => json_response(&result),
540                    None => Err(PluginError::method_not_found(
541                        "Unsupported MCP method 'tasks/get'",
542                    )),
543                }
544            }
545            "tasks/result" => {
546                let params: GetTaskResultParams = parse_rpc_params(&request)?;
547                match self.get_task_result(params, context).await? {
548                    Some(result) => json_response(&result),
549                    None => Err(PluginError::method_not_found(
550                        "Unsupported MCP method 'tasks/result'",
551                    )),
552                }
553            }
554            "tasks/cancel" => {
555                let params: CancelTaskParams = parse_rpc_params(&request)?;
556                match self.cancel_task(params, context).await? {
557                    Some(result) => json_response(&result),
558                    None => Err(PluginError::method_not_found(
559                        "Unsupported MCP method 'tasks/cancel'",
560                    )),
561                }
562            }
563            _ => Err(PluginError::method_not_found(format!(
564                "Unsupported MCP method '{}'",
565                request.method
566            ))),
567        }
568    }
569
570    async fn on_rpc_notification(
571        &mut self,
572        _notification: proto::RpcNotification,
573        _context: &mut PluginContext<'_>,
574    ) -> Result<()> {
575        Ok(())
576    }
577
578    async fn on_channel_message(
579        &mut self,
580        _message: proto::ChannelMessage,
581        _context: &mut PluginContext<'_>,
582    ) -> Result<()> {
583        Ok(())
584    }
585
586    async fn on_bulk_transfer_message(
587        &mut self,
588        _message: proto::BulkTransferMessage,
589        _context: &mut PluginContext<'_>,
590    ) -> Result<()> {
591        Ok(())
592    }
593
594    async fn on_mesh_event(
595        &mut self,
596        _event: proto::MeshEvent,
597        _context: &mut PluginContext<'_>,
598    ) -> Result<()> {
599        Ok(())
600    }
601
602    async fn open_stream(
603        &mut self,
604        _request: proto::OpenStreamRequest,
605        _context: &mut PluginContext<'_>,
606    ) -> PluginResult<Option<proto::OpenStreamResponse>> {
607        Ok(None)
608    }
609
610    async fn on_cancel_stream(
611        &mut self,
612        _notification: proto::CancelStreamNotification,
613        _context: &mut PluginContext<'_>,
614    ) -> Result<()> {
615        Ok(())
616    }
617
618    async fn on_close_stream(
619        &mut self,
620        _notification: proto::CloseStreamNotification,
621        _context: &mut PluginContext<'_>,
622    ) -> Result<()> {
623        Ok(())
624    }
625
626    async fn on_stream_error(
627        &mut self,
628        _error: proto::StreamError,
629        _context: &mut PluginContext<'_>,
630    ) -> Result<()> {
631        Ok(())
632    }
633
634    async fn on_host_error(
635        &mut self,
636        error: proto::ErrorResponse,
637        _context: &mut PluginContext<'_>,
638    ) -> Result<()> {
639        bail!("host error: {}", error.message)
640    }
641}
642
643pub struct PluginRuntime;
644
645struct RuntimeState<P> {
646    plugin: Arc<RwLock<P>>,
647    plugin_id: String,
648    outbound_tx: mpsc::Sender<proto::Envelope>,
649    pending_host_responses: PendingHostResponses,
650}
651
652struct OrderedPayload {
653    request_id: u64,
654    payload: proto::envelope::Payload,
655}
656
657impl PluginRuntime {
658    pub async fn run<P: Plugin + Clone + Sync + 'static>(plugin: P) -> Result<()> {
659        let stream = connect_from_env().await?;
660        Self::run_with_stream(plugin, stream).await
661    }
662
663    pub async fn run_with_stream<P: Plugin + Clone + Sync + 'static>(
664        plugin: P,
665        stream: LocalStream,
666    ) -> Result<()> {
667        let plugin_id = plugin.plugin_id().to_string();
668        let (read, write) = stream.into_split();
669        let (outbound_tx, outbound_rx) = mpsc::channel(256);
670        let (ordered_tx, ordered_rx) = mpsc::channel(256);
671        let (shutdown_tx, shutdown_rx) = watch::channel(false);
672        let state = Arc::new(RuntimeState {
673            plugin: Arc::new(RwLock::new(plugin)),
674            plugin_id,
675            outbound_tx,
676            pending_host_responses: Arc::new(Mutex::new(HashMap::new())),
677        });
678        let mut writer = tokio::spawn(Self::write_loop(
679            write,
680            outbound_rx,
681            state.pending_host_responses.clone(),
682            shutdown_tx.clone(),
683            shutdown_rx.clone(),
684        ));
685        let ordered_handlers = tokio::spawn(Self::ordered_handler_loop(
686            state.clone(),
687            ordered_rx,
688            shutdown_rx,
689        ));
690
691        let read_result = Self::read_loop(state.clone(), read, ordered_tx);
692        tokio::pin!(read_result);
693
694        let result = tokio::select! {
695            read_result = &mut read_result => read_result,
696            writer_result = &mut writer => match writer_result {
697                Ok(Ok(())) => Ok(()),
698                Ok(Err(err)) => Err(err),
699                Err(err) => Err(err.into()),
700            },
701        };
702
703        let shutdown_reason = match &result {
704            Ok(()) => "plugin host connection is closed".to_string(),
705            Err(err) => format!("plugin host connection is closed: {err}"),
706        };
707        Self::shutdown_runtime(&shutdown_tx, &state.pending_host_responses, shutdown_reason);
708        if !writer.is_finished() {
709            let _ = writer.await;
710        }
711        if !ordered_handlers.is_finished() {
712            ordered_handlers.abort();
713        }
714
715        match ordered_handlers.await {
716            Ok(()) => result,
717            Err(err) if err.is_cancelled() => result,
718            Err(err) if result.is_ok() => Err(err.into()),
719            Err(_) => result,
720        }
721    }
722
723    fn shutdown_runtime(
724        shutdown_tx: &watch::Sender<bool>,
725        pending_host_responses: &PendingHostResponses,
726        reason: String,
727    ) {
728        let _ = shutdown_tx.send(true);
729        Self::fail_pending_host_responses(pending_host_responses, reason);
730    }
731
732    fn fail_pending_host_responses(pending_host_responses: &PendingHostResponses, reason: String) {
733        for sender in drain_pending_host_responses(pending_host_responses) {
734            let _ = sender.send(Err(anyhow::anyhow!(reason.clone())));
735        }
736    }
737
738    async fn ordered_handler_loop<P: Plugin + Clone + Sync + 'static>(
739        state: Arc<RuntimeState<P>>,
740        mut ordered_rx: mpsc::Receiver<OrderedPayload>,
741        mut shutdown_rx: watch::Receiver<bool>,
742    ) {
743        loop {
744            tokio::select! {
745                _ = shutdown_rx.changed() => {
746                    if *shutdown_rx.borrow() {
747                        break;
748                    }
749                }
750                payload = ordered_rx.recv() => {
751                    let Some(payload) = payload else {
752                        break;
753                    };
754                    let _ = Self::handle_payload(state.clone(), payload.request_id, payload.payload).await;
755                }
756            }
757        }
758    }
759
760    async fn read_loop<P: Plugin + Clone + Sync + 'static>(
761        state: Arc<RuntimeState<P>>,
762        mut read: LocalReadHalf,
763        ordered_tx: mpsc::Sender<OrderedPayload>,
764    ) -> Result<()> {
765        loop {
766            let envelope = read_envelope_from(&mut *read).await?;
767            if Self::complete_pending_host_response(&state, &envelope).await {
768                continue;
769            }
770            if !Self::handle_envelope(state.clone(), &ordered_tx, envelope).await? {
771                break;
772            }
773        }
774        Ok(())
775    }
776
777    async fn write_loop(
778        mut write: LocalWriteHalf,
779        mut outbound_rx: mpsc::Receiver<proto::Envelope>,
780        pending_host_responses: PendingHostResponses,
781        shutdown_tx: watch::Sender<bool>,
782        mut shutdown_rx: watch::Receiver<bool>,
783    ) -> Result<()> {
784        loop {
785            tokio::select! {
786                _ = shutdown_rx.changed() => {
787                    if *shutdown_rx.borrow() {
788                        Self::drain_outbound_before_shutdown(&mut write, &mut outbound_rx).await?;
789                        return Ok(());
790                    }
791                }
792                envelope = outbound_rx.recv() => {
793                    let Some(envelope) = envelope else {
794                        return Ok(());
795                    };
796                    if let Err(err) = write_envelope_to(&mut *write, &envelope).await {
797                        let reason = format!("plugin host write failed: {err}");
798                        Self::shutdown_runtime(&shutdown_tx, &pending_host_responses, reason);
799                        return Err(err);
800                    }
801                }
802            }
803        }
804    }
805
806    async fn drain_outbound_before_shutdown(
807        write: &mut LocalWriteHalf,
808        outbound_rx: &mut mpsc::Receiver<proto::Envelope>,
809    ) -> Result<()> {
810        while let Ok(envelope) = outbound_rx.try_recv() {
811            write_envelope_to(&mut **write, &envelope).await?;
812        }
813        Ok(())
814    }
815
816    async fn complete_pending_host_response<P: Plugin + Clone + Sync>(
817        state: &RuntimeState<P>,
818        envelope: &proto::Envelope,
819    ) -> bool {
820        let Some(sender) =
821            remove_pending_host_response(&state.pending_host_responses, envelope.request_id)
822        else {
823            return false;
824        };
825        let _ = sender.send(Ok(envelope.clone()));
826        true
827    }
828
829    async fn handle_envelope<P: Plugin + Clone + Sync + 'static>(
830        state: Arc<RuntimeState<P>>,
831        ordered_tx: &mpsc::Sender<OrderedPayload>,
832        envelope: proto::Envelope,
833    ) -> Result<bool> {
834        let request_id = envelope.request_id;
835        let Some(payload) = envelope.payload else {
836            return Ok(true);
837        };
838
839        match payload {
840            proto::envelope::Payload::InitializeRequest(request) => {
841                Self::handle_initialize(state, request_id, request).await
842            }
843            proto::envelope::Payload::HealthRequest(_) => Self::handle_health(state, request_id),
844            proto::envelope::Payload::ShutdownRequest(_) => {
845                Self::write_payload(
846                    &state,
847                    request_id,
848                    proto::envelope::Payload::ShutdownResponse(proto::ShutdownResponse {}),
849                )
850                .await?;
851                Ok(false)
852            }
853            payload if Self::is_ordered_payload(&payload) => {
854                Self::enqueue_ordered_payload(ordered_tx, request_id, payload).await
855            }
856            payload => Self::spawn_payload_handler(state, request_id, payload),
857        }
858    }
859
860    fn is_ordered_payload(payload: &proto::envelope::Payload) -> bool {
861        matches!(
862            payload,
863            proto::envelope::Payload::RpcNotification(_)
864                | proto::envelope::Payload::ChannelMessage(_)
865                | proto::envelope::Payload::BulkTransferMessage(_)
866                | proto::envelope::Payload::MeshEvent(_)
867                | proto::envelope::Payload::CancelStreamNotification(_)
868                | proto::envelope::Payload::CloseStreamNotification(_)
869                | proto::envelope::Payload::StreamError(_)
870                | proto::envelope::Payload::ErrorResponse(_)
871        )
872    }
873
874    async fn enqueue_ordered_payload(
875        ordered_tx: &mpsc::Sender<OrderedPayload>,
876        request_id: u64,
877        payload: proto::envelope::Payload,
878    ) -> Result<bool> {
879        ordered_tx
880            .send(OrderedPayload {
881                request_id,
882                payload,
883            })
884            .await
885            .map_err(|_| anyhow::anyhow!("plugin ordered handler is closed"))?;
886        Ok(true)
887    }
888
889    fn spawn_payload_handler<P: Plugin + Clone + Sync + 'static>(
890        state: Arc<RuntimeState<P>>,
891        request_id: u64,
892        payload: proto::envelope::Payload,
893    ) -> Result<bool> {
894        tokio::spawn(async move {
895            let _ = Self::handle_payload(state, request_id, payload).await;
896        });
897        Ok(true)
898    }
899
900    async fn handle_payload<P: Plugin + Clone + Sync>(
901        state: Arc<RuntimeState<P>>,
902        request_id: u64,
903        payload: proto::envelope::Payload,
904    ) -> Result<()> {
905        match payload {
906            proto::envelope::Payload::RpcRequest(request) => {
907                let payload = Self::rpc_payload(&state, request).await;
908                Self::write_payload(&state, request_id, payload).await
909            }
910            proto::envelope::Payload::InvokeServiceRequest(request) => {
911                let payload = Self::invoke_service_payload(&state, request).await;
912                Self::write_payload(&state, request_id, payload).await
913            }
914            proto::envelope::Payload::OpenStreamRequest(request) => {
915                let payload = Self::open_stream_payload(&state, request).await;
916                Self::write_payload(&state, request_id, payload).await
917            }
918            proto::envelope::Payload::RpcNotification(notification) => {
919                Self::handle_rpc_notification(&state, notification)
920                    .await
921                    .map(|_| ())
922            }
923            proto::envelope::Payload::ChannelMessage(message) => {
924                Self::handle_channel_message(&state, message)
925                    .await
926                    .map(|_| ())
927            }
928            proto::envelope::Payload::BulkTransferMessage(message) => {
929                Self::handle_bulk_transfer_message(&state, message)
930                    .await
931                    .map(|_| ())
932            }
933            proto::envelope::Payload::MeshEvent(event) => {
934                Self::handle_mesh_event(&state, event).await.map(|_| ())
935            }
936            proto::envelope::Payload::CancelStreamNotification(notification) => {
937                Self::handle_cancel_stream(&state, notification)
938                    .await
939                    .map(|_| ())
940            }
941            proto::envelope::Payload::CloseStreamNotification(notification) => {
942                Self::handle_close_stream(&state, notification)
943                    .await
944                    .map(|_| ())
945            }
946            proto::envelope::Payload::StreamError(error) => {
947                Self::handle_stream_error(&state, error).await.map(|_| ())
948            }
949            proto::envelope::Payload::ErrorResponse(error) => {
950                Self::handle_host_error(&state, error).await.map(|_| ())
951            }
952            _ => Ok(()),
953        }?;
954        Ok(())
955    }
956
957    async fn handle_initialize<P: Plugin + Clone + Sync>(
958        state: Arc<RuntimeState<P>>,
959        request_id: u64,
960        request: proto::InitializeRequest,
961    ) -> Result<bool> {
962        let mut plugin = state.plugin.write().await;
963        let mut context = Self::context(&state);
964        let init_result = plugin
965            .initialize(PluginInitializeRequest::from(request), &mut context)
966            .await;
967        if let Err(err) = init_result {
968            Self::write_payload(
969                &state,
970                request_id,
971                proto::envelope::Payload::ErrorResponse(err.into_error_response()),
972            )
973            .await?;
974            return Ok(false);
975        }
976
977        Self::write_payload(
978            &state,
979            request_id,
980            proto::envelope::Payload::InitializeResponse(proto::InitializeResponse {
981                plugin_id: state.plugin_id.clone(),
982                plugin_protocol_version: PROTOCOL_VERSION,
983                plugin_version: plugin.plugin_version(),
984                server_info_json: serde_json::to_string(&plugin.server_info())?,
985                capabilities: plugin.capabilities(),
986                manifest: plugin.manifest(),
987            }),
988        )
989        .await?;
990
991        let mut context = Self::context(&state);
992        plugin.on_initialized(&mut context).await?;
993        Ok(true)
994    }
995
996    fn handle_health<P: Plugin + Clone + Sync + 'static>(
997        state: Arc<RuntimeState<P>>,
998        request_id: u64,
999    ) -> Result<bool> {
1000        tokio::spawn(async move {
1001            let mut plugin = Self::plugin_for_request(&state).await;
1002            let mut context = Self::context(&state);
1003            let payload = match plugin.health(&mut context).await {
1004                Ok(detail) => proto::envelope::Payload::HealthResponse(proto::HealthResponse {
1005                    status: proto::health_response::Status::Ok as i32,
1006                    detail,
1007                }),
1008                Err(err) => proto::envelope::Payload::ErrorResponse(
1009                    PluginError::internal(format!("health check failed: {err}"))
1010                        .into_error_response(),
1011                ),
1012            };
1013            let _ = Self::write_payload(&state, request_id, payload).await;
1014        });
1015        Ok(true)
1016    }
1017
1018    async fn plugin_for_request<P: Plugin + Clone + Sync>(state: &RuntimeState<P>) -> P {
1019        state.plugin.read().await.clone()
1020    }
1021
1022    async fn rpc_payload<P: Plugin + Clone + Sync>(
1023        state: &RuntimeState<P>,
1024        request: proto::RpcRequest,
1025    ) -> proto::envelope::Payload {
1026        let mut plugin = Self::plugin_for_request(state).await;
1027        let mut context = Self::context(state);
1028        match plugin.handle_rpc(request, &mut context).await {
1029            Ok(payload) => payload,
1030            Err(err) => proto::envelope::Payload::ErrorResponse(err.into_error_response()),
1031        }
1032    }
1033
1034    async fn invoke_service_payload<P: Plugin + Clone + Sync>(
1035        state: &RuntimeState<P>,
1036        request: proto::InvokeServiceRequest,
1037    ) -> proto::envelope::Payload {
1038        let mut plugin = Self::plugin_for_request(state).await;
1039        let mut context = Self::context(state);
1040        match plugin.invoke_service(request, &mut context).await {
1041            Ok(Some(response)) => proto::envelope::Payload::InvokeServiceResponse(response),
1042            Ok(None) => proto::envelope::Payload::ErrorResponse(
1043                PluginError::method_not_found("Unsupported service invocation")
1044                    .into_error_response(),
1045            ),
1046            Err(err) => proto::envelope::Payload::ErrorResponse(err.into_error_response()),
1047        }
1048    }
1049
1050    async fn open_stream_payload<P: Plugin + Clone + Sync>(
1051        state: &RuntimeState<P>,
1052        request: proto::OpenStreamRequest,
1053    ) -> proto::envelope::Payload {
1054        let mut plugin = Self::plugin_for_request(state).await;
1055        let mut context = Self::context(state);
1056        match plugin.open_stream(request, &mut context).await {
1057            Ok(Some(response)) => proto::envelope::Payload::OpenStreamResponse(response),
1058            Ok(None) => proto::envelope::Payload::ErrorResponse(
1059                PluginError::method_not_found("Unsupported stream control message 'open_stream'")
1060                    .into_error_response(),
1061            ),
1062            Err(err) => proto::envelope::Payload::ErrorResponse(err.into_error_response()),
1063        }
1064    }
1065
1066    async fn handle_rpc_notification<P: Plugin + Clone + Sync>(
1067        state: &RuntimeState<P>,
1068        notification: proto::RpcNotification,
1069    ) -> Result<bool> {
1070        let mut plugin = Self::plugin_for_request(state).await;
1071        let mut context = Self::context(state);
1072        plugin
1073            .on_rpc_notification(notification, &mut context)
1074            .await?;
1075        Ok(true)
1076    }
1077
1078    async fn handle_channel_message<P: Plugin + Clone + Sync>(
1079        state: &RuntimeState<P>,
1080        message: proto::ChannelMessage,
1081    ) -> Result<bool> {
1082        let mut plugin = Self::plugin_for_request(state).await;
1083        let mut context = Self::context(state);
1084        plugin.on_channel_message(message, &mut context).await?;
1085        Ok(true)
1086    }
1087
1088    async fn handle_bulk_transfer_message<P: Plugin + Clone + Sync>(
1089        state: &RuntimeState<P>,
1090        message: proto::BulkTransferMessage,
1091    ) -> Result<bool> {
1092        let mut plugin = Self::plugin_for_request(state).await;
1093        let mut context = Self::context(state);
1094        plugin
1095            .on_bulk_transfer_message(message, &mut context)
1096            .await?;
1097        Ok(true)
1098    }
1099
1100    async fn handle_mesh_event<P: Plugin + Clone + Sync>(
1101        state: &RuntimeState<P>,
1102        event: proto::MeshEvent,
1103    ) -> Result<bool> {
1104        let mut plugin = Self::plugin_for_request(state).await;
1105        let mut context = Self::context(state);
1106        plugin.on_mesh_event(event, &mut context).await?;
1107        Ok(true)
1108    }
1109
1110    async fn handle_cancel_stream<P: Plugin + Clone + Sync>(
1111        state: &RuntimeState<P>,
1112        notification: proto::CancelStreamNotification,
1113    ) -> Result<bool> {
1114        let mut plugin = Self::plugin_for_request(state).await;
1115        let mut context = Self::context(state);
1116        plugin.on_cancel_stream(notification, &mut context).await?;
1117        Ok(true)
1118    }
1119
1120    async fn handle_close_stream<P: Plugin + Clone + Sync>(
1121        state: &RuntimeState<P>,
1122        notification: proto::CloseStreamNotification,
1123    ) -> Result<bool> {
1124        let mut plugin = Self::plugin_for_request(state).await;
1125        let mut context = Self::context(state);
1126        plugin.on_close_stream(notification, &mut context).await?;
1127        Ok(true)
1128    }
1129
1130    async fn handle_stream_error<P: Plugin + Clone + Sync>(
1131        state: &RuntimeState<P>,
1132        error: proto::StreamError,
1133    ) -> Result<bool> {
1134        let mut plugin = Self::plugin_for_request(state).await;
1135        let mut context = Self::context(state);
1136        plugin.on_stream_error(error, &mut context).await?;
1137        Ok(true)
1138    }
1139
1140    async fn handle_host_error<P: Plugin + Clone + Sync>(
1141        state: &RuntimeState<P>,
1142        error: proto::ErrorResponse,
1143    ) -> Result<bool> {
1144        let mut plugin = Self::plugin_for_request(state).await;
1145        let mut context = Self::context(state);
1146        plugin.on_host_error(error, &mut context).await?;
1147        Ok(true)
1148    }
1149
1150    fn context<P: Plugin + Clone + Sync>(state: &RuntimeState<P>) -> PluginContext<'static> {
1151        PluginContext::new(
1152            state.plugin_id.clone(),
1153            state.outbound_tx.clone(),
1154            state.pending_host_responses.clone(),
1155        )
1156    }
1157
1158    async fn write_payload<P: Plugin + Clone + Sync>(
1159        state: &RuntimeState<P>,
1160        request_id: u64,
1161        payload: proto::envelope::Payload,
1162    ) -> Result<()> {
1163        state
1164            .outbound_tx
1165            .send(proto::Envelope {
1166                protocol_version: PROTOCOL_VERSION,
1167                plugin_id: state.plugin_id.clone(),
1168                request_id,
1169                payload: Some(payload),
1170            })
1171            .await
1172            .map_err(|_| anyhow::anyhow!("plugin host connection is closed"))
1173    }
1174}
1175
1176fn parse_service_input<T: DeserializeOwned>(input_json: &str) -> PluginResult<T> {
1177    let input = if input_json.trim().is_empty() {
1178        "null"
1179    } else {
1180        input_json
1181    };
1182    serde_json::from_str(input)
1183        .map_err(|err| PluginError::invalid_params(format!("Invalid service input JSON: {err}")))
1184}
1185
1186fn serialize_service_output<T: Serialize>(value: &T) -> PluginResult<String> {
1187    serde_json::to_string(value)
1188        .map_err(|err| PluginError::internal(format!("Serialize service output: {err}")))
1189}
1190
1191fn normalize_call_tool_output(result: &CallToolResult) -> PluginResult<String> {
1192    if let Some(value) = &result.structured_content {
1193        return serialize_service_output(value);
1194    }
1195    if let Some(text) = result.content.first().and_then(|content| content.as_text()) {
1196        return Ok(text.text.clone());
1197    }
1198    serialize_service_output(&result.content)
1199}
1200
1201#[cfg(test)]
1202mod tests {
1203    use super::*;
1204    use crate::{mcp, plugin, plugin_server_info};
1205    use crate::{read_envelope, write_envelope};
1206    use rmcp::model::{
1207        ArgumentInfo, PromptMessage, PromptMessageContent, PromptMessageRole, Reference,
1208    };
1209    use serde_json::json;
1210    use tokio::sync::{Barrier, Notify};
1211    use tokio::time::{Duration, timeout};
1212
1213    #[derive(Clone, Debug, Default, Deserialize, Serialize, schemars::JsonSchema)]
1214    struct DemoArgs {
1215        #[serde(default)]
1216        message: String,
1217    }
1218
1219    fn test_context() -> PluginContext<'static> {
1220        let (outbound_tx, _outbound_rx) = mpsc::channel(8);
1221        PluginContext::new(
1222            "demo".into(),
1223            outbound_tx,
1224            Arc::new(Mutex::new(HashMap::new())),
1225        )
1226    }
1227
1228    fn test_channel_message(message_kind: &str) -> proto::ChannelMessage {
1229        proto::ChannelMessage {
1230            channel: "events".into(),
1231            source_peer_id: "peer-a".into(),
1232            target_peer_id: "peer-b".into(),
1233            content_type: "text/plain".into(),
1234            body: Vec::new(),
1235            message_kind: message_kind.into(),
1236            correlation_id: String::new(),
1237            metadata_json: String::new(),
1238        }
1239    }
1240
1241    #[tokio::test]
1242    async fn invoke_service_dispatches_operation_prompt_resource_and_completion() {
1243        let mut plugin = plugin! {
1244            metadata: PluginMetadata::new(
1245                "demo",
1246                "1.0.0",
1247                plugin_server_info("demo", "1.0.0", "Demo", "Demo plugin", None::<String>),
1248            ),
1249            mcp: [
1250                mcp::tool("echo")
1251                    .description("Echo input")
1252                    .input::<DemoArgs>()
1253                    .handle(|args, _context| Box::pin(async move {
1254                        Ok(json!({ "echo": args.message }))
1255                    })),
1256                mcp::resource("demo://state")
1257                    .name("State")
1258                    .handle(|request, _context| Box::pin(async move {
1259                        Ok(crate::read_resource_result(vec![
1260                            rmcp::model::ResourceContents::text("state", request.uri),
1261                        ]))
1262                    })),
1263                mcp::prompt("brief")
1264                    .description("Brief")
1265                    .handle(|request, _context| Box::pin(async move {
1266                        Ok(crate::get_prompt_result(vec![PromptMessage::new(
1267                            PromptMessageRole::User,
1268                            PromptMessageContent::text(format!("brief:{}", request.name)),
1269                        )]))
1270                    })),
1271                mcp::completion("prompt.brief.topic")
1272                    .handle(|_request, _context| Box::pin(async move {
1273                        crate::complete_result(vec!["alpha".into()])
1274                    })),
1275            ],
1276        };
1277
1278        let mut context = test_context();
1279
1280        let op = plugin
1281            .invoke_service(
1282                proto::InvokeServiceRequest {
1283                    kind: proto::ServiceKind::Operation as i32,
1284                    service_name: "echo".into(),
1285                    input_json: json!({ "message": "hello" }).to_string(),
1286                },
1287                &mut context,
1288            )
1289            .await
1290            .unwrap()
1291            .unwrap();
1292        assert_eq!(
1293            serde_json::from_str::<serde_json::Value>(&op.output_json).unwrap(),
1294            json!({"echo": "hello"})
1295        );
1296
1297        let prompt = plugin
1298            .invoke_service(
1299                proto::InvokeServiceRequest {
1300                    kind: proto::ServiceKind::Prompt as i32,
1301                    service_name: "brief".into(),
1302                    input_json: serde_json::to_string(&GetPromptRequestParams::new("brief"))
1303                        .unwrap(),
1304                },
1305                &mut context,
1306            )
1307            .await
1308            .unwrap()
1309            .unwrap();
1310        let prompt_result: GetPromptResult = serde_json::from_str(&prompt.output_json).unwrap();
1311        assert_eq!(prompt_result.messages.len(), 1);
1312
1313        let resource = plugin
1314            .invoke_service(
1315                proto::InvokeServiceRequest {
1316                    kind: proto::ServiceKind::Resource as i32,
1317                    service_name: "demo://state".into(),
1318                    input_json: serde_json::to_string(&ReadResourceRequestParams::new(
1319                        "demo://state",
1320                    ))
1321                    .unwrap(),
1322                },
1323                &mut context,
1324            )
1325            .await
1326            .unwrap()
1327            .unwrap();
1328        let resource_result: ReadResourceResult =
1329            serde_json::from_str(&resource.output_json).unwrap();
1330        assert_eq!(resource_result.contents.len(), 1);
1331
1332        let completion = plugin
1333            .invoke_service(
1334                proto::InvokeServiceRequest {
1335                    kind: proto::ServiceKind::Completion as i32,
1336                    service_name: "brief".into(),
1337                    input_json: serde_json::to_string(&CompleteRequestParams::new(
1338                        Reference::for_prompt("brief"),
1339                        ArgumentInfo {
1340                            name: "topic".into(),
1341                            value: "a".into(),
1342                        },
1343                    ))
1344                    .unwrap(),
1345                },
1346                &mut context,
1347            )
1348            .await
1349            .unwrap()
1350            .unwrap();
1351        let completion_result: CompleteResult =
1352            serde_json::from_str(&completion.output_json).unwrap();
1353        assert_eq!(
1354            completion_result.completion.values,
1355            vec![String::from("alpha")]
1356        );
1357    }
1358
1359    #[tokio::test]
1360    async fn health_request_returns_while_operation_is_running() {
1361        let started = Arc::new(Notify::new());
1362        let started_for_tool = started.clone();
1363        let plugin = plugin! {
1364            metadata: PluginMetadata::new(
1365                "demo",
1366                "1.0.0",
1367                plugin_server_info("demo", "1.0.0", "Demo", "Demo plugin", None::<String>),
1368            ),
1369            mcp: [
1370                mcp::tool("slow")
1371                    .description("Slow operation")
1372                    .input::<DemoArgs>()
1373                    .handle(move |_args, _context| {
1374                        let started = started_for_tool.clone();
1375                        Box::pin(async move {
1376                            started.notify_one();
1377                            tokio::time::sleep(Duration::from_millis(300)).await;
1378                            Ok(json!({ "done": true }))
1379                        })
1380                    }),
1381            ],
1382        };
1383
1384        #[cfg(unix)]
1385        let (plugin_stream, host_stream) = tokio::net::UnixStream::pair().unwrap();
1386        #[cfg(not(unix))]
1387        panic!("runtime stream tests are only implemented for unix");
1388
1389        let runtime = tokio::spawn(PluginRuntime::run_with_stream(
1390            plugin,
1391            LocalStream::Unix(plugin_stream),
1392        ));
1393        let mut host_stream = LocalStream::Unix(host_stream);
1394        write_envelope(
1395            &mut host_stream,
1396            &proto::Envelope {
1397                protocol_version: PROTOCOL_VERSION,
1398                plugin_id: "demo".into(),
1399                request_id: 1,
1400                payload: Some(proto::envelope::Payload::InvokeServiceRequest(
1401                    proto::InvokeServiceRequest {
1402                        kind: proto::ServiceKind::Operation as i32,
1403                        service_name: "slow".into(),
1404                        input_json: "{}".into(),
1405                    },
1406                )),
1407            },
1408        )
1409        .await
1410        .unwrap();
1411
1412        started.notified().await;
1413        write_envelope(
1414            &mut host_stream,
1415            &proto::Envelope {
1416                protocol_version: PROTOCOL_VERSION,
1417                plugin_id: "demo".into(),
1418                request_id: 2,
1419                payload: Some(proto::envelope::Payload::HealthRequest(
1420                    proto::HealthRequest {},
1421                )),
1422            },
1423        )
1424        .await
1425        .unwrap();
1426
1427        let health = timeout(Duration::from_millis(150), read_envelope(&mut host_stream))
1428            .await
1429            .expect("health response should not wait for slow operation")
1430            .unwrap();
1431        assert_eq!(health.request_id, 2);
1432        assert!(matches!(
1433            health.payload,
1434            Some(proto::envelope::Payload::HealthResponse(_))
1435        ));
1436
1437        let invoke = timeout(Duration::from_secs(1), read_envelope(&mut host_stream))
1438            .await
1439            .expect("slow operation should still complete")
1440            .unwrap();
1441        assert_eq!(invoke.request_id, 1);
1442
1443        runtime.abort();
1444    }
1445
1446    #[tokio::test]
1447    async fn invokes_service_requests_concurrently() {
1448        let barrier = Arc::new(Barrier::new(2));
1449        let barrier_for_tool = barrier.clone();
1450        let plugin = plugin! {
1451            metadata: PluginMetadata::new(
1452                "demo",
1453                "1.0.0",
1454                plugin_server_info("demo", "1.0.0", "Demo", "Demo plugin", None::<String>),
1455            ),
1456            mcp: [
1457                mcp::tool("barrier")
1458                    .description("Waits for another request")
1459                    .input::<DemoArgs>()
1460                    .handle(move |_args, _context| {
1461                        let barrier = barrier_for_tool.clone();
1462                        Box::pin(async move {
1463                            barrier.wait().await;
1464                            Ok(json!({ "done": true }))
1465                        })
1466                    }),
1467            ],
1468        };
1469
1470        #[cfg(unix)]
1471        let (plugin_stream, host_stream) = tokio::net::UnixStream::pair().unwrap();
1472        #[cfg(not(unix))]
1473        panic!("runtime stream tests are only implemented for unix");
1474
1475        let runtime = tokio::spawn(PluginRuntime::run_with_stream(
1476            plugin,
1477            LocalStream::Unix(plugin_stream),
1478        ));
1479        let mut host_stream = LocalStream::Unix(host_stream);
1480
1481        for request_id in [1, 2] {
1482            write_envelope(
1483                &mut host_stream,
1484                &proto::Envelope {
1485                    protocol_version: PROTOCOL_VERSION,
1486                    plugin_id: "demo".into(),
1487                    request_id,
1488                    payload: Some(proto::envelope::Payload::InvokeServiceRequest(
1489                        proto::InvokeServiceRequest {
1490                            kind: proto::ServiceKind::Operation as i32,
1491                            service_name: "barrier".into(),
1492                            input_json: "{}".into(),
1493                        },
1494                    )),
1495                },
1496            )
1497            .await
1498            .unwrap();
1499        }
1500
1501        let first = timeout(Duration::from_secs(1), read_envelope(&mut host_stream))
1502            .await
1503            .expect("first invoke response should not block behind another request")
1504            .unwrap();
1505        let second = timeout(Duration::from_secs(1), read_envelope(&mut host_stream))
1506            .await
1507            .expect("second invoke response should not block behind another request")
1508            .unwrap();
1509        let mut request_ids = vec![first.request_id, second.request_id];
1510        request_ids.sort_unstable();
1511        assert_eq!(request_ids, vec![1, 2]);
1512
1513        runtime.abort();
1514    }
1515
1516    #[tokio::test]
1517    async fn dropped_open_mesh_stream_request_removes_pending_response() {
1518        let (outbound_tx, mut outbound_rx) = mpsc::channel(8);
1519        let pending_host_responses = Arc::new(Mutex::new(HashMap::new()));
1520        let mut context =
1521            PluginContext::new("demo".into(), outbound_tx, pending_host_responses.clone());
1522
1523        let request = tokio::spawn(async move {
1524            let _ = context
1525                .open_mesh_stream(proto::OpenMeshStreamRequest::default())
1526                .await;
1527        });
1528
1529        let outbound = outbound_rx
1530            .recv()
1531            .await
1532            .expect("request should be sent before awaiting host response");
1533        assert_ne!(outbound.request_id, 0);
1534        assert_eq!(
1535            pending_host_responses
1536                .lock()
1537                .expect("pending map should not be poisoned")
1538                .len(),
1539            1
1540        );
1541
1542        request.abort();
1543        let _ = request.await;
1544
1545        assert!(
1546            pending_host_responses
1547                .lock()
1548                .expect("pending map should not be poisoned")
1549                .is_empty()
1550        );
1551    }
1552
1553    #[tokio::test]
1554    async fn ordered_notifications_do_not_overtake_each_other() {
1555        let first_started = Arc::new(Notify::new());
1556        let release_first = Arc::new(Notify::new());
1557        let handled = Arc::new(Mutex::new(Vec::<String>::new()));
1558        let first_started_for_handler = first_started.clone();
1559        let release_first_for_handler = release_first.clone();
1560        let handled_for_handler = handled.clone();
1561        let plugin = SimplePlugin::new(PluginMetadata::new(
1562            "demo",
1563            "1.0.0",
1564            plugin_server_info("demo", "1.0.0", "Demo", "Demo plugin", None::<String>),
1565        ))
1566        .on_channel_message(move |message, _context| {
1567            let first_started = first_started_for_handler.clone();
1568            let release_first = release_first_for_handler.clone();
1569            let handled = handled_for_handler.clone();
1570            Box::pin(async move {
1571                if message.message_kind == "first" {
1572                    first_started.notify_one();
1573                    release_first.notified().await;
1574                }
1575                handled
1576                    .lock()
1577                    .expect("handled list should not be poisoned")
1578                    .push(message.message_kind);
1579                Ok(())
1580            })
1581        });
1582
1583        #[cfg(unix)]
1584        let (plugin_stream, host_stream) = tokio::net::UnixStream::pair().unwrap();
1585        #[cfg(not(unix))]
1586        panic!("runtime stream tests are only implemented for unix");
1587
1588        let runtime = tokio::spawn(PluginRuntime::run_with_stream(
1589            plugin,
1590            LocalStream::Unix(plugin_stream),
1591        ));
1592        let mut host_stream = LocalStream::Unix(host_stream);
1593
1594        for (request_id, message_kind) in [(1, "first"), (2, "second")] {
1595            write_envelope(
1596                &mut host_stream,
1597                &proto::Envelope {
1598                    protocol_version: PROTOCOL_VERSION,
1599                    plugin_id: "demo".into(),
1600                    request_id,
1601                    payload: Some(proto::envelope::Payload::ChannelMessage(
1602                        test_channel_message(message_kind),
1603                    )),
1604                },
1605            )
1606            .await
1607            .unwrap();
1608        }
1609
1610        first_started.notified().await;
1611        tokio::time::sleep(Duration::from_millis(50)).await;
1612        assert!(
1613            handled
1614                .lock()
1615                .expect("handled list should not be poisoned")
1616                .is_empty(),
1617            "second message should wait behind the first ordered handler"
1618        );
1619
1620        release_first.notify_one();
1621        timeout(Duration::from_secs(1), async {
1622            loop {
1623                if handled
1624                    .lock()
1625                    .expect("handled list should not be poisoned")
1626                    .len()
1627                    == 2
1628                {
1629                    break;
1630                }
1631                tokio::time::sleep(Duration::from_millis(10)).await;
1632            }
1633        })
1634        .await
1635        .expect("ordered handlers should finish");
1636
1637        assert_eq!(
1638            *handled.lock().expect("handled list should not be poisoned"),
1639            vec![String::from("first"), String::from("second")]
1640        );
1641
1642        runtime.abort();
1643    }
1644}
mesh_llm_plugin/runtime.rs

mesh_llm_plugin/
runtime.rs
