From c10329563d8cb29c7e22df447c458e79cd4f6fbf Mon Sep 17 00:00:00 2001 From: Juan Antonio Osorio Date: Fri, 24 Jul 2026 07:05:54 +0000 Subject: [PATCH] Drain backend notifications before per-call client close TestForwarding_Progress_RealBackend (and _Logging_) flaked because a fire-and-forget backend notification emitted mid tools/call was lost under load. The notification and the tool result travel on separate streams (standalone SSE vs the tools/call response), read through one FIFO receive loop; CallTool returns as soon as the result is read and its deferred Close cancels that loop, so a not-yet-drained notification is discarded and never forwarded downstream -- a permanent loss, which is why bumping the downstream wait timeout never fixed it. Blocking server->client requests (elicitation/sampling) are immune because the backend tool blocks on the response. Before Close, when server->client forwarding is bound, issue a synchronous ping as a drain barrier: its round-trip lets the receive loop drain the already-buffered notification off the shared channel and enqueue it for handling before teardown. Best-effort (the result is already in hand) and skipped entirely when forwarding is not bound, so non-forwarding calls keep the fast path. This is the minimal in-repo mitigation; the robust fix is upstream in the SDK (route mid-call notifications on the correlated request stream, or drain gracefully on close). Tracked for follow-up. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_016WG3mSjVGWNc8nfgbkdd79 --- pkg/vmcp/client/client.go | 6 +++++ pkg/vmcp/client/forwarding.go | 45 +++++++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/pkg/vmcp/client/client.go b/pkg/vmcp/client/client.go index c33826e0ac..9c029f1ce3 100644 --- a/pkg/vmcp/client/client.go +++ b/pkg/vmcp/client/client.go @@ -1064,6 +1064,12 @@ func (h *httpBackendClient) CallTool( return nil, fmt.Errorf("%w: tool call failed on backend %s: %w", vmcp.ErrBackendUnavailable, target.WorkloadID, err) } + // Flush the backend's server->client stream before the deferred Close tears + // down this per-call client, so a fire-and-forget notification the backend + // emitted mid-call (progress/logging) is relayed downstream instead of being + // dropped. See drainServerToClientNotifications for the lost-notification race. + h.drainServerToClientNotifications(ctx, c) + // Extract _meta field from backend response responseMeta := conversion.FromMCPMeta(result.Meta) diff --git a/pkg/vmcp/client/forwarding.go b/pkg/vmcp/client/forwarding.go index 7c62ed4198..d61c324add 100644 --- a/pkg/vmcp/client/forwarding.go +++ b/pkg/vmcp/client/forwarding.go @@ -62,6 +62,51 @@ func (h *httpBackendClient) enableBackendLogging( } } +// drainServerToClientNotifications flushes any in-flight backend->downstream +// notification off the per-call backend client before it is closed, closing a +// lost-notification race that otherwise drops fire-and-forget notifications +// (notifications/progress, notifications/message) under load. +// +// The race: a backend emits such a notification mid tools/call and then returns +// the result. The notification and the result travel on SEPARATE streams — the +// notification on the standalone SSE stream, the result on the tools/call +// response stream — and the client reads every stream through a single FIFO +// channel feeding one receive loop. CallTool returns as soon as the RESULT is +// read; its deferred Close then cancels the receive loop. If the standalone +// stream's notification has not yet been read and enqueued for handling when +// Close fires, it is discarded and never reaches newNotificationForwarder, so +// the downstream client never sees it (a permanent loss, not a slow delivery). +// Blocking server->client REQUESTS (elicitation/sampling) are immune: the +// backend tool blocks on the response, so the client is guaranteed to still be +// reading when they arrive. +// +// The fix is a synchronous ping used purely as a drain barrier. The backend +// flushes the notification onto the wire before returning the result, so by the +// time CallTool returns the notification bytes are already buffered on the +// client's standalone connection. A ping is a full backend round-trip; while it +// is in flight the receive loop drains the buffered notification off the shared +// FIFO channel (a local buffered read completes far ahead of the ping's network +// round-trip), enqueuing it for handling. Because the ping response arrives on +// that same FIFO channel AFTER the notification, the ping returning proves the +// notification was already read and enqueued. The subsequent Close then waits +// for enqueued handlers to finish (jsonrpc2.Connection.Close drains the handler +// queue), so the forward completes before teardown. +// +// Only runs when server->client forwarding is bound (fwd.notifier != nil, the +// same condition under which the notification forwarder is registered); other +// deployments keep the pre-forwarding fast path with no extra round-trip. +// Best-effort: a ping failure only forgoes the barrier — the tool result is +// already in hand, so it must not fail the call. +func (h *httpBackendClient) drainServerToClientNotifications(ctx context.Context, c *client.Client) { + fwd := h.forwarders.Load() + if fwd == nil || fwd.notifier == nil { + return + } + if err := c.Ping(ctx); err != nil { + slog.Debug("notification drain ping failed; forwarded notifications may be lost", "error", err) + } +} + // forwardingClientOptions builds the client-level options that install the // elicitation and sampling handlers on a backend client, each closing over the // captured per-call downstream context so the handler can relay to the right