688 lines
25 KiB
Go
688 lines
25 KiB
Go
package task
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"io"
|
||
"net/http"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
|
||
"github.com/qianfree/team-api/internal/logic/billing"
|
||
"github.com/qianfree/team-api/internal/logic/monitor"
|
||
"github.com/qianfree/team-api/internal/logic/relay"
|
||
"github.com/qianfree/team-api/relay/common"
|
||
"github.com/qianfree/team-api/relay/constant"
|
||
"github.com/qianfree/team-api/relay/helper"
|
||
"github.com/qianfree/team-api/relay/taskchannel"
|
||
)
|
||
|
||
const (
|
||
pollInterval = 15 * time.Second
|
||
taskTimeout = 30 * time.Minute
|
||
pollBatchSize = 500
|
||
timeoutBatchSize = 100
|
||
)
|
||
|
||
var (
|
||
pollingWg sync.WaitGroup
|
||
pollStop chan struct{}
|
||
)
|
||
|
||
// isSyncImageTask 判断任务是否为 sync_image 类型（通过 private_data.task_type 标识）
|
||
func isSyncImageTask(t *common.AsyncTask) bool {
|
||
if t.PrivateData == nil {
|
||
return false
|
||
}
|
||
var privateData map[string]any
|
||
if err := json.Unmarshal(t.PrivateData, &privateData); err != nil {
|
||
return false
|
||
}
|
||
taskType, ok := privateData["task_type"].(string)
|
||
if !ok {
|
||
return false
|
||
}
|
||
return taskType == string(constant.TaskPlatformSyncImage)
|
||
}
|
||
|
||
// StartAsyncPolling 启动异步任务轮询 goroutine
|
||
func StartAsyncPolling(ctx context.Context) {
|
||
g.Log().Info(ctx, "Starting async task polling...")
|
||
pollStop = make(chan struct{})
|
||
pollingWg.Add(1)
|
||
go func() {
|
||
defer pollingWg.Done()
|
||
ticker := time.NewTicker(pollInterval)
|
||
defer ticker.Stop()
|
||
|
||
for {
|
||
select {
|
||
case <-pollStop:
|
||
g.Log().Info(ctx, "Async task polling stopped")
|
||
return
|
||
case <-ticker.C:
|
||
pollOnce(ctx)
|
||
}
|
||
}
|
||
}()
|
||
}
|
||
|
||
// StopAsyncPolling 通知轮询 goroutine 停止并等待其退出，在服务停机时调用。
|
||
func StopAsyncPolling() {
|
||
close(pollStop)
|
||
pollingWg.Wait()
|
||
}
|
||
|
||
// pollOnce 执行一次轮询
|
||
func pollOnce(ctx context.Context) {
|
||
if !HasActiveTasks() {
|
||
return
|
||
}
|
||
|
||
// 1. 处理超时任务
|
||
handleTimedOutTasks(ctx)
|
||
|
||
// 2. 重试未结算的终态任务
|
||
handleUnsettledTasks(ctx)
|
||
|
||
// 3. 轮询非终态任务
|
||
tasks, err := DefaultAsyncProvider.GetNonTerminalTasks(ctx, pollBatchSize)
|
||
if err != nil {
|
||
g.Log().Warningf(ctx, "poll: get non-terminal tasks: %v", err)
|
||
return
|
||
}
|
||
|
||
if len(tasks) == 0 {
|
||
return
|
||
}
|
||
|
||
g.Log().Debugf(ctx, "poll: processing %d active tasks", len(tasks))
|
||
|
||
// 3. 按 platform 分组处理
|
||
platformGroups := make(map[string][]*common.AsyncTask)
|
||
for _, t := range tasks {
|
||
// sync_image 任务通过 private_data.task_type 标识，由 worker 池自行推进；此处跳过，
|
||
// 避免 processPlatformTasks → taskchannel.GetAdaptor 报错刷日志。
|
||
// 其超时/未结算仍由 handleTimedOutTasks / handleUnsettledTasks 兜底。
|
||
if isSyncImageTask(t) {
|
||
g.Log().Debugf(ctx, "poll: skipping sync_image task %s (channel_id=%d, platform=%s)",
|
||
t.PublicTaskID, t.ChannelID, t.Platform)
|
||
continue
|
||
}
|
||
// 记录 channel_id=0 的非 sync_image 任务（这些任务配置有误）
|
||
if t.ChannelID == 0 {
|
||
g.Log().Warningf(ctx, "poll: found task with channel_id=0 but not sync_image: public_id=%s, platform=%s, status=%s",
|
||
t.PublicTaskID, t.Platform, t.Status)
|
||
}
|
||
platformGroups[t.Platform] = append(platformGroups[t.Platform], t)
|
||
}
|
||
|
||
for platform, platformTasks := range platformGroups {
|
||
processPlatformTasks(ctx, platform, platformTasks)
|
||
}
|
||
}
|
||
|
||
// buildTimeoutFailure 把超时任务标记为失败，并返回**覆盖前**的真实状态供 CAS 谓词使用。
|
||
//
|
||
// 独立成纯函数以便回归测试。历史 bug：调用方先 `t.Status = "FAILURE"` 再把 `t.Status`
|
||
// 当作 UpdateTaskCAS 的 oldStatus 传入，使 CAS 变成 `WHERE status='FAILURE' SET status='FAILURE'`，
|
||
// 对任何非终态行匹配 0 行 → 超时兜底网从未生效（超时任务不退款、active 计数不减）。
|
||
func buildTimeoutFailure(t *common.AsyncTask, now time.Time) (oldStatus string) {
|
||
oldStatus = t.Status
|
||
t.Status = "FAILURE"
|
||
t.FailReason = "task timed out"
|
||
t.FinishTime = &now
|
||
return oldStatus
|
||
}
|
||
|
||
// handleTimedOutTasks 处理超时任务
|
||
func handleTimedOutTasks(ctx context.Context) {
|
||
cutoff := time.Now().Add(-taskTimeout).Unix()
|
||
tasks, err := DefaultAsyncProvider.GetTimedOutTasks(ctx, cutoff, timeoutBatchSize)
|
||
if err != nil {
|
||
g.Log().Warningf(ctx, "poll: get timed-out tasks: %v", err)
|
||
return
|
||
}
|
||
|
||
for _, t := range tasks {
|
||
now := time.Now()
|
||
// 必须用覆盖前的真实状态（IN_PROGRESS/QUEUED/SUBMITTED…）作为 CAS 谓词。
|
||
oldStatus := buildTimeoutFailure(t, now)
|
||
|
||
if err := DefaultAsyncProvider.UpdateTaskCAS(ctx, t, oldStatus); err != nil {
|
||
g.Log().Warningf(ctx, "poll: mark timeout task %s: %v", t.PublicTaskID, err)
|
||
continue
|
||
}
|
||
DecrActiveTask()
|
||
monitor.UnregisterRequestByTaskID(t.PublicTaskID)
|
||
|
||
// 退还预扣费用
|
||
if t.PreDeductAmount.GreaterThan(billing.Zero) {
|
||
taskBilling := billing.NewTaskBillingProvider()
|
||
if err := taskBilling.SettleTaskFailed(ctx, t.TenantID, t.RequestID, t.PreDeductAmount); err != nil {
|
||
g.Log().Warningf(ctx, "poll: refund timed-out task %s: %v", t.PublicTaskID, err)
|
||
}
|
||
billing.CleanupPreDeduct(ctx, t.TenantID, t.RequestID+"_adjust")
|
||
}
|
||
g.Log().Infof(ctx, "poll: task %s timed out", t.PublicTaskID)
|
||
|
||
// 查询渠道信息并记录用量日志
|
||
var ch *common.ChannelBasicInfo
|
||
if t.ChannelID > 0 {
|
||
ch, _ = DefaultAsyncProvider.GetChannelByID(ctx, t.ChannelID)
|
||
}
|
||
recordTaskUsage(t, ch, false, "task timed out", nil)
|
||
|
||
// 更新审计记录
|
||
recordTaskCompletionAudit(t, "TIMEOUT", "", nil)
|
||
}
|
||
}
|
||
|
||
// handleUnsettledTasks 重试终态但未结算的任务
|
||
func handleUnsettledTasks(ctx context.Context) {
|
||
tasks, err := DefaultAsyncProvider.GetUnsettledTasks(ctx, timeoutBatchSize)
|
||
if err != nil {
|
||
g.Log().Warningf(ctx, "poll: get unsettled tasks: %v", err)
|
||
return
|
||
}
|
||
|
||
for _, t := range tasks {
|
||
// sync_image 无 upstream_task_id，不能落入下面「UpstreamTaskID=="" → 退款」的分支，
|
||
// 否则崩溃窗口内的成功任务会被误退款（用户出图却全额退回，资损）。单独处理。
|
||
if isSyncImageTask(t) {
|
||
handleUnsettledSyncImage(ctx, t)
|
||
continue
|
||
}
|
||
|
||
var pd privateData
|
||
if err := json.Unmarshal(t.PrivateData, &pd); err != nil || pd.UpstreamTaskID == "" {
|
||
// 无法恢复，直接退还预扣
|
||
g.Log().Warningf(ctx, "poll: unsettled task %s has invalid private_data, refunding", t.PublicTaskID)
|
||
taskBilling := billing.NewTaskBillingProvider()
|
||
if err := taskBilling.SettleTaskFailed(ctx, t.TenantID, t.RequestID, t.PreDeductAmount); err != nil {
|
||
g.Log().Errorf(ctx, "poll: refund unsettled task %s: %v", t.PublicTaskID, err)
|
||
} else {
|
||
t.BillingSettled = true
|
||
DefaultAsyncProvider.UpdateTask(ctx, t)
|
||
}
|
||
billing.CleanupPreDeduct(ctx, t.TenantID, t.RequestID+"_adjust")
|
||
continue
|
||
}
|
||
|
||
if t.Status == "SUCCESS" {
|
||
// 成功任务：用 ActualCost（已在上次轮询中计算）结算
|
||
actualCost := t.ActualCost
|
||
if actualCost.LessThanOrEqual(billing.Zero) {
|
||
actualCost = t.PreDeductAmount
|
||
}
|
||
taskBilling := billing.NewTaskBillingProvider()
|
||
settleResult, err := taskBilling.SettleTaskSuccess(ctx, t.TenantID, t.UserID, t.ApiKeyID, t.ChannelID,
|
||
t.ModelName, t.RequestID, actualCost, t.PreDeductAmount,
|
||
0, 0, pd.BillingContext.Ratios, t.PublicTaskID, t.CreatedAt)
|
||
if err != nil {
|
||
g.Log().Warningf(ctx, "poll: retry settle task %s: %v", t.PublicTaskID, err)
|
||
} else {
|
||
t.BillingSettled = true
|
||
t.ActualCost = actualCost
|
||
// 幂等重复结算（DuplicateSkip）不得重复累加 Key 额度
|
||
if settleResult == nil || !settleResult.DuplicateSkip {
|
||
taskBilling.IncrApiKeyQuotaUsed(ctx, t.ApiKeyID, actualCost)
|
||
}
|
||
DefaultAsyncProvider.UpdateTask(ctx, t)
|
||
g.Log().Infof(ctx, "poll: retried settlement for task %s", t.PublicTaskID)
|
||
}
|
||
} else {
|
||
// 失败任务：退还预扣
|
||
taskBilling := billing.NewTaskBillingProvider()
|
||
if err := taskBilling.SettleTaskFailed(ctx, t.TenantID, t.RequestID, t.PreDeductAmount); err != nil {
|
||
g.Log().Warningf(ctx, "poll: retry refund task %s: %v", t.PublicTaskID, err)
|
||
} else {
|
||
t.BillingSettled = true
|
||
DefaultAsyncProvider.UpdateTask(ctx, t)
|
||
g.Log().Infof(ctx, "poll: retried refund for task %s", t.PublicTaskID)
|
||
}
|
||
billing.CleanupPreDeduct(ctx, t.TenantID, t.RequestID+"_adjust")
|
||
}
|
||
}
|
||
}
|
||
|
||
// handleUnsettledSyncImage 结算重试网中 sync_image 任务的专用处理。
|
||
//
|
||
// sync_image 无 upstream_task_id，通用分支会把它当「不可恢复」直接退款——但对 worker 已判
|
||
// SUCCESS、仅 billing_settled 未落库的崩溃窗口任务而言，退款即资损（用户出图却被全额退回）。
|
||
// 因此这里按终态区分：SUCCESS 用 private_data 的 billing_context 重放结算，FAILURE 才退款。
|
||
func handleUnsettledSyncImage(ctx context.Context, t *common.AsyncTask) {
|
||
taskBilling := billing.NewTaskBillingProvider()
|
||
|
||
if t.Status == "SUCCESS" {
|
||
var pd privateData
|
||
_ = json.Unmarshal(t.PrivateData, &pd)
|
||
actualCost := t.ActualCost
|
||
if actualCost.LessThanOrEqual(billing.Zero) {
|
||
actualCost = t.PreDeductAmount
|
||
}
|
||
settleResult, err := taskBilling.SettleTaskSuccess(ctx, t.TenantID, t.UserID, t.ApiKeyID, t.ChannelID,
|
||
t.ModelName, t.RequestID, actualCost, t.PreDeductAmount, 0, 0, pd.BillingContext.Ratios, t.PublicTaskID, t.CreatedAt)
|
||
if err != nil {
|
||
g.Log().Warningf(ctx, "poll: retry settle sync_image task %s: %v", t.PublicTaskID, err)
|
||
return
|
||
}
|
||
t.BillingSettled = true
|
||
t.ActualCost = actualCost
|
||
// 幂等重复结算（DuplicateSkip）不得重复累加 Key 额度
|
||
if settleResult == nil || !settleResult.DuplicateSkip {
|
||
taskBilling.IncrApiKeyQuotaUsed(ctx, t.ApiKeyID, actualCost)
|
||
}
|
||
DefaultAsyncProvider.UpdateTask(ctx, t)
|
||
g.Log().Infof(ctx, "poll: retried settlement for sync_image task %s", t.PublicTaskID)
|
||
return
|
||
}
|
||
|
||
// FAILURE：退还预扣
|
||
if err := taskBilling.SettleTaskFailed(ctx, t.TenantID, t.RequestID, t.PreDeductAmount); err != nil {
|
||
g.Log().Warningf(ctx, "poll: retry refund sync_image task %s: %v", t.PublicTaskID, err)
|
||
return
|
||
}
|
||
t.BillingSettled = true
|
||
DefaultAsyncProvider.UpdateTask(ctx, t)
|
||
billing.CleanupPreDeduct(ctx, t.TenantID, t.RequestID+"_adjust")
|
||
}
|
||
|
||
// processPlatformTasks 处理同一平台的任务
|
||
func processPlatformTasks(ctx context.Context, platform string, tasks []*common.AsyncTask) {
|
||
// 按 channel_id 分组
|
||
channelGroups := make(map[int64][]*common.AsyncTask)
|
||
for _, t := range tasks {
|
||
channelGroups[t.ChannelID] = append(channelGroups[t.ChannelID], t)
|
||
}
|
||
|
||
for channelID, channelTasks := range channelGroups {
|
||
processChannelTasks(ctx, channelID, channelTasks)
|
||
}
|
||
}
|
||
|
||
// processChannelTasks 处理同一渠道的任务
|
||
func processChannelTasks(ctx context.Context, channelID int64, tasks []*common.AsyncTask) {
|
||
// 获取渠道信息
|
||
channel, err := DefaultAsyncProvider.GetChannelByID(ctx, channelID)
|
||
if err != nil || channel == nil {
|
||
g.Log().Warningf(ctx, "poll: get channel %d: %v", channelID, err)
|
||
return
|
||
}
|
||
|
||
providerType := constant.ProviderType(channel.Type)
|
||
adaptor, err := taskchannel.GetAdaptor(providerType)
|
||
if err != nil {
|
||
g.Log().Warningf(ctx, "poll: get adaptor for channel %d: %v", channelID, err)
|
||
return
|
||
}
|
||
|
||
for _, t := range tasks {
|
||
pollSingleTask(ctx, adaptor, channel, t)
|
||
}
|
||
}
|
||
|
||
// privateData PrivateData 反序列化结构
|
||
type privateData struct {
|
||
UpstreamTaskID string `json:"upstream_task_id"`
|
||
UpstreamRequestID string `json:"upstream_request_id"` // 上游提交调用追踪 ID（部分上游返回，可空）
|
||
TaskType string `json:"task_type"`
|
||
BillingContext struct {
|
||
Ratios map[string]any `json:"ratios"`
|
||
ModelName string `json:"model_name"`
|
||
PreDeduct float64 `json:"pre_deduct"`
|
||
} `json:"billing_context"`
|
||
}
|
||
|
||
// pollSingleTask 轮询单个任务
|
||
func pollSingleTask(ctx context.Context, adaptor common.TaskAdaptor, channel *common.ChannelBasicInfo, task *common.AsyncTask) {
|
||
// 从 PrivateData 提取上游任务 ID 和计费上下文
|
||
var pd privateData
|
||
if err := json.Unmarshal(task.PrivateData, &pd); err != nil || pd.UpstreamTaskID == "" {
|
||
g.Log().Warningf(ctx, "poll: invalid private data for task %s", task.PublicTaskID)
|
||
return
|
||
}
|
||
|
||
// 查询上游状态（model 供供应商适配器按模型分派协议端点，如 MiniMax v1/v2；
|
||
// 其余适配器按结构体解析 taskData，自动忽略未知键）
|
||
upstreamModel := task.UpstreamModel
|
||
if upstreamModel == "" {
|
||
upstreamModel = task.ModelName
|
||
}
|
||
taskData, _ := json.Marshal(map[string]any{
|
||
"task_id": pd.UpstreamTaskID,
|
||
"use_proxy": parseChannelUseProxy(channel.Settings),
|
||
"model": upstreamModel,
|
||
})
|
||
|
||
// 渠道调试日志（轮询记录，仅段2/段3——后台轮询无客户端交互）：
|
||
// 节流规则 = 只在「状态/进度相对上一次轮询发生变化」时落库（任务行 CAS 前的值即上一次
|
||
// 持久化状态，无需额外存储），防止 15s 轮询周期灌爆 chn_debug_logs；
|
||
// fetch/parse 错误无「上一次」可对比且是轮询排查的主要目标，必记。
|
||
// request_id 复用提交请求原 ID，管理后台按 request_id 筛选可看到提交 + 全部状态转移。
|
||
var dbgSess *common.DebugSession
|
||
var dbgAttempt *common.DebugAttempt
|
||
if pollDebugLogEnabled(channel.Settings, task) {
|
||
dbgSess = common.NewDebugSession(task.RequestID, task.TenantID, task.UserID, task.ApiKeyID, "task_poll:"+task.PublicTaskID)
|
||
dbgAttempt = dbgSess.BeginTaskAttempt(&common.ChannelMeta{
|
||
ChannelID: channel.ID,
|
||
ChannelName: channel.Name,
|
||
ChannelType: channel.Type,
|
||
UpstreamModelName: upstreamModel,
|
||
}, task.ModelName, task.Platform, 0)
|
||
}
|
||
pollCtx := ctx
|
||
if dbgAttempt != nil {
|
||
pollCtx = common.WithDebugAttempt(ctx, dbgAttempt.Capture)
|
||
}
|
||
|
||
resp, err := adaptor.FetchTask(pollCtx, channel.BaseURL, channel.ApiKey, taskData)
|
||
if err != nil {
|
||
g.Log().Warningf(ctx, "poll: fetch task %s: %v", task.PublicTaskID, err)
|
||
dbgAttempt.MarkFinal(err)
|
||
dbgSess.FinalizeAndSubmit(0, 0)
|
||
return
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
body, err := io.ReadAll(resp.Body)
|
||
if err != nil {
|
||
g.Log().Warningf(ctx, "poll: read response for task %s: %v", task.PublicTaskID, err)
|
||
dbgAttempt.MarkFinal(err)
|
||
dbgSess.FinalizeAndSubmit(0, 0)
|
||
return
|
||
}
|
||
|
||
// 解析结果
|
||
taskInfo, err := adaptor.ParseTaskResult(body)
|
||
if err != nil {
|
||
g.Log().Warningf(ctx, "poll: parse result for task %s: %v", task.PublicTaskID, err)
|
||
dbgAttempt.MarkFinal(err)
|
||
dbgSess.FinalizeAndSubmit(0, 0)
|
||
return
|
||
}
|
||
|
||
// 状态变化判定 + 调试记录提交（CAS 更新 task.Status 之前，task 行内即上一次轮询值）
|
||
if pollStateChanged(taskInfo, task) {
|
||
dbgAttempt.MarkFinal(nil)
|
||
dbgSess.FinalizeAndSubmit(0, 0)
|
||
}
|
||
|
||
// CAS 更新状态
|
||
oldStatus := task.Status
|
||
task.Status = string(taskInfo.Status)
|
||
task.Progress = taskInfo.Progress
|
||
task.Data = body
|
||
|
||
if taskInfo.FailReason != "" {
|
||
task.FailReason = taskInfo.FailReason
|
||
}
|
||
|
||
now := time.Now()
|
||
switch {
|
||
case taskInfo.Status == common.TaskStatusInProgress:
|
||
task.StartTime = &now
|
||
case taskInfo.Status.IsTerminal():
|
||
task.FinishTime = &now
|
||
if task.StartTime == nil {
|
||
task.StartTime = &now
|
||
}
|
||
// 保留上游返回的 token 用量
|
||
if taskInfo.TotalTokens > 0 {
|
||
task.CompletionTokens = taskInfo.CompletionTokens
|
||
task.PromptTokens = taskInfo.PromptTokens
|
||
task.TotalTokens = taskInfo.TotalTokens
|
||
}
|
||
}
|
||
|
||
if err := DefaultAsyncProvider.UpdateTaskCAS(ctx, task, oldStatus); err != nil {
|
||
g.Log().Debugf(ctx, "poll: CAS conflict for task %s (status changed by another process): %v", task.PublicTaskID, err)
|
||
return
|
||
}
|
||
|
||
if taskInfo.Status.IsTerminal() {
|
||
DecrActiveTask()
|
||
monitor.UnregisterRequestByTaskID(task.PublicTaskID)
|
||
}
|
||
|
||
// 处理终态
|
||
if taskInfo.Status == common.TaskStatusSuccess {
|
||
task.ResultURL = taskInfo.ResultURL
|
||
DefaultAsyncProvider.UpdateTask(ctx, task) // 更新 result_url
|
||
|
||
// 结算计费
|
||
var settleResult *common.SettlementResult
|
||
if task.PreDeductAmount.GreaterThan(billing.Zero) && !task.BillingSettled {
|
||
taskBilling := billing.NewTaskBillingProvider()
|
||
actualCost := task.PreDeductAmount
|
||
|
||
// 优先用上游返回的 ActualCost
|
||
if taskInfo.ActualCost > 0 {
|
||
actualCost = billing.NewFromFloat(taskInfo.ActualCost)
|
||
} else if taskInfo.MaterialUsage != nil {
|
||
// 素材计费方案：按官方 usage 计量（输入视频时长/图片张数/输出秒数）重算。
|
||
// 链接素材的时长提交时不可知，此为最终计费依据；模型未配置素材方案时 ok=false 保持预扣
|
||
if materialCost, ok, mErr := taskBilling.RecalculateByMaterials(ctx, task.TenantID, task.ModelName, pd.BillingContext.Ratios, taskInfo.MaterialUsage, task.CreatedAt); mErr != nil {
|
||
g.Log().Warningf(ctx, "poll: recalculate by materials task %s: %v", task.PublicTaskID, mErr)
|
||
} else if ok {
|
||
actualCost = materialCost
|
||
}
|
||
} else if taskInfo.TotalTokens > 0 && pd.BillingContext.Ratios != nil {
|
||
// 用上游 total_tokens + 保存的 ratios 重算
|
||
if tokenCost, err := taskBilling.RecalculateByTokens(ctx, task.TenantID, task.ModelName, taskInfo.TotalTokens, pd.BillingContext.Ratios, task.CreatedAt); err == nil && tokenCost.GreaterThan(billing.Zero) {
|
||
actualCost = tokenCost
|
||
}
|
||
}
|
||
|
||
task.ActualCost = actualCost
|
||
settleResult, err = taskBilling.SettleTaskSuccess(ctx, task.TenantID, task.UserID, task.ApiKeyID, task.ChannelID, task.ModelName, task.RequestID, actualCost, task.PreDeductAmount, taskInfo.TotalTokens, taskInfo.CompletionTokens, pd.BillingContext.Ratios, task.PublicTaskID, task.CreatedAt)
|
||
if err != nil {
|
||
g.Log().Warningf(ctx, "poll: settle task %s: %v", task.PublicTaskID, err)
|
||
} else {
|
||
task.BillingSettled = true
|
||
// 幂等重复结算（DuplicateSkip）不得重复累加 Key 额度
|
||
if settleResult == nil || !settleResult.DuplicateSkip {
|
||
taskBilling.IncrApiKeyQuotaUsed(ctx, task.ApiKeyID, actualCost)
|
||
}
|
||
DefaultAsyncProvider.UpdateTask(ctx, task)
|
||
}
|
||
billing.CleanupPreDeduct(ctx, task.TenantID, task.RequestID+"_adjust")
|
||
}
|
||
g.Log().Infof(ctx, "poll: task %s completed", task.PublicTaskID)
|
||
|
||
// 记录用量日志
|
||
recordTaskUsage(task, channel, true, "", settleResult)
|
||
|
||
// 更新审计记录
|
||
recordTaskCompletionAudit(task, "SUCCESS", string(body), upstreamRespHeaders(resp))
|
||
|
||
} else if taskInfo.Status == common.TaskStatusFailure {
|
||
// 退还预扣费用
|
||
if task.PreDeductAmount.GreaterThan(billing.Zero) && !task.BillingSettled {
|
||
taskBilling := billing.NewTaskBillingProvider()
|
||
if err := taskBilling.SettleTaskFailed(ctx, task.TenantID, task.RequestID, task.PreDeductAmount); err != nil {
|
||
g.Log().Warningf(ctx, "poll: refund failed task %s: %v", task.PublicTaskID, err)
|
||
} else {
|
||
task.BillingSettled = true
|
||
DefaultAsyncProvider.UpdateTask(ctx, task)
|
||
}
|
||
billing.CleanupPreDeduct(ctx, task.TenantID, task.RequestID+"_adjust")
|
||
}
|
||
g.Log().Infof(ctx, "poll: task %s failed: %s", task.PublicTaskID, task.FailReason)
|
||
|
||
// 记录用量日志
|
||
recordTaskUsage(task, channel, false, task.FailReason, nil)
|
||
|
||
// 更新审计记录
|
||
recordTaskCompletionAudit(task, "FAILURE", string(body), upstreamRespHeaders(resp))
|
||
}
|
||
}
|
||
|
||
// recordTaskUsage 异步记录视频/音乐等异步任务的用量日志
|
||
func recordTaskUsage(task *common.AsyncTask, channel *common.ChannelBasicInfo, success bool, errMsg string, settleResult *common.SettlementResult) {
|
||
latencyMs := 0
|
||
if task.SubmitTime != nil && task.FinishTime != nil {
|
||
latencyMs = int(task.FinishTime.Sub(*task.SubmitTime).Milliseconds())
|
||
}
|
||
|
||
// 提取渠道名称和类型
|
||
var channelName string
|
||
var channelType int
|
||
if channel != nil {
|
||
channelName = channel.Name
|
||
channelType = channel.Type
|
||
}
|
||
|
||
status := "success"
|
||
if !success {
|
||
status = "error"
|
||
}
|
||
|
||
// 按秒计费任务的视频时长：取提交时计费上下文里的 spec.duration（真实秒数）。
|
||
// 只认 spec.* 事实键——旧 duration 键在部分 adaptor 是乘数语义（如 gemini 1.3），不可作为时长
|
||
// 提取上游请求 ID 落用量日志（排障时凭它在上游定位任务）：
|
||
// 优先 upstream_request_id（上游调用追踪 ID，如 DashScope 顶层 request_id），
|
||
// 未记录时回退 upstream_task_id（任务句柄，轮询同款标识）
|
||
durationSeconds := 0
|
||
upstreamRequestID := ""
|
||
if len(task.PrivateData) > 0 {
|
||
var pdDur struct {
|
||
UpstreamTaskID string `json:"upstream_task_id"`
|
||
UpstreamRequestID string `json:"upstream_request_id"`
|
||
BillingContext struct {
|
||
Ratios map[string]any `json:"ratios"`
|
||
} `json:"billing_context"`
|
||
}
|
||
if json.Unmarshal(task.PrivateData, &pdDur) == nil {
|
||
upstreamRequestID = pdDur.UpstreamRequestID
|
||
if upstreamRequestID == "" {
|
||
upstreamRequestID = pdDur.UpstreamTaskID
|
||
}
|
||
if v, ok := pdDur.BillingContext.Ratios["spec.duration"].(float64); ok && v > 0 {
|
||
durationSeconds = int(v)
|
||
}
|
||
}
|
||
}
|
||
|
||
record := &common.UsageRecord{
|
||
TenantID: task.TenantID,
|
||
UserID: task.UserID,
|
||
ApiKeyID: task.ApiKeyID,
|
||
ChannelID: task.ChannelID,
|
||
ChannelName: channelName,
|
||
ChannelType: channelType,
|
||
ModelName: task.ModelName,
|
||
RelayMode: int(constant.RelayModeVideoGenerations),
|
||
RequestType: 3, // async
|
||
LatencyMs: float64(latencyMs),
|
||
IsStream: false,
|
||
Success: success,
|
||
RequestID: task.RequestID,
|
||
Status: status,
|
||
// errMsg 可能来自上游 FailReason 原文，租户端用量日志可见，抹除其中的 URL/IP
|
||
ErrorMessage: helper.RedactMessage(errMsg),
|
||
PromptTokens: task.PromptTokens,
|
||
CompletionTokens: task.CompletionTokens,
|
||
TotalTokens: task.TotalTokens,
|
||
TotalCost: billing.InexactFloat64(task.ActualCost),
|
||
ActualCost: billing.InexactFloat64(task.ActualCost),
|
||
PreDeductAmount: billing.InexactFloat64(task.PreDeductAmount),
|
||
BillingSource: "task",
|
||
TaskID: task.PublicTaskID,
|
||
|
||
DurationSeconds: durationSeconds,
|
||
|
||
// 任务的上游请求 ID = 调用追踪 ID（上游返回时），否则为任务句柄（轮询同款标识）
|
||
UpstreamRequestID: upstreamRequestID,
|
||
}
|
||
|
||
// 从结算结果填充计费快照
|
||
if settleResult != nil {
|
||
record.BillingSnapshot = settleResult.BillingSnapshot
|
||
record.BillingSummary = settleResult.BillingSummary
|
||
record.BillingMode = settleResult.BillingMode
|
||
record.BillingSource = settleResult.BillingSource
|
||
record.RateMultiplier = settleResult.RateMultiplier
|
||
record.InputCost = settleResult.InputCost
|
||
record.OutputCost = settleResult.OutputCost
|
||
record.RefundAmount = settleResult.RefundAmount
|
||
record.SupplementAmount = settleResult.SupplementAmount
|
||
record.Currency = billing.Currency(context.Background())
|
||
} else {
|
||
// 失败/超时任务：从定价中获取计费模式
|
||
billingMode := "per_request"
|
||
if pricing, err := billing.GetModelPrice(context.Background(), task.TenantID, task.ModelName); err == nil {
|
||
if pricing.BillingMode != "" {
|
||
billingMode = pricing.BillingMode
|
||
}
|
||
}
|
||
record.BillingMode = billingMode
|
||
}
|
||
|
||
relay.NewDataProvider().RecordUsage(context.Background(), record)
|
||
}
|
||
|
||
// recordTaskCompletionAudit 更新提交阶段写入的审计记录，补充异步任务最终结果
|
||
func recordTaskCompletionAudit(task *common.AsyncTask, status string, resultBody string, upstreamHeaders map[string]string) {
|
||
now := time.Now()
|
||
|
||
// 计算从任务提交到完成的端到端延迟
|
||
latencyMs := 0
|
||
if task.SubmitTime != nil {
|
||
latencyMs = int(now.Sub(*task.SubmitTime).Milliseconds())
|
||
}
|
||
|
||
relay.NewDataProvider().UpdateTaskAudit(context.Background(), &common.AuditRecord{
|
||
TenantID: task.TenantID,
|
||
TaskID: task.PublicTaskID,
|
||
TaskStatus: status,
|
||
TaskResult: resultBody,
|
||
TaskUpstreamHeaders: upstreamHeaders,
|
||
TaskCompletedAt: &now,
|
||
LatencyMs: latencyMs,
|
||
})
|
||
}
|
||
|
||
// upstreamRespHeaders 从上游 HTTP 响应中提取响应头
|
||
func upstreamRespHeaders(resp *http.Response) map[string]string {
|
||
if resp == nil || resp.Header == nil {
|
||
return nil
|
||
}
|
||
headers := make(map[string]string)
|
||
for k, vals := range resp.Header {
|
||
if len(vals) > 0 {
|
||
headers[k] = vals[0]
|
||
}
|
||
}
|
||
return headers
|
||
}
|
||
|
||
// parseChannelUseProxy 从渠道 Settings JSONB 中提取 use_proxy 字段
|
||
func parseChannelUseProxy(settings json.RawMessage) bool {
|
||
var s struct {
|
||
UseProxy bool `json:"use_proxy"`
|
||
}
|
||
if len(settings) > 0 {
|
||
json.Unmarshal(settings, &s)
|
||
}
|
||
return s.UseProxy
|
||
}
|
||
|
||
// pollDebugLogEnabled 判定轮询渠道是否开启调试日志且匹配任务的租户/成员/密钥过滤
|
||
// （ChannelBasicInfo.Settings 为 JSONB 原文，按需解析；解析失败视为未开启）
|
||
func pollDebugLogEnabled(settings json.RawMessage, task *common.AsyncTask) bool {
|
||
var s common.ChannelSettings
|
||
if len(settings) == 0 || json.Unmarshal(settings, &s) != nil {
|
||
return false
|
||
}
|
||
return s.DebugLogEnabled && s.DebugTargetMatch(task.TenantID, task.UserID, task.ApiKeyID)
|
||
}
|
||
|
||
// pollStateChanged 判定本次轮询结果相对任务行当前持久化状态是否发生变化（状态或进度任一变化）。
|
||
// 独立成纯函数以便回归测试；在 CAS 更新 task 之前调用，task 内即上一次轮询持久化的值
|
||
func pollStateChanged(taskInfo *common.TaskInfo, task *common.AsyncTask) bool {
|
||
return string(taskInfo.Status) != task.Status || taskInfo.Progress != task.Progress
|
||
}
|