Python: Fix checkpoint pickling of runtime raw representations

python/packages/core/agent_framework/_serialization.py

@@ -151,6 +151,43 @@ 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] | 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:
+        if field_name in _iter_instance_fields(instance) or hasattr(instance, "__dict__"):
+            object.__setattr__(instance, field_name, None)
+
+
 class SerializationMixin:
     """Mixin class providing comprehensive serialization and deserialization capabilities.
 
@@ -284,6 +321,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,13 +342,21 @@ 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:
                 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."""
+        return get_pickle_state(self, self._PICKLE_OMIT_FIELDS)
+
+    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)
+
     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.
 

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,
@@ -610,6 +611,28 @@ 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)
+        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] | 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)
+
     @classmethod
     def from_text(
         cls: type[ContentT],

python/packages/core/tests/core/test_serializable_mixin.py

@@ -572,6 +572,59 @@ 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_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):
+            _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."""
 

python/packages/core/tests/core/test_types.py

@@ -2733,6 +2733,29 @@ 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_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

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
 # ---------------------------------------------------------------------------