diff --git a/CLAUDE.md b/CLAUDE.md index ff3c01f0c766..97ba749380e3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -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`. \ No newline at end of file +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. diff --git a/THIRD-PARTY-LICENSES.md b/THIRD-PARTY-LICENSES.md index e04a8cd7028b..e67e61395f15 100644 --- a/THIRD-PARTY-LICENSES.md +++ b/THIRD-PARTY-LICENSES.md @@ -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 | diff --git a/common/api_type.go b/common/api_type.go index 82b088cc8fc9..b4ca7062b614 100644 --- a/common/api_type.go +++ b/common/api_type.go @@ -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 diff --git a/common/api_type_task_plugin_test.go b/common/api_type_task_plugin_test.go new file mode 100644 index 000000000000..3a13e083bec3 --- /dev/null +++ b/common/api_type_task_plugin_test.go @@ -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) +} diff --git a/common/init.go b/common/init.go index 4d4c62b27cac..323fd207dddd 100644 --- a/common/init.go +++ b/common/init.go @@ -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 != "" { diff --git a/common/trusted_proxies.go b/common/trusted_proxies.go new file mode 100644 index 000000000000..363dc019fc28 --- /dev/null +++ b/common/trusted_proxies.go @@ -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 +} diff --git a/constant/channel.go b/constant/channel.go index 2a6c4a31c138..354dd0e14b2b 100644 --- a/constant/channel.go +++ b/constant/channel.go @@ -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 { diff --git a/constant/channel_test.go b/constant/channel_test.go new file mode 100644 index 000000000000..92faa26d075d --- /dev/null +++ b/constant/channel_test.go @@ -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)) +} diff --git a/constant/context_key.go b/constant/context_key.go index ccb8010f9476..93a18ba9af01 100644 --- a/constant/context_key.go +++ b/constant/context_key.go @@ -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" diff --git a/constant/env.go b/constant/env.go index 512bfc31126b..a6de36bce60c 100644 --- a/constant/env.go +++ b/constant/env.go @@ -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 diff --git a/constant/task.go b/constant/task.go index ecccf4dfe119..aee856831156 100644 --- a/constant/task.go +++ b/constant/task.go @@ -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 } diff --git a/constant/task_test.go b/constant/task_test.go new file mode 100644 index 000000000000..9b0065df69f1 --- /dev/null +++ b/constant/task_test.go @@ -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)) + }) + } +} diff --git a/controller/billing_option_test.go b/controller/billing_option_test.go new file mode 100644 index 000000000000..c02334d52ab0 --- /dev/null +++ b/controller/billing_option_test.go @@ -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") +} diff --git a/controller/channel-billing.go b/controller/channel-billing.go index 5974628d01ef..36cd887b805f 100644 --- a/controller/channel-billing.go +++ b/controller/channel-billing.go @@ -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, diff --git a/controller/channel-test.go b/controller/channel-test.go index b294979d5877..00268af26559 100644 --- a/controller/channel-test.go +++ b/controller/channel-test.go @@ -83,6 +83,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) diff --git a/controller/channel.go b/controller/channel.go index 3a1e58328923..19ddca8e6a07 100644 --- a/controller/channel.go +++ b/controller/channel.go @@ -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() } diff --git a/controller/channel_pin_retry_test.go b/controller/channel_pin_retry_test.go new file mode 100644 index 000000000000..f79b7b8b1123 --- /dev/null +++ b/controller/channel_pin_retry_test.go @@ -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 +} diff --git a/controller/channel_task_plugin_bind_test.go b/controller/channel_task_plugin_bind_test.go new file mode 100644 index 000000000000..b8f326661594 --- /dev/null +++ b/controller/channel_task_plugin_bind_test.go @@ -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`) +} diff --git a/controller/channel_task_plugin_validation_test.go b/controller/channel_task_plugin_validation_test.go new file mode 100644 index 000000000000..6dfe065040f5 --- /dev/null +++ b/controller/channel_task_plugin_validation_test.go @@ -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") +} diff --git a/controller/channel_upstream_update.go b/controller/channel_upstream_update.go index e1918c257c19..6817a0086ac6 100644 --- a/controller/channel_upstream_update.go +++ b/controller/channel_upstream_update.go @@ -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() } diff --git a/controller/log.go b/controller/log.go index 470c759fc1a1..18ebb10b2bec 100644 --- a/controller/log.go +++ b/controller/log.go @@ -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) diff --git a/controller/model.go b/controller/model.go index 1d759301bc7e..779739477fe1 100644 --- a/controller/model.go +++ b/controller/model.go @@ -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, }) } diff --git a/controller/option.go b/controller/option.go index 940bb3069023..70a5f1921894 100644 --- a/controller/option.go +++ b/controller/option.go @@ -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 { diff --git a/controller/plugin_endpoint_test.go b/controller/plugin_endpoint_test.go new file mode 100644 index 000000000000..5a0aa69f7341 --- /dev/null +++ b/controller/plugin_endpoint_test.go @@ -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) +} diff --git a/controller/plugin_native_e2e_test.go b/controller/plugin_native_e2e_test.go new file mode 100644 index 000000000000..f0eb4eef0416 --- /dev/null +++ b/controller/plugin_native_e2e_test.go @@ -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) +} diff --git a/controller/plugin_protocol.go b/controller/plugin_protocol.go new file mode 100644 index 000000000000..7b9248aeffdb --- /dev/null +++ b/controller/plugin_protocol.go @@ -0,0 +1,1276 @@ +package controller + +import ( + "context" + "errors" + "fmt" + "hash/fnv" + "net/http" + "strconv" + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin" + "github.com/QuantumNous/new-api/relay" + taskjsplugin "github.com/QuantumNous/new-api/relay/channel/task/jsplugin" + relaycommon "github.com/QuantumNous/new-api/relay/common" + relayconstant "github.com/QuantumNous/new-api/relay/constant" + "github.com/QuantumNous/new-api/relay/helper" + "github.com/QuantumNous/new-api/relaykit/types" + "github.com/QuantumNous/new-api/service" + "github.com/gin-gonic/gin" +) + +type pluginProtocolBridgeDeps struct { + submit func(*gin.Context, *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) + loadTask func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) + now func() time.Time + admissions *pluginProtocolObservationLimiter + protocolLimits relay.PluginProtocolLimits + artifactContentURL func(taskID, artifactKey string) (string, error) + submissionTimeout time.Duration + observationTimeout time.Duration + loadTimeout time.Duration + tickInterval time.Duration + tickJitter time.Duration + heartbeatInterval time.Duration + admissionTimeout time.Duration + getByTaskId func(int, string) (*model.Task, bool, error) + resolvePlugin func(constant.TaskPlatform) (*pluginruntime.LoadedPlugin, *pluginruntime.RoutingGeneration, bool) +} + +func defaultPluginProtocolBridgeDeps() pluginProtocolBridgeDeps { + timeout := time.Duration(constant.TaskPluginProtocolTimeoutSeconds) * time.Second + if timeout <= 0 { + timeout = 10 * time.Minute + } + tick := time.Duration(constant.TaskPluginProtocolTickMilliseconds) * time.Millisecond + if tick <= 0 { + tick = 2 * time.Second + } + jitter := max(time.Duration(constant.TaskPluginProtocolTickJitterMilliseconds)*time.Millisecond, 0) + heartbeat := time.Duration(constant.TaskPluginProtocolHeartbeatSeconds) * time.Second + if heartbeat <= 0 { + heartbeat = 15 * time.Second + } + loadTimeout := 5 * time.Second + if halfHeartbeat := heartbeat / 2; halfHeartbeat > 0 && halfHeartbeat < loadTimeout { + loadTimeout = halfHeartbeat + } + return pluginProtocolBridgeDeps{ + submit: executeTaskSubmission, + loadTask: model.GetTaskForProtocolObservation, + now: time.Now, + admissions: pluginProtocolObservationAdmissions, + protocolLimits: relay.DefaultPluginProtocolLimits(), + artifactContentURL: service.BuildTaskArtifactContentURL, + submissionTimeout: timeout, + observationTimeout: timeout, + loadTimeout: loadTimeout, + tickInterval: tick, + tickJitter: jitter, + heartbeatInterval: heartbeat, + admissionTimeout: pluginruntime.DefaultCallTimeout, + getByTaskId: model.GetByTaskId, + resolvePlugin: resolveTaskPluginForProtocolRetrieve, + } +} + +func (d pluginProtocolBridgeDeps) withDefaults() pluginProtocolBridgeDeps { + defaults := defaultPluginProtocolBridgeDeps() + if d.submit == nil { + d.submit = defaults.submit + } + if d.loadTask == nil { + d.loadTask = defaults.loadTask + } + if d.now == nil { + d.now = defaults.now + } + if d.admissions == nil { + d.admissions = defaults.admissions + } + if d.artifactContentURL == nil { + d.artifactContentURL = defaults.artifactContentURL + } + if d.submissionTimeout <= 0 { + d.submissionTimeout = defaults.submissionTimeout + } + if d.observationTimeout <= 0 { + d.observationTimeout = defaults.observationTimeout + } + if d.loadTimeout <= 0 { + d.loadTimeout = defaults.loadTimeout + } + if d.tickInterval <= 0 { + d.tickInterval = defaults.tickInterval + } + if d.tickJitter < 0 { + d.tickJitter = 0 + } + if d.heartbeatInterval <= 0 { + d.heartbeatInterval = defaults.heartbeatInterval + } + if halfHeartbeat := d.heartbeatInterval / 2; halfHeartbeat > 0 && d.loadTimeout > halfHeartbeat { + d.loadTimeout = halfHeartbeat + } + if d.admissionTimeout <= 0 { + d.admissionTimeout = defaults.admissionTimeout + } + if d.getByTaskId == nil { + d.getByTaskId = defaults.getByTaskId + } + if d.resolvePlugin == nil { + d.resolvePlugin = defaults.resolvePlugin + } + return d +} + +func resolveTaskPluginForProtocolRetrieve(platform constant.TaskPlatform) (*pluginruntime.LoadedPlugin, *pluginruntime.RoutingGeneration, bool) { + generation := pluginruntime.DefaultRegistry.Generation() + plugin, ok := relay.ResolveTaskPluginForPlatform(generation, platform) + return plugin, generation, ok +} + +func serveTaskPluginProtocol( + c *gin.Context, + pinned pluginruntime.PinnedEndpoint, + deps pluginProtocolBridgeDeps, +) { + deps = deps.withDefaults() + generation := uint64(0) + if pinned.Generation != nil { + generation = pinned.Generation.Number + } + pluginKey := "" + if pinned.Plugin != nil { + pluginKey = pinned.Plugin.Meta.Key + } + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=bridge_start generation=%d plugin=%q protocol=%q model=%q", + generation, + pluginKey, + pinned.Protocol, + c.GetString("resolved_task_model"), + ) + if !pluginruntime.SupportsHostProtocol(pinned.Protocol) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=bridge_rejected generation=%d plugin=%q reason=unsupported_protocol", generation, pluginKey) + respondPluginProtocolError(c, http.StatusNotImplemented, "task_protocol_not_available", "Task protocol bridge is not available") + return + } + requestValue, exists := c.Get(pluginruntime.ContextKeyProtocolRequest) + protocolRequest, ok := requestValue.(pluginruntime.ProtocolRequestContext) + if !exists || !ok || protocolRequest.Protocol != pinned.Protocol { + logger.LogDebug(c, "task_plugin subsystem=protocol event=bridge_rejected generation=%d plugin=%q reason=invalid_protocol_context", generation, pluginKey) + respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed") + return + } + if definition, known := pluginruntime.HostProtocol(pinned.Protocol); known && len(definition.DefinedModes()) > 0 && pinned.Plugin != nil { + background := false + if body, ok := protocolRequest.Body.(map[string]any); ok && body["kind"] == string(pluginruntime.BodyJSON) { + if requestBody, ok := body["value"].(map[string]any); ok { + background, _ = requestBody["background"].(bool) + } + } + missing := false + if protocolRequest.Stream && !pinned.Plugin.Meta.ProtocolSupports(pinned.Protocol, "stream") { + missing = true + } + if background && !pinned.Plugin.Meta.ProtocolSupports(pinned.Protocol, "background") { + missing = true + } + if !protocolRequest.Stream && !background && !pinned.Plugin.Meta.ProtocolSupports(pinned.Protocol, "sync") { + missing = true + } + if missing { + logger.LogError(c, "pinned task plugin does not support the requested protocol form") + respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed") + return + } + } + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=request_ready generation=%d plugin=%q protocol=%q stream=%t", + generation, + pluginKey, + protocolRequest.Protocol, + protocolRequest.Stream, + ) + + release, admissionErr := deps.admissions.acquire( + pinned.Plugin.Meta.Key, + common.GetContextKeyInt(c, constant.ContextKeyUserId), + common.GetContextKeyInt(c, constant.ContextKeyTokenId), + ) + if admissionErr != nil { + if errors.Is(admissionErr, errPluginProtocolObservationLimitExceeded) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=admission_rejected generation=%d plugin=%q reason=observation_limit", generation, pluginKey) + respondPluginProtocolError(c, http.StatusTooManyRequests, "rate_limit_exceeded", "Too many active task observations") + return + } + logger.LogDebug(c, "task_plugin subsystem=protocol event=admission_rejected generation=%d plugin=%q reason=invalid_identity", generation, pluginKey) + respondPluginProtocolError(c, http.StatusUnauthorized, "authentication_error", "Authentication failed") + return + } + defer release() + logger.LogDebug(c, "task_plugin subsystem=protocol event=admission_acquired generation=%d plugin=%q", generation, pluginKey) + + clientRequest := c.Request + var relayInfo *relaycommon.RelayInfo + var outcome *taskSubmissionOutcome + var taskErr *dto.TaskError + var relayInfoErr error + submissionStage := "relay_info" + // A Responses client only observes an asynchronous task. Once admitted, + // disconnecting that observer must not cancel submission, persistence, or + // billing settlement; the submission keeps its own bounded lifetime. + func() { + submissionContext, cancelSubmission := context.WithTimeout( + context.WithoutCancel(clientRequest.Context()), + deps.submissionTimeout, + ) + c.Request = clientRequest.Clone(submissionContext) + defer func() { + c.Request = clientRequest + cancelSubmission() + }() + + relayInfo, relayInfoErr = relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil) + if relayInfoErr != nil { + return + } + relayInfo.RelayMode = relayconstant.RelayModeVideoSubmit + relayInfo.IsStream = false + relayInfo.OriginModelName = c.GetString("resolved_task_model") + if action := c.GetString("task_action"); action != "" { + relayInfo.Action = action + } + submissionStage = "origin_task" + if taskErr = relay.ResolveOriginTask(c, relayInfo); taskErr != nil { + return + } + if taskErr = relay.ApplyOriginTaskAffinity(c, relayInfo); taskErr != nil { + return + } + + submissionStage = "submission" + logger.LogDebug(c, "task_plugin subsystem=protocol event=submission_start generation=%d plugin=%q protocol=%q stream=%t", generation, pluginKey, protocolRequest.Protocol, protocolRequest.Stream) + outcome, taskErr = deps.submit(c, relayInfo) + }() + + if clientRequest.Context().Err() != nil { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_disconnected generation=%d plugin=%q stage=%s", generation, pluginKey, submissionStage) + return + } + if relayInfoErr != nil { + err := relayInfoErr + logger.LogError(c, "build task protocol relay info failed: "+err.Error()) + logger.LogDebug(c, "task_plugin subsystem=protocol event=bridge_failed generation=%d plugin=%q stage=relay_info reason=invalid_context", generation, pluginKey) + respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed") + return + } + if submissionStage == "origin_task" && taskErr != nil { + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=bridge_failed generation=%d plugin=%q stage=origin_task code=%q status=%d", + generation, + pluginKey, + taskErr.Code, + taskErr.StatusCode, + ) + respondPluginProtocolSubmissionError(c, taskErr) + return + } + + if taskErr != nil { + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=submission_failed generation=%d plugin=%q code=%q status=%d local=%t", + generation, + pluginKey, + taskErr.Code, + taskErr.StatusCode, + taskErr.LocalError, + ) + respondPluginProtocolSubmissionError(c, taskErr) + return + } + if outcome == nil || outcome.Task == nil || outcome.RelayInfo == nil || + outcome.Task.UserId != relayInfo.UserId || + outcome.Task.Platform != constant.TaskPlatform(pinned.Plugin.Meta.Key) { + logger.LogError(c, "task protocol submission returned an invalid durable outcome") + logger.LogDebug(c, "task_plugin subsystem=protocol event=submission_failed generation=%d plugin=%q reason=invalid_durable_outcome", generation, pluginKey) + respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed") + return + } + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=submission_durable generation=%d plugin=%q public_task_id=%q status=%q stream=%t", + generation, + pluginKey, + outcome.Task.TaskID, + taskPluginDebugStatus(string(outcome.Task.Status)), + protocolRequest.Stream, + ) + + createdAt := outcome.Task.CreatedAt + if createdAt == 0 { + createdAt = outcome.Task.SubmitTime + } + if createdAt == 0 { + createdAt = deps.now().Unix() + } + machine := relay.NewPluginResponsesMachine( + outcome.Task.TaskID, + outcome.RelayInfo.OriginModelName, + createdAt, + deps.protocolLimits, + ) + background := false + if body, ok := protocolRequest.Body.(map[string]any); ok && body["kind"] == string(pluginruntime.BodyJSON) { + if requestBody, ok := body["value"].(map[string]any); ok { + background, _ = requestBody["background"].(bool) + } + } + if background { + outcome.Task.PrivateData.ResponsesBackground = true + if outcome.Task.ID != 0 { + if err := model.DB.Model(outcome.Task).Update("private_data", outcome.Task.PrivateData).Error; err != nil { + logger.LogError(c, "persist task background flag failed: "+err.Error()) + } + } + machine.SetBackground(true) + if !protocolRequest.Stream { + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=background_return generation=%d plugin=%q public_task_id=%q status=%q", + generation, + pluginKey, + outcome.Task.TaskID, + taskPluginDebugStatus(string(outcome.Task.Status)), + ) + c.JSON(http.StatusOK, machine.PendingResponse(string(outcome.Task.Status))) + return + } + logger.LogDebug(c, "task_plugin subsystem=protocol event=background_stream generation=%d plugin=%q public_task_id=%q", generation, pluginKey, outcome.Task.TaskID) + } + if protocolRequest.Stream { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_enter generation=%d plugin=%q mode=stream public_task_id=%q", generation, pluginKey, outcome.Task.TaskID) + streamTaskPluginProtocol(c, pinned, protocolRequest, outcome.Task.TaskID, machine, deps) + return + } + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_enter generation=%d plugin=%q mode=nonstream public_task_id=%q", generation, pluginKey, outcome.Task.TaskID) + waitTaskPluginProtocol(c, pinned, protocolRequest, outcome.Task.TaskID, machine, deps) +} + +func streamTaskPluginProtocol( + c *gin.Context, + pinned pluginruntime.PinnedEndpoint, + protocolRequest pluginruntime.ProtocolRequestContext, + taskID string, + machine *relay.PluginResponsesMachine, + deps pluginProtocolBridgeDeps, +) { + generation := pinned.Generation.Number + pluginKey := pinned.Plugin.Meta.Key + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=observation_start generation=%d plugin=%q mode=stream public_task_id=%q timeout_ms=%d tick_ms=%d heartbeat_ms=%d", + generation, + pluginKey, + taskID, + deps.observationTimeout.Milliseconds(), + deps.tickInterval.Milliseconds(), + deps.heartbeatInterval.Milliseconds(), + ) + created, err := machine.CreatedEvent() + if err != nil { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_failed generation=%d plugin=%q mode=stream stage=created_event reason=state_machine_error", generation, pluginKey) + respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed") + return + } + helper.SetEventStreamHeaders(c) + if err = writeTaskPluginProtocolEvent(c, created); err != nil { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_write_failed generation=%d plugin=%q mode=stream stage=created_event", generation, pluginKey) + return + } + + observationContext, cancelObservation := context.WithTimeout(c.Request.Context(), deps.observationTimeout) + defer cancelObservation() + heartbeatTicker := time.NewTicker(deps.heartbeatInterval) + defer heartbeatTicker.Stop() + + var previous relay.ProtocolState + tickNumber := uint64(0) + lastStatus := "" + for { + loadStarted := deps.now() + loadContext, cancelLoad := context.WithTimeout(observationContext, deps.loadTimeout) + task, exists, loadErr := deps.loadTask( + loadContext, + common.GetContextKeyInt(c, constant.ContextKeyUserId), + constant.TaskPlatform(pinned.Plugin.Meta.Key), + taskID, + ) + loadContextErr := loadContext.Err() + cancelLoad() + loadElapsed := deps.now().Sub(loadStarted) + if errors.Is(loadContextErr, context.DeadlineExceeded) && + observationContext.Err() == nil && + c.Request.Context().Err() == nil { + logger.LogWarn(c, fmt.Sprintf( + "task protocol database observation overloaded; plugin=%s task=%s", + pinned.Plugin.Meta.Key, + taskID, + )) + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=observation_tick generation=%d plugin=%q mode=stream tick=%d load_ms=%d overloaded=true", + generation, + pluginKey, + tickNumber, + loadElapsed.Milliseconds(), + ) + delay := pluginProtocolTickDelay(taskID, tickNumber, deps.tickInterval, deps.tickJitter) + deps.tickInterval + tickNumber++ + if !waitForTaskPluginProtocolTick(c, observationContext, heartbeatTicker, delay) { + if errors.Is(observationContext.Err(), context.DeadlineExceeded) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_timeout generation=%d plugin=%q mode=stream last_status=%q", generation, pluginKey, taskPluginDebugStatus(lastStatus)) + writeTaskPluginProtocolTimeout(c, machine, lastStatus) + } else if c.Request.Context().Err() != nil { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_disconnected generation=%d plugin=%q mode=stream stage=backoff_wait", generation, pluginKey) + } else { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_write_failed generation=%d plugin=%q mode=stream stage=heartbeat", generation, pluginKey) + } + return + } + continue + } + if loadErr != nil || !exists || task == nil { + if errors.Is(observationContext.Err(), context.DeadlineExceeded) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_timeout generation=%d plugin=%q mode=stream last_status=%q", generation, pluginKey, taskPluginDebugStatus(lastStatus)) + writeTaskPluginProtocolTimeout(c, machine, lastStatus) + return + } + if loadErr != nil && !errors.Is(loadErr, context.Canceled) { + logger.LogError(c, "task protocol database observation failed") + } + if c.Request.Context().Err() == nil { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_failed generation=%d plugin=%q mode=stream stage=load reason=task_unavailable", generation, pluginKey) + writeTaskPluginProtocolFailure(c, machine, lastStatus) + } else { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_disconnected generation=%d plugin=%q mode=stream stage=load", generation, pluginKey) + } + return + } + previousStatus := lastStatus + lastStatus = string(task.Status) + if lastStatus != previousStatus { + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=status_transition generation=%d plugin=%q mode=stream tick=%d previous=%q status=%q load_ms=%d", + generation, + pluginKey, + tickNumber, + taskPluginDebugStatus(previousStatus), + taskPluginDebugStatus(lastStatus), + loadElapsed.Milliseconds(), + ) + } + view, viewErr := service.BuildTaskPluginView(task) + if viewErr != nil { + logger.LogError(c, "build task protocol view failed: "+viewErr.Error()) + writeTaskPluginProtocolFailure(c, machine, lastStatus) + return + } + viewValue, viewErr := taskPluginProtocolJSONValue(view) + if viewErr != nil { + logger.LogError(c, "encode task protocol view failed: "+viewErr.Error()) + writeTaskPluginProtocolFailure(c, machine, lastStatus) + return + } + hookStarted := deps.now() + rendererContext, contextErr := taskPluginProtocolRendererContext(protocolRequest, pinned, task, deps.artifactContentURL) + if contextErr != nil { + logger.LogError(c, "build task protocol renderer context failed") + writeTaskPluginProtocolFailure(c, machine, lastStatus) + return + } + args := []any{rendererContext, viewValue} + if previous.Present { + previousValue, stateErr := previous.PluginValue() + if stateErr != nil { + logger.LogError(c, "decode task protocol state failed: "+stateErr.Error()) + writeTaskPluginProtocolFailure(c, machine, lastStatus) + return + } + args = append(args, previousValue) + } + value, callErr := pinned.Plugin.Engine.CallPathWithAdmissionTimeout(observationContext, deps.admissionTimeout, "protocols", []string{pinned.Protocol, "renderEvents"}, args...) + hookElapsed := deps.now().Sub(hookStarted) + overloaded := false + if callErr != nil { + if errors.Is(observationContext.Err(), context.DeadlineExceeded) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_timeout generation=%d plugin=%q mode=stream stage=render_events last_status=%q", generation, pluginKey, taskPluginDebugStatus(lastStatus)) + writeTaskPluginProtocolTimeout(c, machine, lastStatus) + return + } + if c.Request.Context().Err() != nil { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_disconnected generation=%d plugin=%q mode=stream stage=render_events", generation, pluginKey) + return + } + if errors.Is(callErr, pluginruntime.ErrCallAdmissionTimeout) { + overloaded = true + logger.LogWarn(c, fmt.Sprintf( + "task protocol render hook overloaded; plugin=%s task=%s", + pinned.Plugin.Meta.Key, + taskID, + )) + } else { + logger.LogError(c, "task protocol render hook failed") + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=observation_failed generation=%d plugin=%q mode=stream stage=render_events reason=hook_failed elapsed_ms=%d", + generation, + pluginKey, + hookElapsed.Milliseconds(), + ) + writeTaskPluginProtocolFailure(c, machine, lastStatus) + return + } + } + if !overloaded { + result, decodeErr := relay.DecodePluginProtocolEventResult(value, deps.protocolLimits) + if decodeErr != nil { + logger.LogError(c, "task protocol render result invalid: "+decodeErr.Error()) + writeTaskPluginProtocolFailure(c, machine, lastStatus) + return + } + events, applyErr := machine.ApplyTick(result, lastStatus) + if applyErr != nil { + logger.LogError(c, "task protocol state transition failed: "+applyErr.Error()) + writeTaskPluginProtocolFailure(c, machine, lastStatus) + return + } + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=render_events generation=%d plugin=%q mode=stream tick=%d status=%q semantic_events=%d wire_events=%d done=%t state_present=%t elapsed_ms=%d", + generation, + pluginKey, + tickNumber, + taskPluginDebugStatus(lastStatus), + len(result.Events), + len(events), + result.Done, + result.State.Present, + hookElapsed.Milliseconds(), + ) + for _, event := range events { + if err = writeTaskPluginProtocolEvent(c, event); err != nil { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_write_failed generation=%d plugin=%q mode=stream stage=event event_type=%q sequence=%d", generation, pluginKey, event.Type, event.SequenceNumber) + return + } + } + if taskPluginProtocolEventsTerminal(events) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_complete generation=%d plugin=%q mode=stream reason=terminal status=%q ticks=%d", generation, pluginKey, taskPluginDebugStatus(lastStatus), tickNumber+1) + return + } + previous = result.State + } + + delay := pluginProtocolTickDelay(taskID, tickNumber, deps.tickInterval, deps.tickJitter) + tickNumber++ + if overloaded { + delay += deps.tickInterval + } else if hookElapsed > deps.tickInterval { + delay += deps.tickInterval + logger.LogWarn(c, fmt.Sprintf( + "task protocol render hook slow; plugin=%s task=%s elapsed_ms=%d", + pinned.Plugin.Meta.Key, + taskID, + hookElapsed.Milliseconds(), + )) + } + if !waitForTaskPluginProtocolTick(c, observationContext, heartbeatTicker, delay) { + if errors.Is(observationContext.Err(), context.DeadlineExceeded) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_timeout generation=%d plugin=%q mode=stream last_status=%q", generation, pluginKey, taskPluginDebugStatus(lastStatus)) + writeTaskPluginProtocolTimeout(c, machine, lastStatus) + } else if c.Request.Context().Err() != nil { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_disconnected generation=%d plugin=%q mode=stream stage=tick_wait", generation, pluginKey) + } else { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_write_failed generation=%d plugin=%q mode=stream stage=heartbeat", generation, pluginKey) + } + return + } + } +} + +func waitForTaskPluginProtocolTick( + c *gin.Context, + observationContext context.Context, + heartbeatTicker *time.Ticker, + delay time.Duration, +) bool { + tickTimer := time.NewTimer(delay) + defer tickTimer.Stop() + for { + select { + case <-c.Request.Context().Done(): + return false + case <-observationContext.Done(): + return false + case <-heartbeatTicker.C: + helper.ExtendWriteDeadline(c) + if err := writeTaskPluginProtocolHeartbeat(c); err != nil { + return false + } + case <-tickTimer.C: + return true + } + } +} + +func waitTaskPluginProtocol( + c *gin.Context, + pinned pluginruntime.PinnedEndpoint, + protocolRequest pluginruntime.ProtocolRequestContext, + taskID string, + machine *relay.PluginResponsesMachine, + deps pluginProtocolBridgeDeps, +) { + generation := pinned.Generation.Number + pluginKey := pinned.Plugin.Meta.Key + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=observation_start generation=%d plugin=%q mode=nonstream public_task_id=%q timeout_ms=%d tick_ms=%d", + generation, + pluginKey, + taskID, + deps.observationTimeout.Milliseconds(), + deps.tickInterval.Milliseconds(), + ) + observationContext, cancelObservation := context.WithTimeout(c.Request.Context(), deps.observationTimeout) + defer cancelObservation() + tickNumber := uint64(0) + lastStatus := "" + for { + loadStarted := deps.now() + loadContext, cancelLoad := context.WithTimeout(observationContext, deps.loadTimeout) + task, exists, err := deps.loadTask( + loadContext, + common.GetContextKeyInt(c, constant.ContextKeyUserId), + constant.TaskPlatform(pinned.Plugin.Meta.Key), + taskID, + ) + loadContextErr := loadContext.Err() + cancelLoad() + loadElapsed := deps.now().Sub(loadStarted) + loadOverloaded := errors.Is(loadContextErr, context.DeadlineExceeded) && + observationContext.Err() == nil && + c.Request.Context().Err() == nil + if loadOverloaded { + logger.LogWarn(c, fmt.Sprintf( + "task protocol database observation overloaded; plugin=%s task=%s", + pinned.Plugin.Meta.Key, + taskID, + )) + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=observation_tick generation=%d plugin=%q mode=nonstream tick=%d load_ms=%d overloaded=true", + generation, + pluginKey, + tickNumber, + loadElapsed.Milliseconds(), + ) + } else if err != nil || !exists || task == nil { + if errors.Is(observationContext.Err(), context.DeadlineExceeded) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_timeout generation=%d plugin=%q mode=nonstream last_status=%q", generation, pluginKey, taskPluginDebugStatus(lastStatus)) + writeTaskPluginProtocolTimeoutResponse(c, machine, lastStatus) + return + } + if err != nil && !errors.Is(err, context.Canceled) { + logger.LogError(c, "task protocol database observation failed") + } + if c.Request.Context().Err() == nil { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_failed generation=%d plugin=%q mode=nonstream stage=load reason=task_unavailable", generation, pluginKey) + writeTaskPluginProtocolFailureResponse(c, machine, lastStatus) + } else { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_disconnected generation=%d plugin=%q mode=nonstream stage=load", generation, pluginKey) + } + return + } + overloaded := loadOverloaded + if !loadOverloaded { + previousStatus := lastStatus + lastStatus = string(task.Status) + if lastStatus != previousStatus { + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=status_transition generation=%d plugin=%q mode=nonstream tick=%d previous=%q status=%q load_ms=%d", + generation, + pluginKey, + tickNumber, + taskPluginDebugStatus(previousStatus), + taskPluginDebugStatus(lastStatus), + loadElapsed.Milliseconds(), + ) + } + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=observation_tick generation=%d plugin=%q mode=nonstream tick=%d status=%q load_ms=%d overloaded=false", + generation, + pluginKey, + tickNumber, + taskPluginDebugStatus(lastStatus), + loadElapsed.Milliseconds(), + ) + } + if !loadOverloaded && (task.Status == model.TaskStatusSuccess || task.Status == model.TaskStatusFailure) { + if task.Status == model.TaskStatusFailure { + writeTaskPluginProtocolFailureResponse(c, machine, string(task.Status)) + return + } + response, hookElapsed, callErr := renderTaskPluginProtocolFinalResponse( + observationContext, + pinned, + protocolRequest, + task, + machine, + deps, + ) + if callErr != nil { + if errors.Is(observationContext.Err(), context.DeadlineExceeded) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_timeout generation=%d plugin=%q mode=nonstream stage=render_final last_status=%q", generation, pluginKey, taskPluginDebugStatus(lastStatus)) + writeTaskPluginProtocolTimeoutResponse(c, machine, lastStatus) + return + } + if c.Request.Context().Err() != nil { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_disconnected generation=%d plugin=%q mode=nonstream stage=render_final", generation, pluginKey) + return + } + if errors.Is(callErr, pluginruntime.ErrCallAdmissionTimeout) { + overloaded = true + logger.LogWarn(c, fmt.Sprintf( + "task protocol final hook overloaded; plugin=%s task=%s", + pinned.Plugin.Meta.Key, + taskID, + )) + } else { + logger.LogError(c, "task protocol final hook failed") + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=observation_failed generation=%d plugin=%q mode=nonstream stage=render_final reason=hook_failed elapsed_ms=%d", + generation, + pluginKey, + hookElapsed.Milliseconds(), + ) + writeTaskPluginProtocolFailureResponse(c, machine, lastStatus) + return + } + } else { + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=render_final generation=%d plugin=%q mode=nonstream status=%q elapsed_ms=%d", + generation, + pluginKey, + taskPluginDebugStatus(lastStatus), + hookElapsed.Milliseconds(), + ) + c.JSON(http.StatusOK, response) + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=observation_complete generation=%d plugin=%q mode=nonstream reason=terminal status=%q ticks=%d", + generation, + pluginKey, + taskPluginDebugStatus(lastStatus), + tickNumber+1, + ) + return + } + } + + delay := pluginProtocolTickDelay(taskID, tickNumber, deps.tickInterval, deps.tickJitter) + tickNumber++ + if overloaded { + delay += deps.tickInterval + } + tickTimer := time.NewTimer(delay) + select { + case <-c.Request.Context().Done(): + if !tickTimer.Stop() { + select { + case <-tickTimer.C: + default: + } + } + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_disconnected generation=%d plugin=%q mode=nonstream stage=tick_wait", generation, pluginKey) + return + case <-observationContext.Done(): + if !tickTimer.Stop() { + select { + case <-tickTimer.C: + default: + } + } + if errors.Is(observationContext.Err(), context.DeadlineExceeded) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_timeout generation=%d plugin=%q mode=nonstream last_status=%q", generation, pluginKey, taskPluginDebugStatus(lastStatus)) + writeTaskPluginProtocolTimeoutResponse(c, machine, lastStatus) + } + return + case <-tickTimer.C: + } + } +} + +func renderTaskPluginProtocolFinalResponse( + ctx context.Context, + pinned pluginruntime.PinnedEndpoint, + protocolRequest pluginruntime.ProtocolRequestContext, + task *model.Task, + machine *relay.PluginResponsesMachine, + deps pluginProtocolBridgeDeps, +) (map[string]any, time.Duration, error) { + view, err := service.BuildTaskPluginView(task) + if err != nil { + return nil, 0, err + } + viewValue, err := taskPluginProtocolJSONValue(view) + if err != nil { + return nil, 0, err + } + rendererContext, err := taskPluginProtocolRendererContext( + protocolRequest, + pinned, + task, + deps.artifactContentURL, + ) + if err != nil { + return nil, 0, err + } + hookStarted := deps.now() + payload, err := pinned.Plugin.Engine.CallPathWithAdmissionTimeout( + ctx, + deps.admissionTimeout, + "protocols", + []string{pinned.Protocol, "renderFinal"}, + rendererContext, + viewValue, + ) + hookElapsed := deps.now().Sub(hookStarted) + if err != nil { + return nil, hookElapsed, err + } + response, err := machine.FinalResponse(payload, string(task.Status)) + if err != nil { + return nil, hookElapsed, err + } + return response, hookElapsed, nil +} + +func renderTaskPluginProtocolEventsResponse( + ctx context.Context, + pinned pluginruntime.PinnedEndpoint, + protocolRequest pluginruntime.ProtocolRequestContext, + task *model.Task, + machine *relay.PluginResponsesMachine, + deps pluginProtocolBridgeDeps, +) (map[string]any, time.Duration, error) { + view, err := service.BuildTaskPluginView(task) + if err != nil { + return nil, 0, err + } + viewValue, err := taskPluginProtocolJSONValue(view) + if err != nil { + return nil, 0, err + } + rendererContext, err := taskPluginProtocolRendererContext( + protocolRequest, + pinned, + task, + deps.artifactContentURL, + ) + if err != nil { + return nil, 0, err + } + hookStarted := deps.now() + value, err := pinned.Plugin.Engine.CallPathWithAdmissionTimeout( + ctx, + deps.admissionTimeout, + "protocols", + []string{pinned.Protocol, "renderEvents"}, + rendererContext, + viewValue, + ) + hookElapsed := deps.now().Sub(hookStarted) + if err != nil { + return nil, hookElapsed, err + } + result, err := relay.DecodePluginProtocolEventResult(value, deps.protocolLimits) + if err != nil { + return nil, hookElapsed, err + } + response, err := machine.FinalFromEvents(result, string(task.Status)) + if err != nil { + return nil, hookElapsed, err + } + return response, hookElapsed, nil +} + +func RetrieveTaskPluginResponse(c *gin.Context) { + retrieveTaskPluginResponse(c, defaultPluginProtocolBridgeDeps()) +} + +func retrieveTaskPluginResponse(c *gin.Context, deps pluginProtocolBridgeDeps) { + deps = deps.withDefaults() + responseID := strings.TrimSpace(c.Param("response_id")) + if !strings.HasPrefix(responseID, "resp_") { + writeTaskPluginResponseNotFound(c, responseID, "bad_prefix") + return + } + taskID := "task_" + strings.TrimPrefix(responseID, "resp_") + userID := common.GetContextKeyInt(c, constant.ContextKeyUserId) + logger.LogDebug(c, "task_plugin subsystem=protocol event=retrieve_start response_id=%q public_task_id=%q", responseID, taskID) + + task, exists, err := deps.getByTaskId(userID, taskID) + if err != nil { + logger.LogError(c, "task protocol retrieve lookup failed") + logger.LogDebug(c, "task_plugin subsystem=protocol event=retrieve_failed reason=lookup_error public_task_id=%q", taskID) + respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed") + return + } + if !exists || task == nil { + writeTaskPluginResponseNotFound(c, responseID, "missing") + return + } + + plugin, generation, ok := deps.resolvePlugin(task.Platform) + if !ok || plugin == nil { + writeTaskPluginResponseNotFound(c, responseID, "no_plugin") + return + } + claimsProtocol := false + for _, claim := range plugin.Meta.Protocols { + if claim.Name == "openai_responses" { + claimsProtocol = true + break + } + } + if !claimsProtocol { + writeTaskPluginResponseNotFound(c, responseID, "no_claim") + return + } + + generationNumber := uint64(0) + if generation != nil { + generationNumber = generation.Number + } + createdAt := task.CreatedAt + if createdAt == 0 { + createdAt = task.SubmitTime + } + if createdAt == 0 { + createdAt = deps.now().Unix() + } + machine := relay.NewPluginResponsesMachine( + task.TaskID, + task.Properties.OriginModelName, + createdAt, + deps.protocolLimits, + ) + machine.SetBackground(task.PrivateData.ResponsesBackground) + pinned := pluginruntime.PinnedEndpoint{ + Generation: generation, + Plugin: plugin, + Protocol: "openai_responses", + Model: task.Properties.OriginModelName, + } + protocolRequest := pluginruntime.ProtocolRequestContext{ + RouteRequestContext: pluginruntime.RouteRequestContext{ + Path: c.Request.URL.Path, + Method: http.MethodGet, + Params: map[string]string{"response_id": responseID}, + Query: c.Request.URL.Query(), + Body: map[string]any{"kind": string(pluginruntime.BodyNone)}, + }, + Protocol: "openai_responses", + Operation: "retrieve", + Model: task.Properties.OriginModelName, + } + + if task.Status == model.TaskStatusFailure { + logger.LogDebug(c, "task_plugin subsystem=protocol event=retrieve_final generation=%d plugin=%q public_task_id=%q status=%q", generationNumber, plugin.Meta.Key, task.TaskID, taskPluginDebugStatus(string(task.Status))) + writeTaskPluginProtocolFailureResponse(c, machine, string(task.Status)) + return + } + if task.Status != model.TaskStatusSuccess { + logger.LogDebug(c, "task_plugin subsystem=protocol event=retrieve_pending generation=%d plugin=%q public_task_id=%q status=%q", generationNumber, plugin.Meta.Key, task.TaskID, taskPluginDebugStatus(string(task.Status))) + c.JSON(http.StatusOK, machine.PendingResponse(string(task.Status))) + return + } + + var ( + response map[string]any + hookElapsed time.Duration + renderErr error + ) + if plugin.Meta.ProtocolSupports("openai_responses", "sync") || plugin.Meta.ProtocolSupports("openai_responses", "background") { + response, hookElapsed, renderErr = renderTaskPluginProtocolFinalResponse( + c.Request.Context(), + pinned, + protocolRequest, + task, + machine, + deps, + ) + } else { + response, hookElapsed, renderErr = renderTaskPluginProtocolEventsResponse( + c.Request.Context(), + pinned, + protocolRequest, + task, + machine, + deps, + ) + } + if renderErr != nil { + logger.LogError(c, "task protocol retrieve render failed") + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=retrieve_failed generation=%d plugin=%q public_task_id=%q stage=render_final elapsed_ms=%d", + generationNumber, + plugin.Meta.Key, + task.TaskID, + hookElapsed.Milliseconds(), + ) + writeTaskPluginProtocolFailureResponse(c, machine, string(task.Status)) + return + } + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=retrieve_final generation=%d plugin=%q public_task_id=%q status=%q elapsed_ms=%d", + generationNumber, + plugin.Meta.Key, + task.TaskID, + taskPluginDebugStatus(string(task.Status)), + hookElapsed.Milliseconds(), + ) + c.JSON(http.StatusOK, response) +} + +func writeTaskPluginResponseNotFound(c *gin.Context, responseID, reason string) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=retrieve_not_found reason=%s response_id=%q", reason, responseID) + respondPluginProtocolError(c, http.StatusNotFound, "not_found", "No response found with id '"+responseID+"'.") +} + +func writeTaskPluginProtocolHeartbeat(c *gin.Context) error { + if _, err := c.Writer.Write([]byte(": PING\n")); err != nil { + return err + } + return helper.FlushWriter(c) +} + +func writeTaskPluginProtocolFailure( + c *gin.Context, + machine *relay.PluginResponsesMachine, + taskStatus string, +) { + failed, err := machine.FailureEvent(taskStatus) + if err != nil { + logger.LogError(c, "task protocol failure event failed: "+err.Error()) + return + } + _ = writeTaskPluginProtocolEvent(c, failed) +} + +func writeTaskPluginProtocolTimeout( + c *gin.Context, + machine *relay.PluginResponsesMachine, + taskStatus string, +) { + incomplete, err := machine.TimeoutEvent(taskStatus) + if err != nil { + logger.LogError(c, "task protocol timeout event failed: "+err.Error()) + return + } + _ = writeTaskPluginProtocolEvent(c, incomplete) +} + +func writeTaskPluginProtocolFailureResponse( + c *gin.Context, + machine *relay.PluginResponsesMachine, + taskStatus string, +) { + if taskStatus == string(model.TaskStatusFailure) { + response, err := machine.FinalResponse(nil, taskStatus) + if err != nil { + logger.LogError(c, "task protocol terminal failure response failed: "+err.Error()) + respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed") + return + } + c.JSON(http.StatusOK, response) + return + } + response, err := machine.FailureResponse(taskStatus) + if err != nil { + logger.LogError(c, "task protocol failure response failed: "+err.Error()) + respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed") + return + } + c.JSON(http.StatusOK, response) +} + +func writeTaskPluginProtocolTimeoutResponse( + c *gin.Context, + machine *relay.PluginResponsesMachine, + lastStatus string, +) { + response, err := machine.TimeoutResponse(lastStatus) + if err != nil { + logger.LogError(c, "task protocol timeout response failed: "+err.Error()) + respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed") + return + } + c.JSON(http.StatusOK, response) +} + +func taskPluginProtocolJSONValue(value any) (any, error) { + encoded, err := common.Marshal(value) + if err != nil { + return nil, err + } + var decoded any + if err = common.Unmarshal(encoded, &decoded); err != nil { + return nil, err + } + return decoded, nil +} + +func taskPluginProtocolRendererContext( + request pluginruntime.ProtocolRequestContext, + pinned pluginruntime.PinnedEndpoint, + task *model.Task, + artifactContentURL func(taskID, artifactKey string) (string, error), +) (map[string]any, error) { + rendererContext := request.JSValue() + if task == nil || task.Status != model.TaskStatusSuccess { + return rendererContext, nil + } + if pinned.Plugin == nil { + return nil, errors.New("task artifact projection is unavailable") + } + + artifacts, err := taskjsplugin.New(pinned.Plugin).ListArtifacts(task) + if err != nil { + return nil, fmt.Errorf("project task artifacts: %w", err) + } + artifacts, err = validateProjectedTaskArtifacts(artifacts) + if err != nil { + return nil, err + } + if len(artifacts) > 0 && artifactContentURL == nil { + return nil, errors.New("task artifact projection is unavailable") + } + + rendererArtifacts := make(map[string]any, len(artifacts)) + for _, artifact := range artifacts { + contentURL, buildErr := artifactContentURL(task.TaskID, artifact.Key) + if buildErr != nil { + return nil, fmt.Errorf("build task artifact content URL: %w", buildErr) + } + item := map[string]any{ + "key": artifact.Key, + "type": artifact.Type, + "url": contentURL, + } + if artifact.MimeType != "" { + item["mimeType"] = artifact.MimeType + } + rendererArtifacts[artifact.Key] = item + } + rendererContext["artifacts"] = rendererArtifacts + return rendererContext, nil +} + +func pluginProtocolTickDelay(taskID string, tick uint64, base, jitter time.Duration) time.Duration { + if jitter <= 0 { + return base + } + hash := fnv.New64a() + _, _ = hash.Write([]byte(taskID)) + _, _ = hash.Write([]byte(":")) + _, _ = hash.Write([]byte(strconv.FormatUint(tick, 10))) + return base + time.Duration(hash.Sum64()%uint64(jitter+1)) +} + +func taskPluginProtocolEventsTerminal(events []dto.PluginResponsesStreamEvent) bool { + for _, event := range events { + switch event.Type { + case "response.completed", "response.failed", "response.incomplete": + return true + } + } + return false +} + +func writeTaskPluginProtocolEvent(c *gin.Context, event dto.PluginResponsesStreamEvent) error { + encoded, err := common.Marshal(event) + if err != nil { + return err + } + helper.ExtendWriteDeadline(c) + if _, err = c.Writer.Write([]byte("event: " + event.Type + "\n")); err != nil { + return err + } + if _, err = c.Writer.Write([]byte("data: " + string(encoded) + "\n\n")); err != nil { + return err + } + if err = helper.FlushWriter(c); err != nil { + return err + } + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=sse_event_sent event_type=%q sequence=%d", + event.Type, + event.SequenceNumber, + ) + return nil +} + +func respondPluginProtocolSubmissionError(c *gin.Context, taskErr *dto.TaskError) { + status := http.StatusInternalServerError + if taskErr != nil && taskErr.StatusCode >= 400 && taskErr.StatusCode <= 599 { + status = taskErr.StatusCode + } + switch status { + case http.StatusBadRequest: + message := "Invalid task protocol request" + if taskErr != nil && taskErr.Message != "" && (taskErr.Code == "invalid_request" || strings.HasPrefix(taskErr.Code, "invalid_request")) { + message = taskErr.Message + } + respondPluginProtocolError(c, status, "invalid_request_error", message) + case http.StatusUnauthorized: + respondPluginProtocolError(c, status, "authentication_error", "Authentication failed") + case http.StatusForbidden: + respondPluginProtocolError(c, status, "permission_denied", "Task protocol request was denied") + case http.StatusTooManyRequests: + respondPluginProtocolError(c, status, "rate_limit_exceeded", "Too many requests") + default: + respondPluginProtocolError(c, status, "task_protocol_error", "Task protocol request failed") + } +} + +func respondPluginProtocolError(c *gin.Context, status int, code, message string) { + c.JSON(status, gin.H{ + "error": gin.H{ + "message": message, + "type": "new_api_error", + "code": code, + }, + }) +} diff --git a/controller/plugin_protocol_limiter.go b/controller/plugin_protocol_limiter.go new file mode 100644 index 000000000000..d88d177ee8de --- /dev/null +++ b/controller/plugin_protocol_limiter.go @@ -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 +} diff --git a/controller/plugin_protocol_limiter_test.go b/controller/plugin_protocol_limiter_test.go new file mode 100644 index 000000000000..a6b48cf65d43 --- /dev/null +++ b/controller/plugin_protocol_limiter_test.go @@ -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) +} diff --git a/controller/plugin_protocol_sdk_test.go b/controller/plugin_protocol_sdk_test.go new file mode 100644 index 000000000000..971794c9b249 --- /dev/null +++ b/controller/plugin_protocol_sdk_test.go @@ -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) + })) +} diff --git a/controller/plugin_protocol_test.go b/controller/plugin_protocol_test.go new file mode 100644 index 000000000000..a6a5fdc1bdbe --- /dev/null +++ b/controller/plugin_protocol_test.go @@ -0,0 +1,1555 @@ +package controller + +import ( + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "strconv" + "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/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func TestServeTaskPluginProtocolWaitsForDurableSubmissionBeforeWriting(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "durable-barrier", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [], done: false}; }, + renderFinal: function() { return {}; } + }}; + `) + c, recorder := newPluginProtocolTestContext(true, false) + submitStarted := make(chan struct{}) + releaseSubmit := make(chan struct{}) + done := make(chan struct{}) + + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + close(submitStarted) + <-releaseSubmit + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_durable", map[string]any{ + "must_not": "be_written", + }), nil + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + return nil, false, errors.New("observation failed after durable barrier") + } + + go func() { + defer close(done) + serveTaskPluginProtocol(c, pinned, deps) + }() + select { + case <-submitStarted: + case <-time.After(2 * time.Second): + require.FailNow(t, "submission did not start") + } + + assert.Empty(t, recorder.Header().Get("Content-Type")) + assert.Empty(t, recorder.Body.String()) + assert.False(t, recorder.Flushed) + + close(releaseSubmit) + select { + case <-done: + case <-time.After(2 * time.Second): + require.FailNow(t, "protocol handler did not finish") + } + assert.Equal(t, "text/event-stream", recorder.Header().Get("Content-Type")) + assert.Contains(t, recorder.Body.String(), "event: response.created\n") + assert.NotContains(t, recorder.Body.String(), "must_not") +} + +func TestServeTaskPluginProtocolDisconnectDuringSubmissionFinishesDurableWithoutWriting(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "disconnect-during-submit", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [], done: false}; }, + renderFinal: function() { return {}; } + }}; + `) + c, recorder := newPluginProtocolTestContext(true, true) + requestContext, cancel := context.WithCancel(c.Request.Context()) + c.Request = c.Request.WithContext(requestContext) + submitStarted := make(chan struct{}) + checkSubmissionContext := make(chan struct{}) + submissionContextActive := make(chan struct{}) + releaseSubmit := make(chan struct{}) + observationStarted := make(chan struct{}, 1) + done := make(chan struct{}) + deps := pluginProtocolTestDeps() + deps.submit = func(c *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + close(submitStarted) + <-checkSubmissionContext + select { + case <-c.Request.Context().Done(): + return nil, service.TaskErrorWrapperLocal(c.Request.Context().Err(), "request_cancelled", http.StatusRequestTimeout) + default: + close(submissionContextActive) + } + <-releaseSubmit + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_disconnect_durable", nil), nil + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + observationStarted <- struct{}{} + return nil, false, errors.New("observation must not start after disconnect") + } + + go func() { + defer close(done) + serveTaskPluginProtocol(c, pinned, deps) + }() + select { + case <-submitStarted: + case <-time.After(2 * time.Second): + require.FailNow(t, "submission did not start") + } + cancel() + close(checkSubmissionContext) + select { + case <-submissionContextActive: + case <-time.After(2 * time.Second): + require.FailNow(t, "submission context was canceled with the client") + } + select { + case <-done: + require.FailNow(t, "protocol handler stopped before submission became durable") + default: + } + close(releaseSubmit) + select { + case <-done: + case <-time.After(2 * time.Second): + require.FailNow(t, "protocol handler did not finish after durable submission") + } + + assert.Empty(t, recorder.Header().Get("Content-Type")) + assert.Empty(t, recorder.Body.String()) + assert.False(t, recorder.Flushed) + select { + case <-observationStarted: + require.FailNow(t, "protocol observation started after client disconnect") + default: + } +} + +func TestServeTaskPluginProtocolDisconnectBeforeDurableBarrierPersistsAndSettlesWithoutRefund(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 }) + + pinned := compilePluginProtocolTestEndpoint(t, "disconnect-before-durable", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [], done: false}; }, + renderFinal: function() { return {}; } + }}; + `) + c, recorder := newPluginProtocolTestContext(true, true) + requestContext, cancel := context.WithCancel(c.Request.Context()) + c.Request = c.Request.WithContext(requestContext) + billing := &taskSubmissionTestBilling{events: &events} + submitStarted := make(chan struct{}) + releaseSubmit := make(chan struct{}) + observationStarted := make(chan struct{}, 1) + done := make(chan struct{}) + + deps := pluginProtocolTestDeps() + deps.submit = func(c *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + info.Billing = billing + info.TaskRelayInfo.PublicTaskID = "task_disconnect_persisted" + info.TaskRelayInfo.LockedChannel = &model.Channel{ + Id: 1, + Type: constant.ChannelTypeTaskPlugin, + Name: "disconnect-before-durable", + } + info.ChannelMeta = &relaycommon.ChannelMeta{ + ChannelId: 1, + ChannelType: constant.ChannelTypeTaskPlugin, + } + return executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) { + close(submitStarted) + <-releaseSubmit + return &relay.TaskSubmitResult{ + UpstreamTaskID: "upstream_disconnect_persisted", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + Quota: 7, + }, nil + }) + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + observationStarted <- struct{}{} + return nil, false, errors.New("observation must not start after disconnect") + } + + go func() { + defer close(done) + serveTaskPluginProtocol(c, pinned, deps) + }() + select { + case <-submitStarted: + case <-time.After(2 * time.Second): + require.FailNow(t, "submission did not start") + } + cancel() + close(releaseSubmit) + select { + case <-done: + case <-time.After(2 * time.Second): + require.FailNow(t, "detached submission did not finish") + } + + assert.Equal(t, []string{"reserve", "insert", "settle"}, events) + assert.Zero(t, billing.refunds) + var persisted model.Task + require.NoError(t, database.Where("task_id = ?", "task_disconnect_persisted").First(&persisted).Error) + assert.Equal(t, model.TaskStatus(model.TaskStatusNotStart), persisted.Status) + assert.Equal(t, 7, persisted.Quota) + assert.Equal(t, "upstream_disconnect_persisted", persisted.PrivateData.UpstreamTaskID) + assert.Empty(t, recorder.Header().Get("Content-Type")) + assert.Empty(t, recorder.Body.String()) + assert.False(t, recorder.Flushed) + select { + case <-observationStarted: + require.FailNow(t, "protocol observation started after client disconnect") + default: + } +} + +func TestServeTaskPluginProtocolDisconnectDuringTerminalSettlementStopsOnlyObservation(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "disconnect-terminal-settlement", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [], done: false}; }, + renderFinal: function() { return {}; } + }}; + `) + + previousDB := model.DB + previousMemoryCache := common.MemoryCacheEnabled + database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, database.AutoMigrate(&model.Channel{}, &model.Task{})) + model.DB = database + common.MemoryCacheEnabled = false + t.Cleanup(func() { + model.DB = previousDB + common.MemoryCacheEnabled = previousMemoryCache + }) + baseURL := "https://example.com" + channel := model.Channel{ + Type: constant.ChannelTypeTaskPlugin, + Name: "terminal-settlement", + Key: "test-key", + BaseURL: &baseURL, + Status: common.ChannelStatusEnabled, + } + require.NoError(t, database.Create(&channel).Error) + task := model.Task{ + TaskID: "task_terminal_disconnect", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + ChannelId: channel.Id, + Quota: 10, + Status: model.TaskStatusSubmitted, + PrivateData: model.TaskPrivateData{ + UpstreamTaskID: "upstream-terminal", + }, + } + require.NoError(t, database.Create(&task).Error) + + c, recorder := newPluginProtocolTestContext(true, true) + requestContext, cancel := context.WithCancel(c.Request.Context()) + c.Request = c.Request.WithContext(requestContext) + billingEvents := make([]string, 0) + billing := &taskSubmissionTestBilling{events: &billingEvents} + observationStarted := make(chan struct{}) + settlementStarted := make(chan struct{}) + releaseSettlement := make(chan struct{}) + t.Cleanup(func() { + select { + case <-releaseSettlement: + default: + close(releaseSettlement) + } + }) + pollingDone := make(chan struct{}) + adaptor := &terminalSettlementPollingAdaptor{ + started: settlementStarted, + release: releaseSettlement, + } + previousAdaptorFactory := service.GetTaskAdaptorFunc + service.GetTaskAdaptorFunc = func(constant.TaskPlatform) service.TaskPollingAdaptor { + return adaptor + } + t.Cleanup(func() { service.GetTaskAdaptorFunc = previousAdaptorFactory }) + + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + info.Billing = billing + return &taskSubmissionOutcome{ + Result: &relay.TaskSubmitResult{}, + Task: &task, + RelayInfo: info, + }, nil + } + deps.loadTask = func(ctx context.Context, _ int, _ constant.TaskPlatform, _ string) (*model.Task, bool, error) { + close(observationStarted) + <-ctx.Done() + return nil, false, ctx.Err() + } + done := make(chan struct{}) + + go func() { + <-observationStarted + defer close(pollingDone) + service.DispatchPlatformUpdate( + context.Background(), + task.Platform, + map[int][]string{channel.Id: {"upstream-terminal"}}, + map[string]*model.Task{"upstream-terminal": &task}, + ) + }() + go func() { + defer close(done) + serveTaskPluginProtocol(c, pinned, deps) + }() + select { + case <-settlementStarted: + case <-time.After(2 * time.Second): + require.FailNow(t, "terminal settlement did not start") + } + cancel() + select { + case <-done: + case <-time.After(2 * time.Second): + require.FailNow(t, "protocol observation did not stop after terminal disconnect") + } + assert.Equal(t, []string{"response.created"}, pluginProtocolTestSSEEventTypes(recorder.Body.String())) + assert.Zero(t, billing.refunds) + + close(releaseSettlement) + select { + case <-pollingDone: + case <-time.After(2 * time.Second): + require.FailNow(t, "terminal settlement was canceled with the client observation") + } + + var persisted model.Task + require.NoError(t, database.Where("task_id = ?", task.TaskID).First(&persisted).Error) + assert.Equal(t, model.TaskStatus(model.TaskStatusSuccess), persisted.Status) + assert.Equal(t, "100%", persisted.Progress) + assert.Equal(t, 10, persisted.Quota) + assert.True(t, adaptor.completed) + assert.Empty(t, billingEvents) +} + +type terminalSettlementPollingAdaptor struct { + started chan struct{} + release chan struct{} + completed bool +} + +func (a *terminalSettlementPollingAdaptor) Init(*relaycommon.RelayInfo) {} + +func (a *terminalSettlementPollingAdaptor) FetchTask(string, string, map[string]any, string) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(`{}`)), + }, nil +} + +func (a *terminalSettlementPollingAdaptor) ParseTaskResult([]byte) (*relaycommon.TaskInfo, error) { + return &relaycommon.TaskInfo{ + Status: model.TaskStatusSuccess, + Progress: "100%", + }, nil +} + +func (a *terminalSettlementPollingAdaptor) AdjustBillingOnComplete(task *model.Task, _ *relaycommon.TaskInfo) int { + close(a.started) + <-a.release + a.completed = true + return task.Quota +} + +func TestPluginProtocolBridgeBoundsDatabaseReadBelowHeartbeat(t *testing.T) { + deps := pluginProtocolBridgeDeps{ + observationTimeout: time.Minute, + loadTimeout: 10 * time.Second, + tickInterval: time.Second, + heartbeatInterval: 4 * time.Second, + admissionTimeout: time.Second, + }.withDefaults() + + assert.Equal(t, 2*time.Second, deps.loadTimeout) +} + +func TestServeTaskPluginProtocolPostDurableObservationFailureUsesCanonicalResponse(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "observation-failure", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [], done: false}; }, + renderFinal: function() { return {}; } + }}; + `) + c, recorder := newPluginProtocolTestContext(false, false) + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_observation_failure", nil), nil + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + return nil, false, errors.New("database-secret https://database.invalid/?token=hidden") + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, http.StatusOK, recorder.Code) + var response dto.PluginResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "failed", response.Status) + require.NotNil(t, response.Error) + assert.Equal(t, "server_error", response.Error.Code) + assert.Equal(t, "The task could not be observed.", response.Error.Message) + assert.Equal(t, "queued", response.Metadata["task_status"]) + assert.Equal(t, "/v1/responses/resp_observation_failure", response.Metadata["retrieval_path"]) + assert.NotContains(t, recorder.Body.String(), "secret") + assert.NotContains(t, recorder.Body.String(), "database.invalid") +} + +func TestServeTaskPluginProtocolStreamsPinnedGenerationWithHostFraming(t *testing.T) { + oldPinned := compilePluginProtocolTestEndpoint(t, "generation-pinned", ` + export const protocols = {openai_responses: { + renderEvents: function(ctx, task, previousState) { + if (ctx.stream !== true || ctx.body.value.stream !== true) { + throw new Error("host did not preserve parsed stream mode"); + } + if (arguments.length === 2) { + return {events: [], state: null, done: false}; + } + if (arguments.length !== 3 || previousState !== null) { + throw new Error("explicit null state was not supplied on the next tick"); + } + return {events: [{type: "output", data: "old-generation"}], done: true}; + }, + renderFinal: function() { throw new Error("stream called renderFinal"); } + }}; + `) + newPinned := compilePluginProtocolTestEndpoint(t, "generation-pinned", ` + export const protocols = {openai_responses: { + renderEvents: function() { + return {events: [{type: "output", data: "new-generation"}], done: true}; + }, + renderFinal: function() { return "new-generation"; } + }}; + `) + require.NotSame(t, oldPinned.Plugin.Engine, newPinned.Plugin.Engine) + + c, recorder := newPluginProtocolTestContext(true, true) + c.Set(pluginruntime.ContextKeyRouteRequest, pluginruntime.RouteRequestContext{ + Path: "/v1/responses", + Method: http.MethodPost, + RequestBody: map[string]any{"model": "video-model", "stream": false}, + }) + loadCount := 0 + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, oldPinned.Plugin.Meta.Key, "task_generation", map[string]any{ + "client_response": "ignored", + }), nil + } + deps.loadTask = func(_ context.Context, userID int, platform constant.TaskPlatform, taskID string) (*model.Task, bool, error) { + loadCount++ + assert.Equal(t, 71, userID) + assert.Equal(t, constant.TaskPlatform(oldPinned.Plugin.Meta.Key), platform) + assert.Equal(t, "task_generation", taskID) + status := model.TaskStatus(model.TaskStatusInProgress) + if loadCount == 2 { + status = model.TaskStatus(model.TaskStatusSuccess) + } + return &model.Task{ + TaskID: taskID, + UserId: userID, + Platform: platform, + Status: status, + }, true, nil + } + + serveTaskPluginProtocol(c, oldPinned, deps) + + assert.Equal(t, 2, loadCount) + assert.True(t, recorder.Flushed) + assert.Equal(t, "text/event-stream", recorder.Header().Get("Content-Type")) + 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", + }, pluginProtocolTestSSEEventTypes(recorder.Body.String())) + assert.True(t, strings.HasPrefix(recorder.Body.String(), "event: response.created\ndata: {")) + assert.Contains(t, recorder.Body.String(), `"sequence_number":0`) + assert.Contains(t, recorder.Body.String(), `"sequence_number":7`) + assert.Contains(t, recorder.Body.String(), "old-generation") + assert.NotContains(t, recorder.Body.String(), "new-generation") +} + +func TestServeTaskPluginProtocolStreamMissingRenderEventsUsesFailureEnvelope(t *testing.T) { + tests := []struct { + name string + status model.TaskStatus + }{ + {name: "success", status: model.TaskStatusSuccess}, + {name: "failure", status: model.TaskStatusFailure}, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "default-events-"+testCase.name, ` + export const protocols = {openai_responses: { + renderFinal: function() { throw new Error("stream must not call renderFinal"); } + }}; + `) + c, recorder := newPluginProtocolTestContext(true, true) + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_default_events"), nil + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + return &model.Task{TaskID: "task_default_events", Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), UserId: 71, Status: testCase.status}, true, nil + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, []string{"response.created", "response.failed"}, pluginProtocolTestSSEEventTypes(recorder.Body.String())) + assert.NotContains(t, recorder.Body.String(), "stream must not call") + }) + } +} + +func TestServeTaskPluginProtocolStreamInjectsHostArtifactCapabilities(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "stream-artifacts", ` + export function listArtifacts(task) { + if (task.data.output.video_url !== "https://upstream.invalid/video.mp4?secret=hidden") { + throw new Error("listArtifacts did not receive raw Task.Data"); + } + return [{key: "video", type: "video", mimeType: "video/mp4"}]; + } + export function buildContentRequest() { + throw new Error("rendering must not resolve provider content"); + } + export const protocols = {openai_responses: { + renderEvents: function(ctx, task) { + const artifact = ctx.artifacts && ctx.artifacts.video; + if (!artifact || artifact.key !== "video" || artifact.type !== "video" || + artifact.mimeType !== "video/mp4") { + throw new Error("host artifact context is invalid"); + } + return {events: [{type: "output", data: artifact.url}], done: true}; + }, + renderFinal: function() { throw new Error("stream called renderFinal"); } + }}; + `) + c, recorder := newPluginProtocolTestContext(true, true) + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_stream_artifact", nil), nil + } + task := &model.Task{ + TaskID: "task_stream_artifact", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusSuccess, + } + task.SetData(map[string]any{ + "output": map[string]any{ + "video_url": "https://upstream.invalid/video.mp4?secret=hidden", + }, + }) + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + return task, true, nil + } + deps.artifactContentURL = func(taskID, artifactKey string) (string, error) { + assert.Equal(t, "task_stream_artifact", taskID) + assert.Equal(t, "video", artifactKey) + return "https://gateway.example/v1/tasks/task_stream_artifact/artifacts/video/content?access=host-capability", nil + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, model.TaskStatus(model.TaskStatusSuccess), task.Status) + 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", + }, pluginProtocolTestSSEEventTypes(recorder.Body.String())) + assert.Contains(t, recorder.Body.String(), "host-capability") + assert.NotContains(t, recorder.Body.String(), "upstream.invalid") + assert.NotContains(t, recorder.Body.String(), "secret") +} + +func TestTaskPluginProtocolHeartbeatDoesNotDispatchEmptySDKEvent(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + + require.NoError(t, writeTaskPluginProtocolHeartbeat(c)) + + assert.Equal(t, ": PING\n", recorder.Body.String()) + assert.True(t, recorder.Flushed) +} + +func TestServeTaskPluginProtocolNonStreamUsesFinalHookAndHostEnvelope(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "final-response", ` + export const protocols = {openai_responses: { + renderEvents: function() { throw new Error("non-stream called renderEvents"); }, + renderFinal: function(ctx, task) { + if (ctx.stream !== false || ctx.body.value.stream !== false) { + throw new Error("host did not preserve parsed non-stream mode"); + } + return { + id: "plugin-controlled-id", + status: "plugin-controlled-status", + metadata: {plugin_field: "kept", task_id: "plugin-controlled-task"}, + output: [{ + id: "plugin-controlled-item", + type: "message", + status: "plugin-controlled-item-status", + role: "assistant", + content: [{ + id: "plugin-controlled-content", + type: "output_text", + text: task.data.value, + annotations: [], + logprobs: [] + }] + }], + custom_field: "kept" + }; + } + }}; + `) + c, recorder := newPluginProtocolTestContext(false, false) + c.Set(pluginruntime.ContextKeyRouteRequest, pluginruntime.RouteRequestContext{ + Path: "/v1/responses", + Method: http.MethodPost, + RequestBody: map[string]any{"model": "video-model", "stream": true}, + }) + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_final", map[string]any{ + "client_response_secret": "must-be-ignored", + }), nil + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + task := &model.Task{ + TaskID: "task_final", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusSuccess, + } + task.SetData(map[string]any{"value": "plugin-semantic-result"}) + return task, true, nil + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, http.StatusOK, recorder.Code) + assert.NotEqual(t, "text/event-stream", recorder.Header().Get("Content-Type")) + var response dto.PluginResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "resp_final", response.ID) + assert.Equal(t, "response", response.Object) + assert.Equal(t, "completed", response.Status) + assert.Equal(t, "video-model", response.Model) + assert.Equal(t, "task_final", response.Metadata["task_id"]) + assert.Equal(t, "kept", response.Metadata["plugin_field"]) + require.Len(t, response.Output, 1) + assert.Equal(t, "item_task_final_0", response.Output[0].ID) + assert.Equal(t, "completed", response.Output[0].Status) + require.Len(t, response.Output[0].Content, 1) + assert.Equal(t, "content_task_final_0_0", response.Output[0].Content[0].ID) + assert.Equal(t, "plugin-semantic-result", response.Output[0].Content[0].Text) + assert.Contains(t, recorder.Body.String(), `"custom_field":"kept"`) + assert.NotContains(t, recorder.Body.String(), "plugin-controlled-id") + assert.NotContains(t, recorder.Body.String(), "client_response_secret") +} + +func TestServeTaskPluginProtocolNonStreamInjectsHostArtifactCapabilities(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "final-artifacts", ` + export function listArtifacts() { + return [{key: "video", type: "video"}]; + } + export function buildContentRequest() { + throw new Error("rendering must not resolve provider content"); + } + export const protocols = {openai_responses: { + renderEvents: function() { throw new Error("non-stream called renderEvents"); }, + renderFinal: function(ctx) { + const artifact = ctx.artifacts && ctx.artifacts.video; + if (!artifact || artifact.key !== "video" || artifact.type !== "video" || + Object.prototype.hasOwnProperty.call(artifact, "mimeType")) { + throw new Error("host artifact context is invalid"); + } + return { + output: [{ + type: "message", + status: "completed", + role: "assistant", + content: [{ + type: "output_text", + text: artifact.url, + annotations: [], + logprobs: [] + }] + }] + }; + } + }}; + `) + c, recorder := newPluginProtocolTestContext(false, false) + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_final_artifact", nil), nil + } + task := &model.Task{ + TaskID: "task_final_artifact", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusSuccess, + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + return task, true, nil + } + deps.artifactContentURL = func(taskID, artifactKey string) (string, error) { + assert.Equal(t, "task_final_artifact", taskID) + assert.Equal(t, "video", artifactKey) + return "https://gateway.example/v1/tasks/task_final_artifact/artifacts/video/content?access=host-capability", nil + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, model.TaskStatus(model.TaskStatusSuccess), task.Status) + var response dto.PluginResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "completed", response.Status) + require.Len(t, response.Output, 1) + require.Len(t, response.Output[0].Content, 1) + assert.Contains(t, response.Output[0].Content[0].Text, "host-capability") +} + +func TestServeTaskPluginProtocolArtifactURLFailureOnlyFailsCurrentRendering(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(strconv.FormatBool(stream), func(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "artifact-url-failure-"+strconv.FormatBool(stream), ` + export function listArtifacts() { + return [{key: "video", type: "video"}]; + } + export function buildContentRequest() { + throw new Error("unused"); + } + export const protocols = {openai_responses: { + renderEvents: function() { + return {events: [{type: "output", data: "must-not-render"}], done: true}; + }, + renderFinal: function() { + return {output: []}; + } + }}; + `) + c, recorder := newPluginProtocolTestContext(stream, stream) + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_capability_failure", nil), nil + } + task := &model.Task{ + TaskID: "task_capability_failure", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusSuccess, + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + return task, true, nil + } + deps.artifactContentURL = func(string, string) (string, error) { + return "", errors.New("public address is unavailable") + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, model.TaskStatus(model.TaskStatusSuccess), task.Status) + assert.NotContains(t, recorder.Body.String(), "must-not-render") + assert.NotContains(t, recorder.Body.String(), "public address") + if stream { + assert.Equal(t, []string{"response.created", "response.failed"}, pluginProtocolTestSSEEventTypes(recorder.Body.String())) + } else { + var response dto.PluginResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "failed", response.Status) + require.NotNil(t, response.Error) + assert.Equal(t, "server_error", response.Error.Code) + assert.Equal(t, "completed", response.Metadata["task_status"]) + } + }) + } +} + +func TestServeTaskPluginProtocolNonStreamTaskFailureSkipsFinalHook(t *testing.T) { + logs := make([]string, 0, 1) + pinned := compilePluginProtocolTestEndpointWithOptions(t, "failed-final", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [], done: false}; }, + renderFinal: function() { + console.log("renderFinal called"); + return { + output: [{id: "secret-id", content: [{text: "plugin-secret"}]}], + secret: "https://secret.invalid/" + }; + } + }}; + `, pluginruntime.Options{ + Log: func(message string) { logs = append(logs, message) }, + }) + c, recorder := newPluginProtocolTestContext(false, false) + deps := pluginProtocolTestDeps() + deps.artifactContentURL = func(string, string) (string, error) { + require.FailNow(t, "failed tasks must not project artifact URLs") + return "", nil + } + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_failed", map[string]any{ + "credential": "client-response-secret", + }), nil + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + task := &model.Task{ + TaskID: "task_failed", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusFailure, + FailReason: "upstream credential at https://secret.invalid/", + } + task.SetData(map[string]any{"secret": "database-secret"}) + return task, true, nil + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, http.StatusOK, recorder.Code) + var response dto.PluginResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "failed", response.Status) + require.NotNil(t, response.Error) + assert.Equal(t, "server_error", response.Error.Code) + assert.Equal(t, "The task failed.", response.Error.Message) + assert.Empty(t, response.Output) + assert.NotContains(t, recorder.Body.String(), "secret") + assert.NotContains(t, recorder.Body.String(), "credential") + assert.Empty(t, logs) +} + +func TestServeTaskPluginProtocolStreamTaskFailureSuppressesPluginAndDatabaseDetails(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "failed-stream", ` + export const protocols = {openai_responses: { + renderEvents: function() { + return { + events: [{type: "output", data: "plugin-secret https://plugin.invalid/?key=hidden"}], + done: true + }; + }, + renderFinal: function() { return "unused-secret"; } + }}; + `) + c, recorder := newPluginProtocolTestContext(true, false) + deps := pluginProtocolTestDeps() + deps.artifactContentURL = func(string, string) (string, error) { + require.FailNow(t, "failed tasks must not project artifact URLs") + return "", nil + } + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_stream_failed", nil), nil + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + task := &model.Task{ + TaskID: "task_stream_failed", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusFailure, + FailReason: "database-secret https://database.invalid/?token=hidden", + } + task.SetData(map[string]any{"secret": "private-result"}) + return task, true, nil + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, []string{"response.created", "response.failed"}, pluginProtocolTestSSEEventTypes(recorder.Body.String())) + assert.Contains(t, recorder.Body.String(), `"code":"server_error"`) + assert.Contains(t, recorder.Body.String(), `"message":"The task failed."`) + assert.Contains(t, recorder.Body.String(), `"task_status":"failed"`) + assert.NotContains(t, recorder.Body.String(), "secret") + assert.NotContains(t, recorder.Body.String(), "invalid") + assert.NotContains(t, recorder.Body.String(), "hidden") +} + +func TestServeTaskPluginProtocolRejectsUnsupportedProtocolBeforeSubmission(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "unsupported-protocol", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [], done: false}; }, + renderFinal: function() { return {}; } + }}; + `) + pinned.Protocol = "unsupported" + c, recorder := newPluginProtocolTestContext(false, false) + c.Set(pluginruntime.ContextKeyProtocolRequest, pluginruntime.ProtocolRequestContext{ + RouteRequestContext: pluginruntime.RouteRequestContext{ + Path: "/v1/videos", + Method: http.MethodPost, + RequestBody: map[string]any{"model": "video-model"}, + }, + Protocol: pinned.Protocol, + }) + submitted := false + deps := pluginProtocolTestDeps() + deps.submit = func(*gin.Context, *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + submitted = true + return nil, nil + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.False(t, submitted) + assert.Equal(t, http.StatusNotImplemented, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"code":"task_protocol_not_available"`) +} + +func TestServeTaskPluginProtocolRejectsObservationAdmissionBeforeSubmission(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "admission-limit", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [], done: false}; }, + renderFinal: function() { return {}; } + }}; + `) + for _, stream := range []bool{false, true} { + t.Run(map[bool]string{false: "non-stream", true: "stream"}[stream], func(t *testing.T) { + c, recorder := newPluginProtocolTestContext(stream, stream) + submitted := false + deps := pluginProtocolTestDeps() + deps.admissions = newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{ + global: 0, + perPlugin: 1, + perUser: 1, + perToken: 1, + }) + deps.submit = func(*gin.Context, *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + submitted = true + return nil, nil + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.False(t, submitted) + assert.Equal(t, http.StatusTooManyRequests, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"code":"rate_limit_exceeded"`) + assert.NotEqual(t, "text/event-stream", recorder.Header().Get("Content-Type")) + }) + } +} + +func TestServeTaskPluginProtocolBackgroundNonStreamReturnsPendingWithoutObservation(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "background-create", ` + export const protocols = {openai_responses: { + renderEvents: function() { throw new Error("background create called renderEvents"); }, + renderFinal: function() { throw new Error("background create called renderFinal"); } + }}; + `) + c, recorder := newPluginProtocolTestContext(false, false) + setProtocolRequestBackground(c, true) + loadCalls := 0 + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_background", nil), nil + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + loadCalls++ + return nil, false, errors.New("observation must not start for background create") + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, 0, loadCalls) + assert.Equal(t, http.StatusOK, recorder.Code) + var response map[string]any + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "resp_background", response["id"]) + assert.Equal(t, "response", response["object"]) + assert.Equal(t, "queued", response["status"]) + assert.Equal(t, true, response["background"]) + assert.Nil(t, response["completed_at"]) + assert.Nil(t, response["error"]) + assert.Nil(t, response["usage"]) + assert.Empty(t, response["output"]) + metadata, ok := response["metadata"].(map[string]any) + require.True(t, ok) + assert.Equal(t, "task_background", metadata["task_id"]) + assert.Equal(t, "queued", metadata["task_status"]) + assert.Equal(t, "/v1/responses/resp_background", metadata["retrieval_path"]) + assert.NotEqual(t, "text/event-stream", recorder.Header().Get("Content-Type")) +} + +func TestServeTaskPluginProtocolBackgroundStreamEntersObservation(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "background-stream", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [{type: "output", data: "streamed"}], done: true}; }, + renderFinal: function() { throw new Error("stream called renderFinal"); } + }}; + `) + c, recorder := newPluginProtocolTestContext(true, true) + setProtocolRequestBackground(c, true) + loadCalls := 0 + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_background_stream", nil), nil + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + loadCalls++ + return &model.Task{ + TaskID: "task_background_stream", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusSuccess, + }, true, nil + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Greater(t, loadCalls, 0) + 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", + }, pluginProtocolTestSSEEventTypes(recorder.Body.String())) +} + +func TestRetrieveTaskPluginResponsePendingSkipsRenderFinal(t *testing.T) { + logs := make([]string, 0, 1) + pinned := compilePluginProtocolRetrieveEndpoint(t, "retrieve-pending", ` + export const protocols = {openai_responses: { + renderEvents: function() { throw new Error("retrieve pending called renderEvents"); }, + renderFinal: function() { + console.log("renderFinal called"); + return {}; + } + }}; + `, logsAppender(&logs)) + c, recorder := newPluginProtocolRetrieveContext("resp_retrieve_pending") + deps := pluginProtocolRetrieveDeps(pinned, &model.Task{ + TaskID: "task_retrieve_pending", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusInProgress, + PrivateData: model.TaskPrivateData{ + ResponsesBackground: true, + }, + Properties: model.Properties{OriginModelName: "video-model"}, + CreatedAt: 1_710_000_000, + }, true, nil) + + retrieveTaskPluginResponse(c, deps) + + assert.Empty(t, logs) + assert.Equal(t, http.StatusOK, recorder.Code) + var response map[string]any + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "resp_retrieve_pending", response["id"]) + assert.Equal(t, "in_progress", response["status"]) + assert.Equal(t, true, response["background"]) + assert.Nil(t, response["completed_at"]) + assert.Empty(t, response["output"]) + metadata, ok := response["metadata"].(map[string]any) + require.True(t, ok) + assert.Equal(t, "/v1/responses/resp_retrieve_pending", metadata["retrieval_path"]) +} + +func TestRetrieveTaskPluginResponseSuccessRendersFinal(t *testing.T) { + logs := make([]string, 0, 1) + pinned := compilePluginProtocolRetrieveEndpoint(t, "retrieve-success", ` + export const protocols = {openai_responses: { + renderEvents: function() { throw new Error("retrieve success called renderEvents"); }, + renderFinal: function() { + console.log("renderFinal called"); + return { + output: [{ + type: "message", + status: "completed", + role: "assistant", + content: [{type: "output_text", text: "retrieved-final", annotations: [], logprobs: []}] + }] + }; + } + }}; + `, logsAppender(&logs)) + c, recorder := newPluginProtocolRetrieveContext("resp_retrieve_success") + deps := pluginProtocolRetrieveDeps(pinned, &model.Task{ + TaskID: "task_retrieve_success", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusSuccess, + Properties: model.Properties{OriginModelName: "video-model"}, + CreatedAt: 1_710_000_000, + }, true, nil) + + retrieveTaskPluginResponse(c, deps) + + require.NotEmpty(t, logs) + assert.Contains(t, logs[0], "renderFinal called") + assert.Equal(t, http.StatusOK, recorder.Code) + var response dto.PluginResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "completed", response.Status) + assert.Equal(t, "resp_retrieve_success", response.ID) + require.Len(t, response.Output, 1) + require.Len(t, response.Output[0].Content, 1) + assert.Equal(t, "retrieved-final", response.Output[0].Content[0].Text) +} + +func TestRetrieveTaskPluginResponseStreamOnlySuccessSynthesizesFromEvents(t *testing.T) { + logs := make([]string, 0, 1) + pinned := compilePluginProtocolRetrieveEndpoint(t, "retrieve-stream-only", ` + export const protocols = {openai_responses: { + renderEvents: function() { + console.log("renderEvents called"); + return {events: [{type: "output", data: "synthesized-retrieve"}], done: true}; + }, + renderFinal: function() { throw new Error("stream-only retrieve called renderFinal"); } + }}; + `, logsAppender(&logs)) + pinned.Plugin.Meta.Protocols = []pluginruntime.ProtocolClaim{{Name: "openai_responses", Supports: []string{"stream"}}} + c, recorder := newPluginProtocolRetrieveContext("resp_retrieve_stream") + deps := pluginProtocolRetrieveDeps(pinned, &model.Task{ + TaskID: "task_retrieve_stream", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusSuccess, + Properties: model.Properties{OriginModelName: "video-model"}, + CreatedAt: 1_710_000_000, + }, true, nil) + + retrieveTaskPluginResponse(c, deps) + + require.NotEmpty(t, logs) + assert.Contains(t, logs[0], "renderEvents called") + assert.Equal(t, http.StatusOK, recorder.Code) + var response dto.PluginResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "completed", response.Status) + require.NotEmpty(t, response.Output) +} + +func TestRetrieveTaskPluginResponseStreamOnlyPendingAndFailureStayHostEnvelopes(t *testing.T) { + logs := make([]string, 0, 1) + pinned := compilePluginProtocolRetrieveEndpoint(t, "retrieve-stream-envelope", ` + export const protocols = {openai_responses: { + renderEvents: function() { throw new Error("envelope retrieve called renderEvents"); }, + renderFinal: function() { throw new Error("envelope retrieve called renderFinal"); } + }}; + `, logsAppender(&logs)) + pinned.Plugin.Meta.Protocols = []pluginruntime.ProtocolClaim{{Name: "openai_responses", Supports: []string{"stream"}}} + + t.Run("pending", func(t *testing.T) { + logs = logs[:0] + c, recorder := newPluginProtocolRetrieveContext("resp_retrieve_stream_pending") + deps := pluginProtocolRetrieveDeps(pinned, &model.Task{ + TaskID: "task_retrieve_stream_pending", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusInProgress, + Properties: model.Properties{OriginModelName: "video-model"}, + CreatedAt: 1_710_000_000, + }, true, nil) + retrieveTaskPluginResponse(c, deps) + assert.Empty(t, logs) + assert.Equal(t, http.StatusOK, recorder.Code) + var response map[string]any + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "in_progress", response["status"]) + assert.Empty(t, response["output"]) + }) + + t.Run("failure", func(t *testing.T) { + logs = logs[:0] + c, recorder := newPluginProtocolRetrieveContext("resp_retrieve_stream_failure") + deps := pluginProtocolRetrieveDeps(pinned, &model.Task{ + TaskID: "task_retrieve_stream_failure", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusFailure, + Properties: model.Properties{OriginModelName: "video-model"}, + }, true, nil) + retrieveTaskPluginResponse(c, deps) + assert.Empty(t, logs) + assert.Equal(t, http.StatusOK, recorder.Code) + var response dto.PluginResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "failed", response.Status) + require.NotNil(t, response.Error) + assert.Equal(t, "The task failed.", response.Error.Message) + }) +} + +func TestRetrieveTaskPluginResponseStreamOnlyRenderErrorUsesFailureEnvelope(t *testing.T) { + pinned := compilePluginProtocolRetrieveEndpoint(t, "retrieve-stream-throw", ` + export const protocols = {openai_responses: { + renderEvents: function() { throw new Error("retrieve boom"); }, + renderFinal: function() { throw new Error("stream-only retrieve called renderFinal"); } + }}; + `, pluginruntime.Options{}) + pinned.Plugin.Meta.Protocols = []pluginruntime.ProtocolClaim{{Name: "openai_responses", Supports: []string{"stream"}}} + c, recorder := newPluginProtocolRetrieveContext("resp_retrieve_stream_throw") + deps := pluginProtocolRetrieveDeps(pinned, &model.Task{ + TaskID: "task_retrieve_stream_throw", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusSuccess, + Properties: model.Properties{OriginModelName: "video-model"}, + CreatedAt: 1_710_000_000, + }, true, nil) + + retrieveTaskPluginResponse(c, deps) + + assert.Equal(t, http.StatusOK, recorder.Code) + assert.NotContains(t, recorder.Body.String(), "retrieve boom") + var response dto.PluginResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "failed", response.Status) + require.NotNil(t, response.Error) + assert.Equal(t, "server_error", response.Error.Code) + assert.Equal(t, "The task could not be observed.", response.Error.Message) +} + +func TestRetrieveTaskPluginResponseFailureUsesFailedEnvelope(t *testing.T) { + logs := make([]string, 0, 1) + pinned := compilePluginProtocolRetrieveEndpoint(t, "retrieve-failure", ` + export const protocols = {openai_responses: { + renderEvents: function() { throw new Error("retrieve failure called renderEvents"); }, + renderFinal: function() { + console.log("renderFinal called"); + return {}; + } + }}; + `, logsAppender(&logs)) + c, recorder := newPluginProtocolRetrieveContext("resp_retrieve_failure") + deps := pluginProtocolRetrieveDeps(pinned, &model.Task{ + TaskID: "task_retrieve_failure", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusFailure, + Properties: model.Properties{OriginModelName: "video-model"}, + }, true, nil) + + retrieveTaskPluginResponse(c, deps) + + assert.Empty(t, logs) + assert.Equal(t, http.StatusOK, recorder.Code) + var response dto.PluginResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "failed", response.Status) + require.NotNil(t, response.Error) + assert.Equal(t, "server_error", response.Error.Code) + assert.Equal(t, "The task failed.", response.Error.Message) +} + +func TestRetrieveTaskPluginResponseNotFound(t *testing.T) { + pinned := compilePluginProtocolRetrieveEndpoint(t, "retrieve-404", ` + export const protocols = {openai_responses: { + renderFinal: function() { return {}; } + }}; + `, pluginruntime.Options{}) + owned := &model.Task{ + TaskID: "task_owned", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusInProgress, + Properties: model.Properties{OriginModelName: "video-model"}, + } + + tests := []struct { + name string + responseID string + userID int + task *model.Task + exists bool + plugin *pluginruntime.LoadedPlugin + claims []pluginruntime.ProtocolClaim + }{ + {name: "bad prefix", responseID: "task_owned", userID: 71, task: owned, exists: true, plugin: pinned.Plugin}, + {name: "missing", responseID: "resp_missing", userID: 71, exists: false, plugin: pinned.Plugin}, + {name: "other user", responseID: "resp_owned", userID: 99, task: owned, exists: false, plugin: pinned.Plugin}, + {name: "no plugin", responseID: "resp_owned", userID: 71, task: owned, exists: true}, + {name: "plugin does not claim protocol", responseID: "resp_owned", userID: 71, task: owned, exists: true, plugin: pinned.Plugin, claims: []pluginruntime.ProtocolClaim{{Name: "openai_video"}}}, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + if testCase.plugin != nil { + if testCase.claims != nil { + testCase.plugin.Meta.Protocols = testCase.claims + } else { + testCase.plugin.Meta.Protocols = []pluginruntime.ProtocolClaim{{Name: "openai_responses", Supports: []string{"stream", "sync", "background"}}} + } + } + c, recorder := newPluginProtocolRetrieveContext(testCase.responseID) + common.SetContextKey(c, constant.ContextKeyUserId, testCase.userID) + deps := pluginProtocolRetrieveDeps(pinned, testCase.task, testCase.exists, nil) + if testCase.plugin == nil { + deps.resolvePlugin = func(constant.TaskPlatform) (*pluginruntime.LoadedPlugin, *pluginruntime.RoutingGeneration, bool) { + return nil, nil, false + } + } + + retrieveTaskPluginResponse(c, deps) + + assert.Equal(t, http.StatusNotFound, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"code":"not_found"`) + assert.Contains(t, recorder.Body.String(), "No response found with id '"+testCase.responseID+"'.") + }) + } +} + +func TestRespondPluginProtocolSubmissionErrorPassesValidationMessage(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "protocol-validation-detail", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [], done: false}; }, + renderFinal: function() { return {}; } + }}; + `) + c, recorder := newPluginProtocolTestContext(false, false) + deps := pluginProtocolTestDeps() + deps.submit = func(*gin.Context, *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return nil, &dto.TaskError{ + Code: "invalid_request", + Message: "model is required", + StatusCode: http.StatusBadRequest, + LocalError: true, + } + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, http.StatusBadRequest, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"message":"model is required"`) + assert.Contains(t, recorder.Body.String(), `"code":"invalid_request_error"`) + assert.NotContains(t, recorder.Body.String(), "Invalid task protocol request") +} + +func TestRespondPluginProtocolSubmissionErrorKeepsGenericNonValidation400(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "protocol-generic-400", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [], done: false}; }, + renderFinal: function() { return {}; } + }}; + `) + c, recorder := newPluginProtocolTestContext(false, false) + deps := pluginProtocolTestDeps() + deps.submit = func(*gin.Context, *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return nil, &dto.TaskError{ + Code: "task_not_exist", + Message: "task_origin_not_exist", + StatusCode: http.StatusBadRequest, + LocalError: true, + } + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, http.StatusBadRequest, recorder.Code) + assert.Contains(t, recorder.Body.String(), "Invalid task protocol request") + assert.NotContains(t, recorder.Body.String(), "task_origin_not_exist") +} + +func compilePluginProtocolTestEndpoint(t *testing.T, key, source string) pluginruntime.PinnedEndpoint { + t.Helper() + return compilePluginProtocolTestEndpointWithOptions(t, key, source, pluginruntime.Options{}) +} + +func compilePluginProtocolTestEndpointWithOptions( + t *testing.T, + key string, + source string, + options pluginruntime.Options, +) pluginruntime.PinnedEndpoint { + t.Helper() + options.Key = key + options.Version = "1.0.0" + options.Concurrency = 1 + engine, err := pluginruntime.Compile(source, options) + require.NoError(t, err) + return pluginruntime.PinnedEndpoint{ + Generation: &pluginruntime.RoutingGeneration{Number: 41}, + Plugin: &pluginruntime.LoadedPlugin{ + Meta: pluginruntime.Meta{ + Key: key, + Version: "1.0.0", + Protocols: []pluginruntime.ProtocolClaim{{ + Name: "openai_responses", + Supports: []string{"stream", "sync", "background"}, + }}, + }, + Engine: engine, + }, + Protocol: "openai_responses", + Operation: pluginruntime.HostProtocolOperation{Name: "create", Methods: []string{http.MethodPost}, Path: "/v1/responses", BodyKinds: []pluginruntime.BodyKind{pluginruntime.BodyJSON}, ModelField: "model"}, + Model: "video-model", + } +} + +func newPluginProtocolTestContext(stream, requestBodyStream bool) (*gin.Context, *httptest.ResponseRecorder) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{}`)) + common.SetContextKey(c, constant.ContextKeyUserId, 71) + common.SetContextKey(c, constant.ContextKeyTokenId, 81) + common.SetContextKey(c, constant.ContextKeyUsingGroup, "default") + c.Set("resolved_task_model", "video-model") + c.Set(pluginruntime.ContextKeyProtocolRequest, pluginruntime.ProtocolRequestContext{ + RouteRequestContext: pluginruntime.RouteRequestContext{ + Path: "/v1/responses", + Method: http.MethodPost, + Params: map[string]string{}, + Query: map[string][]string{}, + Body: map[string]any{ + "kind": "json", + "value": map[string]any{"model": "video-model", "stream": requestBodyStream}, + }, + RequestBody: map[string]any{ + "model": "video-model", + "stream": requestBodyStream, + }, + }, + Protocol: "openai_responses", + Stream: stream, + }) + return c, recorder +} + +func pluginProtocolTestDeps() pluginProtocolBridgeDeps { + return pluginProtocolBridgeDeps{ + now: func() time.Time { return time.Unix(1_710_000_000, 0) }, + admissions: newPluginProtocolObservationLimiter(defaultPluginProtocolObservationLimits), + protocolLimits: relay.DefaultPluginProtocolLimits(), + observationTimeout: time.Hour, + tickInterval: time.Nanosecond, + tickJitter: 0, + heartbeatInterval: time.Hour, + admissionTimeout: time.Second, + } +} + +func pluginProtocolTestOutcome( + info *relaycommon.RelayInfo, + pluginKey string, + taskID string, + _ ...any, +) *taskSubmissionOutcome { + return &taskSubmissionOutcome{ + Result: &relay.TaskSubmitResult{}, + Task: &model.Task{ + TaskID: taskID, + Platform: constant.TaskPlatform(pluginKey), + UserId: info.UserId, + Status: model.TaskStatusSubmitted, + CreatedAt: 1_710_000_000, + }, + RelayInfo: info, + } +} + +func pluginProtocolTestSSEEventTypes(body string) []string { + lines := strings.Split(body, "\n") + events := make([]string, 0) + for _, line := range lines { + if after, ok := strings.CutPrefix(line, "event: "); ok { + events = append(events, after) + } + } + return events +} + +func setProtocolRequestBackground(c *gin.Context, background bool) { + request := c.MustGet(pluginruntime.ContextKeyProtocolRequest).(pluginruntime.ProtocolRequestContext) + if body, ok := request.Body.(map[string]any); ok { + if value, ok := body["value"].(map[string]any); ok { + value["background"] = background + } + } +} + +func compilePluginProtocolRetrieveEndpoint(t *testing.T, key, source string, options pluginruntime.Options) pluginruntime.PinnedEndpoint { + t.Helper() + pinned := compilePluginProtocolTestEndpointWithOptions(t, key, source, options) + pinned.Plugin.Meta.Protocols = []pluginruntime.ProtocolClaim{{Name: "openai_responses", Supports: []string{"stream", "sync", "background"}}} + return pinned +} + +func logsAppender(logs *[]string) pluginruntime.Options { + if logs == nil { + return pluginruntime.Options{} + } + return pluginruntime.Options{ + Log: func(message string) { *logs = append(*logs, message) }, + } +} + +func newPluginProtocolRetrieveContext(responseID string) (*gin.Context, *httptest.ResponseRecorder) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses/"+responseID, nil) + c.Params = gin.Params{{Key: "response_id", Value: responseID}} + common.SetContextKey(c, constant.ContextKeyUserId, 71) + common.SetContextKey(c, constant.ContextKeyTokenId, 81) + return c, recorder +} + +func pluginProtocolRetrieveDeps(pinned pluginruntime.PinnedEndpoint, task *model.Task, exists bool, err error) pluginProtocolBridgeDeps { + deps := pluginProtocolTestDeps() + deps.getByTaskId = func(userId int, taskId string) (*model.Task, bool, error) { + if !exists { + return nil, false, err + } + if task != nil && (userId != task.UserId || taskId != task.TaskID) { + return nil, false, err + } + return task, task != nil, err + } + deps.resolvePlugin = func(constant.TaskPlatform) (*pluginruntime.LoadedPlugin, *pluginruntime.RoutingGeneration, bool) { + if pinned.Plugin == nil { + return nil, nil, false + } + return pinned.Plugin, pinned.Generation, true + } + return deps +} diff --git a/controller/relay.go b/controller/relay.go index 8dccfe76dddd..0f7792efd970 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -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()) + if taskErr != nil { + 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() } - service.LogTaskConsumption(c, relayInfo) + 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) - 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, + 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()) + } } - 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 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, + }) +} - if taskErr != nil { - respondTaskError(c, taskErr) +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 { diff --git a/controller/relay_task_plugin_test.go b/controller/relay_task_plugin_test.go new file mode 100644 index 000000000000..e9405e2e6280 --- /dev/null +++ b/controller/relay_task_plugin_test.go @@ -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}, + } +} diff --git a/controller/task.go b/controller/task.go index a80f1a687aab..c514bdb325b5 100644 --- a/controller/task.go +++ b/controller/task.go @@ -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" ) +type taskArtifactResponse struct { + Key string `json:"key"` + Type string `json:"type"` + MimeType string `json:"mime_type,omitempty"` + ContentURL string `json:"content_url"` +} + +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"), - } - + 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:") +} diff --git a/controller/task_generic_test.go b/controller/task_generic_test.go new file mode 100644 index 000000000000..a3160b62ceb9 --- /dev/null +++ b/controller/task_generic_test.go @@ -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")) +} diff --git a/controller/task_log_view_test.go b/controller/task_log_view_test.go new file mode 100644 index 000000000000..b9ab2480c2ea --- /dev/null +++ b/controller/task_log_view_test.go @@ -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) +} diff --git a/controller/task_plugin.go b/controller/task_plugin.go new file mode 100644 index 000000000000..d88f44114bef --- /dev/null +++ b/controller/task_plugin.go @@ -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() + } +} diff --git a/controller/task_plugin_debug.go b/controller/task_plugin_debug.go new file mode 100644 index 000000000000..450f6b4b9972 --- /dev/null +++ b/controller/task_plugin_debug.go @@ -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" + } +} diff --git a/controller/task_plugin_debug_test.go b/controller/task_plugin_debug_test.go new file mode 100644 index 000000000000..ff6cc56dd97a --- /dev/null +++ b/controller/task_plugin_debug_test.go @@ -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) + } +} diff --git a/controller/task_plugin_test.go b/controller/task_plugin_test.go new file mode 100644 index 000000000000..e78db097b6c7 --- /dev/null +++ b/controller/task_plugin_test.go @@ -0,0 +1,1148 @@ +package controller + +import ( + "crypto/sha256" + "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/plugins" + "github.com/QuantumNous/new-api/setting" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupTaskPluginControllerTest(t *testing.T) { + t.Helper() + 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.Ability{}, &model.Task{}, &model.Option{})) + model.DB = database + t.Cleanup(func() { model.DB = originalDB }) +} + +const lifecyclePluginSource = ` +export const meta = {apiVersion: 1, key: "lifecycle-only", name: "Lifecycle", 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 {}; } +` + +func taskPluginControllerTestSource(key, version string) string { + return fmt.Sprintf(` +export const meta = {apiVersion: 1, key: %q, name: "Test", version: %q, author: {name: "Test"}, models: ["doc-1"], fetchMode: "per_task"}; +export function buildSubmitRequest() { return {}; } +export function parseSubmitResponse() { return {}; } +export function buildQueryRequest() { return {}; } +export function parseTaskResult() { return {}; } +`, key, version) +} + +func taskPluginControllerChannelSource(key, version string, channelType int) string { + return fmt.Sprintf(` +export const meta = {apiVersion: 1, key: %q, name: "Test", version: %q, author: {name: "Test"}, channelTypes: [%d], models: ["doc-1"], fetchMode: "per_task"}; +export function buildSubmitRequest() { return {}; } +export function parseSubmitResponse() { return {}; } +export function buildQueryRequest() { return {}; } +export function parseTaskResult() { return {}; } +`, key, version, channelType) +} + +func cleanupTaskPluginControllerRuntime(t *testing.T, key string) { + t.Helper() + t.Cleanup(func() { + jsplugin.DefaultRegistry.Unregister(key) + taskPluginSyncState.Lock() + delete(taskPluginSyncState.hashes, key) + delete(taskPluginSyncState.errors, key) + taskPluginSyncState.Unlock() + }) +} + +func TestDeleteThirdPartyPluginReportsAssociatedChannelsAndInFlightTasks(t *testing.T) { + setupTaskPluginControllerTest(t) + loaded, err := jsplugin.DefaultRegistry.Register(lifecyclePluginSource, jsplugin.Options{}) + require.NoError(t, err) + t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister("lifecycle-only") }) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{Key: loaded.Meta.Key, APIVersion: 1, Version: "1", Source: lifecyclePluginSource, SourceHash: "hash", Enabled: true})) + baseURL := "https://example.com" + setting := `{"task_plugin_key":"lifecycle-only"}` + channel := model.Channel{Type: constant.ChannelTypeTaskPlugin, Status: common.ChannelStatusEnabled, Name: "linked", Models: "doc", Group: "default", BaseURL: &baseURL, Setting: &setting} + require.NoError(t, channel.Insert()) + require.NoError(t, model.DB.Create(&model.Task{Platform: "lifecycle-only", Status: model.TaskStatusInProgress}).Error) + + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Params = gin.Params{{Key: "key", Value: "lifecycle-only"}, {Key: "version", Value: "1"}} + context.Request = httptest.NewRequest(http.MethodDelete, "/api/plugin/task/lifecycle-only/versions/1", nil) + DeleteTaskPluginVersion(context) + + assert.Contains(t, recorder.Body.String(), `"name":"linked"`) + assert.Contains(t, recorder.Body.String(), `"in_flight_count":1`) +} + +func TestDisableThirdPartyPluginSupportsCascadeAndForce(t *testing.T) { + setupTaskPluginControllerTest(t) + loaded, err := jsplugin.DefaultRegistry.Register(lifecyclePluginSource, jsplugin.Options{}) + require.NoError(t, err) + t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister("lifecycle-only") }) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{Key: loaded.Meta.Key, APIVersion: 1, Version: "1", Source: lifecyclePluginSource, SourceHash: "hash", Enabled: true})) + baseURL := "https://example.com" + setting := `{"task_plugin_key":"lifecycle-only"}` + channel := model.Channel{Type: constant.ChannelTypeTaskPlugin, Status: common.ChannelStatusEnabled, Name: "linked", Models: "doc", Group: "default", BaseURL: &baseURL, Setting: &setting} + require.NoError(t, channel.Insert()) + require.NoError(t, model.DB.Create(&model.Task{Platform: "lifecycle-only", Status: model.TaskStatusSubmitted}).Error) + + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Params = gin.Params{{Key: "key", Value: "lifecycle-only"}} + context.Request = httptest.NewRequest(http.MethodPost, "/api/plugin/task/lifecycle-only/status?cascade=true&force=true", strings.NewReader(`{"enabled":false}`)) + context.Request.Header.Set("Content-Type", "application/json") + SetTaskPluginStatus(context) + + assert.Contains(t, recorder.Body.String(), `"success":true`) + updated, err := model.GetChannelById(channel.Id, true) + require.NoError(t, err) + assert.Equal(t, common.ChannelStatusManuallyDisabled, updated.Status) +} + +func setupTaskPluginFactoryDisableTest(t *testing.T) { + t.Helper() + setupTaskPluginControllerTest(t) + originalMap := common.OptionMap + common.OptionMapRWMutex.Lock() + common.OptionMap = map[string]string{} + common.OptionMapRWMutex.Unlock() + t.Cleanup(func() { + jsplugin.DefaultRegistry.SetDisabledFactoryKeys(nil) + common.OptionMapRWMutex.Lock() + common.OptionMap = originalMap + common.OptionMapRWMutex.Unlock() + }) +} + +func postTaskPluginStatus(t *testing.T, key, query, body string) *httptest.ResponseRecorder { + t.Helper() + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Params = gin.Params{{Key: "key", Value: key}} + path := "/api/plugin/task/" + key + "/status" + if query != "" { + path += "?" + query + } + context.Request = httptest.NewRequest(http.MethodPost, path, strings.NewReader(body)) + context.Request.Header.Set("Content-Type", "application/json") + SetTaskPluginStatus(context) + return recorder +} + +func listTaskPluginItem(t *testing.T, key string) taskPluginListItem { + t.Helper() + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodGet, "/api/plugin/task", nil) + ListTaskPlugins(context) + var response struct { + Success bool `json:"success"` + Data []taskPluginListItem `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success) + for _, item := range response.Data { + if item.Meta.Key == key { + return item + } + } + t.Fatalf("task plugin %q not found", key) + return taskPluginListItem{} +} + +func taskPluginOptionsHasKey(t *testing.T, key string) bool { + t.Helper() + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodGet, "/api/task_plugin_options", nil) + GetTaskPluginOptions(context) + var response struct { + Success bool `json:"success"` + Data []struct { + Key string `json:"key"` + } `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success) + for _, option := range response.Data { + if option.Key == key { + return true + } + } + return false +} + +func TestDisableFactoryPluginPersistsOptionAndHidesFromBindOptions(t *testing.T) { + setupTaskPluginFactoryDisableTest(t) + const key = "kling" + + recorder := postTaskPluginStatus(t, key, "", `{"enabled":false}`) + assert.Contains(t, recorder.Body.String(), `"success":true`) + assert.Equal(t, []string{key}, setting.GetTaskPluginDisabledFactoryKeys()) + var stored model.Option + require.NoError(t, model.DB.Where("key = ?", setting.TaskPluginDisabledFactoryKeysKey).First(&stored).Error) + assert.Equal(t, `["kling"]`, stored.Value) + + item := listTaskPluginItem(t, key) + assert.Equal(t, "factory", item.Source) + assert.False(t, item.Enabled) + assert.Equal(t, "disabled", item.RuntimeStatus) + assert.False(t, taskPluginOptionsHasKey(t, key)) + _, ok := jsplugin.DefaultRegistry.Get(key) + assert.False(t, ok) + + recorder = postTaskPluginStatus(t, key, "", `{"enabled":true}`) + assert.Contains(t, recorder.Body.String(), `"success":true`) + assert.Empty(t, setting.GetTaskPluginDisabledFactoryKeys()) + item = listTaskPluginItem(t, key) + assert.True(t, item.Enabled) + assert.Equal(t, "registered", item.RuntimeStatus) + assert.True(t, taskPluginOptionsHasKey(t, key)) + _, ok = jsplugin.DefaultRegistry.Get(key) + assert.True(t, ok) +} + +func TestDisableFactoryPluginRespectsInUseGuard(t *testing.T) { + setupTaskPluginFactoryDisableTest(t) + const key = "kling" + baseURL := "https://example.com" + channelSetting := `{"task_plugin_key":"kling"}` + channel := model.Channel{Type: constant.ChannelTypeTaskPlugin, Status: common.ChannelStatusEnabled, Name: "linked-factory", Models: "doc", Group: "default", BaseURL: &baseURL, Setting: &channelSetting} + require.NoError(t, channel.Insert()) + + recorder := postTaskPluginStatus(t, key, "", `{"enabled":false}`) + assert.Contains(t, recorder.Body.String(), `"success":false`) + assert.Contains(t, recorder.Body.String(), "task plugin is still in use") + assert.Empty(t, setting.GetTaskPluginDisabledFactoryKeys()) + _, ok := jsplugin.DefaultRegistry.Get(key) + assert.True(t, ok) +} + +func TestDisableFactoryOverrideRowKeepsEnabledFlagPath(t *testing.T) { + setupTaskPluginFactoryDisableTest(t) + factorySource, err := plugins.Source("kling") + require.NoError(t, err) + overrideSource := strings.Replace(factorySource, `version: "1.0.0"`, `version: "1.0.0-test-factory-status"`, 1) + loaded, err := jsplugin.DefaultRegistry.Register(overrideSource, jsplugin.Options{}) + require.NoError(t, err) + t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister("kling") }) + plugin := model.TaskPlugin{ + Key: "kling", APIVersion: loaded.Meta.APIVersion, Version: loaded.Meta.Version, + Source: overrideSource, SourceHash: "test-hash", Enabled: true, + } + require.NoError(t, model.SaveTaskPlugin(&plugin)) + require.NoError(t, syncTaskPluginsOnce()) + + recorder := postTaskPluginStatus(t, "kling", "", `{"enabled":false}`) + assert.Contains(t, recorder.Body.String(), `"success":true`) + assert.Empty(t, setting.GetTaskPluginDisabledFactoryKeys()) + + row, err := model.GetTaskPluginVersion("kling", "") + require.NoError(t, err) + assert.False(t, row.Enabled) + + item := listTaskPluginItem(t, "kling") + assert.Equal(t, "override_over_factory", item.Source) + assert.False(t, item.Enabled) + assert.Equal(t, "disabled_fallback", item.RuntimeStatus) + assert.True(t, taskPluginOptionsHasKey(t, "kling")) + got, ok := jsplugin.DefaultRegistry.Get("kling") + require.True(t, ok) + assert.Equal(t, "1.0.0", got.Meta.Version) +} + +func TestListTaskPluginsIncludesFactoryWithoutDatabaseRows(t *testing.T) { + setupTaskPluginControllerTest(t) + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodGet, "/api/plugin/task", nil) + + ListTaskPlugins(context) + + var response struct { + Success bool `json:"success"` + Data []taskPluginListItem `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success) + var factoryItem *taskPluginListItem + for i := range response.Data { + if response.Data[i].Meta.Key == "kling" { + factoryItem = &response.Data[i] + break + } + } + require.NotNil(t, factoryItem) + assert.Equal(t, "factory", factoryItem.Source) + assert.Equal(t, "registered", factoryItem.RuntimeStatus) + assert.NotEmpty(t, factoryItem.SourceHash) +} + +func TestMasterSwitchEmptiesOptionsAndKeepsList(t *testing.T) { + setupTaskPluginControllerTest(t) + originalEnabled := constant.TaskPluginEnabled + jsplugin.DefaultRegistry.SetEnabled(false) + t.Cleanup(func() { + constant.TaskPluginEnabled = originalEnabled + jsplugin.DefaultRegistry.SetEnabled(originalEnabled) + }) + + assert.False(t, taskPluginOptionsHasKey(t, "kling")) + item := listTaskPluginItem(t, "kling") + assert.Equal(t, "factory", item.Source) + assert.Equal(t, "kling", item.Meta.Key) +} + +func TestGetTaskPluginOptionsIncludesUsageSchema(t *testing.T) { + const key = "usage-options-probe" + source := ` +export const meta = { + apiVersion: 1, key: "usage-options-probe", name: "Usage Options", version: "1.0.0", author: {name: "Test"}, + models: ["usage-options-model"], fetchMode: "per_task", + usageSchema: {seconds: {type: "number", unit: "second", description: "Generated media duration."}} +}; +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) }) + + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodGet, "/api/task_plugin_options", nil) + + GetTaskPluginOptions(context) + + var response struct { + Success bool `json:"success"` + Data []struct { + Key string `json:"key"` + UsageSchema map[string]jsplugin.UsageFieldSchema `json:"usageSchema"` + } `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success) + for _, option := range response.Data { + if option.Key != key { + continue + } + assert.Equal(t, "second", option.UsageSchema["seconds"].Unit) + assert.Equal(t, "Generated media duration.", option.UsageSchema["seconds"].Description["en"]) + return + } + t.Fatal("task plugin option not found") +} + +func TestListTaskPluginsShowsDisabledFallbackWhenOverridesAreDisabled(t *testing.T) { + setupTaskPluginControllerTest(t) + factorySource, err := plugins.Source("kling") + require.NoError(t, err) + overrideSource := strings.Replace(factorySource, `version: "1.0.0"`, `version: "1.0.0-test-disabled-override"`, 1) + loaded, err := jsplugin.DefaultRegistry.Register(overrideSource, jsplugin.Options{}) + require.NoError(t, err) + plugin := model.TaskPlugin{ + Key: "kling", APIVersion: loaded.Meta.APIVersion, Version: loaded.Meta.Version, + Source: overrideSource, SourceHash: "test-hash", Enabled: true, + } + require.NoError(t, model.SaveTaskPlugin(&plugin)) + originalEnabled := constant.TaskPluginOverrideEnabled + constant.TaskPluginOverrideEnabled = false + jsplugin.DefaultRegistry.SetOverrideEnabled(false) + t.Cleanup(func() { + constant.TaskPluginOverrideEnabled = originalEnabled + jsplugin.DefaultRegistry.SetOverrideEnabled(originalEnabled) + jsplugin.DefaultRegistry.Unregister("kling") + }) + + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodGet, "/api/plugin/task", nil) + ListTaskPlugins(context) + + var response struct { + Success bool `json:"success"` + Data []taskPluginListItem `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success) + for _, item := range response.Data { + if item.Meta.Key == "kling" { + assert.Equal(t, "disabled_fallback", item.RuntimeStatus) + return + } + } + t.Fatal("kling plugin not found") +} + +func TestDeleteActiveOverrideFallsBackToFactoryAndDeletesRecord(t *testing.T) { + setupTaskPluginControllerTest(t) + factorySource, err := plugins.Source("kling") + require.NoError(t, err) + overrideSource := strings.Replace(factorySource, `version: "1.0.0"`, `version: "1.0.0-test-override"`, 1) + loaded, err := jsplugin.DefaultRegistry.Register(overrideSource, jsplugin.Options{Key: "kling", Version: "test-override"}) + require.NoError(t, err) + t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister("kling") }) + plugin := model.TaskPlugin{ + Key: "kling", APIVersion: loaded.Meta.APIVersion, Version: loaded.Meta.Version, + Source: overrideSource, SourceHash: "test-hash", Enabled: true, + } + require.NoError(t, model.SaveTaskPlugin(&plugin)) + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Params = gin.Params{{Key: "key", Value: "kling"}, {Key: "version", Value: loaded.Meta.Version}} + context.Request = httptest.NewRequest(http.MethodDelete, "/api/plugin/task/kling/versions/"+loaded.Meta.Version, nil) + + DeleteTaskPluginVersion(context) + + assert.Contains(t, recorder.Body.String(), `"success":true`) + versions, err := model.ListTaskPluginVersions("kling") + require.NoError(t, err) + assert.Empty(t, versions) + runtimePlugin, ok := jsplugin.DefaultRegistry.Get("kling") + require.True(t, ok) + assert.NotEqual(t, loaded.Meta.Version, runtimePlugin.Meta.Version) +} + +func TestDeleteActiveTaskPluginPromotesEnabledVersionInRuntime(t *testing.T) { + setupTaskPluginControllerTest(t) + key := "delete-promote-probe" + cleanupTaskPluginControllerRuntime(t, key) + v1Source := taskPluginControllerTestSource(key, "1.0.0") + v2Source := taskPluginControllerTestSource(key, "2.0.0") + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: key, APIVersion: 1, Version: "1.0.0", Source: v1Source, SourceHash: "hash-v1", Enabled: true, + })) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: key, APIVersion: 1, Version: "2.0.0", Source: v2Source, SourceHash: "hash-v2", Enabled: true, + })) + _, err := jsplugin.DefaultRegistry.Register(v1Source, jsplugin.Options{Key: key, Version: "1.0.0"}) + require.NoError(t, err) + taskPluginSyncState.Lock() + taskPluginSyncState.hashes[key] = "hash-v1" + taskPluginSyncState.errors[key] = "stale compile error" + taskPluginSyncState.Unlock() + + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Params = gin.Params{{Key: "key", Value: key}, {Key: "version", Value: "1.0.0"}} + context.Request = httptest.NewRequest(http.MethodDelete, "/api/plugin/task/"+key+"/versions/1.0.0", nil) + DeleteTaskPluginVersion(context) + + assert.Contains(t, recorder.Body.String(), `"success":true`) + active, err := model.GetTaskPluginVersion(key, "") + require.NoError(t, err) + assert.Equal(t, "2.0.0", active.Version) + runtimePlugin, ok := jsplugin.DefaultRegistry.Get(key) + require.True(t, ok) + assert.Equal(t, "2.0.0", runtimePlugin.Meta.Version) + taskPluginSyncState.Lock() + syncedHash := taskPluginSyncState.hashes[key] + _, hasSyncError := taskPluginSyncState.errors[key] + taskPluginSyncState.Unlock() + assert.Equal(t, "hash-v2", syncedHash) + assert.False(t, hasSyncError) +} + +func TestUploadTaskPluginRefreshesRuntimeSyncState(t *testing.T) { + setupTaskPluginControllerTest(t) + key := "upload-sync-probe" + cleanupTaskPluginControllerRuntime(t, key) + source := taskPluginControllerTestSource(key, "1.0.0") + taskPluginSyncState.Lock() + taskPluginSyncState.hashes[key] = "stale-hash" + taskPluginSyncState.errors[key] = "stale compile error" + taskPluginSyncState.Unlock() + body, err := common.Marshal(map[string]any{"source": source}) + require.NoError(t, err) + + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodPost, "/api/plugin/task", strings.NewReader(string(body))) + context.Request.Header.Set("Content-Type", "application/json") + UploadTaskPlugin(context) + + assert.Contains(t, recorder.Body.String(), `"success":true`) + stored, err := model.GetTaskPluginVersion(key, "") + require.NoError(t, err) + taskPluginSyncState.Lock() + syncedHash := taskPluginSyncState.hashes[key] + _, hasSyncError := taskPluginSyncState.errors[key] + taskPluginSyncState.Unlock() + assert.Equal(t, stored.SourceHash, syncedHash) + assert.False(t, hasSyncError) +} + +func TestActivateTaskPluginRefreshesRuntimeSyncState(t *testing.T) { + setupTaskPluginControllerTest(t) + key := "activate-sync-probe" + cleanupTaskPluginControllerRuntime(t, key) + v1Source := taskPluginControllerTestSource(key, "1.0.0") + v2Source := taskPluginControllerTestSource(key, "2.0.0") + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: key, APIVersion: 1, Version: "1.0.0", Source: v1Source, SourceHash: "hash-v1", Enabled: true, + })) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: key, APIVersion: 1, Version: "2.0.0", Source: v2Source, SourceHash: "hash-v2", Enabled: true, + })) + _, err := jsplugin.DefaultRegistry.Register(v1Source, jsplugin.Options{Key: key, Version: "1.0.0"}) + require.NoError(t, err) + taskPluginSyncState.Lock() + taskPluginSyncState.hashes[key] = "hash-v1" + taskPluginSyncState.errors[key] = "stale compile error" + taskPluginSyncState.Unlock() + + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Params = gin.Params{{Key: "key", Value: key}} + context.Request = httptest.NewRequest(http.MethodPost, "/api/plugin/task/"+key+"/activate", strings.NewReader(`{"version":"2.0.0"}`)) + context.Request.Header.Set("Content-Type", "application/json") + ActivateTaskPlugin(context) + + assert.Contains(t, recorder.Body.String(), `"success":true`) + runtimePlugin, ok := jsplugin.DefaultRegistry.Get(key) + require.True(t, ok) + assert.Equal(t, "2.0.0", runtimePlugin.Meta.Version) + taskPluginSyncState.Lock() + syncedHash := taskPluginSyncState.hashes[key] + _, hasSyncError := taskPluginSyncState.errors[key] + taskPluginSyncState.Unlock() + assert.Equal(t, "hash-v2", syncedHash) + assert.False(t, hasSyncError) +} + +func TestSyncTaskPluginsPublishesOneGenerationForWholeBatch(t *testing.T) { + setupTaskPluginControllerTest(t) + firstKey := "batch-sync-first" + secondKey := "batch-sync-second" + cleanupTaskPluginControllerRuntime(t, firstKey) + cleanupTaskPluginControllerRuntime(t, secondKey) + firstSource := taskPluginControllerTestSource(firstKey, "1.0.0") + secondSource := taskPluginControllerTestSource(secondKey, "1.0.0") + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: firstKey, APIVersion: 1, Version: "1.0.0", Source: firstSource, SourceHash: "first-hash", Enabled: true, + })) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: secondKey, APIVersion: 1, Version: "1.0.0", Source: secondSource, SourceHash: "second-hash", Enabled: true, + })) + before := jsplugin.DefaultRegistry.Generation().Number + + SyncTaskPluginsOnce() + + assert.Equal(t, before+1, jsplugin.DefaultRegistry.Generation().Number) + _, firstRegistered := jsplugin.DefaultRegistry.Get(firstKey) + _, secondRegistered := jsplugin.DefaultRegistry.Get(secondKey) + assert.True(t, firstRegistered) + assert.True(t, secondRegistered) + + published := jsplugin.DefaultRegistry.Generation() + SyncTaskPluginsOnce() + assert.Same(t, published, jsplugin.DefaultRegistry.Generation()) +} + +func TestTaskPluginRuntimeExposesDatabaseRevisionAheadOfLocalGeneration(t *testing.T) { + setupTaskPluginControllerTest(t) + key := "runtime-revision-probe" + cleanupTaskPluginControllerRuntime(t, key) + taskPluginSyncState.Lock() + previousRebuild := taskPluginSyncState.lastRebuild + taskPluginSyncState.Unlock() + t.Cleanup(func() { + taskPluginSyncState.Lock() + taskPluginSyncState.lastRebuild = previousRebuild + taskPluginSyncState.Unlock() + }) + + v1Source := taskPluginControllerTestSource(key, "1.0.0") + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: key, APIVersion: 1, Version: "1.0.0", + Source: v1Source, SourceHash: "runtime-v1", Enabled: true, + })) + require.NoError(t, syncTaskPluginsOnce()) + localGeneration := jsplugin.DefaultRegistry.Generation().Number + taskPluginSyncState.Lock() + syncedRevision := taskPluginSyncState.lastRebuild.DatabaseRevision + taskPluginSyncState.Unlock() + + v2Source := taskPluginControllerTestSource(key, "2.0.0") + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: key, APIVersion: 1, Version: "2.0.0", + Source: v2Source, SourceHash: "runtime-v2", Enabled: true, + })) + require.NoError(t, model.ActivateTaskPlugin(key, "2.0.0")) + + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodGet, "/api/plugin/task/runtime/status", nil) + GetTaskPluginRuntime(context) + + var response struct { + Success bool `json:"success"` + Data taskPluginRuntimeStatus `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success) + assert.Equal(t, localGeneration, response.Data.CurrentGeneration) + assert.NotZero(t, response.Data.GenerationPublishedAt) + assert.NotEqual(t, syncedRevision, response.Data.DatabaseRevision) + assert.Equal(t, "success", response.Data.LastRebuild.Status) + assert.Equal(t, syncedRevision, response.Data.LastRebuild.DatabaseRevision) + assert.Empty(t, response.Data.PluginErrors) + + active, ok := jsplugin.DefaultRegistry.Get(key) + require.True(t, ok) + assert.Equal(t, "1.0.0", active.Meta.Version) +} + +func TestTaskPluginRuntimeReportsPluginLevelCompileErrors(t *testing.T) { + setupTaskPluginControllerTest(t) + key := "runtime-error-probe" + cleanupTaskPluginControllerRuntime(t, key) + taskPluginSyncState.Lock() + previousRebuild := taskPluginSyncState.lastRebuild + taskPluginSyncState.Unlock() + t.Cleanup(func() { + taskPluginSyncState.Lock() + taskPluginSyncState.lastRebuild = previousRebuild + taskPluginSyncState.Unlock() + }) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: key, APIVersion: 1, Version: "1.0.0", + Source: "export const meta = {", SourceHash: "broken-source", Enabled: true, + })) + require.NoError(t, syncTaskPluginsOnce()) + + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodGet, "/api/plugin/task/runtime/status", nil) + GetTaskPluginRuntime(context) + + var response struct { + Success bool `json:"success"` + Data taskPluginRuntimeStatus `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success) + assert.Equal(t, "partial", response.Data.LastRebuild.Status) + assert.Equal(t, response.Data.DatabaseRevision, response.Data.LastRebuild.DatabaseRevision) + assert.GreaterOrEqual(t, response.Data.LastRebuild.PluginErrorCount, 1) + assert.NotEmpty(t, response.Data.PluginErrors[key]) + _, registered := jsplugin.DefaultRegistry.Get(key) + assert.False(t, registered) +} + +func TestTaskPluginRuntimeSurvivesDatabaseSyncFailure(t *testing.T) { + setupTaskPluginControllerTest(t) + key := "runtime-database-failure" + cleanupTaskPluginControllerRuntime(t, key) + taskPluginSyncState.Lock() + previousRebuild := taskPluginSyncState.lastRebuild + taskPluginSyncState.Unlock() + t.Cleanup(func() { + taskPluginSyncState.Lock() + taskPluginSyncState.lastRebuild = previousRebuild + taskPluginSyncState.Unlock() + }) + + source := taskPluginControllerTestSource(key, "1.0.0") + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: key, APIVersion: 1, Version: "1.0.0", + Source: source, SourceHash: "runtime-database-v1", Enabled: true, + })) + require.NoError(t, syncTaskPluginsOnce()) + generation := jsplugin.DefaultRegistry.Generation().Number + taskPluginSyncState.Lock() + syncedRevision := taskPluginSyncState.lastRebuild.DatabaseRevision + taskPluginSyncState.Unlock() + + sqlDatabase, err := model.DB.DB() + require.NoError(t, err) + require.NoError(t, sqlDatabase.Close()) + require.Error(t, syncTaskPluginsOnce()) + + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodGet, "/api/plugin/task/runtime/status", nil) + GetTaskPluginRuntime(context) + + var response struct { + Success bool `json:"success"` + Data taskPluginRuntimeStatus `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success) + assert.Equal(t, generation, response.Data.CurrentGeneration) + assert.Equal(t, syncedRevision, response.Data.DatabaseRevision) + assert.Equal(t, "database snapshot unavailable", response.Data.DatabaseError) + assert.Equal(t, "failed", response.Data.LastRebuild.Status) + assert.Equal(t, syncedRevision, response.Data.LastRebuild.DatabaseRevision) + assert.Contains(t, response.Data.LastRebuild.Error, "sync task plugins") +} + +func TestSyncTaskPluginsCachesRejectedDesiredSourceWithoutLosingIncumbent(t *testing.T) { + setupTaskPluginControllerTest(t) + pluginKey := "sync-retained-plugin" + ownerKey := "sync-retained-owner" + cleanupTaskPluginControllerRuntime(t, pluginKey) + cleanupTaskPluginControllerRuntime(t, ownerKey) + v1Source := taskPluginControllerChannelSource(pluginKey, "1.0.0", 9001) + v2Source := taskPluginControllerChannelSource(pluginKey, "2.0.0", 9002) + ownerSource := taskPluginControllerChannelSource(ownerKey, "1.0.0", 9002) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: pluginKey, APIVersion: 1, Version: "1.0.0", Source: v1Source, SourceHash: "retained-v1", Enabled: true, + })) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: ownerKey, APIVersion: 1, Version: "1.0.0", Source: ownerSource, SourceHash: "owner-v1", Enabled: true, + })) + require.NoError(t, syncTaskPluginsOnce()) + + incumbent, ok := jsplugin.DefaultRegistry.Get(pluginKey) + require.True(t, ok) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: pluginKey, APIVersion: 1, Version: "2.0.0", Source: v2Source, SourceHash: "retained-v2", Enabled: true, + })) + require.NoError(t, model.ActivateTaskPlugin(pluginKey, "2.0.0")) + + require.NoError(t, syncTaskPluginsOnce()) + rejectedGeneration := jsplugin.DefaultRegistry.Generation() + for range 2 { + require.NoError(t, syncTaskPluginsOnce()) + active, found := jsplugin.DefaultRegistry.Get(pluginKey) + require.True(t, found) + assert.Same(t, incumbent, active) + assert.Equal(t, "2.0.0", jsplugin.DefaultRegistry.OverridePlugins()[pluginKey].Meta.Version) + assert.Same(t, incumbent, jsplugin.DefaultRegistry.ActiveOverridePlugins()[pluginKey]) + assert.Contains(t, jsplugin.DefaultRegistry.RoutingErrors()[pluginKey], "channelType 9002 conflicts") + taskPluginSyncState.Lock() + assert.Equal(t, "retained-v2", taskPluginSyncState.hashes[pluginKey]) + taskPluginSyncState.Unlock() + assert.Same(t, rejectedGeneration, jsplugin.DefaultRegistry.Generation()) + } + + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: ownerKey, APIVersion: 1, Version: "2.0.0", + Source: taskPluginControllerChannelSource(ownerKey, "2.0.0", 9003), SourceHash: "owner-v2", Enabled: true, + })) + require.NoError(t, model.ActivateTaskPlugin(ownerKey, "2.0.0")) + require.NoError(t, syncTaskPluginsOnce()) + active, ok := jsplugin.DefaultRegistry.Get(pluginKey) + require.True(t, ok) + assert.Equal(t, "2.0.0", active.Meta.Version) + assert.NotContains(t, jsplugin.DefaultRegistry.RoutingErrors(), pluginKey) +} + +func TestSyncTaskPluginsPreservesLastCompiledOverrideWhileOverridesAreDisabled(t *testing.T) { + setupTaskPluginControllerTest(t) + key := "sync-disabled-override" + cleanupTaskPluginControllerRuntime(t, key) + v1Source := taskPluginControllerTestSource(key, "1.0.0") + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: key, APIVersion: 1, Version: "1.0.0", Source: v1Source, SourceHash: "disabled-v1", Enabled: true, + })) + require.NoError(t, syncTaskPluginsOnce()) + jsplugin.DefaultRegistry.SetOverrideEnabled(false) + t.Cleanup(func() { jsplugin.DefaultRegistry.SetOverrideEnabled(true) }) + + disabledGeneration := jsplugin.DefaultRegistry.Generation() + require.NoError(t, syncTaskPluginsOnce()) + assert.Same(t, disabledGeneration, jsplugin.DefaultRegistry.Generation()) + + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: key, APIVersion: 1, Version: "2.0.0", + Source: "export const meta = {", SourceHash: "disabled-v2", Enabled: true, + })) + require.NoError(t, model.ActivateTaskPlugin(key, "2.0.0")) + require.NoError(t, syncTaskPluginsOnce()) + assert.Equal(t, "1.0.0", jsplugin.DefaultRegistry.OverridePlugins()[key].Meta.Version) + taskPluginSyncState.Lock() + assert.Equal(t, "disabled-v1", taskPluginSyncState.hashes[key]) + assert.NotEmpty(t, taskPluginSyncState.errors[key]) + taskPluginSyncState.Unlock() + + jsplugin.DefaultRegistry.SetOverrideEnabled(true) + active, ok := jsplugin.DefaultRegistry.Get(key) + require.True(t, ok) + assert.Equal(t, "1.0.0", active.Meta.Version) +} + +const dryRunPluginSource = ` +export const meta = {apiVersion: 1, key: "dryrun-probe", name: "DryRun", version: "1.0.0", author: {name: "Test"}, models: ["doc-1"], fetchMode: "per_task"}; +export function buildSubmitRequest(payload) { + if (!payload || !payload.model) { throw new Error("model required"); } + return {model: payload.model}; +} +export function parseSubmitResponse() { return {}; } +export function buildQueryRequest() { return {}; } +export function parseTaskResult() { return {}; } +export const native = { info: function(ctx, task) { return "task:" + task.id; } }; +` + +func runTaskPluginDryRun(t *testing.T, body string) *httptest.ResponseRecorder { + t.Helper() + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Params = gin.Params{{Key: "key", Value: "dryrun-probe"}} + context.Request = httptest.NewRequest(http.MethodPost, "/api/plugin/task/dryrun-probe/dryrun", strings.NewReader(body)) + context.Request.Header.Set("Content-Type", "application/json") + DryRunTaskPlugin(context) + return recorder +} + +func TestDryRunTaskPluginExecutesHookAndRendererMember(t *testing.T) { + setupTaskPluginControllerTest(t) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{Key: "dryrun-probe", APIVersion: 1, Version: "1.0.0", Source: dryRunPluginSource, SourceHash: "hash", Enabled: true})) + + hookRecorder := runTaskPluginDryRun(t, `{"hook":"buildSubmitRequest","args":[{"model":"doc-1"}]}`) + assert.Contains(t, hookRecorder.Body.String(), `"success":true`) + assert.Contains(t, hookRecorder.Body.String(), `"model":"doc-1"`) + + memberRecorder := runTaskPluginDryRun(t, `{"hook":"native","member":"info","args":[{}, {"id":"t-1"}]}`) + assert.Contains(t, memberRecorder.Body.String(), `"success":true`) + assert.Contains(t, memberRecorder.Body.String(), "task:t-1") +} + +func TestDryRunTaskPluginReportsUnknownHook(t *testing.T) { + setupTaskPluginControllerTest(t) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{Key: "dryrun-probe", APIVersion: 1, Version: "1.0.0", Source: dryRunPluginSource, SourceHash: "hash", Enabled: true})) + + recorder := runTaskPluginDryRun(t, `{"hook":"missingHook"}`) + + assert.Contains(t, recorder.Body.String(), `"success":false`) + assert.Contains(t, recorder.Body.String(), `plugin export \"missingHook\" not found`) +} + +func TestDryRunTaskPluginSurfacesBadArgumentErrors(t *testing.T) { + setupTaskPluginControllerTest(t) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{Key: "dryrun-probe", APIVersion: 1, Version: "1.0.0", Source: dryRunPluginSource, SourceHash: "hash", Enabled: true})) + + malformedRecorder := runTaskPluginDryRun(t, `{"hook":"buildSubmitRequest","args":[{`) + assert.Contains(t, malformedRecorder.Body.String(), `"success":false`) + + rejectedRecorder := runTaskPluginDryRun(t, `{"hook":"buildSubmitRequest","args":[{}]}`) + assert.Contains(t, rejectedRecorder.Body.String(), `"success":false`) + assert.Contains(t, rejectedRecorder.Body.String(), "model required") +} + +func TestUploadTaskPluginPreflightConflict(t *testing.T) { + setupTaskPluginControllerTest(t) + enabledFalse := false + tests := []struct { + name string + key string + enabled *bool + force bool + wantSuccess bool + wantError string + }{ + { + name: "enabled conflict rejected", + key: "preflight-reject", + wantError: "channelType 50 conflicts", + }, + { + name: "force saves despite conflict", + key: "preflight-force", + force: true, + wantSuccess: true, + }, + { + name: "disabled skips preflight", + key: "preflight-disabled", + enabled: &enabledFalse, + wantSuccess: true, + }, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + cleanupTaskPluginControllerRuntime(t, testCase.key) + source := taskPluginControllerChannelSource(testCase.key, "1.0.0", 50) + payload := map[string]any{"source": source} + if testCase.enabled != nil { + payload["enabled"] = *testCase.enabled + } + if testCase.force { + payload["force"] = true + } + body, err := common.Marshal(payload) + require.NoError(t, err) + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodPost, "/api/plugin/task", strings.NewReader(string(body))) + context.Request.Header.Set("Content-Type", "application/json") + + UploadTaskPlugin(context) + + if testCase.wantSuccess { + assert.Contains(t, recorder.Body.String(), `"success":true`) + _, err = model.GetTaskPluginVersion(testCase.key, "") + require.NoError(t, err) + return + } + assert.Contains(t, recorder.Body.String(), `"success":false`) + assert.Contains(t, recorder.Body.String(), testCase.wantError) + assert.Contains(t, recorder.Body.String(), "kling") + var count int64 + require.NoError(t, model.DB.Model(&model.TaskPlugin{}).Where("key = ?", testCase.key).Count(&count).Error) + assert.Zero(t, count) + }) + } +} + +func TestUploadTaskPluginRejectsMetaViolatingV1Schema(t *testing.T) { + setupTaskPluginControllerTest(t) + cases := []struct { + name string + meta string + expectedError string + }{ + { + name: "key with uppercase characters", + meta: `{apiVersion: 1, key: "Bad-Key", name: "Bad", version: "1.0.0", author: {name: "Test"}, models: ["doc-1"], fetchMode: "per_task"}`, + expectedError: "plugin meta key must match", + }, + { + name: "version that is not semver", + meta: `{apiVersion: 1, key: "bad-plugin", name: "Bad", version: "one", author: {name: "Test"}, models: ["doc-1"], fetchMode: "per_task"}`, + expectedError: "plugin meta version must be semver", + }, + { + name: "unsupported fetch mode", + meta: `{apiVersion: 1, key: "bad-plugin", name: "Bad", version: "1.0.0", author: {name: "Test"}, models: ["doc-1"], fetchMode: "sometimes"}`, + expectedError: "plugin meta fetchMode must be per_task or batch", + }, + { + name: "empty model list", + meta: `{apiVersion: 1, key: "bad-plugin", name: "Bad", version: "1.0.0", author: {name: "Test"}, models: [], fetchMode: "per_task"}`, + expectedError: "plugin meta models must contain at least one model", + }, + } + for _, testCase := range cases { + t.Run(testCase.name, func(t *testing.T) { + source := "export const meta = " + testCase.meta + `; +export function buildSubmitRequest() { return {}; } +export function parseSubmitResponse() { return {}; } +export function buildQueryRequest() { return {}; } +export function parseTaskResult() { return {}; } +` + body, err := common.Marshal(map[string]any{"source": source}) + require.NoError(t, err) + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodPost, "/api/plugin/task", strings.NewReader(string(body))) + context.Request.Header.Set("Content-Type", "application/json") + + UploadTaskPlugin(context) + + assert.Contains(t, recorder.Body.String(), `"success":false`) + assert.Contains(t, recorder.Body.String(), testCase.expectedError) + var count int64 + require.NoError(t, model.DB.Model(&model.TaskPlugin{}).Count(&count).Error) + assert.Zero(t, count) + }) + } +} + +func TestDeletePureFactoryPluginIsRejected(t *testing.T) { + setupTaskPluginControllerTest(t) + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Params = gin.Params{{Key: "key", Value: "kling"}, {Key: "version", Value: "1.0.0"}} + context.Request = httptest.NewRequest(http.MethodDelete, "/api/plugin/task/kling/versions/1.0.0", nil) + + DeleteTaskPluginVersion(context) + + assert.Contains(t, recorder.Body.String(), `"success":false`) + assert.Contains(t, recorder.Body.String(), "factory plugins cannot be deleted") + _, ok := jsplugin.DefaultRegistry.Get("kling") + assert.True(t, ok) +} + +func TestUploadTaskPluginSourceSha256(t *testing.T) { + setupTaskPluginControllerTest(t) + tests := []struct { + name string + key string + withHash bool + hash string + wantSuccess bool + wantError string + }{ + { + name: "matching hash succeeds", + key: "sha256-match", + withHash: true, + wantSuccess: true, + }, + { + name: "mismatching hash rejected", + key: "sha256-mismatch", + withHash: true, + hash: "deadbeef", + wantError: "plugin source sha256 mismatch", + }, + { + name: "absent field unchanged", + key: "sha256-absent", + wantSuccess: true, + }, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + cleanupTaskPluginControllerRuntime(t, testCase.key) + source := taskPluginControllerTestSource(testCase.key, "1.0.0") + payload := map[string]any{"source": source} + if testCase.withHash { + hash := testCase.hash + if hash == "" { + hash = " " + strings.ToUpper(fmt.Sprintf("%x", sha256.Sum256([]byte(source)))) + " " + } + payload["sourceSha256"] = hash + } + body, err := common.Marshal(payload) + require.NoError(t, err) + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodPost, "/api/plugin/task", strings.NewReader(string(body))) + context.Request.Header.Set("Content-Type", "application/json") + + UploadTaskPlugin(context) + + if testCase.wantSuccess { + assert.Contains(t, recorder.Body.String(), `"success":true`) + _, err = model.GetTaskPluginVersion(testCase.key, "") + require.NoError(t, err) + return + } + assert.Contains(t, recorder.Body.String(), `"success":false`) + assert.Contains(t, recorder.Body.String(), testCase.wantError) + var count int64 + require.NoError(t, model.DB.Model(&model.TaskPlugin{}).Where("key = ?", testCase.key).Count(&count).Error) + assert.Zero(t, count) + }) + } +} + +func setupTaskPluginMarketplaceSourcesTest(t *testing.T) { + t.Helper() + setupTaskPluginControllerTest(t) + originalMap := common.OptionMap + common.OptionMapRWMutex.Lock() + common.OptionMap = map[string]string{} + common.OptionMapRWMutex.Unlock() + t.Cleanup(func() { + common.OptionMapRWMutex.Lock() + common.OptionMap = originalMap + common.OptionMapRWMutex.Unlock() + }) +} + +func TestGetTaskPluginMarketplaceSourcesDefaultWhenUnset(t *testing.T) { + setupTaskPluginMarketplaceSourcesTest(t) + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodGet, "/api/plugin/task/marketplace/sources", nil) + + GetTaskPluginMarketplaceSources(context) + + var response struct { + Success bool `json:"success"` + Data []setting.TaskPluginMarketplaceSource `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success) + require.Equal(t, []setting.TaskPluginMarketplaceSource{ + {Name: "Official", IndexURL: "https://www.newapi.ai/api/v1/plugins/index.json"}, + {Name: "GitHub", IndexURL: "https://raw.githubusercontent.com/QuantumNous/new-api-plugins/main/index.json"}, + }, response.Data) + var count int64 + require.NoError(t, model.DB.Model(&model.Option{}).Where("key = ?", setting.TaskPluginMarketplaceSourcesKey).Count(&count).Error) + assert.Zero(t, count) +} + +func TestUpdateTaskPluginMarketplaceSourcesRoundTrip(t *testing.T) { + setupTaskPluginMarketplaceSourcesTest(t) + payload := []setting.TaskPluginMarketplaceSource{ + {Name: "Mirror", IndexURL: "https://example.com/plugins/index.json"}, + {Name: "Official", IndexURL: "https://www.newapi.ai/api/v1/plugins/index.json"}, + } + body, err := common.Marshal(payload) + require.NoError(t, err) + putRecorder := httptest.NewRecorder() + putContext, _ := gin.CreateTestContext(putRecorder) + putContext.Request = httptest.NewRequest(http.MethodPut, "/api/plugin/task/marketplace/sources", strings.NewReader(string(body))) + putContext.Request.Header.Set("Content-Type", "application/json") + + UpdateTaskPluginMarketplaceSources(putContext) + + assert.Contains(t, putRecorder.Body.String(), `"success":true`) + + getRecorder := httptest.NewRecorder() + getContext, _ := gin.CreateTestContext(getRecorder) + getContext.Request = httptest.NewRequest(http.MethodGet, "/api/plugin/task/marketplace/sources", nil) + GetTaskPluginMarketplaceSources(getContext) + + var response struct { + Success bool `json:"success"` + Data []setting.TaskPluginMarketplaceSource `json:"data"` + } + require.NoError(t, common.Unmarshal(getRecorder.Body.Bytes(), &response)) + require.True(t, response.Success) + assert.Equal(t, payload, response.Data) +} + +func TestUpdateTaskPluginMarketplaceSourcesValidation(t *testing.T) { + setupTaskPluginMarketplaceSourcesTest(t) + tests := []struct { + name string + body string + wantErr string + }{ + { + name: "empty name", + body: `[{"name":"","index_url":"https://example.com/index.json"}]`, + wantErr: "marketplace source name is required", + }, + { + name: "invalid URL", + body: `[{"name":"Local","index_url":"not-a-url"}]`, + wantErr: "marketplace source index_url must be an absolute http(s) URL", + }, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodPut, "/api/plugin/task/marketplace/sources", strings.NewReader(testCase.body)) + context.Request.Header.Set("Content-Type", "application/json") + + UpdateTaskPluginMarketplaceSources(context) + + assert.Contains(t, recorder.Body.String(), `"success":false`) + assert.Contains(t, recorder.Body.String(), testCase.wantErr) + assert.Empty(t, common.OptionMap[setting.TaskPluginMarketplaceSourcesKey]) + var count int64 + require.NoError(t, model.DB.Model(&model.Option{}).Where("key = ?", setting.TaskPluginMarketplaceSourcesKey).Count(&count).Error) + assert.Zero(t, count) + }) + } +} diff --git a/controller/video_proxy.go b/controller/video_proxy.go index 996d084d88fa..230d2bddfd38 100644 --- a/controller/video_proxy.go +++ b/controller/video_proxy.go @@ -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 } + 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)) + } + } + 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 + } + descriptor = &relaychannel.TaskContentRequest{ + URL: resultURL, + Method: c.Request.Method, + Credentialless: true, + } + } + 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 { - 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 + return &taskMediaProxyError{ + status: http.StatusServiceUnavailable, code: "artifact_plugin_unavailable", + message: "Artifact channel is unavailable", err: err, + } } - baseURL := channel.GetBaseURL() - if baseURL == "" { - baseURL = "https://api.openai.com" + 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, + } } - 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") - return + return &taskMediaProxyError{ + status: http.StatusInternalServerError, code: "artifact_internal_error", + message: "Failed to create artifact proxy client", err: err, + } } } + if client == nil { + client = http.DefaultClient + } - ctx, cancel := context.WithTimeout(c.Request.Context(), 60*time.Second) - defer cancel() - req, err := http.NewRequestWithContext(ctx, http.MethodGet, "", nil) + req, err := http.NewRequestWithContext(c.Request.Context(), method, parsedURL.String(), bytes.NewReader(descriptor.Body)) 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 + 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) } - 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 + 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, + } } - 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 + 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, + } } - 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 + 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", } - 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() + return &taskMediaProxyError{ + status: http.StatusBadGateway, code: "artifact_upstream_error", + message: fmt.Sprintf("Artifact upstream returned status %d", resp.StatusCode), + } } +} - 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 +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} + }() - 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") + 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 } - return + 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 +} - var validateErr error +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 == "" { - 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 + 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, + ) +} - 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 +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)) +} - resp, err := client.Do(req) +func normalizeTaskMediaHost(scheme, host string) string { + host = strings.ToLower(strings.TrimSpace(host)) + hostname, port, err := net.SplitHostPort(host) 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 + return strings.TrimSuffix(host, ".") } - defer resp.Body.Close() + 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) +} - 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 +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") +} - for key, values := range resp.Header { - for _, value := range values { - c.Writer.Header().Add(key, value) +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") +} - 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())) +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 } diff --git a/controller/video_proxy_gemini.go b/controller/video_proxy_gemini.go deleted file mode 100644 index 0c76e33c709a..000000000000 --- a/controller/video_proxy_gemini.go +++ /dev/null @@ -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) -} diff --git a/docs/plugin-api/README.md b/docs/plugin-api/README.md new file mode 100644 index 000000000000..bb188cec040d --- /dev/null +++ b/docs/plugin-api/README.md @@ -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. diff --git a/docs/plugin-api/v1.d.ts b/docs/plugin-api/v1.d.ts new file mode 100644 index 000000000000..11911727e5b8 --- /dev/null +++ b/docs/plugin-api/v1.d.ts @@ -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>}> + | Readonly<{kind: "multipart"; fields: Readonly>; files: readonly FileReference[]}> + | Readonly<{kind: "none"}>; + +export interface NativeDecodeContext {method: string; path: string; params: Readonly>; query: Readonly>; 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); +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>}; +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>; 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} +export interface DriverContext {requestBody: unknown; requestHeaders: Readonly>; 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; /** 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>; 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 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> | null; +export declare function extractUsageOnSubmit(ctx: DriverContext, taskData: unknown): Readonly> | null; +export declare function extractUsageOnComplete(task: TaskView, result: NormalizedTaskResult, data: unknown): Readonly> | 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>}}): RequestDescriptor; diff --git a/docs/plugin-api/v1.md b/docs/plugin-api/v1.md new file mode 100644 index 000000000000..814513b06ffa --- /dev/null +++ b/docs/plugin-api/v1.md @@ -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: