From 871887ac0ea77663525801535444c28a84463ced Mon Sep 17 00:00:00 2001 From: Manohar Paturi <186662190+ManoharPaturi@users.noreply.github.com> Date: Fri, 18 Sep 2026 17:51:46 +0530 Subject: [PATCH] Python: preserve dict subclasses in workflow checkpoints encode_checkpoint_value used isinstance checks for dict and list, so dict subclasses (defaultdict, Counter, OrderedDict) nested in checkpoint state were silently flattened to plain dicts on save. The save-time validation passed because decoding still succeeded, but restored workflows got a plain dict: a defaultdict access pattern like state["missing"].append(x) then raises KeyError on resume, and Counter/OrderedDict types are lost. Only plain dict/list values now take the JSON path; subclasses are pickled like tuples and sets already were, preserving type and behavior through the round trip. collections.Counter joins OrderedDict, defaultdict, and deque in the built-in restricted-unpickler allowlist so the common stdlib dict subclasses restore under restricted decoding too. --- .../_workflows/_checkpoint_encoding.py | 14 ++-- .../tests/workflow/test_checkpoint_encode.py | 66 +++++++++++++++++++ 2 files changed, 76 insertions(+), 4 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py index 1d46253b396..71dcefacb55 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py @@ -145,6 +145,7 @@ def register_checkpoint_type(cls: type[Any]) -> None: "collections:OrderedDict", "collections:defaultdict", "collections:deque", + "collections:Counter", }) _GETATTR_GLOBAL_KEYS: frozenset[str] = frozenset({ @@ -310,8 +311,11 @@ def _encode(value: Any) -> Any: if isinstance(value, _JSON_NATIVE_TYPES): return value - # Recursively encode dict values (keys become strings) - if isinstance(value, dict): + # Recursively encode dict values (keys become strings). Only plain dicts + # take the JSON path: subclasses such as ``defaultdict``, ``Counter``, and + # ``OrderedDict`` carry behavior/type that a plain JSON object cannot + # represent, so they are pickled to preserve object fidelity. + if type(value) is dict: typed_dict = cast(dict[Any, Any], value) # Stringify each key once so reserved-key checks, collision detection, and # the encoded mapping all observe the same strings (stateful ``__str__``). @@ -325,8 +329,10 @@ def _encode(value: Any) -> Any: encoded_dict: dict[str, Any] = {key: _encode(v) for key, v in stringified_items} return encoded_dict - # Recursively encode list items (lists are JSON-native collections) - if isinstance(value, list): + # Recursively encode list items (lists are JSON-native collections). + # As with dicts, only plain lists take the JSON path so list subclasses + # keep their type through a round trip. + if type(value) is list: return [_encode(item) for item in value] # type: ignore # Everything else (tuples, sets, dataclasses, custom objects, etc.): pickle and base64 encode diff --git a/python/packages/core/tests/workflow/test_checkpoint_encode.py b/python/packages/core/tests/workflow/test_checkpoint_encode.py index 35eedaa66d0..3b2a9149b9e 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_encode.py +++ b/python/packages/core/tests/workflow/test_checkpoint_encode.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. import json +from collections import Counter, OrderedDict, defaultdict from dataclasses import dataclass from datetime import datetime, timezone from typing import Any, cast @@ -175,6 +176,71 @@ def test_encode_set() -> None: assert _TYPE_MARKER in result +def test_encode_dict_subclasses_are_pickled() -> None: + """Test that dict subclasses are pickled instead of flattened to plain dicts. + + defaultdict, Counter, and OrderedDict carry behavior that a plain JSON + object cannot represent, so they must take the pickle path like tuples + and sets do. + """ + cases: list[Any] = [ + defaultdict(list, {"todos": ["a"]}), + Counter({"x": 2, "y": 1}), + OrderedDict([("b", 2), ("a", 1)]), + ] + for value in cases: + result = encode_checkpoint_value(value) + assert isinstance(result, dict), type(value) + assert _PICKLE_MARKER in result + assert _TYPE_MARKER in result + assert result[_TYPE_MARKER] == f"collections:{type(value).__name__}" + + +def test_encode_nested_dict_subclass_is_pickled() -> None: + """Test that dict subclasses nested in containers are also pickled.""" + inner: defaultdict[str, list[int]] = defaultdict(list) + inner["k"].append(1) + result = encode_checkpoint_value({"state": [inner]}) + + assert isinstance(result, dict) + nested = result["state"][0] + assert isinstance(nested, dict) + assert _PICKLE_MARKER in nested + assert nested[_TYPE_MARKER] == "collections:defaultdict" + + +def test_round_trip_dict_subclasses_preserve_type_and_behavior() -> None: + """Test that dict subclasses survive a JSON round trip with type and behavior intact.""" + original: defaultdict[str, list[str]] = defaultdict(list) + original["todos"].append("a") + + encoded = json.loads(json.dumps(encode_checkpoint_value(original))) + restored = decode_checkpoint_value(encoded, allowed_types=frozenset()) + + assert type(restored) is defaultdict + assert restored == original + # The default factory must survive so state access patterns keep working on resume. + restored["new_key"].append("x") + assert restored["new_key"] == ["x"] + + +def test_round_trip_counter_and_ordered_dict_preserve_type() -> None: + """Test Counter and OrderedDict round trips under the restricted unpickler.""" + counter = Counter({"a": 2}) + restored_counter = decode_checkpoint_value( + json.loads(json.dumps(encode_checkpoint_value(counter))), allowed_types=frozenset() + ) + assert type(restored_counter) is Counter + assert restored_counter == counter + + ordered = OrderedDict([("b", 2), ("a", 1)]) + restored_ordered = decode_checkpoint_value( + json.loads(json.dumps(encode_checkpoint_value(ordered))), allowed_types=frozenset() + ) + assert type(restored_ordered) is OrderedDict + assert restored_ordered == ordered + + def test_encode_nested_dict() -> None: """Test encoding a nested dictionary structure.""" data = {"outer": {"inner": {"value": 42}}}