fix: make event queue sink removal idempotent

src/a2a/server/events/event_queue_v2.py

@@ -179,9 +179,9 @@ async def tap(
             return sink
 
     async def remove_sink(self, sink: 'EventQueueSink') -> None:
-        """Removes a sink from the source's internal list."""
+        """Removes a sink from the source's internal list if present."""
         async with self._lock:
-            self._sinks.remove(sink)
+            self._sinks.discard(sink)
 
     async def enqueue_event(self, event: Event) -> None:
         """Enqueues an event to this queue and all its children."""
@@ -349,9 +349,7 @@ async def close(self, immediate: bool = False) -> None:
             self._is_closed = True
             self._queue.shutdown(immediate=immediate)
 
-        # Ignore KeyError (close have to be idempotent).
-        with contextlib.suppress(KeyError):
-            await self._parent.remove_sink(self)
+        await self._parent.remove_sink(self)
 
         if not immediate:
             await self._queue.join()

tests/server/events/test_event_queue_v2.py

@@ -2,6 +2,7 @@
 import logging
 
 from typing import Any
+from unittest import mock
 
 import pytest
 import pytest_asyncio
@@ -266,6 +267,27 @@ async def test_close_idempotent(event_queue: EventQueueSource) -> None:
     assert event_queue.is_closed() is True
 
 
+@pytest.mark.asyncio
+async def test_sink_close_idempotent_does_not_record_exception() -> None:
+    span = mock.MagicMock()
+    tracer = mock.MagicMock()
+    tracer.start_as_current_span.return_value.__enter__.return_value = span
+    tracer.start_as_current_span.return_value.__exit__.return_value = False
+
+    with mock.patch('opentelemetry.trace.get_tracer', return_value=tracer):
+        event_queue = EventQueueSource(create_default_sink=False)
+        sink = await event_queue.tap()
+
+        await sink.close(immediate=True)
+        await sink.close(immediate=True)
+        await event_queue.close(immediate=True)
+
+    tracer.start_as_current_span.assert_called()
+    span.record_exception.assert_not_called()
+    for call in span.set_status.call_args_list:
+        assert 'description' not in call.kwargs
+
+
 @pytest.mark.asyncio
 async def test_is_closed_reflects_state(event_queue: EventQueueSource) -> None:
     """Test that is_closed() returns the correct state before and after closing."""