Propagate tracing in system Nexus request headers

Author: tconley1428

scripts/nex_gen_support.py

@@ -167,6 +167,21 @@ def memo_to_proto(
     return message
 
 
+def header_from_proto(
+    proto: common_pb2.Header,
+) -> collections.abc.Mapping[str, object]:
+    return {key: _payload_to_value(value) for key, value in proto.fields.items()}
+
+
+def header_to_proto(
+    header: collections.abc.Mapping[str, object],
+) -> common_pb2.Header:
+    message = common_pb2.Header()
+    for key, value in header.items():
+        message.fields[key].CopyFrom(_value_to_payload(value))
+    return message
+
+
 def duration_from_proto(proto: google.protobuf.duration_pb2.Duration) -> timedelta:
     return proto.ToTimedelta()
 

temporalio/contrib/opentelemetry/_interceptor.py

@@ -833,11 +833,13 @@ async def start_nexus_operation(
     async def start_system_nexus_operation(
         self, input: temporalio.worker.StartSystemNexusOperationInput[Any, Any]
     ) -> temporalio.workflow.NexusOperationHandle[Any]:
-        self.root._completed_span(
-            f"StartNexusOperation:{input.service}/{input.operation_name}",
-            kind=opentelemetry.trace.SpanKind.CLIENT,
-            add_to_outbound=input,
-        )
+        if hasattr(input.input, "headers"):
+            input.input.headers = input.input.headers or {}
+            self.root._completed_span(
+                f"StartNexusOperation:{input.service}/{input.operation_name}",
+                kind=opentelemetry.trace.SpanKind.CLIENT,
+                add_to_outbound=cast(_InputWithHeaders, input.input),
+            )
         return await super().start_system_nexus_operation(input)
 
 

temporalio/contrib/opentelemetry/_otel_interceptor.py

@@ -608,5 +608,6 @@ async def start_system_nexus_operation(
             f"StartNexusOperation:{input.service}/{input.operation_name}",
             kind=opentelemetry.trace.SpanKind.CLIENT,
         ):
-            input.headers = _context_to_headers(input.headers)
+            if hasattr(input.input, "headers"):
+                input.input.headers = _context_to_headers(input.input.headers or {})
             return await super().start_system_nexus_operation(input)

temporalio/converter/_payload_converter.py

@@ -613,26 +613,14 @@ def wrap(payload_converter: PayloadConverter) -> PayloadConverter:
         return _TemporalTransferTypePayloadConverter(payload_converter)
 
     def to_payloads(
-        self,
-        values: Sequence[Any],
-        headers: Mapping[str, temporalio.api.common.v1.Payload] | None = None,
+        self, values: Sequence[Any]
     ) -> 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

@@ -8,7 +8,7 @@
 
 import contextlib
 import contextvars
-from collections.abc import Iterator, Mapping, Sequence
+from collections.abc import Iterator, Sequence
 from dataclasses import dataclass
 from typing import Any
 
@@ -91,31 +91,28 @@ class _SystemNexusPayloadConverter(temporalio.converter.PayloadConverter):
     """Payload converter for system Nexus outer envelopes."""
 
     _user_converters: _SystemNexusUserConverters
-    _outer_payload_converter: _TemporalTransferTypePayloadConverter
-    _headers: Mapping[str, temporalio.api.common.v1.Payload] | None
+    _outer_payload_converter: temporalio.converter.PayloadConverter
 
     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(
+        self._outer_payload_converter = _TemporalTransferTypePayloadConverter.wrap(
             _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, self._headers)
+            return self._outer_payload_converter.to_payloads(values)
 
     def from_payloads(
         self,
@@ -167,12 +164,9 @@ 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, headers
-    )
+    return _SystemNexusPayloadConverter(user_payload_converter, user_failure_converter)
 
 
 def _get_serialization_context(  # pyright: ignore[reportUnusedFunction]

temporalio/nexus/system/workflow_service/__init__.py

@@ -1,4 +1,4 @@
-# Generated by nex-gen. DO NOT EDIT!
+# Generated by nexgen. DO NOT EDIT!
 
 from __future__ import annotations
 

temporalio/nexus/system/workflow_service/_support/__init__.py

@@ -1,4 +1,4 @@
-# Generated by nex-gen. DO NOT EDIT!
+# Generated by nexgen. DO NOT EDIT!
 
 from __future__ import annotations
 

temporalio/nexus/system/workflow_service/_support/nex_gen_support.py

@@ -167,6 +167,21 @@ def memo_to_proto(
     return message
 
 
+def header_from_proto(
+    proto: common_pb2.Header,
+) -> collections.abc.Mapping[str, object]:
+    return {key: _payload_to_value(value) for key, value in proto.fields.items()}
+
+
+def header_to_proto(
+    header: collections.abc.Mapping[str, object],
+) -> common_pb2.Header:
+    message = common_pb2.Header()
+    for key, value in header.items():
+        message.fields[key].CopyFrom(_value_to_payload(value))
+    return message
+
+
 def duration_from_proto(proto: google.protobuf.duration_pb2.Duration) -> timedelta:
     return proto.ToTimedelta()
 

temporalio/nexus/system/workflow_service/models.py

@@ -1,4 +1,4 @@
-# Generated by nex-gen. DO NOT EDIT!
+# Generated by nexgen. DO NOT EDIT!
 
 from __future__ import annotations
 
@@ -17,6 +17,8 @@
 from ._support import (
     duration_from_proto,
     duration_to_proto,
+    header_from_proto,
+    header_to_proto,
     memo_from_proto,
     memo_to_proto,
     payload_from_proto,
@@ -133,6 +135,9 @@ def from_transfer_type(
             if proto.HasField("user_metadata")
             else None,
             namespace=proto.namespace,
+            headers=header_from_proto(proto.header)
+            if proto.HasField("header")
+            else None,
         )
 
     @typing_extensions.override
@@ -191,6 +196,8 @@ def to_transfer_type(
                 )
             )
         message.namespace = value.namespace
+        if value.headers is not None:
+            message.header.CopyFrom(header_to_proto(value.headers))
         return message
 
 
@@ -226,6 +233,7 @@ class SignalWithStartWorkflowRequest:
     start_delay: datetime.timedelta | None = None
     user_metadata: UserMetadata | None = None
     namespace: str = dataclasses.field(default_factory=workflow_namespace)
+    headers: collections.abc.Mapping[str, typing.Any] | None = None
 
 
 class _UserMetadataTransferTypeConverter(

temporalio/nexus/system/workflow_service/operations/__init__.py

@@ -1,3 +1,3 @@
-# Generated by nex-gen. DO NOT EDIT!
+# Generated by nexgen. DO NOT EDIT!
 
 from __future__ import annotations