ttstd-infinite-canvas/backend/internal/handler/comfy_bridge.go

package handler

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

	"infinite-canvas/backend/internal/model"
	"infinite-canvas/backend/internal/service"

	"github.com/gin-gonic/gin"
)

// Bridge 结果通常包含 base64 媒体；保持和 service 层相同的硬上限，避免请求体无限增长。
const maxComfyBridgeRequestBytes int64 = 64 << 20

// Bridge 心跳会携带本机工作流字段和拓扑，允许与 service 层一致的 4MB 上限。
const maxComfyBridgeCapabilitiesBytes int64 = 4 << 20

type comfyBridgeResultRequest struct {
	RequestID string         `json:"requestId"`
	ID        string         `json:"id,omitempty"`
	Status    string         `json:"status"`
	Result    map[string]any `json:"result,omitempty"`
	Error     string         `json:"error,omitempty"`
}

// RegisterComfyBridgeRoutes 同时注册用户管理接口和 Bridge 专用轮询接口。
// 两套接口使用不同认证方式：用户接口依赖登录 Cookie，Bridge 接口只接受专用 Header Token。
func RegisterComfyBridgeRoutes(r *gin.RouterGroup, svc *service.Service) {
	r.POST("/comfy-bridges", func(c *gin.Context) {
		user, err := currentUser(c, svc)
		if err != nil {
			failService(c, err)
			return
		}
		if err := svc.RequireWorkflowPluginForUser(user.ID, "comfyui-bridge-image"); err != nil {
			failService(c, err)
			return
		}
		c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 64<<10)
		var req service.CreateComfyBridgeRequest
		if err := c.ShouldBindJSON(&req); err != nil {
			fail(c, http.StatusBadRequest, err)
			return
		}
		result, err := svc.CreateComfyBridge(user.ID, req)
		if err != nil {
			failService(c, err)
			return
		}
		ok(c, result)
	})
	r.GET("/comfy-bridges", func(c *gin.Context) {
		user, err := currentUser(c, svc)
		if err != nil {
			failService(c, err)
			return
		}
		if err := svc.RequireWorkflowPluginForUser(user.ID, "comfyui-bridge-image"); err != nil {
			failService(c, err)
			return
		}
		bridges, err := svc.ComfyBridges(user.ID)
		if err != nil {
			fail(c, http.StatusInternalServerError, err)
			return
		}
		ok(c, bridges)
	})
	revoke := func(c *gin.Context) {
		user, err := currentUser(c, svc)
		if err != nil {
			failService(c, err)
			return
		}
		if err := svc.RevokeComfyBridge(user.ID, c.Param("id")); err != nil {
			failService(c, err)
			return
		}
		ok(c, gin.H{"revoked": true})
	}
	r.DELETE("/comfy-bridges/:id", revoke)
	r.POST("/comfy-bridges/:id/revoke", revoke)

	r.GET("/comfy-bridge/poll", func(c *gin.Context) {
		bridge, err := authenticateComfyBridgeRequest(c, svc)
		if err != nil {
			failService(c, err)
			return
		}
		wait, err := comfyBridgeWait(c.Query("wait"))
		if err != nil {
			fail(c, http.StatusBadRequest, err)
			return
		}
		request, err := svc.PollComfyBridgeRequest(c.Request.Context(), bridge, wait)
		if err != nil {
			if errors.Is(err, context.Canceled) {
				return
			}
			failService(c, err)
			return
		}
		ok(c, gin.H{"request": request})
	})
	r.POST("/comfy-bridge/result", func(c *gin.Context) {
		bridge, err := authenticateComfyBridgeRequest(c, svc)
		if err != nil {
			failService(c, err)
			return
		}
		if c.Request.ContentLength > maxComfyBridgeRequestBytes {
			fail(c, http.StatusRequestEntityTooLarge, errors.New("Bridge 结果请求体超过 64MB"))
			return
		}
		c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, maxComfyBridgeRequestBytes)
		var req comfyBridgeResultRequest
		if err := c.ShouldBindJSON(&req); err != nil {
			var maxBytesErr *http.MaxBytesError
			if errors.As(err, &maxBytesErr) {
				fail(c, http.StatusRequestEntityTooLarge, errors.New("Bridge 结果请求体超过 64MB"))
				return
			}
			fail(c, http.StatusBadRequest, err)
			return
		}
		requestID := strings.TrimSpace(req.RequestID)
		if requestID == "" {
			requestID = strings.TrimSpace(req.ID)
		}
		err = svc.CompleteComfyBridgeRequest(bridge.ID, service.ComfyBridgeCompletion{
			RequestID: requestID,
			Status:    req.Status,
			Result:    req.Result,
			Error:     req.Error,
		})
		if err != nil {
			failService(c, err)
			return
		}
		ok(c, gin.H{"accepted": true})
	})
	r.POST("/comfy-bridge/heartbeat", func(c *gin.Context) {
		bridge, err := authenticateComfyBridgeRequest(c, svc)
		if err != nil {
			failService(c, err)
			return
		}
		var req struct {
			Capabilities map[string]any `json:"capabilities,omitempty"`
		}
		c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, maxComfyBridgeCapabilitiesBytes)
		if c.Request.ContentLength != 0 {
			if err := c.ShouldBindJSON(&req); err != nil && !errors.Is(err, io.EOF) {
				fail(c, http.StatusBadRequest, err)
				return
			}
		}
		if err := svc.TouchComfyBridgeHeartbeat(bridge.ID, req.Capabilities); err != nil {
			failService(c, err)
			return
		}
		ok(c, gin.H{"ok": true})
	})
}

func authenticateComfyBridgeRequest(c *gin.Context, svc *service.Service) (*model.ComfyBridge, error) {
	token := strings.TrimSpace(c.GetHeader("X-Canvas-Comfy-Bridge-Token"))
	if token == "" {
		token = strings.TrimSpace(c.GetHeader("X-Canvas-Bridge-Token"))
	}
	if token == "" {
		authorization := strings.TrimSpace(c.GetHeader("Authorization"))
		if len(authorization) >= 7 && strings.EqualFold(authorization[:7], "Bearer ") {
			token = strings.TrimSpace(authorization[7:])
		}
	}
	return svc.AuthenticateComfyBridge(token)
}

func comfyBridgeWait(raw string) (time.Duration, error) {
	if strings.TrimSpace(raw) == "" {
		return 20 * time.Second, nil
	}
	seconds, err := strconv.Atoi(strings.TrimSpace(raw))
	if err != nil || seconds < 0 || seconds > 60 {
		return 0, errors.New("wait 必须是 0 到 60 之间的整数秒")
	}
	return time.Duration(seconds) * time.Second, nil
}