Python: introduce workflow checkpointing (#366)

* Add workflow checkpointing functionality. * Reintroduce protocol that went missing during merge * Checkpoint updates * Fix ordering of checkpointing * Cleanup * Cleanup - thanks Copilot * Cleanup - thanks Copilot * State reset updates * State reset updates 2 * Workflow fixes and updates. Addressed PR feedback * A few updates
2026-06-16 21:04:09 +08:00 · 2025-08-12 07:33:46 +09:00
parent bbc07931c1
commit 19676978e9
14 changed files with 1693 additions and 79 deletions
@@ -0,0 +1,334 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import json
+import tempfile
+from datetime import datetime, timezone
+from pathlib import Path
+
+from agent_framework_workflow._checkpoint import (
+    FileCheckpointStorage,
+    InMemoryCheckpointStorage,
+    WorkflowCheckpoint,
+)
+
+
+def test_workflow_checkpoint_default_values():
+    checkpoint = WorkflowCheckpoint()
+
+    assert checkpoint.checkpoint_id != ""
+    assert checkpoint.workflow_id == ""
+    assert checkpoint.timestamp != ""
+    assert checkpoint.messages == {}
+    assert checkpoint.shared_state == {}
+    assert checkpoint.executor_states == {}
+    assert checkpoint.iteration_count == 0
+    assert checkpoint.max_iterations == 100
+    assert checkpoint.metadata == {}
+    assert checkpoint.version == "1.0"
+
+
+def test_workflow_checkpoint_custom_values():
+    custom_timestamp = datetime.now(timezone.utc).isoformat()
+    checkpoint = WorkflowCheckpoint(
+        checkpoint_id="test-checkpoint-123",
+        workflow_id="test-workflow-456",
+        timestamp=custom_timestamp,
+        messages={"executor1": [{"data": "test"}]},
+        shared_state={"key": "value"},
+        executor_states={"executor1": {"state": "active"}},
+        iteration_count=5,
+        max_iterations=50,
+        metadata={"test": True},
+        version="2.0",
+    )
+
+    assert checkpoint.checkpoint_id == "test-checkpoint-123"
+    assert checkpoint.workflow_id == "test-workflow-456"
+    assert checkpoint.timestamp == custom_timestamp
+    assert checkpoint.messages == {"executor1": [{"data": "test"}]}
+    assert checkpoint.shared_state == {"key": "value"}
+    assert checkpoint.executor_states == {"executor1": {"state": "active"}}
+    assert checkpoint.iteration_count == 5
+    assert checkpoint.max_iterations == 50
+    assert checkpoint.metadata == {"test": True}
+    assert checkpoint.version == "2.0"
+
+
+async def test_memory_checkpoint_storage_save_and_load():
+    storage = InMemoryCheckpointStorage()
+    checkpoint = WorkflowCheckpoint(workflow_id="test-workflow", messages={"executor1": [{"data": "hello"}]})
+
+    # Save checkpoint
+    saved_id = await storage.save_checkpoint(checkpoint)
+    assert saved_id == checkpoint.checkpoint_id
+
+    # Load checkpoint
+    loaded_checkpoint = await storage.load_checkpoint(checkpoint.checkpoint_id)
+    assert loaded_checkpoint is not None
+    assert loaded_checkpoint.checkpoint_id == checkpoint.checkpoint_id
+    assert loaded_checkpoint.workflow_id == checkpoint.workflow_id
+    assert loaded_checkpoint.messages == checkpoint.messages
+
+
+async def test_memory_checkpoint_storage_load_nonexistent():
+    storage = InMemoryCheckpointStorage()
+
+    result = await storage.load_checkpoint("nonexistent-id")
+    assert result is None
+
+
+async def test_memory_checkpoint_storage_list_checkpoints():
+    storage = InMemoryCheckpointStorage()
+
+    # Create checkpoints for different workflows
+    checkpoint1 = WorkflowCheckpoint(workflow_id="workflow-1")
+    checkpoint2 = WorkflowCheckpoint(workflow_id="workflow-1")
+    checkpoint3 = WorkflowCheckpoint(workflow_id="workflow-2")
+
+    await storage.save_checkpoint(checkpoint1)
+    await storage.save_checkpoint(checkpoint2)
+    await storage.save_checkpoint(checkpoint3)
+
+    # Test list_checkpoint_ids for workflow-1
+    workflow1_checkpoint_ids = await storage.list_checkpoint_ids("workflow-1")
+    assert len(workflow1_checkpoint_ids) == 2
+    assert checkpoint1.checkpoint_id in workflow1_checkpoint_ids
+    assert checkpoint2.checkpoint_id in workflow1_checkpoint_ids
+
+    # Test list_checkpoints for workflow-1 (returns objects)
+    workflow1_checkpoints = await storage.list_checkpoints("workflow-1")
+    assert len(workflow1_checkpoints) == 2
+    assert all(isinstance(cp, WorkflowCheckpoint) for cp in workflow1_checkpoints)
+    assert {cp.checkpoint_id for cp in workflow1_checkpoints} == {checkpoint1.checkpoint_id, checkpoint2.checkpoint_id}
+
+    # Test list_checkpoint_ids for workflow-2
+    workflow2_checkpoint_ids = await storage.list_checkpoint_ids("workflow-2")
+    assert len(workflow2_checkpoint_ids) == 1
+    assert checkpoint3.checkpoint_id in workflow2_checkpoint_ids
+
+    # Test list_checkpoints for workflow-2 (returns objects)
+    workflow2_checkpoints = await storage.list_checkpoints("workflow-2")
+    assert len(workflow2_checkpoints) == 1
+    assert workflow2_checkpoints[0].checkpoint_id == checkpoint3.checkpoint_id
+
+    # Test list_checkpoint_ids for non-existent workflow
+    empty_checkpoint_ids = await storage.list_checkpoint_ids("nonexistent-workflow")
+    assert len(empty_checkpoint_ids) == 0
+
+    # Test list_checkpoints for non-existent workflow
+    empty_checkpoints = await storage.list_checkpoints("nonexistent-workflow")
+    assert len(empty_checkpoints) == 0
+
+    # Test list_checkpoint_ids without workflow filter (all checkpoints)
+    all_checkpoint_ids = await storage.list_checkpoint_ids()
+    assert len(all_checkpoint_ids) == 3
+    expected_ids = {checkpoint1.checkpoint_id, checkpoint2.checkpoint_id, checkpoint3.checkpoint_id}
+    assert expected_ids.issubset(set(all_checkpoint_ids))
+
+    # Test list_checkpoints without workflow filter (all checkpoints)
+    all_checkpoints = await storage.list_checkpoints()
+    assert len(all_checkpoints) == 3
+    assert all(isinstance(cp, WorkflowCheckpoint) for cp in all_checkpoints)
+
+
+async def test_memory_checkpoint_storage_delete():
+    storage = InMemoryCheckpointStorage()
+    checkpoint = WorkflowCheckpoint(workflow_id="test-workflow")
+
+    # Save checkpoint
+    await storage.save_checkpoint(checkpoint)
+    assert await storage.load_checkpoint(checkpoint.checkpoint_id) is not None
+
+    # Delete checkpoint
+    result = await storage.delete_checkpoint(checkpoint.checkpoint_id)
+    assert result is True
+
+    # Verify deletion
+    assert await storage.load_checkpoint(checkpoint.checkpoint_id) is None
+
+    # Try to delete again
+    result = await storage.delete_checkpoint(checkpoint.checkpoint_id)
+    assert result is False
+
+
+async def test_file_checkpoint_storage_save_and_load():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+        checkpoint = WorkflowCheckpoint(
+            workflow_id="test-workflow",
+            messages={"executor1": [{"data": "hello", "source_id": "test", "target_id": None}]},
+            shared_state={"key": "value"},
+        )
+
+        # Save checkpoint
+        saved_id = await storage.save_checkpoint(checkpoint)
+        assert saved_id == checkpoint.checkpoint_id
+
+        # Verify file was created
+        file_path = Path(temp_dir) / f"{checkpoint.checkpoint_id}.json"
+        assert file_path.exists()
+
+        # Load checkpoint
+        loaded_checkpoint = await storage.load_checkpoint(checkpoint.checkpoint_id)
+        assert loaded_checkpoint is not None
+        assert loaded_checkpoint.checkpoint_id == checkpoint.checkpoint_id
+        assert loaded_checkpoint.workflow_id == checkpoint.workflow_id
+        assert loaded_checkpoint.messages == checkpoint.messages
+        assert loaded_checkpoint.shared_state == checkpoint.shared_state
+
+
+async def test_file_checkpoint_storage_load_nonexistent():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        result = await storage.load_checkpoint("nonexistent-id")
+        assert result is None
+
+
+async def test_file_checkpoint_storage_list_checkpoints():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Create checkpoints for different workflows
+        checkpoint1 = WorkflowCheckpoint(workflow_id="workflow-1")
+        checkpoint2 = WorkflowCheckpoint(workflow_id="workflow-1")
+        checkpoint3 = WorkflowCheckpoint(workflow_id="workflow-2")
+
+        await storage.save_checkpoint(checkpoint1)
+        await storage.save_checkpoint(checkpoint2)
+        await storage.save_checkpoint(checkpoint3)
+
+        # Test list_checkpoint_ids for workflow-1
+        workflow1_checkpoint_ids = await storage.list_checkpoint_ids("workflow-1")
+        assert len(workflow1_checkpoint_ids) == 2
+        assert checkpoint1.checkpoint_id in workflow1_checkpoint_ids
+        assert checkpoint2.checkpoint_id in workflow1_checkpoint_ids
+
+        # Test list_checkpoints for workflow-1 (returns objects)
+        workflow1_checkpoints = await storage.list_checkpoints("workflow-1")
+        assert len(workflow1_checkpoints) == 2
+        assert all(isinstance(cp, WorkflowCheckpoint) for cp in workflow1_checkpoints)
+        checkpoint_ids = {cp.checkpoint_id for cp in workflow1_checkpoints}
+        assert checkpoint_ids == {checkpoint1.checkpoint_id, checkpoint2.checkpoint_id}
+
+        # Test list_checkpoint_ids for workflow-2
+        workflow2_checkpoint_ids = await storage.list_checkpoint_ids("workflow-2")
+        assert len(workflow2_checkpoint_ids) == 1
+        assert checkpoint3.checkpoint_id in workflow2_checkpoint_ids
+
+        # Test list_checkpoints for workflow-2 (returns objects)
+        workflow2_checkpoints = await storage.list_checkpoints("workflow-2")
+        assert len(workflow2_checkpoints) == 1
+        assert workflow2_checkpoints[0].checkpoint_id == checkpoint3.checkpoint_id
+
+        # Test list all checkpoints
+        all_checkpoint_ids = await storage.list_checkpoint_ids()
+        assert len(all_checkpoint_ids) == 3
+
+        all_checkpoints = await storage.list_checkpoints()
+        assert len(all_checkpoints) == 3
+        assert all(isinstance(cp, WorkflowCheckpoint) for cp in all_checkpoints)
+
+
+async def test_file_checkpoint_storage_delete():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+        checkpoint = WorkflowCheckpoint(workflow_id="test-workflow")
+
+        # Save checkpoint
+        await storage.save_checkpoint(checkpoint)
+        file_path = Path(temp_dir) / f"{checkpoint.checkpoint_id}.json"
+        assert file_path.exists()
+
+        # Delete checkpoint
+        result = await storage.delete_checkpoint(checkpoint.checkpoint_id)
+        assert result is True
+        assert not file_path.exists()
+
+        # Try to delete again
+        result = await storage.delete_checkpoint(checkpoint.checkpoint_id)
+        assert result is False
+
+
+async def test_file_checkpoint_storage_directory_creation():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        nested_path = Path(temp_dir) / "nested" / "checkpoint" / "storage"
+        storage = FileCheckpointStorage(nested_path)
+
+        # Directory should be created
+        assert nested_path.exists()
+        assert nested_path.is_dir()
+
+        # Should be able to save checkpoints
+        checkpoint = WorkflowCheckpoint(workflow_id="test")
+        await storage.save_checkpoint(checkpoint)
+
+        file_path = nested_path / f"{checkpoint.checkpoint_id}.json"
+        assert file_path.exists()
+
+
+async def test_file_checkpoint_storage_corrupted_file():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Create a corrupted JSON file
+        corrupted_file = Path(temp_dir) / "corrupted.json"
+        with open(corrupted_file, "w") as f:  # noqa: ASYNC230
+            f.write("{ invalid json }")
+
+        # list_checkpoints should handle the corrupted file gracefully
+        checkpoints = await storage.list_checkpoints("any-workflow")
+        assert checkpoints == []
+
+
+async def test_file_checkpoint_storage_json_serialization():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Create checkpoint with complex nested data
+        checkpoint = WorkflowCheckpoint(
+            workflow_id="complex-workflow",
+            messages={"executor1": [{"data": {"nested": {"value": 42}}, "source_id": "test", "target_id": None}]},
+            shared_state={"list": [1, 2, 3], "dict": {"a": "b", "c": {"d": "e"}}, "bool": True, "null": None},
+            executor_states={"executor1": {"state": "active", "config": {"timeout": 30, "retries": 3}}},
+        )
+
+        # Save and load
+        await storage.save_checkpoint(checkpoint)
+        loaded = await storage.load_checkpoint(checkpoint.checkpoint_id)
+
+        assert loaded is not None
+        assert loaded.messages == checkpoint.messages
+        assert loaded.shared_state == checkpoint.shared_state
+        assert loaded.executor_states == checkpoint.executor_states
+
+        # Verify the JSON file is properly formatted
+        file_path = Path(temp_dir) / f"{checkpoint.checkpoint_id}.json"
+        with open(file_path) as f:  # noqa: ASYNC230
+            data = json.load(f)
+
+        assert data["messages"]["executor1"][0]["data"]["nested"]["value"] == 42
+        assert data["shared_state"]["list"] == [1, 2, 3]
+        assert data["shared_state"]["bool"] is True
+        assert data["shared_state"]["null"] is None
+
+
+def test_checkpoint_storage_protocol_compliance():
+    # This test ensures both implementations have all required methods
+    memory_storage = InMemoryCheckpointStorage()
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        file_storage = FileCheckpointStorage(temp_dir)
+
+        for storage in [memory_storage, file_storage]:
+            # Test that all protocol methods exist and are callable
+            assert hasattr(storage, "save_checkpoint")
+            assert callable(storage.save_checkpoint)
+            assert hasattr(storage, "load_checkpoint")
+            assert callable(storage.load_checkpoint)
+            assert hasattr(storage, "list_checkpoint_ids")
+            assert callable(storage.list_checkpoint_ids)
+            assert hasattr(storage, "list_checkpoints")
+            assert callable(storage.list_checkpoints)
+            assert hasattr(storage, "delete_checkpoint")
+            assert callable(storage.delete_checkpoint)
@@ -1,10 +1,12 @@
 # Copyright (c) Microsoft. All rights reserved.

+import tempfile
 from dataclasses import dataclass

 import pytest
 from agent_framework.workflow import (
    Executor,
+    FileCheckpointStorage,
    RequestInfoEvent,
    RequestInfoExecutor,
    RequestInfoMessage,
@@ -15,6 +17,8 @@ from agent_framework.workflow import (
    handler,
 )

+from agent_framework_workflow import Message
+

@dataclass
 class MockMessage:
@@ -275,3 +279,278 @@ async def test_fan_in():

    completed_event = events.get_completed_event()
    assert completed_event is not None and completed_event.data == 4
+
+
+@pytest.fixture
+def simple_executor() -> Executor:
+    class SimpleExecutor(Executor):
+        @handler
+        async def handle_message(self, message: Message, context: WorkflowContext) -> None:
+            pass
+
+    return SimpleExecutor("test_executor")
+
+
+async def test_workflow_with_checkpointing_enabled(simple_executor: Executor):
+    """Test that a workflow can be built with checkpointing enabled."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Build workflow with checkpointing - should not raise any errors
+        workflow = (
+            WorkflowBuilder()
+            .add_edge(simple_executor, simple_executor)  # Self-loop to satisfy graph requirements
+            .set_start_executor(simple_executor)
+            .with_checkpointing(storage)
+            .build()
+        )
+
+        # Verify workflow was created and can run
+        test_message = Message(data="test message", source_id="test", target_id=None)
+        result = await workflow.run(test_message)
+        assert result is not None
+
+
+async def test_workflow_checkpointing_not_enabled_for_external_restore(simple_executor: Executor):
+    """Test that external checkpoint restoration fails when workflow doesn't support checkpointing."""
+    # Build workflow WITHOUT checkpointing
+    workflow = (
+        WorkflowBuilder()
+        .add_edge(simple_executor, simple_executor)  # Self-loop to satisfy graph requirements
+        .set_start_executor(simple_executor)
+        .build()
+    )
+
+    # Attempt to restore from checkpoint without providing external storage should fail
+    try:
+        [event async for event in workflow.run_streaming_from_checkpoint("fake-checkpoint-id")]
+        raise AssertionError("Expected ValueError to be raised")
+    except ValueError as e:
+        assert "Cannot restore from checkpoint" in str(e)
+        assert "either provide checkpoint_storage parameter" in str(e)
+
+
+async def test_workflow_run_stream_from_checkpoint_no_checkpointing_enabled(simple_executor: Executor):
+    # Build workflow WITHOUT checkpointing
+    workflow = (
+        WorkflowBuilder()
+        .add_edge(simple_executor, simple_executor)  # Self-loop to satisfy graph requirements
+        .set_start_executor(simple_executor)
+        .build()
+    )
+
+    # Attempt to run from checkpoint should fail
+    try:
+        async for _ in workflow.run_streaming_from_checkpoint("fake_checkpoint_id"):
+            pass
+        raise AssertionError("Expected ValueError to be raised")
+    except ValueError as e:
+        assert "Cannot restore from checkpoint" in str(e)
+        assert "either provide checkpoint_storage parameter" in str(e)
+
+
+async def test_workflow_run_stream_from_checkpoint_invalid_checkpoint(simple_executor: Executor):
+    """Test that attempting to restore from a non-existent checkpoint fails appropriately."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Build workflow with checkpointing
+        workflow = (
+            WorkflowBuilder()
+            .add_edge(simple_executor, simple_executor)  # Self-loop to satisfy graph requirements
+            .set_start_executor(simple_executor)
+            .with_checkpointing(storage)
+            .build()
+        )
+
+        # Attempt to run from non-existent checkpoint should fail
+        try:
+            async for _ in workflow.run_streaming_from_checkpoint("nonexistent_checkpoint_id"):
+                pass
+            raise AssertionError("Expected RuntimeError to be raised")
+        except RuntimeError as e:
+            assert "Failed to restore from checkpoint" in str(e)
+
+
+async def test_workflow_run_stream_from_checkpoint_with_external_storage(simple_executor: Executor):
+    """Test that external checkpoint storage can be provided for restoration."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Create a test checkpoint manually in storage
+        from agent_framework_workflow._checkpoint import WorkflowCheckpoint
+
+        test_checkpoint = WorkflowCheckpoint(
+            workflow_id="test-workflow",
+            messages={},
+            shared_state={},
+            executor_states={},
+            iteration_count=0,
+            max_iterations=100,
+        )
+        checkpoint_id = await storage.save_checkpoint(test_checkpoint)
+
+        # Create a workflow WITHOUT checkpointing
+        workflow_without_checkpointing = (
+            WorkflowBuilder().add_edge(simple_executor, simple_executor).set_start_executor(simple_executor).build()
+        )
+
+        # Resume from checkpoint using external storage parameter
+        try:
+            events: list[WorkflowEvent] = []
+            async for event in workflow_without_checkpointing.run_streaming_from_checkpoint(
+                checkpoint_id, checkpoint_storage=storage
+            ):
+                events.append(event)
+                if len(events) >= 2:  # Limit to avoid infinite loops
+                    break
+        except Exception:
+            # Expected since we have minimal setup, but method should accept the parameters
+            pass
+
+
+async def test_workflow_run_from_checkpoint_non_streaming(simple_executor: Executor):
+    """Test the non-streaming run_from_checkpoint method."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Create a test checkpoint manually in storage
+        from agent_framework_workflow._checkpoint import WorkflowCheckpoint
+
+        test_checkpoint = WorkflowCheckpoint(
+            workflow_id="test-workflow",
+            messages={},
+            shared_state={},
+            executor_states={},
+            iteration_count=0,
+            max_iterations=100,
+        )
+        checkpoint_id = await storage.save_checkpoint(test_checkpoint)
+
+        # Build workflow with checkpointing
+        workflow = (
+            WorkflowBuilder()
+            .add_edge(simple_executor, simple_executor)
+            .set_start_executor(simple_executor)
+            .with_checkpointing(storage)
+            .build()
+        )
+
+        # Test non-streaming run_from_checkpoint method
+        result = await workflow.run_from_checkpoint(checkpoint_id)
+        assert isinstance(result, list)  # Should return WorkflowRunResult which extends list
+        assert hasattr(result, "get_completed_event")  # Should have WorkflowRunResult methods
+
+
+async def test_workflow_run_stream_from_checkpoint_with_responses(simple_executor: Executor):
+    """Test that run_streaming_from_checkpoint accepts responses parameter."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Create a test checkpoint manually in storage
+        from agent_framework_workflow._checkpoint import WorkflowCheckpoint
+
+        test_checkpoint = WorkflowCheckpoint(
+            workflow_id="test-workflow",
+            messages={},
+            shared_state={},
+            executor_states={},
+            iteration_count=0,
+            max_iterations=100,
+        )
+        checkpoint_id = await storage.save_checkpoint(test_checkpoint)
+
+        # Build workflow with checkpointing
+        workflow = (
+            WorkflowBuilder()
+            .add_edge(simple_executor, simple_executor)
+            .set_start_executor(simple_executor)
+            .with_checkpointing(storage)
+            .build()
+        )
+
+        # Test that run_stream_from_checkpoint accepts responses parameter
+        responses = {"request_123": {"data": "test_response"}}
+
+        try:
+            events: list[WorkflowEvent] = []
+            async for event in workflow.run_streaming_from_checkpoint(checkpoint_id, responses=responses):
+                events.append(event)
+                if len(events) >= 2:  # Limit to avoid infinite loops
+                    break
+        except Exception:
+            # Expected since we have minimal setup, but method should accept the parameters
+            pass
+
+
+@dataclass
+class StateTrackingMessage:
+    """A message that tracks state for testing context reset behavior."""
+
+    data: str
+    run_id: str
+
+
+class StateTrackingExecutor(Executor):
+    """An executor that tracks state in shared state to test context reset behavior."""
+
+    @handler(output_types=[])
+    async def handle_message(self, message: StateTrackingMessage, ctx: WorkflowContext) -> None:
+        """Handle the message and track it in shared state."""
+        # Get existing messages from shared state
+        try:
+            existing_messages = await ctx.get_shared_state("processed_messages")
+        except KeyError:
+            existing_messages = []
+
+        # Record this message
+        message_record = f"{message.run_id}:{message.data}"
+        existing_messages.append(message_record)  # type: ignore
+
+        # Update shared state
+        await ctx.set_shared_state("processed_messages", existing_messages)
+
+        # Complete workflow with current shared state
+        await ctx.add_event(WorkflowCompletedEvent(data=existing_messages.copy()))  # type: ignore
+
+
+async def test_workflow_multiple_runs_no_state_collision():
+    """Test that running the same workflow instance multiple times doesn't have state collision."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+
+        # Create executor that tracks state in shared state
+        state_executor = StateTrackingExecutor("state_executor")
+
+        # Build workflow with checkpointing
+        workflow = (
+            WorkflowBuilder()
+            .add_edge(state_executor, state_executor)  # Self-loop to satisfy graph requirements
+            .set_start_executor(state_executor)
+            .with_checkpointing(storage)
+            .build()
+        )
+
+        # Run 1: Should only see messages from run 1
+        result1 = await workflow.run(StateTrackingMessage(data="message1", run_id="run1"))
+        completed1 = result1.get_completed_event()
+        assert completed1 is not None
+        assert completed1.data == ["run1:message1"]
+
+        # Run 2: Should only see messages from run 2, not run 1
+        result2 = await workflow.run(StateTrackingMessage(data="message2", run_id="run2"))
+        completed2 = result2.get_completed_event()
+        assert completed2 is not None
+        assert completed2.data == ["run2:message2"]  # Should NOT contain run1 data
+
+        # Run 3: Should only see messages from run 3
+        result3 = await workflow.run(StateTrackingMessage(data="message3", run_id="run3"))
+        completed3 = result3.get_completed_event()
+        assert completed3 is not None
+        assert completed3.data == ["run3:message3"]  # Should NOT contain run1 or run2 data
+
+        # Verify that each run only processed its own message
+        # This confirms that the checkpointable context properly resets between runs
+        assert completed1.data != completed2.data
+        assert completed2.data != completed3.data
+        assert completed1.data != completed3.data