Refactor plugin manager and gRPC server (#96)

src-tauri/yaak_plugin_runtime/src/error.rs

@@ -5,22 +5,35 @@ use crate::server::plugin_runtime::EventStreamEvent;
 
 #[derive(Error, Debug)]
 pub enum Error {
-    #[error("IO error")]
+    #[error("IO error: {0}")]
     IoErr(#[from] io::Error),
-    #[error("Tauri error")]
+    
+    #[error("Tauri error: {0}")]
     TauriErr(#[from] tauri::Error),
-    #[error("Tauri shell error")]
+    
+    #[error("Tauri shell error: {0}")]
     TauriShellErr(#[from] tauri_plugin_shell::Error),
-    #[error("Grpc transport error")]
+    
+    #[error("Grpc transport error: {0}")]
     GrpcTransportErr(#[from] tonic::transport::Error),
-    #[error("Grpc send error")]
+    
+    #[error("Grpc send error: {0}")]
     GrpcSendErr(#[from] SendError<tonic::Result<EventStreamEvent>>),
-    #[error("JSON error")]
+    
+    #[error("JSON error: {0}")]
     JsonErr(#[from] serde_json::Error),
+    
     #[error("Plugin not found: {0}")]
     PluginNotFoundErr(String),
+    
     #[error("Plugin error: {0}")]
     PluginErr(String),
+    
+    #[error("Client not initialized error")]
+    ClientNotInitializedErr,
+    
+    #[error("Unknown event received")]
+    UnknownEventErr,
 }
 
 impl Into<String> for Error {

src-tauri/yaak_plugin_runtime/src/events.rs

@@ -27,6 +27,9 @@ pub enum InternalEventPayload {
     ReloadRequest,
     ReloadResponse,
 
+    TerminateRequest,
+    TerminateResponse,
+
     ImportRequest(ImportRequest),
     ImportResponse(ImportResponse),
 

src-tauri/yaak_plugin_runtime/src/lib.rs

@@ -4,5 +4,5 @@ pub mod manager;
 mod nodejs;
 pub mod plugin;
 mod server;
-pub mod handle;
+pub mod plugin_handle;
 mod util;

src-tauri/yaak_plugin_runtime/src/manager.rs

@@ -1,57 +1,236 @@
+use crate::error::Error::{ClientNotInitializedErr, PluginErr, PluginNotFoundErr, UnknownEventErr};
 use crate::error::Result;
 use crate::events::{
-    BootResponse, CallHttpRequestActionRequest, CallTemplateFunctionArgs,
+    BootRequest, CallHttpRequestActionRequest, CallTemplateFunctionArgs,
     CallTemplateFunctionRequest, CallTemplateFunctionResponse, FilterRequest, FilterResponse,
     GetHttpRequestActionsRequest, GetHttpRequestActionsResponse, GetTemplateFunctionsResponse,
     ImportRequest, ImportResponse, InternalEvent, InternalEventPayload, RenderPurpose,
 };
-use std::collections::HashMap;
-
-use crate::error::Error::PluginErr;
 use crate::nodejs::start_nodejs_plugin_runtime;
-use crate::plugin::start_server;
-use crate::server::PluginRuntimeGrpcServer;
+use crate::plugin_handle::PluginHandle;
+use crate::server::plugin_runtime::plugin_runtime_server::PluginRuntimeServer;
+use crate::server::PluginRuntimeServerImpl;
+use log::{info, warn};
+use std::collections::HashMap;
+use std::path::PathBuf;
+use std::sync::Arc;
 use std::time::Duration;
-use tauri::{AppHandle, Runtime};
-use tokio::sync::mpsc;
-use tokio::sync::watch::Sender;
-use crate::handle::PluginHandle;
+use tauri::path::BaseDirectory;
+use tauri::{AppHandle, Manager, Runtime};
+use tokio::fs::read_dir;
+use tokio::net::TcpListener;
+use tokio::sync::{mpsc, Mutex};
+use tonic::codegen::tokio_stream;
+use tonic::transport::Server;
+use yaak_models::queries::{generate_id, list_plugins};
 
+#[derive(Clone)]
 pub struct PluginManager {
-    kill_tx: Sender<bool>,
-    server: PluginRuntimeGrpcServer,
+    subscribers: Arc<Mutex<HashMap<String, mpsc::Sender<InternalEvent>>>>,
+    plugins: Arc<Mutex<Vec<PluginHandle>>>,
+    kill_tx: tokio::sync::watch::Sender<bool>,
+    server: Arc<PluginRuntimeServerImpl>,
 }
 
 impl PluginManager {
-    pub async fn new<R: Runtime>(
-        app_handle: &AppHandle<R>,
-        plugin_dirs: Vec<String>,
-    ) -> PluginManager {
-        let (server, addr) = start_server(plugin_dirs)
-            .await
-            .expect("Failed to start plugin runtime server");
+    pub fn new<R: Runtime>(app_handle: AppHandle<R>) -> PluginManager {
+        let (events_tx, mut events_rx) = mpsc::channel(128);
+        let (kill_server_tx, kill_server_rx) = tokio::sync::watch::channel(false);
 
-        let (kill_tx, kill_rx) = tokio::sync::watch::channel(false);
-        start_nodejs_plugin_runtime(app_handle, addr, &kill_rx)
-            .await
-            .expect("Failed to start plugin runtime");
+        let (client_disconnect_tx, mut client_disconnect_rx) = mpsc::channel(128);
+        let (client_connect_tx, mut client_connect_rx) = tokio::sync::watch::channel(false);
+        let server =
+            PluginRuntimeServerImpl::new(events_tx, client_disconnect_tx, client_connect_tx);
 
-        PluginManager { kill_tx, server }
+        let plugin_manager = PluginManager {
+            plugins: Arc::new(Mutex::new(Vec::new())),
+            subscribers: Arc::new(Mutex::new(HashMap::new())),
+            server: Arc::new(server.clone()),
+            kill_tx: kill_server_tx,
+        };
+
+        // Forward events to subscribers
+        let subscribers = plugin_manager.subscribers.clone();
+        tauri::async_runtime::spawn(async move {
+            while let Some(event) = events_rx.recv().await {
+                for (tx_id, tx) in subscribers.lock().await.iter_mut() {
+                    if let Err(e) = tx.try_send(event.clone()) {
+                        warn!("Failed to send event to subscriber {tx_id} {e:?}");
+                    }
+                }
+            }
+        });
+
+        // Handle when client plugin runtime disconnects
+        tauri::async_runtime::spawn(async move {
+            while let Some(_) = client_disconnect_rx.recv().await {
+                info!("Plugin runtime client disconnected! TODO: Handle this case");
+            }
+        });
+
+        info!("Starting plugin server");
+
+        let svc = PluginRuntimeServer::new(server.to_owned());
+        let listen_addr = match option_env!("PORT") {
+            None => "localhost:0".to_string(),
+            Some(port) => format!("localhost:{port}"),
+        };
+        let listener = tauri::async_runtime::block_on(async move {
+            TcpListener::bind(listen_addr)
+                .await
+                .expect("Failed to bind TCP listener")
+        });
+        let addr = listener.local_addr().expect("Failed to get local address");
+
+        // 1. Reload all plugins when the Node.js runtime connects
+        {
+            let plugin_manager = plugin_manager.clone();
+            let app_handle = app_handle.clone();
+            tauri::async_runtime::spawn(async move {
+                match client_connect_rx.changed().await {
+                    Ok(_) => {
+                        info!("Plugin runtime client connected!");
+                        plugin_manager
+                            .initialize_all_plugins(&app_handle)
+                            .await
+                            .expect("Failed to reload plugins");
+                    }
+                    Err(e) => {
+                        warn!("Failed to receive from client connection rx {e:?}");
+                    }
+                }
+            });
+        };
+
+        // 1. Spawn server in the background
+        info!("Starting gRPC plugin server on {addr}");
+        tauri::async_runtime::spawn(async move {
+            Server::builder()
+                .timeout(Duration::from_secs(10))
+                .add_service(svc)
+                .serve_with_incoming(tokio_stream::wrappers::TcpListenerStream::new(listener))
+                .await
+                .expect("grpc plugin runtime server failed to start");
+        });
+
+        // 2. Start Node.js runtime and initialize plugins
+        tauri::async_runtime::block_on(async move {
+            start_nodejs_plugin_runtime(&app_handle, addr, &kill_server_rx)
+                .await
+                .unwrap();
+        });
+
+        plugin_manager
     }
 
-    pub async fn reload_all(&self) {
-        self.server.reload_plugins().await
+    pub async fn list_plugin_dirs<R: Runtime>(&self, app_handle: &AppHandle<R>) -> Vec<String> {
+        let plugins_dir = app_handle
+            .path()
+            .resolve("plugins", BaseDirectory::Resource)
+            .expect("failed to resolve plugin directory resource");
+
+        let bundled_plugin_dirs = read_plugins_dir(&plugins_dir)
+            .await
+            .expect(format!("Failed to read plugins dir: {:?}", plugins_dir).as_str());
+
+        let plugins = list_plugins(app_handle).await.unwrap_or_default();
+        let installed_plugin_dirs = plugins
+            .iter()
+            .map(|p| p.directory.to_owned())
+            .collect::<Vec<String>>();
+
+        let plugin_dirs = [bundled_plugin_dirs, installed_plugin_dirs].concat();
+        plugin_dirs
+    }
+
+    pub async fn uninstall(&self, dir: &str) -> Result<()> {
+        let plugin = self
+            .get_plugin_by_dir(dir)
+            .await
+            .ok_or(PluginNotFoundErr(dir.to_string()))?;
+        self.remove_plugin(&plugin).await
+    }
+
+    async fn remove_plugin(&self, plugin: &PluginHandle) -> Result<()> {
+        let mut plugins = self.plugins.lock().await;
+
+        // Terminate the plugin
+        plugin.terminate().await?;
+
+        // Remove the plugin from the list
+        let pos = plugins.iter().position(|p| p.ref_id == plugin.ref_id);
+        if let Some(pos) = pos {
+            plugins.remove(pos);
+        }
+
+        Ok(())
+    }
+
+    pub async fn add_plugin_by_dir(&self, dir: &str) -> Result<()> {
+        info!("Adding plugin by dir {dir}");
+        let maybe_tx = self.server.app_to_plugin_events_tx.lock().await;
+        let tx = match &*maybe_tx {
+            None => return Err(ClientNotInitializedErr),
+            Some(tx) => tx,
+        };
+        let ph = PluginHandle::new(dir, tx.clone());
+        self.plugins.lock().await.push(ph.clone());
+        let plugin = self
+            .get_plugin_by_dir(dir)
+            .await
+            .ok_or(PluginNotFoundErr(dir.to_string()))?;
+
+        // Boot the plugin
+        let event = self
+            .send_to_plugin_and_wait(
+                &plugin,
+                &InternalEventPayload::BootRequest(BootRequest {
+                    dir: dir.to_string(),
+                }),
+            )
+            .await?;
+
+        let resp = match event.payload {
+            InternalEventPayload::BootResponse(resp) => resp,
+            _ => return Err(UnknownEventErr),
+        };
+
+        plugin.set_boot_response(&resp).await;
+
+        Ok(())
+    }
+
+    pub async fn initialize_all_plugins<R: Runtime>(
+        &self,
+        app_handle: &AppHandle<R>,
+    ) -> Result<()> {
+        for dir in self.list_plugin_dirs(app_handle).await {
+            // First remove the plugin if it exists
+            if let Some(plugin) = self.get_plugin_by_dir(dir.as_str()).await {
+                if let Err(e) = self.remove_plugin(&plugin).await {
+                    warn!("Failed to remove plugin {dir} {e:?}");
+                }
+            }
+            if let Err(e) = self.add_plugin_by_dir(dir.as_str()).await {
+                warn!("Failed to add plugin {dir} {e:?}");
+            }
+        }
+
+        Ok(())
     }
 
     pub async fn subscribe(&self) -> (String, mpsc::Receiver<InternalEvent>) {
-        self.server.subscribe().await
+        let (tx, rx) = mpsc::channel(128);
+        let rx_id = generate_id();
+        self.subscribers.lock().await.insert(rx_id.clone(), tx);
+        (rx_id, rx)
     }
 
     pub async fn unsubscribe(&self, rx_id: &str) {
-        self.server.unsubscribe(rx_id).await
+        self.subscribers.lock().await.remove(rx_id);
     }
 
-    pub async fn cleanup(&self) {
+    pub async fn terminate(&self) {
         self.kill_tx.send_replace(true);
 
         // Give it a bit of time to kill
@@ -64,22 +243,115 @@ impl PluginManager {
         payload: &InternalEventPayload,
     ) -> Result<()> {
         let reply_id = Some(source_event.clone().id);
-        self.server
-            .send(&payload, source_event.plugin_ref_id.as_str(), reply_id)
+        let plugin = self
+            .get_plugin_by_ref_id(source_event.plugin_ref_id.as_str())
+            .await
+            .ok_or(PluginNotFoundErr(source_event.plugin_ref_id.to_string()))?;
+        let event = plugin.build_event_to_send(&payload, reply_id);
+        plugin.send(&event).await
+    }
+
+    pub async fn get_plugin_by_ref_id(&self, ref_id: &str) -> Option<PluginHandle> {
+        self.plugins
+            .lock()
+            .await
+            .iter()
+            .find(|p| p.ref_id == ref_id)
+            .cloned()
+    }
+
+    pub async fn get_plugin_by_dir(&self, dir: &str) -> Option<PluginHandle> {
+        self.plugins
+            .lock()
+            .await
+            .iter()
+            .find(|p| p.dir == dir)
+            .cloned()
+    }
+
+    pub async fn get_plugin_by_name(&self, name: &str) -> Option<PluginHandle> {
+        for plugin in self.plugins.lock().await.iter().cloned() {
+            let info = plugin.info().await?;
+            if info.name == name {
+                return Some(plugin);
+            }
+        }
+        None
+    }
+
+    async fn send_to_plugin_and_wait(
+        &self,
+        plugin: &PluginHandle,
+        payload: &InternalEventPayload,
+    ) -> Result<InternalEvent> {
+        let events = self
+            .send_to_plugins_and_wait(payload, vec![plugin.to_owned()])
+            .await?;
+        Ok(events.first().unwrap().to_owned())
+    }
+
+    async fn send_and_wait(&self, payload: &InternalEventPayload) -> Result<Vec<InternalEvent>> {
+        self.send_to_plugins_and_wait(payload, self.plugins.lock().await.clone())
             .await
     }
 
-    pub async fn get_plugin_info(&self, dir: &str) -> Option<BootResponse> {
-        self.server.plugin_by_dir(dir).await.ok()?.info().await
-    }
+    async fn send_to_plugins_and_wait(
+        &self,
+        payload: &InternalEventPayload,
+        plugins: Vec<PluginHandle>,
+    ) -> Result<Vec<InternalEvent>> {
+        let (rx_id, mut rx) = self.subscribe().await;
 
-    pub async fn get_plugin(&self, ref_id: &str) -> Result<PluginHandle> {
-        self.server.plugin_by_ref_id(ref_id).await
+        // 1. Build the events with IDs and everything
+        let events_to_send = plugins
+            .iter()
+            .map(|p| p.build_event_to_send(payload, None))
+            .collect::<Vec<InternalEvent>>();
+
+        // 2. Spawn thread to subscribe to incoming events and check reply ids
+        let send_events_fut = {
+            let events_to_send = events_to_send.clone();
+
+            tokio::spawn(async move {
+                let mut found_events = Vec::new();
+
+                while let Some(event) = rx.recv().await {
+                    if events_to_send
+                        .iter()
+                        .find(|e| Some(e.id.to_owned()) == event.reply_id)
+                        .is_some()
+                    {
+                        found_events.push(event.clone());
+                    };
+                    if found_events.len() == events_to_send.len() {
+                        break;
+                    }
+                }
+
+                found_events
+            })
+        };
+
+        // 3. Send the events
+        for event in events_to_send {
+            let plugin = plugins
+                .iter()
+                .find(|p| p.ref_id == event.plugin_ref_id)
+                .expect("Didn't find plugin in list");
+            plugin.send(&event).await?
+        }
+
+        // 4. Join on the spawned thread
+        let events = send_events_fut.await.expect("Thread didn't succeed");
+
+        // 5. Unsubscribe
+        self.unsubscribe(rx_id.as_str()).await;
+
+        Ok(events)
     }
 
     pub async fn get_http_request_actions(&self) -> Result<Vec<GetHttpRequestActionsResponse>> {
         let reply_events = self
-            .server
             .send_and_wait(&InternalEventPayload::GetHttpRequestActionsRequest(
                 GetHttpRequestActionsRequest {},
             ))
@@ -97,7 +369,6 @@ impl PluginManager {
 
     pub async fn get_template_functions(&self) -> Result<Vec<GetTemplateFunctionsResponse>> {
         let reply_events = self
-            .server
             .send_and_wait(&InternalEventPayload::GetTemplateFunctionsRequest)
             .await?;
 
@@ -112,10 +383,11 @@ impl PluginManager {
     }
 
     pub async fn call_http_request_action(&self, req: CallHttpRequestActionRequest) -> Result<()> {
+        let ref_id = req.plugin_ref_id.clone();
         let plugin = self
-            .server
-            .plugin_by_ref_id(req.plugin_ref_id.as_str())
-            .await?;
+            .get_plugin_by_ref_id(ref_id.as_str())
+            .await
+            .ok_or(PluginNotFoundErr(ref_id))?;
         let event = plugin.build_event_to_send(
             &InternalEventPayload::CallHttpRequestActionRequest(req),
             None,
@@ -139,7 +411,6 @@ impl PluginManager {
         };
 
         let events = self
-            .server
             .send_and_wait(&InternalEventPayload::CallTemplateFunctionRequest(req))
             .await?;
 
@@ -155,7 +426,6 @@ impl PluginManager {
 
     pub async fn import_data(&self, content: &str) -> Result<(ImportResponse, String)> {
         let reply_events = self
-            .server
             .send_and_wait(&InternalEventPayload::ImportRequest(ImportRequest {
                 content: content.to_string(),
             }))
@@ -172,9 +442,12 @@ impl PluginManager {
                 "No importers found for file contents".to_string(),
             )),
             Some((resp, ref_id)) => {
-                let plugin = self.server.plugin_by_ref_id(ref_id.as_str()).await?;
-                let plugin_name = plugin.name().await;
-                Ok((resp, plugin_name))
+                let plugin = self
+                    .get_plugin_by_ref_id(ref_id.as_str())
+                    .await
+                    .ok_or(PluginNotFoundErr(ref_id))?;
+                let info = plugin.info().await.unwrap();
+                Ok((resp, info.name))
             }
         }
     }
@@ -191,10 +464,14 @@ impl PluginManager {
             "filter-xpath"
         };
 
+        let plugin = self
+            .get_plugin_by_dir(plugin_name)
+            .await
+            .ok_or(PluginNotFoundErr(plugin_name.to_string()))?;
+
         let event = self
-            .server
             .send_to_plugin_and_wait(
-                plugin_name,
+                &plugin,
                 &InternalEventPayload::FilterRequest(FilterRequest {
                     filter: filter.to_string(),
                     content: content.to_string(),
@@ -211,3 +488,39 @@ impl PluginManager {
         }
     }
 }
+
+async fn read_plugins_dir(dir: &PathBuf) -> Result<Vec<String>> {
+    let mut result = read_dir(dir).await?;
+    let mut dirs: Vec<String> = vec![];
+    while let Ok(Some(entry)) = result.next_entry().await {
+        if entry.path().is_dir() {
+            #[cfg(target_os = "windows")]
+            dirs.push(fix_windows_paths(&entry.path()));
+            #[cfg(not(target_os = "windows"))]
+            dirs.push(entry.path().to_string_lossy().to_string());
+        }
+    }
+    Ok(dirs)
+}
+
+#[cfg(target_os = "windows")]
+fn fix_windows_paths(p: &PathBuf) -> String {
+    use dunce;
+    use path_slash::PathBufExt;
+    use regex::Regex;
+
+    // 1. Remove UNC prefix for Windows paths to pass to sidecar
+    let safe_path = dunce::simplified(p.as_path()).to_string_lossy().to_string();
+
+    // 2. Remove the drive letter
+    let safe_path = Regex::new("^[a-zA-Z]:")
+        .unwrap()
+        .replace(safe_path.as_str(), "");
+
+    // 3. Convert backslashes to forward
+    let safe_path = PathBuf::from(safe_path.to_string())
+        .to_slash_lossy()
+        .to_string();
+
+    safe_path
+}

src-tauri/yaak_plugin_runtime/src/nodejs.rs

@@ -39,7 +39,7 @@ pub async fn start_nodejs_plugin_runtime<R: Runtime>(
         .args(&[plugin_runtime_main]);
 
     let (mut child_rx, child) = cmd.spawn()?;
-    println!("Spawned plugin runtime");
+    info!("Spawned plugin runtime");
 
     let mut kill_rx = kill_rx.clone();
 

src-tauri/yaak_plugin_runtime/src/plugin.rs

@@ -1,46 +1,16 @@
-use crate::error::Result;
-use crate::events::{InternalEvent, InternalEventPayload};
 use crate::manager::PluginManager;
-use crate::server::plugin_runtime::plugin_runtime_server::PluginRuntimeServer;
-use crate::server::PluginRuntimeGrpcServer;
 use log::info;
-use std::net::SocketAddr;
-use std::path::PathBuf;
 use std::process::exit;
-use std::time::Duration;
-use tauri::path::BaseDirectory;
 use tauri::plugin::{Builder, TauriPlugin};
 use tauri::{Manager, RunEvent, Runtime, State};
-use tokio::fs::read_dir;
-use tokio::net::TcpListener;
-use tonic::codegen::tokio_stream;
-use tonic::transport::Server;
-use yaak_models::queries::list_plugins;
 
 pub fn init<R: Runtime>() -> TauriPlugin<R> {
     Builder::new("yaak_plugin_runtime")
         .setup(|app_handle, _| {
-            let plugins_dir = app_handle
-                .path()
-                .resolve("plugins", BaseDirectory::Resource)
-                .expect("failed to resolve plugin directory resource");
+            let manager = PluginManager::new(app_handle.clone());
+            app_handle.manage(manager.clone());
 
-            tauri::async_runtime::block_on(async move {
-                let bundled_plugin_dirs = read_plugins_dir(&plugins_dir)
-                    .await
-                    .expect(format!("Failed to read plugins dir: {:?}", plugins_dir).as_str());
-
-                let plugins = list_plugins(app_handle).await.unwrap_or_default();
-                let installed_plugin_dirs = plugins
-                    .iter()
-                    .map(|p| p.directory.to_owned())
-                    .collect::<Vec<String>>();
-
-                let plugin_dirs = [installed_plugin_dirs, bundled_plugin_dirs].concat();
-                let manager = PluginManager::new(&app_handle, plugin_dirs).await;
-                app_handle.manage(manager);
-                Ok(())
-            })
+            Ok(())
         })
         .on_event(|app, e| match e {
             // TODO: Also exit when app is force-quit (eg. cmd+r in IntelliJ runner)
@@ -49,94 +19,11 @@ pub fn init<R: Runtime>() -> TauriPlugin<R> {
                 tauri::async_runtime::block_on(async move {
                     info!("Exiting plugin runtime due to app exit");
                     let manager: State<PluginManager> = app.state();
-                    manager.cleanup().await;
+                    manager.terminate().await;
                     exit(0);
                 });
             }
             _ => {}
         })
         .build()
-}
-
-pub async fn start_server(
-    plugin_dirs: Vec<String>,
-) -> Result<(PluginRuntimeGrpcServer, SocketAddr)> {
-    println!("Starting plugin server with {plugin_dirs:?}");
-    let server = PluginRuntimeGrpcServer::new(plugin_dirs);
-
-    let svc = PluginRuntimeServer::new(server.clone());
-    let listen_addr = match option_env!("PORT") {
-        None => "localhost:0".to_string(),
-        Some(port) => format!("localhost:{port}"),
-    };
-
-    {
-        let server = server.clone();
-        tokio::spawn(async move {
-            let (rx_id, mut rx) = server.subscribe().await;
-            while let Some(event) = rx.recv().await {
-                match event.clone() {
-                    InternalEvent {
-                        payload: InternalEventPayload::BootResponse(resp),
-                        plugin_ref_id,
-                        ..
-                    } => {
-                        server.boot_plugin(plugin_ref_id.as_str(), &resp).await;
-                    }
-                    _ => {}
-                };
-            }
-            server.unsubscribe(rx_id.as_str()).await;
-        });
-    };
-
-    let listener = TcpListener::bind(listen_addr).await?;
-    let addr = listener.local_addr()?;
-    println!("Starting gRPC plugin server on {addr}");
-    tokio::spawn(async move {
-        Server::builder()
-            .timeout(Duration::from_secs(10))
-            .add_service(svc)
-            .serve_with_incoming(tokio_stream::wrappers::TcpListenerStream::new(listener))
-            .await
-            .expect("grpc plugin runtime server failed to start");
-    });
-
-    Ok((server, addr))
-}
-
-async fn read_plugins_dir(dir: &PathBuf) -> Result<Vec<String>> {
-    let mut result = read_dir(dir).await?;
-    let mut dirs: Vec<String> = vec![];
-    while let Ok(Some(entry)) = result.next_entry().await {
-        if entry.path().is_dir() {
-            #[cfg(target_os = "windows")]
-            dirs.push(fix_windows_paths(&entry.path()));
-            #[cfg(not(target_os = "windows"))]
-            dirs.push(entry.path().to_string_lossy().to_string());
-        }
-    }
-    Ok(dirs)
-}
-
-#[cfg(target_os = "windows")]
-fn fix_windows_paths(p: &PathBuf) -> String {
-    use dunce;
-    use path_slash::PathBufExt;
-    use regex::Regex;
-
-    // 1. Remove UNC prefix for Windows paths to pass to sidecar
-    let safe_path = dunce::simplified(p.as_path()).to_string_lossy().to_string();
-
-    // 2. Remove the drive letter
-    let safe_path = Regex::new("^[a-zA-Z]:")
-        .unwrap()
-        .replace(safe_path.as_str(), "");
-
-    // 3. Convert backslashes to forward
-    let safe_path = PathBuf::from(safe_path.to_string())
-        .to_slash_lossy()
-        .to_string();
-
-    safe_path
-}
+}

src-tauri/yaak_plugin_runtime/src/plugin_handle.rs

@@ -1,7 +1,9 @@
+use crate::error::Result;
 use crate::events::{BootResponse, InternalEvent, InternalEventPayload};
 use crate::server::plugin_runtime::EventStreamEvent;
 use crate::util::gen_id;
 use std::sync::Arc;
+use log::info;
 use tokio::sync::{mpsc, Mutex};
 
 #[derive(Clone)]
@@ -13,10 +15,14 @@ pub struct PluginHandle {
 }
 
 impl PluginHandle {
-    pub async fn name(&self) -> String {
-        match &*self.boot_resp.lock().await {
-            None => "__NOT_BOOTED__".to_string(),
-            Some(r) => r.name.to_owned(),
+    pub fn new(dir: &str, tx: mpsc::Sender<tonic::Result<EventStreamEvent>>) -> Self {
+        let ref_id = gen_id();
+
+        PluginHandle {
+            ref_id: ref_id.clone(),
+            dir: dir.to_string(),
+            to_plugin_tx: Arc::new(Mutex::new(tx)),
+            boot_resp: Arc::new(Mutex::new(None)),
         }
     }
 
@@ -38,28 +44,33 @@ impl PluginHandle {
         }
     }
 
-    pub async fn reload(&self) -> crate::error::Result<()> {
-        let event = self.build_event_to_send(&InternalEventPayload::ReloadRequest, None);
+    pub async fn terminate(&self) -> Result<()> {
+        info!("Terminating plugin {}", self.dir);
+        let event = self.build_event_to_send(&InternalEventPayload::TerminateRequest, None);
         self.send(&event).await
     }
 
-    pub async fn send(&self, event: &InternalEvent) -> crate::error::Result<()> {
-        // info!(
-        //     "Sending event to plugin {} {:?}",
-        //     event.id,
-        //     self.name().await
-        // );
+    pub async fn send(&self, event: &InternalEvent) -> Result<()> {
         self.to_plugin_tx
             .lock()
             .await
             .send(Ok(EventStreamEvent {
-                event: serde_json::to_string(&event)?,
+                event: serde_json::to_string(event)?,
             }))
             .await?;
         Ok(())
     }
 
-    pub async fn boot(&self, resp: &BootResponse) {
+    pub async fn send_payload(
+        &self,
+        payload: &InternalEventPayload,
+        reply_id: Option<String>,
+    ) -> Result<()> {
+        let event = self.build_event_to_send(payload, reply_id);
+        self.send(&event).await
+    }
+
+    pub async fn set_boot_response(&self, resp: &BootResponse) {
         let mut boot_resp = self.boot_resp.lock().await;
         *boot_resp = Some(resp.clone());
     }

src-tauri/yaak_plugin_runtime/src/server.rs

@@ -1,293 +1,47 @@
-use std::collections::HashMap;
+use log::warn;
 use std::pin::Pin;
 use std::sync::Arc;
-use log::warn;
-use tokio::sync::mpsc::Receiver;
 use tokio::sync::{mpsc, Mutex};
 use tonic::codegen::tokio_stream::wrappers::ReceiverStream;
 use tonic::codegen::tokio_stream::{Stream, StreamExt};
 use tonic::{Request, Response, Status, Streaming};
 
-use crate::error::Error::PluginNotFoundErr;
-use crate::error::Result;
-use crate::events::{InternalEvent, InternalEventPayload, BootRequest, BootResponse};
-use crate::handle::PluginHandle;
+use crate::events::InternalEvent;
 use crate::server::plugin_runtime::plugin_runtime_server::PluginRuntime;
-use crate::util::gen_id;
 use plugin_runtime::EventStreamEvent;
-use yaak_models::queries::generate_id;
 
 pub mod plugin_runtime {
     tonic::include_proto!("yaak.plugins.runtime");
 }
 
-type ResponseStream =
-    Pin<Box<dyn Stream<Item = std::result::Result<EventStreamEvent, Status>> + Send>>;
+type ResponseStream = Pin<Box<dyn Stream<Item = Result<EventStreamEvent, Status>> + Send>>;
 
 #[derive(Clone)]
-pub struct PluginRuntimeGrpcServer {
-    plugin_ref_to_plugin: Arc<Mutex<HashMap<String, PluginHandle>>>,
-    callback_to_plugin_ref: Arc<Mutex<HashMap<String, String>>>,
-    subscribers: Arc<Mutex<HashMap<String, mpsc::Sender<InternalEvent>>>>,
-    plugin_dirs: Vec<String>,
+pub(crate) struct PluginRuntimeServerImpl {
+    pub(crate) app_to_plugin_events_tx:
+        Arc<Mutex<Option<mpsc::Sender<tonic::Result<EventStreamEvent>>>>>,
+    client_disconnect_tx: mpsc::Sender<bool>,
+    client_connect_tx: tokio::sync::watch::Sender<bool>,
+    plugin_to_app_events_tx: mpsc::Sender<InternalEvent>,
 }
 
-impl PluginRuntimeGrpcServer {
-    pub fn new(plugin_dirs: Vec<String>) -> Self {
-        PluginRuntimeGrpcServer {
-            plugin_ref_to_plugin: Arc::new(Mutex::new(HashMap::new())),
-            callback_to_plugin_ref: Arc::new(Mutex::new(HashMap::new())),
-            subscribers: Arc::new(Mutex::new(HashMap::new())),
-            plugin_dirs,
+impl PluginRuntimeServerImpl {
+    pub fn new(
+        events_tx: mpsc::Sender<InternalEvent>,
+        disconnect_tx: mpsc::Sender<bool>,
+        connect_tx: tokio::sync::watch::Sender<bool>,
+    ) -> Self {
+        PluginRuntimeServerImpl {
+            app_to_plugin_events_tx: Arc::new(Mutex::new(None)),
+            client_disconnect_tx: disconnect_tx,
+            client_connect_tx: connect_tx,
+            plugin_to_app_events_tx: events_tx,
         }
     }
-
-    pub async fn plugins(&self) -> Vec<PluginHandle> {
-        self.plugin_ref_to_plugin
-            .lock()
-            .await
-            .iter()
-            .map(|p| p.1.to_owned())
-            .collect::<Vec<PluginHandle>>()
-    }
-
-    pub async fn subscribe(&self) -> (String, Receiver<InternalEvent>) {
-        let (tx, rx) = mpsc::channel(128);
-        let rx_id = generate_id();
-        self.subscribers.lock().await.insert(rx_id.clone(), tx);
-        (rx_id, rx)
-    }
-
-    pub async fn unsubscribe(&self, rx_id: &str) {
-        self.subscribers.lock().await.remove(rx_id);
-    }
-
-    pub async fn remove_plugins(&self, plugin_ids: Vec<String>) {
-        for plugin_id in plugin_ids {
-            self.remove_plugin(plugin_id.as_str()).await;
-        }
-    }
-
-    pub async fn remove_plugin(&self, id: &str) {
-        match self.plugin_ref_to_plugin.lock().await.remove(id) {
-            None => println!("Tried to remove non-existing plugin {}", id),
-            Some(plugin) => println!("Removed plugin {} {}", id, plugin.name().await),
-        };
-    }
-
-    pub async fn boot_plugin(&self, id: &str, resp: &BootResponse) {
-        match self.plugin_ref_to_plugin.lock().await.get(id) {
-            None => println!("Tried booting non-existing plugin {}", id),
-            Some(plugin) => plugin.clone().boot(resp).await,
-        }
-    }
-
-    pub async fn add_plugin(
-        &self,
-        dir: &str,
-        tx: mpsc::Sender<tonic::Result<EventStreamEvent>>,
-    ) -> PluginHandle {
-        let ref_id = gen_id();
-        let plugin_handle = PluginHandle {
-            ref_id: ref_id.clone(),
-            dir: dir.to_string(),
-            to_plugin_tx: Arc::new(Mutex::new(tx)),
-            boot_resp: Arc::new(Mutex::new(None)),
-        };
-        let _ = self
-            .plugin_ref_to_plugin
-            .lock()
-            .await
-            .insert(ref_id, plugin_handle.clone());
-        plugin_handle
-    }
-
-    pub async fn plugin_by_ref_id(&self, ref_id: &str) -> Result<PluginHandle> {
-        let plugins = self.plugin_ref_to_plugin.lock().await;
-        match plugins.get(ref_id) {
-            None => Err(PluginNotFoundErr(ref_id.into())),
-            Some(p) => Ok(p.to_owned()),
-        }
-    }
-
-    pub async fn plugin_by_dir(&self, dir: &str) -> Result<PluginHandle> {
-        let plugins = self.plugin_ref_to_plugin.lock().await;
-        for p in plugins.values() {
-            if p.dir == dir {
-                return Ok(p.to_owned());
-            }
-        }
-
-        Err(PluginNotFoundErr(dir.into()))
-    }
-
-    pub async fn plugin_by_name(&self, plugin_name: &str) -> Result<PluginHandle> {
-        let plugins = self.plugin_ref_to_plugin.lock().await;
-        for p in plugins.values() {
-            if p.name().await == plugin_name {
-                return Ok(p.to_owned());
-            }
-        }
-
-        Err(PluginNotFoundErr(plugin_name.into()))
-    }
-
-    pub async fn send(
-        &self,
-        payload: &InternalEventPayload,
-        plugin_ref_id: &str,
-        reply_id: Option<String>,
-    ) -> Result<()> {
-        let plugin = self.plugin_by_ref_id(plugin_ref_id).await?;
-        let event = plugin.build_event_to_send(payload, reply_id);
-        plugin.send(&event).await
-    }
-
-    pub async fn send_to_plugin(
-        &self,
-        plugin_name: &str,
-        payload: InternalEventPayload,
-    ) -> Result<InternalEvent> {
-        let plugins = self.plugin_ref_to_plugin.lock().await;
-        if plugins.is_empty() {
-            return Err(PluginNotFoundErr(plugin_name.into()));
-        }
-
-        let mut plugin = None;
-        for p in plugins.values() {
-            if p.name().await == plugin_name {
-                plugin = Some(p);
-                break;
-            }
-        }
-
-        match plugin {
-            Some(plugin) => {
-                let event = plugin.build_event_to_send(&payload, None);
-                plugin.send(&event).await?;
-                Ok(event)
-            }
-            None => Err(PluginNotFoundErr(plugin_name.into())),
-        }
-    }
-
-    pub async fn send_to_plugin_and_wait(
-        &self,
-        plugin_name: &str,
-        payload: &InternalEventPayload,
-    ) -> Result<InternalEvent> {
-        let plugin = self.plugin_by_name(plugin_name).await?;
-        let events = self.send_to_plugins_and_wait(payload, vec![plugin]).await?;
-        Ok(events.first().unwrap().to_owned())
-    }
-
-    pub async fn send_and_wait(
-        &self,
-        payload: &InternalEventPayload,
-    ) -> Result<Vec<InternalEvent>> {
-        let plugins = self
-            .plugin_ref_to_plugin
-            .lock()
-            .await
-            .values()
-            .cloned()
-            .collect();
-        self.send_to_plugins_and_wait(payload, plugins).await
-    }
-
-    async fn send_to_plugins_and_wait(
-        &self,
-        payload: &InternalEventPayload,
-        plugins: Vec<PluginHandle>,
-    ) -> Result<Vec<InternalEvent>> {
-        // 1. Build the events with IDs and everything
-        let events_to_send = plugins
-            .iter()
-            .map(|p| p.build_event_to_send(payload, None))
-            .collect::<Vec<InternalEvent>>();
-
-        // 2. Spawn thread to subscribe to incoming events and check reply ids
-        let server = self.clone();
-        let send_events_fut = {
-            let events_to_send = events_to_send.clone();
-            tokio::spawn(async move {
-                let (rx_id, mut rx) = server.subscribe().await;
-                let mut found_events = Vec::new();
-
-                while let Some(event) = rx.recv().await {
-                    if events_to_send
-                        .iter()
-                        .find(|e| Some(e.id.to_owned()) == event.reply_id)
-                        .is_some()
-                    {
-                        found_events.push(event.clone());
-                    };
-                    if found_events.len() == events_to_send.len() {
-                        break;
-                    }
-                }
-                server.unsubscribe(rx_id.as_str()).await;
-
-                found_events
-            })
-        };
-
-        // 3. Send the events
-        for event in events_to_send {
-            let plugin = plugins
-                .iter()
-                .find(|p| p.ref_id == event.plugin_ref_id)
-                .expect("Didn't find plugin in list");
-            plugin.send(&event).await?
-        }
-
-        // 4. Join on the spawned thread
-        let events = send_events_fut.await.expect("Thread didn't succeed");
-        Ok(events)
-    }
-
-    pub async fn reload_plugins(&self) {
-        for (_, plugin) in self.plugin_ref_to_plugin.lock().await.clone() {
-            if let Err(e) = plugin.reload().await {
-                warn!("Failed to reload plugin {} {}", plugin.dir, e)
-            }
-        }
-    }
-
-    async fn load_plugins(
-        &self,
-        to_plugin_tx: mpsc::Sender<tonic::Result<EventStreamEvent>>,
-        plugin_dirs: Vec<String>,
-    ) -> Vec<String> {
-        let mut plugin_ids = Vec::new();
-
-        for dir in plugin_dirs {
-            let plugin = self.add_plugin(dir.as_str(), to_plugin_tx.clone()).await;
-            plugin_ids.push(plugin.clone().ref_id);
-
-            let event = plugin.build_event_to_send(
-                &InternalEventPayload::BootRequest(BootRequest {
-                    dir: dir.to_string(),
-                }),
-                None,
-            );
-            if let Err(e) = plugin.send(&event).await {
-                // TODO: Error handling
-                println!(
-                    "Failed boot plugin {} at {} -> {}",
-                    plugin.ref_id, plugin.dir, e
-                )
-            } else {
-                println!("Loaded plugin {} at {}", plugin.ref_id, plugin.dir)
-            }
-        }
-
-        plugin_ids
-    }
 }
 
 #[tonic::async_trait]
-impl PluginRuntime for PluginRuntimeGrpcServer {
+impl PluginRuntime for PluginRuntimeServerImpl {
     type EventStreamStream = ResponseStream;
 
     async fn event_stream(
@@ -296,51 +50,48 @@ impl PluginRuntime for PluginRuntimeGrpcServer {
     ) -> tonic::Result<Response<Self::EventStreamStream>> {
         let mut in_stream = req.into_inner();
 
-        let (to_plugin_tx, to_plugin_rx) = mpsc::channel(128);
+        let (to_plugin_tx, to_plugin_rx) = mpsc::channel::<tonic::Result<EventStreamEvent>>(128);
+        let mut app_to_plugin_events_tx = self.app_to_plugin_events_tx.lock().await;
+        *app_to_plugin_events_tx = Some(to_plugin_tx);
+        println!("GRPC CLIENT CONNECTED");
 
-        let plugin_ids = self
-            .load_plugins(to_plugin_tx, self.plugin_dirs.clone())
-            .await;
+        let plugin_to_app_events_tx = self.plugin_to_app_events_tx.clone();
+        let client_disconnect_tx = self.client_disconnect_tx.clone();
+
+        self.client_connect_tx
+            .send(true)
+            .expect("Failed to send client ready event");
 
-        let callbacks = self.callback_to_plugin_ref.clone();
-        let server = self.clone();
         tokio::spawn(async move {
             while let Some(result) = in_stream.next().await {
+                // Received event from plugin runtime
                 match result {
                     Ok(v) => {
                         let event: InternalEvent = match serde_json::from_str(v.event.as_str()) {
                             Ok(pe) => pe,
                             Err(e) => {
-                                println!("Failed to deserialize event {e:?} -> {}", v.event);
+                                warn!("Failed to deserialize event {e:?} -> {}", v.event);
                                 continue;
                             }
                         };
 
-                        let plugin_ref_id = event.plugin_ref_id.clone();
-                        let reply_id = event.reply_id.clone();
-
-                        let subscribers = server.subscribers.lock().await;
-                        for tx in subscribers.values() {
-                            // Emit event to the channel for server to handle
-                            if let Err(e) = tx.try_send(event.clone()) {
-                                println!("Failed to send to server channel (n={}). Receiver probably isn't listening: {:?}", subscribers.len(), e);
-                            }
-                        }
-
-                        // Add to callbacks if there's a reply_id
-                        if let Some(reply_id) = reply_id {
-                            callbacks.lock().await.insert(reply_id, plugin_ref_id);
+                        // Send event to subscribers
+                        // Emit event to the channel for server to handle
+                        if let Err(e) = plugin_to_app_events_tx.try_send(event.clone()) {
+                            warn!("Failed to send to channel. Receiver probably isn't listening: {:?}", e);
                         }
                     }
                     Err(err) => {
                         // TODO: Better error handling
-                        println!("gRPC server error {err}");
+                        warn!("gRPC server error {err}");
                         break;
                     }
                 };
             }
 
-            server.remove_plugins(plugin_ids).await;
+            if let Err(e) = client_disconnect_tx.send(true).await {
+                warn!("Failed to send killed event {:?}", e);
+            }
         });
 
         // Write the same data that was received