Changes before error encountered

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

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

@@ -10,6 +10,7 @@ from typing import Any, cast
 # Checkpoint serialization helpers
 MODEL_MARKER = "__af_model__"
 DATACLASS_MARKER = "__af_dataclass__"
+PRESERVED_MARKER = "__af_preserved__"
 
 # Guards to prevent runaway recursion while encoding arbitrary user data
 _MAX_ENCODE_DEPTH = 100
@@ -92,9 +93,22 @@ def encode_checkpoint_value(value: Any) -> Any:
                 return _CYCLE_SENTINEL
             stack.add(oid)
             try:
-                json_dict: dict[str, Any] = {}
+                # Check if this dict looks like a marker pattern (has marker key + "value" key)
+                # If so, preserve it to prevent confusion during deserialization
+                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:
+                    # This is user data that looks like a marker pattern - preserve it
+                    json_dict: dict[str, Any] = {}
+                    for k_any, val_any in v_dict.items():  # type: ignore[assignment]
+                        k_str: str = str(k_any)
+                        json_dict[k_str] = _enc(val_any, stack, depth + 1)
+                    return {PRESERVED_MARKER: True, "value": json_dict}
+
+                json_dict = {}
                 for k_any, val_any in v_dict.items():  # type: ignore[assignment]
-                    k_str: str = str(k_any)
+                    k_str = str(k_any)
                     json_dict[k_str] = _enc(val_any, stack, depth + 1)
                 return json_dict
             finally:
@@ -132,6 +146,13 @@ def decode_checkpoint_value(value: Any) -> Any:
     """Recursively decode values previously encoded by encode_checkpoint_value."""
     if isinstance(value, dict):
         value_dict = cast(dict[str, Any], value)  # encoded form always uses string keys
+
+        # Handle preserved marker first - this was user data that looked like a marker pattern
+        if PRESERVED_MARKER in value_dict and "value" in value_dict:
+            preserved_value = value_dict.get("value")
+            # Decode as a regular dict without marker interpretation
+            return _decode_as_regular_dict(preserved_value)
+
         # Structured model marker handling
         if MODEL_MARKER in value_dict and "value" in value_dict:
             type_key: str | None = value_dict.get(MODEL_MARKER)  # type: ignore[assignment]
@@ -196,6 +217,23 @@ def decode_checkpoint_value(value: Any) -> Any:
     return value
 
 
+def _decode_as_regular_dict(value: Any) -> Any:
+    """Decode value as a regular dict/list without marker interpretation.
+
+    Used to recover preserved user data that looked like marker patterns.
+    """
+    if isinstance(value, dict):
+        value_dict = cast(dict[str, Any], value)
+        decoded: dict[str, Any] = {}
+        for k_any, v_any in value_dict.items():
+            decoded[k_any] = _decode_as_regular_dict(v_any)
+        return decoded
+    if isinstance(value, list):
+        value_list: list[Any] = value  # type: ignore[assignment]
+        return [_decode_as_regular_dict(v_any) for v_any in value_list]
+    return value
+
+
 def _class_supports_model_protocol(cls: type[Any]) -> bool:
     """Check if a class type supports the model serialization protocol.
 

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"