fix(sdk): preserve queued events on stream close

Signed-off-by: Lanre Shittu <136805224+Shizoqua@users.noreply.github.com>
This commit is contained in:
Lanre Shittu
2026-09-03 16:47:23 +08:00
committed by Xubin Ren
parent 972cdde8da
commit 816a999cac
3 changed files with 27 additions and 7 deletions
+1 -1
View File
@@ -302,7 +302,7 @@ class Nanobot:
)) ))
raise raise
finally: finally:
emitter.close() await emitter.close()
task = asyncio.create_task(_run()) task = asyncio.create_task(_run())
return RunStream(task, queue) return RunStream(task, queue)
+2 -6
View File
@@ -158,15 +158,11 @@ class SDKStreamEmitter:
resuming=resuming, resuming=resuming,
)) ))
def close(self) -> None: async def close(self) -> None:
if self._closed: if self._closed:
return return
self._closed = True self._closed = True
if self._queue.full(): await self._queue.put(_STREAM_SENTINEL)
with suppress(asyncio.QueueEmpty):
self._queue.get_nowait()
with suppress(asyncio.QueueFull):
self._queue.put_nowait(_STREAM_SENTINEL)
class SDKStreamingHook(AgentHook): class SDKStreamingHook(AgentHook):
+24
View File
@@ -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