fix: fetch rpc response manually

bec_widgets/cli/rpc/rpc_base.py

@@ -289,15 +289,29 @@ class RPCBase:
 
         if wait_for_rpc_response:
             try:
-                finished = self._msg_wait_event.wait(timeout)
-                if not finished:
-                    logger.error(
-                        "GUI RPC response timeout "
-                        f"request_id={request_id} method={method} receiver={receiver} "
-                        f"target_gui_id={target_gui_id} object_name={self.object_name} "
-                        f"timeout={timeout}"
-                    )
-                    raise RPCResponseTimeoutError(request_id, timeout)
+                start = time.monotonic()
+                while True:
+                    if timeout is None:
+                        wait_timeout = 0.5
+                    else:
+                        remaining = timeout - (time.monotonic() - start)
+                        wait_timeout = max(0.0, min(0.5, remaining))
+
+                    finished = self._msg_wait_event.wait(wait_timeout)
+                    if finished:
+                        break
+                    self._fetch_rpc_response(request_id)
+                    if self._msg_wait_event.is_set():
+                        break
+                    if timeout is not None and (time.monotonic() - start) > timeout:
+                        logger.error(
+                            "GUI RPC response timeout "
+                            f"request_id={request_id} method={method} receiver={receiver} "
+                            f"target_gui_id={target_gui_id} object_name={self.object_name} "
+                            f"timeout={timeout}"
+                        )
+                        raise RPCResponseTimeoutError(request_id, timeout)
+
             finally:
                 self._msg_wait_event.clear()
                 self._client.connector.unregister(
@@ -326,6 +340,13 @@ class RPCBase:
         self._rpc_response = msg
         self._msg_wait_event.set()
 
+    def _fetch_rpc_response(self, request_id: str) -> None:
+        msg = self._client.connector.get(MessageEndpoints.gui_instruction_response(request_id))
+        if msg is None:
+            return
+        self._rpc_response = cast(messages.RequestResponseMessage, msg)
+        self._msg_wait_event.set()
+
     def _create_widget_from_msg_result(self, msg_result):
         if msg_result is None:
             return None

tests/unit_tests/test_rpc_base.py

@@ -1,6 +1,7 @@
 from unittest.mock import MagicMock
 
 import pytest
+from bec_lib import messages
 from bec_lib.device import DeviceBaseWithConfig, Signal
 
 from bec_widgets.cli.rpc import rpc_base as rpc_base_module
@@ -59,6 +60,7 @@ def test_run_rpc_logs_response_timeout(monkeypatch):
     rpc = RPCBase(gui_id="progress_widget", object_name="progressbar")
     rpc._rpc_timeout = 0
     rpc._client = MagicMock()
+    rpc._client.connector.get.return_value = None
 
     info_mock = MagicMock()
     error_mock = MagicMock()
@@ -83,3 +85,55 @@ def test_run_rpc_logs_response_timeout(monkeypatch):
     assert "target_gui_id=progress_widget" in error_message
     assert "object_name=progressbar" in error_message
     assert "timeout=0" in error_message
+
+
+def test_run_rpc_waits_indefinitely_when_timeout_is_none(monkeypatch):
+    rpc = RPCBase(gui_id="progress_widget", object_name="progressbar")
+    rpc._client = MagicMock()
+    rpc._create_widget_from_msg_result = MagicMock(return_value="done")
+
+    response = messages.RequestResponseMessage(accepted=True, message={"result": "rpc-result"})
+
+    wait_calls = {"count": 0}
+
+    def wait_side_effect(timeout):
+        wait_calls["count"] += 1
+        if wait_calls["count"] == 1:
+            return False
+        rpc._rpc_response = response
+        rpc._msg_wait_event.set()
+        return True
+
+    monkeypatch.setattr(rpc._msg_wait_event, "wait", wait_side_effect)
+    fetch_mock = MagicMock()
+    monkeypatch.setattr(rpc, "_fetch_rpc_response", fetch_mock)
+
+    result = rpc._run_rpc("set_value", timeout=None)
+
+    assert result == "done"
+    assert fetch_mock.call_count == 1
+    assert wait_calls["count"] == 2
+    assert rpc._create_widget_from_msg_result.call_args.args == ("rpc-result",)
+
+
+def test_run_rpc_fetches_response_before_timeout(monkeypatch):
+    rpc = RPCBase(gui_id="progress_widget", object_name="progressbar")
+    rpc._client = MagicMock()
+    rpc._create_widget_from_msg_result = MagicMock(return_value="done")
+
+    response = messages.RequestResponseMessage(accepted=True, message={"result": "rpc-result"})
+
+    monkeypatch.setattr(rpc._msg_wait_event, "wait", MagicMock(return_value=False))
+
+    def fetch_side_effect(_request_id):
+        rpc._rpc_response = response
+        rpc._msg_wait_event.set()
+
+    fetch_mock = MagicMock(side_effect=fetch_side_effect)
+    monkeypatch.setattr(rpc, "_fetch_rpc_response", fetch_mock)
+
+    result = rpc._run_rpc("set_value", timeout=0)
+
+    assert result == "done"
+    fetch_mock.assert_called_once()
+    rpc._create_widget_from_msg_result.assert_called_once_with("rpc-result")