feat: ensure Full message closes coalescing buffer to prevent stale duplicates

Author: eberki-scale

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

@@ -518,21 +518,23 @@ async def stream_update(self, update: TaskMessageUpdate) -> TaskMessageUpdate |
                 await self._buffer.add(update)
                 return update
 
+        # A full message supersedes the streamed deltas and ends the stream.
+        # Drain and stop the coalescing buffer BEFORE publishing the Full, so any
+        # leftover buffered deltas land on the stream in order (deltas -> Full)
+        # rather than after the terminal Full — a consumer treating Full as the
+        # final message would otherwise see those trailing deltas as a stale
+        # duplicate tail. Closing here also stops the ticker, so it can't be
+        # orphaned when __aexit__'s close() later short-circuits on _is_closed.
+        if isinstance(update, 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):
             await self.close()
             return update
         elif isinstance(update, StreamTaskMessageFull):
-            # A full message supersedes any buffered deltas and ends the stream.
-            # Close the coalescing buffer (stopping its ticker) BEFORE marking
-            # the context done — otherwise __aexit__'s close() early-returns on
-            # _is_closed and the ticker is never stopped, polling forever and
-            # leaking CPU (mirrors the StreamTaskMessageDone branch, which closes
-            # via self.close()).
-            if self._buffer is not None:
-                await self._buffer.close()
-                self._buffer = None
             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

@@ -613,6 +613,35 @@ async def test_full_message_stops_ticker(self) -> None:
         assert ctx._buffer is None, "Full message left the buffer un-closed"
         assert task.done(), "coalescing-buffer ticker still running after Full (orphaned)"
 
+    @pytest.mark.asyncio
+    async def test_full_is_terminal_publish_no_trailing_deltas(self) -> None:
+        # Leftover buffered deltas must be drained BEFORE the Full hits the
+        # stream (deltas -> Full), never after it — a consumer treating Full as
+        # the final message would see a trailing delta as a stale duplicate tail.
+        ctx, svc, tm = await _make_context("coalesced")
+        # First delta flushes immediately; the second stays in the coalescing
+        # window, so it is still buffered when the Full arrives.
+        await ctx.stream_update(_text(tm, "alpha"))
+        await ctx.stream_update(_text(tm, "beta"))
+
+        full = StreamTaskMessageFull(
+            parent_task_message=tm,
+            content=TextContent(author="agent", content="alphabeta", format="markdown"),
+            type="full",
+        )
+        await ctx.stream_update(full)
+
+        # Every publish (delta flushes + the Full) goes through the service mock.
+        published = [c.args[0] for c in svc.stream_update.await_args_list]
+        assert published, "nothing was published"
+        assert published[-1] is full, (
+            f"Full must be the terminal publish; saw trailing "
+            f"{type(published[-1]).__name__} after it (stale duplicate tail)"
+        )
+        assert any(isinstance(u, StreamTaskMessageDelta) for u in published[:-1]), (
+            "expected the buffered deltas to be published before the Full"
+        )
+
     @pytest.mark.asyncio
     async def test_close_reaps_buffer_even_if_already_marked_closed(self) -> None:
         # Defense-in-depth: if any path marks the context closed without closing