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 connects an MCP client and auto-registers its discovered tools into
// the shared capability registry, making them available to every character.
type Registry struct {
	reg    *registry.Registry
	mcp    *Client
	log    *slog.Logger
	mu     sync.Mutex
	names  []string
	closed bool
}

// NewRegistry builds a Registry that registers cfg's tools into reg. Options
// apply to the underlying Client (logger, injected transport). It is named
// NewRegistry rather than New because this package already exports New for
// Client. A tool-list-changed handler is wired automatically so the registry
// re-syncs when the server announces tools/list_changed.
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 connects the MCP server, discovers its tools, and registers each one
// into the underlying 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 re-syncs the registered tool set with the server after a
// notifications/tools/list_changed notification. It re-discovers the current
// tools and applies only the diff: it removes names that are gone and registers
// newly appeared ones, so tools registered by other providers (which collide by
// name) are preserved. It is invoked (off the transport read path) by the
// client's tool-list-changed handler and coalesced by the client.
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
	}

	// Drop names the server no longer exposes. Tools registered by other
	// providers were skipped at registration time and are not in r.names, so they
	// are never removed here.
	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

	// Register newly appeared tools. Skips names already present (own or
	// foreign) so a name collision with another provider is not doubled.
	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 returns a copy of the tool names this registry registered, in
// registration order. It is empty before Register succeeds.
func (r *Registry) ToolNames() []string {
	r.mu.Lock()
	defer r.mu.Unlock()
	return append([]string(nil), r.names...)
}

// Unregister removes every tool this registry previously registered and closes
// the underlying client. It is safe to call multiple times.
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
}