fix(streaming): stop suppressing context errors and publish DONE once

StreamingTaskMessageContext.__aexit__ returned close()'s TaskMessage,
and a truthy __aexit__ return value suppresses the exception. An error
raised inside `async with streaming_task_message_context(...)` was
therefore swallowed: the caller carried on past the block, the
half-written message was persisted as DONE, and a Temporal activity saw
a success and never retried. __aexit__ now closes the context and
returns False, matching the convention the inference-call context
manager already follows.

Separately, stream_update() published an explicit StreamTaskMessageDone
itself and then called close(), which published a second DONE because
_is_closed was still False. Two DONE frames reached the stream, and
since close() drains the coalescing buffer before publishing, buffered
deltas could land after the first one. The Done branch now delegates to
close(), which reaps the buffer, publishes exactly one DONE and
persists. Full and delta handling are unchanged.

Verified with `uv run pytest -q -n 0` on
tests/lib/core/services/adk/test_streaming.py: the three added
regression tests fail before the fix (two "DID NOT RAISE RuntimeError",
one "expected exactly one DONE publish, got 2") and the file is 37
passed after. tests/lib/adk, tests/lib/core/harness,
tests/lib/test_claude_agents_hooks.py and the openai_agents
test_streaming_model.py suite stay green.

Author: michaelxu2288

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

@@ -395,7 +395,14 @@ async def __aenter__(self) -> "StreamingTaskMessageContext":
         return await self.open()
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        return await self.close()
+        """Close the context, then let any exception from the body propagate.
+
+        Returning ``close()``'s truthy ``TaskMessage`` here would suppress it,
+        silently persisting a half-written message as DONE and hiding the
+        failure from the caller (and from Temporal's activity retries).
+        """
+        await self.close()
+        return False
 
     async def open(self) -> "StreamingTaskMessageContext":
         self._is_closed = False
@@ -484,7 +491,9 @@ async def stream_update(self, update: TaskMessageUpdate) -> TaskMessageUpdate |
 
         ``StreamTaskMessageDone`` and ``StreamTaskMessageFull`` updates always
         publish synchronously regardless of mode so consumers and persistence
-        stay in sync.
+        stay in sync. A Done delegates its publish to ``close()``, which drains
+        the buffer first and emits exactly one DONE; publishing it here as well
+        would put a second DONE on the stream, ahead of the buffered deltas.
         """
         if self._is_closed:
             raise ValueError("Context is already done")
@@ -501,6 +510,10 @@ async def stream_update(self, update: TaskMessageUpdate) -> TaskMessageUpdate |
                 await self._buffer.add(update)
                 return update
 
+        if isinstance(update, StreamTaskMessageDone):
+            await self.close()
+            return update
+
         # A Full ends the stream and supersedes buffered deltas. Drain and stop
         # the buffer BEFORE publishing the Full, so leftover deltas land in order
         # (deltas -> Full) instead of trailing the terminal Full as a stale
@@ -511,10 +524,7 @@ async def stream_update(self, update: TaskMessageUpdate) -> TaskMessageUpdate |
 
         result = await self._streaming_service.stream_update(update)
 
-        if isinstance(update, StreamTaskMessageDone):
-            await self.close()
-            return update
-        elif isinstance(update, StreamTaskMessageFull):
+        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,
 )
@@ -62,7 +63,8 @@ def _reasoning_summary(tm: TaskMessage, idx: int, s: str) -> StreamTaskMessageDe
     )
 
 
-async def _make_context(streaming_mode: str) -> tuple[StreamingTaskMessageContext, MagicMock, TaskMessage]:
+def _make_unopened_context(streaming_mode: str) -> tuple[StreamingTaskMessageContext, MagicMock, TaskMessage]:
+    """Wired-up context that has not been opened yet, for ``async with`` tests."""
     tm = TaskMessage(
         id="m1",
         task_id="t1",
@@ -81,6 +83,11 @@ async def _make_context(streaming_mode: str) -> tuple[StreamingTaskMessageContex
         streaming_service=svc,
         streaming_mode=streaming_mode,  # type: ignore[arg-type]
     )
+    return ctx, svc, tm
+
+
+async def _make_context(streaming_mode: str) -> tuple[StreamingTaskMessageContext, MagicMock, TaskMessage]:
+    ctx, svc, tm = _make_unopened_context(streaming_mode)
     await ctx.open()
     return ctx, svc, tm
 
@@ -597,3 +604,60 @@ async def test_full_is_terminal_publish_no_trailing_deltas(self) -> None:
         assert any(isinstance(u, StreamTaskMessageDelta) for u in published[:-1]), (
             "expected the buffered deltas to be published before the Full"
         )
+
+
+class TestContextDoesNotSuppressErrors:
+    """``__aexit__`` returned ``close()``'s TaskMessage, which is truthy, so an
+    exception raised inside ``async with`` was swallowed: the caller carried on
+    past the block, the half-written message was persisted DONE, and a Temporal
+    activity saw a success and never retried."""
+
+    @pytest.mark.asyncio
+    async def test_exception_inside_context_propagates(self) -> None:
+        ctx, _svc, tm = _make_unopened_context("off")
+
+        with pytest.raises(RuntimeError, match="model blew up"):
+            async with ctx as entered:
+                await entered.stream_update(_text(tm, "partial"))
+                raise RuntimeError("model blew up")
+
+    @pytest.mark.asyncio
+    async def test_context_is_still_closed_when_body_raises(self) -> None:
+        """Not suppressing must not mean leaking: DONE is still published and
+        persisted so consumers and the buffer ticker are not left hanging."""
+        ctx, svc, _tm = _make_unopened_context("off")
+
+        with pytest.raises(RuntimeError):
+            async with ctx:
+                raise RuntimeError("model blew up")
+
+        assert ctx._is_closed
+        published = [c.args[0] for c in svc.stream_update.await_args_list]
+        assert isinstance(published[-1], StreamTaskMessageDone)
+        update_kwargs = ctx._agentex_client.messages.update.call_args.kwargs
+        assert update_kwargs["streaming_status"] == "DONE"
+
+
+class TestExplicitDonePublishesOnce:
+    """An explicit ``StreamTaskMessageDone`` used to be published by
+    ``stream_update`` and then published a second time by the ``close()`` it
+    triggers, putting two DONE frames on the stream. Routing the terminal
+    publish through ``close()`` also keeps buffered deltas ahead of the DONE."""
+
+    @pytest.mark.asyncio
+    async def test_explicit_done_publishes_exactly_one_done(self) -> None:
+        ctx, svc, tm = await _make_context("coalesced")
+        await ctx.stream_update(_text(tm, "hello"))
+
+        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"expected exactly one DONE publish, got {len(dones)}"
+        assert published[-1] is dones[0], (
+            f"DONE must be the terminal publish; saw trailing {type(published[-1]).__name__} after it"
+        )
+        assert ctx._agentex_client.messages.update.call_count == 1
+        update_kwargs = ctx._agentex_client.messages.update.call_args.kwargs
+        assert update_kwargs["content"]["content"] == "hello"
+        assert update_kwargs["streaming_status"] == "DONE"