regenerate init file when adding tools

Author: ilackarms

cmd/kmcp/cmd/add_tool.go

@@ -175,6 +175,15 @@ func generateTool(toolName, toolPath string) error {
 
 	fmt.Printf("✅ Successfully created tool: %s\n", toolName)
 	fmt.Printf("📁 Generated file: %s\n", toolPath)
+	
+	// Regenerate __init__.py file
+	toolsDir := filepath.Dir(toolPath)
+	if err := generator.RegenerateToolsInit(toolsDir); err != nil {
+		fmt.Printf("⚠️  Warning: Failed to regenerate __init__.py: %v\n", err)
+	} else {
+		fmt.Printf("🔄 Updated tools/__init__.py with new tool import\n")
+	}
+
 	fmt.Printf("📝 Edit the file to implement your tool logic\n")
 	fmt.Printf("🚀 The tool will be automatically loaded when the server starts\n")
 

cmd/kmcp/cmd/regenerate_init.go

@@ -0,0 +1,72 @@
+package cmd
+
+import (
+	"fmt"
+	"path/filepath"
+
+	"github.com/spf13/cobra"
+	"kagent.dev/kmcp/pkg/tools"
+)
+
+var regenerateInitCmd = &cobra.Command{
+	Use:   "regenerate-init",
+	Short: "Regenerate the tools/__init__.py file",
+	Long: `Regenerate the tools/__init__.py file based on the current tools in the directory.
+
+This command scans the src/tools/ directory for Python files and regenerates the __init__.py
+file with the appropriate imports and __all__ list. This is useful when:
+
+- You've manually added tool files
+- The __init__.py file is out of sync  
+- You need to refresh the tool imports
+
+The command will:
+1. Scan src/tools/ for all .py files (excluding __init__.py)
+2. Generate import statements for each tool
+3. Update the __all__ list with all discovered tools
+4. Overwrite the existing __init__.py file
+
+Examples:
+  kmcp regenerate-init                    # Regenerate in current directory
+  kmcp regenerate-init --tools-dir custom-tools  # Use custom tools directory`,
+	RunE: runRegenerateInit,
+}
+
+var regenerateInitToolsDir string
+
+func init() {
+	rootCmd.AddCommand(regenerateInitCmd)
+	regenerateInitCmd.Flags().StringVar(&regenerateInitToolsDir, "tools-dir", "src/tools", "Tools directory to regenerate __init__.py for")
+}
+
+func runRegenerateInit(cmd *cobra.Command, args []string) error {
+	// Check if we're in a valid KMCP project
+	if !isKMCPProject() {
+		return fmt.Errorf("not in a KMCP project directory. Run 'kmcp init' first")
+	}
+
+	// Check if tools directory exists
+	if !fileExists(regenerateInitToolsDir) {
+		return fmt.Errorf("tools directory '%s' does not exist", regenerateInitToolsDir)
+	}
+
+	// Get absolute path
+	absToolsDir, err := filepath.Abs(regenerateInitToolsDir)
+	if err != nil {
+		return fmt.Errorf("failed to get absolute path: %w", err)
+	}
+
+	if verbose {
+		fmt.Printf("Regenerating __init__.py in: %s\n", absToolsDir)
+	}
+
+	// Create generator and regenerate
+	generator := tools.NewGenerator()
+	if err := generator.RegenerateToolsInit(absToolsDir); err != nil {
+		return fmt.Errorf("failed to regenerate __init__.py: %w", err)
+	}
+
+	fmt.Printf("✅ Successfully regenerated %s/__init__.py\n", regenerateInitToolsDir)
+
+	return nil
+} 

pkg/tools/generator.go

@@ -57,9 +57,96 @@ func (g *Generator) GenerateToolFile(filePath, toolName string, config map[strin
 		return fmt.Errorf("failed to execute template: %w", err)
 	}
 
+	// After generating the tool file, regenerate the __init__.py file
+	toolsDir := filepath.Dir(filePath)
+	if err := g.RegenerateToolsInit(toolsDir); err != nil {
+		return fmt.Errorf("failed to regenerate __init__.py: %w", err)
+	}
+
+	return nil
+}
+
+// RegenerateToolsInit regenerates the __init__.py file in the tools directory
+func (g *Generator) RegenerateToolsInit(toolsDir string) error {
+	// Scan the tools directory for Python files
+	tools, err := g.ScanToolsDirectory(toolsDir)
+	if err != nil {
+		return fmt.Errorf("failed to scan tools directory: %w", err)
+	}
+
+	// Generate the __init__.py content
+	content := g.generateInitContent(tools)
+
+	// Write the __init__.py file
+	initPath := filepath.Join(toolsDir, "__init__.py")
+	if err := os.WriteFile(initPath, []byte(content), 0644); err != nil {
+		return fmt.Errorf("failed to write __init__.py: %w", err)
+	}
+
 	return nil
 }
 
+// ScanToolsDirectory scans the tools directory and returns a list of tool names
+func (g *Generator) ScanToolsDirectory(toolsDir string) ([]string, error) {
+	var tools []string
+
+	// Read the directory
+	entries, err := os.ReadDir(toolsDir)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read tools directory: %w", err)
+	}
+
+	// Find all Python files (excluding __init__.py)
+	for _, entry := range entries {
+		if entry.IsDir() {
+			continue
+		}
+
+		name := entry.Name()
+		if strings.HasSuffix(name, ".py") && name != "__init__.py" {
+			// Extract tool name (filename without .py extension)
+			toolName := strings.TrimSuffix(name, ".py")
+			tools = append(tools, toolName)
+		}
+	}
+
+	return tools, nil
+}
+
+// generateInitContent generates the content for the __init__.py file
+func (g *Generator) generateInitContent(tools []string) string {
+	var content strings.Builder
+
+	// Add the header comment
+	content.WriteString(`"""Tools package for knowledge-assistant MCP server.
+
+This file is automatically generated by the dynamic loading system.
+Do not edit manually - it will be overwritten when tools are loaded.
+"""
+
+`)
+
+	// Add import statements
+	for _, tool := range tools {
+		content.WriteString(fmt.Sprintf("from .%s import %s\n", tool, tool))
+	}
+
+	// Add empty line
+	content.WriteString("\n")
+
+	// Add __all__ list
+	content.WriteString("__all__ = [")
+	for i, tool := range tools {
+		if i > 0 {
+			content.WriteString(", ")
+		}
+		content.WriteString(fmt.Sprintf(`"%s"`, tool))
+	}
+	content.WriteString("]\n")
+
+	return content.String()
+}
+
 // getUnifiedTemplate returns the unified tool template with commented examples
 func (g *Generator) getUnifiedTemplate() string {
 	return `"""{{.ToolNameTitle}} tool for MCP server.{{if .description}}