diff --git a/temporalio/nexus/system/workflow_service/_system_nexus_interceptor.py b/temporalio/nexus/system/workflow_service/_system_nexus_interceptor.py index 2f5e54ee2..0b59b1d42 100644 --- a/temporalio/nexus/system/workflow_service/_system_nexus_interceptor.py +++ b/temporalio/nexus/system/workflow_service/_system_nexus_interceptor.py @@ -5,8 +5,6 @@ import abc import typing -from temporalio.nexus.system import TEMPORAL_SYSTEM_ENDPOINT - from . import models if typing.TYPE_CHECKING: @@ -66,7 +64,7 @@ async def start_signal_with_start_workflow( class _SystemNexusWorkflowOutboundInterceptorTerminal(abc.ABC): @abc.abstractmethod - async def _outbound_start_nexus_operation( + async def _intercept_system_nexus_operation( self, input: StartNexusOperationInput[_InputT, _OutputT], ) -> temporalio.workflow.NexusOperationHandle[_OutputT]: ... @@ -76,10 +74,11 @@ async def start_signal_with_start_workflow( ) -> temporalio.workflow.NexusOperationHandle[ models.SignalWithStartWorkflowResponse ]: + from temporalio.nexus.system import TEMPORAL_SYSTEM_ENDPOINT from temporalio.worker._interceptor import StartNexusOperationInput from temporalio.workflow import NexusOperationCancellationType - return await self._outbound_start_nexus_operation( + return await self._intercept_system_nexus_operation( StartNexusOperationInput( endpoint=TEMPORAL_SYSTEM_ENDPOINT, service="temporal.api.workflowservice.v1.WorkflowService", diff --git a/temporalio/worker/_interceptor.py b/temporalio/worker/_interceptor.py index 1bac84673..f59b534c2 100644 --- a/temporalio/worker/_interceptor.py +++ b/temporalio/worker/_interceptor.py @@ -491,6 +491,12 @@ async def start_nexus_operation( """Called for every :py:func:`temporalio.workflow.NexusClient.start_operation` call.""" return await self.next.start_nexus_operation(input) + async def start_system_nexus_operation( + self, input: StartNexusOperationInput[InputT, OutputT] + ) -> temporalio.workflow.NexusOperationHandle[OutputT]: + """Intercept a Temporal System Nexus operation started by a workflow.""" + return await self.next.start_system_nexus_operation(input) + @dataclass class ExecuteNexusOperationStartInput: diff --git a/temporalio/worker/_workflow_instance.py b/temporalio/worker/_workflow_instance.py index fca73c266..fc5a6efe3 100644 --- a/temporalio/worker/_workflow_instance.py +++ b/temporalio/worker/_workflow_instance.py @@ -2199,6 +2199,16 @@ async def operation_handle_fn() -> OutputT: ) return handle + async def _intercept_system_nexus_operation( + self, input: StartNexusOperationInput[Any, OutputT] + ) -> temporalio.workflow.NexusOperationHandle[OutputT]: + return await self._outbound.start_system_nexus_operation(input) + + async def _schedule_system_nexus_operation( + self, input: StartNexusOperationInput[Any, OutputT] + ) -> _NexusOperationHandle[OutputT]: + return await self._outbound_start_nexus_operation(input) + #### Miscellaneous helpers #### # These are in alphabetical order. @@ -3192,6 +3202,16 @@ async def start_nexus_operation( ) -> _NexusOperationHandle[OutputT]: return await self._instance._outbound_start_nexus_operation(input) + async def _intercept_system_nexus_operation( + self, input: StartNexusOperationInput[InputT, OutputT] + ) -> temporalio.workflow.NexusOperationHandle[OutputT]: + return await self._instance._intercept_system_nexus_operation(input) + + async def start_system_nexus_operation( + self, input: StartNexusOperationInput[Any, OutputT] + ) -> _NexusOperationHandle[OutputT]: + return await self._instance._schedule_system_nexus_operation(input) + def start_local_activity( self, input: StartLocalActivityInput ) -> temporalio.workflow.ActivityHandle[Any]: diff --git a/tests/nexus/test_temporal_system_nexus.py b/tests/nexus/test_temporal_system_nexus.py index 69bed1a9c..751e5760e 100644 --- a/tests/nexus/test_temporal_system_nexus.py +++ b/tests/nexus/test_temporal_system_nexus.py @@ -36,6 +36,7 @@ from temporalio.testing import WorkflowEnvironment from temporalio.worker import ( Interceptor, + StartNexusOperationInput, Worker, WorkflowInboundInterceptor, WorkflowInterceptorClassInput, @@ -261,6 +262,13 @@ async def start_signal_with_start_workflow( ) return await super().start_signal_with_start_workflow(request) + async def start_system_nexus_operation( + self, + input: StartNexusOperationInput[Any, Any], + ) -> workflow.NexusOperationHandle[Any]: + interceptor_traces.append(("workflow.start_system_nexus_operation", input)) + return await super().start_system_nexus_operation(input) + def _assert_stored_payloads_include( driver: InMemoryTestDriver, expected_payload_data: set[bytes] @@ -275,14 +283,19 @@ def _assert_stored_payloads_include( def _assert_signal_with_start_workflow_interceptor_trace() -> None: - assert len(interceptor_traces) == 1 - trace_name, trace_value = interceptor_traces.pop() + assert len(interceptor_traces) == 2 + trace_name, trace_value = interceptor_traces.pop(0) assert trace_name == "workflow.start_signal_with_start_workflow" request = cast(workflow_service_models.SignalWithStartWorkflowRequest, trace_value) assert request.id == "system-nexus-workflow-id" assert request.signal == "test-signal" assert request.workflow == "test-workflow" assert request.headers == {"interceptor-header": "value"} + trace_name, trace_value = interceptor_traces.pop() + assert trace_name == "workflow.start_system_nexus_operation" + system_input = cast(StartNexusOperationInput[Any, Any], trace_value) + assert system_input.input is request + assert system_input.headers is None class _MarkingPayloadVisitor(VisitorFunctions):