fix(streaming): mark _is_closing in close() and roll back on terminal failure

Greptile follow-ups on the _is_closing guard:
- close() sets _is_closing before reaping the buffer, so a delta racing a direct
  close() (e.g. via __aexit__) can't slip past the None buffer and publish after DONE.
- A failed terminal write (e.g. messages.update raising) rolls _is_closing back to
  False, so the still-open context doesn't silently drop later deltas.

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

Author: eberki-scale

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

@@ -439,41 +439,51 @@ async def close(self) -> TaskMessage:
         if not self.task_message:
             raise ValueError("Context not properly initialized - no task message")
 
-        # Reap the buffer (stopping its ticker) before the _is_closed
-        # short-circuit, so a context already marked done by a Full update can't
-        # leave the ticker orphaned. Draining here also lets consumers see the
-        # full delta sequence in order before DONE.
-        await self._reap_buffer()
+        # Mark closing before reaping so a concurrent coalesced delta can't slip
+        # past the now-None buffer and publish after DONE. Rolled back below if
+        # the terminal write fails, so a failed close doesn't wedge the context
+        # into silently dropping later deltas.
+        self._is_closing = True
+        try:
+            # Reap the buffer (stopping its ticker) before the _is_closed
+            # short-circuit, so a context already marked done by a Full update
+            # can't leave the ticker orphaned. Draining here also lets consumers
+            # see the full delta sequence in order before DONE.
+            await self._reap_buffer()
 
-        if self._is_closed:
-            return self.task_message  # Already done (buffer reaped above)
+            if self._is_closed:
+                return self.task_message  # Already done (buffer reaped above)
 
-        # Send the DONE event
-        done_event = StreamTaskMessageDone(
-            parent_task_message=self.task_message,
-            type="done",
-        )
-        await self._streaming_service.stream_update(done_event)
+            # Send the DONE event
+            done_event = StreamTaskMessageDone(
+                parent_task_message=self.task_message,
+                type="done",
+            )
+            await self._streaming_service.stream_update(done_event)
 
-        # Update the task message with the final content
-        has_deltas = (
-            self._delta_accumulator._accumulated_deltas
-            or self._delta_accumulator._reasoning_summaries
-            or self._delta_accumulator._reasoning_contents
-        )
-        if has_deltas:
-            self.task_message.content = self._delta_accumulator.convert_to_content()
+            # Update the task message with the final content
+            has_deltas = (
+                self._delta_accumulator._accumulated_deltas
+                or self._delta_accumulator._reasoning_summaries
+                or self._delta_accumulator._reasoning_contents
+            )
+            if has_deltas:
+                self.task_message.content = self._delta_accumulator.convert_to_content()
 
-        await self._agentex_client.messages.update(
-            task_id=self.task_id,
-            message_id=self.task_message.id,
-            content=self.task_message.content.model_dump(),
-            streaming_status="DONE",
-        )
+            await self._agentex_client.messages.update(
+                task_id=self.task_id,
+                message_id=self.task_message.id,
+                content=self.task_message.content.model_dump(),
+                streaming_status="DONE",
+            )
 
-        # Mark the context as done
-        self._is_closed = True
-        return self.task_message
+            # Mark the context as done
+            self._is_closed = True
+            return self.task_message
+        except Exception:
+            if not self._is_closed:
+                self._is_closing = False
+            raise
 
     async def stream_update(self, update: TaskMessageUpdate) -> TaskMessageUpdate | None:
         """Stream an update to the repository.
@@ -512,32 +522,38 @@ async def stream_update(self, update: TaskMessageUpdate) -> TaskMessageUpdate |
 
         # From here the update is terminal. Mark closing so a concurrent delta
         # can't fall through and publish after the terminal once the buffer is
-        # reaped below.
-        if isinstance(update, (StreamTaskMessageFull, StreamTaskMessageDone)):
+        # reaped below. Rolled back if the terminal path fails, so a failed
+        # terminal doesn't wedge the context into silently dropping later deltas.
+        is_terminal = isinstance(update, (StreamTaskMessageFull, StreamTaskMessageDone))
+        if is_terminal:
             self._is_closing = True
-
-        # 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
-        # duplicate tail. This also stops the ticker, which would otherwise be
-        # orphaned when __aexit__'s close() short-circuits on _is_closed.
-        if isinstance(update, StreamTaskMessageFull):
-            await self._reap_buffer()
-
-        result = await self._streaming_service.stream_update(update)
-
-        if isinstance(update, StreamTaskMessageDone):
-            await self.close()
-            return update
-        elif 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]
-                content=update.content.model_dump(),
-                streaming_status="DONE",
-            )
-            self._is_closed = True
-        return result
+        try:
+            # 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 duplicate tail. This also stops the ticker, which would
+            # otherwise be orphaned when __aexit__'s close() short-circuits.
+            if isinstance(update, StreamTaskMessageFull):
+                await self._reap_buffer()
+
+            result = await self._streaming_service.stream_update(update)
+
+            if isinstance(update, StreamTaskMessageDone):
+                await self.close()
+                return update
+            elif 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]
+                    content=update.content.model_dump(),
+                    streaming_status="DONE",
+                )
+                self._is_closed = True
+            return result
+        except Exception:
+            if not self._is_closed:
+                self._is_closing = False
+            raise
 
 
 class StreamingService:

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

@@ -621,3 +621,40 @@ async def test_late_delta_during_close_window_is_not_published_after_full(self)
         await ctx.stream_update(_text(tm, "late"))
 
         assert svc.stream_update.await_count == 0, "late delta published after the terminal Full"
+
+    @pytest.mark.asyncio
+    async def test_close_marks_closing_before_reaping_buffer(self) -> None:
+        """close() must set _is_closing before reaping, so a concurrent delta in
+        the post-reap / pre-_is_closed window is dropped, not published after DONE."""
+        ctx, _svc, tm = await _make_context("coalesced")
+        await ctx.stream_update(_text(tm, "hi"))
+
+        closing_at_reap = {}
+        orig_reap = ctx._reap_buffer
+
+        async def spy() -> None:
+            closing_at_reap["value"] = ctx._is_closing
+            await orig_reap()
+
+        ctx._reap_buffer = spy  # type: ignore[method-assign]
+        await ctx.close()
+        assert closing_at_reap.get("value") is True, "close() must mark _is_closing before reaping"
+
+    @pytest.mark.asyncio
+    async def test_failed_terminal_rolls_back_is_closing(self) -> None:
+        """If the terminal persistence fails, _is_closing must roll back so the
+        still-open context doesn't silently drop later coalesced deltas."""
+        ctx, _svc, tm = await _make_context("coalesced")
+        await ctx.stream_update(_text(tm, "hi"))
+        ctx._agentex_client.messages.update.side_effect = RuntimeError("persist boom")  # type: ignore[attr-defined]
+
+        full = StreamTaskMessageFull(
+            parent_task_message=tm,
+            content=TextContent(author="agent", content="hi", format="markdown"),
+            type="full",
+        )
+        with pytest.raises(RuntimeError):
+            await ctx.stream_update(full)
+
+        assert ctx._is_closed is False
+        assert ctx._is_closing is False, "failed terminal left _is_closing stuck on"