extensions: support agent method registration

Create a websocket endpoint that allows clients identified as
agents to register remote methods with Klipper.

Signed-off-by:  Eric Callahan <arksine.code@gmail.com>

moonraker/common.py

@@ -206,23 +206,36 @@ class BaseRemoteConnection(Subscribable):
             'method': "notify_status_update",
             'params': [status, eventtime]})
 
-    def call_method(
+    def call_method_with_response(
         self,
         method: str,
-        params: Optional[Union[List, Dict[str, Any]]] = None
+        params: Optional[Union[List, Dict[str, Any]]] = None,
     ) -> Awaitable:
         fut = self.eventloop.create_future()
-        msg = {
+        msg: Dict[str, Any] = {
             'jsonrpc': "2.0",
             'method': method,
             'id': id(fut)
         }
-        if params is not None:
+        if params:
             msg["params"] = params
         self.pending_responses[id(fut)] = fut
         self.queue_message(msg)
         return fut
 
+    def call_method(
+        self,
+        method: str,
+        params: Optional[Union[List, Dict[str, Any]]] = None
+    ) -> None:
+        msg: Dict[str, Any] = {
+            "jsonrpc": "2.0",
+            "method": method
+        }
+        if params:
+            msg["params"] = params
+        self.queue_message(msg)
+
     def send_notification(self, name: str, data: List) -> None:
         self.wsm.notify_clients(name, data, [self._uid])
 

moonraker/components/extensions.py

@@ -32,7 +32,13 @@ class ExtensionManager:
     def __init__(self, config: ConfigHelper) -> None:
         self.server = config.get_server()
         self.agents: Dict[str, BaseRemoteConnection] = {}
+        self.agent_methods: Dict[int, List[str]] = {}
         self.uds_server: Optional[asyncio.AbstractServer] = None
+        self.server.register_endpoint(
+            "/connection/register_remote_method", ["POST"],
+            self._register_agent_method,
+            transports=["websocket"]
+        )
         self.server.register_endpoint(
             "/connection/send_event", ["POST"], self._handle_agent_event,
             transports=["websocket"]
@@ -66,6 +72,10 @@ class ExtensionManager:
     def remove_agent(self, connection: BaseRemoteConnection) -> None:
         name = connection.client_data["name"]
         if name in self.agents:
+            klippy: Klippy = self.server.lookup_component("klippy_connection")
+            registered_methods = self.agent_methods.pop(connection.uid, [])
+            for method in registered_methods:
+                klippy.unregister_method(method)
             del self.agents[name]
             evt: Dict[str, Any] = {"agent": name, "event": "disconnected"}
             connection.send_notification("agent_event", [evt])
@@ -90,6 +100,16 @@ class ExtensionManager:
         conn.send_notification("agent_event", [evt])
         return "ok"
 
+    async def _register_agent_method(self, web_request: WebRequest) -> str:
+        conn = web_request.get_client_connection()
+        if conn is None:
+            raise self.server.error("No connection detected")
+        method_name = web_request.get_str("method_name")
+        klippy: Klippy = self.server.lookup_component("klippy_connection")
+        klippy.register_method_from_agent(conn, method_name)
+        self.agent_methods.setdefault(conn.uid, []).append(method_name)
+        return "ok"
+
     async def _handle_list_extensions(
         self, web_request: WebRequest
     ) -> Dict[str, List[Dict[str, Any]]]:
@@ -109,7 +129,7 @@ class ExtensionManager:
         if agent not in self.agents:
             raise self.server.error(f"Agent {agent} not connected")
         conn = self.agents[agent]
-        return await conn.call_method(method, args)
+        return await conn.call_method_with_response(method, args)
 
     async def start_unix_server(self) -> None:
         sockfile: str = self.server.get_app_args()["unix_socket_path"]

moonraker/klippy_connection.py

@@ -31,7 +31,7 @@ from typing import (
 if TYPE_CHECKING:
     from .server import Server
     from .app import MoonrakerApp
-    from .common import WebRequest, Subscribable
+    from .common import WebRequest, Subscribable, BaseRemoteConnection
     from .confighelper import ConfigHelper
     from .components.klippy_apis import KlippyAPI
     from .components.file_manager.file_manager import FileManager
@@ -71,6 +71,7 @@ class KlippyConnection:
         self._klippy_identified: bool = False
         self._klippy_initializing: bool = False
         self._klippy_started: bool = False
+        self._methods_registered: bool = False
         self._klipper_version: str = ""
         self._missing_reqs: Set[str] = set()
         self._peer_cred: Dict[str, int] = {}
@@ -223,6 +224,34 @@ class KlippyConnection:
             # These methods need to be registered with Klippy
             self.klippy_reg_methods.append(method_name)
 
+    def register_method_from_agent(
+        self, connection: BaseRemoteConnection, method_name: str
+    ) -> Optional[Awaitable]:
+        if connection.client_data["type"] != "agent":
+            raise self.server.error(
+                "Only connections of the 'agent' type can register methods"
+            )
+        if method_name in self.remote_methods:
+            raise self.server.error(
+                f"Remote method ({method_name}) already registered"
+            )
+
+        def _on_agent_method_received(**kwargs) -> None:
+            connection.call_method(method_name, kwargs)
+        self.remote_methods[method_name] = _on_agent_method_received
+        self.klippy_reg_methods.append(method_name)
+        if self._methods_registered and self._state != "disconnected":
+            coro = self.klippy_apis.register_method(method_name)
+            return self.event_loop.create_task(coro)
+        return None
+
+    def unregister_method(self, method_name: str):
+        self.remote_methods.pop(method_name, None)
+        try:
+            self.klippy_reg_methods.remove(method_name)
+        except ValueError:
+            pass
+
     def connect(self) -> Awaitable[bool]:
         if (
             self.is_connected() or
@@ -298,6 +327,7 @@ class KlippyConnection:
         self._klippy_identified = False
         self._klippy_started = False
         self._klippy_initializing = True
+        self._methods_registered = False
         self._missing_reqs.clear()
         self.init_attempts = 0
         self._state = "startup"
@@ -393,6 +423,7 @@ class KlippyConnection:
                     except ServerError:
                         logging.exception(
                             f"Unable to register method '{method}'")
+                self._methods_registered = True
                 if self._state == "ready":
                     logging.info("Klippy ready")
                     await self.server.send_event("server:klippy_ready")
@@ -669,6 +700,7 @@ class KlippyConnection:
         self._klippy_identified = False
         self._klippy_initializing = False
         self._klippy_started = False
+        self._methods_registered = False
         self._state = "disconnected"
         self._state_message = "Klippy Disconnected"
         for request in self.pending_requests.values():

moonraker/server.py

@@ -48,7 +48,7 @@ if TYPE_CHECKING:
     FlexCallback = Callable[..., Optional[Coroutine]]
     _T = TypeVar("_T", Sentinel, Any)
 
-API_VERSION = (1, 3, 0)
+API_VERSION = (1, 4, 0)
 CORE_COMPONENTS = [
     'dbus_manager', 'database', 'file_manager', 'klippy_apis',
     'machine', 'data_store', 'shell_command', 'proc_stats',