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
	// 工具数超过这个值且没有显式暴露级别时，整台服务器改为延迟暴露。
	deferredToolThreshold = 8
)

// 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
}

// ApplyOverrides 用项目配置覆盖启用状态和暴露级别。空字符串不改。
func (m *Manager) ApplyOverrides(raw string) error {
	if m == nil || strings.TrimSpace(raw) == "" {
		return nil
	}
	var over map[string]struct {
		Enabled  *bool  `json:"enabled"`
		Exposure string `json:"exposure"`
	}
	if err := json.Unmarshal([]byte(raw), &over); err != nil {
		return err
	}
	for i, cfg := range m.configs {
		o, ok := over[cfg.Name]
		if !ok {
			continue
		}
		if o.Enabled != nil {
			m.configs[i].Enabled = o.Enabled
		}
		if o.Exposure != "" {
			m.configs[i].Exposure = o.Exposure
		}
		if err := m.configs[i].Validate(); err != nil {
			return err
		}
	}
	return nil
}

// Ready 报告已启用的服务器是否都已经连上。已连上时注册只是枚举工具，不再握手。
func (m *Manager) Ready() bool {
	if m == nil {
		return true
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	for _, cfg := range m.configs {
		if !cfg.IsEnabled() {
			continue
		}
		if _, ok := m.clients[cfg.Name]; !ok {
			return false
		}
	}
	return true
}
func (m *Manager) StartBackground(ctx context.Context, reg *tools.Registry) {
	if m == nil || reg == nil {
		return
	}
	for _, cfg := range m.configs {
		cfg := cfg
		if !cfg.IsEnabled() {
			continue
		}
		go func() { _ = m.registerServer(ctx, reg, cfg) }()
	}
}

// RegisterInto 连接全部已启用的 server 并注册工具。未启用的跳过。
func (m *Manager) RegisterInto(ctx context.Context, reg *tools.Registry) error {
	for _, cfg := range m.configs {
		if !cfg.IsEnabled() {
			continue
		}
		if err := m.registerServer(ctx, reg, cfg); err != nil {
			return err
		}
	}
	return nil
}

func (m *Manager) registerServer(ctx context.Context, reg *tools.Registry, cfg ServerConfig) error {
	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)
	}
	exposure := chooseExposure(cfg.Exposure, len(list))
	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 = tools.CutUTF8(desc, maxDescriptionBytes) + "…[truncated]"
		}
		schema := ti.InputSchema
		if len(schema) == 0 {
			schema = json.RawMessage(`{"$schema":"https://json-schema.org/draft/2020-12/schema","type":"object","properties":{}}`)
		}
		reg.Register(&bridgeTool{
			client:      c,
			serverCfg:   cfg,
			toolName:    ti.Name,
			declName:    name,
			description: desc,
			schema:      schema,
			exposure:    exposure,
			redial: func(ctx context.Context) (Client, error) {
				m.mu.Lock()
				delete(m.clients, cfg.Name)
				m.mu.Unlock()
				return m.client(ctx, cfg)
			},
		})
	}
	return nil
}

func chooseExposure(explicit string, n int) string {
	switch explicit {
	case tools.ExposureDirect, tools.ExposureDeferred:
		return explicit
	default:
		if n > deferredToolThreshold {
			return tools.ExposureDeferred
		}
		return tools.ExposureDirect
	}
}

// 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。CallTool 失败时重连一次再试。
type bridgeTool struct {
	client      Client
	serverCfg   ServerConfig
	toolName    string
	declName    string
	description string
	schema      json.RawMessage
	exposure    string
	redial      func(ctx context.Context) (Client, error)
}

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) Exposure() string            { return t.exposure }

func (t *bridgeTool) Execute(ctx context.Context, args json.RawMessage) (tools.Result, error) {
	res, err := t.call(ctx, args)
	if err != nil && t.redial != nil {
		redial := t.redial
		t.redial = nil
		if c, rerr := redial(ctx); rerr == nil {
			t.client = c
			res, err = t.call(ctx, args)
		}
	}
	if err != nil {
		return tools.Result{Content: fmt.Sprintf("mcp tool %s failed: %v", t.declName, err), IsError: true}, nil
	}
	// input_required：2026-07-28 协议级结果（取应答后先 tasks/update 再
	// 用原参数重试 tools/call）；旧版内容块路径保持原样。
	if question, ok := inputQuestion(res, t.client.NegotiatedVersion()); ok && res.ResultType == "input_required" && SupportsResultType(t.client.NegotiatedVersion()) {
		return t.answerAndRetry(ctx, args, res, question)
	} else if ok {
		ask := askFrom(ctx)
		if ask == nil {
			return tools.Result{Content: "需要用户输入：" + question, IsError: true}, nil
		}
		answer, aerr := ask(ctx, question)
		if aerr != nil {
			return tools.Result{Content: aerr.Error(), IsError: true}, nil
		}
		return tools.Result{Content: answer}, 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 = tools.CutUTF8(out, maxTextResultBytes) + "\n[truncated]"
	}
	if res.IsError {
		return tools.Result{Content: out, IsError: true}, nil
	}
	return tools.Result{Content: out}, nil
}

// maxInputRounds 是 input_required 应答重试的上限（防服务端循环要输入）。
const maxInputRounds = 3

// answerAndRetry 处理 2026-07-28 的协议级 input_required：经 ask 端口取
// 应答，TaskID 非空时先 tasks/update 投递，然后用原参数重试 tools/call；
// 服务端连续要输入最多 maxInputRounds 轮。
func (t *bridgeTool) answerAndRetry(ctx context.Context, args json.RawMessage, res callToolResult, question string) (tools.Result, error) {
	for round := 0; round < maxInputRounds; round++ {
		ask := askFrom(ctx)
		if ask == nil {
			return tools.Result{Content: "需要用户输入：" + question, IsError: true}, nil
		}
		answer, aerr := ask(ctx, question)
		if aerr != nil {
			return tools.Result{Content: aerr.Error(), IsError: true}, nil
		}
		if res.TaskID != "" {
			if uerr := t.client.TasksUpdate(ctx, res.TaskID, answer); uerr != nil {
				return tools.Result{Content: fmt.Sprintf("mcp task update failed: %v", uerr), IsError: true}, nil
			}
		}
		var err error
		res, err = t.call(ctx, args)
		if err != nil {
			return tools.Result{Content: fmt.Sprintf("mcp tool %s failed: %v", t.declName, err), IsError: true}, nil
		}
		q, ok := inputQuestion(res, t.client.NegotiatedVersion())
		if !ok || res.ResultType != "input_required" || !SupportsResultType(t.client.NegotiatedVersion()) {
			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 res.IsError {
				return tools.Result{Content: out, IsError: true}, nil
			}
			return tools.Result{Content: strings.TrimSpace(out)}, nil
		}
		question = q
	}
	return tools.Result{Content: "需要用户输入：" + question, IsError: true}, nil
}

func (t *bridgeTool) call(ctx context.Context, args json.RawMessage) (callToolResult, error) {
	callCtx, cancel := context.WithTimeout(ctx, requestTimeout(t.serverCfg))
	defer cancel()
	res, err := t.client.CallTool(callCtx, t.toolName, args)
	if err == nil && res.TaskID != "" {
		trackerFrom(ctx).note(t.client, res.TaskID)
	}
	return res, err
}