Wire through additional context type in union

Author: dandavison

temporalio/worker/_interceptor.py

@@ -29,6 +29,7 @@
 import temporalio.common
 import temporalio.nexus
 import temporalio.workflow
+from temporalio.nexus import WorkflowRunOperationContext
 from temporalio.workflow import VersioningIntent
 
 
@@ -301,7 +302,14 @@ class StartNexusOperationInput(Generic[InputT, OutputT]):
         nexusrpc.Operation[InputT, OutputT],
         Callable[[Any], nexusrpc.handler.OperationHandler[InputT, OutputT]],
         Callable[
-            [Any, nexusrpc.handler.StartOperationContext, InputT],
+            [
+                Any,
+                Union[
+                    nexusrpc.handler.StartOperationContext,
+                    WorkflowRunOperationContext,
+                ],
+                InputT,
+            ],
             Awaitable[temporalio.nexus.WorkflowHandle[OutputT]],
         ],
         str,

temporalio/worker/_workflow_instance.py

@@ -61,6 +61,7 @@
 import temporalio.exceptions
 import temporalio.nexus
 import temporalio.workflow
+from temporalio.nexus import WorkflowRunOperationContext
 from temporalio.service import __version__
 
 from ._interceptor import (
@@ -1502,7 +1503,14 @@ async def workflow_start_nexus_operation(
             nexusrpc.Operation[I, O],
             Callable[[Any], nexusrpc.handler.OperationHandler[I, O]],
             Callable[
-                [Any, nexusrpc.handler.StartOperationContext, I],
+                [
+                    Any,
+                    Union[
+                        nexusrpc.handler.StartOperationContext,
+                        WorkflowRunOperationContext,
+                    ],
+                    I,
+                ],
                 Awaitable[temporalio.nexus.WorkflowHandle[O]],
             ],
             str,

temporalio/workflow.py

@@ -58,6 +58,7 @@
 import temporalio.exceptions
 import temporalio.nexus
 import temporalio.workflow
+from temporalio.nexus import WorkflowRunOperationContext
 
 from .types import (
     AnyType,
@@ -858,7 +859,14 @@ async def workflow_start_nexus_operation(
             nexusrpc.Operation[I, O],
             Callable[[Any], nexusrpc.handler.OperationHandler[I, O]],
             Callable[
-                [S, nexusrpc.handler.StartOperationContext, I],
+                [
+                    S,
+                    Union[
+                        nexusrpc.handler.StartOperationContext,
+                        WorkflowRunOperationContext,
+                    ],
+                    I,
+                ],
                 Awaitable[temporalio.nexus.WorkflowHandle[O]],
             ],
             str,
@@ -4428,7 +4436,13 @@ async def start_nexus_operation(
         nexusrpc.Operation[I, O],
         Callable[[Any], nexusrpc.handler.OperationHandler[I, O]],
         Callable[
-            [S, nexusrpc.handler.StartOperationContext, I],
+            [
+                S,
+                Union[
+                    nexusrpc.handler.StartOperationContext, WorkflowRunOperationContext
+                ],
+                I,
+            ],
             Awaitable[temporalio.nexus.WorkflowHandle[O]],
         ],
         str,
@@ -5214,7 +5228,14 @@ async def start_operation(
             nexusrpc.Operation[I, O],
             Callable[[S], nexusrpc.handler.OperationHandler[I, O]],
             Callable[
-                [S, nexusrpc.handler.StartOperationContext, I],
+                [
+                    S,
+                    Union[
+                        nexusrpc.handler.StartOperationContext,
+                        WorkflowRunOperationContext,
+                    ],
+                    I,
+                ],
                 Awaitable[temporalio.nexus.WorkflowHandle[O]],
             ],
             str,
@@ -5242,7 +5263,14 @@ async def execute_operation(
             nexusrpc.Operation[I, O],
             Callable[[S], nexusrpc.handler.OperationHandler[I, O]],
             Callable[
-                [S, nexusrpc.handler.StartOperationContext, I],
+                [
+                    S,
+                    Union[
+                        nexusrpc.handler.StartOperationContext,
+                        WorkflowRunOperationContext,
+                    ],
+                    I,
+                ],
                 Awaitable[temporalio.nexus.WorkflowHandle[O]],
             ],
             str,