fix: wait for req socket send to make sure we dont stay stuck - if there's no... #23

Merged
9828273 merged 19 commits from feat/cb2ui-robot-connections into dev 2025-11-18 12:24:15 +00:00
3 changed files with 67 additions and 75 deletions
Showing only changes of commit debc87c0bb - Show all commits

View File

@@ -56,28 +56,29 @@ class RICommunicationAgent(BaseAgent):
"we probably dont have any receivers... but let's check!"
)
# Wait up to {seconds_to_wait_total/2} seconds for a reply:)
# 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 :(
# 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 soft killing myself."
"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.error(
"communication agent pub socket not correctly initialized."
"Communication agent pub socket not correctly initialized."
)
else:
try:
@@ -85,17 +86,20 @@ class RICommunicationAgent(BaseAgent):
self.agent.pub_socket.send_multipart([topic, data]), 5
)
except TimeoutError:
self.agent.logger.error(
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
@@ -107,7 +111,7 @@ class RICommunicationAgent(BaseAgent):
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."
)
@@ -143,16 +147,20 @@ class RICommunicationAgent(BaseAgent):
if self._req_socket is None:
continue
# Send our message and receive one back:)
# 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=1.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,
)
@@ -160,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
@@ -194,7 +202,7 @@ class RICommunicationAgent(BaseAgent):
if addr != self._address:
if not bind:
self._req_socket.connect(addr)
else: # TODO: Should this ever be the case?
else:
self._req_socket.bind(addr)
case "actuation":
ri_commands_agent = RICommandAgent(
@@ -210,31 +218,32 @@ 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.error("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 >:)
# Let UI know that we're connected
topic = b"ping"
data = json.dumps(True).encode()
if self.pub_socket is None:
self.logger.error("communication agent pub socket not correctly initialized.")
self.logger.error("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.error(
self.logger.warning(
"Initial connection ping for router timed out in ri_communication_agent."
)

View File

@@ -21,7 +21,6 @@ async def receive_command(command: SpeechCommand, request: Request):
SpeechCommand.model_validate(command)
topic = b"command"
# TODO: Check with Kasper
pub_socket: Socket = request.app.state.endpoints_pub_socket
await pub_socket.send_multipart([topic, command.model_dump_json().encode()])
@@ -48,8 +47,8 @@ async def ping_stream(request: Request):
ping_frequency = 2
# Even though its most likely the updates should alternate
# So, True - False - True - False for connectivity.
# Let's still check:)
# (So, True - False - True - False for connectivity),
# let's still check.
while True:
try:
topic, body = await asyncio.wait_for(
@@ -58,11 +57,11 @@ async def ping_stream(request: Request):
connected = json.loads(body)
except TimeoutError:
logger.debug("got timeout error in ping loop in ping router")
await asyncio.sleep(0.1)
connected = False
# Stop if client disconnected
if await request.is_disconnected():
print("Client disconnected from SSE")
logger.info("Client disconnected from SSE")
break
logger.debug(f"Yielded new connection event in robot ping router: {str(connected)}")

View File

@@ -196,7 +196,7 @@ async def test_setup_creates_socket_and_negotiate_3(zmq_context, caplog):
fake_agent_instance = MockCommandAgent.return_value
fake_agent_instance.start = AsyncMock()
# --- Act ---
with caplog.at_level("ERROR"):
agent = RICommunicationAgent(
"test@server",
"password",
@@ -211,7 +211,6 @@ async def test_setup_creates_socket_and_negotiate_3(zmq_context, caplog):
# Since it failed, there should not be any command agent.
fake_agent_instance.start.assert_not_awaited()
assert "Failed to set up RICommunicationAgent" in caplog.text
# Ensure the agent did not attach a ListenBehaviour
assert not any(isinstance(b, agent.ListenBehaviour) for b in agent.behaviours)
@@ -362,7 +361,7 @@ async def test_setup_creates_socket_and_negotiate_7(zmq_context, caplog):
fake_agent_instance.start = AsyncMock()
# --- Act ---
with caplog.at_level("WARNING"):
agent = RICommunicationAgent(
"test@server",
"password",
@@ -377,7 +376,6 @@ async def test_setup_creates_socket_and_negotiate_7(zmq_context, caplog):
# Since it failed, there should not be any command agent.
fake_agent_instance.start.assert_not_awaited()
assert "Unhandled negotiation id:" in caplog.text
@pytest.mark.asyncio
@@ -398,7 +396,7 @@ async def test_setup_creates_socket_and_negotiate_timeout(zmq_context, caplog):
fake_agent_instance.start = AsyncMock()
# --- Act ---
with caplog.at_level("WARNING"):
agent = RICommunicationAgent(
"test@server",
"password",
@@ -412,7 +410,6 @@ async def test_setup_creates_socket_and_negotiate_timeout(zmq_context, caplog):
# Since it failed, there should not be any command agent.
fake_agent_instance.start.assert_not_awaited()
assert "No connection established in 20 seconds" in caplog.text
# Ensure the agent did not attach a ListenBehaviour
assert not any(isinstance(b, agent.ListenBehaviour) for b in agent.behaviours)
@@ -425,7 +422,6 @@ async def test_listen_behaviour_ping_correct(caplog):
fake_socket.recv_json = AsyncMock(return_value={"endpoint": "ping", "data": {}})
fake_socket.send_multipart = AsyncMock()
# TODO: Integration test between actual server and password needed for spade agents
agent = RICommunicationAgent("test@server", "password")
agent._req_socket = fake_socket
agent.connected = True
@@ -433,13 +429,10 @@ async def test_listen_behaviour_ping_correct(caplog):
behaviour = agent.ListenBehaviour()
agent.add_behaviour(behaviour)
# Run once (CyclicBehaviour normally loops)
with caplog.at_level("DEBUG"):
await behaviour.run()
fake_socket.send_json.assert_awaited()
fake_socket.recv_json.assert_awaited()
assert "Received message" in caplog.text
@pytest.mark.asyncio
@@ -470,10 +463,9 @@ async def test_listen_behaviour_ping_wrong_endpoint(caplog):
agent.add_behaviour(behaviour)
# Run once (CyclicBehaviour normally loops)
with caplog.at_level("INFO"):
await behaviour.run()
assert "Received message with topic different than ping, while ping expected." in caplog.text
fake_socket.send_json.assert_awaited()
fake_socket.recv_json.assert_awaited()
@@ -493,10 +485,9 @@ async def test_listen_behaviour_timeout(zmq_context, caplog):
behaviour = agent.ListenBehaviour()
agent.add_behaviour(behaviour)
with caplog.at_level("INFO"):
await behaviour.run()
assert "No ping" in caplog.text
assert not any(isinstance(b, agent.ListenBehaviour) for b in agent.behaviours)
assert not agent.connected
@pytest.mark.asyncio
@@ -522,11 +513,8 @@ async def test_listen_behaviour_ping_no_endpoint(caplog):
behaviour = agent.ListenBehaviour()
agent.add_behaviour(behaviour)
# Run once (CyclicBehaviour normally loops)
with caplog.at_level("ERROR"):
await behaviour.run()
assert "No received endpoint in message, excepted ping endpoint." in caplog.text
fake_socket.send_json.assert_awaited()
fake_socket.recv_json.assert_awaited()
@@ -546,11 +534,10 @@ async def test_setup_unexpected_exception(zmq_context, caplog):
bind=False,
)
with caplog.at_level("ERROR"):
await agent.setup(max_retries=1)
# Ensure that the error was logged
assert "Unexpected error during negotiation: boom!" in caplog.text
assert not any(isinstance(b, agent.ListenBehaviour) for b in agent.behaviours)
assert not agent.connected
@pytest.mark.asyncio
@@ -582,11 +569,8 @@ async def test_setup_unpacking_exception(zmq_context, caplog):
)
# --- Act & Assert ---
with caplog.at_level("ERROR"):
await agent.setup(max_retries=1)
# Ensure the unpacking exception was logged
assert "Error unpacking negotiation data" in caplog.text
await agent.setup(max_retries=1)
# Ensure no command agent was started
fake_agent_instance.start.assert_not_awaited()