async def test_overlapping_citation_handling():
    """Test that overlapping citations are handled correctly."""
    # Create a custom agent configuration
    config = MagicMock()
    config.stream = True
    config.max_iterations = 3

    # Create providers
    llm_provider = MockLLMProvider(
        response_content="This is a test response with overlapping citations",
        citations=["abc1234", "def5678"]
    )
    db_provider = MockDatabaseProvider()

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

    # Replace the search results collector with our mock
    agent.search_results_collector = MockSearchResultsCollector({
        "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"}
        }
    })

    # Modify the arun method for overlapping citations
    original_arun = agent.arun

    async def custom_arun(*args, **kwargs):
        """Custom arun that includes overlapping citations."""
        # Setup like the original
        await agent._setup(kwargs.get('system_instruction'))

        messages = kwargs.get('messages', [])
        if messages:
            for m in messages:
                await agent.conversation.add_message(m)

        # Initialize citation tracker
        citation_tracker = CitationTracker()
        citation_payloads = {}

        # Track streaming citations for final persistence
        agent.streaming_citations = []

        # Create text with overlapping citations (citation spans that overlap)
        response_content = "This text has overlapping citations [abc1234] part of which [def5678] overlap."

        # Yield the message event
        yield agent._format_sse_event("message", {"content": response_content})

        # Manually create overlapping citation spans
        # For simplicity, we'll define the spans directly rather than using regex
        citation_spans = {
            "abc1234": [(30, 39)],  # This span includes "[abc1234]"
            "def5678": [(55, 64)]   # This span includes "[def5678]"
        }

        # Process the citations
        for cid, spans in citation_spans.items():
            for span in spans:
                # Mark as processed in the tracker
                citation_tracker.is_new_span(cid, span)

                # Look up the source document for this citation
                source_doc = agent.search_results_collector.find_by_short_id(cid)

                # Create citation payload
                citation_payload = {
                    "document_id": source_doc.get("document_id", f"doc_{cid}"),
                    "text": source_doc.get("text", f"This is document text for {cid}"),
                    "metadata": source_doc.get("metadata", {"source": f"source_{cid}"}),
                }

                # Store the payload by citation ID
                citation_payloads[cid] = citation_payload

                # Track for persistence
                agent.streaming_citations.append({
                    "id": cid,
                    "span": {"start": span[0], "end": span[1]},
                    "payload": citation_payload
                })

                # Emit citation event
                citation_event = {
                    "id": cid,
                    "object": "citation",
                    "span": {"start": span[0], "end": span[1]},
                    "payload": citation_payload
                }

                yield agent._format_sse_event("citation", citation_event)

        # Add assistant message with citation metadata to conversation
        await agent.conversation.add_message(
            Message(
                role="assistant",
                content=response_content,
                metadata={"citations": agent.streaming_citations}
            )
        )

        # Prepare consolidated citations for final answer
        consolidated_citations = []

        # Group citations by ID with all their spans
        for cid, spans in citation_tracker.get_all_spans().items():
            if cid in citation_payloads:
                consolidated_citations.append({
                    "id": cid,
                    "object": "citation",
                    "spans": [{"start": s[0], "end": s[1]} for s in spans],
                    "payload": citation_payloads[cid]
                })

        # Create and emit final answer event
        final_evt_payload = {
            "id": "msg_final",
            "object": "agent.final_answer",
            "generated_answer": response_content,
            "citations": consolidated_citations
        }

        # Emit final answer event
        yield agent._format_sse_event("agent.final_answer", final_evt_payload)

        # Signal the end of the SSE stream
        yield "event: done\ndata: {}\n\n"

    # Replace the arun method
    with patch.object(agent, 'arun', custom_arun):
        messages = [Message(role="user", content="Test query")]

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

        # Check that both citations were emitted
        citation_abc = any('abc1234' in event for event in output if 'event: citation' in event)
        citation_def = any('def5678' in event for event in output if 'event: citation' in event)

        assert citation_abc, "Citation abc1234 should be emitted"
        assert citation_def, "Citation def5678 should be emitted"

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

        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:
                    citation_ids = [citation.get('id') for citation in data['citations']]
                    assert 'abc1234' in citation_ids, "abc1234 should be in final answer citations"
                    assert 'def5678' in citation_ids, "def5678 should be in final answer citations"
            except json.JSONDecodeError:
                continue
Test that overlapping citations are handled correctly.
py/tests/unit/agent/test_agent_citations_old.py