exception in intervention (#82)

* exception in intervention

* log

Author: ekzhu

src/agnext/application/_single_threaded_agent_runtime.py

@@ -272,7 +272,11 @@ async def process_next(self) -> None:
         match message_envelope:
             case SendMessageEnvelope(message=message, sender=sender, recipient=recipient, future=future):
                 if self._before_send is not None:
-                    temp_message = await self._before_send.on_send(message, sender=sender, recipient=recipient)
+                    try:
+                        temp_message = await self._before_send.on_send(message, sender=sender, recipient=recipient)
+                    except BaseException as e:
+                        future.set_exception(e)
+                        return
                     if temp_message is DropMessage or isinstance(temp_message, DropMessage):
                         future.set_exception(MessageDroppedException())
                         return
@@ -285,7 +289,12 @@ async def process_next(self) -> None:
                 sender=sender,
             ):
                 if self._before_send is not None:
-                    temp_message = await self._before_send.on_publish(message, sender=sender)
+                    try:
+                        temp_message = await self._before_send.on_publish(message, sender=sender)
+                    except BaseException as e:
+                        # TODO: we should raise the intervention exception to the publisher.
+                        logger.error(f"Exception raised in in intervention handler: {e}", exc_info=True)
+                        return
                     if temp_message is DropMessage or isinstance(temp_message, DropMessage):
                         # TODO log message dropped
                         return
@@ -295,7 +304,12 @@ async def process_next(self) -> None:
                 asyncio.create_task(self._process_publish(message_envelope))
             case ResponseMessageEnvelope(message=message, sender=sender, recipient=recipient, future=future):
                 if self._before_send is not None:
-                    temp_message = await self._before_send.on_response(message, sender=sender, recipient=recipient)
+                    try:
+                        temp_message = await self._before_send.on_response(message, sender=sender, recipient=recipient)
+                    except BaseException as e:
+                        # TODO: should we raise the exception to sender of the response instead?
+                        future.set_exception(e)
+                        return
                     if temp_message is DropMessage or isinstance(temp_message, DropMessage):
                         future.set_exception(MessageDroppedException())
                         return

tests/test_intervention.py

@@ -35,13 +35,13 @@ async def on_send(self, message: MessageType, *, sender: Agent | None, recipient
             return message
 
     handler = DebugInterventionHandler()
-    router = SingleThreadedAgentRuntime(before_send=handler)
+    runtime = SingleThreadedAgentRuntime(before_send=handler)
 
-    long_running = LoopbackAgent("name", router)
-    response = router.send_message(MessageType(), recipient=long_running)
+    long_running = LoopbackAgent("name", runtime)
+    response = runtime.send_message(MessageType(), recipient=long_running)
 
     while not response.done():
-        await router.process_next()
+        await runtime.process_next()
 
     assert handler.num_messages == 1
     assert long_running.num_calls == 1
@@ -54,13 +54,13 @@ async def on_send(self, message: MessageType, *, sender: Agent | None, recipient
             return DropMessage # type: ignore
 
     handler = DropSendInterventionHandler()
-    router = SingleThreadedAgentRuntime(before_send=handler)
+    runtime = SingleThreadedAgentRuntime(before_send=handler)
 
-    long_running = LoopbackAgent("name", router)
-    response = router.send_message(MessageType(), recipient=long_running)
+    long_running = LoopbackAgent("name", runtime)
+    response = runtime.send_message(MessageType(), recipient=long_running)
 
     while not response.done():
-        await router.process_next()
+        await runtime.process_next()
 
     with pytest.raises(MessageDroppedException):
         await response
@@ -76,15 +76,63 @@ async def on_response(self, message: MessageType, *, sender: Agent, recipient: A
             return DropMessage # type: ignore
 
     handler = DropResponseInterventionHandler()
-    router = SingleThreadedAgentRuntime(before_send=handler)
+    runtime = SingleThreadedAgentRuntime(before_send=handler)
 
-    long_running = LoopbackAgent("name", router)
-    response = router.send_message(MessageType(), recipient=long_running)
+    long_running = LoopbackAgent("name", runtime)
+    response = runtime.send_message(MessageType(), recipient=long_running)
 
     while not response.done():
-        await router.process_next()
+        await runtime.process_next()
 
     with pytest.raises(MessageDroppedException):
         await response
 
     assert long_running.num_calls == 1
+
+@pytest.mark.asyncio
+async def test_intervention_raise_exception_on_send() -> None:
+
+    class InterventionException(Exception):
+        pass
+
+    class ExceptionInterventionHandler(DefaultInterventionHandler): # type: ignore
+        async def on_send(self, message: MessageType, *, sender: Agent | None, recipient: Agent) -> MessageType | type[DropMessage]: # type: ignore
+            raise InterventionException
+
+    handler = ExceptionInterventionHandler()
+    runtime = SingleThreadedAgentRuntime(before_send=handler)
+
+    long_running = LoopbackAgent("name", runtime)
+    response = runtime.send_message(MessageType(), recipient=long_running)
+
+    while not response.done():
+        await runtime.process_next()
+
+    with pytest.raises(InterventionException):
+        await response
+
+    assert long_running.num_calls == 0
+
+@pytest.mark.asyncio
+async def test_intervention_raise_exception_on_respond() -> None:
+
+    class InterventionException(Exception):
+        pass
+
+    class ExceptionInterventionHandler(DefaultInterventionHandler): # type: ignore
+        async def on_response(self, message: MessageType, *, sender: Agent, recipient: Agent | None) -> MessageType | type[DropMessage]: # type: ignore
+            raise InterventionException
+
+    handler = ExceptionInterventionHandler()
+    runtime = SingleThreadedAgentRuntime(before_send=handler)
+
+    long_running = LoopbackAgent("name", runtime)
+    response = runtime.send_message(MessageType(), recipient=long_running)
+
+    while not response.done():
+        await runtime.process_next()
+
+    with pytest.raises(InterventionException):
+        await response
+
+    assert long_running.num_calls == 1