Fix WorkflowAgent to include thread convo history. Enable checkpointing. (#2774)

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

@@ -9,7 +9,9 @@ from agent_framework import (
     AgentRunResponse,
     AgentRunResponseUpdate,
     AgentRunUpdateEvent,
+    AgentThread,
     ChatMessage,
+    ChatMessageStore,
     Executor,
     FunctionApprovalRequestContent,
     FunctionApprovalResponseContent,
@@ -75,6 +77,31 @@ class RequestingExecutor(Executor):
         await ctx.add_event(AgentRunUpdateEvent(executor_id=self.id, data=update))
 
 
+class ConversationHistoryCapturingExecutor(Executor):
+    """Executor that captures the received conversation history for verification."""
+
+    def __init__(self, id: str):
+        super().__init__(id=id)
+        self.received_messages: list[ChatMessage] = []
+
+    @handler
+    async def handle_message(self, messages: list[ChatMessage], ctx: WorkflowContext[list[ChatMessage]]) -> None:
+        # Capture all received messages
+        self.received_messages = list(messages)
+
+        # Count messages by role for the response
+        message_count = len(messages)
+        response_text = f"Received {message_count} messages"
+
+        response_message = ChatMessage(role=Role.ASSISTANT, contents=[TextContent(text=response_text)])
+
+        streaming_update = AgentRunResponseUpdate(
+            contents=[TextContent(text=response_text)], role=Role.ASSISTANT, message_id=str(uuid.uuid4())
+        )
+        await ctx.add_event(AgentRunUpdateEvent(executor_id=self.id, data=streaming_update))
+        await ctx.send_message([response_message])
+
+
 class TestWorkflowAgent:
     """Test cases for WorkflowAgent end-to-end functionality."""
 
@@ -257,6 +284,105 @@ class TestWorkflowAgent:
         with pytest.raises(ValueError, match="Workflow's start executor cannot handle list\\[ChatMessage\\]"):
             workflow.as_agent()
 
+    async def test_thread_conversation_history_included_in_workflow_run(self) -> None:
+        """Test that conversation history from thread is included when running WorkflowAgent.
+
+        This verifies that when a thread with existing messages is provided to agent.run(),
+        the workflow receives the complete conversation history (thread history + new messages).
+        """
+        # Create an executor that captures all received messages
+        capturing_executor = ConversationHistoryCapturingExecutor(id="capturing")
+        workflow = WorkflowBuilder().set_start_executor(capturing_executor).build()
+        agent = WorkflowAgent(workflow=workflow, name="Thread History Test Agent")
+
+        # Create a thread with existing conversation history
+        history_messages = [
+            ChatMessage(role=Role.USER, text="Previous user message"),
+            ChatMessage(role=Role.ASSISTANT, text="Previous assistant response"),
+        ]
+        message_store = ChatMessageStore(messages=history_messages)
+        thread = AgentThread(message_store=message_store)
+
+        # Run the agent with the thread and a new message
+        new_message = "New user question"
+        await agent.run(new_message, thread=thread)
+
+        # Verify the executor received both history AND new message
+        assert len(capturing_executor.received_messages) == 3
+
+        # Verify the order: history first, then new message
+        assert capturing_executor.received_messages[0].text == "Previous user message"
+        assert capturing_executor.received_messages[1].text == "Previous assistant response"
+        assert capturing_executor.received_messages[2].text == "New user question"
+
+    async def test_thread_conversation_history_included_in_workflow_stream(self) -> None:
+        """Test that conversation history from thread is included when streaming WorkflowAgent.
+
+        This verifies that run_stream also includes thread history.
+        """
+        # Create an executor that captures all received messages
+        capturing_executor = ConversationHistoryCapturingExecutor(id="capturing_stream")
+        workflow = WorkflowBuilder().set_start_executor(capturing_executor).build()
+        agent = WorkflowAgent(workflow=workflow, name="Thread Stream Test Agent")
+
+        # Create a thread with existing conversation history
+        history_messages = [
+            ChatMessage(role=Role.SYSTEM, text="You are a helpful assistant"),
+            ChatMessage(role=Role.USER, text="Hello"),
+            ChatMessage(role=Role.ASSISTANT, text="Hi there!"),
+        ]
+        message_store = ChatMessageStore(messages=history_messages)
+        thread = AgentThread(message_store=message_store)
+
+        # Stream from the agent with the thread and a new message
+        async for _ in agent.run_stream("How are you?", thread=thread):
+            pass
+
+        # Verify the executor received all messages (3 from history + 1 new)
+        assert len(capturing_executor.received_messages) == 4
+
+        # Verify the order
+        assert capturing_executor.received_messages[0].text == "You are a helpful assistant"
+        assert capturing_executor.received_messages[1].text == "Hello"
+        assert capturing_executor.received_messages[2].text == "Hi there!"
+        assert capturing_executor.received_messages[3].text == "How are you?"
+
+    async def test_empty_thread_works_correctly(self) -> None:
+        """Test that an empty thread (no message store) works correctly."""
+        capturing_executor = ConversationHistoryCapturingExecutor(id="empty_thread_test")
+        workflow = WorkflowBuilder().set_start_executor(capturing_executor).build()
+        agent = WorkflowAgent(workflow=workflow, name="Empty Thread Test Agent")
+
+        # Create an empty thread
+        thread = AgentThread()
+
+        # Run with the empty thread
+        await agent.run("Just a new message", thread=thread)
+
+        # Should only receive the new message
+        assert len(capturing_executor.received_messages) == 1
+        assert capturing_executor.received_messages[0].text == "Just a new message"
+
+    async def test_checkpoint_storage_passed_to_workflow(self) -> None:
+        """Test that checkpoint_storage parameter is passed through to the workflow."""
+        from agent_framework import InMemoryCheckpointStorage
+
+        capturing_executor = ConversationHistoryCapturingExecutor(id="checkpoint_test")
+        workflow = WorkflowBuilder().set_start_executor(capturing_executor).build()
+        agent = WorkflowAgent(workflow=workflow, name="Checkpoint Test Agent")
+
+        # Create checkpoint storage
+        checkpoint_storage = InMemoryCheckpointStorage()
+
+        # Run with checkpoint storage enabled
+        async for _ in agent.run_stream("Test message", checkpoint_storage=checkpoint_storage):
+            pass
+
+        # Drain workflow events to get checkpoint
+        # The workflow should have created checkpoints
+        checkpoints = await checkpoint_storage.list_checkpoints(workflow.id)
+        assert len(checkpoints) > 0, "Checkpoints should have been created when checkpoint_storage is provided"
+
 
 class TestWorkflowAgentMergeUpdates:
     """Test cases specifically for the WorkflowAgent.merge_updates static method."""