Add validation for reserved keywords in checkpoint encoding/decoding

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

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)