feat: add System Nexus operation interceptors

CHANGELOG.md

@@ -22,6 +22,10 @@ to include examples, links to docs, or any other relevant information.
 
 ### Changed
 
+- System Nexus Signal-with-Start Workflow operations now use the typed
+  `WorkflowOutboundInterceptor.start_signal_with_start_workflow` interception point instead of
+  the generic `WorkflowOutboundInterceptor.start_nexus_operation` method.
+
 ### Deprecated
 
 ### :boom: Breaking Changes

scripts/gen_nexus_system_api.py

@@ -131,6 +131,7 @@ def generate_nexus_system_api() -> None:
                 str(wit_path),
                 str(wit_deps_dir),
                 "--native-api",
+                "--system-nexus",
                 "--support-file",
                 str(python_support_path),
                 "--descriptors",

temporalio/contrib/opentelemetry/_interceptor.py

@@ -33,6 +33,7 @@
 import temporalio.client
 import temporalio.converter
 import temporalio.exceptions
+import temporalio.nexus.system.workflow_service.models
 import temporalio.worker
 import temporalio.workflow
 from temporalio.exceptions import ApplicationError, ApplicationErrorCategory
@@ -433,6 +434,10 @@ class _InputWithStringHeaders(Protocol):
     headers: Mapping[str, str] | None
 
 
+class _InputWithModelHeaders(Protocol):
+    headers: Mapping[str, Any] | None
+
+
 class _InputWithOperationContext(Generic[_ContextT], Protocol):
     ctx: _ContextT
 
@@ -684,13 +689,26 @@ def _context_carrier_to_headers(
             }
         return headers
 
+    def _context_carrier_to_model_headers(
+        self,
+        carrier: _CarrierDict,
+        headers: Mapping[str, Any] | None,
+    ) -> Mapping[str, Any]:
+        if carrier:
+            return {
+                **(headers or {}),
+                self.header_key: carrier,
+            }
+        return headers or {}
+
     def _completed_span(
         self,
         span_name: str,
         *,
         link_context_carrier: _CarrierDict | None = None,
         add_to_outbound: _InputWithHeaders | None = None,
         add_to_outbound_str: _InputWithStringHeaders | None = None,
+        add_to_outbound_model: _InputWithModelHeaders | None = None,
         new_span_even_on_replay: bool = False,
         additional_attributes: opentelemetry.util.types.Attributes = None,
         exception: Exception | None = None,
@@ -742,6 +760,11 @@ def _completed_span(
                     updated_context_carrier, add_to_outbound_str.headers
                 )
 
+            if add_to_outbound_model:
+                add_to_outbound_model.headers = self._context_carrier_to_model_headers(
+                    updated_context_carrier, add_to_outbound_model.headers
+                )
+
     def _set_on_context(
         self, context: opentelemetry.context.Context
     ) -> opentelemetry.context.Context:
@@ -830,6 +853,19 @@ async def start_nexus_operation(
 
         return await super().start_nexus_operation(input)
 
+    async def start_signal_with_start_workflow(
+        self,
+        request: temporalio.nexus.system.workflow_service.models.SignalWithStartWorkflowRequest,
+    ) -> temporalio.workflow.NexusOperationHandle[
+        temporalio.nexus.system.workflow_service.models.SignalWithStartWorkflowResponse
+    ]:
+        self.root._completed_span(
+            "SignalWithStartWorkflow",
+            kind=opentelemetry.trace.SpanKind.CLIENT,
+            add_to_outbound_model=request,
+        )
+        return await super().start_signal_with_start_workflow(request)
+
 
 def _carrier_to_nexus_headers(
     carrier: _CarrierDict, initial: Mapping[str, str] | None = None

temporalio/contrib/opentelemetry/_otel_interceptor.py

@@ -32,6 +32,7 @@
 import temporalio.api.common.v1
 import temporalio.client
 import temporalio.converter
+import temporalio.nexus.system.workflow_service.models
 import temporalio.worker
 import temporalio.workflow
 from temporalio.contrib.opentelemetry._tracer_provider import (
@@ -600,3 +601,15 @@ async def start_nexus_operation(
         ):
             input.headers = _context_to_nexus_headers(input.headers or {})
             return await super().start_nexus_operation(input)
+
+    async def start_signal_with_start_workflow(
+        self,
+        request: temporalio.nexus.system.workflow_service.models.SignalWithStartWorkflowRequest,
+    ) -> temporalio.workflow.NexusOperationHandle[
+        temporalio.nexus.system.workflow_service.models.SignalWithStartWorkflowResponse
+    ]:
+        with self._workflow_maybe_span(
+            "SignalWithStartWorkflow", kind=opentelemetry.trace.SpanKind.CLIENT
+        ):
+            request.headers = _context_to_headers(request.headers or {})
+            return await super().start_signal_with_start_workflow(request)

temporalio/nexus/system/workflow_service/_system_nexus_interceptor.py

@@ -0,0 +1,96 @@
+# Generated by nexgen v0.2.2. DO NOT EDIT!
+
+from __future__ import annotations
+
+import abc
+import typing
+
+from temporalio.nexus.system import TEMPORAL_SYSTEM_ENDPOINT
+
+from . import models
+
+if typing.TYPE_CHECKING:
+    import temporalio.workflow
+    from temporalio.worker._interceptor import StartNexusOperationInput
+
+
+__all__ = [
+    "_start_system_nexus_operation",
+    "_SystemNexusWorkflowOutboundInterceptorBase",
+    "_SystemNexusWorkflowOutboundInterceptorTerminal",
+]
+
+
+_InputT = typing.TypeVar("_InputT")
+_OutputT = typing.TypeVar("_OutputT")
+
+
+async def _start_system_nexus_operation(
+    interceptor: _SystemNexusWorkflowOutboundInterceptorBase,
+    input: StartNexusOperationInput[_InputT, _OutputT],
+) -> temporalio.workflow.NexusOperationHandle[_OutputT]:
+    if (
+        input.service == "temporal.api.workflowservice.v1.WorkflowService"
+        and input.operation_name == "SignalWithStartWorkflowExecution"
+    ):
+        typed_input = typing.cast(
+            "StartNexusOperationInput[models.SignalWithStartWorkflowRequest, models.SignalWithStartWorkflowResponse]",
+            input,
+        )
+        # The dispatch check above establishes that this operation's response type is _OutputT.
+        return typing.cast(
+            "temporalio.workflow.NexusOperationHandle[_OutputT]",
+            await interceptor.start_signal_with_start_workflow(typed_input.input),
+        )
+    raise ValueError(
+        f"unsupported System Nexus operation: {input.service}/{input.operation_name}"
+    )
+
+
+class _SystemNexusWorkflowOutboundInterceptorBase(abc.ABC):
+    @abc.abstractmethod
+    def _next_system_nexus_interceptor(
+        self,
+    ) -> _SystemNexusWorkflowOutboundInterceptorBase: ...
+
+    async def start_signal_with_start_workflow(
+        self, request: models.SignalWithStartWorkflowRequest
+    ) -> temporalio.workflow.NexusOperationHandle[
+        models.SignalWithStartWorkflowResponse
+    ]:
+        """Intercept the SignalWithStartWorkflow operation."""
+        return await self._next_system_nexus_interceptor().start_signal_with_start_workflow(
+            request
+        )
+
+
+class _SystemNexusWorkflowOutboundInterceptorTerminal(abc.ABC):
+    @abc.abstractmethod
+    async def _outbound_start_nexus_operation(
+        self,
+        input: StartNexusOperationInput[_InputT, _OutputT],
+    ) -> temporalio.workflow.NexusOperationHandle[_OutputT]: ...
+
+    async def start_signal_with_start_workflow(
+        self, request: models.SignalWithStartWorkflowRequest
+    ) -> temporalio.workflow.NexusOperationHandle[
+        models.SignalWithStartWorkflowResponse
+    ]:
+        from temporalio.worker._interceptor import StartNexusOperationInput
+        from temporalio.workflow import NexusOperationCancellationType
+
+        return await self._outbound_start_nexus_operation(
+            StartNexusOperationInput(
+                endpoint=TEMPORAL_SYSTEM_ENDPOINT,
+                service="temporal.api.workflowservice.v1.WorkflowService",
+                operation="SignalWithStartWorkflowExecution",
+                input=request,
+                output_type=models.SignalWithStartWorkflowResponse,
+                schedule_to_close_timeout=None,
+                schedule_to_start_timeout=None,
+                start_to_close_timeout=None,
+                cancellation_type=NexusOperationCancellationType.WAIT_COMPLETED,
+                headers=None,
+                summary=None,
+            )
+        )

temporalio/worker/_interceptor.py

@@ -21,6 +21,9 @@
 import temporalio.nexus
 import temporalio.nexus._util
 import temporalio.workflow
+from temporalio.nexus.system.workflow_service._system_nexus_interceptor import (
+    _SystemNexusWorkflowOutboundInterceptorBase,
+)
 from temporalio.workflow import ContinueAsNewVersioningBehavior, VersioningIntent
 
 
@@ -414,12 +417,14 @@ async def handle_update_handler(self, input: HandleUpdateInput) -> Any:
         return await self.next.handle_update_handler(input)
 
 
-class WorkflowOutboundInterceptor:
+class WorkflowOutboundInterceptor(_SystemNexusWorkflowOutboundInterceptorBase):
     """Outbound interceptor to wrap calls made from within workflows.
 
     This should be extended by any workflow outbound interceptors.
     """
 
+    next: WorkflowOutboundInterceptor
+
     def __init__(self, next: WorkflowOutboundInterceptor) -> None:
         """Create the outbound interceptor.
 
@@ -429,6 +434,11 @@ def __init__(self, next: WorkflowOutboundInterceptor) -> None:
         """
         self.next = next
 
+    def _next_system_nexus_interceptor(
+        self,
+    ) -> _SystemNexusWorkflowOutboundInterceptorBase:
+        return self.next
+
     def continue_as_new(self, input: ContinueAsNewInput) -> NoReturn:
         """Called for every :py:func:`temporalio.workflow.continue_as_new` call."""
         self.next.continue_as_new(input)

temporalio/worker/_workflow_instance.py

@@ -60,6 +60,10 @@
 import temporalio.nexus.system
 import temporalio.workflow
 from temporalio.converter import StorageDriverStoreContext, StorageDriverWorkflowInfo
+from temporalio.nexus.system.workflow_service._system_nexus_interceptor import (
+    _start_system_nexus_operation,
+    _SystemNexusWorkflowOutboundInterceptorTerminal,
+)
 from temporalio.service import __version__
 
 from ..api.failure.v1.message_pb2 import Failure
@@ -1732,7 +1736,23 @@ async def workflow_start_nexus_operation(
         headers: Mapping[str, str] | None,
         summary: str | None,
     ) -> temporalio.workflow.NexusOperationHandle[OutputT]:
-        # start_nexus_operation
+        if temporalio.nexus.system.is_system_endpoint(endpoint):
+            return await _start_system_nexus_operation(
+                self._outbound,
+                StartNexusOperationInput(
+                    endpoint=temporalio.nexus.system.TEMPORAL_SYSTEM_ENDPOINT,
+                    service=service,
+                    operation=operation,
+                    input=input,
+                    output_type=output_type,
+                    schedule_to_close_timeout=schedule_to_close_timeout,
+                    schedule_to_start_timeout=schedule_to_start_timeout,
+                    start_to_close_timeout=start_to_close_timeout,
+                    cancellation_type=cancellation_type,
+                    headers=None,
+                    summary=summary,
+                ),
+            )
         return await self._outbound.start_nexus_operation(
             StartNexusOperationInput(
                 endpoint=endpoint,
@@ -3131,11 +3151,18 @@ async def handle_update_handler(self, input: HandleUpdateInput) -> Any:
             return handler(*input.args)
 
 
-class _WorkflowOutboundImpl(WorkflowOutboundInterceptor):
+class _WorkflowOutboundImpl(
+    _SystemNexusWorkflowOutboundInterceptorTerminal, WorkflowOutboundInterceptor
+):
     def __init__(self, instance: _WorkflowInstanceImpl) -> None:  # type: ignore
         # We are intentionally not calling the base class's __init__ here
         self._instance = instance
 
+    async def _outbound_start_nexus_operation(
+        self, input: StartNexusOperationInput[InputT, OutputT]
+    ) -> temporalio.workflow.NexusOperationHandle[OutputT]:
+        return await self._instance._outbound_start_nexus_operation(input)
+
     def continue_as_new(self, input: ContinueAsNewInput) -> NoReturn:
         self._instance._outbound_continue_as_new(input)
 

tests/contrib/opentelemetry/test_opentelemetry.py

@@ -88,6 +88,34 @@ class TracingWorkflowActionActivity:
     fail_on_non_replay_before_complete: bool = False
 
 
+@workflow.defn
+class SignalWithStartHeaderWorkflow:
+    def __init__(self) -> None:
+        self._signaled = False
+
+    @workflow.run
+    async def run(self) -> bool:
+        await workflow.wait_condition(lambda: self._signaled)
+        return "_tracer-data" in workflow.info().headers
+
+    @workflow.signal
+    def notify(self) -> None:
+        self._signaled = True
+
+
+@workflow.defn
+class SignalWithStartCallerWorkflow:
+    @workflow.run
+    async def run(self, target_id: str, task_queue: str) -> str:
+        handle = await workflow.signal_with_start_workflow(
+            SignalWithStartHeaderWorkflow.run,
+            id=target_id,
+            task_queue=task_queue,
+            signal=SignalWithStartHeaderWorkflow.notify,
+        )
+        return handle.id
+
+
 @dataclass
 class TracingWorkflowActionContinueAsNew:
     param: TracingWorkflowParam
@@ -229,6 +257,41 @@ def update_validator(self) -> None:
         pass
 
 
+@pytest.mark.requires_local_server
+async def test_workflow_signal_with_start_propagates_trace_headers(
+    client: Client, env: WorkflowEnvironment
+):
+    if env.supports_time_skipping:
+        pytest.skip("Nexus tests don't work with the Java test server")
+    provider = TracerProvider()
+    exporter = InMemorySpanExporter()
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    tracer = get_tracer(__name__, tracer_provider=provider)
+    config = client.config()
+    config["interceptors"] = [TracingInterceptor(tracer)]
+    client = Client(**config)
+
+    async with Worker(
+        client,
+        task_queue=f"signal-with-start-{uuid.uuid4()}",
+        workflows=[SignalWithStartCallerWorkflow, SignalWithStartHeaderWorkflow],
+        workflow_runner=UnsandboxedWorkflowRunner(),
+    ) as worker:
+        target_id = f"signal-with-start-target-{uuid.uuid4()}"
+        with tracer.start_as_current_span("signal-with-start"):
+            caller = await client.start_workflow(
+                SignalWithStartCallerWorkflow.run,
+                args=[target_id, worker.task_queue],
+                id=f"signal-with-start-caller-{uuid.uuid4()}",
+                task_queue=worker.task_queue,
+            )
+            assert await caller.result() == target_id
+        assert await client.get_workflow_handle(target_id).result() is True
+    assert any(
+        span.name == "SignalWithStartWorkflow" for span in exporter.get_finished_spans()
+    )
+
+
 async def test_opentelemetry_tracing(client: Client, env: WorkflowEnvironment):
     # TODO(cretz): Fix
     if env.supports_time_skipping: