async def custom_arun(*args, **kwargs):
            """Custom arun that includes repeated citations for the same source."""
            # Setup like the original
            await custom_agent._setup(kwargs.get('system_instruction'))

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

            # Initialize payloads dict for tracking
            citation_payloads = {}

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

            # Create text with multiple citations to the same source
            response_content = "This text has multiple citations to the same source: [abc1234] and again here [abc1234]."

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

            # Manually extract and emit citation events
            # This is a simpler approach than the character-by-character approach
            citation_spans = extract_citation_spans(response_content)

            # 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 = custom_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
                    citation_payloads[cid] = citation_payload

                    # Track for persistence
                    custom_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 custom_agent._format_sse_event("citation", citation_event)

            # Add assistant message with citation metadata to conversation
            await custom_agent.conversation.add_message(
                Message(
                    role="assistant",
                    content=response_content,
                    metadata={"citations": custom_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
            }

            yield custom_agent._format_sse_event("agent.final_answer", final_evt_payload)

            # Signal the end of the SSE stream
            yield "event: done\ndata: {}\n\n"
Custom arun that includes repeated citations for the same source.
py/tests/unit/agent/test_agent_citations.py