async def test_citation_consolidation_logic(mock_streaming_agent):
    """Test that citation consolidation properly groups spans by citation ID."""
    # Patch the get_all_spans method to return a controlled set of spans
    citation_tracker = CitationTracker()

    # Add spans for multiple citations
    citation_tracker.is_new_span("abc1234", (10, 20))
    citation_tracker.is_new_span("abc1234", (30, 40))
    citation_tracker.is_new_span("def5678", (50, 60))
    citation_tracker.is_new_span("ghi9012", (70, 80))
    citation_tracker.is_new_span("ghi9012", (90, 100))

    # Create a custom mock agent that uses our pre-populated citation tracker
    with patch('core.utils.CitationTracker', return_value=citation_tracker):
        # Create a fresh agent with our mocked citation tracker
        new_agent = mock_streaming_agent

        messages = [Message(role="user", content="Test query")]

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

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

        # Verify consolidation in final answer
        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:
                    # There should be at least 2 citations (from our mock agent implementation)
                    assert len(data['citations']) >= 2, "Should include multiple citation objects"

                    # Check spans for each citation
                    for citation in data['citations']:
                        cid = citation.get('id')
                        if cid == 'abc1234':
                            # Spans should be consolidated for abc1234
                            spans = citation.get('spans', [])
                            assert len(spans) >= 1, f"Citation {cid} should have spans"
            except json.JSONDecodeError:
                continue
Test that citation consolidation properly groups spans by citation ID.
py/tests/unit/agent/test_agent_citations.py