mirror of
https://gh-proxy.org/https://github.com/klinsmaya/new-api.git
synced 2026-11-08 01:17:09 +08:00
* fix(task): honor channel header overrides for submission and polling Apply the existing Request Header Override configuration after task plugin headers, including per-task, batch and realtime queries. Resolve api_key from the credential used for that request and skip client-only placeholders when polling has no original client request. Reuse the HTTP header resolver without changing plugin contracts or adding provider-specific settings. Keep polling redirects visible, matching submission and utils.fetch, so a custom credential header cannot be forwarded to a redirect destination. Document configuration and cover precedence, saved keys, absent client context, invalid values, redirects and all polling initialization paths. Validation with Go 1.26.4 and GOWORK=off: - make test: root, relaykit and tokenkit passed before the realtime follow-up - go test ./relay ./relay/channel ./relay/channel/task/jsplugin ./service -count=1: passed after the follow-up - go vet ./... and go build ./...: passed after the follow-up - Independent read-only review: passed Scoped security review used OWASP ASVS 5.0.0 and the Authentication, Session Management and REST Security cheat sheets. Credentials remain in headers, client headers are not persisted, and redirects are not followed. No live provider calls or production deployment were performed. https://owasp.org/www-project-application-security-verification-standard/ https://cheatsheetseries.owasp.org/cheatsheets/Authentication_Cheat_Sheet.html https://cheatsheetseries.owasp.org/cheatsheets/Session_Management_Cheat_Sheet.html https://cheatsheetseries.owasp.org/cheatsheets/REST_Security_Cheat_Sheet.html * fix(task): follow poll redirects and override content request headers Polling follows redirects again through the shared HTTP client, which checks every hop against the SSRF protection settings. Refusing them turned each poll of a redirecting upstream into a failure, and after enough failures the task was failed and refunded. Credentialed content requests now carry the channel header overrides, as submission and polling do. Credentialless links, utils.fetch and the Host header stay as the plugin builds them. * fix(task): settle tasks that a realtime query sees finish Fetching a Gemini or Vertex task queries the upstream directly. When that query saw the task finish, it wrote the final status itself, so the poller skipped the task and never settled or refunded it: a failed task kept its charge. Hand that transition to the poller's single-task update, which records the status and settles exactly once. --------- Co-authored-by: CaIon <[email protected]>
841 lines
30 KiB
Go
841 lines
30 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
taskdto "github.com/QuantumNous/new-api/dto"
|
|
"github.com/QuantumNous/new-api/logger"
|
|
"github.com/QuantumNous/new-api/model"
|
|
perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
|
|
"github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
|
|
"github.com/bytedance/gopkg/util/gopool"
|
|
"github.com/samber/lo"
|
|
)
|
|
|
|
// TaskPollingAdaptor 定义轮询所需的最小适配器接口,避免 service -> relay 的循环依赖
|
|
type TaskPollingAdaptor interface {
|
|
Init(info *relaycommon.RelayInfo)
|
|
FetchTask(baseURL string, key string, task *model.Task, proxy string) (*http.Response, error)
|
|
ParseTaskResult(task *model.Task, resp *http.Response, body []byte) (*relaycommon.TaskInfo, error)
|
|
// AdjustBillingOnComplete 在任务到达终态(成功/失败)时由轮询循环调用。
|
|
// 返回正数触发差额结算(补扣/退还),返回 0 保持预扣费金额不变。
|
|
AdjustBillingOnComplete(task *model.Task, taskResult *relaycommon.TaskInfo) int
|
|
}
|
|
|
|
type BatchTaskPollingAdaptor interface {
|
|
TaskPollingAdaptor
|
|
FetchMode() string
|
|
FetchBatchTasks(baseURL, key string, tasks []*model.Task, proxy string) (*http.Response, error)
|
|
ParseBatchResult(tasks []*model.Task, resp *http.Response, body []byte) (map[string]*BatchTaskResult, error)
|
|
}
|
|
|
|
const (
|
|
pollClassOK = "ok"
|
|
pollClassOtherClient = "other_client"
|
|
pollClassNotFound = "not_found"
|
|
pollClassAuth = "auth"
|
|
pollClassTransient = "transient"
|
|
pollClassUnrecognized = "unrecognized"
|
|
pollClassHookError = "hook_error"
|
|
pollClassTransport = "transport_error"
|
|
)
|
|
|
|
type BatchTaskResult struct {
|
|
TaskInfo relaycommon.TaskInfo
|
|
Action string
|
|
SubmitTime int64
|
|
StartTime int64
|
|
FinishTime int64
|
|
Data any
|
|
}
|
|
|
|
// GetTaskAdaptorFunc 由 main 包注入,用于获取指定平台的任务适配器。
|
|
// 打破 service -> relay -> relay/channel -> service 的循环依赖。
|
|
var GetTaskAdaptorFunc func(platform constant.TaskPlatform) TaskPollingAdaptor
|
|
|
|
// sweepTimedOutTasks 在主轮询之前独立清理超时任务。
|
|
// 每次最多处理 100 条,剩余的下个周期继续处理。
|
|
// 使用 per-task CAS (UpdateWithStatus) 防止覆盖被正常轮询已推进的任务。
|
|
func sweepTimedOutTasks(ctx context.Context) {
|
|
if constant.TaskTimeoutMinutes <= 0 {
|
|
return
|
|
}
|
|
cutoff := time.Now().Unix() - int64(constant.TaskTimeoutMinutes)*60
|
|
tasks := model.GetTimedOutUnfinishedTasks(cutoff, 100)
|
|
if len(tasks) == 0 {
|
|
return
|
|
}
|
|
|
|
reason := "Task timed out"
|
|
legacyReason := "Task timed out (a task left by the old system is not refunded; contact the administrator)"
|
|
now := time.Now().Unix()
|
|
timedOutCount := 0
|
|
|
|
for _, task := range tasks {
|
|
isLegacy := task.SubmitTime > 0 && task.SubmitTime < model.TaskRefundLegacyCutoff
|
|
|
|
oldStatus := task.Status
|
|
task.Status = model.TaskStatusFailure
|
|
task.Progress = "100%"
|
|
task.FinishTime = now
|
|
if isLegacy {
|
|
task.FailReason = legacyReason
|
|
// 旧系统任务明确不退款,随终态 CAS 一并清掉 quota,
|
|
// 避免留下可再次退款的计费状态。
|
|
task.Quota = 0
|
|
} else {
|
|
task.FailReason = reason
|
|
}
|
|
|
|
won, err := task.UpdateWithStatus(oldStatus)
|
|
if err != nil {
|
|
logger.LogError(ctx, common.LogText("sweepTimedOutTasks CAS update error for task %s: %v", task.TaskID, err))
|
|
continue
|
|
}
|
|
if !won {
|
|
logger.LogInfo(ctx, common.LogText("sweepTimedOutTasks: task %s already transitioned, skip", task.TaskID))
|
|
continue
|
|
}
|
|
timedOutCount++
|
|
if !isLegacy && task.Quota != 0 {
|
|
RefundTaskQuota(ctx, task, reason)
|
|
}
|
|
}
|
|
|
|
if timedOutCount > 0 {
|
|
logger.LogInfo(ctx, common.LogText("sweepTimedOutTasks: timed out %d tasks", timedOutCount))
|
|
}
|
|
}
|
|
|
|
// TaskPollSummary is the result recorded on an async_task_poll system task row,
|
|
// summarizing one polling pass.
|
|
type TaskPollSummary struct {
|
|
UnfinishedTasks int `json:"unfinished_tasks"`
|
|
PlatformsScanned int `json:"platforms_scanned"`
|
|
NullTasksFailed int `json:"null_tasks_failed"`
|
|
}
|
|
|
|
// RunTaskPollingOnce performs one async-task (Suno/video) polling pass
|
|
// synchronously. It honors ctx cancellation (the system-task runner cancels it
|
|
// when the lease is lost) and, when report is non-nil, reports progress as
|
|
// (processedPlatforms, totalPlatforms). It returns immediately if the task
|
|
// adaptor factory has not been wired yet, to avoid a nil call during startup.
|
|
func RunTaskPollingOnce(ctx context.Context, report func(processed, total int)) TaskPollSummary {
|
|
summary := TaskPollSummary{}
|
|
if GetTaskAdaptorFunc == nil {
|
|
return summary
|
|
}
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
|
|
common.SysLog(common.LogText("task progress polling started"))
|
|
sweepTimedOutTasks(ctx)
|
|
allTasks := model.GetAllUnFinishSyncTasks(constant.TaskQueryLimit)
|
|
summary.UnfinishedTasks = len(allTasks)
|
|
platformTask := make(map[constant.TaskPlatform][]*model.Task)
|
|
for _, t := range allTasks {
|
|
platformTask[t.Platform] = append(platformTask[t.Platform], t)
|
|
}
|
|
|
|
totalPlatforms := len(platformTask)
|
|
processedPlatforms := 0
|
|
for platform, tasks := range platformTask {
|
|
if ctx.Err() != nil {
|
|
break
|
|
}
|
|
if report != nil {
|
|
report(processedPlatforms, totalPlatforms)
|
|
}
|
|
processedPlatforms++
|
|
if len(tasks) == 0 {
|
|
continue
|
|
}
|
|
summary.PlatformsScanned++
|
|
taskChannelM := make(map[int][]string)
|
|
taskM := make(map[string]*model.Task)
|
|
nullTaskIds := make([]int64, 0)
|
|
for _, task := range tasks {
|
|
upstreamID := task.GetUpstreamTaskID()
|
|
if upstreamID == "" {
|
|
// 统计失败的未完成任务
|
|
nullTaskIds = append(nullTaskIds, task.ID)
|
|
continue
|
|
}
|
|
taskM[upstreamID] = task
|
|
taskChannelM[task.ChannelId] = append(taskChannelM[task.ChannelId], upstreamID)
|
|
}
|
|
if len(nullTaskIds) > 0 {
|
|
summary.NullTasksFailed += len(nullTaskIds)
|
|
err := model.TaskBulkUpdateByID(nullTaskIds, map[string]any{
|
|
"status": "FAILURE",
|
|
"progress": "100%",
|
|
})
|
|
if err != nil {
|
|
logger.LogError(ctx, common.LogText("Fix null task_id task error: %v", err))
|
|
} else {
|
|
logger.LogInfo(ctx, common.LogText("Fix null task_id task success: %v", nullTaskIds))
|
|
}
|
|
}
|
|
if len(taskChannelM) == 0 {
|
|
continue
|
|
}
|
|
|
|
DispatchPlatformUpdate(ctx, platform, taskChannelM, taskM)
|
|
}
|
|
if report != nil && ctx.Err() == nil {
|
|
report(totalPlatforms, totalPlatforms)
|
|
}
|
|
common.SysLog(common.LogText("task progress polling finished"))
|
|
return summary
|
|
}
|
|
|
|
// DispatchPlatformUpdate 按平台分发轮询更新
|
|
func DispatchPlatformUpdate(ctx context.Context, platform constant.TaskPlatform, taskChannelM map[int][]string, taskM map[string]*model.Task) {
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
if platform == constant.TaskPlatformMidjourney {
|
|
// MJ 轮询由其自身处理,这里预留入口
|
|
return
|
|
}
|
|
adaptor := GetTaskAdaptorFunc(platform)
|
|
if batchAdaptor, ok := adaptor.(BatchTaskPollingAdaptor); ok && batchAdaptor.FetchMode() == "batch" {
|
|
if err := UpdateBatchTasks(ctx, batchAdaptor, taskChannelM, taskM); err != nil {
|
|
common.SysLog(common.LogText("UpdateBatchTasks fail: %s", err))
|
|
}
|
|
return
|
|
}
|
|
if err := UpdateVideoTasks(ctx, platform, taskChannelM, taskM); err != nil {
|
|
common.SysLog(common.LogText("UpdateVideoTasks fail: %s", err))
|
|
}
|
|
}
|
|
|
|
func UpdateBatchTasks(ctx context.Context, adaptor BatchTaskPollingAdaptor, taskChannelM map[int][]string, taskM map[string]*model.Task) error {
|
|
for channelId, taskIds := range taskChannelM {
|
|
if ctx.Err() != nil {
|
|
return ctx.Err()
|
|
}
|
|
err := updateBatchTasks(ctx, adaptor, channelId, taskIds, taskM)
|
|
if err != nil {
|
|
logger.LogError(ctx, common.LogText("channel #%d failed to update async tasks: %s", channelId, err.Error()))
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func updateBatchTasks(ctx context.Context, adaptor BatchTaskPollingAdaptor, channelId int, taskIds []string, taskM map[string]*model.Task) error {
|
|
logger.LogInfo(ctx, common.LogText("channel #%d has %d unfinished tasks", channelId, len(taskIds)))
|
|
if ctx.Err() != nil {
|
|
return ctx.Err()
|
|
}
|
|
if len(taskIds) == 0 {
|
|
return nil
|
|
}
|
|
ch, err := model.CacheGetChannel(channelId)
|
|
if err != nil {
|
|
common.SysLog(common.LogText("CacheGetChannel: %v", err))
|
|
// Collect DB primary key IDs for bulk update (taskIds are upstream IDs, not task_id column values)
|
|
var failedIDs []int64
|
|
for _, upstreamID := range taskIds {
|
|
if t, ok := taskM[upstreamID]; ok {
|
|
failedIDs = append(failedIDs, t.ID)
|
|
}
|
|
}
|
|
err = model.TaskBulkUpdateByID(failedIDs, map[string]any{
|
|
"fail_reason": "Failed to get the channel information. Please contact the administrator",
|
|
"status": "FAILURE",
|
|
"progress": "100%",
|
|
})
|
|
if err != nil {
|
|
common.SysLog(common.LogText("UpdateSunoTask error: %v", err))
|
|
}
|
|
return err
|
|
}
|
|
proxy := ch.GetSetting().Proxy
|
|
baseURL := ch.GetBaseURL()
|
|
if baseURL == "" {
|
|
baseURL = constant.GetChannelBaseURL(ch.Type)
|
|
}
|
|
tasks := make([]*model.Task, 0, len(taskIds))
|
|
for _, upstreamID := range taskIds {
|
|
if task := taskM[upstreamID]; task != nil {
|
|
tasks = append(tasks, task)
|
|
}
|
|
}
|
|
// The channel type tells plugin adaptors whether the upstream is another
|
|
// New API gateway, the same signal submission derives from the request.
|
|
info := &relaycommon.RelayInfo{}
|
|
info.ChannelMeta = &relaycommon.ChannelMeta{
|
|
ChannelType: ch.Type,
|
|
ChannelId: ch.Id,
|
|
ChannelBaseUrl: baseURL,
|
|
ChannelSetting: ch.GetSetting(),
|
|
HeadersOverride: ch.GetHeaderOverride(),
|
|
}
|
|
info.ApiKey = ch.Key
|
|
adaptor.Init(info)
|
|
resp, err := adaptor.FetchBatchTasks(baseURL, ch.Key, tasks, proxy)
|
|
if err != nil {
|
|
common.SysLog(common.LogText("Get Task Do req error: %v", err))
|
|
return recordPollFailureForTasks(ctx, adaptor, tasks, pollClassTransport, 0, err.Error())
|
|
}
|
|
defer resp.Body.Close()
|
|
responseBody, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
common.SysLog(common.LogText("Get Suno Task parse body error: %v", err))
|
|
return recordPollFailureForTasks(ctx, adaptor, tasks, pollClassTransport, resp.StatusCode, err.Error())
|
|
}
|
|
switch classifyPollHTTP(resp.StatusCode) {
|
|
case pollClassNotFound:
|
|
for _, task := range tasks {
|
|
if task != nil {
|
|
logger.LogWarn(ctx, common.LogText("task %s poll %s (failures=%d, http=%d): %s", task.TaskID, pollClassNotFound, task.PrivateData.PollFailures, resp.StatusCode, unrecognizedPollDetail("", responseBody)))
|
|
}
|
|
}
|
|
return failTasksFromPoll(ctx, adaptor, tasks, "Upstream task not found")
|
|
case pollClassAuth:
|
|
logger.LogWarn(ctx, common.LogText("task poll auth failure channel_id=%d http=%d", channelId, resp.StatusCode))
|
|
return recordPollFailureForTasks(ctx, adaptor, tasks, pollClassAuth, resp.StatusCode, "")
|
|
case pollClassTransient:
|
|
return recordPollFailureForTasks(ctx, adaptor, tasks, pollClassTransient, resp.StatusCode, "")
|
|
}
|
|
responseItems, err := adaptor.ParseBatchResult(tasks, resp, responseBody)
|
|
if err != nil {
|
|
return recordPollFailureForTasks(ctx, adaptor, tasks, pollClassHookError, resp.StatusCode, err.Error())
|
|
}
|
|
for upstreamID, responseItem := range responseItems {
|
|
if ctx.Err() != nil {
|
|
return ctx.Err()
|
|
}
|
|
task := taskM[upstreamID]
|
|
if task == nil {
|
|
logger.LogWarn(ctx, common.LogText("Batch task response ignored: unknown task_id=%s", upstreamID))
|
|
continue
|
|
}
|
|
snap := task.Snapshot()
|
|
httpClass := classifyPollHTTP(resp.StatusCode)
|
|
parsedStatus := model.TaskStatus(responseItem.TaskInfo.Status)
|
|
if parsedStatus == model.TaskStatusUnknown || parsedStatus == "" || !knownPollStatus(parsedStatus) {
|
|
if err := recordPollFailure(ctx, adaptor, task, snap.Status, pollClassUnrecognized, resp.StatusCode, responseItem.TaskInfo.Reason); err != nil {
|
|
common.SysLog(common.LogText("UpdateSunoTask task error: %s", err.Error()))
|
|
}
|
|
continue
|
|
}
|
|
if httpClass == pollClassOtherClient && isNonTerminalPollStatus(parsedStatus) {
|
|
if err := recordPollFailure(ctx, adaptor, task, snap.Status, pollClassUnrecognized, resp.StatusCode, responseItem.TaskInfo.Reason); err != nil {
|
|
common.SysLog(common.LogText("UpdateSunoTask task error: %s", err.Error()))
|
|
}
|
|
continue
|
|
}
|
|
if isNonTerminalPollStatus(parsedStatus) {
|
|
task.PrivateData.PollFailures = 0
|
|
}
|
|
if len(responseItem.TaskInfo.PluginState) > 0 {
|
|
task.PrivateData.PluginState = responseItem.TaskInfo.PluginState
|
|
}
|
|
task.Status = lo.If(parsedStatus != "", parsedStatus).Else(task.Status)
|
|
task.FailReason = lo.If(responseItem.TaskInfo.Reason != "", responseItem.TaskInfo.Reason).Else(task.FailReason)
|
|
task.SubmitTime = lo.If(responseItem.SubmitTime != 0, responseItem.SubmitTime).Else(task.SubmitTime)
|
|
task.StartTime = lo.If(responseItem.StartTime != 0, responseItem.StartTime).Else(task.StartTime)
|
|
task.FinishTime = lo.If(responseItem.FinishTime != 0, responseItem.FinishTime).Else(task.FinishTime)
|
|
if responseItem.TaskInfo.Progress != "" {
|
|
task.Progress = responseItem.TaskInfo.Progress
|
|
}
|
|
if responseItem.TaskInfo.Reason != "" || task.Status == model.TaskStatusFailure {
|
|
logger.LogInfo(ctx, common.LogText("%s failed: %s", task.TaskID, task.FailReason))
|
|
task.Status = model.TaskStatusFailure
|
|
task.Progress = "100%"
|
|
}
|
|
if responseItem.TaskInfo.Status == model.TaskStatusSuccess {
|
|
task.Progress = "100%"
|
|
}
|
|
if responseItem.Data != nil {
|
|
task.SetData(responseItem.Data)
|
|
} else if task.Status == model.TaskStatusSuccess || task.Status == model.TaskStatusFailure {
|
|
logger.LogWarn(ctx, common.LogText(
|
|
"Batch task %s reached terminal status without data; preserving existing task data",
|
|
task.TaskID,
|
|
))
|
|
}
|
|
if responseItem.TaskInfo.Url != "" {
|
|
task.PrivateData.ResultURL = responseItem.TaskInfo.Url
|
|
}
|
|
|
|
isDone := task.Status == model.TaskStatusSuccess || task.Status == model.TaskStatusFailure
|
|
terminalTransition := isDone && snap.Status != task.Status
|
|
won, updateErr := task.UpdateWithStatus(snap.Status)
|
|
if updateErr != nil {
|
|
common.SysLog(common.LogText("UpdateSunoTask task error: %s", updateErr.Error()))
|
|
continue
|
|
}
|
|
if !won {
|
|
logger.LogWarn(ctx, common.LogText("Batch task %s already transitioned by another process, skip billing", task.TaskID))
|
|
continue
|
|
}
|
|
if terminalTransition {
|
|
finalizeTerminalTask(ctx, adaptor, task, &responseItem.TaskInfo)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// UpdateVideoTasks 按渠道更新所有视频任务
|
|
func UpdateVideoTasks(ctx context.Context, platform constant.TaskPlatform, taskChannelM map[int][]string, taskM map[string]*model.Task) error {
|
|
channelIDs := make([]int, 0, len(taskChannelM))
|
|
for channelID := range taskChannelM {
|
|
channelIDs = append(channelIDs, channelID)
|
|
}
|
|
sort.Ints(channelIDs)
|
|
|
|
var wg sync.WaitGroup
|
|
for _, channelId := range channelIDs {
|
|
taskIds := taskChannelM[channelId]
|
|
if len(taskIds) == 0 {
|
|
continue
|
|
}
|
|
taskIds = append([]string(nil), taskIds...)
|
|
|
|
wg.Add(1)
|
|
gopool.Go(func() {
|
|
defer wg.Done()
|
|
if err := updateVideoTasks(ctx, platform, channelId, taskIds, taskM); err != nil {
|
|
logger.LogError(ctx, common.LogText("Channel #%d failed to update video async tasks: %s", channelId, err.Error()))
|
|
}
|
|
})
|
|
}
|
|
wg.Wait()
|
|
if ctx.Err() != nil {
|
|
return ctx.Err()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func updateVideoTasks(ctx context.Context, platform constant.TaskPlatform, channelId int, taskIds []string, taskM map[string]*model.Task) error {
|
|
logger.LogInfo(ctx, common.LogText("Channel #%d pending video tasks: %d", channelId, len(taskIds)))
|
|
if ctx.Err() != nil {
|
|
return ctx.Err()
|
|
}
|
|
if len(taskIds) == 0 {
|
|
return nil
|
|
}
|
|
cacheGetChannel, err := model.CacheGetChannel(channelId)
|
|
if err != nil {
|
|
// Collect DB primary key IDs for bulk update (taskIds are upstream IDs, not task_id column values)
|
|
var failedIDs []int64
|
|
for _, upstreamID := range taskIds {
|
|
if t, ok := taskM[upstreamID]; ok {
|
|
failedIDs = append(failedIDs, t.ID)
|
|
}
|
|
}
|
|
errUpdate := model.TaskBulkUpdateByID(failedIDs, map[string]any{
|
|
"fail_reason": "Failed to get the channel information. Please contact the administrator",
|
|
"status": "FAILURE",
|
|
"progress": "100%",
|
|
})
|
|
if errUpdate != nil {
|
|
common.SysLog(common.LogText("UpdateVideoTask error: %v", errUpdate))
|
|
}
|
|
return fmt.Errorf("CacheGetChannel failed: %w", err)
|
|
}
|
|
adaptor := GetTaskAdaptorFunc(platform)
|
|
if adaptor == nil {
|
|
return fmt.Errorf("video adaptor not found")
|
|
}
|
|
info := &relaycommon.RelayInfo{}
|
|
info.ChannelMeta = &relaycommon.ChannelMeta{
|
|
ChannelType: cacheGetChannel.Type,
|
|
ChannelId: cacheGetChannel.Id,
|
|
ChannelBaseUrl: cacheGetChannel.GetBaseURL(),
|
|
ChannelSetting: cacheGetChannel.GetSetting(),
|
|
HeadersOverride: cacheGetChannel.GetHeaderOverride(),
|
|
}
|
|
info.ApiKey = cacheGetChannel.Key
|
|
adaptor.Init(info)
|
|
disablePollingSleep := cacheGetChannel.GetOtherSettings().DisableTaskPollingSleep
|
|
for i, taskId := range taskIds {
|
|
if ctx.Err() != nil {
|
|
return ctx.Err()
|
|
}
|
|
if err := updateVideoSingleTask(ctx, adaptor, cacheGetChannel, taskId, taskM); err != nil {
|
|
logger.LogError(ctx, common.LogText("Failed to update video task %s: %s", taskId, err.Error()))
|
|
}
|
|
if disablePollingSleep || i == len(taskIds)-1 {
|
|
continue
|
|
}
|
|
|
|
// sleep 1 second between tasks for this channel only.
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case <-time.After(1 * time.Second):
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch *model.Channel, taskId string, taskM map[string]*model.Task) error {
|
|
if ctx.Err() != nil {
|
|
return ctx.Err()
|
|
}
|
|
baseURL := constant.GetChannelBaseURL(ch.Type)
|
|
if ch.GetBaseURL() != "" {
|
|
baseURL = ch.GetBaseURL()
|
|
}
|
|
proxy := ch.GetSetting().Proxy
|
|
|
|
task := taskM[taskId]
|
|
if task == nil {
|
|
logger.LogError(ctx, common.LogText("Task %s not found in taskM", taskId))
|
|
return fmt.Errorf("task %s not found", taskId)
|
|
}
|
|
key := ch.Key
|
|
|
|
privateData := task.PrivateData
|
|
if privateData.Key != "" {
|
|
key = privateData.Key
|
|
}
|
|
snap := task.Snapshot()
|
|
resp, err := adaptor.FetchTask(baseURL, key, task, proxy)
|
|
if err != nil {
|
|
return recordPollFailure(ctx, adaptor, task, snap.Status, pollClassTransport, 0, err.Error())
|
|
}
|
|
defer resp.Body.Close()
|
|
responseBody, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
return recordPollFailure(ctx, adaptor, task, snap.Status, pollClassTransport, resp.StatusCode, err.Error())
|
|
}
|
|
|
|
logger.LogDebug(ctx, "updateVideoSingleTask response: %s", responseBody)
|
|
|
|
switch classifyPollHTTP(resp.StatusCode) {
|
|
case pollClassNotFound:
|
|
logger.LogWarn(ctx, common.LogText("task %s poll %s (failures=%d, http=%d): %s", task.TaskID, pollClassNotFound, task.PrivateData.PollFailures, resp.StatusCode, unrecognizedPollDetail("", responseBody)))
|
|
return failTaskFromPoll(ctx, adaptor, task, snap.Status, "Upstream task not found")
|
|
case pollClassAuth:
|
|
logger.LogWarn(ctx, common.LogText("task poll auth failure channel_id=%d task=%s http=%d", ch.Id, task.TaskID, resp.StatusCode))
|
|
return recordPollFailure(ctx, adaptor, task, snap.Status, pollClassAuth, resp.StatusCode, "")
|
|
case pollClassTransient:
|
|
return recordPollFailure(ctx, adaptor, task, snap.Status, pollClassTransient, resp.StatusCode, "")
|
|
}
|
|
|
|
taskResult := &relaycommon.TaskInfo{}
|
|
// try parse as New API response format
|
|
var responseItems taskdto.TaskResponse[model.Task]
|
|
if err = common.Unmarshal(responseBody, &responseItems); err == nil && responseItems.IsSuccess() {
|
|
logger.LogDebug(ctx, "updateVideoSingleTask parsed as new api response format: %+v", responseItems)
|
|
t := responseItems.Data
|
|
taskResult.TaskID = t.TaskID
|
|
taskResult.Status = string(t.Status)
|
|
taskResult.Url = t.GetResultURL()
|
|
taskResult.Progress = t.Progress
|
|
taskResult.Reason = t.FailReason
|
|
task.Data = t.Data
|
|
} else if taskResult, err = adaptor.ParseTaskResult(task, resp, responseBody); err != nil {
|
|
return recordPollFailure(ctx, adaptor, task, snap.Status, pollClassHookError, resp.StatusCode, err.Error())
|
|
}
|
|
|
|
logger.LogDebug(ctx, "updateVideoSingleTask taskResult: %+v", taskResult)
|
|
|
|
parsedStatus := model.TaskStatus(taskResult.Status)
|
|
if parsedStatus == model.TaskStatusUnknown || parsedStatus == "" || !knownPollStatus(parsedStatus) {
|
|
return recordPollFailure(ctx, adaptor, task, snap.Status, pollClassUnrecognized, resp.StatusCode, unrecognizedPollDetail(taskResult.Reason, responseBody))
|
|
}
|
|
if classifyPollHTTP(resp.StatusCode) == pollClassOtherClient && isNonTerminalPollStatus(parsedStatus) {
|
|
return recordPollFailure(ctx, adaptor, task, snap.Status, pollClassUnrecognized, resp.StatusCode, unrecognizedPollDetail(taskResult.Reason, responseBody))
|
|
}
|
|
|
|
task.Data = redactVideoResponseBody(responseBody)
|
|
if len(taskResult.PluginState) > 0 {
|
|
task.PrivateData.PluginState = taskResult.PluginState
|
|
}
|
|
if isNonTerminalPollStatus(parsedStatus) {
|
|
task.PrivateData.PollFailures = 0
|
|
}
|
|
|
|
now := time.Now().Unix()
|
|
shouldFinalizeBilling := false
|
|
|
|
task.Status = parsedStatus
|
|
switch parsedStatus {
|
|
case model.TaskStatusSubmitted:
|
|
task.Progress = taskcommon.ProgressSubmitted
|
|
case model.TaskStatusQueued:
|
|
task.Progress = taskcommon.ProgressQueued
|
|
case model.TaskStatusInProgress:
|
|
task.Progress = taskcommon.ProgressInProgress
|
|
if task.StartTime == 0 {
|
|
task.StartTime = now
|
|
}
|
|
case model.TaskStatusSuccess:
|
|
task.Progress = taskcommon.ProgressComplete
|
|
if task.FinishTime == 0 {
|
|
task.FinishTime = now
|
|
}
|
|
if strings.HasPrefix(taskResult.Url, "data:") {
|
|
// data: URI (e.g. Vertex base64 encoded video) — keep in Data, not in ResultURL
|
|
task.PrivateData.ResultURL = taskcommon.BuildProxyURL(task.TaskID)
|
|
} else if taskResult.Url != "" {
|
|
// Direct upstream URL (e.g. Kling, Ali, Doubao, etc.)
|
|
task.PrivateData.ResultURL = taskResult.Url
|
|
} else {
|
|
// No URL from adaptor — construct proxy URL using public task ID
|
|
task.PrivateData.ResultURL = taskcommon.BuildProxyURL(task.TaskID)
|
|
}
|
|
shouldFinalizeBilling = true
|
|
case model.TaskStatusFailure:
|
|
logger.LogJson(ctx, common.LogText("Task %s failed", taskId), task)
|
|
task.Status = model.TaskStatusFailure
|
|
task.Progress = taskcommon.ProgressComplete
|
|
if task.FinishTime == 0 {
|
|
task.FinishTime = now
|
|
}
|
|
task.FailReason = taskResult.Reason
|
|
logger.LogInfo(ctx, common.LogText("Task %s failed: %s", task.TaskID, task.FailReason))
|
|
taskResult.Progress = taskcommon.ProgressComplete
|
|
shouldFinalizeBilling = true
|
|
}
|
|
if taskResult.Progress != "" {
|
|
task.Progress = taskResult.Progress
|
|
}
|
|
|
|
isDone := task.Status == model.TaskStatusSuccess || task.Status == model.TaskStatusFailure
|
|
if isDone && snap.Status != task.Status {
|
|
won, err := task.UpdateWithStatus(snap.Status)
|
|
if err != nil {
|
|
logger.LogError(ctx, common.LogText("UpdateWithStatus failed for task %s: %s", task.TaskID, err.Error()))
|
|
shouldFinalizeBilling = false
|
|
} else if !won {
|
|
logger.LogWarn(ctx, common.LogText("Task %s CAS lost or no-op update, skip billing", task.TaskID))
|
|
shouldFinalizeBilling = false
|
|
}
|
|
} else if !snap.Equal(task.Snapshot()) {
|
|
if _, err := task.UpdateWithStatus(snap.Status); err != nil {
|
|
logger.LogError(ctx, common.LogText("Failed to update task %s: %s", task.TaskID, err.Error()))
|
|
}
|
|
} else {
|
|
// No changes, skip update
|
|
logger.LogDebug(ctx, "No update needed for task %s", task.TaskID)
|
|
}
|
|
|
|
if shouldFinalizeBilling {
|
|
finalizeTerminalTask(ctx, adaptor, task, taskResult)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// finalizeTerminalTask 终态统一收尾(状态 CAS 赢家调用,恰好一次):采样 + 结算 + 失败兜底退款。
|
|
func finalizeTerminalTask(ctx context.Context, adaptor TaskPollingAdaptor, task *model.Task, taskResult *relaycommon.TaskInfo) {
|
|
perfmetrics.RecordTaskResult(task, taskResult)
|
|
billingSettled := settleTaskBillingOnComplete(ctx, adaptor, task, taskResult)
|
|
if task.Status == model.TaskStatusFailure && !billingSettled && task.Quota != 0 {
|
|
RefundTaskQuota(ctx, task, task.FailReason)
|
|
}
|
|
}
|
|
|
|
func redactVideoResponseBody(body []byte) []byte {
|
|
var m map[string]any
|
|
if err := common.Unmarshal(body, &m); err != nil {
|
|
return body
|
|
}
|
|
resp, _ := m["response"].(map[string]any)
|
|
if resp != nil {
|
|
delete(resp, "bytesBase64Encoded")
|
|
if v, ok := resp["video"].(string); ok {
|
|
resp["video"] = truncateBase64(v)
|
|
}
|
|
if vs, ok := resp["videos"].([]any); ok {
|
|
for i := range vs {
|
|
if vm, ok := vs[i].(map[string]any); ok {
|
|
delete(vm, "bytesBase64Encoded")
|
|
}
|
|
}
|
|
}
|
|
}
|
|
b, err := common.Marshal(m)
|
|
if err != nil {
|
|
return body
|
|
}
|
|
return b
|
|
}
|
|
|
|
func truncateBase64(s string) string {
|
|
const maxKeep = 256
|
|
if len(s) <= maxKeep {
|
|
return s
|
|
}
|
|
return s[:maxKeep] + "..."
|
|
}
|
|
|
|
// settleTaskBillingOnComplete 任务完成时的统一计费调整。
|
|
// 返回 true 表示用量结算路径已接管最终计费;失败任务仅在返回 false 时补做全额退款。
|
|
// 优先级:1. tiered snapshot → 2. adaptor 调整 → 3. token 重算。
|
|
//
|
|
// 表达式求值失败会保留预扣额度,因此也视为已接管,避免错误全退。
|
|
func settleTaskBillingOnComplete(ctx context.Context, adaptor TaskPollingAdaptor, task *model.Task, taskResult *relaycommon.TaskInfo) bool {
|
|
if bc := task.PrivateData.BillingContext; bc != nil && bc.TieredSnapshot != nil {
|
|
// 用量表达式结算只适用于成功任务;失败任务由调用方全额退款。
|
|
if task.Status == model.TaskStatusFailure {
|
|
return false
|
|
}
|
|
result, usageFacts, err := EvaluateTaskCompletionUsage(bc.TieredSnapshot, taskResult.UsageFacts)
|
|
if err != nil {
|
|
logger.LogWarn(ctx, common.LogText("task %s expression settlement failed, keeping the pre-consumed quota: %v", task.TaskID, err))
|
|
return true
|
|
}
|
|
if result.Clamp != nil {
|
|
logger.LogWarn(ctx, common.LogText("task %s expression settlement quota was clamped: %+v", task.TaskID, result.Clamp))
|
|
}
|
|
bc.TieredSnapshot.UsageFacts = usageFacts
|
|
bc.TieredSnapshot.EstimatedTier = result.MatchedTier
|
|
RecalculateTaskQuota(ctx, task, result.ActualQuotaAfterGroup, common.NewMessage("Settled by the task usage expression"), result.Clamp)
|
|
return true
|
|
}
|
|
// 按次计费的成功任务保持预扣;失败任务由调用方全额退款。
|
|
if bc := task.PrivateData.BillingContext; bc != nil && bc.PerCallBilling {
|
|
logger.LogInfo(ctx, common.LogText("task %s is billed per request, skipping difference settlement", task.TaskID))
|
|
return false
|
|
}
|
|
// 优先让 adaptor 决定最终额度。
|
|
if actualQuota := adaptor.AdjustBillingOnComplete(task, taskResult); actualQuota > 0 {
|
|
RecalculateTaskQuota(ctx, task, actualQuota, common.NewMessage("Adjusted by the channel adaptor"))
|
|
return true
|
|
}
|
|
// 回退到 token 重算。
|
|
tokens := taskResult.TotalTokens
|
|
if tokens == 0 && taskResult.CompletionTokens > 0 {
|
|
tokens = taskResult.CompletionTokens
|
|
}
|
|
if tokens > 0 {
|
|
return RecalculateTaskQuotaByTokens(ctx, task, tokens)
|
|
}
|
|
return false
|
|
}
|
|
|
|
func classifyPollHTTP(statusCode int) string {
|
|
switch {
|
|
case statusCode >= 200 && statusCode < 300:
|
|
return pollClassOK
|
|
case statusCode == http.StatusNotFound || statusCode == http.StatusGone:
|
|
return pollClassNotFound
|
|
case statusCode == http.StatusUnauthorized || statusCode == http.StatusForbidden:
|
|
return pollClassAuth
|
|
case statusCode == http.StatusTooManyRequests || statusCode >= 500:
|
|
return pollClassTransient
|
|
case statusCode >= 400 && statusCode < 500:
|
|
return pollClassOtherClient
|
|
default:
|
|
return pollClassTransient
|
|
}
|
|
}
|
|
|
|
func knownPollStatus(status model.TaskStatus) bool {
|
|
switch status {
|
|
case model.TaskStatusNotStart, model.TaskStatusSubmitted, model.TaskStatusQueued, model.TaskStatusInProgress, model.TaskStatusSuccess, model.TaskStatusFailure:
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func isNonTerminalPollStatus(status model.TaskStatus) bool {
|
|
switch status {
|
|
case model.TaskStatusNotStart, model.TaskStatusSubmitted, model.TaskStatusQueued, model.TaskStatusInProgress:
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
// unrecognizedPollDetail pairs the plugin's reason with a bounded copy of the
|
|
// upstream body so the WARN line is enough to diagnose a parser gap.
|
|
func unrecognizedPollDetail(reason string, body []byte) string {
|
|
const maxBodyChars = 512
|
|
redacted := string(redactVideoResponseBody(body))
|
|
if len(redacted) > maxBodyChars {
|
|
redacted = redacted[:maxBodyChars] + "…"
|
|
}
|
|
if strings.TrimSpace(reason) == "" {
|
|
return "body=" + redacted
|
|
}
|
|
return reason + "; body=" + redacted
|
|
}
|
|
|
|
func recordPollFailure(ctx context.Context, adaptor TaskPollingAdaptor, task *model.Task, fromStatus model.TaskStatus, class string, statusCode int, detail string) error {
|
|
task.PrivateData.PollFailures++
|
|
// TASK_POLL_MAX_FAILURES <= 0 disables the consecutive-failure cutoff, matching
|
|
// TASK_TIMEOUT_MINUTES semantics; the 24h sweep remains the only backstop.
|
|
giveUp := constant.TaskPollMaxFailures > 0 && task.PrivateData.PollFailures >= constant.TaskPollMaxFailures
|
|
if giveUp || class == pollClassUnrecognized || class == pollClassHookError {
|
|
// The redacted body is intentionally not persisted to Task.Data on these
|
|
// paths, and the fail reason of a task that is given up is a fixed
|
|
// sentence, so the WARN line is the only operator-visible copy of the
|
|
// class, the status and what the plugin could not interpret.
|
|
logger.LogWarn(ctx, common.LogText("task %s poll %s (failures=%d, http=%d): %s", task.TaskID, class, task.PrivateData.PollFailures, statusCode, detail))
|
|
}
|
|
if giveUp {
|
|
return failTaskFromPoll(ctx, adaptor, task, fromStatus, "Failed to query the upstream task status")
|
|
}
|
|
if _, err := task.UpdateWithStatus(fromStatus); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func recordPollFailureForTasks(ctx context.Context, adaptor TaskPollingAdaptor, tasks []*model.Task, class string, statusCode int, detail string) error {
|
|
var firstErr error
|
|
for _, task := range tasks {
|
|
if task == nil {
|
|
continue
|
|
}
|
|
if err := recordPollFailure(ctx, adaptor, task, task.Status, class, statusCode, detail); err != nil && firstErr == nil {
|
|
firstErr = err
|
|
}
|
|
}
|
|
return firstErr
|
|
}
|
|
|
|
func failTaskFromPoll(ctx context.Context, adaptor TaskPollingAdaptor, task *model.Task, fromStatus model.TaskStatus, reason string) error {
|
|
now := time.Now().Unix()
|
|
task.Status = model.TaskStatusFailure
|
|
task.Progress = taskcommon.ProgressComplete
|
|
if task.FinishTime == 0 {
|
|
task.FinishTime = now
|
|
}
|
|
task.FailReason = reason
|
|
won, err := task.UpdateWithStatus(fromStatus)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !won {
|
|
return nil
|
|
}
|
|
finalizeTerminalTask(ctx, adaptor, task, relaycommon.FailTaskInfo(reason))
|
|
return nil
|
|
}
|
|
|
|
func failTasksFromPoll(ctx context.Context, adaptor TaskPollingAdaptor, tasks []*model.Task, reason string) error {
|
|
var firstErr error
|
|
for _, task := range tasks {
|
|
if task == nil {
|
|
continue
|
|
}
|
|
if err := failTaskFromPoll(ctx, adaptor, task, task.Status, reason); err != nil && firstErr == nil {
|
|
firstErr = err
|
|
}
|
|
}
|
|
return firstErr
|
|
}
|