dex-python-sdk 0.2.5__tar.gz → 0.2.6__tar.gz

{dex_python_sdk-0.2.5 → dex_python_sdk-0.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dex-python-sdk
-Version: 0.2.5
+Version: 0.2.6
 Requires-Dist: grpcio>=1.83.0
 Requires-Dist: grpcio-status>=1.83.0
 Requires-Dist: protobuf>=7.35.1
@@ -29,6 +29,7 @@ the shared Rust Core is used only for BlobCache.
 
 ```python
 from datetime import timedelta
+from typing import Generator
 
 import dex
 
@@ -46,8 +47,9 @@ class Run(dex.Step[str]):
 
     def execute(
         self, context: dex.Context, input: str
-    ) -> dex.StepDecision:
-        progress.write(context, "running")
+    ) -> Generator[dex.StepOutput, None, dex.StepDecision]:
+        yield progress.write(context, "running")
+        yield dex.heartbeat({"phase": "running"})
         return dex.graceful_complete(input)
 
 class CounterFlow(dex.Flow[str]):
@@ -78,23 +80,6 @@ explicit codec only for a custom encoding or a type Registry cannot derive.
 `PersistenceSchema.of(...)` accepts attributes, channels, and streams together and
 partitions them by definition type.
 
-Streams provide best-effort resumable progress messages. Their approximate byte
-budget is shared by all instances of the owning Flow type. Client keys cannot
-contain `#`; Step writes generate `runID#stepExecutionID` and allow one write per
-Stream per invocation.
-
-```python
-client.write_stream(flow_id, progress, "frontend/1", "starting")
-message = client.read_stream(
-    flow_id, progress, resume_token, timeout=timedelta(seconds=30)
-)
-resume_token = message.resume_token
-```
-
-Async Step handlers await **Stream.write**, and **AsyncClient** exposes matching
-async methods. Reads return the decoded value, resume token, creation time, and
-idempotency key.
-
 `Worker` and `AsyncWorker` synchronize all registered Indexed Attributes with
 Dex Server before opening their listener. Existing indexes return immediately;
 failure or the default two-minute deadline aborts startup. An indexed
@@ -144,9 +129,10 @@ Applications implement two generic interfaces from [`dex`](dex/):
   binds the Flow input to the starting Step input. Use `StepList.empty()` when
   a Flow has no Steps.
 - `Step[INPUT]` implements `execute` and optionally `wait_for`. The default
-  Worker path requires synchronous handlers. With `AsyncWorker` and
-  `Registry(..., allow_async_handlers=True)`, handlers may be `async def` and
-  `await` an `AsyncClient`.
+  Worker accepts ordinary synchronous handlers and generator handlers. A generator
+  yields `StepOutput` progress frames and returns its final `Wait` or `StepDecision`.
+  With `AsyncWorker` and `Registry(..., allow_async_handlers=True)`, Step
+  coroutines use `AsyncContext`; async RPCs keep `Context`.
 
 `StepOptions.wait_for_method_timeout` and `execute_method_timeout` bound the
 two handler calls. Timer and channel conditions determine how long a Step waits.
@@ -157,6 +143,58 @@ attempts, total duration, and 1-based attempt numbers. Fallback starts
 immediately; later regular retries continue the backoff sequence at the
 cumulative attempt.
 
+The default Step durability is synchronous. A Flow configuration can select
+asynchronous durability, and a Step method override has highest precedence. The
+default retry total duration is four hours. Regular attempts default to a two-hour
+method timeout and one-minute heartbeat timeout; an explicit heartbeat timeout must
+meet the server minimum, which defaults to ten seconds. Asynchronous durability
+first allows at most three local attempts in seven seconds. The local phase ignores
+method timeouts and heartbeat frames before falling back to a regular activity.
+
+### Step progress and heartbeat recovery
+
+A synchronous handler yields every heartbeat and Stream write. The generator return
+value is the only final result:
+
+```python
+def execute(
+    self, context: dex.Context, input: str
+) -> Generator[dex.StepOutput, None, dex.StepDecision]:
+    yield dex.heartbeat({"offset": 10})
+    yield progress.write(context, "processed 10 items")
+    return dex.graceful_complete(input)
+```
+
+An asynchronous handler keeps its normal coroutine return. Stream writes enqueue
+without waiting for Stream Store acknowledgement; heartbeat waits only for the
+Worker output queue:
+
+```python
+async def execute(
+    self, context: dex.AsyncContext, input: str
+) -> dex.StepDecision:
+    progress.write(context, "started")
+    await context.heartbeat({"offset": 10})
+    return dex.graceful_complete(input)
+```
+
+Call `heartbeat()` or `await context.heartbeat()` without a value to clear previous
+details. Passing Python `None` persists a present null Value. On a later regular
+attempt, use `context.has_last_heartbeat_value()` before decoding with
+`context.get_last_heartbeat_value(ExpectedType)`. A Stream frame is also an implicit
+heartbeat, but it preserves the latest explicit heartbeat state.
+
+A Step may write the same Stream any number of times. Dex assigns
+`#<stepExecutionID>` as each Step message's `StreamMessage.source`. Client writes
+provide their own non-empty source; duplicate sources and `#` are allowed and every
+write appends:
+
+```python
+client.write_stream(flow_id, progress, "frontend#preview", "rendering")
+message = client.read_stream(flow_id, progress)
+print(message.source)
+```
+
 ### Canceling Step executions
 
 A successful Step can cancel queued or active executions while continuing with
@@ -181,14 +219,6 @@ already-canceled, and absent targets are no-ops. Next Steps created by the same
 decision are outside the snapshot. Dex immediately applies the next or close
 action; late decisions, writes, retries, and recovery Steps are discarded.
 
-Set `StepOptions.heartbeat_timeout` on long-running regular Steps so
-cancellation reaches the Worker promptly. It applies to `wait_for` and
-`execute`; local activities ignore it, while an ASYNC fallback uses it. `None`
-and zero disable heartbeats, and positive values must be whole seconds in the
-signed int32 range. `AsyncWorker` cancels the handler's asyncio task. A handler
-may catch `asyncio.CancelledError` for cleanup; synchronous CPU-bound handlers
-may check `Context.is_cancellation_requested()` at natural boundaries.
-
 `RPCResult.with_canceling_steps` provides the Flow-wide selector for RPCs.
 RPCs do not support sibling selection because they have no Step execution
 lineage.
@@ -298,14 +328,14 @@ serialization, and invalid handler returns use `FlowDefinitionError`,
 ### Sync vs asyncio
 
 - **Sync (default):** `Client` and `Worker` use blocking gRPC and a thread-pool
-  Worker. Blocking `Client` calls inside `Step.execute` are safe (one pool
-  thread is occupied; other RPCs still run).
+  Worker. A progress generator cooperatively hands each yielded frame to gRPC;
+  `Stream.write` must therefore be yielded. Blocking Client calls inside
+  `Step.execute` are safe while other pool threads remain available.
 - **Asyncio:** `AsyncClient` and `AsyncWorker` use `grpc.aio`. Use
   `Registry(..., allow_async_handlers=True)` when Steps/RPCs are coroutines.
-  Inside async `execute`, inject `AsyncClient` — do not call sync `Client` on
-  the Worker event loop. Sync `Worker` still rejects coroutine handlers at
-  registry construction unless `allow_async_handlers=True` (and even then the
-  sync Worker dispatcher rejects awaitable return values).
+  Step coroutines annotate `AsyncContext`, call `Stream.write` without `await`,
+  and await `context.heartbeat`. Inside async `execute`, inject `AsyncClient` —
+  do not call sync `Client` on the Worker event loop. Async generators are rejected.
 
 Integration scenarios live under
 [`tests/integ`](tests/integ/README.md). They exercise the same workflows,
@@ -318,7 +348,8 @@ The strongly typed contracts, registry, synchronous Client/Worker, optional
 `AsyncClient`/`AsyncWorker` (`grpc.aio`), and Rust-backed BlobCache are
 implemented. Python owns its gRPC transport; the native bridge is limited to
 the shared BlobCache. Design notes:
-[`docs/design/plan/python-sdk-async-apis.md`](../docs/design/plan/python-sdk-async-apis.md).
+[`python-sdk-async-apis.md`](../docs/design/plan/python-sdk-async-apis.md) and
+[`python-sdk-step-streaming.md`](../docs/design/plan/python-sdk-step-streaming.md).
 
 ## Running Dex locally
 

{dex_python_sdk-0.2.5 → dex_python_sdk-0.2.6}/README.md

@@ -13,6 +13,7 @@ the shared Rust Core is used only for BlobCache.
 
 ```python
 from datetime import timedelta
+from typing import Generator
 
 import dex
 
@@ -30,8 +31,9 @@ class Run(dex.Step[str]):
 
     def execute(
         self, context: dex.Context, input: str
-    ) -> dex.StepDecision:
-        progress.write(context, "running")
+    ) -> Generator[dex.StepOutput, None, dex.StepDecision]:
+        yield progress.write(context, "running")
+        yield dex.heartbeat({"phase": "running"})
         return dex.graceful_complete(input)
 
 class CounterFlow(dex.Flow[str]):
@@ -62,23 +64,6 @@ explicit codec only for a custom encoding or a type Registry cannot derive.
 `PersistenceSchema.of(...)` accepts attributes, channels, and streams together and
 partitions them by definition type.
 
-Streams provide best-effort resumable progress messages. Their approximate byte
-budget is shared by all instances of the owning Flow type. Client keys cannot
-contain `#`; Step writes generate `runID#stepExecutionID` and allow one write per
-Stream per invocation.
-
-```python
-client.write_stream(flow_id, progress, "frontend/1", "starting")
-message = client.read_stream(
-    flow_id, progress, resume_token, timeout=timedelta(seconds=30)
-)
-resume_token = message.resume_token
-```
-
-Async Step handlers await **Stream.write**, and **AsyncClient** exposes matching
-async methods. Reads return the decoded value, resume token, creation time, and
-idempotency key.
-
 `Worker` and `AsyncWorker` synchronize all registered Indexed Attributes with
 Dex Server before opening their listener. Existing indexes return immediately;
 failure or the default two-minute deadline aborts startup. An indexed
@@ -128,9 +113,10 @@ Applications implement two generic interfaces from [`dex`](dex/):
   binds the Flow input to the starting Step input. Use `StepList.empty()` when
   a Flow has no Steps.
 - `Step[INPUT]` implements `execute` and optionally `wait_for`. The default
-  Worker path requires synchronous handlers. With `AsyncWorker` and
-  `Registry(..., allow_async_handlers=True)`, handlers may be `async def` and
-  `await` an `AsyncClient`.
+  Worker accepts ordinary synchronous handlers and generator handlers. A generator
+  yields `StepOutput` progress frames and returns its final `Wait` or `StepDecision`.
+  With `AsyncWorker` and `Registry(..., allow_async_handlers=True)`, Step
+  coroutines use `AsyncContext`; async RPCs keep `Context`.
 
 `StepOptions.wait_for_method_timeout` and `execute_method_timeout` bound the
 two handler calls. Timer and channel conditions determine how long a Step waits.
@@ -141,6 +127,58 @@ attempts, total duration, and 1-based attempt numbers. Fallback starts
 immediately; later regular retries continue the backoff sequence at the
 cumulative attempt.
 
+The default Step durability is synchronous. A Flow configuration can select
+asynchronous durability, and a Step method override has highest precedence. The
+default retry total duration is four hours. Regular attempts default to a two-hour
+method timeout and one-minute heartbeat timeout; an explicit heartbeat timeout must
+meet the server minimum, which defaults to ten seconds. Asynchronous durability
+first allows at most three local attempts in seven seconds. The local phase ignores
+method timeouts and heartbeat frames before falling back to a regular activity.
+
+### Step progress and heartbeat recovery
+
+A synchronous handler yields every heartbeat and Stream write. The generator return
+value is the only final result:
+
+```python
+def execute(
+    self, context: dex.Context, input: str
+) -> Generator[dex.StepOutput, None, dex.StepDecision]:
+    yield dex.heartbeat({"offset": 10})
+    yield progress.write(context, "processed 10 items")
+    return dex.graceful_complete(input)
+```
+
+An asynchronous handler keeps its normal coroutine return. Stream writes enqueue
+without waiting for Stream Store acknowledgement; heartbeat waits only for the
+Worker output queue:
+
+```python
+async def execute(
+    self, context: dex.AsyncContext, input: str
+) -> dex.StepDecision:
+    progress.write(context, "started")
+    await context.heartbeat({"offset": 10})
+    return dex.graceful_complete(input)
+```
+
+Call `heartbeat()` or `await context.heartbeat()` without a value to clear previous
+details. Passing Python `None` persists a present null Value. On a later regular
+attempt, use `context.has_last_heartbeat_value()` before decoding with
+`context.get_last_heartbeat_value(ExpectedType)`. A Stream frame is also an implicit
+heartbeat, but it preserves the latest explicit heartbeat state.
+
+A Step may write the same Stream any number of times. Dex assigns
+`#<stepExecutionID>` as each Step message's `StreamMessage.source`. Client writes
+provide their own non-empty source; duplicate sources and `#` are allowed and every
+write appends:
+
+```python
+client.write_stream(flow_id, progress, "frontend#preview", "rendering")
+message = client.read_stream(flow_id, progress)
+print(message.source)
+```
+
 ### Canceling Step executions
 
 A successful Step can cancel queued or active executions while continuing with
@@ -165,14 +203,6 @@ already-canceled, and absent targets are no-ops. Next Steps created by the same
 decision are outside the snapshot. Dex immediately applies the next or close
 action; late decisions, writes, retries, and recovery Steps are discarded.
 
-Set `StepOptions.heartbeat_timeout` on long-running regular Steps so
-cancellation reaches the Worker promptly. It applies to `wait_for` and
-`execute`; local activities ignore it, while an ASYNC fallback uses it. `None`
-and zero disable heartbeats, and positive values must be whole seconds in the
-signed int32 range. `AsyncWorker` cancels the handler's asyncio task. A handler
-may catch `asyncio.CancelledError` for cleanup; synchronous CPU-bound handlers
-may check `Context.is_cancellation_requested()` at natural boundaries.
-
 `RPCResult.with_canceling_steps` provides the Flow-wide selector for RPCs.
 RPCs do not support sibling selection because they have no Step execution
 lineage.
@@ -282,14 +312,14 @@ serialization, and invalid handler returns use `FlowDefinitionError`,
 ### Sync vs asyncio
 
 - **Sync (default):** `Client` and `Worker` use blocking gRPC and a thread-pool
-  Worker. Blocking `Client` calls inside `Step.execute` are safe (one pool
-  thread is occupied; other RPCs still run).
+  Worker. A progress generator cooperatively hands each yielded frame to gRPC;
+  `Stream.write` must therefore be yielded. Blocking Client calls inside
+  `Step.execute` are safe while other pool threads remain available.
 - **Asyncio:** `AsyncClient` and `AsyncWorker` use `grpc.aio`. Use
   `Registry(..., allow_async_handlers=True)` when Steps/RPCs are coroutines.
-  Inside async `execute`, inject `AsyncClient` — do not call sync `Client` on
-  the Worker event loop. Sync `Worker` still rejects coroutine handlers at
-  registry construction unless `allow_async_handlers=True` (and even then the
-  sync Worker dispatcher rejects awaitable return values).
+  Step coroutines annotate `AsyncContext`, call `Stream.write` without `await`,
+  and await `context.heartbeat`. Inside async `execute`, inject `AsyncClient` —
+  do not call sync `Client` on the Worker event loop. Async generators are rejected.
 
 Integration scenarios live under
 [`tests/integ`](tests/integ/README.md). They exercise the same workflows,
@@ -302,7 +332,8 @@ The strongly typed contracts, registry, synchronous Client/Worker, optional
 `AsyncClient`/`AsyncWorker` (`grpc.aio`), and Rust-backed BlobCache are
 implemented. Python owns its gRPC transport; the native bridge is limited to
 the shared BlobCache. Design notes:
-[`docs/design/plan/python-sdk-async-apis.md`](../docs/design/plan/python-sdk-async-apis.md).
+[`python-sdk-async-apis.md`](../docs/design/plan/python-sdk-async-apis.md) and
+[`python-sdk-step-streaming.md`](../docs/design/plan/python-sdk-step-streaming.md).
 
 ## Running Dex locally
 

{dex_python_sdk-0.2.5 → dex_python_sdk-0.2.6}/dex/__init__.py

@@ -35,7 +35,7 @@ from dex.codec import (
     WireKind,
 )
 from dex.condition import ConditionCombination
-from dex.context import Context
+from dex.context import AsyncContext, Context
 from dex.flow import Flow, PersistenceSchema, Registry, RPCResult, rpc
 from dex.flow_config import ActiveStepSearchMode, FlowConfig
 from dex.flow_info import (
@@ -80,6 +80,7 @@ from dex.step import (
     StepList,
     StepMovement,
     StepOptions,
+    StepOutput,
     WaitForFailurePolicy,
     dead_end,
     force_complete,
@@ -88,6 +89,7 @@ from dex.step import (
     go_to,
     go_to_many,
     graceful_complete,
+    heartbeat,
 )
 from dex.step_execution import StepExecutionId, TimerId
 from dex.stream import Stream, StreamMessage
@@ -108,6 +110,7 @@ __all__ = [
     "AttributeIndex",
     "AttributeLock",
     "AttributeMap",
+    "AsyncContext",
     "AsyncClient",
     "AsyncWorker",
     "BlobCache",
@@ -161,6 +164,7 @@ __all__ = [
     "StepList",
     "StepDurability",
     "StepMovement",
+    "StepOutput",
     "StepOptions",
     "Stream",
     "StreamMessage",
@@ -182,6 +186,7 @@ __all__ = [
     "force_complete_if_channels_empty",
     "force_fail",
     "graceful_complete",
+    "heartbeat",
     "go_to",
     "go_to_many",
     "open_blob_cache",

{dex_python_sdk-0.2.5 → dex_python_sdk-0.2.6}/dex/_async_value_hydrator.py

@@ -75,10 +75,16 @@ class AsyncValueHydrator:
     ) -> pb.InvokeWaitForMethodRequest:
         result = pb.InvokeWaitForMethodRequest()
         result.CopyFrom(request)
-        values = [request.step_input, *(entry.value for entry in request.attributes)]
-        hydrated = await self.hydrate_all(values)
-        result.step_input.CopyFrom(hydrated[0])
-        for entry, value in zip(result.attributes, hydrated[1:]):
+        has_heartbeat = request.context.HasField("last_heartbeat_value")
+        values = [request.step_input]
+        if has_heartbeat:
+            values.append(request.context.last_heartbeat_value)
+        values.extend(entry.value for entry in request.attributes)
+        hydrated = iter(await self.hydrate_all(values))
+        result.step_input.CopyFrom(next(hydrated))
+        if has_heartbeat:
+            result.context.last_heartbeat_value.CopyFrom(next(hydrated))
+        for entry, value in zip(result.attributes, hydrated):
             entry.value.CopyFrom(value)
         return result
 
@@ -89,7 +95,10 @@ class AsyncValueHydrator:
         result = pb.InvokeExecuteMethodRequest()
         result.CopyFrom(request)
         has_step_input = request.HasField("step_input")
+        has_heartbeat = request.context.HasField("last_heartbeat_value")
         values = [request.step_input] if has_step_input else []
+        if has_heartbeat:
+            values.append(request.context.last_heartbeat_value)
         values.extend(entry.value for entry in request.attributes)
         values.extend(entry.value for entry in request.step_exe_locals)
         for channel_result in request.condition_results.channel_results:
@@ -97,6 +106,8 @@ class AsyncValueHydrator:
         hydrated = iter(await self.hydrate_all(values))
         if has_step_input:
             result.step_input.CopyFrom(next(hydrated))
+        if has_heartbeat:
+            result.context.last_heartbeat_value.CopyFrom(next(hydrated))
         for entry in result.attributes:
             entry.value.CopyFrom(next(hydrated))
         for entry in result.step_exe_locals:

{dex_python_sdk-0.2.5 → dex_python_sdk-0.2.6}/dex/_async_worker_dispatcher.py

@@ -8,7 +8,10 @@
 
 from __future__ import annotations
 
-from inspect import isawaitable
+import asyncio
+from collections.abc import AsyncGenerator, AsyncIterator
+from contextlib import suppress
+from inspect import isawaitable, isgenerator
 from typing import Any, Callable
 
 from dex._async_value_hydrator import AsyncValueHydrator
@@ -18,27 +21,64 @@ from dex._worker_dispatcher import _TIMEOUT_HANDLER_STEP_TYPE, WorkerDispatcher
 from dex.dexpb import dex_pb2 as pb
 from dex.flow import Registry, RPCResult
 from dex.runtime_errors import InvalidStepResultError, ValueMappingError
-from dex.step import StepDecision
+from dex.step import StepDecision, StepOutput
 from dex.wait import Wait
 
 
+class _AsyncStepOutputEmitter:
+    def __init__(self) -> None:
+        self.queue: asyncio.Queue[StepOutput] = asyncio.Queue()
+        self._is_closed = False
+
+    def emit_nowait(self, output: StepOutput) -> None:
+        if self._is_closed:
+            return
+        self.queue.put_nowait(output)
+
+    async def emit(self, output: StepOutput) -> None:
+        if self._is_closed:
+            raise asyncio.CancelledError
+        self.queue.put_nowait(output)
+        await asyncio.sleep(0)
+        if self._is_closed:
+            raise asyncio.CancelledError
+
+    def close(self) -> None:
+        self._is_closed = True
+
+
 class AsyncWorkerDispatcher(WorkerDispatcher):
     def __init__(
         self,
         registry: Registry,
         values: ValueMapper,
         hydrator: AsyncValueHydrator,
-        stream_writer: Callable[[pb.WriteStreamRequest], Any],
     ) -> None:
         self._registry = registry
         self._values = values
         self._async_hydrator = hydrator
-        self._stream_writer = stream_writer
 
     async def invoke_wait_for(  # type: ignore[override]
         self,
         original: pb.InvokeWaitForMethodRequest,
         is_active: Callable[[], bool] | None = None,
+    ) -> AsyncGenerator[pb.InvokeWaitForMethodOutput, None]:
+        emitter = _AsyncStepOutputEmitter()
+        handler = asyncio.create_task(
+            self._invoke_wait_for_result(original, emitter, is_active)
+        )
+        try:
+            async for output in self._drain_outputs(emitter, handler):
+                yield self._map_wait_for_output(output)
+            yield pb.InvokeWaitForMethodOutput(result=await handler)
+        finally:
+            await self._close_invocation(emitter, handler)
+
+    async def _invoke_wait_for_result(
+        self,
+        original: pb.InvokeWaitForMethodRequest,
+        emitter: _AsyncStepOutputEmitter,
+        is_active: Callable[[], bool] | None,
     ) -> pb.InvokeWaitForMethodResponse:
         request = await self._async_hydrator.wait_for_request(original)
         flow = self._registry._flow_by_type(request.flow_type)
@@ -48,12 +88,20 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
             flow,
             request.context,
             self._values,
-            self._stream_writer,
             request.attributes,
             is_active=is_active,
+            output_emitter=emitter,
         )
         input = self._values.decode(request.step_input, step.input_codec)
         wait = step.step.wait_for(context, input)
+        if isgenerator(wait):
+            wait.close()
+            raise InvalidStepResultError(
+                flow.name,
+                step.name,
+                "wait_for",
+                "synchronous generators require Worker",
+            )
         if isawaitable(wait):
             wait = await wait
         try:
@@ -78,11 +126,30 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
         self,
         original: pb.InvokeExecuteMethodRequest,
         is_active: Callable[[], bool] | None = None,
+    ) -> AsyncGenerator[pb.InvokeExecuteMethodOutput, None]:
+        if original.step_type == _TIMEOUT_HANDLER_STEP_TYPE:
+            response = await self._invoke_timeout_handler_async(original)
+            yield pb.InvokeExecuteMethodOutput(result=response)
+            return
+        emitter = _AsyncStepOutputEmitter()
+        handler = asyncio.create_task(
+            self._invoke_execute_result(original, emitter, is_active)
+        )
+        try:
+            async for output in self._drain_outputs(emitter, handler):
+                yield self._map_execute_output(output)
+            yield pb.InvokeExecuteMethodOutput(result=await handler)
+        finally:
+            await self._close_invocation(emitter, handler)
+
+    async def _invoke_execute_result(
+        self,
+        original: pb.InvokeExecuteMethodRequest,
+        emitter: _AsyncStepOutputEmitter,
+        is_active: Callable[[], bool] | None,
     ) -> pb.InvokeExecuteMethodResponse:
         request = await self._async_hydrator.execute_request(original)
         flow = self._registry._flow_by_type(request.flow_type)
-        if request.step_type == _TIMEOUT_HANDLER_STEP_TYPE:
-            return await self._invoke_timeout_handler_async(request, flow)
         step = flow.step(request.step_type)
         condition_results = (
             request.condition_results if request.HasField("condition_results") else None
@@ -92,14 +159,22 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
             flow,
             request.context,
             self._values,
-            self._stream_writer,
             request.attributes,
             request.step_exe_locals,
             condition_results,
             is_active=is_active,
+            output_emitter=emitter,
         )
         input = self._values.decode(request.step_input, step.input_codec)
         decision: Any = step.step.execute(context, input)
+        if isgenerator(decision):
+            decision.close()
+            raise InvalidStepResultError(
+                flow.name,
+                step.name,
+                "execute",
+                "synchronous generators require Worker",
+            )
         if isawaitable(decision):
             decision = await decision
         try:
@@ -119,9 +194,10 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
 
     async def _invoke_timeout_handler_async(
         self,
-        request: pb.InvokeExecuteMethodRequest,
-        flow: Any,
+        original: pb.InvokeExecuteMethodRequest,
     ) -> pb.InvokeExecuteMethodResponse:
+        request = await self._async_hydrator.execute_request(original)
+        flow = self._registry._flow_by_type(request.flow_type)
         if request.HasField("step_input"):
             raise InvalidStepResultError(
                 flow.name, _TIMEOUT_HANDLER_STEP_TYPE, "execute", "input must be absent"
@@ -137,11 +213,10 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
             request.condition_results if request.HasField("condition_results") else None
         )
         context = InvocationContext(
-            InvocationMethod.EXECUTE,
+            InvocationMethod.TIMEOUT,
             flow,
             request.context,
             self._values,
-            self._stream_writer,
             request.attributes,
             request.step_exe_locals,
             condition_results,
@@ -177,7 +252,6 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
             flow,
             request.context,
             self._values,
-            self._stream_writer,
             request.attributes,
             channel_infos=dict(request.channel_infos),
             is_active=is_active,
@@ -218,3 +292,38 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
             raise
         except (TypeError, ValueError) as error:
             raise InvalidStepResultError(flow.name, None, "rpc", str(error)) from error
+
+    async def _drain_outputs(
+        self,
+        emitter: _AsyncStepOutputEmitter,
+        handler: asyncio.Task[Any],
+    ) -> AsyncIterator[StepOutput]:
+        while not handler.done() or not emitter.queue.empty():
+            if not emitter.queue.empty():
+                yield emitter.queue.get_nowait()
+                continue
+            output = asyncio.create_task(emitter.queue.get())
+            try:
+                done, _ = await asyncio.wait(
+                    (handler, output),
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+                if output in done:
+                    yield output.result()
+            finally:
+                if not output.done():
+                    output.cancel()
+                    with suppress(asyncio.CancelledError):
+                        await output
+
+    async def _close_invocation(
+        self,
+        emitter: _AsyncStepOutputEmitter,
+        handler: asyncio.Task[Any],
+    ) -> None:
+        emitter.close()
+        if handler.done():
+            return
+        handler.cancel()
+        with suppress(asyncio.CancelledError):
+            await handler