Skip to main content

zerolaunch_plugin_sdk_rust/
host_proxy.rs

1//! HostProxy — 第三方插件调用 host/* API 的代理。
2//!
3//! 每个方法经共享出站通道发送 LSP 帧 JSON-RPC 请求,并通过注册在共享
4//! pending map 中的 oneshot 等待响应。该设计将 stdin 读取与 stdout 写入
5//! 集中到独立任务,避免旧同步 stdin 锁方案的死锁。
6use dashmap::DashMap;
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::sync::Arc;
9use std::time::Duration;
10use tokio::sync::{mpsc, oneshot};
11use zerolaunch_plugin_api::services::model::{
12    ModelChatRequest, ModelChatResponse, ModelEmbeddingRequest, ModelEmbeddingResponse, ModelInfo,
13    ModelSimilarityRequest, ModelSimilarityResponse,
14};
15use zerolaunch_plugin_protocol::JsonRpcError;
16
17use base64::Engine as _;
18
19/// Proxy for calling host-side APIs from a plugin subprocess.
20/// Does NOT access stdin/stdout directly — uses channel-based I/O.
21pub struct HostProxy {
22    /// 请求 id 分配器(单调递增,从 1 开始)。
23    next_id: AtomicU64,
24    /// 在途请求表:id → 响应通道。read_task 完成响应后移除条目。
25    /// 值类型为 `Result`:宿主错误经 Err(JsonRpcError) 返回,调用方区分错误与结果。
26    pending: Arc<DashMap<u64, oneshot::Sender<Result<serde_json::Value, JsonRpcError>>>>,
27    /// 出站帧通道:write_task 独占消费并写 stdout。
28    outbound_tx: mpsc::Sender<Vec<u8>>,
29}
30
31impl HostProxy {
32    pub fn new(
33        pending: Arc<DashMap<u64, oneshot::Sender<Result<serde_json::Value, JsonRpcError>>>>,
34        outbound_tx: mpsc::Sender<Vec<u8>>,
35    ) -> Self {
36        Self {
37            next_id: AtomicU64::new(1),
38            pending,
39            outbound_tx,
40        }
41    }
42
43    /// Send a host/* request via the shared stdout channel and await the response
44    /// through the shared pending map.
45    async fn send_request(
46        &self,
47        method: &str,
48        params: serde_json::Value,
49    ) -> Result<serde_json::Value, String> {
50        // 模型类方法耗时可能远超普通 host/*(LLM 生成 / 大批量 embedding),
51        // 单独放宽超时;其余方法维持 30s。
52        let is_model_call = method.starts_with("host/model.");
53        let timeout = if is_model_call {
54            Duration::from_secs(300)
55        } else {
56            Duration::from_secs(30)
57        };
58        self.send_request_inner(method, params, timeout).await
59    }
60
61    /// 发送请求并等待响应;超时/通道关闭返回错误字符串。
62    async fn send_request_inner(
63        &self,
64        method: &str,
65        params: serde_json::Value,
66        timeout: Duration,
67    ) -> Result<serde_json::Value, String> {
68        let id = self.next_id.fetch_add(1, Ordering::SeqCst);
69        let request = serde_json::json!({
70            "jsonrpc": "2.0",
71            "id": id,
72            "method": method,
73            "params": params,
74        });
75
76        let payload = serde_json::to_vec(&request).map_err(|e| e.to_string())?;
77
78        // Register pending response
79        let (tx, rx) = oneshot::channel();
80        self.pending.insert(id, tx);
81
82        // Send via the shared channel (write_task writes to stdout).
83        // 注意:此处只投递干净 JSON,分帧统一由 write_task 完成,
84        // 若在此 encode_frame 会导致双重分帧(帧内嵌套帧),宿主无法解析。
85        if self.outbound_tx.send(payload).await.is_err() {
86            self.pending.remove(&id);
87            return Err("write channel closed".to_string());
88        }
89
90        // Await the response (read_task completes the oneshot with Result).
91        // Apply a timeout so the plugin doesn't hang forever if
92        // the host crashes during request processing.
93        match tokio::time::timeout(timeout, rx).await {
94            Ok(Ok(Ok(value))) => Ok(value),
95            Ok(Ok(Err(err))) => Err(format!(
96                "host call failed (code {}): {}",
97                err.code, err.message
98            )),
99            Ok(Err(_)) => Err("response channel closed".to_string()),
100            Err(_) => {
101                // 超时:移除 pending 条目,防止长驻插件无界泄漏。
102                self.pending.remove(&id);
103                Err("host call timed out".to_string())
104            }
105        }
106    }
107
108    pub async fn log(&self, level: &str, message: &str) -> Result<(), String> {
109        self.send_request(
110            "host/log",
111            serde_json::json!({ "level": level, "message": message }),
112        )
113        .await?;
114        Ok(())
115    }
116
117    /// 发送 host/log 请求但不等待响应(fire-and-forget)。
118    ///
119    /// pending 条目在宿主响应到达时由 read_task 自动清理。
120    /// 若 outbound 通道已满,日志被静默丢弃并从 pending 中移除,
121    /// 避免阻塞调用者(通常来自 tracing subscriber 的回调)。
122    pub fn log_no_wait(&self, level: &str, message: &str) {
123        let id = self.next_id.fetch_add(1, Ordering::SeqCst);
124        let Ok(payload) = serde_json::to_vec(&serde_json::json!({
125            "jsonrpc": "2.0",
126            "id": id,
127            "method": "host/log",
128            "params": { "level": level, "message": message },
129        })) else {
130            return;
131        };
132
133        let (tx, _rx) = oneshot::channel(); // _rx 立即 drop → fire-and-forget
134        self.pending.insert(id, tx);
135
136        // 非阻塞投递:通道满了则丢弃并清理 pending。
137        // 同 send_request:只投递干净 JSON,分帧由 write_task 统一完成。
138        if self.outbound_tx.try_send(payload).is_err() {
139            self.pending.remove(&id);
140        }
141    }
142
143    pub async fn shell_open(&self, target: &str) -> Result<(), String> {
144        self.send_request("host/shell.open", serde_json::json!({ "target": target }))
145            .await?;
146        Ok(())
147    }
148
149    /// 获取图标字节(base64 字符串)。
150    /// 返回:图标字节的 base64(WebP,回退可能为 PNG);失败为空字符串。
151    pub async fn get_icon(&self, path: &str) -> Result<String, String> {
152        let result = self
153            .send_request(
154                "host/icon.get",
155                serde_json::json!({ "request": { "path": path }, "level": "Full" }),
156            )
157            .await?;
158        Ok(result.as_str().unwrap_or("").to_string())
159    }
160
161    pub async fn shell_execute_command(&self, cmd: &str) -> Result<(), String> {
162        self.send_request(
163            "host/shell.execute_command",
164            serde_json::json!({ "cmd": cmd }),
165        )
166        .await?;
167        Ok(())
168    }
169
170    pub async fn shell_open_folder(&self, path: &str) -> Result<(), String> {
171        self.send_request(
172            "host/shell.open_folder",
173            serde_json::json!({ "path": path }),
174        )
175        .await?;
176        Ok(())
177    }
178
179    pub async fn shell_execute_elevation(&self, path: &str) -> Result<(), String> {
180        self.send_request(
181            "host/shell.execute_elevation",
182            serde_json::json!({ "path": path }),
183        )
184        .await?;
185        Ok(())
186    }
187
188    pub async fn notify(&self, title: &str, message: &str) -> Result<(), String> {
189        self.send_request(
190            "host/notify",
191            serde_json::json!({ "title": title, "message": message }),
192        )
193        .await?;
194        Ok(())
195    }
196
197    /// 获取宿主当前界面语言(如 "zh-Hans"),供插件生成本地化文本。
198    pub async fn get_locale(&self) -> Result<String, String> {
199        let result = self
200            .send_request("host/i18n.get_locale", serde_json::json!(null))
201            .await?;
202        Ok(result.as_str().unwrap_or("").to_string())
203    }
204
205    /// 查询宿主当前实际生效主题,返回 `light` 或 `dark`。
206    pub async fn get_theme(&self) -> Result<String, String> {
207        let result = self
208            .send_request("host/theme.get", serde_json::Value::Null)
209            .await?;
210        result
211            .as_str()
212            .map(str::to_string)
213            .ok_or_else(|| "host theme response is not a string".to_string())
214    }
215
216    /// 全网模型清单(聚合所有提供方)。
217    pub async fn model_list(&self) -> Result<Vec<ModelInfo>, String> {
218        let result = self
219            .send_request("host/model.list", serde_json::Value::Null)
220            .await?;
221        serde_json::from_value(result).map_err(|e| e.to_string())
222    }
223
224    /// 按 model_id 调用文本生成。
225    pub async fn model_chat(&self, req: ModelChatRequest) -> Result<ModelChatResponse, String> {
226        let params = serde_json::to_value(req).map_err(|e| e.to_string())?;
227        let result = self.send_request("host/model.chat", params).await?;
228        serde_json::from_value(result).map_err(|e| e.to_string())
229    }
230
231    /// 按 model_id 调用文本向量化(task_type 必填,宿主对缺失/未知值返回错误)。
232    pub async fn model_embedding(
233        &self,
234        req: ModelEmbeddingRequest,
235    ) -> Result<ModelEmbeddingResponse, String> {
236        let params = serde_json::to_value(req).map_err(|e| e.to_string())?;
237        let result = self.send_request("host/model.embedding", params).await?;
238        serde_json::from_value(result).map_err(|e| e.to_string())
239    }
240
241    /// 按 model_id 计算查询向量与多个目标向量的相似度。
242    pub async fn model_similarity(
243        &self,
244        req: ModelSimilarityRequest,
245    ) -> Result<ModelSimilarityResponse, String> {
246        let params = serde_json::to_value(req).map_err(|e| e.to_string())?;
247        let result = self.send_request("host/model.similarity", params).await?;
248        serde_json::from_value(result).map_err(|e| e.to_string())
249    }
250
251    pub async fn enumerate_apps(&self) -> Result<serde_json::Value, String> {
252        self.send_request("host/app.enumerate", serde_json::json!(null))
253            .await
254    }
255
256    pub async fn resolve_path(&self, kind: &str) -> Result<String, String> {
257        let result = self
258            .send_request("host/path.resolve", serde_json::json!({ "kind": kind }))
259            .await?;
260        Ok(result.as_str().unwrap_or("").to_string())
261    }
262
263    /// 上传插件本地文件到宿主资源空间。
264    /// 直接传递文件路径,由宿主负责读取。
265    pub async fn resource_upload(
266        &self,
267        resource_id: &str,
268        file_path: &str,
269        max_size: Option<u64>,
270    ) -> Result<String, String> {
271        let result = self
272            .send_request(
273                "host/resource.upload",
274                serde_json::json!({
275                    "resourceId": resource_id,
276                    "filePath": file_path,
277                    "maxSize": max_size,
278                }),
279            )
280            .await?;
281        Ok(result.as_str().unwrap_or("").to_string())
282    }
283
284    pub async fn resource_get(&self, resource_id: &str) -> Result<Vec<u8>, String> {
285        let result = self
286            .send_request(
287                "host/resource.get",
288                serde_json::json!({
289                    "resourceId": resource_id,
290                }),
291            )
292            .await?;
293        let b64 = result.as_str().unwrap_or("");
294        base64::engine::general_purpose::STANDARD
295            .decode(b64)
296            .map_err(|e| format!("base64 decode failed: {}", e))
297    }
298
299    /// 直接写入资源字节数据(无需临时文件),base64 编解码由 SDK 内部处理。
300    pub async fn resource_put(&self, resource_id: &str, data: &[u8]) -> Result<(), String> {
301        let b64 = base64::engine::general_purpose::STANDARD.encode(data);
302        self.send_request(
303            "host/resource.put",
304            serde_json::json!({
305                "resourceId": resource_id,
306                "bytesB64": b64,
307            }),
308        )
309        .await?;
310        Ok(())
311    }
312
313    /// 删除资源文件。
314    pub async fn resource_delete(&self, resource_id: &str) -> Result<(), String> {
315        self.send_request(
316            "host/resource.delete",
317            serde_json::json!({
318                "resourceId": resource_id,
319            }),
320        )
321        .await?;
322        Ok(())
323    }
324
325    /// 列出本插件的所有资源标识符。
326    pub async fn resource_list(&self) -> Result<Vec<String>, String> {
327        let result = self
328            .send_request("host/resource.list", serde_json::json!({}))
329            .await?;
330        serde_json::from_value(result).map_err(|e| format!("parse resource list failed: {}", e))
331    }
332}
333
334#[cfg(test)]
335mod tests {
336    use super::*;
337
338    fn make_proxy() -> (Arc<HostProxy>, mpsc::Receiver<Vec<u8>>) {
339        let (tx, rx) = mpsc::channel::<Vec<u8>>(16);
340        let pending: Arc<DashMap<u64, oneshot::Sender<Result<serde_json::Value, JsonRpcError>>>> =
341            Arc::new(DashMap::new());
342        (Arc::new(HostProxy::new(pending, tx)), rx)
343    }
344
345    /// 回归测试:host/* 请求必须以**干净 JSON** 投递到 outbound 通道,
346    /// 分帧由 write_task 统一完成。
347    ///
348    /// 修复前 send_request/log_no_wait 在投递前先 encode_frame,
349    /// 导致双重分帧(帧内嵌帧):写出的字节为
350    /// `Content-Length: M\r\n\r\nContent-Length: N\r\n\r\n{JSON}`,
351    /// 宿主 read_frame 后 body 以 `Content-Length:` 开头,serde_json 解析失败
352    /// (`expected value at line 1 column 1`),host/* 调用全部超时。
353    #[tokio::test]
354    async fn send_request_posts_clean_json_without_embedded_frame() {
355        let (proxy, mut rx) = make_proxy();
356
357        // log_no_wait 是同步 fire-and-forget:直接检查通道里的字节。
358        proxy.log_no_wait("warn", "test message");
359        let bytes = rx.recv().await.expect("log_no_wait 应投递一条消息");
360        let text = String::from_utf8(bytes.clone()).expect("UTF-8");
361        assert!(
362            !text.contains("Content-Length"),
363            "outbound 通道中的消息不得包含帧头(双重分帧): {:?}",
364            text
365        );
366        // 必须是合法 JSON-RPC 消息
367        let v: serde_json::Value = serde_json::from_slice(&bytes).expect("干净 JSON 可解析");
368        assert_eq!(v["method"], "host/log");
369        assert_eq!(v["params"]["message"], "test message");
370    }
371
372    /// send_request 投递的也必须是干净 JSON(await 路径)。
373    #[tokio::test]
374    async fn send_request_await_path_posts_clean_json() {
375        let (proxy, mut rx) = make_proxy();
376        // send_request 会 await 响应(oneshot 永不完成→挂起),
377        // 因此只在通道上取一条消息验证负载,然后丢弃任务。
378        let task = tokio::spawn(async move {
379            let _ = proxy.model_list().await; // 挂起直到超时(30s),任务随测试结束
380        });
381        let bytes = tokio::time::timeout(Duration::from_secs(2), rx.recv())
382            .await
383            .expect("send_request 应投递一条消息")
384            .expect("通道未关闭");
385        let text = String::from_utf8_lossy(&bytes);
386        assert!(
387            !text.contains("Content-Length"),
388            "outbound 通道中的消息不得包含帧头(双重分帧): {:?}",
389            text
390        );
391        let v: serde_json::Value = serde_json::from_slice(&bytes).expect("干净 JSON 可解析");
392        assert_eq!(v["method"], "host/model.list");
393        // 清理:中止挂起的任务
394        task.abort();
395    }
396}