async def test_agent_handles_multiple_messages():
    """Test agent handles conversation with multiple messages."""
    # Create mock config
    config = MagicMock()
    config.stream = True

    # Create mock providers
    llm_provider = MockLLMProvider(
        response_content="This is a response to multiple messages",
        citations=[]
    )
    db_provider = MockDatabaseProvider()

    # Create mock search results collector
    search_results = {
        "abc1234": {
            "document_id": "doc_abc1234",
            "text": "This is document text for abc1234",
            "metadata": {"source": "source_abc1234"}
        },
        "def5678": {
            "document_id": "doc_def5678",
            "text": "This is document text for def5678",
            "metadata": {"source": "source_def5678"}
        }
    }
    search_results_collector = MockSearchResultsCollector(search_results)

    # Create agent
    agent = MockR2RStreamingAgent(
        database_provider=db_provider,
        llm_provider=llm_provider,
        config=config,
        rag_generation_config=GenerationConfig(model="test/model")
    )

    # Set the search results collector
    agent.search_results_collector = search_results_collector

    # Test with multiple messages
    messages = [
        Message(role="system", content="You are a helpful assistant"),
        Message(role="user", content="First question"),
        Message(role="assistant", content="First answer"),
        Message(role="user", content="Follow-up question")
    ]

    # Run the agent
    stream = agent.arun(messages=messages)
    output = await collect_stream_output(stream)

    # Verify response
    message_events = [line for line in output if 'event: message' in line]
    assert len(message_events) > 0, "Message event should be emitted"

    # After running, check that conversation has the new assistant response
    # Note: MockR2RStreamingAgent._setup adds a default system message
    # and then our messages are added, plus the agent's response
    assert len(agent.conversation.messages) == 6, "Conversation should have correct number of messages"

    # The last message should be the assistant's response
    assert agent.conversation.messages[-1].role == "assistant", "Last message should be from assistant"

    # We should have two system messages (default + our custom one)
    system_messages = [m for m in agent.conversation.messages if m.role == "system"]
    assert len(system_messages) == 2, "Should have two system messages"
Test agent handles conversation with multiple messages.
py/tests/unit/agent/test_agent.py