diff --git a/internal/node/gate_test.go b/internal/node/gate_test.go index c187bbf6..b67d06e5 100644 --- a/internal/node/gate_test.go +++ b/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") + } +} diff --git a/internal/node/mcp_service.go b/internal/node/mcp_service.go index 32b11c26..9cf376e4 100644 --- a/internal/node/mcp_service.go +++ b/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) + } }