Files
WingStoneandCaIon 9a42033a61 fix(task): 修复任务提交与轮询的请求头覆盖 (#7918)
* 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]>
2026-10-20 20:50:02 +08:00

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
}