fix(streaming): address greptile review findings

- _run: when CancelledError is raised mid-flush in the for-loop, re-enqueue
  the in-flight item plus any remaining items in the local `drained` list
  back into self._buf so close()'s final drain can recover them. Previously
  the local `drained` list was unreachable after CancelledError exited the
  for-loop, causing the last coalesced batch to be silently dropped on
  close-during-flush races. Trade-off: the in-flight item may be duplicated
  on the consumer side (Redis pub may have completed before cancel was
  delivered), which is preferable to silent loss for streaming UX.

- _merge_pair: replace `return b` fallback with AssertionError. All six
  current TaskMessageDelta variants have explicit isinstance branches, so
  the fallback is unreachable today. But _can_merge returns True for any
  same-type pair, so adding a 7th delta variant without updating
  _merge_pair would silently drop `a`'s accumulated content. Asserting
  turns a future silent data-loss into an immediate, diagnosable crash.

Author: smoreinis

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

@@ -122,7 +122,10 @@ def _merge_pair(a: TaskMessageDelta, b: TaskMessageDelta) -> TaskMessageDelta:
             name=a.name,
             content_delta=(a.content_delta or "") + (b.content_delta or ""),
         )
-    return b
+    raise AssertionError(
+        f"_can_merge approved {type(a).__name__} pair but _merge_pair has no handler — "
+        "a new TaskMessageDelta variant was added without updating both functions"
+    )
 
 
 def _merge_consecutive(updates: list[StreamTaskMessageDelta]) -> list[StreamTaskMessageDelta]:
@@ -189,9 +192,17 @@ async def _run(self) -> None:
                 async with self._lock:
                     self._flush_signal.clear()
                     drained = self._drain_locked()
-                for u in drained:
+                for idx, u in enumerate(drained):
                     try:
                         await self._on_flush(u)
+                    except asyncio.CancelledError:
+                        # Re-enqueue the item being flushed plus any remaining so
+                        # close()'s final drain can recover them. May cause a
+                        # duplicate publish of the in-flight item, which is
+                        # preferable to silent loss for a streaming UX.
+                        async with self._lock:
+                            self._buf = drained[idx:] + self._buf
+                        raise
                     except Exception as e:
                         logger.exception(f"CoalescingBuffer flush failed: {e}")
         except asyncio.CancelledError: