mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-09-11 22:49:57 +00:00
feat(task): replace built-in task adaptors with a sandboxed JS plugin system (#7076)
This commit is contained in:
Calcium-Ion
@@ -105,6 +105,9 @@ func TestSetUserPermissionsStoresOnlyOverrides(t *testing.T) {
|
||||
ActionSensitiveWrite: true,
|
||||
ActionSecretView: false,
|
||||
},
|
||||
ResourceTaskPlugin: {
|
||||
ActionBind: false,
|
||||
},
|
||||
}, ExplicitUserPermissions(42))
|
||||
assert.Equal(t, PermissionsMap{
|
||||
ResourceChannel: {
|
||||
@@ -133,6 +136,9 @@ func TestSetUserPermissionsStoresOnlyOverrides(t *testing.T) {
|
||||
ActionSensitiveWrite: false,
|
||||
ActionSecretView: false,
|
||||
},
|
||||
ResourceTaskPlugin: {
|
||||
ActionBind: false,
|
||||
},
|
||||
}, ExplicitUserPermissions(42))
|
||||
assert.Empty(t, ExplicitUserOverrides(42))
|
||||
}
|
||||
@@ -226,4 +232,40 @@ func TestCapabilitiesUseCatalogShape(t *testing.T) {
|
||||
assert.True(t, capabilities[ResourceChannel][ActionWrite])
|
||||
assert.False(t, capabilities[ResourceChannel][ActionSensitiveWrite])
|
||||
assert.False(t, capabilities[ResourceChannel][ActionSecretView])
|
||||
assert.False(t, capabilities[ResourceTaskPlugin][ActionBind])
|
||||
}
|
||||
|
||||
func TestTaskPluginBindIsRootOnlyUntilGranted(t *testing.T) {
|
||||
db := newAuthzTestDB(t)
|
||||
require.NoError(t, Init(db))
|
||||
|
||||
var bindAction *ActionDefinition
|
||||
for _, resource := range Catalog() {
|
||||
if resource.Resource != ResourceTaskPlugin {
|
||||
continue
|
||||
}
|
||||
assert.Equal(t, "Task Plugin", resource.LabelKey)
|
||||
for i := range resource.Actions {
|
||||
if resource.Actions[i].Action == ActionBind {
|
||||
bindAction = &resource.Actions[i]
|
||||
}
|
||||
}
|
||||
}
|
||||
require.NotNil(t, bindAction)
|
||||
assert.Equal(t, "Bind task plugins", bindAction.LabelKey)
|
||||
assert.Equal(t, "List registered task plugins and bind them when creating or editing task plugin channels.", bindAction.DescriptionKey)
|
||||
assert.Empty(t, bindAction.DefaultRoles)
|
||||
|
||||
assert.False(t, Can(2, common.RoleAdminUser, TaskPluginBind))
|
||||
assert.True(t, Can(1, common.RoleRootUser, TaskPluginBind))
|
||||
|
||||
enforcer := currentEnforcer()
|
||||
require.NotNil(t, enforcer)
|
||||
_, err := enforcer.AddPolicy(RoleSubject(BuiltInRoleAdmin), ResourceTaskPlugin, ActionBind, EffectAllow)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, Can(2, common.RoleAdminUser, TaskPluginBind))
|
||||
|
||||
_, err = enforcer.RemovePolicy(RoleSubject(BuiltInRoleAdmin), ResourceTaskPlugin, ActionBind, EffectAllow)
|
||||
require.NoError(t, err)
|
||||
assert.False(t, Can(2, common.RoleAdminUser, TaskPluginBind))
|
||||
}
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
package authz
|
||||
|
||||
const (
|
||||
ResourceTaskPlugin = "task_plugin"
|
||||
|
||||
ActionBind = "bind"
|
||||
)
|
||||
|
||||
var TaskPluginBind = Permission{Resource: ResourceTaskPlugin, Action: ActionBind}
|
||||
|
||||
func init() {
|
||||
RegisterResource(ResourceDefinition{
|
||||
Resource: ResourceTaskPlugin,
|
||||
LabelKey: "Task Plugin",
|
||||
Actions: []ActionDefinition{
|
||||
{
|
||||
Action: ActionBind,
|
||||
LabelKey: "Bind task plugins",
|
||||
DescriptionKey: "List registered task plugins and bind them when creating or editing task plugin channels.",
|
||||
},
|
||||
},
|
||||
})
|
||||
}
|
||||
@@ -5,11 +5,36 @@ import (
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func GetChannelConstraints(c *gin.Context) *dto.ChannelConstraints {
|
||||
if c == nil {
|
||||
return &dto.ChannelConstraints{}
|
||||
}
|
||||
if existing, ok := common.GetContextKeyType[*dto.ChannelConstraints](c, constant.ContextKeyChannelConstraints); ok && existing != nil {
|
||||
return existing
|
||||
}
|
||||
constraints := &dto.ChannelConstraints{}
|
||||
common.SetContextKey(c, constant.ContextKeyChannelConstraints, constraints)
|
||||
return constraints
|
||||
}
|
||||
|
||||
func AppendTaskPluginIdentityFilter(c *gin.Context, pluginKey string) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
GetChannelConstraints(c).AddFilter(dto.ChannelFilter{
|
||||
Kind: dto.FilterTaskPluginIdentity,
|
||||
TaskPluginKey: pluginKey,
|
||||
TaskPluginChannelTypes: pinnedTaskPluginChannelTypes(c, pluginKey),
|
||||
})
|
||||
}
|
||||
|
||||
type RetryParam struct {
|
||||
Ctx *gin.Context
|
||||
TokenGroup string
|
||||
@@ -85,6 +110,7 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string,
|
||||
var err error
|
||||
selectGroup := param.TokenGroup
|
||||
userGroup := common.GetContextKeyString(param.Ctx, constant.ContextKeyUserGroup)
|
||||
filters := GetChannelConstraints(param.Ctx).Filters
|
||||
|
||||
if param.TokenGroup == "auto" {
|
||||
autoGroups := GetRequestAutoGroups(param.Ctx, userGroup)
|
||||
@@ -115,7 +141,12 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string,
|
||||
}
|
||||
logger.LogDebug(param.Ctx, "Auto selecting group: %s, priorityRetry: %d", autoGroup, priorityRetry)
|
||||
|
||||
channel, _ = model.GetRandomSatisfiedChannel(autoGroup, param.ModelName, priorityRetry, param.RequestPath)
|
||||
channel, _ = model.GetRandomSatisfiedChannel(
|
||||
autoGroup,
|
||||
param.ModelName,
|
||||
priorityRetry,
|
||||
filters,
|
||||
)
|
||||
if channel == nil {
|
||||
// Current group has no available channel for this model, try next group
|
||||
// 当前分组没有该模型的可用渠道，尝试下一个分组
|
||||
@@ -153,10 +184,68 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string,
|
||||
break
|
||||
}
|
||||
} else {
|
||||
channel, err = model.GetRandomSatisfiedChannel(param.TokenGroup, param.ModelName, param.GetRetry(), param.RequestPath)
|
||||
channel, err = model.GetRandomSatisfiedChannel(
|
||||
param.TokenGroup,
|
||||
param.ModelName,
|
||||
param.GetRetry(),
|
||||
filters,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, param.TokenGroup, err
|
||||
}
|
||||
}
|
||||
return channel, selectGroup, nil
|
||||
}
|
||||
|
||||
func pinnedTaskPluginChannelTypes(c *gin.Context, expected string) []int {
|
||||
if c == nil || expected == "" {
|
||||
return nil
|
||||
}
|
||||
if value, exists := c.Get(jsplugin.ContextKeyPinnedEndpoint); exists {
|
||||
pinned, ok := value.(jsplugin.PinnedEndpoint)
|
||||
if ok && pinned.Generation != nil && len(pinned.Candidates) > 1 {
|
||||
expectedFound := false
|
||||
channelTypes := make([]int, 0, len(pinned.Candidates))
|
||||
seen := make(map[int]struct{}, len(pinned.Candidates))
|
||||
for _, candidate := range pinned.Candidates {
|
||||
if candidate.Plugin == nil {
|
||||
continue
|
||||
}
|
||||
if candidate.Plugin.Meta.Key == expected {
|
||||
expectedFound = true
|
||||
}
|
||||
for _, channelType := range candidate.Plugin.Meta.ChannelTypes {
|
||||
if channelType == 0 || channelType == constant.ChannelTypeTaskPlugin {
|
||||
continue
|
||||
}
|
||||
if _, duplicate := seen[channelType]; duplicate {
|
||||
continue
|
||||
}
|
||||
if plugin, indexed := pinned.Generation.GetByChannelType(channelType); indexed && plugin == candidate.Plugin {
|
||||
seen[channelType] = struct{}{}
|
||||
channelTypes = append(channelTypes, channelType)
|
||||
}
|
||||
}
|
||||
}
|
||||
if expectedFound {
|
||||
return channelTypes
|
||||
}
|
||||
}
|
||||
}
|
||||
value, exists := c.Get(jsplugin.ContextKeyPinnedPlugin)
|
||||
pinned, ok := value.(jsplugin.PinnedPlugin)
|
||||
if !exists || !ok || pinned.Generation == nil || pinned.Plugin == nil || pinned.Plugin.Meta.Key != expected {
|
||||
return nil
|
||||
}
|
||||
channelTypes := make([]int, 0, len(pinned.Plugin.Meta.ChannelTypes))
|
||||
for _, channelType := range pinned.Plugin.Meta.ChannelTypes {
|
||||
if channelType == 0 || channelType == constant.ChannelTypeTaskPlugin {
|
||||
continue
|
||||
}
|
||||
channelTypes = append(channelTypes, channelType)
|
||||
}
|
||||
if len(channelTypes) == 0 {
|
||||
return nil
|
||||
}
|
||||
return channelTypes
|
||||
}
|
||||
|
||||
@@ -0,0 +1,152 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestPinnedTaskPluginChannelTypesUsesPinnedGenerationIndex(t *testing.T) {
|
||||
registry := jsplugin.NewRegistry()
|
||||
plugin, err := registry.Register(channelSelectTaskPluginSource("legacy-select", constant.ChannelTypeKling), jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
|
||||
c, _ := gin.CreateTestContext(nil)
|
||||
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{
|
||||
Generation: registry.Generation(),
|
||||
Plugin: plugin,
|
||||
})
|
||||
|
||||
assert.Equal(t, []int{constant.ChannelTypeKling}, pinnedTaskPluginChannelTypes(c, "legacy-select"))
|
||||
assert.Empty(t, pinnedTaskPluginChannelTypes(c, "another-plugin"))
|
||||
assert.Empty(t, pinnedTaskPluginChannelTypes(nil, "legacy-select"))
|
||||
}
|
||||
|
||||
func TestPinnedTaskPluginChannelTypesLeavesGenericChannelsKeyed(t *testing.T) {
|
||||
registry := jsplugin.NewRegistry()
|
||||
plugin, err := registry.Register(channelSelectTaskPluginSource("generic-select", constant.ChannelTypeTaskPlugin), jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
|
||||
c, _ := gin.CreateTestContext(nil)
|
||||
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{
|
||||
Generation: registry.Generation(),
|
||||
Plugin: plugin,
|
||||
})
|
||||
|
||||
assert.Empty(t, pinnedTaskPluginChannelTypes(c, "generic-select"))
|
||||
}
|
||||
|
||||
func TestPinnedTaskPluginChannelTypesIncludesSharedEndpointProviders(t *testing.T) {
|
||||
registry := jsplugin.NewRegistry()
|
||||
_, err := registry.Register(channelSelectEndpointPluginSource("gemini-select", constant.ChannelTypeGemini), jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
_, err = registry.Register(channelSelectEndpointPluginSource("vertex-select", constant.ChannelTypeVertexAi), jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
candidates := registry.Generation().LookupEndpointCandidates("POST", "/v1/responses", "task-model")
|
||||
require.Len(t, candidates, 2)
|
||||
|
||||
c, _ := gin.CreateTestContext(nil)
|
||||
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{
|
||||
Generation: registry.Generation(),
|
||||
Plugin: candidates[0].Plugin,
|
||||
})
|
||||
c.Set(jsplugin.ContextKeyPinnedEndpoint, jsplugin.PinnedEndpoint{
|
||||
Generation: registry.Generation(),
|
||||
Plugin: candidates[0].Plugin,
|
||||
Protocol: candidates[0].Protocol,
|
||||
Operation: candidates[0].Operation,
|
||||
Model: "task-model",
|
||||
Candidates: candidates,
|
||||
})
|
||||
|
||||
assert.Equal(t, []int{constant.ChannelTypeGemini, constant.ChannelTypeVertexAi}, pinnedTaskPluginChannelTypes(c, candidates[0].Plugin.Meta.Key))
|
||||
}
|
||||
|
||||
func channelSelectTaskPluginSource(key string, channelType int) string {
|
||||
return fmt.Sprintf(`
|
||||
export const meta = {
|
||||
apiVersion: 1,
|
||||
key: %q,
|
||||
name: %q,
|
||||
version: "1.0.0",
|
||||
author: {name: "Test"},
|
||||
%s
|
||||
models: ["task-model"],
|
||||
fetchMode: "per_task",
|
||||
};
|
||||
export function buildSubmitRequest() { return {}; }
|
||||
export function parseSubmitResponse() { return {taskId: "task"}; }
|
||||
export function buildQueryRequest() { return {}; }
|
||||
export function parseTaskResult() { return {status: "SUCCESS"}; }
|
||||
`, key, key, channelSelectChannelTypesField(channelType))
|
||||
}
|
||||
|
||||
func channelSelectEndpointPluginSource(key string, channelType int) string {
|
||||
return fmt.Sprintf(`
|
||||
export const meta = {
|
||||
apiVersion: 1,
|
||||
key: %q,
|
||||
name: %q,
|
||||
version: "1.0.0",
|
||||
author: {name: "Test"},
|
||||
%s
|
||||
models: ["task-model"],
|
||||
fetchMode: "per_task",
|
||||
protocols: [{name: "openai_responses", supports: ["stream", "sync", "background"]}],
|
||||
};
|
||||
export function buildSubmitRequest() { return {}; }
|
||||
export function parseSubmitResponse() { return {taskId: "task"}; }
|
||||
export function buildQueryRequest() { return {}; }
|
||||
export function parseTaskResult() { return {status: "SUCCESS"}; }
|
||||
export const protocols = {openai_responses: {
|
||||
decodeRequest: function(ctx) { return {kind: "submit", model: "task-model", requestBody: ctx.body.value}; },
|
||||
renderEvents: function() { return {events: [], state: null, done: false}; },
|
||||
renderFinal: function() { return {output: []}; },
|
||||
}};
|
||||
`, key, key, channelSelectChannelTypesField(channelType))
|
||||
}
|
||||
|
||||
func channelSelectChannelTypesField(channelType int) string {
|
||||
if channelType <= 0 || channelType == constant.ChannelTypeTaskPlugin {
|
||||
return ""
|
||||
}
|
||||
return fmt.Sprintf("channelTypes: [%d],", channelType)
|
||||
}
|
||||
|
||||
func TestPinnedTaskPluginChannelTypesIncludesCompatibleTypes(t *testing.T) {
|
||||
registry := jsplugin.NewRegistry()
|
||||
plugin, err := registry.Register(channelSelectCompatiblePluginSource("sora-select", constant.ChannelTypeSora, constant.ChannelTypeOpenAI), jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
|
||||
c, _ := gin.CreateTestContext(nil)
|
||||
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{
|
||||
Generation: registry.Generation(),
|
||||
Plugin: plugin,
|
||||
})
|
||||
|
||||
assert.Equal(t, []int{constant.ChannelTypeSora, constant.ChannelTypeOpenAI}, pinnedTaskPluginChannelTypes(c, "sora-select"))
|
||||
}
|
||||
|
||||
func channelSelectCompatiblePluginSource(key string, channelType, compatibleType int) string {
|
||||
return fmt.Sprintf(`
|
||||
export const meta = {
|
||||
apiVersion: 1,
|
||||
key: %q,
|
||||
name: %q,
|
||||
version: "1.0.0",
|
||||
author: {name: "Test"},
|
||||
channelTypes: [%d, %d],
|
||||
models: ["task-model"],
|
||||
fetchMode: "per_task",
|
||||
};
|
||||
export function buildSubmitRequest() { return {}; }
|
||||
export function parseSubmitResponse() { return {taskId: "task"}; }
|
||||
export function buildQueryRequest() { return {}; }
|
||||
export function parseTaskResult() { return {status: "SUCCESS"}; }
|
||||
`, key, key, channelType, compatibleType)
|
||||
}
|
||||
@@ -33,7 +33,7 @@ func FetchCodexChannelModels(channel *model.Channel) ([]string, error) {
|
||||
|
||||
baseURL := channel.GetBaseURL()
|
||||
if baseURL == "" {
|
||||
baseURL = constant.ChannelBaseURLs[constant.ChannelTypeCodex]
|
||||
baseURL = constant.GetChannelBaseURL(constant.ChannelTypeCodex)
|
||||
}
|
||||
return fetchCodexChannelModels(ctx, channel, baseURL, client, clientVersion)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,137 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"crypto/hmac"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
)
|
||||
|
||||
const (
|
||||
TaskArtifactAccessQueryParameter = "access"
|
||||
taskArtifactAccessVersion = "v1"
|
||||
taskArtifactAccessLength = 43
|
||||
maxTaskArtifactTaskIDLength = 191
|
||||
maxTaskArtifactKeyLength = 128
|
||||
)
|
||||
|
||||
var ErrTaskArtifactAccessInvalid = errors.New("task artifact access is invalid")
|
||||
|
||||
func taskArtifactAccessMessage(taskID, artifactKey string) []byte {
|
||||
return []byte(taskArtifactAccessVersion + "\x00" + taskID + "\x00" + artifactKey)
|
||||
}
|
||||
|
||||
// IssueTaskArtifactAccess creates a stable capability bound to exactly one
|
||||
// public task ID and artifact key. It contains no user or upstream data.
|
||||
func IssueTaskArtifactAccess(taskID, artifactKey string) (string, error) {
|
||||
taskID = strings.TrimSpace(taskID)
|
||||
artifactKey = strings.TrimSpace(artifactKey)
|
||||
if taskID == "" || len(taskID) > maxTaskArtifactTaskIDLength ||
|
||||
artifactKey == "" || len(artifactKey) > maxTaskArtifactKeyLength ||
|
||||
common.CryptoSecret == "" {
|
||||
return "", ErrTaskArtifactAccessInvalid
|
||||
}
|
||||
|
||||
mac := hmac.New(sha256.New, []byte(common.CryptoSecret))
|
||||
_, _ = mac.Write(taskArtifactAccessMessage(taskID, artifactKey))
|
||||
return base64.RawURLEncoding.EncodeToString(mac.Sum(nil)), nil
|
||||
}
|
||||
|
||||
// VerifyTaskArtifactAccess verifies the route binding without reading task,
|
||||
// user, or token state. Signature comparison is constant-time.
|
||||
func VerifyTaskArtifactAccess(access, taskID, artifactKey string) bool {
|
||||
taskID = strings.TrimSpace(taskID)
|
||||
artifactKey = strings.TrimSpace(artifactKey)
|
||||
if len(access) != taskArtifactAccessLength ||
|
||||
taskID == "" || len(taskID) > maxTaskArtifactTaskIDLength ||
|
||||
artifactKey == "" || len(artifactKey) > maxTaskArtifactKeyLength ||
|
||||
common.CryptoSecret == "" {
|
||||
return false
|
||||
}
|
||||
|
||||
actualSignature, err := base64.RawURLEncoding.Strict().DecodeString(access)
|
||||
if err != nil || len(actualSignature) != sha256.Size {
|
||||
return false
|
||||
}
|
||||
|
||||
mac := hmac.New(sha256.New, []byte(common.CryptoSecret))
|
||||
_, _ = mac.Write(taskArtifactAccessMessage(taskID, artifactKey))
|
||||
return hmac.Equal(actualSignature, mac.Sum(nil))
|
||||
}
|
||||
|
||||
// ValidateTaskArtifactBaseURL validates configuration syntax only. It
|
||||
// deliberately performs no DNS lookup or reachability probe.
|
||||
func ValidateTaskArtifactBaseURL(raw string) error {
|
||||
trimmed := strings.TrimSpace(raw)
|
||||
if trimmed == "" {
|
||||
return errors.New("task artifact base URL is empty")
|
||||
}
|
||||
if raw != trimmed {
|
||||
return errors.New("task artifact base URL must not contain surrounding whitespace")
|
||||
}
|
||||
raw = trimmed
|
||||
parsed, err := url.Parse(raw)
|
||||
if err != nil || parsed == nil {
|
||||
return errors.New("task artifact base URL is invalid")
|
||||
}
|
||||
if !strings.EqualFold(parsed.Scheme, "http") && !strings.EqualFold(parsed.Scheme, "https") {
|
||||
return errors.New("task artifact base URL must use http or https")
|
||||
}
|
||||
if parsed.Host == "" || parsed.User != nil || parsed.Opaque != "" {
|
||||
return errors.New("task artifact base URL must contain a host and no userinfo")
|
||||
}
|
||||
if parsed.RawQuery != "" || parsed.ForceQuery || parsed.Fragment != "" || strings.Contains(raw, "#") {
|
||||
return errors.New("task artifact base URL must not contain a query or fragment")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// BuildTaskArtifactContentURL returns an absolute, long-lived capability URL.
|
||||
// TaskPublicAddress wins when configured; ServerAddress is the only fallback.
|
||||
// Request Host headers are intentionally not involved.
|
||||
func BuildTaskArtifactContentURL(taskID, artifactKey string) (string, error) {
|
||||
taskID = strings.TrimSpace(taskID)
|
||||
artifactKey = strings.TrimSpace(artifactKey)
|
||||
if taskID == "" || len(taskID) > maxTaskArtifactTaskIDLength ||
|
||||
artifactKey == "" || len(artifactKey) > maxTaskArtifactKeyLength {
|
||||
return "", ErrTaskArtifactAccessInvalid
|
||||
}
|
||||
|
||||
baseAddress := strings.TrimSpace(system_setting.TaskPublicAddress)
|
||||
if baseAddress == "" {
|
||||
baseAddress = strings.TrimSpace(system_setting.ServerAddress)
|
||||
}
|
||||
if err := ValidateTaskArtifactBaseURL(baseAddress); err != nil {
|
||||
return "", err
|
||||
}
|
||||
baseURL, err := url.Parse(baseAddress)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
access, err := IssueTaskArtifactAccess(taskID, artifactKey)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
basePath := strings.TrimRight(baseURL.Path, "/")
|
||||
escapedBasePath := strings.TrimRight(baseURL.EscapedPath(), "/")
|
||||
suffixPath := fmt.Sprintf("/v1/tasks/%s/artifacts/%s/content", taskID, artifactKey)
|
||||
escapedSuffixPath := fmt.Sprintf(
|
||||
"/v1/tasks/%s/artifacts/%s/content",
|
||||
url.PathEscape(taskID),
|
||||
url.PathEscape(artifactKey),
|
||||
)
|
||||
baseURL.Path = basePath + suffixPath
|
||||
baseURL.RawPath = escapedBasePath + escapedSuffixPath
|
||||
query := baseURL.Query()
|
||||
query.Set(TaskArtifactAccessQueryParameter, access)
|
||||
baseURL.RawQuery = query.Encode()
|
||||
return baseURL.String(), nil
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"net/url"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestTaskArtifactAccessBindsTaskAndKey(t *testing.T) {
|
||||
previousSecret := common.CryptoSecret
|
||||
common.CryptoSecret = "task-artifact-access-test-secret"
|
||||
t.Cleanup(func() { common.CryptoSecret = previousSecret })
|
||||
|
||||
access, err := IssueTaskArtifactAccess("task-1", "video-main")
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, access, 43)
|
||||
assert.NotContains(t, access, ".")
|
||||
assert.True(t, VerifyTaskArtifactAccess(access, "task-1", "video-main"))
|
||||
assert.False(t, VerifyTaskArtifactAccess(access, "task-2", "video-main"))
|
||||
assert.False(t, VerifyTaskArtifactAccess(access, "task-1", "video-other"))
|
||||
assert.False(t, VerifyTaskArtifactAccess(access+"x", "task-1", "video-main"))
|
||||
|
||||
common.CryptoSecret = "another-node-secret"
|
||||
assert.False(t, VerifyTaskArtifactAccess(access, "task-1", "video-main"))
|
||||
}
|
||||
|
||||
func TestBuildTaskArtifactContentURLUsesConfiguredAddressAndPreservesPrefix(t *testing.T) {
|
||||
previousSecret := common.CryptoSecret
|
||||
previousPublicAddress := system_setting.TaskPublicAddress
|
||||
previousServerAddress := system_setting.ServerAddress
|
||||
common.CryptoSecret = "task-artifact-url-test-secret"
|
||||
system_setting.TaskPublicAddress = "https://media.example/gateway/prefix/"
|
||||
system_setting.ServerAddress = "https://fallback.invalid"
|
||||
t.Cleanup(func() {
|
||||
common.CryptoSecret = previousSecret
|
||||
system_setting.TaskPublicAddress = previousPublicAddress
|
||||
system_setting.ServerAddress = previousServerAddress
|
||||
})
|
||||
|
||||
contentURL, err := BuildTaskArtifactContentURL("task-public", "video-main")
|
||||
require.NoError(t, err)
|
||||
parsed, err := url.Parse(contentURL)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "media.example", parsed.Host)
|
||||
assert.Equal(t, "/gateway/prefix/v1/tasks/task-public/artifacts/video-main/content", parsed.Path)
|
||||
assert.True(t, VerifyTaskArtifactAccess(
|
||||
parsed.Query().Get(TaskArtifactAccessQueryParameter),
|
||||
"task-public",
|
||||
"video-main",
|
||||
))
|
||||
}
|
||||
|
||||
func TestBuildTaskArtifactContentURLFallsBackOnlyToServerAddress(t *testing.T) {
|
||||
previousSecret := common.CryptoSecret
|
||||
previousPublicAddress := system_setting.TaskPublicAddress
|
||||
previousServerAddress := system_setting.ServerAddress
|
||||
common.CryptoSecret = "task-artifact-fallback-test-secret"
|
||||
system_setting.TaskPublicAddress = ""
|
||||
system_setting.ServerAddress = "https://gateway.example/root"
|
||||
t.Cleanup(func() {
|
||||
common.CryptoSecret = previousSecret
|
||||
system_setting.TaskPublicAddress = previousPublicAddress
|
||||
system_setting.ServerAddress = previousServerAddress
|
||||
})
|
||||
|
||||
contentURL, err := BuildTaskArtifactContentURL("task-fallback", "audio")
|
||||
require.NoError(t, err)
|
||||
assert.Contains(t, contentURL, "https://gateway.example/root/v1/tasks/task-fallback/artifacts/audio/content")
|
||||
|
||||
system_setting.TaskPublicAddress = "not-a-url"
|
||||
_, err = BuildTaskArtifactContentURL("task-fallback", "audio")
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
func TestValidateTaskArtifactBaseURLOnlyAcceptsSafeAbsoluteHTTPURLs(t *testing.T) {
|
||||
for _, valid := range []string{
|
||||
"http://localhost:3000",
|
||||
"https://gateway.example",
|
||||
"https://gateway.example/prefix/path/",
|
||||
} {
|
||||
assert.NoError(t, ValidateTaskArtifactBaseURL(valid), valid)
|
||||
}
|
||||
for _, invalid := range []string{
|
||||
"",
|
||||
"/relative",
|
||||
"ftp://gateway.example",
|
||||
"https://user:secret@gateway.example",
|
||||
"https://gateway.example/path?tenant=1",
|
||||
"https://gateway.example/path#fragment",
|
||||
" https://gateway.example",
|
||||
"https://gateway.example ",
|
||||
} {
|
||||
assert.Error(t, ValidateTaskArtifactBaseURL(invalid), invalid)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
"github.com/QuantumNous/new-api/types"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// StoredArtifactRef describes a persisted artifact object. No reference is
|
||||
// produced until a concrete storage backend is implemented.
|
||||
type StoredArtifactRef struct {
|
||||
Backend string
|
||||
Bucket string
|
||||
ObjectKey string
|
||||
MimeType string
|
||||
Size int64
|
||||
}
|
||||
|
||||
// TaskArtifactStore is the persistence boundary for generated artifact bytes.
|
||||
// types.TaskArtifact is re-exported by relay/channel as channel.TaskArtifact.
|
||||
type TaskArtifactStore interface {
|
||||
Enabled() bool
|
||||
Resolve(task *model.Task, artifactKey string) (*StoredArtifactRef, error)
|
||||
Persist(ctx context.Context, task *model.Task, artifact types.TaskArtifact, content io.Reader) (*StoredArtifactRef, error)
|
||||
Serve(c *gin.Context, task *model.Task, ref *StoredArtifactRef) error
|
||||
}
|
||||
|
||||
var ErrTaskArtifactStoreDisabled = errors.New("task artifact store is disabled")
|
||||
|
||||
type disabledArtifactStore struct{}
|
||||
|
||||
func (disabledArtifactStore) Enabled() bool {
|
||||
return false
|
||||
}
|
||||
|
||||
func (disabledArtifactStore) Resolve(*model.Task, string) (*StoredArtifactRef, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (disabledArtifactStore) Persist(context.Context, *model.Task, types.TaskArtifact, io.Reader) (*StoredArtifactRef, error) {
|
||||
return nil, ErrTaskArtifactStoreDisabled
|
||||
}
|
||||
|
||||
func (disabledArtifactStore) Serve(*gin.Context, *model.Task, *StoredArtifactRef) error {
|
||||
return ErrTaskArtifactStoreDisabled
|
||||
}
|
||||
|
||||
var taskArtifactStore TaskArtifactStore = &disabledArtifactStore{}
|
||||
|
||||
func init() {
|
||||
_ = system_setting.LoadTaskArtifactStoreConfig()
|
||||
}
|
||||
|
||||
// GetTaskArtifactStore returns the process-wide artifact storage backend. This
|
||||
// release always returns the disabled implementation.
|
||||
func GetTaskArtifactStore() TaskArtifactStore {
|
||||
return taskArtifactStore
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/types"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestDisabledTaskArtifactStoreHasNoStorageBehavior(t *testing.T) {
|
||||
store := GetTaskArtifactStore()
|
||||
require.NotNil(t, store)
|
||||
assert.False(t, store.Enabled())
|
||||
|
||||
task := &model.Task{TaskID: "task-disabled-store"}
|
||||
ref, err := store.Resolve(task, "video")
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, ref)
|
||||
|
||||
ref, err = store.Persist(t.Context(), task, types.TaskArtifact{Key: "video", Type: "video"}, strings.NewReader("content"))
|
||||
assert.Nil(t, ref)
|
||||
assert.ErrorIs(t, err, ErrTaskArtifactStoreDisabled)
|
||||
assert.ErrorIs(t, store.Serve(&gin.Context{}, task, &StoredArtifactRef{Backend: "s3"}), ErrTaskArtifactStoreDisabled)
|
||||
assert.Same(t, store, GetTaskArtifactStore())
|
||||
}
|
||||
|
||||
var _ TaskArtifactStore = disabledArtifactStore{}
|
||||
+60
-9
@@ -2,6 +2,7 @@ package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
@@ -17,24 +18,29 @@ import (
|
||||
|
||||
// LogTaskConsumption 记录任务消费日志和统计信息（仅记录，不涉及实际扣费）。
|
||||
// 实际扣费已由 BillingSession（PreConsumeBilling + SettleBilling）完成。
|
||||
func LogTaskConsumption(c *gin.Context, info *relaycommon.RelayInfo) {
|
||||
func LogTaskConsumption(c *gin.Context, info *relaycommon.RelayInfo, task *model.Task) {
|
||||
tokenName := c.GetString("token_name")
|
||||
logContent := fmt.Sprintf("操作 %s", info.Action)
|
||||
// 支持任务仅按次计费
|
||||
if common.StringsContains(constant.TaskPricePatches, info.OriginModelName) {
|
||||
logContent = fmt.Sprintf("%s，按次计费", logContent)
|
||||
} else {
|
||||
var contents []string
|
||||
if otherRatios := info.PriceData.OtherRatios(); len(otherRatios) > 0 {
|
||||
var contents []string
|
||||
for key, ra := range otherRatios {
|
||||
if 1.0 != ra {
|
||||
contents = append(contents, fmt.Sprintf("%s: %.2f", key, ra))
|
||||
}
|
||||
}
|
||||
if len(contents) > 0 {
|
||||
logContent = fmt.Sprintf("%s, 计算参数：%s", logContent, strings.Join(contents, ", "))
|
||||
}
|
||||
if snap := info.TieredBillingSnapshot; snap != nil {
|
||||
for key, value := range snap.UsageFacts {
|
||||
contents = append(contents, fmt.Sprintf("%s: %v", key, value))
|
||||
}
|
||||
}
|
||||
if len(contents) > 0 {
|
||||
logContent = fmt.Sprintf("%s, 计算参数：%s", logContent, strings.Join(contents, ", "))
|
||||
}
|
||||
}
|
||||
other := make(map[string]interface{})
|
||||
other["is_task"] = true
|
||||
@@ -51,6 +57,15 @@ func LogTaskConsumption(c *gin.Context, info *relaycommon.RelayInfo) {
|
||||
other["is_model_mapped"] = true
|
||||
other["upstream_model_name"] = info.UpstreamModelName
|
||||
}
|
||||
if snap := info.TieredBillingSnapshot; snap != nil {
|
||||
other["billing_mode"] = "tiered_expr"
|
||||
other["expr_b64"] = base64.StdEncoding.EncodeToString([]byte(snap.ExprString))
|
||||
other["matched_tier"] = snap.EstimatedTier
|
||||
if len(snap.UsageFacts) > 0 {
|
||||
other["usage_facts"] = snap.UsageFacts
|
||||
}
|
||||
}
|
||||
appendTaskLogInfo(task, other)
|
||||
attachQuotaSaturation(c, info, other)
|
||||
model.RecordConsumeLog(c, info.UserId, model.RecordConsumeLogParams{
|
||||
ChannelId: info.ChannelId,
|
||||
@@ -132,15 +147,50 @@ func taskBillingOther(task *model.Task) map[string]interface{} {
|
||||
other[k] = v
|
||||
}
|
||||
}
|
||||
if snap := bc.TieredSnapshot; snap != nil {
|
||||
other["billing_mode"] = "tiered_expr"
|
||||
other["expr_b64"] = base64.StdEncoding.EncodeToString([]byte(snap.ExprString))
|
||||
other["matched_tier"] = snap.EstimatedTier
|
||||
if len(snap.UsageFacts) > 0 {
|
||||
other["usage_facts"] = snap.UsageFacts
|
||||
}
|
||||
}
|
||||
}
|
||||
props := task.Properties
|
||||
if props.UpstreamModelName != "" && props.UpstreamModelName != props.OriginModelName {
|
||||
other["is_model_mapped"] = true
|
||||
other["upstream_model_name"] = props.UpstreamModelName
|
||||
}
|
||||
appendTaskLogInfo(task, other)
|
||||
return other
|
||||
}
|
||||
|
||||
func appendTaskLogInfo(task *model.Task, other map[string]interface{}) {
|
||||
if task == nil || other == nil {
|
||||
return
|
||||
}
|
||||
if task.TaskID != "" {
|
||||
other["task_id"] = task.TaskID
|
||||
}
|
||||
if task.PrivateData.Execution != nil {
|
||||
AppendTaskPluginAuditInfo(other, task.PrivateData.Execution.TaskPlugin)
|
||||
}
|
||||
if task.PrivateData.UpstreamTaskID == "" && task.PrivateData.NodeName == "" {
|
||||
return
|
||||
}
|
||||
rootInfo, ok := other["root_info"].(map[string]interface{})
|
||||
if !ok || rootInfo == nil {
|
||||
rootInfo = map[string]interface{}{}
|
||||
other["root_info"] = rootInfo
|
||||
}
|
||||
if task.PrivateData.UpstreamTaskID != "" {
|
||||
rootInfo["upstream_task_id"] = task.PrivateData.UpstreamTaskID
|
||||
}
|
||||
if task.PrivateData.NodeName != "" {
|
||||
rootInfo["node_name"] = task.PrivateData.NodeName
|
||||
}
|
||||
}
|
||||
|
||||
func taskBillingContextPriceData(bc *model.TaskBillingContext) *types.PriceData {
|
||||
if bc == nil || len(bc.OtherRatios) == 0 {
|
||||
return nil
|
||||
@@ -212,7 +262,7 @@ func RefundTaskQuota(ctx context.Context, task *model.Task, reason string) bool
|
||||
// reason 用于日志记录（例如 "token重算" 或 "adaptor调整"）。
|
||||
// clamps 可选：若计算 actualQuota 时发生额度饱和，将其记入日志 admin_info（仅管理员可见）。
|
||||
func RecalculateTaskQuota(ctx context.Context, task *model.Task, actualQuota int, reason string, clamps ...*common.QuotaClamp) {
|
||||
if actualQuota <= 0 {
|
||||
if actualQuota < 0 {
|
||||
return
|
||||
}
|
||||
preConsumedQuota := task.Quota
|
||||
@@ -283,9 +333,9 @@ func RecalculateTaskQuota(ctx context.Context, task *model.Task, actualQuota int
|
||||
// RecalculateTaskQuotaByTokens 根据实际 token 消耗重新计费（异步差额结算）。
|
||||
// 当任务成功且返回了 totalTokens 时，根据模型倍率和分组倍率重新计算实际扣费额度，
|
||||
// 与预扣费的差额进行补扣或退还。支持钱包和订阅计费来源。
|
||||
func RecalculateTaskQuotaByTokens(ctx context.Context, task *model.Task, totalTokens int) {
|
||||
func RecalculateTaskQuotaByTokens(ctx context.Context, task *model.Task, totalTokens int) bool {
|
||||
if totalTokens <= 0 {
|
||||
return
|
||||
return false
|
||||
}
|
||||
|
||||
modelName := taskModelName(task)
|
||||
@@ -294,7 +344,7 @@ func RecalculateTaskQuotaByTokens(ctx context.Context, task *model.Task, totalTo
|
||||
modelRatio, hasRatioSetting, _ := ratio_setting.GetModelRatio(modelName)
|
||||
// 只有配置了倍率(非固定价格)时才按 token 重新计费
|
||||
if !hasRatioSetting || modelRatio <= 0 {
|
||||
return
|
||||
return false
|
||||
}
|
||||
|
||||
// 获取用户和组的倍率信息
|
||||
@@ -306,7 +356,7 @@ func RecalculateTaskQuotaByTokens(ctx context.Context, task *model.Task, totalTo
|
||||
}
|
||||
}
|
||||
if group == "" {
|
||||
return
|
||||
return false
|
||||
}
|
||||
|
||||
groupRatio := ratio_setting.GetGroupRatio(group)
|
||||
@@ -330,4 +380,5 @@ func RecalculateTaskQuotaByTokens(ctx context.Context, task *model.Task, totalTo
|
||||
|
||||
reason := fmt.Sprintf("token重算：tokens=%d, modelRatio=%.2f, groupRatio=%.2f, otherMultiplier=%.4f", totalTokens, modelRatio, finalGroupRatio, otherMultiplier)
|
||||
RecalculateTaskQuota(ctx, task, actualQuota, reason, clamp)
|
||||
return true
|
||||
}
|
||||
|
||||
@@ -2,17 +2,22 @@ package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"math"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
"github.com/QuantumNous/new-api/types"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/shopspring/decimal"
|
||||
"github.com/stretchr/testify/assert"
|
||||
@@ -230,6 +235,194 @@ func TestTaskBillingOtherFiltersHistoricalOtherRatios(t *testing.T) {
|
||||
assert.NotContains(t, other, "negative")
|
||||
assert.NotContains(t, other, "nan")
|
||||
assert.NotContains(t, other, "inf")
|
||||
assert.NotContains(t, other, "billing_mode")
|
||||
assert.NotContains(t, other, "expr_b64")
|
||||
assert.NotContains(t, other, "matched_tier")
|
||||
assert.NotContains(t, other, "usage_facts")
|
||||
}
|
||||
|
||||
func TestTaskBillingOtherIncludesTieredSnapshotAndKeepsUsageFactsNested(t *testing.T) {
|
||||
task := makeTask(1, 1, 100, 0, BillingSourceWallet, 0)
|
||||
expression := `tier("720P", u("seconds") * 5)`
|
||||
task.PrivateData.BillingContext.TieredSnapshot = &billingexpr.BillingSnapshot{
|
||||
ExprString: expression,
|
||||
EstimatedTier: "720P",
|
||||
UsageFacts: map[string]any{
|
||||
"resolution": "720P",
|
||||
"seconds": 5,
|
||||
},
|
||||
}
|
||||
|
||||
other := taskBillingOther(task)
|
||||
|
||||
assert.Equal(t, "tiered_expr", other["billing_mode"])
|
||||
assert.Equal(t, base64.StdEncoding.EncodeToString([]byte(expression)), other["expr_b64"])
|
||||
assert.Equal(t, "720P", other["matched_tier"])
|
||||
facts, ok := other["usage_facts"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, map[string]any{
|
||||
"resolution": "720P",
|
||||
"seconds": 5,
|
||||
}, facts)
|
||||
assert.NotContains(t, other, "resolution")
|
||||
assert.NotContains(t, other, "seconds")
|
||||
}
|
||||
|
||||
func TestTaskBillingOtherOmitsEmptyUsageFacts(t *testing.T) {
|
||||
task := makeTask(1, 1, 100, 0, BillingSourceWallet, 0)
|
||||
expression := `tier("base", 1)`
|
||||
task.PrivateData.BillingContext.TieredSnapshot = &billingexpr.BillingSnapshot{
|
||||
ExprString: expression,
|
||||
EstimatedTier: "base",
|
||||
UsageFacts: map[string]any{},
|
||||
}
|
||||
|
||||
other := taskBillingOther(task)
|
||||
|
||||
assert.Equal(t, "tiered_expr", other["billing_mode"])
|
||||
assert.Equal(t, base64.StdEncoding.EncodeToString([]byte(expression)), other["expr_b64"])
|
||||
assert.Equal(t, "base", other["matched_tier"])
|
||||
assert.NotContains(t, other, "usage_facts")
|
||||
}
|
||||
|
||||
func callLogTaskConsumption(t *testing.T, info *relaycommon.RelayInfo, task *model.Task) *model.Log {
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", nil)
|
||||
ctx.Set("token_name", "test_token")
|
||||
LogTaskConsumption(ctx, info, task)
|
||||
log := getLastLog(t)
|
||||
require.NotNil(t, log)
|
||||
return log
|
||||
}
|
||||
|
||||
func TestLogTaskConsumptionIncludesTieredSnapshotUsageFacts(t *testing.T) {
|
||||
truncate(t)
|
||||
const userID, channelID = 40, 40
|
||||
seedUser(t, userID, 10_000)
|
||||
seedChannel(t, channelID)
|
||||
|
||||
expression := `tier("720P", u("seconds") * 5)`
|
||||
task := makeTask(userID, channelID, 100, 0, BillingSourceWallet, 0)
|
||||
info := &relaycommon.RelayInfo{
|
||||
UserId: userID,
|
||||
TokenId: 0,
|
||||
OriginModelName: "wan2.5-i2v-preview",
|
||||
UsingGroup: "default",
|
||||
ChannelMeta: &relaycommon.ChannelMeta{ChannelId: channelID},
|
||||
TaskRelayInfo: &relaycommon.TaskRelayInfo{Action: "GENERATE"},
|
||||
PriceData: types.PriceData{
|
||||
ModelPrice: 0.02,
|
||||
Quota: 100,
|
||||
GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1},
|
||||
},
|
||||
TieredBillingSnapshot: &billingexpr.BillingSnapshot{
|
||||
ExprString: expression,
|
||||
EstimatedTier: "720P",
|
||||
UsageFacts: map[string]any{
|
||||
"resolution": "720P",
|
||||
"seconds": 5,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
log := callLogTaskConsumption(t, info, task)
|
||||
|
||||
var other map[string]any
|
||||
require.NoError(t, common.UnmarshalJsonStr(log.Other, &other))
|
||||
assert.Equal(t, "tiered_expr", other["billing_mode"])
|
||||
assert.Equal(t, base64.StdEncoding.EncodeToString([]byte(expression)), other["expr_b64"])
|
||||
assert.Equal(t, "720P", other["matched_tier"])
|
||||
facts, ok := other["usage_facts"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, "720P", facts["resolution"])
|
||||
assert.Equal(t, float64(5), facts["seconds"])
|
||||
assert.NotContains(t, other, "resolution")
|
||||
assert.NotContains(t, other, "seconds")
|
||||
assert.Contains(t, log.Content, "计算参数：")
|
||||
assert.Contains(t, log.Content, "resolution: 720P")
|
||||
assert.Contains(t, log.Content, "seconds: 5")
|
||||
}
|
||||
|
||||
func TestLogTaskConsumptionWithoutSnapshotKeepsRatioMode(t *testing.T) {
|
||||
truncate(t)
|
||||
const userID, channelID = 41, 41
|
||||
seedUser(t, userID, 10_000)
|
||||
seedChannel(t, channelID)
|
||||
|
||||
priceData := types.PriceData{
|
||||
ModelPrice: 0.02,
|
||||
Quota: 100,
|
||||
GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1},
|
||||
}
|
||||
priceData.AddOtherRatio("size", 2)
|
||||
task := makeTask(userID, channelID, 100, 0, BillingSourceWallet, 0)
|
||||
info := &relaycommon.RelayInfo{
|
||||
UserId: userID,
|
||||
TokenId: 0,
|
||||
OriginModelName: "test-model",
|
||||
UsingGroup: "default",
|
||||
ChannelMeta: &relaycommon.ChannelMeta{ChannelId: channelID},
|
||||
TaskRelayInfo: &relaycommon.TaskRelayInfo{Action: "GENERATE"},
|
||||
PriceData: priceData,
|
||||
}
|
||||
|
||||
log := callLogTaskConsumption(t, info, task)
|
||||
|
||||
var other map[string]any
|
||||
require.NoError(t, common.UnmarshalJsonStr(log.Other, &other))
|
||||
assert.Equal(t, true, other["is_task"])
|
||||
assert.Equal(t, "/v1/videos", other["request_path"])
|
||||
assert.NotContains(t, other, "billing_mode")
|
||||
assert.NotContains(t, other, "expr_b64")
|
||||
assert.NotContains(t, other, "matched_tier")
|
||||
assert.NotContains(t, other, "usage_facts")
|
||||
assert.Contains(t, log.Content, "计算参数：")
|
||||
assert.Contains(t, log.Content, "size: 2.00")
|
||||
}
|
||||
|
||||
func TestTaskBillingOtherSeparatesPluginAndRootDiagnostics(t *testing.T) {
|
||||
task := makeTask(1, 1, 100, 0, BillingSourceWallet, 0)
|
||||
task.TaskID = "task_public"
|
||||
task.PrivateData.UpstreamTaskID = "upstream-private"
|
||||
task.PrivateData.NodeName = "node-a"
|
||||
task.PrivateData.Execution = &model.TaskExecutionSnapshot{
|
||||
TaskPlugin: &model.TaskPluginSnapshot{
|
||||
Key: "document-parser",
|
||||
Name: "Document Parser",
|
||||
Version: "1.2.3",
|
||||
Author: &model.TaskPluginAuthorSnapshot{
|
||||
Name: "Community Author",
|
||||
URL: "https://plugins.example/author",
|
||||
},
|
||||
APIVersion: 1,
|
||||
Generation: 42,
|
||||
},
|
||||
}
|
||||
|
||||
other := taskBillingOther(task)
|
||||
|
||||
assert.Equal(t, "task_public", other["task_id"])
|
||||
adminInfo, ok := other["admin_info"].(map[string]interface{})
|
||||
require.True(t, ok)
|
||||
pluginInfo, ok := adminInfo["task_plugin"].(map[string]interface{})
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, "document-parser", pluginInfo["key"])
|
||||
assert.Equal(t, "1.2.3", pluginInfo["version"])
|
||||
assert.Equal(t, map[string]interface{}{
|
||||
"name": "Community Author",
|
||||
"url": "https://plugins.example/author",
|
||||
}, pluginInfo["author"])
|
||||
|
||||
rootInfo, ok := other["root_info"].(map[string]interface{})
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, "upstream-private", rootInfo["upstream_task_id"])
|
||||
assert.Equal(t, "node-a", rootInfo["node_name"])
|
||||
runtimeInfo, ok := rootInfo["task_plugin"].(map[string]interface{})
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, uint64(42), runtimeInfo["generation"])
|
||||
assert.NotContains(t, runtimeInfo, "author")
|
||||
}
|
||||
|
||||
func TestTaskBillingContextPriceDataFiltersMultiplier(t *testing.T) {
|
||||
@@ -873,17 +1066,37 @@ func TestRecalculate_ActualQuotaZero(t *testing.T) {
|
||||
truncate(t)
|
||||
ctx := context.Background()
|
||||
|
||||
const userID = 13
|
||||
const userID, preConsumed = 13, 5000
|
||||
const initQuota = 10000
|
||||
|
||||
seedUser(t, userID, initQuota)
|
||||
|
||||
task := makeTask(userID, 0, 5000, 0, BillingSourceWallet, 0)
|
||||
task := makeTask(userID, 0, preConsumed, 0, BillingSourceWallet, 0)
|
||||
require.NoError(t, model.DB.Create(task).Error)
|
||||
|
||||
RecalculateTaskQuota(ctx, task, 0, "zero actual")
|
||||
|
||||
// No change (early return)
|
||||
assert.Equal(t, initQuota+preConsumed, getUserQuota(t, userID))
|
||||
assert.Zero(t, task.Quota)
|
||||
log := getLastLog(t)
|
||||
require.NotNil(t, log)
|
||||
assert.Equal(t, model.LogTypeRefund, log.Type)
|
||||
assert.Equal(t, preConsumed, log.Quota)
|
||||
}
|
||||
|
||||
func TestRecalculate_RejectsNegativeActualQuota(t *testing.T) {
|
||||
truncate(t)
|
||||
ctx := context.Background()
|
||||
|
||||
const userID, preConsumed = 34, 5000
|
||||
const initQuota = 10000
|
||||
seedUser(t, userID, initQuota)
|
||||
task := makeTask(userID, 0, preConsumed, 0, BillingSourceWallet, 0)
|
||||
|
||||
RecalculateTaskQuota(ctx, task, -1, "invalid negative actual")
|
||||
|
||||
assert.Equal(t, initQuota, getUserQuota(t, userID))
|
||||
assert.Equal(t, preConsumed, task.Quota)
|
||||
assert.Equal(t, int64(0), countLogs(t))
|
||||
}
|
||||
|
||||
@@ -1160,9 +1373,10 @@ func TestSettle_PerCallBilling_SkipsAdaptorAdjust(t *testing.T) {
|
||||
adaptor := &mockAdaptor{adjustReturn: 2000}
|
||||
taskResult := &relaycommon.TaskInfo{Status: model.TaskStatusSuccess}
|
||||
|
||||
settleTaskBillingOnComplete(ctx, adaptor, task, taskResult)
|
||||
settled := settleTaskBillingOnComplete(ctx, adaptor, task, taskResult)
|
||||
|
||||
// Per-call: no adjustment despite adaptor returning 2000
|
||||
assert.False(t, settled)
|
||||
assert.Equal(t, initQuota, getUserQuota(t, userID))
|
||||
assert.Equal(t, tokenRemain, getTokenRemainQuota(t, tokenID))
|
||||
assert.Equal(t, preConsumed, task.Quota)
|
||||
@@ -1187,9 +1401,10 @@ func TestSettle_PerCallBilling_SkipsTotalTokens(t *testing.T) {
|
||||
adaptor := &mockAdaptor{adjustReturn: 0}
|
||||
taskResult := &relaycommon.TaskInfo{Status: model.TaskStatusSuccess, TotalTokens: 9999}
|
||||
|
||||
settleTaskBillingOnComplete(ctx, adaptor, task, taskResult)
|
||||
settled := settleTaskBillingOnComplete(ctx, adaptor, task, taskResult)
|
||||
|
||||
// Per-call: no recalculation by tokens
|
||||
assert.False(t, settled)
|
||||
assert.Equal(t, initQuota, getUserQuota(t, userID))
|
||||
assert.Equal(t, tokenRemain, getTokenRemainQuota(t, tokenID))
|
||||
assert.Equal(t, preConsumed, task.Quota)
|
||||
@@ -1215,9 +1430,10 @@ func TestSettle_NonPerCallBilling_AppliesAdaptorAdjustment(t *testing.T) {
|
||||
adaptor := &mockAdaptor{adjustReturn: adaptorQuota}
|
||||
taskResult := &relaycommon.TaskInfo{Status: model.TaskStatusSuccess}
|
||||
|
||||
settleTaskBillingOnComplete(ctx, adaptor, task, taskResult)
|
||||
settled := settleTaskBillingOnComplete(ctx, adaptor, task, taskResult)
|
||||
|
||||
// Non-per-call: adaptor adjustment applies (refund 2000)
|
||||
assert.True(t, settled)
|
||||
assert.Equal(t, initQuota+(preConsumed-adaptorQuota), getUserQuota(t, userID))
|
||||
assert.Equal(t, tokenRemain+(preConsumed-adaptorQuota), getTokenRemainQuota(t, tokenID))
|
||||
assert.Equal(t, adaptorQuota, task.Quota)
|
||||
@@ -1226,3 +1442,303 @@ func TestSettle_NonPerCallBilling_AppliesAdaptorAdjustment(t *testing.T) {
|
||||
require.NotNil(t, log)
|
||||
assert.Equal(t, model.LogTypeRefund, log.Type)
|
||||
}
|
||||
|
||||
func TestSettle_TieredEvaluationFailureKeepsPreConsumedCharge(t *testing.T) {
|
||||
truncate(t)
|
||||
ctx := context.Background()
|
||||
|
||||
const userID, preConsumed = 33, 5_000
|
||||
const initialQuota = 10_000
|
||||
seedUser(t, userID, initialQuota)
|
||||
|
||||
task := makeTask(userID, 0, preConsumed, 0, BillingSourceWallet, 0)
|
||||
task.PrivateData.BillingContext.TieredSnapshot = &billingexpr.BillingSnapshot{
|
||||
ExprString: `tier("broken",`,
|
||||
ExprHash: billingexpr.ExprHashString(`tier("broken",`),
|
||||
GroupRatio: 1,
|
||||
QuotaPerUnit: 1_000,
|
||||
ExprVersion: 1,
|
||||
TaskUsageBilling: true,
|
||||
}
|
||||
|
||||
settled := settleTaskBillingOnComplete(ctx, &mockAdaptor{}, task, &relaycommon.TaskInfo{Status: model.TaskStatusFailure})
|
||||
|
||||
assert.True(t, settled)
|
||||
assert.Equal(t, preConsumed, task.Quota)
|
||||
assert.Equal(t, initialQuota, getUserQuota(t, userID))
|
||||
assert.Equal(t, int64(0), countLogs(t))
|
||||
}
|
||||
|
||||
func TestSettle_TieredFailureReturnsFalseForCallerRefund(t *testing.T) {
|
||||
truncate(t)
|
||||
ctx := context.Background()
|
||||
|
||||
const userID = 37
|
||||
const initialQuota, preConsumed = 10_000, 25
|
||||
seedUser(t, userID, initialQuota)
|
||||
|
||||
expression := `tier("base", u("seconds") + u("clips") * 10)`
|
||||
task := makeTask(userID, 0, preConsumed, 0, BillingSourceWallet, 0)
|
||||
task.Status = model.TaskStatusFailure
|
||||
task.PrivateData.BillingContext.TieredSnapshot = &billingexpr.BillingSnapshot{
|
||||
ExprString: expression,
|
||||
ExprHash: billingexpr.ExprHashString(expression),
|
||||
GroupRatio: 1,
|
||||
QuotaPerUnit: 1,
|
||||
ExprVersion: 1,
|
||||
TaskUsageBilling: true,
|
||||
UsageFacts: map[string]any{"seconds": float64(5), "clips": float64(2)},
|
||||
EstimatedTier: "base",
|
||||
}
|
||||
|
||||
settled := settleTaskBillingOnComplete(
|
||||
ctx,
|
||||
&mockAdaptor{adjustReturn: 1},
|
||||
task,
|
||||
&relaycommon.TaskInfo{Status: model.TaskStatusFailure, UsageFacts: map[string]any{"seconds": float64(8)}},
|
||||
)
|
||||
|
||||
assert.False(t, settled)
|
||||
assert.Equal(t, preConsumed, task.Quota)
|
||||
assert.Equal(t, map[string]any{"seconds": float64(5), "clips": float64(2)}, task.PrivateData.BillingContext.TieredSnapshot.UsageFacts)
|
||||
assert.Equal(t, "base", task.PrivateData.BillingContext.TieredSnapshot.EstimatedTier)
|
||||
assert.Equal(t, initialQuota, getUserQuota(t, userID))
|
||||
assert.Equal(t, int64(0), countLogs(t))
|
||||
}
|
||||
|
||||
func TestSettle_TieredSuccessStillRecomputes(t *testing.T) {
|
||||
truncate(t)
|
||||
ctx := context.Background()
|
||||
|
||||
const userID = 38
|
||||
const initialQuota, preConsumed = 10_000, 50
|
||||
seedUser(t, userID, initialQuota)
|
||||
|
||||
expression := `tier("base", u("seconds") + u("clips") * 10)`
|
||||
task := makeTask(userID, 0, preConsumed, 0, BillingSourceWallet, 0)
|
||||
task.Status = model.TaskStatusSuccess
|
||||
task.PrivateData.BillingContext.TieredSnapshot = &billingexpr.BillingSnapshot{
|
||||
ExprString: expression,
|
||||
ExprHash: billingexpr.ExprHashString(expression),
|
||||
GroupRatio: 1,
|
||||
QuotaPerUnit: 1,
|
||||
ExprVersion: 1,
|
||||
TaskUsageBilling: true,
|
||||
UsageFacts: map[string]any{"seconds": float64(5), "clips": float64(2)},
|
||||
EstimatedTier: "base",
|
||||
}
|
||||
|
||||
settled := settleTaskBillingOnComplete(
|
||||
ctx,
|
||||
&mockAdaptor{adjustReturn: 1},
|
||||
task,
|
||||
&relaycommon.TaskInfo{Status: model.TaskStatusSuccess, UsageFacts: map[string]any{"seconds": float64(8)}},
|
||||
)
|
||||
|
||||
assert.True(t, settled)
|
||||
assert.Equal(t, 28, task.Quota)
|
||||
assert.Equal(t, map[string]any{"seconds": float64(8), "clips": float64(2)}, task.PrivateData.BillingContext.TieredSnapshot.UsageFacts)
|
||||
assert.Equal(t, "base", task.PrivateData.BillingContext.TieredSnapshot.EstimatedTier)
|
||||
assert.Equal(t, initialQuota+(preConsumed-28), getUserQuota(t, userID))
|
||||
|
||||
log := getLastLog(t)
|
||||
require.NotNil(t, log)
|
||||
assert.Equal(t, model.LogTypeRefund, log.Type)
|
||||
var other map[string]any
|
||||
require.NoError(t, common.UnmarshalJsonStr(log.Other, &other))
|
||||
assert.Equal(t, "tiered_expr", other["billing_mode"])
|
||||
assert.Equal(t, "base", other["matched_tier"])
|
||||
facts, ok := other["usage_facts"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, map[string]any{"seconds": float64(8), "clips": float64(2)}, facts)
|
||||
}
|
||||
|
||||
func TestSettle_TieredUsageFactsMergeCompletionOverSubmission(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
completionFacts map[string]any
|
||||
expectedQuota int
|
||||
expectedFacts map[string]any
|
||||
}{
|
||||
{
|
||||
name: "submission facts survive missing completion facts",
|
||||
expectedQuota: 25,
|
||||
expectedFacts: map[string]any{"seconds": float64(5), "clips": float64(2)},
|
||||
},
|
||||
{
|
||||
name: "completion facts partially override submission facts",
|
||||
completionFacts: map[string]any{"seconds": float64(8)},
|
||||
expectedQuota: 28,
|
||||
expectedFacts: map[string]any{"seconds": float64(8), "clips": float64(2)},
|
||||
},
|
||||
{
|
||||
name: "completion facts fully override submission facts",
|
||||
completionFacts: map[string]any{"seconds": float64(8), "clips": float64(3)},
|
||||
expectedQuota: 38,
|
||||
expectedFacts: map[string]any{"seconds": float64(8), "clips": float64(3)},
|
||||
},
|
||||
}
|
||||
|
||||
for _, testCase := range tests {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
truncate(t)
|
||||
const userID = 34
|
||||
const initialQuota = 10_000
|
||||
const preConsumed = 50
|
||||
seedUser(t, userID, initialQuota)
|
||||
|
||||
expression := `tier("base", u("seconds") + u("clips") * 10)`
|
||||
submissionFacts := map[string]any{"seconds": float64(5), "clips": float64(2)}
|
||||
task := makeTask(userID, 0, preConsumed, 0, BillingSourceWallet, 0)
|
||||
task.PrivateData.BillingContext.TieredSnapshot = &billingexpr.BillingSnapshot{
|
||||
ExprString: expression,
|
||||
ExprHash: billingexpr.ExprHashString(expression),
|
||||
GroupRatio: 1,
|
||||
QuotaPerUnit: 1,
|
||||
ExprVersion: 1,
|
||||
TaskUsageBilling: true,
|
||||
UsageFacts: submissionFacts,
|
||||
EstimatedTier: "base",
|
||||
}
|
||||
|
||||
settled := settleTaskBillingOnComplete(
|
||||
context.Background(),
|
||||
&mockAdaptor{},
|
||||
task,
|
||||
&relaycommon.TaskInfo{Status: model.TaskStatusSuccess, UsageFacts: testCase.completionFacts},
|
||||
)
|
||||
|
||||
assert.True(t, settled)
|
||||
assert.Equal(t, testCase.expectedQuota, task.Quota)
|
||||
assert.Equal(t, map[string]any{"seconds": float64(5), "clips": float64(2)}, submissionFacts)
|
||||
require.NotNil(t, task.PrivateData.BillingContext.TieredSnapshot)
|
||||
assert.Equal(t, testCase.expectedFacts, task.PrivateData.BillingContext.TieredSnapshot.UsageFacts)
|
||||
assert.Equal(t, "base", task.PrivateData.BillingContext.TieredSnapshot.EstimatedTier)
|
||||
|
||||
log := getLastLog(t)
|
||||
require.NotNil(t, log)
|
||||
var other map[string]any
|
||||
require.NoError(t, common.UnmarshalJsonStr(log.Other, &other))
|
||||
assert.Equal(t, "tiered_expr", other["billing_mode"])
|
||||
assert.Equal(t, "base", other["matched_tier"])
|
||||
facts, ok := other["usage_facts"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, testCase.expectedFacts, facts)
|
||||
assert.NotContains(t, other, "seconds")
|
||||
assert.NotContains(t, other, "clips")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSettle_TieredSnapshotWriteBackUsesSettledFactsAndMatchedTier(t *testing.T) {
|
||||
truncate(t)
|
||||
const userID = 36
|
||||
const initialQuota = 10_000
|
||||
const preConsumed = 25
|
||||
seedUser(t, userID, initialQuota)
|
||||
|
||||
expression := `u("resolution") == "1080P" ? tier("1080P", u("seconds") * 10) : tier("720P", u("seconds") * 5)`
|
||||
task := makeTask(userID, 0, preConsumed, 0, BillingSourceWallet, 0)
|
||||
task.PrivateData.BillingContext.TieredSnapshot = &billingexpr.BillingSnapshot{
|
||||
ExprString: expression,
|
||||
ExprHash: billingexpr.ExprHashString(expression),
|
||||
GroupRatio: 1,
|
||||
QuotaPerUnit: 1,
|
||||
ExprVersion: 1,
|
||||
TaskUsageBilling: true,
|
||||
UsageFacts: map[string]any{"resolution": "720P", "seconds": float64(5)},
|
||||
EstimatedTier: "720P",
|
||||
}
|
||||
|
||||
settled := settleTaskBillingOnComplete(
|
||||
context.Background(),
|
||||
&mockAdaptor{},
|
||||
task,
|
||||
&relaycommon.TaskInfo{
|
||||
Status: model.TaskStatusSuccess,
|
||||
UsageFacts: map[string]any{"resolution": "1080P"},
|
||||
},
|
||||
)
|
||||
|
||||
require.True(t, settled)
|
||||
snap := task.PrivateData.BillingContext.TieredSnapshot
|
||||
require.NotNil(t, snap)
|
||||
assert.Equal(t, map[string]any{"resolution": "1080P", "seconds": float64(5)}, snap.UsageFacts)
|
||||
assert.Equal(t, "1080P", snap.EstimatedTier)
|
||||
assert.Equal(t, 50, task.Quota)
|
||||
|
||||
log := getLastLog(t)
|
||||
require.NotNil(t, log)
|
||||
var other map[string]any
|
||||
require.NoError(t, common.UnmarshalJsonStr(log.Other, &other))
|
||||
assert.Equal(t, "tiered_expr", other["billing_mode"])
|
||||
assert.Equal(t, "1080P", other["matched_tier"])
|
||||
facts, ok := other["usage_facts"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, "1080P", facts["resolution"])
|
||||
assert.Equal(t, float64(5), facts["seconds"])
|
||||
assert.NotContains(t, other, "resolution")
|
||||
assert.NotContains(t, other, "seconds")
|
||||
}
|
||||
|
||||
func TestSettle_TokenRecalcFallsBackToCompletionTokens(t *testing.T) {
|
||||
previousRatios := ratio_setting.ModelRatio2JSONString()
|
||||
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"test-model":1}`))
|
||||
t.Cleanup(func() {
|
||||
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(previousRatios))
|
||||
})
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
totalTokens int
|
||||
completionTokens int
|
||||
wantSettled bool
|
||||
wantQuota int
|
||||
}{
|
||||
{
|
||||
name: "total tokens still win when both are present",
|
||||
totalTokens: 80,
|
||||
completionTokens: 20,
|
||||
wantSettled: true,
|
||||
wantQuota: 80,
|
||||
},
|
||||
{
|
||||
name: "completion tokens trigger recalc when total is zero",
|
||||
totalTokens: 0,
|
||||
completionTokens: 80,
|
||||
wantSettled: true,
|
||||
wantQuota: 80,
|
||||
},
|
||||
{
|
||||
name: "neither token count skips recalc",
|
||||
wantSettled: false,
|
||||
wantQuota: 50,
|
||||
},
|
||||
}
|
||||
|
||||
for _, testCase := range tests {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
truncate(t)
|
||||
const userID, tokenID, channelID = 35, 35, 35
|
||||
const initialQuota, preConsumed, tokenRemain = 10_000, 50, 8_000
|
||||
seedUser(t, userID, initialQuota)
|
||||
seedToken(t, tokenID, userID, "sk-completion-fallback", tokenRemain)
|
||||
seedChannel(t, channelID)
|
||||
|
||||
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
|
||||
settled := settleTaskBillingOnComplete(
|
||||
context.Background(),
|
||||
&mockAdaptor{},
|
||||
task,
|
||||
&relaycommon.TaskInfo{
|
||||
Status: model.TaskStatusSuccess,
|
||||
TotalTokens: testCase.totalTokens,
|
||||
CompletionTokens: testCase.completionTokens,
|
||||
},
|
||||
)
|
||||
|
||||
assert.Equal(t, testCase.wantSettled, settled)
|
||||
assert.Equal(t, testCase.wantQuota, task.Quota)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// TaskExecutionSnapshotFromContext captures immutable request and plugin
|
||||
// provenance at submission time. It never copies plugin source or payloads.
|
||||
func TaskExecutionSnapshotFromContext(ctx *gin.Context) *model.TaskExecutionSnapshot {
|
||||
if ctx == nil {
|
||||
return nil
|
||||
}
|
||||
snapshot := &model.TaskExecutionSnapshot{
|
||||
RequestID: ctx.GetString(common.RequestIdKey),
|
||||
}
|
||||
if ctx.Request != nil && ctx.Request.URL != nil {
|
||||
snapshot.RequestPath = ctx.Request.URL.Path
|
||||
}
|
||||
|
||||
pinnedValue, exists := ctx.Get(pluginruntime.ContextKeyPinnedPlugin)
|
||||
if exists {
|
||||
pinned, ok := pinnedValue.(pluginruntime.PinnedPlugin)
|
||||
if ok && pinned.Plugin != nil {
|
||||
generation := uint64(0)
|
||||
if pinned.Generation != nil {
|
||||
generation = pinned.Generation.Number
|
||||
}
|
||||
meta := pinned.Plugin.Meta
|
||||
snapshot.TaskPlugin = &model.TaskPluginSnapshot{
|
||||
Key: meta.Key,
|
||||
Name: meta.Name,
|
||||
Version: meta.Version,
|
||||
Author: &model.TaskPluginAuthorSnapshot{
|
||||
Name: meta.Author.Name,
|
||||
URL: meta.Author.URL,
|
||||
},
|
||||
APIVersion: meta.APIVersion,
|
||||
Generation: generation,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if snapshot.RequestID == "" && snapshot.RequestPath == "" && snapshot.TaskPlugin == nil {
|
||||
return nil
|
||||
}
|
||||
return snapshot
|
||||
}
|
||||
|
||||
// AppendTaskPluginAuditInfo writes role-separated, credential-free plugin
|
||||
// provenance into a usage log.
|
||||
func AppendTaskPluginAuditInfo(other map[string]interface{}, snapshot *model.TaskPluginSnapshot) {
|
||||
if other == nil || snapshot == nil || snapshot.Key == "" {
|
||||
return
|
||||
}
|
||||
adminInfo, ok := other["admin_info"].(map[string]interface{})
|
||||
if !ok || adminInfo == nil {
|
||||
adminInfo = map[string]interface{}{}
|
||||
other["admin_info"] = adminInfo
|
||||
}
|
||||
taskPlugin := map[string]interface{}{
|
||||
"key": snapshot.Key,
|
||||
"name": snapshot.Name,
|
||||
"version": snapshot.Version,
|
||||
}
|
||||
if snapshot.Author != nil && snapshot.Author.Name != "" {
|
||||
author := map[string]interface{}{"name": snapshot.Author.Name}
|
||||
if snapshot.Author.URL != "" {
|
||||
author["url"] = snapshot.Author.URL
|
||||
}
|
||||
taskPlugin["author"] = author
|
||||
}
|
||||
adminInfo["task_plugin"] = taskPlugin
|
||||
|
||||
rootInfo, ok := other["root_info"].(map[string]interface{})
|
||||
if !ok || rootInfo == nil {
|
||||
rootInfo = map[string]interface{}{}
|
||||
other["root_info"] = rootInfo
|
||||
}
|
||||
rootInfo["task_plugin"] = map[string]interface{}{
|
||||
"key": snapshot.Key,
|
||||
"version": snapshot.Version,
|
||||
"api_version": snapshot.APIVersion,
|
||||
"generation": snapshot.Generation,
|
||||
}
|
||||
}
|
||||
|
||||
// AppendTaskPluginContextAuditInfo is used before a task row exists, such as
|
||||
// an upstream submission error log.
|
||||
func AppendTaskPluginContextAuditInfo(ctx *gin.Context, other map[string]interface{}) {
|
||||
execution := TaskExecutionSnapshotFromContext(ctx)
|
||||
if execution == nil {
|
||||
return
|
||||
}
|
||||
AppendTaskPluginAuditInfo(other, execution.TaskPlugin)
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
)
|
||||
|
||||
// BuildTaskPluginView converts a persisted task into the deliberately narrow
|
||||
// public shape permitted at JavaScript plugin boundaries.
|
||||
func BuildTaskPluginView(task *model.Task) (dto.TaskView, error) {
|
||||
createdAt := task.CreatedAt
|
||||
if createdAt == 0 {
|
||||
createdAt = task.SubmitTime
|
||||
}
|
||||
view := dto.TaskView{
|
||||
TaskID: task.TaskID,
|
||||
Platform: string(task.Platform),
|
||||
Status: string(task.Status),
|
||||
Progress: task.Progress,
|
||||
FailReason: task.FailReason,
|
||||
CreatedAt: createdAt,
|
||||
UpdatedAt: task.UpdatedAt,
|
||||
FinishedAt: task.FinishTime,
|
||||
}
|
||||
if len(task.Data) > 0 {
|
||||
if err := common.Unmarshal(task.Data, &view.Data); err != nil {
|
||||
return dto.TaskView{}, err
|
||||
}
|
||||
view.Data = replacePrivateTaskID(view.Data, task.PrivateData.UpstreamTaskID, task.TaskID)
|
||||
}
|
||||
return view, nil
|
||||
}
|
||||
|
||||
// replacePrivateTaskID rewrites exact private IDs only in known task-ID fields.
|
||||
// Map keys and opaque strings, including URLs containing the ID, are preserved.
|
||||
func replacePrivateTaskID(value any, privateTaskID, publicTaskID string) any {
|
||||
if privateTaskID == "" || privateTaskID == publicTaskID {
|
||||
return value
|
||||
}
|
||||
switch typed := value.(type) {
|
||||
case []any:
|
||||
replaced := make([]any, len(typed))
|
||||
for index, item := range typed {
|
||||
replaced[index] = replacePrivateTaskID(item, privateTaskID, publicTaskID)
|
||||
}
|
||||
return replaced
|
||||
case map[string]any:
|
||||
replaced := make(map[string]any, len(typed))
|
||||
for key, item := range typed {
|
||||
if (key == "id" || key == "task_id" || key == "taskId") && item == privateTaskID {
|
||||
replaced[key] = publicTaskID
|
||||
continue
|
||||
}
|
||||
replaced[key] = replacePrivateTaskID(item, privateTaskID, publicTaskID)
|
||||
}
|
||||
return replaced
|
||||
default:
|
||||
return value
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
)
|
||||
|
||||
func TestBuildTaskPluginViewRewritesOnlyStructuredTaskIDFields(t *testing.T) {
|
||||
const (
|
||||
privateTaskID = "upstream-task-123"
|
||||
publicTaskID = "task_public_123"
|
||||
resultURL = "https://cdn.example.com/results/upstream-task-123/video.mp4"
|
||||
)
|
||||
|
||||
taskData, err := common.Marshal(map[string]any{
|
||||
"task_id": privateTaskID,
|
||||
"id": privateTaskID,
|
||||
"taskId": privateTaskID,
|
||||
"url": resultURL,
|
||||
"message": "completed upstream-task-123",
|
||||
"nested": []any{
|
||||
map[string]any{
|
||||
"task_id": privateTaskID,
|
||||
"url": resultURL,
|
||||
},
|
||||
privateTaskID,
|
||||
},
|
||||
privateTaskID: "opaque map key",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
task := &model.Task{
|
||||
TaskID: publicTaskID,
|
||||
PrivateData: model.TaskPrivateData{
|
||||
UpstreamTaskID: privateTaskID,
|
||||
},
|
||||
Data: taskData,
|
||||
}
|
||||
|
||||
view, err := BuildTaskPluginView(task)
|
||||
require.NoError(t, err)
|
||||
|
||||
data, ok := view.Data.(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, publicTaskID, data["task_id"])
|
||||
assert.Equal(t, publicTaskID, data["id"])
|
||||
assert.Equal(t, publicTaskID, data["taskId"])
|
||||
assert.Equal(t, resultURL, data["url"])
|
||||
assert.Equal(t, "completed upstream-task-123", data["message"])
|
||||
assert.Equal(t, "opaque map key", data[privateTaskID])
|
||||
|
||||
nested, ok := data["nested"].([]any)
|
||||
require.True(t, ok)
|
||||
nestedData, ok := nested[0].(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, publicTaskID, nestedData["task_id"])
|
||||
assert.Equal(t, resultURL, nestedData["url"])
|
||||
assert.Equal(t, privateTaskID, nested[1])
|
||||
|
||||
}
|
||||
+129
-119
@@ -2,7 +2,6 @@ package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
@@ -16,6 +15,7 @@ import (
|
||||
taskdto "github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
"github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
@@ -34,6 +34,22 @@ type TaskPollingAdaptor interface {
|
||||
AdjustBillingOnComplete(task *model.Task, taskResult *relaycommon.TaskInfo) int
|
||||
}
|
||||
|
||||
type BatchTaskPollingAdaptor interface {
|
||||
TaskPollingAdaptor
|
||||
FetchMode() string
|
||||
FetchBatchTasks(baseURL, key string, taskIDs []string, proxy string) (*http.Response, error)
|
||||
ParseBatchResult(body []byte) (map[string]*BatchTaskResult, 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
|
||||
@@ -180,25 +196,28 @@ func DispatchPlatformUpdate(ctx context.Context, platform constant.TaskPlatform,
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
switch platform {
|
||||
case constant.TaskPlatformMidjourney:
|
||||
if platform == constant.TaskPlatformMidjourney {
|
||||
// MJ 轮询由其自身处理，这里预留入口
|
||||
case constant.TaskPlatformSuno:
|
||||
_ = UpdateSunoTasks(ctx, taskChannelM, taskM)
|
||||
default:
|
||||
if err := UpdateVideoTasks(ctx, platform, taskChannelM, taskM); err != nil {
|
||||
common.SysLog(fmt.Sprintf("UpdateVideoTasks fail: %s", err))
|
||||
return
|
||||
}
|
||||
adaptor := GetTaskAdaptorFunc(platform)
|
||||
if batchAdaptor, ok := adaptor.(BatchTaskPollingAdaptor); ok && batchAdaptor.FetchMode() == "batch" {
|
||||
if err := UpdateBatchTasks(ctx, batchAdaptor, taskChannelM, taskM); err != nil {
|
||||
common.SysLog(fmt.Sprintf("UpdateBatchTasks fail: %s", err))
|
||||
}
|
||||
return
|
||||
}
|
||||
if err := UpdateVideoTasks(ctx, platform, taskChannelM, taskM); err != nil {
|
||||
common.SysLog(fmt.Sprintf("UpdateVideoTasks fail: %s", err))
|
||||
}
|
||||
}
|
||||
|
||||
// UpdateSunoTasks 按渠道更新所有 Suno 任务
|
||||
func UpdateSunoTasks(ctx context.Context, taskChannelM map[int][]string, taskM map[string]*model.Task) error {
|
||||
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 := updateSunoTasks(ctx, channelId, taskIds, taskM)
|
||||
err := updateBatchTasks(ctx, adaptor, channelId, taskIds, taskM)
|
||||
if err != nil {
|
||||
logger.LogError(ctx, fmt.Sprintf("渠道 #%d 更新异步任务失败: %s", channelId, err.Error()))
|
||||
}
|
||||
@@ -206,7 +225,7 @@ func UpdateSunoTasks(ctx context.Context, taskChannelM map[int][]string, taskM m
|
||||
return nil
|
||||
}
|
||||
|
||||
func updateSunoTasks(ctx context.Context, channelId int, taskIds []string, taskM map[string]*model.Task) error {
|
||||
func updateBatchTasks(ctx context.Context, adaptor BatchTaskPollingAdaptor, channelId int, taskIds []string, taskM map[string]*model.Task) error {
|
||||
logger.LogInfo(ctx, fmt.Sprintf("渠道 #%d 未完成的任务有: %d", channelId, len(taskIds)))
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
@@ -234,14 +253,12 @@ func updateSunoTasks(ctx context.Context, channelId int, taskIds []string, taskM
|
||||
}
|
||||
return err
|
||||
}
|
||||
adaptor := GetTaskAdaptorFunc(constant.TaskPlatformSuno)
|
||||
if adaptor == nil {
|
||||
return errors.New("adaptor not found")
|
||||
}
|
||||
proxy := ch.GetSetting().Proxy
|
||||
resp, err := adaptor.FetchTask(*ch.BaseURL, ch.Key, map[string]any{
|
||||
"ids": taskIds,
|
||||
}, proxy)
|
||||
baseURL := ch.GetBaseURL()
|
||||
if baseURL == "" {
|
||||
baseURL = constant.GetChannelBaseURL(ch.Type)
|
||||
}
|
||||
resp, err := adaptor.FetchBatchTasks(baseURL, ch.Key, taskIds, proxy)
|
||||
if err != nil {
|
||||
common.SysLog(fmt.Sprintf("Get Task Do req error: %v", err))
|
||||
return err
|
||||
@@ -256,98 +273,69 @@ func updateSunoTasks(ctx context.Context, channelId int, taskIds []string, taskM
|
||||
common.SysLog(fmt.Sprintf("Get Suno Task parse body error: %v", err))
|
||||
return err
|
||||
}
|
||||
var responseItems taskdto.TaskResponse[[]taskdto.SunoDataResponse]
|
||||
err = common.Unmarshal(responseBody, &responseItems)
|
||||
responseItems, err := adaptor.ParseBatchResult(responseBody)
|
||||
if err != nil {
|
||||
logger.LogError(ctx, fmt.Sprintf("Get Suno Task parse body error2: %v, body: %s", err, string(responseBody)))
|
||||
return err
|
||||
return fmt.Errorf("parse batch result: %w", err)
|
||||
}
|
||||
if !responseItems.IsSuccess() {
|
||||
common.SysLog(fmt.Sprintf("渠道 #%d 未完成的任务有: %d, 成功获取到任务数: %s", channelId, len(taskIds), string(responseBody)))
|
||||
return err
|
||||
}
|
||||
|
||||
for _, responseItem := range responseItems.Data {
|
||||
for upstreamID, responseItem := range responseItems {
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
task := taskM[responseItem.TaskID]
|
||||
task := taskM[upstreamID]
|
||||
if task == nil {
|
||||
logger.LogWarn(ctx, fmt.Sprintf("Suno task response ignored: unknown task_id=%s", responseItem.TaskID))
|
||||
logger.LogWarn(ctx, fmt.Sprintf("Batch task response ignored: unknown task_id=%s", upstreamID))
|
||||
continue
|
||||
}
|
||||
if !taskNeedsUpdate(task, responseItem) {
|
||||
continue
|
||||
}
|
||||
|
||||
prevStatus := task.Status
|
||||
task.Status = lo.If(model.TaskStatus(responseItem.Status) != "", model.TaskStatus(responseItem.Status)).Else(task.Status)
|
||||
task.FailReason = lo.If(responseItem.FailReason != "", responseItem.FailReason).Else(task.FailReason)
|
||||
snap := task.Snapshot()
|
||||
task.Status = lo.If(model.TaskStatus(responseItem.TaskInfo.Status) != "", model.TaskStatus(responseItem.TaskInfo.Status)).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)
|
||||
isFailure := responseItem.FailReason != "" || task.Status == model.TaskStatusFailure
|
||||
if isFailure {
|
||||
if responseItem.TaskInfo.Progress != "" {
|
||||
task.Progress = responseItem.TaskInfo.Progress
|
||||
}
|
||||
if responseItem.TaskInfo.Reason != "" || task.Status == model.TaskStatusFailure {
|
||||
logger.LogInfo(ctx, task.TaskID+" 构建失败，"+task.FailReason)
|
||||
task.Status = model.TaskStatusFailure
|
||||
task.Progress = "100%"
|
||||
}
|
||||
if responseItem.Status == model.TaskStatusSuccess {
|
||||
if responseItem.TaskInfo.Status == model.TaskStatusSuccess {
|
||||
task.Progress = "100%"
|
||||
}
|
||||
task.Data = responseItem.Data
|
||||
if responseItem.Data != nil {
|
||||
task.SetData(responseItem.Data)
|
||||
} else if task.Status == model.TaskStatusSuccess || task.Status == model.TaskStatusFailure {
|
||||
logger.LogWarn(ctx, fmt.Sprintf(
|
||||
"Batch task %s reached terminal status without data; preserving existing task data",
|
||||
task.TaskID,
|
||||
))
|
||||
}
|
||||
if responseItem.TaskInfo.Url != "" {
|
||||
task.PrivateData.ResultURL = responseItem.TaskInfo.Url
|
||||
}
|
||||
|
||||
// 持久化走 CAS，防止重叠轮询/sweep/多实例/持久化失败重试导致重复退款或覆盖终态。
|
||||
won, err := task.UpdateWithStatus(prevStatus)
|
||||
if err != nil {
|
||||
logger.LogError(ctx, fmt.Sprintf("UpdateSunoTask task %s error: %v", task.TaskID, err))
|
||||
} else if !won {
|
||||
logger.LogWarn(ctx, fmt.Sprintf("Task %s CAS lost or no-op update, skip billing", task.TaskID))
|
||||
} else if isFailure && prevStatus != model.TaskStatusFailure && task.Quota != 0 {
|
||||
RefundTaskQuota(ctx, task, task.FailReason)
|
||||
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("UpdateSunoTask task error: " + updateErr.Error())
|
||||
continue
|
||||
}
|
||||
if !won {
|
||||
logger.LogWarn(ctx, fmt.Sprintf("Batch task %s already transitioned by another process, skip billing", task.TaskID))
|
||||
continue
|
||||
}
|
||||
if terminalTransition {
|
||||
billingSettled := settleTaskBillingOnComplete(ctx, adaptor, task, &responseItem.TaskInfo)
|
||||
if task.Status == model.TaskStatusFailure && !billingSettled && task.Quota != 0 {
|
||||
RefundTaskQuota(ctx, task, task.FailReason)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// taskNeedsUpdate 检查 Suno 任务是否需要更新
|
||||
func taskNeedsUpdate(oldTask *model.Task, newTask taskdto.SunoDataResponse) bool {
|
||||
if oldTask.SubmitTime != newTask.SubmitTime {
|
||||
return true
|
||||
}
|
||||
if oldTask.StartTime != newTask.StartTime {
|
||||
return true
|
||||
}
|
||||
if oldTask.FinishTime != newTask.FinishTime {
|
||||
return true
|
||||
}
|
||||
if string(oldTask.Status) != newTask.Status {
|
||||
return true
|
||||
}
|
||||
if oldTask.FailReason != newTask.FailReason {
|
||||
return true
|
||||
}
|
||||
|
||||
if (oldTask.Status == model.TaskStatusFailure || oldTask.Status == model.TaskStatusSuccess) && oldTask.Progress != "100%" {
|
||||
return true
|
||||
}
|
||||
|
||||
oldData, _ := common.Marshal(oldTask.Data)
|
||||
newData, _ := common.Marshal(newTask.Data)
|
||||
|
||||
sort.Slice(oldData, func(i, j int) bool {
|
||||
return oldData[i] < oldData[j]
|
||||
})
|
||||
sort.Slice(newData, func(i, j int) bool {
|
||||
return newData[i] < newData[j]
|
||||
})
|
||||
|
||||
if string(oldData) != string(newData) {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// 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))
|
||||
@@ -442,7 +430,7 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch *
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
baseURL := constant.ChannelBaseURLs[ch.Type]
|
||||
baseURL := constant.GetChannelBaseURL(ch.Type)
|
||||
if ch.GetBaseURL() != "" {
|
||||
baseURL = ch.GetBaseURL()
|
||||
}
|
||||
@@ -461,7 +449,7 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch *
|
||||
}
|
||||
resp, err := adaptor.FetchTask(baseURL, key, map[string]any{
|
||||
"task_id": task.GetUpstreamTaskID(),
|
||||
"action": task.Action,
|
||||
"action": constant.NormalizeTaskAction(task.Action),
|
||||
}, proxy)
|
||||
if err != nil {
|
||||
return fmt.Errorf("fetchTask failed for task %s: %w", taskId, err)
|
||||
@@ -519,9 +507,7 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch *
|
||||
}
|
||||
}
|
||||
|
||||
shouldRefund := false
|
||||
shouldSettle := false
|
||||
quota := task.Quota
|
||||
shouldFinalizeBilling := false
|
||||
|
||||
task.Status = model.TaskStatus(taskResult.Status)
|
||||
switch taskResult.Status {
|
||||
@@ -549,7 +535,7 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch *
|
||||
// No URL from adaptor — construct proxy URL using public task ID
|
||||
task.PrivateData.ResultURL = taskcommon.BuildProxyURL(task.TaskID)
|
||||
}
|
||||
shouldSettle = true
|
||||
shouldFinalizeBilling = true
|
||||
case model.TaskStatusFailure:
|
||||
logger.LogJson(ctx, fmt.Sprintf("Task %s failed", taskId), task)
|
||||
task.Status = model.TaskStatusFailure
|
||||
@@ -560,9 +546,7 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch *
|
||||
task.FailReason = taskResult.Reason
|
||||
logger.LogInfo(ctx, fmt.Sprintf("Task %s failed: %s", task.TaskID, task.FailReason))
|
||||
taskResult.Progress = taskcommon.ProgressComplete
|
||||
if quota != 0 {
|
||||
shouldRefund = true
|
||||
}
|
||||
shouldFinalizeBilling = true
|
||||
default:
|
||||
return fmt.Errorf("unknown task status %s for task %s", taskResult.Status, task.TaskID)
|
||||
}
|
||||
@@ -575,12 +559,10 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch *
|
||||
won, err := task.UpdateWithStatus(snap.Status)
|
||||
if err != nil {
|
||||
logger.LogError(ctx, fmt.Sprintf("UpdateWithStatus failed for task %s: %s", task.TaskID, err.Error()))
|
||||
shouldRefund = false
|
||||
shouldSettle = false
|
||||
shouldFinalizeBilling = false
|
||||
} else if !won {
|
||||
logger.LogWarn(ctx, fmt.Sprintf("Task %s CAS lost or no-op update, skip billing", task.TaskID))
|
||||
shouldRefund = false
|
||||
shouldSettle = false
|
||||
shouldFinalizeBilling = false
|
||||
}
|
||||
} else if !snap.Equal(task.Snapshot()) {
|
||||
if _, err := task.UpdateWithStatus(snap.Status); err != nil {
|
||||
@@ -591,11 +573,11 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch *
|
||||
logger.LogDebug(ctx, "No update needed for task %s", task.TaskID)
|
||||
}
|
||||
|
||||
if shouldSettle {
|
||||
settleTaskBillingOnComplete(ctx, adaptor, task, taskResult)
|
||||
}
|
||||
if shouldRefund {
|
||||
RefundTaskQuota(ctx, task, task.FailReason)
|
||||
if shouldFinalizeBilling {
|
||||
billingSettled := settleTaskBillingOnComplete(ctx, adaptor, task, taskResult)
|
||||
if task.Status == model.TaskStatusFailure && !billingSettled && task.Quota != 0 {
|
||||
RefundTaskQuota(ctx, task, task.FailReason)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
@@ -636,25 +618,53 @@ func truncateBase64(s string) string {
|
||||
}
|
||||
|
||||
// settleTaskBillingOnComplete 任务完成时的统一计费调整。
|
||||
// 优先级：1. adaptor.AdjustBillingOnComplete 返回正数 → 使用 adaptor 计算的额度
|
||||
// 返回 true 表示用量结算路径已接管最终计费；失败任务仅在返回 false 时补做全额退款。
|
||||
// 优先级：1. tiered snapshot → 2. adaptor 调整 → 3. token 重算。
|
||||
//
|
||||
// 2. taskResult.TotalTokens > 0 → 按 token 重算
|
||||
// 3. 都不满足 → 保持预扣额度不变
|
||||
func settleTaskBillingOnComplete(ctx context.Context, adaptor TaskPollingAdaptor, task *model.Task, taskResult *relaycommon.TaskInfo) {
|
||||
// 0. 按次计费的任务不做差额结算
|
||||
// 表达式求值失败会保留预扣额度，因此也视为已接管，避免错误全退。
|
||||
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
|
||||
}
|
||||
usageFacts := make(map[string]any, len(bc.TieredSnapshot.UsageFacts)+len(taskResult.UsageFacts))
|
||||
for key, value := range bc.TieredSnapshot.UsageFacts {
|
||||
usageFacts[key] = value
|
||||
}
|
||||
for key, value := range taskResult.UsageFacts {
|
||||
usageFacts[key] = value
|
||||
}
|
||||
result, err := billingexpr.ComputeTieredQuotaWithRequest(bc.TieredSnapshot, billingexpr.TokenParams{}, billingexpr.RequestInput{Usage: usageFacts})
|
||||
if err != nil {
|
||||
logger.LogWarn(ctx, fmt.Sprintf("任务 %s 表达式结算失败，保留预扣额度: %v", task.TaskID, err))
|
||||
return true
|
||||
}
|
||||
if result.Clamp != nil {
|
||||
logger.LogWarn(ctx, fmt.Sprintf("任务 %s 表达式结算额度发生饱和: %+v", task.TaskID, result.Clamp))
|
||||
}
|
||||
bc.TieredSnapshot.UsageFacts = usageFacts
|
||||
bc.TieredSnapshot.EstimatedTier = result.MatchedTier
|
||||
RecalculateTaskQuota(ctx, task, result.ActualQuotaAfterGroup, "任务用量表达式结算", result.Clamp)
|
||||
return true
|
||||
}
|
||||
// 按次计费的成功任务保持预扣；失败任务由调用方全额退款。
|
||||
if bc := task.PrivateData.BillingContext; bc != nil && bc.PerCallBilling {
|
||||
logger.LogInfo(ctx, fmt.Sprintf("任务 %s 按次计费，跳过差额结算", task.TaskID))
|
||||
return
|
||||
return false
|
||||
}
|
||||
// 1. 优先让 adaptor 决定最终额度
|
||||
// 优先让 adaptor 决定最终额度。
|
||||
if actualQuota := adaptor.AdjustBillingOnComplete(task, taskResult); actualQuota > 0 {
|
||||
RecalculateTaskQuota(ctx, task, actualQuota, "adaptor计费调整")
|
||||
return
|
||||
return true
|
||||
}
|
||||
// 2. 回退到 token 重算
|
||||
if taskResult.TotalTokens > 0 {
|
||||
RecalculateTaskQuotaByTokens(ctx, task, taskResult.TotalTokens)
|
||||
return
|
||||
// 回退到 token 重算。
|
||||
tokens := taskResult.TotalTokens
|
||||
if tokens == 0 && taskResult.CompletionTokens > 0 {
|
||||
tokens = taskResult.CompletionTokens
|
||||
}
|
||||
// 3. 无调整，保持预扣额度
|
||||
if tokens > 0 {
|
||||
return RecalculateTaskQuotaByTokens(ctx, task, tokens)
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
+262
-37
@@ -5,6 +5,7 @@ import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -13,6 +14,7 @@ import (
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
taskdto "github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/bytedance/gopkg/util/gopool"
|
||||
@@ -30,43 +32,28 @@ type taskPollingFetchAdaptor struct {
|
||||
blockOnce sync.Once
|
||||
}
|
||||
|
||||
type sunoFailurePollingAdaptor struct {
|
||||
failReason string
|
||||
type batchPollingAdaptor struct {
|
||||
taskPollingFetchAdaptor
|
||||
batchCalls int
|
||||
batchIDs []string
|
||||
results map[string]*BatchTaskResult
|
||||
}
|
||||
|
||||
func (a *sunoFailurePollingAdaptor) Init(_ *relaycommon.RelayInfo) {}
|
||||
|
||||
func (a *sunoFailurePollingAdaptor) FetchTask(_ string, _ string, body map[string]any, _ string) (*http.Response, error) {
|
||||
taskIDs, _ := body["ids"].([]string)
|
||||
items := make([]taskdto.SunoDataResponse, 0, len(taskIDs))
|
||||
for _, taskID := range taskIDs {
|
||||
items = append(items, taskdto.SunoDataResponse{
|
||||
TaskID: taskID,
|
||||
Status: string(model.TaskStatusFailure),
|
||||
FailReason: a.failReason,
|
||||
FinishTime: time.Now().Unix(),
|
||||
})
|
||||
func (a *batchPollingAdaptor) FetchMode() string { return "batch" }
|
||||
func (a *batchPollingAdaptor) FetchBatchTasks(_ string, _ string, taskIDs []string, _ string) (*http.Response, error) {
|
||||
a.batchCalls++
|
||||
a.batchIDs = append([]string(nil), taskIDs...)
|
||||
return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(bytes.NewReader([]byte(`{}`)))}, nil
|
||||
}
|
||||
func (a *batchPollingAdaptor) ParseBatchResult([]byte) (map[string]*BatchTaskResult, error) {
|
||||
if a.results != nil {
|
||||
return a.results, nil
|
||||
}
|
||||
|
||||
responseBody, err := common.Marshal(taskdto.TaskResponse[[]taskdto.SunoDataResponse]{
|
||||
Code: taskdto.TaskSuccessCode,
|
||||
Data: items,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
results := make(map[string]*BatchTaskResult, len(a.batchIDs))
|
||||
for _, taskID := range a.batchIDs {
|
||||
results[taskID] = &BatchTaskResult{TaskInfo: relaycommon.TaskInfo{TaskID: taskID, Status: model.TaskStatusInProgress, Url: "https://example.com/result"}}
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(bytes.NewReader(responseBody)),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (a *sunoFailurePollingAdaptor) ParseTaskResult([]byte) (*relaycommon.TaskInfo, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (a *sunoFailurePollingAdaptor) AdjustBillingOnComplete(_ *model.Task, _ *relaycommon.TaskInfo) int {
|
||||
return 0
|
||||
return results, nil
|
||||
}
|
||||
|
||||
func (a *taskPollingFetchAdaptor) Init(_ *relaycommon.RelayInfo) {}
|
||||
@@ -130,6 +117,43 @@ func (a *taskPollingFetchAdaptor) fetchedTaskIDs() []string {
|
||||
return append([]string(nil), a.taskIDs...)
|
||||
}
|
||||
|
||||
func TestRedactVideoResponseBodyPreservesPollingPayloadShape(t *testing.T) {
|
||||
rawVideo := strings.Repeat("a", 300)
|
||||
body, err := common.Marshal(map[string]any{
|
||||
"done": true,
|
||||
"name": "operations/provider-task",
|
||||
"response": map[string]any{
|
||||
"bytesBase64Encoded": "secret-bytes",
|
||||
"video": rawVideo,
|
||||
"videos": []any{
|
||||
map[string]any{
|
||||
"bytesBase64Encoded": "other-secret-bytes",
|
||||
"mimeType": "video/mp4",
|
||||
"uri": "https://media.example/video.mp4",
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
var stored map[string]any
|
||||
require.NoError(t, common.Unmarshal(redactVideoResponseBody(body), &stored))
|
||||
assert.Equal(t, true, stored["done"])
|
||||
assert.Equal(t, "operations/provider-task", stored["name"])
|
||||
response, ok := stored["response"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.NotContains(t, response, "bytesBase64Encoded")
|
||||
assert.Equal(t, strings.Repeat("a", 256)+"...", response["video"])
|
||||
videos, ok := response["videos"].([]any)
|
||||
require.True(t, ok)
|
||||
require.Len(t, videos, 1)
|
||||
video, ok := videos[0].(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.NotContains(t, video, "bytesBase64Encoded")
|
||||
assert.Equal(t, "video/mp4", video["mimeType"])
|
||||
assert.Equal(t, "https://media.example/video.mp4", video["uri"])
|
||||
}
|
||||
|
||||
func seedTaskPollingChannel(t *testing.T, id int, disableSleep bool) {
|
||||
t.Helper()
|
||||
ch := &model.Channel{
|
||||
@@ -152,7 +176,7 @@ func seedPollingTask(t *testing.T, channelID int, publicID string, upstreamID st
|
||||
Platform: constant.TaskPlatform("kling"),
|
||||
UserId: 1,
|
||||
ChannelId: channelID,
|
||||
Action: constant.TaskActionGenerate,
|
||||
Action: constant.TaskActionImageToVideo,
|
||||
Status: model.TaskStatusInProgress,
|
||||
Progress: "30%",
|
||||
CreatedAt: time.Now().Unix(),
|
||||
@@ -195,6 +219,201 @@ func TestUpdateVideoTasksDefaultSleepWaitsBetweenTasks(t *testing.T) {
|
||||
assert.Equal(t, 1, adaptor.fetchCount())
|
||||
}
|
||||
|
||||
func TestDispatchPlatformUpdateUsesFetchMode(t *testing.T) {
|
||||
truncate(t)
|
||||
const channelID = 109
|
||||
seedTaskPollingChannel(t, channelID, true)
|
||||
task := seedPollingTask(t, channelID, "task_batch", "upstream_batch")
|
||||
taskChannels := map[int][]string{channelID: {task.GetUpstreamTaskID()}}
|
||||
tasks := map[string]*model.Task{task.GetUpstreamTaskID(): task}
|
||||
|
||||
batch := &batchPollingAdaptor{}
|
||||
previousFactory := GetTaskAdaptorFunc
|
||||
GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor { return batch }
|
||||
DispatchPlatformUpdate(context.Background(), "batch-plugin", taskChannels, tasks)
|
||||
assert.Equal(t, 1, batch.batchCalls)
|
||||
assert.Equal(t, 0, batch.fetchCount())
|
||||
var persisted model.Task
|
||||
require.NoError(t, model.DB.First(&persisted, task.ID).Error)
|
||||
assert.Equal(t, "https://example.com/result", persisted.GetResultURL())
|
||||
|
||||
perTask := &taskPollingFetchAdaptor{}
|
||||
GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor { return perTask }
|
||||
DispatchPlatformUpdate(context.Background(), "per-task-plugin", taskChannels, tasks)
|
||||
assert.Equal(t, 1, perTask.fetchCount())
|
||||
|
||||
GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor { return nil }
|
||||
assert.NotPanics(t, func() { DispatchPlatformUpdate(context.Background(), "missing-plugin", taskChannels, tasks) })
|
||||
GetTaskAdaptorFunc = previousFactory
|
||||
}
|
||||
|
||||
func TestUpdateBatchTasksSettlesTieredUsageForTerminalStates(t *testing.T) {
|
||||
testCases := []struct {
|
||||
name string
|
||||
status model.TaskStatus
|
||||
units float64
|
||||
actualQuota int
|
||||
}{
|
||||
{name: "success with usage", status: model.TaskStatusSuccess, units: 3, actualQuota: 3_000},
|
||||
{name: "failure with usage", status: model.TaskStatusFailure, units: 3, actualQuota: 0},
|
||||
{name: "success with zero usage", status: model.TaskStatusSuccess, units: 0, actualQuota: 0},
|
||||
{name: "failure with zero usage", status: model.TaskStatusFailure, units: 0, actualQuota: 0},
|
||||
}
|
||||
for _, testCase := range testCases {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
truncate(t)
|
||||
|
||||
const userID, tokenID, channelID = 41, 41, 141
|
||||
const initialQuota, preConsumedQuota = 10_000, 5_000
|
||||
const tokenRemain = 8_000
|
||||
seedUser(t, userID, initialQuota)
|
||||
seedToken(t, tokenID, userID, "sk-batch-tiered", tokenRemain)
|
||||
seedTaskPollingChannel(t, channelID, true)
|
||||
|
||||
expression := `tier("actual", u("units"))`
|
||||
task := makeTask(userID, channelID, preConsumedQuota, tokenID, BillingSourceWallet, 0)
|
||||
task.TaskID = "task_batch_tiered_" + string(testCase.status)
|
||||
task.Platform = "batch-plugin"
|
||||
task.PrivateData.UpstreamTaskID = "upstream_batch_tiered_" + string(testCase.status)
|
||||
task.SetData(map[string]any{"provider_payload": "must-be-preserved"})
|
||||
task.PrivateData.BillingContext.TieredSnapshot = &billingexpr.BillingSnapshot{
|
||||
ExprString: expression,
|
||||
ExprHash: billingexpr.ExprHashString(expression),
|
||||
GroupRatio: 1,
|
||||
QuotaPerUnit: 1_000,
|
||||
ExprVersion: 1,
|
||||
TaskUsageBilling: true,
|
||||
}
|
||||
require.NoError(t, model.DB.Create(task).Error)
|
||||
|
||||
upstreamID := task.GetUpstreamTaskID()
|
||||
reason := ""
|
||||
if testCase.status == model.TaskStatusFailure {
|
||||
reason = "upstream failed"
|
||||
}
|
||||
result := &BatchTaskResult{TaskInfo: relaycommon.TaskInfo{
|
||||
TaskID: upstreamID,
|
||||
Status: string(testCase.status),
|
||||
Reason: reason,
|
||||
UsageFacts: map[string]any{"units": testCase.units},
|
||||
}}
|
||||
adaptor := &batchPollingAdaptor{results: map[string]*BatchTaskResult{upstreamID: result}}
|
||||
taskIDs := []string{upstreamID}
|
||||
taskMap := map[string]*model.Task{upstreamID: task}
|
||||
|
||||
require.NoError(t, UpdateBatchTasks(context.Background(), adaptor, map[int][]string{channelID: taskIDs}, taskMap))
|
||||
|
||||
var persisted model.Task
|
||||
require.NoError(t, model.DB.First(&persisted, task.ID).Error)
|
||||
assert.Equal(t, testCase.status, persisted.Status)
|
||||
assert.Equal(t, testCase.actualQuota, persisted.Quota)
|
||||
var persistedData map[string]any
|
||||
require.NoError(t, common.Unmarshal(persisted.Data, &persistedData))
|
||||
assert.Equal(t, "must-be-preserved", persistedData["provider_payload"])
|
||||
assert.Equal(t, initialQuota+(preConsumedQuota-testCase.actualQuota), getUserQuota(t, userID))
|
||||
assert.Equal(t, tokenRemain+(preConsumedQuota-testCase.actualQuota), getTokenRemainQuota(t, tokenID))
|
||||
assert.Equal(t, int64(1), countLogs(t))
|
||||
if testCase.status == model.TaskStatusFailure {
|
||||
log := getLastLog(t)
|
||||
require.NotNil(t, log)
|
||||
assert.Equal(t, model.LogTypeRefund, log.Type)
|
||||
}
|
||||
|
||||
// A duplicate terminal response must not settle the same task twice.
|
||||
require.NoError(t, UpdateBatchTasks(context.Background(), adaptor, map[int][]string{channelID: taskIDs}, taskMap))
|
||||
assert.Equal(t, initialQuota+(preConsumedQuota-testCase.actualQuota), getUserQuota(t, userID))
|
||||
assert.Equal(t, int64(1), countLogs(t))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpdateBatchTasksRefundsFailedTieredTask(t *testing.T) {
|
||||
truncate(t)
|
||||
|
||||
const userID, tokenID, channelID = 43, 43, 143
|
||||
const initialQuota, preConsumedQuota, tokenRemain = 10_000, 5_000, 8_000
|
||||
seedUser(t, userID, initialQuota)
|
||||
seedToken(t, tokenID, userID, "sk-batch-tiered-refund", tokenRemain)
|
||||
seedTaskPollingChannel(t, channelID, true)
|
||||
|
||||
expression := `tier("actual", u("units"))`
|
||||
task := makeTask(userID, channelID, preConsumedQuota, tokenID, BillingSourceWallet, 0)
|
||||
task.TaskID = "task_batch_tiered_refund"
|
||||
task.Platform = "batch-plugin"
|
||||
task.PrivateData.UpstreamTaskID = "upstream_batch_tiered_refund"
|
||||
task.PrivateData.BillingContext.TieredSnapshot = &billingexpr.BillingSnapshot{
|
||||
ExprString: expression,
|
||||
ExprHash: billingexpr.ExprHashString(expression),
|
||||
GroupRatio: 1,
|
||||
QuotaPerUnit: 1_000,
|
||||
ExprVersion: 1,
|
||||
TaskUsageBilling: true,
|
||||
UsageFacts: map[string]any{"units": float64(5)},
|
||||
EstimatedTier: "actual",
|
||||
}
|
||||
require.NoError(t, model.DB.Create(task).Error)
|
||||
|
||||
upstreamID := task.GetUpstreamTaskID()
|
||||
adaptor := &batchPollingAdaptor{results: map[string]*BatchTaskResult{
|
||||
upstreamID: {TaskInfo: relaycommon.TaskInfo{
|
||||
TaskID: upstreamID,
|
||||
Status: model.TaskStatusFailure,
|
||||
Reason: "upstream failed",
|
||||
UsageFacts: map[string]any{"units": float64(5)},
|
||||
}},
|
||||
}}
|
||||
require.NoError(t, UpdateBatchTasks(context.Background(), adaptor, map[int][]string{channelID: {upstreamID}}, map[string]*model.Task{upstreamID: task}))
|
||||
|
||||
var persisted model.Task
|
||||
require.NoError(t, model.DB.First(&persisted, task.ID).Error)
|
||||
assert.EqualValues(t, model.TaskStatusFailure, persisted.Status)
|
||||
assert.Zero(t, persisted.Quota)
|
||||
assert.Equal(t, initialQuota+preConsumedQuota, getUserQuota(t, userID))
|
||||
assert.Equal(t, tokenRemain+preConsumedQuota, getTokenRemainQuota(t, tokenID))
|
||||
|
||||
log := getLastLog(t)
|
||||
require.NotNil(t, log)
|
||||
assert.Equal(t, model.LogTypeRefund, log.Type)
|
||||
assert.Equal(t, preConsumedQuota, log.Quota)
|
||||
var other map[string]any
|
||||
require.NoError(t, common.UnmarshalJsonStr(log.Other, &other))
|
||||
assert.Equal(t, "tiered_expr", other["billing_mode"])
|
||||
assert.Equal(t, "actual", other["matched_tier"])
|
||||
facts, ok := other["usage_facts"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, map[string]any{"units": float64(5)}, facts)
|
||||
}
|
||||
|
||||
func TestUpdateBatchTasksRefundsFailedTaskWithoutUsageSettlement(t *testing.T) {
|
||||
truncate(t)
|
||||
|
||||
const userID, tokenID, channelID = 42, 42, 142
|
||||
const initialQuota, preConsumedQuota, tokenRemain = 10_000, 4_000, 7_000
|
||||
seedUser(t, userID, initialQuota)
|
||||
seedToken(t, tokenID, userID, "sk-batch-refund", tokenRemain)
|
||||
seedTaskPollingChannel(t, channelID, true)
|
||||
|
||||
task := makeTask(userID, channelID, preConsumedQuota, tokenID, BillingSourceWallet, 0)
|
||||
task.TaskID = "task_batch_refund"
|
||||
task.Platform = "batch-plugin"
|
||||
task.Properties.OriginModelName = "missing-batch-token-price"
|
||||
task.PrivateData.UpstreamTaskID = "upstream_batch_refund"
|
||||
task.PrivateData.BillingContext.OriginModelName = "missing-batch-token-price"
|
||||
require.NoError(t, model.DB.Create(task).Error)
|
||||
|
||||
upstreamID := task.GetUpstreamTaskID()
|
||||
adaptor := &batchPollingAdaptor{results: map[string]*BatchTaskResult{
|
||||
upstreamID: {TaskInfo: relaycommon.TaskInfo{TaskID: upstreamID, Status: model.TaskStatusFailure, Reason: "upstream failed", TotalTokens: 123}},
|
||||
}}
|
||||
require.NoError(t, UpdateBatchTasks(context.Background(), adaptor, map[int][]string{channelID: {upstreamID}}, map[string]*model.Task{upstreamID: task}))
|
||||
|
||||
assert.Equal(t, initialQuota+preConsumedQuota, getUserQuota(t, userID))
|
||||
assert.Equal(t, tokenRemain+preConsumedQuota, getTokenRemainQuota(t, tokenID))
|
||||
log := getLastLog(t)
|
||||
require.NotNil(t, log)
|
||||
assert.Equal(t, model.LogTypeRefund, log.Type)
|
||||
}
|
||||
|
||||
func TestUpdateVideoTasksCanSkipPollingSleepPerChannel(t *testing.T) {
|
||||
truncate(t)
|
||||
|
||||
@@ -405,15 +624,21 @@ func TestUpdateSunoTasksStalePollsRefundExactlyOnce(t *testing.T) {
|
||||
require.NoError(t, model.DB.First(&firstPollTask, task.ID).Error)
|
||||
require.NoError(t, model.DB.First(&staleSecondPollTask, task.ID).Error)
|
||||
|
||||
adaptor := &sunoFailurePollingAdaptor{failReason: "upstream failed"}
|
||||
adaptor := &batchPollingAdaptor{results: map[string]*BatchTaskResult{
|
||||
upstreamTaskID: {TaskInfo: relaycommon.TaskInfo{
|
||||
TaskID: upstreamTaskID,
|
||||
Status: model.TaskStatusFailure,
|
||||
Reason: "upstream failed",
|
||||
}},
|
||||
}}
|
||||
previousFactory := GetTaskAdaptorFunc
|
||||
GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor { return adaptor }
|
||||
t.Cleanup(func() { GetTaskAdaptorFunc = previousFactory })
|
||||
|
||||
require.NoError(t, updateSunoTasks(context.Background(), channelID, []string{upstreamTaskID}, map[string]*model.Task{
|
||||
require.NoError(t, updateBatchTasks(context.Background(), adaptor, channelID, []string{upstreamTaskID}, map[string]*model.Task{
|
||||
upstreamTaskID: &firstPollTask,
|
||||
}))
|
||||
require.NoError(t, updateSunoTasks(context.Background(), channelID, []string{upstreamTaskID}, map[string]*model.Task{
|
||||
require.NoError(t, updateBatchTasks(context.Background(), adaptor, channelID, []string{upstreamTaskID}, map[string]*model.Task{
|
||||
upstreamTaskID: &staleSecondPollTask,
|
||||
}))
|
||||
|
||||
|
||||
Reference in New Issue
Block a user