Skip to content

fix: mcp_service - don't drop the last response on backend EOF (#375) - #379

Merged
aojea merged 3 commits into
google:mainfrom
fer-marino:fix/mcp-passthrough-eof-race
Sep 9, 2026
Merged

fix: mcp_service - don't drop the last response on backend EOF (#375)#379
aojea merged 3 commits into
google:mainfrom
fer-marino:fix/mcp-passthrough-eof-race

Conversation

@fer-marino

Copy link
Copy Markdown
Contributor

Summary

Fixes #375.

HandleStreamPassThrough proxies JSON-RPC messages between a client-facing P2P stream and a backend MCP connection. It used to shut the whole thing down - s.Close() - the instant either leg reported any error, including the completely normal case of a one-shot HTTP-style backend answering one request and then closing its side (a clean EOF, not a failure).

network.Stream.Close's own documented contract:

Close may be asynchronous and does not guarantee receipt of the data... If acknowledgment is required, the caller must call CloseWrite, then wait on the stream for a response (or an EOF), then call Close().

The old code skipped exactly that: the backend-read goroutine would Write the response to the client, loop back, get EOF from the backend, and immediately trigger the deferred s.Close() a few instructions later - closing right behind a write that had already reported success. On the QUIC transport, Close() additionally cancels the read side over the wire (CancelRead sends a STOP_SENDING frame to the peer), which a plain TCP+yamux stream's CloseRead explicitly does not do (its doc says "Remote is not notified"). That's consistent with the original report: the producer's write reports success, but the response never fully reaches the consumer, which sees EOF instead.

Fix

The backend leg ending now only half-closes our write side to the client (s.CloseWrite()) and stops the backend->client relay goroutine; it no longer touches the read side or frees the stream. The client leg - the client itself finishing its read and hanging up, or a genuine transport error - is what triggers the final s.Close(). A 5s drain timeout (passThroughDrainTimeout) bounds the wait so a client that never hangs up can't leak the stream indefinitely.

Test plan

  • Added TestHandleStreamPassThrough_BackendEOFDoesNotDropInFlightResponse in gate_test.go, which drives 50 real CallTool round trips over a real local QUIC connection between two SamNodes (a new startBareQUICNode helper - startBareNode's existing TCP+yamux transport can't exercise this: yamux's CloseRead is local-only per its own doc comment, so it structurally can't reproduce the wire-level race). It passes reliably with the fix.
  • Being transparent about the test's limits: run against the pre-fix code, it did not reproduce the race on loopback across 1500+ iterations (with and without GOMAXPROCS=1), so it is not proof this fixes the exact wire-level interaction seen in the original report - only that the new code follows network.Stream's documented safe-shutdown sequence and behaves correctly across many round trips. Localhost loopback has effectively no latency/jitter, which likely closes the race window that a real multi-host deployment (or a relayed connection) would have. go build ./..., go vet ./..., and the full internal/node suite pass (the 6 pre-existing unrelated failures in that package are present identically on unmodified main).
  • Maintainer confirmation against the original multi-host/production repro would be the strongest signal this actually closes out sam-node: tools/call succeeds on the backend but consumer sees EOF (producer->consumer relay drops the response) #375.

…e#375)

HandleStreamPassThrough tore the whole client-facing stream down the
instant either leg reported any error, including the normal EOF a
one-shot backend gives after answering a single request. That
immediate network.Stream.Close call cancels the read side and, per
its own documented contract, does not guarantee the response this
same goroutine had just handed to Write actually reached the peer -
closing right behind a successful write can still lose it in flight,
which is the root cause of google#375.

The backend leg ending now only half-closes the write side to the
client (CloseWrite) and stops relaying backend->client; the read side
is left alone until the client itself finishes reading and hangs up,
or a genuine transport error occurs, bounded by a 5s drain timeout so
a client that never hangs up can't leak the stream.

Added a regression test exercising this over a real local QUIC
connection (the race is specific to QUIC's stream Close, which
cancels the read side over the wire - a TCP+yamux stream's CloseRead
is documented as local-only). It passes reliably with the fix; it did
not reproduce the race against the old code on loopback across 1500+
iterations, which the PR description explains.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request addresses a race condition where in-flight responses could be dropped when a backend closes its connection immediately after responding. It introduces a half-close mechanism on the client-facing stream, allowing the client to finish reading before the stream is fully closed, bounded by a 5-second timeout. A regression test using a real QUIC connection has also been added. Feedback suggests that if writing to the client fails, the error should be propagated to the main goroutine to avoid waiting for the timeout unnecessarily.

Comment thread internal/node/mcp_service.go Outdated
Comment thread internal/node/mcp_service.go
… ended

aojea caught a real bug: passThroughDrainTimeout was timed from the
start of the whole exchange (the select was reached right after
launching both goroutines), not from when the backend leg actually
finished - so any session, healthy or not, that happened to run
longer than the timeout got killed mid-flight, which is worse than
the bug this PR fixes.

Adds a backendDone channel closed when the backend-relay goroutine
returns; the drain wait now only begins after that signal. Also
switches passThroughDrainTimeout from a const to a var so a test can
shrink it, and adds a regression test with a backend slower than the
(shrunk) timeout that must still succeed.
@aojea
aojea merged commit 30a9d89 into google:main Sep 9, 2026
18 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

sam-node: tools/call succeeds on the backend but consumer sees EOF (producer->consumer relay drops the response)

2 participants