Report user-raised SdkInternalBaseException to the VM instead of stalling

python/restate/server_context.py

@@ -534,13 +534,20 @@ def __init__(
         self.request_finished_event = asyncio.Event()
         self.tasks = Tasks()
         self.extension_data: Dict[str, Any] = {}
+        self._sdk_internal_abort = False
 
     def _encode(self, buffer: bytes) -> bytes:
         """Apply the journal value codec encode step, if a codec is configured."""
         if self.journal_codec is None:
             return buffer
         return self.journal_codec.encode(buffer)
 
+    def _handle_internal_exception(self, error: BaseException) -> None:
+        if not self._sdk_internal_abort:
+            logger.warning("Internal SDK exception raised without a VM abort", exc_info=error)
+            stacktrace = "".join(traceback.format_exception(error))
+            self.vm.notify_error(repr(error), stacktrace)
+
     async def enter(self):
         """Invoke the user code."""
         update_restate_context_is_replaying(self.vm)
@@ -570,8 +577,8 @@ async def enter(self):
         # pylint: disable=W0718
         except asyncio.CancelledError:
             pass
-        except SdkInternalBaseException:
-            pass
+        except SdkInternalBaseException as e:
+            self._handle_internal_exception(e)
         except DisconnectedException:
             raise
         except Exception as e:
@@ -591,6 +598,7 @@ async def enter(self):
                     self.vm.notify_error(cause.message, stacktrace, cause.retry_after)
                     break
                 elif isinstance(cause, SdkInternalBaseException):
+                    self._handle_internal_exception(e)
                     break
                 cause = cause.__cause__
             else:
@@ -667,10 +675,12 @@ async def must_take_notification(self, handle):
             # We might need to write out something at this point.
             await self.take_and_send_output()
             logger.exception("Exception in take_notification", exc_info=res)
+            self._sdk_internal_abort = True
             raise SdkInternalException() from res
         if isinstance(res, Suspended):
             # We might need to write out something at this point.
             await self.take_and_send_output()
+            self._sdk_internal_abort = True
             raise SuspendedException()
         if isinstance(res, NotReady):
             raise ValueError(f"Unexpected value error: {handle}")
@@ -687,8 +697,10 @@ async def create_poll_or_cancel_coroutine(self, unresolved_future: UnresolvedFut
             do_progress_response = self.vm.do_progress(unresolved_future)
             if isinstance(do_progress_response, BaseException):
                 logger.exception("Exception in do_progress", exc_info=do_progress_response)
+                self._sdk_internal_abort = True
                 raise SdkInternalException() from do_progress_response
             if isinstance(do_progress_response, Suspended):
+                self._sdk_internal_abort = True
                 raise SuspendedException()
             if isinstance(do_progress_response, DoProgressAnyCompleted):
                 # One of the handles completed

tests/disconnect_hotloop.py

@@ -18,11 +18,13 @@
 
 import asyncio
 from typing import cast
-from unittest.mock import MagicMock
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
-
+from restate.exceptions import SdkInternalBaseException, SdkInternalException, SuspendedException
+from restate.server_context import ServerInvocationContext
 from restate.server_types import ASGIReceiveEvent, ReceiveChannel
+from restate.vm import SingleUnresolvedFuture, Suspended
 
 
 @pytest.fixture(scope="session")
@@ -254,3 +256,59 @@ async def waiter(idx: int):
         assert sorted(results) == [0, 1, 2]
     finally:
         await channel.close()
+
+
+@pytest.mark.parametrize("wrapped", [False, True])
+@pytest.mark.parametrize(
+    "origin", ["user", "notification_error", "notification_suspended", "progress_error", "progress_suspended"]
+)
+async def test_internal_exception_reports_only_without_vm_abort(monkeypatch, caplog, wrapped, origin):
+    vm = MagicMock()
+    vm.take_output.return_value = None
+    handler = MagicMock()
+    handler.context_managers = []
+    invocation = MagicMock(random_seed=0, input_buffer=b"")
+    ctx = ServerInvocationContext(vm, handler, invocation, {}, AsyncMock(), MagicMock())
+    raised: list[BaseException] = []
+
+    async def invoke(**kwargs):
+        try:
+            if origin == "user":
+                raise SdkInternalException()
+            if origin.startswith("notification"):
+                vm.take_notification.return_value = (
+                    Suspended() if origin.endswith("suspended") else ValueError("VM notification failed")
+                )
+                await ctx.must_take_notification(0)
+            else:
+                vm.do_progress.return_value = (
+                    Suspended() if origin.endswith("suspended") else ValueError("VM progress failed")
+                )
+                await ctx.create_poll_or_cancel_coroutine(SingleUnresolvedFuture(0))
+        except SdkInternalBaseException as e:
+            assert isinstance(e, SuspendedException if origin.endswith("suspended") else SdkInternalException)
+            if wrapped:
+                error = RuntimeError("wrapper failed")
+                raised.append(error)
+                raise error from e
+            raised.append(e)
+            raise
+
+    monkeypatch.setattr("restate.server_context.invoke_handler", invoke)
+    await ctx.enter()
+
+    assert len(raised) == 1
+    vm.sys_end.assert_not_called()
+    vm.sys_write_output_success.assert_not_called()
+    vm.sys_write_output_failure.assert_not_called()
+    if origin == "user":
+        vm.notify_error.assert_called_once()
+        message, stacktrace = vm.notify_error.call_args.args
+        assert message == repr(raised[0])
+        assert "SdkInternalException" in stacktrace
+        if wrapped:
+            assert "RuntimeError: wrapper failed" in stacktrace
+        assert "Internal SDK exception raised without a VM abort" in caplog.text
+    else:
+        vm.notify_error.assert_not_called()
+        assert "Internal SDK exception raised without a VM abort" not in caplog.text