Add reset to workflow

python/packages/orchestrations/tests/test_group_chat.py

@@ -1097,3 +1097,139 @@ def test_group_chat_orchestrator_factory_invalid_return_type():
 
 
 # endregion
+
+
+# region Reset
+
+
+async def test_base_orchestrator_reset_clears_conversation_and_round_index() -> None:
+    """reset() clears the conversation history and the round counter."""
+    from agent_framework.orchestrations import GroupChatOrchestrator
+
+    from agent_framework_orchestrations._base_group_chat_orchestrator import ParticipantRegistry
+
+    selector = make_sequence_selector()
+    orchestrator = GroupChatOrchestrator(
+        id="orch",
+        participant_registry=ParticipantRegistry([]),
+        selection_func=selector,
+        max_rounds=2,
+    )
+    orchestrator._full_conversation = [Message(role="user", contents=["hi"], author_name="user")]
+    orchestrator._round_index = 4
+
+    await orchestrator.reset()
+
+    assert orchestrator._full_conversation == []
+    assert orchestrator._round_index == 0
+
+
+async def test_base_orchestrator_reset_invokes_pattern_state_hook() -> None:
+    """reset() calls _reset_pattern_state() so subclasses can clean up their own state."""
+    from agent_framework.orchestrations import GroupChatOrchestrator
+
+    from agent_framework_orchestrations._base_group_chat_orchestrator import ParticipantRegistry
+
+    selector = make_sequence_selector()
+
+    class TrackingOrchestrator(GroupChatOrchestrator):
+        reset_calls: int = 0
+
+        def _reset_pattern_state(self) -> None:
+            type(self).reset_calls += 1
+
+    orchestrator = TrackingOrchestrator(
+        id="orch",
+        participant_registry=ParticipantRegistry([]),
+        selection_func=selector,
+        max_rounds=2,
+    )
+
+    await orchestrator.reset()
+    await orchestrator.reset()
+
+    assert TrackingOrchestrator.reset_calls == 2
+
+
+async def test_agent_based_orchestrator_reset_clears_cache_and_rotates_session() -> None:
+    """When the session was not supplied by the caller, reset() rotates the session and clears the cache."""
+    from agent_framework.orchestrations import AgentBasedGroupChatOrchestrator
+
+    from agent_framework_orchestrations._base_group_chat_orchestrator import ParticipantRegistry
+
+    agent = cast(Agent, StubManagerAgent())
+    orchestrator = AgentBasedGroupChatOrchestrator(
+        agent=agent,
+        participant_registry=ParticipantRegistry([]),
+        max_rounds=2,
+    )
+    original_session = orchestrator._session
+    orchestrator._cache = [Message(role="assistant", contents=["x"], author_name="agent")]
+    orchestrator._full_conversation = [Message(role="user", contents=["x"], author_name="user")]
+    orchestrator._round_index = 3
+
+    await orchestrator.reset()
+
+    assert orchestrator._cache == []
+    assert orchestrator._full_conversation == []
+    assert orchestrator._round_index == 0
+    assert orchestrator._session is not original_session
+
+
+async def test_agent_based_orchestrator_reset_warns_when_session_supplied(caplog: pytest.LogCaptureFixture) -> None:
+    """When the caller supplied a session, reset() preserves it and logs a warning."""
+    import logging
+
+    from agent_framework.orchestrations import AgentBasedGroupChatOrchestrator
+
+    from agent_framework_orchestrations._base_group_chat_orchestrator import ParticipantRegistry
+
+    agent = cast(Agent, StubManagerAgent())
+    supplied_session = agent.create_session()
+    orchestrator = AgentBasedGroupChatOrchestrator(
+        agent=agent,
+        participant_registry=ParticipantRegistry([]),
+        session=supplied_session,
+        max_rounds=2,
+    )
+    orchestrator._cache = [Message(role="assistant", contents=["x"], author_name="agent")]
+
+    with caplog.at_level(logging.WARNING, logger="agent_framework_orchestrations._group_chat"):
+        await orchestrator.reset()
+
+    assert orchestrator._cache == []
+    # The caller-owned session must be preserved.
+    assert orchestrator._session is supplied_session
+    warnings = [
+        r for r in caplog.records if r.levelno == logging.WARNING and "Session was supplied by the caller" in r.message
+    ]
+    assert warnings, f"expected a warning about caller-supplied session, got: {[r.message for r in caplog.records]}"
+
+
+async def test_workflow_reset_resets_group_chat_orchestrator() -> None:
+    """End-to-end: workflow.reset_for_new_run() resets the orchestrator's conversation state."""
+    selector = make_sequence_selector()
+    alpha = StubAgent("alpha", "ack from alpha")
+    beta = StubAgent("beta", "ack from beta")
+
+    workflow = GroupChatBuilder(
+        participants=[alpha, beta],
+        max_rounds=2,
+        selection_func=selector,
+        orchestrator_name="manager",
+    ).build()
+
+    async for _ in workflow.run("first task", stream=True):
+        pass
+
+    orchestrator = cast(BaseGroupChatOrchestrator, workflow.executors[GroupChatBuilder.DEFAULT_ORCHESTRATOR_ID])
+    assert orchestrator._full_conversation, "orchestrator should have accumulated conversation after first run"
+    assert orchestrator._round_index > 0
+
+    await workflow.reset_for_new_run()
+
+    assert orchestrator._full_conversation == []
+    assert orchestrator._round_index == 0
+
+
+# endregion