From 4fc484159716630e6c935333710f8c0f1a8c6d13 Mon Sep 17 00:00:00 2001 From: Ehsan-hash Date: Sat, 17 Oct 2026 00:16:12 -0700 Subject: [PATCH 1/2] test(sender): reproduce sends stranded after loop exit Developed with AI assistance using Codex. Regression tests intentionally fail until the sender lifecycle fix follows. --- tests/test_connection_recovery.py | 56 ++++++++++ tests/test_message_sender.py | 166 ++++++++++++++++++++++++++++++ 2 files changed, 222 insertions(+) create mode 100644 tests/test_message_sender.py diff --git a/tests/test_connection_recovery.py b/tests/test_connection_recovery.py index 67d58d3..4e0a769 100644 --- a/tests/test_connection_recovery.py +++ b/tests/test_connection_recovery.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import contextlib import json from typing import Any from unittest.mock import AsyncMock, MagicMock @@ -29,6 +30,20 @@ def _make_connection( return conn, reader +class _FailingWriter: + def __init__(self) -> None: + self.drain_started = asyncio.Event() + self.release = asyncio.Event() + + def write(self, _payload: bytes) -> None: + pass + + async def drain(self) -> None: + self.drain_started.set() + await self.release.wait() + raise BrokenPipeError("writer failed") + + @pytest.mark.asyncio async def test_receive_loop_handles_oversized_frame(caplog: pytest.LogCaptureFixture) -> None: conn, reader = _make_connection(limit=128) @@ -149,3 +164,44 @@ def tracking_process(message: dict[str, Any]) -> None: await conn.close() assert processed == ["survivor"] + + +@pytest.mark.asyncio +async def test_send_request_removes_pending_after_writer_failure() -> None: + reader = asyncio.StreamReader() + writer = _FailingWriter() + conn = Connection(_noop_handler, writer, reader, listening=False) + requests = [ + asyncio.create_task(conn.send_request("first"), name="request-first"), + asyncio.create_task(conn.send_request("second"), name="request-second"), + ] + + try: + await asyncio.wait_for(writer.drain_started.wait(), timeout=1) + for _ in range(100): + if conn._transport._sender._queue.qsize() == 1: # type: ignore[attr-defined] + break + await asyncio.sleep(0) + else: + raise AssertionError("second request was not queued") + + writer.release.set() + done, pending = await asyncio.wait(requests, timeout=1) + try: + assert not pending, f"requests still pending: {pending}" + for request in done: + with pytest.raises(BrokenPipeError, match="writer failed"): + request.result() + assert conn._pending == {} + finally: + for request in pending: + request.cancel() + await asyncio.gather(*requests, return_exceptions=True) + finally: + writer.release.set() + for request in requests: + if not request.done(): + request.cancel() + await asyncio.gather(*requests, return_exceptions=True) + with contextlib.suppress(BrokenPipeError): + await conn.close() diff --git a/tests/test_message_sender.py b/tests/test_message_sender.py new file mode 100644 index 0000000..a17ca3d --- /dev/null +++ b/tests/test_message_sender.py @@ -0,0 +1,166 @@ +from __future__ import annotations + +import asyncio +import json +from typing import Any + +import pytest + +from acp.task import MessageSender, TaskSupervisor + + +class _ControlledWriter: + def __init__(self, *, error: BaseException | None = None) -> None: + self.error = error + self.drain_started = asyncio.Event() + self.release = asyncio.Event() + self.writes: list[bytes] = [] + + def write(self, payload: bytes) -> None: + self.writes.append(payload) + + async def drain(self) -> None: + self.drain_started.set() + await self.release.wait() + if self.error is not None: + raise self.error + + +async def _wait_for_queue_size(sender: MessageSender, size: int) -> None: + async def queue_reached_size() -> None: + while sender._queue.qsize() < size: + await asyncio.sleep(0) + + await asyncio.wait_for(queue_reached_size(), timeout=1) + + +async def _assert_tasks_finish(tasks: list[asyncio.Task[Any]]) -> None: + _done, pending = await asyncio.wait(tasks, timeout=1) + try: + assert not pending, f"send tasks still pending: {pending}" + finally: + for task in pending: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + +async def _close_supervisor(supervisor: TaskSupervisor, tasks: list[asyncio.Task[Any]]) -> None: + await supervisor.shutdown() + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_writer_failure_settles_active_queued_and_late_sends() -> None: + error = BrokenPipeError("writer failed") + writer = _ControlledWriter(error=error) + supervisor = TaskSupervisor(source="test") + sender = MessageSender(writer, supervisor) + tasks: list[asyncio.Task[Any]] = [] + + try: + first = asyncio.create_task(sender.send({"id": 1}), name="send-1") + tasks.append(first) + await asyncio.wait_for(writer.drain_started.wait(), timeout=1) + + second = asyncio.create_task(sender.send({"id": 2}), name="send-2") + tasks.append(second) + await _wait_for_queue_size(sender, 1) + + writer.release.set() + with pytest.raises(BrokenPipeError, match="writer failed"): + await first + + late = asyncio.create_task(sender.send({"id": 3}), name="send-3") + tasks.append(late) + await _assert_tasks_finish(tasks) + + for task in (second, late): + assert isinstance(task.exception(), BrokenPipeError) + assert str(task.exception()) == "writer failed" + finally: + await _close_supervisor(supervisor, tasks) + + +@pytest.mark.asyncio +async def test_supervisor_shutdown_settles_active_queued_and_late_sends() -> None: + writer = _ControlledWriter() + supervisor = TaskSupervisor(source="test") + sender = MessageSender(writer, supervisor) + tasks: list[asyncio.Task[Any]] = [] + + try: + first = asyncio.create_task(sender.send({"id": 1}), name="send-1") + tasks.append(first) + await asyncio.wait_for(writer.drain_started.wait(), timeout=1) + + second = asyncio.create_task(sender.send({"id": 2}), name="send-2") + tasks.append(second) + await _wait_for_queue_size(sender, 1) + + await supervisor.shutdown() + late = asyncio.create_task(sender.send({"id": 3}), name="send-3") + tasks.append(late) + await _assert_tasks_finish(tasks) + + for task in tasks: + assert isinstance(task.exception(), ConnectionError) + finally: + await _close_supervisor(supervisor, tasks) + + +@pytest.mark.asyncio +async def test_shutdown_before_sender_loop_starts_rejects_later_send() -> None: + writer = _ControlledWriter() + supervisor = TaskSupervisor(source="test") + sender = MessageSender(writer, supervisor) + tasks: list[asyncio.Task[Any]] = [] + + try: + await supervisor.shutdown() + late = asyncio.create_task(sender.send({"id": 1}), name="send-after-shutdown") + tasks.append(late) + await _assert_tasks_finish(tasks) + assert isinstance(late.exception(), ConnectionError) + finally: + await _close_supervisor(supervisor, tasks) + + +@pytest.mark.asyncio +async def test_healthy_sends_preserve_order_and_graceful_close() -> None: + writer = _ControlledWriter() + supervisor = TaskSupervisor(source="test") + sender = MessageSender(writer, supervisor) + tasks: list[asyncio.Task[Any]] = [] + closing: asyncio.Task[None] | None = None + + try: + tasks.extend([ + asyncio.create_task(sender.send({"id": 1}), name="send-1"), + asyncio.create_task(sender.send({"id": 2}), name="send-2"), + ]) + await asyncio.wait_for(writer.drain_started.wait(), timeout=1) + await _wait_for_queue_size(sender, 1) + + closing = asyncio.create_task(sender.close(), name="sender-close") + await asyncio.sleep(0) + assert not closing.done() + assert [json.loads(payload)["id"] for payload in writer.writes] == [1] + + writer.release.set() + await asyncio.gather(*tasks) + await asyncio.wait_for(closing, timeout=1) + assert [json.loads(payload)["id"] for payload in writer.writes] == [1, 2] + + late = asyncio.create_task(sender.send({"id": 3}), name="send-after-close") + tasks.append(late) + await _assert_tasks_finish([late]) + assert isinstance(late.exception(), ConnectionError) + finally: + if closing is not None and not closing.done(): + closing.cancel() + if closing is not None: + await asyncio.gather(closing, return_exceptions=True) + await _close_supervisor(supervisor, tasks) From a4c3ecca44db324de10329fd05e3360396c12510 Mon Sep 17 00:00:00 2001 From: Ehsan-hash Date: Sat, 17 Oct 2026 00:27:14 -0700 Subject: [PATCH 2/2] fix(sender): settle sends when the loop terminates Retain writer failures for active, queued and later sends. Preserve cancellation during shutdown and reject sends after closure, including cancellation before the sender loop starts. Developed and reviewed with AI assistance using Codex. --- src/acp/task/sender.py | 65 ++++++++++++++++++++++++++++++- tests/test_connection_recovery.py | 4 +- tests/test_message_sender.py | 19 ++++----- 3 files changed, 75 insertions(+), 13 deletions(-) diff --git a/src/acp/task/sender.py b/src/acp/task/sender.py index 613535b..f367050 100644 --- a/src/acp/task/sender.py +++ b/src/acp/task/sender.py @@ -23,9 +23,27 @@ def __init__(self, writer: asyncio.StreamWriter, supervisor: TaskSupervisor) -> self._writer = writer self._queue: asyncio.Queue[_PendingSend | None] = asyncio.Queue() self._closed = False + self._active: _PendingSend | None = None + self._terminal_error: BaseException | None = None + self._terminal_cancelled = False self._task = supervisor.create(self._loop(), name="acp.Sender.loop", on_error=self._on_error) + self._task.add_done_callback(self._on_task_done) async def send(self, payload: dict[str, Any]) -> None: + if self._terminal_error is not None: + if self._terminal_cancelled: + raise asyncio.CancelledError + error = self._terminal_error + raise error + if self._closed: + raise asyncio.CancelledError + if self._task.done(): + if self._task.cancelled(): + self._set_terminal(ConnectionError("Sender closed"), cancel_pending=True) + raise asyncio.CancelledError + error = self._task.exception() or ConnectionError("Sender closed") + self._set_terminal(error) + raise error data = (json.dumps(payload, separators=(",", ":")) + "\n").encode("utf-8") future: asyncio.Future[None] = asyncio.get_running_loop().create_future() await self._queue.put(_PendingSend(data, future)) @@ -46,18 +64,61 @@ async def _loop(self) -> None: item = await self._queue.get() if item is None: return + self._active = item try: self._writer.write(item.payload) await self._writer.drain() + except asyncio.CancelledError: + self._set_terminal(ConnectionError("Sender closed"), cancel_pending=True) + raise except Exception as exc: - if not item.future.done(): - item.future.set_exception(exc) + self._set_terminal(exc) raise else: if not item.future.done(): item.future.set_result(None) + finally: + self._active = None except asyncio.CancelledError: + self._set_terminal(ConnectionError("Sender closed"), cancel_pending=True) + raise + + def _set_terminal(self, error: BaseException, *, cancel_pending: bool = False) -> None: + if self._terminal_error is not None: return + self._terminal_error = error + self._terminal_cancelled = cancel_pending + if cancel_pending: + self._cancel_pending() + return + if self._active is not None and not self._active.future.done(): + self._active.future.set_exception(error) + while True: + try: + item = self._queue.get_nowait() + except asyncio.QueueEmpty: + break + if item is not None and not item.future.done(): + item.future.set_exception(error) + + def _cancel_pending(self) -> None: + if self._active is not None and not self._active.future.done(): + self._active.future.cancel() + while True: + try: + item = self._queue.get_nowait() + except asyncio.QueueEmpty: + break + if item is not None and not item.future.done(): + item.future.cancel() + + def _on_task_done(self, task: asyncio.Task[Any]) -> None: + if task.cancelled(): + self._set_terminal(ConnectionError("Sender closed"), cancel_pending=True) + else: + error = task.exception() + if error is not None: + self._set_terminal(error) def _on_error(self, task: asyncio.Task[Any], exc: BaseException) -> None: logging.exception("Send loop failed", exc_info=exc) diff --git a/tests/test_connection_recovery.py b/tests/test_connection_recovery.py index 4e0a769..5505add 100644 --- a/tests/test_connection_recovery.py +++ b/tests/test_connection_recovery.py @@ -3,7 +3,7 @@ import asyncio import contextlib import json -from typing import Any +from typing import Any, cast from unittest.mock import AsyncMock, MagicMock import pytest @@ -170,7 +170,7 @@ def tracking_process(message: dict[str, Any]) -> None: async def test_send_request_removes_pending_after_writer_failure() -> None: reader = asyncio.StreamReader() writer = _FailingWriter() - conn = Connection(_noop_handler, writer, reader, listening=False) + conn = Connection(_noop_handler, cast(asyncio.StreamWriter, writer), reader, listening=False) requests = [ asyncio.create_task(conn.send_request("first"), name="request-first"), asyncio.create_task(conn.send_request("second"), name="request-second"), diff --git a/tests/test_message_sender.py b/tests/test_message_sender.py index a17ca3d..4808ae5 100644 --- a/tests/test_message_sender.py +++ b/tests/test_message_sender.py @@ -2,7 +2,7 @@ import asyncio import json -from typing import Any +from typing import Any, cast import pytest @@ -57,7 +57,7 @@ async def test_writer_failure_settles_active_queued_and_late_sends() -> None: error = BrokenPipeError("writer failed") writer = _ControlledWriter(error=error) supervisor = TaskSupervisor(source="test") - sender = MessageSender(writer, supervisor) + sender = MessageSender(cast(asyncio.StreamWriter, writer), supervisor) tasks: list[asyncio.Task[Any]] = [] try: @@ -88,7 +88,7 @@ async def test_writer_failure_settles_active_queued_and_late_sends() -> None: async def test_supervisor_shutdown_settles_active_queued_and_late_sends() -> None: writer = _ControlledWriter() supervisor = TaskSupervisor(source="test") - sender = MessageSender(writer, supervisor) + sender = MessageSender(cast(asyncio.StreamWriter, writer), supervisor) tasks: list[asyncio.Task[Any]] = [] try: @@ -105,8 +105,9 @@ async def test_supervisor_shutdown_settles_active_queued_and_late_sends() -> Non tasks.append(late) await _assert_tasks_finish(tasks) - for task in tasks: - assert isinstance(task.exception(), ConnectionError) + for task in tasks[:2]: + assert task.cancelled() + assert late.cancelled() finally: await _close_supervisor(supervisor, tasks) @@ -115,7 +116,7 @@ async def test_supervisor_shutdown_settles_active_queued_and_late_sends() -> Non async def test_shutdown_before_sender_loop_starts_rejects_later_send() -> None: writer = _ControlledWriter() supervisor = TaskSupervisor(source="test") - sender = MessageSender(writer, supervisor) + sender = MessageSender(cast(asyncio.StreamWriter, writer), supervisor) tasks: list[asyncio.Task[Any]] = [] try: @@ -123,7 +124,7 @@ async def test_shutdown_before_sender_loop_starts_rejects_later_send() -> None: late = asyncio.create_task(sender.send({"id": 1}), name="send-after-shutdown") tasks.append(late) await _assert_tasks_finish(tasks) - assert isinstance(late.exception(), ConnectionError) + assert late.cancelled() finally: await _close_supervisor(supervisor, tasks) @@ -132,7 +133,7 @@ async def test_shutdown_before_sender_loop_starts_rejects_later_send() -> None: async def test_healthy_sends_preserve_order_and_graceful_close() -> None: writer = _ControlledWriter() supervisor = TaskSupervisor(source="test") - sender = MessageSender(writer, supervisor) + sender = MessageSender(cast(asyncio.StreamWriter, writer), supervisor) tasks: list[asyncio.Task[Any]] = [] closing: asyncio.Task[None] | None = None @@ -157,7 +158,7 @@ async def test_healthy_sends_preserve_order_and_graceful_close() -> None: late = asyncio.create_task(sender.send({"id": 3}), name="send-after-close") tasks.append(late) await _assert_tasks_finish([late]) - assert isinstance(late.exception(), ConnectionError) + assert late.cancelled() finally: if closing is not None and not closing.done(): closing.cancel()