rpi_server/
lib.rs

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

rpi_server/
lib.rs
