package middleware import ( "bytes" "context" "encoding/json" "fmt" "io" "mime/multipart" "net/http" "net/http/httptest" "net/textproto" "os" "strings" "testing" "testing/iotest" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/dto" appI18n "github.com/QuantumNous/new-api/i18n" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/pkg/jsplugin" builtinplugins "github.com/QuantumNous/new-api/plugins" kitdto "github.com/QuantumNous/new-api/relaykit/dto" "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" "gorm.io/gorm/schema" ) const genericTaskPluginSource = ` export const meta = {apiVersion: 1, key: "generic-entry-test", name: "Generic", version: "1.0.0", author: {name: "Test"}, models: ["doc"], fetchMode: "per_task"}; export function buildSubmitRequest() { return {url: "https://example.com"}; } export function parseSubmitResponse() { return {taskId: "one"}; } export function buildQueryRequest() { return {url: "https://example.com"}; } export function parseTaskResult() { return {status: "SUCCESS"}; } ` func TestPrepareTaskPluginSubmitRejectsMissingModel(t *testing.T) { _, err := jsplugin.DefaultRegistry.Register(genericTaskPluginSource, jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister("generic-entry-test") }) recorder := httptest.NewRecorder() c, _ := gin.CreateTestContext(recorder) c.Params = gin.Params{{Key: "key", Value: "generic-entry-test"}} c.Request = httptest.NewRequest(http.MethodPost, "/v1/tasks/generic-entry-test", strings.NewReader(`{"prompt":"x"}`)) c.Request.Header.Set("Content-Type", "application/json") PrepareTaskPluginSubmit()(c) assert.Equal(t, http.StatusBadRequest, recorder.Code) assert.Contains(t, recorder.Body.String(), "model is required") } func TestPrepareTaskPluginRouteUsesCanonicalContextAndResolvedSubmit(t *testing.T) { plugin := compileTaskRoutePlugin(t, ` export const meta = { apiVersion: 1, key: "route-submit-test", name: "Submit", version: "1.0.0", author: {name: "Test"}, models: ["resolved-model"], fetchMode: "per_task", routes: [{method: "POST", path: "/vendor/jobs/:category", type: "submit", action: "static-action", decode: "decodeJob", render: "jobCreated"}], }; export const native = { decodeJob: function(ctx) { if (ctx.path !== "/vendor/jobs/video" || ctx.method !== "POST") throw new Error("bad path"); if (ctx.params.category !== "video") throw new Error("bad params"); if (ctx.query.tag.length !== 2 || ctx.query.tag[0] !== "first" || ctx.query.tag[1] !== "second") throw new Error("bad query"); if (ctx.body.kind !== "json" || !Array.isArray(ctx.body.value) || ctx.body.value[0] !== "prompt") throw new Error("bad body"); // Decoders read the client's headers, without its gateway credentials. if (ctx.requestHeaders["X-Dashscope-Async"] !== "enable" || ctx.requestHeaders.Authorization !== undefined) throw new Error("bad headers"); return {kind: "submit", model: "resolved-model", action: "resolved-action", requestBody: {prompt: "normalized"}}; }, jobCreated: function(ctx, task) { return task; }, }; export function buildSubmitRequest() { return {url: "https://example.com"}; } export function parseSubmitResponse() { return {taskId: "one"}; } export function buildQueryRequest() { return {url: "https://example.com"}; } export function parseTaskResult() { return {status: "SUCCESS"}; } `) router := gin.New() reachedSubmit := false // The inner plugin router names wildcards by position. router.POST("/vendor/jobs/:s2", pinTaskPluginRoute(plugin, 0), PrepareTaskPluginRoute(), func(c *gin.Context) { reachedSubmit = true requestContext, ok := c.MustGet(jsplugin.ContextKeyRouteRequest).(jsplugin.RouteRequestContext) require.True(t, ok) assert.Equal(t, map[string]any{"prompt": "normalized"}, requestContext.RequestBody) assert.Equal(t, "resolved-model", c.GetString("resolved_task_model")) assert.Equal(t, "resolved-action", c.GetString("task_action")) assert.Equal(t, "route-submit-test", c.GetString("expected_task_plugin_key")) assert.Equal(t, "route-submit-test", c.GetString("task_plugin_key")) assert.Equal(t, "route-submit-test", c.GetString("platform")) c.Status(http.StatusNoContent) }) request := httptest.NewRequest(http.MethodPost, "/vendor/jobs/video?tag=first&tag=second", strings.NewReader(`["prompt",2]`)) request.Header.Set("Content-Type", "application/json") request.Header.Set("X-DashScope-Async", "enable") request.Header.Set("Authorization", "Bearer sk-gateway") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.True(t, reachedSubmit) assert.Equal(t, http.StatusNoContent, recorder.Code) } // Native routes and the generic submit entry hand a plugin declaring // json-order@1 the client's member order and keep its decoded requestBody as // ordered JSON text beside the Go value; other plugins see sorted members. func TestTaskPluginEntriesKeepJSONOrderForDeclaringPlugins(t *testing.T) { source := ` export const meta = { apiVersion: 1, key: "route-order-test", name: "Order", version: "1.0.0", author: {name: "Test"}, models: ["m"], fetchMode: "per_task", requiredCapabilities: ["json-order@1"], routes: [{method: "POST", path: "/vendor/eval", type: "submit", decode: "decodeEval", render: "evalDone"}], }; export const native = { decodeEval: function(ctx) { const state = ctx.body.value.state; return {kind: "submit", model: ctx.body.value.model, requestBody: {keys: Object.keys(state), state: state}}; }, evalDone: function(ctx, task) { return task; }, }; export function buildSubmitRequest() { return {url: "https://example.com"}; } export function parseSubmitResponse() { return {taskId: "one"}; } export function buildQueryRequest() { return {url: "https://example.com"}; } export function parseTaskResult() { return {status: "SUCCESS"}; } ` cases := []struct { name string source string wantText string wantKeys []any }{ {name: "declared", source: source, wantText: `{"keys":["zeta","alpha"],"state":{"zeta":1,"alpha":2}}`, wantKeys: []any{"zeta", "alpha"}}, {name: "undeclared", source: strings.Replace(source, `requiredCapabilities: ["json-order@1"],`, "", 1), wantKeys: []any{"alpha", "zeta"}}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { plugin := compileTaskRoutePlugin(t, tc.source) router := gin.New() reached := false router.POST("/vendor/eval", pinTaskPluginRoute(plugin, 0), PrepareTaskPluginRoute(), func(c *gin.Context) { reached = true request, ok := c.MustGet("task_request").(map[string]any) require.True(t, ok) assert.Equal(t, tc.wantKeys, request["keys"]) text, exists := c.Get(jsplugin.ContextKeyRequestBodyText) if tc.wantText == "" { assert.False(t, exists) } else { assert.Equal(t, json.RawMessage(tc.wantText), text) } c.Status(http.StatusNoContent) }) request := httptest.NewRequest(http.MethodPost, "/vendor/eval", strings.NewReader(`{"model":"m","state":{"zeta":1,"alpha":2}}`)) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.True(t, reached, recorder.Body.String()) }) } t.Run("generic submit entry", func(t *testing.T) { declaring := strings.Replace(genericTaskPluginSource, `fetchMode: "per_task"`, `fetchMode: "per_task", requiredCapabilities: ["json-order@1"]`, 1) _, err := jsplugin.DefaultRegistry.Register(declaring, jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister("generic-entry-test") }) recorder := httptest.NewRecorder() c, _ := gin.CreateTestContext(recorder) c.Params = gin.Params{{Key: "key", Value: "generic-entry-test"}} c.Request = httptest.NewRequest(http.MethodPost, "/v1/tasks/generic-entry-test", strings.NewReader(`{"model":"doc","zeta":1,"alpha":2}`)) c.Request.Header.Set("Content-Type", "application/json") PrepareTaskPluginSubmit()(c) require.False(t, c.IsAborted(), recorder.Body.String()) text, exists := c.Get(jsplugin.ContextKeyRequestBodyText) require.True(t, exists) assert.Equal(t, json.RawMessage(`{"model":"doc","zeta":1,"alpha":2}`), text) }) } func TestPrepareTaskPluginNativeRoutePassesMultipartToDecoder(t *testing.T) { plugin := compileTaskRoutePlugin(t, ` export const meta = { apiVersion: 1, key: "route-multipart-test", name: "Multipart", version: "1.0.0", author: {name: "Test"}, models: ["multipart-model"], fetchMode: "per_task", routes: [ {method: "POST", path: "/vendor/uploads", type: "submit", decode: "decodeUpload", render: "created"}, {method: "POST", path: "/vendor/scoped", type: "submit", decode: "decodeUpload", render: "created", models: ["multipart-model"]}, ], }; export const native = { decodeUpload: function(ctx) { if (ctx.body.kind !== "multipart") throw new Error("multipart body expected"); return {kind: "submit", model: "multipart-model", requestBody: { caption: ctx.body.fields.caption[0], tags: ctx.body.fields.tag.slice(), files: ctx.body.files.map(function(file) { return [file.ref, file.filename, file.size]; }), }}; }, created: function(ctx, task) { return task; }, }; export function buildSubmitRequest() { return {url: "https://example.com"}; } export function parseSubmitResponse() { return {taskId: "one"}; } export function buildQueryRequest() { return {url: "https://example.com"}; } export function parseTaskResult() { return {status: "SUCCESS"}; } `) var body bytes.Buffer writer := multipart.NewWriter(&body) require.NoError(t, writer.WriteField("caption", "hello")) require.NoError(t, writer.WriteField("tag", "one")) require.NoError(t, writer.WriteField("tag", "two")) file, err := writer.CreateFormFile("media", "clip.bin") require.NoError(t, err) _, err = file.Write([]byte("opaque-file")) require.NoError(t, err) require.NoError(t, writer.Close()) router := gin.New() var taskRequest any router.POST("/vendor/uploads", pinTaskPluginRoute(plugin, 0), PrepareTaskPluginRoute(), func(c *gin.Context) { taskRequest, _ = c.Get("task_request") c.Status(http.StatusNoContent) }) request := httptest.NewRequest(http.MethodPost, "/vendor/uploads", bytes.NewReader(body.Bytes())) request.Header.Set("Content-Type", writer.FormDataContentType()) recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) require.Equal(t, http.StatusNoContent, recorder.Code, recorder.Body.String()) assert.Equal(t, map[string]any{ "caption": "hello", "tags": []any{"one", "two"}, "files": []any{[]any{"request_file:media", "clip.bin", int64(11)}}, }, taskRequest) form := httptest.NewRequest(http.MethodPost, "/vendor/uploads", strings.NewReader("caption=hello")) form.Header.Set("Content-Type", "application/x-www-form-urlencoded") recorder = httptest.NewRecorder() router.ServeHTTP(recorder, form) assert.Equal(t, http.StatusUnsupportedMediaType, recorder.Code, "submit routes still refuse url-encoded forms") // A route scoped to models reads the model from a JSON body. router.POST("/vendor/scoped", pinTaskPluginRoute(plugin, 1), PrepareTaskPluginRoute(), func(c *gin.Context) { c.Status(http.StatusNoContent) }) scoped := httptest.NewRequest(http.MethodPost, "/vendor/scoped", bytes.NewReader(body.Bytes())) scoped.Header.Set("Content-Type", writer.FormDataContentType()) recorder = httptest.NewRecorder() router.ServeHTTP(recorder, scoped) assert.Equal(t, http.StatusUnsupportedMediaType, recorder.Code) assert.Contains(t, recorder.Body.String(), "this route requires a JSON body") } func TestPrepareTaskPluginRouteModelScope(t *testing.T) { pluginSource := ` export const meta = { apiVersion: 1, key: "route-model-scope-test", name: "Scoped", version: "1.0.0", author: {name: "Test"}, models: ["gpt-5.5", "gpt-5.6"], fetchMode: "per_task", routes: [{method: "POST", path: "/vendor/batch", type: "submit", models: ["gpt-5.5"], decode: "decodeBatch", render: "batchCreated"}], }; export const native = { decodeBatch: function(ctx) { return {kind: "submit", model: ctx.body.value.model, requestBody: ctx.body.value}; }, batchCreated: function(ctx, task) { return task; }, }; export function buildSubmitRequest() { return {url: "https://example.com"}; } export function parseSubmitResponse() { return {taskId: "one"}; } export function buildQueryRequest() { return {url: "https://example.com"}; } export function parseTaskResult() { return {status: "SUCCESS"}; } ` tests := []struct { name string body string wantStatus int wantResolved bool }{ {name: "listed model passes to decode", body: `{"model":"gpt-5.5","input":"x"}`, wantStatus: http.StatusNoContent, wantResolved: true}, {name: "unlisted model rejected before JS", body: `{"model":"gpt-5.6","input":"x"}`, wantStatus: http.StatusBadRequest}, {name: "missing model rejected before JS", body: `{"input":"x"}`, wantStatus: http.StatusBadRequest}, {name: "non-string model rejected before JS", body: `{"model":7,"input":"x"}`, wantStatus: http.StatusBadRequest}, } for _, testCase := range tests { t.Run(testCase.name, func(t *testing.T) { decodeRan := false source := pluginSource if !testCase.wantResolved { // The core invariant: rejected requests never reach the JS engine. source = strings.Replace(source, `decodeBatch: function(ctx) { return {kind: "submit", model: ctx.body.value.model, requestBody: ctx.body.value}; },`, `decodeBatch: function() { throw new Error("decoder must not run"); },`, 1) } plugin := compileTaskRoutePlugin(t, source) router := gin.New() router.POST("/vendor/batch", pinTaskPluginRoute(plugin, 0), PrepareTaskPluginRoute(), func(c *gin.Context) { decodeRan = true assert.Equal(t, "gpt-5.5", c.GetString("resolved_task_model")) c.Status(http.StatusNoContent) }) request := httptest.NewRequest(http.MethodPost, "/vendor/batch", strings.NewReader(testCase.body)) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.Equal(t, testCase.wantStatus, recorder.Code) assert.Equal(t, testCase.wantResolved, decodeRan) }) } } func TestPrepareTaskPluginRouteRejectsResolvedModelOutsideRouteScope(t *testing.T) { plugin := compileTaskRoutePlugin(t, ` export const meta = { apiVersion: 1, key: "route-resolved-scope-test", name: "Scoped", version: "1.0.0", author: {name: "Test"}, models: ["gpt-5.5", "gpt-5.6"], fetchMode: "per_task", routes: [{method: "POST", path: "/vendor/batch", type: "submit", models: ["gpt-5.5"], decode: "decodeBatch", render: "batchCreated"}], }; export const native = { decodeBatch: function() { return {kind: "submit", model: "gpt-5.6", requestBody: {model: "gpt-5.6"}}; }, batchCreated: function(ctx, task) { return task; }, }; export function buildSubmitRequest() { return {url: "https://example.com"}; } export function parseSubmitResponse() { return {taskId: "one"}; } export function buildQueryRequest() { return {url: "https://example.com"}; } export function parseTaskResult() { return {status: "SUCCESS"}; } `) reached := false router := gin.New() router.POST("/vendor/batch", pinTaskPluginRoute(plugin, 0), PrepareTaskPluginRoute(), func(c *gin.Context) { reached = true c.Status(http.StatusNoContent) }) request := httptest.NewRequest(http.MethodPost, "/vendor/batch", strings.NewReader(`{"model":"gpt-5.5"}`)) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.Equal(t, http.StatusBadRequest, recorder.Code) assert.False(t, reached, "decode may run, but an out-of-scope resolved model must not continue into submit") } func TestPrepareTaskPluginRouteWithoutModelScopeIsUnrestricted(t *testing.T) { plugin := compileTaskRoutePlugin(t, ` export const meta = { apiVersion: 1, key: "route-unscoped-test", name: "Unscoped", version: "1.0.0", author: {name: "Test"}, models: ["gpt-5.5", "gpt-5.6"], fetchMode: "per_task", routes: [{method: "POST", path: "/vendor/batch", type: "submit", decode: "decodeBatch", render: "batchCreated"}], }; export const native = { decodeBatch: function(ctx) { return {kind: "submit", model: ctx.body.value.model, requestBody: ctx.body.value}; }, batchCreated: function(ctx, task) { return task; }, }; export function buildSubmitRequest() { return {url: "https://example.com"}; } export function parseSubmitResponse() { return {taskId: "one"}; } export function buildQueryRequest() { return {url: "https://example.com"}; } export function parseTaskResult() { return {status: "SUCCESS"}; } `) router := gin.New() router.POST("/vendor/batch", pinTaskPluginRoute(plugin, 0), PrepareTaskPluginRoute(), func(c *gin.Context) { c.Status(http.StatusNoContent) }) request := httptest.NewRequest(http.MethodPost, "/vendor/batch", strings.NewReader(`{"model":"gpt-5.6"}`)) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.Equal(t, http.StatusNoContent, recorder.Code) } func TestBuildTaskPluginRouteRequestBodyUnion(t *testing.T) { tests := []struct { name string contentType string body string kind jsplugin.BodyKind assertBody func(*testing.T, map[string]any) }{ {name: "none", kind: jsplugin.BodyNone}, {name: "json", contentType: "application/problem+json; charset=utf-8", body: `{"model":"m"}`, kind: jsplugin.BodyJSON, assertBody: func(t *testing.T, body map[string]any) { assert.Equal(t, map[string]any{"model": "m"}, body["value"]) }}, {name: "form preserves repeated values", contentType: "application/x-www-form-urlencoded", body: "tag=one&tag=two", kind: jsplugin.BodyForm, assertBody: func(t *testing.T, body map[string]any) { assert.Equal(t, []string{"one", "two"}, body["fields"].(map[string][]string)["tag"]) }}, } for _, testCase := range tests { t.Run(testCase.name, func(t *testing.T) { c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/body", strings.NewReader(testCase.body)) if testCase.contentType != "" { c.Request.Header.Set("Content-Type", testCase.contentType) } requestContext, err := buildTaskPluginRouteRequest(c) require.NoError(t, err) decodedBody := requestContext.Body.(map[string]any) assert.Equal(t, string(testCase.kind), decodedBody["kind"]) if testCase.assertBody != nil { testCase.assertBody(t, decodedBody) } }) } var multipartBody bytes.Buffer writer := multipart.NewWriter(&multipartBody) require.NoError(t, writer.WriteField("tag", "one")) require.NoError(t, writer.WriteField("tag", "two")) file, err := writer.CreateFormFile("input", "image.png") require.NoError(t, err) _, err = file.Write([]byte("file bytes stay in Go")) require.NoError(t, err) require.NoError(t, writer.Close()) c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/body", bytes.NewReader(multipartBody.Bytes())) c.Request.Header.Set("Content-Type", writer.FormDataContentType()) requestContext, err := buildTaskPluginRouteRequest(c) require.NoError(t, err) decodedBody := requestContext.Body.(map[string]any) assert.Equal(t, string(jsplugin.BodyMultipart), decodedBody["kind"]) assert.Equal(t, []string{"one", "two"}, decodedBody["fields"].(map[string][]string)["tag"]) files := decodedBody["files"].([]map[string]any) require.Len(t, files, 1) assert.Equal(t, "image.png", files[0]["filename"]) assert.NotContains(t, fmt.Sprint(files[0]), "file bytes stay in Go") } func TestBuildTaskPluginRouteRequestRejectsUnsafeBodies(t *testing.T) { tests := []struct { name string contentType string body []byte errorText string }{ {name: "invalid json UTF-8", contentType: "application/json", body: []byte{'{', '"', 'x', '"', ':', '"', 0xff, '"', '}'}, errorText: "valid UTF-8"}, {name: "invalid form UTF-8", contentType: "application/x-www-form-urlencoded", body: []byte("x=%FF"), errorText: "valid UTF-8"}, {name: "too many repeated form fields", contentType: "application/x-www-form-urlencoded", body: []byte(strings.Repeat("x=v&", maxTaskPluginFormFields) + "x=v"), errorText: "exceeds 256 fields"}, {name: "missing multipart boundary", contentType: "multipart/form-data", body: []byte("body"), errorText: "boundary is required"}, } for _, testCase := range tests { t.Run(testCase.name, func(t *testing.T) { c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/body", bytes.NewReader(testCase.body)) c.Request.Header.Set("Content-Type", testCase.contentType) _, err := buildTaskPluginRouteRequest(c) require.ErrorContains(t, err, testCase.errorText) }) } t.Run("conflicting Content-Type", func(t *testing.T) { c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/body", strings.NewReader(`{}`)) c.Request.Header["Content-Type"] = []string{"application/json", "application/x-www-form-urlencoded"} _, err := buildTaskPluginRouteRequest(c) require.ErrorContains(t, err, "conflicting Content-Type") }) t.Run("oversized total body", func(t *testing.T) { previous := constant.MaxRequestBodyMB constant.MaxRequestBodyMB = 1 t.Cleanup(func() { constant.MaxRequestBodyMB = previous }) c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/body", strings.NewReader(strings.Repeat(" ", (1<<20)+1))) c.Request.Header.Set("Content-Type", "application/json") _, err := buildTaskPluginRouteRequest(c) require.ErrorContains(t, err, "request body exceeds 1 MB") }) multipartCases := []struct { name string build func(*testing.T, *multipart.Writer) errorText string }{ {name: "too many parts", build: func(t *testing.T, writer *multipart.Writer) { for i := 0; i <= maxTaskPluginMultipartParts; i++ { require.NoError(t, writer.WriteField("field", "value")) } }, errorText: "exceeds 256 parts"}, {name: "too many files", build: func(t *testing.T, writer *multipart.Writer) { for i := 0; i <= maxTaskPluginFiles; i++ { _, err := writer.CreateFormFile("file", fmt.Sprintf("%d.bin", i)) require.NoError(t, err) } }, errorText: "exceeds 32 files"}, {name: "invalid UTF-8 field", build: func(t *testing.T, writer *multipart.Writer) { part, err := writer.CreateFormField("field") require.NoError(t, err) _, err = part.Write([]byte{0xff}) require.NoError(t, err) }, errorText: "valid UTF-8"}, {name: "nested multipart", build: func(t *testing.T, writer *multipart.Writer) { header := make(textproto.MIMEHeader) header.Set("Content-Disposition", `form-data; name="nested"`) header.Set("Content-Type", "multipart/mixed; boundary=inner") _, err := writer.CreatePart(header) require.NoError(t, err) }, errorText: "nested multipart"}, {name: "invalid UTF-8 filename", build: func(t *testing.T, writer *multipart.Writer) { header := make(textproto.MIMEHeader) header.Set("Content-Disposition", "form-data; name=\"file\"; filename=\""+string([]byte{0xff})+"\"") _, err := writer.CreatePart(header) require.NoError(t, err) }, errorText: "invalid multipart"}, {name: "injected disposition name", build: func(t *testing.T, writer *multipart.Writer) { header := make(textproto.MIMEHeader) header.Set("Content-Disposition", "form-data; name=\"prompt"+"\r\n"+"X-Injected: yes\"") part, err := writer.CreatePart(header) require.NoError(t, err) _, err = part.Write([]byte("hello")) require.NoError(t, err) }, errorText: "invalid multipart field name"}, {name: "injected disposition filename", build: func(t *testing.T, writer *multipart.Writer) { header := make(textproto.MIMEHeader) header.Set("Content-Disposition", "form-data; name=\"file\"; filename=\"safe.png"+"\r\n"+"X-Injected: yes\"") part, err := writer.CreatePart(header) require.NoError(t, err) _, err = part.Write([]byte("file")) require.NoError(t, err) }, errorText: "invalid multipart"}, } for _, testCase := range multipartCases { t.Run(testCase.name, func(t *testing.T) { var body bytes.Buffer writer := multipart.NewWriter(&body) testCase.build(t, writer) require.NoError(t, writer.Close()) c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/body", bytes.NewReader(body.Bytes())) c.Request.Header.Set("Content-Type", writer.FormDataContentType()) _, err := buildTaskPluginRouteRequest(c) require.ErrorContains(t, err, testCase.errorText) }) } t.Run("multipart fields are bounded by the request body limit only", func(t *testing.T) { previous := constant.MaxRequestBodyMB t.Cleanup(func() { constant.MaxRequestBodyMB = previous }) var body bytes.Buffer writer := multipart.NewWriter(&body) // A JSON request sent as a field, larger than the old 1 MiB field cap. require.NoError(t, writer.WriteField("request", strings.Repeat("a", 2<<20))) require.NoError(t, writer.Close()) send := func(limitMB int) (map[string]any, error) { constant.MaxRequestBodyMB = limitMB c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/body", bytes.NewReader(body.Bytes())) c.Request.Header.Set("Content-Type", writer.FormDataContentType()) request, err := buildTaskPluginRouteRequest(c) decoded, _ := request.Body.(map[string]any) return decoded, err } decoded, err := send(4) require.NoError(t, err) assert.Len(t, decoded["fields"].(map[string][]string)["request"][0], 2<<20) _, err = send(1) require.ErrorContains(t, err, "request body exceeds 1 MB") }) t.Run("oversized multipart file", func(t *testing.T) { previous := constant.MaxFileDownloadMB constant.MaxFileDownloadMB = 1 t.Cleanup(func() { constant.MaxFileDownloadMB = previous }) var body bytes.Buffer writer := multipart.NewWriter(&body) file, err := writer.CreateFormFile("file", "large.bin") require.NoError(t, err) _, err = file.Write(bytes.Repeat([]byte{'x'}, (1<<20)+1)) require.NoError(t, err) require.NoError(t, writer.Close()) c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/body", bytes.NewReader(body.Bytes())) c.Request.Header.Set("Content-Type", writer.FormDataContentType()) _, err = buildTaskPluginRouteRequest(c) require.ErrorContains(t, err, "multipart file exceeds 1 MB") }) } func TestPrepareTaskPluginEndpointPinsGenerationBeforeParseAndDistribution(t *testing.T) { const key = "endpoint-pin-test" firstSource := taskProtocolPluginSource( key, "1.0.0", `["claimed-model"]`, "/v1/responses", `if (ctx.protocol !== "openai_responses" || ctx.operation !== "create" || ctx.model !== "claimed-model" || ctx.path !== "/v1/responses" || ctx.stream !== false) throw new Error("bad context"); if (ctx.query.trace[0] !== "one" || ctx.requestBody.prompt !== "hello") throw new Error("bad request"); ctx.requestBody.prompt = "plugin-local-mutation"; return {model: ctx.model, action: "first-action", requestBody: {prompt: "normalized"}};`, ) first, err := jsplugin.DefaultRegistry.Register(firstSource, jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(key)) }) router := gin.New() reachedDistribution := false router.POST( "/v1/responses", PinTaskPluginEndpoint(), func(c *gin.Context) { _, updateErr := jsplugin.DefaultRegistry.Register(taskProtocolPluginSource( key, "2.0.0", `["claimed-model"]`, "/v1/responses", `return {model: ctx.model, action: "second-action"};`, ), jsplugin.Options{}) require.NoError(t, updateErr) c.Next() }, PrepareTaskPluginEndpoint(), func(c *gin.Context) { reachedDistribution = true pinned := c.MustGet(jsplugin.ContextKeyPinnedEndpoint).(jsplugin.PinnedEndpoint) assert.Same(t, first, pinned.Plugin) assert.Equal(t, "claimed-model", pinned.Model) assert.Equal(t, "first-action", c.GetString("task_action")) assert.Equal(t, "claimed-model", c.GetString("resolved_task_model")) assert.Equal(t, key, c.GetString("expected_task_plugin_key")) assert.Equal(t, key, c.GetString("platform")) protocolRequest := c.MustGet(jsplugin.ContextKeyProtocolRequest).(jsplugin.ProtocolRequestContext) assert.Equal(t, map[string]any{"kind": "json", "value": map[string]any{"model": "claimed-model", "prompt": "hello"}}, protocolRequest.Body) assert.False(t, protocolRequest.Stream) routeRequest := c.MustGet(jsplugin.ContextKeyRouteRequest).(jsplugin.RouteRequestContext) assert.Equal(t, map[string]any{"prompt": "normalized"}, routeRequest.RequestBody) current, found := jsplugin.DefaultRegistry.Get(key) require.True(t, found) assert.NotSame(t, first, current) c.Status(http.StatusNoContent) }, ) request := httptest.NewRequest( http.MethodPost, "/v1/responses?trace=one", strings.NewReader(`{"model":"claimed-model","prompt":"hello"}`), ) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.True(t, reachedDistribution) assert.Equal(t, http.StatusNoContent, recorder.Code) } func TestPrepareTaskPluginEndpointClientDisconnectDoesNotCancelParseHook(t *testing.T) { const key = "endpoint-detached-parse-test" _, err := jsplugin.DefaultRegistry.Register(taskProtocolPluginSource( key, "1.0.0", `["claimed-model"]`, "/v1/responses", `return {model: "claimed-model"};`, ), jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(key)) }) router := gin.New() reachedDistribution := false router.POST( "/v1/responses", PinTaskPluginEndpoint(), func(c *gin.Context) { requestContext, cancel := context.WithCancel(c.Request.Context()) cancel() c.Request = c.Request.WithContext(requestContext) c.Next() }, PrepareTaskPluginEndpoint(), func(c *gin.Context) { reachedDistribution = true assert.Equal(t, "claimed-model", c.GetString("resolved_task_model")) c.Status(http.StatusNoContent) }, ) request := httptest.NewRequest( http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"claimed-model"}`), ) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.True(t, reachedDistribution) assert.Equal(t, http.StatusNoContent, recorder.Code) } func TestPrepareTaskPluginEndpointUsesStrictOriginalStreamFlag(t *testing.T) { const key = "endpoint-stream-test" _, err := jsplugin.DefaultRegistry.Register(taskProtocolPluginSource( key, "1.0.0", `["claimed-model"]`, "/v1/responses", `if (ctx.stream !== true) throw new Error("stream mode was not preserved"); return {model: "claimed-model", requestBody: {stream: false}};`, ), jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(key)) }) router := gin.New() reachedDistribution := false router.POST( "/v1/responses", PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) { reachedDistribution = true protocolRequest := c.MustGet(jsplugin.ContextKeyProtocolRequest).(jsplugin.ProtocolRequestContext) assert.True(t, protocolRequest.Stream) normalizedRequest := c.MustGet(jsplugin.ContextKeyRouteRequest).(jsplugin.RouteRequestContext) assert.Equal(t, map[string]any{"stream": false}, normalizedRequest.RequestBody) c.Status(http.StatusNoContent) }, ) request := httptest.NewRequest( http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"claimed-model","stream":true}`), ) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.True(t, reachedDistribution) assert.Equal(t, http.StatusNoContent, recorder.Code) for _, invalid := range []string{`null`, `"true"`, `1`, `"yes"`} { request = httptest.NewRequest( http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"claimed-model","stream":`+invalid+`}`), ) request.Header.Set("Content-Type", "application/json") recorder = httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.Equal(t, http.StatusBadRequest, recorder.Code) } } func TestTaskPluginEndpointMissPreservesOrdinaryRequestBody(t *testing.T) { const key = "endpoint-miss-test" _, err := jsplugin.DefaultRegistry.Register(taskProtocolPluginSource( key, "1.0.0", `["claimed-model"]`, "/v1/responses", `return {model: "claimed-model"};`, ), jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(key)) }) router := gin.New() router.POST( "/v1/responses", PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) { _, pinned := c.Get(jsplugin.ContextKeyPinnedEndpoint) assert.False(t, pinned) storage, storageErr := common.GetBodyStorage(c) require.NoError(t, storageErr) raw, bytesErr := storage.Bytes() require.NoError(t, bytesErr) assert.Equal(t, []byte(`{"model":"ordinary-model","input":"hello"}`), raw) c.Status(http.StatusNoContent) }, ) request := httptest.NewRequest( http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"ordinary-model","input":"hello"}`), ) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.Equal(t, http.StatusNoContent, recorder.Code) } func TestTaskPluginEndpointUsesOneCanonicalModelForDuplicateJSONKeys(t *testing.T) { require.NoError(t, appI18n.Init()) const key = "endpoint-duplicate-model-test" _, err := jsplugin.DefaultRegistry.Register(taskProtocolPluginSource( key, "1.0.0", `["claimed-model"]`, "/v1/responses", `if (ctx.requestBody.model !== "claimed-model") throw new Error("noncanonical model"); return {model: ctx.requestBody.model};`, ), jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(key)) }) tests := []struct { name string body string }{ { name: "claimed model first", body: `{"model":"claimed-model","model":"ordinary-model"}`, }, { name: "claimed model second", body: `{"model":"ordinary-model","model":"claimed-model"}`, }, } for _, testCase := range tests { t.Run(testCase.name, func(t *testing.T) { reachedDownstream := false router := gin.New() router.POST( "/v1/responses", PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) { reachedDownstream = true c.Status(http.StatusNoContent) }, ) request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(testCase.body)) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.False(t, reachedDownstream) assert.Equal(t, http.StatusBadRequest, recorder.Code) }) } } func TestRequestBodyFailuresReportTheirCause(t *testing.T) { require.NoError(t, appI18n.Init()) require.NotNil(t, jsplugin.DefaultRegistry.Generation(), "the built-in plugins keep a routing generation, as in production") _, err := jsplugin.DefaultRegistry.Register(genericTaskPluginSource, jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister("generic-entry-test") }) routePlugin := compileTaskRoutePlugin(t, genericTaskPluginSource) previous := constant.MaxRequestBodyMB constant.MaxRequestBodyMB = 1 t.Cleanup(func() { constant.MaxRequestBodyMB = previous }) reached := func(c *gin.Context) { c.Status(http.StatusNoContent) } router := gin.New() router.POST("/v1/responses", PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), Distribute(), reached) router.POST("/v1/video/generations", PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), Distribute(), reached) router.POST("/v1/chat/completions", Distribute(), reached) router.POST("/mj/submit/imagine", Distribute(), reached) router.POST("/v1/tasks/:key", PrepareTaskPluginSubmit(), reached) router.POST("/vendor/jobs", func(c *gin.Context) { c.Set(jsplugin.ContextKeyPinnedRoute, jsplugin.PinnedRoute{Plugin: routePlugin}) }, PrepareTaskPluginRoute(), reached) oversized := `{"model":"gpt-x","prompt":"` + strings.Repeat("a", 1<<20) + `"}` tests := []struct { name string path string body io.Reader status int message string }{ { name: "oversized body on a plugin protocol path", path: "/v1/responses", body: strings.NewReader(oversized), status: http.StatusRequestEntityTooLarge, message: "Invalid request: request body exceeds 1 MB: request body too large", }, { name: "body cut off on a plugin protocol path", path: "/v1/responses", body: io.MultiReader(strings.NewReader(`{"model":"gpt-x","input":"`), iotest.ErrReader(io.ErrUnexpectedEOF)), status: http.StatusBadRequest, message: "Invalid request: unexpected EOF", }, { // The endpoint check lets this path through and the distributor reads the body again. name: "oversized body on the legacy video path", path: "/v1/video/generations", body: strings.NewReader(oversized), status: http.StatusRequestEntityTooLarge, message: "Invalid request: request body exceeds 1 MB: request body too large", }, { name: "oversized Midjourney body", path: "/mj/submit/imagine", body: strings.NewReader(oversized), status: http.StatusRequestEntityTooLarge, message: "Invalid request: request body exceeds 1 MB: request body too large", }, { name: "oversized task plugin submission", path: "/v1/tasks/generic-entry-test", body: strings.NewReader(oversized), status: http.StatusRequestEntityTooLarge, message: "request body exceeds 1 MB: request body too large", }, { name: "oversized plugin route body", path: "/vendor/jobs", body: strings.NewReader(oversized), status: http.StatusRequestEntityTooLarge, message: "request body exceeds 1 MB: request body too large", }, { name: "invalid JSON on an ordinary path", path: "/v1/chat/completions", body: strings.NewReader(`{"model":"gpt-x",`), status: http.StatusBadRequest, message: "Invalid request: invalid JSON request body", }, } for _, testCase := range tests { t.Run(testCase.name, func(t *testing.T) { request := httptest.NewRequest(http.MethodPost, testCase.path, testCase.body) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.Equal(t, testCase.status, recorder.Code) assert.Contains(t, recorder.Body.String(), `"message":"`+testCase.message) }) } } func TestTaskPluginEndpointOnlyPreservesConditionalMiddlewareSemantics(t *testing.T) { tests := []struct { name string claimed bool wrappedAborts bool expectedWrapped int expectedDownstream bool expectedStatus int }{ { name: "unclaimed skips wrapper", expectedDownstream: true, expectedStatus: http.StatusNoContent, }, { name: "claimed invokes wrapper", claimed: true, expectedWrapped: 1, expectedDownstream: true, expectedStatus: http.StatusNoContent, }, { name: "claimed wrapper aborts", claimed: true, wrappedAborts: true, expectedWrapped: 1, expectedStatus: http.StatusTooManyRequests, }, } for _, testCase := range tests { t.Run(testCase.name, func(t *testing.T) { wrappedCalls := 0 reachedDownstream := false router := gin.New() router.POST( "/v1/videos", func(c *gin.Context) { if testCase.claimed { c.Set(jsplugin.ContextKeyPinnedEndpoint, jsplugin.PinnedEndpoint{}) } c.Next() }, TaskPluginEndpointOnly(func(c *gin.Context) { wrappedCalls++ if testCase.wrappedAborts { c.AbortWithStatus(http.StatusTooManyRequests) return } c.Next() }), func(c *gin.Context) { reachedDownstream = true c.Status(http.StatusNoContent) }, ) recorder := httptest.NewRecorder() router.ServeHTTP(recorder, httptest.NewRequest(http.MethodPost, "/v1/videos", nil)) assert.Equal(t, testCase.expectedWrapped, wrappedCalls) assert.Equal(t, testCase.expectedDownstream, reachedDownstream) assert.Equal(t, testCase.expectedStatus, recorder.Code) }) } } func TestPrepareTaskPluginEndpointRejectsModelDriftBeforeDistribution(t *testing.T) { const key = "endpoint-drift-test" _, err := jsplugin.DefaultRegistry.Register(taskProtocolPluginSource( key, "1.0.0", `["claimed-model"]`, "/v1/responses", `return {model: "outside-model"};`, ), jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(key)) }) reachedDistribution := false router := gin.New() router.POST( "/v1/responses", PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) { reachedDistribution = true c.Status(http.StatusNoContent) }, ) request := httptest.NewRequest( http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"claimed-model"}`), ) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.False(t, reachedDistribution) assert.Equal(t, http.StatusBadRequest, recorder.Code) var payload map[string]any require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &payload)) errObj, _ := payload["error"].(map[string]any) require.NotNil(t, errObj) assert.Contains(t, fmt.Sprint(errObj["message"]), `model "outside-model" is not served by this plugin`) } func TestPrepareTaskPluginEndpointAcceptsRegisteredVideoMultipartBody(t *testing.T) { const key = "endpoint-multipart-test" _, err := jsplugin.DefaultRegistry.Register(taskProtocolPluginSource( key, "1.0.0", `["video-model"]`, "/v1/videos", `if (ctx.body.kind !== "multipart" || ctx.body.fields.prompt[0] !== "hello") throw new Error("bad prompt"); if (ctx.body.files.length !== 1 || ctx.body.files[0].field !== "input_reference") throw new Error("bad file ref"); return {model: ctx.model, requestBody: {prompt: ctx.body.fields.prompt[0]}};`, ), jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(key)) }) var body bytes.Buffer writer := multipart.NewWriter(&body) require.NoError(t, writer.WriteField("model", "video-model")) require.NoError(t, writer.WriteField("prompt", "hello")) file, err := writer.CreateFormFile("input_reference", "reference.bin") require.NoError(t, err) _, err = file.Write([]byte("opaque-video-reference")) require.NoError(t, err) require.NoError(t, writer.Close()) router := gin.New() reachedDistribution := false router.POST( "/v1/videos", PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) { reachedDistribution = true c.Status(http.StatusNoContent) }, ) request := httptest.NewRequest(http.MethodPost, "/v1/videos", bytes.NewReader(body.Bytes())) request.Header.Set("Content-Type", writer.FormDataContentType()) recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.True(t, reachedDistribution) assert.Equal(t, http.StatusNoContent, recorder.Code) } // The OpenAI Images edits endpoint accepts multipart uploads. Every file of a // repeated image[] field is exposed to the decoder with its own ref, and an // unclaimed model on the shared endpoint still reaches the ordinary relay. func TestPrepareTaskPluginEndpointExposesIndexedRefsForRepeatedImageFiles(t *testing.T) { const key = "endpoint-image-edits-test" _, err := jsplugin.DefaultRegistry.Register(taskProtocolPluginSource( key, "1.0.0", `["image-edit-model"]`, "/v1/images/edits", `if (ctx.body.kind !== "multipart" || ctx.body.fields.prompt[0] !== "watercolor") throw new Error("bad prompt"); return {model: ctx.model, action: "image_to_image", requestBody: {prompt: ctx.body.fields.prompt[0], refs: ctx.body.files.map(function(file) { return file.ref; })}};`, ), jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(key)) }) var body bytes.Buffer writer := multipart.NewWriter(&body) require.NoError(t, writer.WriteField("model", "image-edit-model")) require.NoError(t, writer.WriteField("prompt", "watercolor")) for _, name := range []string{"first.png", "second.png"} { file, createErr := writer.CreateFormFile("image[]", name) require.NoError(t, createErr) _, err = file.Write([]byte(name)) require.NoError(t, err) } require.NoError(t, writer.Close()) router := gin.New() var pinnedProtocol string var taskRequest any router.POST("/v1/images/edits", PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) { if pinned, ok := c.MustGet(jsplugin.ContextKeyPinnedEndpoint).(jsplugin.PinnedEndpoint); ok { pinnedProtocol = pinned.Protocol } taskRequest, _ = c.Get("task_request") c.Status(http.StatusNoContent) }) request := httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(body.Bytes())) request.Header.Set("Content-Type", writer.FormDataContentType()) recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) require.Equal(t, http.StatusNoContent, recorder.Code, recorder.Body.String()) assert.Equal(t, jsplugin.ProtocolOpenAIImage, pinnedProtocol) decoded, ok := taskRequest.(map[string]any) require.True(t, ok) assert.Equal(t, []any{"request_file:image[]", "request_file:image[]#1"}, decoded["refs"]) unclaimed := gin.New() reachedRelay := false unclaimed.POST("/v1/images/generations", PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) { _, pinned := c.Get(jsplugin.ContextKeyPinnedEndpoint) reachedRelay = !pinned c.Status(http.StatusNoContent) }) plain := httptest.NewRequest(http.MethodPost, "/v1/images/generations", strings.NewReader(`{"model":"dall-e-3","prompt":"a cat"}`)) plain.Header.Set("Content-Type", "application/json") recorder = httptest.NewRecorder() unclaimed.ServeHTTP(recorder, plain) assert.Equal(t, http.StatusNoContent, recorder.Code) assert.True(t, reachedRelay, "an unclaimed image model keeps the built-in relay") } func TestVideoGenerationsIsNotClaimedByOpenAIVideoProtocol(t *testing.T) { const key = "endpoint-video-gen-test" _, err := jsplugin.DefaultRegistry.Register(taskProtocolPluginSource( key, "1.0.0", `["generation-model"]`, "/v1/video/generations", `return {model: ctx.requestBody.model, action: "generate"};`, ), jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(key)) }) router := gin.New() reachedDistribution := false router.POST( "/v1/video/generations", PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) { reachedDistribution = true c.Status(http.StatusNoContent) }, ) request := httptest.NewRequest( http.MethodPost, "/v1/video/generations", strings.NewReader(`{"model":"generation-model"}`), ) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.True(t, reachedDistribution) assert.Equal(t, http.StatusNoContent, recorder.Code) } func TestPrepareTaskPluginRouteKeepsRawBinaryOpaque(t *testing.T) { plugin := compileTaskRoutePlugin(t, ` export const meta = { apiVersion: 1, key: "route-binary-test", name: "Binary", version: "1.0.0", author: {name: "Test"}, models: ["binary-model"], fetchMode: "per_task", routes: [{method: "POST", path: "/vendor/binary", type: "submit", decode: "decodeBinary", render: "created"}], }; export const native = {decodeBinary: function() { throw new Error("decoder must not run"); }, created: function(ctx, task) { return task; }}; export function buildSubmitRequest() { return {url: "https://example.com"}; } export function parseSubmitResponse() { return {taskId: "one"}; } export function buildQueryRequest() { return {url: "https://example.com"}; } export function parseTaskResult() { return {status: "SUCCESS"}; } `) router := gin.New() router.POST("/vendor/binary", pinTaskPluginRoute(plugin, 0), PrepareTaskPluginRoute(), func(c *gin.Context) { body, err := io.ReadAll(c.Request.Body) require.NoError(t, err) assert.Equal(t, []byte{0, 1, 2, 3}, body) c.Status(http.StatusNoContent) }) request := httptest.NewRequest(http.MethodPost, "/vendor/binary", bytes.NewReader([]byte{0, 1, 2, 3})) request.Header.Set("Content-Type", "application/octet-stream") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.Equal(t, http.StatusUnsupportedMediaType, recorder.Code) } func TestPrepareTaskPluginDynamicQueryPreservesOrderAndSkipsNextHandlers(t *testing.T) { setupTaskPluginRouteDB(t) plugin := compileTaskRoutePlugin(t, ` export const meta = { apiVersion: 1, key: "route-query-test", name: "Query", version: "1.0.0", author: {name: "Test"}, models: ["query-model"], fetchMode: "per_task", routes: [{method: "POST", path: "/vendor/query", type: "dynamic", decode: "decodeQuery", render: "renderQuery"}], }; export const native = { decodeQuery: function() { return {kind: "query", taskIds: ["task-b", "task-a", "task-b"]}; }, renderQuery: function(ctx, tasks) { return {ids: tasks.map(function(task) { return task.task_id; }), keys: Object.keys(tasks[0]).sort(), data: tasks[0].data}; }, }; export function buildSubmitRequest() { return {url: "https://example.com"}; } export function parseSubmitResponse() { return {taskId: "one"}; } export function buildQueryRequest() { return {url: "https://example.com"}; } export function parseTaskResult() { return {status: "SUCCESS"}; } `) insertTaskPluginRouteTask(t, &model.Task{ TaskID: "task-a", UserId: 7, Platform: constant.TaskPlatform("route-query-test"), Status: model.TaskStatusSuccess, Progress: "100%", CreatedAt: 10, }) taskData, err := common.Marshal(map[string]any{ "task_id": "private-upstream-id", "nested": map[string]any{"url": "https://upstream.invalid/tasks/private-upstream-id"}, }) require.NoError(t, err) insertTaskPluginRouteTask(t, &model.Task{ TaskID: "task-b", UserId: 7, Platform: constant.TaskPlatform("route-query-test"), Status: model.TaskStatusInProgress, Progress: "50%", CreatedAt: 20, Data: taskData, ChannelId: 999, Quota: 12345, PrivateData: model.TaskPrivateData{UpstreamTaskID: "private-upstream-id", Key: "secret"}, }) nextHandlerCalled := false router := gin.New() router.POST("/vendor/query", pinTaskPluginRoute(plugin, 0), func(c *gin.Context) { c.Set("id", 7) c.Next() }, PrepareTaskPluginRoute(), func(c *gin.Context) { nextHandlerCalled = true c.Status(http.StatusTeapot) }, ) recorder := httptest.NewRecorder() request := httptest.NewRequest(http.MethodPost, "/vendor/query", strings.NewReader(`{}`)) request.Header.Set("Content-Type", "application/json") router.ServeHTTP(recorder, request) assert.Equal(t, http.StatusOK, recorder.Code) assert.False(t, nextHandlerCalled) var response map[string]any require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) assert.Equal(t, []any{"task-b", "task-a", "task-b"}, response["ids"]) data, ok := response["data"].(map[string]any) require.True(t, ok) assert.Equal(t, "task-b", data["task_id"]) nested, ok := data["nested"].(map[string]any) require.True(t, ok) assert.Equal(t, "https://upstream.invalid/tasks/private-upstream-id", nested["url"]) keys, ok := response["keys"].([]any) require.True(t, ok) assert.NotContains(t, keys, "user_id") assert.NotContains(t, keys, "channel_id") assert.NotContains(t, keys, "quota") assert.NotContains(t, keys, "private_data") assert.NotContains(t, recorder.Body.String(), "secret") } func TestPrepareTaskPluginDynamicDecoderRejectsRendererField(t *testing.T) { plugin := compileTaskRoutePlugin(t, ` export const meta = {apiVersion:1,key:"dynamic-renderer",name:"Dynamic",version:"1.0.0",author:{name:"Test"},models:["model"],fetchMode:"per_task",routes:[{method:"POST",path:"/vendor/query",type:"dynamic",decode:"decode",render:"show"}]}; export const native = {decode:function(){return {kind:"query",taskIds:[],renderer:"legacy"};},show:function(){return {};}}; export function buildSubmitRequest(){return {}} export function parseSubmitResponse(){return {taskId:"one"}} export function buildQueryRequest(){return {}} export function parseTaskResult(){return {status:"SUCCESS"}} `) router := gin.New() reached := false router.POST("/vendor/query", pinTaskPluginRoute(plugin, 0), PrepareTaskPluginRoute(), func(c *gin.Context) { reached = true }) request := httptest.NewRequest(http.MethodPost, "/vendor/query", strings.NewReader(`{}`)) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.False(t, reached) assert.Equal(t, http.StatusBadRequest, recorder.Code) } func TestPrepareTaskPluginSubmitDecoderRejectsRendererField(t *testing.T) { plugin := compileTaskRoutePlugin(t, ` export const meta = {apiVersion:1,key:"submit-renderer",name:"Submit",version:"1.0.0",author:{name:"Test"},models:["model"],fetchMode:"per_task",routes:[{method:"POST",path:"/vendor/submit",type:"submit",decode:"decode",render:"created"}]}; export const native = {decode:function(){return {kind:"submit",model:"model",requestBody:{},renderer:"legacy"};},created:function(){return {};}}; export function buildSubmitRequest(){return {}} export function parseSubmitResponse(){return {taskId:"one"}} export function buildQueryRequest(){return {}} export function parseTaskResult(){return {status:"SUCCESS"}} `) router := gin.New() reached := false router.POST("/vendor/submit", pinTaskPluginRoute(plugin, 0), PrepareTaskPluginRoute(), func(c *gin.Context) { reached = true }) request := httptest.NewRequest(http.MethodPost, "/vendor/submit", strings.NewReader(`{}`)) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.False(t, reached) assert.Equal(t, http.StatusBadRequest, recorder.Code) } func TestPrepareTaskPluginStaticQueryHidesTaskExistenceAndSanitizesErrors(t *testing.T) { setupTaskPluginRouteDB(t) plugin := compileTaskRoutePlugin(t, ` export const meta = { apiVersion: 1, key: "route-static-query-test", name: "Static Query", version: "1.0.0", author: {name: "Test"}, channelTypes: [651], models: ["query-model"], fetchMode: "per_task", routes: [{method: "GET", path: "/vendor/jobs/:id", type: "query", taskIdParam: "id", render: "status"}], }; export function buildSubmitRequest() { return {url: "https://example.com"}; } export function parseSubmitResponse() { return {taskId: "one"}; } export function buildQueryRequest() { return {url: "https://example.com"}; } export function parseTaskResult() { return {status: "SUCCESS"}; } export const native = {status: function(ctx, task) { return {id: task.task_id}; }, error: function(ctx, error) { return {vendor_error: {code: error.code, message: error.message, status: error.httpStatus, retryable: error.retryable}}; }}; `) insertTaskPluginRouteTask(t, &model.Task{ TaskID: "foreign-task", UserId: 8, Platform: constant.TaskPlatform("route-static-query-test"), }) insertTaskPluginRouteTask(t, &model.Task{ TaskID: "wrong-platform", UserId: 7, Platform: constant.TaskPlatform("another-plugin"), }) insertTaskPluginRouteTask(t, &model.Task{ TaskID: "wrong-legacy-platform", UserId: 7, Platform: constant.TaskPlatform("652"), }) insertTaskPluginRouteTask(t, &model.Task{ TaskID: "legacy-task", UserId: 7, Platform: constant.TaskPlatform("651"), }) // Submitted through a route with retainResult: false; owned by the caller // but not retrievable, so it must be indistinguishable from a missing task. insertTaskPluginRouteTask(t, &model.Task{ TaskID: "discarded-task", UserId: 7, Platform: constant.TaskPlatform("route-static-query-test"), Status: model.TaskStatusSuccess, PrivateData: model.TaskPrivateData{ResultDiscarded: true}, }) router := gin.New() router.GET("/vendor/jobs/:s2", pinTaskPluginRoute(plugin, 0), func(c *gin.Context) { c.Set("id", 7) c.Next() }, PrepareTaskPluginRoute(), ) legacyRecorder := httptest.NewRecorder() router.ServeHTTP(legacyRecorder, httptest.NewRequest(http.MethodGet, "/vendor/jobs/legacy-task", nil)) assert.Equal(t, http.StatusOK, legacyRecorder.Code) assert.JSONEq(t, `{"id":"legacy-task"}`, legacyRecorder.Body.String()) var firstBody string for _, taskID := range []string{"missing-task", "foreign-task", "wrong-platform", "wrong-legacy-platform", "discarded-task"} { recorder := httptest.NewRecorder() router.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/vendor/jobs/"+taskID, nil)) assert.Equal(t, http.StatusNotFound, recorder.Code) if firstBody == "" { firstBody = recorder.Body.String() } else { assert.JSONEq(t, firstBody, recorder.Body.String()) } assert.Contains(t, recorder.Body.String(), `"code":"task_not_found"`) assert.NotContains(t, recorder.Body.String(), taskID) } } func TestRespondTaskPluginErrorNeverExposesInternalDetails(t *testing.T) { plugin := compileTaskRoutePlugin(t, ` export const meta = { apiVersion: 1, key: "route-error-test", name: "Error", version: "1.0.0", author: {name: "Test"}, models: ["error-model"], fetchMode: "per_task", }; export function buildSubmitRequest() { return {url: "https://example.com"}; } export function parseSubmitResponse() { return {taskId: "one"}; } export function buildQueryRequest() { return {url: "https://example.com"}; } export function parseTaskResult() { return {status: "SUCCESS"}; } export const native = {error: function(ctx, error) { return {path: ctx.path, code: error.code, message: error.message, status: error.httpStatus, retryable: error.retryable, keys: Object.keys(error).sort()}; }}; `) recorder := httptest.NewRecorder() c, _ := gin.CreateTestContext(recorder) c.Request = httptest.NewRequest(http.MethodPost, "/vendor/failure", strings.NewReader(`{"credential":"do-not-copy"}`)) c.Set(jsplugin.ContextKeyPinnedRoute, jsplugin.PinnedRoute{Plugin: plugin}) c.Set(jsplugin.ContextKeyRouteRequest, jsplugin.RouteRequestContext{ Path: "/vendor/failure", Method: http.MethodPost, Params: map[string]string{}, Query: map[string][]string{}, }) handled := RespondTaskPluginError(c, &dto.TaskError{ Code: "upstream_credential_failure", Message: "https://user:password@upstream.invalid?token=secret", Data: map[string]any{"authorization": "Bearer secret"}, StatusCode: http.StatusBadGateway, Error: assert.AnError, }) assert.True(t, handled) assert.Equal(t, http.StatusBadGateway, recorder.Code) assert.JSONEq(t, `{ "path": "/vendor/failure", "code": "server_error", "message": "Task request failed", "status": 502, "retryable": true, "keys": ["code", "httpStatus", "message", "requestId", "retryable"] }`, recorder.Body.String()) assert.NotContains(t, recorder.Body.String(), "upstream.invalid") assert.NotContains(t, recorder.Body.String(), "secret") assert.NotContains(t, recorder.Body.String(), "password") } func TestTaskPluginErrorFallbackIsSanitized(t *testing.T) { for _, testCase := range []struct { name string renderHook string }{ {name: "missing hook"}, {name: "throwing hook", renderHook: `export const native = {error: function() { throw new Error("renderer secret"); }};`}, } { t.Run(testCase.name, func(t *testing.T) { plugin := compileTaskRoutePlugin(t, genericTaskPluginSource+"\n"+testCase.renderHook) recorder := httptest.NewRecorder() c, _ := gin.CreateTestContext(recorder) c.Request = httptest.NewRequest(http.MethodPost, "/vendor/failure", nil) c.Set(common.RequestIdKey, "fallback-req") c.Set(jsplugin.ContextKeyPinnedRoute, jsplugin.PinnedRoute{Plugin: plugin}) c.Set(jsplugin.ContextKeyRouteRequest, jsplugin.RouteRequestContext{ Path: "/vendor/failure", Method: http.MethodPost, Params: map[string]string{}, Query: map[string][]string{}, }) abortWithOpenAiMessage(c, http.StatusBadGateway, "https://user:password@upstream.invalid?token=secret") assert.Equal(t, http.StatusBadGateway, recorder.Code) assert.JSONEq(t, `{"code":"server_error","message":"Task request failed (request id: fallback-req)","data":null}`, recorder.Body.String()) assert.NotContains(t, recorder.Body.String(), "upstream.invalid") assert.NotContains(t, recorder.Body.String(), "renderer secret") assert.NotContains(t, recorder.Body.String(), "password") }) } } func TestResolvedTaskPluginIDsEnforcesPublicQueryContract(t *testing.T) { valid, ok := resolvedTaskPluginIDs([]any{"task-a", "task-b", "task-a"}) require.True(t, ok) assert.Equal(t, []string{"task-a", "task-b", "task-a"}, valid) empty, ok := resolvedTaskPluginIDs([]any{}) require.True(t, ok) assert.Empty(t, empty) tooMany := make([]any, 101) for index := range tooMany { tooMany[index] = "task" } for _, invalid := range []any{ []any{"task-a", ""}, []any{"task-a", " "}, []any{"task-a", 2}, tooMany, "task-a", } { _, ok = resolvedTaskPluginIDs(invalid) assert.False(t, ok) } } func TestSunoFetchEmptyIDsReturnsSuccessfulEmptyArray(t *testing.T) { source, err := builtinplugins.Source("sunoapi") require.NoError(t, err) plugin := compileTaskRoutePlugin(t, source) nextHandlerCalled := false router := gin.New() router.POST( "/suno/fetch", pinTaskPluginRoute(plugin, 1), PrepareTaskPluginRoute(), func(c *gin.Context) { nextHandlerCalled = true c.Status(http.StatusTeapot) }, ) request := httptest.NewRequest(http.MethodPost, "/suno/fetch", strings.NewReader(`{"ids":[]}`)) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.False(t, nextHandlerCalled) assert.Equal(t, http.StatusOK, recorder.Code) assert.JSONEq(t, `{"code":"success","message":"","data":[]}`, recorder.Body.String()) } func TestPrepareTaskPluginRouteSurfacesDecodeHookMessage(t *testing.T) { plugin := compileTaskRoutePlugin(t, ` export const meta = { apiVersion: 1, key: "route-decode-detail-test", name: "Decode", version: "1.0.0", author: {name: "Test"}, models: ["detail-model"], fetchMode: "per_task", routes: [{method: "POST", path: "/vendor/jobs", type: "submit", decode: "createTask", render: "created"}], }; export const native = {createTask: function() { throw new Error("model is required"); }, created: function(ctx, task) { return task; }}; export function buildSubmitRequest() { return {url: "https://example.com"}; } export function parseSubmitResponse() { return {taskId: "one"}; } export function buildQueryRequest() { return {url: "https://example.com"}; } export function parseTaskResult() { return {status: "SUCCESS"}; } `) router := gin.New() router.POST("/vendor/jobs", pinTaskPluginRoute(plugin, 0), PrepareTaskPluginRoute(), func(c *gin.Context) { c.Status(http.StatusNoContent) }) request := httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{"prompt":"x"}`)) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.Equal(t, http.StatusBadRequest, recorder.Code) var body dto.TaskError require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &body)) assert.Equal(t, "invalid_request", body.Code) assert.Equal(t, "model is required", body.Message) } func TestPrepareTaskPluginRouteNativeErrorReceivesHookMessage(t *testing.T) { plugin := compileTaskRoutePlugin(t, ` export const meta = { apiVersion: 1, key: "route-error-detail-test", name: "Error", version: "1.0.0", author: {name: "Test"}, models: ["detail-model"], fetchMode: "per_task", routes: [{method: "POST", path: "/vendor/jobs", type: "submit", decode: "createTask", render: "created"}], }; export const native = { createTask: function() { throw new Error("model is required"); }, created: function(ctx, task) { return task; }, error: function(ctx, error) { return {code: error.code, message: error.message, requestId: error.requestId}; }, }; export function buildSubmitRequest() { return {url: "https://example.com"}; } export function parseSubmitResponse() { return {taskId: "one"}; } export function buildQueryRequest() { return {url: "https://example.com"}; } export function parseTaskResult() { return {status: "SUCCESS"}; } `) router := gin.New() router.Use(func(c *gin.Context) { c.Set(common.RequestIdKey, "native-error-req") c.Next() }) router.POST("/vendor/jobs", pinTaskPluginRoute(plugin, 0), PrepareTaskPluginRoute(), func(c *gin.Context) { c.Status(http.StatusNoContent) }) request := httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{"prompt":"x"}`)) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.Equal(t, http.StatusBadRequest, recorder.Code) assert.JSONEq(t, `{"code":"invalid_request","message":"model is required","requestId":"native-error-req"}`, recorder.Body.String()) } func TestPrepareTaskPluginRouteRejectsNonObjectResultWithFixedMessage(t *testing.T) { plugin := compileTaskRoutePlugin(t, ` export const meta = { apiVersion: 1, key: "route-result-object-test", name: "Result", version: "1.0.0", author: {name: "Test"}, models: ["detail-model"], fetchMode: "per_task", routes: [{method: "POST", path: "/vendor/jobs", type: "submit", decode: "createTask", render: "created"}], }; export const native = {createTask: function() { return "not-an-object"; }, created: function(ctx, task) { return task; }}; export function buildSubmitRequest() { return {url: "https://example.com"}; } export function parseSubmitResponse() { return {taskId: "one"}; } export function buildQueryRequest() { return {url: "https://example.com"}; } export function parseTaskResult() { return {status: "SUCCESS"}; } `) router := gin.New() router.POST("/vendor/jobs", pinTaskPluginRoute(plugin, 0), PrepareTaskPluginRoute(), func(c *gin.Context) { c.Status(http.StatusNoContent) }) request := httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{"model":"detail-model"}`)) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.Equal(t, http.StatusBadRequest, recorder.Code) var body dto.TaskError require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &body)) assert.Equal(t, "invalid_request", body.Code) assert.Equal(t, "plugin returned an invalid route result", body.Message) } func TestPrepareTaskPluginRouteSurfacesRequestDecodeDetail(t *testing.T) { plugin := compileTaskRoutePlugin(t, ` export const meta = { apiVersion: 1, key: "route-decode-body-test", name: "Decode", version: "1.0.0", author: {name: "Test"}, models: ["detail-model"], fetchMode: "per_task", routes: [{method: "POST", path: "/vendor/jobs", type: "submit", decode: "createTask", render: "created"}], }; export const native = {createTask: function() { throw new Error("decoder must not run"); }, created: function(ctx, task) { return task; }}; export function buildSubmitRequest() { return {url: "https://example.com"}; } export function parseSubmitResponse() { return {taskId: "one"}; } export function buildQueryRequest() { return {url: "https://example.com"}; } export function parseTaskResult() { return {status: "SUCCESS"}; } `) router := gin.New() router.POST("/vendor/jobs", pinTaskPluginRoute(plugin, 0), PrepareTaskPluginRoute(), func(c *gin.Context) { c.Status(http.StatusNoContent) }) request := httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{}`)) request.Header["Content-Type"] = []string{"application/json", "application/x-www-form-urlencoded"} recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.Equal(t, http.StatusBadRequest, recorder.Code) var body dto.TaskError require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &body)) assert.Equal(t, "invalid_request", body.Code) assert.Contains(t, body.Message, "conflicting Content-Type") } func TestSanitizedTaskPluginErrorIgnoresDetailOn5xx(t *testing.T) { got := sanitizedTaskPluginError(http.StatusInternalServerError, "database secret") assert.Equal(t, "server_error", got.Code) assert.Equal(t, "Task request failed", got.Message) assert.Equal(t, http.StatusInternalServerError, got.HTTPStatus) got = sanitizedTaskPluginError(http.StatusBadGateway, "https://user:password@upstream.invalid") assert.Equal(t, "server_error", got.Code) assert.Equal(t, "Task request failed", got.Message) } func TestTaskPluginErrorFallbackMessageIncludesRequestID(t *testing.T) { plugin := compileTaskRoutePlugin(t, genericTaskPluginSource) recorder := httptest.NewRecorder() c, _ := gin.CreateTestContext(recorder) c.Request = httptest.NewRequest(http.MethodPost, "/vendor/failure", nil) c.Set(common.RequestIdKey, "req-fallback-1") c.Set(jsplugin.ContextKeyPinnedRoute, jsplugin.PinnedRoute{Plugin: plugin}) c.Set(jsplugin.ContextKeyRouteRequest, jsplugin.RouteRequestContext{ Path: "/vendor/failure", Method: http.MethodPost, Params: map[string]string{}, Query: map[string][]string{}, }) abortTaskPluginRouteErrorDetail(c, http.StatusBadRequest, "model is required") assert.Equal(t, http.StatusBadRequest, recorder.Code) var body dto.TaskError require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &body)) assert.Equal(t, "invalid_request", body.Code) assert.Equal(t, "model is required (request id: req-fallback-1)", body.Message) assert.True(t, strings.HasSuffix(body.Message, "(request id: req-fallback-1)")) } func TestPrepareTaskPluginEndpointSurfacesDecodeHookMessage(t *testing.T) { const key = "endpoint-decode-detail-test" _, err := jsplugin.DefaultRegistry.Register(taskProtocolPluginSource( key, "1.0.0", `["claimed-model"]`, "/v1/responses", `throw new Error("model is required");`, ), jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(key)) }) router := gin.New() router.POST("/v1/responses", PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) { c.Status(http.StatusNoContent) }) request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"claimed-model"}`)) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.Equal(t, http.StatusBadRequest, recorder.Code) assert.Contains(t, recorder.Body.String(), "model is required") assert.NotContains(t, recorder.Body.String(), "Invalid task protocol request") } func TestPinTaskPluginEndpointRejectsUnsupportedRequestForms(t *testing.T) { setupTaskPluginRouteDB(t) tests := []struct { name string key string supports string hooks string body string message string }{ { name: "stream against sync and background", key: "form-gate-final-only", supports: `["sync", "background"]`, hooks: `renderFinal: function() { return {}; }`, body: `{"model":"form-gate-model","stream":true}`, message: `Streaming is not supported for this model. Set "stream": false, or use "background": true and retrieve the response later.`, }, { name: "sync against stream only", key: "form-gate-stream-only", supports: `["stream"]`, hooks: `renderEvents: function() { return {events: [], done: false}; }`, body: `{"model":"form-gate-model"}`, message: `Synchronous non-streaming requests are not supported for this model. Set "stream": true.`, }, { name: "background against stream and sync", key: "form-gate-no-background", supports: `["stream", "sync"]`, hooks: `renderEvents: function() { return {events: [], done: false}; }, renderFinal: function() { return {}; }`, body: `{"model":"form-gate-model","background":true}`, message: `Background mode is not supported for this model. Remove "background": true.`, }, { name: "background plus stream reports stream first", key: "form-gate-no-stream", supports: `["sync", "background"]`, hooks: `renderFinal: function() { return {}; }`, body: `{"model":"form-gate-model","stream":true,"background":true}`, message: `Streaming is not supported for this model. Set "stream": false, or use "background": true and retrieve the response later.`, }, } for _, testCase := range tests { t.Run(testCase.name, func(t *testing.T) { _, err := jsplugin.DefaultRegistry.Register(taskResponsesPluginSource( testCase.key, 0, `["form-gate-model"]`, testCase.supports, testCase.hooks, `return {model: ctx.model};`, ), jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(testCase.key)) }) reachedPrepare := false quotaConsumed := false router := gin.New() router.POST("/v1/responses", PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) { reachedPrepare = true quotaConsumed = true require.NoError(t, model.DB.Create(&model.Task{TaskID: "should-not-exist", UserId: 1}).Error) c.Status(http.StatusNoContent) }) request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(testCase.body)) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.Equal(t, http.StatusBadRequest, recorder.Code) var payload map[string]any require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &payload)) errObj, ok := payload["error"].(map[string]any) require.True(t, ok) message, _ := errObj["message"].(string) assert.Contains(t, message, testCase.message) assert.False(t, reachedPrepare) assert.False(t, quotaConsumed) var count int64 require.NoError(t, model.DB.Model(&model.Task{}).Count(&count).Error) assert.Zero(t, count) }) } } func TestPinTaskPluginEndpointMalformedStreamStillFailsInPrepare(t *testing.T) { const key = "form-gate-malformed-stream" _, err := jsplugin.DefaultRegistry.Register(taskResponsesPluginSource( key, 0, `["form-gate-bool-model"]`, `["sync", "background"]`, `renderFinal: function() { return {}; }`, `return {model: ctx.model};`, ), jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(key)) }) reachedNext := false router := gin.New() router.POST("/v1/responses", PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) { reachedNext = true c.Status(http.StatusNoContent) }) request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"form-gate-bool-model","stream":"yes"}`)) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.Equal(t, http.StatusBadRequest, recorder.Code) assert.Contains(t, recorder.Body.String(), "stream must be a boolean") assert.False(t, reachedNext) } func TestPinTaskPluginEndpointMovesParserToSurvivingSharedCandidate(t *testing.T) { streamOnly, err := jsplugin.DefaultRegistry.Register(taskResponsesPluginSource( "alpha-stream", constant.ChannelTypeReplicate, `["shared-form-model"]`, `["stream"]`, `renderEvents: function() { return {events: [], done: false}; }`, `return {model: ctx.model, action: "stream-parser"};`, ), jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister("alpha-stream")) }) full, err := jsplugin.DefaultRegistry.Register(taskResponsesPluginSource( "bravo-full", constant.ChannelTypeCodex, `["shared-form-model"]`, `["stream", "sync", "background"]`, `renderEvents: function() { return {events: [], done: false}; }, renderFinal: function() { return {}; }`, `return {model: ctx.model, action: "full-parser"};`, ), jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister("bravo-full")) }) generation := jsplugin.DefaultRegistry.Generation() unfiltered := generation.LookupEndpointCandidates("POST", "/v1/responses", "shared-form-model") require.Len(t, unfiltered, 2) assert.Equal(t, "alpha-stream", unfiltered[0].Plugin.Meta.Key) t.Run("sync moves pin to second candidate", func(t *testing.T) { var pinned jsplugin.PinnedEndpoint var action string router := gin.New() router.POST("/v1/responses", PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) { pinned = c.MustGet(jsplugin.ContextKeyPinnedEndpoint).(jsplugin.PinnedEndpoint) action = c.GetString("task_action") c.Status(http.StatusNoContent) }) request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"shared-form-model"}`)) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.Equal(t, http.StatusNoContent, recorder.Code) assert.Same(t, full, pinned.Plugin) require.Len(t, pinned.Candidates, 1) assert.Same(t, full, pinned.Candidates[0].Plugin) assert.Equal(t, "full-parser", action) }) t.Run("stream keeps first candidate", func(t *testing.T) { var pinned jsplugin.PinnedEndpoint var action string router := gin.New() router.POST("/v1/responses", PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) { pinned = c.MustGet(jsplugin.ContextKeyPinnedEndpoint).(jsplugin.PinnedEndpoint) action = c.GetString("task_action") c.Status(http.StatusNoContent) }) request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"shared-form-model","stream":true}`)) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) assert.Equal(t, http.StatusNoContent, recorder.Code) assert.Same(t, streamOnly, pinned.Plugin) require.Len(t, pinned.Candidates, 2) assert.Same(t, streamOnly, pinned.Candidates[0].Plugin) assert.Equal(t, "stream-parser", action) }) } func compileTaskRoutePlugin(t *testing.T, source string) *jsplugin.LoadedPlugin { t.Helper() plugin, err := jsplugin.CompilePlugin(source, jsplugin.Options{}) require.NoError(t, err) return plugin } func taskResponsesPluginSource(key string, channelType int, models, supports, hooks, parseRequestBody string) string { channelField := "" if channelType > 0 { channelField = fmt.Sprintf("channelTypes: [%d],", channelType) } return fmt.Sprintf(` export const meta = { apiVersion: 1, key: %q, name: %q, version: "1.0.0", author: {name: "Test"}, %s models: %s, fetchMode: "per_task", protocols: [{name: "openai_responses", supports: %s}], }; export function buildSubmitRequest() { return {url: "https://example.com"}; } export function parseSubmitResponse() { return {taskId: "one"}; } export function buildQueryRequest() { return {url: "https://example.com"}; } export function parseTaskResult() { return {status: "SUCCESS"}; } export function listArtifacts() { return []; } export function buildContentRequest() { throw new Error("artifact_not_found"); } export const protocols = {openai_responses: { decodeRequest: function(ctx) { ctx.requestBody = ctx.body.value; const decode = function() { %s }; const result = decode(); if (!result.kind) result.kind = "submit"; return result; }, %s }}; `, key, key, channelField, models, supports, parseRequestBody, hooks) } func taskProtocolPluginSource(key, version, models, endpoint, parseRequestBody string) string { protocol := "openai_responses" protocolClaim := `{name: "openai_responses", supports: ["stream", "sync", "background"]}` presenters := `renderEvents: function() { return {events: [], done: false}; }, renderFinal: function() { return {}; },` if endpoint == "/v1/videos" { protocol = "openai_video" protocolClaim = `"openai_video"` presenters = `render: function() { return {}; },` } if endpoint == "/v1/images/edits" || endpoint == "/v1/images/generations" { protocol = "openai_image" protocolClaim = `"openai_image"` presenters = `render: function() { return {data: []}; },` } return fmt.Sprintf(` export const meta = { apiVersion: 1, key: %q, name: %q, version: %q, author: {name: "Test"}, models: %s, fetchMode: "per_task", protocols: [%s], }; export function buildSubmitRequest() { return {url: "https://example.com"}; } export function parseSubmitResponse() { return {taskId: "one"}; } export function buildQueryRequest() { return {url: "https://example.com"}; } export function parseTaskResult() { return {status: "SUCCESS"}; } export function listArtifacts() { return []; } export function buildContentRequest() { throw new Error("artifact_not_found"); } export const protocols = { %s: { decodeRequest: function(ctx) { ctx.requestBody = ctx.body.value; if (ctx.body.kind === "multipart") { ctx.requestBody = {}; for (const key of Object.keys(ctx.body.fields || {})) ctx.requestBody[key] = ctx.body.fields[key][0]; } const decode = function() { %s }; const result = decode(); if (!result.kind) result.kind = "submit"; return result; }, %s }, }; `, key, key, version, models, protocolClaim, protocol, parseRequestBody, presenters) } func pinTaskPluginRoute(plugin *jsplugin.LoadedPlugin, routeIndex int) gin.HandlerFunc { return func(c *gin.Context) { c.Set(jsplugin.ContextKeyPinnedRoute, jsplugin.PinnedRoute{Plugin: plugin, Route: plugin.Meta.Routes[routeIndex]}) c.Next() } } func setupTaskPluginRouteDB(t *testing.T) { t.Helper() previousDB := model.DB previousType := common.MainDatabaseType() database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) require.NoError(t, err) require.NoError(t, database.AutoMigrate(&model.Task{})) model.DB = database common.SetMainDatabaseType(common.DatabaseTypeSQLite) t.Cleanup(func() { model.DB = previousDB common.SetMainDatabaseType(previousType) }) } func insertTaskPluginRouteTask(t *testing.T, task *model.Task) { t.Helper() require.NoError(t, model.DB.Create(task).Error) } func TestPrepareTaskPluginEndpointFiltersEachSharedCandidate(t *testing.T) { for _, tc := range []struct { name, alpha, beta string wantKeys []string wantError string }{ {name: "first decoder rejects", alpha: `throw new Error("alpha only accepts 720p")`, beta: `return {model:ctx.model,action:"beta",requestBody:{resolution:"1080p"}}`, wantKeys: []string{"decode-beta"}}, {name: "second decoder rejects", alpha: `return {model:ctx.model,action:"alpha"}`, beta: `throw new Error("beta rejects")`, wantKeys: []string{"decode-alpha"}}, {name: "both decoders accept", alpha: `return {model:ctx.model,action:"alpha"}`, beta: `return {model:ctx.model,action:"beta"}`, wantKeys: []string{"decode-alpha", "decode-beta"}}, {name: "all decoders reject", alpha: `throw new Error("alpha rejects first")`, beta: `throw new Error("beta rejects second")`, wantError: "decode-alpha: alpha rejects first; decode-beta: beta rejects second"}, {name: "duplicate failures are grouped", alpha: `throw new Error("unsupported resolution")`, beta: `throw new Error("unsupported resolution")`, wantError: "decode-alpha, decode-beta: unsupported resolution"}, {name: "invalid result is excluded", alpha: `return {kind:"query",model:ctx.model}`, beta: `return {model:ctx.model,action:"beta"}`, wantKeys: []string{"decode-beta"}}, {name: "rewritten model is excluded", alpha: `return {model:"another-model"}`, beta: `return {model:ctx.model,action:"beta"}`, wantKeys: []string{"decode-beta"}}, } { t.Run(tc.name, func(t *testing.T) { for _, spec := range []struct{ key, decode string }{{"decode-alpha", tc.alpha}, {"decode-beta", tc.beta}} { _, err := jsplugin.DefaultRegistry.Register(taskResponsesPluginSource(spec.key, 0, `["decode-shared-model"]`, `["sync"]`, `renderFinal:function(){return {};}`, spec.decode), jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(spec.key)) }) } var gotKeys []string router := gin.New() router.POST("/v1/responses", PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) { pinned := c.MustGet(jsplugin.ContextKeyPinnedEndpoint).(jsplugin.PinnedEndpoint) for _, candidate := range pinned.Candidates { gotKeys = append(gotKeys, candidate.Plugin.Meta.Key) } assert.Equal(t, tc.wantKeys[0], pinned.Plugin.Meta.Key) assert.Equal(t, tc.wantKeys[0], c.GetString("task_plugin_key")) assert.Same(t, pinned.Plugin, c.MustGet(jsplugin.ContextKeyPinnedPlugin).(jsplugin.PinnedPlugin).Plugin) assert.Equal(t, tc.wantKeys, service.GetChannelConstraints(c).Filters[0].TaskPluginKeys) c.Status(http.StatusNoContent) }) request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"decode-shared-model","resolution":"1080p"}`)) request.Header.Set("Content-Type", "application/json") recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) if tc.wantError != "" { assert.Equal(t, http.StatusBadRequest, recorder.Code) assert.Contains(t, recorder.Body.String(), tc.wantError) assert.Empty(t, gotKeys) } else { assert.Equal(t, http.StatusNoContent, recorder.Code, recorder.Body.String()) assert.Equal(t, tc.wantKeys, gotKeys) } }) } } func TestRouteTaskPluginEndpointUsesSelectedChannel(t *testing.T) { require.NoError(t, appI18n.Init()) for _, dialect := range []string{"sqlite", "mysql", "postgres"} { t.Run(dialect, func(t *testing.T) { dsn := "" if dialect != "sqlite" { dsn = os.Getenv("TEST_" + strings.ToUpper(dialect) + "_DSN") if dsn == "" { t.Skip("database test DSN not configured") } } previousDB, previousMain, previousLog := model.DB, common.MainDatabaseType(), common.LogDatabaseType() previousPath, previousMaster := common.SQLitePath, common.IsMasterNode previousCache, previousRedis := common.MemoryCacheEnabled, common.RedisEnabled common.SQLitePath, common.IsMasterNode = t.TempDir()+"/routing.db", false common.RedisEnabled = false t.Setenv("SQL_DSN", dsn) require.NoError(t, model.InitDB()) database := model.DB // Every table belongs to this run, including when a developer points // a test DSN at a database shared with another test suite. database.Config.NamingStrategy = schema.NamingStrategy{TablePrefix: "endpoint_route_" + common.GetUUID() + "_"} require.NoError(t, database.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.Task{})) t.Cleanup(func() { for _, table := range []any{&model.Ability{}, &model.Channel{}, &model.Task{}} { require.NoError(t, database.Session(&gorm.Session{AllowGlobalUpdate: true}).Delete(table).Error) } common.MemoryCacheEnabled = true model.InitChannelCache() require.NoError(t, database.Migrator().DropTable(&model.Ability{}, &model.Channel{}, &model.Task{})) model.DB = previousDB common.SetDatabaseTypes(previousMain, previousLog) common.SQLitePath, common.IsMasterNode = previousPath, previousMaster common.MemoryCacheEnabled, common.RedisEnabled = previousCache, previousRedis if previousDB != nil { model.InitChannelCache() } sqlDB, err := database.DB() require.NoError(t, err) require.NoError(t, sqlDB.Close()) }) query := "SELECT version()" if dialect == "sqlite" { query = "SELECT sqlite_version()" } var version string require.NoError(t, database.Raw(query).Scan(&version).Error) t.Logf("database version: %s", version) for _, memoryCache := range []bool{false, true} { t.Run(fmt.Sprintf("memory_cache=%t", memoryCache), func(t *testing.T) { common.MemoryCacheEnabled = memoryCache for _, tc := range []struct { name, path, model string pluginType int ordinaryType int ordinaryPriority int64 pluginPriority int64 pluginOnly, wantPlugin bool multipart, stream bool pinOrdinary, pinPlugin bool wrongGroup, wrongModel bool disabled, unrelated bool decoderRejects, syncOnly bool originPin, retry bool legacy, mapped, shared bool wantStatus int }{ {name: "ordinary images bypass plugin decoder", decoderRejects: true}, {name: "ordinary images win by priority", ordinaryPriority: 20}, {name: "plugin images win by priority", pluginPriority: 20, wantPlugin: true}, {name: "type 61 plugin remains usable", pluginOnly: true, wantPlugin: true}, {name: "Alibaba legacy channel remains usable", legacy: true, pluginOnly: true, wantPlugin: true}, {name: "OpenAI can win against Alibaba for qwen-image-3.0", legacy: true, ordinaryPriority: 20}, {name: "Alibaba can win against OpenAI for qwen-image-3.0", legacy: true, pluginPriority: 20, wantPlugin: true}, {name: "bound New API uses plugin", pluginType: constant.ChannelTypeNewAPI, pluginPriority: 20, wantPlugin: true}, {name: "unbound New API uses ordinary relay", ordinaryType: constant.ChannelTypeNewAPI, decoderRejects: true}, {name: "unrelated type 61 cannot steal request", unrelated: true, pluginPriority: 20}, {name: "disabled plugin channel is ignored", disabled: true, pluginPriority: 20}, {name: "plugin in another group is ignored", wrongGroup: true, pluginPriority: 20}, {name: "plugin without requested model is ignored", wrongModel: true, pluginPriority: 20}, {name: "ordinary channel pin wins", pinOrdinary: true, pluginPriority: 20}, {name: "plugin channel pin wins", pinPlugin: true, wantPlugin: true}, {name: "ordinary model spelling is preserved", model: "CHANNEL-ROUTE-MODEL", decoderRejects: true}, {name: "mapped plugin alias remains usable", model: "friendly-image", mapped: true, pluginOnly: true, wantPlugin: true}, {name: "ordinary alias is not rewritten by plugin", model: "friendly-image", mapped: true, decoderRejects: true}, {name: "ordinary multipart edits are preserved", path: "/v1/images/edits", multipart: true, decoderRejects: true}, {name: "plugin multipart edits remain usable", path: "/v1/images/edits", multipart: true, pluginOnly: true, wantPlugin: true}, {name: "ordinary Responses bypass unsupported plugin mode", path: "/v1/responses", stream: true, syncOnly: true}, {name: "plugin Responses remain usable", path: "/v1/responses", pluginOnly: true, wantPlugin: true}, {name: "ordinary videos bypass plugin decoder", path: "/v1/videos", multipart: true, decoderRejects: true}, {name: "plugin videos remain usable", path: "/v1/videos", multipart: true, pluginOnly: true, wantPlugin: true}, {name: "selected plugin retains validation errors", pluginOnly: true, decoderRejects: true, wantStatus: http.StatusBadRequest}, {name: "selected plugin retains mode errors", path: "/v1/responses", pluginOnly: true, stream: true, syncOnly: true, wantStatus: http.StatusBadRequest}, {name: "origin task still pins its channel", pluginOnly: true, originPin: true, wantPlugin: true}, {name: "ordinary retries stay on ordinary channels", ordinaryPriority: 20, pluginPriority: 10, retry: true}, {name: "plugin retries stay on plugin channels", pluginPriority: 20, retry: true, wantPlugin: true}, {name: "rejected shared candidate still allows another plugin", pluginOnly: true, shared: true, decoderRejects: true, pluginPriority: 20, wantPlugin: true}, } { t.Run(tc.name, func(t *testing.T) { for _, table := range []any{&model.Ability{}, &model.Channel{}, &model.Task{}} { require.NoError(t, database.Session(&gorm.Session{AllowGlobalUpdate: true}).Delete(table).Error) } key := "channel-route-test" path := common.GetStringIfEmpty(tc.path, "/v1/images/generations") modelName := common.GetStringIfEmpty(tc.model, "channel-route-model") if tc.legacy { key, modelName = "alibaba", "qwen-image-3.0" } decode := `return {model: ctx.model, action: "generate", requestBody: {normalized: true}};` if tc.decoderRejects { decode = `throw new Error("plugin rejected this request");` } else if tc.originPin { decode = `return {model: ctx.model, originTaskIds: ["task-route-origin"], requestBody: {normalized: true}};` } source := taskProtocolPluginSource(key, "1.0.0", `["channel-route-model"]`, path, decode) if tc.syncOnly { source = taskResponsesPluginSource(key, 0, `["channel-route-model"]`, `["sync"]`, `renderFinal: function() { return {}; },`, decode) } if !tc.legacy { _, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(key)) }) } ordinary := &model.Channel{Type: constant.ChannelTypeOpenAI, Status: common.ChannelStatusEnabled, Name: "ordinary", Key: "ordinary-key", Group: "route-test", Models: modelName, Priority: common.GetPointer(tc.ordinaryPriority + 1)} if tc.ordinaryType != 0 { ordinary.Type = tc.ordinaryType } plugin := &model.Channel{Type: constant.ChannelTypeTaskPlugin, Status: common.ChannelStatusEnabled, Name: "plugin", Key: "plugin-key", Group: "route-test", Models: modelName, Priority: &tc.pluginPriority} if tc.pluginType != 0 { plugin.Type = tc.pluginType } plugin.SetSetting(kitdto.ChannelSettings{TaskPluginKey: key}) if tc.legacy { plugin.Type = constant.ChannelTypeAli } if tc.mapped { plugin.ModelMapping = common.GetPointer(`{"friendly-image":"channel-route-model"}`) } if tc.unrelated { plugin.SetSetting(kitdto.ChannelSettings{TaskPluginKey: "unrelated"}) } if tc.wrongGroup { plugin.Group = "another-group" } if tc.wrongModel { plugin.Models = "another-model" } if tc.disabled { plugin.Status = common.ChannelStatusManuallyDisabled } channels := []*model.Channel{plugin} if !tc.pluginOnly { channels = append(channels, ordinary) } if tc.originPin || tc.retry || tc.shared { alternate := &model.Channel{Type: ordinary.Type, Status: common.ChannelStatusEnabled, Name: "alternate", Key: "alternate-key", Group: "route-test", Models: modelName} if tc.wantPlugin { alternate.Type = plugin.Type alternate.SetSetting(plugin.GetSetting()) } if tc.shared { const alternateKey = "channel-route-ztest" alternateSource := taskProtocolPluginSource(alternateKey, "1.0.0", `["channel-route-model"]`, path, `return {model: ctx.model, requestBody: {normalized: true}};`) _, err := jsplugin.DefaultRegistry.Register(alternateSource, jsplugin.Options{}) require.NoError(t, err) t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(alternateKey)) }) alternate.SetSetting(kitdto.ChannelSettings{TaskPluginKey: alternateKey}) } channels = append(channels, alternate) } for _, channel := range channels { require.NoError(t, channel.Insert()) } if tc.originPin { require.NoError(t, database.Create(&model.Task{TaskID: "task-route-origin", UserId: 7, ChannelId: channels[len(channels)-1].Id, Platform: constant.TaskPlatform(key), Status: model.TaskStatusSuccess}).Error) } model.InitChannelCache() body := []byte(fmt.Sprintf(`{"model":%q,"prompt":"a cat","size":"1024x1024","stream":%t}`, modelName, tc.stream)) contentType := "application/json" if tc.multipart { var form bytes.Buffer writer := multipart.NewWriter(&form) require.NoError(t, writer.WriteField("model", modelName)) require.NoError(t, writer.WriteField("prompt", "a cat")) file, err := writer.CreateFormFile("image", "reference.png") require.NoError(t, err) _, err = file.Write([]byte("image bytes")) require.NoError(t, err) require.NoError(t, writer.Close()) body, contentType = form.Bytes(), writer.FormDataContentType() } reached := false router := gin.New() router.POST(path, func(c *gin.Context) { common.SetContextKey(c, constant.ContextKeyUsingGroup, "route-test") common.SetContextKey(c, constant.ContextKeyUserId, 7) if tc.pinOrdinary || tc.pinPlugin { id := ordinary.Id if tc.pinPlugin { id = plugin.Id } service.GetChannelConstraints(c).AddPin(dto.ChannelPin{ChannelId: id, Source: dto.PinSourceToken, Rank: dto.PinRankToken}) } }, RouteTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), Distribute(), func(c *gin.Context) { reached = true _, pinned := c.Get(jsplugin.ContextKeyPinnedEndpoint) assert.Equal(t, tc.wantPlugin, pinned) _, relayEndpoint := c.Get(jsplugin.ContextKeyRelayEndpoint) assert.Equal(t, !tc.wantPlugin, relayEndpoint) expected := ordinary if tc.wantPlugin { expected = plugin if tc.shared { assert.Equal(t, "channel-route-ztest", c.GetString("task_plugin_key")) } else { assert.Equal(t, key, c.GetString("task_plugin_key")) } if tc.legacy { assert.NotNil(t, c.MustGet("task_request")) } else { assert.Equal(t, map[string]any{"normalized": true}, c.MustGet("task_request")) } if tc.mapped { assert.Equal(t, "channel-route-model", c.MustGet(jsplugin.ContextKeyPinnedEndpoint).(jsplugin.PinnedEndpoint).MappedModel) } if tc.originPin || tc.shared { expected = channels[len(channels)-1] } } else { assert.Empty(t, c.GetString("expected_task_plugin_key")) assert.Empty(t, c.GetString("resolved_task_model")) storage, err := common.GetBodyStorage(c) require.NoError(t, err) got, err := storage.Bytes() require.NoError(t, err) assert.Equal(t, body, got, "ordinary relay must receive the original body") } assert.Equal(t, expected.Id, common.GetContextKeyInt(c, constant.ContextKeyChannelId)) assert.Equal(t, expected.Key, common.GetContextKeyString(c, constant.ContextKeyChannelKey)) if tc.retry { next, _, err := service.CacheGetRandomSatisfiedChannel(&service.RetryParam{Ctx: c, TokenGroup: "route-test", ModelName: modelName, Retry: common.GetPointer(1)}) require.NoError(t, err) require.NotNil(t, next) assert.Equal(t, channels[len(channels)-1].Id, next.Id) } c.Status(http.StatusNoContent) }) request := httptest.NewRequest(http.MethodPost, path, bytes.NewReader(body)) request.Header.Set("Content-Type", contentType) recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) status := tc.wantStatus if status == 0 { status = http.StatusNoContent } assert.Equal(t, status, recorder.Code, recorder.Body.String()) assert.Equal(t, status == http.StatusNoContent, reached) if tc.decoderRejects && status == http.StatusBadRequest { assert.Contains(t, recorder.Body.String(), "plugin rejected this request") } }) } }) } }) } } // sharedRouteTestPluginSource declares the probe routes two plugins share; // decoder and presenter bodies name the plugin, and extra adds members to // native. func sharedRouteTestPluginSource(key, models, routes, extra string) string { return fmt.Sprintf(` export const meta = {apiVersion: 1, key: %q, name: %q, version: "1.0.0", author: {name: "Test"}, models: %s, fetchMode: "per_task", routes: %s}; export const native = { createJob(ctx) { const body = ctx.body.value || {}; if (body.reject === %q) throw new Error(%q + " rejects this job"); return {kind: "submit", model: body.model, requestBody: {by: %q, params: ctx.params}}; }, fetchJobs(ctx) { return {kind: "query", taskIds: ctx.body.value.ids}; }, jobStatus(ctx, task) { return {by: %q, tasks: Array.isArray(task) ? task.map(function(t) { return t.task_id; }) : [task.task_id], params: ctx.params, query: ctx.query, artifacts: ctx.artifacts}; }, %s }; export function buildSubmitRequest() { return {url: "https://example.com"}; } export function parseSubmitResponse() { return {taskId: "one"}; } export function buildQueryRequest() { return {url: "https://example.com"}; } export function parseTaskResult() { return {status: "SUCCESS"}; } `, key, key, models, routes, key, key, key, key, extra) } // sharedTaskPluginRouteRouter serves the plugins' route groups as the inner // plugin router does, each pinned to its candidates, for caller 7. func sharedTaskPluginRouteRouter(t *testing.T, plugins ...*jsplugin.LoadedPlugin) *gin.Engine { t.Helper() registry := jsplugin.NewRegistry() require.NoError(t, registry.ReplaceOverrides(plugins)) generation := registry.Generation() router := gin.New() for _, group := range generation.RouteGroups() { router.Handle(group.Method, group.GinPath, func(c *gin.Context) { c.Set(jsplugin.ContextKeyPinnedRoute, jsplugin.PinnedRoute{Generation: generation, Plugin: group.Bindings[0].Plugin, Route: group.Bindings[0].Route, Candidates: group.Bindings}) c.Set("id", 7) c.Next() }, PrepareTaskPluginRoute(), func(c *gin.Context) { pinned := c.MustGet(jsplugin.ContextKeyPinnedRoute).(jsplugin.PinnedRoute) candidates := make([]string, 0, len(pinned.Candidates)) for _, candidate := range pinned.Candidates { candidates = append(candidates, candidate.Plugin.Meta.Key) } c.JSON(http.StatusAccepted, gin.H{"plugin": pinned.Plugin.Meta.Key, "candidates": candidates, "model": c.GetString("resolved_task_model"), "request": c.MustGet("task_request")}) }) } return router } func performSharedRouteRequest(router http.Handler, method, path, body string) *httptest.ResponseRecorder { request := httptest.NewRequest(method, path, strings.NewReader(body)) if body != "" { request.Header.Set("Content-Type", "application/json") } recorder := httptest.NewRecorder() router.ServeHTTP(recorder, request) return recorder } // Every plugin declaring a submit path decodes the request; those accepting // the first acceptor's model stay candidates for distribution, and /{key} // keeps one plugin. func TestPrepareTaskPluginRouteDecodesEverySharedSubmitCandidate(t *testing.T) { alpha := compileTaskRoutePlugin(t, sharedRouteTestPluginSource("share-alpha", `["probe-model", "alpha-only"]`, `[{method: "POST", path: "/probe/v1/jobs/:kind", type: "submit", decode: "createJob", render: "jobStatus"}]`, "")) beta := compileTaskRoutePlugin(t, sharedRouteTestPluginSource("share-beta", `["probe-model"]`, `[{method: "POST", path: "/probe/v1/jobs/:variant", type: "submit", decode: "createJob", render: "jobStatus", models: ["probe-model"]}]`, "")) router := sharedTaskPluginRouteRouter(t, beta, alpha) for _, testCase := range []struct { name, path, body string status int want string }{ {name: "both accept", path: "/probe/v1/jobs/video", body: `{"model":"probe-model"}`, status: http.StatusAccepted, want: `{"plugin":"share-alpha","candidates":["share-alpha","share-beta"],"model":"probe-model","request":{"by":"share-alpha","params":{"kind":"video"}}}`}, {name: "first rejects", path: "/probe/v1/jobs/video", body: `{"model":"probe-model","reject":"share-alpha"}`, status: http.StatusAccepted, want: `{"plugin":"share-beta","candidates":["share-beta"],"model":"probe-model","request":{"by":"share-beta","params":{"variant":"video"}}}`}, {name: "route models scope each candidate", path: "/probe/v1/jobs/video", body: `{"model":"alpha-only"}`, status: http.StatusAccepted, want: `{"plugin":"share-alpha","candidates":["share-alpha"],"model":"alpha-only","request":{"by":"share-alpha","params":{"kind":"video"}}}`}, {name: "every candidate rejects", path: "/probe/v1/jobs/video", body: `{"model":"other-model"}`, status: http.StatusBadRequest, want: `{"code":"invalid_request","message":"share-alpha: model \"other-model\" is not served by this plugin; share-beta: model \"other-model\" is not allowed on this route","data":null}`}, {name: "the key address keeps one plugin", path: "/share-beta/probe/v1/jobs/video", body: `{"model":"probe-model","reject":"share-beta"}`, status: http.StatusBadRequest, want: `{"code":"invalid_request","message":"share-beta rejects this job","data":null}`}, } { t.Run(testCase.name, func(t *testing.T) { recorder := performSharedRouteRequest(router, http.MethodPost, testCase.path, testCase.body) assert.Equal(t, testCase.status, recorder.Code, recorder.Body.String()) assert.JSONEq(t, testCase.want, recorder.Body.String()) }) } } // A case variant or a channel alias of a declared model is matched on a // native path as on the host protocols: the task takes the declared spelling // or keeps the alias, and a rejection names the client's model. func TestPrepareTaskPluginRouteMatchesChannelAliases(t *testing.T) { setupTaskPluginRouteDB(t) require.NoError(t, model.DB.AutoMigrate(&model.Channel{})) require.NoError(t, model.DB.Create(&model.Channel{Type: constant.ChannelTypeTaskPlugin, Status: common.ChannelStatusEnabled, Key: "key", Models: "scoped-alias,outside-alias,shared-alias", ModelMapping: common.GetPointer(`{"scoped-alias":"scoped-model","outside-alias":"other-model","shared-alias":"shared-model"}`)}).Error) // Replace any cached alias view, so the routes' generation builds its own // from this channel. model.InitChannelCache() t.Cleanup(model.InitChannelCache) shared := `[{method: "POST", path: "/probe/v1/shared", type: "submit", decode: "createJob", render: "jobStatus"}]` router := sharedTaskPluginRouteRouter(t, compileTaskRoutePlugin(t, sharedRouteTestPluginSource("alias-scoped", `["scoped-model", "other-model"]`, `[ {method: "POST", path: "/probe/v1/scoped", type: "submit", decode: "createJob", render: "jobStatus", models: ["scoped-model"]}, {method: "POST", path: "/probe/v1/single", type: "submit", decode: "createJob", render: "jobStatus"}]`, "")), compileTaskRoutePlugin(t, sharedRouteTestPluginSource("alias-a", `["shared-model"]`, shared, "")), compileTaskRoutePlugin(t, sharedRouteTestPluginSource("alias-b", `["shared-model"]`, shared, ""))) type result struct { Plugin string `json:"plugin"` Candidates []string `json:"candidates"` Model string `json:"model"` Message string `json:"message"` } for _, testCase := range []struct { name, path, model string status int want result }{ {name: "alias on a route with models", path: "/probe/v1/scoped", model: "scoped-alias", status: http.StatusAccepted, want: result{Plugin: "alias-scoped", Candidates: []string{"alias-scoped"}, Model: "scoped-alias"}}, {name: "case variant on a route with models", path: "/probe/v1/scoped", model: "SCOPED-MODEL", status: http.StatusAccepted, want: result{Plugin: "alias-scoped", Candidates: []string{"alias-scoped"}, Model: "scoped-model"}}, {name: "alias on a route without models", path: "/probe/v1/single", model: "scoped-alias", status: http.StatusAccepted, want: result{Plugin: "alias-scoped", Candidates: []string{"alias-scoped"}, Model: "scoped-alias"}}, {name: "alias on a shared path", path: "/probe/v1/shared", model: "shared-alias", status: http.StatusAccepted, want: result{Plugin: "alias-a", Candidates: []string{"alias-a", "alias-b"}, Model: "shared-alias"}}, {name: "case variant on a shared path", path: "/probe/v1/shared", model: "Shared-Model", status: http.StatusAccepted, want: result{Plugin: "alias-a", Candidates: []string{"alias-a", "alias-b"}, Model: "shared-model"}}, {name: "alias of a model outside the route models", path: "/probe/v1/scoped", model: "outside-alias", status: http.StatusBadRequest, want: result{Message: `model "outside-alias" is not allowed on this route`}}, {name: "declared model outside the route models", path: "/probe/v1/scoped", model: "other-model", status: http.StatusBadRequest, want: result{Message: `model "other-model" is not allowed on this route`}}, {name: "alias of another plugin's model", path: "/probe/v1/single", model: "shared-alias", status: http.StatusBadRequest, want: result{Message: `model "shared-alias" is not served by this plugin`}}, {name: "case variant of another plugin's model", path: "/probe/v1/single", model: "Shared-Model", status: http.StatusBadRequest, want: result{Message: `model "Shared-Model" is not served by this plugin`}}, } { t.Run(testCase.name, func(t *testing.T) { recorder := performSharedRouteRequest(router, http.MethodPost, testCase.path, fmt.Sprintf(`{"model":%q}`, testCase.model)) assert.Equal(t, testCase.status, recorder.Code, recorder.Body.String()) var got result require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &got)) assert.Equal(t, testCase.want, got) }) } } // A query on a path several plugins declare is answered by the task's own // plugin; a task none of them owns is as unknown as a missing one. func TestPrepareTaskPluginRouteAnswersSharedQueriesByTaskOwner(t *testing.T) { setupTaskPluginRouteDB(t) alpha := compileTaskRoutePlugin(t, sharedRouteTestPluginSource("share-alpha", `["probe-model"]`, `[ {method: "GET", path: "/probe/v1/jobs/:task_id", type: "query", render: "jobStatus"}, {method: "POST", path: "/probe/v1/fetch", type: "dynamic", decode: "fetchJobs", render: "jobStatus"} ]`, "")) beta := compileTaskRoutePlugin(t, sharedRouteTestPluginSource("share-beta", `["probe-model"]`, `[ {method: "GET", path: "/probe/v1/jobs/:job_id", type: "query", taskIdParam: "job_id", render: "jobStatus"}, {method: "POST", path: "/probe/v1/fetch", type: "dynamic", decode: "fetchJobs", render: "jobStatus"} ]`, "")) insertTaskPluginRouteTask(t, &model.Task{TaskID: "task_alpha", UserId: 7, Platform: constant.TaskPlatform("share-alpha"), Status: model.TaskStatusInProgress}) insertTaskPluginRouteTask(t, &model.Task{TaskID: "task_beta", UserId: 7, Platform: constant.TaskPlatform("share-beta"), Status: model.TaskStatusInProgress}) insertTaskPluginRouteTask(t, &model.Task{TaskID: "task_foreign", UserId: 8, Platform: constant.TaskPlatform("share-beta"), Status: model.TaskStatusInProgress}) router := sharedTaskPluginRouteRouter(t, alpha, beta) for _, testCase := range []struct{ method, path, body, want string }{ {http.MethodGet, "/probe/v1/jobs/task_alpha", "", `{"by":"share-alpha","tasks":["task_alpha"],"params":{"task_id":"task_alpha"},"query":{},"artifacts":{}}`}, {http.MethodGet, "/probe/v1/jobs/task_beta", "", `{"by":"share-beta","tasks":["task_beta"],"params":{"job_id":"task_beta"},"query":{},"artifacts":{}}`}, {http.MethodPost, "/probe/v1/fetch", `{"ids":["task_beta"]}`, `{"by":"share-beta","tasks":["task_beta"],"params":{},"query":{},"artifacts":{}}`}, } { recorder := performSharedRouteRequest(router, testCase.method, testCase.path, testCase.body) assert.Equal(t, http.StatusOK, recorder.Code, testCase.path) assert.JSONEq(t, testCase.want, recorder.Body.String(), testCase.path) } missing := performSharedRouteRequest(router, http.MethodGet, "/probe/v1/jobs/task_missing", "") require.Equal(t, http.StatusNotFound, missing.Code) for _, testCase := range []struct{ method, path, body string }{ {http.MethodGet, "/share-alpha/probe/v1/jobs/task_beta", ""}, {http.MethodGet, "/probe/v1/jobs/task_foreign", ""}, {http.MethodPost, "/probe/v1/fetch", `{"ids":["task_foreign"]}`}, } { recorder := performSharedRouteRequest(router, testCase.method, testCase.path, testCase.body) assert.Equal(t, http.StatusNotFound, recorder.Code, testCase.path) assert.JSONEq(t, missing.Body.String(), recorder.Body.String(), testCase.path) assert.NotContains(t, recorder.Body.String(), "task_beta") assert.NotContains(t, recorder.Body.String(), "task_foreign") } } // A query route whose taskIdParam is not a path parameter reads the task id // from the query string, as MiniMax's v1 query does. func TestPrepareTaskPluginRouteReadsQueryStringTaskID(t *testing.T) { setupTaskPluginRouteDB(t) plugin := compileTaskRoutePlugin(t, sharedRouteTestPluginSource("query-string-probe", `["probe-model"]`, `[{method: "GET", path: "/probe/v1/query/job", type: "query", taskIdParam: "task_id", render: "jobStatus"}]`, "")) insertTaskPluginRouteTask(t, &model.Task{TaskID: "task_qs", UserId: 7, Platform: constant.TaskPlatform("query-string-probe"), Status: model.TaskStatusInProgress}) router := sharedTaskPluginRouteRouter(t, plugin) for _, path := range []string{"/probe/v1/query/job?task_id=task_qs", "/query-string-probe/probe/v1/query/job?task_id=task_qs"} { recorder := performSharedRouteRequest(router, http.MethodGet, path, "") require.Equal(t, http.StatusOK, recorder.Code, path) assert.JSONEq(t, `{"by":"query-string-probe","tasks":["task_qs"],"params":{},"query":{"task_id":["task_qs"]},"artifacts":{}}`, recorder.Body.String()) } missing := performSharedRouteRequest(router, http.MethodGet, "/probe/v1/query/job", "") assert.Equal(t, http.StatusBadRequest, missing.Code) assert.Contains(t, missing.Body.String(), "task_id is required") unknown := performSharedRouteRequest(router, http.MethodGet, "/probe/v1/query/job?task_id=task_unknown", "") assert.Equal(t, http.StatusNotFound, unknown.Code) } // A query presenter receives the artifacts of a finished task as the gateway // addresses Responses renderers get; a presenter with nothing to deliver for // the task answers like a missing one. func TestPrepareTaskPluginRouteGivesQueryPresentersArtifacts(t *testing.T) { setupTaskPluginRouteDB(t) previousSecret, previousAddress := common.CryptoSecret, system_setting.ServerAddress common.CryptoSecret, system_setting.ServerAddress = "artifact-test-secret", "https://gateway.example" t.Cleanup(func() { common.CryptoSecret, system_setting.ServerAddress = previousSecret, previousAddress }) listing := compileTaskRoutePlugin(t, sharedRouteTestPluginSource("artifact-probe", `["probe-model"]`, `[ {method: "GET", path: "/probe/v1/jobs/:task_id", type: "query", render: "jobStatus"}, {method: "GET", path: "/probe/v1/files/:task_id", type: "query", render: "fileOnly"} ]`, `fileOnly(ctx, task) { if (!ctx.artifacts.video) throw new Error("artifact_not_found"); return {url: ctx.artifacts.video.url}; },`)+ `export function listArtifacts(task) { return task.status === "SUCCESS" ? [{key: "video", type: "video", mimeType: "video/mp4"}] : []; } export function buildContentRequest() { return {url: "https://example.com/video.mp4"}; }`) silent := compileTaskRoutePlugin(t, sharedRouteTestPluginSource("artifact-silent", `["probe-model"]`, `[{method: "GET", path: "/silent/v1/jobs/:task_id", type: "query", render: "jobStatus"}]`, "")) for taskID, platform := range map[string]string{"task_done": "artifact-probe", "task_running": "artifact-probe", "task_silent": "artifact-silent"} { status := model.TaskStatus(model.TaskStatusSuccess) if taskID == "task_running" { status = model.TaskStatusInProgress } insertTaskPluginRouteTask(t, &model.Task{TaskID: taskID, UserId: 7, Platform: constant.TaskPlatform(platform), Status: status}) } router := sharedTaskPluginRouteRouter(t, listing, silent) done := performSharedRouteRequest(router, http.MethodGet, "/probe/v1/jobs/task_done", "") require.Equal(t, http.StatusOK, done.Code) videoURL, err := service.BuildTaskArtifactContentURL("task_done", "video") require.NoError(t, err) var presented map[string]any require.NoError(t, common.Unmarshal(done.Body.Bytes(), &presented)) assert.Equal(t, map[string]any{"video": map[string]any{"key": "video", "type": "video", "url": videoURL, "mimeType": "video/mp4"}}, presented["artifacts"]) file := performSharedRouteRequest(router, http.MethodGet, "/probe/v1/files/task_done", "") assert.JSONEq(t, fmt.Sprintf(`{"url":%q}`, videoURL), file.Body.String()) for _, path := range []string{"/probe/v1/jobs/task_running", "/silent/v1/jobs/task_silent"} { recorder := performSharedRouteRequest(router, http.MethodGet, path, "") require.Equal(t, http.StatusOK, recorder.Code, path) assert.Contains(t, recorder.Body.String(), `"artifacts":{}`, path) } notReady := performSharedRouteRequest(router, http.MethodGet, "/probe/v1/files/task_running", "") assert.Equal(t, http.StatusNotFound, notReady.Code) assert.Contains(t, notReady.Body.String(), `"code":"task_not_found"`) } // A native presenter answers with the vendor's own status by returning a reply // wrapper; anything else it returns is a 200 JSON body. func TestNativePresentersChooseTheVendorsStatus(t *testing.T) { setupTaskPluginRouteDB(t) plugin := compileTaskRoutePlugin(t, sharedRouteTestPluginSource("reply-probe", `["probe-model"]`, `[ {method: "GET", path: "/probe/v1/jobs/:task_id", type: "query", render: "pending"}, {method: "GET", path: "/probe/v1/moved/:task_id", type: "query", render: "moved"}, {method: "GET", path: "/probe/v1/broken/:task_id", type: "query", render: "broken"}, {method: "GET", path: "/probe/v1/plain/:task_id", type: "query", render: "jobStatus"} ]`, `pending(ctx, task) { return {__status: 202, __headers: {"retry-after": "5"}, __body: {status: "pending"}}; }, moved(ctx, task) { return {__status: 302, __headers: {Location: "https://cdn.example/v.mp4"}}; }, broken(ctx, task) { return {__status: 202, __headers: {"Set-Cookie": "a=b"}}; }, error(ctx, error) { return error.httpStatus === 404 ? {__status: 200, __body: {status: "not_found"}} : {message: error.message}; },`)) insertTaskPluginRouteTask(t, &model.Task{TaskID: "task_job", UserId: 7, Platform: "reply-probe", Status: model.TaskStatusInProgress}) router := sharedTaskPluginRouteRouter(t, plugin) pending := performSharedRouteRequest(router, http.MethodGet, "/probe/v1/jobs/task_job", "") assert.Equal(t, http.StatusAccepted, pending.Code) assert.Equal(t, "5", pending.Header().Get("Retry-After")) assert.JSONEq(t, `{"status":"pending"}`, pending.Body.String()) moved := performSharedRouteRequest(router, http.MethodGet, "/probe/v1/moved/task_job", "") assert.Equal(t, http.StatusFound, moved.Code) assert.Equal(t, "https://cdn.example/v.mp4", moved.Header().Get("Location")) assert.Empty(t, moved.Body.String()) broken := performSharedRouteRequest(router, http.MethodGet, "/probe/v1/broken/task_job", "") assert.Equal(t, http.StatusInternalServerError, broken.Code, "a wrapper may set only Location and Retry-After") assert.Empty(t, broken.Header().Get("Set-Cookie")) plain := performSharedRouteRequest(router, http.MethodGet, "/probe/v1/plain/task_job", "") assert.Equal(t, http.StatusOK, plain.Code) missing := performSharedRouteRequest(router, http.MethodGet, "/probe/v1/plain/task_missing", "") assert.Equal(t, http.StatusOK, missing.Code, "the error presenter answers with the vendor's status too") assert.JSONEq(t, `{"status":"not_found"}`, missing.Body.String()) } func TestNativeRouteCredentialRewritesOnlyTheVendorsOwnForm(t *testing.T) { for _, testCase := range []struct { name, plugin, path string headers map[string]string wantAuthorization string wantRequestURI string wantGoogleKeyRemains bool }{ {name: "Vidu Token form", plugin: "vidu", path: "/ent/v2/text2video", headers: map[string]string{"Authorization": "Token gateway-token"}, wantAuthorization: "Bearer gateway-token"}, {name: "Token form elsewhere", plugin: "doubao", path: "/api/v3/contents/generations/tasks", headers: map[string]string{"Authorization": "Token gateway-token"}, wantAuthorization: "Token gateway-token"}, {name: "Gemini header form", plugin: "google", path: "/v1beta/models/m:predictLongRunning", headers: map[string]string{"X-Goog-Api-Key": "gateway-token"}, wantAuthorization: "Bearer gateway-token"}, {name: "Gemini query form", plugin: "google", path: "/google/v1beta/models/m/operations/t?alt=json&key=gateway-token", wantAuthorization: "Bearer gateway-token", wantRequestURI: "/google/v1beta/models/m/operations/t?alt=json"}, {name: "Gemini forms do not replace Authorization", plugin: "google", path: "/v1beta/models/m:predictLongRunning?key=query-token", headers: map[string]string{"Authorization": "Bearer chosen-token", "X-Goog-Api-Key": "header-token"}, wantAuthorization: "Bearer chosen-token", wantRequestURI: "/v1beta/models/m:predictLongRunning"}, {name: "Gemini header form elsewhere", plugin: "vidu", path: "/vidu/ent/v2/text2video", headers: map[string]string{"X-Goog-Api-Key": "gateway-token"}, wantGoogleKeyRemains: true}, } { t.Run(testCase.name, func(t *testing.T) { router := gin.New() var captured *http.Request var query map[string][]string router.Use(NativeRouteCredential(testCase.plugin), func(c *gin.Context) { captured = c.Request requestContext, err := buildTaskPluginRouteRequest(c) require.NoError(t, err) query = requestContext.Query c.Status(http.StatusNoContent) }) router.NoRoute(func(c *gin.Context) {}) request := httptest.NewRequest(http.MethodGet, testCase.path, nil) for name, value := range testCase.headers { request.Header.Set(name, value) } router.ServeHTTP(httptest.NewRecorder(), request) require.NotNil(t, captured) assert.Equal(t, testCase.wantAuthorization, captured.Header.Get("Authorization")) assert.Equal(t, testCase.wantGoogleKeyRemains, captured.Header.Get("X-Goog-Api-Key") != "") assert.NotContains(t, query, "key") if testCase.wantRequestURI != "" { assert.Equal(t, testCase.wantRequestURI, captured.RequestURI) assert.NotContains(t, captured.URL.RawQuery, "key=") } }) } }