dex-python-sdk 0.2.10__tar.gz → 0.3.1__tar.gz

{dex_python_sdk-0.2.10 → dex_python_sdk-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dex-python-sdk
-Version: 0.2.10
+Version: 0.3.1
 Requires-Dist: grpcio>=1.83.0
 Requires-Dist: grpcio-status>=1.83.0
 Requires-Dist: protobuf>=7.35.1
@@ -49,7 +49,8 @@ def snapshot(self, context: dex.Context) -> dex.RPCResult[Snapshot]:
 Put an AttributeMap or ChannelMap directly in its plural load option when every
 current instance is needed. Use the singular instance options for the less common
 exact-instance case. A selected empty queue returns an empty tuple; reading an
-unselected map entry or pending-message snapshot raises **StateNotLoadedError**.
+unselected map entry raises **AttributeMapNotLoadedError**. An unselected pending-message
+snapshot raises **ChannelMessagesNotLoadedError**.
 Pending messages preserve FIFO order and include the server-assigned message ID.
 The snapshot does not change after the handler stages a publish or deletion.
 

{dex_python_sdk-0.2.10 → dex_python_sdk-0.3.1}/README.md

@@ -33,7 +33,8 @@ def snapshot(self, context: dex.Context) -> dex.RPCResult[Snapshot]:
 Put an AttributeMap or ChannelMap directly in its plural load option when every
 current instance is needed. Use the singular instance options for the less common
 exact-instance case. A selected empty queue returns an empty tuple; reading an
-unselected map entry or pending-message snapshot raises **StateNotLoadedError**.
+unselected map entry raises **AttributeMapNotLoadedError**. An unselected pending-message
+snapshot raises **ChannelMessagesNotLoadedError**.
 Pending messages preserve FIFO order and include the server-assigned message ID.
 The snapshot does not change after the handler stages a publish or deletion.
 

{dex_python_sdk-0.2.10 → dex_python_sdk-0.3.1}/dex/__init__.py

@@ -47,6 +47,7 @@ from dex.flow_info import (
     SearchFlowsPage,
 )
 from dex.flow_options import (
+    FlowTimeoutHandlerOptions,
     FlowTimeoutPolicy,
     IdReusePolicy,
     StartFlowOptions,
@@ -62,6 +63,8 @@ from dex.flow_result import FlowResult, StepCompletion
 from dex.runtime_errors import (
     DexServiceError,
     ErrorSubStatus,
+    AttributeMapNotLoadedError,
+    ChannelMessagesNotLoadedError,
     FlowAlreadyStartedError,
     FlowDefinitionError,
     FlowErrorType,
@@ -71,7 +74,6 @@ from dex.runtime_errors import (
     LongPollTimeoutError,
     ChannelMessageNotFoundError,
     RpcLockConflictError,
-    StateNotLoadedError,
     ValueMappingError,
     WorkerInvocationError,
 )
@@ -119,6 +121,7 @@ __all__ = [
     "AttributeLock",
     "AttributeMap",
     "AttributeMapLoad",
+    "AttributeMapNotLoadedError",
     "AsyncContext",
     "AsyncClient",
     "AsyncBufferedTextStream",
@@ -131,6 +134,7 @@ __all__ = [
     "ChannelMapLoad",
     "ChannelMessage",
     "ChannelMessageNotFoundError",
+    "ChannelMessagesNotLoadedError",
     "Client",
     "ClientOptions",
     "Codec",
@@ -149,6 +153,7 @@ __all__ = [
     "FlowNotActiveError",
     "FlowNotFoundError",
     "FlowStatus",
+    "FlowTimeoutHandlerOptions",
     "FlowTimeoutPolicy",
     "HealthInfo",
     "IdReusePolicy",
@@ -160,7 +165,6 @@ __all__ = [
     "RPCResult",
     "RetryAfterError",
     "RpcLockConflictError",
-    "StateNotLoadedError",
     "Registry",
     "TimeTravelOptions",
     "TimeTravelStepMethod",

{dex_python_sdk-0.2.10 → dex_python_sdk-0.3.1}/dex/_async_value_hydrator.py

@@ -80,12 +80,25 @@ class AsyncValueHydrator:
         if has_heartbeat:
             values.append(request.context.last_heartbeat_value)
         values.extend(entry.value for entry in request.attributes)
+        channel_messages = [
+            message
+            for channel_values in request.loaded_channel_messages.values()
+            for message in channel_values.messages
+        ]
+        values.extend(message.value for message in channel_messages)
         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)
+        for entry in result.attributes:
+            entry.value.CopyFrom(next(hydrated))
+        result_messages = [
+            message
+            for channel_values in result.loaded_channel_messages.values()
+            for message in channel_values.messages
+        ]
+        for message in result_messages:
+            message.value.CopyFrom(next(hydrated))
         return result
 
     async def execute_request(
@@ -101,6 +114,12 @@ class AsyncValueHydrator:
             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)
+        channel_messages = [
+            message
+            for channel_values in request.loaded_channel_messages.values()
+            for message in channel_values.messages
+        ]
+        values.extend(message.value for message in channel_messages)
         for channel_result in request.condition_results.channel_results:
             values.extend(channel_result.values)
         hydrated = iter(await self.hydrate_all(values))
@@ -112,6 +131,13 @@ class AsyncValueHydrator:
             entry.value.CopyFrom(next(hydrated))
         for entry in result.step_exe_locals:
             entry.value.CopyFrom(next(hydrated))
+        result_messages = [
+            message
+            for channel_values in result.loaded_channel_messages.values()
+            for message in channel_values.messages
+        ]
+        for message in result_messages:
+            message.value.CopyFrom(next(hydrated))
         for channel_result in result.condition_results.channel_results:
             for value in channel_result.values:
                 value.CopyFrom(next(hydrated))

{dex_python_sdk-0.2.10 → dex_python_sdk-0.3.1}/dex/_async_worker_dispatcher.py

@@ -100,6 +100,11 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
             request.context,
             self._values,
             request.attributes,
+            channel_infos=dict(request.channel_infos),
+            loaded_channel_messages=dict(request.loaded_channel_messages),
+            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,
             is_active=is_active,
             output_emitter=emitter,
         )
@@ -126,6 +131,7 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
                 upsert_step_exe_locals=list(context.local_writes.values()),
                 record_events=context.events,
                 publish_to_channel=context.publications,
+                delete_from_channel=context.channel_deletions,
             )
             waiting = self._map_wait(flow, wait)
             if waiting is not None:
@@ -153,8 +159,16 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
         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)
+            emitter = _AsyncStepOutputEmitter()
+            handler = asyncio.create_task(
+                self._invoke_timeout_handler_async(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)
             return
         emitter = _AsyncStepOutputEmitter()
         handler = asyncio.create_task(
@@ -187,6 +201,11 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
             request.attributes,
             request.step_exe_locals,
             condition_results,
+            channel_infos=dict(request.channel_infos),
+            loaded_channel_messages=dict(request.loaded_channel_messages),
+            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,
             is_active=is_active,
             output_emitter=emitter,
         )
@@ -214,6 +233,7 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
                 record_events=context.events,
                 upsert_step_exe_locals=list(context.local_writes.values()),
                 publish_to_channel=context.publications,
+                delete_from_channel=context.channel_deletions,
             )
         except InvalidStepResultError as error:
             failure = error
@@ -235,6 +255,8 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
     async def _invoke_timeout_handler_async(
         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)
@@ -260,24 +282,55 @@ class AsyncWorkerDispatcher(WorkerDispatcher):
             request.attributes,
             request.step_exe_locals,
             condition_results,
+            channel_infos=dict(request.channel_infos),
+            loaded_channel_messages=dict(request.loaded_channel_messages),
+            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,
+            is_active=is_active,
+            output_emitter=emitter,
         )
-        decision: Any = flow.flow.handle_timeout(context)
-        if isawaitable(decision):
-            decision = await decision
+        response: pb.InvokeExecuteMethodResponse | None = None
+        failure: BaseException | None = None
+        cause: BaseException | None = None
         try:
+            decision: Any = flow.flow.handle_timeout(context)
+            if isgenerator(decision):
+                decision.close()
+                raise InvalidStepResultError(
+                    flow.name,
+                    _TIMEOUT_HANDLER_STEP_TYPE,
+                    "execute",
+                    "synchronous generators require Worker",
+                )
+            if isawaitable(decision):
+                decision = await decision
             if not isinstance(decision, StepDecision):
                 raise TypeError("handle_timeout must return StepDecision")
-            return pb.InvokeExecuteMethodResponse(
+            response = 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,
+                delete_from_channel=context.channel_deletions,
             )
+        except InvalidStepResultError as error:
+            failure = error
         except (TypeError, ValueError) as error:
-            raise InvalidStepResultError(
+            failure = InvalidStepResultError(
                 flow.name, _TIMEOUT_HANDLER_STEP_TYPE, "execute", str(error)
-            ) from error
+            )
+            cause = error
+        except BaseException as error:
+            failure = error
+        combined = context._finalize_step_outputs(failure)
+        if combined is not None:
+            if cause is not None:
+                raise combined from cause
+            raise combined
+        assert response is not None
+        return response
 
     async def invoke_rpc(  # type: ignore[override]
         self,

{dex_python_sdk-0.2.10 → dex_python_sdk-0.3.1}/dex/_invocation_context.py

@@ -20,7 +20,10 @@ from dex.codec import Codec
 from dex.dexpb import dex_pb2 as pb
 from dex.flow import Registry, _RegisteredFlow
 from dex.flow_result import FlowResult, flow_result_from_proto
-from dex.runtime_errors import StateNotLoadedError
+from dex.runtime_errors import (
+    AttributeMapNotLoadedError,
+    ChannelMessagesNotLoadedError,
+)
 from dex.step import (
     _NO_HEARTBEAT_VALUE,
     StepOutput,
@@ -136,7 +139,12 @@ class InvocationContext:
 
     async def heartbeat(self, value: object = _NO_HEARTBEAT_VALUE) -> None:
         if (
-            self._method not in (InvocationMethod.WAIT_FOR, InvocationMethod.EXECUTE)
+            self._method
+            not in (
+                InvocationMethod.WAIT_FOR,
+                InvocationMethod.EXECUTE,
+                InvocationMethod.TIMEOUT,
+            )
             or self._output_emitter is None
         ):
             raise ValueError("heartbeat requires an asynchronous Step Context")
@@ -223,7 +231,11 @@ class InvocationContext:
         definition: Stream[ValueT],
         value: ValueT,
     ) -> StepOutput | None:
-        if self._method not in (InvocationMethod.WAIT_FOR, InvocationMethod.EXECUTE):
+        if self._method not in (
+            InvocationMethod.WAIT_FOR,
+            InvocationMethod.EXECUTE,
+            InvocationMethod.TIMEOUT,
+        ):
             raise ValueError("Stream writes require a Step Context")
         self._require_registered(definition)
         output = _StreamStepOutput(
@@ -242,7 +254,11 @@ class InvocationContext:
         return None
 
     def _prepare_buffered_stream(self, definition: Stream[object]) -> bool:
-        if self._method not in (InvocationMethod.WAIT_FOR, InvocationMethod.EXECUTE):
+        if self._method not in (
+            InvocationMethod.WAIT_FOR,
+            InvocationMethod.EXECUTE,
+            InvocationMethod.TIMEOUT,
+        ):
             raise ValueError("Buffered Streams require a Step Context")
         self._require_registered(definition)
         return self._output_emitter is not None
@@ -354,11 +370,10 @@ class InvocationContext:
                 ),
             )
         )
-        if self._method is InvocationMethod.RPC:
-            current = self._channel_infos.get(name)
-            self._channel_infos[name] = pb.ChannelInfo(
-                size=(current.size if current is not None else 0) + 1
-            )
+        current = self._channel_infos.get(name)
+        self._channel_infos[name] = pb.ChannelInfo(
+            size=(current.size if current is not None else 0) + 1
+        )
 
     def _delete_channel_message(
         self,
@@ -366,8 +381,6 @@ class InvocationContext:
         instance: str | None,
         message_id: str,
     ) -> None:
-        if self._method is not InvocationMethod.RPC:
-            raise ValueError("Channel message deletion requires an RPC Context")
         self._require_registered(definition)
         name = self._physical_name(definition, instance)
         self.channel_deletions.append(
@@ -407,8 +420,6 @@ class InvocationContext:
         definition: ChannelMap[object],
     ) -> tuple[str, ...]:
         self._require_registered(definition)
-        if self._method is not InvocationMethod.RPC:
-            raise ValueError("ChannelMap introspection requires an RPC invocation")
         prefix = f"{definition.name}/"
         return tuple(
             sorted(
@@ -432,8 +443,6 @@ class InvocationContext:
         definition: Channel[ValueT] | ChannelMap[ValueT],
         instance: str | None,
     ) -> tuple[ChannelMessage[ValueT], ...]:
-        if self._method is not InvocationMethod.RPC:
-            raise ValueError("pending Channel messages require an RPC Context")
         self._require_registered(definition)
         if isinstance(definition, ChannelMap):
             channel_name = self._physical_name(definition, instance)
@@ -445,8 +454,8 @@ class InvocationContext:
             channel_name = self._physical_name(definition, instance)
             is_loaded = definition.name in self._loaded_channel_names
         if not is_loaded:
-            raise StateNotLoadedError(
-                f"Channel messages were not loaded for RPC: {definition.name}"
+            raise ChannelMessagesNotLoadedError(
+                f"Channel messages were not loaded for this invocation: {definition.name}"
             )
         values = self._loaded_channel_messages.get(channel_name)
         if values is None:
@@ -486,29 +495,24 @@ class InvocationContext:
         definition: Attribute[Any] | AttributeMap[Any],
         instance: str | None,
     ) -> None:
-        if self._method is not InvocationMethod.RPC or not isinstance(
-            definition, AttributeMap
-        ):
+        if not isinstance(definition, AttributeMap):
             return
         physical_name = self._physical_name(definition, instance)
         if (
             f"{definition.name}/" not in self._loaded_attribute_map_instances
             and physical_name not in self._loaded_attribute_map_instances
         ):
-            raise StateNotLoadedError(
-                f"AttributeMap instance was not loaded for RPC: {physical_name}"
+            raise AttributeMapNotLoadedError(
+                f"AttributeMap instance was not loaded for this invocation: {physical_name}"
             )
 
     def _require_attribute_map_all_loaded(
         self,
         definition: AttributeMap[object],
     ) -> None:
-        if (
-            self._method is InvocationMethod.RPC
-            and f"{definition.name}/" not in self._loaded_attribute_map_instances
-        ):
-            raise StateNotLoadedError(
-                f"all AttributeMap instances were not loaded for RPC: {definition.name}"
+        if f"{definition.name}/" not in self._loaded_attribute_map_instances:
+            raise AttributeMapNotLoadedError(
+                f"all AttributeMap instances were not loaded for this invocation: {definition.name}"
             )
 
     def _require_registered(self, definition: _Definition) -> None:

{dex_python_sdk-0.2.10 → dex_python_sdk-0.3.1}/dex/_value_hydrator.py

@@ -80,12 +80,25 @@ class ValueHydrator:
         if has_heartbeat:
             values.append(request.context.last_heartbeat_value)
         values.extend(entry.value for entry in request.attributes)
+        channel_messages = [
+            message
+            for channel_values in request.loaded_channel_messages.values()
+            for message in channel_values.messages
+        ]
+        values.extend(message.value for message in channel_messages)
         hydrated = iter(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)
+        for entry in result.attributes:
+            entry.value.CopyFrom(next(hydrated))
+        result_messages = [
+            message
+            for channel_values in result.loaded_channel_messages.values()
+            for message in channel_values.messages
+        ]
+        for message in result_messages:
+            message.value.CopyFrom(next(hydrated))
         return result
 
     def execute_request(
@@ -101,6 +114,12 @@ class ValueHydrator:
             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)
+        channel_messages = [
+            message
+            for channel_values in request.loaded_channel_messages.values()
+            for message in channel_values.messages
+        ]
+        values.extend(message.value for message in channel_messages)
         for channel_result in request.condition_results.channel_results:
             values.extend(channel_result.values)
         for flow_result in request.condition_results.sub_flow_results:
@@ -116,6 +135,13 @@ class ValueHydrator:
             entry.value.CopyFrom(next(hydrated))
         for entry in result.step_exe_locals:
             entry.value.CopyFrom(next(hydrated))
+        result_messages = [
+            message
+            for channel_values in result.loaded_channel_messages.values()
+            for message in channel_values.messages
+        ]
+        for message in result_messages:
+            message.value.CopyFrom(next(hydrated))
         for channel_result in result.condition_results.channel_results:
             for value in channel_result.values:
                 value.CopyFrom(next(hydrated))