hunjing/mcp/web-search/cmd/server/main.go

package main

import (
	"flag"
	"fmt"
	"log/slog"
	"net/http"
	"os"

	"github.com/prometheus/client_golang/prometheus/promhttp"

	"github.com/chenxin/hunjing/mcp/web-search/internal/balancer"
	"github.com/chenxin/hunjing/mcp/web-search/internal/cache"
	"github.com/chenxin/hunjing/mcp/web-search/internal/config"
	"github.com/chenxin/hunjing/mcp/web-search/internal/logger"
	"github.com/chenxin/hunjing/mcp/web-search/internal/mcp"
	"github.com/chenxin/hunjing/mcp/web-search/internal/metrics"
	"github.com/chenxin/hunjing/mcp/web-search/internal/models"
	"github.com/chenxin/hunjing/mcp/web-search/internal/providers"
)

func main() {
	configPath := flag.String("config", "", "YAML 配置文件路径（不传则按项目惯例自动探测）")
	port := flag.Int("port", 0, "HTTP SSE 模式监听端口（缺省时使用配置文件的 port）")
	stdioMode := flag.Bool("stdio", false, "以 stdio 模式运行（通过 stdin/stdout 走 JSON-RPC，用于子进程集成）")
	flag.Parse()

	// 未指定 --config 时按项目惯例自动探测
	if *configPath == "" {
		for _, p := range []string{
			"../../config/mcp/web-search.yaml", // 开发模式（mcp/web-search/ 下 go run）
			"config/mcp/web-search.yaml",       // 项目根目录下 go run mcp/... 或 make
			"/etc/bowen/config.yaml",           // Docker 容器内默认路径
		} {
			if _, err := os.Stat(p); err == nil {
				*configPath = p
				break
			}
		}
	}

	// 加载配置
	cfg, err := config.Load(*configPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "配置加载失败: %v\n", err)
		os.Exit(1)
	}

	// 设置日志
	logCfg := logger.Config{
		Level:     slog.LevelInfo,
		AddSource: false,
		JSON:      os.Getenv("ENV") == "production",
	}
	slog.SetDefault(logger.New(logCfg))

	// 端口覆盖
	if *port > 0 {
		cfg.Server.Port = *port
	}

	// 初始化 Provider 注册中心
	reg, err := providers.SetupFromConfig(cfg)
	if err != nil {
		slog.Error("provider 初始化失败", "error", err)
		os.Exit(1)
	}

	// 初始化负载均衡器
	lb := balancer.NewLoadBalancer(&cfg.Balancer)

	// 为每个启用的 provider 及其密钥注册到负载均衡器
	type providerEntry struct {
		kind models.ProviderKind
		keys []string
	}
	entries := []providerEntry{
		{models.ProviderTavily, cfg.Providers.Tavily.APIKeys},
		{models.ProviderBrave, cfg.Providers.Brave.APIKeys},
	}

	var providerKinds []models.ProviderKind
	for _, e := range entries {
		if len(e.keys) == 0 {
			continue
		}
		prov, ok := reg.Get(e.kind)
		if !ok {
			continue
		}
		lb.AddProvider(e.kind, e.keys, 1, prov)
		providerKinds = append(providerKinds, e.kind)
		metrics.ActiveKeys.WithLabelValues(string(e.kind)).Set(float64(len(e.keys)))
		slog.Info("provider 已注册",
			"provider", e.kind,
			"active_keys", len(e.keys),
		)
	}

	// 初始化缓存
	cacheMgr, err := cache.New(cfg.Cache)
	if err != nil {
		slog.Error("缓存初始化失败", "error", err)
		os.Exit(1)
	}

	// 创建 handler
	handler := mcp.NewHandler(lb, cacheMgr, providerKinds)

	// 启动后台健康检查（使用第一个启用的 provider 的配置参数）
	if len(providerKinds) > 0 && cfg.Providers.Tavily.HealthCheck.Interval > 0 {
		checker := balancer.NewHealthChecker(lb, cfg.Providers.Tavily.HealthCheck)
		checker.Start()
		defer checker.Stop()
	}

	// 配置热重载（仅在指定了配置文件时启用）
	if *configPath != "" {
		onChange := func(old, new *config.Config) error {
			handler.ReloadConfig(old, new)
			return nil
		}
		reloader, err := config.NewReloader(*configPath, cfg, onChange)
		if err != nil {
			slog.Warn("配置热重载初始化失败（不影响启动）", "error", err)
		} else {
			defer reloader.Stop()
		}
	}

	// 启动 Metrics 服务
	if cfg.Metrics.Enabled {
		go func() {
			metricsMux := http.NewServeMux()
			metricsMux.Handle(cfg.Metrics.Path, promhttp.Handler())
			metricsMux.HandleFunc("/", handler.HandleMonitor)
			metricsMux.HandleFunc("/health", handler.HandleHealth)
			metricsMuxWithLog := mcp.LoggingMiddleware(metricsMux)
			metricsAddr := fmt.Sprintf(":%d", cfg.Metrics.Port)
			slog.Info("metrics 服务启动", "addr", metricsAddr)
			if err := http.ListenAndServe(metricsAddr, metricsMuxWithLog); err != nil {
				slog.Error("metrics 服务异常", "error", err)
			}
		}()
	}

	// 根据 CLI 标志选择传输模式：
	//   --stdio ：stdin/stdout JSON-RPC（用于 Claude Desktop / Cursor 等 MCP Client 子进程集成）
	//   默认    ：Streamable HTTP 模式（用于灵台注册 / 远程 MCP Client 直连）
	if *stdioMode {
		handler.ServeStdio()
	} else {
		addr := fmt.Sprintf(":%d", cfg.Server.Port)
		slog.Info("博闻 HTTP 服务启动", "addr", addr)
		if err := handler.ServeHTTP(addr); err != nil {
			slog.Error("博闻异常退出", "error", err)
			os.Exit(1)
		}
	}
}