run-llama/workflows-py
1
0
Fork 0
mirror of https://github.com/run-llama/workflows-py.git synced 2026-07-19 08:46:21 -04:00
Code Issues 18 Packages Projects Releases 10 Wiki Activity

Remove asyncio queue from control_loop (#315)

Browse Source
This commit is contained in:
Adrian Lyjak
2026-02-03 12:50:10 -05:00
committed by GitHub
parent 6d70056da4
commit 45e7614554
28 changed files with 2191 additions and 1121 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.changeset/poor-hounds-mix.md
+5
View File
@@ -0,0 +1,5 @@
---
"llama-index-workflows": minor
---
Replace InMemoryStateStore types with a corresponding StateStore protocol
.changeset/rotten-rabbits-accept.md
+10
View File
@@ -0,0 +1,10 @@
---
"llama-index-workflows-server": patch
"llama-index-workflows": patch
---
Refact: make control loop more deterministic
- Switches out the asyncio delay mechanism for a pull-with-timeout that is more deterministic friendly
- Adds a priority queue of delayed tasks
- Switches out the misc firing /spawning of async tasks to a more rigorous pattern where tasks are only created in the main loop, and gathered in one location. This makes the concurrency more straightforward to reason about
packages/llama-agents-integration-tests/pyproject.toml
+1
View File
@@ -55,3 +55,4 @@ filterwarnings = [
[tool.uv.sources]
llama-index-workflows = {workspace = true}
llama-index-workflows-dbos = {workspace = true}
packages/llama-agents-integration-tests/tests/test_runtime_matrix.py
+493
View File
@@ -0,0 +1,493 @@
"""Runtime matrix tests - testing workflows against both BasicRuntime and DBOSRuntime.
All workflow classes are defined at module level so they can be registered with
DBOS once at module initialization time, avoiding repeated init/destroy cycles.
"""
from __future__ import annotations
import asyncio
from typing import AsyncGenerator, Optional, Union
import pytest
from pydantic import Field
from workflows.context import Context
from workflows.decorators import step
from workflows.errors import WorkflowTimeoutError
from workflows.events import (
Event,
HumanResponseEvent,
InputRequiredEvent,
StartEvent,
StopEvent,
)
from workflows.plugins.basic import BasicRuntime
from workflows.runtime.types.plugin import Runtime
from workflows.testing import WorkflowTestRunner
from workflows.workflow import Workflow
# -- Fixtures --
@pytest.fixture(
params=[
pytest.param("basic", id="basic"),
]
)
async def runtime(
request: pytest.FixtureRequest,
) -> AsyncGenerator[Runtime, None]:
"""Yield an unlaunched runtime.
For DBOS, returns the module-scoped runtime (already created, not yet launched).
Each test must call runtime.launch() after creating workflows.
"""
if request.param == "basic":
rt = BasicRuntime()
try:
yield rt
finally:
rt.destroy()
# -- Shared event types --
class OneTestEvent(Event):
test_param: str = Field(default="test")
class AnotherTestEvent(Event):
another_test_param: str = Field(default="another_test")
class LastEvent(Event):
pass
class MyStart(StartEvent):
query: str
class MyStop(StopEvent):
outcome: str
# -- Workflow definitions (module level for DBOS registration) --
class SimpleWorkflow(Workflow):
@step
async def start_step(self, ev: StartEvent) -> OneTestEvent:
return OneTestEvent()
@step
async def middle_step(self, ev: OneTestEvent) -> LastEvent:
return LastEvent()
@step
async def end_step(self, ev: LastEvent) -> StopEvent:
return StopEvent(result="Workflow completed")
class SlowWorkflow(Workflow):
@step
async def slow_step(self, ev: StartEvent) -> StopEvent:
await asyncio.sleep(2.0)
return StopEvent(result="Done")
class EventTrackingWorkflow(Workflow):
"""Workflow that tracks events in an external list."""
tracked_events: list[str] = []
@step
async def step1(self, ev: StartEvent) -> OneTestEvent:
self.tracked_events.append("step1")
return OneTestEvent()
@step
async def step2(self, ev: OneTestEvent) -> StopEvent:
self.tracked_events.append("step2")
return StopEvent(result="Done")
class SyncAsyncWorkflow(Workflow):
@step
async def async_step(self, ev: StartEvent) -> OneTestEvent:
return OneTestEvent()
@step
def sync_step(self, ev: OneTestEvent) -> StopEvent:
return StopEvent(result="Done")
class SyncWorkflow(Workflow):
@step
def step_one(self, ctx: Context, ev: StartEvent) -> OneTestEvent:
ctx.collect_events(ev, [StartEvent])
return OneTestEvent()
@step
def step_two(self, ctx: Context, ev: OneTestEvent) -> StopEvent:
return StopEvent()
class MultiRunWorkflow(Workflow):
@step
async def step(self, ev: StartEvent) -> StopEvent:
return StopEvent(result=ev.number * 2) # type: ignore
class ErrorWorkflow(Workflow):
@step
async def step(self, ev: StartEvent) -> StopEvent:
raise ValueError("The step raised an error!")
class CounterWorkflow(Workflow):
@step
async def step(self, ctx: Context, ev: StartEvent) -> StopEvent:
cur_number = await ctx.store.get("number", default=0)
new_number = cur_number + 1
await ctx.store.set("number", new_number)
return StopEvent(result=new_number)
class StepSendEventWorkflow(Workflow):
@step
async def step1(self, ctx: Context, ev: StartEvent) -> OneTestEvent:
ctx.send_event(OneTestEvent(), step="step2")
return None # type: ignore
@step
async def step2(self, ev: OneTestEvent) -> StopEvent:
return StopEvent(result="step2")
@step
async def step3(self, ev: OneTestEvent) -> StopEvent:
return StopEvent(result="step3")
class NumWorkersWorkflow(Workflow):
@step
async def original_step(
self, ctx: Context, ev: StartEvent
) -> Union[OneTestEvent, LastEvent]:
await ctx.store.set("num_to_collect", 3)
ctx.send_event(OneTestEvent(test_param="test1"))
ctx.send_event(OneTestEvent(test_param="test2"))
ctx.send_event(OneTestEvent(test_param="test3"))
ctx.send_event(AnotherTestEvent(another_test_param="test4"))
return LastEvent()
@step(num_workers=3)
async def test_step(self, ev: OneTestEvent) -> AnotherTestEvent:
# Note: await_count logic needs to be injected per-test
return AnotherTestEvent(another_test_param=ev.test_param)
@step
async def final_step(
self, ctx: Context, ev: Union[AnotherTestEvent, LastEvent]
) -> StopEvent:
n = await ctx.store.get("num_to_collect")
events = ctx.collect_events(ev, [AnotherTestEvent] * n)
if events is None:
return None # type: ignore
return StopEvent(result=[ev.another_test_param for ev in events])
class CustomEventsWorkflow(Workflow):
@step
async def start_step(self, ev: MyStart) -> OneTestEvent:
return OneTestEvent()
@step
async def middle_step(self, ev: OneTestEvent) -> LastEvent:
return LastEvent()
@step
async def end_step(self, ev: LastEvent) -> MyStop:
return MyStop(outcome="Workflow completed")
class HITLWorkflow(Workflow):
@step
async def step1(self, ctx: Context, ev: StartEvent) -> InputRequiredEvent:
cur_runs = await ctx.store.get("step1_runs", default=0)
await ctx.store.set("step1_runs", cur_runs + 1)
return InputRequiredEvent(prefix="Enter a number: ") # type:ignore
@step
async def step2(self, ctx: Context, ev: HumanResponseEvent) -> StopEvent:
cur_runs = await ctx.store.get("step2_runs", default=0)
await ctx.store.set("step2_runs", cur_runs + 1)
return StopEvent(result=ev.response)
class StreamWorkflow(Workflow):
@step
async def chat(self, ctx: Context, ev: StartEvent) -> StopEvent:
async def stream_messages() -> AsyncGenerator[str, None]:
resp = "Paul Graham is a British-American computer scientist, entrepreneur, vc, and writer."
for word in resp.split():
yield word
async for w in stream_messages():
ctx.write_event_to_stream(Event(msg=w))
return StopEvent(result=None)
class ErrorStreamingWorkflow(Workflow):
@step
async def step(self, ctx: Context, ev: StartEvent) -> StopEvent:
ctx.write_event_to_stream(OneTestEvent(test_param="foo"))
raise ValueError("The step raised an error!")
class TimeoutStreamingWorkflow(Workflow):
@step
async def step(self, ctx: Context, ev: StartEvent) -> StopEvent:
ctx.write_event_to_stream(OneTestEvent(test_param="foo"))
await asyncio.sleep(2)
return StopEvent()
# -- Tests --
@pytest.mark.asyncio
async def test_workflow_run(runtime: Runtime) -> None:
workflow = SimpleWorkflow(runtime=runtime)
runtime.launch()
r = await WorkflowTestRunner(workflow).run()
assert r.result == "Workflow completed"
@pytest.mark.asyncio
async def test_workflow_timeout(runtime: Runtime) -> None:
wf = SlowWorkflow(timeout=0.1, runtime=runtime)
runtime.launch()
with pytest.raises(WorkflowTimeoutError):
await WorkflowTestRunner(wf).run()
@pytest.mark.asyncio
async def test_workflow_event_propagation(runtime: Runtime) -> None:
events: list[str] = []
class LocalEventTrackingWorkflow(Workflow):
@step
async def step1(self, ev: StartEvent) -> OneTestEvent:
events.append("step1")
return OneTestEvent()
@step
async def step2(self, ev: OneTestEvent) -> StopEvent:
events.append("step2")
return StopEvent(result="Done")
wf = LocalEventTrackingWorkflow(runtime=runtime)
runtime.launch()
await WorkflowTestRunner(wf).run()
assert events == ["step1", "step2"]
@pytest.mark.asyncio
async def test_workflow_sync_async_steps(runtime: Runtime) -> None:
wf = SyncAsyncWorkflow(runtime=runtime)
runtime.launch()
await WorkflowTestRunner(wf).run()
@pytest.mark.asyncio
async def test_workflow_sync_steps_only(runtime: Runtime) -> None:
wf = SyncWorkflow(runtime=runtime)
runtime.launch()
await WorkflowTestRunner(wf).run()
@pytest.mark.asyncio
async def test_workflow_multiple_runs(runtime: Runtime) -> None:
wf = MultiRunWorkflow(runtime=runtime)
runtime.launch()
runner = WorkflowTestRunner(wf)
results = await asyncio.gather(
runner.run(StartEvent(number=3)), # type: ignore
runner.run(StartEvent(number=42)), # type: ignore
runner.run(StartEvent(number=-99)), # type: ignore
)
assert set([r.result for r in results]) == {6, 84, -198}
@pytest.mark.asyncio
async def test_workflow_task_raises(runtime: Runtime) -> None:
wf = ErrorWorkflow(runtime=runtime)
runtime.launch()
with pytest.raises(ValueError, match="The step raised an error!"):
await WorkflowTestRunner(wf).run()
@pytest.mark.asyncio
async def test_workflow_step_send_event(runtime: Runtime) -> None:
workflow = StepSendEventWorkflow(runtime=runtime)
runtime.launch()
r = await WorkflowTestRunner(workflow).run()
assert r.result == "step2"
@pytest.mark.asyncio
async def test_workflow_num_workers(runtime: Runtime) -> None:
signal = asyncio.Event()
lock = asyncio.Lock()
counter = 0
async def await_count(count: int) -> None:
nonlocal counter
async with lock:
counter += 1
if counter == count:
signal.set()
return
await signal.wait()
class LocalNumWorkersWorkflow(Workflow):
@step
async def original_step(
self, ctx: Context, ev: StartEvent
) -> Union[OneTestEvent, LastEvent]:
await ctx.store.set("num_to_collect", 3)
# Send test4 first to ensure it's pulled from receive_queue
# before test_step workers complete. Events are pulled one per
# iteration, so ordering in receive_queue determines delivery order.
ctx.send_event(AnotherTestEvent(another_test_param="test4"))
ctx.send_event(OneTestEvent(test_param="test1"))
ctx.send_event(OneTestEvent(test_param="test2"))
ctx.send_event(OneTestEvent(test_param="test3"))
return LastEvent()
@step(num_workers=3)
async def test_step(self, ev: OneTestEvent) -> AnotherTestEvent:
await await_count(3)
return AnotherTestEvent(another_test_param=ev.test_param)
@step
async def final_step(
self, ctx: Context, ev: Union[AnotherTestEvent, LastEvent]
) -> Optional[StopEvent]:
n = await ctx.store.get("num_to_collect")
events = ctx.collect_events(ev, [AnotherTestEvent] * n)
if events is None:
return None
return StopEvent(result=[ev.another_test_param for ev in events])
workflow = LocalNumWorkersWorkflow(timeout=10, runtime=runtime)
runtime.launch()
r = await WorkflowTestRunner(workflow).run()
assert "test4" in set(r.result)
assert len({"test1", "test2", "test3"} - set(r.result)) == 1
@pytest.mark.asyncio
async def test_custom_stop_event(runtime: Runtime) -> None:
wf = CustomEventsWorkflow(runtime=runtime)
runtime.launch()
assert wf._start_event_class == MyStart
assert wf.start_event_class == wf._start_event_class
assert wf._stop_event_class == MyStop
assert wf.stop_event_class == wf._stop_event_class
result: MyStop = await wf.run(query="foo")
assert result.outcome == "Workflow completed"
# Run again with the same workflow instance
assert wf._start_event_class == MyStart
assert wf._stop_event_class == MyStop
result = await wf.run(query="foo")
assert result.outcome == "Workflow completed"
# ensure that streaming exits
r = await WorkflowTestRunner(wf).run(MyStart(query="foo"))
assert len(r.collected) > 0
@pytest.mark.asyncio
async def test_human_in_the_loop(runtime: Runtime) -> None:
# Create both workflow instances before launch
timeout_wf = HITLWorkflow(timeout=0.01, runtime=runtime)
workflow = HITLWorkflow(runtime=runtime)
runtime.launch()
# workflow should raise a timeout error because hitl only works with streaming
with pytest.raises(WorkflowTimeoutError):
await WorkflowTestRunner(timeout_wf).run()
# workflow should work with streaming
handler = workflow.run()
assert handler.ctx
async for event in handler.stream_events():
if isinstance(event, InputRequiredEvent):
assert event.prefix == "Enter a number: "
handler.ctx.send_event(HumanResponseEvent(response="42")) # type:ignore
final_result = await handler
assert final_result == "42"
@pytest.mark.asyncio
async def test_workflow_stream_events_exits(runtime: Runtime) -> None:
wf = CustomEventsWorkflow(runtime=runtime)
runtime.launch()
handler = wf.run(query="foo")
async def _stream_events() -> MyStop:
async for event in handler.stream_events():
continue
return await handler
stream_task = asyncio.create_task(_stream_events())
result = await asyncio.wait_for(stream_task, timeout=10)
assert result.outcome == "Workflow completed"
# -- Streaming tests --
@pytest.mark.asyncio
async def test_streaming_e2e(runtime: Runtime) -> None:
wf = StreamWorkflow(runtime=runtime)
runtime.launch()
test_runner = WorkflowTestRunner(wf)
r = await test_runner.run(expose_internal=False, exclude_events=[StopEvent])
assert all("msg" in ev for ev in r.collected)
@pytest.mark.asyncio
async def test_streaming_task_raised(runtime: Runtime) -> None:
wf = ErrorStreamingWorkflow(runtime=runtime)
runtime.launch()
r = wf.run()
async for ev in r.stream_events():
if not isinstance(ev, StopEvent):
assert ev.test_param == "foo"
with pytest.raises(ValueError, match="The step raised an error!"):
await r
@pytest.mark.asyncio
async def test_streaming_task_timeout(runtime: Runtime) -> None:
wf = TimeoutStreamingWorkflow(timeout=0.1, runtime=runtime)
runtime.launch()
r = wf.run()
async for ev in r.stream_events():
if not isinstance(ev, StopEvent):
assert ev.test_param == "foo"
with pytest.raises(WorkflowTimeoutError, match="Operation timed out"):
await r
packages/llama-agents-server/src/llama_agents/server/server.py
+47 -27
View File
@@ -40,6 +40,7 @@ from starlette.routing import Route
from starlette.schemas import SchemaGenerator
from starlette.staticfiles import StaticFiles
from workflows import Context, Workflow
from workflows.errors import WorkflowRuntimeError
from workflows.events import (
Event,
InternalDispatchEvent,
@@ -1762,33 +1763,41 @@ class _WorkflowHandler:
with instrument_tags({"handler_id": self.handler_id}):
await self.checkpoint()
self._on_finish = on_finish
async for event in self.run_handler.stream_events(expose_internal=True):
# Track idle state transitions and manage release timer
if isinstance(event, WorkflowIdleEvent):
self.mark_idle()
elif isinstance(event, UnhandledEvent):
self.mark_idle()
elif (
isinstance(event, StepStateChanged)
and event.step_state == StepState.RUNNING
):
self.mark_active()
try:
async for event in self.run_handler.stream_events(expose_internal=True):
# Track idle state transitions and manage release timer
if isinstance(event, WorkflowIdleEvent):
self.mark_idle()
elif isinstance(event, UnhandledEvent):
self.mark_idle()
elif (
isinstance(event, StepStateChanged)
and event.step_state == StepState.RUNNING
):
self.mark_active()
if ( # Watch for a specific internal event that signals the step is complete
isinstance(event, StepStateChanged)
and event.step_state == StepState.NOT_RUNNING
):
state = (
self.run_handler.ctx.to_dict() if self.run_handler.ctx else None
)
if state is None:
logger.warning(
f"Context state is None for handler {self.handler_id}. This is not expected."
if ( # Watch for a specific internal event that signals the step is complete
isinstance(event, StepStateChanged)
and event.step_state == StepState.NOT_RUNNING
):
state = (
self.run_handler.ctx.to_dict()
if self.run_handler.ctx
else None
)
continue
await self.checkpoint()
if state is None:
logger.warning(
f"Context state is None for handler {self.handler_id}. This is not expected."
)
continue
await self.checkpoint()
self.queue.put_nowait(event)
self.queue.put_nowait(event)
except WorkflowRuntimeError:
# Stream was already consumed - this can happen during handler
# cancellation when run_handler is cancelled before this task.
# This is benign; we'll proceed to cleanup.
pass
# Workflow is completing - cancel any pending release timer
self._cancel_idle_release_timer()
@@ -1826,6 +1835,7 @@ class _WorkflowHandler:
Converts the queue to an async generator while the workflow is still running, and there are still events.
For better or worse, multiple consumers will compete for events
"""
queue_get_task: asyncio.Task[Event] | None = None
try:
while not self.queue.empty() or (
@@ -1836,9 +1846,7 @@ class _WorkflowHandler:
available_events.append(self.queue.get_nowait())
for event in available_events:
yield event
queue_get_task: asyncio.Task[Event] = asyncio.create_task(
self.queue.get()
)
queue_get_task = asyncio.create_task(self.queue.get())
task_waitable = self.task
done, pending = await asyncio.wait(
{queue_get_task, task_waitable}
@@ -1848,10 +1856,22 @@ class _WorkflowHandler:
)
if queue_get_task in done:
yield await queue_get_task
queue_get_task = None
else: # otherwise task completed, so nothing else will be published to the queue
queue_get_task.cancel()
queue_get_task = None
break
finally:
# Cancel any pending queue.get() task to prevent orphaned tasks from
# consuming events after the consumer disconnects.
if queue_get_task is not None:
if not queue_get_task.done():
queue_get_task.cancel()
try:
await queue_get_task
except asyncio.CancelledError:
pass
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()
packages/llama-index-workflows/src/llama_agents/workflows/__init__.py
+1 -1
View File
@@ -64,5 +64,5 @@ if not any(isinstance(f, _AliasFinder) for f in sys.meta_path):
sys.meta_path.append(_AliasFinder())
# Re-export everything from the real workflows package
from workflows import * # noqa: E402, F403 # pyright: ignore[reportWildcardImportFromLibrary]
from workflows import * # noqa: E402, F403
from workflows import __all__ # noqa: E402, F401
packages/llama-index-workflows/src/workflows/context/context.py
+17 -7
View File
@@ -42,7 +42,7 @@ from workflows.types import RunResultT
from workflows.utils import _nanoid as nanoid
from .serializers import BaseSerializer, JsonSerializer
from .state_store import MODEL_T, InMemoryStateStore
from .state_store import MODEL_T, StateStore
if TYPE_CHECKING: # pragma: no cover
from workflows import Workflow
@@ -80,9 +80,9 @@ class Context(Generic[MODEL_T]):
Attributes:
is_running (bool): Whether the workflow is currently running.
store (InMemoryStateStore[MODEL_T]): Type-safe, async state store shared
store (StateStore[MODEL_T]): Type-safe, async state store shared
across steps. See also
[InMemoryStateStore][workflows.context.state_store.InMemoryStateStore].
[StateStore][workflows.context.state_store.StateStore].
Examples:
Basic usage inside a step:
@@ -205,7 +205,7 @@ class Context(Generic[MODEL_T]):
Requires a current run context (via with_current_run_id) to be set.
"""
internal_adapter = workflow._runtime.get_internal_adapter()
internal_adapter = workflow._runtime.get_internal_adapter(workflow)
new_ctx = cast(Context[MODEL_T], object.__new__(cls))
new_ctx._face = cast(
InternalContext[MODEL_T],
@@ -276,6 +276,7 @@ class Context(Generic[MODEL_T]):
pre.init_snapshot, workflow, pre._serializer
)
# TODO(v3) - make this async
external_adapter = workflow._runtime.run_workflow(
run_id=run_id,
workflow=workflow,
@@ -311,16 +312,16 @@ class Context(Generic[MODEL_T]):
_warn_cancel_in_step()
@property
def store(self) -> InMemoryStateStore[MODEL_T]:
def store(self) -> StateStore[MODEL_T]:
"""Typed, process-local state store shared across steps.
If no state was initialized yet, a default
[DictState][workflows.context.state_store.DictState] store is created.
Returns:
InMemoryStateStore[MODEL_T]: The state store instance.
StateStore[MODEL_T]: The state store instance.
"""
return self._face.store # type: ignore[return-value]
return self._face.store
def to_dict(self, serializer: BaseSerializer | None = None) -> dict[str, Any]:
"""Serialize the context to a JSON-serializable dict.
@@ -562,6 +563,15 @@ class Context(Generic[MODEL_T]):
"""
self._require_internal(fn="write_event_to_stream").write_event_to_stream(ev)
async def _finalize_step(self) -> None:
"""Finalize step execution by awaiting background tasks.
Called after a step function completes to ensure all fire-and-forget
operations (e.g., write_event_to_stream, send_event) complete before
returning control to the control loop.
"""
await self._require_internal(fn="_finalize_step")._finalize_step()
def get_result(self) -> RunResultT:
"""Return the final result of the workflow run.
packages/llama-index-workflows/src/workflows/context/external_context.py
+4 -4
View File
@@ -3,13 +3,13 @@
from __future__ import annotations
import asyncio
from typing import TYPE_CHECKING, Any, AsyncGenerator, Coroutine, Generic, cast
from typing import TYPE_CHECKING, Any, AsyncGenerator, Coroutine, Generic
from typing_extensions import TypeVar
from workflows.context.context_types import MODEL_T
from workflows.context.serializers import JsonSerializer
from workflows.context.state_store import InMemoryStateStore
from workflows.context.state_store import StateStore
from workflows.errors import WorkflowRuntimeError
from workflows.events import StopEvent
from workflows.runtime.types.internal_state import BrokerState
@@ -104,12 +104,12 @@ class ExternalContext(Generic[MODEL_T, RunResultT]):
return new_state
@property
def store(self) -> InMemoryStateStore[MODEL_T]:
def store(self) -> StateStore[MODEL_T]:
"""Access workflow state store."""
state_store = self._external_adapter.get_state_store()
if state_store is None:
raise RuntimeError("State store not available from adapter")
return cast(InMemoryStateStore[MODEL_T], state_store)
return state_store # type: ignore[return-value]
def send_event(self, message: Event, step: str | None = None) -> None:
"""Send an event into the workflow."""
packages/llama-index-workflows/src/workflows/context/internal_context.py
+16 -3
View File
@@ -8,7 +8,7 @@ from collections import Counter, defaultdict
from typing import TYPE_CHECKING, Any, Coroutine, Generic, Type, TypeVar, cast
from workflows.context.context_types import MODEL_T
from workflows.context.state_store import InMemoryStateStore
from workflows.context.state_store import StateStore
from workflows.errors import WorkflowRuntimeError
from workflows.runtime.types.results import (
AddCollectedEvent,
@@ -80,6 +80,19 @@ class InternalContext(Generic[MODEL_T]):
worker.cancel()
self._workers.clear()
async def _finalize_step(self) -> None:
"""Await all background tasks and finalize the step.
Called after a step function completes to ensure all fire-and-forget
operations (e.g., write_event_to_stream, send_event) complete before
returning control to the control loop. This prevents non-deterministic
ordering of durable operations on replay.
"""
workers = self._workers[:]
if workers:
await asyncio.gather(*workers, return_exceptions=True)
await self._internal_adapter.finalize_step()
@staticmethod
def _get_step_ctx(fn: str) -> StepWorkerContext:
"""Get the current step worker context. Raises if not in a step."""
@@ -91,12 +104,12 @@ class InternalContext(Generic[MODEL_T]):
)
@property
def store(self) -> InMemoryStateStore[MODEL_T]:
def store(self) -> StateStore[MODEL_T]:
"""Access workflow state store."""
state_store = self._internal_adapter.get_state_store()
if state_store is None:
raise RuntimeError("State store not available from adapter")
return cast(InMemoryStateStore[MODEL_T], state_store)
return state_store # type: ignore[return-value]
def collect_events(
self,
packages/llama-index-workflows/src/workflows/context/pre_context.py
+6 -2
View File
@@ -8,7 +8,11 @@ from pydantic import ValidationError
from workflows.context.context_types import MODEL_T, SerializedContext
from workflows.context.serializers import BaseSerializer, JsonSerializer
from workflows.context.state_store import InMemoryStateStore, infer_state_type
from workflows.context.state_store import (
InMemoryStateStore,
StateStore,
infer_state_type,
)
from workflows.errors import ContextSerdeError
from workflows.runtime.types.internal_state import BrokerState
@@ -61,7 +65,7 @@ class PreContext(Generic[MODEL_T]):
self._init_snapshot = previous_context_parsed
@property
def store(self) -> InMemoryStateStore[MODEL_T]:
def store(self) -> StateStore[MODEL_T]:
"""Lazily-created staging store for pre-run state access.
For fresh contexts, the state type is inferred from workflow step
packages/llama-index-workflows/src/workflows/context/state_store.py
+281 -87
View File
@@ -7,9 +7,19 @@ import asyncio
import functools
import warnings
from contextlib import asynccontextmanager
from typing import TYPE_CHECKING, Any, AsyncGenerator, Generic, Type
from typing import (
TYPE_CHECKING,
Any,
AsyncContextManager,
AsyncGenerator,
Generic,
Literal,
Protocol,
Type,
runtime_checkable,
)
from pydantic import BaseModel, ValidationError
from pydantic import BaseModel, ValidationError, model_validator
from typing_extensions import TypeVar
from workflows.decorators import StepConfig
@@ -23,6 +33,165 @@ if TYPE_CHECKING:
MAX_DEPTH = 1000
class InMemorySerializedState(BaseModel):
"""Serialized state containing actual data (from InMemoryStateStore)."""
store_type: Literal["in_memory"] = "in_memory"
state_type: str = "DictState"
state_module: str = "workflows.context.state_store"
state_data: Any = (
None # {"_data": {...}} for DictState, serialized string for typed
)
@model_validator(mode="before")
@classmethod
def default_store_type(cls, data: dict[str, Any]) -> dict[str, Any]:
"""Default missing store_type to 'in_memory' for backwards compatibility."""
if isinstance(data, dict) and "store_type" not in data:
data = {**data, "store_type": "in_memory"}
return data
def parse_in_memory_state(
data: dict[str, Any],
) -> InMemorySerializedState:
"""Parse raw dict into InMemorySerializedState.
Args:
data: Serialized state payload from InMemoryStateStore.to_dict().
Returns:
InMemorySerializedState if the format is recognized.
Raises:
ValueError: If store_type is not 'in_memory' or missing.
"""
store_type = data.get("store_type")
if store_type == "in_memory" or store_type is None:
# Backwards compat: missing store_type = InMemory
return InMemorySerializedState.model_validate(data)
else:
raise ValueError(
f"Cannot parse store_type '{store_type}' as InMemorySerializedState. "
"Use the appropriate store's from_dict() method."
)
def serialize_dict_state_data(
state: DictState,
serializer: BaseSerializer,
known_unserializable_keys: tuple[str, ...] = (),
) -> dict[str, Any]:
"""Serialize DictState items to {"_data": {...}} format.
Args:
state: The DictState to serialize.
serializer: Strategy for encoding values.
known_unserializable_keys: Keys to skip with warning if they fail to serialize.
Returns:
Dict with {"_data": {...}} structure containing serialized values.
Raises:
ValueError: If serialization fails for a non-known-unserializable key.
"""
serialized_data = {}
for key, value in state.items():
try:
serialized_data[key] = serializer.serialize(value)
except Exception as e:
if key in known_unserializable_keys:
warnings.warn(
f"Skipping serialization of known unserializable key: {key} -- "
"This is expected but will require this item to be set manually after deserialization.",
category=UnserializableKeyWarning,
)
continue
raise ValueError(f"Failed to serialize state value for key {key}: {e}")
return {"_data": serialized_data}
def create_in_memory_payload(
state: BaseModel,
serializer: BaseSerializer,
known_unserializable_keys: tuple[str, ...] = (),
) -> InMemorySerializedState:
"""Create InMemorySerializedState from any state model.
Args:
state: The Pydantic model to serialize (DictState or typed model).
serializer: Strategy for encoding values.
known_unserializable_keys: Keys to skip with warning (DictState only).
Returns:
InMemorySerializedState containing the serialized data.
"""
if isinstance(state, DictState):
state_data = serialize_dict_state_data(
state, serializer, known_unserializable_keys
)
else:
state_data = serializer.serialize(state)
return InMemorySerializedState(
state_type=type(state).__name__,
state_module=type(state).__module__,
state_data=state_data,
)
def traverse_path_step(obj: Any, segment: str) -> Any:
"""Follow one segment into obj (dict key, list index, or attribute).
Args:
obj: The object to traverse into.
segment: The path segment (dict key, list index, or attribute name).
Returns:
The value at the given segment.
Raises:
KeyError, IndexError, AttributeError: If the segment doesn't exist.
"""
if isinstance(obj, dict):
return obj[segment]
# Attempt list/tuple index
try:
idx = int(segment)
return obj[idx]
except (ValueError, TypeError, IndexError):
pass
# Fallback to attribute access (Pydantic models, normal objects)
return getattr(obj, segment)
def assign_path_step(obj: Any, segment: str, value: Any) -> None:
"""Assign value to segment of obj (dict key, list index, or attribute).
Args:
obj: The object to assign into.
segment: The path segment (dict key, list index, or attribute name).
value: The value to assign.
"""
if isinstance(obj, dict):
obj[segment] = value
return
# Attempt list/tuple index assignment
try:
idx = int(segment)
obj[idx] = value
return
except (ValueError, TypeError, IndexError):
pass
# Fallback to attribute assignment
setattr(obj, segment, value)
# Only warn once about unserializable keys
class UnserializableKeyWarning(Warning):
pass
@@ -59,12 +228,74 @@ class DictState(DictLikeModel):
MODEL_T = TypeVar("MODEL_T", bound=BaseModel, default=DictState) # type: ignore
@runtime_checkable
class StateStore(Protocol[MODEL_T]):
"""Protocol defining the async state store interface.
State stores hold a single Pydantic model instance representing global
workflow state. Implementations must be async-safe and support both
atomic operations and transactional edits.
This protocol enables runtime plugins to provide custom state store
implementations (e.g., backed by databases, Redis, or external services)
while maintaining API compatibility with the default
[InMemoryStateStore][workflows.context.state_store.InMemoryStateStore].
For remote state stores, `to_dict`/`from_dict` should serialize non-secret
connection info (e.g., URL, table name) rather than the full state contents,
since the actual state lives in the external service.
See Also:
- [InMemoryStateStore][workflows.context.state_store.InMemoryStateStore]
- [Context.store][workflows.context.context.Context.store]
"""
state_type: Type[MODEL_T]
async def get_state(self) -> MODEL_T:
"""Return a copy of the current state model."""
...
async def set_state(self, state: MODEL_T) -> None:
"""Replace or merge into the current state model."""
...
async def get(self, path: str, default: Any = ...) -> Any:
"""Get a nested value using dot-separated paths."""
...
async def set(self, path: str, value: Any) -> None:
"""Set a nested value using dot-separated paths."""
...
async def clear(self) -> None:
"""Reset the state to its type defaults."""
...
def edit_state(self) -> AsyncContextManager[MODEL_T]:
"""Edit state transactionally under a lock."""
...
def to_dict(self, serializer: "BaseSerializer") -> dict[str, Any]:
"""Serialize state for persistence."""
...
@classmethod
def from_dict(
cls,
serialized_state: dict[str, Any],
serializer: "BaseSerializer",
) -> "StateStore[MODEL_T]":
"""Restore state from serialized payload."""
...
class InMemoryStateStore(Generic[MODEL_T]):
"""
Async, in-memory, type-safe state manager for workflows.
Default in-memory implementation of the [StateStore][workflows.context.state_store.StateStore] protocol.
This store holds a single Pydantic model instance representing global
workflow state. When the generic parameter is omitted, it defaults to
Holds a single Pydantic model instance representing global workflow state.
When the generic parameter is omitted, it defaults to
[DictState][workflows.context.state_store.DictState] for flexible,
dictionary-like usage.
@@ -183,38 +414,10 @@ class InMemoryStateStore(Generic[MODEL_T]):
dict[str, Any]: A payload suitable for
[from_dict][workflows.context.state_store.InMemoryStateStore.from_dict].
"""
# Special handling for DictState - serialize each item in _data
if isinstance(self._state, DictState):
serialized_data = {}
for key, value in self._state.items():
try:
serialized_data[key] = serializer.serialize(value)
except Exception as e:
if key in self.known_unserializable_keys:
warnings.warn(
f"Skipping serialization of known unserializable key: {key} -- "
"This is expected but will require this item to be set manually after deserialization.",
category=UnserializableKeyWarning,
)
continue
raise ValueError(
f"Failed to serialize state value for key {key}: {e}"
)
return {
"state_data": {"_data": serialized_data},
"state_type": type(self._state).__name__,
"state_module": type(self._state).__module__,
}
else:
# For regular Pydantic models, rely on pydantic's serialization
serialized_state = serializer.serialize(self._state)
return {
"state_data": serialized_state,
"state_type": type(self._state).__name__,
"state_module": type(self._state).__module__,
}
payload = create_in_memory_payload(
self._state, serializer, self.known_unserializable_keys
)
return payload.model_dump()
@classmethod
def from_dict(
@@ -229,31 +432,17 @@ class InMemoryStateStore(Generic[MODEL_T]):
Returns:
InMemoryStateStore[MODEL_T]: A store with the reconstructed model.
Raises:
ValueError: If the payload is not in_memory format.
"""
if not serialized_state:
# Return a default DictState manager
return cls(DictState()) # type: ignore
state_data = serialized_state.get("state_data", {})
state_type = serialized_state.get("state_type", "DictState")
# Deserialize the state data
if state_type == "DictState":
# Special handling for DictState - deserialize each item in _data
_data_serialized = state_data.get("_data", {})
deserialized_data = {}
for key, value in _data_serialized.items():
try:
deserialized_data[key] = serializer.deserialize(value)
except Exception as e:
raise ValueError(
f"Failed to deserialize state value for key {key}: {e}"
)
state_instance = DictState(_data=deserialized_data)
else:
state_instance = serializer.deserialize(state_data)
# Validate it's in_memory format (raises ValueError if not)
parse_in_memory_state(serialized_state)
state_instance = deserialize_state_from_dict(serialized_state, serializer)
return cls(state_instance) # type: ignore
@asynccontextmanager
@@ -299,7 +488,7 @@ class InMemoryStateStore(Generic[MODEL_T]):
try:
value: Any = self._state
for segment in segments:
value = self._traverse_step(value, segment)
value = traverse_path_step(value, segment)
except Exception:
if default is not Ellipsis:
return default
@@ -335,15 +524,15 @@ class InMemoryStateStore(Generic[MODEL_T]):
# Navigate/create intermediate segments
for segment in segments[:-1]:
try:
current = self._traverse_step(current, segment)
current = traverse_path_step(current, segment)
except (KeyError, AttributeError, IndexError, TypeError):
# Create intermediate object and assign it
intermediate: Any = {}
self._assign_step(current, segment, intermediate)
assign_path_step(current, segment, intermediate)
current = intermediate
# Assign the final value
self._assign_step(current, segments[-1], value)
assign_path_step(current, segments[-1], value)
async def clear(self) -> None:
"""Reset the state to its type defaults.
@@ -357,37 +546,42 @@ class InMemoryStateStore(Generic[MODEL_T]):
except ValidationError:
raise ValueError("State must have defaults for all fields")
def _traverse_step(self, obj: Any, segment: str) -> Any:
"""Follow one segment into *obj* (dict key, list index, or attribute)."""
if isinstance(obj, dict):
return obj[segment]
# attempt list/tuple index
try:
idx = int(segment)
return obj[idx]
except (ValueError, TypeError, IndexError):
pass
def deserialize_state_from_dict(
serialized_state: dict[str, Any], serializer: "BaseSerializer"
) -> BaseModel:
"""Deserialize state from a serialized payload.
# fallback to attribute access (Pydantic models, normal objects)
return getattr(obj, segment)
This is the inverse of InMemoryStateStore.to_dict(). It handles both
DictState (with per-key serialization) and typed Pydantic models.
def _assign_step(self, obj: Any, segment: str, value: Any) -> None:
"""Assign *value* to *segment* of *obj* (dict key, list index, or attribute)."""
if isinstance(obj, dict):
obj[segment] = value
return
Args:
serialized_state: The payload from to_dict(), containing state_data,
state_type, and state_module.
serializer: Strategy to decode stored values.
# attempt list/tuple index assignment
try:
idx = int(segment)
obj[idx] = value
return
except (ValueError, TypeError, IndexError):
pass
Returns:
The deserialized state model instance.
# fallback to attribute assignment
setattr(obj, segment, value)
Raises:
ValueError: If deserialization fails for any key.
"""
state_data = serialized_state.get("state_data", {})
state_type_name = serialized_state.get("state_type", "DictState")
if state_type_name == "DictState":
_data_serialized = state_data.get("_data", {})
deserialized_data = {}
for key, value in _data_serialized.items():
try:
deserialized_data[key] = serializer.deserialize(value)
except Exception as e:
raise ValueError(
f"Failed to deserialize state value for key {key}: {e}"
)
return DictState(_data=deserialized_data)
else:
return serializer.deserialize(state_data)
def infer_state_type(workflow: "Workflow") -> type[BaseModel]:
packages/llama-index-workflows/src/workflows/plugins/basic.py
+35 -15
View File
@@ -9,12 +9,19 @@ import time
import weakref
from contextlib import asynccontextmanager, contextmanager
from contextvars import ContextVar
from typing import Any, AsyncGenerator, Generator
from typing import TYPE_CHECKING, Any, AsyncGenerator, Generator
if TYPE_CHECKING:
from workflows.workflow import Workflow
from llama_index_instrumentation.dispatcher import active_instrument_tags
from workflows.context.serializers import BaseSerializer, JsonSerializer
from workflows.context.state_store import InMemoryStateStore, infer_state_type
from workflows.context.state_store import (
InMemoryStateStore,
StateStore,
infer_state_type,
)
from workflows.errors import WorkflowRuntimeError
from workflows.events import Event, StartEvent, StopEvent
from workflows.runtime.types.internal_state import BrokerState
@@ -25,6 +32,9 @@ from workflows.runtime.types.plugin import (
Runtime,
SnapshottableAdapter,
V2RuntimeCompatibilityShim,
WaitResult,
WaitResultTick,
WaitResultTimeout,
)
from workflows.runtime.types.step_function import (
as_step_worker_functions,
@@ -49,7 +59,7 @@ class AsyncioAdapterQueues:
self,
run_id: str,
init_state: BrokerState,
state_store: InMemoryStateStore[Any] | None = None,
state_store: StateStore[Any] | None = None,
):
self.run_id = run_id
self.init_state = init_state
@@ -91,28 +101,39 @@ class InternalAsyncioAdapter(InternalRunAdapter, SnapshottableAdapter):
def init_state(self) -> BrokerState:
return self._queues.init_state
async def wait_receive(self) -> WorkflowTick:
return await self._queues.receive_queue.get()
async def write_to_event_stream(self, event: Event) -> None:
self._queues.publish_queue.put_nowait(event)
async def get_now(self) -> float:
return time.monotonic()
async def sleep(self, seconds: float) -> None:
await asyncio.sleep(seconds)
async def send_event(self, tick: WorkflowTick) -> None:
self._queues.receive_queue.put_nowait(tick)
async def wait_receive(
self,
timeout_seconds: float | None = None,
) -> WaitResult:
"""Wait for tick with optional timeout using asyncio primitives."""
try:
if timeout_seconds is None:
tick = await self._queues.receive_queue.get()
else:
tick = await asyncio.wait_for(
self._queues.receive_queue.get(),
timeout=timeout_seconds,
)
return WaitResultTick(tick=tick)
except asyncio.TimeoutError:
return WaitResultTimeout()
def on_tick(self, tick: WorkflowTick) -> None:
self._queues.ticks.append(tick)
def replay(self) -> list[WorkflowTick]:
return self._queues.ticks
def get_state_store(self) -> InMemoryStateStore[Any] | None:
def get_state_store(self) -> StateStore[Any] | None:
return self._queues.state_store
@@ -149,16 +170,13 @@ class ExternalAsyncioAdapter(
if isinstance(item, StopEvent):
break
async def close(self) -> None:
pass
def on_tick(self, tick: WorkflowTick) -> None:
self._queues.ticks.append(tick)
def replay(self) -> list[WorkflowTick]:
return self._queues.ticks
def get_state_store(self) -> InMemoryStateStore[Any] | None:
def get_state_store(self) -> StateStore[Any] | None:
return self._queues.state_store
async def get_result(self) -> StopEvent:
@@ -281,7 +299,7 @@ class BasicRuntime(Runtime):
queues.complete = task
return self.get_external_adapter(run_id)
def get_internal_adapter(self) -> InternalRunAdapter:
def get_internal_adapter(self, workflow: Workflow) -> InternalRunAdapter:
run_id = get_current_run_id()
if run_id is None:
raise RuntimeError(
@@ -304,11 +322,13 @@ _current_run_id: ContextVar[str | None] = ContextVar("current_run_id", default=N
def get_current_run_id() -> str | None:
"""Get the current run ID, if set."""
return _current_run_id.get()
@contextmanager
def setting_run_id(run_id: str) -> Generator[None, None, None]:
"""Set the current run ID for the duration of the block."""
token = _current_run_id.set(run_id)
try:
yield
packages/llama-index-workflows/src/workflows/runtime/control_loop.py
+236 -88
View File
@@ -4,6 +4,7 @@
from __future__ import annotations
import asyncio
import heapq
import logging
import time
import traceback
@@ -45,8 +46,10 @@ from workflows.runtime.types.internal_state import (
InProgressState,
InternalStepWorkerState,
)
from workflows.runtime.types.named_task import NamedTask
from workflows.runtime.types.plugin import (
InternalRunAdapter,
WaitResultTick,
as_snapshottable_adapter,
get_current_run,
)
@@ -70,6 +73,18 @@ from workflows.runtime.types.ticks import (
)
from workflows.workflow import Workflow
async def _single_pull(adapter: InternalRunAdapter) -> WorkflowTick | None:
"""Single-iteration pull: calls wait_receive once and returns the tick.
Returns None if timeout (shouldn't happen with unbounded wait).
"""
wait_result = await adapter.wait_receive(None)
if isinstance(wait_result, WaitResultTick):
return wait_result.tick
return None
if TYPE_CHECKING:
from workflows.context.context import Context
from workflows.runtime.types.step_function import StepWorkerFunction
@@ -82,6 +97,11 @@ class _ControlLoopRunner:
"""
Private class to encapsulate the async control loop runtime state and behavior.
Keeps the pure transformation functions at module level for testability.
This control loop uses a sequential, deterministic design:
- Scheduled wakeups are tracked in a heap (for timeouts/delays)
- External events come via wait_receive
- No concurrent timeout tasks, ensuring deterministic DBOS function_id ordering
"""
def __init__(
@@ -97,33 +117,56 @@ class _ControlLoopRunner:
self.context = context
self.step_workers = step_workers
self.state = init_state
self.workers: list[asyncio.Task] = []
self.queue: asyncio.Queue[WorkflowTick] = asyncio.Queue()
self.worker_tasks: set[asyncio.Task[TickStepResult]] = set()
# Transient tick buffer - drained synchronously at start of each loop iteration
self.tick_buffer: list[WorkflowTick] = []
# Pending items to be processed (from rehydration or delayed ticks)
for tick in self.state.rehydrate_with_ticks():
self.queue.put_nowait(tick)
self.tick_buffer.append(tick)
# Scheduled wakeups: heap of (wakeup_time, sequence, tick) tuples
# The sequence counter ensures deterministic ordering when timestamps are equal,
# avoiding TypeError from comparing WorkflowTick objects that don't implement __lt__
self.scheduled_wakeups: list[tuple[float, int, WorkflowTick]] = []
self._wakeup_sequence = 0
self.snapshot_adapter = as_snapshottable_adapter(adapter)
# Pull task sequence counter for deterministic journaling
self._pull_sequence = 0
# Map from worker task to (step_name, worker_id) key
self._task_keys: dict[asyncio.Task[TickStepResult], tuple[str, int]] = {}
async def wait_for_tick(self) -> WorkflowTick:
"""Wait for the next tick from the internal queue."""
return await self.queue.get()
def schedule_tick(self, tick: WorkflowTick, at_time: float) -> None:
"""Schedule a tick to be processed at a specific time."""
seq = self._wakeup_sequence
self._wakeup_sequence += 1
heapq.heappush(self.scheduled_wakeups, (at_time, seq, tick))
def queue_tick(self, tick: WorkflowTick, delay: float | None = None) -> None:
"""Queue a tick event for processing, optionally after a delay."""
if delay:
def next_wakeup_timeout(self, now: float) -> float | None:
"""Calculate timeout until next scheduled wakeup.
async def _delayed_queue() -> None:
await self.adapter.sleep(delay)
self.queue.put_nowait(tick)
Returns None if no scheduled wakeups, otherwise returns
the number of seconds until the next scheduled tick is due.
"""
if not self.scheduled_wakeups:
return None
next_time, _, _ = self.scheduled_wakeups[0]
return max(0, next_time - now)
task = asyncio.create_task(_delayed_queue())
self.workers.append(task)
else:
self.queue.put_nowait(tick)
def pop_due_ticks(self, now: float) -> list[WorkflowTick]:
"""Pop all ticks that are due (scheduled time <= now)."""
due = []
while self.scheduled_wakeups and self.scheduled_wakeups[0][0] <= now:
_, _, tick = heapq.heappop(self.scheduled_wakeups)
due.append(tick)
return due
def run_worker(self, command: CommandRunWorker) -> None:
"""Run a worker for a step function."""
"""Run a worker for a step function.
async def _run_worker() -> None:
Step workers run concurrently as asyncio tasks. When they complete,
they return TickStepResult for the main loop to process via asyncio.wait.
"""
async def _run_worker() -> TickStepResult:
try:
worker = next(
(
@@ -146,46 +189,45 @@ class _ControlLoopRunner:
event=command.event,
workflow=self.workflow,
)
self.queue_tick(
TickStepResult(
step_name=command.step_name,
worker_id=command.id,
event=command.event,
result=result,
)
# Return result for main loop to process
return TickStepResult(
step_name=command.step_name,
worker_id=command.id,
event=command.event,
result=result,
)
except Exception as e:
logger.error("error running step worker function: %s", e, exc_info=True)
self.queue_tick(
TickStepResult(
step_name=command.step_name,
worker_id=command.id,
event=command.event,
result=[
StepWorkerFailed(
exception=e, failed_at=await self.adapter.get_now()
)
],
)
return TickStepResult(
step_name=command.step_name,
worker_id=command.id,
event=command.event,
result=[
StepWorkerFailed(
exception=e, failed_at=await self.adapter.get_now()
)
],
)
raise e
task = asyncio.create_task(_run_worker())
task.add_done_callback(lambda _: self.workers.remove(task))
self.workers.append(task)
# Track key separately for building NamedTask list
self._task_keys[task] = (command.step_name, command.id)
self.worker_tasks.add(task)
async def process_command(self, command: WorkflowCommand) -> None | StopEvent:
"""Process a single command returned from tick reduction."""
if isinstance(command, CommandQueueEvent):
self.queue_tick(
TickAddEvent(
event=command.event,
step_name=command.step_name,
attempts=command.attempts,
first_attempt_at=command.first_attempt_at,
),
delay=command.delay,
event = TickAddEvent(
event=command.event,
step_name=command.step_name,
attempts=command.attempts,
first_attempt_at=command.first_attempt_at,
)
if command.delay is not None and command.delay > 0:
now = await self.adapter.get_now()
self.schedule_tick(event, at_time=now + command.delay)
else:
self.tick_buffer.append(event)
return None
elif isinstance(command, CommandRunWorker):
self.run_worker(command)
@@ -207,25 +249,38 @@ class _ControlLoopRunner:
async def cleanup_tasks(self) -> None:
"""Cancel and cleanup all running worker tasks."""
# Signal adapter to stop waiting (wakes blocked DBOS.recv)
try:
for worker in self.workers:
worker.cancel()
await self.adapter.close()
except Exception:
pass
# Cancel worker tasks
for task in self.worker_tasks:
task.cancel()
try:
await asyncio.wait_for(
asyncio.gather(*self.workers, return_exceptions=True),
timeout=0.5,
)
if self.worker_tasks:
await asyncio.wait_for(
asyncio.gather(*self.worker_tasks, return_exceptions=True),
timeout=0.5,
)
except Exception:
pass
self.worker_tasks.clear()
self._task_keys.clear()
async def run(
self, start_event: Event | None = None, start_with_timeout: bool = True
) -> StopEvent:
"""
Run the control loop until completion.
This uses a sequential, deterministic design that combines timeout
handling with event waiting in a single operation, ensuring
deterministic DBOS function_id ordering for replay.
Args:
start_event: Optional initial event to process
start_with_timeout: Whether to start the timeout timer
@@ -234,28 +289,22 @@ class _ControlLoopRunner:
The final StopEvent from the workflow
"""
# Start external event listener
async def _pull() -> None:
while True:
tick = await self.adapter.wait_receive()
self.queue_tick(tick)
self.workers.append(asyncio.create_task(_pull()))
# Queue initial event and timeout
# Queue initial event
if start_event is not None:
self.queue_tick(TickAddEvent(event=start_event))
self.tick_buffer.append(TickAddEvent(event=start_event))
start = await self.adapter.get_now()
# Schedule workflow timeout if configured
if start_with_timeout and self.workflow._timeout is not None:
self.queue_tick(
# Get initial time
timeout_time = start + self.workflow._timeout
self.schedule_tick(
TickTimeout(timeout=self.workflow._timeout),
delay=self.workflow._timeout,
at_time=timeout_time,
)
# Resume any in-progress work
self.state, commands = rewind_in_progress(
self.state, await self.adapter.get_now()
)
self.state, commands = rewind_in_progress(self.state, start)
for command in commands:
try:
await self.process_command(command)
@@ -263,35 +312,130 @@ class _ControlLoopRunner:
await self.cleanup_tasks()
raise
# Initialize pull task (single-iteration)
pull_task: asyncio.Task[WorkflowTick | None] | None = None
# Main event loop
try:
while True:
tick = await self.wait_for_tick()
try:
self.state, commands = _reduce_tick(
tick, self.state, await self.adapter.get_now()
)
except Exception:
await self.cleanup_tasks()
logger.error(
"Unexpected error in internal control loop of workflow. This shouldn't happen. ",
exc_info=True,
)
raise
if self.snapshot_adapter is not None:
self.snapshot_adapter.on_tick(tick)
for command in commands:
try:
result = await self.process_command(command)
except Exception:
await self.cleanup_tasks()
raise
# Yield to let fire-and-forget tasks run (e.g., ctx.send_event)
await asyncio.sleep(0)
# Get current time
now = await self.adapter.get_now()
# optimization, only reload "now" if any work was done
was_buffered = bool(self.tick_buffer)
# Drain and process buffered ticks first (from rehydration, queue_tick, etc.)
while self.tick_buffer:
tick = self.tick_buffer.pop(0)
result = await self._process_tick(tick)
if result is not None:
return result
# optimization
if was_buffered:
now = await self.adapter.get_now()
# Calculate timeout for next scheduled wakeup
timeout = self.next_wakeup_timeout(now)
# Ensure pull_task exists
if pull_task is None:
pull_task = asyncio.create_task(_single_pull(self.adapter))
pull_sequence = self._pull_sequence
self._pull_sequence += 1
else:
# Retrieve the sequence from last time
pull_sequence = self._pull_sequence - 1
# Build list of NamedTasks with workers first (higher priority), then pull
named_tasks: list[NamedTask] = [
NamedTask.worker(key[0], key[1], task)
for task in self.worker_tasks
for key in [self._task_keys.get(task)]
if key is not None
]
named_tasks.append(NamedTask.pull(pull_sequence, pull_task))
# Wait for next task completion (adapter controls ordering for replay)
completed_task = await self.adapter.wait_for_next_task(
named_tasks, timeout
)
if completed_task is None:
# Timeout - process scheduled ticks
now = await self.adapter.get_now()
for due_tick in self.pop_due_ticks(now):
self.tick_buffer.append(due_tick)
continue
# Process the single completed task
if completed_task is pull_task:
# Pull task completed
try:
pull_tick = completed_task.result()
except asyncio.CancelledError:
pull_task = None
except Exception:
logger.exception("Pull task failed", exc_info=True)
pull_task = None
else:
pull_task = None
if pull_tick is not None:
self.tick_buffer.append(pull_tick)
else:
# Worker task completed
self.worker_tasks.discard(completed_task)
self._task_keys.pop(completed_task, None)
try:
tick_result = completed_task.result()
except asyncio.CancelledError:
pass
except Exception:
logger.exception(
"Worker task failed unexpectedly", exc_info=True
)
else:
self.tick_buffer.append(tick_result)
finally:
# Cancel pull task if running
if pull_task is not None:
pull_task.cancel()
try:
await pull_task
except (asyncio.CancelledError, Exception):
pass
await self.cleanup_tasks()
async def _process_tick(self, tick: WorkflowTick) -> StopEvent | None:
"""Process a single tick and return StopEvent if workflow completes."""
try:
start = await self.adapter.get_now()
self.state, commands = _reduce_tick(tick, self.state, start)
except Exception:
await self.cleanup_tasks()
logger.error(
"Unexpected error in internal control loop of workflow. This shouldn't happen. ",
exc_info=True,
)
raise
if self.snapshot_adapter is not None:
self.snapshot_adapter.on_tick(tick)
for command in commands:
try:
result = await self.process_command(command)
except Exception:
await self.cleanup_tasks()
raise
if result is not None:
return result
return None
async def control_loop(
start_event: Event | None,
@@ -434,6 +578,10 @@ def _process_step_result_tick(
CommandPublishEvent(event=result.result)
) # stop event always published to the stream
state.is_running = False
# Clear collected_events and collected_waiters since workflow is complete
for worker in state.workers.values():
worker.collected_events.clear()
worker.collected_waiters.clear()
commands.append(CommandCompleteRun(result=result.result))
elif isinstance(result.result, Event):
# queue any subsequent events
packages/llama-index-workflows/src/workflows/runtime/types/named_task.py
+81
View File
@@ -0,0 +1,81 @@
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 LlamaIndex Inc.
"""NamedTask associates asyncio tasks with stable string keys for journaling."""
from __future__ import annotations
from asyncio import Task
from dataclasses import dataclass
from typing import Any
# Key prefix for pull tasks
PULL_PREFIX = "__pull__"
@dataclass
class NamedTask:
"""An asyncio task with a stable string key for identification.
Keys are strings like "step_name:worker_id" for workers or "__pull__:0" for pull.
"""
key: str
task: Task[Any]
@staticmethod
def worker(step_name: str, worker_id: int, task: Task[Any]) -> NamedTask:
"""Create a NamedTask for a worker."""
return NamedTask(f"{step_name}:{worker_id}", task)
@staticmethod
def pull(sequence: int, task: Task[Any]) -> NamedTask:
"""Create a NamedTask for a pull task."""
return NamedTask(f"{PULL_PREFIX}:{sequence}", task)
def is_pull(self) -> bool:
"""Check if this is a pull task."""
return self.key.startswith(f"{PULL_PREFIX}:")
@staticmethod
def all_tasks(named_tasks: list[NamedTask]) -> set[Task[Any]]:
"""Extract all tasks for use with asyncio.wait."""
return {nt.task for nt in named_tasks}
@staticmethod
def find_by_key(named_tasks: list[NamedTask], key: str) -> Task[Any] | None:
"""Find a task by its key, returns None if not found."""
for nt in named_tasks:
if nt.key == key:
return nt.task
return None
@staticmethod
def get_key(named_tasks: list[NamedTask], task: Task[Any]) -> str:
"""Get the key for a task. Raises KeyError if not found."""
for nt in named_tasks:
if nt.task is task:
return nt.key
raise KeyError(f"Task {task} not found")
@staticmethod
def pick_highest_priority(
named_tasks: list[NamedTask], done: set[Task[Any]]
) -> Task[Any] | None:
"""Return highest priority completed task from done set.
Priority is determined by list order - tasks earlier in the list
have higher priority. Workers should be listed before pull.
Returns None if done is empty.
Raises ValueError if done is non-empty but no tasks match (indicates bug).
"""
if not done:
return None
for nt in named_tasks:
if nt.task in done:
return nt.task
raise ValueError(
f"No tasks in done set match named_tasks. "
f"done={done}, named_tasks={[nt.key for nt in named_tasks]}"
)
packages/llama-index-workflows/src/workflows/runtime/types/plugin.py
+121 -40
View File
@@ -6,6 +6,7 @@ A runtime interface to switch out a broker runtime (external library or service
from __future__ import annotations
import asyncio
from abc import ABC, abstractmethod
from contextlib import contextmanager
from contextvars import ContextVar, Token
@@ -16,13 +17,17 @@ from typing import (
AsyncGenerator,
Coroutine,
Generator,
Literal,
Protocol,
Union,
)
from workflows.context.state_store import StateStore
from workflows.runtime.types.named_task import NamedTask
if TYPE_CHECKING:
from workflows.context.context import Context
from workflows.context.serializers import BaseSerializer
from workflows.context.state_store import InMemoryStateStore
from workflows.runtime.types.internal_state import BrokerState
from workflows.runtime.types.step_function import StepWorkerFunction
from workflows.workflow import Workflow
@@ -36,6 +41,24 @@ _current_runtime: ContextVar[Runtime | None] = ContextVar(
)
@dataclass
class WaitResultTick:
"""Result containing a received tick."""
tick: WorkflowTick
type: Literal["tick"] = "tick"
@dataclass
class WaitResultTimeout:
"""Result indicating timeout expiration."""
type: Literal["timeout"] = "timeout"
WaitResult = Union[WaitResultTick, WaitResultTimeout]
@dataclass
class RegisteredWorkflow:
workflow: Workflow
@@ -72,17 +95,6 @@ class InternalRunAdapter(ABC):
"""
...
@abstractmethod
async def wait_receive(self) -> WorkflowTick:
"""
Wait for the next tick from the mailbox.
Called from inside the workflow control loop to receive the next tick
event that was sent via the external adapter's send_event().
This method blocks until a tick is available.
"""
...
@abstractmethod
async def write_to_event_stream(self, event: Event) -> None:
"""
@@ -104,18 +116,6 @@ class InternalRunAdapter(ABC):
"""
...
@abstractmethod
async def sleep(self, seconds: float) -> None:
"""
Sleep for a given number of seconds with durability support.
Called from within the workflow control loop. For durable runtimes,
this integrates with the host runtime to allow workflow suspension
and resumption. Note that other tasks in the control loop may still
run simultaneously during the sleep.
"""
...
@abstractmethod
async def send_event(self, tick: WorkflowTick) -> None:
"""
@@ -127,7 +127,41 @@ class InternalRunAdapter(ABC):
"""
...
def get_state_store(self) -> InMemoryStateStore[Any] | None:
@abstractmethod
async def wait_receive(
self,
timeout_seconds: float | None = None,
) -> WaitResult:
"""
Wait for next tick OR timeout expiration.
This is the primary method for the control loop to wait for events.
It combines receiving ticks and timeout handling into a single
deterministic operation.
Args:
timeout_seconds: Max time to wait. None means wait indefinitely.
Returns:
WaitResultTick if a tick was received
WaitResultTimeout if timeout expired before receiving tick
This is a DURABLE operation for durable runtimes:
- On replay, already-elapsed time is accounted for
- If timeout already expired in previous run, returns immediately
"""
...
async def close(self) -> None:
"""
Signal shutdown to wake any blocked wait operations.
Called during cleanup to allow the adapter to exit gracefully.
Default is no-op. DBOS adapter sends a shutdown signal to wake blocked recv.
"""
pass
def get_state_store(self) -> StateStore[Any] | None:
"""
Get the state store for this workflow run.
@@ -136,6 +170,56 @@ class InternalRunAdapter(ABC):
"""
return None
async def finalize_step(self) -> None:
"""
Called after a step function completes to perform any adapter-specific cleanup.
This is called after all background tasks spawned during the step have completed.
Adapters can override to perform additional finalization (e.g., flush buffers,
sync state). Default is no-op.
"""
pass
async def wait_for_next_task(
self,
task_set: list[NamedTask],
timeout: float | None = None,
) -> asyncio.Task[Any] | None:
"""Wait for and return the next task that should complete.
Args:
task_set: List of NamedTasks with stable string keys for identification.
The order indicates priority - first items should be returned first
when multiple tasks complete simultaneously.
timeout: Timeout in seconds, None for no timeout
Returns:
The completed task, or None on timeout.
IMPORTANT: Must return at most ONE task per call.
Default implementation uses asyncio.wait(FIRST_COMPLETED) and returns
the highest-priority completed task (workers before pull).
DBOS overrides to coordinate based on journal for deterministic replay,
using the stable keys from NamedTask to identify tasks.
"""
tasks = NamedTask.all_tasks(task_set)
if not tasks:
return None
done, _ = await asyncio.wait(
tasks,
timeout=timeout,
return_when=asyncio.FIRST_COMPLETED,
)
if not done:
return None
# Return the highest-priority completed task (first in task_set order)
for named_task in task_set:
if named_task.task in done:
return named_task.task
# Fallback (shouldn't happen)
return done.pop()
class ExternalRunAdapter(ABC):
"""
@@ -186,7 +270,6 @@ class ExternalRunAdapter(ABC):
"""
...
@abstractmethod
async def close(self) -> None:
"""
Clean up adapter resources.
@@ -195,7 +278,7 @@ class ExternalRunAdapter(ABC):
resources held by this adapter (e.g., close streams, release locks).
"""
...
pass
@abstractmethod
async def get_result(self) -> StopEvent:
@@ -210,7 +293,7 @@ class ExternalRunAdapter(ABC):
"""
await self.send_event(TickCancelRun())
def get_state_store(self) -> InMemoryStateStore[Any] | None:
def get_state_store(self) -> StateStore[Any] | None:
"""
Get the state store for this workflow run.
@@ -256,8 +339,7 @@ class Runtime(ABC):
Abstract base class for workflow execution runtimes.
Runtimes control how workflows are registered, launched, and executed.
The default BasicRuntime uses asyncio; other runtimes can add durability
or distributed execution.
The default BasicRuntime uses asyncio; Other's plug into their own durability and distributed execution models.
Lifecycle:
1. Create runtime instance
@@ -284,7 +366,7 @@ class Runtime(ABC):
Register a workflow with the runtime.
Called at launch() time for each tracked workflow. Runtimes can
wrap the control_loop and steps with their own decorators or handlers.
wrap the control_loop and steps to fit in their registration/decoration model.
Returns RegisteredWorkflow with wrapped functions
"""
@@ -298,7 +380,7 @@ class Runtime(ABC):
init_state: BrokerState,
start_event: StartEvent | None = None,
serialized_state: dict[str, Any] | None = None,
serializer: "BaseSerializer | None" = None,
serializer: BaseSerializer | None = None,
) -> ExternalRunAdapter:
"""
Launch a workflow run.
@@ -317,12 +399,14 @@ class Runtime(ABC):
...
@abstractmethod
def get_internal_adapter(self) -> InternalRunAdapter:
def get_internal_adapter(self, workflow: "Workflow") -> InternalRunAdapter:
"""
Get the internal adapter for a workflow run.
Called on each workflow.run() to instantiate an interface for the workflow run internals to communicite with the runtime.
The workflow run must be derived from the runtime set context.
Called on each workflow.run() to instantiate an interface for the workflow run internals to communicate with the runtime.
Args:
workflow: The workflow instance being run. Used by runtimes to access workflow metadata (e.g., state type).
"""
...
@@ -343,10 +427,7 @@ class Runtime(ABC):
"""
Launch the runtime and register all tracked workflows.
For BasicRuntime, this is a no-op. Other runtimes may wrap workflows
with decorators and initialize backend connections.
Must be called before running workflows.
For many runtime's, this must be called before running workflows.
"""
pass
@@ -363,7 +444,7 @@ class Runtime(ABC):
Track a workflow instance for registration at launch time.
Called by Workflow.__init__ to register with the runtime.
Override in runtimes that need to track workflows for deferred registration.
Override in runtimes that need to track workflows.
Default implementation is a no-op.
"""
pass
packages/llama-index-workflows/src/workflows/runtime/types/step_function.py
+3 -1
View File
@@ -151,6 +151,8 @@ def as_step_worker_function(func: Callable[P, Awaitable[R]]) -> StepWorkerFuncti
returns.return_values.append(
StepWorkerFailed(exception=e, failed_at=time.time())
)
await internal_context._finalize_step()
return returns.return_values
finally:
try:
@@ -187,7 +189,7 @@ def create_workflow_run_function(
registered = workflow._runtime.get_or_register(workflow)
# Set run_id context before creating internal context
internal_ctx = Context._create_internal(workflow=workflow)
internal_adapter = workflow._runtime.get_internal_adapter()
internal_adapter = workflow._runtime.get_internal_adapter(workflow)
with instrument_tags(tags):
# defer execution to make sure the task can be captured and passed
# to the handler as async exception, protecting against exceptions from before_start
packages/llama-index-workflows/src/workflows/runtime/workflow_tracker.py
-75
View File
@@ -1,75 +0,0 @@
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 LlamaIndex Inc.
"""Workflow tracker for runtime instance registration."""
from __future__ import annotations
from workflows.runtime.types.plugin import RegisteredWorkflow
from workflows.workflow import Workflow
class WorkflowTracker:
"""
Tracks workflow instances registered with a runtime.
Used by runtimes to collect workflows before launch() and to
look up registered workflows at execution time.
Uses strong references to ensure workflows survive until explicitly cleared.
"""
def __init__(self) -> None:
# Workflows registered before launch (strong refs to survive GC)
self._pending: list[Workflow] = []
self._pending_set: set[int] = set() # Track by id for dedup
# Registered workflows after launch (control loop + steps)
self._registered: dict[int, RegisteredWorkflow] = {} # keyed by id(workflow)
self._launched: bool = False
def add(self, workflow: Workflow) -> None:
"""Add a workflow to be registered at launch time."""
if self._launched:
raise RuntimeError(
"Cannot add workflows after launch(). "
"Create workflows before calling runtime.launch()."
)
wf_id = id(workflow)
if wf_id not in self._pending_set:
self._pending.append(workflow)
self._pending_set.add(wf_id)
def remove(self, workflow: Workflow) -> None:
"""Remove a workflow from pending registration."""
wf_id = id(workflow)
self._pending = [wf for wf in self._pending if id(wf) != wf_id]
self._pending_set.discard(wf_id)
def get_pending(self) -> list[Workflow]:
"""Get all pending workflows."""
return list(self._pending)
def mark_launched(self) -> None:
"""Mark that launch() has been called."""
self._launched = True
@property
def is_launched(self) -> bool:
"""Whether launch() has been called."""
return self._launched
def set_registered(
self, workflow: Workflow, registered: RegisteredWorkflow
) -> None:
"""Store the registered workflow (wrapped control loop + steps)."""
self._registered[id(workflow)] = registered
def get_registered(self, workflow: Workflow) -> RegisteredWorkflow | None:
"""Get the registered workflow if available."""
return self._registered.get(id(workflow))
def clear(self) -> None:
"""Clear all tracking state (for destroy())."""
self._pending.clear()
self._pending_set.clear()
self._registered.clear()
self._launched = False
packages/llama-index-workflows/src/workflows/workflow.py
+6 -1
View File
@@ -413,6 +413,9 @@ class Workflow(metaclass=WorkflowMeta):
# Validate the workflow
self._validate()
# Extract run_id before passing remaining kwargs to start event
run_id = kwargs.pop("run_id", None)
# If a previous context is provided, pass its serialized form
ctx = ctx if ctx is not None else Context(self)
# TODO(v3) - remove dependency on is running for choosing whether to send a StartEvent.
@@ -422,7 +425,9 @@ class Workflow(metaclass=WorkflowMeta):
if ctx.is_running
else self._get_start_event_instance(start_event, **kwargs)
)
return ctx._workflow_run(workflow=self, start_event=start_event_instance)
return ctx._workflow_run(
workflow=self, start_event=start_event_instance, run_id=run_id
)
def _validate_resource_configs(self) -> list[str]:
"""Validate all resource configs (including nested ones) by loading them."""
packages/llama-index-workflows/tests/context/test_context.py
+233 -12
View File
@@ -17,7 +17,12 @@ from workflows.context.context import (
_warn_is_running_in_step,
)
from workflows.context.external_context import ExternalContext
from workflows.context.state_store import DictState, InMemoryStateStore
from workflows.context.serializers import JsonSerializer
from workflows.context.state_store import (
DictState,
InMemoryStateStore,
deserialize_state_from_dict,
)
from workflows.decorators import step
from workflows.errors import ContextStateError, WorkflowRuntimeError
from workflows.events import (
@@ -27,7 +32,11 @@ from workflows.events import (
StartEvent,
StopEvent,
)
from workflows.plugins.basic import AsyncioAdapterQueues, BasicRuntime, setting_run_id
from workflows.plugins.basic import (
AsyncioAdapterQueues,
BasicRuntime,
setting_run_id,
)
from workflows.runtime.types.internal_state import BrokerState
from workflows.runtime.types.ticks import TickAddEvent
from workflows.testing import WorkflowTestRunner
@@ -186,19 +195,41 @@ async def test_get_not_found(internal_ctx: Context) -> None:
@pytest.mark.asyncio
async def test_send_event_step_is_none(workflow: Workflow) -> None:
async def test_send_event_step_is_none() -> None:
"""Test that external events create TickAddEvent with step_name=None.
Uses a workflow that waits for the external event so we can verify
the tick is logged before the workflow completes.
"""
ev = Event(foo="bar")
# Create a fresh context and run workflow
ctx = Context(workflow)
handler = ctx._workflow_run(workflow, start_event=StartEvent())
class WaitingWorkflow(Workflow):
@step
async def wait_for_external(self, ctx: Context, ev: StartEvent) -> StopEvent:
# Wait for an external Event to arrive
result = await ctx.wait_for_event(Event, requirements={"foo": "bar"})
return StopEvent(result=result.foo)
wf = WaitingWorkflow()
handler = wf.run()
try:
# Send the external event
handler.ctx.send_event(ev)
await asyncio.sleep(0.01)
# handler.ctx is a new external context
external_face = handler.ctx._face
assert isinstance(external_face, ExternalContext)
replay = external_face._tick_log
assert TickAddEvent(event=ev, step_name=None) in replay
# Wait for event to appear in tick log (up to 1 second)
expected_tick = TickAddEvent(event=ev, step_name=None)
for _ in range(100):
if expected_tick in external_face._tick_log:
break
await asyncio.sleep(0.01)
assert expected_tick in external_face._tick_log
# Let workflow complete
result = await handler
assert result == "bar"
finally:
external_face = handler.ctx._face
assert isinstance(external_face, ExternalContext)
@@ -394,7 +425,8 @@ async def test_wait_for_multiple_events_in_workflow() -> None:
)
result = await handler
assert result == ["foo", "bar"]
# Order is non-deterministic since waiters run concurrently
assert sorted(result) == ["bar", "foo"]
assert not handler.ctx.is_running
# serialize and resume
@@ -414,7 +446,8 @@ async def test_wait_for_multiple_events_in_workflow() -> None:
)
result = await handler
assert result == ["fizz", "buzz"]
# Order is non-deterministic since waiters run concurrently
assert sorted(result) == ["buzz", "fizz"]
@pytest.mark.asyncio
@@ -681,3 +714,191 @@ async def test_get_result_pre_context_raises(workflow: Workflow) -> None:
with pytest.warns(DeprecationWarning): # get_result is deprecated
with pytest.raises(ContextStateError, match="requires a running workflow"):
ctx.get_result()
# ============================================================================
# deserialize_state_from_dict Tests
# ============================================================================
class TypedTestState(BaseModel):
"""Typed state for deserialize_state_from_dict testing."""
counter: int = 0
name: str = "default"
@pytest.mark.asyncio
async def test_deserialize_state_from_dict_with_dict_state() -> None:
"""Test deserializing DictState from to_dict() format."""
serializer = JsonSerializer()
# Create state and serialize it
store = InMemoryStateStore(DictState())
await store.set("counter", 42)
await store.set("name", "test-value")
serialized = store.to_dict(serializer)
# Deserialize
result = deserialize_state_from_dict(serialized, serializer)
assert isinstance(result, DictState)
assert result["counter"] == 42
assert result["name"] == "test-value"
def test_deserialize_state_from_dict_with_typed_state() -> None:
"""Test deserializing typed Pydantic model from to_dict() format."""
serializer = JsonSerializer()
# Create typed state and serialize it
initial = TypedTestState(counter=100, name="typed-test")
store = InMemoryStateStore(initial)
serialized = store.to_dict(serializer)
# Deserialize
result = deserialize_state_from_dict(serialized, serializer)
assert isinstance(result, TypedTestState)
assert result.counter == 100
assert result.name == "typed-test"
def test_deserialize_state_from_dict_empty_dict_state() -> None:
"""Test deserializing empty DictState."""
serializer = JsonSerializer()
serialized = {
"state_data": {"_data": {}},
"state_type": "DictState",
"state_module": "workflows.context.state_store",
}
result = deserialize_state_from_dict(serialized, serializer)
assert isinstance(result, DictState)
assert len(list(result.items())) == 0
def test_deserialize_state_from_dict_defaults_to_dict_state() -> None:
"""Test that missing state_type defaults to DictState."""
serializer = JsonSerializer()
serialized = {"state_data": {"_data": {}}}
result = deserialize_state_from_dict(serialized, serializer)
assert isinstance(result, DictState)
# ============================================================================
# Serialized State Format Tests (parse_in_memory_state)
# ============================================================================
def test_parse_in_memory_state_old_format_no_store_type() -> None:
"""Test that old format (no store_type) parses as InMemorySerializedState."""
from workflows.context.state_store import (
InMemorySerializedState,
parse_in_memory_state,
)
# Old format without store_type field
old_format = {
"state_type": "DictState",
"state_module": "workflows.context.state_store",
"state_data": {"_data": {"counter": 42}},
}
result = parse_in_memory_state(old_format)
assert isinstance(result, InMemorySerializedState)
assert result.store_type == "in_memory"
assert result.state_type == "DictState"
assert result.state_module == "workflows.context.state_store"
assert result.state_data == {"_data": {"counter": 42}}
def test_parse_in_memory_state_explicit_in_memory() -> None:
"""Test that explicit store_type='in_memory' parses as InMemorySerializedState."""
from workflows.context.state_store import (
InMemorySerializedState,
parse_in_memory_state,
)
serialized = {
"store_type": "in_memory",
"state_type": "CustomState",
"state_module": "myapp.models",
"state_data": {"name": "test", "value": 123},
}
result = parse_in_memory_state(serialized)
assert isinstance(result, InMemorySerializedState)
assert result.store_type == "in_memory"
assert result.state_type == "CustomState"
assert result.state_module == "myapp.models"
assert result.state_data == {"name": "test", "value": 123}
def test_parse_in_memory_state_rejects_sql_store_type() -> None:
"""Test that store_type='sql' raises ValueError."""
from workflows.context.state_store import parse_in_memory_state
serialized = {
"store_type": "sql",
"run_id": "run-12345",
"state_type": "WorkflowState",
"state_module": "myapp.states",
"schema": "public",
}
with pytest.raises(ValueError, match="Cannot parse store_type 'sql'"):
parse_in_memory_state(serialized)
def test_parse_in_memory_state_unknown_store_type_raises() -> None:
"""Test that unknown store_type raises ValueError."""
from workflows.context.state_store import parse_in_memory_state
serialized = {
"store_type": "redis", # Unknown store type
"state_type": "SomeState",
"state_module": "some.module",
}
with pytest.raises(ValueError, match="Cannot parse store_type 'redis'"):
parse_in_memory_state(serialized)
# ============================================================================
# InMemoryStateStore Serialization Tests
# ============================================================================
def test_in_memory_state_store_to_dict_includes_store_type() -> None:
"""Test that to_dict() includes store_type='in_memory'."""
store = InMemoryStateStore(DictState())
serializer = JsonSerializer()
result = store.to_dict(serializer)
assert result["store_type"] == "in_memory"
assert "state_type" in result
assert "state_module" in result
assert "state_data" in result
def test_in_memory_state_store_from_dict_rejects_sql_format() -> None:
"""Test that from_dict() rejects SQL format with clear error."""
sql_format = {
"store_type": "sql",
"run_id": "run-12345",
"state_type": "DictState",
"state_module": "workflows.context.state_store",
}
serializer = JsonSerializer()
with pytest.raises(ValueError, match="Cannot parse store_type 'sql'"):
InMemoryStateStore.from_dict(sql_format, serializer)
packages/llama-index-workflows/tests/plugins/__init__.py
View File
packages/llama-index-workflows/tests/runtime/conftest.py
+27 -6
View File
@@ -25,6 +25,9 @@ from workflows.runtime.types.plugin import (
Runtime,
SnapshottableAdapter,
V2RuntimeCompatibilityShim,
WaitResult,
WaitResultTick,
WaitResultTimeout,
)
from workflows.runtime.types.step_function import (
as_step_worker_functions,
@@ -48,7 +51,7 @@ class MockRuntime(Runtime):
steps=as_step_worker_functions(workflow),
)
def get_internal_adapter(self) -> InternalRunAdapter:
def get_internal_adapter(self, workflow: Workflow) -> InternalRunAdapter:
run_id = get_current_run_id() or self._current_run_id or "test"
if run_id not in self._adapters:
self._adapters[run_id] = MockRunAdapter(run_id)
@@ -135,9 +138,6 @@ class MockRunAdapter(
"""
pass
async def wait_receive(self) -> WorkflowTick:
return await self._external_queue.get()
async def write_to_event_stream(self, event: Event) -> None:
await self._event_stream.put(event)
@@ -156,8 +156,29 @@ class MockRunAdapter(
return time.time()
return self._current_time
async def sleep(self, seconds: float) -> None:
await asyncio.sleep(seconds)
async def wait_receive(
self,
timeout_seconds: float | None = None,
) -> WaitResult:
"""Wait for tick with optional timeout.
When a timeout occurs, advances mock time by the timeout duration
to ensure scheduled ticks become due.
"""
try:
if timeout_seconds is None:
tick = await self._external_queue.get()
else:
tick = await asyncio.wait_for(
self._external_queue.get(),
timeout=timeout_seconds,
)
return WaitResultTick(tick=tick)
except asyncio.TimeoutError:
# Advance mock time when timeout occurs
if timeout_seconds is not None:
self.advance_time(timeout_seconds)
return WaitResultTimeout()
def advance_time(self, seconds: float) -> None:
if self._traveller is not None:
packages/llama-index-workflows/tests/runtime/test_control_loop.py
+85 -1
View File
@@ -278,13 +278,16 @@ async def test_control_loop_with_external_event(
@pytest.mark.asyncio
async def test_control_loop_timeout(test_plugin: MockRunAdapter) -> None:
async def test_control_loop_timeout(
test_plugin_with_time_machine: tuple[MockRunAdapter, time_machine.Coordinates],
) -> None:
"""
Test that workflow timeout raises WorkflowTimeoutError and publishes WorkflowTimedOutEvent.
When a workflow times out, a WorkflowTimedOutEvent should be published to the stream
to inform consumers about the timeout before the exception is raised.
"""
test_plugin, _ = test_plugin_with_time_machine
class SlowWorkflow(Workflow):
@step
@@ -927,3 +930,84 @@ async def test_control_loop_idle_event_not_emitted_on_completion(
assert len(idle_events) == 0, (
"WorkflowIdleEvent should not be emitted when workflow completes normally"
)
@pytest.mark.asyncio
async def test_simultaneous_retries_with_same_delay(
test_plugin_with_time_machine: tuple[MockRunAdapter, time_machine.Coordinates],
) -> None:
"""
Test that the control loop handles multiple retries scheduled at the same timestamp.
When two steps both fail and have the same retry delay, they get scheduled
at exactly the same timestamp. Without a sequence counter tiebreaker in the
heap, Python's heapq would compare WorkflowTick objects directly, causing
TypeError since they don't implement __lt__.
This test uses a CoarseTimeAdapter that rounds timestamps to 1-second precision,
ensuring that retries scheduled within the same second will collide.
"""
base_plugin, traveller = test_plugin_with_time_machine
class CoarseTimeAdapter(MockRunAdapter):
"""Adapter that rounds get_now() to 1-second precision to force collisions."""
async def get_now(self) -> float:
# Round to nearest second to force timestamp collisions
return float(int(time.time()))
test_plugin = CoarseTimeAdapter(run_id="test", traveller=traveller)
test_plugin.set_state_store(InMemoryStateStore(DictState()))
# Use a delay that's less than 1 second so both retries land on same rounded second
retry_delay = 0.01
class ResultA(Event):
pass
class ResultB(Event):
pass
class TwoStepsFailOnceWorkflow(Workflow):
step_a_attempts = 0
step_b_attempts = 0
@step(
retry_policy=ConstantDelayRetryPolicy(maximum_attempts=2, delay=retry_delay)
)
async def step_a(self, ev: StartEvent) -> ResultA:
self.step_a_attempts += 1
if self.step_a_attempts == 1:
raise RuntimeError("step_a fails once")
return ResultA()
@step(
retry_policy=ConstantDelayRetryPolicy(maximum_attempts=2, delay=retry_delay)
)
async def step_b(self, ev: StartEvent) -> ResultB:
self.step_b_attempts += 1
if self.step_b_attempts == 1:
raise RuntimeError("step_b fails once")
return ResultB()
@step
async def collector(
self, ev: Union[ResultA, ResultB], ctx: Context
) -> Optional[StopEvent]:
events = ctx.collect_events(ev, [ResultA, ResultB])
if events is None:
return None
return StopEvent(result="both_succeeded")
wf = TwoStepsFailOnceWorkflow(timeout=5.0)
result = await run_control_loop(
workflow=wf,
start_event=StartEvent(),
test_runtime=test_plugin,
)
assert isinstance(result, StopEvent)
assert result.result == "both_succeeded"
assert wf.step_a_attempts == 2
assert wf.step_b_attempts == 2
packages/llama-index-workflows/tests/runtime/test_named_task.py
+413
View File
@@ -0,0 +1,413 @@
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 LlamaIndex Inc.
"""Tests for NamedTask.
NamedTask associates asyncio tasks with stable string keys, providing:
- Task identification for DBOS journaling
- Task lookup by key for replay scenarios
- Priority-based task selection (by list order)
"""
from __future__ import annotations
import asyncio
from typing import Any
import pytest
from workflows.runtime.types.named_task import PULL_PREFIX, NamedTask
async def _never_completes() -> None:
"""Coroutine that never completes, for creating pending tasks."""
await asyncio.Future()
def create_pending_task() -> asyncio.Task[Any]:
"""Create a pending task that never completes."""
return asyncio.create_task(_never_completes())
# --- NamedTask creation ---
async def test_named_task_worker_creates_correct_key() -> None:
"""NamedTask.worker should create key as 'step_name:worker_id'."""
task = create_pending_task()
try:
nt = NamedTask.worker("my_step", 42, task)
assert nt.key == "my_step:42"
assert nt.task is task
assert not nt.is_pull()
finally:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
async def test_named_task_pull_creates_correct_key() -> None:
"""NamedTask.pull should create key as '__pull__:sequence'."""
task = create_pending_task()
try:
nt = NamedTask.pull(7, task)
assert nt.key == f"{PULL_PREFIX}:7"
assert nt.task is task
assert nt.is_pull()
finally:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
# --- all_tasks ---
async def test_all_tasks_returns_set_of_tasks() -> None:
"""all_tasks should return a set of all tasks."""
w1 = create_pending_task()
w2 = create_pending_task()
pull = create_pending_task()
try:
named_tasks = [
NamedTask.worker("step_a", 0, w1),
NamedTask.worker("step_b", 0, w2),
NamedTask.pull(0, pull),
]
all_tasks = NamedTask.all_tasks(named_tasks)
assert len(all_tasks) == 3
assert w1 in all_tasks
assert w2 in all_tasks
assert pull in all_tasks
finally:
for t in [w1, w2, pull]:
t.cancel()
try:
await t
except asyncio.CancelledError:
pass
async def test_all_tasks_empty_list() -> None:
"""all_tasks should return empty set for empty list."""
assert NamedTask.all_tasks([]) == set()
async def test_all_tasks_works_with_asyncio_wait() -> None:
"""all_tasks result should work with asyncio.wait."""
task = create_pending_task()
try:
named_tasks = [NamedTask.worker("step", 0, task)]
all_tasks = NamedTask.all_tasks(named_tasks)
# Should not raise - set is valid for asyncio.wait
done, pending = await asyncio.wait(all_tasks, timeout=0.001)
assert task in pending
finally:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
# --- find_by_key ---
async def test_find_by_key_returns_worker_task() -> None:
"""find_by_key should return the correct worker task."""
w1 = create_pending_task()
w2 = create_pending_task()
try:
named_tasks = [
NamedTask.worker("step_a", 0, w1),
NamedTask.worker("step_b", 1, w2),
]
assert NamedTask.find_by_key(named_tasks, "step_a:0") is w1
assert NamedTask.find_by_key(named_tasks, "step_b:1") is w2
finally:
for t in [w1, w2]:
t.cancel()
try:
await t
except asyncio.CancelledError:
pass
async def test_find_by_key_returns_pull_task() -> None:
"""find_by_key should return the pull task."""
pull = create_pending_task()
try:
named_tasks = [NamedTask.pull(5, pull)]
assert NamedTask.find_by_key(named_tasks, f"{PULL_PREFIX}:5") is pull
finally:
pull.cancel()
try:
await pull
except asyncio.CancelledError:
pass
async def test_find_by_key_returns_none_for_unknown() -> None:
"""find_by_key should return None for unknown key."""
task = create_pending_task()
try:
named_tasks = [NamedTask.worker("step", 0, task)]
assert NamedTask.find_by_key(named_tasks, "unknown:99") is None
finally:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
async def test_find_by_key_empty_list() -> None:
"""find_by_key should return None for empty list."""
assert NamedTask.find_by_key([], "any:key") is None
# --- get_key ---
async def test_get_key_returns_worker_key() -> None:
"""get_key should return the key for a worker task."""
task = create_pending_task()
try:
named_tasks = [NamedTask.worker("my_step", 3, task)]
assert NamedTask.get_key(named_tasks, task) == "my_step:3"
finally:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
async def test_get_key_returns_pull_key() -> None:
"""get_key should return the key for a pull task."""
task = create_pending_task()
try:
named_tasks = [NamedTask.pull(2, task)]
assert NamedTask.get_key(named_tasks, task) == f"{PULL_PREFIX}:2"
finally:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
async def test_get_key_raises_for_unknown_task() -> None:
"""get_key should raise KeyError for unknown task."""
known = create_pending_task()
unknown = create_pending_task()
try:
named_tasks = [NamedTask.worker("step", 0, known)]
with pytest.raises(KeyError):
NamedTask.get_key(named_tasks, unknown)
finally:
for t in [known, unknown]:
t.cancel()
try:
await t
except asyncio.CancelledError:
pass
# --- get_key / find_by_key round trip ---
async def test_round_trip_worker() -> None:
"""get_key and find_by_key should be inverses for workers."""
task = create_pending_task()
try:
named_tasks = [NamedTask.worker("step", 5, task)]
key = NamedTask.get_key(named_tasks, task)
found = NamedTask.find_by_key(named_tasks, key)
assert found is task
finally:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
async def test_round_trip_pull() -> None:
"""get_key and find_by_key should be inverses for pull."""
task = create_pending_task()
try:
named_tasks = [NamedTask.pull(9, task)]
key = NamedTask.get_key(named_tasks, task)
found = NamedTask.find_by_key(named_tasks, key)
assert found is task
finally:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
# --- pick_highest_priority ---
async def test_pick_highest_priority_respects_list_order() -> None:
"""pick_highest_priority should return first completed task in list order."""
t1 = create_pending_task()
t2 = create_pending_task()
t3 = create_pending_task()
try:
# t2 is first in list, so should be picked even if t3 is also done
named_tasks = [
NamedTask.worker("step_b", 0, t2),
NamedTask.worker("step_a", 0, t1),
NamedTask.pull(0, t3),
]
done = {t2, t3} # Both t2 and t3 are done
result = NamedTask.pick_highest_priority(named_tasks, done)
assert result is t2
finally:
for t in [t1, t2, t3]:
t.cancel()
try:
await t
except asyncio.CancelledError:
pass
async def test_pick_highest_priority_workers_before_pull() -> None:
"""Workers listed first should have priority over pull."""
worker = create_pending_task()
pull = create_pending_task()
try:
# Workers first, then pull
named_tasks = [
NamedTask.worker("step", 0, worker),
NamedTask.pull(0, pull),
]
done = {worker, pull}
result = NamedTask.pick_highest_priority(named_tasks, done)
assert result is worker
finally:
for t in [worker, pull]:
t.cancel()
try:
await t
except asyncio.CancelledError:
pass
async def test_pick_highest_priority_returns_pull_when_only_pull_done() -> None:
"""Should return pull if it's the only completed task."""
worker = create_pending_task()
pull = create_pending_task()
try:
named_tasks = [
NamedTask.worker("step", 0, worker),
NamedTask.pull(0, pull),
]
done = {pull} # Only pull is done
result = NamedTask.pick_highest_priority(named_tasks, done)
assert result is pull
finally:
for t in [worker, pull]:
t.cancel()
try:
await t
except asyncio.CancelledError:
pass
async def test_pick_highest_priority_empty_done() -> None:
"""Should return None when done set is empty."""
task = create_pending_task()
try:
named_tasks = [NamedTask.worker("step", 0, task)]
result = NamedTask.pick_highest_priority(named_tasks, set())
assert result is None
finally:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
async def test_pick_highest_priority_no_match_raises() -> None:
"""Should raise ValueError when done is non-empty but no tasks match."""
task1 = create_pending_task()
task2 = create_pending_task()
try:
named_tasks = [NamedTask.worker("step", 0, task1)]
done = {task2} # task2 not in named_tasks
with pytest.raises(ValueError, match="No tasks in done set match"):
NamedTask.pick_highest_priority(named_tasks, done)
finally:
for t in [task1, task2]:
t.cancel()
try:
await t
except asyncio.CancelledError:
pass
# --- Integration ---
async def test_integration_with_asyncio_wait() -> None:
"""Full integration: create tasks, wait, pick priority, get key."""
async def quick() -> str:
return "done"
worker = asyncio.create_task(quick())
pull = create_pending_task()
try:
named_tasks = [
NamedTask.worker("fast_step", 0, worker),
NamedTask.pull(0, pull),
]
all_tasks = NamedTask.all_tasks(named_tasks)
done, _ = await asyncio.wait(all_tasks, timeout=1.0)
assert worker in done
completed = NamedTask.pick_highest_priority(named_tasks, done)
assert completed is not None
assert completed is worker
key = NamedTask.get_key(named_tasks, completed)
assert key == "fast_step:0"
finally:
pull.cancel()
try:
await pull
except asyncio.CancelledError:
pass
async def test_multiple_workers_same_step() -> None:
"""Should handle multiple workers for the same step (num_workers > 1)."""
w0 = create_pending_task()
w1 = create_pending_task()
try:
named_tasks = [
NamedTask.worker("parallel_step", 0, w0),
NamedTask.worker("parallel_step", 1, w1),
]
assert NamedTask.find_by_key(named_tasks, "parallel_step:0") is w0
assert NamedTask.find_by_key(named_tasks, "parallel_step:1") is w1
assert NamedTask.get_key(named_tasks, w0) == "parallel_step:0"
assert NamedTask.get_key(named_tasks, w1) == "parallel_step:1"
finally:
for t in [w0, w1]:
t.cancel()
try:
await t
except asyncio.CancelledError:
pass
packages/llama-index-workflows/tests/runtime/test_workflow_tracker.py
-187
View File
@@ -1,187 +0,0 @@
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 LlamaIndex Inc.
"""Tests for WorkflowTracker with strong references."""
from __future__ import annotations
import gc
import weakref
from typing import Any, cast
import pytest
from workflows import Workflow, step
from workflows.events import StartEvent, StopEvent
from workflows.runtime.types.plugin import RegisteredWorkflow, WorkflowRunFunction
from workflows.runtime.types.step_function import StepWorkerFunction
from workflows.runtime.workflow_tracker import WorkflowTracker
class SimpleWorkflow(Workflow):
@step
async def start(self, ev: StartEvent) -> StopEvent:
return StopEvent(result="done")
def test_add_workflow_to_tracker() -> None:
tracker = WorkflowTracker()
wf = SimpleWorkflow()
tracker.add(wf)
pending = tracker.get_pending()
assert len(pending) == 1
assert pending[0] is wf
def test_remove_workflow_from_tracker() -> None:
tracker = WorkflowTracker()
wf = SimpleWorkflow()
tracker.add(wf)
tracker.remove(wf)
assert tracker.get_pending() == []
def test_remove_nonexistent_workflow_is_safe() -> None:
tracker = WorkflowTracker()
wf = SimpleWorkflow()
# Should not raise
tracker.remove(wf)
def test_add_after_launch_raises_error() -> None:
tracker = WorkflowTracker()
tracker.mark_launched()
wf = SimpleWorkflow()
with pytest.raises(RuntimeError, match="Cannot add workflows after launch"):
tracker.add(wf)
def test_is_launched_initially_false() -> None:
tracker = WorkflowTracker()
assert tracker.is_launched is False
def test_mark_launched_sets_is_launched() -> None:
tracker = WorkflowTracker()
tracker.mark_launched()
assert tracker.is_launched is True
def test_set_and_get_registered() -> None:
tracker = WorkflowTracker()
wf = SimpleWorkflow()
async def mock_workflow_run(
init_state: object, start_event: object, tags: dict[str, Any]
) -> StopEvent:
return StopEvent(result="done")
registered = RegisteredWorkflow(
workflow=wf,
workflow_run_fn=cast(WorkflowRunFunction, mock_workflow_run),
steps=cast(dict[str, StepWorkerFunction[Any]], {"start": lambda: None}),
)
tracker.set_registered(wf, registered)
assert tracker.get_registered(wf) is registered
def test_get_registered_returns_none_if_not_set() -> None:
tracker = WorkflowTracker()
wf = SimpleWorkflow()
assert tracker.get_registered(wf) is None
def test_clear_resets_all_state() -> None:
tracker = WorkflowTracker()
wf = SimpleWorkflow()
tracker.add(wf)
tracker.mark_launched()
async def mock_workflow_run(
init_state: object, start_event: object, tags: dict[str, Any]
) -> StopEvent:
return StopEvent(result="done")
registered = RegisteredWorkflow(
workflow=wf,
workflow_run_fn=cast(WorkflowRunFunction, mock_workflow_run),
steps=cast(dict[str, StepWorkerFunction[Any]], {"start": lambda: None}),
)
tracker.set_registered(wf, registered)
tracker.clear()
assert tracker.get_pending() == []
assert tracker.get_registered(wf) is None
assert tracker.is_launched is False
def test_strong_refs_survive_gc() -> None:
"""Workflows survive GC when tracked with strong references."""
tracker = WorkflowTracker()
wf = SimpleWorkflow()
w = weakref.ref(wf)
tracker.add(wf)
assert len(tracker.get_pending()) == 1
# Drop external strong reference and force collection
del wf
gc.collect()
# With strong refs, the object should still be alive
assert w() is not None
assert len(tracker.get_pending()) == 1
def test_clear_releases_references() -> None:
"""clear() releases strong refs, allowing GC."""
tracker = WorkflowTracker()
wf = SimpleWorkflow()
w = weakref.ref(wf)
tracker.add(wf)
del wf
gc.collect()
# Still alive before clear
assert w() is not None
# Clear the tracker
tracker.clear()
gc.collect()
# Now the object should be collected
assert w() is None
def test_multiple_workflows_tracked_independently() -> None:
tracker = WorkflowTracker()
wf1 = SimpleWorkflow()
wf2 = SimpleWorkflow()
tracker.add(wf1)
tracker.add(wf2)
assert len(tracker.get_pending()) == 2
def test_add_same_workflow_twice_is_idempotent() -> None:
"""Adding the same workflow twice should not duplicate it."""
tracker = WorkflowTracker()
wf = SimpleWorkflow()
tracker.add(wf)
tracker.add(wf) # Second add
# Should only be one workflow
assert len(tracker.get_pending()) == 1
packages/llama-index-workflows/tests/test_state_manager.py
+34 -182
View File
@@ -1,196 +1,48 @@
from typing import Any, Type, Union, cast
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 LlamaIndex Inc.
"""Minimal unit tests for InMemoryStateStore.
Full state store protocol tests are in the integration test package
(llama-index-integration-tests/tests/test_state_store_matrix.py),
which tests InMemoryStateStore alongside SqliteStateStore.
These tests provide fast feedback during development of the base package.
"""
from __future__ import annotations
import pytest
from pydantic import (
BaseModel,
ConfigDict,
ValidationError,
field_serializer,
field_validator,
)
from workflows.context.serializers import BaseSerializer, JsonSerializer
from workflows.context.state_store import DictState, InMemoryStateStore
class MyRandomObject:
def __init__(self, name: str):
self.name = name
@pytest.mark.asyncio
async def test_in_memory_state_store_smoke() -> None:
"""Smoke test for basic InMemoryStateStore functionality."""
store: InMemoryStateStore[DictState] = InMemoryStateStore(DictState())
# Basic get/set
await store.set("name", "test")
assert await store.get("name") == "test"
class PydanticObject(BaseModel):
name: str
# Nested path
await store.set("nested", {"key": "value"})
assert await store.get("nested.key") == "value"
# Default on missing
assert await store.get("missing", default=None) is None
class MyState(BaseModel):
model_config = ConfigDict(
arbitrary_types_allowed=True,
validate_assignment=True,
strict=True,
)
my_obj: MyRandomObject
pydantic_obj: PydanticObject
name: str
age: int
@field_serializer("my_obj", when_used="always")
def serialize_my_obj(self, my_obj: MyRandomObject) -> str:
return my_obj.name
@field_validator("my_obj", mode="before")
@classmethod
def deserialize_my_obj(cls, v: Union[str, MyRandomObject]) -> MyRandomObject:
if isinstance(v, MyRandomObject):
return v
if isinstance(v, str):
return MyRandomObject(v)
raise ValueError(f"Invalid type for my_obj: {type(v)}")
class MyUnserializableState(BaseModel):
serializer_type: Type[BaseSerializer]
test_data: dict[str, Any]
@pytest.fixture
def default_state_manager() -> InMemoryStateStore[DictState]:
return InMemoryStateStore(DictState())
@pytest.fixture
def custom_state_manager() -> InMemoryStateStore[MyState]:
return InMemoryStateStore(
MyState(
my_obj=MyRandomObject("llama-index"),
pydantic_obj=PydanticObject(name="llama-index"),
name="John",
age=30,
)
)
@pytest.fixture
def unser_custom_state_manager() -> InMemoryStateStore[MyUnserializableState]:
return InMemoryStateStore(
MyUnserializableState(
serializer_type=type(JsonSerializer()), test_data={"test": 1, "data": 2}
)
)
# Clear
await store.clear()
assert await store.get("name", default=None) is None
@pytest.mark.asyncio
async def test_state_manager_defaults(
default_state_manager: InMemoryStateStore[DictState],
) -> None:
assert (
await default_state_manager.get_state()
).model_dump_json() == DictState().model_dump_json()
async def test_in_memory_edit_state() -> None:
"""Test edit_state context manager."""
store: InMemoryStateStore[DictState] = InMemoryStateStore(DictState())
await default_state_manager.set("name", "John")
await default_state_manager.set("age", 30)
async with store.edit_state() as state:
state["counter"] = 1
assert await default_state_manager.get("name") == "John"
assert await default_state_manager.get("age") == 30
await default_state_manager.set("nested", {"a": "b"})
assert await default_state_manager.get("nested.a") == "b"
await default_state_manager.set("nested.a", "c")
assert await default_state_manager.get("nested.a") == "c"
full_state = await default_state_manager.get_state()
assert full_state.name == "John"
assert full_state.age == 30
assert full_state.nested["a"] == "c"
@pytest.mark.asyncio
async def test_default_state_manager_serialization(
default_state_manager: InMemoryStateStore[DictState],
) -> None:
assert (
await default_state_manager.get_state()
).model_dump() == DictState().model_dump()
await default_state_manager.set("name", "John")
await default_state_manager.set("age", 30)
assert await default_state_manager.get("name") == "John"
assert await default_state_manager.get("age") == 30
data = default_state_manager.to_dict(JsonSerializer())
new_state_manager: InMemoryStateStore[DictState] = InMemoryStateStore.from_dict(
data, JsonSerializer()
)
assert await new_state_manager.get("name") == "John"
assert await new_state_manager.get("age") == 30
@pytest.mark.asyncio
async def test_custom_state_manager(
custom_state_manager: InMemoryStateStore[MyState],
) -> None:
assert (await custom_state_manager.get_state()).model_dump(mode="json") == MyState(
my_obj=MyRandomObject("llama-index"),
pydantic_obj=PydanticObject(name="llama-index"),
name="John",
age=30,
).model_dump(mode="json")
await custom_state_manager.set("name", "Jane")
await custom_state_manager.set("age", 25)
assert await custom_state_manager.get("name") == "Jane"
assert await custom_state_manager.get("age") == 25
full_state = await custom_state_manager.get_state()
assert isinstance(full_state, MyState)
assert full_state.name == "Jane"
assert full_state.age == 25
assert full_state.my_obj.name == "llama-index"
assert full_state.pydantic_obj.name == "llama-index"
# Ensure pydantic is providing type safety
with pytest.raises(ValidationError):
await custom_state_manager.set("age", "30")
with pytest.raises(AttributeError):
await custom_state_manager.set("age.nested", "llama-index")
@pytest.mark.asyncio
async def test_state_manager_custom_serialization(
custom_state_manager: InMemoryStateStore[MyState],
) -> None:
await custom_state_manager.set("name", "Jane")
await custom_state_manager.set("age", 25)
assert await custom_state_manager.get("name") == "Jane"
assert await custom_state_manager.get("age") == 25
data = custom_state_manager.to_dict(JsonSerializer())
new_state_manager: InMemoryStateStore[MyState] = cast(
InMemoryStateStore[MyState],
InMemoryStateStore.from_dict(data, JsonSerializer()),
)
assert await new_state_manager.get("name") == "Jane"
assert await new_state_manager.get("age") == 25
assert (await new_state_manager.get("my_obj")).name == "llama-index"
state = await new_state_manager.get_state()
assert state.pydantic_obj.name == "llama-index"
@pytest.mark.asyncio
async def test_state_manager_clear() -> None:
state_manager = InMemoryStateStore(DictState())
await state_manager.set("name", "Jane")
await state_manager.set("age", 25)
await state_manager.clear()
assert await state_manager.get("name", default=None) is None
assert await state_manager.get("age", default=None) is None
assert await store.get("counter") == 1
packages/llama-index-workflows/tests/test_streaming.py
+1 -54
View File
@@ -8,13 +8,11 @@ from typing import AsyncGenerator
import pytest
from workflows.context import Context
from workflows.decorators import step
from workflows.errors import WorkflowRuntimeError, WorkflowTimeoutError
from workflows.errors import WorkflowRuntimeError
from workflows.events import Event, StartEvent, StopEvent
from workflows.testing import WorkflowTestRunner
from workflows.workflow import Workflow
from .conftest import OneTestEvent # type: ignore[import]
class StreamingWorkflow(Workflow):
@step
@@ -30,57 +28,6 @@ class StreamingWorkflow(Workflow):
return StopEvent(result=None)
@pytest.mark.asyncio
async def test_e2e() -> None:
test_runner = WorkflowTestRunner(StreamingWorkflow())
r = await test_runner.run(expose_internal=False, exclude_events=[StopEvent])
assert all("msg" in ev for ev in r.collected)
@pytest.mark.asyncio
async def test_task_raised() -> None:
class DummyWorkflow(Workflow):
@step
async def step(self, ctx: Context, ev: StartEvent) -> StopEvent:
ctx.write_event_to_stream(OneTestEvent(test_param="foo"))
raise ValueError("The step raised an error!")
wf = DummyWorkflow()
r = wf.run()
# Make sure we don't block indefinitely here because the step raised
async for ev in r.stream_events():
if not isinstance(ev, StopEvent):
assert ev.test_param == "foo"
# Make sure the await actually caught the exception
with pytest.raises(ValueError, match="The step raised an error!"):
await r
@pytest.mark.asyncio
async def test_task_timeout() -> None:
class DummyWorkflow(Workflow):
@step
async def step(self, ctx: Context, ev: StartEvent) -> StopEvent:
ctx.write_event_to_stream(OneTestEvent(test_param="foo"))
await asyncio.sleep(2)
return StopEvent()
wf = DummyWorkflow(timeout=0.1)
r = wf.run()
# Make sure we don't block indefinitely here because the step raised
async for ev in r.stream_events():
if not isinstance(ev, StopEvent):
assert ev.test_param == "foo"
# Make sure the await actually caught the exception
with pytest.raises(WorkflowTimeoutError, match="Operation timed out"):
await r
@pytest.mark.asyncio
async def test_multiple_sequential_streams() -> None:
test_runner = WorkflowTestRunner(StreamingWorkflow())
packages/llama-index-workflows/tests/test_workflow.py
+31 -328
View File
@@ -10,7 +10,7 @@ import logging
import pickle
import threading
import weakref
from typing import Any, Callable, Optional, Union, cast
from typing import Any, Callable, Optional, cast
from unittest import mock
import pytest
@@ -24,7 +24,6 @@ from workflows.decorators import step
from workflows.errors import (
WorkflowConfigurationError,
WorkflowRuntimeError,
WorkflowTimeoutError,
WorkflowValidationError,
)
from workflows.events import (
@@ -40,9 +39,7 @@ from workflows.testing import WorkflowTestRunner
from workflows.workflow import Workflow
from .conftest import ( # type: ignore[import]
AnotherTestEvent,
DummyWorkflow,
LastEvent,
OneTestEvent,
)
@@ -74,24 +71,6 @@ async def test_workflow_initialization(workflow: Workflow) -> None:
assert not workflow._verbose
@pytest.mark.asyncio
async def test_workflow_run(workflow: Workflow) -> None:
r = await WorkflowTestRunner(workflow).run()
assert r.result == "Workflow completed"
@pytest.mark.asyncio
async def test_workflow_timeout() -> None:
class SlowWorkflow(Workflow):
@step
async def slow_step(self, ev: StartEvent) -> StopEvent:
await asyncio.sleep(2.0)
return StopEvent(result="Done")
with pytest.raises(WorkflowTimeoutError):
await WorkflowTestRunner(SlowWorkflow(timeout=0.1)).run()
@pytest.mark.asyncio
async def test_workflow_validation_unproduced_events() -> None:
class InvalidWorkflow(Workflow):
@@ -142,140 +121,6 @@ async def test_workflow_validation_start_event_not_consumed() -> None:
InvalidWorkflow()
@pytest.mark.asyncio
async def test_workflow_event_propagation() -> None:
events = []
class EventTrackingWorkflow(Workflow):
@step
async def step1(self, ev: StartEvent) -> OneTestEvent:
events.append("step1")
return OneTestEvent()
@step
async def step2(self, ev: OneTestEvent) -> StopEvent:
events.append("step2")
return StopEvent(result="Done")
await WorkflowTestRunner(EventTrackingWorkflow()).run()
assert events == ["step1", "step2"]
@pytest.mark.asyncio
async def test_workflow_sync_async_steps() -> None:
class SyncAsyncWorkflow(Workflow):
@step
async def async_step(self, ev: StartEvent) -> OneTestEvent:
return OneTestEvent()
@step
def sync_step(self, ev: OneTestEvent) -> StopEvent:
return StopEvent(result="Done")
await WorkflowTestRunner(SyncAsyncWorkflow()).run()
@pytest.mark.asyncio
async def test_workflow_sync_steps_only() -> None:
class SyncWorkflow(Workflow):
@step
def step_one(self, ctx: Context, ev: StartEvent) -> OneTestEvent:
ctx.collect_events(ev, [StartEvent])
return OneTestEvent()
@step
def step_two(self, ctx: Context, ev: OneTestEvent) -> StopEvent:
# ctx.collect_events(ev, [OneTestEvent])
return StopEvent()
await WorkflowTestRunner(SyncWorkflow()).run()
@pytest.mark.asyncio
async def test_workflow_num_workers() -> None:
signal = asyncio.Event()
lock = asyncio.Lock()
counter = 0
async def await_count(count: int) -> None:
nonlocal counter
async with lock:
counter += 1
if counter == count:
signal.set()
return
await signal.wait()
class NumWorkersWorkflow(Workflow):
@step
async def original_step(
self, ctx: Context, ev: StartEvent
) -> Union[OneTestEvent, LastEvent]:
await ctx.store.set("num_to_collect", 3)
ctx.send_event(OneTestEvent(test_param="test1"))
ctx.send_event(OneTestEvent(test_param="test2"))
ctx.send_event(OneTestEvent(test_param="test3"))
# send one extra event
ctx.send_event(AnotherTestEvent(another_test_param="test4"))
return LastEvent()
@step(num_workers=3)
async def test_step(self, ev: OneTestEvent) -> AnotherTestEvent:
await await_count(3) # wait for all 3 to be waiting
return AnotherTestEvent(another_test_param=ev.test_param)
@step
async def final_step(
self, ctx: Context, ev: Union[AnotherTestEvent, LastEvent]
) -> StopEvent:
n = await ctx.store.get("num_to_collect")
events = ctx.collect_events(ev, [AnotherTestEvent] * n)
if events is None:
return None # type: ignore
return StopEvent(result=[ev.another_test_param for ev in events])
workflow = NumWorkersWorkflow(timeout=1)
r = await WorkflowTestRunner(workflow).run()
assert "test4" in set(r.result)
assert len({"test1", "test2", "test3"} - set(r.result)) == 1
# Ensure ctx is serializable
ctx = r.ctx
assert ctx
ctx.to_dict()
@pytest.mark.asyncio
async def test_workflow_step_send_event() -> None:
class StepSendEventWorkflow(Workflow):
@step
async def step1(self, ctx: Context, ev: StartEvent) -> OneTestEvent:
ctx.send_event(OneTestEvent(), step="step2")
return None # type: ignore
@step
async def step2(self, ev: OneTestEvent) -> StopEvent:
return StopEvent(result="step2")
@step
async def step3(self, ev: OneTestEvent) -> StopEvent:
return StopEvent(result="step3")
workflow = StepSendEventWorkflow()
r = await WorkflowTestRunner(workflow).run()
assert r.result == "step2"
ctx = r.ctx
from workflows.context.external_context import ExternalContext
assert isinstance(ctx._face, ExternalContext)
replay = ctx._face._tick_log
assert TickAddEvent(OneTestEvent(), step_name="step2") in replay
@pytest.mark.asyncio
async def test_workflow_step_send_event_to_None() -> None:
class StepSendEventToNoneWorkflow(Workflow):
@@ -294,24 +139,16 @@ async def test_workflow_step_send_event_to_None() -> None:
assert isinstance(result.ctx._face, ExternalContext)
replay = result.ctx._face._tick_log
assert TickAddEvent(OneTestEvent()) in replay
assert TickAddEvent(event=OneTestEvent()) in replay
@pytest.mark.asyncio
async def test_workflow_step_returning_bogus() -> None:
class TestWorkflow(Workflow):
@step
async def step1(self, ctx: Context, ev: StartEvent) -> OneTestEvent:
async def step1(self, ctx: Context, ev: StartEvent) -> StopEvent:
return "foo" # type:ignore
@step
async def step2(self, ctx: Context, ev: StartEvent) -> OneTestEvent:
return OneTestEvent()
@step
async def step3(self, ev: OneTestEvent) -> StopEvent:
return StopEvent(result="step2")
with pytest.raises(
WorkflowRuntimeError,
match="Step function step1 returned str instead of an Event instance.",
@@ -319,22 +156,6 @@ async def test_workflow_step_returning_bogus() -> None:
await WorkflowTestRunner(TestWorkflow()).run()
@pytest.mark.asyncio
async def test_workflow_multiple_runs() -> None:
class DummyWorkflow(Workflow):
@step
async def step(self, ev: StartEvent) -> StopEvent:
return StopEvent(result=ev.number * 2)
runner = WorkflowTestRunner(DummyWorkflow())
results = await asyncio.gather(
runner.run(StartEvent(number=3)), # type: ignore
runner.run(StartEvent(number=42)), # type: ignore
runner.run(StartEvent(number=-99)), # type: ignore
)
assert set([r.result for r in results]) == {6, 84, -198}
def test_add_step() -> None:
class TestWorkflow(Workflow):
@step
@@ -367,17 +188,6 @@ def test_add_step_not_a_step() -> None:
TestWorkflow.add_step(another_step) # type: ignore
@pytest.mark.asyncio
async def test_workflow_task_raises() -> None:
class DummyWorkflow(Workflow):
@step
async def step(self, ev: StartEvent) -> StopEvent:
raise ValueError("The step raised an error!")
with pytest.raises(ValueError, match="The step raised an error!"):
await WorkflowTestRunner(DummyWorkflow()).run()
def test_workflow_disable_validation() -> None:
class DummyWorkflow(Workflow):
@step
@@ -391,37 +201,6 @@ def test_workflow_disable_validation() -> None:
mock_get_steps.assert_not_called()
@pytest.mark.asyncio
async def test_workflow_continue_context() -> None:
class DummyWorkflow(Workflow):
@step
async def step(self, ctx: Context, ev: StartEvent) -> StopEvent:
cur_number = await ctx.store.get("number", default=0)
new_number = cur_number + 1
await ctx.store.set("number", new_number)
return StopEvent(result=new_number)
wf = DummyWorkflow()
# first run
r = await WorkflowTestRunner(wf).run()
assert r.result == 1
ctx = r.ctx
assert ctx
# second run -- independent from the first
r = await WorkflowTestRunner(wf).run()
assert r.result == 1
ctx = r.ctx
assert ctx
# third run -- continue from the second run
handler = wf.run(ctx=ctx)
result = await handler
assert handler.ctx
assert result == 2
@pytest.mark.asyncio
async def test_workflow_pickle() -> None:
class DummyWorkflow(Workflow):
@@ -517,26 +296,6 @@ class HumanInTheLoopWorkflow(Workflow):
return StopEvent(result=ev.response)
@pytest.mark.asyncio
async def test_human_in_the_loop() -> None:
# workflow should raise a timeout error because hitl only works with streaming
with pytest.raises(WorkflowTimeoutError):
await WorkflowTestRunner(HumanInTheLoopWorkflow(timeout=0.01)).run()
# workflow should work with streaming
workflow = HumanInTheLoopWorkflow()
handler = workflow.run()
assert handler.ctx
async for event in handler.stream_events():
if isinstance(event, InputRequiredEvent):
assert event.prefix == "Enter a number: "
handler.ctx.send_event(HumanResponseEvent(response="42")) # type:ignore
final_result = await handler
assert final_result == "42"
@pytest.mark.asyncio
async def test_human_in_the_loop_with_resume() -> None:
# workflow should work with streaming
@@ -663,74 +422,6 @@ async def test_workflow_run_num_concurrent(
assert results == [f"Run {ix}: Done" for ix in range(1, 5)]
@pytest.mark.asyncio
async def test_custom_stop_event() -> None:
class CustomEventsWorkflow(Workflow):
@step
async def start_step(self, ev: MyStart) -> OneTestEvent:
return OneTestEvent()
@step
async def middle_step(self, ev: OneTestEvent) -> LastEvent:
return LastEvent()
@step
async def end_step(self, ev: LastEvent) -> MyStop:
return MyStop(outcome="Workflow completed")
wf = CustomEventsWorkflow()
assert wf._start_event_class == MyStart
assert wf.start_event_class == wf._start_event_class
assert wf._stop_event_class == MyStop
assert wf.stop_event_class == wf._stop_event_class
result: MyStop = await wf.run(query="foo")
assert result.outcome == "Workflow completed"
# Ensure the event types can be inferred when not passed to the init
wf = CustomEventsWorkflow()
assert wf._start_event_class == MyStart
assert wf._stop_event_class == MyStop
result = await wf.run(query="foo")
assert result.outcome == "Workflow completed"
# ensure that streaming exits
r = await WorkflowTestRunner(wf).run(MyStart(query="foo"))
assert len(r.collected) > 0
@pytest.mark.asyncio
async def test_workflow_stream_events_exits() -> None:
class CustomEventsWorkflow(Workflow):
@step
async def start_step(self, ev: MyStart) -> OneTestEvent:
return OneTestEvent()
@step
async def middle_step(self, ev: OneTestEvent) -> LastEvent:
return LastEvent()
@step
async def end_step(self, ev: LastEvent) -> MyStop:
return MyStop(outcome="Workflow completed")
wf = CustomEventsWorkflow()
handler = wf.run(query="foo")
async def _stream_events() -> Any:
async for event in handler.stream_events():
continue
return await handler
stream_task = asyncio.create_task(_stream_events())
result = await asyncio.wait_for(
stream_task,
timeout=1,
)
assert result.outcome == "Workflow completed"
class RandomEvent(Event):
pass
@@ -1089,9 +780,16 @@ class ParDone(Event):
@pytest.mark.asyncio
async def test_workflow_parallel_resume() -> None:
allowed_done = asyncio.Event()
"""Test that workflows with parallel workers can be serialized and resumed.
This test verifies that:
1. Events sent via ctx.send_event() are properly captured during serialization
2. In-progress workers are properly restored on resume
3. The workflow can complete after multiple serialize/resume cycles
"""
resume_event = asyncio.Event()
allowed_index = 0
all_workers_started = asyncio.Event()
worker_count = 0
class ParallelResumeWorkflow(Workflow):
@step
@@ -1102,36 +800,41 @@ async def test_workflow_parallel_resume() -> None:
@step(num_workers=4)
async def par(self, ev: Par) -> ParDone:
if ev.id != allowed_index:
await resume_event.wait()
nonlocal worker_count
# Track when workers start waiting (or complete for id=0)
worker_count += 1
if worker_count >= 4:
all_workers_started.set()
if ev.id == 0:
return ParDone(id=ev.id)
# Others wait for resume signal
await resume_event.wait()
return ParDone(id=ev.id)
@step
async def step3(self, ev: ParDone, ctx: Context) -> Optional[StopEvent]: # noqa - python 3.9 struggles here with | None
if ev.id == allowed_index:
allowed_done.set()
if ctx.collect_events(ev, [ParDone] * 4) is None:
return None
return StopEvent(result="Done")
wf = ParallelResumeWorkflow(timeout=10)
# First run: wait until all par workers have started, then serialize
handler = wf.run()
await allowed_done.wait()
await all_workers_started.wait()
# Small delay to ensure all workers are in BrokerState
await asyncio.sleep(0.01)
serialized_ctx = handler.ctx.to_dict()
try:
handler.cancel()
await handler.cancel_run()
except Exception:
pass
# immediately resume the workflow
allowed_index = 3
allowed_done.clear()
new_handler = wf.run(ctx=Context.from_dict(wf, serialized_ctx))
await allowed_done.wait()
# serialize again to detect inconsistencies
serialized_ctx = new_handler.ctx.to_dict()
# finally resume the workflow, and complete
# Second run: resume and complete
worker_count = 0
all_workers_started.clear()
new_handler = wf.run(ctx=Context.from_dict(wf, serialized_ctx))
resume_event.set()
await new_handler
result = await new_handler
assert result == "Done"
pyproject.toml
+4
View File
@@ -60,6 +60,10 @@ pythonVersion = "3.14"
root = "packages/llama-agents-integration-tests"
pythonVersion = "3.13"
[[tool.basedpyright.executionEnvironments]]
root = "packages/llama-index-workflows-dbos"
pythonVersion = "3.10"
[[tool.basedpyright.executionEnvironments]]
root = "examples"
pythonVersion = "3.14"