Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 3 additions & 5 deletions src/a2a/server/events/event_queue_v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down Expand Up @@ -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()
Expand Down
21 changes: 21 additions & 0 deletions tests/server/events/test_event_queue_v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import logging

from typing import Any
from unittest import mock

import pytest
import pytest_asyncio
Expand Down Expand Up @@ -266,6 +267,26 @@ 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)

span.record_exception.assert_not_called()
Comment thread
dorlugasigal marked this conversation as resolved.
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."""
Expand Down
Loading