Merge branch 'dev' of https://git.science.uu.nl/ics/sp/2025/n25b/pepperplus-cb into feat/recieve-programs-ui

src/control_backend/agents/ri_command_agent.py

@@ -72,7 +72,7 @@ class RICommandAgent(BaseAgent):
 
         # To the robot
         self.pubsocket = context.socket(zmq.PUB)
-        if self.bind:
+        if self.bind:  # TODO: Should this ever be the case?
             self.pubsocket.bind(self.address)
         else:
             self.pubsocket.connect(self.address)

src/control_backend/agents/ri_communication_agent.py

@@ -1,6 +1,7 @@
 import asyncio
+import json
 
-import zmq
+import zmq.asyncio
 from spade.behaviour import CyclicBehaviour
 from zmq.asyncio import Context
 
@@ -14,6 +15,7 @@ class RICommunicationAgent(BaseAgent):
     req_socket: zmq.Socket
     _address = ""
     _bind = True
+    connected = False
 
     def __init__(
         self,
@@ -27,6 +29,8 @@ class RICommunicationAgent(BaseAgent):
         super().__init__(jid, password, port, verify_security)
         self._address = address
         self._bind = bind
+        self._req_socket: zmq.asyncio.Socket | None = None
+        self.pub_socket: zmq.asyncio.Socket | None = None
 
     class ListenBehaviour(CyclicBehaviour):
         async def run(self):
@@ -35,59 +39,128 @@ class RICommunicationAgent(BaseAgent):
             """
             assert self.agent is not None
 
+            if not self.agent.connected:
+                await asyncio.sleep(1)
+                return
+
             # We need to listen and sent pings.
             message = {"endpoint": "ping", "data": {"id": "e.g. some reference id"}}
-            await self.agent.req_socket.send_json(message)
-
-            # Wait up to three seconds for a reply:)
+            seconds_to_wait_total = 1.0
             try:
-                message = await asyncio.wait_for(self.agent.req_socket.recv_json(), timeout=3.0)
-
-            # We didnt get a reply :(
+                await asyncio.wait_for(
+                    self.agent._req_socket.send_json(message), timeout=seconds_to_wait_total / 2
+                )
             except TimeoutError:
-                self.agent.logger.info("No ping retrieved in 3 seconds, killing myself.")
-                self.kill()
+                self.agent.logger.debug(
+                    "Waited too long to send message - "
+                    "we probably dont have any receivers... but let's check!"
+                )
 
-            self.agent.logger.debug('Received message "%s"', message)
+            # Wait up to {seconds_to_wait_total/2} seconds for a reply
+            try:
+                message = await asyncio.wait_for(
+                    self.agent._req_socket.recv_json(), timeout=seconds_to_wait_total / 2
+                )
+
+            # We didnt get a reply
+            except TimeoutError:
+                self.agent.logger.info(
+                    f"No ping retrieved in {seconds_to_wait_total} seconds, "
+                    "sending UI disconnection event and attempting to restart."
+                )
+
+                # Make sure we dont retry receiving messages untill we're setup.
+                self.agent.connected = False
+                self.agent.remove_behaviour(self)
+
+                # Tell UI we're disconnected.
+                topic = b"ping"
+                data = json.dumps(False).encode()
+                if self.agent.pub_socket is None:
+                    self.agent.logger.warning(
+                        "Communication agent pub socket not correctly initialized."
+                    )
+                else:
+                    try:
+                        await asyncio.wait_for(
+                            self.agent.pub_socket.send_multipart([topic, data]), 5
+                        )
+                    except TimeoutError:
+                        self.agent.logger.warning(
+                            "Initial connection ping for router timed"
+                            " out in ri_communication_agent."
+                        )
+
+                # Try to reboot.
+                self.agent.logger.debug("Restarting communication agent.")
+                await self.agent.setup()
+
+            self.agent.logger.debug(f'Received message "{message}" from RI.')
             if "endpoint" not in message:
-                self.agent.logger.error("No received endpoint in message, excepted ping endpoint.")
+                self.agent.logger.warning(
+                    "No received endpoint in message, expected ping endpoint."
+                )
                 return
 
             # See what endpoint we received
             match message["endpoint"]:
                 case "ping":
+                    topic = b"ping"
+                    data = json.dumps(True).encode()
+                    if self.agent.pub_socket is not None:
+                        await self.agent.pub_socket.send_multipart([topic, data])
                     await asyncio.sleep(1)
                 case _:
-                    self.agent.logger.info(
+                    self.agent.logger.debug(
                         "Received message with topic different than ping, while ping expected."
                     )
 
-    async def setup(self, max_retries: int = 5):
+    async def setup_sockets(self, force=False):
+        """
+        Sets up request socket for communication agent.
+        """
+        # Bind request socket
+        if self._req_socket is None or force:
+            self._req_socket = Context.instance().socket(zmq.REQ)
+            if self._bind:
+                self._req_socket.bind(self._address)
+            else:
+                self._req_socket.connect(self._address)
+
+        if self.pub_socket is None or force:
+            self.pub_socket = Context.instance().socket(zmq.PUB)
+            self.pub_socket.connect(settings.zmq_settings.internal_pub_address)
+
+    async def setup(self, max_retries: int = 100):
         """
         Try to setup the communication agent, we have 5 retries in case we dont have a response yet.
         """
         self.logger.info("Setting up %s", self.jid)
-        retries = 0
 
+        # Bind request socket
+        await self.setup_sockets()
+
+        retries = 0
         # Let's try a certain amount of times before failing connection
         while retries < max_retries:
-            # Bind request socket
-            self.req_socket = Context.instance().socket(zmq.REQ)
-            if self._bind:
-                self.req_socket.bind(self._address)
-            else:
-                self.req_socket.connect(self._address)
+            # Make sure the socket is properly setup.
+            if self._req_socket is None:
+                continue
 
-            # Send our message and receive one back:)
-            message = {"endpoint": "negotiate/ports", "data": None}
-            await self.req_socket.send_json(message)
+            # Send our message and receive one back
+            message = {"endpoint": "negotiate/ports", "data": {}}
+            await self._req_socket.send_json(message)
 
+            retry_frequency = 1.0
             try:
-                received_message = await asyncio.wait_for(self.req_socket.recv_json(), timeout=20.0)
+                received_message = await asyncio.wait_for(
+                    self._req_socket.recv_json(), timeout=retry_frequency
+                )
 
             except TimeoutError:
                 self.logger.warning(
-                    "No connection established in 20 seconds (attempt %d/%d)",
+                    "No connection established in %d seconds (attempt %d/%d)",
+                    retries * retry_frequency,
                     retries + 1,
                     max_retries,
                 )
@@ -95,21 +168,21 @@ class RICommunicationAgent(BaseAgent):
                 continue
 
             except Exception as e:
-                self.logger.error("Unexpected error during negotiation: %s", e)
+                self.logger.warning("Unexpected error during negotiation: %s", e)
                 retries += 1
                 continue
 
             # Validate endpoint
             endpoint = received_message.get("endpoint")
             if endpoint != "negotiate/ports":
-                # TODO: Should this send a message back?
-                self.logger.error(
+                self.logger.warning(
                     "Invalid endpoint '%s' received (attempt %d/%d)",
                     endpoint,
                     retries + 1,
                     max_retries,
                 )
                 retries += 1
+                await asyncio.sleep(1)
                 continue
 
             # At this point, we have a valid response
@@ -128,9 +201,9 @@ class RICommunicationAgent(BaseAgent):
                         case "main":
                             if addr != self._address:
                                 if not bind:
-                                    self.req_socket.connect(addr)
+                                    self._req_socket.connect(addr)
                                 else:
-                                    self.req_socket.bind(addr)
+                                    self._req_socket.bind(addr)
                         case "actuation":
                             ri_commands_agent = RICommandAgent(
                                 settings.agent_settings.ri_command_agent_name
@@ -145,18 +218,35 @@ class RICommunicationAgent(BaseAgent):
                             self.logger.warning("Unhandled negotiation id: %s", id)
 
             except Exception as e:
-                self.logger.error("Error unpacking negotiation data: %s", e)
+                self.logger.warning("Error unpacking negotiation data: %s", e)
                 retries += 1
+                await asyncio.sleep(1)
                 continue
 
             # setup succeeded
             break
 
         else:
-            self.logger.error("Failed to set up RICommunicationAgent after %d retries", max_retries)
+            self.logger.warning("Failed to set up %s after %d retries", self.name, max_retries)
             return
 
         # Set up ping behaviour
         listen_behaviour = self.ListenBehaviour()
         self.add_behaviour(listen_behaviour)
+
+        # Let UI know that we're connected
+        topic = b"ping"
+        data = json.dumps(True).encode()
+        if self.pub_socket is None:
+            self.logger.warning("Communication agent pub socket not correctly initialized.")
+        else:
+            try:
+                await asyncio.wait_for(self.pub_socket.send_multipart([topic, data]), 5)
+            except TimeoutError:
+                self.logger.warning(
+                    "Initial connection ping for router timed out in ri_communication_agent."
+                )
+
+        # Make sure to start listening now that we're connected.
+        self.connected = True
         self.logger.info("Finished setting up %s", self.jid)

src/control_backend/api/v1/endpoints/command.py

@@ -1,20 +0,0 @@
-import logging
-
-from fastapi import APIRouter, Request
-
-from control_backend.schemas.ri_message import SpeechCommand
-
-logger = logging.getLogger(__name__)
-
-router = APIRouter()
-
-
-@router.post("/command", status_code=202)
-async def receive_command(command: SpeechCommand, request: Request):
-    # Validate and retrieve data.
-    SpeechCommand.model_validate(command)
-    topic = b"command"
-    pub_socket = request.app.state.endpoints_pub_socket
-    await pub_socket.send_multipart([topic, command.model_dump_json().encode()])
-
-    return {"status": "Command received"}

src/control_backend/api/v1/endpoints/robot.py

@@ -0,0 +1,71 @@
+import asyncio
+import json
+import logging
+
+import zmq.asyncio
+from fastapi import APIRouter, Request
+from fastapi.responses import StreamingResponse
+from zmq.asyncio import Context, Socket
+
+from control_backend.core.config import settings
+from control_backend.schemas.ri_message import SpeechCommand
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter()
+
+
+@router.post("/command", status_code=202)
+async def receive_command(command: SpeechCommand, request: Request):
+    # Validate and retrieve data.
+    SpeechCommand.model_validate(command)
+    topic = b"command"
+
+    pub_socket: Socket = request.app.state.endpoints_pub_socket
+    await pub_socket.send_multipart([topic, command.model_dump_json().encode()])
+
+    return {"status": "Command received"}
+
+
+@router.get("/ping_check")
+async def ping(request: Request):
+    pass
+
+
+@router.get("/ping_stream")
+async def ping_stream(request: Request):
+    """Stream live updates whenever the device state changes."""
+
+    async def event_stream():
+        # Set up internal socket to receive ping updates
+
+        sub_socket = Context.instance().socket(zmq.SUB)
+        sub_socket.connect(settings.zmq_settings.internal_sub_address)
+        sub_socket.setsockopt(zmq.SUBSCRIBE, b"ping")
+        connected = False
+
+        ping_frequency = 2
+
+        # Even though its most likely the updates should alternate
+        # (So, True - False - True - False for connectivity),
+        # let's still check.
+        while True:
+            try:
+                topic, body = await asyncio.wait_for(
+                    sub_socket.recv_multipart(), timeout=ping_frequency
+                )
+                connected = json.loads(body)
+            except TimeoutError:
+                logger.debug("got timeout error in ping loop in ping router")
+                connected = False
+
+            # Stop if client disconnected
+            if await request.is_disconnected():
+                logger.info("Client disconnected from SSE")
+                break
+
+            logger.debug(f"Yielded new connection event in robot ping router: {str(connected)}")
+            connectedJson = json.dumps(connected)
+            yield (f"data: {connectedJson}\n\n")
+
+    return StreamingResponse(event_stream(), media_type="text/event-stream")

src/control_backend/api/v1/router.py

@@ -1,6 +1,6 @@
 from fastapi.routing import APIRouter
 
-from control_backend.api.v1.endpoints import command, logs, message, program, sse
+from control_backend.api.v1.endpoints import logs, message, program, robot, sse
 
 api_router = APIRouter()
 
@@ -8,7 +8,7 @@ api_router.include_router(message.router, tags=["Messages"])
 
 api_router.include_router(sse.router, tags=["SSE"])
 
-api_router.include_router(command.router, tags=["Commands"])
+api_router.include_router(robot.router, prefix="/robot", tags=["Pings", "Commands"])
 
 api_router.include_router(logs.router, tags=["Logs"])