Cyrene_v2/internal/capability/mcp/registry.go

package mcp

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"sync"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
)

// Registry 连接一台 MCP 客户端，把自动发现的工具注册进共享的 capability
// registry，使所有角色都能使用。
type Registry struct {
	reg    *registry.Registry
	mcp    *Client
	log    *slog.Logger
	mu     sync.Mutex
	names  []string
	closed bool
}

// NewRegistry 构建一个 Registry，把 cfg 的工具注册进 reg。opts 作用在底层
// Client 上（日志器、注入的传输）。之所以叫 NewRegistry 而不是 New，是因为
// 本包已为 Client 导出了 New。这里会自动接上工具清单变更回调，
// 服务器发 tools/list_changed 时 registry 会重新同步。
func NewRegistry(reg *registry.Registry, cfg MCPConfig, opts ...Option) *Registry {
	r := &Registry{reg: reg}
	opts = append(opts, WithToolListChangedHandler(r.refresh))
	c := New(cfg, opts...)
	r.mcp = c
	r.log = c.log
	return r
}

// Register 连接 MCP 服务器、发现其工具，并把每把工具注册进底层 registry。
func (r *Registry) Register(ctx context.Context) error {
	if r.reg == nil {
		return fmt.Errorf("mcp registry: nil registry")
	}
	if err := r.mcp.Connect(ctx); err != nil {
		return fmt.Errorf("mcp registry %q: %w", r.mcp.cfg.Name, err)
	}
	tools, err := r.mcp.DiscoverTools(ctx)
	if err != nil {
		return fmt.Errorf("mcp registry %q: %w", r.mcp.cfg.Name, err)
	}

	r.mu.Lock()
	defer r.mu.Unlock()
	if r.closed {
		return fmt.Errorf("mcp registry %q: %w", r.mcp.cfg.Name, ErrClosed)
	}
	for _, t := range tools {
		if err := r.reg.RegisterTool(t); err != nil {
			return fmt.Errorf("mcp register %q: %w", t.Name(), err)
		}
		r.names = append(r.names, t.Name())
	}
	r.log.Info("mcp tools registered", "name", r.mcp.cfg.Name, "count", len(tools))
	return nil
}

// refresh 在收到 notifications/tools/list_changed 通知后，把已注册工具集与
// 服务器重新对齐。它重新发现当前工具并只应用差异：移除已消失的名字、
// 注册新出现的名字，因此其它提供者注册的同名工具（名字冲突）会被保留。
// 它由客户端的工具清单变更回调调用（不在传输读路径上），并由客户端合并触发。
func (r *Registry) refresh() {
	if r.reg == nil {
		return
	}
	tools, err := r.mcp.DiscoverTools(context.Background())
	if err != nil {
		r.log.Warn("mcp tool refresh failed", "name", r.mcp.cfg.Name, "err", err)
		return
	}

	incoming := make(map[string]bool, len(tools))
	for _, t := range tools {
		incoming[t.Name()] = true
	}

	r.mu.Lock()
	defer r.mu.Unlock()
	if r.closed {
		return
	}

	// 丢弃服务器不再暴露的名字。其它提供者注册的工具在注册时就被跳过、
	// 不在 r.names 里，因此这里永远不会把它们删掉。
	kept := r.names[:0]
	for _, name := range r.names {
		if !incoming[name] {
			if err := r.reg.Remove(name); err != nil && !errors.Is(err, registry.ErrNotFound) {
				r.log.Warn("mcp refresh remove failed", "name", name, "err", err)
			}
			continue
		}
		kept = append(kept, name)
	}
	r.names = kept

	// 注册新出现的工具。跳过已存在的名字（无论是自己的还是外来的），
	// 避免与其它提供者同名时重复注册。
	existing := make(map[string]bool, len(r.names))
	for _, name := range r.names {
		existing[name] = true
	}
	for _, t := range tools {
		if existing[t.Name()] {
			continue
		}
		if err := r.reg.RegisterTool(t); err != nil {
			if !errors.Is(err, registry.ErrDuplicate) {
				r.log.Warn("mcp refresh register failed", "name", t.Name(), "err", err)
			}
			continue
		}
		r.names = append(r.names, t.Name())
	}
	r.log.Info("mcp tools refreshed", "name", r.mcp.cfg.Name, "count", len(r.names))
}

// ToolNames 返回本 registry 注册过的工具名副本（按注册顺序）。
// Register 成功之前为空。
func (r *Registry) ToolNames() []string {
	r.mu.Lock()
	defer r.mu.Unlock()
	return append([]string(nil), r.names...)
}

// Unregister 移除本 registry 之前注册的全部工具，并关闭底层客户端。
// 可安全重复调用。
func (r *Registry) Unregister(ctx context.Context) error {
	r.mu.Lock()
	names := r.names
	r.names = nil
	r.closed = true
	r.mu.Unlock()

	var errs error
	for _, name := range names {
		if err := r.reg.Remove(name); err != nil && !errors.Is(err, registry.ErrNotFound) {
			errs = errors.Join(errs, fmt.Errorf("mcp unregister %q: %w", name, err))
		}
	}
	if err := r.mcp.Close(); err != nil {
		errs = errors.Join(errs, fmt.Errorf("mcp close %q: %w", r.mcp.cfg.Name, err))
	}
	return errs
}