Python: Add more specific exceptions to Workflow (#3188)

* Add more specifc workflow exceptions

* Fix tests

* AI comments

* Misc

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

@@ -3,7 +3,14 @@
 import pytest
 from typing_extensions import Never
 
-from agent_framework import WorkflowBuilder, WorkflowContext, WorkflowRunState, WorkflowStatusEvent, handler
+from agent_framework import (
+    WorkflowBuilder,
+    WorkflowCheckpointException,
+    WorkflowContext,
+    WorkflowRunState,
+    WorkflowStatusEvent,
+    handler,
+)
 from agent_framework._workflows._checkpoint import InMemoryCheckpointStorage
 from agent_framework._workflows._executor import Executor
 
@@ -43,7 +50,7 @@ async def test_resume_fails_when_graph_mismatch() -> None:
     # Build a structurally different workflow (different finish executor id)
     mismatched_workflow = build_workflow(storage, finish_id="finish_alt")
 
-    with pytest.raises(ValueError, match="Workflow graph has changed"):
+    with pytest.raises(WorkflowCheckpointException, match="Workflow graph has changed"):
         _ = [
             event
             async for event in mismatched_workflow.run_stream(

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

@@ -31,6 +31,7 @@ from agent_framework import (
     TextContent,
     Workflow,
     WorkflowCheckpoint,
+    WorkflowCheckpointException,
     WorkflowContext,
     WorkflowEvent,
     WorkflowOutputEvent,
@@ -341,7 +342,8 @@ async def test_magentic_orchestrator_round_limit_produces_partial_result():
         events.append(ev)
 
     idle_status = next(
-        (e for e in events if isinstance(e, WorkflowStatusEvent) and e.state == WorkflowRunState.IDLE), None
+        (e for e in events if isinstance(e, WorkflowStatusEvent) and e.state == WorkflowRunState.IDLE),
+        None,
     )
     assert idle_status is not None
     # Check that we got workflow output via WorkflowOutputEvent
@@ -584,7 +586,9 @@ async def _collect_agent_responses_setup(participant: AgentProtocol) -> list[Cha
         if isinstance(ev, AgentRunUpdateEvent):
             captured.append(
                 ChatMessage(
-                    role=ev.data.role or Role.ASSISTANT, text=ev.data.text or "", author_name=ev.data.author_name
+                    role=ev.data.role or Role.ASSISTANT,
+                    text=ev.data.text or "",
+                    author_name=ev.data.author_name,
                 )
             )
 
@@ -604,7 +608,9 @@ async def test_agent_executor_invoke_with_assistants_client_messages():
     assert any((m.author_name == agent.name and "ok" in (m.text or "")) for m in captured)
 
 
-async def _collect_checkpoints(storage: InMemoryCheckpointStorage) -> list[WorkflowCheckpoint]:
+async def _collect_checkpoints(
+    storage: InMemoryCheckpointStorage,
+) -> list[WorkflowCheckpoint]:
     checkpoints = await storage.list_checkpoints()
     assert checkpoints
     checkpoints.sort(key=lambda cp: cp.timestamp)
@@ -719,7 +725,7 @@ async def test_magentic_checkpoint_resume_rejects_participant_renames():
         .build()
     )
 
-    with pytest.raises(ValueError, match="Workflow graph has changed"):
+    with pytest.raises(WorkflowCheckpointException, match="Workflow graph has changed"):
         async for _ in renamed_workflow.run_stream(
             checkpoint_id=target_checkpoint.checkpoint_id,  # type: ignore[reportUnknownMemberType]
         ):
@@ -760,7 +766,8 @@ async def test_magentic_stall_and_reset_reach_limits():
         events.append(ev)
 
     idle_status = next(
-        (e for e in events if isinstance(e, WorkflowStatusEvent) and e.state == WorkflowRunState.IDLE), None
+        (e for e in events if isinstance(e, WorkflowStatusEvent) and e.state == WorkflowRunState.IDLE),
+        None,
     )
     assert idle_status is not None
     output_event = next((e for e in events if isinstance(e, WorkflowOutputEvent)), None)
@@ -799,7 +806,10 @@ async def test_magentic_checkpoint_runtime_overrides_buildtime() -> None:
     """Test that runtime checkpoint storage overrides build-time configuration."""
     import tempfile
 
-    with tempfile.TemporaryDirectory() as temp_dir1, tempfile.TemporaryDirectory() as temp_dir2:
+    with (
+        tempfile.TemporaryDirectory() as temp_dir1,
+        tempfile.TemporaryDirectory() as temp_dir2,
+    ):
         from agent_framework._workflows._checkpoint import FileCheckpointStorage
 
         buildtime_storage = FileCheckpointStorage(temp_dir1)

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

@@ -10,15 +10,21 @@ from agent_framework import (
     AgentResponse,
     Executor,
     WorkflowContext,
+    WorkflowConvergenceException,
     WorkflowEvent,
     WorkflowOutputEvent,
+    WorkflowRunnerException,
     WorkflowRunState,
     WorkflowStatusEvent,
     handler,
 )
 from agent_framework._workflows._edge import SingleEdgeGroup
 from agent_framework._workflows._runner import Runner
-from agent_framework._workflows._runner_context import InProcRunnerContext, Message, RunnerContext
+from agent_framework._workflows._runner_context import (
+    InProcRunnerContext,
+    Message,
+    RunnerContext,
+)
 from agent_framework._workflows._shared_state import SharedState
 
 
@@ -52,7 +58,10 @@ def test_create_runner():
         SingleEdgeGroup(executor_b.id, executor_a.id),
     ]
 
-    executors: dict[str, Executor] = {executor_a.id: executor_a, executor_b.id: executor_b}
+    executors: dict[str, Executor] = {
+        executor_a.id: executor_a,
+        executor_b.id: executor_b,
+    }
 
     runner = Runner(edge_groups, executors, shared_state=SharedState(), ctx=InProcRunnerContext())
 
@@ -70,7 +79,10 @@ async def test_runner_run_until_convergence():
         SingleEdgeGroup(executor_b.id, executor_a.id),
     ]
 
-    executors: dict[str, Executor] = {executor_a.id: executor_a, executor_b.id: executor_b}
+    executors: dict[str, Executor] = {
+        executor_a.id: executor_a,
+        executor_b.id: executor_b,
+    }
     shared_state = SharedState()
     ctx = InProcRunnerContext()
 
@@ -90,6 +102,9 @@ async def test_runner_run_until_convergence():
 
     assert result is not None and result == 10
 
+    # iteration count shouldn't be reset after convergence
+    assert runner._iteration == 10  # type: ignore
+
 
 async def test_runner_run_until_convergence_not_completed():
     """Test running the runner with a simple workflow."""
@@ -102,7 +117,10 @@ async def test_runner_run_until_convergence_not_completed():
         SingleEdgeGroup(executor_b.id, executor_a.id),
     ]
 
-    executors: dict[str, Executor] = {executor_a.id: executor_a, executor_b.id: executor_b}
+    executors: dict[str, Executor] = {
+        executor_a.id: executor_a,
+        executor_b.id: executor_b,
+    }
     shared_state = SharedState()
     ctx = InProcRunnerContext()
 
@@ -114,7 +132,10 @@ async def test_runner_run_until_convergence_not_completed():
         shared_state,  # shared_state
         ctx,  # runner_context
     )
-    with pytest.raises(RuntimeError, match="Runner did not converge after 5 iterations."):
+    with pytest.raises(
+        WorkflowConvergenceException,
+        match="Runner did not converge after 5 iterations.",
+    ):
         async for event in runner.run_until_convergence():
             assert not isinstance(event, WorkflowStatusEvent) or event.state != WorkflowRunState.IDLE
 
@@ -130,7 +151,10 @@ async def test_runner_already_running():
         SingleEdgeGroup(executor_b.id, executor_a.id),
     ]
 
-    executors: dict[str, Executor] = {executor_a.id: executor_a, executor_b.id: executor_b}
+    executors: dict[str, Executor] = {
+        executor_a.id: executor_a,
+        executor_b.id: executor_b,
+    }
     shared_state = SharedState()
     ctx = InProcRunnerContext()
 
@@ -143,7 +167,7 @@ async def test_runner_already_running():
         ctx,  # runner_context
     )
 
-    with pytest.raises(RuntimeError, match="Runner is already running."):
+    with pytest.raises(WorkflowRunnerException, match="Runner is already running."):
 
         async def _run():
             async for _ in runner.run_until_convergence():

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

@@ -25,7 +25,9 @@ from agent_framework import (
     Role,
     TextContent,
     WorkflowBuilder,
+    WorkflowCheckpointException,
     WorkflowContext,
+    WorkflowConvergenceException,
     WorkflowEvent,
     WorkflowOutputEvent,
     WorkflowRunState,
@@ -143,7 +145,7 @@ async def test_workflow_run_stream_not_completed():
         .build()
     )
 
-    with pytest.raises(RuntimeError):
+    with pytest.raises(WorkflowConvergenceException):
         async for _ in workflow.run_stream(NumberMessage(data=0)):
             pass
 
@@ -181,7 +183,7 @@ async def test_workflow_run_not_completed():
         .build()
     )
 
-    with pytest.raises(RuntimeError):
+    with pytest.raises(WorkflowConvergenceException):
         await workflow.run(NumberMessage(data=0))
 
 
@@ -289,7 +291,9 @@ async def test_workflow_with_checkpointing_enabled(simple_executor: Executor):
         assert result is not None
 
 
-async def test_workflow_checkpointing_not_enabled_for_external_restore(simple_executor: Executor):
+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 = (
@@ -308,7 +312,9 @@ async def test_workflow_checkpointing_not_enabled_for_external_restore(simple_ex
         assert "either provide checkpoint_storage parameter" in str(e)
 
 
-async def test_workflow_run_stream_from_checkpoint_no_checkpointing_enabled(simple_executor: Executor):
+async def test_workflow_run_stream_from_checkpoint_no_checkpointing_enabled(
+    simple_executor: Executor,
+):
     # Build workflow WITHOUT checkpointing
     workflow = (
         WorkflowBuilder()
@@ -327,7 +333,9 @@ async def test_workflow_run_stream_from_checkpoint_no_checkpointing_enabled(simp
         assert "either provide checkpoint_storage parameter" in str(e)
 
 
-async def test_workflow_run_stream_from_checkpoint_invalid_checkpoint(simple_executor: Executor):
+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)
@@ -345,12 +353,14 @@ async def test_workflow_run_stream_from_checkpoint_invalid_checkpoint(simple_exe
         try:
             async for _ in workflow.run_stream(checkpoint_id="nonexistent_checkpoint_id"):
                 pass
-            raise AssertionError("Expected RuntimeError to be raised")
-        except RuntimeError as e:
-            assert "Failed to restore from checkpoint" in str(e)
+            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(simple_executor: Executor):
+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)
@@ -416,7 +426,9 @@ async def test_workflow_run_from_checkpoint_non_streaming(simple_executor: Execu
         assert hasattr(result, "get_outputs")  # Should have WorkflowRunResult methods
 
 
-async def test_workflow_run_stream_from_checkpoint_with_responses(simple_executor: Executor):
+async def test_workflow_run_stream_from_checkpoint_with_responses(
+    simple_executor: Executor,
+):
     """Test that workflow can be resumed from checkpoint with pending RequestInfoEvents."""
     with tempfile.TemporaryDirectory() as temp_dir:
         storage = FileCheckpointStorage(temp_dir)
@@ -475,7 +487,9 @@ class StateTrackingExecutor(Executor):
 
     @handler
     async def handle_message(
-        self, message: StateTrackingMessage, ctx: WorkflowContext[StateTrackingMessage, list[str]]
+        self,
+        message: StateTrackingMessage,
+        ctx: WorkflowContext[StateTrackingMessage, list[str]],
     ) -> None:
         """Handle the message and track it in shared state."""
         # Get existing messages from shared state
@@ -537,7 +551,9 @@ async def test_workflow_multiple_runs_no_state_collision():
         assert outputs1[0] != outputs3[0]
 
 
-async def test_workflow_checkpoint_runtime_only_configuration(simple_executor: Executor):
+async def test_workflow_checkpoint_runtime_only_configuration(
+    simple_executor: Executor,
+):
     """Test that checkpointing can be configured ONLY at runtime, not at build time."""
     with tempfile.TemporaryDirectory() as temp_dir:
         storage = FileCheckpointStorage(temp_dir)
@@ -574,12 +590,20 @@ async def test_workflow_checkpoint_runtime_only_configuration(simple_executor: E
             checkpoint_id=resume_checkpoint.checkpoint_id, checkpoint_storage=storage
         )
         assert result_resumed is not None
-        assert result_resumed.get_final_state() in (WorkflowRunState.IDLE, WorkflowRunState.IDLE_WITH_PENDING_REQUESTS)
+        assert result_resumed.get_final_state() in (
+            WorkflowRunState.IDLE,
+            WorkflowRunState.IDLE_WITH_PENDING_REQUESTS,
+        )
 
 
-async def test_workflow_checkpoint_runtime_overrides_buildtime(simple_executor: Executor):
+async def test_workflow_checkpoint_runtime_overrides_buildtime(
+    simple_executor: Executor,
+):
     """Test that runtime checkpoint storage overrides build-time configuration."""
-    with tempfile.TemporaryDirectory() as temp_dir1, tempfile.TemporaryDirectory() as temp_dir2:
+    with (
+        tempfile.TemporaryDirectory() as temp_dir1,
+        tempfile.TemporaryDirectory() as temp_dir2,
+    ):
         buildtime_storage = FileCheckpointStorage(temp_dir1)
         runtime_storage = FileCheckpointStorage(temp_dir2)
 
@@ -740,7 +764,10 @@ async def test_workflow_concurrent_execution_prevention():
     await asyncio.sleep(0.01)
 
     # Try to start a second concurrent execution - this should fail
-    with pytest.raises(RuntimeError, match="Workflow is already running. Concurrent executions are not allowed."):
+    with pytest.raises(
+        RuntimeError,
+        match="Workflow is already running. Concurrent executions are not allowed.",
+    ):
         await workflow.run(NumberMessage(data=0))
 
     # Wait for the first task to complete
@@ -773,7 +800,10 @@ async def test_workflow_concurrent_execution_prevention_streaming():
     await asyncio.sleep(0.02)
 
     # Try to start a second concurrent execution - this should fail
-    with pytest.raises(RuntimeError, match="Workflow is already running. Concurrent executions are not allowed."):
+    with pytest.raises(
+        RuntimeError,
+        match="Workflow is already running. Concurrent executions are not allowed.",
+    ):
         await workflow.run(NumberMessage(data=0))
 
     # Wait for the first task to complete
@@ -803,10 +833,16 @@ async def test_workflow_concurrent_execution_prevention_mixed_methods():
     await asyncio.sleep(0.02)  # Let it start
 
     # Try different execution methods - all should fail
-    with pytest.raises(RuntimeError, match="Workflow is already running. Concurrent executions are not allowed."):
+    with pytest.raises(
+        RuntimeError,
+        match="Workflow is already running. Concurrent executions are not allowed.",
+    ):
         await workflow.run(NumberMessage(data=0))
 
-    with pytest.raises(RuntimeError, match="Workflow is already running. Concurrent executions are not allowed."):
+    with pytest.raises(
+        RuntimeError,
+        match="Workflow is already running. Concurrent executions are not allowed.",
+    ):
         async for _ in workflow.run_stream(NumberMessage(data=0)):
             break
 
@@ -923,7 +959,9 @@ async def test_workflow_run_parameter_validation(simple_executor: Executor) -> N
             pass
 
 
-async def test_workflow_run_stream_parameter_validation(simple_executor: Executor) -> None:
+async def test_workflow_run_stream_parameter_validation(
+    simple_executor: Executor,
+) -> None:
     """Test run_stream() specific parameter validation scenarios."""
     workflow = WorkflowBuilder().add_edge(simple_executor, simple_executor).set_start_executor(simple_executor).build()