wonder
bb5d3431c1
- 在 Task 模型中添加 ExternalID 字段用于关联内存中的 taskID - 修改 Generate 函数，将任务记录同时存入内存和数据库 - 添加 saveTaskToDB 和 updateTaskInDB 函数处理任务持久化 - 同步更新任务状态到数据库，确保 GetProjectTasks 能正确查询 - 处理 default 工程的特殊情况（跳过数据库存储）
291 lines
7.5 KiB
Go
Executable File
291 lines
7.5 KiB
Go
Executable File
package handler
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"strconv"
|
|
"sync"
|
|
"time"
|
|
|
|
"gen2d/internal/db"
|
|
"gen2d/internal/logger"
|
|
"gen2d/internal/model"
|
|
"gen2d/internal/service"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
// GenerateRequest 素材生成请求。
|
|
type GenerateRequest struct {
|
|
ProjectID string `json:"projectId"`
|
|
Prompt string `json:"prompt"`
|
|
AssetType string `json:"assetType" binding:"required"`
|
|
Tags []string `json:"tags"`
|
|
UserNote string `json:"userNote"`
|
|
ProjectStyle map[string]string `json:"projectStyle"`
|
|
TaskStyle map[string]string `json:"taskStyle"`
|
|
Resolution int `json:"resolution"`
|
|
Directions int `json:"directions"`
|
|
FramesPerDir int `json:"framesPerDir"`
|
|
Format string `json:"format"`
|
|
}
|
|
|
|
// GenerateResponse 素材生成响应体。
|
|
type GenerateResponse struct {
|
|
TaskID string `json:"taskId"`
|
|
}
|
|
|
|
// AssetResponse 单个素材响应。
|
|
type AssetResponse struct {
|
|
Key string `json:"key"`
|
|
URL string `json:"url"`
|
|
Format string `json:"format"`
|
|
}
|
|
|
|
// TaskResponse 任务查询响应。
|
|
type TaskResponse struct {
|
|
ID string `json:"id"`
|
|
ProjectID string `json:"projectId"`
|
|
Prompt string `json:"prompt"`
|
|
AssetType string `json:"assetType"`
|
|
Status string `json:"status"`
|
|
Stage string `json:"stage,omitempty"`
|
|
Progress int `json:"progress"`
|
|
RetryCount int `json:"retryCount"`
|
|
Error string `json:"error,omitempty"`
|
|
CreatedAt string `json:"createdAt"`
|
|
}
|
|
|
|
// AssetsResponse 素材列表响应。
|
|
type AssetsResponse struct {
|
|
Assets []AssetResponse `json:"assets"`
|
|
Metadata service.AssetMetadata `json:"metadata"`
|
|
}
|
|
|
|
// taskRecord 内存中的任务记录。
|
|
type taskRecord struct {
|
|
task TaskResponse
|
|
assets []AssetResponse
|
|
metadata service.AssetMetadata
|
|
}
|
|
|
|
var (
|
|
taskStore = sync.Map{} // taskID → *taskRecord
|
|
)
|
|
|
|
// Generate 素材生成接口（异步）。
|
|
// 立即返回 taskId，后台执行管线，前端通过 GET /tasks/:taskId 轮询进度。
|
|
func Generate(c *gin.Context) {
|
|
var req GenerateRequest
|
|
if err := c.ShouldBindJSON(&req); err != nil {
|
|
c.JSON(http.StatusBadRequest, model.Fail(http.StatusBadRequest, "参数错误: "+err.Error()))
|
|
return
|
|
}
|
|
|
|
projectID := req.ProjectID
|
|
if projectID == "" {
|
|
projectID = "default"
|
|
}
|
|
taskID := fmt.Sprintf("task-%d", time.Now().UnixMilli())
|
|
createdAt := time.Now().Format(time.RFC3339)
|
|
|
|
// 存入 pending 状态
|
|
taskStore.Store(taskID, &taskRecord{
|
|
task: TaskResponse{
|
|
ID: taskID,
|
|
ProjectID: projectID,
|
|
Prompt: req.Prompt,
|
|
AssetType: req.AssetType,
|
|
Status: "pending",
|
|
Progress: 0,
|
|
CreatedAt: createdAt,
|
|
},
|
|
})
|
|
|
|
// 将任务记录写入数据库（用于任务列表查询）
|
|
go saveTaskToDB(c.Request.Context(), projectID, taskID, req)
|
|
|
|
// 返回 taskId
|
|
c.JSON(http.StatusOK, model.OK(GenerateResponse{TaskID: taskID}))
|
|
|
|
// 后台执行管线
|
|
go runPipelineBg(projectID, taskID, req)
|
|
}
|
|
|
|
// runPipelineBg 后台执行生成管线，更新任务状态。
|
|
func runPipelineBg(projectID, taskID string, req GenerateRequest) {
|
|
// 注入进度上报回调
|
|
ctx := service.WithProgressReporter(context.Background(), func(stage string, progress int) {
|
|
updateTaskProgress(taskID, "running", stage, progress)
|
|
})
|
|
|
|
l := logger.With("task_id", taskID, "project_id", projectID)
|
|
|
|
updateTaskProgress(taskID, "running", "prompt_builder", 5)
|
|
|
|
in := service.PipelineInput{
|
|
ProjectID: projectID,
|
|
TaskID: taskID,
|
|
Prompt: req.Prompt,
|
|
AssetType: req.AssetType,
|
|
Tags: req.Tags,
|
|
UserNote: req.UserNote,
|
|
ProjectStyle: req.ProjectStyle,
|
|
TaskStyle: req.TaskStyle,
|
|
Params: service.AssetParams{
|
|
Resolution: req.Resolution,
|
|
Frames: service.FrameParams{
|
|
Directions: req.Directions,
|
|
FramesPerDirection: req.FramesPerDir,
|
|
},
|
|
Format: req.Format,
|
|
},
|
|
}
|
|
|
|
output, err := service.RunPipeline(ctx, in)
|
|
if err != nil {
|
|
l.Error("task pipeline failed", "error", err)
|
|
updateFailed(taskID, "生成管线执行失败")
|
|
return
|
|
}
|
|
|
|
updateTaskProgress(taskID, "saving", "format_adapter", 90)
|
|
|
|
assets := make([]AssetResponse, len(output.Assets))
|
|
for i, a := range output.Assets {
|
|
key := fmt.Sprintf("generation/%s/%s/%d.%s", projectID, taskID, i, a.Format)
|
|
cdnURL, err := storageSvc.Upload(ctx, key, a.Data)
|
|
if err != nil {
|
|
l.Error("upload asset failed", "index", i, "error", err)
|
|
updateFailed(taskID, "上传素材失败")
|
|
return
|
|
}
|
|
assets[i] = AssetResponse{
|
|
Key: key,
|
|
URL: cdnURL,
|
|
Format: a.Format,
|
|
}
|
|
}
|
|
|
|
// 更新为完成状态
|
|
taskStore.Store(taskID, &taskRecord{
|
|
task: TaskResponse{
|
|
ID: taskID,
|
|
ProjectID: projectID,
|
|
Prompt: req.Prompt,
|
|
AssetType: req.AssetType,
|
|
Status: "completed",
|
|
Progress: 100,
|
|
CreatedAt: time.Now().Format(time.RFC3339),
|
|
},
|
|
assets: assets,
|
|
metadata: output.Metadata,
|
|
})
|
|
|
|
l.Info("task completed", "asset_count", len(assets))
|
|
}
|
|
|
|
func updateTaskProgress(taskID, status, stage string, progress int) {
|
|
rec, ok := taskStore.Load(taskID)
|
|
if !ok {
|
|
return
|
|
}
|
|
r := rec.(*taskRecord)
|
|
r.task.Status = status
|
|
r.task.Stage = stage
|
|
r.task.Progress = progress
|
|
taskStore.Store(taskID, r)
|
|
|
|
// 同步更新数据库状态
|
|
updateTaskInDB(taskID, status, stage, "", progress)
|
|
}
|
|
|
|
func updateFailed(taskID, errMsg string) {
|
|
rec, ok := taskStore.Load(taskID)
|
|
if !ok {
|
|
return
|
|
}
|
|
r := rec.(*taskRecord)
|
|
r.task.Status = "failed"
|
|
r.task.Error = errMsg
|
|
taskStore.Store(taskID, r)
|
|
|
|
// 同步更新数据库状态
|
|
updateTaskInDB(taskID, "failed", r.task.Stage, errMsg, r.task.Progress)
|
|
}
|
|
|
|
// GetTask 查询任务信息。
|
|
func GetTask(c *gin.Context) {
|
|
taskID := c.Param("taskId")
|
|
rec, ok := taskStore.Load(taskID)
|
|
if !ok {
|
|
c.JSON(http.StatusNotFound, model.Fail(http.StatusNotFound, "任务不存在"))
|
|
return
|
|
}
|
|
r := rec.(*taskRecord)
|
|
c.JSON(http.StatusOK, model.OK(r.task))
|
|
}
|
|
|
|
// GetAssets 查询任务素材列表。
|
|
func GetAssets(c *gin.Context) {
|
|
taskID := c.Param("taskId")
|
|
rec, ok := taskStore.Load(taskID)
|
|
if !ok {
|
|
c.JSON(http.StatusNotFound, model.Fail(http.StatusNotFound, "任务不存在"))
|
|
return
|
|
}
|
|
r := rec.(*taskRecord)
|
|
if r.task.Status != "completed" {
|
|
c.JSON(http.StatusBadRequest, model.Fail(http.StatusBadRequest, "任务尚未完成，当前状态: "+r.task.Status))
|
|
return
|
|
}
|
|
c.JSON(http.StatusOK, model.OK(AssetsResponse{
|
|
Assets: r.assets,
|
|
Metadata: r.metadata,
|
|
}))
|
|
}
|
|
|
|
// saveTaskToDB 将任务记录保存到数据库。
|
|
func saveTaskToDB(ctx context.Context, projectID, taskID string, req GenerateRequest) {
|
|
// "default" 工程没有数据库记录，跳过
|
|
if projectID == "default" {
|
|
return
|
|
}
|
|
|
|
projectIDUint, err := strconv.ParseUint(projectID, 10, 32)
|
|
if err != nil {
|
|
logger.FromCtx(ctx).Error("failed to parse projectID", "projectID", projectID, "error", err)
|
|
return
|
|
}
|
|
|
|
task := &model.Task{
|
|
ExternalID: taskID,
|
|
ProjectID: uint(projectIDUint),
|
|
Prompt: req.Prompt,
|
|
AssetType: req.AssetType,
|
|
Status: "pending",
|
|
Progress: 0,
|
|
RetryCount: 0,
|
|
CreatedAt: time.Now(),
|
|
UpdatedAt: time.Now(),
|
|
}
|
|
|
|
if err := db.GetDB().WithContext(ctx).Create(task).Error; err != nil {
|
|
logger.FromCtx(ctx).Error("failed to save task to db", "taskID", taskID, "error", err)
|
|
}
|
|
}
|
|
|
|
// updateTaskInDB 更新数据库中的任务状态。
|
|
func updateTaskInDB(taskID string, status, stage, error string, progress int) {
|
|
db.GetDB().Model(&model.Task{}).
|
|
Where("external_id = ?", taskID).
|
|
Updates(map[string]interface{}{
|
|
"status": status,
|
|
"stage": stage,
|
|
"progress": progress,
|
|
"error": error,
|
|
"updated_at": time.Now(),
|
|
})
|
|
}
|