From fe6d66421095de7c013e70931a68d8eaa0f8a90e Mon Sep 17 00:00:00 2001 From: droideronline Date: Thu, 3 Sep 2026 13:04:10 +0530 Subject: [PATCH 1/5] Fix checkpoint pickling of runtime raw representations --- .../core/agent_framework/_serialization.py | 13 +++++++++++++ .../packages/core/agent_framework/_types.py | 13 +++++++++++++ .../tests/workflow/test_agent_executor.py | 19 +++++++++++++++++++ 3 files changed, 45 insertions(+) diff --git a/python/packages/core/agent_framework/_serialization.py b/python/packages/core/agent_framework/_serialization.py index 384fa05a5be..3653a3d1f61 100644 --- a/python/packages/core/agent_framework/_serialization.py +++ b/python/packages/core/agent_framework/_serialization.py @@ -303,6 +303,19 @@ def __deepcopy__(self, memo: dict[int, Any]) -> SerializationMixin: object.__setattr__(result, k, copy.deepcopy(v, memo)) return result + def __getstate__(self) -> dict[str, Any]: + """Return pickle state without runtime-only shallow-copy fields.""" + state = dict(self.__dict__) + for field_name in self._SHALLOW_COPY_FIELDS: + state.pop(field_name, None) + return state + + def __setstate__(self, state: dict[str, Any]) -> None: + """Restore pickle state and reset runtime-only shallow-copy fields.""" + self.__dict__.update(state) + for field_name in self._SHALLOW_COPY_FIELDS: + self.__dict__.setdefault(field_name, None) + def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: """Convert the instance and any nested objects to a dictionary. diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 9134dc6a046..ac6690d9682 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -610,6 +610,19 @@ def __deepcopy__(self, memo: dict[int, Any]) -> Content: object.__setattr__(result, k, deepcopy(v, memo)) return result + def __getstate__(self) -> dict[str, Any]: + """Return pickle state without runtime-only shallow-copy fields.""" + state = dict(self.__dict__) + for field_name in self._SHALLOW_COPY_FIELDS: + state.pop(field_name, None) + return state + + def __setstate__(self, state: dict[str, Any]) -> None: + """Restore pickle state and reset runtime-only shallow-copy fields.""" + self.__dict__.update(state) + for field_name in self._SHALLOW_COPY_FIELDS: + self.__dict__.setdefault(field_name, None) + @classmethod def from_text( cls: type[ContentT], diff --git a/python/packages/core/tests/workflow/test_agent_executor.py b/python/packages/core/tests/workflow/test_agent_executor.py index ccb1e9425bf..2cc2ed2ce6e 100644 --- a/python/packages/core/tests/workflow/test_agent_executor.py +++ b/python/packages/core/tests/workflow/test_agent_executor.py @@ -1,5 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. +import pickle + from collections.abc import AsyncIterable, Awaitable from typing import Any, Literal, overload @@ -336,6 +338,9 @@ class _NonCopyableRaw: def __deepcopy__(self, memo: dict) -> Any: raise TypeError("Cannot deepcopy this object") + def __reduce__(self) -> Any: + raise TypeError("Cannot pickle this object") + class _AgentWithRawRepr(BaseAgent): """Agent that returns responses with a non-copyable raw_representation.""" @@ -387,6 +392,20 @@ async def test_agent_executor_workflow_with_non_copyable_raw_representation() -> assert agent_responses[0].raw_representation is raw +def test_serialization_mixin_omits_non_pickleable_raw_representation() -> None: + """Pickling framework objects should not include runtime-only raw representations.""" + raw = _NonCopyableRaw() + response = AgentResponse( + messages=[Message("assistant", [Content.from_text(text="reply", raw_representation=raw)])], + raw_representation=raw, + ) + + restored = pickle.loads(pickle.dumps(response)) + + assert restored.raw_representation is None + assert restored.messages[0].contents[0].raw_representation is None + + # --------------------------------------------------------------------------- # Context mode tests # --------------------------------------------------------------------------- From e47014d19aa0be4b02231387e62d0a24c9f38c23 Mon Sep 17 00:00:00 2001 From: droideronline Date: Thu, 3 Sep 2026 13:50:28 +0530 Subject: [PATCH 2/5] Fix checkpoint serialization review feedback --- .../core/agent_framework/_serialization.py | 49 ++++++++++++++++--- .../packages/core/agent_framework/_types.py | 16 +++--- .../tests/core/test_serializable_mixin.py | 39 +++++++++++++++ python/packages/core/tests/core/test_types.py | 13 +++++ 4 files changed, 102 insertions(+), 15 deletions(-) diff --git a/python/packages/core/agent_framework/_serialization.py b/python/packages/core/agent_framework/_serialization.py index 3653a3d1f61..bf0adeb0d3d 100644 --- a/python/packages/core/agent_framework/_serialization.py +++ b/python/packages/core/agent_framework/_serialization.py @@ -151,6 +151,35 @@ def _is_serialization_protocol(value: Any) -> TypeGuard[SerializationProtocol]: return callable(getattr(value, "to_dict", None)) and callable(getattr(value, "from_dict", None)) +def _iter_instance_fields(instance: Any) -> dict[str, Any]: + """Return attributes stored in ``__dict__`` or slots.""" + fields = dict(getattr(instance, "__dict__", {})) + for cls in type(instance).__mro__: + slots = cls.__dict__.get("__slots__", ()) + if isinstance(slots, str): + slots = (slots,) + for field_name in slots: + if field_name not in {"__dict__", "__weakref__"} and hasattr(instance, field_name): + fields[field_name] = getattr(instance, field_name) + return fields + + +def _get_pickle_state(instance: Any, omitted_fields: set[str]) -> dict[str, Any]: + """Build pickle state while omitting runtime-only fields.""" + state = _iter_instance_fields(instance) + for field_name in omitted_fields: + state.pop(field_name, None) + return state + + +def _restore_pickle_state(instance: Any, state: dict[str, Any], omitted_fields: set[str]) -> None: + """Restore dict- and slot-backed pickle state.""" + for field_name, value in state.items(): + object.__setattr__(instance, field_name, value) + for field_name in omitted_fields: + object.__setattr__(instance, field_name, None) + + class SerializationMixin: """Mixin class providing comprehensive serialization and deserialization capabilities. @@ -284,6 +313,15 @@ def __init__(self, **kwargs): DEFAULT_EXCLUDE: ClassVar[set[str]] = set() INJECTABLE: ClassVar[set[str]] = set() _SHALLOW_COPY_FIELDS: ClassVar[set[str]] = {"raw_representation"} + _PICKLE_OMIT_FIELDS: ClassVar[set[str]] = {"raw_representation"} + + def __copy__(self) -> SerializationMixin: + """Create a shallow copy without invoking pickle state hooks.""" + cls = type(self) + result = cls.__new__(cls) + for field_name, value in _iter_instance_fields(self).items(): + object.__setattr__(result, field_name, value) + return result def __deepcopy__(self, memo: dict[int, Any]) -> SerializationMixin: """Create a deep copy, preserving ``_SHALLOW_COPY_FIELDS`` by reference. @@ -296,7 +334,7 @@ def __deepcopy__(self, memo: dict[int, Any]) -> SerializationMixin: cls = type(self) result = cls.__new__(cls) memo[id(self)] = result - for k, v in self.__dict__.items(): + for k, v in _iter_instance_fields(self).items(): if k in cls._SHALLOW_COPY_FIELDS: object.__setattr__(result, k, v) else: @@ -305,16 +343,11 @@ def __deepcopy__(self, memo: dict[int, Any]) -> SerializationMixin: def __getstate__(self) -> dict[str, Any]: """Return pickle state without runtime-only shallow-copy fields.""" - state = dict(self.__dict__) - for field_name in self._SHALLOW_COPY_FIELDS: - state.pop(field_name, None) - return state + return _get_pickle_state(self, self._PICKLE_OMIT_FIELDS) def __setstate__(self, state: dict[str, Any]) -> None: """Restore pickle state and reset runtime-only shallow-copy fields.""" - self.__dict__.update(state) - for field_name in self._SHALLOW_COPY_FIELDS: - self.__dict__.setdefault(field_name, None) + _restore_pickle_state(self, state, self._PICKLE_OMIT_FIELDS) def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: """Convert the instance and any nested objects to a dictionary. diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index ac6690d9682..ac94ed8d70a 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -30,7 +30,7 @@ from typing_extensions import TypedDict from ._feature_stage import ExperimentalFeature, experimental -from ._serialization import SerializationMixin +from ._serialization import SerializationMixin, _get_pickle_state, _restore_pickle_state from .exceptions import AdditionItemMismatch, ContentError if sys.version_info >= (3, 13): @@ -483,6 +483,7 @@ class Content: """ _SHALLOW_COPY_FIELDS: ClassVar[set[str]] = {"raw_representation"} + _PICKLE_OMIT_FIELDS: ClassVar[set[str]] = {"raw_representation"} def __init__( self, @@ -612,16 +613,17 @@ def __deepcopy__(self, memo: dict[int, Any]) -> Content: def __getstate__(self) -> dict[str, Any]: """Return pickle state without runtime-only shallow-copy fields.""" - state = dict(self.__dict__) - for field_name in self._SHALLOW_COPY_FIELDS: - state.pop(field_name, None) + state = _get_pickle_state(self, self._PICKLE_OMIT_FIELDS) + if self.annotations is not None: + state["annotations"] = [ + {key: value for key, value in annotation.items() if key != "raw_representation"} + for annotation in self.annotations + ] return state def __setstate__(self, state: dict[str, Any]) -> None: """Restore pickle state and reset runtime-only shallow-copy fields.""" - self.__dict__.update(state) - for field_name in self._SHALLOW_COPY_FIELDS: - self.__dict__.setdefault(field_name, None) + _restore_pickle_state(self, state, self._PICKLE_OMIT_FIELDS) @classmethod def from_text( diff --git a/python/packages/core/tests/core/test_serializable_mixin.py b/python/packages/core/tests/core/test_serializable_mixin.py index 03853e83868..0cc498793b7 100644 --- a/python/packages/core/tests/core/test_serializable_mixin.py +++ b/python/packages/core/tests/core/test_serializable_mixin.py @@ -572,6 +572,45 @@ def __init__(self, items: list, opaque: Any = None, additional_properties: dict assert cloned.items is not obj.items assert cloned.items == ["a"] + def test_shallow_copy_preserves_pickle_omitted_fields(self): + """Shallow copies retain runtime fields that pickle omits.""" + + class TestClass(SerializationMixin): + def __init__(self, raw_representation: Any): + self.raw_representation = raw_representation + + raw = object() + cloned = copy.copy(TestClass(raw)) + + assert cloned.raw_representation is raw + + def test_pickle_restores_slot_fields(self): + """Pickle state should include fields declared in slots.""" + class TestClass(SerializationMixin): + __slots__ = ("value",) + + def __init__(self, value: str): + self.value = value + + original = TestClass("value") + restored = TestClass.__new__(TestClass) + restored.__setstate__(original.__getstate__()) + + assert restored.value == "value" + + def test_pickle_omission_is_separate_from_shallow_copy_policy(self): + """Fields shallow-copied by default remain persistent unless explicitly omitted.""" + class TestClass(SerializationMixin): + _PICKLE_OMIT_FIELDS = set() + + def __init__(self, raw_representation: Any): + self.raw_representation = raw_representation + + raw = {"provider": "value"} + state = TestClass(raw).__getstate__() + + assert state["raw_representation"] == raw + def test_dependency_dict_merge_does_not_mutate_input(self): """Test that dict dependency merging does not mutate the caller's input dictionary.""" diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index a4d5957f587..b96fb0abf5e 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -2733,6 +2733,19 @@ def test_content_deepcopy_discards_raw_representation(caplog: pytest.LogCaptureF assert caplog.messages == ["Discarding field 'raw_representation' while deep-copying Content."] +def test_content_pickle_discards_nested_annotation_raw_representation() -> None: + """Pickle should omit provider objects stored on annotations.""" + import pickle + + raw = object() + annotation: Annotation = {"type": "citation", "url": "https://example.com", "raw_representation": raw} + content = Content.from_text("hello", annotations=[annotation]) + + restored = pickle.loads(pickle.dumps(content)) + + assert restored.annotations == [{"type": "citation", "url": "https://example.com"}] + + def test_message_deepcopy_preserves_raw_representation(): """Test that deepcopy of Message keeps raw_representation by reference.""" import copy From 7b2fc12b56ca33bc6d4b0852e6f34fc7278e70b6 Mon Sep 17 00:00:00 2001 From: droideronline Date: Thu, 3 Sep 2026 16:22:51 +0530 Subject: [PATCH 3/5] Address checkpoint serialization review follow-up --- .../core/agent_framework/_serialization.py | 12 ++++++++++-- python/packages/core/agent_framework/_types.py | 8 ++++++++ .../core/tests/core/test_serializable_mixin.py | 14 ++++++++++++++ python/packages/core/tests/core/test_types.py | 10 ++++++++++ 4 files changed, 42 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_serialization.py b/python/packages/core/agent_framework/_serialization.py index bf0adeb0d3d..3a9ebd0fd6f 100644 --- a/python/packages/core/agent_framework/_serialization.py +++ b/python/packages/core/agent_framework/_serialization.py @@ -172,12 +172,20 @@ def _get_pickle_state(instance: Any, omitted_fields: set[str]) -> dict[str, Any] return state -def _restore_pickle_state(instance: Any, state: dict[str, Any], omitted_fields: set[str]) -> None: +def _restore_pickle_state( + instance: Any, + state: dict[str, Any] | tuple[dict[str, Any], dict[str, Any]], + omitted_fields: set[str], +) -> None: """Restore dict- and slot-backed pickle state.""" + if isinstance(state, tuple): + dict_state, slot_state = state + state = {**dict_state, **slot_state} for field_name, value in state.items(): object.__setattr__(instance, field_name, value) for field_name in omitted_fields: - object.__setattr__(instance, field_name, None) + if field_name in _iter_instance_fields(instance) or hasattr(instance, "__dict__"): + object.__setattr__(instance, field_name, None) class SerializationMixin: diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index ac94ed8d70a..eb3e11fba0d 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -611,6 +611,14 @@ def __deepcopy__(self, memo: dict[int, Any]) -> Content: object.__setattr__(result, k, deepcopy(v, memo)) return result + def __copy__(self) -> Content: + """Create a shallow copy while preserving provider runtime fields.""" + cls = type(self) + result = cls.__new__(cls) + for field_name, value in self.__dict__.items(): + object.__setattr__(result, field_name, value) + return result + def __getstate__(self) -> dict[str, Any]: """Return pickle state without runtime-only shallow-copy fields.""" state = _get_pickle_state(self, self._PICKLE_OMIT_FIELDS) diff --git a/python/packages/core/tests/core/test_serializable_mixin.py b/python/packages/core/tests/core/test_serializable_mixin.py index 0cc498793b7..d5b14cb0015 100644 --- a/python/packages/core/tests/core/test_serializable_mixin.py +++ b/python/packages/core/tests/core/test_serializable_mixin.py @@ -598,6 +598,20 @@ def __init__(self, value: str): assert restored.value == "value" + def test_pickle_restores_legacy_tuple_state(self): + """Pickle restoration should accept the legacy dict-and-slots tuple.""" + + class TestClass(SerializationMixin): + __slots__ = ("value",) + + def __init__(self): + self.value = "new" + + restored = TestClass.__new__(TestClass) + restored.__setstate__(({"other": "dict"}, {"value": "legacy"})) + + assert restored.value == "legacy" + def test_pickle_omission_is_separate_from_shallow_copy_policy(self): """Fields shallow-copied by default remain persistent unless explicitly omitted.""" class TestClass(SerializationMixin): diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index b96fb0abf5e..e936da808af 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -2746,6 +2746,16 @@ def test_content_pickle_discards_nested_annotation_raw_representation() -> None: assert restored.annotations == [{"type": "citation", "url": "https://example.com"}] +def test_content_shallow_copy_preserves_raw_representation() -> None: + """Shallow copies of Content retain provider runtime fields.""" + import copy + + raw = _NonCopyableRaw() + cloned = copy.copy(Content.from_text("hello", raw_representation=raw)) + + assert cloned.raw_representation is raw + + def test_message_deepcopy_preserves_raw_representation(): """Test that deepcopy of Message keeps raw_representation by reference.""" import copy From 1f8c5586724ef3e6dcad837dac5f780b08c629c2 Mon Sep 17 00:00:00 2001 From: droideronline Date: Thu, 3 Sep 2026 16:50:26 +0530 Subject: [PATCH 4/5] Fix package check private imports --- python/packages/core/agent_framework/_serialization.py | 8 ++++---- python/packages/core/agent_framework/_types.py | 6 +++--- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/python/packages/core/agent_framework/_serialization.py b/python/packages/core/agent_framework/_serialization.py index 3a9ebd0fd6f..78ec9ed0506 100644 --- a/python/packages/core/agent_framework/_serialization.py +++ b/python/packages/core/agent_framework/_serialization.py @@ -164,7 +164,7 @@ def _iter_instance_fields(instance: Any) -> dict[str, Any]: return fields -def _get_pickle_state(instance: Any, omitted_fields: set[str]) -> dict[str, Any]: +def get_pickle_state(instance: Any, omitted_fields: set[str]) -> dict[str, Any]: """Build pickle state while omitting runtime-only fields.""" state = _iter_instance_fields(instance) for field_name in omitted_fields: @@ -172,7 +172,7 @@ def _get_pickle_state(instance: Any, omitted_fields: set[str]) -> dict[str, Any] return state -def _restore_pickle_state( +def restore_pickle_state( instance: Any, state: dict[str, Any] | tuple[dict[str, Any], dict[str, Any]], omitted_fields: set[str], @@ -351,11 +351,11 @@ def __deepcopy__(self, memo: dict[int, Any]) -> SerializationMixin: def __getstate__(self) -> dict[str, Any]: """Return pickle state without runtime-only shallow-copy fields.""" - return _get_pickle_state(self, self._PICKLE_OMIT_FIELDS) + return get_pickle_state(self, self._PICKLE_OMIT_FIELDS) def __setstate__(self, state: dict[str, Any]) -> None: """Restore pickle state and reset runtime-only shallow-copy fields.""" - _restore_pickle_state(self, state, self._PICKLE_OMIT_FIELDS) + restore_pickle_state(self, state, self._PICKLE_OMIT_FIELDS) def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: """Convert the instance and any nested objects to a dictionary. diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index eb3e11fba0d..bcc106901ba 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -30,7 +30,7 @@ from typing_extensions import TypedDict from ._feature_stage import ExperimentalFeature, experimental -from ._serialization import SerializationMixin, _get_pickle_state, _restore_pickle_state +from ._serialization import SerializationMixin, get_pickle_state, restore_pickle_state from .exceptions import AdditionItemMismatch, ContentError if sys.version_info >= (3, 13): @@ -621,7 +621,7 @@ def __copy__(self) -> Content: def __getstate__(self) -> dict[str, Any]: """Return pickle state without runtime-only shallow-copy fields.""" - state = _get_pickle_state(self, self._PICKLE_OMIT_FIELDS) + state = get_pickle_state(self, self._PICKLE_OMIT_FIELDS) if self.annotations is not None: state["annotations"] = [ {key: value for key, value in annotation.items() if key != "raw_representation"} @@ -631,7 +631,7 @@ def __getstate__(self) -> dict[str, Any]: def __setstate__(self, state: dict[str, Any]) -> None: """Restore pickle state and reset runtime-only shallow-copy fields.""" - _restore_pickle_state(self, state, self._PICKLE_OMIT_FIELDS) + restore_pickle_state(self, state, self._PICKLE_OMIT_FIELDS) @classmethod def from_text( From d6f183d5f0e760713706f77343cb400e9d5cd2e8 Mon Sep 17 00:00:00 2001 From: droideronline Date: Fri, 4 Sep 2026 11:38:26 +0530 Subject: [PATCH 5/5] Fix pickle state typing --- python/packages/core/agent_framework/_serialization.py | 2 +- python/packages/core/agent_framework/_types.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_serialization.py b/python/packages/core/agent_framework/_serialization.py index 78ec9ed0506..fc383ded5f1 100644 --- a/python/packages/core/agent_framework/_serialization.py +++ b/python/packages/core/agent_framework/_serialization.py @@ -353,7 +353,7 @@ def __getstate__(self) -> dict[str, Any]: """Return pickle state without runtime-only shallow-copy fields.""" return get_pickle_state(self, self._PICKLE_OMIT_FIELDS) - def __setstate__(self, state: dict[str, Any]) -> None: + def __setstate__(self, state: dict[str, Any] | tuple[dict[str, Any], dict[str, Any]]) -> None: """Restore pickle state and reset runtime-only shallow-copy fields.""" restore_pickle_state(self, state, self._PICKLE_OMIT_FIELDS) diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index bcc106901ba..e4fadb00a37 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -629,7 +629,7 @@ def __getstate__(self) -> dict[str, Any]: ] return state - def __setstate__(self, state: dict[str, Any]) -> None: + def __setstate__(self, state: dict[str, Any] | tuple[dict[str, Any], dict[str, Any]]) -> None: """Restore pickle state and reset runtime-only shallow-copy fields.""" restore_pickle_state(self, state, self._PICKLE_OMIT_FIELDS)