1use anyhow::{Context, Result, bail};
2use prost::Message;
3#[cfg(unix)]
4use std::path::PathBuf;
5use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
6
7use crate::{PROTOCOL_VERSION, proto};
8
9pub enum LocalStream {
10 #[cfg(unix)]
11 Unix(tokio::net::UnixStream),
12 #[cfg(windows)]
13 PipeClient(tokio::net::windows::named_pipe::NamedPipeClient),
14 #[cfg(windows)]
15 PipeServer(tokio::net::windows::named_pipe::NamedPipeServer),
16}
17
18pub type LocalReadHalf = Box<dyn AsyncRead + Send + Unpin>;
19pub type LocalWriteHalf = Box<dyn AsyncWrite + Send + Unpin>;
20
21impl LocalStream {
22 pub fn into_split(self) -> (LocalReadHalf, LocalWriteHalf) {
23 match self {
24 #[cfg(unix)]
25 LocalStream::Unix(stream) => {
26 let (read, write) = stream.into_split();
27 (Box::new(read), Box::new(write))
28 }
29 #[cfg(windows)]
30 LocalStream::PipeClient(stream) => {
31 let (read, write) = tokio::io::split(stream);
32 (Box::new(read), Box::new(write))
33 }
34 #[cfg(windows)]
35 LocalStream::PipeServer(stream) => {
36 let (read, write) = tokio::io::split(stream);
37 (Box::new(read), Box::new(write))
38 }
39 }
40 }
41
42 async fn write_all(&mut self, bytes: &[u8]) -> Result<()> {
43 match self {
44 #[cfg(unix)]
45 LocalStream::Unix(stream) => stream.write_all(bytes).await?,
46 #[cfg(windows)]
47 LocalStream::PipeClient(stream) => stream.write_all(bytes).await?,
48 #[cfg(windows)]
49 LocalStream::PipeServer(stream) => stream.write_all(bytes).await?,
50 }
51 Ok(())
52 }
53
54 async fn read_exact(&mut self, bytes: &mut [u8]) -> Result<()> {
55 match self {
56 #[cfg(unix)]
57 LocalStream::Unix(stream) => {
58 let _ = stream.read_exact(bytes).await?;
59 }
60 #[cfg(windows)]
61 LocalStream::PipeClient(stream) => {
62 let _ = stream.read_exact(bytes).await?;
63 }
64 #[cfg(windows)]
65 LocalStream::PipeServer(stream) => {
66 let _ = stream.read_exact(bytes).await?;
67 }
68 }
69 Ok(())
70 }
71
72 pub async fn write_all_bytes(&mut self, bytes: &[u8]) -> Result<()> {
73 self.write_all(bytes).await
74 }
75
76 pub async fn read_exact_bytes(&mut self, bytes: &mut [u8]) -> Result<()> {
77 self.read_exact(bytes).await
78 }
79}
80
81pub enum LocalListener {
82 #[cfg(unix)]
83 Unix(tokio::net::UnixListener, PathBuf),
84 #[cfg(windows)]
85 Pipe(String, tokio::net::windows::named_pipe::NamedPipeServer),
86}
87
88impl LocalListener {
89 pub async fn accept(self) -> Result<LocalStream> {
90 match self {
91 #[cfg(unix)]
92 LocalListener::Unix(listener, path) => {
93 let (stream, _) = listener.accept().await?;
94 let _ = std::fs::remove_file(path);
95 Ok(LocalStream::Unix(stream))
96 }
97 #[cfg(windows)]
98 LocalListener::Pipe(_name, server) => {
99 server.connect().await?;
100 Ok(LocalStream::PipeServer(server))
101 }
102 }
103 }
104
105 pub fn endpoint(&self) -> String {
106 match self {
107 #[cfg(unix)]
108 LocalListener::Unix(_, path) => path.display().to_string(),
109 #[cfg(windows)]
110 LocalListener::Pipe(name, _) => name.clone(),
111 }
112 }
113
114 pub fn transport_kind(&self) -> i32 {
115 #[cfg(unix)]
116 {
117 proto::StreamTransportKind::StreamUnixSocket as i32
118 }
119 #[cfg(windows)]
120 {
121 proto::StreamTransportKind::StreamNamedPipe as i32
122 }
123 }
124
125 pub fn open_stream_response(
126 &self,
127 request: &proto::OpenStreamRequest,
128 ) -> proto::OpenStreamResponse {
129 proto::OpenStreamResponse {
130 stream_id: request.stream_id.clone(),
131 accepted: true,
132 transport_kind: self.transport_kind(),
133 endpoint: Some(self.endpoint()),
134 token: None,
135 expires_at_unix_ms: None,
136 message: None,
137 }
138 }
139}
140
141pub async fn connect_from_env() -> Result<LocalStream> {
142 let endpoint = std::env::var("MESH_LLM_PLUGIN_ENDPOINT")
143 .context("MESH_LLM_PLUGIN_ENDPOINT is not set for plugin process")?;
144 let transport =
145 std::env::var("MESH_LLM_PLUGIN_TRANSPORT").unwrap_or_else(|_| default_transport().into());
146
147 match transport.as_str() {
148 #[cfg(unix)]
149 "unix" => Ok(LocalStream::Unix(
150 tokio::net::UnixStream::connect(&endpoint).await?,
151 )),
152 #[cfg(windows)]
153 "pipe" => Ok(LocalStream::PipeClient(
154 tokio::net::windows::named_pipe::ClientOptions::new().open(&endpoint)?,
155 )),
156 _ => bail!("Unsupported plugin transport '{transport}'"),
157 }
158}
159
160pub async fn connect_side_stream(endpoint: &str, transport_kind: i32) -> Result<LocalStream> {
161 match proto::StreamTransportKind::try_from(transport_kind)
162 .unwrap_or(proto::StreamTransportKind::Unspecified)
163 {
164 #[cfg(unix)]
165 proto::StreamTransportKind::StreamUnixSocket => Ok(LocalStream::Unix(
166 tokio::net::UnixStream::connect(endpoint)
167 .await
168 .with_context(|| format!("Failed to connect side stream socket {endpoint}"))?,
169 )),
170 #[cfg(windows)]
171 proto::StreamTransportKind::StreamNamedPipe => Ok(LocalStream::PipeClient(
172 tokio::net::windows::named_pipe::ClientOptions::new()
173 .open(endpoint)
174 .with_context(|| format!("Failed to connect side stream pipe {endpoint}"))?,
175 )),
176 _ => bail!("Unsupported side stream transport kind '{transport_kind}'"),
177 }
178}
179
180pub async fn bind_side_stream(plugin_id: &str, stream_id: &str) -> Result<LocalListener> {
181 #[cfg(unix)]
182 {
183 let path = std::env::temp_dir().join(format!(
184 "mesh-llm-side-{}-{}.sock",
185 sanitize_component(plugin_id),
186 sanitize_component(stream_id)
187 ));
188 if path.exists() {
189 let _ = std::fs::remove_file(&path);
190 }
191 let listener = tokio::net::UnixListener::bind(&path)
192 .with_context(|| format!("Failed to bind side stream socket {}", path.display()))?;
193 Ok(LocalListener::Unix(listener, path))
194 }
195 #[cfg(windows)]
196 {
197 let endpoint = format!(
198 r"\\.\pipe\mesh-llm-side-{}-{}",
199 sanitize_component(plugin_id),
200 sanitize_component(stream_id)
201 );
202 let server = tokio::net::windows::named_pipe::ServerOptions::new()
203 .create(&endpoint)
204 .with_context(|| format!("Failed to create side stream pipe {endpoint}"))?;
205 return Ok(LocalListener::Pipe(endpoint, server));
206 }
207}
208
209pub async fn write_envelope_to<W>(stream: &mut W, envelope: &proto::Envelope) -> Result<()>
210where
211 W: AsyncWrite + Unpin + ?Sized,
212{
213 let mut body = Vec::new();
214 envelope.encode(&mut body)?;
215 stream.write_all(&(body.len() as u32).to_le_bytes()).await?;
216 stream.write_all(&body).await?;
217 Ok(())
218}
219
220pub async fn write_envelope(stream: &mut LocalStream, envelope: &proto::Envelope) -> Result<()> {
221 let mut body = Vec::new();
222 envelope.encode(&mut body)?;
223 stream.write_all(&(body.len() as u32).to_le_bytes()).await?;
224 stream.write_all(&body).await?;
225 Ok(())
226}
227
228pub async fn read_envelope_from<R>(stream: &mut R) -> Result<proto::Envelope>
229where
230 R: AsyncRead + Unpin + ?Sized,
231{
232 let mut len_buf = [0u8; 4];
233 stream.read_exact(&mut len_buf).await?;
234 let len = u32::from_le_bytes(len_buf) as usize;
235 if len > 16 * 1024 * 1024 {
236 bail!("Plugin frame too large");
237 }
238 let mut body = vec![0u8; len];
239 stream.read_exact(&mut body).await?;
240 Ok(proto::Envelope::decode(body.as_slice())?)
241}
242
243pub async fn read_envelope(stream: &mut LocalStream) -> Result<proto::Envelope> {
244 let mut len_buf = [0u8; 4];
245 stream.read_exact(&mut len_buf).await?;
246 let len = u32::from_le_bytes(len_buf) as usize;
247 if len > 16 * 1024 * 1024 {
248 bail!("Plugin frame too large");
249 }
250 let mut body = vec![0u8; len];
251 stream.read_exact(&mut body).await?;
252 Ok(proto::Envelope::decode(body.as_slice())?)
253}
254
255pub async fn send_channel_message(
256 stream: &mut LocalStream,
257 plugin_id: &str,
258 message: proto::ChannelMessage,
259) -> Result<()> {
260 write_envelope(
261 stream,
262 &proto::Envelope {
263 protocol_version: PROTOCOL_VERSION,
264 plugin_id: plugin_id.to_string(),
265 request_id: 0,
266 payload: Some(proto::envelope::Payload::ChannelMessage(message)),
267 },
268 )
269 .await
270}
271
272pub async fn send_bulk_transfer_message(
273 stream: &mut LocalStream,
274 plugin_id: &str,
275 message: proto::BulkTransferMessage,
276) -> Result<()> {
277 write_envelope(
278 stream,
279 &proto::Envelope {
280 protocol_version: PROTOCOL_VERSION,
281 plugin_id: plugin_id.to_string(),
282 request_id: 0,
283 payload: Some(proto::envelope::Payload::BulkTransferMessage(message)),
284 },
285 )
286 .await
287}
288
289fn default_transport() -> &'static str {
290 #[cfg(unix)]
291 {
292 "unix"
293 }
294 #[cfg(windows)]
295 {
296 "pipe"
297 }
298}
299
300fn sanitize_component(value: &str) -> String {
301 value
302 .chars()
303 .map(|ch| {
304 if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' {
305 ch
306 } else {
307 '_'
308 }
309 })
310 .collect()
311}