yangzhaohan
fb2c3a6cae
- Tool 接口 + Registry 注册中心，工具模块化插拔 - system 工具：health / info - database 工具：db_query / db_tables / db_table_info / db_explain（四层 SQL 安全） - docker 工具：docker_ps / docker_logs / docker_inspect（白名单 + 大小限制） - middleware：auth（JWT+APIKey）/ ratelimit（令牌桶）/ audit（slog 结构化） - stdio / SSE 双传输模式，Viper 多环境配置 - Dockerfile 多阶段构建（golang:alpine → scratch），~15MB 镜像 - docker-compose.yml 一键部署 + 安全加固（read_only / no-new-privileges / cap_drop） - go build ./... / go vet ./... / go test ./... 全部通过 Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
84 lines
2.1 KiB
Go
84 lines
2.1 KiB
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)}
|
|
}
|
|
|
|
// Register 注册一个工具。如果工具名重复则报错。
|
|
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, "desc", t.Description())
|
|
return nil
|
|
}
|
|
|
|
// InitializeAll 调用所有工具的 Initialize，任一失败则终止。
|
|
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
|
|
}
|
|
|
|
// RegisterAll 将所有工具注册到 MCP Server。
|
|
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)
|
|
}
|
|
slog.Info("tool handlers registered", "name", name)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ShutdownAll 调用所有工具的 Shutdown，收集所有错误。
|
|
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
|
|
}
|
|
|
|
// HealthCheckAll 检查所有工具的连通性，用于 system health 工具。
|
|
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
|
|
}
|
|
|
|
// List 返回所有已注册工具的名称。
|
|
func (r *Registry) List() []string {
|
|
names := make([]string, 0, len(r.tools))
|
|
for name := range r.tools {
|
|
names = append(names, name)
|
|
}
|
|
return names
|
|
}
|