feat: centrally mount and dispatch agent MCP services

2026-09-28 13:09:51 +08:00
parent 6449de0cdb
commit 0c8ac02f82
14 changed files with 1097 additions and 8 deletions
@@ -0,0 +1,329 @@
+package cluster
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"os/exec"
+	"strings"
+	"sync"
+)
+
+type MCPService struct {
+	ID          string            `json:"id"`
+	Name        string            `json:"name"`
+	Description string            `json:"description,omitempty"`
+	Transport   string            `json:"transport"`
+	Command     string            `json:"command,omitempty"`
+	Args        []string          `json:"args,omitempty"`
+	Env         map[string]string `json:"env,omitempty"`
+	URL         string            `json:"url,omitempty"`
+	Headers     map[string]string `json:"headers,omitempty"`
+}
+
+type MCPCallResult struct {
+	Server     string `json:"server"`
+	Tool       string `json:"tool"`
+	Result     any    `json:"result,omitempty"`
+	Error      string `json:"error,omitempty"`
+	DurationMS int64  `json:"duration_ms"`
+}
+
+type mcpRPC struct {
+	JSONRPC string `json:"jsonrpc"`
+	ID      any    `json:"id,omitempty"`
+	Method  string `json:"method"`
+	Params  any    `json:"params,omitempty"`
+}
+
+type mcpRPCResponse struct {
+	JSONRPC string          `json:"jsonrpc"`
+	ID      any             `json:"id,omitempty"`
+	Result  json.RawMessage `json:"result,omitempty"`
+	Error   *struct {
+		Code    int    `json:"code"`
+		Message string `json:"message"`
+		Data    any    `json:"data,omitempty"`
+	} `json:"error,omitempty"`
+}
+
+type MCPRuntime struct {
+	mu       sync.RWMutex
+	services map[string]MCPService
+}
+
+func NewMCPRuntime() *MCPRuntime { return &MCPRuntime{services: make(map[string]MCPService)} }
+
+func (r *MCPRuntime) Sync(services []MCPService) {
+	next := make(map[string]MCPService, len(services))
+	for _, s := range services {
+		if s.ID == "" || s.Transport == "" {
+			continue
+		}
+		next[s.ID] = s
+	}
+	r.mu.Lock()
+	r.services = next
+	r.mu.Unlock()
+}
+
+func (r *MCPRuntime) List() []MCPService {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+	out := make([]MCPService, 0, len(r.services))
+	for _, s := range r.services {
+		out = append(out, s)
+	}
+	return out
+}
+
+func (r *MCPRuntime) Get(id string) (MCPService, bool) {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+	s, ok := r.services[id]
+	return s, ok
+}
+
+func (r *MCPRuntime) Call(ctx context.Context, server, tool string, args any) (any, error) {
+	s, ok := r.Get(server)
+	if !ok {
+		return nil, fmt.Errorf("MCP service %q is not mounted", server)
+	}
+	switch strings.ToLower(s.Transport) {
+	case "stdio":
+		return callStdio(ctx, s, tool, args)
+	case "http", "sse":
+		return callHTTP(ctx, s, tool, args)
+	default:
+		return nil, fmt.Errorf("unsupported MCP transport %q", s.Transport)
+	}
+}
+
+func callStdio(ctx context.Context, s MCPService, tool string, args any) (any, error) {
+	if s.Command == "" {
+		return nil, errors.New("MCP stdio command is empty")
+	}
+	cmd := exec.CommandContext(ctx, s.Command, s.Args...)
+	if len(s.Env) > 0 {
+		cmd.Env = append(os.Environ(), envList(s.Env)...)
+	}
+	stdin, err := cmd.StdinPipe()
+	if err != nil {
+		return nil, err
+	}
+	stdout, err := cmd.StdoutPipe()
+	if err != nil {
+		return nil, err
+	}
+	cmd.Stderr = io.Discard
+	if err := cmd.Start(); err != nil {
+		return nil, err
+	}
+	defer func() { _ = cmd.Process.Kill(); _ = cmd.Wait() }()
+
+	enc := json.NewEncoder(stdin)
+	dec := bufio.NewReaderSize(stdout, 64*1024)
+	id := int64(1)
+	if err := enc.Encode(mcpRPC{JSONRPC: "2.0", ID: id, Method: "initialize", Params: map[string]any{
+		"protocolVersion": "2024-11-05", "capabilities": map[string]any{},
+		"clientInfo": map[string]string{"name": "vps-commander-agent", "version": "1.0"},
+	}}); err != nil {
+		return nil, err
+	}
+	if _, err := readRPC(ctx, dec, id); err != nil {
+		return nil, err
+	}
+	if err := enc.Encode(mcpRPC{JSONRPC: "2.0", Method: "notifications/initialized", Params: map[string]any{}}); err != nil {
+		return nil, err
+	}
+	id++
+	if err := enc.Encode(mcpRPC{JSONRPC: "2.0", ID: id, Method: "tools/call", Params: map[string]any{"name": tool, "arguments": args}}); err != nil {
+		return nil, err
+	}
+	return readRPC(ctx, dec, id)
+}
+
+func readRPC(ctx context.Context, r *bufio.Reader, wantID int64) (any, error) {
+	for {
+		lineCh := make(chan []byte, 1)
+		errCh := make(chan error, 1)
+		go func() {
+			b, err := r.ReadBytes('\n')
+			if err != nil && len(b) == 0 {
+				errCh <- err
+				return
+			}
+			lineCh <- bytes.TrimSpace(b)
+		}()
+		var line []byte
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case err := <-errCh:
+			return nil, err
+		case line = <-lineCh:
+		}
+		if len(line) == 0 {
+			continue
+		}
+		var resp mcpRPCResponse
+		if json.Unmarshal(line, &resp) != nil {
+			continue
+		}
+		if resp.ID == nil {
+			continue
+		}
+		if n, ok := resp.ID.(float64); !ok || int64(n) != wantID {
+			continue
+		}
+		if resp.Error != nil {
+			return nil, fmt.Errorf("MCP error %d: %s", resp.Error.Code, resp.Error.Message)
+		}
+		var out any
+		if len(resp.Result) > 0 {
+			if err := json.Unmarshal(resp.Result, &out); err != nil {
+				return nil, err
+			}
+		}
+		return out, nil
+	}
+}
+
+func callHTTP(ctx context.Context, s MCPService, tool string, args any) (any, error) {
+	if s.URL == "" {
+		return nil, errors.New("MCP HTTP URL is empty")
+	}
+	client := &http.Client{Timeout: 0}
+	var session string
+	post := func(req mcpRPC) (any, http.Header, error) {
+		b, _ := json.Marshal(req)
+		h := make(http.Header)
+		h.Set("Content-Type", "application/json")
+		h.Set("Accept", "application/json, text/event-stream")
+		for k, v := range s.Headers {
+			h.Set(k, v)
+		}
+		if session != "" {
+			h.Set("Mcp-Session-Id", session)
+		}
+		hr, err := http.NewRequestWithContext(ctx, http.MethodPost, s.URL, bytes.NewReader(b))
+		if err != nil {
+			return nil, nil, err
+		}
+		hr.Header = h
+		resp, err := client.Do(hr)
+		if err != nil {
+			return nil, nil, err
+		}
+		defer resp.Body.Close()
+		if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+			return nil, resp.Header, fmt.Errorf("MCP HTTP status %s", resp.Status)
+		}
+		if v := resp.Header.Get("Mcp-Session-Id"); v != "" {
+			session = v
+		}
+		data, err := io.ReadAll(io.LimitReader(resp.Body, 4*1024*1024))
+		if err != nil {
+			return nil, nil, err
+		}
+		if len(bytes.TrimSpace(data)) == 0 {
+			return nil, resp.Header, nil
+		}
+		if bytes.Contains(data, []byte("event:")) {
+			data = parseSSEData(data)
+		}
+		var rr mcpRPCResponse
+		if err := json.Unmarshal(bytes.TrimSpace(data), &rr); err != nil {
+			return nil, resp.Header, err
+		}
+		if rr.Error != nil {
+			return nil, resp.Header, fmt.Errorf("MCP error %d: %s", rr.Error.Code, rr.Error.Message)
+		}
+		var out any
+		if len(rr.Result) > 0 {
+			if err := json.Unmarshal(rr.Result, &out); err != nil {
+				return nil, resp.Header, err
+			}
+		}
+		return out, resp.Header, nil
+	}
+	if strings.EqualFold(s.Transport, "sse") {
+		endpoint, err := discoverSSEEndpoint(ctx, s)
+		if err != nil {
+			return nil, err
+		}
+		old := s.URL
+		s.URL = endpoint
+		defer func() { s.URL = old }()
+	}
+	if _, _, err := post(mcpRPC{JSONRPC: "2.0", ID: int64(1), Method: "initialize", Params: map[string]any{"protocolVersion": "2024-11-05", "capabilities": map[string]any{}, "clientInfo": map[string]string{"name": "vps-commander-agent", "version": "1.0"}}}); err != nil {
+		return nil, err
+	}
+	if _, _, err := post(mcpRPC{JSONRPC: "2.0", Method: "notifications/initialized", Params: map[string]any{}}); err != nil {
+		return nil, err
+	}
+	return func() (any, error) {
+		v, _, err := post(mcpRPC{JSONRPC: "2.0", ID: int64(2), Method: "tools/call", Params: map[string]any{"name": tool, "arguments": args}})
+		return v, err
+	}()
+}
+
+func discoverSSEEndpoint(ctx context.Context, s MCPService) (string, error) {
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, s.URL, nil)
+	if err != nil {
+		return "", err
+	}
+	for k, v := range s.Headers {
+		req.Header.Set(k, v)
+	}
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return "", err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		return "", fmt.Errorf("MCP SSE status %s", resp.Status)
+	}
+	sc := bufio.NewScanner(io.LimitReader(resp.Body, 256*1024))
+	isEndpoint := false
+	for sc.Scan() {
+		line := sc.Text()
+		if strings.HasPrefix(line, "event:") {
+			isEndpoint = strings.TrimSpace(strings.TrimPrefix(line, "event:")) == "endpoint"
+			continue
+		}
+		if isEndpoint && strings.HasPrefix(line, "data:") {
+			u := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
+			if strings.HasPrefix(u, "/") {
+				base := *resp.Request.URL
+				base.Path = u
+				base.RawQuery = ""
+				return base.String(), nil
+			}
+			return u, nil
+		}
+	}
+	return "", errors.New("MCP SSE endpoint event not received")
+}
+
+func parseSSEData(b []byte) []byte {
+	for _, line := range strings.Split(string(b), "\n") {
+		if strings.HasPrefix(line, "data:") {
+			return []byte(strings.TrimSpace(strings.TrimPrefix(line, "data:")))
+		}
+	}
+	return b
+}
+func envList(m map[string]string) []string {
+	out := make([]string, 0, len(m))
+	for k, v := range m {
+		out = append(out, k+"="+v)
+	}
+	return out
+}