Add example of customising context

Author: sd2k

examples/custom_context/main.go

@@ -0,0 +1,164 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"flag"
+	"fmt"
+	"io"
+	"log"
+	"net/http"
+	"os"
+
+	"github.com/mark3labs/mcp-go/mcp"
+	"github.com/mark3labs/mcp-go/server"
+)
+
+// authKey is a custom context key for storing the auth token.
+type authKey struct{}
+
+// withAuthKey adds an auth key to the context.
+func withAuthKey(ctx context.Context, auth string) context.Context {
+	return context.WithValue(ctx, authKey{}, auth)
+}
+
+// authFromRequest extracts the auth token from the request headers.
+func authFromRequest(ctx context.Context, r *http.Request) context.Context {
+	return withAuthKey(ctx, r.Header.Get("Authorization"))
+}
+
+// authFromEnv extracts the auth token from the environment
+func authFromEnv(ctx context.Context) context.Context {
+	return withAuthKey(ctx, os.Getenv("API_KEY"))
+}
+
+// tokenFromContext extracts the auth token from the context.
+// This can be used by tools to extract the token regardless of the
+// transport being used by the server.
+func tokenFromContext(ctx context.Context) (string, error) {
+	auth, ok := ctx.Value(authKey{}).(string)
+	if !ok {
+		return "", fmt.Errorf("missing auth")
+	}
+	return auth, nil
+}
+
+type response struct {
+	Args    map[string]interface{} `json:"args"`
+	Headers map[string]string      `json:"headers"`
+}
+
+// makeRequest makes a request to httpbin.org including the auth token in the request
+// headers and the message in the query string.
+func makeRequest(ctx context.Context, message, token string) (*response, error) {
+	req, err := http.NewRequestWithContext(ctx, "GET", "https://httpbin.org/anything", nil)
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Set("Authorization", token)
+	req.URL.Query().Add("message", message)
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, err
+	}
+	var r *response
+	if err := json.Unmarshal(body, r); err != nil {
+		return nil, err
+	}
+	return r, nil
+}
+
+// handleMakeAuthenticatedRequestTool is a tool that makes an authenticated request
+// using the token from the context.
+func handleMakeAuthenticatedRequestTool(
+	ctx context.Context,
+	request mcp.CallToolRequest,
+) (*mcp.CallToolResult, error) {
+	message, ok := request.Params.Arguments["message"].(string)
+	if !ok {
+		return nil, fmt.Errorf("missing message")
+	}
+	token, err := tokenFromContext(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("missing token: %v", err)
+	}
+	// Now our tool can make a request with the token, irrespective of where it came from.
+	resp, err := makeRequest(ctx, message, token)
+	if err != nil {
+		return nil, err
+	}
+	return mcp.NewToolResultText(fmt.Sprintf("%+v", resp)), nil
+}
+
+type MCPServer struct {
+	server *server.MCPServer
+}
+
+func NewMCPServer() *MCPServer {
+	mcpServer := server.NewMCPServer(
+		"example-server",
+		"1.0.0",
+		server.WithResourceCapabilities(true, true),
+		server.WithPromptCapabilities(true),
+		server.WithToolCapabilities(true),
+	)
+	mcpServer.AddTool(mcp.NewTool("make_authenticated_request",
+		mcp.WithDescription("Makes an authenticated request"),
+		mcp.WithString("message",
+			mcp.Description("Message to echo"),
+			mcp.Required(),
+		),
+	), handleMakeAuthenticatedRequestTool)
+
+	return &MCPServer{
+		server: mcpServer,
+	}
+}
+
+func (s *MCPServer) ServeSSE(addr string) *server.SSEServer {
+	return server.NewSSEServer(s.server,
+		server.WithBaseURL(fmt.Sprintf("http://%s", addr)),
+		server.WithSSEContextFunc(authFromRequest),
+	)
+}
+
+func (s *MCPServer) ServeStdio() error {
+	return server.ServeStdio(s.server, server.WithStdioContextFunc(authFromEnv))
+}
+
+func main() {
+	var transport string
+	flag.StringVar(&transport, "t", "stdio", "Transport type (stdio or sse)")
+	flag.StringVar(
+		&transport,
+		"transport",
+		"stdio",
+		"Transport type (stdio or sse)",
+	)
+	flag.Parse()
+
+	s := NewMCPServer()
+
+	switch transport {
+	case "stdio":
+		if err := s.ServeStdio(); err != nil {
+			log.Fatalf("Server error: %v", err)
+		}
+	case "sse":
+		sseServer := s.ServeSSE("localhost:8080")
+		log.Printf("SSE server listening on :8080")
+		if err := sseServer.Start(":8080"); err != nil {
+			log.Fatalf("Server error: %v", err)
+		}
+	default:
+		log.Fatalf(
+			"Invalid transport type: %s. Must be 'stdio' or 'sse'",
+			transport,
+		)
+	}
+}