Fix/mqtt publish hanging (#81)

* Fix robot status hang issue with best-effort publishing

Replace blocking _send_robot_status with non-blocking best-effort approach:
- Remove wait_for_publish() calls that could hang indefinitely
- Use direct client.publish() with exception handling
- Maintain status accuracy through InOrbit's multi-layer detection system
- Eliminate threading complexity while preserving functionality

* Add get_state protocol handler with callback mechanism to EdgeSDK

- Add get_state handling to _handle_in_cmd method for InOrbit protocol compliance
- Add set_online_status_callback() method for connectors to provide status logic
- Add _handle_get_state() method that uses callback or defaults to online

* Fix get_state tests to use dynamic module version

---------

Co-authored-by: Jehan Shah <jehan@inorbit.ai>

Author: jehanshah8

inorbit_edge/robot.py

@@ -309,6 +309,9 @@ def __init__(self, robot_id, robot_name, api_key=None, **kwargs) -> None:
         self.map_data_mutex = threading.Lock()
         self.map_files: dict[str, RobotMap] = {}  # label to map data
 
+        # Callback for determining robot online status
+        self._online_status_callback = None
+
         self.message_handlers[MQTT_INITIAL_POSE] = self._handle_initial_pose
         self.message_handlers[MQTT_CUSTOM_COMMAND] = self._handle_custom_command
         self.message_handlers[MQTT_CUSTOM_COMMAND_MESSAGE] = self._handle_custom_message
@@ -473,11 +476,10 @@ def _on_connect(self, client, userdata, flags, reason_code, properties):
             )
             return
 
-        # Send robot online status.
-        # This method is blocking so do it on a separate thread just in case.
-        threading.Thread(
-            target=self._send_robot_status, kwargs={"robot_status": "1"}
-        ).start()
+        # Send robot online status (best effort)
+        # If this fails and InOrbit is getting data, it will detect discrepancy
+        # and request a status update via get_state command.
+        self._send_robot_status(online=True)
 
         # Subscribe to interesting topics
         self.client.subscribe(
@@ -620,6 +622,8 @@ def _handle_in_cmd(self, msg):
             self._handle_load_module(args[1], args[2])
         if args[0] == "unload_module" and len(args) >= 2:
             self._handle_unload_module(args[1])
+        if args[0] == "get_state":
+            self._handle_get_state()
 
     def _handle_load_module(self, module_name, run_level):
         """Handles a load_module command"""
@@ -631,6 +635,22 @@ def _handle_unload_module(self, module_name):
         if module_name == INORBIT_MODULE_CAMERAS:
             self._stop_cameras_streaming()
 
+    def _handle_get_state(self):
+        """Handle get_state command from InOrbit."""
+        is_online = True  # Default assumption
+
+        if self._online_status_callback:
+            try:
+                is_online = self._online_status_callback()
+            except Exception as e:
+                self.logger.error(f"Online status callback failed: {e}")
+                # Fall back to default (True) on callback error
+
+        self._send_robot_status(online=is_online)
+        self.logger.debug(
+            f"Responded to get_state: robot {'online' if is_online else 'offline'}"
+        )
+
     def _start_cameras_streaming(self):
         """Start streaming on all registered cameras"""
         with self.camera_streaming_mutex:
@@ -893,6 +913,17 @@ def register_command_callback(self, callback):
 
         self.command_callbacks.append(callback)
 
+    def set_online_status_callback(self, callback):
+        """Set callback to determine robot online status.
+
+        Args:
+            callback: A callable that returns bool indicating if robot is online.
+                Should return True if robot is online, False otherwise.
+                Will be called when InOrbit requests status via get_state command.
+        """
+        if callable(callback):
+            self._online_status_callback = callback
+
     def unregister_command_callback(self, callback):
         """Unregisters the specified callback"""
         # TODO: Implement
@@ -922,48 +953,36 @@ def _resend_modules(self):
             qos=1,
         )
 
-    def _send_robot_status(self, robot_status):
-        """Sends robot online/offline status message.
-
-        This method blocks until either the message
-        is sent or the client errors out.
+    def _send_robot_status(self, online=True):
+        """Send robot online/offline status (best effort, non-blocking).
 
         Args:
-            robot_status (Union[bool,str]): Connection status
-                It supports ``bool`` and ``str`` values ("0" or "1")
-
-        Raises:
-            ValueError: on invalid ``robot_status``
+            online (bool): True for online status, False for offline status.
         """
+        status_value = "1" if online else "0"
+        status_str = "online" if online else "offline"
 
-        # Validate ``robot_status`` parameter.
-        if isinstance(robot_status, bool):
-            robot_status = "1" if robot_status else "0"
-
-        if robot_status not in ["0", "1"]:
-            raise ValueError("Robot status must be boolean, '0' or '1'")
-
-        # Every time we connect or disconnect to the service, send
-        # updated status including online/offline bit
-        status_message = "{}|{}|{}|{}".format(
-            robot_status, self.robot_api_key, self.agent_version, self.robot_name
-        )
-        ret = self.publish(
-            self._get_robot_subtopic(subtopic=MQTT_SUBTOPIC_STATE),
-            status_message,
-            qos=1,
-            retain=True,
-        )
-        self.logger.info("Publishing status {}. ret = {}.".format(robot_status, ret))
-
-        # TODO: handle errors while waiting for publish. Consider that
-        # this method would typically run on a separate thread.
-        ret.wait_for_publish()
-        published = ret.is_published()
-
-        self.logger.info(
-            "Robot status '{}' published: {:b}.".format(robot_status, published)
-        )
+        try:
+            status_message = (
+                f"{status_value}|{self.robot_api_key}|"
+                f"{self.agent_version}|{self.robot_name}"
+            )
+            self.client.publish(
+                self._get_robot_subtopic(subtopic=MQTT_SUBTOPIC_STATE),
+                status_message,
+                qos=1,
+                retain=True,
+            )
+            self.logger.debug(f"{status_str.capitalize()} status sent successfully")
+        except Exception as e:
+            fallback_msg = (
+                "InOrbit will detect via data messages"
+                if online
+                else "InOrbit will detect via data absence"
+            )
+            self.logger.debug(
+                f"{status_str.capitalize()} status failed: {e} - {fallback_msg}"
+            )
 
     def _is_connected(self):
         return self.client.is_connected()
@@ -1036,7 +1055,10 @@ def disconnect(self):
         """Ends session, disconnecting from cloud services"""
         self.logger.info("Ending robot session")
         self._stop_cameras_streaming()
-        self._send_robot_status(robot_status="0")
+
+        # Send offline status (best effort, non-blocking)
+        # InOrbit will detect offline via data absence if this fails
+        self._send_robot_status(online=False)
 
         # TODO: Unsubscribe from topics
 

inorbit_edge/tests/test_robot_session.py

@@ -80,10 +80,8 @@ def test_robot_session_connect(mock_mqtt_client, mock_inorbit_api, mock_sleep):
     assert robot_session.robot_api_key == "robot_apikey_123"
     # check publish state was called with the correct API key
     robot_session.client.publish.assert_any_call(
-        topic="r/id_123/state",
-        payload="1|robot_apikey_123|{}.edgesdk_py|name_123".format(
-            get_module_version()
-        ),
+        "r/id_123/state",
+        "1|robot_apikey_123|{}.edgesdk_py|name_123".format(get_module_version()),
         qos=1,
         retain=True,
     )

inorbit_edge/tests/test_robot_session_callbacks.py

@@ -15,6 +15,7 @@
     MapRequest,
 )
 from inorbit_edge.robot import RobotSession
+from inorbit_edge import get_module_version
 from inorbit_edge.tests.utils.helpers import test_robot_session_connect_helper
 
 
@@ -38,6 +39,96 @@ def test_builtin_callbacks(mock_mqtt_client, mock_inorbit_api, mock_sleep):
     robot_session.client.subscribe.assert_any_call(topic="r/id_123/ros/loc/mapreq")
 
 
+def test_responds_to_get_state_with_default_online_status(
+    mock_mqtt_client, mock_inorbit_api, mock_sleep
+):
+    """Test that get_state command responds with default online status."""
+    robot_session = RobotSession(
+        robot_id="id_123", robot_name="name_123", api_key="apikey_123"
+    )
+    robot_session.connect()
+    robot_session._on_connect(None, None, None, 0, None)
+
+    # Simulate get_state command
+    robot_session._handle_in_cmd(b"get_state")
+
+    # Verify online status was published
+    robot_session.client.publish.assert_any_call(
+        "r/id_123/state",
+        "1|robot_apikey_123|{}.edgesdk_py|name_123".format(get_module_version()),
+        qos=1,
+        retain=True,
+    )
+
+
+def test_responds_to_get_state_with_callback_online_status(
+    mock_mqtt_client, mock_inorbit_api, mock_sleep
+):
+    """Test that get_state command uses callback for online status."""
+    robot_session = RobotSession(
+        robot_id="id_123", robot_name="name_123", api_key="apikey_123"
+    )
+    robot_session.connect()
+    robot_session._on_connect(None, None, None, 0, None)
+
+    # Set callback that returns False (offline)
+    robot_session.set_online_status_callback(lambda: False)
+
+    # Simulate get_state command
+    robot_session._handle_in_cmd(b"get_state")
+
+    # Verify offline status was published
+    robot_session.client.publish.assert_any_call(
+        "r/id_123/state",
+        "0|robot_apikey_123|{}.edgesdk_py|name_123".format(get_module_version()),
+        qos=1,
+        retain=True,
+    )
+
+
+def test_get_state_handles_callback_exception(
+    mock_mqtt_client, mock_inorbit_api, mock_sleep
+):
+    """Test that get_state handles callback exceptions gracefully."""
+    robot_session = RobotSession(
+        robot_id="id_123", robot_name="name_123", api_key="apikey_123"
+    )
+    robot_session.connect()
+    robot_session._on_connect(None, None, None, 0, None)
+
+    # Set callback that raises exception
+    def failing_callback():
+        raise RuntimeError("Test error")
+
+    robot_session.set_online_status_callback(failing_callback)
+
+    # Simulate get_state command - should not raise exception
+    robot_session._handle_in_cmd(b"get_state")
+
+    # Verify default online status was published despite callback error
+    robot_session.client.publish.assert_any_call(
+        "r/id_123/state",
+        "1|robot_apikey_123|{}.edgesdk_py|name_123".format(get_module_version()),
+        qos=1,
+        retain=True,
+    )
+
+
+def test_set_online_status_callback_ignores_non_callable(
+    mock_mqtt_client, mock_inorbit_api, mock_sleep
+):
+    """Test that set_online_status_callback ignores non-callable values."""
+    robot_session = RobotSession(
+        robot_id="id_123", robot_name="name_123", api_key="apikey_123"
+    )
+
+    # Try to set non-callable
+    robot_session.set_online_status_callback("not_callable")
+
+    # Should remain None
+    assert robot_session._online_status_callback is None
+
+
 def test_robot_session_register_command_callback(
     mock_mqtt_client, mock_inorbit_api, mock_sleep
 ):