diff --git a/src/a2a/server/events/event_queue_v2.py b/src/a2a/server/events/event_queue_v2.py index 3386a732e..fc65de173 100644 --- a/src/a2a/server/events/event_queue_v2.py +++ b/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() diff --git a/tests/server/events/test_event_queue_v2.py b/tests/server/events/test_event_queue_v2.py index 27bceea4c..0484396f9 100644 --- a/tests/server/events/test_event_queue_v2.py +++ b/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."""