Scaled MCP Server

A horizontally scalable MCP (Message Context Protocol) server implementation that supports load-balanced deployments.

Overview

The Scaled MCP Server is a Go library that implements the MCP 2025-03 specification with support for horizontal scaling. It's designed to be embedded in your application and provides flexible configuration options.

Features

• HTTP Transport: Flexible HTTP transport with main /mcp endpoint, optional SSE endpoint, and capabilities negotiation
• Session Management: Distributed session management with Redis or in-memory options
• Actor System: Uses an actor-based architecture for handling sessions and message routing
• Horizontal Scaling: Support for load-balanced deployments across multiple nodes

Installation

go get github.com/traego/scaled-mcp@latest

Note: This library requires Go 1.24 or higher.

Usage

Basic Server with Static Tool

package main

import (
	"context"
	"fmt"
	"log/slog"
	"os"
	"os/signal"
	"syscall"

	"github.com/traego/scaled-mcp/pkg/config"
	"github.com/traego/scaled-mcp/pkg/resources"
	"github.com/traego/scaled-mcp/pkg/server"
)

func main() {
	// Configure logging
	logHandler := slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelDebug})
	slog.SetDefault(slog.New(logHandler))

	// Create a server with default configuration
	cfg := config.DefaultConfig()
	
	// Use in-memory session store for simplicity
	cfg.Session.UseInMemory = true
	
	// Create a static tool registry
	registry := resources.NewStaticToolRegistry()
	
	// Define and register a simple calculator tool
	calculatorTool := resources.NewTool("calculator").
		WithDescription("Performs basic arithmetic operations").
		WithInputs([]resources.ToolInput{
			{
				Name:        "operation",
				Type:        "string",
				Description: "Operation to perform (add, subtract, multiply, divide)",
				Required:    true,
			},
			{
				Name:        "a",
				Type:        "number",
				Description: "First operand",
				Required:    true,
			},
			{
				Name:        "b",
				Type:        "number",
				Description: "Second operand",
				Required:    true,
			},
		}).
		Build()
	
	// Register the tool with the registry
	registry.RegisterTool(calculatorTool)
	
	// Define a prompt for the server
	prompt := "You are a helpful AI assistant that can perform calculations using the calculator tool."
	
	// Create the server with the tool registry and prompt
	srv, err := server.NewMcpServer(cfg,
		server.WithToolRegistry(registry),
		server.WithServerInfo("Example MCP Server", "1.0.0"),
		server.WithPrompt(prompt),
	)
	if err != nil {
		slog.Error("Failed to create server", "error", err)
		os.Exit(1)
	}
	
	// Set up the tool handler
	registry.SetToolHandler("calculator", func(ctx context.Context, params map[string]interface{}) (interface{}, error) {
		// Extract parameters
		operation, ok := params["operation"].(string)
		if !ok {
			return nil, fmt.Errorf("%w: operation must be a string", resources.ErrInvalidParams)
		}
		
		a, ok := params["a"].(float64)
		if !ok {
			return nil, fmt.Errorf("%w: a must be a number", resources.ErrInvalidParams)
		}
		
		b, ok := params["b"].(float64)
		if !ok {
			return nil, fmt.Errorf("%w: b must be a number", resources.ErrInvalidParams)
		}
		
		// Perform the calculation
		var result float64
		switch operation {
		case "add":
			result = a + b
		case "subtract":
			result = a - b
		case "multiply":
			result = a * b
		case "divide":
			if b == 0 {
				return nil, fmt.Errorf("%w: division by zero", resources.ErrInvalidParams)
			}
			result = a / b
		default:
			return nil, fmt.Errorf("%w: unknown operation %s", resources.ErrInvalidParams, operation)
		}
		
		return map[string]interface{}{
			"result": result,
		}, nil
	})
	
	// Start the server in a goroutine
	go func() {
		if err := srv.Start(context.Background()); err != nil {
			slog.Error("Failed to start server", "error", err)
			os.Exit(1)
		}
	}()
	
	slog.Info("Server started", "host", cfg.HTTP.Host, "port", cfg.HTTP.Port)
	
	// Wait for termination signal
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
	<-sig
	
	// Shutdown the server
	slog.Info("Shutting down server...")
	shutdownCtx, cancel := context.WithTimeout(context.Background(), cfg.HTTP.ShutdownTimeout)
	defer cancel()
	
	if err := srv.Stop(shutdownCtx); err != nil {
		slog.Error("Failed to stop server", "error", err)
	}
	
	slog.Info("Server stopped")
}

### Using an External HTTP Server

You can use your own HTTP server with the MCP transport:

```go
package main

import (
	"context"
	"log"
	"net/http"

	"github.com/go-chi/chi/v5"
	"github.com/go-chi/chi/v5/middleware"
	"github.com/go-chi/cors"
	"github.com/traego/scaled-mcp/pkg/config"
	"github.com/traego/scaled-mcp/pkg/server"
	"github.com/traego/scaled-mcp/pkg/transport"
)

func main() {
	// Create a server with default configuration
	cfg := config.DefaultConfig()
	cfg.Session.UseInMemory = true
	
	// Create the MCP server but don't start the HTTP server
	srv, err := server.NewServer(cfg)
	if err != nil {
		log.Fatalf("Failed to create server: %v", err)
	}
	
	// Create a custom router
	r := chi.NewRouter()
	
	// Add middleware
	r.Use(middleware.RequestID)
	r.Use(middleware.RealIP)
	r.Use(middleware.Logger)
	r.Use(middleware.Recoverer)
	
	// Add CORS middleware - important when using an external server
	r.Use(cors.Handler(cors.Options{
		AllowedOrigins:   []string{"*"},
		AllowedMethods:   []string{"GET", "POST", "PUT", "DELETE", "OPTIONS"},
		AllowedHeaders:   []string{"Accept", "Authorization", "Content-Type", "X-CSRF-Token"},
		ExposedHeaders:   []string{"Link"},
		AllowCredentials: true,
		MaxAge:           300,
	}))
	
	// Add your custom routes
	r.Get("/", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("Welcome to the MCP server!"))
	})
	
	// Create the HTTP transport with the custom router
	httpTransport := transport.NewHTTPTransport(
		cfg,
		srv.GetActorSystem(),
		srv.GetSessionManager(),
		transport.WithExternalRouter(r),
	)
	
	// Start the MCP server without HTTP
	if err := srv.Start(context.Background()); err != nil {
		log.Fatalf("Failed to start server: %v", err)
	}
	
	// Start the HTTP transport
	if err := httpTransport.Start(); err != nil {
		log.Fatalf("Failed to start HTTP transport: %v", err)
	}
	
	// Start your HTTP server
	log.Printf("Starting HTTP server on %s:%d", cfg.HTTP.Host, cfg.HTTP.Port)
	if err := http.ListenAndServe(":8080", r); err != nil {
		log.Fatalf("HTTP server error: %v", err)
	}
}

## Dynamic Tool Registry Example

This library supports both static and dynamic tool registries. Here's an client_example of using a dynamic tool registry:

```go
// Create a custom tool provider
toolProvider := NewExampleToolProvider()

// Create a dynamic tool registry with the provider
registry := resources.NewDynamicToolRegistry(toolProvider)

// Create server with the dynamic tool registry
cfg := config.DefaultConfig()
mcpServer, err := server.NewMcpServer(cfg,
    server.WithToolRegistry(registry),
)

Tool Definition

The library provides three ways to define tool inputs:

1. Using Struct-Based Reflection (Recommended)

// Define input struct with mcp tags
type CalculatorInput struct {
    Operation string  `mcp:"operation,The operation to perform (add subtract multiply divide),required"`
    A         float64 `mcp:"a,First operand,required"`
    B         float64 `mcp:"b,Second operand,required"`
}

// Type-safe handler function
func calculatorHandler(ctx context.Context, input *CalculatorInput) (interface{}, error) {
    // Direct access to typed fields: input.Operation, input.A, input.B
    switch input.Operation {
    case "add":
        return map[string]interface{}{"result": input.A + input.B}, nil
    // ...
    }
}

// Register with automatic schema generation (method 1)
calculatorType := reflect.TypeOf(CalculatorInput{})
registry.RegisterStructToolWithHandler("calculator", "Performs arithmetic operations", calculatorType, calculatorHandler)

// Or use the convenience function (method 2)
resources.RegisterStructTool(registry, "calculator", "Performs arithmetic operations", calculatorHandler)

### Generics-Based Tool Registration (New in June 2025)

For even better type safety, you can use the new generics-based approach with two type parameters for input and output:

```go
type CalculatorInput struct {
    Operation string  `mcp:"operation,The operation to perform,required"`
    A         float64 `mcp:"a,First operand,required"`
    B         float64 `mcp:"b,Second operand,required"`
}

type CalculatorOutput struct {
    Result    float64 `mcp:"result,The calculation result,required"`
    Operation string  `mcp:"operation,The operation performed,required"`
}

func calculatorHandler(ctx context.Context, input *CalculatorInput) (*CalculatorOutput, error) {
    result := input.A + input.B
    return &CalculatorOutput{
        Result:    result,
        Operation: input.Operation,
    }, nil
}

// Register with compile-time type safety for both input and output
err := resources.RegisterStructToolWithTypes(registry, "calculator", "Performs arithmetic", calculatorHandler)

This approach provides:

• Compile-time type safety for both input and output
• Automatic schema generation for both input and output types
• Zero reflection at registration time (uses generics instead)
• Better IDE support with full type checking

#### 2. Using WithInputs

```go
weatherTool := resources.NewTool("weather").
    WithDescription("Get weather information for a location").
    WithInputs([]resources.ToolInput{
        {
            Name:        "location",
            Type:        "string",
            Description: "The location to get weather for",
            Required:    true,
        },
        {
            Name:        "units",
            Type:        "string",
            Description: "Temperature units (celsius or fahrenheit)",
            Default:     "celsius",
        },
    }).
    Build()

3. Using Individual Parameter Methods

calculatorTool := resources.NewTool("calculator").
    WithDescription("Performs basic arithmetic operations").
    WithString("operation").
    Required().
    Description("Operation to perform (add, subtract, multiply, divide)").
    Add().
    WithNumber("a").
    Required().
    Description("First operand").
    Add().
    WithNumber("b").
    Required().
    Description("Second operand").
    Add().
    Build()

Struct Tag Format

The mcp struct tag follows the format: mcp:"name,description,required,default=value"

• name: Field name in the schema (defaults to lowercase field name)
• description: Human-readable description of the field
• required: Mark field as required
• default=value: Set a default value for optional fields
• -: Skip field (not included in schema)

Supported Go Types:

• string → "string"
• int, int64, etc. → "integer"
• float64, float32 → "number"
• bool → "boolean"
• []T, [N]T → "array"
• struct, map → "object"
• *T → Same as T (pointer types)

Important Notes

CORS Configuration

When using an external HTTP server with the MCP transport, you need to configure CORS