Handle partial flushes for application/json (#2422)

This adds code fixing the issue #2403 for `application/json` content
type.

Author: blkt

pkg/mcp/tool_filter.go

@@ -336,7 +336,7 @@ func (rw *toolFilterWriter) Flush() {
 
 		var b bytes.Buffer
 		err := processBuffer(rw.config, rw.buffer, mimeType, &b)
-		if err == errKeepBuffering {
+		if errors.Is(err, errKeepBuffering) {
 			logger.Debugf("Buffered %d so far, keep buffering...", len(rw.buffer))
 			return
 		}
@@ -386,7 +386,11 @@ func processBuffer(
 	switch mimeType {
 	case "application/json":
 		var toolsListResponse toolsListResponse
+		var syntaxError *json.SyntaxError
 		err := json.Unmarshal(buffer, &toolsListResponse)
+		if errors.As(err, &syntaxError) {
+			return fmt.Errorf("%w: %v", errKeepBuffering, err)
+		}
 		if err == nil && toolsListResponse.Result.Tools != nil {
 			return processToolsListResponse(config, toolsListResponse, w)
 		}
@@ -416,7 +420,7 @@ func processEventStream(
 	w io.Writer,
 ) error {
 	if len(buffer) > 1 && buffer[len(buffer)-1] != '\n' && buffer[len(buffer)-1] != '\r' {
-		return errKeepBuffering
+		return fmt.Errorf("%w: %v", errKeepBuffering, "event separator not found")
 	}
 
 	// NOTE: this looks uglier, but is more efficient than scanning the whole buffer

pkg/mcp/tool_middleware_test.go

@@ -117,23 +117,6 @@ func TestNewListToolsMappingMiddleware_Scenarios(t *testing.T) {
 				{"name": "MyFoo", "description": "Foo tool"},
 			},
 		},
-		{
-			name: "Filter MyFoo, Override Foo -> MyFoo with connection hang",
-			serverOpts: []testkit.TestMCPServerOption{
-				//nolint:goconst
-				testkit.WithTool("Foo", "Foo tool", func() string { return "Foo" }),
-				//nolint:goconst
-				testkit.WithTool("Bar", "Bar tool", func() string { return "Bar" }),
-				testkit.WithConnectionHang(10 * time.Second),
-			},
-			opts: &[]ToolMiddlewareOption{
-				WithToolsFilter("MyFoo"),
-				WithToolsOverride("Foo", "MyFoo", ""),
-			},
-			expected: &[]map[string]any{
-				{"name": "MyFoo", "description": "Foo tool"},
-			},
-		},
 	}
 
 	for _, tt := range tests {
@@ -631,7 +614,7 @@ func TestNewToolCallMappingMiddleware_ErrorCases(t *testing.T) {
 	}
 }
 
-func TestNewListToolsMappingMiddleware_ConnectionHang(t *testing.T) {
+func TestSSEBufferFlushes(t *testing.T) {
 	t.Parallel()
 	middlewares := []func(http.Handler) http.Handler{}
 
@@ -684,4 +667,61 @@ func TestNewListToolsMappingMiddleware_ConnectionHang(t *testing.T) {
 	require.NoError(t, err)
 	require.NotNil(t, response.Result)
 	require.NotNil(t, response.Result.Tools)
+	require.Len(t, *response.Result.Tools, 1)
+}
+
+func TestJSONBufferFlushes(t *testing.T) {
+	t.Parallel()
+	middlewares := []func(http.Handler) http.Handler{}
+
+	opts := []ToolMiddlewareOption{
+		WithToolsFilter("MyFoo"),
+		WithToolsOverride("Foo", "MyFoo", ""),
+	}
+
+	// Create the middleware
+	toolsListmiddleware, err := NewListToolsMappingMiddleware(opts...)
+	assert.NoError(t, err)
+	toolsCallMiddleware, err := NewToolCallMappingMiddleware(opts...)
+	assert.NoError(t, err)
+
+	middlewares = append(middlewares,
+		toolsCallMiddleware,
+		toolsListmiddleware,
+	)
+
+	// Create test server
+	serverOpts := []testkit.TestMCPServerOption{
+		testkit.WithJSONClientType(),
+		testkit.WithConnectionHang(10 * time.Second),
+		testkit.WithMiddlewares(middlewares...),
+		testkit.WithWithProxy(),
+		testkit.WithTool("Foo", "Foo tool", func() string { return "Foo" }),
+	}
+
+	for i := 0; i < 100; i++ {
+		opt := testkit.WithTool(
+			fmt.Sprintf("Foo%d", i),
+			strings.Repeat("A", 10*1024),
+			func() string { return fmt.Sprintf("Foo%d", i) },
+		)
+		serverOpts = append(serverOpts, opt)
+	}
+
+	server, client, err := testkit.NewStreamableTestServer(
+		serverOpts...,
+	)
+	require.NoError(t, err)
+	defer server.Close()
+
+	// Make request
+	respBody, err := client.ToolsList()
+	require.NoError(t, err)
+
+	var response toolsListResponse
+	err = json.NewDecoder(bytes.NewReader(respBody)).Decode(&response)
+	require.NoError(t, err)
+	require.NotNil(t, response.Result)
+	require.NotNil(t, response.Result.Tools)
+	require.Len(t, *response.Result.Tools, 1)
 }

test/testkit/sse_server.go

@@ -9,8 +9,6 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
-	"net/http/httputil"
-	"net/url"
 	"time"
 
 	"github.com/go-chi/chi/v5"
@@ -225,28 +223,10 @@ func NewSSETestServer(
 	// If the server is configured to use a proxy,create a reverse proxy to
 	// the backend test server.
 	if server.withProxy {
-		backendURL, err := url.Parse(backendServer.URL)
+		proxyServer, err := wrapBackendWithProxy(backendServer.URL, allMiddlewares)
 		if err != nil {
-			return nil, nil, fmt.Errorf("failed to parse backend URL: %w", err)
+			return nil, nil, fmt.Errorf("failed to wrap backend with proxy: %w", err)
 		}
-
-		// Create a reverse proxy to the backend test server.
-		// Ideally, this would use ToolHive reverse proxy, but
-		// it is too tightly coupled with containers and needs
-		// to be refactored.
-		proxy := httputil.NewSingleHostReverseProxy(backendURL)
-		proxy.FlushInterval = -1
-		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			proxy.ServeHTTP(w, r)
-		})
-
-		// Apply middleware chain in reverse order (last middleware is applied first)
-		var finalHandler http.Handler = handler
-		for _, mw := range allMiddlewares {
-			finalHandler = mw(finalHandler)
-		}
-
-		proxyServer := httptest.NewServer(finalHandler)
 		testServer = proxyServer
 	}
 
@@ -319,7 +299,7 @@ func (s *sseServer) sseHandler(w http.ResponseWriter, _ *http.Request) {
 	w.Header().Set("Content-Type", "text/event-stream")
 
 	// Get flusher for streaming responses
-	flusher, ok := w.(http.Flusher)
+	_, ok := w.(http.Flusher)
 	if !ok {
 		http.Error(w, "Streaming unsupported", http.StatusInternalServerError)
 		return
@@ -356,18 +336,12 @@ func (s *sseServer) sseHandler(w http.ResponseWriter, _ *http.Request) {
 					response = "failed to generate response"
 				}
 
-				if _, err := w.Write([]byte("event: random-stuff\ndata: " + response + "\n\n")); err != nil {
-					http.Error(w, "Error writing response", http.StatusInternalServerError)
-					return
+				if s.connHangDuration == 0 {
+					singleFlushResponse([]byte("event: random-stuff\ndata: "+response+"\n\n"), w)
+				} else {
+					staggeredFlushResponse([]byte("event: random-stuff\ndata: "+response+"\n\n"), w, s.connHangDuration)
 				}
 			}
-
-			// Flush the response immediately
-			flusher.Flush()
-
-			if s.connHangDuration != 0 {
-				time.Sleep(s.connHangDuration)
-			}
 		}
 	}
 }

test/testkit/streamable_server.go

@@ -9,8 +9,6 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
-	"net/http/httputil"
-	"net/url"
 	"time"
 
 	"github.com/go-chi/chi/v5"
@@ -202,18 +200,19 @@ func NewStreamableTestServer(
 	// This precompiles the tools list response based on the provided tools
 	server.toolsListResponse = makeToolsList(server.tools)
 
+	allMiddlewares := append(
+		[]func(http.Handler) http.Handler{
+			middleware.RequestID,
+			middleware.Recoverer,
+		},
+		server.middlewares...,
+	)
+
 	router := chi.NewRouter()
 
 	// If the server is not configured to use a proxy, apply the middlewares to
 	// the router directly.
 	if !server.withProxy {
-		allMiddlewares := append(
-			[]func(http.Handler) http.Handler{
-				middleware.RequestID,
-				middleware.Recoverer,
-			},
-			server.middlewares...,
-		)
 		router.Use(allMiddlewares...)
 	}
 
@@ -229,35 +228,10 @@ func NewStreamableTestServer(
 	// If the server is configured to use a proxy,create a reverse proxy to
 	// the backend test server.
 	if server.withProxy {
-		backendURL, err := url.Parse(backendServer.URL)
+		proxyServer, err := wrapBackendWithProxy(backendServer.URL, allMiddlewares)
 		if err != nil {
-			return nil, nil, fmt.Errorf("failed to parse backend URL: %w", err)
-		}
-
-		// Create a reverse proxy to the backend test server.
-		// Ideally, this would use ToolHive reverse proxy, but
-		// it is too tightly coupled with containers and needs
-		// to be refactored.
-		proxy := httputil.NewSingleHostReverseProxy(backendURL)
-		proxy.FlushInterval = -1
-		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			proxy.ServeHTTP(w, r)
-		})
-
-		// Apply middleware chain in reverse order (last middleware is applied first)
-		allMiddlewares := append(
-			[]func(http.Handler) http.Handler{
-				middleware.RequestID,
-				middleware.Recoverer,
-			},
-			server.middlewares...,
-		)
-		var finalHandler http.Handler = handler
-		for _, mw := range allMiddlewares {
-			finalHandler = mw(finalHandler)
+			return nil, nil, fmt.Errorf("failed to wrap backend with proxy: %w", err)
 		}
-
-		proxyServer := httptest.NewServer(finalHandler)
 		testServer = proxyServer
 	}
 
@@ -277,7 +251,10 @@ func NewStreamableTestServer(
 	}
 }
 
-func (s *streamableServer) mcpJSONHandler(w http.ResponseWriter, r *http.Request) {
+func (s *streamableServer) mcpJSONHandler(
+	w http.ResponseWriter,
+	r *http.Request,
+) {
 	// Read the request body
 	body, err := io.ReadAll(r.Body)
 	if err != nil {
@@ -319,23 +296,18 @@ func (s *streamableServer) mcpJSONHandler(w http.ResponseWriter, r *http.Request
 
 	w.Header().Set("Content-Type", "application/json")
 	w.WriteHeader(http.StatusOK)
-	if _, err := w.Write([]byte(response)); err != nil {
-		http.Error(w, "Error writing response", http.StatusInternalServerError)
-		return
-	}
 
-	// Flush if available
-	if flusher, ok := w.(http.Flusher); ok {
-		flusher.Flush()
-	}
-
-	if s.connHangDuration != 0 {
-		time.Sleep(s.connHangDuration)
+	if s.connHangDuration == 0 {
+		singleFlushResponse([]byte(response), w)
+	} else {
+		staggeredFlushResponse([]byte(response), w, s.connHangDuration)
 	}
 }
 
-//nolint:gocyclo
-func (s *streamableServer) mcpEventStreamHandler(w http.ResponseWriter, r *http.Request) {
+func (s *streamableServer) mcpEventStreamHandler(
+	w http.ResponseWriter,
+	r *http.Request,
+) {
 	// Read the request body
 	body, err := io.ReadAll(r.Body)
 	if err != nil {
@@ -378,41 +350,11 @@ func (s *streamableServer) mcpEventStreamHandler(w http.ResponseWriter, r *http.
 		response = "failed to generate response"
 	}
 
-	if _, err := w.Write([]byte("event: random-stuff\ndata: " + response)); err != nil {
-		http.Error(w, "Error writing response", http.StatusInternalServerError)
-		return
-	}
+	response = "event: random-stuff\ndata: " + response + "\n\n"
 
 	if s.connHangDuration == 0 {
-		_, err := w.Write([]byte("\n\n"))
-		if err != nil {
-			http.Error(w, "Error writing response", http.StatusInternalServerError)
-			return
-		}
-
-		// Flush if available
-		if flusher, ok := w.(http.Flusher); ok {
-			flusher.Flush()
-		}
+		singleFlushResponse([]byte(response), w)
 	} else {
-		// Flush if available
-		if flusher, ok := w.(http.Flusher); ok {
-			flusher.Flush()
-		}
-
-		if s.connHangDuration != 0 {
-			time.Sleep(s.connHangDuration)
-		}
-
-		_, err := w.Write([]byte("\n\n"))
-		if err != nil {
-			http.Error(w, "Error writing response", http.StatusInternalServerError)
-			return
-		}
-
-		// Flush if available
-		if flusher, ok := w.(http.Flusher); ok {
-			flusher.Flush()
-		}
+		staggeredFlushResponse([]byte(response), w, s.connHangDuration)
 	}
 }