MonocoEditor-With-Lsp-Backend/backend/internal/api/handler.go

package api

import (
	"context"
	"errors"
	"net/http"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"go.uber.org/zap"

	"monica-go-completion-backend/internal/completion"
)

type CompletionService interface {
	Complete(ctx context.Context, req completion.Request) (completion.Response, error)
}

// SessionStatsProvider 暴露会话统计信息，供就绪探针输出。
type SessionStatsProvider interface {
	ActiveSessions() map[string]int
}

// RouteOptions 控制 HTTP/WS 接口的超时与请求体上限。
type RouteOptions struct {
	RequestTimeout time.Duration // 单次补全调用超时时间。
	MaxBodyBytes   int64         // 请求体最大字节数（HTTP/WS 共用）。
}

// RegisterRoutes 注册健康检查、HTTP 补全接口和 WebSocket 补全接口。
func RegisterRoutes(router *gin.Engine, service CompletionService, options ...RouteOptions) {
	opts := RouteOptions{
		RequestTimeout: 10 * time.Second,
		MaxBodyBytes:   2 << 20, // 2MB
	}
	if len(options) > 0 {
		if options[0].RequestTimeout > 0 {
			opts.RequestTimeout = options[0].RequestTimeout
		}
		if options[0].MaxBodyBytes > 0 {
			opts.MaxBodyBytes = options[0].MaxBodyBytes
		}
	}

	router.GET("/health", func(c *gin.Context) {
		c.JSON(http.StatusOK, gin.H{"status": "ok", "service": "lsp-gateway"})
	})

	router.GET("/health/live", func(c *gin.Context) {
		c.JSON(http.StatusOK, gin.H{"status": "alive"})
	})

	router.GET("/health/ready", func(c *gin.Context) {
		sessions := map[string]int{}
		if provider, ok := service.(SessionStatsProvider); ok {
			sessions = provider.ActiveSessions()
		}
		c.JSON(http.StatusOK, gin.H{
			"status":   "ready",
			"sessions": sessions,
		})
	})

	registerWSRoutes(router, service, opts)

	handleCompletion := func(c *gin.Context) {
		// 为单次请求限制 body 大小，避免异常大包占满内存。
		c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, opts.MaxBodyBytes)
		var req completion.Request
		if err := c.ShouldBindJSON(&req); err != nil {
			zap.L().Warn(
				"invalid completion json payload",
				zap.String("traceId", traceIDFromGin(c)),
				zap.String("routeLanguage", c.Param("language")),
				zap.String("clientIP", c.ClientIP()),
				zap.Error(err),
			)
			c.JSON(http.StatusBadRequest, gin.H{"error": "invalid JSON payload"})
			return
		}

		routeLang := c.Param("language")
		// 若 body 未显式给出 language，则使用路由参数。
		if req.Language == "" {
			req.Language = routeLang
		}

		// 对下游补全调用增加超时保护，防止请求长时间悬挂。
		ctx, cancel := context.WithTimeout(c.Request.Context(), opts.RequestTimeout)
		defer cancel()

		resp, err := service.Complete(ctx, req)
		if err != nil {
			if errors.Is(err, completion.ErrInvalidRequest) {
				zap.L().Warn(
					"completion request invalid",
					zap.String("traceId", traceIDFromGin(c)),
					zap.String("clientIP", c.ClientIP()),
					zap.String("language", req.Language),
					zap.String("uri", req.URI),
					zap.Error(err),
				)
				c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
				return
			}
			if errors.Is(err, completion.ErrUnsupportedLanguage) {
				zap.L().Warn(
					"completion language not supported",
					zap.String("traceId", traceIDFromGin(c)),
					zap.String("clientIP", c.ClientIP()),
					zap.String("language", req.Language),
					zap.Error(err),
				)
				c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
				return
			}
			if errors.Is(err, completion.ErrTooManySessions) {
				zap.L().Warn(
					"completion rejected due to session limit",
					zap.String("traceId", traceIDFromGin(c)),
					zap.String("clientIP", c.ClientIP()),
					zap.String("language", req.Language),
					zap.String("sessionId", req.SessionID),
					zap.Error(err),
				)
				c.JSON(http.StatusTooManyRequests, gin.H{"error": err.Error()})
				return
			}
			var ownedErr *completion.ErrSessionOwnedByOtherInstance
			if errors.As(err, &ownedErr) {
				zap.L().Info(
					"completion routed to session owner",
					zap.String("traceId", traceIDFromGin(c)),
					zap.String("clientIP", c.ClientIP()),
					zap.String("language", req.Language),
					zap.String("sessionId", req.SessionID),
					zap.String("ownerId", ownedErr.OwnerID),
					zap.String("routeTo", ownedErr.OwnerEndpoint),
				)
				if ownedErr.OwnerEndpoint != "" {
					c.Header("X-LSP-Route-To", ownedErr.OwnerEndpoint)
				}
				c.JSON(http.StatusConflict, gin.H{
					"error":   err.Error(),
					"routeTo": ownedErr.OwnerEndpoint,
					"ownerId": ownedErr.OwnerID,
				})
				return
			}
			zap.L().Error(
				"completion failed",
				zap.String("traceId", traceIDFromGin(c)),
				zap.String("clientIP", c.ClientIP()),
				zap.String("language", req.Language),
				zap.String("sessionId", req.SessionID),
				zap.String("uri", req.URI),
				zap.Error(err),
			)
			c.JSON(http.StatusInternalServerError, gin.H{"error": "completion failed"})
			return
		}
		zap.L().Info(
			"completion success",
			zap.String("traceId", traceIDFromGin(c)),
			zap.String("clientIP", c.ClientIP()),
			zap.String("language", req.Language),
			zap.String("sessionId", req.SessionID),
			zap.String("uri", req.URI),
			zap.Int("items", len(resp.Items)),
			zap.Bool("isIncomplete", resp.IsIncomplete),
		)

		c.JSON(http.StatusOK, resp)
	}

	router.POST("/api/v1/completions/:language", func(c *gin.Context) {
		handleCompletion(c)
	})
}

func traceIDFromGin(c *gin.Context) string {
	if c == nil {
		return ""
	}
	traceID := strings.TrimSpace(c.GetString("traceId"))
	if traceID != "" {
		return traceID
	}
	return strings.TrimSpace(c.GetString("requestId"))
}