refact: introduce runtime plugin architecture and context separation (#296)

* refact: introduce runtime plugin architecture and context separation

- Add Runtime abstract base class with BasicRuntime as default implementation
- Split context into external, internal, and pre-context modules
- Replace workflow_registry with workflow_tracker for workflow lifecycle management
- Remove broker.py in favor of runtime-managed execution
- Add plugins package with get_current_runtime() for context-aware runtime access
- Update handler and workflow to use new runtime interfaces

.changeset/nine-hairs-relax.md

@@ -0,0 +1,10 @@
+---
+"llama-index-utils-workflow": minor
+"llama-index-workflows": minor
+---
+
+refactor: expand runtime plugin architecture
+
+- Refactoring to better support alternate distributed backends
+- Some `Context` methods may now raise errors if used in an unexpected context
+- `WorkflowHandler` is no longer a future. Retains compatibility methods for main use cases (exception, cancel, etc)

AGENTS.md

@@ -74,3 +74,13 @@ Linting, typechecking, and formatting:
 ```bash
 uv run --directory packages/llama-index-workflows pre-commit run -a
 ```
+
+## Testing Patterns
+
+We use **pytest** with idiomatic pytest patterns. Follow these guidelines:
+
+- **No Test Classes**: Do not use test classes to organize tests. Write tests as standalone functions. Achieve organization through descriptive function names (e.g., `test_create_job_with_invalid_input_raises_error`) or by splitting into separate test files.
+- **Pytest Fixtures**: Use fixtures for setup/teardown and shared test dependencies. Prefer fixtures over manual setup code repeated across tests.
+- **Prefer Real Objects Over Mocks**: Use simple dataclasses and real objects directly when available rather than mocking them. Only mock external dependencies or things that are truly difficult to instantiate.
+- **DRY Test Setup**: Do not repeat patches or setup code. Create reusable abstractions—fixtures, helper functions, or module-level constants—that can be shared across tests. Tests can easily be overwhelmed with setup; start from a rich suite of testing utilities to enable many small, expressive tests.
+- **Simple Testing Utilities**: Testing utilities should be basic—just functions, fixtures, and global variables. Avoid over-engineering test infrastructure.

packages/llama-index-integration-tests/tests/test_context_store.py

@@ -45,13 +45,15 @@ async def test_initial_state_accessible_in_tool(
 async def test_state_modification_persists(create_workflow: WorkflowFactory) -> None:
     """Test that state modifications in tools persist across calls."""
     call_count = 0
+    final_counter_value = None
 
     async def increment_counter(ctx: Context) -> str:
-        nonlocal call_count
+        nonlocal call_count, final_counter_value
         call_count += 1
         state = await ctx.store.get("state")
         state["counter"] = state.get("counter", 0) + 1
         await ctx.store.set("state", state)
+        final_counter_value = state["counter"]
         return f"Counter: {state['counter']}"
 
     workflow = create_workflow(
@@ -72,20 +74,22 @@ async def test_state_modification_persists(create_workflow: WorkflowFactory) ->
     # Verify tool was called twice
     assert call_count == 2
 
-    # Verify final state
-    final_state = await handler.ctx.store.get("state")
-    assert final_state["counter"] == 2
+    # Verify final state via the last captured value
+    assert final_counter_value == 2
 
 
 async def test_state_survives_handler_access(
     create_workflow: WorkflowFactory,
 ) -> None:
     """Test that state can be read from handler.ctx after workflow completes."""
+    captured_result = None
 
     async def set_result(ctx: Context) -> str:
+        nonlocal captured_result
         state = await ctx.store.get("state")
         state["result"] = "computation_complete"
         await ctx.store.set("state", state)
+        captured_result = state["result"]
         return "Done"
 
     workflow = create_workflow(
@@ -102,19 +106,25 @@ async def test_state_survives_handler_access(
         pass
     await handler
 
-    # Access state after workflow completion
-    state = await handler.ctx.store.get("state")
-    assert state["result"] == "computation_complete"
+    # Verify state was set correctly via captured value
+    assert captured_result == "computation_complete"
 
 
 async def test_complex_state_types(create_workflow: WorkflowFactory) -> None:
     """Test that complex state types (nested dicts, lists) work correctly."""
+    captured_state: dict | None = None
 
     async def modify_complex_state(ctx: Context) -> str:
+        nonlocal captured_state
         state = await ctx.store.get("state")
         state["items"].append("new_item")
         state["nested"]["count"] += 1
         await ctx.store.set("state", state)
+        # Capture a copy of the state for verification
+        captured_state = {
+            "items": list(state["items"]),
+            "nested": dict(state["nested"]),
+        }
         return "Modified"
 
     workflow = create_workflow(
@@ -134,14 +144,15 @@ async def test_complex_state_types(create_workflow: WorkflowFactory) -> None:
         pass
     await handler
 
-    state = await handler.ctx.store.get("state")
-    assert state["items"] == ["initial", "new_item"]
-    assert state["nested"]["count"] == 1
-    assert state["nested"]["name"] == "test"  # Unchanged
+    assert captured_state is not None
+    assert captured_state["items"] == ["initial", "new_item"]
+    assert captured_state["nested"]["count"] == 1
+    assert captured_state["nested"]["name"] == "test"  # Unchanged
 
 
 async def test_multiple_tools_share_state(create_workflow: WorkflowFactory) -> None:
     """Test that multiple different tools can share state."""
+    final_state: dict | None = None
 
     async def tool_a(ctx: Context) -> str:
         state = await ctx.store.get("state")
@@ -150,11 +161,13 @@ async def test_multiple_tools_share_state(create_workflow: WorkflowFactory) -> N
         return "A done"
 
     async def tool_b(ctx: Context) -> str:
+        nonlocal final_state
         state = await ctx.store.get("state")
         state["from_b"] = True
         # Verify tool_a's change is visible
         assert state.get("from_a") is True
         await ctx.store.set("state", state)
+        final_state = dict(state)
         return "B done"
 
     workflow = create_workflow(
@@ -172,6 +185,6 @@ async def test_multiple_tools_share_state(create_workflow: WorkflowFactory) -> N
         pass
     await handler
 
-    state = await handler.ctx.store.get("state")
-    assert state["from_a"] is True
-    assert state["from_b"] is True
+    assert final_state is not None
+    assert final_state["from_a"] is True
+    assert final_state["from_b"] is True

packages/llama-index-integration-tests/tests/test_human_in_the_loop.py

@@ -142,8 +142,10 @@ async def test_state_preserved_across_pause_resume(
     create_workflow: WorkflowFactory,
 ) -> None:
     """Test that workflow state is preserved when pausing and resuming."""
+    final_state: dict | None = None
 
     async def stateful_pause(ctx: Context) -> str:
+        nonlocal final_state
         # Modify state before pausing
         state = await ctx.store.get("state")
         state["before_pause"] = True
@@ -158,6 +160,7 @@ async def test_state_preserved_across_pause_resume(
         state = await ctx.store.get("state")
         state["after_resume"] = True
         await ctx.store.set("state", state)
+        final_state = dict(state)
 
         return "Done"
 
@@ -189,8 +192,8 @@ async def test_state_preserved_across_pause_resume(
 
     await handler
 
-    # Check all state was preserved
-    state = await handler.ctx.store.get("state")
-    assert state["initial"] is True
-    assert state["before_pause"] is True
-    assert state["after_resume"] is True
+    # Check all state was preserved via captured final state
+    assert final_state is not None
+    assert final_state["initial"] is True
+    assert final_state["before_pause"] is True
+    assert final_state["after_resume"] is True

packages/llama-index-utils-workflow/pyproject.toml

@@ -19,7 +19,7 @@ authors = [{name = "Adrian Lyjak", email = "adrianlyjak@gmail.com"}]
 requires-python = ">=3.9"
 dependencies = [
   "llama-index-core>=0.14,<0.15.0",
-  "llama-index-workflows>=2.12.0,<3.0.0",
+  "llama-index-workflows>=2.13.0,<3.0.0",
   "pyvis>=0.3.2"
 ]
 

packages/llama-index-utils-workflow/src/llama_index/utils/workflow/__init__.py

@@ -15,6 +15,7 @@ from llama_index.core.agent.workflow import (
 from llama_index.core.tools import AsyncBaseTool, BaseTool
 from pyvis.network import Network
 from workflows import Workflow
+from workflows.context.external_context import ExternalContext
 from workflows.events import (
     Event,
     StartEvent,
@@ -54,15 +55,15 @@ def _get_node_color(node: WorkflowGraphNode) -> str:
     elif node.node_type == "resource_config":
         return "#B2DFDB"  # Light teal for resource configs
     elif node.node_type == "event":
-        if node.is_subclass_of("StartEvent"):
+        if node.is_subclass_of("StartEvent"):  # type: ignore[possibly-missing-attribute]
             return "#E27AFF"  # Pink for start events
-        elif node.is_subclass_of("StopEvent"):
+        elif node.is_subclass_of("StopEvent"):  # type: ignore[possibly-missing-attribute]
             return "#FFA07A"  # Orange for stop events
         return "#90EE90"  # Light green for other events
     elif node.node_type == "agent":
-        if node.is_subclass_of("ReActAgent"):
+        if node.is_subclass_of("ReActAgent"):  # type: ignore[possibly-missing-attribute]
             return "#E27AFF"
-        elif node.is_subclass_of("CodeActAgent"):
+        elif node.is_subclass_of("CodeActAgent"):  # type: ignore[possibly-missing-attribute]
             return "#66ccff"
         return "#90EE90"
     elif node.node_type == "tool":
@@ -197,15 +198,15 @@ def _get_mermaid_css_class(node: WorkflowGraphNode) -> str:
     elif node.node_type == "resource_config":
         return "resourceConfigStyle"
     elif node.node_type == "event":
-        if node.is_subclass_of("StartEvent"):
+        if node.is_subclass_of("StartEvent"):  # type: ignore[possibly-missing-attribute]
             return "startEventStyle"
-        elif node.is_subclass_of("StopEvent"):
+        elif node.is_subclass_of("StopEvent"):  # type: ignore[possibly-missing-attribute]
             return "stopEventStyle"
         return "defaultEventStyle"
     elif node.node_type == "agent":
-        if node.is_subclass_of("ReActAgent"):
+        if node.is_subclass_of("ReActAgent"):  # type: ignore[possibly-missing-attribute]
             return "reactAgentStyle"
-        elif node.is_subclass_of("CodeActAgent"):
+        elif node.is_subclass_of("CodeActAgent"):  # type: ignore[possibly-missing-attribute]
             return "codeActAgentStyle"
         return "defaultAgentStyle"
     elif node.node_type == "tool":
@@ -462,10 +463,12 @@ def _extract_execution_graph(
     handler: WorkflowHandler, max_label_length: int | None = None
 ) -> Tuple[Dict[str, Tuple[str, str, type | None]], List[Tuple[str, str]]]:
     """Helper to extract nodes and edges from the workflow handler's tick log."""
-    if handler.ctx is None or handler.ctx._broker_run is None:
-        raise ValueError("No context/run info in this handler. Has it been run yet?")
 
-    ticks: List[WorkflowTick] = handler.ctx._broker_run._tick_log
+    ticks: List[WorkflowTick] = []
+    if handler.ctx is not None:
+        face = handler.ctx._face
+        if isinstance(face, ExternalContext):
+            ticks = face._tick_log
     nodes: Dict[str, Tuple[str, str, type | None]] = {}
     edges: List[Tuple[str, str]] = []
     event_node_by_identity: Dict[int, str] = {}

packages/llama-index-workflows-server/tests/server/test_handler_serialization.py

@@ -5,15 +5,60 @@ from __future__ import annotations
 import asyncio
 import json
 from datetime import datetime, timezone
+from typing import AsyncGenerator
 from unittest.mock import MagicMock
 
 import pytest
 from llama_agents.client.protocol import HandlerData
 from llama_agents.server.memory_workflow_store import MemoryWorkflowStore
 from llama_agents.server.server import _WorkflowHandler
-from workflows.context import Context
 from workflows.events import Event, StopEvent
 from workflows.handler import WorkflowHandler
+from workflows.runtime.types.plugin import ExternalRunAdapter
+from workflows.runtime.types.ticks import WorkflowTick
+from workflows.workflow import Workflow
+
+
+class MockRunAdapter(ExternalRunAdapter):
+    """Minimal mock adapter for testing."""
+
+    def __init__(self, run_id: str) -> None:
+        self._run_id = run_id
+        self._result: asyncio.Future[StopEvent] = asyncio.Future()
+
+    @property
+    def run_id(self) -> str:
+        return self._run_id
+
+    @property
+    def is_running(self) -> bool:
+        return not self._result.done()
+
+    async def get_result(self) -> StopEvent:
+        return await self._result
+
+    def get_result_or_none(self) -> StopEvent | None:
+        if self._result.done() and not self._result.cancelled():
+            return self._result.result()
+        return None
+
+    async def stream_published_events(self) -> AsyncGenerator[Event, None]:
+        result = await self._result
+        yield result
+
+    def abort(self) -> None:
+        if not self._result.done():
+            self._result.cancel()
+
+    def set_result(self, result: StopEvent) -> None:
+        if not self._result.done():
+            self._result.set_result(result)
+
+    async def send_event(self, tick: WorkflowTick) -> None:
+        pass
+
+    async def close(self) -> None:
+        pass
 
 
 class MyStopEvent(StopEvent):
@@ -21,10 +66,19 @@ class MyStopEvent(StopEvent):
 
 
 @pytest.mark.asyncio
-async def test__workflow_handler_to_dict_json_roundtrip() -> None:
-    ctx = MagicMock(spec=Context)
-    handler: WorkflowHandler = WorkflowHandler(ctx=ctx)
-    handler.set_result(MyStopEvent(message="ok"))
+async def test_workflow_handler_to_dict_json_roundtrip() -> None:
+    workflow = MagicMock(spec=Workflow)
+    workflow.workflow_name = "TestWorkflow"
+    adapter = MockRunAdapter(run_id="test-run-id")
+    # Set the result on the adapter
+    stop_event = MyStopEvent(message="ok")
+    adapter.set_result(stop_event)
+
+    handler: WorkflowHandler = WorkflowHandler(
+        workflow=workflow, external_adapter=adapter
+    )
+    # Wait for the result task to complete
+    await asyncio.sleep(0)
 
     queue: asyncio.Queue[Event] = asyncio.Queue()
 

packages/llama-index-workflows-server/tests/server/test_idle_release.py

@@ -18,6 +18,7 @@ from llama_agents.server.memory_workflow_store import MemoryWorkflowStore
 from llama_agents.server.server import WorkflowServer, _WorkflowHandler
 from server_test_fixtures import async_yield, wait_for_passing
 from workflows import Context, Workflow, step
+from workflows.context.context_types import SerializedContext
 from workflows.events import HumanResponseEvent, StartEvent, StopEvent
 
 
@@ -242,7 +243,7 @@ async def test_idle_release_restores_idle_since_on_reload(
 
     # Seed a persisted handler with idle_since
     idle_time = datetime.now(timezone.utc) - timedelta(minutes=2)
-    ctx = Context(waiting_workflow).to_dict()
+    ctx = SerializedContext().model_dump(mode="python")
     await seed_persistent_handler(
         memory_store,
         "idle-restore-1",
@@ -326,7 +327,7 @@ async def test_idle_handlers_not_resumed_on_server_start(
     """Test that idle handlers are not loaded into memory on server start."""
     # Seed the store with an idle handler
     idle_time = datetime.now(timezone.utc) - timedelta(minutes=2)
-    ctx = Context(waiting_workflow).to_dict()
+    ctx = SerializedContext().model_dump(mode="python")
     await seed_persistent_handler(
         memory_store,
         "idle-on-start-1",
@@ -536,7 +537,7 @@ async def test_try_reload_is_singleton_under_concurrency(
 
     # Seed a persisted idle handler
     idle_time = datetime.now(timezone.utc) - timedelta(minutes=2)
-    ctx = Context(waiting_workflow).to_dict()
+    ctx = SerializedContext().model_dump(mode="python")
     handler_id = "concurrent-reload-test"
     await seed_persistent_handler(
         memory_store,
@@ -617,9 +618,12 @@ async def test_release_handler_cancels_runtime_on_checkpoint_failure(
             ctx = run_handler.ctx
             assert ctx is not None
 
-            # Store a non-serializable object in context store
+            # Store a non-serializable object in context store directly
+            # (bypassing the Context.store property which requires internal context)
             # This will cause checkpoint() to fail when it tries to serialize
-            await ctx.store.set("bad_data", lambda x: x)  # lambdas can't be serialized
+            state_store = ctx._face._external_adapter.get_state_store()  # type: ignore[union-attr]
+            assert state_store is not None
+            await state_store.set("bad_data", lambda x: x)
 
             # Advance time past the idle timeout to trigger the timer
             await advance_time(traveller, timedelta(milliseconds=200), iterations=20)

packages/llama-index-workflows-server/tests/server/test_idle_release_live_http.py

@@ -34,7 +34,9 @@ async def test_fast_idle_timeout_does_not_drop_valid_event() -> None:
     def make_server() -> WorkflowServer:
         server = WorkflowServer(
             workflow_store=MemoryWorkflowStore(),
-            idle_release_timeout=timedelta(milliseconds=1),
+            # 50ms is still fast enough to test idle release behavior while
+            # avoiding race conditions in test infrastructure (especially with xdist)
+            idle_release_timeout=timedelta(milliseconds=50),
         )
         server.add_workflow(
             "waiting",

packages/llama-index-workflows-server/tests/server/test_server_endpoints.py

@@ -28,7 +28,9 @@ from server_test_fixtures import (
 from workflows import Context, step
 
 # Prepare the event to send
+from workflows.context.context_types import SerializedContext
 from workflows.context.serializers import JsonSerializer
+from workflows.context.state_store import DictState, InMemoryStateStore
 from workflows.events import StartEvent, StopEvent
 from workflows.workflow import Workflow
 
@@ -43,6 +45,15 @@ class CustomStopWorkflow(Workflow):
         return CustomStopEvent(message="custom-completed")
 
 
+async def serialize_context(state_dict: dict[str, Any]) -> SerializedContext:
+    ser_context = SerializedContext()
+    state = InMemoryStateStore(DictState())
+    for key, value in state_dict.items():
+        await state.set(key, value)
+    ser_context.state = state.to_dict(JsonSerializer())
+    return ser_context
+
+
 @pytest.fixture
 def server(
     simple_test_workflow: Workflow,
@@ -196,9 +207,9 @@ async def test_run_workflow_no_kwargs(client: AsyncClient) -> None:
 async def test_run_workflow_with_context(
     client: AsyncClient, server: WorkflowServer
 ) -> None:
-    ctx: Context = Context(server._workflows["test"])
-    await ctx.store.set("test_param", "message from context")
-    ctx_dict = ctx.to_dict()
+    ctx_dict = (
+        await serialize_context({"test_param": "message from context"})
+    ).model_dump(mode="python")
     response = await client.post("/workflows/test/run", json={"context": ctx_dict})
     assert response.status_code == 200
     data = response.json()
@@ -356,9 +367,9 @@ async def test_run_workflow_nowait_not_found(client: AsyncClient) -> None:
 @pytest.mark.asyncio
 async def test_get_workflow_result(client: AsyncClient, server: WorkflowServer) -> None:
     # Setup a context to test all the code paths
-    ctx: Context = Context(server._workflows["test"])
-    await ctx.store.set("test_param", "message from context")
-    ctx_dict = ctx.to_dict()
+    ctx_dict = (
+        await serialize_context({"test_param": "message from context"})
+    ).model_dump(mode="python")
     # run no-wait
     response = await client.post(
         "/workflows/test/run-nowait", json={"context": ctx_dict}

packages/llama-index-workflows-server/tests/server/test_server_persistence.py

@@ -16,7 +16,7 @@ from server_test_fixtures import (  # type: ignore[import]
     RequestedExternalEvent,
     wait_for_passing,  # type: ignore[import]
 )
-from workflows import Context
+from workflows.context.context_types import SerializedContext
 from workflows.events import Event, InternalDispatchEvent, StopEvent
 from workflows.workflow import Workflow
 
@@ -99,7 +99,7 @@ async def test_resume_active_handlers_across_server_restart(
     # Seed the store with a valid serialized Context using public API
 
     handler_id = "resume-1"
-    initial_ctx = Context(simple_test_workflow).to_dict()
+    initial_ctx = SerializedContext().model_dump(mode="python")
     await memory_store.update(
         PersistentHandler(
             handler_id=handler_id,
@@ -212,7 +212,7 @@ async def test_startup_ignores_unregistered_workflows(
             handler_id="known-1",
             workflow_name="test",
             status="running",
-            ctx=Context(simple_test_workflow).to_dict(),
+            ctx=SerializedContext().model_dump(mode="python"),
         )
     )
 
@@ -394,7 +394,7 @@ async def test_result_for_completed_persisted_handler_without_runtime_registrati
             workflow_name="test",
             status="completed",
             result=StopEvent(result="processed: default"),
-            ctx=Context(simple_test_workflow).to_dict(),
+            ctx=SerializedContext().model_dump(mode="python"),
         )
     )
 

packages/llama-index-workflows/src/llama_agents/workflows/__init__.py

@@ -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
+from workflows import *  # noqa: E402, F403  # pyright: ignore[reportWildcardImportFromLibrary]
 from workflows import __all__  # noqa: E402, F401

packages/llama-index-workflows/src/workflows/context/context.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 
 import asyncio
 import functools
-import uuid
 import warnings
 from typing import (
     TYPE_CHECKING,
@@ -17,13 +16,17 @@ from typing import (
     cast,
 )
 
-from pydantic import BaseModel, ValidationError
+from llama_index_instrumentation.dispatcher import (
+    active_instrument_tags,
+    instrument_tags,
+)
 
-from workflows.context.context_types import SerializedContext
-from workflows.decorators import StepConfig
+from workflows.context.external_context import ExternalContext
+from workflows.context.internal_context import InternalContext
+from workflows.context.pre_context import PreContext
 from workflows.errors import (
     ContextSerdeError,
-    WorkflowRuntimeError,
+    ContextStateError,
 )
 from workflows.events import (
     Event,
@@ -31,14 +34,15 @@ from workflows.events import (
     StopEvent,
 )
 from workflows.handler import WorkflowHandler
-from workflows.plugins.basic import basic_runtime
-from workflows.runtime.broker import WorkflowBroker
 from workflows.runtime.types.internal_state import BrokerState
-from workflows.runtime.types.plugin import Plugin, WorkflowRuntime
+from workflows.runtime.types.plugin import (
+    ExternalRunAdapter,
+)
 from workflows.types import RunResultT
+from workflows.utils import _nanoid as nanoid
 
 from .serializers import BaseSerializer, JsonSerializer
-from .state_store import MODEL_T, DictState, InMemoryStateStore
+from .state_store import MODEL_T, InMemoryStateStore
 
 if TYPE_CHECKING:  # pragma: no cover
     from workflows import Workflow
@@ -56,6 +60,7 @@ class UnserializableKeyWarning(Warning):
 warnings.simplefilter("once", UnserializableKeyWarning)
 
 
+# TODO(v3) remove this class, and replace with direct references to the pre/internal/external contexts
 class Context(Generic[MODEL_T]):
     """
     Global, per-run context for a `Workflow`. Provides an interface into the
@@ -120,163 +125,190 @@ class Context(Generic[MODEL_T]):
     # are known to be unserializable in some cases.
     known_unserializable_keys = ("memory",)
 
-    # Backing state store; serialized as `state`
-    _state_store: InMemoryStateStore[MODEL_T]
-    _broker_run: WorkflowBroker[MODEL_T] | None
-    _plugin: Plugin
-    _workflow: Workflow
+    # Current face - context is in exactly one state at a time
+    _face: (
+        PreContext[MODEL_T] | ExternalContext[MODEL_T, Any] | InternalContext[MODEL_T]
+    )
 
     def __init__(
         self,
         workflow: Workflow,
         previous_context: dict[str, Any] | None = None,
         serializer: BaseSerializer | None = None,
-        plugin: Plugin = basic_runtime,
     ) -> None:
-        self._serializer = serializer or JsonSerializer()
-        self._broker_run = None
-        self._plugin = plugin
-        self._workflow = workflow
+        # Start in pre-run (config) state - PreContext handles deserialization
+        pre_context: PreContext[MODEL_T] = PreContext(
+            workflow=workflow,
+            previous_context=previous_context,
+            serializer=serializer,
+        )
+        self._face = pre_context
 
-        # parse the serialized context
-        serializer = serializer or JsonSerializer()
-        if previous_context is not None:
-            try:
-                # Auto-detect and convert V0 to V1 if needed
-                previous_context_parsed = SerializedContext.from_dict_auto(
-                    previous_context
-                )
-                # validate it fully parses synchronously to avoid delayed validation errors
-                BrokerState.from_serialized(
-                    previous_context_parsed, workflow, serializer
-                )
-            except ValidationError as e:
-                raise ContextSerdeError(
-                    f"Context dict specified in an invalid format: {e}"
-                ) from e
-        else:
-            previous_context_parsed = SerializedContext()
-
-        self._init_snapshot = previous_context_parsed
-
-        # initialization of the state store is a bit complex, due to inferring and validating its type from the
-        # provided workflow context args
-
-        state_types: set[Type[BaseModel]] = set()
-        for _, step_func in workflow._get_steps().items():
-            step_config: StepConfig = step_func._step_config
-            if (
-                step_config.context_state_type is not None
-                and step_config.context_state_type != DictState
-                and issubclass(step_config.context_state_type, BaseModel)
-            ):
-                state_types.add(step_config.context_state_type)
-
-        # Find the most derived state type from the inheritance hierarchy
-        # This allows base workflows to use Context[BaseState] and child workflows
-        # to use Context[ChildState] where ChildState extends BaseState
-        state_type: Type[BaseModel]
-        if state_types:
-            state_type = _find_most_derived_state_type(state_types)
-        else:
-            state_type = DictState
-        if previous_context_parsed.state:
-            # perhaps offer a way to clear on invalid
-            store_state = InMemoryStateStore.from_dict(
-                previous_context_parsed.state, serializer
-            )
-            if store_state.state_type != state_type:
-                raise ValueError(
-                    f"State type mismatch. Workflow context expected {state_type.__name__}, got {store_state.state_type.__name__}"
-                )
-            self._state_store = cast(InMemoryStateStore[MODEL_T], store_state)
-        else:
-            try:
-                state_instance = cast(MODEL_T, state_type())
-                self._state_store = InMemoryStateStore(state_instance)
-            except Exception as e:
-                raise WorkflowRuntimeError(
-                    f"Failed to initialize state of type {state_type}. Does your state define defaults for all fields? Original error:\n{e}"
-                ) from e
+    @classmethod
+    def _create_face(
+        cls,
+        face: PreContext[MODEL_T]
+        | ExternalContext[MODEL_T, Any]
+        | InternalContext[MODEL_T],
+    ) -> Context[MODEL_T]:
+        new_ctx = cast(Context[MODEL_T], object.__new__(cls))
+        new_ctx._face = face
+        return new_ctx
 
     @property
     def is_running(self) -> bool:
         """Whether the workflow is currently running."""
-        if self._broker_run is None:
-            return self._init_snapshot.is_running
+        if isinstance(self._face, PreContext):
+            return self._face.is_running
+        elif isinstance(self._face, ExternalContext):
+            return self._face.is_running
         else:
-            return self._broker_run.is_running
+            _warn_is_running_in_step()
+            return True
 
-    def _init_broker(
-        self, workflow: Workflow, plugin: WorkflowRuntime | None = None
-    ) -> WorkflowBroker[MODEL_T]:
-        if self._broker_run is not None:
-            raise WorkflowRuntimeError("Broker already initialized")
-        # Initialize a runtime plugin (asyncio-based by default)
-        runtime: WorkflowRuntime = plugin or self._plugin.new_runtime(str(uuid.uuid4()))
-        # Initialize the new broker implementation (broker2)
-        broker: WorkflowBroker[MODEL_T] = WorkflowBroker(
-            workflow=workflow,
-            context=cast("Context[MODEL_T]", self),
-            runtime=runtime,
-            plugin=self._plugin,
+    def _require_pre(self, fn: str) -> PreContext[MODEL_T]:
+        """Require context to be in pre-run state. Raises ContextStateError if not."""
+        if isinstance(self._face, PreContext):
+            return self._face  # type: ignore[invalid-return-type]
+        raise ContextStateError(
+            f"{fn} requires a pre-run context. The workflow has already started."
         )
-        self._broker_run = broker
-        return broker
+
+    def _require_external(self, fn: str) -> ExternalContext[MODEL_T, Any]:
+        """Require context to be in external state. Raises ContextStateError if not."""
+        if isinstance(self._face, ExternalContext):
+            return self._face
+        if isinstance(self._face, PreContext):
+            raise ContextStateError(
+                f"{fn} requires a running workflow. Call workflow.run() first."
+            )
+        raise ContextStateError(
+            f"{fn} is only available from handler code, not from within steps."
+        )
+
+    def _require_internal(self, fn: str) -> InternalContext[MODEL_T]:
+        """Require context to be in internal state. Raises ContextStateError if not."""
+        if isinstance(self._face, InternalContext):
+            return self._face  # type: ignore[invalid-return-type]
+        if isinstance(self._face, PreContext):
+            raise ContextStateError(
+                f"{fn} requires a running workflow. Call workflow.run() first."
+            )
+        raise ContextStateError(f"{fn} is only available from within step functions.")
+
+    @classmethod
+    def _create_internal(
+        cls,
+        workflow: Workflow,
+    ) -> Context[MODEL_T]:
+        """Create a Context directly in internal face state.
+
+        Requires a current run context (via with_current_run_id) to be set.
+        """
+        internal_adapter = workflow._runtime.get_internal_adapter()
+        new_ctx = cast(Context[MODEL_T], object.__new__(cls))
+        new_ctx._face = cast(
+            InternalContext[MODEL_T],
+            InternalContext(
+                internal_adapter=internal_adapter,
+                workflow=workflow,
+            ),
+        )
+        return new_ctx
+
+    @classmethod
+    def _create_external(
+        cls,
+        workflow: Workflow,
+        external_adapter: ExternalRunAdapter,
+        serializer: BaseSerializer = JsonSerializer(),
+    ) -> Context[MODEL_T]:
+        """Create a Context directly in external face state with a broker."""
+
+        new_ctx = cast(Context[MODEL_T], object.__new__(cls))
+
+        # Set external face
+        new_ctx._face = cast(
+            ExternalContext[MODEL_T, Any],
+            ExternalContext(
+                workflow=workflow,
+                external_adapter=external_adapter,
+                serializer=serializer,
+            ),
+        )
+        return new_ctx
 
     def _workflow_run(
         self,
         workflow: Workflow,
-        start_event: StartEvent | None = None,
-        semaphore: asyncio.Semaphore | None = None,
+        start_event: StartEvent
+        | None,  # None only when resuming a workflow from a snapshotted context
+        run_id: str | None = None,
     ) -> WorkflowHandler:
         """
         called by package internally from the workflow to run it
         """
-        prev_broker: WorkflowBroker[MODEL_T] | None = None
-        if self._broker_run is not None:
-            prev_broker = self._broker_run
-            self._broker_run = None
+        run_id = run_id or nanoid()
+        with instrument_tags({**active_instrument_tags.get(), "run_id": run_id}):
+            # Get or create PreContext for initialization
+            if isinstance(self._face, PreContext):
+                pre = self._face
+            elif isinstance(self._face, ExternalContext):
+                # Check for concurrent run
+                if self._face.is_running:
+                    raise ContextStateError(
+                        "Cannot start a new run while context is already running. "
+                        "Wait for completion or use a new Context."
+                    )
+                # Continuation: create fresh PreContext from current state
+                pre = PreContext(
+                    workflow=workflow,
+                    previous_context=self._face.to_dict(),
+                    serializer=self._face._serializer,
+                )
+            else:
+                raise ContextStateError(
+                    "Cannot start workflow from a step function context"
+                )
 
-        self._broker_run = self._init_broker(workflow)
+            # Compute state from serialized snapshot
+            init_state = BrokerState.from_serialized(
+                pre.init_snapshot, workflow, pre._serializer
+            )
 
-        async def before_start() -> None:
-            if prev_broker is not None:
-                try:
-                    await prev_broker.shutdown()
-                except Exception:
-                    pass
-            if semaphore is not None:
-                await semaphore.acquire()
+            external_adapter = workflow._runtime.run_workflow(
+                run_id=run_id,
+                workflow=workflow,
+                init_state=init_state,
+                start_event=start_event,
+                serialized_state=pre.serialized_state,
+                serializer=pre.serializer,
+            )
 
-        async def after_complete() -> None:
-            if semaphore is not None:
-                semaphore.release()
+            # TODO(v3): Remove mutation. Handler will just be the external face.
+            self._face = cast(
+                ExternalContext[MODEL_T, Any],
+                ExternalContext(
+                    workflow=workflow,
+                    external_adapter=external_adapter,
+                    serializer=pre._serializer,
+                ),
+            )
 
-        state = BrokerState.from_serialized(
-            self._init_snapshot, workflow, self._serializer
-        )
-        return self._broker_run.start(
-            workflow=workflow,
-            previous=state,
-            start_event=start_event,
-            before_start=before_start,
-            after_complete=after_complete,
-        )
+            return WorkflowHandler(
+                workflow=workflow,
+                external_adapter=external_adapter,
+                ctx=self,
+            )
 
     def _workflow_cancel_run(self) -> None:
-        """
-        Called internally from the handler to cancel a context's run
-        """
-        self._running_broker.cancel_run()
-
-    @property
-    def _running_broker(self) -> WorkflowBroker[MODEL_T]:
-        if self._broker_run is None:
-            raise WorkflowRuntimeError(
-                "Workflow run is not yet running. Make sure to only call this method after the context has been passed to a workflow.run call."
-            )
-        return self._broker_run
+        """Called internally from the handler to cancel a context's run."""
+        if isinstance(self._face, ExternalContext):
+            self._face.cancel()
+        elif isinstance(self._face, PreContext):
+            _warn_cancel_before_start()
+        else:
+            _warn_cancel_in_step()
 
     @property
     def store(self) -> InMemoryStateStore[MODEL_T]:
@@ -288,7 +320,7 @@ class Context(Generic[MODEL_T]):
         Returns:
             InMemoryStateStore[MODEL_T]: The state store instance.
         """
-        return self._state_store
+        return self._face.store  # type: ignore[return-value]
 
     def to_dict(self, serializer: BaseSerializer | None = None) -> dict[str, Any]:
         """Serialize the context to a JSON-serializable dict.
@@ -320,35 +352,15 @@ class Context(Generic[MODEL_T]):
             result = await my_workflow.run(..., ctx=restored_ctx)
             ```
         """
-        serializer = serializer or self._serializer
-
-        # Serialize state using the state manager's method
-        state_data = {}
-        if self._state_store is not None:
-            state_data = self._state_store.to_dict(serializer)
-
-        # Get the broker state - either from the running broker or from the init snapshot
-        if self._broker_run is not None:
-            broker_state = self._broker_run._state
-        else:
-            # Deserialize the init snapshot to get a BrokerState, then re-serialize it
-            # This ensures we always output the current format
-            broker_state = BrokerState.from_serialized(
-                self._init_snapshot, self._workflow, serializer
-            )
-
-        context = broker_state.to_serialized(serializer)
-        context.state = state_data
-        # mode="python" to support pickling over json if one so chooses. This should perhaps be moved into the serializers
-        return context.model_dump(mode="python")
+        return self._require_external(fn="to_dict").to_dict(serializer)
 
     @classmethod
     def from_dict(
         cls,
-        workflow: "Workflow",
+        workflow: Workflow,
         data: dict[str, Any],
         serializer: BaseSerializer | None = None,
-    ) -> "Context[MODEL_T]":
+    ) -> Context[MODEL_T]:
         """Reconstruct a `Context` from a serialized payload.
 
         Args:
@@ -389,7 +401,7 @@ class Context(Generic[MODEL_T]):
         Returns:
             list[str]: Names of steps that have at least one active worker.
         """
-        return await self._running_broker.running_steps()
+        return await self._require_external(fn="running_steps").running_steps()
 
     def collect_events(
         self, ev: Event, expected: list[Type[Event]], buffer_id: str | None = None
@@ -428,7 +440,9 @@ class Context(Generic[MODEL_T]):
         See Also:
             - [Event][workflows.events.Event]
         """
-        return self._running_broker.collect_events(ev, expected, buffer_id)
+        return self._require_internal(fn="collect_events").collect_events(
+            ev, expected, buffer_id
+        )
 
     def send_event(self, message: Event, step: str | None = None) -> None:
         """Dispatch an event to one or all workflow steps.
@@ -468,7 +482,16 @@ class Context(Generic[MODEL_T]):
             result = await handler
             ```
         """
-        return self._running_broker.send_event(message, step)
+        # send_event can be called from internal (steps) or external (handler) contexts
+        if isinstance(self._face, InternalContext):
+            self._face.send_event(message, step)
+        elif isinstance(self._face, ExternalContext):
+            self._face.send_event(message, step)
+        else:
+            raise ContextStateError(
+                "send_event() called before workflow started. "
+                "Call workflow.run() first."
+            )
 
     async def wait_for_event(
         self,
@@ -518,7 +541,7 @@ class Context(Generic[MODEL_T]):
                 return StopEvent(result=response.response)
             ```
         """
-        return await self._running_broker.wait_for_event(
+        return await self._require_internal(fn="wait_for_event").wait_for_event(
             event_type, waiter_event, waiter_id, requirements, timeout
         )
 
@@ -537,7 +560,7 @@ class Context(Generic[MODEL_T]):
                 return StopEvent(result="ok")
             ```
         """
-        self._running_broker.write_event_to_stream(ev)
+        self._require_internal(fn="write_event_to_stream").write_event_to_stream(ev)
 
     def get_result(self) -> RunResultT:
         """Return the final result of the workflow run.
@@ -558,15 +581,29 @@ class Context(Generic[MODEL_T]):
 
         Returns:
             RunResultT: The value provided via a `StopEvent`.
+
+        Raises:
+            ContextStateError: If called before the workflow is running or
+                from within a step function.
         """
         _warn_get_result()
-        if self._running_broker._handler is None:
-            raise WorkflowRuntimeError("Workflow handler is not set")
-        return self._running_broker._handler.result()
+        stop_event = self._require_external(fn="get_result").get_result()
+        return stop_event.result if type(stop_event) is StopEvent else stop_event
 
     def stream_events(self) -> AsyncGenerator[Event, None]:
-        """The internal queue used for streaming events to callers."""
-        return self._running_broker.stream_published_events()
+        """Stream events published by the workflow.
+
+        Returns an async generator that yields events as they are published
+        by steps via `write_event_to_stream()`.
+
+        Returns:
+            AsyncGenerator[Event, None]: Stream of published events.
+
+        Raises:
+            ContextStateError: If called before the workflow is running or
+                from within a step function.
+        """
+        return self._require_external(fn="stream_events").stream_events()
 
     @property
     def streaming_queue(self) -> asyncio.Queue:
@@ -576,6 +613,7 @@ class Context(Generic[MODEL_T]):
         stream_events(). A deprecation warning is emitted once per process.
         """
         _warn_streaming_queue()
+        self._require_external(fn="streaming_queue")
         q: asyncio.Queue[Event] = asyncio.Queue()
 
         async def _pump() -> None:
@@ -618,48 +656,27 @@ def _warn_streaming_queue() -> None:
     )
 
 
-def _find_most_derived_state_type(state_types: set[Type[BaseModel]]) -> Type[BaseModel]:
-    """Find the most derived (most specific) state type from a set of types.
+@functools.lru_cache(maxsize=1)
+def _warn_is_running_in_step() -> None:
+    warnings.warn(
+        "is_running called from within a step; the workflow is always "
+        "running inside a step. This usage is deprecated.",
+        DeprecationWarning,
+        stacklevel=3,
+    )
 
-    All types must be in a single inheritance chain, i.e., one type must be
-    a subclass of all other types (the most derived type).
 
-    Args:
-        state_types: Set of state types to analyze.
+@functools.lru_cache(maxsize=1)
+def _warn_cancel_before_start() -> None:
+    warnings.warn(
+        "cancel() called before workflow started; nothing to cancel.",
+        stacklevel=3,
+    )
 
-    Returns:
-        The most derived type in the inheritance hierarchy.
 
-    Raises:
-        ValueError: If types are not in a compatible inheritance hierarchy.
-    """
-    type_list = list(state_types)
-
-    if len(type_list) == 1:
-        return type_list[0]
-
-    # Find the most derived type - it should be a subclass of all others
-    most_derived: Type[BaseModel] | None = None
-
-    for candidate in type_list:
-        is_most_derived = True
-        for other in type_list:
-            if other is candidate:
-                continue
-            # candidate must be a subclass of other (or equal to it)
-            if not issubclass(candidate, other):
-                is_most_derived = False
-                break
-        if is_most_derived:
-            most_derived = candidate
-            break
-
-    if most_derived is None:
-        # No single type is a subclass of all others - incompatible hierarchy
-        raise ValueError(
-            "Multiple state types are not in a compatible inheritance hierarchy. "
-            "All state types must share a common inheritance chain. Found: "
-            + ", ".join([st.__name__ for st in state_types])
-        )
-
-    return most_derived
+@functools.lru_cache(maxsize=1)
+def _warn_cancel_in_step() -> None:
+    warnings.warn(
+        "cancel() called from within a step; use send_event() instead.",
+        stacklevel=3,
+    )

packages/llama-index-workflows/src/workflows/context/context_types.py

@@ -3,6 +3,11 @@ from typing import Any, Optional
 
 from pydantic import BaseModel, Field
 from pydantic.functional_validators import model_validator
+from typing_extensions import TypeVar
+
+from workflows.context.state_store import DictState
+
+MODEL_T = TypeVar("MODEL_T", bound=BaseModel, default=DictState)  # type: ignore
 
 
 class SerializedContextV0(BaseModel):

packages/llama-index-workflows/src/workflows/context/external_context.py

@@ -0,0 +1,186 @@
+"""External context - for handlers and code outside workflow steps."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Coroutine, Generic, cast
+
+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.errors import WorkflowRuntimeError
+from workflows.events import StopEvent
+from workflows.runtime.types.internal_state import BrokerState
+from workflows.runtime.types.plugin import (
+    ExternalRunAdapter,
+    SnapshottableAdapter,
+    V2RuntimeCompatibilityShim,
+    as_snapshottable_adapter,
+    as_v2_runtime_compatibility_shim,
+)
+from workflows.runtime.types.ticks import TickAddEvent, TickCancelRun, WorkflowTick
+
+if TYPE_CHECKING:
+    from workflows.context.serializers import BaseSerializer
+    from workflows.events import Event
+    from workflows.workflow import Workflow
+
+RunResultT = TypeVar("RunResultT", default=Any)  # type: ignore[misc]
+
+
+class ExternalContext(Generic[MODEL_T, RunResultT]):
+    """Context for handler code and external workflow interaction.
+
+    Used by WorkflowHandler to send events into the workflow,
+    stream events out, and retrieve the final result.
+    """
+
+    _workflow: Workflow
+    _external_adapter: ExternalRunAdapter
+    _workers: list[asyncio.Task[Any]]
+    _serializer: BaseSerializer
+
+    def __init__(
+        self,
+        workflow: Workflow,
+        external_adapter: ExternalRunAdapter,
+        serializer: BaseSerializer = JsonSerializer(),
+    ) -> None:
+        self._workflow = workflow
+        self._external_adapter = external_adapter
+        self._serializer = serializer
+        self._workers = []
+
+    @property
+    def is_running(self) -> bool:
+        """Whether the workflow is currently running."""
+        as_shim = as_v2_runtime_compatibility_shim(self._external_adapter)
+        if as_shim is None:
+            # Assume running if not v2 runtime compatible. This is mainly just used for resuming
+            # an interrupted serialized context, which is not supported the same in distributed runtimes
+            return True
+
+        return as_shim.is_running
+
+    def _execute_task(self, coro: Coroutine[Any, Any, Any]) -> asyncio.Task[Any]:
+        """Execute a coroutine as a background task."""
+        task = asyncio.create_task(coro)
+        self._workers.append(task)
+
+        def _remove_task(_: asyncio.Task[Any]) -> None:
+            try:
+                self._workers.remove(task)
+            except ValueError:
+                # Task was already cleared during shutdown or cleanup.
+                pass
+
+        task.add_done_callback(_remove_task)
+        return task
+
+    @property
+    def _tick_log(self) -> list[WorkflowTick]:
+        """Get the tick log from the snapshottable adapter."""
+        return self._require_snapshottable().replay()
+
+    def _require_snapshottable(self) -> SnapshottableAdapter:
+        snapshottable = as_snapshottable_adapter(self._external_adapter)
+        if snapshottable is None:
+            raise WorkflowRuntimeError(
+                f"Runtime of type {self._workflow.runtime.__class__.__qualname__} is not snapshottable"
+            )
+        return snapshottable
+
+    @property
+    def _state(self) -> BrokerState:
+        """Compute current state from init state and tick log."""
+        from workflows.runtime.control_loop import rebuild_state_from_ticks
+
+        ticks = self._tick_log
+        snapshottable = self._require_snapshottable()
+        state = snapshottable.init_state
+        new_state = rebuild_state_from_ticks(state, ticks)
+        return new_state
+
+    @property
+    def store(self) -> InMemoryStateStore[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)
+
+    def send_event(self, message: Event, step: str | None = None) -> None:
+        """Send an event into the workflow."""
+        if step is not None:
+            self._workflow._validate_valid_step_message(step, message)
+
+        self._execute_task(
+            self._external_adapter.send_event(
+                TickAddEvent(event=message, step_name=step)
+            )
+        )
+
+    async def running_steps(self) -> list[str]:
+        """Get list of currently running step names."""
+        state = self._state
+        return [
+            step for step in state.workers.keys() if state.workers[step].in_progress
+        ]
+
+    def _require_v2_runtime_compatibility(self) -> V2RuntimeCompatibilityShim:
+        v2_shim = as_v2_runtime_compatibility_shim(self._external_adapter)
+        if v2_shim is None:
+            raise WorkflowRuntimeError(
+                f"Runtime of type {self._workflow.runtime.__class__.__qualname__} is not v2 runtime compatible"
+            )
+        return v2_shim
+
+    def get_result(self) -> StopEvent:
+        """Get the workflow's final result. Raises if not yet complete."""
+        result = self._require_v2_runtime_compatibility().get_result_or_none()
+        if result is None:
+            raise WorkflowRuntimeError(
+                f"Workflow run with run_id {self._external_adapter.run_id} is not complete"
+            )
+        return result
+
+    def stream_events(self) -> AsyncGenerator[Event, None]:
+        """Stream events published by the workflow."""
+        return self._external_adapter.stream_published_events()
+
+    def to_dict(self, serializer: BaseSerializer | None = None) -> dict[str, Any]:
+        """Serialize context state for persistence."""
+        active_serializer = serializer or self._serializer
+
+        # Fetch state store from adapter and serialize
+        state_data = {}
+        state_store = self._external_adapter.get_state_store()
+        if state_store is not None:
+            state_data = state_store.to_dict(active_serializer)
+
+        # Get the broker state
+        broker_state = self._state
+
+        context = broker_state.to_serialized(active_serializer)
+        context.state = state_data
+        return context.model_dump(mode="python")
+
+    def cancel(self) -> None:
+        """Request workflow cancellation."""
+        self._execute_task(self._external_adapter.send_event(TickCancelRun()))
+
+    async def shutdown(self) -> None:
+        """Cancel the running workflow and clean up resources.
+
+        Sends a cancel signal, cancels all outstanding workers (both external
+        and broker workers), and closes the adapter. State remains available
+        for inspection.
+        """
+        await self._external_adapter.send_event(TickCancelRun())
+        # Clean up external context workers
+        for worker in self._workers:
+            worker.cancel()
+        self._workers.clear()
+        await self._external_adapter.close()

packages/llama-index-workflows/src/workflows/context/internal_context.py

@@ -0,0 +1,187 @@
+"""Internal context - for step execution within workflows."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+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.errors import WorkflowRuntimeError
+from workflows.runtime.types.results import (
+    AddCollectedEvent,
+    AddWaiter,
+    DeleteCollectedEvent,
+    DeleteWaiter,
+    StepWorkerContext,
+    StepWorkerStateContextVar,
+    WaitingForEvent,
+)
+from workflows.runtime.types.ticks import TickAddEvent
+
+if TYPE_CHECKING:
+    from workflows.events import Event
+    from workflows.runtime.types.plugin import InternalRunAdapter
+    from workflows.workflow import Workflow
+
+T = TypeVar("T", bound="Event")
+
+logger = logging.getLogger(__name__)
+
+
+class InternalContext(Generic[MODEL_T]):
+    """Context for code running inside workflow step functions.
+
+    Provides access to state store, event collection, waiting for events,
+    and publishing to the event stream.
+    """
+
+    _internal_adapter: InternalRunAdapter
+    _workflow: Workflow
+    _workers: list[asyncio.Task[Any]]
+
+    def __init__(
+        self,
+        internal_adapter: InternalRunAdapter,
+        workflow: Workflow,
+    ) -> None:
+        self._internal_adapter = internal_adapter
+        self._workflow = workflow
+        self._workers = []
+
+    def _execute_task(self, coro: Coroutine[Any, Any, Any]) -> asyncio.Task[Any]:
+        """Execute a coroutine as a tracked background task."""
+        task = asyncio.create_task(coro)
+        self._workers.append(task)
+
+        def _on_done(t: asyncio.Task[Any]) -> None:
+            try:
+                self._workers.remove(t)
+            except ValueError:
+                # Task was already cleared during shutdown or cleanup.
+                pass
+            # Log exceptions from fire-and-forget tasks (cancelled is not an error)
+            if not t.cancelled():
+                exc = t.exception()
+                if exc is not None:
+                    logger.error(
+                        "Background task failed with exception",
+                        exc_info=(type(exc), exc, exc.__traceback__),
+                    )
+
+        task.add_done_callback(_on_done)
+        return task
+
+    def cancel_background_tasks(self) -> None:
+        """Cancel all tracked background tasks."""
+        for worker in self._workers:
+            worker.cancel()
+        self._workers.clear()
+
+    @staticmethod
+    def _get_step_ctx(fn: str) -> StepWorkerContext:
+        """Get the current step worker context. Raises if not in a step."""
+        try:
+            return StepWorkerStateContextVar.get()
+        except LookupError:
+            raise WorkflowRuntimeError(
+                f"{fn} may only be called from within a step function"
+            )
+
+    @property
+    def store(self) -> InMemoryStateStore[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)
+
+    def collect_events(
+        self,
+        ev: Event,
+        expected: list[Type[Event]],
+        buffer_id: str | None = None,
+    ) -> list[Event] | None:
+        """Collect events until all expected types are received."""
+        step_ctx = self._get_step_ctx(fn="collect_events")
+
+        # If no events are expected, return an empty list immediately
+        if not expected:
+            return []
+
+        buffer_id = buffer_id or "default"
+        collected_events = step_ctx.state.collected_events.get(buffer_id, [])
+
+        remaining_event_types = Counter(expected) - Counter(
+            [type(e) for e in collected_events]
+        )
+
+        if remaining_event_types != Counter([type(ev)]):
+            if type(ev) in remaining_event_types:
+                step_ctx.returns.return_values.append(
+                    AddCollectedEvent(event_id=buffer_id, event=ev)
+                )
+            return None
+
+        total = []
+        by_type: dict[Type[Event], list[Event]] = defaultdict(list)
+        for e in collected_events + [ev]:
+            by_type[type(e)].append(e)
+        # order by expected type
+        for e_type in expected:
+            total.append(by_type[e_type].pop(0))
+        # Clear the collected events when the step is complete
+        step_ctx.returns.return_values.append(DeleteCollectedEvent(event_id=buffer_id))
+        return total
+
+    def send_event(self, message: Event, step: str | None = None) -> None:
+        """Send an event to trigger another step."""
+        if step is not None:
+            self._workflow._validate_valid_step_message(step, message)
+
+        self._execute_task(
+            self._internal_adapter.send_event(
+                TickAddEvent(event=message, step_name=step)
+            )
+        )
+
+    async def wait_for_event(
+        self,
+        event_type: Type[T],
+        waiter_event: Event | None = None,
+        waiter_id: str | None = None,
+        requirements: dict[str, Any] | None = None,
+        timeout: float | None = 2000,
+    ) -> T:
+        """Wait for an event of the specified type."""
+        step_ctx = self._get_step_ctx(fn="wait_for_event")
+
+        collected_waiters = step_ctx.state.collected_waiters
+        requirements = requirements or {}
+
+        # Generate a unique key for the waiter
+        event_str = f"{event_type.__module__}.{event_type.__name__}"
+        requirements_str = str(requirements)
+        waiter_id = waiter_id or f"waiter_{event_str}_{requirements_str}"
+
+        waiter = next((w for w in collected_waiters if w.waiter_id == waiter_id), None)
+        if waiter is None or waiter.resolved_event is None:
+            raise WaitingForEvent(
+                AddWaiter(
+                    waiter_id=waiter_id,
+                    requirements=requirements,
+                    timeout=timeout,
+                    event_type=event_type,
+                    waiter_event=waiter_event,
+                )
+            )
+        else:
+            step_ctx.returns.return_values.append(DeleteWaiter(waiter_id=waiter_id))
+            return cast(T, waiter.resolved_event)
+
+    def write_event_to_stream(self, ev: Event | None) -> None:
+        """Write an event to the published event stream."""
+        if ev is not None:
+            self._execute_task(self._internal_adapter.write_to_event_stream(ev))

packages/llama-index-workflows/src/workflows/context/pre_context.py

@@ -0,0 +1,115 @@
+"""Pre-run context - configuration face before workflow execution."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Generic, cast
+
+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.errors import ContextSerdeError
+from workflows.runtime.types.internal_state import BrokerState
+
+if TYPE_CHECKING:
+    from workflows.workflow import Workflow
+
+
+class PreContext(Generic[MODEL_T]):
+    """Context state before workflow starts.
+
+    Provides access to workflow configuration and serialization
+    for persistence/restoration. A staging store is lazily created
+    on first `.store` access and carried into the runtime when the
+    workflow starts.
+    """
+
+    _init_snapshot: SerializedContext
+    _serializer: BaseSerializer
+    _workflow: "Workflow"
+    _store: InMemoryStateStore[MODEL_T] | None
+
+    def __init__(
+        self,
+        workflow: "Workflow",
+        previous_context: dict[str, Any] | None = None,
+        serializer: BaseSerializer | None = None,
+    ) -> None:
+        self._serializer = serializer or JsonSerializer()
+        self._workflow = workflow
+        self._store = None
+
+        # Parse the serialized context
+        if previous_context is not None:
+            try:
+                # Auto-detect and convert V0 to V1 if needed
+                previous_context_parsed = SerializedContext.from_dict_auto(
+                    previous_context
+                )
+                # Validate it fully parses synchronously to avoid delayed validation errors
+                BrokerState.from_serialized(
+                    previous_context_parsed, workflow, self._serializer
+                )
+            except ValidationError as e:
+                raise ContextSerdeError(
+                    f"Context dict specified in an invalid format: {e}"
+                ) from e
+        else:
+            previous_context_parsed = SerializedContext()
+
+        self._init_snapshot = previous_context_parsed
+
+    @property
+    def store(self) -> InMemoryStateStore[MODEL_T]:
+        """Lazily-created staging store for pre-run state access.
+
+        For fresh contexts, the state type is inferred from workflow step
+        annotations. For deserialized contexts, the store is restored from
+        the serialized state data.
+        """
+        if self._store is None:
+            serialized_state = self._init_snapshot.state
+            if serialized_state:
+                self._store = cast(
+                    InMemoryStateStore[MODEL_T],
+                    InMemoryStateStore.from_dict(serialized_state, self._serializer),
+                )
+            else:
+                state_type = infer_state_type(self._workflow)
+                self._store = cast(
+                    InMemoryStateStore[MODEL_T],
+                    InMemoryStateStore(state_type()),
+                )
+        return self._store
+
+    @property
+    def serialized_state(self) -> dict[str, Any] | None:
+        """Return the serialized state for handoff to the runtime.
+
+        If the staging store was accessed, its current contents are
+        serialized.  Otherwise the snapshot's original state is returned
+        unchanged, avoiding unnecessary work.
+        """
+        if self._store is not None:
+            return self._store.to_dict(self._serializer)
+        return self._init_snapshot.state
+
+    @property
+    def is_running(self) -> bool:
+        """Whether the workflow is currently running.
+
+        Returns the is_running state from the init snapshot, which may be True
+        if restoring a context that was previously mid-run.
+        """
+        return self._init_snapshot.is_running
+
+    @property
+    def init_snapshot(self) -> SerializedContext:
+        """The initial serialized context snapshot."""
+        return self._init_snapshot
+
+    @property
+    def serializer(self) -> BaseSerializer:
+        """The serializer used for this context."""
+        return self._serializer

packages/llama-index-workflows/src/workflows/context/state_store.py

@@ -1,15 +1,25 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 LlamaIndex Inc.
+
+from __future__ import annotations
+
 import asyncio
+import functools
 import warnings
 from contextlib import asynccontextmanager
-from typing import Any, AsyncGenerator, Generic, Optional, Type
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Generic, Type
 
 from pydantic import BaseModel, ValidationError
 from typing_extensions import TypeVar
 
+from workflows.decorators import StepConfig
 from workflows.events import DictLikeModel
 
 from .serializers import BaseSerializer
 
+if TYPE_CHECKING:
+    from workflows.workflow import Workflow
+
 MAX_DEPTH = 1000
 
 
@@ -99,9 +109,17 @@ class InMemoryStateStore(Generic[MODEL_T]):
 
     def __init__(self, initial_state: MODEL_T):
         self._state = initial_state
-        self._lock = asyncio.Lock()
         self.state_type = type(initial_state)
 
+    @functools.cached_property
+    def _lock(self) -> asyncio.Lock:
+        """Lazy lock initialization for Python 3.14+ compatibility.
+
+        asyncio.Lock() requires a running event loop in Python 3.14+.
+        Using cached_property defers creation to first use in async context.
+        """
+        return asyncio.Lock()
+
     async def get_state(self) -> MODEL_T:
         """Return a shallow copy of the current state model.
 
@@ -255,7 +273,7 @@ class InMemoryStateStore(Generic[MODEL_T]):
 
             self._state = state
 
-    async def get(self, path: str, default: Optional[Any] = Ellipsis) -> Any:
+    async def get(self, path: str, default: Any = Ellipsis) -> Any:
         """Get a nested value using dot-separated paths.
 
         Supports dict keys, list indices, and attribute access transparently at
@@ -370,3 +388,84 @@ class InMemoryStateStore(Generic[MODEL_T]):
 
         # fallback to attribute assignment
         setattr(obj, segment, value)
+
+
+def infer_state_type(workflow: "Workflow") -> type[BaseModel]:
+    """Infer the state type from workflow step configs.
+
+    Looks at Context[T] annotations in step functions to determine
+    the expected state type. Returns DictState if no typed state is found.
+
+    Args:
+        workflow: The workflow to inspect for state type annotations.
+
+    Returns:
+        The inferred state type, or DictState if none found.
+
+    Raises:
+        ValueError: If multiple different state types are found.
+    """
+    state_types: set[type[BaseModel]] = set()
+    for _, step_func in workflow._get_steps().items():
+        step_config: StepConfig = step_func._step_config
+        if (
+            step_config.context_state_type is not None
+            and step_config.context_state_type != DictState
+            and issubclass(step_config.context_state_type, BaseModel)
+        ):
+            state_types.add(step_config.context_state_type)
+
+    state_type: Type[BaseModel]
+    if state_types:
+        state_type = _find_most_derived_state_type(state_types)
+    else:
+        state_type = DictState
+
+    return state_type
+
+
+def _find_most_derived_state_type(state_types: set[Type[BaseModel]]) -> Type[BaseModel]:
+    """Find the most derived (most specific) state type from a set of types.
+
+    All types must be in a single inheritance chain, i.e., one type must be
+    a subclass of all other types (the most derived type).
+
+    Args:
+        state_types: Set of state types to analyze.
+
+    Returns:
+        The most derived type in the inheritance hierarchy.
+
+    Raises:
+        ValueError: If types are not in a compatible inheritance hierarchy.
+    """
+    type_list = list(state_types)
+
+    if len(type_list) == 1:
+        return type_list[0]
+
+    # Find the most derived type - it should be a subclass of all others
+    most_derived: Type[BaseModel] | None = None
+
+    for candidate in type_list:
+        is_most_derived = True
+        for other in type_list:
+            if other is candidate:
+                continue
+            # candidate must be a subclass of other (or equal to it)
+            if not issubclass(candidate, other):
+                is_most_derived = False
+                break
+        if is_most_derived:
+            most_derived = candidate
+            break
+
+    if most_derived is None:
+        # No single type is a subclass of all others - incompatible hierarchy
+        raise ValueError(
+            "Multiple state types are not in a compatible inheritance hierarchy. "
+            "All state types must share a common inheritance chain. Found: "
+            + ", ".join([st.__name__ for st in state_types])
+        )
+
+    return most_derived

packages/llama-index-workflows/src/workflows/errors.py

@@ -32,3 +32,13 @@ class WorkflowConfigurationError(Exception):
 
 class ContextSerdeError(Exception):
     """Raised when serializing/deserializing a `Context` fails."""
+
+
+class ContextStateError(Exception):
+    """Raised when a context method is called in the wrong state.
+
+    Context transitions between three states:
+    - PreContext: Before workflow starts (configuration)
+    - ExternalContext: During run, for handler/external code
+    - InternalContext: During run, for step execution
+    """

packages/llama-index-workflows/src/workflows/handler.py

@@ -4,47 +4,82 @@
 from __future__ import annotations
 
 import asyncio
-from typing import TYPE_CHECKING, Any, AsyncGenerator
+import functools
+import warnings
+from collections.abc import Generator
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Awaitable
 
-from .errors import WorkflowRuntimeError
-from .events import Event, InternalDispatchEvent, StopEvent
+from workflows.runtime.types.plugin import (
+    ExternalRunAdapter,
+    as_v2_runtime_compatibility_shim,
+)
+
+from .errors import WorkflowCancelledByUser, WorkflowRuntimeError
+from .events import Event, InternalDispatchEvent, StopEvent, WorkflowCancelledEvent
 from .types import RunResultT
 
 if TYPE_CHECKING:
     from .context import Context
+    from .workflow import Workflow
 
 
-class WorkflowHandler(asyncio.Future[RunResultT]):
+class WorkflowHandler(Awaitable[RunResultT]):
     """
-    Handle a running workflow: await results, stream events, access context, or cancel.
-
-    Instances are returned by [Workflow.run][workflows.workflow.Workflow.run].
-    They can be awaited for the final result and support streaming intermediate
-    events via [stream_events][workflows.handler.WorkflowHandler.stream_events].
-
-    See Also:
-        - [Context][workflows.context.context.Context]
-        - [StopEvent][workflows.events.StopEvent]
+    Stable interface for communicating with a running workflow. Is awaitable and streamable, and supports things like cancellation.
     """
 
     _ctx: Context
-    _run_task: asyncio.Task[None] | None
-    _all_events_consumed: bool
-    _stop_event: StopEvent | None
+
+    async def _await_result(self) -> RunResultT:
+        stop_event = await self.stop_event_result()
+        return stop_event.result if type(stop_event) is StopEvent else stop_event
+
+    def __await__(self) -> Generator[Any, Any, RunResultT]:
+        return self._await_result().__await__()
 
     def __init__(
         self,
-        *args: Any,
-        ctx: Context,
-        run_id: str | None = None,
-        run_task: asyncio.Task[None] | None = None,
-        **kwargs: Any,
+        workflow: "Workflow",
+        external_adapter: ExternalRunAdapter,
+        ctx: "Context[Any] | None" = None,
     ) -> None:
-        super().__init__(*args, **kwargs)
-        self.run_id = run_id
-        self._ctx = ctx
-        self._run_task = run_task
+        from .context import Context
+
+        self._workflow = workflow
+        self._external_adapter = external_adapter
+        # TODO(v3): Remove ctx parameter. The handler will just be the external face.
+        self._ctx = (
+            ctx
+            if ctx is not None
+            else Context._create_external(
+                workflow=workflow, external_adapter=external_adapter
+            )
+        )
+        self.run_id = external_adapter.run_id
         self._all_events_consumed = False
+        self._result: StopEvent | None = None
+        self._result_exception: BaseException | None = None
+        self._result_task = asyncio.create_task(self._wait_for_result())
+        self._result_task.add_done_callback(self._handle_result_task_done)
+
+    async def _wait_for_result(self) -> StopEvent:
+        result = await self._external_adapter.get_result()
+        self._result = result
+        return result
+
+    def _handle_result_task_done(self, task: asyncio.Task[StopEvent]) -> None:
+        if task.cancelled():
+            return
+        try:
+            exc = task.exception()
+        except asyncio.CancelledError:
+            return
+        if exc is None:
+            return
+        self._result_exception = exc
+        if isinstance(exc, WorkflowCancelledByUser) and self._result is None:
+            # Preserve cancellation in handler state without changing await semantics.
+            self._result = WorkflowCancelledEvent()
 
     @property
     def ctx(self) -> Context:
@@ -53,31 +88,18 @@ class WorkflowHandler(asyncio.Future[RunResultT]):
 
     def get_stop_event(self) -> StopEvent | None:
         """The stop event for this run. Always defined once the future is done. In a future major release, this will be removed, and the result will be the stop event itself."""
-        return self._stop_event
+        return self._result
 
     async def stop_event_result(self) -> StopEvent:
         """Get the stop event for this run. Always defined once the future is done. In a future major release, this will be removed, and the result will be the stop event itself."""
-        await self.result()
-        assert self._stop_event is not None, (
-            "Stop event must be defined once the future is done."
-        )
-        return self._stop_event
-
-    def _set_stop_event(self, stop_event: StopEvent) -> None:
-        self._stop_event = stop_event
-        # sad but necessary legacy behavior:
-        # set the result to the stop event result. To be removed in a future major release,
-        # and justuse the stop event directly.
-        self.set_result(
-            stop_event.result if type(stop_event) is StopEvent else stop_event
-        )
+        return await self._result_task
 
     def __str__(self) -> str:
-        return str(self.result())
+        return f"WorkflowHandler(workflow={self._workflow.workflow_name}, run_id={self.run_id}, result={self._result})"
 
     def is_done(self) -> bool:
         """Return True when the workflow has completed."""
-        return self.done()
+        return self._result_task.done()
 
     async def stream_events(
         self, expose_internal: bool = False
@@ -139,7 +161,45 @@ class WorkflowHandler(asyncio.Future[RunResultT]):
                 self._all_events_consumed = True
                 break
 
-    async def cancel_run(self) -> None:
+    def done(self) -> bool:
+        """Return True when the workflow has completed."""
+        _warn_done_deprecated()
+        return self._result_task.done()
+
+    def cancel(self) -> None:
+        """Cancel the running workflow."""
+        _warn_cancel_deprecated()
+        shim = as_v2_runtime_compatibility_shim(self._external_adapter)
+        if shim is None:
+            raise NotImplementedError(
+                "Hard cancel is not supported by this runtime. "
+                "Use await handler.cancel_run() for graceful cancellation."
+            )
+        shim.abort()
+        self._result_task.cancel()
+
+    def exception(self) -> BaseException | None:
+        """Get the exception for this run. Always defined once the future is done."""
+        _warn_exception_deprecated()
+        try:
+            return self._result_task.exception()
+        except asyncio.CancelledError:
+            return None
+
+    def cancelled(self) -> bool:
+        """Return True when the underlying workflow has been cancelled."""
+        _warn_cancelled_deprecated()
+        if self._result_task.cancelled():
+            return True
+        exc = self.exception()
+        if exc is not None and isinstance(exc, WorkflowCancelledByUser):
+            return True
+        stop_event = self.get_stop_event()
+        if stop_event is not None and isinstance(stop_event, WorkflowCancelledEvent):
+            return True
+        return False
+
+    async def cancel_run(self, *, timeout: float = 5.0) -> None:
         """Cancel the running workflow.
 
         Signals the underlying context to raise
@@ -152,10 +212,60 @@ class WorkflowHandler(asyncio.Future[RunResultT]):
             await handler.cancel_run()
             ```
         """
-        if self.ctx:
-            self.ctx._workflow_cancel_run()
-            if self._run_task is not None:
-                try:
-                    await self._run_task
-                except Exception:
-                    pass
+        try:
+            await self._external_adapter.cancel()
+        except Exception:
+            pass
+        try:
+            await asyncio.wait_for(self._result_task, timeout=timeout)
+        except asyncio.TimeoutError:
+            pass
+        except asyncio.CancelledError:
+            pass
+        except Exception:
+            pass
+
+    async def send_event(self, event: Event, step: str | None = None) -> None:
+        """Send an event into the workflow.
+
+        Args:
+            event: The event to send into the workflow.
+            step: Optional step name to target. If None, broadcasts to all.
+        """
+        self.ctx.send_event(event, step)
+
+
+@functools.lru_cache(maxsize=1)
+def _warn_done_deprecated() -> None:
+    warnings.warn(
+        "WorkflowHandler.done() is deprecated and will be removed in a future release",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+
+
+@functools.lru_cache(maxsize=1)
+def _warn_cancel_deprecated() -> None:
+    warnings.warn(
+        "WorkflowHandler.cancel() is deprecated and will be removed in a future release. Prefer to cancel the underlying workflow with await handler.cancel_run(), and then awaiting the result with await handler to obtain the cancellation exception.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+
+
+@functools.lru_cache(maxsize=1)
+def _warn_exception_deprecated() -> None:
+    warnings.warn(
+        "WorkflowHandler.exception() is deprecated and will be removed in a future release",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+
+
+@functools.lru_cache(maxsize=1)
+def _warn_cancelled_deprecated() -> None:
+    warnings.warn(
+        "WorkflowHandler.cancelled() is deprecated and will be removed in a future release",
+        DeprecationWarning,
+        stacklevel=2,
+    )

packages/llama-index-workflows/src/workflows/plugins/__init__.py

@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 LlamaIndex Inc.
+"""Workflow runtime implementations."""
+
+from workflows.plugins._context import get_current_runtime
+from workflows.plugins.basic import BasicRuntime, basic_runtime
+
+__all__ = [
+    "get_current_runtime",
+    "basic_runtime",
+    "BasicRuntime",
+]

packages/llama-index-workflows/src/workflows/plugins/_context.py

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 LlamaIndex Inc.
+"""Context-scoped runtime access."""
+
+from __future__ import annotations
+
+from workflows.runtime.types.plugin import Runtime, _current_runtime
+
+
+def get_current_runtime() -> Runtime:
+    """
+    Get the current runtime from context or fall back to basic_runtime.
+
+    Returns the context-scoped runtime if set, otherwise returns basic_runtime.
+    """
+    # Inline import to avoid circular dependency (basic -> runtime -> workflow)
+    from workflows.plugins.basic import basic_runtime
+
+    runtime = _current_runtime.get()
+    return runtime if runtime is not None else basic_runtime

packages/llama-index-workflows/src/workflows/plugins/basic.py

@@ -4,85 +4,98 @@
 from __future__ import annotations
 
 import asyncio
+import functools
 import time
-from typing import AsyncGenerator, Callable
+import weakref
+from contextlib import asynccontextmanager, contextmanager
+from contextvars import ContextVar
+from typing import Any, AsyncGenerator, Generator
 
-from workflows.decorators import P, R
-from workflows.events import Event, StopEvent
+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.errors import WorkflowRuntimeError
+from workflows.events import Event, StartEvent, StopEvent
+from workflows.runtime.types.internal_state import BrokerState
 from workflows.runtime.types.plugin import (
-    ControlLoopFunction,
-    Plugin,
+    ExternalRunAdapter,
+    InternalRunAdapter,
     RegisteredWorkflow,
-    SnapshottableRuntime,
-    WorkflowRuntime,
+    Runtime,
+    SnapshottableAdapter,
+    V2RuntimeCompatibilityShim,
+)
+from workflows.runtime.types.step_function import (
+    as_step_worker_functions,
+    create_workflow_run_function,
 )
-from workflows.runtime.types.step_function import StepWorkerFunction
 from workflows.runtime.types.ticks import WorkflowTick
 from workflows.workflow import Workflow
 
 
-class BasicRuntime:
-    def register(
-        self,
-        workflow: Workflow,
-        workflow_function: ControlLoopFunction,
-        steps: dict[str, StepWorkerFunction],
-    ) -> None | RegisteredWorkflow:
-        return None
+class AsyncioAdapterQueues:
+    """Shared state between internal and external adapters.
 
-    def new_runtime(self, run_id: str) -> WorkflowRuntime:
-        snapshottable: SnapshottableRuntime = AsyncioWorkflowRuntime(run_id)
-        return snapshottable
-
-
-basic_runtime: Plugin = BasicRuntime()
-
-
-class AsyncioWorkflowRuntime:
-    """
-    A plugin interface to switch out a broker runtime (external library or service that manages durable/distributed step execution)
+    The `complete` task is set by run_workflow() after instantiation due to
+    circular dependency: the task closure captures this object to prevent
+    premature GC from the WeakValueDictionary.
     """
 
+    # Set by run_workflow() after task creation
+    complete: asyncio.Task[StopEvent]
+
     def __init__(
         self,
         run_id: str,
-    ) -> None:
+        init_state: BrokerState,
+        state_store: InMemoryStateStore[Any] | None = None,
+    ):
         self.run_id = run_id
-        self.receive_queue: asyncio.Queue[WorkflowTick] = asyncio.Queue()
-        self.publish_queue: asyncio.Queue[Event] = asyncio.Queue()
+        self.init_state = init_state
         self.ticks: list[WorkflowTick] = []
+        self.state_store = state_store
 
-    def on_tick(self, tick: WorkflowTick) -> None:
-        self.ticks.append(tick)
+    # created lazily via cached_property for Python 3.14+ compatibility (they require a running event loop)
+    @functools.cached_property
+    def receive_queue(self) -> asyncio.Queue[WorkflowTick]:
+        return asyncio.Queue[WorkflowTick]()
 
-    def replay(self) -> list[WorkflowTick]:
-        return self.ticks
+    # created lazily via cached_property for Python 3.14+ compatibility (they require a running event loop)
+    @functools.cached_property
+    def publish_queue(self) -> asyncio.Queue[Event]:
+        return asyncio.Queue[Event]()
+
+    # created lazily via cached_property for Python 3.14+ compatibility (they require a running event loop)
+    @functools.cached_property
+    def stream_lock(self) -> asyncio.Lock:
+        return asyncio.Lock()
+
+
+class InternalAsyncioAdapter(InternalRunAdapter, SnapshottableAdapter):
+    """
+    Internal adapter for asyncio-based workflow execution.
+
+    Used by the workflow control loop to receive ticks, publish events,
+    and manage timing. Also supports snapshotting for debugging/replay.
+    """
+
+    def __init__(self, queues: AsyncioAdapterQueues) -> None:
+        self._queues = queues
+
+    @property
+    def run_id(self) -> str:
+        return self._queues.run_id
+
+    @property
+    def init_state(self) -> BrokerState:
+        return self._queues.init_state
 
     async def wait_receive(self) -> WorkflowTick:
-        return await self.receive_queue.get()
+        return await self._queues.receive_queue.get()
 
     async def write_to_event_stream(self, event: Event) -> None:
-        self.publish_queue.put_nowait(event)
-
-    async def stream_published_events(self) -> AsyncGenerator[Event, None]:
-        while True:
-            item = await self.publish_queue.get()
-            yield item
-            if isinstance(item, StopEvent):
-                break
-
-    async def send_event(self, tick: WorkflowTick) -> None:
-        self.receive_queue.put_nowait(tick)
-
-    async def register_step_worker(
-        self, step_name: str, step_worker: StepWorkerFunction[R]
-    ) -> StepWorkerFunction[R]:
-        return step_worker
-
-    async def register_workflow_function(
-        self, workflow_function: Callable[P, R]
-    ) -> Callable[P, R]:
-        return workflow_function
+        self._queues.publish_queue.put_nowait(event)
 
     async def get_now(self) -> float:
         return time.monotonic()
@@ -90,5 +103,217 @@ class AsyncioWorkflowRuntime:
     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)
+
+    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:
+        return self._queues.state_store
+
+
+class ExternalAsyncioAdapter(
+    ExternalRunAdapter, SnapshottableAdapter, V2RuntimeCompatibilityShim
+):
+    """
+    External adapter for asyncio-based workflow execution.
+
+    Used by external code to send events into the workflow
+    and stream events published by the workflow.
+    """
+
+    def __init__(self, queues: AsyncioAdapterQueues) -> None:
+        self._queues = queues
+
+    @property
+    def run_id(self) -> str:
+        return self._queues.run_id
+
+    async def send_event(self, tick: WorkflowTick) -> None:
+        self._queues.receive_queue.put_nowait(tick)
+
+    async def stream_published_events(self) -> AsyncGenerator[Event, None]:
+        async with self._queues.stream_lock:
+            if self._queues.complete.done() and self._queues.publish_queue.empty():
+                raise WorkflowRuntimeError(
+                    "Event stream already consumed. "
+                    "Events can only be streamed once per workflow run."
+                )
+            while True:
+                item = await self._queues.publish_queue.get()
+                yield item
+                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:
+        return self._queues.state_store
+
+    async def get_result(self) -> StopEvent:
+        return await self._queues.complete
+
+    def get_result_or_none(self) -> StopEvent | None:
+        if not self._queues.complete.done():
+            return None
+        return self._queues.complete.result()
+
+    @property
+    def is_running(self) -> bool:
+        return not self._queues.complete.done()
+
+    def abort(self) -> None:
+        """Abort by cancelling the control loop task."""
+        if not self._queues.complete.done():
+            self._queues.complete.cancel()
+
+    @property
+    def init_state(self) -> BrokerState:
+        return self._queues.init_state
+
+
+class BasicRuntime(Runtime):
+    """Default asyncio-based runtime with no durability."""
+
+    def __init__(self) -> None:
+        # WeakValueDictionary allows queues to be GC'd when no adapters reference them.
+        # The task closure in run_workflow() captures a strong reference, keeping
+        # queues alive for fire-and-forget workflows even if the external adapter is dropped.
+        self._queues: weakref.WeakValueDictionary[str, AsyncioAdapterQueues] = (
+            weakref.WeakValueDictionary()
+        )
+        # Keyed by id(workflow) so each instance has independent concurrency limits
+        self._max_concurrent_runs: weakref.WeakValueDictionary[
+            int, asyncio.Semaphore
+        ] = weakref.WeakValueDictionary()
+
+    def register(self, workflow: Workflow) -> RegisteredWorkflow:
+        return RegisteredWorkflow(
+            workflow=workflow,
+            workflow_run_fn=create_workflow_run_function(workflow),
+            steps=as_step_worker_functions(workflow),
+        )
+
+    def _get_or_create_queues(
+        self, run_id: str, init_state: BrokerState
+    ) -> AsyncioAdapterQueues:
+        """Get existing queues or create new ones for a run_id."""
+        queues = self._queues.get(run_id)
+        if queues is None:
+            queues = AsyncioAdapterQueues(run_id=run_id, init_state=init_state)
+            self._queues[run_id] = queues
+        return queues
+
+    @asynccontextmanager
+    async def _maybe_acquire_max_concurrent_runs(
+        self, workflow: Workflow, run_id: str
+    ) -> AsyncGenerator[None, None]:
+        if workflow._num_concurrent_runs is None:
+            yield
+        else:
+            # Key by instance id so each workflow instance has independent concurrency limits
+            workflow_id = id(workflow)
+            if workflow_id in self._max_concurrent_runs:
+                sem = self._max_concurrent_runs[workflow_id]
+            else:
+                sem = asyncio.Semaphore(workflow._num_concurrent_runs)
+                self._max_concurrent_runs[workflow_id] = sem
+            async with sem:
+                yield
+
+    def run_workflow(
+        self,
+        run_id: str,
+        workflow: Workflow,
+        init_state: BrokerState,
+        start_event: StartEvent | None = None,
+        serialized_state: dict[str, Any] | None = None,
+        serializer: BaseSerializer | None = None,
+    ) -> ExternalRunAdapter:
+        """Set up a workflow run. Currently only creates state store.
+
+        Note: Execution is still managed by the broker for now. This will
+        change as we refactor to have the runtime fully own execution.
+        """
+        if run_id in self._queues:
+            # not supported in any way right now. Might make sense to support run as new, or some other idempotency semantics
+            raise RuntimeError(f"Workflow run with run_id '{run_id}' already exists.")
+
+        registered = self.get_or_register(workflow)
+
+        # Create state store from serialized state or infer type from workflow
+        active_serializer = serializer or JsonSerializer()
+        if serialized_state:
+            state_store = InMemoryStateStore.from_dict(
+                serialized_state, active_serializer
+            )
+        else:
+            # Infer state type from workflow step configs
+            state_type = infer_state_type(registered.workflow)
+            state_store = InMemoryStateStore(state_type())
+        # might want to lock this better. Unlikely race condition if you spam with the same run_id.
+        queues = self._get_or_create_queues(run_id, init_state)
+        queues.state_store = state_store
+
+        async def run_with_concurrency_limit() -> StopEvent:
+            # Capture strong reference to queues for the task's lifetime,
+            # enabling fire-and-forget even if the caller drops the external adapter.
+            _ = queues
+            async with self._maybe_acquire_max_concurrent_runs(workflow, run_id):
+                return await registered.workflow_run_fn(
+                    init_state, start_event, active_instrument_tags.get()
+                )
+
+        with setting_run_id(run_id):
+            # actually pump the task through the runtime
+            task = asyncio.create_task(run_with_concurrency_limit())
+            queues.complete = task
+            return self.get_external_adapter(run_id)
+
+    def get_internal_adapter(self) -> InternalRunAdapter:
+        run_id = get_current_run_id()
+        if run_id is None:
+            raise RuntimeError(
+                "No current run id. Must be called within a workflow run."
+            )
+        if run_id not in self._queues:
+            raise RuntimeError(
+                f"No queues found for run_id '{run_id}'. Must be called within a workflow run."
+            )
+        queues = self._queues[run_id]
+        return InternalAsyncioAdapter(queues)
+
+    def get_external_adapter(self, run_id: str) -> ExternalRunAdapter:
+        if run_id not in self._queues:
+            raise RuntimeError(f"No active workflow with run_id '{run_id}'. ")
+        return ExternalAsyncioAdapter(self._queues[run_id])
+
+
+_current_run_id: ContextVar[str | None] = ContextVar("current_run_id", default=None)
+
+
+def get_current_run_id() -> str | None:
+    return _current_run_id.get()
+
+
+@contextmanager
+def setting_run_id(run_id: str) -> Generator[None, None, None]:
+    token = _current_run_id.set(run_id)
+    try:
+        yield
+    finally:
+        _current_run_id.reset(token)
+
+
+basic_runtime = BasicRuntime()

packages/llama-index-workflows/src/workflows/resource.py

@@ -161,10 +161,19 @@ class _Resource(Generic[T]):
         return deps
 
 
-@functools.lru_cache(maxsize=1)
 def _get_resource_config_data(
     config_file: str,
     path_selector: str | None,
+) -> dict[str, Any]:
+    # Resolve to absolute path for cache key to handle different working directories
+    abs_path = str(Path(config_file).resolve())
+    return _get_resource_config_data_cached(abs_path, path_selector)
+
+
+@functools.lru_cache(maxsize=128)
+def _get_resource_config_data_cached(
+    config_file: str,
+    path_selector: str | None,
 ) -> dict[str, Any]:
     with open(config_file, "r") as f:
         data = json.load(f)

packages/llama-index-workflows/src/workflows/runtime/broker.py

@@ -1,353 +0,0 @@
-# SPDX-License-Identifier: MIT
-# Copyright (c) 2026 LlamaIndex Inc.
-
-from __future__ import annotations
-
-import asyncio
-import logging
-from collections import Counter, defaultdict
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    AsyncGenerator,
-    Awaitable,
-    Callable,
-    Coroutine,
-    Generic,
-    Type,
-    TypeVar,
-    cast,
-)
-
-from llama_index_instrumentation.dispatcher import (
-    active_instrument_tags,
-    instrument_tags,
-)
-from workflows.errors import WorkflowRuntimeError
-from workflows.events import (
-    Event,
-    StartEvent,
-)
-from workflows.handler import WorkflowHandler
-from workflows.runtime.control_loop import control_loop, rebuild_state_from_ticks
-from workflows.runtime.types.internal_state import BrokerState
-from workflows.runtime.types.plugin import Plugin, WorkflowRuntime, as_snapshottable
-from workflows.runtime.types.results import (
-    AddCollectedEvent,
-    AddWaiter,
-    DeleteCollectedEvent,
-    DeleteWaiter,
-    StepWorkerContext,
-    StepWorkerStateContextVar,
-    WaitingForEvent,
-)
-from workflows.runtime.types.step_function import (
-    StepWorkerFunction,
-    as_step_worker_function,
-)
-from workflows.runtime.types.ticks import TickAddEvent, TickCancelRun, WorkflowTick
-from workflows.runtime.workflow_registry import workflow_registry
-from workflows.utils import _nanoid as nanoid
-
-from ..context.state_store import MODEL_T
-
-if TYPE_CHECKING:
-    from workflows import Workflow
-    from workflows.context.context import Context
-
-
-T = TypeVar("T", bound=Event)
-EventBuffer = dict[str, list[Event]]
-
-logger = logging.getLogger()
-
-
-# Only warn once about unserializable keys
-class UnserializableKeyWarning(Warning):
-    pass
-
-
-class WorkflowBroker(Generic[MODEL_T]):
-    """
-    The workflow broker manages starting up and connecting a workflow handler, a runtime, and triggering the
-    execution of the workflow. From there it manages communication between the workflow and the outside world.
-    """
-
-    _context: Context[MODEL_T]
-    _runtime: WorkflowRuntime
-    _plugin: Plugin
-    _is_running: bool
-    _handler: WorkflowHandler | None
-    _workflow: Workflow
-    # transient tasks to run async ops in background, exposing sync interfaces
-    _workers: list[asyncio.Task]
-    _init_state: BrokerState | None
-
-    def __init__(
-        self,
-        workflow: Workflow,
-        context: Context[MODEL_T],
-        runtime: WorkflowRuntime,
-        plugin: Plugin,
-    ) -> None:
-        self._context = context
-        self._runtime = runtime
-        self._plugin = plugin
-        self._is_running = False
-        self._handler = None
-        self._workflow = workflow
-        self._workers = []
-        self._init_state = None
-
-    def _execute_task(self, coro: Coroutine[Any, Any, Any]) -> asyncio.Task[Any]:
-        task = asyncio.create_task(coro)
-        self._workers.append(task)
-
-        def _remove_task(_: asyncio.Task[Any]) -> None:
-            try:
-                self._workers.remove(task)
-            except ValueError:
-                # Handle Task was already cleared during shutdown or cleanup.
-                pass
-
-        task.add_done_callback(_remove_task)
-        return task
-
-    # context API only
-    def start(
-        self,
-        workflow: Workflow,
-        previous: BrokerState | None = None,
-        start_event: StartEvent | None = None,
-        before_start: Callable[[], Awaitable[None]] | None = None,
-        after_complete: Callable[[], Awaitable[None]] | None = None,
-    ) -> WorkflowHandler:
-        """Start the workflow run. Can only be called once."""
-        if self._handler is not None:
-            raise WorkflowRuntimeError(
-                "this WorkflowBroker already run or running. Cannot start again."
-            )
-        self._init_state = previous
-
-        async def _run_workflow(run_id: str, tags: dict[str, Any]) -> None:
-            with instrument_tags({"run_id": run_id, **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
-                self._is_running = True
-                await asyncio.sleep(0)
-                if before_start is not None:
-                    await before_start()
-                try:
-                    init_state = previous or BrokerState.from_workflow(workflow)
-
-                    try:
-                        exception_raised = None
-
-                        step_workers: dict[str, StepWorkerFunction] = {}
-                        for name, step_func in workflow._get_steps().items():
-                            # Avoid capturing a bound method (which retains the instance).
-                            # If it's a bound method, extract the unbound function from the class.
-                            unbound = getattr(step_func, "__func__", step_func)
-                            step_workers[name] = as_step_worker_function(unbound)
-
-                        registered = workflow_registry.get_registered_workflow(
-                            workflow, self._plugin, control_loop, step_workers
-                        )
-
-                        # Register run context prior to invoking control loop
-                        workflow_registry.register_run(
-                            run_id=run_id,
-                            workflow=workflow,
-                            plugin=self._runtime,
-                            context=self._context,  # type: ignore
-                            steps=registered.steps,
-                        )
-
-                        try:
-                            workflow_result = await registered.workflow_function(
-                                start_event,
-                                init_state,
-                                run_id,
-                            )
-                        finally:
-                            # ensure run context is cleaned up even on failure
-                            workflow_registry.delete_run(run_id)
-                        result._set_stop_event(workflow_result)
-                    except Exception as e:
-                        exception_raised = e
-
-                    if exception_raised:
-                        # cancel the stream
-                        if not result.done():
-                            result.set_exception(exception_raised)
-                finally:
-                    if after_complete is not None:
-                        await after_complete()
-                    self._is_running = False
-
-        # Start the machinery in a new Context or use the provided one
-        run_id = nanoid()
-
-        # If a previous context is provided, pass its serialized form
-
-        run_task = self._execute_task(
-            _run_workflow(run_id, tags=active_instrument_tags.get())
-        )
-        result = WorkflowHandler(
-            ctx=self._context,  # type: ignore
-            run_id=run_id,
-            run_task=run_task,
-        )
-        self._handler = result
-        return result
-
-    # outer handler API to cancel the workflow run
-    def cancel_run(self) -> None:
-        self._execute_task(self._runtime.send_event(TickCancelRun()))
-
-    @property
-    def is_running(self) -> bool:
-        return self._is_running
-
-    @property
-    def _state(self) -> BrokerState:
-        ticks = self._tick_log
-        state = self._init_state or BrokerState.from_workflow(self._workflow)
-        new_state = rebuild_state_from_ticks(state, ticks)
-        return new_state
-
-    @property
-    def _tick_log(self) -> list[WorkflowTick]:
-        snapshottable = as_snapshottable(self._runtime)
-        if snapshottable is None:
-            raise WorkflowRuntimeError("Plugin is not snapshottable")
-        return snapshottable.replay()
-
-    # mostly a debug API. May be removed in the future.
-    async def running_steps(self) -> list[str]:
-        return [
-            step
-            for step in self._state.workers.keys()
-            if self._state.workers[step].in_progress
-        ]
-
-    # step api only
-    def collect_events(
-        self, ev: Event, expected: list[Type[Event]], buffer_id: str | None = None
-    ) -> list[Event] | None:
-        step_ctx = self._get_step_ctx(fn="collect_events")
-
-        # If no events are expected, return an empty list immediately
-        if not expected:
-            return []
-
-        buffer_id = buffer_id or "default"
-
-        collected_events = step_ctx.state.collected_events.get(buffer_id, [])
-
-        remaining_event_types = Counter(expected) - Counter(
-            [type(e) for e in collected_events]
-        )
-
-        if remaining_event_types != Counter([type(ev)]):
-            if type(ev) in remaining_event_types:
-                step_ctx.returns.return_values.append(
-                    AddCollectedEvent(event_id=buffer_id, event=ev)
-                )
-            return None
-
-        total = []
-        by_type = defaultdict(list)
-        for e in collected_events + [ev]:
-            by_type[type(e)].append(e)
-        # order by expected type
-        for e_type in expected:
-            total.append(by_type[e_type].pop(0))
-        # if we got here, it means the collection is fulfilled. Clear the collected events when the step is complete
-        step_ctx.returns.return_values.append(DeleteCollectedEvent(event_id=buffer_id))
-        return total
-
-    # may be called from both step API and outer handler API
-    def send_event(self, message: Event, step: str | None = None) -> None:
-        if step is not None:
-            if step not in self._workflow._get_steps():
-                raise WorkflowRuntimeError(f"Step {step} does not exist")
-
-            # Validate that the step accepts this event type
-            step_func = self._workflow._get_steps()[step]
-            step_config = step_func._step_config
-            if type(message) not in step_config.accepted_events:
-                raise WorkflowRuntimeError(
-                    f"Step {step} does not accept event of type {type(message)}"
-                )
-
-        self._execute_task(
-            self._runtime.send_event(TickAddEvent(event=message, step_name=step))
-        )
-
-    def _get_step_ctx(self, fn: str) -> StepWorkerContext:
-        try:
-            return StepWorkerStateContextVar.get()
-        except LookupError:
-            raise WorkflowRuntimeError(
-                f"{fn} may only be called from within a step function"
-            )
-
-    # step api only
-    async def wait_for_event(
-        self,
-        event_type: Type[T],
-        waiter_event: Event | None = None,
-        waiter_id: str | None = None,
-        requirements: dict[str, Any] | None = None,
-        timeout: float | None = 2000,
-    ) -> T:
-        step_ctx = self._get_step_ctx(fn="wait_for_event")
-
-        collected_waiters = step_ctx.state.collected_waiters
-        requirements = requirements or {}
-
-        # Generate a unique key for the waiter
-        event_str = self._get_full_path(event_type)
-        requirements_str = str(requirements)
-        waiter_id = waiter_id or f"waiter_{event_str}_{requirements_str}"
-
-        waiter = next((w for w in collected_waiters if w.waiter_id == waiter_id), None)
-        if waiter is None or waiter.resolved_event is None:
-            raise WaitingForEvent(
-                AddWaiter(
-                    waiter_id=waiter_id,
-                    requirements=requirements,
-                    timeout=timeout,
-                    event_type=event_type,
-                    waiter_event=waiter_event,
-                )
-            )
-        else:
-            step_ctx.returns.return_values.append(DeleteWaiter(waiter_id=waiter_id))
-            return cast(T, waiter.resolved_event)
-
-    def _get_full_path(self, ev_type: Type[Event]) -> str:
-        return f"{ev_type.__module__}.{ev_type.__name__}"
-
-    def stream_published_events(self) -> AsyncGenerator[Event, None]:
-        """The internal queue used for streaming events to callers."""
-        return self._runtime.stream_published_events()
-
-    # step API only
-    def write_event_to_stream(self, ev: Event | None) -> None:
-        if ev is not None:
-            self._execute_task(self._runtime.write_to_event_stream(ev))
-
-    async def shutdown(self) -> None:
-        """Cancels the running workflow loop
-
-        Cancels all outstanding workers, waits for them to finish, and marks the
-        broker as not running. Queues and state remain available so callers can
-        inspect or drain leftover events.
-        """
-        await self._runtime.send_event(TickCancelRun())
-        for worker in self._workers:
-            worker.cancel()
-        self._workers.clear()
-        await self._runtime.close()

packages/llama-index-workflows/src/workflows/runtime/control_loop.py

@@ -46,8 +46,9 @@ from workflows.runtime.types.internal_state import (
     InternalStepWorkerState,
 )
 from workflows.runtime.types.plugin import (
-    WorkflowRuntime,
-    as_snapshottable,
+    InternalRunAdapter,
+    as_snapshottable_adapter,
+    get_current_run,
 )
 from workflows.runtime.types.results import (
     AddCollectedEvent,
@@ -59,9 +60,6 @@ from workflows.runtime.types.results import (
     StepWorkerState,
     StepWorkerWaiter,
 )
-from workflows.runtime.types.step_function import (
-    StepWorkerFunction,
-)
 from workflows.runtime.types.ticks import (
     TickAddEvent,
     TickCancelRun,
@@ -70,11 +68,11 @@ from workflows.runtime.types.ticks import (
     TickTimeout,
     WorkflowTick,
 )
-from workflows.runtime.workflow_registry import workflow_registry
 from workflows.workflow import Workflow
 
 if TYPE_CHECKING:
     from workflows.context.context import Context
+    from workflows.runtime.types.step_function import StepWorkerFunction
 
 
 logger = logging.getLogger()
@@ -89,13 +87,13 @@ class _ControlLoopRunner:
     def __init__(
         self,
         workflow: Workflow,
-        plugin: WorkflowRuntime,
+        adapter: InternalRunAdapter,
         context: Context,
         step_workers: dict[str, StepWorkerFunction],
         init_state: BrokerState,
     ):
         self.workflow = workflow
-        self.plugin = plugin
+        self.adapter = adapter
         self.context = context
         self.step_workers = step_workers
         self.state = init_state
@@ -103,7 +101,7 @@ class _ControlLoopRunner:
         self.queue: asyncio.Queue[WorkflowTick] = asyncio.Queue()
         for tick in self.state.rehydrate_with_ticks():
             self.queue.put_nowait(tick)
-        self.snapshot_plugin = as_snapshottable(plugin)
+        self.snapshot_adapter = as_snapshottable_adapter(adapter)
 
     async def wait_for_tick(self) -> WorkflowTick:
         """Wait for the next tick from the internal queue."""
@@ -114,7 +112,7 @@ class _ControlLoopRunner:
         if delay:
 
             async def _delayed_queue() -> None:
-                await self.plugin.sleep(delay)
+                await self.adapter.sleep(delay)
                 self.queue.put_nowait(tick)
 
             task = asyncio.create_task(_delayed_queue())
@@ -146,7 +144,6 @@ class _ControlLoopRunner:
                     state=snapshot,
                     step_name=command.step_name,
                     event=command.event,
-                    context=self.context,
                     workflow=self.workflow,
                 )
                 self.queue_tick(
@@ -166,7 +163,7 @@ class _ControlLoopRunner:
                         event=command.event,
                         result=[
                             StepWorkerFailed(
-                                exception=e, failed_at=await self.plugin.get_now()
+                                exception=e, failed_at=await self.adapter.get_now()
                             )
                         ],
                     )
@@ -200,7 +197,7 @@ class _ControlLoopRunner:
         elif isinstance(command, CommandCompleteRun):
             return command.result
         elif isinstance(command, CommandPublishEvent):
-            await self.plugin.write_to_event_stream(command.event)
+            await self.adapter.write_to_event_stream(command.event)
             return None
         elif isinstance(command, CommandFailWorkflow):
             await self.cleanup_tasks()
@@ -240,7 +237,7 @@ class _ControlLoopRunner:
         # Start external event listener
         async def _pull() -> None:
             while True:
-                tick = await self.plugin.wait_receive()
+                tick = await self.adapter.wait_receive()
                 self.queue_tick(tick)
 
         self.workers.append(asyncio.create_task(_pull()))
@@ -257,7 +254,7 @@ class _ControlLoopRunner:
 
         # Resume any in-progress work
         self.state, commands = rewind_in_progress(
-            self.state, await self.plugin.get_now()
+            self.state, await self.adapter.get_now()
         )
         for command in commands:
             try:
@@ -272,7 +269,7 @@ class _ControlLoopRunner:
                 tick = await self.wait_for_tick()
                 try:
                     self.state, commands = _reduce_tick(
-                        tick, self.state, await self.plugin.get_now()
+                        tick, self.state, await self.adapter.get_now()
                     )
                 except Exception:
                     await self.cleanup_tasks()
@@ -281,8 +278,8 @@ class _ControlLoopRunner:
                         exc_info=True,
                     )
                     raise
-                if self.snapshot_plugin is not None:
-                    self.snapshot_plugin.on_tick(tick)
+                if self.snapshot_adapter is not None:
+                    self.snapshot_adapter.on_tick(tick)
                 for command in commands:
                     try:
                         result = await self.process_command(command)
@@ -304,13 +301,10 @@ async def control_loop(
     """
     The main async control loop for a workflow run.
     """
-    # Prefer run-scoped context if available (set by broker)
-    current = workflow_registry.get_run(run_id)
-    if current is None:
-        raise WorkflowRuntimeError("Run context not found for control loop")
+    current = get_current_run()
     state = init_state or BrokerState.from_workflow(current.workflow)
     runner = _ControlLoopRunner(
-        current.workflow, current.plugin, current.context, current.steps, state
+        current.workflow, current.run_adapter, current.context, current.steps, state
     )
     return await runner.run(start_event=start_event)
 

packages/llama-index-workflows/src/workflows/runtime/types/_identity_weak_ref.py

@@ -1,64 +0,0 @@
-from __future__ import annotations
-
-import weakref
-from typing import Callable, Generic, TypeVar, overload
-
-K = TypeVar("K")
-V = TypeVar("V")
-
-
-class _IdentityWeakRef(weakref.ref, Generic[K]):
-    __slots__ = ("_hash",)
-
-    _hash: int
-
-    def __new__(
-        cls, obj: K, callback: Callable[[_IdentityWeakRef[K]], None] | None = None
-    ) -> _IdentityWeakRef[K]:
-        self = super().__new__(cls, obj, callback)
-        self._hash = id(
-            obj
-        )  # cache identity-based hash; works even if obj is unhashable
-        return self
-
-    def __hash__(self) -> int:
-        return self._hash
-
-    def __eq__(self, other: object) -> bool:
-        if not isinstance(other, _IdentityWeakRef):
-            return NotImplemented
-        return self() is other()
-
-
-class IdentityWeakKeyDict(Generic[K, V]):
-    _d: dict[_IdentityWeakRef[K], V]
-
-    def __init__(self) -> None:
-        self._d = {}
-
-    def _mk(self, obj: K) -> _IdentityWeakRef[K]:
-        def _cb(wr: _IdentityWeakRef[K]) -> None:
-            self._d.pop(wr)
-
-        return _IdentityWeakRef(obj, _cb)
-
-    def __setitem__(self, obj: K, value: V) -> None:
-        self._d[self._mk(obj)] = value
-
-    def __getitem__(self, obj: K) -> V:
-        return self._d[_IdentityWeakRef(obj)]
-
-    @overload
-    def get(self, obj: K) -> V | None: ...
-
-    @overload
-    def get(self, obj: K, default: V) -> V: ...
-
-    def get(self, obj: K, default: V | None = None) -> V | None:
-        return self._d.get(_IdentityWeakRef(obj), default)
-
-    def __contains__(self, obj: K) -> bool:
-        return _IdentityWeakRef(obj) in self._d
-
-    def __delitem__(self, obj: K) -> None:
-        del self._d[_IdentityWeakRef(obj)]

packages/llama-index-workflows/src/workflows/runtime/types/plugin.py

@@ -1,123 +1,507 @@
 # SPDX-License-Identifier: MIT
 # Copyright (c) 2026 LlamaIndex Inc.
 """
-A plugin interface to switch out a broker runtime (external library or service that manages durable/distributed step execution).
+A runtime interface to switch out a broker runtime (external library or service that manages durable/distributed step execution).
 """
 
 from __future__ import annotations
 
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from contextvars import ContextVar, Token
 from dataclasses import dataclass
 from typing import (
     TYPE_CHECKING,
+    Any,
     AsyncGenerator,
     Coroutine,
+    Generator,
     Protocol,
-    cast,
 )
 
-from workflows.events import Event, StopEvent
-from workflows.runtime.types.internal_state import BrokerState
-from workflows.runtime.types.step_function import StepWorkerFunction
-from workflows.runtime.types.ticks import WorkflowTick
-
 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
 
+from workflows.events import Event, StartEvent, StopEvent
+from workflows.runtime.types.ticks import TickCancelRun, WorkflowTick
+
+# Context variable for implicit runtime scoping
+_current_runtime: ContextVar[Runtime | None] = ContextVar(
+    "current_runtime", default=None
+)
+
 
 @dataclass
 class RegisteredWorkflow:
-    workflow_function: ControlLoopFunction
-    steps: dict[str, StepWorkerFunction]
+    workflow: Workflow
+    workflow_run_fn: WorkflowRunFunction
+    steps: dict[str, StepWorkerFunction[Any]]
 
 
-class Plugin(Protocol):
-    def register(
-        self,
-        workflow: Workflow,
-        workflow_function: ControlLoopFunction,
-        steps: dict[str, StepWorkerFunction],
-    ) -> None | RegisteredWorkflow:
-        """
-        Called on a workflow before the first time each workflow instance is run, in order to register it within the plugin's runtime.
-
-        Provides an opportunity to modify the workflow function and steps, e.g. to wrap the workflow_function, or StepWorkerFunction's within the steps a decorator.
-        """
-        ...
-
-    def new_runtime(self, run_id: str) -> WorkflowRuntime:
-        """
-        Called on each workflow run, in order to create a new runtime instance for driving the workflow via the plugin's runtime.
-        """
-        ...
-
-
-class WorkflowRuntime(Protocol):
+class InternalRunAdapter(ABC):
     """
-    A LlamaIndex workflow's internal state is managed via an event-sourced reducer that triggers step executions. It communicates
-    with the outside world via messages. Messages may be sent to it to update its event log, and it in turn publishes messages that are made
-    available via the event stream.
+    Adapter interface for use INSIDE a workflow's control loop.
+
+    This adapter is used by the workflow execution engine (broker) to receive
+    ticks from external sources, publish events to listeners, manage timing,
+    and perform durable sleeps.
+
+    The InternalRunAdapter is created by Runtime.new_internal_adapter() for each
+    workflow run and is passed to the control loop function. It provides the
+    internal-facing side of workflow communication:
+    - Receiving ticks from the external mailbox (wait_receive)
+    - Publishing events that external code can stream (write_to_event_stream)
+    - Getting current time with durability support (get_now)
+    - Sleeping with durability support (sleep)
+
+    The run_id is always available and required at construction time.
     """
 
-    async def send_event(self, tick: WorkflowTick) -> None:
-        """Called from outside of the workflow to modify the workflow execution. WorkflowTick events are appended to a mailbox and processed sequentially"""
+    @property
+    @abstractmethod
+    def run_id(self) -> str:
+        """
+        The unique identifier for this workflow run.
+
+        Always available - required at adapter construction time.
+        """
         ...
 
+    @abstractmethod
     async def wait_receive(self) -> WorkflowTick:
-        """called from inside of the workflow control loop function to add a tick event from `send_event` to the mailbox. Function waits until a tick event is received."""
+        """
+        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:
-        """Called from inside of a workflow function to write / emit events to listeners outside of the workflow"""
-        ...
-
-    def stream_published_events(self) -> AsyncGenerator[Event, None]:
-        """Called from outside of a workflow, reads event stream published by the workflow"""
+        """
+        Publish an event to external listeners.
+
+        Called from inside the workflow to emit events that can be observed
+        by external code via the ExternalRunAdapter's stream_published_events().
+        """
         ...
 
+    @abstractmethod
     async def get_now(self) -> float:
-        """Called from within the workflow control loop function to get the current time in seconds since epoch. If workflow is durable via replay, it should return a cached value from the first call. (e.g. this should be implemented similar to a regular durable step)"""
+        """
+        Get the current time in seconds since epoch.
+
+        Called from within the workflow control loop. For durable workflows,
+        this should return a memoized/replayed value to ensure deterministic
+        replay behavior.
+        """
         ...
 
+    @abstractmethod
     async def sleep(self, seconds: float) -> None:
-        """Called from within the workflow control loop function to sleep for a given number of seconds. This should integrate with the host plugin for cases where an inactive workflow may be paused, and awoken later via memoized replay. Note that other tasks in the control loop may still be running simultaneously."""
+        """
+        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:
+        """
+        Send a tick into the workflow's own mailbox from within the control loop.
+
+        Called from inside the workflow (e.g., from step functions via ctx.send_event)
+        to inject events back into the workflow's execution. The tick will be
+        received by wait_receive() on the next iteration.
+        """
+        ...
+
+    def get_state_store(self) -> InMemoryStateStore[Any] | None:
+        """
+        Get the state store for this workflow run.
+
+        Returns the state store from the runtime, or None if not initialized.
+        Default implementation returns None.
+        """
+        return None
+
+
+class ExternalRunAdapter(ABC):
+    """
+    Adapter interface for use OUTSIDE a workflow's control loop.
+
+    This adapter is used by external code (e.g., HTTP handlers, client code)
+    to interact with a running workflow - sending events into the workflow
+    and streaming events published by the workflow.
+
+    The ExternalRunAdapter is created by Runtime.new_external_adapter() and
+    provides the external-facing side of workflow communication:
+    - Sending ticks into the workflow mailbox (send_event)
+    - Streaming events published by the workflow (stream_published_events)
+    - Cleaning up resources when done (close)
+
+    The run_id is always available and matches the internal adapter's run_id.
+    """
+
+    @property
+    @abstractmethod
+    def run_id(self) -> str:
+        """
+        The unique identifier for this workflow run.
+
+        Always available - matches the InternalRunAdapter's run_id.
+        """
+        ...
+
+    @abstractmethod
+    async def send_event(self, tick: WorkflowTick) -> None:
+        """
+        Send a tick into the workflow mailbox.
+
+        Called from outside the workflow to inject events into the workflow's
+        execution. The tick will be received by the internal adapter's
+        wait_receive() method.
+        """
+        ...
+
+    @abstractmethod
+    def stream_published_events(self) -> AsyncGenerator[Event, None]:
+        """
+        Stream events published by the workflow.
+
+        Called from outside the workflow to observe events emitted by the
+        workflow via the internal adapter's write_to_event_stream().
+        Returns an async generator that yields events as they are published.
+        """
+        ...
+
+    @abstractmethod
     async def close(self) -> None:
-        """API that the broker calls to close the plugin after a workflow run is fully complete"""
+        """
+        Clean up adapter resources.
+
+        Called when done interacting with the workflow to release any
+        resources held by this adapter (e.g., close streams, release locks).
+        """
+
         ...
 
+    @abstractmethod
+    async def get_result(self) -> StopEvent:
+        """
+        Get the result of the workflow run, if completed. Will raise if the workflow failed or was cancelled
+        """
+        ...
 
-class SnapshottableRuntime(WorkflowRuntime, Protocol):
+    async def cancel(self) -> None:
+        """
+        Cancel the workflow run if it is still running.
+        """
+        await self.send_event(TickCancelRun())
+
+    def get_state_store(self) -> InMemoryStateStore[Any] | None:
+        """
+        Get the state store for this workflow run.
+
+        Returns the state store if this adapter owns it, or None if state
+        is managed externally. Default implementation returns None.
+        """
+        return None
+
+
+@dataclass
+class RunContext:
+    """Context for an active workflow run, available via get_current_run()."""
+
+    workflow: Workflow
+    run_adapter: InternalRunAdapter
+    context: Context
+    steps: dict[str, StepWorkerFunction[Any]]
+
+
+_current_run: ContextVar[RunContext | None] = ContextVar("current_run", default=None)
+
+
+@contextmanager
+def run_context(ctx: RunContext) -> Generator[RunContext, None, None]:
+    """Set the current run context for the duration of a workflow run."""
+    token = _current_run.set(ctx)
+    try:
+        yield ctx
+    finally:
+        _current_run.reset(token)
+
+
+def get_current_run() -> RunContext:
+    """Get the current run context. Raises if not in a workflow run."""
+    ctx = _current_run.get()
+    if ctx is None:
+        raise RuntimeError("Not in a workflow run context")
+    return ctx
+
+
+class Runtime(ABC):
     """
-    Snapshot API. Optional mix in to a WorkflowRuntime. When implemented, plugin should offer a replay function to return the recorded
-    ticks so that callers can debug or replay the workflow. `on_tick` is called whenever a tick event is received externally OR as a result
-    from an internal command (e.g. a step function completing, a timeout occurring, etc.)
+    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.
+
+    Lifecycle:
+    1. Create runtime instance
+    2. Create workflow instances (auto-register with runtime via registering())
+    3. Call launch() to start workers/register with backend
+    4. Run workflows
+    5. Call destroy() to clean up
+
+    Use registering() context manager for implicit workflow registration.
     """
 
+    _token: Token[Runtime | None]
+
+    def get_or_register(self, workflow: "Workflow") -> RegisteredWorkflow:
+        """Get the registered workflow if available, otherwise register it."""
+        registered = self.get_registered(workflow)
+        if registered is None:
+            registered = self.register(workflow)
+        return registered
+
+    @abstractmethod
+    def register(self, workflow: "Workflow") -> RegisteredWorkflow:
+        """
+        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.
+
+        Returns RegisteredWorkflow with wrapped functions
+        """
+        ...
+
+    @abstractmethod
+    def run_workflow(
+        self,
+        run_id: str,
+        workflow: Workflow,
+        init_state: BrokerState,
+        start_event: StartEvent | None = None,
+        serialized_state: dict[str, Any] | None = None,
+        serializer: "BaseSerializer | None" = None,
+    ) -> ExternalRunAdapter:
+        """
+        Launch a workflow run.
+
+        The runtime creates and owns the state store based on serialized_state.
+        Returns the external adapter for the workflow run.
+
+        Args:
+            run_id: Unique identifier for this workflow run.
+            registered: The registered workflow to run.
+            init_state: Initial broker state (queues, workers, etc).
+            start_event: Optional start event to begin the workflow.
+            serialized_state: Serialized state store data to restore from.
+            serializer: Serializer to use for deserializing state.
+        """
+        ...
+
+    @abstractmethod
+    def get_internal_adapter(self) -> 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.
+        """
+        ...
+
+    @abstractmethod
+    def get_external_adapter(self, run_id: str) -> ExternalRunAdapter:
+        """
+        Get the external adapter for a workflow run.
+
+        Called after launching a workflow run, or when getting a handle for an existing workflow run.
+        Used to send events into the workflow and stream published events.
+
+        The run_id must match the internal adapter's run_id for the same run.
+        The external adapter is used by client code interacting with the workflow.
+        """
+        ...
+
+    def launch(self) -> None:
+        """
+        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.
+        """
+        pass
+
+    def destroy(self) -> None:
+        """
+        Clean up runtime resources.
+
+        Called when done with the runtime. Stops workers, closes connections.
+        """
+        pass
+
+    def track_workflow(self, workflow: "Workflow") -> None:
+        """
+        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.
+        Default implementation is a no-op.
+        """
+        pass
+
+    def get_registered(self, workflow: "Workflow") -> RegisteredWorkflow | None:
+        """
+        Get the registered workflow if available.
+
+        Returns the pre-registered workflow from launch(), or None if not tracked.
+        """
+        return None
+
+    @contextmanager
+    def registering(self) -> Generator[Runtime, None, None]:
+        """
+        Context manager for implicit workflow registration.
+
+        Workflows created inside this block will automatically register
+        with this runtime. Does NOT call launch() on exit.
+        """
+        token = _current_runtime.set(self)
+        try:
+            yield self
+        finally:
+            _current_runtime.reset(token)
+
+
+class SnapshottableAdapter(ABC):
+    """
+    Mixin interface that adds snapshot/replay capabilities to adapters.
+
+    This is a standalone mixin (not inheriting from InternalRunAdapter or
+    ExternalRunAdapter) that can be combined with adapter implementations
+    to add tick recording for debugging or replay purposes.
+
+    Adapters that implement this interface can record ticks as they occur
+    and replay them later. `on_tick` is called whenever a tick event is
+    received externally OR as a result from an internal command (e.g., a
+    step function completing, a timeout occurring, etc.)
+
+    Use `as_snapshottable_adapter()` to check if an adapter supports snapshotting.
+    """
+
+    @property
+    @abstractmethod
+    def init_state(self) -> "BrokerState":
+        """
+        Get the initial state of the adapter.
+        """
+        ...
+
+    @abstractmethod
     def on_tick(self, tick: WorkflowTick) -> None:
-        """Called whenever a tick event is received"""
+        """
+        Called whenever a tick event is received.
+
+        This method is invoked for both external ticks (sent via send_event)
+        and internal ticks (generated by step completions, timeouts, etc.).
+        Implementations should record the tick for later replay.
+        """
         ...
 
+    @abstractmethod
     def replay(self) -> list[WorkflowTick]:
-        """return the recorded ticks for replay"""
+        """
+        Return the recorded ticks for replay.
+
+        Returns all ticks that were recorded via on_tick(), in the order
+        they were received. Used for debugging and workflow replay.
+        """
         ...
 
 
-def as_snapshottable(runtime: WorkflowRuntime) -> SnapshottableRuntime | None:
-    """Check if a runtime is snapshottable."""
-    if (
-        getattr(runtime, "on_tick", None) is not None
-        and getattr(runtime, "replay", None) is not None
-    ):
-        return cast(SnapshottableRuntime, runtime)
+def as_snapshottable_adapter(
+    adapter: ExternalRunAdapter | InternalRunAdapter,
+) -> SnapshottableAdapter | None:
+    """
+    Check if an internal adapter supports snapshotting.
+
+    Returns the adapter cast to SnapshottableAdapter if it implements
+    the snapshotting interface, or None otherwise.
+    """
+    if isinstance(adapter, SnapshottableAdapter):
+        return adapter
+    return None
+
+
+class V2RuntimeCompatibilityShim(ABC):
+    """
+    This interface will be deleted in V3. Temporary shim to support deprecated v2 functionality
+    """
+
+    @abstractmethod
+    def get_result_or_none(self) -> StopEvent | None:
+        """
+        Get the result of the workflow run, if completed. Will raise if the workflow failed or was cancelled, otherwise return None if still running
+        """
+        ...
+
+    @property
+    @abstractmethod
+    def is_running(self) -> bool:
+        """
+        Check if the workflow run is still running.
+        """
+        ...
+
+    @abstractmethod
+    def abort(self) -> None:
+        """
+        Forcefully abort the workflow execution (ungraceful hard cancel).
+
+        This immediately terminates execution by cancelling the underlying task.
+        Unlike cancel() which sends a graceful cancellation signal:
+        - In-flight step work is cancelled immediately
+        - No WorkflowCancelledEvent is emitted
+        - The workflow does not finalize gracefully
+
+        This is deprecated v2 behavior - prefer cancel_run() for graceful cancellation.
+        """
+        ...
+
+
+def as_v2_runtime_compatibility_shim(
+    adapter: ExternalRunAdapter,
+) -> V2RuntimeCompatibilityShim | None:
+    """
+    Check if an adapter supports the V2 runtime compatibility shim.
+    """
+    if isinstance(adapter, V2RuntimeCompatibilityShim):
+        return adapter
     return None
 
 
 class ControlLoopFunction(Protocol):
     """
     Protocol for a function that starts and runs the internal control loop for a workflow run.
-    Plugin decorators to the control loop function must maintain this signature.
+    Runtime decorators to the control loop function must maintain this signature.
     """
 
     def __call__(
@@ -126,3 +510,16 @@ class ControlLoopFunction(Protocol):
         init_state: BrokerState | None,
         run_id: str,
     ) -> Coroutine[None, None, StopEvent]: ...
+
+
+class WorkflowRunFunction(Protocol):
+    """
+    Protocol for a function that runs a workflow. Wraps a control loop function with glue to the runtime.
+    """
+
+    def __call__(
+        self,
+        init_state: BrokerState,
+        start_event: StartEvent | None = None,
+        tags: dict[str, Any] = {},
+    ) -> Coroutine[None, None, StopEvent]: ...

packages/llama-index-workflows/src/workflows/runtime/types/step_function.py

@@ -9,10 +9,21 @@ import time
 from contextvars import copy_context
 from typing import TYPE_CHECKING, Any, Awaitable, Callable, Generic, Protocol
 
+from llama_index_instrumentation.dispatcher import instrument_tags
 from workflows.decorators import P, R, StepConfig
 from workflows.errors import WorkflowRuntimeError
 from workflows.events import (
     Event,
+    StartEvent,
+    StopEvent,
+)
+from workflows.runtime.control_loop import control_loop
+from workflows.runtime.types.internal_state import BrokerState
+from workflows.runtime.types.plugin import (
+    ControlLoopFunction,
+    RunContext,
+    WorkflowRunFunction,
+    run_context,
 )
 from workflows.runtime.types.results import (
     Returns,
@@ -36,7 +47,6 @@ class StepWorkerFunction(Protocol, Generic[R]):
         state: StepWorkerState,
         step_name: str,
         event: Event,
-        context: Context,  # TODO - pass an identifier and re-hydrate from the plugin for distributed step workers
         workflow: Workflow,
     ) -> Awaitable[list[StepFunctionResult[R]]]: ...
 
@@ -51,6 +61,7 @@ async def partial(
     kwargs: dict[str, Any] = {}
     kwargs[step_config.event_name] = event
     if step_config.context_parameter:
+        # Convert to internal face for step execution
         kwargs[step_config.context_parameter] = context
     with workflow._resource_manager.resolution_scope():
         for resource_def in step_config.resources:
@@ -62,6 +73,15 @@ async def partial(
     return functools.partial(func, **kwargs)
 
 
+def as_step_worker_functions(workflow: Workflow) -> dict[str, StepWorkerFunction]:
+    step_funcs = workflow._get_steps()
+    step_workers: dict[str, StepWorkerFunction[Any]] = {
+        name: as_step_worker_function(getattr(func, "__func__", func))
+        for name, func in step_funcs.items()
+    }
+    return step_workers
+
+
 def as_step_worker_function(func: Callable[P, Awaitable[R]]) -> StepWorkerFunction[R]:
     """
     Wrap a step function, setting context variables and handling exceptions to instead
@@ -79,9 +99,11 @@ def as_step_worker_function(func: Callable[P, Awaitable[R]]) -> StepWorkerFuncti
         state: StepWorkerState,
         step_name: str,
         event: Event,
-        context: Context,
         workflow: Workflow,
     ) -> list[StepFunctionResult[R]]:
+        from workflows.context.context import Context
+
+        internal_context = Context._create_internal(workflow=workflow)
         returns = Returns[R](return_values=[])
 
         token = StepWorkerStateContextVar.set(
@@ -102,7 +124,7 @@ def as_step_worker_function(func: Callable[P, Awaitable[R]]) -> StepWorkerFuncti
                 func=workflow._dispatcher.span(call_func),
                 step_config=config,
                 event=event,
-                context=context,
+                context=internal_context,
                 workflow=workflow,
             )
 
@@ -149,3 +171,45 @@ def as_step_worker_function(func: Callable[P, Awaitable[R]]) -> StepWorkerFuncti
         pass
 
     return wrapper
+
+
+def create_workflow_run_function(
+    workflow: Workflow, control_loop_fn: ControlLoopFunction = control_loop
+) -> WorkflowRunFunction:
+    async def run_workflow(
+        init_state: BrokerState,
+        start_event: StartEvent | None = None,
+        tags: dict[str, Any] = {},
+    ) -> StopEvent:
+        from workflows.context.context import Context
+        from workflows.context.internal_context import InternalContext
+
+        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()
+        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
+            await asyncio.sleep(0)
+
+            run_ctx = RunContext(
+                workflow=workflow,
+                run_adapter=internal_adapter,
+                context=internal_ctx,
+                steps=registered.steps,
+            )
+            try:
+                with run_context(run_ctx):
+                    result = await control_loop_fn(
+                        start_event,
+                        init_state,
+                        internal_adapter.run_id,
+                    )
+                    return result
+            finally:
+                # Cancel any background tasks from InternalContext on completion or cancellation
+                if isinstance(internal_ctx._face, InternalContext):
+                    internal_ctx._face.cancel_background_tasks()
+
+    return run_workflow

packages/llama-index-workflows/src/workflows/runtime/workflow_registry.py

@@ -1,93 +0,0 @@
-from dataclasses import dataclass
-from threading import Lock
-from typing import TYPE_CHECKING, Optional
-
-from workflows.runtime.types._identity_weak_ref import IdentityWeakKeyDict
-from workflows.runtime.types.plugin import (
-    ControlLoopFunction,
-    Plugin,
-    RegisteredWorkflow,
-    WorkflowRuntime,
-)
-from workflows.runtime.types.step_function import StepWorkerFunction
-from workflows.workflow import Workflow
-
-if TYPE_CHECKING:
-    from workflows.context.context import Context
-
-
-class WorkflowPluginRegistry:
-    """
-    Ensures that plugins register each workflow once and only once for each plugin.
-    """
-
-    def __init__(self) -> None:
-        # Map each workflow instance to its plugin registrations.
-        # Weakly references workflow keys so entries are GC'd when workflows are.
-        self.workflows: IdentityWeakKeyDict[
-            Workflow, dict[type[Plugin], RegisteredWorkflow]
-        ] = IdentityWeakKeyDict()
-        self.lock = Lock()
-        self.run_contexts: dict[str, RegisteredRunContext] = {}
-
-    def get_registered_workflow(
-        self,
-        workflow: Workflow,
-        plugin: Plugin,
-        workflow_function: ControlLoopFunction,
-        steps: dict[str, StepWorkerFunction],
-    ) -> RegisteredWorkflow:
-        plugin_type = type(plugin)
-
-        # Fast path without lock
-        plugin_map = self.workflows.get(workflow)
-        if plugin_map is not None and plugin_type in plugin_map:
-            return plugin_map[plugin_type]
-        with self.lock:
-            # Double-check after acquiring lock
-            plugin_map = self.workflows.get(workflow)
-            if plugin_map is not None and plugin_type in plugin_map:
-                return plugin_map[plugin_type]
-
-            registered_workflow = plugin.register(workflow, workflow_function, steps)
-            if registered_workflow is None:
-                registered_workflow = RegisteredWorkflow(workflow_function, steps)
-            if plugin_map is None:
-                plugin_map = {}
-                self.workflows[workflow] = plugin_map
-            plugin_map[plugin_type] = registered_workflow
-            return registered_workflow
-
-    def register_run(
-        self,
-        run_id: str,
-        workflow: Workflow,
-        plugin: WorkflowRuntime,
-        context: "Context",
-        steps: dict[str, StepWorkerFunction],
-    ) -> None:
-        self.run_contexts[run_id] = RegisteredRunContext(
-            run_id=run_id,
-            workflow=workflow,
-            plugin=plugin,
-            context=context,
-            steps=steps,
-        )
-
-    def get_run(self, run_id: str) -> Optional["RegisteredRunContext"]:
-        return self.run_contexts.get(run_id)
-
-    def delete_run(self, run_id: str) -> None:
-        self.run_contexts.pop(run_id, None)
-
-
-workflow_registry = WorkflowPluginRegistry()
-
-
-@dataclass
-class RegisteredRunContext:
-    run_id: str
-    workflow: Workflow
-    plugin: WorkflowRuntime
-    context: "Context"
-    steps: dict[str, StepWorkerFunction]

packages/llama-index-workflows/src/workflows/runtime/workflow_tracker.py

@@ -0,0 +1,75 @@
+# 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

@@ -17,6 +17,7 @@ from pydantic import ValidationError
 
 if TYPE_CHECKING:  # pragma: no cover
     from .context import Context
+    from .runtime.types.plugin import Runtime
 from .decorators import StepConfig, StepFunction
 from .errors import (
     WorkflowConfigurationError,
@@ -100,6 +101,9 @@ class Workflow(metaclass=WorkflowMeta):
     # Populated by the metaclass; declared here for type checkers.
     _step_functions: dict[str, StepFunction]
 
+    _runtime: Runtime
+    _workflow_name: str | None
+
     def __init__(
         self,
         timeout: float | None = 45.0,
@@ -107,6 +111,8 @@ class Workflow(metaclass=WorkflowMeta):
         verbose: bool = False,
         resource_manager: ResourceManager | None = None,
         num_concurrent_runs: int | None = None,
+        runtime: Runtime | None = None,
+        workflow_name: str | None = None,
     ) -> None:
         """
         Initialize a workflow instance.
@@ -120,24 +126,78 @@ class Workflow(metaclass=WorkflowMeta):
             resource_manager (ResourceManager | None): Custom resource manager
                 for dependency injection.
             num_concurrent_runs (int | None): Limit on concurrent `run()` calls.
+            runtime (Runtime | None): Optional runtime to use for this workflow.
+                If not provided, uses the current context-scoped runtime or
+                falls back to basic_runtime.
+            workflow_name (str | None): Optional explicit name for this workflow.
+                If not provided, a module-qualified name is computed from
+                the class's `__module__` and `__qualname__` attributes.
         """
         # Configuration
         self._timeout = timeout
         self._verbose = verbose
         self._disable_validation = disable_validation
         self._num_concurrent_runs = num_concurrent_runs
+        # Store explicit name (None means use computed name)
+        self._workflow_name = workflow_name
         # Detect StartEvent issues before StopEvent for clearer guidance
         self._start_event_class = self._ensure_start_event_class()
         self._stop_event_class = self._ensure_stop_event_class()
         self._events = self._ensure_events_collected()
-        self._sem = (
-            asyncio.Semaphore(num_concurrent_runs) if num_concurrent_runs else None
-        )
         # Resource management
         self._resource_manager = resource_manager or ResourceManager()
         # Instrumentation
         self._dispatcher = dispatcher
 
+        # Runtime registration: explicit > context-scoped > basic_runtime
+        from workflows.plugins._context import get_current_runtime
+
+        if runtime is not None:
+            self._runtime = runtime
+        else:
+            # get_current_runtime() falls back to basic_runtime
+            self._runtime = get_current_runtime()
+
+        # Register with runtime for tracking (no-op for BasicRuntime)
+        self._runtime.track_workflow(self)
+
+    def _validate_valid_step_message(self, step: str, message: Event) -> None:
+        """Validate that a step name exists in the workflow."""
+        if step not in self._get_steps():
+            raise WorkflowRuntimeError(f"Step {step} does not exist")
+
+        step_func = self._get_steps()[step]
+        step_config = step_func._step_config
+        if type(message) not in step_config.accepted_events:
+            raise WorkflowRuntimeError(
+                f"Step {step} does not accept event of type {type(message)}"
+            )
+
+    @property
+    def runtime(self) -> Runtime:
+        """The runtime this workflow is registered with."""
+        return self._runtime
+
+    @property
+    def workflow_name(self) -> str:
+        """
+        The workflow name.
+
+        If an explicit name was provided at construction, returns that.
+        Otherwise, returns a module-qualified name based on the class's
+        __module__ and __qualname__ attributes.
+
+        Examples:
+            - Explicit: `Workflow(workflow_name="my-workflow")` → `"my-workflow"`
+            - Module-level class: `"mymodule.MyWorkflow"`
+            - Nested class: `"mymodule.Outer.Inner"`
+            - Function-scoped: `"mymodule.func.<locals>.LocalWorkflow"`
+        """
+        if self._workflow_name is not None:
+            return self._workflow_name
+        cls = self.__class__
+        return f"{cls.__module__}.{cls.__qualname__}"
+
     def _ensure_start_event_class(self) -> type[StartEvent]:
         """
         Returns the StartEvent type used in this workflow.
@@ -355,14 +415,14 @@ class Workflow(metaclass=WorkflowMeta):
 
         # 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.
+        # Is not an easily synchronously queryable property.
         start_event_instance: StartEvent | None = (
             None
             if ctx.is_running
             else self._get_start_event_instance(start_event, **kwargs)
         )
-        return ctx._workflow_run(
-            workflow=self, start_event=start_event_instance, semaphore=self._sem
-        )
+        return ctx._workflow_run(workflow=self, start_event=start_event_instance)
 
     def _validate_resource_configs(self) -> list[str]:
         """Validate all resource configs (including nested ones) by loading them."""

packages/llama-index-workflows/tests/conftest.py

@@ -1,13 +1,16 @@
 # SPDX-License-Identifier: MIT
 # Copyright (c) 2026 LlamaIndex Inc.
 
-from typing import AsyncGenerator
+from typing import Any, AsyncGenerator
 
 import pytest
 from pydantic import Field
 from workflows.context import Context
+from workflows.context.state_store import DictState, InMemoryStateStore
 from workflows.decorators import step
 from workflows.events import Event, StartEvent, StopEvent
+from workflows.plugins.basic import AsyncioAdapterQueues, ExternalAsyncioAdapter
+from workflows.runtime.types.internal_state import BrokerState
 from workflows.workflow import Workflow
 
 
@@ -48,10 +51,20 @@ def events() -> list:
 
 
 @pytest.fixture()
-async def ctx(workflow: Workflow) -> AsyncGenerator[Context, None]:
-    ctx = Context(workflow=workflow)
-    broker = ctx._init_broker(workflow)
+async def ctx(workflow: Workflow) -> AsyncGenerator[Context[Any], None]:
+    from workflows.context.external_context import ExternalContext
+
+    queues = AsyncioAdapterQueues(
+        run_id="test-run",
+        init_state=BrokerState.from_workflow(workflow),
+        state_store=InMemoryStateStore(DictState()),
+    )
+    ctx = Context._create_external(
+        workflow=workflow,
+        external_adapter=ExternalAsyncioAdapter(queues=queues),
+    )
+    assert isinstance(ctx._face, ExternalContext)
     try:
         yield ctx
     finally:
-        await broker.shutdown()
+        await ctx._face.shutdown()

packages/llama-index-workflows/tests/context/test_context.py

@@ -4,14 +4,22 @@
 from __future__ import annotations
 
 import asyncio
+import json
 from typing import Optional, Union
 
 import pytest
 from pydantic import BaseModel
 from workflows.context import Context
-from workflows.context.state_store import DictState
+from workflows.context.context import (
+    _warn_cancel_before_start,
+    _warn_cancel_in_step,
+    _warn_get_result,
+    _warn_is_running_in_step,
+)
+from workflows.context.external_context import ExternalContext
+from workflows.context.state_store import DictState, InMemoryStateStore
 from workflows.decorators import step
-from workflows.errors import WorkflowRuntimeError
+from workflows.errors import ContextStateError, WorkflowRuntimeError
 from workflows.events import (
     Event,
     HumanResponseEvent,
@@ -19,6 +27,8 @@ from workflows.events import (
     StartEvent,
     StopEvent,
 )
+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
 from workflows.workflow import Workflow
@@ -30,6 +40,25 @@ from ..conftest import (  # type: ignore[import]
 )
 
 
+@pytest.fixture()
+def internal_ctx(workflow: Workflow) -> Context:
+    """Create a context directly in internal face for testing store operations."""
+    # Set up a runtime with state store for this workflow
+    runtime = BasicRuntime()
+    run_id = "test-run"
+    init_state = BrokerState.from_workflow(workflow)
+    # Create queues with state store so get_internal_adapter() returns adapter with store
+    queues = AsyncioAdapterQueues(
+        run_id=run_id,
+        init_state=init_state,
+        state_store=InMemoryStateStore(DictState()),
+    )
+    runtime._queues[run_id] = queues
+    workflow._runtime = runtime
+    with setting_run_id(run_id):
+        return Context._create_internal(workflow=workflow)
+
+
 @pytest.mark.asyncio
 async def test_collect_events() -> None:
     ev1 = OneTestEvent()
@@ -131,39 +160,49 @@ async def test_collect_events_with_extra_event_type() -> None:
     # Verify the collector was called multiple times (once for each event)
     ctx = r.ctx
     assert ctx is not None
-    calls = await ctx.store.get("calls")
+    ctx_dict = ctx.to_dict()
+    # State is serialized as JSON strings under state_data._data
+    calls = json.loads(ctx_dict["state"]["state_data"]["_data"]["calls"])
     # Should be called at least 3 times: once for LastEvent (returns None),
     # once for OneTestEvent (returns None), once for AnotherTestEvent (returns result)
     assert calls >= 3
 
 
 @pytest.mark.asyncio
-async def test_get_default(workflow: Workflow) -> None:
-    c1: Context[DictState] = Context(workflow)
-    assert await c1.store.get("test_key", default=42) == 42
+async def test_get_default(internal_ctx: Context) -> None:
+    assert await internal_ctx.store.get("test_key", default=42) == 42
 
 
 @pytest.mark.asyncio
-async def test_get(ctx: Context) -> None:
-    await ctx.store.set("foo", 42)
-    assert await ctx.store.get("foo") == 42
+async def test_get(internal_ctx: Context) -> None:
+    await internal_ctx.store.set("foo", 42)
+    assert await internal_ctx.store.get("foo") == 42
 
 
 @pytest.mark.asyncio
-async def test_get_not_found(ctx: Context) -> None:
+async def test_get_not_found(internal_ctx: Context) -> None:
     with pytest.raises(ValueError):
-        await ctx.store.get("foo")
+        await internal_ctx.store.get("foo")
 
 
 @pytest.mark.asyncio
-async def test_send_event_step_is_none(workflow: Workflow, ctx: Context) -> None:
+async def test_send_event_step_is_none(workflow: Workflow) -> None:
     ev = Event(foo="bar")
-    ctx._workflow_run(workflow, start_event=StartEvent())
-    ctx.send_event(ev)
-    await asyncio.sleep(0.01)
-    assert ctx._broker_run is not None
-    replay = ctx._broker_run._tick_log
-    assert TickAddEvent(event=ev, step_name=None) in replay
+    # Create a fresh context and run workflow
+    ctx = Context(workflow)
+    handler = ctx._workflow_run(workflow, start_event=StartEvent())
+    try:
+        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
+    finally:
+        external_face = handler.ctx._face
+        assert isinstance(external_face, ExternalContext)
+        await external_face.shutdown()
 
 
 @pytest.mark.asyncio
@@ -190,9 +229,9 @@ async def test_empty_inprogress_when_workflow_done(workflow: Workflow) -> None:
 
     # there shouldn't be any in progress events
     assert ctx is not None
-    assert ctx._broker_run is not None
+    assert isinstance(ctx._face, ExternalContext)
     # After workflow completion, in_progress should be empty for all steps
-    state = ctx._broker_run._state
+    state = ctx._face._state
     for step_name, worker_state in state.workers.items():
         assert len(worker_state.in_progress) == 0, (
             f"Step {step_name} has {len(worker_state.in_progress)} in-progress events"
@@ -264,9 +303,9 @@ async def test_wait_for_event_in_workflow_serialization() -> None:
     # verify creating a new context has the correct state
     new_ctx = Context.from_dict(workflow, ctx_dict)
     new_handler = workflow.run(ctx=new_ctx)
-    assert new_ctx._broker_run
+    assert isinstance(new_handler.ctx._face, ExternalContext)
     # Check that the waiters are properly restored
-    state = new_ctx._broker_run._state
+    state = new_handler.ctx._face._state
     total_waiters = sum(
         len(worker.collected_waiters) for worker in state.workers.values()
     )
@@ -277,9 +316,9 @@ async def test_wait_for_event_in_workflow_serialization() -> None:
     new_handler.ctx.send_event(Event(msg="bar"))
     result = await new_handler
     assert result == "bar"
-    assert new_handler.ctx._broker_run
+    assert isinstance(new_handler.ctx._face, ExternalContext)
     # After workflow completion, there should be no more waiters
-    state = new_handler.ctx._broker_run._state
+    state = new_handler.ctx._face._state
     total_waiters = sum(
         len(worker.collected_waiters) for worker in state.workers.values()
     )
@@ -379,8 +418,266 @@ async def test_wait_for_multiple_events_in_workflow() -> None:
 
 
 @pytest.mark.asyncio
-async def test_clear(ctx: Context) -> None:
-    await ctx.store.set("test_key", 42)
-    await ctx.store.clear()
-    res = await ctx.store.get("test_key", default=None)
+async def test_clear(internal_ctx: Context) -> None:
+    await internal_ctx.store.set("test_key", 42)
+    await internal_ctx.store.clear()
+    res = await internal_ctx.store.get("test_key", default=None)
     assert res is None
+
+
+@pytest.mark.asyncio
+async def test_running_steps_before_run_raises(workflow: Workflow) -> None:
+    """Calling running_steps() before workflow.run() should raise ContextStateError."""
+    ctx = Context(workflow)
+    with pytest.raises(ContextStateError, match="requires a running workflow"):
+        await ctx.running_steps()
+
+
+@pytest.mark.asyncio
+async def test_store_access_outside_step_works() -> None:
+    """Accessing ctx.store from handler code (outside step) should work."""
+
+    class SimpleWorkflow(Workflow):
+        @step
+        async def only(self, ev: StartEvent) -> StopEvent:
+            return StopEvent(result="done")
+
+    wf = SimpleWorkflow()
+    handler = wf.run()
+
+    # Access store from external context (handler code) should work
+    store = handler.ctx.store
+    assert isinstance(store, InMemoryStateStore)
+
+    # Verify reads and writes work
+    await store.set("key", "value")
+    assert await store.get("key") == "value"
+
+    await handler
+
+
+@pytest.mark.asyncio
+async def test_store_access_before_run_works() -> None:
+    """Accessing ctx.store before workflow.run() should return a staging store."""
+
+    class SimpleWorkflow(Workflow):
+        @step
+        async def only(self, ev: StartEvent) -> StopEvent:
+            return StopEvent(result="done")
+
+    wf = SimpleWorkflow()
+    ctx = Context(wf)
+
+    store = ctx.store
+    assert isinstance(store, InMemoryStateStore)
+    await store.set("key", "value")
+    assert await store.get("key") == "value"
+
+
+@pytest.mark.asyncio
+async def test_store_seed_before_run_visible_in_step() -> None:
+    """State seeded via ctx.store before run should be visible inside steps."""
+
+    class SeededWorkflow(Workflow):
+        @step
+        async def read_seed(self, ctx: Context, ev: StartEvent) -> StopEvent:
+            val = await ctx.store.get("seeded_key")
+            return StopEvent(result=val)
+
+    wf = SeededWorkflow()
+    ctx = Context(wf)
+    await ctx.store.set("seeded_key", "hello")
+
+    result = await wf.run(ctx=ctx)
+    assert result == "hello"
+
+
+@pytest.mark.asyncio
+async def test_store_seed_on_deserialized_context() -> None:
+    """Seeding state on a deserialized context should merge with existing state."""
+
+    class StatefulWorkflow(Workflow):
+        @step
+        async def first_step(self, ctx: Context, ev: StartEvent) -> StopEvent:
+            # Just pass through; state is set externally
+            return StopEvent(result="done")
+
+    class ReadWorkflow(Workflow):
+        @step
+        async def check(self, ctx: Context, ev: StartEvent) -> StopEvent:
+            old = await ctx.store.get("existing")
+            new = await ctx.store.get("added")
+            return StopEvent(result=f"{old}-{new}")
+
+    wf = StatefulWorkflow()
+
+    # First run to create some state
+    ctx = Context(wf)
+    await ctx.store.set("existing", "orig")
+    handler = wf.run(ctx=ctx)
+    await handler
+
+    # Serialize and restore (use ReadWorkflow for second run)
+    ctx_dict = ctx.to_dict()
+    read_wf = ReadWorkflow()
+    restored_ctx = Context.from_dict(read_wf, ctx_dict)
+
+    # Seed additional state before next run
+    await restored_ctx.store.set("added", "new")
+
+    result = await read_wf.run(ctx=restored_ctx)
+    assert result == "orig-new"
+
+
+@pytest.mark.asyncio
+async def test_store_continuation_with_pre_run_seeding() -> None:
+    """Continuation runs with pre-run seeding should carry state through."""
+
+    class CountWorkflow(Workflow):
+        @step
+        async def increment(self, ctx: Context, ev: StartEvent) -> StopEvent:
+            count = await ctx.store.get("count", default=0)
+            count += 1
+            await ctx.store.set("count", count)
+            return StopEvent(result=count)
+
+    wf = CountWorkflow()
+    ctx = Context(wf)
+
+    # Seed initial count
+    await ctx.store.set("count", 10)
+
+    # First run: 10 + 1 = 11
+    result = await wf.run(ctx=ctx)
+    assert result == 11
+
+    # Second run (continuation): 11 + 1 = 12
+    result = await wf.run(ctx=ctx)
+    assert result == 12
+
+
+@pytest.mark.asyncio
+async def test_to_dict_before_run_raises(workflow: Workflow) -> None:
+    """Calling to_dict() before workflow.run() should raise ContextStateError."""
+    ctx = Context(workflow)
+    with pytest.raises(ContextStateError, match="requires a running workflow"):
+        ctx.to_dict()
+
+
+@pytest.mark.asyncio
+async def test_stream_events_before_run_raises(workflow: Workflow) -> None:
+    """Calling stream_events() before workflow.run() should raise ContextStateError."""
+    ctx = Context(workflow)
+    with pytest.raises(ContextStateError, match="requires a running workflow"):
+        ctx.stream_events()
+
+
+# ============================================================================
+# Warning Tests
+# ============================================================================
+
+
+@pytest.mark.asyncio
+async def test_cancel_before_start_warns(workflow: Workflow) -> None:
+    """Calling cancel before run() should emit warning."""
+    # Clear the lru_cache to ensure warning fires
+    _warn_cancel_before_start.cache_clear()
+
+    ctx = Context(workflow)
+    with pytest.warns(UserWarning, match="cancel.*called before workflow started"):
+        ctx._workflow_cancel_run()
+
+
+@pytest.mark.asyncio
+async def test_send_event_before_start_raises(workflow: Workflow) -> None:
+    """Sending event before run() should raise ContextStateError."""
+    ctx = Context(workflow)
+    with pytest.raises(
+        ContextStateError, match="send_event.*called before workflow started"
+    ):
+        ctx.send_event(Event())
+
+
+@pytest.mark.asyncio
+async def test_is_running_in_step_warns() -> None:
+    """Calling is_running from within a step should emit deprecation warning."""
+    # Clear the lru_cache to ensure warning fires
+    _warn_is_running_in_step.cache_clear()
+
+    is_running_value = None
+
+    class TestWorkflow(Workflow):
+        @step
+        async def check_running(self, ctx: Context, ev: StartEvent) -> StopEvent:
+            nonlocal is_running_value
+            is_running_value = ctx.is_running
+            return StopEvent(result="done")
+
+    wf = TestWorkflow()
+    with pytest.warns(DeprecationWarning, match="is_running called from within a step"):
+        await wf.run()
+
+    # Should still return True despite the warning
+    assert is_running_value is True
+
+
+@pytest.mark.asyncio
+async def test_cancel_in_step_warns() -> None:
+    """Calling cancel from within a step should emit warning."""
+    # Clear the lru_cache to ensure warning fires
+    _warn_cancel_in_step.cache_clear()
+
+    class TestWorkflow(Workflow):
+        @step
+        async def cancel_self(self, ctx: Context, ev: StartEvent) -> StopEvent:
+            ctx._workflow_cancel_run()
+            return StopEvent(result="done")
+
+    wf = TestWorkflow()
+    with pytest.warns(UserWarning, match="cancel.*called from within a step"):
+        await wf.run()
+
+
+@pytest.mark.asyncio
+async def test_get_result_before_complete_raises() -> None:
+    """Calling get_result() while workflow still running should raise WorkflowRuntimeError."""
+    # Clear the lru_cache to ensure deprecation warning fires
+    _warn_get_result.cache_clear()
+
+    step_started = asyncio.Event()
+    step_continue = asyncio.Event()
+
+    class SlowWorkflow(Workflow):
+        @step
+        async def slow(self, ev: StartEvent) -> StopEvent:
+            step_started.set()
+            await step_continue.wait()
+            return StopEvent(result="done")
+
+    wf = SlowWorkflow()
+    handler = wf.run()
+
+    # Wait for step to start
+    await step_started.wait()
+
+    # Try to get result before workflow completes - should raise
+    with pytest.warns(DeprecationWarning):  # get_result is deprecated
+        with pytest.raises(WorkflowRuntimeError, match="is not complete"):
+            handler.ctx.get_result()
+
+    # Let workflow complete
+    step_continue.set()
+    await handler
+
+
+@pytest.mark.asyncio
+async def test_get_result_pre_context_raises(workflow: Workflow) -> None:
+    """Calling get_result() before run() should raise ContextStateError."""
+    # Clear the lru_cache to ensure deprecation warning fires
+    _warn_get_result.cache_clear()
+
+    ctx = Context(workflow)
+
+    with pytest.warns(DeprecationWarning):  # get_result is deprecated
+        with pytest.raises(ContextStateError, match="requires a running workflow"):
+            ctx.get_result()

packages/llama-index-workflows/tests/context/test_context_preservation.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 LlamaIndex Inc.
+
+"""Tests for context state preservation when passing ctx to workflow.run().
+
+These tests verify that when a Context is passed to Workflow.run(ctx=ctx),
+the original context object is updated with run state and can be used for
+subsequent operations.
+
+See: thoughts/shared/plans/2026-01-23-context-state-preservation-fix.md
+"""
+
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+from workflows.context import Context
+from workflows.decorators import step
+from workflows.errors import ContextStateError, WorkflowRuntimeError
+from workflows.events import StartEvent, StopEvent
+from workflows.workflow import Workflow
+
+
+class CounterWorkflow(Workflow):
+    """Simple workflow that increments a counter in state."""
+
+    @step
+    async def count(self, ctx: Context, ev: StartEvent) -> StopEvent:
+        count = await ctx.store.get("count", default=0)
+        count += 1
+        await ctx.store.set("count", count)
+        return StopEvent(result=count)
+
+
+@pytest.mark.asyncio
+async def test_original_ctx_is_handler_ctx() -> None:
+    """ctx passed to run() should be the same object as handler.ctx"""
+    wf = CounterWorkflow()
+    ctx = Context(wf)
+    handler = wf.run(ctx=ctx)
+    await handler
+    assert ctx is handler.ctx
+
+
+@pytest.mark.asyncio
+async def test_original_ctx_to_dict_works() -> None:
+    """ctx.to_dict() should work after run (not just handler.ctx.to_dict())"""
+    wf = CounterWorkflow()
+    ctx = Context(wf)
+    handler = wf.run(ctx=ctx)
+    await handler
+    ctx_dict = ctx.to_dict()  # Should NOT raise
+    assert "state" in ctx_dict
+
+
+@pytest.mark.asyncio
+async def test_sequential_runs_accumulate_state() -> None:
+    """Three sequential runs should produce 1, 2, 3"""
+    wf = CounterWorkflow()
+    ctx = Context(wf)
+    r1 = await wf.run(ctx=ctx)  # count: 0 -> 1
+    r2 = await wf.run(ctx=ctx)  # count: 1 -> 2
+    r3 = await wf.run(ctx=ctx)  # count: 2 -> 3
+    assert (r1, r2, r3) == (1, 2, 3)
+
+
+@pytest.mark.asyncio
+async def test_concurrent_runs_same_context_raises() -> None:
+    """Starting a second run while first is running should raise"""
+
+    class SlowWorkflow(Workflow):
+        @step
+        async def slow_step(self, ev: StartEvent) -> StopEvent:
+            await asyncio.sleep(0.1)
+            return StopEvent(result="done")
+
+    wf = SlowWorkflow()
+    ctx = Context(wf)
+    handler1 = wf.run(ctx=ctx)
+    with pytest.raises(ContextStateError, match="already running"):
+        wf.run(ctx=ctx)
+    await handler1  # Clean up
+
+
+@pytest.mark.asyncio
+async def test_concurrent_runs_different_contexts_ok() -> None:
+    """Different contexts can run concurrently"""
+
+    class SlowWorkflow(Workflow):
+        @step
+        async def slow_step(self, ev: StartEvent) -> StopEvent:
+            await asyncio.sleep(0.01)
+            return StopEvent(result="done")
+
+    wf = SlowWorkflow()
+    ctx1, ctx2 = Context(wf), Context(wf)
+    h1 = wf.run(ctx=ctx1)
+    h2 = wf.run(ctx=ctx2)  # Should NOT raise
+    await h1
+    await h2
+
+
+@pytest.mark.asyncio
+async def test_from_dict_then_sequential_runs() -> None:
+    """Restored context should work for sequential runs"""
+    wf = CounterWorkflow()
+    ctx1 = Context(wf)
+    await wf.run(ctx=ctx1)  # count -> 1
+    ctx2 = Context.from_dict(wf, ctx1.to_dict())
+    r = await wf.run(ctx=ctx2)  # count -> 2
+    assert r == 2
+
+
+@pytest.mark.asyncio
+async def test_stream_events_twice_raises() -> None:
+    """Streaming events twice on same handler should raise"""
+    wf = CounterWorkflow()
+    handler = wf.run()
+
+    async for _ in handler.stream_events():
+        pass
+
+    with pytest.raises(WorkflowRuntimeError, match="already been consumed"):
+        async for _ in handler.stream_events():
+            pass

packages/llama-index-workflows/tests/runtime/conftest.py

@@ -3,39 +3,91 @@
 
 """Test fixtures and utilities for runtime tests."""
 
+from __future__ import annotations
+
 import asyncio
 import time
-from typing import Any, AsyncGenerator, Optional
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional
 
 import pytest
 import time_machine
-from workflows.events import Event, StopEvent
-from workflows.runtime.types.plugin import ControlLoopFunction, Plugin, WorkflowRuntime
-from workflows.runtime.types.step_function import StepWorkerFunction
+from workflows.events import Event, StartEvent, StopEvent
+from workflows.runtime.types.internal_state import BrokerConfig, BrokerState
+
+if TYPE_CHECKING:
+    from workflows.context.serializers import BaseSerializer
+    from workflows.context.state_store import InMemoryStateStore
+from workflows.plugins.basic import get_current_run_id
+from workflows.runtime.types.plugin import (
+    ExternalRunAdapter,
+    InternalRunAdapter,
+    RegisteredWorkflow,
+    Runtime,
+    SnapshottableAdapter,
+    V2RuntimeCompatibilityShim,
+)
+from workflows.runtime.types.step_function import (
+    as_step_worker_functions,
+    create_workflow_run_function,
+)
 from workflows.runtime.types.ticks import WorkflowTick
 from workflows.workflow import Workflow
 
 
-class MockPlugin(Plugin):
-    def register(
+class MockRuntime(Runtime):
+    """Mock runtime that stores adapters for test access."""
+
+    def __init__(self) -> None:
+        self._adapters: dict[str, MockRunAdapter] = {}
+        self._current_run_id: str | None = None
+
+    def register(self, workflow: Workflow) -> RegisteredWorkflow:
+        return RegisteredWorkflow(
+            workflow=workflow,
+            workflow_run_fn=create_workflow_run_function(workflow),
+            steps=as_step_worker_functions(workflow),
+        )
+
+    def get_internal_adapter(self) -> 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)
+        return self._adapters[run_id]
+
+    def get_external_adapter(self, run_id: str) -> ExternalRunAdapter:
+        if run_id not in self._adapters:
+            self._adapters[run_id] = MockRunAdapter(run_id)
+        return self._adapters[run_id]
+
+    def run_workflow(
         self,
+        run_id: str,
         workflow: Workflow,
-        workflow_function: ControlLoopFunction,
-        steps: dict[str, StepWorkerFunction],
-    ) -> None:
-        return
+        init_state: BrokerState,
+        start_event: StartEvent | None = None,
+        serialized_state: dict[str, Any] | None = None,
+        serializer: "BaseSerializer | None" = None,
+    ) -> ExternalRunAdapter:
+        self._current_run_id = run_id
+        return self.get_external_adapter(run_id)
 
-    def new_runtime(self, run_id: str) -> WorkflowRuntime:
-        return MockRuntimePlugin(run_id)
+    def set_adapter(self, run_id: str, adapter: "MockRunAdapter") -> None:
+        """Set a specific adapter for a run_id (for test setup)."""
+        self._adapters[run_id] = adapter
 
 
-class MockRuntimePlugin(WorkflowRuntime):
-    """Mock WorkflowRuntime for testing control loops."""
+class MockRunAdapter(
+    InternalRunAdapter,
+    ExternalRunAdapter,
+    SnapshottableAdapter,
+    V2RuntimeCompatibilityShim,
+):
+    """Mock RunAdapter for testing control loops. Supports snapshot/replay."""
 
     def __init__(
         self, run_id: str, traveller: Optional[time_machine.Coordinates] = None
     ) -> None:
-        self.run_id = run_id
+        self._run_id = run_id
         # Queue for events sent from external sources (e.g., via send_event)
         self._external_queue: asyncio.Queue[WorkflowTick] = asyncio.Queue()
         # Queue for events published to the event stream (e.g., for UI/callbacks)
@@ -44,10 +96,42 @@ class MockRuntimePlugin(WorkflowRuntime):
         self._traveller = traveller
         # Current time in seconds, can be advanced manually for testing
         self._current_time: float = time.time()
+        # Recorded ticks for snapshot/replay
+        self._ticks: list[WorkflowTick] = []
+        # State store for context
+        self._state_store: "InMemoryStateStore[Any] | None" = None
+        # Result tracking for get_result/cancel
+        self._result: asyncio.Future[StopEvent] = asyncio.Future()
+        self._cancelled: bool = False
+
+    @property
+    def run_id(self) -> str:
+        return self._run_id
+
+    @property
+    def tags(self) -> dict[str, Any]:
+        return {"run_id": self._run_id}
+
+    @property
+    def init_state(self) -> BrokerState:
+        # Return a minimal BrokerState for testing
+        return BrokerState(
+            is_running=False,
+            config=BrokerConfig(steps={}, timeout=None),
+            workers={},
+        )
+
+    def on_tick(self, tick: WorkflowTick) -> None:
+        """Record a tick for replay."""
+        self._ticks.append(tick)
+
+    def replay(self) -> list[WorkflowTick]:
+        """Return recorded ticks for replay."""
+        return self._ticks
 
     async def close(self) -> None:
         """
-        Close the plugin.
+        Close the adapter.
         """
         pass
 
@@ -67,12 +151,6 @@ class MockRuntimePlugin(WorkflowRuntime):
     async def send_event(self, tick: WorkflowTick) -> None:
         await self._external_queue.put(tick)
 
-    async def register_step_worker(self, step_name: str, step_worker: Any) -> Any:
-        return step_worker
-
-    async def register_workflow_function(self, workflow_function: Any) -> Any:
-        return workflow_function
-
     async def get_now(self) -> float:
         if self._traveller is not None:
             return time.time()
@@ -93,16 +171,48 @@ class MockRuntimePlugin(WorkflowRuntime):
     def has_stream_events(self) -> bool:
         return not self._event_stream.empty()
 
+    def get_state_store(self) -> "InMemoryStateStore[Any] | None":
+        return self._state_store
+
+    def set_state_store(self, state_store: "InMemoryStateStore[Any]") -> None:
+        self._state_store = state_store
+
+    async def get_result(self) -> StopEvent:
+        """Get the result of the workflow run."""
+        return await self._result
+
+    def get_result_or_none(self) -> StopEvent | None:
+        """Get the result if completed, otherwise None."""
+        if self._result.done() and not self._result.cancelled():
+            return self._result.result()
+        return None
+
+    @property
+    def is_running(self) -> bool:
+        """Check if the workflow run is still running."""
+        return not self._result.done() and not self._cancelled
+
+    def abort(self) -> None:
+        """Abort by cancelling the result future."""
+        if not self._result.done():
+            self._result.cancel()
+            self._cancelled = True
+
+    def set_result(self, result: StopEvent) -> None:
+        """Set the result (for test setup)."""
+        if not self._result.done():
+            self._result.set_result(result)
+
 
 @pytest.fixture
-async def test_plugin() -> MockRuntimePlugin:
-    return MockRuntimePlugin(run_id="test")
+async def test_plugin() -> MockRunAdapter:
+    return MockRunAdapter(run_id="test")
 
 
 @pytest.fixture
 async def test_plugin_with_time_machine() -> AsyncGenerator[
-    tuple[MockRuntimePlugin, time_machine.Coordinates], None
+    tuple[MockRunAdapter, time_machine.Coordinates], None
 ]:
-    """Plugin with time-machine at epoch 1000.0, tick=True."""
+    """Adapter with time-machine at epoch 1000.0, tick=True."""
     with time_machine.travel("2026-01-07T12:27:00.000-08:00", tick=True) as traveller:
-        yield MockRuntimePlugin(run_id="test", traveller=traveller), traveller
+        yield MockRunAdapter(run_id="test", traveller=traveller), traveller

packages/llama-index-workflows/tests/runtime/test_control_loop.py

@@ -19,6 +19,7 @@ from typing import Coroutine, Optional, Union
 import pytest
 import time_machine
 from workflows.context import Context
+from workflows.context.state_store import DictState, InMemoryStateStore
 from workflows.decorators import step
 from workflows.errors import WorkflowCancelledByUser, WorkflowTimeoutError
 from workflows.events import (
@@ -33,15 +34,16 @@ from workflows.events import (
     WorkflowIdleEvent,
     WorkflowTimedOutEvent,
 )
+from workflows.plugins.basic import setting_run_id
 from workflows.retry_policy import ConstantDelayRetryPolicy, RetryPolicy
 from workflows.runtime.control_loop import control_loop
 from workflows.runtime.types.internal_state import BrokerState
+from workflows.runtime.types.plugin import RunContext, run_context
 from workflows.runtime.types.step_function import as_step_worker_function
 from workflows.runtime.types.ticks import TickAddEvent, TickCancelRun
-from workflows.runtime.workflow_registry import workflow_registry
 from workflows.workflow import Workflow
 
-from .conftest import MockRuntimePlugin  # type: ignore[import]
+from .conftest import MockRunAdapter, MockRuntime
 
 
 class IntermediateEvent(Event):
@@ -144,44 +146,50 @@ class CollectWorkflow(Workflow):
 
 
 def run_control_loop(
-    workflow: Workflow, start_event: Optional[StartEvent], plugin: MockRuntimePlugin
+    workflow: Workflow,
+    start_event: Optional[StartEvent],
+    test_runtime: MockRunAdapter,
 ) -> Coroutine[None, None, StopEvent]:
     step_workers = {}
     for name, step_func in workflow._get_steps().items():
         unbound = getattr(step_func, "__func__", step_func)
         step_workers[name] = as_step_worker_function(unbound)
-    ctx = Context(workflow=workflow)
-    ctx._broker_run = ctx._init_broker(workflow, plugin=plugin)
     run_id = str(uuid.uuid4())
-    workflow_registry.register_run(
-        run_id=run_id,
-        workflow=workflow,
-        plugin=plugin,
-        context=ctx,
-        steps=step_workers,
-    )
+    # Set up mock runtime with the test adapter
+    mock_runtime = MockRuntime()
+    test_runtime.set_state_store(InMemoryStateStore(DictState()))
+    mock_runtime.set_adapter(run_id, test_runtime)
+    # Override workflow's runtime to use mock
+    workflow._runtime = mock_runtime
+    with setting_run_id(run_id):
+        ctx = Context._create_internal(workflow=workflow)
 
     async def _run() -> StopEvent:
-        try:
-            return await control_loop(
-                start_event=start_event,
-                init_state=BrokerState.from_workflow(workflow),
-                run_id=run_id,
+        with setting_run_id(run_id):
+            run_ctx = RunContext(
+                workflow=workflow,
+                run_adapter=test_runtime,
+                context=ctx,
+                steps=step_workers,
             )
-        finally:
-            workflow_registry.delete_run(run_id)
+            with run_context(run_ctx):
+                return await control_loop(
+                    start_event=start_event,
+                    init_state=BrokerState.from_workflow(workflow),
+                    run_id=run_id,
+                )
 
     return _run()
 
 
 async def wait_for_stop_event(
-    plugin: MockRuntimePlugin, timeout: float = 1.0
+    plugin: MockRunAdapter, timeout: float = 1.0
 ) -> Optional[StopEvent]:
     """
     Helper to wait for a StopEvent in the event stream.
 
     Args:
-        plugin: The MockRuntimePlugin to read events from
+        plugin: The MockRunAdapter to read events from
         timeout: Maximum time to wait for StopEvent (default: 1.0 seconds)
 
     Returns:
@@ -202,7 +210,7 @@ async def wait_for_stop_event(
 
 
 @pytest.mark.asyncio
-async def test_control_loop_happy_path(test_plugin: MockRuntimePlugin) -> None:
+async def test_control_loop_happy_path(test_plugin: MockRunAdapter) -> None:
     """
     Test the happy path through the control loop.
 
@@ -217,7 +225,7 @@ async def test_control_loop_happy_path(test_plugin: MockRuntimePlugin) -> None:
     result = await run_control_loop(
         workflow=SimpleWorkflow(timeout=1.0),
         start_event=StartEvent(),
-        plugin=test_plugin,
+        test_runtime=test_plugin,
     )
 
     # Verify the workflow completed with expected result
@@ -226,7 +234,9 @@ async def test_control_loop_happy_path(test_plugin: MockRuntimePlugin) -> None:
 
 
 @pytest.mark.asyncio
-async def test_control_loop_with_external_event(test_plugin: MockRuntimePlugin) -> None:
+async def test_control_loop_with_external_event(
+    test_plugin: MockRunAdapter,
+) -> None:
     """
     Test that external events can be sent to a running workflow.
 
@@ -252,7 +262,7 @@ async def test_control_loop_with_external_event(test_plugin: MockRuntimePlugin)
         run_control_loop(
             workflow=workflow,
             start_event=None,
-            plugin=test_plugin,
+            test_runtime=test_plugin,
         )
     )
 
@@ -268,7 +278,7 @@ async def test_control_loop_with_external_event(test_plugin: MockRuntimePlugin)
 
 
 @pytest.mark.asyncio
-async def test_control_loop_timeout(test_plugin: MockRuntimePlugin) -> None:
+async def test_control_loop_timeout(test_plugin: MockRunAdapter) -> None:
     """
     Test that workflow timeout raises WorkflowTimeoutError and publishes WorkflowTimedOutEvent.
 
@@ -288,7 +298,7 @@ async def test_control_loop_timeout(test_plugin: MockRuntimePlugin) -> None:
         run_control_loop(
             workflow=wf,
             start_event=StartEvent(),
-            plugin=test_plugin,
+            test_runtime=test_plugin,
         )
     )
 
@@ -311,7 +321,7 @@ async def test_control_loop_timeout(test_plugin: MockRuntimePlugin) -> None:
 
 
 @pytest.mark.asyncio
-async def test_control_loop_retry_policy(test_plugin: MockRuntimePlugin) -> None:
+async def test_control_loop_retry_policy(test_plugin: MockRunAdapter) -> None:
     """
     Test that retry policy works correctly when a step fails initially but succeeds on retry.
     """
@@ -331,7 +341,7 @@ async def test_control_loop_retry_policy(test_plugin: MockRuntimePlugin) -> None
     result = await run_control_loop(
         workflow=wf,
         start_event=StartEvent(),
-        plugin=test_plugin,
+        test_runtime=test_plugin,
     )
 
     assert isinstance(result, StopEvent)
@@ -340,7 +350,7 @@ async def test_control_loop_retry_policy(test_plugin: MockRuntimePlugin) -> None
 
 @pytest.mark.asyncio
 async def test_control_loop_step_failure_publishes_stop_event(
-    test_plugin: MockRuntimePlugin,
+    test_plugin: MockRunAdapter,
 ) -> None:
     """
     Test that when a step fails permanently (retries exhausted),
@@ -359,7 +369,7 @@ async def test_control_loop_step_failure_publishes_stop_event(
         run_control_loop(
             workflow=wf,
             start_event=StartEvent(),
-            plugin=test_plugin,
+            test_runtime=test_plugin,
         )
     )
 
@@ -394,7 +404,7 @@ async def test_control_loop_step_failure_publishes_stop_event(
 
 
 @pytest.mark.asyncio
-async def test_control_loop_waiter_resolution(test_plugin: MockRuntimePlugin) -> None:
+async def test_control_loop_waiter_resolution(test_plugin: MockRunAdapter) -> None:
     class Awaited(Event):
         tag: str
 
@@ -414,7 +424,7 @@ async def test_control_loop_waiter_resolution(test_plugin: MockRuntimePlugin) ->
         run_control_loop(
             workflow=wf,
             start_event=StartEvent(),
-            plugin=test_plugin,
+            test_runtime=test_plugin,
         )
     )
 
@@ -437,7 +447,7 @@ async def test_control_loop_waiter_resolution(test_plugin: MockRuntimePlugin) ->
 
 @pytest.mark.asyncio
 async def test_control_loop_input_required_published_to_stream(
-    test_plugin: MockRuntimePlugin,
+    test_plugin: MockRunAdapter,
 ) -> None:
     """
     Test that InputRequiredEvent is automatically published to the outward stream.
@@ -466,7 +476,7 @@ async def test_control_loop_input_required_published_to_stream(
         run_control_loop(
             workflow=wf,
             start_event=StartEvent(),
-            plugin=test_plugin,
+            test_runtime=test_plugin,
         )
     )
 
@@ -494,14 +504,14 @@ async def test_control_loop_input_required_published_to_stream(
 
 @pytest.mark.asyncio
 async def test_control_loop_collect_events_same_type(
-    test_plugin: MockRuntimePlugin,
+    test_plugin: MockRunAdapter,
 ) -> None:
     wf = CollectWorkflow(timeout=1.0)
     result = await asyncio.create_task(
         run_control_loop(
             workflow=wf,
             start_event=StartEvent(),
-            plugin=test_plugin,
+            test_runtime=test_plugin,
         )
     )
 
@@ -511,26 +521,26 @@ async def test_control_loop_collect_events_same_type(
 
 @pytest.mark.asyncio
 async def test_control_loop_collect_events_multiple_types(
-    test_plugin: MockRuntimePlugin,
+    test_plugin: MockRunAdapter,
 ) -> None:
     wf = CollectMultipleEventTypesWorkflow(timeout=1.0)
     result = await run_control_loop(
         workflow=wf,
         start_event=StartEvent(),
-        plugin=test_plugin,
+        test_runtime=test_plugin,
     )
     assert isinstance(result, StopEvent)
     assert result.result == "sum_33"
 
 
 @pytest.mark.asyncio
-async def test_control_loop_stream_events(test_plugin: MockRuntimePlugin) -> None:
+async def test_control_loop_stream_events(test_plugin: MockRunAdapter) -> None:
     wf = SimpleWorkflow(timeout=5.0)
     task = asyncio.create_task(
         run_control_loop(
             workflow=wf,
             start_event=StartEvent(),
-            plugin=test_plugin,
+            test_runtime=test_plugin,
         )
     )
 
@@ -563,7 +573,7 @@ class SomeEvent(HumanResponseEvent):
 
 
 @pytest.mark.asyncio
-async def test_control_loop_per_step_routing(test_plugin: MockRuntimePlugin) -> None:
+async def test_control_loop_per_step_routing(test_plugin: MockRunAdapter) -> None:
     class RouteWorkflow(Workflow):
         @step
         async def starter(self, ev: StartEvent) -> Optional[StopEvent]:
@@ -582,7 +592,7 @@ async def test_control_loop_per_step_routing(test_plugin: MockRuntimePlugin) ->
         run_control_loop(
             workflow=wf,
             start_event=StartEvent(),
-            plugin=test_plugin,
+            test_runtime=test_plugin,
         )
     )
 
@@ -596,7 +606,7 @@ async def test_control_loop_per_step_routing(test_plugin: MockRuntimePlugin) ->
 
 @pytest.mark.asyncio
 async def test_control_loop_concurrency_queueing(
-    test_plugin: MockRuntimePlugin,
+    test_plugin: MockRunAdapter,
 ) -> None:
     class LimitedWorkflow(Workflow):
         @step(num_workers=1)
@@ -610,7 +620,7 @@ async def test_control_loop_concurrency_queueing(
     task = asyncio.create_task(
         run_control_loop(
             workflow=wf,
-            plugin=test_plugin,
+            test_runtime=test_plugin,
             start_event=None,
         )
     )
@@ -636,7 +646,7 @@ async def test_control_loop_concurrency_queueing(
 
 
 @pytest.mark.asyncio
-async def test_control_loop_user_cancellation(test_plugin: MockRuntimePlugin) -> None:
+async def test_control_loop_user_cancellation(test_plugin: MockRunAdapter) -> None:
     """
     Test that user cancellation raises WorkflowCancelledByUser and publishes WorkflowCancelledEvent.
 
@@ -655,7 +665,7 @@ async def test_control_loop_user_cancellation(test_plugin: MockRuntimePlugin) ->
     task = asyncio.create_task(
         run_control_loop(
             workflow=wf,
-            plugin=test_plugin,
+            test_runtime=test_plugin,
             start_event=StartEvent(),
         )
     )
@@ -682,7 +692,7 @@ async def test_control_loop_user_cancellation(test_plugin: MockRuntimePlugin) ->
 
 @pytest.mark.asyncio
 async def test_control_loop_retry_with_delay(
-    test_plugin_with_time_machine: tuple[MockRuntimePlugin, time_machine.Coordinates],
+    test_plugin_with_time_machine: tuple[MockRunAdapter, time_machine.Coordinates],
 ) -> None:
     """Test that retry delay is enforced between attempts."""
     test_plugin, _ = test_plugin_with_time_machine
@@ -705,7 +715,7 @@ async def test_control_loop_retry_with_delay(
     result = await run_control_loop(
         workflow=wf,
         start_event=StartEvent(),
-        plugin=test_plugin,
+        test_runtime=test_plugin,
     )
 
     assert isinstance(result, StopEvent)
@@ -723,7 +733,7 @@ async def test_control_loop_retry_with_delay(
 
 @pytest.mark.asyncio
 async def test_control_loop_retry_gives_up_after_max_attempts(
-    test_plugin_with_time_machine: tuple[MockRuntimePlugin, time_machine.Coordinates],
+    test_plugin_with_time_machine: tuple[MockRunAdapter, time_machine.Coordinates],
 ) -> None:
     """Test that workflow fails after exhausting maximum_attempts."""
     test_plugin, _ = test_plugin_with_time_machine
@@ -747,7 +757,7 @@ async def test_control_loop_retry_gives_up_after_max_attempts(
         await run_control_loop(
             workflow=wf,
             start_event=StartEvent(),
-            plugin=test_plugin,
+            test_runtime=test_plugin,
         )
 
     assert wf.attempt_count == max_attempts
@@ -755,7 +765,7 @@ async def test_control_loop_retry_gives_up_after_max_attempts(
 
 @pytest.mark.asyncio
 async def test_control_loop_retry_exhaustion_respects_total_time(
-    test_plugin_with_time_machine: tuple[MockRuntimePlugin, time_machine.Coordinates],
+    test_plugin_with_time_machine: tuple[MockRunAdapter, time_machine.Coordinates],
 ) -> None:
     """Test that retry policy receives correct elapsed_time across retries."""
     test_plugin, _ = test_plugin_with_time_machine
@@ -794,7 +804,7 @@ async def test_control_loop_retry_exhaustion_respects_total_time(
     result = await run_control_loop(
         workflow=wf,
         start_event=StartEvent(),
-        plugin=test_plugin,
+        test_runtime=test_plugin,
     )
 
     assert isinstance(result, StopEvent)
@@ -826,7 +836,7 @@ async def test_control_loop_retry_exhaustion_respects_total_time(
 
 @pytest.mark.asyncio
 async def test_control_loop_emits_idle_event_when_waiting(
-    test_plugin: MockRuntimePlugin,
+    test_plugin: MockRunAdapter,
 ) -> None:
     """
     Test that WorkflowIdleEvent is emitted when workflow becomes idle waiting for event.
@@ -852,7 +862,7 @@ async def test_control_loop_emits_idle_event_when_waiting(
         run_control_loop(
             workflow=wf,
             start_event=StartEvent(),
-            plugin=test_plugin,
+            test_runtime=test_plugin,
         )
     )
 
@@ -888,7 +898,7 @@ async def test_control_loop_emits_idle_event_when_waiting(
 
 @pytest.mark.asyncio
 async def test_control_loop_idle_event_not_emitted_on_completion(
-    test_plugin: MockRuntimePlugin,
+    test_plugin: MockRunAdapter,
 ) -> None:
     """
     Test that WorkflowIdleEvent is NOT emitted when workflow completes normally.
@@ -900,7 +910,7 @@ async def test_control_loop_idle_event_not_emitted_on_completion(
     result = await run_control_loop(
         workflow=SimpleWorkflow(timeout=1.0),
         start_event=StartEvent(),
-        plugin=test_plugin,
+        test_runtime=test_plugin,
     )
 
     # Verify the workflow completed

packages/llama-index-workflows/tests/runtime/test_identity_weak_ref.py

@@ -1,37 +0,0 @@
-from __future__ import annotations
-
-import gc
-import weakref
-
-import pytest
-from workflows.runtime.types._identity_weak_ref import IdentityWeakKeyDict
-
-
-class Unhashable:
-    def __hash__(self) -> int:
-        raise TypeError("Unhashable")
-
-
-def test_identity_weak_key_dict_removes_entry_when_object_unreferenced() -> None:
-    with pytest.raises(TypeError):
-        hash(Unhashable())
-
-    d: IdentityWeakKeyDict[Unhashable, str] = IdentityWeakKeyDict()
-
-    obj = Unhashable()
-    d[obj] = "value"
-
-    # While the object is strongly referenced, it should be present
-    assert obj in d
-    assert d.get(obj) == "value"
-
-    # Keep a weak reference to verify collection happened
-    w = weakref.ref(obj)
-
-    # Drop strong reference and force collection
-    del obj
-    gc.collect()
-
-    # Object should be collected and the dict entry removed via callback
-    assert w() is None
-    assert d._d == {}

packages/llama-index-workflows/tests/runtime/test_runtime_lifecycle.py

@@ -0,0 +1,203 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 LlamaIndex Inc.
+"""Tests for Runtime lifecycle (registering context manager, context variable)."""
+
+from __future__ import annotations
+
+import pytest
+from workflows import Workflow, step
+from workflows.events import StartEvent, StopEvent
+from workflows.plugins import BasicRuntime, basic_runtime, get_current_runtime
+from workflows.runtime.types.plugin import (
+    Runtime,
+    _current_runtime,
+)
+
+
+class SimpleWorkflow(Workflow):
+    @step
+    async def start(self, ev: StartEvent) -> StopEvent:
+        return StopEvent(result="done")
+
+
+def test_basic_runtime_is_runtime_instance() -> None:
+    """BasicRuntime extends Runtime ABC."""
+    assert isinstance(basic_runtime, Runtime)
+
+
+def test_basic_runtime_has_lifecycle_methods() -> None:
+    """BasicRuntime has launch() and destroy() methods."""
+    runtime = BasicRuntime()
+    # Should not raise
+    runtime.launch()
+    runtime.destroy()
+
+
+def test_get_current_runtime_returns_basic_runtime_by_default() -> None:
+    """When no context-scoped runtime, get_current_runtime returns basic_runtime."""
+    runtime = get_current_runtime()
+    assert runtime is basic_runtime
+
+
+def test_registering_sets_current_runtime() -> None:
+    """registering() context manager sets the context-scoped runtime."""
+    custom_runtime = BasicRuntime()
+
+    with custom_runtime.registering():
+        current = _current_runtime.get()
+        assert current is custom_runtime
+
+
+def test_registering_resets_on_exit() -> None:
+    """registering() context manager resets the context variable on exit."""
+    custom_runtime = BasicRuntime()
+
+    with custom_runtime.registering():
+        pass
+
+    current = _current_runtime.get()
+    assert current is None
+
+
+def test_registering_returns_runtime() -> None:
+    """registering() context manager returns the runtime when entering."""
+    custom_runtime = BasicRuntime()
+
+    with custom_runtime.registering() as r:
+        assert r is custom_runtime
+
+
+def test_registering_does_not_call_launch_on_exit() -> None:
+    """registering() context manager does NOT call launch() on exit."""
+    launched = False
+
+    class TrackingRuntime(BasicRuntime):
+        def launch(self) -> None:
+            nonlocal launched
+            launched = True
+
+    with TrackingRuntime().registering():
+        pass
+
+    # Key change: launch() should NOT be called
+    assert not launched
+
+
+def test_registering_does_not_call_launch_on_exception() -> None:
+    """registering() context manager does NOT call launch() when exception is raised."""
+    launched = False
+
+    class TrackingRuntime(BasicRuntime):
+        def launch(self) -> None:
+            nonlocal launched
+            launched = True
+
+    with pytest.raises(ValueError):
+        with TrackingRuntime().registering():
+            raise ValueError("test error")
+
+    assert not launched
+
+
+def test_get_current_runtime_returns_context_runtime_when_set() -> None:
+    """get_current_runtime returns context-scoped runtime when set."""
+    custom_runtime = BasicRuntime()
+
+    with custom_runtime.registering():
+        assert get_current_runtime() is custom_runtime
+
+
+def test_nested_registering_contexts() -> None:
+    """Nested registering() context managers restore correctly."""
+    outer_runtime = BasicRuntime()
+    inner_runtime = BasicRuntime()
+
+    with outer_runtime.registering():
+        assert get_current_runtime() is outer_runtime
+
+        with inner_runtime.registering():
+            assert get_current_runtime() is inner_runtime
+
+        # After inner exits, should restore to outer
+        assert get_current_runtime() is outer_runtime
+
+    # After outer exits, should restore to basic_runtime
+    assert get_current_runtime() is basic_runtime
+
+
+def test_explicit_runtime_parameter() -> None:
+    """Workflow with runtime= uses that runtime."""
+    custom_runtime = BasicRuntime()
+    wf = SimpleWorkflow(runtime=custom_runtime)
+    assert wf.runtime is custom_runtime
+
+
+def test_registering_context_manager() -> None:
+    """Workflows inside registering() use context-scoped runtime."""
+    custom_runtime = BasicRuntime()
+
+    with custom_runtime.registering():
+        wf = SimpleWorkflow()
+        assert wf.runtime is custom_runtime
+
+
+def test_explicit_overrides_registering() -> None:
+    """Explicit runtime= takes precedence over registering() context."""
+    context_runtime = BasicRuntime()
+    explicit_runtime = BasicRuntime()
+
+    with context_runtime.registering():
+        wf = SimpleWorkflow(runtime=explicit_runtime)
+        assert wf.runtime is explicit_runtime
+
+
+def test_fallback_to_basic_runtime() -> None:
+    """Workflow without runtime uses basic_runtime."""
+    wf = SimpleWorkflow()
+    assert wf.runtime is basic_runtime
+
+
+def test_workflow_runs_after_registering_exit() -> None:
+    """Workflow can run after registering() exits (requires explicit launch())."""
+    custom_runtime = BasicRuntime()
+
+    with custom_runtime.registering():
+        SimpleWorkflow()
+
+    # Explicit launch() should be called before running
+    custom_runtime.launch()
+    # BasicRuntime doesn't actually need launch(), but the pattern should work
+
+
+def test_registering_yields_runtime() -> None:
+    """registering() context manager yields the runtime."""
+    custom_runtime = BasicRuntime()
+
+    with custom_runtime.registering() as r:
+        assert r is custom_runtime
+        assert isinstance(r, Runtime)
+
+
+def test_empty_registering_block() -> None:
+    """Empty registering() block is valid no-op."""
+    custom_runtime = BasicRuntime()
+
+    # Should not raise
+    with custom_runtime.registering():
+        pass
+
+    # Context should be reset
+    assert get_current_runtime() is basic_runtime
+
+
+def test_registering_with_exception_still_resets_context() -> None:
+    """Context reset even on exception inside registering()."""
+    custom_runtime = BasicRuntime()
+
+    with pytest.raises(RuntimeError):
+        with custom_runtime.registering():
+            assert get_current_runtime() is custom_runtime
+            raise RuntimeError("test")
+
+    # Context should still be reset after exception
+    assert get_current_runtime() is basic_runtime

packages/llama-index-workflows/tests/runtime/test_workflow_tracker.py

@@ -0,0 +1,187 @@
+# 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_child_state_inheritance.py

@@ -76,7 +76,9 @@ class ChildWorkflowWithChildState(BaseWorkflowWithBaseState):
     async def child_step(self, ctx: Context[ChildState], ev: MiddleEvent) -> StopEvent:
         # Child step can access both base and child fields
         await ctx.store.set("child_field", "set_by_child_step")
-        return StopEvent()
+        # Return state as result for testing
+        state = await ctx.store.get_state()
+        return StopEvent(result=state)
 
 
 @pytest.mark.asyncio
@@ -95,9 +97,7 @@ async def test_subtype_inheritance_works() -> None:
 
     result = await test_runner.run()
 
-    ctx = result.ctx
-    assert ctx is not None
-    state = await ctx.store.get_state()
+    state = result.result
 
     # Verify state is ChildState
     assert isinstance(state, ChildState)
@@ -128,7 +128,9 @@ class WorkflowWithBaseStateSetState(Workflow):
         # and sets it. This should merge, not obliterate child fields.
         new_state = BaseState(base_field="modified_by_base_step")
         await ctx.store.set_state(new_state)  # type: ignore[arg-type]
-        return StopEvent()
+        # Return state as result for testing
+        state = await ctx.store.get_state()
+        return StopEvent(result=state)
 
 
 @pytest.mark.asyncio
@@ -145,9 +147,7 @@ async def test_set_state_with_parent_type_merges_fields() -> None:
 
     result = await test_runner.run()
 
-    ctx = result.ctx
-    assert ctx is not None
-    state = await ctx.store.get_state()
+    state = result.result
 
     # The base field was modified
     assert state.base_field == "modified_by_base_step"
@@ -277,7 +277,9 @@ class ChildWorkflowConsistent(BaseWorkflowConsistent):
         state = await ctx.store.get_state()
         await ctx.store.set("child_field", "modified_by_child_step")
         await ctx.store.set("extra_counter", state.extra_counter + 1)
-        return StopEvent()
+        # Return state as result for testing
+        final_state = await ctx.store.get_state()
+        return StopEvent(result=final_state)
 
 
 @pytest.mark.asyncio
@@ -294,9 +296,7 @@ async def test_consistent_child_state_works() -> None:
 
     result = await test_runner.run()
 
-    ctx = result.ctx
-    assert ctx is not None
-    state = await ctx.store.get_state()
+    state = result.result
 
     # Both base and child fields should be properly set
     assert state.base_field == "modified_by_base_step"
@@ -328,7 +328,9 @@ class SetStateWorkflow(Workflow):
         state = await ctx.store.get_state()
         state.base_field = "modified_base"
         await ctx.store.set_state(state)
-        return StopEvent()
+        # Return state as result for testing
+        final_state = await ctx.store.get_state()
+        return StopEvent(result=final_state)
 
 
 @pytest.mark.asyncio
@@ -344,9 +346,7 @@ async def test_set_state_same_type_preserves_fields() -> None:
 
     result = await test_runner.run()
 
-    ctx = result.ctx
-    assert ctx is not None
-    state = await ctx.store.get_state()
+    state = result.result
 
     # The base field was modified
     assert state.base_field == "modified_base"
@@ -433,7 +433,10 @@ class ChildWorkflowDictState(BaseWorkflowDictState):
     @step
     async def end_step(self, ctx: Context, ev: MiddleEvent) -> StopEvent:
         await ctx.store.set("child_field", "set_by_child")
-        return StopEvent()
+        # Return field values as result for testing
+        base_field = await ctx.store.get("base_field")
+        child_field = await ctx.store.get("child_field")
+        return StopEvent(result={"base_field": base_field, "child_field": child_field})
 
 
 @pytest.mark.asyncio
@@ -450,15 +453,10 @@ async def test_dict_state_allows_flexible_inheritance() -> None:
 
     result = await test_runner.run()
 
-    ctx = result.ctx
-    assert ctx is not None
+    fields = result.result
 
-    # Both fields can be retrieved
-    base_field = await ctx.store.get("base_field")
-    child_field = await ctx.store.get("child_field")
-
-    assert base_field == "set_by_base"
-    assert child_field == "set_by_child"
+    assert fields["base_field"] == "set_by_base"
+    assert fields["child_field"] == "set_by_child"
 
 
 # ============================================================================
@@ -483,7 +481,9 @@ class EditStateWorkflow(Workflow):
         async with ctx.store.edit_state() as state:
             state.base_field = "edited_base"
             # Not touching child fields
-        return StopEvent()
+        # Return state as result for testing
+        final_state = await ctx.store.get_state()
+        return StopEvent(result=final_state)
 
 
 @pytest.mark.asyncio
@@ -499,9 +499,7 @@ async def test_edit_state_preserves_child_fields() -> None:
 
     result = await test_runner.run()
 
-    ctx = result.ctx
-    assert ctx is not None
-    state = await ctx.store.get_state()
+    state = result.result
 
     # Base field was modified
     assert state.base_field == "edited_base"
@@ -549,7 +547,9 @@ class GrandchildWorkflowThreeLevel(ChildWorkflowThreeLevel):
         self, ctx: Context[GrandchildState], ev: MiddleEvent
     ) -> StopEvent:
         await ctx.store.set("grandchild_field", "set_at_level3")
-        return StopEvent()
+        # Return state as result for testing
+        state = await ctx.store.get_state()
+        return StopEvent(result=state)
 
 
 @pytest.mark.asyncio
@@ -566,9 +566,7 @@ async def test_three_level_inheritance_works() -> None:
 
     result = await test_runner.run()
 
-    ctx = result.ctx
-    assert ctx is not None
-    state = await ctx.store.get_state()
+    state = result.result
 
     # Verify state is GrandchildState
     assert isinstance(state, GrandchildState)

packages/llama-index-workflows/tests/test_handler.py

@@ -4,24 +4,32 @@
 from unittest import mock
 
 import pytest
-from workflows.context import Context
 from workflows.errors import WorkflowRuntimeError
 from workflows.handler import WorkflowHandler
+from workflows.workflow import Workflow
+
+from tests.runtime.conftest import MockRunAdapter
+
+
+def _create_mock_handler() -> WorkflowHandler:
+    """Create a WorkflowHandler with MockRunAdapter."""
+    workflow = mock.MagicMock(spec=Workflow)
+    workflow.workflow_name = "TestWorkflow"
+    adapter = MockRunAdapter(run_id="test-run-id")
+    return WorkflowHandler(workflow=workflow, external_adapter=adapter)
 
 
 @pytest.mark.asyncio
 async def test_str() -> None:
-    ctx = mock.MagicMock(spec=Context)
-    h = WorkflowHandler(ctx=ctx)
-    h.set_result([])
-    assert str(h) == "[]"
+    h = _create_mock_handler()
+    # The str representation shows workflow name, run_id, and result
+    assert "TestWorkflow" in str(h)
+    assert "test-run-id" in str(h)
 
 
 @pytest.mark.asyncio
 async def test_stream_events_consume_only_once() -> None:
-    ctx = mock.MagicMock(spec=Context)
-
-    h = WorkflowHandler(ctx=ctx)
+    h = _create_mock_handler()
     h._all_events_consumed = True
 
     with pytest.raises(

packages/llama-index-workflows/tests/test_runtime_integration.py

@@ -0,0 +1,296 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 LlamaIndex Inc.
+"""Integration tests for runtime lifecycle with registering() context manager."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+from workflows import Context, Workflow, step
+from workflows.events import StartEvent, StopEvent
+from workflows.plugins import BasicRuntime, basic_runtime, get_current_runtime
+from workflows.testing import WorkflowTestRunner
+
+
+class SimpleWorkflow(Workflow):
+    """Simple test workflow."""
+
+    @step
+    async def start(self, ev: StartEvent) -> StopEvent:
+        return StopEvent(result="done")
+
+
+class CountingWorkflow(Workflow):
+    """Workflow that counts invocations."""
+
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.run_count = 0
+
+    @step
+    async def start(self, ev: StartEvent) -> StopEvent:
+        self.run_count += 1
+        return StopEvent(result=self.run_count)
+
+
+class StatefulWorkflow(Workflow):
+    """Workflow that preserves state across runs."""
+
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.value = 0
+
+    @step
+    async def start(self, ctx: Context, ev: StartEvent) -> StopEvent:
+        if hasattr(ev, "value"):
+            self.value = getattr(ev, "value")
+        return StopEvent(result=self.value)
+
+
+# Basic Runtime Tests
+
+
+async def test_basic_runtime_no_explicit_launch() -> None:
+    """BasicRuntime works without explicit launch()."""
+    wf = SimpleWorkflow()
+    result = await WorkflowTestRunner(wf).run()
+    assert result.result == "done"
+
+
+async def test_basic_runtime_with_explicit_launch() -> None:
+    """launch()/destroy() are no-ops but don't error for BasicRuntime."""
+    runtime = BasicRuntime()
+    runtime.launch()
+
+    wf = SimpleWorkflow(runtime=runtime)
+    result = await WorkflowTestRunner(wf).run()
+    assert result.result == "done"
+
+    runtime.destroy()
+
+
+async def test_basic_runtime_launch_idempotent() -> None:
+    """Multiple launch() calls are allowed."""
+    runtime = BasicRuntime()
+    runtime.launch()
+    runtime.launch()
+
+    wf = SimpleWorkflow(runtime=runtime)
+    result = await WorkflowTestRunner(wf).run()
+    assert result.result == "done"
+
+
+# Registering Context Manager Tests
+
+
+async def test_registering_multiple_workflows() -> None:
+    """Multiple workflows register in same block."""
+    runtime = BasicRuntime()
+
+    with runtime.registering():
+        wf1 = SimpleWorkflow()
+        wf2 = CountingWorkflow()
+
+    assert wf1.runtime is runtime
+    assert wf2.runtime is runtime
+
+    runtime.launch()
+
+    result1 = await WorkflowTestRunner(wf1).run()
+    result2 = await WorkflowTestRunner(wf2).run()
+
+    assert result1.result == "done"
+    assert result2.result == 1
+
+
+async def test_registering_preserves_workflow_state() -> None:
+    """Workflow state preserved through launch()."""
+    runtime = BasicRuntime()
+
+    with runtime.registering():
+        wf = CountingWorkflow()
+
+    runtime.launch()
+
+    # Run multiple times, state should persist
+    result1 = await WorkflowTestRunner(wf).run()
+    result2 = await WorkflowTestRunner(wf).run()
+
+    assert result1.result == 1
+    assert result2.result == 2
+
+
+async def test_registering_with_exception_still_resets_context() -> None:
+    """Context reset even on exception."""
+    runtime = BasicRuntime()
+
+    with pytest.raises(ValueError):
+        with runtime.registering():
+            raise ValueError("test")
+
+    # Context should be reset
+    assert get_current_runtime() is basic_runtime
+
+
+async def test_mixed_explicit_and_implicit_registration() -> None:
+    """Explicit runtime= overrides context."""
+    context_runtime = BasicRuntime()
+    explicit_runtime = BasicRuntime()
+
+    with context_runtime.registering():
+        wf_implicit = SimpleWorkflow()
+        wf_explicit = SimpleWorkflow(runtime=explicit_runtime)
+
+    assert wf_implicit.runtime is context_runtime
+    assert wf_explicit.runtime is explicit_runtime
+
+    context_runtime.launch()
+    explicit_runtime.launch()
+
+    result1 = await WorkflowTestRunner(wf_implicit).run()
+    result2 = await WorkflowTestRunner(wf_explicit).run()
+
+    assert result1.result == "done"
+    assert result2.result == "done"
+
+
+# Workflow Execution Tests
+
+
+async def test_workflow_can_run_multiple_times() -> None:
+    """Same workflow runs multiple times after launch()."""
+    runtime = BasicRuntime()
+
+    with runtime.registering():
+        wf = CountingWorkflow()
+
+    runtime.launch()
+
+    result1 = await WorkflowTestRunner(wf).run()
+    result2 = await WorkflowTestRunner(wf).run()
+    result3 = await WorkflowTestRunner(wf).run()
+
+    assert result1.result == 1
+    assert result2.result == 2
+    assert result3.result == 3
+
+
+async def test_destroy_allows_reuse() -> None:
+    """Runtime can create new workflows after destroy()."""
+    runtime = BasicRuntime()
+
+    with runtime.registering():
+        wf1 = SimpleWorkflow()
+
+    runtime.launch()
+    result1 = await WorkflowTestRunner(wf1).run()
+
+    runtime.destroy()
+
+    # Can register new workflows after destroy
+    with runtime.registering():
+        wf2 = SimpleWorkflow()
+
+    runtime.launch()
+    result2 = await WorkflowTestRunner(wf2).run()
+
+    assert result1.result == "done"
+    assert result2.result == "done"
+
+
+# Edge Cases
+
+
+async def test_workflow_without_any_runtime_context() -> None:
+    """Falls back to basic_runtime."""
+    wf = SimpleWorkflow()
+    assert wf.runtime is basic_runtime
+
+    result = await WorkflowTestRunner(wf).run()
+    assert result.result == "done"
+
+
+async def test_empty_registering_block() -> None:
+    """Empty block is valid no-op."""
+    runtime = BasicRuntime()
+
+    with runtime.registering():
+        pass
+
+    # Context should be reset
+    assert get_current_runtime() is basic_runtime
+
+
+async def test_registering_yields_runtime() -> None:
+    """Context manager yields the runtime."""
+    runtime = BasicRuntime()
+
+    with runtime.registering() as r:
+        assert r is runtime
+
+    # Context should be reset after exit
+    assert get_current_runtime() is basic_runtime
+
+
+async def test_nested_registering_preserves_workflows() -> None:
+    """Nested registering blocks correctly assign workflows."""
+    outer_runtime = BasicRuntime()
+    inner_runtime = BasicRuntime()
+
+    with outer_runtime.registering():
+        wf_outer = SimpleWorkflow()
+
+        with inner_runtime.registering():
+            wf_inner = SimpleWorkflow()
+
+        wf_outer_again = SimpleWorkflow()
+
+    assert wf_outer.runtime is outer_runtime
+    assert wf_inner.runtime is inner_runtime
+    assert wf_outer_again.runtime is outer_runtime
+
+    outer_runtime.launch()
+    inner_runtime.launch()
+
+    r1 = await WorkflowTestRunner(wf_outer).run()
+    r2 = await WorkflowTestRunner(wf_inner).run()
+    r3 = await WorkflowTestRunner(wf_outer_again).run()
+
+    assert r1.result == "done"
+    assert r2.result == "done"
+    assert r3.result == "done"
+
+
+async def test_workflow_runs_without_registering() -> None:
+    """Workflows created outside registering() use basic_runtime."""
+    wf = SimpleWorkflow()
+    assert wf.runtime is basic_runtime
+
+    result = await WorkflowTestRunner(wf).run()
+    assert result.result == "done"
+
+
+async def test_multiple_concurrent_workflows() -> None:
+    """Multiple workflows can run concurrently."""
+    import asyncio
+
+    runtime = BasicRuntime()
+
+    with runtime.registering():
+        wf1 = SimpleWorkflow()
+        wf2 = SimpleWorkflow()
+        wf3 = SimpleWorkflow()
+
+    runtime.launch()
+
+    # Run all concurrently
+    results = await asyncio.gather(
+        WorkflowTestRunner(wf1).run(),
+        WorkflowTestRunner(wf2).run(),
+        WorkflowTestRunner(wf3).run(),
+    )
+
+    assert results[0].result == "done"
+    assert results[1].result == "done"
+    assert results[2].result == "done"

packages/llama-index-workflows/tests/test_streaming.py

@@ -2,6 +2,7 @@
 # Copyright (c) 2026 LlamaIndex Inc.
 
 import asyncio
+import json
 from typing import AsyncGenerator
 
 import pytest
@@ -145,4 +146,6 @@ async def test_resume_streams() -> None:
     ctx2 = result2.ctx
 
     assert ctx2
-    assert await ctx2.store.get("cur_count") == 2
+    ctx_dict = ctx2.to_dict()
+    # State is serialized as JSON strings under state_data._data for DictState
+    assert json.loads(ctx_dict["state"]["state_data"]["_data"]["cur_count"]) == 2

packages/llama-index-workflows/tests/test_workflow.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 
 import asyncio
 import gc
+import json
 import logging
 import pickle
 import threading
@@ -13,7 +14,10 @@ from typing import Any, Callable, Optional, Union, cast
 from unittest import mock
 
 import pytest
-from llama_index_instrumentation.dispatcher import active_instrument_tags
+from llama_index_instrumentation.dispatcher import (
+    active_instrument_tags,
+    instrument_tags,
+)
 from pydantic import PrivateAttr
 from workflows.context import Context, PickleSerializer
 from workflows.decorators import step
@@ -242,7 +246,6 @@ async def test_workflow_num_workers() -> None:
     # Ensure ctx is serializable
     ctx = r.ctx
     assert ctx
-    assert ctx._broker_run is not None
     ctx.to_dict()
 
 
@@ -266,7 +269,10 @@ async def test_workflow_step_send_event() -> None:
     r = await WorkflowTestRunner(workflow).run()
     assert r.result == "step2"
     ctx = r.ctx
-    replay = ctx._broker_run._runtime.replay()  # type:ignore
+    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
 
 
@@ -284,8 +290,10 @@ async def test_workflow_step_send_event_to_None() -> None:
 
     workflow = StepSendEventToNoneWorkflow(verbose=True)
     result = await WorkflowTestRunner(workflow).run()
-    assert result.ctx._broker_run is not None
-    replay = result.ctx._broker_run._runtime.replay()  # type:ignore
+    from workflows.context.external_context import ExternalContext
+
+    assert isinstance(result.ctx._face, ExternalContext)
+    replay = result.ctx._face._tick_log
     assert TickAddEvent(OneTestEvent()) in replay
 
 
@@ -389,31 +397,29 @@ async def test_workflow_continue_context() -> None:
         @step
         async def step(self, ctx: Context, ev: StartEvent) -> StopEvent:
             cur_number = await ctx.store.get("number", default=0)
-            await ctx.store.set("number", cur_number + 1)
-            return StopEvent(result="Done")
+            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 == "Done"
+    assert r.result == 1
     ctx = r.ctx
     assert ctx
-    assert await ctx.store.get("number") == 1
 
     # second run -- independent from the first
     r = await WorkflowTestRunner(wf).run()
-    assert r.result == "Done"
+    assert r.result == 1
     ctx = r.ctx
     assert ctx
-    assert await ctx.store.get("number") == 1
 
     # third run -- continue from the second run
     handler = wf.run(ctx=ctx)
     result = await handler
     assert handler.ctx
-    assert result == "Done"
-    assert await handler.ctx.store.get("number") == 2
+    assert result == 2
 
 
 @pytest.mark.asyncio
@@ -433,24 +439,22 @@ async def test_workflow_pickle() -> None:
 
     # by default, we can't pickle the LLM/embedding object
     with pytest.raises(ValueError):
-        state_dict = ctx.to_dict()
+        ctx.to_dict()
 
     # if we allow pickle, then we can pickle the LLM/embedding object
     state_dict = ctx.to_dict(serializer=PickleSerializer())
+    # Verify step count via serialized dict (integers are JSON serialized)
+    assert json.loads(state_dict["state"]["state_data"]["_data"]["step"]) == 1
+
+    # Restore and run again to verify deserialization works
     new_ctx = Context.from_dict(wf, state_dict, serializer=PickleSerializer())
-    assert await new_ctx.store.get("step") == 1
-    new_handler = WorkflowHandler(ctx=new_ctx)
-    assert new_handler.ctx
-    assert await new_handler.ctx.store.get("step") == 1
+    r2 = await WorkflowTestRunner(wf).run(ctx=new_ctx)
+    ctx2 = r2.ctx
+    assert ctx2
 
-    # check that the step count is the same
-    cur_step = await ctx.store.get("step")
-    new_step = await new_handler.ctx.store.get("step")
-    assert new_step == cur_step
-    await WorkflowTestRunner(wf).run(ctx=new_handler.ctx)
-
-    # check that the step count is incremented
-    assert await new_handler.ctx.store.get("step") == cur_step + 1
+    # check that the step count is incremented after second run
+    state_dict2 = ctx2.to_dict(serializer=PickleSerializer())
+    assert json.loads(state_dict2["state"]["state_data"]["_data"]["step"]) == 2
 
 
 @pytest.mark.asyncio
@@ -476,6 +480,8 @@ async def test_workflow_context_to_dict() -> None:
             await signal_continue.wait()
             return StopEvent(result="Done")
 
+    from workflows.context.external_context import ExternalContext
+
     workflow = StallableWorkflow()
     try:
         handler = workflow.run()
@@ -491,10 +497,10 @@ async def test_workflow_context_to_dict() -> None:
         signal_continue.set()
         await handler2
     finally:
-        if ctx is not None and ctx._broker_run is not None:
-            await ctx._broker_run.shutdown()
-        if new_ctx is not None and new_ctx._broker_run is not None:
-            await new_ctx._broker_run.shutdown()
+        if ctx is not None and isinstance(ctx._face, ExternalContext):
+            await ctx._face.shutdown()
+        if new_ctx is not None and isinstance(new_ctx._face, ExternalContext):
+            await new_ctx._face.shutdown()
 
 
 class HumanInTheLoopWorkflow(Workflow):
@@ -555,10 +561,10 @@ async def test_human_in_the_loop_with_resume() -> None:
     assert final_result == "42"
 
     # ensure the workflow ran each step once
-    step1_runs = await new_handler.ctx.store.get("step1_runs")  # type:ignore
-    step2_runs = await new_handler.ctx.store.get("step2_runs")  # type:ignore
-    assert step1_runs == 1
-    assert step2_runs == 1
+    assert new_handler.ctx is not None
+    ctx_dict = new_handler.ctx.to_dict()
+    assert json.loads(ctx_dict["state"]["state_data"]["_data"]["step1_runs"]) == 1  # type: ignore[index]
+    assert json.loads(ctx_dict["state"]["state_data"]["_data"]["step2_runs"]) == 1  # type: ignore[index]
 
 
 @pytest.mark.asyncio
@@ -601,10 +607,9 @@ async def test_human_in_the_loop_resume_custom_start_event_inactive_ctx() -> Non
     final_result = await resumed_handler
     assert final_result == "42"
 
-    ask_runs = await resumed_handler.ctx.store.get("ask_runs")  # type: ignore[arg-type]
-    complete_runs = await resumed_handler.ctx.store.get("complete_runs")  # type: ignore[arg-type]
-    assert ask_runs == 1
-    assert complete_runs == 1
+    ctx_dict = resumed_handler.ctx.to_dict()  # type: ignore[union-attr]
+    assert json.loads(ctx_dict["state"]["state_data"]["_data"]["ask_runs"]) == 1
+    assert json.loads(ctx_dict["state"]["state_data"]["_data"]["complete_runs"]) == 1
 
 
 class DummyWorkflowForConcurrentRunsTest(Workflow):
@@ -844,6 +849,28 @@ def test__get_start_event_instance(caplog: Any) -> None:
         d._get_start_event_instance(None, wrong_field="test")
 
 
+def test_run_with_invalid_start_event_raises() -> None:
+    """Passing wrong type to start_event argument should raise ValueError."""
+
+    class SimpleWorkflow(Workflow):
+        @step
+        async def only(self, ev: StartEvent) -> StopEvent:
+            return StopEvent(result="done")
+
+    wf = SimpleWorkflow()
+
+    # Pass a non-StartEvent object
+    with pytest.raises(ValueError, match="must be an instance of 'StartEvent'"):
+        wf.run(start_event="not a StartEvent")  # type: ignore[arg-type]
+
+    # Also test with other invalid types
+    with pytest.raises(ValueError, match="must be an instance of 'StartEvent'"):
+        wf.run(start_event=42)  # type: ignore[arg-type]
+
+    with pytest.raises(ValueError, match="must be an instance of 'StartEvent'"):
+        wf.run(start_event={"topic": "test"})  # type: ignore[arg-type]
+
+
 def test__ensure_start_event_class_multiple_types() -> None:
     class DummyWorkflow(Workflow):
         @step
@@ -1028,24 +1055,28 @@ async def test_workflow_non_picklable_event() -> None:
 
 
 @pytest.mark.asyncio
-async def test_inner_step_can_access_run_id_from_instrument_tags() -> None:
+async def test_inner_step_can_access_run_id_and_others_from_instrument_tags() -> None:
     # container to mutate rather than deal with nonlocal
-    run_id: dict[str, str | None] = {"run_id": None}
+    run_id: dict[str, str | None] = {"run_id": None, "foo": None}
 
     class TagReadingWorkflow(Workflow):
         @step
         async def read_tags(self, ctx: Context, ev: StartEvent) -> StopEvent:
             tags = active_instrument_tags.get()
             run_id["run_id"] = tags.get("run_id")
+            run_id["foo"] = tags.get("foo")
             return StopEvent()
 
     wf = TagReadingWorkflow()
-    handler = wf.run()
+    with instrument_tags({"foo": "bar"}):
+        handler = wf.run()
     await handler
 
     assert handler.run_id is not None
     assert run_id["run_id"] is not None
     assert run_id["run_id"] == handler.run_id
+    assert run_id["foo"] is not None
+    assert run_id["foo"] == "bar"
 
 
 class Par(Event):

packages/llama-index-workflows/tests/test_workflow_naming.py

@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 LlamaIndex Inc.
+"""Tests for Workflow.workflow_name property."""
+
+from workflows import Workflow, step
+from workflows.events import StartEvent, StopEvent
+
+
+class SimpleWorkflow(Workflow):
+    @step
+    async def start(self, ev: StartEvent) -> StopEvent:
+        return StopEvent(result="done")
+
+
+def test_explicit_name() -> None:
+    """Workflow with explicit name uses that name."""
+    wf = SimpleWorkflow(workflow_name="my-custom-name")
+    assert wf.workflow_name == "my-custom-name"
+
+
+def test_default_name_is_module_qualified() -> None:
+    """Workflow without name uses module.qualname."""
+    wf = SimpleWorkflow()
+    # Should be tests.test_workflow_naming.SimpleWorkflow or similar
+    assert wf.workflow_name.endswith("SimpleWorkflow")
+    assert "." in wf.workflow_name  # Has module prefix
+
+
+def test_nested_class_qualname() -> None:
+    """Nested class workflow has correct qualname."""
+
+    class Outer:
+        class InnerWorkflow(Workflow):
+            @step
+            async def start(self, ev: StartEvent) -> StopEvent:
+                return StopEvent(result="done")
+
+    wf = Outer.InnerWorkflow()
+    assert "Outer.InnerWorkflow" in wf.workflow_name
+
+
+def test_function_scoped_workflow_has_locals_in_name() -> None:
+    """Function-scoped workflow includes <locals> in name."""
+
+    def create_workflow() -> Workflow:
+        class LocalWorkflow(Workflow):
+            @step
+            async def start(self, ev: StartEvent) -> StopEvent:
+                return StopEvent(result="done")
+
+        return LocalWorkflow()
+
+    wf = create_workflow()
+    assert "<locals>" in wf.workflow_name
+    assert "LocalWorkflow" in wf.workflow_name
+
+
+def test_explicit_name_overrides_default() -> None:
+    """Explicit name takes precedence over computed name."""
+
+    def create_workflow() -> Workflow:
+        class LocalWorkflow(Workflow):
+            @step
+            async def start(self, ev: StartEvent) -> StopEvent:
+                return StopEvent(result="done")
+
+        return LocalWorkflow(workflow_name="explicit-name")
+
+    wf = create_workflow()
+    assert wf.workflow_name == "explicit-name"
+    assert "<locals>" not in wf.workflow_name

packages/llama-index-workflows/tests/test_workflow_typed_state.py

@@ -1,4 +1,5 @@
 import asyncio
+import json
 from typing import Optional, Union
 
 import pytest
@@ -35,11 +36,15 @@ async def test_typed_state() -> None:
 
     result = await test_runner.run()
 
-    # Check final state
+    # Check final state via to_dict()
     ctx = result.ctx
     assert ctx is not None
-    state = await ctx.store.get_state()
-    assert state.model_dump() == MyState(name="John", age=31).model_dump()
+    ctx_dict = ctx.to_dict()
+    # For typed Pydantic models, state_data is a JSON string with {"__is_pydantic": True, "value": {...}}
+    state_data = json.loads(ctx_dict["state"]["state_data"])
+    assert state_data["__is_pydantic"] is True
+    assert state_data["value"]["name"] == "John"
+    assert state_data["value"]["age"] == 31
 
 
 class SomeState(BaseModel):