Update scheduled tasks with new serialization fixes

Author: andystaples

durabletask/internal/client_helpers.py

@@ -205,14 +205,16 @@ def build_signal_entity_req(
         entity_instance_id: EntityInstanceId,
         operation_name: str,
         input: Any | None = None,
+        signal_time: datetime | None = None,
         data_converter: DataConverter | None = None) -> pb.SignalEntityRequest:
     """Build a SignalEntityRequest for signaling an entity."""
+    scheduled_time = helpers.new_timestamp(signal_time) if signal_time is not None else None
     return pb.SignalEntityRequest(
         instanceId=str(entity_instance_id),
         name=operation_name,
         input=helpers.get_string_value(_serialize(input, data_converter)),
         requestId=str(uuid.uuid4()),
-        scheduledTime=None,
+        scheduledTime=scheduled_time,
         parentTraceContext=None,
         requestTime=helpers.new_timestamp(datetime.now(timezone.utc))
     )

durabletask/internal/helpers.py

@@ -255,11 +255,13 @@ def new_signal_entity_action(id: int,
                              entity_id: EntityInstanceId,
                              operation: str,
                              encoded_input: str | None,
-                             request_id: str) -> pb.OrchestratorAction:
+                             request_id: str,
+                             signal_time: datetime | None = None) -> pb.OrchestratorAction:
+    scheduled_time = new_timestamp(signal_time) if signal_time is not None else None
     return pb.OrchestratorAction(id=id, sendEntityMessage=pb.SendEntityMessageAction(entityOperationSignaled=pb.EntityOperationSignaledEvent(
         requestId=request_id,
         operation=operation,
-        scheduledTime=None,
+        scheduledTime=scheduled_time,
         input=get_string_value(encoded_input),
         targetInstanceId=get_string_value(str(entity_id)),
     )))

durabletask/scheduled/models.py

@@ -277,20 +277,22 @@ def to_json(self) -> dict[str, Any]:
         }
 
     @classmethod
-    def from_json(cls, data: dict[str, Any], converter: Any) -> "ScheduleState":
-        # The nested configuration is reconstructed through the converter, which
-        # routes it to ``ScheduleConfiguration``'s own ``from_json`` hook (and
-        # honors a custom converter). Only this type's datetime leaves are
-        # rebuilt by hand.
+    def from_json(cls, data: dict[str, Any]) -> "ScheduleState":
+        # The nested configuration is reconstructed by calling its own
+        # ``from_json`` hook directly. ``ScheduleConfiguration`` is an internal
+        # type, so there is no need to route it through a (possibly custom)
+        # converter -- keeping this hook converter-free means it round-trips
+        # under any code path, not only the worker's threaded converter.
         state = cls()
         state.status = ScheduleStatus(data["status"])
         state.execution_token = data["execution_token"]
         state.last_run_at = _from_iso(data.get("last_run_at"))
         state.next_run_at = _from_iso(data.get("next_run_at"))
         state.schedule_created_at = _from_iso(data.get("schedule_created_at"))
         state.schedule_last_modified_at = _from_iso(data.get("schedule_last_modified_at"))
-        state.schedule_configuration = converter.coerce(
-            data.get("schedule_configuration"), ScheduleConfiguration)
+        config_data = data.get("schedule_configuration")
+        state.schedule_configuration = (
+            ScheduleConfiguration.from_json(config_data) if config_data is not None else None)
         return state
 
     def to_description(self) -> ScheduleDescription:

durabletask/scheduled/orchestrator.py

@@ -13,10 +13,10 @@
 class ScheduleOperationRequest:
     """Request describing an operation to execute against a schedule entity.
 
-    A plain dataclass: the serializer round-trips it (and its ``input`` payload)
-    automatically. ``input`` stays an ``Any`` here -- it is reconstructed into the
-    concrete options type at the entity-method boundary from that method's
-    parameter annotation.
+    A plain dataclass: the serializer round-trips it automatically. ``input`` is
+    typed ``Any``, so it is reconstructed as the raw deserialized payload; the
+    concrete options type is rebuilt later, at the entity-method boundary, from
+    that method's parameter annotation.
     """
 
     entity_id: str

durabletask/serialization.py

@@ -360,6 +360,9 @@ def _coerce_to_type(value: Any, expected_type: Any, converter: DataConverter | N
     if expected_type is None or value is None:
         return value
 
+    if expected_type is typing.Any:
+        return value
+
     origin = typing.get_origin(expected_type)
     if origin is not None:
         return _coerce_generic(value, expected_type, origin, converter)
@@ -447,6 +450,11 @@ def _coerce_generic(value: Any, expected_type: Any, origin: Any,
         # If the value already matches a member type, keep it as-is.
         non_none = [a for a in args if a is not type(None)]
         for arg in non_none:
+            # ``Any`` imposes no constraint, so the value already satisfies the
+            # union. (Checked explicitly because ``isinstance(value, Any)`` would
+            # raise -- ``typing.Any`` is a class on 3.11+ but not isinstance-able.)
+            if arg is typing.Any:
+                return value
             if isinstance(arg, type) and isinstance(value, arg):
                 return value
         # ``Optional[T]`` (exactly one non-None member): coerce to that member.

tests/durabletask/scheduled/test_models.py

@@ -7,12 +7,14 @@
 
 import pytest
 
-from durabletask.internal import shared
+from durabletask.serialization import JsonDataConverter
 from durabletask.scheduled.models import (ScheduleConfiguration,
                                           ScheduleCreationOptions,
                                           ScheduleState, ScheduleUpdateOptions)
 from durabletask.scheduled.schedule_status import ScheduleStatus
 
+converter = JsonDataConverter()
+
 
 class TestCreationOptionsValidation:
     def test_requires_schedule_id(self):
@@ -51,8 +53,8 @@ def test_round_trip_through_json(self):
             orchestration_input={"key": "value"}, orchestration_instance_id="inst-1",
             start_at=start, end_at=end, start_immediately_if_late=True)
 
-        encoded = shared.to_json(options)
-        decoded = shared.from_json(encoded, ScheduleCreationOptions)
+        encoded = converter.serialize(options)
+        decoded = converter.deserialize(encoded, ScheduleCreationOptions)
 
         assert decoded.schedule_id == "s1"
         assert decoded.orchestration_name == "orch"
@@ -71,7 +73,7 @@ def test_interval_validation(self):
 
     def test_round_trip_through_json(self):
         options = ScheduleUpdateOptions(orchestration_name="orch2", interval=timedelta(seconds=10))
-        decoded = shared.from_json(shared.to_json(options), ScheduleUpdateOptions)
+        decoded = converter.deserialize(converter.serialize(options), ScheduleUpdateOptions)
         assert decoded.orchestration_name == "orch2"
         assert decoded.interval == timedelta(seconds=10)
         assert decoded.start_at is None
@@ -108,7 +110,7 @@ def test_config_round_trip(self):
             ScheduleCreationOptions(schedule_id="s1", orchestration_name="orch",
                                     interval=timedelta(seconds=5),
                                     start_at=datetime(2026, 1, 1, tzinfo=timezone.utc)))
-        restored = shared.from_json(shared.to_json(config), ScheduleConfiguration)
+        restored = converter.deserialize(converter.serialize(config), ScheduleConfiguration)
         assert restored.schedule_id == "s1"
         assert restored.interval == timedelta(seconds=5)
         assert restored.start_at == datetime(2026, 1, 1, tzinfo=timezone.utc)
@@ -124,7 +126,7 @@ def test_round_trip_and_description(self):
                                     interval=timedelta(seconds=5)))
 
         # The nested ``ScheduleConfiguration`` round-trips automatically.
-        restored = shared.from_json(shared.to_json(state), ScheduleState)
+        restored = converter.deserialize(converter.serialize(state), ScheduleState)
         assert restored.status == ScheduleStatus.ACTIVE
         assert restored.schedule_created_at == datetime(2026, 1, 1, tzinfo=timezone.utc)
         assert restored.schedule_configuration is not None

tests/durabletask/scheduled/test_schedule_entity.py

@@ -5,66 +5,76 @@
 
 import logging
 from datetime import datetime, timedelta, timezone
+from typing import Any
 
 import pytest
 
+import durabletask.internal.orchestrator_service_pb2 as pb
 from durabletask.entities import EntityInstanceId
-from durabletask.internal import shared
 from durabletask.internal.entity_state_shim import StateShim
 from durabletask.scheduled.exceptions import ScheduleInvalidTransitionError
-from durabletask.scheduled.models import (ScheduleCreationOptions,
+from durabletask.scheduled.models import (ScheduleCreationOptions, ScheduleState,
                                           ScheduleUpdateOptions)
 from durabletask.scheduled.schedule_entity import (ENTITY_NAME, Schedule)
 from durabletask.scheduled.schedule_status import ScheduleStatus
+from durabletask.serialization import JsonDataConverter
 from durabletask.worker import _EntityExecutor, _Registry
 
 SCHEDULE_ID = "sched-1"
 
 
 class Harness:
-    """Drives Schedule entity operations against a persistent in-memory state."""
+    """Drives Schedule entity operations against a persistent in-memory state.
+
+    Mirrors the worker's entity-batch lifecycle: the ``StateShim`` holds the
+    serialized state between operations, so each ``run`` deserializes on read and
+    serializes on write through the data converter -- exactly the wire round-trip
+    a real entity experiences across batches.
+    """
 
     def __init__(self):
         registry = _Registry()
         registry.add_entity(Schedule, ENTITY_NAME)
-        self.executor = _EntityExecutor(registry, logging.getLogger("test"))
-        self.state = StateShim(None)
+        self.converter = JsonDataConverter()
+        self.executor = _EntityExecutor(registry, logging.getLogger("test"), self.converter)
+        self.shim = StateShim(None, self.converter)
         self.entity_id = EntityInstanceId(ENTITY_NAME, SCHEDULE_ID)
 
-    def run(self, operation, input=None):
-        before = len(self.state.get_operation_actions())
-        encoded = shared.to_json(input) if input is not None else None
-        result = self.executor.execute("orch-1", self.entity_id, operation, self.state, encoded)
-        self.state.commit()
-        # Mimic the wire round-trip: the worker serializes the entity state at
-        # the end of each batch, and the next batch receives it as deserialized
-        # JSON (a plain dict). This exercises the state ``to_json``/``from_json``
-        # hooks between operations and keeps assertions dict-based.
-        current = self.state._current_state  # pyright: ignore[reportPrivateUsage]
-        if current is not None:
-            self.state._current_state = shared.from_json(shared.to_json(current))  # pyright: ignore[reportPrivateUsage]
-        actions = self.state.get_operation_actions()[before:]
+    def run(self, operation: str, input: Any = None) -> tuple[str | None, list[pb.OperationAction]]:
+        before = len(self.shim.get_operation_actions())
+        encoded = self.converter.serialize(input) if input is not None else None
+        result = self.executor.execute("orch-1", self.entity_id, operation, self.shim, encoded)
+        self.shim.commit()
+        actions = self.shim.get_operation_actions()[before:]
         return result, actions
 
+    def state(self) -> ScheduleState | None:
+        """Reconstruct the typed state object, the way the entity itself would."""
+        return self.shim.get_state(ScheduleState)
+
     @property
-    def state_dict(self):
-        return self.state._current_state  # pyright: ignore[reportPrivateUsage]
+    def current(self) -> ScheduleState:
+        """Like :meth:`state` but asserts the state exists (most operations)."""
+        state = self.state()
+        assert state is not None
+        return state
 
     @property
-    def token(self):
-        return self.state_dict["execution_token"]
+    def token(self) -> str:
+        return self.current.execution_token
 
 
-def _signal_actions(actions):
+def _signal_actions(actions: list[pb.OperationAction]) -> list[pb.OperationAction]:
     return [a for a in actions if a.HasField("sendSignal")]
 
 
-def _start_actions(actions):
+def _start_actions(actions: list[pb.OperationAction]) -> list[pb.OperationAction]:
     return [a for a in actions if a.HasField("startNewOrchestration")]
 
 
-def _creation_options(**kwargs):
-    base = dict(schedule_id=SCHEDULE_ID, orchestration_name="my_orch", interval=timedelta(seconds=30))
+def _creation_options(**kwargs: Any) -> ScheduleCreationOptions:
+    base: dict[str, Any] = dict(
+        schedule_id=SCHEDULE_ID, orchestration_name="my_orch", interval=timedelta(seconds=30))
     base.update(kwargs)
     return ScheduleCreationOptions(**base)
 
@@ -74,8 +84,9 @@ def test_create_activates_and_signals_run(self):
         h = Harness()
         _, actions = h.run("create_schedule", _creation_options())
 
-        assert h.state_dict["status"] == ScheduleStatus.ACTIVE.value
-        assert h.state_dict["schedule_created_at"] is not None
+        state = h.current
+        assert state.status == ScheduleStatus.ACTIVE
+        assert state.schedule_created_at is not None
         signals = _signal_actions(actions)
         assert len(signals) == 1
         assert signals[0].sendSignal.name == "run_schedule"
@@ -88,7 +99,7 @@ def test_create_twice_updates_in_place(self):
         h.run("create_schedule", _creation_options(interval=timedelta(seconds=60)))
         # Re-creation refreshes the execution token.
         assert h.token != first_token
-        assert h.state_dict["status"] == ScheduleStatus.ACTIVE.value
+        assert h.current.status == ScheduleStatus.ACTIVE
 
 
 class TestPauseResume:
@@ -97,11 +108,11 @@ def test_pause_then_resume(self):
         h.run("create_schedule", _creation_options())
 
         h.run("pause_schedule")
-        assert h.state_dict["status"] == ScheduleStatus.PAUSED.value
-        assert h.state_dict["next_run_at"] is None
+        assert h.current.status == ScheduleStatus.PAUSED
+        assert h.current.next_run_at is None
 
         _, actions = h.run("resume_schedule")
-        assert h.state_dict["status"] == ScheduleStatus.ACTIVE.value
+        assert h.current.status == ScheduleStatus.ACTIVE
         assert len(_signal_actions(actions)) == 1
 
     def test_pause_when_not_active_raises(self):
@@ -118,7 +129,9 @@ def test_update_changes_config_and_resignals(self):
         h.run("create_schedule", _creation_options())
         _, actions = h.run("update_schedule",
                            ScheduleUpdateOptions(interval=timedelta(seconds=120)))
-        assert abs(h.state_dict["schedule_configuration"]["interval_seconds"] - 120) < 0.001
+        config = h.current.schedule_configuration
+        assert config is not None
+        assert config.interval == timedelta(seconds=120)
         assert len(_signal_actions(actions)) == 1
 
     def test_update_no_change_does_not_signal(self):
@@ -140,7 +153,7 @@ def test_runs_orchestration_when_due_and_rearms(self):
         starts = _start_actions(actions)
         assert len(starts) == 1
         assert starts[0].startNewOrchestration.name == "my_orch"
-        assert h.state_dict["last_run_at"] is not None
+        assert h.current.last_run_at is not None
 
         # Re-arm signal should carry a future scheduled time.
         signals = _signal_actions(actions)
@@ -180,4 +193,4 @@ def test_delete_clears_state(self):
         h = Harness()
         h.run("create_schedule", _creation_options())
         h.run("delete")
-        assert h.state_dict is None
+        assert h.state() is None

tests/durabletask/test_serialization.py

@@ -12,7 +12,7 @@
 from collections import namedtuple
 from dataclasses import dataclass
 from types import SimpleNamespace
-from typing import List, Optional, Union, get_args
+from typing import Any, List, Optional, Union, get_args
 
 import pytest
 
@@ -422,6 +422,33 @@ def test_coerce_dict_values_recursively():
     assert isinstance(result["home"], Address)
 
 
+def test_coerce_bare_any_returns_value_unchanged():
+    # ``Any`` carries no type info; the parsed value is already the result.
+    value = {"k": "v"}
+    assert coerce_to_type(value, Any) is value
+
+
+def test_coerce_optional_any_returns_value_unchanged():
+    # ``Any | None`` must pass the value through rather than raising on the
+    # ``isinstance(value, Any)`` check inside the union loop.
+    value = {"k": "v"}
+    assert coerce_to_type(value, Optional[Any]) == {"k": "v"}
+
+
+def test_coerce_dataclass_with_any_field_round_trips():
+    # A plain dataclass whose field is annotated ``Any | None`` must reconstruct
+    # with that field left as the raw parsed JSON, not crash.
+    @dataclass
+    class Envelope:
+        name: str
+        payload: Any | None = None
+
+    restored = from_json(to_json(Envelope("x", {"nested": [1, 2]})), Envelope)
+    assert isinstance(restored, Envelope)
+    assert restored.name == "x"
+    assert restored.payload == {"nested": [1, 2]}
+
+
 # ----- from_json converter hook (PR #154 follow-up) -----