MuyueWorkspace/internal/mcpserver/server.go

package mcpserver

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"time"
)

type Tool struct {
	Name        string                 `json:"name"`
	Description string                 `json:"description"`
	InputSchema map[string]interface{} `json:"inputSchema"`
}

type ToolCall struct {
	Name string          `json:"name"`
	Args json.RawMessage `json:"arguments"`
}

type ToolResult struct {
	Content []ContentBlock `json:"content"`
	IsError bool           `json:"isError,omitempty"`
}

type ContentBlock struct {
	Type string `json:"type"`
	Text string `json:"text,omitempty"`
}

type JSONRPCRequest struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id,omitempty"`
	Method  string          `json:"method"`
	Params  json.RawMessage `json:"params,omitempty"`
}

type JSONRPCResponse struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id"`
	Result  interface{}     `json:"result,omitempty"`
	Error   *RPCError       `json:"error,omitempty"`
}

type RPCError struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

var tools = []Tool{
	{
		Name:        "terminal_exec",
		Description: "Execute a command in the terminal and return the output",
		InputSchema: map[string]interface{}{
			"type": "object",
			"properties": map[string]interface{}{
				"command": map[string]interface{}{"type": "string", "description": "The command to execute"},
				"cwd":     map[string]interface{}{"type": "string", "description": "Working directory (optional)"},
				"timeout": map[string]interface{}{"type": "integer", "description": "Timeout in seconds (default 30)"},
			},
			"required": []string{"command"},
		},
	},
	{
		Name:        "file_read",
		Description: "Read the contents of a file",
		InputSchema: map[string]interface{}{
			"type": "object",
			"properties": map[string]interface{}{
				"path":    map[string]interface{}{"type": "string", "description": "Path to the file"},
				"offset":  map[string]interface{}{"type": "integer", "description": "Line offset to start reading from (0-based)"},
				"limit":   map[string]interface{}{"type": "integer", "description": "Maximum number of lines to read"},
			},
			"required": []string{"path"},
		},
	},
	{
		Name:        "file_write",
		Description: "Write content to a file, creating it if needed",
		InputSchema: map[string]interface{}{
			"type": "object",
			"properties": map[string]interface{}{
				"path":    map[string]interface{}{"type": "string", "description": "Path to the file"},
				"content": map[string]interface{}{"type": "string", "description": "Content to write"},
			},
			"required": []string{"path", "content"},
		},
	},
	{
		Name:        "search",
		Description: "Search for files by name pattern",
		InputSchema: map[string]interface{}{
			"type": "object",
			"properties": map[string]interface{}{
				"path":    map[string]interface{}{"type": "string", "description": "Directory to search in"},
				"pattern": map[string]interface{}{"type": "string", "description": "Glob pattern to match filenames"},
			},
			"required": []string{"path", "pattern"},
		},
	},
	{
		Name:        "grep",
		Description: "Search file contents for a pattern",
		InputSchema: map[string]interface{}{
			"type": "object",
			"properties": map[string]interface{}{
				"path":    map[string]interface{}{"type": "string", "description": "Directory to search in"},
				"pattern": map[string]interface{}{"type": "string", "description": "Text or regex pattern to search for"},
			},
			"required": []string{"path", "pattern"},
		},
	},
	{
		Name:        "system_info",
		Description: "Get system information (OS, CPU, memory, disk)",
		InputSchema: map[string]interface{}{
			"type":       "object",
			"properties": map[string]interface{}{},
		},
	},
}

type MCPServer struct {
	port   int
	server *http.Server
	mu     sync.Mutex

	sseClients   map[string]chan SSEEvent
	sseClientsMu sync.Mutex
}

type SSEEvent struct {
	Event string
	Data  string
}

func New(port int) *MCPServer {
	return &MCPServer{
		port:       port,
		sseClients: make(map[string]chan SSEEvent),
	}
}

func (m *MCPServer) Start() error {
	mux := http.NewServeMux()
	mux.HandleFunc("/", m.handleSSE)
	mux.HandleFunc("/message", m.handleHTTPMessage)
	mux.HandleFunc("/mcp", m.handleStreamableHTTP)

	m.server = &http.Server{
		Addr:    fmt.Sprintf("127.0.0.1:%d", m.port),
		Handler: mux,
	}

	go func() {
		if err := m.server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			fmt.Printf("[MCP Server] Error: %v\n", err)
		}
	}()

	return nil
}

func (m *MCPServer) Stop() error {
	if m.server != nil {
		return m.server.Close()
	}
	return nil
}

func (m *MCPServer) Port() int {
	return m.port
}

func (m *MCPServer) handleSSE(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("Access-Control-Allow-Origin", "*")

	clientID := fmt.Sprintf("%d", time.Now().UnixNano())
	ch := make(chan SSEEvent, 32)

	m.sseClientsMu.Lock()
	m.sseClients[clientID] = ch
	m.sseClientsMu.Unlock()

	defer func() {
		m.sseClientsMu.Lock()
		delete(m.sseClients, clientID)
		m.sseClientsMu.Unlock()
		close(ch)
	}()

	fmt.Fprintf(w, "event: endpoint\ndata: /message?clientId=%s\n\n", clientID)
	if f, ok := w.(http.Flusher); ok {
		f.Flush()
	}

	for {
		select {
		case evt, ok := <-ch:
			if !ok {
				return
			}
			fmt.Fprintf(w, "event: %s\ndata: %s\n\n", evt.Event, evt.Data)
			if f, ok := w.(http.Flusher); ok {
				f.Flush()
			}
		case <-r.Context().Done():
			return
		}
	}
}

func (m *MCPServer) handleHTTPMessage(w http.ResponseWriter, r *http.Request) {
	if r.Method == "OPTIONS" {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type")
		w.WriteHeader(http.StatusOK)
		return
	}

	if r.Method != "POST" {
		http.Error(w, "POST only", http.StatusMethodNotAllowed)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Access-Control-Allow-Origin", "*")

	body, err := io.ReadAll(r.Body)
	if err != nil {
		m.writeRPCError(w, nil, -32700, "Parse error")
		return
	}

	resp := m.handleJSONRPC(body)
	json.NewEncoder(w).Encode(resp)
}

func (m *MCPServer) handleStreamableHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method == "OPTIONS" {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type")
		w.WriteHeader(http.StatusOK)
		return
	}

	if r.Method == "GET" {
		m.handleSSE(w, r)
		return
	}

	if r.Method != "POST" {
		http.Error(w, "POST only", http.StatusMethodNotAllowed)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Access-Control-Allow-Origin", "*")

	body, err := io.ReadAll(r.Body)
	if err != nil {
		m.writeRPCError(w, nil, -32700, "Parse error")
		return
	}

	resp := m.handleJSONRPC(body)
	json.NewEncoder(w).Encode(resp)
}

func (m *MCPServer) handleJSONRPC(body []byte) JSONRPCResponse {
	var req JSONRPCRequest
	if err := json.Unmarshal(body, &req); err != nil {
		return JSONRPCResponse{
			JSONRPC: "2.0",
			Error:   &RPCError{Code: -32700, Message: "Parse error"},
		}
	}

	switch req.Method {
	case "initialize":
		return JSONRPCResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Result: map[string]interface{}{
				"protocolVersion": "2024-11-05",
				"capabilities": map[string]interface{}{
					"tools": map[string]interface{}{},
				},
				"serverInfo": map[string]interface{}{
					"name":    "muyue",
					"version": "0.9.0",
				},
			},
		}

	case "notifications/initialized":
		return JSONRPCResponse{JSONRPC: "2.0", ID: req.ID}

	case "tools/list":
		return JSONRPCResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Result: map[string]interface{}{
				"tools": tools,
			},
		}

	case "tools/call":
		var params struct {
			Name      string          `json:"name"`
			Arguments json.RawMessage `json:"arguments"`
		}
		if err := json.Unmarshal(req.Params, &params); err != nil {
			return JSONRPCResponse{
				JSONRPC: "2.0",
				ID:      req.ID,
				Error:   &RPCError{Code: -32602, Message: "Invalid params"},
			}
		}
		result := m.executeTool(params.Name, params.Arguments)
		return JSONRPCResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Result:  result,
		}

	default:
		return JSONRPCResponse{
			JSONRPC: "2.0",
			ID:      req.ID,
			Error:   &RPCError{Code: -32601, Message: fmt.Sprintf("Method not found: %s", req.Method)},
		}
	}
}

func (m *MCPServer) executeTool(name string, args json.RawMessage) ToolResult {
	switch name {
	case "terminal_exec":
		return m.toolTerminalExec(args)
	case "file_read":
		return m.toolFileRead(args)
	case "file_write":
		return m.toolFileWrite(args)
	case "search":
		return m.toolSearch(args)
	case "grep":
		return m.toolGrep(args)
	case "system_info":
		return m.toolSystemInfo()
	default:
		return ToolResult{
			Content: []ContentBlock{{Type: "text", Text: fmt.Sprintf("Unknown tool: %s", name)}},
			IsError: true,
		}
	}
}

func (m *MCPServer) toolTerminalExec(args json.RawMessage) ToolResult {
	var params struct {
		Command string `json:"command"`
		Cwd     string `json:"cwd"`
		Timeout int    `json:"timeout"`
	}
	if err := json.Unmarshal(args, &params); err != nil {
		return ToolResult{Content: []ContentBlock{{Type: "text", Text: err.Error()}}, IsError: true}
	}

	timeout := params.Timeout
	if timeout <= 0 {
		timeout = 30
	}

	ctx := exec.Command("sh", "-c", params.Command)
	if params.Cwd != "" {
		ctx.Dir = params.Cwd
	}

	var stdout, stderr strings.Builder
	ctx.Stdout = &stdout
	ctx.Stderr = &stderr

	done := make(chan error, 1)
	go func() { done <- ctx.Run() }()

	select {
	case err := <-done:
		output := stdout.String()
		if errMsg := stderr.String(); errMsg != "" {
			output += "\n" + errMsg
		}
		if err != nil {
			output += fmt.Sprintf("\nExit error: %v", err)
		}
		return ToolResult{Content: []ContentBlock{{Type: "text", Text: output}}}
	case <-time.After(time.Duration(timeout) * time.Second):
		ctx.Process.Kill()
		return ToolResult{
			Content: []ContentBlock{{Type: "text", Text: fmt.Sprintf("Command timed out after %ds\n%s%s", timeout, stdout.String(), stderr.String())}},
			IsError: true,
		}
	}
}

func (m *MCPServer) toolFileRead(args json.RawMessage) ToolResult {
	var params struct {
		Path   string `json:"path"`
		Offset int    `json:"offset"`
		Limit  int    `json:"limit"`
	}
	if err := json.Unmarshal(args, &params); err != nil {
		return ToolResult{Content: []ContentBlock{{Type: "text", Text: err.Error()}}, IsError: true}
	}

	home, _ := os.UserHomeDir()
	path := strings.ReplaceAll(params.Path, "~", home)

	data, err := os.ReadFile(path)
	if err != nil {
		return ToolResult{Content: []ContentBlock{{Type: "text", Text: fmt.Sprintf("Error reading file: %v", err)}}, IsError: true}
	}

	lines := strings.Split(string(data), "\n")
	start := params.Offset
	if start < 0 {
		start = 0
	}
	end := len(lines)
	if params.Limit > 0 && start+params.Limit < end {
		end = start + params.Limit
	}

	if start > len(lines) {
		start = len(lines)
	}

	return ToolResult{Content: []ContentBlock{{Type: "text", Text: strings.Join(lines[start:end], "\n")}}}
}

func (m *MCPServer) toolFileWrite(args json.RawMessage) ToolResult {
	var params struct {
		Path    string `json:"path"`
		Content string `json:"content"`
	}
	if err := json.Unmarshal(args, &params); err != nil {
		return ToolResult{Content: []ContentBlock{{Type: "text", Text: err.Error()}}, IsError: true}
	}

	home, _ := os.UserHomeDir()
	path := strings.ReplaceAll(params.Path, "~", home)

	dir := filepath.Dir(path)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return ToolResult{Content: []ContentBlock{{Type: "text", Text: fmt.Sprintf("Error creating directory: %v", err)}}, IsError: true}
	}

	if err := os.WriteFile(path, []byte(params.Content), 0644); err != nil {
		return ToolResult{Content: []ContentBlock{{Type: "text", Text: fmt.Sprintf("Error writing file: %v", err)}}, IsError: true}
	}

	return ToolResult{Content: []ContentBlock{{Type: "text", Text: fmt.Sprintf("Successfully wrote %d bytes to %s", len(params.Content), path)}}}
}

func (m *MCPServer) toolSearch(args json.RawMessage) ToolResult {
	var params struct {
		Path    string `json:"path"`
		Pattern string `json:"pattern"`
	}
	if err := json.Unmarshal(args, &params); err != nil {
		return ToolResult{Content: []ContentBlock{{Type: "text", Text: err.Error()}}, IsError: true}
	}

	home, _ := os.UserHomeDir()
	basePath := strings.ReplaceAll(params.Path, "~", home)

	cmd := exec.Command("find", basePath, "-name", params.Pattern, "-type", "f", "-not", "-path", "*/node_modules/*", "-not", "-path", "*/.git/*")
	output, err := cmd.CombinedOutput()
	if err != nil {
		return ToolResult{Content: []ContentBlock{{Type: "text", Text: string(output)}}}
	}

	lines := strings.Split(strings.TrimSpace(string(output)), "\n")
	if len(lines) > 100 {
		lines = lines[:100]
	}

	return ToolResult{Content: []ContentBlock{{Type: "text", Text: strings.Join(lines, "\n")}}}
}

func (m *MCPServer) toolGrep(args json.RawMessage) ToolResult {
	var params struct {
		Path    string `json:"path"`
		Pattern string `json:"pattern"`
	}
	if err := json.Unmarshal(args, &params); err != nil {
		return ToolResult{Content: []ContentBlock{{Type: "text", Text: err.Error()}}, IsError: true}
	}

	home, _ := os.UserHomeDir()
	basePath := strings.ReplaceAll(params.Path, "~", home)

	cmd := exec.Command("grep", "-rn", "--include=*", params.Pattern, basePath)
	output, _ := cmd.CombinedOutput()

	lines := strings.Split(strings.TrimSpace(string(output)), "\n")
	if len(lines) > 50 {
		lines = lines[:50]
		lines = append(lines, fmt.Sprintf("... (%d more results truncated)", len(strings.Split(string(output), "\n"))-50))
	}

	return ToolResult{Content: []ContentBlock{{Type: "text", Text: strings.Join(lines, "\n")}}}
}

func (m *MCPServer) toolSystemInfo() ToolResult {
	var info strings.Builder
	info.WriteString(fmt.Sprintf("OS: %s/%s\n", runtime.GOOS, runtime.GOARCH))
	info.WriteString(fmt.Sprintf("CPUs: %d\n", runtime.NumCPU()))

	if out, err := exec.Command("uname", "-a").Output(); err == nil {
		info.WriteString(fmt.Sprintf("Kernel: %s", string(out)))
	}

	if out, err := exec.Command("free", "-h").Output(); err == nil {
		info.WriteString(fmt.Sprintf("Memory:\n%s", string(out)))
	}

	if out, err := exec.Command("df", "-h", "/").Output(); err == nil {
		info.WriteString(fmt.Sprintf("Disk:\n%s", string(out)))
	}

	if out, err := exec.Command("uptime").Output(); err == nil {
		info.WriteString(fmt.Sprintf("Uptime: %s", string(out)))
	}

	return ToolResult{Content: []ContentBlock{{Type: "text", Text: info.String()}}}
}

func (m *MCPServer) writeRPCError(w http.ResponseWriter, id json.RawMessage, code int, msg string) {
	resp := JSONRPCResponse{
		JSONRPC: "2.0",
		ID:      id,
		Error:   &RPCError{Code: code, Message: msg},
	}
	json.NewEncoder(w).Encode(resp)
}