fix(harness): close the turn's event source when delivery stops early

The delivery adapters (auto_send for the async push path, yield_events
and UnifiedEmitter.yield_turn for the sync yield path) closed their
streaming contexts and flushed span derivation in `finally`, but never
closed the event source they were iterating. Every tap relies on that
close: the Claude Code and Codex taps close their raw line iterator in a
`finally`, which is what terminates the CLI subprocess in the scaffolds.

When a turn is cancelled while delivery is awaiting the backend
(stream_update or a context close) rather than the source, or a sync
client disconnects mid-stream, the source stays suspended at a yield.
The turn object keeps a reference to its event generator, so garbage
collection does not finalize it while the caller still holds the turn,
and the CLI subprocess and its stdout handle leak.

Both adapters now close the source in `finally`, after the existing
cleanup and guarded so a failing aclose cannot mask the original
exception. yield_turn closes its delivery generator explicitly instead
of leaving it to async-generator finalization. Plain async iterators
without aclose and already-exhausted generators are unaffected.

Three new tests (cancel during a blocked stream_update; early close of
yield_events; early close of yield_turn with a turn that pins its
generator) fail on main with `assert [] == [True]` and pass here.
tests/lib/core/harness and tests/lib/adk: 571 passed, 1 skipped.

Author: michaelxu2288

src/agentex/lib/core/harness/auto_send.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import contextlib
 from typing import Any, AsyncIterator
 from datetime import datetime
 
@@ -46,7 +47,13 @@ async def auto_send(
     Index-keyed routing: each Start(index=i) opens a context stored in
     ctx_map[i]; Delta(index=i) routes to ctx_map.get(i); Done(index=i) closes
     and removes ctx_map[i]. Events with index is None are skipped. The finally
-    block closes all remaining open contexts.
+    block closes all remaining open contexts, then closes `events` itself (when
+    it exposes `aclose`) so delivery stopping early — cancelled while awaiting
+    the backend rather than the source — tears the tap down instead of leaving it
+    suspended at a yield: the turn object pins its event generator, so GC does
+    not rescue it and the harness subprocess leaks. Closing an exhausted
+    generator is a no-op, and a failure there is suppressed so it cannot mask the
+    original exception.
 
     final_text last-segment semantics: a new Start(TextContent) resets
     final_text_parts so that multi-step turns return the LAST text segment.
@@ -148,9 +155,15 @@ async def _close_all() -> None:
                     pass
 
     finally:
-        await _close_all()
-        if deriver is not None and tracer is not None:
-            for signal in deriver.flush():
-                await tracer.handle(signal)
+        try:
+            await _close_all()
+            if deriver is not None and tracer is not None:
+                for signal in deriver.flush():
+                    await tracer.handle(signal)
+        finally:
+            aclose = getattr(events, "aclose", None)
+            if aclose is not None:
+                with contextlib.suppress(Exception):
+                    await aclose()
 
     return TurnResult(final_text="".join(final_text_parts), usage=usage or TurnUsage())

src/agentex/lib/core/harness/emitter.py

@@ -53,9 +53,19 @@ def __init__(
             self.tracer = None
 
     async def yield_turn(self, turn: HarnessTurn) -> AsyncGenerator[StreamTaskMessage, None]:
-        """Sync HTTP ACP delivery: forward events, trace as side effect."""
-        async for event in yield_events(turn.events, tracer=self.tracer):
-            yield event
+        """Sync HTTP ACP delivery: forward events, trace as side effect.
+
+        The finally closes the delivery generator, which closes the turn's event
+        source in turn, so a consumer that stops early (client disconnect) tears
+        the tap down here rather than leaving it to async-generator finalization,
+        which never runs promptly while the turn still pins its generator.
+        """
+        delivery = yield_events(turn.events, tracer=self.tracer)
+        try:
+            async for event in delivery:
+                yield event
+        finally:
+            await delivery.aclose()
 
     async def auto_send_turn(self, turn: HarnessTurn, created_at: datetime | None = None) -> TurnResult:
         """Async/temporal delivery: push to the task stream, return TurnResult.

src/agentex/lib/core/harness/yield_delivery.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import contextlib
 from typing import AsyncIterator, AsyncGenerator
 
 from agentex.lib.core.harness.types import StreamTaskMessage
@@ -17,6 +18,13 @@ async def yield_events(
 
     For sync HTTP ACP agents that yield events back over the response. When
     `tracer` is None, this is a pure passthrough.
+
+    The finally also closes `events` (when it exposes `aclose`), so a consumer
+    that stops early — a client disconnect closes this generator — tears the tap
+    down instead of leaving it suspended at a yield: the turn object pins its
+    event generator, so GC does not rescue it and the harness subprocess leaks.
+    Closing an exhausted generator is a no-op, and a failure there is suppressed
+    so it cannot mask the original exception.
     """
     deriver = SpanDeriver() if tracer is not None else None
     try:
@@ -26,6 +34,12 @@ async def yield_events(
                     await tracer.handle(signal)
             yield event
     finally:
-        if deriver is not None and tracer is not None:
-            for signal in deriver.flush():
-                await tracer.handle(signal)
+        try:
+            if deriver is not None and tracer is not None:
+                for signal in deriver.flush():
+                    await tracer.handle(signal)
+        finally:
+            aclose = getattr(events, "aclose", None)
+            if aclose is not None:
+                with contextlib.suppress(Exception):
+                    await aclose()

tests/lib/core/harness/test_auto_send.py

@@ -9,6 +9,8 @@
 This mirrors _langgraph_async.py lines 62-78 and 100-127.
 """
 
+import asyncio
+from typing import override
 from datetime import datetime
 
 import pytest
@@ -478,3 +480,116 @@ async def test_auto_send_created_at_forwarded():
     await auto_send(_gen(events), task_id="task1", tracer=None, streaming=streaming, created_at=dt)
 
     assert all(ts == dt for ts in streaming.recorded_created_at)
+
+
+class _BlockingCtx(_FakeCtx):
+    """A context whose stream_update never returns (a stalled backend).
+
+    Sets `blocked` once stream_update is awaited so the test can cancel exactly
+    while delivery is suspended on the backend rather than on the source.
+    """
+
+    def __init__(self, sink, content_type, initial_content, blocked):
+        super().__init__(sink, content_type, initial_content)
+        self.blocked = blocked
+
+    @override
+    async def stream_update(self, update):
+        self.sink.append(("update", update))
+        self.blocked.set()
+        await asyncio.Event().wait()
+
+
+class _BlockingStreaming(_FakeStreaming):
+    """_FakeStreaming whose contexts block forever inside stream_update."""
+
+    def __init__(self):
+        super().__init__()
+        self.blocked = asyncio.Event()
+
+    @override
+    def streaming_task_message_context(self, task_id, initial_content, streaming_mode="coalesced", created_at=None):
+        ctype = getattr(initial_content, "type", None)
+        self.sink.append(("ctx", ctype))
+        self.recorded_created_at.append(created_at)
+        return _BlockingCtx(self.sink, ctype, initial_content, self.blocked)
+
+
+class _PlainAsyncIterator:
+    """An async iterator with no aclose (the AsyncIterator contract minimum)."""
+
+    def __init__(self, events):
+        self._events = iter(events)
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        try:
+            return next(self._events)
+        except StopIteration:
+            raise StopAsyncIteration from None
+
+
+@pytest.mark.asyncio
+async def test_auto_send_closes_source_when_cancelled_mid_delivery():
+    """Cancelling delivery while the backend blocks must close the event source.
+
+    The turn object pins its event generator, so GC cannot rescue it: when
+    auto_send returns with the source still suspended at a yield, the tap's
+    finally never runs and the harness CLI subprocess leaks. The source is held
+    by a local here for the whole test, and nothing calls gc.collect().
+    """
+    streaming = _BlockingStreaming()
+    closed: list[bool] = []
+
+    async def _recording_source():
+        try:
+            yield StreamTaskMessageStart(
+                type="start",
+                index=0,
+                content=TextContent(type="text", author="agent", content=""),
+            )
+            yield StreamTaskMessageDelta(
+                type="delta",
+                index=0,
+                delta=TextDelta(type="text", text_delta="hi"),
+            )
+            yield StreamTaskMessageDone(type="done", index=0)
+        finally:
+            closed.append(True)
+
+    source = _recording_source()
+    task = asyncio.create_task(auto_send(source, task_id="task1", tracer=None, streaming=streaming))
+    await asyncio.wait_for(streaming.blocked.wait(), timeout=5)
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    assert closed == [True]
+    assert ("close", "text") in [(s[0], s[1]) for s in streaming.sink]
+
+
+@pytest.mark.asyncio
+async def test_auto_send_accepts_source_without_aclose():
+    """A plain async iterator (no aclose) must deliver exactly as before."""
+    streaming = _FakeStreaming()
+    events = [
+        StreamTaskMessageStart(
+            type="start",
+            index=0,
+            content=TextContent(type="text", author="agent", content=""),
+        ),
+        StreamTaskMessageDelta(
+            type="delta",
+            index=0,
+            delta=TextDelta(type="text", text_delta="Hi"),
+        ),
+        StreamTaskMessageDone(type="done", index=0),
+    ]
+    result = await auto_send(_PlainAsyncIterator(events), task_id="task1", tracer=None, streaming=streaming)
+
+    assert result.final_text == "Hi"
+    kinds = [s[0] for s in streaming.sink]
+    assert kinds.count("open") == 1
+    assert kinds.count("close") == 1

tests/lib/core/harness/test_emitter.py

@@ -140,3 +140,54 @@ async def test_emitter_auto_send_turn_reads_usage_after_exhaustion():
     result = await emitter.auto_send_turn(turn)
     assert result.usage == real_usage
     assert result.usage.input_tokens == 11 and result.usage.total_tokens == 33
+
+
+class _PinnedTurn:
+    """A turn that pins its event generator, as the real CLI taps do.
+
+    `closed` records that the generator's finally ran (what terminates the CLI
+    subprocess in the scaffolds).
+    """
+
+    def __init__(self, events_list):
+        self._events_list = events_list
+        self.closed: list[bool] = []
+        self._gen = None
+
+    @property
+    def events(self):
+        if self._gen is None:
+            self._gen = self._stream()
+        return self._gen
+
+    async def _stream(self):
+        try:
+            for e in self._events_list:
+                yield e
+        finally:
+            self.closed.append(True)
+
+    def usage(self):
+        return TurnUsage()
+
+
+@pytest.mark.asyncio
+async def test_emitter_yield_turn_closes_source_on_early_close():
+    """Closing the delivery generator must close the turn's event source.
+
+    No gc.collect() and no sleep: the turn keeps the generator referenced, and
+    async-generator finalization is too late for a disconnected client anyway.
+    """
+    events = [
+        StreamTaskMessageStart(type="start", index=0, content=TextContent(type="text", author="agent", content="")),
+        StreamTaskMessageDelta(type="delta", index=0, delta=TextDelta(type="text", text_delta="hi")),
+        StreamTaskMessageDone(type="done", index=0),
+    ]
+    turn = _PinnedTurn(events)
+    emitter = UnifiedEmitter(task_id="t", trace_id=None, parent_span_id=None)
+    gen = emitter.yield_turn(turn)
+    first = await gen.__anext__()
+    await gen.aclose()
+
+    assert first.index == 0
+    assert turn.closed == [True]

tests/lib/core/harness/test_yield_delivery.py

@@ -76,3 +76,30 @@ async def test_flush_runs_on_early_close():
     await gen.aclose()  # triggers the finally -> flush()
     assert fake.started_names == ["Bash"]
     assert fake.ended_outputs == [None]  # flush closed the unpaired span (incomplete, no output)
+
+
+@pytest.mark.asyncio
+async def test_source_closed_when_consumer_closes_early():
+    """Closing the delivery generator must close the upstream event source.
+
+    A client disconnect (or any early break) closes the generator handed to the
+    caller. The source is pinned by the turn object, so if it is left suspended
+    at a yield the tap's finally never runs and the CLI subprocess leaks. No
+    gc.collect() here: the source stays referenced for the whole test.
+    """
+    closed: list[bool] = []
+
+    async def _recording_source():
+        try:
+            yield StreamTaskMessageDone(type="done", index=0)
+            yield StreamTaskMessageDone(type="done", index=1)
+        finally:
+            closed.append(True)
+
+    source = _recording_source()
+    gen = yield_events(source, tracer=None)
+    first = await gen.__anext__()
+    await gen.aclose()
+
+    assert first.index == 0
+    assert closed == [True]