Commit 519f3fb

mo khan <mo@mokhan.ca>
2025-08-15 15:23:59
refactor: add handler to Tool
1 parent dc8d81e
Changed files (3)
pkg/filesystem/server.go
@@ -16,8 +16,6 @@ type Server struct {
 }
 
 func New(allowedDirs []string) *Server {
-  server := mcp.NewServer("filesystem", "0.2.0")
-
   normalizedDirs := make([]string, len(allowedDirs))
   for i, dir := range allowedDirs {
     absPath, err := filepath.Abs(expandHome(dir))
@@ -28,11 +26,76 @@ func New(allowedDirs []string) *Server {
   }
 
   fsServer := &Server{
-    Server:             server,
     allowedDirectories: normalizedDirs,
   }
 
-  fsServer.registerTools()
+  tools := []mcp.Tool{
+    mcp.NewTool("read_file", "Read the contents of a file", map[string]interface{}{
+      "type": "object",
+      "properties": map[string]interface{}{
+        "path": map[string]interface{}{
+          "type": "string",
+        },
+      },
+      "required": []string{"path"},
+    }, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
+      pathStr, ok := req.Arguments["path"].(string)
+      if !ok {
+        return mcp.NewToolError("path is required"), nil
+      }
+
+      validPath, err := fsServer.validatePath(pathStr)
+      if err != nil {
+        return mcp.NewToolError(err.Error()), nil
+      }
+
+      content, err := os.ReadFile(validPath)
+      if err != nil {
+        return mcp.NewToolError(fmt.Sprintf("Failed to read file: %v", err)), nil
+      }
+
+      return mcp.NewToolResult(mcp.NewTextContent(string(content))), nil
+    }),
+
+    mcp.NewTool("write_file", "Write content to a file", map[string]interface{}{
+      "type": "object",
+      "properties": map[string]interface{}{
+        "path": map[string]interface{}{
+          "type": "string",
+        },
+        "content": map[string]interface{}{
+          "type": "string",
+        },
+      },
+      "required": []string{"path", "content"},
+    }, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
+      pathStr, ok := req.Arguments["path"].(string)
+      if !ok {
+        return mcp.NewToolError("path is required"), nil
+      }
+
+      content, ok := req.Arguments["content"].(string)
+      if !ok {
+        return mcp.NewToolError("content is required"), nil
+      }
+
+      validPath, err := fsServer.validatePath(pathStr)
+      if err != nil {
+        return mcp.NewToolError(err.Error()), nil
+      }
+
+      err = os.WriteFile(validPath, []byte(content), 0644)
+      if err != nil {
+        return mcp.NewToolError(fmt.Sprintf("Failed to write file: %v", err)), nil
+      }
+
+      return mcp.NewToolResult(mcp.NewTextContent(fmt.Sprintf("Successfully wrote to %s", pathStr))), nil
+    }),
+  }
+
+  server := mcp.NewServer("filesystem", "0.2.0", tools...)
+  fsServer.Server = server
+
   fsServer.registerResources()
   fsServer.registerRoots()
   fsServer.setupResourceHandling()
@@ -40,22 +103,6 @@ func New(allowedDirs []string) *Server {
   return fsServer
 }
 
-func (fs *Server) registerTools() {
-  tools := fs.ListTools()
-  
-  for _, tool := range tools {
-    var handler mcp.ToolHandler
-    switch tool.Name {
-    case "read_file":
-      handler = fs.HandleReadFile
-    case "write_file":
-      handler = fs.HandleWriteFile
-    default:
-      continue
-    }
-    fs.RegisterToolWithDefinition(tool, handler)
-  }
-}
 
 func (fs *Server) registerResources() {
   for _, dir := range fs.allowedDirectories {
@@ -252,80 +299,6 @@ func isBinaryContent(content []byte) bool {
   return false
 }
 
-func (fs *Server) ListTools() []mcp.Tool {
-  return []mcp.Tool{
-    {
-      Name: "read_file",
-      InputSchema: map[string]interface{}{
-        "type": "object",
-        "properties": map[string]interface{}{
-          "path": map[string]interface{}{
-            "type": "string",
-          },
-        },
-        "required": []string{"path"},
-      },
-    },
-    {
-      Name: "write_file",
-      InputSchema: map[string]interface{}{
-        "type": "object",
-        "properties": map[string]interface{}{
-          "path": map[string]interface{}{
-            "type": "string",
-          },
-          "content": map[string]interface{}{
-            "type": "string",
-          },
-        },
-        "required": []string{"path", "content"},
-      },
-    },
-  }
-}
-
-func (fs *Server) HandleReadFile(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
-  pathStr, ok := req.Arguments["path"].(string)
-  if !ok {
-    return mcp.NewToolError("path is required"), nil
-  }
-
-  validPath, err := fs.validatePath(pathStr)
-  if err != nil {
-    return mcp.NewToolError(err.Error()), nil
-  }
-
-  content, err := os.ReadFile(validPath)
-  if err != nil {
-    return mcp.NewToolError(fmt.Sprintf("Failed to read file: %v", err)), nil
-  }
-
-  return mcp.NewToolResult(mcp.NewTextContent(string(content))), nil
-}
-
-func (fs *Server) HandleWriteFile(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
-  pathStr, ok := req.Arguments["path"].(string)
-  if !ok {
-    return mcp.NewToolError("path is required"), nil
-  }
-
-  content, ok := req.Arguments["content"].(string)
-  if !ok {
-    return mcp.NewToolError("content is required"), nil
-  }
-
-  validPath, err := fs.validatePath(pathStr)
-  if err != nil {
-    return mcp.NewToolError(err.Error()), nil
-  }
-
-  err = os.WriteFile(validPath, []byte(content), 0644)
-  if err != nil {
-    return mcp.NewToolError(fmt.Sprintf("Failed to write file: %v", err)), nil
-  }
-
-  return mcp.NewToolResult(mcp.NewTextContent(fmt.Sprintf("Successfully wrote to %s", pathStr))), nil
-}
 
 func (fs *Server) validatePath(requestedPath string) (string, error) {
   expandedPath := expandHome(requestedPath)
pkg/mcp/server.go
@@ -42,8 +42,8 @@ type PromptHandler func(GetPromptRequest) (GetPromptResult, error)
 type ResourceHandler func(ReadResourceRequest) (ReadResourceResult, error)
 
 // NewServer creates a new MCP server
-func NewServer(name, version string) *Server {
-	return &Server{
+func NewServer(name, version string, tools ...Tool) *Server {
+	server := &Server{
 		name:                  name,
 		version:               version,
 		toolHandlers:          make(map[string]ToolHandler),
@@ -62,6 +62,13 @@ func NewServer(name, version string) *Server {
 			Logging:   &LoggingCapability{},
 		},
 	}
+
+	for _, tool := range tools {
+		server.toolHandlers[tool.Name] = tool.Handler
+		server.toolDefinitions[tool.Name] = tool
+	}
+
+	return server
 }
 
 
@@ -441,6 +448,16 @@ func NewToolError(message string) CallToolResult {
 	}
 }
 
+// Helper function to create a new tool with all fields
+func NewTool(name, description string, inputSchema interface{}, handler ToolHandler) Tool {
+	return Tool{
+		Name:        name,
+		Description: description,
+		InputSchema: inputSchema,
+		Handler:     handler,
+	}
+}
+
 // Helper function to extract resource name from URI
 func extractResourceName(uri string) string {
 	// Find the last "/" in the URI and extract the part after it
pkg/mcp/types.go
@@ -95,6 +95,7 @@ type Tool struct {
 	Name        string      `json:"name"`
 	Description string      `json:"description"`
 	InputSchema interface{} `json:"inputSchema"`
+	Handler     ToolHandler `json:"-"`
 }
 
 type ListToolsResult struct {