Use endpoint for system Nexus detection

Author: tconley1428

scripts/gen_payload_visitor.py

@@ -188,13 +188,11 @@ async def visit(
     async def _visit_nexus_operation_input_payload(
         self,
         fs: VisitorFunctions,
-        service: str,
-        operation: str,
+        endpoint: str,
         payload: Payload,
     ) -> None:
         new_payload = await temporalio.nexus.system.maybe_visit_payload(
-            service,
-            operation,
+            endpoint,
             payload,
             fs,
             self.skip_search_attributes,
@@ -292,8 +290,7 @@ def walk(self, desc: Descriptor) -> bool:
                     (
                         "system_nexus",
                         field.name,
-                        "o.service",
-                        "o.operation",
+                        "o.endpoint",
                         "o.input",
                     )
                 )
@@ -406,11 +403,11 @@ def walk(self, desc: Descriptor) -> bool:
                         )
                     )
                 elif item[0] == "system_nexus":
-                    _, field_name, service_expr, operation_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, {service_expr}, {operation_expr}, {payload_expr}\n"
+                        f"                fs, {endpoint_expr}, {payload_expr}\n"
                         "            )"
                     )
                 else:  # oneof_group

temporalio/bridge/_visitor.py

@@ -58,13 +58,11 @@ async def visit(self, fs: VisitorFunctions, root: Any) -> None:
     async def _visit_nexus_operation_input_payload(
         self,
         fs: VisitorFunctions,
-        service: str,
-        operation: str,
+        endpoint: str,
         payload: Payload,
     ) -> None:
         new_payload = await temporalio.nexus.system.maybe_visit_payload(
-            service,
-            operation,
+            endpoint,
             payload,
             fs,
             self.skip_search_attributes,
@@ -476,9 +474,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.service, o.operation, o.input
-            )
+            await self._visit_nexus_operation_input_payload(fs, o.endpoint, o.input)
 
     async def _visit_coresdk_workflow_commands_WorkflowCommand(
         self, fs: VisitorFunctions, o: Any

temporalio/nexus/system/__init__.py

@@ -2,16 +2,12 @@
 
 from __future__ import annotations
 
-import typing
-
-import google.protobuf.message
-import nexusrpc
-
 import temporalio.api.common.v1
 import temporalio.converter
 from temporalio.bridge._visitor_functions import VisitorFunctions
 from temporalio.converter import BinaryProtoPayloadConverter, CompositePayloadConverter
-from temporalio.nexus.system import workflow_service
+
+TEMPORAL_SYSTEM_ENDPOINT = "__temporal_system"
 
 
 class SystemNexusPayloadConverter(CompositePayloadConverter):
@@ -22,32 +18,23 @@ def __init__(self) -> None:
         super().__init__(BinaryProtoPayloadConverter())
 
 
-def _operation(
-    service: str, operation: str
-) -> nexusrpc.Operation[typing.Any, typing.Any] | None:
-    return workflow_service.__nexus_operation_registry__.get((service, operation))
+def is_system_endpoint(endpoint: str) -> bool:
+    """Return whether a Nexus endpoint is the Temporal system endpoint."""
+    return endpoint == TEMPORAL_SYSTEM_ENDPOINT
 
 
 async def maybe_visit_payload(
-    service: str,
-    operation: str,
+    endpoint: str,
     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 a recognized system Nexus envelope."""
-    operation_def = _operation(service, operation)
-    if operation_def is None:
-        return None
-    input_type = operation_def.input_type
-    if not (
-        isinstance(input_type, type)
-        and issubclass(input_type, google.protobuf.message.Message)
-    ):
+    """Visit nested payloads if the payload is for the Temporal system endpoint."""
+    if not is_system_endpoint(endpoint):
         return None
 
     payload_converter = get_payload_converter()
-    value = payload_converter.from_payload(payload, input_type)
+    value = payload_converter.from_payload(payload)
     from ._payload_visitor import PayloadVisitor
 
     await PayloadVisitor(skip_search_attributes=skip_search_attributes).visit(
@@ -56,19 +43,15 @@ async def maybe_visit_payload(
     return payload_converter.to_payload(value)
 
 
-def is_system_operation(service: str, operation: str) -> bool:
-    """Return whether a Nexus operation uses a generated system envelope."""
-    return _operation(service, operation) is not None
-
-
 def get_payload_converter() -> temporalio.converter.PayloadConverter:
     """Return the fixed payload converter for system Nexus outer envelopes."""
     return SystemNexusPayloadConverter()
 
 
 __all__ = [
+    "TEMPORAL_SYSTEM_ENDPOINT",
     "get_payload_converter",
-    "is_system_operation",
+    "is_system_endpoint",
     "maybe_visit_payload",
     "SystemNexusPayloadConverter",
 ]

temporalio/nexus/system/_payload_visitor.py

@@ -58,13 +58,11 @@ async def visit(self, fs: VisitorFunctions, root: Any) -> None:
     async def _visit_nexus_operation_input_payload(
         self,
         fs: VisitorFunctions,
-        service: str,
-        operation: str,
+        endpoint: str,
         payload: Payload,
     ) -> None:
         new_payload = await temporalio.nexus.system.maybe_visit_payload(
-            service,
-            operation,
+            endpoint,
             payload,
             fs,
             self.skip_search_attributes,

temporalio/worker/_workflow_instance.py

@@ -2090,9 +2090,7 @@ async def operation_handle_fn() -> OutputT:
 
         payload_converter = (
             temporalio.nexus.system.get_payload_converter()
-            if temporalio.nexus.system.is_system_operation(
-                input.service, input.operation_name
-            )
+            if temporalio.nexus.system.is_system_endpoint(input.endpoint)
             else self._context_free_payload_converter
         )
         handle = _NexusOperationHandle(

tests/nexus/test_temporal_system_nexus.py

@@ -15,6 +15,10 @@
 import temporalio.converter
 import temporalio.nexus.system as nexus_system
 from temporalio import workflow
+from temporalio.bridge._visitor import PayloadVisitor
+from temporalio.bridge.proto.workflow_completion.workflow_completion_pb2 import (
+    WorkflowActivationCompletion,
+)
 from temporalio.client import Client
 from temporalio.converter import ExternalStorage, PayloadCodec
 from temporalio.testing import WorkflowEnvironment
@@ -139,6 +143,89 @@ def _assert_start_nexus_operation_interceptor_trace() -> None:
     assert request.workflow_type.name == "test-workflow"
 
 
+class _MarkingPayloadVisitor:
+    def __init__(self) -> None:
+        self.visited_payload_count = 0
+        self.system_envelope_count = 0
+
+    async def visit_payload(self, payload: temporalio.api.common.v1.Payload) -> None:
+        self.visited_payload_count += 1
+        payload.metadata["visited"] = b"true"
+
+    async def visit_payloads(
+        self, payloads: Sequence[temporalio.api.common.v1.Payload]
+    ) -> None:
+        for payload in payloads:
+            await self.visit_payload(payload)
+
+    async def visit_system_nexus_envelope(
+        self, payload: temporalio.api.common.v1.Payload
+    ) -> None:
+        _ = payload
+        self.system_envelope_count += 1
+
+
+def _new_schedule_nexus_completion(
+    endpoint: str, payload: temporalio.api.common.v1.Payload
+) -> WorkflowActivationCompletion:
+    completion = WorkflowActivationCompletion()
+    command = completion.successful.commands.add()
+    schedule = command.schedule_nexus_operation
+    schedule.seq = 1
+    schedule.endpoint = endpoint
+    schedule.service = "not-a-registered-system-service"
+    schedule.operation = "NotARegisteredSystemOperation"
+    schedule.input.CopyFrom(payload)
+    return completion
+
+
+def _new_system_nexus_request_payload() -> temporalio.api.common.v1.Payload:
+    nested_payload = temporalio.converter.PayloadConverter.default.to_payload(
+        "workflow-input"
+    )
+    assert nested_payload is not None
+    request = workflowservice_pb2.SignalWithStartWorkflowExecutionRequest()
+    request.input.payloads.add().CopyFrom(nested_payload)
+    payload = nexus_system.get_payload_converter().to_payload(request)
+    assert payload is not None
+    return payload
+
+
+async def test_schedule_system_nexus_endpoint_ignores_operation_registry() -> None:
+    completion = _new_schedule_nexus_completion(
+        nexus_system.TEMPORAL_SYSTEM_ENDPOINT,
+        _new_system_nexus_request_payload(),
+    )
+    visitor = _MarkingPayloadVisitor()
+
+    await PayloadVisitor().visit(visitor, completion)
+
+    schedule = completion.successful.commands[0].schedule_nexus_operation
+    decoded = nexus_system.get_payload_converter().from_payload(schedule.input)
+    assert isinstance(
+        decoded, workflowservice_pb2.SignalWithStartWorkflowExecutionRequest
+    )
+    assert decoded.input.payloads[0].metadata["visited"] == b"true"
+    assert "visited" not in schedule.input.metadata
+    assert visitor.visited_payload_count == 1
+    assert visitor.system_envelope_count == 1
+
+
+async def test_schedule_non_system_nexus_visits_input_as_regular_payload() -> None:
+    completion = _new_schedule_nexus_completion(
+        "not-the-system-endpoint",
+        _new_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"
+    assert visitor.visited_payload_count == 1
+    assert visitor.system_envelope_count == 0
+
+
 def _build_proto_sample(message_type: type[Message]) -> Message:
     message = message_type()
     _populate_proto_sample(message)