mcp: configure heartbeats to be less chatty (#1294)

**Description**

We were emitting heartbeats every 1s and that was very chatty,
especially when we're generating traces for each heartbeat. This PR:

* Makes the heartbeat configurable and allows even disabling it by
setting the `MCP_PROXY_HEARTBEAT_INTERVAL=0`.
* Changes the default heartbeat interval to 1 minute.
* Eagerly sends an initial ping to tackle the Goose issue. With this,
heartbeats can be disabled and Goose works fine as well.

**Related Issues/PRs (if applicable)**

N/A

**Special notes for reviewers (if applicable)**

N/A

---------

Signed-off-by: Ignasi Barrera <[email protected]>

Author: nacx

internal/mcpproxy/session.go

@@ -15,6 +15,7 @@ import (
 	"io"
 	"log/slog"
 	"net/http"
+	"os"
 	"strings"
 	"sync"
 	"time"
@@ -179,8 +180,22 @@ func (s *session) streamNotifications(ctx context.Context, w http.ResponseWriter
 	//
 	// TODO: no idea exactly why this is necessary. Goose shouldn't block on the first event.
 
-	heartbeatTicker := time.NewTicker(1 * time.Second)
-	defer heartbeatTicker.Stop()
+	var (
+		heartbeats      <-chan time.Time
+		heartbeatTicker *time.Ticker
+	)
+	if heartbeatInterval > 0 {
+		heartbeatTicker = time.NewTicker(heartbeatInterval)
+		defer heartbeatTicker.Stop()
+		heartbeats = heartbeatTicker.C
+	} else {
+		heartbeats = make(chan time.Time) // never ticks
+	}
+
+	// Eagerly send an initial heartbeat event to unblock Goose
+	heartBeatEvent := &sseEvent{event: "message", messages: []jsonrpc.Message{newHeartBeatPingMessage()}}
+	heartBeatEvent.writeAndMaybeFlush(w)
+
 	for {
 		select {
 		case event, ok := <-backendMsgs:
@@ -208,7 +223,12 @@ func (s *session) streamNotifications(ctx context.Context, w http.ResponseWriter
 				s.proxy.recordResponse(ctx, _msg)
 			}
 			event.writeAndMaybeFlush(w)
-		case <-heartbeatTicker.C:
+			// Reset the heartbeat ticker so that the next heartbeat will be sent after the full interval.
+			// This avoids sending heartbeats too frequently when there are events.
+			if heartbeatTicker != nil {
+				heartbeatTicker.Reset(heartbeatInterval)
+			}
+		case <-heartbeats:
 			heartBeatEvent := &sseEvent{event: "message", messages: []jsonrpc.Message{newHeartBeatPingMessage()}}
 			heartBeatEvent.writeAndMaybeFlush(w)
 		case <-ctx.Done():
@@ -218,6 +238,20 @@ func (s *session) streamNotifications(ctx context.Context, w http.ResponseWriter
 	}
 }
 
+// heartbeatInterval is computed at startup to avoid the locks in os.Getenv() to be called on the request path.
+var heartbeatInterval = getHeartbeatInterval(1 * time.Minute)
+
+// getHeartbeatInterval returns the heartbeat interval configured via the MCP_HEARTBEAT_INTERVAL environment variable.
+// If the environment variable is not set or invalid, it returns the default value of 1 minute.
+// This value is intentionally hidden under an environment variable as it is unclear if it is generally useful.
+func getHeartbeatInterval(def time.Duration) time.Duration {
+	hbi, err := time.ParseDuration(os.Getenv("MCP_PROXY_HEARTBEAT_INTERVAL"))
+	if err != nil {
+		return def
+	}
+	return hbi
+}
+
 // sendToAllBackends sends an HTTP request to all backends in this session and returns a channel that streams
 // the response events from all backends.
 func (s *session) sendToAllBackends(ctx context.Context, httpMethod string, request *jsonrpc.Request, span tracing.MCPSpan) <-chan *sseEvent {

internal/mcpproxy/session_test.go

@@ -13,6 +13,7 @@ import (
 	"log/slog"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -150,51 +151,83 @@ func TestHandleNotificationsPerBackend_SSE(t *testing.T) {
 }
 
 func TestSession_StreamNotifications(t *testing.T) {
-	// Single backend streaming two events with valid messages.
-	id1, _ := jsonrpc.MakeID("1")
-	id2, _ := jsonrpc.MakeID("2")
-	msg1, _ := jsonrpc.EncodeMessage(&jsonrpc.Request{Method: "a1", ID: id1})
-	msg2, _ := jsonrpc.EncodeMessage(&jsonrpc.Request{Method: "a2", ID: id2})
-	body := "event: a1\n" + "data: " + string(msg1) + "\n\n" + "event: a2\n" + "data: " + string(msg2) + "\n\n"
-	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if r.Method != http.MethodGet {
-			w.WriteHeader(http.StatusBadRequest)
-			return
-		}
-		if r.Header.Get(internalapi.MCPBackendHeader) != "backend1" {
-			w.WriteHeader(http.StatusBadRequest)
-			return
-		}
-		w.Header().Set("Content-Type", "text/event-stream")
-		for _, b := range []byte(body) {
-			_, _ = w.Write([]byte{b})
-			if f, ok := w.(http.Flusher); ok {
-				f.Flush()
+	tests := []struct {
+		name              string
+		eventInterval     time.Duration
+		deadline          time.Duration
+		heartbeatInterval time.Duration
+		wantHeartbeats    bool
+	}{
+		// the default heartbeat interval is 1 second, but the events will come faster, so
+		// we don't expect any heartbeats.
+		{"fast events", 10 * time.Millisecond, 5 * time.Second, 10 * time.Second, false},
+		// configure a heartbeat interval faster than the event interval, so we expect heartbeats.
+		{"slow events", 20 * time.Millisecond, 5 * time.Second, 10 * time.Millisecond, true},
+		// disable heartbeats. Even though events come in slowly, we don't expect heartbeats.
+		{"no heartbeats", 50 * time.Millisecond, 25 * time.Second, 0, false},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			// Override the default heartbeat interval for testing.
+			originalHeartbeatInterval := heartbeatInterval
+			heartbeatInterval = tc.heartbeatInterval
+			t.Cleanup(func() { heartbeatInterval = originalHeartbeatInterval })
+
+			// Single backend streaming two events with valid messages.
+			id1, _ := jsonrpc.MakeID("1")
+			id2, _ := jsonrpc.MakeID("2")
+			msg1, _ := jsonrpc.EncodeMessage(&jsonrpc.Request{Method: "a1", ID: id1})
+			msg2, _ := jsonrpc.EncodeMessage(&jsonrpc.Request{Method: "a2", ID: id2})
+			body := "event: a1\n" + "data: " + string(msg1) + "\n\n" + "event: a2\n" + "data: " + string(msg2) + "\n\n"
+			srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				if r.Method != http.MethodGet {
+					w.WriteHeader(http.StatusBadRequest)
+					return
+				}
+				if r.Header.Get(internalapi.MCPBackendHeader) != "backend1" {
+					w.WriteHeader(http.StatusBadRequest)
+					return
+				}
+				w.Header().Set("Content-Type", "text/event-stream")
+				for _, b := range []byte(body) {
+					_, _ = w.Write([]byte{b})
+					if f, ok := w.(http.Flusher); ok {
+						f.Flush()
+					}
+					time.Sleep(tc.eventInterval)
+				}
+			}))
+			defer srv.Close()
+			proxy := newTestMCPProxy()
+			proxy.backendListenerAddr = srv.URL
+
+			s := &session{
+				proxy: proxy,
+				perBackendSessions: map[filterapi.MCPBackendName]*compositeSessionEntry{
+					"backend1": {
+						sessionID: "s1",
+					},
+				},
+				route: "test-route",
 			}
-			time.Sleep(10 * time.Millisecond)
-		}
-	}))
-	defer srv.Close()
-	proxy := newTestMCPProxy()
-	proxy.backendListenerAddr = srv.URL
+			rr := httptest.NewRecorder()
+			ctx, cancel := context.WithTimeout(t.Context(), tc.deadline)
+			defer cancel()
+			err2 := s.streamNotifications(ctx, rr)
+			require.NoError(t, err2)
+			out := rr.Body.String()
+			require.Contains(t, out, "event: a1")
+			require.Contains(t, out, "event: a2")
+			heartbeatCount := strings.Count(out, `"method":"ping"`)
 
-	s := &session{
-		proxy: proxy,
-		perBackendSessions: map[filterapi.MCPBackendName]*compositeSessionEntry{
-			"backend1": {
-				sessionID: "s1",
-			},
-		},
-		route: "test-route",
+			if tc.wantHeartbeats {
+				require.Greater(t, heartbeatCount, 1, "expected some heartbeats after the initial one")
+			} else {
+				require.Equal(t, 1, heartbeatCount, "expected only the initial heartbeat")
+			}
+		})
 	}
-	rr := httptest.NewRecorder()
-	ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second)
-	defer cancel()
-	err2 := s.streamNotifications(ctx, rr)
-	require.NoError(t, err2)
-	out := rr.Body.String()
-	require.Contains(t, out, "event: a1")
-	require.Contains(t, out, "event: a2")
 }
 
 func TestSendRequestPerBackend_ErrorStatus(t *testing.T) {
@@ -231,3 +264,27 @@ func TestSendRequestPerBackend_EOF(t *testing.T) {
 	}, http.MethodGet, nil)
 	require.True(t, err2 == nil || errors.Is(err2, io.EOF), "unexpected error: %v", err2)
 }
+
+func TestGetHeartbeatInterval(t *testing.T) {
+	defaultInterval := 1 * time.Minute
+
+	tests := []struct {
+		name string
+		env  string
+		want time.Duration
+	}{
+		{"unset", "", defaultInterval},
+		{"invalid", "invalid", defaultInterval},
+		{"zero", "0s", 0},
+		{"value", "5m", 5 * time.Minute},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.env != "" {
+				t.Setenv("MCP_PROXY_HEARTBEAT_INTERVAL", tt.env)
+			}
+			require.Equal(t, tt.want, getHeartbeatInterval(defaultInterval))
+		})
+	}
+}

tests/extproc/mcp/mcp_test.go

@@ -79,15 +79,6 @@ func TestMCP_forceJSONResponse(t *testing.T) {
 	}
 }
 
-func TestMCP_customWriteTimeout(t *testing.T) {
-	env := requireNewMCPEnv(t, false, 2*time.Second, defaultMCPPath)
-	for _, tc := range tests {
-		t.Run(tc.name+"/custom_write_timeout", func(t *testing.T) {
-			tc.testFn(t, env)
-		})
-	}
-}
-
 func TestMCP_differentPath(t *testing.T) {
 	env := requireNewMCPEnv(t, false, 1200*time.Second, "/mcp/yet/another/path")
 	t.Run("call", func(t *testing.T) {

tests/internal/testmcp/server.go

@@ -10,6 +10,7 @@ import (
 	"errors"
 	"fmt"
 	"log"
+	"net"
 	"net/http"
 	"os"
 	"time"
@@ -112,6 +113,9 @@ func NewServer(opts *Options) *http.Server {
 		// Allow long-lived connections.
 		WriteTimeout: opts.WriteTimeout,
 		Handler:      handler,
+		ConnState: func(conn net.Conn, state http.ConnState) {
+			log.Printf("MCP SERVER connection [%s] %s -> %s\n", state, conn.RemoteAddr(), conn.LocalAddr())
+		},
 	}
 	go func() {
 		log.Printf("starting MCP Streamable-HTTP server on :%d at /mcp", opts.Port)