rpi_server/
lib.rs

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

rpi_server/
lib.rs
