Refactor: don't use default PayloadVisitor to set command context

Author: dandavison

scripts/gen_payload_visitor.py

@@ -62,55 +62,6 @@ def emit_singular(
         await self._visit_{child_method}(fs, {access_expr})"""
 
 
-def emit_singular_with_seq(
-    field_name: str, access_expr: str, child_method: str, presence_word: str
-) -> str:
-    # Helper to emit a singular field visit that sets the seq contextvar, with presence check but
-    # without headers guard since this is used for commands only.
-
-    # Map field names to command types
-    field_to_command_type = {
-        # Commands
-        "schedule_activity": "COMMAND_TYPE_SCHEDULE_ACTIVITY_TASK",
-        "schedule_local_activity": "COMMAND_TYPE_SCHEDULE_ACTIVITY_TASK",
-        "start_child_workflow_execution": "COMMAND_TYPE_START_CHILD_WORKFLOW_EXECUTION",
-        "signal_external_workflow_execution": "COMMAND_TYPE_SIGNAL_EXTERNAL_WORKFLOW_EXECUTION",
-        "schedule_nexus_operation": "COMMAND_TYPE_SCHEDULE_NEXUS_OPERATION",
-        "request_cancel_activity": "COMMAND_TYPE_REQUEST_CANCEL_ACTIVITY_TASK",
-        "request_cancel_local_activity": "COMMAND_TYPE_REQUEST_CANCEL_ACTIVITY_TASK",
-        "request_cancel_external_workflow_execution": "COMMAND_TYPE_REQUEST_CANCEL_EXTERNAL_WORKFLOW_EXECUTION",
-        "request_cancel_nexus_operation": "COMMAND_TYPE_REQUEST_CANCEL_NEXUS_OPERATION",
-        "cancel_timer": "COMMAND_TYPE_CANCEL_TIMER",
-        "cancel_signal_workflow": "COMMAND_TYPE_CANCEL_SIGNAL_WORKFLOW",
-        "start_timer": "COMMAND_TYPE_START_TIMER",
-        # Resolutions (use the corresponding command type)
-        "resolve_activity": "COMMAND_TYPE_SCHEDULE_ACTIVITY_TASK",
-        "resolve_child_workflow_execution_start": "COMMAND_TYPE_START_CHILD_WORKFLOW_EXECUTION",
-        "resolve_child_workflow_execution": "COMMAND_TYPE_START_CHILD_WORKFLOW_EXECUTION",
-        "resolve_signal_external_workflow": "COMMAND_TYPE_SIGNAL_EXTERNAL_WORKFLOW_EXECUTION",
-        "resolve_request_cancel_external_workflow": "COMMAND_TYPE_REQUEST_CANCEL_EXTERNAL_WORKFLOW_EXECUTION",
-        "resolve_nexus_operation_start": "COMMAND_TYPE_SCHEDULE_NEXUS_OPERATION",
-        "resolve_nexus_operation": "COMMAND_TYPE_SCHEDULE_NEXUS_OPERATION",
-    }
-
-    command_type = field_to_command_type.get(field_name)
-    if not command_type:
-        raise ValueError(f"Unknown field with seq: {field_name}")
-
-    return f"""\
-        {presence_word} o.HasField("{field_name}"):
-            token = current_command_info.set(
-                CommandInfo(
-                    command_type=temporalio.api.enums.v1.command_type_pb2.CommandType.{command_type},
-                    command_seq={access_expr}.seq,
-                )
-            )
-            try:
-                await self._visit_{child_method}(fs, {access_expr})
-            finally:
-                current_command_info.reset(token)"""
-
-
 class VisitorGenerator:
     def generate(self, roots: list[Descriptor]) -> str:
         """
@@ -129,26 +80,11 @@ def generate(self, roots: list[Descriptor]) -> str:
         header = """
 # This file is generated by gen_payload_visitor.py. Changes should be made there.
 import abc
-import contextvars
-from dataclasses import dataclass
-from typing import Any, MutableSequence, Optional
+from typing import Any, MutableSequence
 
-import temporalio.api.enums.v1.command_type_pb2
 from temporalio.api.common.v1.message_pb2 import Payload
 
 
-@dataclass(frozen=True)
-class CommandInfo:
-    \"\"\"Information identifying a specific command instance.\"\"\"
-
-    command_type: temporalio.api.enums.v1.command_type_pb2.CommandType
-    command_seq: int
-
-
-current_command_info: contextvars.ContextVar[Optional[CommandInfo]] = (
-    contextvars.ContextVar("current_command_info", default=None)
-)
-
 class VisitorFunctions(abc.ABC):
     \"\"\"Set of functions which can be called by the visitor. 
     Allows handling payloads as a sequence.
@@ -318,29 +254,6 @@ def walk(self, desc: Descriptor) -> bool:
                         )
                     )
 
-        commands_with_seq = {
-            "cancel_signal_workflow",
-            "cancel_timer",
-            "request_cancel_activity",
-            "request_cancel_external_workflow_execution",
-            "request_cancel_local_activity",
-            "request_cancel_nexus_operation",
-            "schedule_activity",
-            "schedule_local_activity",
-            "schedule_nexus_operation",
-            "signal_external_workflow_execution",
-            "start_child_workflow_execution",
-            "start_timer",
-        }
-        activation_jobs_with_seq = {
-            "resolve_activity",
-            "resolve_child_workflow_execution_start",
-            "resolve_child_workflow_execution",
-            "resolve_nexus_operation_start",
-            "resolve_nexus_operation",
-            "resolve_request_cancel_external_workflow",
-            "resolve_signal_external_workflow",
-        }
         # Process oneof fields as if/elif chains
         for oneof_idx, fields in oneof_fields.items():
             oneof_lines = []
@@ -352,25 +265,9 @@ def walk(self, desc: Descriptor) -> bool:
                 if child_has_payload:
                     if_word = "if" if first else "elif"
                     first = False
-                    if (
-                        desc.full_name == "coresdk.workflow_commands.WorkflowCommand"
-                        and field.name in commands_with_seq
-                    ):
-                        line = emit_singular_with_seq(
-                            field.name, f"o.{field.name}", name_for(child_desc), if_word
-                        )
-                    elif (
-                        desc.full_name
-                        == "coresdk.workflow_activation.WorkflowActivationJob"
-                        and field.name in activation_jobs_with_seq
-                    ):
-                        line = emit_singular_with_seq(
-                            field.name, f"o.{field.name}", name_for(child_desc), if_word
-                        )
-                    else:
-                        line = emit_singular(
-                            field.name, f"o.{field.name}", name_for(child_desc), if_word
-                        )
+                    line = emit_singular(
+                        field.name, f"o.{field.name}", name_for(child_desc), if_word
+                    )
                     oneof_lines.append(line)
             if oneof_lines:
                 lines.extend(oneof_lines)

temporalio/bridge/_visitor.py

@@ -1,26 +1,10 @@
 # This file is generated by gen_payload_visitor.py. Changes should be made there.
 import abc
-import contextvars
-from dataclasses import dataclass
-from typing import Any, MutableSequence, Optional
+from typing import Any, MutableSequence
 
-import temporalio.api.enums.v1.command_type_pb2
 from temporalio.api.common.v1.message_pb2 import Payload
 
 
-@dataclass(frozen=True)
-class CommandInfo:
-    """Information identifying a specific command instance."""
-
-    command_type: temporalio.api.enums.v1.command_type_pb2.CommandType.ValueType
-    command_seq: int
-
-
-current_command_info: contextvars.ContextVar[Optional[CommandInfo]] = (
-    contextvars.ContextVar("current_command_info", default=None)
-)
-
-
 class VisitorFunctions(abc.ABC):
     """Set of functions which can be called by the visitor.
     Allows handling payloads as a sequence.
@@ -263,100 +247,35 @@ async def _visit_coresdk_workflow_activation_WorkflowActivationJob(self, fs, o):
                 fs, o.signal_workflow
             )
         elif o.HasField("resolve_activity"):
-            token = current_command_info.set(
-                CommandInfo(
-                    command_type=temporalio.api.enums.v1.command_type_pb2.CommandType.COMMAND_TYPE_SCHEDULE_ACTIVITY_TASK,
-                    command_seq=o.resolve_activity.seq,
-                )
+            await self._visit_coresdk_workflow_activation_ResolveActivity(
+                fs, o.resolve_activity
             )
-            try:
-                await self._visit_coresdk_workflow_activation_ResolveActivity(
-                    fs, o.resolve_activity
-                )
-            finally:
-                current_command_info.reset(token)
         elif o.HasField("resolve_child_workflow_execution_start"):
-            token = current_command_info.set(
-                CommandInfo(
-                    command_type=temporalio.api.enums.v1.command_type_pb2.CommandType.COMMAND_TYPE_START_CHILD_WORKFLOW_EXECUTION,
-                    command_seq=o.resolve_child_workflow_execution_start.seq,
-                )
+            await self._visit_coresdk_workflow_activation_ResolveChildWorkflowExecutionStart(
+                fs, o.resolve_child_workflow_execution_start
             )
-            try:
-                await self._visit_coresdk_workflow_activation_ResolveChildWorkflowExecutionStart(
-                    fs, o.resolve_child_workflow_execution_start
-                )
-            finally:
-                current_command_info.reset(token)
         elif o.HasField("resolve_child_workflow_execution"):
-            token = current_command_info.set(
-                CommandInfo(
-                    command_type=temporalio.api.enums.v1.command_type_pb2.CommandType.COMMAND_TYPE_START_CHILD_WORKFLOW_EXECUTION,
-                    command_seq=o.resolve_child_workflow_execution.seq,
-                )
+            await self._visit_coresdk_workflow_activation_ResolveChildWorkflowExecution(
+                fs, o.resolve_child_workflow_execution
             )
-            try:
-                await self._visit_coresdk_workflow_activation_ResolveChildWorkflowExecution(
-                    fs, o.resolve_child_workflow_execution
-                )
-            finally:
-                current_command_info.reset(token)
         elif o.HasField("resolve_signal_external_workflow"):
-            token = current_command_info.set(
-                CommandInfo(
-                    command_type=temporalio.api.enums.v1.command_type_pb2.CommandType.COMMAND_TYPE_SIGNAL_EXTERNAL_WORKFLOW_EXECUTION,
-                    command_seq=o.resolve_signal_external_workflow.seq,
-                )
+            await self._visit_coresdk_workflow_activation_ResolveSignalExternalWorkflow(
+                fs, o.resolve_signal_external_workflow
             )
-            try:
-                await self._visit_coresdk_workflow_activation_ResolveSignalExternalWorkflow(
-                    fs, o.resolve_signal_external_workflow
-                )
-            finally:
-                current_command_info.reset(token)
         elif o.HasField("resolve_request_cancel_external_workflow"):
-            token = current_command_info.set(
-                CommandInfo(
-                    command_type=temporalio.api.enums.v1.command_type_pb2.CommandType.COMMAND_TYPE_REQUEST_CANCEL_EXTERNAL_WORKFLOW_EXECUTION,
-                    command_seq=o.resolve_request_cancel_external_workflow.seq,
-                )
+            await self._visit_coresdk_workflow_activation_ResolveRequestCancelExternalWorkflow(
+                fs, o.resolve_request_cancel_external_workflow
             )
-            try:
-                await self._visit_coresdk_workflow_activation_ResolveRequestCancelExternalWorkflow(
-                    fs, o.resolve_request_cancel_external_workflow
-                )
-            finally:
-                current_command_info.reset(token)
         elif o.HasField("do_update"):
             await self._visit_coresdk_workflow_activation_DoUpdate(fs, o.do_update)
         elif o.HasField("resolve_nexus_operation_start"):
-            token = current_command_info.set(
-                CommandInfo(
-                    command_type=temporalio.api.enums.v1.command_type_pb2.CommandType.COMMAND_TYPE_SCHEDULE_NEXUS_OPERATION,
-                    command_seq=o.resolve_nexus_operation_start.seq,
-                )
+            await self._visit_coresdk_workflow_activation_ResolveNexusOperationStart(
+                fs, o.resolve_nexus_operation_start
             )
-            try:
-                await (
-                    self._visit_coresdk_workflow_activation_ResolveNexusOperationStart(
-                        fs, o.resolve_nexus_operation_start
-                    )
-                )
-            finally:
-                current_command_info.reset(token)
         elif o.HasField("resolve_nexus_operation"):
-            token = current_command_info.set(
-                CommandInfo(
-                    command_type=temporalio.api.enums.v1.command_type_pb2.CommandType.COMMAND_TYPE_SCHEDULE_NEXUS_OPERATION,
-                    command_seq=o.resolve_nexus_operation.seq,
-                )
+            await self._visit_coresdk_workflow_activation_ResolveNexusOperation(
+                fs, o.resolve_nexus_operation
             )
-            try:
-                await self._visit_coresdk_workflow_activation_ResolveNexusOperation(
-                    fs, o.resolve_nexus_operation
-                )
-            finally:
-                current_command_info.reset(token)
 
     async def _visit_coresdk_workflow_activation_WorkflowActivation(self, fs, o):
         for v in o.jobs:
@@ -452,18 +371,9 @@ async def _visit_coresdk_workflow_commands_WorkflowCommand(self, fs, o):
         if o.HasField("user_metadata"):
             await self._visit_temporal_api_sdk_v1_UserMetadata(fs, o.user_metadata)
         if o.HasField("schedule_activity"):
-            token = current_command_info.set(
-                CommandInfo(
-                    command_type=temporalio.api.enums.v1.command_type_pb2.CommandType.COMMAND_TYPE_SCHEDULE_ACTIVITY_TASK,
-                    command_seq=o.schedule_activity.seq,
-                )
+            await self._visit_coresdk_workflow_commands_ScheduleActivity(
+                fs, o.schedule_activity
             )
-            try:
-                await self._visit_coresdk_workflow_commands_ScheduleActivity(
-                    fs, o.schedule_activity
-                )
-            finally:
-                current_command_info.reset(token)
         elif o.HasField("respond_to_query"):
             await self._visit_coresdk_workflow_commands_QueryResult(
                 fs, o.respond_to_query
@@ -481,44 +391,17 @@ async def _visit_coresdk_workflow_commands_WorkflowCommand(self, fs, o):
                 fs, o.continue_as_new_workflow_execution
             )
         elif o.HasField("start_child_workflow_execution"):
-            token = current_command_info.set(
-                CommandInfo(
-                    command_type=temporalio.api.enums.v1.command_type_pb2.CommandType.COMMAND_TYPE_START_CHILD_WORKFLOW_EXECUTION,
-                    command_seq=o.start_child_workflow_execution.seq,
-                )
+            await self._visit_coresdk_workflow_commands_StartChildWorkflowExecution(
+                fs, o.start_child_workflow_execution
             )
-            try:
-                await self._visit_coresdk_workflow_commands_StartChildWorkflowExecution(
-                    fs, o.start_child_workflow_execution
-                )
-            finally:
-                current_command_info.reset(token)
         elif o.HasField("signal_external_workflow_execution"):
-            token = current_command_info.set(
-                CommandInfo(
-                    command_type=temporalio.api.enums.v1.command_type_pb2.CommandType.COMMAND_TYPE_SIGNAL_EXTERNAL_WORKFLOW_EXECUTION,
-                    command_seq=o.signal_external_workflow_execution.seq,
-                )
+            await self._visit_coresdk_workflow_commands_SignalExternalWorkflowExecution(
+                fs, o.signal_external_workflow_execution
             )
-            try:
-                await self._visit_coresdk_workflow_commands_SignalExternalWorkflowExecution(
-                    fs, o.signal_external_workflow_execution
-                )
-            finally:
-                current_command_info.reset(token)
         elif o.HasField("schedule_local_activity"):
-            token = current_command_info.set(
-                CommandInfo(
-                    command_type=temporalio.api.enums.v1.command_type_pb2.CommandType.COMMAND_TYPE_SCHEDULE_ACTIVITY_TASK,
-                    command_seq=o.schedule_local_activity.seq,
-                )
+            await self._visit_coresdk_workflow_commands_ScheduleLocalActivity(
+                fs, o.schedule_local_activity
             )
-            try:
-                await self._visit_coresdk_workflow_commands_ScheduleLocalActivity(
-                    fs, o.schedule_local_activity
-                )
-            finally:
-                current_command_info.reset(token)
         elif o.HasField("upsert_workflow_search_attributes"):
             await self._visit_coresdk_workflow_commands_UpsertWorkflowSearchAttributes(
                 fs, o.upsert_workflow_search_attributes
@@ -532,18 +415,9 @@ async def _visit_coresdk_workflow_commands_WorkflowCommand(self, fs, o):
                 fs, o.update_response
             )
         elif o.HasField("schedule_nexus_operation"):
-            token = current_command_info.set(
-                CommandInfo(
-                    command_type=temporalio.api.enums.v1.command_type_pb2.CommandType.COMMAND_TYPE_SCHEDULE_NEXUS_OPERATION,
-                    command_seq=o.schedule_nexus_operation.seq,
-                )
+            await self._visit_coresdk_workflow_commands_ScheduleNexusOperation(
+                fs, o.schedule_nexus_operation
             )
-            try:
-                await self._visit_coresdk_workflow_commands_ScheduleNexusOperation(
-                    fs, o.schedule_nexus_operation
-                )
-            finally:
-                current_command_info.reset(token)
 
     async def _visit_coresdk_workflow_completion_Success(self, fs, o):
         for v in o.commands: