mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-10-28 08:35:33 +08:00
feat(task): replace built-in task adaptors with a sandboxed JS plugin system (#7076)
This commit is contained in:
1 parent
7037ac15bd
commit
eb48396d5f
336 files changed
+52333
-6369
No files matched your search
@@ -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.
|
||||
@@ -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 |
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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 != "" {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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"
|
||||
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
@@ -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))
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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
@@ -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()
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -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
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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:")
|
||||
}
|
||||
@@ -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"))
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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"
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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.
|
||||
Vendored
+48
@@ -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;
|
||||
@@ -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.
|
||||
@@ -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"]}}]}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
@@ -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
@@ -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 {
|
||||
|
||||
@@ -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"`
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
@@ -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
@@ -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, "普通用户不支持指定渠道")
|
||||
|
||||
@@ -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
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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",
|
||||
},
|
||||
})
|
||||
}
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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) {
|
||||
|
||||
@@ -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")
|
||||
})
|
||||
}
|
||||
@@ -323,6 +323,7 @@ func migrateDB() error {
|
||||
&TopUp{},
|
||||
&QuotaData{},
|
||||
&Task{},
|
||||
&TaskPlugin{},
|
||||
&Model{},
|
||||
&Vendor{},
|
||||
&PrefillGroup{},
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -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
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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,
|
||||
}}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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() {
|
||||
|
||||
@@ -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",`))
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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.
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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(); }
|
||||
`
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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())
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
Reference in new issue
Block a user