async def test_citation_event_format(mock_streaming_agent):
    """Test that citation events follow the expected format."""
    messages = [Message(role="user", content="Test query")]

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

    # Extract citation events
    citation_events = [
        line for line in output
        if 'event: citation' in line
    ]

    assert len(citation_events) > 0, "Citation events should be emitted"

    # Check the format of each citation event
    for event in citation_events:
        # Should have 'event: citation' and 'data: {...}'
        assert 'event: citation' in event, "Event type should be 'citation'"
        assert 'data: ' in event, "Event should have data payload"

        # Parse the data payload
        data_part = event.split('data: ')[1] if 'data: ' in event else event
        try:
            data = json.loads(data_part)

            # Check required fields
            assert 'id' in data, "Citation event should have an 'id'"
            assert 'object' in data and data['object'] == 'citation', "Event object should be 'citation'"
            assert 'span' in data, "Citation event should have a 'span'"
            assert 'start' in data['span'] and 'end' in data['span'], "Span should have 'start' and 'end'"
            assert 'payload' in data, "Citation event should have a 'payload'"

            # Check payload fields
            assert 'document_id' in data['payload'], "Payload should have 'document_id'"
            assert 'text' in data['payload'], "Payload should have 'text'"
            assert 'metadata' in data['payload'], "Payload should have 'metadata'"

        except json.JSONDecodeError:
            pytest.fail(f"Citation event data is not valid JSON: {data_part}")
Test that citation events follow the expected format.
py/tests/unit/agent/test_agent_citations_old.py