From 1a8567758f8db649a5d72be2c849ce674bf874c9 Mon Sep 17 00:00:00 2001
From: feitianbubu <feitianbubu@qq.com>
Date: Wed, 28 Jan 2026 18:33:14 +0800
Subject: [PATCH] feat: kling cost quota support use FinalUnitDeduction as
 totalToken

---
 relay/channel/task/kling/adaptor.go | 42 ++++++++++++++++++++++-------
 1 file changed, 32 insertions(+), 10 deletions(-)

diff --git a/relay/channel/task/kling/adaptor.go b/relay/channel/task/kling/adaptor.go
index 5fb853481..e2b4a9f9d 100644
--- a/relay/channel/task/kling/adaptor.go
+++ b/relay/channel/task/kling/adaptor.go
@@ -5,7 +5,9 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"math"
 	"net/http"
+	"strconv"
 	"strings"
 	"time"
 
@@ -80,15 +82,28 @@ type responsePayload struct {
 		TaskId        string `json:"task_id"`
 		TaskStatus    string `json:"task_status"`
 		TaskStatusMsg string `json:"task_status_msg"`
-		TaskResult    struct {
+		TaskInfo      struct {
+			ExternalTaskId string `json:"external_task_id"`
+		} `json:"task_info"`
+		WatermarkInfo struct {
+			Enabled bool `json:"enabled"`
+		} `json:"watermark_info"`
+		TaskResult struct {
 			Videos []struct {
-				Id       string `json:"id"`
-				Url      string `json:"url"`
-				Duration string `json:"duration"`
+				Id           string `json:"id"`
+				Url          string `json:"url"`
+				WatermarkUrl string `json:"watermark_url"`
+				Duration     string `json:"duration"`
 			} `json:"videos"`
+			Images []struct {
+				Index        int    `json:"index"`
+				Url          string `json:"url"`
+				WatermarkUrl string `json:"watermark_url"`
+			} `json:"images"`
 		} `json:"task_result"`
-		CreatedAt int64 `json:"created_at"`
-		UpdatedAt int64 `json:"updated_at"`
+		CreatedAt          int64  `json:"created_at"`
+		UpdatedAt          int64  `json:"updated_at"`
+		FinalUnitDeduction string `json:"final_unit_deduction"`
 	} `json:"data"`
 }
 
@@ -356,15 +371,22 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
 		taskInfo.Status = model.TaskStatusInProgress
 	case "succeed":
 		taskInfo.Status = model.TaskStatusSuccess
+		if videos := resPayload.Data.TaskResult.Videos; len(videos) > 0 {
+			video := videos[0]
+			taskInfo.Url = video.Url
+		}
+		if tokens, err := strconv.ParseFloat(resPayload.Data.FinalUnitDeduction, 64); err == nil {
+			rounded := int(math.Ceil(tokens))
+			if rounded > 0 {
+				taskInfo.CompletionTokens = rounded
+				taskInfo.TotalTokens = rounded
+			}
+		}
 	case "failed":
 		taskInfo.Status = model.TaskStatusFailure
 	default:
 		return nil, fmt.Errorf("unknown task status: %s", status)
 	}
-	if videos := resPayload.Data.TaskResult.Videos; len(videos) > 0 {
-		video := videos[0]
-		taskInfo.Url = video.Url
-	}
 	return taskInfo, nil
 }