feat(mcp): add SSE server and multi-server client

michiosw · michiosw · commit 603603a593fa · 2025-04-12T12:31:22.000-07:00
Implements MCP server with hello_world tool and client with multi-server connection support
diff --git a/go.mod b/go.mod
@@ -5,6 +5,7 @@ go 1.24.1
 require (
 	github.com/a-h/templ v0.3.857
 	github.com/mattn/go-sqlite3 v1.14.27
+	github.com/mark3labs/mcp-go v0.18.0
 	github.com/rs/zerolog v1.34.0
 	go.uber.org/fx v1.23.0
 )
@@ -16,4 +17,6 @@ require (
 	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
 )
diff --git a/go.sum b/go.sum
@@ -6,6 +6,11 @@ github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs
 github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
 github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
+github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/mark3labs/mcp-go v0.18.0 h1:YuhgIVjNlTG2ZOwmrkORWyPTp0dz1opPEqvsPtySXao=
+github.com/mark3labs/mcp-go v0.18.0/go.mod h1:KmJndYv7GIgcPVwEKJjNcbhVQ+hJGJhrCCB/9xITzpE=
+github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
 github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
 github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE=
 github.com/mattn/go-colorable v0.1.14/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stgPZH1UqBm1s8=
@@ -21,10 +26,12 @@ github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZN
 github.com/rs/xid v1.6.0/go.mod h1:7XoLgs4eV+QndskICGsho+ADou8ySMSjJKDIan90Nz0=
 github.com/rs/zerolog v1.34.0 h1:k43nTLIwcTVQAncfCw4KZ2VY6ukYoZaBPNOE8txlOeY=
 github.com/rs/zerolog v1.34.0/go.mod h1:bJsvje4Z08ROH4Nhs5iH600c3IkWhwp44iRc54W6wYQ=
-github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
-github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
-go.uber.org/dig v1.18.1 h1:rLww6NuajVjeQn+49u5NcezUJEGwd5uXmyoCKW2g5Es=
-go.uber.org/dig v1.18.1/go.mod h1:Us0rSJiThwCv2GteUN0Q7OKvU7n5J4dxZ9JKUXozFdE=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
+github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
+go.uber.org/dig v1.18.0 h1:imUL1UiY0Mg4bqbFfsRQO5G4CGRBec/ZujWTvSVp3pw=
+go.uber.org/dig v1.18.0/go.mod h1:Us0rSJiThwCv2GteUN0Q7OKvU7n5J4dxZ9JKUXozFdE=
 go.uber.org/fx v1.23.0 h1:lIr/gYWQGfTwGcSXWXu4vP5Ws6iqnNEIY+F/aFzCKTg=
 go.uber.org/fx v1.23.0/go.mod h1:o/D9n+2mLP6v1EG+qsdT1O8wKopYAsqZasju97SDFCU=
 go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
diff --git a/internal/app/app.go b/internal/app/app.go
@@ -11,8 +11,7 @@ import (
 	"github.com/co-browser/agent-browser/internal/config"
 	"github.com/co-browser/agent-browser/internal/events"
 	"github.com/co-browser/agent-browser/internal/log"
-	"github.com/co-browser/agent-browser/internal/mcp/client"
-	"github.com/co-browser/agent-browser/internal/updater"
+	"github.com/co-browser/agent-browser/internal/mcp"
 	"github.com/co-browser/agent-browser/internal/web"
 	"github.com/co-browser/agent-browser/internal/web/handlers"
 
@@ -87,8 +86,11 @@ var BackendModule = fx.Module("backend",
 // MCPClientModule provides the MCP client (distinct from server frontend).
 var MCPClientModule = fx.Module("mcp_client",
 	fx.Provide(
-		client.NewHTTPClient,
+		mcp.NewServer,
 	),
+	fx.Invoke(func(server *mcp.Server, lc fx.Lifecycle) error {
+		return server.Start(lc)
+	}),
 )
 
 // UpdaterModule provides the tool and server updater service.
diff --git a/internal/mcp/client/client.go b/internal/mcp/client/client.go
@@ -2,76 +2,11 @@
 package client
 
 import (
-	"encoding/json"
-	"fmt"
-	"net/http"
-	"time"
-
-	"github.com/co-browser/agent-browser/internal/backend/models"
-	"github.com/co-browser/agent-browser/internal/log" // Import log
+	"github.com/mark3labs/mcp-go/client"
 )
 
-// Client defines the interface for fetching tool metadata from an MCP server.
-type Client interface {
-	// FetchTools retrieves the list of tools from the given MCP server URL.
-	// It should return an error if the fetch fails or the response is invalid.
-	FetchTools(serverURL string) ([]models.FetchedTool, error)
-}
-
-// httpClient implements the Client interface using standard net/http.
-type httpClient struct {
-	client *http.Client
-	logger log.Logger // Inject logger
-}
-
-// NewHTTPClient creates a new MCP client with a default HTTP timeout.
-func NewHTTPClient(logger log.Logger) Client { // Accept logger
-	return &httpClient{
-		client: &http.Client{
-			Timeout: 15 * time.Second, // Sensible default timeout
-		},
-		logger: logger, // Store logger
-	}
-}
-
-// FetchTools implements the Client interface.
-// This assumes the MCP server exposes a simple GET endpoint (e.g., /tools) returning JSON.
-// This might need significant adjustments based on the actual MCP protocol.
-func (c *httpClient) FetchTools(serverURL string) ([]models.FetchedTool, error) {
-	// Assume a standard endpoint like /tools, trim trailing slash if present
-	baseURL := serverURL
-	if baseURL[len(baseURL)-1] == '/' {
-		baseURL = baseURL[:len(baseURL)-1]
-	}
-	targetURL := baseURL + "/tools" // TODO: Make this configurable or discoverable?
-
-	c.logger.Debug().Str("url", targetURL).Msg("Fetching tools") // Use logger
-	req, err := http.NewRequest("GET", targetURL, nil)
-	if err != nil {
-		// Log error if desired, but return wrapped error regardless
-		return nil, fmt.Errorf("failed to create request for %s: %w", targetURL, err)
-	}
-	// TODO: Add any necessary headers (e.g., Accept: application/json)
-	req.Header.Set("Accept", "application/json")
-
-	resp, err := c.client.Do(req)
-	if err != nil {
-		return nil, fmt.Errorf("failed to fetch tools from %s: %w", targetURL, err)
-	}
-	defer func() {
-		_ = resp.Body.Close() // Explicitly ignore Close error
-	}()
-
-	if resp.StatusCode != http.StatusOK {
-		// TODO: Read response body for more error details?
-		return nil, fmt.Errorf("failed to fetch tools from %s: received status code %d", targetURL, resp.StatusCode)
-	}
-
-	var fetchedTools []models.FetchedTool
-	if err := json.NewDecoder(resp.Body).Decode(&fetchedTools); err != nil {
-		return nil, fmt.Errorf("failed to decode tools response from %s: %w", targetURL, err)
-	}
+// Re-export the SSEMCPClient from mark3labs/mcp-go/client
+type SSEMCPClient = client.SSEMCPClient
 
-	c.logger.Info().Int("count", len(fetchedTools)).Str("url", targetURL).Msg("Successfully fetched tools") // Use logger
-	return fetchedTools, nil
-}
+// NewSSEMCPClient creates a new SSE MCP client
+var NewSSEMCPClient = client.NewSSEMCPClient
diff --git a/internal/mcp/client/client_test.go b/internal/mcp/client/client_test.go
@@ -0,0 +1,107 @@
+package client
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/mark3labs/mcp-go/mcp"
+)
+
+// MCPConnection represents a connection to an MCP server
+type MCPConnection struct {
+	client *SSEMCPClient
+	url    string
+	ctx    context.Context
+	cancel context.CancelFunc
+}
+
+func setupMCPClient(t *testing.T, url string) *MCPConnection {
+	// Create a context for this specific connection
+	ctx, cancel := context.WithCancel(context.Background())
+
+	client, err := NewSSEMCPClient(url)
+	if err != nil {
+		cancel()
+		t.Fatalf("Failed to create client for %s: %v", url, err)
+	}
+
+	if err := client.Start(ctx); err != nil {
+		cancel()
+		t.Fatalf("Failed to start client for %s: %v", url, err)
+	}
+
+	// Initialize the client
+	initRequest := mcp.InitializeRequest{}
+	initRequest.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION
+	initRequest.Params.ClientInfo = mcp.Implementation{
+		Name:    "test-client",
+		Version: "1.0.0",
+	}
+
+	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)
+
+	return &MCPConnection{
+		client: client,
+		url:    url,
+		ctx:    ctx,
+		cancel: cancel,
+	}
+}
+
+func TestMultipleMCPServers(t *testing.T) {
+	// Define our MCP servers
+	servers := []string{
+		"https://0.0.0.0:8000/sse",
+	}
+
+	// Create connections to all servers
+	connections := make([]*MCPConnection, 0, len(servers))
+	for _, url := range servers {
+		conn := setupMCPClient(t, url)
+		connections = append(connections, conn)
+		// Ensure we clean up each connection properly
+		defer func(conn *MCPConnection) {
+			conn.cancel()       // Cancel the context first
+			conn.client.Close() // Then close the client
+		}(conn)
+	}
+
+	// Create a WaitGroup to wait for all concurrent operations
+	var wg sync.WaitGroup
+
+	// Test ListTools on all servers concurrently
+	for _, conn := range connections {
+		wg.Add(1)
+		go func(conn *MCPConnection) {
+			defer wg.Done()
+
+			// Create a timeout context for this specific operation
+			opCtx, opCancel := context.WithTimeout(conn.ctx, 5*time.Second)
+			defer opCancel()
+
+			// Test ListTools
+			toolsRequest := mcp.ListToolsRequest{}
+			toolsResult, err := conn.client.ListTools(opCtx, toolsRequest)
+			if err != nil {
+				t.Errorf("Failed to list tools for %s: %v", conn.url, err)
+				return
+			}
+
+			// Log the available tools
+			t.Logf("Available tools for %s:", conn.url)
+			for _, tool := range toolsResult.Tools {
+				t.Logf("- %s: %s", tool.Name, tool.Description)
+			}
+		}(conn)
+	}
+
+	// Wait for all operations to complete
+	wg.Wait()
+}
diff --git a/internal/mcp/server.go b/internal/mcp/server.go
@@ -1,2 +1,95 @@
 // Package mcp implements the MCP server logic.
 package mcp
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"github.com/co-browser/agent-browser/internal/log"
+	"github.com/mark3labs/mcp-go/mcp"
+	"github.com/mark3labs/mcp-go/server"
+	"go.uber.org/fx"
+)
+
+// Server wraps the MCP server instance
+type Server struct {
+	mcpServer *server.MCPServer
+	sseServer *server.SSEServer
+	logger    log.Logger
+	addr      string
+}
+
+// ServerParams contains the parameters for creating a new Server
+type ServerParams struct {
+	fx.In
+	Logger log.Logger
+}
+
+// NewServer creates a new MCP server with configured tools
+func NewServer(p ServerParams) *Server {
+	// Create MCP server
+	mcpServer := server.NewMCPServer(
+		"CoBrowser Agent 🚀",
+		"1.0.0",
+		server.WithToolCapabilities(true),
+	)
+
+	// Add hello world tool
+	tool := mcp.NewTool("hello_world",
+		mcp.WithDescription("Say hello to someone"),
+		mcp.WithString("name",
+			mcp.Required(),
+			mcp.Description("Name of the person to greet"),
+		),
+	)
+
+	// Add tool handler
+	mcpServer.AddTool(tool, helloHandler)
+
+	// Create SSE server
+	sseServer := server.NewSSEServer(mcpServer)
+
+	return &Server{
+		mcpServer: mcpServer,
+		sseServer: sseServer,
+		logger:    p.Logger,
+		addr:      ":8087", // Default to same port as web server
+	}
+}
+
+// SetAddr sets the address for the server to listen on
+func (s *Server) SetAddr(addr string) {
+	s.addr = addr
+}
+
+// Start starts the MCP server
+func (s *Server) Start(lc fx.Lifecycle) error {
+	lc.Append(fx.Hook{
+		OnStart: func(ctx context.Context) error {
+			// Start the server in a goroutine
+			go func() {
+				s.logger.Info().Str("addr", s.addr).Msg("Starting MCP SSE server")
+				if err := s.sseServer.Start(s.addr); err != nil {
+					s.logger.Fatal().Err(err).Msg("MCP SSE server Start() failed")
+				}
+			}()
+			return nil
+		},
+		OnStop: func(ctx context.Context) error {
+			s.logger.Info().Msg("Shutting down MCP server...")
+			// The SSE server doesn't provide a way to gracefully shutdown yet
+			return nil
+		},
+	})
+	return nil
+}
+
+func helloHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+	name, ok := request.Params.Arguments["name"].(string)
+	if !ok {
+		return nil, errors.New("name must be a string")
+	}
+
+	return mcp.NewToolResultText(fmt.Sprintf("Hello, %s!", name)), nil
+}

Original file line number	Diff line number	Diff line change
`@@ -5,6 +5,7 @@ go 1.24.1`
`5`	`5`	`require (`
`6`	`6`	`github.com/a-h/templ v0.3.857`
`7`	`7`	`github.com/mattn/go-sqlite3 v1.14.27`
	`8`	`+ github.com/mark3labs/mcp-go v0.18.0`
`8`	`9`	`github.com/rs/zerolog v1.34.0`
`9`	`10`	`go.uber.org/fx v1.23.0`
`10`	`11`	`)`
`@@ -16,4 +17,6 @@ require (`
`16`	`17`	`go.uber.org/multierr v1.11.0 // indirect`
`17`	`18`	`go.uber.org/zap v1.27.0 // indirect`
`18`	`19`	`golang.org/x/sys v0.32.0 // indirect`
	`20`	`+ github.com/google/uuid v1.6.0 // indirect`
	`21`	`+ github.com/yosida95/uritemplate/v3 v3.0.2 // indirect`
`19`	`22`	`)`