Skip to content

Commit 8a99407

Browse files
frostmingclaude
andauthored
fix(task): serialize sends with a lock so send() cannot hang (#190)
MessageSender used a queue drained by a background task. Once that loop stopped (write error, cancellation, or a send after close), queued and later sends awaited futures that were never settled. Replace the queue and loop with an asyncio.Lock: writes run in the caller's task, so errors and cancellation surface directly. After a failed write or close(), further sends raise ConnectionError. Also: - Connection._run_request drops a response that can't be sent because the connection already closed, instead of failing the task. - TaskSupervisor.shutdown() no longer re-raises errors from tasks that already failed and were reported by the error handlers. Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
1 parent 43ffda0 commit 8a99407

4 files changed

Lines changed: 124 additions & 50 deletions

File tree

‎src/acp/connection.py‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ def __init__(
7676
if reader is None:
7777
self._transport: Transport = cast("Transport", writer)
7878
else:
79-
sender = MessageSender(cast("asyncio.StreamWriter", writer), self._tasks)
79+
sender = MessageSender(cast("asyncio.StreamWriter", writer))
8080
self._transport = NdjsonTransport(reader, sender, receive_timeout=receive_timeout)
8181
self._observers: list[StreamObserver] = list(observers or [])
8282
if listening:
@@ -248,7 +248,14 @@ async def _run_request(
248248
payload = {"jsonrpc": "2.0", "id": message["id"], "error": RequestError.request_cancelled().to_error_obj()}
249249
else:
250250
payload = handler.result()
251-
await self._transport.send(payload)
251+
try:
252+
await self._transport.send(payload)
253+
except ConnectionError:
254+
if self._closed or self._disconnected:
255+
# The connection went away while the handler was running; nobody can receive it.
256+
logging.debug("Dropping response for request %r: connection closed", message["id"])
257+
return
258+
raise
252259
self._notify_observers(StreamDirection.OUTGOING, payload)
253260

254261
async def _execute_request(self, message: dict[str, Any]) -> dict[str, Any]:

‎src/acp/task/sender.py‎

Lines changed: 21 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,63 +1,40 @@
11
from __future__ import annotations
22

33
import asyncio
4-
import contextlib
54
import json
6-
import logging
7-
from dataclasses import dataclass
85
from typing import Any
96

10-
from .supervisor import TaskSupervisor
11-
127
__all__ = ["MessageSender"]
138

149

15-
@dataclass(slots=True)
16-
class _PendingSend:
17-
payload: bytes
18-
future: asyncio.Future[None]
10+
class MessageSender:
11+
"""Serialise NDJSON frames onto a stream writer.
1912
13+
``StreamWriter.write`` buffers a whole frame synchronously, so frames never interleave; the
14+
lock only serialises ``drain()`` so backpressure is applied in FIFO order. Writes happen in
15+
the caller's task, so errors and cancellation surface directly in ``send()``.
16+
"""
2017

21-
class MessageSender:
22-
def __init__(self, writer: asyncio.StreamWriter, supervisor: TaskSupervisor) -> None:
18+
def __init__(self, writer: asyncio.StreamWriter) -> None:
2319
self._writer = writer
24-
self._queue: asyncio.Queue[_PendingSend | None] = asyncio.Queue()
20+
self._lock = asyncio.Lock()
2521
self._closed = False
26-
self._task = supervisor.create(self._loop(), name="acp.Sender.loop", on_error=self._on_error)
22+
self._failure: BaseException | None = None
2723

2824
async def send(self, payload: dict[str, Any]) -> None:
2925
data = (json.dumps(payload, separators=(",", ":")) + "\n").encode("utf-8")
30-
future: asyncio.Future[None] = asyncio.get_running_loop().create_future()
31-
await self._queue.put(_PendingSend(data, future))
32-
await future
26+
async with self._lock:
27+
if self._closed:
28+
raise ConnectionError("Sender closed") from self._failure
29+
try:
30+
self._writer.write(data)
31+
await self._writer.drain()
32+
except Exception as exc:
33+
# The stream is unusable after a failed write; fail later sends explicitly
34+
# instead of silently writing into a closed transport.
35+
self._failure = exc
36+
self._closed = True
37+
raise
3338

3439
async def close(self) -> None:
35-
if self._closed:
36-
return
3740
self._closed = True
38-
await self._queue.put(None)
39-
if self._task is not None:
40-
with contextlib.suppress(asyncio.CancelledError):
41-
await self._task
42-
43-
async def _loop(self) -> None:
44-
try:
45-
while True:
46-
item = await self._queue.get()
47-
if item is None:
48-
return
49-
try:
50-
self._writer.write(item.payload)
51-
await self._writer.drain()
52-
except Exception as exc:
53-
if not item.future.done():
54-
item.future.set_exception(exc)
55-
raise
56-
else:
57-
if not item.future.done():
58-
item.future.set_result(None)
59-
except asyncio.CancelledError:
60-
return
61-
62-
def _on_error(self, task: asyncio.Task[Any], exc: BaseException) -> None:
63-
logging.exception("Send loop failed", exc_info=exc)

‎src/acp/task/supervisor.py‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
import asyncio
44
import logging
55
from collections.abc import Awaitable, Callable
6-
from contextlib import suppress
76
from typing import Any
87

98
__all__ = ["TaskSupervisor"]
@@ -74,7 +73,7 @@ async def shutdown(self) -> None:
7473
tasks = list(self._tasks)
7574
for task in tasks:
7675
task.cancel()
77-
for task in tasks:
78-
with suppress(asyncio.CancelledError):
79-
await task
76+
# Failures were already routed to the error handlers by ``_on_done``; shutdown only
77+
# needs to wait for the tasks to finish, so don't let them escape from here.
78+
await asyncio.gather(*tasks, return_exceptions=True)
8079
self._tasks.clear()

‎tests/test_message_sender.py‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
from __future__ import annotations
2+
3+
import asyncio
4+
5+
import pytest
6+
7+
from acp.task import MessageSender
8+
9+
10+
class _FakeWriter:
11+
def __init__(self, *, fail_with: Exception | None = None) -> None:
12+
self.fail_with = fail_with
13+
self.written: list[bytes] = []
14+
self.release = asyncio.Event()
15+
self.block = False
16+
17+
def write(self, data: bytes) -> None:
18+
self.written.append(data)
19+
20+
async def drain(self) -> None:
21+
if self.block:
22+
await self.release.wait()
23+
if self.fail_with is not None:
24+
raise self.fail_with
25+
26+
27+
def _make_sender(writer: _FakeWriter) -> MessageSender:
28+
return MessageSender(writer) # type: ignore[arg-type]
29+
30+
31+
@pytest.mark.asyncio
32+
async def test_sends_are_written_in_call_order() -> None:
33+
writer = _FakeWriter()
34+
writer.block = True
35+
sender = _make_sender(writer)
36+
37+
tasks = [asyncio.create_task(sender.send({"n": i})) for i in range(3)]
38+
await asyncio.sleep(0)
39+
writer.release.set()
40+
await asyncio.wait_for(asyncio.gather(*tasks), 1)
41+
assert writer.written == [b'{"n":0}\n', b'{"n":1}\n', b'{"n":2}\n']
42+
43+
44+
@pytest.mark.asyncio
45+
async def test_send_after_close_raises() -> None:
46+
writer = _FakeWriter()
47+
sender = _make_sender(writer)
48+
await sender.send({"a": 1})
49+
await sender.close()
50+
with pytest.raises(ConnectionError):
51+
await asyncio.wait_for(sender.send({"a": 2}), 1)
52+
assert writer.written == [b'{"a":1}\n']
53+
54+
55+
@pytest.mark.asyncio
56+
async def test_write_failure_fails_waiting_and_later_sends() -> None:
57+
writer = _FakeWriter(fail_with=BrokenPipeError("boom"))
58+
writer.block = True
59+
sender = _make_sender(writer)
60+
61+
first = asyncio.create_task(sender.send({"n": 1}))
62+
waiting = asyncio.create_task(sender.send({"n": 2}))
63+
await asyncio.sleep(0)
64+
writer.release.set()
65+
66+
with pytest.raises(BrokenPipeError):
67+
await asyncio.wait_for(first, 1)
68+
with pytest.raises(ConnectionError) as excinfo:
69+
await asyncio.wait_for(waiting, 1)
70+
assert isinstance(excinfo.value.__cause__, BrokenPipeError)
71+
with pytest.raises(ConnectionError):
72+
await asyncio.wait_for(sender.send({"n": 3}), 1)
73+
assert len(writer.written) == 1
74+
75+
76+
@pytest.mark.asyncio
77+
async def test_cancelled_send_releases_lock() -> None:
78+
writer = _FakeWriter()
79+
writer.block = True
80+
sender = _make_sender(writer)
81+
82+
stuck = asyncio.create_task(sender.send({"n": 1}))
83+
waiting = asyncio.create_task(sender.send({"n": 2}))
84+
await asyncio.sleep(0)
85+
stuck.cancel()
86+
with pytest.raises(asyncio.CancelledError):
87+
await stuck
88+
89+
writer.release.set()
90+
await asyncio.wait_for(waiting, 1)
91+
assert writer.written == [b'{"n":1}\n', b'{"n":2}\n']

0 commit comments

Comments
 (0)