refactor(task): extract billing and polling logic from controller to service layer

Restructure the task relay system for better separation of concerns:
- Extract task billing into service/task_billing.go with unified settlement flow
- Move task polling loop from controller to service/task_polling.go (supports Suno + video platforms)
- Split RelayTask into fetch/submit paths with dedicated retry logic (taskSubmitWithRetry)
- Add TaskDto, TaskResponse generics, and FetchReq to dto/task.go
- Add taskcommon/helpers.go for shared task adaptor utilities
- Remove controller/task_video.go (logic consolidated into service layer)
- Update all task adaptors (ali, doubao, gemini, hailuo, jimeng, kling, sora, suno, vertex, vidu)
- Simplify frontend task logs to use new TaskDto response format
This commit is contained in:
CaIon
2026-02-10 20:40:33 +08:00
parent e0a6ee1cb8
commit 9e3954428d
34 changed files with 1465 additions and 1191 deletions

View File

@@ -384,7 +384,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
// 转换为 OpenAI 格式响应
openAIResp := dto.NewOpenAIVideo()
openAIResp.ID = aliResp.Output.TaskID
openAIResp.ID = info.PublicTaskID
openAIResp.TaskID = info.PublicTaskID
openAIResp.Model = c.GetString("model")
if openAIResp.Model == "" && info != nil {
openAIResp.Model = info.OriginModelName

View File

@@ -2,7 +2,6 @@ package doubao
import (
"bytes"
"encoding/json"
"fmt"
"io"
"net/http"
@@ -14,6 +13,7 @@ import (
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay/channel"
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
@@ -131,7 +131,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
return nil, errors.Wrap(err, "convert request payload failed")
}
info.UpstreamModelName = body.Model
data, err := json.Marshal(body)
data, err := common.Marshal(body)
if err != nil {
return nil, err
}
@@ -154,7 +154,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
// Parse Doubao response
var dResp responsePayload
if err := json.Unmarshal(responseBody, &dResp); err != nil {
if err := common.Unmarshal(responseBody, &dResp); err != nil {
taskErr = service.TaskErrorWrapper(errors.Wrapf(err, "body: %s", responseBody), "unmarshal_response_body_failed", http.StatusInternalServerError)
return
}
@@ -165,8 +165,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
}
ov := dto.NewOpenAIVideo()
ov.ID = dResp.ID
ov.TaskID = dResp.ID
ov.ID = info.PublicTaskID
ov.TaskID = info.PublicTaskID
ov.CreatedAt = time.Now().Unix()
ov.Model = info.OriginModelName
@@ -234,12 +234,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
}
metadata := req.Metadata
medaBytes, err := json.Marshal(metadata)
if err != nil {
return nil, errors.Wrap(err, "metadata marshal metadata failed")
}
err = json.Unmarshal(medaBytes, &r)
if err != nil {
if err := taskcommon.UnmarshalMetadata(metadata, &r); err != nil {
return nil, errors.Wrap(err, "unmarshal metadata failed")
}
@@ -248,7 +243,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
resTask := responseTask{}
if err := json.Unmarshal(respBody, &resTask); err != nil {
if err := common.Unmarshal(respBody, &resTask); err != nil {
return nil, errors.Wrap(err, "unmarshal task result failed")
}
@@ -286,7 +281,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
var dResp responseTask
if err := json.Unmarshal(originTask.Data, &dResp); err != nil {
if err := common.Unmarshal(originTask.Data, &dResp); err != nil {
return nil, errors.Wrap(err, "unmarshal doubao task data failed")
}
@@ -307,6 +302,5 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, erro
}
}
jsonData, _ := common.Marshal(openAIVideo)
return jsonData, nil
return common.Marshal(openAIVideo)
}

View File

@@ -2,8 +2,6 @@ package gemini
import (
"bytes"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net/http"
@@ -16,10 +14,10 @@ import (
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay/channel"
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/QuantumNous/new-api/setting/system_setting"
"github.com/gin-gonic/gin"
"github.com/pkg/errors"
)
@@ -145,16 +143,11 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
}
metadata := req.Metadata
medaBytes, err := json.Marshal(metadata)
if err != nil {
return nil, errors.Wrap(err, "metadata marshal metadata failed")
}
err = json.Unmarshal(medaBytes, &body.Parameters)
if err != nil {
if err := taskcommon.UnmarshalMetadata(metadata, &body.Parameters); err != nil {
return nil, errors.Wrap(err, "unmarshal metadata failed")
}
data, err := json.Marshal(body)
data, err := common.Marshal(body)
if err != nil {
return nil, err
}
@@ -175,16 +168,16 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
_ = resp.Body.Close()
var s submitResponse
if err := json.Unmarshal(responseBody, &s); err != nil {
if err := common.Unmarshal(responseBody, &s); err != nil {
return "", nil, service.TaskErrorWrapper(err, "unmarshal_response_failed", http.StatusInternalServerError)
}
if strings.TrimSpace(s.Name) == "" {
return "", nil, service.TaskErrorWrapper(fmt.Errorf("missing operation name"), "invalid_response", http.StatusInternalServerError)
}
taskID = encodeLocalTaskID(s.Name)
taskID = taskcommon.EncodeLocalTaskID(s.Name)
ov := dto.NewOpenAIVideo()
ov.ID = taskID
ov.TaskID = taskID
ov.ID = info.PublicTaskID
ov.TaskID = info.PublicTaskID
ov.CreatedAt = time.Now().Unix()
ov.Model = info.OriginModelName
c.JSON(http.StatusOK, ov)
@@ -206,7 +199,7 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
return nil, fmt.Errorf("invalid task_id")
}
upstreamName, err := decodeLocalTaskID(taskID)
upstreamName, err := taskcommon.DecodeLocalTaskID(taskID)
if err != nil {
return nil, fmt.Errorf("decode task_id failed: %w", err)
}
@@ -232,7 +225,7 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
var op operationResponse
if err := json.Unmarshal(respBody, &op); err != nil {
if err := common.Unmarshal(respBody, &op); err != nil {
return nil, fmt.Errorf("unmarshal operation response failed: %w", err)
}
@@ -254,9 +247,8 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
ti.Status = model.TaskStatusSuccess
ti.Progress = "100%"
taskID := encodeLocalTaskID(op.Name)
ti.TaskID = taskID
ti.Url = fmt.Sprintf("%s/v1/videos/%s/content", system_setting.ServerAddress, taskID)
ti.TaskID = taskcommon.EncodeLocalTaskID(op.Name)
// Url intentionally left empty — the caller constructs the proxy URL using the public task ID
// Extract URL from generateVideoResponse if available
if len(op.Response.GenerateVideoResponse.GeneratedSamples) > 0 {
@@ -269,7 +261,10 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
}
func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
upstreamName, err := decodeLocalTaskID(task.TaskID)
// Use GetUpstreamTaskID() to get the real upstream operation name for model extraction.
// task.TaskID is now a public task_xxxx ID, no longer a base64-encoded upstream name.
upstreamTaskID := task.GetUpstreamTaskID()
upstreamName, err := taskcommon.DecodeLocalTaskID(upstreamTaskID)
if err != nil {
upstreamName = ""
}
@@ -297,18 +292,6 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
// helpers
// ============================
func encodeLocalTaskID(name string) string {
return base64.RawURLEncoding.EncodeToString([]byte(name))
}
func decodeLocalTaskID(local string) (string, error) {
b, err := base64.RawURLEncoding.DecodeString(local)
if err != nil {
return "", err
}
return string(b), nil
}
var modelRe = regexp.MustCompile(`models/([^/]+)/operations/`)
func extractModelFromOperationName(name string) string {

View File

@@ -2,7 +2,6 @@ package hailuo
import (
"bytes"
"encoding/json"
"fmt"
"io"
"net/http"
@@ -65,7 +64,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
return nil, errors.Wrap(err, "convert request payload failed")
}
data, err := json.Marshal(body)
data, err := common.Marshal(body)
if err != nil {
return nil, err
}
@@ -86,7 +85,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
_ = resp.Body.Close()
var hResp VideoResponse
if err := json.Unmarshal(responseBody, &hResp); err != nil {
if err := common.Unmarshal(responseBody, &hResp); err != nil {
taskErr = service.TaskErrorWrapper(errors.Wrapf(err, "body: %s", responseBody), "unmarshal_response_body_failed", http.StatusInternalServerError)
return
}
@@ -101,8 +100,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
}
ov := dto.NewOpenAIVideo()
ov.ID = hResp.TaskID
ov.TaskID = hResp.TaskID
ov.ID = info.PublicTaskID
ov.TaskID = info.PublicTaskID
ov.CreatedAt = time.Now().Unix()
ov.Model = info.OriginModelName
@@ -182,7 +181,7 @@ func (a *TaskAdaptor) parseResolutionFromSize(size string, modelConfig ModelConf
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
resTask := QueryTaskResponse{}
if err := json.Unmarshal(respBody, &resTask); err != nil {
if err := common.Unmarshal(respBody, &resTask); err != nil {
return nil, errors.Wrap(err, "unmarshal task result failed")
}
@@ -224,7 +223,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
var hailuoResp QueryTaskResponse
if err := json.Unmarshal(originTask.Data, &hailuoResp); err != nil {
if err := common.Unmarshal(originTask.Data, &hailuoResp); err != nil {
return nil, errors.Wrap(err, "unmarshal hailuo task data failed")
}
@@ -271,7 +270,7 @@ func (a *TaskAdaptor) buildVideoURL(_, fileID string) string {
}
var retrieveResp RetrieveFileResponse
if err := json.Unmarshal(responseBody, &retrieveResp); err != nil {
if err := common.Unmarshal(responseBody, &retrieveResp); err != nil {
return ""
}

View File

@@ -6,7 +6,6 @@ import (
"crypto/sha256"
"encoding/base64"
"encoding/hex"
"encoding/json"
"fmt"
"io"
"net/http"
@@ -25,6 +24,7 @@ import (
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/relay/channel"
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
)
@@ -168,7 +168,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
if err != nil {
return nil, errors.Wrap(err, "convert request payload failed")
}
data, err := json.Marshal(body)
data, err := common.Marshal(body)
if err != nil {
return nil, err
}
@@ -191,7 +191,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
// Parse Jimeng response
var jResp responsePayload
if err := json.Unmarshal(responseBody, &jResp); err != nil {
if err := common.Unmarshal(responseBody, &jResp); err != nil {
taskErr = service.TaskErrorWrapper(errors.Wrapf(err, "body: %s", responseBody), "unmarshal_response_body_failed", http.StatusInternalServerError)
return
}
@@ -202,8 +202,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
}
ov := dto.NewOpenAIVideo()
ov.ID = jResp.Data.TaskID
ov.TaskID = jResp.Data.TaskID
ov.ID = info.PublicTaskID
ov.TaskID = info.PublicTaskID
ov.CreatedAt = time.Now().Unix()
ov.Model = info.OriginModelName
c.JSON(http.StatusOK, ov)
@@ -225,7 +225,7 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
"req_key": "jimeng_vgfm_t2v_l20", // This is fixed value from doc: https://www.volcengine.com/docs/85621/1544774
"task_id": taskID,
}
payloadBytes, err := json.Marshal(payload)
payloadBytes, err := common.Marshal(payload)
if err != nil {
return nil, errors.Wrap(err, "marshal fetch task payload failed")
}
@@ -398,13 +398,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
r.BinaryDataBase64 = req.Images
}
}
metadata := req.Metadata
medaBytes, err := json.Marshal(metadata)
if err != nil {
return nil, errors.Wrap(err, "metadata marshal metadata failed")
}
err = json.Unmarshal(medaBytes, &r)
if err != nil {
if err := taskcommon.UnmarshalMetadata(req.Metadata, &r); err != nil {
return nil, errors.Wrap(err, "unmarshal metadata failed")
}
@@ -432,7 +426,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
resTask := responseTask{}
if err := json.Unmarshal(respBody, &resTask); err != nil {
if err := common.Unmarshal(respBody, &resTask); err != nil {
return nil, errors.Wrap(err, "unmarshal task result failed")
}
taskResult := relaycommon.TaskInfo{}
@@ -458,7 +452,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
var jimengResp responseTask
if err := json.Unmarshal(originTask.Data, &jimengResp); err != nil {
if err := common.Unmarshal(originTask.Data, &jimengResp); err != nil {
return nil, errors.Wrap(err, "unmarshal jimeng task data failed")
}
@@ -477,8 +471,7 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, erro
}
}
jsonData, _ := common.Marshal(openAIVideo)
return jsonData, nil
return common.Marshal(openAIVideo)
}
func isNewAPIRelay(apiKey string) bool {

View File

@@ -2,7 +2,6 @@ package kling
import (
"bytes"
"encoding/json"
"fmt"
"io"
"net/http"
@@ -21,6 +20,7 @@ import (
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/relay/channel"
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
)
@@ -156,7 +156,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
if body.Image == "" && body.ImageTail == "" {
c.Set("action", constant.TaskActionTextGenerate)
}
data, err := json.Marshal(body)
data, err := common.Marshal(body)
if err != nil {
return nil, err
}
@@ -180,7 +180,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
}
var kResp responsePayload
err = json.Unmarshal(responseBody, &kResp)
err = common.Unmarshal(responseBody, &kResp)
if err != nil {
taskErr = service.TaskErrorWrapper(err, "unmarshal_response_failed", http.StatusInternalServerError)
return
@@ -190,8 +190,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
return
}
ov := dto.NewOpenAIVideo()
ov.ID = kResp.Data.TaskId
ov.TaskID = kResp.Data.TaskId
ov.ID = info.PublicTaskID
ov.TaskID = info.PublicTaskID
ov.CreatedAt = time.Now().Unix()
ov.Model = info.OriginModelName
c.JSON(http.StatusOK, ov)
@@ -251,8 +251,8 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
r := requestPayload{
Prompt: req.Prompt,
Image: req.Image,
Mode: defaultString(req.Mode, "std"),
Duration: fmt.Sprintf("%d", defaultInt(req.Duration, 5)),
Mode: taskcommon.DefaultString(req.Mode, "std"),
Duration: fmt.Sprintf("%d", taskcommon.DefaultInt(req.Duration, 5)),
AspectRatio: a.getAspectRatio(req.Size),
ModelName: req.Model,
Model: req.Model, // Keep consistent with model_name, double writing improves compatibility
@@ -266,13 +266,7 @@ func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*
if r.ModelName == "" {
r.ModelName = "kling-v1"
}
metadata := req.Metadata
medaBytes, err := json.Marshal(metadata)
if err != nil {
return nil, errors.Wrap(err, "metadata marshal metadata failed")
}
err = json.Unmarshal(medaBytes, &r)
if err != nil {
if err := taskcommon.UnmarshalMetadata(req.Metadata, &r); err != nil {
return nil, errors.Wrap(err, "unmarshal metadata failed")
}
return &r, nil
@@ -291,20 +285,6 @@ func (a *TaskAdaptor) getAspectRatio(size string) string {
}
}
func defaultString(s, def string) string {
if strings.TrimSpace(s) == "" {
return def
}
return s
}
func defaultInt(v int, def int) int {
if v == 0 {
return def
}
return v
}
// ============================
// JWT helpers
// ============================
@@ -340,7 +320,7 @@ func (a *TaskAdaptor) createJWTTokenWithKey(apiKey string) (string, error) {
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
taskInfo := &relaycommon.TaskInfo{}
resPayload := responsePayload{}
err := json.Unmarshal(respBody, &resPayload)
err := common.Unmarshal(respBody, &resPayload)
if err != nil {
return nil, errors.Wrap(err, "failed to unmarshal response body")
}
@@ -374,7 +354,7 @@ func isNewAPIRelay(apiKey string) bool {
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
var klingResp responsePayload
if err := json.Unmarshal(originTask.Data, &klingResp); err != nil {
if err := common.Unmarshal(originTask.Data, &klingResp); err != nil {
return nil, errors.Wrap(err, "unmarshal kling task data failed")
}
@@ -401,6 +381,5 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, erro
Code: fmt.Sprintf("%d", klingResp.Code),
}
}
jsonData, _ := common.Marshal(openAIVideo)
return jsonData, nil
return common.Marshal(openAIVideo)
}

View File

@@ -13,7 +13,6 @@ import (
"github.com/QuantumNous/new-api/relay/channel"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/system_setting"
"github.com/gin-gonic/gin"
"github.com/pkg/errors"
@@ -116,7 +115,7 @@ func (a *TaskAdaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, req
}
// DoResponse handles upstream response, returns taskID etc.
func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, _ *relaycommon.RelayInfo) (taskID string, taskData []byte, taskErr *dto.TaskError) {
func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (taskID string, taskData []byte, taskErr *dto.TaskError) {
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
taskErr = service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError)
@@ -131,17 +130,20 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, _ *relayco
return
}
if dResp.ID == "" {
if dResp.TaskID == "" {
taskErr = service.TaskErrorWrapper(fmt.Errorf("task_id is empty"), "invalid_response", http.StatusInternalServerError)
return
}
dResp.ID = dResp.TaskID
dResp.TaskID = ""
upstreamID := dResp.ID
if upstreamID == "" {
upstreamID = dResp.TaskID
}
if upstreamID == "" {
taskErr = service.TaskErrorWrapper(fmt.Errorf("task_id is empty"), "invalid_response", http.StatusInternalServerError)
return
}
// 使用公开 task_xxxx ID 返回给客户端
dResp.ID = info.PublicTaskID
dResp.TaskID = info.PublicTaskID
c.JSON(http.StatusOK, dResp)
return dResp.ID, responseBody, nil
return upstreamID, responseBody, nil
}
// FetchTask fetch task status
@@ -192,7 +194,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
taskResult.Status = model.TaskStatusInProgress
case "completed":
taskResult.Status = model.TaskStatusSuccess
taskResult.Url = fmt.Sprintf("%s/v1/videos/%s/content", system_setting.ServerAddress, resTask.ID)
// Url intentionally left empty — the caller constructs the proxy URL using the public task ID
case "failed", "cancelled":
taskResult.Status = model.TaskStatusFailure
if resTask.Error != nil {

View File

@@ -3,7 +3,6 @@ package suno
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
@@ -24,8 +23,12 @@ type TaskAdaptor struct {
ChannelType int
}
// ParseTaskResult is not used for Suno tasks.
// Suno polling uses a dedicated batch-fetch path (service.UpdateSunoTasks) that
// receives dto.TaskResponse[[]dto.SunoDataResponse] from the upstream /fetch API.
// This differs from the per-task polling used by video adaptors.
func (a *TaskAdaptor) ParseTaskResult([]byte) (*relaycommon.TaskInfo, error) {
return nil, fmt.Errorf("not implement") // todo implement this method if needed
return nil, fmt.Errorf("suno uses batch polling via UpdateSunoTasks, ParseTaskResult is not applicable")
}
func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) {
@@ -81,7 +84,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
return nil, err
}
}
data, err := json.Marshal(sunoRequest)
data, err := common.Marshal(sunoRequest)
if err != nil {
return nil, err
}
@@ -99,7 +102,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
return
}
var sunoResponse dto.TaskResponse[string]
err = json.Unmarshal(responseBody, &sunoResponse)
err = common.Unmarshal(responseBody, &sunoResponse)
if err != nil {
taskErr = service.TaskErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError)
return
@@ -109,17 +112,13 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
return
}
for k, v := range resp.Header {
c.Writer.Header().Set(k, v[0])
}
c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode)
_, err = io.Copy(c.Writer, bytes.NewBuffer(responseBody))
if err != nil {
taskErr = service.TaskErrorWrapper(err, "copy_response_body_failed", http.StatusInternalServerError)
return
// 使用公开 task_xxxx ID 替换上游 ID 返回给客户端
publicResponse := dto.TaskResponse[string]{
Code: sunoResponse.Code,
Message: sunoResponse.Message,
Data: info.PublicTaskID,
}
c.JSON(http.StatusOK, publicResponse)
return sunoResponse.Data, nil, nil
}
@@ -134,7 +133,7 @@ func (a *TaskAdaptor) GetChannelName() string {
func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy string) (*http.Response, error) {
requestUrl := fmt.Sprintf("%s/suno/fetch", baseUrl)
byteBody, err := json.Marshal(body)
byteBody, err := common.Marshal(body)
if err != nil {
return nil, err
}

View File

@@ -0,0 +1,70 @@
package taskcommon
import (
"encoding/base64"
"fmt"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/setting/system_setting"
)
// UnmarshalMetadata converts a map[string]any metadata to a typed struct via JSON round-trip.
// This replaces the repeated pattern: json.Marshal(metadata) → json.Unmarshal(bytes, &target).
func UnmarshalMetadata(metadata map[string]any, target any) error {
if metadata == nil {
return nil
}
metaBytes, err := common.Marshal(metadata)
if err != nil {
return fmt.Errorf("marshal metadata failed: %w", err)
}
if err := common.Unmarshal(metaBytes, target); err != nil {
return fmt.Errorf("unmarshal metadata failed: %w", err)
}
return nil
}
// DefaultString returns val if non-empty, otherwise fallback.
func DefaultString(val, fallback string) string {
if val == "" {
return fallback
}
return val
}
// DefaultInt returns val if non-zero, otherwise fallback.
func DefaultInt(val, fallback int) int {
if val == 0 {
return fallback
}
return val
}
// EncodeLocalTaskID encodes an upstream operation name to a URL-safe base64 string.
// Used by Gemini/Vertex to store upstream names as task IDs.
func EncodeLocalTaskID(name string) string {
return base64.RawURLEncoding.EncodeToString([]byte(name))
}
// DecodeLocalTaskID decodes a base64-encoded upstream operation name.
func DecodeLocalTaskID(id string) (string, error) {
b, err := base64.RawURLEncoding.DecodeString(id)
if err != nil {
return "", err
}
return string(b), nil
}
// BuildProxyURL constructs the video proxy URL using the public task ID.
// e.g., "https://your-server.com/v1/videos/task_xxxx/content"
func BuildProxyURL(taskID string) string {
return fmt.Sprintf("%s/v1/videos/%s/content", system_setting.ServerAddress, taskID)
}
// Status-to-progress mapping constants for polling updates.
const (
ProgressSubmitted = "10%"
ProgressQueued = "20%"
ProgressInProgress = "30%"
ProgressComplete = "100%"
)

View File

@@ -2,13 +2,12 @@ package vertex
import (
"bytes"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net/http"
"regexp"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
@@ -17,6 +16,7 @@ import (
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/relay/channel"
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
vertexcore "github.com/QuantumNous/new-api/relay/channel/vertex"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
@@ -82,7 +82,7 @@ func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycom
// BuildRequestURL constructs the upstream URL.
func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) {
adc := &vertexcore.Credentials{}
if err := json.Unmarshal([]byte(a.apiKey), adc); err != nil {
if err := common.Unmarshal([]byte(a.apiKey), adc); err != nil {
return "", fmt.Errorf("failed to decode credentials: %w", err)
}
modelName := info.OriginModelName
@@ -116,7 +116,7 @@ func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info
req.Header.Set("Accept", "application/json")
adc := &vertexcore.Credentials{}
if err := json.Unmarshal([]byte(a.apiKey), adc); err != nil {
if err := common.Unmarshal([]byte(a.apiKey), adc); err != nil {
return fmt.Errorf("failed to decode credentials: %w", err)
}
@@ -184,7 +184,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
// info.PriceData.OtherRatios["durationSeconds"] = float64(v.(int))
// }
data, err := json.Marshal(body)
data, err := common.Marshal(body)
if err != nil {
return nil, err
}
@@ -205,14 +205,19 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
_ = resp.Body.Close()
var s submitResponse
if err := json.Unmarshal(responseBody, &s); err != nil {
if err := common.Unmarshal(responseBody, &s); err != nil {
return "", nil, service.TaskErrorWrapper(err, "unmarshal_response_failed", http.StatusInternalServerError)
}
if strings.TrimSpace(s.Name) == "" {
return "", nil, service.TaskErrorWrapper(fmt.Errorf("missing operation name"), "invalid_response", http.StatusInternalServerError)
}
localID := encodeLocalTaskID(s.Name)
c.JSON(http.StatusOK, gin.H{"task_id": localID})
localID := taskcommon.EncodeLocalTaskID(s.Name)
ov := dto.NewOpenAIVideo()
ov.ID = info.PublicTaskID
ov.TaskID = info.PublicTaskID
ov.CreatedAt = time.Now().Unix()
ov.Model = info.OriginModelName
c.JSON(http.StatusOK, ov)
return localID, responseBody, nil
}
@@ -225,7 +230,7 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
if !ok {
return nil, fmt.Errorf("invalid task_id")
}
upstreamName, err := decodeLocalTaskID(taskID)
upstreamName, err := taskcommon.DecodeLocalTaskID(taskID)
if err != nil {
return nil, fmt.Errorf("decode task_id failed: %w", err)
}
@@ -245,12 +250,12 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
url = fmt.Sprintf("https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/publishers/google/models/%s:fetchPredictOperation", region, project, region, modelName)
}
payload := map[string]string{"operationName": upstreamName}
data, err := json.Marshal(payload)
data, err := common.Marshal(payload)
if err != nil {
return nil, err
}
adc := &vertexcore.Credentials{}
if err := json.Unmarshal([]byte(key), adc); err != nil {
if err := common.Unmarshal([]byte(key), adc); err != nil {
return nil, fmt.Errorf("failed to decode credentials: %w", err)
}
token, err := vertexcore.AcquireAccessToken(*adc, proxy)
@@ -274,7 +279,7 @@ func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
var op operationResponse
if err := json.Unmarshal(respBody, &op); err != nil {
if err := common.Unmarshal(respBody, &op); err != nil {
return nil, fmt.Errorf("unmarshal operation response failed: %w", err)
}
ti := &relaycommon.TaskInfo{}
@@ -338,7 +343,10 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
}
func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
upstreamName, err := decodeLocalTaskID(task.TaskID)
// Use GetUpstreamTaskID() to get the real upstream operation name for model extraction.
// task.TaskID is now a public task_xxxx ID, no longer a base64-encoded upstream name.
upstreamTaskID := task.GetUpstreamTaskID()
upstreamName, err := taskcommon.DecodeLocalTaskID(upstreamTaskID)
if err != nil {
upstreamName = ""
}
@@ -353,8 +361,8 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
v.SetProgressStr(task.Progress)
v.CreatedAt = task.CreatedAt
v.CompletedAt = task.UpdatedAt
if strings.HasPrefix(task.FailReason, "data:") && len(task.FailReason) > 0 {
v.SetMetadata("url", task.FailReason)
if resultURL := task.GetResultURL(); strings.HasPrefix(resultURL, "data:") && len(resultURL) > 0 {
v.SetMetadata("url", resultURL)
}
return common.Marshal(v)
@@ -364,18 +372,6 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(task *model.Task) ([]byte, error) {
// helpers
// ============================
func encodeLocalTaskID(name string) string {
return base64.RawURLEncoding.EncodeToString([]byte(name))
}
func decodeLocalTaskID(local string) (string, error) {
b, err := base64.RawURLEncoding.DecodeString(local)
if err != nil {
return "", err
}
return string(b), nil
}
var regionRe = regexp.MustCompile(`locations/([a-z0-9-]+)/`)
func extractRegionFromOperationName(name string) string {

View File

@@ -2,7 +2,6 @@ package vidu
import (
"bytes"
"encoding/json"
"fmt"
"io"
"net/http"
@@ -16,6 +15,7 @@ import (
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay/channel"
taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
@@ -127,7 +127,7 @@ func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayIn
}
}
data, err := json.Marshal(body)
data, err := common.Marshal(body)
if err != nil {
return nil, err
}
@@ -168,7 +168,7 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
}
var vResp responsePayload
err = json.Unmarshal(responseBody, &vResp)
err = common.Unmarshal(responseBody, &vResp)
if err != nil {
taskErr = service.TaskErrorWrapper(errors.Wrap(err, fmt.Sprintf("%s", responseBody)), "unmarshal_response_failed", http.StatusInternalServerError)
return
@@ -180,8 +180,8 @@ func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *rela
}
ov := dto.NewOpenAIVideo()
ov.ID = vResp.TaskId
ov.TaskID = vResp.TaskId
ov.ID = info.PublicTaskID
ov.TaskID = info.PublicTaskID
ov.CreatedAt = time.Now().Unix()
ov.Model = info.OriginModelName
c.JSON(http.StatusOK, ov)
@@ -225,45 +225,25 @@ func (a *TaskAdaptor) GetChannelName() string {
func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*requestPayload, error) {
r := requestPayload{
Model: defaultString(req.Model, "viduq1"),
Model: taskcommon.DefaultString(req.Model, "viduq1"),
Images: req.Images,
Prompt: req.Prompt,
Duration: defaultInt(req.Duration, 5),
Resolution: defaultString(req.Size, "1080p"),
Duration: taskcommon.DefaultInt(req.Duration, 5),
Resolution: taskcommon.DefaultString(req.Size, "1080p"),
MovementAmplitude: "auto",
Bgm: false,
}
metadata := req.Metadata
medaBytes, err := json.Marshal(metadata)
if err != nil {
return nil, errors.Wrap(err, "metadata marshal metadata failed")
}
err = json.Unmarshal(medaBytes, &r)
if err != nil {
if err := taskcommon.UnmarshalMetadata(req.Metadata, &r); err != nil {
return nil, errors.Wrap(err, "unmarshal metadata failed")
}
return &r, nil
}
func defaultString(value, defaultValue string) string {
if value == "" {
return defaultValue
}
return value
}
func defaultInt(value, defaultValue int) int {
if value == 0 {
return defaultValue
}
return value
}
func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
taskInfo := &relaycommon.TaskInfo{}
var taskResp taskResultResponse
err := json.Unmarshal(respBody, &taskResp)
err := common.Unmarshal(respBody, &taskResp)
if err != nil {
return nil, errors.Wrap(err, "failed to unmarshal response body")
}
@@ -293,7 +273,7 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) {
var viduResp taskResultResponse
if err := json.Unmarshal(originTask.Data, &viduResp); err != nil {
if err := common.Unmarshal(originTask.Data, &viduResp); err != nil {
return nil, errors.Wrap(err, "unmarshal vidu task data failed")
}
@@ -315,6 +295,5 @@ func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, erro
}
}
jsonData, _ := common.Marshal(openAIVideo)
return jsonData, nil
return common.Marshal(openAIVideo)
}

View File

@@ -118,8 +118,12 @@ type RelayInfo struct {
SendResponseCount int
ReceivedResponseCount int
FinalPreConsumedQuota int // 最终预消耗的配额
// ForcePreConsume 为 true 时禁用 BillingSession 的信任额度旁路,
// 强制预扣全额。用于异步任务(视频/音乐生成等),因为请求返回后任务仍在运行,
// 必须在提交前锁定全额。
ForcePreConsume bool
// Billing 是计费会话,封装了预扣费/结算/退款的统一生命周期。
// 免费模型和按次计费MJ/Task时为 nil。
// 免费模型时为 nil。
Billing BillingSettler
// BillingSource indicates whether this request is billed from wallet quota or subscription.
// "" or "wallet" => wallet; "subscription" => subscription
@@ -525,8 +529,10 @@ func GenRelayInfo(c *gin.Context, relayFormat types.RelayFormat, request dto.Req
return nil, errors.New("request is not a OpenAIResponsesCompactionRequest")
case types.RelayFormatTask:
info = genBaseRelayInfo(c, nil)
info.TaskRelayInfo = &TaskRelayInfo{}
case types.RelayFormatMjProxy:
info = genBaseRelayInfo(c, nil)
info.TaskRelayInfo = &TaskRelayInfo{}
default:
err = errors.New("invalid relay format")
}
@@ -608,6 +614,9 @@ func (info *RelayInfo) HasSendResponse() bool {
type TaskRelayInfo struct {
Action string
OriginTaskID string
// PublicTaskID 是提交时预生成的 task_xxxx 格式公开 ID
// 供 DoResponse 在返回给客户端时使用(避免暴露上游真实 ID
PublicTaskID string
ConsumeQuota bool
}
@@ -667,11 +676,11 @@ func (t *TaskSubmitReq) UnmarshalJSON(data []byte) error {
func (t *TaskSubmitReq) UnmarshalMetadata(v any) error {
metadata := t.Metadata
if metadata != nil {
metadataBytes, err := json.Marshal(metadata)
metadataBytes, err := common.Marshal(metadata)
if err != nil {
return fmt.Errorf("marshal metadata failed: %w", err)
}
err = json.Unmarshal(metadataBytes, v)
err = common.Unmarshal(metadataBytes, v)
if err != nil {
return fmt.Errorf("unmarshal metadata to target failed: %w", err)
}

View File

@@ -140,7 +140,7 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens
}
// ModelPriceHelperPerCall 按次计费的 PriceHelper (MJ、Task)
func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) types.PerCallPriceData {
func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) types.PriceData {
groupRatioInfo := HandleGroupRatio(c, info)
modelPrice, success := ratio_setting.GetModelPrice(info.OriginModelName, true)
@@ -154,7 +154,18 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) types.
}
}
quota := int(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
priceData := types.PerCallPriceData{
// 免费模型检测(与 ModelPriceHelper 对齐)
freeModel := false
if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume {
if groupRatioInfo.GroupRatio == 0 || modelPrice == 0 {
quota = 0
freeModel = true
}
}
priceData := types.PriceData{
FreeModel: freeModel,
ModelPrice: modelPrice,
Quota: quota,
GroupRatioInfo: groupRatioInfo,

View File

@@ -2,7 +2,6 @@ package relay
import (
"bytes"
"encoding/json"
"errors"
"fmt"
"io"
@@ -15,29 +14,33 @@ import (
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/gin-gonic/gin"
)
/*
Task 任务通过平台、Action 区分任务
*/
func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (taskErr *dto.TaskError) {
info.InitChannelMeta(c)
// ensure TaskRelayInfo is initialized to avoid nil dereference when accessing embedded fields
if info.TaskRelayInfo == nil {
info.TaskRelayInfo = &relaycommon.TaskRelayInfo{}
}
type TaskSubmitResult struct {
UpstreamTaskID string
TaskData []byte
Platform constant.TaskPlatform
ModelName string
Quota int
//PerCallPrice types.PriceData
}
// ResolveOriginTask 处理基于已有任务的提交remix / continuation
// 查找原始任务、从中提取模型名称、将渠道锁定到原始任务的渠道(并通过
// specific_channel_id 禁止重试),以及提取 OtherRatios时长、分辨率
// 该函数在控制器的重试循环之前调用一次,其结果通过 info 字段和上下文持久化。
func ResolveOriginTask(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskError {
// 检测 remix action
path := c.Request.URL.Path
if strings.Contains(path, "/v1/videos/") && strings.HasSuffix(path, "/remix") {
info.Action = constant.TaskActionRemix
}
// 提取 remix 任务的 video_id
if info.Action == constant.TaskActionRemix {
videoID := c.Param("video_id")
if strings.TrimSpace(videoID) == "" {
@@ -46,241 +49,164 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (taskErr *dto.
info.OriginTaskID = videoID
}
platform := constant.TaskPlatform(c.GetString("platform"))
if info.OriginTaskID == "" {
return nil
}
// 获取原始任务信息
if info.OriginTaskID != "" {
originTask, exist, err := model.GetByTaskId(info.UserId, info.OriginTaskID)
if err != nil {
taskErr = service.TaskErrorWrapper(err, "get_origin_task_failed", http.StatusInternalServerError)
return
}
if !exist {
taskErr = service.TaskErrorWrapperLocal(errors.New("task_origin_not_exist"), "task_not_exist", http.StatusBadRequest)
return
}
if info.OriginModelName == "" {
if originTask.Properties.OriginModelName != "" {
info.OriginModelName = originTask.Properties.OriginModelName
} else if originTask.Properties.UpstreamModelName != "" {
info.OriginModelName = originTask.Properties.UpstreamModelName
} else {
var taskData map[string]interface{}
_ = json.Unmarshal(originTask.Data, &taskData)
if m, ok := taskData["model"].(string); ok && m != "" {
info.OriginModelName = m
platform = originTask.Platform
}
}
}
if originTask.ChannelId != info.ChannelId {
channel, err := model.GetChannelById(originTask.ChannelId, true)
if err != nil {
taskErr = service.TaskErrorWrapperLocal(err, "channel_not_found", http.StatusBadRequest)
return
}
if channel.Status != common.ChannelStatusEnabled {
taskErr = service.TaskErrorWrapperLocal(errors.New("the channel of the origin task is disabled"), "task_channel_disable", http.StatusBadRequest)
return
}
key, _, newAPIError := channel.GetNextEnabledKey()
if newAPIError != nil {
taskErr = service.TaskErrorWrapper(newAPIError, "channel_no_available_key", newAPIError.StatusCode)
return
}
common.SetContextKey(c, constant.ContextKeyChannelKey, key)
common.SetContextKey(c, constant.ContextKeyChannelType, channel.Type)
common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, channel.GetBaseURL())
common.SetContextKey(c, constant.ContextKeyChannelId, originTask.ChannelId)
// 查找原始任务
originTask, exist, err := model.GetByTaskId(info.UserId, info.OriginTaskID)
if err != nil {
return service.TaskErrorWrapper(err, "get_origin_task_failed", http.StatusInternalServerError)
}
if !exist {
return service.TaskErrorWrapperLocal(errors.New("task_origin_not_exist"), "task_not_exist", http.StatusBadRequest)
}
info.ChannelBaseUrl = channel.GetBaseURL()
info.ChannelId = originTask.ChannelId
info.ChannelType = channel.Type
info.ApiKey = key
platform = originTask.Platform
}
// 使用原始任务的参数
if info.Action == constant.TaskActionRemix {
// 从原始任务推导模型名称
if info.OriginModelName == "" {
if originTask.Properties.OriginModelName != "" {
info.OriginModelName = originTask.Properties.OriginModelName
} else if originTask.Properties.UpstreamModelName != "" {
info.OriginModelName = originTask.Properties.UpstreamModelName
} else {
var taskData map[string]interface{}
_ = json.Unmarshal(originTask.Data, &taskData)
secondsStr, _ := taskData["seconds"].(string)
seconds, _ := strconv.Atoi(secondsStr)
if seconds <= 0 {
seconds = 4
}
sizeStr, _ := taskData["size"].(string)
if info.PriceData.OtherRatios == nil {
info.PriceData.OtherRatios = map[string]float64{}
}
info.PriceData.OtherRatios["seconds"] = float64(seconds)
info.PriceData.OtherRatios["size"] = 1
if sizeStr == "1792x1024" || sizeStr == "1024x1792" {
info.PriceData.OtherRatios["size"] = 1.666667
_ = common.Unmarshal(originTask.Data, &taskData)
if m, ok := taskData["model"].(string); ok && m != "" {
info.OriginModelName = m
}
}
}
// 锁定到原始任务的渠道(如果与当前选中的不同)
if originTask.ChannelId != info.ChannelId {
ch, err := model.GetChannelById(originTask.ChannelId, true)
if err != nil {
return service.TaskErrorWrapperLocal(err, "channel_not_found", http.StatusBadRequest)
}
if ch.Status != common.ChannelStatusEnabled {
return service.TaskErrorWrapperLocal(errors.New("the channel of the origin task is disabled"), "task_channel_disable", http.StatusBadRequest)
}
key, _, newAPIError := ch.GetNextEnabledKey()
if newAPIError != nil {
return service.TaskErrorWrapper(newAPIError, "channel_no_available_key", newAPIError.StatusCode)
}
common.SetContextKey(c, constant.ContextKeyChannelKey, key)
common.SetContextKey(c, constant.ContextKeyChannelType, ch.Type)
common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, ch.GetBaseURL())
common.SetContextKey(c, constant.ContextKeyChannelId, originTask.ChannelId)
info.ChannelBaseUrl = ch.GetBaseURL()
info.ChannelId = originTask.ChannelId
info.ChannelType = ch.Type
info.ApiKey = key
}
// 渠道已锁定到原始任务 → 禁止重试切换到其他渠道
c.Set("specific_channel_id", fmt.Sprintf("%d", originTask.ChannelId))
// 提取 remix 参数(时长、分辨率 → OtherRatios
if info.Action == constant.TaskActionRemix {
var taskData map[string]interface{}
_ = common.Unmarshal(originTask.Data, &taskData)
secondsStr, _ := taskData["seconds"].(string)
seconds, _ := strconv.Atoi(secondsStr)
if seconds <= 0 {
seconds = 4
}
sizeStr, _ := taskData["size"].(string)
if info.PriceData.OtherRatios == nil {
info.PriceData.OtherRatios = map[string]float64{}
}
info.PriceData.OtherRatios["seconds"] = float64(seconds)
info.PriceData.OtherRatios["size"] = 1
if sizeStr == "1792x1024" || sizeStr == "1024x1792" {
info.PriceData.OtherRatios["size"] = 1.666667
}
}
return nil
}
// RelayTaskSubmit 完成 task 提交的全部流程(每次尝试调用一次):
// 刷新渠道元数据 → 确定 platform/adaptor → 验证请求 → 计算价格 →
// 预扣费(仅首次,通过 info.Billing==nil 守卫)→ 构建/发送/解析上游请求。
// 控制器负责 defer Refund 和成功后 Settle。
func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitResult, *dto.TaskError) {
info.InitChannelMeta(c)
// 1. 确定 platform → 创建适配器 → 验证请求
platform := constant.TaskPlatform(c.GetString("platform"))
if platform == "" {
platform = GetTaskPlatform(c)
}
info.InitChannelMeta(c)
adaptor := GetTaskAdaptor(platform)
if adaptor == nil {
return service.TaskErrorWrapperLocal(fmt.Errorf("invalid api platform: %s", platform), "invalid_api_platform", http.StatusBadRequest)
return nil, service.TaskErrorWrapperLocal(fmt.Errorf("invalid api platform: %s", platform), "invalid_api_platform", http.StatusBadRequest)
}
adaptor.Init(info)
// get & validate taskRequest 获取并验证文本请求
taskErr = adaptor.ValidateRequestAndSetAction(c, info)
if taskErr != nil {
return
if taskErr := adaptor.ValidateRequestAndSetAction(c, info); taskErr != nil {
return nil, taskErr
}
// 2. 确定模型名称
modelName := info.OriginModelName
if modelName == "" {
modelName = service.CoverTaskActionToModelName(platform, info.Action)
}
modelPrice, success := ratio_setting.GetModelPrice(modelName, true)
if !success {
defaultPrice, ok := ratio_setting.GetDefaultModelPriceMap()[modelName]
if !ok {
modelPrice = float64(common.PreConsumedQuota) / common.QuotaPerUnit
} else {
modelPrice = defaultPrice
}
// 3. 预生成公开 task ID仅首次
if info.PublicTaskID == "" {
info.PublicTaskID = model.GenerateTaskID()
}
// 处理 auto 分组:从 context 获取实际选中的分组
// 当使用 auto 分组时Distribute 中间件会将实际选中的分组存储在 ContextKeyAutoGroup 中
if autoGroup, exists := common.GetContextKey(c, constant.ContextKeyAutoGroup); exists {
if groupStr, ok := autoGroup.(string); ok && groupStr != "" {
info.UsingGroup = groupStr
}
}
// 4. 价格计算
info.OriginModelName = modelName
info.PriceData = helper.ModelPriceHelperPerCall(c, info)
// 预扣
groupRatio := ratio_setting.GetGroupRatio(info.UsingGroup)
var ratio float64
userGroupRatio, hasUserGroupRatio := ratio_setting.GetGroupGroupRatio(info.UserGroup, info.UsingGroup)
if hasUserGroupRatio {
ratio = modelPrice * userGroupRatio
} else {
ratio = modelPrice * groupRatio
}
// FIXME: 临时修补,支持任务仅按次计费
if !common.StringsContains(constant.TaskPricePatches, modelName) {
if len(info.PriceData.OtherRatios) > 0 {
for _, ra := range info.PriceData.OtherRatios {
if 1.0 != ra {
ratio *= ra
}
for _, ra := range info.PriceData.OtherRatios {
if ra != 1.0 {
info.PriceData.Quota = int(float64(info.PriceData.Quota) * ra)
}
}
}
println(fmt.Sprintf("model: %s, model_price: %.4f, group: %s, group_ratio: %.4f, final_ratio: %.4f", modelName, modelPrice, info.UsingGroup, groupRatio, ratio))
userQuota, err := model.GetUserQuota(info.UserId, false)
if err != nil {
taskErr = service.TaskErrorWrapper(err, "get_user_quota_failed", http.StatusInternalServerError)
return
}
quota := int(ratio * common.QuotaPerUnit)
if userQuota-quota < 0 {
taskErr = service.TaskErrorWrapperLocal(errors.New("user quota is not enough"), "quota_not_enough", http.StatusForbidden)
return
// 5. 预扣费(仅首次 — 重试时 info.Billing 已存在,跳过)
if info.Billing == nil && !info.PriceData.FreeModel {
info.ForcePreConsume = true
if apiErr := service.PreConsumeBilling(c, info.PriceData.Quota, info); apiErr != nil {
return nil, service.TaskErrorFromAPIError(apiErr)
}
}
// build body
// 6. 构建请求体
requestBody, err := adaptor.BuildRequestBody(c, info)
if err != nil {
taskErr = service.TaskErrorWrapper(err, "build_request_failed", http.StatusInternalServerError)
return
return nil, service.TaskErrorWrapper(err, "build_request_failed", http.StatusInternalServerError)
}
// do request
// 7. 发送请求
resp, err := adaptor.DoRequest(c, info, requestBody)
if err != nil {
taskErr = service.TaskErrorWrapper(err, "do_request_failed", http.StatusInternalServerError)
return
return nil, service.TaskErrorWrapper(err, "do_request_failed", http.StatusInternalServerError)
}
// handle response
if resp != nil && resp.StatusCode != http.StatusOK {
responseBody, _ := io.ReadAll(resp.Body)
taskErr = service.TaskErrorWrapper(fmt.Errorf("%s", string(responseBody)), "fail_to_fetch_task", resp.StatusCode)
return
return nil, service.TaskErrorWrapper(fmt.Errorf("%s", string(responseBody)), "fail_to_fetch_task", resp.StatusCode)
}
defer func() {
// release quota
if info.ConsumeQuota && taskErr == nil {
err := service.PostConsumeQuota(info, quota, 0, true)
if err != nil {
common.SysLog("error consuming token remain quota: " + err.Error())
}
if quota != 0 {
tokenName := c.GetString("token_name")
//gRatio := groupRatio
//if hasUserGroupRatio {
// gRatio = userGroupRatio
//}
logContent := fmt.Sprintf("操作 %s", info.Action)
// FIXME: 临时修补,支持任务仅按次计费
if common.StringsContains(constant.TaskPricePatches, modelName) {
logContent = fmt.Sprintf("%s按次计费", logContent)
} else {
if len(info.PriceData.OtherRatios) > 0 {
var contents []string
for key, ra := range info.PriceData.OtherRatios {
if 1.0 != ra {
contents = append(contents, fmt.Sprintf("%s: %.2f", key, ra))
}
}
if len(contents) > 0 {
logContent = fmt.Sprintf("%s, 计算参数:%s", logContent, strings.Join(contents, ", "))
}
}
}
other := make(map[string]interface{})
if c != nil && c.Request != nil && c.Request.URL != nil {
other["request_path"] = c.Request.URL.Path
}
other["model_price"] = modelPrice
other["group_ratio"] = groupRatio
if hasUserGroupRatio {
other["user_group_ratio"] = userGroupRatio
}
model.RecordConsumeLog(c, info.UserId, model.RecordConsumeLogParams{
ChannelId: info.ChannelId,
ModelName: modelName,
TokenName: tokenName,
Quota: quota,
Content: logContent,
TokenId: info.TokenId,
Group: info.UsingGroup,
Other: other,
})
model.UpdateUserUsedQuotaAndRequestCount(info.UserId, quota)
model.UpdateChannelUsedQuota(info.ChannelId, quota)
}
}
}()
taskID, taskData, taskErr := adaptor.DoResponse(c, resp, info)
// 8. 解析响应
upstreamTaskID, taskData, taskErr := adaptor.DoResponse(c, resp, info)
if taskErr != nil {
return
return nil, taskErr
}
info.ConsumeQuota = true
// insert task
task := model.InitTask(platform, info)
task.TaskID = taskID
task.Quota = quota
task.Data = taskData
task.Action = info.Action
err = task.Insert()
if err != nil {
taskErr = service.TaskErrorWrapper(err, "insert_task_failed", http.StatusInternalServerError)
return
}
return nil
return &TaskSubmitResult{
UpstreamTaskID: upstreamTaskID,
TaskData: taskData,
Platform: platform,
ModelName: modelName,
}, nil
}
var fetchRespBuilders = map[int]func(c *gin.Context) (respBody []byte, taskResp *dto.TaskError){
@@ -336,7 +262,7 @@ func sunoFetchRespBodyBuilder(c *gin.Context) (respBody []byte, taskResp *dto.Ta
} else {
tasks = make([]any, 0)
}
respBody, err = json.Marshal(dto.TaskResponse[[]any]{
respBody, err = common.Marshal(dto.TaskResponse[[]any]{
Code: "success",
Data: tasks,
})
@@ -357,7 +283,7 @@ func sunoFetchByIDRespBodyBuilder(c *gin.Context) (respBody []byte, taskResp *dt
return
}
respBody, err = json.Marshal(dto.TaskResponse[any]{
respBody, err = common.Marshal(dto.TaskResponse[any]{
Code: "success",
Data: TaskModel2Dto(originTask),
})
@@ -381,97 +307,16 @@ func videoFetchByIDRespBodyBuilder(c *gin.Context) (respBody []byte, taskResp *d
return
}
func() {
channelModel, err2 := model.GetChannelById(originTask.ChannelId, true)
if err2 != nil {
return
}
if channelModel.Type != constant.ChannelTypeVertexAi && channelModel.Type != constant.ChannelTypeGemini {
return
}
baseURL := constant.ChannelBaseURLs[channelModel.Type]
if channelModel.GetBaseURL() != "" {
baseURL = channelModel.GetBaseURL()
}
proxy := channelModel.GetSetting().Proxy
adaptor := GetTaskAdaptor(constant.TaskPlatform(strconv.Itoa(channelModel.Type)))
if adaptor == nil {
return
}
resp, err2 := adaptor.FetchTask(baseURL, channelModel.Key, map[string]any{
"task_id": originTask.TaskID,
"action": originTask.Action,
}, proxy)
if err2 != nil || resp == nil {
return
}
defer resp.Body.Close()
body, err2 := io.ReadAll(resp.Body)
if err2 != nil {
return
}
ti, err2 := adaptor.ParseTaskResult(body)
if err2 == nil && ti != nil {
if ti.Status != "" {
originTask.Status = model.TaskStatus(ti.Status)
}
if ti.Progress != "" {
originTask.Progress = ti.Progress
}
if ti.Url != "" {
if strings.HasPrefix(ti.Url, "data:") {
} else {
originTask.FailReason = ti.Url
}
}
_ = originTask.Update()
var raw map[string]any
_ = json.Unmarshal(body, &raw)
format := "mp4"
if respObj, ok := raw["response"].(map[string]any); ok {
if vids, ok := respObj["videos"].([]any); ok && len(vids) > 0 {
if v0, ok := vids[0].(map[string]any); ok {
if mt, ok := v0["mimeType"].(string); ok && mt != "" {
if strings.Contains(mt, "mp4") {
format = "mp4"
} else {
format = mt
}
}
}
}
}
status := "processing"
switch originTask.Status {
case model.TaskStatusSuccess:
status = "succeeded"
case model.TaskStatusFailure:
status = "failed"
case model.TaskStatusQueued, model.TaskStatusSubmitted:
status = "queued"
}
if !strings.HasPrefix(c.Request.RequestURI, "/v1/videos/") {
out := map[string]any{
"error": nil,
"format": format,
"metadata": nil,
"status": status,
"task_id": originTask.TaskID,
"url": originTask.FailReason,
}
respBody, _ = json.Marshal(dto.TaskResponse[any]{
Code: "success",
Data: out,
})
}
}
}()
isOpenAIVideoAPI := strings.HasPrefix(c.Request.RequestURI, "/v1/videos/")
if len(respBody) != 0 {
// Gemini/Vertex 支持实时查询:用户 fetch 时直接从上游拉取最新状态
if realtimeResp := tryRealtimeFetch(originTask, isOpenAIVideoAPI); len(realtimeResp) > 0 {
respBody = realtimeResp
return
}
if strings.HasPrefix(c.Request.RequestURI, "/v1/videos/") {
// OpenAI Video API 格式: 走各 adaptor 的 ConvertToOpenAIVideo
if isOpenAIVideoAPI {
adaptor := GetTaskAdaptor(originTask.Platform)
if adaptor == nil {
taskResp = service.TaskErrorWrapperLocal(fmt.Errorf("invalid channel id: %d", originTask.ChannelId), "invalid_channel_id", http.StatusBadRequest)
@@ -486,10 +331,12 @@ func videoFetchByIDRespBodyBuilder(c *gin.Context) (respBody []byte, taskResp *d
respBody = openAIVideoData
return
}
taskResp = service.TaskErrorWrapperLocal(errors.New(fmt.Sprintf("not_implemented:%s", originTask.Platform)), "not_implemented", http.StatusNotImplemented)
taskResp = service.TaskErrorWrapperLocal(fmt.Errorf("not_implemented:%s", originTask.Platform), "not_implemented", http.StatusNotImplemented)
return
}
respBody, err = json.Marshal(dto.TaskResponse[any]{
// 通用 TaskDto 格式
respBody, err = common.Marshal(dto.TaskResponse[any]{
Code: "success",
Data: TaskModel2Dto(originTask),
})
@@ -499,16 +346,145 @@ func videoFetchByIDRespBodyBuilder(c *gin.Context) (respBody []byte, taskResp *d
return
}
// tryRealtimeFetch 尝试从上游实时拉取 Gemini/Vertex 任务状态。
// 仅当渠道类型为 Gemini 或 Vertex 时触发;其他渠道或出错时返回 nil。
// 当非 OpenAI Video API 时,还会构建自定义格式的响应体。
func tryRealtimeFetch(task *model.Task, isOpenAIVideoAPI bool) []byte {
channelModel, err := model.GetChannelById(task.ChannelId, true)
if err != nil {
return nil
}
if channelModel.Type != constant.ChannelTypeVertexAi && channelModel.Type != constant.ChannelTypeGemini {
return nil
}
baseURL := constant.ChannelBaseURLs[channelModel.Type]
if channelModel.GetBaseURL() != "" {
baseURL = channelModel.GetBaseURL()
}
proxy := channelModel.GetSetting().Proxy
adaptor := GetTaskAdaptor(constant.TaskPlatform(strconv.Itoa(channelModel.Type)))
if adaptor == nil {
return nil
}
resp, err := adaptor.FetchTask(baseURL, channelModel.Key, map[string]any{
"task_id": task.GetUpstreamTaskID(),
"action": task.Action,
}, proxy)
if err != nil || resp == nil {
return nil
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil
}
ti, err := adaptor.ParseTaskResult(body)
if err != nil || ti == nil {
return nil
}
// 将上游最新状态更新到 task
if ti.Status != "" {
task.Status = model.TaskStatus(ti.Status)
}
if ti.Progress != "" {
task.Progress = ti.Progress
}
if strings.HasPrefix(ti.Url, "data:") {
// data: URI — kept in Data, not ResultURL
} else if ti.Url != "" {
task.PrivateData.ResultURL = ti.Url
} else if task.Status == model.TaskStatusSuccess {
// No URL from adaptor — construct proxy URL using public task ID
task.PrivateData.ResultURL = taskcommon.BuildProxyURL(task.TaskID)
}
_ = task.Update()
// OpenAI Video API 由调用者的 ConvertToOpenAIVideo 分支处理
if isOpenAIVideoAPI {
return nil
}
// 非 OpenAI Video API: 构建自定义格式响应
format := detectVideoFormat(body)
out := map[string]any{
"error": nil,
"format": format,
"metadata": nil,
"status": mapTaskStatusToSimple(task.Status),
"task_id": task.TaskID,
"url": task.GetResultURL(),
}
respBody, _ := common.Marshal(dto.TaskResponse[any]{
Code: "success",
Data: out,
})
return respBody
}
// detectVideoFormat 从 Gemini/Vertex 原始响应中探测视频格式
func detectVideoFormat(rawBody []byte) string {
var raw map[string]any
if err := common.Unmarshal(rawBody, &raw); err != nil {
return "mp4"
}
respObj, ok := raw["response"].(map[string]any)
if !ok {
return "mp4"
}
vids, ok := respObj["videos"].([]any)
if !ok || len(vids) == 0 {
return "mp4"
}
v0, ok := vids[0].(map[string]any)
if !ok {
return "mp4"
}
mt, ok := v0["mimeType"].(string)
if !ok || mt == "" || strings.Contains(mt, "mp4") {
return "mp4"
}
return mt
}
// mapTaskStatusToSimple 将内部 TaskStatus 映射为简化状态字符串
func mapTaskStatusToSimple(status model.TaskStatus) string {
switch status {
case model.TaskStatusSuccess:
return "succeeded"
case model.TaskStatusFailure:
return "failed"
case model.TaskStatusQueued, model.TaskStatusSubmitted:
return "queued"
default:
return "processing"
}
}
func TaskModel2Dto(task *model.Task) *dto.TaskDto {
return &dto.TaskDto{
ID: task.ID,
CreatedAt: task.CreatedAt,
UpdatedAt: task.UpdatedAt,
TaskID: task.TaskID,
Platform: string(task.Platform),
UserId: task.UserId,
Group: task.Group,
ChannelId: task.ChannelId,
Quota: task.Quota,
Action: task.Action,
Status: string(task.Status),
FailReason: task.FailReason,
ResultURL: task.GetResultURL(),
SubmitTime: task.SubmitTime,
StartTime: task.StartTime,
FinishTime: task.FinishTime,
Progress: task.Progress,
Properties: task.Properties,
Username: task.Username,
Data: task.Data,
}
}