Add validation for reserved keywords in checkpoint encoding/decoding

Co-authored-by: TaoChenOSU <12570346+TaoChenOSU@users.noreply.github.com>

python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py

@@ -92,6 +92,16 @@ def encode_checkpoint_value(value: Any) -> Any:
                 return _CYCLE_SENTINEL
             stack.add(oid)
             try:
+                # Check for reserved marker keys that could cause deserialization issues
+                has_model_marker = MODEL_MARKER in v_dict
+                has_dataclass_marker = DATACLASS_MARKER in v_dict
+                has_value_key = "value" in v_dict
+                if (has_model_marker or has_dataclass_marker) and has_value_key:
+                    marker = MODEL_MARKER if has_model_marker else DATACLASS_MARKER
+                    raise ValueError(
+                        f"Cannot encode dict containing reserved checkpoint marker key '{marker}' "
+                        f"with 'value' key. These keys are reserved for polymorphic serialization."
+                    )
                 json_dict: dict[str, Any] = {}
                 for k_any, val_any in v_dict.items():  # type: ignore[assignment]
                     k_str: str = str(k_any)
@@ -146,6 +156,12 @@ def decode_checkpoint_value(value: Any) -> Any:
                     cls = None
 
                 if cls is not None:
+                    # Verify the class actually supports the model protocol
+                    if not _class_supports_model_protocol(cls):
+                        logger.debug(
+                            f"Class {type_key} does not support model protocol; returning raw value"
+                        )
+                        return decoded_payload
                     if strategy == "to_dict" and hasattr(cls, "from_dict"):
                         with contextlib.suppress(Exception):
                             return cls.from_dict(decoded_payload)
@@ -169,6 +185,12 @@ def decode_checkpoint_value(value: Any) -> Any:
                     if module is None:
                         module = importlib.import_module(module_name)
                     cls_dc: Any = getattr(module, class_name)
+                    # Verify the class is actually a dataclass
+                    if not is_dataclass(cls_dc):
+                        logger.debug(
+                            f"Class {type_key_dc} is not a dataclass; returning raw value"
+                        )
+                        return decoded_raw
                     constructed = _instantiate_checkpoint_dataclass(cls_dc, decoded_raw)
                     if constructed is not None:
                         return constructed
@@ -204,6 +226,21 @@ def _supports_model_protocol(obj: object) -> bool:
     return (has_to_dict and has_from_dict) or (has_to_json and has_from_json)
 
 
+def _class_supports_model_protocol(cls: type[Any]) -> bool:
+    """Check if a class type supports the model serialization protocol.
+
+    This is similar to _supports_model_protocol but works on classes directly,
+    used for validation during deserialization.
+    """
+    has_to_dict = hasattr(cls, "to_dict") and callable(getattr(cls, "to_dict", None))
+    has_from_dict = hasattr(cls, "from_dict") and callable(getattr(cls, "from_dict", None))
+
+    has_to_json = hasattr(cls, "to_json") and callable(getattr(cls, "to_json", None))
+    has_from_json = hasattr(cls, "from_json") and callable(getattr(cls, "from_json", None))
+
+    return (has_to_dict and has_from_dict) or (has_to_json and has_from_json)
+
+
 def _import_qualified_name(qualname: str) -> type[Any] | None:
     if ":" not in qualname:
         return None

python/packages/core/tests/workflow/test_checkpoint_decode.py

@@ -3,7 +3,11 @@
 from dataclasses import dataclass  # noqa: I001
 from typing import Any, cast
 
+import pytest
+
 from agent_framework._workflows._checkpoint_encoding import (
+    DATACLASS_MARKER,
+    MODEL_MARKER,
     decode_checkpoint_value,
     encode_checkpoint_value,
 )
@@ -126,3 +130,113 @@ def test_encode_decode_nested_structures() -> None:
     assert response.data == "first response"
     assert isinstance(response.original_request, SampleRequest)
     assert response.original_request.request_id == "req-1"
+
+
+def test_encode_raises_error_for_reserved_model_marker_with_value() -> None:
+    """Test that encoding a dict with MODEL_MARKER and 'value' keys raises an error."""
+    malicious_dict = {
+        MODEL_MARKER: "some.module:FakeClass",
+        "value": {"data": "test"},
+        "strategy": "to_dict",
+    }
+
+    with pytest.raises(ValueError, match="Cannot encode dict containing reserved checkpoint marker key"):
+        encode_checkpoint_value(malicious_dict)
+
+
+def test_encode_raises_error_for_reserved_dataclass_marker_with_value() -> None:
+    """Test that encoding a dict with DATACLASS_MARKER and 'value' keys raises an error."""
+    malicious_dict = {
+        DATACLASS_MARKER: "some.module:FakeClass",
+        "value": {"field1": "test"},
+    }
+
+    with pytest.raises(ValueError, match="Cannot encode dict containing reserved checkpoint marker key"):
+        encode_checkpoint_value(malicious_dict)
+
+
+def test_encode_allows_marker_key_without_value_key() -> None:
+    """Test that encoding a dict with only the marker key (no 'value') is allowed."""
+    # This should not raise, as it doesn't match the marker pattern
+    dict_with_marker_only = {
+        MODEL_MARKER: "some.module:FakeClass",
+        "other_key": "test",
+    }
+    encoded = encode_checkpoint_value(dict_with_marker_only)
+    assert MODEL_MARKER in encoded
+    assert "other_key" in encoded
+
+
+def test_encode_allows_value_key_without_marker_key() -> None:
+    """Test that encoding a dict with only 'value' key (no marker) is allowed."""
+    dict_with_value_only = {
+        "value": {"data": "test"},
+        "other_key": "test",
+    }
+    encoded = encode_checkpoint_value(dict_with_value_only)
+    assert "value" in encoded
+    assert "other_key" in encoded
+
+
+class NotADataclass:
+    """A regular class that is not a dataclass."""
+
+    def __init__(self, value: str) -> None:
+        self.value = value
+
+    def get_value(self) -> str:
+        return self.value
+
+
+class NotAModel:
+    """A regular class that does not support the model protocol."""
+
+    def __init__(self, value: str) -> None:
+        self.value = value
+
+    def get_value(self) -> str:
+        return self.value
+
+
+def test_decode_rejects_non_dataclass_with_dataclass_marker() -> None:
+    """Test that decode returns raw value when marked class is not a dataclass."""
+    # Manually construct a payload that claims NotADataclass is a dataclass
+    fake_payload = {
+        DATACLASS_MARKER: f"{NotADataclass.__module__}:{NotADataclass.__name__}",
+        "value": {"value": "test_value"},
+    }
+
+    decoded = decode_checkpoint_value(fake_payload)
+
+    # Should return the raw decoded value, not an instance of NotADataclass
+    assert isinstance(decoded, dict)
+    assert decoded["value"] == "test_value"
+
+
+def test_decode_rejects_non_model_with_model_marker() -> None:
+    """Test that decode returns raw value when marked class doesn't support model protocol."""
+    # Manually construct a payload that claims NotAModel supports the model protocol
+    fake_payload = {
+        MODEL_MARKER: f"{NotAModel.__module__}:{NotAModel.__name__}",
+        "strategy": "to_dict",
+        "value": {"value": "test_value"},
+    }
+
+    decoded = decode_checkpoint_value(fake_payload)
+
+    # Should return the raw decoded value, not an instance of NotAModel
+    assert isinstance(decoded, dict)
+    assert decoded["value"] == "test_value"
+
+
+def test_encode_raises_for_nested_dict_with_reserved_keys() -> None:
+    """Test that encoding fails for nested dicts containing reserved marker patterns."""
+    nested_data = {
+        "outer": {
+            MODEL_MARKER: "some.module:FakeClass",
+            "value": {"data": "test"},
+        }
+    }
+
+    with pytest.raises(ValueError, match="Cannot encode dict containing reserved checkpoint marker key"):
+        encode_checkpoint_value(nested_data)