Serialization context

.github/workflows/ci.yml

@@ -148,6 +148,8 @@ jobs:
         with:
           submodules: recursive
       - uses: dtolnay/rust-toolchain@stable
+        with:
+          components: "clippy"
       - uses: Swatinem/rust-cache@v2
         with:
           workspaces: temporalio/bridge -> target

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
 
@@ -89,6 +89,7 @@ def generate(self, roots: list[Descriptor]) -> str:
 
 from temporalio.api.common.v1.message_pb2 import Payload
 
+
 class VisitorFunctions(abc.ABC):
     \"\"\"Set of functions which can be called by the visitor. 
     Allows handling payloads as a sequence.

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,12 +32,13 @@
 import temporalio.bridge.temporal_sdk_bridge
 import temporalio.converter
 import temporalio.exceptions
-from temporalio.api.common.v1.message_pb2 import Payload, Payloads
-from temporalio.bridge._visitor import PayloadVisitor, VisitorFunctions
+from temporalio.api.common.v1.message_pb2 import Payload
+from temporalio.bridge._visitor import VisitorFunctions
 from temporalio.bridge.temporal_sdk_bridge import (
     CustomSlotSupplier as BridgeCustomSlotSupplier,
 )
 from temporalio.bridge.temporal_sdk_bridge import PollShutdownError  # type: ignore
+from temporalio.worker._command_aware_visitor import CommandAwarePayloadVisitor
 
 
 @dataclass
@@ -299,22 +297,22 @@ async def visit_payloads(self, payloads: MutableSequence[Payload]) -> None:
 
 
 async def decode_activation(
-    act: temporalio.bridge.proto.workflow_activation.WorkflowActivation,
+    activation: temporalio.bridge.proto.workflow_activation.WorkflowActivation,
     codec: temporalio.converter.PayloadCodec,
     decode_headers: bool,
 ) -> None:
-    """Decode the given activation with the codec."""
-    await PayloadVisitor(
+    """Decode all payloads in the activation."""
+    await CommandAwarePayloadVisitor(
         skip_search_attributes=True, skip_headers=not decode_headers
-    ).visit(_Visitor(codec.decode), act)
+    ).visit(_Visitor(codec.decode), activation)
 
 
 async def encode_completion(
-    comp: temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion,
+    completion: temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion,
     codec: temporalio.converter.PayloadCodec,
     encode_headers: bool,
 ) -> None:
-    """Recursively encode the given completion with the codec."""
-    await PayloadVisitor(
+    """Encode all payloads in the completion."""
+    await CommandAwarePayloadVisitor(
         skip_search_attributes=True, skip_headers=not encode_headers
-    ).visit(_Visitor(codec.encode), comp)
+    ).visit(_Visitor(codec.encode), completion)