fix bugs (#178)

docs/src/content/docs/llamaagents/workflows/deployment.md

@@ -95,7 +95,7 @@ The `WorkflowServer` exposes the following RESTful endpoints:
 | `GET`  | `/workflows`                   | Lists the names of all registered workflows.                                                            |
 | `POST` | `/workflows/{name}/run`        | Runs the specified workflow synchronously and returns the final result.                                 |
 | `POST` | `/workflows/{name}/run-nowait` | Starts the specified workflow asynchronously and returns a `handler_id`.                                |
-| `GET`  | `/results/{handler_id}`        | Retrieves the result of an asynchronously run workflow. Returns `202 Accepted` if still running.       |
+| `GET`  | `/handlers/{handler_id}`        | Retrieves the result of an asynchronously run workflow. Returns `202 Accepted` if still running, `500` if the workflow failed, `200` if the workflow completed.       |
 | `GET`  | `/events/{handler_id}`         | Streams all events from a running workflow as newline-delimited JSON (`application/x-ndjson` and `text/event-stream` if SSE are enabled).          |
 | `POST`  | `/events/{handler_id}`         | Sends an event to a workflow during its execution (useful for human-in-the-loop)         |
 | `GET`  | `/handlers`         |  Get all the workflow handlers (running and completed)        |
@@ -306,7 +306,7 @@ async def main()
 
     async for event in client.get_workflow_events(handler_id=handler_id):
         print("Received data:", event)
-    result = await client.get_result(handler_id)
+    result = await client.get_handler(handler_id)
 
     print(f"Final result: {result.result} (status: {result.status})")
 
@@ -373,7 +373,7 @@ async for event in client.get_workflow_events(handler_id=handler_id):
         )
         msg = "Event has been sent" if sent_event else "Event failed to send"
         print(msg)
-result = await client.get_result(handler_id)
+result = await client.get_handler(handler_id)
 print(f"Workflow complete with status: {result.status})")
 res = OutEvent.model_validate(result.result)
 print("Received final message:", res.output)

examples/client/base/workflow_client.py

@@ -32,7 +32,7 @@ async def main() -> None:
 
     async for event in client.get_workflow_events(handler_id=handler_id):
         print(f"Received event type={event.type} data={event.value}")
-    result = await client.get_result(handler_id)
+    result = await client.get_handler(handler_id)
 
     print(f"Final result: {result.result} (status: {result.status})")
 

examples/client/human_in_the_loop/workflow_client_hitl.py

@@ -33,7 +33,7 @@ async def main() -> None:
             )
             msg = "Event has been sent" if sent_event else "Event failed to send"
             print(msg)
-    result = await client.get_result(handler_id)
+    result = await client.get_handler(handler_id)
     print(f"Workflow complete with status: {result.status})")
     res = OutEvent.model_validate(result.result)
     print("Received final message:", res.output)

src/workflows/client/client.py

@@ -287,19 +287,9 @@ class WorkflowClient:
 
     async def get_result(self, handler_id: str) -> HandlerData:
         """
-        Get the result of the workflow associated with the specified handler ID.
-
-        Args:
-            handler_id (str): ID of the handler running the workflow
-
-        Returns:
-            HandlerData: Complete handler data for the workflow
+        Deprecated. Use get_handler instead.
         """
-        async with self._get_client() as client:
-            response = await client.get(f"/handlers/{handler_id}")
-            _raise_for_status_with_body(response)
-
-            return HandlerData.model_validate(response.json())
+        return await self.get_handler(handler_id)
 
     async def get_handlers(self) -> HandlersListResponse:
         """
@@ -314,6 +304,22 @@ class WorkflowClient:
 
             return HandlersListResponse.model_validate(response.json())
 
+    async def get_handler(self, handler_id: str) -> HandlerData:
+        """
+        Get a single workflow handler by identifier.
+
+        Args:
+            handler_id (str): ID of the handler associated with the workflow run
+
+        Returns:
+            HandlerData: Handler metadata persisted by the server.
+        """
+        async with self._get_client() as client:
+            response = await client.get(f"/handlers/{handler_id}")
+            _raise_for_status_with_body(response)
+
+            return HandlerData.model_validate(response.json())
+
     async def cancel_handler(
         self, handler_id: str, purge: bool = False
     ) -> CancelHandlerResponse:

src/workflows/server/server.py

@@ -9,7 +9,7 @@ import json
 import logging
 from importlib.metadata import version
 from pathlib import Path
-from typing import Any, AsyncGenerator
+from typing import Any, AsyncGenerator, Callable, Awaitable
 from datetime import datetime, timezone
 
 import uvicorn
@@ -514,6 +514,8 @@ class WorkflowServer:
             )
             try:
                 await handler
+                if wrapper.task is not None:
+                    await wrapper.task
                 status = 200
             except Exception as e:
                 status = 500
@@ -927,6 +929,13 @@ class WorkflowServer:
 
         handler = self._handlers.get(handler_id)
         if handler is None:
+            persisted = await self._workflow_store.query(
+                HandlerQuery(handler_id_in=[handler_id])
+            )
+            if persisted:
+                status = persisted[0].status
+                if status in {"completed", "failed", "cancelled"}:
+                    raise HTTPException(detail="Handler is completed", status_code=204)
             raise HTTPException(detail="Handler not found", status_code=404)
         if handler.queue.empty() and handler.task is not None and handler.task.done():
             # https://html.spec.whatwg.org/multipage/server-sent-events.html
@@ -1058,14 +1067,19 @@ class WorkflowServer:
         # Check if handler exists
         wrapper = self._handlers.get(handler_id)
         if wrapper is None:
-            raise HTTPException(detail="Handler not found", status_code=404)
+            handler_data = await self._load_handler(handler_id)
+            if handler_data.status in {"completed", "failed", "cancelled"}:
+                raise HTTPException(
+                    detail="Workflow already completed", status_code=409
+                )
+            else:
+                # this branch is for cases where handler status is running but somehow not in memory
+                # Ideally, this should never happen. We probably need to revisit when we add pause/expire functionality.
+                logger.warning(f"Handler {handler_id} is running but not in memory.")
+                raise HTTPException(detail="Handler expired", status_code=409)
 
         handler = wrapper.run_handler
 
-        # Check if workflow is still running
-        if handler.done():
-            raise HTTPException(detail="Workflow already completed", status_code=409)
-
         # Get the context
         ctx = handler.ctx
         if ctx is None:
@@ -1275,7 +1289,12 @@ class WorkflowServer:
         await wrapper.checkpoint()
         # Now register and start streaming
         self._handlers[handler_id] = wrapper
-        wrapper.start_streaming()
+
+        async def on_finish() -> None:
+            self._handlers.pop(handler_id, None)
+            self._results.pop(handler_id, None)
+
+        wrapper.start_streaming(on_finish=on_finish)
 
         return wrapper
 
@@ -1451,11 +1470,11 @@ class _WorkflowHandler:
         except Exception:
             return None
 
-    def start_streaming(self) -> None:
+    def start_streaming(self, on_finish: Callable[[], Awaitable[None]]) -> None:
         """Start streaming events from the handler and managing state."""
-        self.task = asyncio.create_task(self._stream_events())
+        self.task = asyncio.create_task(self._stream_events(on_finish=on_finish))
 
-    async def _stream_events(self) -> 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()
         async for event in self.run_handler.stream_events(expose_internal=True):
@@ -1483,6 +1502,7 @@ class _WorkflowHandler:
             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

src/workflows/server/utils.py

@@ -4,6 +4,7 @@
 import secrets
 import string
 
+
 alphabet = string.ascii_letters + string.digits  # A-Z, a-z, 0-9
 
 

tests/client/test_client.py

@@ -71,7 +71,25 @@ async def test_get_result_for_handler(client: WorkflowClient) -> None:
 
     # Result should be retrievable again and reference the same handler
     result_again = await client.get_result(handler_id)
-    assert result_again.handler_id == handler_id
+    assert result_again == result
+
+
+@pytest.mark.asyncio
+async def test_get_handler(client: WorkflowClient) -> None:
+    handler = await client.run_workflow(
+        "greeting", start_event=InputEvent(greeting="hello", name="John")
+    )
+    assert handler.status == "completed"
+    handler_id = handler.handler_id
+
+    handler_data = await client.get_handler(handler_id)
+    assert handler_data.handler_id == handler_id
+    assert handler_data.workflow_name == "greeting"
+    assert handler_data.run_id == handler.run_id
+    assert handler_data.status == "completed"
+    assert handler_data.started_at is not None
+    assert handler_data.updated_at is not None
+    assert handler_data.completed_at is not None
 
 
 @pytest.mark.asyncio

tests/server/conftest.py

@@ -65,6 +65,8 @@ class InteractiveWorkflow(Workflow):
 
     @step
     async def end(self, ctx: Context, ev: ExternalEvent) -> StopEvent:
+        if ev.response == "error":
+            raise RuntimeError("Error response received")
         return StopEvent(result=f"received: {ev.response}")