fix: stop streaming events when a step raises (#15714)

* fix: stop streaming events when a step raises

* remove session

llama-index-core/llama_index/core/workflow/workflow.py

@@ -247,6 +247,8 @@ async def run(self, **kwargs: Any) -> str:
 
         # Bubble up the error if any step raised an exception
         if exception_raised:
+            # Make sure to stop streaming, in case the workflow terminated abnormally
+            ctx.write_event_to_stream(StopEvent())
             raise exception_raised
 
         # Raise WorkflowTimeoutError if the workflow timed out

llama-index-core/tests/workflow/test_streaming.py

@@ -8,6 +8,8 @@
 from llama_index.core.workflow.workflow import Workflow
 from llama_index.core.workflow.errors import WorkflowRuntimeError
 
+from .conftest import OneTestEvent
+
 
 class StreamingWorkflow(Workflow):
     @step
@@ -45,3 +47,23 @@ async def test_too_many_runs():
         async for ev in wf.stream_events():
             pass
     await r
+
+
+@pytest.mark.asyncio()
+async def test_task_raised():
+    class DummyWorkflow(Workflow):
+        @step
+        async def step(self, ctx: Context, ev: StartEvent) -> StopEvent:
+            ctx.write_event_to_stream(OneTestEvent(test_param="foo"))
+            raise ValueError("The step raised an error!")
+
+    wf = DummyWorkflow()
+    r = asyncio.create_task(wf.run())
+
+    # Make sure we don't block indefinitely here because the step raised
+    async for ev in wf.stream_events():
+        assert ev.test_param == "foo"
+
+    # Make sure the await actually caught the exception
+    with pytest.raises(ValueError, match="The step raised an error!"):
+        await r