Command-aware serialization context during payload visitor traversal

- Change visitor API
- Add command seq IDs to async context during payload visitor traversal
- Store data_converter on WorkflowInstanceDetails instead of converter classes
- Command-aware codec for payload visitor

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
-from typing import Any, MutableSequence
+import contextvars
+from typing import Any, MutableSequence, Optional
 
 from temporalio.api.common.v1.message_pb2 import Payload
 
+# Current workflow command sequence number
+current_command_seq: contextvars.ContextVar[Optional[int]] = 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,29 @@ 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",
+        }
+        activation_jobs_with_seq = {
+            "resolve_activity",
+            "resolve_child_workflow_execution",
+            "resolve_child_workflow_execution_start",
+            "resolve_request_cancel_external_workflow",
+            "resolve_signal_external_workflow",
+            "resolve_nexus_operation",
+            "resolve_nexus_operation_start",
+        }
         # Process oneof fields as if/elif chains
         for oneof_idx, fields in oneof_fields.items():
             oneof_lines = []
@@ -264,9 +307,25 @@ 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
+                        )
+                    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
+                        )
                     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
-from typing import Any, MutableSequence
+import contextvars
+from typing import Any, MutableSequence, Optional
 
 from temporalio.api.common.v1.message_pb2 import Payload
 
+# Current workflow command sequence number
+current_command_seq: contextvars.ContextVar[Optional[int]] = contextvars.ContextVar(
+    "current_command_seq", default=None
+)
+
 
 class VisitorFunctions(abc.ABC):
     """Set of functions which can be called by the visitor.
@@ -247,35 +253,69 @@ async def _visit_coresdk_workflow_activation_WorkflowActivationJob(self, fs, o):
                 fs, o.signal_workflow
             )
         elif o.HasField("resolve_activity"):
-            await self._visit_coresdk_workflow_activation_ResolveActivity(
-                fs, o.resolve_activity
-            )
+            token = current_command_seq.set(o.resolve_activity.seq)
+            try:
+                await self._visit_coresdk_workflow_activation_ResolveActivity(
+                    fs, o.resolve_activity
+                )
+            finally:
+                current_command_seq.reset(token)
         elif o.HasField("resolve_child_workflow_execution_start"):
-            await self._visit_coresdk_workflow_activation_ResolveChildWorkflowExecutionStart(
-                fs, o.resolve_child_workflow_execution_start
+            token = current_command_seq.set(
+                o.resolve_child_workflow_execution_start.seq
             )
+            try:
+                await self._visit_coresdk_workflow_activation_ResolveChildWorkflowExecutionStart(
+                    fs, o.resolve_child_workflow_execution_start
+                )
+            finally:
+                current_command_seq.reset(token)
         elif o.HasField("resolve_child_workflow_execution"):
-            await self._visit_coresdk_workflow_activation_ResolveChildWorkflowExecution(
-                fs, o.resolve_child_workflow_execution
-            )
+            token = current_command_seq.set(o.resolve_child_workflow_execution.seq)
+            try:
+                await self._visit_coresdk_workflow_activation_ResolveChildWorkflowExecution(
+                    fs, o.resolve_child_workflow_execution
+                )
+            finally:
+                current_command_seq.reset(token)
         elif o.HasField("resolve_signal_external_workflow"):
-            await self._visit_coresdk_workflow_activation_ResolveSignalExternalWorkflow(
-                fs, o.resolve_signal_external_workflow
-            )
+            token = current_command_seq.set(o.resolve_signal_external_workflow.seq)
+            try:
+                await self._visit_coresdk_workflow_activation_ResolveSignalExternalWorkflow(
+                    fs, o.resolve_signal_external_workflow
+                )
+            finally:
+                current_command_seq.reset(token)
         elif o.HasField("resolve_request_cancel_external_workflow"):
-            await self._visit_coresdk_workflow_activation_ResolveRequestCancelExternalWorkflow(
-                fs, o.resolve_request_cancel_external_workflow
+            token = current_command_seq.set(
+                o.resolve_request_cancel_external_workflow.seq
             )
+            try:
+                await self._visit_coresdk_workflow_activation_ResolveRequestCancelExternalWorkflow(
+                    fs, o.resolve_request_cancel_external_workflow
+                )
+            finally:
+                current_command_seq.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"):
-            await self._visit_coresdk_workflow_activation_ResolveNexusOperationStart(
-                fs, o.resolve_nexus_operation_start
-            )
+            token = current_command_seq.set(o.resolve_nexus_operation_start.seq)
+            try:
+                await (
+                    self._visit_coresdk_workflow_activation_ResolveNexusOperationStart(
+                        fs, o.resolve_nexus_operation_start
+                    )
+                )
+            finally:
+                current_command_seq.reset(token)
         elif o.HasField("resolve_nexus_operation"):
-            await self._visit_coresdk_workflow_activation_ResolveNexusOperation(
-                fs, o.resolve_nexus_operation
-            )
+            token = current_command_seq.set(o.resolve_nexus_operation.seq)
+            try:
+                await self._visit_coresdk_workflow_activation_ResolveNexusOperation(
+                    fs, o.resolve_nexus_operation
+                )
+            finally:
+                current_command_seq.reset(token)
 
     async def _visit_coresdk_workflow_activation_WorkflowActivation(self, fs, o):
         for v in o.jobs:
@@ -371,9 +411,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 +435,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 +471,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:

temporalio/bridge/worker.py

@@ -7,11 +7,9 @@
 
 from dataclasses import dataclass
 from typing import (
-    TYPE_CHECKING,
     Awaitable,
     Callable,
     List,
-    Mapping,
     MutableSequence,
     Optional,
     Sequence,
@@ -20,7 +18,6 @@
     Union,
 )
 
-import google.protobuf.internal.containers
 from typing_extensions import TypeAlias
 
 import temporalio.api.common.v1
@@ -35,7 +32,7 @@
 import temporalio.bridge.temporal_sdk_bridge
 import temporalio.converter
 import temporalio.exceptions
-from temporalio.api.common.v1.message_pb2 import Payload, Payloads
+from temporalio.api.common.v1.message_pb2 import Payload
 from temporalio.bridge._visitor import PayloadVisitor, VisitorFunctions
 from temporalio.bridge.temporal_sdk_bridge import (
     CustomSlotSupplier as BridgeCustomSlotSupplier,
@@ -299,22 +296,22 @@ async def visit_payloads(self, payloads: MutableSequence[Payload]) -> None:
 
 
 async def decode_activation(
-    act: temporalio.bridge.proto.workflow_activation.WorkflowActivation,
-    codec: temporalio.converter.PayloadCodec,
+    activation: temporalio.bridge.proto.workflow_activation.WorkflowActivation,
+    decode: Callable[[Sequence[Payload]], Awaitable[List[Payload]]],
     decode_headers: bool,
 ) -> None:
-    """Decode the given activation with the codec."""
+    """Decode all payloads in the activation."""
     await PayloadVisitor(
         skip_search_attributes=True, skip_headers=not decode_headers
-    ).visit(_Visitor(codec.decode), act)
+    ).visit(_Visitor(decode), activation)
 
 
 async def encode_completion(
-    comp: temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion,
-    codec: temporalio.converter.PayloadCodec,
+    completion: temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion,
+    encode: Callable[[Sequence[Payload]], Awaitable[List[Payload]]],
     encode_headers: bool,
 ) -> None:
-    """Recursively encode the given completion with the codec."""
+    """Encode all payloads in the completion."""
     await PayloadVisitor(
         skip_search_attributes=True, skip_headers=not encode_headers
-    ).visit(_Visitor(codec.encode), comp)
+    ).visit(_Visitor(encode), completion)

temporalio/client.py

@@ -5974,6 +5974,7 @@ async def _populate_start_workflow_execution_request(
         req.workflow_type.name = input.workflow
         req.task_queue.name = input.task_queue
         if input.args:
+            # client encode wf input
             req.input.payloads.extend(await data_converter.encode(input.args))
         if input.execution_timeout is not None:
             req.workflow_execution_timeout.FromTimedelta(input.execution_timeout)