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.

Author: Ehsan-hash

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

@@ -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"),

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,15 +116,15 @@ 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:
         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)
+        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()