db_view/service/internal/bootstrap/mcp.go

package bootstrap

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"time"

	"dbview/service/internal/common/databases/meta"
	"dbview/service/internal/common/logger"
	"dbview/service/internal/common/manager/connection"
	"dbview/service/internal/common/manager/storage"
	"dbview/service/internal/common/manager/storage/types"
	"dbview/service/internal/common/manager/task"
	"dbview/service/internal/common/mcp"
	"dbview/service/internal/config"
	dq "dbview/service/internal/modules/data_query/service"

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

// initializeMCP creates and starts the MCP component when enabled in config.
// It registers a few useful tools backed by existing storage and query services.
func initializeMCP(ctx context.Context, cfg *config.AppConfig, log logger.Logger, storageMgr storage.StorageInterface, pool *connection.ConnectionPool, taskMgr *task.Manager) *mcp.Component {
	if cfg == nil || !cfg.MCP.Enable {
		if log != nil {
			log.Info("MCP 组件未启用，跳过初始化", zap.Bool("enabled", cfg != nil && cfg.MCP.Enable))
		}
		return nil
	}

	comp := mcp.NewComponent(mcp.Config{
		Enable:        cfg.MCP.Enable,
		ServerName:    cfg.MCP.ServerName,
		ServerVersion: cfg.MCP.ServerVersion,
	}, log)

	// 基础工具
	comp.RegisterHealthTool()
	comp.RegisterEchoTool()

	// 与当前业务相关的工具
	registerConnectionTools(comp, storageMgr, log)
	registerExecuteSQLTool(comp, pool, taskMgr, log)
	registerAIChatTool(comp, cfg.AI, log)

	go func() {
		if err := comp.ServeStdio(ctx); err != nil && log != nil {
			log.Error("MCP ServeStdio 启动失败", zap.Error(err))
		}
	}()

	if log != nil {
		log.Info("MCP 组件已启用", zap.String("server_name", cfg.MCP.ServerName), zap.String("version", cfg.MCP.ServerVersion))
	}

	return comp
}

// registerConnectionTools exposes a simple connection listing backed by storage.
func registerConnectionTools(comp *mcp.Component, storageMgr storage.StorageInterface, log logger.Logger) {
	if storageMgr == nil {
		if log != nil {
			log.Warn("未找到存储管理器，跳过 MCP 连接工具注册")
		}
		return
	}

	tool := mcpgo.NewTool("connections_list",
		mcpgo.WithDescription("列出已存储的连接（不含敏感字段）"),
	)

	comp.Server().AddTool(tool, func(ctx context.Context, req mcpgo.CallToolRequest) (*mcpgo.CallToolResult, error) {
		conns, err := storageMgr.GetAllConnections()
		if err != nil {
			return nil, err
		}

		summaries := make([]connectionSummary, 0, len(conns))
		for _, c := range conns {
			summaries = append(summaries, summarizeConnection(c))
		}

		payload, err := json.MarshalIndent(summaries, "", "  ")
		if err != nil {
			return nil, err
		}

		return &mcpgo.CallToolResult{Content: []mcpgo.Content{mcpgo.TextContent{Text: string(payload)}}}, nil
	})
}

// registerExecuteSQLTool exposes a synchronous SQL execution tool backed by the connection pool.
func registerExecuteSQLTool(comp *mcp.Component, pool *connection.ConnectionPool, taskMgr *task.Manager, log logger.Logger) {
	if pool == nil {
		if log != nil {
			log.Warn("未找到连接池，跳过 MCP SQL 执行工具注册")
		}
		return
	}

	dataSvc := dq.NewDataService(
		dq.WithConnectionPool(pool),
		dq.WithTaskManager(taskMgr),
	)

	tool := mcpgo.NewTool("execute_sql",
		mcpgo.WithDescription("在指定连接上同步执行 SQL"),
		mcpgo.WithString("connection_id", mcpgo.Description("连接ID，需已在连接池中就绪")),
		mcpgo.WithString("sql", mcpgo.Description("要执行的 SQL 语句")),
	)

	comp.Server().AddTool(tool, func(ctx context.Context, req mcpgo.CallToolRequest) (*mcpgo.CallToolResult, error) {
		args, _ := req.Params.Arguments.(map[string]any)
		connID, _ := args["connection_id"].(string)
		sqlText, _ := args["sql"].(string)

		if connID == "" || sqlText == "" {
			return nil, fmt.Errorf("connection_id 与 sql 不能为空")
		}

		result, err := dataSvc.ExecuteSQL(ctx, connID, "", meta.ObjectPath{}, sqlText, nil, false, false)
		if err != nil {
			return nil, err
		}

		execRes, ok := result.(meta.ExecuteResult)
		if !ok {
			return nil, fmt.Errorf("unexpected execute result type %T", result)
		}

		payload, err := json.MarshalIndent(execRes, "", "  ")
		if err != nil {
			return nil, err
		}

		return &mcpgo.CallToolResult{Content: []mcpgo.Content{mcpgo.TextContent{Text: string(payload)}}}, nil
	})
}

// registerAIChatTool exposes a simple AI chat tool backed by OpenAI HTTP API.
func registerAIChatTool(comp *mcp.Component, aiCfg config.AIConfig, log logger.Logger) {
	if !aiCfg.Enable {
		if log != nil {
			log.Info("AI 交互未启用，跳过 MCP ai_chat 注册")
		}
		return
	}

	if aiCfg.Provider != "openai" {
		if log != nil {
			log.Warn("AI Provider 未支持，跳过 MCP ai_chat 注册", zap.String("provider", aiCfg.Provider))
		}
		return
	}

	apiKeyEnv := aiCfg.APIKeyEnv
	if apiKeyEnv == "" {
		apiKeyEnv = "OPENAI_API_KEY"
	}
	apiKey := os.Getenv(apiKeyEnv)
	if apiKey == "" {
		if log != nil {
			log.Warn("AI 交互已启用但未找到 API Key 环境变量", zap.String("env", apiKeyEnv))
		}
		return
	}

	defaultModel := aiCfg.Model
	if defaultModel == "" {
		defaultModel = "gpt-4o-mini"
	}

	tool := mcpgo.NewTool("ai_chat",
		mcpgo.WithDescription("调用外部 AI 模型生成回复 (OpenAI)"),
		mcpgo.WithString("prompt", mcpgo.Description("用户提示，必填")),
		mcpgo.WithString("system", mcpgo.Description("系统指令，可选")),
		mcpgo.WithString("model", mcpgo.Description("覆盖默认模型，可选")),
	)

	comp.Server().AddTool(tool, func(ctx context.Context, req mcpgo.CallToolRequest) (*mcpgo.CallToolResult, error) {
		args, _ := req.Params.Arguments.(map[string]any)
		prompt, _ := args["prompt"].(string)
		if prompt == "" {
			return nil, fmt.Errorf("prompt 不能为空")
		}
		systemPrompt, _ := args["system"].(string)
		model := defaultModel
		if override, ok := args["model"].(string); ok && override != "" {
			model = override
		}

		callCtx, cancel := context.WithTimeout(ctx, 45*time.Second)
		defer cancel()

		respText, err := callOpenAIChat(callCtx, apiKey, model, systemPrompt, prompt)
		if err != nil {
			return nil, err
		}

		return &mcpgo.CallToolResult{Content: []mcpgo.Content{mcpgo.TextContent{Text: respText}}}, nil
	})

	if log != nil {
		log.Info("MCP ai_chat 工具已注册", zap.String("model", defaultModel), zap.String("provider", aiCfg.Provider))
	}
}

// callOpenAIChat performs a minimal chat completion request.
func callOpenAIChat(ctx context.Context, apiKey, model, systemPrompt, userPrompt string) (string, error) {
	if apiKey == "" {
		return "", fmt.Errorf("missing api key")
	}

	messages := []openAIMessage{}
	if systemPrompt != "" {
		messages = append(messages, openAIMessage{Role: "system", Content: systemPrompt})
	}
	messages = append(messages, openAIMessage{Role: "user", Content: userPrompt})

	body := openAIChatRequest{Model: model, Messages: messages}
	b, err := json.Marshal(body)
	if err != nil {
		return "", err
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://api.openai.com/v1/chat/completions", bytes.NewReader(b))
	if err != nil {
		return "", err
	}
	req.Header.Set("Authorization", "Bearer "+apiKey)
	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	var out openAIChatResponse
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return "", err
	}

	if resp.StatusCode >= 300 {
		if out.Error != nil && out.Error.Message != "" {
			return "", fmt.Errorf("openai error: %s", out.Error.Message)
		}
		return "", fmt.Errorf("openai request failed: status %d", resp.StatusCode)
	}

	if len(out.Choices) == 0 {
		return "", fmt.Errorf("openai 返回为空")
	}

	return out.Choices[0].Message.Content, nil
}

type openAIChatRequest struct {
	Model    string          `json:"model"`
	Messages []openAIMessage `json:"messages"`
}

type openAIMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type openAIChatResponse struct {
	Choices []struct {
		Message openAIMessage `json:"message"`
	} `json:"choices"`
	Error *struct {
		Message string `json:"message"`
	} `json:"error,omitempty"`
}

// connectionSummary holds non-sensitive fields for MCP responses.
type connectionSummary struct {
	ID          string `json:"id"`
	Name        string `json:"name"`
	Description string `json:"description,omitempty"`
	Kind        string `json:"kind"`
	Type        string `json:"type,omitempty"`
	Version     string `json:"version,omitempty"`
	Server      string `json:"server,omitempty"`
	Port        int    `json:"port,omitempty"`
	Database    string `json:"database,omitempty"`
	Color       string `json:"color,omitempty"`
	AutoConnect bool   `json:"auto_connect"`
}

func summarizeConnection(c types.ConnectionWithDetails) connectionSummary {
	summary := connectionSummary{
		ID:          c.ID,
		Name:        c.Name,
		Description: c.Description,
		Kind:        c.Kind,
		Color:       c.Color,
		AutoConnect: c.AutoConnect,
	}

	if c.DBDetail != nil {
		summary.Type = c.DBDetail.Type
		summary.Version = c.DBDetail.Version
		summary.Server = c.DBDetail.Server
		summary.Port = c.DBDetail.Port
		summary.Database = c.DBDetail.DatabaseName
	}

	if c.ServerDetail != nil {
		summary.Type = c.ServerDetail.Type
		summary.Version = c.ServerDetail.Version
		summary.Server = c.ServerDetail.Server
		summary.Port = c.ServerDetail.Port
	}

	return summary
}