dex-python-sdk 1.26.0__tar.gz → 1.28.0__tar.gz

{dex_python_sdk-1.26.0 → dex_python_sdk-1.28.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dex-python-sdk
-Version: 1.26.0
+Version: 1.28.0
 Requires-Dist: grpcio>=1.83.0
 Requires-Dist: grpcio-status>=1.83.0
 Requires-Dist: protobuf>=7.35.1
@@ -394,6 +394,7 @@ registry = dex.Registry((flow,))
 Registry derives codecs from declared Python types and handler annotations.
 Built-in primitive types and dataclasses need no codec arguments. Register an
 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.
 
@@ -465,6 +466,56 @@ pip install dex-python-sdk==0.1.0
 
 See [samples](../examples/python) for use case examples.
 
+### JSON dataclasses
+
+Automatic JSON codecs support parameterized generic dataclasses such as
+`Result[Receipt]`, including nested generics, typed containers, `Optional[T]`, and
+`T | None`. Decoding restores nested dataclass instances. Missing keys use the
+field's `default` or `default_factory`; a missing required key raises `ValueError`
+with the dataclass and wire field name. Unknown keys are ignored.
+
+Field metadata controls the wire representation:
+
+- `field(metadata={"dex_json_name": "callId"})` uses `callId` instead of the
+  Python field name for encoding and decoding. Wire names must be non-empty and
+  unique within a dataclass.
+- `field(default=None, metadata={"dex_json_omit_default": True})` omits a field
+  only when it equals its declared default. This also supports `default_factory`;
+  the factory is called on each encode to obtain the comparison value, so use a
+  deterministic factory without side effects. Without this opt-in, default
+  values, empty containers, and `None` remain present in JSON.
+- `field(metadata={"dex_json_format": "rfc3339"})` formats a datetime field like
+  Go's `encoding/json`: UTC uses `Z` and trailing fractional zeroes are removed.
+  This requires a datetime or optional datetime field with an aware value and a
+  whole-minute timezone offset. Ordinary datetime fields keep their existing
+  `isoformat()` encoding. All datetime fields accept Go RFC3339 strings on
+  decode, including `Z`, numeric offsets, and fractional seconds. Python retains
+  microsecond precision; additional fractional digits are truncated.
+
+Metadata options can be combined on one field. Existing dataclasses retain their
+field names and encoded values unless metadata opts into a change.
+
+Unions select by JSON kind: null, boolean, number, string, array, or object.
+A union declared directly, such as a Step input of `Receipt | None`, needs
+alternatives with different kinds. `Receipt | None`, `int | str`, and
+`int | bool` are supported. `int | float`, `str | datetime`, two dataclasses,
+or two list types are ambiguous and raise `TypeError` when the codec is
+resolved. Use an explicit `JsonCodec` with a discriminator for these shapes.
+Union values always use JSON on the wire, including primitive alternatives.
+
+Unions inside dataclass fields and containers keep the decoding of earlier
+releases and never raise when the codec is resolved. Such a union decodes a
+value into a dataclass, Enum, or datetime alternative, or a container of those,
+when the value's JSON kind matches only that alternative. It returns any other
+value unchanged, without a type check. Fields typed by a type variable and
+`Mapping` or `Sequence` types follow the same rule: `Result[Receipt]` decodes a
+`Receipt`, but `Result[int]` returns the JSON value of its type-variable field
+unchanged. A `float` accepts a JSON integer and keeps it as an `int`.
+
+See `help(CodecRegistry)` for a complete generic dataclass example and
+[`tests/fixtures/json_codec`](tests/fixtures/json_codec) for a reproducible
+Go-to-Python wire fixture.
+
 ## Requirements
 
 - Python 3.11+
@@ -484,6 +535,37 @@ Applications implement two generic interfaces from [`dex`](dex/):
   With `AsyncWorker` and `Registry(..., allow_async_handlers=True)`, Step
   coroutines use `AsyncContext`; async RPCs keep `Context`.
 
+Steps are registered by `get_step_type()`, which defaults to the class name.
+Multiple instances of one class may return different Step types in the same Flow.
+A duplicate Step type fails registration. `go_to`, `StepMovement.of`,
+`StepOptions.on_execute_failure_proceed_to`,
+`FlowTimeoutHandlerOptions.on_failure_proceed_to`, and decision or RPC
+cancellation selections accept a registered instance, `StepRef("StepType")`, or
+a class. A class must have exactly one registered instance; otherwise select an
+instance or use `StepRef`. Instances match by identity and must belong to the
+containing Flow, so a new instance of a registered class is rejected.
+
+Declared execute-failure targets are validated when the Registry is constructed.
+The Registry also rejects an execute-failure target whose input type differs from
+the failing Step's input type. It skips this check when either type is `Any`,
+`object`, or an unresolved TypeVar. Timeout-handler failure targets are validated
+when Flow start options are mapped. Handler-created movements and cancellation
+selections, including StepRef targets, are validated when the Worker applies the
+returned decision; unknown Step types and ambiguous classes raise
+`FlowDefinitionError` with the Flow and target details. StepRef contains only a
+Step type and cannot execute or be registered in a StepList. Its input is checked
+by the destination codec at runtime.
+
+For a generic Step whose input annotation is a TypeVar, override
+`get_input_type()` to return a concrete type supported by `CodecRegistry`. The
+override takes precedence over both `execute` and `wait_for` input annotations;
+Context and return annotations are still required. The default returns `None`,
+which preserves annotation inference; return `type(None)` for a Step that accepts
+no input. The Registry resolves the type through its codec registrations, so
+custom codecs work too. See the runnable
+[generic Step instances test](tests/integ/step_instances_flow.py), which registers
+integer and string instances and routes to both by instance and StepRef.
+
 `StepOptions.wait_for_method_timeout` and `execute_method_timeout` bound the
 two handler calls. Timer and channel conditions determine how long a Step waits.
 
@@ -501,6 +583,26 @@ 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.
 
+### Flow run start time
+
+`Context.flow_started_at` returns the current Flow run's server start time as a
+timezone-aware UTC `datetime` with one-second precision. It stays the same across
+Steps and retries in that run. It is available in `wait_for`, `execute`, Flow
+timeout handlers, and RPC handlers for both Worker and AsyncWorker. Compare it
+with `datetime.now(timezone.utc)` to measure the run's age. It matches
+`Client.describe_flow(flow_id).started_at` to the second.
+
+### Step attempt deadline
+
+`Context.deadline` is the absolute deadline supplied by the Dex Server for the
+current `wait_for`, `execute`, or Flow timeout handler attempt. It is a
+timezone-aware UTC `datetime` with microsecond precision, not a duration.
+Subtract `datetime.now(timezone.utc)` and call `total_seconds()` to budget
+remaining work after checking for `None`. A shorter remaining
+`RetryPolicy.total_duration` limits this deadline; each retry receives its own
+effective deadline. RPC handlers and tasks without a supplied deadline return
+`None`.
+
 ### Step progress and heartbeat recovery
 
 A synchronous handler yields every heartbeat and Stream write. The generator return

{dex_python_sdk-1.26.0 → dex_python_sdk-1.28.0}/README.md

@@ -378,6 +378,7 @@ registry = dex.Registry((flow,))
 Registry derives codecs from declared Python types and handler annotations.
 Built-in primitive types and dataclasses need no codec arguments. Register an
 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.
 
@@ -449,6 +450,56 @@ pip install dex-python-sdk==0.1.0
 
 See [samples](../examples/python) for use case examples.
 
+### JSON dataclasses
+
+Automatic JSON codecs support parameterized generic dataclasses such as
+`Result[Receipt]`, including nested generics, typed containers, `Optional[T]`, and
+`T | None`. Decoding restores nested dataclass instances. Missing keys use the
+field's `default` or `default_factory`; a missing required key raises `ValueError`
+with the dataclass and wire field name. Unknown keys are ignored.
+
+Field metadata controls the wire representation:
+
+- `field(metadata={"dex_json_name": "callId"})` uses `callId` instead of the
+  Python field name for encoding and decoding. Wire names must be non-empty and
+  unique within a dataclass.
+- `field(default=None, metadata={"dex_json_omit_default": True})` omits a field
+  only when it equals its declared default. This also supports `default_factory`;
+  the factory is called on each encode to obtain the comparison value, so use a
+  deterministic factory without side effects. Without this opt-in, default
+  values, empty containers, and `None` remain present in JSON.
+- `field(metadata={"dex_json_format": "rfc3339"})` formats a datetime field like
+  Go's `encoding/json`: UTC uses `Z` and trailing fractional zeroes are removed.
+  This requires a datetime or optional datetime field with an aware value and a
+  whole-minute timezone offset. Ordinary datetime fields keep their existing
+  `isoformat()` encoding. All datetime fields accept Go RFC3339 strings on
+  decode, including `Z`, numeric offsets, and fractional seconds. Python retains
+  microsecond precision; additional fractional digits are truncated.
+
+Metadata options can be combined on one field. Existing dataclasses retain their
+field names and encoded values unless metadata opts into a change.
+
+Unions select by JSON kind: null, boolean, number, string, array, or object.
+A union declared directly, such as a Step input of `Receipt | None`, needs
+alternatives with different kinds. `Receipt | None`, `int | str`, and
+`int | bool` are supported. `int | float`, `str | datetime`, two dataclasses,
+or two list types are ambiguous and raise `TypeError` when the codec is
+resolved. Use an explicit `JsonCodec` with a discriminator for these shapes.
+Union values always use JSON on the wire, including primitive alternatives.
+
+Unions inside dataclass fields and containers keep the decoding of earlier
+releases and never raise when the codec is resolved. Such a union decodes a
+value into a dataclass, Enum, or datetime alternative, or a container of those,
+when the value's JSON kind matches only that alternative. It returns any other
+value unchanged, without a type check. Fields typed by a type variable and
+`Mapping` or `Sequence` types follow the same rule: `Result[Receipt]` decodes a
+`Receipt`, but `Result[int]` returns the JSON value of its type-variable field
+unchanged. A `float` accepts a JSON integer and keeps it as an `int`.
+
+See `help(CodecRegistry)` for a complete generic dataclass example and
+[`tests/fixtures/json_codec`](tests/fixtures/json_codec) for a reproducible
+Go-to-Python wire fixture.
+
 ## Requirements
 
 - Python 3.11+
@@ -468,6 +519,37 @@ Applications implement two generic interfaces from [`dex`](dex/):
   With `AsyncWorker` and `Registry(..., allow_async_handlers=True)`, Step
   coroutines use `AsyncContext`; async RPCs keep `Context`.
 
+Steps are registered by `get_step_type()`, which defaults to the class name.
+Multiple instances of one class may return different Step types in the same Flow.
+A duplicate Step type fails registration. `go_to`, `StepMovement.of`,
+`StepOptions.on_execute_failure_proceed_to`,
+`FlowTimeoutHandlerOptions.on_failure_proceed_to`, and decision or RPC
+cancellation selections accept a registered instance, `StepRef("StepType")`, or
+a class. A class must have exactly one registered instance; otherwise select an
+instance or use `StepRef`. Instances match by identity and must belong to the
+containing Flow, so a new instance of a registered class is rejected.
+
+Declared execute-failure targets are validated when the Registry is constructed.
+The Registry also rejects an execute-failure target whose input type differs from
+the failing Step's input type. It skips this check when either type is `Any`,
+`object`, or an unresolved TypeVar. Timeout-handler failure targets are validated
+when Flow start options are mapped. Handler-created movements and cancellation
+selections, including StepRef targets, are validated when the Worker applies the
+returned decision; unknown Step types and ambiguous classes raise
+`FlowDefinitionError` with the Flow and target details. StepRef contains only a
+Step type and cannot execute or be registered in a StepList. Its input is checked
+by the destination codec at runtime.
+
+For a generic Step whose input annotation is a TypeVar, override
+`get_input_type()` to return a concrete type supported by `CodecRegistry`. The
+override takes precedence over both `execute` and `wait_for` input annotations;
+Context and return annotations are still required. The default returns `None`,
+which preserves annotation inference; return `type(None)` for a Step that accepts
+no input. The Registry resolves the type through its codec registrations, so
+custom codecs work too. See the runnable
+[generic Step instances test](tests/integ/step_instances_flow.py), which registers
+integer and string instances and routes to both by instance and StepRef.
+
 `StepOptions.wait_for_method_timeout` and `execute_method_timeout` bound the
 two handler calls. Timer and channel conditions determine how long a Step waits.
 
@@ -485,6 +567,26 @@ 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.
 
+### Flow run start time
+
+`Context.flow_started_at` returns the current Flow run's server start time as a
+timezone-aware UTC `datetime` with one-second precision. It stays the same across
+Steps and retries in that run. It is available in `wait_for`, `execute`, Flow
+timeout handlers, and RPC handlers for both Worker and AsyncWorker. Compare it
+with `datetime.now(timezone.utc)` to measure the run's age. It matches
+`Client.describe_flow(flow_id).started_at` to the second.
+
+### Step attempt deadline
+
+`Context.deadline` is the absolute deadline supplied by the Dex Server for the
+current `wait_for`, `execute`, or Flow timeout handler attempt. It is a
+timezone-aware UTC `datetime` with microsecond precision, not a duration.
+Subtract `datetime.now(timezone.utc)` and call `total_seconds()` to budget
+remaining work after checking for `None`. A shorter remaining
+`RetryPolicy.total_duration` limits this deadline; each retry receives its own
+effective deadline. RPC handlers and tasks without a supplied deadline return
+`None`.
+
 ### Step progress and heartbeat recovery
 
 A synchronous handler yields every heartbeat and Stream write. The generator return

{dex_python_sdk-1.26.0 → dex_python_sdk-1.28.0}/dex/__init__.py

@@ -101,6 +101,8 @@ from dex.step import (
     StepMovement,
     StepOptions,
     StepOutput,
+    StepRef,
+    StepTarget,
     WaitForFailurePolicy,
     dead_end,
     force_complete,
@@ -210,6 +212,8 @@ __all__ = [
     "StepDurability",
     "StepMovement",
     "StepOutput",
+    "StepRef",
+    "StepTarget",
     "StepOptions",
     "Stream",
     "StreamMessage",

{dex_python_sdk-1.26.0 → dex_python_sdk-1.28.0}/dex/_async_worker_dispatcher.py

@@ -11,6 +11,7 @@ from __future__ import annotations
 import asyncio
 from collections.abc import AsyncGenerator, AsyncIterator
 from contextlib import suppress
+from datetime import datetime
 from inspect import isawaitable, isgenerator
 from typing import Any, Callable
 
@@ -80,10 +81,11 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
         self,
         original: pb.InvokeWaitForMethodRequest,
         is_active: Callable[[], bool] | None = None,
+        deadline: datetime | None = None,
     ) -> AsyncGenerator[pb.InvokeWaitForMethodOutput, None]:
         emitter = _AsyncStepOutputEmitter()
         handler = asyncio.create_task(
-            self._invoke_wait_for_result(original, emitter, is_active)
+            self._invoke_wait_for_result(original, emitter, is_active, deadline)
         )
         try:
             async for output in self._drain_outputs(emitter, handler):
@@ -97,6 +99,7 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
         original: pb.InvokeWaitForMethodRequest,
         emitter: _AsyncStepOutputEmitter,
         is_active: Callable[[], bool] | None,
+        deadline: datetime | None,
     ) -> pb.InvokeWaitForMethodResponse:
         request = await self._async_hydrator.wait_for_request(original)
         flow = self._registry._flow_by_type(request.flow_type)
@@ -113,6 +116,7 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
             loaded_channel_names=request.loaded_channel_names,
             loaded_channel_map_instances=request.loaded_channel_map_instances,
             is_active=is_active,
+            deadline=deadline,
             output_emitter=emitter,
         )
         input = self._values.decode(request.step_input, step.input_codec)
@@ -164,11 +168,14 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
         self,
         original: pb.InvokeExecuteMethodRequest,
         is_active: Callable[[], bool] | None = None,
+        deadline: datetime | None = None,
     ) -> AsyncGenerator[pb.InvokeExecuteMethodOutput, None]:
         if original.step_type == _TIMEOUT_HANDLER_STEP_TYPE:
             emitter = _AsyncStepOutputEmitter()
             handler = asyncio.create_task(
-                self._invoke_timeout_handler_async(original, emitter, is_active)
+                self._invoke_timeout_handler_async(
+                    original, emitter, is_active, deadline
+                )
             )
             try:
                 async for output in self._drain_outputs(emitter, handler):
@@ -179,7 +186,7 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
             return
         emitter = _AsyncStepOutputEmitter()
         handler = asyncio.create_task(
-            self._invoke_execute_result(original, emitter, is_active)
+            self._invoke_execute_result(original, emitter, is_active, deadline)
         )
         try:
             async for output in self._drain_outputs(emitter, handler):
@@ -193,6 +200,7 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
         original: pb.InvokeExecuteMethodRequest,
         emitter: _AsyncStepOutputEmitter,
         is_active: Callable[[], bool] | None,
+        deadline: datetime | None,
     ) -> pb.InvokeExecuteMethodResponse:
         request = await self._async_hydrator.execute_request(original)
         flow = self._registry._flow_by_type(request.flow_type)
@@ -214,6 +222,7 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
             loaded_channel_names=request.loaded_channel_names,
             loaded_channel_map_instances=request.loaded_channel_map_instances,
             is_active=is_active,
+            deadline=deadline,
             output_emitter=emitter,
         )
         input = self._values.decode(request.step_input, step.input_codec)
@@ -264,6 +273,7 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
         original: pb.InvokeExecuteMethodRequest,
         emitter: _AsyncStepOutputEmitter,
         is_active: Callable[[], bool] | None,
+        deadline: datetime | None,
     ) -> pb.InvokeExecuteMethodResponse:
         request = await self._async_hydrator.execute_request(original)
         flow = self._registry._flow_by_type(request.flow_type)
@@ -295,6 +305,7 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
             loaded_channel_names=request.loaded_channel_names,
             loaded_channel_map_instances=request.loaded_channel_map_instances,
             is_active=is_active,
+            deadline=deadline,
             output_emitter=emitter,
         )
         response: pb.InvokeExecuteMethodResponse | None = None

{dex_python_sdk-1.26.0 → dex_python_sdk-1.28.0}/dex/_invocation_context.py

@@ -83,7 +83,9 @@ class InvocationContext:
         loaded_channel_map_instances: Sequence[str] = (),
         is_active: Callable[[], bool] | None = None,
         output_emitter: _StepOutputEmitter | None = None,
+        deadline: datetime | None = None,
     ) -> None:
+        self._deadline = None if method is InvocationMethod.RPC else deadline
         self._method = method
         self._flow = flow
         self._metadata = metadata
@@ -115,6 +117,10 @@ class InvocationContext:
     def run_id(self) -> str:
         return self._metadata.run_id
 
+    @property
+    def flow_started_at(self) -> datetime:
+        return datetime.fromtimestamp(self._metadata.flow_started_timestamp, UTC)
+
     @property
     def step_execution_id(self) -> str:
         return self._metadata.step_execution_id
@@ -136,6 +142,10 @@ class InvocationContext:
     def first_attempt_at(self) -> datetime:
         return datetime.fromtimestamp(self._metadata.first_attempt_timestamp, UTC)
 
+    @property
+    def deadline(self) -> datetime | None:
+        return self._deadline
+
     @property
     def attempt(self) -> int:
         return self._metadata.attempt

{dex_python_sdk-1.26.0 → dex_python_sdk-1.28.0}/dex/_pull_worker.py

@@ -705,16 +705,25 @@ class PullStream:
             await self._send_error(task.task_id, failure, status)
 
     async def _execute_async(self, task: pb.PollWorkerTask, kind: str | None) -> None:
+        deadline = (
+            task.deadline.ToDatetime(tzinfo=timezone.utc)
+            if task.HasField("deadline")
+            else None
+        )
         if kind == "wait_for":
             async for output in self.runtime._dispatcher.invoke_wait_for(
-                task.wait_for, self._active.is_set
+                task.wait_for,
+                self._active.is_set,
+                deadline,
             ):
                 await self._send_frame(
                     pb.PollWorkerTaskFrame(task_id=task.task_id, wait_for=output)
                 )
         elif kind == "execute":
             async for output in self.runtime._dispatcher.invoke_execute(
-                task.execute, self._active.is_set
+                task.execute,
+                self._active.is_set,
+                deadline,
             ):
                 await self._send_frame(
                     pb.PollWorkerTaskFrame(task_id=task.task_id, execute=output)
@@ -730,9 +739,16 @@ class PullStream:
             raise ValueError(f"unsupported Worker task request {kind}")
 
     def _execute_sync(self, task: pb.PollWorkerTask, kind: str | None) -> None:
+        deadline = (
+            task.deadline.ToDatetime(tzinfo=timezone.utc)
+            if task.HasField("deadline")
+            else None
+        )
         if kind == "wait_for":
             outputs = self.runtime._dispatcher.invoke_wait_for(
-                task.wait_for, self._active.is_set
+                task.wait_for,
+                self._active.is_set,
+                deadline,
             )
             for output in outputs:
                 self._send_from_thread(
@@ -740,7 +756,9 @@ class PullStream:
                 )
         elif kind == "execute":
             outputs = self.runtime._dispatcher.invoke_execute(
-                task.execute, self._active.is_set
+                task.execute,
+                self._active.is_set,
+                deadline,
             )
             for output in outputs:
                 self._send_from_thread(

{dex_python_sdk-1.26.0 → dex_python_sdk-1.28.0}/dex/_value_mapper.py

@@ -9,7 +9,7 @@
 from __future__ import annotations
 
 import json
-from dataclasses import asdict, is_dataclass
+from dataclasses import is_dataclass
 from datetime import datetime
 from enum import Enum
 from typing import Any, cast
@@ -17,7 +17,14 @@ from typing import Any, cast
 from google.protobuf import struct_pb2
 
 from dex.attribute import AttributeIndex, IndexType
-from dex.codec import Codec, CodecRegistry, Value, WireKind
+from dex.codec import (
+    Codec,
+    CodecRegistry,
+    Value,
+    WireKind,
+    _encode_json_value,
+    _validate_json_value_metadata,
+)
 from dex.dexpb import dex_pb2 as pb
 from dex.runtime_errors import ValueMappingError
 
@@ -87,6 +94,7 @@ class ValueMapper:
             codec = self._codecs.resolve(value_type)
         except TypeError:
             if isinstance(value, (list, tuple, dict)):
+                _validate_json_value_metadata(value, set())
                 return pb.Value(
                     obj_value=pb.EncodedObject(
                         encoding=_JSON_ENCODING,
@@ -213,7 +221,7 @@ class ValueMapper:
     def _encode_json_collection(value: object) -> bytes:
         def default(item: object) -> object:
             if is_dataclass(item) and not isinstance(item, type):
-                return asdict(item)
+                return _encode_json_value(item)
             if isinstance(item, datetime):
                 return item.isoformat()
             if isinstance(item, Enum):

{dex_python_sdk-1.26.0 → dex_python_sdk-1.28.0}/dex/_worker_dispatcher.py

@@ -8,7 +8,7 @@
 
 from __future__ import annotations
 
-from datetime import timedelta
+from datetime import datetime, timedelta
 from inspect import isawaitable, isgenerator
 from typing import Any, Callable, Generator, cast
 
@@ -40,12 +40,12 @@ from dex.runtime_errors import InvalidStepResultError, ValueMappingError
 from dex.step import (
     DecisionKind,
     RetryPolicy,
-    Step,
     StepDecision,
     StepDurability,
     StepMovement,
     StepOptions,
     StepOutput,
+    StepTarget,
     WaitForFailurePolicy,
     _HeartbeatStepOutput,
     _StreamStepOutput,
@@ -74,6 +74,7 @@ class WorkerDispatcher:
         self,
         original: pb.InvokeWaitForMethodRequest,
         is_active: Callable[[], bool] | None = None,
+        deadline: datetime | None = None,
     ) -> Generator[pb.InvokeWaitForMethodOutput, None, None]:
         request = self._hydrator.wait_for_request(original)
         flow = self._registry._flow_by_type(request.flow_type)
@@ -90,6 +91,7 @@ class WorkerDispatcher:
             loaded_channel_names=request.loaded_channel_names,
             loaded_channel_map_instances=request.loaded_channel_map_instances,
             is_active=is_active,
+            deadline=deadline,
         )
         input = self._values.decode(request.step_input, step.input_codec)
         wait = step.step.wait_for(context, input)
@@ -125,11 +127,12 @@ class WorkerDispatcher:
         self,
         original: pb.InvokeExecuteMethodRequest,
         is_active: Callable[[], bool] | None = None,
+        deadline: datetime | None = None,
     ) -> Generator[pb.InvokeExecuteMethodOutput, None, None]:
         request = self._hydrator.execute_request(original)
         flow = self._registry._flow_by_type(request.flow_type)
         if request.step_type == _TIMEOUT_HANDLER_STEP_TYPE:
-            response = yield from self._invoke_timeout_handler(request, flow)
+            response = yield from self._invoke_timeout_handler(request, flow, deadline)
             yield pb.InvokeExecuteMethodOutput(result=response)
             return
         step = flow.step(request.step_type)
@@ -150,6 +153,7 @@ class WorkerDispatcher:
             loaded_channel_names=request.loaded_channel_names,
             loaded_channel_map_instances=request.loaded_channel_map_instances,
             is_active=is_active,
+            deadline=deadline,
         )
         input = self._values.decode(request.step_input, step.input_codec)
         decision = step.step.execute(context, input)
@@ -184,6 +188,7 @@ class WorkerDispatcher:
         self,
         request: pb.InvokeExecuteMethodRequest,
         flow: _RegisteredFlow,
+        deadline: datetime | None,
     ) -> Generator[pb.InvokeExecuteMethodOutput, None, pb.InvokeExecuteMethodResponse]:
         if request.HasField("step_input"):
             raise InvalidStepResultError(
@@ -212,6 +217,7 @@ class WorkerDispatcher:
             loaded_attribute_map_instances=request.loaded_attribute_map_instances,
             loaded_channel_names=request.loaded_channel_names,
             loaded_channel_map_instances=request.loaded_channel_map_instances,
+            deadline=deadline,
         )
         decision = flow.flow.handle_timeout(context)
         if isgenerator(decision):
@@ -840,7 +846,7 @@ class WorkerDispatcher:
     def _map_cancellation_steps(
         self,
         flow: _RegisteredFlow,
-        steps: tuple[type[Step[Any]], ...],
+        steps: tuple[StepTarget[Any], ...],
     ) -> list[str]:
         step_types: list[str] = []
         seen: set[str] = set()
@@ -896,9 +902,9 @@ class WorkerDispatcher:
     @staticmethod
     def _registered_movement_target(
         flow: _RegisteredFlow,
-        step: type[Step[Any]],
+        step: StepTarget[Any],
     ) -> _RegisteredStep:
-        return flow.step_class(step)
+        return flow.step_target(step)
 
     @staticmethod
     def _require_persistence_identity(