fix(sender): settle sends when the loop terminates

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)

tests/test_connection_recovery.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import json
-from typing import Any
+from typing import Any, cast
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
@@ -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, 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"),
+    ]
+
+    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()

tests/test_message_sender.py

@@ -0,0 +1,167 @@
+from __future__ import annotations
+
+import asyncio
+import json
+from typing import Any, cast
+
+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(cast(asyncio.StreamWriter, 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(cast(asyncio.StreamWriter, 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[:2]:
+            assert task.cancelled()
+        assert late.cancelled()
+    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(cast(asyncio.StreamWriter, 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 late.cancelled()
+    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(cast(asyncio.StreamWriter, 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 late.cancelled()
+    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)