Address PR feedback

Author: andystaples

CHANGELOG.md

@@ -69,6 +69,14 @@ CHANGED
   across an upgrade.
 - JSON serialization failures now raise a `TypeError` that chains the original
   error (`__cause__`) and names the offending type.
+- `EntityContext.get_state()` / `DurableEntity.get_state()` now return a freshly
+  reconstructed value on every call rather than a reference to a single cached
+  object. As a result, mutating a value returned by `get_state()` in place no
+  longer affects the persisted entity state — write the change back with
+  `set_state()` to persist it. The entity's state is also serialized eagerly at
+  `set_state()` time, so a value that cannot be serialized surfaces the error
+  inside the failing operation (which rolls back) instead of after the batch has
+  run.
 
 FIXED
 

durabletask/internal/entity_state_shim.py

@@ -14,35 +14,50 @@
 class StateShim:
     """In-memory view of an entity's state during a batch.
 
-    The state arriving from the wire is held as its raw serialized JSON string
-    and is **not** deserialized in the constructor: deserialization is deferred
-    until :meth:`get_state` is called, so the caller's requested type reaches the
-    data converter together with the original payload (a custom converter can
-    then deserialize the string directly into the target type). Once the state
-    has been read into a Python value or replaced via :meth:`set_state`, it is
-    held as that live object instead.
-
-    Tracking whether the current value is still the raw serialized string also
-    lets :meth:`encode_state` pass an unmodified payload straight back to the
-    wire instead of re-serializing it, which would double-encode the JSON.
+    The state is held internally as its serialized JSON string at all times.
+    The raw payload off the wire is stored verbatim; a live value supplied via
+    :meth:`set_state` (or as a non-serialized constructor argument) is
+    serialized immediately. Keeping a single, always-serialized representation
+    has two consequences worth noting:
+
+    * Deserialization is deferred to :meth:`get_state`, so the caller's
+      requested type reaches the data converter together with the original
+      payload (a custom converter can deserialize the string directly into the
+      target type), and the unmodified wire payload is handed back by
+      :meth:`encode_state` without being re-encoded.
+    * Serialization errors surface inside the failing operation (at
+      :meth:`set_state`) rather than after the batch has run, so a bad write
+      rolls back just that operation.
+
+    Because the held value is always the serialized form, :meth:`get_state`
+    returns a freshly reconstructed object on every call; it does **not** return
+    a reference to a stored live object. Mutating a value read from
+    :meth:`get_state` therefore has no effect on the persisted state unless it
+    is written back with :meth:`set_state`.
     """
 
     def __init__(self, start_state: Any, data_converter: "DataConverter | None" = None,
                  *, is_serialized: bool = False):
-        # ``is_serialized`` marks ``start_state`` as a raw serialized payload
-        # (the value off the wire) whose deserialization should be deferred. A
-        # ``None`` state is never treated as serialized.
-        serialized = is_serialized and start_state is not None
-        self._current_state: Any = start_state
-        self._current_is_serialized: bool = serialized
-        self._checkpoint_state: Any = start_state
-        self._checkpoint_is_serialized: bool = serialized
-        self._operation_actions: list[pb.OperationAction] = []
-        self._actions_checkpoint_state: int = 0
         if data_converter is None:
             from durabletask.serialization import JsonDataConverter
             data_converter = JsonDataConverter()
         self._data_converter = data_converter
+        # The state is normalized to its serialized string form. ``is_serialized``
+        # marks ``start_state`` as a raw payload already off the wire (stored
+        # verbatim); otherwise a live value is serialized now. ``None`` stays
+        # ``None`` (no persisted state).
+        serialized_start = self._serialize(start_state, is_serialized)
+        self._current_state: str | None = serialized_start
+        self._checkpoint_state: str | None = serialized_start
+        self._operation_actions: list[pb.OperationAction] = []
+        self._actions_checkpoint_state: int = 0
+
+    def _serialize(self, state: Any, is_serialized: bool = False) -> str | None:
+        if state is None:
+            return None
+        if is_serialized:
+            return state
+        return self._data_converter.serialize(state)
 
     @overload
     def get_state(self, intended_type: type[TState], default: TState) -> TState:
@@ -60,16 +75,11 @@ def get_state(self, intended_type: type[TState] | None = None, default: TState |
         if self._current_state is None:
             return default
 
-        if self._current_is_serialized:
-            # Deferred deserialization: the converter receives the raw payload
-            # together with the requested type.
-            if intended_type is None:
-                return self._data_converter.deserialize(self._current_state)
-            result = self._data_converter.deserialize(self._current_state, intended_type)
-        else:
-            if intended_type is None:
-                return self._current_state
-            result = self._data_converter.coerce(self._current_state, intended_type)
+        # Deferred deserialization: the converter receives the raw payload
+        # together with the requested type.
+        if intended_type is None:
+            return self._data_converter.deserialize(self._current_state)
+        result = self._data_converter.deserialize(self._current_state, intended_type)
 
         # An explicit ``intended_type`` is a request to receive that type. The
         # default converter is best-effort and would silently return the raw
@@ -80,30 +90,25 @@ def get_state(self, intended_type: type[TState] | None = None, default: TState |
         if (isinstance(intended_type, type)  # pyright: ignore[reportUnnecessaryIsInstance]
                 and not isinstance(result, intended_type)):
             raise TypeError(
-                f"Could not convert state of type '{type(self._current_state).__name__}' to '{intended_type.__name__}'"
+                f"Could not convert state of type '{type(result).__name__}' to '{intended_type.__name__}'"
             )
 
         return result
 
     def set_state(self, state: Any) -> None:
-        # A value set in-process is a live Python object, not a serialized payload.
-        self._current_state = state
-        self._current_is_serialized = False
+        # Serialize eagerly so the held value is always the wire form and any
+        # serialization error surfaces here, inside the failing operation.
+        self._current_state = self._serialize(state)
 
     def encode_state(self) -> str | None:
-        """Serialize the current state for persistence back to the wire.
+        """Return the serialized current state for persistence back to the wire.
 
-        Returns ``None`` only when the state is actually ``None`` (which clears
-        the persisted entity state). When the current value is still the raw
-        serialized payload (the state was never modified), it is returned
-        unchanged to avoid double-encoding; otherwise the live value is
-        serialized.
+        The state is already held in serialized form, so this is the stored
+        value verbatim: ``None`` when there is no state (which clears the
+        persisted entity state), otherwise the JSON string. No re-encoding
+        occurs, so a payload that was never modified round-trips unchanged.
         """
-        if self._current_state is None:
-            return None
-        if self._current_is_serialized:
-            return self._current_state
-        return self._data_converter.serialize(self._current_state)
+        return self._current_state
 
     def add_operation_action(self, action: pb.OperationAction) -> None:
         self._operation_actions.append(action)
@@ -113,18 +118,14 @@ def get_operation_actions(self) -> list[pb.OperationAction]:
 
     def commit(self) -> None:
         self._checkpoint_state = self._current_state
-        self._checkpoint_is_serialized = self._current_is_serialized
         self._actions_checkpoint_state = len(self._operation_actions)
 
     def rollback(self) -> None:
         self._current_state = self._checkpoint_state
-        self._current_is_serialized = self._checkpoint_is_serialized
         self._operation_actions = self._operation_actions[:self._actions_checkpoint_state]
 
     def reset(self) -> None:
         self._current_state = None
-        self._current_is_serialized = False
         self._checkpoint_state = None
-        self._checkpoint_is_serialized = False
         self._operation_actions = []
         self._actions_checkpoint_state = 0

durabletask/serialization.py

@@ -417,10 +417,13 @@ def _coerce_generic(value: Any, expected_type: Any, origin: Any,
         if len(args) == 2 and args[1] is Ellipsis:
             return tuple(_coerce_to_type(item, args[0], converter) for item in cast(list[Any], value))
         # Fixed-length ``tuple[T1, T2, ...]`` -- coerce element-wise by position.
-        return tuple(
-            _coerce_to_type(item, t, converter)
-            for item, t in zip(cast(list[Any], value), args)
-        )
+        arr = cast(list[Any], value)
+        if len(arr) != len(args):
+            raise TypeError(
+                f"Could not coerce JSON array of length {len(arr)} to "
+                f"tuple of length {len(args)}"
+            )
+        return tuple(_coerce_to_type(item, t, converter) for item, t in zip(arr, args))
     # Other generics are returned as parsed JSON.
     return value
 

tests/durabletask/test_entity_executor.py

@@ -269,14 +269,14 @@ def test_encode_state_none_when_state_is_none(self):
         state = StateShim(None, is_serialized=True)
         assert state.encode_state() is None
 
-    def test_commit_preserves_serialized_flag(self):
+    def test_commit_preserves_unmodified_payload(self):
         state = StateShim('{"value": 7}', is_serialized=True)
         state.commit()
         # After commit, the (unmodified) state still round-trips without
         # double-encoding.
         assert state.encode_state() == '{"value": 7}'
 
-    def test_rollback_restores_serialized_flag(self):
+    def test_rollback_restores_unmodified_payload(self):
         state = StateShim('{"value": 7}', is_serialized=True)
         state.commit()
         state.set_state({"value": 99})

tests/durabletask/test_serialization.py

@@ -364,6 +364,20 @@ def test_from_json_coerces_homogeneous_tuple():
     assert result[1] == Address("c", "d")
 
 
+def test_coerce_to_type_fixed_length_tuple_too_long_raises():
+    # A JSON array longer than the fixed-length tuple type must fail fast rather
+    # than silently dropping the trailing element(s).
+    with pytest.raises(TypeError):
+        coerce_to_type([1, 2, 3], tuple[int, int])
+
+
+def test_coerce_to_type_fixed_length_tuple_too_short_raises():
+    # A JSON array shorter than the fixed-length tuple type must fail fast rather
+    # than silently producing a short tuple.
+    with pytest.raises(TypeError):
+        coerce_to_type([1], tuple[int, int])
+
+
 # ----- coerce_to_type -----