fix(streaming): drain coalesced buffer before terminal Done, dedupe Done

The Done path published the terminal update before reaping the buffer, so a
delta still in the coalescing window arrived after Done, and close() emitted a
second Done — consumers saw a stale trailing delta and a duplicate terminal.
Drain/stop the buffer before the terminal for Done as well as Full, and let
close() be the sole Done emitter so it publishes exactly once.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: eberki-scale

src/agentex/lib/core/services/adk/streaming.py

@@ -505,20 +505,23 @@ async def stream_update(self, update: TaskMessageUpdate) -> TaskMessageUpdate |
                 await self._buffer.add(update)
                 return update
 
-        # Drain and stop the buffer BEFORE publishing the Full, so leftover
-        # deltas land in order (deltas -> Full); publishing them after the
-        # terminal Full would look like a stale duplicate tail. This also stops
-        # the ticker so it isn't orphaned past __aexit__'s close().
-        if isinstance(update, StreamTaskMessageFull) and self._buffer is not None:
+        # Terminal Done/Full updates must drain and stop the buffer BEFORE the
+        # terminal event reaches consumers, so leftover deltas land in order
+        # (deltas -> terminal) instead of trailing it as a stale duplicate tail.
+        # This also stops the ticker so it isn't orphaned past __aexit__'s close().
+        if isinstance(update, (StreamTaskMessageDone, StreamTaskMessageFull)) and self._buffer is not None:
             await self._buffer.close()
             self._buffer = None
 
-        result = await self._streaming_service.stream_update(update)
-
         if isinstance(update, StreamTaskMessageDone):
+            # close() publishes the single terminal Done, persists, and marks the
+            # context closed — don't publish here too, that would duplicate it.
             await self.close()
             return update
-        elif isinstance(update, StreamTaskMessageFull):
+
+        result = await self._streaming_service.stream_update(update)
+
+        if isinstance(update, StreamTaskMessageFull):
             await self._agentex_client.messages.update(
                 task_id=self.task_id,
                 message_id=update.parent_task_message.id,  # type: ignore[union-attr]

tests/lib/core/services/adk/test_streaming.py

@@ -23,6 +23,7 @@
     ReasoningSummaryDelta,
 )
 from agentex.types.task_message_update import (
+    StreamTaskMessageDone,
     StreamTaskMessageFull,
     StreamTaskMessageDelta,
 )
@@ -621,6 +622,27 @@ async def test_full_is_terminal_publish_no_trailing_deltas(self) -> None:
             "expected the buffered deltas to be published before the Full"
         )
 
+    @pytest.mark.asyncio
+    async def test_done_is_single_terminal_publish_no_trailing_deltas(self) -> None:
+        # Same guarantee as Full: buffered deltas publish BEFORE the terminal
+        # Done, and Done is published exactly once (not duplicated by close()).
+        ctx, svc, tm = await _make_context("coalesced")
+        # "alpha" flushes immediately; "beta" stays buffered in the window.
+        await ctx.stream_update(_text(tm, "alpha"))
+        await ctx.stream_update(_text(tm, "beta"))
+
+        await ctx.stream_update(StreamTaskMessageDone(parent_task_message=tm, type="done"))
+
+        published = [c.args[0] for c in svc.stream_update.await_args_list]
+        dones = [u for u in published if isinstance(u, StreamTaskMessageDone)]
+        assert len(dones) == 1, f"Done must publish exactly once, saw {len(dones)}"
+        assert isinstance(published[-1], StreamTaskMessageDone), (
+            f"Done must be the terminal publish; saw trailing {type(published[-1]).__name__}"
+        )
+        assert any(isinstance(u, StreamTaskMessageDelta) for u in published[:-1]), (
+            "expected the buffered deltas to be published before the Done"
+        )
+
     @pytest.mark.asyncio
     async def test_close_reaps_buffer_even_if_already_marked_closed(self) -> None:
         # close() must stop the ticker even when _is_closed is already set.