edgecrab_plugins/tool_server/
client.rs

1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::time::{Duration, Instant};
5
6use serde_json::{Value, json};
7use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
8use tokio::process::{Child, ChildStdin, ChildStdout, Command};
9use tokio::sync::Mutex;
10
11use crate::error::PluginError;
12use crate::host_api::{handle_host_request, is_host_method};
13use crate::manifest::{PluginCapabilities, PluginExecConfig, PluginRestartPolicy};
14use edgecrab_tools::registry::ToolContext;
15
16struct ProcessState {
17    child: Child,
18    stdin: ChildStdin,
19    stdout: BufReader<ChildStdout>,
20    restart_count: u32,
21    last_used_at: Instant,
22}
23
24#[derive(Clone)]
25pub struct ToolServerClient {
26    plugin_dir: PathBuf,
27    plugin_name: String,
28    config: PluginExecConfig,
29    capabilities: PluginCapabilities,
30    state: Arc<Mutex<Option<ProcessState>>>,
31    next_id: Arc<AtomicU64>,
32}
33
34impl ToolServerClient {
35    pub fn new(
36        plugin_dir: PathBuf,
37        plugin_name: String,
38        config: PluginExecConfig,
39        capabilities: PluginCapabilities,
40    ) -> Self {
41        Self {
42            plugin_dir,
43            plugin_name,
44            config,
45            capabilities,
46            state: Arc::new(Mutex::new(None)),
47            next_id: Arc::new(AtomicU64::new(1)),
48        }
49    }
50
51    pub async fn shutdown(&self) -> Result<(), PluginError> {
52        let mut state = self.state.lock().await;
53        if let Some(process) = state.as_mut() {
54            let request = json!({
55                "jsonrpc": "2.0",
56                "id": self.next_id.fetch_add(1, Ordering::Relaxed),
57                "method": "shutdown",
58                "params": {},
59            });
60            let _ = process
61                .stdin
62                .write_all(request.to_string().as_bytes())
63                .await;
64            let _ = process.stdin.write_all(b"\n").await;
65            let _ = process.stdin.flush().await;
66            let _ = tokio::time::timeout(Duration::from_secs(2), process.child.wait()).await;
67            let _ = process.child.kill().await;
68        }
69        *state = None;
70        Ok(())
71    }
72
73    pub async fn tool_list(&self) -> Result<Vec<serde_json::Value>, PluginError> {
74        let result = self.rpc("tools/list", json!({}), None).await?;
75        Ok(result
76            .as_array()
77            .cloned()
78            .or_else(|| {
79                result
80                    .get("tools")
81                    .and_then(|value| value.as_array().cloned())
82            })
83            .unwrap_or_default())
84    }
85
86    pub async fn tool_call(
87        &self,
88        name: &str,
89        arguments: serde_json::Value,
90        ctx: &ToolContext,
91    ) -> Result<serde_json::Value, PluginError> {
92        self.rpc(
93            "tools/call",
94            json!({
95                "name": name,
96                "arguments": arguments,
97            }),
98            Some(ctx),
99        )
100        .await
101    }
102
103    pub async fn call_method(
104        &self,
105        method: &str,
106        params: serde_json::Value,
107        ctx: Option<&ToolContext>,
108    ) -> Result<serde_json::Value, PluginError> {
109        self.rpc(method, params, ctx).await
110    }
111
112    async fn rpc(
113        &self,
114        method: &str,
115        params: serde_json::Value,
116        ctx: Option<&ToolContext>,
117    ) -> Result<serde_json::Value, PluginError> {
118        let mut state = self.state.lock().await;
119        self.ensure_process(&mut state).await?;
120
121        let process = state
122            .as_mut()
123            .ok_or_else(|| PluginError::Process("plugin process unavailable".into()))?;
124        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
125        let request = json!({
126            "jsonrpc": "2.0",
127            "id": id,
128            "method": method,
129            "params": params,
130        });
131        process
132            .stdin
133            .write_all(request.to_string().as_bytes())
134            .await?;
135        process.stdin.write_all(b"\n").await?;
136        process.stdin.flush().await?;
137
138        let mut line = String::new();
139        loop {
140            line.clear();
141            let read = tokio::time::timeout(
142                Duration::from_secs(self.config.call_timeout_secs.max(1)),
143                process.stdout.read_line(&mut line),
144            )
145            .await
146            .map_err(|_| PluginError::Rpc(format!("timeout waiting for {method} response")))??;
147            if read == 0 {
148                self.handle_process_failure(&mut state, method).await?;
149                return Err(PluginError::Process(format!(
150                    "plugin process closed stdout during {method}"
151                )));
152            }
153            process.last_used_at = Instant::now();
154            let response: serde_json::Value = serde_json::from_str(line.trim())?;
155            if let Some(host_method) = response.get("method").and_then(|value| value.as_str()) {
156                if is_host_method(host_method) {
157                    let Some(ctx) = ctx else {
158                        return Err(PluginError::Rpc(format!(
159                            "plugin requested {host_method} without a host context"
160                        )));
161                    };
162                    let host_response =
163                        handle_host_request(&self.plugin_name, &self.capabilities, &response, ctx)
164                            .await;
165                    process
166                        .stdin
167                        .write_all(host_response.to_string().as_bytes())
168                        .await?;
169                    process.stdin.write_all(b"\n").await?;
170                    process.stdin.flush().await?;
171                    continue;
172                }
173                tracing::debug!(plugin = %self.plugin_name, method = host_method, "ignoring plugin notification");
174                continue;
175            }
176            if response.get("id") != Some(&json!(id)) {
177                continue;
178            }
179            if let Some(error) = response.get("error") {
180                return Err(PluginError::Rpc(error.to_string()));
181            }
182            return Ok(extract_result_payload(
183                response.get("result").cloned().unwrap_or(Value::Null),
184            ));
185        }
186    }
187
188    async fn ensure_process(&self, state: &mut Option<ProcessState>) -> Result<(), PluginError> {
189        let should_respawn = match state.as_mut() {
190            None => true,
191            Some(process) => {
192                let idle_timeout = self.config.idle_timeout_secs.max(1);
193                if process.last_used_at.elapsed() >= Duration::from_secs(idle_timeout) {
194                    let _ = process.child.kill().await;
195                    true
196                } else {
197                    process.child.try_wait()?.is_some()
198                }
199            }
200        };
201
202        if should_respawn {
203            let restart_count = state.as_ref().map_or(0, |process| process.restart_count);
204            *state = Some(self.spawn(restart_count).await?);
205        }
206        Ok(())
207    }
208
209    async fn handle_process_failure(
210        &self,
211        state: &mut Option<ProcessState>,
212        method: &str,
213    ) -> Result<(), PluginError> {
214        let restart_count = state.as_ref().map_or(0, |process| process.restart_count);
215        if !self.should_restart(restart_count) {
216            *state = None;
217            return Err(PluginError::Process(format!(
218                "plugin process crashed during {method} and restart policy prevents recovery"
219            )));
220        }
221        *state = Some(self.spawn(restart_count + 1).await?);
222        Ok(())
223    }
224
225    fn should_restart(&self, restart_count: u32) -> bool {
226        match self.config.restart_policy {
227            PluginRestartPolicy::Never => false,
228            PluginRestartPolicy::Once => restart_count == 0,
229            PluginRestartPolicy::Always => restart_count < self.config.restart_max_attempts,
230        }
231    }
232
233    async fn spawn(&self, restart_count: u32) -> Result<ProcessState, PluginError> {
234        let mut command = Command::new(resolve_command(&self.plugin_dir, &self.config.command));
235        command.args(&self.config.args);
236        command.current_dir(resolve_cwd(&self.plugin_dir, self.config.cwd.as_deref()));
237        command.envs(&self.config.env);
238        command.stdin(std::process::Stdio::piped());
239        command.stdout(std::process::Stdio::piped());
240        command.stderr(std::process::Stdio::inherit());
241
242        let mut child = command.spawn()?;
243        let stdin = child
244            .stdin
245            .take()
246            .ok_or_else(|| PluginError::Process("missing child stdin".into()))?;
247        let stdout = child
248            .stdout
249            .take()
250            .ok_or_else(|| PluginError::Process("missing child stdout".into()))?;
251        let mut state = ProcessState {
252            child,
253            stdin,
254            stdout: BufReader::new(stdout),
255            restart_count,
256            last_used_at: Instant::now(),
257        };
258
259        let init = json!({
260            "jsonrpc": "2.0",
261            "id": 0,
262            "method": "initialize",
263            "params": {
264                "protocolVersion": "2024-11-05",
265                "capabilities": {},
266                "clientInfo": {
267                    "name": "edgecrab",
268                    "version": env!("CARGO_PKG_VERSION"),
269                }
270            },
271        });
272        state.stdin.write_all(init.to_string().as_bytes()).await?;
273        state.stdin.write_all(b"\n").await?;
274        state.stdin.flush().await?;
275
276        let mut line = String::new();
277        let read = tokio::time::timeout(
278            Duration::from_secs(self.config.startup_timeout_secs.max(1)),
279            state.stdout.read_line(&mut line),
280        )
281        .await
282        .map_err(|_| PluginError::Rpc("tool server initialize timeout".into()))??;
283        if read == 0 {
284            return Err(PluginError::Process(
285                "tool server exited before initialize completed".into(),
286            ));
287        }
288        let init_response: serde_json::Value = serde_json::from_str(line.trim())?;
289        if let Some(error) = init_response.get("error") {
290            return Err(PluginError::Rpc(error.to_string()));
291        }
292        let initialized = json!({
293            "jsonrpc": "2.0",
294            "method": "notifications/initialized",
295        });
296        state
297            .stdin
298            .write_all(initialized.to_string().as_bytes())
299            .await?;
300        state.stdin.write_all(b"\n").await?;
301        state.stdin.flush().await?;
302        state.last_used_at = Instant::now();
303        Ok(state)
304    }
305}
306
307fn extract_result_payload(result: serde_json::Value) -> serde_json::Value {
308    if result
309        .get("isError")
310        .and_then(|value| value.as_bool())
311        .unwrap_or(false)
312    {
313        return json!({
314            "isError": true,
315            "content": result.get("content").cloned().unwrap_or(Value::Null),
316        });
317    }
318    if let Some(content) = result.get("content").and_then(|value| value.as_array()) {
319        let text = content
320            .iter()
321            .filter(|item| item.get("type").and_then(|value| value.as_str()) == Some("text"))
322            .map(|item| {
323                item.get("text")
324                    .and_then(|value| value.as_str())
325                    .unwrap_or_default()
326            })
327            .collect::<Vec<_>>()
328            .join("\n");
329        if !text.is_empty() {
330            return Value::String(text);
331        }
332    }
333    result
334}
335
336fn resolve_command(plugin_dir: &Path, command: &str) -> PathBuf {
337    let command_path = PathBuf::from(command);
338    if command_path.is_absolute() {
339        command_path
340    } else if command.starts_with("./") || command.starts_with("../") {
341        plugin_dir.join(command)
342    } else {
343        command_path
344    }
345}
346
347fn resolve_cwd(plugin_dir: &Path, cwd: Option<&str>) -> PathBuf {
348    match cwd {
349        Some(".") | None => plugin_dir.to_path_buf(),
350        Some(cwd) => plugin_dir.join(cwd),
351    }
352}
353
354#[cfg(test)]
355mod tests {
356    use super::*;
357
358    #[test]
359    fn resolve_relative_command_against_plugin_dir() {
360        let dir = PathBuf::from("/tmp/plugin");
361        assert_eq!(
362            resolve_command(&dir, "./bin/server"),
363            dir.join("./bin/server")
364        );
365        assert_eq!(resolve_command(&dir, "python3"), PathBuf::from("python3"));
366    }
367}
edgecrab_plugins/tool_server/client.rs

edgecrab_plugins/tool_server/
client.rs
