diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py index 8c52c7a521..d5a1a7f970 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/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 diff --git a/python/packages/core/tests/workflow/test_checkpoint_decode.py b/python/packages/core/tests/workflow/test_checkpoint_decode.py index b126eafacf..64b29e2fac 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_decode.py +++ b/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)