diff --git a/src/acp/connection.py b/src/acp/connection.py index f60e6b37..27993547 100644 --- a/src/acp/connection.py +++ b/src/acp/connection.py @@ -76,7 +76,7 @@ def __init__( if reader is None: self._transport: Transport = cast("Transport", writer) else: - sender = MessageSender(cast("asyncio.StreamWriter", writer), self._tasks) + sender = MessageSender(cast("asyncio.StreamWriter", writer)) self._transport = NdjsonTransport(reader, sender, receive_timeout=receive_timeout) self._observers: list[StreamObserver] = list(observers or []) if listening: @@ -248,7 +248,14 @@ async def _run_request( payload = {"jsonrpc": "2.0", "id": message["id"], "error": RequestError.request_cancelled().to_error_obj()} else: payload = handler.result() - await self._transport.send(payload) + try: + await self._transport.send(payload) + except ConnectionError: + if self._closed or self._disconnected: + # The connection went away while the handler was running; nobody can receive it. + logging.debug("Dropping response for request %r: connection closed", message["id"]) + return + raise self._notify_observers(StreamDirection.OUTGOING, payload) async def _execute_request(self, message: dict[str, Any]) -> dict[str, Any]: diff --git a/src/acp/task/sender.py b/src/acp/task/sender.py index 613535bf..69e47783 100644 --- a/src/acp/task/sender.py +++ b/src/acp/task/sender.py @@ -1,63 +1,40 @@ from __future__ import annotations import asyncio -import contextlib import json -import logging -from dataclasses import dataclass from typing import Any -from .supervisor import TaskSupervisor - __all__ = ["MessageSender"] -@dataclass(slots=True) -class _PendingSend: - payload: bytes - future: asyncio.Future[None] +class MessageSender: + """Serialise NDJSON frames onto a stream writer. + ``StreamWriter.write`` buffers a whole frame synchronously, so frames never interleave; the + lock only serialises ``drain()`` so backpressure is applied in FIFO order. Writes happen in + the caller's task, so errors and cancellation surface directly in ``send()``. + """ -class MessageSender: - def __init__(self, writer: asyncio.StreamWriter, supervisor: TaskSupervisor) -> None: + def __init__(self, writer: asyncio.StreamWriter) -> None: self._writer = writer - self._queue: asyncio.Queue[_PendingSend | None] = asyncio.Queue() + self._lock = asyncio.Lock() self._closed = False - self._task = supervisor.create(self._loop(), name="acp.Sender.loop", on_error=self._on_error) + self._failure: BaseException | None = None async def send(self, payload: dict[str, Any]) -> None: 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)) - await future + async with self._lock: + if self._closed: + raise ConnectionError("Sender closed") from self._failure + try: + self._writer.write(data) + await self._writer.drain() + except Exception as exc: + # The stream is unusable after a failed write; fail later sends explicitly + # instead of silently writing into a closed transport. + self._failure = exc + self._closed = True + raise async def close(self) -> None: - if self._closed: - return self._closed = True - await self._queue.put(None) - if self._task is not None: - with contextlib.suppress(asyncio.CancelledError): - await self._task - - async def _loop(self) -> None: - try: - while True: - item = await self._queue.get() - if item is None: - return - try: - self._writer.write(item.payload) - await self._writer.drain() - except Exception as exc: - if not item.future.done(): - item.future.set_exception(exc) - raise - else: - if not item.future.done(): - item.future.set_result(None) - except asyncio.CancelledError: - return - - def _on_error(self, task: asyncio.Task[Any], exc: BaseException) -> None: - logging.exception("Send loop failed", exc_info=exc) diff --git a/src/acp/task/supervisor.py b/src/acp/task/supervisor.py index 7ad7e6d7..2ed7ab41 100644 --- a/src/acp/task/supervisor.py +++ b/src/acp/task/supervisor.py @@ -3,7 +3,6 @@ import asyncio import logging from collections.abc import Awaitable, Callable -from contextlib import suppress from typing import Any __all__ = ["TaskSupervisor"] @@ -74,7 +73,7 @@ async def shutdown(self) -> None: tasks = list(self._tasks) for task in tasks: task.cancel() - for task in tasks: - with suppress(asyncio.CancelledError): - await task + # Failures were already routed to the error handlers by ``_on_done``; shutdown only + # needs to wait for the tasks to finish, so don't let them escape from here. + await asyncio.gather(*tasks, return_exceptions=True) self._tasks.clear() diff --git a/tests/test_message_sender.py b/tests/test_message_sender.py new file mode 100644 index 00000000..9a6604f8 --- /dev/null +++ b/tests/test_message_sender.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +import asyncio + +import pytest + +from acp.task import MessageSender + + +class _FakeWriter: + def __init__(self, *, fail_with: Exception | None = None) -> None: + self.fail_with = fail_with + self.written: list[bytes] = [] + self.release = asyncio.Event() + self.block = False + + def write(self, data: bytes) -> None: + self.written.append(data) + + async def drain(self) -> None: + if self.block: + await self.release.wait() + if self.fail_with is not None: + raise self.fail_with + + +def _make_sender(writer: _FakeWriter) -> MessageSender: + return MessageSender(writer) # type: ignore[arg-type] + + +@pytest.mark.asyncio +async def test_sends_are_written_in_call_order() -> None: + writer = _FakeWriter() + writer.block = True + sender = _make_sender(writer) + + tasks = [asyncio.create_task(sender.send({"n": i})) for i in range(3)] + await asyncio.sleep(0) + writer.release.set() + await asyncio.wait_for(asyncio.gather(*tasks), 1) + assert writer.written == [b'{"n":0}\n', b'{"n":1}\n', b'{"n":2}\n'] + + +@pytest.mark.asyncio +async def test_send_after_close_raises() -> None: + writer = _FakeWriter() + sender = _make_sender(writer) + await sender.send({"a": 1}) + await sender.close() + with pytest.raises(ConnectionError): + await asyncio.wait_for(sender.send({"a": 2}), 1) + assert writer.written == [b'{"a":1}\n'] + + +@pytest.mark.asyncio +async def test_write_failure_fails_waiting_and_later_sends() -> None: + writer = _FakeWriter(fail_with=BrokenPipeError("boom")) + writer.block = True + sender = _make_sender(writer) + + first = asyncio.create_task(sender.send({"n": 1})) + waiting = asyncio.create_task(sender.send({"n": 2})) + await asyncio.sleep(0) + writer.release.set() + + with pytest.raises(BrokenPipeError): + await asyncio.wait_for(first, 1) + with pytest.raises(ConnectionError) as excinfo: + await asyncio.wait_for(waiting, 1) + assert isinstance(excinfo.value.__cause__, BrokenPipeError) + with pytest.raises(ConnectionError): + await asyncio.wait_for(sender.send({"n": 3}), 1) + assert len(writer.written) == 1 + + +@pytest.mark.asyncio +async def test_cancelled_send_releases_lock() -> None: + writer = _FakeWriter() + writer.block = True + sender = _make_sender(writer) + + stuck = asyncio.create_task(sender.send({"n": 1})) + waiting = asyncio.create_task(sender.send({"n": 2})) + await asyncio.sleep(0) + stuck.cancel() + with pytest.raises(asyncio.CancelledError): + await stuck + + writer.release.set() + await asyncio.wait_for(waiting, 1) + assert writer.written == [b'{"n":1}\n', b'{"n":2}\n']