sausage/cli/src/routing.rs

//! JSON-RPC routing and responses between callers and processes.
use crate::{
    Host, NativeFuture, RpcResult, error,
    process::{Process, ProcessInfo},
};
use serde_json::{Value, json};
use tokio::sync::{mpsc, oneshot, watch};
use uuid::Uuid;

impl Host {
    /// Queue an incoming response or notification for its destination.
    pub(crate) fn send(&self, pid: &str, value: Value) {
        if pid == "console" {
            let _ = self.0.notifications.send(value);
        } else if let Some(tx) = self
            .0
            .state
            .lock()
            .unwrap()
            .processes
            .get(pid)
            .and_then(|p| p.input.as_ref())
        {
            let _ = tx.send(value);
        }
    }
    /// Remap progress to the caller's original request ID.
    pub(crate) fn notify(&self, pid: &str, name: &str, params: Value) {
        let target = self
            .0
            .state
            .lock()
            .unwrap()
            .processes
            .get(pid)
            .and_then(|p| p.answer.as_ref().map(|_| (p.caller.clone(), p.id.clone())));
        if let Some((caller, id)) = target {
            let notification = json!({
                "jsonrpc": "2.0",
                "method": "core.notify",
                "params": { "call": id, "name": name, "params": params }
            });
            self.send(&caller, notification);
        }
    }
    /// Validate and route a request, then start its native or WASI process.
    pub(crate) fn route(
        &self,
        caller: String,
        id: Option<Value>,
        method: String,
        params: Value,
        detached: bool,
    ) -> NativeFuture {
        self.route_from(caller, id, method, params, detached, None)
    }
    /// Like `route`, but a WASI process runs the component at `wasm` instead
    /// of the current `plugin.wasm` of its plugin. Activation uses this to
    /// call `rpc.discover` on a candidate file.
    pub(crate) fn route_from(
        &self,
        caller: String,
        id: Option<Value>,
        method: String,
        params: Value,
        detached: bool,
        wasm: Option<std::path::PathBuf>,
    ) -> NativeFuture {
        let host = self.clone();
        Box::pin(async move {
            if !params.is_object() {
                return Err(error(-32602, "params must be an object"));
            }
            if caller != "console" && caller != "core" {
                let state = host.0.state.lock().unwrap();
                let Some(process) = state.processes.get(&caller) else {
                    return Err(error(-32600, "process ended"));
                };
                if process.info.state != "running" && method != "core.notify" {
                    return Err(error(-32600, "process is stopping"));
                }
            }
            let (plugin, method) = method
                .split_once('.')
                .ok_or_else(|| error(-32601, "method not found"))?;
            if plugin == "core" {
                return host.core(&caller, id, method, params).await;
            }
            if plugin.is_empty()
                || !plugin
                    .bytes()
                    .all(|b| b.is_ascii_lowercase() || b.is_ascii_digit() || b == b'-' || b == b'_')
            {
                return Err(error(-32601, "invalid plugin name"));
            }
            let admission = host.0.admission.read().await;
            let native = host.0.state.lock().unwrap().native.get(plugin).cloned();
            let (component, version) = if native.is_some() {
                (None, "native".into())
            } else {
                let (component, hash) = host
                    .0
                    .components
                    .load(
                        wasm.unwrap_or_else(|| host.0.world.plugin_dir(plugin).join("plugin.wasm")),
                        host.0.runtime.clone(),
                    )
                    .await?;
                (Some(component), hash)
            };
            let pid = format!("{plugin}-{}", Uuid::now_v7());
            let (tx, rx) = oneshot::channel();
            let (input, incoming) = mpsc::unbounded_channel::<Value>();
            let info = ProcessInfo {
                plugin: plugin.into(),
                version,
                method: method.into(),
                params: params.clone(),
                parent: if detached { None } else { Some(caller.clone()) },
                detached,
                start_time: std::time::SystemTime::now()
                    .duration_since(std::time::UNIX_EPOCH)
                    .unwrap()
                    .as_millis(),
                state: "running".into(),
            };
            {
                let mut state = host.0.state.lock().unwrap();
                if caller != "console"
                    && caller != "core"
                    && !state
                        .processes
                        .get(&caller)
                        .is_some_and(|p| p.info.state == "running")
                {
                    return Err(error(-32600, "caller stopped or ended"));
                }
                host.0
                    .journal
                    .write(&pid, plugin, "core", &format!("start {method} {params}"));
                state.processes.insert(
                    pid.clone(),
                    Process {
                        info: info.clone(),
                        memory: sausage_vfs::Memory::default(),
                        files: sausage_vfs::Process::default(),
                        store_ended: watch::channel(true).0,
                        input: Some(input),
                        abort: None,
                        caller,
                        id: id.clone(),
                        answer: Some(tx),
                        ended: watch::channel(false).0,
                    },
                );
            }
            let h = host.clone();
            let p = pid.clone();
            let has_id = id.is_some();
            let (start, ready) = oneshot::channel();
            let task = tokio::spawn(async move {
                let _ = ready.await;
                h.run_process(p, info, component, native, incoming, has_id)
                    .await
            });
            if let Some(p) = host.0.state.lock().unwrap().processes.get_mut(&pid) {
                p.abort = Some(task.abort_handle());
                if p.info.state == "stopping" {
                    task.abort();
                }
            } else {
                task.abort();
            }
            drop(admission);
            let cleanup = host.clone();
            let finished_pid = pid.clone();
            tokio::spawn(async move {
                let status = task.await.unwrap_or_else(|error| {
                    if error.is_cancelled() {
                        "killed".into()
                    } else {
                        format!("host task failed: {error}")
                    }
                });
                cleanup.finish(&finished_pid, &status).await;
            });
            let _ = start.send(());
            if id.is_none() {
                return Ok(Value::Null);
            }
            rx.await
                .unwrap_or_else(|_| Err(error(-32002, "process ended")))
        })
    }
    /// Complete the initial call once and discard late responses.
    pub(crate) fn answer(&self, pid: &str, result: RpcResult) {
        if let Some(tx) = self
            .0
            .state
            .lock()
            .unwrap()
            .processes
            .get_mut(pid)
            .and_then(|p| p.answer.take())
        {
            let _ = tx.send(result);
        }
    }
    /// Classify a stdout line and route requests or the initial response.
    pub(crate) fn output(&self, pid: &str, line: String) {
        if line.trim().is_empty() {
            return;
        }
        let Ok(v) = serde_json::from_str::<Value>(&line) else {
            self.discard(pid);
            return;
        };
        if !v.is_object() || v["jsonrpc"] != "2.0" {
            self.discard(pid);
            return;
        }
        if let Some(method) = v["method"].as_str() {
            let h = self.clone();
            let caller = pid.to_string();
            let method = method.to_string();
            tokio::spawn(async move {
                let id = v.get("id").cloned();
                let result = h
                    .route(
                        caller.clone(),
                        id.clone(),
                        method,
                        v.get("params").cloned().unwrap_or(json!({})),
                        false,
                    )
                    .await;
                if let Some(id) = id {
                    let value = match result {
                        Ok(result) => json!({
                            "jsonrpc": "2.0",
                            "id": id,
                            "result": result
                        }),
                        Err(error) => json!({
                            "jsonrpc": "2.0",
                            "id": id,
                            "error": error
                        }),
                    };
                    h.send(&caller, value);
                }
            });
        } else if v["id"] == pid {
            self.answer(
                pid,
                if let Some(e) = v.get("error") {
                    Err(e.clone())
                } else {
                    Ok(v["result"].clone())
                },
            );
        }
    }
    pub(crate) fn discard(&self, pid: &str) {
        // The journal deduplicates this warning for each process.
        if let Some(info) = self.processes().get(pid) {
            self.0.journal.warning(pid, &info.plugin);
        }
    }
}