sausage/cli/src/lib.rs

//! Process ownership, JSON-RPC routing, native plugins, and host lifetime.
pub mod cargo;
mod journal;
pub mod runtime;
pub mod world;
use runtime::{InstanceConfig, Runtime};
use serde_json::{Value, json};
use sha2::{Digest, Sha256};
use std::{
    collections::{HashMap, VecDeque},
    future::Future,
    path::Path,
    pin::Pin,
    sync::{Arc, Mutex},
};
use tokio::sync::{mpsc, oneshot, watch};
use uuid::Uuid;
use world::World;
pub type RpcResult = Result<Value, Value>;
pub type NativeFuture = Pin<Box<dyn Future<Output = RpcResult> + Send>>;
pub type NativePlugin = Arc<dyn Fn(NativeContext, String, Value) -> NativeFuture + Send + Sync>;
#[derive(Clone)]
pub struct NativeContext {
    pub host: Host,
    pub pid: String,
}
impl NativeContext {
    pub async fn call(&self, method: &str, params: Value) -> RpcResult {
        let exit = requests_exit(method, &params);
        let result = self
            .host
            .route(
                self.pid.clone(),
                Some(json!(Uuid::now_v7().to_string())),
                method.into(),
                params,
                false,
            )
            .await;
        if exit && result.is_ok() {
            self.host.0.exit.send_replace(true);
        }
        result
    }
    pub fn notify(&self, name: &str, params: Value) {
        self.host.notify(&self.pid, name, params);
    }
}
#[derive(Clone, serde::Serialize)]
pub struct ProcessInfo {
    pub plugin: String,
    pub version: String,
    pub method: String,
    pub params: Value,
    pub parent: Option<String>,
    pub detached: bool,
    pub start_time: u128,
    pub state: String,
}
struct Process {
    info: ProcessInfo,
    input: Option<mpsc::UnboundedSender<Value>>,
    abort: Option<tokio::task::AbortHandle>,
    caller: String,
    id: Option<Value>,
    answer: Option<oneshot::Sender<RpcResult>>,
    ended: watch::Sender<bool>,
}
struct Lock {
    owner: String,
    queue: VecDeque<(String, oneshot::Sender<RpcResult>)>,
}
struct State {
    processes: HashMap<String, Process>,
    locks: HashMap<String, Lock>,
    components: HashMap<String, wasmtime::component::Component>,
    native: HashMap<String, NativePlugin>,
}
struct Inner {
    world: World,
    runtime: Arc<Runtime>,
    state: Mutex<State>,
    journal: journal::Journal,
    notifications: tokio::sync::broadcast::Sender<Value>,
    exit: watch::Sender<bool>,
    compilation: tokio::sync::Mutex<()>,
}
#[derive(Clone)]
pub struct Host(Arc<Inner>);
struct ProcessGuard {
    host: Host,
    pid: String,
}
impl Drop for ProcessGuard {
    fn drop(&mut self) {
        self.host.finish(&self.pid, "process task ended");
    }
}
fn error(code: i64, text: &str) -> Value {
    json!({"code":code,"message":text})
}
impl Host {
    pub async fn start(path: &Path) -> anyhow::Result<Self> {
        let base = std::env::var_os("XDG_CACHE_HOME")
            .map(std::path::PathBuf::from)
            .unwrap_or_else(|| {
                std::path::PathBuf::from(std::env::var_os("HOME").unwrap_or_default())
                    .join(".cache")
            });
        Self::start_with_cache(path, &base.join("sausage")).await
    }
    pub async fn start_with_cache(path: &Path, cache: &Path) -> anyhow::Result<Self> {
        let world = World::open(path)?;
        let journal = journal::Journal::open(&world.state_dir("core").join("journal.sqlite"))?;
        let host = Self(Arc::new(Inner {
            world,
            runtime: Runtime::new(cache)?,
            journal,
            state: Mutex::new(State {
                processes: HashMap::new(),
                locks: HashMap::new(),
                components: HashMap::new(),
                native: HashMap::new(),
            }),
            notifications: tokio::sync::broadcast::channel(1024).0,
            exit: watch::channel(false).0,
            compilation: tokio::sync::Mutex::new(()),
        }));
        let boot: Vec<Value> =
            serde_json::from_slice(&std::fs::read(host.0.world.root.join("config/boot.json"))?)?;
        for service in boot {
            let result = host
                .route(
                    "core".into(),
                    Some(json!(Uuid::now_v7().to_string())),
                    service["method"].as_str().unwrap_or("").into(),
                    service["params"].clone(),
                    true,
                )
                .await;
            if let Err(e) = result {
                host.0.journal.write("core", "core", "core", &e.to_string());
            }
        }
        Ok(host)
    }
    pub fn register_native(&self, name: &str, plugin: NativePlugin) {
        self.0
            .state
            .lock()
            .unwrap()
            .native
            .insert(name.into(), plugin);
    }
    pub fn subscribe(&self) -> tokio::sync::broadcast::Receiver<Value> {
        self.0.notifications.subscribe()
    }
    pub fn processes(&self) -> HashMap<String, ProcessInfo> {
        self.0
            .state
            .lock()
            .unwrap()
            .processes
            .iter()
            .map(|(p, s)| (p.clone(), s.info.clone()))
            .collect()
    }
    pub async fn call(&self, method: &str, params: Value) -> RpcResult {
        let exit = requests_exit(method, &params);
        let result = self
            .route(
                "console".into(),
                Some(json!(Uuid::now_v7().to_string())),
                method.into(),
                params,
                false,
            )
            .await;
        if exit && result.is_ok() {
            self.0.exit.send_replace(true);
        }
        result
    }
    pub async fn exit_requested(&self) {
        let mut exit = self.0.exit.subscribe();
        let _ = exit.wait_for(|requested| *requested).await;
    }
    pub fn restart_requested(&self) -> bool {
        *self.0.exit.borrow()
    }
    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);
        }
    }
    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 {
            self.send(&caller,json!({"jsonrpc":"2.0","method":"core.notify","params":{"call":id,"name":name,"params":params}}));
        }
    }
    fn route(
        &self,
        caller: String,
        id: Option<Value>,
        method: String,
        params: Value,
        detached: bool,
    ) -> 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 native = host.0.state.lock().unwrap().native.get(plugin).cloned();
            let (component, version) = if native.is_some() {
                (None, "native".into())
            } else {
                let _compilation = host.0.compilation.lock().await;
                let bytes = std::fs::read(host.0.world.plugin_dir(plugin).join("plugin.wasm"))
                    .map_err(|_| error(-32001, "no such plugin"))?;
                let hash = format!("{:x}", Sha256::digest(&bytes));
                let cached = host.0.state.lock().unwrap().components.get(&hash).cloned();
                let component = match cached {
                    Some(c) => c,
                    None => {
                        let rt = host.0.runtime.clone();
                        let c = tokio::task::spawn_blocking(move || rt.compile(&bytes))
                            .await
                            .map_err(|e| error(-32603, &e.to_string()))?
                            .map_err(|e| error(-32603, &e.to_string()))?;
                        host.0
                            .state
                            .lock()
                            .unwrap()
                            .components
                            .insert(hash.clone(), c.clone());
                        c
                    }
                };
                (Some(component), hash)
            };
            let pid = format!("{plugin}-{}", Uuid::now_v7());
            let (tx, rx) = oneshot::channel();
            let (input, mut 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 parent_name = info.parent.clone().unwrap_or_else(|| "core".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,
                        input: Some(input),
                        abort: None,
                        caller,
                        id: id.clone(),
                        answer: Some(tx),
                        ended: watch::channel(false).0,
                    },
                );
            }
            let plugin_name = plugin.to_string();
            let h = host.clone();
            let p = pid.clone();
            let method = method.to_string();
            let has_id = id.is_some();
            let task = tokio::spawn(async move {
                let _guard = ProcessGuard {
                    host: h.clone(),
                    pid: p.clone(),
                };
                if let Some(native) = native {
                    let result = native(
                        NativeContext {
                            host: h.clone(),
                            pid: p.clone(),
                        },
                        method,
                        params,
                    )
                    .await;
                    h.answer(&p, result);
                    h.finish(&p, "exit 0");
                    return;
                }
                let cfg = InstanceConfig {
                    world: h.0.world.root.clone(),
                    network: true,
                    env: vec![
                        ("SAUSAGE_PID".into(), p.clone()),
                        ("SAUSAGE_PLUGIN".into(), plugin_name.clone()),
                        ("SAUSAGE_PARENT".into(), parent_name),
                    ],
                };
                let result = async {
                    let mut instance = h.0.runtime.start(&component.unwrap(),cfg)?;
                    let initial = if has_id { json!({"jsonrpc":"2.0","id":p,"method":method,"params":params}) } else {json!({"jsonrpc":"2.0","method":method,"params":params})};
                    instance.send_line(&initial.to_string(),|line|h.output(&p,line)).await?;
                    let mut stdin_open = true;
                    let mut stderr_open = true;
                    loop { tokio::select! {
                        line = instance.stderr.recv(), if stderr_open => match line { Some(line) => h.0.journal.write(&p,&plugin_name,"stderr",&line), None => stderr_open=false },
                        line = instance.stdout.recv() => match line { Some(line) => h.output(&p,line), None => break },
                        value = incoming.recv(), if stdin_open => match value { Some(v) => { instance.send_line(&v.to_string(),|line|h.output(&p,line)).await?; }, None => { instance.stdin.take(); stdin_open=false; } }
                    }}
                    while let Some(line) = instance.stderr.recv().await { h.0.journal.write(&p,&plugin_name,"stderr",&line); }
                    Ok::<_,anyhow::Error>(instance.wait().await.to_string())
                }.await;
                h.finish(&p, &result.unwrap_or_else(|e| e.to_string()));
            });
            if let Some(p) = host.0.state.lock().unwrap().processes.get_mut(&pid) {
                p.abort = Some(task.abort_handle());
            } else {
                task.abort();
            }
            if id.is_none() {
                return Ok(Value::Null);
            }
            rx.await
                .unwrap_or_else(|_| Err(error(-32002, "process ended")))
        })
    }
    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);
        }
    }
    fn output(&self, pid: &str, line: String) {
        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 exit = requests_exit(method, &v["params"]);
            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;
                let success = result.is_ok();
                if let Some(id) = id {
                    let value = match result {
                        Ok(r) => json!({"jsonrpc":"2.0","id":id,"result":r}),
                        Err(e) => json!({"jsonrpc":"2.0","id":id,"error":e}),
                    };
                    h.send(&caller, value);
                }
                if exit && success {
                    h.0.exit.send_replace(true);
                }
            });
        } else if v["id"] == pid {
            self.answer(
                pid,
                if let Some(e) = v.get("error") {
                    Err(e.clone())
                } else {
                    Ok(v["result"].clone())
                },
            );
        }
    }
    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);
        }
    }
    fn finish(&self, pid: &str, status: &str) {
        let mut state = self.0.state.lock().unwrap();
        let process = state.processes.remove(pid);
        if let Some(p) = process {
            self.0.journal.write(pid, &p.info.plugin, "core", status);
            if let Some(tx) = p.answer {
                let _=tx.send(Err(json!({"code":-32002,"message":"process ended without an answer","data":{"status":status,"lines":self.0.journal.query(&json!({"pid":pid,"limit":20}))}})));
            }
            Self::release_locks(&mut state, pid);
            let _ = p.ended.send(true);
            let children: Vec<_> = state
                .processes
                .iter()
                .filter(|(_, p)| p.info.parent.as_deref() == Some(pid))
                .map(|(id, _)| id.clone())
                .collect();
            drop(state);
            for child in children {
                self.begin_stop(&child);
            }
        }
    }
    fn begin_stop(&self, pid: &str) -> Option<watch::Receiver<bool>> {
        let mut state = self.0.state.lock().unwrap();
        let p = state.processes.get_mut(pid)?;
        let ended = p.ended.subscribe();
        if p.info.state == "stopping" {
            return Some(ended);
        }
        p.info.state = "stopping".into();
        p.input.take();
        self.0.journal.write(pid, &p.info.plugin, "core", "stop");
        let children: Vec<_> = state
            .processes
            .iter()
            .filter(|(_, p)| p.info.parent.as_deref() == Some(pid))
            .map(|(id, _)| id.clone())
            .collect();
        drop(state);
        for child in children {
            self.begin_stop(&child);
        }
        let h = self.clone();
        let pid = pid.to_string();
        tokio::spawn(async move {
            tokio::time::sleep(std::time::Duration::from_secs(5)).await;
            h.kill(&pid);
        });
        Some(ended)
    }
    fn kill(&self, pid: &str) {
        if let Some(p) = self.0.state.lock().unwrap().processes.get(pid) {
            self.0.journal.write(pid, &p.info.plugin, "core", "kill");
            if let Some(a) = &p.abort {
                a.abort();
            }
        }
        self.finish(pid, "killed");
    }
    pub async fn stop(&self) {
        let ids: Vec<_> = self.processes().into_keys().collect();
        let mut waits = Vec::new();
        for pid in ids {
            if let Some(n) = self.begin_stop(&pid) {
                waits.push(n);
            }
        }
        for mut n in waits {
            let _ = n.wait_for(|ended| *ended).await;
        }
    }
    fn release_locks(state: &mut State, pid: &str) {
        for lock in state.locks.values_mut() {
            lock.queue.retain(|(p, _)| p != pid);
        }
        let paths: Vec<_> = state
            .locks
            .iter()
            .filter(|(_, l)| l.owner == pid)
            .map(|(p, _)| p.clone())
            .collect();
        for path in paths {
            grant(state, &path);
        }
    }
    async fn core(&self, caller: &str, id: Option<Value>, method: &str, p: Value) -> RpcResult {
        match method {
            "call" => {
                self.route(
                    caller.into(),
                    id,
                    required_str(&p, "method")?.into(),
                    p.get("params").cloned().unwrap_or(json!({})),
                    p["detach"].as_bool().unwrap_or(false),
                )
                .await
            }
            "notify" => {
                self.notify(
                    caller,
                    required_str(&p, "name")?,
                    p.get("params").cloned().unwrap_or(json!({})),
                );
                Ok(Value::Null)
            }
            "kill" => {
                self.kill(required_str(&p, "pid")?);
                Ok(json!({}))
            }
            "stop" => {
                if let Some(mut n) = self.begin_stop(required_str(&p, "pid")?) {
                    let _ = n.wait_for(|ended| *ended).await;
                }
                Ok(json!({}))
            }
            "lock" => {
                let path = p["path"]
                    .as_str()
                    .ok_or_else(|| error(-32602, "missing path"))?;
                let rx = {
                    let mut state = self.0.state.lock().unwrap();
                    if let Some(lock) = state.locks.get_mut(path) {
                        if p["wait"] == false {
                            return Ok(json!({"held":false}));
                        }
                        let (tx, rx) = oneshot::channel();
                        lock.queue.push_back((caller.into(), tx));
                        Some(rx)
                    } else {
                        state.locks.insert(
                            path.into(),
                            Lock {
                                owner: caller.into(),
                                queue: VecDeque::new(),
                            },
                        );
                        None
                    }
                };
                match rx {
                    Some(rx) => rx
                        .await
                        .unwrap_or_else(|_| Err(error(-32603, "lock canceled"))),
                    None => Ok(json!({"held":true})),
                }
            }
            "unlock" => {
                let mut state = self.0.state.lock().unwrap();
                let path = required_str(&p, "path")?;
                if state.locks.get(path).is_some_and(|l| l.owner == caller) {
                    grant(&mut state, path);
                }
                Ok(Value::Null)
            }
            "log" => Ok(self.0.journal.query(&p)),
            "exit" => Ok(json!({})),
            _ => Err(error(-32601, "method not found")),
        }
    }
}
fn grant(state: &mut State, path: &str) {
    if let Some(mut lock) = state.locks.remove(path) {
        while let Some((owner, tx)) = lock.queue.pop_front() {
            if tx.send(Ok(json!({"held":true}))).is_ok() {
                lock.owner = owner;
                state.locks.insert(path.into(), lock);
                break;
            }
        }
    }
}

fn requests_exit(method: &str, params: &Value) -> bool {
    method == "core.exit"
        || (method == "core.call"
            && requests_exit(params["method"].as_str().unwrap_or(""), &params["params"]))
}

fn required_str<'a>(params: &'a Value, name: &str) -> Result<&'a str, Value> {
    params[name]
        .as_str()
        .ok_or_else(|| error(-32602, &format!("{name} must be a string")))
}