diff --git a/nanobot/nanobot.py b/nanobot/nanobot.py index 6c9951ed8..3c1654cb0 100644 --- a/nanobot/nanobot.py +++ b/nanobot/nanobot.py @@ -302,7 +302,7 @@ class Nanobot: )) raise finally: - emitter.close() + await emitter.close() task = asyncio.create_task(_run()) return RunStream(task, queue) diff --git a/nanobot/sdk/streaming.py b/nanobot/sdk/streaming.py index 1abfd64fc..b8036bd67 100644 --- a/nanobot/sdk/streaming.py +++ b/nanobot/sdk/streaming.py @@ -158,15 +158,11 @@ class SDKStreamEmitter: resuming=resuming, )) - def close(self) -> None: + async def close(self) -> None: if self._closed: return self._closed = True - if self._queue.full(): - with suppress(asyncio.QueueEmpty): - self._queue.get_nowait() - with suppress(asyncio.QueueFull): - self._queue.put_nowait(_STREAM_SENTINEL) + await self._queue.put(_STREAM_SENTINEL) class SDKStreamingHook(AgentHook): diff --git a/tests/test_sdk_streaming.py b/tests/test_sdk_streaming.py new file mode 100644 index 000000000..a292315d8 --- /dev/null +++ b/tests/test_sdk_streaming.py @@ -0,0 +1,24 @@ +"""Tests for SDK streaming primitives.""" + +import asyncio + +import pytest + +from nanobot.sdk.streaming import SDKStreamEmitter +from nanobot.sdk.types import STREAM_EVENT_TEXT_DELTA, StreamEvent + + +@pytest.mark.asyncio +async def test_close_preserves_events_when_queue_is_full(): + queue: asyncio.Queue[StreamEvent | object] = asyncio.Queue(maxsize=1) + emitter = SDKStreamEmitter(queue) + event = StreamEvent(type=STREAM_EVENT_TEXT_DELTA, delta="kept") + await emitter.emit(event) + + close_task = asyncio.create_task(emitter.close()) + await asyncio.sleep(0) + + assert not close_task.done() + assert queue.get_nowait() is event + await close_task + assert queue.qsize() == 1