Fix custom serialization gaps from #154 (#161)

* Fix custom serialization gaps from #154

Close several round-tripping gaps left by the type-aware custom
serialization work in #154, without introducing new breaking changes
versus 1.5.0 or any serialization-related security concerns.

Serialize side:
- Prefer a to_json() hook over the built-in dataclass / SimpleNamespace
  handling so a dataclass (or namespace) with a non-serializable field can
  opt in, mirroring the decode side which already prefers from_json().
- Encode dataclasses via a shallow field mapping instead of
  dataclasses.asdict(), so nested to_json() hooks are honored and leaf
  values are not deep-copied.
- Serialize enum.Enum values to their underlying .value so non-int enums
  round-trip (IntEnum already serialized as integers).

Deserialize side:
- Recurse type-directed reconstruction into dict/Mapping values and tuple
  elements, in addition to the existing list / Optional / Union / dataclass
  recursion.
- Optionally pass the active DataConverter to a from_json(cls, value,
  converter) hook so it can rebuild nested typed values the built-in
  recursion does not cover.

Entity state:
- Defer deserialization of an entity's wire state until get_state() is
  called, so the caller's requested type reaches the converter together
  with the raw payload. Track whether the held value is still the raw
  serialized string and pass it back through unchanged on persist to avoid
  double-encoding.
- Replace a redundant serialize/deserialize round-trip in the legacy entity
  event path with converter.coerce().

Module structure / deprecation:
- Merge the internal json_codec module into durabletask.serialization and
  make the codec functions private; the supported surface is the pluggable
  DataConverter.
- Deprecate durabletask.internal.shared.to_json / from_json with a
  DeprecationWarning; they continue to work for backwards compatibility.

Adds a comprehensive JsonDataConverter round-trip test suite plus targeted
tests for each fix, and documents intentional limitations (multi-member
Union, types needing a custom converter such as datetime/Decimal/set).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address PR feedback

* Fix annotation discovery in 3.10

* Add pydantic example, fix reconstructibility concern

* More fixes:
- Rename is_reconstructible to can_reconstruct
- Correct ownership of _can_reconstruct
- Required DataConverter for internal classes

* CHANGELOG summarization

* No more silent fallbacks to JsonDataConverter

* PR feedback

* Final CHANGELOG tuneups

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: andystaples

CHANGELOG.md

@@ -12,54 +12,44 @@ ADDED
 - Added a pluggable `DataConverter` (`durabletask.serialization`) accepted by
   `TaskHubGrpcWorker`, `TaskHubGrpcClient`, and `AsyncTaskHubGrpcClient` via a
   `data_converter` argument. Every payload boundary (inputs, outputs, events,
-  custom status, entity state) routes through it. The default
+  custom status, entity state) routes through it, so one converter controls how
+  Python values become JSON and how they are reconstructed. The default
   `JsonDataConverter` preserves existing behavior, so a custom converter (for
-  example one backed by pydantic) is opt-in. Custom objects can opt in via a
-  `to_json()` hook and a `from_json(value)` classmethod.
-- `OrchestrationContext.call_activity`, `call_sub_orchestrator`, and
-  `call_entity` accept an optional `return_type`, and `wait_for_external_event`
-  accepts an optional `data_type`. When provided, the result/event payload is
-  reconstructed as that type (dataclasses — including nested dataclass,
-  `Optional`, and `list` fields — and `from_json()`-capable types) and the
-  returned task is typed accordingly (e.g. `call_activity(..., return_type=Foo)`
-  yields `CompletableTask[Foo]`). When omitted, the raw deserialized JSON is
-  returned as before.
-- Inbound payloads are reconstructed from function type annotations. When an
-  orchestrator, activity, or entity operation annotates its input parameter (or
-  an activity its return value) with a dataclass or `from_json()`-capable type,
-  the payload is reconstructed as that type. Builtins and unannotated/unknown
-  types are passed through unchanged. An explicit `return_type` takes precedence
-  over a discovered annotation.
-- Added typed accessors to `client.OrchestrationState`: `get_input()`,
-  `get_output()`, and `get_custom_status()` each accept an optional
-  `expected_type` and deserialize the corresponding payload, reconstructing
-  dataclasses and `from_json()`-capable types. The raw `serialized_*` fields are
-  retained.
-- Objects exposing a `to_json()` method are now JSON-serializable when passed as
-  activity/orchestrator inputs or outputs.
-- Added `EntityMetadata.get_typed_state(intended_type=...)`, which deserializes
-  the entity's persisted state and reconstructs dataclasses and
-  `from_json()`-capable types. The existing `get_state()` is unchanged: with no
-  argument it returns the raw serialized JSON payload, and `get_state(some_type)`
-  applies constructor-based coercion (`some_type(raw)`).
-- Entity runtime state retrieval (`EntityContext.get_state(intended_type=...)` /
-  `DurableEntity.get_state(...)`) now also reconstructs dataclasses and
-  `from_json()`-capable types, in addition to the existing constructor-based
-  coercion.
+  example one backed by pydantic) is fully opt-in.
+- Custom objects can participate in serialization by exposing a `to_json()`
+  method and a `from_json(value)` classmethod. Both are honored recursively, so
+  nested custom objects round-trip through their own hooks.
+- Payloads are reconstructed into a caller-supplied type — dataclasses
+  (including nested fields), `from_json()`-capable types, and `enum.Enum`
+  members, recursing through `list`, `dict`, `tuple`, and `Optional`/`Union`
+  hints. The type comes from a function's annotations, from an explicit
+  `return_type` on `call_activity` / `call_sub_orchestrator` / `call_entity`
+  (or `data_type` on `wait_for_external_event`), or from the typed accessors
+  `get_input()` / `get_output()` / `get_custom_status()` on
+  `client.OrchestrationState` and `EntityMetadata.get_typed_state(...)`. It is
+  never inferred from the payload. Which annotated types are eligible is decided
+  by the converter via the overridable `DataConverter.can_reconstruct(...)`; a
+  custom converter can override it to recognize its own types (for example
+  `pydantic.BaseModel` subclasses).
 
 CHANGED
 
 - Custom objects (dataclasses, `SimpleNamespace`, namedtuples) are now
   serialized as plain JSON. Decoding such a payload *without* a type hint now
   yields a plain `dict` (previously a `SimpleNamespace`; a namedtuple now
-  round-trips as a JSON array). To get the original type back, pass the new
-  `return_type` / `data_type` arguments, annotate the consuming function's
-  parameter or return type, or use the typed client accessors. Payloads produced
-  by older SDK versions still deserialize — including into a `SimpleNamespace`
-  when no type is supplied — so in-flight orchestrations continue to replay
-  across an upgrade.
+  round-trips as a JSON array). To get the original type back, supply a type via
+  one of the mechanisms above. Payloads produced by older SDK versions still
+  deserialize — including into a `SimpleNamespace` when no type is supplied — so
+  in-flight orchestrations continue to replay 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. This changes v1.6.0 behavior: mutating the returned value in place no
+  longer affects persisted state — write it back with `set_state()`. State is
+  also serialized eagerly at `set_state()` time, so a non-serializable value
+  fails inside the operation (which rolls back) instead of after the batch has
+  run.
 
 FIXED
 
@@ -68,19 +58,31 @@ FIXED
   "no state" and written as `None`, effectively deleting it; only an actual
   `None` state now clears the persisted entity state.
 
-BREAKING CHANGES (type-level only — no runtime impact for typical users)
+DEPRECATED
+
+- `durabletask.internal.shared.to_json` and `durabletask.internal.shared.from_json`
+  are deprecated and now emit a `DeprecationWarning`. Use a
+  `durabletask.serialization.DataConverter` (for example the default
+  `JsonDataConverter`) instead. The functions continue to work for backwards
+  compatibility.
+
+BREAKING CHANGES (no runtime impact for typical users)
 
-These changes do not alter runtime behavior, but because the package ships
-`py.typed`, consumers running strict type checkers (pyright/mypy) — or
-subclassing the public abstract types — may need to update their code:
+Most of these are type-level only: because the package ships `py.typed`,
+consumers running strict type checkers (pyright/mypy) — or subclassing the
+public abstract types — may need to update their code. The constructor change
+below also affects callers who *directly* construct the named classes, which is
+uncommon since they are normally handed to you by the SDK.
 
 - `OrchestrationContext.call_activity`, `call_sub_orchestrator`, `call_entity`,
   and `wait_for_external_event` gained new keyword-only parameters
   (`return_type` / `data_type`). Subclasses overriding these methods should add
   the parameter to match the base signature.
-- `client.OrchestrationState` gained a non-public `_data_converter` field
-  (excluded from equality and `repr`). Code constructing `OrchestrationState`
-  positionally should pass it via the new field or rely on its default.
+- `EntityContext` and `EntityMetadata` (and its `from_entity_metadata` /
+  `from_entity_response` factories) now require a `data_converter` argument.
+  These objects are normally constructed by the SDK — you receive an
+  `EntityContext` in an entity function and an `EntityMetadata` from the client —
+  so this only affects code that constructs them directly.
 
 ## v1.6.0
 

durabletask-azuremanaged/CHANGELOG.md

@@ -7,7 +7,12 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## Unreleased
 
-N/A
+ADDED
+
+- `DurableTaskSchedulerWorker`, `DurableTaskSchedulerClient`, and the async
+  client now accept a `data_converter` argument and forward it to the base
+  worker/client, so a custom `durabletask.serialization.DataConverter` (for
+  example a pydantic-backed one) can be used with the Durable Task Scheduler.
 
 ## v1.6.0
 

durabletask-azuremanaged/durabletask/azuremanaged/client.py

@@ -20,6 +20,7 @@
 )
 import durabletask.internal.shared as shared
 from durabletask.payload.store import PayloadStore
+from durabletask.serialization import DataConverter
 
 
 # Client class used for Durable Task Scheduler (DTS)
@@ -35,6 +36,7 @@ def __init__(self, *,
                  resiliency_options: GrpcClientResiliencyOptions | None = None,
                  default_version: str | None = None,
                  payload_store: PayloadStore | None = None,
+                 data_converter: DataConverter | None = None,
                  log_handler: logging.Handler | None = None,
                  log_formatter: logging.Formatter | None = None):
 
@@ -59,7 +61,8 @@ def __init__(self, *,
             channel_options=channel_options,
             resiliency_options=resiliency_options,
             default_version=default_version,
-            payload_store=payload_store)
+            payload_store=payload_store,
+            data_converter=data_converter)
 
 
 # Async client class used for Durable Task Scheduler (DTS)
@@ -113,6 +116,7 @@ def __init__(self, *,
                  resiliency_options: GrpcClientResiliencyOptions | None = None,
                  default_version: str | None = None,
                  payload_store: PayloadStore | None = None,
+                 data_converter: DataConverter | None = None,
                  log_handler: logging.Handler | None = None,
                  log_formatter: logging.Formatter | None = None):
 
@@ -137,4 +141,5 @@ def __init__(self, *,
             channel_options=channel_options,
             resiliency_options=resiliency_options,
             default_version=default_version,
-            payload_store=payload_store)
+            payload_store=payload_store,
+            data_converter=data_converter)

durabletask-azuremanaged/durabletask/azuremanaged/worker.py

@@ -18,6 +18,7 @@
 )
 import durabletask.internal.shared as shared
 from durabletask.payload.store import PayloadStore
+from durabletask.serialization import DataConverter
 from durabletask.worker import ConcurrencyOptions, TaskHubGrpcWorker
 
 
@@ -81,6 +82,7 @@ def __init__(self, *,
                  resiliency_options: GrpcWorkerResiliencyOptions | None = None,
                  concurrency_options: ConcurrencyOptions | None = None,
                  payload_store: PayloadStore | None = None,
+                 data_converter: DataConverter | None = None,
                  log_handler: logging.Handler | None = None,
                  log_formatter: logging.Formatter | None = None):
 
@@ -110,5 +112,6 @@ def __init__(self, *,
             concurrency_options=concurrency_options,
             # DTS natively supports long timers so chunking is unnecessary
             maximum_timer_interval=None,
-            payload_store=payload_store
+            payload_store=payload_store,
+            data_converter=data_converter
         )

durabletask/entities/entity_context.py

@@ -16,14 +16,11 @@
 
 class EntityContext:
     def __init__(self, orchestration_id: str, operation: str, state: StateShim,
-                 entity_id: EntityInstanceId, data_converter: "DataConverter | None" = None):
+                 entity_id: EntityInstanceId, data_converter: "DataConverter"):
         self._orchestration_id = orchestration_id
         self._operation = operation
         self._state = state
         self._entity_id = entity_id
-        if data_converter is None:
-            from durabletask.serialization import JsonDataConverter
-            data_converter = JsonDataConverter()
         self._data_converter = data_converter
 
     @property

durabletask/entities/entity_metadata.py

@@ -36,7 +36,7 @@ def __init__(self,
                  locked_by: str,
                  includes_state: bool,
                  state: Any | None,
-                 data_converter: "DataConverter | None" = None):
+                 data_converter: "DataConverter"):
         """Initializes a new instance of the EntityMetadata class.
 
         Args:
@@ -48,20 +48,17 @@ def __init__(self,
         self._locked_by = locked_by
         self.includes_state = includes_state
         self._state = state
-        if data_converter is None:
-            from durabletask.serialization import JsonDataConverter
-            data_converter = JsonDataConverter()
         self._data_converter = data_converter
 
     @staticmethod
     def from_entity_response(entity_response: pb.GetEntityResponse, includes_state: bool,
-                             data_converter: "DataConverter | None" = None):
+                             data_converter: "DataConverter"):
         return EntityMetadata.from_entity_metadata(
             entity_response.entity, includes_state, data_converter)
 
     @staticmethod
     def from_entity_metadata(entity: pb.EntityMetadata, includes_state: bool,
-                             data_converter: "DataConverter | None" = None):
+                             data_converter: "DataConverter"):
         try:
             entity_id = EntityInstanceId.parse(entity.instanceId)
         except ValueError:

durabletask/internal/entity_state_shim.py

@@ -12,15 +12,49 @@
 
 
 class StateShim:
-    def __init__(self, start_state: Any, data_converter: "DataConverter | None" = None):
-        self._current_state: Any = start_state
-        self._checkpoint_state: Any = start_state
+    """In-memory view of an entity's state during a batch.
+
+    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",
+                 *, is_serialized: bool = False):
+        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
-        if data_converter is None:
-            from durabletask.serialization import JsonDataConverter
-            data_converter = JsonDataConverter()
-        self._data_converter = data_converter
+
+    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:
@@ -35,31 +69,43 @@ def get_state(self, intended_type: None = None, default: Any = None) -> Any:
         ...
 
     def get_state(self, intended_type: type[TState] | None = None, default: TState | None = None) -> TState | Any | None:
-        if self._current_state is None and default is not None:
+        if self._current_state is None:
             return default
 
+        # Deferred deserialization: the converter receives the raw payload
+        # together with the requested type.
         if intended_type is None:
-            return self._current_state
-
-        coerced = self._data_converter.coerce(self._current_state, intended_type)
+            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
         # value on a failed coercion; restore the stricter contract here by
         # raising when a non-None state could not be coerced to a concrete type.
         # ``intended_type`` may be a typing generic (e.g. ``list[int]``) at
         # runtime, which is not a ``type`` instance, so the guard is required.
-        if (self._current_state is not None
-                and isinstance(intended_type, type)  # pyright: ignore[reportUnnecessaryIsInstance]
-                and not isinstance(coerced, intended_type)):
+        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 coerced
+        return result
 
     def set_state(self, state: Any) -> None:
-        self._current_state = state
+        # 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:
+        """Return the serialized current state for persistence back to the wire.
+
+        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.
+        """
+        return self._current_state
 
     def add_operation_action(self, action: pb.OperationAction) -> None:
         self._operation_actions.append(action)