chat-deploy/internal/api/bot/start.go

package bot

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	chatmw "git.imall.cloud/openim/chat/internal/api/mw"
	"git.imall.cloud/openim/chat/internal/api/util"
	"git.imall.cloud/openim/chat/pkg/common/config"
	"git.imall.cloud/openim/chat/pkg/common/kdisc"
	disetcd "git.imall.cloud/openim/chat/pkg/common/kdisc/etcd"
	adminclient "git.imall.cloud/openim/chat/pkg/protocol/admin"
	botclient "git.imall.cloud/openim/chat/pkg/protocol/bot"
	"github.com/gin-gonic/gin"
	"github.com/openimsdk/tools/discovery/etcd"
	"github.com/openimsdk/tools/errs"
	"github.com/openimsdk/tools/mw"
	"github.com/openimsdk/tools/system/program"
	"github.com/openimsdk/tools/utils/datautil"
	"github.com/openimsdk/tools/utils/runtimeenv"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"strings"
)

type Config struct {
	ApiConfig config.APIBot
	Discovery config.Discovery
	Share     config.Share
	Redis     config.Redis

	RuntimeEnv string
}

func Start(ctx context.Context, index int, cfg *Config) error {
	cfg.RuntimeEnv = runtimeenv.PrintRuntimeEnvironment()
	apiPort, err := datautil.GetElemByIndex(cfg.ApiConfig.Api.Ports, index)
	if err != nil {
		return err
	}
	client, err := kdisc.NewDiscoveryRegister(&cfg.Discovery, cfg.RuntimeEnv, nil)
	if err != nil {
		return err
	}

	botConn, err := client.GetConn(ctx, cfg.Discovery.RpcService.Bot, grpc.WithTransportCredentials(insecure.NewCredentials()), mw.GrpcClient())
	if err != nil {
		return err
	}
	adminConn, err := client.GetConn(ctx, cfg.Discovery.RpcService.Admin, grpc.WithTransportCredentials(insecure.NewCredentials()), mw.GrpcClient())
	if err != nil {
		return err
	}
	adminClient := adminclient.NewAdminClient(adminConn)
	botClient := botclient.NewBotClient(botConn)
	base := util.Api{
		ImUserID:        cfg.Share.OpenIM.AdminUserID,
		ProxyHeader:     cfg.Share.ProxyHeader,
		ChatAdminUserID: cfg.Share.ChatAdmin[0],
	}
	botApi := New(botClient, &base)
	mwApi := chatmw.New(adminClient)
	gin.SetMode(gin.ReleaseMode)
	engine := gin.New()
	engine.Use(gin.Recovery(), mw.CorsHandler(), mw.GinParseOperationID(), func(c *gin.Context) {
		// 确保 operationID 被正确设置到 context 中
		operationID := c.GetHeader("operationid")
		if operationID != "" {
			c.Set("operationID", operationID)
		}
		c.Next()
	})
	
	// 可选的 Prometheus metrics 端点
	// 支持配置文件或环境变量控制（环境变量优先）
	prometheusEnable := cfg.ApiConfig.Prometheus.Enable
	if envEnable := os.Getenv("PROMETHEUS_ENABLE"); envEnable != "" {
		prometheusEnable = strings.ToLower(envEnable) == "true" || envEnable == "1"
	}
	if prometheusEnable {
		metricsPath := cfg.ApiConfig.Prometheus.Path
		if envPath := os.Getenv("PROMETHEUS_PATH"); envPath != "" {
			metricsPath = envPath
		}
		if metricsPath == "" {
			metricsPath = "/metrics"
		}
		engine.GET(metricsPath, gin.WrapH(promhttp.Handler()))
	}

	SetBotRoute(engine, botApi, mwApi)

	var (
		netDone = make(chan struct{}, 1)
		netErr  error
	)
	server := http.Server{Addr: fmt.Sprintf(":%d", apiPort), Handler: engine}
	go func() {
		err = server.ListenAndServe()
		if err != nil && !errors.Is(err, http.ErrServerClosed) {
			netErr = errs.WrapMsg(err, fmt.Sprintf("api start err: %s", server.Addr))
			netDone <- struct{}{}
		}
	}()
	if cfg.Discovery.Enable == kdisc.ETCDCONST {
		cm := disetcd.NewConfigManager(client.(*etcd.SvcDiscoveryRegistryImpl).GetClient(),
			[]string{
				config.ChatAPIBotCfgFileName,
				config.DiscoveryConfigFileName,
				config.ShareFileName,
				config.LogConfigFileName,
			},
		)
		cm.Watch(ctx)
	}
	shutdown := func() error {
		ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
		defer cancel()
		err := server.Shutdown(ctx)
		if err != nil {
			return errs.WrapMsg(err, "shutdown err")
		}
		return nil
	}
	disetcd.RegisterShutDown(shutdown)

	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGTERM)
	select {
	case <-sigs:
		program.SIGTERMExit()
		if err := shutdown(); err != nil {
			return err
		}
	case <-netDone:
		close(netDone)
		return netErr
	}
	return nil
}

func SetBotRoute(router gin.IRouter, bot *Api, mw *chatmw.MW) {
	account := router.Group("/agent")
	account.POST("/create", mw.CheckAdmin, bot.CreateAgent)
	account.POST("/delete", mw.CheckAdmin, bot.DeleteAgent)
	account.POST("/update", mw.CheckAdmin, bot.UpdateAgent)
	account.POST("/page", mw.CheckToken, bot.PageFindAgent)

	imwebhook := router.Group("/im_callback")
	imwebhook.POST("/callbackAfterSendSingleMsgCommand", bot.AfterSendSingleMsg)
	imwebhook.POST("/callbackAfterSendGroupMsgCommand", bot.AfterSendGroupMsg)
}