Separate Temporal context for each operation verb

Author: dandavison

temporalio/nexus/__init__.py

@@ -4,10 +4,10 @@
     WorkflowRunOperationContext as WorkflowRunOperationContext,
 )
 from ._operation_context import (
-    _temporal_operation_context as _temporal_operation_context,
+    _TemporalCancelOperationContext as _TemporalCancelOperationContext,
 )
 from ._operation_context import (
-    _TemporalNexusOperationContext as _TemporalNexusOperationContext,
+    _TemporalStartOperationContext as _TemporalStartOperationContext,
 )
 from ._operation_context import client as client
 from ._operation_context import info as info

temporalio/nexus/_decorators.py

@@ -16,7 +16,10 @@
     StartOperationContext,
 )
 
-from temporalio.nexus._operation_context import WorkflowRunOperationContext
+from temporalio.nexus._operation_context import (
+    WorkflowRunOperationContext,
+    _TemporalStartOperationContext,
+)
 from temporalio.nexus._operation_handlers import (
     WorkflowRunOperationHandler,
 )
@@ -112,7 +115,8 @@ def operation_handler_factory(
             async def _start(
                 ctx: StartOperationContext, input: InputT
             ) -> WorkflowHandle[OutputT]:
-                return await start(self, WorkflowRunOperationContext(ctx), input)
+                tctx = _TemporalStartOperationContext.get()
+                return await start(self, WorkflowRunOperationContext(tctx), input)
 
             _start.__doc__ = start.__doc__
             return WorkflowRunOperationHandler(_start, input_type, output_type)

temporalio/nexus/_operation_context.py

@@ -10,6 +10,7 @@
     Any,
     Callable,
     Mapping,
+    MutableMapping,
     Optional,
     Sequence,
     Union,
@@ -30,9 +31,16 @@
     SelfType,
 )
 
+# The Temporal Nexus worker always builds a nexusrpc StartOperationContext or
+# CancelOperationContext and passes it as the first parameter to the nexusrpc operation
+# handler. In addition, it sets one of the following context vars.
 
-_temporal_operation_context: ContextVar[_TemporalNexusOperationContext] = ContextVar(
-    "temporal-operation-context"
+_temporal_start_operation_context: ContextVar[_TemporalStartOperationContext] = (
+    ContextVar("temporal-start-operation-context")
+)
+
+_temporal_cancel_operation_context: ContextVar[_TemporalCancelOperationContext] = (
+    ContextVar("temporal-cancel-operation-context")
 )
 
 
@@ -51,59 +59,126 @@ def info() -> Info:
     """
     Get the current Nexus operation information.
     """
-    return _TemporalNexusOperationContext.get().info()
+    return _temporal_context().info()
 
 
 def client() -> temporalio.client.Client:
     """
     Get the Temporal client used by the worker handling the current Nexus operation.
     """
-    return _TemporalNexusOperationContext.get().client
+    return _temporal_context().client
+
+
+def _temporal_context() -> (
+    Union[_TemporalStartOperationContext, _TemporalCancelOperationContext]
+):
+    ctx = _try_temporal_context()
+    if ctx is None:
+        raise RuntimeError("Not in Nexus operation context.")
+    return ctx
+
+
+def _try_temporal_context() -> (
+    Optional[Union[_TemporalStartOperationContext, _TemporalCancelOperationContext]]
+):
+    start_ctx = _temporal_start_operation_context.get(None)
+    cancel_ctx = _temporal_cancel_operation_context.get(None)
+    if start_ctx and cancel_ctx:
+        raise RuntimeError("Cannot be in both start and cancel operation contexts.")
+    return start_ctx or cancel_ctx
 
 
 @dataclass
-class _TemporalNexusOperationContext:
+class _TemporalStartOperationContext:
     """
-    Context for a Nexus operation being handled by a Temporal Nexus Worker.
+    Context for a Nexus start operation being handled by a Temporal Nexus Worker.
     """
 
-    info: Callable[[], Info]
-    """Information about the running Nexus operation."""
+    nexus_context: StartOperationContext
+    """Nexus-specific start operation context."""
 
-    nexus_operation_context: Union[StartOperationContext, CancelOperationContext]
+    info: Callable[[], Info]
+    """Temporal information about the running Nexus operation."""
 
     client: temporalio.client.Client
     """The Temporal client in use by the worker handling this Nexus operation."""
 
     @classmethod
-    def get(cls) -> _TemporalNexusOperationContext:
-        ctx = _temporal_operation_context.get(None)
+    def get(cls) -> _TemporalStartOperationContext:
+        ctx = _temporal_start_operation_context.get(None)
         if ctx is None:
             raise RuntimeError("Not in Nexus operation context.")
         return ctx
 
-    @property
-    def _temporal_start_operation_context(
+    def set(self) -> None:
+        _temporal_start_operation_context.set(self)
+
+    def get_completion_callbacks(
         self,
-    ) -> Optional[_TemporalStartOperationContext]:
-        ctx = self.nexus_operation_context
-        if not isinstance(ctx, StartOperationContext):
-            return None
-        return _TemporalStartOperationContext(ctx)
+    ) -> list[temporalio.client.NexusCompletionCallback]:
+        ctx = self.nexus_context
+        return (
+            [
+                # TODO(nexus-prerelease): For WorkflowRunOperation, when it handles the Nexus
+                # request, it needs to copy the links to the callback in
+                # StartWorkflowRequest.CompletionCallbacks and to StartWorkflowRequest.Links
+                # (for backwards compatibility). PR reference in Go SDK:
+                # https://github.com/temporalio/sdk-go/pull/1945
+                temporalio.client.NexusCompletionCallback(
+                    url=ctx.callback_url,
+                    header=ctx.callback_headers,
+                )
+            ]
+            if ctx.callback_url
+            else []
+        )
 
-    @property
-    def _temporal_cancel_operation_context(
+    def get_workflow_event_links(
         self,
-    ) -> Optional[_TemporalCancelOperationContext]:
-        ctx = self.nexus_operation_context
-        if not isinstance(ctx, CancelOperationContext):
-            return None
-        return _TemporalCancelOperationContext(ctx)
+    ) -> list[temporalio.api.common.v1.Link.WorkflowEvent]:
+        event_links = []
+        for inbound_link in self.nexus_context.inbound_links:
+            if link := _nexus_link_to_workflow_event(inbound_link):
+                event_links.append(link)
+        return event_links
+
+    def add_outbound_links(
+        self, workflow_handle: temporalio.client.WorkflowHandle[Any, Any]
+    ):
+        try:
+            link = _workflow_event_to_nexus_link(
+                _workflow_handle_to_workflow_execution_started_event_link(
+                    workflow_handle
+                )
+            )
+        except Exception as e:
+            logger.warning(
+                f"Failed to create WorkflowExecutionStarted event link for workflow {id}: {e}"
+            )
+        else:
+            self.nexus_context.outbound_links.append(
+                # TODO(nexus-prerelease): Before, WorkflowRunOperation was generating an EventReference
+                # link to send back to the caller. Now, it checks if the server returned
+                # the link in the StartWorkflowExecutionResponse, and if so, send the link
+                # from the response to the caller. Fallback to generating the link for
+                # backwards compatibility. PR reference in Go SDK:
+                # https://github.com/temporalio/sdk-go/pull/1934
+                link
+            )
+        return workflow_handle
 
 
 @dataclass
 class WorkflowRunOperationContext:
-    start_operation_context: StartOperationContext
+    temporal_context: _TemporalStartOperationContext
+
+    @property
+    def nexus_context(self) -> StartOperationContext:
+        return self.temporal_context.nexus_context
+
+    @classmethod
+    def get(cls) -> WorkflowRunOperationContext:
+        return cls(_TemporalStartOperationContext.get())
 
     # Overload for single-param workflow
     # TODO(nexus-prerelease): bring over other overloads
@@ -164,14 +239,6 @@ async def start_workflow(
             Nexus caller is itself a workflow, this means that the workflow in the caller
             namespace web UI will contain links to the started workflow, and vice versa.
         """
-        tctx = _TemporalNexusOperationContext.get()
-        start_operation_context = tctx._temporal_start_operation_context
-        if not start_operation_context:
-            raise RuntimeError(
-                "WorkflowRunOperationContext.start_workflow() must be called from "
-                "within a Nexus start operation context"
-            )
-
         # TODO(nexus-preview): When sdk-python supports on_conflict_options, Typescript does this:
         # if (workflowOptions.workflowIdConflictPolicy === 'USE_EXISTING') {
         #     internalOptions.onConflictOptions = {
@@ -184,11 +251,11 @@ async def start_workflow(
         # We must pass nexus_completion_callbacks, workflow_event_links, and request_id,
         # but these are deliberately not exposed in overloads, hence the type-check
         # violation.
-        wf_handle = await tctx.client.start_workflow(  # type: ignore
+        wf_handle = await self.temporal_context.client.start_workflow(  # type: ignore
             workflow=workflow,
             arg=arg,
             id=id,
-            task_queue=task_queue or tctx.info().task_queue,
+            task_queue=task_queue or self.temporal_context.info().task_queue,
             execution_timeout=execution_timeout,
             run_timeout=run_timeout,
             task_timeout=task_timeout,
@@ -208,78 +275,40 @@ async def start_workflow(
             request_eager_start=request_eager_start,
             priority=priority,
             versioning_override=versioning_override,
-            nexus_completion_callbacks=start_operation_context.get_completion_callbacks(),
-            workflow_event_links=start_operation_context.get_workflow_event_links(),
-            request_id=start_operation_context.nexus_operation_context.request_id,
+            nexus_completion_callbacks=self.temporal_context.get_completion_callbacks(),
+            workflow_event_links=self.temporal_context.get_workflow_event_links(),
+            request_id=self.temporal_context.nexus_context.request_id,
         )
 
-        start_operation_context.add_outbound_links(wf_handle)
+        self.temporal_context.add_outbound_links(wf_handle)
 
         return WorkflowHandle[ReturnType]._unsafe_from_client_workflow_handle(wf_handle)
 
 
 @dataclass
-class _TemporalStartOperationContext:
-    nexus_operation_context: StartOperationContext
+class _TemporalCancelOperationContext:
+    """
+    Context for a Nexus cancel operation being handled by a Temporal Nexus Worker.
+    """
 
-    def get_completion_callbacks(
-        self,
-    ) -> list[temporalio.client.NexusCompletionCallback]:
-        ctx = self.nexus_operation_context
-        return (
-            [
-                # TODO(nexus-prerelease): For WorkflowRunOperation, when it handles the Nexus
-                # request, it needs to copy the links to the callback in
-                # StartWorkflowRequest.CompletionCallbacks and to StartWorkflowRequest.Links
-                # (for backwards compatibility). PR reference in Go SDK:
-                # https://github.com/temporalio/sdk-go/pull/1945
-                temporalio.client.NexusCompletionCallback(
-                    url=ctx.callback_url,
-                    header=ctx.callback_headers,
-                )
-            ]
-            if ctx.callback_url
-            else []
-        )
+    nexus_context: CancelOperationContext
+    """Nexus-specific cancel operation context."""
 
-    def get_workflow_event_links(
-        self,
-    ) -> list[temporalio.api.common.v1.Link.WorkflowEvent]:
-        event_links = []
-        for inbound_link in self.nexus_operation_context.inbound_links:
-            if link := _nexus_link_to_workflow_event(inbound_link):
-                event_links.append(link)
-        return event_links
+    info: Callable[[], Info]
+    """Temporal information about the running Nexus cancel operation."""
 
-    def add_outbound_links(
-        self, workflow_handle: temporalio.client.WorkflowHandle[Any, Any]
-    ):
-        try:
-            link = _workflow_event_to_nexus_link(
-                _workflow_handle_to_workflow_execution_started_event_link(
-                    workflow_handle
-                )
-            )
-        except Exception as e:
-            logger.warning(
-                f"Failed to create WorkflowExecutionStarted event link for workflow {id}: {e}"
-            )
-        else:
-            self.nexus_operation_context.outbound_links.append(
-                # TODO(nexus-prerelease): Before, WorkflowRunOperation was generating an EventReference
-                # link to send back to the caller. Now, it checks if the server returned
-                # the link in the StartWorkflowExecutionResponse, and if so, send the link
-                # from the response to the caller. Fallback to generating the link for
-                # backwards compatibility. PR reference in Go SDK:
-                # https://github.com/temporalio/sdk-go/pull/1934
-                link
-            )
-        return workflow_handle
+    client: temporalio.client.Client
+    """The Temporal client in use by the worker handling the current Nexus operation."""
 
+    @classmethod
+    def get(cls) -> _TemporalCancelOperationContext:
+        ctx = _temporal_cancel_operation_context.get(None)
+        if ctx is None:
+            raise RuntimeError("Not in Nexus cancel operation context.")
+        return ctx
 
-@dataclass
-class _TemporalCancelOperationContext:
-    nexus_operation_context: CancelOperationContext
+    def set(self) -> None:
+        _temporal_cancel_operation_context.set(self)
 
 
 def _workflow_handle_to_workflow_execution_started_event_link(
@@ -376,9 +405,9 @@ def process(
         self, msg: Any, kwargs: MutableMapping[str, Any]
     ) -> tuple[Any, MutableMapping[str, Any]]:
         extra = dict(self.extra or {})
-        if tctx := _temporal_operation_context.get(None):
-            extra["service"] = tctx.nexus_operation_context.service
-            extra["operation"] = tctx.nexus_operation_context.operation
+        if tctx := _try_temporal_context():
+            extra["service"] = tctx.nexus_context.service
+            extra["operation"] = tctx.nexus_context.operation
             extra["task_queue"] = tctx.info().task_queue
         kwargs["extra"] = extra | kwargs.get("extra", {})
         return msg, kwargs

temporalio/nexus/_operation_handlers.py

@@ -26,7 +26,7 @@
 
 from temporalio import client
 from temporalio.nexus._operation_context import (
-    _temporal_operation_context,
+    _temporal_start_operation_context,
 )
 from temporalio.nexus._token import WorkflowHandle
 
@@ -114,7 +114,7 @@ async def fetch_result(
                 type=HandlerErrorType.NOT_FOUND,
                 cause=err,
             )
-        ctx = _temporal_operation_context.get()
+        ctx = _temporal_start_operation_context.get()
         try:
             client_handle = nexus_handle.to_workflow_handle(
                 ctx.client, result_type=self._output_type
@@ -148,7 +148,7 @@ async def cancel_operation(
             cause=err,
         )
 
-    ctx = _temporal_operation_context.get()
+    ctx = _temporal_start_operation_context.get()
     try:
         client_workflow_handle = nexus_workflow_handle._to_client_workflow_handle(
             ctx.client