ops-mcp/internal/server/registry.go

package server

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

	"ops-mcp/internal/tool"

	"github.com/mark3labs/mcp-go/server"
)

// Registry 管理所有工具的生命周期。
type Registry struct {
	tools map[string]tool.Tool
}

func NewRegistry() *Registry {
	return &Registry{tools: make(map[string]tool.Tool)}
}

func (r *Registry) Register(t tool.Tool) error {
	name := t.Name()
	if _, exists := r.tools[name]; exists {
		return fmt.Errorf("tool %q already registered", name)
	}
	r.tools[name] = t
	slog.Info("tool registered", "name", name)
	return nil
}

func (r *Registry) InitializeAll(ctx context.Context) error {
	for name, t := range r.tools {
		slog.Info("initializing tool", "name", name)
		if err := t.Initialize(ctx); err != nil {
			return fmt.Errorf("init tool %q: %w", name, err)
		}
	}
	return nil
}

func (r *Registry) RegisterAll(mcpServer *server.MCPServer) error {
	for name, t := range r.tools {
		if err := t.Register(mcpServer); err != nil {
			return fmt.Errorf("register tool %q: %w", name, err)
		}
	}
	return nil
}

func (r *Registry) ShutdownAll(ctx context.Context) []error {
	var errs []error
	for name, t := range r.tools {
		slog.Info("shutting down tool", "name", name)
		if err := t.Shutdown(ctx); err != nil {
			errs = append(errs, fmt.Errorf("shutdown tool %q: %w", name, err))
		}
	}
	return errs
}

func (r *Registry) HealthCheckAll(ctx context.Context) map[string]error {
	result := make(map[string]error, len(r.tools))
	for name, t := range r.tools {
		result[name] = t.HealthCheck(ctx)
	}
	return result
}