Mark system Nexus envelope payloads

CHANGELOG.md

@@ -69,6 +69,9 @@ to include examples, links to docs, or any other relevant information.
 
 ### Fixed
 
+- Marked system Nexus envelope payloads so nested payloads can be detected and
+  visited after the envelope is already stored as a payload.
+
 ### Security
 
 ## [1.30.0] - 2026-07-01

scripts/gen_payload_visitor.py

@@ -188,22 +188,9 @@ async def visit(
     async def _visit_nexus_operation_input_payload(
         self,
         fs: VisitorFunctions,
-        endpoint: str,
         payload: Payload,
     ) -> None:
-        new_payload = await temporalio.nexus.system._maybe_visit_payload(
-            endpoint,
-            payload,
-            fs,
-            self.skip_search_attributes,
-        )
-        if new_payload is None:
-            await self._visit_temporal_api_common_v1_Payload(fs, payload)
-            return
-
-        if new_payload is not payload:
-            payload.CopyFrom(new_payload)
-        await fs.visit_system_nexus_envelope(payload)
+        await self._visit_temporal_api_common_v1_Payload(fs, payload)
 
 """
 
@@ -218,8 +205,21 @@ def __init__(self):
         self.in_progress: set[str] = set()
         self.methods: list[str] = [
             """\
-    async def _visit_temporal_api_common_v1_Payload(self, fs: VisitorFunctions, o: Payload):
-        await fs.visit_payload(o)
+    async def _visit_temporal_api_common_v1_Payload(
+        self, fs: VisitorFunctions, payload: Payload
+    ) -> None:
+        new_payload = await temporalio.nexus.system.maybe_visit_payload(
+            payload,
+            fs,
+            self.skip_search_attributes,
+        )
+        if new_payload is None:
+            await fs.visit_payload(payload)
+            return
+
+        if new_payload is not payload:
+            payload.CopyFrom(new_payload)
+        await fs.visit_system_nexus_envelope(payload)
     """,
             """\
     async def _visit_temporal_api_common_v1_Payloads(self, fs: VisitorFunctions, o: Any):
@@ -403,11 +403,11 @@ def walk(self, desc: Descriptor) -> bool:
                         )
                     )
                 elif item[0] == "system_nexus":
-                    _, field_name, endpoint_expr, payload_expr = item
+                    _, field_name, _endpoint_expr, payload_expr = item
                     lines.append(
                         f'        if o.HasField("{field_name}"):\n'
                         "            await self._visit_nexus_operation_input_payload(\n"
-                        f"                fs, {endpoint_expr}, {payload_expr}\n"
+                        f"                fs, {payload_expr}\n"
                         "            )"
                     )
                 else:  # oneof_group

temporalio/bridge/_visitor.py

@@ -58,28 +58,26 @@ async def visit(self, fs: VisitorFunctions, root: Any) -> None:
     async def _visit_nexus_operation_input_payload(
         self,
         fs: VisitorFunctions,
-        endpoint: str,
         payload: Payload,
     ) -> None:
-        new_payload = await temporalio.nexus.system._maybe_visit_payload(
-            endpoint,
+        await self._visit_temporal_api_common_v1_Payload(fs, payload)
+
+    async def _visit_temporal_api_common_v1_Payload(
+        self, fs: VisitorFunctions, payload: Payload
+    ) -> None:
+        new_payload = await temporalio.nexus.system.maybe_visit_payload(
             payload,
             fs,
             self.skip_search_attributes,
         )
         if new_payload is None:
-            await self._visit_temporal_api_common_v1_Payload(fs, payload)
+            await fs.visit_payload(payload)
             return
 
         if new_payload is not payload:
             payload.CopyFrom(new_payload)
         await fs.visit_system_nexus_envelope(payload)
 
-    async def _visit_temporal_api_common_v1_Payload(
-        self, fs: VisitorFunctions, o: Payload
-    ):
-        await fs.visit_payload(o)
-
     async def _visit_temporal_api_common_v1_Payloads(
         self, fs: VisitorFunctions, o: Any
     ):
@@ -474,7 +472,7 @@ async def _visit_coresdk_workflow_commands_ScheduleNexusOperation(
         self, fs: VisitorFunctions, o: Any
     ):
         if o.HasField("input"):
-            await self._visit_nexus_operation_input_payload(fs, o.endpoint, o.input)
+            await self._visit_nexus_operation_input_payload(fs, o.input)
 
     async def _visit_coresdk_workflow_commands_WorkflowCommand(
         self, fs: VisitorFunctions, o: Any

temporalio/bridge/_visitor_functions.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import asyncio
-from typing import Protocol
+from abc import ABC, abstractmethod
 
 from google.protobuf.internal.containers import RepeatedCompositeFieldContainer
 
@@ -10,21 +10,31 @@
 PayloadSequence = list[Payload] | RepeatedCompositeFieldContainer[Payload]
 
 
-class VisitorFunctions(Protocol):
+class VisitorFunctions(ABC):
     """Functions invoked by generated payload visitors."""
 
+    @abstractmethod
     async def visit_payload(self, payload: Payload) -> None:
         """Visit a single payload."""
         ...
 
+    @abstractmethod
     async def visit_payloads(self, payloads: PayloadSequence) -> None:
         """Visit a sequence of payloads together."""
         ...
 
-    async def visit_system_nexus_envelope(self, payload: Payload) -> None:
+    async def visit_system_nexus_envelope(self, _payload: Payload) -> None:
         """Visit a recognized system Nexus envelope payload."""
         return None
 
+    def checkpoint(self) -> int | None:
+        """Return a marker for visits scheduled after this point, if supported."""
+        return None
+
+    async def drain_since(self, _checkpoint: int) -> None:
+        """Wait for visits scheduled after ``checkpoint`` to finish."""
+        return None
+
 
 class BoundedVisitorFunctions(VisitorFunctions):
     """Wraps VisitorFunctions to cap concurrent payload visits via a semaphore.
@@ -74,16 +84,33 @@ async def _run() -> None:
 
         self._tasks.append(asyncio.create_task(_run()))
 
+    def checkpoint(self) -> int:
+        """Return a marker for tasks scheduled after this point."""
+        return len(self._tasks)
+
+    async def drain_since(self, checkpoint: int) -> None:
+        """Wait for tasks scheduled after ``checkpoint`` to finish.
+
+        This lets system-envelope traversal finish mutating its decoded value
+        before that value is serialized again, without waiting for unrelated
+        visits that were already in progress.
+        """
+        await self._drain_tasks(self._tasks[checkpoint:])
+
     async def drain(self) -> None:
         """Wait for all in-flight background tasks to complete.
 
         On cancellation or error, cancels all remaining tasks and awaits
         them so their finally blocks run before this coroutine returns.
         """
-        if not self._tasks:
+        await self._drain_tasks(self._tasks)
+
+    async def _drain_tasks(self, tasks: list[asyncio.Task[None]]) -> None:
+        """Wait for the given tasks, cancelling all tasks if one fails."""
+        if not tasks:
             return
         try:
-            await asyncio.gather(*self._tasks)
+            await asyncio.gather(*tasks)
         except BaseException:
             for task in self._tasks:
                 task.cancel()

temporalio/nexus/system/__init__.py

@@ -12,6 +12,7 @@
 from typing import Any
 
 import temporalio.api.common.v1
+import temporalio.common
 import temporalio.converter
 from temporalio.bridge._visitor_functions import VisitorFunctions
 from temporalio.converter import BinaryProtoPayloadConverter, CompositePayloadConverter
@@ -23,6 +24,8 @@
 _user_payload_converter: contextvars.ContextVar[
     temporalio.converter.PayloadConverter | None
 ] = contextvars.ContextVar("temporal-system-nexus-user-payload-converter", default=None)
+_SYSTEM_PAYLOAD_METADATA_KEY = "__temporal_system_payload"
+_SYSTEM_PAYLOAD_METADATA_VALUE = b"true"
 
 
 @contextlib.contextmanager
@@ -52,6 +55,19 @@ def __init__(self) -> None:
         """Create a payload converter for system Nexus outer envelopes."""
         super().__init__(BinaryProtoPayloadConverter())
 
+    def to_payloads(
+        self, values: Sequence[Any]
+    ) -> list[temporalio.api.common.v1.Payload]:
+        """See base class."""
+        payloads = super().to_payloads(values)
+        for value, payload in zip(values, payloads):
+            if isinstance(value, temporalio.common.RawValue):
+                continue
+            payload.metadata[_SYSTEM_PAYLOAD_METADATA_KEY] = (
+                _SYSTEM_PAYLOAD_METADATA_VALUE
+            )
+        return payloads
+
 
 class _SystemNexusPayloadConverter(temporalio.converter.PayloadConverter):
     """Payload converter for system Nexus outer envelopes."""
@@ -94,23 +110,31 @@ def is_system_endpoint(endpoint: str) -> bool:
     return endpoint == TEMPORAL_SYSTEM_ENDPOINT
 
 
-async def _maybe_visit_payload(  # pyright: ignore[reportUnusedFunction]
-    endpoint: str,
+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,
     skip_search_attributes: bool,
 ) -> temporalio.api.common.v1.Payload | None:
-    """Visit nested payloads if the payload is for the Temporal system endpoint."""
-    if not is_system_endpoint(endpoint):
+    """Visit nested payloads if the payload is a Temporal system Nexus envelope."""
+    if not _is_system_payload(payload):
         return None
 
     payload_converter = _SystemNexusOuterPayloadConverter()
     value = payload_converter.from_payload(payload)
     from temporalio.bridge._visitor import PayloadVisitor
 
-    await PayloadVisitor(skip_search_attributes=skip_search_attributes).visit(
-        visitor_functions, value
-    )
+    payload_visitor = PayloadVisitor(skip_search_attributes=skip_search_attributes)
+    checkpoint = visitor_functions.checkpoint()
+    await payload_visitor.visit(visitor_functions, value)
+    if checkpoint is not None:
+        await visitor_functions.drain_since(checkpoint)
     return payload_converter.to_payload(value)
 
 

tests/nexus/test_temporal_system_nexus.py

@@ -16,6 +16,7 @@
 import temporalio.nexus.system as nexus_system
 from temporalio import workflow
 from temporalio.bridge._visitor import PayloadVisitor
+from temporalio.bridge._visitor_functions import VisitorFunctions
 from temporalio.bridge.proto.workflow_completion.workflow_completion_pb2 import (
     WorkflowActivationCompletion,
 )
@@ -34,6 +35,7 @@
 from tests.test_extstore import InMemoryTestDriver
 
 interceptor_traces: list[tuple[str, object]] = []
+SYSTEM_NEXUS_PAYLOAD_METADATA_KEY = "__temporal_system_payload"
 
 
 @workflow.defn
@@ -143,7 +145,7 @@ def _assert_start_nexus_operation_interceptor_trace() -> None:
     assert request.workflow_type.name == "test-workflow"
 
 
-class _MarkingPayloadVisitor:
+class _MarkingPayloadVisitor(VisitorFunctions):
     def __init__(self) -> None:
         self.visited_payload_count = 0
         self.system_envelope_count = 0
@@ -193,9 +195,15 @@ def _new_system_nexus_request_payload() -> temporalio.api.common.v1.Payload:
     return payload
 
 
-async def test_schedule_system_nexus_endpoint_ignores_operation_registry() -> None:
+def _new_unmarked_system_nexus_request_payload() -> temporalio.api.common.v1.Payload:
+    payload = _new_system_nexus_request_payload()
+    del payload.metadata[SYSTEM_NEXUS_PAYLOAD_METADATA_KEY]
+    return payload
+
+
+async def test_schedule_marked_system_nexus_payload_ignores_endpoint() -> None:
     completion = _new_schedule_nexus_completion(
-        nexus_system.TEMPORAL_SYSTEM_ENDPOINT,
+        "not-the-system-endpoint",
         _new_system_nexus_request_payload(),
     )
     visitor = _MarkingPayloadVisitor()
@@ -215,17 +223,26 @@ async def test_schedule_system_nexus_endpoint_ignores_operation_registry() -> No
     assert visitor.system_envelope_count == 1
 
 
-async def test_schedule_non_system_nexus_visits_input_as_regular_payload() -> None:
+async def test_schedule_unmarked_system_nexus_payload_visits_input_as_regular_payload() -> (
+    None
+):
     completion = _new_schedule_nexus_completion(
-        "not-the-system-endpoint",
-        _new_system_nexus_request_payload(),
+        nexus_system.TEMPORAL_SYSTEM_ENDPOINT,
+        _new_unmarked_system_nexus_request_payload(),
     )
     visitor = _MarkingPayloadVisitor()
 
     await PayloadVisitor().visit(visitor, completion)
 
     schedule = completion.successful.commands[0].schedule_nexus_operation
     assert schedule.input.metadata["visited"] == b"true"
+    decoded = nexus_system._get_payload_converter(
+        temporalio.converter.default().payload_converter
+    ).from_payload(schedule.input)
+    assert isinstance(
+        decoded, workflowservice_pb2.SignalWithStartWorkflowExecutionRequest
+    )
+    assert "visited" not in decoded.input.payloads[0].metadata
     assert visitor.visited_payload_count == 1
     assert visitor.system_envelope_count == 0
 
@@ -351,6 +368,7 @@ def test_system_nexus_proto_roundtrip(message_type: type[Message]) -> None:
     assert payload is not None
     assert payload.metadata["encoding"] == b"binary/protobuf"
     assert payload.metadata["messageType"] == message_type.DESCRIPTOR.full_name.encode()
+    assert payload.metadata[SYSTEM_NEXUS_PAYLOAD_METADATA_KEY] == b"true"
     roundtripped = payload_converter.from_payload(payload, message_type)
     assert isinstance(roundtripped, message_type)
     assert roundtripped == proto_value