diff --git a/CHANGELOG.md b/CHANGELOG.md index 6230f2754..356ac3934 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -57,6 +57,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 diff --git a/scripts/gen_payload_visitor.py b/scripts/gen_payload_visitor.py index c2bc15837..8fc707fbd 100644 --- a/scripts/gen_payload_visitor.py +++ b/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 diff --git a/temporalio/bridge/_visitor.py b/temporalio/bridge/_visitor.py index a9956e12b..21c3da88d 100644 --- a/temporalio/bridge/_visitor.py +++ b/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 diff --git a/temporalio/bridge/_visitor_functions.py b/temporalio/bridge/_visitor_functions.py index 548a0ea94..00b25248f 100644 --- a/temporalio/bridge/_visitor_functions.py +++ b/temporalio/bridge/_visitor_functions.py @@ -1,7 +1,7 @@ from __future__ import annotations import asyncio -from typing import Protocol +from typing import Protocol, runtime_checkable from google.protobuf.internal.containers import RepeatedCompositeFieldContainer @@ -26,6 +26,19 @@ async def visit_system_nexus_envelope(self, payload: Payload) -> None: return None +@runtime_checkable +class CheckpointingVisitorFunctions(VisitorFunctions, Protocol): + """Visitor functions that can await a scoped set of deferred visits.""" + + def checkpoint(self) -> int: + """Return a marker for visits scheduled after this point.""" + ... + + async def drain_since(self, checkpoint: int) -> None: + """Wait for visits scheduled after ``checkpoint`` to finish.""" + ... + + class BoundedVisitorFunctions(VisitorFunctions): """Wraps VisitorFunctions to cap concurrent payload visits via a semaphore. @@ -74,16 +87,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() diff --git a/temporalio/nexus/system/__init__.py b/temporalio/nexus/system/__init__.py index 14a43cb72..09f708c29 100644 --- a/temporalio/nexus/system/__init__.py +++ b/temporalio/nexus/system/__init__.py @@ -12,8 +12,12 @@ from typing import Any import temporalio.api.common.v1 +import temporalio.common import temporalio.converter -from temporalio.bridge._visitor_functions import VisitorFunctions +from temporalio.bridge._visitor_functions import ( + CheckpointingVisitorFunctions, + VisitorFunctions, +) from temporalio.converter import BinaryProtoPayloadConverter, CompositePayloadConverter from temporalio.converter._payload_converter import ( _TemporalTransferTypePayloadConverter, @@ -23,6 +27,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 +58,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 +113,33 @@ 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 ._payload_visitor import PayloadVisitor - await PayloadVisitor(skip_search_attributes=skip_search_attributes).visit( - visitor_functions, value - ) + payload_visitor = PayloadVisitor(skip_search_attributes=skip_search_attributes) + if isinstance(visitor_functions, CheckpointingVisitorFunctions): + checkpoint = visitor_functions.checkpoint() + await payload_visitor.visit(visitor_functions, value) + await visitor_functions.drain_since(checkpoint) + else: + await payload_visitor.visit(visitor_functions, value) return payload_converter.to_payload(value) diff --git a/temporalio/nexus/system/_payload_visitor.py b/temporalio/nexus/system/_payload_visitor.py index 4f194168f..1610a06af 100644 --- a/temporalio/nexus/system/_payload_visitor.py +++ b/temporalio/nexus/system/_payload_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 ): diff --git a/tests/nexus/test_temporal_system_nexus.py b/tests/nexus/test_temporal_system_nexus.py index eb8ee603c..76d2f0e35 100644 --- a/tests/nexus/test_temporal_system_nexus.py +++ b/tests/nexus/test_temporal_system_nexus.py @@ -34,6 +34,7 @@ from tests.test_extstore import InMemoryTestDriver interceptor_traces: list[tuple[str, object]] = [] +SYSTEM_NEXUS_PAYLOAD_METADATA_KEY = "__temporal_system_payload" @workflow.defn @@ -193,9 +194,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,10 +222,12 @@ 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() @@ -226,6 +235,13 @@ async def test_schedule_non_system_nexus_visits_input_as_regular_payload() -> No 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 +367,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 diff --git a/tests/worker/test_visitor.py b/tests/worker/test_visitor.py index 3c3df42c1..d9606624e 100644 --- a/tests/worker/test_visitor.py +++ b/tests/worker/test_visitor.py @@ -212,6 +212,58 @@ async def test_visit_payloads_on_other_commands(): assert ur.completed.metadata["visited"] +async def test_system_nexus_envelope_is_detected_in_generic_payload_field(): + class SystemNexusVisitor(Visitor): + def __init__(self) -> None: + self.visited_payload_count = 0 + self.system_envelope_count = 0 + + async def visit_payload(self, payload: Payload) -> None: + self.visited_payload_count += 1 + await super().visit_payload(payload) + + async def visit_payloads(self, payloads: MutableSequence[Payload]) -> None: + self.visited_payload_count += len(payloads) + await super().visit_payloads(payloads) + + async def visit_system_nexus_envelope(self, payload: Payload) -> None: + _ = payload + self.system_envelope_count += 1 + + system_request = workflowservice_pb2.SignalWithStartWorkflowExecutionRequest( + input=Payloads(payloads=[Payload(data=b"workflow-input")]), + ) + payload_converter = nexus_system._get_payload_converter( + temporalio.converter.default().payload_converter + ) + system_payload = payload_converter.to_payload(system_request) + assert system_payload is not None + comp = WorkflowActivationCompletion( + run_id="3", + successful=Success( + commands=[ + WorkflowCommand( + update_response=UpdateResponse(completed=system_payload), + ) + ] + ), + ) + visitor = SystemNexusVisitor() + + await PayloadVisitor(concurrency_limit=3).visit(visitor, comp) + + completed = comp.successful.commands[0].update_response.completed + assert completed.metadata["__temporal_system_payload"] == b"true" + assert "visited" not in completed.metadata + decoded = payload_converter.from_payload(completed) + assert isinstance( + decoded, workflowservice_pb2.SignalWithStartWorkflowExecutionRequest + ) + assert decoded.input.payloads[0].metadata["visited"] == b"True" + assert visitor.visited_payload_count == 1 + assert visitor.system_envelope_count == 1 + + async def test_concurrent_throughput(): """Demonstrate that concurrent visitation is faster than serialized for I/O-bound codecs.""" N_CMDS = 10