db_view/service/internal/common/mcp/component.go

package mcp

import (
	"context"

	"dbview/service/internal/common/logger"

	"github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"
	"go.uber.org/zap"
)

// Config MCP 组件配置（可按需扩展，例如新增 SSE/WS 等传输层参数）。
// 目前仅提供名称/版本控制与开关。
// 若未来需要从全局配置加载，可在 bootstrap 层读取后传入。
type Config struct {
	Enable        bool   // 是否启用 MCP 服务
	ServerName    string // MCP Server 名称（展示给客户端）
	ServerVersion string // MCP Server 版本
}

// DefaultConfig 提供一个安全的默认值（关闭状态）。
func DefaultConfig() Config {
	return Config{
		Enable:        false,
		ServerName:    "dbview-mcp",
		ServerVersion: "0.1.0",
	}
}

// Component 封装 mcp-go Server，便于在项目内统一注册工具、启动 Stdio 服务。
type Component struct {
	cfg    Config
	server *server.MCPServer
	logger logger.Logger
}

// NewComponent 创建 MCP 组件；仅注册 Server 与 Logger，不启动。
func NewComponent(cfg Config, logger logger.Logger) *Component {
	if cfg.ServerName == "" {
		cfg.ServerName = "dbview-mcp"
	}
	if cfg.ServerVersion == "" {
		cfg.ServerVersion = "0.1.0"
	}
	return &Component{
		cfg:    cfg,
		server: server.NewMCPServer(cfg.ServerName, cfg.ServerVersion),
		logger: logger,
	}
}

// RegisterHealthTool 注册一个内置健康检查工具，便于客户端验证 MCP 连接。
func (c *Component) RegisterHealthTool() {
	healthTool := mcp.NewTool("health",
		mcp.WithDescription("返回 MCP 服务健康状态"),
	)

	c.server.AddTool(healthTool, func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		return &mcp.CallToolResult{
			Content: []mcp.Content{mcp.TextContent{Text: "ok"}},
		}, nil
	})
}

// RegisterEchoTool 注册一个示例工具，可用于调试与演示。
// 参数：message（string，可选），返回与输入一致的文本。
func (c *Component) RegisterEchoTool() {
	echoTool := mcp.NewTool("echo",
		mcp.WithDescription("回显传入的 message 文本"),
		mcp.WithString("message", mcp.Description("要回显的文本")),
	)

	c.server.AddTool(echoTool, func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		args, _ := req.Params.Arguments.(map[string]any)
		msg, _ := args["message"].(string)
		if msg == "" {
			msg = "(empty)"
		}
		return &mcp.CallToolResult{
			Content: []mcp.Content{mcp.TextContent{Text: msg}},
		}, nil
	})
}

// Server 返回底层 MCPServer，便于外部注册自定义工具。
func (c *Component) Server() *server.MCPServer {
	return c.server
}

// ServeStdio 启动基于 stdio 的 MCP 服务（常用于被 IDE/CLI 以子进程方式拉起）。
func (c *Component) ServeStdio(ctx context.Context) error {
	if !c.cfg.Enable {
		c.logger.Warn("MCP 组件未启用，跳过 ServeStdio")
		return nil
	}
	c.logger.Info("启动 MCP Stdio 服务", zap.String("server", c.cfg.ServerName), zap.String("version", c.cfg.ServerVersion))
	// mcp-go 当前 Stdio Server 不接受外部 Context；如需退出，可由调用方中断进程。
	return server.ServeStdio(c.server)
}