feat(task): replace built-in task adaptors with a sandboxed JS plugin system (#7076)

This commit is contained in:
Calcium-Ion authored and GitHub committed 2026-08-29 18:51:57 +08:00
1 parent 7037ac15bd
commit eb48396d5f
336 files changed
+52333 -6369

No files matched your search

+11 -3
View File
@@ -1,7 +1,15 @@
# CLAUDE.md — Project Conventions for new-api
@AGENTS.md
## MANDATORY: Read AGENTS.md with the Read tool
## Claude Code
Do not treat `@AGENTS.md` as loaded. Claude Code does not reliably inline that import.
- Follow the shared project instructions imported from `AGENTS.md`.
Before any planning, coding, reviewing, or answering a project question, you MUST call the Read tool on the repo-root file `AGENTS.md` and wait for the full contents. This is the first action of every session and every new task.
Rules:
- Do not start from memory, summaries, or this file alone.
- Do not skip the Read because a previous turn mentioned AGENTS.md.
- Do not replace the Read with a grep, glob, or partial skim.
- After reading, follow every rule in `AGENTS.md` for the rest of the work.
- If the task touches `web/`, also Read `web/AGENTS.md` before editing frontend files.
+3
View File
@@ -34,11 +34,13 @@ Transitive dependencies should be audited before a final external release.
| backend | production | Go | `github.com/google/uuid` | `v1.6.0` | BSD-3-Clause |
| backend | production | Go | `github.com/gorilla/websocket` | `v1.5.0` | BSD-2-Clause |
| backend | production | Go | `github.com/grafana/pyroscope-go` | `v1.2.7` | Apache-2.0 |
| backend | production | Go | `github.com/grafana/sobek` | `v0.0.0-20260708062710-267a0e055bb4` | MIT |
| backend | production | Go | `github.com/jfreymuth/oggvorbis` | `v1.0.5` | MIT |
| backend | production | Go | `github.com/jinzhu/copier` | `v0.4.0` | MIT |
| backend | production | Go | `github.com/joho/godotenv` | `v1.5.1` | MIT |
| backend | production | Go | `github.com/mewkiz/flac` | `v1.0.13` | Unlicense |
| backend | production | Go | `github.com/nicksnyder/go-i18n/v2` | `v2.6.1` | MIT |
| backend | test | Go | `github.com/openai/openai-go` | `v1.12.0` | Apache-2.0 |
| backend | production | Go | `github.com/pkg/errors` | `v0.9.1` | BSD-2-Clause |
| backend | production | Go | `github.com/pquerna/otp` | `v1.5.0` | Apache-2.0 |
| backend | production | Go | `github.com/samber/hot` | `v0.11.0` | MIT |
@@ -66,6 +68,7 @@ Transitive dependencies should be audited before a final external release.
| backend | production | Go | `gorm.io/gorm` | `v1.25.2` | MIT |
| backend | production | Go | `github.com/expr-lang/expr` | `v1.17.8` | MIT |
| web | production | npm | `@base-ui/react` | `1.6.0` | MIT |
| web | production | npm | `@codemirror/lang-javascript` | `6.2.5` | MIT |
| web | production | npm | `@codemirror/lang-markdown` | `6.5.1` | MIT |
| web | production | npm | `@codemirror/language` | `6.12.4` | MIT |
| web | production | npm | `@codemirror/state` | `6.7.1` | MIT |
+5
View File
@@ -83,6 +83,11 @@ func ChannelType2APIType(channelType int) (int, bool) {
apiType = constant.APITypeNewAPI
}
if apiType == -1 {
// Task plugin channels are served by the task relay and must never
// fall back to the OpenAI adaptor.
if channelType == constant.ChannelTypeTaskPlugin {
return -1, false
}
return constant.APITypeOpenAI, false
}
return apiType, true
+14
View File
@@ -0,0 +1,14 @@
package common
import (
"testing"
"github.com/QuantumNous/new-api/constant"
"github.com/stretchr/testify/assert"
)
func TestTaskPluginChannelHasNoOrdinaryAPIType(t *testing.T) {
apiType, ok := ChannelType2APIType(constant.ChannelTypeTaskPlugin)
assert.Equal(t, -1, apiType)
assert.False(t, ok)
}
+8
View File
@@ -187,6 +187,8 @@ func initConstantEnv() {
constant.GetMediaToken = GetEnvOrDefaultBool("GET_MEDIA_TOKEN", true)
constant.GetMediaTokenNotStream = GetEnvOrDefaultBool("GET_MEDIA_TOKEN_NOT_STREAM", false)
constant.UpdateTask = GetEnvOrDefaultBool("UPDATE_TASK", true)
constant.TaskPluginEnabled = GetEnvOrDefaultBool("TASK_PLUGIN_ENABLED", true)
constant.TaskPluginOverrideEnabled = GetEnvOrDefaultBool("TASK_PLUGIN_OVERRIDE_ENABLED", true)
constant.AzureDefaultAPIVersion = GetEnvOrDefaultString("AZURE_DEFAULT_API_VERSION", "2025-04-01-preview")
constant.NotifyLimitCount = GetEnvOrDefault("NOTIFY_LIMIT_COUNT", 2)
constant.NotificationLimitDurationMinute = GetEnvOrDefault("NOTIFICATION_LIMIT_DURATION_MINUTE", 10)
@@ -198,6 +200,12 @@ func initConstantEnv() {
constant.TaskQueryLimit = GetEnvOrDefault("TASK_QUERY_LIMIT", 1000)
// 异步任务超时时间(分钟),超过此时间未完成的任务将被标记为失败并退款。0 表示禁用。
constant.TaskTimeoutMinutes = GetEnvOrDefault("TASK_TIMEOUT_MINUTES", 1440)
// 声明式任务协议桥只观察数据库;这些值控制一次客户端观察连接,
// 不改变后台轮询或结算生命周期。
constant.TaskPluginProtocolTimeoutSeconds = GetEnvOrDefault("TASK_PLUGIN_PROTOCOL_TIMEOUT_SECONDS", 600)
constant.TaskPluginProtocolTickMilliseconds = GetEnvOrDefault("TASK_PLUGIN_PROTOCOL_TICK_MILLISECONDS", 2000)
constant.TaskPluginProtocolTickJitterMilliseconds = GetEnvOrDefault("TASK_PLUGIN_PROTOCOL_TICK_JITTER_MILLISECONDS", 500)
constant.TaskPluginProtocolHeartbeatSeconds = GetEnvOrDefault("TASK_PLUGIN_PROTOCOL_HEARTBEAT_SECONDS", 15)
soraPatchStr := GetEnvOrDefaultString("TASK_PRICE_PATCH", "")
if soraPatchStr != "" {
+54
View File
@@ -0,0 +1,54 @@
package common
import (
"errors"
"fmt"
"strings"
"github.com/gin-gonic/gin"
)
var defaultTrustedProxyCIDRs = []string{
"127.0.0.0/8",
"::1",
"10.0.0.0/8",
"172.16.0.0/12",
"192.168.0.0/16",
"fc00::/7",
}
// ResolveTrustedProxies parses TRUSTED_PROXIES without applying it to an
// engine. The returned slice can be reused by the outer and plugin engines.
func ResolveTrustedProxies(raw string) (trustedProxies []string, usedDefaults bool, err error) {
raw = strings.TrimSpace(raw)
if raw == "" {
return append([]string(nil), defaultTrustedProxyCIDRs...), true, nil
}
if strings.EqualFold(raw, "none") {
return nil, false, nil
}
parts := strings.Split(raw, ",")
trustedProxies = make([]string, 0, len(parts))
for _, part := range parts {
trustedProxy := strings.TrimSpace(part)
if trustedProxy == "" {
continue
}
if strings.EqualFold(trustedProxy, "none") {
return nil, false, errors.New("TRUSTED_PROXIES=none must be used alone")
}
trustedProxies = append(trustedProxies, trustedProxy)
}
if len(trustedProxies) == 0 {
return nil, false, errors.New("TRUSTED_PROXIES does not contain an IP address or CIDR")
}
return trustedProxies, false, nil
}
func ConfigureTrustedProxies(engine *gin.Engine, trustedProxies []string) error {
if err := engine.SetTrustedProxies(trustedProxies); err != nil {
return fmt.Errorf("invalid TRUSTED_PROXIES: %w", err)
}
return nil
}
+10
View File
@@ -58,6 +58,7 @@ const (
ChannelTypeAdvancedCustom = 58
ChannelTypeSub2API = 59
ChannelTypeNewAPI = 60
ChannelTypeTaskPlugin = 61
ChannelTypeDummy // this one is only for count, do not add any channel after this
)
@@ -124,6 +125,14 @@ var ChannelBaseURLs = []string{
"", //58
"", //59
"", //60
"", //61
}
func GetChannelBaseURL(channelType int) string {
if channelType < 0 || channelType >= len(ChannelBaseURLs) {
return ""
}
return ChannelBaseURLs[channelType]
}
var ChannelTypeNames = map[int]string{
@@ -184,6 +193,7 @@ var ChannelTypeNames = map[int]string{
ChannelTypeAdvancedCustom: "Advanced Custom",
ChannelTypeSub2API: "Sub2API",
ChannelTypeNewAPI: "New API",
ChannelTypeTaskPlugin: "Task Plugin",
}
func GetChannelTypeName(channelType int) string {
+12
View File
@@ -0,0 +1,12 @@
package constant
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestGetChannelBaseURLIsBoundsSafe(t *testing.T) {
assert.Empty(t, GetChannelBaseURL(ChannelTypeTaskPlugin))
assert.Empty(t, GetChannelBaseURL(9999))
}
+2 -1
View File
@@ -15,7 +15,8 @@ const (
ContextKeyTokenKey ContextKey = "token_key"
ContextKeyTokenId ContextKey = "token_id"
ContextKeyTokenGroup ContextKey = "token_group"
ContextKeyTokenSpecificChannelId ContextKey = "specific_channel_id"
ContextKeyOriginTasks ContextKey = "origin_tasks"
ContextKeyChannelConstraints ContextKey = "channel_constraints"
ContextKeyTokenModelLimitEnabled ContextKey = "token_model_limit_enabled"
ContextKeyTokenModelLimit ContextKey = "token_model_limit"
ContextKeyTokenCrossGroupRetry ContextKey = "token_cross_group_retry"
+4
View File
@@ -18,6 +18,10 @@ var GenerateDefaultToken bool
var ErrorLogEnabled bool
var TaskQueryLimit int
var TaskTimeoutMinutes int
var TaskPluginProtocolTimeoutSeconds int
var TaskPluginProtocolTickMilliseconds int
var TaskPluginProtocolTickJitterMilliseconds int
var TaskPluginProtocolHeartbeatSeconds int
// temporary variable for sora patch, will be removed in future
var TaskPricePatches []string
+29 -11
View File
@@ -8,17 +8,35 @@ const (
)
const (
SunoActionMusic = "MUSIC"
SunoActionLyrics = "LYRICS"
TaskActionGenerate = "generate"
TaskActionTextGenerate = "textGenerate"
TaskActionFirstTailGenerate = "firstTailGenerate"
TaskActionReferenceGenerate = "referenceGenerate"
TaskActionRemix = "remixGenerate"
TaskActionImageToVideo = "image_to_video"
TaskActionTextToVideo = "text_to_video"
TaskActionFirstTailToVideo = "first_tail_to_video"
TaskActionReferenceToVideo = "reference_to_video"
TaskActionRemix = "remix"
)
var SunoModel2Action = map[string]string{
"suno_music": SunoActionMusic,
"suno_lyrics": SunoActionLyrics,
var legacyTaskActionAliases = map[string]string{
"generate": TaskActionImageToVideo,
"textGenerate": TaskActionTextToVideo,
"firstTailGenerate": TaskActionFirstTailToVideo,
"referenceGenerate": TaskActionReferenceToVideo,
"remixGenerate": TaskActionRemix,
}
// TaskPluginEnabled is the master switch for the whole task-plugin system.
// When disabled, factory and override plugins both stop serving.
var TaskPluginEnabled = true
// TaskPluginOverrideEnabled controls whether the database override layer is
// active. When disabled, uploaded plugins are ignored and factory plugins are
// used instead; the factory layer is unaffected.
var TaskPluginOverrideEnabled = true
// NormalizeTaskAction maps persisted legacy action names to the canonical task
// action vocabulary. Unknown platform-specific actions pass through unchanged.
func NormalizeTaskAction(action string) string {
if canonical, ok := legacyTaskActionAliases[action]; ok {
return canonical
}
return action
}
+27
View File
@@ -0,0 +1,27 @@
package constant
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestNormalizeTaskAction(t *testing.T) {
tests := map[string]string{
"generate": TaskActionImageToVideo,
"textGenerate": TaskActionTextToVideo,
"firstTailGenerate": TaskActionFirstTailToVideo,
"referenceGenerate": TaskActionReferenceToVideo,
"remixGenerate": TaskActionRemix,
TaskActionTextToVideo: TaskActionTextToVideo,
"MUSIC": "MUSIC",
"custom_action": "custom_action",
"": "",
}
for input, expected := range tests {
t.Run(input, func(t *testing.T) {
assert.Equal(t, expected, NormalizeTaskAction(input))
})
}
}
+100
View File
@@ -0,0 +1,100 @@
package controller
import (
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestUpdateOptionRejectsInvalidTaskBillingExpressions(t *testing.T) {
const pluginKey = "billing-save-probe"
const modelName = "billing-save-model"
source := `
export const meta = {
apiVersion: 1, key: "billing-save-probe", name: "Billing Save Probe", version: "1.0.0", author: {name: "Test"},
models: ["billing-save-model"], fetchMode: "per_task",
usageSchema: {seconds: {type: "number", unit: "second"}}
};
export function buildSubmitRequest() { return {}; }
export function parseSubmitResponse() { return {}; }
export function buildQueryRequest() { return {}; }
export function parseTaskResult() { return {}; }
`
_, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{})
require.NoError(t, err)
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(pluginKey) })
tests := []struct {
name string
expression string
errorText string
}{
{
name: "invalid syntax",
expression: `tier("base",`,
errorText: "expr compile error",
},
{
name: "undeclared usage key",
expression: `tier("base", u("clips") * 0.1)`,
errorText: `usage key \"clips\" is not declared`,
},
}
for _, testCase := range tests {
t.Run(testCase.name, func(t *testing.T) {
expressions, marshalErr := common.Marshal(map[string]string{modelName: testCase.expression})
require.NoError(t, marshalErr)
body, marshalErr := common.Marshal(OptionUpdateRequest{
Key: "billing_setting.billing_expr",
Value: string(expressions),
})
require.NoError(t, marshalErr)
recorder := httptest.NewRecorder()
context, _ := gin.CreateTestContext(recorder)
context.Request = httptest.NewRequest(http.MethodPut, "/api/option/", strings.NewReader(string(body)))
UpdateOption(context)
assert.Equal(t, http.StatusOK, recorder.Code)
assert.Contains(t, recorder.Body.String(), `"success":false`)
assert.Contains(t, recorder.Body.String(), modelName)
assert.Contains(t, recorder.Body.String(), testCase.errorText)
})
}
}
func TestUpdateOptionRejectsUsageExpressionWithoutTaskPlugin(t *testing.T) {
const modelName = "billing-save-model-without-plugin"
expressions, err := common.Marshal(map[string]string{
modelName: `u("mode") == "std" ? 1 : 2`,
})
require.NoError(t, err)
body, err := common.Marshal(OptionUpdateRequest{
Key: "billing_setting.billing_expr",
Value: string(expressions),
})
require.NoError(t, err)
recorder := httptest.NewRecorder()
context, _ := gin.CreateTestContext(recorder)
context.Request = httptest.NewRequest(
http.MethodPut,
"/api/option/",
strings.NewReader(string(body)),
)
UpdateOption(context)
assert.Equal(t, http.StatusOK, recorder.Code)
assert.Contains(t, recorder.Body.String(), `"success":false`)
assert.Contains(t, recorder.Body.String(), modelName)
assert.Contains(t, recorder.Body.String(), "mode")
assert.Contains(t, recorder.Body.String(), "no task plugin usage schema")
}
+5 -1
View File
@@ -463,7 +463,7 @@ func updateChannelBalance(channel *model.Channel) (channelBalanceResult, error)
}
func updateStandardChannelBalance(channel *model.Channel) (float64, error) {
baseURL := constant.ChannelBaseURLs[channel.Type]
baseURL := constant.GetChannelBaseURL(channel.Type)
if channel.GetBaseURL() == "" {
channel.BaseURL = &baseURL
}
@@ -538,6 +538,10 @@ func UpdateChannelBalance(c *gin.Context) {
common.ApiError(c, err)
return
}
if channel.Type == constant.ChannelTypeTaskPlugin {
c.JSON(http.StatusOK, gin.H{"success": false, "message": "Task Plugin channels do not support balance queries"})
return
}
if channel.ChannelInfo.IsMultiKey {
c.JSON(http.StatusOK, gin.H{
"success": false,
+1
View File
@@ -82,6 +82,7 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te
constant.ChannelTypeJimeng,
constant.ChannelTypeDoubaoVideo,
constant.ChannelTypeVidu,
constant.ChannelTypeTaskPlugin,
}
if lo.Contains(unsupportedTestChannelTypes, channel.Type) {
channelTypeName := constant.GetChannelTypeName(channel.Type)
+44 -5
View File
@@ -13,6 +13,7 @@ import (
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/jsplugin"
relaychannel "github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/relay/channel/ollama"
relaycommon "github.com/QuantumNous/new-api/relay/common"
@@ -480,6 +481,21 @@ func validateChannel(channel *model.Channel, isAdd bool) error {
if err := channel.ValidateSettings(); err != nil {
return fmt.Errorf("渠道额外设置[channel setting] 格式错误:%s", err.Error())
}
if channel.Type == constant.ChannelTypeTaskPlugin {
pluginKey := strings.TrimSpace(channel.GetSetting().TaskPluginKey)
if pluginKey == "" {
return fmt.Errorf("task plugin key is required")
}
if len(pluginKey) > 30 {
return fmt.Errorf("task plugin key must not exceed 30 characters")
}
if _, ok := jsplugin.DefaultRegistry.Get(pluginKey); !ok {
return fmt.Errorf("task plugin %q is not registered", pluginKey)
}
if channel.BaseURL == nil || strings.TrimSpace(*channel.BaseURL) == "" {
return fmt.Errorf("base URL is required for task plugin channels")
}
}
if channel.Type == constant.ChannelTypeNewAPI && strings.TrimSpace(channel.GetBaseURL()) == "" {
return fmt.Errorf("New API channel base URL cannot be empty")
@@ -617,6 +633,15 @@ func AddChannel(c *gin.Context) {
return
}
if addChannelRequest.Channel != nil && addChannelRequest.Channel.Type == constant.ChannelTypeTaskPlugin &&
!authz.Can(c.GetInt("id"), c.GetInt("role"), authz.TaskPluginBind) {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "task plugin channels require the task_plugin.bind permission",
})
return
}
// 使用统一的校验函数
if err := validateChannel(addChannelRequest.Channel, true); err != nil {
c.JSON(http.StatusOK, gin.H{
@@ -964,6 +989,15 @@ func UpdateChannel(c *gin.Context) {
}
clearChannelReadOnlyFields(&channel, requestData)
if channel.Type == constant.ChannelTypeTaskPlugin &&
!authz.Can(c.GetInt("id"), c.GetInt("role"), authz.TaskPluginBind) {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "task plugin channels require the task_plugin.bind permission",
})
return
}
// 使用统一的校验函数
if err := validateChannel(&channel.Channel, false); err != nil {
c.JSON(http.StatusOK, gin.H{
@@ -1299,7 +1333,7 @@ func FetchModels(c *gin.Context) {
baseURL = strings.TrimSpace(*req.BaseURL)
}
if baseURL == "" {
baseURL = constant.ChannelBaseURLs[req.Type]
baseURL = constant.GetChannelBaseURL(req.Type)
}
key := strings.TrimSpace(req.Key)
@@ -1424,6 +1458,11 @@ func CopyChannel(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"success": false, "message": "获取渠道信息失败,请稍后重试"})
return
}
if origin.Type == constant.ChannelTypeTaskPlugin &&
!authz.Can(c.GetInt("id"), c.GetInt("role"), authz.TaskPluginBind) {
c.JSON(http.StatusOK, gin.H{"success": false, "message": "task plugin channels require the task_plugin.bind permission"})
return
}
// clone channel
clone := *origin // shallow copy is sufficient as we will overwrite primitives
@@ -2010,7 +2049,7 @@ func OllamaPullModel(c *gin.Context) {
return
}
baseURL := constant.ChannelBaseURLs[channel.Type]
baseURL := constant.GetChannelBaseURL(channel.Type)
if channel.GetBaseURL() != "" {
baseURL = channel.GetBaseURL()
}
@@ -2073,7 +2112,7 @@ func OllamaPullModelStream(c *gin.Context) {
return
}
baseURL := constant.ChannelBaseURLs[channel.Type]
baseURL := constant.GetChannelBaseURL(channel.Type)
if channel.GetBaseURL() != "" {
baseURL = channel.GetBaseURL()
}
@@ -2155,7 +2194,7 @@ func OllamaDeleteModel(c *gin.Context) {
return
}
baseURL := constant.ChannelBaseURLs[channel.Type]
baseURL := constant.GetChannelBaseURL(channel.Type)
if channel.GetBaseURL() != "" {
baseURL = channel.GetBaseURL()
}
@@ -2204,7 +2243,7 @@ func OllamaVersion(c *gin.Context) {
return
}
baseURL := constant.ChannelBaseURLs[channel.Type]
baseURL := constant.GetChannelBaseURL(channel.Type)
if channel.GetBaseURL() != "" {
baseURL = channel.GetBaseURL()
}
+95
View File
@@ -0,0 +1,95 @@
package controller
import (
"errors"
"net/http"
"net/http/httptest"
"testing"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestShouldRetryHonorsPinRetryMode(t *testing.T) {
openaiErr := types.NewOpenAIError(errors.New("upstream"), types.ErrorCodeBadResponseStatusCode, http.StatusInternalServerError)
c := newPinRetryContext()
assert.True(t, shouldRetry(c, openaiErr, 1))
origin := newPinRetryContext()
service.GetChannelConstraints(origin).AddPin(dto.ChannelPin{
ChannelId: 2,
Source: dto.PinSourceOriginTask,
Rank: dto.PinRankOriginTask,
RetryMode: dto.PinRetrySameChannel,
})
assert.True(t, shouldRetry(origin, openaiErr, 1), "origin pin retries on the same channel")
token := newPinRetryContext()
service.GetChannelConstraints(token).AddPin(dto.ChannelPin{
ChannelId: 1,
Source: dto.PinSourceToken,
Rank: dto.PinRankToken,
RetryMode: dto.PinRetrySingleAttempt,
})
assert.False(t, shouldRetry(token, openaiErr, 1), "token pin suppresses retry")
}
func TestShouldRetryTaskRelayHonorsPinRetryMode(t *testing.T) {
taskErr := &dto.TaskError{StatusCode: http.StatusInternalServerError}
c := newPinRetryContext()
assert.True(t, shouldRetryTaskRelay(c, 1, taskErr, 1))
origin := newPinRetryContext()
service.GetChannelConstraints(origin).AddPin(dto.ChannelPin{
ChannelId: 2,
Source: dto.PinSourceOriginTask,
Rank: dto.PinRankOriginTask,
RetryMode: dto.PinRetrySameChannel,
})
assert.True(t, shouldRetryTaskRelay(origin, 2, taskErr, 1))
token := newPinRetryContext()
service.GetChannelConstraints(token).AddPin(dto.ChannelPin{
ChannelId: 1,
Source: dto.PinSourceToken,
Rank: dto.PinRankToken,
RetryMode: dto.PinRetrySingleAttempt,
})
assert.False(t, shouldRetryTaskRelay(token, 1, taskErr, 1))
}
func TestSameChannelPinsMergeToStricterRetryMode(t *testing.T) {
c := newPinRetryContext()
constraints := service.GetChannelConstraints(c)
constraints.AddPin(dto.ChannelPin{
ChannelId: 7,
Source: dto.PinSourceOriginTask,
Rank: dto.PinRankOriginTask,
RetryMode: dto.PinRetrySameChannel,
})
constraints.AddPin(dto.ChannelPin{
ChannelId: 7,
Source: dto.PinSourceToken,
Rank: dto.PinRankToken,
RetryMode: dto.PinRetrySingleAttempt,
})
pin, found, overridden := constraints.ResolvedPin()
require.True(t, found)
assert.Equal(t, 7, pin.ChannelId)
assert.Equal(t, dto.PinRetrySingleAttempt, pin.RetryMode)
assert.Empty(t, overridden)
assert.False(t, shouldRetry(c, types.NewOpenAIError(errors.New("upstream"), types.ErrorCodeBadResponseStatusCode, http.StatusInternalServerError), 1))
}
func newPinRetryContext() *gin.Context {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
return c
}
+131
View File
@@ -0,0 +1,131 @@
package controller
import (
"fmt"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/service/authz"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)
func setupTaskPluginBindChannelTest(t *testing.T) {
t.Helper()
wasMaster := common.IsMasterNode
common.IsMasterNode = true
previousRedisEnabled := common.RedisEnabled
common.RedisEnabled = false
originalDB, originalLogDB := model.DB, model.LOG_DB
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
sqlDB, err := database.DB()
require.NoError(t, err)
sqlDB.SetMaxOpenConns(1)
require.NoError(t, database.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.CasbinRule{}, &model.AuthzRole{}, &model.Log{}, &model.User{}))
model.DB = database
model.LOG_DB = database
require.NoError(t, authz.Init(database))
t.Cleanup(func() {
common.IsMasterNode = wasMaster
common.RedisEnabled = previousRedisEnabled
model.DB = originalDB
model.LOG_DB = originalLogDB
})
}
func postAddChannel(t *testing.T, userID, role int, body string) *httptest.ResponseRecorder {
t.Helper()
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
context, _ := gin.CreateTestContext(recorder)
context.Set("id", userID)
context.Set("role", role)
context.Request = httptest.NewRequest(http.MethodPost, "/api/channel", strings.NewReader(body))
context.Request.Header.Set("Content-Type", "application/json")
AddChannel(context)
return recorder
}
func TestAddChannelTaskPluginRequiresBindPermission(t *testing.T) {
setupTaskPluginBindChannelTest(t)
const key = "channel-bind"
source := `
export const meta = {apiVersion: 1, key: "channel-bind", name: "Bind", version: "1.0.0", author: {name: "Test"}, models: ["doc"], fetchMode: "per_task"};
export function buildSubmitRequest() { return {}; }
export function parseSubmitResponse() { return {}; }
export function buildQueryRequest() { return {}; }
export function parseTaskResult() { return {}; }
`
_, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{})
require.NoError(t, err)
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(key) })
taskPluginBody := `{"mode":"single","channel":{"type":61,"name":"plugin-channel","key":"sk","models":"doc","group":"default","base_url":"https://example.com","setting":"{\"task_plugin_key\":\"channel-bind\"}"}}`
openaiBody := `{"mode":"single","channel":{"type":1,"name":"openai-channel","key":"sk","models":"gpt","group":"default"}}`
adminDenied := postAddChannel(t, 2, common.RoleAdminUser, taskPluginBody)
assert.Contains(t, adminDenied.Body.String(), "task plugin channels require the task_plugin.bind permission")
assert.Contains(t, adminDenied.Body.String(), `"success":false`)
rootAllowed := postAddChannel(t, 1, common.RoleRootUser, taskPluginBody)
assert.Contains(t, rootAllowed.Body.String(), `"success":true`)
assert.NotContains(t, rootAllowed.Body.String(), "task_plugin.bind")
adminOtherType := postAddChannel(t, 2, common.RoleAdminUser, openaiBody)
assert.Contains(t, adminOtherType.Body.String(), `"success":true`)
assert.NotContains(t, adminOtherType.Body.String(), "task_plugin.bind")
}
func TestUpdateChannelTaskPluginRequiresBindPermission(t *testing.T) {
setupTaskPluginBindChannelTest(t)
const key = "channel-bind-update"
source := `
export const meta = {apiVersion: 1, key: "channel-bind-update", name: "Bind", version: "1.0.0", author: {name: "Test"}, models: ["doc"], fetchMode: "per_task"};
export function buildSubmitRequest() { return {}; }
export function parseSubmitResponse() { return {}; }
export function buildQueryRequest() { return {}; }
export function parseTaskResult() { return {}; }
`
_, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{})
require.NoError(t, err)
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(key) })
baseURL := "https://example.com"
setting := `{"task_plugin_key":"channel-bind-update"}`
channel := model.Channel{
Type: constant.ChannelTypeTaskPlugin,
Status: common.ChannelStatusEnabled,
Name: "existing-plugin",
Models: "doc",
Group: "default",
Key: "sk",
BaseURL: &baseURL,
Setting: &setting,
}
require.NoError(t, channel.Insert())
payload := fmt.Sprintf(
`{"id":%d,"type":61,"name":"existing-plugin","key":"sk","models":"doc","group":"default","base_url":"https://example.com","setting":"{\"task_plugin_key\":\"channel-bind-update\"}"}`,
channel.Id,
)
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
context, _ := gin.CreateTestContext(recorder)
context.Set("id", 2)
context.Set("role", common.RoleAdminUser)
context.Request = httptest.NewRequest(http.MethodPut, "/api/channel", strings.NewReader(payload))
context.Request.Header.Set("Content-Type", "application/json")
UpdateChannel(context)
assert.Contains(t, recorder.Body.String(), "task plugin channels require the task_plugin.bind permission")
assert.Contains(t, recorder.Body.String(), `"success":false`)
}
@@ -0,0 +1,41 @@
package controller
import (
"strings"
"testing"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/stretchr/testify/require"
)
func TestValidateTaskPluginChannel(t *testing.T) {
source := `
export const meta = {apiVersion: 1, key: "channel-validation", name: "Validation", version: "1.0.0", author: {name: "Test"}, models: ["doc"], fetchMode: "per_task"};
export function buildSubmitRequest() { return {}; }
export function parseSubmitResponse() { return {}; }
export function buildQueryRequest() { return {}; }
export function parseTaskResult() { return {}; }
`
_, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{})
require.NoError(t, err)
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister("channel-validation") })
baseURL := "https://example.com"
channel := &model.Channel{Type: constant.ChannelTypeTaskPlugin, BaseURL: &baseURL}
require.ErrorContains(t, validateChannel(channel, false), "task plugin key is required")
missing := `{"task_plugin_key":"missing"}`
channel.Setting = &missing
require.ErrorContains(t, validateChannel(channel, false), "is not registered")
longKey := `{"task_plugin_key":"` + strings.Repeat("x", 31) + `"}`
channel.Setting = &longKey
require.ErrorContains(t, validateChannel(channel, false), "must not exceed 30")
valid := `{"task_plugin_key":"channel-validation"}`
channel.Setting = &valid
channel.BaseURL = nil
require.ErrorContains(t, validateChannel(channel, false), "base URL is required")
}
+9 -1
View File
@@ -16,6 +16,7 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/relay/channel/advancedcustom"
"github.com/QuantumNous/new-api/relay/channel/gemini"
"github.com/QuantumNous/new-api/relay/channel/ollama"
@@ -361,7 +362,14 @@ func getFetchModelsResponseBody(method string, requestURL string, channel *model
}
func fetchChannelUpstreamModelIDs(channel *model.Channel) ([]string, error) {
baseURL := constant.ChannelBaseURLs[channel.Type]
if channel.Type == constant.ChannelTypeTaskPlugin {
plugin, ok := jsplugin.DefaultRegistry.Get(channel.GetSetting().TaskPluginKey)
if !ok {
return nil, fmt.Errorf("task plugin %q is not registered", channel.GetSetting().TaskPluginKey)
}
return normalizeModelNames(plugin.Meta.Models), nil
}
baseURL := constant.GetChannelBaseURL(channel.Type)
if channel.GetBaseURL() != "" {
baseURL = channel.GetBaseURL()
}
+3
View File
@@ -27,6 +27,9 @@ func GetAllLogs(c *gin.Context) {
common.ApiError(c, err)
return
}
if c.GetInt("role") < common.RoleRootUser {
model.FormatAdminLogs(logs)
}
pageInfo.SetTotal(int(total))
pageInfo.SetItems(logs)
common.ApiSuccess(c, pageInfo)
+19 -1
View File
@@ -9,6 +9,7 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/relay"
"github.com/QuantumNous/new-api/relay/channel/ai360"
"github.com/QuantumNous/new-api/relay/channel/lingyiwanwu"
@@ -97,6 +98,9 @@ func init() {
for i := 1; i <= constant.ChannelTypeDummy; i++ {
apiType, success := common.ChannelType2APIType(i)
if !success || apiType == constant.APITypeAIProxyLibrary {
if plugin, ok := jsplugin.DefaultRegistry.GetByChannelType(i); ok {
channelId2Models[i] = append([]string(nil), plugin.Meta.Models...)
}
continue
}
meta := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{
@@ -105,6 +109,11 @@ func init() {
adaptor := relay.GetAdaptor(apiType)
adaptor.Init(meta)
channelId2Models[i] = adaptor.GetModelList()
if len(channelId2Models[i]) == 0 {
if plugin, ok := jsplugin.DefaultRegistry.GetByChannelType(i); ok {
channelId2Models[i] = append([]string(nil), plugin.Meta.Models...)
}
}
}
openAIModels = lo.UniqBy(openAIModels, func(m dto.OpenAIModels) string {
return m.Id
@@ -314,9 +323,18 @@ func ChannelListModels(c *gin.Context) {
}
func DashboardListModels(c *gin.Context) {
modelsByChannel := make(map[int][]string, len(channelId2Models))
for channelType, models := range channelId2Models {
modelsByChannel[channelType] = append([]string(nil), models...)
}
for channelType := 1; channelType <= constant.ChannelTypeDummy; channelType++ {
if plugin, ok := jsplugin.DefaultRegistry.GetByChannelType(channelType); ok {
modelsByChannel[channelType] = append([]string(nil), plugin.Meta.Models...)
}
}
c.JSON(200, gin.H{
"success": true,
"data": channelId2Models,
"data": modelsByChannel,
})
}
+34
View File
@@ -3,13 +3,17 @@ package controller
import (
"fmt"
"net/http"
"sort"
"strconv"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting"
"github.com/QuantumNous/new-api/setting/billing_setting"
"github.com/QuantumNous/new-api/setting/console_setting"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/QuantumNous/new-api/setting/operation_setting"
@@ -153,6 +157,12 @@ func UpdateOption(c *gin.Context) {
return
}
}
if option.Key == "TaskPublicAddress" && option.Value.(string) != "" {
if err := service.ValidateTaskArtifactBaseURL(option.Value.(string)); err != nil {
common.ApiErrorMsg(c, err.Error())
return
}
}
switch option.Key {
case "GitHubOAuthEnabled":
if option.Value == "true" && common.GitHubClientId == "" {
@@ -326,6 +336,30 @@ func UpdateOption(c *gin.Context) {
})
return
}
case "billing_setting.billing_expr":
expressions := make(map[string]string)
if err = common.UnmarshalJsonStr(option.Value.(string), &expressions); err != nil {
common.ApiErrorMsg(c, "计费表达式配置必须是模型到表达式的 JSON 对象: "+err.Error())
return
}
models := make([]string, 0, len(expressions))
for modelName := range expressions {
models = append(models, modelName)
}
sort.Strings(models)
generation := jsplugin.DefaultRegistry.Generation()
for _, modelName := range models {
expression := expressions[modelName]
if plugin, ok := generation.GetByModel(modelName); ok {
err = billing_setting.SmokeTestTaskExpr(expression, plugin.Meta.UsageSchema)
} else {
err = billing_setting.SmokeTestExpr(expression)
}
if err != nil {
common.ApiErrorMsg(c, fmt.Sprintf("模型 %s 的计费表达式无效: %v", modelName, err))
return
}
}
case "console_setting.api_info":
err = console_setting.ValidateConsoleSettings(option.Value.(string), "ApiInfo")
if err != nil {
+47
View File
@@ -0,0 +1,47 @@
package controller
import (
"net/http"
"net/http/httptest"
"testing"
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
)
func TestRelayTaskPluginEndpointPreservesUnclaimedFallback(t *testing.T) {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
fallbackCalls := 0
RelayTaskPluginEndpoint(c, func(c *gin.Context) {
fallbackCalls++
c.Status(http.StatusNoContent)
c.Writer.WriteHeaderNow()
})
assert.Equal(t, 1, fallbackCalls)
assert.Equal(t, http.StatusNoContent, recorder.Code)
}
func TestRelayTaskPluginEndpointNeverEntersOrdinaryRelayWhenClaimed(t *testing.T) {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Set(jsplugin.ContextKeyPinnedEndpoint, jsplugin.PinnedEndpoint{
Generation: &jsplugin.RoutingGeneration{},
Plugin: &jsplugin.LoadedPlugin{},
Protocol: "openai_responses",
Operation: jsplugin.HostProtocolOperation{Name: "create"},
})
fallbackCalls := 0
RelayTaskPluginEndpoint(c, func(c *gin.Context) {
fallbackCalls++
c.Status(http.StatusNoContent)
c.Writer.WriteHeaderNow()
})
assert.Zero(t, fallbackCalls)
assert.NotEqual(t, http.StatusNoContent, recorder.Code)
}
+244
View File
@@ -0,0 +1,244 @@
package controller
import (
"bytes"
"context"
"io"
"net/http"
"net/http/httptest"
"sync/atomic"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/middleware"
"github.com/QuantumNous/new-api/model"
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/relay"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)
type nativeRouteBilling struct {
events []string
preConsumed int
userID int
settled bool
}
func (b *nativeRouteBilling) Settle(int) error {
b.events = append(b.events, "settle")
b.settled = true
return nil
}
func (b *nativeRouteBilling) Refund(*gin.Context) {
b.events = append(b.events, "refund")
if !b.settled && b.preConsumed > 0 {
_ = model.IncreaseUserQuota(b.userID, b.preConsumed, true)
b.preConsumed = 0
}
}
func (b *nativeRouteBilling) NeedsRefund() bool {
return !b.settled && b.preConsumed > 0
}
func (b *nativeRouteBilling) GetPreConsumedQuota() int {
return b.preConsumed
}
func (b *nativeRouteBilling) Reserve(quota int) error {
b.events = append(b.events, "reserve")
if err := model.DecreaseUserQuota(b.userID, quota, true); err != nil {
return err
}
b.preConsumed = quota
return nil
}
func TestKlingNativeRouteSubmitPollSettleAndQuery(t *testing.T) {
gin.SetMode(gin.TestMode)
service.InitHttpClient()
previousDB := model.DB
previousLogDB := model.LOG_DB
previousMemoryCache := common.MemoryCacheEnabled
previousBatchUpdate := common.BatchUpdateEnabled
previousLogConsume := common.LogConsumeEnabled
previousRedisEnabled := common.RedisEnabled
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
require.NoError(t, database.AutoMigrate(&model.User{}, &model.Channel{}, &model.Task{}, &model.Log{}))
model.DB = database
model.LOG_DB = database
common.MemoryCacheEnabled = false
common.BatchUpdateEnabled = false
common.LogConsumeEnabled = false
common.RedisEnabled = false
previousModelRatios := ratio_setting.ModelRatio2JSONString()
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"kling-v1":1}`))
t.Cleanup(func() {
model.DB = previousDB
model.LOG_DB = previousLogDB
common.MemoryCacheEnabled = previousMemoryCache
common.BatchUpdateEnabled = previousBatchUpdate
common.LogConsumeEnabled = previousLogConsume
common.RedisEnabled = previousRedisEnabled
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(previousModelRatios))
})
require.NoError(t, database.Create(&model.User{
Id: 7,
Username: "native-route-user",
Group: "default",
Quota: 1_000_000,
}).Error)
var submitCalls atomic.Int32
var queryCalls atomic.Int32
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
switch {
case r.Method == http.MethodPost && r.URL.Path == "/kling/v1/videos/text2video":
submitCalls.Add(1)
body, readErr := io.ReadAll(r.Body)
if !assert.NoError(t, readErr) {
http.Error(w, "read request", http.StatusInternalServerError)
return
}
assert.Contains(t, string(body), `"model_name":"kling-v1"`)
_, _ = io.WriteString(w, `{"code":0,"message":"","data":{"task_id":"kling-private-1","task_status":"submitted"}}`)
case r.Method == http.MethodGet && r.URL.Path == "/kling/v1/videos/text2video/kling-private-1":
queryCalls.Add(1)
_, _ = io.WriteString(w, `{"code":0,"message":"","data":{"task_id":"kling-private-1","task_status":"succeed","task_status_msg":"","task_result":{"videos":[{"id":"video-private","url":"https://cdn.example/video.mp4","duration":"5"}]},"final_unit_deduction":"1"}}`)
default:
http.NotFound(w, r)
}
}))
defer upstream.Close()
channel := model.Channel{
Type: constant.ChannelTypeKling,
Name: "kling-native-e2e",
Key: "sk-test",
BaseURL: &upstream.URL,
Status: common.ChannelStatusEnabled,
Models: "kling-v1",
Group: "default",
}
require.NoError(t, database.Create(&channel).Error)
generation := pluginruntime.DefaultRegistry.Generation()
require.NotNil(t, generation)
submitBinding, found := generation.LookupDeclaredRoute(http.MethodPost, "/kling/v1/videos/text2video")
require.True(t, found)
require.Equal(t, "kling", submitBinding.Plugin.Meta.Key)
submitRecorder := httptest.NewRecorder()
submitContext, _ := gin.CreateTestContext(submitRecorder)
submitContext.Request = httptest.NewRequest(
http.MethodPost,
"/kling/v1/videos/text2video",
bytes.NewBufferString(`{"model_name":"kling-v1","prompt":"a lighthouse"}`),
)
submitContext.Request.Header.Set("Content-Type", "application/json")
submitContext.Set(pluginruntime.ContextKeyPinnedRoute, pluginruntime.PinnedRoute{
Generation: generation,
Plugin: submitBinding.Plugin,
Route: submitBinding.Route,
})
common.SetContextKey(submitContext, constant.ContextKeyUserId, 7)
common.SetContextKey(submitContext, constant.ContextKeyUserGroup, "default")
common.SetContextKey(submitContext, constant.ContextKeyUsingGroup, "default")
common.SetContextKey(submitContext, constant.ContextKeyTokenGroup, "default")
common.SetContextKey(submitContext, constant.ContextKeyUserQuota, 1_000_000)
middleware.PrepareTaskPluginRoute()(submitContext)
require.False(t, submitContext.IsAborted(), submitRecorder.Body.String())
require.Equal(t, "kling-v1", submitContext.GetString("resolved_task_model"))
require.Equal(t, "text_to_video", submitContext.GetString("task_action"))
require.Nil(t, middleware.SetupContextForSelectedChannel(submitContext, &channel, "kling-v1"))
billing := &nativeRouteBilling{userID: 7}
relayInfo := &relaycommon.RelayInfo{
UserId: 7,
UserGroup: "default",
UsingGroup: "default",
UserQuota: 1_000_000,
TokenGroup: "default",
OriginModelName: "kling-v1",
Billing: billing,
TaskRelayInfo: &relaycommon.TaskRelayInfo{
Action: submitContext.GetString("task_action"),
PublicTaskID: "task_kling_public",
LockedChannel: &channel,
},
}
outcome, taskErr := executeTaskSubmissionWith(submitContext, relayInfo, relay.RelayTaskSubmit)
require.Nil(t, taskErr)
require.NotNil(t, outcome)
require.Equal(t, []string{"reserve", "settle"}, billing.events)
require.False(t, submitContext.Writer.Written())
presentTaskSubmission(submitContext, outcome)
require.Equal(t, http.StatusOK, submitRecorder.Code)
assert.Contains(t, submitRecorder.Body.String(), `"task_id":"task_kling_public"`)
assert.NotContains(t, submitRecorder.Body.String(), "kling-private-1")
assert.Equal(t, int32(1), submitCalls.Load())
var persisted model.Task
require.NoError(t, database.Where("task_id = ?", "task_kling_public").First(&persisted).Error)
assert.Equal(t, constant.TaskPlatform("kling"), persisted.Platform)
assert.Equal(t, "kling-private-1", persisted.PrivateData.UpstreamTaskID)
assert.Equal(t, model.TaskStatus(model.TaskStatusNotStart), persisted.Status)
previousAdaptorFactory := service.GetTaskAdaptorFunc
service.GetTaskAdaptorFunc = func(platform constant.TaskPlatform) service.TaskPollingAdaptor {
return relay.GetTaskAdaptor(platform)
}
t.Cleanup(func() { service.GetTaskAdaptorFunc = previousAdaptorFactory })
service.DispatchPlatformUpdate(
context.Background(),
persisted.Platform,
map[int][]string{channel.Id: {"kling-private-1"}},
map[string]*model.Task{"kling-private-1": &persisted},
)
require.NoError(t, database.Where("task_id = ?", "task_kling_public").First(&persisted).Error)
assert.Equal(t, model.TaskStatus(model.TaskStatusSuccess), persisted.Status)
assert.Equal(t, "100%", persisted.Progress)
assert.Equal(t, 1, persisted.Quota)
assert.Equal(t, int32(1), queryCalls.Load())
var settledUser model.User
require.NoError(t, database.First(&settledUser, 7).Error)
assert.Equal(t, 999_999, settledUser.Quota)
queryBinding, found := generation.LookupDeclaredRoute(http.MethodGet, "/kling/v1/videos/text2video/:task_id")
require.True(t, found)
queryRecorder := httptest.NewRecorder()
queryContext, _ := gin.CreateTestContext(queryRecorder)
queryContext.Request = httptest.NewRequest(http.MethodGet, "/kling/v1/videos/text2video/task_kling_public", nil)
queryContext.Params = gin.Params{{Key: "task_id", Value: "task_kling_public"}}
queryContext.Set(pluginruntime.ContextKeyPinnedRoute, pluginruntime.PinnedRoute{
Generation: generation,
Plugin: queryBinding.Plugin,
Route: queryBinding.Route,
})
common.SetContextKey(queryContext, constant.ContextKeyUserId, 7)
middleware.PrepareTaskPluginRoute()(queryContext)
require.True(t, queryContext.IsAborted())
require.Equal(t, http.StatusOK, queryRecorder.Code)
assert.Contains(t, queryRecorder.Body.String(), `"task_id":"task_kling_public"`)
assert.Contains(t, queryRecorder.Body.String(), `"task_status":"succeed"`)
assert.NotContains(t, queryRecorder.Body.String(), "kling-private-1")
assert.NotContains(t, queryRecorder.Body.String(), upstream.URL)
}
File diff suppressed because it is too large. Load diff
+157
View File
@@ -0,0 +1,157 @@
package controller
import (
"errors"
"fmt"
"strings"
"sync"
)
var (
errPluginProtocolObservationLimitExceeded = errors.New("plugin protocol observation limit exceeded")
errInvalidPluginProtocolObservationIdentity = errors.New("invalid plugin protocol observation identity")
)
type pluginProtocolObservationLimits struct {
global int
perPlugin int
perUser int
perToken int
}
var defaultPluginProtocolObservationLimits = pluginProtocolObservationLimits{
global: 128,
perPlugin: 32,
perUser: 4,
perToken: 2,
}
var pluginProtocolObservationAdmissions = newPluginProtocolObservationLimiter(
defaultPluginProtocolObservationLimits,
)
type pluginProtocolObservationLimitError struct {
scope string
limit int
}
func (e *pluginProtocolObservationLimitError) Error() string {
return fmt.Sprintf("%s: %s capacity is %d", errPluginProtocolObservationLimitExceeded, e.scope, e.limit)
}
func (e *pluginProtocolObservationLimitError) Unwrap() error {
return errPluginProtocolObservationLimitExceeded
}
type pluginProtocolObservationLimiter struct {
mu sync.Mutex
limits pluginProtocolObservationLimits
global int
plugin map[string]int
user map[int]int
token map[int]int
}
func newPluginProtocolObservationLimiter(limits pluginProtocolObservationLimits) *pluginProtocolObservationLimiter {
return &pluginProtocolObservationLimiter{
limits: limits,
plugin: make(map[string]int),
user: make(map[int]int),
token: make(map[int]int),
}
}
func (l *pluginProtocolObservationLimiter) acquire(
pluginKey string,
userID int,
tokenID int,
) (func(), error) {
pluginKey = strings.TrimSpace(pluginKey)
switch {
case pluginKey == "":
return nil, fmt.Errorf("%w: plugin key is required", errInvalidPluginProtocolObservationIdentity)
case userID <= 0:
return nil, fmt.Errorf("%w: user id must be positive", errInvalidPluginProtocolObservationIdentity)
case tokenID <= 0:
return nil, fmt.Errorf("%w: token id must be positive", errInvalidPluginProtocolObservationIdentity)
}
l.mu.Lock()
if l.global >= l.limits.global {
l.mu.Unlock()
return nil, &pluginProtocolObservationLimitError{
scope: "global",
limit: l.limits.global,
}
}
l.global++
if l.plugin[pluginKey] >= l.limits.perPlugin {
l.global--
l.mu.Unlock()
return nil, &pluginProtocolObservationLimitError{
scope: "plugin",
limit: l.limits.perPlugin,
}
}
l.plugin[pluginKey]++
if l.user[userID] >= l.limits.perUser {
l.global--
l.plugin[pluginKey]--
if l.plugin[pluginKey] == 0 {
delete(l.plugin, pluginKey)
}
l.mu.Unlock()
return nil, &pluginProtocolObservationLimitError{
scope: "user",
limit: l.limits.perUser,
}
}
l.user[userID]++
if l.token[tokenID] >= l.limits.perToken {
l.global--
l.plugin[pluginKey]--
if l.plugin[pluginKey] == 0 {
delete(l.plugin, pluginKey)
}
l.user[userID]--
if l.user[userID] == 0 {
delete(l.user, userID)
}
l.mu.Unlock()
return nil, &pluginProtocolObservationLimitError{
scope: "token",
limit: l.limits.perToken,
}
}
l.token[tokenID]++
l.mu.Unlock()
var once sync.Once
return func() {
once.Do(func() {
l.mu.Lock()
defer l.mu.Unlock()
l.global--
l.plugin[pluginKey]--
if l.plugin[pluginKey] == 0 {
delete(l.plugin, pluginKey)
}
l.user[userID]--
if l.user[userID] == 0 {
delete(l.user, userID)
}
l.token[tokenID]--
if l.token[tokenID] == 0 {
delete(l.token, tokenID)
}
})
}, nil
}
+248
View File
@@ -0,0 +1,248 @@
package controller
import (
"errors"
"sync"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestPluginProtocolObservationLimiterCaps(t *testing.T) {
t.Run("global", func(t *testing.T) {
limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{
global: 2,
perPlugin: 2,
perUser: 2,
perToken: 2,
})
releaseFirst, err := limiter.acquire("first", 1, 1)
require.NoError(t, err)
defer releaseFirst()
releaseSecond, err := limiter.acquire("second", 2, 2)
require.NoError(t, err)
defer releaseSecond()
release, err := limiter.acquire("third", 3, 3)
assert.Nil(t, release)
assertLimitError(t, err, "global", 2)
})
t.Run("plugin", func(t *testing.T) {
limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{
global: 3,
perPlugin: 1,
perUser: 3,
perToken: 3,
})
releaseFirst, err := limiter.acquire("shared", 1, 1)
require.NoError(t, err)
defer releaseFirst()
release, err := limiter.acquire("shared", 2, 2)
assert.Nil(t, release)
assertLimitError(t, err, "plugin", 1)
releaseOther, err := limiter.acquire("other", 2, 2)
require.NoError(t, err)
defer releaseOther()
})
t.Run("user", func(t *testing.T) {
limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{
global: 3,
perPlugin: 3,
perUser: 1,
perToken: 3,
})
releaseFirst, err := limiter.acquire("first", 1, 1)
require.NoError(t, err)
defer releaseFirst()
release, err := limiter.acquire("second", 1, 2)
assert.Nil(t, release)
assertLimitError(t, err, "user", 1)
releaseOther, err := limiter.acquire("second", 2, 2)
require.NoError(t, err)
defer releaseOther()
})
t.Run("token", func(t *testing.T) {
limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{
global: 3,
perPlugin: 3,
perUser: 3,
perToken: 1,
})
releaseFirst, err := limiter.acquire("first", 1, 1)
require.NoError(t, err)
defer releaseFirst()
release, err := limiter.acquire("second", 2, 1)
assert.Nil(t, release)
assertLimitError(t, err, "token", 1)
releaseOther, err := limiter.acquire("second", 2, 2)
require.NoError(t, err)
defer releaseOther()
})
}
func TestPluginProtocolObservationLimiterRollsBackFailedAdmission(t *testing.T) {
t.Run("user failure", func(t *testing.T) {
limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{
global: 2,
perPlugin: 2,
perUser: 1,
perToken: 2,
})
releaseHeld, err := limiter.acquire("first", 1, 1)
require.NoError(t, err)
defer releaseHeld()
release, err := limiter.acquire("second", 1, 2)
assert.Nil(t, release)
assertLimitError(t, err, "user", 1)
releaseReplacement, err := limiter.acquire("second", 2, 2)
require.NoError(t, err)
defer releaseReplacement()
})
t.Run("token failure", func(t *testing.T) {
limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{
global: 2,
perPlugin: 2,
perUser: 2,
perToken: 1,
})
releaseHeld, err := limiter.acquire("first", 1, 1)
require.NoError(t, err)
defer releaseHeld()
release, err := limiter.acquire("second", 2, 1)
assert.Nil(t, release)
assertLimitError(t, err, "token", 1)
releaseReplacement, err := limiter.acquire("second", 2, 2)
require.NoError(t, err)
defer releaseReplacement()
})
}
func TestPluginProtocolObservationLimiterReleaseIsIdempotent(t *testing.T) {
limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{
global: 1,
perPlugin: 1,
perUser: 1,
perToken: 1,
})
release, err := limiter.acquire("plugin", 1, 1)
require.NoError(t, err)
release()
release()
releaseAgain, err := limiter.acquire("plugin", 1, 1)
require.NoError(t, err)
releaseAgain()
}
func TestPluginProtocolObservationLimiterRejectsMissingIdentityWithoutConsumingCapacity(t *testing.T) {
limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{
global: 1,
perPlugin: 1,
perUser: 1,
perToken: 1,
})
for _, testCase := range []struct {
name string
pluginKey string
userID int
tokenID int
}{
{name: "empty plugin", userID: 1, tokenID: 1},
{name: "blank plugin", pluginKey: " \t", userID: 1, tokenID: 1},
{name: "zero user", pluginKey: "plugin", tokenID: 1},
{name: "negative user", pluginKey: "plugin", userID: -1, tokenID: 1},
{name: "zero token", pluginKey: "plugin", userID: 1},
{name: "negative token", pluginKey: "plugin", userID: 1, tokenID: -1},
} {
t.Run(testCase.name, func(t *testing.T) {
release, err := limiter.acquire(testCase.pluginKey, testCase.userID, testCase.tokenID)
assert.Nil(t, release)
assert.ErrorIs(t, err, errInvalidPluginProtocolObservationIdentity)
})
}
release, err := limiter.acquire("plugin", 1, 1)
require.NoError(t, err)
release()
}
func TestPluginProtocolObservationLimiterConcurrentAdmissionsRespectCap(t *testing.T) {
const (
workerCount = 8
globalLimit = 3
)
limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{
global: globalLimit,
perPlugin: workerCount,
perUser: workerCount,
perToken: workerCount,
})
start := make(chan struct{})
releases := make(chan func(), workerCount)
errorsFound := make(chan error, workerCount)
var workers sync.WaitGroup
workers.Add(workerCount)
for worker := 1; worker <= workerCount; worker++ {
go func(id int) {
defer workers.Done()
<-start
release, err := limiter.acquire("plugin", id, id)
if err != nil {
errorsFound <- err
return
}
releases <- release
}(worker)
}
close(start)
workers.Wait()
close(releases)
close(errorsFound)
assert.Len(t, releases, globalLimit)
assert.Len(t, errorsFound, workerCount-globalLimit)
for err := range errorsFound {
assert.ErrorIs(t, err, errPluginProtocolObservationLimitExceeded)
}
for release := range releases {
release()
}
release, err := limiter.acquire("plugin", 1, 1)
require.NoError(t, err)
release()
}
func assertLimitError(
t *testing.T,
err error,
expectedScope string,
expectedLimit int,
) {
t.Helper()
require.Error(t, err)
assert.ErrorIs(t, err, errPluginProtocolObservationLimitExceeded)
var limitError *pluginProtocolObservationLimitError
require.True(t, errors.As(err, &limitError))
assert.Equal(t, expectedScope, limitError.scope)
assert.Equal(t, expectedLimit, limitError.limit)
}
+305
View File
@@ -0,0 +1,305 @@
package controller
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/model"
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
builtinplugins "github.com/QuantumNous/new-api/plugins"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/gin-gonic/gin"
"github.com/openai/openai-go"
"github.com/openai/openai-go/option"
"github.com/openai/openai-go/responses"
"github.com/openai/openai-go/shared"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestTaskPluginResponsesNonStreamDecodesWithOfficialGoSDK(t *testing.T) {
pinned := compilePluginProtocolTestEndpoint(t, "official-sdk-non-stream", `
export const protocols = {openai_responses: {
renderEvents: function() { throw new Error("non-stream called renderEvents"); },
renderFinal: function() {
return {
output: [{
type: "message",
status: "completed",
role: "assistant",
content: [{
type: "output_text",
text: "official-sdk-final",
annotations: [],
logprobs: []
}]
}]
};
}
}};
`)
deps := pluginProtocolTestDeps()
deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) {
return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_sdk_final", nil), nil
}
deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) {
return &model.Task{
TaskID: "task_sdk_final",
UserId: 71,
Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key),
Status: model.TaskStatusSuccess,
}, true, nil
}
server := newPluginProtocolSDKTestServer(t, pinned, deps)
defer server.Close()
client := openai.NewClient(
option.WithAPIKey("test-key"),
option.WithBaseURL(server.URL+"/v1/"),
option.WithMaxRetries(0),
)
requestContext, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
response, err := client.Responses.New(requestContext, responses.ResponseNewParams{
Model: shared.ResponsesModel("video-model"),
Input: responses.ResponseNewParamsInputUnion{
OfString: openai.String("create a video"),
},
})
require.NoError(t, err)
assert.Equal(t, "resp_sdk_final", response.ID)
assert.Equal(t, responses.ResponseStatusCompleted, response.Status)
assert.Equal(t, "video-model", response.Model)
assert.Equal(t, "official-sdk-final", response.OutputText())
assert.Equal(t, "task_sdk_final", response.Metadata["task_id"])
}
func TestTaskPluginResponsesStreamDecodesWithOfficialGoSDK(t *testing.T) {
pinned := compilePluginProtocolTestEndpoint(t, "official-sdk-stream", `
export const protocols = {openai_responses: {
renderEvents: function() {
return {events: [{type: "output", data: "official-sdk-stream"}], done: true};
},
renderFinal: function() { throw new Error("stream called renderFinal"); }
}};
`)
deps := pluginProtocolTestDeps()
deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) {
return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_sdk_stream", nil), nil
}
deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) {
return &model.Task{
TaskID: "task_sdk_stream",
UserId: 71,
Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key),
Status: model.TaskStatusSuccess,
}, true, nil
}
server := newPluginProtocolSDKTestServer(t, pinned, deps)
defer server.Close()
client := openai.NewClient(
option.WithAPIKey("test-key"),
option.WithBaseURL(server.URL+"/v1/"),
option.WithMaxRetries(0),
)
requestContext, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
stream := client.Responses.NewStreaming(requestContext, responses.ResponseNewParams{
Model: shared.ResponsesModel("video-model"),
Input: responses.ResponseNewParamsInputUnion{
OfString: openai.String("create a video"),
},
})
defer stream.Close()
eventTypes := make([]string, 0, 8)
sequenceNumbers := make([]int64, 0, 8)
var completedText string
for stream.Next() {
event := stream.Current()
eventTypes = append(eventTypes, event.Type)
sequenceNumbers = append(sequenceNumbers, event.SequenceNumber)
if event.Type == "response.completed" {
completedText = event.Response.OutputText()
}
}
require.NoError(t, stream.Err())
assert.Equal(t, []string{
"response.created",
"response.output_item.added",
"response.content_part.added",
"response.output_text.delta",
"response.output_text.done",
"response.content_part.done",
"response.output_item.done",
"response.completed",
}, eventTypes)
assert.Equal(t, []int64{0, 1, 2, 3, 4, 5, 6, 7}, sequenceNumbers)
assert.Equal(t, "official-sdk-stream", completedText)
}
func TestBuiltInKlingResponsesNonStreamDecodesWithOfficialGoSDK(t *testing.T) {
pinned, deps := builtInKlingProtocolSDKFixture(t, "task_kling_sdk_final")
server := newPluginProtocolSDKTestServer(t, pinned, deps)
defer server.Close()
client := openai.NewClient(
option.WithAPIKey("test-key"),
option.WithBaseURL(server.URL+"/v1/"),
option.WithMaxRetries(0),
)
requestContext, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
response, err := client.Responses.New(requestContext, responses.ResponseNewParams{
Model: shared.ResponsesModel("kling-v2-master"),
Input: responses.ResponseNewParamsInputUnion{
OfString: openai.String("camera orbit"),
},
})
require.NoError(t, err)
assert.Equal(t, responses.ResponseStatusCompleted, response.Status)
assert.Equal(t, "kling-v2-master", response.Model)
assert.Contains(t, response.OutputText(), "https://gateway.example/v1/tasks/task_kling_sdk_final/artifacts/video/content")
assert.NotContains(t, response.OutputText(), "upstream.example")
assert.Equal(t, "kling", response.Metadata["vendor"])
assert.Equal(t, "task_kling_sdk_final", response.Metadata["task_id"])
}
func TestBuiltInKlingResponsesStreamDecodesWithOfficialGoSDK(t *testing.T) {
pinned, deps := builtInKlingProtocolSDKFixture(t, "task_kling_sdk_stream")
server := newPluginProtocolSDKTestServer(t, pinned, deps)
defer server.Close()
client := openai.NewClient(
option.WithAPIKey("test-key"),
option.WithBaseURL(server.URL+"/v1/"),
option.WithMaxRetries(0),
)
requestContext, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
stream := client.Responses.NewStreaming(requestContext, responses.ResponseNewParams{
Model: shared.ResponsesModel("kling-v2-master"),
Input: responses.ResponseNewParamsInputUnion{
OfString: openai.String("camera orbit"),
},
})
defer stream.Close()
eventTypes := make([]string, 0, 8)
sequenceNumbers := make([]int64, 0, 8)
var completedText string
for stream.Next() {
event := stream.Current()
eventTypes = append(eventTypes, event.Type)
sequenceNumbers = append(sequenceNumbers, event.SequenceNumber)
if event.Type == "response.completed" {
completedText = event.Response.OutputText()
}
}
require.NoError(t, stream.Err())
assert.Equal(t, []string{
"response.created",
"response.output_item.added",
"response.content_part.added",
"response.output_text.delta",
"response.output_text.done",
"response.content_part.done",
"response.output_item.done",
"response.completed",
}, eventTypes)
assert.Equal(t, []int64{0, 1, 2, 3, 4, 5, 6, 7}, sequenceNumbers)
assert.Contains(t, completedText, "https://gateway.example/v1/tasks/task_kling_sdk_stream/artifacts/video/content")
assert.NotContains(t, completedText, "upstream.example")
}
func builtInKlingProtocolSDKFixture(t *testing.T, taskID string) (pluginruntime.PinnedEndpoint, pluginProtocolBridgeDeps) {
t.Helper()
source, err := builtinplugins.Source("kling")
require.NoError(t, err)
registry := pluginruntime.NewRegistry()
plugin, err := registry.RegisterFactory(source, pluginruntime.Options{Key: "kling"})
require.NoError(t, err)
binding, found := registry.Generation().LookupEndpoint(http.MethodPost, "/v1/responses", "kling-v2-master")
require.True(t, found)
pinned := pluginruntime.PinnedEndpoint{
Generation: registry.Generation(),
Plugin: plugin,
Protocol: binding.Protocol,
Operation: binding.Operation,
Model: binding.Model,
Candidates: []pluginruntime.ProtocolBinding{binding},
}
deps := pluginProtocolTestDeps()
deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) {
return pluginProtocolTestOutcome(info, plugin.Meta.Key, taskID, nil), nil
}
deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) {
task := &model.Task{
TaskID: taskID, Platform: constant.TaskPlatform(plugin.Meta.Key), UserId: 71,
Status: model.TaskStatusSuccess, Progress: "100%", CreatedAt: 1_710_000_000,
}
task.SetData(map[string]any{
"code": 0,
"data": map[string]any{
"task_id": "upstream-private",
"task_status": "succeed",
"task_result": map[string]any{
"videos": []any{map[string]any{"url": "https://upstream.example/private-video.mp4"}},
},
},
})
return task, true, nil
}
deps.artifactContentURL = func(publicTaskID, artifactKey string) (string, error) {
return "https://gateway.example/v1/tasks/" + publicTaskID + "/artifacts/" + artifactKey + "/content", nil
}
return pinned, deps
}
func newPluginProtocolSDKTestServer(
t *testing.T,
pinned pluginruntime.PinnedEndpoint,
deps pluginProtocolBridgeDeps,
) *httptest.Server {
t.Helper()
return httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
if request.Method != http.MethodPost || request.URL.Path != "/v1/responses" {
http.NotFound(writer, request)
return
}
var requestBody map[string]any
if err := common.DecodeJson(request.Body, &requestBody); err != nil {
http.Error(writer, "invalid request", http.StatusBadRequest)
return
}
modelName, _ := requestBody["model"].(string)
stream, _ := requestBody["stream"].(bool)
c, _ := gin.CreateTestContext(writer)
c.Request = request
common.SetContextKey(c, constant.ContextKeyUserId, 71)
common.SetContextKey(c, constant.ContextKeyTokenId, 81)
common.SetContextKey(c, constant.ContextKeyUsingGroup, "default")
c.Set("resolved_task_model", modelName)
c.Set(pluginruntime.ContextKeyProtocolRequest, pluginruntime.ProtocolRequestContext{
RouteRequestContext: pluginruntime.RouteRequestContext{
Path: request.URL.Path,
Method: request.Method,
Params: map[string]string{},
Query: request.URL.Query(),
RequestBody: requestBody,
},
Protocol: pinned.Protocol,
Stream: stream,
})
serveTaskPluginProtocol(c, pinned, deps)
}))
}
File diff suppressed because it is too large. Load diff
+250 -37
View File
@@ -15,8 +15,10 @@ import (
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/middleware"
"github.com/QuantumNous/new-api/model"
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
"github.com/QuantumNous/new-api/relay"
"github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relay/helper"
@@ -344,7 +346,7 @@ func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) b
if retryTimes <= 0 {
return false
}
if _, ok := c.Get("specific_channel_id"); ok {
if service.GetChannelConstraints(c).SuppressesRetry() {
return false
}
code := openaiErr.StatusCode
@@ -397,6 +399,7 @@ func processChannelError(c *gin.Context, channelError types.ChannelError, err *t
}
service.AppendChannelAffinityAdminInfo(c, adminInfo)
other["admin_info"] = adminInfo
service.AppendTaskPluginContextAuditInfo(c, other)
startTime := common.GetContextKeyTime(c, constant.ContextKeyRequestStartTime)
if startTime.IsZero() {
startTime = time.Now()
@@ -474,6 +477,33 @@ func RelayNotFound(c *gin.Context) {
})
}
// RelayTaskPluginEndpoint keeps unclaimed shared-endpoint traffic on its
// existing handler while claimed requests enter the generation-pinned
// host-owned protocol bridge.
func RelayTaskPluginEndpoint(c *gin.Context, fallback gin.HandlerFunc) {
pinnedValue, exists := c.Get(pluginruntime.ContextKeyPinnedEndpoint)
if !exists {
fallback(c)
return
}
pinned, ok := pinnedValue.(pluginruntime.PinnedEndpoint)
if !ok || pinned.Plugin == nil || pinned.Generation == nil {
c.JSON(http.StatusInternalServerError, gin.H{
"error": gin.H{
"message": "Task protocol request failed",
"type": "new_api_error",
"code": "task_protocol_error",
},
})
return
}
if pinned.Protocol != "openai_responses" {
fallback(c)
return
}
serveTaskPluginProtocol(c, pinned, defaultPluginProtocolBridgeDeps())
}
func RelayTaskFetch(c *gin.Context) {
relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
if err != nil {
@@ -489,29 +519,76 @@ func RelayTaskFetch(c *gin.Context) {
}
}
type taskSubmissionOutcome struct {
Result *relay.TaskSubmitResult
Task *model.Task
RelayInfo *relaycommon.RelayInfo
}
func RelayTask(c *gin.Context) {
relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
if err != nil {
c.JSON(http.StatusInternalServerError, &taskdto.TaskError{
respondTaskSubmissionError(c, &taskdto.TaskError{
Code: "gen_relay_info_failed",
Message: err.Error(),
StatusCode: http.StatusInternalServerError,
})
return
}
if action := c.GetString("task_action"); action != "" {
relayInfo.Action = action
}
if taskErr := relay.ResolveOriginTask(c, relayInfo); taskErr != nil {
respondTaskError(c, taskErr)
respondTaskSubmissionError(c, taskErr)
return
}
if taskErr := relay.ApplyOriginTaskAffinity(c, relayInfo); taskErr != nil {
respondTaskSubmissionError(c, taskErr)
return
}
outcome, taskErr := executeTaskSubmission(c, relayInfo)
if taskErr != nil {
respondTaskSubmissionError(c, taskErr)
return
}
presentTaskSubmission(c, outcome)
}
// executeTaskSubmission owns the retry, billing, and persistence lifecycle.
// It deliberately performs no client response writes so JSON and protocol
// presenters share the same durable task barrier. Its cancellation semantics
// come from c.Request.Context: native task endpoints use the client context,
// while the Responses bridge supplies an independently bounded context.
func executeTaskSubmission(c *gin.Context, relayInfo *relaycommon.RelayInfo) (*taskSubmissionOutcome, *taskdto.TaskError) {
return executeTaskSubmissionWith(c, relayInfo, relay.RelayTaskSubmit)
}
type taskSubmitAttempt func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *taskdto.TaskError)
func executeTaskSubmissionWith(
c *gin.Context,
relayInfo *relaycommon.RelayInfo,
submit taskSubmitAttempt,
) (*taskSubmissionOutcome, *taskdto.TaskError) {
diagnostics := newTaskPluginSubmitDiagnostics(c)
diagnostics.start(relayInfo)
var result *relay.TaskSubmitResult
var taskErr *taskdto.TaskError
durable := false
stage := "start"
defer func() {
if taskErr != nil && relayInfo.Billing != nil {
if !durable && relayInfo.Billing != nil {
diagnostics.refund(stage)
relayInfo.Billing.Refund(c)
}
}()
stage = "before_attempt"
if requestErr := c.Request.Context().Err(); requestErr != nil {
diagnostics.cancelled("before_attempt", 0)
return nil, service.TaskErrorWrapperLocal(requestErr, "request_cancelled", http.StatusRequestTimeout)
}
retryParam := &service.RetryParam{
Ctx: c,
@@ -522,6 +599,12 @@ func RelayTask(c *gin.Context) {
}
for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
stage = "select_channel"
if requestErr := c.Request.Context().Err(); requestErr != nil {
diagnostics.cancelled("before_attempt", retryParam.GetRetry()+1)
taskErr = service.TaskErrorWrapperLocal(requestErr, "request_cancelled", http.StatusRequestTimeout)
break
}
var channel *model.Channel
if lockedCh, ok := relayInfo.LockedChannel.(*model.Channel); ok && lockedCh != nil {
@@ -541,10 +624,12 @@ func RelayTask(c *gin.Context) {
break
}
}
diagnostics.attempt(retryParam.GetRetry()+1, channel, relayInfo.LockedChannel != nil)
addUsedChannel(c, channel.Id)
bodyStorage, bodyErr := common.GetBodyStorage(c)
if bodyErr != nil {
stage = "read_body"
if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) {
taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusRequestEntityTooLarge)
} else {
@@ -554,8 +639,15 @@ func RelayTask(c *gin.Context) {
}
c.Request.Body = io.NopCloser(bodyStorage)
result, taskErr = relay.RelayTaskSubmit(c, relayInfo)
stage = "submit"
result, taskErr = submit(c, relayInfo)
if requestErr := c.Request.Context().Err(); requestErr != nil {
diagnostics.cancelled("after_submit", retryParam.GetRetry()+1)
taskErr = service.TaskErrorWrapperLocal(requestErr, "request_cancelled", http.StatusRequestTimeout)
break
}
if taskErr == nil {
diagnostics.attemptSucceeded(retryParam.GetRetry()+1, result)
break
}
@@ -566,7 +658,9 @@ func RelayTask(c *gin.Context) {
types.NewOpenAIError(taskErr.Error, types.ErrorCodeBadResponseStatusCode, taskErr.StatusCode))
}
if !shouldRetryTaskRelay(c, channel.Id, taskErr, common.RetryTimes-retryParam.GetRetry()) {
willRetry := shouldRetryTaskRelay(c, channel.Id, taskErr, common.RetryTimes-retryParam.GetRetry())
diagnostics.attemptFailed(retryParam.GetRetry()+1, channel, taskErr, willRetry)
if !willRetry {
break
}
}
@@ -577,38 +671,157 @@ func RelayTask(c *gin.Context) {
logger.LogInfo(c, retryLogStr)
}
// ── 成功:结算 + 日志 + 插入任务 ──
if taskErr == nil {
if settleErr := service.SettleBilling(c, relayInfo, result.Quota); settleErr != nil {
common.SysError("settle task billing error: " + settleErr.Error())
}
service.LogTaskConsumption(c, relayInfo)
task := model.InitTask(result.Platform, relayInfo)
task.PrivateData.UpstreamTaskID = result.UpstreamTaskID
task.PrivateData.BillingSource = relayInfo.BillingSource
task.PrivateData.SubscriptionId = relayInfo.SubscriptionId
task.PrivateData.TokenId = relayInfo.TokenId
task.PrivateData.NodeName = common.NodeName
task.PrivateData.BillingContext = &model.TaskBillingContext{
ModelPrice: relayInfo.PriceData.ModelPrice,
GroupRatio: relayInfo.PriceData.GroupRatioInfo.GroupRatio,
ModelRatio: relayInfo.PriceData.ModelRatio,
OtherRatios: relayInfo.PriceData.OtherRatios(),
OriginModelName: relayInfo.OriginModelName,
PerCallBilling: common.StringsContains(constant.TaskPricePatches, relayInfo.OriginModelName) || relayInfo.PriceData.UsePrice,
}
task.Quota = result.Quota
task.Data = result.TaskData
task.Action = relayInfo.Action
if insertErr := task.Insert(); insertErr != nil {
common.SysError("insert task error: " + insertErr.Error())
}
}
if taskErr != nil {
respondTaskError(c, taskErr)
diagnostics.failed(stage, "task_error", taskErr, false)
return nil, taskErr
}
if result == nil {
taskErr = service.TaskErrorWrapperLocal(errors.New("task submission returned no result"), "task_submit_failed", http.StatusInternalServerError)
diagnostics.failed("submit", "missing_result", taskErr, false)
return nil, taskErr
}
if requestErr := c.Request.Context().Err(); requestErr != nil {
diagnostics.cancelled("before_reserve", retryParam.GetRetry()+1)
return nil, service.TaskErrorWrapperLocal(requestErr, "request_cancelled", http.StatusRequestTimeout)
}
// Reserve any submit-time upward billing adjustment before persistence.
// This keeps insertion failures fully refundable while ensuring settlement
// after the barrier normally has a zero positive delta.
if relayInfo.Billing != nil {
stage = "reserve"
diagnostics.reserve("reserve_start", result.Quota)
if reserveErr := relayInfo.Billing.Reserve(result.Quota); reserveErr != nil {
common.SysError("reserve adjusted task billing error: " + reserveErr.Error())
taskErr = service.TaskErrorWrapperLocal(errors.New("insufficient quota for adjusted task cost"), string(types.ErrorCodeInsufficientUserQuota), http.StatusForbidden)
diagnostics.failed("reserve", "insufficient_quota", taskErr, false)
return nil, taskErr
}
diagnostics.reserve("reserve_complete", result.Quota)
}
if requestErr := c.Request.Context().Err(); requestErr != nil {
diagnostics.cancelled("before_insert", retryParam.GetRetry()+1)
return nil, service.TaskErrorWrapperLocal(requestErr, "request_cancelled", http.StatusRequestTimeout)
}
stage = "insert"
task := model.InitTask(result.Platform, relayInfo)
task.PrivateData.Execution = service.TaskExecutionSnapshotFromContext(c)
task.PrivateData.UpstreamTaskID = result.UpstreamTaskID
task.PrivateData.BillingSource = relayInfo.BillingSource
task.PrivateData.SubscriptionId = relayInfo.SubscriptionId
task.PrivateData.TokenId = relayInfo.TokenId
task.PrivateData.NodeName = common.NodeName
task.PrivateData.BillingContext = &model.TaskBillingContext{
ModelPrice: relayInfo.PriceData.ModelPrice,
GroupRatio: relayInfo.PriceData.GroupRatioInfo.GroupRatio,
ModelRatio: relayInfo.PriceData.ModelRatio,
OtherRatios: relayInfo.PriceData.OtherRatios(),
OriginModelName: relayInfo.OriginModelName,
PerCallBilling: common.StringsContains(constant.TaskPricePatches, relayInfo.OriginModelName) || relayInfo.PriceData.UsePrice,
TieredSnapshot: relayInfo.TieredBillingSnapshot,
}
task.Quota = result.Quota
task.Data = result.TaskData
task.Action = relayInfo.Action
if immediate := result.Immediate; immediate != nil {
task.Status = model.TaskStatus(immediate.Status)
task.Progress = immediate.Progress
if immediate.Status == model.TaskStatusSuccess || immediate.Status == model.TaskStatusFailure {
task.FinishTime = time.Now().Unix()
}
if immediate.Status == model.TaskStatusFailure {
task.FailReason = immediate.Reason
}
if immediate.Url != "" {
task.PrivateData.ResultURL = immediate.Url
} else if immediate.Status == model.TaskStatusSuccess {
task.PrivateData.ResultURL = taskcommon.BuildProxyURL(task.TaskID)
}
}
diagnostics.insertStart(task)
if insertErr := task.InsertWithContext(c.Request.Context()); insertErr != nil {
common.SysError("insert task error: " + insertErr.Error())
taskErr = service.TaskErrorWrapperLocal(errors.New("failed to persist task"), "task_insert_failed", http.StatusInternalServerError)
diagnostics.failed("insert", "database_error", taskErr, false)
return nil, taskErr
}
durable = true
stage = "settle"
diagnostics.durable(task)
diagnostics.settleStart(task, result.Quota)
if settleErr := service.SettleBilling(c, relayInfo, result.Quota); settleErr != nil {
common.SysError("settle task billing error: " + settleErr.Error())
taskErr = service.TaskErrorWrapperLocal(errors.New("failed to settle task billing"), "task_billing_settlement_failed", http.StatusInternalServerError)
diagnostics.failed("settle", "billing_error", taskErr, true)
return nil, taskErr
}
service.LogTaskConsumption(c, relayInfo, task)
diagnostics.complete(task, result.Quota)
return &taskSubmissionOutcome{Result: result, Task: task, RelayInfo: relayInfo}, nil
}
func presentTaskSubmission(c *gin.Context, outcome *taskSubmissionOutcome) {
diagnostics := newTaskPluginSubmitDiagnostics(c)
otherRatios := outcome.RelayInfo.PriceData.OtherRatios()
if otherRatios == nil {
otherRatios = map[string]float64{}
}
if ratiosJSON, err := common.Marshal(otherRatios); err == nil {
c.Header("X-New-Api-Other-Ratios", string(ratiosJSON))
}
if pinnedValue, exists := c.Get(pluginruntime.ContextKeyPinnedRoute); exists {
if pinned, ok := pinnedValue.(pluginruntime.PinnedRoute); ok && pinned.Plugin != nil && pinned.Route.Render != "" {
view, err := service.BuildTaskPluginView(outcome.Task)
requestValue, _ := c.Get(pluginruntime.ContextKeyRouteRequest)
requestContext, _ := requestValue.(pluginruntime.RouteRequestContext)
if err == nil {
viewValue, valueErr := taskPluginProtocolJSONValue(view)
if valueErr == nil {
if body, callErr := pinned.Plugin.Engine.CallPath(c.Request.Context(), "native", []string{pinned.Route.Render}, requestContext.JSValue(), viewValue); callErr == nil {
diagnostics.present(outcome.Task, "native_presenter")
c.JSON(http.StatusOK, body)
return
} else {
logger.LogError(c, "task plugin native submit presenter failed: "+callErr.Error())
}
} else {
logger.LogError(c, "encode task plugin native submit view failed: "+valueErr.Error())
}
} else {
logger.LogError(c, "build task plugin native submit view failed: "+err.Error())
}
}
}
if pinnedValue, exists := c.Get(pluginruntime.ContextKeyPinnedEndpoint); exists {
if pinned, ok := pinnedValue.(pluginruntime.PinnedEndpoint); ok && pinned.Protocol == "openai_video" && pinned.Operation.Name == "create" {
diagnostics.present(outcome.Task, "openai_video_create")
c.JSON(http.StatusOK, outcome.Task.ToOpenAIVideo())
return
}
}
createdAt := outcome.Task.CreatedAt
if createdAt == 0 {
createdAt = outcome.Task.SubmitTime
}
diagnostics.present(outcome.Task, "host_fallback")
c.JSON(http.StatusOK, map[string]any{
"id": outcome.Task.TaskID,
"task_id": outcome.Task.TaskID,
"status": "queued",
"model": outcome.RelayInfo.OriginModelName,
"created_at": createdAt,
})
}
func respondTaskSubmissionError(c *gin.Context, taskErr *taskdto.TaskError) {
newTaskPluginSubmitDiagnostics(c).presentError(taskErr)
if middleware.RespondTaskPluginError(c, taskErr) {
return
}
respondTaskError(c, taskErr)
}
// respondTaskError 统一输出 Task 错误响应(含 429 限流提示改写)
@@ -629,7 +842,7 @@ func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *taskdto.TaskEr
if retryTimes <= 0 {
return false
}
if _, ok := c.Get("specific_channel_id"); ok {
if service.GetChannelConstraints(c).SuppressesRetry() {
return false
}
if taskErr.StatusCode == http.StatusTooManyRequests {
+389
View File
@@ -0,0 +1,389 @@
package controller
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/model"
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/relay"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)
type taskSubmissionTestBilling struct {
events *[]string
settleErr error
onSettle func()
refunds int
}
func (b *taskSubmissionTestBilling) Settle(int) error {
*b.events = append(*b.events, "settle")
if b.onSettle != nil {
b.onSettle()
}
return b.settleErr
}
func (b *taskSubmissionTestBilling) Refund(*gin.Context) {
*b.events = append(*b.events, "refund")
b.refunds++
}
func (b *taskSubmissionTestBilling) NeedsRefund() bool { return b.refunds == 0 }
func (b *taskSubmissionTestBilling) GetPreConsumedQuota() int { return 0 }
func (b *taskSubmissionTestBilling) Reserve(int) error {
*b.events = append(*b.events, "reserve")
return nil
}
func TestPresentTaskSubmissionUsesNativePresenterAfterPersistence(t *testing.T) {
plugin, err := pluginruntime.CompilePlugin(`
export const meta = {apiVersion:1,key:"presenter-test",name:"Presenter",version:"1.0.0",author:{name:"Test"},models:["model"],fetchMode:"per_task",routes:[{method:"POST",path:"/vendor/jobs",type:"submit",decode:"decode",render:"created"}]};
export const native = {decode:function(ctx){return {kind:"submit",model:"model",requestBody:ctx.body.value};},created:function(ctx,task){return {data:{task_id:task.task_id},upstream:task.data};}};
export function buildSubmitRequest(){return {}} export function parseSubmitResponse(){return {taskId:"upstream"}} export function buildQueryRequest(){return {}} export function parseTaskResult(){return {status:"SUCCESS"}}
`, pluginruntime.Options{})
require.NoError(t, err)
priceData := types.PriceData{}
priceData.AddOtherRatio("seconds", 5)
task := &model.Task{TaskID: "task_public", SubmitTime: 123}
task.SetData(map[string]any{"task_id": "upstream_private"})
outcome := &taskSubmissionOutcome{
Result: &relay.TaskSubmitResult{},
Task: task,
RelayInfo: &relaycommon.RelayInfo{PriceData: priceData},
}
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{"model":"model"}`))
c.Set(pluginruntime.ContextKeyPinnedRoute, pluginruntime.PinnedRoute{Plugin: plugin, Route: plugin.Meta.Routes[0]})
c.Set(pluginruntime.ContextKeyRouteRequest, pluginruntime.RouteRequestContext{Path: "/vendor/jobs", Method: http.MethodPost, Body: map[string]any{"kind": "json", "value": map[string]any{"model": "model"}}})
presentTaskSubmission(c, outcome)
assert.JSONEq(t, `{
"data":{"task_id":"task_public"},
"upstream":{"task_id":"upstream_private"}
}`, recorder.Body.String())
assert.JSONEq(t, `{"seconds":5}`, recorder.Header().Get("X-New-Api-Other-Ratios"))
}
func TestPresentTaskSubmissionFallbackUsesPersistedPublicID(t *testing.T) {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
outcome := &taskSubmissionOutcome{
Result: &relay.TaskSubmitResult{},
Task: &model.Task{TaskID: "task_persisted", SubmitTime: 456},
RelayInfo: &relaycommon.RelayInfo{OriginModelName: "video-model"},
}
presentTaskSubmission(c, outcome)
assert.JSONEq(t, `{
"id":"task_persisted",
"task_id":"task_persisted",
"status":"queued",
"model":"video-model",
"created_at":456
}`, recorder.Body.String())
}
func TestPresentTaskSubmissionUsesHostOpenAIVideoCreateReceipt(t *testing.T) {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Set(pluginruntime.ContextKeyPinnedEndpoint, pluginruntime.PinnedEndpoint{
Protocol: "openai_video",
Operation: pluginruntime.HostProtocolOperation{Name: "create"},
})
task := &model.Task{
TaskID: "task_public",
Status: model.TaskStatusSubmitted,
Progress: "0%",
CreatedAt: 456,
Properties: model.Properties{OriginModelName: "video-model"},
}
outcome := &taskSubmissionOutcome{Result: &relay.TaskSubmitResult{}, Task: task, RelayInfo: &relaycommon.RelayInfo{}}
presentTaskSubmission(c, outcome)
assert.JSONEq(t, `{"id":"task_public","object":"video","model":"video-model","status":"queued","progress":0,"created_at":456}`, recorder.Body.String())
assert.NotContains(t, recorder.Body.String(), "task_id")
}
func TestExecuteTaskSubmissionRefundsWhenInsertFails(t *testing.T) {
events := make([]string, 0, 3)
database := setupTaskSubmissionDatabase(t, false, &events)
_ = database
billing := &taskSubmissionTestBilling{events: &events}
c := taskSubmissionTestContext()
info := taskSubmissionRelayInfo(billing)
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
return &relay.TaskSubmitResult{
UpstreamTaskID: "upstream_private",
Platform: constant.TaskPlatform("plugin"),
}, nil
})
assert.Nil(t, outcome)
require.NotNil(t, taskErr)
assert.Equal(t, "task_insert_failed", taskErr.Code)
assert.Equal(t, []string{"reserve", "insert", "refund"}, events)
assert.Equal(t, 1, billing.refunds)
assert.False(t, c.Writer.Written())
}
func TestExecuteTaskSubmissionSettlementFailureStaysDurableAndWritesNothing(t *testing.T) {
events := make([]string, 0, 3)
database := setupTaskSubmissionDatabase(t, true, &events)
billing := &taskSubmissionTestBilling{events: &events, settleErr: errors.New("settlement failed")}
c := taskSubmissionTestContext()
info := taskSubmissionRelayInfo(billing)
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
return &relay.TaskSubmitResult{
UpstreamTaskID: "upstream_private",
Platform: constant.TaskPlatform("plugin"),
}, nil
})
assert.Nil(t, outcome)
require.NotNil(t, taskErr)
assert.Equal(t, "task_billing_settlement_failed", taskErr.Code)
assert.Equal(t, []string{"reserve", "insert", "settle"}, events)
assert.Zero(t, billing.refunds)
var count int64
require.NoError(t, database.Model(&model.Task{}).Where("task_id = ?", "task_public").Count(&count).Error)
assert.Equal(t, int64(1), count)
assert.False(t, c.Writer.Written())
}
func TestExecuteTaskSubmissionPersistsPinnedPluginProvenance(t *testing.T) {
events := make([]string, 0, 3)
database := setupTaskSubmissionDatabase(t, true, &events)
previousLogConsumeEnabled := common.LogConsumeEnabled
common.LogConsumeEnabled = false
t.Cleanup(func() { common.LogConsumeEnabled = previousLogConsumeEnabled })
c := taskSubmissionTestContext()
c.Set(common.RequestIdKey, "request-public")
c.Set(pluginruntime.ContextKeyPinnedPlugin, pluginruntime.PinnedPlugin{
Generation: &pluginruntime.RoutingGeneration{Number: 42},
Plugin: &pluginruntime.LoadedPlugin{Meta: pluginruntime.Meta{
Key: "document-parser",
Name: "Document Parser",
Version: "1.2.3",
APIVersion: 1,
Author: pluginruntime.AuthorMeta{
Name: "Community Author",
URL: "https://plugins.example/author",
},
}},
})
billing := &taskSubmissionTestBilling{events: &events}
info := taskSubmissionRelayInfo(billing)
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
return &relay.TaskSubmitResult{
UpstreamTaskID: "upstream-private",
Platform: constant.TaskPlatform("document-parser"),
}, nil
})
require.Nil(t, taskErr)
require.NotNil(t, outcome)
require.NotNil(t, outcome.Task.PrivateData.Execution)
require.NotNil(t, outcome.Task.PrivateData.Execution.TaskPlugin)
assert.Equal(t, "request-public", outcome.Task.PrivateData.Execution.RequestID)
assert.Equal(t, "/plugin/submit", outcome.Task.PrivateData.Execution.RequestPath)
assert.Equal(t, "1.2.3", outcome.Task.PrivateData.Execution.TaskPlugin.Version)
assert.Equal(t, uint64(42), outcome.Task.PrivateData.Execution.TaskPlugin.Generation)
require.NotNil(t, outcome.Task.PrivateData.Execution.TaskPlugin.Author)
assert.Equal(t, "Community Author", outcome.Task.PrivateData.Execution.TaskPlugin.Author.Name)
assert.Equal(t, "https://plugins.example/author", outcome.Task.PrivateData.Execution.TaskPlugin.Author.URL)
var stored model.Task
require.NoError(t, database.Where("task_id = ?", "task_public").First(&stored).Error)
require.NotNil(t, stored.PrivateData.Execution)
require.NotNil(t, stored.PrivateData.Execution.TaskPlugin)
assert.Equal(t, "document-parser", stored.PrivateData.Execution.TaskPlugin.Key)
require.NotNil(t, stored.PrivateData.Execution.TaskPlugin.Author)
assert.Equal(t, "Community Author", stored.PrivateData.Execution.TaskPlugin.Author.Name)
assert.Equal(t, "upstream-private", stored.PrivateData.UpstreamTaskID)
}
func TestExecuteTaskSubmissionRefundsCancellationBeforeDurableBarrier(t *testing.T) {
events := make([]string, 0, 2)
setupTaskSubmissionDatabase(t, true, &events)
billing := &taskSubmissionTestBilling{events: &events}
c := taskSubmissionTestContext()
requestContext, cancel := context.WithCancel(c.Request.Context())
c.Request = c.Request.WithContext(requestContext)
info := taskSubmissionRelayInfo(billing)
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
cancel()
return &relay.TaskSubmitResult{
UpstreamTaskID: "upstream_private",
Platform: constant.TaskPlatform("plugin"),
}, nil
})
assert.Nil(t, outcome)
require.NotNil(t, taskErr)
assert.Equal(t, "request_cancelled", taskErr.Code)
assert.Equal(t, []string{"refund"}, events)
assert.Equal(t, 1, billing.refunds)
assert.False(t, c.Writer.Written())
}
func TestExecuteTaskSubmissionDisconnectBeforeUpstreamAcceptanceSkipsSubmitAndRefunds(t *testing.T) {
events := make([]string, 0, 1)
setupTaskSubmissionDatabase(t, true, &events)
billing := &taskSubmissionTestBilling{events: &events}
c := taskSubmissionTestContext()
requestContext, cancel := context.WithCancel(c.Request.Context())
cancel()
c.Request = c.Request.WithContext(requestContext)
info := taskSubmissionRelayInfo(billing)
submitted := false
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
submitted = true
return nil, nil
})
assert.Nil(t, outcome)
require.NotNil(t, taskErr)
assert.Equal(t, "request_cancelled", taskErr.Code)
assert.False(t, submitted)
assert.Equal(t, []string{"refund"}, events)
assert.Equal(t, 1, billing.refunds)
assert.False(t, c.Writer.Written())
}
func TestExecuteTaskSubmissionCallerCancellationDuringSubmitRefundsBeforeDurableBarrier(t *testing.T) {
events := make([]string, 0, 1)
setupTaskSubmissionDatabase(t, true, &events)
billing := &taskSubmissionTestBilling{events: &events}
c := taskSubmissionTestContext()
requestContext, cancel := context.WithCancel(c.Request.Context())
c.Request = c.Request.WithContext(requestContext)
info := taskSubmissionRelayInfo(billing)
submitStarted := make(chan struct{})
done := make(chan struct{})
var outcome *taskSubmissionOutcome
var taskErr *dto.TaskError
go func() {
defer close(done)
outcome, taskErr = executeTaskSubmissionWith(c, info, func(c *gin.Context, _ *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
close(submitStarted)
<-c.Request.Context().Done()
return nil, service.TaskErrorWrapperLocal(c.Request.Context().Err(), "do_request_failed", http.StatusInternalServerError)
})
}()
select {
case <-submitStarted:
case <-time.After(2 * time.Second):
require.FailNow(t, "submission did not start")
}
cancel()
select {
case <-done:
case <-time.After(2 * time.Second):
require.FailNow(t, "submission did not stop after disconnect")
}
assert.Nil(t, outcome)
require.NotNil(t, taskErr)
assert.Equal(t, "request_cancelled", taskErr.Code)
assert.Equal(t, []string{"refund"}, events)
assert.Equal(t, 1, billing.refunds)
assert.False(t, c.Writer.Written())
}
func TestExecuteTaskSubmissionDisconnectAfterDurableInsertDoesNotRefund(t *testing.T) {
events := make([]string, 0, 3)
database := setupTaskSubmissionDatabase(t, true, &events)
previousLogConsumeEnabled := common.LogConsumeEnabled
common.LogConsumeEnabled = false
t.Cleanup(func() { common.LogConsumeEnabled = previousLogConsumeEnabled })
c := taskSubmissionTestContext()
requestContext, cancel := context.WithCancel(c.Request.Context())
c.Request = c.Request.WithContext(requestContext)
billing := &taskSubmissionTestBilling{
events: &events,
onSettle: cancel,
}
info := taskSubmissionRelayInfo(billing)
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
return &relay.TaskSubmitResult{
UpstreamTaskID: "upstream_private",
Platform: constant.TaskPlatform("plugin"),
}, nil
})
require.Nil(t, taskErr)
require.NotNil(t, outcome)
assert.Equal(t, "task_public", outcome.Task.TaskID)
assert.Equal(t, []string{"reserve", "insert", "settle"}, events)
assert.Zero(t, billing.refunds)
var count int64
require.NoError(t, database.Model(&model.Task{}).Where("task_id = ?", "task_public").Count(&count).Error)
assert.Equal(t, int64(1), count)
assert.False(t, c.Writer.Written())
}
func setupTaskSubmissionDatabase(t *testing.T, migrate bool, events *[]string) *gorm.DB {
t.Helper()
previousDB := model.DB
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
require.NoError(t, database.Callback().Create().Before("gorm:create").Register("test:task-submit-order", func(*gorm.DB) {
*events = append(*events, "insert")
}))
if migrate {
require.NoError(t, database.AutoMigrate(&model.Task{}))
}
model.DB = database
t.Cleanup(func() { model.DB = previousDB })
return database
}
func taskSubmissionTestContext() *gin.Context {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/plugin/submit", strings.NewReader(`{}`))
return c
}
func taskSubmissionRelayInfo(billing relaycommon.BillingSettler) *relaycommon.RelayInfo {
return &relaycommon.RelayInfo{
UserId: 1,
UsingGroup: "default",
OriginModelName: "plugin-model",
Billing: billing,
TaskRelayInfo: &relaycommon.TaskRelayInfo{
PublicTaskID: "task_public",
LockedChannel: &model.Channel{Id: 1, Type: constant.ChannelTypeTaskPlugin, Name: "plugin"},
},
ChannelMeta: &relaycommon.ChannelMeta{ChannelId: 1, ChannelType: constant.ChannelTypeTaskPlugin},
}
}
+435 -45
View File
@@ -1,88 +1,478 @@
package controller
import (
"errors"
"fmt"
"net/http"
"regexp"
"strconv"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/middleware"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay"
relaychannel "github.com/QuantumNous/new-api/relay/channel"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin"
)
func GetAllTask(c *gin.Context) {
pageInfo := common.GetPageQuery(c)
type taskArtifactResponse struct {
Key string `json:"key"`
Type string `json:"type"`
MimeType string `json:"mime_type,omitempty"`
ContentURL string `json:"content_url"`
}
startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64)
endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64)
// 解析其他查询参数
queryParams := model.SyncTaskQueryParams{
Platform: constant.TaskPlatform(c.Query("platform")),
TaskID: c.Query("task_id"),
Status: c.Query("status"),
Action: c.Query("action"),
StartTimestamp: startTimestamp,
EndTimestamp: endTimestamp,
ChannelID: c.Query("channel_id"),
var (
taskArtifactKeyPattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._~-]{0,127}$`)
errTaskArtifactPluginUnavailable = errors.New("task artifact plugin unavailable")
errTaskArtifactPlugin = errors.New("task artifact plugin error")
)
func GetTask(c *gin.Context) {
task, exists, err := model.GetByTaskId(c.GetInt("id"), c.Param("key"))
if err != nil {
videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to query task")
return
}
if !exists {
videoProxyError(c, http.StatusNotFound, "invalid_request_error", "Task not found")
return
}
createdAt := task.CreatedAt
if createdAt == 0 {
createdAt = task.SubmitTime
}
failReason := task.FailReason
if task.Status == model.TaskStatusSuccess && taskFailReasonIsLegacyResultURL(task.FailReason) {
failReason = ""
}
c.JSON(http.StatusOK, gin.H{
"task_id": task.TaskID,
"platform": task.Platform,
"status": task.Status,
"progress": task.Progress,
"fail_reason": failReason,
"created_at": createdAt,
"finished_at": task.FinishTime,
})
}
func GetTaskArtifacts(c *gin.Context) {
task, exists, err := model.GetByTaskId(c.GetInt("id"), c.Param("key"))
if err != nil {
writeTaskArtifactError(c, http.StatusInternalServerError, "artifact_internal_error", "Failed to query task")
return
}
if !exists || task == nil {
writeTaskArtifactError(c, http.StatusNotFound, "artifact_not_found", "Task or artifact not found")
return
}
writeTaskArtifacts(c, task, false)
}
func GetDashboardTaskArtifacts(c *gin.Context) {
task, exists, err := getTaskForArtifactRequest(c, c.Param("task_id"))
if err != nil {
writeTaskArtifactError(c, http.StatusInternalServerError, "artifact_internal_error", "Failed to query task")
return
}
if !exists || task == nil {
writeTaskArtifactError(c, http.StatusNotFound, "artifact_not_found", "Task or artifact not found")
return
}
writeTaskArtifacts(c, task, true)
}
func writeTaskArtifacts(c *gin.Context, task *model.Task, dashboard bool) {
c.Header("Cache-Control", "private, no-store")
artifacts, err := projectTaskArtifacts(task)
if err != nil {
writeTaskArtifactProjectionError(c, err)
return
}
items := make([]taskArtifactResponse, 0, len(artifacts))
for _, artifact := range artifacts {
contentURL, buildErr := service.BuildTaskArtifactContentURL(task.TaskID, artifact.Key)
if buildErr != nil {
writeTaskArtifactError(c, http.StatusInternalServerError, "artifact_url_error", "Failed to build artifact content URL")
return
}
items = append(items, taskArtifactResponse{
Key: artifact.Key,
Type: artifact.Type,
MimeType: artifact.MimeType,
ContentURL: contentURL,
})
}
response := gin.H{"task_id": task.TaskID, "artifacts": items}
if legacyVideoAvailable(task) {
legacyContentURL, buildErr := service.BuildTaskArtifactContentURL(task.TaskID, "video")
if buildErr != nil {
writeTaskArtifactError(c, http.StatusInternalServerError, "artifact_url_error", "Failed to build artifact content URL")
return
}
response["legacy_content_url"] = legacyContentURL
}
if dashboard {
common.ApiSuccess(c, response)
return
}
c.JSON(http.StatusOK, response)
}
func projectTaskArtifacts(task *model.Task) ([]relaychannel.TaskArtifact, error) {
if task == nil || task.Status != model.TaskStatusSuccess || !taskHasPluginExecution(task) {
return []relaychannel.TaskArtifact{}, nil
}
adaptor := relay.GetTaskAdaptor(task.Platform)
if adaptor == nil {
return nil, errTaskArtifactPluginUnavailable
}
provider, ok := adaptor.(relaychannel.TaskArtifactProvider)
if !ok {
return []relaychannel.TaskArtifact{}, nil
}
artifacts, err := provider.ListArtifacts(task)
if err != nil {
return nil, fmt.Errorf("%w: %v", errTaskArtifactPlugin, err)
}
return validateProjectedTaskArtifacts(artifacts)
}
func validateProjectedTaskArtifacts(artifacts []relaychannel.TaskArtifact) ([]relaychannel.TaskArtifact, error) {
if len(artifacts) > 64 {
return nil, fmt.Errorf("%w: too many artifacts", errTaskArtifactPlugin)
}
seen := make(map[string]struct{}, len(artifacts))
for i := range artifacts {
if artifacts[i].Key != strings.TrimSpace(artifacts[i].Key) ||
artifacts[i].Type != strings.TrimSpace(artifacts[i].Type) {
return nil, fmt.Errorf("%w: invalid artifact identity", errTaskArtifactPlugin)
}
if !taskArtifactKeyPattern.MatchString(artifacts[i].Key) {
return nil, fmt.Errorf("%w: invalid artifact key", errTaskArtifactPlugin)
}
if _, exists := seen[artifacts[i].Key]; exists {
return nil, fmt.Errorf("%w: duplicate artifact key", errTaskArtifactPlugin)
}
seen[artifacts[i].Key] = struct{}{}
switch artifacts[i].Type {
case "video", "audio", "image", "file":
default:
return nil, fmt.Errorf("%w: invalid artifact type", errTaskArtifactPlugin)
}
if len(artifacts[i].MimeType) > 255 || strings.ContainsAny(artifacts[i].MimeType, "\r\n") {
return nil, fmt.Errorf("%w: invalid artifact mime type", errTaskArtifactPlugin)
}
}
return artifacts, nil
}
func initTaskArtifactAdaptor(task *model.Task) (relaychannel.TaskAdaptor, error) {
if task == nil || !taskHasPluginExecution(task) {
return nil, errTaskArtifactPluginUnavailable
}
channelModel, err := model.CacheGetChannel(task.ChannelId)
if err != nil {
return nil, fmt.Errorf("%w: channel unavailable", errTaskArtifactPluginUnavailable)
}
adaptor := relay.GetTaskAdaptor(task.Platform)
if adaptor == nil {
return nil, errTaskArtifactPluginUnavailable
}
pluginKey := task.PrivateData.Key
if pluginKey == "" {
pluginKey = channelModel.Key
}
baseURL := channelModel.GetBaseURL()
if baseURL == "" {
baseURL = constant.GetChannelBaseURL(channelModel.Type)
}
adaptor.Init(&relaycommon.RelayInfo{
ChannelMeta: &relaycommon.ChannelMeta{
ChannelType: channelModel.Type,
ChannelBaseUrl: baseURL,
ApiKey: pluginKey,
ChannelSetting: channelModel.GetSetting(),
},
})
return adaptor, nil
}
func taskHasPluginExecution(task *model.Task) bool {
return task != nil &&
task.PrivateData.Execution != nil &&
task.PrivateData.Execution.TaskPlugin != nil &&
strings.TrimSpace(task.PrivateData.Execution.TaskPlugin.Key) != ""
}
func legacyVideoAvailable(task *model.Task) bool {
if task == nil || task.Status != model.TaskStatusSuccess ||
taskHasPluginExecution(task) || task.Platform == constant.TaskPlatformSuno ||
strings.TrimSpace(task.GetResultURL()) == "" {
return false
}
switch constant.NormalizeTaskAction(task.Action) {
case constant.TaskActionImageToVideo,
constant.TaskActionTextToVideo,
constant.TaskActionFirstTailToVideo,
constant.TaskActionReferenceToVideo,
constant.TaskActionRemix:
return true
default:
return false
}
}
func getTaskForArtifactRequest(c *gin.Context, taskID string) (*model.Task, bool, error) {
if middleware.IsTaskArtifactAccess(c) {
task, exists, err := model.GetUniqueByOnlyTaskId(taskID)
if err != nil || !exists || task == nil {
return task, exists, err
}
owner, err := model.GetUserCache(task.UserId)
if err != nil || owner == nil || owner.Status != common.UserStatusEnabled {
return nil, false, err
}
return task, true, nil
}
if c.GetInt("token_id") == 0 && c.GetInt("role") >= common.RoleAdminUser {
return model.GetByOnlyTaskId(taskID)
}
return model.GetByTaskId(c.GetInt("id"), taskID)
}
func writeTaskArtifactProjectionError(c *gin.Context, err error) {
if errors.Is(err, errTaskArtifactPluginUnavailable) {
writeTaskArtifactError(c, http.StatusServiceUnavailable, "artifact_plugin_unavailable", "Artifact preview plugin is unavailable")
return
}
writeTaskArtifactError(c, http.StatusInternalServerError, "artifact_plugin_error", "Artifact preview plugin failed")
}
func writeTaskArtifactError(c *gin.Context, status int, code, message string) {
c.Header("Cache-Control", "private, no-store")
if middleware.IsTaskArtifactAccess(c) {
status = http.StatusNotFound
code = "artifact_not_found"
message = "Task or artifact not found"
}
if strings.HasPrefix(c.Request.URL.Path, "/api/") {
c.JSON(status, gin.H{"success": false, "code": code, "message": message})
return
}
c.JSON(status, gin.H{
"error": gin.H{
"message": message,
"type": code,
"code": code,
},
})
}
func TaskArtifactContent(c *gin.Context) {
task, exists, err := getTaskForArtifactRequest(c, c.Param("key"))
if err != nil {
writeTaskArtifactError(c, http.StatusInternalServerError, "artifact_internal_error", "Failed to query task")
return
}
if !exists || task == nil {
writeTaskArtifactError(c, http.StatusNotFound, "artifact_not_found", "Task or artifact not found")
return
}
artifactKey := strings.TrimSpace(c.Param("artifact_key"))
if !taskArtifactKeyPattern.MatchString(artifactKey) {
writeTaskArtifactError(c, http.StatusNotFound, "artifact_not_found", "Task or artifact not found")
return
}
if task.Status != model.TaskStatusSuccess {
writeTaskArtifactError(c, http.StatusConflict, "artifact_not_ready", "Task artifacts are not ready")
return
}
if !taskHasPluginExecution(task) {
if artifactKey != "video" || !legacyVideoAvailable(task) {
writeTaskArtifactError(c, http.StatusNotFound, "artifact_not_found", "Task or artifact not found")
return
}
descriptor := &relaychannel.TaskContentRequest{
URL: task.GetResultURL(),
Method: c.Request.Method,
Credentialless: true,
}
if err := proxyTaskMedia(c, task, descriptor); err != nil {
writeTaskMediaProxyError(c, err)
}
return
}
artifacts, err := projectTaskArtifacts(task)
if err != nil {
writeTaskArtifactProjectionError(c, err)
return
}
found := false
for _, artifact := range artifacts {
if artifact.Key == artifactKey {
found = true
break
}
}
if !found {
writeTaskArtifactError(c, http.StatusNotFound, "artifact_not_found", "Task or artifact not found")
return
}
artifactStore := service.GetTaskArtifactStore()
if ref, resolveErr := artifactStore.Resolve(task, artifactKey); resolveErr == nil && ref != nil {
_ = artifactStore.Serve(c, task, ref)
return
}
adaptor, err := initTaskArtifactAdaptor(task)
if err != nil {
writeTaskArtifactProjectionError(c, err)
return
}
provider, ok := adaptor.(relaychannel.TaskContentRequestProvider)
if !ok {
writeTaskArtifactError(c, http.StatusServiceUnavailable, "artifact_plugin_unavailable", "Artifact content plugin is unavailable")
return
}
clientRequest := relaychannel.TaskArtifactClientRequest{
Method: c.Request.Method,
Headers: taskArtifactClientHeaders(c.Request.Header),
}
descriptor, err := provider.BuildContentRequest(task, artifactKey, clientRequest)
if err != nil || descriptor == nil {
writeTaskArtifactError(c, http.StatusInternalServerError, "artifact_plugin_error", "Artifact content plugin failed")
return
}
if err := proxyTaskMedia(c, task, descriptor); err != nil {
writeTaskMediaProxyError(c, err)
}
}
func taskArtifactClientHeaders(headers http.Header) map[string]string {
result := make(map[string]string, 4)
for _, name := range []string{"Range", "If-Range", "If-None-Match", "If-Modified-Since"} {
if value := strings.TrimSpace(headers.Get(name)); value != "" {
result[name] = value
}
}
return result
}
/*
The task list handlers below deliberately do not call projectTaskArtifacts.
Artifact projection is confined to the explicit endpoints above.
*/
func GetAllTask(c *gin.Context) {
pageInfo := common.GetPageQuery(c)
startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64)
endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64)
queryParams := model.SyncTaskQueryParams{Platform: constant.TaskPlatform(c.Query("platform")), TaskID: c.Query("task_id"), Status: c.Query("status"), Action: c.Query("action"), StartTimestamp: startTimestamp, EndTimestamp: endTimestamp, ChannelID: c.Query("channel_id")}
items := model.TaskGetAllTasks(pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams)
total := model.TaskCountAllTasks(queryParams)
pageInfo.SetTotal(int(total))
pageInfo.SetItems(tasksToDto(items, true))
pageInfo.SetTotal(int(model.TaskCountAllTasks(queryParams)))
pageInfo.SetItems(tasksToDto(items, true, c.GetInt("role")))
common.ApiSuccess(c, pageInfo)
}
func GetUserTask(c *gin.Context) {
pageInfo := common.GetPageQuery(c)
userId := c.GetInt("id")
userID := c.GetInt("id")
startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64)
endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64)
queryParams := model.SyncTaskQueryParams{
Platform: constant.TaskPlatform(c.Query("platform")),
TaskID: c.Query("task_id"),
Status: c.Query("status"),
Action: c.Query("action"),
StartTimestamp: startTimestamp,
EndTimestamp: endTimestamp,
}
items := model.TaskGetAllUserTask(userId, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams)
total := model.TaskCountAllUserTask(userId, queryParams)
pageInfo.SetTotal(int(total))
pageInfo.SetItems(tasksToDto(items, false))
queryParams := model.SyncTaskQueryParams{Platform: constant.TaskPlatform(c.Query("platform")), TaskID: c.Query("task_id"), Status: c.Query("status"), Action: c.Query("action"), StartTimestamp: startTimestamp, EndTimestamp: endTimestamp}
items := model.TaskGetAllUserTask(userID, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams)
pageInfo.SetTotal(int(model.TaskCountAllUserTask(userID, queryParams)))
pageInfo.SetItems(tasksToDto(items, false, common.RoleCommonUser))
common.ApiSuccess(c, pageInfo)
}
func tasksToDto(tasks []*model.Task, fillUser bool) []*dto.TaskDto {
var userIdMap map[int]*model.UserBase
func tasksToDto(tasks []*model.Task, fillUser bool, viewerRole int) []*dto.TaskDto {
var userIDMap map[int]*model.UserBase
if fillUser {
userIdMap = make(map[int]*model.UserBase)
userIds := types.NewSet[int]()
userIDMap = make(map[int]*model.UserBase)
userIDs := types.NewSet[int]()
for _, task := range tasks {
userIds.Add(task.UserId)
userIDs.Add(task.UserId)
}
for _, userId := range userIds.Items() {
cacheUser, err := model.GetUserCache(userId)
if err == nil {
userIdMap[userId] = cacheUser
for _, userID := range userIDs.Items() {
if cacheUser, err := model.GetUserCache(userID); err == nil {
userIDMap[userID] = cacheUser
}
}
}
result := make([]*dto.TaskDto, len(tasks))
for i, task := range tasks {
if fillUser {
if user, ok := userIdMap[task.UserId]; ok {
if user, ok := userIDMap[task.UserId]; ok {
task.Username = user.Username
}
}
result[i] = relay.TaskModel2Dto(task)
item := relay.TaskModel2Dto(task)
item.LegacyVideoAvailable = legacyVideoAvailable(task)
if task.Status == model.TaskStatusSuccess {
item.ResultURL = ""
if taskFailReasonIsLegacyResultURL(task.FailReason) {
item.FailReason = ""
}
}
if viewerRole >= common.RoleAdminUser {
adminInfo := &dto.TaskAdminInfo{}
if execution := task.PrivateData.Execution; execution != nil {
adminInfo.RequestID = execution.RequestID
adminInfo.RequestPath = execution.RequestPath
if snapshot := execution.TaskPlugin; snapshot != nil {
adminInfo.TaskPlugin = &dto.TaskPluginInfo{
Key: snapshot.Key,
Name: snapshot.Name,
Version: snapshot.Version,
}
if snapshot.Author != nil {
adminInfo.TaskPlugin.Author = &dto.TaskPluginAuthorInfo{
Name: snapshot.Author.Name,
URL: snapshot.Author.URL,
}
}
}
}
if adminInfo.RequestID != "" || adminInfo.RequestPath != "" || adminInfo.TaskPlugin != nil {
item.AdminInfo = adminInfo
}
}
if viewerRole >= common.RoleRootUser {
rootInfo := &dto.TaskRootInfo{
UpstreamTaskID: task.PrivateData.UpstreamTaskID,
NodeName: task.PrivateData.NodeName,
}
if execution := task.PrivateData.Execution; execution != nil {
if snapshot := execution.TaskPlugin; snapshot != nil {
rootInfo.TaskPlugin = &dto.TaskPluginRuntimeInfo{
Key: snapshot.Key,
Version: snapshot.Version,
APIVersion: snapshot.APIVersion,
Generation: snapshot.Generation,
}
}
}
if rootInfo.TaskPlugin != nil || rootInfo.UpstreamTaskID != "" || rootInfo.NodeName != "" {
item.RootInfo = rootInfo
}
}
result[i] = item
}
return result
}
func taskFailReasonIsLegacyResultURL(value string) bool {
value = strings.TrimSpace(value)
return len(value) >= len("https://") && strings.EqualFold(value[:len("https://")], "https://") ||
len(value) >= len("http://") && strings.EqualFold(value[:len("http://")], "http://") ||
len(value) >= len("data:") && strings.EqualFold(value[:len("data:")], "data:")
}
+590
View File
@@ -0,0 +1,590 @@
package controller
import (
"context"
"io"
"net/http"
"net/http/httptest"
"net/url"
"strings"
"testing"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/middleware"
"github.com/QuantumNous/new-api/model"
relaychannel "github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/system_setting"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)
func setupGenericTaskTest(t *testing.T) *model.Task {
t.Helper()
originalDB := model.DB
previousRedisEnabled := common.RedisEnabled
common.RedisEnabled = false
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
require.NoError(t, database.AutoMigrate(&model.Task{}, &model.Channel{}, &model.User{}))
model.DB = database
t.Cleanup(func() {
model.DB = originalDB
common.RedisEnabled = previousRedisEnabled
})
require.NoError(t, database.Create(&model.User{
Id: 7, Username: "artifact-owner", Status: common.UserStatusEnabled,
Role: common.RoleCommonUser, Group: "default",
}).Error)
baseURL := "https://example.com"
require.NoError(t, database.Create(&model.Channel{
Id: 1, Name: "artifact", Key: "key", BaseURL: &baseURL, Status: common.ChannelStatusEnabled,
}).Error)
task := &model.Task{
TaskID: "task_generic", Platform: "document", UserId: 7, ChannelId: 1,
Status: model.TaskStatusSuccess, Progress: "100%", SubmitTime: 10, FinishTime: 20,
}
require.NoError(t, database.Create(task).Error)
return task
}
func allowPrivateTaskMediaTest(t *testing.T) {
t.Helper()
originalFetchSetting := *system_setting.GetFetchSetting()
system_setting.GetFetchSetting().EnableSSRFProtection = true
system_setting.GetFetchSetting().AllowPrivateIp = true
system_setting.GetFetchSetting().AllowedPorts = []string{"1-65535"}
t.Cleanup(func() { *system_setting.GetFetchSetting() = originalFetchSetting })
service.InitHttpClient()
}
func TestGetTaskDoesNotProjectArtifacts(t *testing.T) {
task := setupGenericTaskTest(t)
task.FailReason = "https://stale-upstream.invalid/video.mp4"
task.PrivateData = model.TaskPrivateData{
ResultURL: "https://private-upstream.invalid/video.mp4",
Execution: &model.TaskExecutionSnapshot{
TaskPlugin: &model.TaskPluginSnapshot{Key: "missing-plugin", Name: "Missing"},
},
}
require.NoError(t, model.DB.Save(task).Error)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Set("id", 7)
c.Params = gin.Params{{Key: "key", Value: task.TaskID}}
c.Request = httptest.NewRequest(http.MethodGet, "/v1/tasks/"+task.TaskID, nil)
GetTask(c)
assert.Equal(t, http.StatusOK, recorder.Code)
var response map[string]any
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
assert.Equal(t, task.TaskID, response["task_id"])
assert.NotContains(t, response, "artifacts")
assert.NotContains(t, recorder.Body.String(), "upstream.invalid")
}
func TestGetTaskArtifactsReturnsEmptyForLegacyTask(t *testing.T) {
task := setupGenericTaskTest(t)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Set("id", task.UserId)
c.Params = gin.Params{{Key: "key", Value: task.TaskID}}
c.Request = httptest.NewRequest(http.MethodGet, "/v1/tasks/"+task.TaskID+"/artifacts", nil)
GetTaskArtifacts(c)
assert.Equal(t, http.StatusOK, recorder.Code)
assert.Equal(t, "private, no-store", recorder.Header().Get("Cache-Control"))
var response struct {
TaskID string `json:"task_id"`
Artifacts []taskArtifactResponse `json:"artifacts"`
}
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
assert.Equal(t, task.TaskID, response.TaskID)
assert.Empty(t, response.Artifacts)
}
func TestTaskArtifactAuthorizationKeepsForeignTasksHidden(t *testing.T) {
task := setupGenericTaskTest(t)
commonUser, _ := gin.CreateTestContext(httptest.NewRecorder())
commonUser.Set("id", 8)
commonUser.Set("role", common.RoleCommonUser)
_, exists, err := getTaskForArtifactRequest(commonUser, task.TaskID)
require.NoError(t, err)
assert.False(t, exists)
admin, _ := gin.CreateTestContext(httptest.NewRecorder())
admin.Set("id", 8)
admin.Set("role", common.RoleAdminUser)
found, exists, err := getTaskForArtifactRequest(admin, task.TaskID)
require.NoError(t, err)
require.True(t, exists)
assert.Equal(t, task.TaskID, found.TaskID)
apiToken, _ := gin.CreateTestContext(httptest.NewRecorder())
apiToken.Set("id", 8)
apiToken.Set("role", common.RoleRootUser)
apiToken.Set("token_id", 99)
_, exists, err = getTaskForArtifactRequest(apiToken, task.TaskID)
require.NoError(t, err)
assert.False(t, exists)
}
func TestDashboardTaskArtifactsReturnsLegacyCapabilityWithoutUpstreamURL(t *testing.T) {
task := setupGenericTaskTest(t)
previousSecret := common.CryptoSecret
previousPublicAddress := system_setting.TaskPublicAddress
common.CryptoSecret = "controller-task-artifact-access-secret"
system_setting.TaskPublicAddress = "https://gateway.example/prefix"
t.Cleanup(func() {
common.CryptoSecret = previousSecret
system_setting.TaskPublicAddress = previousPublicAddress
})
task.Action = constant.TaskActionTextToVideo
task.FailReason = "https://upstream.invalid/private-video.mp4?signature=secret"
require.NoError(t, model.DB.Save(task).Error)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Set("id", task.UserId)
c.Set("role", common.RoleCommonUser)
c.Params = gin.Params{{Key: "task_id", Value: task.TaskID}}
c.Request = httptest.NewRequest(http.MethodGet, "/api/task/"+task.TaskID+"/artifacts", nil)
GetDashboardTaskArtifacts(c)
assert.Equal(t, http.StatusOK, recorder.Code)
assert.Equal(t, "private, no-store", recorder.Header().Get("Cache-Control"))
var response struct {
Success bool `json:"success"`
Data struct {
Artifacts []taskArtifactResponse `json:"artifacts"`
LegacyContentURL string `json:"legacy_content_url"`
} `json:"data"`
}
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
assert.True(t, response.Success)
assert.Empty(t, response.Data.Artifacts)
contentURL, err := url.Parse(response.Data.LegacyContentURL)
require.NoError(t, err)
assert.Equal(t, "/prefix/v1/tasks/"+task.TaskID+"/artifacts/video/content", contentURL.Path)
assert.True(t, service.VerifyTaskArtifactAccess(
contentURL.Query().Get(service.TaskArtifactAccessQueryParameter),
task.TaskID,
"video",
))
assert.NotContains(t, recorder.Body.String(), "upstream.invalid")
assert.NotContains(t, recorder.Body.String(), "signature=secret")
}
func TestTaskArtifactAccessRequiresActiveOwner(t *testing.T) {
task := setupGenericTaskTest(t)
task.Action = constant.TaskActionTextToVideo
task.FailReason = "https://upstream.invalid/private-video.mp4"
require.NoError(t, model.DB.Save(task).Error)
require.NoError(t, model.DB.Model(&model.User{}).
Where("id = ?", task.UserId).
Update("status", common.UserStatusDisabled).Error)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Set(middleware.TaskArtifactAccessContextKey, true)
c.Params = gin.Params{
{Key: "key", Value: task.TaskID},
{Key: "artifact_key", Value: "video"},
}
c.Request = httptest.NewRequest(
http.MethodGet,
"/v1/tasks/"+task.TaskID+"/artifacts/video/content",
nil,
)
TaskArtifactContent(c)
assert.Equal(t, http.StatusNotFound, recorder.Code)
assert.Equal(t, "private, no-store", recorder.Header().Get("Cache-Control"))
}
func TestTaskArtifactAccessRejectsAmbiguousHistoricalTaskID(t *testing.T) {
task := setupGenericTaskTest(t)
task.Action = constant.TaskActionTextToVideo
task.FailReason = "https://first-upstream.invalid/video.mp4"
require.NoError(t, model.DB.Save(task).Error)
require.NoError(t, model.DB.Create(&model.User{
Id: 8, Username: "other-artifact-owner", Status: common.UserStatusEnabled,
Role: common.RoleCommonUser, Group: "default", AffCode: "artifact-owner-8",
}).Error)
require.NoError(t, model.DB.Create(&model.Task{
TaskID: task.TaskID, Platform: task.Platform, UserId: 8, ChannelId: task.ChannelId,
Action: constant.TaskActionTextToVideo, Status: model.TaskStatusSuccess,
FailReason: "https://second-upstream.invalid/video.mp4",
}).Error)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Set(middleware.TaskArtifactAccessContextKey, true)
c.Params = gin.Params{
{Key: "key", Value: task.TaskID},
{Key: "artifact_key", Value: "video"},
}
c.Request = httptest.NewRequest(
http.MethodGet,
"/v1/tasks/"+task.TaskID+"/artifacts/video/content",
nil,
)
TaskArtifactContent(c)
assert.Equal(t, http.StatusNotFound, recorder.Code)
assert.NotContains(t, recorder.Body.String(), "upstream.invalid")
}
func TestLegacyVideoArtifactContentUsesGetResultURL(t *testing.T) {
task := setupGenericTaskTest(t)
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, "bytes=0-3", r.Header.Get("Range"))
w.Header().Set("Content-Type", "video/mp4")
w.Header().Set("Content-Range", "bytes 0-3/4")
w.WriteHeader(http.StatusPartialContent)
_, _ = w.Write([]byte("data"))
}))
defer upstream.Close()
allowPrivateTaskMediaTest(t)
task.Action = constant.TaskActionTextToVideo
task.PrivateData.ResultURL = upstream.URL
task.FailReason = "https://stale.invalid/legacy-fallback.mp4"
require.NoError(t, model.DB.Save(task).Error)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Set(middleware.TaskArtifactAccessContextKey, true)
c.Params = gin.Params{
{Key: "key", Value: task.TaskID},
{Key: "artifact_key", Value: "video"},
}
c.Request = httptest.NewRequest(
http.MethodGet,
"/v1/tasks/"+task.TaskID+"/artifacts/video/content",
nil,
)
c.Request.Header.Set("Range", "bytes=0-3")
TaskArtifactContent(c)
assert.Equal(t, http.StatusPartialContent, recorder.Code)
assert.Equal(t, "data", recorder.Body.String())
assert.Equal(t, "bytes 0-3/4", recorder.Header().Get("Content-Range"))
}
func TestDisabledArtifactStorePreservesPluginUpstreamContent(t *testing.T) {
task := setupGenericTaskTest(t)
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, "provider-key", r.Header.Get("x-goog-api-key"))
assert.Equal(t, "bytes=0-13", r.Header.Get("Range"))
w.Header().Set("Content-Type", "video/mp4")
w.Header().Set("Content-Range", "bytes 0-13/14")
w.WriteHeader(http.StatusPartialContent)
_, _ = w.Write([]byte("artifact-bytes"))
}))
defer upstream.Close()
allowPrivateTaskMediaTest(t)
previousMemoryCache := common.MemoryCacheEnabled
common.MemoryCacheEnabled = false
t.Cleanup(func() { common.MemoryCacheEnabled = previousMemoryCache })
require.NoError(t, model.DB.Model(&model.Channel{}).Where("id = ?", task.ChannelId).Updates(map[string]any{
"type": constant.ChannelTypeGemini,
"key": "provider-key",
"base_url": upstream.URL,
}).Error)
task.Platform = constant.TaskPlatform("google")
task.PrivateData.Execution = &model.TaskExecutionSnapshot{TaskPlugin: &model.TaskPluginSnapshot{
Key: "google", Name: "Google Veo (Gemini API)", Version: "1.0.0", APIVersion: 1,
}}
task.SetData(map[string]any{"response": map[string]any{
"generateVideoResponse": map[string]any{
"generatedVideos": []any{map[string]any{"video": map[string]any{"uri": upstream.URL}}},
},
}})
require.NoError(t, model.DB.Save(task).Error)
require.False(t, service.GetTaskArtifactStore().Enabled())
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Set(middleware.TaskArtifactAccessContextKey, true)
c.Params = gin.Params{
{Key: "key", Value: task.TaskID},
{Key: "artifact_key", Value: "video"},
}
c.Request = httptest.NewRequest(http.MethodGet, "/v1/tasks/"+task.TaskID+"/artifacts/video/content", nil)
c.Request.Header.Set("Range", "bytes=0-13")
TaskArtifactContent(c)
assert.Equal(t, http.StatusPartialContent, recorder.Code)
assert.Equal(t, "artifact-bytes", recorder.Body.String())
assert.Equal(t, "video/mp4", recorder.Header().Get("Content-Type"))
assert.Equal(t, "bytes 0-13/14", recorder.Header().Get("Content-Range"))
}
func TestProjectedTaskArtifactValidationRejectsAmbiguousIdentity(t *testing.T) {
validated, err := validateProjectedTaskArtifacts([]relaychannel.TaskArtifact{
{Key: "video-main", Type: "video", MimeType: "video/mp4"},
{Key: "cover.main", Type: "image", MimeType: "image/png"},
})
require.NoError(t, err)
require.Len(t, validated, 2)
assert.Equal(t, "video-main", validated[0].Key)
for _, artifacts := range [][]relaychannel.TaskArtifact{
{{Key: "../video", Type: "video"}},
{{Key: "video/0", Type: "video"}},
{{Key: "video-main", Type: "video"}, {Key: "video-main", Type: "image"}},
{{Key: "video-main", Type: "unknown"}},
{{Key: "video-main", Type: "video", MimeType: "video/mp4\r\nX-Test: injected"}},
} {
_, err := validateProjectedTaskArtifacts(artifacts)
assert.ErrorIs(t, err, errTaskArtifactPlugin)
}
}
func TestProxyTaskMediaForwardsRangeAndFiltersResponseHeaders(t *testing.T) {
task := setupGenericTaskTest(t)
var receivedRange, receivedAuthorization string
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
receivedRange = r.Header.Get("Range")
receivedAuthorization = r.Header.Get("Authorization")
w.Header().Set("Content-Type", "video/mp4")
w.Header().Set("Content-Range", "bytes 0-3/10")
w.Header().Set("Accept-Ranges", "bytes")
w.Header().Set("Set-Cookie", "provider=secret")
w.Header().Set("WWW-Authenticate", "Bearer provider")
w.Header().Set("X-Provider-Secret", "hidden")
w.Header().Set("Cache-Control", "public, max-age=86400")
w.WriteHeader(http.StatusPartialContent)
_, _ = w.Write([]byte("data"))
}))
defer upstream.Close()
allowPrivateTaskMediaTest(t)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodGet, "/v1/tasks/task_generic/artifacts/video-main/content", nil)
c.Request.Header.Set("Range", "bytes=0-3")
err := proxyTaskMedia(c, task, &relaychannel.TaskContentRequest{
URL: upstream.URL,
Method: http.MethodGet,
Headers: map[string]string{"Authorization": "Bearer provider-secret"},
})
require.NoError(t, err)
assert.Equal(t, http.StatusPartialContent, recorder.Code)
assert.Equal(t, "data", recorder.Body.String())
assert.Equal(t, "bytes=0-3", receivedRange)
assert.Equal(t, "Bearer provider-secret", receivedAuthorization)
assert.Equal(t, "bytes 0-3/10", recorder.Header().Get("Content-Range"))
assert.Equal(t, "bytes", recorder.Header().Get("Accept-Ranges"))
assert.Equal(t, "private, no-store", recorder.Header().Get("Cache-Control"))
assert.Equal(t, "sandbox; default-src 'none'", recorder.Header().Get("Content-Security-Policy"))
assert.Equal(t, "no-referrer", recorder.Header().Get("Referrer-Policy"))
assert.Equal(t, "nosniff", recorder.Header().Get("X-Content-Type-Options"))
assert.Empty(t, recorder.Header().Get("Set-Cookie"))
assert.Empty(t, recorder.Header().Get("WWW-Authenticate"))
assert.Empty(t, recorder.Header().Get("X-Provider-Secret"))
}
func TestProxyTaskMediaPassesThroughUnsatisfiedRange(t *testing.T) {
task := setupGenericTaskTest(t)
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Range", "bytes */10")
w.WriteHeader(http.StatusRequestedRangeNotSatisfiable)
}))
defer upstream.Close()
allowPrivateTaskMediaTest(t)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodGet, "/content", nil)
require.NoError(t, proxyTaskMedia(c, task, &relaychannel.TaskContentRequest{
URL: upstream.URL, Method: http.MethodGet,
}))
assert.Equal(t, http.StatusRequestedRangeNotSatisfiable, recorder.Code)
assert.Equal(t, "bytes */10", recorder.Header().Get("Content-Range"))
assert.Equal(t, "private, no-store", recorder.Header().Get("Cache-Control"))
}
func TestTaskMediaResponseHeaderTimeoutDoesNotTruncateBody(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "video/mp4")
w.WriteHeader(http.StatusOK)
w.(http.Flusher).Flush()
time.Sleep(75 * time.Millisecond)
_, _ = w.Write([]byte("complete-body"))
}))
defer upstream.Close()
request, err := http.NewRequest(http.MethodGet, upstream.URL, nil)
require.NoError(t, err)
response, err := doTaskMediaRequest(upstream.Client(), request, 20*time.Millisecond)
require.NoError(t, err)
defer response.Body.Close()
body, err := io.ReadAll(response.Body)
require.NoError(t, err)
assert.Equal(t, "complete-body", string(body))
}
func TestTaskMediaResponseHeaderTimeoutCancelsBeforeHeaders(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
time.Sleep(75 * time.Millisecond)
_, _ = w.Write([]byte("late"))
}))
defer upstream.Close()
request, err := http.NewRequest(http.MethodGet, upstream.URL, nil)
require.NoError(t, err)
_, err = doTaskMediaRequest(upstream.Client(), request, 10*time.Millisecond)
assert.ErrorIs(t, err, context.DeadlineExceeded)
}
func TestWriteVideoDataURLStreamsAndSupportsHead(t *testing.T) {
const dataURL = "data:video/mp4;base64,Y29tcGxldGUtYm9keQ=="
getRecorder := httptest.NewRecorder()
getContext, _ := gin.CreateTestContext(getRecorder)
getContext.Request = httptest.NewRequest(http.MethodGet, "/content", nil)
require.NoError(t, writeVideoDataURL(getContext, dataURL))
assert.Equal(t, http.StatusOK, getRecorder.Code)
assert.Equal(t, "complete-body", getRecorder.Body.String())
assert.Equal(t, "13", getRecorder.Header().Get("Content-Length"))
headRecorder := httptest.NewRecorder()
headContext, _ := gin.CreateTestContext(headRecorder)
headContext.Request = httptest.NewRequest(http.MethodHead, "/content", nil)
require.NoError(t, writeVideoDataURL(headContext, dataURL))
assert.Equal(t, http.StatusOK, headRecorder.Code)
assert.Empty(t, headRecorder.Body.String())
assert.Equal(t, "13", headRecorder.Header().Get("Content-Length"))
}
func TestWriteVideoDataURLRejectsOversizedPayloadBeforeDecode(t *testing.T) {
previousLimit := taskMediaDataURLMaxEncodedBytes
taskMediaDataURLMaxEncodedBytes = 32
t.Cleanup(func() { taskMediaDataURLMaxEncodedBytes = previousLimit })
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodGet, "/content", nil)
err := writeVideoDataURL(c, "data:video/mp4;base64,"+strings.Repeat("A", 64))
assert.ErrorIs(t, err, errTaskMediaRequestRejected)
assert.Empty(t, recorder.Header().Get("Content-Type"))
}
func TestProxyTaskMediaAllowsOnlyCredentiallessCrossOriginRedirect(t *testing.T) {
task := setupGenericTaskTest(t)
var destinationAuthorization, destinationRange string
destination := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
destinationAuthorization = r.Header.Get("Authorization")
destinationRange = r.Header.Get("Range")
_, _ = w.Write([]byte("redirected"))
}))
defer destination.Close()
source := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
http.Redirect(w, r, destination.URL, http.StatusFound)
}))
defer source.Close()
allowPrivateTaskMediaTest(t)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodGet, "/content", nil)
c.Request.Header.Set("Range", "bytes=0-3")
err := proxyTaskMedia(c, task, &relaychannel.TaskContentRequest{
URL: source.URL, Method: http.MethodGet, Credentialless: true,
})
require.NoError(t, err)
assert.Equal(t, http.StatusOK, recorder.Code)
assert.Equal(t, "redirected", recorder.Body.String())
assert.Empty(t, destinationAuthorization)
assert.Equal(t, "bytes=0-3", destinationRange)
destinationRange = ""
rejectedRecorder := httptest.NewRecorder()
rejectedContext, _ := gin.CreateTestContext(rejectedRecorder)
rejectedContext.Request = httptest.NewRequest(http.MethodGet, "/content", nil)
err = proxyTaskMedia(rejectedContext, task, &relaychannel.TaskContentRequest{
URL: source.URL, Method: http.MethodGet,
Headers: map[string]string{"Authorization": "Bearer provider-secret"},
})
var proxyErr *taskMediaProxyError
require.ErrorAs(t, err, &proxyErr)
assert.Equal(t, "artifact_request_rejected", proxyErr.code)
assert.Empty(t, destinationRange)
}
func TestTaskMediaRequestHeaderPolicy(t *testing.T) {
header := http.Header{}
require.NoError(t, applyTaskMediaRequestHeaders(header, map[string]string{
"Authorization": "Bearer provider-secret",
"X-Signature": "signed",
}))
assert.Equal(t, "Bearer provider-secret", header.Get("Authorization"))
assert.Equal(t, "signed", header.Get("X-Signature"))
for _, name := range []string{"Host", "Content-Length", "Accept-Encoding", "Connection", "Proxy-Authorization", "Transfer-Encoding"} {
t.Run(name, func(t *testing.T) {
assert.ErrorIs(t, applyTaskMediaRequestHeaders(http.Header{}, map[string]string{name: "bad"}), errTaskMediaRequestRejected)
})
}
assert.ErrorIs(t, applyTaskMediaRequestHeaders(http.Header{}, map[string]string{"X-Test": "bad\r\ninjected"}), errTaskMediaRequestRejected)
}
func TestCredentiallessTaskMediaDescriptorRejectsCredentialsAndBody(t *testing.T) {
task := setupGenericTaskTest(t)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodGet, "/content", nil)
for _, descriptor := range []*relaychannel.TaskContentRequest{
{URL: "https://example.com/video", Method: http.MethodPost, Credentialless: true},
{URL: "https://example.com/video", Method: http.MethodGet, Body: []byte("secret"), Credentialless: true},
{URL: "https://example.com/video", Method: http.MethodGet, Headers: map[string]string{"X-Key": "secret"}, Credentialless: true},
} {
err := proxyTaskMedia(c, task, descriptor)
var proxyErr *taskMediaProxyError
require.ErrorAs(t, err, &proxyErr)
assert.Equal(t, "artifact_request_rejected", proxyErr.code)
}
}
func TestSelfTaskMediaURLGuard(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodGet, "https://gateway.example/v1/videos/task-1/content", nil)
c.Request.Host = "gateway.example"
selfURL, err := url.Parse("https://gateway.example/v1/videos/task-1/content")
require.NoError(t, err)
assert.True(t, isSelfTaskMediaURL(c, selfURL))
remoteURL, err := url.Parse("https://cdn.example/v1/videos/task-1/content")
require.NoError(t, err)
assert.False(t, isSelfTaskMediaURL(c, remoteURL))
assert.True(t, isTaskMediaFallbackLoop(remoteURL.String(), "task-1"))
assert.False(t, isTaskMediaFallbackLoop(remoteURL.String(), "task-2"))
}
+137
View File
@@ -0,0 +1,137 @@
package controller
import (
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestTaskLogDTOSeparatesUserAdminAndRootDetails(t *testing.T) {
task := &model.Task{
TaskID: "task_public",
Platform: "document-parser",
PrivateData: model.TaskPrivateData{
Key: "channel-secret-canary",
UpstreamTaskID: "upstream-private",
NodeName: "node-a",
Execution: &model.TaskExecutionSnapshot{
RequestID: "request-public",
RequestPath: "/v1/documents",
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,
},
},
},
}
userView := tasksToDto([]*model.Task{task}, false, common.RoleCommonUser)[0]
assert.Nil(t, userView.AdminInfo)
assert.Nil(t, userView.RootInfo)
adminView := tasksToDto([]*model.Task{task}, false, common.RoleAdminUser)[0]
require.NotNil(t, adminView.AdminInfo)
require.NotNil(t, adminView.AdminInfo.TaskPlugin)
assert.Equal(t, "document-parser", adminView.AdminInfo.TaskPlugin.Key)
assert.Equal(t, "Document Parser", adminView.AdminInfo.TaskPlugin.Name)
assert.Equal(t, "1.2.3", adminView.AdminInfo.TaskPlugin.Version)
require.NotNil(t, adminView.AdminInfo.TaskPlugin.Author)
assert.Equal(t, "Community Author", adminView.AdminInfo.TaskPlugin.Author.Name)
assert.Equal(t, "https://plugins.example/author", adminView.AdminInfo.TaskPlugin.Author.URL)
assert.Equal(t, "request-public", adminView.AdminInfo.RequestID)
assert.Equal(t, "/v1/documents", adminView.AdminInfo.RequestPath)
assert.Nil(t, adminView.RootInfo)
rootView := tasksToDto([]*model.Task{task}, false, common.RoleRootUser)[0]
require.NotNil(t, rootView.AdminInfo)
require.NotNil(t, rootView.RootInfo)
require.NotNil(t, rootView.RootInfo.TaskPlugin)
assert.Equal(t, 1, rootView.RootInfo.TaskPlugin.APIVersion)
assert.Equal(t, uint64(42), rootView.RootInfo.TaskPlugin.Generation)
assert.Equal(t, "upstream-private", rootView.RootInfo.UpstreamTaskID)
assert.Equal(t, "node-a", rootView.RootInfo.NodeName)
adminJSON, err := common.Marshal(adminView)
require.NoError(t, err)
assert.NotContains(t, string(adminJSON), "channel-secret-canary")
assert.NotContains(t, string(adminJSON), "upstream-private")
rootJSON, err := common.Marshal(rootView)
require.NoError(t, err)
assert.NotContains(t, string(rootJSON), "channel-secret-canary")
assert.Contains(t, string(rootJSON), "upstream-private")
}
func TestTaskLogDTODoesNotInventHistoricalPluginProvenance(t *testing.T) {
task := &model.Task{
TaskID: "task_without_snapshot",
Platform: "document-parser",
}
adminView := tasksToDto([]*model.Task{task}, false, common.RoleAdminUser)[0]
assert.Nil(t, adminView.AdminInfo)
assert.Nil(t, adminView.RootInfo)
}
func TestTaskLogDTOReplacesLegacyVideoURLWithAvailabilityFlag(t *testing.T) {
task := &model.Task{
TaskID: "task_legacy_video",
Platform: "jimeng",
Action: constant.TaskActionTextToVideo,
Status: model.TaskStatusSuccess,
FailReason: "https://private-upstream.invalid/video.mp4?signature=secret",
}
view := tasksToDto([]*model.Task{task}, false, common.RoleCommonUser)[0]
assert.True(t, view.LegacyVideoAvailable)
assert.Empty(t, view.ResultURL)
assert.Empty(t, view.FailReason)
encoded, err := common.Marshal(view)
require.NoError(t, err)
assert.NotContains(t, string(encoded), "private-upstream.invalid")
assert.NotContains(t, string(encoded), "result_url")
assert.Contains(t, string(encoded), "legacy_video_available")
}
func TestTaskLogDTOKeepsFailureReasonAndDoesNotMarkPluginTaskLegacy(t *testing.T) {
failed := &model.Task{
TaskID: "task_failed",
Platform: "jimeng",
Action: constant.TaskActionTextToVideo,
Status: model.TaskStatusFailure,
FailReason: "provider rejected the request",
}
failedView := tasksToDto([]*model.Task{failed}, false, common.RoleCommonUser)[0]
assert.Equal(t, "provider rejected the request", failedView.FailReason)
assert.False(t, failedView.LegacyVideoAvailable)
pluginTask := &model.Task{
TaskID: "task_plugin_video",
Platform: "community-video",
Action: constant.TaskActionTextToVideo,
Status: model.TaskStatusSuccess,
FailReason: "https://stale-upstream.invalid/plugin-video.mp4",
PrivateData: model.TaskPrivateData{
ResultURL: "https://private-upstream.invalid/plugin-video.mp4",
Execution: &model.TaskExecutionSnapshot{
TaskPlugin: &model.TaskPluginSnapshot{Key: "community-video"},
},
},
}
pluginView := tasksToDto([]*model.Task{pluginTask}, false, common.RoleCommonUser)[0]
assert.False(t, pluginView.LegacyVideoAvailable)
assert.Empty(t, pluginView.ResultURL)
assert.Empty(t, pluginView.FailReason)
}
+761
View File
@@ -0,0 +1,761 @@
package controller
import (
"context"
"crypto/sha256"
"encoding/json"
"errors"
"fmt"
"net/url"
"sort"
"strings"
"sync"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/plugins"
"github.com/QuantumNous/new-api/setting"
"github.com/gin-gonic/gin"
"gorm.io/gorm"
)
const maxTaskPluginSourceBytes = 1024 * 1024
type taskPluginUploadRequest struct {
Source string `json:"source" binding:"required"`
Enabled *bool `json:"enabled"`
Remark string `json:"remark"`
Force bool `json:"force"`
SourceSha256 string `json:"sourceSha256"`
}
func UploadTaskPlugin(c *gin.Context) {
var request taskPluginUploadRequest
if err := c.ShouldBindJSON(&request); err != nil {
common.ApiErrorMsg(c, err.Error())
return
}
if len(request.Source) > maxTaskPluginSourceBytes {
common.ApiErrorMsg(c, "plugin source exceeds 1 MiB")
return
}
if expected := strings.TrimSpace(request.SourceSha256); expected != "" {
actual := fmt.Sprintf("%x", sha256.Sum256([]byte(request.Source)))
if !strings.EqualFold(actual, expected) {
common.ApiErrorMsg(c, "plugin source sha256 mismatch")
return
}
}
temporary := jsplugin.NewRegistry()
loaded, err := temporary.Register(request.Source, jsplugin.Options{})
if err != nil {
common.ApiErrorMsg(c, err.Error())
return
}
if err = jsplugin.ValidateV1Meta(loaded.Meta); err != nil {
common.ApiErrorMsg(c, err.Error())
return
}
enabled := true
if request.Enabled != nil {
enabled = *request.Enabled
}
if enabled && !request.Force {
if err = jsplugin.PreflightRoutingConflict(jsplugin.DefaultRegistry.Generation(), loaded); err != nil {
common.ApiErrorMsg(c, err.Error())
return
}
}
plugin := model.TaskPlugin{
Key: loaded.Meta.Key, APIVersion: loaded.Meta.APIVersion, Version: loaded.Meta.Version,
Source: request.Source, SourceHash: fmt.Sprintf("%x", sha256.Sum256([]byte(request.Source))),
Enabled: enabled, Remark: request.Remark,
}
if err = model.SaveTaskPlugin(&plugin); err != nil {
common.ApiError(c, err)
return
}
if err = syncTaskPluginsOnceContext(c.Request.Context()); err != nil {
common.ApiError(c, err)
return
}
common.ApiSuccess(c, taskPluginDetail{Plugin: &plugin, Meta: loaded.Meta, Source: plugin.Source, Layer: "override"})
}
func GetTaskPluginVersions(c *gin.Context) {
plugins, err := model.ListTaskPluginVersions(c.Param("key"))
if err != nil {
common.ApiError(c, err)
return
}
common.ApiSuccess(c, plugins)
}
type taskPluginListItem struct {
Meta jsplugin.Meta `json:"meta"`
Source string `json:"source"`
Enabled bool `json:"enabled"`
Active bool `json:"active"`
SourceHash string `json:"source_hash"`
Remark string `json:"remark"`
RuntimeStatus string `json:"runtime_status"`
RuntimeError string `json:"runtime_error,omitempty"`
FactoryMeta *jsplugin.Meta `json:"factory_meta,omitempty"`
ChannelCount int `json:"channel_count"`
InFlightCount int64 `json:"in_flight_count"`
}
type taskPluginRebuildOutcome struct {
Status string `json:"status"`
AttemptedAt time.Time `json:"attempted_at"`
Generation uint64 `json:"generation"`
DatabaseRevision string `json:"database_revision,omitempty"`
PluginErrorCount int `json:"plugin_error_count"`
Error string `json:"error,omitempty"`
}
type taskPluginRuntimeStatus struct {
CurrentGeneration uint64 `json:"current_generation"`
GenerationPublishedAt time.Time `json:"generation_published_at"`
DatabaseRevision string `json:"database_revision"`
DatabaseError string `json:"database_error,omitempty"`
LastRebuild taskPluginRebuildOutcome `json:"last_rebuild"`
PluginErrors map[string]string `json:"plugin_errors"`
}
func ListTaskPlugins(c *gin.Context) {
databasePlugins, err := model.ListTaskPlugins()
if err != nil {
common.ApiError(c, err)
return
}
snapshot := jsplugin.DefaultRegistry.Snapshot()
factory := make(map[string]jsplugin.Meta, len(snapshot.Factory))
override := make(map[string]jsplugin.Meta, len(snapshot.Override))
for _, meta := range snapshot.Factory {
factory[meta.Key] = meta
}
for _, meta := range snapshot.Override {
override[meta.Key] = meta
}
activeRows := make(map[string]model.TaskPlugin)
keys := make(map[string]struct{}, len(factory)+len(databasePlugins))
for key := range factory {
keys[key] = struct{}{}
}
for _, plugin := range databasePlugins {
keys[plugin.Key] = struct{}{}
if plugin.Active {
activeRows[plugin.Key] = plugin
}
}
runtimeErrors := jsplugin.DefaultRegistry.RoutingErrors()
taskPluginSyncState.Lock()
for key, message := range taskPluginSyncState.errors {
runtimeErrors[key] = message
}
taskPluginSyncState.Unlock()
items := make([]taskPluginListItem, 0, len(keys))
for key := range keys {
factoryMeta, hasFactory := factory[key]
row, hasOverride := activeRows[key]
item := taskPluginListItem{Enabled: true, Active: true, RuntimeStatus: "registered"}
if hasOverride {
item.Source = "override"
if hasFactory {
item.Source = "override_over_factory"
factoryCopy := factoryMeta
item.FactoryMeta = &factoryCopy
}
item.Meta = jsplugin.Meta{Key: row.Key, Version: row.Version, APIVersion: row.APIVersion}
if compiled, compileErr := jsplugin.NewRegistry().Register(row.Source, jsplugin.Options{Key: row.Key, Version: row.Version}); compileErr == nil {
item.Meta = compiled.Meta
}
item.Enabled = row.Enabled
item.Active = row.Active
item.SourceHash = row.SourceHash
item.Remark = row.Remark
if !constant.TaskPluginOverrideEnabled {
item.RuntimeStatus = "disabled_fallback"
} else if message := runtimeErrors[key]; message != "" {
item.RuntimeStatus = "compile_failed"
item.RuntimeError = message
} else if runtimeMeta, ok := override[key]; ok {
item.Meta = runtimeMeta
} else if !row.Enabled {
item.RuntimeStatus = "disabled_fallback"
} else {
item.RuntimeStatus = "not_registered"
}
} else {
item.Source = "factory"
item.Meta = factoryMeta
item.Enabled = !setting.IsTaskPluginFactoryDisabled(key)
source, sourceErr := plugins.Source(key)
if sourceErr == nil {
item.SourceHash = fmt.Sprintf("%x", sha256.Sum256([]byte(source)))
}
if !item.Enabled {
item.RuntimeStatus = "disabled"
} else if message := runtimeErrors[key]; message != "" {
item.RuntimeStatus = "compile_failed"
item.RuntimeError = message
}
}
if !hasFactory {
channels, inFlight, usageErr := model.GetTaskPluginUsage(key)
if usageErr != nil {
common.ApiError(c, usageErr)
return
}
item.ChannelCount = len(channels)
item.InFlightCount = inFlight
}
items = append(items, item)
}
sort.Slice(items, func(i, j int) bool { return items[i].Meta.Key < items[j].Meta.Key })
common.ApiSuccess(c, items)
}
func GetTaskPluginRuntime(c *gin.Context) {
routingStatus := jsplugin.DefaultRegistry.RoutingStatus()
pluginErrors := routingStatus.Errors
taskPluginSyncState.Lock()
for key, message := range taskPluginSyncState.errors {
pluginErrors[key] = message
}
lastRebuild := taskPluginSyncState.lastRebuild
lastDatabaseRevision := lastRebuild.DatabaseRevision
taskPluginSyncState.Unlock()
registryRebuild := routingStatus.LastRebuild
if lastRebuild.AttemptedAt.Before(registryRebuild.AttemptedAt) {
lastRebuild = taskPluginRebuildOutcome{
Status: registryRebuild.Status,
AttemptedAt: registryRebuild.AttemptedAt,
Generation: registryRebuild.Generation,
Error: registryRebuild.Error,
}
}
if lastRebuild.Status == "" {
lastRebuild.Status = "never"
}
lastRebuild.PluginErrorCount = len(pluginErrors)
if lastRebuild.Status == "success" && len(pluginErrors) > 0 {
lastRebuild.Status = "partial"
}
status := taskPluginRuntimeStatus{
DatabaseRevision: lastDatabaseRevision,
LastRebuild: lastRebuild,
PluginErrors: pluginErrors,
}
databaseSnapshot, err := model.GetTaskPluginSyncSnapshot()
if err != nil {
status.DatabaseError = "database snapshot unavailable"
} else {
status.DatabaseRevision = databaseSnapshot.Revision
}
if routingStatus.Generation != nil {
status.CurrentGeneration = routingStatus.Generation.Number
status.GenerationPublishedAt = routingStatus.Generation.PublishedAt
}
common.ApiSuccess(c, status)
}
type taskPluginDetail struct {
Plugin *model.TaskPlugin `json:"plugin,omitempty"`
Meta jsplugin.Meta `json:"meta"`
Source string `json:"source"`
Layer string `json:"layer"`
}
func GetTaskPlugin(c *gin.Context) {
key := c.Param("key")
version := c.Query("version")
plugin, err := model.GetTaskPluginVersion(key, version)
if err == nil {
loaded, compileErr := jsplugin.NewRegistry().Register(plugin.Source, jsplugin.Options{Key: plugin.Key, Version: plugin.Version})
if compileErr != nil {
common.ApiErrorMsg(c, compileErr.Error())
return
}
common.ApiSuccess(c, taskPluginDetail{Plugin: plugin, Meta: loaded.Meta, Source: plugin.Source, Layer: "override"})
return
}
if !errors.Is(err, gorm.ErrRecordNotFound) || version != "" {
common.ApiError(c, err)
return
}
source, err := plugins.Source(key)
if err != nil {
common.ApiErrorMsg(c, "task plugin not found")
return
}
loaded, err := jsplugin.NewRegistry().RegisterFactory(source, jsplugin.Options{Key: key})
if err != nil {
common.ApiError(c, err)
return
}
common.ApiSuccess(c, taskPluginDetail{Meta: loaded.Meta, Source: source, Layer: "factory"})
}
type taskPluginDryRunRequest struct {
Hook string `json:"hook" binding:"required"`
Member string `json:"member"`
Args []json.RawMessage `json:"args"`
}
func DryRunTaskPlugin(c *gin.Context) {
var request taskPluginDryRunRequest
if err := c.ShouldBindJSON(&request); err != nil {
common.ApiErrorMsg(c, err.Error())
return
}
detailSource := ""
plugin, err := model.GetTaskPluginVersion(c.Param("key"), "")
if err == nil {
detailSource = plugin.Source
} else if errors.Is(err, gorm.ErrRecordNotFound) {
detailSource, err = plugins.Source(c.Param("key"))
}
if err != nil {
common.ApiErrorMsg(c, "task plugin not found")
return
}
loaded, err := jsplugin.NewRegistry().Register(detailSource, jsplugin.Options{Key: c.Param("key")})
if err != nil {
common.ApiErrorMsg(c, err.Error())
return
}
args := make([]any, len(request.Args))
for index, raw := range request.Args {
if err = common.Unmarshal(raw, &args[index]); err != nil {
common.ApiErrorMsg(c, fmt.Sprintf("invalid argument %d: %v", index+1, err))
return
}
}
var output any
if request.Member == "" {
output, err = loaded.Engine.Call(context.Background(), request.Hook, args...)
} else {
output, err = loaded.Engine.CallMember(context.Background(), request.Hook, request.Member, args...)
}
if err != nil {
common.ApiErrorMsg(c, err.Error())
return
}
common.ApiSuccess(c, output)
}
func DeleteTaskPluginVersion(c *gin.Context) {
key := c.Param("key")
version := c.Param("version")
plugin, lookupErr := model.GetTaskPluginVersion(key, version)
if lookupErr != nil {
if errors.Is(lookupErr, gorm.ErrRecordNotFound) {
common.ApiErrorMsg(c, "override plugin version not found; factory plugins cannot be deleted")
return
}
common.ApiError(c, lookupErr)
return
}
if plugin.Active && !taskPluginHasFactory(key) {
channels, inFlight, usageErr := model.GetTaskPluginUsage(key)
if usageErr != nil {
common.ApiError(c, usageErr)
return
}
if (len(channels) > 0 || inFlight > 0) && c.Query("force") != "true" {
c.JSON(200, gin.H{"success": false, "message": "task plugin is still in use", "data": gin.H{"channels": channels, "in_flight_count": inFlight}})
return
}
}
_, err := model.DeleteTaskPluginVersion(key, version)
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
common.ApiErrorMsg(c, "override plugin version not found; factory plugins cannot be deleted")
return
}
common.ApiError(c, err)
return
}
if err = syncTaskPluginsOnceContext(c.Request.Context()); err != nil {
common.ApiError(c, err)
return
}
common.ApiSuccess(c, nil)
}
type taskPluginActivateRequest struct {
Version string `json:"version" binding:"required"`
}
func ActivateTaskPlugin(c *gin.Context) {
var request taskPluginActivateRequest
if err := c.ShouldBindJSON(&request); err != nil {
common.ApiErrorMsg(c, err.Error())
return
}
versions, err := model.ListTaskPluginVersions(c.Param("key"))
if err != nil {
common.ApiError(c, err)
return
}
var target *model.TaskPlugin
for i := range versions {
if versions[i].Version == request.Version {
target = &versions[i]
break
}
}
if target == nil {
common.ApiErrorMsg(c, "plugin version not found")
return
}
if _, err = jsplugin.NewRegistry().Register(target.Source, jsplugin.Options{Key: target.Key, Version: target.Version}); err != nil {
common.ApiErrorMsg(c, err.Error())
return
}
if err = model.ActivateTaskPlugin(target.Key, target.Version); err != nil {
common.ApiError(c, err)
return
}
if err = syncTaskPluginsOnceContext(c.Request.Context()); err != nil {
common.ApiError(c, err)
return
}
common.ApiSuccess(c, nil)
}
type taskPluginStatusRequest struct {
Enabled *bool `json:"enabled" binding:"required"`
}
func SetTaskPluginStatus(c *gin.Context) {
var request taskPluginStatusRequest
if err := c.ShouldBindJSON(&request); err != nil || request.Enabled == nil {
common.ApiErrorMsg(c, "enabled is required")
return
}
key := c.Param("key")
disabledChannels := 0
if !*request.Enabled {
channels, inFlight, usageErr := model.GetTaskPluginUsage(key)
if usageErr != nil {
common.ApiError(c, usageErr)
return
}
cascade := c.Query("cascade") == "true"
force := c.Query("force") == "true"
if (len(channels) > 0 && !cascade) || (inFlight > 0 && !force) {
c.JSON(200, gin.H{"success": false, "message": "task plugin is still in use", "data": gin.H{"channels": channels, "in_flight_count": inFlight}})
return
}
if cascade {
for _, channel := range channels {
if model.UpdateChannelStatus(channel.Id, "", common.ChannelStatusManuallyDisabled, "task plugin disabled") {
disabledChannels++
}
}
}
}
_, lookupErr := model.GetTaskPluginVersion(key, "")
hasActiveOverride := lookupErr == nil
if lookupErr != nil && !errors.Is(lookupErr, gorm.ErrRecordNotFound) {
common.ApiError(c, lookupErr)
return
}
// The disabled set suppresses only the factory fallback layer. An enabled
// override for the same key keeps serving and is toggled independently.
if taskPluginHasFactory(key) && !hasActiveOverride {
keys := setting.GetTaskPluginDisabledFactoryKeys()
if *request.Enabled {
next := make([]string, 0, len(keys))
for _, item := range keys {
if item != key {
next = append(next, item)
}
}
keys = next
} else {
keys = append(append([]string{}, keys...), key)
}
if err := setting.SetTaskPluginDisabledFactoryKeysOption(keys); err != nil {
common.ApiError(c, err)
return
}
encoded, err := common.Marshal(setting.GetTaskPluginDisabledFactoryKeys())
if err != nil {
common.ApiError(c, err)
return
}
if err = model.UpdateOption(setting.TaskPluginDisabledFactoryKeysKey, string(encoded)); err != nil {
common.ApiError(c, err)
return
}
common.ApiSuccess(c, gin.H{"plugin_enabled": *request.Enabled, "disabled_channels": disabledChannels})
return
}
if err := model.SetTaskPluginEnabled(key, *request.Enabled); err != nil {
common.ApiError(c, err)
return
}
if err := syncTaskPluginsOnceContext(c.Request.Context()); err != nil {
common.ApiError(c, err)
return
}
common.ApiSuccess(c, gin.H{"plugin_enabled": *request.Enabled, "disabled_channels": disabledChannels})
}
func taskPluginHasFactory(key string) bool {
for _, meta := range jsplugin.DefaultRegistry.Snapshot().Factory {
if meta.Key == key {
return true
}
}
return false
}
func GetTaskPluginMarketplaceSources(c *gin.Context) {
common.ApiSuccess(c, setting.GetTaskPluginMarketplaceSources())
}
func UpdateTaskPluginMarketplaceSources(c *gin.Context) {
var sources []setting.TaskPluginMarketplaceSource
if err := c.ShouldBindJSON(&sources); err != nil {
common.ApiErrorMsg(c, err.Error())
return
}
if sources == nil {
sources = []setting.TaskPluginMarketplaceSource{}
}
for i := range sources {
name := strings.TrimSpace(sources[i].Name)
indexURL := strings.TrimSpace(sources[i].IndexURL)
if name == "" {
common.ApiErrorMsg(c, "marketplace source name is required")
return
}
parsed, err := url.Parse(indexURL)
if err != nil || !parsed.IsAbs() || parsed.Host == "" || (!strings.EqualFold(parsed.Scheme, "http") && !strings.EqualFold(parsed.Scheme, "https")) {
common.ApiErrorMsg(c, "marketplace source index_url must be an absolute http(s) URL")
return
}
sources[i].Name = name
sources[i].IndexURL = indexURL
}
encoded, err := common.Marshal(sources)
if err != nil {
common.ApiError(c, err)
return
}
if err = model.UpdateOption(setting.TaskPluginMarketplaceSourcesKey, string(encoded)); err != nil {
common.ApiError(c, err)
return
}
common.ApiSuccess(c, sources)
}
func GetTaskPluginOptions(c *gin.Context) {
snapshot := jsplugin.DefaultRegistry.Snapshot()
seen := make(map[string]bool)
options := make([]gin.H, 0, len(snapshot.Factory)+len(snapshot.Override))
for layer, metas := range [][]jsplugin.Meta{snapshot.Override, snapshot.Factory} {
for _, meta := range metas {
if seen[meta.Key] {
continue
}
// Disabled factory keys are omitted from bind options. The disabled
// set suppresses only the factory fallback; an enabled override for
// the same key is listed in the override pass and still appears.
if layer == 1 && setting.IsTaskPluginFactoryDisabled(meta.Key) {
continue
}
if _, ok := jsplugin.DefaultRegistry.Get(meta.Key); !ok {
continue
}
seen[meta.Key] = true
options = append(options, gin.H{
"key": meta.Key,
"name": meta.Name,
"models": meta.Models,
"usageSchema": meta.UsageSchema,
})
}
}
sort.Slice(options, func(i, j int) bool { return options[i]["key"].(string) < options[j]["key"].(string) })
common.ApiSuccess(c, options)
}
var taskPluginSyncState = struct {
sync.Mutex
hashes map[string]string
errors map[string]string
lastRebuild taskPluginRebuildOutcome
}{hashes: map[string]string{}, errors: map[string]string{}}
func syncTaskPluginsOnce() error {
return syncTaskPluginsOnceContext(context.Background())
}
func syncTaskPluginsOnceContext(ctx context.Context) error {
started := time.Now()
taskPluginSyncState.Lock()
defer taskPluginSyncState.Unlock()
databaseSnapshot, err := model.GetTaskPluginSyncSnapshot()
if err != nil {
syncErr := fmt.Errorf("sync task plugins: %w", err)
taskPluginSyncState.lastRebuild = taskPluginRebuildOutcome{
Status: "failed",
AttemptedAt: time.Now(),
Generation: jsplugin.DefaultRegistry.Generation().Number,
DatabaseRevision: taskPluginSyncState.lastRebuild.DatabaseRevision,
Error: syncErr.Error(),
}
logger.LogDebug(
ctx,
"task_plugin subsystem=sync event=failed stage=database_snapshot retained_generation=%d elapsed_ms=%d",
jsplugin.DefaultRegistry.Generation().Number,
time.Since(started).Milliseconds(),
)
return syncErr
}
databasePlugins := databaseSnapshot.Plugins
sort.Slice(databasePlugins, func(i, j int) bool { return databasePlugins[i].Key < databasePlugins[j].Key })
currentOverrides := jsplugin.DefaultRegistry.OverridePlugins()
generationBefore := jsplugin.DefaultRegistry.Generation().Number
logger.LogDebug(
ctx,
"task_plugin subsystem=sync event=start database_revision=%q generation=%d desired_plugins=%d current_overrides=%d",
databaseSnapshot.Revision,
generationBefore,
len(databasePlugins),
len(currentOverrides),
)
nextOverrides := make([]*jsplugin.LoadedPlugin, 0, len(databasePlugins))
nextHashes := make(map[string]string, len(databasePlugins))
seen := make(map[string]bool, len(databasePlugins))
for _, plugin := range databasePlugins {
seen[plugin.Key] = true
if current := currentOverrides[plugin.Key]; current != nil && taskPluginSyncState.hashes[plugin.Key] == plugin.SourceHash {
nextOverrides = append(nextOverrides, current)
nextHashes[plugin.Key] = plugin.SourceHash
logger.LogDebug(
ctx,
"task_plugin subsystem=sync event=plugin plugin=%q version=%q action=reuse",
plugin.Key,
plugin.Version,
)
continue
}
logger.LogDebug(
ctx,
"task_plugin subsystem=sync event=plugin plugin=%q version=%q action=compile_start",
plugin.Key,
plugin.Version,
)
compiled, compileErr := jsplugin.CompilePlugin(plugin.Source, jsplugin.Options{Key: plugin.Key, Version: plugin.Version})
if compileErr != nil {
retainedIncumbent := false
if current := currentOverrides[plugin.Key]; current != nil {
nextOverrides = append(nextOverrides, current)
retainedIncumbent = true
if currentHash := taskPluginSyncState.hashes[plugin.Key]; currentHash != "" {
nextHashes[plugin.Key] = currentHash
}
}
taskPluginSyncState.errors[plugin.Key] = compileErr.Error()
common.SysError(fmt.Sprintf("compile task plugin %s@%s: %v", plugin.Key, plugin.Version, compileErr))
logger.LogDebug(
ctx,
"task_plugin subsystem=sync event=plugin plugin=%q version=%q action=compile_failed retained_incumbent=%t",
plugin.Key,
plugin.Version,
retainedIncumbent,
)
continue
}
nextOverrides = append(nextOverrides, compiled)
nextHashes[plugin.Key] = plugin.SourceHash
delete(taskPluginSyncState.errors, plugin.Key)
logger.LogDebug(
ctx,
"task_plugin subsystem=sync event=plugin plugin=%q version=%q action=compile_success",
plugin.Key,
plugin.Version,
)
}
if err = jsplugin.DefaultRegistry.ReplaceOverrides(nextOverrides); err != nil {
syncErr := fmt.Errorf("publish task plugin generation: %w", err)
taskPluginSyncState.lastRebuild = taskPluginRebuildOutcome{
Status: "failed",
AttemptedAt: time.Now(),
Generation: jsplugin.DefaultRegistry.Generation().Number,
DatabaseRevision: databaseSnapshot.Revision,
Error: syncErr.Error(),
}
logger.LogDebug(
ctx,
"task_plugin subsystem=sync event=failed stage=publish retained_generation=%d retained_generation_active=true database_revision=%q elapsed_ms=%d",
jsplugin.DefaultRegistry.Generation().Number,
databaseSnapshot.Revision,
time.Since(started).Milliseconds(),
)
return syncErr
}
taskPluginSyncState.hashes = nextHashes
for key := range taskPluginSyncState.errors {
if !seen[key] {
delete(taskPluginSyncState.errors, key)
}
}
pluginErrors := jsplugin.DefaultRegistry.RoutingErrors()
for key, message := range taskPluginSyncState.errors {
pluginErrors[key] = message
}
pluginErrorCount := len(pluginErrors)
status := "success"
if pluginErrorCount > 0 {
status = "partial"
}
taskPluginSyncState.lastRebuild = taskPluginRebuildOutcome{
Status: status,
AttemptedAt: time.Now(),
Generation: jsplugin.DefaultRegistry.Generation().Number,
DatabaseRevision: databaseSnapshot.Revision,
PluginErrorCount: pluginErrorCount,
}
logger.LogDebug(
ctx,
"task_plugin subsystem=sync event=complete database_revision=%q previous_generation=%d generation=%d status=%q active_overrides=%d plugin_errors=%d elapsed_ms=%d",
databaseSnapshot.Revision,
generationBefore,
jsplugin.DefaultRegistry.Generation().Number,
status,
len(jsplugin.DefaultRegistry.ActiveOverridePlugins()),
pluginErrorCount,
time.Since(started).Milliseconds(),
)
return nil
}
func SyncTaskPluginsOnce() {
if err := syncTaskPluginsOnce(); err != nil {
common.SysError(err.Error())
}
}
func SyncTaskPlugins() {
SyncTaskPluginsOnce()
for range time.NewTicker(30 * time.Second).C {
SyncTaskPluginsOnce()
}
}
+269
View File
@@ -0,0 +1,269 @@
package controller
import (
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/relay"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/gin-gonic/gin"
)
// taskPluginSubmitDiagnostics keeps plugin-only lifecycle logging out of the
// ordinary task path. An empty plugin key makes every method a no-op.
type taskPluginSubmitDiagnostics struct {
context *gin.Context
pluginKey string
generation uint64
}
func newTaskPluginSubmitDiagnostics(c *gin.Context) taskPluginSubmitDiagnostics {
diagnostics := taskPluginSubmitDiagnostics{
context: c,
pluginKey: c.GetString("expected_task_plugin_key"),
}
if diagnostics.pluginKey == "" {
return diagnostics
}
if pinnedValue, exists := c.Get(pluginruntime.ContextKeyPinnedPlugin); exists {
if pinned, ok := pinnedValue.(pluginruntime.PinnedPlugin); ok && pinned.Generation != nil {
diagnostics.generation = pinned.Generation.Number
}
}
return diagnostics
}
func (d taskPluginSubmitDiagnostics) start(info *relaycommon.RelayInfo) {
if d.pluginKey == "" {
return
}
logger.LogDebug(
d.context,
"task_plugin subsystem=submit event=start generation=%d plugin=%q model=%q action_present=%t",
d.generation,
d.pluginKey,
info.OriginModelName,
info.Action != "",
)
}
func (d taskPluginSubmitDiagnostics) refund(stage string) {
if d.pluginKey == "" {
return
}
logger.LogDebug(
d.context,
"task_plugin subsystem=submit event=refund_invoked generation=%d plugin=%q stage=%q durable=false",
d.generation,
d.pluginKey,
stage,
)
}
func (d taskPluginSubmitDiagnostics) cancelled(stage string, attempt int) {
if d.pluginKey == "" {
return
}
logger.LogDebug(
d.context,
"task_plugin subsystem=submit event=cancelled generation=%d plugin=%q stage=%q attempt=%d",
d.generation,
d.pluginKey,
stage,
attempt,
)
}
func (d taskPluginSubmitDiagnostics) attempt(attempt int, channel *model.Channel, locked bool) {
if d.pluginKey == "" || channel == nil {
return
}
logger.LogDebug(
d.context,
"task_plugin subsystem=submit event=attempt generation=%d plugin=%q attempt=%d channel_id=%d channel_type=%d locked=%t",
d.generation,
d.pluginKey,
attempt,
channel.Id,
channel.Type,
locked,
)
}
func (d taskPluginSubmitDiagnostics) attemptSucceeded(attempt int, result *relay.TaskSubmitResult) {
if d.pluginKey == "" || result == nil {
return
}
logger.LogDebug(
d.context,
"task_plugin subsystem=submit event=attempt_succeeded generation=%d plugin=%q attempt=%d platform=%q quota=%d task_data_bytes=%d client_response=%t immediate=%t",
d.generation,
d.pluginKey,
attempt,
result.Platform,
result.Quota,
len(result.TaskData),
result.ClientResponse != nil,
result.Immediate != nil,
)
}
func (d taskPluginSubmitDiagnostics) attemptFailed(attempt int, channel *model.Channel, taskErr *dto.TaskError, willRetry bool) {
if d.pluginKey == "" || channel == nil || taskErr == nil {
return
}
logger.LogDebug(
d.context,
"task_plugin subsystem=submit event=attempt_failed generation=%d plugin=%q attempt=%d channel_id=%d channel_type=%d code=%q status=%d local=%t will_retry=%t",
d.generation,
d.pluginKey,
attempt,
channel.Id,
channel.Type,
taskErr.Code,
taskErr.StatusCode,
taskErr.LocalError,
willRetry,
)
}
func (d taskPluginSubmitDiagnostics) failed(stage, reason string, taskErr *dto.TaskError, durable bool) {
if d.pluginKey == "" {
return
}
code := ""
status := 0
local := true
if taskErr != nil {
code = taskErr.Code
status = taskErr.StatusCode
local = taskErr.LocalError
}
logger.LogDebug(
d.context,
"task_plugin subsystem=submit event=failed generation=%d plugin=%q stage=%q reason=%q code=%q status=%d local=%t durable=%t",
d.generation,
d.pluginKey,
stage,
reason,
code,
status,
local,
durable,
)
}
func (d taskPluginSubmitDiagnostics) reserve(event string, quota int) {
if d.pluginKey == "" {
return
}
logger.LogDebug(
d.context,
"task_plugin subsystem=submit event=%s generation=%d plugin=%q quota=%d",
event,
d.generation,
d.pluginKey,
quota,
)
}
func (d taskPluginSubmitDiagnostics) insertStart(task *model.Task) {
if d.pluginKey == "" || task == nil {
return
}
logger.LogDebug(
d.context,
"task_plugin subsystem=submit event=insert_start generation=%d plugin=%q public_task_id=%q platform=%q channel_id=%d quota=%d",
d.generation,
d.pluginKey,
task.TaskID,
task.Platform,
task.ChannelId,
task.Quota,
)
}
func (d taskPluginSubmitDiagnostics) durable(task *model.Task) {
if d.pluginKey == "" || task == nil {
return
}
logger.LogDebug(
d.context,
"task_plugin subsystem=submit event=durable generation=%d plugin=%q public_task_id=%q status=%q durable=true",
d.generation,
d.pluginKey,
task.TaskID,
taskPluginDebugStatus(string(task.Status)),
)
}
func (d taskPluginSubmitDiagnostics) settleStart(task *model.Task, quota int) {
if d.pluginKey == "" || task == nil {
return
}
logger.LogDebug(
d.context,
"task_plugin subsystem=submit event=settle_start generation=%d plugin=%q public_task_id=%q quota=%d durable=true",
d.generation,
d.pluginKey,
task.TaskID,
quota,
)
}
func (d taskPluginSubmitDiagnostics) complete(task *model.Task, quota int) {
if d.pluginKey == "" || task == nil {
return
}
logger.LogDebug(
d.context,
"task_plugin subsystem=submit event=complete generation=%d plugin=%q public_task_id=%q quota=%d durable=true",
d.generation,
d.pluginKey,
task.TaskID,
quota,
)
}
func (d taskPluginSubmitDiagnostics) present(task *model.Task, presenter string) {
if d.pluginKey == "" || task == nil {
return
}
logger.LogDebug(
d.context,
"task_plugin subsystem=submit event=present generation=%d plugin=%q public_task_id=%q presenter=%q durable=true",
d.generation,
d.pluginKey,
task.TaskID,
presenter,
)
}
func (d taskPluginSubmitDiagnostics) presentError(taskErr *dto.TaskError) {
if d.pluginKey == "" || taskErr == nil {
return
}
logger.LogDebug(
d.context,
"task_plugin subsystem=submit event=present_error generation=%d plugin=%q code=%q status=%d local=%t",
d.generation,
d.pluginKey,
taskErr.Code,
taskErr.StatusCode,
taskErr.LocalError,
)
}
func taskPluginDebugStatus(status string) string {
switch model.TaskStatus(status) {
case model.TaskStatusSubmitted,
model.TaskStatusQueued,
model.TaskStatusInProgress,
model.TaskStatusSuccess,
model.TaskStatusFailure:
return status
default:
return "unknown"
}
}
+88
View File
@@ -0,0 +1,88 @@
package controller
import (
"bytes"
"net/http"
"net/http/httptest"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestTaskPluginSubmitDiagnosticsArePluginOnlyAndDoNotLogPayloads(t *testing.T) {
previousDebug := common.DebugEnabled
common.DebugEnabled = true
t.Cleanup(func() { common.DebugEnabled = previousDebug })
var output bytes.Buffer
common.LogWriterMu.Lock()
previousWriter := gin.DefaultErrorWriter
gin.DefaultErrorWriter = &output
common.LogWriterMu.Unlock()
t.Cleanup(func() {
common.LogWriterMu.Lock()
gin.DefaultErrorWriter = previousWriter
common.LogWriterMu.Unlock()
})
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/v1/tasks/debug-plugin", nil)
c.Set(common.RequestIdKey, "plugin-submit-request")
info := &relaycommon.RelayInfo{
OriginModelName: "safe-model",
TaskRelayInfo: &relaycommon.TaskRelayInfo{Action: "https://private-action.invalid/?key=hidden"},
}
newTaskPluginSubmitDiagnostics(c).start(info)
assert.Empty(t, output.String())
c.Set("expected_task_plugin_key", "debug-plugin")
diagnostics := newTaskPluginSubmitDiagnostics(c)
diagnostics.start(info)
diagnostics.attemptSucceeded(1, &relay.TaskSubmitResult{
UpstreamTaskID: "private-upstream-canary",
TaskData: []byte("private-task-data-canary"),
ClientResponse: map[string]any{"secret": "private-client-response-canary"},
Platform: constant.TaskPlatform("debug-plugin"),
Quota: 12,
})
task := &model.Task{
TaskID: "public-task-id",
Platform: constant.TaskPlatform("debug-plugin"),
Status: model.TaskStatus("https://private-status.invalid/?key=hidden"),
PrivateData: model.TaskPrivateData{
UpstreamTaskID: "private-task-record-canary",
ResultURL: "https://private-url.invalid/result",
},
}
diagnostics.insertStart(task)
diagnostics.durable(task)
diagnostics.complete(task, 12)
logOutput := output.String()
require.Contains(t, logOutput, "plugin-submit-request")
assert.Contains(t, logOutput, `plugin="debug-plugin"`)
assert.Contains(t, logOutput, `public_task_id="public-task-id"`)
assert.Contains(t, logOutput, "task_data_bytes=24")
assert.Contains(t, logOutput, "action_present=true")
assert.Contains(t, logOutput, `status="unknown"`)
for _, secret := range []string{
"private-upstream-canary",
"private-task-data-canary",
"private-client-response-canary",
"private-task-record-canary",
"private-url.invalid",
"private-action.invalid",
"private-status.invalid",
"key=hidden",
} {
assert.NotContains(t, logOutput, secret)
}
}
File diff suppressed because it is too large. Load diff
+536 -131
View File
@@ -1,27 +1,56 @@
package controller
import (
"bytes"
"context"
"encoding/base64"
"errors"
"fmt"
"io"
"net"
"net/http"
"net/url"
"strconv"
"strings"
"sync"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
relaychannel "github.com/QuantumNous/new-api/relay/channel"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/system_setting"
"github.com/gin-gonic/gin"
"golang.org/x/net/http/httpguts"
)
var errTaskMediaRequestRejected = errors.New("task media request rejected")
var taskMediaResponseHeaderTimeout = 60 * time.Second
var taskMediaDataURLMaxEncodedBytes = 64 << 20
type taskMediaProxyError struct {
status int
code string
message string
err error
}
func (e *taskMediaProxyError) Error() string {
if e.err == nil {
return e.message
}
return e.message + ": " + e.err.Error()
}
func (e *taskMediaProxyError) Unwrap() error {
return e.err
}
// videoProxyError returns a standardized OpenAI-style error response.
func videoProxyError(c *gin.Context, status int, errType, message string) {
c.Header("Cache-Control", "private, no-store")
c.JSON(status, gin.H{
"error": gin.H{
"message": message,
@@ -37,8 +66,7 @@ func VideoProxy(c *gin.Context) {
return
}
userID := c.GetInt("id")
task, exists, err := model.GetByTaskId(userID, taskID)
task, exists, err := getTaskForArtifactRequest(c, taskID)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to query task %s: %s", taskID, err.Error()))
videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to query task")
@@ -48,141 +76,505 @@ func VideoProxy(c *gin.Context) {
videoProxyError(c, http.StatusNotFound, "invalid_request_error", "Task not found")
return
}
if task.Status != model.TaskStatusSuccess {
videoProxyError(c, http.StatusBadRequest, "invalid_request_error",
fmt.Sprintf("Task is not completed yet, current status: %s", task.Status))
return
}
channel, err := model.CacheGetChannel(task.ChannelId)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to get channel for task %s: %s", taskID, err.Error()))
videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to retrieve channel information")
return
var descriptor *relaychannel.TaskContentRequest
if taskHasPluginExecution(task) {
artifacts, projectionErr := projectTaskArtifacts(task)
if projectionErr == nil {
for _, artifact := range artifacts {
if artifact.Type != "video" {
continue
}
adaptor, adaptorErr := initTaskArtifactAdaptor(task)
if adaptorErr == nil {
if provider, ok := adaptor.(relaychannel.TaskContentRequestProvider); ok {
descriptor, adaptorErr = provider.BuildContentRequest(task, artifact.Key, relaychannel.TaskArtifactClientRequest{
Method: c.Request.Method,
Headers: taskArtifactClientHeaders(c.Request.Header),
})
}
}
if adaptorErr != nil {
logger.LogWarn(c.Request.Context(), fmt.Sprintf("Failed to resolve plugin video content for task %s", taskID))
descriptor = nil
}
break
}
} else {
logger.LogWarn(c.Request.Context(), fmt.Sprintf("Failed to project plugin video for task %s", taskID))
}
}
baseURL := channel.GetBaseURL()
if baseURL == "" {
baseURL = "https://api.openai.com"
}
var videoURL string
proxy := channel.GetSetting().Proxy
client := service.GetSSRFProtectedHTTPClient()
if proxy != "" {
// 渠道代理路径的连接由代理侧建立,无法做拨号时逐 IP 校验,
// 因此后面对 videoURL 保留请求前的一次性 SSRF 校验。
client, err = service.GetHttpClientWithProxy(proxy)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to create proxy client for task %s: %s", taskID, err.Error()))
videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to create proxy client")
if descriptor == nil {
resultURL := task.GetResultURL()
if isTaskMediaFallbackLoop(resultURL, task.TaskID) {
writeTaskMediaProxyError(c, &taskMediaProxyError{
status: http.StatusGone, code: "artifact_gone",
message: "Artifact content is no longer available",
})
return
}
}
ctx, cancel := context.WithTimeout(c.Request.Context(), 60*time.Second)
defer cancel()
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "", nil)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to create request: %s", err.Error()))
videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to create proxy request")
return
}
switch channel.Type {
case constant.ChannelTypeGemini:
apiKey := task.PrivateData.Key
if apiKey == "" {
logger.LogError(c.Request.Context(), fmt.Sprintf("Missing stored API key for Gemini task %s", taskID))
videoProxyError(c, http.StatusInternalServerError, "server_error", "API key not stored for task")
return
}
videoURL, err = getGeminiVideoURL(channel, task, apiKey)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to resolve Gemini video URL for task %s: %s", taskID, err.Error()))
videoProxyError(c, http.StatusBadGateway, "server_error", "Failed to resolve Gemini video URL")
return
}
req.Header.Set("x-goog-api-key", apiKey)
case constant.ChannelTypeVertexAi:
videoURL, err = getVertexVideoURL(channel, task)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to resolve Vertex video URL for task %s: %s", taskID, err.Error()))
videoProxyError(c, http.StatusBadGateway, "server_error", "Failed to resolve Vertex video URL")
return
}
case constant.ChannelTypeOpenAI, constant.ChannelTypeSora:
videoURL = fmt.Sprintf("%s/v1/videos/%s/content", baseURL, task.GetUpstreamTaskID())
req.Header.Set("Authorization", "Bearer "+channel.Key)
default:
// Video URL is stored in PrivateData.ResultURL (fallback to FailReason for old data)
videoURL = task.GetResultURL()
}
videoURL = strings.TrimSpace(videoURL)
if videoURL == "" {
logger.LogError(c.Request.Context(), fmt.Sprintf("Video URL is empty for task %s", taskID))
videoProxyError(c, http.StatusBadGateway, "server_error", "Failed to fetch video content")
return
}
if strings.HasPrefix(videoURL, "data:") {
if err := writeVideoDataURL(c, videoURL); err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to decode video data URL for task %s: %s", taskID, err.Error()))
videoProxyError(c, http.StatusBadGateway, "server_error", "Failed to fetch video content")
}
return
}
var validateErr error
if proxy == "" {
validateErr = service.ValidateSSRFProtectedFetchURL(videoURL)
} else {
fetchSetting := system_setting.GetFetchSetting()
validateErr = common.ValidateURLWithFetchSetting(videoURL, fetchSetting.EnableSSRFProtection, fetchSetting.AllowPrivateIp, fetchSetting.DomainFilterMode, fetchSetting.IpFilterMode, fetchSetting.DomainList, fetchSetting.IpList, fetchSetting.AllowedPorts, fetchSetting.ApplyIPFilterForDomain)
}
if validateErr != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Video URL blocked for task %s: %v", taskID, validateErr))
videoProxyError(c, http.StatusForbidden, "server_error", fmt.Sprintf("request blocked: %v", validateErr))
return
}
req.URL, err = url.Parse(videoURL)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to parse URL %s: %s", videoURL, err.Error()))
videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to create proxy request")
return
}
resp, err := client.Do(req)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to fetch video from %s: %s", videoURL, err.Error()))
videoProxyError(c, http.StatusBadGateway, "server_error", "Failed to fetch video content")
return
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
logger.LogError(c.Request.Context(), fmt.Sprintf("Upstream returned status %d for %s", resp.StatusCode, videoURL))
videoProxyError(c, http.StatusBadGateway, "server_error",
fmt.Sprintf("Upstream service returned status %d", resp.StatusCode))
return
}
for key, values := range resp.Header {
for _, value := range values {
c.Writer.Header().Add(key, value)
descriptor = &relaychannel.TaskContentRequest{
URL: resultURL,
Method: c.Request.Method,
Credentialless: true,
}
}
c.Writer.Header().Set("Cache-Control", "public, max-age=86400")
c.Writer.WriteHeader(resp.StatusCode)
if _, err = io.Copy(c.Writer, resp.Body); err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to stream video content: %s", err.Error()))
if err := proxyTaskMedia(c, task, descriptor); err != nil {
writeTaskMediaProxyError(c, err)
}
}
func proxyTaskMedia(c *gin.Context, task *model.Task, descriptor *relaychannel.TaskContentRequest) error {
if descriptor == nil {
return &taskMediaProxyError{
status: http.StatusInternalServerError, code: "artifact_plugin_error",
message: "Artifact content plugin returned no request",
}
}
rawURL := strings.TrimSpace(descriptor.URL)
if rawURL == "" {
return &taskMediaProxyError{
status: http.StatusGone, code: "artifact_gone",
message: "Artifact content is no longer available",
}
}
if strings.HasPrefix(rawURL, "data:") {
if len(rawURL) > taskMediaDataURLMaxEncodedBytes {
return &taskMediaProxyError{
status: http.StatusBadGateway, code: "artifact_request_rejected",
message: "Artifact request was rejected", err: errTaskMediaRequestRejected,
}
}
if err := writeVideoDataURL(c, rawURL); err != nil {
return &taskMediaProxyError{
status: http.StatusBadGateway, code: "artifact_upstream_error",
message: "Failed to decode artifact content", err: err,
}
}
return nil
}
if len(rawURL) > 64<<10 {
return &taskMediaProxyError{
status: http.StatusBadGateway, code: "artifact_request_rejected",
message: "Artifact request was rejected", err: errTaskMediaRequestRejected,
}
}
parsedURL, err := url.Parse(rawURL)
if err != nil || parsedURL == nil || (parsedURL.Scheme != "http" && parsedURL.Scheme != "https") ||
parsedURL.Host == "" || parsedURL.User != nil || parsedURL.Fragment != "" {
return &taskMediaProxyError{
status: http.StatusBadGateway, code: "artifact_request_rejected",
message: "Artifact request was rejected", err: errTaskMediaRequestRejected,
}
}
if isTaskMediaFallbackLoop(rawURL, task.TaskID) || isSelfTaskMediaURL(c, parsedURL) {
return &taskMediaProxyError{
status: http.StatusBadGateway, code: "artifact_request_rejected",
message: "Artifact proxy loop was rejected", err: errTaskMediaRequestRejected,
}
}
method := strings.ToUpper(strings.TrimSpace(descriptor.Method))
if method == "" {
method = c.Request.Method
}
switch method {
case http.MethodGet, http.MethodHead, http.MethodPost:
default:
return &taskMediaProxyError{
status: http.StatusBadGateway, code: "artifact_request_rejected",
message: "Artifact request method was rejected", err: errTaskMediaRequestRejected,
}
}
if len(descriptor.Body) > 1<<20 {
return &taskMediaProxyError{
status: http.StatusBadGateway, code: "artifact_request_rejected",
message: "Artifact request body was rejected", err: errTaskMediaRequestRejected,
}
}
if descriptor.Credentialless &&
(method != http.MethodGet && method != http.MethodHead ||
descriptor.Body != nil || len(descriptor.Headers) != 0) {
return &taskMediaProxyError{
status: http.StatusBadGateway, code: "artifact_request_rejected",
message: "Credentialless artifact request was rejected", err: errTaskMediaRequestRejected,
}
}
channel, err := model.CacheGetChannel(task.ChannelId)
if err != nil {
return &taskMediaProxyError{
status: http.StatusServiceUnavailable, code: "artifact_plugin_unavailable",
message: "Artifact channel is unavailable", err: err,
}
}
proxy := strings.TrimSpace(channel.GetSetting().Proxy)
if err := validateTaskMediaURL(rawURL, proxy); err != nil {
return &taskMediaProxyError{
status: http.StatusBadGateway, code: "artifact_request_rejected",
message: "Artifact request was rejected", err: err,
}
}
client := service.GetSSRFProtectedHTTPClient()
if proxy != "" {
client, err = service.GetHttpClientWithProxy(proxy)
if err != nil {
return &taskMediaProxyError{
status: http.StatusInternalServerError, code: "artifact_internal_error",
message: "Failed to create artifact proxy client", err: err,
}
}
}
if client == nil {
client = http.DefaultClient
}
req, err := http.NewRequestWithContext(c.Request.Context(), method, parsedURL.String(), bytes.NewReader(descriptor.Body))
if err != nil {
return &taskMediaProxyError{
status: http.StatusInternalServerError, code: "artifact_internal_error",
message: "Failed to create artifact request", err: err,
}
}
if err := applyTaskMediaRequestHeaders(req.Header, descriptor.Headers); err != nil {
return &taskMediaProxyError{
status: http.StatusBadGateway, code: "artifact_request_rejected",
message: "Artifact request headers were rejected", err: err,
}
}
clientHeaders := taskArtifactClientHeaders(c.Request.Header)
for name, value := range clientHeaders {
req.Header.Set(name, value)
}
client = taskMediaRedirectClient(client, proxy, c, clientHeaders, descriptor.Credentialless)
clientWithoutBodyTimeout := *client
clientWithoutBodyTimeout.Timeout = 0
resp, err := doTaskMediaRequest(&clientWithoutBodyTimeout, req, taskMediaResponseHeaderTimeout)
if err != nil {
if errors.Is(err, errTaskMediaRequestRejected) {
return &taskMediaProxyError{
status: http.StatusBadGateway, code: "artifact_request_rejected",
message: "Artifact redirect was rejected", err: err,
}
}
var netErr net.Error
if errors.Is(err, context.DeadlineExceeded) || errors.As(err, &netErr) && netErr.Timeout() {
return &taskMediaProxyError{
status: http.StatusGatewayTimeout, code: "artifact_upstream_timeout",
message: "Artifact upstream request timed out", err: err,
}
}
return &taskMediaProxyError{
status: http.StatusBadGateway, code: "artifact_upstream_error",
message: "Failed to fetch artifact content", err: err,
}
}
defer resp.Body.Close()
switch resp.StatusCode {
case http.StatusOK, http.StatusPartialContent, http.StatusNotModified, http.StatusRequestedRangeNotSatisfiable:
copyTaskMediaResponseHeaders(c.Writer.Header(), resp.Header)
setTaskMediaResponseSecurityHeaders(c.Writer.Header())
c.Status(resp.StatusCode)
c.Writer.WriteHeaderNow()
if c.Request.Method == http.MethodHead || resp.StatusCode == http.StatusNotModified {
return nil
}
if _, err := io.Copy(c.Writer, resp.Body); err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to stream task media: %v", err))
}
return nil
case http.StatusUnauthorized, http.StatusForbidden:
return &taskMediaProxyError{
status: http.StatusBadGateway, code: "artifact_upstream_auth_failed",
message: "Artifact upstream authentication failed",
}
case http.StatusNotFound, http.StatusGone:
return &taskMediaProxyError{
status: http.StatusGone, code: "artifact_gone",
message: "Artifact content is no longer available",
}
case http.StatusTooManyRequests:
if retryAfter := strings.TrimSpace(resp.Header.Get("Retry-After")); retryAfter != "" &&
len(retryAfter) <= 256 && !strings.ContainsAny(retryAfter, "\r\n") {
c.Header("Retry-After", retryAfter)
}
return &taskMediaProxyError{
status: http.StatusServiceUnavailable, code: "artifact_upstream_busy",
message: "Artifact upstream is busy",
}
default:
return &taskMediaProxyError{
status: http.StatusBadGateway, code: "artifact_upstream_error",
message: fmt.Sprintf("Artifact upstream returned status %d", resp.StatusCode),
}
}
}
type taskMediaHTTPResult struct {
response *http.Response
err error
}
type taskMediaCancelBody struct {
io.ReadCloser
cancel context.CancelFunc
once sync.Once
}
func (b *taskMediaCancelBody) Close() error {
b.once.Do(b.cancel)
return b.ReadCloser.Close()
}
func doTaskMediaRequest(client *http.Client, request *http.Request, responseHeaderTimeout time.Duration) (*http.Response, error) {
if client == nil {
client = http.DefaultClient
}
requestContext, cancel := context.WithCancel(request.Context())
request = request.Clone(requestContext)
resultChannel := make(chan taskMediaHTTPResult, 1)
go func() {
response, err := client.Do(request)
resultChannel <- taskMediaHTTPResult{response: response, err: err}
}()
timer := time.NewTimer(responseHeaderTimeout)
defer timer.Stop()
cleanupResult := func() {
go func() {
result := <-resultChannel
if result.response != nil && result.response.Body != nil {
_ = result.response.Body.Close()
}
}()
}
select {
case result := <-resultChannel:
if result.err != nil {
cancel()
if result.response != nil && result.response.Body != nil {
_ = result.response.Body.Close()
}
return nil, result.err
}
if result.response == nil || result.response.Body == nil {
cancel()
return nil, errors.New("artifact upstream returned no response body")
}
result.response.Body = &taskMediaCancelBody{
ReadCloser: result.response.Body,
cancel: cancel,
}
return result.response, nil
case <-timer.C:
cancel()
cleanupResult()
return nil, context.DeadlineExceeded
case <-request.Context().Done():
cancel()
cleanupResult()
return nil, request.Context().Err()
}
}
func applyTaskMediaRequestHeaders(destination http.Header, headers map[string]string) error {
if len(headers) > 64 {
return errTaskMediaRequestRejected
}
for name, value := range headers {
name = strings.TrimSpace(name)
if !httpguts.ValidHeaderFieldName(name) || !httpguts.ValidHeaderFieldValue(value) || len(value) > 8192 {
return errTaskMediaRequestRejected
}
switch strings.ToLower(name) {
case "host", "content-length", "accept-encoding", "connection", "proxy-connection", "keep-alive",
"proxy-authorization", "te", "trailer", "transfer-encoding", "upgrade":
return errTaskMediaRequestRejected
}
destination.Set(name, value)
}
return nil
}
func taskMediaRedirectClient(base *http.Client, proxy string, c *gin.Context, clientHeaders map[string]string, credentialless bool) *http.Client {
cloned := *base
cloned.CheckRedirect = func(req *http.Request, via []*http.Request) error {
if len(via) >= 10 {
return fmt.Errorf("%w: too many redirects", errTaskMediaRequestRejected)
}
if req.URL == nil || (req.URL.Scheme != "http" && req.URL.Scheme != "https") ||
req.URL.Host == "" || req.URL.User != nil || req.URL.Fragment != "" {
return fmt.Errorf("%w: invalid redirect URL", errTaskMediaRequestRejected)
}
if err := validateTaskMediaURL(req.URL.String(), proxy); err != nil {
return fmt.Errorf("%w: %v", errTaskMediaRequestRejected, err)
}
if isSelfTaskMediaURL(c, req.URL) {
return fmt.Errorf("%w: proxy loop", errTaskMediaRequestRejected)
}
if len(via) > 0 && !sameTaskMediaOrigin(via[len(via)-1].URL, req.URL) {
if !credentialless {
return fmt.Errorf("%w: credentialed cross-origin redirect", errTaskMediaRequestRejected)
}
for name := range req.Header {
req.Header.Del(name)
}
req.Body = http.NoBody
req.GetBody = nil
req.ContentLength = 0
}
for name, value := range clientHeaders {
req.Header.Set(name, value)
}
return nil
}
return &cloned
}
func validateTaskMediaURL(rawURL, proxy string) error {
if proxy == "" {
return service.ValidateSSRFProtectedFetchURL(rawURL)
}
fetchSetting := system_setting.GetFetchSetting()
return common.ValidateURLWithFetchSetting(
rawURL,
fetchSetting.EnableSSRFProtection,
fetchSetting.AllowPrivateIp,
fetchSetting.DomainFilterMode,
fetchSetting.IpFilterMode,
fetchSetting.DomainList,
fetchSetting.IpList,
fetchSetting.AllowedPorts,
fetchSetting.ApplyIPFilterForDomain,
)
}
func sameTaskMediaOrigin(left, right *url.URL) bool {
if left == nil || right == nil {
return false
}
return strings.EqualFold(left.Scheme, right.Scheme) &&
strings.EqualFold(normalizeTaskMediaHost(left.Scheme, left.Host), normalizeTaskMediaHost(right.Scheme, right.Host))
}
func normalizeTaskMediaHost(scheme, host string) string {
host = strings.ToLower(strings.TrimSpace(host))
hostname, port, err := net.SplitHostPort(host)
if err != nil {
return strings.TrimSuffix(host, ".")
}
hostname = strings.TrimSuffix(strings.ToLower(hostname), ".")
if (strings.EqualFold(scheme, "http") && port == "80") || (strings.EqualFold(scheme, "https") && port == "443") {
return hostname
}
return net.JoinHostPort(hostname, port)
}
func isSelfTaskMediaURL(c *gin.Context, target *url.URL) bool {
if c == nil || target == nil || !isTaskMediaProxyPath(target.Path) {
return false
}
targetHost := normalizeTaskMediaHost(target.Scheme, target.Host)
if targetHost == "" {
return true
}
scheme := strings.TrimSpace(strings.Split(c.Request.Header.Get("X-Forwarded-Proto"), ",")[0])
if scheme == "" {
scheme = "http"
if c.Request.TLS != nil {
scheme = "https"
}
}
hosts := []string{c.Request.Host}
if forwardedHost := strings.TrimSpace(strings.Split(c.Request.Header.Get("X-Forwarded-Host"), ",")[0]); forwardedHost != "" {
hosts = append(hosts, forwardedHost)
}
for _, host := range hosts {
if strings.EqualFold(targetHost, normalizeTaskMediaHost(scheme, host)) {
return true
}
}
return false
}
func isTaskMediaProxyPath(path string) bool {
if strings.HasPrefix(path, "/v1/videos/") && strings.HasSuffix(path, "/content") {
return true
}
return strings.HasPrefix(path, "/v1/tasks/") &&
strings.Contains(path, "/artifacts/") &&
strings.HasSuffix(path, "/content")
}
func isTaskMediaFallbackLoop(rawURL, taskID string) bool {
parsedURL, err := url.Parse(strings.TrimSpace(rawURL))
if err != nil || parsedURL == nil {
return false
}
path, err := url.PathUnescape(parsedURL.EscapedPath())
if err != nil {
path = parsedURL.Path
}
if path == "/v1/videos/"+taskID+"/content" {
return true
}
artifactPrefix := "/v1/tasks/" + taskID + "/artifacts/"
return strings.HasPrefix(path, artifactPrefix) && strings.HasSuffix(path, "/content")
}
func copyTaskMediaResponseHeaders(destination, source http.Header) {
for _, name := range []string{
"Content-Type",
"Content-Length",
"Content-Range",
"Accept-Ranges",
"ETag",
"Last-Modified",
"Content-Disposition",
} {
for _, value := range source.Values(name) {
destination.Add(name, value)
}
}
}
func setTaskMediaResponseSecurityHeaders(header http.Header) {
header.Set("Cache-Control", "private, no-store")
header.Set("Content-Security-Policy", "sandbox; default-src 'none'")
header.Set("Referrer-Policy", "no-referrer")
header.Set("X-Content-Type-Options", "nosniff")
}
func writeTaskMediaProxyError(c *gin.Context, err error) {
if c.Writer.Written() {
logger.LogError(c.Request.Context(), err.Error())
return
}
var proxyErr *taskMediaProxyError
if !errors.As(err, &proxyErr) {
proxyErr = &taskMediaProxyError{
status: http.StatusBadGateway, code: "artifact_upstream_error",
message: "Failed to fetch artifact content", err: err,
}
}
c.Header("Cache-Control", "private, no-store")
writeTaskArtifactError(c, proxyErr.status, proxyErr.code, proxyErr.message)
}
func writeVideoDataURL(c *gin.Context, dataURL string) error {
if len(dataURL) > taskMediaDataURLMaxEncodedBytes {
return errTaskMediaRequestRejected
}
parts := strings.SplitN(dataURL, ",", 2)
if len(parts) != 2 {
return fmt.Errorf("invalid data url")
@@ -199,18 +591,31 @@ func writeVideoDataURL(c *gin.Context, dataURL string) error {
if mimeType == "" {
mimeType = "video/mp4"
}
if len(mimeType) > 255 || !httpguts.ValidHeaderFieldValue(mimeType) {
return fmt.Errorf("invalid data url media type")
}
videoBytes, err := base64.StdEncoding.DecodeString(payload)
if err != nil {
videoBytes, err = base64.RawStdEncoding.DecodeString(payload)
if err != nil {
return err
var encoding *base64.Encoding
var contentLength int64
for _, candidate := range []*base64.Encoding{base64.StdEncoding, base64.RawStdEncoding} {
decodedLength, err := io.Copy(io.Discard, base64.NewDecoder(candidate, strings.NewReader(payload)))
if err == nil {
encoding = candidate
contentLength = decodedLength
break
}
}
if encoding == nil {
return fmt.Errorf("invalid base64 data")
}
c.Writer.Header().Set("Content-Type", mimeType)
c.Writer.Header().Set("Cache-Control", "public, max-age=86400")
c.Writer.Header().Set("Content-Length", strconv.FormatInt(contentLength, 10))
setTaskMediaResponseSecurityHeaders(c.Writer.Header())
c.Writer.WriteHeader(http.StatusOK)
_, err = c.Writer.Write(videoBytes)
if c.Request.Method == http.MethodHead {
return nil
}
_, err := io.Copy(c.Writer, base64.NewDecoder(encoding, strings.NewReader(payload)))
return err
}
-294
View File
@@ -1,294 +0,0 @@
package controller
import (
"fmt"
"io"
"strconv"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relay"
)
func getGeminiVideoURL(channel *model.Channel, task *model.Task, apiKey string) (string, error) {
if channel == nil || task == nil {
return "", fmt.Errorf("invalid channel or task")
}
if url := extractGeminiVideoURLFromTaskData(task); url != "" {
return ensureAPIKey(url, apiKey), nil
}
baseURL := constant.ChannelBaseURLs[channel.Type]
if channel.GetBaseURL() != "" {
baseURL = channel.GetBaseURL()
}
adaptor := relay.GetTaskAdaptor(constant.TaskPlatform(strconv.Itoa(channel.Type)))
if adaptor == nil {
return "", fmt.Errorf("gemini task adaptor not found")
}
if apiKey == "" {
return "", fmt.Errorf("api key not available for task")
}
proxy := channel.GetSetting().Proxy
resp, err := adaptor.FetchTask(baseURL, apiKey, map[string]any{
"task_id": task.GetUpstreamTaskID(),
"action": task.Action,
}, proxy)
if err != nil {
return "", fmt.Errorf("fetch task failed: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return "", fmt.Errorf("read task response failed: %w", err)
}
taskInfo, parseErr := adaptor.ParseTaskResult(body)
if parseErr == nil && taskInfo != nil && taskInfo.RemoteUrl != "" {
return ensureAPIKey(taskInfo.RemoteUrl, apiKey), nil
}
if url := extractGeminiVideoURLFromPayload(body); url != "" {
return ensureAPIKey(url, apiKey), nil
}
if parseErr != nil {
return "", fmt.Errorf("parse task result failed: %w", parseErr)
}
return "", fmt.Errorf("gemini video url not found")
}
func extractGeminiVideoURLFromTaskData(task *model.Task) string {
if task == nil || len(task.Data) == 0 {
return ""
}
var payload map[string]any
if err := common.Unmarshal(task.Data, &payload); err != nil {
return ""
}
return extractGeminiVideoURLFromMap(payload)
}
func extractGeminiVideoURLFromPayload(body []byte) string {
var payload map[string]any
if err := common.Unmarshal(body, &payload); err != nil {
return ""
}
return extractGeminiVideoURLFromMap(payload)
}
func extractGeminiVideoURLFromMap(payload map[string]any) string {
if payload == nil {
return ""
}
if uri, ok := payload["uri"].(string); ok && uri != "" {
return uri
}
if resp, ok := payload["response"].(map[string]any); ok {
if uri := extractGeminiVideoURLFromResponse(resp); uri != "" {
return uri
}
}
return ""
}
func extractGeminiVideoURLFromResponse(resp map[string]any) string {
if resp == nil {
return ""
}
if gvr, ok := resp["generateVideoResponse"].(map[string]any); ok {
if uri := extractGeminiVideoURLFromGeneratedSamples(gvr); uri != "" {
return uri
}
}
if videos, ok := resp["videos"].([]any); ok {
for _, video := range videos {
if vm, ok := video.(map[string]any); ok {
if uri, ok := vm["uri"].(string); ok && uri != "" {
return uri
}
}
}
}
if uri, ok := resp["video"].(string); ok && uri != "" {
return uri
}
if uri, ok := resp["uri"].(string); ok && uri != "" {
return uri
}
return ""
}
func extractGeminiVideoURLFromGeneratedSamples(gvr map[string]any) string {
if gvr == nil {
return ""
}
if samples, ok := gvr["generatedSamples"].([]any); ok {
for _, sample := range samples {
if sm, ok := sample.(map[string]any); ok {
if video, ok := sm["video"].(map[string]any); ok {
if uri, ok := video["uri"].(string); ok && uri != "" {
return uri
}
}
}
}
}
return ""
}
func getVertexVideoURL(channel *model.Channel, task *model.Task) (string, error) {
if channel == nil || task == nil {
return "", fmt.Errorf("invalid channel or task")
}
if url := strings.TrimSpace(task.GetResultURL()); url != "" && !isTaskProxyContentURL(url, task.TaskID) {
return url, nil
}
if url := extractVertexVideoURLFromTaskData(task); url != "" {
return url, nil
}
baseURL := constant.ChannelBaseURLs[channel.Type]
if channel.GetBaseURL() != "" {
baseURL = channel.GetBaseURL()
}
adaptor := relay.GetTaskAdaptor(constant.TaskPlatform(strconv.Itoa(channel.Type)))
if adaptor == nil {
return "", fmt.Errorf("vertex task adaptor not found")
}
key := getVertexTaskKey(channel, task)
if key == "" {
return "", fmt.Errorf("vertex key not available for task")
}
resp, err := adaptor.FetchTask(baseURL, key, map[string]any{
"task_id": task.GetUpstreamTaskID(),
"action": task.Action,
}, channel.GetSetting().Proxy)
if err != nil {
return "", fmt.Errorf("fetch task failed: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return "", fmt.Errorf("read task response failed: %w", err)
}
taskInfo, parseErr := adaptor.ParseTaskResult(body)
if parseErr == nil && taskInfo != nil && strings.TrimSpace(taskInfo.Url) != "" {
return taskInfo.Url, nil
}
if url := extractVertexVideoURLFromPayload(body); url != "" {
return url, nil
}
if parseErr != nil {
return "", fmt.Errorf("parse task result failed: %w", parseErr)
}
return "", fmt.Errorf("vertex video url not found")
}
func isTaskProxyContentURL(url string, taskID string) bool {
if strings.TrimSpace(url) == "" || strings.TrimSpace(taskID) == "" {
return false
}
return strings.Contains(url, "/v1/videos/"+taskID+"/content")
}
func getVertexTaskKey(channel *model.Channel, task *model.Task) string {
if task != nil {
if key := strings.TrimSpace(task.PrivateData.Key); key != "" {
return key
}
}
if channel == nil {
return ""
}
keys := channel.GetKeys()
for _, key := range keys {
key = strings.TrimSpace(key)
if key != "" {
return key
}
}
return strings.TrimSpace(channel.Key)
}
func extractVertexVideoURLFromTaskData(task *model.Task) string {
if task == nil || len(task.Data) == 0 {
return ""
}
return extractVertexVideoURLFromPayload(task.Data)
}
func extractVertexVideoURLFromPayload(body []byte) string {
var payload map[string]any
if err := common.Unmarshal(body, &payload); err != nil {
return ""
}
resp, ok := payload["response"].(map[string]any)
if !ok || resp == nil {
return ""
}
if videos, ok := resp["videos"].([]any); ok && len(videos) > 0 {
if video, ok := videos[0].(map[string]any); ok && video != nil {
if b64, _ := video["bytesBase64Encoded"].(string); strings.TrimSpace(b64) != "" {
mime, _ := video["mimeType"].(string)
enc, _ := video["encoding"].(string)
return buildVideoDataURL(mime, enc, b64)
}
}
}
if b64, _ := resp["bytesBase64Encoded"].(string); strings.TrimSpace(b64) != "" {
enc, _ := resp["encoding"].(string)
return buildVideoDataURL("", enc, b64)
}
if video, _ := resp["video"].(string); strings.TrimSpace(video) != "" {
if strings.HasPrefix(video, "data:") || strings.HasPrefix(video, "http://") || strings.HasPrefix(video, "https://") {
return video
}
enc, _ := resp["encoding"].(string)
return buildVideoDataURL("", enc, video)
}
return ""
}
func buildVideoDataURL(mimeType string, encoding string, base64Data string) string {
mime := strings.TrimSpace(mimeType)
if mime == "" {
enc := strings.TrimSpace(encoding)
if enc == "" {
enc = "mp4"
}
if strings.Contains(enc, "/") {
mime = enc
} else {
mime = "video/" + enc
}
}
return "data:" + mime + ";base64," + base64Data
}
func ensureAPIKey(uri, key string) string {
if key == "" || uri == "" {
return uri
}
if strings.Contains(uri, "key=") {
return uri
}
if strings.Contains(uri, "?") {
return fmt.Sprintf("%s&key=%s", uri, key)
}
return fmt.Sprintf("%s?key=%s", uri, key)
}
+138
View File
@@ -0,0 +1,138 @@
# Task plugin API v1
Task plugins are single-file synchronous ECMAScript modules. The plugin contract
is currently unreleased; [`v1.schema.json`](./v1.schema.json) and
[`v1.d.ts`](./v1.d.ts) are the authoritative v1 contract.
## Contract and lifecycle
Every plugin exports `meta`, `buildSubmitRequest`, `parseSubmitResponse`, and
`parseTaskResult`. A `per_task` plugin also exports `buildQueryRequest`; a
`batch` plugin exports `buildBatchQueryRequest` and `parseBatchResult`.
`meta.author.name` is required and `meta.author.url`, when present, must be an
absolute HTTP(S) URL. This is self-declared attribution; a future marketplace's
verified publisher identity is a separate host-owned record.
Plugins may declare authenticated vendor-native `meta.routes` and claim
host-owned names through `meta.protocols`. Submit and dynamic routes name a
`native` decoder and presenter; query routes name only a presenter. Protocol
bindings are registered once by the host registry, and protocol decoders receive
the host-parsed `body` union plus the pinned model. Shared protocol hooks are
synchronous transformations; Go owns connections and wire framing.
The host selects a channel, invokes the request-building hook, validates the
returned URL against the channel host, performs HTTP, and gives the decoded
response to the matching parse hook. It owns persistence, retries, polling,
billing, and settlement. Plugins only transform data and report usage facts.
See [v1.d.ts](./v1.d.ts) for signatures and
[v1.schema.json](./v1.schema.json) for machine-readable shapes.
Plugins that expose task outputs export `listArtifacts(task)` and
`buildContentRequest(ctx)` together. Artifacts are projected on explicit reads
from persisted `Task.Data`; they are never stored as a second source of truth.
The list contains only stable `key`, `type`, and optional `mimeType` fields.
The content hook receives the selected key, raw decoded task data, the explicit
private upstream task id, the producer plugin version, channel authentication,
and a safe client Range/conditional-header subset. Its URL and headers exist
only for that proxy request.
When a Responses observation reaches persisted `SUCCESS`, the host also runs
the pinned plugin's `listArtifacts` and injects a read-only
`ctx.artifacts[key] = {key, type, mimeType?, url}` map into `renderEvents` or
`renderFinal`. Each `url` is a long-lived host-signed capability URL, never the
provider URL from `Task.Data`. Nonterminal and failed tasks receive no artifact
map. Capability construction or rendering failure fails only that Responses
observation; it cannot change the task, billing settlement, or refunds.
The absolute URL uses `TaskPublicAddress`, falling back only to
`ServerAddress`; multi-node deployments must share the effective
`CRYPTO_SECRET`.
Dashboard artifact reads return each `content_url` (or the legacy
`legacy_content_url`) directly, without a temporary URL exchange. Capability
generation and verification are stateless and have no expiry; after
verification the host still loads the task, owner, and plugin needed to serve
the artifact. Rotating `CRYPTO_SECRET` invalidates issued URLs. The `access`
query is redacted before request logging.
Deployment boundaries and concurrency environment variables are documented in
[v1.md](./v1.md#generic-task-management-api).
The host treats `protocols.openai_video.render` as a standard DTO, not an arbitrary
JSON passthrough. Unknown top-level fields and legacy `task_id` are removed,
`id` is forced to the public task id, and case-insensitive `url` entries are
removed from metadata. Provider output URLs belong only behind artifact
capabilities.
Provider-authenticated content URLs must use the channel base host or a
plugin-declared `meta.allowedHosts` entry. A public dynamic CDN URL may instead
set `credentialless: true`; the host then permits only GET/HEAD with no
plugin-supplied headers or body and applies SSRF checks to the initial URL and
every redirect.
Registry publication is generation-atomic. A request pins one plugin generation
for its full lifetime, while background polling may use a later active plugin
version. New versions must continue parsing responses for in-flight tasks.
Root administrators can inspect the local node with
`GET /api/plugin/task/runtime/status`. The response includes the node-local generation,
a deterministic revision of the active database overrides, the latest rebuild
outcome, and plugin-level compile or routing errors. Generation numbers are
local to a node; compare database revisions when diagnosing rollout lag between
nodes. If the database snapshot is temporarily unavailable, the endpoint keeps
serving node-local state and the last known revision with `database_error` set.
For live diagnosis, start the process with `DEBUG=true` and filter logs on
`task_plugin`. Plugin registry, routing, endpoint ownership, channel selection,
submit durability, polling adapters, and protocol observation emit safe
key/value lifecycle events. Request-context events carry the request id;
scheduled, background, and context-less work is labeled `SYSTEM`. Plugin
`console.log` output is also forwarded in DEBUG mode. Hook-time output is
prefixed with plugin key/version; module-initialization output may have an empty
identity during initial upload validation. Do not print credentials, headers,
request bodies, upstream payloads, or private URLs from plugin code; free-form
console output cannot be redacted by the host.
## Fixtures and dry runs
A fixture case is `{name?, hook, member?, args, expected?, expectedError?}`.
Keep deterministic cases for every exported hook, its main error branch, batch
behavior, renderers, usage, and content requests. Run a fixture locally with:
```sh
new-api plugin lint plugin.js
new-api plugin test plugin.js --fixture golden.json
```
Root administrators can open the plugin detail Sandbox tab, choose a hook, and
submit an `args` JSON array. `POST /api/plugin/task/:key/dryrun` compiles the
active database source or factory source in a temporary registry and invokes
only that synchronous function. Dry runs never execute a request descriptor and
therefore never contact an upstream service.
## Upload and release
Upload from the root-only task plugin page or `POST /api/plugin/task` with
`{"source":"...","remark":"..."}`. The server compiles the module, validates
v1 metadata and required exports, and rejects invalid source before saving it.
Use semantic plugin versions. Reusing a key/version with different source is
rejected; activate or roll back a stored version through the management page.
For a third-party platform, create a channel of type `Task Plugin`, select the
plugin key, provide an explicit base URL, and configure models. Clients may use
the plugin's declared native routes. The generic management surface remains
`POST /v1/tasks/:pluginKey`, `GET /v1/tasks/:taskId`, and
`GET /v1/tasks/:taskId/artifacts` plus
`GET|HEAD /v1/tasks/:taskId/artifacts/:key/content`.
## Security boundary
Plugins have no `fetch`, filesystem, `require`, imports, async functions, or
environment access. The host limits execution time, concurrency, input size,
allowed request hosts, and resolves OAuth credentials outside JavaScript.
Multipart files enter JavaScript only as opaque references.
This is not a hard memory-isolation boundary. A plugin sees data needed for the
current request and can influence an authenticated upstream request. Uploading a
plugin is an administrator-level trust decision equivalent to configuring a
channel credential. Review source and version diffs before activation. Never run
untrusted plugins merely because they compile.
Usage hooks may return facts such as seconds, resolution, or upstream units, but
must never calculate prices or attempt quota settlement. The host owns all
pricing and clamps billing conversions.
+48
View File
@@ -0,0 +1,48 @@
export type JSONValue = null | boolean | number | string | readonly JSONValue[] | {readonly [key: string]: JSONValue};
export type FileReference = Readonly<{ref: string; field: string; filename: string; mimeType: string; size: number}>;
export type FilePlaceholder = Readonly<{__fileRef: string; encoding: "base64" | "dataUrl"; mimeType?: string; maxBytes?: number}>;
export type DecodedBody =
| Readonly<{kind: "json"; value: JSONValue}>
| Readonly<{kind: "form"; fields: Readonly<Record<string, readonly string[]>>}>
| Readonly<{kind: "multipart"; fields: Readonly<Record<string, readonly string[]>>; files: readonly FileReference[]}>
| Readonly<{kind: "none"}>;
export interface NativeDecodeContext {method: string; path: string; params: Readonly<Record<string, string>>; query: Readonly<Record<string, readonly string[]>>; body: DecodedBody}
export interface ProtocolDecodeContext extends NativeDecodeContext {protocol: "openai_responses" | "openai_video"; operation: string; model: string; stream: boolean}
export type SubmitIntent = {kind: "submit"; model: string; action?: string; requestBody?: unknown; originTaskIds?: readonly string[]};
export type QueryIntent = {kind: "query"; taskIds: readonly string[]};
export type TaskIntent = SubmitIntent | QueryIntent;
export interface NativeRoute {method: "GET" | "POST" | "PUT" | "PATCH" | "DELETE"; path: string; type: "submit" | "query" | "dynamic"; action?: string; taskIdParam?: string; decode?: string; render: string; models?: readonly string[]}
export type ProtocolName = "openai_responses" | "openai_video";
export type ResponsesMode = "stream" | "sync" | "background";
export type ProtocolClaim =
| "openai_video"
| {name: "openai_responses"; supports: readonly ResponsesMode[]; models?: readonly string[]}
| {name: "openai_video"; models?: readonly string[]};
export type LocalizedText = string | ({ en: string } & Record<string, string>);
export type UsageFieldSchema = {type: "number"; unit: "second" | "count" | "token" | "credit"; description?: LocalizedText} | {type: "boolean"; description?: LocalizedText} | {enum: readonly string[]; description?: LocalizedText};
export type UsageExample = {label: string; facts: Readonly<Record<string, string | number>>};
export interface Meta {apiVersion: 1; key: string; name: string; icon?: string; description?: LocalizedText; version: string; author: {name: string; url?: string}; channelTypes?: readonly number[]; models: readonly string[]; fetchMode: "per_task" | "batch"; allowedHosts?: readonly string[]; routes?: readonly NativeRoute[]; protocols?: readonly ProtocolClaim[]; usageSchema?: Readonly<Record<string, UsageFieldSchema>>; usageExamples?: readonly UsageExample[]; auth?: "none" | "api_key" | "vertex_oauth" | {type: "none" | "api_key" | "oauth2_jwt"}}
export interface TaskView {task_id: string; status: string; progress?: string; fail_reason?: string; created_at?: number; updated_at?: number; data?: unknown; properties?: Record<string, unknown>}
export interface DriverContext {requestBody: unknown; requestHeaders: Readonly<Record<string, string>>; action: string; model: string; upstreamModel: string; baseUrl: string; apiKey?: string; authHeader: string; files: readonly FileReference[]; publicTaskId: string; originTasks?: readonly {taskId: string; upstreamTaskId: string; action: string; status: string; data: unknown}[]}
export interface RequestDescriptor {url: string; method?: string; headers?: Record<string, string>; /** JSON body may contain FilePlaceholder objects at any depth; the host replaces each with a Base64 or data-URL string. */ body?: unknown; credentialless?: boolean; action?: string; model?: string; rewriteModel?: string; bodyType?: "json" | "multipart"; parts?: readonly {name: string; value?: unknown; fileRef?: string; filename?: string}[]}
export interface UpstreamResponse {statusCode: number; headers: Readonly<Record<string, readonly string[]>>; body: unknown}
export interface NormalizedTaskResult {taskId?: string; status: "NOT_START" | "SUBMITTED" | "QUEUED" | "IN_PROGRESS" | "SUCCESS" | "FAILURE" | "UNKNOWN"; progress?: string; reason?: string; url?: string; remoteUrl?: string; completionTokens?: number; totalTokens?: number}
export interface TaskArtifact {key: string; type: "video" | "audio" | "image" | "file"; mimeType?: string}
export declare const meta: Meta;
export declare const native: Record<string, ((ctx: NativeDecodeContext) => TaskIntent) | ((ctx: NativeDecodeContext, task: TaskView | readonly TaskView[]) => unknown)> & {error?: (ctx: NativeDecodeContext, error: {code: string; message: string; httpStatus: number; retryable: boolean}) => unknown};
export declare const protocols: {
openai_responses?: {decodeRequest(ctx: ProtocolDecodeContext): SubmitIntent; renderEvents?(ctx: unknown, task: TaskView, previousState: unknown): unknown; renderFinal?(ctx: unknown, task: TaskView): unknown};
openai_video?: {decodeRequest(ctx: ProtocolDecodeContext): SubmitIntent; render(ctx: unknown, task: TaskView): unknown};
};
export declare function buildSubmitRequest(ctx: DriverContext): RequestDescriptor;
export declare function parseSubmitResponse(ctx: DriverContext, response: UpstreamResponse): {taskId: string; taskData?: unknown; immediate?: NormalizedTaskResult};
export declare function buildQueryRequest(ctx: DriverContext & {taskId: string}): RequestDescriptor;
export declare function buildBatchQueryRequest(ctx: DriverContext, taskIds: readonly string[]): RequestDescriptor;
export declare function parseTaskResult(ctx: DriverContext, body: unknown): NormalizedTaskResult;
export declare function parseBatchResult(ctx: DriverContext, body: unknown): readonly (NormalizedTaskResult & {taskId: string; data?: unknown})[];
export declare function extractUsage(ctx: DriverContext & {usagePurpose?: "facts" | "billing_ratios"}): Readonly<Record<string, string | number | boolean>> | null;
export declare function extractUsageOnSubmit(ctx: DriverContext, taskData: unknown): Readonly<Record<string, string | number | boolean>> | null;
export declare function extractUsageOnComplete(task: TaskView, result: NormalizedTaskResult, data: unknown): Readonly<Record<string, string | number | boolean>> | null;
export declare function listArtifacts(task: {taskId: string; status: string; action: string; data: unknown; producerVersion: string}): readonly TaskArtifact[];
export declare function buildContentRequest(ctx: DriverContext & {artifactKey: string; data: unknown; upstreamTaskId: string; clientRequest: {method: "GET" | "HEAD"; headers: Readonly<Record<string, string>>}}): RequestDescriptor;
+140
View File
@@ -0,0 +1,140 @@
# Task Plugin API v1
Task Plugin v1 has two independent entry surfaces. `meta.routes` registers plugin-owned native URLs; `meta.protocols` claims host-owned protocols without registering or copying their URLs. `apiVersion` remains `1`.
## Manifest
```js
export const meta = {
apiVersion: 1,
key: "vendor",
name: "Vendor",
version: "1.0.0",
author: {name: "Author"},
description: {en: "Video generation via the vendor API", zh: "通过厂商接口生成视频"},
models: ["vendor-model"],
fetchMode: "per_task",
routes: [
{method: "POST", path: "/vendor/v1/jobs", type: "submit", decode: "createJob", render: "jobCreated"},
{method: "GET", path: "/vendor/v1/jobs/:task_id", type: "query", render: "jobStatus"},
],
protocols: [{name: "openai_responses", supports: ["stream", "sync", "background"]}, "openai_video"],
};
```
`submit` and `dynamic` routes require `decode` and `render`. `query` routes require `render`, prohibit `decode`, and use `taskIdParam` (default `task_id`). Names refer only to callable members of `native`. Route conflicts use method plus normalized path shape and the route index is published atomically with its plugin generation.
`routes[].models` optionally restricts a `submit` or `dynamic` route to a subset of `meta.models`. The host matches the canonical top-level `model` body field before any JS hook runs; a missing, non-string, or unlisted model is rejected with 400 (plugin routes are exclusive paths, so there is no relay fallback). Declare it only when the model lives at the body top level — vendor formats that nest the model elsewhere should omit `models` and filter inside `decode`. `query` routes have no body and reject the field.
Each `protocols` entry claims a host protocol. A protocol that defines modes must be claimed in object form with an explicit `supports`; the bare-string form is valid only for protocols without modes (`openai_video`). `supports` declares which client request forms the plugin accepts on `openai_responses`: `"stream"` (`stream: true`), `"sync"` (neither flag; the call blocks until the terminal Response), `"background"` (`background: true`; the create returns a pending Response immediately). An unsupported request form is rejected with a 400 at channel selection, before any plugin hook runs and before billing. Retrieval (`GET /v1/responses/:response_id`) is not a mode: every created response is always retrievable.
`{name, models}` still narrows that protocol's endpoint bindings to a subset of `meta.models` and composes with `supports`. Unlisted models never enter the plugin on that protocol path — they fall through to the built-in Go relay. Cross-plugin endpoint conflicts are judged on the narrowed set, so two plugins may share one protocol path by claiming disjoint model sets.
Enabled uploads pre-flight the candidate against the live routing generation and reject the first channel-type, native-route, or protocol-model conflict (the error names the counterpart plugin). Set `force: true` or `enabled: false` to store the plugin anyway.
`endpoints`, `routes[].renderer`, global `resolveRequest`, global `renderError`, and global `renderers` are rejected. `parseSubmitResponse` returns only `{taskId, taskData}` (plus the documented lifecycle fields); `clientResponse` is rejected.
`icon` is an optional LobeHub icon name string (for example `Sora.Color`). The values `text` and `text:<label>` request a generated text avatar instead (label defaults to the first two characters of `name`). It is display-only and does not participate in routing, billing, or admission beyond type and length checks.
| Field | Type | Notes |
|-------|------|-------|
| `key` | string | Required. Canonical plugin id, ≤ 30 characters. |
| `name` | string | Required. Display name. |
| `icon` | string | Optional LobeHub icon or `text` / `text:<label>`. ≤ 128 characters. |
| `description` | LocalizedText | Optional plugin summary. See LocalizedText. ≤ 512 runes per locale. |
| `version` | string | Required semver. |
| `author` | `{name, url?}` | Required name; `url`, when present, must be an absolute HTTP(S) URL. |
`channelTypes` lists the legacy channel types this plugin's driver can drive (for example, sora declares `[55, 1]` because the same OpenAI-type base URL and bearer key serve both chat and video). Every entry equally participates in channel selection, historical `Task.Platform` matching, and the `byChannelType` routing index; the same type value may not appear on two plugins. Third-party plugins normally omit `channelTypes` and live on type-59 "Task Plugin" channels bound by `task_plugin_key`. The previous split identity/compatibility field names are rejected.
Numeric `usageSchema` fields declare a host-owned unit of `second`, `count`, `token`, or `credit`. Boolean fields declare `{type: "boolean"}`.
`meta.usageExamples` is an optional display-only list of pricing examples. Each `label` is a human-readable spec name. Each `facts` object must be a complete vector over `usageSchema`: every declared key present, no undeclared keys, and values that pass the same validation as usage facts. Plugins whose schema contains a `unit: "token"` number field must declare at least one example. Examples never affect billing. Labels are spec names such as `std · 10s`, not storefront prices.
## LocalizedText
`meta.description` and each `usageSchema` field `description` accept LocalizedText. A bare string is equivalent to `{en: <string>}`. A map must include a non-empty `en` value. The host normalizes both forms to a map; API responses always emit an object.
```js
description: "Video generation via the vendor API"
description: {en: "Video generation via the vendor API", zh: "通过厂商接口生成视频"}
```
Rules:
- Locale keys match `^[a-zA-Z]{2,3}(-[a-zA-Z0-9]{2,8})*$`. A map has at most 16 locales.
- Keys are canonicalized to BCP-47 casing (`EN` → `en`, `zh-tw` → `zh-TW`, `zh-hans` → `zh-Hans`); two keys that collide after canonicalization are rejected as duplicates.
- Each value is trimmed, must be non-empty, and must not contain control characters.
- `meta.description` is at most 512 runes per locale. Each `usageSchema` field `description` is at most 256 runes per locale.
- The frontend resolves a locale with exact tag → primary subtag → `en` (for example `zh-TW` → `zh` → `en`).
- Description copy in any language must not include vendor currency prices. The same prohibition applies to `usageExamples` labels.
## Request body
Every decoder receives one host-parsed body:
```ts
{kind: "json", value}
{kind: "form", fields: Record<string, readonly string[]>}
{kind: "multipart", fields: Record<string, readonly string[]>, files: readonly FileReference[]}
{kind: "none"}
```
`FileReference` contains `{ref, field, filename, mimeType, size}`. The ref is opaque and request-scoped; file bytes, readers, temporary paths, and a second inbound `requestBody` representation are never exposed to JavaScript. The decoder creates a new canonical `{kind:"submit", model, action?, requestBody?, originTaskIds?}` intent. Driver hooks consume that normalized `requestBody` and remain surface-independent.
## Request building
`buildSubmitRequest` returns a request descriptor. JSON bodies (the default) may embed a file placeholder at any object or array depth. The host replaces each placeholder with an encoded string before the request is sent; JavaScript never sees file bytes.
```js
{ __fileRef: "request_file:input_reference", encoding: "base64" } // raw Base64
{ __fileRef: "request_file:input_reference", encoding: "dataUrl" } // data:<mime>;base64,<data>
{ __fileRef: "request_file:input_reference", encoding: "dataUrl", mimeType: "image/png", maxBytes: 20971520 }
```
Placeholder objects may contain only `__fileRef`, `encoding`, `mimeType`, and `maxBytes`. `encoding` is `base64` or `dataUrl`. `__fileRef` must resolve to an uploaded multipart file (`request_file:<field>`). `mimeType` optionally overrides the part `Content-Type` for `dataUrl` (otherwise the part header, else `application/octet-stream`). `maxBytes` is an optional vendor cap; the host also applies `MAX_FILE_DOWNLOAD_MB` (default 64MB) to each file and to the total inlined bytes. Multipart descriptors still stream files with `parts[].fileRef`.
## Origin tasks
Decoders normalize vendor references (draft ids, continuation ids, gateway-issued asset references) into `originTaskIds` as public gateway task IDs. The host never parses vendor request bodies.
The host verifies ownership, requires every referenced task to belong to the same plugin platform set and one enabled channel, pins that channel for the submit (including retries), and injects the resolved rows into driver hooks as `ctx.originTasks` (`taskId`, `upstreamTaskId`, `action`, `status`, and `data`, including the private upstream id). Renderers and presenters never receive `originTasks`.
## Native surface
```js
export const native = {
createJob(ctx) { return {kind: "submit", model: ctx.body.value.model, requestBody: ctx.body.value}; },
jobCreated(ctx, task) { return {job_id: task.task_id, state: "queued"}; },
jobStatus(ctx, task) { return task.data; },
error(ctx, error) { return {code: error.code, message: error.message}; },
};
```
The host authenticates, checks ownership, persists the task, builds `TaskView`, and replaces known private task-id fields with the public ID before calling presenters. Native presenters own the vendor envelope. Authenticated native query presenters may intentionally pass through provider URLs from `task.data`.
## Error contract
Hook `Error` messages are surfaced (truncated and sanitized) to API callers. Plugin authors should write validation failures as user-readable sentences.
## Host protocols
The host registry owns these bindings:
| Protocol | Operation | Binding | Body |
|--------------------|-----------|---------------------------|-----------------------------------|
| `openai_responses` | create | `POST /v1/responses` | JSON |
| `openai_responses` | retrieve | `GET /v1/responses/:response_id` | none |
| `openai_video` | create | `POST /v1/videos` | JSON or multipart |
| `openai_video` | retrieve | `GET /v1/videos/:task_id` | none |
| `openai_video` | content | `GET | HEAD /v1/videos/:task_id/content` | none |
`protocols.openai_responses` always requires `decodeRequest`. The remaining hooks are derived from `supports` and verified exactly at load in both directions: `"stream"` requires `renderEvents`; `"sync"` or `"background"` requires `renderFinal`. A hook required by a supported mode but not exported, or an exported hook that no supported mode uses, rejects the plugin. For a plugin supporting only `"stream"`, the host renders retrieval at terminal SUCCESS by calling `renderEvents` once with no previous state — the same first-call-at-terminal semantics streaming already requires. The mode vocabulary is append-only under apiVersion 1: new modes may be added (existing manifests are unaffected; a manifest declaring a mode loads only on gateways that define it), but a mode may never be removed, renamed, or re-mapped to different hooks without an apiVersion bump. `protocols.openai_video` requires `decodeRequest` for create and `render` for retrieve. The host extracts and pins `ctx.model`, normalizes `ctx.stream`, frames SSE, creates failure envelopes, and calls `renderFinal` only for `SUCCESS`.
`background: true` on create returns the pending Response immediately (host-synthesized, no plugin hook). Retrieval renders via `renderFinal` only at terminal status. Plugins declaring `"background"` need no create-time hook; the deliverable is rendered at retrieval via `renderFinal`. Plugins cannot observe the `background` field. Unlike the upstream OpenAI API, a plugin that does not declare `"stream"` rejects `stream: true` with a 400; the error names the supported forms.
Protocol media uses host-injected `ctx.artifacts[key].url`. Provider URLs from `task.data` are not protocol output. OpenAI Video projections are DTO-whitelisted and the host overwrites identity, lifecycle, timestamps, and removes URL-like metadata.
## Persisted data and driver hooks
The persisted field remains `task.data`; there is no `task.raw` alias. Driver hooks (`buildSubmitRequest`, `parseSubmitResponse`, query/result, usage, artifact, and content hooks) stay flat and must not branch on the client path or protocol.
+46
View File
@@ -0,0 +1,46 @@
{
"$schema": "https://json-schema.org/draft/2020-12/schema",
"$id": "https://github.com/QuantumNous/new-api/docs/plugin-api/v1.schema.json",
"title": "new-api Task Plugin v1 manifest",
"type": "object",
"additionalProperties": false,
"required": ["apiVersion", "key", "name", "version", "author", "models", "fetchMode"],
"properties": {
"apiVersion": {"const": 1}, "key": {"type": "string"}, "name": {"type": "string"}, "icon": {"type": "string", "maxLength": 128}, "description": {"$ref": "#/$defs/localizedText"}, "version": {"type": "string"},
"author": {"type": "object", "required": ["name"], "additionalProperties": false, "properties": {"name": {"type": "string"}, "url": {"type": "string", "format": "uri"}}},
"channelTypes": {"type": "array", "uniqueItems": true, "items": {"type": "integer", "minimum": 1}}, "models": {"type": "array", "minItems": 1, "uniqueItems": true, "items": {"type": "string", "minLength": 1}},
"fetchMode": {"enum": ["per_task", "batch"]}, "allowedHosts": {"type": "array", "uniqueItems": true, "items": {"type": "string"}},
"protocols": {"type": "array", "uniqueItems": true, "items": {"oneOf": [
{"enum": ["openai_video"]},
{"type": "object", "additionalProperties": false, "required": ["name", "supports"], "properties": {"name": {"const": "openai_responses"}, "supports": {"type": "array", "minItems": 1, "uniqueItems": true, "items": {"enum": ["stream", "sync", "background"]}}, "models": {"type": "array", "minItems": 1, "uniqueItems": true, "items": {"type": "string", "minLength": 1}}}},
{"type": "object", "additionalProperties": false, "required": ["name"], "properties": {"name": {"const": "openai_video"}, "models": {"type": "array", "minItems": 1, "uniqueItems": true, "items": {"type": "string", "minLength": 1}}}}
]}},
"routes": {"type": "array", "items": {"$ref": "#/$defs/route"}}, "usageSchema": {"type": "object", "additionalProperties": {"type": "object", "properties": {"description": {"$ref": "#/$defs/localizedText"}}}}, "usageExamples": {"type": "array", "items": {"type": "object", "additionalProperties": false, "required": ["label", "facts"], "properties": {"label": {"type": "string"}, "facts": {"type": "object"}}}}, "auth": {}
},
"$defs": {
"localizedText": {
"oneOf": [
{"type": "string", "minLength": 1},
{
"type": "object",
"required": ["en"],
"maxProperties": 16,
"propertyNames": {"pattern": "^[a-zA-Z]{2,3}(-[a-zA-Z0-9]{2,8})*$"},
"additionalProperties": {"type": "string", "minLength": 1}
}
]
},
"filePlaceholder": {
"type": "object",
"additionalProperties": false,
"required": ["__fileRef", "encoding"],
"properties": {
"__fileRef": {"type": "string", "minLength": 1},
"encoding": {"enum": ["base64", "dataUrl"]},
"mimeType": {"type": "string", "minLength": 1},
"maxBytes": {"type": "integer", "exclusiveMinimum": 0}
}
},
"route": {"type": "object", "additionalProperties": false, "required": ["method", "path", "type", "render"], "properties": {"method": {"enum": ["GET", "POST", "PUT", "PATCH", "DELETE"]}, "path": {"type": "string", "pattern": "^/"}, "type": {"enum": ["submit", "query", "dynamic"]}, "action": {"type": "string"}, "taskIdParam": {"type": "string"}, "decode": {"type": "string"}, "render": {"type": "string"}, "models": {"type": "array", "minItems": 1, "uniqueItems": true, "items": {"type": "string", "minLength": 1}}}, "allOf": [{"if": {"properties": {"type": {"const": "query"}}}, "then": {"allOf": [{"not": {"required": ["decode"]}}, {"not": {"required": ["models"]}}]}}, {"if": {"properties": {"type": {"enum": ["submit", "dynamic"]}}}, "then": {"required": ["decode"]}}]}
}
}
+115
View File
@@ -0,0 +1,115 @@
package dto
type ChannelPinSource string
const (
PinSourceToken ChannelPinSource = "token" // Rank 0, highest
PinSourceOriginTask ChannelPinSource = "origin_task" // Rank 10
)
const (
PinRankToken = 0
PinRankOriginTask = 10
)
type PinRetryMode int
const (
PinRetrySameChannel PinRetryMode = iota
PinRetrySingleAttempt
)
func (m PinRetryMode) Stricter(other PinRetryMode) PinRetryMode {
if m == PinRetrySingleAttempt || other == PinRetrySingleAttempt {
return PinRetrySingleAttempt
}
return PinRetrySameChannel
}
type ChannelPin struct {
ChannelId int
Source ChannelPinSource
Rank int
RetryMode PinRetryMode
}
type ChannelFilterKind string
const (
FilterRequestPath ChannelFilterKind = "request_path"
FilterTaskPluginIdentity ChannelFilterKind = "task_plugin_identity"
)
type ChannelFilter struct {
Kind ChannelFilterKind
RequestPath string
TaskPluginKey string
TaskPluginChannelTypes []int
}
type ChannelConstraints struct {
Pins []ChannelPin
Filters []ChannelFilter
}
func (cc *ChannelConstraints) AddPin(p ChannelPin) {
if cc == nil {
return
}
cc.Pins = append(cc.Pins, p)
}
func (cc *ChannelConstraints) AddFilter(f ChannelFilter) {
if cc == nil {
return
}
cc.Filters = append(cc.Filters, f)
}
// ResolvedPin returns the winning pin after priority resolution.
// Lowest Rank wins. Pins that name the same channel are merged (stricter RetryMode).
// overridden lists pins that lost to a different channel (for warn logging).
func (cc *ChannelConstraints) ResolvedPin() (ChannelPin, bool, []ChannelPin) {
if cc == nil || len(cc.Pins) == 0 {
return ChannelPin{}, false, nil
}
merged := make(map[int]ChannelPin, len(cc.Pins))
order := make([]int, 0, len(cc.Pins))
for _, pin := range cc.Pins {
existing, seen := merged[pin.ChannelId]
if !seen {
merged[pin.ChannelId] = pin
order = append(order, pin.ChannelId)
continue
}
existing.RetryMode = existing.RetryMode.Stricter(pin.RetryMode)
if pin.Rank < existing.Rank {
existing.Rank = pin.Rank
existing.Source = pin.Source
}
merged[pin.ChannelId] = existing
}
winner := merged[order[0]]
for _, channelID := range order[1:] {
candidate := merged[channelID]
if candidate.Rank < winner.Rank {
winner = candidate
}
}
var overridden []ChannelPin
for _, channelID := range order {
candidate := merged[channelID]
if candidate.ChannelId != winner.ChannelId {
overridden = append(overridden, candidate)
}
}
return winner, true, overridden
}
func (cc *ChannelConstraints) SuppressesRetry() bool {
pin, found, _ := cc.ResolvedPin()
return found && pin.RetryMode == PinRetrySingleAttempt
}
+46
View File
@@ -0,0 +1,46 @@
package dto
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestResolvedPinPriorityAndMerge(t *testing.T) {
t.Run("token pin beats origin pin on a different channel", func(t *testing.T) {
constraints := &ChannelConstraints{}
constraints.AddPin(ChannelPin{ChannelId: 10, Source: PinSourceOriginTask, Rank: PinRankOriginTask, RetryMode: PinRetrySameChannel})
constraints.AddPin(ChannelPin{ChannelId: 1, Source: PinSourceToken, Rank: PinRankToken, RetryMode: PinRetrySingleAttempt})
pin, found, overridden := constraints.ResolvedPin()
require.True(t, found)
assert.Equal(t, 1, pin.ChannelId)
assert.Equal(t, PinSourceToken, pin.Source)
assert.Equal(t, PinRetrySingleAttempt, pin.RetryMode)
require.Len(t, overridden, 1)
assert.Equal(t, PinSourceOriginTask, overridden[0].Source)
assert.Equal(t, 10, overridden[0].ChannelId)
})
t.Run("same channel pins merge to the stricter retry mode", func(t *testing.T) {
constraints := &ChannelConstraints{}
constraints.AddPin(ChannelPin{ChannelId: 7, Source: PinSourceOriginTask, Rank: PinRankOriginTask, RetryMode: PinRetrySameChannel})
constraints.AddPin(ChannelPin{ChannelId: 7, Source: PinSourceToken, Rank: PinRankToken, RetryMode: PinRetrySingleAttempt})
pin, found, overridden := constraints.ResolvedPin()
require.True(t, found)
assert.Equal(t, 7, pin.ChannelId)
assert.Equal(t, PinSourceToken, pin.Source)
assert.Equal(t, PinRetrySingleAttempt, pin.RetryMode)
assert.Empty(t, overridden)
assert.True(t, constraints.SuppressesRetry())
})
t.Run("empty set has no pin", func(t *testing.T) {
pin, found, overridden := (*ChannelConstraints)(nil).ResolvedPin()
assert.False(t, found)
assert.Zero(t, pin.ChannelId)
assert.Nil(t, overridden)
})
}
+82
View File
@@ -0,0 +1,82 @@
package dto
// PluginResponsesResponse is the host-owned Responses facade used for stream
// snapshots and sanitized terminal failures. Non-stream success objects may
// retain additional validated plugin fields, but identifiers, lifecycle state,
// and retrieval metadata are always populated by the host.
type PluginResponsesResponse struct {
ID string `json:"id"`
Object string `json:"object"`
CreatedAt int64 `json:"created_at"`
Status string `json:"status"`
Error *PluginResponsesError `json:"error"`
IncompleteDetails *PluginResponsesIncompleteDetail `json:"incomplete_details"`
Instructions any `json:"instructions"`
Model string `json:"model"`
Output []PluginResponsesOutput `json:"output"`
ParallelToolCalls bool `json:"parallel_tool_calls"`
Temperature float64 `json:"temperature"`
ToolChoice any `json:"tool_choice"`
Tools []any `json:"tools"`
TopP float64 `json:"top_p"`
Metadata map[string]string `json:"metadata"`
Usage *PluginResponsesUsage `json:"usage"`
}
type PluginResponsesError struct {
Code string `json:"code"`
Message string `json:"message"`
}
type PluginResponsesIncompleteDetail struct {
Reason string `json:"reason"`
}
type PluginResponsesUsage struct {
InputTokens int `json:"input_tokens"`
InputTokensDetails PluginResponsesInputTokenDetails `json:"input_tokens_details"`
OutputTokens int `json:"output_tokens"`
OutputTokensDetails PluginResponsesOutputTokenDetails `json:"output_tokens_details"`
TotalTokens int `json:"total_tokens"`
}
type PluginResponsesInputTokenDetails struct {
CachedTokens int `json:"cached_tokens"`
}
type PluginResponsesOutputTokenDetails struct {
ReasoningTokens int `json:"reasoning_tokens"`
}
type PluginResponsesOutput struct {
ID string `json:"id"`
Type string `json:"type"`
Status string `json:"status"`
Role string `json:"role"`
Content []PluginResponsesContent `json:"content"`
}
type PluginResponsesContent struct {
ID string `json:"id"`
Type string `json:"type"`
Text string `json:"text"`
Annotations []any `json:"annotations"`
Logprobs []any `json:"logprobs"`
}
// PluginResponsesStreamEvent contains the exact event-specific fields used by
// the host Responses state machine. Pointer fields preserve required empty
// strings and arrays on delta/done events while omitting unrelated fields.
type PluginResponsesStreamEvent struct {
Type string `json:"type"`
SequenceNumber int `json:"sequence_number"`
Response *PluginResponsesResponse `json:"response,omitempty"`
OutputIndex *int `json:"output_index,omitempty"`
ContentIndex *int `json:"content_index,omitempty"`
ItemID string `json:"item_id,omitempty"`
Item *PluginResponsesOutput `json:"item,omitempty"`
Part *PluginResponsesContent `json:"part,omitempty"`
Delta *string `json:"delta,omitempty"`
Text *string `json:"text,omitempty"`
Logprobs *[]any `json:"logprobs,omitempty"`
}
+54 -20
View File
@@ -30,26 +30,60 @@ func (t *TaskResponse[T]) IsSuccess() bool {
}
type TaskDto struct {
ID int64 `json:"id"`
CreatedAt int64 `json:"created_at"`
UpdatedAt int64 `json:"updated_at"`
TaskID string `json:"task_id"`
Platform string `json:"platform"`
UserId int `json:"user_id"`
Group string `json:"group"`
ChannelId int `json:"channel_id"`
Quota int `json:"quota"`
Action string `json:"action"`
Status string `json:"status"`
FailReason string `json:"fail_reason"`
ResultURL string `json:"result_url,omitempty"` // 任务结果 URL(视频地址等)
SubmitTime int64 `json:"submit_time"`
StartTime int64 `json:"start_time"`
FinishTime int64 `json:"finish_time"`
Progress string `json:"progress"`
Properties any `json:"properties"`
Username string `json:"username,omitempty"`
Data json.RawMessage `json:"data"`
ID int64 `json:"id"`
CreatedAt int64 `json:"created_at"`
UpdatedAt int64 `json:"updated_at"`
TaskID string `json:"task_id"`
Platform string `json:"platform"`
UserId int `json:"user_id"`
Group string `json:"group"`
ChannelId int `json:"channel_id"`
Quota int `json:"quota"`
Action string `json:"action"`
Status string `json:"status"`
FailReason string `json:"fail_reason"`
ResultURL string `json:"result_url,omitempty"` // 任务结果 URL(视频地址等)
LegacyVideoAvailable bool `json:"legacy_video_available,omitempty"`
SubmitTime int64 `json:"submit_time"`
StartTime int64 `json:"start_time"`
FinishTime int64 `json:"finish_time"`
Progress string `json:"progress"`
Properties any `json:"properties"`
Username string `json:"username,omitempty"`
Data json.RawMessage `json:"data"`
AdminInfo *TaskAdminInfo `json:"admin_info,omitempty"`
RootInfo *TaskRootInfo `json:"root_info,omitempty"`
}
type TaskPluginInfo struct {
Key string `json:"key"`
Name string `json:"name"`
Version string `json:"version,omitempty"`
Author *TaskPluginAuthorInfo `json:"author,omitempty"`
}
type TaskPluginAuthorInfo struct {
Name string `json:"name"`
URL string `json:"url,omitempty"`
}
type TaskPluginRuntimeInfo struct {
Key string `json:"key"`
Version string `json:"version"`
APIVersion int `json:"api_version"`
Generation uint64 `json:"generation"`
}
type TaskAdminInfo struct {
RequestID string `json:"request_id,omitempty"`
RequestPath string `json:"request_path,omitempty"`
TaskPlugin *TaskPluginInfo `json:"task_plugin,omitempty"`
}
type TaskRootInfo struct {
TaskPlugin *TaskPluginRuntimeInfo `json:"task_plugin,omitempty"`
UpstreamTaskID string `json:"upstream_task_id,omitempty"`
NodeName string `json:"node_name,omitempty"`
}
type FetchReq struct {
+23
View File
@@ -0,0 +1,23 @@
package dto
type TaskPluginError struct {
Code string `json:"code"`
Message string `json:"message"`
HTTPStatus int `json:"httpStatus"`
Retryable bool `json:"retryable"`
}
// TaskView is the only persisted-task shape exposed to JavaScript plugins.
// It deliberately excludes ownership, channel, quota, properties, and private
// upstream identifiers.
type TaskView struct {
TaskID string `json:"task_id"`
Platform string `json:"platform"`
Status string `json:"status"`
Progress string `json:"progress"`
FailReason string `json:"fail_reason"`
CreatedAt int64 `json:"created_at"`
UpdatedAt int64 `json:"updated_at,omitempty"`
FinishedAt int64 `json:"finished_at,omitempty"`
Data any `json:"data,omitempty"`
}
+148
View File
@@ -0,0 +1,148 @@
package e2e
import (
"bytes"
"context"
"io"
"net/http"
"net/http/httptest"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/controller"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/middleware"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/relay"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/system_setting"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)
const docParsePluginSource = `export const meta = {apiVersion:1,key:"doc-parse",name:"Document Parser",version:"1.0.0",author:{name:"Test"},models:["doc-parse-v1"],fetchMode:"batch"};
export function buildSubmitRequest(ctx){return {url:ctx.baseUrl+"/submit",method:"POST",headers:{"Content-Type":"application/json"},body:ctx.requestBody,action:"parse_document"};}
export function parseSubmitResponse(ctx,resp){if(!resp.body.id)throw new Error("missing id");return {taskId:resp.body.id,taskData:resp.body};}
export function buildBatchQueryRequest(ctx,taskIds){return {url:ctx.baseUrl+"/batch",method:"POST",headers:{"Content-Type":"application/json"},body:{ids:taskIds}};}
export function parseBatchResult(ctx,body){return body.tasks.map((task)=>({taskId:task.id,status:task.status,progress:"100%",data:task}));}
export function parseTaskResult(ctx,body){return {taskId:body.id,status:body.status};}
export function listArtifacts(task){return task.status==="SUCCESS"?(task.data.artifacts||[]).map((item)=>({key:item.key,type:"file",mimeType:item.mimeType})):[];}
export function buildContentRequest(ctx){const item=(ctx.data.artifacts||[]).find((artifact)=>artifact.key===ctx.artifactKey);if(!item)throw new Error("artifact_not_found");return {url:item.url,method:ctx.clientRequest.method,credentialless:true};}
`
func TestDocumentPluginRunsGenericBatchArtifactChain(t *testing.T) {
service.InitHttpClient()
originalDB := model.DB
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
require.NoError(t, database.AutoMigrate(&model.TaskPlugin{}, &model.Channel{}, &model.Task{}))
model.DB = database
t.Cleanup(func() { model.DB = originalDB; jsplugin.DefaultRegistry.Unregister("doc-parse") })
source := docParsePluginSource
uploadBody, err := common.Marshal(map[string]any{"source": source, "remark": "phase 4 acceptance"})
require.NoError(t, err)
uploadRecorder := httptest.NewRecorder()
uploadContext, _ := gin.CreateTestContext(uploadRecorder)
uploadContext.Request = httptest.NewRequest(http.MethodPost, "/api/plugin/task", bytes.NewReader(uploadBody))
uploadContext.Request.Header.Set("Content-Type", "application/json")
controller.UploadTaskPlugin(uploadContext)
require.Equal(t, http.StatusOK, uploadRecorder.Code, uploadRecorder.Body.String())
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
switch r.URL.Path {
case "/submit":
_, _ = io.WriteString(w, `{"id":"doc-upstream-1"}`)
case "/batch":
_, _ = io.WriteString(w, `{"tasks":[{"id":"doc-upstream-1","status":"SUCCESS","artifacts":[{"key":"text","url":"`+"http://"+r.Host+`/artifact/text","mimeType":"text/plain"},{"key":"json","url":"`+"http://"+r.Host+`/artifact/json","mimeType":"application/json"}]}]}`)
case "/artifact/text":
w.Header().Set("Content-Type", "text/plain")
_, _ = io.WriteString(w, "parsed text")
case "/artifact/json":
_, _ = io.WriteString(w, `{"pages":2}`)
default:
http.NotFound(w, r)
}
}))
defer upstream.Close()
setting := dto.ChannelSettings{TaskPluginKey: "doc-parse"}
channel := model.Channel{Type: constant.ChannelTypeTaskPlugin, Name: "documents", Key: "unused", BaseURL: &upstream.URL, Status: common.ChannelStatusEnabled, Models: "doc-parse-v1", Group: "default"}
channel.SetSetting(setting)
require.NoError(t, database.Create(&channel).Error)
adaptor := relay.GetTaskAdaptor("doc-parse")
require.NotNil(t, adaptor)
info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ChannelType: channel.Type, ChannelBaseUrl: upstream.URL, ApiKey: channel.Key, ChannelSetting: setting, UpstreamModelName: "doc-parse-v1"}, OriginModelName: "doc-parse-v1", TaskRelayInfo: &relaycommon.TaskRelayInfo{PublicTaskID: "task_doc_parse"}}
adaptor.Init(info)
submitRecorder := httptest.NewRecorder()
submitContext, _ := gin.CreateTestContext(submitRecorder)
submitContext.Request = httptest.NewRequest(http.MethodPost, "/v1/tasks/doc-parse", bytes.NewBufferString(`{"model":"doc-parse-v1","document":"opaque-ref"}`))
submitContext.Request.Header.Set("Content-Type", "application/json")
submitContext.Params = gin.Params{{Key: "key", Value: "doc-parse"}}
middleware.PrepareTaskPluginSubmit()(submitContext)
require.Empty(t, submitRecorder.Body.String())
require.Equal(t, "doc-parse-v1", submitContext.GetString("resolved_task_model"))
require.Nil(t, adaptor.ValidateRequestAndSetAction(submitContext, info))
require.Equal(t, "parse_document", info.Action)
requestBody, err := adaptor.BuildRequestBody(submitContext, info)
require.NoError(t, err)
response, err := adaptor.DoRequest(submitContext, info, requestBody)
require.NoError(t, err)
parsed, taskErr := adaptor.ParseResponse(submitContext, response, info)
require.Nil(t, taskErr)
require.NotNil(t, parsed)
require.Equal(t, "doc-upstream-1", parsed.UpstreamTaskID)
task := model.Task{
TaskID: info.PublicTaskID, Platform: "doc-parse", UserId: 7, ChannelId: channel.Id,
Status: model.TaskStatusInProgress, Data: parsed.TaskData,
PrivateData: model.TaskPrivateData{
UpstreamTaskID: parsed.UpstreamTaskID,
Execution: &model.TaskExecutionSnapshot{TaskPlugin: &model.TaskPluginSnapshot{
Key: "doc-parse", Name: "Document Parser", Version: "1.0.0",
Author: &model.TaskPluginAuthorSnapshot{Name: "Test"}, APIVersion: 1,
}},
},
}
require.NoError(t, database.Create(&task).Error)
originalFactory := service.GetTaskAdaptorFunc
service.GetTaskAdaptorFunc = func(platform constant.TaskPlatform) service.TaskPollingAdaptor { return relay.GetTaskAdaptor(platform) }
t.Cleanup(func() { service.GetTaskAdaptorFunc = originalFactory })
service.DispatchPlatformUpdate(context.Background(), "doc-parse", map[int][]string{channel.Id: {parsed.UpstreamTaskID}}, map[string]*model.Task{parsed.UpstreamTaskID: &task})
require.NoError(t, database.First(&task, task.ID).Error)
assert.Equal(t, model.TaskStatus(model.TaskStatusSuccess), task.Status)
queryRecorder := httptest.NewRecorder()
queryContext, _ := gin.CreateTestContext(queryRecorder)
queryContext.Set("id", 7)
queryContext.Params = gin.Params{{Key: "key", Value: task.TaskID}}
queryContext.Request = httptest.NewRequest(http.MethodGet, "/v1/tasks/"+task.TaskID+"/artifacts", nil)
controller.GetTaskArtifacts(queryContext)
require.Equal(t, http.StatusOK, queryRecorder.Code)
var query struct {
Artifacts []map[string]any `json:"artifacts"`
}
require.NoError(t, common.Unmarshal(queryRecorder.Body.Bytes(), &query))
require.Len(t, query.Artifacts, 2)
originalFetch := *system_setting.GetFetchSetting()
system_setting.GetFetchSetting().EnableSSRFProtection = true
system_setting.GetFetchSetting().AllowPrivateIp = true
system_setting.GetFetchSetting().AllowedPorts = []string{"1-65535"}
t.Cleanup(func() { *system_setting.GetFetchSetting() = originalFetch })
contentRecorder := httptest.NewRecorder()
contentContext, _ := gin.CreateTestContext(contentRecorder)
contentContext.Set("id", 7)
contentContext.Params = gin.Params{{Key: "key", Value: task.TaskID}, {Key: "artifact_key", Value: "text"}}
contentContext.Request = httptest.NewRequest(http.MethodGet, "/v1/tasks/"+task.TaskID+"/artifacts/text/content", nil)
controller.TaskArtifactContent(contentContext)
assert.Equal(t, http.StatusOK, contentRecorder.Code)
assert.Equal(t, "parsed text", contentRecorder.Body.String())
}
+6 -1
View File
@@ -70,8 +70,11 @@ require (
github.com/ClickHouse/ch-go v0.65.0 // indirect
github.com/bmatcuk/doublestar/v4 v4.6.1 // indirect
github.com/casbin/govaluate v1.10.0 // indirect
github.com/dlclark/regexp2/v2 v2.2.2 // indirect
github.com/go-faster/city v1.0.1 // indirect
github.com/go-faster/errors v0.7.1 // indirect
github.com/go-sourcemap/sourcemap v2.1.3+incompatible // indirect
github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e // indirect
github.com/hashicorp/go-version v1.7.0 // indirect
github.com/paulmach/orb v0.11.1 // indirect
github.com/pierrec/lz4/v4 v4.1.22 // indirect
@@ -85,6 +88,8 @@ require (
require (
github.com/Azure/go-ntlmssp v0.1.1
github.com/alicebob/miniredis/v2 v2.38.0
github.com/grafana/sobek v0.0.0-20260708062710-267a0e055bb4
github.com/openai/openai-go v1.12.0
)
require (
@@ -146,7 +151,7 @@ require (
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
github.com/samber/go-singleflightx v0.3.2 // indirect
github.com/tidwall/match v1.1.1 // indirect
github.com/tidwall/pretty v1.2.0 // indirect
github.com/tidwall/pretty v1.2.1 // indirect
github.com/tklauser/go-sysconf v0.3.12 // indirect
github.com/tklauser/numcpus v0.6.1 // indirect
github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
+14 -1
View File
@@ -643,6 +643,8 @@ github.com/ClickHouse/clickhouse-go/v2 v2.32.0/go.mod h1:rGFIgeNbJVggBp2C+0FXOdf
github.com/DmitriyVTitov/size v1.5.0 h1:/PzqxYrOyOUX1BXj6J9OuVRVGe+66VL4D9FlUaW515g=
github.com/DmitriyVTitov/size v1.5.0/go.mod h1:le6rNI4CoLQV1b9gzp1+3d7hMAD/uu2QcJ+aYbNgiU0=
github.com/JohnCGriffin/overflow v0.0.0-20211019200055-46fa312c352c/go.mod h1:X0CRv0ky0k6m906ixxpzmDRLvX58TFUKS2eePweuyxk=
github.com/Masterminds/semver/v3 v3.5.0 h1:kQceYJfbupGfZOKZQg0kou0DgAKhzDg2NZPAwZ/2OOE=
github.com/Masterminds/semver/v3 v3.5.0/go.mod h1:4V+yj/TJE1HU9XfppCwVMZq3I84lprf4nC11bSS5beM=
github.com/Microsoft/go-winio v0.4.11/go.mod h1:VhR8bwka0BXejwEJY73c50VrPtXAaKcyvVC4A4RozmA=
github.com/Microsoft/go-winio v0.4.14/go.mod h1:qXqCSQ3Xa7+6tgxaGTIe4Kpcdsi+P8jBhyzoq1bpyYA=
github.com/Microsoft/go-winio v0.4.15-0.20190919025122-fc70bd9a86b5/go.mod h1:tTuCMEN+UleMWgg9dVx4Hu52b1bJo+59jBh3ajtinzw=
@@ -1010,6 +1012,8 @@ github.com/dgryski/go-sip13 v0.0.0-20181026042036-e10d5fee7954/go.mod h1:vAd38F8
github.com/distribution/distribution/v3 v3.0.0-20220526142353-ffbd94cbe269/go.mod h1:28YO/VJk9/64+sTGNuYaBjWxrXTPrj0C0XmgTIOjxX4=
github.com/dlclark/regexp2 v1.11.5 h1:Q/sSnsKerHeCkc/jSTNq1oCm7KiVgUMZRDUoRu0JQZQ=
github.com/dlclark/regexp2 v1.11.5/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
github.com/dlclark/regexp2/v2 v2.2.2 h1:MYWvNYw8okuqNhwTYO587EZMiDruVa2vhV6fsGpfya0=
github.com/dlclark/regexp2/v2 v2.2.2/go.mod h1:avUrQvPaLz2DrFNHJF0taWAFFX2C1GMSSoeiqFjcBmU=
github.com/dmarkham/enumer v1.5.8/go.mod h1:d10o8R3t/gROm2p3BXqTkMt2+HMuxEmWCXzorAruYak=
github.com/dnaeon/go-vcr v1.0.1/go.mod h1:aBB1+wY4s93YsC3HHjMBMrwTj2R9FHDzUr9KyGc8n1E=
github.com/docker/cli v0.0.0-20191017083524-a8ff7f821017/go.mod h1:JLrzqnKDaYBop7H2jaqPtU4hHvMKP+vjCwu2uszcLI8=
@@ -1196,6 +1200,8 @@ github.com/go-playground/validator/v10 v10.20.0 h1:K9ISHbSaI0lyB2eWMPJo+kOS/FBEx
github.com/go-playground/validator/v10 v10.20.0/go.mod h1:dbuPbCMFw/DrkbEynArYaCwl3amGuJotoKCe95atGMM=
github.com/go-redis/redis/v8 v8.11.5 h1:AcZZR7igkdvfVmQTPnu9WE37LRrO/YrBH5zWyjDC0oI=
github.com/go-redis/redis/v8 v8.11.5/go.mod h1:gREzHqY1hg6oD9ngVRbLStwAWKhA0FEgq8Jd4h5lpwo=
github.com/go-sourcemap/sourcemap v2.1.3+incompatible h1:W1iEw64niKVGogNgBN3ePyLFfuisuzeidWPMPWmECqU=
github.com/go-sourcemap/sourcemap v2.1.3+incompatible/go.mod h1:F8jJfvm2KbVjc5NqelyYJmf/v5J0dwNLS2mL4sNA1Jg=
github.com/go-sql-driver/mysql v1.4.0/go.mod h1:zAC/RDZ24gD3HViQzih4MyKcchzm+sOG5ZlKdlhCg5w=
github.com/go-sql-driver/mysql v1.6.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg=
github.com/go-sql-driver/mysql v1.7.0 h1:ueSltNNllEqE3qcWBTD0iQd3IpL/6U+mJxLkazJ7YPc=
@@ -1211,6 +1217,8 @@ github.com/goccy/go-json v0.9.7/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGF
github.com/goccy/go-json v0.9.11/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=
github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU=
github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=
github.com/goccy/go-yaml v1.19.2 h1:PmFC1S6h8ljIz6gMRBopkjP1TVT7xuwrButHID66PoM=
github.com/goccy/go-yaml v1.19.2/go.mod h1:XBurs7gK8ATbW4ZPGKgcbrY1Br56PdM69F7LkFRi1kA=
github.com/godbus/dbus v0.0.0-20151105175453-c7fdd8b5cd55/go.mod h1:/YcGZj5zSblfDWMMoOzV4fas9FZnQYTkDnsGvmh2Grw=
github.com/godbus/dbus v0.0.0-20180201030542-885f9cc04c9c/go.mod h1:/YcGZj5zSblfDWMMoOzV4fas9FZnQYTkDnsGvmh2Grw=
github.com/godbus/dbus v0.0.0-20190422162347-ade71ed3457e/go.mod h1:bBOAhwG1umN6/6ZUMtDFBMQR8jRg9O75tm9K00oMsK4=
@@ -1375,6 +1383,8 @@ github.com/grafana/pyroscope-go v1.2.7 h1:VWBBlqxjyR0Cwk2W6UrE8CdcdD80GOFNutj0Kb
github.com/grafana/pyroscope-go v1.2.7/go.mod h1:o/bpSLiJYYP6HQtvcoVKiE9s5RiNgjYTj1DhiddP2Pc=
github.com/grafana/pyroscope-go/godeltaprof v0.1.9 h1:c1Us8i6eSmkW+Ez05d3co8kasnuOY813tbMN8i/a3Og=
github.com/grafana/pyroscope-go/godeltaprof v0.1.9/go.mod h1:2+l7K7twW49Ct4wFluZD3tZ6e0SjanjcUUBPVD/UuGU=
github.com/grafana/sobek v0.0.0-20260708062710-267a0e055bb4 h1:LJ2pEOxFbfUIhrkpROwZ6hLhuCn6e5GSX08v5GFwN/4=
github.com/grafana/sobek v0.0.0-20260708062710-267a0e055bb4/go.mod h1:BL/2XROA/Wtlb+zGEhkSdSZiMFIgw+D2ZdDfrbccyVE=
github.com/gregjones/httpcache v0.0.0-20180305231024-9cad4c3443a7/go.mod h1:FecbI9+v66THATjSRHfNgh1IVFe/9kFxbXtjV0ctIMA=
github.com/grpc-ecosystem/go-grpc-middleware v1.0.0/go.mod h1:FiyG127CGDf3tlThmgyCl78X/SZQqEOJBCDaAfeWzPs=
github.com/grpc-ecosystem/go-grpc-middleware v1.0.1-0.20190118093823-f849b5445de4/go.mod h1:FiyG127CGDf3tlThmgyCl78X/SZQqEOJBCDaAfeWzPs=
@@ -1677,6 +1687,8 @@ github.com/onsi/gomega v1.24.1/go.mod h1:3AOiACssS3/MajrniINInwbfOOtfZvplPzuRSmv
github.com/onsi/gomega v1.24.2 h1:J/tulyYK6JwBldPViHJReihxxZ+22FHs0piGjQAvoUE=
github.com/onsi/gomega v1.24.2/go.mod h1:gs3J10IS7Z7r7eXRoNJIrNqU4ToQukCJhFtKrWgHWnk=
github.com/open-policy-agent/opa v0.42.2/go.mod h1:MrmoTi/BsKWT58kXlVayBb+rYVeaMwuBm3nYAN3923s=
github.com/openai/openai-go v1.12.0 h1:NBQCnXzqOTv5wsgNC36PrFEiskGfO5wccfCWDo9S1U0=
github.com/openai/openai-go v1.12.0/go.mod h1:g461MYGXEXBVdV5SaR/5tNzNbSfwTBBefwc+LlDCK0Y=
github.com/opencontainers/go-digest v0.0.0-20170106003457-a6d0ee40d420/go.mod h1:cMLVZDEM3+U2I4VmLI6N8jQYUd2OVphdqWwCJHrFt2s=
github.com/opencontainers/go-digest v0.0.0-20180430190053-c9281466c8b2/go.mod h1:cMLVZDEM3+U2I4VmLI6N8jQYUd2OVphdqWwCJHrFt2s=
github.com/opencontainers/go-digest v1.0.0-rc1/go.mod h1:cMLVZDEM3+U2I4VmLI6N8jQYUd2OVphdqWwCJHrFt2s=
@@ -1951,8 +1963,9 @@ github.com/tidwall/gjson v1.19.0/go.mod h1:V37/opeE/JbLUOfH0QTXiNez2l0RUjYUhpT4s
github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA=
github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM=
github.com/tidwall/pretty v1.0.0/go.mod h1:XNkn88O1ChpSDQmQeStsy+sBenx6DDtFZJxhVysOjyk=
github.com/tidwall/pretty v1.2.0 h1:RWIZEg2iJ8/g6fDDYzMpobmaoGh5OLl4AXtGUGPcqCs=
github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4=
github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY=
github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28=
github.com/tiktoken-go/tokenizer v0.6.2 h1:t0GN2DvcUZSFWT/62YOgoqb10y7gSXBGs0A+4VCQK+g=
+4 -1
View File
@@ -77,7 +77,10 @@ func LogInfo(ctx context.Context, msg string) {
logHelper(ctx, loggerINFO, msg)
}
func LogWarn(ctx context.Context, msg string) {
func LogWarn(ctx context.Context, msg string, args ...any) {
if len(args) > 0 {
msg = fmt.Sprintf(msg, args...)
}
logHelper(ctx, loggerWarn, msg)
}
+5
View File
@@ -23,6 +23,7 @@ import (
"github.com/QuantumNous/new-api/middleware"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/oauth"
"github.com/QuantumNous/new-api/pkg/jsplugin"
perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
"github.com/QuantumNous/new-api/relay"
kitutil "github.com/QuantumNous/new-api/relaykit/relayconvert/kitutil"
@@ -46,6 +47,9 @@ var buildFS embed.FS
var indexPage []byte
func main() {
if len(os.Args) > 1 && os.Args[1] == "plugin" {
os.Exit(jsplugin.RunCLI(os.Args[2:], os.Stdout, os.Stderr))
}
startTime := time.Now()
kitutil.SetLogging(common.SysLog, func(message string) {
logger.LogError(nil, message)
@@ -107,6 +111,7 @@ func main() {
// 热更新配置
go model.SyncOptions(common.SyncFrequency)
go controller.SyncTaskPlugins()
// 周期性重载授权策略,保证多节点/多 master 部署下权限变更能传播到每个实例
go authz.StartPolicySync(common.SyncFrequency)
+13 -1
View File
@@ -5,10 +5,12 @@ import (
"fmt"
"net"
"net/http"
"strconv"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
@@ -515,7 +517,17 @@ func SetupContextForToken(c *gin.Context, token *model.Token, parts ...string) e
}
if len(parts) > 1 {
if model.IsAdmin(token.UserId) {
c.Set("specific_channel_id", parts[1])
id, err := strconv.Atoi(parts[1])
if err != nil {
abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidChannelId))
return fmt.Errorf("invalid specific channel id")
}
service.GetChannelConstraints(c).AddPin(dto.ChannelPin{
ChannelId: id,
Source: dto.PinSourceToken,
Rank: dto.PinRankToken,
RetryMode: dto.PinRetrySingleAttempt,
})
} else {
c.Header("specific_channel_version", "701e3ae1dc3f7975556d354e0675168d004891c8")
abortWithOpenAiMessage(c, http.StatusForbidden, "普通用户不支持指定渠道")
+7 -7
View File
@@ -10,13 +10,13 @@ import (
// 在请求处理完成后自动清理磁盘/内存缓存
func BodyStorageCleanup() gin.HandlerFunc {
return func(c *gin.Context) {
// 处理请求
defer func() {
// 请求结束后清理存储
common.CleanupBodyStorage(c)
// 清理文件缓存(URL 下载的文件等)
service.CleanupFileSources(c)
}()
c.Next()
// 请求结束后清理存储
common.CleanupBodyStorage(c)
// 清理文件缓存(URL 下载的文件等)
service.CleanupFileSources(c)
}
}
+212 -35
View File
@@ -6,7 +6,6 @@ import (
"io"
"net/http"
"slices"
"strconv"
"strings"
"time"
@@ -14,7 +13,9 @@ import (
"github.com/QuantumNous/new-api/constant"
taskdto "github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/jsplugin"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
@@ -33,25 +34,46 @@ type ModelRequest struct {
func Distribute() func(c *gin.Context) {
return func(c *gin.Context) {
var channel *model.Channel
channelId, ok := common.GetContextKey(c, constant.ContextKeyTokenSpecificChannelId)
constraints := service.GetChannelConstraints(c)
constraints.AddFilter(taskdto.ChannelFilter{
Kind: taskdto.FilterRequestPath,
RequestPath: c.Request.URL.Path,
})
service.AppendTaskPluginIdentityFilter(c, c.GetString("expected_task_plugin_key"))
modelRequest, shouldSelectChannel, err := getModelRequest(c)
if err != nil {
abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidRequest, map[string]any{"Error": err.Error()}))
return
}
if ok {
id, err := strconv.Atoi(channelId.(string))
if err != nil {
abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidChannelId))
return
if pin, found, overridden := constraints.ResolvedPin(); found {
for _, lost := range overridden {
logger.LogWarn(c, fmt.Sprintf(
"channel pin overridden: winning_source=%s winning_channel_id=%d overridden_source=%s overridden_channel_id=%d",
pin.Source, pin.ChannelId, lost.Source, lost.ChannelId,
))
}
channel, err = model.GetChannelById(id, true)
channel, err = model.CacheGetChannel(pin.ChannelId)
if err != nil {
abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidChannelId))
if pin.Source == taskdto.PinSourceOriginTask {
abortWithOpenAiMessage(c, http.StatusBadRequest, "origin_task_channel_disabled", types.ErrorCode("origin_task_channel_disabled"))
} else {
abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidChannelId))
}
return
}
if channel.Status != common.ChannelStatusEnabled {
abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorChannelDisabled))
if pin.Source == taskdto.PinSourceOriginTask {
abortWithOpenAiMessage(c, http.StatusBadRequest, "origin_task_channel_disabled", types.ErrorCode("origin_task_channel_disabled"))
} else {
abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorChannelDisabled))
}
return
}
if ok, kind := model.ChannelSatisfiesFilters(channel, modelRequest.Model, constraints.Filters); !ok {
if kind == taskdto.FilterTaskPluginIdentity {
logTaskPluginChannelDecision(c, channel, modelRequest.Model, "channel_rejected", "identity_mismatch")
}
abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorNoAvailableChannel, map[string]any{"Group": common.GetContextKeyString(c, constant.ContextKeyUsingGroup), "Model": modelRequest.Model}), types.ErrorCode(kind))
return
}
} else {
@@ -105,8 +127,11 @@ func Distribute() func(c *gin.Context) {
if preferredChannelID, found := service.GetPreferredChannelByAffinity(c, modelRequest.Model, usingGroup); found {
affinityUsable := false
preferred, err := model.CacheGetChannel(preferredChannelID)
if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled &&
channelSupportsRequestPath(preferred, c.Request.URL.Path, modelRequest.Model) {
affinitySatisfied := false
if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled {
affinitySatisfied, _ = model.ChannelSatisfiesFilters(preferred, modelRequest.Model, constraints.Filters)
}
if affinitySatisfied {
if usingGroup == "auto" {
userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
autoGroups := service.GetRequestAutoGroups(c, userGroup)
@@ -161,6 +186,15 @@ func Distribute() func(c *gin.Context) {
}
}
}
if channel != nil {
if ok, kind := model.ChannelSatisfiesFilters(channel, modelRequest.Model, constraints.Filters); !ok {
if kind == taskdto.FilterTaskPluginIdentity {
logTaskPluginChannelDecision(c, channel, modelRequest.Model, "channel_rejected", "identity_mismatch")
}
abortWithOpenAiMessage(c, http.StatusServiceUnavailable, i18n.T(c, i18n.MsgDistributorNoAvailableChannel, map[string]any{"Group": common.GetContextKeyString(c, constant.ContextKeyUsingGroup), "Model": modelRequest.Model}), types.ErrorCodeModelNotFound)
return
}
}
common.SetContextKey(c, constant.ContextKeyRequestStartTime, time.Now())
SetupContextForSelectedChannel(c, channel, modelRequest.Model)
c.Next()
@@ -170,18 +204,68 @@ func Distribute() func(c *gin.Context) {
}
}
// channelSupportsRequestPath reports whether a channel can serve the request path.
// Only Advanced Custom (type 58) channels are path-checked; all other channel types
// always pass. A type-58 channel is usable only when one of its routes matches.
func channelSupportsRequestPath(channel *model.Channel, requestPath string, requestModel string) bool {
func channelMatchesExpectedTaskPlugin(c *gin.Context, channel *model.Channel, expected string) bool {
if channel == nil {
return false
}
if channel.Type != constant.ChannelTypeAdvancedCustom {
if c != nil {
if _, matched := pinnedEndpointCandidateForChannel(c, channel, expected); matched {
return true
}
}
if channel.Type == constant.ChannelTypeTaskPlugin {
return expected != "" && channel.GetSetting().TaskPluginKey == expected
}
if expected == "" {
return true
}
config := channel.GetOtherSettings().AdvancedCustom
return config != nil && config.SupportsPathForModel(requestPath, requestModel)
if c == nil {
return false
}
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 false
}
plugin, ok := pinned.Generation.GetByChannelType(channel.Type)
return ok && plugin == pinned.Plugin
}
func pinnedEndpointCandidateForChannel(c *gin.Context, channel *model.Channel, expected string) (jsplugin.ProtocolBinding, bool) {
if c == nil || channel == nil || expected == "" {
return jsplugin.ProtocolBinding{}, false
}
value, exists := c.Get(jsplugin.ContextKeyPinnedEndpoint)
pinned, ok := value.(jsplugin.PinnedEndpoint)
if !exists || !ok || pinned.Generation == nil || pinned.Plugin == nil {
return jsplugin.ProtocolBinding{}, false
}
candidates := pinned.Candidates
if len(candidates) == 0 {
candidates = []jsplugin.ProtocolBinding{{Plugin: pinned.Plugin, Protocol: pinned.Protocol, Operation: pinned.Operation, Model: pinned.Model}}
}
expectedOwned := false
selected := jsplugin.ProtocolBinding{}
for _, candidate := range candidates {
if candidate.Plugin == nil {
continue
}
if candidate.Plugin.Meta.Key == expected {
expectedOwned = true
}
if channel.Type == constant.ChannelTypeTaskPlugin {
if channel.GetSetting().TaskPluginKey == candidate.Plugin.Meta.Key {
selected = candidate
}
continue
}
plugin, indexed := pinned.Generation.GetByChannelType(channel.Type)
if indexed && plugin == candidate.Plugin {
selected = candidate
}
}
return selected, expectedOwned && selected.Plugin != nil
}
// getModelFromRequest 从请求中读取模型信息
@@ -190,6 +274,12 @@ func channelSupportsRequestPath(channel *model.Channel, requestPath string, requ
// - application/x-www-form-urlencoded
// - multipart/form-data
func getModelFromRequest(c *gin.Context) (*ModelRequest, error) {
if cached, exists := c.Get(contextKeyTaskPluginEndpointModel); exists {
if modelRequest, ok := cached.(ModelRequest); ok {
cachedRequest := modelRequest
return &cachedRequest, nil
}
}
if strings.HasPrefix(c.Request.Header.Get("Content-Type"), "application/json") {
modelRequest, err := getModelFromJSONBody(c)
if err != nil {
@@ -218,6 +308,9 @@ func getModelFromJSONBody(c *gin.Context) (*ModelRequest, error) {
if !gjson.ValidBytes(requestBody) {
return nil, errors.New("invalid JSON request body")
}
if countTopLevelJSONKey(requestBody, "model") > 1 {
return nil, errors.New("model must be provided once")
}
values := gjson.GetManyBytes(requestBody, "model", "group")
model, err := getJSONStringValue(values[0], "model")
@@ -240,6 +333,64 @@ func getModelFromJSONBody(c *gin.Context) (*ModelRequest, error) {
}, nil
}
func countTopLevelJSONKey(data []byte, target string) int {
depth := 0
inString := false
escaped := false
stringStart := 0
expectingKey := false
count := 0
for index, current := range data {
if inString {
if escaped {
escaped = false
continue
}
if current == '\\' {
escaped = true
continue
}
if current != '"' {
continue
}
inString = false
if depth == 1 && expectingKey {
key := string(data[stringStart:index])
var decodedKey string
if common.Unmarshal(data[stringStart-1:index+1], &decodedKey) == nil {
key = decodedKey
}
cursor := index + 1
for cursor < len(data) && (data[cursor] == ' ' || data[cursor] == '\t' || data[cursor] == '\r' || data[cursor] == '\n') {
cursor++
}
if cursor < len(data) && data[cursor] == ':' && key == target {
count++
}
expectingKey = false
}
continue
}
switch current {
case '"':
inString = true
stringStart = index + 1
case '{':
depth++
if depth == 1 {
expectingKey = true
}
case '}':
depth--
case ',':
if depth == 1 {
expectingKey = true
}
}
}
return count
}
func getJSONStringValue(result gjson.Result, field string) (string, error) {
if !result.Exists() || result.Type == gjson.Null {
return "", nil
@@ -254,7 +405,9 @@ func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) {
var modelRequest ModelRequest
shouldSelectChannel := true
var err error
if strings.Contains(c.Request.URL.Path, "/mj/") {
if modelName := c.GetString("resolved_task_model"); modelName != "" {
modelRequest.Model = modelName
} else if strings.Contains(c.Request.URL.Path, "/mj/") {
relayMode := relayconstant.Path2RelayModeMidjourney(c.Request.URL.Path)
if relayMode == relayconstant.RelayModeMidjourneyTaskFetch ||
relayMode == relayconstant.RelayModeMidjourneyTaskFetchByCondition ||
@@ -282,17 +435,6 @@ func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) {
modelRequest.Model = midjourneyModel
}
c.Set("relay_mode", relayMode)
} else if strings.Contains(c.Request.URL.Path, "/suno/") {
relayMode := relayconstant.Path2RelaySuno(c.Request.Method, c.Request.URL.Path)
if relayMode == relayconstant.RelayModeSunoFetch ||
relayMode == relayconstant.RelayModeSunoFetchByID {
shouldSelectChannel = false
} else {
modelName := service.CoverTaskActionToModelName(constant.TaskPlatformSuno, c.Param("action"))
modelRequest.Model = modelName
}
c.Set("platform", string(constant.TaskPlatformSuno))
c.Set("relay_mode", relayMode)
} else if strings.Contains(c.Request.URL.Path, "/v1/videos/") && strings.HasSuffix(c.Request.URL.Path, "/remix") {
relayMode := relayconstant.RelayModeVideoSubmit
c.Set("relay_mode", relayMode)
@@ -423,10 +565,6 @@ func getTaskOriginModelName(c *gin.Context) string {
}
taskId := c.Param("task_id")
if taskId == "" {
// jimeng adapter
taskId = c.GetString("task_id")
}
if taskId == "" {
return ""
}
@@ -440,15 +578,54 @@ func getTaskOriginModelName(c *gin.Context) string {
func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, modelName string) *types.NewAPIError {
c.Set("original_model", modelName) // for retry
expectedPlugin := c.GetString("expected_task_plugin_key")
if channel == nil {
logTaskPluginChannelDecision(c, nil, modelName, "channel_rejected", "nil_channel")
return types.NewError(errors.New("channel is nil"), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
}
if expectedPlugin != "" && !channelMatchesExpectedTaskPlugin(c, channel, expectedPlugin) {
logTaskPluginChannelDecision(c, channel, modelName, "channel_rejected", "identity_mismatch")
return types.NewError(
errors.New("selected channel does not match the pinned task plugin"),
types.ErrorCodeGetChannelFailed,
types.ErrOptionWithSkipRetry(),
)
}
if candidate, matched := pinnedEndpointCandidateForChannel(c, channel, expectedPlugin); matched {
if value, exists := c.Get(jsplugin.ContextKeyPinnedEndpoint); exists {
if pinned, ok := value.(jsplugin.PinnedEndpoint); ok && candidate.Plugin != nil && candidate.Plugin != pinned.Plugin {
previousPlugin := pinned.Plugin.Meta.Key
pinned.Plugin = candidate.Plugin
pinned.Protocol = candidate.Protocol
pinned.Operation = candidate.Operation
c.Set(jsplugin.ContextKeyPinnedEndpoint, pinned)
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{Generation: pinned.Generation, Plugin: candidate.Plugin})
c.Set("expected_task_plugin_key", candidate.Plugin.Meta.Key)
c.Set("task_plugin_key", candidate.Plugin.Meta.Key)
c.Set("platform", candidate.Plugin.Meta.Key)
logger.LogDebug(
c,
"task_plugin subsystem=endpoint event=provider_selected generation=%d previous_plugin=%q plugin=%q model=%q channel_id=%d channel_type=%d",
pinned.Generation.Number,
previousPlugin,
candidate.Plugin.Meta.Key,
modelName,
channel.Id,
channel.Type,
)
}
}
}
common.SetContextKey(c, constant.ContextKeyChannelId, channel.Id)
common.SetContextKey(c, constant.ContextKeyChannelName, channel.Name)
common.SetContextKey(c, constant.ContextKeyChannelType, channel.Type)
common.SetContextKey(c, constant.ContextKeyChannelCreateTime, channel.CreatedTime)
common.SetContextKey(c, constant.ContextKeyChannelSetting, channel.GetSetting())
common.SetContextKey(c, constant.ContextKeyChannelOtherSetting, channel.GetOtherSettings())
if channel.Type == constant.ChannelTypeTaskPlugin {
c.Set("task_plugin_key", channel.GetSetting().TaskPluginKey)
}
logTaskPluginChannelDecision(c, channel, modelName, "channel_selected", "")
paramOverride := channel.GetParamOverride()
headerOverride := channel.GetHeaderOverride()
if mergedParam, applied := service.ApplyChannelAffinityOverrideTemplate(c, paramOverride); applied {
+149
View File
@@ -0,0 +1,149 @@
package middleware
import (
"fmt"
"testing"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestChannelMatchesExpectedTaskPluginUsesGenericChannelSetting(t *testing.T) {
channel := &model.Channel{Type: constant.ChannelTypeTaskPlugin}
channel.SetSetting(dto.ChannelSettings{TaskPluginKey: "generic-alpha"})
assert.True(t, channelMatchesExpectedTaskPlugin(nil, channel, "generic-alpha"))
assert.False(t, channelMatchesExpectedTaskPlugin(nil, channel, "generic-beta"))
assert.False(t, channelMatchesExpectedTaskPlugin(nil, channel, ""))
}
func TestChannelMatchesExpectedTaskPluginUsesPinnedLegacyIndex(t *testing.T) {
registry := jsplugin.NewRegistry()
alpha, err := registry.Register(distributorTaskPluginSource("legacy-alpha", constant.ChannelTypeKling), jsplugin.Options{})
require.NoError(t, err)
pinnedGeneration := registry.Generation()
require.NoError(t, registry.Unregister("legacy-alpha"))
_, err = registry.Register(distributorTaskPluginSource("legacy-beta", constant.ChannelTypeKling), jsplugin.Options{})
require.NoError(t, err)
c, _ := gin.CreateTestContext(nil)
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{
Generation: pinnedGeneration,
Plugin: alpha,
})
channel := &model.Channel{Type: constant.ChannelTypeKling}
assert.True(t, channelMatchesExpectedTaskPlugin(c, channel, "legacy-alpha"))
assert.False(t, channelMatchesExpectedTaskPlugin(c, channel, "legacy-beta"))
assert.False(t, channelMatchesExpectedTaskPlugin(c, &model.Channel{Type: constant.ChannelTypeJimeng}, "legacy-alpha"))
}
func TestChannelMatchesExpectedTaskPluginRejectsUnindexedLegacyChannel(t *testing.T) {
registry := jsplugin.NewRegistry()
plugin, err := registry.Register(distributorTaskPluginSource("legacy-alpha", constant.ChannelTypeKling), jsplugin.Options{})
require.NoError(t, err)
c, _ := gin.CreateTestContext(nil)
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{
Generation: registry.Generation(),
Plugin: plugin,
})
assert.False(t, channelMatchesExpectedTaskPlugin(c, &model.Channel{Type: constant.ChannelTypeJimeng}, "legacy-alpha"))
assert.False(t, channelMatchesExpectedTaskPlugin(c, &model.Channel{Type: 0}, "legacy-alpha"))
assert.True(t, channelMatchesExpectedTaskPlugin(c, &model.Channel{Type: constant.ChannelTypeJimeng}, ""))
assert.False(t, channelMatchesExpectedTaskPlugin(nil, &model.Channel{Type: constant.ChannelTypeKling}, "legacy-alpha"))
c.Set("expected_task_plugin_key", "legacy-alpha")
setupErr := SetupContextForSelectedChannel(c, &model.Channel{Type: constant.ChannelTypeJimeng}, "task-model")
require.NotNil(t, setupErr)
assert.Contains(t, setupErr.Error(), "does not match")
}
func TestSharedEndpointRebindsToSelectedLegacyProvider(t *testing.T) {
registry := jsplugin.NewRegistry()
_, err := registry.Register(distributorEndpointPluginSource("gemini-shared", constant.ChannelTypeGemini), jsplugin.Options{})
require.NoError(t, err)
_, err = registry.Register(distributorEndpointPluginSource("vertex-shared", 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,
})
c.Set("expected_task_plugin_key", candidates[0].Plugin.Meta.Key)
geminiChannel := &model.Channel{Id: 1, Type: constant.ChannelTypeGemini}
vertexChannel := &model.Channel{Id: 2, Type: constant.ChannelTypeVertexAi}
assert.True(t, channelMatchesExpectedTaskPlugin(c, geminiChannel, candidates[0].Plugin.Meta.Key))
assert.True(t, channelMatchesExpectedTaskPlugin(c, vertexChannel, candidates[0].Plugin.Meta.Key))
assert.False(t, channelMatchesExpectedTaskPlugin(c, &model.Channel{Type: constant.ChannelTypeKling}, candidates[0].Plugin.Meta.Key))
require.Nil(t, SetupContextForSelectedChannel(c, vertexChannel, "task-model"))
pinnedValue, exists := c.Get(jsplugin.ContextKeyPinnedEndpoint)
require.True(t, exists)
pinned, ok := pinnedValue.(jsplugin.PinnedEndpoint)
require.True(t, ok)
assert.Equal(t, "vertex-shared", pinned.Plugin.Meta.Key)
assert.Equal(t, "vertex-shared", c.GetString("expected_task_plugin_key"))
assert.Equal(t, "vertex-shared", c.GetString("task_plugin_key"))
assert.True(t, channelMatchesExpectedTaskPlugin(c, geminiChannel, "vertex-shared"), "a retry may select another declared provider")
}
func distributorTaskPluginSource(key string, channelType int) string {
return fmt.Sprintf(`
export const meta = {
apiVersion: 1,
key: %q,
name: %q,
version: "1.0.0",
author: {name: "Test"},
channelTypes: [%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)
}
func distributorEndpointPluginSource(key string, channelType int) string {
return fmt.Sprintf(`
export const meta = {
apiVersion: 1,
key: %q,
name: %q,
version: "1.0.0",
author: {name: "Test"},
channelTypes: [%d],
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, channelType)
}
-67
View File
@@ -1,67 +0,0 @@
package middleware
import (
"bytes"
"encoding/json"
"io"
"net/http"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/gin-gonic/gin"
)
func JimengRequestConvert() func(c *gin.Context) {
return func(c *gin.Context) {
action := c.Query("Action")
if action == "" {
abortWithOpenAiMessage(c, http.StatusBadRequest, "Action query parameter is required")
return
}
// Handle Jimeng official API request
var originalReq map[string]interface{}
if err := common.UnmarshalBodyReusable(c, &originalReq); err != nil {
abortWithOpenAiMessage(c, http.StatusBadRequest, "Invalid request body")
return
}
model, _ := originalReq["req_key"].(string)
prompt, _ := originalReq["prompt"].(string)
unifiedReq := map[string]interface{}{
"model": model,
"prompt": prompt,
"metadata": originalReq,
}
jsonData, err := json.Marshal(unifiedReq)
if err != nil {
abortWithOpenAiMessage(c, http.StatusInternalServerError, "Failed to marshal request body")
return
}
// Update request body
c.Request.Body = io.NopCloser(bytes.NewBuffer(jsonData))
c.Set(common.KeyRequestBody, jsonData)
if image, ok := originalReq["image"]; !ok || image == "" {
c.Set("action", constant.TaskActionTextGenerate)
}
c.Request.URL.Path = "/v1/video/generations"
if action == "CVSync2AsyncGetResult" {
taskId, ok := originalReq["task_id"].(string)
if !ok || taskId == "" {
abortWithOpenAiMessage(c, http.StatusBadRequest, "task_id is required for CVSync2AsyncGetResult")
return
}
c.Request.URL.Path = "/v1/video/generations/" + taskId
c.Request.Method = http.MethodGet
c.Set("task_id", taskId)
c.Set("relay_mode", relayconstant.RelayModeVideoFetchByID)
}
c.Next()
}
}
-52
View File
@@ -1,52 +0,0 @@
package middleware
import (
"bytes"
"encoding/json"
"io"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/gin-gonic/gin"
)
func KlingRequestConvert() func(c *gin.Context) {
return func(c *gin.Context) {
var originalReq map[string]interface{}
if err := common.UnmarshalBodyReusable(c, &originalReq); err != nil {
c.Next()
return
}
// Support both model_name and model fields
model, _ := originalReq["model_name"].(string)
if model == "" {
model, _ = originalReq["model"].(string)
}
prompt, _ := originalReq["prompt"].(string)
unifiedReq := map[string]interface{}{
"model": model,
"prompt": prompt,
"metadata": originalReq,
}
jsonData, err := json.Marshal(unifiedReq)
if err != nil {
c.Next()
return
}
// Rewrite request body and path
c.Request.Body = io.NopCloser(bytes.NewBuffer(jsonData))
c.Request.URL.Path = "/v1/video/generations"
if image, ok := originalReq["image"]; !ok || image == "" {
c.Set("action", constant.TaskActionTextGenerate)
}
// We have to reset the request body for the next handlers
c.Set(common.KeyRequestBody, jsonData)
c.Next()
}
}
+1
View File
@@ -17,6 +17,7 @@ func RouteTag(tag string) gin.HandlerFunc {
}
func SetUpLogger(server *gin.Engine) {
server.Use(redactTaskArtifactAccessQuery())
server.Use(gin.LoggerWithFormatter(func(param gin.LogFormatterParams) string {
var requestID string
if param.Keys != nil {
+245
View File
@@ -0,0 +1,245 @@
package middleware
import (
"net/http"
"net/url"
"strings"
"sync"
"time"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/system_setting"
"github.com/gin-gonic/gin"
)
const TaskArtifactAccessContextKey = "task_artifact_access"
const (
taskArtifactAccessRawContextKey = "task_artifact_access_raw"
taskArtifactAccessPresentContextKey = "task_artifact_access_present"
taskArtifactAccessInvalidContextKey = "task_artifact_access_invalid"
taskArtifactAccessRateWindow = time.Minute
taskArtifactAccessCleanupInterval = time.Minute
maxEncodedTaskArtifactAccessQuerySize = 128
)
type taskArtifactRateEntry struct {
windowStart time.Time
count int
}
type taskArtifactAccessLimiter struct {
mutex sync.Mutex
global int
byIP map[string]int
byObject map[string]int
rates map[string]taskArtifactRateEntry
nextCleanup time.Time
limits system_setting.TaskArtifactAccessLimits
}
var taskArtifactAnonymousLimiter = newTaskArtifactAccessLimiter(
system_setting.LoadTaskArtifactAccessLimits(),
)
func newTaskArtifactAccessLimiter(limits system_setting.TaskArtifactAccessLimits) *taskArtifactAccessLimiter {
return &taskArtifactAccessLimiter{
byIP: make(map[string]int),
byObject: make(map[string]int),
rates: make(map[string]taskArtifactRateEntry),
limits: limits,
}
}
func (l *taskArtifactAccessLimiter) invalidAttempt(now time.Time, ip string) bool {
l.mutex.Lock()
defer l.mutex.Unlock()
if l.nextCleanup.IsZero() || !now.Before(l.nextCleanup) {
for key, entry := range l.rates {
if now.Sub(entry.windowStart) >= taskArtifactAccessRateWindow {
delete(l.rates, key)
}
}
l.nextCleanup = now.Add(taskArtifactAccessCleanupInterval)
}
rate := l.rates[ip]
if rate.windowStart.IsZero() || now.Sub(rate.windowStart) >= taskArtifactAccessRateWindow {
rate = taskArtifactRateEntry{windowStart: now}
}
if rate.count >= l.limits.InvalidRatePerMinute {
return false
}
rate.count++
l.rates[ip] = rate
return true
}
func (l *taskArtifactAccessLimiter) acquire(ip, taskID, artifactKey string) (func(), bool) {
l.mutex.Lock()
defer l.mutex.Unlock()
objectKey := taskID + "\x00" + artifactKey
if l.global >= l.limits.GlobalConcurrency ||
l.byIP[ip] >= l.limits.IPConcurrency ||
l.byObject[objectKey] >= l.limits.ObjectConcurrency {
return nil, false
}
l.global++
l.byIP[ip]++
l.byObject[objectKey]++
var once sync.Once
return func() {
once.Do(func() {
l.mutex.Lock()
defer l.mutex.Unlock()
l.global--
l.byIP[ip]--
l.byObject[objectKey]--
if l.byIP[ip] == 0 {
delete(l.byIP, ip)
}
if l.byObject[objectKey] == 0 {
delete(l.byObject, objectKey)
}
})
}, true
}
func redactTaskArtifactAccessQuery() gin.HandlerFunc {
return func(c *gin.Context) {
path := c.Request.URL.Path
isArtifactContent := strings.HasPrefix(path, "/v1/tasks/") &&
strings.Contains(path, "/artifacts/") &&
strings.HasSuffix(path, "/content")
isLegacyVideoContent := strings.HasPrefix(path, "/v1/videos/") &&
strings.HasSuffix(path, "/content")
if !isArtifactContent && !isLegacyVideoContent {
c.Next()
return
}
rawAccess, present, invalid := popTaskArtifactAccessQuery(c.Request)
if present {
c.Set(taskArtifactAccessRawContextKey, rawAccess)
c.Set(taskArtifactAccessPresentContextKey, true)
c.Set(taskArtifactAccessInvalidContextKey, invalid)
}
c.Next()
}
}
func popTaskArtifactAccessQuery(request *http.Request) (string, bool, bool) {
if request == nil || request.URL == nil {
return "", false, false
}
rawAccess := ""
count := 0
invalid := false
kept := make([]string, 0)
for _, part := range strings.Split(request.URL.RawQuery, "&") {
rawKey, rawValue, _ := strings.Cut(part, "=")
key, err := url.QueryUnescape(rawKey)
if err != nil || key != service.TaskArtifactAccessQueryParameter {
kept = append(kept, part)
continue
}
count++
if len(rawValue) > maxEncodedTaskArtifactAccessQuerySize {
invalid = true
continue
}
if count == 1 {
value, decodeErr := url.QueryUnescape(rawValue)
if decodeErr != nil {
invalid = true
} else {
rawAccess = value
}
}
}
if count == 0 {
return "", false, false
}
invalid = invalid || count != 1
request.URL.RawQuery = strings.Join(kept, "&")
request.RequestURI = request.URL.RequestURI()
return rawAccess, true, invalid
}
// TokenOrTaskArtifactAccessAuth accepts the normal relay API Bearer token or a
// route-bound capability. Capabilities are verified before any database read.
func TokenOrTaskArtifactAccessAuth(taskParam, artifactParam string) gin.HandlerFunc {
return func(c *gin.Context) {
c.Header("Cache-Control", "private, no-store")
rawAccess := c.GetString(taskArtifactAccessRawContextKey)
present := c.GetBool(taskArtifactAccessPresentContextKey)
invalid := c.GetBool(taskArtifactAccessInvalidContextKey)
if queryAccess, queryPresent, queryInvalid := popTaskArtifactAccessQuery(c.Request); queryPresent {
present = true
if rawAccess == "" {
rawAccess = queryAccess
}
invalid = invalid || queryInvalid
}
if !present {
TokenAuth()(c)
return
}
taskID := c.Param(taskParam)
artifactKey := c.Param(artifactParam)
ip := c.ClientIP()
if ip == "" {
ip = "unknown"
}
if invalid || !service.VerifyTaskArtifactAccess(rawAccess, taskID, artifactKey) {
if !taskArtifactAnonymousLimiter.invalidAttempt(time.Now(), ip) {
writeTaskArtifactAccessLimited(c)
return
}
writeTaskArtifactAccessNotFound(c)
return
}
release, ok := taskArtifactAnonymousLimiter.acquire(ip, taskID, artifactKey)
if !ok {
writeTaskArtifactAccessLimited(c)
return
}
defer release()
c.Set(TaskArtifactAccessContextKey, true)
c.Next()
}
}
func IsTaskArtifactAccess(c *gin.Context) bool {
return c != nil && c.GetBool(TaskArtifactAccessContextKey)
}
func writeTaskArtifactAccessNotFound(c *gin.Context) {
c.Header("Cache-Control", "private, no-store")
c.AbortWithStatusJSON(http.StatusNotFound, gin.H{
"error": gin.H{
"message": "Task or artifact not found",
"type": "artifact_not_found",
"code": "artifact_not_found",
},
})
}
func writeTaskArtifactAccessLimited(c *gin.Context) {
c.Header("Cache-Control", "private, no-store")
c.Header("Retry-After", "60")
c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
"error": gin.H{
"message": "Artifact access limit exceeded",
"type": "rate_limit_error",
"code": "artifact_access_limited",
},
})
}
+158
View File
@@ -0,0 +1,158 @@
package middleware
import (
"bytes"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/system_setting"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestTaskArtifactAccessIsRedactedAndVerifiedBeforeHandler(t *testing.T) {
gin.SetMode(gin.TestMode)
previousSecret := common.CryptoSecret
common.CryptoSecret = "task-artifact-middleware-secret"
t.Cleanup(func() { common.CryptoSecret = previousSecret })
access, err := service.IssueTaskArtifactAccess("task-1", "video-main")
require.NoError(t, err)
router := gin.New()
router.Use(redactTaskArtifactAccessQuery())
router.GET(
"/v1/tasks/:key/artifacts/:artifact_key/content",
TokenOrTaskArtifactAccessAuth("key", "artifact_key"),
func(c *gin.Context) {
assert.True(t, IsTaskArtifactAccess(c))
assert.NotContains(t, c.Request.URL.RawQuery, service.TaskArtifactAccessQueryParameter)
assert.Equal(t, "kept", c.Query("keep"))
c.Status(http.StatusNoContent)
},
)
request := httptest.NewRequest(
http.MethodGet,
"/v1/tasks/task-1/artifacts/video-main/content?access="+urlQueryEscape(access)+"&keep=kept",
nil,
)
request.RemoteAddr = "192.0.2.1:1234"
recorder := httptest.NewRecorder()
router.ServeHTTP(recorder, request)
assert.Equal(t, http.StatusNoContent, recorder.Code)
}
func TestTaskArtifactAccessRejectsTamperedAndEmptyCapabilitiesAsNotFound(t *testing.T) {
gin.SetMode(gin.TestMode)
previousSecret := common.CryptoSecret
common.CryptoSecret = "task-artifact-middleware-reject-secret"
t.Cleanup(func() { common.CryptoSecret = previousSecret })
router := gin.New()
router.Use(redactTaskArtifactAccessQuery())
router.GET(
"/v1/tasks/:key/artifacts/:artifact_key/content",
TokenOrTaskArtifactAccessAuth("key", "artifact_key"),
func(c *gin.Context) { c.Status(http.StatusNoContent) },
)
for _, query := range []string{
"?access=",
"?access=invalid",
"?access=first&access=second",
"?access=" + strings.Repeat("x", 1024),
"?access=%20" + strings.Repeat("A", 43) + "%20",
} {
request := httptest.NewRequest(
http.MethodGet,
"/v1/tasks/task-1/artifacts/video-main/content"+query,
nil,
)
recorder := httptest.NewRecorder()
router.ServeHTTP(recorder, request)
assert.Equal(t, http.StatusNotFound, recorder.Code)
}
}
func TestTaskArtifactAccessLimiterDefaults(t *testing.T) {
limits := system_setting.TaskArtifactAccessLimits{
InvalidRatePerMinute: system_setting.DefaultTaskArtifactInvalidRateLimitPerMinute,
GlobalConcurrency: system_setting.DefaultTaskArtifactGlobalConcurrency,
IPConcurrency: system_setting.DefaultTaskArtifactIPConcurrency,
ObjectConcurrency: system_setting.DefaultTaskArtifactObjectConcurrency,
}
limiter := newTaskArtifactAccessLimiter(limits)
now := time.Unix(1000, 0)
releases := make([]func(), 0, limits.ObjectConcurrency)
for i := 0; i < limits.ObjectConcurrency; i++ {
release, ok := limiter.acquire("192.0.2.1", "task-1", "video")
require.True(t, ok)
releases = append(releases, release)
}
_, ok := limiter.acquire("192.0.2.2", "task-1", "video")
assert.False(t, ok, "task+key concurrency is shared across IPs")
for _, release := range releases {
release()
}
rateLimiter := newTaskArtifactAccessLimiter(limits)
for i := 0; i < limits.InvalidRatePerMinute; i++ {
assert.True(t, rateLimiter.invalidAttempt(now, "192.0.2.10"))
}
assert.False(t, rateLimiter.invalidAttempt(now, "192.0.2.10"))
assert.True(t, rateLimiter.invalidAttempt(now.Add(time.Minute), "192.0.2.10"))
}
func TestRedactTaskArtifactAccessAlsoCoversLegacyVideoRoute(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
router.Use(redactTaskArtifactAccessQuery())
router.GET("/v1/videos/:task_id/content", func(c *gin.Context) {
assert.NotContains(t, c.Request.URL.RawQuery, "access")
assert.NotContains(t, c.Request.RequestURI, "secret-capability")
assert.Equal(t, "ok", c.Query("keep"))
c.Status(http.StatusNoContent)
})
request := httptest.NewRequest(
http.MethodGet,
"/v1/videos/task-1/content?access=secret-capability&keep=ok",
nil,
)
recorder := httptest.NewRecorder()
router.ServeHTTP(recorder, request)
assert.Equal(t, http.StatusNoContent, recorder.Code)
}
func TestSetUpLoggerNeverWritesTaskArtifactAccess(t *testing.T) {
gin.SetMode(gin.TestMode)
previousWriter := gin.DefaultWriter
var output bytes.Buffer
gin.DefaultWriter = &output
t.Cleanup(func() { gin.DefaultWriter = previousWriter })
router := gin.New()
SetUpLogger(router)
router.GET("/v1/tasks/:key/artifacts/:artifact_key/content", func(c *gin.Context) {
c.Status(http.StatusNoContent)
})
request := httptest.NewRequest(
http.MethodGet,
"/v1/tasks/task-1/artifacts/video/content?access=never-log-this&keep=ok",
nil,
)
router.ServeHTTP(httptest.NewRecorder(), request)
assert.False(t, strings.Contains(output.String(), "never-log-this"))
}
func urlQueryEscape(value string) string {
replacer := strings.NewReplacer("+", "%2B", "=", "%3D")
return replacer.Replace(value)
}
File diff suppressed because it is too large. Load diff
+500
View File
@@ -0,0 +1,500 @@
package middleware
import (
"bytes"
"fmt"
"net/http"
"net/http/httptest"
"strconv"
"strings"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
appI18n "github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/relay"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)
func setupOriginTaskDB(t *testing.T) {
t.Helper()
previousDB := model.DB
previousType := common.MainDatabaseType()
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
require.NoError(t, database.AutoMigrate(&model.Task{}, &model.Channel{}))
model.DB = database
common.SetMainDatabaseType(common.DatabaseTypeSQLite)
t.Cleanup(func() {
model.DB = previousDB
common.SetMainDatabaseType(previousType)
})
}
func insertOriginTaskChannel(t *testing.T, status int) *model.Channel {
t.Helper()
channel := &model.Channel{
Name: "origin-channel",
Key: "sk-origin",
Status: status,
Type: constant.ChannelTypeDoubaoVideo,
}
require.NoError(t, model.DB.Create(channel).Error)
return channel
}
func insertOriginOwnedTask(t *testing.T, taskID string, userID, channelID int, platform constant.TaskPlatform) *model.Task {
t.Helper()
task := &model.Task{
TaskID: taskID,
UserId: userID,
ChannelId: channelID,
Platform: platform,
Action: "text_to_video",
Status: model.TaskStatusSuccess,
PrivateData: model.TaskPrivateData{
UpstreamTaskID: "upstream-" + taskID,
},
}
data, err := common.Marshal(map[string]any{"id": "upstream-" + taskID})
require.NoError(t, err)
task.Data = data
require.NoError(t, model.DB.Create(task).Error)
return task
}
func originTaskTestContext(userID int) *gin.Context {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/vendor/jobs", nil)
common.SetContextKey(c, constant.ContextKeyUserId, userID)
return c
}
func resolvedOriginPin(c *gin.Context) (int, bool) {
pin, found, _ := service.GetChannelConstraints(c).ResolvedPin()
if !found {
return 0, false
}
return pin.ChannelId, true
}
func TestApplyOriginTaskIntent(t *testing.T) {
setupOriginTaskDB(t)
enabled := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
otherEnabled := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
disabled := insertOriginTaskChannel(t, common.ChannelStatusManuallyDisabled)
insertOriginOwnedTask(t, "task-own", 7, enabled.Id, "origin-plugin")
insertOriginOwnedTask(t, "task-own-b", 7, enabled.Id, "origin-plugin")
insertOriginOwnedTask(t, "task-other-channel", 7, otherEnabled.Id, "origin-plugin")
insertOriginOwnedTask(t, "task-foreign", 8, enabled.Id, "origin-plugin")
insertOriginOwnedTask(t, "task-wrong-platform", 7, enabled.Id, "other-plugin")
insertOriginOwnedTask(t, "task-legacy", 7, enabled.Id, constant.TaskPlatform(strconv.Itoa(constant.ChannelTypeDoubaoVideo)))
insertOriginOwnedTask(t, "task-disabled", 7, disabled.Id, "origin-plugin")
tooMany := make([]any, 17)
for i := range tooMany {
tooMany[i] = fmt.Sprintf("task-%d", i)
}
tests := []struct {
name string
userID int
intent map[string]any
channelType int
wantCode string
wantPinned int
wantIDs []string
}{
{
name: "valid single origin id pins channel",
userID: 7,
intent: map[string]any{"originTaskIds": []any{"task-own"}},
wantPinned: enabled.Id,
wantIDs: []string{"task-own"},
},
{
name: "dedupes preserving first order",
userID: 7,
intent: map[string]any{"originTaskIds": []any{"task-own", " task-own ", "task-own-b"}},
wantPinned: enabled.Id,
wantIDs: []string{"task-own", "task-own-b"},
},
{
name: "legacy platform matches channelType",
userID: 7,
intent: map[string]any{"originTaskIds": []any{"task-legacy"}},
channelType: constant.ChannelTypeDoubaoVideo,
wantPinned: enabled.Id,
wantIDs: []string{"task-legacy"},
},
{
name: "unknown id",
userID: 7,
intent: map[string]any{"originTaskIds": []any{"task-missing"}},
wantCode: "origin_task_not_found",
},
{
name: "other user's task is not found",
userID: 7,
intent: map[string]any{"originTaskIds": []any{"task-foreign"}},
wantCode: "origin_task_not_found",
},
{
name: "platform mismatch",
userID: 7,
intent: map[string]any{"originTaskIds": []any{"task-wrong-platform"}},
wantCode: "origin_task_platform_mismatch",
},
{
name: "two ids on different channels",
userID: 7,
intent: map[string]any{"originTaskIds": []any{"task-own", "task-other-channel"}},
wantCode: "origin_task_channel_conflict",
},
{
name: "disabled channel",
userID: 7,
intent: map[string]any{"originTaskIds": []any{"task-disabled"}},
wantCode: "origin_task_channel_disabled",
},
{
name: "more than 16 ids",
userID: 7,
intent: map[string]any{"originTaskIds": tooMany},
wantCode: "invalid_origin_task_ids",
},
{
name: "non-array originTaskIds",
userID: 7,
intent: map[string]any{"originTaskIds": "task-own"},
wantCode: "invalid_origin_task_ids",
},
{
name: "empty string entry",
userID: 7,
intent: map[string]any{"originTaskIds": []any{"task-own", " "}},
wantCode: "invalid_origin_task_ids",
},
}
for _, testCase := range tests {
t.Run(testCase.name, func(t *testing.T) {
c := originTaskTestContext(testCase.userID)
intentErr := applyOriginTaskIntent(c, testCase.intent, jsplugin.Meta{Key: "origin-plugin", ChannelTypes: []int{testCase.channelType}})
if testCase.wantCode != "" {
require.NotNil(t, intentErr)
assert.Equal(t, testCase.wantCode, intentErr.Code)
assert.Equal(t, http.StatusBadRequest, intentErr.StatusCode)
_, pinned := resolvedOriginPin(c)
assert.False(t, pinned)
return
}
require.Nil(t, intentErr)
pinnedID, ok := resolvedOriginPin(c)
require.True(t, ok)
assert.Equal(t, testCase.wantPinned, pinnedID)
tasks, ok := common.GetContextKeyType[[]*model.Task](c, constant.ContextKeyOriginTasks)
require.True(t, ok)
require.Len(t, tasks, len(testCase.wantIDs))
for i, wantID := range testCase.wantIDs {
assert.Equal(t, wantID, tasks[i].TaskID)
}
})
}
}
func TestApplyOriginTaskIntentAbsentAndEmptyAreNoop(t *testing.T) {
setupOriginTaskDB(t)
c := originTaskTestContext(7)
require.Nil(t, applyOriginTaskIntent(c, map[string]any{}, jsplugin.Meta{Key: "origin-plugin"}))
require.Nil(t, applyOriginTaskIntent(c, map[string]any{"originTaskIds": []any{}}, jsplugin.Meta{Key: "origin-plugin"}))
_, pinned := resolvedOriginPin(c)
assert.False(t, pinned)
}
func TestApplyOriginTaskAffinitySetsLockedChannel(t *testing.T) {
setupOriginTaskDB(t)
channel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
insertOriginOwnedTask(t, "task-lock", 7, channel.Id, "origin-plugin")
c := originTaskTestContext(7)
require.Nil(t, applyOriginTaskIntent(c, map[string]any{"originTaskIds": []any{"task-lock"}}, jsplugin.Meta{Key: "origin-plugin"}))
info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
taskErr := relay.ApplyOriginTaskAffinity(c, info)
require.Nil(t, taskErr)
locked, ok := info.LockedChannel.(*model.Channel)
require.True(t, ok)
require.NotNil(t, locked)
assert.Equal(t, channel.Id, locked.Id)
require.Len(t, info.OriginTasks, 1)
assert.Equal(t, "task-lock", info.OriginTasks[0].TaskID)
assert.Equal(t, "upstream-task-lock", info.OriginTasks[0].UpstreamTaskID)
assert.Equal(t, "text_to_video", info.OriginTasks[0].Action)
assert.Equal(t, string(model.TaskStatusSuccess), info.OriginTasks[0].Status)
}
func TestPrepareTaskPluginRoutePinsOriginTaskChannel(t *testing.T) {
setupOriginTaskDB(t)
channel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
insertOriginOwnedTask(t, "task-route", 7, channel.Id, "origin-route")
plugin := compileTaskRoutePlugin(t, `
export const meta = {
apiVersion: 1, key: "origin-route", name: "Origin", version: "1.0.0",
author: {name: "Test"},
models: ["resolved-model"], fetchMode: "per_task",
routes: [{method: "POST", path: "/vendor/jobs", type: "submit", decode: "decodeJob", render: "jobCreated"}],
};
export const native = {
decodeJob: function() { return {kind: "submit", model: "resolved-model", originTaskIds: ["task-route"], requestBody: {prompt: "ok"}}; },
jobCreated: function(ctx, task) { return task; },
};
export function buildSubmitRequest() { return {url: "https://example.com"}; }
export function parseSubmitResponse() { return {taskId: "one"}; }
export function buildQueryRequest() { return {url: "https://example.com"}; }
export function parseTaskResult() { return {status: "SUCCESS"}; }
`)
reached := false
router := gin.New()
router.POST("/vendor/jobs", pinTaskPluginRoute(plugin, 0), func(c *gin.Context) {
common.SetContextKey(c, constant.ContextKeyUserId, 7)
c.Next()
}, PrepareTaskPluginRoute(), func(c *gin.Context) {
reached = true
pinnedID, ok := resolvedOriginPin(c)
require.True(t, ok)
assert.Equal(t, channel.Id, pinnedID)
c.Status(http.StatusNoContent)
})
request := httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{"model":"resolved-model"}`))
request.Header.Set("Content-Type", "application/json")
recorder := httptest.NewRecorder()
router.ServeHTTP(recorder, request)
assert.True(t, reached)
assert.Equal(t, http.StatusNoContent, recorder.Code)
}
func TestPrepareTaskPluginRouteRejectsUnknownOriginTask(t *testing.T) {
setupOriginTaskDB(t)
plugin := compileTaskRoutePlugin(t, `
export const meta = {
apiVersion: 1, key: "origin-route-missing", name: "Origin", version: "1.0.0",
author: {name: "Test"},
models: ["resolved-model"], fetchMode: "per_task",
routes: [{method: "POST", path: "/vendor/jobs", type: "submit", decode: "decodeJob", render: "jobCreated"}],
};
export const native = {
decodeJob: function() { return {kind: "submit", model: "resolved-model", originTaskIds: ["missing"], requestBody: {}}; },
jobCreated: function(ctx, task) { return task; },
};
export function buildSubmitRequest() { return {url: "https://example.com"}; }
export function parseSubmitResponse() { return {taskId: "one"}; }
export function buildQueryRequest() { return {url: "https://example.com"}; }
export function parseTaskResult() { return {status: "SUCCESS"}; }
`)
reached := false
router := gin.New()
router.POST("/vendor/jobs", pinTaskPluginRoute(plugin, 0), func(c *gin.Context) {
common.SetContextKey(c, constant.ContextKeyUserId, 7)
c.Next()
}, PrepareTaskPluginRoute(), func(c *gin.Context) { reached = true })
request := httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{"model":"resolved-model"}`))
request.Header.Set("Content-Type", "application/json")
recorder := httptest.NewRecorder()
router.ServeHTTP(recorder, request)
assert.False(t, reached)
assert.Equal(t, http.StatusBadRequest, recorder.Code)
}
func TestPrepareTaskPluginEndpointPinsOriginTaskChannel(t *testing.T) {
setupOriginTaskDB(t)
channel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
insertOriginOwnedTask(t, "task-endpoint", 7, channel.Id, "origin-endpoint")
const key = "origin-endpoint"
_, err := jsplugin.DefaultRegistry.Register(taskProtocolPluginSource(
key,
"1.0.0",
`["claimed-model"]`,
"/v1/responses",
`return {model: ctx.model, originTaskIds: ["task-endpoint"], requestBody: {prompt: "ok"}};`,
), jsplugin.Options{})
require.NoError(t, err)
t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(key)) })
reached := false
router := gin.New()
router.POST("/v1/responses", func(c *gin.Context) {
common.SetContextKey(c, constant.ContextKeyUserId, 7)
c.Next()
}, PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) {
reached = true
pinnedID, ok := resolvedOriginPin(c)
require.True(t, ok)
assert.Equal(t, channel.Id, pinnedID)
c.Status(http.StatusNoContent)
})
request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"claimed-model","input":"hello"}`))
request.Header.Set("Content-Type", "application/json")
recorder := httptest.NewRecorder()
router.ServeHTTP(recorder, request)
assert.True(t, reached)
assert.Equal(t, http.StatusNoContent, recorder.Code)
}
func TestPrepareTaskPluginEndpointRejectsUnknownOriginTask(t *testing.T) {
setupOriginTaskDB(t)
const key = "origin-endpoint-missing"
_, err := jsplugin.DefaultRegistry.Register(taskProtocolPluginSource(
key,
"1.0.0",
`["claimed-model"]`,
"/v1/responses",
`return {model: ctx.model, originTaskIds: ["missing"], requestBody: {prompt: "ok"}};`,
), jsplugin.Options{})
require.NoError(t, err)
t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(key)) })
reached := false
router := gin.New()
router.POST("/v1/responses", func(c *gin.Context) {
common.SetContextKey(c, constant.ContextKeyUserId, 7)
c.Next()
}, PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) { reached = true })
request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"claimed-model","input":"hello"}`))
request.Header.Set("Content-Type", "application/json")
recorder := httptest.NewRecorder()
router.ServeHTTP(recorder, request)
assert.False(t, reached)
assert.Equal(t, http.StatusBadRequest, recorder.Code)
assert.Contains(t, recorder.Body.String(), "origin_task_not_found")
}
func TestDistributeHonorsOriginTaskChannelPin(t *testing.T) {
require.NoError(t, appI18n.Init())
setupOriginTaskDB(t)
channel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{}`))
c.Request.Header.Set("Content-Type", "application/json")
c.Set("resolved_task_model", "resolved-model")
service.GetChannelConstraints(c).AddPin(dto.ChannelPin{
ChannelId: channel.Id,
Source: dto.PinSourceOriginTask,
Rank: dto.PinRankOriginTask,
RetryMode: dto.PinRetrySameChannel,
})
nextCalled := false
handler := Distribute()
handler(c)
if !c.IsAborted() {
nextCalled = true
}
assert.True(t, nextCalled)
assert.Equal(t, channel.Id, common.GetContextKeyInt(c, constant.ContextKeyChannelId))
}
func TestDistributeTokenPinBeatsOriginPin(t *testing.T) {
require.NoError(t, appI18n.Init())
setupOriginTaskDB(t)
tokenChannel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
originChannel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
var warnBuf bytes.Buffer
previousWriter := gin.DefaultErrorWriter
gin.DefaultErrorWriter = &warnBuf
t.Cleanup(func() { gin.DefaultErrorWriter = previousWriter })
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{}`))
c.Request.Header.Set("Content-Type", "application/json")
c.Set("resolved_task_model", "resolved-model")
constraints := service.GetChannelConstraints(c)
constraints.AddPin(dto.ChannelPin{
ChannelId: tokenChannel.Id,
Source: dto.PinSourceToken,
Rank: dto.PinRankToken,
RetryMode: dto.PinRetrySingleAttempt,
})
constraints.AddPin(dto.ChannelPin{
ChannelId: originChannel.Id,
Source: dto.PinSourceOriginTask,
Rank: dto.PinRankOriginTask,
RetryMode: dto.PinRetrySameChannel,
})
nextCalled := false
Distribute()(c)
if !c.IsAborted() {
nextCalled = true
}
assert.True(t, nextCalled)
assert.Equal(t, tokenChannel.Id, common.GetContextKeyInt(c, constant.ContextKeyChannelId))
warn := warnBuf.String()
assert.Contains(t, warn, "winning_source=token")
assert.Contains(t, warn, fmt.Sprintf("winning_channel_id=%d", tokenChannel.Id))
assert.Contains(t, warn, "overridden_source=origin_task")
assert.Contains(t, warn, fmt.Sprintf("overridden_channel_id=%d", originChannel.Id))
}
func TestDistributePinViolatingIdentityFilterErrors(t *testing.T) {
require.NoError(t, appI18n.Init())
setupOriginTaskDB(t)
channel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{}`))
c.Request.Header.Set("Content-Type", "application/json")
c.Set("resolved_task_model", "resolved-model")
c.Set("expected_task_plugin_key", "alpha")
service.GetChannelConstraints(c).AddPin(dto.ChannelPin{
ChannelId: channel.Id,
Source: dto.PinSourceOriginTask,
Rank: dto.PinRankOriginTask,
RetryMode: dto.PinRetrySameChannel,
})
Distribute()(c)
assert.True(t, c.IsAborted())
assert.Equal(t, http.StatusBadRequest, recorder.Code)
assert.Contains(t, recorder.Body.String(), string(dto.FilterTaskPluginIdentity))
}
func TestApplyChannelPinLocksOnlySameChannelRetry(t *testing.T) {
setupOriginTaskDB(t)
channel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
insertOriginOwnedTask(t, "task-lock-mode", 7, channel.Id, "origin-plugin")
c := originTaskTestContext(7)
require.Nil(t, applyOriginTaskIntent(c, map[string]any{"originTaskIds": []any{"task-lock-mode"}}, jsplugin.Meta{Key: "origin-plugin"}))
info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
require.Nil(t, relay.ApplyChannelPin(c, info))
locked, ok := info.LockedChannel.(*model.Channel)
require.True(t, ok)
assert.Equal(t, channel.Id, locked.Id)
tokenOnly := originTaskTestContext(7)
service.GetChannelConstraints(tokenOnly).AddPin(dto.ChannelPin{
ChannelId: channel.Id,
Source: dto.PinSourceToken,
Rank: dto.PinRankToken,
RetryMode: dto.PinRetrySingleAttempt,
})
tokenInfo := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
require.Nil(t, relay.ApplyChannelPin(tokenOnly, tokenInfo))
assert.Nil(t, tokenInfo.LockedChannel)
}
File diff suppressed because it is too large. Load diff
+7 -38
View File
@@ -1,51 +1,20 @@
package middleware
import (
"errors"
"fmt"
"log"
"os"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/gin-gonic/gin"
)
var defaultTrustedProxyCIDRs = []string{
"127.0.0.0/8",
"::1",
"10.0.0.0/8",
"172.16.0.0/12",
"192.168.0.0/16",
"fc00::/7",
}
func ConfigureTrustedProxies(engine *gin.Engine) error {
rawTrustedProxies := strings.TrimSpace(os.Getenv("TRUSTED_PROXIES"))
if rawTrustedProxies == "" {
trustedProxies, usedDefaults, err := common.ResolveTrustedProxies(os.Getenv("TRUSTED_PROXIES"))
if err != nil {
return err
}
if usedDefaults {
log.Print("WARNING: TRUSTED_PROXIES is unset or blank; trusting loopback, RFC 1918, and IPv6 ULA proxy addresses for compatibility. Set TRUSTED_PROXIES=none to trust no proxies, or configure explicit proxy IPs/CIDRs to replace these defaults.")
return engine.SetTrustedProxies(defaultTrustedProxyCIDRs)
}
if strings.EqualFold(rawTrustedProxies, "none") {
return engine.SetTrustedProxies(nil)
}
parts := strings.Split(rawTrustedProxies, ",")
trustedProxies := make([]string, 0, len(parts))
for _, part := range parts {
trustedProxy := strings.TrimSpace(part)
if trustedProxy == "" {
continue
}
if strings.EqualFold(trustedProxy, "none") {
return errors.New("TRUSTED_PROXIES=none must be used alone")
}
trustedProxies = append(trustedProxies, trustedProxy)
}
if len(trustedProxies) == 0 {
return errors.New("TRUSTED_PROXIES does not contain an IP address or CIDR")
}
if err := engine.SetTrustedProxies(trustedProxies); err != nil {
return fmt.Errorf("invalid TRUSTED_PROXIES: %w", err)
}
return nil
return common.ConfigureTrustedProxies(engine, trustedProxies)
}
+16 -7
View File
@@ -4,7 +4,9 @@ import (
"fmt"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/logger"
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
)
@@ -15,13 +17,20 @@ func abortWithOpenAiMessage(c *gin.Context, statusCode int, message string, code
codeStr = string(code[0])
}
userId := c.GetInt("id")
c.JSON(statusCode, gin.H{
"error": gin.H{
"message": common.MessageWithRequestId(message, c.GetString(common.RequestIdKey)),
"type": "new_api_error",
"code": codeStr,
},
})
_, preparedPluginRoute := c.Get(pluginruntime.ContextKeyRouteRequest)
if !preparedPluginRoute || !RespondTaskPluginError(c, &dto.TaskError{
Code: codeStr,
Message: message,
StatusCode: statusCode,
}) {
c.JSON(statusCode, gin.H{
"error": gin.H{
"message": common.MessageWithRequestId(message, c.GetString(common.RequestIdKey)),
"type": "new_api_error",
"code": codeStr,
},
})
}
c.Abort()
logger.LogError(c.Request.Context(), fmt.Sprintf("user %d | %s", userId, message))
}
+53 -33
View File
@@ -3,12 +3,12 @@ package model
import (
"errors"
"fmt"
"sort"
"strings"
"sync"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/dto"
"github.com/samber/lo"
"gorm.io/gorm"
@@ -105,23 +105,40 @@ func getChannelQuery(group string, model string, retry int) (*gorm.DB, error) {
return channelQuery, nil
}
func GetChannel(group string, model string, retry int, requestPath string) (*Channel, error) {
func GetChannel(
group string,
model string,
retry int,
filters []dto.ChannelFilter,
) (*Channel, error) {
var abilities []Ability
var err error = nil
channelQuery, err := getChannelQuery(group, model, retry)
err := DB.Where(commonGroupCol+" = ? and model = ? and enabled = ?", group, model, true).Order("priority DESC, weight DESC").Find(&abilities).Error
if err != nil {
return nil, err
}
if common.UsingMainDatabase(common.DatabaseTypeSQLite) || common.UsingMainDatabase(common.DatabaseTypePostgreSQL) {
err = channelQuery.Order("weight DESC").Find(&abilities).Error
} else {
err = channelQuery.Order("weight DESC").Find(&abilities).Error
abilities = filterAbilitiesByConstraints(abilities, model, filters)
if len(abilities) > 0 {
priorities := make([]int64, 0)
seen := make(map[int64]bool)
for _, ability := range abilities {
priority := int64(0)
if ability.Priority != nil {
priority = *ability.Priority
}
if !seen[priority] {
seen[priority] = true
priorities = append(priorities, priority)
}
}
sort.Slice(priorities, func(i, j int) bool { return priorities[i] > priorities[j] })
if retry >= len(priorities) {
retry = len(priorities) - 1
}
targetPriority := priorities[retry]
abilities = lo.Filter(abilities, func(ability Ability, _ int) bool {
return ability.Priority == nil && targetPriority == 0 || ability.Priority != nil && *ability.Priority == targetPriority
})
}
if err != nil {
return nil, err
}
abilities = filterAbilitiesByRequestPathAndModel(abilities, requestPath, model)
channel := Channel{}
if len(abilities) > 0 {
// Randomly choose one
@@ -146,14 +163,12 @@ func GetChannel(group string, model string, retry int, requestPath string) (*Cha
return &channel, err
}
// filterAbilitiesByRequestPathAndModel restricts candidates by request path and
// model for the DB (non-memory-cache) selection path. Only Advanced Custom
// (type 58) channels are path-checked: kept only when one of their routes matches
// requestPath and model; all other channel types always pass. When requestPath is
// empty, filtering is skipped.
func filterAbilitiesByRequestPathAndModel(abilities []Ability, requestPath string, model string) []Ability {
if requestPath == "" || len(abilities) == 0 {
return abilities
// filterAbilitiesByConstraints applies the same ChannelSatisfiesFilters
// predicate used by the memory-cache path. A failed channel lookup fails
// closed when a task-plugin identity is required and fails open otherwise.
func filterAbilitiesByConstraints(abilities []Ability, modelName string, filters []dto.ChannelFilter) []Ability {
if len(abilities) == 0 {
return nil
}
channelIds := make([]int, 0, len(abilities))
@@ -168,31 +183,36 @@ func filterAbilitiesByRequestPathAndModel(abilities []Ability, requestPath strin
var channels []*Channel
if err := DB.Where("id IN ?", channelIds).Find(&channels).Error; err != nil {
// On error, fall back to unfiltered candidates to avoid blocking selection
if identityFilterRequiresKey(filters) {
return nil
}
return abilities
}
advancedConfigs := make(map[int]*dto.AdvancedCustomConfig)
channelsByID := make(map[int]*Channel, len(channels))
for _, channel := range channels {
if channel.Type == constant.ChannelTypeAdvancedCustom {
advancedConfigs[channel.Id] = channel.GetOtherSettings().AdvancedCustom
}
channelsByID[channel.Id] = channel
}
filtered := make([]Ability, 0, len(abilities))
for _, ability := range abilities {
config, isAdvancedCustom := advancedConfigs[ability.ChannelId]
if !isAdvancedCustom {
filtered = append(filtered, ability)
continue
}
if config != nil && config.SupportsPathForModel(requestPath, model) {
channel := channelsByID[ability.ChannelId]
if ok, _ := ChannelSatisfiesFilters(channel, modelName, filters); ok {
filtered = append(filtered, ability)
}
}
return filtered
}
func identityFilterRequiresKey(filters []dto.ChannelFilter) bool {
for _, filter := range filters {
if filter.Kind == dto.FilterTaskPluginIdentity && filter.TaskPluginKey != "" {
return true
}
}
return false
}
func (channel *Channel) AddAbilities(tx *gorm.DB) error {
models_ := strings.Split(channel.Models, ",")
groups_ := strings.Split(channel.Group, ",")
+10 -1
View File
@@ -419,6 +419,15 @@ func SearchChannels(keyword string, group string, model string, idSort bool, sor
return channels, nil
}
// GetChannelById loads a channel directly from the database, bypassing the
// in-memory channel cache.
//
// WARNING: do NOT call this on request hot paths (middleware, distribution,
// relay submit/retry, polling). Every call is a synchronous DB query and will
// not see cache-only state. Use CacheGetChannel instead: it serves from the
// in-memory cache and falls back to this function automatically when
// MemoryCacheEnabled is false. Direct use is appropriate only where fresh DB
// state is required, e.g. admin CRUD, channel testing, or cache (re)building.
func GetChannelById(id int, selectAll bool) (*Channel, error) {
channel := &Channel{Id: id}
var err error = nil
@@ -510,7 +519,7 @@ func (channel *Channel) GetBaseURL() string {
}
url := *channel.BaseURL
if url == "" {
url = constant.ChannelBaseURLs[channel.Type]
url = constant.GetChannelBaseURL(channel.Type)
}
return url
}
+14 -36
View File
@@ -12,7 +12,8 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/dto"
kitdto "github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/setting/ratio_setting"
)
@@ -20,7 +21,7 @@ var group2model2channels map[string]map[string][]int // enabled channel
var channelsIDM map[int]*Channel // all channels include disabled
// channel2advancedCustomConfig caches parsed Advanced Custom (type 58) configs so
// path-aware selection avoids re-parsing JSON per request. Refreshed on full sync.
var channel2advancedCustomConfig map[int]*dto.AdvancedCustomConfig
var channel2advancedCustomConfig map[int]*kitdto.AdvancedCustomConfig
var channelSyncLock sync.RWMutex
func InitChannelCache() {
@@ -29,7 +30,7 @@ func InitChannelCache() {
return
}
newChannelId2channel := make(map[int]*Channel)
newChannel2advancedCustomConfig := make(map[int]*dto.AdvancedCustomConfig)
newChannel2advancedCustomConfig := make(map[int]*kitdto.AdvancedCustomConfig)
var channels []*Channel
DB.Find(&channels)
for _, channel := range channels {
@@ -111,22 +112,27 @@ func SyncChannelCache(frequency int) {
}
}
func GetRandomSatisfiedChannel(group string, model string, retry int, requestPath string) (*Channel, error) {
func GetRandomSatisfiedChannel(
group string,
model string,
retry int,
filters []dto.ChannelFilter,
) (*Channel, error) {
// if memory cache is disabled, get channel directly from database
if !common.MemoryCacheEnabled {
return GetChannel(group, model, retry, requestPath)
return GetChannel(group, model, retry, filters)
}
channelSyncLock.RLock()
defer channelSyncLock.RUnlock()
// First, try to find channels with the exact model name.
channels := filterChannelsByRequestPathAndModel(group2model2channels[group][model], requestPath, model)
channels, _ := filterCandidateIDs(group2model2channels[group][model], model, filters)
// If no channels found, try to find channels with the normalized model name.
if len(channels) == 0 {
normalizedModel := ratio_setting.FormatMatchingModelName(model)
channels = filterChannelsByRequestPathAndModel(group2model2channels[group][normalizedModel], requestPath, model)
channels, _ = filterCandidateIDs(group2model2channels[group][normalizedModel], model, filters)
}
if len(channels) == 0 {
@@ -208,34 +214,6 @@ func GetRandomSatisfiedChannel(group string, model string, retry int, requestPat
return nil, errors.New("channel not found")
}
// filterChannelsByRequestPathAndModel restricts candidates by request path and
// model. Only Advanced Custom (type 58) channels are path-checked: they are kept
// only when one of their configured routes matches requestPath and model. All
// other channel types always pass. When requestPath is empty, filtering is skipped.
// Caller must hold channelSyncLock (read lock). The cached slice is never mutated.
func filterChannelsByRequestPathAndModel(channels []int, requestPath string, model string) []int {
if requestPath == "" || len(channels) == 0 {
return channels
}
filtered := make([]int, 0, len(channels))
for _, channelId := range channels {
channel, ok := channelsIDM[channelId]
if !ok {
// keep it so the downstream consistency error is raised as before
filtered = append(filtered, channelId)
continue
}
if channel.Type != constant.ChannelTypeAdvancedCustom {
filtered = append(filtered, channelId)
continue
}
if config := channel2advancedCustomConfig[channelId]; config != nil && config.SupportsPathForModel(requestPath, model) {
filtered = append(filtered, channelId)
}
}
return filtered
}
func CacheGetChannel(id int) (*Channel, error) {
if !common.MemoryCacheEnabled {
return GetChannelById(id, true)
@@ -311,7 +289,7 @@ func CacheUpdateChannel(channel *Channel) {
}
channelsIDM[channel.Id] = channel
if channel2advancedCustomConfig == nil {
channel2advancedCustomConfig = make(map[int]*dto.AdvancedCustomConfig)
channel2advancedCustomConfig = make(map[int]*kitdto.AdvancedCustomConfig)
}
delete(channel2advancedCustomConfig, channel.Id)
if channel.Type == constant.ChannelTypeAdvancedCustom {
+108
View File
@@ -0,0 +1,108 @@
package model
import (
"slices"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
)
var filterEvalOrder = []dto.ChannelFilterKind{
dto.FilterRequestPath,
dto.FilterTaskPluginIdentity,
}
// ChannelSatisfiesFilters reports whether ch passes every filter.
// On false, it returns the kind of the first violated filter (request_path
// then task_plugin_identity) for error attribution.
func ChannelSatisfiesFilters(ch *Channel, modelName string, filters []dto.ChannelFilter) (bool, dto.ChannelFilterKind) {
if ch == nil {
return false, ""
}
for _, kind := range filterEvalOrder {
for _, filter := range filters {
if filter.Kind != kind {
continue
}
if !channelMatchesFilter(ch, modelName, filter) {
return false, kind
}
}
}
return true, ""
}
// filterCandidateIDs applies filters to a cached candidate id list.
// Caller must hold channelSyncLock (read lock). The input slice is never mutated.
// A missing id in channelsIDM is kept for request_path (downstream consistency
// error) and dropped for task_plugin_identity, matching the previous filters.
func filterCandidateIDs(ids []int, modelName string, filters []dto.ChannelFilter) (kept []int, emptiedBy dto.ChannelFilterKind) {
if len(ids) == 0 {
return ids, ""
}
kept = ids
for _, kind := range filterEvalOrder {
kindFilters := filtersByKind(filters, kind)
if len(kindFilters) == 0 {
continue
}
next := make([]int, 0, len(kept))
for _, id := range kept {
channel, exists := channelsIDM[id]
if candidatePassesKindFilters(channel, exists, modelName, kind, kindFilters) {
next = append(next, id)
}
}
if len(kept) > 0 && len(next) == 0 {
return next, kind
}
kept = next
}
return kept, ""
}
func filtersByKind(filters []dto.ChannelFilter, kind dto.ChannelFilterKind) []dto.ChannelFilter {
var matched []dto.ChannelFilter
for _, filter := range filters {
if filter.Kind == kind {
matched = append(matched, filter)
}
}
return matched
}
func candidatePassesKindFilters(ch *Channel, exists bool, modelName string, kind dto.ChannelFilterKind, filters []dto.ChannelFilter) bool {
if kind == dto.FilterRequestPath && !exists {
return true
}
if !exists || ch == nil {
return false
}
for _, filter := range filters {
if !channelMatchesFilter(ch, modelName, filter) {
return false
}
}
return true
}
func channelMatchesFilter(ch *Channel, modelName string, filter dto.ChannelFilter) bool {
switch filter.Kind {
case dto.FilterRequestPath:
if filter.RequestPath == "" {
return true
}
if ch.Type != constant.ChannelTypeAdvancedCustom {
return true
}
config := ch.GetOtherSettings().AdvancedCustom
return config != nil && config.SupportsPathForModel(filter.RequestPath, modelName)
case dto.FilterTaskPluginIdentity:
if ch.Type == constant.ChannelTypeTaskPlugin {
return filter.TaskPluginKey != "" && ch.GetSetting().TaskPluginKey == filter.TaskPluginKey
}
return filter.TaskPluginKey == "" || slices.Contains(filter.TaskPluginChannelTypes, ch.Type)
default:
return true
}
}
+218
View File
@@ -0,0 +1,218 @@
package model
import (
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
kitdto "github.com/QuantumNous/new-api/relaykit/dto"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestFilterCandidateIDs(t *testing.T) {
alphaSetting := `{"task_plugin_key":"alpha"}`
betaSetting := `{"task_plugin_key":"beta"}`
alpha := &Channel{Id: 900001, Type: constant.ChannelTypeTaskPlugin, Status: common.ChannelStatusEnabled, Setting: &alphaSetting}
beta := &Channel{Id: 900002, Type: constant.ChannelTypeTaskPlugin, Status: common.ChannelStatusEnabled, Setting: &betaSetting}
ordinary := &Channel{Id: 900003, Type: constant.ChannelTypeOpenAI, Status: common.ChannelStatusEnabled}
kling := &Channel{Id: 900004, Type: constant.ChannelTypeKling, Status: common.ChannelStatusEnabled}
jimeng := &Channel{Id: 900005, Type: constant.ChannelTypeJimeng, Status: common.ChannelStatusEnabled}
matchingCustom := &Channel{Id: 900010, Type: constant.ChannelTypeAdvancedCustom, Status: common.ChannelStatusEnabled}
matchingCustom.SetOtherSettings(kitdto.ChannelOtherSettings{
AdvancedCustom: &kitdto.AdvancedCustomConfig{
Routes: []kitdto.AdvancedCustomRoute{{
IncomingPath: "/v1/chat/completions",
Models: []string{"gpt-4"},
}},
},
})
otherCustom := &Channel{Id: 900011, Type: constant.ChannelTypeAdvancedCustom, Status: common.ChannelStatusEnabled}
otherCustom.SetOtherSettings(kitdto.ChannelOtherSettings{
AdvancedCustom: &kitdto.AdvancedCustomConfig{
Routes: []kitdto.AdvancedCustomRoute{{
IncomingPath: "/v1/responses",
Models: []string{"gpt-4"},
}},
},
})
pathFilter := dto.ChannelFilter{Kind: dto.FilterRequestPath, RequestPath: "/v1/chat/completions"}
emptyPathFilter := dto.ChannelFilter{Kind: dto.FilterRequestPath, RequestPath: ""}
tests := []struct {
name string
ids []int
modelName string
filters []dto.ChannelFilter
wantKept []int
wantEmpty dto.ChannelFilterKind
}{
{
name: "identity keeps matching type-59 key",
ids: []int{900001, 900002},
modelName: "shared",
filters: identityFilters("alpha", nil),
wantKept: []int{900001},
},
{
name: "identity empty key drops all type-59",
ids: []int{900001, 900002},
modelName: "shared",
filters: identityFilters("", nil),
wantKept: []int{},
wantEmpty: dto.FilterTaskPluginIdentity,
},
{
name: "identity empty key keeps ordinary channel",
ids: []int{900003},
modelName: "ordinary",
filters: identityFilters("", nil),
wantKept: []int{900003},
},
{
name: "identity keeps matching legacy type",
ids: []int{900004, 900005},
modelName: "legacy",
filters: identityFilters("legacy-alpha", []int{constant.ChannelTypeKling}),
wantKept: []int{900004},
},
{
name: "identity keeps all listed legacy types",
ids: []int{900004, 900005},
modelName: "legacy",
filters: identityFilters("legacy-alpha", []int{constant.ChannelTypeKling, constant.ChannelTypeJimeng}),
wantKept: []int{900004, 900005},
},
{
name: "identity keyed with no types drops legacy",
ids: []int{900004, 900005},
modelName: "legacy",
filters: identityFilters("legacy-alpha", nil),
wantKept: []int{},
wantEmpty: dto.FilterTaskPluginIdentity,
},
{
name: "identity drops missing cache entry",
ids: []int{900004, 999999},
modelName: "legacy",
filters: identityFilters("legacy-alpha", []int{constant.ChannelTypeKling}),
wantKept: []int{900004},
},
{
name: "empty request path is a passthrough including missing ids",
ids: []int{900003, 900010, 999999},
modelName: "gpt-4",
filters: []dto.ChannelFilter{emptyPathFilter},
wantKept: []int{900003, 900010, 999999},
},
{
name: "request path keeps missing cache entry for consistency",
ids: []int{900003, 999999},
modelName: "gpt-4",
filters: []dto.ChannelFilter{pathFilter},
wantKept: []int{900003, 999999},
},
{
name: "request path keeps matching type-58 and ordinary",
ids: []int{900003, 900010, 900011},
modelName: "gpt-4",
filters: []dto.ChannelFilter{pathFilter},
wantKept: []int{900003, 900010},
},
{
name: "request path empties when only unmatched type-58 remains",
ids: []int{900011},
modelName: "gpt-4",
filters: []dto.ChannelFilter{pathFilter},
wantKept: []int{},
wantEmpty: dto.FilterRequestPath,
},
{
name: "intersection attributes empty set to identity after path keeps candidates",
ids: []int{900001, 900010},
modelName: "gpt-4",
filters: []dto.ChannelFilter{pathFilter, identityFilters("missing", nil)[0]},
wantKept: []int{},
wantEmpty: dto.FilterTaskPluginIdentity,
},
{
name: "intersection attributes empty set to path when path runs first",
ids: []int{900011},
modelName: "gpt-4",
filters: []dto.ChannelFilter{identityFilters("", nil)[0], pathFilter},
wantKept: []int{},
wantEmpty: dto.FilterRequestPath,
},
}
channelSyncLock.Lock()
previous := channelsIDM
channelsIDM = map[int]*Channel{
900001: alpha,
900002: beta,
900003: ordinary,
900004: kling,
900005: jimeng,
900010: matchingCustom,
900011: otherCustom,
}
t.Cleanup(func() {
channelsIDM = previous
channelSyncLock.Unlock()
})
for _, testCase := range tests {
t.Run(testCase.name, func(t *testing.T) {
kept, emptiedBy := filterCandidateIDs(testCase.ids, testCase.modelName, testCase.filters)
if testCase.wantKept == nil {
assert.Nil(t, kept)
} else {
assert.Equal(t, testCase.wantKept, kept)
}
assert.Equal(t, testCase.wantEmpty, emptiedBy)
})
}
}
func TestChannelSatisfiesFilters(t *testing.T) {
alphaSetting := `{"task_plugin_key":"alpha"}`
alpha := &Channel{Id: 1, Type: constant.ChannelTypeTaskPlugin, Setting: &alphaSetting}
ordinary := &Channel{Id: 2, Type: constant.ChannelTypeOpenAI}
custom := &Channel{Id: 3, Type: constant.ChannelTypeAdvancedCustom}
custom.SetOtherSettings(kitdto.ChannelOtherSettings{
AdvancedCustom: &kitdto.AdvancedCustomConfig{
Routes: []kitdto.AdvancedCustomRoute{{
IncomingPath: "/v1/chat/completions",
Models: []string{"gpt-4"},
}},
},
})
ok, kind := ChannelSatisfiesFilters(nil, "gpt-4", nil)
assert.False(t, ok)
assert.Equal(t, dto.ChannelFilterKind(""), kind)
ok, kind = ChannelSatisfiesFilters(alpha, "shared", identityFilters("alpha", nil))
require.True(t, ok)
assert.Equal(t, dto.ChannelFilterKind(""), kind)
ok, kind = ChannelSatisfiesFilters(alpha, "shared", identityFilters("beta", nil))
assert.False(t, ok)
assert.Equal(t, dto.FilterTaskPluginIdentity, kind)
ok, kind = ChannelSatisfiesFilters(ordinary, "gpt-4", []dto.ChannelFilter{{
Kind: dto.FilterRequestPath,
RequestPath: "/v1/chat/completions",
}})
require.True(t, ok)
assert.Equal(t, dto.ChannelFilterKind(""), kind)
ok, kind = ChannelSatisfiesFilters(custom, "gpt-4", []dto.ChannelFilter{{
Kind: dto.FilterRequestPath,
RequestPath: "/v1/responses",
}})
assert.False(t, ok)
assert.Equal(t, dto.FilterRequestPath, kind)
}
+15
View File
@@ -121,6 +121,8 @@ func formatUserLogs(logs []*Log, startIdx int) {
if otherMap != nil {
// Remove admin-only debug fields.
delete(otherMap, "admin_info")
// Remove diagnostics reserved for root.
delete(otherMap, "root_info")
// Remove operation-audit details (operator/route info), admin-only.
delete(otherMap, "audit_info")
// delete(otherMap, "reject_reason")
@@ -131,6 +133,19 @@ func formatUserLogs(logs []*Log, startIdx int) {
assignDisplayLogIds(logs, startIdx)
}
// FormatAdminLogs removes root-only diagnostics while retaining operational
// admin_info. Root callers must not pass their results through this formatter.
func FormatAdminLogs(logs []*Log) {
for i := range logs {
otherMap, _ := common.StrToMap(logs[i].Other)
if otherMap == nil {
continue
}
delete(otherMap, "root_info")
logs[i].Other = common.MapToJsonStr(otherMap)
}
}
func GetLogByTokenId(tokenId int) (logs []*Log, err error) {
order := "id desc"
if common.UsingLogDatabase(common.DatabaseTypeClickHouse) {
+48
View File
@@ -5,6 +5,7 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
@@ -33,3 +34,50 @@ func TestFormatUserLogsStripsQuotaSaturation(t *testing.T) {
// Non-admin billing fields remain visible.
require.Contains(t, parsed, "model_price")
}
func TestTaskPluginLogVisibilityIsRoleSeparated(t *testing.T) {
other := common.MapToJsonStr(map[string]interface{}{
"model_price": 1.25,
"admin_info": map[string]interface{}{
"task_plugin": map[string]interface{}{
"key": "document-parser",
"name": "Document Parser",
"version": "1.2.3",
},
},
"root_info": map[string]interface{}{
"upstream_task_id": "upstream-private",
"task_plugin": map[string]interface{}{
"generation": 42,
},
},
})
t.Run("user", func(t *testing.T) {
logs := []*Log{{Other: other}}
formatUserLogs(logs, 0)
parsed, err := common.StrToMap(logs[0].Other)
require.NoError(t, err)
assert.NotContains(t, parsed, "admin_info")
assert.NotContains(t, parsed, "root_info")
assert.Equal(t, 1.25, parsed["model_price"])
})
t.Run("admin", func(t *testing.T) {
logs := []*Log{{Other: other}}
FormatAdminLogs(logs)
parsed, err := common.StrToMap(logs[0].Other)
require.NoError(t, err)
assert.Contains(t, parsed, "admin_info")
assert.NotContains(t, parsed, "root_info")
})
t.Run("root", func(t *testing.T) {
parsed, err := common.StrToMap(other)
require.NoError(t, err)
assert.Contains(t, parsed, "admin_info")
assert.Contains(t, parsed, "root_info")
})
}
+1
View File
@@ -323,6 +323,7 @@ func migrateDB() error {
&TopUp{},
&QuotaData{},
&Task{},
&TaskPlugin{},
&Model{},
&Vendor{},
&PrefillGroup{},
+21
View File
@@ -6,6 +6,8 @@ import (
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/setting"
"github.com/QuantumNous/new-api/setting/config"
"github.com/QuantumNous/new-api/setting/operation_setting"
@@ -52,6 +54,13 @@ func InitOptionMap() {
common.OptionMap["DisplayTokenStatEnabled"] = strconv.FormatBool(common.DisplayTokenStatEnabled)
common.OptionMap["DrawingEnabled"] = strconv.FormatBool(common.DrawingEnabled)
common.OptionMap["TaskEnabled"] = strconv.FormatBool(common.TaskEnabled)
common.OptionMap["TaskPluginEnabled"] = strconv.FormatBool(constant.TaskPluginEnabled)
jsplugin.DefaultRegistry.SetEnabled(constant.TaskPluginEnabled)
common.OptionMap["TaskPluginOverrideEnabled"] = strconv.FormatBool(constant.TaskPluginOverrideEnabled)
jsplugin.DefaultRegistry.SetOverrideEnabled(constant.TaskPluginOverrideEnabled)
common.OptionMap[setting.TaskPluginMarketplaceSourcesKey] = setting.TaskPluginMarketplaceSources2JsonString()
common.OptionMap[setting.TaskPluginDisabledFactoryKeysKey] = "[]"
jsplugin.DefaultRegistry.SetDisabledFactoryKeys(nil)
common.OptionMap["DataExportEnabled"] = strconv.FormatBool(common.DataExportEnabled)
common.OptionMap["ChannelDisableThreshold"] = strconv.FormatFloat(common.ChannelDisableThreshold, 'f', -1, 64)
common.OptionMap["EmailDomainRestrictionEnabled"] = strconv.FormatBool(common.EmailDomainRestrictionEnabled)
@@ -73,6 +82,7 @@ func InitOptionMap() {
common.OptionMap["SystemName"] = common.SystemName
common.OptionMap["Logo"] = common.Logo
common.OptionMap["ServerAddress"] = ""
common.OptionMap["TaskPublicAddress"] = system_setting.TaskPublicAddress
common.OptionMap["WorkerUrl"] = system_setting.WorkerUrl
common.OptionMap["WorkerValidKey"] = system_setting.WorkerValidKey
common.OptionMap["WorkerAllowHttpImageRequestEnabled"] = strconv.FormatBool(system_setting.WorkerAllowHttpImageRequestEnabled)
@@ -352,6 +362,12 @@ func updateOptionMap(key string, value string) (err error) {
common.DrawingEnabled = boolValue
case "TaskEnabled":
common.TaskEnabled = boolValue
case "TaskPluginEnabled":
constant.TaskPluginEnabled = boolValue
jsplugin.DefaultRegistry.SetEnabled(boolValue)
case "TaskPluginOverrideEnabled":
constant.TaskPluginOverrideEnabled = boolValue
jsplugin.DefaultRegistry.SetOverrideEnabled(boolValue)
case "DataExportEnabled":
common.DataExportEnabled = boolValue
case "DefaultCollapseSidebar":
@@ -394,6 +410,9 @@ func updateOptionMap(key string, value string) (err error) {
ratio_setting.SetExposeRatioEnabled(boolValue)
}
}
if key == setting.TaskPluginDisabledFactoryKeysKey {
jsplugin.DefaultRegistry.SetDisabledFactoryKeys(setting.ParseTaskPluginDisabledFactoryKeys(value))
}
switch key {
case "EmailDomainWhitelist":
common.EmailDomainWhitelist = strings.Split(value, ",")
@@ -410,6 +429,8 @@ func updateOptionMap(key string, value string) (err error) {
common.SMTPToken = value
case "ServerAddress":
system_setting.ServerAddress = value
case "TaskPublicAddress":
system_setting.TaskPublicAddress = value
case "WorkerUrl":
system_setting.WorkerUrl = value
case "WorkerValidKey":
+92
View File
@@ -0,0 +1,92 @@
package model
import (
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/setting"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestTaskPluginEnabledOptionUpdatesRegistry(t *testing.T) {
originalEnabled := constant.TaskPluginEnabled
originalMap := common.OptionMap
common.OptionMap = map[string]string{}
const key = "option-master-off"
source := `
export const meta = {apiVersion: 1, key: "option-master-off", name: "Option Master", version: "1.0.0", author: {name: "Test"}, models: ["option-master-model"], fetchMode: "per_task"};
export function buildSubmitRequest() { return {}; }
export function parseSubmitResponse() { return {}; }
export function buildQueryRequest() { return {}; }
export function parseTaskResult() { return {}; }
`
_, err := jsplugin.DefaultRegistry.RegisterFactory(source, jsplugin.Options{})
require.NoError(t, err)
t.Cleanup(func() {
constant.TaskPluginEnabled = originalEnabled
jsplugin.DefaultRegistry.SetEnabled(originalEnabled)
common.OptionMap = originalMap
})
_, ok := jsplugin.DefaultRegistry.Get(key)
require.True(t, ok)
require.NoError(t, updateOptionMap("TaskPluginEnabled", "false"))
assert.False(t, constant.TaskPluginEnabled)
assert.Equal(t, "false", common.OptionMap["TaskPluginEnabled"])
_, ok = jsplugin.DefaultRegistry.Get(key)
assert.False(t, ok)
require.NoError(t, updateOptionMap("TaskPluginEnabled", "true"))
_, ok = jsplugin.DefaultRegistry.Get(key)
assert.True(t, ok)
}
func TestTaskPluginOverrideEnabledOptionUpdatesRuntimeSwitch(t *testing.T) {
originalEnabled := constant.TaskPluginOverrideEnabled
originalMap := common.OptionMap
common.OptionMap = map[string]string{}
t.Cleanup(func() {
constant.TaskPluginOverrideEnabled = originalEnabled
jsplugin.DefaultRegistry.SetOverrideEnabled(originalEnabled)
common.OptionMap = originalMap
})
require.NoError(t, updateOptionMap("TaskPluginOverrideEnabled", "false"))
assert.False(t, constant.TaskPluginOverrideEnabled)
assert.Equal(t, "false", common.OptionMap["TaskPluginOverrideEnabled"])
}
func TestTaskPluginDisabledFactoryKeysOptionUpdatesRegistry(t *testing.T) {
originalMap := common.OptionMap
common.OptionMap = map[string]string{}
const key = "option-factory-off"
source := `
export const meta = {apiVersion: 1, key: "option-factory-off", name: "Option Factory", version: "1.0.0", author: {name: "Test"}, models: ["option-factory-model"], fetchMode: "per_task"};
export function buildSubmitRequest() { return {}; }
export function parseSubmitResponse() { return {}; }
export function buildQueryRequest() { return {}; }
export function parseTaskResult() { return {}; }
`
_, err := jsplugin.DefaultRegistry.RegisterFactory(source, jsplugin.Options{})
require.NoError(t, err)
t.Cleanup(func() {
jsplugin.DefaultRegistry.SetDisabledFactoryKeys(nil)
common.OptionMap = originalMap
})
_, ok := jsplugin.DefaultRegistry.Get(key)
require.True(t, ok)
require.NoError(t, updateOptionMap(setting.TaskPluginDisabledFactoryKeysKey, `["option-factory-off"]`))
assert.Equal(t, `["option-factory-off"]`, common.OptionMap[setting.TaskPluginDisabledFactoryKeysKey])
_, ok = jsplugin.DefaultRegistry.Get(key)
assert.False(t, ok)
assert.Equal(t, []string{key}, jsplugin.DefaultRegistry.Snapshot().DisabledFactory)
}
+46 -20
View File
@@ -2,6 +2,7 @@ package model
import (
"fmt"
"maps"
"strings"
"sync"
@@ -9,6 +10,7 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/setting/billing_setting"
"github.com/QuantumNous/new-api/setting/ratio_setting"
@@ -16,26 +18,28 @@ import (
)
type Pricing struct {
ModelName string `json:"model_name"`
Description string `json:"description,omitempty"`
Icon string `json:"icon,omitempty"`
Tags string `json:"tags,omitempty"`
VendorID int `json:"vendor_id,omitempty"`
QuotaType int `json:"quota_type"`
ModelRatio float64 `json:"model_ratio"`
ModelPrice float64 `json:"model_price"`
OwnerBy string `json:"owner_by"`
CompletionRatio float64 `json:"completion_ratio"`
CacheRatio *float64 `json:"cache_ratio,omitempty"`
CreateCacheRatio *float64 `json:"create_cache_ratio,omitempty"`
ImageRatio *float64 `json:"image_ratio,omitempty"`
AudioRatio *float64 `json:"audio_ratio,omitempty"`
AudioCompletionRatio *float64 `json:"audio_completion_ratio,omitempty"`
EnableGroup []string `json:"enable_groups"`
SupportedEndpointTypes []constant.EndpointType `json:"supported_endpoint_types"`
BillingMode string `json:"billing_mode,omitempty"`
BillingExpr string `json:"billing_expr,omitempty"`
PricingVersion string `json:"pricing_version,omitempty"`
ModelName string `json:"model_name"`
Description string `json:"description,omitempty"`
Icon string `json:"icon,omitempty"`
Tags string `json:"tags,omitempty"`
VendorID int `json:"vendor_id,omitempty"`
QuotaType int `json:"quota_type"`
ModelRatio float64 `json:"model_ratio"`
ModelPrice float64 `json:"model_price"`
OwnerBy string `json:"owner_by"`
CompletionRatio float64 `json:"completion_ratio"`
CacheRatio *float64 `json:"cache_ratio,omitempty"`
CreateCacheRatio *float64 `json:"create_cache_ratio,omitempty"`
ImageRatio *float64 `json:"image_ratio,omitempty"`
AudioRatio *float64 `json:"audio_ratio,omitempty"`
AudioCompletionRatio *float64 `json:"audio_completion_ratio,omitempty"`
EnableGroup []string `json:"enable_groups"`
SupportedEndpointTypes []constant.EndpointType `json:"supported_endpoint_types"`
BillingMode string `json:"billing_mode,omitempty"`
BillingExpr string `json:"billing_expr,omitempty"`
BillingUsageSchema map[string]jsplugin.UsageFieldSchema `json:"billing_usage_schema,omitempty"`
BillingUsageExamples []jsplugin.UsageExample `json:"billing_usage_examples,omitempty"`
PricingVersion string `json:"pricing_version,omitempty"`
}
type PricingVendor struct {
@@ -355,6 +359,7 @@ func updatePricing() {
}
pricingMap = make([]Pricing, 0)
pluginGeneration := jsplugin.DefaultRegistry.Generation()
for model, groups := range modelGroupsMap {
pricing := Pricing{
ModelName: model,
@@ -406,6 +411,27 @@ func updatePricing() {
pricing.BillingExpr = expr
}
}
if plugin, ok := pluginGeneration.GetByModel(model); ok && len(plugin.Meta.UsageSchema) > 0 {
pricing.BillingUsageSchema = make(map[string]jsplugin.UsageFieldSchema, len(plugin.Meta.UsageSchema))
for key, field := range plugin.Meta.UsageSchema {
field.Enum = append([]string(nil), field.Enum...)
field.Description = maps.Clone(field.Description)
pricing.BillingUsageSchema[key] = field
}
if len(plugin.Meta.UsageExamples) > 0 {
pricing.BillingUsageExamples = make([]jsplugin.UsageExample, len(plugin.Meta.UsageExamples))
for index, example := range plugin.Meta.UsageExamples {
facts := make(map[string]any, len(example.Facts))
for key, value := range example.Facts {
facts[key] = value
}
pricing.BillingUsageExamples[index] = jsplugin.UsageExample{
Label: example.Label,
Facts: facts,
}
}
}
}
pricingMap = append(pricingMap, pricing)
}
+69
View File
@@ -0,0 +1,69 @@
package model
import (
"fmt"
"testing"
"time"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func pricingUsagePluginSource(version, usageSchema string) string {
return fmt.Sprintf(`
export const meta = {
apiVersion: 1, key: "pricing-usage-probe", name: "Pricing Usage Probe", version: %q, author: {name: "Test"},
models: ["pricing-usage-model"], fetchMode: "per_task", usageSchema: %s
};
export function buildSubmitRequest() { return {}; }
export function parseSubmitResponse() { return {}; }
export function buildQueryRequest() { return {}; }
export function parseTaskResult() { return {}; }
`, version, usageSchema)
}
func TestPricingCarriesTaskUsageSchemaAndRefreshesWithPluginGeneration(t *testing.T) {
resetPricingEndpointTestTables(t)
const pluginKey = "pricing-usage-probe"
initialSource := pricingUsagePluginSource("1.0.0", `{
seconds: {type: "number", unit: "second", description: "Estimated duration."}
}`)
_, err := jsplugin.DefaultRegistry.Register(initialSource, jsplugin.Options{})
require.NoError(t, err)
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(pluginKey) })
insertPricingEndpointChannel(t, 901, constant.ChannelTypeTaskPlugin, dto.ChannelOtherSettings{})
insertPricingEndpointAbility(t, 901, "pricing-usage-model")
insertPricingEndpointAbility(t, 901, "ordinary-model")
initialPricing := pricingByModel(GetPricing())
require.Contains(t, initialPricing, "pricing-usage-model")
require.Contains(t, initialPricing, "ordinary-model")
assert.Equal(t, "second", initialPricing["pricing-usage-model"].BillingUsageSchema["seconds"].Unit)
assert.Equal(t, "Estimated duration.", initialPricing["pricing-usage-model"].BillingUsageSchema["seconds"].Description["en"])
assert.Nil(t, initialPricing["ordinary-model"].BillingUsageSchema)
updatedSource := pricingUsagePluginSource("1.1.0", `{
seconds: {type: "number", unit: "second", description: "Measured duration."},
clips: {type: "number", unit: "count", description: "Generated clip count."}
}`)
_, err = jsplugin.DefaultRegistry.Register(updatedSource, jsplugin.Options{})
require.NoError(t, err)
lastGetPricingTime = time.Now().Add(-2 * time.Minute)
refreshedPricing := pricingByModel(GetPricing())
require.Len(t, refreshedPricing["pricing-usage-model"].BillingUsageSchema, 2)
assert.Equal(t, "Measured duration.", refreshedPricing["pricing-usage-model"].BillingUsageSchema["seconds"].Description["en"])
assert.Equal(t, "count", refreshedPricing["pricing-usage-model"].BillingUsageSchema["clips"].Unit)
}
func pricingByModel(pricings []Pricing) map[string]Pricing {
result := make(map[string]Pricing, len(pricings))
for _, pricing := range pricings {
result[pricing.ModelName] = pricing
}
return result
}
+109 -19
View File
@@ -2,12 +2,14 @@ package model
import (
"bytes"
"context"
"database/sql/driver"
"encoding/json"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/pkg/billingexpr"
commonRelay "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
)
@@ -17,7 +19,7 @@ type TaskStatus string
func (t TaskStatus) ToVideoStatus() string {
var status string
switch t {
case TaskStatusQueued, TaskStatusSubmitted:
case TaskStatusNotStart, TaskStatusQueued, TaskStatusSubmitted:
status = dto.VideoStatusQueued
case TaskStatusInProgress:
status = dto.VideoStatusInProgress
@@ -104,22 +106,53 @@ type TaskPrivateData struct {
Key string `json:"key,omitempty"`
UpstreamTaskID string `json:"upstream_task_id,omitempty"` // 上游真实 task ID
ResultURL string `json:"result_url,omitempty"` // 任务成功后的结果 URL(视频地址等)
// Execution records safe, immutable request provenance. It lives next to
// other private task state so public task DTOs cannot expose it by accident.
Execution *TaskExecutionSnapshot `json:"execution,omitempty"`
// 计费上下文:用于异步退款/差额结算(轮询阶段读取)
BillingSource string `json:"billing_source,omitempty"` // "wallet" 或 "subscription"
SubscriptionId int `json:"subscription_id,omitempty"` // 订阅 ID,用于订阅退款
TokenId int `json:"token_id,omitempty"` // 令牌 ID,用于令牌额度退款
NodeName string `json:"node_name,omitempty"` // 发起任务的节点名,轮询结算阶段据此归属日志而非最后查询节点
BillingContext *TaskBillingContext `json:"billing_context,omitempty"` // 计费参数快照(用于轮询阶段重新计算)
// ResponsesBackground records that the openai_responses create request
// asked for background:true. Every task is durable and survives client
// disconnect regardless; this only echoes the protocol-level request
// attribute back on retrieval snapshots.
ResponsesBackground bool `json:"responses_background,omitempty"`
}
type TaskExecutionSnapshot struct {
RequestID string `json:"request_id,omitempty"`
RequestPath string `json:"request_path,omitempty"`
TaskPlugin *TaskPluginSnapshot `json:"task_plugin,omitempty"`
}
// TaskPluginSnapshot contains credential-free identity only. Plugin source,
// request/response payloads, and channel secrets must never be added here.
type TaskPluginSnapshot struct {
Key string `json:"key"`
Name string `json:"name"`
Version string `json:"version"`
Author *TaskPluginAuthorSnapshot `json:"author,omitempty"`
APIVersion int `json:"api_version"`
Generation uint64 `json:"generation"`
}
type TaskPluginAuthorSnapshot struct {
Name string `json:"name"`
URL string `json:"url,omitempty"`
}
// TaskBillingContext 记录任务提交时的计费参数,以便轮询阶段可以重新计算额度。
type TaskBillingContext struct {
ModelPrice float64 `json:"model_price,omitempty"` // 模型单价
GroupRatio float64 `json:"group_ratio,omitempty"` // 分组倍率
ModelRatio float64 `json:"model_ratio,omitempty"` // 模型倍率
OtherRatios map[string]float64 `json:"other_ratios,omitempty"` // 附加倍率(时长、分辨率等)
OriginModelName string `json:"origin_model_name,omitempty"` // 模型名称,必须为OriginModelName
PerCallBilling bool `json:"per_call_billing,omitempty"` // 按次计费:跳过轮询阶段的差额结算
ModelPrice float64 `json:"model_price,omitempty"` // 模型单价
GroupRatio float64 `json:"group_ratio,omitempty"` // 分组倍率
ModelRatio float64 `json:"model_ratio,omitempty"` // 模型倍率
OtherRatios map[string]float64 `json:"other_ratios,omitempty"` // 附加倍率(时长、分辨率等)
OriginModelName string `json:"origin_model_name,omitempty"` // 模型名称,必须为OriginModelName
PerCallBilling bool `json:"per_call_billing,omitempty"` // 按次计费:跳过轮询阶段的差额结算
TieredSnapshot *billingexpr.BillingSnapshot `json:"tiered_snapshot,omitempty"`
}
// GetUpstreamTaskID 获取上游真实 task ID(用于与 provider 通信)
@@ -334,6 +367,38 @@ func HasUnfinishedSyncTasks() bool {
return err == nil && id != 0
}
func GetByOnlyTaskId(taskId string) (*Task, bool, error) {
if taskId == "" {
return nil, false, nil
}
var task *Task
var err error
err = DB.Where("task_id = ?", taskId).First(&task).Error
exist, err := RecordExist(err)
if err != nil {
return nil, false, err
}
return task, exist, err
}
// GetUniqueByOnlyTaskId resolves a public task identifier only when exactly one
// row owns it. Historical task identifiers were not globally unique, so
// capability-based reads must fail closed instead of selecting an arbitrary
// tenant's row.
func GetUniqueByOnlyTaskId(taskId string) (*Task, bool, error) {
if taskId == "" {
return nil, false, nil
}
var tasks []*Task
if err := DB.Where("task_id = ?", taskId).Order("id").Limit(2).Find(&tasks).Error; err != nil {
return nil, false, err
}
if len(tasks) != 1 {
return nil, false, nil
}
return tasks[0], true, nil
}
func GetByTaskId(userId int, taskId string) (*Task, bool, error) {
if taskId == "" {
return nil, false, nil
@@ -349,24 +414,44 @@ func GetByTaskId(userId int, taskId string) (*Task, bool, error) {
return task, exist, err
}
func GetByTaskIds(userId int, taskIds []any) ([]*Task, error) {
if len(taskIds) == 0 {
func GetByTaskIdsForPlatforms(userID int, platforms []constant.TaskPlatform, taskIDs []string) ([]*Task, error) {
if len(platforms) == 0 || len(taskIDs) == 0 {
return nil, nil
}
var task []*Task
var err error
err = DB.Where("user_id = ? and task_id in (?)", userId, taskIds).
Find(&task).Error
var tasks []*Task
err := DB.
Where("user_id = ? AND platform IN ? AND task_id IN ?", userID, platforms, taskIDs).
Find(&tasks).Error
if err != nil {
return nil, err
}
return task, nil
return tasks, nil
}
// GetTaskForProtocolObservation reloads one public task through the ownership
// boundary used by long-lived plugin protocol observers. A missing task,
// foreign user, and wrong plugin platform are deliberately indistinguishable.
func GetTaskForProtocolObservation(ctx context.Context, userID int, platform constant.TaskPlatform, taskID string) (*Task, bool, error) {
if taskID == "" {
return nil, false, nil
}
var task Task
err := DB.WithContext(ctx).
Where("user_id = ? AND platform = ? AND task_id = ?", userID, platform, taskID).
First(&task).Error
exists, err := RecordExist(err)
if err != nil || !exists {
return nil, exists, err
}
return &task, true, nil
}
func (Task *Task) Insert() error {
var err error
err = DB.Create(Task).Error
return err
return Task.InsertWithContext(context.Background())
}
func (Task *Task) InsertWithContext(ctx context.Context) error {
return DB.WithContext(ctx).Create(Task).Error
}
type taskSnapshot struct {
@@ -514,7 +599,12 @@ func (t *Task) ToOpenAIVideo() *dto.OpenAIVideo {
openAIVideo.Model = t.Properties.OriginModelName
openAIVideo.SetProgressStr(t.Progress)
openAIVideo.CreatedAt = t.CreatedAt
openAIVideo.CompletedAt = t.UpdatedAt
openAIVideo.SetMetadata("url", t.GetResultURL())
if t.Status == TaskStatusSuccess {
if t.FinishTime != 0 {
openAIVideo.CompletedAt = t.FinishTime
} else {
openAIVideo.CompletedAt = t.UpdatedAt
}
}
return openAIVideo
}
+36
View File
@@ -1,6 +1,7 @@
package model
import (
"context"
"encoding/json"
"os"
"sync"
@@ -8,6 +9,7 @@ import (
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
@@ -99,6 +101,40 @@ func insertTask(t *testing.T, task *Task) {
require.NoError(t, DB.Create(task).Error)
}
func TestGetTaskForProtocolObservationScopesOwnerAndPlatform(t *testing.T) {
truncateTables(t)
task := &Task{
TaskID: "task_protocol_scope",
UserId: 7,
Platform: "plugin-a",
Status: TaskStatusInProgress,
}
insertTask(t, task)
got, exists, err := GetTaskForProtocolObservation(context.Background(), 7, "plugin-a", task.TaskID)
require.NoError(t, err)
require.True(t, exists)
assert.Equal(t, task.ID, got.ID)
for _, query := range []struct {
userID int
platform string
}{
{userID: 8, platform: "plugin-a"},
{userID: 7, platform: "plugin-b"},
} {
got, exists, err = GetTaskForProtocolObservation(context.Background(), query.userID, constant.TaskPlatform(query.platform), task.TaskID)
require.NoError(t, err)
assert.False(t, exists)
assert.Nil(t, got)
}
cancelled, cancel := context.WithCancel(context.Background())
cancel()
_, _, err = GetTaskForProtocolObservation(cancelled, 7, "plugin-a", task.TaskID)
require.ErrorIs(t, err, context.Canceled)
}
// ---------------------------------------------------------------------------
// Snapshot / Equal — pure logic tests (no DB)
// ---------------------------------------------------------------------------
+93
View File
@@ -0,0 +1,93 @@
package model
import (
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestTaskToOpenAIVideoDoesNotExposeResultURL(t *testing.T) {
task := &Task{
TaskID: "task_public",
Status: TaskStatusSuccess,
Progress: "100%",
CreatedAt: 10,
UpdatedAt: 20,
Properties: Properties{
OriginModelName: "video-model",
},
PrivateData: TaskPrivateData{
ResultURL: "https://upstream.example/video.mp4?signature=secret",
},
}
video := task.ToOpenAIVideo()
assert.Equal(t, "task_public", video.ID)
assert.Equal(t, "video", video.Object)
assert.Equal(t, "completed", video.Status)
assert.Nil(t, video.Metadata)
encoded, err := common.Marshal(video)
require.NoError(t, err)
assert.NotContains(t, string(encoded), "upstream.example")
assert.NotContains(t, string(encoded), "signature")
}
func TestTaskToOpenAIVideoStatusAndCompletedAt(t *testing.T) {
tests := []struct {
name string
task Task
wantStatus string
wantCompletedAt int64
wantCompletedAtJSON bool
}{
{
name: "not start maps to queued",
task: Task{TaskID: "t1", Status: TaskStatusNotStart, CreatedAt: 10, UpdatedAt: 20},
wantStatus: "queued",
wantCompletedAt: 0,
},
{
name: "in progress omits completed at",
task: Task{TaskID: "t2", Status: TaskStatusInProgress, CreatedAt: 10, UpdatedAt: 20, FinishTime: 30},
wantStatus: "in_progress",
wantCompletedAt: 0,
},
{
name: "success uses finish time",
task: Task{TaskID: "t3", Status: TaskStatusSuccess, CreatedAt: 10, UpdatedAt: 20, FinishTime: 30},
wantStatus: "completed",
wantCompletedAt: 30,
wantCompletedAtJSON: true,
},
{
name: "success falls back to updated at",
task: Task{TaskID: "t4", Status: TaskStatusSuccess, CreatedAt: 10, UpdatedAt: 20},
wantStatus: "completed",
wantCompletedAt: 20,
wantCompletedAtJSON: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
video := tt.task.ToOpenAIVideo()
assert.Equal(t, tt.wantStatus, video.Status)
assert.Equal(t, tt.wantCompletedAt, video.CompletedAt)
encoded, err := common.Marshal(video)
require.NoError(t, err)
var fields map[string]any
require.NoError(t, common.Unmarshal(encoded, &fields))
if tt.wantCompletedAtJSON {
require.Contains(t, fields, "completed_at")
assert.Equal(t, float64(tt.wantCompletedAt), fields["completed_at"])
} else {
assert.NotContains(t, fields, "completed_at")
}
})
}
}
+233
View File
@@ -0,0 +1,233 @@
package model
import (
"crypto/sha256"
"encoding/hex"
"errors"
"sort"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"gorm.io/gorm"
"gorm.io/gorm/clause"
)
type TaskPluginChannelRef struct {
Id int `json:"id"`
Name string `json:"name"`
}
func GetTaskPluginUsage(key string) ([]TaskPluginChannelRef, int64, error) {
var channels []Channel
if err := DB.Where("type = ? AND status = ?", constant.ChannelTypeTaskPlugin, common.ChannelStatusEnabled).Find(&channels).Error; err != nil {
return nil, 0, err
}
refs := make([]TaskPluginChannelRef, 0)
for _, channel := range channels {
if channel.GetSetting().TaskPluginKey == key {
refs = append(refs, TaskPluginChannelRef{Id: channel.Id, Name: channel.Name})
}
}
var inFlight int64
err := DB.Model(&Task{}).Where("platform = ? AND status NOT IN ?", key, []TaskStatus{TaskStatusSuccess, TaskStatusFailure}).Count(&inFlight).Error
return refs, inFlight, err
}
type TaskPlugin struct {
Id int64 `json:"id"`
Key string `json:"key" gorm:"size:128;not null;uniqueIndex:uk_task_plugin_key_version,priority:1"`
APIVersion int `json:"api_version" gorm:"not null"`
Version string `json:"version" gorm:"size:64;not null;uniqueIndex:uk_task_plugin_key_version,priority:2"`
Source string `json:"source" gorm:"type:text;not null"`
SourceHash string `json:"source_hash" gorm:"size:64;not null"`
Enabled bool `json:"enabled" gorm:"not null"`
Active bool `json:"active" gorm:"not null;index"`
CreatedAt int64 `json:"created_at" gorm:"not null"`
Remark string `json:"remark" gorm:"type:text"`
}
func SaveTaskPlugin(plugin *TaskPlugin) error {
return DB.Transaction(func(tx *gorm.DB) error {
var existing TaskPlugin
err := tx.Where(&TaskPlugin{Key: plugin.Key, Version: plugin.Version}).First(&existing).Error
if err == nil {
if existing.SourceHash != plugin.SourceHash {
return errors.New("plugin key and version already exist with different source")
}
if err = tx.Model(&existing).Updates(map[string]any{"enabled": plugin.Enabled, "remark": plugin.Remark}).Error; err != nil {
return err
}
existing.Enabled = plugin.Enabled
existing.Remark = plugin.Remark
*plugin = existing
return nil
}
if !errors.Is(err, gorm.ErrRecordNotFound) {
return err
}
plugin.CreatedAt = time.Now().Unix()
var count int64
if err = tx.Model(&TaskPlugin{}).Where(&TaskPlugin{Key: plugin.Key, Active: true}).Count(&count).Error; err != nil {
return err
}
plugin.Active = count == 0
return tx.Create(plugin).Error
})
}
func ListTaskPluginVersions(key string) ([]TaskPlugin, error) {
var plugins []TaskPlugin
err := DB.Where(&TaskPlugin{Key: key}).Order("created_at DESC, id DESC").Find(&plugins).Error
return plugins, err
}
func ListTaskPlugins() ([]TaskPlugin, error) {
var plugins []TaskPlugin
err := DB.
Order(clause.OrderByColumn{Column: clause.Column{Name: "key"}}).
Order(clause.OrderByColumn{Column: clause.Column{Name: "created_at"}, Desc: true}).
Order(clause.OrderByColumn{Column: clause.Column{Name: "id"}, Desc: true}).
Find(&plugins).Error
return plugins, err
}
func GetTaskPluginVersion(key, version string) (*TaskPlugin, error) {
var plugin TaskPlugin
query := DB.Where(&TaskPlugin{Key: key})
if version == "" {
query = query.Where(&TaskPlugin{Active: true})
} else {
query = query.Where(&TaskPlugin{Version: version})
}
if err := query.First(&plugin).Error; err != nil {
return nil, err
}
return &plugin, nil
}
func ListActiveTaskPlugins() ([]TaskPlugin, error) {
snapshot, err := GetTaskPluginSyncSnapshot()
return snapshot.Plugins, err
}
type TaskPluginSyncSnapshot struct {
Plugins []TaskPlugin
Revision string
}
// GetTaskPluginSyncSnapshot returns the enabled override set together with a
// deterministic revision of every active database override. Nodes can compare
// the revision even though their local routing-generation counters differ.
func GetTaskPluginSyncSnapshot() (TaskPluginSyncSnapshot, error) {
var activePlugins []TaskPlugin
if err := DB.Where(&TaskPlugin{Active: true}).
Order(clause.OrderByColumn{Column: clause.Column{Name: "key"}}).
Order(clause.OrderByColumn{Column: clause.Column{Name: "version"}}).
Order(clause.OrderByColumn{Column: clause.Column{Name: "id"}}).
Find(&activePlugins).Error; err != nil {
return TaskPluginSyncSnapshot{}, err
}
type revisionEntry struct {
Key string `json:"key"`
APIVersion int `json:"api_version"`
Version string `json:"version"`
SourceHash string `json:"source_hash"`
Enabled bool `json:"enabled"`
}
entries := make([]revisionEntry, 0, len(activePlugins))
enabledPlugins := make([]TaskPlugin, 0, len(activePlugins))
for _, plugin := range activePlugins {
entries = append(entries, revisionEntry{
Key: plugin.Key,
APIVersion: plugin.APIVersion,
Version: plugin.Version,
SourceHash: plugin.SourceHash,
Enabled: plugin.Enabled,
})
if plugin.Enabled {
enabledPlugins = append(enabledPlugins, plugin)
}
}
sort.Slice(entries, func(i, j int) bool {
if entries[i].Key != entries[j].Key {
return entries[i].Key < entries[j].Key
}
return entries[i].Version < entries[j].Version
})
payload, err := common.Marshal(entries)
if err != nil {
return TaskPluginSyncSnapshot{}, err
}
digest := sha256.Sum256(payload)
return TaskPluginSyncSnapshot{
Plugins: enabledPlugins,
Revision: hex.EncodeToString(digest[:]),
}, nil
}
func ActivateTaskPlugin(key, version string) error {
return DB.Transaction(func(tx *gorm.DB) error {
var target TaskPlugin
if err := tx.Where(&TaskPlugin{Key: key, Version: version}).First(&target).Error; err != nil {
return err
}
if err := tx.Model(&TaskPlugin{}).Where(&TaskPlugin{Key: key}).Update("active", false).Error; err != nil {
return err
}
return tx.Model(&target).Updates(map[string]any{"active": true, "enabled": true}).Error
})
}
func SetTaskPluginEnabled(key string, enabled bool) error {
result := DB.Model(&TaskPlugin{}).Where(&TaskPlugin{Key: key, Active: true}).Update("enabled", enabled)
if result.Error != nil {
return result.Error
}
if result.RowsAffected == 0 {
return gorm.ErrRecordNotFound
}
return nil
}
type TaskPluginDeleteResult struct {
DeletedActive bool
Promoted *TaskPlugin
}
func DeleteTaskPluginVersion(key, version string) (TaskPluginDeleteResult, error) {
result := TaskPluginDeleteResult{}
err := DB.Transaction(func(tx *gorm.DB) error {
var plugin TaskPlugin
if err := lockForUpdate(tx).Where(&TaskPlugin{Key: key, Version: version}).First(&plugin).Error; err != nil {
return err
}
result.DeletedActive = plugin.Active
if err := tx.Delete(&plugin).Error; err != nil {
return err
}
if !plugin.Active {
return nil
}
var promoted TaskPlugin
err := lockForUpdate(tx).
Where(&TaskPlugin{Key: key}).
Order("created_at DESC, id DESC").
First(&promoted).Error
if errors.Is(err, gorm.ErrRecordNotFound) {
return nil
}
if err != nil {
return err
}
if err = tx.Model(&promoted).Update("active", true).Error; err != nil {
return err
}
promoted.Active = true
result.Promoted = &promoted
return nil
})
return result, err
}
+58
View File
@@ -0,0 +1,58 @@
package model
import (
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestTaskPluginChannelSelectionFiltersBothCachePaths(t *testing.T) {
truncateTables(t)
priority := int64(0)
weight := uint(1)
baseURL := "https://example.com"
alphaSetting := `{"task_plugin_key":"alpha"}`
betaSetting := `{"task_plugin_key":"beta"}`
channels := []Channel{
{Id: 900001, Type: constant.ChannelTypeTaskPlugin, Status: common.ChannelStatusEnabled, Name: "alpha", Models: "shared", Group: "default", Priority: &priority, Weight: &weight, BaseURL: &baseURL, Setting: &alphaSetting},
{Id: 900002, Type: constant.ChannelTypeTaskPlugin, Status: common.ChannelStatusEnabled, Name: "beta", Models: "shared", Group: "default", Priority: &priority, Weight: &weight, BaseURL: &baseURL, Setting: &betaSetting},
{Id: 900003, Type: constant.ChannelTypeOpenAI, Status: common.ChannelStatusEnabled, Name: "ordinary", Models: "ordinary", Group: "default", Priority: &priority, Weight: &weight},
{Id: 900004, Type: constant.ChannelTypeKling, Status: common.ChannelStatusEnabled, Name: "legacy-alpha", Models: "legacy", Group: "default", Priority: &priority, Weight: &weight},
{Id: 900005, Type: constant.ChannelTypeJimeng, Status: common.ChannelStatusEnabled, Name: "legacy-beta", Models: "legacy", Group: "default", Priority: &priority, Weight: &weight},
}
for i := range channels {
require.NoError(t, channels[i].Insert())
}
selected, err := GetChannel("default", "shared", 0, identityFilters("alpha", nil))
require.NoError(t, err)
require.NotNil(t, selected)
assert.Equal(t, "alpha", selected.Name)
selected, err = GetChannel("default", "shared", 0, identityFilters("", nil))
require.NoError(t, err)
assert.Nil(t, selected)
selected, err = GetChannel("default", "ordinary", 0, identityFilters("", nil))
require.NoError(t, err)
require.NotNil(t, selected)
assert.Equal(t, "ordinary", selected.Name)
selected, err = GetChannel("default", "legacy", 0, identityFilters("legacy-alpha", []int{constant.ChannelTypeKling}))
require.NoError(t, err)
require.NotNil(t, selected)
assert.Equal(t, "legacy-alpha", selected.Name)
selected, err = GetChannel("default", "legacy", 0, identityFilters("legacy-alpha", []int{constant.ChannelTypeKling, constant.ChannelTypeJimeng}))
require.NoError(t, err)
require.NotNil(t, selected)
assert.Contains(t, []string{"legacy-alpha", "legacy-beta"}, selected.Name)
}
func identityFilters(key string, channelTypes []int) []dto.ChannelFilter {
return []dto.ChannelFilter{{
Kind: dto.FilterTaskPluginIdentity,
TaskPluginKey: key,
TaskPluginChannelTypes: channelTypes,
}}
}
+168
View File
@@ -0,0 +1,168 @@
package model
import (
"testing"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
"gorm.io/gorm/utils/tests"
)
func setupTaskPluginModelTest(t *testing.T) {
t.Helper()
originalDB := DB
t.Cleanup(func() { DB = originalDB })
var err error
DB, err = gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
require.NoError(t, DB.AutoMigrate(&TaskPlugin{}))
}
func TestTaskPluginVersionActivationAndSourceImmutability(t *testing.T) {
setupTaskPluginModelTest(t)
v1 := TaskPlugin{Key: "mock", APIVersion: 1, Version: "1.0.0", Source: "v1", SourceHash: "hash-v1", Enabled: true}
require.NoError(t, SaveTaskPlugin(&v1))
assert.True(t, v1.Active)
v2 := TaskPlugin{Key: "mock", APIVersion: 1, Version: "2.0.0", Source: "v2", SourceHash: "hash-v2", Enabled: true}
require.NoError(t, SaveTaskPlugin(&v2))
assert.False(t, v2.Active)
require.NoError(t, ActivateTaskPlugin("mock", "2.0.0"))
active, err := ListActiveTaskPlugins()
require.NoError(t, err)
require.Len(t, active, 1)
assert.Equal(t, "2.0.0", active[0].Version)
conflict := TaskPlugin{Key: "mock", APIVersion: 1, Version: "2.0.0", Source: "changed", SourceHash: "different", Enabled: true}
err = SaveTaskPlugin(&conflict)
require.Error(t, err)
assert.Contains(t, err.Error(), "different source")
require.NoError(t, SetTaskPluginEnabled("mock", false))
active, err = ListActiveTaskPlugins()
require.NoError(t, err)
assert.Empty(t, active)
all, err := ListTaskPlugins()
require.NoError(t, err)
assert.Len(t, all, 2)
deleteResult, err := DeleteTaskPluginVersion("mock", "2.0.0")
require.NoError(t, err)
assert.True(t, deleteResult.DeletedActive)
require.NotNil(t, deleteResult.Promoted)
assert.Equal(t, "1.0.0", deleteResult.Promoted.Version)
versions, err := ListTaskPluginVersions("mock")
require.NoError(t, err)
require.Len(t, versions, 1)
assert.Equal(t, "1.0.0", versions[0].Version)
assert.True(t, versions[0].Active)
}
func TestDeleteActiveTaskPluginPromotesNewestRemainingVersion(t *testing.T) {
setupTaskPluginModelTest(t)
plugins := []*TaskPlugin{
{Key: "promote", APIVersion: 1, Version: "1.0.0", Source: "v1", SourceHash: "hash-v1", Enabled: true},
{Key: "promote", APIVersion: 1, Version: "2.0.0", Source: "v2", SourceHash: "hash-v2", Enabled: false},
{Key: "promote", APIVersion: 1, Version: "3.0.0", Source: "v3", SourceHash: "hash-v3", Enabled: true},
{Key: "promote", APIVersion: 1, Version: "4.0.0", Source: "v4", SourceHash: "hash-v4", Enabled: true},
}
for _, plugin := range plugins {
require.NoError(t, SaveTaskPlugin(plugin))
}
require.NoError(t, DB.Model(plugins[1]).Update("created_at", 200).Error)
require.NoError(t, DB.Model(plugins[2]).Update("created_at", 100).Error)
require.NoError(t, DB.Model(plugins[3]).Update("created_at", 100).Error)
deleteResult, err := DeleteTaskPluginVersion("promote", "1.0.0")
require.NoError(t, err)
assert.True(t, deleteResult.DeletedActive)
require.NotNil(t, deleteResult.Promoted)
assert.Equal(t, "2.0.0", deleteResult.Promoted.Version)
assert.False(t, deleteResult.Promoted.Enabled)
deleteResult, err = DeleteTaskPluginVersion("promote", "2.0.0")
require.NoError(t, err)
assert.True(t, deleteResult.DeletedActive)
require.NotNil(t, deleteResult.Promoted)
assert.Equal(t, "4.0.0", deleteResult.Promoted.Version)
active, err := GetTaskPluginVersion("promote", "")
require.NoError(t, err)
assert.Equal(t, "4.0.0", active.Version)
}
func TestTaskPluginSyncSnapshotRevisionTracksDesiredRuntimeState(t *testing.T) {
setupTaskPluginModelTest(t)
empty, err := GetTaskPluginSyncSnapshot()
require.NoError(t, err)
assert.Empty(t, empty.Plugins)
require.NotEmpty(t, empty.Revision)
v1 := TaskPlugin{
Key: "revision-probe", APIVersion: 1, Version: "1.0.0",
Source: "v1", SourceHash: "hash-v1", Enabled: true,
}
require.NoError(t, SaveTaskPlugin(&v1))
v1Snapshot, err := GetTaskPluginSyncSnapshot()
require.NoError(t, err)
require.Len(t, v1Snapshot.Plugins, 1)
assert.NotEqual(t, empty.Revision, v1Snapshot.Revision)
v2 := TaskPlugin{
Key: "revision-probe", APIVersion: 1, Version: "2.0.0",
Source: "v2", SourceHash: "hash-v2", Enabled: true,
}
require.NoError(t, SaveTaskPlugin(&v2))
inactiveAdded, err := GetTaskPluginSyncSnapshot()
require.NoError(t, err)
assert.Equal(t, v1Snapshot.Revision, inactiveAdded.Revision)
require.NoError(t, DB.Model(&v1).Update("remark", "operator note").Error)
remarkChanged, err := GetTaskPluginSyncSnapshot()
require.NoError(t, err)
assert.Equal(t, v1Snapshot.Revision, remarkChanged.Revision)
require.NoError(t, ActivateTaskPlugin("revision-probe", "2.0.0"))
v2Snapshot, err := GetTaskPluginSyncSnapshot()
require.NoError(t, err)
require.Len(t, v2Snapshot.Plugins, 1)
assert.Equal(t, "2.0.0", v2Snapshot.Plugins[0].Version)
assert.NotEqual(t, v1Snapshot.Revision, v2Snapshot.Revision)
require.NoError(t, SetTaskPluginEnabled("revision-probe", false))
disabled, err := GetTaskPluginSyncSnapshot()
require.NoError(t, err)
assert.Empty(t, disabled.Plugins)
assert.NotEqual(t, v2Snapshot.Revision, disabled.Revision)
}
func TestTaskPluginOrderSQLQuotesMySQLKeyColumn(t *testing.T) {
db, err := gorm.Open(tests.DummyDialector{}, &gorm.Config{DryRun: true})
require.NoError(t, err)
var sqls []string
require.NoError(t, db.Callback().Query().After("gorm:query").Register("test:capture_task_plugin_sql", func(tx *gorm.DB) {
sqls = append(sqls, tx.Statement.SQL.String())
}))
originalDB := DB
t.Cleanup(func() { DB = originalDB })
DB = db
_, err = ListTaskPlugins()
require.NoError(t, err)
_, err = GetTaskPluginSyncSnapshot()
require.NoError(t, err)
require.Len(t, sqls, 2)
for _, sql := range sqls {
assert.Contains(t, sql, "`key`")
assert.NotRegexp(t, `(?i)ORDER BY[[:space:]]+key([[:space:],]|$)`, sql)
}
}
+59 -8
View File
@@ -109,10 +109,11 @@ func usesRequestProbe(node ast.Node) bool {
}
type cachedEntry struct {
prog *vm.Program
usedVars map[string]bool
requestRules []RequestRuleTrace
version int
prog *vm.Program
usedVars map[string]bool
usedUsageKeys map[string]bool
requestRules []RequestRuleTrace
version int
}
var (
@@ -137,6 +138,7 @@ var compileEnvPrototypeV1 = map[string]interface{}{
"_trace_int": func(int, bool, int) int { return 1 },
"header": func(string) string { return "" },
"param": func(string) interface{} { return nil },
"u": func(string) interface{} { return nil },
"has": func(interface{}, string) bool { return false },
"hour": func(string) int { return 0 },
"minute": func(string) int { return 0 },
@@ -196,10 +198,11 @@ func compileEntryFromCacheByHash(exprStr, hash string) (*cachedEntry, error) {
}
entry := &cachedEntry{
prog: prog,
usedVars: extractUsedVars(prog),
requestRules: patcher.requestRules,
version: version,
prog: prog,
usedVars: extractUsedVars(prog),
usedUsageKeys: extractUsedUsageKeys(prog),
requestRules: patcher.requestRules,
version: version,
}
cacheMu.Lock()
if len(cache) >= maxCacheSize {
@@ -244,6 +247,27 @@ func extractUsedVars(prog *vm.Program) map[string]bool {
return vars
}
func extractUsedUsageKeys(prog *vm.Program) map[string]bool {
keys := make(map[string]bool)
ast.Find(prog.Node(), func(node ast.Node) bool {
call, ok := node.(*ast.CallNode)
if !ok || len(call.Arguments) != 1 {
return false
}
callee, ok := call.Callee.(*ast.IdentifierNode)
if !ok || callee.Value != "u" {
return false
}
literal, ok := call.Arguments[0].(*ast.StringNode)
if !ok {
return false
}
keys[strings.TrimSpace(literal.Value)] = true
return false
})
return keys
}
// UsedVars returns the set of identifier names referenced by an expression.
// The result is cached alongside the compiled program. Returns nil for empty input.
func UsedVars(exprStr string) map[string]bool {
@@ -271,6 +295,33 @@ func UsedVars(exprStr string) map[string]bool {
return nil
}
// UsedUsageKeys returns literal keys referenced by u("...") calls. Calls with
// dynamic arguments are intentionally omitted because they cannot be
// validated statically.
func UsedUsageKeys(exprStr string) map[string]bool {
if exprStr == "" {
return nil
}
hash := ExprHashString(exprStr)
cacheMu.RLock()
if entry, ok := cache[hash]; ok {
cacheMu.RUnlock()
return entry.usedUsageKeys
}
cacheMu.RUnlock()
if _, err := compileFromCacheByHash(exprStr, hash); err != nil {
return nil
}
cacheMu.RLock()
entry, ok := cache[hash]
cacheMu.RUnlock()
if ok {
return entry.usedUsageKeys
}
return nil
}
// InvalidateCache clears the compiled-expression cache.
// Called when billing rules are updated.
func InvalidateCache() {
+21
View File
@@ -0,0 +1,21 @@
package billingexpr_test
import (
"testing"
"github.com/QuantumNous/new-api/pkg/billingexpr"
"github.com/stretchr/testify/assert"
)
func TestUsedUsageKeysExtractsLiteralCallsAndSkipsDynamicArguments(t *testing.T) {
expression := `tier("base", u(" seconds ") * 0.4 + (u("clips") > 0 ? u("clips") * 0.1 : 0)) + (u(header("usage-key")) == nil ? 0 : 0)`
keys := billingexpr.UsedUsageKeys(expression)
assert.Equal(t, map[string]bool{"seconds": true, "clips": true}, keys)
}
func TestUsedUsageKeysReturnsNilForEmptyOrInvalidExpressions(t *testing.T) {
assert.Nil(t, billingexpr.UsedUsageKeys(""))
assert.Nil(t, billingexpr.UsedUsageKeys(`tier("broken",`))
}
+79
View File
@@ -144,6 +144,85 @@ The usage-log UI treats `request_rules` as the authoritative rule list and rende
---
## Task Usage Expressions
Task plugins can expose validated, provider-specific billing facts through
`meta.usageSchema`. Expressions read those facts with `u("key")`. A literal key
must be declared by the plugin schema before the expression can be saved.
Numeric facts are finite, non-negative values in the declared canonical unit
(`second`, `count`, `token`, or `credit`); enum facts are exact strings from the
declared value list. The schema description is display-only metadata and never
affects evaluation. `token` is the host unit for upstream billing tokens (for
example doubao `usage.completion_tokens`). `credit` is the host unit for vendor
resource-pack units (for example kling `final_unit_deduction`). Both share the
int32 quota bound, not the 3600-second / 128-count limits.
Task usage billing has a deliberately different conversion rule from token
billing:
```
task quota = expression output in USD * QuotaPerUnit * groupRatio
token quota = expression output in $/1M tokens / 1,000,000 * QuotaPerUnit * groupRatio
```
In other words, a task expression already returns the request's dollar cost.
For example, `u("seconds") * 0.4` means $0.40 per second. Engine semantics do
not divide task output by one million.
The visual editor generates, and public pricing displays recognize, exactly
these canonical task shapes. Expressions outside these shapes remain valid in
raw mode but fall back to the special-expression display:
```
# Flat unit pricing
tier("base", u("seconds") * 0.4)
# Enum tiers (conditions may combine enum comparisons with &&)
u("mode") == "pro"
? tier("pro", u("seconds") * 0.8)
: tier("std", u("seconds") * 0.4)
# Optional constant plus multiple numeric usage terms
tier("base", 0.1 + u("seconds") * 0.4 + u("clips") * 0.05)
# Upstream token overlay (doubao Seedance tokens)
# The editor takes a $/1M token input and emits the / 1000000 literal.
# Engine semantics are unchanged: the expression still returns USD.
tier("base", u("tokens") * 9.8 / 1000000)
# Vendor credit overlay (kling resource-pack units)
# The coefficient is the real $/credit price; no /1M scale.
tier("base", u("units") * 0.14)
```
The tier body is an optional non-negative constant plus one or more
`u("<number field>") * <unit price>` terms. Token-unit fields use the
canonical scaled shape `u("<field>") * <dollars per 1M tokens> / 1000000`.
Credit, second, and count fields keep the bare `u("<field>") * <unit price>`
shape. Tier conditions are equality checks
between declared enum fields and values, optionally joined by `&&`, with
chained ternaries following the same ordering rules as token tiers. Numeric
range tiers are not part of the current canonical shape. Request rules after
`|||` remain orthogonal and use the same syntax as token expressions.
Before saving, the host compiles every expression, rejects literal `u()` keys
that the selected task plugin did not declare, and smoke-tests usage vectors.
Every numeric field is exercised at 0, 1, and its host-owned unit ceiling
(`second` 3600, `count` 128, `token`/`credit` int32 max). Enum values are exercised as a
Cartesian product. Smoke vectors are capped at
64, reducing oversized enum dimensions to their first and last values. Every
evaluated result must be finite and non-negative.
Submission evaluates the expression with request-derived usage facts and
freezes both the expression and those facts in the billing snapshot. On
completion, the host overlays completion facts on the frozen submission facts
key by key, so measured values replace estimates while facts omitted by the
completion hook retain their submission values. The same expression is then
evaluated again; a changed fact can therefore produce a settlement delta and a
different matched tier. Evaluation failure keeps the pre-consumed charge.
---
## Architecture
### Data Flow
+6
View File
@@ -102,6 +102,12 @@ func runProgram(prog *vm.Program, requestRules []RequestRuleTrace, params TokenP
}
return result.Value()
},
"u": func(name string) interface{} {
if request.Usage == nil {
return nil
}
return request.Usage[strings.TrimSpace(name)]
},
"has": func(source interface{}, substr string) bool {
if source == nil || substr == "" {
return false
+3
View File
@@ -6,6 +6,9 @@ import "github.com/QuantumNous/new-api/common"
// expression version. This is the central dispatch point for future versions
// that may use a different conversion formula.
func quotaConversion(exprOutput float64, snap *BillingSnapshot) float64 {
if snap.TaskUsageBilling {
return exprOutput * snap.QuotaPerUnit
}
switch snap.ExprVersion {
default: // v1: coefficients are $/1M tokens prices
return exprOutput / 1_000_000 * snap.QuotaPerUnit
+19
View File
@@ -0,0 +1,19 @@
package billingexpr
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestTaskUsageExpressionUsesFactsAndTaskQuotaConversion(t *testing.T) {
expression := `tier("1080p", u("seconds") * (u("resolution") == "1080p" ? 0.4 : 0.2))`
cost, trace, err := RunExprWithRequest(expression, TokenParams{}, RequestInput{Usage: map[string]any{"seconds": 10.0, "resolution": "1080p"}})
require.NoError(t, err)
assert.Equal(t, 4.0, cost)
assert.Equal(t, "1080p", trace.MatchedTier)
result, err := ComputeTieredQuotaWithRequest(&BillingSnapshot{ExprString: expression, ExprHash: ExprHashString(expression), GroupRatio: 2, QuotaPerUnit: 500000, ExprVersion: 1, TaskUsageBilling: true}, TokenParams{}, RequestInput{Usage: map[string]any{"seconds": 10.0, "resolution": "1080p"}})
require.NoError(t, err)
assert.Equal(t, 4_000_000, result.ActualQuotaAfterGroup)
}
+15 -12
View File
@@ -10,6 +10,7 @@ import (
type RequestInput struct {
Headers map[string]string
Body []byte
Usage map[string]any
}
// TokenParams holds all token dimensions passed into an Expr evaluation.
@@ -47,18 +48,20 @@ type TraceResult struct {
// auto-group retry and settlement. It is fully serializable and contains no
// compiled program pointers.
type BillingSnapshot struct {
BillingMode string `json:"billing_mode"`
ModelName string `json:"model_name"`
ExprString string `json:"expr_string"`
ExprHash string `json:"expr_hash"`
GroupRatio float64 `json:"group_ratio"`
EstimatedPromptTokens int `json:"estimated_prompt_tokens"`
EstimatedCompletionTokens int `json:"estimated_completion_tokens"`
EstimatedQuotaBeforeGroup float64 `json:"estimated_quota_before_group"`
EstimatedQuotaAfterGroup int `json:"estimated_quota_after_group"`
EstimatedTier string `json:"estimated_tier"`
QuotaPerUnit float64 `json:"quota_per_unit"`
ExprVersion int `json:"expr_version"`
BillingMode string `json:"billing_mode"`
ModelName string `json:"model_name"`
ExprString string `json:"expr_string"`
ExprHash string `json:"expr_hash"`
GroupRatio float64 `json:"group_ratio"`
EstimatedPromptTokens int `json:"estimated_prompt_tokens"`
EstimatedCompletionTokens int `json:"estimated_completion_tokens"`
EstimatedQuotaBeforeGroup float64 `json:"estimated_quota_before_group"`
EstimatedQuotaAfterGroup int `json:"estimated_quota_after_group"`
EstimatedTier string `json:"estimated_tier"`
QuotaPerUnit float64 `json:"quota_per_unit"`
ExprVersion int `json:"expr_version"`
TaskUsageBilling bool `json:"task_usage_billing,omitempty"`
UsageFacts map[string]any `json:"usage_facts,omitempty"`
}
// TieredResult holds everything needed after running tiered settlement.
+59
View File
@@ -0,0 +1,59 @@
package jsplugin
import (
"context"
"fmt"
"io"
"os"
"strings"
)
// RunCLI implements the `new-api plugin` subcommand: linting a plugin source
// and replaying a golden fixture against it. It returns a process exit code.
func RunCLI(args []string, stdout, stderr io.Writer) int {
if len(args) < 2 {
fmt.Fprintln(stderr, "usage: new-api plugin lint <plugin.js> | new-api plugin test <plugin.js> --fixture <fixture.json>")
return 2
}
command, sourcePath := args[0], args[1]
source, err := os.ReadFile(sourcePath)
if err != nil {
fmt.Fprintf(stderr, "read plugin: %v\n", err)
return 1
}
switch command {
case "lint":
if len(args) != 2 {
fmt.Fprintln(stderr, "usage: new-api plugin lint <plugin.js>")
return 2
}
plugin, compileErr := NewRegistry().Register(string(source), Options{Key: sourcePath, Version: "lint"})
if compileErr != nil {
fmt.Fprintf(stderr, "plugin lint failed: %v\n", compileErr)
return 1
}
fmt.Fprintf(stdout, "plugin %s@%s is valid\n", plugin.Meta.Key, plugin.Meta.Version)
return 0
case "test":
if len(args) != 4 || args[2] != "--fixture" || strings.TrimSpace(args[3]) == "" {
fmt.Fprintln(stderr, "usage: new-api plugin test <plugin.js> --fixture <fixture.json>")
return 2
}
fixture, readErr := os.ReadFile(args[3])
if readErr != nil {
fmt.Fprintf(stderr, "read fixture: %v\n", readErr)
return 1
}
report, replayErr := ReplayFixture(context.Background(), string(source), fixture)
if replayErr != nil {
fmt.Fprintf(stderr, "plugin fixture failed after %d/%d cases: %v\n", report.Passed, report.Total, replayErr)
return 1
}
fmt.Fprintf(stdout, "plugin fixture passed: %d/%d cases\n", report.Passed, report.Total)
return 0
default:
fmt.Fprintf(stderr, "unknown plugin command %q\n", command)
return 2
}
}
+38
View File
@@ -0,0 +1,38 @@
package jsplugin
import (
"bytes"
"os"
"path/filepath"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestPluginCLI(t *testing.T) {
tempDir := t.TempDir()
pluginPath := filepath.Join(tempDir, "fixture.js")
fixturePath := filepath.Join(tempDir, "fixture.json")
require.NoError(t, os.WriteFile(pluginPath, []byte(cliFixturePluginSource), 0o600))
require.NoError(t, os.WriteFile(fixturePath, []byte(`{"unixNow":42,"cases":[{"hook":"value","args":[],"expected":42}]}`), 0o600))
var stdout bytes.Buffer
var stderr bytes.Buffer
assert.Equal(t, 0, RunCLI([]string{"lint", pluginPath}, &stdout, &stderr))
assert.Contains(t, stdout.String(), "plugin cli-fixture@1.0.0 is valid")
assert.Empty(t, stderr.String())
stdout.Reset()
assert.Equal(t, 0, RunCLI([]string{"test", pluginPath, "--fixture", fixturePath}, &stdout, &stderr))
assert.Contains(t, stdout.String(), "1/1 cases")
}
const cliFixturePluginSource = `
export const meta = { apiVersion: 1, key: "cli-fixture", name: "CLI Fixture", version: "1.0.0", author: {name: "Test"}, channelTypes: [1003], models: ["fixture-model"], fetchMode: "per_task" };
export function buildSubmitRequest(ctx) { return {url: ctx.baseUrl}; }
export function parseSubmitResponse(ctx, resp) { return {taskId: "task"}; }
export function buildQueryRequest(ctx) { return {url: ctx.baseUrl}; }
export function parseTaskResult(ctx, body) { return body; }
export function value() { return utils.unixNow(); }
`
+579
View File
@@ -0,0 +1,579 @@
package jsplugin
import (
"context"
"errors"
"fmt"
"regexp"
"strings"
"sync"
"time"
"github.com/QuantumNous/new-api/logger"
"github.com/grafana/sobek"
"github.com/grafana/sobek/parser"
)
const (
DefaultCallTimeout = 5 * time.Second
DefaultConcurrency = 8
)
var ErrCallAdmissionTimeout = errors.New("plugin call admission timed out")
const hookErrorMessageLimit = 512
// HookError reports a JavaScript exception thrown by a plugin hook. Message
// is the sanitized JS error message with engine prefixes stripped; it is safe
// to surface to API callers.
type HookError struct {
Hook string
Message string
wrapped error
}
func (e *HookError) Error() string {
if e == nil || e.wrapped == nil {
return "plugin hook failed"
}
return e.wrapped.Error()
}
func (e *HookError) Unwrap() error {
if e == nil {
return nil
}
return e.wrapped
}
func newHookError(hook, rawMessage string, wrapped error) *HookError {
var b strings.Builder
b.Grow(len(rawMessage))
count := 0
for _, r := range rawMessage {
if count >= hookErrorMessageLimit {
break
}
if r < 0x20 || (r >= 0x80 && r <= 0x9F) {
r = ' '
}
b.WriteRune(r)
count++
}
message := b.String()
if message == "" {
message = "plugin hook failed"
}
return &HookError{Hook: hook, Message: message, wrapped: wrapped}
}
func hookErrorFromException(hook string, exc *sobek.Exception, wrapped error) (hookErr *HookError) {
// Reading message/toString executes plugin getters, which can throw again
// and panic sobek. By this point the caller's recover is already consumed,
// so a second panic would crash the process; fall back to a blank message.
defer func() {
if recover() != nil {
hookErr = newHookError(hook, "", wrapped)
}
}()
raw := ""
if exc != nil {
if val := exc.Value(); val != nil && !sobek.IsUndefined(val) && !sobek.IsNull(val) {
gotMessage := false
if obj, ok := val.(*sobek.Object); ok {
if msg := obj.Get("message"); msg != nil && !sobek.IsUndefined(msg) && !sobek.IsNull(msg) {
raw = msg.String()
gotMessage = true
}
}
if !gotMessage {
if exported, ok := val.Export().(string); ok {
raw = exported
} else {
raw = val.String()
}
}
}
}
return newHookError(hook, raw, wrapped)
}
var forbiddenSyntax = regexp.MustCompile(`(?m)(^|[^A-Za-z0-9_$])(async|await|import)([^A-Za-z0-9_$]|$)`)
type Options struct {
Key string
Version string
Timeout time.Duration
Concurrency int
Now func() time.Time
Log func(string)
}
type Engine struct {
key string
version string
timeout time.Duration
now func() time.Time
log func(string)
module *sobek.SourceTextModuleRecord
pool sync.Pool
semaphore chan struct{}
}
type runtimeInstance struct {
runtime *sobek.Runtime
module sobek.ModuleInstance
logContext *runtimeLogContext
}
type runtimeLogContext struct {
context context.Context
}
// Compile performs upload-time syntax checks and compiles an ESM plugin once.
// All Sobek-specific module and runtime handling is intentionally kept here.
func Compile(source string, options Options) (*Engine, error) {
if match := forbiddenSyntax.FindString(sourceWithoutCommentsAndStrings(source)); match != "" {
return nil, fmt.Errorf("unsupported plugin syntax %q: plugins must be synchronous and cannot import modules", strings.TrimSpace(match))
}
resolve := func(_ interface{}, specifier string) (sobek.ModuleRecord, error) {
return nil, fmt.Errorf("plugin imports are disabled: %s", specifier)
}
// Plugin source is untrusted; without this option a sourceMappingURL
// comment makes the parser read arbitrary server files via os.ReadFile.
module, err := sobek.ParseModule(options.Key+".js", source, resolve, parser.WithDisableSourceMaps)
if err != nil {
return nil, fmt.Errorf("compile plugin: %w", err)
}
if err = module.Link(); err != nil {
return nil, fmt.Errorf("link plugin: %w", err)
}
timeout := options.Timeout
if timeout <= 0 {
timeout = DefaultCallTimeout
}
concurrency := options.Concurrency
if concurrency <= 0 {
concurrency = DefaultConcurrency
}
now := options.Now
if now == nil {
now = time.Now
}
engine := &Engine{
key: options.Key,
version: options.Version,
timeout: timeout,
now: now,
log: options.Log,
module: module,
semaphore: make(chan struct{}, concurrency),
}
instance, err := engine.newRuntime(context.Background())
if err != nil {
return nil, err
}
instance.logContext.context = nil
engine.pool.Put(instance)
return engine, nil
}
// Export returns one module export without exposing Sobek values outside the
// engine boundary. It is used for declarative exports such as meta.
func (e *Engine) Export(ctx context.Context, exportName string) (result any, err error) {
select {
case e.semaphore <- struct{}{}:
defer func() { <-e.semaphore }()
case <-ctx.Done():
return nil, ctx.Err()
}
instance, err := e.getRuntime(ctx)
if err != nil {
return nil, err
}
reusable := true
defer func() {
instance.runtime.ClearInterrupt()
instance.logContext.context = nil
if reusable {
e.pool.Put(instance)
}
}()
timedOut := errors.New("plugin export timed out")
timer := time.AfterFunc(e.timeout, func() { instance.runtime.Interrupt(timedOut) })
stopContext := context.AfterFunc(ctx, func() { instance.runtime.Interrupt(ctx.Err()) })
defer stopContext()
defer timer.Stop()
defer func() {
if recovered := recover(); recovered != nil {
reusable = false
switch value := recovered.(type) {
case *sobek.InterruptedError:
err = fmt.Errorf("plugin %s@%s export %s interrupted: %v", e.key, e.version, exportName, value.Value())
case *sobek.Exception:
err = fmt.Errorf("plugin %s@%s export %s failed: %v", e.key, e.version, exportName, value)
default:
panic(recovered)
}
}
}()
value := instance.module.GetBindingValue(exportName)
if value == nil || sobek.IsUndefined(value) {
return nil, fmt.Errorf("plugin export %q not found", exportName)
}
return value.Export(), nil
}
// HasExport reports whether a module export exists. Optional contract hooks
// should be detected with this method instead of relying on engine errors.
func (e *Engine) HasExport(ctx context.Context, exportName string) (bool, error) {
select {
case e.semaphore <- struct{}{}:
defer func() { <-e.semaphore }()
case <-ctx.Done():
return false, ctx.Err()
}
instance, err := e.getRuntime(ctx)
if err != nil {
return false, err
}
defer func() {
instance.logContext.context = nil
e.pool.Put(instance)
}()
value := instance.module.GetBindingValue(exportName)
return value != nil && !sobek.IsUndefined(value), nil
}
// HasCallablePath reports whether an exported value, or a nested member below
// it, exists and is callable.
func (e *Engine) HasCallablePath(ctx context.Context, exportName string, members ...string) (found bool, err error) {
select {
case e.semaphore <- struct{}{}:
defer func() { <-e.semaphore }()
case <-ctx.Done():
return false, ctx.Err()
}
instance, err := e.getRuntime(ctx)
if err != nil {
return false, err
}
reusable := true
defer func() {
instance.runtime.ClearInterrupt()
instance.logContext.context = nil
if reusable {
e.pool.Put(instance)
}
}()
timedOut := errors.New("plugin inspection timed out")
timer := time.AfterFunc(e.timeout, func() { instance.runtime.Interrupt(timedOut) })
stopContext := context.AfterFunc(ctx, func() { instance.runtime.Interrupt(ctx.Err()) })
defer stopContext()
defer timer.Stop()
defer func() {
if recovered := recover(); recovered != nil {
reusable = false
hookName := strings.Join(append([]string{exportName}, members...), ".")
switch value := recovered.(type) {
case *sobek.InterruptedError:
err = fmt.Errorf("plugin %s@%s hook %s inspection interrupted: %v", e.key, e.version, hookName, value.Value())
case *sobek.Exception:
err = fmt.Errorf("plugin %s@%s hook %s inspection failed: %v", e.key, e.version, hookName, value)
default:
panic(recovered)
}
}
}()
value, _, found := resolveExportPath(instance, exportName, members)
if !found {
return false, nil
}
_, callable := sobek.AssertFunction(value)
return callable, nil
}
// Call invokes one named module export and returns its JSON-compatible value.
func (e *Engine) Call(ctx context.Context, exportName string, args ...any) (result any, err error) {
return e.call(ctx, 0, exportName, nil, args...)
}
// CallMember invokes a function stored on an exported object, such as a
// renderer in the renderers export.
func (e *Engine) CallMember(ctx context.Context, exportName, memberName string, args ...any) (result any, err error) {
return e.call(ctx, 0, exportName, []string{memberName}, args...)
}
// CallPath invokes a function nested below an exported object. It is used for
// protocol hooks such as protocols.openai_responses.renderEvents.
func (e *Engine) CallPath(ctx context.Context, exportName string, members []string, args ...any) (result any, err error) {
return e.call(ctx, 0, exportName, members, args...)
}
// CallPathWithAdmissionTimeout gives long-lived observers a separate bound for
// waiting on JavaScript capacity. Once admitted, the hook receives the
// engine's full execution timeout instead of inheriting time already spent in
// the semaphore queue.
func (e *Engine) CallPathWithAdmissionTimeout(
ctx context.Context,
admissionTimeout time.Duration,
exportName string,
members []string,
args ...any,
) (result any, err error) {
return e.call(ctx, admissionTimeout, exportName, members, args...)
}
func (e *Engine) call(
ctx context.Context,
admissionTimeout time.Duration,
exportName string,
members []string,
args ...any,
) (result any, err error) {
if err = e.acquireCallSlot(ctx, admissionTimeout); err != nil {
return nil, err
}
defer func() { <-e.semaphore }()
instance, err := e.getRuntime(ctx)
if err != nil {
return nil, err
}
reusable := true
defer func() {
instance.runtime.ClearInterrupt()
instance.logContext.context = nil
if reusable {
e.pool.Put(instance)
}
}()
hookName := strings.Join(append([]string{exportName}, members...), ".")
timedOut := errors.New("plugin call timed out")
timer := time.AfterFunc(e.timeout, func() { instance.runtime.Interrupt(timedOut) })
stopContext := context.AfterFunc(ctx, func() { instance.runtime.Interrupt(ctx.Err()) })
defer stopContext()
defer timer.Stop()
defer func() {
if recovered := recover(); recovered != nil {
reusable = false
switch value := recovered.(type) {
case *sobek.InterruptedError:
err = fmt.Errorf("plugin %s@%s hook %s interrupted: %v", e.key, e.version, hookName, value.Value())
case *sobek.Exception:
wrapped := fmt.Errorf("plugin %s@%s hook %s failed: %v", e.key, e.version, hookName, value)
err = hookErrorFromException(hookName, value, wrapped)
default:
panic(recovered)
}
}
}()
value, resolvedHookName, found := resolveExportPath(instance, exportName, members)
hookName = resolvedHookName
if !found {
if len(members) == 0 {
return nil, fmt.Errorf("plugin export %q not found", exportName)
}
return nil, fmt.Errorf("plugin hook %q not found", hookName)
}
if value == nil || sobek.IsUndefined(value) {
return nil, fmt.Errorf("plugin export %q not found", exportName)
}
callable, ok := sobek.AssertFunction(value)
if !ok {
return nil, fmt.Errorf("plugin hook %q is not a function", hookName)
}
callArgs := make([]sobek.Value, len(args))
for i, arg := range args {
callArgs[i] = instance.runtime.ToValue(arg)
}
value, err = callable(sobek.Undefined(), callArgs...)
if err != nil {
var interrupted *sobek.InterruptedError
if errors.As(err, &interrupted) {
reusable = false
return nil, fmt.Errorf("plugin %s@%s hook %s failed: %w", e.key, e.version, hookName, err)
}
wrapped := fmt.Errorf("plugin %s@%s hook %s failed: %w", e.key, e.version, hookName, err)
var exc *sobek.Exception
if errors.As(err, &exc) {
return nil, hookErrorFromException(hookName, exc, wrapped)
}
return nil, wrapped
}
return value.Export(), nil
}
func (e *Engine) acquireCallSlot(ctx context.Context, admissionTimeout time.Duration) error {
if admissionTimeout <= 0 {
select {
case e.semaphore <- struct{}{}:
return nil
case <-ctx.Done():
return ctx.Err()
}
}
timer := time.NewTimer(admissionTimeout)
defer timer.Stop()
select {
case e.semaphore <- struct{}{}:
return nil
case <-ctx.Done():
return ctx.Err()
case <-timer.C:
if err := ctx.Err(); err != nil {
return err
}
return fmt.Errorf("%w: plugin %s@%s", ErrCallAdmissionTimeout, e.key, e.version)
}
}
func resolveExportPath(instance *runtimeInstance, exportName string, members []string) (sobek.Value, string, bool) {
value := instance.module.GetBindingValue(exportName)
hookName := exportName
if value == nil || sobek.IsUndefined(value) || sobek.IsNull(value) {
return nil, hookName, false
}
for _, member := range members {
hookName += "." + member
object := value.ToObject(instance.runtime)
own := false
for _, name := range object.GetOwnPropertyNames() {
if name == member {
own = true
break
}
}
if !own {
return nil, hookName, false
}
value = object.Get(member)
if value == nil || sobek.IsUndefined(value) || sobek.IsNull(value) {
return nil, hookName, false
}
}
return value, hookName, true
}
func (e *Engine) getRuntime(ctx context.Context) (*runtimeInstance, error) {
if pooled := e.pool.Get(); pooled != nil {
instance := pooled.(*runtimeInstance)
instance.logContext.context = ctx
return instance, nil
}
return e.newRuntime(ctx)
}
func (e *Engine) newRuntime(ctx context.Context) (instance *runtimeInstance, err error) {
runtime := sobek.New()
logContext := &runtimeLogContext{context: ctx}
logOutput := e.log
if logOutput == nil {
logOutput = func(message string) {
logger.LogDebug(logContext.context, "task_plugin subsystem=runtime event=console message=%q", message)
}
}
if err := injectGlobals(runtime, func() string {
return fmt.Sprintf("[plugin:%s@%s]", e.key, e.version)
}, e.now, logOutput); err != nil {
return nil, fmt.Errorf("inject plugin utils: %w", err)
}
timedOut := errors.New("plugin initialization timed out")
timer := time.AfterFunc(e.timeout, func() { runtime.Interrupt(timedOut) })
defer timer.Stop()
defer func() {
if recovered := recover(); recovered != nil {
if interrupted, ok := recovered.(*sobek.InterruptedError); ok {
instance = nil
err = fmt.Errorf("initialize plugin %s@%s: %v", e.key, e.version, interrupted.Value())
return
}
panic(recovered)
}
}()
promise := runtime.CyclicModuleRecordEvaluate(e.module, func(_ interface{}, specifier string) (sobek.ModuleRecord, error) {
return nil, fmt.Errorf("plugin imports are disabled: %s", specifier)
})
if promise.State() != sobek.PromiseStateFulfilled {
return nil, fmt.Errorf("evaluate plugin: %v", promise.Result().Export())
}
return &runtimeInstance{
runtime: runtime,
module: runtime.GetModuleInstance(e.module),
logContext: logContext,
}, nil
}
func sourceWithoutCommentsAndStrings(source string) string {
var output strings.Builder
output.Grow(len(source))
quote := byte(0)
escaped := false
lineComment := false
blockComment := false
for i := 0; i < len(source); i++ {
current := source[i]
next := byte(0)
if i+1 < len(source) {
next = source[i+1]
}
if lineComment {
if current == '\n' {
lineComment = false
output.WriteByte('\n')
} else {
output.WriteByte(' ')
}
continue
}
if blockComment {
if current == '*' && next == '/' {
blockComment = false
output.WriteString(" ")
i++
} else {
output.WriteByte(' ')
}
continue
}
if quote != 0 {
output.WriteByte(' ')
if escaped {
escaped = false
} else if current == '\\' {
escaped = true
} else if current == quote {
quote = 0
}
continue
}
if current == '/' && next == '/' {
lineComment = true
output.WriteString(" ")
i++
continue
}
if current == '/' && next == '*' {
blockComment = true
output.WriteString(" ")
i++
continue
}
if current == '\'' || current == '"' || current == '`' {
quote = current
output.WriteByte(' ')
continue
}
output.WriteByte(current)
}
return output.String()
}
+372
View File
@@ -0,0 +1,372 @@
package jsplugin
import (
"bytes"
"context"
"errors"
"strings"
"testing"
"time"
"unicode/utf8"
"github.com/QuantumNous/new-api/common"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestEngineCallsESMExportWithInjectedUtils(t *testing.T) {
t.Parallel()
logs := make([]string, 0, 1)
engine, err := Compile(`
export function sign(ctx) {
console.log("called", ctx.name);
return {
now: utils.unixNow(),
digest: utils.hmacSHA256(ctx.message, ctx.secret),
encoded: utils.base64(ctx.message),
};
}
export const meta = { apiVersion: 1, key: "mock" };
`, Options{
Key: "mock", Version: "1.0.0",
Now: func() time.Time { return time.Unix(1234, 0) },
Log: func(message string) { logs = append(logs, message) },
})
require.NoError(t, err)
result, err := engine.Call(context.Background(), "sign", map[string]any{
"name": "fixture", "message": "hello", "secret": "secret",
})
require.NoError(t, err)
assert.Equal(t, map[string]any{
"now": int64(1234), "digest": "88aab3ede8d3adf94d26ab90d3bafd4a2083070c3bcce9c014ee04a443847c0b", "encoded": "aGVsbG8=",
}, result)
assert.Equal(t, []string{"[plugin:mock@1.0.0] called fixture"}, logs)
meta, err := engine.Export(context.Background(), "meta")
require.NoError(t, err)
assert.Equal(t, map[string]any{"apiVersion": int64(1), "key": "mock"}, meta)
}
func TestEngineConsoleLogUsesDebugLoggerAndRequestContext(t *testing.T) {
previousDebug := common.DebugEnabled
common.DebugEnabled = false
t.Cleanup(func() { common.DebugEnabled = previousDebug })
var output bytes.Buffer
common.LogWriterMu.Lock()
previousWriter := gin.DefaultErrorWriter
gin.DefaultErrorWriter = &output
common.LogWriterMu.Unlock()
t.Cleanup(func() {
common.LogWriterMu.Lock()
gin.DefaultErrorWriter = previousWriter
common.LogWriterMu.Unlock()
})
plugin, err := CompilePlugin(`
export const meta = {
apiVersion: 1,
key: "console-debug",
name: "Console debug",
version: "1.2.3",
author: {name: "Test"},
models: ["debug-model"],
fetchMode: "per_task",
};
export function run(label) {
console.log("checkpoint", label);
return true;
}
export function buildSubmitRequest() { return {url: "https://example.com"}; }
export function parseSubmitResponse() { return {taskId: "one"}; }
export function buildQueryRequest() { return {url: "https://example.com"}; }
export function parseTaskResult() { return {status: "SUCCESS"}; }
`, Options{})
require.NoError(t, err)
engine := plugin.Engine
disabledContext := context.WithValue(context.Background(), common.RequestIdKey, "plugin-console-disabled")
_, err = engine.Call(disabledContext, "run", "disabled")
require.NoError(t, err)
assert.Empty(t, output.String())
common.DebugEnabled = true
contextA := context.WithValue(context.Background(), common.RequestIdKey, "plugin-console-request-a")
_, err = engine.Call(contextA, "run", "context-a")
require.NoError(t, err)
logA := output.String()
output.Reset()
contextB := context.WithValue(context.Background(), common.RequestIdKey, "plugin-console-request-b")
_, err = engine.Call(contextB, "run", "context-b")
require.NoError(t, err)
logB := output.String()
output.Reset()
_, err = engine.Call(context.Background(), "run", "background")
require.NoError(t, err)
logBackground := output.String()
assert.Contains(t, logA, "plugin-console-request-a")
assert.NotContains(t, logA, "plugin-console-request-b")
assert.Contains(t, logA, "task_plugin subsystem=runtime event=console")
assert.Contains(t, logA, "[plugin:console-debug@1.2.3] checkpoint context-a")
assert.NotContains(t, logA, "disabled")
assert.Contains(t, logB, "plugin-console-request-b")
assert.NotContains(t, logB, "plugin-console-request-a")
assert.Contains(t, logB, "[plugin:console-debug@1.2.3] checkpoint context-b")
assert.Contains(t, logBackground, "| SYSTEM |")
assert.NotContains(t, logBackground, "plugin-console-request-a")
assert.NotContains(t, logBackground, "plugin-console-request-b")
assert.Contains(t, logBackground, "[plugin:console-debug@1.2.3] checkpoint background")
}
func TestCompileRejectsAsynchronousAndImportedPlugins(t *testing.T) {
t.Parallel()
for name, source := range map[string]string{
"async": `export async function run() {}`,
"static import": `import value from "dependency"; export function run() { return value; }`,
"dynamic import": `export function run() { return import("dependency"); }`,
"top-level await": `const value = await work(); export function run() { return value; }`,
} {
t.Run(name, func(t *testing.T) {
_, err := Compile(source, Options{Key: "invalid"})
require.Error(t, err)
assert.Contains(t, err.Error(), "unsupported plugin syntax")
})
}
_, err := Compile(`export function run() { return "import async await"; }`, Options{Key: "valid"})
require.NoError(t, err)
}
func TestCompileIgnoresSourceMapDirectives(t *testing.T) {
t.Parallel()
// A sourceMappingURL comment must stay inert. Sobek's default loader
// os.ReadFiles the referenced server path during Compile and turns any
// load failure into a compile error, so an unresolvable path compiling
// cleanly proves the loader is disabled.
engine, err := Compile("export function run() { return 1; }\n//# sourceMappingURL=/nonexistent/leak-probe.map\n", Options{Key: "sourcemap"})
require.NoError(t, err)
value, err := engine.Call(context.Background(), "run")
require.NoError(t, err)
assert.Equal(t, int64(1), value)
}
func TestEngineInterruptsLongRunningHook(t *testing.T) {
t.Parallel()
engine, err := Compile(`export function run() { while (true) {} }`, Options{
Key: "loop", Version: "1", Timeout: 20 * time.Millisecond,
})
require.NoError(t, err)
_, err = engine.Call(context.Background(), "run")
require.Error(t, err)
assert.Contains(t, err.Error(), "timed out")
var hookErr *HookError
assert.False(t, errors.As(err, &hookErr), "timeouts must not be HookError")
}
func TestEngineHookErrorExtractsSanitizedJSMessage(t *testing.T) {
t.Parallel()
tests := []struct {
name string
source string
wantMessage string
wantLen int
}{
{
name: "Error object",
source: `export function run() { throw new Error("model is required"); }`,
wantMessage: "model is required",
},
{
name: "raw string throw",
source: `export function run() { throw "raw string"; }`,
wantMessage: "raw string",
},
{
name: "truncates to 512 runes",
source: `export function run() { throw new Error("x".repeat(2000)); }`,
wantLen: 512,
wantMessage: strings.Repeat("x", 512),
},
{
name: "scrubs control characters",
source: "export function run() { throw new Error(\"line1\\nline2\\x1b[31mred\"); }",
wantMessage: "line1 line2 [31mred",
},
{
name: "throwing message getter falls back without crashing",
source: `export function run() { throw {get message() { throw {get message() { return "deep"; }}; }}; }`,
wantMessage: "plugin hook failed",
},
}
for _, testCase := range tests {
t.Run(testCase.name, func(t *testing.T) {
t.Parallel()
engine, err := Compile(testCase.source, Options{Key: "diag", Version: "1.0.0"})
require.NoError(t, err)
_, err = engine.Call(context.Background(), "run")
require.Error(t, err)
var hookErr *HookError
require.True(t, errors.As(err, &hookErr))
assert.Equal(t, "run", hookErr.Hook)
assert.Equal(t, testCase.wantMessage, hookErr.Message)
if testCase.wantLen > 0 {
assert.Equal(t, testCase.wantLen, utf8.RuneCountInString(hookErr.Message))
}
assert.Contains(t, hookErr.Error(), "plugin diag@1.0.0")
assert.NotContains(t, hookErr.Message, "Error:")
assert.NotContains(t, hookErr.Message, "plugin diag@")
})
}
}
func TestEngineReportsProtocolAdmissionTimeoutSeparately(t *testing.T) {
engine, err := Compile(`
export const protocols = {
responses: {renderEvents: function() { return {events: [], done: false}; }},
};
`, Options{Key: "admission", Version: "1", Concurrency: 1})
require.NoError(t, err)
engine.semaphore <- struct{}{}
_, err = engine.CallPathWithAdmissionTimeout(
context.Background(),
time.Nanosecond,
"protocols",
[]string{"responses", "renderEvents"},
)
<-engine.semaphore
require.ErrorIs(t, err, ErrCallAdmissionTimeout)
result, err := engine.CallPathWithAdmissionTimeout(
context.Background(),
time.Second,
"protocols",
[]string{"responses", "renderEvents"},
)
require.NoError(t, err)
assert.Equal(t, map[string]any{"events": []any{}, "done": false}, result)
}
func TestEngineExportInterruptsLongRunningGetter(t *testing.T) {
t.Parallel()
engine, err := Compile(`
export const meta = {
apiVersion: 1,
get name() { while (true) {} },
};
`, Options{Key: "meta-loop", Version: "1.0.0", Timeout: 20 * time.Millisecond})
require.NoError(t, err)
_, err = engine.Export(context.Background(), "meta")
require.ErrorContains(t, err, "export meta interrupted")
}
func TestEngineExportReturnsThrownGetterError(t *testing.T) {
t.Parallel()
engine, err := Compile(`
export const meta = {
apiVersion: 1,
get name() { throw new Error("getter failed"); },
};
`, Options{Key: "meta-throw", Version: "1.0.0"})
require.NoError(t, err)
_, err = engine.Export(context.Background(), "meta")
require.ErrorContains(t, err, "export meta failed")
assert.Contains(t, err.Error(), "getter failed")
}
func TestEngineNestedHooksRequireOwnProperties(t *testing.T) {
t.Parallel()
engine, err := Compile(`
const inheritedRenderers = {
inherited: function(value) { return value; },
constructor: function(value) { return value; },
toString: function(value) { return value; },
["__proto__"]: function(value) { return value; },
};
export const renderers = Object.create(inheritedRenderers);
renderers.own = function(value) { return {id: value.id}; };
const inheritedProtocol = {
renderFinal: function(value) { return value; },
};
export const protocols = {
responses: Object.create(inheritedProtocol),
};
`, Options{Key: "own-hooks", Version: "1.0.0"})
require.NoError(t, err)
found, err := engine.HasCallablePath(context.Background(), "renderers", "own")
require.NoError(t, err)
assert.True(t, found)
result, err := engine.CallMember(context.Background(), "renderers", "own", map[string]any{"id": "task-1"})
require.NoError(t, err)
assert.Equal(t, map[string]any{"id": "task-1"}, result)
for _, member := range []string{"inherited", "constructor", "toString", "__proto__"} {
t.Run(member, func(t *testing.T) {
found, err := engine.HasCallablePath(context.Background(), "renderers", member)
require.NoError(t, err)
assert.False(t, found)
_, err = engine.CallMember(context.Background(), "renderers", member)
require.ErrorContains(t, err, "not found")
})
}
found, err = engine.HasCallablePath(context.Background(), "protocols", "responses", "renderFinal")
require.NoError(t, err)
assert.False(t, found)
_, err = engine.CallPath(context.Background(), "protocols", []string{"responses", "renderFinal"})
require.ErrorContains(t, err, "not found")
}
func TestCompileInterruptsLongRunningInitialization(t *testing.T) {
t.Parallel()
_, err := Compile(`while (true) {}; export function run() {}`, Options{
Key: "loop", Version: "1", Timeout: 20 * time.Millisecond,
})
require.Error(t, err)
assert.Contains(t, err.Error(), "initialization timed out")
}
func TestValidateRequestURL(t *testing.T) {
t.Parallel()
tests := []struct {
name string
requestURL string
baseURL string
allowedHosts []string
wantError string
}{
{name: "same host", requestURL: "https://api.example.com/v1/task", baseURL: "https://api.example.com/v1"},
{name: "default port", requestURL: "https://api.example.com:443/v1/task", baseURL: "https://api.example.com"},
{name: "approved host", requestURL: "https://upload.example.com/task", baseURL: "https://api.example.com", allowedHosts: []string{"upload.example.com"}},
{name: "subdomain is not implicit", requestURL: "https://evil.api.example.com/task", baseURL: "https://api.example.com", wantError: "not allowed"},
{name: "userinfo trick", requestURL: "https://api.example.com@evil.example/task", baseURL: "https://api.example.com", wantError: "not allowed"},
{name: "relative URL", requestURL: "/v1/task", baseURL: "https://api.example.com", wantError: "absolute"},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
err := ValidateRequestURL(test.requestURL, test.baseURL, test.allowedHosts)
if test.wantError == "" {
require.NoError(t, err)
return
}
require.Error(t, err)
assert.True(t, strings.Contains(err.Error(), test.wantError), err.Error())
})
}
}
+116
View File
@@ -0,0 +1,116 @@
package jsplugin
import (
"context"
"encoding/json"
"fmt"
"reflect"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
)
// Fixture describes deterministic calls into a plugin. Arguments and expected
// values stay as JSON so fixtures remain portable across engine implementations.
type Fixture struct {
UnixNow *int64 `json:"unixNow"`
Cases []FixtureCase `json:"cases"`
}
type FixtureCase struct {
Name string `json:"name"`
Hook string `json:"hook"`
Member string `json:"member,omitempty"`
Path []string `json:"path,omitempty"`
Args []json.RawMessage `json:"args"`
Expected json.RawMessage `json:"expected"`
ExpectedError string `json:"expectedError,omitempty"`
}
type FixtureReport struct {
Total int
Passed int
}
// ReplayFixture compiles a plugin and runs every fixture case in declaration
// order. unixNow is fixed by the fixture to keep signing and timestamp hooks
// reproducible.
func ReplayFixture(ctx context.Context, source string, data []byte) (FixtureReport, error) {
var fixture Fixture
if err := common.Unmarshal(data, &fixture); err != nil {
return FixtureReport{}, fmt.Errorf("decode fixture: %w", err)
}
if len(fixture.Cases) == 0 {
return FixtureReport{}, fmt.Errorf("fixture must contain at least one case")
}
options := Options{Key: "fixture", Version: "fixture"}
if fixture.UnixNow != nil {
fixed := time.Unix(*fixture.UnixNow, 0)
options.Now = func() time.Time { return fixed }
}
plugin, err := NewRegistry().Register(source, options)
if err != nil {
return FixtureReport{}, fmt.Errorf("compile plugin: %w", err)
}
report := FixtureReport{Total: len(fixture.Cases)}
for index, testCase := range fixture.Cases {
caseName := strings.TrimSpace(testCase.Name)
if caseName == "" {
caseName = fmt.Sprintf("case %d", index+1)
}
if strings.TrimSpace(testCase.Hook) == "" {
return report, fmt.Errorf("%s: hook is required", caseName)
}
args := make([]any, len(testCase.Args))
for argumentIndex, raw := range testCase.Args {
if err = common.Unmarshal(raw, &args[argumentIndex]); err != nil {
return report, fmt.Errorf("%s: decode argument %d: %w", caseName, argumentIndex+1, err)
}
}
var result any
if testCase.Member != "" && len(testCase.Path) > 0 {
return report, fmt.Errorf("%s: member and path are mutually exclusive", caseName)
}
if len(testCase.Path) > 0 {
result, err = plugin.Engine.CallPath(ctx, testCase.Hook, testCase.Path, args...)
} else if testCase.Member == "" {
result, err = plugin.Engine.Call(ctx, testCase.Hook, args...)
} else {
result, err = plugin.Engine.CallMember(ctx, testCase.Hook, testCase.Member, args...)
}
if testCase.ExpectedError != "" {
if err == nil || !strings.Contains(err.Error(), testCase.ExpectedError) {
return report, fmt.Errorf("%s: expected error containing %q, got %v", caseName, testCase.ExpectedError, err)
}
report.Passed++
continue
}
if err != nil {
return report, fmt.Errorf("%s: %w", caseName, err)
}
if len(testCase.Expected) == 0 {
return report, fmt.Errorf("%s: expected is required when expectedError is empty", caseName)
}
var expected any
if err = common.Unmarshal(testCase.Expected, &expected); err != nil {
return report, fmt.Errorf("%s: decode expected value: %w", caseName, err)
}
actualData, marshalErr := common.Marshal(result)
if marshalErr != nil {
return report, fmt.Errorf("%s: encode actual value: %w", caseName, marshalErr)
}
var actual any
if err = common.Unmarshal(actualData, &actual); err != nil {
return report, fmt.Errorf("%s: normalize actual value: %w", caseName, err)
}
if !reflect.DeepEqual(expected, actual) {
return report, fmt.Errorf("%s: result mismatch: expected %s, got %s", caseName, testCase.Expected, actualData)
}
report.Passed++
}
return report, nil
}
+49
View File
@@ -0,0 +1,49 @@
package jsplugin
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
const fixturePlugin = `
export const meta = { apiVersion: 1, key: "fixture", name: "Fixture", version: "1.0.0", author: {name: "Test"}, channelTypes: [1002], models: ["fixture-model"], fetchMode: "per_task", protocols: [{name: "openai_responses", supports: ["sync", "background"]}] };
export function buildSubmitRequest(ctx) { return {url: ctx.baseUrl + "/submit"}; }
export function parseSubmitResponse(ctx, resp) { return {taskId: resp.body.id}; }
export function buildQueryRequest(ctx) { return {url: ctx.baseUrl + "/task"}; }
export function parseTaskResult(ctx, body) { return body; }
export function stamp(value) { return {value: value, now: utils.unixNow()}; }
export function fail() { throw new Error("fixture failure"); }
export const native = { compact: function(value) { return {id: value.task_id}; } };
export const protocols = { openai_responses: {
decodeRequest: function(value) { return {kind: "submit", model: value.model}; },
renderFinal: function(ctx, task) { return task; }
} };
`
func TestReplayFixture(t *testing.T) {
t.Parallel()
report, err := ReplayFixture(context.Background(), fixturePlugin, []byte(`{
"unixNow": 1700000000,
"cases": [
{"name":"deterministic time","hook":"stamp","args":["ok"],"expected":{"value":"ok","now":1700000000}},
{"name":"member call","hook":"native","member":"compact","args":[{"task_id":"task-1"}],"expected":{"id":"task-1"}},
{"name":"nested path call","hook":"protocols","path":["openai_responses","decodeRequest"],"args":[{"model":"video-1"}],"expected":{"kind":"submit","model":"video-1"}},
{"name":"expected failure","hook":"fail","args":[],"expectedError":"fixture failure"}
]
}`))
require.NoError(t, err)
assert.Equal(t, FixtureReport{Total: 4, Passed: 4}, report)
}
func TestReplayFixtureReportsMismatch(t *testing.T) {
t.Parallel()
report, err := ReplayFixture(context.Background(), fixturePlugin, []byte(`{
"cases": [{"name":"wrong output","hook":"stamp","args":["ok"],"expected":{"value":"different"}}]
}`))
require.Error(t, err)
assert.Contains(t, err.Error(), "wrong output: result mismatch")
assert.Equal(t, FixtureReport{Total: 1}, report)
}
Loaded 100 of 336 files, more files were not shown because too many files have changed in this diff. Show more