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>
112 lines
2.5 KiB
Go
112 lines
2.5 KiB
Go
package server
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log/slog"
|
|
"net/http"
|
|
"os"
|
|
"os/signal"
|
|
"syscall"
|
|
|
|
"ops-mcp/internal/tool"
|
|
|
|
"github.com/mark3labs/mcp-go/server"
|
|
)
|
|
|
|
const version = "0.1.0"
|
|
|
|
// Run 启动 MCP Server，处理信号优雅关闭。
|
|
func Run(cfg *Config, tools ...tool.Tool) error {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
// 监听终止信号
|
|
sigCh := make(chan os.Signal, 1)
|
|
signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
|
|
go func() {
|
|
sig := <-sigCh
|
|
slog.Info("received signal, shutting down", "signal", sig)
|
|
cancel()
|
|
}()
|
|
|
|
// 创建注册中心
|
|
registry := NewRegistry()
|
|
for _, t := range tools {
|
|
if err := registry.Register(t); err != nil {
|
|
return fmt.Errorf("register: %w", err)
|
|
}
|
|
}
|
|
|
|
// 系统工具总是最后注册，确保它能看到所有工具
|
|
sysTool := tool.NewSystemTool(version, registry)
|
|
if err := registry.Register(sysTool); err != nil {
|
|
return fmt.Errorf("register system tool: %w", err)
|
|
}
|
|
|
|
// 初始化所有工具
|
|
if err := registry.InitializeAll(ctx); err != nil {
|
|
return fmt.Errorf("initialize: %w", err)
|
|
}
|
|
defer func() {
|
|
for _, err := range registry.ShutdownAll(ctx) {
|
|
slog.Error("shutdown error", "err", err)
|
|
}
|
|
}()
|
|
|
|
// 创建 MCP Server
|
|
mcpServer := server.NewMCPServer(
|
|
"ops-mcp",
|
|
version,
|
|
server.WithLogging(),
|
|
)
|
|
|
|
// 注册所有工具处理器
|
|
if err := registry.RegisterAll(mcpServer); err != nil {
|
|
return fmt.Errorf("register all: %w", err)
|
|
}
|
|
|
|
// 选择传输模式
|
|
switch cfg.Server.Transport {
|
|
case "stdio":
|
|
return runStdio(ctx, mcpServer)
|
|
case "sse":
|
|
return runSSE(ctx, mcpServer, cfg)
|
|
default:
|
|
return fmt.Errorf("unknown transport: %s (expect stdio or sse)", cfg.Server.Transport)
|
|
}
|
|
}
|
|
|
|
func runStdio(ctx context.Context, mcpServer *server.MCPServer) error {
|
|
slog.Info("starting MCP server", "transport", "stdio")
|
|
go func() {
|
|
<-ctx.Done()
|
|
slog.Info("shutting down stdio server")
|
|
}()
|
|
return server.ServeStdio(mcpServer)
|
|
}
|
|
|
|
func runSSE(ctx context.Context, mcpServer *server.MCPServer, cfg *Config) error {
|
|
sseServer := server.NewSSEServer(mcpServer)
|
|
|
|
addr := cfg.Server.Addr
|
|
if addr == "" {
|
|
addr = ":8080"
|
|
}
|
|
|
|
slog.Info("starting MCP server", "transport", "sse", "addr", addr)
|
|
|
|
httpServer := &http.Server{Addr: addr, Handler: sseServer}
|
|
|
|
go func() {
|
|
<-ctx.Done()
|
|
slog.Info("shutting down HTTP server")
|
|
httpServer.Shutdown(context.Background())
|
|
}()
|
|
|
|
if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|