Adapt System Nexus tracing to generated request types

Author: tconley1428

temporalio/converter/_payload_converter.py

@@ -613,14 +613,26 @@ def wrap(payload_converter: PayloadConverter) -> PayloadConverter:
         return _TemporalTransferTypePayloadConverter(payload_converter)
 
     def to_payloads(
-        self, values: Sequence[Any]
+        self,
+        values: Sequence[Any],
+        headers: Mapping[str, temporalio.api.common.v1.Payload] | None = None,
     ) -> list[temporalio.api.common.v1.Payload]:
         """See base class."""
         transfer_type_values: list[Any] = []
         for value in values:
             converter = _get_transfer_type_converter(type(value))
             if converter is not None:
                 value = converter.to_transfer_type(value)
+            if (
+                headers
+                and isinstance(value, google.protobuf.message.Message)
+                and "header" in value.DESCRIPTOR.fields_by_name
+            ):
+                # System Nexus starts with generated models, so headers can only be
+                # applied after conversion to a request protobuf and before encoding.
+                temporalio.common._apply_headers(
+                    headers, getattr(value, "header").fields
+                )
             transfer_type_values.append(value)
         return self._inner_payload_converter.to_payloads(transfer_type_values)
 

temporalio/nexus/system/__init__.py

@@ -12,8 +12,6 @@
 from dataclasses import dataclass
 from typing import Any
 
-from google.protobuf.message import Message
-
 import temporalio.api.common.v1
 import temporalio.common
 import temporalio.converter
@@ -93,27 +91,31 @@ class _SystemNexusPayloadConverter(temporalio.converter.PayloadConverter):
     """Payload converter for system Nexus outer envelopes."""
 
     _user_converters: _SystemNexusUserConverters
-    _outer_payload_converter: temporalio.converter.PayloadConverter
+    _outer_payload_converter: _TemporalTransferTypePayloadConverter
+    _headers: Mapping[str, temporalio.api.common.v1.Payload] | None
 
     def __init__(
         self,
         user_payload_converter: temporalio.converter.PayloadConverter,
         user_failure_converter: temporalio.converter.FailureConverter,
+        headers: Mapping[str, temporalio.api.common.v1.Payload] | None = None,
     ) -> None:
         """Create a payload converter for system Nexus outer envelopes."""
         self._user_converters = _SystemNexusUserConverters(
             user_payload_converter, user_failure_converter
         )
-        self._outer_payload_converter = _TemporalTransferTypePayloadConverter.wrap(
+
+        self._outer_payload_converter = _TemporalTransferTypePayloadConverter(
             _SystemNexusOuterPayloadConverter()
         )
+        self._headers = headers
 
     def to_payloads(
         self, values: Sequence[Any]
     ) -> list[temporalio.api.common.v1.Payload]:
         """See base class."""
         with _user_converter_context(self._user_converters):
-            return self._outer_payload_converter.to_payloads(values)
+            return self._outer_payload_converter.to_payloads(values, self._headers)
 
     def from_payloads(
         self,
@@ -134,23 +136,13 @@ def is_system_endpoint(endpoint: str) -> bool:
     return endpoint == TEMPORAL_SYSTEM_ENDPOINT
 
 
-def _apply_headers_to_request(
-    request: Message,
-    headers: Mapping[str, temporalio.api.common.v1.Payload],
-) -> None:
-    """Apply headers to a system request when it supports Temporal headers."""
-    if not headers or "header" not in request.DESCRIPTOR.fields_by_name:
-        return
-    request_header = getattr(request, "header")
-    for key, payload in headers.items():
-        request_header.fields[key].CopyFrom(payload)
-
-
 def _is_system_payload(payload: temporalio.api.common.v1.Payload) -> bool:
     return (
         payload.metadata.get(_SYSTEM_PAYLOAD_METADATA_KEY)
         == _SYSTEM_PAYLOAD_METADATA_VALUE
     )
+
+
 async def maybe_visit_payload(
     payload: temporalio.api.common.v1.Payload,
     visitor_functions: VisitorFunctions,
@@ -175,9 +167,12 @@ async def maybe_visit_payload(
 def _get_payload_converter(  # pyright: ignore[reportUnusedFunction]
     user_payload_converter: temporalio.converter.PayloadConverter,
     user_failure_converter: temporalio.converter.FailureConverter,
+    headers: Mapping[str, temporalio.api.common.v1.Payload] | None = None,
 ) -> temporalio.converter.PayloadConverter:
     """Return the fixed payload converter for system Nexus outer envelopes."""
-    return _SystemNexusPayloadConverter(user_payload_converter, user_failure_converter)
+    return _SystemNexusPayloadConverter(
+        user_payload_converter, user_failure_converter, headers
+    )
 
 
 def _get_serialization_context(  # pyright: ignore[reportUnusedFunction]

temporalio/worker/_workflow_instance.py

@@ -2189,22 +2189,51 @@ async def operation_handle_fn() -> OutputT:
     async def _outbound_start_system_nexus_operation(
         self, input: StartSystemNexusOperationInput[Any, OutputT]
     ) -> _NexusOperationHandle[OutputT]:
-        temporalio.nexus.system._apply_headers_to_request(input.input, input.headers)
-        return await self._outbound_start_nexus_operation(
-            StartNexusOperationInput(
-                endpoint=temporalio.nexus.system.TEMPORAL_SYSTEM_ENDPOINT,
-                service=input.service,
-                operation=input.operation,
-                input=input.input,
-                output_type=input.output_type,
-                schedule_to_close_timeout=input.schedule_to_close_timeout,
-                schedule_to_start_timeout=input.schedule_to_start_timeout,
-                start_to_close_timeout=input.start_to_close_timeout,
-                cancellation_type=input.cancellation_type,
-                headers=None,
-                summary=input.summary,
+        nexus_input = StartNexusOperationInput(
+            endpoint=temporalio.nexus.system.TEMPORAL_SYSTEM_ENDPOINT,
+            service=input.service,
+            operation=input.operation,
+            input=input.input,
+            output_type=input.output_type,
+            schedule_to_close_timeout=input.schedule_to_close_timeout,
+            schedule_to_start_timeout=input.schedule_to_start_timeout,
+            start_to_close_timeout=input.start_to_close_timeout,
+            cancellation_type=input.cancellation_type,
+            headers=None,
+            summary=input.summary,
+        )
+        handle: _NexusOperationHandle[OutputT]
+
+        async def operation_handle_fn() -> OutputT:
+            return cast(
+                OutputT,
+                await self._await_temporal_operation(
+                    handle._result_fut,
+                    lambda _err, command: handle._apply_cancel_command(command),
+                ),
             )
+
+        payload_converter = temporalio.nexus.system._get_payload_converter(
+            self._workflow_context_payload_converter,
+            self._workflow_context_failure_converter,
+            input.headers,
         )
+        handle = _NexusOperationHandle(
+            self,
+            self._next_seq("nexus_operation"),
+            nexus_input,
+            operation_handle_fn(),
+            payload_converter,
+        )
+        handle._apply_schedule_command()
+        self._pending_nexus_operations[handle._seq] = handle
+
+        await self._await_temporal_operation(
+            handle._start_fut,
+            lambda _err, command: handle._apply_cancel_command(command),
+            reraise_on_workflow_cancellation=True,
+        )
+        return handle
 
     #### Miscellaneous helpers ####
     # These are in alphabetical order.

tests/contrib/opentelemetry/test_opentelemetry.py

@@ -88,6 +88,34 @@ class TracingWorkflowActionActivity:
     fail_on_non_replay_before_complete: bool = False
 
 
+@workflow.defn
+class LegacySignalWithStartHeaderWorkflow:
+    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 LegacySignalWithStartCallerWorkflow:
+    @workflow.run
+    async def run(self, target_id: str, task_queue: str) -> str:
+        handle = await workflow.signal_with_start_workflow(
+            LegacySignalWithStartHeaderWorkflow.run,
+            id=target_id,
+            task_queue=task_queue,
+            signal=LegacySignalWithStartHeaderWorkflow.notify,
+        )
+        return handle.id
+
+
 @dataclass
 class TracingWorkflowActionContinueAsNew:
     param: TracingWorkflowParam
@@ -229,6 +257,38 @@ def update_validator(self) -> None:
         pass
 
 
+async def test_legacy_otel_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()
+    tracer = provider.get_tracer(__name__)
+    config = client.config()
+    config["interceptors"] = [TracingInterceptor(tracer)]
+    client = Client(**config)
+
+    async with Worker(
+        client,
+        task_queue=f"signal-with-start-{uuid.uuid4()}",
+        workflows=[
+            LegacySignalWithStartCallerWorkflow,
+            LegacySignalWithStartHeaderWorkflow,
+        ],
+        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(
+                LegacySignalWithStartCallerWorkflow.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
+
+
 async def test_opentelemetry_tracing(client: Client, env: WorkflowEnvironment):
     # TODO(cretz): Fix
     if env.supports_time_skipping:

tests/contrib/opentelemetry/test_opentelemetry_plugin.py

@@ -122,6 +122,34 @@ async def run(self):
         return
 
 
+@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
+
+
 async def test_otel_tracing_basic(client: Client, reset_otel_tracer_provider: Any):  # type: ignore[reportUnusedParameter]
     exporter = InMemorySpanExporter()
     provider = create_tracer_provider()
@@ -169,6 +197,35 @@ async def test_otel_tracing_basic(client: Client, reset_otel_tracer_provider: An
     )
 
 
+async def test_otel_workflow_signal_with_start_propagates_trace_headers(
+    client: Client,
+    env: WorkflowEnvironment,
+    reset_otel_tracer_provider: Any,  # type: ignore[reportUnusedParameter]
+):
+    if env.supports_time_skipping:
+        pytest.skip("Nexus tests don't work with the Java test server")
+    provider = create_tracer_provider()
+    opentelemetry.trace.set_tracer_provider(provider)
+    config = client.config()
+    config["plugins"] = [OpenTelemetryPlugin()]
+    client = Client(**config)
+
+    async with new_worker(
+        client, SignalWithStartCallerWorkflow, SignalWithStartHeaderWorkflow
+    ) as worker:
+        target_id = f"signal-with-start-target-{uuid.uuid4()}"
+        with get_tracer(__name__).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,
+                execution_timeout=timedelta(seconds=3),
+            )
+            assert await caller.result() == target_id
+        assert await client.get_workflow_handle(target_id).result() is True
+
+
 @workflow.defn
 class ComprehensiveWorkflow:
     def __init__(self) -> None:

tests/nexus/test_temporal_system_nexus.py

@@ -36,6 +36,7 @@
 from temporalio.worker import (
     Interceptor,
     StartNexusOperationInput,
+    StartSystemNexusOperationInput,
     Worker,
     WorkflowInboundInterceptor,
     WorkflowInterceptorClassInput,
@@ -253,6 +254,12 @@ async def start_nexus_operation(
         interceptor_traces.append(("workflow.start_nexus_operation", input))
         return await super().start_nexus_operation(input)
 
+    async def start_system_nexus_operation(
+        self, input: StartSystemNexusOperationInput[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]
@@ -269,8 +276,8 @@ def _assert_stored_payloads_include(
 def _assert_start_nexus_operation_interceptor_trace() -> None:
     assert len(interceptor_traces) == 1
     trace_name, trace_value = interceptor_traces.pop()
-    assert trace_name == "workflow.start_nexus_operation"
-    trace_input = cast(StartNexusOperationInput[Any, Any], trace_value)
+    assert trace_name == "workflow.start_system_nexus_operation"
+    trace_input = cast(StartSystemNexusOperationInput[Any, Any], trace_value)
     request = trace_input.input
     assert request.id == "system-nexus-workflow-id"
     assert request.signal == "test-signal"