Add tests

python/packages/core/agent_framework/_workflows/_runner.py

@@ -97,7 +97,7 @@ class Runner:
                 yield event
 
         # Create a checkpoint before a run starts. Checkpoints are usually considered to be created at the
-        # end of an iteration, we can think of this checkpoint as being created at the end of a "superstep 0"
+        # end of an iteration, we can think of this checkpoint as being created at the end of "superstep 0"
         # which captures the states after which the start executor has run. Note that we execute the start
         # executor outside of the main iteration loop.
         if await self._ctx.has_messages() and not self._resumed_from_checkpoint:

python/packages/core/agent_framework/_workflows/_workflow.py

@@ -788,8 +788,7 @@ class Workflow(DictConvertible):
         # fully consumed, ``_run_core``'s finally clears the attribute. When the
         # caller drops the stream without iterating, garbage collection invalidates
         # the weakref, so a subsequent ``run`` is permitted.
-        existing_stream = self._active_run() if self._active_run is not None else None
-        if existing_stream is not None:
+        if self._is_run_active():
             raise WorkflowException(
                 "Workflow is already running; concurrent runs are not allowed on the same instance."
             )

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

@@ -0,0 +1,72 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tests for `InProcRunnerContext`."""
+
+import pytest
+
+from agent_framework import (
+    InProcRunnerContext,
+    WorkflowEvent,
+    WorkflowMessage,
+)
+
+
+def _make_request_info_event(request_id: str, source_executor_id: str = "executor") -> WorkflowEvent[str]:
+    return WorkflowEvent.request_info(
+        request_id=request_id,
+        source_executor_id=source_executor_id,
+        request_data="please respond",
+        response_type=str,
+    )
+
+
+class TestInProcRunnerContextResetForNewRun:
+    """Verify `reset_for_new_run` clears per-run state, including pending request_info events."""
+
+    async def test_reset_clears_pending_request_info_events(self) -> None:
+        ctx = InProcRunnerContext()
+
+        await ctx.add_request_info_event(_make_request_info_event("req-1"))
+        await ctx.add_request_info_event(_make_request_info_event("req-2"))
+
+        assert set((await ctx.get_pending_request_info_events()).keys()) == {"req-1", "req-2"}
+
+        ctx.reset_for_new_run()
+
+        assert await ctx.get_pending_request_info_events() == {}
+
+    async def test_reset_clears_pending_request_info_events_when_already_empty(self) -> None:
+        ctx = InProcRunnerContext()
+
+        assert await ctx.get_pending_request_info_events() == {}
+
+        ctx.reset_for_new_run()
+
+        assert await ctx.get_pending_request_info_events() == {}
+
+    async def test_reset_after_pending_event_blocks_response_correlation(self) -> None:
+        """After `reset_for_new_run`, prior request ids must no longer correlate to a response."""
+        ctx = InProcRunnerContext()
+        await ctx.add_request_info_event(_make_request_info_event("req-1"))
+
+        ctx.reset_for_new_run()
+
+        with pytest.raises(ValueError, match="No pending request found for request_id: req-1"):
+            await ctx.send_request_info_response("req-1", "answer")
+
+    async def test_reset_clears_messages_events_and_streaming_flag(self) -> None:
+        """Sanity-check the other state `reset_for_new_run` is documented to clear."""
+        ctx = InProcRunnerContext()
+        await ctx.send_message(WorkflowMessage(data="hello", source_id="executor"))
+        await ctx.add_event(WorkflowEvent("status", data="running"))
+        ctx.set_streaming(True)
+
+        assert await ctx.has_messages() is True
+        assert await ctx.has_events() is True
+        assert ctx.is_streaming() is True
+
+        ctx.reset_for_new_run()
+
+        assert await ctx.has_messages() is False
+        assert await ctx.has_events() is False
+        assert ctx.is_streaming() is False

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

@@ -20,6 +20,7 @@ from agent_framework import (
     Content,
     Executor,
     FileCheckpointStorage,
+    InMemoryCheckpointStorage,
     Message,
     ResponseStream,
     WorkflowBuilder,
@@ -1353,3 +1354,112 @@ async def test_output_executors_filtering_with_run_responses_streaming() -> None
 
 
 # endregion
+
+
+# region Workflow.create_checkpoint
+
+
+class TestWorkflowCreateCheckpoint:
+    """Tests for :meth:`Workflow.create_checkpoint`."""
+
+    async def test_returns_checkpoint_id_with_runtime_storage(self, simple_executor: Executor) -> None:
+        """Calling `create_checkpoint` with a runtime storage persists a checkpoint and returns its id."""
+        storage = InMemoryCheckpointStorage()
+        workflow = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, simple_executor).build()
+
+        checkpoint_id = await workflow.create_checkpoint(storage)
+
+        assert checkpoint_id
+        loaded = await storage.load(checkpoint_id)
+        assert loaded is not None
+        assert loaded.checkpoint_id == checkpoint_id
+        assert loaded.workflow_name == workflow.name
+        assert loaded.graph_signature_hash == workflow.graph_signature_hash
+
+    async def test_uses_buildtime_storage_when_none_provided(self, simple_executor: Executor) -> None:
+        """When called with `None`, the build-time storage is used."""
+        storage = InMemoryCheckpointStorage()
+        workflow = (
+            WorkflowBuilder(start_executor=simple_executor, checkpoint_storage=storage)
+            .add_edge(simple_executor, simple_executor)
+            .build()
+        )
+
+        checkpoint_id = await workflow.create_checkpoint(None)
+
+        loaded = await storage.load(checkpoint_id)
+        assert loaded is not None
+        assert loaded.checkpoint_id == checkpoint_id
+
+    async def test_raises_when_no_storage_available(self, simple_executor: Executor) -> None:
+        """Without build-time or runtime storage, `create_checkpoint(None)` raises."""
+        workflow = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, simple_executor).build()
+
+        with pytest.raises(WorkflowCheckpointException, match="Checkpoint storage must be provided"):
+            await workflow.create_checkpoint(None)
+
+    async def test_raises_while_run_active(self, simple_executor: Executor) -> None:
+        """`create_checkpoint` must reject while a workflow run is still active."""
+        storage = InMemoryCheckpointStorage()
+        workflow = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, simple_executor).build()
+
+        # Hold a live reference to a streaming run without iterating it so that
+        # ``_is_run_active`` remains True (the active-run weakref still resolves).
+        active_stream = workflow.run(WorkflowMessage(data="hi", source_id="test"), stream=True)
+        try:
+            with pytest.raises(WorkflowException, match="Cannot create checkpoint while a workflow run is active"):
+                await workflow.create_checkpoint(storage)
+        finally:
+            # Drain the stream so the run completes cleanly and the active-run
+            # weakref is cleared; otherwise pytest's asyncio teardown can leak
+            # the unconsumed generator.
+            async for _ in active_stream:
+                pass
+
+    async def test_clears_runtime_storage_after_call(self, simple_executor: Executor) -> None:
+        """The runtime storage override must not leak past the call."""
+        storage = InMemoryCheckpointStorage()
+        workflow = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, simple_executor).build()
+
+        await workflow.create_checkpoint(storage)
+
+        assert workflow._runner.context.has_checkpointing() is False
+        assert workflow._runner.context._runtime_checkpoint_storage is None  # type: ignore[attr-defined]
+
+    async def test_clears_runtime_storage_after_failure(self, simple_executor: Executor) -> None:
+        """The runtime storage override must be cleared even if checkpoint creation fails."""
+        from unittest.mock import AsyncMock
+
+        storage = InMemoryCheckpointStorage()
+        workflow = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, simple_executor).build()
+
+        # The runner logs-and-swallows storage save errors, so a failed save
+        # surfaces as the "Failed to create checkpoint." path when
+        # ``previous_checkpoint_id`` remains ``None``. Either way, the
+        # ``finally`` cleanup must still clear the runtime override.
+        storage.save = AsyncMock(side_effect=RuntimeError("boom"))  # type: ignore[method-assign]
+
+        with pytest.raises(WorkflowCheckpointException, match="Failed to create checkpoint"):
+            await workflow.create_checkpoint(storage)
+
+        assert workflow._runner.context._runtime_checkpoint_storage is None  # type: ignore[attr-defined]
+
+    async def test_alters_lineage_for_next_checkpoint(self, simple_executor: Executor) -> None:
+        """A manually created checkpoint becomes the parent of the next checkpoint."""
+        storage = InMemoryCheckpointStorage()
+        workflow = (
+            WorkflowBuilder(start_executor=simple_executor, checkpoint_storage=storage)
+            .add_edge(simple_executor, simple_executor)
+            .build()
+        )
+
+        first_id = await workflow.create_checkpoint(None)
+        second_id = await workflow.create_checkpoint(None)
+
+        assert first_id != second_id
+        second = await storage.load(second_id)
+        assert second is not None
+        assert second.previous_checkpoint_id == first_id
+
+
+# endregion