Python: fix(workflows): preserve scalar subclasses and IntEnum in checkpoint encoding

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

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 ---