ai-agent/internal/bootstrap/server.go

package bootstrap

import (
	"log/slog"
	"net/http"
	"strings"
	"time"

	_ "code.tczkiot.com/wlw/ai-agent/internal/ai/runtime"
	"code.tczkiot.com/wlw/ai-agent/internal/handlers/api"
	"code.tczkiot.com/wlw/ai-agent/internal/middleware"
	"code.tczkiot.com/wlw/ai-agent/internal/pkg/config"
	"code.tczkiot.com/wlw/ai-agent/internal/pkg/ginx"
	"code.tczkiot.com/wlw/ai-agent/internal/pkg/httpx"
	"code.tczkiot.com/wlw/ai-agent/internal/pkg/i18nx"
	"code.tczkiot.com/wlw/ai-agent/internal/pkg/tracex"
	"code.tczkiot.com/wlw/ai-agent/internal/services"
	"code.tczkiot.com/wlw/ai-agent/internal/web/supportchat"

	"github.com/gin-gonic/gin"
	"github.com/mlogclub/simple/web"

	_ "code.tczkiot.com/wlw/ai-agent/internal/services/wx_callback_handlers"
)

func NewServer() (*gin.Engine, error) {
	cfg := config.Current()
	i18nx.SetDefaultLocale(cfg.LanguageOrDefault())

	gin.SetMode(gin.ReleaseMode)
	printBanner()

	app := gin.New()
	app.Use(requestIDMiddleware())
	app.Use(corsMiddleware())
	app.Use(gin.Recovery())
	app.Use(requestLogMiddleware())
	app.Use(maxBodySizeMiddleware())
	app.Use(i18nx.Middleware())

	addRouter(app)

	notFoundPrefixes := []string{"/api/"}
	if baseURL := strings.TrimRight(cfg.Storage.Local.BaseURL, "/"); baseURL != "" {
		notFoundPrefixes = append(notFoundPrefixes, baseURL+"/")
	}
	app.StaticFS(cfg.Storage.Local.BaseURL, ginx.StaticFiles(cfg.Storage.Local.Root))
	app.NoRoute(func(ctx *gin.Context) {
		for _, prefix := range notFoundPrefixes {
			if strings.HasPrefix(ctx.Request.URL.Path, prefix) {
				httpx.WriteHttpStatusJSON(ctx, http.StatusNotFound, web.JsonErrorCode(http.StatusNotFound, i18nx.T(ctx, "error.notFound")))
				return
			}
		}
		httpx.WriteHttpStatusJSON(ctx, http.StatusNotFound, web.JsonErrorCode(http.StatusNotFound, i18nx.T(ctx, "error.notFound")))
	})

	return app, nil
}

func corsMiddleware() gin.HandlerFunc {
	allowedOrigins := config.Current().Server.CORS.AllowedOrigins
	allowHeaders := "Origin, Content-Type, Accept, Authorization, X-Requested-With, X-Channel-Id, X-External-Id, X-External-Name"
	exposeHeaders := "Content-Length, Content-Type, Authorization"
	allowMethods := "GET, POST, PUT, PATCH, DELETE, OPTIONS"
	allowedOriginSet := make(map[string]struct{}, len(allowedOrigins))
	for _, origin := range allowedOrigins {
		origin = strings.TrimRight(strings.TrimSpace(origin), "/")
		if origin == "" {
			continue
		}
		allowedOriginSet[origin] = struct{}{}
	}
	return func(ctx *gin.Context) {
		if isWebsocketUpgrade(ctx) {
			ctx.Next()
			return
		}
		origin := strings.TrimRight(strings.TrimSpace(ctx.GetHeader("Origin")), "/")
		if origin != "" {
			ctx.Header("Vary", "Origin")
			if _, ok := allowedOriginSet[origin]; !ok {
				if ctx.Request.Method == http.MethodOptions {
					ctx.AbortWithStatus(http.StatusForbidden)
					return
				}
				ctx.Next()
				return
			}
			ctx.Header("Access-Control-Allow-Origin", origin)
			ctx.Header("Access-Control-Allow-Methods", allowMethods)
			ctx.Header("Access-Control-Allow-Headers", allowHeaders)
			ctx.Header("Access-Control-Expose-Headers", exposeHeaders)
			ctx.Header("Access-Control-Max-Age", "600")
		}
		if ctx.Request.Method == http.MethodOptions {
			ctx.AbortWithStatus(http.StatusNoContent)
			return
		}
		ctx.Next()
	}
}

func requestIDMiddleware() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		requestID := tracex.EnsureRequestID(ctx.GetHeader(tracex.RequestIDHeader))
		ctx.Set(tracex.GinRequestIDKey, requestID)
		if requestID != "" {
			ctx.Header(tracex.RequestIDHeader, requestID)
		}
		ctx.Next()
	}
}

func requestLogMiddleware() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		start := time.Now()
		path := ctx.Request.URL.Path
		method := ctx.Request.Method
		requestID, _ := ctx.Get(tracex.GinRequestIDKey)
		ctx.Next()

		slog.Info("http request",
			"requestId", requestID,
			"method", method,
			"path", path,
			"status", ctx.Writer.Status(),
			"elapsed", time.Since(start).Milliseconds(),
			"clientIp", ctx.ClientIP(),
		)
	}
}

func maxBodySizeMiddleware() gin.HandlerFunc {
	limit := config.Current().Storage.MaxRequestBodySizeBytes()
	return func(ctx *gin.Context) {
		ctx.Request.Body = http.MaxBytesReader(ctx.Writer, ctx.Request.Body, limit)
		ctx.Next()
	}
}

func isWebsocketUpgrade(ctx *gin.Context) bool {
	if !strings.EqualFold(ctx.GetHeader("Upgrade"), "websocket") {
		return false
	}
	return strings.Contains(strings.ToLower(ctx.GetHeader("Connection")), "upgrade")
}

func addRouter(app *gin.Engine) {
	supportchat.RegisterRoutes(app)

	apiGroup := app.Group("/api")
	apiGroup.GET("/health", api.Health)
	apiGroup.GET("/config", api.PublicConfig)
	registerApiChannelRoutes(apiGroup.Group("/channel"))
	registerApiConversationRoutes(apiGroup.Group("/conversation", middleware.ExternalUserMiddleware))
	registerApiMessageRoutes(apiGroup.Group("/message", middleware.ExternalUserMiddleware))

	wsGroup := app.Group("/api/ws")
	wsGroup.GET("/dashboard", middleware.AuthMiddleware, services.WsService.HandleDashboardWS)
	wsGroup.GET("/dashboard/notification", middleware.AuthMiddleware, services.WsService.HandleDashboardNotificationWS)
	wsGroup.GET("/open", services.WsService.HandleOpenWS)

	dashboardGroup := app.Group("/api/dashboard", middleware.AuthMiddleware)
	registerDashboardDashboardRoutes(dashboardGroup.Group("/dashboard"))
	registerDashboardConversationRoutes(dashboardGroup.Group("/conversation"))
	registerDashboardNotificationRoutes(dashboardGroup.Group("/notification"))
	registerDashboardQuickReplyRoutes(dashboardGroup.Group("/quick-reply"))
	registerDashboardChannelRoutes(dashboardGroup.Group("/channel"))
	registerDashboardAgentRoutes(dashboardGroup.Group("/agent"))
	registerDashboardAgentTeamRoutes(dashboardGroup.Group("/agent-team"))
	registerDashboardAgentTeamScheduleRoutes(dashboardGroup.Group("/agent-team-schedule"))
	registerDashboardAIAgentRoutes(dashboardGroup.Group("/ai-agent"))
	registerDashboardAgentRunRoutes(dashboardGroup.Group("/agent-run"))
	registerDashboardAIConfigRoutes(dashboardGroup.Group("/ai-config"))
	registerDashboardPlatformAIRoutes(dashboardGroup.Group("/platform-ai"))
	registerDashboardAssetRoutes(dashboardGroup.Group("/asset"))
	registerDashboardKnowledgeBaseRoutes(dashboardGroup.Group("/knowledge-base"))
	registerDashboardKnowledgeDirectoryRoutes(dashboardGroup.Group("/knowledge-directory"))
	registerDashboardKnowledgeDocumentRoutes(dashboardGroup.Group("/knowledge-document"))
	registerDashboardKnowledgeFAQRoutes(dashboardGroup.Group("/knowledge-faq"))
	registerDashboardKnowledgeRetrieveRoutes(dashboardGroup.Group("/knowledge-retrieve"))
	registerDashboardKnowledgeRetrieveLogRoutes(dashboardGroup.Group("/knowledge-retrieve-log"))

	thirdGroup := app.Group("/api/third")
	registerThirdWechatRoutes(thirdGroup.Group("/wechat"))
}