Merge pull request #2147 from shaangill025/issue_2111

Fix: messages stuck in mediator

Author: ianco

aries_cloudagent/admin/server.py

@@ -124,7 +124,9 @@ def __init__(
         self._profile = weakref.ref(profile)
         self._send = send
 
-    async def send_outbound(self, message: OutboundMessage) -> OutboundSendStatus:
+    async def send_outbound(
+        self, message: OutboundMessage, **kwargs
+    ) -> OutboundSendStatus:
         """
         Send outbound message.
 

aries_cloudagent/core/dispatcher.py

@@ -15,15 +15,14 @@
 
 from aiohttp.web import HTTPException
 
-
 from ..connections.models.conn_record import ConnRecord
 from ..core.profile import Profile
 from ..messaging.agent_message import AgentMessage
 from ..messaging.base_message import BaseMessage
 from ..messaging.error import MessageParseError
 from ..messaging.models.base import BaseModelError
 from ..messaging.request_context import RequestContext
-from ..messaging.responder import BaseResponder
+from ..messaging.responder import BaseResponder, SKIP_ACTIVE_CONN_CHECK_MSG_TYPES
 from ..messaging.util import datetime_now
 from ..protocols.connections.v1_0.manager import ConnectionManager
 from ..protocols.problem_report.v1_0.message import ProblemReport
@@ -377,7 +376,9 @@ async def create_outbound(
 
         return await super().create_outbound(message, **kwargs)
 
-    async def send_outbound(self, message: OutboundMessage) -> OutboundSendStatus:
+    async def send_outbound(
+        self, message: OutboundMessage, **kwargs
+    ) -> OutboundSendStatus:
         """
         Send outbound message.
 
@@ -388,6 +389,23 @@ async def send_outbound(self, message: OutboundMessage) -> OutboundSendStatus:
         if not context:
             raise RuntimeError("weakref to context has expired")
 
+        msg_type = kwargs.get("message_type")
+        msg_id = kwargs.get("message_id")
+
+        if (
+            message.connection_id
+            and msg_type
+            and msg_type not in SKIP_ACTIVE_CONN_CHECK_MSG_TYPES
+            and not await super().conn_rec_active_state_check(
+                profile=context.profile,
+                connection_id=message.connection_id,
+            )
+        ):
+            raise RuntimeError(
+                f"Connection {message.connection_id} is not ready"
+                " which is required for sending outbound"
+                f" message {msg_id} of type {msg_type}."
+            )
         return await self._send(context.profile, message, self._inbound_message)
 
     async def send_webhook(self, topic: str, payload: dict):

aries_cloudagent/core/tests/test_dispatcher.py

@@ -6,6 +6,8 @@
 
 from marshmallow import EXCLUDE
 
+from ...cache.base import BaseCache
+from ...cache.in_memory import InMemoryCache
 from ...config.injection_context import InjectionContext
 from ...core.event_bus import EventBus
 from ...core.in_memory import InMemoryProfile
@@ -413,12 +415,85 @@ async def test_create_send_outbound(self):
             profile,
             settings={"timing.enabled": True},
         )
+        registry = profile.inject(ProtocolRegistry)
+        registry.register_message_types(
+            {
+                pfx.qualify(StubAgentMessage.Meta.message_type): StubAgentMessage
+                for pfx in DIDCommPrefix
+            }
+        )
         message = StubAgentMessage()
         responder = test_module.DispatcherResponder(context, message, None)
-        outbound_message = await responder.create_outbound(message)
-        with async_mock.patch.object(responder, "_send", async_mock.AsyncMock()):
+        outbound_message = await responder.create_outbound(
+            json.dumps(message.serialize())
+        )
+        with async_mock.patch.object(
+            responder, "_send", async_mock.AsyncMock()
+        ), async_mock.patch.object(
+            test_module.BaseResponder,
+            "conn_rec_active_state_check",
+            async_mock.AsyncMock(return_value=True),
+        ):
             await responder.send_outbound(outbound_message)
 
+    async def test_create_send_outbound_with_msg_attrs(self):
+        profile = make_profile()
+        context = RequestContext(
+            profile,
+            settings={"timing.enabled": True},
+        )
+        registry = profile.inject(ProtocolRegistry)
+        registry.register_message_types(
+            {
+                pfx.qualify(StubAgentMessage.Meta.message_type): StubAgentMessage
+                for pfx in DIDCommPrefix
+            }
+        )
+        message = StubAgentMessage()
+        responder = test_module.DispatcherResponder(context, message, None)
+        outbound_message = await responder.create_outbound(message)
+        with async_mock.patch.object(
+            responder, "_send", async_mock.AsyncMock()
+        ), async_mock.patch.object(
+            test_module.BaseResponder,
+            "conn_rec_active_state_check",
+            async_mock.AsyncMock(return_value=True),
+        ):
+            await responder.send_outbound(
+                message=outbound_message,
+                message_type=message._message_type,
+                message_id=message._id,
+            )
+
+    async def test_create_send_outbound_with_msg_attrs_x(self):
+        profile = make_profile()
+        context = RequestContext(
+            profile,
+            settings={"timing.enabled": True},
+        )
+        registry = profile.inject(ProtocolRegistry)
+        registry.register_message_types(
+            {
+                pfx.qualify(StubAgentMessage.Meta.message_type): StubAgentMessage
+                for pfx in DIDCommPrefix
+            }
+        )
+        message = StubAgentMessage()
+        responder = test_module.DispatcherResponder(context, message, None)
+        outbound_message = await responder.create_outbound(message)
+        outbound_message.connection_id = "123"
+        with async_mock.patch.object(
+            test_module.BaseResponder,
+            "conn_rec_active_state_check",
+            async_mock.AsyncMock(return_value=False),
+        ):
+            with self.assertRaises(RuntimeError):
+                await responder.send_outbound(
+                    message=outbound_message,
+                    message_type=message._message_type,
+                    message_id=message._id,
+                )
+
     async def test_create_send_webhook(self):
         profile = make_profile()
         context = RequestContext(profile)
@@ -427,16 +502,81 @@ async def test_create_send_webhook(self):
         with pytest.deprecated_call():
             await responder.send_webhook("topic", {"pay": "load"})
 
+    async def test_conn_rec_active_state_check_a(self):
+        profile = make_profile()
+        profile.context.injector.bind_instance(BaseCache, InMemoryCache())
+        context = RequestContext(profile)
+        message = StubAgentMessage()
+        responder = test_module.DispatcherResponder(context, message, None)
+        with async_mock.patch.object(
+            test_module.ConnRecord, "retrieve_by_id", async_mock.AsyncMock()
+        ) as mock_conn_ret_by_id:
+            conn_rec = test_module.ConnRecord()
+            conn_rec.state = test_module.ConnRecord.State.COMPLETED
+            mock_conn_ret_by_id.return_value = conn_rec
+            check_flag = await responder.conn_rec_active_state_check(
+                profile,
+                "conn-id",
+            )
+            assert check_flag
+            check_flag = await responder.conn_rec_active_state_check(
+                profile,
+                "conn-id",
+            )
+            assert check_flag
+
+    async def test_conn_rec_active_state_check_b(self):
+        profile = make_profile()
+        profile.context.injector.bind_instance(BaseCache, InMemoryCache())
+        profile.context.injector.bind_instance(
+            EventBus, async_mock.MagicMock(notify=async_mock.AsyncMock())
+        )
+        context = RequestContext(profile)
+        message = StubAgentMessage()
+        responder = test_module.DispatcherResponder(context, message, None)
+        with async_mock.patch.object(
+            test_module.ConnRecord, "retrieve_by_id", async_mock.AsyncMock()
+        ) as mock_conn_ret_by_id:
+            conn_rec_a = test_module.ConnRecord()
+            conn_rec_a.state = test_module.ConnRecord.State.REQUEST
+            conn_rec_b = test_module.ConnRecord()
+            conn_rec_b.state = test_module.ConnRecord.State.COMPLETED
+            mock_conn_ret_by_id.side_effect = [conn_rec_a, conn_rec_b]
+            check_flag = await responder.conn_rec_active_state_check(
+                profile,
+                "conn-id",
+            )
+            assert check_flag
+
     async def test_create_enc_outbound(self):
         profile = make_profile()
         context = RequestContext(profile)
-        message = b"abc123xyz7890000"
+        message = StubAgentMessage()
         responder = test_module.DispatcherResponder(context, message, None)
         with async_mock.patch.object(
             responder, "send_outbound", async_mock.AsyncMock()
         ) as mock_send_outbound:
             await responder.send(message)
             assert mock_send_outbound.called_once()
+        msg_json = json.dumps(StubAgentMessage().serialize())
+        message = msg_json.encode("utf-8")
+        with async_mock.patch.object(
+            responder, "send_outbound", async_mock.AsyncMock()
+        ) as mock_send_outbound:
+            await responder.send(message)
+
+        message = StubAgentMessage()
+        with async_mock.patch.object(
+            responder, "send_outbound", async_mock.AsyncMock()
+        ) as mock_send_outbound:
+            await responder.send_reply(message)
+            assert mock_send_outbound.called_once()
+
+        message = json.dumps(StubAgentMessage().serialize())
+        with async_mock.patch.object(
+            responder, "send_outbound", async_mock.AsyncMock()
+        ) as mock_send_outbound:
+            await responder.send_reply(message)
 
     async def test_expired_context_x(self):
         def _smaller_scope():

aries_cloudagent/messaging/responder.py

@@ -4,18 +4,30 @@
 The responder is provided to message handlers to enable them to send a new message
 in response to the message being handled.
 """
+import asyncio
+import json
 
 from abc import ABC, abstractmethod
-import json
-from typing import Sequence, Union
+from typing import Sequence, Union, Optional, Tuple
 
+from ..cache.base import BaseCache
 from ..connections.models.connection_target import ConnectionTarget
+from ..connections.models.conn_record import ConnRecord
 from ..core.error import BaseError
+from ..core.profile import Profile
 from ..transport.outbound.message import OutboundMessage
 
 from .base_message import BaseMessage
 from ..transport.outbound.status import OutboundSendStatus
 
+SKIP_ACTIVE_CONN_CHECK_MSG_TYPES = [
+    "didexchange/1.0/request",
+    "didexchange/1.0/response",
+    "connections/1.0/invitation",
+    "connections/1.0/request",
+    "connections/1.0/response",
+]
+
 
 class ResponderError(BaseError):
     """Responder error."""
@@ -79,7 +91,18 @@ async def send(
     ) -> OutboundSendStatus:
         """Convert a message to an OutboundMessage and send it."""
         outbound = await self.create_outbound(message, **kwargs)
-        return await self.send_outbound(outbound)
+        if isinstance(message, BaseMessage):
+            msg_type = message._message_type
+            msg_id = message._id
+        else:
+            msg_dict = json.loads(message)
+            msg_type = msg_dict.get("@type")
+            msg_id = msg_dict.get("@id")
+        return await self.send_outbound(
+            message=outbound,
+            message_type=msg_type,
+            message_id=msg_id,
+        )
 
     async def send_reply(
         self,
@@ -109,10 +132,59 @@ async def send_reply(
             target=target,
             target_list=target_list,
         )
-        return await self.send_outbound(outbound)
+        if isinstance(message, BaseMessage):
+            msg_type = message._message_type
+            msg_id = message._id
+        else:
+            msg_dict = json.loads(message)
+            msg_type = msg_dict.get("@type")
+            msg_id = msg_dict.get("@id")
+        return await self.send_outbound(
+            message=outbound, message_type=msg_type, message_id=msg_id
+        )
+
+    async def conn_rec_active_state_check(
+        self, profile: Profile, connection_id: str, timeout: int = 7
+    ) -> bool:
+        """Check if the connection record is ready for sending outbound message."""
+
+        async def _wait_for_state() -> Tuple[bool, Optional[str]]:
+            while True:
+                async with profile.session() as session:
+                    conn_record = await ConnRecord.retrieve_by_id(
+                        session, connection_id
+                    )
+                    if conn_record.is_ready:
+                        # if ConnRecord.State.get(conn_record.state) in (
+                        #     ConnRecord.State.COMPLETED,
+                        # ):
+                        return (True, conn_record.state)
+                    await asyncio.sleep(1)
+
+        try:
+            cache_key = f"conn_rec_state::{connection_id}"
+            connection_state = None
+            cache = profile.inject_or(BaseCache)
+            if cache:
+                connection_state = await cache.get(cache_key)
+            if connection_state and ConnRecord.State.get(connection_state) in (
+                ConnRecord.State.COMPLETED,
+                ConnRecord.State.RESPONSE,
+            ):
+                return True
+            check_flag, connection_state = await asyncio.wait_for(
+                _wait_for_state(), timeout
+            )
+            if cache and connection_state:
+                await cache.set(cache_key, connection_state)
+            return check_flag
+        except asyncio.TimeoutError:
+            return False
 
     @abstractmethod
-    async def send_outbound(self, message: OutboundMessage) -> OutboundSendStatus:
+    async def send_outbound(
+        self, message: OutboundMessage, **kwargs
+    ) -> OutboundSendStatus:
         """
         Send an outbound message.
 
@@ -152,7 +224,9 @@ async def send_reply(
         self.messages.append((message, kwargs))
         return OutboundSendStatus.QUEUED_FOR_DELIVERY
 
-    async def send_outbound(self, message: OutboundMessage) -> OutboundSendStatus:
+    async def send_outbound(
+        self, message: OutboundMessage, **kwargs
+    ) -> OutboundSendStatus:
         """Send an outbound message."""
         self.messages.append((message, None))
         return OutboundSendStatus.QUEUED_FOR_DELIVERY

aries_cloudagent/protocols/connections/v1_0/routes.py

@@ -15,6 +15,7 @@
 
 from ....admin.request_context import AdminRequestContext
 from ....connections.models.conn_record import ConnRecord, ConnRecordSchema
+from ....cache.base import BaseCache
 from ....messaging.models.base import BaseModelError
 from ....messaging.models.openapi import OpenAPISchema
 from ....messaging.valid import (
@@ -739,6 +740,9 @@ async def connections_remove(request: web.BaseRequest):
         async with profile.session() as session:
             connection = await ConnRecord.retrieve_by_id(session, connection_id)
             await connection.delete_record(session)
+            cache = session.inject_or(BaseCache)
+            if cache:
+                await cache.clear(f"conn_rec_state::{connection_id}")
     except StorageNotFoundError as err:
         raise web.HTTPNotFound(reason=err.roll_up) from err
     except StorageError as err: