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

[michaelxu2288]

Problem

Two bugs in StreamingTaskMessageContext (src/agentex/lib/core/services/adk/streaming.py).

1. Errors raised inside the context are swallowed. __aexit__ was return await self.close(), and close() returns the TaskMessage. A truthy return from __aexit__ suppresses the exception, so in

async with adk.streaming.streaming_task_message_context(task_id, initial) as ctx:
    await ctx.stream_update(...)
    await call_the_model()   # raises: rate limit, tool crash, bad JSON

the exception disappears. Execution continues after the block, the half-written message is persisted as DONE, and a Temporal activity reports success, so its retry policy never runs.

2. An explicit StreamTaskMessageDone is published twice. stream_update() published the Done and then called close(). _is_closed was still False at that point, so close() published a second Done. The message was also updated twice. And because close() drains the coalescing buffer before publishing, buffered deltas could land after the first Done.

Fix

• __aexit__ closes the context and returns False. On the error path the context is still closed and DONE is still persisted, but the exception is no longer suppressed. The inference-call context manager already follows this convention.
• A StreamTaskMessageDone passed to stream_update() now goes straight to close(), which reaps the buffer, publishes exactly one Done and persists once. Delta and Full handling are unchanged.

Verification

• uv run pytest -n 0 tests/lib/core/services/adk/test_streaming.py
  • Three new tests, run before the fix:
    • test_exception_inside_context_propagates: DID NOT RAISE.
    • test_context_is_still_closed_when_body_raises: DID NOT RAISE.
    • test_explicit_done_publishes_exactly_one_done: expected exactly one DONE publish, got 2.
  • After the fix: 37 passed.
• These suites still pass: tests/lib/adk, tests/lib/core/harness, tests/lib/test_claude_agents_hooks.py, and the openai_agents test_streaming_model.py.
• ruff check is clean on both files.

Confidence Score: 5/5

The PR appears safe to merge; the earlier Done-index issue is fixed, and this review found no new issue.

What we checked:

• Plain close keeps its Done index: No. StreamTaskMessageDone.index already defaults to None, so plain close still sends the same index value.

Summary

The streaming context now lets errors from its body reach the caller while still closing the message. Explicit DONE updates now finish after queued deltas and publish and save the final message once.

• Keeps cleanup running when a context body raises.
• Routes explicit DONE updates through the shared finish path and preserves their index.
• Adds tests for error propagation, cleanup, and single-DONE publishing.

_{Reviews (2) · Last reviewed commit: "fix(streaming): keep the caller's index ..."}