Fix sandboxed datetime decoding with replay compatibility

changelog/fixed/tiptoeing-clockwork.md

@@ -0,0 +1 @@
+Fix datetime decoding failures in sandboxed workflows, including datetime values nested in collections or dataclasses and datetime-typed signals being dropped. Existing workflow histories retain their original decoding behavior during replay.

temporalio/_workflow_context.py

@@ -0,0 +1,59 @@
+"""Private workflow compatibility flags and their runtime context."""
+
+from __future__ import annotations
+
+from collections.abc import Generator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass
+from enum import IntEnum
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from temporalio.worker._workflow_instance import _WorkflowInstanceImpl
+
+
+class _WorkflowLogicFlag(IntEnum):
+    """Flags recorded in workflow task metadata to preserve replay behavior."""
+
+    RAISE_ON_CANCELLING_COMPLETED_ACTIVITY = 1
+    PROCESS_WORKFLOW_ACTIVATION_JOBS_AS_SINGLE_BATCH = 2
+    DECODE_DATETIME_PROXIES = 3
+
+
+_DEFAULT_ENABLED_WORKFLOW_LOGIC_FLAGS: frozenset[_WorkflowLogicFlag] = frozenset(
+    {
+        _WorkflowLogicFlag.PROCESS_WORKFLOW_ACTIVATION_JOBS_AS_SINGLE_BATCH,
+        _WorkflowLogicFlag.DECODE_DATETIME_PROXIES,
+    }
+)
+
+
+@dataclass(frozen=True)
+class _WorkflowFlagContext:
+    _runtime: _WorkflowInstanceImpl
+
+    def is_flag_enabled(self, flag: _WorkflowLogicFlag) -> bool:
+        """Consult live runtime state, not the state at context creation."""
+        return self._runtime._workflow_logic_flag_enabled(flag)
+
+
+_current_workflow_flags: ContextVar[_WorkflowFlagContext | None] = ContextVar(
+    "temporal_workflow_flags", default=None
+)
+
+
+@contextmanager
+def _workflow_flag_scope(  # type:ignore[reportUnusedFunction]
+    *,
+    runtime: _WorkflowInstanceImpl,
+) -> Generator[None]:
+    """Expose live flag decisions without inheriting another workflow's state."""
+    context = _WorkflowFlagContext(
+        _runtime=runtime,
+    )
+    token = _current_workflow_flags.set(context)
+    try:
+        yield
+    finally:
+        _current_workflow_flags.reset(token)

temporalio/converter/_payload_converter.py

@@ -43,6 +43,10 @@
     # Python's datetime.fromisoformat doesn't support certain formats pre-3.11
     from dateutil import parser  # type: ignore
 
+from temporalio._workflow_context import (
+    _current_workflow_flags,
+    _WorkflowLogicFlag,
+)
 from temporalio.converter._serialization_context import (
     SerializationContext,
     WithSerializationContext,
@@ -800,6 +804,21 @@ def from_payload(
             raise RuntimeError("Failed parsing") from err
 
 
+def _is_datetime_hint(hint: type) -> bool:
+    if hint is datetime:
+        return True
+    context = _current_workflow_flags.get()
+    if context is None or not context.is_flag_enabled(
+        _WorkflowLogicFlag.DECODE_DATETIME_PROXIES
+    ):
+        return False
+
+    # Importing the worker package during converter initialization creates a cycle.
+    from temporalio.worker.workflow_sandbox._restrictions import RestrictionContext
+
+    return RestrictionContext.unwrap_if_proxied(hint) is datetime
+
+
 def _get_iso_datetime_parser() -> Callable[[str], datetime]:
     """Isolates system version check and returns relevant datetime passer
 
@@ -846,7 +865,7 @@ def value_to_type(
     # Any or primitives
     if hint is Any:
         return value
-    elif hint is datetime:
+    elif _is_datetime_hint(hint):
         if isinstance(value, str):
             try:
                 return _get_iso_datetime_parser()(value)

temporalio/worker/_replayer.py

@@ -20,15 +20,17 @@
 import temporalio.worker
 import temporalio.workflow
 
+from .._workflow_context import (
+    _DEFAULT_ENABLED_WORKFLOW_LOGIC_FLAGS,
+    _WorkflowLogicFlag,
+)
 from ..common import HeaderCodecBehavior
 from ._interceptor import Interceptor
 from ._worker import load_default_build_id
 from ._workflow import _WorkflowWorker
 from ._workflow_instance import (
-    _DEFAULT_ENABLED_WORKFLOW_LOGIC_FLAGS,
     UnsandboxedWorkflowRunner,
     WorkflowRunner,
-    _WorkflowLogicFlag,
 )
 from .workflow_sandbox import SandboxedWorkflowRunner
 

temporalio/worker/_worker.py

@@ -30,6 +30,7 @@
     WorkerDeploymentVersion,
 )
 
+from .._workflow_context import _WorkflowLogicFlag
 from ._activity import SharedStateManager, _ActivityWorker
 from ._interceptor import Interceptor
 from ._nexus import _NexusWorker
@@ -43,7 +44,6 @@
     PatchActivationInput,
     UnsandboxedWorkflowRunner,
     WorkflowRunner,
-    _WorkflowLogicFlag,
 )
 from .workflow_sandbox import SandboxedWorkflowRunner
 

temporalio/worker/_workflow.py

@@ -30,6 +30,10 @@
 from temporalio.converter import StorageDriverStoreContext, StorageDriverWorkflowInfo
 from temporalio.worker.workflow_sandbox._runner import SandboxedWorkflowRunner
 
+from .._workflow_context import (
+    _DEFAULT_ENABLED_WORKFLOW_LOGIC_FLAGS,
+    _WorkflowLogicFlag,
+)
 from . import _command_aware_visitor
 from ._debugger import (
     _install_workflow_breakpoint_hook,
@@ -41,13 +45,11 @@
     WorkflowInterceptorClassInput,
 )
 from ._workflow_instance import (
-    _DEFAULT_ENABLED_WORKFLOW_LOGIC_FLAGS,
     PatchActivationInput,
     WorkflowInstance,
     WorkflowInstanceDetails,
     WorkflowRunner,
     _WorkflowExternFunctions,
-    _WorkflowLogicFlag,
 )
 
 logger = logging.getLogger(__name__)

temporalio/worker/_workflow_instance.py

@@ -29,7 +29,6 @@
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from datetime import timedelta
-from enum import IntEnum
 from typing import (
     Any,
     Generic,
@@ -67,6 +66,11 @@
 )
 from temporalio.service import __version__
 
+from .._workflow_context import (
+    _DEFAULT_ENABLED_WORKFLOW_LOGIC_FLAGS,
+    _workflow_flag_scope,
+    _WorkflowLogicFlag,
+)
 from ..api.failure.v1.message_pb2 import Failure
 from ..workflow._event_groups import (
     _capture_event_group_markers,
@@ -477,6 +481,12 @@ def get_thread_id(self) -> int | None:
 
     def activate(
         self, act: temporalio.bridge.proto.workflow_activation.WorkflowActivation
+    ) -> temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion:
+        with _workflow_flag_scope(runtime=self):
+            return self._activate(act)
+
+    def _activate(
+        self, act: temporalio.bridge.proto.workflow_activation.WorkflowActivation
     ) -> temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion:
         # Reset current completion, time, and whether replaying
         self._current_completion = (
@@ -504,6 +514,8 @@ def activate(
         self._single_batch_activation = self._workflow_logic_flag_enabled(
             _WorkflowLogicFlag.PROCESS_WORKFLOW_ACTIVATION_JOBS_AS_SINGLE_BATCH
         )
+        # Record this before queries can arrive, even when no input needs decoding.
+        self._workflow_logic_flag_enabled(_WorkflowLogicFlag.DECODE_DATETIME_PROXIES)
         activation_err: Exception | None = None
         try:
             # Split into job sets with patches, then signals + updates, then
@@ -4251,15 +4263,3 @@ def _make_unfinished_signal_handler_message(
             [{"name": name, "count": count} for name, count in names.most_common()]
         )
     )
-
-
-class _WorkflowLogicFlag(IntEnum):
-    """Flags that may be set on task/activation completion to differentiate new from old workflow behavior."""
-
-    RAISE_ON_CANCELLING_COMPLETED_ACTIVITY = 1
-    PROCESS_WORKFLOW_ACTIVATION_JOBS_AS_SINGLE_BATCH = 2
-
-
-_DEFAULT_ENABLED_WORKFLOW_LOGIC_FLAGS: frozenset[_WorkflowLogicFlag] = frozenset(
-    {_WorkflowLogicFlag.PROCESS_WORKFLOW_ACTIVATION_JOBS_AS_SINGLE_BATCH}
-)

tests/test_workflow_context.py

@@ -0,0 +1,101 @@
+from datetime import datetime
+from typing import Any
+from unittest.mock import create_autospec
+
+import pytest
+
+import temporalio.converter
+from temporalio._workflow_context import (
+    _current_workflow_flags,
+    _workflow_flag_scope,
+    _WorkflowLogicFlag,
+)
+from temporalio.worker._workflow_instance import _WorkflowInstanceImpl
+from temporalio.worker.workflow_sandbox._importer import Importer
+from temporalio.worker.workflow_sandbox._restrictions import (
+    RestrictionContext,
+    SandboxRestrictions,
+)
+
+
+class FlagRuntime:
+    def __init__(self, enabled: bool) -> None:
+        self.enabled = enabled
+        self.checked: list[_WorkflowLogicFlag] = []
+        instance = create_autospec(_WorkflowInstanceImpl, instance=True)
+        instance._workflow_logic_flag_enabled.side_effect = (
+            self._workflow_logic_flag_enabled
+        )
+        self.instance: _WorkflowInstanceImpl = instance
+
+    def _workflow_logic_flag_enabled(self, flag: _WorkflowLogicFlag) -> bool:
+        self.checked.append(flag)
+        return self.enabled
+
+
+def datetime_hint() -> type:
+    with Importer(SandboxRestrictions.default, RestrictionContext()).applied():
+        from datetime import datetime as sandboxed_datetime
+
+    return sandboxed_datetime
+
+
+def test_workflow_flag_scope_restores_context_after_exception() -> None:
+    outer, inner = FlagRuntime(True), FlagRuntime(False)
+    assert _current_workflow_flags.get() is None
+    with _workflow_flag_scope(runtime=outer.instance):
+        context = _current_workflow_flags.get()
+        assert context is not None
+        assert context.is_flag_enabled(_WorkflowLogicFlag.DECODE_DATETIME_PROXIES)
+        with (
+            pytest.raises(ValueError),
+            _workflow_flag_scope(runtime=inner.instance),
+        ):
+            nested = _current_workflow_flags.get()
+            assert nested is not None
+            assert not nested.is_flag_enabled(
+                _WorkflowLogicFlag.DECODE_DATETIME_PROXIES
+            )
+            raise ValueError("test cleanup")
+        assert _current_workflow_flags.get() is context
+    assert _current_workflow_flags.get() is None
+
+
+def test_datetime_proxy_conversion_policy() -> None:
+    hint = datetime_hint()
+    assert hint is not datetime
+    value = "2021-01-02T00:00:00+00:00"
+    with pytest.raises(TypeError):
+        temporalio.converter.value_to_type(hint, value)
+
+    disabled = FlagRuntime(False)
+    with _workflow_flag_scope(runtime=disabled.instance):
+        with pytest.raises(TypeError):
+            temporalio.converter.value_to_type(hint, value)
+    assert disabled.checked == [_WorkflowLogicFlag.DECODE_DATETIME_PROXIES]
+
+    enabled = FlagRuntime(True)
+    with _workflow_flag_scope(runtime=enabled.instance):
+        assert temporalio.converter.value_to_type(
+            hint, value
+        ) == datetime.fromisoformat(value)
+    assert enabled.checked == [_WorkflowLogicFlag.DECODE_DATETIME_PROXIES]
+
+
+def test_datetime_proxy_custom_converter_precedence() -> None:
+    expected_hint = datetime_hint()
+
+    class CustomConverter(temporalio.converter.JSONTypeConverter):
+        def to_typed_value(self, hint: type, value: Any) -> Any:
+            assert hint is expected_hint
+            return "custom"
+
+    runtime = FlagRuntime(True)
+    with _workflow_flag_scope(runtime=runtime.instance):
+        assert (
+            temporalio.converter.value_to_type(
+                expected_hint, "input", [CustomConverter()]
+            )
+            == "custom"
+        )
+    assert not runtime.checked

tests/worker/test_replayer.py

@@ -10,8 +10,8 @@
 
 import pytest
 
-import temporalio.worker._workflow_instance
 from temporalio import activity, workflow
+from temporalio._workflow_context import _WorkflowLogicFlag
 from temporalio.client import Client, WorkflowFailureError, WorkflowHistory
 from temporalio.exceptions import ApplicationError
 from temporalio.testing import WorkflowEnvironment
@@ -82,7 +82,6 @@ def waiting(self) -> bool:
         return self._waiting
 
 
-_WorkflowLogicFlag = temporalio.worker._workflow_instance._WorkflowLogicFlag
 _SINGLE_BATCH_WORKFLOW_LOGIC_FLAG = (
     _WorkflowLogicFlag.PROCESS_WORKFLOW_ACTIVATION_JOBS_AS_SINGLE_BATCH
 )

tests/worker/test_workflow.py

@@ -45,6 +45,10 @@
 import temporalio.worker._workflow_instance
 import temporalio.workflow
 from temporalio import activity, workflow
+from temporalio._workflow_context import (
+    _DEFAULT_ENABLED_WORKFLOW_LOGIC_FLAGS,
+    _WorkflowLogicFlag,
+)
 from temporalio.api.common.v1 import Payload, Payloads, WorkflowExecution
 from temporalio.api.enums.v1 import EventType
 from temporalio.api.failure.v1 import Failure
@@ -1474,17 +1478,13 @@ async def run(self) -> str:
             return "cancelled"
 
 
-_WorkflowLogicFlag = temporalio.worker._workflow_instance._WorkflowLogicFlag
 _SINGLE_BATCH_WORKFLOW_LOGIC_FLAG = (
     _WorkflowLogicFlag.PROCESS_WORKFLOW_ACTIVATION_JOBS_AS_SINGLE_BATCH
 )
 
 
 def test_single_batch_workflow_activation_jobs_default_enabled() -> None:
-    assert (
-        _SINGLE_BATCH_WORKFLOW_LOGIC_FLAG
-        in temporalio.worker._workflow_instance._DEFAULT_ENABLED_WORKFLOW_LOGIC_FLAGS
-    )
+    assert _SINGLE_BATCH_WORKFLOW_LOGIC_FLAG in _DEFAULT_ENABLED_WORKFLOW_LOGIC_FLAGS
 
 
 @workflow.defn