openvcs_core/

plugin_runtime.rs

#![cfg(feature = "plugin-protocol")]

use crate::models::VcsEvent;
use crate::plugin_protocol::{PluginMessage, RpcRequest};
use crate::plugin_stdio::{PluginError, receive_message, respond_shared, send_message_shared};
use crate::plugin_stdio::ok_null;
use std::collections::VecDeque;
use std::collections::HashMap;
use std::io::{self, BufReader, LineWriter};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::Duration;

pub type HandlerResult = Result<serde_json::Value, PluginError>;
pub type EventHandlerResult = Result<(), PluginError>;

const DEFAULT_HOST_TIMEOUT_MS: u64 = 60_000;
const HOST_TIMEOUT_ENV: &str = "OPENVCS_PLUGIN_HOST_TIMEOUT_MS";

fn host_timeout() -> Duration {
    std::env::var(HOST_TIMEOUT_ENV)
        .ok()
        .and_then(|s| s.parse::<u64>().ok())
        .filter(|&ms| ms > 0)
        .map(Duration::from_millis)
        .unwrap_or(Duration::from_millis(DEFAULT_HOST_TIMEOUT_MS))
}

pub struct PluginCtx {
    stdout: Arc<Mutex<LineWriter<io::Stdout>>>,
}

impl PluginCtx {
    pub fn emit(&self, event: VcsEvent) {
        send_message_shared(&self.stdout, &PluginMessage::Event { event });
    }

    pub fn stdout(&self) -> Arc<Mutex<LineWriter<io::Stdout>>> {
        Arc::clone(&self.stdout)
    }
}

fn next_request(
    queue: &Arc<Mutex<VecDeque<RpcRequest>>>,
    stdin: &Arc<Mutex<BufReader<io::Stdin>>>,
) -> Option<RpcRequest> {
    if let Ok(mut q) = queue.lock() {
        if let Some(req) = q.pop_front() {
            return Some(req);
        }
    }

    loop {
        let msg = {
            let mut lock = stdin.lock().ok()?;
            receive_message(&mut *lock)?
        };
        match msg {
            PluginMessage::Request(req) => return Some(req),
            PluginMessage::Response(_) | PluginMessage::Event { .. } => continue,
        }
    }
}

pub struct PluginRuntime {
    ctx: PluginCtx,
    stdin: Arc<Mutex<BufReader<io::Stdin>>>,
    queue: Arc<Mutex<VecDeque<RpcRequest>>>,
}

impl PluginRuntime {
    pub fn init() -> Self {
        #[cfg(target_arch = "wasm32")]
        let next_id = 1u64 << 63;

        #[cfg(not(target_arch = "wasm32"))]
        let next_id = 1u64;

        let timeout = host_timeout();

        let stdout = Arc::new(Mutex::new(LineWriter::new(io::stdout())));
        let stdin = Arc::new(Mutex::new(BufReader::new(io::stdin())));
        let queue: Arc<Mutex<VecDeque<RpcRequest>>> = Arc::new(Mutex::new(VecDeque::new()));
        let ids = Arc::new(Mutex::new(crate::plugin_stdio::RequestIdState { next_id }));

        crate::host::init_default_stdio_host(
            Arc::clone(&stdout),
            Arc::clone(&stdin),
            Arc::clone(&queue),
            Arc::clone(&ids),
            timeout,
        );

        Self {
            ctx: PluginCtx { stdout },
            stdin,
            queue,
        }
    }

    pub fn ctx(&mut self) -> &mut PluginCtx {
        &mut self.ctx
    }

    pub(crate) fn tick(
        &mut self,
        mut handle: impl FnMut(&mut PluginCtx, RpcRequest) -> HandlerResult,
    ) -> io::Result<bool> {
        let Some(req) = next_request(&self.queue, &self.stdin) else {
            return Ok(false);
        };
        let id = req.id;
        let res = handle(&mut self.ctx, req);
        respond_shared(&self.ctx.stdout, id, res);
        Ok(true)
    }
}

type RpcHandler = Box<dyn FnMut(&mut PluginCtx, RpcRequest) -> HandlerResult + Send + 'static>;
type EventHandler =
    Box<dyn FnMut(&mut PluginCtx, serde_json::Value) -> EventHandlerResult + Send + 'static>;

struct Registry {
    rpc: HashMap<String, RpcHandler>,
    event: HashMap<String, Vec<EventHandler>>,
    fallback: Option<RpcHandler>,
}

static REGISTRY: OnceLock<Mutex<Registry>> = OnceLock::new();

fn registry() -> &'static Mutex<Registry> {
    REGISTRY.get_or_init(|| {
        Mutex::new(Registry {
            rpc: HashMap::new(),
            event: HashMap::new(),
            fallback: None,
        })
    })
}

fn register_rpc_impl(
    method: &str,
    handler: impl FnMut(&mut PluginCtx, RpcRequest) -> HandlerResult + Send + 'static,
) {
    if let Ok(mut lock) = registry().lock() {
        lock.rpc.insert(method.to_string(), Box::new(handler));
    }
}

pub fn register_delegate(
    method: &str,
    handler: impl FnMut(&mut PluginCtx, RpcRequest) -> HandlerResult + Send + 'static,
) {
    register_rpc_impl(method, handler)
}

/// Backward-compatible alias (will be removed once callers migrate).
#[deprecated(note = "renamed to register_delegate")]
pub fn on_rpc(
    method: &str,
    handler: impl FnMut(&mut PluginCtx, RpcRequest) -> HandlerResult + Send + 'static,
) {
    register_rpc_impl(method, handler)
}

#[deprecated(note = "renamed to register_delegate")]
pub fn register_rpc(
    method: &str,
    handler: impl FnMut(&mut PluginCtx, RpcRequest) -> HandlerResult + Send + 'static,
) {
    register_rpc_impl(method, handler)
}

pub fn on_event(
    name: &str,
    handler: impl FnMut(&mut PluginCtx, serde_json::Value) -> EventHandlerResult + Send + 'static,
) {
    if let Ok(mut lock) = registry().lock() {
        lock.event.entry(name.to_string()).or_default().push(Box::new(handler));
    }
}

pub fn set_fallback_rpc(
    handler: impl FnMut(&mut PluginCtx, RpcRequest) -> HandlerResult + Send + 'static,
) {
    if let Ok(mut lock) = registry().lock() {
        lock.fallback = Some(Box::new(handler));
    }
}

fn dispatch_registered(ctx: &mut PluginCtx, req: RpcRequest) -> HandlerResult {
    if req.method == "event.dispatch" {
        #[derive(serde::Deserialize)]
        struct P {
            name: String,
            #[serde(default)]
            payload: serde_json::Value,
        }
        let p: P = crate::plugin_stdio::parse_json_params(req.params)
            .map_err(PluginError::message)?;

        if let Ok(mut lock) = registry().lock() {
            if let Some(handlers) = lock.event.get_mut(&p.name) {
                for h in handlers.iter_mut() {
                    h(ctx, p.payload.clone())?;
                }
            }
        }
        return ok_null();
    }

    if let Ok(mut lock) = registry().lock() {
        if let Some(h) = lock.rpc.get_mut(&req.method) {
            return h(ctx, req);
        }
        if let Some(fallback) = lock.fallback.as_mut() {
            return fallback(ctx, req);
        }
    }

    Err(PluginError::code(
        "plugin.unknown_method",
        format!("unknown method '{}'", req.method),
    ))
}

pub fn run_registered() -> io::Result<()> {
    let mut rt = PluginRuntime::init();
    while rt.tick(|ctx, req| dispatch_registered(ctx, req))? {}
    Ok(())
}

pub fn run(mut handle: impl FnMut(&mut PluginCtx, RpcRequest) -> HandlerResult) -> io::Result<()> {
    let mut rt = PluginRuntime::init();
    while rt.tick(|ctx, req| handle(ctx, req))? {}
    Ok(())
}