rpi_server/
lib.rs

1mod server;
2
3use rpi_plugin_sdk::{
4    register_entrypoint, EventTag, FreeStringFn, PluginApiVt, StablePluginEvent, StableToolSchema,
5    StbString, StbStringRef, StepHandle, StepResult, ToolPartialCb,
6};
7use serde_json::{json, Value};
8use std::collections::HashMap;
9use std::sync::atomic::{AtomicBool, Ordering};
10use std::sync::{Arc, Mutex, OnceLock};
11use std::sync::mpsc;
12use std::thread;
13use std::time::{Duration, Instant};
14use server::{generate_token, launch_args, start_rpc_server, RpcServerHandle};
15
16const DEFAULT_TIMEOUT: u64 = 30;
17
18// ── Data ──────────────────────────────────────────────────────────────────────
19
20static SERVERS: OnceLock<Mutex<HashMap<String, Arc<RpcServerHandle>>>> = OnceLock::new();
21
22fn servers() -> &'static Mutex<HashMap<String, Arc<RpcServerHandle>>> {
23    SERVERS.get_or_init(|| Mutex::new(HashMap::new()))
24}
25
26struct Drive {
27    receiver: mpsc::Receiver<Result<String, String>>,
28    cancelled: Arc<AtomicBool>,
29    done: bool,
30}
31
32// ── Param parsing ─────────────────────────────────────────────────────────────
33
34fn string_param(params: &Value, key: &str) -> Result<String, String> {
35    params
36        .get(key)
37        .and_then(Value::as_str)
38        .map(str::to_owned)
39        .filter(|s| !s.trim().is_empty())
40        .ok_or_else(|| format!("{key} is required"))
41}
42
43fn optional_string(params: &Value, key: &str) -> Result<Option<String>, String> {
44    match params.get(key) {
45        None | Some(Value::Null) => Ok(None),
46        Some(Value::String(s)) if !s.trim().is_empty() => Ok(Some(s.clone())),
47        Some(Value::String(_)) => Err(format!("{key} must not be empty")),
48        Some(_) => Err(format!("{key} must be a string")),
49    }
50}
51
52fn timeout_param(params: &Value) -> Result<Duration, String> {
53    match params.get("timeoutSeconds") {
54        None | Some(Value::Null) => Ok(Duration::from_secs(DEFAULT_TIMEOUT)),
55        Some(v) => v
56            .as_u64()
57            .filter(|v| (1..=300).contains(v))
58            .map(Duration::from_secs)
59            .ok_or_else(|| "timeoutSeconds must be an integer from 1 to 300".into()),
60    }
61}
62
63// ── Server lifecycle ──────────────────────────────────────────────────────────
64
65fn do_start(params: &Value) -> Result<String, String> {
66    let bind_addr = optional_string(params, "bind")?.unwrap_or_else(|| "127.0.0.1".into());
67    let port: u16 = match params.get("port") {
68        Some(Value::Number(n)) => n.as_u64()
69            .and_then(|n| u16::try_from(n).ok())
70            .ok_or("port must be an integer 0–65535")?,
71        None => 9800,
72        Some(_) => return Err("port must be an integer".into()),
73    };
74
75    let (executable, args) = launch_args(params)?;
76    let token = optional_string(params, "token")?.unwrap_or_else(generate_token);
77
78    let exec_clone = executable.clone();
79    let args_clone = args.clone();
80    let token_clone = token.clone();
81    let bind_clone = bind_addr.clone();
82
83    let (ready_tx, ready_rx) = mpsc::channel();
84    thread::spawn(move || {
85        let runtime = match tokio::runtime::Builder::new_multi_thread()
86            .enable_all()
87            .worker_threads(2)
88            .build()
89        {
90            Ok(r) => r,
91            Err(e) => {
92                let _ = ready_tx.send(Err(format!("failed to create runtime: {e}")));
93                return;
94            }
95        };
96
97        let result = runtime.block_on(async {
98            start_rpc_server(&bind_clone, port, token_clone, exec_clone, args_clone).await
99        });
100
101        match result {
102            Ok(server) => {
103                if let Ok(mut registry) = servers().lock() {
104                    registry.insert(server.id.clone(), server.clone());
105                }
106                let _ = ready_tx.send(Ok(server));
107            }
108            Err(e) => {
109                let _ = ready_tx.send(Err(e));
110            }
111        }
112
113        runtime.block_on(async { loop { tokio::time::sleep(Duration::from_secs(3600)).await; } });
114    });
115
116    let server = match ready_rx.recv_timeout(Duration::from_secs(10)) {
117        Ok(Ok(s)) => s,
118        Ok(Err(e)) => return Err(e),
119        Err(_) => return Err("timed out waiting for server startup".into()),
120    };
121
122    Ok(json!({
123        "id": server.id,
124        "address": format!("ws://{}", server.address),
125        "token": token,
126        "state": "running",
127        "protocol": "jsonrpc-2.0",
128        "transport": "websocket",
129    }).to_string())
130}
131
132fn do_stop(id: &str) -> Result<String, String> {
133    let server = servers()
134        .lock()
135        .map_err(|_| "server registry poisoned")?
136        .remove(id)
137        .ok_or("RPC server not found")?;
138
139    let (tx, rx) = mpsc::channel();
140    let server_clone = Arc::clone(&server);
141    thread::spawn(move || {
142        let runtime = tokio::runtime::Builder::new_current_thread()
143            .enable_all()
144            .build()
145            .unwrap();
146        let result = runtime.block_on(server_clone.stop());
147        let _ = tx.send(result);
148    });
149
150    match rx.recv_timeout(Duration::from_secs(5)) {
151        Ok(Ok(())) => Ok(json!({"id": id, "state": "stopped"}).to_string()),
152        Ok(Err(e)) => Err(e),
153        Err(_) => Err("timed out waiting for server stop".into()),
154    }
155}
156
157async fn server_status(server: &Arc<RpcServerHandle>) -> Result<Value, String> {
158    let info = server.info().await;
159    Ok(json!({
160        "id": info.id,
161        "address": format!("ws://{}", info.address),
162        "token": info.token,
163        "state": info.state,
164        "protocol": "jsonrpc-2.0",
165        "transport": "websocket",
166        "sessions": info.sessions,
167    }))
168}
169
170async fn do_status(id: &str) -> Result<String, String> {
171    let server = servers()
172        .lock()
173        .map_err(|_| "server registry poisoned")?
174        .get(id)
175        .cloned()
176        .ok_or("RPC server not found")?;
177    Ok(server_status(&server).await?.to_string())
178}
179
180async fn do_list() -> Result<String, String> {
181    let registry = servers().lock().map_err(|_| "server registry poisoned")?;
182    let mut list: Vec<Value> = Vec::new();
183    for server in registry.values() {
184        list.push(server_status(server).await?);
185    }
186    Ok(json!(list).to_string())
187}
188
189async fn server_tool(params: &Value) -> Result<String, String> {
190    match params.get("action").and_then(Value::as_str).unwrap_or("start") {
191        "start" => do_start(params),
192        "stop" => {
193            let id = string_param(params, "id")?;
194            do_stop(&id)
195        }
196        "status" => {
197            let id = string_param(params, "id")?;
198            do_status(&id).await
199        }
200        "list" => do_list().await,
201        other => Err(format!("action must be one of: start, stop, status, list (got {other})")),
202    }
203}
204
205fn stop_all_servers() {
206    let ids: Vec<String> = match servers().lock() {
207        Ok(registry) => registry.keys().cloned().collect(),
208        Err(_) => return,
209    };
210    for id in ids {
211        let _ = do_stop(&id);
212    }
213}
214
215extern "C" fn on_session_shutdown(event: StablePluginEvent, _: *mut std::ffi::c_void) -> i32 {
216    if event.tag == EventTag::SessionShutdown {
217        stop_all_servers();
218    }
219    0
220}
221
222// ── Client ────────────────────────────────────────────────────────────────────
223
224fn client_tool(params: &Value, cancelled: &AtomicBool) -> Result<String, String> {
225    let server_id = string_param(params, "serverId")?;
226    let method = string_param(params, "method")?;
227    let rpc_params = params.get("params").cloned().unwrap_or(json!({}));
228
229    let server = servers()
230        .lock()
231        .map_err(|_| "server registry poisoned")?
232        .get(&server_id)
233        .cloned()
234        .ok_or("RPC server not found")?;
235
236    if !server.is_running() {
237        return Err("RPC server is not running".into());
238    }
239
240    let wait = timeout_param(params)?;
241    let subscribe = params.get("subscribe").and_then(Value::as_bool).unwrap_or(false);
242
243    let (tx, rx) = mpsc::channel();
244    let cancelled_clone = Arc::new(AtomicBool::new(cancelled.load(Ordering::Acquire)));
245    let cancelled_inner = Arc::clone(&cancelled_clone);
246
247    thread::spawn(move || {
248        let runtime = tokio::runtime::Builder::new_current_thread()
249            .enable_all()
250            .build()
251            .unwrap();
252
253        let result = runtime.block_on(async {
254            use jsonrpsee::ws_client::WsClientBuilder;
255            use jsonrpsee::core::client::ClientT;
256
257            let ws_url = format!("ws://{}", server.address);
258            let client = WsClientBuilder::default()
259                .build(&ws_url)
260                .await
261                .map_err(|e| format!("failed to connect: {e}"))?;
262
263            if subscribe {
264                use jsonrpsee::core::client::SubscriptionClientT;
265                use futures::StreamExt;
266
267                // Convert Value to Vec<Value> for ToRpcParams
268                let params_vec = if rpc_params.is_array() {
269                    rpc_params.as_array().unwrap().clone()
270                } else {
271                    vec![rpc_params.clone()]
272                };
273
274                let mut sub = client
275                    .subscribe::<Value, _>(&method, params_vec, &method)
276                    .await
277                    .map_err(|e| format!("subscribe failed: {e}"))?;
278
279                let mut events = Vec::new();
280                let deadline = Instant::now() + wait;
281
282                loop {
283                    if cancelled_inner.load(Ordering::Acquire) { break; }
284                    if Instant::now() >= deadline { break; }
285
286                    match tokio::time::timeout(Duration::from_millis(100), sub.next()).await {
287                        Ok(Some(Ok(event))) => { events.push(event); }
288                        Ok(Some(Err(e))) => { return Err(format!("subscription error: {e}")); }
289                        Ok(None) => break,
290                        Err(_) => continue,
291                    }
292                }
293
294                Ok(json!({
295                    "serverId": server_id,
296                    "method": method,
297                    "events": events,
298                    "eventCount": events.len(),
299                }).to_string())
300            } else {
301                // Convert Value to Vec<Value> for ToRpcParams
302                let params_vec = if rpc_params.is_array() {
303                    rpc_params.as_array().unwrap().clone()
304                } else {
305                    vec![rpc_params]
306                };
307
308                let result: Value = client
309                    .request(&method, params_vec)
310                    .await
311                    .map_err(|e| format!("RPC call failed: {e}"))?;
312
313                Ok(json!({
314                    "serverId": server_id,
315                    "method": method,
316                    "result": result,
317                }).to_string())
318            }
319        });
320
321        let _ = tx.send(result);
322    });
323
324    let deadline = Instant::now() + wait;
325    loop {
326        if cancelled.load(Ordering::Acquire) {
327            return Err("request cancelled".into());
328        }
329        if Instant::now() >= deadline {
330            return Err("request timed out".into());
331        }
332        match rx.recv_timeout(Duration::from_millis(100)) {
333            Ok(result) => return result,
334            Err(mpsc::RecvTimeoutError::Timeout) => continue,
335            Err(mpsc::RecvTimeoutError::Disconnected) => {
336                return Err("worker stopped without result".into());
337            }
338        }
339    }
340}
341
342// ── SDK integration ───────────────────────────────────────────────────────────
343
344extern "C" fn execute_server(
345    tool_call_id: StbStringRef,
346    params: StbString,
347    free: Option<FreeStringFn>,
348) -> StepHandle {
349    let _ = tool_call_id;
350    let text = params.to_string_lossy();
351    params.free_with(free);
352    let cancelled = Arc::new(AtomicBool::new(false));
353    let (tx, rx) = mpsc::channel();
354    thread::spawn(move || {
355        let result = serde_json::from_str::<Value>(&text)
356            .map_err(|e| format!("invalid parameters: {e}"))
357            .and_then(|p| {
358                let runtime = tokio::runtime::Builder::new_current_thread()
359                    .enable_all()
360                    .build()
361                    .unwrap();
362                runtime.block_on(async {
363                    server_tool(&p).await
364                })
365            });
366        let _ = tx.send(result);
367    });
368    Box::into_raw(Box::new(Drive { receiver: rx, cancelled, done: false })) as StepHandle
369}
370
371extern "C" fn execute_client(
372    tool_call_id: StbStringRef,
373    params: StbString,
374    free: Option<FreeStringFn>,
375) -> StepHandle {
376    let _ = tool_call_id;
377    let text = params.to_string_lossy();
378    params.free_with(free);
379    let cancelled = Arc::new(AtomicBool::new(false));
380    let cancelled_clone = Arc::clone(&cancelled);
381    let (tx, rx) = mpsc::channel();
382    thread::spawn(move || {
383        let result = serde_json::from_str::<Value>(&text)
384            .map_err(|e| format!("invalid parameters: {e}"))
385            .and_then(|p| client_tool(&p, &cancelled_clone));
386        let _ = tx.send(result);
387    });
388    Box::into_raw(Box::new(Drive { receiver: rx, cancelled, done: false })) as StepHandle
389}
390
391extern "C" fn poll(
392    handle: StepHandle,
393    _: Option<ToolPartialCb>,
394    _: *mut std::ffi::c_void,
395) -> StepResult {
396    if handle.is_null() {
397        return StepResult::err(StbString::from_string("null handle".into()));
398    }
399    let drive = unsafe { &mut *(handle as *mut Drive) };
400    if drive.done {
401        return StepResult::err(StbString::from_string("handle polled after completion".into()));
402    }
403    if drive.cancelled.load(Ordering::Acquire) {
404        drive.done = true;
405        return StepResult::err(StbString::from_string("request cancelled".into()));
406    }
407    match drive.receiver.try_recv() {
408        Ok(Ok(text)) => {
409            drive.done = true;
410            StepResult::done(StbString::from_string(
411                json!({"content": [{"type": "text", "text": text}]}).to_string(),
412            ))
413        }
414        Ok(Err(e)) => {
415            drive.done = true;
416            StepResult::err(StbString::from_string(e))
417        }
418        Err(mpsc::TryRecvError::Empty) => StepResult::pending(StbString::empty()),
419        Err(mpsc::TryRecvError::Disconnected) => {
420            drive.done = true;
421            StepResult::err(StbString::from_string("worker stopped without result".into()))
422        }
423    }
424}
425
426extern "C" fn cancel(handle: StepHandle) {
427    if !handle.is_null() {
428        unsafe { (&*(handle as *mut Drive)).cancelled.store(true, Ordering::Release); }
429    }
430}
431
432extern "C" fn destroy(handle: StepHandle) {
433    if !handle.is_null() {
434        unsafe {
435            let drive = Box::from_raw(handle as *mut Drive);
436            drive.cancelled.store(true, Ordering::Release);
437        }
438    }
439}
440
441extern "C" fn free_string(s: StbString) {
442    if !s.is_empty() && !s.ptr.is_null() {
443        unsafe {
444            let bytes = std::slice::from_raw_parts(s.ptr as *const u8, s.len);
445            let _ = Box::from_raw(bytes as *const [u8] as *mut [u8]);
446        }
447    }
448}
449
450// ── Schemas ───────────────────────────────────────────────────────────────────
451
452const SERVER_PARAMETERS: &str = r#"{
453  "type":"object",
454  "properties":{
455    "action":{"type":"string","enum":["start","stop","status","list"]},
456    "id":{"type":"string"},
457    "bind":{"type":"string"},
458    "port":{"type":"integer","minimum":0,"maximum":65535},
459    "token":{"type":"string"},
460    "provider":{"type":"string"},
461    "model":{"type":"string"},
462    "baseUrl":{"type":"string"},
463    "systemPrompt":{"type":"string"},
464    "appendSystemPrompt":{"type":"array","items":{"type":"string"}},
465    "thinking":{"type":"string","enum":["off","minimal","low","medium","high","xhigh","max"]},
466    "name":{"type":"string"},
467    "session":{"type":"string"},
468    "sessionId":{"type":"string"},
469    "sessionDir":{"type":"string"},
470    "noSession":{"type":"boolean"},
471    "tools":{"type":"array","items":{"type":"string"}},
472    "excludeTools":{"type":"array","items":{"type":"string"}},
473    "noTools":{"type":"boolean"},
474    "noBuiltinTools":{"type":"boolean"},
475    "noSkills":{"type":"boolean"},
476    "noPromptTemplates":{"type":"boolean"},
477    "noContextFiles":{"type":"boolean"},
478    "noExtensions":{"type":"boolean"},
479    "enablePiPackages":{"type":"boolean"},
480    "extensionsDir":{"type":"array","items":{"type":"string"}},
481    "extension":{"type":"array","items":{"type":"string"}},
482    "skill":{"type":"array","items":{"type":"string"}},
483    "promptTemplate":{"type":"array","items":{"type":"string"}},
484    "offline":{"type":"boolean"},
485    "approve":{"type":"boolean"}
486  }
487}"#;
488
489const CLIENT_PARAMETERS: &str = r#"{
490  "type":"object",
491  "properties":{
492    "serverId":{"type":"string"},
493    "method":{"type":"string"},
494    "params":{"type":"object"},
495    "subscribe":{"type":"boolean","description":"Enable streaming subscription"},
496    "timeoutSeconds":{"type":"integer","minimum":1,"maximum":300}
497  },
498  "required":["serverId","method"]
499}"#;
500
501// ── Plugin entry ──────────────────────────────────────────────────────────────
502
503#[no_mangle]
504pub extern "C" fn rpi_plugin_register_v2(api: *const PluginApiVt, abi: u32) -> i32 {
505    register_entrypoint(api, abi, |api| {
506        let Some(register) = api.register_tool else {
507            return 1;
508        };
509
510        let server_schema = Box::new(StableToolSchema {
511            name: StbString::from_string("rpc_server".into()),
512            description: StbString::from_string(
513                "Start and manage a JSON-RPC 2.0 server with WebSocket transport and streaming subscriptions.".into(),
514            ),
515            parameters: StbString::from_string(SERVER_PARAMETERS.into()),
516        });
517
518        let client_schema = Box::new(StableToolSchema {
519            name: StbString::from_string("rpc_client".into()),
520            description: StbString::from_string(
521                "Send a JSON-RPC request to a running server. Set subscribe=true for streaming responses.".into(),
522            ),
523            parameters: StbString::from_string(CLIENT_PARAMETERS.into()),
524        });
525
526        let first = register(
527            &*server_schema,
528            execute_server,
529            poll,
530            cancel,
531            destroy,
532            free_string,
533        );
534
535        let second = if first == 0 {
536            register(
537                &*client_schema,
538                execute_client,
539                poll,
540                cancel,
541                destroy,
542                free_string,
543            )
544        } else {
545            first
546        };
547
548        let third = if second == 0 {
549            api.register_event_handler
550                .map(|register| {
551                    register(
552                        EventTag::SessionShutdown,
553                        on_session_shutdown,
554                        std::ptr::null_mut(),
555                    )
556                })
557                .unwrap_or(0)
558        } else {
559            second
560        };
561
562        drop(server_schema);
563        drop(client_schema);
564        third
565    })
566}
567
568#[cfg(test)]
569mod tests {
570    use super::*;
571
572    #[test]
573    fn launch_args_basic() {
574        let (_, args) = launch_args(&json!({
575            "provider": "anthropic",
576            "model": "claude-3",
577            "noSession": true,
578        }))
579        .unwrap();
580        assert_eq!(&args[0..2], ["--mode", "rpc"]);
581        assert!(args.windows(2).any(|w| w == ["--provider", "anthropic"]));
582        assert!(args.windows(2).any(|w| w == ["--model", "claude-3"]));
583        assert!(args.contains(&"--no-session".to_string()));
584    }
585
586    #[test]
587    fn schemas_are_valid_json() {
588        assert!(serde_json::from_str::<Value>(SERVER_PARAMETERS).is_ok());
589        assert!(serde_json::from_str::<Value>(CLIENT_PARAMETERS).is_ok());
590    }
591
592    #[test]
593    fn generate_token_is_32_hex_chars() {
594        let t = generate_token();
595        assert_eq!(t.len(), 32);
596        assert!(t.chars().all(|c| c.is_ascii_hexdigit()));
597    }
598}
rpi_server/lib.rs

rpi_server/
lib.rs
