Make WorkflowRunOperationContext subclass StartOperationContext

Author: dandavison

temporalio/nexus/_decorators.py

@@ -17,7 +17,6 @@
 )
 
 from temporalio.nexus._operation_context import (
-    TemporalStartOperationContext,
     WorkflowRunOperationContext,
 )
 from temporalio.nexus._operation_handlers import (
@@ -115,8 +114,11 @@ 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,
+                    WorkflowRunOperationContext.from_start_operation_context(ctx),
+                    input,
+                )
 
             _start.__doc__ = start.__doc__
             return WorkflowRunOperationHandler(_start, input_type, output_type)

temporalio/nexus/_operation_context.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import dataclasses
 import logging
 import re
 import urllib.parse
@@ -168,17 +169,28 @@ def add_outbound_links(
         return workflow_handle
 
 
-@dataclass
-class WorkflowRunOperationContext:
-    temporal_context: TemporalStartOperationContext
+@dataclass(frozen=True)
+class WorkflowRunOperationContext(StartOperationContext):
+    _temporal_context: Optional[TemporalStartOperationContext] = None
+
+    @property
+    def temporal_context(self) -> TemporalStartOperationContext:
+        if not self._temporal_context:
+            raise RuntimeError("Temporal context not set")
+        return self._temporal_context
 
     @property
     def nexus_context(self) -> StartOperationContext:
         return self.temporal_context.nexus_context
 
     @classmethod
-    def get(cls) -> WorkflowRunOperationContext:
-        return cls(TemporalStartOperationContext.get())
+    def from_start_operation_context(
+        cls, ctx: StartOperationContext
+    ) -> WorkflowRunOperationContext:
+        return cls(
+            _temporal_context=TemporalStartOperationContext.get(),
+            **{f.name: getattr(ctx, f.name) for f in dataclasses.fields(ctx)},
+        )
 
     # Overload for single-param workflow
     # TODO(nexus-prerelease): bring over other overloads

tests/nexus/test_workflow_caller.py

@@ -160,7 +160,8 @@ 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(
+            tctx = WorkflowRunOperationContext.from_start_operation_context(ctx)
+            handle = await tctx.start_workflow(
                 HandlerWorkflow.run,
                 HandlerWfInput(op_input=input),
                 id=input.response_type.operation_workflow_id,

tests/nexus/test_workflow_run_operation.py

@@ -49,7 +49,8 @@ def __init__(self):
     async def start(
         self, ctx: StartOperationContext, input: Input
     ) -> StartOperationResultAsync:
-        handle = await WorkflowRunOperationContext.get().start_workflow(
+        tctx = WorkflowRunOperationContext.from_start_operation_context(ctx)
+        handle = await tctx.start_workflow(
             EchoWorkflow.run,
             input.value,
             id=str(uuid.uuid4()),