fix: fix serialization of State and PipelineSnapshot (#12113)

Author: sjrl

haystack/components/agents/state/state.py

@@ -7,7 +7,8 @@
 from typing import Any, get_args
 
 from haystack.dataclasses import ChatMessage
-from haystack.utils import _deserialize_value_with_schema, _serialize_value_with_schema
+from haystack.utils import _deserialize_value_with_schema
+from haystack.utils.base_serialization import _serialize_with_field_fallback
 from haystack.utils.callable_serialization import deserialize_callable, serialize_callable
 from haystack.utils.type_serialization import deserialize_type, serialize_type
 
@@ -194,7 +195,9 @@ def to_dict(self) -> dict[str, Any]:
         """
         serialized = {}
         serialized["schema"] = _schema_to_dict(self.schema)
-        serialized["data"] = _serialize_value_with_schema(self._data)
+        # Field-level fallback so a single non-serializable value omits only that field instead of
+        # failing the whole State serialization.
+        serialized["data"] = _serialize_with_field_fallback(self._data, description="the agent's State data")
         return serialized
 
     @classmethod

haystack/core/pipeline/breakpoint.py

@@ -13,9 +13,8 @@
 
 from haystack import logging
 from haystack.core.errors import PipelineInvalidPipelineSnapshotError
-from haystack.core.pipeline.utils import _deepcopy_with_exceptions
 from haystack.dataclasses.breakpoints import Breakpoint, PipelineSnapshot, PipelineState
-from haystack.utils.base_serialization import _serialize_value_with_schema
+from haystack.utils.base_serialization import _serialize_with_field_fallback
 
 logger = logging.getLogger(__name__)
 
@@ -292,54 +291,3 @@ def _transform_json_structure(data: dict[str, Any] | list[Any] | Any) -> Any:
 
     # For other data types, just return the value as is.
     return data
-
-
-def _serialize_with_field_fallback(payload: Any, *, description: str) -> dict[str, Any]:
-    """
-    Serialize a payload and, on failure, retry field-by-field to preserve resumable fields.
-
-    If the whole payload serializes, the result is returned as-is. Otherwise, and if the payload is a
-    mapping, each top-level field is serialized individually and only the failing fields are omitted.
-    When the payload is not a mapping, or when every field fails to serialize, the helper returns a
-    structurally valid empty-object payload so that the downstream ``_deserialize_value_with_schema``
-    can still load it back instead of raising ``DeserializationError`` on a bare ``{}``.
-
-    :param payload: The value to serialize.
-    :param description: Short human-readable label used in warning messages, for example
-        ``"the agent's chat_generator inputs"`` or ``"the inputs of the current pipeline state"``.
-    :returns: A dict of the form ``{"serialization_schema": ..., "serialized_data": ...}``.
-    """
-    try:
-        return _serialize_value_with_schema(_deepcopy_with_exceptions(payload))
-    except Exception as error:
-        logger.warning(
-            "Failed to serialize {description}. "
-            "Haystack will omit only the non-serializable fields when possible. Error: {e}",
-            description=description,
-            e=error,
-        )
-
-    serialized_properties: dict[str, Any] = {}
-    serialized_data: dict[str, Any] = {}
-
-    if isinstance(payload, dict):
-        for field_name, value in payload.items():
-            try:
-                serialized_value = _serialize_value_with_schema(_deepcopy_with_exceptions(value))
-            except Exception as field_error:
-                logger.warning(
-                    "Failed to serialize the '{field_name}' field of {description}. "
-                    "The field will be omitted from the snapshot. Error: {e}",
-                    field_name=field_name,
-                    description=description,
-                    e=field_error,
-                )
-                continue
-
-            serialized_properties[field_name] = serialized_value["serialization_schema"]
-            serialized_data[field_name] = serialized_value["serialized_data"]
-
-    return {
-        "serialization_schema": {"type": "object", "properties": serialized_properties},
-        "serialized_data": serialized_data,
-    }

haystack/utils/base_serialization.py

@@ -8,7 +8,7 @@
 import pydantic
 
 from haystack import logging
-from haystack.core.errors import DeserializationError
+from haystack.core.errors import DeserializationError, SerializationError
 from haystack.core.serialization import generate_qualified_class_name, import_class_by_name
 from haystack.utils import deserialize_callable, serialize_callable
 
@@ -60,7 +60,7 @@ def _serialize_value_with_schema(payload: Any) -> dict[str, Any]:  # noqa: PLR09
         return {"serialization_schema": {"type": "object", "properties": schema}, "serialized_data": data}
 
     # Handle array case - iterate through elements
-    if isinstance(payload, (list, tuple, set)):
+    if isinstance(payload, (list, tuple, set, frozenset)):
         # Serialize each item in the array
         serialized_list = []
         for item in payload:
@@ -76,9 +76,12 @@ def _serialize_value_with_schema(payload: Any) -> dict[str, Any]:  # noqa: PLR09
         else:
             base_schema = {"type": "array", "items": {}}
 
-        # Add JSON Schema properties to infer sets and tuples
-        if isinstance(payload, set):
+        # Add JSON Schema properties to infer sets, frozensets and tuples
+        if isinstance(payload, (set, frozenset)):
             base_schema["uniqueItems"] = True
+            # `frozen` distinguishes frozenset from set on deserialization
+            if isinstance(payload, frozenset):
+                base_schema["frozen"] = True
         elif isinstance(payload, tuple):
             base_schema["minItems"] = len(payload)
             base_schema["maxItems"] = len(payload)
@@ -101,19 +104,71 @@ def _serialize_value_with_schema(payload: Any) -> dict[str, Any]:  # noqa: PLR09
         type_name = generate_qualified_class_name(type(payload))
         return {"serialization_schema": {"type": type_name}, "serialized_data": payload.name}
 
-    # Handle arbitrary objects with __dict__
-    if hasattr(payload, "__dict__"):
-        type_name = generate_qualified_class_name(type(payload))
-        schema = {"type": type_name}
-        serialized_data = {}
-        for key, value in vars(payload).items():
-            serialized_value = _serialize_value_with_schema(value)
-            serialized_data[key] = serialized_value["serialized_data"]
-        return {"serialization_schema": schema, "serialized_data": serialized_data}
-
     # Handle primitives
-    schema = {"type": _primitive_schema_type(payload)}
-    return {"serialization_schema": schema, "serialized_data": payload}
+    if payload is None or isinstance(payload, (bool, int, float, str)):
+        return {"serialization_schema": {"type": _primitive_schema_type(payload)}, "serialized_data": payload}
+
+    # Unsupported type: raise so callers can omit the field (see `_serialize_with_field_fallback`)
+    # instead of embedding a live, non-portable object in the output.
+    raise SerializationError(
+        f"Cannot serialize value of type '{type(payload).__name__}'. Supported values are primitives "
+        f"(str, int, float, bool, None), lists/tuples/sets/frozensets/dicts of supported values, Enums, "
+        f"callables, pydantic models, and objects exposing a 'to_dict' method."
+    )
+
+
+def _serialize_with_field_fallback(payload: Any, *, description: str) -> dict[str, Any]:
+    """
+    Serialize a payload and, on failure, retry field-by-field to preserve serializable fields.
+
+    If the whole payload serializes, the result is returned as-is. Otherwise, and if the payload is a
+    mapping, each top-level field is serialized individually and only the failing fields are omitted.
+    When the payload is not a mapping, or when every field fails to serialize, the helper returns a
+    structurally valid empty-object payload so that the downstream `_deserialize_value_with_schema`
+    can still load it back instead of raising `DeserializationError` on a bare `{}`.
+
+    :param payload: The value to serialize.
+    :param description: Short human-readable label used in warning messages, for example
+        `"the agent's State data"` or `"the inputs of the current pipeline state"`.
+    :returns: A dict of the form `{"serialization_schema": ..., "serialized_data": ...}`.
+    """
+    # Local import to avoid a circular import at module load time.
+    from haystack.core.pipeline.utils import _deepcopy_with_exceptions
+
+    try:
+        return _serialize_value_with_schema(_deepcopy_with_exceptions(payload))
+    except Exception as error:
+        logger.warning(
+            "Failed to serialize {description}. "
+            "Haystack will omit only the non-serializable fields when possible. Error: {e}",
+            description=description,
+            e=error,
+        )
+
+    serialized_properties: dict[str, Any] = {}
+    serialized_data: dict[str, Any] = {}
+
+    if isinstance(payload, dict):
+        for field_name, value in payload.items():
+            try:
+                serialized_value = _serialize_value_with_schema(_deepcopy_with_exceptions(value))
+            except Exception as field_error:
+                logger.warning(
+                    "Failed to serialize the '{field_name}' field of {description}. "
+                    "The field will be omitted from the snapshot. Error: {e}",
+                    field_name=field_name,
+                    description=description,
+                    e=field_error,
+                )
+                continue
+
+            serialized_properties[field_name] = serialized_value["serialization_schema"]
+            serialized_data[field_name] = serialized_value["serialized_data"]
+
+    return {
+        "serialization_schema": {"type": "object", "properties": serialized_properties},
+        "serialized_data": serialized_data,
+    }
 
 
 def _primitive_schema_type(value: Any) -> str:
@@ -183,10 +238,10 @@ def _deserialize_value_with_schema(serialized: dict[str, Any]) -> Any:
             _deserialize_value_with_schema({"serialization_schema": schema["items"], "serialized_data": item})
             for item in data
         ]
-        final_array: list | set | tuple
-        # Is a set if uniqueItems is True
+        final_array: list | set | frozenset | tuple
+        # Is a set or frozenset if uniqueItems is True
         if schema.get("uniqueItems") is True:
-            final_array = set(deserialized_items)
+            final_array = frozenset(deserialized_items) if schema.get("frozen") is True else set(deserialized_items)
         # Is a tuple if minItems and maxItems are set
         elif schema.get("minItems") is not None and schema.get("maxItems") is not None:
             final_array = tuple(deserialized_items)

releasenotes/notes/fix-serialization-passthrough-9f1c2a4b7e6d3f08.yaml

@@ -0,0 +1,13 @@
+---
+fixes:
+  - |
+    Fixed the schema-aware serialization helper used for pipeline snapshots and ``Agent`` ``State``
+    (``_serialize_value_with_schema``) so it no longer silently passes unsupported objects through as
+    if they were serialized. Values such as ``datetime``, ``bytes``, ``complex`` and arbitrary objects
+    without a ``to_dict`` method were previously stored unchanged and mislabeled as strings, which broke
+    JSON storage and round-tripping of snapshots. Unsupported values now raise a ``SerializationError``,
+    and the callers that build snapshots (pipeline breakpoints and ``State.to_dict``) catch it to omit
+    only the offending field while keeping the rest of the payload resumable.
+  - |
+    Added support for serializing and deserializing ``frozenset`` values in ``_serialize_value_with_schema``.
+    A ``frozenset`` now round-trips back to a ``frozenset`` instead of being dropped.

test/components/agents/test_state_class.py

@@ -3,7 +3,9 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import inspect
+import logging
 from dataclasses import dataclass
+from datetime import datetime
 from typing import Dict, Generic, List, Optional, TypeVar, Union
 
 import pytest
@@ -512,6 +514,16 @@ def test_state_to_dict_typing_list(self):
             },
         }
 
+    def test_state_to_dict_omits_non_serializable_field(self, caplog):
+        # A non-serializable value must not crash serialization: only that field is omitted,
+        # and a warning is logged.
+        state = State({"good": {"type": int}, "bad": {"type": object}}, {"good": 1, "bad": datetime(2024, 1, 1)})
+        with caplog.at_level(logging.WARNING):
+            state_dict = state.to_dict()
+        assert state_dict["data"]["serialized_data"] == {"good": 1}
+        assert "bad" not in state_dict["data"]["serialization_schema"]["properties"]
+        assert any("Failed to serialize the 'bad' field of the agent's State data" in msg for msg in caplog.messages)
+
     def test_state_from_dict_typing_list(self):
         state_dict = {
             "schema": {

test/core/pipeline/test_breakpoint.py

@@ -349,19 +349,24 @@ def to_dict(self):
 def test_save_pipeline_snapshot_raises_on_failure(tmp_path, caplog, monkeypatch):
     monkeypatch.setenv(HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED, "true")
 
+    # Point the snapshot directory below an existing file so creating it fails with a filesystem
+    # error, exercising the raise_on_failure contract.
+    blocking_file = tmp_path / "not_a_dir"
+    blocking_file.write_text("i am a file")
+    snapshot_path = blocking_file / "snapshots"
+
     snapshot = _create_pipeline_snapshot(
         inputs={},
         component_inputs={},
-        break_point=Breakpoint(component_name="comp2", snapshot_file_path=str(tmp_path)),
+        break_point=Breakpoint(component_name="comp2", snapshot_file_path=str(snapshot_path)),
         component_visits={"comp1": 1, "comp2": 0},
         original_input_data={},
         ordered_component_names=["comp1", "comp2"],
         include_outputs_from={"comp1"},
-        # We use a non-serializable type (bytes) directly in pipeline outputs to trigger the error
-        pipeline_outputs={"comp1": {"result": b"test"}},
+        pipeline_outputs={"comp1": {"result": "test"}},
     )
 
-    with pytest.raises(TypeError):
+    with pytest.raises(OSError):
         _save_pipeline_snapshot(snapshot)
 
     with caplog.at_level(logging.ERROR):