Add message emit mechanism

2 files changed, 39 insertions, 24 deletions

diff --git a/backend/src/plugin/plugin.py b/backend/src/plugin/plugin.py
index d5f79d42..a1b43ba9 100644
--- a/backend/src/plugin/plugin.py
+++ b/backend/src/plugin/plugin.py
@@ -2,12 +2,13 @@ from asyncio import Task, create_task
 from json import dumps, load, loads
 from logging import getLogger
 from os import path
+from multiprocessing import Process
 
 from .sandboxed_plugin import SandboxedPlugin
 from .method_call_request import MethodCallRequest
 from ..localplatform.localsocket import LocalSocket
 
-from typing import Any, Dict
+from typing import Any, Callable, Coroutine, Dict
 
 class PluginWrapper:
     def __init__(self, file: str, plugin_directory: str, plugin_path: str) -> None:
@@ -28,40 +29,52 @@ class PluginWrapper:
         self.passive = not path.isfile(self.file)
 
         self.log = getLogger("plugin")
-        self.method_call_requests: Dict[str, MethodCallRequest] = {}
+
         self.sandboxed_plugin = SandboxedPlugin(self.name, self.passive, self.flags, self.file, self.plugin_directory, self.plugin_path, self.version, self.author)
-        #TODO: Maybe somehow make LocalSocket not require on_new_message to make this more clear
-        self.socket = LocalSocket(self.sandboxed_plugin.on_new_message)
-        self.listener_task: Task[Any]
+        #TODO: Maybe make LocalSocket not require on_new_message to make this cleaner
+        self._socket = LocalSocket(self.sandboxed_plugin.on_new_message)
+        self._listener_task: Task[Any]
+        self._method_call_requests: Dict[str, MethodCallRequest] = {}
+
+        self.emitted_message_callback: Callable[[Dict[Any, Any]], Coroutine[Any, Any, Any]]
 
     def __str__(self) -> str:
         return self.name
     
     async def _response_listener(self):
         while True:
-            line = await self.socket.read_single_line()
+            line = await self._socket.read_single_line()
             if line != None:
                 res = loads(line)
-                self.method_call_requests.pop(res["id"]).set_result(res)
+                if res["id"] == 0:
+                    create_task(self.emitted_message_callback(res["payload"]))
+                    return
+                self._method_call_requests.pop(res["id"]).set_result(res)
+
+    async def set_emitted_message_callback(self, callback: Callable[[Dict[Any, Any]], Coroutine[Any, Any, Any]]):
+        self.emitted_message_callback = callback
 
     async def execute_method(self, method_name: str, kwargs: Dict[Any, Any]):
         if self.passive:
             raise RuntimeError("This plugin is passive (aka does not implement main.py)")
         
         request = MethodCallRequest()
-        await self.socket.get_socket_connection()
-        await self.socket.write_single_line(dumps({ "method": method_name, "args": kwargs, "id": request.id }, ensure_ascii=False))
-        self.method_call_requests[request.id] = request
+        await self._socket.get_socket_connection()
+        await self._socket.write_single_line(dumps({ "method": method_name, "args": kwargs, "id": request.id }, ensure_ascii=False))
+        self._method_call_requests[request.id] = request
 
         return await request.wait_for_result()
     
     async def start(self):
-        self.sandboxed_plugin.start(self.socket)
+        if self.passive:
+            return self
+        Process(target=self.sandboxed_plugin.initialize, args=[self._socket]).start()
         self.listener_task = create_task(self._response_listener())
+        return self
 
     async def stop(self):
-        self.listener_task.cancel()
+        self._listener_task.cancel()
         async def _(self: PluginWrapper):
-            await self.socket.write_single_line(dumps({ "stop": True }, ensure_ascii=False))
-            await self.socket.close_socket_connection()
+            await self._socket.write_single_line(dumps({ "stop": True }, ensure_ascii=False))
+            await self._socket.close_socket_connection()
         create_task(_(self))
\ No newline at end of file
diff --git a/backend/src/plugin/sandboxed_plugin.py b/backend/src/plugin/sandboxed_plugin.py
index fea02ac2..e4b5f7df 100644
--- a/backend/src/plugin/sandboxed_plugin.py
+++ b/backend/src/plugin/sandboxed_plugin.py
@@ -3,7 +3,6 @@ from signal import SIGINT, signal
 from importlib.util import module_from_spec, spec_from_file_location
 from json import dumps, loads
 from logging import getLogger
-import multiprocessing
 from sys import exit, path as syspath
 from traceback import format_exc
 from asyncio import (get_event_loop, new_event_loop,
@@ -15,7 +14,7 @@ from ..localplatform.localplatform import setgid, setuid, get_username, get_home
 from ..customtypes import UserType
 from .. import helpers
 
-from typing import List
+from typing import Any, Dict, List
 
 class SandboxedPlugin:
     def __init__(self,
@@ -38,7 +37,9 @@ class SandboxedPlugin:
 
         self.log = getLogger("plugin")
 
-    def _init(self, socket: LocalSocket):
+    def initialize(self, socket: LocalSocket):
+        self._socket = socket
+
         try:
             signal(SIGINT, lambda s, f: exit(0))
 
@@ -84,6 +85,9 @@ class SandboxedPlugin:
             spec.loader.exec_module(module)
             self.Plugin = module.Plugin
 
+            setattr(self.Plugin, "emit_message", self.emit_message)
+            #TODO: Find how to put emit_message on global namespace so it doesn't pollute Plugin
+
             if hasattr(self.Plugin, "_migration"):
                 get_event_loop().run_until_complete(self.Plugin._migration(self.Plugin))
             if hasattr(self.Plugin, "_main"):
@@ -118,7 +122,6 @@ class SandboxedPlugin:
             get_event_loop().close()
             raise Exception("Closing message listener")
 
-        # TODO there is definitely a better way to type this
         d: SocketResponseDict = {"res": None, "success": True, "id": data["id"]}
         try:
             d["res"] = await getattr(self.Plugin, data["method"])(self.Plugin, **data["args"])
@@ -128,9 +131,8 @@ class SandboxedPlugin:
         finally:
             return dumps(d, ensure_ascii=False)
         
-
-    def start(self, socket: LocalSocket):
-        if self.passive:
-            return self
-        multiprocessing.Process(target=self._init, args=[socket]).start()
-        return self
\ No newline at end of file
+    async def emit_message(self, message: Dict[Any, Any]):
+        await self._socket.write_single_line(dumps({
+            "id": 0,
+            "payload": message
+        }))
\ No newline at end of file