MuyueWorkspace/internal/mcp/discover.go

package mcp

import (
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
)

type DiscoveredMCPServer struct {
	Name       string   `json:"name"`
	Command    string   `json:"command"`
	Source     string   `json:"source"`
	Args       []string `json:"args,omitempty"`
	Installed  bool     `json:"installed"`
	Running    bool     `json:"running"`
	Category   string   `json:"category,omitempty"`
}

type DiscoveryResult struct {
	Servers    []DiscoveredMCPServer `json:"servers"`
	ScanPaths  []string              `json:"scan_paths"`
	TotalFound int                   `json:"total_found"`
	NewServers int                   `json:"new_servers"`
}

type ToolDiscovery struct {
	Name        string          `json:"name"`
	Description string          `json:"description"`
	InputSchema json.RawMessage `json:"input_schema"`
}

type ServerCapabilities struct {
	Name      string           `json:"name"`
	Tools     []ToolDiscovery  `json:"tools"`
	Version   string           `json:"version,omitempty"`
	Raw       json.RawMessage  `json:"raw,omitempty"`
}

var (
	capCache     map[string]*ServerCapabilities
	capCacheMu   sync.RWMutex
)

func init() {
	capCache = make(map[string]*ServerCapabilities)
}

func DiscoverSystemServers() *DiscoveryResult {
	result := &DiscoveryResult{}

	knownNames := make(map[string]bool)
	for _, s := range knownMCPServers {
		knownNames[s.Name] = true
	}

	reg, _ := LoadRegistry()
	if reg != nil {
		for _, s := range reg.Servers {
			knownNames[s.Name] = true
		}
	}

	var servers []DiscoveredMCPServer

	npmServers := discoverNpmGlobalServers(knownNames)
	servers = append(servers, npmServers...)

	pipServers := discoverPipServers(knownNames)
	servers = append(servers, pipServers...)

	pathServers := discoverPathServers(knownNames)
	servers = append(servers, pathServers...)

	result.Servers = servers
	result.TotalFound = len(servers)
	result.NewServers = countNew(servers, knownNames)

	paths := []string{}
	if path := os.Getenv("PATH"); path != "" {
		paths = strings.Split(path, ":")
	}
	if home, err := os.UserHomeDir(); err == nil {
		paths = append(paths,
			filepath.Join(home, ".local", "bin"),
			filepath.Join(home, ".npm-global", "bin"),
		)
	}
	result.ScanPaths = paths

	return result
}

func discoverNpmGlobalServers(known map[string]bool) []DiscoveredMCPServer {
	var servers []DiscoveredMCPServer

	npx, err := exec.LookPath("npx")
	if err != nil {
		return servers
	}

	patterns := []struct {
		pkg   string
		name  string
		cat   string
	}{
		{"@anthropic/mcp-server-fetch", "anthropic-fetch", "web"},
		{"@anthropic/mcp-server-sqlite", "anthropic-sqlite", "database"},
		{"@anthropic/mcp-server-brave-search", "anthropic-brave-search", "web"},
		{"@anthropic/mcp-server-filesystem", "anthropic-filesystem", "core"},
		{"@anthropic/mcp-server-github", "anthropic-github", "vcs"},
		{"@anthropic/mcp-server-memory", "anthropic-memory", "core"},
		{"@anthropic/mcp-server-puppeteer", "anthropic-puppeteer", "web"},
		{"@anthropic/mcp-server-sequential-thinking", "anthropic-thinking", "ai"},
	}

	for _, p := range patterns {
		if known[p.name] {
			continue
		}

		servers = append(servers, DiscoveredMCPServer{
			Name:      p.name,
			Command:   npx,
			Source:    "npm-global",
			Args:      []string{"-y", p.pkg},
			Installed: true,
			Category:  p.cat,
		})
	}

	return servers
}

func discoverPipServers(known map[string]bool) []DiscoveredMCPServer {
	var servers []DiscoveredMCPServer

	pipCmds := []string{"pip", "pip3", "uv"}
	for _, pip := range pipCmds {
		if _, err := exec.LookPath(pip); err != nil {
			continue
		}

		cmd := exec.Command(pip, "list", "--format=json")
		output, err := cmd.CombinedOutput()
		if err != nil {
			continue
		}

		var packages []struct {
			Name    string `json:"name"`
			Version string `json:"version"`
		}
		if err := json.Unmarshal(output, &packages); err != nil {
			continue
		}

		for _, pkg := range packages {
			nameLower := strings.ToLower(pkg.Name)
			if !strings.Contains(nameLower, "mcp") {
				continue
			}

			serverName := strings.ReplaceAll(nameLower, "_", "-")
			if strings.HasPrefix(serverName, "mcp-") {
				serverName = serverName[4:]
			}

			if known[serverName] {
				continue
			}

			binName := strings.ReplaceAll(pkg.Name, "-", "_")
			if _, err := exec.LookPath(binName); err != nil {
				binName = pkg.Name
				if _, err := exec.LookPath(binName); err != nil {
					continue
				}
			}

			servers = append(servers, DiscoveredMCPServer{
				Name:      serverName,
				Command:   binName,
				Source:    "pip",
				Installed: true,
				Category:  "python",
			})
		}
		break
	}

	return servers
}

func discoverPathServers(known map[string]bool) []DiscoveredMCPServer {
	var servers []DiscoveredMCPServer

	home, _ := os.UserHomeDir()
	searchDirs := []string{}

	if home != "" {
		searchDirs = append(searchDirs,
			filepath.Join(home, ".local", "bin"),
			filepath.Join(home, ".muyue", "mcp-servers"),
		)
	}

	for _, dir := range searchDirs {
		entries, err := os.ReadDir(dir)
		if err != nil {
			continue
		}

		for _, entry := range entries {
			if entry.IsDir() {
				continue
			}

			name := entry.Name()
			if !strings.Contains(strings.ToLower(name), "mcp") {
				continue
			}

			serverName := strings.ToLower(name)
			serverName = strings.TrimPrefix(serverName, "mcp-")
			serverName = strings.TrimPrefix(serverName, "mcp_")
			serverName = strings.TrimSuffix(serverName, ".sh")

			if known[serverName] {
				continue
			}

			fullPath := filepath.Join(dir, name)
			if info, err := os.Stat(fullPath); err == nil && info.Mode()&0111 != 0 {
				servers = append(servers, DiscoveredMCPServer{
					Name:      serverName,
					Command:   fullPath,
					Source:    "path",
					Installed: true,
					Category:  "local",
				})
			}
		}
	}

	return servers
}

func DiscoverServerTools(serverName string) (*ServerCapabilities, error) {
	capCacheMu.RLock()
	if caps, ok := capCache[serverName]; ok {
		capCacheMu.RUnlock()
		return caps, nil
	}
	capCacheMu.RUnlock()

	server, err := findServerConfig(serverName)
	if err != nil {
		return nil, err
	}

	script := buildListToolsScript(server)
	if script == "" {
		return &ServerCapabilities{
			Name:  serverName,
			Tools: []ToolDiscovery{},
		}, nil
	}

	cmd := exec.Command(server.Command, append(server.Args, "--list-tools")...)
	output, err := cmd.CombinedOutput()
	_ = script

	if err != nil {
		return discoverToolsFallback(serverName, server)
	}

	var caps ServerCapabilities
	if jsonErr := json.Unmarshal(output, &caps); jsonErr != nil {
		caps = ServerCapabilities{
			Name: serverName,
			Tools: []ToolDiscovery{
				{
					Name:        serverName,
					Description: "MCP server: " + serverName,
				},
			},
		}
	}

	capCacheMu.Lock()
	capCache[serverName] = &caps
	capCacheMu.Unlock()

	return &caps, nil
}

func discoverToolsFallback(name string, server *RegistryServer) (*ServerCapabilities, error) {
	caps := &ServerCapabilities{
		Name: name,
		Tools: []ToolDiscovery{
			{
				Name:        name,
				Description: server.Description,
			},
		},
	}

	capCacheMu.Lock()
	capCache[name] = caps
	capCacheMu.Unlock()

	return caps, nil
}

func findServerConfig(name string) (*RegistryServer, error) {
	reg, err := LoadRegistry()
	if err != nil {
		return nil, err
	}

	for i := range reg.Servers {
		if reg.Servers[i].Name == name {
			return &reg.Servers[i], nil
		}
	}

	for _, s := range knownMCPServers {
		if s.Name == name {
			return &RegistryServer{
				Name:    s.Name,
				Command: s.Command,
				Args:    s.Args,
				Env:     s.Env,
			}, nil
		}
	}

	return nil, fmt.Errorf("server %q not found", name)
}

func buildListToolsScript(server *RegistryServer) string {
	return ""
}

func InvalidateCapabilitiesCache() {
	capCacheMu.Lock()
	defer capCacheMu.Unlock()
	capCache = make(map[string]*ServerCapabilities)
}

func GetCachedCapabilities(name string) *ServerCapabilities {
	capCacheMu.RLock()
	defer capCacheMu.RUnlock()
	return capCache[name]
}

func countNew(servers []DiscoveredMCPServer, known map[string]bool) int {
	count := 0
	for _, s := range servers {
		if !known[s.Name] {
			count++
		}
	}
	return count
}