Commit 207a777

mo khan <mo@mokhan.ca>
2025-06-23 13:59:04
fix: implement lazy loading for filesystem resources to prevent memory issues
Replace eager file discovery at startup with on-demand resource listing. This fixes a fatal design flaw where the filesystem MCP would load all file metadata into memory during initialization, causing performance issues and potential crashes with large directory trees. Key changes: - Remove recursive file discovery from registerResources() - Add dynamic ListResources() method that discovers files on request - Enhance base MCP server with pattern-based resource handlers - Add custom request handler system for proper method resolution 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
1 parent 611a48e
Changed files (2)
pkg
pkg/filesystem/server.go
@@ -1,6 +1,7 @@
 package filesystem
 
 import (
+	"encoding/json"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -42,9 +43,42 @@ func New(allowedDirs []string) *Server {
 	fsServer.registerResources()
 	fsServer.registerRoots()
 
+	// Override the base server's resource listing handler
+	fsServer.setupResourceHandling()
+
 	return fsServer
 }
 
+// setupResourceHandling configures custom resource handling to ensure proper method resolution
+func (fs *Server) setupResourceHandling() {
+	// We need to create a wrapper function that calls our ListResources method
+	originalHandlers := make(map[string]func(mcp.JSONRPCRequest) mcp.JSONRPCResponse)
+	
+	// Store reference to filesystem server for closure
+	fsServer := fs
+	
+	// Create custom handler that calls our ListResources method
+	customListResourcesHandler := func(req mcp.JSONRPCRequest) mcp.JSONRPCResponse {
+		resources := fsServer.ListResources()
+		result := mcp.ListResourcesResult{Resources: resources}
+		id := req.ID
+		var resultBytes *json.RawMessage
+		bytes, _ := json.Marshal(result)
+		rawMsg := json.RawMessage(bytes)
+		resultBytes = &rawMsg
+		return mcp.JSONRPCResponse{
+			JSONRPC: "2.0",
+			ID:      id,
+			Result:  resultBytes,
+		}
+	}
+	
+	originalHandlers["resources/list"] = customListResourcesHandler
+	
+	// Override the base server's request handling
+	fs.SetCustomRequestHandler(originalHandlers)
+}
+
 // registerTools registers all Filesystem tools with the server
 func (fs *Server) registerTools() {
 	fs.RegisterTool("read_file", fs.HandleReadFile)
@@ -88,13 +122,12 @@ func (fs *Server) registerPrompts() {
 	fs.RegisterPrompt(editFilePrompt, fs.HandleEditFilePrompt)
 }
 
-// registerResources registers filesystem resources (file:// URIs)
+// registerResources sets up resource handling (lazy loading) 
 func (fs *Server) registerResources() {
-	// Register resources for all files in allowed directories
-	// This provides discovery of available files as resources
-	for _, dir := range fs.allowedDirectories {
-		fs.discoverFilesInDirectory(dir)
-	}
+	// Register a generic file:// resource handler for pattern matching
+	// This avoids loading all files into memory at startup
+	// We register it but it won't appear in ListResources since we override that method
+	fs.Server.RegisterResource("file://", fs.HandleFileResource)
 }
 
 // registerRoots registers filesystem allowed directories as roots
@@ -115,8 +148,20 @@ func (fs *Server) registerRoots() {
 	}
 }
 
-// discoverFilesInDirectory recursively discovers files and registers them as resources
-func (fs *Server) discoverFilesInDirectory(dirPath string) {
+// ListResources dynamically discovers and returns file resources from allowed directories
+func (fs *Server) ListResources() []mcp.Resource {
+	var resources []mcp.Resource
+
+	// Dynamically discover files in allowed directories
+	for _, dir := range fs.allowedDirectories {
+		fs.discoverFilesInDirectory(dir, &resources)
+	}
+
+	return resources
+}
+
+// discoverFilesInDirectory recursively discovers files and adds them to the resources slice
+func (fs *Server) discoverFilesInDirectory(dirPath string, resources *[]mcp.Resource) {
 	filepath.Walk(dirPath, func(path string, info os.FileInfo, err error) error {
 		if err != nil {
 			return nil // Skip files with errors
@@ -160,8 +205,8 @@ func (fs *Server) discoverFilesInDirectory(dirPath string) {
 			MimeType:    mimeType,
 		}
 
-		// Register resource with handler
-		fs.RegisterResourceWithDefinition(resource, fs.HandleFileResource)
+		// Add to resources slice instead of registering
+		*resources = append(*resources, resource)
 
 		return nil
 	})
pkg/mcp/server.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"log"
 	"os"
+	"strings"
 	"sync"
 )
 
@@ -28,6 +29,9 @@ type Server struct {
 	initializeHandler func(InitializeRequest) (InitializeResult, error)
 	shutdownHandler   func() error
 
+	// Custom request handlers for overriding default behavior
+	customRequestHandlers map[string]func(JSONRPCRequest) JSONRPCResponse
+
 	mu sync.RWMutex
 }
 
@@ -39,14 +43,15 @@ type ResourceHandler func(ReadResourceRequest) (ReadResourceResult, error)
 // NewServer creates a new MCP server
 func NewServer(name, version string) *Server {
 	return &Server{
-		name:                name,
-		version:             version,
-		toolHandlers:        make(map[string]ToolHandler),
-		promptHandlers:      make(map[string]PromptHandler),
-		promptDefinitions:   make(map[string]Prompt),
-		resourceHandlers:    make(map[string]ResourceHandler),
-		resourceDefinitions: make(map[string]Resource),
-		rootDefinitions:     make(map[string]Root),
+		name:                  name,
+		version:               version,
+		toolHandlers:          make(map[string]ToolHandler),
+		promptHandlers:        make(map[string]PromptHandler),
+		promptDefinitions:     make(map[string]Prompt),
+		resourceHandlers:      make(map[string]ResourceHandler),
+		resourceDefinitions:   make(map[string]Resource),
+		rootDefinitions:       make(map[string]Root),
+		customRequestHandlers: make(map[string]func(JSONRPCRequest) JSONRPCResponse),
 		capabilities: ServerCapabilities{
 			Tools:     &ToolsCapability{},
 			Prompts:   &PromptsCapability{},
@@ -106,6 +111,15 @@ func (s *Server) SetShutdownHandler(handler func() error) {
 	s.shutdownHandler = handler
 }
 
+// SetCustomRequestHandler sets custom request handlers for overriding default behavior
+func (s *Server) SetCustomRequestHandler(handlers map[string]func(JSONRPCRequest) JSONRPCResponse) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	for method, handler := range handlers {
+		s.customRequestHandlers[method] = handler
+	}
+}
+
 // ListTools returns all registered tools
 func (s *Server) ListTools() []Tool {
 	s.mu.RLock()
@@ -201,6 +215,15 @@ func (s *Server) Run(ctx context.Context) error {
 
 // handleRequest processes a JSON-RPC request
 func (s *Server) handleRequest(req JSONRPCRequest) JSONRPCResponse {
+	// Check for custom handlers first
+	s.mu.RLock()
+	if customHandler, exists := s.customRequestHandlers[req.Method]; exists {
+		s.mu.RUnlock()
+		return customHandler(req)
+	}
+	s.mu.RUnlock()
+
+	// Default handlers
 	switch req.Method {
 	case "initialize":
 		return s.handleInitialize(req)
@@ -343,6 +366,18 @@ func (s *Server) handleReadResource(req JSONRPCRequest) JSONRPCResponse {
 
 	s.mu.RLock()
 	handler, exists := s.resourceHandlers[readReq.URI]
+	if !exists {
+		// Try to find a pattern-based handler (e.g., for "file://" prefix)
+		for pattern, h := range s.resourceHandlers {
+			if pattern != "" && readReq.URI != pattern && 
+			   ((pattern == "file://" && strings.HasPrefix(readReq.URI, "file://")) ||
+			    (strings.HasSuffix(pattern, "*") && strings.HasPrefix(readReq.URI, strings.TrimSuffix(pattern, "*")))) {
+				handler = h
+				exists = true
+				break
+			}
+		}
+	}
 	s.mu.RUnlock()
 
 	if !exists {