Enhance entity command handler with client connection

zehnm · zehnm · commit 41aa9c449ce9 · 2025-12-16T12:23:14.000+01:00
Pass the WS connection as an optional command parameter to the client.
This allows the client to send directed WS events back. Otherwise, the
only option is to use event broadcasts to all connected clients.
Implementation is backward compatible to existing clients without the
new websocket parameter in the CommandHandler callback, or extending a
concrete Entity class.

Import cleanup to avoid circular imports through the ucapi definition.
diff --git a/.idea/misc.xml b/.idea/misc.xml
diff --git a/examples/hello_integration.py b/examples/hello_integration.py
@@ -11,7 +11,7 @@
 
 
 async def cmd_handler(
-    entity: ucapi.Button, cmd_id: str, _params: dict[str, Any] | None
+    entity: ucapi.Button, cmd_id: str, _params: dict[str, Any] | None, websocket: Any
 ) -> ucapi.StatusCodes:
     """
     Push button command handler.
@@ -21,6 +21,7 @@ async def cmd_handler(
     :param entity: button entity
     :param cmd_id: command
     :param _params: optional command parameters
+    :param websocket: optional client connection for sending directed events
     :return: status of the command
     """
     print(f"Got {entity.id} command request: {cmd_id}")
diff --git a/examples/remote.py b/examples/remote.py
@@ -46,7 +46,7 @@
 
 
 async def cmd_handler(
-    entity: ucapi.Remote, cmd_id: str, params: dict[str, Any] | None
+    entity: ucapi.Remote, cmd_id: str, params: dict[str, Any] | None, websocket: Any
 ) -> ucapi.StatusCodes:
     """
     Remote command handler.
@@ -56,6 +56,7 @@ async def cmd_handler(
     :param entity: remote entity
     :param cmd_id: command
     :param params: optional command parameters
+    :param websocket: optional client connection for sending directed events
     :return: status of the command
     """
     print(f"Got {entity.id} command request: {cmd_id}")
diff --git a/examples/setup_flow.py b/examples/setup_flow.py
@@ -157,7 +157,7 @@ async def handle_user_data_response(msg: ucapi.UserDataResponse) -> ucapi.SetupA
 
 
 async def cmd_handler(
-    entity: ucapi.Button, cmd_id: str, _params: dict[str, Any] | None
+    entity: ucapi.Button, cmd_id: str, _params: dict[str, Any] | None, websocket: Any
 ) -> ucapi.StatusCodes:
     """
     Push button command handler.
@@ -167,6 +167,7 @@ async def cmd_handler(
     :param entity: button entity
     :param cmd_id: command
     :param _params: optional command parameters
+    :param websocket: optional client connection for sending directed events
     :return: status of the command
     """
     print(f"Got {entity.id} command request: {cmd_id}")
diff --git a/examples/voice.py b/examples/voice.py
@@ -45,7 +45,10 @@ async def on_subscribe_entities(entity_ids: list[str]) -> None:
 
 
 async def on_voice_cmd(
-    entity: ucapi.VoiceAssistant, cmd_id: str, params: dict[str, Any] | None
+    entity: ucapi.VoiceAssistant,
+    cmd_id: str,
+    params: dict[str, Any] | None,
+    websocket: Any,
 ) -> ucapi.StatusCodes:
     """
     Voice assistant command handler.
@@ -55,6 +58,7 @@ async def on_voice_cmd(
     :param entity: voice assistant entity
     :param cmd_id: command
     :param params: optional command parameters
+    :param websocket: optional client connection for sending directed events
     :return: status of the command
     """
     # HACK until core is fixed
diff --git a/ucapi/api.py b/ucapi/api.py
@@ -589,10 +589,10 @@ async def broadcast_assistant_event(self, event: uc.AssistantEvent) -> None:
             uc.EventCategory.ENTITY,
         )
 
-    async def send_assistant_event(self, websocket, event: uc.AssistantEvent) -> None:
-        """Send an assistant event to the given WebSocket client."""
+    async def send_assistant_event(self, client, event: uc.AssistantEvent) -> None:
+        """Send an assistant event to the given client connection."""
         await self._send_ws_event(
-            websocket,
+            client,
             uc.WsMsgEvents.ASSISTANT_EVENT,
             asdict(event),
             uc.EventCategory.ENTITY,
@@ -888,9 +888,22 @@ async def _entity_command(
             # Start timeout immediately on session creation
             self._schedule_voice_timeout(key)
 
-        result = await entity.command(
-            cmd_id, msg_data["params"] if "params" in msg_data else None
-        )
+        try:
+            result = await entity.command(
+                cmd_id,
+                msg_data["params"] if "params" in msg_data else None,
+                websocket=websocket,
+            )
+        except TypeError:
+            # Entity command method likely doesn't accept 'websocket' kwarg -> legacy handler signature.
+            _LOG.warning(
+                "Old Entity.command signature detected for %s, trying old signature. Please update the command signature.",
+                entity.id,
+            )
+            result = await entity.command(
+                cmd_id, msg_data["params"] if "params" in msg_data else None
+            )
+
         await self.acknowledge_command(websocket, req_id, result)
 
     async def _setup_driver(
diff --git a/ucapi/api_definitions.py b/ucapi/api_definitions.py
@@ -205,8 +205,19 @@ class SetupComplete(SetupAction):
 
 
 CommandHandler: TypeAlias = Callable[
-    [Any, str, dict[str, Any] | None], Awaitable[StatusCodes]
+    [Any, str, dict[str, Any] | None, Any | None], Awaitable[StatusCodes]
 ]
+"""Entity command handler signature.
+
+Parameters:
+
+- entity: entity instance
+- cmd_id: command identifier
+- params: optional command parameters
+- websocket: optional client connection for sending directed events
+
+Returns: status code
+"""
 
 
 SetupHandler: TypeAlias = Callable[[SetupDriver], Awaitable[SetupAction]]
diff --git a/ucapi/button.py b/ucapi/button.py
@@ -7,8 +7,8 @@
 
 from enum import Enum
 
-from ucapi.api_definitions import CommandHandler
-from ucapi.entity import Entity, EntityTypes
+from .api_definitions import CommandHandler
+from .entity import Entity, EntityTypes
 
 
 class States(str, Enum):
diff --git a/ucapi/climate.py b/ucapi/climate.py
@@ -8,8 +8,8 @@
 from enum import Enum
 from typing import Any
 
-from ucapi.api_definitions import CommandHandler
-from ucapi.entity import Entity, EntityTypes
+from .api_definitions import CommandHandler
+from .entity import Entity, EntityTypes
 
 
 class States(str, Enum):
@@ -111,8 +111,8 @@ def __init__(
             EntityTypes.CLIMATE,
             features,
             attributes,
-            device_class,
-            options,
-            area,
-            cmd_handler,
+            device_class=device_class,
+            options=options,
+            area=area,
+            cmd_handler=cmd_handler,
         )
diff --git a/ucapi/cover.py b/ucapi/cover.py
@@ -8,8 +8,8 @@
 from enum import Enum
 from typing import Any
 
-from ucapi.api_definitions import CommandHandler
-from ucapi.entity import Entity, EntityTypes
+from .api_definitions import CommandHandler
+from .entity import Entity, EntityTypes
 
 
 class States(str, Enum):
@@ -110,8 +110,8 @@ def __init__(
             EntityTypes.COVER,
             features,
             attributes,
-            device_class,
-            options,
-            area,
-            cmd_handler,
+            device_class=device_class,
+            options=options,
+            area=area,
+            cmd_handler=cmd_handler,
         )
diff --git a/ucapi/entities.py b/ucapi/entities.py
@@ -11,8 +11,8 @@
 
 from pyee.asyncio import AsyncIOEventEmitter
 
-from ucapi.api_definitions import Events
-from ucapi.entity import Entity
+from .api_definitions import Events
+from .entity import Entity
 
 _LOG = logging.getLogger(__name__)
 _LOG.setLevel(logging.DEBUG)
diff --git a/ucapi/entity.py b/ucapi/entity.py
@@ -5,11 +5,12 @@
 :license: MPL-2.0, see LICENSE for more details.
 """
 
+import inspect
 import logging
 from enum import Enum
 from typing import Any
 
-from ucapi.api_definitions import CommandHandler, StatusCodes
+from .api_definitions import CommandHandler, StatusCodes
 
 _LOG = logging.getLogger(__name__)
 _LOG.setLevel(logging.DEBUG)
@@ -36,6 +37,9 @@ class Entity:
 
     See https://github.com/unfoldedcircle/core-api/blob/main/doc/entities/
     for more information.
+
+    A client can either use a ``CommandHandler`` callback or extend a concrete Entity
+    class and override the ``command()`` method.
     """
 
     # pylint: disable=R0917
@@ -46,6 +50,7 @@ def __init__(
         entity_type: EntityTypes,
         features: list[str],
         attributes: dict[str, Any],
+        *,
         device_class: str | None = None,
         options: dict[str, Any] | None = None,
         area: str | None = None,
@@ -62,6 +67,7 @@ def __init__(
         :param device_class: entity device class
         :param options: entity options
         :param area: optional area name
+        :param cmd_handler: optional handler for entity commands
         """
         self.id = identifier
         self.name = {"en": name} if isinstance(name, str) else name
@@ -73,28 +79,60 @@ def __init__(
         self.options = options
         self.area = area
         self._cmd_handler = cmd_handler
+        self._ext_cmd_handler = False
+
+        # check if we got an extended command handler with the websocket kwarg
+        if cmd_handler:
+            try:
+                sig = inspect.signature(cmd_handler)
+                params_map = sig.parameters
+                self._ext_cmd_handler = "websocket" in params_map or any(
+                    p.kind is inspect.Parameter.VAR_KEYWORD for p in params_map.values()
+                )
+            except (TypeError, ValueError):
+                # Some callables (builtins, C-extensions, functools.partial edge cases)
+                # may not have an introspectable signature. Fall back to safe call.
+                _LOG.warning(
+                    "Old CommandHandler signature detected for %s, using old signature. Please update the command handler.",
+                    self.id,
+                )
 
         _LOG.debug("Created %s entity: %s", self.entity_type.value, self.id)
 
     async def command(
-        self, cmd_id: str, params: dict[str, Any] | None = None
+        self,
+        cmd_id: str,
+        params: dict[str, Any] | None = None,
+        *,
+        websocket: Any,
     ) -> StatusCodes:
         """
         Execute entity command with the installed command handler.
 
+        Backward compatible:
+        - Existing handlers usually accept (entity, cmd_id, params)
+        - New handlers may optionally accept websocket as kw-only / kwarg
+
         Returns NOT_IMPLEMENTED if no command handler is installed.
 
         :param cmd_id: the command
         :param params: optional command parameters
+        :param websocket: optional websocket connection. Allows for directed event
+                          callbacks instead of broadcasts.
         :return: command status code to acknowledge to UCR2
         """
-        if self._cmd_handler:
-            return await self._cmd_handler(self, cmd_id, params)
-
-        _LOG.warning(
-            "No command handler for %s: cannot execute command '%s' %s",
-            self.id,
-            cmd_id,
-            params if params else "",
-        )
-        return StatusCodes.NOT_IMPLEMENTED
+        handler = self._cmd_handler
+        if not handler:
+            _LOG.warning(
+                "No command handler for %s: cannot execute command '%s' %s",
+                self.id,
+                cmd_id,
+                params if params else "",
+            )
+            return StatusCodes.NOT_IMPLEMENTED
+
+        # Decide if we can pass websocket as a keyword argument
+        if self._ext_cmd_handler:
+            return await handler(self, cmd_id, params, websocket=websocket)
+
+        return await handler(self, cmd_id, params)
diff --git a/ucapi/light.py b/ucapi/light.py
@@ -8,8 +8,8 @@
 from enum import Enum
 from typing import Any
 
-from ucapi.api_definitions import CommandHandler
-from ucapi.entity import Entity, EntityTypes
+from .api_definitions import CommandHandler
+from .entity import Entity, EntityTypes
 
 
 class States(str, Enum):
@@ -97,8 +97,8 @@ def __init__(
             EntityTypes.LIGHT,
             features,
             attributes,
-            device_class,
-            options,
-            area,
-            cmd_handler,
+            device_class=device_class,
+            options=options,
+            area=area,
+            cmd_handler=cmd_handler,
         )
diff --git a/ucapi/media_player.py b/ucapi/media_player.py
@@ -8,8 +8,8 @@
 from enum import Enum
 from typing import Any
 
-from ucapi.api_definitions import CommandHandler
-from ucapi.entity import Entity, EntityTypes
+from .api_definitions import CommandHandler
+from .entity import Entity, EntityTypes
 
 
 class States(str, Enum):
@@ -226,8 +226,8 @@ def __init__(
             EntityTypes.MEDIA_PLAYER,
             features,
             attributes,
-            device_class,
-            options,
-            area,
-            cmd_handler,
+            device_class=device_class,
+            options=options,
+            area=area,
+            cmd_handler=cmd_handler,
         )
diff --git a/ucapi/remote.py b/ucapi/remote.py
@@ -9,9 +9,9 @@
 from enum import Enum
 from typing import Any
 
-from ucapi.api_definitions import CommandHandler
-from ucapi.entity import Entity, EntityTypes
-from ucapi.ui import DeviceButtonMapping, EntityCommand, UiPage
+from .api_definitions import CommandHandler
+from .entity import Entity, EntityTypes
+from .ui import DeviceButtonMapping, EntityCommand, UiPage
 
 
 class States(str, Enum):
@@ -135,7 +135,7 @@ def __init__(
         cmd_handler: CommandHandler = None,
     ):
         """
-        Create remote entity instance.
+        Create remote-entity instance.
 
         :param identifier: entity identifier
         :param name: friendly name
diff --git a/ucapi/sensor.py b/ucapi/sensor.py
diff --git a/ucapi/switch.py b/ucapi/switch.py
diff --git a/ucapi/voice_assistant.py b/ucapi/voice_assistant.py
