Skip to main content

rpi_server/
lib.rs

1#![forbid(unsafe_code)]
2
3mod server;
4mod transport;
5
6use rpi_plugin_sdk::handles::ToolHandles;
7use rpi_plugin_sdk::{
8    EventTag, FreeStringFn, RuntimeActionFn, RuntimeActionId, StablePluginEvent, StableToolSchema,
9    StbString, StbStringRef, StepHandle, StepResult, ToolPartialCb,
10};
11use serde_json::{json, Value};
12use server::{generate_token, launch_args, run_server, TcpServerHandle};
13use std::collections::HashMap;
14use std::sync::atomic::{AtomicBool, AtomicPtr, Ordering};
15use std::sync::mpsc;
16use std::sync::{Arc, Mutex, OnceLock};
17use std::thread;
18use std::time::{Duration, Instant};
19use transport::{ManagedConnection, TcpTransport};
20
21const DEFAULT_TIMEOUT: u64 = 30;
22
23// ── Data ──────────────────────────────────────────────────────────────────────
24
25static SERVERS: OnceLock<Mutex<HashMap<String, Arc<TcpServerHandle>>>> = OnceLock::new();
26
27fn servers() -> &'static Mutex<HashMap<String, Arc<TcpServerHandle>>> {
28    SERVERS.get_or_init(|| Mutex::new(HashMap::new()))
29}
30
31// The host's `runtime_action` entry point + its opaque `user_data` (the host's
32// `ActionBridge`), copied OUT of the vtable during `register`. The vtable is the
33// host's to own — we must not retain the `api` pointer (that is a
34// use-after-free: the host builds it on the register stack). `GetCliFlag` needs
35// both the entry point and the `user_data`.
36static RUNTIME_ACTION: OnceLock<RuntimeActionFn> = OnceLock::new();
37static RUNTIME_USER_DATA: AtomicPtr<std::ffi::c_void> = AtomicPtr::new(std::ptr::null_mut());
38// The host's `free_string`, for reclaiming host-produced `out` strings.
39static HOST_FREE_STRING: OnceLock<FreeStringFn> = OnceLock::new();
40
41/// Call the host's `GetCliFlag` runtime action to read a parsed CLI flag value.
42fn get_cli_flag(name: &str) -> Option<Value> {
43    let runtime_action = *RUNTIME_ACTION.get()?;
44    let user_data = RUNTIME_USER_DATA.load(Ordering::Acquire);
45    if user_data.is_null() {
46        return None;
47    }
48
49    let args = json!({"name": name});
50    let args_str = args.to_string();
51    let args_ref = StbStringRef::from_str(&args_str);
52    let mut out = StbString::empty();
53
54    let rc = runtime_action(
55        RuntimeActionId::GetCliFlag as u32,
56        args_ref,
57        &mut out,
58        user_data,
59    );
60
61    // The host may have produced `out` even on an error path; reclaim it via the
62    // host's free so we don't leak the host allocation.
63    let result_str = if rc == 0 {
64        Some(out.to_string_lossy())
65    } else {
66        None
67    };
68    if let Some(free) = HOST_FREE_STRING.get() {
69        free(out);
70    }
71
72    let result: Value = serde_json::from_str(&result_str?).ok()?;
73    result.get("value").cloned()
74}
75
76static HANDLES: ToolHandles<Drive> = ToolHandles::new();
77
78struct Drive {
79    receiver: mpsc::Receiver<Result<String, String>>,
80    done: bool,
81}
82
83// ── Param parsing ─────────────────────────────────────────────────────────────
84
85fn string_param(params: &Value, key: &str) -> Result<String, String> {
86    params
87        .get(key)
88        .and_then(Value::as_str)
89        .map(str::to_owned)
90        .filter(|s| !s.trim().is_empty())
91        .ok_or_else(|| format!("{key} is required"))
92}
93
94fn optional_string(params: &Value, key: &str) -> Result<Option<String>, String> {
95    match params.get(key) {
96        None | Some(Value::Null) => Ok(None),
97        Some(Value::String(s)) if !s.trim().is_empty() => Ok(Some(s.clone())),
98        Some(Value::String(_)) => Err(format!("{key} must not be empty")),
99        Some(_) => Err(format!("{key} must be a string")),
100    }
101}
102
103fn timeout_param(params: &Value) -> Result<Duration, String> {
104    match params.get("timeoutSeconds") {
105        None | Some(Value::Null) => Ok(Duration::from_secs(DEFAULT_TIMEOUT)),
106        Some(v) => v
107            .as_u64()
108            .filter(|v| (1..=300).contains(v))
109            .map(Duration::from_secs)
110            .ok_or_else(|| "timeoutSeconds must be an integer from 1 to 300".into()),
111    }
112}
113
114// ── Server lifecycle ──────────────────────────────────────────────────────────
115
116fn do_start(params: &Value) -> Result<String, String> {
117    let bind_addr = optional_string(params, "bind")?.unwrap_or_else(|| "127.0.0.1".into());
118    let port: u16 = match params.get("port") {
119        Some(Value::Number(n)) => n
120            .as_u64()
121            .and_then(|n| u16::try_from(n).ok())
122            .ok_or("port must be an integer 0–65535")?,
123        None => 9800,
124        Some(_) => return Err("port must be an integer".into()),
125    };
126
127    let (executable, args) = launch_args(params)?;
128    let no_token = params
129        .get("noToken")
130        .and_then(Value::as_bool)
131        .unwrap_or(false);
132    let token = if no_token {
133        // Authentication disabled: any client may connect (trusted local use).
134        String::new()
135    } else {
136        params
137            .get("token")
138            .and_then(Value::as_str)
139            .map(str::to_owned)
140            .filter(|t| !t.is_empty())
141            .unwrap_or_else(generate_token)
142    };
143
144    let exec_clone = executable.clone();
145    let args_clone = args.clone();
146    let token_clone = token.clone();
147    let bind_clone = bind_addr.clone();
148
149    let (ready_tx, ready_rx) = mpsc::channel();
150    thread::spawn(move || {
151        let runtime = match tokio::runtime::Builder::new_multi_thread()
152            .enable_all()
153            .worker_threads(2)
154            .build()
155        {
156            Ok(r) => r,
157            Err(e) => {
158                let _ = ready_tx.send(Err(format!("failed to create runtime: {e}")));
159                return;
160            }
161        };
162
163        let result = runtime.block_on(async {
164            let transport = Arc::new(TcpTransport::new());
165            run_server(
166                transport,
167                &bind_clone,
168                port,
169                token_clone,
170                exec_clone,
171                args_clone,
172            )
173            .await
174        });
175
176        match result {
177            Ok(server) => {
178                if let Ok(mut registry) = servers().lock() {
179                    registry.insert(server.id.clone(), server.clone());
180                }
181                let _ = ready_tx.send(Ok(server));
182            }
183            Err(e) => {
184                let _ = ready_tx.send(Err(e));
185            }
186        }
187
188        runtime.block_on(async {
189            loop {
190                tokio::time::sleep(Duration::from_secs(3600)).await;
191            }
192        });
193    });
194
195    let server = match ready_rx.recv_timeout(Duration::from_secs(10)) {
196        Ok(Ok(s)) => s,
197        Ok(Err(e)) => return Err(e),
198        Err(_) => return Err("timed out waiting for server startup".into()),
199    };
200
201    Ok(json!({
202        "id": server.id,
203        "address": server.address,
204        "token": token,
205        "state": "running",
206        "protocol": "jsonl",
207        "transport": "tcp",
208    })
209    .to_string())
210}
211
212fn do_stop(id: &str) -> Result<String, String> {
213    let server = servers()
214        .lock()
215        .map_err(|_| "server registry poisoned")?
216        .remove(id)
217        .ok_or("RPC server not found")?;
218
219    let (tx, rx) = mpsc::channel();
220    let server_clone = Arc::clone(&server);
221    thread::spawn(move || {
222        let runtime = tokio::runtime::Builder::new_current_thread()
223            .enable_all()
224            .build()
225            .unwrap();
226        let result = runtime.block_on(server_clone.stop());
227        let _ = tx.send(result);
228    });
229
230    match rx.recv_timeout(Duration::from_secs(5)) {
231        Ok(Ok(())) => Ok(json!({"id": id, "state": "stopped"}).to_string()),
232        Ok(Err(e)) => Err(e),
233        Err(_) => Err("timed out waiting for server stop".into()),
234    }
235}
236
237async fn server_status(server: &Arc<TcpServerHandle>) -> Result<Value, String> {
238    let info = server.info().await;
239    Ok(json!({
240        "id": info.id,
241        "address": info.address,
242        "token": info.token,
243        "state": info.state,
244        "protocol": "jsonl",
245        "transport": "tcp",
246        "sessions": info.sessions,
247    }))
248}
249
250async fn do_status(id: &str) -> Result<String, String> {
251    let server = servers()
252        .lock()
253        .map_err(|_| "server registry poisoned")?
254        .get(id)
255        .cloned()
256        .ok_or("RPC server not found")?;
257    Ok(server_status(&server).await?.to_string())
258}
259
260async fn do_list() -> Result<String, String> {
261    let registry = servers().lock().map_err(|_| "server registry poisoned")?;
262    let mut list: Vec<Value> = Vec::new();
263    for server in registry.values() {
264        list.push(server_status(server).await?);
265    }
266    Ok(json!(list).to_string())
267}
268
269async fn server_tool(params: &Value) -> Result<String, String> {
270    match params
271        .get("action")
272        .and_then(Value::as_str)
273        .unwrap_or("start")
274    {
275        "start" => do_start(params),
276        "stop" => {
277            let id = string_param(params, "id")?;
278            do_stop(&id)
279        }
280        "status" => {
281            let id = string_param(params, "id")?;
282            do_status(&id).await
283        }
284        "list" => do_list().await,
285        other => Err(format!(
286            "action must be one of: start, stop, status, list (got {other})"
287        )),
288    }
289}
290
291fn stop_all_servers() {
292    let ids: Vec<String> = match servers().lock() {
293        Ok(registry) => registry.keys().cloned().collect(),
294        Err(_) => return,
295    };
296    for id in ids {
297        let _ = do_stop(&id);
298    }
299}
300
301extern "C" fn on_session_shutdown(event: StablePluginEvent, _: *mut std::ffi::c_void) -> i32 {
302    if event.tag == EventTag::SessionShutdown {
303        stop_all_servers();
304    }
305    0
306}
307
308/// Whether a `GetCliFlag` result means the switch is ON. `GetCliFlag` yields
309/// `true` for a bare `--flag`, a string when `--flag=v`, and `null` when unset.
310fn flag_is_on(value: Option<&Value>) -> bool {
311    match value {
312        Some(Value::Bool(b)) => *b,
313        Some(Value::String(s)) => !s.is_empty() && !s.eq_ignore_ascii_case("false"),
314        _ => false,
315    }
316}
317
318/// On SessionStart, check if --server CLI flag was passed and auto-start the TCP server.
319extern "C" fn on_session_start(event: StablePluginEvent, _: *mut std::ffi::c_void) -> i32 {
320    if event.tag != EventTag::SessionStart {
321        return 0;
322    }
323
324    // Check if --server flag was passed. `GetCliFlag` returns `{"value": null}`
325    // for an *unset* flag, so a bare `is_none()` would wrongly treat that as
326    // present — require an explicit truthy value.
327    if !flag_is_on(get_cli_flag("server").as_ref()) {
328        return 0;
329    }
330
331    // Read optional --port flag (default 9800)
332    let port: u16 = match get_cli_flag("port") {
333        Some(Value::Number(n)) => n
334            .as_u64()
335            .and_then(|n| u16::try_from(n).ok())
336            .unwrap_or(9800),
337        Some(Value::String(s)) => s.parse().unwrap_or(9800),
338        _ => 9800,
339    };
340
341    // Read optional --bind flag (default 127.0.0.1)
342    let bind = match get_cli_flag("bind") {
343        Some(Value::String(s)) if !s.is_empty() => s,
344        _ => "127.0.0.1".to_string(),
345    };
346
347    eprintln!("[rpi-server] --server detected, starting TCP server on {bind}:{port}");
348
349    let token = match get_cli_flag("token") {
350        Some(Value::String(s)) if !s.is_empty() => Some(s),
351        _ => None,
352    };
353    let no_token = flag_is_on(get_cli_flag("no-token").as_ref());
354
355    let params = json!({
356        "bind": bind,
357        "port": port,
358        "token": token,
359        "noToken": no_token,
360    });
361
362    match do_start(&params) {
363        Ok(info) => eprintln!("[rpi-server] TCP server started: {info}"),
364        Err(e) => eprintln!("[rpi-server] failed to start: {e}"),
365    }
366
367    0
368}
369
370// ── Client ────────────────────────────────────────────────────────────────────
371
372fn client_tool(params: &Value, cancelled: &AtomicBool) -> Result<String, String> {
373    let server_id = string_param(params, "serverId")?;
374    let method = string_param(params, "method")?;
375    let rpc_params = params.get("params").cloned().unwrap_or(json!({}));
376
377    let server = servers()
378        .lock()
379        .map_err(|_| "server registry poisoned")?
380        .get(&server_id)
381        .cloned()
382        .ok_or("RPC server not found")?;
383
384    if !server.is_running() {
385        return Err("RPC server is not running".into());
386    }
387
388    let wait = timeout_param(params)?;
389    let subscribe = params
390        .get("subscribe")
391        .and_then(Value::as_bool)
392        .unwrap_or(false);
393
394    let (tx, rx) = mpsc::channel();
395    let cancelled_clone = Arc::new(AtomicBool::new(cancelled.load(Ordering::Acquire)));
396    let cancelled_inner = Arc::clone(&cancelled_clone);
397
398    thread::spawn(move || {
399        let runtime = tokio::runtime::Builder::new_current_thread()
400            .enable_all()
401            .build()
402            .unwrap();
403
404        let result = runtime.block_on(async {
405            let transport = Arc::new(TcpTransport::new());
406            let managed = ManagedConnection::new(transport);
407
408            // Connect with retry
409            managed
410                .connect_with_retry(&server.address)
411                .await
412                .map_err(|e| format!("failed to connect: {e}"))?;
413
414            // Authenticate when the server requires a token.
415            if !server.token.is_empty() {
416                let auth = json!({
417                    "jsonrpc": "2.0",
418                    "id": 0,
419                    "method": "authenticate",
420                    "params": {"token": server.token},
421                });
422                managed
423                    .send(auth)
424                    .await
425                    .map_err(|e| format!("auth send failed: {e}"))?;
426                match tokio::time::timeout(Duration::from_secs(5), managed.recv()).await {
427                    Ok(Ok(Some(response))) => {
428                        if let Some(error) = response.get("error") {
429                            return Err(format!("authentication failed: {error}"));
430                        }
431                    }
432                    Ok(Ok(None)) => return Err("authentication failed: connection closed".into()),
433                    Ok(Err(e)) => return Err(format!("authentication failed: {e}")),
434                    Err(_) => return Err("authentication failed: timed out".into()),
435                }
436            }
437
438            // Send request
439            let request = json!({
440                "jsonrpc": "2.0",
441                "method": method,
442                "params": rpc_params,
443                "id": 1,
444            });
445            managed
446                .send(request)
447                .await
448                .map_err(|e| format!("send failed: {e}"))?;
449
450            if subscribe {
451                // Collect streaming events
452                let mut events = Vec::new();
453                let deadline = Instant::now() + wait;
454
455                loop {
456                    if cancelled_inner.load(Ordering::Acquire) {
457                        break;
458                    }
459                    if Instant::now() >= deadline {
460                        break;
461                    }
462
463                    match tokio::time::timeout(Duration::from_millis(100), managed.recv()).await {
464                        Ok(Ok(Some(event))) => {
465                            // Check for session_end before pushing
466                            let is_end =
467                                event.get("type").and_then(Value::as_str) == Some("session_end");
468                            events.push(event);
469                            if is_end {
470                                break;
471                            }
472                        }
473                        Ok(Ok(None)) => break,
474                        Ok(Err(e)) => return Err(format!("recv error: {e}")),
475                        Err(_) => continue, // timeout, loop again
476                    }
477                }
478
479                Ok(json!({
480                    "serverId": server_id,
481                    "method": method,
482                    "events": events,
483                    "eventCount": events.len(),
484                })
485                .to_string())
486            } else {
487                // Wait for response
488                let deadline = Instant::now() + wait;
489                loop {
490                    if cancelled_inner.load(Ordering::Acquire) {
491                        return Err("request cancelled".into());
492                    }
493                    if Instant::now() >= deadline {
494                        return Err("request timed out".into());
495                    }
496
497                    match tokio::time::timeout(Duration::from_millis(100), managed.recv()).await {
498                        Ok(Ok(Some(response))) => {
499                            return Ok(json!({
500                                "serverId": server_id,
501                                "method": method,
502                                "result": response,
503                            })
504                            .to_string());
505                        }
506                        Ok(Ok(None)) => return Err("connection closed".into()),
507                        Ok(Err(e)) => return Err(format!("recv error: {e}")),
508                        Err(_) => continue, // timeout, loop again
509                    }
510                }
511            }
512        });
513
514        let _ = tx.send(result);
515    });
516
517    let deadline = Instant::now() + wait;
518    loop {
519        if cancelled.load(Ordering::Acquire) {
520            return Err("request cancelled".into());
521        }
522        if Instant::now() >= deadline {
523            return Err("request timed out".into());
524        }
525        match rx.recv_timeout(Duration::from_millis(100)) {
526            Ok(result) => return result,
527            Err(mpsc::RecvTimeoutError::Timeout) => continue,
528            Err(mpsc::RecvTimeoutError::Disconnected) => {
529                return Err("worker stopped without result".into());
530            }
531        }
532    }
533}
534
535// ── SDK integration ───────────────────────────────────────────────────────────
536
537extern "C" fn execute_server(
538    tool_call_id: StbStringRef,
539    params: StbString,
540    free: Option<FreeStringFn>,
541) -> StepHandle {
542    let _ = tool_call_id;
543    let text = params.to_string_lossy();
544    params.free_with(free);
545    let cancelled = Arc::new(AtomicBool::new(false));
546    let (tx, rx) = mpsc::channel();
547    thread::spawn(move || {
548        let result = serde_json::from_str::<Value>(&text)
549            .map_err(|e| format!("invalid parameters: {e}"))
550            .and_then(|p| {
551                let runtime = tokio::runtime::Builder::new_current_thread()
552                    .enable_all()
553                    .build()
554                    .unwrap();
555                runtime.block_on(async { server_tool(&p).await })
556            });
557        let _ = tx.send(result);
558    });
559    HANDLES.insert_with_cancel(
560        Drive {
561            receiver: rx,
562            done: false,
563        },
564        cancelled,
565    )
566}
567
568extern "C" fn execute_client(
569    tool_call_id: StbStringRef,
570    params: StbString,
571    free: Option<FreeStringFn>,
572) -> StepHandle {
573    let _ = tool_call_id;
574    let text = params.to_string_lossy();
575    params.free_with(free);
576    let cancelled = Arc::new(AtomicBool::new(false));
577    let cancelled_clone = Arc::clone(&cancelled);
578    let (tx, rx) = mpsc::channel();
579    thread::spawn(move || {
580        let result = serde_json::from_str::<Value>(&text)
581            .map_err(|e| format!("invalid parameters: {e}"))
582            .and_then(|p| client_tool(&p, &cancelled_clone));
583        let _ = tx.send(result);
584    });
585    HANDLES.insert_with_cancel(
586        Drive {
587            receiver: rx,
588            done: false,
589        },
590        cancelled,
591    )
592}
593
594extern "C" fn poll(
595    handle: StepHandle,
596    _: Option<ToolPartialCb>,
597    _: *mut std::ffi::c_void,
598) -> StepResult {
599    if handle.is_null() {
600        return StepResult::err(StbString::from_string("null handle".into()));
601    }
602    let Some(entry) = HANDLES.get(handle) else {
603        return StepResult::err(StbString::from_string(
604            "invalid or destroyed tool handle".into(),
605        ));
606    };
607    let mut drive = entry.lock();
608    if drive.done {
609        return StepResult::err(StbString::from_string(
610            "handle polled after completion".into(),
611        ));
612    }
613    if entry.cancelled.load(Ordering::Acquire) {
614        drive.done = true;
615        return StepResult::err(StbString::from_string("request cancelled".into()));
616    }
617    match drive.receiver.try_recv() {
618        Ok(Ok(text)) => {
619            drive.done = true;
620            StepResult::done(StbString::from_string(
621                json!({"content": [{"type": "text", "text": text}]}).to_string(),
622            ))
623        }
624        Ok(Err(e)) => {
625            drive.done = true;
626            StepResult::err(StbString::from_string(e))
627        }
628        Err(mpsc::TryRecvError::Empty) => StepResult::pending(StbString::empty()),
629        Err(mpsc::TryRecvError::Disconnected) => {
630            drive.done = true;
631            StepResult::err(StbString::from_string(
632                "worker stopped without result".into(),
633            ))
634        }
635    }
636}
637
638extern "C" fn cancel(handle: StepHandle) {
639    HANDLES.cancel(handle);
640}
641
642extern "C" fn destroy(handle: StepHandle) {
643    drop(HANDLES.remove(handle));
644}
645
646use rpi_plugin_sdk::callbacks::free_string;
647
648// ── Schemas ───────────────────────────────────────────────────────────────────
649
650const SERVER_PARAMETERS: &str = r#"{
651  "type":"object",
652  "properties":{
653    "action":{"type":"string","enum":["start","stop","status","list"]},
654    "id":{"type":"string"},
655    "bind":{"type":"string"},
656    "port":{"type":"integer","minimum":0,"maximum":65535},
657    "token":{"type":"string"},
658    "provider":{"type":"string"},
659    "model":{"type":"string"},
660    "baseUrl":{"type":"string"},
661    "systemPrompt":{"type":"string"},
662    "appendSystemPrompt":{"type":"array","items":{"type":"string"}},
663    "thinking":{"type":"string","enum":["off","minimal","low","medium","high","xhigh","max"]},
664    "rpcMessageUpdates":{"type":"string","enum":["snapshot","compact-deltas"],"description":"Native RPC delta payload mode; omitted uses the executable default (compatible with older hosts)."},
665    "name":{"type":"string"},
666    "session":{"type":"string"},
667    "sessionId":{"type":"string"},
668    "sessionDir":{"type":"string"},
669    "noSession":{"type":"boolean"},
670    "tools":{"type":"array","items":{"type":"string"}},
671    "excludeTools":{"type":"array","items":{"type":"string"}},
672    "noTools":{"type":"boolean"},
673    "noBuiltinTools":{"type":"boolean"},
674    "noSkills":{"type":"boolean"},
675    "noPromptTemplates":{"type":"boolean"},
676    "noContextFiles":{"type":"boolean"},
677    "noExtensions":{"type":"boolean"},
678    "enablePiPackages":{"type":"boolean"},
679    "extensionsDir":{"type":"array","items":{"type":"string"}},
680    "extension":{"type":"array","items":{"type":"string"}},
681    "skill":{"type":"array","items":{"type":"string"}},
682    "promptTemplate":{"type":"array","items":{"type":"string"}},
683    "offline":{"type":"boolean"},
684    "approve":{"type":"boolean"}
685  }
686}"#;
687
688const CLIENT_PARAMETERS: &str = r#"{
689  "type":"object",
690  "properties":{
691    "serverId":{"type":"string"},
692    "method":{"type":"string"},
693    "params":{"type":"object"},
694    "subscribe":{"type":"boolean","description":"Enable streaming subscription"},
695    "timeoutSeconds":{"type":"integer","minimum":1,"maximum":300}
696  },
697  "required":["serverId","method"]
698}"#;
699
700// ── Plugin entry ──────────────────────────────────────────────────────────────
701rpi_plugin_sdk::export_plugin!(|api| {
702    // Copy the host function pointers/handles OUT of the vtable during
703    // register — the `api` pointer is only valid for this call.
704    let _ = RUNTIME_ACTION.set(api.runtime_action);
705    RUNTIME_USER_DATA.store(api.user_data, Ordering::Release);
706    let _ = HOST_FREE_STRING.set(api.free_string);
707    let Some(register) = api.register_tool else {
708        return 1;
709    };
710
711    let server_schema = Box::new(StableToolSchema {
712        name: StbString::from_string("rpc_server".into()),
713        description: StbString::from_string(
714            "Start and manage a TCP JSONL server with streaming subscriptions.".into(),
715        ),
716        parameters: StbString::from_string(SERVER_PARAMETERS.into()),
717    });
718
719    let client_schema = Box::new(StableToolSchema {
720    name: StbString::from_string("rpc_client".into()),
721    description: StbString::from_string(
722        "Send a JSON-RPC request to a running TCP server. Set subscribe=true for streaming responses.".into(),
723    ),
724    parameters: StbString::from_string(CLIENT_PARAMETERS.into()),
725});
726
727    let first = register(
728        &*server_schema,
729        execute_server,
730        poll,
731        cancel,
732        destroy,
733        free_string,
734    );
735
736    let second = if first == 0 {
737        register(
738            &*client_schema,
739            execute_client,
740            poll,
741            cancel,
742            destroy,
743            free_string,
744        )
745    } else {
746        first
747    };
748
749    let third = if second == 0 {
750        api.register_event_handler
751            .map(|register| {
752                register(
753                    EventTag::SessionShutdown,
754                    on_session_shutdown,
755                    std::ptr::null_mut(),
756                )
757            })
758            .unwrap_or(0)
759    } else {
760        second
761    };
762
763    // Register SessionStart handler to auto-start server when --server flag is passed
764    let fourth = if third == 0 {
765        api.register_event_handler
766            .map(|register| {
767                register(
768                    EventTag::SessionStart,
769                    on_session_start,
770                    std::ptr::null_mut(),
771                )
772            })
773            .unwrap_or(0)
774    } else {
775        third
776    };
777
778    drop(server_schema);
779    drop(client_schema);
780    fourth
781});
782
783#[cfg(test)]
784mod tests {
785    use super::*;
786
787    #[test]
788    fn launch_args_basic() {
789        let (_, args) = launch_args(&json!({
790            "provider": "anthropic",
791            "model": "claude-3",
792            "noSession": true,
793        }))
794        .unwrap();
795        assert_eq!(&args[0..2], ["--mode", "rpc"]);
796        assert!(args.windows(2).any(|w| w == ["--provider", "anthropic"]));
797        assert!(args.windows(2).any(|w| w == ["--model", "claude-3"]));
798        assert!(args.contains(&"--no-session".to_string()));
799    }
800
801    #[test]
802    fn rpc_message_updates_is_explicit_and_strict() {
803        let (_, args) = launch_args(&json!({})).unwrap();
804        assert!(!args
805            .iter()
806            .any(|arg| arg.starts_with("--rpc-message-updates")));
807        for mode in ["snapshot", "compact-deltas"] {
808            let (_, args) = launch_args(&json!({"rpcMessageUpdates":mode})).unwrap();
809            assert_eq!(
810                args,
811                vec!["--mode", "rpc", &format!("--rpc-message-updates={mode}")]
812            );
813        }
814        for invalid in [
815            json!(""),
816            json!("compact"),
817            json!("SNAPSHOT"),
818            json!(null),
819            json!(true),
820            json!(1),
821            json!([]),
822            json!({}),
823        ] {
824            assert_eq!(
825                launch_args(&json!({"rpcMessageUpdates":invalid})).unwrap_err(),
826                "rpcMessageUpdates must be snapshot or compact-deltas"
827            );
828        }
829        let schema: Value = serde_json::from_str(SERVER_PARAMETERS).unwrap();
830        assert_eq!(
831            schema["properties"]["rpcMessageUpdates"]["enum"],
832            json!(["snapshot", "compact-deltas"])
833        );
834        assert!(schema["properties"]["rpcMessageUpdates"]
835            .get("default")
836            .is_none());
837    }
838
839    #[test]
840    fn schemas_are_valid_json() {
841        assert!(serde_json::from_str::<Value>(SERVER_PARAMETERS).is_ok());
842        assert!(serde_json::from_str::<Value>(CLIENT_PARAMETERS).is_ok());
843    }
844
845    #[test]
846    fn generate_token_is_32_hex_chars() {
847        let t = generate_token();
848        assert_eq!(t.len(), 32);
849        assert!(t.chars().all(|c| c.is_ascii_hexdigit()));
850    }
851}