Fix reconnect (#199)

* Fix reconnect

* Bump version to 2.11.1

Author: adrianlyjak

pyproject.toml

@@ -19,7 +19,7 @@ dev = [
 
 [project]
 name = "llama-index-workflows"
-version = "2.11.0"
+version = "2.11.1"
 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/client/client.py

@@ -228,6 +228,8 @@ async def get_workflow_events(
                         "include_internal": incl_inter,
                         "acquire_timeout": lock_timeout,
                     },
+                    headers={"Connection": "keep-alive"},
+                    timeout=None,
                 ) as response:
                     # Handle different response codes
                     if response.status_code == 404:

src/workflows/server/server.py

@@ -522,6 +522,9 @@ async def _run_workflow(self, request: Request) -> JSONResponse:
                     await wrapper.task
                 except Exception:
                     pass
+            # explicitly close handlers from this synchronous api so they don't linger with events
+            # that no-one is listening for
+            await self._close_handler(wrapper)
 
             return JSONResponse(
                 wrapper.to_response_model().model_dump(), status_code=status
@@ -1424,7 +1427,7 @@ class _WorkflowHandler:
     # Dependencies for persistence
     _workflow_store: AbstractWorkflowStore
     _persistence_backoff: list[float]
-    _on_finsh: Callable[[], Awaitable[None]] | None = None
+    _on_finish: Callable[[], Awaitable[None]] | None = None
 
     def _as_persistent(self) -> PersistentHandler:
         """Persist the current handler state immediately to the workflow store."""
@@ -1641,7 +1644,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:
+            if self._on_finish is not None and self.run_handler.done():
+                # clean up the resources if the stream has been consumed
                 await self._on_finish()
             self.consumer_mutex.release()
 

tests/server/test_server_endpoints.py

@@ -1238,3 +1238,17 @@ async def _wait_done() -> dict[str, Any]:
             assert data["status"] == "completed"
             assert seen_handler_id["handler_id"] is not None
             assert seen_handler_id["handler_id"] == handler_id
+
+
+@pytest.mark.asyncio
+async def test_run_sync_removes_handler_even_with_unconsumed_events(
+    client: AsyncClient, server: WorkflowServer
+) -> None:
+    # Run a streaming workflow synchronously; it emits user events but we don't consume them here.
+    resp = await client.post("/workflows/streaming/run", json={"kwargs": {"count": 2}})
+    assert resp.status_code == 200
+    data = resp.json()
+    assert data["status"] == "completed"
+
+    # The synchronous run path should clean up the handler from memory even if events remain
+    assert len(server._handlers) == 0

tests/server/test_server_live_http.py

@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2025 LlamaIndex Inc.
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import socket
+from contextlib import closing
+from typing import AsyncGenerator
+
+import httpx
+import pytest
+import uvicorn
+
+from workflows import Workflow
+from workflows.events import StopEvent
+from workflows.server import WorkflowServer
+from workflows.client.client import WorkflowClient
+from .conftest import ExternalEvent, RequestedExternalEvent
+from .util import wait_for_passing
+
+
+def _get_free_port() -> int:
+    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
+        s.bind(("127.0.0.1", 0))
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        return int(s.getsockname()[1])
+
+
+@pytest.fixture
+async def live_server(
+    simple_test_workflow: Workflow,
+    streaming_workflow: Workflow,
+    interactive_workflow: Workflow,
+    error_workflow: Workflow,
+) -> AsyncGenerator[tuple[str, WorkflowServer], None]:
+    port = _get_free_port()
+    server = WorkflowServer()
+    server.add_workflow("test", simple_test_workflow)
+    server.add_workflow("streaming", streaming_workflow)
+    server.add_workflow("interactive", interactive_workflow)
+    server.add_workflow("error", error_workflow)
+
+    config = uvicorn.Config(
+        server.app,
+        host="127.0.0.1",
+        port=port,
+        log_level="error",
+        loop="asyncio",
+    )
+    uv_server = uvicorn.Server(config)
+
+    # Start server in background task (lifespan will start workflows)
+    task = asyncio.create_task(uv_server.serve())
+
+    # Wait until server responds on /health or timeout
+    base_url = f"http://127.0.0.1:{port}"
+    async with httpx.AsyncClient(base_url=base_url, timeout=1.0) as client:
+        for _ in range(50):  # ~0.5s max wait
+            try:
+                resp = await client.get("/health")
+                if resp.status_code == 200:
+                    break
+            except Exception:
+                pass
+            await asyncio.sleep(0.01)
+        else:
+            uv_server.should_exit = True
+            await task
+            raise RuntimeError("Live server did not start in time")
+
+    try:
+        yield base_url, server
+    finally:
+        uv_server.should_exit = True
+        try:
+            await task
+        finally:
+            # Ensure graceful shutdown of workflow server
+            await server.stop()
+
+
+@pytest.mark.asyncio
+async def test_streaming_over_real_http(
+    live_server: tuple[str, WorkflowServer],
+) -> None:
+    base_url, _server = live_server
+
+    client = WorkflowClient(base_url=base_url)
+
+    # 1) Start interactive workflow (no-wait)
+    started = await client.run_workflow_nowait("interactive")
+    handler_id = started.handler_id
+
+    # Stream until we see the RequestedExternalEvent, then respond and stop streaming
+    saw_prompt = False
+    async for ev in client.get_workflow_events(handler_id):
+        event = ev.load_event()
+        if isinstance(event, RequestedExternalEvent):
+            saw_prompt = True
+            sent = await client.send_event(handler_id, ExternalEvent(response="pong"))
+            assert sent.status == "sent"
+            break
+    assert saw_prompt
+
+    data = await client.get_handler(handler_id)
+    assert data.status == "completed"
+    assert data.result is not None
+    assert data.result.value.get("result") == "received: pong"
+
+
+@pytest.mark.asyncio
+async def test_reconnect_stream_and_send_event_succeeds(
+    live_server: tuple[str, WorkflowServer],
+) -> None:
+    base_url, _server = live_server
+
+    client = WorkflowClient(base_url=base_url)
+
+    # Start the interactive workflow (no-wait)
+    started = await client.run_workflow_nowait("interactive")
+    handler_id = started.handler_id
+
+    # 1) Connect and read until the first RequestedExternalEvent, then disconnect
+    saw_prompt = False
+    async for ev in client.get_workflow_events(handler_id):
+        event = ev.load_event()
+        if isinstance(event, RequestedExternalEvent):
+            saw_prompt = True
+            break  # simulate client disconnect
+    assert saw_prompt
+
+    # 2) Reconnect to stream; this should succeed and allow streaming again
+    stop_seen = asyncio.Event()
+
+    async def _consume_again() -> None:
+        async for ev in client.get_workflow_events(handler_id):
+            event = ev.load_event()
+            if isinstance(event, StopEvent):
+                stop_seen.set()
+                break
+
+    consumer_task = asyncio.create_task(_consume_again())
+    try:
+        # 3) After reconnect, post the human response; this should succeed
+        sent = await client.send_event(handler_id, ExternalEvent(response="pong"))
+        assert sent.status == "sent"
+
+        # 4) Wait for completion
+        async def validate_result_response() -> None:
+            data = await client.get_handler(handler_id)
+            assert data.status == "completed"
+
+        await wait_for_passing(validate_result_response)
+
+        await asyncio.wait_for(stop_seen.wait(), timeout=2.0)
+    finally:
+        if not consumer_task.done():
+            consumer_task.cancel()
+            with contextlib.suppress(Exception):
+                await consumer_task