Skip to main content

Model Context Protocol (MCP) Implementation

Overview

This package implements the Model Context Protocol (MCP) server specification in Go, providing a framework for real-time communication between AI models and clients using Server-Sent Events (SSE). The implementation follows the standardized protocol for building AI-assisted applications with bidirectional communication capabilities.

Core Components

Server-Sent Events (SSE) Communication

  • Real-time Communication: Robust SSE-based communication system that maintains persistent connections with clients
  • Connection Management: Client registration, message broadcasting, and client cleanup mechanisms
  • Event Handling: Event types for tools, prompts, and resources changes

JSON-RPC Implementation

  • Request Processing: Complete JSON-RPC request processor for handling MCP protocol methods
  • Response Formatting: Proper response formatting according to JSON-RPC specifications
  • Error Handling: Comprehensive error handling with appropriate error codes

Tool Management

  • Tool Registration: System to register custom tools with handlers
  • Tool Execution: Mechanism to execute tool functions with proper timeout handling
  • Result Handling: Flexible result handling supporting various return types (string, JSON, images)

Prompt System

  • Prompt Registration: System for registering both static and dynamic prompts
  • Argument Validation: Validation for required arguments and default values for optional ones
  • Message Generation: Handlers that generate properly formatted conversation messages

Resource Management

  • Resource Registration: System for managing and accessing external resources
  • Content Delivery: Handlers for delivering resource content to clients on demand
  • Resource Subscription: Mechanisms for clients to subscribe to resource updates

Protocol Features

  • Initialization Sequence: Proper handshaking with capability negotiation
  • Notification Handling: Support for both standard and client-specific notifications
  • Message Routing: Intelligent routing of requests to appropriate handlers

Technical Highlights

Configuration System

  • Flexible Configuration: Configuration system with sensible defaults and customization options
  • CORS Support: Configurable CORS settings for cross-origin requests
  • Server Information: Proper server identification and versioning

Client Session Management

  • Session Tracking: Client session tracking with unique identifiers
  • Connection Health: Ping/pong mechanism to maintain connection health
  • Initialization State: Client initialization state tracking

Content Handling

  • Multi-format Content: Support for text, code, and binary content
  • MIME Type Support: Proper MIME type identification for various content types
  • Audience Annotations: Content audience annotations for user/assistant targeting

Usage

Setting Up an MCP Server

To create and start an MCP server:

package main

import (
    "github.com/zeromicro/go-zero/core/conf"
    "github.com/zeromicro/go-zero/core/logx"
    "github.com/zeromicro/go-zero/mcp"
)

func main() {
    // Load configuration from YAML file
    var c mcp.McpConf
    conf.MustLoad("config.yaml", &c)

    // Optional: Disable stats logging
    logx.DisableStat()

    // Create MCP server
    server := mcp.NewMcpServer(c)

    // Register tools, prompts, and resources (examples below)

    // Start the server and ensure it's stopped on exit
    defer server.Stop()
    server.Start()
}

Sample configuration file (config.yaml):

name: mcp-server
host: localhost
port: 8080
mcp:
  name: my-mcp-server
  messageTimeout: 30s # Timeout for tool calls
  cors:
    - http://localhost:3000 # Optional CORS configuration

Registering Tools

Tools allow AI models to execute custom code through the MCP protocol.

Basic Tool Example:

// Register a simple echo tool
echoTool := mcp.Tool{
    Name:        "echo",
    Description: "Echoes back the message provided by the user",
    InputSchema: mcp.InputSchema{
        Properties: map[string]any{
            "message": map[string]any{
                "type":        "string",
                "description": "The message to echo back",
            },
            "prefix": map[string]any{
                "type":        "string",
                "description": "Optional prefix to add to the echoed message",
                "default":     "Echo: ",
            },
        },
        Required: []string{"message"},
    },
    Handler: func(ctx context.Context, params map[string]any) (any, error) {
        var req struct {
            Message string `json:"message"`
            Prefix  string `json:"prefix,optional"`
        }

        if err := mcp.ParseArguments(params, &req); err != nil {
            return nil, fmt.Errorf("failed to parse params: %w", err)
        }

        prefix := "Echo: "
        if len(req.Prefix) > 0 {
            prefix = req.Prefix
        }

        return prefix + req.Message, nil
    },
}

server.RegisterTool(echoTool)

Tool with Different Response Types:

// Tool returning JSON data
dataTool := mcp.Tool{
    Name:        "data.generate",
    Description: "Generates sample data in various formats",
    InputSchema: mcp.InputSchema{
        Properties: map[string]any{
            "format": map[string]any{
                "type":        "string",
                "description": "Format of data (json, text)",
                "enum":        []string{"json", "text"},
            },
        },
    },
    Handler: func(ctx context.Context, params map[string]any) (any, error) {
        var req struct {
            Format string `json:"format"`
        }

        if err := mcp.ParseArguments(params, &req); err != nil {
            return nil, fmt.Errorf("failed to parse params: %w", err)
        }

        if req.Format == "json" {
            // Return structured data
            return map[string]any{
                "items": []map[string]any{
                    {"id": 1, "name": "Item 1"},
                    {"id": 2, "name": "Item 2"},
                },
                "count": 2,
            }, nil
        }

        // Default to text
        return "Sample text data", nil
    },
}

server.RegisterTool(dataTool)

Image Generation Tool Example:

// Tool returning image content
imageTool := mcp.Tool{
    Name:        "image.generate",
    Description: "Generates a simple image",
    InputSchema: mcp.InputSchema{
        Properties: map[string]any{
            "type": map[string]any{
                "type":        "string",
                "description": "Type of image to generate",
                "default":     "placeholder",
            },
        },
    },
    Handler: func(ctx context.Context, params map[string]any) (any, error) {
        // Return image content directly
        return mcp.ImageContent{
            Data:     "base64EncodedImageData...", // Base64 encoded image data
            MimeType: "image/png",
        }, nil
    },
}

server.RegisterTool(imageTool)

Using ToolResult for Custom Outputs:

// Tool that returns a custom ToolResult type
customResultTool := mcp.Tool{
    Name:        "custom.result",
    Description: "Returns a custom formatted result",
    InputSchema: mcp.InputSchema{
        Properties: map[string]any{
            "resultType": map[string]any{
                "type": "string",
                "enum": []string{"text", "image"},
            },
        },
    },
    Handler: func(ctx context.Context, params map[string]any) (any, error) {
        var req struct {
            ResultType string `json:"resultType"`
        }

        if err := mcp.ParseArguments(params, &req); err != nil {
            return nil, fmt.Errorf("failed to parse params: %w", err)
        }

        if req.ResultType == "image" {
            return mcp.ToolResult{
                Type: mcp.ContentTypeImage,
                Content: map[string]any{
                    "data":     "base64EncodedImageData...",
                    "mimeType": "image/jpeg",
                },
            }, nil
        }

        // Default to text
        return mcp.ToolResult{
            Type:    mcp.ContentTypeText,
            Content: "This is a text result from ToolResult",
        }, nil
    },
}

server.RegisterTool(customResultTool)

Registering Prompts

Prompts are reusable conversation templates for AI models.

Static Prompt Example:

// Register a simple static prompt with placeholders
server.RegisterPrompt(mcp.Prompt{
    Name:        "hello",
    Description: "A simple hello prompt",
    Arguments: []mcp.PromptArgument{
        {
            Name:        "name",
            Description: "The name to greet",
            Required:    false,
        },
    },
    Content: "Say hello to {{name}} and introduce yourself as an AI assistant.",
})

Dynamic Prompt with Handler Function:

// Register a prompt with a dynamic handler function
server.RegisterPrompt(mcp.Prompt{
    Name:        "dynamic-prompt",
    Description: "A prompt that uses a handler to generate dynamic content",
    Arguments: []mcp.PromptArgument{
        {
            Name:        "username",
            Description: "User's name for personalized greeting",
            Required:    true,
        },
        {
            Name:        "topic",
            Description: "Topic of expertise",
            Required:    true,
        },
    },
    Handler: func(ctx context.Context, args map[string]string) ([]mcp.PromptMessage, error) {
        var req struct {
            Username string `json:"username"`
            Topic    string `json:"topic"`
        }

        if err := mcp.ParseArguments(args, &req); err != nil {
            return nil, fmt.Errorf("failed to parse args: %w", err)
        }

        // Create a user message
        userMessage := mcp.PromptMessage{
            Role: mcp.RoleUser,
            Content: mcp.TextContent{
                Text: fmt.Sprintf("Hello, I'm %s and I'd like to learn about %s.", req.Username, req.Topic),
            },
        }

        // Create an assistant response with current time
        currentTime := time.Now().Format(time.RFC1123)
        assistantMessage := mcp.PromptMessage{
            Role: mcp.RoleAssistant,
            Content: mcp.TextContent{
                Text: fmt.Sprintf("Hello %s! I'm an AI assistant and I'll help you learn about %s. The current time is %s.",
                    req.Username, req.Topic, currentTime),
            },
        }

        // Return both messages as a conversation
        return []mcp.PromptMessage{userMessage, assistantMessage}, nil
    },
})

Multi-Message Prompt with Code Examples:

// Register a prompt that provides code examples in different programming languages
server.RegisterPrompt(mcp.Prompt{
    Name:        "code-example",
    Description: "Provides code examples in different programming languages",
    Arguments: []mcp.PromptArgument{
        {
            Name:        "language",
            Description: "Programming language for the example",
            Required:    true,
        },
        {
            Name:        "complexity",
            Description: "Complexity level (simple, medium, advanced)",
        },
    },
    Handler: func(ctx context.Context, args map[string]string) ([]mcp.PromptMessage, error) {
        var req struct {
            Language   string `json:"language"`
            Complexity string `json:"complexity,optional"`
        }

        if err := mcp.ParseArguments(args, &req); err != nil {
            return nil, fmt.Errorf("failed to parse args: %w", err)
        }

        // Validate language
        supportedLanguages := map[string]bool{"go": true, "python": true, "javascript": true, "rust": true}
        if !supportedLanguages[req.Language] {
            return nil, fmt.Errorf("unsupported language: %s", req.Language)
        }

        // Generate code example based on language and complexity
        var codeExample string

        switch req.Language {
        case "go":
            if req.Complexity == "simple" {
                codeExample = `
package main

import "fmt"

func main() {
    fmt.Println("Hello, World!")
}`
            } else {
                codeExample = `
package main

import (
    "fmt"
    "time"
)

func main() {
    now := time.Now()
    fmt.Printf("Hello, World! Current time is %s\n", now.Format(time.RFC3339))
}`
            }
        case "python":
            // Python example code
            if req.Complexity == "simple" {
                codeExample = `
def greet(name):
    return f"Hello, {name}!"

print(greet("World"))`
            } else {
                codeExample = `
import datetime

def greet(name, include_time=False):
    message = f"Hello, {name}!"
    if include_time:
        message += f" Current time is {datetime.datetime.now().isoformat()}"
    return message

print(greet("World", include_time=True))`
            }
        }

        // Create messages array according to MCP spec
        messages := []mcp.PromptMessage{
            {
                Role: mcp.RoleAssistant,
                Content: mcp.TextContent{
                    Text: fmt.Sprintf("You are a helpful coding assistant specialized in %s programming.", req.Language),
                },
            },
            {
                Role: mcp.RoleUser,
                Content: mcp.TextContent{
                    Text: fmt.Sprintf("Show me a %s example of a Hello World program in %s.", req.Complexity, req.Language),
                },
            },
            {
                Role: mcp.RoleAssistant,
                Content: mcp.TextContent{
                    Text: fmt.Sprintf("Here's a %s example in %s:\n\n```%s%s\n```\n\nHow can I help you implement this?",
                        req.Complexity, req.Language, req.Language, codeExample),
                },
            },
        }

        return messages, nil
    },
})

Registering Resources

Resources provide access to external content such as files or generated data.

Basic Resource Example:

// Register a static resource
server.RegisterResource(mcp.Resource{
    Name:        "example-document",
    URI:         "file:///example/document.txt",
    Description: "An example document",
    MimeType:    "text/plain",
    Handler: func(ctx context.Context) (mcp.ResourceContent, error) {
        return mcp.ResourceContent{
            URI:      "file:///example/document.txt",
            MimeType: "text/plain",
            Text:     "This is an example document content.",
        }, nil
    },
})

Dynamic Resource with Code Example:

// Register a Go code resource with dynamic handler
server.RegisterResource(mcp.Resource{
    Name:        "go-example",
    URI:         "file:///project/src/main.go",
    Description: "A simple Go example with multiple files",
    MimeType:    "text/x-go",
    Handler: func(ctx context.Context) (mcp.ResourceContent, error) {
        // Return ResourceContent with all required fields
        return mcp.ResourceContent{
            URI:      "file:///project/src/main.go",
            MimeType: "text/x-go",
            Text:     "package main\n\nimport (\n\t\"fmt\"\n\t\"./greeting\"\n)\n\nfunc main() {\n\tfmt.Println(greeting.Hello(\"world\"))\n}",
        }, nil
    },
})

// Register a companion file for the above example
server.RegisterResource(mcp.Resource{
    Name:        "go-greeting",
    URI:         "file:///project/src/greeting/greeting.go",
    Description: "A greeting package for the Go example",
    MimeType:    "text/x-go",
    Handler: func(ctx context.Context) (mcp.ResourceContent, error) {
        return mcp.ResourceContent{
            URI:      "file:///project/src/greeting/greeting.go",
            MimeType: "text/x-go",
            Text:     "package greeting\n\nfunc Hello(name string) string {\n\treturn \"Hello, \" + name + \"!\"\n}",
        }, nil
    },
})

Binary Resource Example:

// Register a binary resource (like an image)
server.RegisterResource(mcp.Resource{
    Name:        "example-image",
    URI:         "file:///example/image.png",
    Description: "An example image",
    MimeType:    "image/png",
    Handler: func(ctx context.Context) (mcp.ResourceContent, error) {
        // Read image from file or generate it
        imageData := "base64EncodedImageData..." // Base64 encoded image data

        return mcp.ResourceContent{
            URI:      "file:///example/image.png",
            MimeType: "image/png",
            Blob:     imageData, // For binary data
        }, nil
    },
})

Using Resources in Prompts

You can embed resources in prompt responses to create rich interactions with proper MCP-compliant structure:

// Register a prompt that embeds a resource
server.RegisterPrompt(mcp.Prompt{
    Name:        "resource-example",
    Description: "A prompt that embeds a resource",
    Arguments: []mcp.PromptArgument{
        {
            Name:        "file_type",
            Description: "Type of file to show (rust or go)",
            Required:    true,
        },
    },
    Handler: func(ctx context.Context, args map[string]string) ([]mcp.PromptMessage, error) {
        var req struct {
            FileType string `json:"file_type"`
        }

        if err := mcp.ParseArguments(args, &req); err != nil {
            return nil, fmt.Errorf("failed to parse args: %w", err)
        }

        var resourceURI, mimeType, fileContent string
        if req.FileType == "rust" {
            resourceURI = "file:///project/src/main.rs"
            mimeType = "text/x-rust"
            fileContent = "fn main() {\n    println!(\"Hello world!\");\n}"
        } else {
            resourceURI = "file:///project/src/main.go"
            mimeType = "text/x-go"
            fileContent = "package main\n\nimport \"fmt\"\n\nfunc main() {\n\tfmt.Println(\"Hello, world!\")\n}"
        }

        // Create message with embedded resource using proper MCP format
        return []mcp.PromptMessage{
            {
                Role: mcp.RoleUser,
                Content: mcp.TextContent{
                    Text: fmt.Sprintf("Can you explain this %s code?", req.FileType),
                },
            },
            {
                Role: mcp.RoleAssistant,
                Content: mcp.EmbeddedResource{
                    Type: mcp.ContentTypeResource,
                    Resource: struct {
                        URI      string `json:"uri"`
                        MimeType string `json:"mimeType"`
                        Text     string `json:"text,omitempty"`
                        Blob     string `json:"blob,omitempty"`
                    }{
                        URI:      resourceURI,
                        MimeType: mimeType,
                        Text:     fileContent,
                    },
                },
            },
            {
                Role: mcp.RoleAssistant,
                Content: mcp.TextContent{
                    Text: fmt.Sprintf("Above is a simple Hello World example in %s. Let me explain how it works.", req.FileType),
                },
            },
        }, nil
    },
})

Multiple File Resources Example

// Register a prompt that demonstrates embedding multiple resource files
server.RegisterPrompt(mcp.Prompt{
    Name:        "go-code-example",
    Description: "A prompt that correctly embeds multiple resource files",
    Arguments: []mcp.PromptArgument{
        {
            Name:        "format",
            Description: "How to format the code display",
        },
    },
    Handler: func(ctx context.Context, args map[string]string) ([]mcp.PromptMessage, error) {
        var req struct {
            Format string `json:"format,optional"`
        }

        if err := mcp.ParseArguments(args, &req); err != nil {
            return nil, fmt.Errorf("failed to parse args: %w", err)
        }

        // Get the Go code for multiple files
        var mainGoText string = "package main\n\nimport (\n\t\"fmt\"\n\t\"./greeting\"\n)\n\nfunc main() {\n\tfmt.Println(greeting.Hello(\"world\"))\n}"
        var greetingGoText string = "package greeting\n\nfunc Hello(name string) string {\n\treturn \"Hello, \" + name + \"!\"\n}"

        // Create message with properly formatted embedded resource per MCP spec
        messages := []mcp.PromptMessage{
            {
                Role: mcp.RoleUser,
                Content: mcp.TextContent{
                    Text: "Show me a simple Go example with proper imports.",
                },
            },
            {
                Role: mcp.RoleAssistant,
                Content: mcp.TextContent{
                    Text: "Here's a simple Go example project:",
                },
            },
            {
                Role: mcp.RoleAssistant,
                Content: mcp.EmbeddedResource{
                    Type: mcp.ContentTypeResource,
                    Resource: struct {
                        URI      string `json:"uri"`
                        MimeType string `json:"mimeType"`
                        Text     string `json:"text,omitempty"`
                        Blob     string `json:"blob,omitempty"`
                    }{
                        URI:      "file:///project/src/main.go",
                        MimeType: "text/x-go",
                        Text:     mainGoText,
                    },
                },
            },
        }

        // Add explanation and additional file if requested
        if req.Format == "with_explanation" {
            messages = append(messages, mcp.PromptMessage{
                Role: mcp.RoleAssistant,
                Content: mcp.TextContent{
                    Text: "This example demonstrates a simple Go application with modular structure. The main.go file imports from a local 'greeting' package that provides the Hello function.",
                },
            })

            // Also show the greeting.go file with correct resource format
            messages = append(messages, mcp.PromptMessage{
                Role: mcp.RoleAssistant,
                Content: mcp.EmbeddedResource{
                    Type: mcp.ContentTypeResource,
                    Resource: struct {
                        URI      string `json:"uri"`
                        MimeType string `json:"mimeType"`
                        Text     string `json:"text,omitempty"`
                        Blob     string `json:"blob,omitempty"`
                    }{
                        URI:      "file:///project/src/greeting/greeting.go",
                        MimeType: "text/x-go",
                        Text:     greetingGoText,
                    },
                },
            })
        }

        return messages, nil
    },
})

Complete Application Example

Here's a complete example demonstrating all the components:

package main

import (
    "context"
    "fmt"
    "log"
    "time"

    "github.com/zeromicro/go-zero/core/conf"
    "github.com/zeromicro/go-zero/core/logx"
    "github.com/zeromicro/go-zero/mcp"
)

func main() {
    // Load configuration
    var c mcp.McpConf
    if err := conf.Load("config.yaml", &c); err != nil {
        log.Fatalf("Failed to load config: %v", err)
    }

    // Set up logging
    logx.DisableStat()

    // Create MCP server
    server := mcp.NewMcpServer(c)
    defer server.Stop()

    // Register a simple echo tool
    echoTool := mcp.Tool{
        Name:        "echo",
        Description: "Echoes back the message provided by the user",
        InputSchema: mcp.InputSchema{
            Properties: map[string]any{
                "message": map[string]any{
                    "type":        "string",
                    "description": "The message to echo back",
                },
                "prefix": map[string]any{
                    "type":        "string",
                    "description": "Optional prefix to add to the echoed message",
                    "default":     "Echo: ",
                },
            },
            Required: []string{"message"},
        },
        Handler: func(ctx context.Context, params map[string]any) (any, error) {
            var req struct {
                Message string `json:"message"`
                Prefix  string `json:"prefix,optional"`
            }

            if err := mcp.ParseArguments(params, &req); err != nil {
                return nil, fmt.Errorf("failed to parse args: %w", err)
            }

            prefix := "Echo: "
            if len(req.Prefix) > 0 {
                prefix = req.Prefix
            }

            return prefix + req.Message, nil
        },
    }
    server.RegisterTool(echoTool)

    // Register a static prompt
    server.RegisterPrompt(mcp.Prompt{
        Name:        "greeting",
        Description: "A simple greeting prompt",
        Arguments: []mcp.PromptArgument{
            {
                Name:        "name",
                Description: "The name to greet",
                Required:    true,
            },
        },
        Content: "Hello {{name}}! How can I assist you today?",
    })

    // Register a dynamic prompt
    server.RegisterPrompt(mcp.Prompt{
        Name:        "dynamic-prompt",
        Description: "A prompt that uses a handler to generate dynamic content",
        Arguments: []mcp.PromptArgument{
            {
                Name:        "username",
                Description: "User's name for personalized greeting",
                Required:    true,
            },
            {
                Name:        "topic",
                Description: "Topic of expertise",
                Required:    true,
            },
        },
        Handler: func(ctx context.Context, args map[string]string) ([]mcp.PromptMessage, error) {
            var req struct {
                Username string `json:"username"`
                Topic    string `json:"topic"`
            }

            if err := mcp.ParseArguments(args, &req); err != nil {
                return nil, fmt.Errorf("failed to parse args: %w", err)
            }

            // Create messages with current time
            currentTime := time.Now().Format(time.RFC1123)
            return []mcp.PromptMessage{
                {
                    Role: mcp.RoleUser,
                    Content: mcp.TextContent{
                        Text: fmt.Sprintf("Hello, I'm %s and I'd like to learn about %s.", req.Username, req.Topic),
                    },
                },
                {
                    Role: mcp.RoleAssistant,
                    Content: mcp.TextContent{
                        Text: fmt.Sprintf("Hello %s! I'm an AI assistant and I'll help you learn about %s. The current time is %s.",
                            req.Username, req.Topic, currentTime),
                    },
                },
            }, nil
        },
    })

    // Register a resource
    server.RegisterResource(mcp.Resource{
        Name:        "example-doc",
        URI:         "file:///example/doc.txt",
        Description: "An example document",
        MimeType:    "text/plain",
        Handler: func(ctx context.Context) (mcp.ResourceContent, error) {
            return mcp.ResourceContent{
                URI:      "file:///example/doc.txt",
                MimeType: "text/plain",
                Text:     "This is the content of the example document.",
            }, nil
        },
    })

    // Start the server
    fmt.Printf("Starting MCP server on %s:%d\n", c.Host, c.Port)
    server.Start()
}

Error Handling

The MCP implementation provides comprehensive error handling:

  • Tool execution errors are properly reported back to clients
  • Missing or invalid parameters are detected and reported with appropriate error codes
  • Resource and prompt lookup failures are handled gracefully
  • Timeout handling for long-running tool executions using context
  • Panic recovery to prevent server crashes

Advanced Features

  • Annotations: Add audience and priority metadata to content
  • Content Types: Support for text, images, audio, and other content formats
  • Embedded Resources: Include file resources directly in prompt responses
  • Context Awareness: All handlers receive context.Context for timeout and cancellation support
  • Progress Tokens: Support for tracking progress of long-running operations
  • Customizable Timeouts: Configure execution timeouts for tools and operations

Performance Considerations

  • Tool execution runs with configurable timeouts to prevent blocking
  • Efficient client tracking and cleanup to prevent resource leaks
  • Proper concurrency handling with mutex protection for shared resources
  • Buffered message channels to prevent blocking on client message delivery

Contents