feat(mcp): implement tool aggregation proxy and fix SSE stability

michiosw · michiosw · commit 9ead95c5f7e8 · 2025-04-12T12:34:23.000-07:00
- Add MCP server proxy that aggregates tools from remote servers
- Implement transparent proxying of tool calls to remote servers
- Copy tool schemas and parameters correctly from remote tools
- Set SSE read timeout to 1200 seconds for long-running operations
- Remove complex reconnection logic in favor of native SSE persistence

The MCP server now acts as a proxy that:
1. Connects to configured remote MCP servers
2. Aggregates their available tools
3. Exposes these tools locally with proper parameter schemas
4. Proxies tool calls to the appropriate remote server

Additionally, fixes SSE connection stability by properly configuring
timeouts and trusting the protocol's built-in persistence, rather than
implementing custom reconnection logic.
diff --git a/go.mod b/go.mod
@@ -7,16 +7,22 @@ require (
 	github.com/mattn/go-sqlite3 v1.14.27
 	github.com/mark3labs/mcp-go v0.18.0
 	github.com/rs/zerolog v1.34.0
+	github.com/stretchr/testify v1.9.0
 	go.uber.org/fx v1.23.0
 )
 
 require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/google/uuid v1.6.0 // indirect
 	github.com/mattn/go-colorable v0.1.14 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
 	go.uber.org/dig v1.18.1 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
 	go.uber.org/zap v1.27.0 // indirect
 	golang.org/x/sys v0.32.0 // indirect
 	github.com/google/uuid v1.6.0 // indirect
 	github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
diff --git a/go.sum b/go.sum
@@ -45,5 +45,7 @@ golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.32.0 h1:s77OFDvIQeibCmezSnk/q6iAfkdiQaJi4VzroCFrN20=
 golang.org/x/sys v0.32.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
diff --git a/internal/app/app.go b/internal/app/app.go
@@ -86,11 +86,10 @@ var BackendModule = fx.Module("backend",
 // MCPClientModule provides the MCP client (distinct from server frontend).
 var MCPClientModule = fx.Module("mcp_client",
 	fx.Provide(
-		mcp.NewServer,
+		mcp.NewMCPServer,
+		mcp.NewSSEServer,
 	),
-	fx.Invoke(func(server *mcp.Server, lc fx.Lifecycle) error {
-		return server.Start(lc)
-	}),
+	fx.Invoke(mcp.RegisterMCPServerHooks),
 )
 
 // UpdaterModule provides the tool and server updater service.
diff --git a/internal/mcp/client/client.go b/internal/mcp/client/client.go
@@ -2,11 +2,17 @@
 package client
 
 import (
+	"time"
+
 	"github.com/mark3labs/mcp-go/client"
 )
 
 // Re-export the SSEMCPClient from mark3labs/mcp-go/client
 type SSEMCPClient = client.SSEMCPClient
 
-// NewSSEMCPClient creates a new SSE MCP client
-var NewSSEMCPClient = client.NewSSEMCPClient
+// NewSSEMCPClient creates a new SSE MCP client with a 60 second timeout
+func NewSSEMCPClient(url string) (*SSEMCPClient, error) {
+	return client.NewSSEMCPClient(url,
+		client.WithSSEReadTimeout(1200*time.Second),
+	)
+}
diff --git a/internal/mcp/client/client_test.go b/internal/mcp/client/client_test.go
@@ -2,6 +2,9 @@ package client
 
 import (
 	"context"
+	"encoding/json"
+	"regexp"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -21,12 +24,14 @@ func setupMCPClient(t *testing.T, url string) *MCPConnection {
 	// Create a context for this specific connection
 	ctx, cancel := context.WithCancel(context.Background())
 
+	t.Logf("Creating new SSE client for %s", url)
 	client, err := NewSSEMCPClient(url)
 	if err != nil {
 		cancel()
 		t.Fatalf("Failed to create client for %s: %v", url, err)
 	}
 
+	t.Logf("Starting SSE client...")
 	if err := client.Start(ctx); err != nil {
 		cancel()
 		t.Fatalf("Failed to start client for %s: %v", url, err)
@@ -40,12 +45,13 @@ func setupMCPClient(t *testing.T, url string) *MCPConnection {
 		Version: "1.0.0",
 	}
 
+	t.Logf("Initializing SSE client...")
 	result, err := client.Initialize(ctx, initRequest)
 	if err != nil {
 		cancel()
 		t.Fatalf("Failed to initialize client for %s: %v", url, err)
 	}
-	t.Logf("Connected to server at %s: %s", url, result.ServerInfo.Name)
+	t.Logf("Successfully connected to server at %s: %s", url, result.ServerInfo.Name)
 
 	return &MCPConnection{
 		client: client,
@@ -55,11 +61,27 @@ func setupMCPClient(t *testing.T, url string) *MCPConnection {
 	}
 }
 
+func checkConnection(t *testing.T, conn *MCPConnection) error {
+	t.Log("Checking connection status...")
+	// Create a context with a short timeout for the check
+	checkCtx, cancel := context.WithTimeout(conn.ctx, 5*time.Second)
+	defer cancel()
+
+	// Try a simple ListTools request to verify connection
+	toolsRequest := mcp.ListToolsRequest{}
+	_, err := conn.client.ListTools(checkCtx, toolsRequest)
+	if err != nil {
+		t.Logf("Connection check failed: %v", err)
+		return err
+	}
+	t.Log("Connection check successful")
+	return nil
+}
+
 func TestMultipleMCPServers(t *testing.T) {
 	// Define our MCP servers
 	servers := []string{
-		"https://0.0.0.0:8000/sse",
-	}
+		"https://0.0.0.0:8000/sse"}
 
 	// Create connections to all servers
 	connections := make([]*MCPConnection, 0, len(servers))
@@ -105,3 +127,121 @@ func TestMultipleMCPServers(t *testing.T) {
 	// Wait for all operations to complete
 	wg.Wait()
 }
+
+func TestBrowserTools(t *testing.T) {
+	// Setup connection to the MCP server
+	conn := setupMCPClient(t, "https://0.0.0.0:8000/sse")
+	defer func() {
+		t.Log("Cleaning up connection...")
+		conn.cancel()       // Cancel the context first
+		conn.client.Close() // Then close the client
+		t.Log("Connection cleanup complete")
+	}()
+
+	// Create a context with timeout for the first operation
+	ctx1, cancel1 := context.WithTimeout(conn.ctx, 60*time.Second)
+	defer cancel1()
+
+	// First tool call: browser_use
+	t.Log("Making browser_use tool call...")
+	browserUseRequest := mcp.CallToolRequest{
+		Params: struct {
+			Name      string                 `json:"name"`
+			Arguments map[string]interface{} `json:"arguments,omitempty"`
+			Meta      *struct {
+				ProgressToken mcp.ProgressToken `json:"progressToken,omitempty"`
+			} `json:"_meta,omitempty"`
+		}{
+			Name: "browser_use",
+			Arguments: map[string]interface{}{
+				"action": "top 5 headlines",
+				"url":    "https://news.google.com/home?hl=en-US&gl=US&ceid=US:en",
+			},
+		},
+	}
+
+	browserUseResult, err := conn.client.CallTool(ctx1, browserUseRequest)
+	if err != nil {
+		t.Fatalf("Failed to call browser_use: %v", err)
+	}
+
+	// Log the raw result for debugging
+	t.Logf("Raw browser_use result: %+v", browserUseResult)
+
+	// Extract task ID from the result
+	var taskID string
+	if len(browserUseResult.Content) > 0 {
+		content := browserUseResult.Content[0]
+		// Convert the content to string for regex matching
+		contentBytes, err := json.Marshal(content)
+		if err != nil {
+			t.Fatalf("Failed to marshal content: %v", err)
+		}
+		contentStr := string(contentBytes)
+
+		// First extract the text field content
+		textRegex := regexp.MustCompile(`"text":\s*"({[^}]+})"`)
+		textMatches := textRegex.FindStringSubmatch(contentStr)
+		if len(textMatches) < 2 {
+			t.Fatalf("Could not find text content in response: %s", contentStr)
+		}
+
+		// The text field contains escaped JSON, so we need to unescape it
+		textContent := textMatches[1]
+		textContent = strings.ReplaceAll(textContent, `\n`, "\n")
+		textContent = strings.ReplaceAll(textContent, `\"`, `"`)
+
+		// Now extract the task_id from the unescaped JSON
+		taskIDRegex := regexp.MustCompile(`"task_id":\s*"([^"]+)"`)
+		matches := taskIDRegex.FindStringSubmatch(textContent)
+		if len(matches) < 2 {
+			t.Fatalf("Could not find task_id in text content: %s", textContent)
+		}
+		taskID = matches[1]
+		t.Logf("Got task ID: %s", taskID)
+	} else {
+		t.Fatalf("No content in result")
+	}
+
+	// Wait 40 seconds before checking result
+	t.Log("Waiting 40 seconds before checking result...")
+	time.Sleep(40 * time.Second)
+
+	// Create a fresh context for the second operation
+	ctx2, cancel2 := context.WithTimeout(conn.ctx, 60*time.Second)
+	defer cancel2()
+
+	// Second tool call: browser_get_result
+	t.Log("Making browser_get_result tool call...")
+	browserGetResultRequest := mcp.CallToolRequest{
+		Params: struct {
+			Name      string                 `json:"name"`
+			Arguments map[string]interface{} `json:"arguments,omitempty"`
+			Meta      *struct {
+				ProgressToken mcp.ProgressToken `json:"progressToken,omitempty"`
+			} `json:"_meta,omitempty"`
+		}{
+			Name: "browser_get_result",
+			Arguments: map[string]interface{}{
+				"task_id": taskID,
+			},
+		},
+	}
+
+	browserGetResultResult, err := conn.client.CallTool(ctx2, browserGetResultRequest)
+	if err != nil {
+		t.Fatalf("Failed to call browser_get_result: %v", err)
+	}
+
+	// Log the final result
+	if len(browserGetResultResult.Content) > 0 {
+		content := browserGetResultResult.Content[0]
+		contentBytes, err := json.Marshal(content)
+		if err != nil {
+			t.Fatalf("Failed to marshal content: %v", err)
+		}
+		t.Logf("Final result: %s", string(contentBytes))
+	} else {
+		t.Fatalf("No content in final result")
+	}
+}
diff --git a/internal/mcp/server.go b/internal/mcp/server.go
