From 3d512ac8e2f91196304e78801bf383be03bbee2c Mon Sep 17 00:00:00 2001 From: Terry Zhao Date: Thu, 30 Oct 2025 05:55:52 -0700 Subject: [PATCH] fix bugs (#178) --- .../docs/llamaagents/workflows/deployment.md | 6 +-- examples/client/base/workflow_client.py | 2 +- .../human_in_the_loop/workflow_client_hitl.py | 2 +- src/workflows/client/client.py | 30 ++++++++------ src/workflows/server/server.py | 40 ++++++++++++++----- src/workflows/server/utils.py | 1 + tests/client/test_client.py | 20 +++++++++- tests/server/conftest.py | 2 + 8 files changed, 75 insertions(+), 28 deletions(-) diff --git a/docs/src/content/docs/llamaagents/workflows/deployment.md b/docs/src/content/docs/llamaagents/workflows/deployment.md index b3ac004d..77609bfd 100644 --- a/docs/src/content/docs/llamaagents/workflows/deployment.md +++ b/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) diff --git a/examples/client/base/workflow_client.py b/examples/client/base/workflow_client.py index cd43ca8d..08c7015f 100644 --- a/examples/client/base/workflow_client.py +++ b/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})") diff --git a/examples/client/human_in_the_loop/workflow_client_hitl.py b/examples/client/human_in_the_loop/workflow_client_hitl.py index 04e44cf5..5aa30bcb 100644 --- a/examples/client/human_in_the_loop/workflow_client_hitl.py +++ b/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) diff --git a/src/workflows/client/client.py b/src/workflows/client/client.py index ce52440c..eb944b52 100644 --- a/src/workflows/client/client.py +++ b/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: diff --git a/src/workflows/server/server.py b/src/workflows/server/server.py index 29b17832..f20e9090 100644 --- a/src/workflows/server/server.py +++ b/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 diff --git a/src/workflows/server/utils.py b/src/workflows/server/utils.py index ff2beabb..ee13f505 100644 --- a/src/workflows/server/utils.py +++ b/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 diff --git a/tests/client/test_client.py b/tests/client/test_client.py index 6f9b652d..e0bc59b4 100644 --- a/tests/client/test_client.py +++ b/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 diff --git a/tests/server/conftest.py b/tests/server/conftest.py index f9823615..cfd9e286 100644 --- a/tests/server/conftest.py +++ b/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}")