[bug fix] Ensure that StopEvent gets cleared from `Context._in_prog…

…ress["_done"]` after a Workflow run (#17300)

* remove StopEvent from in progress

* add new and update stepwise test

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

@@ -247,6 +247,7 @@ async def _task(
                                 new_ev = await instrumented_step(**kwargs)
                                 break  # exit the retrying loop
                             except WorkflowDone:
+                                await ctx.remove_from_in_progress(name=name, ev=ev)
                                 raise
                             except Exception as e:
                                 if config.retry_policy is None:
@@ -277,6 +278,7 @@ async def _task(
                                 None, run_task
                             )
                         except WorkflowDone:
+                            await ctx.remove_from_in_progress(name=name, ev=ev)
                             raise
                         except Exception as e:
                             raise WorkflowRuntimeError(

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

@@ -120,3 +120,13 @@ async def test_deprecated_params(ctx):
         DeprecationWarning, match="`make_private` is deprecated and will be ignored"
     ):
         await ctx.set("foo", 42, make_private=True)
+
+
+@pytest.mark.asyncio()
+async def test_empty_inprogress_when_workflow_done(workflow):
+    h = workflow.run()
+    _ = await h
+
+    # there shouldn't be any in progress events
+    for inprogress_list in h.ctx._in_progress.values():
+        assert len(inprogress_list) == 0

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

@@ -76,6 +76,9 @@ async def test_workflow_run_step(workflow):
     result = await handler
     assert handler.is_done()
     assert result == "Workflow completed"
+    # there shouldn't be any in progress events
+    for inprogress_list in handler.ctx._in_progress.values():
+        assert len(inprogress_list) == 0
 
 
 @pytest.mark.asyncio()