526 lines
15 KiB
Go
526 lines
15 KiB
Go
// Package mcpclient manages MCP server configuration and short-lived discovery sessions.
|
||
package mcpclient
|
||
|
||
import (
|
||
"context"
|
||
"crypto/rand"
|
||
"encoding/hex"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"io"
|
||
"net"
|
||
"net/http"
|
||
"net/url"
|
||
"os"
|
||
"path/filepath"
|
||
"reflect"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/zalando/go-keyring"
|
||
)
|
||
|
||
type ServerConfig struct {
|
||
ID string `json:"id"`
|
||
Name string `json:"name"`
|
||
Enabled bool `json:"enabled"`
|
||
Transport string `json:"transport"`
|
||
URL string `json:"url"`
|
||
Command string `json:"command"`
|
||
Args []string `json:"args"`
|
||
WorkingDirectory string `json:"workingDirectory"`
|
||
TimeoutSeconds int `json:"timeoutSeconds"`
|
||
}
|
||
|
||
// Secrets are write-only and never included in configuration files or views.
|
||
type Secrets struct {
|
||
Env map[string]string `json:"env"`
|
||
Headers map[string]string `json:"headers"`
|
||
}
|
||
type ServerInput struct {
|
||
Config ServerConfig `json:"config"`
|
||
Secrets Secrets `json:"secrets"`
|
||
ReplaceSecrets bool `json:"replaceSecrets"`
|
||
}
|
||
type ServerView struct {
|
||
Config ServerConfig `json:"config"`
|
||
HasSecrets bool `json:"hasSecrets"`
|
||
}
|
||
type ToolInfo struct {
|
||
Name string `json:"name"`
|
||
Description string `json:"description"`
|
||
InputSchema string `json:"inputSchema"`
|
||
}
|
||
type TestResult struct {
|
||
OK bool `json:"ok"`
|
||
Message string `json:"message"`
|
||
ElapsedMs int64 `json:"elapsedMs"`
|
||
Tools []ToolInfo `json:"tools"`
|
||
}
|
||
|
||
type secretStore interface {
|
||
Get(string) (string, error)
|
||
Set(string, string) error
|
||
Delete(string) error
|
||
}
|
||
type systemSecrets struct{}
|
||
|
||
func (systemSecrets) Get(id string) (string, error) { return keyring.Get("com.ituoke.mcp", id) }
|
||
func (systemSecrets) Set(id, value string) error { return keyring.Set("com.ituoke.mcp", id, value) }
|
||
func (systemSecrets) Delete(id string) error { return keyring.Delete("com.ituoke.mcp", id) }
|
||
|
||
type Service struct {
|
||
mu sync.Mutex
|
||
path string
|
||
secrets secretStore
|
||
servers []ServerConfig
|
||
loadErr error
|
||
agentGate chan struct{}
|
||
}
|
||
|
||
func NewService() *Service {
|
||
dir, err := os.UserConfigDir()
|
||
s := newService(filepath.Join(dir, "iTuoke", "mcp.json"), systemSecrets{})
|
||
if err != nil {
|
||
s.loadErr = fmt.Errorf("无法定位 MCP 配置目录")
|
||
}
|
||
if s.loadErr == nil {
|
||
if err := s.ensureMobile(); err != nil {
|
||
s.loadErr = err
|
||
}
|
||
}
|
||
if s.loadErr == nil {
|
||
if err := s.ensureFilesystem(); err != nil {
|
||
s.loadErr = err
|
||
}
|
||
}
|
||
return s
|
||
}
|
||
func newService(path string, secrets secretStore) *Service {
|
||
s := &Service{path: path, secrets: secrets, servers: []ServerConfig{}}
|
||
raw, err := os.ReadFile(path)
|
||
if errors.Is(err, os.ErrNotExist) {
|
||
return s
|
||
}
|
||
if err == nil {
|
||
err = json.Unmarshal(raw, &s.servers)
|
||
}
|
||
if err == nil {
|
||
seen := map[string]bool{}
|
||
for i, c := range s.servers {
|
||
if c.ID == "" || seen[c.ID] {
|
||
err = fmt.Errorf("duplicate ID")
|
||
break
|
||
}
|
||
seen[c.ID] = true
|
||
s.servers[i], err = normalize(c)
|
||
if err != nil {
|
||
break
|
||
}
|
||
}
|
||
}
|
||
if err != nil {
|
||
s.loadErr = fmt.Errorf("MCP 配置读取失败，请检查 mcp.json 后重新启动应用")
|
||
}
|
||
return s
|
||
}
|
||
|
||
func normalize(c ServerConfig) (ServerConfig, error) {
|
||
c.Name = strings.TrimSpace(c.Name)
|
||
c.URL = strings.TrimSpace(c.URL)
|
||
c.Command = strings.TrimSpace(c.Command)
|
||
c.WorkingDirectory = strings.TrimSpace(c.WorkingDirectory)
|
||
if c.Name == "" || len(c.Name) > 100 {
|
||
return c, fmt.Errorf("服务名称应为 1–100 字节")
|
||
}
|
||
if c.TimeoutSeconds < 3 || c.TimeoutSeconds > 120 {
|
||
return c, fmt.Errorf("连接超时应为 3–120 秒")
|
||
}
|
||
if len(c.ID) > 64 || strings.ContainsAny(c.ID, "/\\\x00") {
|
||
return c, fmt.Errorf("服务标识无效")
|
||
}
|
||
if c.ID == filesystemServerID && c.Transport != "filesystem" {
|
||
return c, fmt.Errorf("系统内置 Filesystem 的连接方式不可修改")
|
||
}
|
||
if c.ID == mobileServerID && c.Transport != "mobile" {
|
||
return c, fmt.Errorf("系统内置 Mobile Next 的连接方式不可修改")
|
||
}
|
||
switch c.Transport {
|
||
case "mobile":
|
||
c.Command, c.WorkingDirectory, c.URL = "", "", ""
|
||
c.Args = []string{}
|
||
case "filesystem":
|
||
if len(c.Args) == 0 || len(c.Args) > 20 {
|
||
return c, fmt.Errorf("请指定 1–20 个允许访问的目录")
|
||
}
|
||
for _, dir := range c.Args {
|
||
if !filepath.IsAbs(dir) || strings.ContainsAny(dir, "\x00\r\n") {
|
||
return c, fmt.Errorf("允许目录必须为有效的绝对路径")
|
||
}
|
||
}
|
||
c.Command, c.WorkingDirectory, c.URL = "", "", ""
|
||
case "http":
|
||
u, err := url.Parse(c.URL)
|
||
if err != nil || u.Hostname() == "" || u.User != nil || u.RawQuery != "" || u.Fragment != "" || (u.Scheme != "https" && u.Scheme != "http") {
|
||
return c, fmt.Errorf("请填写有效的 MCP HTTP 地址，认证信息请放入请求头")
|
||
}
|
||
ip := net.ParseIP(u.Hostname())
|
||
if u.Scheme == "http" && u.Hostname() != "localhost" && (ip == nil || (!ip.IsLoopback() && !ip.IsPrivate())) {
|
||
return c, fmt.Errorf("公网 MCP 地址请使用 HTTPS")
|
||
}
|
||
c.Command, c.WorkingDirectory, c.Args = "", "", []string{}
|
||
case "stdio":
|
||
if c.Command == "" || len(c.Command) > 4096 || strings.ContainsAny(c.Command+c.WorkingDirectory, "\x00\r\n") {
|
||
return c, fmt.Errorf("请填写有效的可执行文件路径或命令名称")
|
||
}
|
||
if c.WorkingDirectory != "" && !filepath.IsAbs(c.WorkingDirectory) {
|
||
return c, fmt.Errorf("工作目录请使用绝对路径")
|
||
}
|
||
if len(c.Args) > 128 {
|
||
return c, fmt.Errorf("启动参数不能超过 128 项")
|
||
}
|
||
for _, arg := range c.Args {
|
||
if len(arg) > 8192 || strings.ContainsRune(arg, 0) {
|
||
return c, fmt.Errorf("启动参数无效或过长")
|
||
}
|
||
}
|
||
if c.Args == nil {
|
||
c.Args = []string{}
|
||
}
|
||
c.URL = ""
|
||
default:
|
||
return c, fmt.Errorf("请选择 stdio 或 Streamable HTTP")
|
||
}
|
||
return c, nil
|
||
}
|
||
|
||
func validateSecrets(c ServerConfig, sec Secrets) error {
|
||
if (c.Transport == "filesystem" || c.Transport == "mobile") && len(sec.Env)+len(sec.Headers) > 0 {
|
||
return fmt.Errorf("内置 MCP 服务不需要凭据")
|
||
}
|
||
raw, _ := json.Marshal(sec)
|
||
if len(raw) > 16000 || len(sec.Env)+len(sec.Headers) > 100 {
|
||
return fmt.Errorf("环境变量或请求头配置过大")
|
||
}
|
||
if c.Transport == "http" && len(sec.Env) > 0 {
|
||
return fmt.Errorf("HTTP 服务请配置请求头")
|
||
}
|
||
if c.Transport == "stdio" && len(sec.Headers) > 0 {
|
||
return fmt.Errorf("stdio 服务请配置环境变量")
|
||
}
|
||
for k, v := range sec.Env {
|
||
if k == "" || strings.ContainsAny(k, "=\x00\r\n") || strings.ContainsRune(v, 0) {
|
||
return fmt.Errorf("环境变量名称或值无效")
|
||
}
|
||
}
|
||
for k, v := range sec.Headers {
|
||
if k == "" || strings.TrimSpace(k) != k || strings.ContainsAny(v, "\x00\r\n") {
|
||
return fmt.Errorf("请求头名称或值无效")
|
||
}
|
||
for _, r := range k {
|
||
if !strings.ContainsRune("!#$%&'*+-.^_`|~0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ", r) {
|
||
return fmt.Errorf("请求头名称无效")
|
||
}
|
||
}
|
||
switch strings.ToLower(k) {
|
||
case "host", "content-type", "accept", "content-length", "connection", "transfer-encoding":
|
||
return fmt.Errorf("请勿覆盖 HTTP 协议请求头")
|
||
}
|
||
if strings.HasPrefix(strings.ToLower(k), "mcp-") {
|
||
return fmt.Errorf("请勿覆盖 MCP 协议请求头")
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (s *Service) index(id string) int {
|
||
for i, c := range s.servers {
|
||
if c.ID == id {
|
||
return i
|
||
}
|
||
}
|
||
return -1
|
||
}
|
||
func (s *Service) readSecret(id string) (string, error) {
|
||
if id == "" {
|
||
return "", nil
|
||
}
|
||
raw, err := s.secrets.Get(id)
|
||
if errors.Is(err, keyring.ErrNotFound) {
|
||
return "", nil
|
||
}
|
||
if err != nil {
|
||
return "", fmt.Errorf("无法访问系统凭据库")
|
||
}
|
||
return raw, nil
|
||
}
|
||
func sameTarget(a, b ServerConfig) bool {
|
||
return a.Transport == b.Transport && a.URL == b.URL && a.Command == b.Command && a.WorkingDirectory == b.WorkingDirectory && reflect.DeepEqual(a.Args, b.Args)
|
||
}
|
||
func (s *Service) resolve(input ServerInput) (ServerConfig, Secrets, string, error) {
|
||
c, err := normalize(input.Config)
|
||
if err != nil {
|
||
return c, Secrets{}, "", err
|
||
}
|
||
if i := s.index(c.ID); i >= 0 && systemTransport(s.servers[i].Transport) && c.Transport != s.servers[i].Transport {
|
||
return c, Secrets{}, "", fmt.Errorf("系统内置 MCP 的连接方式不可修改")
|
||
}
|
||
if err := validateFilesystemDirs(c); err != nil {
|
||
return c, Secrets{}, "", err
|
||
}
|
||
i := s.index(c.ID)
|
||
if c.ID != "" && i < 0 {
|
||
return c, Secrets{}, "", fmt.Errorf("MCP 服务不存在，请重新加载")
|
||
}
|
||
old, err := s.readSecret(c.ID)
|
||
if err != nil {
|
||
return c, Secrets{}, "", err
|
||
}
|
||
sec := input.Secrets
|
||
if !input.ReplaceSecrets {
|
||
if len(sec.Env)+len(sec.Headers) > 0 {
|
||
return c, sec, old, fmt.Errorf("填写新凭据时请选择替换凭据")
|
||
}
|
||
if old != "" && !sameTarget(s.servers[i], c) {
|
||
return c, sec, old, fmt.Errorf("连接目标已更改，请重新填写或清除凭据")
|
||
}
|
||
if old != "" && json.Unmarshal([]byte(old), &sec) != nil {
|
||
return c, sec, old, fmt.Errorf("MCP 凭据读取失败，请重新填写")
|
||
}
|
||
}
|
||
return c, sec, old, validateSecrets(c, sec)
|
||
}
|
||
|
||
func (s *Service) ListServers() ([]ServerView, error) {
|
||
s.mu.Lock()
|
||
defer s.mu.Unlock()
|
||
if s.loadErr != nil {
|
||
return nil, s.loadErr
|
||
}
|
||
views := make([]ServerView, 0, len(s.servers))
|
||
for _, c := range s.servers {
|
||
raw, err := s.readSecret(c.ID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
c.Args = append([]string{}, c.Args...)
|
||
views = append(views, ServerView{Config: c, HasSecrets: raw != ""})
|
||
}
|
||
return views, nil
|
||
}
|
||
|
||
func (s *Service) persist(servers []ServerConfig) error {
|
||
if err := os.MkdirAll(filepath.Dir(s.path), 0700); err != nil {
|
||
return fmt.Errorf("无法创建 MCP 配置目录")
|
||
}
|
||
f, err := os.CreateTemp(filepath.Dir(s.path), "mcp-*.tmp")
|
||
if err != nil {
|
||
return fmt.Errorf("无法写入 MCP 配置")
|
||
}
|
||
defer os.Remove(f.Name())
|
||
data, _ := json.MarshalIndent(servers, "", " ")
|
||
_, writeErr := f.Write(data)
|
||
closeErr := f.Close()
|
||
if writeErr != nil || closeErr != nil {
|
||
return fmt.Errorf("MCP 配置写入失败")
|
||
}
|
||
if err := os.Rename(f.Name(), s.path); err != nil {
|
||
return fmt.Errorf("MCP 配置保存失败")
|
||
}
|
||
return nil
|
||
}
|
||
func (s *Service) setSecret(id, raw string) error {
|
||
var err error
|
||
if raw == "" {
|
||
err = s.secrets.Delete(id)
|
||
} else {
|
||
err = s.secrets.Set(id, raw)
|
||
}
|
||
if err != nil && !errors.Is(err, keyring.ErrNotFound) {
|
||
return fmt.Errorf("MCP 凭据保存失败，请检查系统凭据库")
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (s *Service) SaveServer(input ServerInput) (ServerView, error) {
|
||
s.mu.Lock()
|
||
defer s.mu.Unlock()
|
||
if s.loadErr != nil {
|
||
return ServerView{}, s.loadErr
|
||
}
|
||
c, sec, old, err := s.resolve(input)
|
||
if err != nil {
|
||
return ServerView{}, err
|
||
}
|
||
if c.ID == "" {
|
||
if len(s.servers) >= 50 {
|
||
return ServerView{}, fmt.Errorf("最多添加 50 个 MCP 服务")
|
||
}
|
||
var id [16]byte
|
||
if _, err := rand.Read(id[:]); err != nil {
|
||
return ServerView{}, fmt.Errorf("无法创建服务标识")
|
||
}
|
||
c.ID = hex.EncodeToString(id[:])
|
||
}
|
||
for _, other := range s.servers {
|
||
if other.ID != c.ID && strings.EqualFold(other.Name, c.Name) {
|
||
return ServerView{}, fmt.Errorf("服务名称已存在")
|
||
}
|
||
}
|
||
raw := ""
|
||
if len(sec.Env)+len(sec.Headers) > 0 {
|
||
b, _ := json.Marshal(sec)
|
||
raw = string(b)
|
||
}
|
||
if raw != old {
|
||
if err := s.setSecret(c.ID, raw); err != nil {
|
||
return ServerView{}, err
|
||
}
|
||
}
|
||
next := append([]ServerConfig{}, s.servers...)
|
||
if i := s.index(c.ID); i >= 0 {
|
||
next[i] = c
|
||
} else {
|
||
next = append(next, c)
|
||
}
|
||
if err := s.persist(next); err != nil {
|
||
if raw != old {
|
||
if rollback := s.setSecret(c.ID, old); rollback != nil {
|
||
return ServerView{}, fmt.Errorf("配置保存失败，凭据恢复失败，请重新填写凭据")
|
||
}
|
||
}
|
||
return ServerView{}, err
|
||
}
|
||
s.servers = next
|
||
return ServerView{Config: c, HasSecrets: raw != ""}, nil
|
||
}
|
||
|
||
func (s *Service) SetEnabled(id string, enabled bool) error {
|
||
s.mu.Lock()
|
||
defer s.mu.Unlock()
|
||
if s.loadErr != nil {
|
||
return s.loadErr
|
||
}
|
||
i := s.index(id)
|
||
if i < 0 {
|
||
return fmt.Errorf("MCP 服务不存在")
|
||
}
|
||
next := append([]ServerConfig{}, s.servers...)
|
||
next[i].Enabled = enabled
|
||
if err := s.persist(next); err != nil {
|
||
return err
|
||
}
|
||
s.servers = next
|
||
return nil
|
||
}
|
||
|
||
func (s *Service) DeleteServer(id string) error {
|
||
s.mu.Lock()
|
||
defer s.mu.Unlock()
|
||
if s.loadErr != nil {
|
||
return s.loadErr
|
||
}
|
||
if id == mobileServerID {
|
||
return fmt.Errorf("系统内置 Mobile Next 可停用，无需删除")
|
||
}
|
||
i := s.index(id)
|
||
if i < 0 {
|
||
return fmt.Errorf("MCP 服务不存在")
|
||
}
|
||
if systemTransport(s.servers[i].Transport) {
|
||
return fmt.Errorf("系统内置 MCP 可停用或编辑，不允许删除")
|
||
}
|
||
old, err := s.readSecret(id)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if old != "" {
|
||
if err := s.setSecret(id, ""); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
next := append([]ServerConfig{}, s.servers[:i]...)
|
||
next = append(next, s.servers[i+1:]...)
|
||
if err := s.persist(next); err != nil {
|
||
if old != "" {
|
||
if rollback := s.setSecret(id, old); rollback != nil {
|
||
return fmt.Errorf("删除失败，凭据恢复失败，请重新填写凭据")
|
||
}
|
||
}
|
||
return err
|
||
}
|
||
s.servers = next
|
||
return nil
|
||
}
|
||
|
||
// TestServer discovers tools using the draft without saving or executing tools.
|
||
// Sessions are closed after discovery; enabled is a saved availability preference.
|
||
func (s *Service) TestServer(ctx context.Context, input ServerInput) TestResult {
|
||
start := time.Now()
|
||
result := TestResult{Tools: []ToolInfo{}}
|
||
s.mu.Lock()
|
||
c, sec, _, err := s.resolve(input)
|
||
if s.loadErr != nil {
|
||
err = s.loadErr
|
||
}
|
||
s.mu.Unlock()
|
||
if err == nil {
|
||
var cancel context.CancelFunc
|
||
ctx, cancel = context.WithTimeout(ctx, time.Duration(c.TimeoutSeconds)*time.Second)
|
||
defer cancel()
|
||
result.Tools, err = discover(ctx, c, sec)
|
||
if err != nil {
|
||
switch {
|
||
case errors.Is(ctx.Err(), context.Canceled):
|
||
err = fmt.Errorf("连接测试已取消")
|
||
case errors.Is(ctx.Err(), context.DeadlineExceeded):
|
||
err = fmt.Errorf("连接超时，请检查服务是否启动以及连接地址")
|
||
default:
|
||
err = fmt.Errorf("连接或工具发现失败，请检查启动命令、运行环境、MCP 地址和认证配置")
|
||
}
|
||
}
|
||
}
|
||
result.ElapsedMs = time.Since(start).Milliseconds()
|
||
result.OK = err == nil
|
||
if err != nil {
|
||
result.Message = err.Error()
|
||
result.Tools = []ToolInfo{}
|
||
} else {
|
||
result.Message = fmt.Sprintf("连接成功，发现 %d 个工具", len(result.Tools))
|
||
}
|
||
return result
|
||
}
|
||
|
||
// Request headers are attached only to this service's own HTTP client.
|
||
type headerTransport struct {
|
||
base http.RoundTripper
|
||
headers map[string]string
|
||
ctx context.Context
|
||
}
|
||
|
||
func (t headerTransport) RoundTrip(req *http.Request) (*http.Response, error) {
|
||
ctx, cancel := context.WithCancel(req.Context())
|
||
stop := context.AfterFunc(t.ctx, cancel)
|
||
cleanup := func() { stop(); cancel() }
|
||
clone := req.Clone(ctx)
|
||
for k, v := range t.headers {
|
||
clone.Header.Set(k, v)
|
||
}
|
||
res, err := t.base.RoundTrip(clone)
|
||
if err != nil {
|
||
cleanup()
|
||
return nil, err
|
||
}
|
||
res.Body = &cleanupBody{ReadCloser: res.Body, cleanup: cleanup}
|
||
return res, nil
|
||
}
|
||
|
||
type cleanupBody struct {
|
||
io.ReadCloser
|
||
cleanup func()
|
||
}
|
||
|
||
func (b *cleanupBody) Close() error { defer b.cleanup(); return b.ReadCloser.Close() }
|