diff --git a/mcp/example_middleware_test.go b/mcp/example_middleware_test.go new file mode 100644 index 0000000..1328473 --- /dev/null +++ b/mcp/example_middleware_test.go @@ -0,0 +1,141 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package mcp_test + +import ( + "context" + "fmt" + "log/slog" + "os" + "time" + + "github.com/modelcontextprotocol/go-sdk/jsonschema" + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +// This example demonstrates server side logging using the mcp.Middleware system. +func Example_loggingMiddleware() { + // Create a logger for demonstration purposes. + logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{ + Level: slog.LevelInfo, + ReplaceAttr: func(_ []string, a slog.Attr) slog.Attr { + // Simplify timestamp format for consistent output. + if a.Key == slog.TimeKey { + return slog.String("time", "2025-01-01T00:00:00Z") + } + return a + }, + })) + + loggingMiddleware := func(next mcp.MethodHandler[*mcp.ServerSession]) mcp.MethodHandler[*mcp.ServerSession] { + return func( + ctx context.Context, + session *mcp.ServerSession, + method string, + params mcp.Params, + ) (mcp.Result, error) { + logger.Info("MCP method started", + "method", method, + "session_id", session.ID(), + "has_params", params != nil, + ) + + start := time.Now() + + result, err := next(ctx, session, method, params) + + duration := time.Since(start) + + if err != nil { + logger.Error("MCP method failed", + "method", method, + "session_id", session.ID(), + "duration_ms", duration.Milliseconds(), + "err", err, + ) + } else { + logger.Info("MCP method completed", + "method", method, + "session_id", session.ID(), + "duration_ms", duration.Milliseconds(), + "has_result", result != nil, + ) + } + + return result, err + } + } + + // Create server with middleware + server := mcp.NewServer(&mcp.Implementation{Name: "logging-example"}, nil) + server.AddReceivingMiddleware(loggingMiddleware) + + // Add a simple tool + server.AddTool( + &mcp.Tool{ + Name: "greet", + Description: "Greet someone with logging.", + InputSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "name": { + Type: "string", + Description: "Name to greet", + }, + }, + Required: []string{"name"}, + }, + }, + func( + ctx context.Context, + ss *mcp.ServerSession, + params *mcp.CallToolParamsFor[map[string]any], + ) (*mcp.CallToolResultFor[any], error) { + name, ok := params.Arguments["name"].(string) + if !ok { + return nil, fmt.Errorf("name parameter is required and must be a string") + } + + message := fmt.Sprintf("Hello, %s!", name) + return &mcp.CallToolResultFor[any]{ + Content: []mcp.Content{ + &mcp.TextContent{Text: message}, + }, + }, nil + }, + ) + + // Create client-server connection for demonstration + client := mcp.NewClient(&mcp.Implementation{Name: "test-client"}, nil) + clientTransport, serverTransport := mcp.NewInMemoryTransports() + + ctx := context.Background() + + // Connect server and client + serverSession, _ := server.Connect(ctx, serverTransport) + defer serverSession.Close() + + clientSession, _ := client.Connect(ctx, clientTransport) + defer clientSession.Close() + + // Call the tool to demonstrate logging + result, _ := clientSession.CallTool(ctx, &mcp.CallToolParams{ + Name: "greet", + Arguments: map[string]any{ + "name": "World", + }, + }) + + fmt.Printf("Tool result: %s\n", result.Content[0].(*mcp.TextContent).Text) + + // Output: + // time=2025-01-01T00:00:00Z level=INFO msg="MCP method started" method=initialize session_id="" has_params=true + // time=2025-01-01T00:00:00Z level=INFO msg="MCP method completed" method=initialize session_id="" duration_ms=0 has_result=true + // time=2025-01-01T00:00:00Z level=INFO msg="MCP method started" method=notifications/initialized session_id="" has_params=true + // time=2025-01-01T00:00:00Z level=INFO msg="MCP method completed" method=notifications/initialized session_id="" duration_ms=0 has_result=false + // time=2025-01-01T00:00:00Z level=INFO msg="MCP method started" method=tools/call session_id="" has_params=true + // time=2025-01-01T00:00:00Z level=INFO msg="MCP method completed" method=tools/call session_id="" duration_ms=0 has_result=true + // Tool result: Hello, World! +}