Skip to content

Commit 8935437

Browse files
committed
Add secure, type-aware custom object serialization
Rewrite the JSON codec in shared.py to emit plain JSON (no internal type marker) and add type-directed deserialization via an optional expected_type. Custom objects round-trip everywhere: - call_activity/call_sub_orchestrator/call_entity gain return_type; wait_for_external_event gains data_type; these also refine the returned task's static type via overloads. - Inbound payloads (orchestrator/activity/entity inputs) and call_activity results are reconstructed from function type annotations (new internal type_discovery module), best-effort and conservative. - Entity get_state and new client OrchestrationState.get_input/get_output/get_custom_status accessors route through the shared codec. - Fix nested-dataclass round-trip bug; chain serialization errors with the original cause. Legacy AUTO_SERIALIZED payloads still deserialize for in-flight replay.
1 parent 5189878 commit 8935437

13 files changed

Lines changed: 1476 additions & 93 deletions

CHANGELOG.md

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,57 @@ All notable changes to this project will be documented in this file.
55
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project
66
adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
77

8+
## Unreleased
9+
10+
### Added
11+
12+
- Type-aware deserialization of payloads. `OrchestrationContext.call_activity`,
13+
`call_sub_orchestrator`, and `call_entity` accept an optional `return_type`,
14+
and `wait_for_external_event` accepts an optional `data_type`. When provided,
15+
the result/event payload is coerced to that type: dataclasses are
16+
reconstructed from their dict payloads (including nested dataclass, `Optional`,
17+
and `list` fields), and types exposing a `from_json()` classmethod are rebuilt
18+
via that hook. When omitted, the raw deserialized JSON is returned as before.
19+
The `return_type` / `data_type` argument also refines the static type of the
20+
returned task (e.g. `call_activity(..., return_type=Foo)` is typed as
21+
`CompletableTask[Foo]`).
22+
- Inbound payloads are reconstructed from function type annotations. When an
23+
orchestrator, activity, or entity operation annotates its input parameter with
24+
a dataclass or a `from_json()`-capable type, the incoming payload is
25+
automatically coerced to that type. Discovery is best-effort and conservative:
26+
builtins and unannotated/unknown types are passed through unchanged, and a
27+
payload that cannot be coerced falls back to the raw value.
28+
- `call_activity` results are reconstructed from the activity's return
29+
annotation. When an activity function reference is passed (not a string name)
30+
and its return type is annotated with a dataclass or `from_json()`-capable
31+
type, the result is automatically coerced to that type. An explicit
32+
`return_type` argument takes precedence over the discovered annotation.
33+
- Added typed accessors to `client.OrchestrationState`: `get_input()`,
34+
`get_output()`, and `get_custom_status()` each accept an optional
35+
`expected_type` and deserialize the corresponding `serialized_*` payload,
36+
reconstructing dataclasses and `from_json()`-capable types. The raw
37+
`serialized_input` / `serialized_output` / `serialized_custom_status` string
38+
fields are retained.
39+
- Objects exposing a `to_json()` method are now JSON-serializable when passed as
40+
activity/orchestrator inputs or outputs.
41+
- Entity state retrieval (`get_state(intended_type=...)`) now reconstructs
42+
dataclasses from their stored dict payloads and supports types exposing a
43+
`from_json()` classmethod, in addition to the existing constructor-based
44+
coercion.
45+
46+
### Changed
47+
48+
- Custom objects (dataclasses, `SimpleNamespace`) are now serialized as plain
49+
JSON without an internal type marker. Decoding without a `return_type` /
50+
`data_type` therefore yields a plain `dict` (previously a `SimpleNamespace`
51+
for marked payloads). Pass the new type arguments to reconstruct the original
52+
type. Payloads produced by older SDK versions (carrying the legacy marker)
53+
continue to deserialize, including into a `SimpleNamespace` when no type is
54+
supplied.
55+
- JSON serialization failures now raise a `TypeError` that chains the original
56+
error (`__cause__`) and names the offending type, making serialization issues
57+
easier to diagnose.
58+
859
## v1.5.0
960

1061
BREAKING CHANGES (type-level only — no runtime impact for typical users)

durabletask/client.py

Lines changed: 92 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
from dataclasses import dataclass
1111
from datetime import datetime
1212
from enum import Enum
13-
from typing import Any, Generic, Protocol, TypeVar, cast
13+
from typing import Any, Generic, Protocol, TypeVar, cast, overload
1414

1515
import grpc
1616
import grpc.aio
@@ -54,6 +54,7 @@
5454
TInput = TypeVar('TInput')
5555
TOutput = TypeVar('TOutput')
5656
TItem = TypeVar('TItem')
57+
T = TypeVar('T')
5758

5859

5960
class OrchestrationStatus(Enum):
@@ -82,6 +83,96 @@ class OrchestrationState:
8283
serialized_custom_status: str | None
8384
failure_details: task.FailureDetails | None
8485

86+
@overload
87+
def get_input(self, expected_type: type[T]) -> T | None:
88+
...
89+
90+
@overload
91+
def get_input(self, expected_type: None = ...) -> Any:
92+
...
93+
94+
def get_input(self, expected_type: type | None = None) -> Any:
95+
"""Deserialize the orchestration's input.
96+
97+
Parameters
98+
----------
99+
expected_type : type | None
100+
Optional type used to reconstruct the input. When provided, the
101+
payload is coerced to this type (dataclasses are constructed from
102+
their dict payloads, types exposing a ``from_json()`` classmethod
103+
are reconstructed via that hook) and the return value is typed as
104+
``expected_type | None``. When omitted, the raw deserialized JSON is
105+
returned.
106+
107+
Returns
108+
-------
109+
Any
110+
The deserialized input, or None if there is no input.
111+
"""
112+
if self.serialized_input is None:
113+
return None
114+
return shared.from_json(self.serialized_input, expected_type)
115+
116+
@overload
117+
def get_output(self, expected_type: type[T]) -> T | None:
118+
...
119+
120+
@overload
121+
def get_output(self, expected_type: None = ...) -> Any:
122+
...
123+
124+
def get_output(self, expected_type: type | None = None) -> Any:
125+
"""Deserialize the orchestration's output.
126+
127+
Parameters
128+
----------
129+
expected_type : type | None
130+
Optional type used to reconstruct the output. When provided, the
131+
payload is coerced to this type (dataclasses are constructed from
132+
their dict payloads, types exposing a ``from_json()`` classmethod
133+
are reconstructed via that hook) and the return value is typed as
134+
``expected_type | None``. When omitted, the raw deserialized JSON is
135+
returned.
136+
137+
Returns
138+
-------
139+
Any
140+
The deserialized output, or None if there is no output.
141+
"""
142+
if self.serialized_output is None:
143+
return None
144+
return shared.from_json(self.serialized_output, expected_type)
145+
146+
@overload
147+
def get_custom_status(self, expected_type: type[T]) -> T | None:
148+
...
149+
150+
@overload
151+
def get_custom_status(self, expected_type: None = ...) -> Any:
152+
...
153+
154+
def get_custom_status(self, expected_type: type | None = None) -> Any:
155+
"""Deserialize the orchestration's custom status.
156+
157+
Parameters
158+
----------
159+
expected_type : type | None
160+
Optional type used to reconstruct the custom status. When provided,
161+
the payload is coerced to this type (dataclasses are constructed
162+
from their dict payloads, types exposing a ``from_json()``
163+
classmethod are reconstructed via that hook) and the return value is
164+
typed as ``expected_type | None``. When omitted, the raw
165+
deserialized JSON is returned.
166+
167+
Returns
168+
-------
169+
Any
170+
The deserialized custom status, or None if there is no custom status.
171+
"""
172+
if self.serialized_custom_status is None:
173+
return None
174+
return shared.from_json(self.serialized_custom_status, expected_type)
175+
85176
def raise_if_failed(self):
86177
if self.failure_details is not None:
87178
raise OrchestrationFailedError(

durabletask/entities/entity_metadata.py

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from durabletask.entities.entity_instance_id import EntityInstanceId
44

55
import durabletask.internal.orchestrator_service_pb2 as pb
6+
from durabletask.internal import shared
67

78
TState = TypeVar("TState")
89

@@ -77,15 +78,7 @@ def get_state(self, intended_type: type[TState] | None = None) -> TState | Any |
7778
if intended_type is None or self._state is None:
7879
return self._state
7980

80-
if isinstance(self._state, intended_type):
81-
return self._state
82-
83-
try:
84-
return intended_type(self._state) # type: ignore[call-arg]
85-
except Exception as ex:
86-
raise TypeError(
87-
f"Could not convert state of type '{type(self._state).__name__}' to '{intended_type.__name__}'"
88-
) from ex
81+
return shared.coerce_to_type(self._state, intended_type)
8982

9083
def get_locked_by(self) -> EntityInstanceId | None:
9184
"""Get the identifier of the worker that currently holds the lock on the entity.

durabletask/internal/entity_state_shim.py

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from typing import Any, TypeVar, overload
22

33
import durabletask.internal.orchestrator_service_pb2 as pb
4+
from durabletask.internal import shared
45

56
TState = TypeVar("TState")
67

@@ -31,15 +32,7 @@ def get_state(self, intended_type: type[TState] | None = None, default: TState |
3132
if intended_type is None:
3233
return self._current_state
3334

34-
if isinstance(self._current_state, intended_type):
35-
return self._current_state
36-
37-
try:
38-
return intended_type(self._current_state) # type: ignore[call-arg]
39-
except Exception as ex:
40-
raise TypeError(
41-
f"Could not convert state of type '{type(self._current_state).__name__}' to '{intended_type.__name__}'"
42-
) from ex
35+
return shared.coerce_to_type(self._current_state, intended_type)
4336

4437
def set_state(self, state: Any) -> None:
4538
self._current_state = state

0 commit comments

Comments
 (0)