async def test_final_answer_includes_consolidated_citations(mock_streaming_agent):
    """Test that the final answer includes consolidated citations."""
    messages = [Message(role="user", content="Test query")]

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

    # Look for final answer event in the output
    final_answer_events = [
        line for line in output
        if 'event: agent.final_answer' in line
    ]

    assert len(final_answer_events) > 0, "Final answer event should be emitted"

    # Parse the event to check for citations
    for event in final_answer_events:
        data_part = event.split('data: ')[1] if 'data: ' in event else event
        try:
            data = json.loads(data_part)
            if 'citations' in data:
                assert len(data['citations']) > 0, "Final answer should include citations"
                citation_ids = [citation.get('id') for citation in data['citations']]
                assert 'abc1234' in citation_ids or 'def5678' in citation_ids, "Known citation IDs should be included"
        except json.JSONDecodeError:
            continue
Test that the final answer includes consolidated citations.
py/tests/unit/agent/test_agent_citations_old.py