fix(transport): make AsyncRwTransport::receive cancel-safe (#941) (#947)

`receive()` reads incoming lines with `read_until`, which is not
cancellation-safe, and it is polled inside the service loop's `select!`.
When another branch (e.g. an outgoing response) becomes ready while a
request line is only partially read, the `receive()` future is dropped
and the next call's `self.line_buf.clear()` discards the partial bytes,
so that incoming request is silently lost and never gets a response.
Under many concurrent large responses this intermittently drops requests.

Keep the partially-read bytes in `line_buf` across calls and clear it
only after a whole line has been consumed, so a cancelled read resumes
the same line instead of dropping it. The buffer is cleared (retaining
capacity) rather than reallocated per message.

Adds a regression test that fires 200 concurrent 64 KiB tool responses
over real stdio pipes and asserts no response id goes missing.

This also covers the child-process client transport, which reuses
`AsyncRwTransport`.
This commit is contained in:
Amey Pawar 2026-07-02 18:38:24 +05:30 • committed by GitHub
commit 4833ec7b63
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -124,8 +124,19 @@ where
async fn receive(&mut self) -> Option<RxJsonRpcMessage<Role>> {
loop {
self.line_buf.clear();
// `read_until` is not cancellation-safe on its own, and `receive` is
// polled inside a `select!` in the service loop: an in-progress line
// read is dropped whenever another branch (e.g. an outgoing response)
// becomes ready. We rely on `read_until` appending into `self.line_buf`
// and only returning at a delimiter or EOF, so a cancelled read leaves
// its partial bytes in `self.line_buf`. Keeping that buffer across
// calls lets the next read resume the same line; it is cleared only
// after a whole line has been consumed. Clearing at the top of the
// loop (the previous behaviour) discarded the partial read and so
// dropped incoming requests under concurrent response load.
match self.read.read_until(b'\n', &mut self.line_buf).await {
// EOF. Any bytes still in `line_buf` are an incomplete trailing
// message with no delimiter, so there is nothing to deliver.
Ok(0) => return None,
Ok(_) => {}
Err(e) => {
@ -133,13 +144,21 @@ where
return None;
}
}
let line = without_carriage_return(
self.line_buf.strip_suffix(b"\n").unwrap_or(&self.line_buf),
);
if line.is_empty() {
continue;
}
match try_parse_with_compatibility::<RxJsonRpcMessage<Role>>(line, "receive") {
// A returned `read_until` means a full line is buffered. Parse it
// (borrowing `line_buf`), then clear the buffer — retaining its
// capacity for the next read — before handling the parse result.
let parsed = {
let line = without_carriage_return(
self.line_buf.strip_suffix(b"\n").unwrap_or(&self.line_buf),
);
if line.is_empty() {
self.line_buf.clear();
continue;
}
try_parse_with_compatibility::<RxJsonRpcMessage<Role>>(line, "receive")
};
self.line_buf.clear();
match parsed {
Ok(Some(msg)) => return Some(msg),
Ok(None) => continue,
Err(JsonRpcMessageCodecError::Serde(e)) => {