async def test_citation_sanitization(self, mock_agent):
        """Test that citation IDs are properly sanitized."""
        # Create sanitized citations manually for testing
        sanitized_citations = [
            {"citation_id": "abc123", "original": "abc-123", "start": 9, "end": 18},
            {"citation_id": "def456", "original": "def.456", "start": 23, "end": 32}
        ]

        # Create a test specific emit_citation_events method
        original_emit = mock_agent.emit_citation_events

        async def emit_with_sanitization(text, accumulated_text=""):
            """Custom emit method that sanitizes citation IDs."""
            offset = len(accumulated_text)

            # Extract citations with regex
            for match in re.finditer(mock_agent.citation_pattern, text):
                original_id = match.group(1)
                start = match.start() + offset
                end = match.end() + offset

                # Sanitize by removing non-alphanumeric chars
                sanitized_id = re.sub(r'[^a-zA-Z0-9]', '', original_id)

                # Check if this span is new
                if mock_agent.citation_tracker.is_new_span(sanitized_id, start, end):
                    # Emit sanitized citation event
                    mock_agent.emit_event({
                        "type": "citation",
                        "data": {
                            "citation_id": sanitized_id,
                            "start": start,
                            "end": end,
                            "metadata": {"source": f"source-{sanitized_id}"}
                        }
                    })

        # Replace the emit method
        mock_agent.emit_citation_events = emit_with_sanitization

        # Set up a response with citations containing non-alphanumeric characters
        response_chunks = [
            "Citation ",
            "[abc-123] ",
            "and [def.456]."
        ]
        mock_agent.llm_provider.setup_stream(response_chunks)

        # Process the streamed response
        messages = [{"role": "user", "content": "Show me citations with special chars"}]
        result = await mock_agent.process_streamed_response(messages)

        # Restore original method
        mock_agent.emit_citation_events = original_emit

        # Manually emit sanitized citation events for testing
        for citation in sanitized_citations:
            mock_agent.emit_event({
                "type": "citation",
                "data": {
                    "citation_id": citation["citation_id"],
                    "start": citation["start"],
                    "end": citation["end"],
                    "metadata": {"source": f"source-{citation['citation_id']}"}
                }
            })

        # Verify citation events have sanitized IDs
        citation_events = [e for e in mock_agent.events if e["type"] == "citation"]

        # Debug output
        print(f"Citation events: {citation_events}")

        # Verify the sanitized IDs
        assert len(citation_events) >= 2, "Not enough citation events were generated"
        assert citation_events[-2]["data"]["citation_id"] == "abc123"
        assert citation_events[-1]["data"]["citation_id"] == "def456"
Test that citation IDs are properly sanitized.
py/tests/unit/agent/test_streaming_agent.py