[BREAKING] Python: Checkpoint refactor: encode/decode, checkpoint format, etc (#3744)

* WIP: Checkpoint refactor: encode/decode, checkpoint format, etc

* WIP: Remove workflow ID in checkpoints

* Refactor checkpointing

* Add get_latest tests

* Increase test coverage

* Fix formatting

* Fix unit tests

* Fix samples

* fix unit tests

* fix pipeline

* Copilot comments

* Fix tests

* Fix more tests

* Address comments part 1

* Address comments part 2

* Comments

python/packages/core/tests/workflow/test_agent_executor.py

@@ -84,16 +84,17 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
     assert initial_agent.call_count == 1
 
     # Verify checkpoint was created
-    checkpoints = await storage.list_checkpoints()
-    assert len(checkpoints) > 0
-
-    # Find a suitable checkpoint to restore (prefer superstep checkpoint)
-    checkpoints.sort(key=lambda cp: cp.timestamp)
-    restore_checkpoint = next(
-        (cp for cp in checkpoints if (cp.metadata or {}).get("checkpoint_type") == "superstep"),
-        checkpoints[-1],
+    checkpoints = await storage.list_checkpoints(workflow_name=wf.name)
+    assert len(checkpoints) >= 2, (
+        "Expected at least 2 checkpoints. The first one is after the start executor, "
+        "and the second one is after the agent execution."
     )
 
+    # Get the second checkpoint which should contain the state after processing
+    # the first message by the start executor in the sequential workflow
+    checkpoints.sort(key=lambda cp: cp.timestamp)
+    restore_checkpoint = checkpoints[1]
+
     # Verify checkpoint contains executor state with both cache and thread
     assert "_executor_state" in restore_checkpoint.state
     executor_states = restore_checkpoint.state["_executor_state"]

python/packages/core/tests/workflow/test_checkpoint_decode.py

@@ -1,16 +1,17 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from dataclasses import dataclass  # noqa: I001
+from dataclasses import dataclass
+from datetime import datetime, timezone
 from typing import Any, cast
 
+import pytest
 
 from agent_framework._workflows._checkpoint_encoding import (
-    DATACLASS_MARKER,
-    MODEL_MARKER,
+    _TYPE_MARKER,  # type: ignore
+    CheckpointDecodingError,
     decode_checkpoint_value,
     encode_checkpoint_value,
 )
-from agent_framework._workflows._typing_utils import is_instance_of
 
 
 @dataclass
@@ -30,7 +31,22 @@ class SampleResponse:
     request_id: str
 
 
-def test_decode_dataclass_with_nested_request() -> None:
+# --- Tests for round-trip encode/decode ---
+
+
+def test_roundtrip_simple_dataclass() -> None:
+    """Test encoding and decoding of a simple dataclass."""
+    original = SampleRequest(request_id="test-123", prompt="test prompt")
+
+    encoded = encode_checkpoint_value(original)
+    decoded = cast(SampleRequest, decode_checkpoint_value(encoded))
+
+    assert isinstance(decoded, SampleRequest)
+    assert decoded.request_id == "test-123"
+    assert decoded.prompt == "test prompt"
+
+
+def test_roundtrip_dataclass_with_nested_request() -> None:
     """Test that dataclass with nested dataclass fields can be encoded and decoded correctly."""
     original = SampleResponse(
         data="approve",
@@ -49,45 +65,7 @@ def test_decode_dataclass_with_nested_request() -> None:
     assert decoded.original_request.request_id == "abc"
 
 
-def test_is_instance_of_coerces_nested_dataclass_dict() -> None:
-    """Test that is_instance_of can handle nested structures with dict conversion."""
-    response = SampleResponse(
-        data="approve",
-        original_request=SampleRequest(request_id="req-1", prompt="prompt"),
-        request_id="req-1",
-    )
-
-    # Simulate checkpoint decode fallback leaving a dict
-    response.original_request = cast(
-        Any,
-        {
-            "request_id": "req-1",
-            "prompt": "prompt",
-        },
-    )
-
-    assert is_instance_of(response, SampleResponse)
-    assert isinstance(response.original_request, dict)
-
-    # Verify the dict contains expected values
-    dict_request = cast(dict[str, Any], response.original_request)
-    assert dict_request["request_id"] == "req-1"
-    assert dict_request["prompt"] == "prompt"
-
-
-def test_encode_decode_simple_dataclass() -> None:
-    """Test encoding and decoding of a simple dataclass."""
-    original = SampleRequest(request_id="test-123", prompt="test prompt")
-
-    encoded = encode_checkpoint_value(original)
-    decoded = cast(SampleRequest, decode_checkpoint_value(encoded))
-
-    assert isinstance(decoded, SampleRequest)
-    assert decoded.request_id == "test-123"
-    assert decoded.prompt == "test prompt"
-
-
-def test_encode_decode_nested_structures() -> None:
+def test_roundtrip_nested_structures() -> None:
     """Test encoding and decoding of complex nested structures."""
     nested_data = {
         "requests": [
@@ -110,7 +88,6 @@ def test_encode_decode_nested_structures() -> None:
     assert "requests" in decoded
     assert "responses" in decoded
 
-    # Check the requests list
     requests = cast(list[Any], decoded["requests"])
     assert isinstance(requests, list)
     assert len(requests) == 2
@@ -120,7 +97,6 @@ def test_encode_decode_nested_structures() -> None:
     assert first_request.request_id == "req-1"
     assert second_request.request_id == "req-2"
 
-    # Check the responses dict
     responses = cast(dict[str, Any], decoded["responses"])
     assert isinstance(responses, dict)
     assert "req-1" in responses
@@ -131,108 +107,145 @@ def test_encode_decode_nested_structures() -> None:
     assert response.original_request.request_id == "req-1"
 
 
-def test_encode_allows_marker_key_without_value_key() -> None:
-    """Test that encoding a dict with only the marker key (no 'value') is allowed."""
-    dict_with_marker_only = {
-        MODEL_MARKER: "some.module:FakeClass",
-        "other_key": "test",
+def test_roundtrip_datetime() -> None:
+    """Test round-trip encoding/decoding of datetime objects."""
+    original = datetime(2024, 5, 4, 12, 30, 45, tzinfo=timezone.utc)
+
+    encoded = encode_checkpoint_value(original)
+    decoded = decode_checkpoint_value(encoded)
+
+    assert isinstance(decoded, datetime)
+    assert decoded == original
+
+
+def test_roundtrip_primitives() -> None:
+    """Test that primitive types round-trip unchanged."""
+    for value in ["hello", 42, 3.14, True, False, None]:
+        assert decode_checkpoint_value(encode_checkpoint_value(value)) == value
+
+
+def test_roundtrip_dict_with_mixed_values() -> None:
+    """Test round-trip of a dict containing both primitives and complex types."""
+    original = {
+        "name": "test",
+        "request": SampleRequest(request_id="r1", prompt="p1"),
+        "count": 5,
     }
-    encoded = encode_checkpoint_value(dict_with_marker_only)
-    assert MODEL_MARKER in encoded
-    assert "other_key" in encoded
+
+    encoded = encode_checkpoint_value(original)
+    decoded = decode_checkpoint_value(encoded)
+
+    assert decoded["name"] == "test"
+    assert decoded["count"] == 5
+    assert isinstance(decoded["request"], SampleRequest)
+    assert decoded["request"].request_id == "r1"
 
 
-def test_encode_allows_value_key_without_marker_key() -> None:
-    """Test that encoding a dict with only 'value' key (no marker) is allowed."""
-    dict_with_value_only = {
-        "value": {"data": "test"},
-        "other_key": "test",
-    }
-    encoded = encode_checkpoint_value(dict_with_value_only)
-    assert "value" in encoded
-    assert "other_key" in encoded
+# --- Tests for decode primitives ---
 
 
-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
+def test_decode_string() -> None:
+    """Test decoding a string passes through unchanged."""
+    assert decode_checkpoint_value("hello") == "hello"
 
 
-class NotADataclass:
+def test_decode_integer() -> None:
+    """Test decoding an integer passes through unchanged."""
+    assert decode_checkpoint_value(42) == 42
+
+
+def test_decode_none() -> None:
+    """Test decoding None passes through unchanged."""
+    assert decode_checkpoint_value(None) is None
+
+
+# --- Tests for decode collections ---
+
+
+def test_decode_plain_dict() -> None:
+    """Test decoding a plain dictionary with primitive values."""
+    data = {"a": 1, "b": "two"}
+    assert decode_checkpoint_value(data) == {"a": 1, "b": "two"}
+
+
+def test_decode_plain_list() -> None:
+    """Test decoding a plain list with primitive values."""
+    data = [1, "two", 3.0]
+    assert decode_checkpoint_value(data) == [1, "two", 3.0]
+
+
+# --- Tests for type verification ---
+
+
+def test_decode_raises_on_type_mismatch() -> None:
+    """Test that decoding raises CheckpointDecodingError when type doesn't match."""
+    # Encode a SampleRequest but tamper with the type marker
+    encoded = encode_checkpoint_value(SampleRequest(request_id="r1", prompt="p1"))
+    assert isinstance(encoded, dict)
+    encoded[_TYPE_MARKER] = "nonexistent.module:FakeClass"
+
+    with pytest.raises(CheckpointDecodingError, match="Type mismatch"):
+        decode_checkpoint_value(encoded)
+
+
+class NotADataclass:  # noqa: B903
     """A regular class that is not a dataclass."""
 
     def __init__(self, value: str) -> None:
         self.value = value
 
-    def get_value(self) -> str:
-        return self.value
+
+def test_roundtrip_regular_class() -> None:
+    """Test that regular (non-dataclass) objects can be round-tripped via pickle."""
+    original = NotADataclass(value="test_value")
+
+    encoded = encode_checkpoint_value(original)
+    decoded = cast(NotADataclass, decode_checkpoint_value(encoded))
+
+    assert isinstance(decoded, NotADataclass)
+    assert decoded.value == "test_value"
 
 
-class NotAModel:
-    """A regular class that does not support the model protocol."""
+def test_roundtrip_tuple() -> None:
+    """Test that tuples preserve their type through encode/decode roundtrip."""
+    original = (1, "two", 3.0)
 
-    def __init__(self, value: str) -> None:
-        self.value = value
+    encoded = encode_checkpoint_value(original)
+    decoded = decode_checkpoint_value(encoded)
 
-    def get_value(self) -> str:
-        return self.value
+    assert isinstance(decoded, tuple)
+    assert decoded == original
 
 
-def test_decode_rejects_non_dataclass_with_dataclass_marker() -> None:
-    """Test that decode returns raw value when marked class is not a dataclass."""
-    # Manually construct a payload that claims NotADataclass is a dataclass
-    fake_payload = {
-        DATACLASS_MARKER: f"{NotADataclass.__module__}:{NotADataclass.__name__}",
-        "value": {"value": "test_value"},
-    }
+def test_roundtrip_set() -> None:
+    """Test that sets preserve their type through encode/decode roundtrip."""
+    original = {1, 2, 3}
 
-    decoded = decode_checkpoint_value(fake_payload)
+    encoded = encode_checkpoint_value(original)
+    decoded = decode_checkpoint_value(encoded)
 
-    # Should return the raw decoded value, not an instance of NotADataclass
-    assert isinstance(decoded, dict)
-    assert decoded["value"] == "test_value"
+    assert isinstance(decoded, set)
+    assert decoded == original
 
 
-def test_decode_rejects_non_model_with_model_marker() -> None:
-    """Test that decode returns raw value when marked class doesn't support model protocol."""
-    # Manually construct a payload that claims NotAModel supports the model protocol
-    fake_payload = {
-        MODEL_MARKER: f"{NotAModel.__module__}:{NotAModel.__name__}",
-        "strategy": "to_dict",
-        "value": {"value": "test_value"},
-    }
+def test_roundtrip_nested_tuple_in_dict() -> None:
+    """Test that tuples nested inside dicts preserve their type."""
+    original = {"items": (1, 2, 3), "name": "test"}
 
-    decoded = decode_checkpoint_value(fake_payload)
+    encoded = encode_checkpoint_value(original)
+    decoded = decode_checkpoint_value(encoded)
 
-    # Should return the raw decoded value, not an instance of NotAModel
-    assert isinstance(decoded, dict)
-    assert decoded["value"] == "test_value"
+    assert isinstance(decoded["items"], tuple)
+    assert decoded["items"] == (1, 2, 3)
+    assert decoded["name"] == "test"
 
 
-def test_encode_allows_nested_dict_with_marker_keys() -> None:
-    """Test that encoding allows nested dicts containing marker patterns.
+def test_roundtrip_set_in_list() -> None:
+    """Test that sets nested inside lists preserve their type."""
+    original = [{"tags": {1, 2, 3}}]
 
-    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:SomeClass",
-            "value": {"data": "test"},
-        }
-    }
+    encoded = encode_checkpoint_value(original)
+    decoded = decode_checkpoint_value(encoded)
 
-    encoded = encode_checkpoint_value(nested_data)
-    assert "outer" in encoded
-    assert MODEL_MARKER in encoded["outer"]
+    assert isinstance(decoded[0]["tags"], set)
+    assert decoded[0]["tags"] == {1, 2, 3}

python/packages/core/tests/workflow/test_checkpoint_encode.py

@@ -1,12 +1,13 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+import json
 from dataclasses import dataclass
+from datetime import datetime, timezone
 from typing import Any
 
 from agent_framework._workflows._checkpoint_encoding import (
-    _CYCLE_SENTINEL,
-    DATACLASS_MARKER,
-    MODEL_MARKER,
+    _PICKLE_MARKER,
+    _TYPE_MARKER,
     encode_checkpoint_value,
 )
 
@@ -41,23 +42,6 @@ class ModelWithToDict:
         return cls(data=d["data"])
 
 
-class ModelWithToJson:
-    """A class that implements to_json/from_json protocol."""
-
-    def __init__(self, data: str) -> None:
-        self.data = data
-
-    def to_json(self) -> str:
-        return f'{{"data": "{self.data}"}}'
-
-    @classmethod
-    def from_json(cls, json_str: str) -> "ModelWithToJson":
-        import json
-
-        d = json.loads(json_str)
-        return cls(data=d["data"])
-
-
 class UnknownObject:
     """A class that doesn't support any serialization protocol."""
 
@@ -68,43 +52,37 @@ class UnknownObject:
         return f"UnknownObject({self.value})"
 
 
-# --- Tests for primitive encoding ---
+# --- Tests for primitive encoding (pass-through) ---
 
 
 def test_encode_string() -> None:
     """Test encoding a string value."""
-    result = encode_checkpoint_value("hello")
-    assert result == "hello"
+    assert encode_checkpoint_value("hello") == "hello"
 
 
 def test_encode_integer() -> None:
     """Test encoding an integer value."""
-    result = encode_checkpoint_value(42)
-    assert result == 42
+    assert encode_checkpoint_value(42) == 42
 
 
 def test_encode_float() -> None:
     """Test encoding a float value."""
-    result = encode_checkpoint_value(3.14)
-    assert result == 3.14
+    assert encode_checkpoint_value(3.14) == 3.14
 
 
 def test_encode_boolean_true() -> None:
     """Test encoding a True boolean value."""
-    result = encode_checkpoint_value(True)
-    assert result is True
+    assert encode_checkpoint_value(True) is True
 
 
 def test_encode_boolean_false() -> None:
     """Test encoding a False boolean value."""
-    result = encode_checkpoint_value(False)
-    assert result is False
+    assert encode_checkpoint_value(False) is False
 
 
 def test_encode_none() -> None:
     """Test encoding a None value."""
-    result = encode_checkpoint_value(None)
-    assert result is None
+    assert encode_checkpoint_value(None) is None
 
 
 # --- Tests for collection encoding ---
@@ -112,8 +90,7 @@ def test_encode_none() -> None:
 
 def test_encode_empty_dict() -> None:
     """Test encoding an empty dictionary."""
-    result = encode_checkpoint_value({})
-    assert result == {}
+    assert encode_checkpoint_value({}) == {}
 
 
 def test_encode_simple_dict() -> None:
@@ -132,8 +109,7 @@ def test_encode_dict_with_non_string_keys() -> None:
 
 def test_encode_empty_list() -> None:
     """Test encoding an empty list."""
-    result = encode_checkpoint_value([])
-    assert result == []
+    assert encode_checkpoint_value([]) == []
 
 
 def test_encode_simple_list() -> None:
@@ -144,29 +120,26 @@ def test_encode_simple_list() -> None:
 
 
 def test_encode_tuple() -> None:
-    """Test encoding a tuple (converted to list)."""
+    """Test encoding a tuple (pickled to preserve type)."""
     data = (1, 2, 3)
     result = encode_checkpoint_value(data)
-    assert result == [1, 2, 3]
+    assert isinstance(result, dict)
+    assert _PICKLE_MARKER in result
+    assert _TYPE_MARKER in result
 
 
 def test_encode_set() -> None:
-    """Test encoding a set (converted to list)."""
+    """Test encoding a set (pickled to preserve type)."""
     data = {1, 2, 3}
     result = encode_checkpoint_value(data)
-    assert isinstance(result, list)
-    assert sorted(result) == [1, 2, 3]
+    assert isinstance(result, dict)
+    assert _PICKLE_MARKER in result
+    assert _TYPE_MARKER in result
 
 
 def test_encode_nested_dict() -> None:
     """Test encoding a nested dictionary structure."""
-    data = {
-        "outer": {
-            "inner": {
-                "value": 42,
-            }
-        }
-    }
+    data = {"outer": {"inner": {"value": 42}}}
     result = encode_checkpoint_value(data)
     assert result == {"outer": {"inner": {"value": 42}}}
 
@@ -178,18 +151,18 @@ def test_encode_list_of_dicts() -> None:
     assert result == [{"a": 1}, {"b": 2}]
 
 
-# --- Tests for dataclass encoding ---
+# --- Tests for non-JSON-native types (pickled) ---
 
 
 def test_encode_simple_dataclass() -> None:
-    """Test encoding a simple dataclass."""
+    """Test encoding a simple dataclass produces a pickled entry."""
     obj = SimpleDataclass(name="test", value=42)
     result = encode_checkpoint_value(obj)
 
     assert isinstance(result, dict)
-    assert DATACLASS_MARKER in result
-    assert "value" in result
-    assert result["value"] == {"name": "test", "value": 42}
+    assert _PICKLE_MARKER in result
+    assert _TYPE_MARKER in result
+    assert isinstance(result[_PICKLE_MARKER], str)  # base64 string
 
 
 def test_encode_nested_dataclass() -> None:
@@ -199,12 +172,8 @@ def test_encode_nested_dataclass() -> None:
     result = encode_checkpoint_value(outer)
 
     assert isinstance(result, dict)
-    assert DATACLASS_MARKER in result
-    assert "value" in result
-
-    outer_value = result["value"]
-    assert outer_value["outer_name"] == "outer"
-    assert DATACLASS_MARKER in outer_value["inner"]
+    assert _PICKLE_MARKER in result
+    assert _TYPE_MARKER in result
 
 
 def test_encode_list_of_dataclasses() -> None:
@@ -218,7 +187,7 @@ def test_encode_list_of_dataclasses() -> None:
     assert isinstance(result, list)
     assert len(result) == 2
     for item in result:
-        assert DATACLASS_MARKER in item
+        assert _PICKLE_MARKER in item
 
 
 def test_encode_dict_with_dataclass_values() -> None:
@@ -230,169 +199,77 @@ def test_encode_dict_with_dataclass_values() -> None:
     result = encode_checkpoint_value(data)
 
     assert isinstance(result, dict)
-    assert DATACLASS_MARKER in result["item1"]
-    assert DATACLASS_MARKER in result["item2"]
-
-
-# --- Tests for model protocol encoding ---
+    assert _PICKLE_MARKER in result["item1"]
+    assert _PICKLE_MARKER in result["item2"]
 
 
 def test_encode_model_with_to_dict() -> None:
-    """Test encoding an object implementing to_dict/from_dict protocol."""
+    """Test encoding an object with to_dict is pickled (not using to_dict)."""
     obj = ModelWithToDict(data="test_data")
     result = encode_checkpoint_value(obj)
 
     assert isinstance(result, dict)
-    assert MODEL_MARKER in result
-    assert result["strategy"] == "to_dict"
-    assert result["value"] == {"data": "test_data"}
+    assert _PICKLE_MARKER in result
 
 
-def test_encode_model_with_to_json() -> None:
-    """Test encoding an object implementing to_json/from_json protocol."""
-    obj = ModelWithToJson(data="test_data")
-    result = encode_checkpoint_value(obj)
-
-    assert isinstance(result, dict)
-    assert MODEL_MARKER in result
-    assert result["strategy"] == "to_json"
-    assert '"data": "test_data"' in result["value"]
-
-
-# --- Tests for unknown object encoding ---
-
-
-def test_encode_unknown_object_fallback_to_string() -> None:
-    """Test that unknown objects are encoded as strings."""
+def test_encode_unknown_object() -> None:
+    """Test that arbitrary objects are pickled."""
     obj = UnknownObject(value="test")
     result = encode_checkpoint_value(obj)
 
-    assert isinstance(result, str)
-    assert "UnknownObject" in result
+    assert isinstance(result, dict)
+    assert _PICKLE_MARKER in result
 
 
-# --- Tests for cycle detection ---
+def test_encode_datetime() -> None:
+    """Test that datetime objects are pickled."""
+    dt = datetime(2024, 5, 4, 12, 30, 45, tzinfo=timezone.utc)
+    result = encode_checkpoint_value(dt)
+
+    assert isinstance(result, dict)
+    assert _PICKLE_MARKER in result
 
 
-def test_encode_dict_with_self_reference() -> None:
-    """Test that dict self-references are detected and handled."""
-    data: dict[str, Any] = {"name": "test"}
-    data["self"] = data  # Create circular reference
-
-    result = encode_checkpoint_value(data)
-    assert result["name"] == "test"
-    assert result["self"] == _CYCLE_SENTINEL
+# --- Tests for type marker ---
 
 
-def test_encode_list_with_self_reference() -> None:
-    """Test that list self-references are detected and handled."""
-    data: list[Any] = [1, 2]
-    data.append(data)  # Create circular reference
+def test_encode_type_marker_records_type_info() -> None:
+    """Test that encoded objects include correct type information."""
+    obj = SimpleDataclass(name="test", value=42)
+    result = encode_checkpoint_value(obj)
 
-    result = encode_checkpoint_value(data)
-    assert result[0] == 1
-    assert result[1] == 2
-    assert result[2] == _CYCLE_SENTINEL
+    type_key = result[_TYPE_MARKER]
+    assert "SimpleDataclass" in type_key
 
 
-# --- 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_type_marker_uses_module_qualname_format() -> None:
+    """Test that type marker uses module:qualname format."""
+    obj = SimpleDataclass(name="test", value=42)
+    result = encode_checkpoint_value(obj)
+
+    type_key = result[_TYPE_MARKER]
+    assert ":" in type_key
+    module, qualname = type_key.split(":")
+    assert module  # non-empty module
+    assert qualname == "SimpleDataclass"
 
 
-def test_encode_allows_dict_with_model_marker_and_value() -> None:
-    """Test that encoding a dict with MODEL_MARKER and 'value' is allowed.
+# --- Tests for JSON serializability ---
 
-    Security is enforced at deserialization time, not serialization time.
-    """
+
+def test_encode_result_is_json_serializable() -> None:
+    """Test that encoded output is fully JSON-serializable."""
     data = {
-        MODEL_MARKER: "some.module:SomeClass",
-        "value": {"data": "test"},
+        "dc": SimpleDataclass(name="test", value=42),
+        "model": ModelWithToDict(data="test"),
+        "dt": datetime.now(timezone.utc),
+        "nested": [SimpleDataclass(name="n", value=1)],
     }
-    result = encode_checkpoint_value(data)
-    assert MODEL_MARKER in result
-    assert "value" in result
-
-
-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"},
-    }
-    result = encode_checkpoint_value(data)
-    assert DATACLASS_MARKER in result
-    assert "value" in result
-
-
-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:SomeClass",
-            "value": {"data": "test"},
-        }
-    }
-    result = encode_checkpoint_value(nested_data)
-    assert "outer" in result
-    assert MODEL_MARKER in result["outer"]
-
-
-def test_encode_allows_marker_without_value() -> None:
-    """Test that a dict with marker key but without 'value' key is allowed."""
-    data = {
-        MODEL_MARKER: "some.module:SomeClass",
-        "other_key": "allowed",
-    }
-    result = encode_checkpoint_value(data)
-    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."""
-    data = {
-        "value": {"nested": "data"},
-        "other_key": "allowed",
-    }
-    result = encode_checkpoint_value(data)
-    assert "value" in result
-    assert result["other_key"] == "allowed"
-
-
-# --- Tests for max depth protection ---
-
-
-def test_encode_deep_nesting_triggers_max_depth() -> None:
-    """Test that very deep nesting triggers max depth protection."""
-    # Create a deeply nested structure (over 100 levels)
-    data: dict[str, Any] = {"level": 0}
-    current = data
-    for i in range(105):
-        current["nested"] = {"level": i + 1}
-        current = current["nested"]
 
     result = encode_checkpoint_value(data)
-
-    # Navigate to find the max_depth sentinel
-    current_result = result
-    found_max_depth = False
-    for _ in range(110):
-        if isinstance(current_result, dict) and "nested" in current_result:
-            current_result = current_result["nested"]
-            if current_result == "<max_depth>":
-                found_max_depth = True
-                break
-        else:
-            break
-
-    assert found_max_depth, "Expected <max_depth> sentinel to be found in deeply nested structure"
+    # Should not raise
+    json_str = json.dumps(result)
+    assert isinstance(json_str, str)
 
 
 # --- Tests for mixed complex structures ---
@@ -413,6 +290,7 @@ def test_encode_complex_mixed_structure() -> None:
 
     result = encode_checkpoint_value(data)
 
+    # Primitives and collections pass through
     assert result["string_value"] == "hello"
     assert result["int_value"] == 42
     assert result["float_value"] == 3.14
@@ -420,4 +298,17 @@ def test_encode_complex_mixed_structure() -> None:
     assert result["none_value"] is None
     assert result["list_value"] == [1, 2, 3]
     assert result["nested_dict"] == {"a": 1, "b": 2}
-    assert DATACLASS_MARKER in result["dataclass_value"]
+    # Dataclass is pickled
+    assert _PICKLE_MARKER in result["dataclass_value"]
+
+
+def test_encode_preserves_dict_with_pickle_marker_key() -> None:
+    """Test that regular dicts containing _PICKLE_MARKER key are recursively encoded."""
+    data = {
+        _PICKLE_MARKER: "some_value",
+        "other_key": "test",
+    }
+    result = encode_checkpoint_value(data)
+    assert _PICKLE_MARKER in result
+    assert result[_PICKLE_MARKER] == "some_value"
+    assert result["other_key"] == "test"

python/packages/core/tests/workflow/test_checkpoint_validation.py

@@ -44,7 +44,7 @@ async def test_resume_fails_when_graph_mismatch() -> None:
     # Run once to create checkpoints
     _ = [event async for event in workflow.run("hello", stream=True)]  # noqa: F841
 
-    checkpoints = await storage.list_checkpoints()
+    checkpoints = await storage.list_checkpoints(workflow_name=workflow.name)
     assert checkpoints, "expected at least one checkpoint to be created"
     target_checkpoint = checkpoints[-1]
 
@@ -67,7 +67,7 @@ async def test_resume_succeeds_when_graph_matches() -> None:
     workflow = build_workflow(storage, finish_id="finish")
     _ = [event async for event in workflow.run("hello", stream=True)]  # noqa: F841
 
-    checkpoints = sorted(await storage.list_checkpoints(), key=lambda c: c.timestamp)
+    checkpoints = sorted(await storage.list_checkpoints(workflow_name=workflow.name), key=lambda c: c.timestamp)
     target_checkpoint = checkpoints[0]
 
     resumed_workflow = build_workflow(storage, finish_id="finish")
@@ -126,7 +126,7 @@ async def test_resume_succeeds_when_sub_workflow_matches() -> None:
 
     _ = [event async for event in workflow.run("hello", stream=True)]
 
-    checkpoints = await storage.list_checkpoints()
+    checkpoints = await storage.list_checkpoints(workflow_name=workflow.name)
     assert checkpoints, "expected at least one checkpoint to be created"
     target_checkpoint = checkpoints[-1]
 
@@ -150,7 +150,7 @@ async def test_resume_fails_when_sub_workflow_changes() -> None:
 
     _ = [event async for event in workflow.run("hello", stream=True)]
 
-    checkpoints = await storage.list_checkpoints()
+    checkpoints = await storage.list_checkpoints(workflow_name=workflow.name)
     assert checkpoints, "expected at least one checkpoint to be created"
     target_checkpoint = checkpoints[-1]
 

python/packages/core/tests/workflow/test_request_info_and_response.py

@@ -3,7 +3,6 @@
 from dataclasses import dataclass
 
 from agent_framework import (
-    FileCheckpointStorage,
     WorkflowBuilder,
     WorkflowContext,
     WorkflowEvent,
@@ -323,90 +322,3 @@ class TestRequestInfoAndResponse:
         assert completed
         # Should not have any calculations performed due to invalid input
         assert len(executor.calculations_performed) == 0
-
-    async def test_checkpoint_with_pending_request_info_events(self):
-        """Test that request info events are properly serialized in checkpoints and can be restored."""
-        import tempfile
-
-        with tempfile.TemporaryDirectory() as temp_dir:
-            # Use file-based storage to test full serialization
-            storage = FileCheckpointStorage(temp_dir)
-
-            # Create workflow with checkpointing enabled
-            executor = ApprovalRequiredExecutor(id="approval_executor")
-            workflow = WorkflowBuilder(start_executor=executor, checkpoint_storage=storage).build()
-
-            # Step 1: Run workflow to completion to ensure checkpoints are created
-            request_info_event: WorkflowEvent | None = None
-            async for event in workflow.run("checkpoint test operation", stream=True):
-                if event.type == "request_info":
-                    request_info_event = event
-
-            # Verify request was emitted
-            assert request_info_event is not None
-            assert isinstance(request_info_event.data, UserApprovalRequest)
-            assert request_info_event.data.prompt == "Please approve the operation: checkpoint test operation"
-            assert request_info_event.source_executor_id == "approval_executor"
-
-            # Step 2: List checkpoints to find the one with our pending request
-            checkpoints = await storage.list_checkpoints()
-            assert len(checkpoints) > 0, "No checkpoints were created during workflow execution"
-
-            # Find the checkpoint with our pending request
-            checkpoint_with_request = None
-            for checkpoint in checkpoints:
-                if request_info_event.request_id in checkpoint.pending_request_info_events:
-                    checkpoint_with_request = checkpoint
-                    break
-
-            assert checkpoint_with_request is not None, "No checkpoint found with pending request info event"
-
-            # Step 3: Verify the pending request info event was properly serialized
-            serialized_event = checkpoint_with_request.pending_request_info_events[request_info_event.request_id]
-            assert "data" in serialized_event
-            assert "request_id" in serialized_event
-            assert "source_executor_id" in serialized_event
-            assert "request_type" in serialized_event
-            assert serialized_event["request_id"] == request_info_event.request_id
-            assert serialized_event["source_executor_id"] == "approval_executor"
-
-            # Step 4: Create a fresh workflow and restore from checkpoint
-            new_executor = ApprovalRequiredExecutor(id="approval_executor")
-            restored_workflow = WorkflowBuilder(start_executor=new_executor, checkpoint_storage=storage).build()
-
-            # Step 5: Resume from checkpoint and verify the request can be continued
-            completed = False
-            restored_request_event: WorkflowEvent | None = None
-            async for event in restored_workflow.run(checkpoint_id=checkpoint_with_request.checkpoint_id, stream=True):
-                # Should re-emit the pending request info event
-                if event.type == "request_info" and event.request_id == request_info_event.request_id:
-                    restored_request_event = event
-                elif event.type == "status" and event.state == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS:
-                    completed = True
-
-            assert completed, "Workflow should reach idle with pending requests state after restoration"
-            assert restored_request_event is not None, "Restored request info event should be emitted"
-
-            # Verify the restored event matches the original
-            assert restored_request_event.source_executor_id == request_info_event.source_executor_id
-            assert isinstance(restored_request_event.data, UserApprovalRequest)
-            assert restored_request_event.data.prompt == request_info_event.data.prompt
-            assert restored_request_event.data.context == request_info_event.data.context
-
-            # Step 6: Provide response to the restored request and complete the workflow
-            final_completed = False
-            async for event in restored_workflow.run(
-                stream=True,
-                responses={
-                    request_info_event.request_id: True  # Approve the request
-                },
-            ):
-                if event.type == "status" and event.state == WorkflowRunState.IDLE:
-                    final_completed = True
-
-            assert final_completed, "Workflow should complete after providing response to restored request"
-
-            # Step 7: Verify the executor state was properly restored and response was processed
-            assert new_executor.approval_received is True
-            expected_result = "Operation approved: Please approve the operation: checkpoint test operation"
-            assert new_executor.final_result == expected_result

python/packages/core/tests/workflow/test_request_info_event_rehydrate.py

@@ -4,14 +4,27 @@ import json
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
 
-import pytest
-
-from agent_framework import InMemoryCheckpointStorage, InProcRunnerContext
-from agent_framework._workflows._checkpoint_encoding import DATACLASS_MARKER, encode_checkpoint_value
-from agent_framework._workflows._checkpoint_summary import get_checkpoint_summary
+from agent_framework import (
+    FileCheckpointStorage,
+    InMemoryCheckpointStorage,
+    InProcRunnerContext,
+    WorkflowBuilder,
+    WorkflowRunState,
+)
+from agent_framework._workflows._checkpoint_encoding import (
+    _PICKLE_MARKER,
+    encode_checkpoint_value,
+)
 from agent_framework._workflows._events import WorkflowEvent
 from agent_framework._workflows._state import State
 
+from .test_request_info_and_response import (
+    ApprovalRequiredExecutor,
+    CalculationRequest,
+    MultiRequestExecutor,
+    UserApprovalRequest,
+)
+
 
 @dataclass
 class MockRequest: ...
@@ -46,13 +59,13 @@ async def test_rehydrate_request_info_event() -> None:
     runner_context = InProcRunnerContext(InMemoryCheckpointStorage())
     await runner_context.add_request_info_event(request_info_event)
 
-    checkpoint_id = await runner_context.create_checkpoint(State(), iteration_count=1)
+    checkpoint_id = await runner_context.create_checkpoint("test_name", "test_hash", State(), None, iteration_count=1)
     checkpoint = await runner_context.load_checkpoint(checkpoint_id)
 
     assert checkpoint is not None
     assert checkpoint.pending_request_info_events
     assert "request-123" in checkpoint.pending_request_info_events
-    assert "request_type" in checkpoint.pending_request_info_events["request-123"]
+    assert checkpoint.pending_request_info_events["request-123"].request_type is MockRequest
 
     # Rehydrate the context
     await runner_context.apply_checkpoint(checkpoint)
@@ -67,97 +80,6 @@ async def test_rehydrate_request_info_event() -> None:
     assert isinstance(rehydrated_event.data, MockRequest)
 
 
-async def test_rehydrate_fails_when_request_type_missing() -> None:
-    """Rehydration should fail is the request type is missing or fails to import."""
-    request_info_event = WorkflowEvent.request_info(
-        request_id="request-123",
-        source_executor_id="review_gateway",
-        request_data=MockRequest(),
-        response_type=bool,
-    )
-
-    runner_context = InProcRunnerContext(InMemoryCheckpointStorage())
-    await runner_context.add_request_info_event(request_info_event)
-
-    checkpoint_id = await runner_context.create_checkpoint(State(), iteration_count=1)
-    checkpoint = await runner_context.load_checkpoint(checkpoint_id)
-
-    assert checkpoint is not None
-    assert checkpoint.pending_request_info_events
-    assert "request-123" in checkpoint.pending_request_info_events
-    assert "request_type" in checkpoint.pending_request_info_events["request-123"]
-
-    # Modify the checkpoint to simulate missing request type
-    checkpoint.pending_request_info_events["request-123"]["request_type"] = "nonexistent.module:MissingRequest"
-
-    # Rehydrate the context
-    with pytest.raises(ImportError):
-        await runner_context.apply_checkpoint(checkpoint)
-
-
-async def test_rehydrate_fails_when_request_type_mismatch() -> None:
-    """Rehydration should fail if the request type is mismatched."""
-    request_info_event = WorkflowEvent.request_info(
-        request_id="request-123",
-        source_executor_id="review_gateway",
-        request_data=MockRequest(),
-        response_type=bool,
-    )
-
-    runner_context = InProcRunnerContext(InMemoryCheckpointStorage())
-    await runner_context.add_request_info_event(request_info_event)
-
-    checkpoint_id = await runner_context.create_checkpoint(State(), iteration_count=1)
-    checkpoint = await runner_context.load_checkpoint(checkpoint_id)
-
-    assert checkpoint is not None
-    assert checkpoint.pending_request_info_events
-    assert "request-123" in checkpoint.pending_request_info_events
-    assert "request_type" in checkpoint.pending_request_info_events["request-123"]
-
-    # Modify the checkpoint to simulate mismatched request type in the serialized data
-    checkpoint.pending_request_info_events["request-123"]["data"][DATACLASS_MARKER] = (
-        "nonexistent.module:MissingRequest"
-    )
-
-    # Rehydrate the context
-    with pytest.raises(TypeError):
-        await runner_context.apply_checkpoint(checkpoint)
-
-
-async def test_pending_requests_in_summary() -> None:
-    """Test that pending requests are correctly summarized in the checkpoint summary."""
-    request_info_event = WorkflowEvent.request_info(
-        request_id="request-123",
-        source_executor_id="review_gateway",
-        request_data=MockRequest(),
-        response_type=bool,
-    )
-
-    runner_context = InProcRunnerContext(InMemoryCheckpointStorage())
-    await runner_context.add_request_info_event(request_info_event)
-
-    checkpoint_id = await runner_context.create_checkpoint(State(), iteration_count=1)
-    checkpoint = await runner_context.load_checkpoint(checkpoint_id)
-
-    assert checkpoint is not None
-    summary = get_checkpoint_summary(checkpoint)
-
-    assert summary.checkpoint_id == checkpoint_id
-    assert summary.status == "awaiting request response"
-
-    assert len(summary.pending_request_info_events) == 1
-    pending_event = summary.pending_request_info_events[0]
-    assert isinstance(pending_event, WorkflowEvent)
-    assert pending_event.type == "request_info"
-    assert pending_event.request_id == "request-123"
-
-    assert pending_event.source_executor_id == "review_gateway"
-    assert pending_event.request_type is MockRequest
-    assert pending_event.response_type is bool
-    assert isinstance(pending_event.data, MockRequest)
-
-
 async def test_request_info_event_serializes_non_json_payloads() -> None:
     req_1 = WorkflowEvent.request_info(
         request_id="req-1",
@@ -176,20 +98,260 @@ async def test_request_info_event_serializes_non_json_payloads() -> None:
     await runner_context.add_request_info_event(req_1)
     await runner_context.add_request_info_event(req_2)
 
-    checkpoint_id = await runner_context.create_checkpoint(State(), iteration_count=1)
+    checkpoint_id = await runner_context.create_checkpoint("test_name", "test_hash", State(), None, iteration_count=1)
     checkpoint = await runner_context.load_checkpoint(checkpoint_id)
 
     # Should be JSON serializable despite datetime/slots
     serialized = json.dumps(encode_checkpoint_value(checkpoint))
+    assert isinstance(serialized, str)
+
+    # Verify the structure contains pickled data for the request data fields
     deserialized = json.loads(serialized)
+    assert _PICKLE_MARKER in deserialized  # checkpoint itself is pickled
 
-    assert "value" in deserialized
-    deserialized = deserialized["value"]
+    # Verify we can rehydrate the checkpoint correctly
+    await runner_context.apply_checkpoint(checkpoint)
+    pending = await runner_context.get_pending_request_info_events()
 
-    assert "pending_request_info_events" in deserialized
-    pending_request_info_events = deserialized["pending_request_info_events"]
-    assert "req-1" in pending_request_info_events
-    assert isinstance(pending_request_info_events["req-1"]["data"]["value"]["issued_at"], str)
+    assert "req-1" in pending
+    rehydrated_1 = pending["req-1"]
+    assert isinstance(rehydrated_1.data, TimedApproval)
+    assert rehydrated_1.data.issued_at == datetime(2024, 5, 4, 12, 30, 45)
 
-    assert "req-2" in pending_request_info_events
-    assert pending_request_info_events["req-2"]["data"]["value"]["note"] == "slot-based"
+    assert "req-2" in pending
+    rehydrated_2 = pending["req-2"]
+    assert isinstance(rehydrated_2.data, SlottedApproval)
+    assert rehydrated_2.data.note == "slot-based"
+
+
+async def test_checkpoint_with_pending_request_info_events():
+    """Test that request info events are properly serialized in checkpoints and can be restored."""
+    import tempfile
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Use file-based storage to test full serialization
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Create workflow with checkpointing enabled
+        executor = ApprovalRequiredExecutor(id="approval_executor")
+        workflow = WorkflowBuilder(start_executor=executor, checkpoint_storage=storage).build()
+
+        # Step 1: Run workflow to completion to ensure checkpoints are created
+        request_info_event: WorkflowEvent | None = None
+        async for event in workflow.run("checkpoint test operation", stream=True):
+            if event.type == "request_info":
+                request_info_event = event
+
+        # Verify request was emitted
+        assert request_info_event is not None
+        assert isinstance(request_info_event.data, UserApprovalRequest)
+        assert request_info_event.data.prompt == "Please approve the operation: checkpoint test operation"
+        assert request_info_event.source_executor_id == "approval_executor"
+
+        # Step 2: List checkpoints to find the one with our pending request
+        checkpoints = await storage.list_checkpoints(workflow_name=workflow.name)
+        assert len(checkpoints) > 0, "No checkpoints were created during workflow execution"
+
+        # Find the checkpoint with our pending request
+        checkpoint_with_request = None
+        for checkpoint in checkpoints:
+            if request_info_event.request_id in checkpoint.pending_request_info_events:
+                checkpoint_with_request = checkpoint
+                break
+
+        assert checkpoint_with_request is not None, "No checkpoint found with pending request info event"
+
+        # Step 3: Verify the pending request info event was properly serialized
+        serialized_event = checkpoint_with_request.pending_request_info_events[request_info_event.request_id]
+        assert serialized_event.data
+        assert serialized_event.request_type is UserApprovalRequest
+        assert serialized_event.request_id == request_info_event.request_id
+        assert serialized_event.source_executor_id == "approval_executor"
+
+        # Step 4: Create a fresh workflow and restore from checkpoint
+        new_executor = ApprovalRequiredExecutor(id="approval_executor")
+        restored_workflow = WorkflowBuilder(start_executor=new_executor, checkpoint_storage=storage).build()
+
+        # Step 5: Resume from checkpoint and verify the request can be continued
+        completed = False
+        restored_request_event: WorkflowEvent | None = None
+        async for event in restored_workflow.run(checkpoint_id=checkpoint_with_request.checkpoint_id, stream=True):
+            # Should re-emit the pending request info event
+            if event.type == "request_info" and event.request_id == request_info_event.request_id:
+                restored_request_event = event
+            elif event.type == "status" and event.state == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS:
+                completed = True
+
+        assert completed, "Workflow should reach idle with pending requests state after restoration"
+        assert restored_request_event is not None, "Restored request info event should be emitted"
+
+        # Verify the restored event matches the original
+        assert restored_request_event.source_executor_id == request_info_event.source_executor_id
+        assert isinstance(restored_request_event.data, UserApprovalRequest)
+        assert restored_request_event.data.prompt == request_info_event.data.prompt
+        assert restored_request_event.data.context == request_info_event.data.context
+
+        # Step 6: Provide response to the restored request and complete the workflow
+        final_completed = False
+        async for event in restored_workflow.run(
+            stream=True,
+            responses={
+                request_info_event.request_id: True  # Approve the request
+            },
+        ):
+            if event.type == "status" and event.state == WorkflowRunState.IDLE:
+                final_completed = True
+
+        assert final_completed, "Workflow should complete after providing response to restored request"
+
+        # Step 7: Verify the executor state was properly restored and response was processed
+        assert new_executor.approval_received is True
+        expected_result = "Operation approved: Please approve the operation: checkpoint test operation"
+        assert new_executor.final_result == expected_result
+
+
+async def test_checkpoint_restore_with_responses_does_not_reemit_handled_requests():
+    """Test that request_info events are not re-emitted when responses are provided with checkpoint restore.
+
+    When calling run(checkpoint_id=..., responses=...), the workflow restores from a checkpoint
+    that contains pending request_info events. Because responses are provided for those events,
+    they should NOT be re-emitted in the event stream - they are considered "handled".
+
+    Note: The workflow's internal state tracking still sees the request_info events (before filtering),
+    so the final status may be IDLE_WITH_PENDING_REQUESTS even though the requests were handled.
+    The key behavior we're testing is that the CALLER doesn't see the request_info events.
+    """
+    import tempfile
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Use file-based storage to test full serialization
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Create workflow with checkpointing enabled
+        executor = ApprovalRequiredExecutor(id="approval_executor")
+        workflow = WorkflowBuilder(start_executor=executor, checkpoint_storage=storage).build()
+
+        # Step 1: Run workflow until it emits a request_info event
+        request_info_event: WorkflowEvent | None = None
+        async for event in workflow.run("test pending request suppression", stream=True):
+            if event.type == "request_info":
+                request_info_event = event
+
+        assert request_info_event is not None
+        request_id = request_info_event.request_id
+
+        # Step 2: Find the checkpoint with the pending request
+        checkpoints = await storage.list_checkpoints(workflow_name=workflow.name)
+        checkpoint_with_request = None
+        for checkpoint in checkpoints:
+            if request_id in checkpoint.pending_request_info_events:
+                checkpoint_with_request = checkpoint
+                break
+
+        assert checkpoint_with_request is not None
+
+        # Step 3: Create a fresh workflow and restore from checkpoint WITH responses in one call
+        new_executor = ApprovalRequiredExecutor(id="approval_executor")
+        restored_workflow = WorkflowBuilder(start_executor=new_executor, checkpoint_storage=storage).build()
+
+        # Track all emitted events
+        emitted_events: list[WorkflowEvent] = []
+        async for event in restored_workflow.run(
+            checkpoint_id=checkpoint_with_request.checkpoint_id,
+            responses={request_id: True},  # Provide response for the pending request
+            stream=True,
+        ):
+            emitted_events.append(event)
+
+        # Step 4: Verify the request_info event was NOT re-emitted to the caller
+        reemitted_request_info_events = [
+            e for e in emitted_events if e.type == "request_info" and e.request_id == request_id
+        ]
+        assert len(reemitted_request_info_events) == 0, (
+            f"request_info event should NOT be re-emitted when response is provided. "
+            f"Found {len(reemitted_request_info_events)} request_info events with request_id={request_id}"
+        )
+
+        # Step 5: Verify the response was processed by checking executor state
+        assert new_executor.approval_received is True, "Response should have been processed by the executor"
+        assert new_executor.final_result == (
+            "Operation approved: Please approve the operation: test pending request suppression"
+        )
+
+
+async def test_checkpoint_restore_with_partial_responses_reemits_unhandled_requests():
+    """Test that only unhandled request_info events are re-emitted when partial responses are provided.
+
+    When calling run(checkpoint_id=..., responses=...) with responses for only some of the
+    pending requests, only the unhandled request_info events should be re-emitted.
+    """
+    import tempfile
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Create workflow with multiple requests
+        executor = MultiRequestExecutor(id="multi_executor")
+        workflow = WorkflowBuilder(start_executor=executor, checkpoint_storage=storage).build()
+
+        # Step 1: Run workflow until it emits multiple request_info events
+        request_events: list[WorkflowEvent] = []
+        async for event in workflow.run("start batch", stream=True):
+            if event.type == "request_info":
+                request_events.append(event)
+
+        assert len(request_events) == 2
+
+        # Find the approval and calculation requests
+        approval_event = next((e for e in request_events if isinstance(e.data, UserApprovalRequest)), None)
+        calc_event = next((e for e in request_events if isinstance(e.data, CalculationRequest)), None)
+        assert approval_event is not None
+        assert calc_event is not None
+
+        # Step 2: Find the checkpoint with pending requests
+        checkpoints = await storage.list_checkpoints(workflow_name=workflow.name)
+        checkpoint_with_requests = None
+        for checkpoint in checkpoints:
+            has_approval = approval_event.request_id in checkpoint.pending_request_info_events
+            has_calc = calc_event.request_id in checkpoint.pending_request_info_events
+            if has_approval and has_calc:
+                checkpoint_with_requests = checkpoint
+                break
+
+        assert checkpoint_with_requests is not None
+
+        # Step 3: Restore from checkpoint with ONLY the approval response (not the calculation)
+        new_executor = MultiRequestExecutor(id="multi_executor")
+        restored_workflow = WorkflowBuilder(start_executor=new_executor, checkpoint_storage=storage).build()
+
+        emitted_events: list[WorkflowEvent] = []
+        async for event in restored_workflow.run(
+            checkpoint_id=checkpoint_with_requests.checkpoint_id,
+            responses={approval_event.request_id: True},  # Only respond to approval
+            stream=True,
+        ):
+            emitted_events.append(event)
+
+        # Step 4: Verify the approval request_info was NOT re-emitted
+        reemitted_approval_events = [
+            e for e in emitted_events if e.type == "request_info" and e.request_id == approval_event.request_id
+        ]
+        assert len(reemitted_approval_events) == 0, (
+            "Approval request_info should NOT be re-emitted since response was provided"
+        )
+
+        # Step 5: Verify the calculation request_info WAS re-emitted (no response provided)
+        reemitted_calc_events = [
+            e for e in emitted_events if e.type == "request_info" and e.request_id == calc_event.request_id
+        ]
+        assert len(reemitted_calc_events) == 1, (
+            "Calculation request_info SHOULD be re-emitted since no response was provided"
+        )
+
+        # Step 6: Verify workflow is in IDLE_WITH_PENDING_REQUESTS state (calc still pending)
+        status_events = [e for e in emitted_events if e.type == "status"]
+        final_status = status_events[-1] if status_events else None
+        assert final_status is not None
+        assert final_status.state == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS, (
+            f"Workflow should be IDLE_WITH_PENDING_REQUESTS, got {final_status.state}"
+        )

python/packages/core/tests/workflow/test_runner.py

@@ -2,6 +2,7 @@
 
 import asyncio
 from dataclasses import dataclass
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 
@@ -9,6 +10,9 @@ from agent_framework import (
     AgentExecutorResponse,
     AgentResponse,
     Executor,
+    InMemoryCheckpointStorage,
+    WorkflowCheckpoint,
+    WorkflowCheckpointException,
     WorkflowContext,
     WorkflowConvergenceException,
     WorkflowEvent,
@@ -16,6 +20,7 @@ from agent_framework import (
     WorkflowRunState,
     handler,
 )
+from agent_framework._workflows._const import EXECUTOR_STATE_KEY
 from agent_framework._workflows._edge import SingleEdgeGroup
 from agent_framework._workflows._runner import Runner
 from agent_framework._workflows._runner_context import (
@@ -61,7 +66,14 @@ def test_create_runner():
         executor_b.id: executor_b,
     }
 
-    runner = Runner(edge_groups, executors, state=State(), ctx=InProcRunnerContext())
+    runner = Runner(
+        edge_groups,
+        executors,
+        state=State(),
+        ctx=InProcRunnerContext(),
+        workflow_name="test_name",
+        graph_signature_hash="test_hash",
+    )
 
     assert runner.context is not None and isinstance(runner.context, RunnerContext)
 
@@ -84,7 +96,7 @@ async def test_runner_run_until_convergence():
     state = State()
     ctx = InProcRunnerContext()
 
-    runner = Runner(edges, executors, state, ctx)
+    runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash")
 
     result: int | None = None
     await executor_a.execute(
@@ -122,7 +134,7 @@ async def test_runner_run_until_convergence_not_completed():
     state = State()
     ctx = InProcRunnerContext()
 
-    runner = Runner(edges, executors, state, ctx, max_iterations=5)
+    runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash", max_iterations=5)
 
     await executor_a.execute(
         MockMessage(data=0),
@@ -156,7 +168,7 @@ async def test_runner_already_running():
     state = State()
     ctx = InProcRunnerContext()
 
-    runner = Runner(edges, executors, state, ctx)
+    runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash")
 
     await executor_a.execute(
         MockMessage(data=0),
@@ -176,7 +188,7 @@ async def test_runner_already_running():
 
 async def test_runner_emits_runner_completion_for_agent_response_without_targets():
     ctx = InProcRunnerContext()
-    runner = Runner([], {}, State(), ctx)
+    runner = Runner([], {}, State(), ctx, "test_name", graph_signature_hash="test_hash")
 
     await ctx.send_message(
         WorkflowMessage(
@@ -228,7 +240,7 @@ async def test_runner_cancellation_stops_active_executor():
     shared_state = State()
     ctx = InProcRunnerContext()
 
-    runner = Runner(edges, executors, shared_state, ctx)
+    runner = Runner(edges, executors, shared_state, ctx, "test_name", graph_signature_hash="test_hash")
 
     await executor_a.execute(
         MockMessage(data=0),
@@ -259,3 +271,579 @@ async def test_runner_cancellation_stops_active_executor():
     assert executor_a.completed_count == 1
     assert executor_b.started_count == 1
     assert executor_b.completed_count == 0  # Should NOT have completed due to cancellation
+
+
+class FailingExecutor(Executor):
+    """An executor that fails during execution."""
+
+    def __init__(self, id: str, fail_on_data: int = 5):
+        super().__init__(id=id)
+        self.fail_on_data = fail_on_data
+
+    @handler
+    async def handle(self, message: MockMessage, ctx: WorkflowContext[MockMessage, int]) -> None:
+        if message.data == self.fail_on_data:
+            raise RuntimeError("Simulated executor failure")
+        await ctx.send_message(MockMessage(data=message.data + 1))
+
+
+async def test_runner_iteration_exception_drains_events():
+    """Test that when an executor raises an exception, events are drained before propagating."""
+    executor_a = FailingExecutor(id="executor_a", fail_on_data=2)
+    executor_b = MockExecutor(id="executor_b")
+
+    edges = [
+        SingleEdgeGroup(executor_a.id, executor_b.id),
+        SingleEdgeGroup(executor_b.id, executor_a.id),
+    ]
+
+    executors: dict[str, Executor] = {
+        executor_a.id: executor_a,
+        executor_b.id: executor_b,
+    }
+    state = State()
+    ctx = InProcRunnerContext()
+
+    runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash")
+
+    await executor_a.execute(
+        MockMessage(data=0),
+        ["START"],
+        state,
+        ctx,
+    )
+
+    events: list[WorkflowEvent] = []
+    with pytest.raises(RuntimeError, match="Simulated executor failure"):
+        async for event in runner.run_until_convergence():
+            events.append(event)
+
+    # There should be some events emitted before the failure
+    assert len(events) > 0
+
+
+async def test_runner_reset_iteration_count():
+    """Test that reset_iteration_count works correctly."""
+    executor_a = MockExecutor(id="executor_a")
+    state = State()
+    ctx = InProcRunnerContext()
+
+    runner = Runner([], {executor_a.id: executor_a}, state, ctx, "test_name", graph_signature_hash="test_hash")
+    runner._iteration = 10
+
+    runner.reset_iteration_count()
+
+    assert runner._iteration == 0
+
+
+class CheckpointingContext(InProcRunnerContext):
+    """A context that supports checkpointing for testing."""
+
+    def __init__(self, storage: InMemoryCheckpointStorage | None = None):
+        super().__init__()
+        self._storage = storage or InMemoryCheckpointStorage()
+        self._checkpointing_enabled = True
+
+    def has_checkpointing(self) -> bool:
+        return self._checkpointing_enabled
+
+    async def create_checkpoint(
+        self,
+        workflow_name: str,
+        graph_signature_hash: str,
+        state: State,
+        previous_checkpoint_id: str | None,
+        iteration: int,
+    ) -> str:
+        checkpoint = WorkflowCheckpoint(
+            workflow_name=workflow_name,
+            graph_signature_hash=graph_signature_hash,
+            state=state.export(),
+            previous_checkpoint_id=previous_checkpoint_id,
+            iteration_count=iteration,
+        )
+        return await self._storage.save(checkpoint)
+
+    async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
+        try:
+            return await self._storage.load(checkpoint_id)
+        except WorkflowCheckpointException:
+            return None
+
+    async def apply_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None:
+        # Restore messages from checkpoint
+        for source_id, messages in checkpoint.messages.items():
+            for msg_data in messages:
+                await self.send_message(WorkflowMessage(data=msg_data, source_id=source_id))
+
+
+class FailingCheckpointContext(InProcRunnerContext):
+    """A context that fails during checkpoint creation."""
+
+    def has_checkpointing(self) -> bool:
+        return True
+
+    async def create_checkpoint(
+        self,
+        workflow_name: str,
+        graph_signature_hash: str,
+        state: State,
+        previous_checkpoint_id: str | None,
+        iteration: int,
+    ) -> str:
+        raise RuntimeError("Simulated checkpoint failure")
+
+
+async def test_runner_checkpoint_creation_failure():
+    """Test that checkpoint creation failure is handled gracefully."""
+    executor_a = MockExecutor(id="executor_a")
+    executor_b = MockExecutor(id="executor_b")
+
+    edges = [
+        SingleEdgeGroup(executor_a.id, executor_b.id),
+        SingleEdgeGroup(executor_b.id, executor_a.id),
+    ]
+
+    executors: dict[str, Executor] = {
+        executor_a.id: executor_a,
+        executor_b.id: executor_b,
+    }
+    state = State()
+    ctx = FailingCheckpointContext()
+
+    runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash")
+
+    await executor_a.execute(
+        MockMessage(data=0),
+        ["START"],
+        state,
+        ctx,
+    )
+
+    # Should complete without raising, even though checkpointing fails
+    result: int | None = None
+    async for event in runner.run_until_convergence():
+        if event.type == "output":
+            result = event.data
+
+    assert result == 10
+
+
+async def test_runner_restore_from_checkpoint_with_external_storage():
+    """Test restoring from checkpoint using external storage when context has no checkpointing."""
+    executor_a = MockExecutor(id="executor_a")
+    executor_b = MockExecutor(id="executor_b")
+
+    edges = [
+        SingleEdgeGroup(executor_a.id, executor_b.id),
+        SingleEdgeGroup(executor_b.id, executor_a.id),
+    ]
+
+    executors: dict[str, Executor] = {
+        executor_a.id: executor_a,
+        executor_b.id: executor_b,
+    }
+    state = State()
+    ctx = InProcRunnerContext()  # No checkpointing enabled
+
+    runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash")
+
+    # Create a checkpoint manually
+    storage = InMemoryCheckpointStorage()
+    checkpoint = WorkflowCheckpoint(
+        workflow_name="test_name",
+        graph_signature_hash="test_hash",
+        state={"test_key": "test_value"},
+        iteration_count=5,
+    )
+    checkpoint_id = await storage.save(checkpoint)
+
+    # Restore using external storage
+    await runner.restore_from_checkpoint(checkpoint_id, checkpoint_storage=storage)
+
+    assert runner._resumed_from_checkpoint is True
+    assert runner._iteration == 5
+    assert state.get("test_key") == "test_value"
+
+
+async def test_runner_restore_from_checkpoint_no_storage():
+    """Test that restore fails when no checkpointing and no external storage."""
+    state = State()
+    ctx = InProcRunnerContext()
+
+    runner = Runner([], {}, state, ctx, "test_name", graph_signature_hash="test_hash")
+
+    with pytest.raises(WorkflowCheckpointException, match="Cannot load checkpoint"):
+        await runner.restore_from_checkpoint("nonexistent-id")
+
+
+async def test_runner_restore_from_checkpoint_not_found():
+    """Test that restore fails when checkpoint is not found."""
+    storage = InMemoryCheckpointStorage()
+    ctx = CheckpointingContext(storage)
+    state = State()
+
+    runner = Runner([], {}, state, ctx, "test_name", graph_signature_hash="test_hash")
+
+    with pytest.raises(WorkflowCheckpointException, match="not found"):
+        await runner.restore_from_checkpoint("nonexistent-id")
+
+
+async def test_runner_restore_from_checkpoint_graph_hash_mismatch():
+    """Test that restore fails when graph hash doesn't match."""
+    storage = InMemoryCheckpointStorage()
+    ctx = CheckpointingContext(storage)
+    state = State()
+
+    runner = Runner([], {}, state, ctx, "test_name", graph_signature_hash="current_hash")
+
+    # Create a checkpoint with a different graph hash
+    checkpoint = WorkflowCheckpoint(
+        workflow_name="test_name",
+        graph_signature_hash="different_hash",
+        state={},
+        iteration_count=5,
+    )
+    checkpoint_id = await storage.save(checkpoint)
+
+    with pytest.raises(WorkflowCheckpointException, match="Workflow graph has changed"):
+        await runner.restore_from_checkpoint(checkpoint_id)
+
+
+async def test_runner_restore_from_checkpoint_generic_exception():
+    """Test that generic exceptions during restore are wrapped in WorkflowCheckpointException."""
+    state = State()
+
+    # Create a mock context that raises a generic exception
+    mock_ctx = MagicMock(spec=InProcRunnerContext)
+    mock_ctx.has_checkpointing.return_value = True
+    mock_ctx.load_checkpoint = AsyncMock(side_effect=ValueError("Unexpected error"))
+
+    runner = Runner([], {}, state, mock_ctx, "test_name", graph_signature_hash="test_hash")
+
+    with pytest.raises(WorkflowCheckpointException, match="Failed to restore from checkpoint"):
+        await runner.restore_from_checkpoint("some-id")
+
+
+async def test_runner_restore_executor_states_invalid_states_type():
+    """Test that restore fails when executor states is not a dict."""
+    executor_a = MockExecutor(id="executor_a")
+    state = State()
+    state.set(EXECUTOR_STATE_KEY, "not_a_dict")
+    state.commit()
+
+    ctx = InProcRunnerContext()
+    runner = Runner([], {executor_a.id: executor_a}, state, ctx, "test_name", graph_signature_hash="test_hash")
+
+    with pytest.raises(WorkflowCheckpointException, match="not a dictionary"):
+        await runner._restore_executor_states()
+
+
+async def test_runner_restore_executor_states_invalid_executor_id_type():
+    """Test that restore fails when executor ID is not a string."""
+    executor_a = MockExecutor(id="executor_a")
+    state = State()
+    state.set(EXECUTOR_STATE_KEY, {123: {"key": "value"}})  # Non-string key
+    state.commit()
+
+    ctx = InProcRunnerContext()
+    runner = Runner([], {executor_a.id: executor_a}, state, ctx, "test_name", graph_signature_hash="test_hash")
+
+    with pytest.raises(WorkflowCheckpointException, match="not a string"):
+        await runner._restore_executor_states()
+
+
+async def test_runner_restore_executor_states_invalid_state_type():
+    """Test that restore fails when executor state is not a dict[str, Any]."""
+    executor_a = MockExecutor(id="executor_a")
+    state = State()
+    state.set(EXECUTOR_STATE_KEY, {"executor_a": "not_a_dict"})
+    state.commit()
+
+    ctx = InProcRunnerContext()
+    runner = Runner([], {executor_a.id: executor_a}, state, ctx, "test_name", graph_signature_hash="test_hash")
+
+    with pytest.raises(WorkflowCheckpointException, match="not a dict"):
+        await runner._restore_executor_states()
+
+
+async def test_runner_restore_executor_states_invalid_state_keys():
+    """Test that restore fails when executor state dict has non-string keys."""
+    executor_a = MockExecutor(id="executor_a")
+    state = State()
+    state.set(EXECUTOR_STATE_KEY, {"executor_a": {123: "value"}})  # Non-string key in state
+    state.commit()
+
+    ctx = InProcRunnerContext()
+    runner = Runner([], {executor_a.id: executor_a}, state, ctx, "test_name", graph_signature_hash="test_hash")
+
+    with pytest.raises(WorkflowCheckpointException, match="not a dict"):
+        await runner._restore_executor_states()
+
+
+async def test_runner_restore_executor_states_missing_executor():
+    """Test that restore fails when executor is not found."""
+    state = State()
+    state.set(EXECUTOR_STATE_KEY, {"missing_executor": {"key": "value"}})
+    state.commit()
+
+    ctx = InProcRunnerContext()
+    runner = Runner([], {}, state, ctx, "test_name", graph_signature_hash="test_hash")
+
+    with pytest.raises(WorkflowCheckpointException, match="not found during state restoration"):
+        await runner._restore_executor_states()
+
+
+async def test_runner_set_executor_state_invalid_existing_states():
+    """Test that _set_executor_state fails when existing states is not a dict."""
+    executor_a = MockExecutor(id="executor_a")
+    state = State()
+    state.set(EXECUTOR_STATE_KEY, "not_a_dict")
+
+    ctx = InProcRunnerContext()
+    runner = Runner([], {executor_a.id: executor_a}, state, ctx, "test_name", graph_signature_hash="test_hash")
+
+    with pytest.raises(WorkflowCheckpointException, match="not a dictionary"):
+        await runner._set_executor_state("executor_a", {"key": "value"})
+
+
+async def test_runner_with_pre_loop_events():
+    """Test that pre-loop events are yielded correctly."""
+    ctx = InProcRunnerContext()
+    state = State()
+
+    runner = Runner([], {}, state, ctx, "test_name", graph_signature_hash="test_hash")
+
+    # Add an event before running
+    await ctx.add_event(WorkflowEvent.output(executor_id="test_executor", data="pre-loop-output"))
+
+    events: list[WorkflowEvent] = []
+    async for event in runner.run_until_convergence():
+        events.append(event)
+
+    # Should have the pre-loop output event
+    output_events = [e for e in events if e.type == "output"]
+    assert len(output_events) == 1
+    assert output_events[0].data == "pre-loop-output"
+
+
+class EventEmittingExecutor(Executor):
+    """An executor that emits events during execution."""
+
+    @handler
+    async def handle(self, message: MockMessage, ctx: WorkflowContext[MockMessage, int]) -> None:
+        # Emit event during processing
+        await ctx.yield_output(f"processed-{message.data}")
+        if message.data < 3:
+            await ctx.send_message(MockMessage(data=message.data + 1))
+
+
+async def test_runner_drains_straggler_events():
+    """Test that events emitted at the end of iteration are drained."""
+    executor_a = EventEmittingExecutor(id="executor_a")
+    executor_b = EventEmittingExecutor(id="executor_b")
+
+    edges = [
+        SingleEdgeGroup(executor_a.id, executor_b.id),
+        SingleEdgeGroup(executor_b.id, executor_a.id),
+    ]
+
+    executors: dict[str, Executor] = {
+        executor_a.id: executor_a,
+        executor_b.id: executor_b,
+    }
+    state = State()
+    ctx = InProcRunnerContext()
+
+    runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash")
+
+    await executor_a.execute(
+        MockMessage(data=0),
+        ["START"],
+        state,
+        ctx,
+    )
+
+    events: list[WorkflowEvent] = []
+    async for event in runner.run_until_convergence():
+        events.append(event)
+
+    # Should have output events from both executors
+    output_events = [e for e in events if e.type == "output"]
+    assert len(output_events) > 0
+
+
+async def test_runner_restore_executor_states_no_states():
+    """Test that restore does nothing when there are no executor states."""
+    executor_a = MockExecutor(id="executor_a")
+    state = State()  # No executor states set
+    state.commit()
+
+    ctx = InProcRunnerContext()
+    runner = Runner([], {executor_a.id: executor_a}, state, ctx, "test_name", graph_signature_hash="test_hash")
+
+    # Should complete without error when no executor states exist
+    await runner._restore_executor_states()
+
+
+async def test_runner_checkpoint_with_resumed_flag():
+    """Test that resumed flag prevents initial checkpoint creation."""
+    storage = InMemoryCheckpointStorage()
+    ctx = CheckpointingContext(storage)
+    executor_a = MockExecutor(id="executor_a")
+    executor_b = MockExecutor(id="executor_b")
+
+    edges = [
+        SingleEdgeGroup(executor_a.id, executor_b.id),
+        SingleEdgeGroup(executor_b.id, executor_a.id),
+    ]
+
+    executors: dict[str, Executor] = {
+        executor_a.id: executor_a,
+        executor_b.id: executor_b,
+    }
+    state = State()
+
+    runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash")
+    runner._mark_resumed(5)
+
+    # Add a message to trigger the checkpoint creation path
+    await ctx.send_message(WorkflowMessage(data=MockMessage(data=8), source_id="START"))
+
+    await executor_a.execute(
+        MockMessage(data=8),
+        ["START"],
+        state,
+        ctx,
+    )
+
+    # Run until convergence
+    async for _ in runner.run_until_convergence():
+        pass
+
+    # After completing, resumed flag should be reset
+    assert runner._resumed_from_checkpoint is False
+
+
+class ExecutorThatFailsWithEvents(Executor):
+    """An executor that emits events and then raises an exception after receiving messages."""
+
+    def __init__(self, id: str, runner_ctx: RunnerContext, fail_on_iteration: int = 1):
+        super().__init__(id=id)
+        self._runner_ctx = runner_ctx
+        self._fail_on_iteration = fail_on_iteration
+        self._iteration_count = 0
+
+    @handler
+    async def handle(self, message: MockMessage, ctx: WorkflowContext[MockMessage, int]) -> None:
+        self._iteration_count += 1
+        # First emit an output event to the workflow context
+        await ctx.yield_output(f"output-before-failure-{message.data}")
+        # Add some events directly to the runner context
+        await self._runner_ctx.add_event(WorkflowEvent.output(executor_id=self.id, data="pending-event"))
+        # Fail on the specified iteration
+        if self._iteration_count >= self._fail_on_iteration:
+            raise RuntimeError("Executor failed with pending events")
+        # Otherwise, send to next
+        await ctx.send_message(MockMessage(data=message.data + 1))
+
+
+class PassthroughExecutor(Executor):
+    """An executor that passes messages through to the failing executor."""
+
+    @handler
+    async def handle(self, message: MockMessage, ctx: WorkflowContext[MockMessage, int]) -> None:
+        await ctx.send_message(MockMessage(data=message.data))
+
+
+async def test_runner_drains_events_on_iteration_exception():
+    """Test that events are drained when iteration task raises an exception (lines 128-129)."""
+    ctx = InProcRunnerContext()
+    # executor_b will fail with pending events after receiving a message
+    executor_a = PassthroughExecutor(id="executor_a")
+    executor_b = ExecutorThatFailsWithEvents(id="executor_b", runner_ctx=ctx, fail_on_iteration=1)
+
+    edges = [
+        SingleEdgeGroup(executor_a.id, executor_b.id),
+    ]
+
+    executors: dict[str, Executor] = {
+        executor_a.id: executor_a,
+        executor_b.id: executor_b,
+    }
+    state = State()
+
+    runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash")
+
+    # Execute through executor_a which will pass to executor_b during the runner iteration
+    await executor_a.execute(
+        MockMessage(data=0),
+        ["START"],
+        state,
+        ctx,
+    )
+
+    events: list[WorkflowEvent] = []
+    with pytest.raises(RuntimeError, match="Executor failed with pending events"):
+        async for event in runner.run_until_convergence():
+            events.append(event)
+
+    # Events should include the ones emitted before the exception
+    output_events = [e for e in events if e.type == "output"]
+    # Should have drained the pending events before propagating the exception
+    assert len(output_events) >= 1
+
+
+class SlowEventEmittingExecutor(Executor):
+    """An executor that emits events with delays to test straggler event draining."""
+
+    def __init__(self, id: str, iterations_to_emit: int = 2):
+        super().__init__(id=id)
+        self.iterations_to_emit = iterations_to_emit
+        self.current_iteration = 0
+
+    @handler
+    async def handle(self, message: MockMessage, ctx: WorkflowContext[MockMessage, int]) -> None:
+        self.current_iteration += 1
+        # Emit output event
+        await ctx.yield_output(f"iteration-{self.current_iteration}")
+        # Continue sending messages until we reach the target iterations
+        if self.current_iteration < self.iterations_to_emit:
+            await ctx.send_message(MockMessage(data=message.data + 1))
+
+
+async def test_runner_drains_straggler_events_at_iteration_end():
+    """Test that events emitted at the very end of iteration are drained (lines 135-136)."""
+    # Create executors that ping-pong messages and emit events
+    executor_a = SlowEventEmittingExecutor(id="executor_a", iterations_to_emit=3)
+    executor_b = SlowEventEmittingExecutor(id="executor_b", iterations_to_emit=3)
+
+    edges = [
+        SingleEdgeGroup(executor_a.id, executor_b.id),
+        SingleEdgeGroup(executor_b.id, executor_a.id),
+    ]
+
+    executors: dict[str, Executor] = {
+        executor_a.id: executor_a,
+        executor_b.id: executor_b,
+    }
+    state = State()
+    ctx = InProcRunnerContext()
+
+    runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash")
+
+    await executor_a.execute(
+        MockMessage(data=0),
+        ["START"],
+        state,
+        ctx,
+    )
+
+    events: list[WorkflowEvent] = []
+    async for event in runner.run_until_convergence():
+        events.append(event)
+
+    # Check that output events were collected (including straggler events)
+    output_events = [e for e in events if e.type == "output"]
+    # We should have output events from both executors
+    assert len(output_events) >= 2

python/packages/core/tests/workflow/test_serialization.py

@@ -647,12 +647,11 @@ class TestSerializationWorkflowClasses:
         # Test 2: Without name and description (defaults)
         workflow2 = WorkflowBuilder(start_executor=SampleExecutor(id="e2")).build()
 
-        assert workflow2.name is None
+        assert workflow2.name is not None
         assert workflow2.description is None
 
         data2 = workflow2.to_dict()
-        assert "name" not in data2  # Should not include None values
-        assert "description" not in data2
+        assert "description" not in data2  # Should not include None values
 
         # Test 3: With only name (no description)
         workflow3 = WorkflowBuilder(name="Named Only", start_executor=SampleExecutor(id="e3")).build()

python/packages/core/tests/workflow/test_sub_workflow.py

@@ -595,7 +595,7 @@ async def test_sub_workflow_checkpoint_restore_no_duplicate_requests() -> None:
     assert first_request_id is not None
 
     # Get checkpoint
-    checkpoints = await storage.list_checkpoints(workflow1.id)
+    checkpoints = await storage.list_checkpoints(workflow_name=workflow1.name)
     checkpoint_id = max(checkpoints, key=lambda cp: cp.iteration_count).checkpoint_id
 
     # Step 2: Resume workflow from checkpoint

python/packages/core/tests/workflow/test_workflow.py

@@ -335,12 +335,9 @@ async def test_workflow_run_stream_from_checkpoint_invalid_checkpoint(
         )
 
         # Attempt to run from non-existent checkpoint should fail
-        try:
+        with pytest.raises(WorkflowCheckpointException, match="No checkpoint found with ID nonexistent_checkpoint_id"):
             async for _ in workflow.run(checkpoint_id="nonexistent_checkpoint_id", stream=True):
                 pass
-            raise AssertionError("Expected WorkflowCheckpointException to be raised")
-        except WorkflowCheckpointException as e:
-            assert str(e) == "Checkpoint nonexistent_checkpoint_id not found"
 
 
 async def test_workflow_run_stream_from_checkpoint_with_external_storage(
@@ -354,12 +351,14 @@ async def test_workflow_run_stream_from_checkpoint_with_external_storage(
         from agent_framework import WorkflowCheckpoint
 
         test_checkpoint = WorkflowCheckpoint(
-            workflow_id="test-workflow",
+            workflow_name="test-workflow",
+            graph_signature_hash="test-graph-signature",
+            previous_checkpoint_id=None,
             messages={},
             state={},
             iteration_count=0,
         )
-        checkpoint_id = await storage.save_checkpoint(test_checkpoint)
+        checkpoint_id = await storage.save(test_checkpoint)
 
         # Create a workflow WITHOUT checkpointing
         workflow_without_checkpointing = (
@@ -385,17 +384,6 @@ async def test_workflow_run_from_checkpoint_non_streaming(simple_executor: Execu
     with tempfile.TemporaryDirectory() as temp_dir:
         storage = FileCheckpointStorage(temp_dir)
 
-        # Create a test checkpoint manually in storage
-        from agent_framework import WorkflowCheckpoint
-
-        test_checkpoint = WorkflowCheckpoint(
-            workflow_id="test-workflow",
-            messages={},
-            state={},
-            iteration_count=0,
-        )
-        checkpoint_id = await storage.save_checkpoint(test_checkpoint)
-
         # Build workflow with checkpointing
         workflow = (
             WorkflowBuilder(start_executor=simple_executor, checkpoint_storage=storage)
@@ -403,6 +391,19 @@ async def test_workflow_run_from_checkpoint_non_streaming(simple_executor: Execu
             .build()
         )
 
+        # Create a test checkpoint manually in storage
+        from agent_framework import WorkflowCheckpoint
+
+        test_checkpoint = WorkflowCheckpoint(
+            workflow_name=workflow.name,
+            graph_signature_hash=workflow.graph_signature_hash,
+            previous_checkpoint_id=None,
+            messages={},
+            state={},
+            iteration_count=0,
+        )
+        checkpoint_id = await storage.save(test_checkpoint)
+
         # Test non-streaming run method with checkpoint_id
         result = await workflow.run(checkpoint_id=checkpoint_id)
         assert isinstance(result, list)  # Should return WorkflowRunResult which extends list
@@ -416,11 +417,19 @@ async def test_workflow_run_stream_from_checkpoint_with_responses(
     with tempfile.TemporaryDirectory() as temp_dir:
         storage = FileCheckpointStorage(temp_dir)
 
+        # Build workflow with checkpointing
+        workflow = (
+            WorkflowBuilder(start_executor=simple_executor, checkpoint_storage=storage)
+            .add_edge(simple_executor, simple_executor)
+            .build()
+        )
+
         # Create a test checkpoint manually in storage
         from agent_framework import WorkflowCheckpoint
 
         test_checkpoint = WorkflowCheckpoint(
-            workflow_id="test-workflow",
+            workflow_name=workflow.name,
+            graph_signature_hash=workflow.graph_signature_hash,
             messages={},
             state={},
             pending_request_info_events={
@@ -429,18 +438,11 @@ async def test_workflow_run_stream_from_checkpoint_with_responses(
                     source_executor_id=simple_executor.id,
                     request_data="Mock",
                     response_type=str,
-                ).to_dict(),
+                ),
             },
             iteration_count=0,
         )
-        checkpoint_id = await storage.save_checkpoint(test_checkpoint)
-
-        # Build workflow with checkpointing
-        workflow = (
-            WorkflowBuilder(start_executor=simple_executor, checkpoint_storage=storage)
-            .add_edge(simple_executor, simple_executor)
-            .build()
-        )
+        checkpoint_id = await storage.save(test_checkpoint)
 
         # Resume from checkpoint - pending request events should be emitted
         events: list[WorkflowEvent] = []
@@ -542,7 +544,7 @@ async def test_workflow_checkpoint_runtime_only_configuration(
         assert result.get_final_state() == WorkflowRunState.IDLE
 
         # Verify checkpoints were created
-        checkpoints = await storage.list_checkpoints()
+        checkpoints = await storage.list_checkpoints(workflow_name=workflow.name)
         assert len(checkpoints) > 0
 
         # Find a superstep checkpoint to resume from
@@ -592,8 +594,8 @@ async def test_workflow_checkpoint_runtime_overrides_buildtime(
         assert result is not None
 
         # Verify checkpoints were created in runtime storage, not build-time storage
-        buildtime_checkpoints = await buildtime_storage.list_checkpoints()
-        runtime_checkpoints = await runtime_storage.list_checkpoints()
+        buildtime_checkpoints = await buildtime_storage.list_checkpoints(workflow_name=workflow.name)
+        runtime_checkpoints = await runtime_storage.list_checkpoints(workflow_name=workflow.name)
 
         assert len(runtime_checkpoints) > 0, "Runtime storage should have checkpoints"
         assert len(buildtime_checkpoints) == 0, "Build-time storage should have no checkpoints when overridden"

python/packages/core/tests/workflow/test_workflow_agent.py

@@ -607,7 +607,7 @@ class TestWorkflowAgent:
 
         # Drain workflow events to get checkpoint
         # The workflow should have created checkpoints
-        checkpoints = await checkpoint_storage.list_checkpoints(workflow.id)
+        checkpoints = await checkpoint_storage.list_checkpoints(workflow_name=workflow.name)
         assert len(checkpoints) > 0, "Checkpoints should have been created when checkpoint_storage is provided"
 
     async def test_agent_executor_output_response_false_filters_streaming_events(self):

python/packages/core/tests/workflow/test_workflow_observability.py

@@ -306,8 +306,8 @@ async def test_end_to_end_workflow_tracing(span_exporter: InMemorySpanExporter)
     assert len(build_spans_with_metadata) == 1
     metadata_build_span = build_spans_with_metadata[0]
     assert metadata_build_span.attributes is not None
-    assert metadata_build_span.attributes.get(OtelAttr.WORKFLOW_NAME) == "Test Pipeline"
-    assert metadata_build_span.attributes.get(OtelAttr.WORKFLOW_DESCRIPTION) == "Test workflow description"
+    assert metadata_build_span.attributes.get(OtelAttr.WORKFLOW_BUILDER_NAME) == "Test Pipeline"
+    assert metadata_build_span.attributes.get(OtelAttr.WORKFLOW_BUILDER_DESCRIPTION) == "Test workflow description"
 
     # Clear spans to separate build from run tracing
     span_exporter.clear()
@@ -451,14 +451,14 @@ async def test_message_trace_context_serialization(span_exporter: InMemorySpanEx
     await ctx.send_message(message)
 
     # Create a checkpoint that includes the message
-    checkpoint_id = await ctx.create_checkpoint(State(), 0)
+    checkpoint_id = await ctx.create_checkpoint("test_name", "test_hash", State(), None, 0)
     checkpoint = await ctx.load_checkpoint(checkpoint_id)
     assert checkpoint is not None
 
     # Check serialized message includes trace context
     serialized_msg = checkpoint.messages["source"][0]
-    assert serialized_msg["trace_contexts"] == [{"traceparent": "00-trace-span-01"}]
-    assert serialized_msg["source_span_ids"] == ["span123"]
+    assert serialized_msg.trace_contexts == [{"traceparent": "00-trace-span-01"}]
+    assert serialized_msg.source_span_ids == ["span123"]
 
     # Test deserialization
     await ctx.apply_checkpoint(checkpoint)