Use TemporalStartOperationContext instead of WorkflowRunOperationContext

Author: dandavison

temporalio/nexus/__init__.py

@@ -6,9 +6,6 @@
 from ._operation_context import (
     _TemporalCancelOperationContext as _TemporalCancelOperationContext,
 )
-from ._operation_context import (
-    WorkflowRunOperationContext as WorkflowRunOperationContext,
-)
 from ._operation_context import client as client
 from ._operation_context import info as info
 from ._operation_context import logger as logger

temporalio/nexus/_decorators.py

@@ -18,7 +18,6 @@
 
 from temporalio.nexus._operation_context import (
     TemporalStartOperationContext,
-    WorkflowRunOperationContext,
 )
 from temporalio.nexus._operation_handlers import (
     WorkflowRunOperationHandler,
@@ -37,11 +36,11 @@
 @overload
 def workflow_run_operation(
     start: Callable[
-        [ServiceHandlerT, WorkflowRunOperationContext, InputT],
+        [ServiceHandlerT, TemporalStartOperationContext, InputT],
         Awaitable[WorkflowHandle[OutputT]],
     ],
 ) -> Callable[
-    [ServiceHandlerT, WorkflowRunOperationContext, InputT],
+    [ServiceHandlerT, TemporalStartOperationContext, InputT],
     Awaitable[WorkflowHandle[OutputT]],
 ]: ...
 
@@ -53,12 +52,12 @@ def workflow_run_operation(
 ) -> Callable[
     [
         Callable[
-            [ServiceHandlerT, WorkflowRunOperationContext, InputT],
+            [ServiceHandlerT, TemporalStartOperationContext, InputT],
             Awaitable[WorkflowHandle[OutputT]],
         ]
     ],
     Callable[
-        [ServiceHandlerT, WorkflowRunOperationContext, InputT],
+        [ServiceHandlerT, TemporalStartOperationContext, InputT],
         Awaitable[WorkflowHandle[OutputT]],
     ],
 ]: ...
@@ -67,26 +66,26 @@ def workflow_run_operation(
 def workflow_run_operation(
     start: Optional[
         Callable[
-            [ServiceHandlerT, WorkflowRunOperationContext, InputT],
+            [ServiceHandlerT, TemporalStartOperationContext, InputT],
             Awaitable[WorkflowHandle[OutputT]],
         ]
     ] = None,
     *,
     name: Optional[str] = None,
 ) -> Union[
     Callable[
-        [ServiceHandlerT, WorkflowRunOperationContext, InputT],
+        [ServiceHandlerT, TemporalStartOperationContext, InputT],
         Awaitable[WorkflowHandle[OutputT]],
     ],
     Callable[
         [
             Callable[
-                [ServiceHandlerT, WorkflowRunOperationContext, InputT],
+                [ServiceHandlerT, TemporalStartOperationContext, InputT],
                 Awaitable[WorkflowHandle[OutputT]],
             ]
         ],
         Callable[
-            [ServiceHandlerT, WorkflowRunOperationContext, InputT],
+            [ServiceHandlerT, TemporalStartOperationContext, InputT],
             Awaitable[WorkflowHandle[OutputT]],
         ],
     ],
@@ -97,11 +96,11 @@ def workflow_run_operation(
 
     def decorator(
         start: Callable[
-            [ServiceHandlerT, WorkflowRunOperationContext, InputT],
+            [ServiceHandlerT, TemporalStartOperationContext, InputT],
             Awaitable[WorkflowHandle[OutputT]],
         ],
     ) -> Callable[
-        [ServiceHandlerT, WorkflowRunOperationContext, InputT],
+        [ServiceHandlerT, TemporalStartOperationContext, InputT],
         Awaitable[WorkflowHandle[OutputT]],
     ]:
         (
@@ -115,8 +114,7 @@ def operation_handler_factory(
             async def _start(
                 ctx: StartOperationContext, input: InputT
             ) -> WorkflowHandle[OutputT]:
-                tctx = TemporalStartOperationContext.get()
-                return await start(self, WorkflowRunOperationContext(tctx), input)
+                return await start(self, TemporalStartOperationContext.get(), input)
 
             _start.__doc__ = start.__doc__
             return WorkflowRunOperationHandler(_start, input_type, output_type)

temporalio/nexus/_operation_context.py

@@ -167,19 +167,6 @@ def add_outbound_links(
             )
         return workflow_handle
 
-
-@dataclass
-class WorkflowRunOperationContext:
-    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
     async def start_workflow(
@@ -251,11 +238,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 self.temporal_context.client.start_workflow(  # type: ignore
+        wf_handle = await self.client.start_workflow(  # type: ignore
             workflow=workflow,
             arg=arg,
             id=id,
-            task_queue=task_queue or self.temporal_context.info().task_queue,
+            task_queue=task_queue or self.info().task_queue,
             execution_timeout=execution_timeout,
             run_timeout=run_timeout,
             task_timeout=task_timeout,
@@ -275,12 +262,12 @@ async def start_workflow(
             request_eager_start=request_eager_start,
             priority=priority,
             versioning_override=versioning_override,
-            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,
+            nexus_completion_callbacks=self.get_completion_callbacks(),
+            workflow_event_links=self.get_workflow_event_links(),
+            request_id=self.nexus_context.request_id,
         )
 
-        self.temporal_context.add_outbound_links(wf_handle)
+        self.add_outbound_links(wf_handle)
 
         return WorkflowHandle[ReturnType]._unsafe_from_client_workflow_handle(wf_handle)
 

temporalio/nexus/_operation_handlers.py

@@ -41,7 +41,7 @@ class WorkflowRunOperationHandler(OperationHandler[InputT, OutputT]):
 
     Use this class to create an operation handler that starts a workflow by passing your
     ``start`` method to the constructor. Your ``start`` method must use
-    :py:func:`temporalio.nexus.WorkflowRunOperationContext.start_workflow` to start the
+    :py:func:`temporalio.nexus.TemporalStartOperationContext.start_workflow` to start the
     workflow.
     """
 
@@ -77,7 +77,7 @@ async def start(
             if isinstance(handle, client.WorkflowHandle):
                 raise RuntimeError(
                     f"Expected {handle} to be a nexus.WorkflowHandle, but got a client.WorkflowHandle. "
-                    f"You must use WorkflowRunOperationContext.start_workflow "
+                    f"You must use TemporalStartOperationContext.start_workflow "
                     "to start a workflow that will deliver the result of the Nexus operation, "
                     "not client.Client.start_workflow."
                 )

temporalio/nexus/_util.py

@@ -19,7 +19,7 @@
     OutputT,
 )
 
-from temporalio.nexus._operation_context import WorkflowRunOperationContext
+from temporalio.nexus._operation_context import TemporalStartOperationContext
 
 from ._token import (
     WorkflowHandle as WorkflowHandle,
@@ -30,7 +30,7 @@
 
 def get_workflow_run_start_method_input_and_output_type_annotations(
     start: Callable[
-        [ServiceHandlerT, WorkflowRunOperationContext, InputT],
+        [ServiceHandlerT, TemporalStartOperationContext, InputT],
         Awaitable[WorkflowHandle[OutputT]],
     ],
 ) -> tuple[
@@ -70,7 +70,7 @@ def get_workflow_run_start_method_input_and_output_type_annotations(
 
 def _get_start_method_input_and_output_type_annotations(
     start: Callable[
-        [ServiceHandlerT, WorkflowRunOperationContext, InputT],
+        [ServiceHandlerT, TemporalStartOperationContext, InputT],
         Union[OutputT, Awaitable[OutputT]],
     ],
 ) -> tuple[
@@ -102,11 +102,11 @@ def _get_start_method_input_and_output_type_annotations(
         input_type = None
     else:
         ctx_type, input_type = type_annotations.values()
-        if not issubclass(ctx_type, WorkflowRunOperationContext):
+        if not issubclass(ctx_type, TemporalStartOperationContext):
             # TODO(preview): stacklevel
             warnings.warn(
                 f"Expected first parameter of {start} to be an instance of "
-                f"WorkflowRunOperationContext, but is {ctx_type}."
+                f"TemporalStartOperationContext, but is {ctx_type}."
             )
             input_type = None
 

tests/nexus/test_handler.py

@@ -45,7 +45,7 @@
 from temporalio.client import Client
 from temporalio.common import WorkflowIDReusePolicy
 from temporalio.exceptions import ApplicationError
-from temporalio.nexus import WorkflowRunOperationContext, workflow_run_operation
+from temporalio.nexus import TemporalStartOperationContext, workflow_run_operation
 from temporalio.testing import WorkflowEnvironment
 from temporalio.worker import Worker
 from tests.helpers.nexus import (
@@ -208,7 +208,7 @@ async def log(self, ctx: StartOperationContext, input: Input) -> Output:
 
     @workflow_run_operation
     async def workflow_run_operation_happy_path(
-        self, ctx: WorkflowRunOperationContext, input: Input
+        self, ctx: TemporalStartOperationContext, input: Input
     ) -> nexus.WorkflowHandle[Output]:
         return await ctx.start_workflow(
             MyWorkflow.run,
@@ -266,7 +266,7 @@ async def workflow_run_operation_without_type_annotations(self, ctx, input):
 
     @workflow_run_operation
     async def workflow_run_op_link_test(
-        self, ctx: WorkflowRunOperationContext, input: Input
+        self, ctx: TemporalStartOperationContext, input: Input
     ) -> nexus.WorkflowHandle[Output]:
         assert any(
             link.url == "http://inbound-link/"
@@ -1022,7 +1022,7 @@ async def run(self, input: Input) -> Output:
 class ServiceHandlerForRequestIdTest:
     @workflow_run_operation
     async def operation_backed_by_a_workflow(
-        self, ctx: WorkflowRunOperationContext, input: Input
+        self, ctx: TemporalStartOperationContext, input: Input
     ) -> nexus.WorkflowHandle[Output]:
         return await ctx.start_workflow(
             EchoWorkflow.run,
@@ -1033,7 +1033,7 @@ async def operation_backed_by_a_workflow(
 
     @workflow_run_operation
     async def operation_that_executes_a_workflow_before_starting_the_backing_workflow(
-        self, ctx: WorkflowRunOperationContext, input: Input
+        self, ctx: TemporalStartOperationContext, input: Input
     ) -> nexus.WorkflowHandle[Output]:
         await nexus.client().start_workflow(
             EchoWorkflow.run,

tests/nexus/test_handler_interface_implementation.py

@@ -6,7 +6,7 @@
 from nexusrpc.handler import StartOperationContext, sync_operation
 
 from temporalio import nexus
-from temporalio.nexus import WorkflowRunOperationContext, workflow_run_operation
+from temporalio.nexus import TemporalStartOperationContext, workflow_run_operation
 
 HTTP_PORT = 7243
 
@@ -37,7 +37,7 @@ class Interface:
     class Impl:
         @workflow_run_operation
         async def op(
-            self, ctx: WorkflowRunOperationContext, input: str
+            self, ctx: TemporalStartOperationContext, input: str
         ) -> nexus.WorkflowHandle[int]: ...
 
     error_message = None

tests/nexus/test_handler_operation_definitions.py

@@ -10,7 +10,7 @@
 import pytest
 
 from temporalio import nexus
-from temporalio.nexus import WorkflowRunOperationContext, workflow_run_operation
+from temporalio.nexus import TemporalStartOperationContext, workflow_run_operation
 
 
 @dataclass
@@ -34,7 +34,7 @@ class NotCalled(_TestCase):
     class Service:
         @workflow_run_operation
         async def my_workflow_run_operation_handler(
-            self, ctx: WorkflowRunOperationContext, input: Input
+            self, ctx: TemporalStartOperationContext, input: Input
         ) -> nexus.WorkflowHandle[Output]: ...
 
     expected_operations = {
@@ -52,7 +52,7 @@ class CalledWithoutArgs(_TestCase):
     class Service:
         @workflow_run_operation
         async def my_workflow_run_operation_handler(
-            self, ctx: WorkflowRunOperationContext, input: Input
+            self, ctx: TemporalStartOperationContext, input: Input
         ) -> nexus.WorkflowHandle[Output]: ...
 
     expected_operations = NotCalled.expected_operations
@@ -63,7 +63,7 @@ class CalledWithNameOverride(_TestCase):
     class Service:
         @workflow_run_operation(name="operation-name")
         async def workflow_run_operation_with_name_override(
-            self, ctx: WorkflowRunOperationContext, input: Input
+            self, ctx: TemporalStartOperationContext, input: Input
         ) -> nexus.WorkflowHandle[Output]: ...
 
     expected_operations = {

tests/nexus/test_workflow_caller.py

@@ -38,7 +38,7 @@
 )
 from temporalio.common import WorkflowIDConflictPolicy
 from temporalio.exceptions import CancelledError, NexusHandlerError, NexusOperationError
-from temporalio.nexus import WorkflowRunOperationContext, workflow_run_operation
+from temporalio.nexus import TemporalStartOperationContext, workflow_run_operation
 from temporalio.service import RPCError, RPCStatusCode
 from temporalio.worker import Worker
 from tests.helpers.nexus import create_nexus_endpoint, make_nexus_endpoint_name
@@ -160,7 +160,7 @@ async def start(
             # TODO(nexus-preview): what do we want the DX to be for a user who is
             # starting a Nexus backing workflow from a custom start method? (They may
             # need to do this in order to customize the cancel method).
-            handle = await WorkflowRunOperationContext.get().start_workflow(
+            handle = await TemporalStartOperationContext.get().start_workflow(
                 HandlerWorkflow.run,
                 HandlerWfInput(op_input=input),
                 id=input.response_type.operation_workflow_id,
@@ -206,7 +206,7 @@ async def sync_operation(
 
     @workflow_run_operation
     async def async_operation(
-        self, ctx: WorkflowRunOperationContext, input: OpInput
+        self, ctx: TemporalStartOperationContext, input: OpInput
     ) -> nexus.WorkflowHandle[HandlerWfOutput]:
         assert isinstance(input.response_type, AsyncResponse)
         if input.response_type.exception_in_operation_start:
@@ -912,7 +912,7 @@ async def run(self, input: str) -> str:
 class ServiceImplWithOperationsThatExecuteWorkflowBeforeStartingBackingWorkflow:
     @workflow_run_operation
     async def my_workflow_run_operation(
-        self, ctx: WorkflowRunOperationContext, input: None
+        self, ctx: TemporalStartOperationContext, input: None
     ) -> nexus.WorkflowHandle[str]:
         result_1 = await nexus.client().execute_workflow(
             EchoWorkflow.run,

tests/nexus/test_workflow_run_operation.py

@@ -13,7 +13,7 @@
 from nexusrpc.handler._decorators import operation_handler
 
 from temporalio import workflow
-from temporalio.nexus import WorkflowRunOperationContext
+from temporalio.nexus import TemporalStartOperationContext
 from temporalio.nexus._operation_handlers import WorkflowRunOperationHandler
 from temporalio.testing import WorkflowEnvironment
 from temporalio.worker import Worker
@@ -49,7 +49,7 @@ def __init__(self):
     async def start(
         self, ctx: StartOperationContext, input: Input
     ) -> StartOperationResultAsync:
-        handle = await WorkflowRunOperationContext.get().start_workflow(
+        handle = await TemporalStartOperationContext.get().start_workflow(
             EchoWorkflow.run,
             input.value,
             id=str(uuid.uuid4()),