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

internal/node/gate_test.go

@@ -16,7 +16,9 @@ package node
 
 import (
 	"context"
+	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 	"time"
 
@@ -313,3 +315,216 @@ func TestHandleMCPStream_ForwarderRoutesCalls(t *testing.T) {
 		t.Errorf("forwarder did not pass un-namespaced name; got %q", tc.Text)
 	}
 }
+
+// startBareQUICNode is startBareNode over the QUIC transport instead of TCP.
+//
+// The distinction matters for TestHandleStreamPassThrough_BackendEOFDoesNotDropInFlightResponse
+// below: go-libp2p's QUIC-backed network.Stream.Close cancels the read side
+// over the wire (CancelRead sends a QUIC STOP_SENDING frame to the peer),
+// while the yamux stream used over startBareNode's plain TCP transport does
+// the equivalent locally only - its CloseRead doc says plainly "Remote is
+// not notified." A stream backed by TCP+yamux structurally cannot exercise
+// the wire-level race this test is after; only a real QUIC connection can.
+func startBareQUICNode(t *testing.T, ctx context.Context) (*SamNode, func()) {
+	t.Helper()
+	dir := t.TempDir()
+	store, err := NewStore(dir)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	priv, _, err := crypto.GenerateKeyPair(crypto.Ed25519, -1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	node, err := NewSamNode(Options{
+		PrivKey:           priv,
+		RouterAddrs:       nil,
+		Store:             store,
+		MeshID:            "test-mesh",
+		DiscoveryInterval: "1s",
+		ListenAddrs:       []string{"/ip4/127.0.0.1/udp/0/quic-v1"},
+		EnableRelay:       false,
+		NodeConfig:        &NodeConfigComplete{},
+		KeyGracePeriod:    24 * time.Hour,
+		AllowLoopback:     true,
+		MonitorBootstrap:  2 * time.Minute,
+		MonitorInterval:   1 * time.Minute,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	node.BiscuitTimeout = 500 * time.Millisecond
+	if err := node.Start(ctx); err != nil {
+		t.Fatal(err)
+	}
+
+	cleanup := func() {
+		_ = node.Teardown()
+		_ = store.Close()
+	}
+	return node, cleanup
+}
+
+// TestHandleStreamPassThrough_BackendEOFDoesNotDropInFlightResponse is a
+// regression test for google/sam#375.
+//
+// A backend that answers a single request and then closes its side of the
+// connection - a normal EOF for a one-shot HTTP-style backend, not a
+// failure - used to make HandleStreamPassThrough tear the whole
+// client-facing stream down immediately via network.Stream.Close. Close's
+// own documented contract says it "does not guarantee receipt of the data";
+// closing right behind a Write that had just reported success raced the
+// response off the wire, and the caller in the original report sometimes
+// saw EOF in its place. See the comment on the drain logic in
+// HandleStreamPassThrough for the fix.
+//
+// The race is timing-dependent, so this repeats the call many times over a
+// real QUIC connection (see startBareQUICNode) rather than asserting on a
+// single attempt.
+func TestHandleStreamPassThrough_BackendEOFDoesNotDropInFlightResponse(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
+	defer cancel()
+
+	// A payload closer in size to the 834-byte response in the original
+	// report than a trivial "ok" would be - small messages are far more
+	// likely to always win the race regardless of the bug.
+	wantText := strings.Repeat("x", 700)
+	upstream := httptest.NewServer(mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server {
+		srv := mcp.NewServer(&mcp.Implementation{Name: "fake", Version: "0.0.1"}, nil)
+		srv.AddTool(&mcp.Tool{Name: "echo", Description: "echo", InputSchema: map[string]any{"type": "object"}},
+			func(context.Context, *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+				return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: wantText}}}, nil
+			})
+		return srv
+	}, nil))
+	defer upstream.Close()
+
+	nodeA, cleanupA := startBareQUICNode(t, ctx)
+	defer cleanupA()
+	nodeB, cleanupB := startBareQUICNode(t, ctx)
+	defer cleanupB()
+
+	svc := &MCPService{baseService: baseService{
+		info:    &api.ServiceInfo{Type: api.ServiceType_SERVICE_TYPE_MCP, Name: "svc"},
+		backend: &api.RegisterServiceRequest_TargetUrl{TargetUrl: upstream.URL},
+	}}
+	if err := svc.Init(ctx); err != nil {
+		t.Fatalf("MCPService.Init: %v", err)
+	}
+	nodeA.services.insertService(svc)
+	t.Cleanup(func() { _ = svc.Teardown() })
+
+	nodeA.Host.SetStreamHandler(testMCPProtocol, func(s network.Stream) {
+		nodeA.HandleMCPStream(s, RequestContext{Target: "svc"})
+	})
+
+	if err := nodeB.Host.Connect(ctx, peer.AddrInfo{ID: nodeA.Host.ID(), Addrs: nodeA.Host.Addrs()}); err != nil {
+		t.Fatalf("connect: %v", err)
+	}
+
+	const iterations = 50
+	for i := 0; i < iterations; i++ {
+		func() {
+			s, err := nodeB.Host.NewStream(ctx, nodeA.Host.ID(), testMCPProtocol)
+			if err != nil {
+				t.Fatalf("iteration %d: NewStream: %v", i, err)
+			}
+			defer func() { _ = s.Close() }()
+
+			client := mcp.NewClient(&mcp.Implementation{Name: "tc", Version: "0.0.1"}, nil)
+			session, err := client.Connect(ctx, NewStreamTransport(s), nil)
+			if err != nil {
+				t.Fatalf("iteration %d: client.Connect: %v", i, err)
+			}
+			defer func() { _ = session.Close() }()
+
+			res, err := session.CallTool(ctx, &mcp.CallToolParams{Name: "echo", Arguments: map[string]any{}})
+			if err != nil {
+				t.Fatalf("iteration %d: CallTool: %v", i, err)
+			}
+			tc, ok := res.Content[0].(*mcp.TextContent)
+			if !ok || tc.Text != wantText {
+				t.Fatalf("iteration %d: got %v, want text of length %d", i, res.Content, len(wantText))
+			}
+		}()
+	}
+}
+
+// TestHandleStreamPassThrough_SlowBackendDoesNotHitDrainTimeout is a
+// regression test for review feedback from aojea on google/sam#379: the
+// drain wait was timed from the start of the whole exchange, not from when
+// the backend leg actually finished, so any session - healthy or not -
+// that happened to run longer than passThroughDrainTimeout got killed
+// mid-flight. Shrinks passThroughDrainTimeout for the test so proving this
+// doesn't need a multi-second sleep; startBareNode's plain TCP transport is
+// enough here, since this is about goroutine timing, not the QUIC-specific
+// Close semantics TestHandleStreamPassThrough_BackendEOFDoesNotDropInFlightResponse
+// covers.
+func TestHandleStreamPassThrough_SlowBackendDoesNotHitDrainTimeout(t *testing.T) {
+	oldTimeout := passThroughDrainTimeout
+	passThroughDrainTimeout = 50 * time.Millisecond
+	t.Cleanup(func() { passThroughDrainTimeout = oldTimeout })
+
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
+	// Several times passThroughDrainTimeout: under the bug, the stream was
+	// torn down well before the backend ever answered.
+	const backendDelay = 300 * time.Millisecond
+	upstream := httptest.NewServer(mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server {
+		srv := mcp.NewServer(&mcp.Implementation{Name: "slow", Version: "0.0.1"}, nil)
+		srv.AddTool(&mcp.Tool{Name: "slow_echo", Description: "slow", InputSchema: map[string]any{"type": "object"}},
+			func(context.Context, *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+				time.Sleep(backendDelay)
+				return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: "ok"}}}, nil
+			})
+		return srv
+	}, nil))
+	defer upstream.Close()
+
+	nodeA, cleanupA := startBareNode(t, ctx)
+	defer cleanupA()
+	nodeB, cleanupB := startBareNode(t, ctx)
+	defer cleanupB()
+
+	svc := &MCPService{baseService: baseService{
+		info:    &api.ServiceInfo{Type: api.ServiceType_SERVICE_TYPE_MCP, Name: "slow-svc"},
+		backend: &api.RegisterServiceRequest_TargetUrl{TargetUrl: upstream.URL},
+	}}
+	if err := svc.Init(ctx); err != nil {
+		t.Fatalf("MCPService.Init: %v", err)
+	}
+	nodeA.services.insertService(svc)
+	t.Cleanup(func() { _ = svc.Teardown() })
+
+	nodeA.Host.SetStreamHandler(testMCPProtocol, func(s network.Stream) {
+		nodeA.HandleMCPStream(s, RequestContext{Target: "slow-svc"})
+	})
+
+	if err := nodeB.Host.Connect(ctx, peer.AddrInfo{ID: nodeA.Host.ID(), Addrs: nodeA.Host.Addrs()}); err != nil {
+		t.Fatalf("connect: %v", err)
+	}
+
+	s, err := nodeB.Host.NewStream(ctx, nodeA.Host.ID(), testMCPProtocol)
+	if err != nil {
+		t.Fatalf("NewStream: %v", err)
+	}
+	defer func() { _ = s.Close() }()
+
+	client := mcp.NewClient(&mcp.Implementation{Name: "tc", Version: "0.0.1"}, nil)
+	session, err := client.Connect(ctx, NewStreamTransport(s), nil)
+	if err != nil {
+		t.Fatalf("client.Connect: %v", err)
+	}
+	defer func() { _ = session.Close() }()
+
+	res, err := session.CallTool(ctx, &mcp.CallToolParams{Name: "slow_echo", Arguments: map[string]any{}})
+	if err != nil {
+		t.Fatalf("CallTool: %v (a healthy exchange slower than passThroughDrainTimeout must not be killed for that alone)", err)
+	}
+	tc, ok := res.Content[0].(*mcp.TextContent)
+	if !ok || tc.Text != "ok" {
+		t.Fatalf("got %v, want text %q", res.Content, "ok")
+	}
+}

internal/node/mcp_service.go

@@ -140,6 +140,18 @@ var preflightMethodsUnsupportedByPassThrough = map[string]bool{
 	"server/discover": true,
 }
 
+// passThroughDrainTimeout bounds how long HandleStreamPassThrough keeps a
+// client-facing stream open after the backend leg has ended, waiting for the
+// client to finish reading the last response and hang up on its own. See the
+// comment on the drain logic in HandleStreamPassThrough for why this wait
+// exists at all.
+//
+// A var, not a const, so a test can shrink it - the countdown only starts
+// once the backend leg is done, so shrinking it does not affect an
+// in-progress exchange, only how long a stalled client is tolerated for
+// after that.
+var passThroughDrainTimeout = 5 * time.Second
+
 // HandleStreamPassThrough connects to the backend and proxies JSON-RPC messages.
 func (m *MCPService) HandleStreamPassThrough(s network.Stream) {
 	defer func() {
@@ -148,23 +160,14 @@ func (m *MCPService) HandleStreamPassThrough(s network.Stream) {
 		}
 	}()
 
-	var backendTransport mcp.Transport
-	var closeTransport func()
-
 	backendTransport, err := m.backendTransport()
 	if err != nil {
 		logger.Errorf("[MCPService] %s: %v", m.info.Name, err)
 		return
 	}
-	closeTransport = func() {} // fresh per stream for URL; shared bridge is never closed
 
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
-	defer func() {
-		if closeTransport != nil {
-			closeTransport()
-		}
-	}()
 
 	backendConn, err := backendTransport.Connect(ctx)
 	if err != nil {
@@ -180,49 +183,98 @@ func (m *MCPService) HandleStreamPassThrough(s network.Stream) {
 		return
 	}
 
-	// Dumb pipe: Proxy JSON-RPC messages between client and backend
-	errc := make(chan error, 2)
+	// Dumb pipe: proxy JSON-RPC messages between client and backend.
+	//
+	// The two legs are not symmetric on shutdown. A backend answering one
+	// request and then hanging up (a clean EOF, typical of a one-shot
+	// HTTP-style backend transport) is normal completion, not a failure of
+	// the client-facing side of the pipe - but closing s in reaction to it
+	// used to tear down both directions immediately (network.Stream.Close
+	// implies CancelRead), including the read side and, per that method's
+	// own documented contract, without waiting for the response this same
+	// goroutine had just handed to Write to actually reach the peer. Close
+	// "does not guarantee receipt of the data"; the documented safe sequence
+	// is CloseWrite, then wait for the peer to finish reading (or hang up),
+	// then Close. That race is the root cause of google/sam#375: the
+	// producer's write reports success, but the immediate teardown right
+	// behind it can still lose the response in flight, and the consumer
+	// sees EOF instead.
+	//
+	// So the backend leg ending only half-closes our write side to the
+	// client (CloseWrite: no more responses are coming, but nothing already
+	// in flight is discarded) and stops relaying backend->client; it leaves
+	// the read side alone. The client leg - the client itself finishing the
+	// read and hanging up, or a genuine transport error - is what triggers
+	// the final s.Close() at the top of this function. passThroughDrainTimeout
+	// bounds that wait, but only once the backend leg is actually done
+	// (backendDone below) - it is not a cap on the whole exchange, or a
+	// slow-but-healthy session would be killed mid-flight for no better
+	// reason than having taken a while.
+	//
+	// clientErrc is buffered for 2, not 1: a client write failure below is
+	// also a client-leg error (the stream to the client is dead, so there is
+	// nothing left to drain for), and with both goroutines able to send, an
+	// unlucky interleaving where the main select has already consumed one
+	// value could otherwise leave the second sender blocked forever.
+	clientErrc := make(chan error, 2)
+	backendDone := make(chan struct{})
+
+	go func() {
+		defer close(backendDone)
+		for {
+			msg, err := backendConn.Read(ctx)
+			if err != nil {
+				logger.Debugf("[MCPService] %s: backend read error: %v", m.info.Name, err)
+				if cwErr := s.CloseWrite(); cwErr != nil {
+					logger.Debugf("[MCPService] %s: failed to close write side to client: %v", m.info.Name, cwErr)
+				}
+				return
+			}
+			if err := clientConn.Write(ctx, msg); err != nil {
+				logger.Debugf("[MCPService] %s: client write error: %v", m.info.Name, err)
+				clientErrc <- err
+				return
+			}
+		}
+	}()
 
 	go func() {
 		for {
 			msg, err := clientConn.Read(ctx)
 			if err != nil {
 				logger.Debugf("[MCPService] %s: client read error: %v", m.info.Name, err)
-				errc <- err
+				clientErrc <- err
 				return
 			}
 			if req, ok := msg.(*jsonrpc.Request); ok && preflightMethodsUnsupportedByPassThrough[req.Method] {
 				resp := &jsonrpc.Response{ID: req.ID, Error: &jsonrpc.Error{Code: jsonrpc.CodeMethodNotFound, Message: req.Method + " is not supported by this pass-through proxy"}}
 				if werr := clientConn.Write(ctx, resp); werr != nil {
 					logger.Debugf("[MCPService] %s: failed to reject %s: %v", m.info.Name, req.Method, werr)
-					errc <- werr
+					clientErrc <- werr
 					return
 				}
 				continue
 			}
 			if err := backendConn.Write(ctx, msg); err != nil {
 				logger.Debugf("[MCPService] %s: backend write error: %v", m.info.Name, err)
-				errc <- err
+				clientErrc <- err
 				return
 			}
 		}
 	}()
 
-	go func() {
-		for {
-			msg, err := backendConn.Read(ctx)
-			if err != nil {
-				logger.Debugf("[MCPService] %s: backend read error: %v", m.info.Name, err)
-				errc <- err
-				return
-			}
-			if err := clientConn.Write(ctx, msg); err != nil {
-				logger.Debugf("[MCPService] %s: client write error: %v", m.info.Name, err)
-				errc <- err
-				return
-			}
-		}
-	}()
+	// No timeout here: the exchange runs for as long as both legs are
+	// making progress. Only once the backend leg ends (backendDone) does a
+	// bounded wait for the client to also finish begin, below.
+	select {
+	case <-clientErrc:
+		return
+	case <-backendDone:
+	}
 
-	<-errc
+	select {
+	case <-clientErrc:
+	case <-time.After(passThroughDrainTimeout):
+		logger.Debugf("[MCPService] %s: client did not hang up within %v of the backend finishing; closing", m.info.Name, passThroughDrainTimeout)
+	}
 }