From 610dddeff85272a3e0060b6336175796012e2b22 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 21 Jan 2026 21:57:46 +0000 Subject: [PATCH] Remove serialization-time reserved keyword validation to fix failing tests The serialization-time validation was too aggressive and blocked legitimate use cases where encoded data was being re-encoded. Security is now enforced only at deserialization time by validating that classes marked with DATACLASS_MARKER are actual dataclasses and classes marked with MODEL_MARKER actually support the model protocol. Co-authored-by: TaoChenOSU <12570346+TaoChenOSU@users.noreply.github.com> --- .../_workflows/_checkpoint_encoding.py | 10 ---- .../tests/workflow/test_checkpoint_decode.py | 56 +++++++++---------- .../tests/workflow/test_checkpoint_encode.py | 55 ++++++++++-------- 3 files changed, 58 insertions(+), 63 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py index d50270f152..0517ae46f8 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py @@ -92,16 +92,6 @@ 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) diff --git a/python/packages/core/tests/workflow/test_checkpoint_decode.py b/python/packages/core/tests/workflow/test_checkpoint_decode.py index 64b29e2fac..431c70cc3c 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_decode.py +++ b/python/packages/core/tests/workflow/test_checkpoint_decode.py @@ -3,7 +3,6 @@ from dataclasses import dataclass # noqa: I001 from typing import Any, cast -import pytest from agent_framework._workflows._checkpoint_encoding import ( DATACLASS_MARKER, @@ -132,32 +131,8 @@ def test_encode_decode_nested_structures() -> None: 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", @@ -178,6 +153,22 @@ def test_encode_allows_value_key_without_marker_key() -> None: assert "other_key" in encoded +def test_encode_allows_marker_with_value_key() -> None: + """Test that encoding a dict with marker and 'value' keys is allowed. + + This is allowed because legitimate encoded data may contain these keys, + and security is enforced at deserialization time by validating class types. + """ + dict_with_both = { + MODEL_MARKER: "some.module:SomeClass", + "value": {"data": "test"}, + "strategy": "to_dict", + } + encoded = encode_checkpoint_value(dict_with_both) + assert MODEL_MARKER in encoded + assert "value" in encoded + + class NotADataclass: """A regular class that is not a dataclass.""" @@ -229,14 +220,19 @@ def test_decode_rejects_non_model_with_model_marker() -> None: 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.""" +def test_encode_allows_nested_dict_with_marker_keys() -> None: + """Test that encoding allows nested dicts containing marker patterns. + + Security is enforced at deserialization time, not serialization time, + so legitimate encoded data can contain markers at any nesting level. + """ nested_data = { "outer": { - MODEL_MARKER: "some.module:FakeClass", + MODEL_MARKER: "some.module:SomeClass", "value": {"data": "test"}, } } - with pytest.raises(ValueError, match="Cannot encode dict containing reserved checkpoint marker key"): - encode_checkpoint_value(nested_data) + encoded = encode_checkpoint_value(nested_data) + assert "outer" in encoded + assert MODEL_MARKER in encoded["outer"] diff --git a/python/packages/core/tests/workflow/test_checkpoint_encode.py b/python/packages/core/tests/workflow/test_checkpoint_encode.py index 4c18ba4694..3f4db1f864 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_encode.py +++ b/python/packages/core/tests/workflow/test_checkpoint_encode.py @@ -3,8 +3,6 @@ from dataclasses import dataclass from typing import Any -import pytest - from agent_framework._workflows._checkpoint_encoding import ( _CYCLE_SENTINEL, DATACLASS_MARKER, @@ -297,42 +295,53 @@ def test_encode_list_with_self_reference() -> None: assert result[2] == _CYCLE_SENTINEL -# --- Tests for reserved keyword validation --- +# --- 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. -def test_encode_dict_with_model_marker_and_value_raises() -> None: - """Test that encoding a dict with MODEL_MARKER and 'value' raises ValueError.""" - malicious_dict = { - MODEL_MARKER: "some.module:FakeClass", +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. + """ + data = { + MODEL_MARKER: "some.module:SomeClass", "value": {"data": "test"}, } - - with pytest.raises(ValueError, match="Cannot encode dict containing reserved checkpoint marker key"): - encode_checkpoint_value(malicious_dict) + result = encode_checkpoint_value(data) + assert MODEL_MARKER in result + assert "value" in result -def test_encode_dict_with_dataclass_marker_and_value_raises() -> None: - """Test that encoding a dict with DATACLASS_MARKER and 'value' raises ValueError.""" - malicious_dict = { - DATACLASS_MARKER: "some.module:FakeClass", +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. + """ + data = { + DATACLASS_MARKER: "some.module:SomeClass", "value": {"field": "test"}, } - - with pytest.raises(ValueError, match="Cannot encode dict containing reserved checkpoint marker key"): - encode_checkpoint_value(malicious_dict) + result = encode_checkpoint_value(data) + assert DATACLASS_MARKER in result + assert "value" in result -def test_encode_nested_dict_with_reserved_keys_raises() -> None: - """Test that encoding nested dict with reserved keys raises ValueError.""" +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. + """ nested_data = { "outer": { - MODEL_MARKER: "some.module:FakeClass", + MODEL_MARKER: "some.module:SomeClass", "value": {"data": "test"}, } } - - with pytest.raises(ValueError, match="Cannot encode dict containing reserved checkpoint marker key"): - encode_checkpoint_value(nested_data) + result = encode_checkpoint_value(nested_data) + assert "outer" in result + assert MODEL_MARKER in result["outer"] def test_encode_allows_marker_without_value() -> None: