fix(streaming): StreamTaskMessageFull closes the coalescing buffer

[eberki-scale]

Summary

A StreamTaskMessageFull ends the stream and marks the StreamingTaskMessageContext done, but it did not close the coalescing buffer. When the context later exits, __aexit__ → close() hits its if self._is_closed: return guard and never stops the buffer's background ticker — leaving an orphaned task per such stream. Separately, any deltas still buffered when the Full arrives could be published after the terminal Full, which a consumer treating Full as the final message reads as a stale duplicate tail.

Fix

Two small changes to StreamingTaskMessageContext in agentex/lib/core/services/adk/streaming.py:

1. stream_update — when a StreamTaskMessageFull arrives, drain and close the buffer before publishing the Full, so leftover deltas land in order (deltas → Full) and the ticker is stopped.
2. close() — reap the buffer before the _is_closed short-circuit, so a context already marked done (by a Full on another path) can never leave the ticker orphaned.

No change to CoalescingBuffer itself.

Tests

New TestFullMessageClosesBuffer:

• test_full_message_stops_ticker — after a Full, the buffer is reaped and its ticker task is done() (not orphaned).
• test_full_is_terminal_publish_no_trailing_deltas — buffered deltas publish before the Full; the Full is the terminal publish.

Full streaming suite: 33 passed (tests/lib/core/services/adk/test_streaming.py).

Scope

Deliberately narrow — just the StreamTaskMessageFull orphan + ordering fix, split out from the broader buffer work in #418 for easier review.

Greptile Summary

• Updates StreamingTaskMessageContext so a terminal StreamTaskMessageFull drains and closes the coalescing buffer before publishing the full message.
• Reaps the coalescing buffer before close() returns early for an already-closed context.
• Adds streaming tests covering ticker cleanup, terminal full-message ordering, and add/close race behavior.

Confidence Score: 4/5

The change fixes the intended buffer cleanup path, but the terminal full-message path still needs serialization with close().

The modified streaming code and focused tests make the intended behavior clear, and the remaining issue is localized to concurrent terminal handling.

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

T-Rex Logs

What T-Rex did

• Ran a focused async repro script against the real StreamingTaskMessageContext and StreamTaskMessageFull types, using test doubles for publishing and persistence, which reproduced the terminal update race.
• Inspected the before/after buffer and ticker state to confirm that the buffer was reaped and the ticker stopped after Full and close/exit.
• Verified the non-concurrent terminal publish order: Start -> Delta('AB') -> Full('AB'), with trailing deltas after Full cleared and exit code 0.
• Compared the base and head logs for the non-DONE short-circuit scenario and observed cleanup completed with the buffer closed, the original buffer cleared, and the lifecycle matching the expected behavior.
• Observed the add-close race path where the racing delta is not appended or flushed after the buffer is closed.

_{Ran code and verified through T-Rex}

Comments Outside Diff (1)

1. src/agentex/lib/core/services/adk/streaming.py, line 529-530 (link)

   Drain before Done

   StreamTaskMessageDone still publishes the terminal event before the coalescing buffer is drained. In coalesced mode, any buffered deltas remain in self._buffer when stream_update(update) sends Done; only afterwards does close() reap the buffer, so consumers can see Done followed by a delta. Because _is_closed is still false, close() can also emit a second Done. Please drain the buffer before publishing Done and avoid re-emitting it on this path.

   Artifacts

   
   

   Repro: async harness that buffers a coalesced delta before explicit Done

   • Contains supporting evidence from the run (text/x-python; charset=utf-8).

   Repro: failing run output showing Done before buffered delta and duplicate Done

   • Keeps the command output available without making the summary code-heavy.

   

   _{Ran code and verified through T-Rex}

   Prompt To Fix With AI

   This is a comment left during a code review.
   Path: src/agentex/lib/core/services/adk/streaming.py
   Line: 529-530
   
   Comment:
   **Drain before Done**
   
   `StreamTaskMessageDone` still publishes the terminal event before the coalescing buffer is drained. In coalesced mode, any buffered deltas remain in `self._buffer` when `stream_update(update)` sends `Done`; only afterwards does `close()` reap the buffer, so consumers can see `Done` followed by a delta. Because `_is_closed` is still false, `close()` can also emit a second `Done`. Please drain the buffer before publishing `Done` and avoid re-emitting it on this path.
   
   How can I resolve this? If you propose a fix, please make it concise.

   

Prompt To Fix All With AI

Fix the following 1 code review issue. Work through them one at a time, proposing concise fixes.

---

### Issue 1 of 1
src/agentex/lib/core/services/adk/streaming.py:509-512
**Serialize terminal updates**

`StreamTaskMessageFull` drains the buffer before publishing the terminal `Full`, but the context is still not marked closed or closing while that publish is awaiting. If `close()` runs in that window, it sees `_is_closed == False` and publishes/persists `DONE`; when this coroutine resumes, it then publishes the `Full` after `DONE`. Consumers can receive two terminal events in the wrong order, and the persisted message can be updated by two terminal paths. Please guard the `Full` path and `close()` with the same terminal state or a shared lock so only one terminal path can run.

_{Reviews (7): Last reviewed commit: "refactor(streaming): drop _is_closing co..." | Re-trigger Greptile}

Greptile also left 1 inline comment on this PR.

[smoreinis]

I saw that Greptile also commented about this and looks like it resolved its own comment but it looks to me like it's still possible that we could have a race here? e.g.

1. This code block closes the buffer and sets it to none
2. stream_update (or messages.update) below is being awaited
3. another delta comes in and skips this block since buffer is already none but _is_closed is not yet set to True

In this case, consumers will still see a delta after full since it just falls through to publishing directly. If this is something we care about (and my logic checks out and you also believe there's a race), we could have a terminal-in-progress flag like _is_closing set before awaiting buffer close to reject deltas?

[eberki-scale]

Good point, started introducing _is_closing, but it qot quite complex. For now, I would keep this PR more focused on the deterministic issue with the buffer leak, and make the more general imrovements in a follow-up PR (most likely #418 once that is rebased)

[smoreinis]

if i'm reading this correctly, this code block still has a latent bug similar to the one we fixed for the Full case above, right? curious about just abstracting the code above into something like _reap_buffer and calling that here as well for readability and more defensiveness going forward.

[eberki-scale]

Done did not have this problem, as the self.close() called here will close the buffer. close() is unfortunately specific to Done, so Full cannot call it directly, but agree that the common _reap_buffer can help. Added!

[smoreinis]

thanks for this fix - overall the PR is solid and I'm fine with it going in as-is.
just a couple of minor nits in the comments that I think could be nice to haves, but happy to leave these at the author discretion.

[eberki-scale]

@smoreinis thanks for the review!
commented on the specific comment as well - I think improving robustness with _is_closing is a great improvement, but it got quite complex and wanted to keep this PR focused on the deterministic buffer issue with the StreamTaskMessageFull closer. Concurrency issues could be improved later in the follow up #426 more generally.

[eberki-scale]

Done did not have this problem, as the self.close() called here will close the buffer. close() is unfortunately specific to Done, so Full cannot call it directly, but agree that the common _reap_buffer can help. Added!

[eberki-scale]

Good point, started introducing _is_closing, but it qot quite complex. For now, I would keep this PR more focused on the deterministic issue with the buffer leak, and make the more general imrovements in a follow-up PR (most likely #418 once that is rebased)

[smoreinis]

@eberki-scale this all makes sense! LGTM