WIP: WAIT_REQUESTED

Author: dandavison

tests/nexus/test_workflow_caller_cancellation_types.py

@@ -41,11 +41,28 @@ class TestContext:
 test_context: TestContext
 
 
-@workflow.defn
+@workflow.defn(sandboxed=False)
 class HandlerWorkflow:
+    def __init__(self):
+        self.caller_op_future_resolved = asyncio.Event()
+
     @workflow.run
     async def run(self) -> None:
-        await asyncio.Future()
+        try:
+            await asyncio.Future()
+        except asyncio.CancelledError:
+            if (
+                test_context.cancellation_type
+                == workflow.NexusOperationCancellationType.WAIT_REQUESTED
+            ):
+                # We want to prove that the caller op future can be resolved before the operation
+                # (i.e. its backing workflow) is cancelled.
+                await self.caller_op_future_resolved.wait()
+            raise
+
+    @workflow.signal
+    def set_caller_op_future_resolved(self) -> None:
+        self.caller_op_future_resolved.set()
 
 
 @nexusrpc.service
@@ -137,6 +154,7 @@ async def run(self, input: Input) -> CancellationResult:
             else self.nexus_client.start_operation(Service.workflow_op, input=None)
         )
         self.operation_token = op_handle.operation_token
+        assert self.operation_token
         # Request cancellation of the asyncio task representing the nexus operation. When the handle
         # task is awaited, the resulting asyncio.CancelledError is caught, and a
         # RequestCancelNexusOperation command is emitted instead (see
@@ -167,6 +185,17 @@ async def run(self, input: Input) -> CancellationResult:
             # Block the thread for a bit to avoid flakiness in asserting that the future is
             # unblocked before the WFT completion is sent in the TryCancel case.
             time.sleep(1e-2)
+
+            # Notify handler workflow that caller op future has been resolved
+            await workflow.get_external_workflow_handle_for(
+                HandlerWorkflow.run,
+                workflow_id=(
+                    nexus.WorkflowHandle[None]
+                    .from_token(self.operation_token)
+                    .workflow_id
+                ),
+            ).signal(HandlerWorkflow.set_caller_op_future_resolved)
+
             await workflow.wait_condition(lambda: self.released)
             return CancellationResult(
                 operation_token=op_handle.operation_token,