Changes before error encountered

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

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

@@ -7,6 +7,7 @@ from typing import Any, cast
 from agent_framework._workflows._checkpoint_encoding import (
     DATACLASS_MARKER,
     MODEL_MARKER,
+    PRESERVED_MARKER,
     decode_checkpoint_value,
     encode_checkpoint_value,
 )
@@ -154,10 +155,10 @@ def test_encode_allows_value_key_without_marker_key() -> None:
 
 
 def test_encode_allows_marker_with_value_key() -> None:
-    """Test that encoding a dict with marker and 'value' keys is allowed.
+    """Test that encoding a dict with marker and 'value' keys wraps it in preservation envelope.
 
-    This is allowed because legitimate encoded data may contain these keys,
-    and security is enforced at deserialization time by validating class types.
+    User data that looks like a marker pattern is preserved to prevent confusion
+    during deserialization.
     """
     dict_with_both = {
         MODEL_MARKER: "some.module:SomeClass",
@@ -165,8 +166,12 @@ def test_encode_allows_marker_with_value_key() -> None:
         "strategy": "to_dict",
     }
     encoded = encode_checkpoint_value(dict_with_both)
-    assert MODEL_MARKER in encoded
+    # Should be wrapped in preservation envelope
+    assert PRESERVED_MARKER in encoded
+    assert encoded[PRESERVED_MARKER] is True
     assert "value" in encoded
+    # Original dict is inside the envelope
+    assert MODEL_MARKER in encoded["value"]
 
 
 class NotADataclass:
@@ -221,10 +226,10 @@ def test_decode_rejects_non_model_with_model_marker() -> None:
 
 
 def test_encode_allows_nested_dict_with_marker_keys() -> None:
-    """Test that encoding allows nested dicts containing marker patterns.
+    """Test that encoding wraps nested dicts containing marker patterns in preservation envelope.
 
-    Security is enforced at deserialization time, not serialization time,
-    so legitimate encoded data can contain markers at any nesting level.
+    User data that looks like marker patterns is preserved at any nesting level
+    to prevent confusion during deserialization.
     """
     nested_data = {
         "outer": {
@@ -235,4 +240,6 @@ def test_encode_allows_nested_dict_with_marker_keys() -> None:
 
     encoded = encode_checkpoint_value(nested_data)
     assert "outer" in encoded
-    assert MODEL_MARKER in encoded["outer"]
+    # Nested dict should be wrapped in preservation envelope
+    assert PRESERVED_MARKER in encoded["outer"]
+    assert encoded["outer"][PRESERVED_MARKER] is True

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

@@ -7,6 +7,8 @@ from agent_framework._workflows._checkpoint_encoding import (
     _CYCLE_SENTINEL,
     DATACLASS_MARKER,
     MODEL_MARKER,
+    PRESERVED_MARKER,
+    decode_checkpoint_value,
     encode_checkpoint_value,
 )
 
@@ -296,43 +298,44 @@ def test_encode_list_with_self_reference() -> None:
 
 
 # --- Tests for reserved keyword handling ---
-# Note: Security is enforced at deserialization time by validating class types,
-# not at serialization time. This allows legitimate encoded data to be re-encoded.
+# User data containing marker keys + "value" is preserved in a special envelope
+# during serialization and recovered during deserialization.
 
 
-def test_encode_allows_dict_with_model_marker_and_value() -> None:
-    """Test that encoding a dict with MODEL_MARKER and 'value' is allowed.
-
-    Security is enforced at deserialization time, not serialization time.
-    """
+def test_encode_preserves_dict_with_model_marker_and_value() -> None:
+    """Test that user dict with MODEL_MARKER and 'value' is preserved in envelope."""
     data = {
         MODEL_MARKER: "some.module:SomeClass",
         "value": {"data": "test"},
     }
     result = encode_checkpoint_value(data)
-    assert MODEL_MARKER in result
+    # Should be wrapped in preservation envelope
+    assert PRESERVED_MARKER in result
+    assert result[PRESERVED_MARKER] is True
     assert "value" in result
+    # The inner value should contain the original dict
+    assert MODEL_MARKER in result["value"]
+    assert result["value"]["value"] == {"data": "test"}
 
 
-def test_encode_allows_dict_with_dataclass_marker_and_value() -> None:
-    """Test that encoding a dict with DATACLASS_MARKER and 'value' is allowed.
-
-    Security is enforced at deserialization time, not serialization time.
-    """
+def test_encode_preserves_dict_with_dataclass_marker_and_value() -> None:
+    """Test that user dict with DATACLASS_MARKER and 'value' is preserved in envelope."""
     data = {
         DATACLASS_MARKER: "some.module:SomeClass",
         "value": {"field": "test"},
     }
     result = encode_checkpoint_value(data)
-    assert DATACLASS_MARKER in result
+    # Should be wrapped in preservation envelope
+    assert PRESERVED_MARKER in result
+    assert result[PRESERVED_MARKER] is True
     assert "value" in result
+    # The inner value should contain the original dict
+    assert DATACLASS_MARKER in result["value"]
+    assert result["value"]["value"] == {"field": "test"}
 
 
-def test_encode_allows_nested_dict_with_marker_keys() -> None:
-    """Test that encoding nested dict with marker keys is allowed.
-
-    Security is enforced at deserialization time, not serialization time.
-    """
+def test_encode_preserves_nested_dict_with_marker_keys() -> None:
+    """Test that nested dict with marker keys is preserved in envelope."""
     nested_data = {
         "outer": {
             MODEL_MARKER: "some.module:SomeClass",
@@ -341,27 +344,61 @@ def test_encode_allows_nested_dict_with_marker_keys() -> None:
     }
     result = encode_checkpoint_value(nested_data)
     assert "outer" in result
-    assert MODEL_MARKER in result["outer"]
+    # The nested dict should be wrapped in preservation envelope
+    assert PRESERVED_MARKER in result["outer"]
+    assert result["outer"][PRESERVED_MARKER] is True
+
+
+def test_decode_recovers_preserved_dict_with_model_marker() -> None:
+    """Test that preserved dict with MODEL_MARKER is recovered correctly."""
+    original_data = {
+        MODEL_MARKER: "some.module:SomeClass",
+        "value": {"data": "test"},
+    }
+    encoded = encode_checkpoint_value(original_data)
+    decoded = decode_checkpoint_value(encoded)
+    # Should recover the original dict structure
+    assert MODEL_MARKER in decoded
+    assert decoded[MODEL_MARKER] == "some.module:SomeClass"
+    assert decoded["value"] == {"data": "test"}
+
+
+def test_decode_recovers_preserved_dict_with_dataclass_marker() -> None:
+    """Test that preserved dict with DATACLASS_MARKER is recovered correctly."""
+    original_data = {
+        DATACLASS_MARKER: "some.module:SomeClass",
+        "value": {"field": "test"},
+    }
+    encoded = encode_checkpoint_value(original_data)
+    decoded = decode_checkpoint_value(encoded)
+    # Should recover the original dict structure
+    assert DATACLASS_MARKER in decoded
+    assert decoded[DATACLASS_MARKER] == "some.module:SomeClass"
+    assert decoded["value"] == {"field": "test"}
 
 
 def test_encode_allows_marker_without_value() -> None:
-    """Test that a dict with marker key but without 'value' key is allowed."""
+    """Test that a dict with marker key but without 'value' key is NOT preserved."""
     data = {
         MODEL_MARKER: "some.module:SomeClass",
         "other_key": "allowed",
     }
     result = encode_checkpoint_value(data)
+    # Should NOT be wrapped (no "value" key present)
+    assert PRESERVED_MARKER not in result
     assert MODEL_MARKER in result
     assert result["other_key"] == "allowed"
 
 
 def test_encode_allows_value_without_marker() -> None:
-    """Test that a dict with 'value' key but without marker is allowed."""
+    """Test that a dict with 'value' key but without marker is NOT preserved."""
     data = {
         "value": {"nested": "data"},
         "other_key": "allowed",
     }
     result = encode_checkpoint_value(data)
+    # Should NOT be wrapped (no marker key present)
+    assert PRESERVED_MARKER not in result
     assert "value" in result
     assert result["other_key"] == "allowed"