diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py index 516a4547a0..00920d25ed 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/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. diff --git a/python/packages/core/tests/workflow/test_checkpoint_decode.py b/python/packages/core/tests/workflow/test_checkpoint_decode.py index 431c70cc3c..224bf74ec8 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_decode.py +++ b/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 diff --git a/python/packages/core/tests/workflow/test_checkpoint_encode.py b/python/packages/core/tests/workflow/test_checkpoint_encode.py index 3f4db1f864..7f2d4bfb9e 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_encode.py +++ b/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"