yaak/crates-tauri/yaak-app/src/plugin_events.rs

use crate::error::Result;
use crate::http_request::send_http_request_with_context;
use crate::models_ext::BlobManagerExt;
use crate::models_ext::QueryManagerExt;
use crate::render::{render_grpc_request, render_http_request, render_json_value};
use crate::window::{CreateWindowConfig, create_window};
use crate::{
    call_frontend, cookie_jar_from_window, environment_from_window, get_window_from_plugin_context,
    workspace_from_window,
};
use chrono::Utc;
use cookie::Cookie;
use log::error;
use std::sync::Arc;
use tauri::{AppHandle, Emitter, Listener, Manager, Runtime};
use tauri_plugin_clipboard_manager::ClipboardExt;
use tauri_plugin_opener::OpenerExt;
use yaak::plugin_events::{
    GroupedPluginEvent, HostRequest, SharedPluginEventContext, handle_shared_plugin_event,
};
use yaak_crypto::manager::EncryptionManager;
use yaak_models::models::{HttpResponse, Plugin};
use yaak_models::queries::any_request::AnyRequest;
use yaak_models::util::UpdateSource;
use yaak_plugins::error::Error::PluginErr;
use yaak_plugins::events::{
    Color, EmptyPayload, ErrorResponse, GetCookieValueResponse, Icon, InternalEvent,
    InternalEventPayload, ListCookieNamesResponse, ListOpenWorkspacesResponse,
    RenderGrpcRequestResponse, RenderHttpRequestResponse, SendHttpRequestResponse,
    ShowToastRequest, TemplateRenderResponse, WindowInfoResponse, WindowNavigateEvent,
    WorkspaceInfo,
};
use yaak_plugins::manager::PluginManager;
use yaak_plugins::plugin_handle::PluginHandle;
use yaak_plugins::template_callback::PluginTemplateCallback;
use yaak_tauri_utils::window::WorkspaceWindowTrait;
use yaak_templates::{RenderErrorBehavior, RenderOptions};

pub(crate) async fn handle_plugin_event<R: Runtime>(
    app_handle: &AppHandle<R>,
    event: &InternalEvent,
    plugin_handle: &PluginHandle,
) -> Result<Option<InternalEventPayload>> {
    // log::debug!("Got event to app {event:?}");
    let plugin_context = event.context.to_owned();
    let plugin_name = plugin_handle.info().name;
    let fallback_workspace_id = plugin_context.workspace_id.clone().or_else(|| {
        plugin_context
            .label
            .as_ref()
            .and_then(|label| app_handle.get_webview_window(label))
            .and_then(|window| workspace_from_window(&window).map(|workspace| workspace.id))
    });

    match handle_shared_plugin_event(
        app_handle.db_manager().inner(),
        &event.payload,
        SharedPluginEventContext {
            plugin_name: &plugin_name,
            workspace_id: fallback_workspace_id.as_deref(),
        },
    ) {
        GroupedPluginEvent::Handled(payload) => Ok(payload),
        GroupedPluginEvent::ToHandle(host_request) => {
            handle_host_plugin_request(
                app_handle,
                event,
                plugin_handle,
                &plugin_context,
                host_request,
            )
            .await
        }
    }
}

async fn handle_host_plugin_request<R: Runtime>(
    app_handle: &AppHandle<R>,
    event: &InternalEvent,
    plugin_handle: &PluginHandle,
    plugin_context: &yaak_plugins::events::PluginContext,
    host_request: HostRequest<'_>,
) -> Result<Option<InternalEventPayload>> {
    match host_request {
        HostRequest::ErrorResponse(resp) => {
            error!("Plugin error: {}: {:?}", resp.error, resp);
            let toast_event = plugin_handle.build_event_to_send(
                plugin_context,
                &InternalEventPayload::ShowToastRequest(ShowToastRequest {
                    message: format!(
                        "Plugin error from {}: {}",
                        plugin_handle.info().name,
                        resp.error
                    ),
                    color: Some(Color::Danger),
                    timeout: Some(30000),
                    ..Default::default()
                }),
                None,
            );
            Box::pin(handle_plugin_event(app_handle, &toast_event, plugin_handle)).await
        }
        HostRequest::ReloadResponse(req) => {
            let plugins = app_handle.db().list_plugins()?;
            for plugin in plugins {
                if plugin.directory != plugin_handle.dir {
                    continue;
                }

                let new_plugin = Plugin { updated_at: Utc::now().naive_utc(), ..plugin };
                app_handle.db().upsert_plugin(&new_plugin, &UpdateSource::Plugin)?;
            }

            if !req.silent {
                let info = plugin_handle.info();
                let toast_event = plugin_handle.build_event_to_send(
                    plugin_context,
                    &InternalEventPayload::ShowToastRequest(ShowToastRequest {
                        message: format!("Reloaded plugin {}@{}", info.name, info.version),
                        icon: Some(Icon::Info),
                        timeout: Some(5000),
                        ..Default::default()
                    }),
                    None,
                );
                Box::pin(handle_plugin_event(app_handle, &toast_event, plugin_handle)).await
            } else {
                Ok(None)
            }
        }
        HostRequest::CopyText(req) => {
            app_handle.clipboard().write_text(req.text.as_str())?;
            Ok(Some(InternalEventPayload::CopyTextResponse(EmptyPayload {})))
        }
        HostRequest::ShowToast(req) => {
            match &plugin_context.label {
                Some(label) => app_handle.emit_to(label, "show_toast", req)?,
                None => app_handle.emit("show_toast", req)?,
            };
            Ok(Some(InternalEventPayload::ShowToastResponse(EmptyPayload {})))
        }
        HostRequest::PromptText(_) => {
            let window = get_window_from_plugin_context(app_handle, plugin_context)?;
            Ok(call_frontend(&window, event).await)
        }
        HostRequest::PromptForm(_) => {
            let window = get_window_from_plugin_context(app_handle, plugin_context)?;
            if event.reply_id.is_some() {
                window.emit_to(window.label(), "plugin_event", event.clone())?;
                Ok(None)
            } else {
                window.emit_to(window.label(), "plugin_event", event.clone()).unwrap();

                let event_id = event.id.clone();
                let plugin_handle = plugin_handle.clone();
                let plugin_context = plugin_context.clone();
                let window = window.clone();

                tauri::async_runtime::spawn(async move {
                    let (tx, mut rx) = tokio::sync::mpsc::channel::<InternalEvent>(128);

                    let listener_id = window.listen(event_id, move |ev: tauri::Event| {
                        let resp: InternalEvent = serde_json::from_str(ev.payload()).unwrap();
                        let _ = tx.try_send(resp);
                    });

                    while let Some(resp) = rx.recv().await {
                        let is_done = matches!(
                            &resp.payload,
                            InternalEventPayload::PromptFormResponse(r) if r.done.unwrap_or(false)
                        );

                        let event_to_send = plugin_handle.build_event_to_send(
                            &plugin_context,
                            &resp.payload,
                            Some(resp.reply_id.unwrap_or_default()),
                        );
                        if let Err(e) = plugin_handle.send(&event_to_send).await {
                            log::warn!("Failed to forward form response to plugin: {:?}", e);
                        }

                        if is_done {
                            break;
                        }
                    }

                    window.unlisten(listener_id);
                });

                Ok(None)
            }
        }
        HostRequest::RenderGrpcRequest(req) => {
            let window = get_window_from_plugin_context(app_handle, plugin_context)?;

            let workspace =
                workspace_from_window(&window).expect("Failed to get workspace_id from window URL");
            let environment_id = environment_from_window(&window).map(|e| e.id);
            let environment_chain = window.db().resolve_environments(
                &workspace.id,
                req.grpc_request.folder_id.as_deref(),
                environment_id.as_deref(),
            )?;
            let plugin_manager = Arc::new((*app_handle.state::<PluginManager>()).clone());
            let encryption_manager = Arc::new((*app_handle.state::<EncryptionManager>()).clone());
            let cb = PluginTemplateCallback::new(
                plugin_manager,
                encryption_manager,
                plugin_context,
                req.purpose.clone(),
            );
            let opt = RenderOptions { error_behavior: RenderErrorBehavior::Throw };
            let grpc_request =
                render_grpc_request(&req.grpc_request, environment_chain, &cb, &opt).await?;
            Ok(Some(InternalEventPayload::RenderGrpcRequestResponse(RenderGrpcRequestResponse {
                grpc_request,
            })))
        }
        HostRequest::RenderHttpRequest(req) => {
            let window = get_window_from_plugin_context(app_handle, plugin_context)?;

            let workspace =
                workspace_from_window(&window).expect("Failed to get workspace_id from window URL");
            let environment_id = environment_from_window(&window).map(|e| e.id);
            let environment_chain = window.db().resolve_environments(
                &workspace.id,
                req.http_request.folder_id.as_deref(),
                environment_id.as_deref(),
            )?;
            let plugin_manager = Arc::new((*app_handle.state::<PluginManager>()).clone());
            let encryption_manager = Arc::new((*app_handle.state::<EncryptionManager>()).clone());
            let cb = PluginTemplateCallback::new(
                plugin_manager,
                encryption_manager,
                plugin_context,
                req.purpose.clone(),
            );
            let opt = &RenderOptions { error_behavior: RenderErrorBehavior::Throw };
            let http_request =
                render_http_request(&req.http_request, environment_chain, &cb, opt).await?;
            Ok(Some(InternalEventPayload::RenderHttpRequestResponse(RenderHttpRequestResponse {
                http_request,
            })))
        }
        HostRequest::TemplateRender(req) => {
            let window = get_window_from_plugin_context(app_handle, plugin_context)?;

            let workspace =
                workspace_from_window(&window).expect("Failed to get workspace_id from window URL");
            let environment_id = environment_from_window(&window).map(|e| e.id);
            let folder_id = if let Some(id) = window.request_id() {
                match window.db().get_any_request(&id) {
                    Ok(AnyRequest::HttpRequest(r)) => r.folder_id,
                    Ok(AnyRequest::GrpcRequest(r)) => r.folder_id,
                    Ok(AnyRequest::WebsocketRequest(r)) => r.folder_id,
                    Err(_) => None,
                }
            } else {
                None
            };
            let environment_chain = window.db().resolve_environments(
                &workspace.id,
                folder_id.as_deref(),
                environment_id.as_deref(),
            )?;
            let plugin_manager = Arc::new((*app_handle.state::<PluginManager>()).clone());
            let encryption_manager = Arc::new((*app_handle.state::<EncryptionManager>()).clone());
            let cb = PluginTemplateCallback::new(
                plugin_manager,
                encryption_manager,
                plugin_context,
                req.purpose.clone(),
            );
            let opt = RenderOptions { error_behavior: RenderErrorBehavior::Throw };
            let data = render_json_value(req.data.clone(), environment_chain, &cb, &opt).await?;
            Ok(Some(InternalEventPayload::TemplateRenderResponse(TemplateRenderResponse { data })))
        }
        HostRequest::SendHttpRequest(req) => {
            let window = get_window_from_plugin_context(app_handle, plugin_context)?;
            let mut http_request = req.http_request.clone();
            let workspace =
                workspace_from_window(&window).expect("Failed to get workspace_id from window URL");
            let cookie_jar = cookie_jar_from_window(&window);
            let environment = environment_from_window(&window);

            if http_request.workspace_id.is_empty() {
                http_request.workspace_id = workspace.id;
            }

            let http_response = if http_request.id.is_empty() {
                HttpResponse::default()
            } else {
                let blobs = window.blob_manager();
                window.db().upsert_http_response(
                    &HttpResponse {
                        request_id: http_request.id.clone(),
                        workspace_id: http_request.workspace_id.clone(),
                        ..Default::default()
                    },
                    &UpdateSource::from_window_label(window.label()),
                    &blobs,
                )?
            };

            let http_response = send_http_request_with_context(
                &window,
                &http_request,
                &http_response,
                environment,
                cookie_jar,
                &mut tokio::sync::watch::channel(false).1,
                plugin_context,
            )
            .await?;

            Ok(Some(InternalEventPayload::SendHttpRequestResponse(SendHttpRequestResponse {
                http_response,
            })))
        }
        HostRequest::OpenWindow(req) => {
            let (navigation_tx, mut navigation_rx) = tokio::sync::mpsc::channel(128);
            let (close_tx, mut close_rx) = tokio::sync::mpsc::channel(128);
            let win_config = CreateWindowConfig {
                url: &req.url,
                label: &req.label,
                title: &req.title.clone().unwrap_or_default(),
                navigation_tx: Some(navigation_tx),
                close_tx: Some(close_tx),
                inner_size: req.size.clone().map(|s| (s.width, s.height)),
                data_dir_key: req.data_dir_key.clone(),
                ..Default::default()
            };
            if let Err(e) = create_window(app_handle, win_config) {
                let error_event = plugin_handle.build_event_to_send(
                    plugin_context,
                    &InternalEventPayload::ErrorResponse(ErrorResponse {
                        error: format!("Failed to create window: {:?}", e),
                    }),
                    None,
                );
                return Box::pin(handle_plugin_event(app_handle, &error_event, plugin_handle))
                    .await;
            }

            {
                let event_id = event.id.clone();
                let plugin_handle = plugin_handle.clone();
                let plugin_context = plugin_context.clone();
                tauri::async_runtime::spawn(async move {
                    while let Some(url) = navigation_rx.recv().await {
                        let url = url.to_string();
                        let event_to_send = plugin_handle.build_event_to_send(
                            &plugin_context,
                            &InternalEventPayload::WindowNavigateEvent(WindowNavigateEvent { url }),
                            Some(event_id.clone()),
                        );
                        plugin_handle.send(&event_to_send).await.unwrap();
                    }
                });
            }

            {
                let event_id = event.id.clone();
                let plugin_handle = plugin_handle.clone();
                let plugin_context = plugin_context.clone();
                tauri::async_runtime::spawn(async move {
                    while close_rx.recv().await.is_some() {
                        let event_to_send = plugin_handle.build_event_to_send(
                            &plugin_context,
                            &InternalEventPayload::WindowCloseEvent,
                            Some(event_id.clone()),
                        );
                        plugin_handle.send(&event_to_send).await.unwrap();
                    }
                });
            }

            Ok(None)
        }
        HostRequest::CloseWindow(req) => {
            if let Some(window) = app_handle.webview_windows().get(&req.label) {
                window.close()?;
            }
            Ok(None)
        }
        HostRequest::OpenExternalUrl(req) => {
            app_handle.opener().open_url(&req.url, None::<&str>)?;
            Ok(Some(InternalEventPayload::OpenExternalUrlResponse(EmptyPayload {})))
        }
        HostRequest::ListOpenWorkspaces(_) => {
            let mut workspaces = Vec::new();
            for (_, window) in app_handle.webview_windows() {
                if let Some(workspace) = workspace_from_window(&window) {
                    workspaces.push(WorkspaceInfo {
                        id: workspace.id.clone(),
                        name: workspace.name.clone(),
                        label: window.label().to_string(),
                    });
                }
            }
            Ok(Some(InternalEventPayload::ListOpenWorkspacesResponse(ListOpenWorkspacesResponse {
                workspaces,
            })))
        }
        HostRequest::ListCookieNames(_) => {
            let window = get_window_from_plugin_context(app_handle, plugin_context)?;
            let names = match cookie_jar_from_window(&window) {
                None => Vec::new(),
                Some(j) => j
                    .cookies
                    .into_iter()
                    .filter_map(|c| Cookie::parse(c.raw_cookie).ok().map(|c| c.name().to_string()))
                    .collect(),
            };
            Ok(Some(InternalEventPayload::ListCookieNamesResponse(ListCookieNamesResponse {
                names,
            })))
        }
        HostRequest::GetCookieValue(req) => {
            let window = get_window_from_plugin_context(app_handle, plugin_context)?;
            let value = match cookie_jar_from_window(&window) {
                None => None,
                Some(j) => j.cookies.into_iter().find_map(|c| match Cookie::parse(c.raw_cookie) {
                    Ok(c) if c.name().to_string().eq(&req.name) => {
                        Some(c.value_trimmed().to_string())
                    }
                    _ => None,
                }),
            };
            Ok(Some(InternalEventPayload::GetCookieValueResponse(GetCookieValueResponse { value })))
        }
        HostRequest::WindowInfo(req) => {
            let w = app_handle
                .get_webview_window(&req.label)
                .ok_or(PluginErr(format!("Failed to find window for {}", req.label)))?;

            let environment_id = environment_from_window(&w).map(|m| m.id);
            let workspace_id = workspace_from_window(&w).map(|m| m.id);
            let request_id =
                match app_handle.db().get_any_request(&w.request_id().unwrap_or_default()) {
                    Ok(AnyRequest::HttpRequest(r)) => Some(r.id),
                    Ok(AnyRequest::WebsocketRequest(r)) => Some(r.id),
                    Ok(AnyRequest::GrpcRequest(r)) => Some(r.id),
                    Err(_) => None,
                };

            Ok(Some(InternalEventPayload::WindowInfoResponse(WindowInfoResponse {
                label: w.label().to_string(),
                request_id,
                workspace_id,
                environment_id,
            })))
        }
        HostRequest::OtherRequest(req) => {
            Ok(Some(InternalEventPayload::ErrorResponse(ErrorResponse {
                error: format!(
                    "Unsupported plugin request in app host handler: {}",
                    req.type_name()
                ),
            })))
        }
    }
}