diff --git a/python/restate/server_context.py b/python/restate/server_context.py index 80fbec3..7ecd442 100644 --- a/python/restate/server_context.py +++ b/python/restate/server_context.py @@ -534,6 +534,7 @@ 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.""" @@ -541,6 +542,12 @@ def _encode(self, buffer: bytes) -> bytes: 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 diff --git a/tests/disconnect_hotloop.py b/tests/disconnect_hotloop.py index 46a7b57..15a21e4 100644 --- a/tests/disconnect_hotloop.py +++ b/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