validate that workflow is not cleared until events are consumed (#186)

* validate that workflow is not cleared until events are consumed

* version bump

* move completion check to function located with status union

Author: adrianlyjak

pyproject.toml

@@ -18,7 +18,7 @@ dev = [
 
 [project]
 name = "llama-index-workflows"
-version = "2.10.1"
+version = "2.10.2"
 description = "An event-driven, async-first, step-based way to control the execution flow of AI applications like Agents."
 readme = "README.md"
 license = "MIT"

src/workflows/protocol/__init__.py

@@ -11,6 +11,10 @@
 Status = Literal["running", "completed", "failed", "cancelled"]
 
 
+def is_status_completed(status: Status) -> bool:
+    return status in {"completed", "failed", "cancelled"}
+
+
 class HandlerData(BaseModel):
     handler_id: str
     workflow_name: str

src/workflows/server/server.py

@@ -44,6 +44,7 @@
     WorkflowEventsListResponse,
     WorkflowGraphResponse,
     WorkflowSchemaResponse,
+    is_status_completed,
 )
 from workflows.server.abstract_workflow_store import (
     AbstractWorkflowStore,
@@ -1125,9 +1126,11 @@ async def _post_event(self, request: Request) -> JSONResponse:
 
         # Check if handler exists
         wrapper = self._handlers.get(handler_id)
+        if wrapper is not None and is_status_completed(wrapper.status):
+            raise HTTPException(detail="Workflow already completed", status_code=409)
         if wrapper is None:
             handler_data = await self._load_handler(handler_id)
-            if handler_data.status in {"completed", "failed", "cancelled"}:
+            if is_status_completed(handler_data.status):
                 raise HTTPException(
                     detail="Workflow already completed", status_code=409
                 )
@@ -1407,6 +1410,7 @@ class _WorkflowHandler:
     # Dependencies for persistence
     _workflow_store: AbstractWorkflowStore
     _persistence_backoff: list[float]
+    _on_finsh: Callable[[], Awaitable[None]] | None = None
 
     async def persist(self) -> None:
         """Persist the current handler state immediately to the workflow store."""
@@ -1536,6 +1540,7 @@ def start_streaming(self, on_finish: Callable[[], Awaitable[None]]) -> None:
     async def _stream_events(self, on_finish: Callable[[], Awaitable[None]]) -> None:
         """Internal method that streams events, updates status, and persists state."""
         await self.checkpoint()
+        self._on_finish = on_finish
         async for event in self.run_handler.stream_events(expose_internal=True):
             if (  # Watch for a specific internal event that signals the step is complete
                 isinstance(event, StepStateChanged)
@@ -1561,7 +1566,6 @@ async def _stream_events(self, on_finish: Callable[[], Awaitable[None]]) -> None
             logger.error(f"Workflow run {self.handler_id} failed! {e}", exc_info=True)
 
         await self.checkpoint()
-        await on_finish()
 
     async def acquire_events_stream(
         self, timeout: float = 1
@@ -1608,6 +1612,8 @@ async def _iter_events(self, timeout: float = 1) -> AsyncGenerator[Event, None]:
                     queue_get_task.cancel()
                     break
         finally:
+            if self._on_finish is not None:
+                await self._on_finish()
             self.consumer_mutex.release()
 
 

tests/server/test_server_endpoints.py

@@ -970,7 +970,9 @@ async def test_post_event_context_not_available(
 ) -> None:
     # Dumb test for code coverage. Inject a dummy handler with no context to trigger 500 path
     wrapper = SimpleNamespace(
-        run_handler=SimpleNamespace(done=lambda: False, ctx=None), workflow_name="test"
+        run_handler=SimpleNamespace(done=lambda: False, ctx=None),
+        workflow_name="test",
+        status="running",
     )
 
     handler_id = "noctx-1"
@@ -1162,3 +1164,39 @@ async def _wait_done() -> Response:
     assert isinstance(data, dict)
     assert data.get("handler_id") == handler_id
     assert data.get("result") is not None
+
+
+@pytest.mark.asyncio
+async def test_stream_events_after_completion_should_return_unconsumed_events(
+    client: AsyncClient,
+) -> None:
+    # Start streaming workflow that emits 3 events and completes
+    start_resp = await client.post(
+        "/workflows/streaming/run-nowait", json={"kwargs": {"count": 3}}
+    )
+    assert start_resp.status_code == 200
+    handler_id = start_resp.json()["handler_id"]
+
+    # Wait for completion via results endpoint
+    async def _wait_done() -> Response:
+        r = await client.get(f"/handlers/{handler_id}")
+        if r.status_code == 200:
+            return r
+        raise AssertionError("not done")
+
+    await wait_for_passing(_wait_done)
+
+    # Now fetch events AFTER completion. Expect the unconsumed events to still be retrievable.
+    # Use NDJSON for easier parsing.
+    resp = await client.get(f"/events/{handler_id}?sse=false")
+    assert resp.status_code == 200
+    assert resp.headers["content-type"].startswith("application/x-ndjson")
+
+    # Collect NDJSON lines
+    lines: list[str] = []
+    async for line in resp.aiter_lines():
+        data = line.strip()
+        if data:
+            lines.append(data)
+
+    assert len(lines) == 4