dex-python-sdk 0.1.3__tar.gz → 0.1.5__tar.gz

{dex_python_sdk-0.1.3 → dex_python_sdk-0.1.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dex-python-sdk
-Version: 0.1.3
+Version: 0.1.5
 Requires-Dist: grpcio>=1.83.0
 Requires-Dist: grpcio-status>=1.83.0
 Requires-Dist: protobuf>=7.35.1
@@ -71,7 +71,7 @@ registry = dex.Registry((flow,))
 ```
 
 Registry derives codecs from declared Python types and handler annotations.
-Built-in scalar types and dataclasses need no codec arguments. Register an
+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 and channels together and
 partitions them by definition type.
@@ -131,6 +131,71 @@ Applications implement two generic interfaces from [`dex`](dex/):
 `StepOptions.wait_for_method_timeout` and `execute_method_timeout` bound the
 two handler calls. Timer and channel conditions determine how long a Step waits.
 
+`wait_for_retry` and `execute_retry` limit one logical handler execution. With
+`StepDurability.ASYNC`, local and fallback regular activities share maximum
+attempts, total duration, and 1-based attempt numbers. Fallback starts
+immediately; later regular retries continue the backoff sequence at the
+cumulative attempt.
+
+### Canceling Step executions
+
+A successful Step can cancel queued or active executions while continuing with
+its normal decision:
+
+```python
+return (
+    dex.go_to(self.record_quote, quote)
+    .with_canceling_sibling_steps(self.carrier_a, self.carrier_b)
+    .with_canceling_steps(self.global_quote_timeout)
+)
+```
+
+`with_canceling_steps` selects every current execution of each registered Step
+type. `with_canceling_sibling_steps` selects only executions with the same
+`Context.from_step_execution_id` as the current execution. Decisions are
+immutable; repeated calls form a union, and Flow-wide selection wins for the
+same Step type. Unregistered selectors produce an invalid Step result.
+
+Dex resolves one snapshot after the current execution succeeds. Completed,
+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.
+
+### Soft Flow timeout
+
+Override `Flow.handle_timeout` to make a positive timeout use handler policy
+by default. Both synchronous and async Workers support the hook:
+
+```python
+class Orders(dex.Flow[str]):
+    async def handle_timeout(self, context: dex.Context) -> dex.StepDecision:
+        await notify_expiration(context)
+        return dex.force_complete("expired")
+
+options = dex.StartFlowOptions(
+    timeout=timedelta(minutes=30),
+    timeout_policy=dex.FlowTimeoutPolicy.HANDLER,
+)
+```
+
+Register async hooks with `allow_async_handlers=True` and run them with
+`AsyncWorker`. `FAIL` produces `FlowErrorType.FLOW_TIMEOUT` and permits Flow
+retry; `CANCEL` cancels without retry. Continue-as-new preserves the deadline,
+while retry and cron runs receive a fresh budget. A zero or absent timeout
+disables the feature.
+
 `Registry` validates every Flow, Step, RPC signature, durable name, lock, and
 codec before Client or Worker startup. `Client` methods use these typed objects
 instead of raw Flow, Step, or RPC strings.
@@ -141,18 +206,54 @@ instead of raw Flow, Step, or RPC strings.
 `Wait.any_combination_of` must have a non-empty user ID; the same Condition
 instance may appear in multiple combinations.
 
-Both `Client` and `AsyncClient` provide `wait_for_attribute_equal` and
-`wait_for_attribute_map_equal`. They target the current run and accept only
-string, bool, int, or float wire values. JSON objects, bytes, and null fail
+Both `Client` and `AsyncClient` provide singleton and AttributeMap-instance
+overloads of `wait_for_attribute_equal`. They target the current run and accept
+only string, bool, int, or float wire values. JSON objects, bytes, and null fail
 before transport. `AttributeMap.get_map_size/get_all_instance_keys` include
 buffered sets and deletes. The matching `ChannelMap` methods are RPC-only,
 include buffered publishes, and omit empty instances. Keys are decoded and
 sorted. Use `force_complete_if_channels_empty(...)` for conditional completion.
 
+`Client.wait_for_flow` and `AsyncClient.wait_for_flow` return a
+`FlowResult` after hydrating every output-bearing completion. Use
+`single_output` only when the Flow contract produces exactly one output:
+
+```python
+output = client.wait_for_flow(flow_id).single_output(OrderResult)
+
+result = client.wait_for_flow(flow_id)
+for completion in result.completions:
+    if completion.step_execution_id == expected_execution_id:
+        output = completion.decode(OrderResult)
+```
+
+`completions` is an immutable tuple in server collection order. Parallel branch
+order is not deterministic, so select by `step_type` or `step_execution_id`.
+No-output Flows return an empty tuple; `single_output` raises `ValueError` for
+zero or multiple completions. Every terminal status returns a `FlowResult`; inspect
+`status`, `error_type`, and `error_message` for unsuccessful completion.
+
+SubFlows are normal, independently addressable Flows used as durable Conditions:
+
+```python
+def wait_for(self, context: Context, input: ChargeInput) -> Wait:
+    return Wait.until(SubFlow.run(self.charge_flow, input))
+
+def execute(self, context: Context, input: ChargeInput) -> StepDecision:
+    del input
+    receipt = SubFlow.get_condition_results(context).single_output(Receipt)
+    return graceful_complete(receipt)
+```
+
+`SubFlow.get_flow_id(context, index=0)` remains available for a running `any_of`
+loser. `SubFlowOptions` configures timing, timeout policy, retry, initial target
+Attributes, Flow config, Condition ID, and reuse. Parent completion does not cancel an unfinished
+SubFlow.
+
 ### Errors
 
 Client calls raise concrete `DexServiceError` subclasses. Existing-Flow reads
-(`get_attribute`, `describe_flow`, `wait_for_flow`, and `reset_flow`) raise
+(`get_attribute`, `describe_flow`, `wait_for_flow`, and `time_travel`) raise
 `FlowNotFoundError` when the Flow does not exist. Mutations, RPCs, timer/Step
 waits, config updates, and continue-as-new triggers raise
 `FlowNotActiveError` when no running Flow can accept the operation.
@@ -199,22 +300,17 @@ 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).
 
-## Running dex-server locally
+## Running Dex locally
 
-### Option 1: use docker compose
-See [dex README](https://github.com/superdurable/dex#using-docker-image--docker-compose)
+Install and start the complete local environment with `dexcli`:
 
-### Option 2: VSCode Dev Container
+```bash
+brew install superdurable/tap/dexcli
+dexcli dev
+```
 
-Dev Container is an easy way to get dex-server running locally. Follow these steps to launch a dev container:
-- Install Docker, VSCode, and [VSCode Dev Container plugin](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-containers).
-- Open the project in VSCode.
-    ```bash
-    cd dex-python-sdk
-    code .
-    ```
-- Launch the Remote-Containers: Reopen in Container command from Command Palette (Ctrl + Shift + P). You can also click in the bottom left corner to access the remote container menu.
-- Once the dev container starts, dex-server will be listening on port 8801.
+Dex Server listens on `127.0.0.1:8801`. See the
+[CLI README](../cli/README.md) for endpoints and persistence options.
 
 ## How To Contribute
 

{dex_python_sdk-0.1.3 → dex_python_sdk-0.1.5}/README.md

@@ -55,7 +55,7 @@ registry = dex.Registry((flow,))
 ```
 
 Registry derives codecs from declared Python types and handler annotations.
-Built-in scalar types and dataclasses need no codec arguments. Register an
+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 and channels together and
 partitions them by definition type.
@@ -115,6 +115,71 @@ Applications implement two generic interfaces from [`dex`](dex/):
 `StepOptions.wait_for_method_timeout` and `execute_method_timeout` bound the
 two handler calls. Timer and channel conditions determine how long a Step waits.
 
+`wait_for_retry` and `execute_retry` limit one logical handler execution. With
+`StepDurability.ASYNC`, local and fallback regular activities share maximum
+attempts, total duration, and 1-based attempt numbers. Fallback starts
+immediately; later regular retries continue the backoff sequence at the
+cumulative attempt.
+
+### Canceling Step executions
+
+A successful Step can cancel queued or active executions while continuing with
+its normal decision:
+
+```python
+return (
+    dex.go_to(self.record_quote, quote)
+    .with_canceling_sibling_steps(self.carrier_a, self.carrier_b)
+    .with_canceling_steps(self.global_quote_timeout)
+)
+```
+
+`with_canceling_steps` selects every current execution of each registered Step
+type. `with_canceling_sibling_steps` selects only executions with the same
+`Context.from_step_execution_id` as the current execution. Decisions are
+immutable; repeated calls form a union, and Flow-wide selection wins for the
+same Step type. Unregistered selectors produce an invalid Step result.
+
+Dex resolves one snapshot after the current execution succeeds. Completed,
+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.
+
+### Soft Flow timeout
+
+Override `Flow.handle_timeout` to make a positive timeout use handler policy
+by default. Both synchronous and async Workers support the hook:
+
+```python
+class Orders(dex.Flow[str]):
+    async def handle_timeout(self, context: dex.Context) -> dex.StepDecision:
+        await notify_expiration(context)
+        return dex.force_complete("expired")
+
+options = dex.StartFlowOptions(
+    timeout=timedelta(minutes=30),
+    timeout_policy=dex.FlowTimeoutPolicy.HANDLER,
+)
+```
+
+Register async hooks with `allow_async_handlers=True` and run them with
+`AsyncWorker`. `FAIL` produces `FlowErrorType.FLOW_TIMEOUT` and permits Flow
+retry; `CANCEL` cancels without retry. Continue-as-new preserves the deadline,
+while retry and cron runs receive a fresh budget. A zero or absent timeout
+disables the feature.
+
 `Registry` validates every Flow, Step, RPC signature, durable name, lock, and
 codec before Client or Worker startup. `Client` methods use these typed objects
 instead of raw Flow, Step, or RPC strings.
@@ -125,18 +190,54 @@ instead of raw Flow, Step, or RPC strings.
 `Wait.any_combination_of` must have a non-empty user ID; the same Condition
 instance may appear in multiple combinations.
 
-Both `Client` and `AsyncClient` provide `wait_for_attribute_equal` and
-`wait_for_attribute_map_equal`. They target the current run and accept only
-string, bool, int, or float wire values. JSON objects, bytes, and null fail
+Both `Client` and `AsyncClient` provide singleton and AttributeMap-instance
+overloads of `wait_for_attribute_equal`. They target the current run and accept
+only string, bool, int, or float wire values. JSON objects, bytes, and null fail
 before transport. `AttributeMap.get_map_size/get_all_instance_keys` include
 buffered sets and deletes. The matching `ChannelMap` methods are RPC-only,
 include buffered publishes, and omit empty instances. Keys are decoded and
 sorted. Use `force_complete_if_channels_empty(...)` for conditional completion.
 
+`Client.wait_for_flow` and `AsyncClient.wait_for_flow` return a
+`FlowResult` after hydrating every output-bearing completion. Use
+`single_output` only when the Flow contract produces exactly one output:
+
+```python
+output = client.wait_for_flow(flow_id).single_output(OrderResult)
+
+result = client.wait_for_flow(flow_id)
+for completion in result.completions:
+    if completion.step_execution_id == expected_execution_id:
+        output = completion.decode(OrderResult)
+```
+
+`completions` is an immutable tuple in server collection order. Parallel branch
+order is not deterministic, so select by `step_type` or `step_execution_id`.
+No-output Flows return an empty tuple; `single_output` raises `ValueError` for
+zero or multiple completions. Every terminal status returns a `FlowResult`; inspect
+`status`, `error_type`, and `error_message` for unsuccessful completion.
+
+SubFlows are normal, independently addressable Flows used as durable Conditions:
+
+```python
+def wait_for(self, context: Context, input: ChargeInput) -> Wait:
+    return Wait.until(SubFlow.run(self.charge_flow, input))
+
+def execute(self, context: Context, input: ChargeInput) -> StepDecision:
+    del input
+    receipt = SubFlow.get_condition_results(context).single_output(Receipt)
+    return graceful_complete(receipt)
+```
+
+`SubFlow.get_flow_id(context, index=0)` remains available for a running `any_of`
+loser. `SubFlowOptions` configures timing, timeout policy, retry, initial target
+Attributes, Flow config, Condition ID, and reuse. Parent completion does not cancel an unfinished
+SubFlow.
+
 ### Errors
 
 Client calls raise concrete `DexServiceError` subclasses. Existing-Flow reads
-(`get_attribute`, `describe_flow`, `wait_for_flow`, and `reset_flow`) raise
+(`get_attribute`, `describe_flow`, `wait_for_flow`, and `time_travel`) raise
 `FlowNotFoundError` when the Flow does not exist. Mutations, RPCs, timer/Step
 waits, config updates, and continue-as-new triggers raise
 `FlowNotActiveError` when no running Flow can accept the operation.
@@ -183,22 +284,17 @@ 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).
 
-## Running dex-server locally
+## Running Dex locally
 
-### Option 1: use docker compose
-See [dex README](https://github.com/superdurable/dex#using-docker-image--docker-compose)
+Install and start the complete local environment with `dexcli`:
 
-### Option 2: VSCode Dev Container
+```bash
+brew install superdurable/tap/dexcli
+dexcli dev
+```
 
-Dev Container is an easy way to get dex-server running locally. Follow these steps to launch a dev container:
-- Install Docker, VSCode, and [VSCode Dev Container plugin](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-containers).
-- Open the project in VSCode.
-    ```bash
-    cd dex-python-sdk
-    code .
-    ```
-- Launch the Remote-Containers: Reopen in Container command from Command Palette (Ctrl + Shift + P). You can also click in the bottom left corner to access the remote container menu.
-- Once the dev container starts, dex-server will be listening on port 8801.
+Dex Server listens on `127.0.0.1:8801`. See the
+[CLI README](../cli/README.md) for endpoints and persistence options.
 
 ## How To Contribute
 

{dex_python_sdk-0.1.3 → dex_python_sdk-0.1.5}/dex/__init__.py

@@ -8,6 +8,8 @@
 # Legacy Materials remain under their original licenses.
 # See LICENSE and LEGACY_NOTICES.md.
 
+from dex.async_client import AsyncClient
+from dex.async_worker import AsyncWorker
 from dex.attribute import (
     Attribute,
     AttributeIndex,
@@ -17,8 +19,6 @@ from dex.attribute import (
 )
 from dex.blob_cache import BlobCache, BlobCacheConfig, open_blob_cache
 from dex.channel import Channel, ChannelMap
-from dex.async_client import AsyncClient
-from dex.async_worker import AsyncWorker
 from dex.client import Client
 from dex.client_options import ClientOptions
 from dex.codec import (
@@ -45,13 +45,18 @@ from dex.flow_info import (
     SearchFlowsPage,
 )
 from dex.flow_options import (
+    FlowTimeoutPolicy,
     IdReusePolicy,
-    ResetFlowOptions,
-    ResetType,
+    TimeTravelOptions,
+    TimeTravelStepMethod,
+    TimeTravelType,
     StartFlowOptions,
+    SubFlowOptions,
+    SubFlowReusePolicy,
     StopFlowOptions,
     StopType,
 )
+from dex.flow_result import FlowResult, StepCompletion
 from dex.runtime_errors import (
     DexServiceError,
     ErrorSubStatus,
@@ -60,7 +65,6 @@ from dex.runtime_errors import (
     FlowErrorType,
     FlowNotActiveError,
     FlowNotFoundError,
-    FlowUncompletedError,
     InvalidStepResultError,
     LongPollTimeoutError,
     RpcLockConflictError,
@@ -86,6 +90,7 @@ from dex.step import (
 )
 from dex.step_execution import StepExecutionId, TimerId
 from dex.timer import Timer
+from dex.subflow import SubFlow
 from dex.wait import Wait
 from dex.worker import Worker
 from dex.worker_options import WorkerOptions, WorkerTarget
@@ -121,10 +126,11 @@ __all__ = [
     "FlowDefinitionError",
     "FlowErrorType",
     "FlowInfo",
+    "FlowResult",
     "FlowNotActiveError",
     "FlowNotFoundError",
     "FlowStatus",
-    "FlowUncompletedError",
+    "FlowTimeoutPolicy",
     "HealthInfo",
     "IdReusePolicy",
     "InvalidStepResultError",
@@ -135,12 +141,17 @@ __all__ = [
     "RPCResult",
     "RpcLockConflictError",
     "Registry",
-    "ResetFlowOptions",
-    "ResetType",
+    "TimeTravelOptions",
+    "TimeTravelStepMethod",
+    "TimeTravelType",
     "RetryPolicy",
     "SearchFlowEntry",
     "SearchFlowsPage",
+    "StepCompletion",
     "StartFlowOptions",
+    "SubFlow",
+    "SubFlowOptions",
+    "SubFlowReusePolicy",
     "StepExecutionId",
     "StepDecision",
     "Step",

{dex_python_sdk-0.1.3 → dex_python_sdk-0.1.5}/dex/_async_value_hydrator.py

@@ -88,13 +88,15 @@ class AsyncValueHydrator:
     ) -> pb.InvokeExecuteMethodRequest:
         result = pb.InvokeExecuteMethodRequest()
         result.CopyFrom(request)
-        values = [request.step_input]
+        has_step_input = request.HasField("step_input")
+        values = [request.step_input] if has_step_input else []
         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:
             values.extend(channel_result.values)
         hydrated = iter(await self.hydrate_all(values))
-        result.step_input.CopyFrom(next(hydrated))
+        if has_step_input:
+            result.step_input.CopyFrom(next(hydrated))
         for entry in result.attributes:
             entry.value.CopyFrom(next(hydrated))
         for entry in result.step_exe_locals:

{dex_python_sdk-0.1.3 → dex_python_sdk-0.1.5}/dex/_async_worker_dispatcher.py

@@ -9,14 +9,14 @@
 from __future__ import annotations
 
 from inspect import isawaitable
-from typing import Any
+from typing import Any, Callable
 
 from dex._async_value_hydrator import AsyncValueHydrator
 from dex._invocation_context import InvocationContext, InvocationMethod
 from dex._value_mapper import ValueMapper
-from dex._worker_dispatcher import WorkerDispatcher
+from dex._worker_dispatcher import _TIMEOUT_HANDLER_STEP_TYPE, WorkerDispatcher
 from dex.dexpb import dex_pb2 as pb
-from dex.flow import RPCResult, Registry
+from dex.flow import Registry, RPCResult
 from dex.runtime_errors import InvalidStepResultError, ValueMappingError
 from dex.step import StepDecision
 from dex.wait import Wait
@@ -36,6 +36,7 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
     async def invoke_wait_for(  # type: ignore[override]
         self,
         original: pb.InvokeWaitForMethodRequest,
+        is_active: Callable[[], bool] | None = None,
     ) -> pb.InvokeWaitForMethodResponse:
         request = await self._async_hydrator.wait_for_request(original)
         flow = self._registry._flow_by_type(request.flow_type)
@@ -46,6 +47,7 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
             request.context,
             self._values,
             request.attributes,
+            is_active=is_active,
         )
         input = self._values.decode(request.step_input, step.input_codec)
         wait = step.step.wait_for(context, input)
@@ -72,9 +74,12 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
     async def invoke_execute(  # type: ignore[override]
         self,
         original: pb.InvokeExecuteMethodRequest,
+        is_active: Callable[[], bool] | None = 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
@@ -87,6 +92,7 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
             request.attributes,
             request.step_exe_locals,
             condition_results,
+            is_active=is_active,
         )
         input = self._values.decode(request.step_input, step.input_codec)
         decision: Any = step.step.execute(context, input)
@@ -107,9 +113,56 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
                 flow.name, step.name, "execute", str(error)
             ) from error
 
+    async def _invoke_timeout_handler_async(
+        self,
+        request: pb.InvokeExecuteMethodRequest,
+        flow: Any,
+    ) -> pb.InvokeExecuteMethodResponse:
+        if request.HasField("step_input"):
+            raise InvalidStepResultError(
+                flow.name, _TIMEOUT_HANDLER_STEP_TYPE, "execute", "input must be absent"
+            )
+        if not flow.has_timeout_handler:
+            raise InvalidStepResultError(
+                flow.name,
+                _TIMEOUT_HANDLER_STEP_TYPE,
+                "execute",
+                "handler is not registered",
+            )
+        condition_results = (
+            request.condition_results if request.HasField("condition_results") else None
+        )
+        context = InvocationContext(
+            InvocationMethod.EXECUTE,
+            flow,
+            request.context,
+            self._values,
+            request.attributes,
+            request.step_exe_locals,
+            condition_results,
+        )
+        decision: Any = flow.flow.handle_timeout(context)
+        if isawaitable(decision):
+            decision = await decision
+        try:
+            if not isinstance(decision, StepDecision):
+                raise TypeError("handle_timeout must return StepDecision")
+            return pb.InvokeExecuteMethodResponse(
+                step_decision=self._map_decision(flow, decision),
+                upsert_attributes=list(context.attribute_writes.values()),
+                record_events=context.events,
+                upsert_step_exe_locals=list(context.local_writes.values()),
+                publish_to_channel=context.publications,
+            )
+        except (TypeError, ValueError) as error:
+            raise InvalidStepResultError(
+                flow.name, _TIMEOUT_HANDLER_STEP_TYPE, "execute", str(error)
+            ) from error
+
     async def invoke_rpc(  # type: ignore[override]
         self,
         original: pb.InvokeWorkerRPCRequest,
+        is_active: Callable[[], bool] | None = None,
     ) -> pb.InvokeWorkerRPCResponse:
         request = await self._async_hydrator.rpc_request(original)
         flow = self._registry._flow_by_type(request.flow_type)
@@ -121,6 +174,7 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
             self._values,
             request.attributes,
             channel_infos=dict(request.channel_infos),
+            is_active=is_active,
         )
         arguments: list[object] = [context]
         if rpc.input_codec is not None:
@@ -144,6 +198,11 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
                     response.step_decision.next_steps.extend(
                         self._map_movements(flow, returned.next_steps)
                     )
+                response.step_decision.cancel_step_types.extend(
+                    self._map_cancellation_steps(flow, returned.canceling_steps)
+                )
+                if not returned.next_steps and not returned.canceling_steps:
+                    response.ClearField("step_decision")
             elif returned is None and rpc.output_codec is None:
                 response.output.CopyFrom(self._values.encode_dynamic(None))
             else:

{dex_python_sdk-0.1.3 → dex_python_sdk-0.1.5}/dex/_async_worker_service.py

@@ -9,8 +9,9 @@
 from __future__ import annotations
 
 import logging
+from asyncio import CancelledError, current_task
 from collections.abc import Awaitable, Callable
-from typing import TypeVar
+from typing import Any, TypeVar
 
 import grpc
 
@@ -27,39 +28,71 @@ class AsyncWorkerService(dex_pb2_grpc.WorkerServiceServicer):
     def __init__(self, dispatcher: AsyncWorkerDispatcher) -> None:
         self._dispatcher = dispatcher
 
-    async def InvokeWaitForMethod(  # type: ignore[override]
+    async def InvokeWaitForMethod(
         self,
         request: pb.InvokeWaitForMethodRequest,
-        context: grpc.aio.ServicerContext,
+        context: grpc.aio.ServicerContext[
+            pb.InvokeWaitForMethodRequest,
+            pb.InvokeWaitForMethodResponse,
+        ],
     ) -> pb.InvokeWaitForMethodResponse:
         return await self._invoke(
-            context, lambda: self._dispatcher.invoke_wait_for(request)
+            context,
+            lambda: self._dispatcher.invoke_wait_for(
+                request, lambda: self._is_active(context)
+            ),
         )
 
-    async def InvokeExecuteMethod(  # type: ignore[override]
+    async def InvokeExecuteMethod(
         self,
         request: pb.InvokeExecuteMethodRequest,
-        context: grpc.aio.ServicerContext,
+        context: grpc.aio.ServicerContext[
+            pb.InvokeExecuteMethodRequest,
+            pb.InvokeExecuteMethodResponse,
+        ],
     ) -> pb.InvokeExecuteMethodResponse:
         return await self._invoke(
-            context, lambda: self._dispatcher.invoke_execute(request)
+            context,
+            lambda: self._dispatcher.invoke_execute(
+                request, lambda: self._is_active(context)
+            ),
         )
 
-    async def InvokeWorkerRPC(  # type: ignore[override]
+    async def InvokeWorkerRPC(
         self,
         request: pb.InvokeWorkerRPCRequest,
-        context: grpc.aio.ServicerContext,
+        context: grpc.aio.ServicerContext[
+            pb.InvokeWorkerRPCRequest,
+            pb.InvokeWorkerRPCResponse,
+        ],
     ) -> pb.InvokeWorkerRPCResponse:
-        return await self._invoke(context, lambda: self._dispatcher.invoke_rpc(request))
+        return await self._invoke(
+            context,
+            lambda: self._dispatcher.invoke_rpc(
+                request, lambda: self._is_active(context)
+            ),
+        )
+
+    @staticmethod
+    def _is_active(context: grpc.aio.ServicerContext[Any, Any]) -> bool:
+        task = current_task()
+        return not context.cancelled() and (task is None or task.cancelling() == 0)
 
     @staticmethod
     async def _invoke(
-        context: grpc.aio.ServicerContext,
+        context: grpc.aio.ServicerContext[Any, Any],
         invocation: Callable[[], Awaitable[ResponseT]],
     ) -> ResponseT:
         try:
             return await invocation()
+        except CancelledError:
+            raise
         except BaseException as error:
+            if not AsyncWorkerService._is_active(context):
+                await context.abort(
+                    grpc.StatusCode.CANCELLED,
+                    "Python AsyncWorker invocation canceled",
+                )
             _LOGGER.exception("Python AsyncWorker invocation failed")
             await async_abort_worker_error(context, error)
             raise RuntimeError("gRPC abort returned unexpectedly") from error