Agent-Native/runtime/mcp/manager.go

package mcp

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"sync"
	"time"

	"agent-native/runtime/tools"
)

// 桥接上限：server 元数据是不可信输入，长度与数量都必须受限。
const (
	maxToolsPerServer    = 64
	maxDescriptionBytes  = 1024
	maxTextResultBytes   = 64 << 10
	maxToolNameBytes     = 128
	callTimeoutDefault   = 30 * time.Second
	registerTimeoutFixed = 30 * time.Second
)

// Manager 维护一组 MCP server 连接，并把它们的工具注册进 tools.Registry。
// 零值不可用；由组合根显式构造并注入（未注入 = 无 MCP 工具，fail-closed）。
type Manager struct {
	configs []ServerConfig
	mu      sync.Mutex
	clients map[string]Client
}

func NewManager(configs []ServerConfig) (*Manager, error) {
	seen := map[string]bool{}
	for _, c := range configs {
		if err := c.Validate(); err != nil {
			return nil, err
		}
		if seen[c.Name] {
			return nil, fmt.Errorf("mcp server %q configured twice", c.Name)
		}
		seen[c.Name] = true
	}
	return &Manager{configs: configs, clients: map[string]Client{}}, nil
}

// Configs 返回声明的 server 清单（只读副本）。
func (m *Manager) Configs() []ServerConfig {
	out := make([]ServerConfig, len(m.configs))
	copy(out, m.configs)
	return out
}

func (m *Manager) client(ctx context.Context, cfg ServerConfig) (Client, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if c, ok := m.clients[cfg.Name]; ok {
		return c, nil
	}
	dialCtx, cancel := context.WithTimeout(ctx, registerTimeoutFixed)
	defer cancel()
	var c Client
	var err error
	if cfg.Command != "" {
		c, err = DialStdio(dialCtx, cfg)
	} else {
		c, err = DialHTTP(dialCtx, cfg)
	}
	if err != nil {
		return nil, fmt.Errorf("mcp server %q unreachable (fail-closed): %w", cfg.Name, err)
	}
	m.clients[cfg.Name] = c
	return c, nil
}

// RegisterInto 连接全部 server、枚举工具并以 mcp__<server>__<tool> 命名注册。
// 任一 server 不可达即返回错误（fail-closed，不静默跳过）。
func (m *Manager) RegisterInto(ctx context.Context, reg *tools.Registry) error {
	for _, cfg := range m.configs {
		c, err := m.client(ctx, cfg)
		if err != nil {
			return err
		}
		listCtx, cancel := context.WithTimeout(ctx, registerTimeoutFixed)
		list, err := c.ListTools(listCtx)
		cancel()
		if err != nil {
			return fmt.Errorf("mcp server %q tools/list failed: %w", cfg.Name, err)
		}
		if len(list) > maxToolsPerServer {
			return fmt.Errorf("mcp server %q exposes %d tools (max %d)", cfg.Name, len(list), maxToolsPerServer)
		}
		for _, ti := range list {
			name := ToolName(cfg.Name, ti.Name)
			if len(name) > maxToolNameBytes {
				return fmt.Errorf("mcp tool name too long: %s", name)
			}
			desc := ti.Description
			if len(desc) > maxDescriptionBytes {
				desc = desc[:maxDescriptionBytes] + "…[truncated]"
			}
			schema := ti.InputSchema
			if len(schema) == 0 {
				schema = json.RawMessage(`{"type":"object","properties":{}}`)
			}
			reg.Register(&bridgeTool{
				client:      c,
				serverCfg:   cfg,
				toolName:    ti.Name,
				declName:    name,
				description: desc,
				schema:      schema,
			})
		}
	}
	return nil
}

// ToolName 生成桥接工具名（命名空间隔离）。
func ToolName(server, tool string) string {
	return "mcp__" + server + "__" + tool
}

// Close 关闭全部连接（进程型 server 随之退出）。
func (m *Manager) Close() {
	m.mu.Lock()
	defer m.mu.Unlock()
	for _, c := range m.clients {
		_ = c.Close()
	}
	m.clients = map[string]Client{}
}

// bridgeTool 把一个 MCP 工具适配为 tools.Tool；调用失败重连一次后重试。
type bridgeTool struct {
	client      Client
	serverCfg   ServerConfig
	toolName    string
	declName    string
	description string
	schema      json.RawMessage
}

func (t *bridgeTool) Name() string        { return t.declName }
func (t *bridgeTool) Description() string { return t.description }
func (t *bridgeTool) Parameters() json.RawMessage {
	return t.schema
}

func (t *bridgeTool) Execute(ctx context.Context, args json.RawMessage) (tools.Result, error) {
	callCtx, cancel := context.WithTimeout(ctx, requestTimeout(t.serverCfg))
	defer cancel()
	res, err := t.client.CallTool(callCtx, t.toolName, args)
	if err != nil {
		return tools.Result{Content: fmt.Sprintf("mcp tool %s failed: %v", t.declName, err), IsError: true}, nil
	}
	var b strings.Builder
	for _, blk := range res.Content {
		if blk.Type == "text" && blk.Text != "" {
			b.WriteString(blk.Text)
			b.WriteString("\n")
		}
	}
	out := b.String()
	if len(out) > maxTextResultBytes {
		out = out[:maxTextResultBytes] + "\n[truncated]"
	}
	if res.IsError {
		return tools.Result{Content: out, IsError: true}, nil
	}
	return tools.Result{Content: out}, nil
}