Add command seq IDs to async context during payload visitor traversal

Author: dandavison

scripts/gen_payload_visitor.py

@@ -1,7 +1,7 @@
 import subprocess
 import sys
 from pathlib import Path
-from typing import Optional, Tuple
+from typing import Optional
 
 from google.protobuf.descriptor import Descriptor, FieldDescriptor
 
@@ -62,6 +62,20 @@ 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.
+    return f"""\
+        {presence_word} o.HasField("{field_name}"):
+            token = current_command_seq.set({access_expr}.seq)
+            try:
+                await self._visit_{child_method}(fs, {access_expr})
+            finally:
+                current_command_seq.reset(token)"""
+
+
 class VisitorGenerator:
     def generate(self, roots: list[Descriptor]) -> str:
         """
@@ -80,10 +94,16 @@ 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 typing import Any, MutableSequence
 
 from temporalio.api.common.v1.message_pb2 import Payload
 
+# Context variable for tracking the current workflow command sequence number
+current_command_seq: contextvars.ContextVar[int | None] = contextvars.ContextVar(
+    "current_command_seq", default=None
+)
+
 class VisitorFunctions(abc.ABC):
     \"\"\"Set of functions which can be called by the visitor. 
     Allows handling payloads as a sequence.
@@ -253,6 +273,20 @@ def walk(self, desc: Descriptor) -> bool:
                         )
                     )
 
+        commands_with_seq = {
+            "start_timer",
+            "cancel_timer",
+            "schedule_activity",
+            "schedule_local_activity",
+            "request_cancel_activity",
+            "request_cancel_local_activity",
+            "start_child_workflow_execution",
+            "request_cancel_external_workflow_execution",
+            "signal_external_workflow_execution",
+            "cancel_signal_workflow",
+            "schedule_nexus_operation",
+            "request_cancel_nexus_operation",
+        }
         # Process oneof fields as if/elif chains
         for oneof_idx, fields in oneof_fields.items():
             oneof_lines = []
@@ -264,9 +298,17 @@ def walk(self, desc: Descriptor) -> bool:
                 if child_has_payload:
                     if_word = "if" if first else "elif"
                     first = False
-                    line = emit_singular(
-                        field.name, f"o.{field.name}", name_for(child_desc), if_word
-                    )
+                    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
+                        )
+                    else:
+                        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,9 +1,15 @@
 # This file is generated by gen_payload_visitor.py. Changes should be made there.
 import abc
+import contextvars
 from typing import Any, MutableSequence
 
 from temporalio.api.common.v1.message_pb2 import Payload
 
+# Context variable for tracking the current workflow command sequence number
+current_command_seq: contextvars.ContextVar[int | None] = contextvars.ContextVar(
+    "current_command_seq", default=None
+)
+
 
 class VisitorFunctions(abc.ABC):
     """Set of functions which can be called by the visitor.
@@ -371,9 +377,13 @@ 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"):
-            await self._visit_coresdk_workflow_commands_ScheduleActivity(
-                fs, o.schedule_activity
-            )
+            token = current_command_seq.set(o.schedule_activity.seq)
+            try:
+                await self._visit_coresdk_workflow_commands_ScheduleActivity(
+                    fs, o.schedule_activity
+                )
+            finally:
+                current_command_seq.reset(token)
         elif o.HasField("respond_to_query"):
             await self._visit_coresdk_workflow_commands_QueryResult(
                 fs, o.respond_to_query
@@ -391,17 +401,29 @@ async def _visit_coresdk_workflow_commands_WorkflowCommand(self, fs, o):
                 fs, o.continue_as_new_workflow_execution
             )
         elif o.HasField("start_child_workflow_execution"):
-            await self._visit_coresdk_workflow_commands_StartChildWorkflowExecution(
-                fs, o.start_child_workflow_execution
-            )
+            token = current_command_seq.set(o.start_child_workflow_execution.seq)
+            try:
+                await self._visit_coresdk_workflow_commands_StartChildWorkflowExecution(
+                    fs, o.start_child_workflow_execution
+                )
+            finally:
+                current_command_seq.reset(token)
         elif o.HasField("signal_external_workflow_execution"):
-            await self._visit_coresdk_workflow_commands_SignalExternalWorkflowExecution(
-                fs, o.signal_external_workflow_execution
-            )
+            token = current_command_seq.set(o.signal_external_workflow_execution.seq)
+            try:
+                await self._visit_coresdk_workflow_commands_SignalExternalWorkflowExecution(
+                    fs, o.signal_external_workflow_execution
+                )
+            finally:
+                current_command_seq.reset(token)
         elif o.HasField("schedule_local_activity"):
-            await self._visit_coresdk_workflow_commands_ScheduleLocalActivity(
-                fs, o.schedule_local_activity
-            )
+            token = current_command_seq.set(o.schedule_local_activity.seq)
+            try:
+                await self._visit_coresdk_workflow_commands_ScheduleLocalActivity(
+                    fs, o.schedule_local_activity
+                )
+            finally:
+                current_command_seq.reset(token)
         elif o.HasField("upsert_workflow_search_attributes"):
             await self._visit_coresdk_workflow_commands_UpsertWorkflowSearchAttributes(
                 fs, o.upsert_workflow_search_attributes
@@ -415,9 +437,13 @@ async def _visit_coresdk_workflow_commands_WorkflowCommand(self, fs, o):
                 fs, o.update_response
             )
         elif o.HasField("schedule_nexus_operation"):
-            await self._visit_coresdk_workflow_commands_ScheduleNexusOperation(
-                fs, o.schedule_nexus_operation
-            )
+            token = current_command_seq.set(o.schedule_nexus_operation.seq)
+            try:
+                await self._visit_coresdk_workflow_commands_ScheduleNexusOperation(
+                    fs, o.schedule_nexus_operation
+                )
+            finally:
+                current_command_seq.reset(token)
 
     async def _visit_coresdk_workflow_completion_Success(self, fs, o):
         for v in o.commands: