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) } } }