mesh_llm_plugin/
context.rs

1use anyhow::{Result, bail};
2use serde::Serialize;
3use std::collections::HashMap;
4use std::marker::PhantomData;
5use std::sync::atomic::{AtomicU64, Ordering};
6use std::sync::{Arc, Mutex};
7use tokio::sync::{mpsc, oneshot};
8
9use crate::{
10    PROTOCOL_VERSION,
11    helpers::{channel_message, json_channel_message},
12    io::{LocalStream, connect_side_stream},
13    proto,
14};
15
16static NEXT_HOST_REQUEST_ID: AtomicU64 = AtomicU64::new(1);
17const PLUGIN_ORIGINATED_REQUEST_BIT: u64 = 1 << 63;
18
19pub(crate) type PendingHostResponses =
20    Arc<Mutex<HashMap<u64, oneshot::Sender<Result<proto::Envelope>>>>>;
21
22struct PendingHostResponseGuard {
23    request_id: u64,
24    pending_host_responses: PendingHostResponses,
25    active: bool,
26}
27
28impl PendingHostResponseGuard {
29    fn new(request_id: u64, pending_host_responses: PendingHostResponses) -> Self {
30        Self {
31            request_id,
32            pending_host_responses,
33            active: true,
34        }
35    }
36
37    fn disarm(&mut self) {
38        self.active = false;
39    }
40}
41
42impl Drop for PendingHostResponseGuard {
43    fn drop(&mut self) {
44        if self.active {
45            remove_pending_host_response(&self.pending_host_responses, self.request_id);
46        }
47    }
48}
49
50pub struct PluginContext<'a> {
51    pub(crate) outbound_tx: mpsc::Sender<proto::Envelope>,
52    pub(crate) pending_host_responses: PendingHostResponses,
53    pub(crate) plugin_id: String,
54    pub(crate) _marker: PhantomData<&'a mut ()>,
55}
56
57impl<'a> PluginContext<'a> {
58    pub(crate) fn new(
59        plugin_id: String,
60        outbound_tx: mpsc::Sender<proto::Envelope>,
61        pending_host_responses: PendingHostResponses,
62    ) -> Self {
63        Self {
64            outbound_tx,
65            pending_host_responses,
66            plugin_id,
67            _marker: PhantomData,
68        }
69    }
70
71    pub async fn send_channel(&mut self, message: proto::ChannelMessage) -> Result<()> {
72        self.send_channel_message(message).await
73    }
74
75    pub async fn send_channel_message(&mut self, message: proto::ChannelMessage) -> Result<()> {
76        self.send_payload(proto::envelope::Payload::ChannelMessage(message), 0)
77            .await
78    }
79
80    pub async fn send_text_channel(
81        &mut self,
82        channel: impl Into<String>,
83        target_peer_id: impl Into<String>,
84        message_kind: impl Into<String>,
85        text: impl Into<String>,
86    ) -> Result<()> {
87        self.send_channel_message(channel_message(
88            channel,
89            target_peer_id,
90            "text/plain",
91            text.into().into_bytes(),
92            message_kind,
93        ))
94        .await
95    }
96
97    pub async fn send_json_channel<T: Serialize>(
98        &mut self,
99        channel: impl Into<String>,
100        target_peer_id: impl Into<String>,
101        message_kind: impl Into<String>,
102        payload: &T,
103    ) -> Result<()> {
104        self.send_channel_message(json_channel_message(
105            channel,
106            target_peer_id,
107            message_kind,
108            payload,
109        )?)
110        .await
111    }
112
113    pub async fn send_bulk(&mut self, message: proto::BulkTransferMessage) -> Result<()> {
114        self.send_bulk_transfer_message(message).await
115    }
116
117    pub async fn send_bulk_transfer_message(
118        &mut self,
119        message: proto::BulkTransferMessage,
120    ) -> Result<()> {
121        self.send_payload(proto::envelope::Payload::BulkTransferMessage(message), 0)
122            .await
123    }
124
125    pub async fn notify_host<P>(&mut self, method: &str, params: P) -> Result<()>
126    where
127        P: Serialize,
128    {
129        self.send_payload(
130            proto::envelope::Payload::RpcNotification(proto::RpcNotification {
131                method: method.to_string(),
132                params_json: serde_json::to_string(&params)?,
133            }),
134            0,
135        )
136        .await
137    }
138
139    pub async fn open_mesh_stream(
140        &mut self,
141        request: proto::OpenMeshStreamRequest,
142    ) -> Result<proto::OpenMeshStreamResponse> {
143        let request_id = next_host_request_id();
144        let (tx, rx) = oneshot::channel();
145        insert_pending_host_response(&self.pending_host_responses, request_id, tx);
146        let mut pending_guard =
147            PendingHostResponseGuard::new(request_id, self.pending_host_responses.clone());
148
149        self.send_payload(
150            proto::envelope::Payload::OpenMeshStreamRequest(request),
151            request_id,
152        )
153        .await?;
154
155        let response = rx.await??;
156        pending_guard.disarm();
157        match response.payload {
158            Some(proto::envelope::Payload::OpenMeshStreamResponse(response)) => Ok(response),
159            Some(proto::envelope::Payload::ErrorResponse(error)) => bail!(error.message),
160            _ => bail!("Host returned an unexpected open_mesh_stream response"),
161        }
162    }
163
164    pub async fn connect_mesh_stream(
165        &mut self,
166        request: proto::OpenMeshStreamRequest,
167    ) -> Result<LocalStream> {
168        let response = self.open_mesh_stream(request).await?;
169        if !response.accepted {
170            bail!(
171                "Host rejected mesh stream: {}",
172                response
173                    .message
174                    .unwrap_or_else(|| "no reason provided".into())
175            );
176        }
177        let endpoint = response
178            .endpoint
179            .as_deref()
180            .ok_or_else(|| anyhow::anyhow!("Host accepted mesh stream without an endpoint"))?;
181        connect_side_stream(endpoint, response.transport_kind).await
182    }
183
184    async fn send_payload(&self, payload: proto::envelope::Payload, request_id: u64) -> Result<()> {
185        self.outbound_tx
186            .send(proto::Envelope {
187                protocol_version: PROTOCOL_VERSION,
188                plugin_id: self.plugin_id.clone(),
189                request_id,
190                payload: Some(payload),
191            })
192            .await
193            .map_err(|_| anyhow::anyhow!("plugin host connection is closed"))
194    }
195}
196
197pub(crate) fn next_host_request_id() -> u64 {
198    PLUGIN_ORIGINATED_REQUEST_BIT | NEXT_HOST_REQUEST_ID.fetch_add(1, Ordering::Relaxed)
199}
200
201pub(crate) fn insert_pending_host_response(
202    pending_host_responses: &PendingHostResponses,
203    request_id: u64,
204    sender: oneshot::Sender<Result<proto::Envelope>>,
205) {
206    pending_host_responses
207        .lock()
208        .expect("pending host response map poisoned")
209        .insert(request_id, sender);
210}
211
212pub(crate) fn remove_pending_host_response(
213    pending_host_responses: &PendingHostResponses,
214    request_id: u64,
215) -> Option<oneshot::Sender<Result<proto::Envelope>>> {
216    pending_host_responses
217        .lock()
218        .expect("pending host response map poisoned")
219        .remove(&request_id)
220}
221
222pub(crate) fn drain_pending_host_responses(
223    pending_host_responses: &PendingHostResponses,
224) -> Vec<oneshot::Sender<Result<proto::Envelope>>> {
225    pending_host_responses
226        .lock()
227        .expect("pending host response map poisoned")
228        .drain()
229        .map(|(_, sender)| sender)
230        .collect()
231}
mesh_llm_plugin/context.rs

mesh_llm_plugin/
context.rs
