test(execution): contract suite, cancellation guarantees, seam docstrings

Robustness pass on the execution seam after using it to integrate an external
executor end-to-end. Four things this pins down:

1. Reusable executor contract suite
   - tests/unit/execution/test_executor_contract.py
   - ExecutorContract with 7 universal seam guarantees: kind shape, execute()
     returns an AsyncIterator, stream ends with exactly one RunComplete,
     instance is reusable across runs, concurrent runs are isolated, consumer
     cancellation does not hang or leak, execute() signature is stable.
   - TestInProcessExecutorContract subclasses and provides fixtures. Downstream
     executors (stepflow, future remote/sandbox) get the same battery by
     subclassing and overriding two fixtures.

2. Cancellation/cleanup correctness
   - tests/unit/execution/test_cancellation.py
   - InProcessExecutor wraps graph.async_start() in try/finally with explicit
     aclose so consumer aclose cascades to the underlying graph generator.
     Previously the inner generator was abandoned and only finalized on a GC
     pass: a real leak for executors holding subprocesses or sockets.
   - Coordinator.run and Coordinator.stream propagate aclose the same way so
     the full chain (consumer -> coordinator -> executor -> graph) cleans up
     deterministically.

3. Documented seam contracts
   - Unit.runtime_options: free-form bag, "_"-prefixed keys reserved for
     executor-internal flags, common in-process keys listed.
   - StepResult.payload: untyped on purpose; each executor defines its own
     event vocabulary; consumers normalize at consumption site.
   - RunComplete.outputs: only the legacy in-process path populates it;
     streaming consumers should collect from StepResult.payload.
   - Executor ABC: lifecycle (shared instance, must be reusable, must tolerate
     concurrent execute() calls and consumer aclose).

4. Coordinator-routed E2E suite
   - tests/unit/execution/test_coordinator_e2e.py
   - Real Graph through Coordinator -> registry -> InProcessExecutor chain.
   - Asserts payload shape, RunComplete never leaks to stream() consumers,
     dispatch routes to configured kind, registry returns same instance.

Full execution suite: 63 passing. flow_executor coordinator test: passing.
This commit is contained in:
ogabrielluiz
2026-05-11 12:09:11 -03:00
parent e320e5df5d
commit 357ea19156
7 changed files with 794 additions and 15 deletions

View File

@ -27,16 +27,32 @@ class InProcessExecutor(Executor):
yield RunComplete(outputs=list(outputs))
return
async for result in graph.async_start(
# Hold a reference to the inner async iterator so we can guarantee
# finalization on consumer cancellation. ``async for`` does NOT call
# ``aclose()`` on its iterator when the loop is interrupted by an exception
# (e.g. ``GeneratorExit`` from this generator being closed). Without the
# explicit try/finally below, the inner generator's ``finally:`` block --
# which is where event managers, connections, or vertex finalizers run --
# would only execute on GC, leaking resources for an unpredictable window.
inner = graph.async_start(
inputs=opts.get("initial_inputs"),
max_iterations=opts.get("max_iterations"),
config=opts.get("config"),
event_manager=opts.get("event_manager"),
reset_output_values=opts.get("reset_output_values", True),
fallback_to_env_vars=opts.get("fallback_to_env_vars", False),
):
yield StepResult(payload=result)
if isinstance(result, Finish):
break
)
try:
async for result in inner:
yield StepResult(payload=result)
if isinstance(result, Finish):
break
yield RunComplete(outputs=[])
yield RunComplete(outputs=[])
finally:
# ``aclose`` is a no-op if the generator already completed. Async
# iterables that aren't generators (test stubs, custom iterators) may
# not implement aclose, so guard with hasattr.
aclose = getattr(inner, "aclose", None)
if aclose is not None:
await aclose()

View File

@ -1,4 +1,21 @@
"""Coordinator: graph in, streaming step results out."""
"""Coordinator: graph in, streaming step results out.
The ``Coordinator`` is the single dispatch point between callers and executors. It
looks up the registered ``Executor`` for the configured ``kind``, partitions the
input graph into ``Unit`` value objects, and forwards each unit's event stream.
Callers pick the view they want:
- ``run()``: full seam stream (``StepResult`` + terminal ``RunComplete``). Use when
you need the terminal signal -- e.g. ``run_to_completion``.
- ``stream()``: payloads only. Drops the terminal ``RunComplete`` silently because
every existing payload-consumer in the codebase (``flow_executor``, the CLI,
``Loop``, ``run/base``) processes events incrementally and has no use for the
terminal envelope.
- ``run_to_completion()``: drains the stream and returns ``RunComplete.outputs``.
Practically useful only with the in-process executor's legacy passthrough; see
the docstring on ``RunComplete.outputs`` for the full semantics.
"""
from __future__ import annotations
@ -14,6 +31,8 @@ if TYPE_CHECKING:
class Coordinator:
"""Routes graph runs through a registered ``Executor`` of the configured kind."""
def __init__(
self,
*,
@ -30,11 +49,30 @@ class Coordinator:
inputs: list[dict[str, Any]],
**runtime_options: Any,
) -> AsyncIterator[StepResult | RunComplete]:
"""Stream the seam-typed event sequence for ``graph``.
Yields the executor's ``StepResult`` items in order followed by the
terminal ``RunComplete``. If you only need payloads, use ``stream()``;
if you only need the final ``outputs`` list, use ``run_to_completion()``.
``runtime_options`` are forwarded as-is into ``Unit.runtime_options``;
consult the active executor's documentation for which keys it honors.
"""
units = identity_partition(graph, inputs=inputs, runtime_options=runtime_options)
executor = self._registry.get(self._executor_kind)
for unit in units:
async for item in executor.execute(unit):
yield item
# Cascade cleanup: if our consumer aclose()s us, we aclose the executor
# stream, which lets InProcessExecutor's try/finally aclose the underlying
# graph generator. Without this, the executor stream is abandoned on
# consumer aclose and only finalizes when CPython gets around to GC.
inner = executor.execute(unit)
try:
async for item in inner:
yield item
finally:
aclose = getattr(inner, "aclose", None)
if aclose is not None:
await aclose()
async def run_to_completion(
self,
@ -43,6 +81,16 @@ class Coordinator:
inputs: list[dict[str, Any]],
**runtime_options: Any,
) -> list[Any]:
"""Drain the stream and return ``RunComplete.outputs``.
Raises ``RuntimeError`` if the executor's stream ends without a terminal
``RunComplete`` -- a contract violation that every executor must avoid.
See ``RunComplete.outputs`` for the per-executor semantics of the returned
list. In particular, the in-process executor's streaming path returns an
empty list here; callers that want final outputs from the streaming path
should use ``stream()`` and collect what they need from payloads.
"""
async for item in self.run(graph, inputs=inputs, **runtime_options):
if isinstance(item, RunComplete):
return item.outputs
@ -56,6 +104,17 @@ class Coordinator:
inputs: list[dict[str, Any]] | None = None,
**runtime_options: Any,
) -> AsyncIterator[Any]:
async for item in self.run(graph, inputs=inputs or [], **runtime_options):
if isinstance(item, StepResult):
yield item.payload
"""Yield ``StepResult.payload`` values only.
The terminal ``RunComplete`` is intentionally NOT yielded -- this helper
exists for consumers that process events incrementally and have no use for
the end-of-run envelope. If you need to know when the run ends, use
``run()`` directly.
"""
inner = self.run(graph, inputs=inputs or [], **runtime_options)
try:
async for item in inner:
if isinstance(item, StepResult):
yield item.payload
finally:
await inner.aclose()

View File

@ -1,4 +1,23 @@
"""Executor abstract base class."""
"""Executor abstract base class.
An ``Executor`` is the pluggable backend that turns a ``Unit`` into a stream of
``StepResult`` events terminated by a single ``RunComplete``. The ``Coordinator``
holds at most one instance per ``kind`` and dispatches every ``run()`` against it,
so implementations must observe the lifecycle contract documented below.
Authoring a new executor:
1. Subclass ``Executor`` and set ``kind`` to a unique, stable string.
2. Implement ``execute`` as an ``async def`` generator. Yield zero or more
``StepResult`` items, then exactly one ``RunComplete``.
3. Register through one of three paths (all equivalent at runtime):
- Explicit: ``get_service(ServiceType.EXECUTOR_SERVICE).register(MyExecutor())``
- Plugin: ``[project.entry-points."lfx.executors"] my_kind = "pkg.mod:MyExecutor"``
- Service replacement: ``@register_service(ServiceType.EXECUTOR_SERVICE)``
4. Validate with the shared contract suite: subclass
``tests.unit.execution.test_executor_contract.ExecutorContract`` in your tests
and provide the two fixtures. All universal seam guarantees are checked there.
"""
from __future__ import annotations
@ -12,8 +31,46 @@ if TYPE_CHECKING:
class Executor(ABC):
"""Pluggable execution backend for the seam.
Class attributes:
kind: Unique stable identifier used by ``ExecutorRegistry`` and by
``settings.executor_kind``. Two executors with the same ``kind`` cannot
coexist in a single registry; entry-point discovery refuses collisions.
Lifecycle contract:
- A single instance is shared across all runs of its ``kind``. The same
instance may receive concurrent ``execute()`` calls (e.g. parallel API
requests, ``Loop`` subgraphs). Implementations MUST be safe under
concurrent calls and MUST NOT carry per-run state on ``self``.
- ``execute()`` is called once per ``Unit``. Each call MUST return a fresh
async iterator; reusing/caching a generator across calls is a bug.
- The stream MUST emit zero or more ``StepResult`` followed by exactly one
terminal ``RunComplete``. ``Coordinator.run_to_completion`` raises if the
terminal item is missing.
- Exceptions raised inside the generator propagate to the consumer. The
generator's finalizer runs as usual, so any ``async with`` / cleanup code
will execute. Consumers MAY drop the iterator mid-stream (``aclose``);
implementations MUST tolerate this without hanging or leaking external
resources (subprocesses, sockets, locks, file handles).
- ``execute()`` MUST be tolerant of arbitrary keys in ``unit.runtime_options``:
ignore keys the executor doesn't recognize. Keys prefixed with ``_`` are
reserved for executor-internal flags; non-owning executors must ignore them.
"""
kind: ClassVar[str]
@abstractmethod
def execute(self, unit: Unit) -> AsyncIterator[StepResult | RunComplete]:
"""Yield StepResult items, end with a RunComplete."""
"""Run ``unit`` and yield the seam-typed event stream.
Args:
unit: The work item. ``unit.graph`` is executor-defined (an in-memory
``Graph`` for in-process; whatever shape a remote/sandboxed executor
expects). ``unit.inputs`` and ``unit.runtime_options`` carry per-run
state.
Yields:
``StepResult`` items for mid-run events (executor-defined payloads),
terminated by exactly one ``RunComplete``.
"""

View File

@ -1,4 +1,11 @@
"""Value objects for the execution layer."""
"""Value objects for the execution layer.
These three dataclasses are the entire vocabulary of the seam: a ``Unit`` flows from
``Coordinator`` into ``Executor.execute()``, which yields a stream of ``StepResult``
items terminated by a single ``RunComplete``. Every consumer downstream of the
coordinator (``run_to_completion``, ``stream``, ``flow_executor``, the CLI, ``Loop``
subgraphs) reads only these types, so changes here ripple widely. Keep them small.
"""
from __future__ import annotations
@ -8,6 +15,30 @@ from typing import Any
@dataclass
class Unit:
"""A self-contained slice of work handed to an ``Executor``.
Attributes:
graph: The graph (or graph-shaped object) the executor should run. Executors
decide what they can accept here -- the in-process executor wants a real
``lfx.Graph``; a remote executor may want a serialized form. There is no
seam-level constraint other than "the executor knows what to do with it."
inputs: Per-run input dicts. Shape is consumer-defined.
runtime_options: A free-form bag of run-scoped options. Stringly typed on
purpose: the seam can't enumerate every option every executor might want
(event managers, session IDs, fallback flags, executor-specific tokens),
and pinning a schema here would force every change into this module.
Conventions worth knowing:
- Keys starting with ``_`` are reserved for executor-internal flags
(e.g. ``_use_arun_legacy`` selects the legacy passthrough in the
in-process executor). Other executors MUST ignore unknown keys.
- Common keys recognized by the in-process executor:
``initial_inputs``, ``max_iterations``, ``config``, ``event_manager``,
``reset_output_values``, ``fallback_to_env_vars``, ``session_id``.
- Other executors are free to define their own keys; document them in
the executor module.
"""
graph: Any
inputs: list[dict[str, Any]] = field(default_factory=list)
runtime_options: dict[str, Any] = field(default_factory=dict)
@ -15,9 +46,54 @@ class Unit:
@dataclass
class StepResult:
"""A single mid-run event from an executor.
``payload`` is intentionally untyped at the seam level. Each executor decides
what mid-run events look like for its backend:
- ``InProcessExecutor`` yields ``Vertex`` / ``Finish`` instances from
``Graph.async_start``.
- A remote executor (e.g. stepflow) might yield protobuf status events.
- A sandboxed executor might yield structured ``StepStarted`` / ``StepCompleted``
records.
Consumers that need a uniform shape across executors should normalize at the
consumer site (or wrap an executor with an adapter), not push a structured event
type down through the seam. Doing the latter would force every executor to
translate its native event vocabulary into a synthetic one and would mask
information the native vocabulary carries.
``Coordinator.stream()`` unwraps ``StepResult.payload`` for callers that want a
raw event stream; ``Coordinator.run()`` yields the full ``StepResult``/``
RunComplete`` union for callers that want the seam shape.
"""
payload: Any
@dataclass
class RunComplete:
"""Terminal item in an executor stream.
``outputs`` is populated only by executors / paths that have a meaningful notion
of "final outputs of the run." In particular:
- The in-process executor's *legacy* path (selected via
``runtime_options["_use_arun_legacy"]``) populates ``outputs`` with the
``list[RunOutputs]`` returned by ``Graph._arun_legacy``. ``Graph.arun`` reads
this via ``Coordinator.run_to_completion``.
- The in-process executor's *streaming* path yields ``RunComplete(outputs=[])``;
consumers of the streaming path are expected to collect what they need from
``StepResult.payload`` events along the way (this is what ``flow_executor``,
the CLI, ``Loop``, and ``run/base`` already do via ``Coordinator.stream``).
- Other executors MAY populate ``outputs`` with whatever terminal-status event
their backend provides (e.g. a stepflow ``RunCompletedEvent``).
Bottom line: the *only* universal contract on ``outputs`` is that it is a list.
If you need final values across all executors, collect them from the stream;
if you specifically need the legacy ``list[RunOutputs]`` shape, use
``Coordinator.run_to_completion`` with ``_use_arun_legacy=True`` against the
in-process executor.
"""
outputs: list[Any]

View File

@ -0,0 +1,187 @@
"""Cancellation and cleanup tests for the execution seam.
The contract suite covers the bare-minimum ``aclose()`` shape via
``test_consumer_cancellation_does_not_hang_or_leak``. This module exercises the
stricter behaviour the seam needs in production: explicit cleanup at the top of
the consumer chain MUST cascade all the way down to the underlying graph iterator.
What we deliberately do NOT test here:
- Implicit GC-based finalization (``async for ... break`` without ``aclosing``,
``task.cancel()`` then immediate assert). CPython finalizes abandoned async
generators on a GC pass; asyncio runs the registered finalizer on event-loop
shutdown. Neither is observable inside a single test, and asserting either
would test CPython, not us. Consumers that need deterministic cleanup MUST
use ``contextlib.aclosing`` or call ``aclose()`` explicitly.
"""
from __future__ import annotations
import asyncio
from contextlib import aclosing
import pytest
from lfx.execution.backends.in_process import InProcessExecutor
from lfx.execution.coordinator import Coordinator
from lfx.execution.registry import ExecutorRegistry
from lfx.execution.types import StepResult, Unit
class _InstrumentedGraph:
"""Minimal graph stub whose ``async_start`` runs forever until closed.
We track entry, items yielded, and finalizer invocation so tests can assert
the iterator was properly closed rather than abandoned.
"""
def __init__(self) -> None:
self.entered: bool = False
self.finalized: bool = False
self.items_yielded: int = 0
async def async_start(self, **_kwargs):
self.entered = True
try:
for i in range(10_000):
self.items_yielded += 1
yield f"item-{i}"
await asyncio.sleep(0)
finally:
self.finalized = True
def _unit(graph: _InstrumentedGraph) -> Unit:
return Unit(graph=graph, inputs=[], runtime_options={})
@pytest.mark.asyncio
async def test_aclose_on_executor_finalizes_underlying_graph_iteration():
"""``aclose()`` on the executor's generator MUST run the underlying graph's finalizer.
Without the explicit try/finally inside ``InProcessExecutor.execute``, the inner
generator is abandoned and the graph's ``finally:`` block (event managers,
connections, vertex teardown) only runs on a future GC pass.
"""
graph = _InstrumentedGraph()
executor = InProcessExecutor()
stream = executor.execute(_unit(graph))
first = await anext(stream)
assert isinstance(first, StepResult)
assert graph.entered
assert not graph.finalized
await stream.aclose()
assert graph.finalized, "underlying async generator was not finalized on aclose()"
@pytest.mark.asyncio
async def test_aclosing_helper_chains_cleanup_through_executor():
"""``contextlib.aclosing`` MUST trigger the same cascade as explicit ``aclose``."""
graph = _InstrumentedGraph()
executor = InProcessExecutor()
async with aclosing(executor.execute(_unit(graph))) as stream:
async for _ in stream:
break
assert graph.finalized
@pytest.mark.asyncio
async def test_coordinator_run_cascades_aclose_to_executor_and_graph():
"""Cleanup must propagate through every seam layer.
Consumer aclose -> Coordinator.run -> Executor.execute -> graph.async_start.
A regression at any layer would leave the graph generator hanging.
"""
graph = _InstrumentedGraph()
registry = ExecutorRegistry()
registry.register(InProcessExecutor())
coordinator = Coordinator(registry=registry)
async with aclosing(coordinator.run(graph, inputs=[])) as stream:
first = await anext(stream)
assert first is not None
assert graph.finalized, "Coordinator.run did not cascade aclose to the graph"
@pytest.mark.asyncio
async def test_coordinator_stream_cascades_aclose_to_graph():
"""``Coordinator.stream`` must cascade cleanup the same way ``Coordinator.run`` does."""
graph = _InstrumentedGraph()
registry = ExecutorRegistry()
registry.register(InProcessExecutor())
coordinator = Coordinator(registry=registry)
async with aclosing(coordinator.stream(graph)) as stream:
async for _ in stream:
break
assert graph.finalized
@pytest.mark.asyncio
async def test_cleanup_bounded_by_timeout():
"""Finalization must complete within a bounded window even on a slow CI box.
A regression that swaps in a non-cancellable resource (blocking subprocess
wait, sync lock) would hang here rather than slow.
"""
graph = _InstrumentedGraph()
executor = InProcessExecutor()
async def drive():
gen = executor.execute(_unit(graph))
await anext(gen)
await gen.aclose()
await asyncio.wait_for(drive(), timeout=2.0)
assert graph.finalized
@pytest.mark.asyncio
async def test_exception_in_consumer_still_finalizes_graph():
"""An exception inside the consumer MUST still trigger cleanup via ``aclosing``."""
graph = _InstrumentedGraph()
executor = InProcessExecutor()
class _ConsumerError(RuntimeError):
pass
async def consume_then_raise() -> None:
async with aclosing(executor.execute(_unit(graph))) as stream:
async for _ in stream:
msg = "downstream bug"
raise _ConsumerError(msg)
with pytest.raises(_ConsumerError):
await consume_then_raise()
assert graph.finalized
@pytest.mark.asyncio
async def test_concurrent_consumers_each_clean_up_independently():
"""Concurrent runs on one executor instance must clean up independently.
Regression guard for any future change that introduces shared cleanup state
on the executor (which would couple unrelated consumers).
"""
graph_a = _InstrumentedGraph()
graph_b = _InstrumentedGraph()
executor = InProcessExecutor()
async def drive(graph: _InstrumentedGraph, items_before_stop: int) -> None:
async with aclosing(executor.execute(_unit(graph))) as stream:
count = 0
async for _ in stream:
count += 1
if count >= items_before_stop:
break
await asyncio.gather(drive(graph_a, 1), drive(graph_b, 3))
assert graph_a.finalized
assert graph_b.finalized

View File

@ -0,0 +1,156 @@
"""End-to-end seam tests: a real Graph driven through Coordinator.
Tests in ``test_coordinator.py`` check Coordinator behaviour at a behavioural level
but defer most of the heavy lifting to the executor or graph layer. This module
verifies the seam contract by running a real ``lfx.Graph`` through the full
``Coordinator -> ExecutorRegistry -> InProcessExecutor -> graph.async_start``
chain and asserting the event stream looks like what consumers actually receive:
- ``stream()`` payloads are vertex-shaped (not the seam envelopes).
- ``run()`` yields ``StepResult`` items in order then exactly one ``RunComplete``.
- ``stream()`` never leaks ``RunComplete`` through the payload channel.
- Switching the registered executor changes what runs (the dispatch is real, not
hard-wired to ``InProcessExecutor``).
- Lifecycle: registry get() returns the same instance for repeated calls.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
import pytest
from lfx.execution.backends.in_process import InProcessExecutor
from lfx.execution.coordinator import Coordinator
from lfx.execution.executor import Executor
from lfx.execution.registry import ExecutorRegistry
from lfx.execution.types import RunComplete, StepResult, Unit
if TYPE_CHECKING:
from collections.abc import AsyncIterator
def _make_coordinator() -> Coordinator:
registry = ExecutorRegistry()
registry.register(InProcessExecutor())
return Coordinator(registry=registry)
@pytest.mark.asyncio
async def test_run_yields_step_results_then_terminal_run_complete(simple_graph):
coordinator = _make_coordinator()
items = [item async for item in coordinator.run(simple_graph, inputs=[{"input_value": "hi"}])]
assert items, "real Graph drove zero events through Coordinator"
assert isinstance(items[-1], RunComplete), "stream must end with RunComplete"
assert sum(isinstance(i, RunComplete) for i in items) == 1, "exactly one RunComplete expected"
assert all(isinstance(i, StepResult) for i in items[:-1]), "pre-terminal items must all be StepResult"
@pytest.mark.asyncio
async def test_stream_yields_payloads_only_no_run_complete_leak(simple_graph):
"""``Coordinator.stream`` MUST NOT yield ``RunComplete`` instances.
It is the payload-only helper; the terminal envelope from the executor must
be stripped before reaching the consumer.
"""
coordinator = _make_coordinator()
payloads = [p async for p in coordinator.stream(simple_graph, initial_inputs=None)]
assert payloads, "stream produced no payloads from a real graph"
assert not any(isinstance(p, RunComplete) for p in payloads), (
"Coordinator.stream() leaked RunComplete into the payload channel"
)
assert not any(isinstance(p, StepResult) for p in payloads), (
"Coordinator.stream() leaked StepResult envelopes; payloads should be unwrapped"
)
@pytest.mark.asyncio
async def test_payload_shape_is_vertex_or_finish_for_in_process(simple_graph):
"""Pin the in-process payload contract.
Changes to graph.async_start can't silently rename or restructure the
events consumers depend on without breaking this test.
"""
from lfx.graph.graph.constants import Finish
coordinator = _make_coordinator()
payloads = [p async for p in coordinator.stream(simple_graph, initial_inputs=None)]
vertex_ids = [getattr(p, "vertex", None) and p.vertex.id for p in payloads]
assert "chat_input" in vertex_ids
assert "chat_output" in vertex_ids
assert any(isinstance(p, Finish) for p in payloads), "expected a terminal Finish payload"
@pytest.mark.asyncio
async def test_dispatch_routes_to_configured_kind(simple_graph):
"""Changing ``executor_kind`` changes where execution lands.
Proves the registry lookup is the real dispatch, not a hard-coded reference.
"""
class _Marker(Executor):
kind = "marker"
ran: bool = False
async def execute(self, unit: Unit) -> AsyncIterator[StepResult | RunComplete]: # noqa: ARG002
type(self).ran = True
yield StepResult(payload="sentinel")
yield RunComplete(outputs=["sentinel"])
registry = ExecutorRegistry()
registry.register(InProcessExecutor())
registry.register(_Marker())
coordinator = Coordinator(registry=registry, executor_kind="marker")
items = [item async for item in coordinator.run(simple_graph, inputs=[])]
assert _Marker.ran is True
assert items[0].payload == "sentinel"
assert isinstance(items[-1], RunComplete)
assert items[-1].outputs == ["sentinel"]
@pytest.mark.asyncio
async def test_run_to_completion_returns_outputs_from_legacy_path(simple_graph):
"""``run_to_completion`` with the legacy flag returns populated outputs.
This is the only path through the seam where ``RunComplete.outputs`` carries
a meaningful value (see ``RunComplete`` docstring).
"""
coordinator = _make_coordinator()
outputs = await coordinator.run_to_completion(
simple_graph,
inputs=[{"input_value": "hi"}],
_use_arun_legacy=True,
)
assert isinstance(outputs, list)
@pytest.mark.asyncio
async def test_registry_returns_same_executor_instance_across_runs(simple_graph):
"""The seam shares one executor instance per kind across all runs.
Two sequential ``coordinator.run`` calls MUST resolve to the same instance;
otherwise executor authors couldn't rely on instance-level setup.
"""
instances: list[Executor] = []
real_inprocess = InProcessExecutor()
class _Recorder(Executor):
kind = "recorder"
async def execute(self, unit: Unit) -> AsyncIterator[StepResult | RunComplete]:
instances.append(self)
async for item in real_inprocess.execute(unit):
yield item
registry = ExecutorRegistry()
registry.register(_Recorder())
coordinator = Coordinator(registry=registry, executor_kind="recorder")
[_ async for _ in coordinator.run(simple_graph, inputs=[{"input_value": "a"}])]
[_ async for _ in coordinator.run(simple_graph, inputs=[{"input_value": "b"}])]
assert len(instances) == 2
assert instances[0] is instances[1], "registry must return the same instance per kind"

View File

@ -0,0 +1,228 @@
"""Reusable contract test suite for any ``lfx.execution.Executor`` implementation.
The seam has a small but load-bearing set of behaviors that any executor MUST satisfy
to be safely swappable behind ``Coordinator``. This module spells them out as a
parametrizable suite so:
- The built-in ``InProcessExecutor`` is checked here against the same suite the
external executors (stepflow, future remote/sandbox, etc.) will be checked against.
- Authors of new executors can subclass ``ExecutorContract`` in their own project,
override the two fixtures, and inherit the full battery of behavioural assertions
rather than re-deriving them from the docstring of the ABC.
The contract intentionally does NOT pin executor-specific event shapes -- only the
universal seam guarantees. If a future executor needs to add stricter checks for its
own payloads, those go in an executor-specific test file, not here.
"""
from __future__ import annotations
import asyncio
import inspect
from collections.abc import AsyncIterator, Callable
from typing import TYPE_CHECKING
import pytest
from lfx.execution.types import RunComplete, StepResult, Unit
if TYPE_CHECKING:
from lfx.execution.executor import Executor
class ExecutorContract:
"""Base class of behavioural assertions every Executor must satisfy.
Subclasses MUST override two fixtures:
- ``contract_executor``: returns a freshly constructed ``Executor`` instance.
- ``contract_unit_factory``: returns a callable ``() -> Unit`` that produces a
fresh, runnable ``Unit`` each call. The unit should drive a minimal but real
execution -- not a stub -- so the test actually exercises the executor's
end-to-end path.
Tests are designed to be independent: each test creates its own unit through
the factory so subclasses don't have to worry about state bleed between cases.
"""
@pytest.fixture
def contract_executor(self) -> Executor:
msg = "Subclasses must override the contract_executor fixture"
raise NotImplementedError(msg)
@pytest.fixture
def contract_unit_factory(self) -> Callable[[], Unit]:
msg = "Subclasses must override the contract_unit_factory fixture"
raise NotImplementedError(msg)
# --- Universal seam guarantees ----------------------------------------------------
def test_kind_is_a_non_empty_string(self, contract_executor: Executor) -> None:
"""``kind`` is a ClassVar on the ABC; it must be set and non-empty."""
assert isinstance(contract_executor.kind, str)
assert contract_executor.kind, "Executor.kind must be a non-empty string"
def test_execute_returns_async_iterator(
self,
contract_executor: Executor,
contract_unit_factory: Callable[[], Unit],
) -> None:
"""``execute(unit)`` MUST return an async iterator.
The ABC types it as such; consumers (Coordinator.run, Coordinator.stream,
run_to_completion) rely on it.
"""
stream = contract_executor.execute(contract_unit_factory())
assert isinstance(stream, AsyncIterator), (
f"{type(contract_executor).__name__}.execute() returned a "
f"{type(stream).__name__}, expected AsyncIterator. Consumers can't iterate "
"this. Common cause: forgetting to make the method an `async def` "
"generator."
)
@pytest.mark.asyncio
async def test_stream_ends_with_exactly_one_run_complete(
self,
contract_executor: Executor,
contract_unit_factory: Callable[[], Unit],
) -> None:
"""Streams emit zero-or-more ``StepResult`` items terminated by one ``RunComplete``."""
items = [item async for item in contract_executor.execute(contract_unit_factory())]
assert items, "executor yielded nothing -- a RunComplete is mandatory"
assert isinstance(items[-1], RunComplete), (
f"final item was {type(items[-1]).__name__}; the seam requires a terminal "
"RunComplete so run_to_completion() and downstream collectors can detect end-of-run."
)
run_completes = [i for i in items if isinstance(i, RunComplete)]
assert len(run_completes) == 1, (
f"executor yielded {len(run_completes)} RunComplete items; the seam allows "
"at most one (and requires exactly one)."
)
non_terminal = items[:-1]
bad = [i for i in non_terminal if not isinstance(i, StepResult)]
assert not bad, (
f"all pre-terminal items must be StepResult; got: {[type(i).__name__ for i in bad]}. "
"Coordinator.stream() relies on this to filter; mixing other types makes "
"the stream untyped at the seam boundary."
)
@pytest.mark.asyncio
async def test_executor_is_reusable_across_runs(
self,
contract_executor: Executor,
contract_unit_factory: Callable[[], Unit],
) -> None:
"""An Executor MUST be reusable: state that carries between runs is a bug.
The registry returns the same instance for every ``coordinator.run()`` call,
so any internal state pinned to ``self`` will silently corrupt later runs.
"""
first = [item async for item in contract_executor.execute(contract_unit_factory())]
second = [item async for item in contract_executor.execute(contract_unit_factory())]
assert isinstance(first[-1], RunComplete)
assert isinstance(second[-1], RunComplete)
# Same shape across runs proves no carry-over short-circuit. We don't assert
# equality of payloads because some executors include timestamps / run IDs.
assert len(first) >= 1
assert len(second) >= 1
@pytest.mark.asyncio
async def test_concurrent_runs_on_one_instance_are_isolated(
self,
contract_executor: Executor,
contract_unit_factory: Callable[[], Unit],
) -> None:
"""Concurrent runs on one executor instance MUST be isolated.
Coordinator may launch concurrent runs (Loop subgraphs, parallel API
requests). An executor that crosses state between them is broken.
"""
async def drain() -> list[object]:
return [item async for item in contract_executor.execute(contract_unit_factory())]
a, b = await asyncio.gather(drain(), drain())
assert isinstance(a[-1], RunComplete)
assert isinstance(b[-1], RunComplete)
@pytest.mark.asyncio
async def test_consumer_cancellation_does_not_hang_or_leak(
self,
contract_executor: Executor,
contract_unit_factory: Callable[[], Unit],
) -> None:
"""Consumer cancellation via ``aclose()`` MUST finalize within a bounded window.
The executor's async generator is closed by the consumer; any internal
context managers, subprocesses, or sockets MUST be released. A strict
timeout surfaces regressions as test failures rather than CI hangs.
"""
gen = contract_executor.execute(contract_unit_factory())
try:
first = await asyncio.wait_for(anext(gen), timeout=10.0)
except StopAsyncIteration:
# Executor produced no items at all -- still a valid stream (just an empty one).
return
assert isinstance(first, (StepResult, RunComplete)), (
f"first yielded item was {type(first).__name__}; seam requires StepResult or RunComplete"
)
# Drop the iterator. If the executor holds a subprocess / network / lock, this
# is the moment of truth: aclose() must finalize without blocking forever.
await asyncio.wait_for(gen.aclose(), timeout=10.0)
@pytest.mark.asyncio
async def test_execute_signature_takes_exactly_one_unit(self, contract_executor: Executor) -> None:
"""Pin the ABC signature so executors don't drift away from a single ``Unit`` parameter.
``execute(*args, **kwargs)`` shapes bypass the ``Unit`` value object and
make the seam opaque to introspection.
"""
sig = inspect.signature(contract_executor.execute)
params = [p for p in sig.parameters.values() if p.name != "self"]
assert len(params) == 1, (
f"{type(contract_executor).__name__}.execute must accept exactly one positional "
f"Unit argument; got signature: {sig}"
)
class TestInProcessExecutorContract(ExecutorContract):
"""Run the contract suite against the built-in InProcessExecutor.
This is the reference implementation: if the contract suite is wrong, this is
where it will fail first. Stepflow / remote / sandbox executors are exercised
against the same suite in their own test files.
"""
@pytest.fixture
def contract_executor(self) -> Executor:
from lfx.execution.backends.in_process import InProcessExecutor
return InProcessExecutor()
@pytest.fixture
def contract_unit_factory(self, simple_graph) -> Callable[[], Unit]: # noqa: ARG002
# ``simple_graph`` is a pytest fixture from conftest.py. We don't capture it
# by closure across runs because Graph instances accumulate state; instead we
# rebuild on every call so each unit is fresh.
from lfx.components.input_output import ChatInput, ChatOutput
from lfx.graph import Graph
from lfx.schema.schema import InputValueRequest
def factory() -> Unit:
ci = ChatInput(_id="chat_input")
ci.set(should_store_message=False)
co = ChatOutput(input_value="test", _id="chat_output")
co.set(sender_name=ci.message_response)
return Unit(
graph=Graph(ci, co),
inputs=[],
runtime_options={"initial_inputs": InputValueRequest(input_value="hi")},
)
return factory