diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py index 71dcefacb55..03e1bb0cd98 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py @@ -307,8 +307,9 @@ class MyState: ... def _encode(value: Any) -> Any: """Recursively encode a value for JSON storage.""" - # JSON-native types pass through - if isinstance(value, _JSON_NATIVE_TYPES): + # JSON-native types pass through only for exact types so scalar subclasses + # (such as IntEnum, StrEnum, or custom scalar subclasses) preserve their type. + if type(value) in _JSON_NATIVE_TYPES: return value # Recursively encode dict values (keys become strings). Only plain dicts diff --git a/python/packages/core/tests/workflow/test_checkpoint_encode.py b/python/packages/core/tests/workflow/test_checkpoint_encode.py index 3b2a9149b9e..de8808a1918 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_encode.py +++ b/python/packages/core/tests/workflow/test_checkpoint_encode.py @@ -4,6 +4,7 @@ from collections import Counter, OrderedDict, defaultdict from dataclasses import dataclass from datetime import datetime, timezone +from enum import IntEnum from typing import Any, cast from agent_framework._workflows._checkpoint_encoding import ( @@ -14,6 +15,19 @@ ) +class StatusEnum(IntEnum): + PENDING = 1 + ACTIVE = 2 + + +class CustomInt(int): + pass + + +class CustomStr(str): + pass + + @dataclass class SimpleDataclass: """A simple dataclass for testing encoding.""" @@ -87,6 +101,33 @@ def test_encode_none() -> None: assert encode_checkpoint_value(None) is None +def test_encode_scalar_subclasses_preserves_type_fidelity() -> None: + """Test that scalar subclasses like IntEnum are pickled rather than flattened to raw primitives.""" + status = StatusEnum.ACTIVE + encoded = encode_checkpoint_value(status) + assert isinstance(encoded, dict) + assert _PICKLE_MARKER in encoded + decoded = decode_checkpoint_value(encoded) + assert decoded == StatusEnum.ACTIVE + assert type(decoded) is StatusEnum + + custom_int = CustomInt(42) + encoded_int = encode_checkpoint_value(custom_int) + assert isinstance(encoded_int, dict) + assert _PICKLE_MARKER in encoded_int + decoded_int = decode_checkpoint_value(encoded_int) + assert decoded_int == 42 + assert type(decoded_int) is CustomInt + + custom_str = CustomStr("custom") + encoded_str = encode_checkpoint_value(custom_str) + assert isinstance(encoded_str, dict) + assert _PICKLE_MARKER in encoded_str + decoded_str = decode_checkpoint_value(encoded_str) + assert decoded_str == "custom" + assert type(decoded_str) is CustomStr + + # --- Tests for collection encoding ---